#!/usr/bin/python3
# La méthode de Gauss-Newton

import numpy as np
import gauss   # pour la fonction gauss() des TD
import matplotlib.pyplot as plt

def gauss_newton(t, y, f, gradf, beta0, epsilon=1.E-4):
    beta = np.copy(beta0)      # les paramètres à ajuster
    delta = np.ones(len(beta)) # la différence entre deux itérations
                                
    while np.sqrt(np.sum(delta**2)) > epsilon:   
        r = f(t, beta) - y   # les résidus
        J = np.array([df(t, beta) for df in gradf]).T  
        delta = gauss.gauss(J.T @ J, - r @ J)
        beta += delta
    
    chi2 = np.sum(r**2)     # chi^2 après minimisation
    return beta, chi2

def f(t, beta):
    return beta[0] * np.sin(beta[1] * t + beta[2])
    
gradf = [ lambda t, beta: np.sin(beta[1] * t + beta[2]),
          lambda t, beta: beta[0] * np.cos(beta[1] * t + beta[2]) * t,
          lambda t, beta: beta[0] * np.cos(beta[1] * t + beta[2])]

data = np.loadtxt('noisysin.txt').T

beta, chi2 = gauss_newton(data[0], data[1], f, gradf, np.array([2., 1., 1.]))

plt.plot(data[0], data[1], 'ro')
t = np.linspace(0, 10, 100)
plt.plot(t, f(t, beta))
plt.xlabel("t")
plt.ylabel("y")
plt.title("A = {:.2f}, omega = {:.2f}, phi = {:.2f}, chi2 = {:.2f}".format(beta[0], beta[1], beta[2] % (2 * np.pi), chi2))
plt.show()