"""Recompute a fixed first-order phase-transition benchmark: sound lifetime and today's spectrum.

Standard library only (Python 3.9 or newer). No network access.

Usage
  python3 recompute_benchmark.py inputs/benchmarks.json > my-output.json
  python3 recompute_benchmark.py inputs/benchmarks.json --compare expected.json

The sound-wave template is the double broken power law of Caprini et al.,
arXiv:2403.03723v2, section 2.2.1 (equations sw_shape, sw_amplitude, Rstar,
Om2_sound). The single-bubble efficiency kappa uses the appendix-B fits of
Espinosa, Konstandin, No and Servant, arXiv:1004.4187v1. The LISA noise uses
the test-mass and optical-metrology model of arXiv:2403.03723v2 section 3 and
appendix A, with the optical-metrology corner frequency taken as 2 mHz (see
README.md for why). The result is a parametric template evaluation, not a
solution of the Einstein or fluid equations.
"""
import json
import math
import sys

SOUND_SPEED = 1.0 / math.sqrt(3.0)
ADIABATIC_INDEX = 4.0 / 3.0
SOUND_AMPLITUDE_CONSTANT = 0.11
BULK_EFFICIENCY_FACTOR = 0.6
SLOPES = {'n1': 3.0, 'n2': 1.0, 'n3': -3.0, 'a1': 2.0, 'a2': 4.0}

# Physical constants (CODATA 2018 / Fixsen 2009) for the first-principles redshift check.
CMB_TEMPERATURE_K = 2.7255
BOLTZMANN_EV_PER_K = 8.617333262e-5
REDUCED_PLANCK_GEV_S = 6.582119569e-25
PLANCK_MASS_GEV = 1.220890e19
TODAY_ENTROPY_DOF = 3.909
SPEED_OF_LIGHT = 299792458.0
NEWTON_CONSTANT = 6.67430e-11
MEGAPARSEC_M = 3.0856775814913673e22
HUBBLE_100_PER_S = 100e3 / MEGAPARSEC_M

LISA_ARM_M = 2.5e9
LISA_TEST_MASS_ACCELERATION = 3e-15      # m s^-2 Hz^-1/2
LISA_METROLOGY_DISPLACEMENT = 15e-12     # m Hz^-1/2
LISA_METROLOGY_CORNER_HZ = 2e-3
SECONDS_PER_YEAR = 365.25 * 86400.0


class OutsideTemplateDomain(ValueError):
    pass


def jouguet_speed(alpha):
    return (math.sqrt(2.0 * alpha / 3.0 + alpha * alpha) + math.sqrt(1.0 / 3.0)) / (1.0 + alpha)


def single_bubble_efficiency(alpha, wall_speed):
    """EKNS appendix-B fits; quoted precision about 15 percent for 1e-3 < alpha < 10."""
    sound = SOUND_SPEED
    kappa_a = wall_speed ** 1.2 * 6.9 * alpha / (1.36 - 0.037 * math.sqrt(alpha) + alpha)
    kappa_b = alpha ** 0.4 / (0.017 + (0.997 + alpha) ** 0.4)
    kappa_c = math.sqrt(alpha) / (0.135 + math.sqrt(0.98 + alpha))
    kappa_d = alpha / (0.73 + 0.083 * math.sqrt(alpha) + alpha)
    critical = jouguet_speed(alpha)
    if wall_speed < sound:
        return sound ** 2.2 * kappa_a * kappa_b / ((sound ** 2.2 - wall_speed ** 2.2) * kappa_b + wall_speed * sound ** 1.2 * kappa_a)
    if wall_speed < critical:
        delta = -0.9 * math.log(math.sqrt(alpha) / (1.0 + math.sqrt(alpha)))
        return kappa_b + (wall_speed - sound) * delta + (wall_speed - sound) ** 3 / (critical - sound) ** 3 * (kappa_c - kappa_b - (critical - sound) * delta)
    return ((critical - 1.0) ** 3 * critical ** 2.5 * wall_speed ** -2.5 * kappa_c * kappa_d
            / (((critical - 1.0) ** 3 - (wall_speed - 1.0) ** 3) * critical ** 2.5 * kappa_c + (wall_speed - 1.0) ** 3 * kappa_d))


def deflagration_alpha_bound(wall_speed, bound_table):
    rows = bound_table['rows']
    if wall_speed <= rows[0]['wallSpeed']:
        return rows[0]['alphaLower']
    for left, right in zip(rows, rows[1:]):
        if wall_speed <= right['wallSpeed']:
            share = (wall_speed - left['wallSpeed']) / (right['wallSpeed'] - left['wallSpeed'])
            return math.exp((1 - share) * math.log(left['alphaLower']) + share * math.log(right['alphaLower']))
    return math.inf


