# Heavy lifting
import numpy as np
import pandas as pd
from scipy import stats

# Mostrare i dati
import ipysheet
from IPython.display import display
import matplotlib.pyplot as plt
import seaborn as sns

# Calcolo simbolico
import sympy as sp

def compat(x1, x2, u1,u2):
    
    diff = x1 - x2
    k = np.abs(diff) / np.sqrt(u1**2 + u2**2)
    return k, diff

class Data:

    def __init__(self,
                 campione,
                 analisi_stat=None, param_reg=None):
        
        self.campione     = campione
        self.analisi_stat = analisi_stat
        self.param_reg    = param_reg

# Probabilità
def p_t_student(valori, err_strumentale):

  n   = len(valori)
  GdL = n - 1

  media = valori.mean()
  s     = valori.std(ddof=1)

  s = np.sqrt(s**2 + err_strumentale**2)
  return 2 * (1 - stats.t.cdf(np.abs(valori - media) / s, df=GdL))


# Indice del peggiore outlier o None
def trova_outlier(valori, err_strumentale):

  n       = len(valori)
  soglia  = 1 / (2 * n)
  p       = p_t_student(valori, err_strumentale)
  idx_min = np.argmin(p)

  if p[idx_min] < soglia:
    return idx_min
  
  return None


# Rimozione outlier
def rimuovi_outlier(valori, err_strumentale):
  
  rimossi  = []
  campione = valori.copy()

  while len(campione) > 2:
    idx = trova_outlier(campione, err_strumentale)

    if idx is None: # nessun outlier: stop
      break

    rimossi.append(campione[idx])
    campione = np.delete(campione, idx)

  return campione, rimossi

# Analisi statistica con criterio di Chauvenet
def stat_chauv(self, prefissi, err_strumentale):

    self.analisi_stat = pd.DataFrame()
    campione = self.campione
    
    
    if isinstance(prefissi, str):
        prefissi = [prefissi]
    if isinstance(err_strumentale, (int, float)):
        err_strumentale = [err_strumentale] * len(prefissi)

    for prefix, err_strum in zip(prefissi, err_strumentale):
        cols = [col for col in campione.columns if col.startswith(prefix)]

        for i, row in campione.iterrows():
            valori            = row[cols].dropna().values.astype(float)
            campione, rimossi = rimuovi_outlier(valori, err_strum)

            n = len(campione)
            s = campione.std(ddof=1)
            s = s / np.sqrt(n)
            u = np.sqrt(s**2 + err_strum**2)

            self.analisi_stat.at[i, prefix]              = campione.mean()
            self.analisi_stat.at[i, f"u{prefix}_strum"]  = err_strum
            self.analisi_stat.at[i, f"u{prefix}_stat"]   = s
            self.analisi_stat.at[i, f"u{prefix}"]        = u
            self.analisi_stat.at[i, f"n{prefix}"]        = n
            self.analisi_stat.at[i, f"out{prefix}"]      = rimossi

    return self


# Aggiunta a classe Data
Data.stat_chauv = stat_chauv

def calc_var(self, funz, nome_risultato):

    # Estrazione e ordinamento variabili simboliche
    vars_list = sorted(list(funz.free_symbols), key=lambda s: s.name)
    
    # Creazione incertezze e contributi simbolici
    u     = {var: sp.Symbol(f'u{var.name}', positive=True) for var in vars_list}
    contr = [sp.diff(funz, var)*u[var] for var in vars_list]
    
    # Calcolo simbolico incertezza propagata
    u_prop = 0
    for contributo in contr:
        u_prop += contributo**2 
    u_prop = sp.sqrt(u_prop)

    # Creazione lista incertezze simboliche
    u_list = [u[var] for var in vars_list]
    
    # Funzioni numeriche
    funz_fn   = sp.lambdify(vars_list, funz, 'numpy')
    u_prop_fn = sp.lambdify((vars_list, u_list), u_prop, 'numpy')
    contr_fn  = sp.lambdify((vars_list, u_list), contr, 'numpy')
    
    # Estrazione dati da analisi_stat
    vars_num = [self.analisi_stat[var.name].values.astype(float) for var in vars_list]
    u_num    = [self.analisi_stat[f'u{var.name}'].values.astype(float) for var in vars_list]

    # Calcolo numerico risultati
    risultato_val   = funz_fn(*vars_num)
    risultato_u     = u_prop_fn(vars_num, u_num)
    risultato_contr = (np.array(contr_fn(vars_num, u_num)) / risultato_u)**2 * 100

    # Aggiunta risultati ad analisi_stat
    self.analisi_stat[nome_risultato]       = risultato_val
    self.analisi_stat[f"u{nome_risultato}"] = risultato_u
    for i, var in enumerate(vars_list):
        self.analisi_stat[f"%u{nome_risultato}_{var.name}"] = risultato_contr[i]
    
    return self

