"""parity.py -- the kind means of residual.py read as counts, log L(1)
and log D.

QUESTION. residual.py reads rho_chi = [PI(x, chi) + m(loglog x + gamma)
- log L*(0, chi)]/PI(x, 1) per class character of a real field and finds
the odd mean falling with log|D| at 250 (R3) and the even rising. By
orthogonality the odd mean of PI(x, chi) is PI_1 - PI_k0, a count of the
trivial class against k0; by the functional equation log L*(0, chi) is
(1/2) log D + log L(1, chi) + c. Which of the three carries the slope?

THE DERIVATION (before code). chi(k0) = +-1, h+/2 odd characters and
h+/2 - 1 even nontrivial (even'):
  (a) mean_odd rho = (resT - resK)/h+, mean_even' rho = (resT + resK)/
      (h+ - 2), and with no k0 mean rho = resT/(h+ - 1), resT and resK
      the trivial and k0 cells of residual.py: the kind means restate the
      cells.
  (b) mean_odd PI(x, chi) = PI_1 - PI_k0, mean_even' PI(x, chi) =
      ((h+/2)(PI_1 + PI_k0) - PI)/(h+/2 - 1), with no k0 (h+ PI_1 -
      PI)/(h+ - 1).
  (c) log L*(0, chi) = (1/2) log D + log L(1, chi) + c, c_o = -2 log pi,
      c_e = -log 4 (Lambda(s) = D^(s/2) Gamma_R(s + a)^2 L(s), root
      number 1), so per field
      PI mean_odd rho = (PI_1 - PI_k0) - B_o - (1/2) log D + 2 log pi,
      B the kind's mean of log L(1, chi).
  (d) At a fixed cut the counts are bounded and B = o(log D)
      (Brauer-Siegel over the Galois fields H+, H, known; with no k0,
      H = H+ and the even' sum is over zeta_H+/zeta_K), so every kind
      mean is -(1/2 + o(1)) log D/PI as D -> oo.
  (e) R3's slope is the field-weighted OLS of mean_odd rho (weight h+/2)
      and splits exactly: y = A + B + C, A = (PI_1 - PI_k0)/PI, B =
      -B_o/PI, C = (-(1/2) log D + 2 log pi)/PI. Hand: PI(250, 1) about
      58, s_C about -0.0088; R3 -0.0035; the gap about +0.0053, +0.31
      counts per unit log D. By hand none of the counts trends with D
      (an inert p gives +1 at every D, a reduced split P lies in 1 or k0
      as +p or -p sits in the principal cycle, alternating), so the
      guess is that B carries it.

PREDICTIONS, frozen before the engine.
  P1 (b) per field in rho units (divided by PI), odd and even', at 250
     and 10^4, to 1e-10; (a) the same.
  P2 s_C at 250 within -0.0110 to -0.0070, and s_y equal to the
     per-character OLS slope (R3) to 1e-12.
  P3 printed: s_A, s_B, s_C with se, the gap in counts; the even' and
     no-k0 means alike; PI_1 - PI_k0, B_o by |D| band.
  KILL COUNTS. "The counts carry the gap": fires if s_A < 0.5 (s_y -
     s_C). The hand guess is that it fires.
  C1 THE BREAK: in (b)'s odd identity k0's count replaced by another
     class of order 2's (2-rank >= 2): fails by more than 1e-3 at every
     such field.
  C2 THE KILL'S CONTROL, read before the verdict: A permuted across
     fields fires KILL COUNTS in at least 190 of 200 seeds; A permuted
     plus a planted slope equal to the gap misses in at least 190.

CHANGED BEFORE THE RUN, nothing read: KILL COUNTS fires iff s_A < s_B
(half of s_A + s_B), so C2 as frozen tested the data's signs, not the
kill. C2 now permutes A and B both, killing both trends, and plants the
data's gap s_A + s_B on B (the kill fires in at least 190 of 200) or on
A (it misses in at least 190). After the rehearsal, C1 took the
largest gap over the cuts per field, as "fails at every such field"
means; it had taken the least, which a cut with PI_k0 = PI_c zeroes.

FINDINGS, the real fields with h+ >= 2 to |D| <= 48000 (12599; 10323
with k0), cut 250 unless named.
  F0  P1 held: the counts' identities to 1.27e-14, the restatement to
      1.59e-14 (rho units); C1 breaks it at 5771 fields, least gap
      0.0090; C2 fires 200/200 with the gap on B, misses 200/200 on A.
  F1  THE SPLIT (P2 held: s_y = R3 = -0.003535; s_C -0.00832, <PI>
      61.36). Odd: s_y -0.00353 +- 0.00065 = counts +0.00508 +- 0.00068
      + log L(1) -0.00030 +- 0.00007 + log D -0.00832 +- 0.00006. KILL
      COUNTS MISSED: the counts carry the gap and more (+0.29 per unit
      log D in PI_1 - PI_k0 against B_o's +0.017); the hand guess that
      log L(1) carries it was wrong. Even': +0.00944 = +0.01915 -
      0.00150 - 0.00821; no k0: +0.00896 = +0.01879 - 0.00144 - 0.00838.
      At 10^4 every s_y within 0.0006 of 0, the counts' part +0.0007 to
      +0.0010, log D's -0.0004; the count per unit log D +0.88 (odd),
      +0.90 (even'), +1.26 (no k0), past (1/2) log D's pace.
  F2  BY |D| BAND (odd, character-weighted): PI_1 - PI_k0 +2.53, +3.17,
      +3.38, +3.41, +3.53 at |D| to 4000, 8000, 16000, 32000, 48000,
      against (1/2) log D 3.73, 4.34, 4.69, 5.04, 5.30 and PI near 61.4
      throughout; B_o +0.165 to +0.216. The count gains 0.640 from the
      first band to the second, against (1/2) log D's 0.608, and 0.366
      from the second to the last, against 0.953. The odd mean reads
      +0.0151, +0.0162, +0.0137, +0.0084, +0.0060, its one clear fall
      from the third band to the fourth (about 3 se by hand; no band se
      is printed), still far above (d)'s -(1/2) log D/PI, -0.061 to
      -0.086 by band (by hand from the band means).
      The parts' band means, means of ratios, rebuild the mean's only to
      within 0.0011, too coarse to carry the first two values' rise.
  TIERS. (a), (b) properties; (d) a theorem given Brauer-Siegel; F1, F2
  observations over this population; why the count rises with |D| is
  not derived.

RUN RECORD. A rehearsal at --d 600: 3/5 (C1 as the least over cuts, C2
at 110 fields 180 and 178 of 200). Full: 5/5, 92.6 s, peak commit
149.3 MB under a memory guard (512 MB). From code/: 5/5, 92.2 s, 149.3 MB,
every print as above.
"""

