#!/usr/bin/env python3
"""Seeded equivalence fuzzer for the Python-driven OpenFOAM stepper."""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import random
import shutil
import subprocess
import sys
from dataclasses import asdict, dataclass
from pathlib import Path


def _drop_ambient_pythonpath() -> None:
    pythonpath = os.environ.pop("PYTHONPATH", "")
    for entry in pythonpath.split(os.pathsep):
        if not entry:
            continue
        while entry in sys.path:
            sys.path.remove(entry)


_drop_ambient_pythonpath()

import numpy as np


def _load_openfoam_env_helpers():
    try:
        from openfoam_env import apply_openfoam_env, openfoam_env
    except ModuleNotFoundError:
        spec = importlib.util.spec_from_file_location("openfoam_env", Path(__file__).with_name("openfoam_env.py"))
        if spec is None or spec.loader is None:
            raise
        module = importlib.util.module_from_spec(spec)
        sys.modules["openfoam_env"] = module
        spec.loader.exec_module(module)
        return module.apply_openfoam_env, module.openfoam_env
    return apply_openfoam_env, openfoam_env


apply_openfoam_env, openfoam_env = _load_openfoam_env_helpers()

ROOT = Path(__file__).resolve().parents[1]
TUTORIAL = ROOT / "OpenFOAM-14/tutorials/incompressibleFluid/venturiTube"
DEFAULT_WORK = ROOT / "tmp/python_stepper_fuzz"
DEFAULT_SEEDS = 10
DEFAULT_ATOL = 1e-10
DEFAULT_RTOL = 1e-10
FIELDS = ("U", "p", "phi")


@dataclass(frozen=True)
class CaseVariant:
    seed: int
    diameter: float
    dia_cells: int
    ven_cells: int
    in_cells: int
    out_cells: int
    box_cells: int
    rad_cells: int
    out_grading: float
    u_inlet: float
    nu: float
    p_relax: float
    u_relax: float
    p_tolerance: float
    u_tolerance: float
    n_non_orthogonal_correctors: int


def scalar(value: float) -> str:
    return f"{value:.17g}"


def variant_for_seed(seed: int) -> CaseVariant:
    rng = random.Random(seed)
    return CaseVariant(
        seed=seed,
        diameter=rng.choice([0.05, 0.075, 0.1, 0.15, 0.2]),
        dia_cells=rng.choice([4, 6, 8]),
        ven_cells=rng.choice([2, 4]),
        in_cells=rng.choice([6, 8, 10]),
        out_cells=rng.choice([8, 10, 12]),
        box_cells=rng.choice([2, 3, 4]),
        rad_cells=rng.choice([4, 6, 8]),
        out_grading=rng.choice([0.5, 0.75, 1.0, 1.25]),
        u_inlet=rng.choice([0.05, 0.1, 0.2, 0.4]),
        nu=rng.choice([2e-5, 4e-5, 8e-5, 1.6e-4]),
        p_relax=rng.choice([0.2, 0.3, 0.5, 0.7]),
        u_relax=rng.choice([0.5, 0.7, 0.9]),
        p_tolerance=rng.choice([1e-6, 1e-7]),
        u_tolerance=rng.choice([1e-7, 1e-8]),
        n_non_orthogonal_correctors=rng.choice([0, 1]),
    )


def replace_required(path: Path, text: str, old: str, new: str) -> str:
    count = text.count(old)
    if count == 0:
        raise AssertionError(f"{path}: missing literal {old!r}")
    if count != 1:
        raise AssertionError(f"{path}: expected one literal {old!r}, found {count}")
    return text.replace(old, new, 1)


def patch_file(path: Path, replacements: list[tuple[str, str]]) -> None:
    text = path.read_text()
    for old, new in replacements:
        text = replace_required(path, text, old, new)
    path.write_text(text)


