import subprocess

subprocess.check_call('pip install -r /opt/ml/processing/input/dependencies/requirements-OLS.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 snowflake.connector
import pandas as pd
import boto3
import re
import logging
from absl import logging

import sys

import numpy as np
import time

import statsmodels.api as sm
from sklearn.linear_model import LinearRegression
from multiprocessing import Pool

##################################################################################################################
############################################# SOURCE TABLE #######################################################
##################################################################################################################

log_level = "DEBUG"
ticket_code = "EXP_1"

logging.set_verbosity(log_level)
logging.debug("READY!!!")

sec_id = 'dev/sagemaker-notebook-instance/SNOWFLAKE_PASSWORD'


def get_secret_value(name, version=None):
    """Gets the value of a secret.

    Version (if defined) is used to retrieve a particular version of
    the secret.

    """
    secrets_client = boto3.client("secretsmanager", region_name='us-east-1')
    kwargs = {'SecretId': name}
    if version is not None:
        kwargs['VersionStage'] = version
    response = secrets_client.get_secret_value(**kwargs)
    return response


def get_snowflake_creds(username="SAGEMAKER", account="orchard",
                        warehouse="DEV_OWS_ENGINEERING"):
    """
    Fetches and returns snowflake creds for connecting to snowflake

    Please use this within the scope of a function if using this on a shared instance
    This is so that the password is in memory only when its needed and gets dropped
    once its no longer required.

    returns:
    - creds (dict) - a dictionary containing user creds

    """
    creds = {
        "user": username,
        "password": get_secret_value(sec_id)['SecretString'],
        "account": "orchard",
        "warehouse": warehouse,
        "protocol": 'https'
    }
    return creds


def snowflake_connector_factory(creds=None):
    """
    A Factory for creating snowflake connectors.

    This returns the cursor after opening a session with snowflake.

    params:
    - creds - snowflake credentials

    returns:
    - cursor - snowflake session cursor
    """
    try:
        if creds:
            _creds = creds
        else:
            _creds = get_snowflake_creds()
        return snowflake.connector.connect(**_creds).cursor()
    except Exception as e:
        logging.error(f"Something went wrong - {str(e)}")


def _is_version_number(s):
    "Check and returns true if its a version number"
    return re.search("^[0-9][.0-9]*[0-9]$", s) is not None


def test_connection():
    """ tests connection to snowflake """
    with snowflake_connector_factory() as cs:
        try:
            cs.execute("SELECT current_version()")
            one_row = cs.fetchone()
            assert len(one_row) == 1
            assert _is_version_number(one_row[0])
            logging.info(f"Your snowflake version - {one_row[0]} PASSED!")
        except Exception as e:
            logging.error(f"Something went wrong - {str(e)}")


##################################################################################################################
############################################# FOURIER ############################################################
##################################################################################################################

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 detect_uptrend(time_series, window=7):
    moving_avg = time_series.rolling(window=window).mean()
    X = np.arange(len(moving_avg)).reshape(-1,1)
    y = moving_avg.dropna().values
    X2 = sm.add_constant(X[len(X) - len(y):])
    model = sm.OLS(y, X2).fit()
    return (model.params[1], model.pvalues[1])

def detect_uptrend_quadratic(df, var='COLLECTION_STREAMS', window=7):
    df['time'] = df['ACTIVITY_DATE'].apply(lambda x: x.toordinal())
    df['time_sq'] = df['time']**2
    df['MA7'] = df['COLLECTION_STREAMS'].rolling(window=7).mean()
    df = df.sort_values(by='ACTIVITY_DATE').reset_index()
    df = df.dropna()
    X = df[['time', 'time_sq']]
    y = df['MA7'].values
    X = sm.add_constant(X)
    model = sm.OLS(y, X).fit()
    coef_time = model.params['time']
    coef_time_sq = model.params['time_sq']
    derivative1 = coef_time + ((2*coef_time_sq)*len(df))
    derivative2 = 2 * coef_time_sq 
    return (model.params['time'], 
            model.pvalues['time'], 
            model.params['time_sq'], 
            model.pvalues['time_sq'],
            derivative1,
            derivative2)


def compile_coeff_table_ols(df, variable1='COLLECTION_STREAMS', variable2='SEARCH_STREAMS', window=7):
    res = pd.DataFrame(columns=['ISRC_KEY', 'COEFF_collection', 'p-value_collection', 'COEFF_search', 'p-value_search'])
    unique_keys = df['ISRC_KEY'].unique()
    for subdf in iterate_group_by_key(df, ['ISRC_KEY']):
        df_isrc = subdf
        isrckey = df_isrc['ISRC_KEY'].iloc[0]
        df_isrc = df_isrc.sort_values(by=['ACTIVITY_DATE'])

        try:
            time_series_collection = df_isrc.sort_values(by='ACTIVITY_DATE')[variable1]
            (coeff_collection, pvalue_collection) = detect_uptrend(time_series_collection, window)
            time_series_search = df_isrc.sort_values(by='ACTIVITY_DATE')[variable2]
            (coeff_search, pvalue_search) = detect_uptrend(time_series_search, window)
            res = pd.concat([res, pd.DataFrame([{
                'ISRC_KEY': isrckey,
                'COEFF_collection': coeff_collection,
                'p-value_collection': pvalue_collection,
                'COEFF_search': coeff_search,
                'p-value_search': pvalue_search
            }], columns=res.columns)])
        except Exception as e:
            logging.error(f"Error processing ISRC_KEY {isrckey}: {e}")
    return res

def compile_coeff_table_ols_quadratic(df, window=7):
    res = pd.DataFrame(columns=['ISRC_KEY', 
                                'COEFF_collection_time_term', 
                                'p-value_collection_time_term', 
                                'COEFF_collection_timesq_term', 
                                'p-value_collection_timesq_term',
                                'first_derivativ_collection',
                                'second_derivative_collection_timesq',
                                'COEFF_search_time_term', 
                                'p-value_search_time_term',
                                'COEFF_search_timesq_term', 
                                'p-value_search_timesq_term',
                                'first_derivativ_search',
                                'second_derivative_search_timesq'
                               ])
    unique_keys = df['ISRC_KEY'].unique()
    for subdf in iterate_group_by_key(df, ['ISRC_KEY']):  
        df_isrc = subdf
        isrckey    = df_isrc['ISRC_KEY'].iloc[0]
        df_isrc = df_isrc.sort_values(by=['ACTIVITY_DATE'])
        
        try:
            #time_series_collection = df_isrc.sort_values(by='ACTIVITY_DATE')[variable1]
            (coeff_collection_t, pvalue_collection_t, coeff_collection_t2, pvalue_collection_t2, first_derivative_collection, second_derivative_collection_t2) = detect_uptrend_quadratic(df_isrc, var= 'COLLECTION_STREAMS', window=7)
            (coeff_search_t, pvalue_search_t, coeff_search_t2, pvalue_search_t2, first_derivative_search, second_derivative_search_t2) = detect_uptrend_quadratic(df_isrc, var= 'SEARCH_STREAMS', window=7)
            res = pd.concat([res, pd.DataFrame([{
                'ISRC_KEY': isrckey, 
                'COEFF_collection_time_term': coeff_collection_t,
                'p-value_collection_time_term': pvalue_collection_t,
                'COEFF_collection_timesq_term': coeff_collection_t2,
                'p-value_collection_timesq_term':pvalue_collection_t2, 
                'first_derivativ_collection': first_derivative_collection,
                'second_derivative_collection_timesq':second_derivative_collection_t2,
                'COEFF_search_time_term': coeff_search_t, 
                'p-value_search_time_term': pvalue_search_t,
                'COEFF_search_timesq_term': coeff_search_t2, 
                'p-value_search_timesq_term': pvalue_search_t2,
                'first_derivativ_searcg': first_derivative_search,
                'second_derivative_search_timesq':second_derivative_search_t2,
                }], columns=res.columns)])
        except Exception as e:
            print('Error')
    return res


def compute_regression_for_chunk(chunk_df, chunk_id, par_id):
    '''
    Compute FastFourierTransform for a the given chunk of data, and save the data to S3

        chunk_df -- current dataframe for running chunk
        folder -- where work is stored
        run_id -- ID of current run
        chunk_id -- ID of current chunk of data running
        par_id -- parallel job number
        dry_run -- bolean flag indicating whether the results are saved.
    '''

    logging.info(f"Starting to compute linear regression on data chunk {chunk_id} and parallel job {par_id,} ...")
    #regression_table = compile_coeff_table_ols(chunk_df, variable1='COLLECTION_STREAMS', variable2='SEARCH_STREAMS', window=7)
    regression_table = compile_coeff_table_ols_quadratic(chunk_df, window=7)
    logging.info(f"Linear regression on data chunk {chunk_id} and parallel job {par_id,}  finished.")
    logging.info(f"Saving linear regression data to S3 ...")
    return regression_table


def iterate_group_by_key(df, col_for_key, sort_data=True):
    '''
       Optimized function (generator) that iterates over the DataFrame and groups it by a specified column.
       It operates on a sorted dataframe, so all rows for a given key (here ISRC) are together.
       It yields one group at a time as a new DataFrame.
    '''

    if sort_data:
        df = df.sort_values(by=col_for_key)

    # helper function that returns the value in the column specified by col_for_key at the index i.
    def key_at(i):
        return np.array(df[col_for_key].iloc[i])

    index, size = (0, df.shape[0])
    while index < size:
        current_key = key_at(index)
        res = []
        while index < size and list(current_key) == list(key_at(index)):
            res.append(df.iloc[index])
            index = index + 1
        resdf = pd.DataFrame(res, columns=df.columns)
        yield resdf


##################################################################################################################
############################################# SAVING TABLES ######################################################
##################################################################################################################


def save_dataframe_s3(df, folder, prefix, run_id):
    '''
        Save the given parallel chunk result into S3:
        df -- dataframe to save
        folder -- where work is stored
        prefix -- assigned prefix for file name
        run_id -- ID of current run
    '''
    s3 = boto3.client('s3')
    bucket_name = 'dev-cucumbers'
    filepath = "{}/{}_{}_merged.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}")


