import os
import math
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import sklearn

from fcmeans import FCM
from kneed import KneeLocator
from tslearn.clustering import TimeSeriesKMeans, silhouette_score
from sklearn_som.som import SOM
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from datetime import datetime, date
from minisom import MiniSom
from sklearn.cluster import KMeans
from password import PRIVATE_KEY_PASSPHRASE

import warnings
warnings.filterwarnings("ignore")

import snowflake.connector
from snowflake.connector.pandas_tools import write_pandas
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import dsa
from cryptography.hazmat.primitives import serialization

# Param to change if you want to run it
key_path = "/Users/impr001/Keys_snowflake/rsa_key.p8" # PLEASE CHANGE HERE
user ='eimpara'
account ='orchard'
role = 'PROD_DATALYTICS_ROLE'
warehouse = 'DEV_OWS_WAREHOUSE'
database='DEV_ENGINEERING'
schema="EIMPARA"
table_name = '!!!_DO_NOT_USE_!!!'

with open(key_path, "rb") as key:
    p_key= serialization.load_pem_private_key(
        key.read(),
        password=PRIVATE_KEY_PASSPHRASE.encode(),
        backend=default_backend()
    )

pkb = p_key.private_bytes(
    encoding=serialization.Encoding.DER,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption())

ctx = snowflake.connector.connect(
    user= user,
    account= account,
    private_key=pkb,
    role= role,
    warehouse = warehouse
    )
try:
    cs = ctx.cursor()
    sql = f"SELECT * from DEV_ENGINEERING.AADAMU_DEBUT_FORECASTING_DBT.DATASET_STREAMS_DAILY_2022;"
    cs.execute(sql)
    df = cs.fetch_pandas_all()

finally:
    cs.close()

ctx.close()


# Functions to transpose dataframe: originally this fataframe is a time series one, with each row correspondoing to a snapshot day.
# For the purpose of SOM, each day will become a column. However, column names will be recorded in a way to indicate the day since release. 
# e.g. D1 corresponds to the day of release, D2 second day since release, etc.


def compute_day_after_release(release_date: date, date_to_convert: date) -> str:
    diff = (date_to_convert - release_date).days
    return "D" + str(diff + 1)

def key_from_row(row):
    isrc = row['ISRC']
    release_format = row['RELEASE_FORMAT']
    store_id = row['STORE_ID']
    feed_id = row['FEED_ID']
    key = (isrc, release_format, store_id, feed_id)
    return key

# Scan the data and keeps (isrc, release format, store id) -> first release date
def first_release_dates_per_key(df) -> dict :
    res = {}
    for _, row in df.iterrows():
        key = key_from_row(row)
        release_date = row['RELEASE_DATE']
        cached = res.get(key)
        if cached is None:
            res[key] = release_date
        elif cached > release_date:
            res[key] = release_date
        else:
            pass
    return res

def find_first_release_date_for_key(cache: dict, key) -> date:
    return cache.get(key)

def find_first_release_date_for_row(cache: dict, row) -> date:
    return find_first_release_date_for_key(cache, key_from_row(row))

def compute_day_from_row(cache: dict, row) -> str:
    first_release_date = find_first_release_date_for_row(cache, row)
    row_release_date = row['SNAPSHOT_DATE']
    day = compute_day_after_release(first_release_date, row_release_date)
    return day

def compute_earliest_release_date(cache: dict, row) -> str:
    first_release_date = find_first_release_date_for_row(cache, row)
    return first_release_date

# Setting time variables as datetime and removing releases for which activity date (snapshot_date) is prior to release

df['RELEASE_DATE'] = pd.to_datetime(df['RELEASE_DATE'])
df['SNAPSHOT_DATE'] = pd.to_datetime(df['SNAPSHOT_DATE'])
df['DOW_activity'] = df['SNAPSHOT_DATE'].dt.day_name()
df['Month_name'] = df['SNAPSHOT_DATE'].dt.month_name()