import os
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

import sys
import time
from math import log, pi

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import numpy as np  # noqa: E402
import residual as R  # noqa: E402  (parses --d, --dreal from argv)
from module_law import check, section, CHECKS  # noqa: E402

P, O = R.P, R.O
CUTS = P.CUTS
J250, J10K = R.J250, R.J10K
SEEDS = 200
BANDS = ((0, 4000), (4000, 8000), (8000, 16000), (16000, 32000),
         (32000, 10 ** 9))
C_O, C_E = -2 * log(pi), -log(4)


def read():
    """Per real field with h+ >= 2 the kind means and their parts, per
    cut; and the identities' worst gaps."""
    rows, worst, brk = [], {"b": 0.0, "a": 0.0}, []
    rho_o_chars = []                       # (log D, rho_o at 250)
    for D in range(5, R.REAL_D + 1):
        if not P.fundamental(D):
            continue
        G = P.Group(D)
        h = G.h
        if h < 2:
            continue
        _, cor = P.counts(G)
        cor = np.array(cor)                # (cuts, h)
        tot = cor.sum(axis=1)
        chis = O.characters(G)
        m, L, _ = O.lstar(G, chis)
        nt = R.nontrivial(chis)
        lx = np.log(np.log(np.array(CUTS, float))) + O.GAMMA
        PIchi = np.real(cor @ chis.T)      # (cuts, h)
        lL0 = np.log(np.abs(L))
        rho = (PIchi + np.outer(lx, m) - lL0[None, :]) / tot[:, None]
        hask0 = G.k0 != G.e
        resT = rho[:, nt].sum(axis=1)
        lD = log(D)
        row = {"D": D, "lD": lD, "h": h, "k0": hask0, "PI": tot}
        if hask0:
            ck = np.real(chis[:, G.k0])
            odd = [i for i in nt if ck[i] < 0]
            ev = [i for i in nt if ck[i] > 0]
            resK = (ck[nt][None, :] * rho[:, nt]).sum(axis=1)
            d_odd = cor[:, G.e] - cor[:, G.k0]
            worst["b"] = max(worst["b"], np.max(np.abs(
                PIchi[:, odd].mean(axis=1) - d_odd) / tot))
            yo = rho[:, odd].mean(axis=1)
            worst["a"] = max(worst["a"], np.max(np.abs(
                yo - (resT - resK) / h)))
            Bo = (lL0[odd] - 0.5 * lD - C_O).mean()
            row["o"] = (len(odd), yo, d_odd / tot, -Bo / tot,
                        (-0.5 * lD - C_O) / tot, d_odd, Bo)
            for i in odd:
                rho_o_chars.append((lD, rho[J250, i]))
            others = [c for c in G.two if c not in (G.e, G.k0)]
            if others:
                c = others[0]
                brk.append(np.max(np.abs(PIchi[:, odd].mean(axis=1)
                                         - (cor[:, G.e] - cor[:, c]))
                                  / tot))
            if ev:
                ne = ((h / 2) * (cor[:, G.e] + cor[:, G.k0]) - tot) / (h / 2 - 1)
                worst["b"] = max(worst["b"], np.max(np.abs(
                    PIchi[:, ev].mean(axis=1) - ne) / tot))
                ye = rho[:, ev].mean(axis=1)
                worst["a"] = max(worst["a"], np.max(np.abs(
                    ye - (resT + resK) / (h - 2))))
                Be = (lL0[ev] - 0.5 * lD - C_E).mean()
                row["e"] = (len(ev), ye, ne / tot, -Be / tot,
                            (2 * lx - 0.5 * lD - C_E) / tot, ne, Be)
        else:
            nn = (h * cor[:, G.e] - tot) / (h - 1)
            worst["b"] = max(worst["b"], np.max(np.abs(
                PIchi[:, nt].mean(axis=1) - nn) / tot))
            yn = rho[:, nt].mean(axis=1)
            worst["a"] = max(worst["a"], np.max(np.abs(yn - resT / (h - 1))))
            Bn = (lL0[nt] - 0.5 * lD - C_E).mean()
            row["n"] = (len(nt), yn, nn / tot, -Bn / tot,
                        (2 * lx - 0.5 * lD - C_E) / tot, nn, Bn)
        rows.append(row)
    return rows, worst, brk, np.array(rho_o_chars)


