import os

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

from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature
from utils.data_utils import get_collated_batch_for_key
from utils.dataset_utils import get_dataset_path
from utils.model_utils import get_model_config_from_file
from utils.training_utils import get_training_config_from_file, get_data_ids, find_max_batch_size


def load_model(model_dir: str,
               results_base_dir: str,
               lmdb_base_dir: str,
               device: str = "cuda") -> tuple:
    """
    Load the model from the model directory.
    :param model_dir: model directory
    :param results_base_dir: results base directory
    :param lmdb_base_dir: LMDB base directory
    :param device: device to use
    :return: model configuration, training configuration, dataset directory
    """

    training_id = None
    model_configuration = get_model_config_from_file(model_dir, results_base_dir, lmdb_base_dir)
    feature_config = model_configuration["feature_config"]
    dataset_dir = get_dataset_path(lmdb_base_dir, feature_config["feature_set_name"])

    # get list of trainings
    training_base_dir = os.path.join(model_dir, "trainings")
    # if no training id is given, use the latest training
    if training_id is None:
        training_dirs = os.listdir(training_base_dir)
        training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
                         os.path.isdir(os.path.join(training_base_dir, dir))]
        training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
        training_dir = training_dirs[0]
    else:
        training_dir = os.path.join(training_base_dir, training_id)

    training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)

    limit = None
    train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, limit)

    dataset_dir = f"{lmdb_base_dir}/{feature_config['feature_set_name']}"
    env = lmdb.open(dataset_dir, readonly=True)
    load_fn = model_configuration["model_load_fn"]
    model = load_fn(model_configuration, training_configuration, test_ids[0], device, lmdb_env=env)

    # estimate batch size
    sample_batch = get_collated_batch_for_key(train_ids[0], model_configuration, lmdb_env=env)
    sample_input = sample_batch[0][0]  # get the first item in the batch
    input_shape = sample_input.shape
    # send model to proper device
    model.to(device)
    batch_size = find_max_batch_size(model,
                                     input_shape,
                                     device=device)
    training_configuration["batch_size"] = batch_size
    return model, model_configuration, training_configuration, train_ids, val_ids, test_ids


def get_training(model_configuration: dict,
                 results_base_dir: str,
                 training_id: str = None) -> dict:
    """
    Get the training configuration from the model directory.
    :param model_configuration: model configuration
    :param training_id: training name
    :return: training configuration
    """
    model_dir = model_configuration["model_dir"]
    # get list of trainings
    training_base_dir = os.path.join(model_dir, "trainings")
    # if no training id is given, use the latest training
    if training_id is None:
        training_dirs = os.listdir(training_base_dir)
        training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
                         os.path.isdir(os.path.join(training_base_dir, dir))]
        training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
        training_dir = training_dirs[0]
    else:
        training_dir = os.path.join(training_base_dir, training_id)

    training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
    return training_configuration


def find_models(base_dir: str, current_dir: str = None) -> dict:
    """
    Find all models in the base directory.
    :param base_dir: base directory
    :return: list of model directories
    """
    model_dirs = dict()
    is_root = base_dir == current_dir
    if current_dir is None:
        current_dir = base_dir
    for sub_dir_name in os.listdir(current_dir):
        sub_dir = os.path.join(current_dir, sub_dir_name)
        if not os.path.isdir(sub_dir):
            continue
        model_config_path = os.path.join(sub_dir, "model_configuration.pickle")

        if current_dir not in model_dirs:
            model_dirs[current_dir] = list()
        if os.path.exists(model_config_path) and is_root:
            model_dirs[current_dir].append(sub_dir)
        elif os.path.exists(model_config_path) and not is_root:
            model_dirs[current_dir].append(sub_dir)
        else:
            model_dirs = model_dirs | find_models(base_dir, sub_dir)

    return model_dirs


def apply_sigmoid_if_necessary(model_outputs: np.ndarray | torch.Tensor,
                               model_configuration: dict) -> np.ndarray:
    """
    Applies a sigmoid to the outputs of a model, if they use a BCE loss.
    Args:
        model_outputs: models outputs
        model_configuration: configuration of the model that produced the outputs

    Returns:
        processed_model_outputs: processed model outputs with sigmoid applied where necessary
    """

    target_features = model_configuration["feature_config"]["target_features"]
    ignored_features = model_configuration["feature_config"]["ignored_features"]
    used_targets = [feature for feature in target_features if feature["name"] not in ignored_features]

    is_3d = len(model_outputs.shape) == 3

    processed_outputs = np.empty_like(model_outputs)
    for i, used_target in enumerate(used_targets):
        if used_target["loss_fn"] == nn.BCEWithLogitsLoss:
            if is_3d:
                target_values = model_outputs[:, 0, i]
                if isinstance(target_values, np.ndarray):
                    target_values = torch.from_numpy(target_values)
                processed_outputs[:, 0, i] = torch.sigmoid(target_values).numpy()
            else:
                target_values = model_outputs[:, i]
                if isinstance(target_values, np.ndarray):
                    target_values = torch.from_numpy(target_values)
                processed_outputs[:, i] = torch.sigmoid(target_values).numpy()
        else:
            if is_3d:
                processed_outputs[:, 0, i] = model_outputs[:, 0, i].numpy() if isinstance(model_outputs,
                                                                                          torch.Tensor) else model_outputs[
                                                                                                             :, 0, i]
            else:
                processed_outputs[:, i] = model_outputs[:, i].numpy() if isinstance(model_outputs,
                                                                                    torch.Tensor) else model_outputs[:,
                                                                                                       i]

    return processed_outputs


def scale_features(features: np.ndarray | torch.Tensor,
                   feature_names: list | None,
                   model_configuration: dict,
                   scalers: dict = None,
                   subset_name: str = "train") -> np.ndarray:
    """
    Scales the given features based on the scalers associated with the given model configuration
    Args:
        features: feature arrays, may be up to 3D (but may only have 2 usable dimensions, (B, 1, F) is allowed)
        feature_names: names of the features in the same order as in the input array, may be None, then no scaling is performed
        model_configuration: model configuration of the model the features belong to
        scalers: dict of scales associated by name, will be fetched from model configuration, if None
        subset_name: subset name, either "train", "test" or "val", used for determining the correct scalers

    Returns:
        features: scaled features
    """

    if features.shape[-1] != len(feature_names):
        raise ValueError("Number of features does not match number of given feature names")

    # remove singleton dimension, if 3D
    if len(features.shape) == 3:
        if features.shape[1] != 1:
            raise ValueError(f"Input has 3 dimensions, expected (B, 1, F), but got {features.shape}")
        features = features[:, 0, :]

    if scalers is None:
        scalers = get_scalers_for_model(model_configuration)

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

    scaled_features = np.empty_like(features)
    for i in range(len(feature_names)):
        # skip scaling, if feature name is None
        if feature_names[i] is not None:
            scaled_output = inverse_scale_feature(features[:, i],
                                                  feature_names[i],
                                                  scalers,
                                                  subset_name=subset_name)
        else:
            scaled_output = features[:, i]

        # make sure to make into 1D array when assigning
        scaled_features[:, i] = scaled_output.ravel()

    return scaled_features