import logging
import pandas as pd
import numpy as np

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder, LabelEncoder, OrdinalEncoder


def prepare_data(data):
    data = data.drop_duplicates(subset='ISRC', keep='first').reset_index(drop=True)
    X = data.drop(columns = ['ISRC','MIN_RELEASE_DATE', 'LATEST_HIT_DATE', 'DAYS_SINCE_LAST_HIT', 'DAYS_TO_CHART', 'FIXED_RELEASE_DATE', 'ARTIST_NAME'])
    y = data[['HIT']]
    # In the paper (2018) the following traits were used: Danceability, Energy, Speechiness, Acousticness, Instrumentalness, Liveness, Valence, Loudness, and Tempo.
    # numeric_features= data[['ACOUSTICNESS', 'DANCEABILITY', 'ENERGY', 'INSTRUMENTALNESS', 
    #                    ''LIVENESS', 'LOUDNESS', 'SPEECHINESS', 'TEMPO', 'VALENCE']]
    numeric_features= data[['ACOUSTICNESS', 'DANCEABILITY', 'DURATION_MS', 'ENERGY', 'INSTRUMENTALNESS', 
                       'KEY', 'LIVENESS', 'LOUDNESS', 'MODE', 'SPEECHINESS', 'TEMPO', 'TIME_SIGNATURE', 'VALENCE']]
    categorical_features1 = data[['Month_release']]
    categorical_features2 = data[['artist_score']]
    
    scaler = StandardScaler() 
    label_encoder1 = OrdinalEncoder()
    label_encoder2 = OrdinalEncoder()
    
    X_processed = pd.DataFrame(scaler.fit_transform(numeric_features), columns = numeric_features.columns)
    
    encoded1 = label_encoder1.fit_transform(categorical_features1)
    X_categorical1 = pd.DataFrame(encoded1, columns = categorical_features1.columns) 

    X_categorical2 = pd.DataFrame(label_encoder2.fit_transform(categorical_features2), columns = categorical_features2.columns)
    
    new_X = pd.concat([X_processed.reset_index(), X_categorical1.reset_index(), X_categorical2.reset_index()], axis=1)
    new_X = new_X.drop(columns=['index'], axis=1)
    return (new_X , y)

def split_train_test_data(X, y):
    
    """
    Cross validation with balanced classes
    
    """
    
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=2023, stratify=y)
    
    # checking if y is balanced
    train_0, train_1 = len(y_train[y_train==0]), len(y_train[y_train==1])
    test_0, test_1 = len(y_test[y_test==0]), len(y_test[y_test==1])
    logging.info('> Train: 0=%d, 1=%d, Test: 0=%d, 1=%d' % (train_0, train_1, test_0, test_1))
    
    return (X_train, y_train, X_test, y_test)