##################################################################################################################
############################################# UTILS ##############################################################
##################################################################################################################


def check_one_track_df(one_track_df):
    '''Checking that the length of the individual time series is 140 days'''

    return len(one_track_df) == 183


def unique_isrc_df(isrc, df):
    '''
       As time series analysis is univariate, decomposing the dataframe as to have one-track dataframes
       at time of analysis.
    '''

    one_track_df = df[df['ISRC_KEY'] == isrc].copy()
    one_track_df['ACTIVITY_DATE'] = pd.to_datetime(one_track_df['ACTIVITY_DATE'])
    one_track_df['ACTIVITY_DATE'] = one_track_df['ACTIVITY_DATE'].dt.date
    one_track_df = one_track_df[['ISRC_KEY', 'ISRC', 'ACTIVITY_DATE', 'COLLECTION_STREAMS', 'SEARCH_STREAMS', 'TRANSACTION_COUNTRY_CODE']]
    return one_track_df.sort_values(by='ACTIVITY_DATE')


def data_for_chunk(df, chunk_id, chunk_count):
    '''
       Function sorting ISRCs and splitting them into n chuncks
       and taking only one of them (chunk_id)
    '''

    IDs = df['ISRC_KEY'].unique().tolist()
    IDs.sort()
    chunk_size = len(IDs) // chunk_count
    chunked = [IDs[n:n + chunk_size] for n in range(0, len(IDs), chunk_size)]
    ids = chunked[chunk_id]
    chunk_df = df[df['ISRC_KEY'].isin(ids)]
    return chunk_df


