# -*- coding: utf-8 -*-
"""
Created on Sun May 30 12:21:19 2021

There is a difference between the theoretical and numerical values
of the sigma² parameter. This is due to numerical sampling and
the forced central frequency to be null.

@author: rfetick
"""

import matplotlib.pyplot as plt
import numpy as np
from maoppy.psfmodel import Psfao
from maoppy.instrument import muse_nfm

npix = 128 # pixel size of PSF
wvl = 600*1e-9 # wavelength [m]

#%% Initialize PSF model
samp = muse_nfm.samp(wvl) # sampling (2.0 for Shannon-Nyquist)
Pmodel = Psfao((npix,npix),system=muse_nfm,samp=samp)

#%% Generate a MUSE-NFM PSF
r0 = 0.15 # Fried parameter [m]
b = 0 # Phase PSD background [rad² m²]
amp = 1.0 # Phase PSD Moffat amplitude [rad²]
alpha = np.logspace(-3,0,num=50) # Phase PSD Moffat alpha [1/m]
ratio = 1.0
theta = 0.0
beta = [1.4,1.5,1.6] # Phase PSD Moffat beta power law

sig2num = np.zeros((len(alpha),len(beta)))

f2D = Pmodel._fxfy
F2 = f2D[0]**2. + f2D[1]**2.
Fao = Pmodel.system.Nact/(2.0*Pmodel.system.D)
df = Pmodel._pix2freq
mask = F2 < Fao**2

for i in range(len(alpha)):
    for j in range(len(beta)):
        param = [r0,b,amp,alpha[i],ratio,theta,beta[j]]
        psd,integral = Pmodel.psd(param)
        sig2num[i,j] = np.sum(psd*mask)*df**2

#%% PLOT
plt.figure(1)
plt.clf()
for j in range(len(beta)):
    plt.semilogx(alpha/df,sig2num[:,j],label='numerical sum ($\\beta=$%.1f)'%beta[j])
plt.semilogx(alpha/df,alpha*0 + amp, label='required')
plt.grid(which='both')
plt.xlabel("$\\alpha$/df")
plt.ylabel("$\\sigma^2$ [rad²]")
plt.legend()
plt.ylim(0,1.2)