import csv
import json
import os
from enum import Enum

from keras.models import load_model
from keras.utils import Sequence

import numpy as np


def get_properties_json(path_to_directory):
    properties_json_path = os.path.join(path_to_directory, "properties.json")
    with open(properties_json_path, "r") as properties_json_file:
        properties_json = json.load(properties_json_file)
    return properties_json


class BatchNumpyFileSequence(Sequence):
    def __init__(self, np_batch_directory, max_base_count, properties_json=None, array_type='train'):
        if properties_json is None:
            properties_json = get_properties_json(np_batch_directory)
        self.properties_json = properties_json
        self.batch_path_and_prefix = os.path.join(np_batch_directory, self.properties_json["batch_prefix"])
        self.max_base_count = max_base_count
        if array_type not in ['train', 'vcf', 'eval']:
            self.array_type = 'train'
        else:
            self.array_type = array_type

    def __len__(self):
        return self.properties_json["total_batches_written"]

    def __getitem__(self, index):
        with np.load("{}_{}.npz".format(self.batch_path_and_prefix, index)) as batch_data_set:
            batch_input = batch_data_set["input"]
            batch_label = batch_data_set["label"]
            # Prepad timesteps, postpad features/labels (if there's a shape mismatch)
            num_base_diff = max(0, self.max_base_count - batch_input.shape[1])
            timestep_padding = (num_base_diff, 0) if self.properties_json["padding"] == "pre" else (0, num_base_diff)
            batch_input = self._pad_batch(batch_input, timestep_padding)
            batch_label = self._pad_batch(batch_label, timestep_padding)
            batch_output = ({"model_input": batch_input}, {"main_output": batch_label})
            if self.array_type == 'vcf':
                batch_ref = self._pad_batch(batch_data_set["ref"], timestep_padding, pad_3d=False)
                batch_location = self._pad_batch(batch_data_set["location"], timestep_padding, pad_3d=False)
                batch_chromosome = batch_data_set["chromosome"]
                return batch_output, batch_ref, batch_location, batch_chromosome
            elif self.array_type == 'eval':
                batch_metadata = self._pad_batch(batch_data_set["metadata"], timestep_padding, pad_3d=False)
                return batch_output, batch_metadata
            else:
                return batch_output

    def _pad_batch(self, batch_array, timestep_padding, pad_3d=True):
        pad_width = ((0, 0), timestep_padding, (0, 0)) if pad_3d else ((0, 0), timestep_padding)
        batch_array_padded = np.pad(batch_array,
                                    pad_width=pad_width,
                                    mode="constant")
        if self.max_base_count < batch_array.shape[1]:
            if self.properties_json["padding"] == "post":
                if pad_3d:
                    batch_array_padded = batch_array_padded[:, :self.max_base_count, :]
                else:
                    batch_array_padded = batch_array_padded[:, :self.max_base_count]
            else:
                start_base = batch_array.shape[1] - self.max_base_count
                if pad_3d:
                    batch_array_padded = batch_array_padded[:, start_base:, :]
                else:
                    batch_array_padded = batch_array_padded[:, start_base:]
        return batch_array_padded

    def on_epoch_end(self):
        pass


class Metadata(Enum):
    REF = 0
    SNP = 1
    INDEL = 2


