# Co-regulation versus co-dysregulation (Butler 2011; Sbarra & Hazan 2008) in a two-variable arousal
# model. a1, a2 = each partner's physiological arousal above their own resting level (arbitrary
# units). On its own, arousal decays back with time constant tau = 10 minutes. Coupling:
#   da1/dt = -a1/tau + c*a2 + k*(a2 - a1) + stress1(t)     (and the same for a2 with 1 and 2 swapped)
#   c > 0: the other's arousal raises mine (mutual escalation, "co-dysregulation")
#   k > 0: each is pulled toward the other's state (a calmer partner helps me settle, "co-regulation")
# Partner 1 meets a stressor: arousal input of 1 per minute for 5 minutes. Illustrative numbers.
TAU = 10.0

def run(c, k, T=120.0, dt=0.01):
    a1 = a2 = 0.0; peak1 = peak2 = 0.0; t_rec = None; hist = []
    for i in range(int(T / dt)):
        t = i * dt
        s1 = 1.0 if t < 5 else 0.0
        d1 = -a1 / TAU + c * a2 + k * (a2 - a1) + s1
        d2 = -a2 / TAU + c * a1 + k * (a1 - a2)
        a1, a2 = a1 + d1 * dt, a2 + d2 * dt
        peak1, peak2 = max(peak1, a1), max(peak2, a2)
        if t > 5 and t_rec is None and a1 < 0.1 * peak1: t_rec = t
        if i % 100 == 0: hist.append((t, a1, a2))
    return peak1, peak2, t_rec, hist

print("coupling                            | peak 1 | peak 2 | partner 1 back below 10 % of peak")
for label, c, k in (("independent", 0.0, 0.0), ("co-regulation, k = 0.10/min", 0.0, 0.10),
                    ("mild mutual escalation, c = 0.05", 0.05, 0.0), ("strong mutual escalation, c = 0.09", 0.09, 0.0),
                    ("past the limit, c = 0.11 > 1/tau", 0.11, 0.0)):
    p1, p2, tr, h = run(c, k)
    end = h[-1][1]
    rec = (f"{tr - 5:.0f} min after the stressor ends" if tr else
           f"not within 2 hours; still {end:.2f} and rising" if end >= h[-2][1] else
           f"not within 2 hours; still {end:.2f}, falling very slowly")
    print(f"{label:35s} | {p1:6.2f} | {p2:6.2f} | {rec}")
# Linear analysis: the 'together' mode a1 + a2 decays at rate 1/tau - c; it grows once c > 1/tau.
print(f"escalation threshold: c = 1/tau = {1 / TAU:.2f} per minute; co-regulation (k) never destabilises")
# coupling                            | peak 1 | peak 2 | partner 1 back below 10 % of peak
# independent                         |   3.94 |   0.00 | 23 min after the stressor ends
# co-regulation, k = 0.10/min         |   3.26 |   0.93 | 18 min after the stressor ends
# mild mutual escalation, c = 0.05    |   3.97 |   0.96 | 35 min after the stressor ends
# strong mutual escalation, c = 0.09  |   4.05 |   2.01 | not within 2 hours; still 0.78, falling very slowly
# past the limit, c = 0.11 > 1/tau    |   8.10 |   8.10 | not within 2 hours; still 8.02 and rising
# escalation threshold: c = 1/tau = 0.10 per minute; co-regulation (k) never destabilises

# Measuring "physiological linkage" (Levenson & Gottman 1983) and why a shared situation can mimic
# it. Synthetic heart rates, one value per second, for a 15-minute conflict discussion (900 s).
#   influence  : partner 2's heart rate follows partner 1's with a 5-second delay
#   shared     : no influence; both respond to the same moment-to-moment heat of the topic
#   independent: neither
# Each series is first "pre-whitened" (its own AR(1) predictability removed), then cross-correlated
# at lags -15..+15 s. Significance: compare the largest correlation with the same statistic after
# circularly shifting partner 2's series by a random amount (500 shifts).
import math, random

def ar1(rng, n, phi, sd):
    x, out = 0.0, []
    for _ in range(n):
        x = phi * x + rng.gauss(0, sd); out.append(x)
    return out

def make(kind, rng, n=900):
    topic = ar1(rng, n, 0.98, 0.6)                      # slowly varying heat of the discussion
    h1 = [70 + v for v in ar1(rng, n, 0.9, 1.0)]
    if kind == "shared": h1 = [a + 1.5 * t for a, t in zip(h1, topic)]
    own2 = ar1(rng, n, 0.9, 1.0)
    if kind == "influence": h2 = [75 + o + 0.6 * (h1[i - 5] - 70 if i >= 5 else 0) for i, o in enumerate(own2)]
    elif kind == "shared": h2 = [75 + o + 1.5 * t for o, t in zip(own2, topic)]
    else: h2 = [75 + o for o in own2]
    return h1, h2

def whiten(x):
    m = sum(x) / len(x); x = [v - m for v in x]
    phi = sum(a * b for a, b in zip(x, x[1:])) / sum(a * a for a in x[:-1])
    return [b - phi * a for a, b in zip(x, x[1:])]

def xcorr(a, b, lag):
    if lag >= 0: a, b = a[:len(a) - lag], b[lag:]
    else: a, b = a[-lag:], b[:len(b) + lag]
    ma, mb = sum(a) / len(a), sum(b) / len(b)
    s = sum((x - ma) * (y - mb) for x, y in zip(a, b))
    return s / math.sqrt(sum((x - ma) ** 2 for x in a) * sum((y - mb) ** 2 for y in b))

def best(a, b):
    return max(((xcorr(a, b, L), L) for L in range(-15, 16)), key=lambda t: abs(t[0]))

rng = random.Random(11)
print("situation    | raw r at lag 0 | after pre-whitening: largest r, at lag | shift-test p")
for kind in ("influence", "shared", "independent"):
    h1, h2 = make(kind, rng)
    raw = xcorr(h1, h2, 0)
    w1, w2 = whiten(h1), whiten(h2)
    r, L = best(w1, w2)
    null = []
    for _ in range(500):
        s = rng.randrange(60, len(w2) - 60); null.append(abs(best(w1, w2[s:] + w2[:s])[0]))
    p = (1 + sum(v >= abs(r) for v in null)) / 501
    print(f"{kind:12s} | {raw:14.2f} | {r:+.2f} at {L:+d} s{'':22s} | {p:.3f}")
print("positive lag = partner 2 follows partner 1")
# situation    | raw r at lag 0 | after pre-whitening: largest r, at lag | shift-test p
# influence    |           0.34 | +0.51 at +5 s                       | 0.002
# shared       |           0.75 | +0.48 at +0 s                       | 0.002
# independent  |           0.01 | +0.07 at +5 s                       | 0.685
# positive lag = partner 2 follows partner 1