using FFTW, LinearAlgebra, Plots, ColorSchemes, StatsBase

## Reference:
#=
    A Phase retrival algorithm for shifting illumination
        J. M. Rodenburg and H. M. L. Faulkner
        15 November 2004
        Applied Physics Letters
        Vol. 85
        No. 20
        https://aip.scitation.org/doi/10.1063/1.1823034
=#

# Setting up system and directory
cd(".../04_RodenburgFaulkner")

include(".../02_ForwardModel/ForwardModel.jl");
#= imports:    
    Ï‰: freq domain
    Î”Ï‰: resolution of Ï‰
    t: time domain
    Î”t: resolution of t

    object: object âˆˆ â„‚
    obj_amp: object Amplitude
    obj_phase: object phase
    Probe_set : set of dimension 1:2001 Ã— -M_max:M_max
    Int_Measured : Set of internsity measurements 2001 Ã— M_max
=#

# 2001Ã—M_max OffsetArray with indices 1:2001 Ã— -M:M
Int_Measured = OffsetArray(Int_Measured, 1:length(Ï‰), -m_end:m_end);

#=
Algorithm
    1. Start with a guess of the object function ğ�‘‚áµ§,â‚™ (ğ��«). Î³,n represents a 
        guessed function at the ğ�‘›th iteration of the algorithm. ğ�‘‚ âˆˆ â„�.   
    2. Multiply current guess at the object function by the illumination function
        at the current position ğ��‘, ğ�‘ƒ(ğ��«-ğ��‘). This produces the guessed exit wave 
        function for position ğ��‘,
            Ïˆáµ§,â‚™(ğ��«,ğ��‘) = ğ�‘‚áµ§,â‚™(ğ��«).ğ�‘ƒ(ğ��«-ğ��‘)
    3. Transform Ïˆáµ§,â‚™(ğ��«,ğ��‘) to obtain the corresponding wave function in the 
        diffraction space plane, for that position ğ��‘,
            Î¨áµ§,â‚™(ğ��¤,ğ��‘) = â„±[Ïˆáµ§,â‚™(ğ��«,ğ��‘)]
        ğ��¤ is the usual reciprocal space coordinate. Î¨ is a "guessed" version of the
        actual wave function in diffraction space. 
        Decompose into Amplitude and Phase:
            Î¨áµ§,â‚™(ğ��¤,ğ��‘) = |Î¨áµ§,â‚™(ğ��¤,ğ��‘)|exp(iÎ¸áµ§,â‚™(ğ��¤,ğ��‘))    
    4. Correct the intensities of the guessed diffraction space wave function to 
        known values,
            Î¨c,n(ğ��¤,ğ��‘) = |Î¨(ğ��¤,ğ��‘)|exp(iÎ¸áµ§,â‚™(ğ��¤,ğ��‘))
        where |Î¨(ğ��¤,ğ��‘)| is known modulus.
    5. Inverse transform back to real space to obtain the new and improved guess at 
        the exit wave function
            Ïˆc,n(ğ��«,ğ��‘) = â„± â�»Â¹[Î¨c,n(ğ��¤,ğ��‘)]
    6. Update the guessed object wave function in the area covered by the aperture 
        or probe, using the update function
            ğ�‘‚áµ§,â‚™â‚Šâ‚�(ğ��«) = ğ�‘‚áµ§,â‚™(ğ��«) + (|P(ğ��«-ğ��‘)|/|Pâ‚˜â‚�â‚“(ğ��«-ğ��‘)|)(P*(ğ��«-ğ��‘)/(|P(ğ��«-ğ��‘)|Â²+Î±)) 
                        Ã— Î²(Ïˆc,n(ğ��«,ğ��‘)- Ïˆáµ§,â‚™(ğ��¤,ğ��‘))
    7. Move to the next position of ğ��‘, for which the illumination in part overlaps
        with that of the previous position.
    8. Repeat (2) - (7) until the sum squared error (SSE) is sufficiently small. 
        The SSE is measured in the diffraction plane as
            SSE = (|Î¨(ğ��¤,ğ��‘)|Â²-|Î¨áµ§,â‚™(ğ��¤,ğ��‘)|Â²)Â²/N
        where N is the number of pixels in the array representing the wave function.  
