# Import packages
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from datetime import datetime
import pandas as pd
import pathlib
from datetime import timedelta
import random


# load min_val and max_val to normalize data & load latent dimension
LOC_VARS = "/important_variables/"
NAME_MAX_VAL = "max_val.npy"
NAME_MIN_VAL = "min_val.npy"
ENCODER_LOCATION = "encoder"

min_val = np.load(LOC_VARS + NAME_MIN_VAL)
max_val = np.load(LOC_VARS + NAME_MAX_VAL)
LATENT_DIM = int(open(LOC_VARS + "/latent_dim.txt", "r").read())


## Normalize data
def normalize(data):
    # Normalize data
    return (data - min_val) / (max_val - min_val)


## Convert a row to a np.datetime64 object
def check_validity_date_elements(date):
    # make sure that the year contains 4 digits and add 0's if necessary
    date[0] = "%04d" %date[0]
    # make sure that the months, days, hours, minutes and seconds contain 2 digits and add 0's if necessary
    date[1:] = ["%02d" %x for x in date[1:]]
    return date

def create_np_datetime_from_date(date):
    # create a numpy datetime64 object from a row in the generated data
    date = date[:6].tolist()
    date = check_validity_date_elements(date)
    result = "-".join(date[:3])
    result += "T"
    result += ":".join(date[3:])
    return np.datetime64(result)


## Vectorize an incoming timestep
def days_per_month(month_number, is_leap_year):
    # February
    if month_number == 2:
        if is_leap_year:
            return 29
        else:
            return 28
    # based on the Knuckle mnemonic
    elif month_number <= 7:
        if month_number % 2 != 0:
            return 31
        else:
            return 30
    else:
        if month_number % 2 == 0:
            return 31
        else:
            return 30

def extract_features_date(date):
    # extract the features of the date
    years = date.astype(object).year
    months = date.astype(object).month
    days = date.astype(object).day
    hours = date.astype(object).hour
    minutes= date.astype(object).minute
    seconds = date.astype(object).second
    day_of_week = date.astype(datetime).isoweekday()
    millennium = round(np.floor(years / 1000))
    century = round(np.floor(years / 100))
    decade = round(np.floor(years / 10))
    last_digit_year = years % 10
    is_leap_year = years % 4 == 0
    # 0 because Sunday is 0
    is_business_day = day_of_week <= 5 
    quarter = round(np.ceil(months / 3))
    days_in_month = days_per_month(months, is_leap_year)
    day_of_year = date.astype(object).timetuple().tm_yday
    is_month_start = days == 1
    is_quarter_start = days == 1 and months % 3 == 0
    is_year_start = days == 1 and months == 1
    is_year_end = days == 31 and months == 12
    # end of the year is the end of the last quarter, 31 March, 30 June, 30 September
    is_quarter_end = is_year_end or (months == 3 and days == 31) or ((months == 6 or 9) and days == 30)
    is_month_end = days == days_in_month
    
    encoding = [years, months, days, hours, minutes, seconds, day_of_week, millennium, century, decade, last_digit_year, is_leap_year, is_business_day, quarter, days_in_month, day_of_year, is_month_start,
            is_quarter_start, is_year_start, is_year_end, is_quarter_end, is_month_end]
    
    return encoding


def date_to_features(date):
    # make a np.datetime64 of the date (in the shape YYYY-MM-DD HH:MM:SS)
    date = np.datetime64(date)
    date = date.astype('datetime64[ms]')
    # obtain features
    features = extract_features_date(date)
    # return all features, except the datetime object
    return features


def get_latent_vars(dates, model):
    # make strings of the dates
    dates = [str(date) for date in dates]
    # obtain the latent vectors
    dates_latent_vec = model(dates)
    return dates_latent_vec


def sample_date(start, end):
    # create datetime objects for start and end to sample
    start = datetime.strptime(start, '%Y-%m-%d %H:%M:%S')
    end = datetime.strptime(end, '%Y-%m-%d %H:%M:%S')
    # calculate the difference between the 2 dates
    delta = end - start
    int_delta = (delta.days * 24 * 60 * 60) + delta.seconds
    # randomly select the number of seconds to add to the start date
    random_second = random.randrange(int_delta)
    # add the random to the start date and return
    return start + timedelta(seconds=random_second)


def validate_sample(sample, interest, start):
    # check whether the sampled date is a valid date (only if the interest is fixed)
    start_date = np.datetime64(start)
    try:
        if interest == "Y":
            sample = sample.replace(year=start_date.astype(object).year)
        elif interest == "M":
            sample = sample.replace(month=start_date.astype(object).month)
        elif interest == "D":
            sample = sample.replace(day=start_date.astype(object).day)
        elif interest == "h":
            sample = sample.replace(hour=start_date.astype(object).hour)
        elif interest == "m":
            sample = sample.replace(minute=start_date.astype(object).minute)
        elif interest == "s":
            sample = sample.replace(second=start_date.astype(object).second)
        else:
            print("Invalid interest type")
            return None
        return True
    except ValueError:
        return False


