__precompile__()
module Models

    # This module stands for modelling the observational data

    include("EoMs.jl")
    using DifferentialEquations
    using Roots
    using .EoMs

    export model_massless, c_to_kms

    rad_to_arcsec = 3600.0 * 1000.0 * 180.0 / π
    c_to_kms = 299792.458
    r̂_to_km = 1476625.03852
    pc_to_km = 3600.0 * 180.0 / π * 149697870.7
    r̂_to_kpc = r̂_to_km / pc_to_km / 1000.0
    t̂_to_sec = 4.92549094921
    yr_to_sec = 365.24219878 * 86400.0
    t̂_to_yr = t̂_to_sec / yr_to_sec

    t_beg = 1992.224 / t̂_to_yr
    t_end = 2019.357 / t̂_to_yr

    function turnorbit(x, y, ω_gr, Ω_gr, i_gr)

        N = size(x)

        r0 = [x, y, 0.0]

        ω = ω_gr / 180.0 * π
        Ω = Ω_gr / 180.0 * π + π / 2.0
        i = i_gr / 180.0 * π

        cω = cos(ω)
        sω = sin(ω)
        cΩ = cos(Ω)
        sΩ = sin(Ω)
        ci = cos(i)
        si = sin(i)

        Ω_turn = [cΩ -sΩ 0.0; sΩ cΩ 0.0; 0.0 0.0 1.0]
        i_turn = [1.0 0.0 0.0; 0.0 ci -si; 0.0 si ci]
        ω_turn = [cω -sω 0.0; sω cω 0.0; 0.0 0.0 1.0]

        r = Ω_turn * i_turn * ω_turn * r0

        return r

    end

    function solve_EoM_massless(x0, y0, vx0, vy0, M, β, γ; tol=(1e-7, 1e-10))

        p = [M, β, γ]
        u0 = [x0, y0, vx0, vy0]

        tspan = (t_beg, t_end)
        prob = ODEProblem(EoM_massless, u0, tspan, p, reltol=tol[1], abstol=tol[2])
        sol = solve(prob)

        return sol

    end

    function Rømer_delay_strict(sol, Ω, i, t_obs)

        z = t -> turnorbit(sol(t)[1], sol(t)[2], 0.0, Ω, i)[3]
        f = t_em_arg -> (t_em_arg + z(t_em_arg) - t_obs)

        convkey = try
            find_zero(f, t_obs)
        catch e
            e
        end

        if typeof(convkey) == Roots.ConvergenceFailed
            return t_obs
        else
            return find_zero(f, t_obs)
        end

    end

    function Rømer_delay(sol, Ω, i, 𝐭_obs)

        z = t -> turnorbit(sol(t)[1], sol(t)[2], 0.0, Ω, i)[3]
        return z.(𝐭_obs) .+ 𝐭_obs

    end

    function model_massless(θ, 𝐭_input)
    
        𝐭_input = 𝐭_input ./ t̂_to_yr

        S2_x0, S2_y0, S2_vx0, S2_vy0, S2_i, S2_Ω, S38_x0, S38_y0, S38_vx0, S38_vy0, S38_i, S38_Ω, S55_x0, S55_y0, S55_vx0, S55_vy0, S55_i, S55_Ω, M, R0, vx_SgrA, vy_SgrA, β, γ = θ

        S2_sol  = solve_EoM_massless(S2_x0, S2_y0, S2_vx0 / c_to_kms, S2_vy0 / c_to_kms, M, β, γ)
        S38_sol = solve_EoM_massless(S38_x0, S38_y0, S38_vx0 / c_to_kms, S38_vy0 / c_to_kms, M, β, γ)
        S55_sol = solve_EoM_massless(S55_x0, S55_y0, S55_vx0 / c_to_kms, S55_vy0 / c_to_kms, M, β, γ)

        𝐭_em_S2 = Rømer_delay(S2_sol, S2_Ω, S2_i, 𝐭_input)
        𝐭_em_S38 = Rømer_delay(S38_sol, S38_Ω, S38_i, 𝐭_input)
        𝐭_em_S55 = Rømer_delay(S55_sol, S55_Ω, S55_i, 𝐭_input)

        S2_𝐱, S2_𝐲, S2_𝐯𝐱, S2_𝐯𝐲 = [hcat(S2_sol.(𝐭_em_S2)...)[i,:] for i in 1:4]
        S38_𝐱, S38_𝐲, S38_𝐯𝐱, S38_𝐯𝐲 = [hcat(S38_sol.(𝐭_em_S38)...)[i,:] for i in 1:4]
        S55_𝐱, S55_𝐲, S55_𝐯𝐱, S55_𝐯𝐲 = [hcat(S55_sol.(𝐭_em_S55)...)[i,:] for i in 1:4]

        S2_𝐫 = turnorbit.(S2_𝐱, S2_𝐲, Ref(0.0), Ref(S2_Ω), Ref(S2_i))
        S2_𝐯 = turnorbit.(S2_𝐯𝐱, S2_𝐯𝐲, Ref(0.0), Ref(S2_Ω), Ref(S2_i))
        S38_𝐫 = turnorbit.(S38_𝐱, S38_𝐲, Ref(0.0), Ref(S38_Ω), Ref(S38_i))
        S38_𝐯 = turnorbit.(S38_𝐯𝐱, S38_𝐯𝐲, Ref(0.0), Ref(S38_Ω), Ref(S38_i))
        S55_𝐫 = turnorbit.(S55_𝐱, S55_𝐲, Ref(0.0), Ref(S55_Ω), Ref(S55_i))
        S55_𝐯 = turnorbit.(S55_𝐯𝐱, S55_𝐯𝐲, Ref(0.0), Ref(S55_Ω), Ref(S55_i))

        Dist = R0 / r̂_to_kpc

        S2_RA = ((hcat(S2_𝐫...)[1,:] .+ vx_SgrA .* (𝐭_em_S2 .- t_beg)) ./ Dist) .* rad_to_arcsec
        S2_Dec = -((hcat(S2_𝐫...)[2,:] .+ vy_SgrA .* (𝐭_em_S2 .- t_beg)) ./ Dist) .* rad_to_arcsec
        S2_RV = (-hcat(S2_𝐯...)[3,:] .- M ./ .√(S2_𝐱.^2 .+ S2_𝐲.^2) .+ (S2_𝐯𝐱.^2 .+ S2_𝐯𝐲.^2) ./ 2.0) .* c_to_kms

        S38_RA = ((hcat(S38_𝐫...)[1,:] .+ vx_SgrA .* (𝐭_em_S38 .- t_beg)) ./ Dist) .* rad_to_arcsec
        S38_Dec = -((hcat(S38_𝐫...)[2,:] .+ vy_SgrA .* (𝐭_em_S38 .- t_beg)) ./ Dist) .* rad_to_arcsec
        S38_RV = (-hcat(S38_𝐯...)[3,:] .- M ./ .√(S38_𝐱.^2 .+ S38_𝐲.^2) .+ (S38_𝐯𝐱.^2 .+ S38_𝐯𝐲.^2) ./ 2.0) .* c_to_kms

        S55_RA = ((hcat(S55_𝐫...)[1,:] .+ vx_SgrA .* (𝐭_em_S55 .- t_beg)) ./ Dist) .* rad_to_arcsec
        S55_Dec = -((hcat(S55_𝐫...)[2,:] .+ vy_SgrA .* (𝐭_em_S55 .- t_beg)) ./ Dist) .* rad_to_arcsec
        S55_RV = (-hcat(S55_𝐯...)[3,:] .- M ./ .√(S55_𝐱.^2 .+ S55_𝐲.^2) .+ (S55_𝐯𝐱.^2 .+ S55_𝐯𝐲.^2) ./ 2.0) .* c_to_kms

        return [[S2_RA[i], S2_Dec[i], S2_RV[i], S38_RA[i], S38_Dec[i], S38_RV[i], S55_RA[i], S55_Dec[i], S55_RV[i]] for i in 1:length(S2_RA)]

    end

end