"""Bead simulation sweep scaffolding for accuracy benchmarking."""

from __future__ import annotations

import argparse
import itertools
import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np

from benchmarks.speed import log_utils
from magtrack.simulation import simulate_beads

_DEFAULT_SWEEP_ROOT = Path(__file__).resolve().parent / "sweeps"


@dataclass
class ParameterSet:
    """Configuration describing a grid of simulation parameters."""

    name: str
    parameters: Mapping[str, Sequence[Any]]
    roi_bead_ratio: float = 6.0
    radius_nm: float = 1500.0
    nm_per_px_1x: float = 5000.0
    magnification: float = 40.0

    def combinations(self) -> list[dict[str, Any]]:
        """Return all parameter combinations expanded from the grid."""

        param_names = list(self.parameters.keys())
        values = [list(self.parameters[name]) for name in param_names]
        combos: list[dict[str, Any]] = []
        for combination in itertools.product(*values):
            combos.append(dict(zip(param_names, combination)))
        return combos


@dataclass
class SweepArtifact:
    """Paths and metadata generated by a simulation sweep."""

    sweep_name: str
    sweep_dir: Path
    images_path: Path
    metadata_path: Path
    metadata: dict[str, Any] = field(default_factory=dict)


class BeadSimulationSweep:
    """Generate bead simulation sweeps and metadata."""

    def __init__(
        self,
        sweep_name: str,
        parameter_sets: Sequence[ParameterSet] | None = None,
        *,
        sweep_root: Path | str = _DEFAULT_SWEEP_ROOT,
    ) -> None:
        self.sweep_name = self._sanitize_name(sweep_name)
        self.parameter_sets = list(parameter_sets or [default_parameter_set()])
        self.sweep_root = Path(sweep_root)

    def generate(self, overwrite: bool = False) -> SweepArtifact:
        """Create the sweep outputs on disk and return their locations."""

        sweep_dir = self.sweep_root / self.sweep_name
        sweep_dir.mkdir(parents=True, exist_ok=True)

        images_path = sweep_dir / "images.npz"
        metadata_path = sweep_dir / "metadata.json"
        if not overwrite and (images_path.exists() or metadata_path.exists()):
            raise FileExistsError(
                f"Sweep artifacts already exist in '{sweep_dir}'. Set overwrite=True to replace them."
            )

        images = {}
        parameter_metadata = []
        for param_set in self.parameter_sets:
            set_name = self._sanitize_name(param_set.name)
            combinations = param_set.combinations()
            combination_metadata = []
            for index, combo in enumerate(combinations):
                roi_bead_ratio = self._validate_positive(
                    float(combo.get("roi_bead_ratio", param_set.roi_bead_ratio)),
                    "roi_bead_ratio",
                )
                radius_nm = self._validate_positive(
                    float(combo.get("radius_nm", param_set.radius_nm)), "radius_nm"
                )
                nm_per_px_1x = self._validate_positive(
                    float(combo.get("nm_per_px_1x", param_set.nm_per_px_1x)),
                    "nm_per_px_1x",
                )
                magnification = self._validate_positive(
                    float(combo.get("magnification", param_set.magnification)),
                    "magnification",
                )
                nm_per_px = self._validate_positive(
                    nm_per_px_1x / magnification, "nm_per_px"
                )
                size_nm = self._validate_positive(
                    roi_bead_ratio * radius_nm, "size_nm"
                )
                size_px = self._validate_positive_int(size_nm / nm_per_px, "size_px")
                x_offset_nm = float(combo.get("x_offset", 0.0))
                y_offset_nm = float(combo.get("y_offset", 0.0))
                z_offset_nm = float(combo.get("z_offset", 0.0))

                xyz_nm = np.array(
                    [[x_offset_nm, y_offset_nm, z_offset_nm]],
                    dtype=np.float64,
                )
                center_px = size_px // 2
                center_nm = center_px * nm_per_px
                true_x_nm = center_nm + x_offset_nm
                true_y_nm = center_nm + y_offset_nm
                true_z_nm = z_offset_nm
                true_x_px = true_x_nm / nm_per_px
                true_y_px = true_y_nm / nm_per_px
                true_z_px = true_z_nm / nm_per_px
                image_key = f"{set_name}__{index:04d}"
                images[image_key] = simulate_beads(
                    xyz_nm,
                    size_px=size_px,
                    nm_per_px=nm_per_px,
                    radius_nm=radius_nm,
                    background_level=float(combo.get("background_level", 0.8)),
                )
                combination_metadata.append(
                    {
                        "key": image_key,
                        "values": self._ensure_json_serializable(combo),
                        "size_px": size_px,
                        "nm_per_px": nm_per_px,
                        "radius_nm": radius_nm,
                        "image_path": f"{images_path.name}::{image_key}",
                        "true_position_nm": {
                            "x": true_x_nm,
                            "y": true_y_nm,
                            "z": true_z_nm,
                        },
                        "true_position_px": {
                            "x": true_x_px,
                            "y": true_y_px,
                            "z": true_z_px,
                        },
                    }
                )
            parameter_metadata.append(
                {
                    "name": set_name,
                    "parameters": self._ensure_json_serializable({
                        key: list(values) for key, values in param_set.parameters.items()
                    }),
                    "combinations": combination_metadata,
                }
            )

        np.savez(images_path, **images)

        system_id, timestamp, system_metadata = log_utils.collect_system_metadata()
        created_at = datetime.now(timezone.utc).isoformat()
        metadata = {
            "sweep_name": self.sweep_name,
            "created_at": created_at,
            "image_file": images_path.name,
            "total_images": len(images),
            "parameter_sets": parameter_metadata,
            "system": {
                "id": system_id,
                "timestamp": timestamp,
                "metadata": system_metadata,
            },
        }

        metadata_path.write_text(json.dumps(metadata, indent=2))

        return SweepArtifact(
            sweep_name=self.sweep_name,
            sweep_dir=sweep_dir,
            images_path=images_path,
            metadata_path=metadata_path,
            metadata=metadata,
        )

    @staticmethod
    def _sanitize_name(name: str) -> str:
        safe = [c if c.isalnum() or c in {"-", "_"} else "-" for c in name]
        sanitized = "".join(safe).strip("-")
        return sanitized or "sweep"

    @staticmethod
    def _ensure_json_serializable(value: Any) -> Any:
        if value is None or isinstance(value, (str, int, float, bool)):
            return value
        if isinstance(value, dict):
            return {str(k): BeadSimulationSweep._ensure_json_serializable(v) for k, v in value.items()}
        if isinstance(value, (list, tuple, set)):
            return [BeadSimulationSweep._ensure_json_serializable(v) for v in value]
        if hasattr(value, "tolist"):
            try:
                return value.tolist()
            except Exception:  # noqa: BLE001
                pass
        try:
            return json.loads(json.dumps(value))
        except Exception:  # noqa: BLE001
            return repr(value)

    @staticmethod
    def _validate_positive(value: float, name: str) -> float:
        if value <= 0:
            raise ValueError(f"{name} must be positive; received {value}.")
        return value

    @staticmethod
    def _validate_positive_int(value: float, name: str) -> int:
        rounded = int(round(value))
        if rounded <= 0:
            raise ValueError(f"{name} must round to a positive integer; received {value}.")
        return rounded

    @staticmethod
    def default_parameter_set() -> ParameterSet:
        return default_parameter_set()


