import copy
import logging
import math
from typing import Callable

import lmdb
import numpy as np
import sklearn
import torch
from torch import nn
from tqdm import tqdm

from utils.inference import scale_features, apply_sigmoid_if_necessary
from utils.utils import get_logger
from utils.data_utils import get_collated_batch_for_key, get_padding_length
from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature
from utils.dataset_utils import load_key_stats
from utils.training_utils import load_model_for_usage

from vsm_datascience_common import constants


def evaluate_model(model_configuration: dict,
                   training_configuration: dict,
                   test_ids: list,
                   evaluation_functions: list,
                   log_dir: str,
                   aggregate: bool = True,
                   logger: logging.Logger = None,
                   print_progress: bool = False) -> dict:
    if logger is None:
        logger = get_logger(__name__, f"{log_dir}/{model_configuration['id']}_{training_configuration['id']}.log")

    model, training_configuration = load_model_for_usage(
        model_configuration,
        training_configuration,
        test_ids[0]
    )

    dataset_dir = model_configuration["feature_config"]["dataset_dir"]
    lmdb_env = lmdb.open(dataset_dir, readonly=True)
    # get computation rank
    if torch.distributed.is_initialized():
        local_rank = torch.distributed.get_rank()
        torch.cuda.set_device(local_rank)
        device = torch.device(f"cuda:{local_rank}")
    else:
        local_rank = 0
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model.to(device)

    batch_size = training_configuration["batch_size"]
    logger.info(
        f"Rank {local_rank}: Estimated batch size: {batch_size} for GPU {torch.cuda.get_device_name(local_rank)}")

    # load key stats to retrieve individual cycles
    keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"])

    predict_fn = model_configuration["predict_fn"]
    actual_fn = model_configuration["actual_fn"]

    target_features = [x["name"] for x in model_configuration["feature_config"]["target_features"]]
    ignored_features = model_configuration["feature_config"]["ignored_features"]
    used_targets = [x for x in target_features if x not in ignored_features]
    scalers = get_scalers_for_model(model_configuration)

    errors = dict()

    for test_id in tqdm(test_ids, disable=not print_progress):
        # fetch stats for key
        current_key_stats = keys_stats["by_key"][test_id]
        cycle_stats = current_key_stats["cycle_stats"]
        # make sure cycle stats are sorted by start time
        # cycle_stats = sorted(cycle_stats, key=lambda x: x["starts_at"])
        for i in range(len(cycle_stats)):
            # compute cutoffs to isolate current cycle
            current_cycle_start_cutoff = sum([x["cycle_length"] for x in cycle_stats[:i]])
            current_cycle_end_cutoff = sum([x["cycle_length"] for x in cycle_stats[:i + 1]])
            try:
                batch = get_collated_batch_for_key(test_id, model_configuration,
                                                   start_cutoff=current_cycle_start_cutoff,
                                                   end_cutoff=current_cycle_end_cutoff,
                                                   lmdb_env=lmdb_env)
            except ValueError:
                print(f"Skipping {test_id}")
                continue
            if batch is None:
                print(f"Empty batch, skipping {test_id}")
                continue

            preds = predict_fn(model, batch,
                               batch_size=batch_size,
                               device=device)
            actuals = actual_fn(batch)
            num_outputs = preds.shape[-1] if len(preds.shape) > 1 else 1
            processed_preds = apply_sigmoid_if_necessary(preds, model_configuration)
            scaled_preds = scale_features(
                processed_preds,
                used_targets,
                model_configuration,
                scalers,
            )

            scaled_actuals = scale_features(
                actuals,
                used_targets,
                model_configuration,
                scalers,
            )

            for eval_fn in evaluation_functions:
                if eval_fn is not None:
                    eval_fn_name = eval_fn["name"]
                    if eval_fn_name not in errors:
                        errors[eval_fn_name] = dict()
                    eval_function = eval_fn["eval_fn"]
                    eval_fn_indices = [eval_fn["input_index"]] if "input_index" in eval_fn else range(num_outputs)
                    # skip error fn if actuals are nan, since they are ignored
                    for output_index in eval_fn_indices:
                        if any(np.isnan(scaled_actuals[:, output_index])):
                            continue
                        error = eval_function(scaled_preds[:, output_index],
                                              scaled_actuals[:, output_index])

                        if np.isnan(error):
                            # skip if error is nan
                            continue

                        if f"after_{i}" not in errors[eval_fn_name]:
                            errors[eval_fn_name][f"after_{i}"] = dict()

                        if output_index not in errors[eval_fn_name][f"after_{i}"]:
                            errors[eval_fn_name][f"after_{i}"][output_index] = list()

                        errors[eval_fn_name][f"after_{i}"][output_index].append(error)

    if aggregate:
        errors = aggregate_errors(errors, evaluation_functions)
    return errors


