import numpy as np
import pandas as pd
from sklearn.decomposition import PCA
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LassoLarsCV
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline, make_union
from sklearn.preprocessing import Imputer
from tpot.builtins import StackingEstimator, ZeroCount

from sklearn.externals import joblib

def tpot_best_model():
    pop_5_rel = pd.read_feather( '../../data/pop_5_with_playlist_normalised.feather')
    pop_5_rel['pop_5'] = pop_5_rel['value']
    pop_5_rel = pop_5_rel.drop('value', axis =1)
    pop_5_rel['log_streams'] = np.log(pop_5_rel['streams'])

    X = pop_5_rel[['first_position', 'playlist_spyid', 'pop_5']].copy()

    T = pd.get_dummies(X)
    feature_columns = T.columns.values
    X = pd.get_dummies(X).values
    y = pop_5_rel[['log_streams']]

    training_features, testing_features, training_target, testing_target = train_test_split(X, y, test_size = 0.20, random_state = 1)

    imputer = Imputer(strategy="median")
    imputer.fit(training_features)
    training_features = imputer.transform(training_features)
    testing_features = imputer.transform(testing_features)

# Average CV score on the training set was:-1.8122937118960807
    exported_pipeline = make_pipeline(
        StackingEstimator(estimator=LassoLarsCV(normalize=True)),
        PCA(iterated_power=5, svd_solver="randomized"),
        StackingEstimator(estimator=LassoLarsCV(normalize=False)),
        ZeroCount(),
        RandomForestRegressor(bootstrap=False, max_features=0.25, min_samples_leaf=12, min_samples_split=19, n_estimators=100)
    )

    exported_pipeline.fit(training_features, training_target)

    joblib.dump({'model': exported_pipeline, 'feature_columns' : feature_columns}, './model_spec/tpot_generated_model.pkl')

    # results = exported_pipeline.predict(testing_features)