from RandomizedLass import RandomizedLasso
import argparse

from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.feature_selection import RFE

from sklearn.preprocessing import MinMaxScaler
import numpy as np

def get_arguments():
    parser = argparse.ArgumentParser()

    parser.add_argument('--top_k', type=int, default=4, help='Кол-во самых выжных признаков')

    args = parser.parse_args()
    return args


def data_gen():
    # --- генерируем исходные данные: 750 строк-наблюдений и 14 столбцов-признаков ---
    np.random.seed(0)
    size = 750
    X = np.random.uniform(0, 1, (size, 14))
    #Задаем функцию-выход: регрессионную проблему Фридмана
    Y = (10 * np.sin(np.pi*X[:,0]*X[:,1]) + 20*(X[:,2] - .5)**2 +
    10*X[:,3] + 5*X[:,4]**5 + np.random.normal(0,1))
    #Добавляем зависимость признаков
    X[:,10:] = X[:,:4] + np.random.normal(0, .025, (size,4))

    return X, Y


def rank_to_dict(ranks, names):
    ranks = np.abs(ranks)
    minmax = MinMaxScaler()
    ranks = minmax.fit_transform(np.array(ranks).reshape(14,1)).ravel()
    ranks = map(lambda x: round(x, 2), ranks)
    return dict(zip(names, ranks)) 


def print_sorted_data(ranks: dict):
    print()
    for key, value in ranks.items():
        ranks[key] = sorted(value.items(), key=lambda item: item[1], reverse=True)
    for key, value in ranks.items():
        print(key)
        print(value)


def estimation(ranks: dict, top_k):
    #Создаем пустой список для данных
    mean = {}
    #«Бежим» по списку ranks
    for key, value in ranks.items():
        #«Пробегаемся» по списку значений ranks, которые являются парой имя:оценка
        for item in value.items():
            #имя будет ключом для нашего mean
            #если элемента с текущим ключем в mean нет - добавляем
            if (item[0] not in mean):
                mean[item[0]] = 0
                #суммируем значения по каждому ключу-имени признака
                mean[item[0]] += item[1]

    #находим среднее по каждому признаку
    for key, value in mean.items():
        res=value/len(ranks)
        mean[key] = round(res, 2)

    #сортируем и распечатываем список
    mean_sorted = dict(sorted(mean.items(), key=lambda x:x[1], reverse=True))
    print("sorted MEAN")
    print(mean_sorted, '---'*25, sep='\n') 

    for item in list(mean_sorted.items())[:top_k]:
        print(f'Параметр - {item[0]}, значение - {item[1]}')
    print('---'*25)


if __name__=="__main__":
    args = get_arguments()

    X,Y = data_gen()

    # Линейная модель
    lr = LinearRegression()
    lr.fit(X, Y)

    # Рекурсивное сокращение признаков
    rfe = RFE(lr)
    rfe.fit(X, Y)

    # Случайное Лассо
    randomized_lasso = RandomizedLasso(alpha=.01)
    randomized_lasso.fit(X, Y)

    names = ["x%s" % i for i in range(1,15)] 
    
    ranks = {}

    ranks["Linear reg"] = rank_to_dict(lr.coef_, names)
    ranks["RFE"] = rank_to_dict(rfe.ranking_, names)
    ranks["RandomizedLasso"] = rank_to_dict(randomized_lasso.coef_, names) 

    estimation(ranks, args.top_k)
    
    print_sorted_data(ranks)