"""Declared sparse-cube processing; not the nineteen-exposure published reduction."""

import math
import numpy as np
from image_tools import bilinear, insert_template, aperture
from imaging_reduction import resample


def finite_mean(values, axis):
    valid = np.isfinite(values)
    count = np.sum(valid, axis=axis)
    total = np.sum(np.where(valid, values, 0), axis=axis)
    return np.divide(total, count, out=np.full_like(total, np.nan), where=count > 0)


def registered_science(dataset, strict_quality=False, size=111):
    rows, columns = np.indices((size, size))
    result, centers = [], []
    for exposure, images in enumerate(dataset["science"]):
        frame, frame_centers = [], []
        for channel, image in enumerate(images):
            locations = np.asarray(
                [spot["measured"] for spot in dataset["spots"][exposure][channel]]
            )
            center = locations.mean(axis=0)
            source = image.copy()
            if strict_quality:
                source[dataset["quality"][exposure, channel] != 0] = np.nan
            frame.append(
                bilinear(
                    source,
                    center[0] + columns - size // 2,
                    center[1] + rows - size // 2,
                )
            )
            frame_centers.append(center)
        result.append(frame)
        centers.append(frame_centers)
    return np.transpose(np.asarray(result), (1, 0, 2, 3)), np.asarray(centers)


def detector_offset(header, east_mas, north_mas):
    matrix = (
        np.asarray(
            [[header["CD1_1"], header["CD1_2"]], [header["CD2_1"], header["CD2_2"]]]
        )
        * 3.6e6
    )
    return np.linalg.solve(
        matrix, np.stack(np.broadcast_arrays(east_mas, north_mas)).reshape(2, -1)
    ).reshape((2,) + np.broadcast(east_mas, north_mas).shape)


def sky_project(image, header, sample_mas=14.14):
    size = image.shape[0]
    rows, columns = np.indices(image.shape)
    east = -(columns - size // 2) * sample_mas
    north = (size // 2 - rows) * sample_mas
    offset = detector_offset(header, east, north)
    return bilinear(image, size // 2 + offset[0], size // 2 + offset[1])


def reduce_registered(data, dataset, algorithm="adi"):
    frames = data.copy()
    wavelengths = np.asarray(
        dataset["manifest"]["files"][0]["observed"]["selected_wavelengths_micron"]
    )
    if algorithm == "asdi":
        for exposure in range(data.shape[1]):
            aligned = np.asarray(
                [
                    resample(frame, scale=1.24 / wavelength)
                    for frame, wavelength in zip(frames[:, exposure], wavelengths)
                ]
            )
            aligned -= aligned.mean(axis=0)
            for channel, wavelength in enumerate(wavelengths):
                frames[channel, exposure] = resample(
                    aligned[channel], scale=wavelength / 1.24
                )
    # All three frames enter the reference, including the source itself.
    mean = finite_mean(frames, axis=1)
    result = []
    for channel in range(data.shape[0]):
        rotated = np.asarray(
            [
                sky_project(frames[channel, exposure] - mean[channel], header)
                for exposure, header in enumerate(dataset["headers"])
            ]
        )
        result.append(finite_mean(rotated, axis=0))
    return np.asarray(result)


def empirical_injection(
    dataset, shape, separation_mas, pa_degrees, template_variant=None
):
    count = len(dataset["manifest"]["channels"])
    output = np.zeros((count, 3, shape, shape))
    east, north = separation_mas * math.sin(
        math.radians(pa_degrees)
    ), separation_mas * math.cos(math.radians(pa_degrees))
    for exposure, header in enumerate(dataset["headers"]):
        offset = detector_offset(header, east, north)
        for channel in range(count):
            template = dataset["templates"][exposure, channel]
            if template_variant is not None:
                template = dataset["individual_templates"][
                    exposure, channel, template_variant
                ]
            output[channel, exposure] = (
                insert_template(
                    (shape, shape),
                    template,
                    shape // 2 + offset[0],
                    shape // 2 + offset[1],
                )
                * dataset["star_aperture_adu"][exposure, channel]
            )
    return output


def sky_measurements(images, separation_mas, pa_degrees, spatial=False):
    radius_pixels = separation_mas / 14.14
    angle = math.radians(pa_degrees)
    x, y = -radius_pixels * math.sin(angle), -radius_pixels * math.cos(angle)
    center = images.shape[-1] // 2
    values = []
    for image in images:
        central = aperture(image, center + x, center + y, 3)
        if not spatial:
            values.append(central)
        else:
            radial_x, radial_y = 2 * x / radius_pixels, 2 * y / radius_pixels
            values.extend(
                [
                    central,
                    aperture(image, center + x + radial_x, center + y + radial_y, 3)
                    - aperture(image, center + x - radial_x, center + y - radial_y, 3),
                    aperture(image, center + x - radial_y, center + y + radial_x, 3)
                    - aperture(image, center + x + radial_y, center + y - radial_x, 3),
                ]
            )
    return np.asarray(values)


def real_comparison(images, separation_mas, pa_degrees, spatial=False):
    radius_pixels = separation_mas / 14.14
    count = int(2 * math.pi * radius_pixels / 3.2)
    samples = []
    accepted_angles = []
    # Three time samples span 19.7246 degrees. Exclude the target and full ADI
    # negative-lobe arc, plus eight pixels for the aperture/background footprint.
    exclusion = 19.7245815655 + math.degrees(8 / radius_pixels)
    for index in range(count):
        angle = index * 360 / count
        distance = abs((angle - pa_degrees + 180) % 360 - 180)
        if distance < exclusion:
            continue
        try:
            sample = sky_measurements(images, separation_mas, angle, spatial)
            if np.all(np.isfinite(sample)):
                samples.append(sample)
                accepted_angles.append(angle)
        except ValueError:
            pass  # Missing/bad pixels exclude an aperture; the count is reported.
    samples = np.asarray(samples)
    if samples.shape[0] < samples.shape[1] + 2:
        raise ValueError(
            "Too few valid real comparison patches for this covariance dimension"
        )
    return (
        samples.mean(axis=0),
        np.cov(samples, rowvar=False) * (1 + 1 / len(samples)),
        accepted_angles,
    )


def real_response(
    data,
    dataset,
    separation,
    angle,
    measurement_separation,
    measurement_angle,
    algorithm="adi",
    spatial=False,
    step=5e-6,
    variant=None,
):
    templates = empirical_injection(dataset, data.shape[-1], separation, angle, variant)
    response = []
    for channel in range(data.shape[0]):
        injection = np.zeros_like(data)
        injection[channel] = templates[channel] * step
        plus_images = reduce_registered(data + injection, dataset, algorithm)
        minus_images = reduce_registered(data - injection, dataset, algorithm)
        plus = sky_measurements(plus_images, measurement_separation, measurement_angle, spatial)
        minus = sky_measurements(minus_images, measurement_separation, measurement_angle, spatial)
        plus_mean = real_comparison(plus_images, measurement_separation, measurement_angle, spatial)[0]
        minus_mean = real_comparison(minus_images, measurement_separation, measurement_angle, spatial)[0]
        plus -= plus_mean
        minus -= minus_mean
        response.append((plus - minus) / (2 * step))
    return np.asarray(response).T