#!/usr/bin/env python3
from __future__ import annotations

import argparse
import math
import random
import re
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List, Tuple


@dataclass
class CellCase:
    ls: int
    name: str
    a: float
    b: float
    c: float
    alpha: float
    beta: float
    gamma: float


@dataclass
class FitResult:
    a: float
    b: float
    c: float
    alpha: float
    beta: float
    gamma: float


def deg2rad(x: float) -> float:
    return x * math.pi / 180.0


def cell_metric(cell: CellCase) -> List[List[float]]:
    a, b, c = cell.a, cell.b, cell.c
    ca = math.cos(deg2rad(cell.alpha))
    cb = math.cos(deg2rad(cell.beta))
    cg = math.cos(deg2rad(cell.gamma))
    return [
        [a * a, a * b * cg, a * c * cb],
        [a * b * cg, b * b, b * c * ca],
        [a * c * cb, b * c * ca, c * c],
    ]


def inv3(m: List[List[float]]) -> List[List[float]]:
    a = m
    det = (
        a[0][0] * (a[1][1] * a[2][2] - a[1][2] * a[2][1])
        - a[0][1] * (a[1][0] * a[2][2] - a[1][2] * a[2][0])
        + a[0][2] * (a[1][0] * a[2][1] - a[1][1] * a[2][0])
    )
    if abs(det) < 1e-14:
        raise ValueError("Singular metric tensor.")
    inv = [[0.0] * 3 for _ in range(3)]
    inv[0][0] = (a[1][1] * a[2][2] - a[1][2] * a[2][1]) / det
    inv[0][1] = (a[0][2] * a[2][1] - a[0][1] * a[2][2]) / det
    inv[0][2] = (a[0][1] * a[1][2] - a[0][2] * a[1][1]) / det
    inv[1][0] = (a[1][2] * a[2][0] - a[1][0] * a[2][2]) / det
    inv[1][1] = (a[0][0] * a[2][2] - a[0][2] * a[2][0]) / det
    inv[1][2] = (a[0][2] * a[1][0] - a[0][0] * a[1][2]) / det
    inv[2][0] = (a[1][0] * a[2][1] - a[1][1] * a[2][0]) / det
    inv[2][1] = (a[0][1] * a[2][0] - a[0][0] * a[2][1]) / det
    inv[2][2] = (a[0][0] * a[1][1] - a[0][1] * a[1][0]) / det
    return inv


def d_spacing(cell: CellCase, h: int, k: int, l: int) -> float:
    g = cell_metric(cell)
    g_star = inv3(g)
    v = [h, k, l]
    dinv2 = 0.0
    for i in range(3):
        for j in range(3):
            dinv2 += v[i] * g_star[i][j] * v[j]
    if dinv2 <= 0.0:
        raise ValueError(f"Non-positive 1/d^2 for hkl=({h},{k},{l})")
    return 1.0 / math.sqrt(dinv2)


def two_theta_from_d(d: float, wavelength: float) -> float | None:
    x = wavelength / (2.0 * d)
    if not (0.0 < x < 1.0):
        return None
    theta = math.asin(x)
    return 2.0 * math.degrees(theta)


def generate_reflections(
    cell: CellCase,
    wavelength: float,
    n_lines: int = 15,
    noise_sigma_deg: float = 0.02,
    hmax: int = 6,
    twotheta_min: float = 10.0,
    twotheta_max: float = 140.0,
    seed: int = 1234,
) -> List[Tuple[int, int, int, float, float]]:
    rng = random.Random(seed)
    peaks: List[Tuple[int, int, int, float]] = []
    seen = set()

    for h in range(0, hmax + 1):
        for k in range(0, hmax + 1):
            for l in range(0, hmax + 1):
                if h == 0 and k == 0 and l == 0:
                    continue
                d = d_spacing(cell, h, k, l)
                tt = two_theta_from_d(d, wavelength)
                if tt is None:
                    continue
                if not (twotheta_min <= tt <= twotheta_max):
                    continue
                key = round(tt, 3)
                if key in seen:
                    continue
                seen.add(key)
                peaks.append((h, k, l, tt))

    peaks.sort(key=lambda x: x[3])
    if len(peaks) < n_lines:
        raise RuntimeError(f"{cell.name}: not enough peaks ({len(peaks)})")

    chosen = peaks[:n_lines]
    noisy: List[Tuple[int, int, int, float, float]] = []
    for h, k, l, tt in chosen:
        tt_noisy = tt + rng.gauss(0.0, noise_sigma_deg)
        noisy.append((h, k, l, tt, tt_noisy))
    return noisy


