import copy
import logging
import os
import random

import lmdb
import numpy as np
from bson import ObjectId
from torch.utils.data import IterableDataset
from tqdm import tqdm

from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features
from utils.dataset_utils import load_key_stats
from utils.lmdb_utils import load_from_lmdb
from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id
from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection

logger = logging.getLogger(__name__)


def get_prepared_sequence_length(input_sequence: np.ndarray, model_configuration: dict) -> int:
    """
    Returns the length of the input sequence after padding and resampling
    """

    initial_length = len(input_sequence)
    take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
    padding_length = get_padding_length(model_configuration, resampled=False)
    raw_length = initial_length + padding_length
    adjusted = int(raw_length // take_every_nth)
    return adjusted


def get_padding_length(model_configuration: dict, resampled: bool = True) -> int:
    """
    Returns the length of the input sequence after padding and resampling
    """
    take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
    padding_length = int(model_configuration["input_window_length"] * (
            1 - model_configuration["preprocessing"]["min_input_length_fraction_for_padding"])) * (
                         take_every_nth if not resampled else 1)
    return padding_length


def get_number_of_windows(base_length: int, model_configuration: dict) -> int:
    """
    Returns the number of windows for a given sequence length
    """
    window_shift = model_configuration["preprocessing"]["window_shift"]
    input_window_length = model_configuration["input_window_length"]
    output_window_length = model_configuration["output_window_length"]
    output_window_offset = model_configuration["output_window_offset"]

    if input_window_length > output_window_length + output_window_offset:
        return (base_length - input_window_length) // window_shift + 1
    else:
        return (base_length - output_window_length - output_window_offset) // window_shift + 1


def get_collated_batch_for_key(sample_key: ObjectId | str,
                               model_configuration: dict,
                               start_cutoff: int = None,
                               end_cutoff: int = None,
                               lmdb_env=None) -> dict:
    sample_batch_for_key = get_batch_for_key(sample_key, model_configuration,
                                             start_cutoff=start_cutoff,
                                             end_cutoff=end_cutoff,
                                             lmdb_env=lmdb_env)
    collated = model_configuration["collate_fn"](sample_batch_for_key)
    return collated


def get_batch_for_key(key,
                      model_configuration: dict,
                      start_cutoff: int = None,
                      end_cutoff: int = None,
                      lmdb_env=None) -> np.ndarray:
    """
    Get the batch for a given key from the lmdb database or compute it directly.

    The return has the same format as the batch produced by the batch_fn in the model configuration.
    Args:
        key: key to get the batch for, can also be object id from database
        model_configuration: model configuration to use
        start_cutoff: start cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points
        end_cutoff: end cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points
        lmdb_env: the lmdb environment to use, if None, the features are computed directly from the database

    Returns:
        batch: batch for the given key, as returned by the batch_fn in the model configuration

    """
    if lmdb_env is None:
        # compute and scale features, ignored features are handled internally
        features = get_scaled_feature_for_key(key, model_configuration)
    else:
        # load features from lmdb
        features = load_from_lmdb(lmdb_env, str(key))
        # ignored features are handled internally, thus we need to remove them here
        ignored_features = model_configuration["feature_config"]["ignored_features"] if "ignored_features" in \
                                                                                        model_configuration[
                                                                                            "feature_config"] else []
        # add scaled versions of features
        ignored_features = ignored_features + [feat + "_scaled" for feat in ignored_features]
        for feature_set in features:
            for feature_name in list(features[feature_set].keys()):
                if feature_name in ignored_features:
                    del features[feature_set][feature_name]

    processed_chunk = process_chunk([key], [features], model_configuration)

    # use cutoff, if provided
    if start_cutoff is not None or end_cutoff is not None:

        item_length = len(processed_chunk[0]["target_features"])

        if start_cutoff is None:
            start_cutoff = 0
        if end_cutoff is None:
            end_cutoff = item_length

        # "normalize" cutoff to adjust for added padding and resampling
        padding_length = get_padding_length(model_configuration, resampled=False)
        take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
        # make sure that there is no padding added to the start cutoff, so that it includes the padding added to the sequence
        start_cutoff_normalized = max(int(start_cutoff // take_every_nth), 0)
        # end cutoff must be adjusted to include the padding added to the sequence
        end_cutoff_normalized = min(int((end_cutoff + padding_length) // take_every_nth), item_length)
        for feature_set in processed_chunk[0]:
            if feature_set in model_configuration["feature_config"]["feature_sets"]:
                processed_chunk[0][feature_set] = processed_chunk[0][feature_set][
                                                  start_cutoff_normalized:end_cutoff_normalized + 1]

    if "batch_fn" in model_configuration and model_configuration["batch_fn"] is not None:
        batch = model_configuration["batch_fn"](processed_chunk, model_configuration)
    else:
        batch = processed_chunk
    if batch is None or len(batch) == 0:
        raise ValueError(f"Batch is empty for key {key}")
    return batch


def get_scaled_feature_for_key(key: str,
                               model_configuration: dict) -> tuple:
    feature_config = model_configuration["feature_config"]

    user_cycles = list(
        get_cycles_collection().find({"user_id": ObjectId(key)} | feature_config["filter_criteria"]).sort("starts_at",
                                                                                                          1))
    cycle_features = list()
    for cycle in user_cycles:
        features = get_features(cycle, feature_config)
        cycle_features.append(features)

    # combine all cycles for the user
    if len(cycle_features) == 0:
        raise ValueError(f"No features found for key {key}")
    features = combine_features(cycle_features, feature_config)
    scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
    scalers = load_scalers(scaler_dir)
    # scale features
    scaled_features = scale_item(features, scalers)
    return scaled_features


def augment_items(identifiers: list,
                  items: list[dict],
                  feature_config: dict,
                  lmdb_env: lmdb.Environment = None) -> list:
    """
    Augments the given sequences by attaching previous sequences
    Args:
        identifiers: identifiers of items for finding previous sequences
        items: actual items
        feature_config: feature config of the dataset
        lmdb_env: lmdb environment for loading previous sequences, can be None, in this case the items are computed directly

    Returns:
        list of augmented items
    """

    if "augmentation" not in feature_config or \
            feature_config["augmentation"]["use_augmentation"] is False:
        return items

    augmented_items = list()
    for i, item in enumerate(items):
        item_id = identifiers[i]
        try:
            item_data = get_cycles_collection().find_one({"_id": ObjectId(item_id)}, {"user_id": 1, "starts_at": 1})
            item_user_id = item_data["user_id"]
            item_starts_at = item_data["starts_at"]
        except KeyError:
            print(f"User ID not found for item {item_id}")
            continue
        previous_item_identifiers = list(get_cycles_collection().aggregate(
            # make sure to filter by user first to significantly reduce the number of items
            [
                {
                    "$match": {
                        "user_id": item_user_id,
                        "starts_at": {"$lt": item_starts_at},
                    }
                }
            ] + feature_config["filter_criteria_pipeline"] + [
                {
                    "$sort": {
                        "starts_at": -1
                    }
                },
                {
                    "$project": {
                        "_id": 1,
                        "starts_at": 1
                    }
                }
            ]
        ))
        max_lookback = feature_config["augmentation"]["max_lookback"]
        previous_items = list()
        for previous_item in previous_item_identifiers:
            if len(previous_items) >= max_lookback:
                break
            previous_item_id = previous_item["_id"]
            if lmdb_env is None:
                # compute item
                cycle = get_cycle_by_id(previous_item_id)
                features = get_features(cycle, feature_config)
                scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
                scalers = load_scalers(scaler_dir)
                previous_item = scale_item(features, scalers)
                previous_items.append(previous_item)
            else:
                # take item from lmdb
                try:
                    # make sure to parse object id
                    previous_item = load_from_lmdb(lmdb_env, str(previous_item_id))
                except KeyError:
                    continue
                previous_items.append(previous_item)

        # merge items into one
        feature_sets = feature_config["feature_sets"]
        augmented_item = copy.deepcopy(item)
        for feature_set in feature_sets:
            if feature_set not in augmented_item:
                continue
            if "static" in feature_set:
                continue
            for feature_name in augmented_item[feature_set]:
                if feature_name not in augmented_item[feature_set]:
                    continue
                augmented_item[feature_set][feature_name] = np.concatenate(
                    [previous_item[feature_set][feature_name] for previous_item in previous_items]
                    + [augmented_item[feature_set][feature_name]])

        augmented_items.append(augmented_item)

    return augmented_items


def process_chunk(ids: list,
                  chunk: list,
                  model_configuration: dict,
                  ignored_features: list = None,
                  pad_sequences: bool = True,
                  statics_as_list: bool = True):
    """
    Creates a chunk of data as dataframe for training and inference of a tft model.
    :param ids: list of ids of the data points in the chunk
    :param chunk: list of data points

    :return: dataframe with the data points in the chunk
    """

    if ignored_features is None:
        ignored_features = list()

    input_features = list(chunk[0]["target_features"].keys())
    feature_config = model_configuration["feature_config"]

    # variables for padding
    padding_value = 0

    # create dataframe from items in chunk
    data = []
    for i, item in enumerate(chunk):
        item_length = item["target_features"][input_features[0]].shape[0]
        take_every_nth = model_configuration["preprocessing"]["take_every_nth"]

        data_item = dict()
        for feature_set in feature_config["feature_sets"]:
            if feature_set in item and len(item[feature_set]) > 0:
                # fill in the data
                source = item[feature_set]
                features_to_use = [feature for feature in source.keys() if "_scaled" in feature]
                feature_set_data = [source[feature] for feature in features_to_use]

                padding_length = get_padding_length(model_configuration, resampled=False)

                if len(feature_set_data) == 0:
                    raise ValueError(f"No features found for key {ids[i]} in feature set {feature_set}")

                if len(feature_set_data[0]) == 1:
                    # if the feature is constant, we need to repeat it for all time points
                    if statics_as_list:
                        if pad_sequences:
                            desired_length = item_length + padding_length
                            feature_set_data = np.array(
                                [np.full((desired_length,), val) for val in feature_set_data]).T[::take_every_nth]
                        else:
                            feature_set_data = np.array(
                                [np.full((item_length,), val) for val in feature_set_data]).T[::take_every_nth]
                    else:
                        feature_set_data = np.array(feature_set_data).flatten()
                else:
                    if pad_sequences:
                        padding_values = np.full((padding_length, len(feature_set_data)),
                                                 [padding_value for vals in feature_set_data]).T
                        feature_set_data = np.concatenate([padding_values, np.array(feature_set_data)], axis=1)
                        # create bins of n values and apply mean

                    # check, if special accumulation function has been specified
                    individual_feature_data = list()
                    for i, feature_to_use in enumerate(features_to_use):
                        try:
                            # get config for feature, make sure to replace scaled to get actual config
                            current_feature_config = \
                                [x for x in feature_config[feature_set] if
                                 x["name"] == feature_to_use.replace("_scaled", "")][0]
                            current_accumulation_fn = current_feature_config["accumulation_fn"]
                        except (KeyError, IndexError):
                            # default to mean
                            current_accumulation_fn = np.mean

                        current_feature_data = apply_fn_to_bins(feature_set_data[i], take_every_nth,
                                                                current_accumulation_fn)
                        individual_feature_data.append(current_feature_data)

                    # stack individual features
                    feature_set_data = np.stack(individual_feature_data).T

                if feature_set not in data_item:
                    data_item[feature_set] = feature_set_data
                else:
                    data_item[feature_set] = np.concatenate([data_item[feature_set], feature_set_data], axis=1)

        data.append(data_item)

    return data


def apply_fn_to_bins(input_sequence: np.ndarray,
                     bin_size: int,
                     fn: callable) -> np.ndarray:
    """
    Applies a function to the bins of the input sequence along the last dimension and returns the results.
    Args:
        input_sequence: sequence to apply the function to (1D or 2D)
        bin_size: size of the bins
        fn: function to apply to the bins

    Returns:
        np.ndarray: result of the function applied to the bins

    """
    if input_sequence.ndim == 1:
        input_sequence = input_sequence.reshape(1, -1)
        dims = 1
    else:
        dims = input_sequence.shape[-1]

    result = []
    for row in input_sequence:
        row_result = []
        for start in range(0, len(row), bin_size):
            end = min(start + bin_size, len(row))
            bin_slice = row[start:end]
            row_result.append(fn(bin_slice))
        result.append(row_result)

    if dims > 1:
        return np.array(result)
    else:
        return np.array(result[0])


def produce_window_batches(data_chunk: list,
                           model_configuration: dict,
                           offsets: list | np.ndarray = None) -> list:
    """
    Produces batches of windows from the data chunk.
    Args:
        data_chunk: data chunk
        model_configuration: configuration to use
        offsets: offsets used to determine the length of data to use, uses item[-offset:] of data

    Returns:
        data: list of windows
    """
    input_window_length = model_configuration["input_window_length"]
    output_window_length = model_configuration["output_window_length"]
    output_window_offset = model_configuration["output_window_offset"]
    window_shift = model_configuration["preprocessing"]["window_shift"]

    if offsets is None:
        offsets = np.zeros(len(data_chunk), dtype=int)

    data = list()
    for i, item in enumerate(data_chunk):
        item_length = item["target_features"][-offsets[i]:].shape[0]
        total_window_length = input_window_length if input_window_length > output_window_length + output_window_offset \
            else output_window_length + output_window_offset

        if item_length < total_window_length:
            continue

        # number of windows is defined by the input window length and output window length with offset
        num_windows = get_number_of_windows(item_length, model_configuration)
        for j in range(num_windows):
            window = dict()
            for key, value in item.items():
                if value is None:
                    window[key] = None
                else:
                    if key == "target_features":
                        window[key] = value[-offsets[i]:][j * window_shift + output_window_offset:
                                                          j * window_shift + output_window_length + output_window_offset]
                    else:
                        window[key] = value[-offsets[i]:][j * window_shift:j * window_shift + input_window_length]

            data.append(window)

    return data


def produce_simple_batches(data_chunk: list,
                           model_configuration: dict,
                           offsets: list | np.ndarray = None) -> list:
    """
    Produces batches without windowing.
    Args:
         data_chunk: data chunk
         model_configuration: configuration to use
         offsets: offsets used to determine the length of data to use, uses item[-offset:] of data

    Returns:
           data: list of items
    """
    pass


class LMDBIterableDataset(IterableDataset):
    def __init__(self,
                 lmdb_env_path: str,
                 lmdb_keys: list[str],
                 model_configuration: dict,
                 batch_size: int = 32):
        self.lmdb_path = lmdb_env_path
        self.lmdb_env = None
        self.lmdb_keys = lmdb_keys
        self.key_subset = None
        self.model_configuration = model_configuration
        self.batch_size = batch_size
        self.keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"])
        self.len = None

        try:
            self.model_configuration["batch_fn"]
        except KeyError:
            raise KeyError("Batch function not found in model configuration")

    def set_key_subset(self, key_subset: list[str]):
        """
        Set the key subset to use for the dataset.
        Args:
            key_subset: list of keys to use
        """
        self.key_subset = key_subset
        # reset length
        self.len = None
        # reset lmdb env
        if self.lmdb_env is not None:
            self.lmdb_env.close()
            self.lmdb_env = None

    def get_length_of_data_subset(self, key_set: list[str]):
        """
        Get the length of the data subset.
        Args:
            key_set: list of keys to use
        """
        # run simplified version of __iter__ to get the length

        num_steps = 0
        i = 0
        num_batch_items = 0
        while True:
            if num_batch_items >= self.batch_size:
                num_steps += 1
                num_batch_items -= self.batch_size
            else:
                if i >= len(key_set):
                    if num_batch_items > 0:
                        num_steps += 1
                    break
                key = key_set[i]
                i += 1

                # try to fetch stats from key_stats
                try:
                    current_key_stats = self.keys_stats["by_key"][key]
                    base_length = current_key_stats["item_length"]
                    take_every_nth = self.model_configuration["preprocessing"]["take_every_nth"]
                    padding_length = get_padding_length(self.model_configuration, resampled=False)
                    raw_length = base_length + padding_length
                    item_length = int(raw_length // take_every_nth)
                except:
                    item = load_from_lmdb(self.lmdb_env, key)
                    item_length = get_prepared_sequence_length(
                        item["target_features"][list(item["target_features"])[0]],
                        self.model_configuration)
                num_windows = get_number_of_windows(item_length, self.model_configuration)
                num_batch_items += num_windows

        return num_steps

    def __iter__(self):

        # if no key subset is set, use all keys
        if self.key_subset is None:
            self.key_subset = self.lmdb_keys

        self.init_lmdb_env()

        random.shuffle(self.key_subset)

        batch = list()
        counter = 0
        collate_fn = self.model_configuration["collate_fn"]
        while True:
            if len(batch) >= self.batch_size:
                if collate_fn is None:
                    yield batch[:self.batch_size]
                else:
                    yield collate_fn(batch[:self.batch_size])
                batch = batch[self.batch_size:]
            else:
                if counter >= len(self.key_subset):
                    if len(batch) > 0:
                        yield collate_fn(batch)
                    break
                key = self.lmdb_keys[counter]
                counter += 1
                try:
                    current_batch = get_batch_for_key(key, self.model_configuration, lmdb_env=self.lmdb_env)
                except ValueError:
                    # if betch is empty, try next
                    continue
                batch.extend(current_batch)

    def init_lmdb_env(self):
        if self.lmdb_env is None:
            self.lmdb_env = lmdb.open(self.lmdb_path,
                                      readonly=True,
                                      lock=False,
                                      readahead=False,
                                      meminit=False)

    def __len__(self):
        # if no key subset is set, use all keys
        if self.key_subset is None:
            self.key_subset = self.lmdb_keys

        self.init_lmdb_env()

        if self.len is None:
            num_steps = self.get_length_of_data_subset(self.key_subset)
            self.len = num_steps

        return self.len
