#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Nov 13 17:42:33 2025

This python script produces the graph displayed in the Section 3.3 of the paper, by implementing the 'approximation of binomial mechanism'

Remark that the code is made for a statistical model given by a Gaussian translation model (the function $\nu_+$ depends on this choice). Also, the code relies on the fact that the law of the model is symmetric w.r.t $0$, and the true value of the parameter must be $theta_0=0$.
Changing the code to deal with a statistical model given by the translation of a symmetric distribution could be made with few modifications. But more general situations would necessitate to change most of the code.

@author: https://arxiv.org/abs/2603.27572  C. Amorino, A. Gloter
"""

import numpy as np
from numba import jit, float64, void, prange, vectorize, guvectorize,int64,types
from numba.experimental import jitclass
import matplotlib.pyplot as plt
from scipy.stats import norm
from scipy import integrate, special
import math
from scipy.optimize import minimize, Bounds
from multiprocessing import Pool, cpu_count
import sys


text_sauv_titre='nu_gauss_0_1'
text_sauv='nu_gauss_0_1'
alpha=np.array([ 0.2, 0.5 , 1, 2, 4])
c_list=np.linspace(0.01,0.99,40)
#alpha_list=np.ones_like(c_list)*alpha
theta_0=0  # seule valeur possible est zero

bounds=Bounds(-4,4)
theta_start=1  # Pour la recherche d'optimum

Nb_data=1000
Nb_Iter=2000

precision=5*10**4
left,right= -100 , 100



@jit()
def abs_p_dot(x) :
    #absolute value must be symmetric w.r.t. zero with value zero possible at zero
    #x=x-1
    #return (np.exp(-np.abs(x))/2)
    #return(1 /  (np.pi * (1+x*x) ) )
    return (np.exp(-x*x/(2))/np.sqrt(2*np.pi))
#    return (np.abs(x)*np.exp(-x*x/(2))/np.sqrt(2*np.pi))

constante_normalisation=integrate.quad(abs_p_dot,-np.inf,np.inf)[0]

@jit()
def nu_plus(x) :
    return( 2* abs_p_dot(x) / constante_normalisation ) 



@vectorize([float64(float64,float64)])
def one_over_nu_ODE(x,t) :
    resu=nu_plus(x)
    if ( (np.abs(x) < 10**(-6)) & (resu < 10**(-6) ) ) : resu=np.float64(10**6)
    else : resu=1/nu_plus(x)
    return (resu)

def repartition_nu_plus (u: np.float64) -> np.float64 :
    F0=integrate.quad(nu_plus,0,u)[0]
    return(F0)

def inverse_repartition_nu_plus(u : np.float64) -> np.float64:
    #F0=integrate.quad(nu,-np.inf,0)[0]
    Finv=integrate.odeint(one_over_nu_ODE,0,np.array([0,u]))
    return(Finv[1][0])



class XI_plus_lookup:
    
    def __init__(self,precision=10**4,left=None,right=None):
        self.precision=precision
        self.table_F=np.zeros(precision)          
        self.table_invF=np.zeros(precision)   
        if (left==None) : self.left=inverse_repartition_nu_plus(1/precision)
        else : self.left=left
        if (right==None): self.right=inverse_repartition_nu_plus(1-1/precision)
        else : self.right=right
        for i in range(precision):
            self.table_invF[i]=inverse_repartition_nu_plus((i+1)/(precision+1))
        for i in range(precision):
            self.table_F[i]=repartition_nu_plus( (i / self.precision) * (self.right-self.left)  + self.left    )
            
            
    def value(self, x):
        if ((x > self.left) & (x < self.right)) :
            index=self.precision *  (x-self.left) / (self.right-self.left) 
            val=self.table_F[ int( index)  ]
            return(val)
        else :
            return(repartition_nu_plus(x))
    
    def inv_value(self, x):
        if ((x > 0) & (x < 1)) :
            index=self.precision *  x 
            val=self.table_invF[ int(index)  ]
            return(val)
        elif (x <=0 ):
            return(-np.inf)
        else :
            return(+np.inf)
        
    def l(self,x,c):
        if (x <= self.inv_value(c)) : return(0)
        else : return (self.inv_value( (-c/2+ self.value(x) ) ) )

    def r(self,x,c):
        if (x >= self.inv_value(1-c)) : return(np.inf)
        else : return (self.inv_value( (c/2+ self.value(x) ) ) )
        
    def inv_l(self,x,c):
            if (x <= self.inv_value(c/2)) : return ( self.inv_value(c))
            elif (x < self.inv_value(1-c/2)) : 
                return ( (self.inv_value( (c/2+ self.value(x) ) ) ) ) 
            else : return(np.inf)                        
            
    def inv_r(self,x,c):
            if (x <= self.inv_value(c/2)) : return (0)
            elif (x < self.inv_value(1-c/2)) : 
                return ( (self.inv_value( (-c/2+ self.value(x) ) ) ) ) 
            else : return(self.inv_value(1-c))            
    
    def density_q (self,c,alpha,x,y):
            indi=1
            if ( ( self.l(y,c) <= x ) & (x<=self.r(y,c))) : indi=np.exp(alpha)
            densi=indi * nu_plus(y) / ( 1 + c * (np.exp(alpha)-1) )
            return(densi)
    
    def density_q_v2 (self,c,alpha,x,y):
            indi=1
            if ( ( self.inv_r(x,c) <= y ) & (y<=self.inv_l(x,c))) : indi=np.exp(alpha)
            densi=indi * nu_plus(y) / ( 1 + c * (np.exp(alpha)-1) )
            return(densi)
        
    
    def plot_densi_q(self,c,alpha,x,y,version=1) :
            if (version==1) : densite=self.density_q
            else : densite=self.density_q_v2 
            ordo=np.zeros(len(y))
            for i in range(len(y)):
                ordo[i]=densite(c,alpha,x,y[i])
            plt.plot(y,ordo)
    
    def check_density_q (self,c,alpha,x,version) :
            if (version==1) : densite=self.density_q
            else : densite=self.density_q_v2 
            def to_integrate(y,c,alpha,x) :
                return ( densite(c,alpha,x,y) )
            area=integrate.quad(to_integrate,0,np.inf,args=(c,alpha,x))[0]
            return(area)

    def plot_l_r(self,c,reso=300) :
             abscisse=np.linspace(self.inv_value(c), self.inv_value(1-c),reso) 
 #            abscisse=np.linspace(0, 10, reso)  
             l,r=np.zeros(reso),np.zeros(reso)
             for i in range(reso):
                 l[i] , r[i] =self.l(abscisse[i],c) , self.r(abscisse[i],c) 
             plt.plot(abscisse,l)
             plt.plot(abscisse,r)
             plt.show()
             return


def tirage_selon_nu_plus():
    return(inverse_repartition_nu_plus(np.random.uniform()))


def randomisation(X,c,alpha,XI) :
    if (X<0) :
        X=-X
        retourne=-1
    else : retourne=1
    Yc=[]
    #Yc=np.zeros_like(c)
    for i in range(len(c)) :
       inv_l,inv_r= XI.inv_l(X,c[i]),XI.inv_r(X,c[i])
       cont=1
       while (cont==1):
           Y=  tirage_selon_nu_plus()*np.sign(np.random.uniform(-1,1))
           U=np.random.uniform(0,1)
           if ( (Y > inv_r) and (Y < inv_l) ) :
               nume=np.exp(alpha)
           else:
               nume=1
           if ( nume / np.exp(alpha)  > U) : cont=0    
           #if ( nume / ( ( 2*np.exp(alpha) ) / (1 +  (np.exp(alpha) - 1)*c[i] ) )  > U) : cont=0
       Y=Y*retourne
       #Yc[i]=Y
       if (c[i] > 0.5) :
           if ( (XI.inv_value(1-c[i]) < np.abs(Y)) and (XI.inv_value(c[i]) > np.abs(Y)) ):
               if (Y>0) : Yc.append('F+')
               else : Yc.append('F-')
           else : Yc.append(Y)
       else : Yc.append(Y)
    return(Yc)

@jit()
def randomisation_pur_Bino(X,alpha) :
    s=np.exp(alpha)/(1+np.exp(alpha))
    U=np.random.uniform(0,1)
    if ( X > 0 )  : 
            if(U>s) : Y='F-'
            else : Y='F+'
    else :
            if(U>s) : Y='F+'
            else : Y='F-'
    return(Y)   


@jit()
def generate_Z_X_pur_Bino(X,alpha) :
    data_Z=[]
    for i in range(len(X)):
        data_Z.append( randomisation_pur_Bino(X[i],alpha) )
    return(data_Z)



#@jit(float64(float64,))
def minus_Log_vraisemblance_c_fixed_Bino(theta,data_Z,alpha):
    LL=0
    for j in range(len(data_Z)):
        if (data_Z[j] == 'F+') : delta_LL= 1 + ( np.exp(alpha) - 1 ) * (1- F (0, theta ) ) 
        else :  delta_LL= 1 + ( np.exp(alpha) - 1 ) * F (0, theta )   
        LL+=np.log(delta_LL) - np.log ( np.exp(alpha) + 1 )
    return(-LL)



def MLE_computation_Bino (data_Z, alpha, theta_start, bounds ):
    res_min=minimize(minus_Log_vraisemblance_c_fixed_Bino, theta_start, \
                      args=(data_Z,alpha),  bounds=bounds)
    return(res_min.x)


    
def MC_estimation_Bino( Nb_data, Nb_iteration ,alpha, theta_start, bounds) :
    np.random.seed()
    resultat=np.zeros(Nb_iteration) 
    for i in range(Nb_iteration):
        X_data=np.random.normal(0,1,Nb_data)
        Z_data =generate_Z_X_pur_Bino(X_data, alpha)  
        MLE=MLE_computation_Bino(Z_data ,alpha , theta_start, bounds)
        resultat[i]=MLE
    return(resultat)
    
 
#@jit()
def plot_MLE_Bino(nb_data, alpha, theta_0,theta_min,theta_max,nb_points_grid ) :
    X_data=generate_data_X(nb_data,theta_0)
    data_Z =generate_Z_X_pur_Bino(X_data, alpha)
    LL_X=np.zeros(nb_points_grid)
    LL_Z=np.zeros(nb_points_grid)
    abscice=np.linspace(theta_min,theta_max,nb_points_grid)
    for i in range(nb_points_grid) :  
        LL_Z[i] = -minus_Log_vraisemblance_c_fixed_Bino(abscice[i],data_Z, alpha)
        LL=0
        for j in range(nb_data):
            LL+= np.log (p (X_data[j],abscice[i]))
        LL_X[i]=LL
    plt.plot( abscice,LL_X)
    plt.plot( abscice,LL_Z)
    return(LL_Z,LL_X)    
   
    
"""     
XI_table=XI_plus_lookup(10**3)
Nb_data=100
Xdata=np.random.normal(0,1,size=Nb_data)
Z=[]
for i in range(Nb_data) :
#    for j in range (len(c_list)):
    Z.append ( randomisation(Xdata[i],c_list,alpha,XI_table) )
