"""Independent exact transpose for fixed ADI/ASDI source-photon covariance."""

import math
import numpy as np
from imaging_reduction import reduce_cube, measurements


def sampling_adjoint(weights, scale=1, angle=0):
    size = weights.shape[0]
    rows, columns = np.indices(weights.shape)
    x, y = columns - size / 2, size / 2 - rows
    cosine, sine = math.cos(angle), math.sin(angle)
    xx, yy = (
        size / 2 + (cosine * x - sine * y) / scale,
        size / 2 - (sine * x + cosine * y) / scale,
    )
    left, top = np.floor(xx).astype(int), np.floor(yy).astype(int)
    fraction_x, fraction_y = xx - left, yy - top
    result = np.zeros(weights.shape)
    for offset_y in (0, 1):
        for offset_x in (0, 1):
            column, row = left + offset_x, top + offset_y
            factor = (fraction_x if offset_x else 1 - fraction_x) * (
                fraction_y if offset_y else 1 - fraction_y
            )
            inside = (column >= 0) & (row >= 0) & (column < size) & (row < size)
            np.add.at(
                result,
                (row[inside], column[inside]),
                (weights * factor / scale**2)[inside],
            )
    return result


def reduction_adjoint(selected_channel, weight, angles, wavelengths, algorithm="adi"):
    result = np.zeros((len(wavelengths), len(angles), *weight.shape))
    angular = np.asarray(
        [sampling_adjoint(weight, angle=angle) / len(angles) for angle in angles]
    )
    angular -= angular.mean(axis=0)
    result[selected_channel] = angular
    if algorithm == "asdi":
        for exposure in range(len(angles)):
            common = np.asarray(
                [
                    sampling_adjoint(result[channel, exposure], scale=wavelength / 1.24)
                    for channel, wavelength in enumerate(wavelengths)
                ]
            )
            common -= common.mean(axis=0)
            for channel, wavelength in enumerate(wavelengths):
                result[channel, exposure] = sampling_adjoint(
                    common[channel], scale=1.24 / wavelength
                )
    return result


def unit_source_covariances(templates, weights, angles, wavelengths, algorithm="adi"):
    adjoints = np.asarray(
        [
            reduction_adjoint(channel, weight, angles, wavelengths, algorithm)
            for channel in range(len(wavelengths))
            for weight in weights
        ]
    )
    result = []
    for channel in range(len(wavelengths)):
        selected = adjoints[:, channel].reshape(len(adjoints), -1)
        variance = templates[channel].ravel() * 3e7
        if not np.all(np.isfinite(selected)) or not np.all(np.isfinite(variance)):
            raise ValueError("Nonfinite source-noise propagation inputs")
        # Direct contraction avoids platform BLAS floating-status artifacts.
        propagated = np.einsum("ik,jk,k->ij", selected, selected, variance, optimize=False)
        if not np.all(np.isfinite(propagated)):
            raise ValueError("Nonfinite source-noise covariance")
        result.append(propagated)
    return np.asarray(result)


def bootstrap_source_covariance(
    data,
    templates,
    contrasts,
    weights,
    angles,
    wavelengths,
    rng,
    algorithm="klip",
    rank=1,
    draws=32,
    fixed_basis=False,
):
    variance = templates * np.maximum(contrasts, 0)[:, None, None, None] * 3e7
    baseline = measurements(
        reduce_cube(data, angles, wavelengths, algorithm, rank), weights
    )
    # Fixed-basis comparator is intentionally distinct from the recomputed basis.
    bases = None
    if fixed_basis:
        bases = []
        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
            )
            bases.append(vectors[:rank])
    samples = []
    for _ in range(draws):
        noise = rng.poisson(variance) - variance
        if fixed_basis:
            from imaging_reduction import resample

            images = []
            for channel, basis in enumerate(bases):
                centered = noise[channel] - noise[channel].mean(axis=0)
                matrix = centered.reshape(data.shape[1], -1)
                residual = (matrix - (matrix @ basis.T) @ basis).reshape(centered.shape)
                images.append(
                    np.mean(
                        [
                            resample(frame, angle=angle)
                            for frame, angle in zip(residual, angles)
                        ],
                        axis=0,
                    )
                )
            samples.append(measurements(np.asarray(images), weights))
        else:
            samples.append(
                measurements(
                    reduce_cube(data + noise, angles, wavelengths, algorithm, rank),
                    weights,
                )
                - baseline
            )
    return np.cov(np.asarray(samples), rowvar=False)