#!/usr/bin/env python3
"""Generate the finite-tilt table and compare the Gaussian integral formulas."""
from __future__ import annotations

import csv
from pathlib import Path
import mpmath as mp

ROOT = Path(__file__).resolve().parents[1]


def close(actual, expected, tolerance) -> None:
    if abs(actual-expected) > tolerance:
        raise AssertionError(f"Numerical difference exceeds {tolerance}: {actual-expected}")


def w_moment(T, power: int):
    """Integral of [v(1-v)]^power exp[-T v(1-v)] over v in [0,1]."""
    return mp.beta(power+1, power+1)*mp.hyp1f1(power+1, power+mp.mpf('1.5'), -T/4)


def time_integrand(T):
    j = {p: w_moment(T, p) for p in range(1, 6)}
    return (3*T**4*j[5]-T**3*(72*j[4]+16*j[3])
            +T*T*(540*j[3]+160*j[2]+16*j[1])
            -T*(1440*j[2]+320*j[1])+1080*j[1])/(2*T**4)


def scaled_difference(epsilon: str, precision: int):
    with mp.workdps(precision):
        eps = mp.mpf(epsilon)
        value = eps**3*mp.quad(time_integrand, [eps, mp.sqrt(eps), 1])
        return +value


def check_kernels() -> None:
    with mp.workdps(55):
        tolerance = mp.mpf('1e-42')
        for tau, k, eta, kap in [('.7', '1.1', '.3', '1.4'),
                                 ('-.4', '.8', '.5', '.9'),
                                 ('.2', '0', '.25', '1.2')]:
            tau, k, eta, kap = map(mp.mpf, (tau, k, eta, kap))
            density = lambda u: mp.exp(-(tau-2*k*u)**2/(4*eta**2*u)-kap**2*u)/mp.sqrt(4*mp.pi*eta**2*u)
            integral = mp.quad(density, [0, mp.mpf('.1'), 1, mp.inf])
            energy = mp.sqrt(k*k+eta*eta*kap*kap)
            formula = mp.exp((k*tau-energy*abs(tau))/eta**2)/(2*energy)
            close(integral, formula, tolerance)
        width, centre = mp.mpf('.18'), mp.mpf('.137')
        for M in (4, 8, 16):
            direct = mp.fsum(mp.exp(-((mp.mpf(j)/M-centre)/width)**2)/(M*mp.sqrt(mp.pi)*width) for j in range(-100, 101))
            transformed = mp.fsum(mp.exp(-mp.pi**2*(M*width*r)**2)*mp.exp(2j*mp.pi*r*M*centre) for r in range(-100, 101))
            close(direct, transformed, tolerance)
        # Check the beta/hypergeometric reduction against real one-dimensional quadrature.
        for T in (mp.mpf('.1'), mp.mpf('.6'), mp.mpf('1')):
            for power in range(1, 6):
                integral = mp.quad(lambda v: (v*(1-v))**power*mp.exp(-T*v*(1-v)), [0, mp.mpf('.5'), 1])
                close(integral, w_moment(T, power), tolerance)


def write_table() -> None:
    epsilons = ('0.1', '0.05', '0.025', '0.0125', '0.00625')
    rows = []
    for epsilon in epsilons:
        low = scaled_difference(epsilon, 40)
        high = scaled_difference(epsilon, 60)
        with mp.workdps(60):
            close(low, high, mp.mpf('1e-34'))
            rows.append({'epsilon': epsilon, 'delta': mp.nstr(mp.mpf(epsilon)**4, 25),
                         'difference_over_epsilon': mp.nstr(high, 32)})
    directory = ROOT/'data'
    directory.mkdir(exist_ok=True)
    with (directory/'limit_values.csv').open('w', newline='', encoding='utf-8') as stream:
        writer = csv.DictWriter(stream, fieldnames=tuple(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    print("Gaussian integral comparisons passed.")
    print("Wrote data/limit_values.csv.")


if __name__ == '__main__':
    check_kernels()
    write_table()