from __future__ import annotations

import sys
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as functional
from torch.utils.data import DataLoader, Dataset


def _import_era5_dataset(onescience_source_dir: str | None = None):
    if onescience_source_dir:
        source_dir = str(Path(onescience_source_dir).expanduser().resolve())
        if source_dir not in sys.path:
            sys.path.insert(0, source_dir)
    from onescience.datapipes.climate import ERA5Dataset

    return ERA5Dataset


class SpatialAdapter(Dataset):
    """Resize OneScience ERA5 samples only for reduced smoke profiles."""

    def __init__(self, dataset: Dataset, output_size: tuple[int, int]) -> None:
        self.dataset = dataset
        self.output_size = output_size

    def __len__(self) -> int:
        return len(self.dataset)

    def _resize(self, tensor: torch.Tensor) -> torch.Tensor:
        if tuple(tensor.shape[-2:]) == self.output_size:
            return tensor
        leading_shape = tensor.shape[:-2]
        resized = functional.interpolate(
            tensor.reshape(-1, 1, *tensor.shape[-2:]),
            size=self.output_size,
            mode="bilinear",
            align_corners=False,
        )
        return resized.reshape(*leading_shape, *self.output_size)

    def __getitem__(self, index: int):
        inputs, targets, cos_zenith, step_idx, time_index = self.dataset[index]
        inputs = self._resize(inputs)
        targets = self._resize(targets)
        cos_zenith = self._resize(cos_zenith)
        return inputs, targets, cos_zenith, step_idx, time_index


def build_dataset(
    config: dict[str, Any],
    years: list[int],
    *,
    output_steps: int = 1,
) -> Dataset:
    from common import active_model_config, resolve_path

    era5_dataset = _import_era5_dataset(config["project"].get("onescience_source_dir"))
    data_config = config["data"]
    dataset = era5_dataset(
        dataset_dir=str(resolve_path(config, data_config["dataset_dir"])),
        used_years=years,
        used_variables=data_config["variables"],
        input_steps=data_config["input_steps"],
        output_steps=output_steps,
        normalize=data_config["normalize"],
    )
    model_size = tuple(active_model_config(config)["img_size"])
    data_size = tuple(data_config["grid_shape"])
    if model_size != data_size:
        dataset = SpatialAdapter(dataset, model_size)
    return dataset


def build_loader(
    config: dict[str, Any],
    years: list[int],
    *,
    train: bool,
    distributed: bool,
    output_steps: int = 1,
) -> tuple[DataLoader, torch.utils.data.Sampler | None]:
    dataset = build_dataset(config, years, output_steps=output_steps)
    sampler = None
    if distributed:
        sampler = torch.utils.data.distributed.DistributedSampler(
            dataset, shuffle=train
        )
    loader = DataLoader(
        dataset,
        batch_size=config["training"]["batch_size"],
        shuffle=train and sampler is None,
        sampler=sampler,
        num_workers=config["training"]["num_workers"],
        pin_memory=True,
        drop_last=False,
    )
    return loader, sampler


def load_statistics(config: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
    import h5py
    import numpy as np

    from common import resolve_path

    data_config = config["data"]
    year = data_config["test_years"][0]
    path = resolve_path(config, data_config["dataset_dir"]) / "data" / f"{year}.h5"
    with h5py.File(path, "r") as handle:
        fields = handle["fields"]
        all_variables = [
            item.decode() if isinstance(item, bytes) else str(item)
            for item in fields.attrs["variables"]
        ]
        indices = [all_variables.index(name) for name in data_config["variables"]]
        if "global_means" in handle:
            means = handle["global_means"][:]
            stds = handle["global_stds"][:]
        else:
            stats_dir = path.parents[1] / "stats"
            means = np.load(stats_dir / "global_means.npy")
            stds = np.load(stats_dir / "global_stds.npy")
    return torch.from_numpy(means[:, indices]), torch.from_numpy(stds[:, indices])