# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Lyric tf-idf DEA #

# ## Imports, Inits, and Method definitions ##

# +
import pandas as pd
import numpy as np
from sqlalchemy import create_engine

import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
sns.set()

# %matplotlib inline

from sklearn.model_selection import train_test_split, cross_val_score, GridSearchCV, cross_validate
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_curve, roc_auc_score
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import GaussianNB
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn import preprocessing

import importlib

import mcnulty_methods
import word_utils
importlib.reload(mcnulty_methods);
importlib.reload(word_utils);
from mcnulty_methods import get_formatted_feature_df, get_lyrics_for_tracks
from word_utils import get_word_counts, generate_word_charts,get_word_total_idf

# +
mpl.rcParams['axes.titlesize'] = 16
mpl.rcParams['axes.labelsize'] = 16
mpl.rcParams['xtick.labelsize'] = 13
mpl.rcParams['ytick.labelsize'] = 13

test_size = 0.2
random_state = 10


# +
def get_artist_term_counts():
    term_counts = pd.read_csv('top_artist_terms.csv', index_col='artist_id', names=['artist_id','term', 'term_count'])
    term_counts = term_counts[~(term_counts['term'] == 'term')]
    del term_counts['term_count']
    return term_counts


def get_term_counts():
    return pd.read_csv('term_counts.csv', names=['term', 'count'])


# -

# ## Fetch Tracks for Particular Genres

conn = create_engine('postgresql://@localhost:5432/mcnulty_songs').raw_connection()
cursor = conn.cursor()

genres = ['hip hop', 'metal']
features = get_formatted_feature_df(conn, genres=genres)
features.set_index('track_id', inplace=True)

features.shape

features.sample(5)

# ## Fetch Lyrics from Tracks ##

# +
genre_labels = genres
unique_words = set()

all_lyrics = None
hiphop_lyrics = None
pop_lyrics = None

lyrics_by_genre = dict()
remaining_features = features

for genre_label in genre_labels:
    genre_df = remaining_features[(remaining_features['term'] == genre_label)]
    remaining_features = remaining_features[~(remaining_features.index.isin(genre_df.index))]
    genre_ids = genre_df.index
    
    genre_lyrics = get_lyrics_for_tracks(conn, genre_ids)
    del genre_lyrics['is_test']
    lyrics_by_genre[genre_label] = genre_lyrics
        
    if all_lyrics is None:
        all_lyrics = genre_lyrics
    else:
        all_lyrics = pd.concat([all_lyrics, genre_lyrics])
        
    #generate_word_charts(genre_lyrics, genre_label)
    

# +
def tf_calculate(series):
    return series / series.count()

total_unique_songs = all_lyrics.index.nunique()

all_lyrics['tf'] = all_lyrics.groupby('track_id')['count'].transform(lambda x: x / x.sum())
all_lyrics['word_document_count'] = all_lyrics.groupby('word')['count'].transform('count')
#all_lyrics['idf'] = all_lyrics.groupby('word')['count'].transform(lambda x: total_unique_songs / x.count())

all_lyrics['idf'] = np.log10(total_unique_songs / all_lyrics['word_document_count'])
all_lyrics['tf_idf'] = all_lyrics['idf'] * all_lyrics['tf']
all_lyrics[all_lyrics['word'] == '&'].head(5)


# -

# ## Word Analysis and Reshaping for Modeling ##

# ## Feature Selection ##
#
# Starting with the top x words found per song in the dataset, we'll add features and record the results from our classification models

def plot_ROC_compute_AUC(model, model_name, X,y):
    X_val, X_val_test, y_val, y_val_test = train_test_split(X, y, test_size=test_size, random_state=random_state)

    model.fit(X_val, y_val)
    y_prob = model.predict_proba(X_val_test)[:,1]
    auc = roc_auc_score(y_val_test, y_prob)
    
    return auc
    #TODO save these
    fpr, tpr, _ = roc_curve(y_val_test, y_prob)
    auc = roc_auc_score(y_val_test, y_prob)

    plt.plot(fpr, tpr)

    x = np.linspace(0,1, 100000)
    plt.plot(x, x, linestyle='--')

    plt.title('ROC Curve (Pop or Hip Hop)')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(['Logistic Regression'])
    return auc


