"""Run the Y-branch example. Install with: python -m pip install --upgrade bamsolve."""
import argparse
import getpass
import os
from pathlib import Path
import sys
import warnings

import numpy as np
from bamsolve import Layer, Stack, bamsolve, pred


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("gds", nargs="?", type=Path,
                        default=Path(__file__).with_name("ebeam_y_1550.gds"),
                        help="GDS file; default: ebeam_y_1550.gds beside this script")
    args = parser.parse_args()
    gds = args.gds.expanduser().resolve()
    if not gds.is_file():
        parser.error(f"GDS file not found: {gds}. Download it beside quickstart.py.")

    key = os.environ.get("BAMSOLVE_API_KEY", "").strip()
    if not key:
        if not sys.stdin.isatty():
            parser.error("Run in a terminal for the hidden key prompt, or set BAMSOLVE_API_KEY.")
        with warnings.catch_warnings():
            warnings.simplefilter("error", getpass.GetPassWarning)
            try:
                key = getpass.getpass("Paste your API key (input is hidden): ").strip()
            except (getpass.GetPassWarning, EOFError):
                parser.error("The hidden prompt is not available. Set BAMSOLVE_API_KEY instead.")
    if not key:
        parser.error("An API key is required. Get one at https://bamsolve.com/#get-key.")
    bamsolve(api_key=key)

    stack = Stack(clad=1.444, layers={(1, 0): Layer(n=3.48, z=(0.0, 0.22))})
    print("Requesting TE at 1.55 um. Wait for the API response.")
    res = pred(stack(gds), inputs={((1, 10), "TE", 1.55)})

    ports = res.S.shape[1]
    touchstone = res.export_s(gds.with_suffix(f".s{ports}p"))
    arrays = gds.with_name(gds.stem + "_result.npz")
    np.savez_compressed(arrays, S=res.S, E=res.E, n_eff=res.n_eff,
                        x=res.x, y=res.y, wavelengths=res.wavelengths)
    print("S shape (wavelength, output port, input port):", res.S.shape)
    print("Port order in the supplied examples:", ", ".join(f"opt{i + 1}" for i in range(ports)))
    print("Power ratios at the first wavelength (rows: output; columns: input):")
    print(np.abs(res.S[0]) ** 2)
    print("Effective index:", res.n_eff)
    print("Field shape (y, x):", res.E.shape)
    print("Saved:", touchstone)
    print("Saved:", arrays)


if __name__ == "__main__":
    main()