from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

import numpy as np
import torch
from torch.utils.data import Dataset

from .io import (
    ensure_time_axis,
    load_image,
    percentile_normalize,
    read_jsonl,
    resolve_manifest_path,
)


@dataclass
class ManifestSample:
    image: Path
    label: Path
    meta: Path | None
    is_real: bool


def _parse_is_real(value: Any, default: bool) -> bool:
    if value is None:
        return bool(default)
    if isinstance(value, bool):
        return bool(value)
    if isinstance(value, (int, np.integer)):
        return bool(int(value) != 0)
    if isinstance(value, str):
        txt = value.strip().lower()
        if txt in {"1", "true", "yes", "y", "real"}:
            return True
        if txt in {"0", "false", "no", "n", "synth", "synthetic"}:
            return False
    return bool(default)


def apply_domain_style_augment(
    image: np.ndarray,
    rng: np.random.Generator,
    *,
    strength: float = 1.0,
) -> np.ndarray:
    """Domain-randomization style augmentation on normalized [0,1] kymographs."""
    out = np.asarray(image, dtype=np.float32).copy()
    if out.size == 0:
        return out

    s = float(max(0.0, strength))
    if s <= 0.0:
        return np.clip(out, 0.0, 1.0).astype(np.float32)

    h, w = int(out.shape[0]), int(out.shape[1])

    # Gamma-like camera response variation.
    if rng.random() < 0.95:
        gamma = float(rng.uniform(max(0.45, 1.0 - (0.35 * s)), min(2.20, 1.0 + (0.50 * s))))
        out = np.clip(out, 0.0, 1.0) ** gamma

    # Global gain/offset drift.
    if rng.random() < 0.90:
        gain = float(rng.uniform(max(0.35, 1.0 - (0.40 * s)), 1.0 + (0.40 * s)))
        bias = float(rng.uniform(-0.12 * s, 0.12 * s))
        out = (out * gain) + bias

    # Smooth illumination/background gradients.
    if rng.random() < 0.85:
        tt = np.linspace(-1.0, 1.0, h, dtype=np.float32)[:, None]
        xx = np.linspace(-1.0, 1.0, w, dtype=np.float32)[None, :]
        grad_t = float(rng.uniform(-0.20 * s, 0.20 * s))
        grad_x = float(rng.uniform(-0.20 * s, 0.20 * s))
        quad_x = float(rng.uniform(-0.08 * s, 0.08 * s))
        shading = 1.0 + (grad_t * tt) + (grad_x * xx) + (quad_x * (xx**2))
        shading = np.clip(shading, 0.35, 1.75)
        out = out * shading

    # Row/column fixed-pattern noise.
    if rng.random() < 0.70:
        row_sigma = float(rng.uniform(0.002, 0.018) * s)
        col_sigma = float(rng.uniform(0.001, 0.010) * s)
        if row_sigma > 0:
            out = out + rng.normal(0.0, row_sigma, size=(h, 1)).astype(np.float32)
        if col_sigma > 0:
            out = out + rng.normal(0.0, col_sigma, size=(1, w)).astype(np.float32)

    # Read noise.
    if rng.random() < 0.75:
        read_sigma = float(rng.uniform(0.003, 0.030) * s)
        out = out + rng.normal(0.0, read_sigma, size=out.shape).astype(np.float32)

    # Sparse hot pixels / cosmic spikes.
    if rng.random() < 0.30:
        n_hot = int(rng.integers(1, max(2, int((0.0015 * h * w * s)) + 2)))
        rr = rng.integers(0, h, size=n_hot)
        cc = rng.integers(0, w, size=n_hot)
        out[rr, cc] += rng.uniform(0.08, 0.35, size=n_hot).astype(np.float32)

    # Mild percentile remapping to vary dynamic range.
    if rng.random() < 0.55:
        q_lo = float(rng.uniform(0.2, 2.5))
        q_hi = float(rng.uniform(97.0, 99.9))
        lo = float(np.percentile(out, q_lo))
        hi = float(np.percentile(out, q_hi))
        if hi > (lo + 1e-6):
            out = (out - lo) / (hi - lo)

    return np.clip(out, 0.0, 1.0).astype(np.float32)


