"""Compare finite-grid optical response bias with independently computed extraction noise.

The same frozen synthetic acquisition supplies the noise covariance. Reference
planet templates use independent NumPy transforms at successively finer pupil
and focal grids. This tests optical calibration bias, not unknown real PSFs.
"""
import argparse
import json
from pathlib import Path
import numpy as np
import synthetic_optics
from imaging_reduction import reduce_cube, background_subtracted_weights, finite_response


def reference_templates(settings, angles, wavelengths, size, diameter):
    original = synthetic_optics.pupil_fields
    def finer(*args, **kwargs):
        return original(*args, **{**kwargs, 'size': size, 'diameter': diameter})
    synthetic_optics.pupil_fields = finer
    angle = np.deg2rad(settings['positionAngleDegrees'])
    try:
        return synthetic_optics.source_templates(settings, angles, wavelengths,
            settings['separation'] * np.cos(angle), settings['separation'] * np.sin(angle))
    finally:
        synthetic_optics.pupil_fields = original


def main(materials, output):
    metadata = json.loads((materials / 'synthetic-input.json').read_text())
    arrays = np.load(materials / 'synthetic-observation.npz', allow_pickle=False)
    settings, angles = metadata['settings'], np.asarray(metadata['angles'])
    wavelengths = metadata['camera']['wavelengthsMicrons']
    weights = background_subtracted_weights(80, settings, False)
    baseline = finite_response(arrays['data'], arrays['source_templates'], angles, wavelengths, weights)
    summary = json.loads((materials / 'turns-summary.json').read_text())['spectrum']
    truth = settings['contrast'] * np.array([1, .9, .62, .48, .76])
    rows=[]
    for size, diameter in [(256,64),(512,64),(1024,128)]:
        templates=reference_templates(settings, angles, wavelengths, size, diameter)
        response=finite_response(arrays['data'], templates, angles, wavelengths, weights)
        fitted=np.linalg.solve(baseline, response @ truth)
        row={'transform_size':size,'pupil_diameter':diameter,
             'reference_response_diagonal':np.diag(response).tolist(),
             'coarse_calibration_contrast_bias':(fitted-truth).tolist(),
             'bias_in_saved_working_sigma':((fitted-truth)/summary['errors']).tolist(),
             'template_relative_l1':float(np.sum(np.abs(templates-arrays['source_templates']))/np.sum(templates))}
        rows.append(row);print(json.dumps(row),flush=True)
    result={'scope':'Planet optical-calibration discrepancy on the frozen camera/background; source optics only. This is not a full continuum limit or covariance calibration.','settings':settings,'reference_rows':rows}
    output.write_text(json.dumps(result,indent=2)+'\n')

if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--materials',type=Path,default=Path('.'))
    parser.add_argument('--output',type=Path,default=Path('extraction-grid.json'))
    args=parser.parse_args();main(args.materials,args.output)