"""Optimise the sum-profile h for F = h(t_1+...+t_k) prod g(t_i) on lattice-aligned cells (spacing DELTA),
so that the result can be certified by hcert.py.  Densities from the band-limited Fourier sum (gexp.Model).
Returns integer edges E (units of DELTA; E[0] = 0, h = 0 on the first cell [0, E_1)) and h rounded to 2^-40, max h = 1."""
import sys, json, math
import numpy as np
import os; sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from gexp import Model

DELTA = 2.0 ** -17

def optimise(k, theta, beta, eps_cells=0, ncell=200, ptail=None, M=None, DELTA=DELTA):
    if M is None:
        M = Model(k, theta, beta)
    fk = M.density(k); fk1 = M.density(k - 1)
    hs = M.h; s = np.arange(M.N) * hs
    if ptail is None:
        ptail = 1e-9 if k < 1e7 else 1e-6
    cdf = np.cumsum(fk) * hs
    s0 = float(s[np.searchsorted(cdf, ptail)])
    ONE = int(round(1 / DELTA))
    top = ONE + eps_cells; jtop = ONE - eps_cells
    i0 = int(math.floor(s0 / DELTA))
    E = np.unique(np.round(np.linspace(i0, top, ncell + 1)).astype(np.int64))
    edges = E * DELTA
    nc = len(E) - 1
    m1 = (s <= jtop * DELTA) & (s >= edges[0] - M.T)
    sv = s[m1]; w1 = fk1[m1] * hs
    A = np.zeros((nc, nc))
    CH = 50000
    for a in range(0, len(sv), CH):
        ss = sv[a:a + CH]; ww = w1[a:a + CH]
        Gv = M.Gt(edges[:, None] - ss[None, :])
        Mx = Gv[1:] - Gv[:-1]
        A += (Mx * ww) @ Mx.T
    m0 = (s <= top * DELTA) & (s >= edges[0])
    cell = np.clip(np.searchsorted(edges, s[m0], side="right") - 1, 0, nc - 1)
    Dg = np.bincount(cell, weights=fk[m0] * hs, minlength=nc)
    keep = Dg > 1e-14 * Dg.max()
    Wd = np.where(keep, 1 / np.sqrt(np.where(keep, Dg, 1)), 0.0)
    B = (A * Wd[:, None]) * Wd[None, :]
    lam, V = np.linalg.eigh(B)
    hv = V[:, -1] * Wd
    hv = hv / hv[np.argmax(np.abs(hv))]
    hv = np.clip(np.round(hv * 2.0 ** 40) / 2.0 ** 40, 0.0, 1.0)
    Q = k / M.m2 * (hv @ A @ hv) / (hv ** 2 @ Dg)
    Efull = np.concatenate([[0], E]); hfull = np.concatenate([[0.0], hv])
    return dict(k=k, theta=theta, beta=beta, eps_cells=eps_cells, E=Efull, h=hfull, Q_fft=Q)

if __name__ == "__main__":
    k, th, be = int(sys.argv[1]), float(sys.argv[2]), float(sys.argv[3])
    r = optimise(k, th, be)
    print(json.dumps(dict(k=k, Q_fft=r["Q_fft"], nc=len(r["h"]), hmin=float(r["h"][1:].min()))))