"""Définition centralisée des 5 domaines tokamak avec leurs samplers 2D et 3D.

Importer ce module dans les scripts de plot ou d'entraînement::

    from tokamak_domains import TOKAMAKS
    name, sampler2d, sampler3d = TOKAMAKS[0]
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

from scimba_jax.domains.tokamak import (
    TokamakSampler2D,
    TokamakSampler3DFrom2D,
    read_eqdsk_wall,
)

_DATA = (
    Path(__file__).parent.parent.parent.parent / "src/scimba_jax/domains/tokamak/data"
)

_EQDSK_FILES = {
    "MAST  g_p49320 (t=0.60s)": _DATA / "g_p49320_t0.60000",
    "MAST  g_p44849 (t=0.20s)": _DATA / "g_p44849_t0.20000",
    "MAST  eqdsk_mast_4": _DATA / "eqdsk_mast_4.dat",
    "MAST-U eqdsk_mast-u": _DATA / "eqdsk_mast-u.dat",
    "JET   eqdsk_jet_compare": _DATA / "eqdsk_jet_compare.dat",
}

# List of (name, sampler2d, sampler3d, R_wall, Z_wall)
TOKAMAKS: list[
    tuple[str, TokamakSampler2D, TokamakSampler3DFrom2D, np.ndarray, np.ndarray]
] = []

for _name, _path in _EQDSK_FILES.items():
    if not _path.exists():
        print(f"[SKIP] {_path.name} not found")
        continue
    _R_wall, _Z_wall = read_eqdsk_wall(str(_path))
    _s2d = TokamakSampler2D(_R_wall, _Z_wall, oversample=5)
    _s3d = TokamakSampler3DFrom2D(_R_wall, _Z_wall, oversample=5)
    TOKAMAKS.append((_name, _s2d, _s3d, _R_wall, _Z_wall))
    print(
        f"[OK]  {_name:<35}  wall={len(_R_wall)} pts  "
        f"R∈[{_R_wall.min():.2f}, {_R_wall.max():.2f}]"
    )

print(f"\n{len(TOKAMAKS)} tokamaks chargés.\n")