import random

import numpy as np
import pandas as pd

from keras import Input
from keras.src.losses import SparseCategoricalCrossentropy
from keras.src.metrics import SparseCategoricalAccuracy
from pandas import DataFrame
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional,GRU
from tensorflow.keras.optimizers import Adam
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix

epochs = 5#50
model_type_gru = 'GRU'
model_type_lstm = 'LSTM'
model_type_bilstm = 'BiLSTM'

def make_sequences(data, sequence_length):
    """
    Converts the data into sequences of the given length

    :param data: Original data
    :param sequence_length: length of intended sequences
    :return: x,y for the sequences
    """
    x, y = [], []
    features = data.drop('user', axis=1).values
    labels = data['user'].values
    for i in range(len(features) - sequence_length+1): # with overlap on days
#    for i in range(0, len(features) - sequence_length + 1, sequence_length): # without overlap on days
        x.append(features[i:i + sequence_length])
        y.append(labels[i + sequence_length-1])
    return x, y

def prepare_data_for_basic_algorithm(user_data, sequence_length):
    """
    Converts the data into a format the sklearn algorithms can work with. Does not change the data, only the structure
    :param user_data: the dict of dataframe with the data
    :param sequence_length: intended sequence length
    :return: the formatted data
    """
    combined = pd.DataFrame()
    for user, data in user_data.items():
        x_new, y_new = make_sequences(data, sequence_length)
        if len(x_new)>0:
            var = [[pd.DataFrame(a[s]) for s in range(sequence_length)] for a in x_new]
            df_var = pd.concat([pd.concat(seq_list).T for seq_list in var])
            df_var['user'] = user
            combined = pd.concat([combined, df_var], ignore_index=True)
    return combined.drop(columns=['user']), combined['user']

def prepare_data_for_neural_model(user_data, sequence_length, print_counts=False):
    """
    Converts the data into a format the neural model can work with. Does not change the data, only the structure

    :param print_counts: Whether to print some additional debug data
    :param user_data: the dict of dataframe with the data
    :param sequence_length: intended sequence length
    :return: the formatted data
    """
    x, y = [], []
    combined = pd.DataFrame()
    for user, data in user_data.items():
        x_new, y_new = make_sequences(data, sequence_length)
        x = x + x_new
        y = y + y_new
        if print_counts and len(x_new)>0:
            var = [[pd.DataFrame(a[s])for s in range(sequence_length)] for a in x_new ]
            df_var = pd.concat([pd.concat(seq_list).T for seq_list in var])
            df_var['user'] = user
            combined = pd.concat([combined, df_var], ignore_index=True)
    if print_counts:
        combined_ohne = combined.drop('user', axis=1)
        print('Alle', len(combined))
        print('Unique mit user', len(combined.drop_duplicates()))
        print('Unique ohne user', len(combined_ohne.drop_duplicates()))
        print('Unique')
        print(combined.drop_duplicates()['user'].value_counts())
        print('Alle')
        print(combined['user'].value_counts())
    random.Random(17).shuffle(x)
    random.Random(17).shuffle(y)
    x = np.array(x)
    y = np.array(y)
    return x,y

def train_one_model(train_data, val_data, n_batch, n_epochs, n_neurons, sequence_length, model_type):
    x, y = train_data
    x_v, y_v = val_data
    users = list(set(y))

    # renumber users
    user_map = {users[i]:i for i in range(len(users))}
    y = np.array([user_map[x] for x in y])
    y_v = np.array([user_map[x] for x in y_v])
    n_features = x.shape[2]
    user_num = len(users)

    # prepare model
    def build_model():
        model = Sequential()
        model.add(Input(shape=(sequence_length, n_features), batch_size=n_batch))
        if model_type == model_type_bilstm:
            model.add(Bidirectional(LSTM(n_neurons)))
        if model_type == model_type_lstm:
            model.add(LSTM(n_neurons))
        if model_type == model_type_gru:
            model.add(GRU(n_neurons))
        model.add(Dense(user_num, activation='softmax'))
        model.compile(
            optimizer=Adam(),
            loss=SparseCategoricalCrossentropy(),
            metrics=[SparseCategoricalAccuracy()],
        )
        return model

    model = build_model()

    # fit model
    train_acc, test_acc, train_p, test_p, train_r, test_r, train_f1, test_f1 = list(), list(),list(), list(),list(), list(),list(), list()
    for i in range(n_epochs):
        model.fit(x, y, batch_size=n_batch, epochs=1, verbose=0, shuffle=False)
        # evaluate model on train data
        acc, p, r, f1 = evaluate(model, (x, y), sequence_length, n_batch)
        train_acc.append(acc)
        train_p.append(p)
        train_r.append(r)
        train_f1.append(f1)
        # evaluate model on test data
        savename = 'cf_matrix_'+get_save_id(n_epochs, n_neurons,  n_batch)+'.json'
        acc, p, r, f1 = evaluate(model, (x_v, y_v), n_batch, save_name=savename)
        test_acc.append(acc)
        test_p.append(p)
        test_r.append(r)
        test_f1.append(f1)

    history = DataFrame()
    history['train_acc'], history['test_acc'] = train_acc, test_acc
    history['train_p'], history['test_p'] = train_p, test_p
    history['train_r'], history['test_r'] = train_r, test_r
    history['train_f1'], history['test_f1'] = train_f1, test_f1
    return history

def get_save_id(n_epochs, n_neurons, n_batch):
    return '_e'+str(n_epochs)+'_n'+str(n_neurons)+'_b'+ str(n_batch)

def evaluate(model, data, batch_size, save_name=None):
    """
    GIven a model, the data is used for prediction and then evaluated.

    :param model: Model to use with a .predict() call
    :param data: x, y_true of the data, already prepared for the model
    :param batch_size: batch size for prediction
    :param save_name: if provided, results will be saved to a json file of that name
    :return: the evaluation results
    """
    x, y_true = data
    y_pred = model.predict(x, verbose=0, batch_size=batch_size)
    y_pred_classes = np.argmax(y_pred, axis=1)
    cf_matrix = pd.DataFrame(confusion_matrix(y_true, y_pred_classes))
    if save_name is not None:
        cf_matrix.to_json('results/'+save_name)
        true_counts = pd.DataFrame(y_true).value_counts()
        print('Top true occurrences', true_counts[:6])
        predicted_counts = pd.DataFrame(y_pred_classes).value_counts()
        print('Top predicted occurrences', predicted_counts[:6])

    return eval_metrics(y_true=y_true, y_pred=y_pred_classes)


def eval_metrics(y_true, y_pred):
    """
    Calculate the evaluation metrics
    :param y_true:
    :param y_pred:
    :return: acc, p, r, f1
    """
    f1 = f1_score(y_true=y_true, y_pred=y_pred, average='weighted')
    p = precision_score(y_true=y_true, y_pred=y_pred, average='weighted')
    r = recall_score(y_true=y_true, y_pred=y_pred, average='weighted')
    acc = accuracy_score(y_true=y_true, y_pred=y_pred)
    return acc, p, r, f1