import numpy as np

from qmrpy.models.t2 import T2DECAESMap


def _load_csv_1d(path: str) -> np.ndarray:
    # writedlm output is a single row or column; np.loadtxt handles both
    x = np.loadtxt(path, delimiter=",")
    return np.atleast_1d(x).astype(np.float64)


def test_decaes_parity_fixed_flip_no_reg() -> None:
    # Reference generated by Julia DECAES.jl (tests/data/*.csv)
    t2times = _load_csv_1d("tests/data/decaes_ref_t2times.csv") * 1000.0
    echotimes = _load_csv_1d("tests/data/decaes_ref_echotimes.csv") * 1000.0
    alpha_ref = float(_load_csv_1d("tests/data/decaes_ref_alpha.csv")[0])
    dist_ref = _load_csv_1d("tests/data/decaes_ref_dist.csv")

    m = T2DECAESMap(
        n_te=len(echotimes),
        te_ms=float(echotimes[0]),
        n_t2=len(t2times),
        t2_range_ms=(float(t2times[0]), float(t2times[-1])),
        reg="none",
        set_flip_angle_deg=180.0,
        epg_backend="decaes",
    )

    signal = _load_csv_1d("tests/data/decaes_ref_signal.csv")
    assert signal.shape == (len(echotimes),)

    out = m.fit(signal)

    assert np.isclose(out["params"]["alpha_deg"], alpha_ref, rtol=0, atol=1e-9)
    assert np.allclose(out["params"]["t2times_ms"], t2times, rtol=0, atol=1e-9)
    assert np.allclose(out["params"]["distribution"], dist_ref, rtol=1e-8, atol=1e-10)


def test_decaes_parity_opt_alpha_gcv() -> None:
    # Reference generated by Julia DECAES.jl with Reg="gcv" and flip-angle optimization.
    alpha_ref = float(_load_csv_1d("tests/data/decaes_ref2_alpha.csv")[0])
    mu_ref = float(_load_csv_1d("tests/data/decaes_ref2_mu.csv")[0])
    dist_ref = _load_csv_1d("tests/data/decaes_ref2_dist.csv")
    signal = _load_csv_1d("tests/data/decaes_ref_signal.csv")

    m = T2DECAESMap(
        n_te=16,
        te_ms=10.0,
        n_t2=30,
        t2_range_ms=(10.0, 2000.0),
        reg="gcv",
        n_ref_angles=64,
        save_reg_param=True,
        epg_backend="decaes",
    )
    out = m.fit(signal)

    assert np.isclose(out["params"]["alpha_deg"], alpha_ref, rtol=0, atol=1e-6)
    assert np.isclose(out["diagnostics"]["mu"], mu_ref, rtol=1e-5, atol=0)
    assert np.allclose(out["params"]["distribution"], dist_ref, rtol=1e-5, atol=1e-7)


def _assert_reg_parity(reg: str, *, chi2_factor: float | None = None, noise_level: float | None = None) -> None:
    alpha_ref = float(_load_csv_1d(f"tests/data/decaes_ref_{reg}_alpha.csv")[0])
    mu_ref = float(_load_csv_1d(f"tests/data/decaes_ref_{reg}_mu.csv")[0])
    chi2_ref = float(_load_csv_1d(f"tests/data/decaes_ref_{reg}_chi2factor.csv")[0])
    dist_ref = _load_csv_1d(f"tests/data/decaes_ref_{reg}_dist.csv")
    signal = _load_csv_1d("tests/data/decaes_ref_signal.csv")

    m = T2DECAESMap(
        n_te=16,
        te_ms=10.0,
        n_t2=30,
        t2_range_ms=(10.0, 2000.0),
        reg=reg,
        n_ref_angles=64,
        save_reg_param=True,
        epg_backend="decaes",
        chi2_factor=chi2_factor,
        noise_level=noise_level,
    )
    out = m.fit(signal)

    assert np.isclose(out["params"]["alpha_deg"], alpha_ref, rtol=0, atol=1e-6)
    assert np.isclose(out["diagnostics"]["mu"], mu_ref, rtol=1e-5, atol=0)
    assert np.isclose(out["diagnostics"]["chi2factor"], chi2_ref, rtol=1e-5, atol=0)
    assert np.allclose(out["params"]["distribution"], dist_ref, rtol=1e-5, atol=1e-7)


def test_decaes_parity_opt_alpha_lcurve() -> None:
    _assert_reg_parity("lcurve")


def test_decaes_parity_opt_alpha_chi2() -> None:
    _assert_reg_parity("chi2", chi2_factor=1.02)


def test_decaes_parity_opt_alpha_mdp() -> None:
    # DECAES reference uses NoiseLevel_raw = 1e-3 * max(signal), which corresponds to noise_level=1e-3
    # in qmrpy because we normalize by max(signal) before calling the MDP search.
    _assert_reg_parity("mdp", noise_level=1e-3)