def check_domain(case, bound_table):
    alpha, wall_speed, beta = case['alpha'], case['wallSpeed'], case['betaOverHubble']
    if not 0.001 < alpha <= 1.0:
        raise OutsideTemplateDomain('alpha outside 0.001 < alpha <= 1 used here for the efficiency fits')
    if not 0.0 < wall_speed < 1.0:
        raise OutsideTemplateDomain('wall speed must lie between 0 and the speed of light')
    if SOUND_SPEED - 0.05 < wall_speed <= jouguet_speed(alpha):
        raise OutsideTemplateDomain('hybrid or near-sonic wall: the template shell rule |xi_w - c_s| is not stated there')
    if wall_speed < SOUND_SPEED and alpha > deflagration_alpha_bound(wall_speed, bound_table):
        raise OutsideTemplateDomain('no subsonic deflagration profile exists for this alpha and wall speed')
    bubble_size = (8.0 * math.pi) ** (1.0 / 3.0) * max(wall_speed, SOUND_SPEED) / beta
    if bubble_size >= 1.0:
        raise OutsideTemplateDomain('mean bubble separation reaches the Hubble radius; flat-space template not applicable')


def first_principles_redshift(temperature_gev, dof):
    """H_{*,0} in Hz and h^2 F_GW,0 from entropy conservation, with g_s* = g_* assumed."""
    hubble_star_gev = math.sqrt(8.0 * math.pi ** 3 * dof / 90.0) * temperature_gev ** 2 / PLANCK_MASS_GEV
    hubble_star_hz = hubble_star_gev / REDUCED_PLANCK_GEV_S
    today_temperature_gev = CMB_TEMPERATURE_K * BOLTZMANN_EV_PER_K * 1e-9
    scale_ratio = (TODAY_ENTROPY_DOF / dof) ** (1.0 / 3.0) * today_temperature_gev / temperature_gev
    joule_per_ev = 1.602176634e-19
    photon_energy_density = (math.pi ** 2 / 15.0) * (CMB_TEMPERATURE_K * BOLTZMANN_EV_PER_K * joule_per_ev) ** 4 / (
        (REDUCED_PLANCK_GEV_S * 1e9 * joule_per_ev) ** 3 * SPEED_OF_LIGHT ** 3)
    critical_energy_density_h2 = 3.0 * HUBBLE_100_PER_S ** 2 * SPEED_OF_LIGHT ** 2 / (8.0 * math.pi * NEWTON_CONSTANT)
    photon_density_h2 = photon_energy_density / critical_energy_density_h2
    dilution_h2 = photon_density_h2 * (dof / 2.0) * (TODAY_ENTROPY_DOF / dof) ** (4.0 / 3.0)
    return {'hubbleRateTodayHz': scale_ratio * hubble_star_hz, 'dilutionH2': dilution_h2, 'photonDensityH2': photon_density_h2}


def shape(frequency, first_break, second_break):
    s = SLOPES
    return ((frequency / first_break) ** s['n1']
            * (1.0 + (frequency / first_break) ** s['a1']) ** ((s['n2'] - s['n1']) / s['a1'])
            * (1.0 + (frequency / second_break) ** s['a2']) ** ((s['n3'] - s['n2']) / s['a2']))


def lisa_noise_density_h2(frequency):
    """h^2 Omega_N of one of the A/E channels (no foregrounds)."""
    x = 2.0 * math.pi * frequency * LISA_ARM_M / SPEED_OF_LIGHT
    test_mass = (LISA_TEST_MASS_ACCELERATION ** 2 * (1.0 + (0.4e-3 / frequency) ** 2) * (1.0 + (frequency / 8e-3) ** 4)
                 * (1.0 / (2.0 * math.pi * frequency * SPEED_OF_LIGHT)) ** 2)
    metrology = (LISA_METROLOGY_DISPLACEMENT ** 2 * (1.0 + (LISA_METROLOGY_CORNER_HZ / frequency) ** 4)
                 * (2.0 * math.pi * frequency / SPEED_OF_LIGHT) ** 2)
    noise_over_response = (4.0 * (1.0 + math.cos(x) + math.cos(x) ** 2) * test_mass + (2.0 + math.cos(x)) * metrology) / (2.0 * x * x)
    return 4.0 * math.pi ** 2 * frequency ** 3 * noise_over_response / (3.0 * HUBBLE_100_PER_S ** 2)


def signal_to_noise(spectrum, years):
    lower, upper, count = math.log(3e-5), math.log(0.5), 1600
    step = (upper - lower) / (count - 1)
    total = 0.0
    for index in range(count):
        frequency = math.exp(lower + index * step)
        ratio = spectrum(frequency) / lisa_noise_density_h2(frequency)
        weight = 0.5 if index in (0, count - 1) else 1.0
        total += weight * ratio * ratio * frequency * step
    return math.sqrt(2.0 * years * SECONDS_PER_YEAR * total)