"""    

 

"ici modele Gaussien"
def F(x,theta):
    "Fonction de répartition du modele stat"
    return(special.ndtr(x-theta))
    #return(norm.cdf(x-theta))
    
def p(x,theta):
    "Fonction de répartition du modele stat"
    return(np.exp( - (x-theta)**2 / 2 ) / np.sqrt( 2 * np. pi)  )
    #return(norm.cdf(x-theta))
  
    
def partial_F_theta ( x, theta) :
    return( - np.exp( - (x-theta)**2 / 2 ) / np.sqrt( 2 * np. pi)    )    
    

def generate_data_X(n,theta_0):
    return(np.random.normal(theta_0,1,size=n))    

def generate_Z_X(X,c,alpha,XI) :
    data_Z=[]
    for i in range(len(X)):
        data_Z.append( randomisation(X[i],c,alpha,XI) )
    data_Z_transpose=[]
    for j in range(len(c)):
        Zj=[]
        for i in range(len(X)):
            Zj.append(data_Z[i][j])
        data_Z_transpose.append(Zj)
    return(data_Z_transpose,data_Z)

    
 

def minus_Log_vraisemblance_c_fixed(theta,data_Z, c, alpha,XI):
    LL=0
    k=( 2 + (np.exp(alpha)-1) * c ) 
    if (c < 0.5) :
       for j in range(len(data_Z)):
            Zj_abs=np.abs(data_Z[j])
            lY , rY = XI.l(Zj_abs,c) , XI.r(Zj_abs,c)
            if (data_Z[j] < 0) : rY , lY = -lY , -rY
            #print(lY,rY)
            s=(1+ (np.exp(alpha) - 1) * (F(rY,theta)-F(lY,theta))) / k
            LL+=np.log(s) + np.log ( nu_plus(Zj_abs) )
            
    if (c >= 0.5) :
       for j in range(len(data_Z)):
            if ( data_Z[j] == 'F+' ) : 
                deltaLL = (np.exp(alpha)-1) * (2*c - 1) / k * (1- F (0, theta ) )
                deltaLL += (2*c-1) / k
            elif ( data_Z[j] == 'F-' ) :
                deltaLL = (np.exp(alpha)-1) * (2*c - 1) / k * F (0, theta  )
                deltaLL += (2*c-1) / k
            elif  ( data_Z[j] >= 0 )  :
                lY , rY = XI.l(data_Z[j],c) , XI.r(data_Z[j],c)
                deltaLL = (1+ (np.exp(alpha) - 1) * (F(rY,theta)-F(lY,theta))) / k
                deltaLL *= nu_plus(data_Z[j])
            else :
                lY , rY = -XI.r(-data_Z[j],c) , -XI.l(-data_Z[j],c) 
                deltaLL = (1+ (np.exp(alpha) - 1) * (F(rY,theta)-F(lY,theta))) / k
                deltaLL *= nu_plus(-data_Z[j])
            LL+=np.log(deltaLL)             
    return(-LL)     



def MLE_computation (data_Z, c_list ,alpha, theta_start, bounds, XI):
     MLE=np.zeros(len(c_list))
     for i in range(len(c_list)):
            res_min=minimize(minus_Log_vraisemblance_c_fixed, theta_start, \
                     args=(data_Z[i],c_list[i],alpha,XI),  bounds=bounds)
            MLE[i]=res_min.x[0]
     return(MLE)
        
   
 
def MC_estimation(theta_0, Nb_data, Nb_iteration,c_list ,alpha, theta_start, bounds, left=None, right=None, precision=10**5) :
    np.random.seed()
    resultat=np.zeros((Nb_iteration,len(c_list))) 
    XI=XI_plus_lookup(precision,left,right)
    for i in range(Nb_iteration):
        X_data=generate_data_X(Nb_data,theta_0)
        Z_data =generate_Z_X(X_data, c_list, alpha,XI)[0]     
        MLE=MLE_computation (Z_data, c_list ,alpha , theta_start, bounds, XI)
        resultat[i,:]=MLE[:]
    return(resultat)
    
  
def plot_MLE(nb_data,c, alpha, theta_0,theta_min,theta_max,nb_points_grid,precision=10**4,left=None,right=None,XI=None) :
    X_data=generate_data_X(nb_data,theta_0)
    if (XI==None) : XI=XI_plus_lookup(precision,left,right)
    data_Z =generate_Z_X(X_data, np.array([c]) ,alpha,XI)[0]
    LL_X=np.zeros(nb_points_grid)
    LL_Z=np.zeros(nb_points_grid)
    abscice=np.linspace(theta_min,theta_max,nb_points_grid)
    for i in range(nb_points_grid) :  
        LL_Z[i] = -minus_Log_vraisemblance_c_fixed(abscice[i],data_Z[0], c, alpha,XI)
        LL=0
        for j in range(nb_data):
            LL+= np.log (p (X_data[j],abscice[i]))
        LL_X[i]=LL
    plt.plot( abscice,LL_X)
    plt.plot( abscice,LL_Z)
    return(LL_Z,LL_X)
  


def MC_estimation_parallel(theta_0, Nb_data, Nb_iteration,c_list ,alpha , theta_start, bounds, precision, left=None, right=None, Nb_worker=0) :
    if (Nb_worker == 0) : Nb_worker=int(cpu_count()*0.8) 
    with Pool(processes=Nb_worker) as pool: 
        parameter=[(theta_0, Nb_data, int(Nb_iteration/Nb_worker), c_list ,alpha , theta_start, bounds, left, right, precision)]*Nb_worker
        resultat=pool.starmap(MC_estimation,parameter)
    return(np.concatenate(resultat))

#MC_resu=MC_estimation_parallel(theta_0, Nb_data, 1000, c_list, alpha, theta_start, bounds, 10**6)
   


def MC_estimation_min_std(resultat):
    #Nb_iter,Nb_c=resultat.shape
    moyenne=resultat.mean(0)
    std=resultat.std(0)
    return(moyenne,std)
 


def Fisher_i(x_0,alpha,theta_0, c,XI) :
    x_0=np.abs(x_0)
    num=(partial_F_theta(XI.r(x_0,c),theta_0) - partial_F_theta(XI.l(x_0,c),theta_0))**2
    deno=1 + ( np.exp(alpha) - 1 ) * ( F(XI.r(x_0,c),theta_0) -F(XI.l(x_0,c),theta_0) )
    return (num / deno)     

def Fisher_i_times_density(x_0,alpha,theta_0, c,XI) :
    x_0=np.abs(x_0)
    num=(partial_F_theta(XI.r(x_0,c),theta_0) - partial_F_theta(XI.l(x_0,c),theta_0))**2
    deno=1 + ( np.exp(alpha) - 1 ) * ( F(XI.r(x_0,c),theta_0) -F(XI.l(x_0,c),theta_0) )
    densite = nu_plus(x_0) / ( 2 + c * ( np.exp(alpha) -1 )   )
    return (2*num  * densite / deno)     
             
         
def Fisher_information (c_list,alpha,theta_0) :
    XI=XI_plus_lookup(10**5) 
    resultat=np.zeros(len(c_list))
    for i in range(len(c_list)):
         resu=integrate.quad(Fisher_i_times_density, 0, + np.inf, \
                                args=(alpha, theta_0, c_list[i],XI))         
         resultat[i]=resu[0] * ( np.exp( alpha ) - 1 )**2 
    return(resultat)    
   
    
"""    
XI_table=XI_plus_lookup(10**4)
Nb_data=1000
plot_MLE(Nb_data, 0.9 , 6, 0 , -2 , 2, 300, XI=XI_table ) 
 """       
    
"""
XI_table=XI_plus_lookup(10**4)
Nb_data=1000
Xdata=np.random.normal(size=Nb_data)
alpha=1
Zdata1=generate_Z_X(Xdata, np.array([0.95]), alpha, XI_table)[0]    
Zdata2=[generate_Z_X_pur_Bino(Xdata, alpha)]   
MLE1=MLE_computation (Zdata2, np.array([0.95]) ,alpha, theta_start, bounds, XI_table)
MLE2=MLE_computation_Bino(Zdata2[0], theta_start, bounds)  
"""    
    
#sys.exit()


if __name__ == '__main__':
    
    
 file_sauve = open('sauv_Nb_data'+str(Nb_data)+text_sauv_titre+'.txt','w')
 file_sauve.write(text_sauv + '\n')
 file_sauve.write('Nb_data : '+str(Nb_data)+'\n')
 file_sauve.write('Nb_iter : '+str(Nb_Iter)+'\n')
 file_sauve.write('Theta_0 : '+str(theta_0)+'\n')
 file_sauve.write('c_list : \n' + str(c_list)+'\n\n')    
    
 Fisher_info_theoB_all=[]
 Moyenne_allB=[]  
 ecart_type_theoB_all=[]
 ecart_typeB_all=[]

 for i in range(len(alpha)) :
      
      alpha_current=alpha[i]
  
      resu2=MC_estimation_parallel(theta_0, Nb_data, Nb_Iter, c_list, alpha_current , theta_start, bounds , precision, left=None, right=None, Nb_worker=0)


      moyenne2B, ecart_type2B=MC_estimation_min_std(resu2)

      Fisher_info_theoB=Fisher_information (c_list,alpha_current,theta_0)

      ecart_type_theoB=np.sqrt(1/( Fisher_info_theoB * Nb_data) )   
  
      resu_Bino_pur=MC_estimation_Bino(Nb_data,Nb_Iter,alpha_current,theta_start,bounds)

      moyenne2B_pur, ecart_type2B_pur=MC_estimation_min_std(resu_Bino_pur)
    
      file_sauve.write('alpha : ' + str(alpha_current)+'\n')

      file_sauve.write('moyenne : \n' + str(moyenne2B)+'\n')
      file_sauve.write('std : \n' + str(ecart_type2B)+'\n')
      file_sauve.write('std_theo : \n' + str(ecart_type_theoB)+'\n')
  
      file_sauve.write('moyenne_B_pur : ' + str(moyenne2B_pur) + '\n')
      file_sauve.write('std_B_pur : ' + str(ecart_type2B_pur) + '\n\n\n')


      Fisher_info_theoB_all.append(Fisher_info_theoB)
      Moyenne_allB.append(moyenne2B)
      ecart_type_theoB_all.append(ecart_type_theoB)
      ecart_typeB_all.append(ecart_type2B)


  
  #plt.rc('text', usetex=True)fig , ax = plt.subplots(layout='constrained')
 fig , ax = plt.subplots(layout='constrained')
 ax.set_xlim(0,1)
 ax.set_ylim(1/np.sqrt(Nb_data)*0.7,0.3)  

 for i in range(len(alpha)) :
    label_c=r"$\alpha=$"+str(alpha[i])
    ax.plot(c_list,ecart_typeB_all[i],label=label_c)
 ax.set_xlabel('value of $c$')
 ax.set_ylabel('Std dev')
 
 ax.legend()
 fig.show() 
 name_fig='sauv_Nb_data'+str(Nb_data)+text_sauv_titre+'std'+'.png'
 fig.savefig(name_fig, transparent=None, dpi='figure')
#plt.title("")

 fig , ax = plt.subplots(layout='constrained')
 ax.set_xlim(0,1)
 ax.set_ylim(1/np.sqrt(Nb_data)*0.7,0.3)  

 for i in range(len(alpha)) :
    label_c=r"$\alpha=$"+str(alpha[i])
    ax.plot(c_list,ecart_type_theoB_all[i],label=label_c)
 ax.set_xlabel('value of $c$')
 ax.set_ylabel('Std dev theo')
 
 ax.legend()
 fig.show() 
 name_fig='sauv_Nb_data'+str(Nb_data)+text_sauv_titre+'std_theo'+'.png'
 fig.savefig(name_fig, transparent=None, dpi='figure')


 fig , ax = plt.subplots(layout='constrained')
 ax.set_xlim(0,1)
 ax.set_ylim(-3/np.sqrt(Nb_data),+3/np.sqrt(Nb_data))
 for i in range(len(alpha)) :
    label_c=r"$\alpha=$"+str(alpha[i])
    ax.plot(c_list,Moyenne_allB[i],label=label_c)
 ax.set_xlabel('value of $c$')
 ax.set_ylabel('Bias')
 
 ax.legend()
 fig.show() 
 name_fig='sauv_Nb_data'+str(Nb_data)+text_sauv_titre+'bias'+'.png'
 fig.savefig(name_fig, transparent=None, dpi='figure')


 file_sauve.close()

 
  

sys.exit()















