#!/usr/bin/env python3
"""Recompute the direct two-saddle major-arc validation in Table 5."""
from __future__ import annotations

import csv
import math
import numpy as np
from numpy.polynomial.legendre import leggauss
from common import (
    DATA, load_compact_counts, partition_numbers, saddle_parameter,
    saddle_variance, fwht_inplace,
)

MILESTONES = [10000, 12000, 19584, 20000]


def complex_boltzmann_xor(z: complex) -> complex:
    t = z.real
    ell = math.log(1.0 / t)
    cutoff = math.ceil((2.0 * ell + math.log(ell) + 32.0) / t)
    size = 1
    while size <= cutoff:
        size *= 2
    weights = np.zeros(size, dtype=np.complex128)
    j = np.arange(1, cutoff + 1, dtype=np.float64)
    weights[1:cutoff + 1] = np.log(np.tanh(z * j / 2.0))
    total = np.sum(weights)
    fwht_inplace(weights)
    return complex(np.mean(np.exp(0.5 * (total - weights))))


def log_partition_product(z: complex) -> complex:
    cutoff = math.ceil(60.0 / z.real)
    j = np.arange(1, cutoff + 1, dtype=np.float64)
    return complex(-np.sum(np.log1p(-np.exp(-z * j))))


def major_arc_ratio(n: int, width_in_sigma: float, nodes: int) -> tuple[complex, float, float]:
    t = saddle_parameter(n)
    sigma = math.sqrt(saddle_variance(t))
    half_width = width_in_sigma / sigma
    x, weights = leggauss(nodes)
    theta = half_width * x
    weights = half_width * weights
    log_p_real = log_partition_product(complex(t, 0.0))

    denominator = 0.0j
    numerator = 0.0j
    for th, weight in zip(theta, weights):
        z = complex(t, -float(th))
        kernel = np.exp(log_partition_product(z) - log_p_real - 1j * n * th)
        denominator += weight * kernel
        numerator += weight * complex_boltzmann_xor(z) * kernel
    return 2.0 * numerator / denominator, t, sigma


def main() -> None:
    counts = load_compact_counts()
    partitions = partition_numbers(max(MILESTONES))
    (DATA / 'tables').mkdir(parents=True, exist_ok=True)
    out = DATA / 'tables' / 'table_05_major_arc_validation.tsv'
    with out.open('w', newline='') as fh:
        w = csv.writer(fh, delimiter='\t', lineterminator='\n')
        w.writerow([
            'total_size_N', 't_N', 'sigma_N', 'exact_density',
            'major_arc_value_5sigma', 'signed_relative_error_5sigma',
            'major_arc_value_10sigma', 'signed_relative_error_10sigma'
        ])
        for n in MILESTONES:
            exact = counts[n] / partitions[n]
            value5, t, sigma = major_arc_ratio(n, 5.0, 256)
            value10, _, _ = major_arc_ratio(n, 10.0, 512)
            w.writerow([n, f'{t:.17g}', f'{sigma:.17g}', f'{exact:.17g}',
                        f'{value5.real:.17g}', f'{value5.real / exact - 1.0:.17g}',
                        f'{value10.real:.17g}', f'{value10.real / exact - 1.0:.17g}'])
    print(out)


if __name__ == '__main__':
    main()
