import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import Imputer

from sklearn.externals import joblib

import prepp_ml as prepp

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

    pipeline = joblib.load('../model_spec/tpot_10_10_rf.pkl')['model']
    tpot_best_model(input_file, 
            pre_process_function=process_to_playlist_location_dummies,
            pipeline=pipeline,
            model_name='tpot_10_10_rf.pkl')


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

    pipeline = joblib.load('../model_spec/tpot_20_20_xgboos.pkl')['model']
    tpot_best_model(input_file, 
            pre_process_function=process_to_first_seen_dummies,
            pipeline=pipeline,
            model_name='tpot_20_20_tree.pkl')


def process_to_first_seen_dummies(input_file):
    return prepp.first_seen_dummies(prepp.preprocess(input_file))


def process_to_playlist_location_dummies(input_file):
    pop_5_rel =  prepp.preprocess(input_file)

    X = pop_5_rel[['first_position', 'playlist_spyid', 'pop_5']].copy()
    X = pd.get_dummies(X).values
    y = pop_5_rel[['log_streams']]

    return X,y


def tpot_best_model(input_file, pre_process_function, pipeline, model_name):
    
    X, y = pre_process_function(input_file)

    feature_columns = X.columns.values
    training_features, testing_features, training_target, testing_target = train_test_split(X, y, test_size = 0.20, random_state = 1)

    # TODO: this ius probably not universal
    imputer = Imputer(strategy="median")
    imputer.fit(training_features)
    training_features = imputer.transform(training_features)
    testing_features = imputer.transform(testing_features)

    pipeline.fit(training_features, training_target)

    # import ipdb
    # ipdb.set_trace()
    predicted = pipeline.predict(imputer.transform(X))
    out = pd.DataFrame({'predicted_streams':predicted, 'actual_streams':y.log_streams})
    out.reset_index().to_feather('../../data/backtesting/first_seen_dummy.feather')

    joblib.dump({'model': pipeline, 'feature_columns' : feature_columns, 'imputer': imputer}, f'../model_spec/sklearn_model_{model_name}.pkl')