from stream_predictor import streams

import pytest
TOL = 0.01

def test_no_playlist():
    pop_5 = 45
    playlists = {}

    assert streams(pop_5, playlists) == pytest.approx(97776.4, TOL)
    
def test_one_playlist():
    pop_5 = 45
    playlists = {'37i9dQZEVXbLRQDuF5jeBp' : 1 }

    assert streams(pop_5, playlists) == pytest.approx(246625, TOL)

def test_not_existing_playlist_raises_exception():
    pop_5 = 45
    playlists = {'doesnt_exist' : 1 }

    with pytest.raises(ValueError):
        streams(pop_5, playlists)

def test_two_playlist():
    pop_5 = 45
    playlists = {'6X35NoAOCzPvpYaq6r2dO3' : 1,'37i9dQZEVXbLRQDuF5jeBp' : 1 }

    assert streams(pop_5, playlists) == pytest.approx(189408, TOL)


def test_estimate_sheet():
    from io import StringIO
    data = StringIO("""
track id	pop_5	pop_5_date
3W1FUFK0mLvXDYRT0YGaiT	54	2018-06-14
4UbOwTSBizbJtTFqrBUGdM	51	2018-06-14
4vEPpdMcDRy3Zmofn40Po0	35	2018-06-14
6Dvxpqv61hyRZJs4qDfiP3	38	2018-06-14
6jaLFu5HvfNrI1JmcM4eXu	66	2018-06-13
6LtFWkakk2no03jeyznFrf	51	2018-06-14
6nj15YcIrPTfHAejsDGRY4	57	2018-06-13
7j12NwkG1VLC0b1P05QrA7	21	2018-06-14
7JBQaXZfEo0KHICuJmULmS	35	2018-06-13
""")
    estimates_sheet = pd.read_csv(data, sep='\t', skiprows=0)
    estimates_sheet['pop_5_date'] = pd.to_datetime(estimates_sheet['pop_5_date'])
    playlists = pd.read_feather('../../data/playlist_data.feather')
    estimates_sheet = pd.merge(estimates_sheet, playlists[['track_spyid', 'playlist_spyid', 'first_seen' ,'first_position']],left_on='track id', right_on='track_spyid')
    estimates_sheet = estimates_sheet[estimates_sheet['pop_5_date'] >= estimates_sheet['first_seen']]
    
    estimates_args = estimates_sheet.groupby('track_spyid').apply( lambda x: (x['pop_5'], dict(zip(x['playlist_spyid'], x['first_position']))))
    
    estimates_args.map(lambda x: streams(*x))