from __future__ import annotations

import json
from typing import Any

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

from .datasets import apply_domain_style_augment, load_manifest
from .io import ensure_time_axis, load_image, percentile_normalize


def _draw_instance_disk(inst: np.ndarray, t: float, x: float, radius: int, track_id: int) -> None:
    t0 = max(0, int(np.floor(t)) - radius)
    t1 = min(inst.shape[0], int(np.ceil(t)) + radius + 1)
    x0 = max(0, int(np.floor(x)) - radius)
    x1 = min(inst.shape[1], int(np.ceil(x)) + radius + 1)
    if t0 >= t1 or x0 >= x1:
        return
    yy, xx = np.ogrid[t0:t1, x0:x1]
    if radius > 0:
        mask = (yy - t) ** 2 + (xx - x) ** 2 <= float(radius**2)
    else:
        mask = np.ones((t1 - t0, x1 - x0), dtype=bool)
    region = inst[t0:t1, x0:x1]
    region[np.logical_and(mask, region == 0)] = int(track_id)


def _draw_segment(inst: np.ndarray, t0: float, x0: float, t1: float, x1: float, radius: int, track_id: int) -> None:
    dt = float(t1 - t0)
    dx = float(x1 - x0)
    steps = int(max(abs(dt), abs(dx))) + 1
    if steps < 2:
        steps = 2
    for i in range(steps):
        a = float(i / max(1, steps - 1))
        tt = float(t0 + a * dt)
        xx = float(x0 + a * dx)
        _draw_instance_disk(inst, tt, xx, radius=radius, track_id=int(track_id))


def _dedupe_points(points: list[tuple[int, float]]) -> list[tuple[int, float]]:
    if not points:
        return []
    points = sorted(points, key=lambda z: z[0])
    out: list[tuple[int, float]] = []
    cur_t = points[0][0]
    xs: list[float] = []
    for t, x in points:
        if t != cur_t:
            out.append((int(cur_t), float(np.mean(xs))))
            cur_t = t
            xs = [x]
        else:
            xs.append(x)
    out.append((int(cur_t), float(np.mean(xs))))
    return out


def _track_is_moving(track: dict[str, Any]) -> bool:
    cls = str(track.get("class", "")).strip().lower()
    if not cls:
        return True
    return cls in {"moving", "diffusive"}


def _parse_track_points(track: dict[str, Any]) -> tuple[list[tuple[int, float]], int | None]:
    pts: list[tuple[int, float]] = []

    raw_path = track.get("path_tx")
    if isinstance(raw_path, list):
        for node in raw_path:
            if not isinstance(node, (list, tuple)) or len(node) < 2:
                continue
            try:
                t = int(round(float(node[0])))
                x = float(node[1])
            except Exception:
                continue
            if np.isfinite(x):
                pts.append((t, x))

    if not pts:
        anchors = track.get("anchors")
        if isinstance(anchors, list):
            for a in anchors:
                if isinstance(a, dict):
                    try:
                        t = int(round(float(a.get("t"))))
                        x = float(a.get("x"))
                    except Exception:
                        continue
                elif isinstance(a, (list, tuple)) and len(a) >= 2:
                    try:
                        t = int(round(float(a[0])))
                        x = float(a[1])
                    except Exception:
                        continue
                else:
                    continue
                if np.isfinite(x):
                    pts.append((t, x))

    if not pts:
        # Synth fallback for metadata that only stores line params.
        try:
            t0 = int(round(float(track.get("t0"))))
            t1 = int(round(float(track.get("t1"))))
            x0 = float(track.get("x0"))
            slope = float(track.get("slope", 0.0))
            if np.isfinite(x0) and np.isfinite(slope) and t1 >= t0:
                pts = [(int(t), float(x0 + slope * (t - t0))) for t in range(t0, t1 + 1)]
        except Exception:
            pts = []

    max_gap = track.get("max_interp_gap", None)
    if max_gap is None:
        max_gap_i = None
    else:
        try:
            max_gap_i = int(max_gap)
        except Exception:
            max_gap_i = None

    return _dedupe_points(pts), max_gap_i