def build_input_text(
    title: str,
    ls: int,
    wavelength: float,
    peaks: List[Tuple[int, int, int, float, float]],
    ip: int = 1,
) -> str:
    lines = [
        title,
        f"{ls} 0 0 0 0 0 2 4 {ip}",
        f"{wavelength:.6f} 0.0",
    ]
    for h, k, l, _tt, tt_noisy in peaks:
        lines.append(f"{h:d} {k:d} {l:d} {tt_noisy:.6f}")
    lines.append("1000 0 0 0")
    lines.append("0")
    return "\n".join(lines) + "\n"


DIRECT_CELL_BLOCK_RE = re.compile(
    r"Direct cell constant\s+"
    r".*?a.*?b.*?c\s+"
    r"([0-9Ee+\-.]+)\(\s*[0-9Ee+\-.]+\)\s+"
    r"([0-9Ee+\-.]+)\(\s*[0-9Ee+\-.]+\)\s+"
    r"([0-9Ee+\-.]+)\(\s*[0-9Ee+\-.]+\)\s+"
    r".*?alpha.*?beta.*?gamma\s+"
    r"([0-9Ee+\-.]+)\(\s*[0-9Ee+\-.]+\)\s+"
    r"([0-9Ee+\-.]+)\(\s*[0-9Ee+\-.]+\)\s+"
    r"([0-9Ee+\-.]+)\(\s*[0-9Ee+\-.]+\)",
    re.S,
)


def parse_output_file(outfile: Path) -> FitResult:
    text = outfile.read_text(encoding="utf-8", errors="replace")
    m = DIRECT_CELL_BLOCK_RE.search(text)
    if not m:
        raise RuntimeError(f"Could not parse Direct cell constant block: {outfile}")
    vals = [float(x) for x in m.groups()]
    return FitResult(*vals)


def all_cases() -> List[CellCase]:
    return [
        CellCase(1, "triclinic",    4.83, 5.27, 6.11, 72.0, 83.0, 76.0),
        CellCase(2, "monoclinic_b", 5.12, 6.35, 7.21, 90.0, 104.0, 90.0),
        CellCase(3, "monoclinic_c", 5.74, 6.22, 7.48, 90.0, 90.0, 111.0),
        CellCase(4, "orthorhombic", 4.91, 5.36, 6.02, 90.0, 90.0, 90.0),
        CellCase(5, "tetragonal",   4.14, 4.14, 6.57, 90.0, 90.0, 90.0),
        CellCase(6, "cubic",        4.07, 4.07, 4.07, 90.0, 90.0, 90.0),
        CellCase(7, "trigonal",     5.21, 5.21, 5.21, 75.0, 75.0, 75.0),
        CellCase(8, "hexagonal",    3.25, 3.25, 5.18, 90.0, 90.0, 120.0),
    ]


def cell_to_dict(cell: CellCase | FitResult) -> dict:
    return {
        "a": cell.a,
        "b": cell.b,
        "c": cell.c,
        "alpha": cell.alpha,
        "beta": cell.beta,
        "gamma": cell.gamma,
    }


def judge_diffs(expected: CellCase, got: FitResult) -> tuple[bool, dict, dict]:
    exp = cell_to_dict(expected)
    out = cell_to_dict(got)
    abs_err = {k: abs(out[k] - exp[k]) for k in exp}
    rel_err = {}
    for k in exp:
        denom = abs(exp[k])
        rel_err[k] = abs_err[k] / denom if denom > 1e-12 else 0.0

    ok = (
        abs_err["a"] < 0.05 and
        abs_err["b"] < 0.05 and
        abs_err["c"] < 0.05 and
        abs_err["alpha"] < 1.0 and
        abs_err["beta"] < 1.0 and
        abs_err["gamma"] < 1.0
    )
    return ok, abs_err, rel_err


def infer_issue(expected: CellCase, got: FitResult, abs_err: dict, rel_err: dict) -> str:
    if expected.name == "trigonal":
        ratio_a = got.a / expected.a if abs(expected.a) > 1e-12 else float("nan")
        if abs(ratio_a - math.sqrt(3.0)) < 0.02 and abs_err["alpha"] < 0.1:
            return (
                "length parameters only: output a,b,c are about sqrt(3) times expected while angles are correct; "
                "this strongly suggests the LS=7 post-processing interprets P0 as 3*a*^2 and divides by 3 once too many or too few."
            )
    worst = max(abs_err, key=abs_err.get)
    if worst in {"a", "b", "c"}:
        return f"length mismatch is dominant, especially {worst}"
    return f"angle mismatch is dominant, especially {worst}"


def format_case_report(expected: CellCase, got: FitResult, abs_err: dict, rel_err: dict, issue: str) -> str:
    lines = []
    lines.append(f"  expected: a={expected.a:.5f}, b={expected.b:.5f}, c={expected.c:.5f}, alpha={expected.alpha:.3f}, beta={expected.beta:.3f}, gamma={expected.gamma:.3f}")
    lines.append(f"  output  : a={got.a:.5f}, b={got.b:.5f}, c={got.c:.5f}, alpha={got.alpha:.3f}, beta={got.beta:.3f}, gamma={got.gamma:.3f}")
    lines.append("  errors  : " + ", ".join(
        f"{k}={abs_err[k]:.5f} (rel={100.0*rel_err[k]:.2f}%)" for k in ["a", "b", "c", "alpha", "beta", "gamma"]
    ))
    lines.append(f"  issue   : {issue}")
    return "\n".join(lines)