df['RELEASE_DATE'] = df['RELEASE_DATE'].dt.date
df['SNAPSHOT_DATE'] = df['SNAPSHOT_DATE'].dt.date
df = df[df['SNAPSHOT_DATE'] >= df['RELEASE_DATE']]
df = df[(df['RELEASE_DATE'] >= date(2022, 1, 1)) & (df['RELEASE_DATE'] <= date(2022, 10, 15))]

# Manipulating the dataframe so to get the right format for analysis 

first_release_date_cache = first_release_dates_per_key(df)
df['day_after_release'] = df.apply(lambda row: compute_day_from_row(first_release_date_cache,row), axis=1)
df['Earliest_release_date'] = df.apply(lambda row: compute_earliest_release_date(first_release_date_cache,row), axis=1)

df['Earliest_release_date'] = pd.to_datetime(df['Earliest_release_date'])

df['Earliest_release_date'] = pd.to_datetime(df['Earliest_release_date'])
df = df[df['STORE_ID']==286].copy()
df = df.drop_duplicates()

grp_key = ['ISRC','day_after_release']
grp_df = df.groupby(grp_key).sum('STREAMS').reset_index()
grp_df = grp_df[grp_key + ['STREAMS']]

days_subset_0 = ["D" + str(x) for x in range(1, 258)]

subset = grp_df[grp_df['day_after_release'].isin(days_subset_0)]

streams_by_day = pd.pivot_table(subset, index =['ISRC'], 
                                     columns = 'day_after_release',
                                     values = ['STREAMS'],
                                     aggfunc = np.sum).reset_index()

streams_by_day.columns = ['_'.join(str(s).strip() for s in col if s) for col in streams_by_day.head().columns]
streams_by_day.reset_index(inplace=True)
streams_by_day  = streams_by_day.drop(columns=['index'])
streams_by_day.columns = streams_by_day.columns.str.removeprefix('STREAMS_') 

# Preparing the dataframe for clustering 
# Taking first 30 days since release and calculating relevant features

sub = streams_by_day.copy()

normal_columns = ['ISRC']
days_subset = ["D" + str(x) for x in range(1, 31)]
sub = sub[normal_columns+days_subset]
sub['tot_streams'] = sub[days_subset].sum(axis=1)
sub = sub[sub['tot_streams']>1]

# assumption: NA = 0 streams.
sub = sub.fillna(0)
sub = sub.set_index('ISRC')

sub['mean'] = sub.loc[:,days_subset].mean(axis=1)
sub['sd'] = sub.loc[:,days_subset].std(axis=1)
sub['var'] = sub.loc[:,days_subset].var(axis=1)
sub['CV'] = sub['sd']/sub['mean']

h = 15
sub['Slope'] = sub['D15'] - sub['D1']/h
sub['D30_15_diff'] = sub['D30'] - sub['D15']/h
sub['Acceleration'] = (sub['D30_15_diff'] - sub['Slope'])/h    
sub = sub[~sub['CV'].isnull()]

conditions = [(sub['Slope'] > 0) & (sub['Acceleration'] > 0),
              (sub['Slope'] > 0) & (sub['Acceleration'] <= 0),
             (sub['Slope'] < 0) & (sub['Acceleration'] > 0),
              (sub['Slope'] < 0) & (sub['Acceleration'] <= 0),
             (sub['Slope'] == 0)]

choices = ['attack - fast', 'attack - slow', 'decay - slow', 'decay - fast', 'sustain']

sub['Phase'] = np.select(conditions, choices)

# Standardising features prior to clustering

features = sub[['Slope', 'Acceleration','CV', 'mean']]
scaler = StandardScaler()
scaled_features = scaler.fit_transform(features)


######################################################################################################################################################### Self-Organising Maps (SOM) ###############################################
##################################################################################################################

# From Kohonen, T. (2001) "Self-Organising Maps":
# Form of the array: for visual inspection hexagonal lattice to be preferred because it doesn't favour horizontal abd vertical directions as much as rectangular array (p. 159).

