"""
Command-line interface for Strelka FDTD.

Usage:
    strelka optimize --ports 8 --res 10
    strelka optimize --ports 64 --res 5 --iters 50
    strelka simulate topology.npy --ports 8
    strelka decompose matrix.npy --block-size 8
    strelka benchmark
    strelka info
"""

import argparse
import sys
import time
import json


def cmd_optimize(args):
    """Run inverse-design optimization."""
    from strelka.core import optimize
    import numpy as np

    pitch = 1.5
    design_um = args.ports * pitch + 2.0

    result = optimize(
        n_ports=args.ports,
        design_um=design_um,
        resolution=args.res,
        n_steps=args.steps,
        n_iters=args.iters,
        lr=args.lr,
    )

    if args.output:
        np.save(args.output, result["rho"])
        print(f"Topology saved: {args.output}")
    else:
        out = f"pnn_{args.ports}x{args.ports}_topology.npy"
        np.save(out, result["rho"])
        print(f"Topology saved: {out}")


def cmd_simulate(args):
    """Run forward simulation on existing topology."""
    from strelka.core import simulate
    import numpy as np

    rho = np.load(args.topology)
    result = simulate(rho, n_ports=args.ports, resolution=args.res, n_steps=args.steps)

    print(f"Simulation results:")
    for i, p in enumerate(result["power"]):
        print(f"  Port {i}: power = {p:.6f}")
    print(f"  Total: {result['total_power']:.6f}")


def cmd_decompose(args):
    """Decompose a weight matrix via Block Clements."""
    from strelka.cascade import GivensBlockCascade
    import numpy as np

    W = np.load(args.matrix)
    N = W.shape[0]
    print(f"Matrix: {N}×{N}, condition number: {np.linalg.cond(W):.0f}")

    gbc = GivensBlockCascade(N=N, block_size=args.block_size)
    result = gbc.inference(W)

    print(f"Reconstruction error: {result['reconstruction_error']:.2e}")
    print(f"Total devices: {result['total_devices']}")
    print(f"  Blocks ({args.block_size}×{args.block_size}): {result['total_blocks']}")
    print(f"  Cross-couplers: {result['total_cross']}")
    print(f"  Clements MZIs: {result['clements_mzis']}")
    print(f"  Reduction: {result['clements_mzis']/result['total_devices']:.1f}×")


def cmd_benchmark(args):
    """Benchmark FDTD performance on this machine."""
    import jax
    import jax.numpy as jnp

    print(f"JAX backend: {jax.default_backend()}")
    print(f"Devices: {jax.devices()}")

    N = 500
    dx, dt_val = 0.1, 0.05

    @jax.jit
    def step(Hz, Ex, Ey, eps):
        Hz = Hz.at[1:-1, 1:-1].set(
            Hz[1:-1, 1:-1] + (dt_val / (2*dx)) * (
                Ey[2:, 1:-1] - Ey[:-2, 1:-1]
                - Ex[1:-1, 2:] + Ex[1:-1, :-2]))
        Ex = Ex.at[1:-1, 1:-1].set(
            Ex[1:-1, 1:-1] - (dt_val / (2*dx)) * (
                Hz[1:-1, 2:] - Hz[1:-1, :-2]) / eps[1:-1, 1:-1])
        Ey = Ey.at[1:-1, 1:-1].set(
            Ey[1:-1, 1:-1] + (dt_val / (2*dx)) * (
                Hz[2:, 1:-1] - Hz[:-2, 1:-1]) / eps[1:-1, 1:-1])
        return Hz, Ex, Ey

    Hz = jnp.zeros((N, N))
    Ex = jnp.zeros((N, N))
    Ey = jnp.zeros((N, N))
    eps = jnp.ones((N, N)) * 2.44**2

    Hz, Ex, Ey = step(Hz, Ex, Ey, eps)
    jax.block_until_ready(Hz)

    t0 = time.perf_counter()
    for _ in range(200):
        Hz, Ex, Ey = step(Hz, Ex, Ey, eps)
    jax.block_until_ready(Hz)
    dt = (time.perf_counter() - t0) / 200

    mcells = N * N / dt / 1e6
    print(f"\n{N}×{N} FDTD step: {dt*1000:.2f} ms")
    print(f"Throughput: {mcells:.0f} Mcells/s")

    est_64x64 = 200 * 3000 * dt * 3 / 60  # 200 iters, 3000 steps, 3× for grad
    print(f"\nEstimated 64×64 optimization: {est_64x64:.0f} min")
    print(f"Meep equivalent: ~7500 min (5 days)")
    print(f"Speedup: ~{7500/est_64x64:.0f}×")


def cmd_info(args):
    """Print system and package info."""
    import strelka
    import jax
    print(f"strelka-fdtd {strelka.__version__}")
    print(f"JAX {jax.__version__} ({jax.default_backend()})")
    print(f"Devices: {jax.devices()}")
    try:
        import meep
        print(f"Meep {meep.__version__} (available)")
    except ImportError:
        print("Meep: not installed")


def main():
    parser = argparse.ArgumentParser(
        prog="strelka",
        description="Strelka FDTD — Fastest open-source differentiable photonic optimizer",
    )
    sub = parser.add_subparsers(dest="command")

    # optimize
    p_opt = sub.add_parser("optimize", help="Inverse-design a photonic device")
    p_opt.add_argument("--ports", type=int, default=8, help="Number of I/O ports")
    p_opt.add_argument("--res", type=float, default=10, help="Resolution (px/µm)")
    p_opt.add_argument("--steps", type=int, default=2000, help="FDTD timesteps")
    p_opt.add_argument("--iters", type=int, default=30, help="Iterations per beta stage")
    p_opt.add_argument("--lr", type=float, default=0.03, help="Learning rate")
    p_opt.add_argument("--output", type=str, default=None, help="Output .npy path")

    # simulate
    p_sim = sub.add_parser("simulate", help="Forward simulate a topology")
    p_sim.add_argument("topology", help="Path to .npy topology file")
    p_sim.add_argument("--ports", type=int, default=8)
    p_sim.add_argument("--res", type=float, default=10)
    p_sim.add_argument("--steps", type=int, default=2000)

    # decompose
    p_dec = sub.add_parser("decompose", help="Block Clements decomposition")
    p_dec.add_argument("matrix", help="Path to .npy weight matrix")
    p_dec.add_argument("--block-size", type=int, default=8)

    # benchmark
    sub.add_parser("benchmark", help="Benchmark FDTD on this machine")

    # info
    sub.add_parser("info", help="System info")

    args = parser.parse_args()

    if args.command == "optimize":
        cmd_optimize(args)
    elif args.command == "simulate":
        cmd_simulate(args)
    elif args.command == "decompose":
        cmd_decompose(args)
    elif args.command == "benchmark":
        cmd_benchmark(args)
    elif args.command == "info":
        cmd_info(args)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()