def patch_block_mesh(case: Path, variant: CaseVariant) -> None:
    patch_file(
        case / "system/blockMeshDict",
        [
            ("diameter   0.1;", f"diameter   {scalar(variant.diameter)};"),
            ("diaCells   16;", f"diaCells   {variant.dia_cells};"),
            ("venCells    8;", f"venCells   {variant.ven_cells};"),
            ("inCells    20;", f"inCells   {variant.in_cells};"),
            ("outCells   40;", f"outCells   {variant.out_cells};"),
            ("boxCells   8;", f"boxCells   {variant.box_cells};"),
            ("radCells   16;", f"radCells   {variant.rad_cells};"),
            ("outGrading 0.5;", f"outGrading {scalar(variant.out_grading)};"),
        ],
    )


def patch_initial_fields(case: Path, variant: CaseVariant) -> None:
    patch_file(case / "0/U", [("Uinlet          0.2;", f"Uinlet          {scalar(variant.u_inlet)};")])


def patch_physical_properties(case: Path, variant: CaseVariant) -> None:
    patch_file(case / "constant/physicalProperties", [("nu      4e-05;", f"nu      {scalar(variant.nu)};")])


def patch_fv_solution(case: Path, variant: CaseVariant) -> None:
    patch_file(
        case / "system/fvSolution",
        [
            ("        tolerance       1e-6;", f"        tolerance       {scalar(variant.p_tolerance)};"),
            ("        tolerance       1e-7;", f"        tolerance       {scalar(variant.u_tolerance)};"),
            (
                "    nNonOrthogonalCorrectors 0;",
                f"    nNonOrthogonalCorrectors {variant.n_non_orthogonal_correctors};",
            ),
            ("        p               0.3;", f"        p               {scalar(variant.p_relax)};"),
            ("        U               0.7;", f"        U               {scalar(variant.u_relax)};"),
        ],
    )


def patch_control_dict_for_seed(case: Path) -> None:
    patch_file(
        case / "system/controlDict",
        [
            ("startFrom       latestTime;", "startFrom       startTime;"),
            ("endTime         1000;", "endTime         1;"),
            ("writeInterval   50;", "writeInterval   1;"),
            ("writePrecision   8;", "writePrecision   17;"),
        ],
    )


def patch_control_dict_to_latest(case: Path) -> None:
    patch_file(case / "system/controlDict", [("startFrom       startTime;", "startFrom       latestTime;")])


def run(cmd: list[str], *, log_path: Path | None = None) -> None:
    env = openfoam_env()
    if log_path is None:
        subprocess.run(cmd, cwd=ROOT, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT, env=env)
        return

    with log_path.open("w") as log:
        subprocess.run(cmd, cwd=ROOT, check=True, stdout=log, stderr=subprocess.STDOUT, env=env)


def prepare_case(dst: Path, variant: CaseVariant) -> None:
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(TUTORIAL, dst, ignore=shutil.ignore_patterns("processor*", "postProcessing", "*.log"))
    for orig in (dst / "0").glob("*.orig"):
        shutil.copyfile(orig, orig.with_suffix(""))

    patch_block_mesh(dst, variant)
    patch_initial_fields(dst, variant)
    patch_physical_properties(dst, variant)
    patch_fv_solution(dst, variant)
    patch_control_dict_for_seed(dst)

    run(["blockMesh", "-case", str(dst)])
    run(["createZones", "-case", str(dst)])


def prepare_seed(seed_dir: Path, variant: CaseVariant) -> dict[str, Path]:
    if seed_dir.exists():
        shutil.rmtree(seed_dir)
    outputs = seed_dir / "outputs"
    outputs.mkdir(parents=True)
    (seed_dir / "variant.json").write_text(json.dumps(asdict(variant), indent=2, sort_keys=True) + "\n")

    cases = {
        "foam": seed_dir / "foam_case",
        "run_one": seed_dir / "python_run_one_case",
        "run_split": seed_dir / "python_split_case",
    }
    for case in cases.values():
        prepare_case(case, variant)
    return cases


def save_npz(out: Path, *, U: np.ndarray, p: np.ndarray, phi: np.ndarray) -> None:
    out.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(out, U=np.asarray(U), p=np.asarray(p), phi=np.asarray(phi))


def child_load_fields(case: Path, out: Path) -> None:
    import foam_stepper as foam

    fields = foam.Case(case).make_stepper().fields()
    save_npz(out, U=fields.U.internal, p=fields.p.internal, phi=fields.phi.internal)