def aggregate_errors(errors: dict, evaluation_functions: list) -> dict:
    errors = copy.deepcopy(errors)
    for eval_fn in evaluation_functions:
        if eval_fn is not None:
            eval_fn_name = eval_fn["name"]
            accumulation_fn = eval_fn["accumulation_fn"]
            if eval_fn_name not in errors:
                continue
            for key in errors[eval_fn_name]:
                for output_index in errors[eval_fn_name][key]:
                    if len(errors[eval_fn_name][key][output_index]) == 0:
                        errors[eval_fn_name][key][output_index] = np.nan
                    else:
                        # only use non nan values
                        without_nans = [val for val in errors[eval_fn_name][key][output_index] if
                                        val is not None and val is not np.nan]
                        errors[eval_fn_name][key][output_index] = accumulation_fn(without_nans)

    return errors


def get_index_offset(offset_in_hours: int,
                     model_configuration: dict, ) -> float:
    step_size = model_configuration["preprocessing"]["window_shift"]
    measurements_per_day = constants.MEASUREMENTS_PER_DAY
    downsampling_factor = model_configuration["preprocessing"]["take_every_nth"]
    shift_hour_factor = 24 / (measurements_per_day / downsampling_factor) * step_size

    if offset_in_hours == 0:
        return 0

    index_offset_in_hours = math.ceil(offset_in_hours / shift_hour_factor)
    return index_offset_in_hours


def pre_ov_error(preds, actuals, *args, **kwargs):
    if any(np.isnan(actuals)):
        return np.nan
    # get ov day index
    try:
        ov_day_index = np.where(actuals == 0)[0][0]
    except IndexError:
        # no ov day in actuals
        return np.nan

    # get pre ov predictions
    pre_ov_preds = preds[:ov_day_index]
    pre_ov_actuals = actuals[:ov_day_index]

    if len(pre_ov_preds) == 0 or len(pre_ov_actuals) == 0:
        return np.nan

    # calculate error
    error = sklearn.metrics.mean_absolute_error(pre_ov_preds, pre_ov_actuals)
    return error


def after_ov_error(preds, actuals, *args, **kwargs):
    if any(np.isnan(actuals)):
        return np.nan
    # get ov day index
    try:
        ov_day_index = np.where(actuals == 0)[0][0]
    except IndexError:
        # no ov day in actuals
        return np.nan
    # get pre ov predictions
    after_ov_preds = preds[ov_day_index:]
    after_ov_actuals = actuals[ov_day_index:]

    if len(after_ov_preds) == 0 or len(after_ov_actuals) == 0:
        return np.nan

    # calculate error
    error = sklearn.metrics.mean_absolute_error(after_ov_preds, after_ov_actuals)
    return error


