import pandas as pd
import numpy as np

from sklearn.model_selection import train_test_split

from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import make_pipeline
from sklearn.grid_search import GridSearchCV
from sklearn.metrics import mean_squared_error

from plotnine import *

pd.set_option('display.max_columns', None)

import clean

from model_evaluation import *


def eda():
    ggplot(pop_5_rel, aes('value')) + geom_histogram(binwidth=3)
    ggplot(pop_5_rel, aes('streams')) + geom_histogram()
    ggplot(pop_5_rel, aes('log_streams')) + geom_histogram()



def basemodel():
    pop_5_rel = pd.read_feather('../data/pop_5_release.feather')
    return(model_pipe(pop_5_rel))

def model_pipe(pop_5_rel):
    pop_5_rel['log_streams'] = np.log(pop_5_rel['streams'])
    # pop_5_rel = pop_5_rel[pop_5_rel['streams'] > 100]

    seed = 1

    x_train, x_test, y_train, y_test = train_test_split(pop_5_rel[['value']], pop_5_rel[['log_streams']], test_size = 0.20, random_state = 1)
    
    param_grid = {'polynomialfeatures__degree': np.arange(2), 'linearregression__normalize': [True, False]}
    grid = GridSearchCV(PolynomialRegression(), param_grid, cv = 5, error_score=mean_squared_error)

    model = fitAndEvaluate(x_train, y_train, x_test, y_test, grid, plot=True)
    return(model)

def linearModel():
    pop_5_rel['log_streams'] = np.log(pop_5_rel['streams'])
    # pop_5_rel = pop_5_rel[pop_5_rel['streams'] > 100]

    seed = 1
    x_train, x_test, y_train, y_test = train_test_split(pop_5_rel[['value']], pop_5_rel[['log_streams']], test_size = 0.20, random_state = 1)
    
    model = LinearRegression(fit_intercept=True, normalize=False)
    model.fit(x_train, y_train)

    model.coef_
    model.intercept_


    printModelEvaluations(model, x_train, y_train, x_test, y_test, plot=True)
    

def firstNonZeroPop():
    return(nthNonZeroPop(1))


def nthNonZeroPop(n):
    pop = pd.read_feather('../data/pop_with_streams.feather').drop('index', axis = 1)
    return(model_pipe(clean.datasetForNthNonZero(pop, n)))


def PolynomialRegression(degree=1, **kwargs):
    return make_pipeline(PolynomialFeatures(degree), LinearRegression(**kwargs))


def fitAndEvaluate(x,y,x_test,y_test,grid, plot=True):
    grid.fit(x,y)
    model = grid.best_estimator_
    printModelEvaluations(model, x, y, x_test, y_test, plot=plot)