def run_one_case(
    lsq_script: Path,
    workdir: Path,
    cell: CellCase,
    wavelength: float,
    n_lines: int,
    noise_sigma_deg: float,
    seed: int,
) -> tuple[bool, FitResult | None, dict | None, dict | None, str]:
    case_dir = workdir / f"LS{cell.ls}_{cell.name}"
    case_dir.mkdir(parents=True, exist_ok=True)

    peaks = generate_reflections(
        cell=cell,
        wavelength=wavelength,
        n_lines=n_lines,
        noise_sigma_deg=noise_sigma_deg,
        seed=seed + cell.ls,
    )

    infile = case_dir / "input.BZ.K"
    outfile = case_dir / "output.txt"
    infile.write_text(
        build_input_text(
            title=f"TEST {cell.name}",
            ls=cell.ls,
            wavelength=wavelength,
            peaks=peaks,
            ip=1,
        ),
        encoding="utf-8",
    )

    cmd = [sys.executable, str(lsq_script), str(infile), str(outfile), "--silent"]
    proc = subprocess.run(
        cmd,
        input="",
        text=True,
        capture_output=True,
        cwd=case_dir,
    )

    if proc.returncode != 0:
        msg = f"execution failed: rc={proc.returncode}\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}"
        return False, None, None, None, msg

    if not outfile.exists():
        msg = f"output file not created\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}"
        return False, None, None, None, msg

    got = parse_output_file(outfile)
    ok, abs_err, rel_err = judge_diffs(cell, got)
    issue = infer_issue(cell, got, abs_err, rel_err)
    return ok, got, abs_err, rel_err, issue


def main() -> int:
    parser = argparse.ArgumentParser(description="lsq_latt2.py full-system test with verbose diagnostics")
    parser.add_argument("lsq_script", nargs="?", default="lsq_latt2.py", help="path to lsq_latt2.py")
    parser.add_argument("--wavelength", type=float, default=1.5406)
    parser.add_argument("--n-lines", type=int, default=15)
    parser.add_argument("--noise-sigma", type=float, default=0.02)
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument("--keep-workdir", action="store_true")
    args = parser.parse_args()

    src = Path(args.lsq_script).resolve()
    if not src.exists():
        print(f"ERROR: script not found: {src}", file=sys.stderr)
        return 2

    if args.keep_workdir:
        tmp_root = Path(tempfile.mkdtemp(prefix="test_lsq_latt2_"))
        cleanup = False
    else:
        tmpctx = tempfile.TemporaryDirectory(prefix="test_lsq_latt2_")
        tmp_root = Path(tmpctx.name)
        cleanup = True

    try:
        work_script = tmp_root / "lsq_latt2_under_test.py"
        work_script.write_text(src.read_text(encoding="utf-8"), encoding="utf-8")

        results = []
        print(f"Work directory : {tmp_root}")
        print(f"Target script  : {src}")
        print()

        for cell in all_cases():
            ok, got, abs_err, rel_err, issue = run_one_case(
                lsq_script=work_script,
                workdir=tmp_root,
                cell=cell,
                wavelength=args.wavelength,
                n_lines=args.n_lines,
                noise_sigma_deg=args.noise_sigma,
                seed=args.seed,
            )
            results.append((cell, ok, got, abs_err, rel_err, issue))
            if got is None:
                print(f"[FAIL] LS={cell.ls} {cell.name}")
                print("  issue   : " + issue)
            else:
                print(f"[{'PASS' if ok else 'FAIL'}] LS={cell.ls} {cell.name}")
                print(format_case_report(cell, got, abs_err, rel_err, issue))
            print()

        n_pass = sum(1 for _c, ok, _g, _a, _r, _i in results if ok)
        n_total = len(results)
        print(f"Summary: {n_pass}/{n_total} passed")

        failed = [(c, g, a, r, i) for c, ok, g, a, r, i in results if not ok]
        if failed:
            print("\nFailed cases:")
            for cell, got, abs_err, rel_err, issue in failed:
                print(f"- LS={cell.ls} {cell.name}")
                if got is None:
                    print(f"  issue   : {issue}")
                else:
                    print(format_case_report(cell, got, abs_err, rel_err, issue))
                print()
            return 1

        if args.keep_workdir:
            print(f"Temporary files kept at: {tmp_root}")
        return 0
    finally:
        if cleanup:
            tmpctx.cleanup()


if __name__ == "__main__":
    raise SystemExit(main())