#!/usr/bin/env python3
"""
Limenka PID gain tuner.

Sweeps (kp, ki, kd, window, ebasis) over the scenarios and prints a
scoreboard.  The tuning goals, in priority order:

  1. bounded difficulty noise in steady state (no limit cycle, d wanders
     < ~+-20% with no hashrate forcing)
  2. fast tracking: time for d to reach 2x on a 3x hashrate step
  3. flood resistance: low d damage from a 30-min 50x flood and a short
     stall window after exit
  4. issuance parity (should be ~1.000 by construction; verify)

Run:  python3 tune_mtp.py --ebasis direct
"""
import argparse
import math
import os
import random
import sys

sys.path.insert(0, os.path.dirname(__file__))
import mtp_sim as M


class Cfg:
    def __init__(self, window, future, kp, ki, kd, alpha, lp, fill, ebasis,
                 seed, blocks):
        self.window = window
        self.future = future
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.alpha = alpha
        self.lp = lp
        self.fill = fill
        self.ebasis = ebasis
        self.seed = seed
        self.blocks = blocks
        self.scenario = "steady"
        self.drop = 0.9
        self.flood_x = 50
        self.flood_min = 30
        self.stall_frac = 0.001
        self.stall_days = 4


def steady_wander(cfg):
    """log-std of d in steady state (no forcing) after warmup."""
    rng = random.Random(cfg.seed)
    world = M.World()
    world.const(0, 1.0)
    chain = M.Chain(cfg)
    chain.add_block(0.0, False, seed=True)
    for _ in range(cfg.blocks):
        h, a, _mode = world.rates_at(chain.t)
        dt = rng.expovariate(h / (M.T * chain.d))
        chain.add_block(chain.t + dt, False)
    dlogs = [math.log(b[3]) for b in chain.blocks[cfg.blocks // 3:]]
    mean = sum(dlogs) / len(dlogs)
    var = sum((x - mean) ** 2 for x in dlogs) / len(dlogs)
    return math.sqrt(var), chain.d


def step_time(cfg, step_x=3.0):
    """hours for d to reach 2.0 after a step to step_x at 20 days."""
    rng = random.Random(cfg.seed)
    world = M.World()
    step_t = 20 * 86400
    world.const(0, 1.0)
    world.const(step_t, step_x)
    chain = M.Chain(cfg)
    chain.add_block(0.0, False, seed=True)
    t50 = None
    for _ in range(cfg.blocks):
        h, a, _mode = world.rates_at(chain.t)
        dt = rng.expovariate(h / (M.T * chain.d))
        arr = chain.t + dt
        nxt = world.next_boundary(chain.t)
        if nxt is not None and arr >= nxt:
            chain.advance_time(nxt)
            continue
        chain.add_block(arr, False)
        if chain.t > step_t and t50 is None and chain.d >= 2.0:
            t50 = (chain.t - step_t) / 3600
    return t50


def flood_metrics(cfg):
    """d damage, attacker reward, stall window for a 30-min 50x flood."""
    rng = random.Random(cfg.seed)
    world = M.World()
    exit_t = 100 * M.T + cfg.flood_min * 60
    world.const(0, 1.0)
    world.const(100 * M.T, 1.0, cfg.flood_x)
    world.const(exit_t, 1.0)
    chain = M.Chain(cfg)
    chain.add_block(0.0, False, seed=True)
    for _ in range(cfg.blocks):
        h, a, _mode = world.rates_at(chain.t)
        total = h + a
        if total <= 0:
            break
        dt = rng.expovariate(total / (M.T * chain.d))
        arr = chain.t + dt
        nxt = world.next_boundary(chain.t)
        if nxt is not None and arr >= nxt:
            chain.advance_time(nxt)
            continue
        chain.add_block(arr, rng.random() < a / total)
    d_max = max(b[3] for b in chain.blocks)
    post = [b for b in chain.blocks if b[0] >= exit_t and not b[1]]
    restore = None
    if post:
        good = 0
        for b in post:
            if abs(b[2] - M.T) <= 0.5 * M.T:
                good += 1
            else:
                good = 0
            if good >= 3:
                restore = (b[0] - exit_t) / 3600
                break
    return d_max, chain.attacker_reward, restore


def main():
    ap = argparse.ArgumentParser(description="Limenka PID tuner")
    ap.add_argument("--window", type=int, default=101)
    ap.add_argument("--ebasis", choices=["mtp", "direct"], default="direct")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("--blocks", type=int, default=9000)
    args = ap.parse_args()

    gains = [
        (5e6, 5e6, 0),
        (10e6, 10e6, 0),
        (20e6, 10e6, 0),
        (50e6, 10e6, 0),
        (50e6, 20e6, 0),
        (20e6, 20e6, 0),
        (100e6, 20e6, 0),
        (50e6, 10e6, 16.7e6),
    ]
    print(f"window={args.window} ebasis={args.ebasis} seed={args.seed}")
    print(f"{'kp':>6} {'ki':>6} {'kd':>6} | {'wander':>6} {'step_t50':>8} | "
          f"{'flood_dmax':>10} {'flood_rew':>9} {'restore':>7}")
    for kp, ki, kd in gains:
        cfg = Cfg(args.window, 60, kp, ki, kd, 0.10, 0.15, 6600,
                  args.ebasis, args.seed, args.blocks)
        wander, _ = steady_wander(cfg)
        t50 = step_time(cfg)
        dmax, arew, restore = flood_metrics(cfg)
        def f(x):
            return "   nan" if x is None else f"{x:7.2f}"
        print(f"{kp / 1e6:6.1f}M {ki / 1e6:6.1f}M {kd / 1e6:6.1f}M | "
              f"{wander:6.3f} {f(t50)} | {dmax:10.2f} {arew:9.2f} {f(restore)}")


if __name__ == "__main__":
    main()