# Aggiunta a classe Data
Data.calc_var = calc_var

# Plot residui
 
def plot_residui(x, r,
                 ux, ur,
                 x_label, r_label, titolo):
    
    fig, ax = plt.subplots(figsize=(8, 5))

    # Residui con barre d'errore
    ax.errorbar(
        x, r,
        xerr=ux, yerr=ur,
        fmt='o', color=sns.color_palette()[0],
        ecolor='gray', elinewidth=1, capsize=3,
        markersize=5, label="Residui"
    )

    # Linea dello zero
    ax.axhline(0, color='red', linestyle='--', linewidth=1)

    # Estetica
    sns.despine(ax=ax)
    ax.set_xlabel(x_label)
    ax.set_ylabel(r_label)
    ax.set_title(titolo)
    ax.legend()
    ax.grid(True, linestyle=':', linewidth=0.5, alpha=0.7)

    plt.tight_layout()
    plt.show()

# Calcolo residui

def residui(analisi_stat,
            A_num, B_num,
            uA_num, uB_num, covAB_num):

  # Variabili simboliche
  A, B, x, y = sp.symbols('A B x y', real=True)
  uA, uB, ux, uy = sp.symbols('uA uB ux uy', positive=True)
  covAB = sp.symbols('covAB', real=True)

  # Residuo: r = y - (Ax + B)
  r = y - (A*x + B)

  # Propagazione errore (con covarianza A,B)
  dr_dA  = sp.diff(r, A)
  dr_dB  = sp.diff(r, B)
  dr_dx  = sp.diff(r, x)
  dr_dy  = sp.diff(r, y)

  u_r  = sp.sqrt(
    (dr_dA * uA)**2 +
    (dr_dB * uB)**2 +
    (dr_dx * ux)**2 +
    (dr_dy * uy)**2 +
    2 * dr_dA * dr_dB * covAB
  )

  # Funzioni numeriche
  r_fn   = sp.lambdify((x , y , A , B ), r, 'numpy')
  u_r_fn = sp.lambdify(
    (x , y , ux , uy , A , B ,  uA , uB , covAB ),
    u_r , 'numpy'
  )

  # Calcolo numerico
  analisi_stat["r"] = r_fn(
    analisi_stat["x"],
    analisi_stat["y"],
    A_num,
    B_num
  )

  analisi_stat["ur"] = u_r_fn(
    analisi_stat["x"], analisi_stat["y"],
    analisi_stat["ux"], analisi_stat["uy"],
    A_num, B_num,
    uA_num, uB_num, covAB_num
  )

  return analisi_stat


# Plot regressione

def plot_reg(x, y, ux, uy,
             A, B, uA, uB, P,
             x_label, y_label, titolo,
             cd_A=4, cd_B=4, scala_barre=1):

    fig, ax = plt.subplots(figsize=(8, 5))

    x_fit = np.linspace(x.min(), x.max(), 300)
    y_fit = A * x_fit + B

    ax.errorbar(
        x, y,
        xerr = scala_barre * ux,
        yerr = scala_barre * uy,
        fmt='o', color=sns.color_palette()[0],
        ecolor='gray', elinewidth=1, capsize=3,
        markersize=5, label=f"Dati (barre errore x{scala_barre})"
    )
    ax.plot(
        x_fit, y_fit,
        color='red', linewidth=1.5,
        label = f"$A={A:.{cd_A}f}\\pm{uA:.{cd_A}f}$\n"
                f"$B={B:.{cd_B}f}\\pm{uB:.{cd_B}f}$\n"
                f"$P(chi², ∞)={P:.4f}$\n"
    )

    sns.despine(ax=ax)
    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)
    ax.set_title(titolo)
    ax.legend(fontsize=9)
    ax.grid(True, linestyle=':', linewidth=0.5, alpha=0.7)

    plt.tight_layout()
    plt.show()


# Funzioni regressione