class ModelEvaluator:
    def __init__(self, test_data_path, log_path, write_header, log_epochs, max_base_count=None, main_model=None):
        self.properties_json = get_properties_json(test_data_path)
        if main_model is not None:
            self.main_model = ModelEvaluator._get_model(main_model)
        else:
            self.main_model = None
        if max_base_count is None:
            if self.main_model is not None and self.main_model.input_shape[1] is not None:
                max_base_count = self.main_model.input_shape[1]
            else:
                max_base_count = self.properties_json["max_base_count"]
        self.test_data = BatchNumpyFileSequence(test_data_path, max_base_count,
                                                self.properties_json, array_type='eval')
        field_names = ["epoch"] if log_epochs else []
        field_names += ["precision_ref", "precision_snp", "precision_indel", "recall_ref", "recall_snp", "recall_indel",
                        "f1_ref", "f1_snp", "f1_indel", "accuracy_overall", "accuracy_ref", "accuracy_snp",
                        "accuracy_indel", "true_or_predicted_ref", "true_or_predicted_snp", "true_or_predicted_indel",
                        "true_ref", "true_snp", "true_indel", "predicted_ref", "predicted_snp", "predicted_indel"]
        self.log_file = open(log_path, "a")
        self.log_writer = csv.DictWriter(self.log_file, fieldnames=field_names, quoting=csv.QUOTE_NONNUMERIC,
                                         delimiter="\t", lineterminator="\n")
        if write_header:
            self.log_writer.writeheader()

    @staticmethod
    def _get_model(model):
        if type(model) == str:
            return load_model(model)
        else:
            return model

    def eval_model(self, model_to_eval=None, epoch=None):
        if self.main_model is not None:
            model = self.main_model
        elif model_to_eval is not None:
            model = ModelEvaluator._get_model(model_to_eval)
        else:
            raise Exception("No model provided to evaluator")
        # Not using predict_generator here because difficult to reconstruct order of inputs,
        # see https://github.com/fchollet/keras/issues/5048
        # Using batch_prediction rather than predicting on each sample to potentially speed up predictions
        count_true_or_predicted_ref = 0
        count_true_or_predicted_snp = 0
        count_true_or_predicted_indel = 0
        count_predicted_ref = 0
        count_predicted_snp = 0
        count_predicted_indel = 0
        count_true_ref = 0
        count_true_snp = 0
        count_true_indel = 0
        correct_overall = 0
        correct_ref = 0
        correct_snp = 0
        correct_indel = 0
        count_overall = 0
        tp_ref = 0
        tp_snp = 0
        tp_indel = 0
        fp_ref = 0
        fp_snp = 0
        fp_indel = 0
        fn_ref = 0
        fn_snp = 0
        fn_indel = 0
        for data_idx in range(len(self.test_data)):
            (batch_input_dict, batch_label_dict), batch_metadata = self.test_data[data_idx]
            batch_input = batch_input_dict["model_input"]
            batch_label = batch_label_dict["main_output"]
            # batch_predictions, _ = model.predict_on_batch(batch_input)
            batch_predictions = model.predict_on_batch(batch_input)
            for segment_in_batch_idx in range(batch_predictions.shape[0]):
                segment_label_categorical = batch_label[segment_in_batch_idx]
                segment_prediction_categorical = batch_predictions[segment_in_batch_idx]
                segment_metadata_with_padding = batch_metadata[segment_in_batch_idx]
                # segment_predicted_metadata_categorical = batch_predicted_metadata[segment_in_batch_idx]
                segment_label_with_padding = np.argmax(segment_label_categorical, axis=1)
                # Get true genotypes from these segment labels
                # Keep track of tp, fp, tn, and fn based on metadata ref, snp, indel, and calculate acc, prec, rec, F1
                segment_prediction_with_padding = np.argmax(segment_prediction_categorical, axis=1)
                # segment_predicted_metadata_with_padding = np.argmax(segment_predicted_metadata_categorical, axis=1)
                # Only use positions where label != 0, as label 0 reserved for padding
                segment_label_non_padding_positions = segment_label_with_padding != 0
                segment_label = np.extract(segment_label_non_padding_positions, segment_label_with_padding)
                segment_prediction = np.extract(segment_label_non_padding_positions, segment_prediction_with_padding)
                segment_metadata = np.extract(segment_label_non_padding_positions, segment_metadata_with_padding)
                # segment_predicted_metadata = np.extract(segment_label_non_padding_positions,
                #                                         segment_predicted_metadata_with_padding)
                # segment_true_genotype_label = [properties_json["genotype.segment.label_plus_one.{}".format(label)]
                #                                for label in segment_label]
                segment_true_genotype_prediction = [
                    self.properties_json["genotype.segment.label_plus_one.{}".format(label)]
                    for label in segment_prediction
                ]

                for base_idx in range(segment_label.shape[0]):
                    count_overall += 1
                    true_or_predicted_ref_at_base = False
                    true_or_predicted_indel_at_base = False
                    true_or_predicted_snp_at_base = False

                    base_metadata_value = segment_metadata[base_idx]
                    if base_metadata_value == 0:
                        base_metadata = Metadata.REF
                        true_or_predicted_ref_at_base = True
                        count_true_ref += 1
                    elif base_metadata_value == 1:
                        base_metadata = Metadata.SNP
                        true_or_predicted_snp_at_base = True
                        count_true_snp += 1
                    elif base_metadata_value == 2:
                        base_metadata = Metadata.INDEL
                        true_or_predicted_indel_at_base = True
                        count_true_indel += 1
                    else:
                        raise Exception("Unknown metadata value")

                    # Check if prediction is for indel, SNP or ref based on true genotype corresponding to label
                    # If it contains a "-", prediction is indel
                    # If heterozygous, prediction is SNP
                    # Otherwise, is ref
                    base_true_genotype_prediction = segment_true_genotype_prediction[base_idx]
                    if "-" in base_true_genotype_prediction:
                        base_predicted_metadata = Metadata.INDEL
                        true_or_predicted_indel_at_base = True
                        count_predicted_indel += 1
                    elif len("".join(set(base_true_genotype_prediction))) > 1:
                        base_predicted_metadata = Metadata.SNP
                        true_or_predicted_snp_at_base = True
                        count_predicted_snp += 1
                    else:
                        base_predicted_metadata = Metadata.REF
                        true_or_predicted_ref_at_base = True
                        count_predicted_ref += 1

                    if true_or_predicted_ref_at_base:
                        count_true_or_predicted_ref += 1
                    if true_or_predicted_snp_at_base:
                        count_true_or_predicted_snp += 1
                    if true_or_predicted_indel_at_base:
                        count_true_or_predicted_indel += 1

                    if segment_prediction[base_idx] == segment_label[base_idx]:
                        correct_overall += 1
                        if true_or_predicted_ref_at_base:
                            correct_ref += 1
                        if true_or_predicted_snp_at_base:
                            correct_snp += 1
                        if true_or_predicted_indel_at_base:
                            correct_indel += 1

                        if base_metadata == Metadata.REF:
                            tp_ref += 1
                        elif base_metadata == Metadata.SNP:
                            tp_snp += 1
                        elif base_metadata == Metadata.INDEL:
                            tp_indel += 1
                        else:
                            raise Exception("Unknown metadata value")
                    else:
                        if base_metadata == Metadata.REF:
                            fn_ref += 1
                        elif base_metadata == Metadata.SNP:
                            fn_snp += 1
                        elif base_metadata == Metadata.INDEL:
                            fn_indel += 1
                        else:
                            raise Exception("Unknown metadata value")

                    if base_predicted_metadata == Metadata.REF and not base_metadata == Metadata.REF:
                        fp_ref += 1
                    elif base_predicted_metadata == Metadata.SNP and not base_metadata == Metadata.SNP:
                        fp_snp += 1
                    elif base_predicted_metadata == Metadata.INDEL and not base_metadata == Metadata.INDEL:
                        fp_indel += 1

        precision_ref = np.divide(tp_ref, (tp_ref + fp_ref))
        precision_snp = np.divide(tp_snp, (tp_snp + fp_snp))
        precision_indel = np.divide(tp_indel, (tp_indel + fp_indel))

        recall_ref = np.divide(tp_ref, (tp_ref + fn_ref))
        recall_snp = np.divide(tp_snp, (tp_snp + fn_snp))
        recall_indel = np.divide(tp_indel, (tp_indel + fn_indel))

        row_dict = {
            "precision_ref": precision_ref,
            "precision_snp": precision_snp,
            "precision_indel": precision_indel,
            "recall_ref": recall_ref,
            "recall_snp": recall_snp,
            "recall_indel": recall_indel,
            "f1_ref": np.divide(2 * precision_ref * recall_ref, precision_ref + recall_ref),
            "f1_snp": np.divide(2 * precision_snp * recall_snp, precision_snp + recall_snp),
            "f1_indel": np.divide(2 * precision_indel * recall_indel, precision_indel + recall_indel),
            "accuracy_overall": np.divide(correct_overall, count_overall),
            "accuracy_ref": np.divide(correct_ref, count_true_or_predicted_ref),
            "accuracy_snp": np.divide(correct_snp, count_true_or_predicted_snp),
            "accuracy_indel": np.divide(correct_indel, count_true_or_predicted_indel),
            "true_or_predicted_ref": count_true_or_predicted_ref,
            "true_or_predicted_snp": count_true_or_predicted_snp,
            "true_or_predicted_indel": count_true_or_predicted_indel,
            "true_ref": count_true_ref,
            "true_snp": count_true_snp,
            "true_indel": count_true_indel,
            "predicted_ref": count_predicted_ref,
            "predicted_snp": count_predicted_snp,
            "predicted_indel": count_predicted_indel,
        }
        if epoch is not None:
            row_dict["epoch"] = epoch
        self.log_writer.writerow(row_dict)
        self.log_file.flush()

    def close_log(self):
        self.log_file.close()