def evaluate(case, bound_table):
    check_domain(case, bound_table)
    alpha, wall_speed, beta = case['alpha'], case['wallSpeed'], case['betaOverHubble']
    temperature, dof = case['temperatureGeV'], case['relativisticDof']
    kappa = single_bubble_efficiency(alpha, wall_speed)
    kinetic_fraction = BULK_EFFICIENCY_FACTOR * kappa * alpha / (1.0 + alpha)
    bubble_size = (8.0 * math.pi) ** (1.0 / 3.0) * max(wall_speed, SOUND_SPEED) / beta
    rms_fluid_speed = math.sqrt(kinetic_fraction / ADIABATIC_INDEX)
    shock_time = bubble_size / rms_fluid_speed
    source_time = min(shock_time, 1.0)
    expansion_factor = 1.0 - 1.0 / math.sqrt(1.0 + 2.0 * shock_time)
    hubble_today = 1.65e-5 * (dof / 100.0) ** (1.0 / 6.0) * (temperature / 100.0)
    dilution = 1.64e-5 * (100.0 / dof) ** (1.0 / 3.0)
    shell = abs(wall_speed - SOUND_SPEED) / max(wall_speed, SOUND_SPEED)
    first_break = 0.2 * hubble_today / bubble_size
    second_break = 0.5 * hubble_today / (shell * bubble_size)
    integrated = dilution * SOUND_AMPLITUDE_CONSTANT * kinetic_fraction ** 2 * source_time * bubble_size
    ratio = second_break / first_break
    amplitude_second_break = (math.sqrt(2.0) + 2.0 * ratio / (1.0 + ratio * ratio)) / math.pi * integrated
    normalisation = shape(second_break, first_break, second_break)

    def spectrum(frequency):
        return amplitude_second_break * shape(frequency, first_break, second_break) / normalisation

    peak_frequency, peak_value = first_break, 0.0
    for index in range(2001):
        frequency = first_break * 0.1 * (100.0 * ratio) ** (index / 2000)
        value = spectrum(frequency)
        if value > peak_value:
            peak_frequency, peak_value = frequency, value
    reference = first_principles_redshift(temperature, dof)
    return {
        'id': case['id'],
        'regime': 'deflagration' if wall_speed < SOUND_SPEED else 'detonation',
        'kappa': kappa, 'kineticFraction': kinetic_fraction,
        'hubbleBubbleSize': bubble_size, 'rmsFluidSpeed': rms_fluid_speed,
        'hubbleShockTime': shock_time, 'hubbleSourceTime': source_time,
        'lifetimeLimitedBy': 'shock formation' if shock_time < 1.0 else 'Hubble expansion',
        'expansionSuppressionUpsilon': expansion_factor,
        'shellThickness': shell,
        'firstBreakHz': first_break, 'secondBreakHz': second_break,
        'integratedAmplitudeH2': integrated, 'secondBreakAmplitudeH2': amplitude_second_break,
        'peakFrequencyHz': peak_frequency, 'peakAmplitudeH2': peak_value,
        'snrLisaFourYearsAE': signal_to_noise(spectrum, 4.0),
        'redshiftCheck': {
            'hubbleRateTodayTemplateHz': hubble_today,
            'hubbleRateTodayFirstPrinciplesHz': reference['hubbleRateTodayHz'],
            'dilutionTemplateH2': dilution,
            'dilutionFirstPrinciplesH2': reference['dilutionH2'],
        },
    }


def run(payload):
    bound_table = payload['deflagrationBound']
    results, refusals = [], []
    for case in payload['benchmarks']:
        results.append(evaluate(case, bound_table))
    for case in payload['failureCases']:
        try:
            evaluate(case, bound_table)
            refusals.append({'id': case['id'], 'refused': False, 'reason': None})
        except OutsideTemplateDomain as error:
            refusals.append({'id': case['id'], 'refused': True, 'reason': str(error)})
    return {'script': 'recompute_benchmark.py', 'python': sys.version.split()[0], 'results': results, 'failureCases': refusals}


def compare(actual, expected, tolerance=1e-9):
    problems = []

    def walk(left, right, path):
        if isinstance(right, dict):
            for key, value in right.items():
                if key == 'python':
                    continue
                walk(left.get(key) if isinstance(left, dict) else None, value, path + '.' + key)
        elif isinstance(right, list):
            for index, value in enumerate(right):
                walk(left[index] if isinstance(left, list) and index < len(left) else None, value, path + '[' + str(index) + ']')
        elif isinstance(right, float):
            if not isinstance(left, (int, float)) or abs(left - right) > tolerance * max(1.0, abs(right)) and abs(left / right - 1.0) > tolerance:
                problems.append(path + ': ' + repr(left) + ' != ' + repr(right))
        elif left != right:
            problems.append(path + ': ' + repr(left) + ' != ' + repr(right))
    walk(actual, expected, '')
    return problems


def main(arguments):
    if not arguments:
        print(__doc__)
        return 2
    with open(arguments[0], encoding='utf-8') as handle:
        payload = json.load(handle)
    output = run(payload)
    if len(arguments) == 3 and arguments[1] == '--compare':
        with open(arguments[2], encoding='utf-8') as handle:
            expected = json.load(handle)
        problems = compare(output, expected)
        print('MATCH' if not problems else 'MISMATCH\n' + '\n'.join(problems))
        return 0 if not problems else 1
    print(json.dumps(output, indent=1))
    return 0


if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))
