from hcipy import *
from matplotlib import pyplot as plt
import numpy as np

def test_pupil_plane_wavefront_sensors():
    # Create the input grid
    D = 1
    pupil_grid = make_pupil_grid(128, D)
    aperture = make_circular_aperture(D)

    # Make zernike basis
    num_zernike = 3
    zernike_basis = make_zernike_basis(num_zernike, D, pupil_grid, 2)

    # Create wavefront sensors
    pyramid_wfs = PyramidWavefrontSensorOptics(pupil_grid)
    zernike_wfs = ZernikeWavefrontSensorOptics(pupil_grid, q=32, num_airy=20)
    god_wfs = gODWavefrontSensorOptics(40, 1/3, pupil_grid, pupil_separation=1.5, q=16, num_airy=40)
    pod_wfs = PolgODWavefrontSensorOptics(40, 1/3, pupil_grid, pupil_separation=1.5, q=16, num_airy=40)

    wavefront_sensors = (pyramid_wfs, zernike_wfs, god_wfs, pod_wfs)
    num_wfs = len(wavefront_sensors)

    # Create reference images
    wf = Wavefront(aperture(pupil_grid))
    wf.total_power = 1
    refs = [wfs.forward(wf).intensity for wfs in wavefront_sensors]
    eff = [wfs.forward(wf).total_power for wfs in wavefront_sensors]
    print("Total efficiency : ", eff)

    for i in range(num_wfs):
        plt.subplot(2,2,i+1)
        imshow_field(refs[i], wavefront_sensors[i].output_grid)
    plt.show()

    for i, zern in enumerate(zernike_basis):
        # create input
        wf = Wavefront(aperture(pupil_grid) * np.exp(1j * 0.01 * zern))
        wf.total_power = 1

        # Measure the response of all wavefront sensors
        for j, wfs in enumerate(wavefront_sensors):
            wf_out = wfs.forward(wf)
            print(wf_out.total_power)
            plt.subplot(num_wfs, num_zernike, i+1 + num_zernike*j)
            imshow_field(wf_out.intensity-refs[j])

    plt.show()



if __name__ == "__main__":
    test_pupil_plane_wavefront_sensors()