def child_run_one(case: Path, out: Path) -> None:
    import foam_stepper as foam

    result = foam.Case(case).make_stepper().run_one_pimple_iteration()
    fields = result.outputs["fields"]
    save_npz(out, U=fields["U"].internal, p=fields["p"].internal, phi=fields["phi"].internal)


def child_run_split(case: Path, out: Path) -> None:
    import foam_stepper as foam

    stepper = foam.Case(case).make_stepper()
    stepper.pre_solve()
    stepper.advance_time()
    begin = stepper.begin_pimple_iteration()
    assert begin.outputs["active"] is True
    stepper.fv_models_correct()
    stepper.pre_predictor()
    stepper.momentum_transport_predictor()
    stepper.assemble_momentum_matrix()
    stepper.relax_matrix()
    stepper.constrain_matrix()
    stepper.solve_momentum()
    stepper.compute_pressure_inputs()
    stepper.assemble_pressure_matrix()
    stepper.solve_pressure()
    stepper.correct_velocity_pressure_flux()
    stepper.momentum_transport_corrector()
    stepper.end_pimple_iteration()
    stepper.post_solve(write=False)
    fields = stepper.fields()
    save_npz(out, U=fields.U.internal, p=fields.p.internal, phi=fields.phi.internal)


def load_npz(path: Path) -> dict[str, np.ndarray]:
    with np.load(path) as data:
        return {field: data[field] for field in FIELDS}


def failure_message(
    *,
    seed: int,
    variant: CaseVariant,
    seed_dir: Path,
    path_label: str,
    field: str,
    baseline_path: Path,
    actual_path: Path,
    expected_shape: tuple[int, ...],
    actual_shape: tuple[int, ...],
    max_abs: float,
    max_rel: float,
) -> str:
    return (
        f"fuzz equivalence failed seed={seed} path={path_label} field={field}\n"
        f"variant={json.dumps(asdict(variant), sort_keys=True)}\n"
        f"variant_json={seed_dir / 'variant.json'}\n"
        f"foamRun_log={seed_dir / 'foamRun.log'}\n"
        f"baseline_output={baseline_path}\n"
        f"actual_output={actual_path}\n"
        f"shape_pair=actual{actual_shape} expected{expected_shape}\n"
        f"max_abs={max_abs:.17g}\n"
        f"max_rel={max_rel:.17g}"
    )


def compare_field(
    *,
    seed: int,
    variant: CaseVariant,
    seed_dir: Path,
    path_label: str,
    field: str,
    expected: np.ndarray,
    actual: np.ndarray,
    baseline_path: Path,
    actual_path: Path,
    atol: float,
    rtol: float,
) -> float:
    if actual.shape != expected.shape:
        raise AssertionError(
            failure_message(
                seed=seed,
                variant=variant,
                seed_dir=seed_dir,
                path_label=path_label,
                field=field,
                baseline_path=baseline_path,
                actual_path=actual_path,
                expected_shape=expected.shape,
                actual_shape=actual.shape,
                max_abs=float("nan"),
                max_rel=float("nan"),
            )
        )

    abs_diff = np.abs(actual - expected)
    max_abs = float(np.max(abs_diff)) if abs_diff.size else 0.0
    max_rel = float(np.max(abs_diff / np.maximum(np.abs(expected), atol))) if abs_diff.size else 0.0
    if not np.allclose(actual, expected, rtol=rtol, atol=atol):
        raise AssertionError(
            failure_message(
                seed=seed,
                variant=variant,
                seed_dir=seed_dir,
                path_label=path_label,
                field=field,
                baseline_path=baseline_path,
                actual_path=actual_path,
                expected_shape=expected.shape,
                actual_shape=actual.shape,
                max_abs=max_abs,
                max_rel=max_rel,
            )
        )
    return max_abs


