# ruff: noqa: D103
"""Paths for entropice data storage."""

import datetime
import os
from pathlib import Path
from typing import Literal

from entropice.utils.types import Grid, Model, TargetDataset, Task, TemporalMode

DATA_DIR = (
    Path(os.environ.get("FAST_DATA_DIR", None) or os.environ.get("DATA_DIR", None) or "data").resolve() / "entropice"
)
DATA_DIR = Path("/raid/scratch/tohoel001/data/entropice")  # Temporary hardcoding for FAST cluster

GRIDS_DIR = DATA_DIR / "grids"
FIGURES_DIR = Path("figures")
DARTS_V1_DIR = DATA_DIR / "darts-v1"
DARTS_V2_DIR = DATA_DIR / "darts-v2"
DARTS_MLLABELS_DIR = DATA_DIR / "darts-mllabels"
ERA5_DIR = DATA_DIR / "era5"
ARCTICDEM_DIR = DATA_DIR / "arcticdem"
EMBEDDINGS_DIR = DATA_DIR / "embeddings"
WATERMASK_DIR = DATA_DIR / "watermask"
TRAINING_DIR = DATA_DIR / "training"
DATASET_ENSEMBLES_DIR = DATA_DIR / "dataset_ensembles"
RESULTS_DIR = DATA_DIR / "results"

GRIDS_DIR.mkdir(parents=True, exist_ok=True)
FIGURES_DIR.mkdir(parents=True, exist_ok=True)
DARTS_V1_DIR.mkdir(parents=True, exist_ok=True)
DARTS_V2_DIR.mkdir(parents=True, exist_ok=True)
DARTS_MLLABELS_DIR.mkdir(parents=True, exist_ok=True)
ERA5_DIR.mkdir(parents=True, exist_ok=True)
ARCTICDEM_DIR.mkdir(parents=True, exist_ok=True)
EMBEDDINGS_DIR.mkdir(parents=True, exist_ok=True)
WATERMASK_DIR.mkdir(parents=True, exist_ok=True)
TRAINING_DIR.mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
DATASET_ENSEMBLES_DIR.mkdir(parents=True, exist_ok=True)


watermask_file = WATERMASK_DIR / "simplified_water_polygons.shp"


def _get_gridname(grid: Grid, level: int) -> str:
    return f"permafrost_{grid}{level}"


def get_grid_file(grid: Grid, level: int) -> Path:
    gridname = _get_gridname(grid, level)
    gridfile = GRIDS_DIR / f"{gridname}_grid.parquet"
    return gridfile


def get_grid_viz_file(grid: Grid, level: int) -> Path:
    gridname = _get_gridname(grid, level)
    vizfile = FIGURES_DIR / f"{gridname}_grid.png"
    return vizfile


def get_darts_file(grid: Grid, level: int, version: Literal["v1", "v1-l3", "v2", "v2-l3", "mllabels"]) -> Path:
    gridname = _get_gridname(grid, level)
    match version:
        case "v1" | "v1-l3":
            darts_file = DARTS_V1_DIR / f"{gridname}_darts-{version}.zarr"
        case "v2" | "v2-l3":
            darts_file = DARTS_V2_DIR / f"{gridname}_darts-{version}.zarr"
        case "mllabels":
            darts_file = DARTS_MLLABELS_DIR / f"{gridname}_darts-{version}.zarr"
        case _:
            raise ValueError(f"Unknown DARTS version: {version}")
    return darts_file


def get_annual_embeddings_file(grid: Grid, level: int, year: int) -> Path:
    gridname = _get_gridname(grid, level)
    embfile = EMBEDDINGS_DIR / f"{gridname}_embeddings-{year}.parquet"
    return embfile


def get_embeddings_store(grid: Grid, level: int, temporal: TemporalMode | None = None) -> Path:
    gridname = _get_gridname(grid, level)
    if temporal == "synopsis":
        gridname += "_synopsis"
    embstore = EMBEDDINGS_DIR / f"{gridname}_embeddings.zarr"
    return embstore


def get_era5_stores(
    agg: Literal["daily", "monthly", "summer", "winter", "yearly", "seasonal", "shoulder"] = "daily",
    grid: Grid | None = None,
    level: int | None = None,
    temporal: TemporalMode | None = None,
) -> Path:
    pdir = ERA5_DIR
    if temporal == "synopsis":
        agg += "_synopsis"  # ty:ignore[invalid-assignment]
    fname = f"{agg}_climate.zarr"

    if grid is None or level is None:
        pdir = pdir / "intermediate"
        pdir.mkdir(parents=True, exist_ok=True)
    else:
        gridname = _get_gridname(grid, level)
        fname = f"{gridname}_{fname}"
    aligned_path = pdir / fname
    return aligned_path


def get_arcticdem_stores(
    grid: Grid | None = None,
    level: int | None = None,
) -> Path:
    if grid is None or level is None:
        return DATA_DIR / "arcticdem32m.icechunk.zarr"
    gridname = _get_gridname(grid, level)
    aligned_path = ARCTICDEM_DIR / f"{gridname}_dem.zarr"
    return aligned_path


def get_train_dataset_file(grid: Grid, level: int) -> Path:
    gridname = _get_gridname(grid, level)
    dataset_file = TRAINING_DIR / f"{gridname}_train_dataset.parquet"
    return dataset_file


def get_dataset_cache(eid: str, subset: str | None = None, batch: tuple[int, int] | None = None) -> Path:
    if batch is not None:
        eid = f"{eid}_batch{batch[0]}-{batch[1]}"
    if subset is None:
        cache_file = DATASET_ENSEMBLES_DIR / f"{eid}_dataset.parquet"
    else:
        cache_file = DATASET_ENSEMBLES_DIR / f"{eid}_{subset}_dataset.parquet"
    return cache_file


def get_features_cache(ensemble_id: str, cells_hash: str) -> Path:
    cache_dir = DATASET_ENSEMBLES_DIR / "cache" / "features" / ensemble_id
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir / f"cells_{cells_hash}.parquet"


def get_dataset_stats_cache() -> Path:
    cache_dir = DATASET_ENSEMBLES_DIR / "cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir / "dataset_stats.pckl"


def get_training_results_dir(
    experiment: str | None,
    name: str,
    grid: Grid,
    level: int,
    task: Task,
    target: TargetDataset,
    model_type: Model | None = None,
) -> Path:
    gridname = _get_gridname(grid, level)
    now = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    if experiment is not None:
        experiment_dir = RESULTS_DIR / experiment
        experiment_dir.mkdir(parents=True, exist_ok=True)
    else:
        experiment_dir = RESULTS_DIR
    parts = [gridname, name, now, task, target]
    if model_type is not None:
        parts.append(model_type)
    results_dir = experiment_dir / "_".join(parts)
    results_dir.mkdir(parents=True, exist_ok=True)
    return results_dir
