"""
absolute_clock_ode.py  (v3)
===========================
Differential equation for an *absolute-time clock* from two parallel
CPU threads, equipped with a small AI AUTOMATON that auto-calibrates
the gain λ and the Laurent regulariser ε.

Theory (continuous)::

    dA/dt = dC̄/dt  −  λ_n · δ(t) / [ ε_n² + δ(t)² ]

where (λ, ε) are NOT constants — they are emitted each accepted step by
a tiny Q-learning finite-state machine (the "AI") that observes:

    δ_n        = C₁ − C₂                       (inter-thread disagreement)
    A − wall   = drift of absolute clock vs system time
    |correction_n|                              (energy spent per step)
    variance of δ over a sliding window

and emits actions on a discrete {-0.10, -0.05, 0, +0.05, +0.10} axis for
Δλ (and a correlated but damped Δε).  An internal Q-table (5 × 5 = 25
entries, seeded by hand-coded priors) is updated online via::

    Q[s,a] ← Q[s,a] + α · ( r + γ·max_a' Q[s',a'] − Q[s,a] )

with reward::

    r = − |A − wall|  −  0.1·|correction|  −  0.01·|δ|       (all in sec.)

Run::

    python absolute_clock_ode.py                           # default λ=0.4 ε=1e-3
    python absolute_clock_ode.py --seconds 30 --lam 0.5
    python absolute_clock_ode.py --ai-off                 # disable the AI; fixed params
    python absolute_clock_ode.py --quiet                  # CSV only

CSV columns::

    step, sys_datetime_iso, abs_datetime_iso,
    C1, C2, delta, C_bar, correction, A,
    state, lam, eps, action, reward
"""
from __future__ import annotations

import argparse
import math
import multiprocessing as mp
import os
import random
import statistics
import sys
import time
from datetime import datetime, timedelta


# ──────────────────────────────────────────────────────────────────────────
#                          calibration constants
# ──────────────────────────────────────────────────────────────────────────
UNIT_TIME        = 0.500      # seconds — pole X (calibrated once)
LAMBDA_DEFAULT   = 0.40
EPSILON_DEFAULT  = 1e-3
LOG_PATH         = "abs_clock_log.csv"


# ──────────────────────────────────────────────────────────────────────────
#                          known compute workload
# ──────────────────────────────────────────────────────────────────────────
def known_compute(unit_index: int) -> float:
    """Deterministic CPU workload whose TRUE wall-clock cost ≈ UNIT_TIME."""
    acc = 0.0
    for i in range(4_000_000):
        acc += math.sin(i * 0.001 + unit_index) * math.cos(i * 0.0007)
        acc -= math.sqrt(i + 1.0) * 0.000001
    return acc


def thread_worker(tid: int, q: mp.Queue, stop_flag) -> None:
    while not stop_flag.value:
        t0 = time.perf_counter()
        known_compute(tid)
        t1 = time.perf_counter()
        q.put((tid, t0, t1))


