from __future__ import annotations

import json
from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np

from .io import RasterData, save_geotiff


def create_run_folder(root: Path) -> Path:
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = root / f"run_{stamp}"
    run_dir.mkdir(parents=True, exist_ok=True)
    return run_dir


def save_run_config(run_dir: Path, config: dict[str, Any]) -> None:
    path = run_dir / "config.json"
    path.write_text(json.dumps(config, indent=2), encoding="utf-8")


def save_segmentation_outputs(
    run_dir: Path,
    reference: RasterData | None,
    labels: np.ndarray,
    uncertainty: np.ndarray,
    probabilities: np.ndarray | None,
) -> None:
    out_dir = run_dir / "outputs"
    out_dir.mkdir(parents=True, exist_ok=True)

    if reference:
        labels_path = out_dir / "labels.tif"
        save_geotiff(labels_path, labels[None, :, :].astype("int32"), reference)

        unc_path = out_dir / "uncertainty.tif"
        save_geotiff(unc_path, uncertainty[None, :, :].astype("float32"), reference)

        if probabilities is not None:
            prob_path = out_dir / "probabilities.tif"
            save_geotiff(prob_path, probabilities.astype("float32"), reference)
    else:
        np.save(out_dir / "labels.npy", labels)
        np.save(out_dir / "uncertainty.npy", uncertainty)
        if probabilities is not None:
            np.save(out_dir / "probabilities.npy", probabilities)
