"""Dynamical sweeps across the N=4 PT-symmetric exceptional point.

Model and EP data are imported from phase_ii_b_c_bisect.py.  The calculation
uses hbar=1, J=1, the conserved twice-Sz=0 block, a linear gamma ramp, and
RK4 integration of i*d|psi>/dt = H(gamma(t))|psi>.  The state is rescaled
after every step to control overflow; the accumulated log norm is retained.
Right-eigenbasis coefficient weights are diagnostics, not Born probabilities,
because the eigenbasis is nonorthogonal.
"""

import csv
from pathlib import Path

import numpy as np
import matplotlib

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

from phase_ii_b_c_bisect import (
    find_ep_N4,
    magnetization_sector_indices,
    puiseux_coefficients_N4,
    xxz_chain_N4,
)


def get_pair_modes(H, energy_ep, gamma, side):
    """Return the EP pair of right eigenmodes with a fixed branch ordering."""
    energies, vectors = np.linalg.eig(H(gamma))
    pair_idx = np.argsort(np.abs(energies - energy_ep))[:2]
    if side == "below":
        pair_idx = sorted(pair_idx, key=lambda k: energies[k].real)  # lower, upper
    elif side == "above":
        pair_idx = sorted(pair_idx, key=lambda k: energies[k].imag, reverse=True)  # gain, loss
    else:
        raise ValueError("side must be 'below' or 'above'")
    return energies[pair_idx], vectors[:, pair_idx]


def integrate_linear_ramp(psi0, H, gamma_start, gamma_stop, duration, dt_max=1e-3):
    """RK4 evolution with per-step state rescaling and accumulated log norm."""
    n_steps = int(np.ceil(duration / dt_max))
    dt = duration / n_steps
    rate = (gamma_stop - gamma_start) / duration
    psi = np.asarray(psi0, dtype=complex).copy()
    psi /= np.linalg.norm(psi)
    log_norm = 0.0

    def rhs(time, state):
        gamma = gamma_start + rate * time
        return -1j * (H(gamma) @ state)

    for step in range(n_steps):
        time = step * dt
        k1 = rhs(time, psi)
        k2 = rhs(time + dt / 2, psi + dt * k1 / 2)
        k3 = rhs(time + dt / 2, psi + dt * k2 / 2)
        k4 = rhs(time + dt, psi + dt * k3)
        psi = psi + (dt / 6) * (k1 + 2 * k2 + 2 * k3 + k4)
        scale = np.linalg.norm(psi)
        if not np.isfinite(scale) or scale == 0:
            raise FloatingPointError("State norm became zero or nonfinite during evolution.")
        log_norm += np.log(scale)
        psi /= scale
    return psi, log_norm


def endpoint_pair_weights(psi, H, gamma, energy_ep, side):
    """Resolve the final state in the full right-eigenvector basis.

    Returns normalized squared coordinate magnitudes for the EP pair and the
    fraction of total squared coordinate magnitude carried by that pair.
    These are basis-coordinate diagnostics, not probabilities.
    """
    energies, vectors = np.linalg.eig(H(gamma))
    coeffs = np.linalg.solve(vectors, psi)
    pair_idx = np.argsort(np.abs(energies - energy_ep))[:2]
    if side == "below":
        pair_idx = sorted(pair_idx, key=lambda k: energies[k].real)
        names = ("lower", "upper")
    else:
        pair_idx = sorted(pair_idx, key=lambda k: energies[k].imag, reverse=True)
        names = ("gain", "loss")
    pair_coeffs = np.abs(coeffs[pair_idx]) ** 2
    pair_share = float(pair_coeffs.sum() / np.sum(np.abs(coeffs) ** 2))
    weights = pair_coeffs / pair_coeffs.sum()
    return names, weights, pair_share


def run_sweeps(
    gamma_ep,
    energy_ep,
    offsets=(0.04,),
    speeds=(0.04, 0.01, 0.002),
    dt_max=1e-3,
):
    """Run both ramp directions, initializing each of the two EP modes."""
    sector_indices = magnetization_sector_indices(4)[0]
    H_full_0 = xxz_chain_N4(0.0)
    H_full_1 = xxz_chain_N4(1.0)
    H0 = H_full_0[np.ix_(sector_indices, sector_indices)]
    H1 = (H_full_1 - H_full_0)[np.ix_(sector_indices, sector_indices)]
    H = lambda gamma: H0 + gamma * H1

    records = []
    for offset in offsets:
        gamma_low = gamma_ep - offset
        gamma_high = gamma_ep + offset
        for speed in speeds:
            duration = (gamma_high - gamma_low) / speed
            for direction in ("up", "down"):
                if direction == "up":
                    gamma_start, gamma_stop = gamma_low, gamma_high
                    _, initial_vectors = get_pair_modes(H, energy_ep, gamma_low, "below")
                    start_names = ("lower", "upper")
                    endpoint_side = "above"
                else:
                    gamma_start, gamma_stop = gamma_high, gamma_low
                    _, initial_vectors = get_pair_modes(H, energy_ep, gamma_high, "above")
                    start_names = ("gain", "loss")
                    endpoint_side = "below"

                direction_records = []
                final_states = []
                for branch in range(2):
                    final_state, log_norm = integrate_linear_ramp(
                        initial_vectors[:, branch], H, gamma_start, gamma_stop, duration, dt_max
                    )
                    final_states.append(final_state)
                    target_names, weights, pair_share = endpoint_pair_weights(
                        final_state, H, gamma_stop, energy_ep, endpoint_side
                    )
                    direction_records.append({
                        "offset": offset,
                        "speed": speed,
                        "duration": duration,
                        "direction": direction,
                        "initial_branch": start_names[branch],
                        "final_branch_1": target_names[0],
                        "final_branch_1_weight": float(weights[0]),
                        "final_branch_2": target_names[1],
                        "final_branch_2_weight": float(weights[1]),
                        "ep_pair_coordinate_share": pair_share,
                        "log_norm_gain": log_norm,
                    })
                ray_fidelity = float(abs(np.vdot(final_states[0], final_states[1])) ** 2)
                for row in direction_records:
                    row["final_ray_fidelity_between_initial_branches"] = ray_fidelity
                records.extend(direction_records)
    return records


