#!/usr/bin/env python3
"""Reprocess the same frozen synthetic camera cube with an independent NumPy implementation."""
import argparse
import json
from pathlib import Path
import numpy as np
from imaging_reduction import (
    reduce_cube,
    background_subtracted_weights,
    spatial_weights,
    measurements,
    comparison_covariance,
    finite_response,
    fit_spectrum,
)
from synthetic_optics import pupil_fields
from noise_covariance import unit_source_covariances


def main(materials, output):
    metadata = json.loads((materials / "synthetic-input.json").read_text())
    arrays = np.load(materials / "synthetic-observation.npz", allow_pickle=False)
    data, templates = arrays["data"], arrays["source_templates"]
    settings, angles = metadata["settings"], np.asarray(metadata["angles"])
    wavelengths = metadata["camera"]["wavelengthsMicrons"]
    images = reduce_cube(data, angles, wavelengths)
    angle = np.deg2rad(settings["positionAngleDegrees"])
    x, y = settings["separation"] * np.cos(angle), settings["separation"] * np.sin(
        angle
    )
    weights = spatial_weights(80, x, y, spatial=False)
    mean, covariance, count = comparison_covariance(images, settings)
    calibration_weights = background_subtracted_weights(80, settings, False)
    response = finite_response(data, templates, angles, wavelengths, calibration_weights)
    fit = fit_spectrum(measurements(images, weights) - mean, response, covariance)
    reference = json.loads((materials / "turns-summary.json").read_text())["spectrum"]
    difference = float(np.max(np.abs(fit["contrast"] - reference["contrast"])))
    if difference > 1e-10:
        raise AssertionError(
            "Independent processing differs from the frozen browser spectrum"
        )
    unit_covariance = unit_source_covariances(templates, calibration_weights, angles, wavelengths)
    source_covariance = np.einsum(
        "c,cij->ij", np.maximum(fit["contrast"], 0), unit_covariance
    )
    corrected_fit = fit_spectrum(
        measurements(images, weights) - mean, response, covariance + source_covariance
    )
    error_difference = float(
        np.max(
            np.abs(np.sqrt(np.diag(corrected_fit["covariance"])) - reference["errors"])
        )
    )
    if error_difference > 1e-10:
        raise AssertionError("Independent source-photon covariance disagrees")
    spatial = spatial_weights(80, x, y, spatial=True)
    spatial_mean, spatial_covariance, spatial_count = comparison_covariance(
        images, settings, spatial=True
    )
    spatial_calibration_weights = background_subtracted_weights(80, settings, True)
    spatial_response = finite_response(data, templates, angles, wavelengths, spatial_calibration_weights)
    spatial_observed = measurements(images, spatial) - spatial_mean
    omitted = fit_spectrum(spatial_observed, spatial_response, spatial_covariance)
    spatial_unit_covariance = unit_source_covariances(
        templates, spatial_calibration_weights, angles, wavelengths
    )
    spatial_source_covariance = np.einsum(
        "c,cij->ij", np.maximum(fit["contrast"], 0), spatial_unit_covariance
    )
    corrected_spatial = fit_spectrum(
        spatial_observed,
        spatial_response,
        spatial_covariance + spatial_source_covariance,
    )
    # Independent Fourier transform and normalization, compared to saved TS samples.
    optical = json.loads((materials / "optical-reference.json").read_text())
    errors = []
    for case in optical["cases"]:
        calculated = pupil_fields(
            settings, case["wavelength"], tuple(case["source"]), case["exposure"]
        )
        errors.append(
            max(
                abs(calculated.ravel()[pixel] - value)
                for pixel, value in case["samples"]
            )
        )
    if max(errors) > 1e-12:
        raise AssertionError("Independent NumPy FFT disagrees")
    result = {
        "comparison_patches": count,
        "independent_contrast": fit["contrast"].tolist(),
        "max_contrast_difference": difference,
        "max_independent_standard_error_difference": error_difference,
        "missing_source_noise_failure": {
            "chi_square": omitted["chi_square"],
            "residual_dimensions": 10,
        },
        "corrected_conditional_spatial_fit": {
            "chi_square": corrected_spatial["chi_square"],
            "negative2loglikelihood": corrected_spatial["negative2loglikelihood"],
            "comparison_count": spatial_count,
            "comparison_condition": float(np.linalg.cond(spatial_covariance)),
            "source_covariance": spatial_source_covariance.tolist(),
        },
        "max_independent_optical_sample_error": max(errors),
        "units": "entrance planet/star photon contrast per finite band",
        "limitations": "This exact linear five-aperture solve alone has zero residual degrees of freedom. Error coverage and the spatial profile are separate validation products. No observed GPI pixels are in this synthetic input.",
    }
    output.write_text(json.dumps(result, indent=2) + "\n")
    print("PASS independent NumPy optical propagation and same-cube extraction", result)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--materials", type=Path, default=Path("."))
    parser.add_argument(
        "--output", type=Path, default=Path("synthetic-reproduction.json")
    )
    args = parser.parse_args()
    main(args.materials, args.output)