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

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

""" 
Contains the functions and class for implementing the continuous staircase mechanism

One can change :
        The function $\nu$ is the distibution serving as parameter for the mechanism 
        The statistical model (a Gaussian shift) can be modified 

"""


import numpy as np
from numba import jit,float64
import matplotlib.pyplot as plt
from scipy import integrate, special
from scipy.optimize import minimize, Bounds
from multiprocessing import Pool, cpu_count
import math
import sys

@jit()
def nu(x) :
    """ The distribution $\nu$ defining the staircase mechanism.
        This density can be changed"""
    #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))


@jit()
def one_over_nu_ODE(x,t) :
    return (1/nu(x))


def repartition_nu (u: np.float64) -> np.float64 :
    F0=integrate.quad(nu,-np.inf,u)[0]
    return(F0)

@jit(float64(float64),forceobj=True)                                 
def inverse_repartition_nu(u : np.float64) -> np.float64:
    F0=integrate.quad(nu,-np.inf,0)[0]
    Finv=integrate.odeint(one_over_nu_ODE,0,np.array([F0,u]))
    return(Finv[1][0])
                                   
  

def tirage_selon_nu():
    """ Draw a random variable with law $\nu$ """
    return(inverse_repartition_nu(np.random.uniform()))


class XI_lookup:
    """ Class which implements a looukup table for quantities related to 
    $\nu$ : its cdf, the quantities $l_c$, $r_c$, 
    the density of the randomisation mecanism
    """
    
    def __init__(self,precision=10**4,left=None,right=None):
        """ Precision is the nb of points in the table for saving the CDF and inverse CDF of $\nu$
         In the interval (left,right) the approximate values of the CDF and inv CDF 
             are stored in a lookup table of length precision
         Outside of the interval (left, right) the  values of the CDF are computed rather than read on
         the lookup table
         Il left, right are not passed as parameter : 
                 they are computed as quantile of 1/precision , 1-1/precision
                 It is a correct choice if $\nu$ is not too heavy tailed
                 For $\nu$ being Cauchy, it is necessary set left, right (e.g. -30,30) 
                    and choose the nb of points in the grid (e.g 10**5)
        
        This Lookup table procedure is to get faster numerical simulations 
        """ 
        self.precision=precision
        self.table_F=np.zeros(precision)          
        self.table_invF=np.zeros(precision)   
        if (left==None) : self.left=inverse_repartition_nu(1/precision)
        else : self.left=left
        if (right==None): self.right=inverse_repartition_nu(1-1/precision)
        else : self.right=right
        for i in range(precision):
            self.table_invF[i]=inverse_repartition_nu((i+1)/(precision+1))
        for i in range(precision):
            self.table_F[i]=repartition_nu( (i / self.precision) * (self.right-self.left)  + self.left    )
            
            
    def value(self, x):
        """ Return the cdf of $\nu$ """
        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(x))
      
    def inv_value(self, x):
        """ Return the inverse of the cdf of $\nu$ """
        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):
        """ return $g_c(x)$ 
            x: real
            c: must be in (0,1)
        """
        if (x <= self.inv_value(c)) : return(-np.inf)
        else : return (self.inv_value( (-c/2+ self.value(x) ) ) )

    def r(self,x,c):
        """ return $d_c(x)$ """
        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):
        """ return the inverse of $ x \to g_c(x)$ """
        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):
        """ return the inverse of $ x \to d_c(x)$"""
        if (x <= self.inv_value(c/2)) : return (-np.inf)
        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 cond_normalisatio_check(self,x,c):
        """ Chech that the conditions related to the normalization of
        the measure on $g$ and $d$ are valid """
        inv_l,inv_r=self.inv_l(x,c),self.inv_r(x,c)
        integrale=integrate.quad(nu,inv_r,inv_l)[0]
        print("inv_r=",inv_r,"inv_l=",inv_l, " int de inv_r à inv_l=",integrale)
    
    def plot_l_r(self,c,reso=300) :
        abscisse=np.linspace(self.inv_value(c), self.inv_value(1-c),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 density_q (self,c,alpha,x,y):
        """ Density of the randomsation procedure $P( Z in y + d y \mid X =x)$"""
        indi=1
        if ( ( self.l(y,c) <= x ) & (x<=self.r(y,c))) : indi=np.exp(alpha)
        densi=indi * nu(y) / ( 1 + c * (np.exp(alpha)-1) )
        return(densi)
    
    def density_q_v2 (self,c,alpha,x,y):
        """ Density of the randomsation procedure i.e. 
         q_x(y)= P( Z in [y,y + d y} \mid X =x) / dy$"""
        indi=1
        if ( ( self.inv_r(x,c) <= y ) & (y<=self.inv_l(x,c))) : indi=np.exp(alpha)
        densi=indi * nu(y) / ( 1 + c * (np.exp(alpha)-1) )
        return(densi)
        
    
    def plot_densi_q(self,c,alpha,x,y,version=1) :
        """ Plot the density of the public variable 
            x : the value of the private data
            y : vector containing the grid of values where the density is computed 
            alpha : real
            c : in (0,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) :
        """ Numerical check that $q$ is indeed a density """
        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,-np.inf,np.inf,args=(c,alpha,x))[0]
        return(area)




def randomisation(X,c,alpha,XI) :
    """ Draw the public data in $Z\in\mathbb{R}$ when the private is $X\in\mathbb{R}$. 
    Here XI is a class containing the lookup table related to $\nu$
        X , alpha : real
        c : a vector of values for $c$. 
                If $c\in(0,1/2]$ then the output $Z\in\matbb{R}$ is correct
                If $c in [1.2,1)$ then the output is not correct, as one should theoretically output a value 
                          in $\mathbb{R} \cap $r$ for $r$ some extra point in the extramal set
    """
    Yc=np.zeros_like(c)
    for i in range(len(Yc)) :
       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() 
           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
       Yc[i]=Y    
    return(Yc)            




""" The statistical model  """  

""" Here it is a Gaussian translation
    Densite, CDF, derivative of CDF are explicitely given here
    The faur functions below should be changed, if the model change
""" 

def F(x,theta):
    """ CDF of the model """
    return(special.ndtr(x-theta))
    
def p(x,theta):
    """ Density of the model """ 
    return(np.exp( - (x-theta)**2 / 2 ) / np.sqrt( 2 * np. pi)  )
  
def partial_F_theta (x,theta) :
    """ Derivative of F with respect to theta"""
    return( - np.exp( - (x-theta)**2 / 2 ) / np.sqrt( 2 * np. pi)    )    
    
def generate_data_X(n,theta_0):
    """ Generate the data
        n : nb of data
        theta_0 : true value of the parameter
    """
    return(np.random.normal(theta_0,1,size=n))    




""" The remaining part can be left unchanged for another model """

def generate_Z_X(X,c,alpha,XI) :
    """ Generate the public view 
        X : vector of the $n$ private data
        c : vector of values for $c$
        alpha : real
    """      
    data_Z=np.zeros((len(X),len(c)))
    for i in range(len(X)):
        data_Z[i,:]=randomisation(X[i],c,alpha,XI)
    return(data_Z)


def minus_Log_vraisemblance_c_fixed(theta,data_Z, c, alpha,XI):
    """ Compute minus MLE 
        theta : value of the parameter for the MLE
        data_Z : public data
        c , alpha : real
    """
    LL=0
    for j in range(len(data_Z)):
            lY , rY = XI.l(data_Z[j],c) , XI.r(data_Z[j],c)
            #print(lY,rY)
            s=1+ (np.exp(alpha) - 1) * (F(rY,theta)-F(lY,theta))
            LL+=np.log(s)    
    return(-LL)     


def MLE_computation (data_Z, c_list ,alpha, theta_start, bounds, XI):
     """ Compute MLE by minimization of the previous function
         c_list : list of the values of $c$
         theta_start : starting point of the parameter theta in the minimization algo
         bounds : bounds for the minimization algo            
     """
     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) :
    """ Monte Carlo procedure to study the MLE defined in the previous function
        theta_0 : true value of the parameter
        Nb_data : $n$ 
        Nb_interation : Nb of repetitions in the Monte Carlo method
        c_list : vector with the values of $c$
        alpha : real
        theta_start : starting point for the minization algo when searching the MLE
        bound :  bounds for the minimization algo     
        left , right : real numbers, same as in the class XI
        precision : integer, same as in the class XI
    """
    np.random.seed()
    resultat=np.zeros((Nb_iteration,len(c_list))) 
    XI=XI_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)     
        MLE=MLE_computation (Z_data, c_list ,alpha , theta_start, bounds, XI)
        resultat[i,:]=MLE[:]
    return(resultat)

"""
XI_table=XI_lookup(10**4) 
Nb_data=1000
Xdata=np.random.normal(0,1,size=Nb_data)
Z_data=generate_Z_X(Xdata, c_list, alpha,XI_table)  
MLE_resu=MLE_computation (Z_data, c_list ,alpha , theta_start, bounds, XI_table)
#MC_resu=MC_estimation(theta_0, Nb_data,100, c_list ,alpha , theta_start, bounds, 10**4)
"""
 


def MC_estimation_parallel(theta_0, Nb_data, Nb_iteration,c_list ,alpha , theta_start, bounds, precision, left=None, right=None, Nb_worker=0) :
    """ Parallel version of the Monte Carlo procedure 
        Nb worker : nb of thread. Better if divide Nb_iteration 
    """
    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):
    """ Return the average means and stds from the results of the Monte Carlo experiement"""
    Nb_iter,Nb_c=resultat.shape
    moyenne=resultat.mean(0)
    std=resultat.std(0)
    return(moyenne,std)
    

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) :
    """ Draw $n$ private variable, then the $n$ public views and plot the log likelihood
        Both the LL computed from the public and the private data are displayed  
        They should match if $alpha$ large, $c$ small, $e^\alpha$ * $c$  large
        theta_min, theta_max: interval of abscissa for the graph 
        nb_point_grid : nb of abscissa points
        XI : the lookup table related to $\nu$
        c : must be in (0,1/2) for the result to be theoretically valid
        precision, left, right : if XI is not passed as an argument, 
    """
    X_data=generate_data_X(nb_data,theta_0)
    if (XI==None) : XI=XI_lookup(precision,left,right)
    data_Z=generate_Z_X(X_data, np.array([c]) ,alpha,XI)
    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=0
        for j in range(len(data_Z)):
                lY , rY = XI.l(data_Z[j],c) , XI.r(data_Z[j],c)
                s=1+ (np.exp(alpha) - 1) * (F(rY,abscice[i])-F(lY,abscice[i]))
                s=s*( nu( data_Z[j] ) ) / ( 1 + c * (np.exp(alpha) - 1) )
                LL+=np.log(s)    
        LL_Z[i] = LL 
        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 Fisher_i(x_0,alpha,theta_0, c,XI) :
    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) :
    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(x_0) / ( 1 + c * ( np.exp(alpha) -1 )   )
    return (num  * densite / deno)     
             
         
def Fisher_information (c_list,alpha,theta_0,XI=None) :
    """ Compute the Fisher information 
        c_list : vector of values for $c$. Here $c$ can be in $(0,1)$
        alpha : real
        theta_0 : true value of the parameter
    """
    if (XI==None) : XI=XI_lookup(10**5) 
    resultat=np.zeros(len(c_list))
    for i in range(len(c_list)):
         resu=integrate.quad(Fisher_i_times_density, -np.inf, + np.inf, \
                                args=(alpha, theta_0, c_list[i],XI))         
         resultat[i]=resu[0] * ( np.exp( alpha ) - 1 )**2 
    return(resultat)    


         