# ══════════════════════════════════════════════════════════════════════════
#                    T H E   A I   A U T O M A T O N
# ══════════════════════════════════════════════════════════════════════════
class ClockwiseAutomaton:
    """
    A tiny online RL agent (5 states × 5 actions, 25 Q-values) that
    calibrates the absolute-clock ODE.

    STATE SPACE  (derived from a sliding window of observations)::

        CALM        — small δ, small drift, low oscillation
        DRIFT       — persistent |A − wall|, trend established
        NOISY       — δ oscillates without settling
        CRITICAL    — max |δ| > 5·ε   (close to the pole)
        CONVERGED   — |A − wall| < 1 ms AND |δ| < ε/2

    ACTION SPACE  (5 actions)::

        0 : Δλ = −0.10 , ε → 0.95·ε   (strong relax, tighten)
        1 : Δλ = −0.05 , ε → 0.99·ε   (soft relax)
        2 : Δλ =  0.00 , ε → 1.00·ε   (hold)
        3 : Δλ = +0.05 , ε → 1.01·ε   (soft push)
        4 : Δλ = +0.10 , ε → 1.05·ε   (strong push, soften)

    REWARD (per step)::

        r = − |A − wall|  −  0.1·|correction|  −  0.01·|δ|       (seconds)

    Q-UPDATE::

        Q[s,a] ← Q[s,a] + α · ( r + γ · max_a' Q[s',a'] − Q[s,a] )
    """

    CALM, DRIFT, NOISY, CRITICAL, CONVERGED = range(5)
    STATE_NAMES = ('CALM', 'DRIFT', 'NOISY', 'CRITICAL', 'CONVERGED')

    ACTIONS = (                                 # (Δλ, ε factor)
        (-0.10, 0.95),
        (-0.05, 0.99),
        ( 0.00, 1.00),
        (+0.05, 1.01),
        (+0.10, 1.05),
    )

    EPS_GREEDY = 0.15
    ALPHA      = 0.20
    GAMMA      = 0.85
    WINDOW     = 32

    # ── constructor ──────────────────────────────────────────────────────
    def __init__(self, lam: float = LAMBDA_DEFAULT, eps: float = EPSILON_DEFAULT):
        self.lam = lam
        self.eps = eps
        self._reset_history()
        self.state = self.CALM
        self.prev_state  = None
        self.prev_action = 2                      # default to "hold"
        self.state_change_count = 0
        self.total_reward = 0.0
        # Q-table: 5 states × 5 actions
        self.Q = {(s, a): 0.0 for s in range(5) for a in range(5)}
        self.last_transition_reason = 'init'
        self._seed_prior()

    def _reset_history(self):
        self.deltas = []
        self.drifts = []         # A − wall_t  (signed)
        self.corrs  = []

    def _push(self, d, dr, c):
        self.deltas.append(d); self.drifts.append(dr); self.corrs.append(c)
        if len(self.deltas) > self.WINDOW:
            self.deltas.pop(0); self.drifts.pop(0); self.corrs.pop(0)

    # ── hand-coded prior on the Q-table (warm start) ─────────────────────
    def _seed_prior(self):
        """Bake reasonable policy beliefs into Q so the agent doesn't flail at start."""
        self.Q[(self.CALM,       2)] = +0.30   # hold in calm
        self.Q[(self.CALM,       0)] = -0.40   # over-relaxing in calm wastes capacity
        self.Q[(self.CALM,       4)] = -1.50   # pushing in calm is wasteful
        self.Q[(self.DRIFT,      3)] = +0.50   # soft push in drift
        self.Q[(self.DRIFT,      4)] = +0.60   # strong push in drift
        self.Q[(self.DRIFT,      0)] = -1.00   # relaxing in drift is bad
        self.Q[(self.NOISY,      0)] = +0.40   # back off when noisy
        self.Q[(self.NOISY,      1)] = +0.20
        self.Q[(self.CRITICAL,   4)] = +0.80   # strong push near pole
        self.Q[(self.CRITICAL,   3)] = +0.30
        self.Q[(self.CONVERGED,  2)] = +0.30   # hold when everything's great

    # ── observe the world, transition state, reward, learn ────────────────
    def observe_and_learn(self, delta, A_drift, correction):
        """Call once per accepted paired sample. Returns action index chosen."""
        self._push(delta, A_drift, correction)
        if len(self.deltas) < 4:
            return 2  # not enough data yet → hold

        # ── features
        mean_abs_d  = sum(abs(d) for d in self.deltas) / len(self.deltas)
        mean_abs_dr = sum(abs(r) for r in self.drifts) / len(self.drifts)
        max_d       = max(abs(d) for d in self.deltas)
        recent      = self.deltas[-min(4, len(self.deltas)):]
        rng_recent  = (max(recent) - min(recent)) if recent else 0.0

        # ── state inference (with reason logging)
        prev_state = self.state
        if max_d > 5.0 * self.eps:
            new_state, reason = self.CRITICAL, "max|δ|>5ε"
        elif mean_abs_dr > 0.05:
            new_state, reason = self.DRIFT,    "persistent |A-wall|>50ms"
        elif rng_recent > 3 * max(mean_abs_d, 1e-9) and mean_abs_d > self.eps:
            new_state, reason = self.NOISY,    "δ oscillation"
        elif mean_abs_d < 0.5 * self.eps and mean_abs_dr < 1e-3:
            new_state, reason = self.CONVERGED, "|A-wall|<1ms & |δ|<ε/2"
        else:
            new_state, reason = self.CALM,      "no-pattern"

        self.state = new_state
        self.last_transition_reason = reason
        if new_state != prev_state:
            self.state_change_count += 1

        # ── reward
        r = -abs(A_drift) - 0.1 * abs(correction) - 0.01 * abs(delta)
        self.total_reward += r

        # ── Q-learning on the previous (state, action) pair
        self.prev_state = prev_state
        if self.prev_state is not None:
            target = r + self.GAMMA * max(self.Q[(new_state, a)] for a in range(5))
            old_q  = self.Q[(prev_state, self.prev_action)]
            self.Q[(prev_state, self.prev_action)] = old_q + self.ALPHA * (target - old_q)

        # ── pick next action
        if random.random() < self.EPS_GREEDY:
            action = random.randrange(5)
        else:
            action = max(range(5), key=lambda a: self.Q[(new_state, a)])
        self.prev_action = action

        d_lam, eps_factor = self.ACTIONS[action]
        self.lam = max(0.05, min(1.0, self.lam + d_lam))
        self.eps = max(1e-5, min(1e-2, self.eps * eps_factor))
        return action

    # ── introspection ─────────────────────────────────────────────────────
    def snapshot(self) -> dict:
        return {
            'state'   : self.STATE_NAMES[self.state],
            'state_ix': self.state,
            'lam'     : self.lam,
            'eps'     : self.eps,
            'action'  : self.prev_action,
            'reason'  : self.last_transition_reason,
            'q_row'   : [round(self.Q[(self.state, a)], 3) for a in range(5)],
            'n_changes': self.state_change_count,
            '∑reward' : round(self.total_reward, 3),
        }


