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 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 prepare_data(df, var = 'STREAMS'):
    df['MA7'] =  df[var].rolling(window=7).mean()
    sub = df.iloc[7:].copy()
    sub = sub.sort_values(by='ACTIVITY_DATE', ascending =True).reset_index()
    sub['time'] = (sub.index)+1
    sub['mapping_curve'] = sub['MA7'] /sub.iloc[0]['MA7']
    return sub

def fit_curve(sub, A = 0.1148, B = 0.0334, C = 1.6014):
    sub['f(x)'] = A*np.exp(B*sub['time']) + C
    sub['error'] = (sub['f(x)'] - sub['mapping_curve'])**2
    return sub['error'].sum()

def compile_errors_table(df):
    res = pd.DataFrame(columns=['ISRC_KEY', 'country', 'error_streams', 'error_collection','error_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'])
        country = df_isrc['TRANSACTION_COUNTRY_CODE'].iloc[0]
        try:
            streams_df = prepare_data(df_isrc, var = 'STREAMS')
            error_streams = fit_curve(streams_df, A = 0.1148, B = 0.0334, C = 1.6014)
            collection_df = prepare_data(df_isrc, var = 'COLLECTION_STREAMS')
            error_collection = fit_curve(collection_df,  A = 0.0965, B = 0.0347, C = 1.3734)
            search_df = prepare_data(df_isrc, var = 'SEARCH_STREAMS')
            error_search= fit_curve(search_df,  A = 0.3223, B = 0.0291, C = 0.9835)
            res = pd.concat([res, pd.DataFrame([{
                'ISRC_KEY': isrckey, 
                'country': country,
                'error_streams': error_streams,
                'error_collection': error_collection,
                'error_search': error_search
                }], columns=res.columns)])
            
        except Exception as e:
            print(f'Error: {e}')
    return res



def compute_for_chunk(chunk_df, chunk_id, par_id):

    logging.info(f"Starting to compute linear regression on data chunk {chunk_id} and parallel job {par_id,} ...")
    error_table = compile_errors_table(chunk_df)
    logging.info(f"Fitting exponential line on data chunk {chunk_id} and parallel job {par_id,}  finished.")
    logging.info(f"Saving linear regression data to S3 ...")
    return error_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_KEY) 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 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 OLS and save the data to S3
    reg_df = compute_for_chunk(df, chunk_id, par_id)
    logging.info(f"Manual calculations data saved.")
    logging.info(f"Error 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 = 'PATTERN_DETECTION_TOTAL_STREAMS'

    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_6_months_test;
            """)
            logging.info("Fetching rows from Snowflake")
            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)