def compare_outputs(
    *,
    seed: int,
    variant: CaseVariant,
    seed_dir: Path,
    baseline_path: Path,
    actual_paths: dict[str, Path],
    atol: float,
    rtol: float,
) -> dict[str, dict[str, float]]:
    baseline = load_npz(baseline_path)
    summary: dict[str, dict[str, float]] = {}
    for path_label, actual_path in actual_paths.items():
        actual = load_npz(actual_path)
        summary[path_label] = {}
        for field in FIELDS:
            summary[path_label][field] = compare_field(
                seed=seed,
                variant=variant,
                seed_dir=seed_dir,
                path_label=path_label,
                field=field,
                expected=baseline[field],
                actual=actual[field],
                baseline_path=baseline_path,
                actual_path=actual_path,
                atol=atol,
                rtol=rtol,
            )
    return summary


def run_child(child: str, case: Path, out: Path) -> None:
    subprocess.run(
        [sys.executable, str(Path(__file__).resolve()), "--child", child, "--case", str(case), "--out", str(out)],
        cwd=ROOT,
        check=True,
        env=openfoam_env(),
    )


def run_seed(seed: int, variant: CaseVariant, work_dir: Path, *, atol: float, rtol: float) -> Path:
    seed_dir = work_dir / f"seed_{seed}"
    cases = prepare_seed(seed_dir, variant)
    outputs = seed_dir / "outputs"
    baseline_path = outputs / "baseline.npz"
    run_one_path = outputs / "run_one.npz"
    run_split_path = outputs / "run_split.npz"

    run(
        ["foamRun", "-case", str(cases["foam"]), "-solver", "incompressibleFluid", "-noFunctionObjects"],
        log_path=seed_dir / "foamRun.log",
    )
    patch_control_dict_to_latest(cases["foam"])
    run_child("load-fields", cases["foam"], baseline_path)
    run_child("run-one", cases["run_one"], run_one_path)
    run_child("run-split", cases["run_split"], run_split_path)

    summary = compare_outputs(
        seed=seed,
        variant=variant,
        seed_dir=seed_dir,
        baseline_path=baseline_path,
        actual_paths={"run_one": run_one_path, "run_split": run_split_path},
        atol=atol,
        rtol=rtol,
    )
    cells = int(load_npz(baseline_path)["U"].shape[0])
    print(
        f"seed={seed} cells={cells} "
        f"run_one: U={summary['run_one']['U']:.3e} p={summary['run_one']['p']:.3e} phi={summary['run_one']['phi']:.3e} "
        f"split: U={summary['run_split']['U']:.3e} p={summary['run_split']['p']:.3e} phi={summary['run_split']['phi']:.3e}",
        flush=True,
    )
    return seed_dir


def run_parent(args: argparse.Namespace) -> None:
    if args.seeds < 1:
        raise SystemExit("--seeds must be at least 1")

    args.work_dir.mkdir(parents=True, exist_ok=True)
    for seed in range(args.seed_start, args.seed_start + args.seeds):
        variant = variant_for_seed(seed)
        seed_dir = run_seed(seed, variant, args.work_dir, atol=args.atol, rtol=args.rtol)
        if not args.keep_passing:
            shutil.rmtree(seed_dir)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seeds", type=int, default=DEFAULT_SEEDS)
    parser.add_argument("--seed-start", type=int, default=0)
    parser.add_argument("--work-dir", type=Path, default=DEFAULT_WORK)
    parser.add_argument("--atol", type=float, default=DEFAULT_ATOL)
    parser.add_argument("--rtol", type=float, default=DEFAULT_RTOL)
    parser.add_argument("--keep-passing", action="store_true")
    parser.add_argument("--child", choices=("load-fields", "run-one", "run-split"))
    parser.add_argument("--case", type=Path)
    parser.add_argument("--out", type=Path)
    args = parser.parse_args()

    if args.child is not None and (args.case is None or args.out is None):
        parser.error("--child requires --case and --out")
    if args.child is None and (args.case is not None or args.out is not None):
        parser.error("--case and --out are only valid with --child")
    return args


def main() -> None:
    args = parse_args()
    if args.child == "load-fields":
        apply_openfoam_env()
        child_load_fields(args.case, args.out)
    elif args.child == "run-one":
        apply_openfoam_env()
        child_run_one(args.case, args.out)
    elif args.child == "run-split":
        apply_openfoam_env()
        child_run_split(args.case, args.out)
    else:
        run_parent(args)


if __name__ == "__main__":
    main()