import numpy as np
import pandas as pd

def first_seen_dummies(pop_5_rel, fill_na = None):
    X = pop_5_rel[['track_spyid', 'first_position', 'playlist_spyid', 'pop_5']].copy()
    p5 = X[['track_spyid', 'pop_5']].copy()
    p5 = p5.drop_duplicates()

    X = X.drop('pop_5', axis = 1)
    aa = pd.pivot_table(X, values='first_position', index=['track_spyid','playlist_spyid'])
    spread = aa.unstack().reset_index()
    spread.columns = spread.columns.droplevel()
    spread = spread.rename(columns={'':'track_spyid'})
    spread = pd.merge(p5, spread, on='track_spyid', how='left')

    y = pop_5_rel[['track_spyid', 'log_streams']].drop_duplicates(subset = 'track_spyid')

    if fill_na is not None:
        spread = spread.fillna(fill_na)
        
    return spread.drop('track_spyid', axis =1), y[['log_streams']]

def preprocess(in_file):
    pop_5_rel = pd.read_feather(in_file)
    pop_5_rel['pop_5'] = pop_5_rel['value']
    pop_5_rel = pop_5_rel.drop('value', axis =1)
    pop_5_rel['log_streams'] = np.log(pop_5_rel['streams'])
    return pop_5_rel

def base_playlist_dummy(pop_5_rel):
    X = pop_5_rel.drop(['log_streams','streams','track_spyid'], axis=1)
    X = sm.add_constant(X)
    return X, pop_5_rel[['log_streams']]