from stream_predictor.stream_predictor import streams_fill_na
from stream_predictor import streams

import pandas as pd
import numpy as np

from plotnine import *
from sklearn.externals import joblib

def sanity_check_streams_increase_with_lower_playlist_num():

    pop_range = np.arange(10, 100, 10)
    playlists = {}

    pop_prediction = [streams_fill_na(pop, playlists) for pop in pop_range]
    df = pd.DataFrame({'pop' : pop_range, 'predicted_streams_log' : np.log(pop_prediction)})
    ggplot(df, aes('pop', 'predicted_streams_log')) + geom_point()

    pop_prediction = [streams(pop, playlists) for pop in pop_range]
    df = pd.DataFrame({'pop' : pop_range, 'predicted_streams_log' : np.log(pop_prediction)})
    ggplot(df, aes('pop', 'predicted_streams_log')) + geom_point()
    

    model_spec_file =  'stream_predictor/models/tpot_20_20_xgboos_fill_na_1000.pkl'
    model_spec = joblib.load(model_spec_file)
    feature_columns = model_spec['feature_cols']

    all_playlists = np.delete(feature_columns,0)
    
    playlist_num_range = np.arange(1, 100, 10)

    # playlist_id = '37i9dQZEVXbLRQDuF5jeBp'

    preds = []

    for playlist_id in all_playlists:
        print(playlist_id)
        for pp in pop_range:
            for pl in playlist_num_range:
                preds.append(pd.DataFrame({ 'pop': pp, 'playlists_num' : pl, 'playlist' : playlist_id, 'streams' : streams_fill_na(pp, {playlist_id : pl})}, index = [0]) )

    all_preds = pd.concat(preds, ignore_index=True)
    all_preds['log_streams'] = np.log(all_preds['streams'])

    all_preds.to_csv('all_preds.csv')

    all_preds = pd.read_csv('all_preds.csv')

    by_pop = all_preds.groupby(['pop', 'playlists_num'])['streams'].mean().reset_index()

    ggplot(by_pop, aes('playlists_num', 'streams')) + geom_point() + facet_wrap('~pop') + scale_y_log10()


def sanity_check_streams_increase_with_higher_count_of_playlist_ids():

    pop_range = np.arange(10, 100, 10)

    model_spec_file =  'stream_predictor/models/tpot_20_20_xgboos_fill_na_1000.pkl'
    model_spec = joblib.load(model_spec_file)
    feature_columns = model_spec['feature_cols']

    all_playlists = np.delete(feature_columns,0)
    number_of_playlists = np.array([1,2,3,4,5,10,15,20,40])
    playlist_position = 1

    samples = 100
    
    preds = []

    import warnings
    from sklearn.exceptions import DataConversionWarning
    warnings.filterwarnings(action='ignore', category=DataConversionWarning)

    for ss in np.arange(samples):
        print(f'Sample: {ss}')
        for num_playlists in number_of_playlists:
            playlists = {x:playlist_position for x in np.random.choice(all_playlists, num_playlists)}
            for pop in pop_range:
                preds.append(pd.DataFrame({ 'pop': pop, 
                    'number_of_playlists' : num_playlists, 
                    'streams' : streams_fill_na(pop, playlists)}, index = [0]) )
                
            
    all_preds = pd.concat(preds, ignore_index=True)
    all_preds['log_streams'] = np.log(all_preds['streams'])

        
    by_pop = all_preds.groupby(['pop', 'number_of_playlists'])['streams'].median().reset_index()
    ggplot(by_pop, aes('number_of_playlists', 'streams')) + geom_point() + facet_wrap('~pop') + scale_y_log10()