#!/usr/bin/env python3
"""
The Edge of Synchronization: collective recognition peaks at critical coupling.

A population of N Kuramoto phase oscillators tries, collectively, to track a
noisy external drive Omega(t) — "the true beat," the territory. Each oscillator
is only weakly and noisily coupled to the drive (epsilon); the oscillators are
coupled to *each other* with strength K. The question the framework asks:

    as internal coupling K rises from 0 (incoherent) to large (full lock),
    how well does the ENSEMBLE track the external drive?

Two failure modes bracket the answer, and naming them is the whole point:

  - K too low:  incoherent. Each oscillator is a noisy, private estimate of the
                drive; nothing pools them. The crowd is a mob, not a loop.
  - K too high: the oscillators lock to *each other*. The synchronized cluster
                behaves as one sluggish effective oscillator whose tiny external
                coupling epsilon can no longer pull it to follow a fluctuating
                drive. Internal consensus replaces external grounding. This is
                the collective Veer: a confident, coherent crowd tracking its
                own echo instead of the territory — groupthink as a phase
                transition.
  - K ~ K_c:    the edge of synchronization. Subgroups form and dissolve;
                the ensemble pools its noisy glimpses of the drive without
                fusing into a single deaf block. Collective tracking is best.

The framework's claim (this is the novel, falsifiable part): collective
recognition — the population's information about the drive, beyond what any
single oscillator carries — is *non-monotone* in K and peaks near K_c, not at
full synchronization. Full lock is the collective form of the equivalence case
the Asymmetric Synergy Bound calls sterile: everyone identical, redundant, and
in this dynamical version, decoupled from the territory.

This connects to a real literature (collective sensing / stochastic resonance
in coupled-oscillator populations); the framework's contribution is the
*reading*: the edge of synchronization is where the collective loop is alive.

Run:  python3 research/collective-loop/collective_loop.py
Deterministic under SEED.
"""

import numpy as np

SEED = 11
N = 120                 # oscillators
OMEGA_DRIVE = 0.0       # mean drive frequency (rotating frame)
CROWD_DETUNE = 0.40     # the crowd's intrinsic tempo, offset from the drive: the key tension.
                        # Below criticality the weak drive cannot organize the crowd; far above
                        # it the crowd locks to its OWN tempo and the weak drive cannot move the
                        # rigid block off it. Recognition lives in between.
DRIVE_NOISE = 0.0       # the drive is a steady external beat at frequency OMEGA_DRIVE
EPS = 0.28              # GLOBAL drive coupling — weak: the drive cannot entrain the crowd alone,
                        # and cannot move a rigidly self-locked crowd off its own tempo.
SIGMA_OMEGA = 0.35      # spread of natural frequencies (sets the critical coupling K_c ~ 0.6)
PHASE_NOISE = 0.15      # per-oscillator phase diffusion
DT = 0.05
T_TRANSIENT = 250.0     # burn-in time
T_SAMPLE = 2400.0       # sampling time
SAMPLE_EVERY = 4        # subsample steps for the time series we analyze


def circ_dist(a, b):
    """Wrapped phase difference in (-pi, pi]."""
    return (a - b + np.pi) % (2 * np.pi) - np.pi


def simulate(K, rng):
    """Integrate the driven Kuramoto system. Returns sampled time series."""
    n_trans = int(T_TRANSIENT / DT)
    n_samp = int(T_SAMPLE / DT)

    theta = rng.uniform(-np.pi, np.pi, N)
    # natural frequencies centered on the crowd's intrinsic tempo, not the drive
    omega = CROWD_DETUNE + SIGMA_OMEGA * rng.standard_normal(N)

    # the drive is an oscillator with a slowly drifting (random-walk) frequency:
    # gentle "territory" the collective is trying to track.
    drive_phase = 0.0
    drive_freq = OMEGA_DRIVE

    sdt = np.sqrt(DT)
    rec_psi, rec_drive, rec_single, rec_r = [], [], [], []

    for step in range(n_trans + n_samp):
        # drive evolves
        drive_freq += -0.05 * drive_freq * DT + DRIVE_NOISE * sdt * rng.standard_normal()
        drive_phase += drive_freq * DT

        # order parameter
        z = np.exp(1j * theta).mean()
        r = np.abs(z)
        psi = np.angle(z)

        # Kuramoto update: internal mean-field coupling + weak GLOBAL drive
        dtheta = (omega
                  + K * r * np.sin(psi - theta)
                  + EPS * np.sin(drive_phase - theta)) * DT
        dtheta += PHASE_NOISE * sdt * rng.standard_normal(N)
        theta = theta + dtheta

        if step >= n_trans and (step % SAMPLE_EVERY == 0):
            rec_psi.append(psi)
            rec_drive.append(drive_phase)
            rec_single.append(theta[0])
            rec_r.append(r)

    return (np.array(rec_psi), np.array(rec_drive),
            np.array(rec_single), np.array(rec_r))


def circ_mi(x, y, nbins=16):
    """Mutual information (bits) between two circular variables by binning."""
    edges = np.linspace(-np.pi, np.pi, nbins + 1)
    xb = np.clip(np.digitize(((x + np.pi) % (2 * np.pi)) - np.pi, edges) - 1, 0, nbins - 1)
    yb = np.clip(np.digitize(((y + np.pi) % (2 * np.pi)) - np.pi, edges) - 1, 0, nbins - 1)
    joint = np.zeros((nbins, nbins))
    for a, b in zip(xb, yb):
        joint[a, b] += 1
    joint /= joint.sum()
    px = joint.sum(1, keepdims=True)
    py = joint.sum(0, keepdims=True)
    nz = joint > 0
    return float((joint[nz] * np.log2(joint[nz] / (px @ py)[nz])).sum())


