"""
Fichier: probastat.py
Auteur: Stéphane Pasquet
Web : https://mathweb.fr
Date: janvier 2022
Version : 0.2
"""

from random import random
from matplotlib.pyplot import bar, show

class binom:
    """
    Objet 'binom'
    Déclaration : X = binom(n,p)
    Méthodes:
        * X.esp(<r>)
        * X.var(<r>)
        * X.ecart(<r>)
        * X.proba(k,<r>)        -> retourne P(X = k)
        * X.proba_cdf(k,<r>)    -> retourne P(X ≤ k)
        * X.proba_icdf(p)       -> retourne la plus grande valeur de h telle que P(X ≤ h) ≤ p
        * X.simul(k, <couleur fond>, <couleur bord>) -> effectue k simulations et affiche le diagramme en barre
        * X.distrib( <couleur fond>, <couleur bord> ) -> distribution de X
        * X.fluctuations( <seuil = 0.95>) -> .tableau, .intervalle --> 'a' et 'b' tels que P(a ≤ X ≤ b) > seuil
        
        --> 'r' : nombre de chiffres après la virgule
    """
    def __init__(self,n,p):
        self.n = n
        self.p = p
        self.fact = [ 1 ]
        for i in range(1,n+1):
            self.fact.append( self.fact[i-1] * i )
        
    def esp(self , r = None):
        if r != None:
            return round(self.n * self.p , r)
        else:
            return self.n * self.p
    
    def var(self , r = None):
        if r != None:
            return round(self.n * self.p * (1 - self.p) , r)
        else:
            return self.n * self.p * (1 - self.p)
    
    def ecart(self , r = None):
        if r != None:
            return round(self.var()**0.5 , r)
        else:
            return self.var()**0.5
    
    def proba(self,k,r = None):
        if r == None:
            return ( self.fact[self.n] / ( self.fact[k] * self.fact[self.n - k] ) ) * self.p**k * (1 - self.p)**(self.n-k)
        else:
            return round( ( self.fact[self.n] / ( self.fact[k] * self.fact[self.n - k] ) ) * self.p**k * (1 - self.p)**(self.n-k) , r)
    
    def proba_cdf(self,k,r = None):
        cdf = 0
        for i in range(k+1):
            cdf += self.proba(i)
        
        if r == None:
            return cdf
        else:
            return round(cdf , r)
        
    def proba_cdfsup(self,k,r = None):
        if r == None:
            return 1-self.proba_cdf(k-1)
        else:
            return round(1-self.proba_cdf(k-1),r)
        
    def proba_icdf(self , pr):
        h = 0
        while self.proba_cdf(h) < pr :
            h += 1
        return h
    
    def simul(self , k , c = None , e = None):
        LX = [ i for i in range( self.n + 1) ]
        LY = [ 0 ] * (self.n + 1)
        for i in range(k):
            p = self.proba_icdf( random() )
            LY[p] += 1
        
        Y = [i/k for i in LY]
        
        if c != None:
            if e != None:
                bar(LX,Y,color=c,edgecolor=e)
            else:
                bar(LX,Y,color=c)
        else:
            bar(LX,Y)
        show()
    
    def distrib(self , c = None , e = None):
        LX = [ i for i in range( self.n + 1) ]
        LY = [ self.proba(k) for k in range( self.n + 1) ]
        
        if c != None:
            if e != None:
                bar(LX,LY,color=c,edgecolor=e)
            else:
                bar(LX,LY,color=c)
        else:
            bar(LX,LY)
        show()
    
    def fluctuations(self , seuil = 0.95):               
        return interv(self.n, self.p, seuil)

class interv:
    def __init__(self,n,p,seuil):
        V = binom(n,p)
        L = [ V.proba_cdf(k) if V.proba_cdf(k)<1 else 1 for k in range(n+1) ]
        self.tableau = '{:>5} | {:^20}\n{:5} | {:>20}'.format('k' , 'P(X ≤ k)' , 0 , V.proba_cdf(0))
            
        a = 0
        
        for k in range(1,n+1):
            if L[k-1] < (1-seuil)/2 and L[k] >= (1-seuil)/2:
                a = k
            elif L[k-1] <= (1+seuil)/2 and L[k] > (1+seuil)/2:
                b = k
                
            self.tableau += '\n{:5} | {:>20}'.format( k , L[k] )
            
        self.intervalle = a, b
        
    
