#!/usr/bin/env python3
"""
edge_phases.py -- complex single-flip amplitudes in the monitored tie model (Remark D, last sentence).

CLAIM.  Give every tie edge its own phase, K = sum_edges g (e^{i phi_e} |s^i><s| + h.c.).  Then the
pointer kernel is still pi-reversible, pi_a T_t(a->b) = pi_b T_t(b->a), even though Theta (complex
conjugation) is no longer a symmetry and the phases carry non-gaugeable flux around square plaquettes.

ARGUMENT CHECKED HERE, step by step.
  (i)   GNS adjoint:  pi_a T_t[H](a->b) = pi_b T_t[-H](b->a)   (needs only [H, pi] = 0 and GNS-symmetric jumps).
  (ii)  W = prod_i Z_i anticommutes with every single-flip term (the configuration hypercube is
        bipartite) and commutes with H0, the monitor and the thermal dissipator:  T_t[H0 + K] = T_t[H0 - K].
  (iii) H0 drops out of the population-coupled sector, because E_s - E_t is conserved there:
        T_t[H0 + K] = T_t[-H0 + K].
  Together:  T_t[-H] = T_t[H], so (i) is reversibility.  A two-spin exchange commutes with W, which is
  why a complex exchange around an odd loop breaks it (see reversibility.py, loop controls).

Reuses the model builders of reversibility.py (same directory).  numpy + scipy only.  No files are written.
"""
import os
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, "1")
import time
import numpy as np
import reversibility as rv

BETA, G, NU = 0.7, 1.0, 1.0
KAPPAS = (0.5, 2.0, 10.0)
SEEDS = (1, 2, 3)


def edge_phase_hamiltonian(m, g, rng, sign_H0=1.0, sign_K=1.0):
    """H0 + K with an independent random phase on every (unordered) tie edge."""
    H = np.diag(sign_H0 * m.E).astype(complex)
    phases = {}
    for a in range(m.d):
        for i in range(m.n):
            if m.h[a, i] == 0:
                b = a ^ (1 << i)
                e = (min(a, b), max(a, b))
                if e not in phases:
                    phases[e] = rng.uniform(-np.pi, np.pi)
                ph = phases[e] if a == e[0] else -phases[e]
                H[b, a] += sign_K * g * np.exp(1j * ph)
    return H, phases


def plaquette_flux(m, phases):
    """Largest |flux| (mod 2 pi) around 4-cycles s -> s^i -> s^ij -> s^j -> s made of tie edges."""
    def ph(a, b):
        e = (min(a, b), max(a, b))
        return phases[e] if a == e[0] else -phases[e]
    worst = 0.0
    for s in range(m.d):
        for i in range(m.n):
            for j in range(i + 1, m.n):
                cyc = [s, s ^ (1 << i), s ^ (1 << i) ^ (1 << j), s ^ (1 << j), s]
                es = [(cyc[k], cyc[k + 1]) for k in range(4)]
                if all((min(x, y), max(x, y)) in phases for x, y in es):
                    f = sum(ph(x, y) for x, y in es)
                    worst = max(worst, abs((f + np.pi) % (2 * np.pi) - np.pi))
    return worst


def kernels_for(m, H, jumps, kappa):
    L = rv.lindbladian(m, H, jumps, kappa)
    return [T.real for T in rv.kernels(rv.Sector(L, m), m.d)], L


def main():
    t0 = time.time()
    print("edge_phases.py -- random complex amplitude on every tie edge (betaJ=%.1f, g=%.1f, nu=%.1f)" % (BETA, G, NU))
    print("columns: max over t in %s and seeds; DB = pi-reversibility residual of the kernel (rel.)" % (rv.TIMES,))
    print(f"{'system':7s}{'conv':5s}{'kappa':>6s}{'edges':>6s}{'max flux':>9s}{'Tid(Theta)':>11s}"
          f"{'DB[T]':>10s}{'|T[H]-T[-H]|':>13s}{'|T[+K]-T[-K]|':>14s}{'|T[+H0]-T[-H0]|':>16s}{'|T[phi]-T[0]|':>14s}")
    print("-" * 111)
    worst = dict(db=0.0, sym=0.0, w=0.0, h0=0.0)
    for key in rv.GEOM:
        m = rv.Model(key)
        convs = ("R1",) if key == "S1" else ("R1", "R2")          # R1 == R2 for S1
        kappas = (2.0, 10.0) if key == "S3" else KAPPAS
        seeds = SEEDS[:1] if key == "S3" else SEEDS
        pw = np.exp(-BETA * (m.E - m.E.min()))
        pi = pw / pw.sum()
        for conv in convs:
            jumps = rv.thermal_jumps(m, BETA, NU, conv)
            for kappa in kappas:
                row = dict(db=0.0, sym=0.0, w=0.0, h0=0.0, real=0.0, tid=0.0, flux=0.0)
                for sd in seeds:
                    Hp, phases = edge_phase_hamiltonian(m, G, np.random.default_rng(sd))
                    Tp, L = kernels_for(m, Hp, jumps, kappa)
                    Tm, _ = kernels_for(m, -Hp, jumps, kappa)
                    Hk, _ = edge_phase_hamiltonian(m, G, np.random.default_rng(sd), sign_K=-1.0)
                    Tk, _ = kernels_for(m, Hk, jumps, kappa)
                    Hh, _ = edge_phase_hamiltonian(m, G, np.random.default_rng(sd), sign_H0=-1.0)
                    Th, _ = kernels_for(m, Hh, jumps, kappa)
                    T0, _ = kernels_for(m, rv.hamiltonian(m, G, np.ones(m.n), np.zeros(m.n)), jumps, kappa)
                    row["db"] = max(row["db"], max(rv.db_residual(T, pi) for T in Tp))
                    row["sym"] = max(row["sym"], max(float(np.max(np.abs(a - b))) for a, b in zip(Tp, Tm)))
                    row["w"] = max(row["w"], max(float(np.max(np.abs(a - b))) for a, b in zip(Tp, Tk)))
                    row["h0"] = max(row["h0"], max(float(np.max(np.abs(a - b))) for a, b in zip(Tp, Th)))
                    row["real"] = max(row["real"], max(float(np.max(np.abs(a - b))) for a, b in zip(Tp, T0)))
                    row["tid"] = max(row["tid"], rv.identity_residuals(L, pi, m.d)[0])
                    row["flux"] = max(row["flux"], plaquette_flux(m, phases))
                for k in worst:
                    worst[k] = max(worst[k], row[k])
                print(f"{key:7s}{conv:5s}{kappa:6.1f}{len(phases):6d}{row['flux']:9.3f}{row['tid']:11.1e}"
                      f"{row['db']:10.1e}{row['sym']:13.1e}{row['w']:14.1e}{row['h0']:16.1e}{row['real']:14.1e}")
    print("-" * 111)
    print("Tid(Theta) = residual of identity (7) with plain complex conjugation: large, so Theta is NOT the symmetry here.")
    print("|T[phi]-T[0]| = kernel change vs the real model: nonzero, so the phases matter (they carry plaquette flux).")
    ok = worst["db"] < 1e-11 and worst["sym"] < 1e-11 and worst["w"] < 1e-11 and worst["h0"] < 1e-11
    print(f"worst: DB {worst['db']:.1e}, T[H]=T[-H] {worst['sym']:.1e}, W-step {worst['w']:.1e}, H0-step {worst['h0']:.1e}"
          f"  ->  {'all checks pass' if ok else 'CHECK FAILED'}   ({time.time() - t0:.0f} s)")


if __name__ == "__main__":
    main()