def default_parameter_set() -> ParameterSet:
    """Return the default parameter set."""

    return ParameterSet(
        name="default",
        parameters={
            "x_offset": [0, 1, 10, 100],
            "y_offset": [0, 1, 10, 100],
            "z_offset": [0, 1, 10, 100, 1000],
            "background_level": [0.4],
            "seed": [0],
            "roi_bead_ratio": [5],
            "radius_nm": [500, 1500],
            "nm_per_px_1x": [5000],
            "magnification": [10, 40, 100],
        },
    )


__all__ = ["BeadSimulationSweep", "ParameterSet", "SweepArtifact", "default_parameter_set"]


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate the default bead simulation sweep.")
    parser.add_argument(
        "--sweep-name",
        default="default",
        help="Name for the generated sweep directory (default: default)",
    )
    parser.add_argument(
        "--sweep-root",
        default=str(_DEFAULT_SWEEP_ROOT),
        help=f"Root directory for sweeps (default: {_DEFAULT_SWEEP_ROOT})",
    )
    parser.add_argument(
        "--overwrite",
        default=True,
        action="store_true",
        help="Overwrite existing sweep artifacts if they already exist.",
    )
    return parser.parse_args()


def _main() -> None:
    args = _parse_args()
    sweep = BeadSimulationSweep(args.sweep_name, sweep_root=Path(args.sweep_root))
    artifact = sweep.generate(overwrite=args.overwrite)
    print(f"Generated sweep '{artifact.sweep_name}' at {artifact.sweep_dir}")
    print(f"Images: {artifact.images_path}")
    print(f"Metadata: {artifact.metadata_path}")


if __name__ == "__main__":
    _main()
