# Dependencies
import os
import sys 

# comment when running locally
import subprocess
subprocess.check_call('pip install -r /opt/ml/processing/input/dependencies/requirements-2.txt', shell=True)

import logging
logging.basicConfig(
   level=logging.INFO,
   format='%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s',
   datefmt='%Y-%m-%d %H:%M:%S',
)

import pandas as pd
import numpy as np
import boto3 
import time
from datetime import datetime, date, timedelta

def distribution_mean(df, var):
    
    ''' function that calculates mean '''
    
    return df[var].describe()[1]
    
def distribution_25th_percentile(df, var):
    
    ''' function that calculates 25th percentile '''
    
    return df[var].describe()[4]

def distribution_median(df, var):
    
    ''' function that calculates median '''
    
    return df[var].describe()[5]

def distribution_75th_percentile(df, var):
    
    ''' function that calculates 75th percentile '''
    
    return df[var].describe()[6]

def descriptive_stats(df, var):
    
    ''' function that calculates descriptive statistics '''
    
    mean = distribution_mean(df, var)
    perc_25 = distribution_25th_percentile(df, var)
    median = distribution_median(df, var)
    perc_75 = distribution_75th_percentile(df, var)
    return (mean, perc_25, median, perc_75)


def manual_clustering(df):
    
    ''' 
    Function that implements logic for manual clustering
    It calculates an overall positive/negative trend for ISRC, based on sum of forecast errors
    It assigns a prediction label
    It assigns a trend label
    It compiles the table to save
    
    '''
    
    # create country variable
    #df['TRANSACTION_COUNTRY_CODE'] = 'GB'
    # create sign to identify whether sum of forecast errors is overall positive or negative
    df['sign'] = np.where(df['sum_forecast_errors'] > 0, 'overall_positive', 'overall_negative')

    (mean_fe, perc_25_fe, median_fe, perc_75_fe) = descriptive_stats(df, 'sum_forecast_errors')

    # Create new variable that classifies how an ISRC is performing 
    conditions_fe = [
        (df['sum_forecast_errors']<= perc_75_fe),
        ((df['sum_forecast_errors']> perc_75_fe) & (df['sum_forecast_errors']<= mean_fe)),
        (df['sum_forecast_errors'] > mean_fe)]

    choices_fe = ['as_predicted','marginally_better_than_predicted','significantly_better_than_predicted']

    df['performance_test_period']  = np.select(conditions_fe, choices_fe)

    (mean_per, perc_25_per, median_per, perc_75_per) = descriptive_stats(df, 'linear_gradient_test')

    conditions_per = [
        (df['linear_gradient_test']<= perc_25_per),
        ((df['linear_gradient_test']> perc_25_per) & (df['linear_gradient_test']<= median_per)),
        ((df['linear_gradient_test']> median_per) & (df['linear_gradient_test']<= perc_75_per)),
        (df['linear_gradient_test'] > perc_75_per)]

    choices_per = ['down_trend','slight_down_trend','slight_up_trend','up_trend']

    df['slope_test_period']  = np.select(conditions_per, choices_per)
    
    # take subset of table to save
    table_to_save = df[['ISRC_KEY','ISRC', 'TRANSACTION_COUNTRY_CODE','avg_streams_train','avg_streams_test','sum_forecast_errors', 'performance_test_period', 'slope_test_period']].copy()

    return table_to_save


def save_dataframe_s3(df, folder, prefix, run_id):
    s3 = boto3.client('s3')
    bucket_name = 'dev-cucumbers'
    filepath = "{}/CLUSTERING_{}_{}.csv".format(folder, prefix, run_id)
    csv_buffer = df.to_csv(index=False).encode('utf-8')
    # Save the CSV file to S3
    s3.put_object(Body=csv_buffer, Bucket=bucket_name, Key=filepath)
    print(f"Table saved to S3 bucket: {bucket_name}, with file name: {filepath}")

def get_argument(args, name, default_value = None):  
    ''' Getting arguments from processing script ''' 
    arg_name = "--" + name
    if arg_name in args:
        index = args.index("--" + name) + 1
        return args[index]
    else:
        return default_value

    
def get_mandatory_argument(args, name):
    ''' Getting mandatory arguments from processing script '''
        
    res = get_argument(args, name)
    if res is None:
        raise "Missing script mandatory argument --{}".format(name)
    else:
        return res

if __name__ == '__main__':

    # Run id
    RUN_ID = get_mandatory_argument(sys.argv, "run-id")
    # Folder
    FOLDER = get_mandatory_argument(sys.argv, "folder")
    PREFIX = 'arima_cross'
    INPUT_FILE = 's3://dev-cucumbers/{}/{}_{}_merged.csv'.format(FOLDER, PREFIX, RUN_ID)
    
    df = pd.read_csv(INPUT_FILE)
    
    cluster_df = manual_clustering(df)
    save_dataframe_s3(cluster_df, FOLDER, PREFIX, RUN_ID)