import argparse
import gc
import logging
import os
import sys
import typing
from dataclasses import dataclass

import keras
import keras.backend as K
import numpy as np
import pandas as pd
import tensorflow as tf
import yaml

import ivus_segmentation.architectures as architectures
import ivus_segmentation.augment as augment
import ivus_segmentation.metrics as metrics
import ivus_segmentation.preprocessing as preprocessing
from ivus_segmentation.dataset import get_data
from ivus_segmentation.model_selection import TRAIN_SIZE, folds
from ivus_segmentation.models import IvusModel

AUGMENTS = {
    "RandomZoom": augment.RandomZoom,
    "RandomRotation": augment.RandomRotation,
    "RandomFlip": augment.RandomFlip,
    "ElasticDeformation": augment.ElasticDeformation,
    "EDMask": augment.ElasticDeformationOnlyMasks,
}

ARCHITECTURES = {
    "unet": architectures.unet,
    "inception": architectures.inception_resnetv2_unet,
    "inception#frozen": architectures.inception_resnetv2_unet_frozen,
}

WEIGHTS = np.array([0.14891244, 0.93363497, 0.91745259])
LOSSES = {
    "categorical_crossentropy": metrics.categorical_crossentropy,
    "jaccard": metrics.jaccard_coef_loss,
    "hausdorff_dt": metrics.HausdorffDTLoss(alpha=2.0, crit="sum", weights=WEIGHTS),
    "combined": metrics.CombinedLoss(weights=WEIGHTS),
}

PREPROCESSORS = {
    "identity": preprocessing.IdentityPreprocessor,
}


@dataclass
class Experiment:
    name: str
    input_shape: tuple
    arch: object
    augments: list
    loss: object
    preprocessor: preprocessing.Preprocessor

    @classmethod
    def from_dict(cls, experiment: dict) -> "Experiment":
        name = experiment["name"]
        input_shape = experiment["input_shape"]
        arch = ARCHITECTURES[experiment["arch"]]
        augments = (
            [
                build(AUGMENTS[name], params)
                for name, params in experiment["augments"].items()
            ]
            if "augments" in experiment and experiment["augments"]
            else []
        )
        loss = LOSSES[experiment["loss"]]

        if "preprocessor" not in experiment:
            preprocessing = None
        elif experiment["preprocessor"] is None:
            preprocessing = None
        else:
            f, *args = experiment["preprocessor"]
            f = PREPROCESSORS[f]
            preprocessing = f(*args)

        return Experiment(
            arch=arch,
            augments=augments,
            loss=loss,
            input_shape=input_shape,
            name=name,
            preprocessor=preprocessing,
        )

    def build(self, seed=None) -> IvusModel:
        K.clear_session()
        if seed is not None:
            tf.random.set_seed(0)
        model = IvusModel(
            arch=self.arch,
            loss=self.loss,
            input_shape=self.input_shape,
            augments=self.augments,
            preprocessor=self.preprocessor,
        )
        model.compile()
        return model


def build(target, params):
    if target is None:
        return None
    elif isinstance(params, list):
        return target(*params)
    elif isinstance(params, dict):
        return target(**params)
    elif params is None:
        return target()
    else:
        return target(params)


CACHE = ".ivus_segmentation__cache"


def train(
    experiment: "Experiment",
    x,
    y,
    epochs=150,
    callbacks=None,
    batch_size=2,
    outpath=".",
    n_splits=3,
    verbose=False,
):
    if callbacks is None:
        callbacks = [
            keras.callbacks.ModelCheckpoint(
                "/tmp/weigts",
                save_weights_only=True,
                save_best_only=True,
                monitor="val_loss",
                mode="min",
            ),
            keras.callbacks.ReduceLROnPlateau(
                monitor="loss",
                factor=0.1,
                patience=15,
                verbose=1,
                min_delta=1e-4,
                mode="min",
            ),
        ]

    rng = np.random.RandomState(0)
    scores = []
    test_x = x[TRAIN_SIZE:]
    test_y = y[TRAIN_SIZE:]
    for i, (train_idx, val_idx) in enumerate(folds("B", n_splits=n_splits,
                                                   rng=rng)):
        tf.keras.backend.clear_session()

        train_x = x[train_idx]
        train_y = y[train_idx]
        valid_x = x[val_idx]
        valid_y = y[val_idx]

        model = experiment.build(seed=0)

        out = model.fit(
            train_x,
            train_y,
            validation_data=(valid_x, valid_y),
            epochs=epochs,
            callbacks=callbacks,
            batch_size=batch_size,
            verbose=verbose,
        )

        modelout, histout, scoreout = (
            os.path.join(outpath, f"{experiment.name}_{i}.h5"),
            os.path.join(outpath, f"{experiment.name}_history_{i}.csv"),
            os.path.join(outpath, f"{experiment.name}_score_{i}.csv"),
        )

        score = model.evaluate(test_x, test_y, verbose=verbose)
        scores.append(score)

        model.save(modelout)
        pd.DataFrame(out.history).to_csv(histout)
        score = pd.DataFrame(score)
        score.to_csv(scoreout)
        print(score.mean())
        del model  # trigger garbage collection
        gc.collect()

    return scores


