#!/usr/bin/python3
# Régression linéaire (chute libre)

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

t = np.array([0., .31, .59, 1.02, 1.32, 1.74]) # les données de temps
y = np.array([1., 3., 4., 4., 3., 0.])         # les données de hauteur

def x1(t):        # les variables prédicteur
    return t
def x2(t):
    return t**2
    
A = np.ones((6, 3)) # la matrice de coefficients
A[:, 1] = x1(t)
A[:, 2] = x2(t)

beta = gauss.gauss(A.T @ A, A.T @ y) # paramètres = solution du système linéaire 
y0, v0, g = beta[0], beta[1], -2 * beta[2]

def f(t):
    return y0 + v0 * t - g/2 * t**2

r = y - f(t)         # les résidus
chi2 = np.sum(r**2)
    
tpoints = np.linspace(0, 1.74, 100)
plt.plot(t, y, 'ro')
plt.plot( tpoints, f(tpoints))
plt.xlabel("t [s]")
plt.ylabel("y [m]")
plt.title("y0 = {:.2f} m, v0 = {:.2f} m/s, g = {:.2f} m/s^2, chi^2 = {:.5f}".format(y0, v0, g, chi2))
plt.show()