using FFTW
using Distributions
using Plots
theme(:dao)
##


##
spectrum = zeros(300)
# spectrum[end÷2-1:end÷2+1] .= 100
# spectrum[end÷2] = 100
# spectrum[47:53] .= 100

spectrum .= 100sin.(axes(spectrum,1) .* 0.5*2π/size(spectrum,1))
spectrum .+= 30sin.(axes(spectrum,1) .* 2.5*2π/size(spectrum,1))
spectrum[end÷2-5:end÷2+5] .= 0
spectrum[end÷2+16:end÷2+20] .+= 100


spectrum[begin] = 0
spectrum[end]= 0

# read_noise = 3
mod_eff = 0.65
read_noise = 5
# mod_eff = 1
# outname = fts-disp-comparison-0eread-0mod
outname = "fts-disp-comparison-5eread-065mod"


p_temp = plot(spectrum, color=3, label="Template", xlabel="Channel", ylabel="Flux -- arb")
# savefig("fts-disp-comparison-3eread-065mod-template.svg")
savefig("$outname-template.svg")
p_temp
##

prepared = spectrum
prepared = vcat(reverse(spectrum), 0, spectrum)

# Fourier-transform the spectrum
f = fftshift(fft(prepared))

# Interfere it with itself, and find the absolute value
ideal_interfere_perfect_modulation = abs.(0.5f .+ 0.5f[end÷2+1])

ideal_interfere = maximum(ideal_interfere_perfect_modulation)/2  .+ (ideal_interfere_perfect_modulation .- maximum(ideal_interfere_perfect_modulation)/2 ).*mod_eff

# normalization = maximum(ideal_interfere)/sum(prepared)/2
sampled_interfere = Float64.(rand.(Poisson.(ideal_interfere)))
# sampled_interfere = ideal_interfere .+ rand.(Normal.(0, sqrt.(ideal_interfere)))

# Not correct anymore if there modulation efficiency is not 1
# if !(sum(spectrum) ≈ maximum(ideal_interfere*mod_eff))
#     error("transform not normalized correctly")
# end

shifted_ideal_interfere = ideal_interfere .-  maximum(sampled_interfere)/2
sampled_interfere .-= maximum(sampled_interfere)/2



plot(shifted_ideal_interfere, label="ideal")
plot!(sampled_interfere, label="sampled")

##
reconstructed_ideal = abs.(fft(shifted_ideal_interfere))
reconstructed_sampled = abs.(fft(sampled_interfere))
# reconstructed_ideal[begin]=0
# reconstructed_sampled[begin]=0
plot()
plot!(0:length(reconstructed_sampled)-1, reconstructed_sampled, label="sampled")
plot!(0:length(reconstructed_ideal)-1, reconstructed_ideal, label="ideal")
xlims!(0, length(spectrum))
# ylims!(-10,110)
# xlims!(140,160)
xlabel!("Channels")
ylabel!("ADU")

## Repeat N times

N = 10000
# zs = zeros(1000)
zs = zeros(0)

reconstructed_sampled_N = zeros(N, (length(reconstructed_sampled)+length(zs))÷ 2)
reconstructed_sampled_N_dispersive = zeros(N, length(spectrum))

for i in 1:N
    sampled_interfere .= rand.(Poisson.(ideal_interfere))  .+ round.(Int, rand(Normal.(0, read_noise), length(ideal_interfere)))
    # sampled_interfere .= ideal_interfere .+ rand.(Normal.(0, sqrt.(ideal_interfere))) .+ rand(Normal.(0, read_noise), length(ideal_interfere))
    sampled_interfere .-= maximum(sampled_interfere)/2
    reconstructed_sampled_N[i,:] .= abs.(@view ifft(vcat(sampled_interfere,zs))[end÷2+2:end])
    reconstructed_sampled_N_dispersive[i,:] .= rand.(Poisson.(spectrum)) .+ round.(Int, rand(Normal.(0, read_noise), length(spectrum)))
end


## Look at statistics across MC runs
noise = std(reconstructed_sampled_N, dims=1)[2:end]
means = mean(reconstructed_sampled_N, dims=1)[2:end]
totals = sum(reconstructed_sampled_N, dims=2)[2:end]
total_counts = mean(totals)
total_counts_noise = std(totals)

noise_disp = std(reconstructed_sampled_N_dispersive, dims=1)[:]
means_disp = mean(reconstructed_sampled_N_dispersive, dims=1)[:]
totals_disp = sum(reconstructed_sampled_N_dispersive, dims=2)[:]
total_counts_disp = mean(totals_disp)
total_counts_noise_disp = std(totals_disp)



## Comparison
p_comp = plot()
# means[begin]=0
plot!(spectrum, label="Template", color=3)
plot!(means_disp, ribbon=noise_disp, label="Dispersive", color=2)
plot!(means.*2 ./mod_eff, ribbon=noise.*2 ./mod_eff, label="FTS", color=1)
# ylims!(-10,NaN)
# # ylims!(-10,15100)
xlabel!("Channels")
ylabel!("Recovered Spectrum - arb.")
# savefig("fts-disp-comparison-3eread-065mod-spec.svg")
savefig("$outname-spec.svg")

p_comp

## Noise comparison
p_res = plot()
plot!(noise.*2  ./mod_eff, label="FTS Noise", color=1)
plot!(noise_disp, label="Dispersive Noise", color=2)
# ylims!(-1,15)
xlabel!("Channels")
ylabel!("Noise - arb.")
# savefig("fts-disp-comparison-3eread-065mod-resid.svg")
savefig("$outname-resid.svg")
p_res

##
p_res2 = plot()
noise_ratio = noise.*2  ./mod_eff ./ noise_disp[begin:length(noise)]
plot!(noise_ratio, label="Ratio", color=1)
ylims!(0, NaN)
# savefig("fts-disp-comparison-3eread-065mod-ratio.svg")
savefig("$outname-ratio.svg")
p_res2