def split_universe(df, n):
    '''
       Function splitting dataframe for parallelism

       df -- initial dataframe
       n -- number of chunks
    '''

    # list of sorted ISRCs into n chunks
    def chunks(l, n):
        for i in range(0, len(l), n):
            yield l[i:i + n]

    isrcs = df['ISRC_KEY'].unique().tolist()
    isrcs.sort()
    isrc_chunks = chunks(isrcs, (len(isrcs) // n) + 1)
    return list(isrc_chunks)


def merge_parallel_results(parallel_results):
    return pd.concat(parallel_results)


@timing_val
def split_work(full_df, n, f, folder, run_id, chunk_id, dry_run):
    '''
        full_df -- full dataframe that will be split into n parallel tasks
        n -- number of parallel tasks
        f -- function handling one parallel task
        folder -- where work is stored
        run_id -- ID of current run
        chunk_id -- ID of current chunk of data running
        dry_run -- bolean flag indicating whether the results are saved.
    '''

    isrc_chunks = split_universe(full_df, n)
    splitted_df = []
    parallel_results = []
    # For each chunk, creating an individual dataframe that only includes rows with the same ID
    for chunk in isrc_chunks:
        df_chunk = full_df[full_df['ISRC_KEY'].isin(chunk)]
        splitted_df.append(df_chunk)
    # Unique run ID generator based on timestamp
    logging.info("ISRCs have been splitted in {} dataframes for run id={}".format(n, run_id))
    # Using multiprocessing pool to apply the function to each chunk
    with Pool(n) as p:
        # Adding indices to the list of chunks
        splitted_df_with_index = list(enumerate(splitted_df))
        # List of arguments for the function
        splitted_df_with_args = map(lambda x: (dry_run, folder, run_id, chunk_id, x[0], x[1]),
                                    splitted_df_with_index)
        parallel_results = p.map(f, splitted_df_with_args)
    logging.info("Regression calculation finished.")
    merged_df = merge_parallel_results(parallel_results)
    return merged_df


def work_load(arg):
    # Result is a shared list between all parallel tasks where computed dataframes will be appended
    (dry_run, folder, run_id, chunk_id, par_id, df) = arg
    logging.info(
        f"Arguments folder={folder}, run-id={run_id}, chunk-id={chunk_id}, par-id={par_id}, df type={type(df)}")
    # Fourier
    logging.info(f"Calculation chunk {CHUNK_ID} subset contains {df.shape[0]} rows.")
    logging.info(f"Splitting work into {parallelism} parts for parallel processing")
    # Compute FastFourierTransform and save the data to S3
    reg_df = compute_regression_for_chunk(df, chunk_id, par_id)
    logging.info(f"Linear regression data saved.")
    logging.info(f"Linear regression table contains {reg_df.shape[0]} rows.")
    return reg_df


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__':

    # Dry run = True => data will NOT be saved
    DRY_RUN = get_argument(sys.argv, "dry-run", 'False') == 'True'
    # Run id
    RUN_ID = get_mandatory_argument(sys.argv, "run-id")
    # Folder
    FOLDER = get_mandatory_argument(sys.argv, "folder")
    # Splits the universe of ids into CHUNK_COUNT stable chunks
    CHUNK_COUNT = int(get_mandatory_argument(sys.argv, "chunk-count"))
    # We are running the script only for the below CHUNK_ID, out of CHUNK_COUNT chunks
    CHUNK_ID = int(get_mandatory_argument(sys.argv, "chunk-id"))

    PARALLELISM = 70  # int(get_argument(sys.argv, "parallelism", str(max(1, cpu_count() // 2))))

    PREFIX = 'regression_table_quadratic'

    with snowflake_connector_factory(get_snowflake_creds()) as cs:
        try:
            cs.execute("USE WAREHOUSE DEV_PERFORMANCE_WAREHOUSE;")
            cs.execute(""" 
            select * from dev_engineering.eimpara.moments_collection_search_6_months
            where TRANSACTION_COUNTRY_CODE = 'GB';
            """)
            rows = cs.fetchall()
        except Exception as e:
            logging.error(f"Something went wrong - {str(e)}")

        data_df = pd.DataFrame(rows, columns=list(map(lambda meta: meta[0], cs.description)))
        df = data_df.drop_duplicates().copy()
        print(data_df.shape)
        print(df.shape)

    if DRY_RUN:
        logging.warning(f"DRY_RUN flag is True, no data will be saved !!!")
    logging.info(f"Processing Moments with run-id='{RUN_ID}'")
    logging.info(f"Computing chunk {CHUNK_ID} out of {CHUNK_COUNT} chunk(s)")
    logging.info(f"Data will be processed with parallelism={PARALLELISM}")

    # Fixing dates that are not read with the correct type in a CSV file
    df['ACTIVITY_DATE'] = pd.to_datetime(df['ACTIVITY_DATE'])
    df['ACTIVITY_DATE'] = df['ACTIVITY_DATE'].dt.date
    logging.info(f"Min activity date: {df['ACTIVITY_DATE'].min()}")
    logging.info(f"Max activity date: {df['ACTIVITY_DATE'].max()}")

    parallelism = PARALLELISM
    (elapsed_time, merged_df, _) = split_work(df, parallelism, work_load, FOLDER, RUN_ID, CHUNK_ID, DRY_RUN)
    logging.info(f"Data processed in {elapsed_time} seconds")

    s3 = boto3.resource('s3')
    bucket_name = 'dev-cucumbers'
    bucket = s3.Bucket(bucket_name)

    save_dataframe_s3(merged_df, FOLDER, PREFIX, RUN_ID)