import os

import lmdb
import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import joblib

from utils.lmdb_utils import load_from_lmdb, get_lmdb_keys


def get_features(cycle: dict,
                 feature_config: dict) -> dict:
    """
    Computes the feature for a given cycle defined by the feature_config
    :param cycle: cycle data as dictionary
    :param feature_config: feature config as dictionary
    :return: dict with feature according feature config
    """
    features = {}
    feature_sets = feature_config["feature_sets"]
    for feature_set in feature_sets:
        for feature_def in feature_config[feature_set]:
            if feature_set not in features:
                features[feature_set] = {}

            # skip features that are marked as ignored
            if "ignored_features" in feature_config and feature_def["name"] in feature_config["ignored_features"]:
                continue

            feature_return = feature_def["fn"](cycle=cycle)
            if isinstance(feature_return, dict):
                if len(feature_return) > 1:
                    for feature_name, feature in feature_return.items():
                        features[feature_set][f"{feature_def['name']}_{feature_name}"] = feature
                else:
                    features[feature_set][feature_def["name"]] = list(feature_return.values())[0]
            else:
                features[feature_set][feature_def["name"]] = feature_return

        # check, if all features have same length
        if feature_set in features and len(features[feature_set]) > 0 and isinstance(
                list(features[feature_set].values())[0], np.ndarray):
            feature_lengths = [len(x) for x in features[feature_set].values()]
            if len(set(feature_lengths)) > 1:
                raise ValueError(f"Feature set {feature_set} has features of different lengths: {feature_lengths}")

    return features


def save_scalers(scalers: dict,
                 scaler_dir: str):
    if not os.path.exists(scaler_dir):
        os.makedirs(scaler_dir)
    for feature_type in scalers:
        for feature_name in scalers[feature_type]:
            joblib.dump(scalers[feature_type][feature_name],
                        f"{scaler_dir}/{feature_name}.pkl")


def load_scalers(scaler_dir: str) -> dict:
    scalers = dict()
    for scaler_file in os.listdir(scaler_dir):
        if scaler_file.endswith(".pkl"):
            feature_name = scaler_file.replace(".pkl", "")
            scalers[feature_name] = joblib.load(f"{scaler_dir}/{scaler_file}")
    return scalers


def get_scalers_for_model(model_configuration: dict):
    """
    Get the scalers for the features of a model configuration
    Args:
        model_configuration: configuration of the model

    Returns:
        scalers: scalers for the model

    """

    scaler_dir = os.path.join(model_configuration["feature_config"]["dataset_dir"], "scalers")
    scalers = load_scalers(scaler_dir)
    return scalers


def get_feature_values(feature_type: str | None,
                       feature_name: str,
                       env: lmdb.Environment,
                       keys: list = None) -> list:
    """
    Get the values of a feature from the lmdb dataset for all keys
    :param feature_type: type of feature, can be None, then the first feature with the given name will be used
    :param feature_name: feature name to extract
    :param env: lmdb environment
    :param keys: keys to extract the feature from, if None, all keys will be used / fetched from the database
    :return: list with feature values
    """
    if keys is None:
        keys = get_lmdb_keys(env)
    feature_values = []
    for key in keys:
        data = load_from_lmdb(env, key)
        if feature_type is None:
            for feature_type, feature_data in data.items():
                if feature_name in feature_data:
                    feature_values.append(feature_data[feature_name])
                    break
        else:
            for current_feature_name, feature_data in data[feature_type].items():
                # catch sub features that have been prefixed with the feature name
                if current_feature_name.startswith(feature_name):
                    feature_values.append(feature_data)
                    break
    return feature_values


def train_scalers(feature_type: str,
                  feature_name: str,
                  scaler_type,
                  sample,
                  env) -> dict:
    scalers = dict()
    all_features = sample[feature_type]
    individual_feature_names = list()
    for individual_feature_name in all_features:
        if individual_feature_name.startswith(feature_name) and not individual_feature_name.endswith("_scaled"):
            individual_feature_names.append(individual_feature_name)

    for individual_feature_name in individual_feature_names:
        feature_values = get_feature_values(feature_type, individual_feature_name, env)
        if len(feature_values) == 0:
            raise ValueError(f"No feature values found for feature {individual_feature_name}")
        if isinstance(feature_values[0], list) or isinstance(feature_values[0], np.ndarray):
            features_reshaped = np.concatenate(feature_values).reshape(-1, 1)
        else:
            features_reshaped = np.array(feature_values).reshape(-1, 1)
        del feature_values

        if scaler_type is not None:
            scaler = scaler_type()
            scaler.fit(features_reshaped)
        else:
            scaler = None

        scalers[individual_feature_name] = scaler
    return scalers


def scale_item(data: dict,
               scalers: dict):
    data_format = dict()
    for feature_set in data:
        if feature_set not in data_format:
            data_format[feature_set] = dict()
        for feature_name in data[feature_set]:
            if feature_name.endswith("_scaled"):
                continue
            data_format[feature_set][feature_name] = data[feature_set][feature_name]

    for feature_set in data_format:
        for feature_name in data_format[feature_set]:
            feature_values = data[feature_set][feature_name]
            if scalers[feature_name] is not None:
                if isinstance(feature_values, list) or isinstance(feature_values, np.ndarray):
                    scaled_feature_values = scalers[feature_name].transform(
                        feature_values.reshape(-1, 1)).flatten()
                else:
                    scaled_feature_values = scalers[feature_name].transform(
                        np.array(feature_values).reshape(-1, 1)).flatten()
            else:
                scaled_feature_values = feature_values
            data[feature_set][f"{feature_name}_scaled"] = scaled_feature_values

    return data


def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor,
                          feature_names: np.ndarray | list | str,
                          scalers: dict) -> np.ndarray:
    """
    Inverse scales the input features, can handle single and multi feature input
    Args:
        input_feature: input feature as a numpy array
        feature_names: names of input features as reference for scalers
        scalers: dict of feature scalers

    Returns:
        scaled input features as numpy array
    """

    if isinstance(feature_names, str):
        feature_names = [feature_names]

    if isinstance(input_feature, np.ndarray) or isinstance(input_feature, torch.Tensor):
        input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1
    elif isinstance(input_feature, float) or isinstance(input_feature, int):
        # handle case, where input is scalar
        input_scaled = scalers[feature_names[0]].inverse_transform([[input_feature]])[0][0]
        return input_scaled
    elif isinstance(input_feature, list):
        input_feature = np.array(input_feature)
        input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1
    else:
        raise ValueError(f"Unsupported input type: {type(input_feature)}")

    # reshape, if input is one dimensional
    if input_dim == 1:
        input_feature = input_feature.reshape(-1, 1)

    input_scaled = input_feature.copy()
    for i in range(input_dim):
        input_scaled[:, i] = scalers[feature_names[i]].inverse_transform(np.array([input_scaled[:, i]]))

    return input_scaled


def combine_features(cycles: list[dict], feature_config: dict) -> dict:
    combined_features = dict()
    for feature_set in feature_config["feature_sets"]:
        combined_features[feature_set] = dict()
        if feature_set not in cycles[0]:
            continue
        features_in_set = cycles[0][feature_set].keys()
        for feature in features_in_set:
            feature_values = [cycles[i][feature_set][feature] for i in range(len(cycles))]
            if isinstance(feature_values[0], dict):
                for key in feature_values[0].keys():
                    feature_array = np.concatenate([feature_values[i][key] for i in range(len(feature_values))])
                    combined_features[feature_set][key] = feature_array
            else:
                feature_array = np.concatenate(feature_values)
                combined_features[feature_set][feature] = feature_array

    return combined_features