using Distributions
using EMpht
using LinearAlgebra: norm
using Random
using Test

# Fit a simple exponential distribution using
# both methods (uniformization and the ODE solver)
trueRate = 10
trueDist = Exponential(1/trueRate)

Random.seed!(1)
sObs = EMpht.Sample(obs=rand(trueDist, 1_000))
ph1unif = empht(sObs, p=1, method=:unif)
ph1ode = empht(sObs, p=1, method=:ode)

@test -ph1unif.T[1,1] ≈ trueRate atol=1
@test -ph1ode.T[1,1] ≈ trueRate atol=1

# Fit binned samples from a gamma distribution.
trueDist = Gamma(20, 2/10)
xGrid = range(0, 8, length=1_000)
truePDFs = pdf.(trueDist, xGrid)

# Random.seed!(1)
# data = rand(trueDist, 1_000)
# int, intweight = bin_observations(data, 15)
int = [1.5 2.0; 2.0 2.5; 2.5 3.0; 3.0 3.5; 3.5 4.0; 4.0 4.5; 4.5 5.0; 5.0 5.5;
        5.5 6.0; 6.0 6.5; 6.5 7.0; 7.0 7.5]
intweight = [4.0, 34.0, 107.0, 170.0, 202.0, 222.0, 140.0, 77.0, 24.0, 14.0,
        4.0, 2.0]
sInt = EMpht.Sample(int=int, intweight=intweight)

phCF1 = empht(sInt, p=100, ph_structure="CanonicalForm1")
fitPDFs = pdf.(phCF1, xGrid)

@test norm(truePDFs .- fitPDFs, Inf) <= 1e-1

# Test that reading settings from a file works.
ph100 = empht("Coxian100.json")

# Test that some of the internal functions work as expected.
π=[0.005632862236468221, 0.00898302181442919, 0.007418596691981033,
        0.020026201266915976, 0.04062880012941406, 0.05812175206666447,
        0.04260242040415281, 0.1654039105298598, 0.3888838419368729,
        0.2622985929232415]
T = [-1.7486539983379394 1.7486539983379394 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 -1.9875507782505892 1.9875507782505892 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 -1.9962263354076826 1.9962263354076826 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 -2.064044688471394 2.064044688471394 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 -2.1489226352793827 2.1489226352793827 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 -2.2809066095317565 2.2809066095317565 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 -2.5379722045000195 2.5379722045000195 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 -2.763746562995887 2.763746562995887 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -3.082614321671955 3.082614321671955;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -3.910958867504505]
t=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.910958867504505]
ph = EMpht.PhaseType(π, T, t)

u = zeros(ph.p * ph.p)
du = similar(u)
EMpht.ode_observations!(du, u, ph, 10.0)

duCorrect = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 5.60666e-10, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.8095e-9, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.0, 2.46331e-8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 1.15854e-7, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        4.26538e-7, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.23682e-6,
        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.71623e-6, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 5.29954e-6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        0.0, 0.0, 0.0, 8.87025e-6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
        1.093e-5]
@test norm(du .- duCorrect, Inf) <= 1e-10

EMpht.ode_censored!(du, u, ph, 10.0)