# SOM parameters options:
# neighborhood_function : string, optional (default='gaussian')
# Function that weights the neighborhood of a position in the map.
# Possible values: 'gaussian', 'mexican_hat', 'bubble', 'triangle'

# topology : string, optional (default='rectangular')
# Topology of the map.
# Possible values: 'rectangular', 'hexagonal'

# activation_distance : string, callable optional (default='euclidean')
# Distance used to activate the map.
# Possible values: 'euclidean', 'cosine', 'manhattan', 'chebyshev'

som_x = 6
som_y = 6
features_n = 4
param_sigma = 1.5 # 1.5
param_learning_rate = 0.001
param_iterations = 50000 # change with over values 
param_activation_distance = 'euclidean'
param_topology = 'rectangular'

som = MiniSom(som_x, som_y, features_n, 
              # decay_function=asymptotic_decay,
              sigma=param_sigma, 
              learning_rate = param_learning_rate, 
              neighborhood_function='gaussian', 
              activation_distance = param_activation_distance,
              topology = param_topology,
              random_seed=2022)

som.pca_weights_init(scaled_features)
som.train_batch(scaled_features, param_iterations)

# SOM evaluation

# From https://www.intechopen.com/chapters/69305
# Quantization error and topographical error are main measurements to assess the quality of SOM. 
# Quantization error is the average difference of the input samples compared to its corresponding winning neurons (BMU). 
# It assesses the accuracy of the represented data, therefore, it is better when the value is smaller.

# Topographical error assesses the topology preservation. 
# It indicates the number of the data samples having the first best matching unit (BMU1) and the second best matching unit (BMU2) being not adjacent. 
# Therefore, the smaller value is better.

print(round(som.quantization_error(scaled_features), 3))
print(round(som.topographic_error(scaled_features), 3))

# 

def cluster_name(som_x, som_y, cluster):
    
    """ given SOM size and cluster co-ordinates, it names the cluster"""
    
    (x,y) = cluster
    cluster_number = x*som_y+y+1
    return f"Cluster_SOM {cluster_number}"

# Adjusted from:
# https://www.kaggle.com/code/izzettunc/introduction-to-time-series-clustering/notebook

def plot_som_series_averaged_center(som_x, som_y, win_map):
    
    """ It plots instances and centroid in each cluster """
    
    fig, axs = plt.subplots(som_x,som_y,figsize=(25,25))
    fig.suptitle('Clusters')
    for x in range(som_x):
        for y in range(som_y):
            cluster = (x,y)
            if cluster in win_map.keys():
                for series in win_map[cluster]:
                    axs[cluster].plot(series,c="gray",alpha=0.5) 
                axs[cluster].plot(np.average(np.vstack(win_map[cluster]),axis=0),c="red")
            name = cluster_name(som_x, som_y, cluster)
            axs[cluster].set_title(f"{name} - {cluster} - n={len(win_map[cluster])}")

    plt.show()
    
# Returns the mapping of the winner nodes and inputs.
# NB. The red line represents the centroid. The grey lines are the observations in each cluster/neuron. However these are not time series data.
# The lines in each cluster/neuron represent the standardised values for 'Slope', 'Acceleration','CV', 'mean' (the variables used for SOM) scored by each observation.

win_map = som.win_map(scaled_features)

#plot_som_series_averaged_center(som_x, som_y, win_map)

# Computes n. of instances in each cluster

cluster_map = []
for idx in range(len(scaled_features)):
    winner_node = som.winner(scaled_features[idx])
    name = cluster_name(som_x, som_y, winner_node)
    cluster_map.append(name)
    
# Creating cluster label for SOM

sub['Cluster_SOM'] = cluster_map

######################################################################################################################################################### k-means clustering #######################################################
##################################################################################################################

# Elbow method

cs = []
MAX_K_ELBOW_METHOD = 35
for i in range(1, MAX_K_ELBOW_METHOD):
    kmeans = KMeans(n_clusters = i, init = 'k-means++', algorithm = 'lloyd', max_iter = 300, n_init = 1, random_state=2023)
    kmeans.fit(scaled_features)
    cs.append(kmeans.inertia_)
