msid1

import numpy as np
from utils import polynomial


def mean_squared_error(x, y, w):
    '''
    :param x: ciag wejsciowy Nx1
    :param y: ciag wyjsciowy Nx1
    :param w: parametry modelu (M+1)x1
    :return: blad sredniokwadratowy pomiedzy wyjsciami y
    oraz wyjsciami uzyskanymi z wielowamiu o parametrach w dla wejsc x
    '''
    y_pred = polynomial(x, w)
    N = x.shape[0]
    return np.sum((y-y_pred)**2)/N
    pass


def design_matrix(x_train, M):
    '''
    :param x_train: ciag treningowy Nx1
    :param M: stopien wielomianu 0,1,2,...
    :return: funkcja wylicza Design Matrix Nx(M+1) dla wielomianu rzedu M
    '''
    N = x_train.shape[0]
    my_matrix = np.zeros((N, M+1))
    for i in range(x_train.shape[0]):
        for j in range(M+1):
            my_matrix[i][j] = (x_train[i]**j)
    return my_matrix
    pass

def least_squares(x_train, y_train, M):
    '''
    :param x_train: ciag treningowy wejscia Nx1
    :param y_train: ciag treningowy wyjscia Nx1
    :param M: rzad wielomianu
    :return: funkcja zwraca krotke (w,err), gdzie w sa parametrami dopasowanego wielomianu, a err blad sredniokwadratowy
    dopasowania
    '''
    my_matrix = design_matrix(x_train,M)
    w = np.linalg.inv((my_matrix.transpose())@my_matrix)@(my_matrix.transpose())@y_train
    err = mean_squared_error(x_train,y_train,w)
    return w, err
    pass



def regularized_least_squares(x_train, y_train, M, regularization_lambda):
    '''
    :param x_train: ciag treningowy wejscia Nx1
    :param y_train: ciag treningowy wyjscia Nx1
    :param M: rzad wielomianu
    :param regularization_lambda: parametr regularyzacji
    :return: funkcja zwraca krotke (w,err), gdzie w sa parametrami dopasowanego wielomianu zgodnie z kryterium z regularyzacja l2,
    a err blad sredniokwadratowy dopasowania
    '''
    my_matrix = design_matrix(x_train,M)
    unit_matrix = np.eye(M+1)
    lambda_matrix = unit_matrix * regularization_lambda
    w = np.linalg.inv(((my_matrix.transpose())@ my_matrix) + lambda_matrix)@(my_matrix.transpose())@y_train
    err = mean_squared_error(x_train, y_train, w)
    return w, err
    pass


def model_selection(x_train, y_train, x_val, y_val, M_values):
    '''
    :param x_train: ciag treningowy wejscia Nx1
    :param y_train: ciag treningowy wyjscia Nx1
    :param x_val: ciag walidacyjny wejscia Nx1
    :param y_val: ciag walidacyjny wyjscia Nx1
    :param M_values: tablica stopni wielomianu, ktore maja byc sprawdzone
    :return: funkcja zwraca krotke (w,train_err,val_err), gdzie w sa parametrami modelu, ktory najlepiej generalizuje dane,
    tj. daje najmniejszy blad na ciagu walidacyjnym, train_err i val_err to bledy na sredniokwadratowe na ciagach treningowym
    i walidacyjnym
    '''
    w_err_train = least_squares(x_train,y_train,0)
    val_err = mean_squared_error(x_val,y_val,w_err_train[0])
    for i in M_values:
        w_i_train = least_squares(x_train,y_train,i)
        w_i_val = mean_squared_error(x_val,y_val,w_i_train[0])
        if w_i_val < val_err:
            w_err_train = w_i_train
            val_err = w_i_val

    w = w_err_train[0]
    train_err = w_err_train[1]
    return w, train_err, val_err
    pass


def regularized_model_selection(x_train, y_train, x_val, y_val, M, lambda_values):
    '''
    :param x_train: ciag treningowy wejscia Nx1
    :param y_train: ciag treningowy wyjscia Nx1
    :param x_val: ciag walidacyjny wejscia Nx1
    :param y_val: ciag walidacyjny wyjscia Nx1
    :param M: stopien wielomianu
    :param lambda_values: lista ze wartosciami roznych parametrow regularyzacji
    :return: funkcja zwraca krotke (w,train_err,val_err,regularization_lambda), gdzie w sa parametrami modelu, ktory najlepiej generalizuje dane,
    tj. daje najmniejszy blad na ciagu walidacyjnym. Wielomian dopasowany jest wg kryterium z regularyzacja. train_err i val_err to
    bledy na sredniokwadratowe na ciagach treningowym i walidacyjnym. regularization_lambda to najlepsza wartosc parametru regularyzacji
    '''
    w_err_train = regularized_least_squares(x_train, y_train, M,0)
    val_err = mean_squared_error(x_val, y_val, w_err_train[0])
    regularization_lambda = 0
    for i in lambda_values:
        w_i_train = regularized_least_squares(x_train,y_train,M,i)
        val_i = mean_squared_error(x_val,y_val,w_i_train[0])
        if val_i < val_err:
            w_err_train = w_i_train
            val_err = val_i
            regularization_lambda = i
    w = w_err_train[0]
    train_err = w_err_train[1]
    return w, train_err, val_err, regularization_lambda
    pass