"""Independent NumPy implementation of the synthetic sampling/reduction contract.

This is science tutorial code, not a browser Engine. Array axes are wavelength,
exposure, detector row, detector column. Coordinates are x right and y up.
"""

import math
import numpy as np
from image_tools import bilinear


def resample(image, scale=1, angle=0):
    size = image.shape[0]
    rows, columns = np.indices(image.shape)
    x, y = columns - size / 2, size / 2 - rows
    cosine, sine = math.cos(angle), math.sin(angle)
    return (
        np.nan_to_num(
            bilinear(
                image,
                size / 2 + (cosine * x - sine * y) / scale,
                size / 2 - (sine * x + cosine * y) / scale,
            ),
            nan=0,
        )
        / scale**2
    )


def reduce_cube(data, angles, wavelengths, algorithm="adi", rank=1):
    frames = data.copy()
    if algorithm == "asdi":
        for exposure in range(data.shape[1]):
            aligned = np.stack(
                [
                    resample(frames[channel, exposure], scale=1.24 / wavelength)
                    for channel, wavelength in enumerate(wavelengths)
                ]
            )
            aligned -= aligned.mean(axis=0)
            for channel, wavelength in enumerate(wavelengths):
                frames[channel, exposure] = resample(
                    aligned[channel], scale=wavelength / 1.24
                )
    output = []
    for channel in range(data.shape[0]):
        centered = frames[channel] - frames[channel].mean(axis=0)
        if algorithm == "klip":
            matrix = centered.reshape(data.shape[1], -1)
            _, _, vectors = np.linalg.svd(matrix, full_matrices=False)
            matrix = matrix - (matrix @ vectors[:rank].T) @ vectors[:rank]
            centered = matrix.reshape(centered.shape)
        output.append(
            np.mean(
                [
                    resample(frame, angle=angle)
                    for frame, angle in zip(centered, angles)
                ],
                axis=0,
            )
        )
    return np.stack(output)


def aperture_weights(size, x, y, radius=0.5, sample=0.25):
    rows, columns = np.indices((size, size))
    distance_x = columns - size / 2 - x / sample
    distance_y = rows - size / 2 + y / sample
    weights = np.zeros((size, size))
    for offset_y in (-0.375, -0.125, 0.125, 0.375):
        for offset_x in (-0.375, -0.125, 0.125, 0.375):
            weights += (
                np.hypot(distance_x + offset_x, distance_y + offset_y)
                <= radius / sample
            )
    return weights / 16


def spatial_weights(size, x, y, spatial=True):
    center = aperture_weights(size, x, y)
    if not spatial:
        return np.asarray([center])
    angle = math.atan2(y, x)
    dx, dy = 0.4 * math.cos(angle), 0.4 * math.sin(angle)
    radial = aperture_weights(size, x + dx, y + dy) - aperture_weights(
        size, x - dx, y - dy
    )
    tangent = aperture_weights(size, x - dy, y + dx) - aperture_weights(
        size, x + dy, y - dx
    )
    return np.asarray([center, radial, tangent])


def background_subtracted_weights(size, settings, spatial=False):
    separation = settings["separation"]
    source_angle = math.radians(settings["positionAngleDegrees"])
    weights = spatial_weights(size, separation * math.cos(source_angle), separation * math.sin(source_angle), spatial)
    count = int(2 * math.pi * separation / 1.1)
    references = []
    for index in range(count):
        angle = 2 * math.pi * index / count
        delta = abs(math.atan2(math.sin(angle - source_angle), math.cos(angle - source_angle)))
        if delta >= math.radians(settings["rotationDegrees"]) + 2 / separation:
            references.append(spatial_weights(size, separation * math.cos(angle), separation * math.sin(angle), spatial))
    return weights - np.mean(references, axis=0)


def measurements(images, weights):
    return np.einsum("cij,kij->ck", images, weights).ravel()


def comparison_covariance(images, settings, spatial=False):
    separation = settings["separation"]
    source_angle = math.radians(settings["positionAngleDegrees"])
    count = int(2 * math.pi * separation / 1.1)
    values = []
    for index in range(count):
        angle = 2 * math.pi * index / count
        delta = abs(
            math.atan2(math.sin(angle - source_angle), math.cos(angle - source_angle))
        )
        if delta < math.radians(settings["rotationDegrees"]) + 2 / separation:
            continue
        values.append(
            measurements(
                images,
                spatial_weights(
                    images.shape[-1],
                    separation * math.cos(angle),
                    separation * math.sin(angle),
                    spatial,
                ),
            )
        )
    values = np.asarray(values)
    if values.shape[0] < values.shape[1] + 2:
        raise ValueError("Insufficient comparison patches for empirical covariance")
    return (
        values.mean(axis=0),
        np.cov(values, rowvar=False) * (1 + 1 / values.shape[0]),
        values.shape[0],
    )


def finite_response(
    data, templates, angles, wavelengths, weights, algorithm="adi", rank=1, step=2e-5
):
    response = []
    for channel in range(data.shape[0]):
        injection = np.zeros_like(data)
        injection[channel] = templates[channel] * 3e7 * step
        plus = reduce_cube(data + injection, angles, wavelengths, algorithm, rank)
        minus = reduce_cube(data - injection, angles, wavelengths, algorithm, rank)
        response.append(measurements(plus - minus, weights) / (2 * step))
    return np.asarray(response).T


def fit_spectrum(observed, response, covariance):
    eigenvalues = np.linalg.eigvalsh(covariance)
    if eigenvalues[0] <= eigenvalues[-1] * 1e-10:
        raise ValueError("Empirical covariance is rank deficient")
    weighted_response = np.linalg.solve(covariance, response)
    normal = response.T @ weighted_response
    normal_eigenvalues = np.linalg.eigvalsh(normal)
    if (
        normal_eigenvalues[0] <= normal_eigenvalues[-1] * 1e-8
        or np.max(np.abs(response)) < 3e7 * 1e-8
    ):
        raise ValueError("Optical response is rank deficient")
    fitted_covariance = np.linalg.inv(normal)
    flux = fitted_covariance @ response.T @ np.linalg.solve(covariance, observed)
    residual = observed - response @ flux
    chi_square = float(residual @ np.linalg.solve(covariance, residual))
    logdet = float(np.linalg.slogdet(covariance)[1])
    return {
        "contrast": flux,
        "covariance": fitted_covariance,
        "chi_square": chi_square,
        "negative2loglikelihood": chi_square
        + logdet
        + len(observed) * math.log(2 * math.pi),
        "logdet": logdet,
        "condition": float(eigenvalues[-1] / eigenvalues[0]),
    }