def sample_and_vectorize_timestamps(start, end, interest, n_samples, single_interst=False):
    """"
    Obtain samples from start to end
    The interest is used for plotting (and possibly the sampling process)
    When single_interst = True, the samples will get the same value as start for a specific interest
    start = start date
    end = end date
    interest = Y:Year, M:Month, D:Day, h:hour, m:minute, s:seconds
    n_samples = the number of samples
    """
    # sample dates between start and end, while keep the interest time component the same
    result = []
    for _ in range(n_samples):
        start_date = np.datetime64(start)
        sample = sample_date(start, end)
        if single_interst:
            while not validate_sample(sample, interest, start):
                sample = sample_date(start, end)
            if interest == "Y":
                sample = sample.replace(year=start_date.astype(object).year)
            elif interest == "M":
                sample = sample.replace(month=start_date.astype(object).month)
            elif interest == "D":
                sample = sample.replace(day=start_date.astype(object).day)
            elif interest == "h":
                sample = sample.replace(hour=start_date.astype(object).hour)
            elif interest == "m":
                sample = sample.replace(minute=start_date.astype(object).minute)
            elif interest == "s":
                sample = sample.replace(second=start_date.astype(object).second)
            else:
                print("Invalid interest type")
                return None
        sample = np.datetime64(sample)
        result.append(sample)
    result = np.array(result)
    return result


def preprocess_and_obtain_latent_variables(start, end, interest, interval, model):
    """"
    Obtain the latent variables of the dates from start till end, using a certain interval (interest)
    start = start date
    end = end date
    interest = Y:Year, M:Month, D:Day, h:hour, m:minute, s:seconds
    """
    # get the dates from <start> to <end> with interval <interest> 
    dates = obtain_date_range(start, end, interest, interval)
    # obtain the labels for plotting
    labels = obtain_labels_date_range(dates, interest)
    # change type to strings
    dates = [str(date) for date in dates]
    # get latent vectors
    dates_latent_vec = model(dates)

    return dates_latent_vec, labels


### Visualization of the latent space
def label_point(x, y, val, ax, diff):
    x, y, val = pd.Series(x), pd.Series(y), pd.Series(val)
    a = pd.concat({'x': x, 'y': y, 'val': val}, axis=1)
    for _, point in a.iterrows():
        ax.text(point['x']+diff, point['y'], str(point['val']))


def plot_latent_with_label(latent_vector, latent_dim, labels, interest):
    """"
    Plot the latent variable with dimension <latent_dim> with corresponding labels 
    """
    ax = sns.scatterplot(labels, latent_vector[:, latent_dim])
    plt.title('Example Plot')
    # Set x-axis label
    plt.xlabel(interest)
    # Set y-axis label
    plt.ylabel('Latent variable')
    label_point(labels, latent_vector[:, latent_dim], labels, plt.gca(), .0002)


def plot_latent_with_label_2dim(latent_vec, latent_dim1, latent_dim2, labels, interest):
    """"
    Plot the latent variable with dimension <latent_dim> with corresponding labels 
    """
    ax = sns.scatterplot(latent_vec[:, latent_dim1], latent_vec[:, latent_dim2])
    plt.title('Example Plot')
    # Set x-axis label
    plt.xlabel(interest)
    # Set y-axis label
    plt.ylabel('Latent variable')
    label_point(latent_vec[:, latent_dim1], latent_vec[:, latent_dim2], labels, plt.gca(), .0002)


def create_plot_range_2dim(start, end, interest, interval, latent_dim1, latent_dim2, model):
    """"
    Create plots from start till end, using a certain interval (interval)
    start = start date
    end = end date
    interest = Y:Year, M:Month, D:Day, h:hour, m:minute, s:seconds
    """

    # get latent vector and labels
    dates_latent_vec, labels = preprocess_and_obtain_latent_variables(start, end, interest, interval, model)
    # call plot_latent_with_label 
    plot_latent_with_label_2dim(dates_latent_vec, latent_dim1, latent_dim2, labels, interest)


def create_plot_range_all_latent_dim(start, end, interest, interval, model):
    """"
    Create plots from start till end, using a certain interval (interval), for all latent dimensions
    start = start date
    end = end date
    interest = Y:Year, M:Month, D:Day, h:hour, m:minute, s:seconds
    """
    # get latent vector and labels
    dates_latent_vec, labels = preprocess_and_obtain_latent_variables(start, end, interest, interval, model)
    # call plot_all_latent_dim
    plot_all_latent_dim(start, end, dates_latent_vec, labels, interest, interval)


def flatten(t):
    return [item for sublist in t for item in sublist]


