#!/usr/bin/env python3
"""Full finite injections, reference-basis changes and deliberately wrong PSFs."""
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,
    resample,
)
from synthetic_optics import source_templates


def fixed_klip_delta(data, injection, angles, rank=1):
    images = []
    for channel in range(data.shape[0]):
        centered = data[channel] - data[channel].mean(axis=0)
        _, _, vectors = np.linalg.svd(
            centered.reshape(data.shape[1], -1), full_matrices=False
        )
        basis = vectors[:rank]
        signal = injection[channel] - injection[channel].mean(axis=0)
        matrix = signal.reshape(data.shape[1], -1)
        signal = (matrix - matrix @ basis.T @ basis).reshape(signal.shape)
        images.append(
            np.mean(
                [resample(frame, angle=angle) for frame, angle in zip(signal, angles)],
                axis=0,
            )
        )
    return np.asarray(images)


def validate(materials, output):
    metadata = json.loads((materials / "synthetic-input.json").read_text())
    arrays = np.load(materials / "synthetic-observation.npz")
    data, templates = arrays["data"], arrays["source_templates"]
    settings, angles = metadata["settings"], metadata["angles"]
    wavelengths = metadata["camera"]["wavelengthsMicrons"]
    angle = np.deg2rad(settings["positionAngleDegrees"])
    x, y = settings["separation"] * np.cos(angle), settings["separation"] * np.sin(
        angle
    )
    weights = background_subtracted_weights(80, settings, True)
    rows = []
    spectral_shape = np.asarray([1, 0.4, 1.2, 0.3, 0.8])
    for algorithm in ("adi", "asdi", "klip"):
        baseline = measurements(
            reduce_cube(data, angles, wavelengths, algorithm), weights
        )
        jacobian = finite_response(
            data, templates, angles, wavelengths, weights, algorithm
        )
        half = finite_response(
            data, templates, angles, wavelengths, weights, algorithm, step=1e-5
        )
        for amplitude in (1e-5, 1e-4, 1e-3):
            contrasts = amplitude * spectral_shape
            injection = templates * contrasts[:, None, None, None] * 3e7
            actual = (
                measurements(
                    reduce_cube(data + injection, angles, wavelengths, algorithm),
                    weights,
                )
                - baseline
            )
            linear = jacobian @ contrasts
            row = {
                "algorithm": algorithm,
                "injected_contrasts": contrasts.tolist(),
                "relative_full_vs_derivative_error": float(
                    np.linalg.norm(actual - linear) / np.linalg.norm(actual)
                ),
                "jacobian_half_step_relative_difference": float(
                    np.linalg.norm(jacobian - half) / np.linalg.norm(half)
                ),
            }
            if algorithm == "klip":
                fixed = measurements(fixed_klip_delta(data, injection, angles), weights)
                row["fixed_vs_recomputed_basis_relative_difference"] = float(
                    np.linalg.norm(fixed - actual) / np.linalg.norm(actual)
                )
            rows.append(row)
    # Same REALIZED synthetic acquisition, with an incorrect calibration OPD.
    # The generating source/star optics stay fixed; only the fit template is wrong.
    wrong_settings = {**settings, "staticOpdNm": 0, "evolvingOpdNm": 0}
    wrong_templates = source_templates(wrong_settings, angles, wavelengths, x, y)
    images = reduce_cube(data, angles, wavelengths)
    simple_weights = spatial_weights(80, x, y, False)
    simple_calibration_weights = background_subtracted_weights(80, settings, False)
    mean, covariance, count = comparison_covariance(images, settings)
    wrong_response = finite_response(
        data, wrong_templates, angles, wavelengths, simple_calibration_weights
    )
    fit = fit_spectrum(
        measurements(images, simple_weights) - mean, wrong_response, covariance
    )
    result = {
        "finite_injections": rows,
        "psf_mismatch": {
            "generating_optics": settings,
            "adopted_wrong_template_optics": wrong_settings,
            "fitted_contrasts": fit["contrast"].tolist(),
            "working_standard_errors": np.sqrt(np.diag(fit["covariance"])).tolist(),
            "comparison_count": count,
        },
        "interpretation": "Injections are deterministic response experiments on a frozen noisy background. They do not add independent source photon noise. Fixed KL references omit basis changes. PSF mismatch changes the calibration model, not the observed pixels.",
    }
    output.write_text(json.dumps(result, indent=2) + "\n")
    print("COMPLETE finite injection and PSF mismatch tests", rows, flush=True)


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