#!/usr/bin/python
# -*- coding: utf-8 -*-

"""
Monte Carlo réutilisable & exemple sur une manip de leçon
"""

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from cycler import cycler
from scipy.optimize import curve_fit

#-------------------------------------------------------------------------------
# Generalites
#-------------------------------------------------------------------------------

mpl.rcParams["font.size"] = 20
mpl.rcParams["lines.linestyle"] = "None"
mpl.rcParams["lines.marker"] = "+"
mpl.rcParams["lines.markersize"] = 10
mpl.rcParams["axes.prop_cycle"] = cycler(color = ["blue", "green", "orange", "purple"])
mpl.rcParams["axes.grid"] = True

rng = np.random.default_rng()

mcN = 1000

def normale(mu, err):
    return rng.normal(mu, err/2, (mcN, *np.shape(mu)))

def uniforme(mu, err):
    return rng.uniform(mu - err, mu + err, (mcN, *np.shape(mu)))

def triangulaire(mu, err):
    return rng.triangular(mu - err, mu, mu + err, (mcN, *np.shape(mu)))

def formatter(V):
    valeur = np.median(V, axis = 0)
    erreur = np.maximum(valeur - np.percentile(V, 2.5, axis = 0), np.percentile(V, 97.5, axis = 0) - valeur)
    exp = np.floor(np.log10(erreur/.9)).astype(int) -1  # A modifier 
    err = np.ceil(erreur / 10.**exp).astype(int)
    val = np.round(valeur / 10.**exp).astype(int)
    if exp == 0:
        return r"({} ± {})".format(val, err)
    else:
        return r"({} ± {})e{}".format(val, err, exp)

def afficher(V):
    print(formatter(V))

def R2(x, y, modele):
    return 1 - np.sum((y - modele.evaluer(x))**2) / np.sum((y - np.mean(y))**2)

def afficher_R2(*args):
    print("R² = {}".format(np.round(R2(*args), decimals = 3)))

def chi2reduit(X, Y, modele):
    x = np.median(X, axis = 0)
    y = np.median(Y, axis = 0)
    sigma2total = np.std(modele.evaluer(X), 0)**2 + np.std(Y, 0)**2
    chi2 = np.sum((y - modele.evaluer(x))**2 / sigma2total)
    return chi2 / (X.shape[1] - modele.degres_de_liberte)

def afficher_chi2(*args):
    print("χ²/n = {}".format(np.round(chi2reduit(*args), decimals = 2)))

def plotXY(X, Y, N = None):
    x = np.median(X, axis = 0)
    xinf = x - np.percentile(X, 5, axis = 0)
    xsup = np.percentile(X, 95, axis = 0) - x
    y = np.median(Y, axis = 0)
    yinf = y - np.percentile(Y, 5, axis = 0)
    ysup = np.percentile(Y, 95, axis = 0) - y
    if N is None:
        plt.errorbar(x, y, xerr = (xinf, xsup), yerr = (yinf, ysup))
    else:
        plt.errorbar(x[:N], y[:N], xerr = (xinf[:N], xsup[:N]), yerr = (yinf[:N], ysup[:N]))
        plt.errorbar(x[N:], y[N:], xerr = (xinf[N:], xsup[N:]), yerr = (yinf[N:], ysup[N:]), color = "red")

class Modele:
    def tracer(self):
        gauche, droite = plt.xlim()
        x = np.linspace(gauche, droite, 500)
        plt.plot(x, self.evaluer(x), linestyle = "dashed", marker = "None", color = "black", label = "Ajustement")
    
    def afficher(self):
      print(self.formatter())

    def ajusterMC(self, X, Y):
        coef0 = self.ajuster(X[0], Y[0])
        self.coef = np.zeros((mcN, *np.shape(coef0)))
        self.coef[0] = coef0
        for i in range(1, mcN):
            self.coef[i] = self.ajuster(X[i], Y[i])

class ModelePolynomial(Modele):
    def __init__(self, deg):
        self.deg = deg
        self.degres_de_liberte = deg + 1

    def ajuster(self, x, y):
        return np.polynomial.polynomial.polyfit(x, y, self.deg)

    def evaluer(self, x):
        return np.polynomial.Polynomial(np.median(self.coef, axis = 0))(x)

    def formatter(self):
        s = "y = "
        for i in range(self.deg + 1):
            if i == 0:
                s += "{} + ".format(formatter(self.coef[..., i]))
            elif i == 1:
                s += "{} · x + ".format(formatter(self.coef[..., i]))
            else:
                s += "{} · x^{} + ".format(formatter(self.coef[..., i]), i)
        return s[:-3]

class ModeleLineaire(Modele):
    def __init__(self):
        self.degres_de_liberte = 1

    def ajuster(self, x, y):
        x = x[:, np.newaxis]
        a, _, _, _ = np.linalg.lstsq(x, y, rcond = None)
        return a[0]

    def evaluer(self, x):
        return np.median(self.coef) * x

    def formatter(self):
        return "y = {} · x".format(formatter(self.coef))

class ModeleFonction(Modele):
    def __init__(self, f, p0 = None):
        self.f = f
        self.p0 = p0
        self.degres_de_liberte = f.__code__.co_argcount - 1

    def ajuster(self, x, y):
        popt, _ = curve_fit(self.f, x, y, p0 = self.p0)
        return popt

    def evaluer(self, x):
        return self.f(x, *np.median(self.coef, axis = 0))

    def formatter(self):
        s = ""
        for i, v in enumerate(self.f.__code__.co_varnames[1:]):
            s += "{} = {}\n".format(v, formatter(self.coef[..., i]))
        return s[:-1]

#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
# Ce code permet de réaliser une regression polynomiale sur des données expérimentales.

# Les incertitudes sont calculées par la méthode de MonteCarlo (Simulation aléatoire et statistiques). 

# Incertitudes (à compléter)

DeltaL = "à compléter"
DeltaD = "à compléter"
Deltaepsilon = 0
# Valeurs expérimentales (à compléter)

L = uniforme(np.array(["à compléter"]),DeltaL)
D = uniforme(np.array(["à compléter"]),DeltaD)
epsilon = uniforme("à compléter"*np.ones(4),Deltaepsilon)

X = 2*D/epsilon
Y = L

plotXY(X,Y)

modele = ModeleLineaire()
modele.ajusterMC(X,L)
modele.tracer()
modele.afficher()

# Légendes
plt.xlabel("2*D/epsilon []", size=14)
plt.ylabel("Largeur L de la tache de diffraction [m]", size=14)

# Grille
plt.grid(True, linestyle=":", alpha=0.7)