"""Independent checks of the distortion shapes used by the Lab.

The closed forms G, Y and M in ac75_model.py are compared with quantities that
do not use those formulas:
  1. number and energy moments, against exact Planck integrals;
  2. Y against a finite-difference Kompaneets diffusion operator acting on
     the Planck occupation number;
  3. M against a numerically perturbed Bose-Einstein spectrum whose
     temperature is re-fitted so the photon number stays fixed.

Usage:  python3 -I check_distortion_kernels.py my-kernel-checks.json
"""
import json
import math
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ac75_model as model  # noqa: E402

ZETA3 = 1.2020569031595942


def simpson(function, low, high, intervals=40000):
    width = (high - low) / intervals
    total = function(low) + function(high)
    for index in range(1, intervals):
        total += (4.0 if index % 2 else 2.0) * function(low + index * width)
    return total * width / 3.0


def occupation(x, mu=0.0):
    return 1.0 / math.expm1(x + mu)


def kompaneets_diffusion(x, h=1.0e-3):
    """x^-2 d/dx [x^4 dn/dx] by nested central differences."""
    def flux(point):
        derivative = (occupation(point + h) - occupation(point - h)) / (2.0 * h)
        return point ** 4 * derivative
    return (flux(x + h) - flux(x - h)) / (2.0 * h) / x ** 2


def number_conserving_bose_einstein(mu=1.0e-6):
    """Delta n / mu for n(x/(1+d); mu) - n(x; 0), d chosen to keep the number.
    Number density scales as T^3, so (1+d)^3 N(mu) = N(0). N(mu) is integrated
    numerically rather than taken from a formula."""
    number = simpson(lambda u: u * u * occupation(u, mu), 1.0e-5, 60.0, 20000)
    reference = simpson(lambda u: u * u * occupation(u), 1.0e-5, 60.0, 20000)
    shift = (reference / number) ** (1.0 / 3.0) - 1.0
    def shape(point):
        return point ** 3 * (occupation(point / (1.0 + shift), mu) - occupation(point)) / mu
    return shape, shift


def main():
    output_path = sys.argv[1]
    low, high = 1.0e-4, 60.0
    planck_energy = math.pi ** 4 / 15.0
    planck_number = 2.0 * ZETA3
    moments = {}
    for name, shape in (('G', model.temperature_shape), ('Y', model.compton_y_shape),
                        ('M', model.chemical_potential_shape)):
        moments[name] = {'energy_over_planck_energy': simpson(shape, low, high) / planck_energy,
                         'number_over_planck_number': simpson(lambda x: shape(x) / x, low, high) / planck_number}
    expected = {'G': (4.0, 3.0), 'Y': (4.0, 0.0), 'M': (1.0 / 1.401, 0.0)}
    for name, (energy, number) in expected.items():
        moments[name]['expected_energy'] = energy
        moments[name]['expected_number'] = number
    grid = [0.2, 0.5, 1.0, 2.0, 3.83, 5.0, 8.0, 12.0]
    kompaneets = [{'x': x, 'finite_difference': x ** 3 * kompaneets_diffusion(x), 'closed_form': model.compton_y_shape(x)}
                  for x in grid]
    worst_y = max(abs(row['finite_difference'] - row['closed_form']) for row in kompaneets) / max(
        abs(model.compton_y_shape(x)) for x in grid)
    bose, shift = number_conserving_bose_einstein()
    chemical = [{'x': x, 'perturbed_spectrum': bose(x), 'closed_form': model.chemical_potential_shape(x)} for x in grid]
    worst_m = max(abs(row['perturbed_spectrum'] - row['closed_form']) for row in chemical) / max(
        abs(model.chemical_potential_shape(x)) for x in grid)
    zero_crossing = {}
    for name, shape, guess in (('Y', model.compton_y_shape, (2.0, 6.0)), ('M', model.chemical_potential_shape, (1.0, 4.0))):
        a, b = guess
        for _ in range(80):
            middle = 0.5 * (a + b)
            if shape(a) * shape(middle) <= 0.0:
                b = middle
            else:
                a = middle
        zero_crossing[name] = {'x': middle, 'GHz_at_T0': middle * model.K_BOLTZMANN * 2.7255 / model.H_PLANCK / 1e9}
    result = {
        'format': 'sandboxphysics-ac75-kernel-checks-v1',
        'python': sys.version.split()[0],
        'beta_mu': model.BETA_MU,
        'implied_energy_factor_1_over_M_energy': 1.0 / moments['M']['energy_over_planck_energy'],
        'moments': moments,
        'kompaneets_check': {'rows': kompaneets, 'largest_error_over_peak': worst_y},
        'bose_einstein_check': {'mu': 1.0e-6, 'temperature_shift_per_mu': shift / 1.0e-6, 'rows': chemical,
                                'largest_error_over_peak': worst_m},
        'zero_crossings': zero_crossing,
    }
    with open(output_path, 'w', encoding='utf-8') as handle:
        json.dump(result, handle, indent=1)
        handle.write('\n')
    print('wrote', output_path)
    for name, values in moments.items():
        print('%s energy %.5f (expected %.5f), number %.2e (expected %.1f)' % (
            name, values['energy_over_planck_energy'], values['expected_energy'],
            values['number_over_planck_number'], values['expected_number']))
    print('Kompaneets Y largest error / peak: %.1e' % worst_y)
    print('Bose-Einstein M largest error / peak: %.1e' % worst_m)
    print('zero crossings:', {k: round(v['GHz_at_T0'], 2) for k, v in zero_crossing.items()})


if __name__ == '__main__':
    main()