"""Replay independent limits and a finite-prior synthetic bolometric inference."""
from pathlib import Path
import json
import numpy as np
from thermal_reference import reference_field, properties, surface_curve, linear_map_eclipse, ORBIT_DAYS, resolution_check, exposure_references

ROOT = Path(__file__).resolve().parent


def main():
    cases = []
    for name, capacity, omega, albedo, longitudes in [
            ('default128', 1e7, 2000/7e7, .2, 128),
            ('default256', 1e7, 2000/7e7, .2, 256),
            ('west', 1e7, -2000/7e7, .2, 128),
            ('zero', 1e7, 0., .2, 128),
            ('near_zero_east', 1e7, 1e-12, .2, 128),
            ('near_zero_west', 1e7, -1e-12, .2, 128),
            ('inertia', 4e7, 2000/7e7, .2, 128)]:
        temperature = reference_field(longitudes, 24, capacity, omega, albedo)
        cases.append(dict(name=name, capacity=capacity, omega=omega, albedo=albedo,
                          longitudes=longitudes, latitudes=24, **properties(temperature, albedo)))
    benchmark = dict(kind='Independent SciPy sparse steady reference; no Engine imports',
                     cases=cases, nonuniform_eclipse=[linear_map_eclipse(phase) for phase in [.465, .47, .475, .48]],
                     exposure=dict(seconds=10.34, period_seconds=ORBIT_DAYS*86400,
                                   first_harmonic_attenuation=float(np.sinc(10.34/(ORBIT_DAYS*86400)))),
                     nonuniform_exposure=exposure_references())
    (ROOT / 'thermal-benchmarks.json').write_text(json.dumps(benchmark, indent=2) + '\n')
    # A deliberately finite discrete prior: the posterior below describes only these nodes.
    phases = np.linspace(.04, .96, 24)
    true_temperature = reference_field(256, 32)
    truth = surface_curve(true_temperature, phases)
    noise_ppm = 12.
    observed = truth + np.random.default_rng(210104).normal(0, noise_ppm, len(phases))
    solved = {}
    comparisons = []
    for subdivision in [1, 4, 8]:
        nodes = []
        albedo_edges = np.linspace(.05, .45, 17)
        transport_edges = np.linspace(40, 800, 39)
        for albedo_low, albedo_high in zip(albedo_edges[:-1], albedo_edges[1:]):
            for transport_low, transport_high in zip(transport_edges[:-1], transport_edges[1:]):
                refine = (albedo_low >= .15-1e-12 and albedo_high <= .25+1e-12
                          and transport_low >= 180-1e-12 and transport_high <= 380+1e-12)
                divisions = subdivision if refine else 1
                for ai in range(divisions):
                    albedo = albedo_low + (ai+.5)*(albedo_high-albedo_low)/divisions
                    for pi in range(divisions):
                        capacity_omega = transport_low + (pi+.5)*(transport_high-transport_low)/divisions
                        key = (float(albedo), float(capacity_omega))
                        if key not in solved:
                            temperature = reference_field(128, 24, 1e7, capacity_omega/1e7, float(albedo))
                            prediction = surface_curve(temperature, phases)
                            solved[key] = float(np.sum(((observed-prediction)/noise_ppm)**2))
                        prior_weight = ((albedo_high-albedo_low)*(transport_high-transport_low)
                                        /divisions**2/(.4*760))
                        nodes.append(dict(albedo=float(albedo), capacity_omega_w_m2_k=float(capacity_omega),
                                          chi_squared=solved[key], prior_weight=float(prior_weight)))
        likelihood = np.array([node['chi_squared'] for node in nodes])
        prior_weights = np.array([node['prior_weight'] for node in nodes])
        assert abs(prior_weights.sum()-1) < 1e-12
        weights = prior_weights*np.exp(-.5*(likelihood-likelihood.min()))
        weights /= weights.sum()
        for node, weight in zip(nodes, weights):
            node['posterior_weight'] = float(weight)
        def interval(key):
            values = np.array([node[key] for node in nodes])
            order = np.argsort(values)
            cumulative = np.cumsum(weights[order])
            return [float(values[order[min(np.searchsorted(cumulative, probability), len(order)-1)]]) for probability in [.16, .5, .84]]
        comparisons.append(dict(subdivision=subdivision, nodes=len(nodes), prior_sum=float(prior_weights.sum()),
                                albedo=interval('albedo'), capacity_omega=interval('capacity_omega_w_m2_k'),
                                effective_nodes=float(1/np.sum(weights**2)),
                                posterior_outside_refined_rectangle=float(sum(weight for node, weight in zip(nodes, weights)
                                    if not (.15 <= node['albedo'] <= .25 and 180 <= node['capacity_omega_w_m2_k'] <= 380)))))
    best = int(np.argmax(weights))
    output = dict(kind='Synthetic bolometric observations and posterior on an explicitly finite area-weighted discrete prior',
                  source='Independent numerical layer model, not observed WASP-43b flux',
                  seed=210104, phases=phases.tolist(), observed_ppm=observed.tolist(), truth_ppm=truth.tolist(), sigma_ppm=noise_ppm,
                  truth=dict(albedo=.2, capacity=1e7, omega=2000/7e7),
                  prior=dict(nodes=len(nodes), albedo_bounds=[.05, .45], capacity_omega_bounds=[40, 800],
                             weights='Cell area in the parameter rectangle; refined cells retain their total prior mass. Midpoint finite prior; not an exact continuous posterior.',
                             base_cells=[16, 38], refined_rectangle=[.15, .25, 180, 380], final_subdivision=8),
                  refinement=comparisons,
                  posterior_intervals_16_50_84=dict(albedo=interval('albedo'), capacity_omega=interval('capacity_omega_w_m2_k')),
                  effective_nodes=float(1/np.sum(weights**2)), best_node=nodes[best], nodes=nodes,
                  numerical_scope='Data use 256 longitude by 32 latitude surface-cell quadrature; inference uses 128 by 24. These grids have finite diffusion and projection error. No continuous posterior or independent C/wind measurement is asserted.',
                  forward_resolution_check=resolution_check(phases, nodes[best]),
                  exact_degenerate_examples=[dict(capacity=1e7, omega=2000/7e7), dict(capacity=2e7, omega=1000/7e7)],
                  failures=['One infrared band is not bolometric absorbed power.', 'Steady circular data cannot separate capacity from relative angular flow.', 'Even the refined area-weighted midpoint grid is a finite prior; quantile changes under subdivision are reported.'])
    (ROOT / 'synthetic-inference.json').write_text(json.dumps(output, indent=2)+'\n')
    print('PASS: independent steady/near-zero/eclipse references and', len(nodes), 'synthetic finite-prior nodes')


if __name__ == '__main__':
    main()