def ov_error(preds, actuals, model_configuration, *args, **kwargs):
    if any(np.isnan(actuals)):
        return np.nan
    # get ov day index
    try:
        ov_day_index = np.where(actuals == 0)[0][0]
    except IndexError:
        # no ov day in actuals
        return np.nan

    # get predicted ov index
    try:
        pred_ov_index = np.where(preds >= 0)[0][0]
    except IndexError:
        # no ov day in actuals
        return np.nan

    # calculate error
    error = abs(pred_ov_index - ov_day_index)

    # scale error to account for step size
    step_size = model_configuration["preprocessing"]["window_shift"]
    measurements_per_day = constants.MEASUREMENTS_PER_DAY
    downsampling_factor = model_configuration["preprocessing"]["take_every_nth"]
    shift_hour_factor = 24 / (measurements_per_day // downsampling_factor) * step_size
    error_in_days = error * shift_hour_factor / 24

    return error_in_days


def day_relative_to_ov_error(preds: np.ndarray | torch.Tensor,
                             actual: np.ndarray | torch.Tensor,
                             day_relative_to_ov: int,
                             model_configuration: dict) -> float:
    """
    Calculate the error of the model predictions relative to the ov day
    Args:
        preds: predictions
        actual: actual values
        day_relative_to_ov: day relative to ov day
        model_configuration: model configuration

    Returns:
        error: error of the model predictions relative to the ov day
    """

    if isinstance(preds, torch.Tensor):
        preds = preds.cpu().numpy()

    if isinstance(actual, torch.Tensor):
        actual = actual.cpu().numpy()

    if any(np.isnan(actual)):
        return np.nan

    # get ov day index
    try:
        ov_day_index = np.where(actual == 0)[0][0]
    except IndexError:
        # no ov day in actuals
        return np.nan

    # get index offset factor -> how much time between each step
    index_offset = get_index_offset(day_relative_to_ov * 24, model_configuration)

    if ov_day_index + index_offset >= preds.shape[0]:
        relative_day_pred = preds[-1] if len(preds.shape) > 1 else preds[-1]
    elif ov_day_index + index_offset < 0:
        relative_day_pred = preds[0][0] if len(preds.shape) > 1 else preds[0]
    else:
        relative_day_pred = preds[ov_day_index + index_offset][0] if len(preds.shape) > 1 else \
            preds[ov_day_index + index_offset][0]

    relative_day_actual = day_relative_to_ov

    # calculate error
    error = np.abs(relative_day_pred - relative_day_actual)
    return error


def get_ov_over_ov_index(input_sequence: np.ndarray) -> int:
    """
    Returns the index of the ovulation based on the ov_over feature curve.

    Ov index is first non-zero index in actual curve
    Args:
        input_sequence: input sequence
    Returns:
        index of the ovulation based on the ov_over feature curve
    """

    try:
        ov_index = int(np.where(input_sequence > 0)[0][0])
    except IndexError:
        return None

    return ov_index


def get_ov_over_pre_ov_error(preds: np.ndarray | torch.Tensor,
                             actual: np.ndarray | torch.Tensor,
                             error_fn: Callable) -> float:
    ov_index = get_ov_over_ov_index(actual)

    if ov_index is None:
        return np.nan

    if any(np.isnan(actual)):
        return np.nan

    pre_ov_preds = preds[:ov_index]
    pre_ov_actuals = actual[:ov_index]

    if len(pre_ov_preds) == 0 or len(pre_ov_actuals) == 0:
        return np.nan

    error = error_fn(pre_ov_preds, pre_ov_actuals)
    return error


def get_ov_over_post_ov_error(preds: np.ndarray | torch.Tensor,
                              actual: np.ndarray | torch.Tensor,
                              error_fn: Callable) -> float:
    ov_index = get_ov_over_ov_index(actual)
    if any(np.isnan(actual)):
        return np.nan
    post_ov_preds = preds[ov_index:]
    post_ov_actuals = actual[ov_index:]
    if len(post_ov_preds) == 0 or len(post_ov_actuals) == 0:
        return np.nan
    error = error_fn(post_ov_preds, post_ov_actuals)
    return error


def get_pre_fertility_error(preds: np.ndarray | torch.Tensor,
                            actual: np.ndarray | torch.Tensor,
                            error_fn: Callable) -> float:
    try:
        fertility_start_index = np.where(actual > 0)[0][0]
    except IndexError:
        return np.nan

    pre_fertility_preds = preds[:fertility_start_index]
    pre_fertility_actuals = actual[:fertility_start_index]

    if len(pre_fertility_preds) == 0 or len(pre_fertility_actuals) == 0:
        return np.nan

    error = error_fn(pre_fertility_preds, pre_fertility_actuals)
    return error


def get_post_fertility_error(preds: np.ndarray | torch.Tensor,
                             actual: np.ndarray | torch.Tensor,
                             error_fn: Callable) -> float:
    try:
        fertility_end_index = np.where(actual > 0)[-1][-1]
    except IndexError:
        return np.nan

    pre_fertility_preds = preds[fertility_end_index:]
    pre_fertility_actuals = actual[fertility_end_index:]

    if len(pre_fertility_preds) == 0 or len(pre_fertility_actuals) == 0:
        return np.nan

    error = error_fn(pre_fertility_preds, pre_fertility_actuals)
    return error


def get_during_fertility_error(preds: np.ndarray | torch.Tensor,
                               actual: np.ndarray | torch.Tensor,
                               error_fn: Callable) -> float:
    try:
        fertility_indices = np.where(actual > 0)[0]
    except IndexError:
        return np.nan

    during_fertility_preds = preds[fertility_indices]
    during_fertility_actuals = actual[fertility_indices]

    if len(during_fertility_preds) == 0 or len(during_fertility_actuals) == 0:
        return np.nan

    error = error_fn(during_fertility_preds, during_fertility_actuals)
    return error


def get_non_fertility_error(preds: np.ndarray | torch.Tensor,
                            actual: np.ndarray | torch.Tensor,
                            error_fn: Callable) -> float:
    """
    Calculate the error of the model predictions for non-fertility periods.
    Args:
        preds: predictions
        actual: actual values
    Returns:
        error: error of the model predictions for non-fertility periods
    """

    if any(np.isnan(actual)):
        return np.nan

    try:
        non_fertility_indices = np.where(actual == 0)[0]
    except IndexError:
        return np.nan

    if len(non_fertility_indices) == 0:
        return np.nan

    # get non-fertility predictions
    non_fertility_preds = preds[non_fertility_indices]
    non_fertility_actuals = actual[non_fertility_indices]
    if len(non_fertility_preds) == 0 or len(non_fertility_actuals) == 0:
        return np.nan

    # calculate error
    error = error_fn(non_fertility_preds, non_fertility_actuals)
    return error


def get_fertility_over_index(input_sequence: np.ndarray) -> int:
    """
    Returns the index of the fertility based on the fertility feature curve.

    Fertility index is first non-zero index in actual curve
    Args:
        input_sequence: input sequence
    Returns:
        index of the fertility based on the fertility feature curve
    """

    try:
        fertility_index = np.where(input_sequence > 0)[-1][-1]
    except IndexError:
        return None

    return fertility_index


def get_error_relative_to_x(preds: np.ndarray | torch.Tensor,
                            actual: np.ndarray | torch.Tensor,
                            x_fn: Callable,
                            offset_in_hours: int,
                            model_configuration: dict) -> float:
    """
    Calculate the error at the index relative
    Args:
        preds: predictions
        actual: actual values
        x_fn: function to get the index of the x value
        offset_in_hours: offset in hours to apply to the x index
        model_configuration: model configuration

    Returns:
        error: error of the model predictions relative to the x value
    """

    if isinstance(preds, torch.Tensor):
        preds = preds.cpu().numpy()

    if isinstance(actual, torch.Tensor):
        actual = actual.cpu().numpy()

    if any(np.isnan(actual)):
        return np.nan

    # get x index
    try:
        x_index = x_fn(actual)
    except Exception:
        return np.nan

    if x_index is None:
        return np.nan

    # get index offset factor -> how much time between each step
    index_offset = get_index_offset(offset_in_hours, model_configuration)

    if x_index + index_offset >= preds.shape[0]:
        offset_pred = preds[-1]
        offset_actual = actual[-1]
    elif x_index + index_offset < 0:
        offset_pred = preds[0]
        offset_actual = actual[0]
    else:
        offset_pred = preds[x_index + index_offset]
        offset_actual = actual[x_index + index_offset]

    # calculate error
    error = np.abs(offset_pred - offset_actual)
    return error


def r2_wrapper(preds: np.ndarray | torch.Tensor,
               actual: np.ndarray | torch.Tensor,
               var_tol: float = 1e-8,
               max_allowed_var: float = 1.0):
    if len(preds) <= 1 or len(actual) <= 1:
        return np.nan

    # make sure to capture the case where actuals have no variance
    if np.var(actual) < var_tol:
        # return variance of predictions here, as it should be at best zero
        # cap at 1.0 to make it comparable
        error = min(np.var(preds), max_allowed_var)
        return error

    # return 1 - r2 to get an error based metric, where lower is better
    error = 1 - sklearn.metrics.r2_score(actual, preds)
    return error