plt.plot(range(1, MAX_K_ELBOW_METHOD), cs)
#plt.show()

knee = KneeLocator(range(1,MAX_K_ELBOW_METHOD), cs, curve='convex', direction='decreasing')
k = knee.elbow

#Choosing inertia value

init_list = ["k-means++", "random"]
inertia_list = []
for init in init_list:
    kmeans_model_par = KMeans(n_clusters=k, init=init)
    kmeans_model_par.fit(scaled_features)
    inertia_list.append(kmeans_model_par.inertia_)

results = pd.DataFrame(data=inertia_list, columns=["Inertia Value"], index=init_list)
#results

kmeans_model = KMeans(n_clusters=k, init = 'k-means++', algorithm = 'lloyd', max_iter = 300, n_init = 1, random_state=2023)
kmeans_model.fit(scaled_features)
clusters = kmeans_model.predict(scaled_features)

# Evaluating k-means

sklearn.metrics.silhouette_score(scaled_features, clusters).round(2)

# Creating cluster label for k-means

sub['Kmeans_Cluster'] = kmeans_model.labels_

######################################################################################################################################################### Fuzzy k-means clustering #################################################
##################################################################################################################

fcm = FCM(n_clusters=k,  random_state=2023)
fcm.fit(scaled_features)
fcm_centers = fcm.centers
fcm_labels = fcm.predict(scaled_features)

# Evaluating fuzzy k-means

fcm.partition_entropy_coefficient
sklearn.metrics.silhouette_score(scaled_features, fcm_labels).round(2)

# Creating cluster label for fuzzy k-means

sub['Cmeans_Cluster'] = fcm_labels

######################################################################################################################################################### attack/decay/sustain #####################################################
##################################################################################################################

# Saving table

table = sub.reset_index()

cols = ['ISRC',
'tot_streams',
'mean',           
'sd',              
'var',             
'CV',              
'Slope',          
'Acceleration',
'Cluster_SOM',
'Kmeans_Cluster',
'Cmeans_Cluster',
'Phase']
days = ["D" + str(x) for x in range(1, 31)]
table = table[cols+days].copy()

def create_fourier_table(cs):
    cs.execute("USE WAREHOUSE DEV_PERFORMANCE_WAREHOUSE")
    cs.execute("USE DATABASE DEV_ENGINEERING")
    cs.execute("USE SCHEMA EIMPARA")
    cs.execute(
        "CREATE TABLE " +
        table_name + "(ISRC string, tot_streams float, mean float, sd float, var float, CV float, Slope float, Acceleration float, Cluster_SOM string, Kmeans_Cluster integer, Cmeans_Cluster integer, Phase string, D1 float, D2 float, D3 float, D4 float, D5 float, D6 float, D7 float, D8 float, D9 float, D10 float, " + 
        "D11 float,D12 float, D13 float, D14 float, D15 float, D16 float, D17 float, D18 float, D19 float, D20 float, " +
          "D21 float, D22 float, D23 float, D24 float, D25 float, D26 float, D27 float, D28 float, D29 float, D30 float)")

def save_to_snowflake(ctx, df):
    return write_pandas(
        conn=ctx, 
        df=df, 
        table_name=table_name,
        database="DEV_ENGINEERING",
        schema="EIMPARA",
        quote_identifiers=False)

ctx = snowflake.connector.connect(
    user='eimpara',
    account='orchard',
    private_key=pkb,
    role= 'PROD_DATALYTICS_ROLE',
    warehouse = "DEV_PERFORMANCE_WAREHOUSE")
print("And here...")

try:
    cs = ctx.cursor()
    print("Here")
    create_fourier_table(cs)
    print("Saving data ...")
    success, num_chunks, num_rows, output = save_to_snowflake(ctx, table)
    print("Data saved: success={},rows={},output={}".format(success, num_rows, output))
finally:
    cs.close()