#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""第111便 — f(4) を押す。第110便の機械（erdos1110f4.py）をそのまま import して、
  (1) 切り所の細かい走査（r_0 = 1 + max K∩[1,Z] で重複を除く）と候補表の p 上限の引き上げ
  (2) 政策：率基準 → ベルマン（対数尺度の格子上の値関数。ロールアウトを包含する）
を試し、効いたものだけ整数厳密で回す。挟み込みの式は第110便と同一（Rows.bracket／Kempner.head）。

  python3 erdos1111f4.py bench                  大きい p の Rows と float_table の費用
  python3 erdos1111f4.py scan  [pmax] [pstep]   細かい切り所 × (率基準, ベルマン) の見積り
  python3 erdos1111f4.py k4 log10Z policy [pmax] [pstep] [cap]   本番（厳密）。policy ∈ {rate, bell}
"""
import sys, time, math
import os
# [配布版の変更点 / changed for distribution] 元は書き手の機体の絶対パスを指していた箇所を、このディレクトリからの相対に変えた。計算の中身は無改変。
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
import erdos1110f4 as M
from erdos1110f4 import Rows, Supp, Kempner, S55, SC, T, float_table, show, rlo, rhi
import numpy as np

# ---------------- 候補表（p > 200 は pstep 刻み） ----------------
def build_cand(pmax=200, pstep=10, qmax=26, lim=10 ** 82, log=print):
    t0 = time.time()
    tab = float_table(3, min(pmax, 200), qmax, lim)
    for p in range(200 + pstep, pmax + 1, pstep):
        tab.update(float_table(p, p, qmax, lim))
    log(f"[候補表] 浮動小数 {len(tab)} 対 (p≤{pmax}, p>200 は {pstep} 刻み, q≤{qmax}) {time.time()-t0:.1f} 秒")
    t0 = time.time(); cand = []; sp = {}
    for (p, q), (r0, n0) in sorted(tab.items()):
        if p not in sp: sp[p] = Supp(p, qmax)
        cand.append((p, q, r0, sp[p].maxelt(q, r0), n0))
    log(f"[候補表] maxelt {time.time()-t0:.1f} 秒")
    return cand

# ---------------- 政策 ----------------
RMIN, RMAX = 1e-3, 1e4        # 第110便は [0.01, 1000]。ベルマンでは広げても害はない

class Bellman:
    """V(x) = max_c [ g_c(z) + V(log10(2z+1+t_c)) ]、x = log10 z を格子 [x0, x1] 刻み h で持つ。
    格子の外（x > x1）は V = 0（下界側に安全：政策が過小評価するだけ）。"""
    def __init__(self, cand, x0=29.0, x1=78.0, h=0.01):
        self.cand = cand
        self.P = np.array([c[0] for c in cand]); self.Q = np.array([c[1] for c in cand])
        self.t = np.array([float(c[3]) for c in cand]); self.n = np.array([float(c[4]) for c in cand])
        self.xs = np.arange(x0, x1 + h / 2, h); self.V = np.zeros(len(self.xs))
        lt = np.log10(self.t)
        for i in range(len(self.xs) - 1, -1, -1):
            x = self.xs[i]; z = 10.0 ** x
            ok = (lt >= x + math.log10(RMIN)) & (lt <= x + math.log10(RMAX))
            if not ok.any(): continue
            a = 2 * z + 1; t = self.t[ok]; n = self.n[ok]
            g = (n / t) * np.log1p(t / a)
            xn = np.log10(a + t)
            Vn = np.interp(xn, self.xs, self.V, right=0.0)
            self.V[i] = max(0.0, float((g + Vn).max()))
    def value(self, z): return float(np.interp(math.log10(z), self.xs, self.V, right=0.0))
    def pick(self, z):
        x = math.log10(z); a = 2 * z + 1; best = None
        for (p, q, r, t, n) in self.cand:
            if not (RMIN * z <= t <= RMAX * z): continue
            g = (n / t) * math.log1p(t / a)
            sc = g + self.value(a + t)
            if best is None or sc > best[0]: best = (sc, p, q, r, t, n, g)
        return best

def pick_rate(z, cand):
    return M.pick4(z, cand)

def chain(r0, cand, policy, nst=60, exact=False, cap=200000, log=None, rows_cache=None):
    """policy: 'rate' か Bellman インスタンス。exact なら各段を厳密に挟む。"""
    z = r0; lo = hi = 0; est = 0.0; stages = []
    for j in range(nst):
        b = pick_rate(z, cand) if policy == 'rate' else policy.pick(z)
        if b is None: break
        _, p, q, r, t, n, g = b
        a = 2 * z + 1; d = None; l = h = 0
        if exact:
            t1 = time.time()
            # p ごとに Rows(p, qmax_p) を一つ作って全 q に使う（LRU 2 本。メモリ 1 GB 以内）
            if rows_cache is not None and p in rows_cache: R = rows_cache.pop(p)
            else:
                qm = max(c[1] for c in cand if c[0] == p)
                R = Rows(p, qm)
            if rows_cache is not None:
                rows_cache[p] = R
                while len(rows_cache) > 2: rows_cache.pop(next(iter(rows_cache)))
            r = R.argmax(q); t = R.maxelt(r, q); n = R.cf(q, r)
            if p ** 3 <= 2 * cap: d = R.depth_for(r, cap, q)
            else: d = 2 if R.leaves(2, q, r) <= cap else 1
            l, h = R.bracket(d, a, q, r)
            lo += l; hi += h
            g = (n / t) * math.log1p(t / a)
            if log: log(f"  段{j+1:2d} B({p},{q},{r}) |B|={n:.3e} t=1e{math.log10(t):.2f} a=1e{math.log10(a):.2f} d={d} "
                        f"[{l/SC:.4e}, {h/SC:.4e}] 幅 {(h-l)/SC:.1e}  {time.time()-t1:.1f}s")
        assert a > 2 * z and t > 0
        est += g
        stages.append(dict(p=p, q=q, r=r, t=t, a=a, n=n, lo=l, hi=h, d=d, zprev=z, g=g))
        z = a + t
    return lo, hi, est, stages, z

# ---------------- 切り所の細かい走査 ----------------
def cutpoints(K, e0=30.0, e1=31.3, step=0.02):
    """Z = 10^e を細かく動かし、r_0 が変わるものだけ残す。"""
    seen = {}; e = e0
    while e <= e1 + 1e-9:
        Zk = int(10 ** e); hl, hh, r0 = K.head(Zk)
        if r0 not in seen: seen[r0] = (e, hl, hh)
        e += step
    return sorted((r0, e, hl, hh) for r0, (e, hl, hh) in seen.items())

def main():
    mode = sys.argv[1] if len(sys.argv) > 1 else 'bench'
    T0 = time.time()
    if mode == 'bench':
        for (p, q) in ((260, 22), (300, 22), (350, 21)):
            t0 = time.time(); R = Rows(p, q); t1 = time.time(); r0 = R.argmax(); t2 = time.time()
            mx = R.maxelt(r0); d = R.depth_for(r0, 200000); lv = R.leaves(d, q, r0); t3 = time.time()
            lo, hi = R.bracket(d, 10 ** 50, q, r0); t4 = time.time()
            mem = sum(x.bit_length() for x in R.packed) // 8 // 2 ** 20
            print(f"[bench] Rows({p},{q}) 詰め {t1-t0:.1f}s argmax {t2-t1:.1f}s maxelt+depth {t3-t2:.1f}s "
                  f"bracket(d={d}, 葉 {lv}) {t4-t3:.1f}s  packed {mem} MB  |B|={R.cf(q,r0):.3e} 幅/値={(hi-lo)/max(lo,1):.1e}", flush=True)
            del R
        t0 = time.time(); tab = float_table(300, 300, 26, 10 ** 82); print(f"[bench] float_table p=300: {time.time()-t0:.1f}s ({len(tab)} 対)")
        t0 = time.time(); tab = float_table(400, 400, 26, 10 ** 82); print(f"[bench] float_table p=400: {time.time()-t0:.1f}s ({len(tab)} 対)")
        print(f"経過 {time.time()-T0:.1f} 秒")
    elif mode == 'scan':
        pmax = int(sys.argv[2]) if len(sys.argv) > 2 else 200
        pstep = int(sys.argv[3]) if len(sys.argv) > 3 else 10
        K = Kempner(S55, 55, 60); Hlo, Hhi = K.full(60)
        print(f"H(K+1) ∈ [{show(Hlo,12)}, {show(Hhi,12)}]")
        cand = build_cand(pmax, pstep)
        t0 = time.time(); B = Bellman(cand); print(f"[ベルマン] 格子 {len(B.xs)} 点 {time.time()-t0:.1f} 秒")
        cps = cutpoints(K); print(f"[切り所] r_0 の異なる点 {len(cps)} 個 {time.time()-T0:.1f} 秒")
        # 第110便の切り所（10^30.5）を基準として先頭に
        print("| log10 Z | r_0 | 頭 lo | 率基準の尾 | 段 | ベルマン V(r_0) | ベルマン経路の尾 | 段 | 総和(ベル) − H_hi |")
        best = None
        for (r0, e, hl, hh) in cps:
            lo, hi, est_r, st_r, z_r = chain(r0, cand, 'rate')
            lo, hi, est_b, st_b, z_b = chain(r0, cand, B)
            tot = hl / SC + est_b
            print(f"| {e:.2f} | 1e{math.log10(r0):.4f} | {show(hl,10)} | {est_r:.5e} | {len(st_r)} | {B.value(r0):.5e} | {est_b:.5e} | {len(st_b)} | {tot - Hhi/SC:+.5e} |")
            if best is None or tot > best[0]: best = (tot, e, r0, est_r, est_b)
        print(f"[最良] log10 Z = {best[1]:.2f}  r_0 = {best[2]}  率基準 {best[3]:.5e} → ベルマン {best[4]:.5e}  (差 {best[4]-best[3]:+.3e})")
        # 最良の切り所での経路
        lo, hi, est_b, st_b, z_b = chain(best[2], cand, B)
        print("[ベルマン経路]", ' '.join(f"({s['p']},{s['q']})" for s in st_b))
        lo, hi, est_r, st_r, z_r = chain(best[2], cand, 'rate')
        print("[率基準経路]  ", ' '.join(f"({s['p']},{s['q']})" for s in st_r))
        print(f"経過 {time.time()-T0:.1f} 秒")
    elif mode == 'k4':
        e = float(sys.argv[2]); pol = sys.argv[3]
        pmax = int(sys.argv[4]) if len(sys.argv) > 4 else 200
        pstep = int(sys.argv[5]) if len(sys.argv) > 5 else 10
        cap = int(sys.argv[6]) if len(sys.argv) > 6 else 200000
        K = Kempner(S55, 55, 90); Hlo, Hhi = K.full(85)
        print(f"[H(K+1)] [{show(Hlo,15)}, {show(Hhi,15)}]")
        cand = build_cand(pmax, pstep)
        policy = 'rate'
        if pol == 'bell':
            t0 = time.time(); policy = Bellman(cand); print(f"[ベルマン] 格子 {len(policy.xs)} 点 {time.time()-t0:.1f} 秒")
        Zk = int(10 ** e); hl, hh, r0 = K.head(Zk)
        print(f"[頭] Zk=1e{e}  rmax = {r0} (1e{math.log10(r0):.4f})  頭 ∈ [{show(hl,15)}, {show(hh,15)}]  捨てた尾 ≤ {(Hhi-hl)/SC:.4e}")
        # 第110便の第 1 段の逐語と照合（同じ切り所・同じ政策なら一致するはず）
        cache = {}
        lo, hi, est, st, z = chain(r0, cand, policy, exact=True, cap=cap, log=print, rows_cache=cache)
        TL, TH = hl + lo, hh + hi
        print()
        print(f"[段] {len(st)} 段、最終 z = 1e{math.log10(z):.2f}、尾の厳密和 ∈ [{lo/SC:.6e}, {hi/SC:.6e}]（見積り {est:.4e}）")
        print(f"[倍加税] 全段 a_j = 2 r_(j−1) + 1 > 2 r_(j−1): {all(s['a'] == 2*s['zprev']+1 for s in st)}")
        print(f"[総和] 厳密な下界 = {show(TL,15)}")
        print(f"[総和] 厳密な上界 = {show(TH,15)}")
        print(f"[総和] 幅 = {(TH-TL)/SC:.3e}")
        print(f"[比較] 下界 − H(K+1) の上界 = {(TL-Hhi)/SC:+.4e}")
        REF110 = 44397534682905455503579557696352024575141216707677669324302415340948610579786889024136402729213136126
        print(f"[比較] 下界 − 第110便の下界 = {(TL-REF110)/SC:+.4e}   超えた: {TL > REF110}")
        tag = f"{sys.argv[2]}-{pol}-{pmax}"
        with open(f'erdos1111f4-stages-{tag}.txt', 'w') as f:
            f.write(f"# Zk=1e{e} policy={pol} pmax={pmax} pstep={pstep} rmax={r0} head_lo={hl} head_hi={hh} H_lo={Hlo} H_hi={Hhi}\n")
            f.write("# j p q r t a n d lo hi\n")
            for j, s in enumerate(st, 1):
                f.write(f"{j} {s['p']} {s['q']} {s['r']} {s['t']} {s['a']} {s['n']} {s['d']} {s['lo']} {s['hi']}\n")
            f.write(f"# TOTAL_LO={TL}\n# TOTAL_HI={TH}\n")
        print(f"経過 {time.time()-T0:.1f} 秒")

if __name__ == '__main__':
    main()