import pandas as pd
import numpy as np

import json

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

from plotnine import *

def main():

    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']]

    playlist_counts = pp5.groupby('playlist_spyid').size().to_frame('size')
    playlist_counts.describe()

    with open('../modelling/model_spec/base_dummy_playlist.json') as pc:
        playlist_coefficients = json.load(pc)

    pc = pd.DataFrame({ 'playlist_spyid':playlist_coefficients.keys(), 'coeff': playlist_coefficients.values() })

    pc = pd.DataFrame({ 'playlist_spyid': list(playlist_coefficients.keys()), 'coeff': list(playlist_coefficients.values())})

    playlist_counts = pd.merge(playlist_counts, pc, on='playlist_spyid')
    playlist_counts.sort_values('size')
    playlist_counts.sort_values('coeff')

    ggplot(playlist_counts, aes('size','coeff')) + geom_point() + scale_x_log10()
    
    pp5[pp5['playlist_spyid'] == '37i9dQZEVXbKM896FDX8L1']