def write_csv(records, path):
    with open(path, "w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(records[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(records)


def plot_summary(records, path):
    plt.rcParams["axes.unicode_minus"] = False
    fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
    styles = {
        ("up", "lower"): ("#1769aa", "o", "-", "up: start lower"),
        ("up", "upper"): ("#4d9de0", "s", "--", "up: start upper"),
        ("down", "gain"): ("#c0392b", "o", "-", "down: start gain"),
        ("down", "loss"): ("#e67e22", "s", "--", "down: start loss"),
    }
    for key, (color, marker, linestyle, label) in styles.items():
        rows = sorted(
            [row for row in records if (row["direction"], row["initial_branch"]) == key],
            key=lambda row: row["speed"],
        )
        if not rows:
            continue
        speeds = [row["speed"] for row in rows]
        weights = [row["final_branch_1_weight"] for row in rows]
        lognorm = [row["log_norm_gain"] for row in rows]
        axes[0].plot(speeds, weights, marker=marker, linestyle=linestyle, color=color, label=label)
        axes[1].plot(speeds, lognorm, marker=marker, linestyle=linestyle, color=color, label=label)

    axes[0].set_xscale("log")
    speed_ticks = sorted({float(row["speed"]) for row in records})
    speed_labels = [f"{speed:g}" for speed in speed_ticks]
    axes[0].set_xticks(speed_ticks, labels=speed_labels)
    axes[0].minorticks_off()
    axes[0].set_ylim(-0.03, 1.03)
    axes[0].set_xlabel("Ramp speed |dγ/dt| (J², with J=ℏ=1)")
    axes[0].set_ylabel("Weight of first endpoint EP-pair mode", labelpad=12)
    axes[0].set_title("Endpoint mode selection")
    axes[0].grid(True, alpha=0.25)

    axes[1].set_xscale("log")
    axes[1].set_xticks(speed_ticks, labels=speed_labels)
    axes[1].minorticks_off()
    axes[1].set_xlabel("Ramp speed |dγ/dt| (J², with J=ℏ=1)")
    axes[1].set_ylabel("Accumulated log ||ψ|| (step-rescaled)", labelpad=12)
    axes[1].set_ylim(-2.0, 3.0)
    axes[1].set_yticks(np.arange(-2.0, 3.1, 1.0))
    axes[1].set_title("Net amplification / attenuation")
    axes[1].grid(True, alpha=0.25)

    handles, labels = axes[0].get_legend_handles_labels()
    fig.legend(handles, labels, loc="center left", bbox_to_anchor=(0.80, 0.5), fontsize=9, frameon=True)
    fig.suptitle("N=4, Sz=0 linear sweeps across the first EP", y=0.97)
    fig.subplots_adjust(left=0.11, right=0.78, bottom=0.20, top=0.86, wspace=0.37)
    fig.savefig(path, dpi=180, bbox_inches="tight", pad_inches=0.2)
    plt.close(fig)


if __name__ == "__main__":
    gamma_bisected, sector, bracket = find_ep_N4(return_details=True)
    gamma_ep, energy_ep, _, residual = puiseux_coefficients_N4(
        gamma_bisected, twice_sz=sector
    )
    records = run_sweeps(gamma_ep, energy_ep)
    output_dir = Path(__file__).resolve().parent
    csv_path = output_dir / "phase_ii_ep_dynamics.csv"
    plot_path = output_dir / "phase_ii_ep_dynamics.png"
    write_csv(records, csv_path)
    plot_summary(records, plot_path)

    print(f"gamma_EP={gamma_ep:.14f}, E_EP={energy_ep:.14f}, 2*Sz={sector}")
    print(f"bisection bracket=[{bracket[0]:.14f}, {bracket[1]:.14f}], double-root residual={residual:.3e}")
    for row in records:
        print(
            f"{row['direction']:>4} init={row['initial_branch']:<5} speed={row['speed']:.3g} "
            f"{row['final_branch_1']}={row['final_branch_1_weight']:.6f} "
            f"{row['final_branch_2']}={row['final_branch_2_weight']:.6f} "
            f"pair_share={row['ep_pair_coordinate_share']:.6f} "
            f"ray_fidelity={row['final_ray_fidelity_between_initial_branches']:.6f} "
            f"log_norm={row['log_norm_gain']:+.6f}"
        )
    print(f"CSV: {csv_path}")
    print(f"Plot: {plot_path}")