# ══════════════════════════════════════════════════════════════════════════
#                    T H E   O B S E R V E R   (ODE driver)
# ══════════════════════════════════════════════════════════════════════════
def absolute_clock_observer(q: mp.Queue,
                            stop_flag,
                            lam_init:      float,
                            eps_init:      float,
                            start_datetime: datetime,
                            enable_ai:     bool = True) -> float:
    """
    Consumes the cross-thread queue, forms paired samples, and updates
    the absolute clock via::

        A_{n+1} = A_n + ΔC̄_n − λ · δ_n / (ε² + δ_n²)

    where (λ, ε) are held CONSTANT if enable_ai=False, OR updated by the
    ClockwiseAutomaton each step if enable_ai=True.
    """
    log = open(LOG_PATH, "w")
    log.write("step,sys_datetime_iso,abs_datetime_iso,C1,C2,delta,"
              "C_bar,correction,A,state,lam,eps,action,reward\n")

    auto  = ClockwiseAutomaton(lam_init, eps_init)
    lam   = lam_init
    eps   = eps_init

    A, step, last_C_bar = 0.0, 0, None
    pending = {0: None, 1: None}

    # ── banner ────────────────────────────────────────────────────────────
    sys.stdout.write("\n")
    sys.stdout.write(f"  ┌─ program START datetime  : {start_datetime.isoformat(timespec='milliseconds')}\n")
    sys.stdout.write(f"  ├─ λ initial               : {lam_init}\n")
    sys.stdout.write(f"  ├─ ε initial               : {eps_init}\n")
    sys.stdout.write(f"  ├─ T_ref (pole X)           : {UNIT_TIME:.3f} s\n")
    sys.stdout.write(f"  ├─ AI enabled              : {enable_ai}\n")
    sys.stdout.write(f"  └─ cores available          : {os.cpu_count()}\n\n")

    hdr = (f"  {'step':>4} │ {'SYSTEM':^11} │ {'ABSOLUTE':^11} │ "
           f"{'δ':>8} │ {'A':>9} │ {'STATE':^10} │ "
           f"{'λ':>5} │ {'ε×1e3':>6} │ {'act':>3} │ {'reward':>7}")
    sep = "  " + "─" * (len(hdr) - 2)
    sys.stdout.write(hdr + "\n" + sep + "\n")

    while not stop_flag.value:
        try:
            msg = q.get(timeout=0.01)
        except Exception:
            continue
        if msg is None:
            continue
        tid, _t0, t1 = msg
        pending[tid] = (t1 - _t0,)
        if pending[0] is None or pending[1] is None:
            continue
        C1, = pending[0];  C2, = pending[1]
        pending = {0: None, 1: None}

        # ── ODE (uses current λ, ε) ──────────────────────────────────────
        delta      = C1 - C2
        C_bar      = 0.5 * (C1 + C2)
        dC_bar     = 0.0 if last_C_bar is None else C_bar - last_C_bar
        correction = lam * delta / (eps * eps + delta * delta)
        A = A + dC_bar - correction
        last_C_bar = C_bar
        step += 1

        # ── AI loop ──────────────────────────────────────────────────────
        absolute_dt  = start_datetime + timedelta(seconds=A)
        A_drift      = (absolute_dt - datetime.now()).total_seconds()

        if enable_ai:
            action = auto.observe_and_learn(delta, A_drift, correction)
            lam, eps = auto.lam, auto.eps
            snap = auto.snapshot()
            state_name  = snap['state']
            state_ix    = snap['state_ix']
            reward_now  = -abs(A_drift) - 0.1 * abs(correction) - 0.01 * abs(delta)
        else:
            action       = 2
            state_name   = '—'
            state_ix     = -1
            reward_now   = 0.0

        # ── print ────────────────────────────────────────────────────────
        sys_dt   = datetime.now()
        sys_iso  = sys_dt.strftime("%H:%M:%S")    + f".{sys_dt.microsecond//1000:03d}"
        abs_iso  = absolute_dt.strftime("%H:%M:%S")+ f".{absolute_dt.microsecond//1000:03d}"

        sys.stdout.write(
            f"  {step:>4d} │ {sys_iso:^11} │ {abs_iso:^11} │ "
            f"{delta:>+8.4f} │ {A:>9.4f} │ {state_name:^10} │ "
            f"{lam:>5.3f} │ {eps*1e3:>6.3f} │ {action:>3d} │ "
            f"{reward_now:>+7.3f}\n"
        )

        # ── log row ──────────────────────────────────────────────────────
        log.write(
            f"{step},{sys_dt.isoformat(timespec='milliseconds')},"
            f"{absolute_dt.isoformat(timespec='milliseconds')},"
            f"{C1:.6f},{C2:.6f},{delta:.6f},{C_bar:.6f},"
            f"{correction:.6f},{A:.6f},"
            f"{state_name},{lam:.6f},{eps:.6f},{action},{reward_now:.6f}\n"
        )
        log.flush()

    sys.stdout.write(sep + "\n")
    final_dt = start_datetime + timedelta(seconds=A)
    sys.stdout.write(
        f"  ▸ final A         = {A:.6f} s\n"
        f"  ▸ final absolute  = {final_dt.isoformat(timespec='milliseconds')}\n"
        f"  ▸ paired samples  = {step}\n"
    )
    if enable_ai:
        sys.stdout.write(
            f"  ▸ AI state changes = {auto.state_change_count}\n"
            f"  ▸ AI ∑reward       = {auto.total_reward:.3f}\n"
        )
    log.close()
    return A