def plot_all_latent_dim(start, end, dates_latent_vec, labels, interest, interval):
    fig, axs = plt.subplots(2, 4, figsize=(LATENT_DIM * 4, LATENT_DIM * 2 + 4))
    dim = 1
    for i in range(2):
        for j in range(4):
            axs[i, j].scatter(labels, dates_latent_vec[:, (i*2 + j*1)])
            axs[i, j].set_title("Latent dimension: " + str(dim))
            axs[i, j].set(xlabel=interest, ylabel='Latent dimension')
            label_point(labels, dates_latent_vec[:, (i*2 + j*1)], labels, axs[i, j], .001)
            dim += 1

    fig.suptitle('All single latent dimensions from ' + str(start) + " till " + str(end) + " with an interval of " + str(interval) + interest, y=0.93, fontsize="xx-large", fontweight="bold")

def obtain_date_range(start, end, interest, interval):
    """"
    Obtain the needed date range, with corresponding interval for the create_plot_range function
    interest = Y:Year, M:Month, D:Day, h:hour, m:minute, s:seconds
    """
    if interest == "Y":
        dates = np.arange(start, end, interval, dtype='datetime64[' + interest + ']')
        return dates
    if interest == "M":
        dates = np.arange(start, end, interval, dtype='datetime64[' + interest + ']')
        return dates
    if interest == "D":
        dates = np.arange(start, end, interval, dtype='datetime64[' + interest + ']')
        return dates
    if interest == "h":
        dates = np.arange(start, end, interval, dtype='datetime64[' + interest + ']')
        return dates
    if interest == "m":
        dates = np.arange(start, end, interval, dtype='datetime64[' + interest + ']')
        return dates
    if interest == "s":
        dates = np.arange(start, end, interval, dtype='datetime64[' + interest + ']')
        return dates
    else:
        print("Invalid interest type")
        return None


def obtain_labels_date_range(dates, interest):
    """"
    Obtain the needed labels for the create_plot_range function
        Year = 0
        Month = 1
        Day = 2
        Hour = 3
        Minute = 4
        Second = 5
    """
    if interest == "Y":
        return [date.astype(object).year for date in dates]
    if interest == "M":
        return [date.astype(object).month for date in dates]
    if interest == "D":
        return [date.astype(object).day for date in dates]
    if interest == "h":
        return [date.astype(object).hour for date in dates]
    if interest == "m":
        return [date.astype(object).minute for date in dates]
    if interest == "s":
        return [date.astype(object).second for date in dates]
    else:
        print("Invalid interest type")
        return None

def create_plot_range(start, end, interest, interval, latent_dim, model):
    """"
    Create plot from start till end, using a certain interval (interest), of latent dim <latent_dim>
    start = start date
    end = end date
    interest = Y:Year, M:Month, D:Day, h:hour, m:minute, s:seconds
    """

    # get latent vector and labels
    dates_latent_vec, labels = preprocess_and_obtain_latent_variables(start, end, interest, interval, model)
    # call plot_latent_with_label 
    plot_latent_with_label(dates_latent_vec, latent_dim, labels, interest)


### Visualization of time dataset
def visualize_distribution(data):
    # Visualize the years
    sns.displot(data[:,0])
    plt.xlabel("Years")
    plt.show()
    # Visualize the months
    sns.displot(data[:,1])
    plt.xlabel("Months")
    plt.show()
    # Visualize the days
    sns.displot(data[:,2])
    plt.xlabel("Days")
    plt.show()
    # Visualize the hours
    sns.displot(data[:,3])
    plt.xlabel("Hours")
    plt.show()
    # Visualize the minutes
    sns.displot(data[:,4])
    plt.xlabel("Minutes")
    plt.show()
    # Visualize the seconds
    sns.displot(data[:,5])
    plt.xlabel("Seconds")
    plt.show()
    # Visualize the day of week
    sns.displot(data[:,6])
    plt.xlabel("Day of week")
    plt.show()
    # Visualize the day of week
    sns.displot(data[:,7])
    plt.xlabel("Millenium")
    plt.show()
    # Visualize the century
    sns.displot(data[:,8])
    plt.xlabel("Century")
    plt.show()
    # Visualize the decade
    sns.displot(data[:,9])
    plt.xlabel("Decade")
    plt.show()


### Visualization of the error between reconstructed and input data
def plot_error_timestep(input, reconstruction):
    plt.scatter(range(len(input)), input, color='blue')
    plt.scatter(range(len(reconstruction)), reconstruction, color='red')
    plt.fill_between(np.arange(20), input, reconstruction, color='lightcoral')
    plt.legend(labels=["Input", "Reconstruction", "Error"])
    plt.show()

def plot_error_timesteps(input, reconstructed, n):
    for i in range(n):
        plot_error_timestep(input[i], reconstructed[i])
    