def build_instance_map_from_meta(
    meta: dict[str, Any],
    *,
    shape: tuple[int, int],
    moving_only: bool = True,
    default_thickness: int = 3,
) -> np.ndarray:
    tdim, xdim = int(shape[0]), int(shape[1])
    inst = np.zeros((tdim, xdim), dtype=np.int64)
    tracks = meta.get("tracks", [])
    if not isinstance(tracks, list) or not tracks:
        return inst

    thickness = int(max(1, int(meta.get("label_thickness", default_thickness))))
    radius = max(0, thickness // 2)

    next_id = 1
    for tr in tracks:
        if not isinstance(tr, dict):
            continue
        if moving_only and not _track_is_moving(tr):
            continue

        points, max_gap = _parse_track_points(tr)
        if len(points) < 2:
            continue

        # Respect converter semantics:
        # max_interp_gap > 0 means skip bridging larger temporal gaps.
        for (t0, x0), (t1, x1) in zip(points[:-1], points[1:]):
            if t1 <= t0:
                continue
            if max_gap is not None and max_gap > 0 and (t1 - t0) > int(max_gap):
                continue
            _draw_segment(
                inst,
                float(np.clip(t0, 0, tdim - 1)),
                float(np.clip(x0, 0.0, xdim - 1.0)),
                float(np.clip(t1, 0, tdim - 1)),
                float(np.clip(x1, 0.0, xdim - 1.0)),
                radius=radius,
                track_id=int(next_id),
            )

        next_id += 1

    return inst


class KymoInstanceDataset(Dataset):
    def __init__(
        self,
        manifest_path: str,
        *,
        is_real: bool = False,
        time_axis: str = "y",
        augment: bool = False,
        seed: int = 0,
        moving_only_instances: bool = True,
        default_instance_thickness: int = 3,
        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.moving_only_instances = bool(moving_only_instances)
        self.default_instance_thickness = int(max(1, default_instance_thickness))
        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))
        self._meta_cache: dict[str, dict[str, Any] | None] = {}

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

    def _read_meta(self, path_str: str | None) -> dict[str, Any] | None:
        if not path_str:
            return None
        hit = self._meta_cache.get(path_str)
        if hit is not None or path_str in self._meta_cache:
            return hit
        try:
            with open(path_str, "r", encoding="utf-8") as f:
                obj = json.load(f)
            if not isinstance(obj, dict):
                obj = None
        except Exception:
            obj = None
        self._meta_cache[path_str] = obj
        return obj

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

        if label_raw.ndim != 2:
            raise ValueError(f"Expected 2D label image, got shape {label_raw.shape!r} for {sample.label}")

        meta_obj = self._read_meta(str(sample.meta) if sample.meta is not None else None)
        if meta_obj is None:
            inst_raw = np.zeros_like(label_raw, dtype=np.int64)
        else:
            inst_raw = build_instance_map_from_meta(
                meta_obj,
                shape=(int(label_raw.shape[0]), int(label_raw.shape[1])),
                moving_only=bool(self.moving_only_instances),
                default_thickness=int(self.default_instance_thickness),
            )

        image = ensure_time_axis(image_raw, time_axis=self.time_axis)
        label = ensure_time_axis(label_raw, time_axis=self.time_axis)
        instance = ensure_time_axis(inst_raw, time_axis=self.time_axis)

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

        if self.augment and self.rng.random() < 0.5:
            image_f = image_f[:, ::-1].copy()
            label_i = label_i[:, ::-1].copy()
            instance_i = instance_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,
                )

        has_instances = int(np.count_nonzero(instance_i) > 0)

        return {
            "image": torch.from_numpy(image_f[None, ...]).float(),
            "label": torch.from_numpy(label_i).long(),
            "instance": torch.from_numpy(instance_i).long(),
            "has_instances": torch.tensor(has_instances, dtype=torch.uint8),
            "is_real": torch.tensor(1 if sample.is_real else 0, dtype=torch.uint8),
            "image_path": str(sample.image),
            "label_path": str(sample.label),
            "meta_path": str(sample.meta) if sample.meta is not None else "",
        }