def wslope(x, y, w):
    """Weighted OLS slope of y on x and its field-level (HC0) se."""
    xm = np.sum(w * x) / np.sum(w)
    ym = np.sum(w * y) / np.sum(w)
    sxx = np.sum(w * (x - xm) ** 2)
    s = np.sum(w * (x - xm) * (y - ym)) / sxx
    e = y - ym - s * (x - xm)
    return s, np.sqrt(np.sum((w * (x - xm) * e) ** 2)) / sxx


def kind_arrays(rows, key, j):
    sel = [r for r in rows if key in r]
    x = np.array([r["lD"] for r in sel])
    w = np.array([r[key][0] for r in sel], float)
    y, A, B, C = (np.array([r[key][k][j] for r in sel]) for k in (1, 2, 3, 4))
    cnt = np.array([r[key][5][j] for r in sel])
    BL = np.array([r[key][6] for r in sel])
    PI = np.array([r["PI"][j] for r in sel])
    return x, w, y, A, B, C, cnt, BL, PI


def kills(sA, sB):
    return sA < 0.5 * (sA + sB)


def main():
    t0 = time.time()
    section("A  THE IDENTITIES")
    rows, worst, brk, ro = read()
    print(f"  {len(rows)} real fields with h+ >= 2, |D| <= {R.REAL_D}"
          f"  [{time.time() - t0:.0f} s]")
    check("P1 (b) the counts: the kind means of PI(x, chi), in rho units",
          worst["b"] < 1e-10, f"worst {worst['b']:.2e}")
    check("P1 (a) the restatement: the kind means of rho are the cells",
          worst["a"] < 1e-10, f"worst {worst['a']:.2e}")
    check("C1 the break: another class of order 2 in place of k0",
          len(brk) > 0 and min(brk) > 1e-3,
          f"{len(brk)} fields, least gap {min(brk) if brk else 0:.4f}")

    section("B  THE SLOPE ON log D, SPLIT")
    out = {}
    for key, name in (("o", "odd"), ("e", "even'"), ("n", "no k0")):
        for j in (J250, J10K):
            x, w, y, A, B, C, cnt, BL, PI = kind_arrays(rows, key, j)
            sy, ey = wslope(x, y, w)
            parts = [wslope(x, v, w) for v in (A, B, C)]
            sA, sB, sC = (p[0] for p in parts)
            add = abs(sy - (sA + sB + sC))
            pm = np.sum(w * PI) / np.sum(w)
            print(f"  {name:6s} cut {CUTS[j]:5d}: {len(x)} fields, {int(w.sum())}"
                  f" characters, <PI> {pm:.2f}; s_y {sy:+.5f}+-{ey:.5f} ="
                  f" counts {sA:+.5f}+-{parts[0][1]:.5f} + log L(1)"
                  f" {sB:+.5f}+-{parts[1][1]:.5f} + log D {sC:+.5f}"
                  f"+-{parts[2][1]:.5f}  (adds to {add:.1e})")
            gc = wslope(x, cnt, w)[0]
            gb = wslope(x, BL, w)[0]
            print(f"         in counts per unit log D: the count {gc:+.4f},"
                  f" B {gb:+.4f}; the gap s_y - s_C {sy - sC:+.5f} x <PI>"
                  f" = {(sy - sC) * pm:+.3f}")
            out[(key, j)] = (x, w, y, A, B, C, cnt, BL, PI, sy, sA, sB, sC)

    x, w, y, A, B, C, cnt, BL, PI, sy, sA, sB, sC = out[("o", J250)]
    X = np.column_stack([np.ones(len(ro)), ro[:, 0]])
    r3 = np.linalg.lstsq(X, ro[:, 1], rcond=None)[0][1]
    check("P2 s_y is the per-character OLS slope (R3)",
          abs(sy - r3) < 1e-12, f"{sy:+.6f} against {r3:+.6f}")
    print(f"  P2 s_C at 250 within -0.0110 to -0.0070: "
          f"{'HELD' if -0.0110 <= sC <= -0.0070 else 'MISSED'} ({sC:+.5f})")

    print("  P3 by |D| band, odd, character-weighted at 250:")
    for lo, hi in BANDS:
        s = (np.exp(x) > lo) & (np.exp(x) <= hi)
        if s.sum() < 2:
            continue
        ww = w[s]

        def mw(v):
            return np.sum(ww * v[s]) / ww.sum()
        print(f"    {lo:6d}-{min(hi, R.REAL_D):6d}: {s.sum():5d} fields;"
              f" PI_1 - PI_k0 {mw(cnt):+.3f}, B_o {mw(BL):+.3f},"
              f" (1/2) log D {mw(0.5 * x):.3f}, PI {mw(PI):.2f},"
              f" mean rho_o {mw(y):+.4f}")

    section("C  KILL COUNTS, ITS CONTROL FIRST")
    rng = np.random.default_rng(1662)
    gap = sA + sB
    xm = np.sum(w * x) / np.sum(w)
    f_perm = f_plant = 0
    for _ in range(SEEDS):
        Ap, Bp = rng.permutation(A), rng.permutation(B)
        tilt = gap * (x - xm)
        f_perm += kills(wslope(x, Ap, w)[0], wslope(x, Bp + tilt, w)[0])
        f_plant += kills(wslope(x, Ap + tilt, w)[0], wslope(x, Bp, w)[0])
    check("C2 the kill's control: the gap planted on B fires, on A misses",
          f_perm >= 190 and SEEDS - f_plant >= 190,
          f"on B fires {f_perm}/{SEEDS}, on A misses"
          f" {SEEDS - f_plant}/{SEEDS}")
    v = kills(sA, sB)
    print(f"  KILL COUNTS: s_A {sA:+.5f} against half the gap"
          f" {0.5 * gap:+.5f}: {'FIRED' if v else 'MISSED'}")

    print(f"\n{sum(CHECKS)}/{len(CHECKS)} checks PASS, "
          f"{time.time() - t0:.1f} s")
    raise SystemExit(0 if all(CHECKS) else 1)


if __name__ == "__main__":
    main()