import numpy as np
import statsmodels.api as sm
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from statsmodels.sandbox.regression.predstd import wls_prediction_std

import pandas as pd

from model_evaluation import *

import json

def main_models():

# Only base 5
    run_model_with(base_model, '../../data/pop_5_release.feather')

    # Baseline model with aggregatet playlist features
    run_model_with(base_playlist_agg_x, '../data/pop_5_release_with_playlist_agg.feather')

    # Baseline model with dummy playlist features
    run_model_with(base_playlist_dummy, '../data/pop5_play_dummy.feather')
    

def plotFittedWithConfidence(x_train, y_train, y_pred, res):

    fig, ax = plt.subplots(figsize=(8,6))

    plt.subplot(2,2,1)
    x = x_train.pop_5
    ax.plot(x, y_train, 'o', label="data")
    ax.plot(x, y_pred, 'r.', label="Model")
    ax.legend(loc='best');







def base_playlist_agg_x(pop_5_rel):
    # TODO: HACK
    pop_5_rel = pop_5_rel[~pd.isna(pop_5_rel['num_pp'])]

    X = pop_5_rel[['pop_5', 'min_pp', 'num_pp']]
    X = sm.add_constant(X)
    return X, pop_5_rel[['log_streams']]

def base_model(pop_5_rel):
    X = pop_5_rel[['pop_5']]
    X = sm.add_constant(X)

    return X, pop_5_rel[['log_streams']]


def run_model_with(transform_fun, input_file):
    pop_5_rel = pd.read_feather(input_file)
    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, y = transform_fun(pop_5_rel) 

    x_train, x_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state = 1)

    model = sm.OLS(y_train, x_train)
    res = model.fit()

    coeffs = res.params.to_dict()

    # with open('../modelling/model_spec/base_dummy_playlist.json', 'w') as mm:
        # json.dump(coeffs, mm)

    print(res.summary())

    y_pred = res.predict(x_train)
    plotFittedWithConfidence(x_train, y_train, y_pred, res)
    printModelEvaluations(res, x_train, y_train, x_test, y_test, plot=True)

    # pop_5_rel['predicted_streams'] = np.exp(res.predict(X))
    # for_backtest = pop_5_rel[['spyid', 'streams','pop_5', 'streams_100_date', 'pop_5_date', 'predicted_streams']]
    # for_backtest.to_feather('../../data/backtesting/pop_5_base_model.feather')