import numpy as np
from astropy.io import fits

from hcipy import Field, Wavefront, DeformableMirror, FraunhoferPropagator
from hcipy import SurfaceApodizer, SurfaceAberrationAtDistance
from hcipy import make_pupil_grid, make_circular_aperture, make_focal_grid
from hcipy import make_xinetics_influence_functions
from hcipy import read_fits

from .config import cpism_refdata, S
from .optics import filter_throughput


# initial psf simulation
apm, apm_header = fits.getdata(
    f'{cpism_refdata}/optics/apm.fits', header=True)
actuator = read_fits(f'{cpism_refdata}/optics/actuator.fits')
surface_aberration = read_fits(
    f'{cpism_refdata}/optics/initial_phase_aberration.fits')

wavelength = 625e-9  # m
pupil_diameter = 2  # m
focal_length = pupil_diameter * 83
pupil_grid = make_pupil_grid(apm.shape[0], apm.shape[0] * apm_header['PHRATE'])
aperture = make_circular_aperture(pupil_diameter)(pupil_grid)
aperture = aperture * Field(apm.flatten(), pupil_grid)


emccd_pixel_size = 13e-6  # m
arcsec2rad = np.radians(1 / 3600)
emccd_pixel_scale = emccd_pixel_size / \
    (arcsec2rad * focal_length)  # arcsec / pixel
spatial_resolution = focal_length * wavelength / \
    pupil_diameter  # meter per airy disk

q = spatial_resolution / emccd_pixel_size  # pixels per airy disk
num_airy = 512 / q  # airy disk per frame (2 * 512 = 1024 pix)
focal_full_frame = make_focal_grid(
    q, num_airy, spatial_resolution=spatial_resolution)
prop_full_frame = FraunhoferPropagator(
    pupil_grid, focal_full_frame, focal_length)

num_airy = 128 / q  # make a small frame for the planets
focal_sub_frame = make_focal_grid(
    q, num_airy, spatial_resolution=spatial_resolution)
prop_sub_frame = FraunhoferPropagator(
    pupil_grid, focal_sub_frame, focal_length)

num_actuators_across = 32
# dm spacing is little smaller than pupil
actuator_spacing = 0.95 / 32 * pupil_diameter
influence_functions = make_xinetics_influence_functions(
    pupil_grid, num_actuators_across, actuator_spacing)
deformable_mirror = DeformableMirror(influence_functions)

aberration = SurfaceApodizer(
    surface_sag=surface_aberration.flatten(), refractive_index=-1)
# arbitrary distance for the aberration to propagate
aberration_distance = 80 * focal_length
aberration = SurfaceAberrationAtDistance(aberration, aberration_distance)


def simulate_psf(error_value, band, spectrum, nsample=5, cstar=True, aber_phase=None):

    filter = filter_throughput(band)
    wave = filter.wave
    throughput = filter.throughput
    min_wave = wave[0]
    max_wave = wave[-1]

    sed = []
    sed_center_wavelength = []
    for i_wave in range(nsample):
        d_wave = (max_wave - min_wave) / nsample
        wave0 = min_wave + i_wave * d_wave
        wave1 = min_wave + (i_wave + 1) * d_wave
        sed_center_wavelength.append((wave0 + wave1) / 2 * 1e-10)

        i_throughput = throughput.copy()
        i_throughput[(wave > wave1) | (wave < wave0)] = 0
        i_band = S.ArrayBandpass(wave, i_throughput, waveunits='angstrom')
        i_sed = (spectrum * i_band).integrate()
        sed.append(i_sed)

    error = np.random.normal(0, error_value*1e-9, actuator.shape)


    imgs = []
    deformable_mirror.actuators = actuator + error

    prop = prop_full_frame
    if not cstar:
        prop = prop_sub_frame

    for i_sample in range(nsample):
        wf = Wavefront(aperture, sed_center_wavelength[i_sample])
        wf = aberration(wf)
        if aber_phase is not None:
            other_error = SurfaceApodizer(
                surface_sag=aber_phase.flatten(), refractive_index=-1)
            wf = other_error(wf)

        img = prop(deformable_mirror(wf)).intensity.shaped
        imgs.append(img / img.sum() * sed[i_sample])

    return np.array(imgs).sum(axis=0)