"""
doe_design.py
---------------
Designs a phase-only Diffractive Optical Element (DOE / kinoform) with an
Iterative Fourier Transform Algorithm (Gerchberg-Saxton), and writes the
complex field leaving the DOE to a Zemax Beam File (.ZBF) that can be used as
the source beam in an OpticStudio Physical Optics Propagation (POP) analysis.

Optical arrangement assumed by the design:

    DOE plane ---- f ---- lens (focal length f) ---- f ---- display plane

The field at the back focal plane is the Fourier transform of the field
leaving the DOE, so the DOE phase is designed directly in the FFT domain.

Scaling at the display plane (lens units = mm):
    pixel pitch      dx_img = lam * f / (N * dx)
    full array width W_img  = lam * f / dx
    max diffraction angle    theta_max = lam / (2 * dx)

ZBF binary format (OpticStudio Help File):
    5 int32 : version(1), nx, ny, ispol, units(0=mm,1=cm,2=in,3=m)
    4 int32 : unused
    20 float64 : dx, dy, zx, Rx, wx, zy, Ry, wy, lambda, index, re, se,
                 + 8 unused
    2*nx*ny float64 : Ex real/imag interleaved, x index fastest
All lengths, including the wavelength, are in the units given by the units flag.

Author: <your name>
"""

import os
import struct
import numpy as np

# ----------------------------------------------------------------------------
# 1. Design parameters (all lengths in mm, i.e. OpticStudio lens units)
# ----------------------------------------------------------------------------
N = 512             # samples across the DOE (must match POP X/Y sampling)
DX = 0.016          # DOE sample pitch [mm]  -> 16 um features
LAM = 633e-6        # wavelength [mm] (0.633 um HeNe)
W0 = 2.0            # illuminating Gaussian beam waist radius [mm]
F_LENS = 50.0       # Fourier-transform lens focal length [mm]
TOTAL_POWER = 1.0   # total power in the beam [W]

N_ITER = 60         # Gerchberg-Saxton iterations
SIGNAL_FRAC = 0.5   # target occupies this fraction of the display array
SEED = 12345        # random seed for the starting phase

# ---- target definition ------------------------------------------------------
# Set IMAGE_FILE to None to use the built-in geometric test pattern, or to a
# path to use your own artwork. Use the r"..." prefix so that the backslashes
# in a Windows path are not read as escape characters.
IMAGE_FILE = "doe_input.png"   # file name, or None for the built-in test pattern

IMAGE_INVERT = False      # True if the artwork is dark on a light background
IMAGE_THRESHOLD = None    # e.g. 0.5 to force a hard binary target, None to keep greyscale
IMAGE_MARGIN = 0.9        # fraction of the signal window the artwork may fill
PREVIEW_ONLY = False      # True: write the target preview and stop (fast check)

OUTNAME = "doe_pattern"   # output file base name

# Derived display-plane scaling
DX_IMG = LAM * F_LENS / (N * DX)
W_IMG = LAM * F_LENS / DX
THETA_MAX = LAM / (2.0 * DX)