=#

# Initial vectors and variables
O_i_amp = peakâ‚‹â‚‚ .+ peakâ‚‹â‚� .+ peakâ‚‚; #initial guess of object
O_i_amp = normalize(O_i_amp);
phase_0 = Ï‰;
N = length(Ï‰);  # number of pixels in the array
probe_offset = -2; #Offset of Probe
P_1 = Probe_set[:,probe_offset];

#= Step 1:
    Start with an initial guess of the object function. ğ�‘‚_1 âˆˆ â„�?  
=#
O_1 = O_i_amp;  

#= Step 2:
    Produce the guessed exit wave function for position ğ��‘,
        Ïˆáµ§,â‚™(ğ��«,ğ��‘) = ğ�‘‚áµ§,â‚™(ğ��«).ğ�‘ƒ(ğ��«-ğ��‘)
=#
Ïˆ_1 = O_1.*P_1;

#= Step 3:
    Transform Ïˆáµ§,â‚™(ğ��«,ğ��‘) to get the wave function in the diffraction space plane
        Î¨áµ§,â‚™(ğ��¤,ğ��‘) = â„±[Ïˆáµ§,â‚™(ğ��«,ğ��‘)]
    Î¨ is a "guessed" version of the actual wave function in diffraction space. 
    Decompose into Amplitude and Phase  
=#

Î¨_Î³_1 = fftshift(fft(Ïˆ_1));   # guessed version of wave function
abs_Î¨_Î³_1 = abs.(Î¨_Î³_1);        # Amplitude of guess
phase_Î¨_Î³_1 = angle.(Î¨_Î³_1);    # phase of guess

#=
4. Correct the intensities of the guessed diffraction space wave function to 
    known values,
        Î¨c,n(ğ��¤,ğ��‘) = |Î¨(ğ��¤,ğ��‘)|exp(iÎ¸áµ§,â‚™(ğ��¤,ğ��‘))
    where |Î¨(ğ��¤,ğ��‘)| is known modulus.
=#
Î¨_c_1 = Int_Measured[:,probe_offset].*exp.(im.*phase_Î¨_Î³_1);

#= Step 5
 Inverse transform back to real space to obtain the new and improved guess at 
the exit wave function
    Ïˆc,n(ğ��«,ğ��‘) = â„± â�»Â¹[Î¨c,n(ğ��¤,ğ��‘)]
=#
Ïˆ_c_1 = ifft(ifftshift(Î¨_c_1));

#= Step 6
6. Update the guessed object wave function in the area covered by the aperture 
or probe, using the update function
    ğ�‘‚áµ§,â‚™â‚Šâ‚�(ğ��«) = ğ�‘‚áµ§,â‚™(ğ��«) + (|P(ğ��«-ğ��‘)|/|Pâ‚˜â‚�â‚“(ğ��«-ğ��‘)|)(P*(ğ��«-ğ��‘)/(|P(ğ��«-ğ��‘)|Â²+Î±)) 
                Ã— Î²(Ïˆc,n(ğ��«,ğ��‘)- Ïˆáµ§,â‚™(ğ��¤,ğ��‘))
=#
# update function
function O_new(O_old, P, Î±, Î², Ïˆ_c, Ïˆ_g)
    P = vec(P);
    O_old = vec(O_old);
    Ïˆ_c = vec(Ïˆ_c);
    Ïˆ_g = vec(Ïˆ_g);
    return O_old .+ ((norm(P)./abs(maximum(abs.(P)))).*(vec(conj.(P))./((norm(P))^2 .+ Î±))) .* (Î².*(Ïˆ_c-Ïˆ_g))
end

O_2 = O_new(O_1, P_1, 0.0, 0.1, Ïˆ_c_1, Î¨_Î³_1);

