"""Fast explorer: F = h(sum t) prod g(t_i), g(t) = 1/(c+bt) + A  on [0,T]  (A = shift, g > 0 required).
Densities of S_n by a band-limited Fourier sum (phi^n is negligible beyond xi_cut), FFT on a fine grid."""
import sys, json, math, time
import numpy as np
from scipy.special import exp1
from scipy.linalg import eigh

def F_parts(xi, c, b, T):
    """F_{-2}, F_{-1}, F_0 at xi != 0 :  int_0^T e^{i xi t} u^m dt, u = c + b t."""
    a = xi / b
    u1, u2 = c, c + b * T
    z1, z2 = -1j * a * u1, -1j * a * u2
    E1d = exp1(z1) - exp1(z2)
    pre = np.exp(-1j * a * c) / b
    Fm2 = pre * ((-np.exp(1j * a * u2) / u2 + np.exp(1j * a * u1) / u1) + 1j * a * E1d)
    Fm1 = pre * E1d
    F0 = (np.exp(1j * xi * T) - 1) / (1j * xi)
    return Fm2, Fm1, F0

class Model:
    def __init__(self, k, theta, beta, A=0.0, P=4.0, N=2 ** 22, cut=1e-22):
        L = math.log(k); c = theta / L; T = beta / L; b = k - 1
        self.k, self.c, self.T, self.b, self.A = k, c, T, b, A
        u2 = c + b * T
        if 1 / u2 + A <= 0:
            raise ValueError("g not positive")
        m2 = (1 / b) * (1 / c - 1 / u2) + (2 * A / b) * math.log(u2 / c) + A * A * T
        self.m2 = m2
        h = P / N; self.h = h; self.P = P; self.N = N
        # frequencies: 2 pi j / P;  find cut-off where |phi|^(k-1) < cut
        J = 64
        while True:
            xi = 2 * np.pi * np.arange(1, J + 1) / P
            Fm2, Fm1, F0 = F_parts(xi, c, b, T)
            ph = (Fm2 + 2 * A * Fm1 + A * A * F0) / m2
            if (k - 1) * math.log(abs(ph[-1])) < math.log(cut) or J > N // 4:
                break
            J *= 2
        self.phi = np.concatenate([[1.0 + 0j], ph])
        self.J = J
    def density(self, n):
        spec = np.zeros(self.N, dtype=complex)
        pn = self.phi ** n
        spec[:len(pn)] = pn
        spec[-(len(pn) - 1):] = np.conj(pn[1:][::-1])
        f = np.fft.fft(spec).real / self.P           # sum_xi phi^n e^{-i xi s} / P
        return np.maximum(f, 0.0)
    def Gt(self, y):
        yy = np.clip(y, 0.0, self.T)
        return np.log1p(self.b * yy / self.c) / self.b + self.A * yy

def solve(model, eps=0.0, ncell=120, ptail=None, return_h=False):
    k = model.k
    fk = model.density(k); fk1 = model.density(k - 1)
    hs = model.h; s = np.arange(model.N) * hs
    cdf = np.cumsum(fk) * hs
    if ptail is None:
        ptail = 1e-9 if k < 1e7 else 1e-6
    s0 = float(s[np.searchsorted(cdf, ptail)])
    top = 1.0 + eps
    edges = np.linspace(s0, top, ncell + 1)          # h = 0 on [0, s0]
    nc = len(edges) - 1
    m1 = (s <= 1.0 - eps) & (s >= s0 - model.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 = model.Gt(edges[:, None] - ss[None, :])
        Mx = Gv[1:] - Gv[:-1]
        A += (Mx * ww) @ Mx.T
    m0 = (s <= top) & (s >= s0)
    cell = np.clip(np.searchsorted(edges, s[m0], side="right") - 1, 0, nc - 1)
    Dg = np.maximum(np.bincount(cell, weights=fk[m0] * hs, minlength=nc), 1e-300)
    lam, V = eigh(A, np.diag(Dg))
    Q = k / model.m2 * lam[-1]
    one = np.ones(nc)
    Q1 = None   # (indicator trial now needs the cell [0, s0]; see gexp_ind)
    out = dict(Q=Q, Q_ind=Q1, s0=s0, J=model.J)
    if return_h:
        hv = V[:, -1]; out["h"] = hv / hv[np.argmax(np.abs(hv))]; out["edges"] = edges
    return out

if __name__ == "__main__":
    k, th, be = int(sys.argv[1]), float(sys.argv[2]), float(sys.argv[3])
    for Aval in [float(x) for x in sys.argv[4].split(",")]:
        t0 = time.time()
        M = Model(k, th, be, A=Aval)
        r = solve(M)
        u2 = M.c + M.b * M.T
        print(json.dumps(dict(k=k, theta=th, beta=be, A=Aval, A_rel=Aval * u2, Q=r["Q"], Q_ind=r["Q_ind"], J=r["J"], secs=round(time.time() - t0, 1))), flush=True)