def analyze(K, rng):
    psi, drive, single, r = simulate(K, rng)

    # phase-locking value to the steady drive: how entrained is the collective
    # mean phase (and a single oscillator) to the external beat? PLV in [0,1].
    plv_collective = np.abs(np.exp(1j * (psi - drive)).mean())
    plv_single = np.abs(np.exp(1j * (single - drive)).mean())

    # collective frequency: does the ensemble rotate at the DRIVE's frequency
    # (entrained, ~OMEGA_DRIVE) or at the CROWD's own tempo (~CROWD_DETUNE)?
    omega_collective = np.mean(np.diff(np.unwrap(psi))) / (DT * SAMPLE_EVERY)

    return dict(K=K, r=r.mean(),
                plv_collective=plv_collective, plv_single=plv_single,
                omega_collective=omega_collective)


def main():
    Ks = np.round(np.concatenate([
        np.linspace(0.0, 1.5, 16),
        np.linspace(1.6, 4.0, 7),
    ]), 3)

    rows = []
    for K in Ks:
        rng = np.random.default_rng(SEED)   # same noise realization across K
        rows.append(analyze(K, rng))

    print("=" * 74)
    print("Driven Kuramoto collective — tracking a weak detuned drive (subthreshold)")
    print(f"N={N}, eps={EPS}, crowd_detune={CROWD_DETUNE}, sigma_omega={SIGMA_OMEGA}, "
          f"phase_noise={PHASE_NOISE}, drive_noise={DRIVE_NOISE}")
    print("=" * 74)
    hdr = f"{'K':>5} {'r':>6} {'PLV_coll':>9} {'PLV_one':>8} {'omega_coll':>11}"
    print(hdr); print("-" * len(hdr))
    for d in rows:
        print(f"{d['K']:5.2f} {d['r']:6.3f} {d['plv_collective']:9.3f} "
              f"{d['plv_single']:8.3f} {d['omega_collective']:11.3f}")

    plv = np.array([d['plv_collective'] for d in rows])
    rbar = np.array([d['r'] for d in rows])
    # the collective-Veer transition: the steepest collapse in drive entrainment
    k_star = Ks[1:][int(np.diff(plv).argmin())]
    print("-" * len(hdr))
    print(f"drive freq = {OMEGA_DRIVE}, crowd tempo = {CROWD_DETUNE}")
    print(f"collective-Veer transition K* ~ {k_star} "
          f"(entrainment collapses; crowd locks to its own tempo)")
    print(f"  below K*: PLV ~ {plv[Ks < k_star].mean():.2f} (entrained to the beat)")
    print(f"  above K*: PLV ~ {plv[Ks > k_star].mean():.2f} (locked to crowd tempo)")
    print(f"  internal sync r at K*: {rbar[Ks == k_star][0]:.2f}; "
          f"r saturates to {rbar.max():.2f}")

    figure(Ks, rows)


def figure(Ks, rows):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    r = np.array([d['r'] for d in rows])
    plv = np.array([d['plv_collective'] for d in rows])
    plv_s = np.array([d['plv_single'] for d in rows])
    om = np.array([d['omega_collective'] for d in rows])
    k_star = Ks[1:][int(np.diff(plv).argmin())]

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))

    ax = axes[0]
    ax.plot(Ks, plv, "o-", color="#c9a84c", label="ensemble (collective)")
    ax.plot(Ks, plv_s, "s--", color="#7aa0c4", label="single oscillator")
    ax.axvline(k_star, color="#b05050", lw=1, ls=":")
    ax.set_xlabel("internal coupling K")
    ax.set_ylabel("entrainment to the drive (PLV)")
    ax.set_title("Entrainment to the beat collapses\nabove a critical coupling K*")
    ax.legend(fontsize=9)

    ax = axes[1]
    ax.plot(Ks, r, "o-", color="#c9a84c")
    ax.axvline(k_star, color="#b05050", lw=1, ls=":")
    ax.annotate("K* (Veer)", (k_star, r[Ks == k_star][0]),
                textcoords="offset points", xytext=(10, -4), color="#b05050", fontsize=9)
    ax.set_xlabel("internal coupling K")
    ax.set_ylabel("internal synchronization r")
    ax.set_title("Internal sync rises smoothly through K*;\nthe loss of grounding is abrupt")

    ax = axes[2]
    ax.plot(Ks, om, "o-", color="#c9a84c")
    ax.axhline(OMEGA_DRIVE, color="#7aa0c4", lw=1, ls="--", label="drive tempo (the world)")
    ax.axhline(CROWD_DETUNE, color="#b05050", lw=1, ls="--", label="crowd's own tempo")
    ax.axvline(k_star, color="#b05050", lw=1, ls=":")
    ax.set_xlabel("internal coupling K")
    ax.set_ylabel("collective frequency")
    ax.set_title("The collective Veer: past K* the crowd\nlocks to its own tempo, not the world")
    ax.legend(fontsize=9)

    fig.tight_layout()
    fig.savefig("collective-loop-results.png", dpi=130)
    print("wrote collective-loop-results.png")


if __name__ == "__main__":
    main()