def create_test_result_df(rows):
    test_results = pd.DataFrame(rows, columns=['num_words', 'use_tfidf', 'model','accuracy','precision','recall','f1','auc'])
    test_results['accuracy'] = test_results['accuracy'].astype(np.float64)
    test_results['use_tfidf'] = test_results['use_tfidf'].astype(np.bool)
    test_results['precision'] = test_results['precision'].astype(np.float64)
    test_results['recall'] = test_results['recall'].astype(np.float64)
    test_results['f1'] = test_results['f1'].astype(np.float64)
    test_results['auc'] = test_results['auc'].astype(np.float64)
    return test_results


def create_feature_df(rows):
    test_results = pd.DataFrame(rows, columns=['num_words','use_tfidf','model','feature','importance'])
    test_results['use_tfidf'] = test_results['use_tfidf'].astype(np.bool)
    test_results['importance'] = test_results['importance'].astype(np.float64)
    return test_results


def get_X_Y(word_sample_size, use_tfidf=False):
    word_song_appearance, total_word_appearance = get_word_counts(all_lyrics)

    if use_tfidf:
        word_song_appearance = get_word_total_idf(all_lyrics)

    
    word_subset = word_song_appearance.iloc[:word_sample_size]

    remaining_lyrics = pd.merge(all_lyrics.reset_index(), word_subset[['word']], how='right', on='word')

    remaining_lyrics.set_index('track_id', inplace=True)
    if use_tfidf:
        del remaining_lyrics['tf']
        del remaining_lyrics['word_document_count']
        del remaining_lyrics['idf']
        pivote_value = 'tf_idf'
    else:
        pivote_value= 'count'
        
    tid_lyrics = remaining_lyrics.pivot(columns='word', values=pivote_value)

    music_features = ['music_duration','music_key','music_loudness', 'music_mode', 'music_tempo', 'music_time_signature']
    
    term_only = features[['term']].reset_index().set_index('track_id')
    feature_names = list(tid_lyrics.columns)
    # complete set,= tid_index -> genre -> word_a -> .... -> word_z
    complete_set = pd.merge(term_only, tid_lyrics, left_index=True, right_index=True, how='right')
    complete_set.fillna(0, inplace=True)


    y_text = np.asarray(complete_set.iloc[:,0])
    y = np.array([1 if val=='hip hop' else 0 for val in y_text])
    X = np.asarray(complete_set.iloc[:,1:])

    return X,y,feature_names

features['term'].value_counts()

'''for i in range (0,2):
    print('Generating X and y')
    X, y, feature_names = get_X_Y(200,use_tfidf=i)
    X_val, X_test, y_val, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state)

    X_val_fit, X_val_test, y_val_fit, y_val_test = train_test_split(X_val, y_val, test_size=test_size, random_state=random_state)

    model = RandomForestClassifier(n_estimators=100, class_weight={1 : 1, 0 : 2})
    print('Fitting Model')
    model.fit(X_val_fit, y_val_fit)
    print('Predicting Values')
    y_test_pred = model.predict(X_val_test)

    print('Accuracy: {}'.format(accuracy_score(y_val_test, y_test_pred)))
    print('Recall: {}'.format(recall_score(y_val_test, y_test_pred)))
    print('Precision: {}'.format(precision_score(y_val_test, y_test_pred)))
    print('F1: {}'.format(f1_score(y_val_test, y_test_pred)))''';



# +
word_chunk_size = 10
word_upper_bound = 3010
# let's stop at 3000 words
feat_results = create_feature_df(None)
score_results = create_test_result_df(None)

tree100 = RandomForestClassifier(n_estimators=100, class_weight={1 : 1, 0 : 2})
tree1000 = RandomForestClassifier(n_estimators=1000, class_weight={1 : 1, 0 : 2})
tree_models = [tree100]
tree_model_names = ['tree100']
log_model = LogisticRegression(penalty='l1', class_weight={1 : 1, 0 : 2})


