"""Read a saved result without an API call. Add --plot to save a field image."""
import argparse
from pathlib import Path
import numpy as np

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("result", nargs="?", type=Path,
                    default=Path(__file__).with_name("ebeam_y_1550_result.npz"),
                    help="saved NPZ; default: ebeam_y_1550_result.npz beside this script")
parser.add_argument("--plot", action="store_true", help="save a PNG; requires matplotlib")
parser.add_argument("--check-touchstone", action="store_true",
                    help="compare the Touchstone file with saved arrays; requires scikit-rf")
args = parser.parse_args()
if not args.result.is_file():
    parser.error(f"Result file not found: {args.result}. Run quickstart.py first or select a saved NPZ file.")

with np.load(args.result, allow_pickle=False) as res:
    print("S shape (wavelength, output port, input port):", res["S"].shape)
    print("Power ratios at the first wavelength:")
    print(np.abs(res["S"][0]) ** 2)
    print("Wavelengths (um):", res["wavelengths"])
    print("Effective index:", res["n_eff"])
    print("Field shape (y, x):", res["E"].shape)
    if args.check_touchstone:
        try:
            import skrf
        except ImportError:
            parser.error("Install the export-check dependency: python -m pip install scikit-rf")
        saved_s = res["S"]
        wavelengths = res["wavelengths"]
        if saved_s.ndim != 3 or saved_s.shape[1] != saved_s.shape[2]:
            parser.error("The saved S array must have shape (wavelength, port, port).")
        if wavelengths.ndim != 1 or len(wavelengths) != len(saved_s):
            parser.error("The saved wavelengths must match the first S dimension.")
        if not np.isfinite(wavelengths).all() or np.any(wavelengths <= 0):
            parser.error("Wavelengths must be finite and greater than zero.")
        name = args.result.stem.removesuffix("_result")
        touchstone = args.result.with_name(f"{name}.s{saved_s.shape[1]}p")
        if not touchstone.is_file():
            parser.error(f"Touchstone file not found: {touchstone}")
        try:
            network = skrf.Network(str(touchstone))
            frequencies = 299792458.0 / (wavelengths * 1e-6)
            expected_order = np.argsort(frequencies)
            actual_order = np.argsort(network.f)
            if network.s.shape != saved_s.shape:
                raise ValueError("S dimensions differ.")
            np.testing.assert_allclose(network.f[actual_order], frequencies[expected_order],
                                       rtol=1e-8, atol=1.0)
            np.testing.assert_allclose(network.s[actual_order], saved_s[expected_order],
                                       rtol=0, atol=8e-9)
        except (OSError, ValueError, AssertionError) as error:
            parser.error(f"Touchstone check failed: {error}")
        print("Touchstone check passed: S and frequencies match the saved arrays.")
    if args.plot:
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            parser.error("Install the plot dependency: python -m pip install matplotlib")
        intensity = np.abs(res["E"]) ** 2
        if not np.isfinite(intensity).all() or intensity.max() <= 0:
            parser.error("The field must be finite and have nonzero intensity.")
        intensity /= intensity.max()
        fig, ax = plt.subplots(figsize=(8, 4), layout="constrained")
        mesh = ax.pcolormesh(res["x"], res["y"], intensity, shading="auto",
                             cmap="magma", vmin=0, vmax=1)
        ax.set(xlabel="x (um)", ylabel="y (um)",
               title="API field: intensity divided by its maximum", aspect="equal")
        fig.colorbar(mesh, ax=ax, label="Normalized |E|²")
        output = args.result.with_suffix(".png")
        fig.savefig(output, dpi=180)
        plt.close(fig)
        print("Saved:", output)