from copy import deepcopy
from itertools import islice
from pathlib import Path

import numpy as np
import pandas as pd
from gluonts.dataset.arrow import ArrowWriter
from gluonts.itertools import Cyclic
from tqdm.auto import tqdm


def ts_sampler(datasets):
    infinite_iterators = list(map(lambda d: iter(Cyclic(d)), datasets))
    while True:
        idx = np.random.randint(len(datasets))
        yield next(infinite_iterators[idx])


def mean_scaler(ts_entry):
    ts_entry = deepcopy(ts_entry)
    mean = np.nanmean(ts_entry["target"])
    mean = mean if mean > 0.0 and mean == mean else 1.0
    ts_entry["target"] = ts_entry["target"] / mean
    return ts_entry


def slice_ts(ts_entry, length):
    ts_length = len(ts_entry["target"])
    start_idx = np.random.randint(ts_length - length + 1)
    return ts_entry["target"][start_idx : start_idx + length]


def validate_entry(entry):
    if np.isnan(entry["target"]).mean() >= 0.9:
        return False
    if (entry["target"] == 0.0).mean() >= 0.9:
        return False
    return True


def make_augmentation(
    ts_generator,
    K=3,
    alpha=1.5,
    min_length=128,
    max_length=1024,
):
    while True:
        k = np.random.randint(1, K + 1)
        length = np.random.randint(min_length, max_length + 1)
        tss = list(islice(filter(lambda x: len(x["target"]) >= length, ts_generator), k))
        tss = np.array([slice_ts(mean_scaler(ts), length) for ts in tss])
        weights = np.random.dirichlet([alpha] * k)
        ts = np.einsum("k, kl -> l", weights, tss)

        entry = {"start": np.datetime64("2000-01-01 00:00", "s"), "target": ts}
        if validate_entry(entry):
            return entry


def make_augmentations(
    datasets,
    num_series=100,
    K=3,
    alpha=1.5,
    min_length=128,
    max_length=2048,
):
    ts_generator = ts_sampler(datasets)
    for _ in tqdm(range(num_series)):
        yield make_augmentation(ts_generator, K=K, alpha=alpha, min_length=min_length, max_length=max_length)


def generate_dataset(num_series):
    """Generates a random GluonTS-style dataset"""
    dataset = []
    for _ in range(num_series):
        series_length = np.random.randint(10_000)
        dataset.append({"start": pd.Period("2024", freq="h"), "target": np.random.randn(series_length)})
    return dataset


# Usage Example

# Let's generate 10 random GluonTS-style datasets
datasets = [generate_dataset(num_series=np.random.randint(10_000)) for i in range(10)]
# Generate TXMixup augmentations
ts_mixed_up_dataset = make_augmentations(datasets, num_series=100_000)
# Save in Arrow format
ArrowWriter(compression="lz4").write_to_file(ts_mixed_up_dataset, Path("tsmixup-example.arrow"))