import pandas as pd
from scipy.spatial.distance import cosine
import numpy as np
from scipy.sparse import csr_matrix
from sklearn.neighbors import NearestNeighbors

# display results to 3 decimal points, not in scientific notation
pd.set_option('display.float_format', lambda x: '%.3f' % x)

user_data = pd.read_table('./lastfm-dataset-360K/usersha1-artmbid-artname-plays_1.tsv',
                          header = None, nrows = 2e7,
                          names = ['users', 'musicbrainz-artist-id', 'artist-name', 'plays'],
                          usecols = ['users', 'artist-name', 'plays'])

user_profiles = pd.read_table('./lastfm-dataset-360K/usersha1-profile.tsv',
                          header = None,
                          names = ['users', 'gender', 'age', 'country', 'signup'],
                          usecols = ['users', 'country'])

# print(user_data.head())
# print(user_profiles.head())

# filter null artist name
if user_data['artist-name'].isnull().sum() > 0:
    user_data = user_data.dropna(axis = 0, subset = ['artist-name'])

# grouping and count the total plays for artists
artist_plays = (user_data.
     groupby(by = ['artist-name'])['plays'].
     sum().
     reset_index().
     rename(columns = {'plays': 'total_artist_plays'})
     [['artist-name', 'total_artist_plays']]
    )

# artist_plays.head()

# # merge with users
user_data_with_artist_plays = user_data.merge(artist_plays, left_on = 'artist-name', right_on = 'artist-name', how = 'left')
# user_data_with_artist_plays.head()

# # keep only 3% top popular artists
popularity_threshold = 40000
user_data_popular_artists = user_data_with_artist_plays.query('total_artist_plays >= @popularity_threshold')
# user_data_popular_artists.head()

# # filter only user in US
combined = user_data_popular_artists.merge(user_profiles, left_on = 'users', right_on = 'users', how = 'left')
usa_data = combined.query('country == \'United States\'')
# usa_data.head()

# # remove rows have same user name and artist-name, dont know why have these cases
if not usa_data[usa_data.duplicated(['users', 'artist-name'])].empty:
    initial_rows = usa_data.shape[0]
    
    print('Initial dataframe shape {0}'.format(usa_data.shape))
    usa_data = usa_data.drop_duplicates(['users', 'artist-name'])
    current_rows = usa_data.shape[0]
    print('New dataframe shape {0}'.format(usa_data.shape))
    print('Removed {0} rows'.format(initial_rows - current_rows))

# # start implement knn model
# # reshape, pivot data to mxn matrix - artitst x user
wide_artist_data = usa_data.pivot(index = 'artist-name', columns = 'users', values = 'plays').fillna(0)
wide_artist_data_sparse = csr_matrix(wide_artist_data.values)

# # save the model for future
from scipy.sparse import csr_matrix

def save_sparse_csr(filename,array):
    np.savez(filename,data = array.data ,indices=array.indices,
             indptr =array.indptr, shape=array.shape )

def load_sparse_csr(filename):
    loader = np.load(filename)
    return csr_matrix((loader['data'], loader['indices'], loader['indptr']),
                         shape = loader['shape'])

save_sparse_csr('./lastfm_sparse_artist_matrix.npz', wide_artist_data_sparse)

    # # model
    # # fitting the model
def recommend(artist):
    model_knn = NearestNeighbors(metric = 'cosine', algorithm = 'brute')
    model_knn.fit(wide_artist_data_sparse)

    # # making the recommendation
    # query_index = np.random.choice(wide_artist_data.shape[0])
    
    results = []

    if artist != None:
        query_index = int(artist)
        print('Recommendations for {0}:\n'.format(wide_artist_data.index[query_index]))
        distances, indices = model_knn.kneighbors(wide_artist_data.iloc[query_index, :].values.reshape(1, -1), n_neighbors = 6)

        for i in range(0, len(distances.flatten())):
            # if i == 0:
            #     print('Recommendations for {0}:\n'.format(wide_artist_data.index[query_index]))
            # else:
            #     print('{0}: {1}, with distance of {2}:'.format(i, wide_artist_data.index[indices.flatten()[i]], distances.flatten()[i]))
            if i > 0:
                results.append(wide_artist_data.index[indices.flatten()[i]])
    print(results)
    return results