duCorrect = [1.43358e-10, 1.43358e-10, 1.43358e-10, 1.43358e-10, 1.43358e-10,
        1.43358e-10, 1.43358e-10, 1.43358e-10, 1.43358e-10, 1.43358e-10,
        9.74057e-10, 9.74057e-10, 9.74057e-10, 9.74057e-10, 9.74057e-10,
        9.74057e-10, 9.74057e-10, 9.74057e-10, 9.74057e-10, 9.74057e-10,
        6.29847e-9, 6.29847e-9, 6.29847e-9, 6.29847e-9, 6.29847e-9, 6.29847e-9,
        6.29847e-9, 6.29847e-9, 6.29847e-9, 6.29847e-9, 2.96229e-8, 2.96229e-8,
        2.96229e-8, 2.96229e-8, 2.96229e-8, 2.96229e-8, 2.96229e-8, 2.96229e-8,
        2.96229e-8, 2.96229e-8, 1.09062e-7, 1.09062e-7, 1.09062e-7, 1.09062e-7,
        1.09062e-7, 1.09062e-7, 1.09062e-7, 1.09062e-7, 1.09062e-7, 1.09062e-7,
        3.16245e-7, 3.16245e-7, 3.16245e-7, 3.16245e-7, 3.16245e-7, 3.16245e-7,
        3.16245e-7, 3.16245e-7, 3.16245e-7, 3.16245e-7, 6.94519e-7, 6.94519e-7,
        6.94519e-7, 6.94519e-7, 6.94519e-7, 6.94519e-7, 6.94519e-7, 6.94519e-7,
        6.94519e-7, 6.94519e-7, 1.35505e-6, 1.35505e-6, 1.35505e-6, 1.35505e-6,
        1.35505e-6, 1.35505e-6, 1.35505e-6, 1.35505e-6, 1.35505e-6, 1.35505e-6,
        2.26805e-6, 2.26805e-6, 2.26805e-6, 2.26805e-6, 2.26805e-6, 2.26805e-6,
        2.26805e-6, 2.26805e-6, 2.26805e-6, 2.26805e-6, 2.7947e-6, 2.7947e-6,
        2.7947e-6, 2.7947e-6, 2.7947e-6, 2.7947e-6, 2.7947e-6, 2.7947e-6,
        2.7947e-6, 2.7947e-6]

@test norm(du .- duCorrect, Inf) <= 1e-10

cIntCorrect = [1.95424e-7, 1.37498e-7, 8.25671e-8, 3.9341e-8, 1.42689e-8,
        3.6283e-9, 5.76897e-10, 5.8746e-11, 2.60493e-12, 1.13063e-14,
        1.09161e-6, 7.68042e-7, 4.61207e-7, 2.19753e-7, 7.97037e-8, 2.02671e-8,
        3.22246e-9, 3.28146e-10, 1.45507e-11, 6.3155e-14, 4.93283e-6,
        3.47068e-6, 2.08413e-6, 9.93034e-7, 3.6017e-7, 9.15845e-8, 1.45619e-8,
        1.48285e-9, 6.57528e-11, 2.85389e-13, 1.56191e-5, 1.09894e-5,
        6.59911e-6, 3.1443e-6, 1.14043e-6, 2.89989e-7, 4.6108e-8, 4.69522e-9,
        2.08197e-10, 9.03644e-13, 3.88217e-5, 2.73145e-5, 1.64023e-5,
        7.81523e-6, 2.83456e-6, 7.20775e-7, 1.14603e-7, 1.16701e-8, 5.17479e-10,
        2.24603e-12, 7.8674e-5, 5.53541e-5, 3.324e-5, 1.58379e-5, 5.74438e-6,
        1.46069e-6, 2.32248e-7, 2.365e-8, 1.0487e-9, 4.55169e-12, 0.000128921,
        9.07072e-5, 5.44694e-5, 2.59532e-5, 9.41315e-6, 2.39359e-6, 3.80578e-7,
        3.87547e-8, 1.71847e-9, 7.45873e-12, 0.000195269, 0.000137389,
        8.25018e-5, 3.93099e-5, 1.42576e-5, 3.62543e-6, 5.7644e-7, 5.86995e-8,
        2.60287e-9, 1.12973e-11, 0.000264903, 0.000186383, 0.000111922,
        5.33279e-5, 1.93419e-5, 4.91827e-6, 7.82001e-7, 7.9632e-8, 3.53106e-9,
        1.5326e-11, 0.000282308, 0.000198628, 0.000119276, 5.68317e-5,
        2.06127e-5, 5.24141e-6, 8.3338e-7, 8.4864e-8, 3.76306e-9, 1.63329e-11]

@test norm(EMpht.c_integrand(5.0, ph, 10.0) .- cIntCorrect, Inf) <= 1e-8

