#!/usr/bin/env python3
"""Compute Mathieu bands from Bloch matrices and ODE monodromy.

The standard Mathieu operator is

    H_q = -d^2/dx^2 + 2*q*cos(2*x),  q >= 0,

with period pi and spectral parameter a.  The script computes periodic and
antiperiodic characteristic values with Fourier--Bloch matrices, samples the
band functions, and independently propagates the fundamental ODE matrix to
check the Hill discriminant Delta(a)=tr M_pi(a).  Its floating-point audit is
calibrated for 0 <= q <= 8; stronger coupling requires a rescaled or
high-precision propagator.

Refinement shifts and cross-representation residuals are empirical
diagnostics, not certified enclosures.

Requirements: CPython 3.9+ and NumPy 1.24+.
"""

from __future__ import annotations

import argparse
import csv
from dataclasses import dataclass
import math
from pathlib import Path
import platform
from typing import Iterable

import numpy as np


PI = math.pi
MAX_CALIBRATED_Q = 8.0


def require(condition: bool, message: str) -> None:
    """Raise an optimization-safe exception when an audit fails."""

    if not condition:
        raise RuntimeError(message)


def bloch_matrix(nu: float, q: float, cutoff: int) -> np.ndarray:
    """Return the Fourier matrix for multiplier exp(i*pi*nu)."""

    require(q >= 0.0, "this real-band calibration uses q >= 0")
    require(cutoff >= 4, "Fourier cutoff is too small")
    # At nu=1 the exact involution is m -> -m-1.  The even-sized endpoint
    # section below preserves it; the common -N,...,N section does not.
    if abs(nu - 1.0) < 8.0 * np.finfo(float).eps:
        modes = np.arange(-cutoff, cutoff, dtype=float)
    else:
        modes = np.arange(-cutoff, cutoff + 1, dtype=float)
    diagonal = (nu + 2.0 * modes) ** 2
    matrix = np.diag(diagonal)
    coupling = np.full(modes.size - 1, q)
    matrix += np.diag(coupling, 1) + np.diag(coupling, -1)
    return matrix


def bloch_levels(nu: float, q: float, cutoff: int, count: int) -> np.ndarray:
    """Return the lowest Bloch-fiber eigenvalues."""

    values = np.linalg.eigvalsh(bloch_matrix(nu, q, cutoff))
    return values[:count]


@dataclass(frozen=True)
class Edge:
    """One named Mathieu characteristic value."""

    name: str
    index: int
    value: float
    discriminant_sign: int
    sector: str


def characteristic_edges(q: float, cutoff: int, band_count: int) -> list[Edge]:
    """Return a_0,b_1,a_1,...,b_band_count in spectral order."""

    require(band_count >= 1, "at least one band is required")
    periodic = bloch_levels(0.0, q, cutoff, band_count + 1)
    antiperiodic = bloch_levels(1.0, q, cutoff, band_count + 1)

    a_values: dict[int, float] = {0: float(periodic[0])}
    b_values: dict[int, float] = {}
    sectors: dict[int, str] = {0: "periodic"}
    for n in range(1, band_count + 1):
        fiber = antiperiodic if n % 2 else periodic
        a_values[n] = float(fiber[n])
        b_values[n] = float(fiber[n - 1])
        sectors[n] = "antiperiodic" if n % 2 else "periodic"

    edges = [
        Edge("a_0", 0, a_values[0], +1, "periodic"),
    ]
    for n in range(1, band_count + 1):
        sign = -1 if n % 2 else +1
        edges.append(Edge(f"b_{n}", n, b_values[n], sign, sectors[n]))
        if n < band_count:
            edges.append(Edge(f"a_{n}", n, a_values[n], sign, sectors[n]))
    return edges


def ode_rhs(x: float, state: np.ndarray, a: float, q: float) -> np.ndarray:
    """Return the first-order system for two fundamental solutions."""

    coefficient = 2.0 * q * math.cos(2.0 * x) - a
    c, cp, s, sp = state
    return np.array((cp, coefficient * c, sp, coefficient * s))


