#!/usr/local/bin/pythonw
# coding: iso-8859-1
# 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 Beta distribution and others derived from it: FRatio and Order
'''

from math import sqrt
from Probability import Probability
from mathlib.Beta import Beta, IBeta, InverseIBeta

class BetaDist(Probability):
	
	def __init__(self, xA, xB):
		#Probability.__init__(self) # doesn't exist.
		self.xA = xA
		self.xB = xB
	
	def __repr__(self): return '%s(%s, %s)' % (self.__class__.__name__, self.xA, self.xB)
	
	def Distribution(self, x):
		return IBeta(self.xA, self.xB, x)
	
	def Density(self, x):
		if (x == 0.) or (x == 1.): return 0.
		return pow(x, self.xA-1.) * pow(1.-x, self.xB-1.) / Beta(self.xA, self.xB)
	
	def InverseDistribution(self, y, xEpsilon=1e-7):
		return InverseIBeta(self.xA, self.xB, y, xEpsilon)

class StudentT(BetaDist):
	'''Student\'s t distribution as a variant on the Beta distribution
	'''
	
	def __init__(self, nDF):
		self.nDF = nDF
		BetaDist.__init__(self, nDF/2., 0.5)
	
	def __repr__(self): return '%s(%s)' % (self.__class__.__name__, self.nDF)
	
	def Distribution(self, x):
		a, b = self.xA, self.xB
		z = 0.5 * BetaDist.Distribution(self, a/(a + b*x*x))
		if x > 0.: z = 1-z
		return z
	
	def Density(self, x):
		n = self.nDF
		return 1.0 / (Beta(n/2.0, 0.5) * sqrt(n) * pow(1+x*x/n, (n+1)/2.0))
	
	def InverseDistribution(self, y, xEpsilon=1e-7):
		nSign = -1
		if 0.5 < y: y, nSign = 1-y, +1
		return nSign * sqrt(self.nDF * (-1+1/BetaDist.InverseDistribution(self, 2*y)))

class FRatio(BetaDist):
	'''Fisher\'s F distribution as a variant on the Beta distribution
	'''
	
	def __init__(self, nM, nN):
		BetaDist.__init__(self, nN/2., nM/2.)
	
	def Distribution(self, x):
		a, b = self.xA, self.xB
		return 1. - BetaDist.Distribution(self, a/(a + b*x))
	
	def Density(self, x):
		if x <= 0.: return 0.
		a, b = self.xA, self.xB
		y = 1.+x*b/a
		return pow(b/a, b) * pow(x, b-1.) / pow(y, a+b) / Beta(a, b)
	
	def InverseDistribution(self, y, xEpsilon=1e-7):
		a, b = self.xA, self.xB
		return (a/b) * ((1/BetaDist.InverseDistribution(self, 1.-y, xEpsilon)) - 1)

class Order(BetaDist):
	'''Given a base probability law governing a random sample of size n, this is the probability governing the r-th largest member of that sample, for r = 1, ..., n.
	'''

	def __init__(self, prob, nR, nN):
		self.probBase = prob
		self.nR = nR
		self.nN = nN
		BetaDist.__init__(self, nR, nN-nR+1)
	
	def Distribution(self, x):
		return BetaDist.Distribution(self, self.probBase.Distribution(x))
	
	def Density(self, x):
		return BetaDist.Density(self, self.probBase.Distribution(x)) * self.probBase.Density(x)
	
	def Sample(self):
		lx = [self.probBase.Sample() for n in range(self.nN)]
		lx.sort()
		return lx[self.nR-1]
	
	def __repr__(self):
		return '%s(%s,%s,%s)' % (self.__class__.__name__, `self.probBase`, self.nR, self.nN)
	def __str__(self): return `self`

if __name__ == '__main__':
	from DistributionTests import Tests
	#Tests(BetaDist(2., 3.), 0., 1., .05)
	Tests(FRatio(1., 9.), 0., 3., 0.1, isLogDim=0, nSelect=0x1E)