#!/usr/local/bin/pythonw
# coding: iso-8859-1
# Probability.py, copyright © 2009 by Amos Newcombe.
# 
# This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along with this program.  If not, see <http://www.gnu.org/licenses/>.

'''the base class for probability laws, and a subclass that con apply a  linear transformation to another probability law
'''

from random import random
from math import sqrt

class Probability: # tag prob
	'the base class for probability laws'
	
	def __str__(self): return self.__class__.__name__
	def __repr__(self): return '%s()' % self
	def html(self): return str(self)

	def Distribution(self, xX):
		'the distribution function of the probability law'
		raise NotImplementedError, '%s.Distribution()' % self.__class__.__name__

	def Density(self, xX, xDelta=5e-7):
		'''the density function of the probability law, equal to the derivative of the distribution function
		
		By default, numerically differentiate the distribution function.
		'''
		return (self.Distribution(xX+xDelta) - self.Distribution(xX-xDelta)) / (xDelta*2.)
	
	def DensityTest(self, xX, **dArgs):
		'''Compare the Density() with the numerically differentiated Distribution().
		
		This doesn't make much sense in the abstract base class. But in general, the Density() method will be redefined, and this will be a real test.
		'''
		return self.Density(xX) - Probability.Density(self, xX, **dArgs)
	
	def InverseDistribution(self, xY, xX=0., xEpsilon=1e-6):
		'''the inverse of the distribution function of the probability law
		
		By default, invert using Newton's method.
		
		Note that using a numerically calculated derivative, as in the default Density(), with Newton's method as here, is numerically problematic. One or the other, or both, should be redefined in any (non-abstract) extension. Fortunately, alternate expressions are available for commonly-used distributions.
		'''
		for i in range(32):
			d = xY - self.Distribution(xX)
			if abs(d) < xEpsilon: return xX
			xX += d / self.Density(xX)
		raise ValueError, 'Can\'t invert %s distribution' % self.__class__.__name__
	
	def InverseTest(self, xY, **dArgs):
		'''Run a y value through InverseDistribution and then through Distribution, and compare the results.
		'''
		return self.Distribution(self.InverseDistribution(xY, **dArgs)) - xY
	
	def Sample(self):
		'''Return a random sample distributed according to the probability law.
		
		By default, pass a uniform random sample to InverseDistribution().
		
		Similarly, passing a sample through Distribution() will give you a uniform sample on [0, 1].
		'''
		return self.InverseDistribution(random())
	
# 	def SampleTest(self, cSample, cRep):
# 		'''Test the Sample() function to see if the samples follow its Distribution() function.
# 		
# 		Collect a number of samples, each of a given size. Use the Kolmogorov-Smirnov-Kuiper test on each sample to get a significance. These numbers should obey a uniform distribution. Use another KSKp test to see if they do.
# 		'''
# 		r = range(cSample)
# 		return KSKp1([
# 			KSKp1([self.Sample() for x in r], self).Calculate()['sig'] 
# 			for i in range(cRep)
# 		], Uniform()).Calculate()['sig']
	
	def getMean(self): return None
	def getVariance(self): return None
	def getStDev(self):
		s2 = self.getVariance()
		return None if s2 == None else sqrt(s2)
	
	def getMedian(self): return self.InverseDistribution(0.5)
	def getIQR(self): return self.InverseDistribution(0.75) - self.InverseDistribution(0.25)

class Affine(Probability):
	'''the probability law governing Z*s + m, if Z obeys probBase.
	'''

	def __init__(self, probBase, xM=0., xS=1.):
		self.probBase = probBase
		self.xM = float(xM)
		self.xS = float(xS)
	
	def __repr__(self):
		return '%s(%s,%s,%s)' % (self.__class__.__name__, `self.probBase`, self.xM, self.xS)
	
	def Standardize(self, xX): return (xX - self.xM) / self.xS
	def Affine     (self, xZ): return (xZ * self.xS) + self.xM
	
	def Distribution(self, x):
		return self.probBase.Distribution(self.Standardize(x))
	
	def Density(self, x):
		return self.probBase.Density(self.Standardize(x)) / self.xS
	
	def InverseDistribution(self, y):
		return self.Affine(self.probBase.InverseDistribution(y))
	
	def Sample(self):
		return self.Affine(self.probBase.Sample())
	
	def getMean(self): return self.Affine(self.probBase.getMean())
	def getVariance(self): return self.probBase.getVariance() * self.xS * self.xS