def load_manifest(manifest_path: str | Path, is_real: bool = False) -> list[ManifestSample]:
    manifest_path = Path(manifest_path)
    base_dir = manifest_path.resolve().parent
    rows = read_jsonl(manifest_path)
    samples: list[ManifestSample] = []
    for row in rows:
        image_path = resolve_manifest_path(base_dir, row["image"]) if "image" in row else None
        label_path = resolve_manifest_path(base_dir, row["label"]) if "label" in row else None
        meta_path = resolve_manifest_path(base_dir, row["meta"]) if row.get("meta") else None
        if image_path is None or label_path is None:
            raise ValueError(f"Manifest row missing image/label fields: {row}")
        row_is_real = _parse_is_real(row.get("is_real"), bool(is_real))
        samples.append(
            ManifestSample(
                image=Path(image_path),
                label=Path(label_path),
                meta=Path(meta_path) if meta_path else None,
                is_real=bool(row_is_real),
            )
        )
    return samples


class KymoSegmentationDataset(Dataset):
    def __init__(
        self,
        manifest_path: str | Path,
        *,
        is_real: bool = False,
        time_axis: str = "y",
        augment: bool = False,
        seed: int = 0,
        synth_style_aug_prob: float = 0.85,
        real_style_aug_prob: float = 0.25,
        style_aug_strength: float = 1.0,
    ):
        self.samples = load_manifest(manifest_path, is_real=is_real)
        self.time_axis = time_axis
        self.augment = augment
        self.rng = np.random.default_rng(seed)
        self.synth_style_aug_prob = float(np.clip(synth_style_aug_prob, 0.0, 1.0))
        self.real_style_aug_prob = float(np.clip(real_style_aug_prob, 0.0, 1.0))
        self.style_aug_strength = float(max(0.0, style_aug_strength))

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

    def __getitem__(self, idx: int) -> dict[str, Any]:
        sample = self.samples[idx]
        image = ensure_time_axis(load_image(sample.image), time_axis=self.time_axis)
        label = ensure_time_axis(load_image(sample.label), time_axis=self.time_axis)

        image_f = percentile_normalize(image)
        label_i = np.asarray(label, dtype=np.int64)

        if self.augment and self.rng.random() < 0.5:
            image_f = image_f[:, ::-1].copy()
            label_i = label_i[:, ::-1].copy()

        if self.augment:
            aug_p = self.real_style_aug_prob if sample.is_real else self.synth_style_aug_prob
            if aug_p > 0.0 and self.rng.random() < aug_p:
                image_f = apply_domain_style_augment(
                    image_f,
                    self.rng,
                    strength=self.style_aug_strength,
                )

        image_t = torch.from_numpy(image_f[None, ...]).float()
        label_t = torch.from_numpy(label_i).long()

        return {
            "image": image_t,
            "label": label_t,
            "is_real": torch.tensor(1 if sample.is_real else 0, dtype=torch.uint8),
            "image_path": str(sample.image),
            "label_path": str(sample.label),
        }


def compute_class_histogram(
    dataset,
    num_classes: int = 3,
    max_samples: int | None = None,
    *,
    progress_every: int = 0,
    progress_prefix: str = "class histogram",
    progress_fn: Callable[[str], None] | None = None,
) -> np.ndarray:
    n = len(dataset) if max_samples is None else min(len(dataset), max_samples)
    hist = np.zeros((num_classes,), dtype=np.float64)
    for i in range(n):
        item = dataset[i]
        y = item["label"].numpy().astype(np.int64)
        vals, counts = np.unique(y, return_counts=True)
        for v, c in zip(vals, counts):
            if 0 <= int(v) < num_classes:
                hist[int(v)] += float(c)
        if progress_every > 0 and ((i + 1) % int(progress_every) == 0 or (i + 1) == n):
            msg = f"[info] {progress_prefix}: {i + 1}/{n} ({(100.0 * (i + 1) / max(1, n)):.1f}%)"
            if progress_fn is None:
                print(msg, flush=True)
            else:
                progress_fn(msg)
    return hist
