import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = (12, 7)

import statsmodels.api as sm
from sklearn.model_selection import train_test_split
from statsmodels.sandbox.regression.predstd import wls_prediction_std

from plotnine import *

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


def main():
    streams = pd.read_feather('../data/archive_as_of_friday_20_july/train_streams_with_release.feather')
    streams = streams[['track_id', 'date', 'all_streams', 'all_streams_cumsum', 'days_from_release']]

    streams_at_day_1 = streams[streams['days_from_release'] == 1][['track_id', 'all_streams_cumsum']].rename(columns = {'all_streams_cumsum' : 'cumsum_day_1'})
    streams_at_day_100 = streams[streams['days_from_release'] == 100][['track_id', 'all_streams_cumsum']].rename(columns = {'all_streams_cumsum' : 'cumsum_day_100'})
    streams_1_100 = pd.merge(streams_at_day_1, streams_at_day_100, on = 'track_id')
    streams_1_100['cumsum_2to100'] = streams_1_100['cumsum_day_100'] - streams_1_100['cumsum_day_1']

    streams_1_100.to_feather('../data/archive_as_of_friday_20_july/streams_1_100.feather')
    model = sm.OLS(np.log(streams_1_100['cumsum_2to100']), np.log(streams_1_100['cumsum_day_1']))
    model.fit().summary()


    pop = pd.read_feather('../data/archive_as_of_friday_20_july/pop_60_cleaned_with_days_after_release.feather')
    pop_5_date = pop[pop['days_after_release'] == 5][['track_id', 'date','popularity']]
    pop_5_date = pop_5_date.rename(columns = {'date':'date_5'})

    ss = pd.merge(streams, pop_5_date, on = 'track_id')

    ss_5_cumsum = ss[ss['date'] == ss['date_5']][['track_id', 'all_streams_cumsum', 'popularity']]
    ss_5_cumsum.columns = ['track_id', 'streams_cumsum_5', 'pop_5']

    ss_100_cumsum = ss[ss['days_from_release'] == 100][['track_id', 'all_streams_cumsum']]

    sp = pd.merge(ss_5_cumsum, ss_100_cumsum, on = 'track_id')

    sp['cum_sum_after_pop_5'] = sp['all_streams_cumsum'] - sp['streams_cumsum_5']
    sp['log_cum_sum_after_pop_5'] = np.log(sp['cum_sum_after_pop_5'])
    sp['log_cum_sum_5'] = np.log(sp['streams_cumsum_5'])
    sp['log_all_streams_cumsum'] = np.log(sp['all_streams_cumsum'])
    sp['fraction_of_stream_pre_5'] = sp['streams_cumsum_5']/(sp['streams_cumsum_5'] + sp.all_streams_cumsum)

    sp.to_feather('../data/archive_as_of_friday_20_july/streams_post_day_5_pop.feather')

    ggplot(sp, aes('pop_5', 'log_cum_sum_after_pop_5')) + geom_point() + geom_point(aes('pop_5', 'log_all_streams_cumsum'), color = "red")

    ggplot(sp, aes('log_cum_sum_5', 'log_cum_sum_after_pop_5')) + geom_point() + geom_abline(slope=1)


    model = sm.OLS(sp['log_cum_sum_after_pop_5'], sp['pop_5'])
    results = model.fit()
    results.summary()
    res = results