# --------------------------------------------------------------------------
# ------------  Metody Systemowe i Decyzyjne w Informatyce  ----------------
# --------------------------------------------------------------------------
#  Zadanie 1: Regresja liniowa
#  autorzy: A. Gonczarek, J. Kaczmar, S. Zareba
#  2017
# --------------------------------------------------------------------------

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
    '''
    return (sum(pow(y - polynomial(x, w), 2))/y.shape[0])[0]



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
    matrix = np.ones(shape=(x_train.shape[0], M+1))
    for col in range(M + 1):
        for row in range(matrix.shape[0]):
            matrix[row][col] = pow(x_train[row][0], col)
    '''
    return np.array([[pow(x_train[row][0], col) for col in range(M+1)] for row in range(x_train.shape[0])])




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
    '''
    #pass
    matrix = design_matrix(x_train, M)
    w = np.linalg.inv(matrix.transpose()@matrix)@matrix.transpose()@y_train
    return w, mean_squared_error(x_train, y_train, w)


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
    '''
    #pass

    matrix = design_matrix(x_train, M)
    w = np.linalg.inv(matrix.transpose()@matrix+regularization_lambda*np.identity(matrix.shape[1]))@matrix.transpose()@y_train
    return w, mean_squared_error(x_train, y_train, w)


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, train_err = least_squares(x_train, y_train, 0)
    val_err = mean_squared_error(x_val, y_val, w)

    for m in M_values:
        w_temp, train_err_temp = least_squares(x_train, y_train, m)
        if mean_squared_error(x_val, y_val, w_temp) < val_err:
            w, train_err, val_err = w_temp, train_err_temp, mean_squared_error(x_val, y_val, w_temp)
    return w, train_err, val_err


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, train_err = regularized_least_squares(x_train, y_train, M, lambda_values[0])
    val_err, lmbd = mean_squared_error(x_val, y_val, w), lambda_values[0]

    for rlambda in lambda_values:
        w_temp, train_err_temp = regularized_least_squares(x_train, y_train, M, rlambda)
        if mean_squared_error(x_val, y_val, w_temp) < val_err:
            w, train_err, val_err, lmbd = w_temp, train_err_temp, mean_squared_error(x_val, y_val, w_temp), rlambda

    return w, train_err, val_err, lmbd