import pandas as pd
import numpy as np
pd.set_option('display.max_columns', 500)
from plotnine import *

def create_ml_data():
    pop = pd.read_feather('../data/pop_with_streams.feather').drop('index', axis = 1)

    release_dates = pd.read_feather('../data/release_dates.feather')

    pop_rel = pd.merge(pop, release_dates, left_on='track_spyid', right_on='spyid')

    pop_rel['days_since_release']=((pop_rel['as_of']-pop_rel['release_date'])/pd.Timedelta('1 d'))

    #TODO: Why are these null?
    pop_rel = pop_rel[~pop_rel['days_since_release'].isnull()]
    pop_rel['days_since_release'] = pop_rel['days_since_release'].astype(int)


    # pop_rel['diff_release_first_seen']=((pop_rel['first_seen']-pop_rel['release_date'])/pd.Timedelta('1 d'))
    # pop_rel[pop_rel['days_since_release'].isnull()]
    # pop_rel[pop_rel['days_since_release'] < 0].head()

    pop_rel_5 = pop_rel[pop_rel['days_since_release'] == 5]
    pop_rel_5 = pop_rel_5.drop_duplicates('spyid', keep='last')

    streams = pd.read_feather('../data/raw_stream.feather')
    pop_rel_5 = pd.merge(pop_rel_5, streams, on = 'spyid')

    pop_rel_5['pop_5_date'] = pop_rel_5['as_of_x'].dt.date
    pop_rel_5['streams_100_date'] = pop_rel_5['as_of_y'].dt.date

    pop_rel_5 = pop_rel_5.drop(['as_of_x', 'as_of_y'], axis =1)

    pop_rel_5.to_feather('../data/pop_5_release.feather')

    
def create_extended_ml_data():
    pop_5 = pd.read_feather('../data/pop_5_release.feather')
    playlist_data = pd.read_feather('../data/playlist_data.feather')

    pp5 = pop_5[['track_spyid', 'pop_5_date']]

    pp5 = pd.merge(pp5, playlist_data, on='track_spyid', how='left')
    pp5 = pp5[pp5['pop_5_date'] >= pp5['first_seen']]
    pp5 = pp5[['track_spyid', 'playlist_spyid', 'first_seen', 'first_position']]

    agg_data = pp5.groupby('track_spyid')['first_position'].agg(['min', 'max', np.mean, np.median, len])
    agg_data.columns = ['min_pp', 'max_pp', 'mean_pp', 'median_pp', 'num_pp']

    playlist_agg = pd.merge(pop_5, agg_data, on='track_spyid', how='left')

    playlist_agg.to_feather('../data/pop_5_release_with_playlist_agg.feather')

    pp5['ss'] = 1

    playlist_dummy = pp5.pivot_table(index='track_spyid', columns='playlist_spyid', values='ss', fill_value = 0).reset_index()
    
    pop_5_play_dummy = pd.merge(pop_5[['track_spyid','value','streams']], playlist_dummy, on='track_spyid', how='left')
    pop_5_play_dummy = pop_5_play_dummy.fillna(0)

    pop_5_play_dummy.to_feather('../data/pop5_play_dummy.feather')
    

    pop_5_unstacked = pd.merge(pop_5[['track_spyid','value','streams']], pp5[['track_spyid', 'first_position', 'playlist_spyid']], on = 'track_spyid', how='left')
    pop_5_unstacked.to_feather('../data/pop_5_with_playlist_normalised.feather')


    # pop_5[pd.isnull(pop_5['min'])]





def charts():
    ggplot(pop_rel, aes('days_since_release', 'value', group = 'spyid')) + geom_line()
    ggplot(pop_rel[pop_rel['days_since_release'] < 110], aes('days_since_release', 'value', group = 'spyid')) + geom_line(aes(color = 'spyid'))
    ggplot(pop_rel[pop_rel['days_since_release'] < 10], aes('days_since_release', 'value', group = 'spyid')) + geom_line(aes(color = 'spyid'))


def calculateDaysAfterNonZeroPop(pop):

    pop_0 = pop.copy()

    pop_0.set_index('as_of', inplace=True)

    first_pop = pop_0.groupby('track_spyid').apply(lambda x: x[:1])
    ggplot(first_pop, aes('value')) + geom_histogram()

    start = 0
    second_pop=  pop_0.groupby('track_spyid').apply(lambda x: x[start:(start + 1)])
    ggplot(second_pop, aes('value')) + geom_histogram() + ggtitle(f'Start: {start}')


def bla():
    pop = pd.read_feather('../data/pop_with_streams.feather').drop('index', axis = 1)
    datasetForNthNonZero(pop, 2)


def datasetForNthNonZero(pop, n):

    pop_0 = pop.copy()

    non_zero_start = pop_0.groupby('track_spyid').apply(lambda x: x[x['value'] > 0]).drop('track_spyid', axis = 1)
    non_zero_start.reset_index(inplace = True)
    nn_zer0_start = non_zero_start.sort_values(['track_spyid', 'as_of']).groupby('track_spyid').nth(n)

    release_dates = pd.read_feather('../data/release_dates.feather')

    pop_rel = pd.merge(nn_zer0_start, release_dates, left_on='track_spyid', right_on='spyid')

    pop_rel['days_since_release']=((pop_rel['as_of']-pop_rel['release_date'])/pd.Timedelta('1 d'))

    #TODO: Why are these null?
    pop_rel = pop_rel[~pop_rel['days_since_release'].isnull()]
    pop_rel['days_since_release'] = pop_rel['days_since_release'].astype(int)

    streams = pd.read_feather('../data/raw_stream.feather')
    pop_rel = pd.merge(pop_rel, streams, on = 'spyid')
    pop_rel.to_feather(f'../data/pop_with_nth_non_zero_{n}.feather')
    return(pop_rel)