import logging
import pandas as pd
import numpy as np
import time
from datetime import date
from imblearn.over_sampling import SMOTE
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 timing_val(func):
    
    '''
       Utility function that calculates the execution time of other functions.
       When @timing_val is present, the function will return a tuple containing result and execution time.
       From source: http://www.daniweb.com/code/snippet368.html
    '''
    
    def wrapper(*arg, **kw):
        t1 = time.time()
        res = func(*arg, **kw)
        t2 = time.time()
        return (t2 - t1), res, func.__name__
    return wrapper


# def fix_time_variables(df):
#     df['FIXED_RELEASE_DATE'] = pd.to_datetime(df['FIXED_RELEASE_DATE'])
#     df['MIN_CHART_DATE'] = pd.to_datetime(df['MIN_CHART_DATE'])
#     df['LATEST_HIT_DATE'] = pd.to_datetime(df['LATEST_HIT_DATE'])
#     df['Month_release'] = df['FIXED_RELEASE_DATE'].dt.month_name()
#     df['Month_chart'] = df['MIN_CHART_DATE'].dt.month_name()
#     df['Month_latest_chart'] = df['LATEST_HIT_DATE'].dt.month_name()
#     return df

def fix_time_variables(df):
    df['FIXED_RELEASE_DATE'] = pd.to_datetime(df['FIXED_RELEASE_DATE'])
    df['MIN_CHART_DATE'] = pd.to_datetime(df['MIN_CHART_DATE'], format='%Y-%m-%d')
    df['LATEST_HIT_DATE'] = pd.to_datetime(df['LATEST_HIT_DATE'], format='%Y-%m-%d')
    df['Month_release'] = df['FIXED_RELEASE_DATE'].dt.strftime('%b')
    df['Month_chart'] = df['MIN_CHART_DATE'].dt.strftime('%b')
    df['Month_latest_chart'] = df['LATEST_HIT_DATE'].dt.strftime('%b')
    return df

def prepare_sample(df):
    prefixes = ['USS1Z', 'QMS1Z', 'ISRC', 'QZS1Z',  'JMK40', 'AB'] 
    mask = df['ISRC'].str.startswith(tuple(prefixes))
    filtered = df[~mask]
    filtered['FIXED_RELEASE_DATE'] = filtered['FIXED_RELEASE_DATE'].dt.date
    sub = filtered[filtered['FIXED_RELEASE_DATE'] <= date(2024,2, 13)]
    country_sample = sub[(sub['COUNTRY_CODE'].isin(['US', 'GB'])) | (sub['CHART_COUNTRY']=='GB')]
    return country_sample

def artist_score_calculation(df):
    conditions = [(df['PREVIOUS_HIT']==0),
                 (df['PREVIOUS_HIT']==1) & (df['DAYS_SINCE_LAST_HIT'] > 365),
                (df['PREVIOUS_HIT']==1) & (df['DAYS_SINCE_LAST_HIT'] <= 365)]
    choices = ['no_previous_hit', 'non_recent_hit', 'recent_hit']
    df['artist_score'] = np.select(conditions, choices)
    return df

def fix_dataframe(df):
    df = fix_time_variables(df)
    #df = extract_country(df)
    df = artist_score_calculation(df)
    sub = prepare_sample(df)
    return sub


def prepare_data(data):
    data = fix_dataframe(data)
    # data = data.drop_duplicates(subset='ISRC', keep='first').reset_index(drop=True)
    #data = data.sort_values(by=['ISRC'])
    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']]
    data['MODE'] = data['MODE'].apply(str)
    numeric_features= data[['ACOUSTICNESS', 'DANCEABILITY', 'DURATION_MS', 'ENERGY', 'INSTRUMENTALNESS', 
                       'KEY', 'LIVENESS', 'LOUDNESS', 'SPEECHINESS', 'TEMPO', 'TIME_SIGNATURE', 'VALENCE']]
    categorical_features = data[['Month_release', 'artist_score']]
    binary_features = data[['MODE']]
    
    scaler = StandardScaler() 
    ordinal_encoder = OrdinalEncoder()
    label_encoder = LabelEncoder()
    
    X_processed = pd.DataFrame(scaler.fit_transform(numeric_features), columns = numeric_features.columns)
    
    encoded = ordinal_encoder.fit_transform(categorical_features)
    X_categorical1 = pd.DataFrame(encoded, columns = categorical_features.columns) 
    X_categorical2 = pd.DataFrame(label_encoder.fit_transform(binary_features), columns = binary_features.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 balancing_minor_class(new_X):
    (X_prepared, y_prepared) = prepare_data(new_X)
    X_resampled, y_resampled = SMOTE(random_state=2024, sampling_strategy='minority').fit_resample(X_prepared, y_prepared)
    return X_resampled, y_resampled

    
def split_train_test_data(new_X):
    
    """
    Cross validation with balanced classes
    
    """
    X_resampled, y_resampled = balancing_minor_class(new_X)
    
    X_train, X_test, y_train, y_test = train_test_split(X_resampled, y_resampled, 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)