#!/usr/bin/env python3
"""Export sparse A/b in OpticsFEM COO text format (matches C++ Test_OutputMatrix)."""

from __future__ import annotations

from pathlib import Path

import numpy as np
from scipy import sparse


def export_ab_coo(a: sparse.spmatrix, b: np.ndarray, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    a = a.tocoo()
    order = np.lexsort((a.col, a.row))
    ai = a.row[order].astype(np.int64)
    aj = a.col[order].astype(np.int64)
    av = a.data[order]

    np.savetxt(out_dir / "Ai.txt", ai, fmt="%d")
    np.savetxt(out_dir / "Aj.txt", aj, fmt="%d")
    with open(out_dir / "Av.txt", "w", encoding="utf-8") as f:
        for v in av:
            f.write(f"({v.real:.12g},{v.imag:.12g})\n")

    b = np.asarray(b).reshape(-1)
    np.savetxt(out_dir / "Bv_real.txt", b.real, fmt="%.12g")
    np.savetxt(out_dir / "Bv_imag.txt", b.imag, fmt="%.12g")

    meta = out_dir / "meta.txt"
    meta.write_text(
        "\n".join(
            [
                "source=export_ab_coo.py",
                f"n={a.shape[0]}",
                f"nnz={ai.size}",
                f"|b|={np.linalg.norm(b):.15g}",
            ]
        )
        + "\n",
        encoding="utf-8",
    )
    print(f"Exported A/b to {out_dir}")
    print(f"  n={a.shape[0]}, nnz={ai.size}, |b|={np.linalg.norm(b):.6g}")