# ──────────────────────────────────────────────────────────────────────────
def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Two-thread absolute clock + AI auto-calibrator")
    p.add_argument("--lam",     type=float, default=LAMBDA_DEFAULT)
    p.add_argument("--eps",     type=float, default=EPSILON_DEFAULT)
    p.add_argument("--seconds", type=float, default=10.0)
    p.add_argument("--ai-off",  action="store_true",
                   help="disable the AI; (λ, ε) fixed")
    p.add_argument("--quiet",   action="store_true",
                   help="suppress per-step output, CSV only")
    return p.parse_args()


def main() -> int:
    args = _parse_args()
    start_datetime   = datetime.now()
    stop_flag        = mp.Value("b", False)
    q                = mp.Queue(maxsize=4096)
    p0 = mp.Process(target=thread_worker, args=(0, q, stop_flag), name="T0")
    p1 = mp.Process(target=thread_worker, args=(1, q, stop_flag), name="T1")
    p0.start(); p1.start()

    try:
        if args.quiet:
            import contextlib, io
            with contextlib.redirect_stdout(io.StringIO()):
                absolute_clock_observer(
                    q, stop_flag, args.lam, args.eps,
                    start_datetime, enable_ai=not args.ai_off
                )
            deadline = time.perf_counter() + args.seconds
            while time.perf_counter() < deadline and stop_flag.value == 0:
                time.sleep(0.05)
        else:
            absolute_clock_observer(
                q, stop_flag, args.lam, args.eps,
                start_datetime, enable_ai=not args.ai_off
            )
            deadline = time.perf_counter() + args.seconds
            while time.perf_counter() < deadline and stop_flag.value == 0:
                time.sleep(0.05)
    except KeyboardInterrupt:
        sys.stdout.write("\n  [ctrl-c] stopping…\n")
    finally:
        stop_flag.value = True
        for p in (p0, p1):
            p.join(timeout=2.0)
            if p.is_alive():
                p.terminate(); p.join(timeout=1.0)
        sys.stdout.write(f"  ✓ done. CSV → {LOG_PATH}\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())