dIntCorrect = [2.49838e-7, 1.38082e-7, 6.89019e-8, 2.75403e-8, 8.48015e-9,
        1.84014e-9, 2.4941e-10, 2.21039e-11, 8.4793e-13, 2.89092e-15,
        1.39556e-6, 7.71302e-7, 3.84875e-7, 1.53836e-7, 4.73688e-8, 1.02787e-8,
        1.39317e-9, 1.23469e-10, 4.7364e-12, 1.61482e-14, 6.30634e-6,
        3.48541e-6, 1.7392e-6, 6.95164e-7, 2.14053e-7, 4.64482e-8, 6.29553e-9,
        5.57939e-10, 2.14032e-11, 7.29717e-14, 1.99681e-5, 1.1036e-5,
        5.50693e-6, 2.20114e-6, 6.77769e-7, 1.47071e-7, 1.99339e-8, 1.76663e-9,
        6.77701e-11, 2.31054e-13, 4.96313e-5, 2.74304e-5, 1.36876e-5,
        5.47098e-6, 1.68461e-6, 3.6555e-7, 4.95462e-8, 4.39101e-9, 1.68444e-10,
        5.74291e-13, 0.00010058, 5.5589e-5, 2.77386e-5, 1.10872e-5, 3.41395e-6,
        7.40805e-7, 1.00408e-7, 8.8986e-9, 3.41361e-10, 1.16383e-12,
        0.000164818, 9.10922e-5, 4.54545e-5, 1.81683e-5, 5.59434e-6, 1.21394e-6,
        1.64535e-7, 1.45819e-8, 5.59378e-10, 1.90714e-12, 0.000249641,
        0.000137972, 6.88474e-5, 2.75185e-5, 8.47344e-6, 1.83868e-6, 2.49213e-7,
        2.20864e-8, 8.47259e-10, 2.88863e-12, 0.000338663, 0.000187174,
        9.33986e-5, 3.73317e-5, 1.14951e-5, 2.49436e-6, 3.38083e-7, 2.99624e-8,
        1.14939e-9, 3.91873e-12, 0.000360914, 0.000199471, 9.95351e-5,
        3.97845e-5, 1.22503e-5, 2.65825e-6, 3.60296e-7, 3.1931e-8, 1.22491e-9,
        4.1762e-12]

@test norm(EMpht.d_integrand(5.0, ph, 10.0) .- dIntCorrect, Inf) <= 1e-8


p = ph.p; Bs = zeros(p); Zs = zeros(p); Ns = zeros(p, p+1)

EMpht.e_step_observed_ode!(sObs, ph, Bs, Zs, Ns)

BsCorrect = [9.02369e-6, 9.24901e-5, 0.00041208, 0.00568032, 0.0558114,
        0.377397, 1.30242, 23.3691, 251.208, 723.681]
ZsCorrect = [7.43742e-7, 8.80107e-6, 4.71973e-5, 0.000600675, 0.00613343,
        0.04273, 0.160008, 2.10905, 21.4158, 71.7403]
NsCorrect = [0.0 9.02399e-6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.000101516 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.000513599 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0061939 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.062005 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.439402 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.74183 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 25.1109 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 276.319 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1000.0]

@test norm(Bs .- BsCorrect, Inf) <= 1e-3
@test norm(Zs .- ZsCorrect, Inf) <= 1e-3
@test norm(Ns .- NsCorrect, Inf) <= 1e-3


p = ph.p; Bs = zeros(p); Zs = zeros(p); Ns = zeros(p, p+1)

EMpht.e_step_censored_ode!(sInt, ph, Bs, Zs, Ns)

BsCorrect = [130.738, 174.776, 108.273, 188.54, 206.27, 126.544, 29.5952,
        26.7001, 8.34855, 0.214886]
ZsCorrect = [79.0644, 173.84, 245.305, 376.65, 533.48, 617.446, 565.593,
        538.451, 475.918, 341.487]
NsCorrect = [0.0 130.737 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 305.513 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 413.787 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 602.329 0.0 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 808.599 0.0 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 935.148 0.0 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 964.712 0.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 991.49 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 999.738 0.0;
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1000.0]

@test norm(Bs .- BsCorrect, Inf) <= 1e-3
@test norm(Zs .- ZsCorrect, Inf) <= 1e-3
@test norm(Ns .- NsCorrect, Inf) <= 1e-3