output_feature_file_name = 'feature_importance_csvs/feature_importance_results.csv'
output_score_file_name = 'feature_importance_csvs/feature_score_results.csv'

for word_sample_size in range(10, word_upper_bound, word_chunk_size):
    print('word sample size: {}'.format(word_sample_size))
    for i in range (1,2):
        print('\tis tfidf: {}'.format(i))
        X, y, feature_names = get_X_Y(word_sample_size,use_tfidf=i)
        X_val, X_test, y_val, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state)
        
        # For the pandas rows we're generating
        number_word_columns = [str(word_sample_size)] * len(feature_names)
        tf_idf_columns = [i] * len(feature_names)
        
        # Handle Logistic Regression
        print('\t\tlog fit')
        log_model.fit(X_val, y_val)
        model_columns = ['log'] * len(feature_names)

        zipped_features = list(zip(number_word_columns, tf_idf_columns, model_columns, feature_names, log_model.coef_[0]))
        new_results = create_feature_df(zipped_features)

        # Merge with our results
        feat_results = feat_results.append(new_results, ignore_index=True)
        feat_results.to_csv(output_feature_file_name)
        
        print('\t\tlog score')
        ## handle scores
        scoring = {'accuracy': 'accuracy',
               'precision': 'precision',
               'recall': 'recall',
               'f1': 'f1',}
        scores = cross_validate(log_model, X_val, y_val, scoring=scoring, cv=5, n_jobs=-1)
        accuracy = np.mean(scores['test_accuracy'])
        precision = np.mean(scores['test_precision'])
        recall = np.mean(scores['test_recall'])
        f1 = np.mean(scores['test_f1'])

        auc = plot_ROC_compute_AUC(log_model, 'log', X_val, y_val)
        
        cv_row = [word_sample_size, i, 'log', accuracy, precision, recall, f1, auc]
    
        score_results = score_results.append(create_test_result_df([cv_row]), ignore_index=True)
        score_results.to_csv(output_score_file_name)
        
        # Now our trees
        cv_rows = []
        for idx, tree in enumerate(tree_models):
            print('\t\t{} fit'.format(tree_model_names[idx]))
            tree.fit(X_val, y_val)
            model_columns = [tree_model_names[idx]] * len(feature_names)
            
            zipped_features = list(zip(number_word_columns, tf_idf_columns, model_columns, feature_names, tree.feature_importances_))
            new_results = create_feature_df(zipped_features)

            feat_results = feat_results.append(new_results,
                                               ignore_index=True)
            feat_results.to_csv(output_feature_file_name)
            print('\t\t{} score'.format(tree_model_names[idx]))
            
            scores = cross_validate(tree, X_val, y_val, scoring=scoring, cv=5, n_jobs=-1)
            accuracy = np.mean(scores['test_accuracy'])
            precision = np.mean(scores['test_precision'])
            recall = np.mean(scores['test_recall'])
            f1 = np.mean(scores['test_f1'])
            
            #accuracy = np.mean(cross_val_score(tree, X_val, y_val, cv=5, n_jobs=-1, scoring='accuracy'))
            #precision = np.mean(cross_val_score(tree, X_val, y_val, cv=5, n_jobs=-1, scoring='precision'))
            #recall = np.mean(cross_val_score(tree, X_val, y_val, cv=5, n_jobs=-1, scoring='recall'))
            #f1 = np.mean(cross_val_score(tree, X_val, y_val, cv=5, n_jobs=-1, scoring='f1'))
            auc = plot_ROC_compute_AUC(tree, tree_model_names[idx], X_val, y_val)

            cv_row = [word_sample_size, i, tree_model_names[idx], accuracy, precision, recall, f1, auc]
            cv_rows.append(cv_row)
            
        score_results = score_results.append(create_test_result_df(cv_rows), ignore_index=True)
        score_results.to_csv(output_score_file_name)
        
            
# -



StarcoderdataJupyterScriptsDedupFiltered