from .wavefront_sensor import WavefrontSensorOptics, WavefrontSensorEstimator
from ..optics import OpticalSystem, MicroLensArray
from ..field import CartesianGrid, Field, SeparatedCoords
from ..propagation import FresnelPropagator

import numpy as np
from scipy import ndimage

class ShackHartmannWavefrontSensorOptics(WavefrontSensorOptics):
    def __init__(self, input_grid, micro_lens_array):
        # Make propagator
        sh_prop = FresnelPropagator(input_grid, micro_lens_array.focal_length)

        # Make optical system
        OpticalSystem.__init__(self, (micro_lens_array, sh_prop))
        self.mla_index = micro_lens_array.mla_index
        self.mla_grid = micro_lens_array.mla_grid
        self.micro_lens_array = micro_lens_array

class SquareShackHartmannWavefrontSensorOptics(ShackHartmannWavefrontSensorOptics):
    ## Helper class to create a Shack-Hartmann WFS with square microlens array
    def __init__(self, input_grid, f_number, num_lenslets, pupil_diameter):
        lenslet_diameter = float(pupil_diameter) / num_lenslets
        x = np.arange(-pupil_diameter, pupil_diameter, lenslet_diameter)
        self.mla_grid = CartesianGrid(SeparatedCoords((x, x)))

        focal_length = f_number * lenslet_diameter
        self.micro_lens_array = MicroLensArray(input_grid, self.mla_grid, focal_length)

        ShackHartmannWavefrontSensorOptics.__init__(self, input_grid, self.micro_lens_array)

class ShackHartmannWavefrontSensorEstimator(WavefrontSensorEstimator):
    def __init__(self, mla_grid, mla_index, estimation_subapertures=None):
        self.mla_grid = mla_grid
        self.mla_index = mla_index
        if estimation_subapertures is None:
            self.estimation_subapertures = np.unique(self.mla_index)
        else:
            self.estimation_subapertures = np.flatnonzero(np.array(estimation_subapertures))
        self.estimation_grid = self.mla_grid.subset(estimation_subapertures)

    def estimate(self, images):
        image = images[0]

        fluxes = ndimage.measurements.sum(image, self.mla_index, self.estimation_subapertures)
        sum_x = ndimage.measurements.sum(image * image.grid.x, self.mla_index, self.estimation_subapertures)
        sum_y = ndimage.measurements.sum(image * image.grid.y, self.mla_index, self.estimation_subapertures)

        centroid_x = sum_x / fluxes
        centroid_y = sum_y / fluxes

        centroids = np.array((centroid_x, centroid_y)) - np.array(self.mla_grid.points[self.estimation_subapertures, :]).T
        return Field(centroids, self.estimation_grid)