# ----------------------------------------------------------------------------
# 2. Target pattern
# ----------------------------------------------------------------------------
def make_target(n, signal_frac=0.5, image_file=None,
                invert=False, threshold=None, margin=0.9):
    """Return the desired far-field irradiance pattern on an n x n grid.

    White means light, black means none. If image_file is given the artwork is
    loaded, flattened to greyscale, scaled to fit the signal window and
    centred. Otherwise a built-in geometric test pattern is generated so that
    the script is fully self-contained.
    """
    tgt = np.zeros((n, n), dtype=float)
    m = int(round(n * signal_frac * margin))   # signal window size in pixels
    m += m % 2

    if image_file:
        if not os.path.exists(image_file):
            raise FileNotFoundError(
                f"IMAGE_FILE does not exist: {image_file}")

        from PIL import Image
        img = Image.open(image_file)

        # flatten transparency onto black, so "transparent" means "no light"
        if img.mode in ("RGBA", "LA", "P"):
            bg = Image.new("RGBA", img.size, (0, 0, 0, 255))
            img = Image.alpha_composite(bg, img.convert("RGBA"))
        img = img.convert("L")

        img.thumbnail((m, m), Image.LANCZOS)
        a = np.asarray(img, dtype=float)
        a = (a - a.min()) / max(a.max() - a.min(), 1e-12)

        if invert:
            a = 1.0 - a
        a[a < 0.02] = 0.0                      # remove resampling ringing
        if threshold is not None:
            a = (a >= threshold).astype(float)

        y0 = (n - a.shape[0]) // 2
        x0 = (n - a.shape[1]) // 2
        tgt[y0:y0 + a.shape[0], x0:x0 + a.shape[1]] = a

        print(f"    loaded {os.path.basename(image_file)}: "
              f"{a.shape[1]} x {a.shape[0]} px in the target grid, "
              f"{100*np.mean(tgt > 0):.1f} % of the array lit")
        return tgt

    # ---- built-in pattern: ring + cross + four corner squares -------------
    yy, xx = np.mgrid[0:n, 0:n].astype(float)
    xx = (xx - n / 2) / (m / 2)                # normalised to signal window
    yy = (yy - n / 2) / (m / 2)
    r = np.hypot(xx, yy)

    ring = (r > 0.80) & (r < 0.92)
    cross = ((np.abs(xx) < 0.055) & (np.abs(yy) < 0.45)) | \
            ((np.abs(yy) < 0.055) & (np.abs(xx) < 0.45))
    dots = np.zeros_like(r, dtype=bool)
    for cx, cy in [(-0.55, -0.55), (0.55, -0.55), (-0.55, 0.55), (0.55, 0.55)]:
        dots |= (np.abs(xx - cx) < 0.09) & (np.abs(yy - cy) < 0.09)
    inner = (r > 0.28) & (r < 0.34)

    tgt[ring | cross | dots | inner] = 1.0
    return tgt


# ----------------------------------------------------------------------------
# 3. Iterative Fourier Transform Algorithm (Gerchberg-Saxton)
# ----------------------------------------------------------------------------
def design_kinoform(source_amp, target_irr, n_iter=60, seed=0, verbose=True):
    """Return (phase, history). phase is the DOE phase in radians."""
    rng = np.random.default_rng(seed)
    tgt_amp = np.sqrt(target_irr)
    tgt_amp /= np.sqrt(np.sum(tgt_amp ** 2))

    signal = target_irr > 0
    phase = rng.uniform(-np.pi, np.pi, source_amp.shape)
    history = []

    for k in range(n_iter):
        # DOE plane -> display plane
        g = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(source_amp * np.exp(1j * phase))))
        g /= np.sqrt(np.sum(np.abs(g) ** 2))

        # metrics
        eff = np.sum(np.abs(g[signal]) ** 2)                 # power in signal
        vals = np.abs(g[signal]) ** 2
        rms = np.std(vals) / np.mean(vals)                   # non-uniformity
        history.append((k, eff, rms))
        if verbose and (k % 10 == 0 or k == n_iter - 1):
            print(f"    iter {k:3d}   efficiency {eff*100:6.2f} %   "
                  f"RMS non-uniformity {rms*100:6.2f} %")

        # apply the target amplitude, keep the phase
        g = tgt_amp * np.exp(1j * np.angle(g))

        # display plane -> DOE plane, keep the phase only
        h = np.fft.fftshift(np.fft.ifft2(np.fft.ifftshift(g)))
        phase = np.angle(h)

    return phase, history


# ----------------------------------------------------------------------------
# 4. Zemax Beam File writer / reader
# ----------------------------------------------------------------------------
def write_zbf(path, field, dx, dy, lam, waist_x, waist_y,
              zx=0.0, zy=0.0, index=1.0, units=0):
    """Write a scalar (unpolarized) complex field to a binary ZBF file.

    field   : complex 2-D array, shape (ny, nx), row index = y
    dx, dy  : sample spacing in lens units
    lam     : wavelength in lens units
    waist_x/y : pilot beam waist radius in lens units
    zx, zy  : distance from the pilot beam waist (0 = beam is at its waist)
    units   : 0 = mm, 1 = cm, 2 = in, 3 = m
    """
    ny, nx = field.shape
    rayleigh_x = np.pi * waist_x ** 2 / lam
    rayleigh_y = np.pi * waist_y ** 2 / lam

    with open(path, "wb") as fp:
        fp.write(struct.pack("<9i", 1, nx, ny, 0, units, 0, 0, 0, 0))
        fp.write(struct.pack("<12d",
                             dx, dy,
                             zx, rayleigh_x, waist_x,
                             zy, rayleigh_y, waist_y,
                             lam, index, 0.0, 0.0))
        fp.write(struct.pack("<8d", *([0.0] * 8)))
        data = np.empty(2 * nx * ny, dtype="<f8")
        data[0::2] = np.real(field).ravel(order="C")   # x index fastest
        data[1::2] = np.imag(field).ravel(order="C")
        fp.write(data.tobytes())


