"""Shared numerical routines for the supplementary computations."""
from __future__ import annotations

import math
from pathlib import Path
import numpy as np

ROOT = Path(__file__).resolve().parents[2]
DATA = ROOT / 'data'
FIGURES = ROOT / 'figures'


def load_compact_counts() -> dict[int, int]:
    values: dict[int, int] = {}
    path = DATA / 'counts' / 'oeis_A048833_k0_k10000.txt'
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        k_s, value_s = line.split()
        values[2 * int(k_s)] = int(value_s)
    return values


def partition_numbers(nmax: int) -> list[int]:
    p = [0] * (nmax + 1)
    p[0] = 1
    for n in range(1, nmax + 1):
        total = 0
        k = 1
        while True:
            g1 = k * (3 * k - 1) // 2
            if g1 > n:
                break
            sign = 1 if k % 2 else -1
            total += sign * p[n - g1]
            g2 = k * (3 * k + 1) // 2
            if g2 <= n:
                total += sign * p[n - g2]
            k += 1
        p[n] = total
    return p


def saddle_mean(t: float) -> float:
    cutoff = max(2000, math.ceil(60.0 / t))
    j = np.arange(1, cutoff + 1, dtype=np.float64)
    return float(np.sum(j / np.expm1(t * j)))


def saddle_parameter(n: int) -> float:
    center = math.pi / math.sqrt(6.0 * n)
    lo, hi = 0.5 * center, 2.0 * center
    for _ in range(80):
        mid = 0.5 * (lo + hi)
        if saddle_mean(mid) > n:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def saddle_variance(t: float) -> float:
    cutoff = max(2000, math.ceil(60.0 / t))
    j = np.arange(1, cutoff + 1, dtype=np.float64)
    et = np.exp(t * j)
    return float(np.sum(j * j * et / (et - 1.0) ** 2))


def saddle_variance_and_k3(t: float) -> tuple[float, float]:
    cutoff = max(2000, math.ceil(50.0 / t))
    j = np.arange(1, cutoff + 1, dtype=np.float64)
    x = t * j
    em1 = np.expm1(x)
    ex = em1 + 1.0
    variance = float(np.sum(j * j * ex / (em1 * em1)))
    k3 = float(np.sum(j**3 * ex * (ex + 1.0) / (em1**3)))
    return variance, k3


def fwht_inplace(values: np.ndarray) -> None:
    h = 1
    n = values.shape[0]
    while h < n:
        blocks = values.reshape(-1, 2 * h)
        left = blocks[:, :h].copy()
        right = blocks[:, h:].copy()
        blocks[:, :h] = left + right
        blocks[:, h:] = left - right
        h *= 2


def boltzmann_derivatives(t: float) -> tuple[float, float, float]:
    ell = math.log(1.0 / t)
    cutoff = math.ceil((2.0 * ell + math.log(ell) + 30.0) / t)
    size = 1
    while size <= cutoff:
        size *= 2

    j = np.arange(1, cutoff + 1, dtype=np.float64)
    x = t * j
    w = np.zeros(size, dtype=np.float64)
    wp = np.zeros(size, dtype=np.float64)
    wpp = np.zeros(size, dtype=np.float64)
    w[1:cutoff + 1] = np.log(np.tanh(x / 2.0))
    wp[1:cutoff + 1] = j / np.sinh(x)
    wpp[1:cutoff + 1] = -(j * j) * np.cosh(x) / (np.sinh(x)**2)

    sw, swp, swpp = float(w.sum()), float(wp.sum()), float(wpp.sum())
    fwht_inplace(w)
    fwht_inplace(wp)
    fwht_inplace(wpp)
    ell_u = 0.5 * (sw - w)
    ell_u_p = 0.5 * (swp - wp)
    ell_u_pp = 0.5 * (swpp - wpp)
    exp_ell = np.exp(ell_u)
    R = float(exp_ell.mean())
    Rp = float((exp_ell * ell_u_p).mean())
    Rpp = float((exp_ell * (ell_u_p * ell_u_p + ell_u_pp)).mean())
    return R, Rp, Rpp


def boltzmann_zero_xor(t: float) -> float:
    return boltzmann_derivatives(t)[0]


def block_product(t: float, block: int, jscale: int) -> float:
    j = np.arange(block * jscale, (block + 1) * jscale, dtype=np.float64)
    return float(np.exp(np.sum(np.log(np.tanh(t * j / 2.0)))))


def one_level(t: float, jscale: int) -> float:
    p1 = block_product(t, 1, jscale)
    return (1.0 + p1) / (2.0 * jscale)


def two_level(t: float, jscale: int) -> float:
    p1 = block_product(t, 1, jscale)
    p2 = block_product(t, 2, jscale)
    p3 = block_product(t, 3, jscale)
    return (1.0 + p1 * p2 + p1 * p3 + p2 * p3) / (4.0 * jscale)


def low_block_scale(t: float, jscale: int) -> float:
    j = np.arange(jscale // 2, jscale, dtype=np.float64)
    return float(np.exp(np.sum(np.log(np.tanh(t * j / 2.0)))))