class Parser:
    def __init__(self):
        self.parser = argparse.ArgumentParser()
        self.parser.add_argument("spec", type=self.yaml_type)

        sub = self.parser.add_subparsers()
        self.train = sub.add_parser("train")
        self.test = sub.add_parser("test")
        self.test.add_argument(
            "-o", "--output", type=argparse.FileType("w+"), required=False
        )

    @staticmethod
    def yaml_type(filename: str) -> dict:
        try:
            with open(filename, "r") as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            raise argparse.ArgumentTypeError(f"Could not find file {filename}")
        except yaml.parser.ParserError:
            raise argparse.ArgumentTypeError(
                f"Could not parser {filename} as yaml file"
            )

    def register(self, name: typing.Literal["train", "test"]):
        def inner(f):
            if name == "train":
                self.train.set_defaults(func=f)
            elif name == "test":
                self.test.set_defaults(func=f)
            return f

        return inner

    def __call__(self, fin=None):
        args = self.parser.parse_args(fin)
        func = args.func
        args = {k: v for k, v in vars(args).items() if k != "func"}
        return func(**args)


def load_results(name, location=".", folds=5):
    df = pd.concat(
        pd.read_csv(os.path.join(
            location, f"{name}_score_{f}.csv"), index_col=0)
        for f in range(folds)
    )
    for col in df.columns:
        if col.startswith("hd"):
            df[col] = np.minimum(df[col], np.sqrt(2 * 192 ** 2))
    return pd.concat([df.mean().add_suffix("_mean"), df.std().add_suffix("_std")])


if __name__ == "__main__":
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

    parser = Parser()

    @parser.register("train")
    def do_train(spec):
        dest = spec["destination"]
        dataset = spec.get('data-location', 'dataset/')
        os.makedirs(dest, exist_ok=True)
        os.makedirs(CACHE, exist_ok=True)
        x, y = get_data("B", dataset)
        x = preprocessing.minmax(x)

        verbose = spec.get("verbose", 1)
        batch_size = spec.get("batch-size", 2)
        n_splits = spec.get("splits", 5)
        epochs = spec["epochs"]

        experiments = [
            Experiment.from_dict({"input_shape": x.shape[1:], **e})
            for e in spec["experiments"]
        ]

        for experiment in experiments:
            if os.path.exists(os.path.join(CACHE, experiment.name)):
                logging.info(f'skipping experiment "{experiment.name}"')
                continue

            print("Running", experiment.name)
            train(experiment, x, y, epochs=epochs,
                  outpath=dest, verbose=verbose, batch_size=batch_size,
                  n_splits=n_splits)

            with open(os.path.join(CACHE, experiment.name), "w+") as f:
                print("OK", file=f)

    @parser.register("test")
    def do_test(spec, output=None):
        if output is None:
            output = sys.stdout
        dest = spec["destination"]
        verbose = spec.get("verbose", 1)
        dataset = spec.get('data-location', 'dataset/')
        x, y = get_data("B", dataset, verbose=verbose)
        x = preprocessing.minmax(x)
        experiments = [
            Experiment.from_dict({"input_shape": x.shape[1:], **e})
            for e in spec["experiments"]
        ]

        epochs = spec["epochs"]
        series = []
        for experiment in experiments:
            if not os.path.exists(os.path.join(CACHE, experiment.name)):
                train(experiment, x, y, epochs=epochs,
                      outpath=dest, verbose=verbose)
            print("Testing", experiment.name)
            s = load_results(experiment.name, location=dest)
            s["name"] = experiment.name
            series.append(s)
        df = pd.DataFrame(series)
        df.to_csv(output, index=False)

    parser()