class variable:
    """
    Objet 'variable'
    Déclaration : X = variable(liste valeurs , liste effectifs ou probabilités)
    Méthodes:
        * X.esp(<r>)
        * X.var(<r>)
        * X.ecart(<r>)
        * X.mediane(k,<r>)        -> retourne P(X = k)
        
        --> 'r' : nombre de chiffres après la virgule
    """
    def __init__(self,X,P):
        self.X = X
        self.P = P
        
    def esp(self,r = None):
        if (len(self.X) != len(self.P)):
            return False
        
        e = 0
        for i in range( len(self.X) ):
            e += self.X[i] * self.P[i]
        
        e /= sum(self.P)
        
        if r != None:
            return round(e , r)
        else:
            return e
    
    def var(self, r = None):
        if (len(self.X) != len(self.P)):
            return False
        
        m = self.esp()
        V = [ ( self.X[i] - m )**2 for i in range( len(self.X) ) ]
        Y = variable( V , self.P )
        return Y.esp(r)
    
    def ecart(self , r = None):
        if r != None:
            return round(self.var()**0.5,r)
        else:
            return self.var()**0.5
        
    def mediane(self):
        if (len(self.X) != len(self.P)):
            return False
        
        total = 0
        for i in range( len(self.X) ):
            total += self.P[i]
            if total > sum(self.P)/2:
                return self.X[i]
            elif total == sum(self.P)/2:
                return ( self.X[i] + self.X[i+1] ) / 2
            

class loigeom:
    """
    Objet 'loigeom'
    Déclaration : X = loigeom(p,n) [tronquée] ou X = loigeom(p)
    Méthodes:
        * X.esp(<r>)
        * X.var(<r>)
        * X.ecart(<r>)
        * X.proba(k,<r>)        -> retourne P(X = k)
        * X.proba_cdf(k,<r>)    -> retourne P(X ≤ k)
        * X.proba_icdf(p)       -> retourne la plus grande valeur de h telle que P(X ≤ h) ≤ p
        * X.simul(k, <couleur fond>, <couleur bord>) -> effectue k simulations et affiche le diagramme en barre
        * X.distrib( <couleur fond>, <couleur bord> ) -> distribution de X
        * X.fluctuations( <seuil = 0.95>) -> .tableau, .intervalle --> 'a' et 'b' tels que P(a ≤ X ≤ b) > seuil
        
        --> 'r' : nombre de chiffres après la virgule
    """
    def __init__(self,p,n=None):
        self.n = n
        self.p = p
        
    def esp(self,r=None):
        if r != None:
            return round(1/self.p, r)
        else:
            return 1/self.p
    
    def var(self,r=None):
        if r != None:
            return round((1-self.p)/(self.p*self.p),r)
        else:
            return (1-self.p)/(self.p*self.p)
        
    def ecart(self,r=None):
        if r!= None:
            return round(self.var()**0.5,r)
        else:
            return self.var()**0.5
        
    def proba(self,k,r=None):
        if r != None:
            return round((1-self.p)**(k-1)*self.p,r)
        else:
            return (1-self.p)**(k-1)*self.p
        
    def proba_cdf(self,k,r = None):
        cdf = 0
        for i in range(k+1):
            cdf += self.proba(i)
        
        if r == None:
            return cdf
        else:
            return round(cdf , r)
    
    def proba_icdf(self , pr):
        h = 0
        while self.proba_cdf(h) < pr :
            h += 1
        return h
    
    def simul(self , k , c = None , e = None):
        if self.n != None:
            LX = [ i for i in range( self.n + 1) ]
            LY = [ 0 ] * (self.n + 1)
        else:
            LX = [ i for i in range( int(8/self.p) + 1) ]
            LY = [ 0 ] * (int(8/self.p) + 1)
        for i in range(k):
            p = self.proba_icdf( random() )
            LY[p] += 1
        
        Y = [i/k for i in LY]
        
        if c != None:
            if e != None:
                bar(LX,Y,color=c,edgecolor=e)
            else:
                bar(LX,Y,color=c)
        else:
            bar(LX,Y)
        show()
    
    def distrib(self , c = None , e = None):
        if self.n != None:
            LX = [ i for i in range( self.n + 1) ]
            LY = [ self.proba(k) for k in range( self.n + 1) ]
        else:
            LX = [ i for i in range( int(8/self.p) + 1) ]
            LY = [ self.proba(k) for k in range( int(8/self.p) + 1) ]
        
        if c != None:
            if e != None:
                bar(LX,LY,color=c,edgecolor=e)
            else:
                bar(LX,LY,color=c)
        else:
            bar(LX,LY)
        show()