import gc
import os
from collections import defaultdict

import numpy as np
import pandas as pd
import tensorflow as tf
from keras import backend as K
from sklearn.metrics import confusion_matrix, jaccard_score
from tensorflow import keras

from ivus_segmentation.metrics import dice_for_class, hausdorff_distance


def get_specificity(col, cm):
    """Return the specificity of a variable in a 3x3 confusion matrix"""
    col_bef = (col - 1) % 3
    col_aft = (col + 1) % 3

    tp = cm[col, col]
    tn = cm[col_bef, col_bef] + cm[col_aft, col_aft]
    fp = cm[col_bef, col] + cm[col_aft, col]
    fn = cm[col, col_bef] + cm[col, col_aft]

    return float(tn) / float(tn + fp)


def get_accuracy(col, cm):
    """Return the accuracy of a variable in a 3x3 confusion matrix"""
    col_bef = (col - 1) % 3
    col_aft = (col + 1) % 3

    tp = cm[col, col]
    tn = cm[col_bef, col_bef] + cm[col_aft, col_aft]
    fp = cm[col_bef, col] + cm[col_aft, col]
    fn = cm[col, col_bef] + cm[col, col_aft]

    return float(tp + tn) / float(tp + tn + fp + fn)


def custom_classification_report(y_true, y_pred,
                                 classes=["media", "lumen"],
                                 verbose=False, pixel_size=0.0352 / 2):
    n, _, _ = y_true.shape
    m, _, _ = y_pred.shape

    assert n == m, "Mismatched number of predictions"

    # Obtain precision, recall, f1-score and support
    report = defaultdict(list)
    for i in range(n):
        if verbose:
            print(f"{i+1}/{n}     ", end="\r")
        gt, pred = y_true[i], y_pred[i]
        js = jaccard_score(
            gt.flatten(),
            pred.flatten(),
            average=None,
        )

        cm = confusion_matrix(gt.flatten(), pred.flatten())
        for j, cls in enumerate(classes):
            j = j + 1
            hd = hausdorff_distance(gt == j, pred == j) * pixel_size
            report["hd_" + cls].append(hd)
            report["dice_" + cls].append(dice_for_class(j))
            report["jacc_" + cls].append(js[j])

            report["specificity_" + cls].append(get_specificity(j, cm))
            report["accuracy_" + cls].append(get_accuracy(j, cm))
    return report


def get_callbacks(filepath, patience_lr):
    mcp_save = keras.callbacks.ModelCheckpoint(
        filepath,
        save_weights_only=True,
        save_best_only=True,
        monitor="val_loss",
        mode="min",
    )
    reduce_lr_loss = keras.callbacks.ReduceLROnPlateau(
        monitor="loss",
        factor=0.1,
        patience=patience_lr,
        verbose=1,
        min_delta=1e-4,
        mode="min",
    )
    return [mcp_save, reduce_lr_loss]


OPTIMIZER_BASE = {
    "learning_rate": 0.001,
    "beta_1": 0.9,
    "beta_2": 0.999,
    "epsilon": 1e-08,
    "decay": 0.0,
}


def k_fold(
    model_factory,
    folds,
    train_data,
    test_data,
    verbose=False,
    batch_size=8,
    epochs=150,
    weight_output=None,
    history_output=None,
):
    data_by_fold = defaultdict(list)
    # model_by_fold = []
    # metrics = ["categorical_accuracy"]

    images, masks = train_data
    images_test, masks_test = test_data
    optimizer = keras.optimizers.Adam(**OPTIMIZER_BASE)
    count = 0
    for j, (train_idx, val_idx) in folds:
        K.clear_session()
        # Free up RAM in case the model definition cells were run multiple times

        # model, optimizer, loss = model_factory(images.shape[1:], 3)
        model = model_factory(images.shape[1:])

        x_train = images[train_idx]
        y_train = masks[train_idx]
        x_valid = images[val_idx]
        y_valid = masks[val_idx]

        model.compile(optimizer=optimizer)

        out = model.fit(
            x_train,
            y_train,
            validation_data=(x_valid, y_valid),
            batch_size=batch_size,
            shuffle=True,
            epochs=epochs,
            verbose=verbose,
            callbacks=get_callbacks("weights", 10),
        )
        history = out.history

        if history_output is not None:
            df = pd.DataFrame(history)
            df.to_csv(f"{history_output}-{count}.csv")

        score = model.evaluate(images_test, masks_test, verbose=verbose)

        for k, v in score.items():
            data_by_fold[k].extend(v)

        if weight_output is not None:
            print("attempting to save", f"{weight_output}-{count}")
            model.model.save(f"{weight_output}-{count}.h5")
        count += 1

        del model
        gc.collect()

        # model_by_fold.append(model)

    return data_by_fold
