"""Independent NumPy FFT reference, with the same continuous circular mask and independent focal-cell quadrature."""

import numpy as np
from imaging_reduction import resample
from focal_mask import circular_transmission, cell_coverage
from functools import lru_cache


def pupil_fields(
    settings, wavelength=1.24, source=(5, 0), exposure=0, size=512, diameter=64, pixel_nodes=2
):
    rows, columns = np.indices((size, size))
    x, y = (columns - size / 2) / diameter, (size / 2 - rows) / diameter
    pupil = circular_aperture(size, diameter, 0.5)
    static = (
        np.cos(2 * np.pi * (4 * x + 2 * y))
        + 0.55 * np.sin(2 * np.pi * (6 * x - y))
        + 0.3 * (x * x - y * y)
    )
    evolving = np.sin(2 * np.pi * (3 * x - 4 * y)) + 0.4 * np.cos(
        2 * np.pi * (2 * x + 5 * y)
    )
    for pattern in (static, evolving):
        pattern -= np.sum(pattern * pupil) / pupil.sum()
        pattern /= np.sqrt(np.sum(pattern ** 2 * pupil) / pupil.sum())
    opd = (
        settings["staticOpdNm"] * static
        + settings["evolvingOpdNm"] * np.sin(2 * np.pi * exposure / 17) * evolving
    ) / 1000
    phase = (
        2
        * np.pi
        * ((source[0] * x + source[1] * y) * 1.24 / wavelength + opd / wavelength)
    )
    entrance = pupil * np.exp(1j * phase)
    fft = lambda field: np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(field)))
    focus = fft(entrance)
    focal_radius = np.hypot(columns - size / 2, rows - size / 2) * diameter / size
    focus *= circular_transmission(size, 2.5 * 1.24 / wavelength, diameter / size)
    lyot = fft(focus) / size**2
    lyot *= circular_aperture(size, diameter, .41)
    final = fft(lyot)
    normalization = np.pi * diameter ** 2 / 4 * size ** 2
    # Independent complex-amplitude cubic sampling, then nonnegative intensity.
    output_size, output_sample = 80, 0.25
    native_sample = wavelength / 1.24 * diameter / size
    output_rows, output_columns = np.indices((output_size, output_size))
    result = np.zeros((output_size, output_size))
    if pixel_nodes == 2:
        for offset_y in (-1 / (2 * np.sqrt(3)), 1 / (2 * np.sqrt(3))):
            for offset_x in (-1 / (2 * np.sqrt(3)), 1 / (2 * np.sqrt(3))):
                horizontal = size / 2 - (output_columns - output_size / 2 + offset_x) * output_sample / native_sample
                vertical = size / 2 - (output_rows - output_size / 2 + offset_y) * output_sample / native_sample
                sampled = complex_cubic(final, horizontal, vertical)
                result += np.abs(sampled) ** 2
        return result * (output_sample / native_sample) ** 2 / (4 * normalization)
    nodes, weights = np.polynomial.legendre.leggauss(pixel_nodes)
    for offset_y, weight_y in zip(nodes / 2, weights / 2):
        for offset_x, weight_x in zip(nodes / 2, weights / 2):
            horizontal = size / 2 - (output_columns - output_size / 2 + offset_x) * output_sample / native_sample
            vertical = size / 2 - (output_rows - output_size / 2 + offset_y) * output_sample / native_sample
            sampled = complex_cubic(final, horizontal, vertical)
            result += weight_x * weight_y * np.abs(sampled) ** 2
    return result * (output_sample / native_sample) ** 2 / normalization



@lru_cache(maxsize=20)
def circular_aperture(size, diameter, radius):
    rows, columns = np.indices((size, size))
    x, y = (columns - size / 2) / diameter, (size / 2 - rows) / diameter
    half = .5 / diameter
    result = (np.hypot(x, y) < radius).astype(float)
    boundary = ((np.hypot(abs(x) + half, abs(y) + half) > radius)
        & (np.hypot(np.maximum(0, abs(x) - half), np.maximum(0, abs(y) - half)) < radius))
    for row, column in np.argwhere(boundary):
        result[row, column] = cell_coverage(x[row, column], y[row, column], half, radius)
    return result


def complex_cubic(image, x, y):
    first_x, first_y = np.floor(x).astype(int), np.floor(y).astype(int)
    result = np.zeros(x.shape, dtype=complex)
    def kernel(distance):
        absolute = abs(distance)
        return np.where(absolute < 1, 1.5 * absolute ** 3 - 2.5 * absolute ** 2 + 1,
            np.where(absolute < 2, -.5 * absolute ** 3 + 2.5 * absolute ** 2 - 4 * absolute + 2, 0))
    for dy in (-1, 0, 1, 2):
        for dx in (-1, 0, 1, 2):
            column, row = first_x + dx, first_y + dy
            valid = (column >= 0) & (row >= 0) & (column < image.shape[1]) & (row < image.shape[0])
            result[valid] += image[row[valid], column[valid]] * kernel(x[valid] - column[valid]) * kernel(y[valid] - row[valid])
    return result


def finite_band(settings, center, source, exposure, nodes=3):
    abscissas, weights = np.polynomial.legendre.leggauss(nodes)
    wavelengths = center + 0.01 * abscissas
    photon_weights = weights * wavelengths
    return (
        sum(
            weight * pupil_fields(settings, wavelength, source, exposure)
            for wavelength, weight in zip(wavelengths, photon_weights)
        )
        / photon_weights.sum()
    )


def source_templates(settings, angles, wavelengths, x, y):
    result = []
    for wavelength in wavelengths:
        frames = []
        for exposure, angle in enumerate(angles):
            detector = (
                np.cos(angle) * x - np.sin(angle) * y,
                np.sin(angle) * x + np.cos(angle) * y,
            )
            frames.append(finite_band(settings, wavelength, detector, exposure))
        result.append(frames)
    return np.asarray(result)


def stellar_expectation(settings, angles, wavelengths):
    return (
        3e7
        * np.asarray(
            [
                [
                    finite_band(settings, wavelength, (0, 0), exposure)
                    for exposure in range(len(angles))
                ]
                for wavelength in wavelengths
            ]
        )
        + 5
    )