"""Replay attributed WASP-43b MIRI products and conditional phase-curve inference."""
from pathlib import Path
import hashlib
import json
import h5py
import numpy as np
from phase_statistics import conditional_posterior, phase_summary, interval, harmonic_design, wild_block_sensitivity
from mapping_operator import observed_window_rank

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


def fit_summary(posterior, errors):
    selected = posterior['selected']
    return dict(rows=int(selected.sum()), mean=posterior['mean'].tolist(), covariance=posterior['covariance'].tolist(),
                phase_peak_before_geometric_secondary_degrees=phase_summary(posterior['mean'])[0],
                reduced_chi_squared=float(np.sum((posterior['residual']/errors[selected])**2)/(selected.sum()-5)))


def main():
    manifest = json.loads((ROOT/'miri-provenance.json').read_text())
    for item in manifest['files']:
        payload = (ROOT/item['name']).read_bytes()
        assert len(payload) == item['bytes'] and hashlib.sha256(payload).hexdigest() == item['sha256'], item['name']
    with h5py.File(ROOT/'extracted-eureka-v1.h5') as source, h5py.File(ROOT/'figure-supp1.h5') as processed:
        times = processed['time'][:]
        extracted_times = source['time'][:]
        assert processed['time'].attrs['time_units'] == 'BMJD_TDB'
        assert np.max(np.abs(times-extracted_times[779:]-2400000.5)) < 1e-8
        values = processed['eureka_v1_flux'][:]
        errors = processed['eureka_v1_err'][:]
        source_quality_bad = source['mask_white'][779:] != 0
        assert source_quality_bad.shape == values.shape
        period, epoch_bmjd = .8134740621723353, 55934.29223313741
        epoch = epoch_bmjd+2400000.5
        posterior = conditional_posterior(times, values, errors, period, epoch, quality_mask=source_quality_bad)
        rng = np.random.default_rng(210101)
        normals = rng.normal(size=(4096, 5))
        draws = posterior['mean'] + np.einsum('ij,nj->ni', np.linalg.cholesky(posterior['covariance']), normals)
        summaries = np.array([phase_summary(draw) for draw in draws])
        curve_phase = np.linspace(0, 1, 241)
        design = harmonic_design(curve_phase)
        mean_curve = np.einsum('ni,i->n', design, posterior['mean'])
        curve_sd = np.sqrt(np.einsum('ni,ij,nj->n', design, posterior['covariance'], design))
        selected = posterior['selected']
        source_indices = np.flatnonzero(selected)
        adjacent = np.diff(source_indices) == 1
        residual = posterior['residual']
        adjacent_correlation = np.corrcoef(residual[:-1][adjacent], residual[1:][adjacent])[0, 1]
        bins = []
        for left in np.linspace(0, 1, 61)[:-1]:
            chosen = selected & (posterior['phase'] >= left) & (posterior['phase'] < left+1/60)
            if chosen.sum():
                weight = 1/errors[chosen]**2
                bins.append(dict(phase=float(np.sum(posterior['phase'][chosen]*weight)/weight.sum()),
                                 flux_ppm=float(np.sum(values[chosen]*weight)/weight.sum()*1e6),
                                 conditional_error_ppm=float(1/np.sqrt(weight.sum())*1e6), rows=int(chosen.sum())))
        reductions = {}
        for reduction in ['eureka_v1', 'eureka_v2', 'sparta', 'teatro']:
            alternative_errors = processed[reduction+'_err'][:]
            result = conditional_posterior(times, processed[reduction+'_flux'][:], alternative_errors, period, epoch,
                                           quality_mask=source_quality_bad if reduction == 'eureka_v1' else None)
            reductions[reduction] = fit_summary(result, alternative_errors)
        prior_centres = conditional_posterior(times, values, errors, .81347437, 55934.292239+2400000.5, quality_mask=source_quality_bad)
        coefficients = posterior['mean']
        longitude = np.linspace(-np.pi, np.pi, 721)
        brightness = (coefficients[0]/2 + coefficients[1]/(np.pi/2)*np.cos(longitude)
                      - coefficients[2]/(np.pi/2)*np.sin(longitude) + coefficients[3]/(2/3)*np.cos(2*longitude)
                      - coefficients[4]/(2/3)*np.sin(2*longitude))
        map_design = np.array([np.ones(len(longitude))/2, np.cos(longitude)/(np.pi/2), -np.sin(longitude)/(np.pi/2),
                               np.cos(2*longitude)/(2/3), -np.sin(2*longitude)/(2/3)]).T
        map_sd = np.sqrt(np.einsum('ni,ij,nj->n', map_design, posterior['covariance'], map_design))
        null_amplitude = float(brightness.min()*.4)
        assert null_amplitude > 0
        output = dict(kind='Conditional posterior for the provided detrended dimensionless MIRI ratios and diagonal errors',
                      object='WASP-43b', source='Zenodo 10525170 v1, CC BY 4.0; Bell and 17 other named creators',
                      reduction_versions=manifest['reduction'],
                      source_rows=9216, trimmed_rows=779, retained_rows=len(times), finite_flux_rows=int(np.isfinite(values).sum()),
                      extracted_masked_rows=int(np.sum(source['mask_white'][:]!=0)),
                      retained_source_masked_rows=int(source_quality_bad.sum()),
                      finite_flux_but_source_masked_rows=int(np.sum(np.isfinite(values)&source_quality_bad)),
                      quality_mask_convention='True/nonzero is rejected. Eureka v1 uses the matching extracted mask plus finite positive-error selection; alternate reductions use their own provided finite processed values.',
                      inherited_processed_time_attribute=str(processed['time'].attrs['time_units']),
                      checked_numerical_time_system='BJD_TDB = extracted BMJD_TDB + 2400000.5, applied exactly once',
                      nominal_exposure_seconds=10.34, white_band_micrometres=[5, 10.5],
                      ephemeris=dict(reference='Supplied EPF Value column, not newly inferred orbital posterior', period_days=period, epoch_bmjd_tdb=epoch_bmjd,
                                     geometric_conjunction=True, original_source_includes_light_travel=True,
                                     nominal_light_travel_seconds=15.113995726072995),
                      proper_prior=dict(distribution='Independent zero-mean Gaussian on all five dimensionless flux coefficients', standard_deviation=.02),
                      posterior={**fit_summary(posterior, errors), 'draw_seed':210101, 'joint_draws':draws.tolist(),
                                 'phase_peak_before_secondary_degrees_16_50_84':interval(summaries[:, 0]),
                                 'maximum_ppm_16_50_84':interval(summaries[:, 1]), 'minimum_ppm_16_50_84':interval(summaries[:, 2]),
                                 'night_at_primary_ppm_16_50_84':interval((draws[:,0]-draws[:,1]+draws[:,3])*1e6),
                                 'day_at_secondary_ppm_16_50_84':interval((draws[:,0]+draws[:,1]+draws[:,3])*1e6)},
                      curve=dict(phase=curve_phase.tolist(), mean_ppm=(mean_curve*1e6).tolist(), conditional_sd_ppm=(curve_sd*1e6).tolist()),
                      observed_bins=bins, adjacent_residual_correlation=float(adjacent_correlation),
                      residual_sensitivity=wild_block_sensitivity(posterior, times, errors),
                      ephemeris_sensitivity=dict(supplied_value_peak=phase_summary(posterior['mean'])[0], prior_centre_peak=phase_summary(prior_centres['mean'])[0],
                                                statement='Two supplied references, not two independent measurements; missing epoch-period covariance remains missing.'),
                      dependent_reductions=reductions,
                      edge_on_map_family=dict(longitude_radians=longitude.tolist(), base=brightness.tolist(), alternative=(brightness+null_amplitude*np.cos(3*longitude)).tolist(),
                                              third_mode_amplitude=null_amplitude, exact_phase_difference=0, conditional_sd=map_sd.tolist(),
                                              scope='Positive longitude-brightness maps for the edge-on, isotropic, phase-only convolution. Not gas temperatures or two thermal-PDE solutions.'),
                      failures=['Processed negative primary transit is system-excess flux; primary and eclipse intervals are excluded here.',
                                'Four reductions and derived spectra share the same photons.',
                                'No detrending covariance is supplied; conditional intervals omit it.',
                                'Phase peak is not directly the longitude of the hottest parcel.',
                                'Finite-band brightness does not measure Bond albedo or independent capacity and wind.'])
        (ROOT/'miri-posterior.json').write_text(json.dumps(output, indent=2)+'\n')
        valid_mapping = np.isfinite(values) & ~source_quality_bad
        finite_times = times[valid_mapping]
        rank = observed_window_rank(finite_times, epoch, period, np.flatnonzero(valid_mapping)+779)
        (ROOT/'mapping-rank.json').write_text(json.dumps(rank, indent=2)+'\n')
        print('PASS:', int(selected.sum()), 'actual selected points, 4096 joint proper-prior draws, residual blocks and two-resolution map operators')


if __name__ == '__main__':
    main()