def read_zbf(path):
    """Read a scalar or polarized ZBF file. Returns (header dict, Ex, Ey)."""
    with open(path, "rb") as fp:
        ints = struct.unpack("<9i", fp.read(36))
        dbls = struct.unpack("<20d", fp.read(160))
        version, nx, ny, ispol, units = ints[:5]
        keys = ["dx", "dy", "zx", "Rx", "wx", "zy", "Ry", "wy",
                "lambda", "index", "re", "se"]
        hdr = dict(zip(keys, dbls[:12]))
        hdr.update(version=version, nx=nx, ny=ny, ispol=ispol, units=units)

        raw = np.frombuffer(fp.read(16 * nx * ny), dtype="<f8")
        ex = (raw[0::2] + 1j * raw[1::2]).reshape(ny, nx)
        ey = None
        if ispol:
            raw = np.frombuffer(fp.read(16 * nx * ny), dtype="<f8")
            ey = (raw[0::2] + 1j * raw[1::2]).reshape(ny, nx)
    return hdr, ex, ey


# ----------------------------------------------------------------------------
# 5. Main
# ----------------------------------------------------------------------------
if __name__ == "__main__":
    here = os.path.dirname(os.path.realpath(__file__))
    out_dir = os.path.join(here, "output")
    os.makedirs(out_dir, exist_ok=True)

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    print("DOE / kinoform design for OpticStudio POP")
    print("=========================================")
    print(f"  wavelength            {LAM*1e6:.1f} nm")
    print(f"  DOE grid              {N} x {N} samples, pitch {DX*1e3:.1f} um")
    print(f"  DOE aperture          {N*DX:.3f} mm")
    print(f"  Gaussian waist w0     {W0:.3f} mm")
    print(f"  FT lens focal length  {F_LENS:.1f} mm")
    print(f"  max diffraction angle {THETA_MAX*1e3:.2f} mrad")
    print(f"  display array width   {W_IMG:.4f} mm  ({DX_IMG*1e3:.3f} um / pixel)")
    print(f"  signal window         {W_IMG*SIGNAL_FRAC:.4f} mm")
    print(f"  target                "
          f"{'built-in test pattern' if not IMAGE_FILE else IMAGE_FILE}\n")

    # 5.1 illumination and target
    x = (np.arange(N) - N / 2) * DX
    X, Y = np.meshgrid(x, x)
    source_amp = np.exp(-(X ** 2 + Y ** 2) / W0 ** 2)   # w0 = 1/e^2 irradiance radius

    img_path = IMAGE_FILE
    if img_path and not os.path.isabs(img_path):
        img_path = os.path.join(here, img_path)   # relative to this script

    print("  Building target:")
    target = make_target(N, SIGNAL_FRAC, img_path,
                         invert=IMAGE_INVERT, threshold=IMAGE_THRESHOLD,
                         margin=IMAGE_MARGIN)

    # The FFT sign convention used here is mirrored in X relative to the
    # OpticStudio display. Flipping the target once, before the design runs,
    # makes the projected pattern appear the right way round in POP.
    target = np.fliplr(target)

    half = W_IMG / 2
    prev = os.path.join(out_dir, "target_preview.png")
    figp, axp = plt.subplots(figsize=(6, 5.5))
    axp.imshow(target, cmap="gray", extent=[-half, half, -half, half])
    axp.set_title("Target pattern at the display plane")
    axp.set_xlabel("x [mm]")
    axp.set_ylabel("y [mm]")
    plt.tight_layout()
    figp.savefig(prev, dpi=130)
    plt.close(figp)
    print(f"    target preview written to {prev}")

    if PREVIEW_ONLY:
        print("\n  PREVIEW_ONLY is set, stopping before the design loop.")
        raise SystemExit(0)

    # 5.2 design
    print("\n  Running Gerchberg-Saxton:")
    phase, hist = design_kinoform(source_amp, target, N_ITER, SEED)

    # 5.3 field leaving the DOE, normalised so that sum(|E|^2) = TOTAL_POWER
    field = source_amp * np.exp(1j * phase)
    field *= np.sqrt(TOTAL_POWER / np.sum(np.abs(field) ** 2))

    # 5.4 write the ZBF
    zbf_path = os.path.join(out_dir, OUTNAME + ".zbf")
    write_zbf(zbf_path, field, DX, DX, LAM, W0, W0)
    print(f"\n  Wrote {zbf_path}  ({os.path.getsize(zbf_path)/1e6:.1f} MB)")

    # 5.5 read it back and verify
    hdr, ex, _ = read_zbf(zbf_path)
    err = np.max(np.abs(ex - field))
    print("  Round-trip check:")
    print(f"    header  nx={hdr['nx']} ny={hdr['ny']} ispol={hdr['ispol']} "
          f"units={hdr['units']}")
    print(f"    dx={hdr['dx']:.6f} mm  lambda={hdr['lambda']:.3e} mm  "
          f"w={hdr['wx']:.4f} mm  R={hdr['Rx']:.4e} mm")
    print(f"    total power = {np.sum(np.abs(ex)**2):.6f} W")
    print(f"    max |E_read - E_written| = {err:.3e}")

    # 5.6 predicted display-plane irradiance (what POP should reproduce)
    recon = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(ex)))
    recon_irr = np.abs(recon) ** 2
    recon_irr /= recon_irr.max()

    signal = target > 0
    g = np.abs(recon) ** 2 / np.sum(np.abs(recon) ** 2)
    print(f"\n  Predicted diffraction efficiency into the pattern: "
          f"{100*np.sum(g[signal]):.1f} %")

    # 5.7 figures
    fig, ax = plt.subplots(2, 2, figsize=(11, 10))

    ax[0, 0].imshow(target, cmap="gray",
                    extent=[-half, half, -half, half])
    ax[0, 0].set_title("Target pattern at the display plane")
    ax[0, 0].set_xlabel("x [mm]")

    im = ax[0, 1].imshow(np.abs(field) ** 2, cmap="inferno",
                         extent=[-N*DX/2, N*DX/2, -N*DX/2, N*DX/2])
    ax[0, 1].set_title("DOE plane irradiance (Gaussian illumination)")
    ax[0, 1].set_xlabel("x [mm]")
    plt.colorbar(im, ax=ax[0, 1], fraction=0.046)

    c = N // 2
    w = 64
    im = ax[1, 0].imshow(np.angle(field)[c-w:c+w, c-w:c+w], cmap="twilight",
                         vmin=-np.pi, vmax=np.pi,
                         extent=[-w*DX, w*DX, -w*DX, w*DX])
    ax[1, 0].set_title("DOE phase, central 128 x 128 samples [rad]")
    ax[1, 0].set_xlabel("x [mm]")
    plt.colorbar(im, ax=ax[1, 0], fraction=0.046)

    ax[1, 1].imshow(recon_irr ** 0.5, cmap="inferno",
                    extent=[-half, half, -half, half])
    ax[1, 1].set_title("Predicted display-plane irradiance (sqrt scale)")
    ax[1, 1].set_xlabel("x [mm]")

    plt.tight_layout()
    fig.savefig(os.path.join(out_dir, "doe_design_summary.png"), dpi=130)

    fig2, ax2 = plt.subplots(figsize=(6, 5.5))
    ax2.imshow(recon_irr ** 0.5, cmap="inferno",
               extent=[-half, half, -half, half])
    ax2.set_title("Expected POP result at the display plane")
    ax2.set_xlabel("x [mm]")
    ax2.set_ylabel("y [mm]")
    plt.tight_layout()
    fig2.savefig(os.path.join(out_dir, "expected_pop_result.png"), dpi=130)

    print(f"  Figures written to {out_dir}")
    print("\n  Next: copy the .ZBF into "
          r"{Documents}\Zemax\POP\BEAMFILES and select it in POP "
          "(Beam Definition > File).")