import pandas as pd
import numpy as np
from tpot import TPOTRegressor
from sklearn.model_selection import train_test_split
from model_evaluation import *
from sklearn.externals import joblib

from prepp_ml import preprocess
from prepp_ml import first_seen_dummies


def main():
    input_file =  '../../data/pop_5_with_playlist_normalised.feather'

    # playlist_dummy
    # run_model_wi


def run_model_with(transform_fun, input_file):
    pop_5_rel = preprocess(input_file)

    X, y = first_seen_dummies(pop_5_rel, fill_na = 1000) 
    feature_cols = X.columns
    x_train, x_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state = 1)

    y_train = np.ravel(y_train)
    y_test = np.ravel(y_test)

    tpot = TPOTRegressor(generations=20, population_size=20, verbosity=2, n_jobs=2)
    tpot.fit(x_train, y_train)
    print(tpot.score(x_test, y_test))

    model_name = 'tpot_20_20_xgboos_fill_na_1000'
    joblib.dump({'model':tpot.fitted_pipeline_, 'feature_cols':feature_cols}, filename=f'../model_spec/{model_name}.pkl')
    tpot.export(f'../model_spec/{model_name}.py')
    
    printModelEvaluations(tpot, x_train, y_train.log_streams.values, x_test, y_test.log_streams.values, plot=True)