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

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 vsm_datascience_common import constants


def evaluate_model(model_configuration: dict,
                   training_configuration: dict,
                   test_ids: list,
                   evaluation_functions: list) -> dict:
    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:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    load_fn = model_configuration["model_load_fn"]
    model = load_fn(model_configuration=model_configuration,
                    training_configuration=training_configuration,
                    sample_key=test_ids[0],
                    device=device,
                    lmdb_env=lmdb_env)
    # load key stats to retrieve individual cycles
    keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"])

    batch_size = training_configuration["batch_size"]
    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):
        # fetch stats for key
        current_key_stats = keys_stats["by_key"][test_id]
        cycle_stats = current_key_stats["cycle_stats"]
        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
            scaled_preds = list()
            for j in range(num_outputs):
                if isinstance(preds, torch.Tensor):
                    preds = preds.cpu().numpy()

                if len(preds.shape) == 3:
                    output = preds[:, 0, j].squeeze()
                else:
                    output = preds[:, j] if num_outputs > 1 else preds

                if "loss_functions" in training_configuration and len(training_configuration["loss_functions"]) > j:
                    if isinstance(training_configuration["loss_functions"][j], nn.BCEWithLogitsLoss):
                        output = torch.sigmoid(torch.tensor(output)).numpy()

                scaled_output = inverse_scale_feature(output,
                                                      used_targets[j],
                                                      scalers)

                scaled_preds.append(scaled_output)

            scaled_actuals = list()
            for j in range(num_outputs):
                if len(actuals.shape) == 3:
                    output = actuals[:, 0, j].squeeze()
                else:
                    output = actuals[:, j] if num_outputs > 1 else actuals
                if isinstance(output, torch.Tensor):
                    output = output.cpu().numpy()
                scaled_output = inverse_scale_feature(output,
                                                      used_targets[j],
                                                      scalers).ravel()

                scaled_actuals.append(scaled_output)

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

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

                    if eval_fn_name not in errors:
                        errors[eval_fn_name] = dict()

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

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

    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]:
                if len(errors[eval_fn_name][key]) == 0:
                    errors[eval_fn_name][key] = np.nan
                else:
                    errors[eval_fn_name][key] = accumulation_fn(errors[eval_fn_name][key])
    return errors


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
    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

    index_offset = int(day_relative_to_ov * (shift_hour_factor // 24))

    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