def uy_equiv(x, y, ux, uy):
  
  # Stima iniziale di A con sola uy
  sy2   = uy**2
  Sw    = np.sum(1 / sy2)
  Sx    = np.sum(x / sy2)
  Sxx   = np.sum(x**2 / sy2)
  Sy    = np.sum(y / sy2)
  Sxy   = np.sum(x * y / sy2)
  delta = Sxx * Sw - Sx**2
  A_est = (Sxy * Sw - Sx * Sy) / delta

  # Propagazione
  u_eq = np.sqrt(uy**2 + A_est**2 * ux**2)
  return u_eq


def reg_lin(self,
            stampa_param=True, plot_regressione=True, calc_residui=True,
            x_label="", y_label="", r_label="",
            titolo_reg="", titolo_residui="",
            cd_A=4, cd_B=4, scala_barre=1):
  
  x  = np.asarray(self.analisi_stat["x"], dtype=float)
  y  = np.asarray(self.analisi_stat["y"], dtype=float)
  ux = np.asarray(self.analisi_stat["ux"], dtype=float)
  uy = np.asarray(self.analisi_stat["uy"], dtype=float)

  # Propagazione errore x → y
  uy_eq = uy_equiv(x, y, ux, uy)

  # Somme pesate
  w     = 1.0 / uy_eq**2
  Sw    = np.sum(w)
  Sx    = np.sum(w * x)
  Sxx   = np.sum(w * x**2)
  Sy    = np.sum(w * y)
  Sxy   = np.sum(w * x * y)
  delta = Sxx * Sw - Sx**2

  # Parametri
  A     = (Sxy * Sw - Sx * Sy)  / delta
  B     = (Sxx * Sy - Sxy * Sx) / delta
  uA    = np.sqrt(Sw  / delta)
  uB    = np.sqrt(Sxx / delta)
  covAB = -Sx / delta

  # Chi quadro
  x2  = np.sum((y - A * x - B)**2 / uy_eq**2)
  dof = len(x) - 2
  P   = stats.chi2.sf(x2, dof)

  # Raccolta parametri
  self.param_reg = pd.DataFrame()
  self.param_reg["A"]     = A
  self.param_reg["B"]     = B
  self.param_reg["uA"]    = uA
  self.param_reg["uB"]    = uB
  self.param_reg["covAB"] = covAB
  self.param_reg["x2"]    = x2
  self.param_reg["P"]     = P

  # Stampa
  if stampa_param == True:
    print("Ax + B : ")
    print(f"A          = {A:.{cd_A}f} ± {uA:.{cd_A}f}")
    print(f"B          = {B:.{cd_B}f} ± {uB:.{cd_B}f}")
    print(f"covAB      = {covAB:.6f}")
    print(f"chi²       = {x2:.2f}")
    print(f"P(chi², ∞) = {P:.2f}")
  
  # Plot
  if plot_regressione == True:
    plot_reg(x, y, ux, uy,
             A, B, uA, uB, P,
             x_label, y_label, titolo_reg,
             cd_A, cd_B, scala_barre)
  
  # Residui
  if calc_residui == True:
    self.analisi_stat = residui(self.analisi_stat, A, B, uA, uB, covAB)
    plot_residui(x, self.analisi_stat["r"],
                 ux, self.analisi_stat["ur"],
                 x_label, r_label, titolo_residui)

  return self


# Aggiunta a classe Data
Data.reg_lin = reg_lin

def plot_gauss(gaussiane):
    
    # Creazione figura
    plt.figure(figsize=(12, 7))

    # Creazione asse x
    xMin = float('inf')
    xMax = float('-inf')

    for mu, sigma, _, _ in gaussiane:
        minimoLocale  = mu - 4 * sigma
        massimoLocale = mu + 4 * sigma
        
        if minimoLocale < xMin:
            xMin = minimoLocale
        
        if massimoLocale > xMax:
            xMax = massimoLocale
    
    x = np.linspace(xMin, xMax, 1000)

    # Ciclo gaussiane
    for mu, sigma, colore, etichetta in gaussiane:
        y = stats.norm.pdf(x, mu, sigma)
        plt.plot(x, y, color=sns.color_palette()[colore], linewidth=1, label=etichetta)
        
        puntiLinee = [mu - sigma, mu, mu + sigma]
        for px in puntiLinee:
            py = stats.norm.pdf(px, mu, sigma) 
            plt.vlines(x=px,
                       ymin=0, ymax=py,
                       colors=sns.color_palette()[colore],
                       linestyles='dashed', linewidth=1)
            
    # Dettagli estetici finali
    plt.ylim(bottom=0)
    plt.title('Confronto dati')
    plt.xlabel('k')
    plt.legend()

    # Mostriamo il grafico
    plt.show()