def monodromy(a: float, q: float, steps: int) -> np.ndarray:
    """Propagate the fundamental matrix over one period with RK4."""

    require(steps >= 64, "too few propagation steps")
    step = PI / steps
    state = np.array((1.0, 0.0, 0.0, 1.0))
    x = 0.0
    for _ in range(steps):
        k1 = ode_rhs(x, state, a, q)
        k2 = ode_rhs(x + 0.5 * step, state + 0.5 * step * k1, a, q)
        k3 = ode_rhs(x + 0.5 * step, state + 0.5 * step * k2, a, q)
        k4 = ode_rhs(x + step, state + step * k3, a, q)
        state += step * (k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0
        x += step
    c, cp, s, sp = state
    return np.array(((c, s), (cp, sp)))


def refined_monodromy(a: float, q: float, steps: int) -> tuple[np.ndarray, float]:
    """Cancel the leading fourth-order propagation error by step halving."""

    coarse = monodromy(a, q, steps)
    fine = monodromy(a, q, 2 * steps)
    extrapolated = (16.0 * fine - coarse) / 15.0
    shift = float(np.max(np.abs(extrapolated - fine)))
    return extrapolated, shift


def discriminant(
    a: float,
    q: float,
    steps: int,
) -> tuple[float, float, float, float]:
    """Return Delta, determinant/symmetry defects, and refinement shift."""

    matrix, shift = refined_monodromy(a, q, steps)
    return (
        float(np.trace(matrix)),
        abs(float(np.linalg.det(matrix)) - 1.0),
        abs(float(matrix[0, 0] - matrix[1, 1])),
        shift,
    )


def sample_bands(
    q: float,
    cutoff: int,
    band_count: int,
    sample_count: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Sample the first band functions for 0 <= nu <= 1."""

    nus = np.linspace(0.0, 1.0, sample_count)
    levels = np.vstack(
        [bloch_levels(float(nu), q, cutoff, band_count) for nu in nus]
    )
    return nus, levels


def leading_gap(q: float, n: int) -> float:
    """Return the leading weak-coupling width a_n(q)-b_n(q)."""

    require(n >= 1, "gap index must be positive")
    return 8.0 * (q / 4.0) ** n / math.factorial(n - 1) ** 2


def edge_dictionary(edges: Iterable[Edge]) -> dict[str, Edge]:
    """Index edge records by their printed names."""

    return {edge.name: edge for edge in edges}


def print_report(
    q: float,
    edges: list[Edge],
    *,
    cutoff_shift: float,
    propagation_steps: int,
    small_q: float,
    small_edges: list[Edge],
) -> None:
    """Print the characteristic values and all numerical audits."""

    print("Mathieu Floquet-band calibration")
    print("H_q = -d^2/dx^2 + 2 q cos(2x), period pi")
    print(f"q = {q:g}; Python {platform.python_version()}; NumPy {np.__version__}")
    print()
    print("Characteristic edges from Fourier--Bloch matrices")
    print(f"{'edge':>8s} {'type':>14s} {'a':>18s} {'Delta target':>14s}")
    discriminant_rows: list[tuple[Edge, float, float, float, float]] = []
    for edge in edges:
        delta, determinant_defect, symmetry_defect, step_shift = discriminant(
            edge.value, q, propagation_steps
        )
        discriminant_rows.append(
            (edge, delta, determinant_defect, symmetry_defect, step_shift)
        )
        print(
            f"{edge.name:>8s} {edge.sector:>14s} {edge.value:18.12f} "
            f"{2 * edge.discriminant_sign:14d}"
        )
    print()
    print("Independent monodromy audit at the Fourier edges")
    print(f"{'edge':>8s} {'Delta':>18s} {'|Delta-target|':>18s} "
          f"{'|det M-1|':>14s} {'|c-s prime|':>14s} {'step shift':>14s}")
    for edge, delta, determinant_defect, symmetry_defect, step_shift in discriminant_rows:
        target = 2.0 * edge.discriminant_sign
        print(
            f"{edge.name:>8s} {delta:18.12f} {abs(delta-target):18.2e} "
            f"{determinant_defect:14.2e} {symmetry_defect:14.2e} "
            f"{step_shift:14.2e}"
        )

    records = edge_dictionary(edges)
    print()
    print("Bands and gaps")
    band_count = (len(edges) + 1) // 2
    for j in range(band_count):
        lower = records[f"a_{j}"].value
        upper = records[f"b_{j + 1}"].value
        delta_mid, _, _, _ = discriminant(
            0.5 * (lower + upper), q, propagation_steps
        )
        print(
            f"  band {j}: [{lower:.10f}, {upper:.10f}], "
            f"|Delta(mid)|={abs(delta_mid):.6f}"
        )
        if j + 1 < band_count:
            gap_lower = records[f"b_{j + 1}"].value
            gap_upper = records[f"a_{j + 1}"].value
            gap_delta, _, _, _ = discriminant(
                0.5 * (gap_lower + gap_upper), q, propagation_steps
            )
            print(
                f"  gap  {j + 1}: ({gap_lower:.10f}, {gap_upper:.10f}), "
                f"|Delta(mid)|={abs(gap_delta):.6f}"
            )

    print()
    print(f"Maximum Fourier-cutoff shift: {cutoff_shift:.3e}")
    print(
        "Maximum edge discriminant residual: "
        f"{max(abs(delta - 2.0 * edge.discriminant_sign) for edge, delta, _, _, _ in discriminant_rows):.3e}"
    )
    print(
        "Maximum monodromy determinant defect: "
        f"{max(defect for _, _, defect, _, _ in discriminant_rows):.3e}"
    )
    print(
        "Maximum Mathieu symmetry defect: "
        f"{max(defect for _, _, _, defect, _ in discriminant_rows):.3e}"
    )
    print(
        "Maximum monodromy step shift: "
        f"{max(shift for _, _, _, _, shift in discriminant_rows):.3e}"
    )

    small_records = edge_dictionary(small_edges)
    print()
    print(f"Weak-coupling gaps at q={small_q:g}")
    print(f"{'n':>5s} {'a_n-b_n':>16s} {'leading':>16s} {'ratio':>12s}")
    for n in range(1, min(4, (len(small_edges) + 1) // 2)):
        width = small_records[f"a_{n}"].value - small_records[f"b_{n}"].value
        leading = leading_gap(small_q, n)
        print(f"{n:5d} {width:16.9e} {leading:16.9e} {width/leading:12.8f}")


def write_csv(path: str, nus: np.ndarray, levels: np.ndarray) -> None:
    """Write a dynamic-width Bloch-dispersion table."""

    with open(path, "w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["nu"] + [f"band_{j}" for j in range(levels.shape[1])])
        for nu, row in zip(nus, levels):
            writer.writerow([float(nu)] + [float(value) for value in row])


def write_figure_data(
    prefix: str,
    q: float,
    cutoff: int,
    propagation_steps: int,
    edges: list[Edge],
) -> tuple[Path, Path, Path]:
    """Export every numerical input needed to rebuild the page figure."""

    base_grid = np.linspace(-0.6, 9.2, 81)
    plot_edges = sorted(
        edge.value for edge in edges if -0.6 <= edge.value <= 9.2
    )
    # Resolve every edge and the narrow intervals that a uniform grid could
    # otherwise miss.  This reproduces the sampling rule used in the TikZ
    # source while remaining deterministic for another calibrated q.
    narrow_midpoints = [
        0.5 * (left + right)
        for left, right in zip(plot_edges, plot_edges[1:])
        if right - left < 0.5
    ]
    energies = np.array(
        sorted(set(float(value) for value in base_grid)
               | set(plot_edges)
               | set(narrow_midpoints))
    )

    prefix_path = Path(prefix)
    discriminant_path = Path(f"{prefix_path}-discriminant.csv")
    bands_path = Path(f"{prefix_path}-bands.csv")
    edges_path = Path(f"{prefix_path}-edges.csv")

    with discriminant_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["a", "Delta"])
        for energy in energies:
            delta, _, _, _ = discriminant(
                float(energy), q, propagation_steps
            )
            writer.writerow([float(energy), delta])

    figure_nus, figure_levels = sample_bands(q, cutoff, 3, 21)
    write_csv(str(bands_path), figure_nus, figure_levels)

    with edges_path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(
            ["name", "index", "a", "Delta_target", "sector"]
        )
        for edge in edges:
            writer.writerow(
                [
                    edge.name,
                    edge.index,
                    edge.value,
                    2 * edge.discriminant_sign,
                    edge.sector,
                ]
            )

    return discriminant_path, bands_path, edges_path


def calibrated_coupling(text: str) -> float:
    """Parse a coupling inside the range covered by the numerical audit."""

    try:
        value = float(text)
    except ValueError as error:
        raise argparse.ArgumentTypeError("q must be a real number") from error
    if not 0.0 <= value <= MAX_CALIBRATED_Q:
        raise argparse.ArgumentTypeError(
            "q must lie in 0 <= q <= 8; stronger coupling needs a "
            "rescaled or high-precision propagator"
        )
    return value


def parse_arguments() -> argparse.Namespace:
    """Parse command-line options."""

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--q",
        type=calibrated_coupling,
        default=1.0,
        help="Mathieu coupling in the calibrated interval 0 <= q <= 8",
    )
    parser.add_argument(
        "--high",
        action="store_true",
        help="use six bands, a larger Fourier cutoff, and finer ODE propagation",
    )
    parser.add_argument(
        "--csv",
        metavar="PATH",
        help="write sampled Bloch band functions to a CSV file",
    )
    parser.add_argument(
        "--figure-data",
        metavar="PREFIX",
        help=(
            "write PREFIX-discriminant.csv, PREFIX-bands.csv, and "
            "PREFIX-edges.csv for the textbook figure"
        ),
    )
    return parser.parse_args()


def main() -> None:
    """Run the two-representation Mathieu calibration."""

    arguments = parse_arguments()
    band_count = 6 if arguments.high else 4
    cutoff = 48 if arguments.high else 32
    propagation_steps = 4096 if arguments.high else 2048
    sample_count = 81 if arguments.high else 41

    edges = characteristic_edges(arguments.q, cutoff, band_count)
    smaller_edges = characteristic_edges(arguments.q, cutoff - 8, band_count)
    cutoff_shift = max(
        abs(left.value - right.value)
        for left, right in zip(edges, smaller_edges)
    )

    # The free discriminant is an exact normalization check.
    free_errors = []
    for a in (0.25, 2.3, 9.7):
        delta, _, _, _ = discriminant(a, 0.0, propagation_steps)
        free_errors.append(abs(delta - 2.0 * math.cos(PI * math.sqrt(a))))
    require(max(free_errors) < 2.0e-10, "free discriminant audit failed")

    if arguments.q > 0.0:
        values = np.array([edge.value for edge in edges])
        require(np.all(np.diff(values) > 0.0), "Mathieu edge ordering audit failed")
        max_residual = 0.0
        max_determinant_defect = 0.0
        max_symmetry_defect = 0.0
        for edge in edges:
            delta, determinant_defect, symmetry_defect, _ = discriminant(
                edge.value, arguments.q, propagation_steps
            )
            max_residual = max(
                max_residual,
                abs(delta - 2.0 * edge.discriminant_sign),
            )
            max_determinant_defect = max(max_determinant_defect, determinant_defect)
            max_symmetry_defect = max(max_symmetry_defect, symmetry_defect)
        require(max_residual < 2.0e-8, "Fourier/monodromy edge audit failed")
        require(max_determinant_defect < 2.0e-10, "Wronskian audit failed")
        require(max_symmetry_defect < 2.0e-10, "Mathieu symmetry audit failed")

    small_q = 0.05
    small_edges = characteristic_edges(small_q, cutoff, 4)
    print_report(
        arguments.q,
        edges,
        cutoff_shift=cutoff_shift,
        propagation_steps=propagation_steps,
        small_q=small_q,
        small_edges=small_edges,
    )

    nus, levels = sample_bands(
        arguments.q,
        cutoff,
        band_count,
        sample_count,
    )
    if arguments.csv:
        write_csv(arguments.csv, nus, levels)
        print(f"\nWrote {arguments.csv}")
    if arguments.figure_data:
        paths = write_figure_data(
            arguments.figure_data,
            arguments.q,
            cutoff,
            propagation_steps,
            edges,
        )
        print("\nWrote figure data:")
        for path in paths:
            print(f"  {path}")


if __name__ == "__main__":
    main()