# using GSL.sf_legendre_sphPlm

# function sphharm(l::Integer, m::Integer, theta::Real, phi::Real)
# 	@assert l >= 0 && abs(m) <= l
# 	if m >= 0
# 		return exp(m*phi*im)*sf_legendre_sphPlm(l,m,cos(theta))
# 	else
#     	return (-1.)^m*exp(m*phi*im)*sf_legendre_sphPlm(l,-m,cos(theta))
#     end
# end

doc"""
`sphharm(l,m,θ,ϕ)`

computes the spherical harmonic function (standard definition used in quantum mechanics)

$ Y_l^m(\theta,\phi) = (-1)^m \sqrt{{2l+1 \over 4\pi} {(l-m)!\over (l+m)!}} P_l^m(cos \theta) e^{im\phi} $

which are normalized such that

$ \int Y_l^m(\theta,\phi)^* Y_{l'}^{m'}(\theta,\phi) d\Omega = \delta_{ll'} \delta_{mm'} $
"""
# sphharm(l::Int, m::Int, θ::Real, ϕ::Real) = sphharmtab[l,m](θ,ϕ)::Complex{Float64}
sphharm(l::Int, m::Int, θ::Real, ϕ::Real) = sphharmtab[i_lm(l,m)](θ,ϕ)::Complex{Float64}

doc"""
`realsphharm(l,m,θ,ϕ)`

computes the real spherical harmonic function

$ Y_{lm} = {i \over \sqrt{2}} \left(Y_l^m - (-1)^m Y_l^{-m}\right), \quad m<0 $

$ Y_{lm} = Y_l^m, \quad m=0 $

$ Y_{lm} = {1 \over \sqrt{2}} \left(Y_l^{-m} + (-1)^m Y_l^m\right), \quad m>0 $

"""
# realsphharm(l::Int, m::Int, θ::Real, ϕ::Real) = realsphharmtab[l,m](θ,ϕ)::Float64
realsphharm(l::Int, m::Int, θ::Real, ϕ::Real) = realsphharmtab[i_lm(l,m)](θ,ϕ)::Float64

sphharm00(θ::Float64, ϕ::Float64) = complex(1/(2*sqrt(pi)))
sphharm1m1(θ::Float64, ϕ::Float64) = ((1/2)*sqrt(3/(2*pi))*sin(θ))/exp(ϕ*im)
sphharm10(θ::Float64, ϕ::Float64) = complex((1/2)*sqrt(3/pi)*cos(θ))
sphharm11(θ::Float64, ϕ::Float64) = (-(1/2))*exp(ϕ*im)*sqrt(3/(2*pi))*sin(θ)
sphharm2m2(θ::Float64, ϕ::Float64) = ((1/4)*sqrt(15/(2*pi))*sin(θ)^2)/exp(2*ϕ*im)
sphharm2m1(θ::Float64, ϕ::Float64) = ((1/2)*sqrt(15/(2*pi))*cos(θ)*sin(θ))/exp(ϕ*im)
sphharm20(θ::Float64, ϕ::Float64) = complex((1/8)*sqrt(5/pi)*(1 + 3*cos(2*θ)))
sphharm21(θ::Float64, ϕ::Float64) = (-(1/2))*exp(ϕ*im)*sqrt(15/(2*pi))*cos(θ)*sin(θ)
sphharm22(θ::Float64, ϕ::Float64) = (1/4)*exp(2*ϕ*im)*sqrt(15/(2*pi))*sin(θ)^2
sphharm3m3(θ::Float64, ϕ::Float64) = ((1/8)*sqrt(35/pi)*sin(θ)^3)/exp(3*ϕ*im)
sphharm3m2(θ::Float64, ϕ::Float64) = ((1/4)*sqrt(105/(2*pi))*cos(θ)*sin(θ)^2)/exp(2*ϕ*im)
sphharm3m1(θ::Float64, ϕ::Float64) = ((1/16)*sqrt(21/pi)*(3 + 5*cos(2*θ))*sin(θ))/exp(ϕ*im)
sphharm30(θ::Float64, ϕ::Float64) = complex((1/16)*sqrt(7/pi)*(3*cos(θ) + 5*cos(3*θ)))
sphharm31(θ::Float64, ϕ::Float64) = (-(1/16))*exp(ϕ*im)*sqrt(21/pi)*(3 + 5*cos(2*θ))*sin(θ)
sphharm32(θ::Float64, ϕ::Float64) = (1/4)*exp(2*ϕ*im)*sqrt(105/(2*pi))*cos(θ)*sin(θ)^2
sphharm33(θ::Float64, ϕ::Float64) = (-(1/8))*exp(3*ϕ*im)*sqrt(35/pi)*sin(θ)^3
sphharm4m4(θ::Float64, ϕ::Float64) = ((3/16)*sqrt(35/(2*pi))*sin(θ)^4)/exp(4*ϕ*im)
sphharm4m3(θ::Float64, ϕ::Float64) = ((3/8)*sqrt(35/pi)*cos(θ)*sin(θ)^3)/exp(3*ϕ*im)
sphharm4m2(θ::Float64, ϕ::Float64) = ((3/16)*sqrt(5/(2*pi))*(5 + 7*cos(2*θ))*sin(θ)^2)/exp(2*ϕ*im)
sphharm4m1(θ::Float64, ϕ::Float64) = ((3/32)*sqrt(5/pi)*(1 + 7*cos(2*θ))*sin(2*θ))/exp(ϕ*im)
sphharm40(θ::Float64, ϕ::Float64) = complex((3*(9 + 20*cos(2*θ) + 35*cos(4*θ)))/(128*sqrt(pi)))
sphharm41(θ::Float64, ϕ::Float64) = (-(3/32))*exp(ϕ*im)*sqrt(5/pi)*(1 + 7*cos(2*θ))*sin(2*θ)
sphharm42(θ::Float64, ϕ::Float64) = (3/16)*exp(2*ϕ*im)*sqrt(5/(2*pi))*(5 + 7*cos(2*θ))*sin(θ)^2
sphharm43(θ::Float64, ϕ::Float64) = (-(3/8))*exp(3*ϕ*im)*sqrt(35/pi)*cos(θ)*sin(θ)^3
sphharm44(θ::Float64, ϕ::Float64) = (3/16)*exp(4*ϕ*im)*sqrt(35/(2*pi))*sin(θ)^4
sphharm5m5(θ::Float64, ϕ::Float64) = ((3/32)*sqrt(77/pi)*sin(θ)^5)/exp(5*ϕ*im)
sphharm5m4(θ::Float64, ϕ::Float64) = ((3/16)*sqrt(385/(2*pi))*cos(θ)*sin(θ)^4)/exp(4*ϕ*im)
sphharm5m3(θ::Float64, ϕ::Float64) = ((1/64)*sqrt(385/pi)*(7 + 9*cos(2*θ))*sin(θ)^3)/exp(3*ϕ*im)
sphharm5m2(θ::Float64, ϕ::Float64) = ((1/16)*sqrt(1155/(2*pi))*cos(θ)*(1 + 3*cos(2*θ))*sin(θ)^2)/exp(2*ϕ*im)
sphharm5m1(θ::Float64, ϕ::Float64) = ((1/16)*sqrt(165/(2*pi))*(1 - 14*cos(θ)^2 +21*cos(θ)^4)*sin(θ))/exp(ϕ*im)
sphharm50(θ::Float64, ϕ::Float64) = complex((1/256)*sqrt(11/pi)*(30*cos(θ) + 35*cos(3*θ) +63*cos(5*θ)))
sphharm51(θ::Float64, ϕ::Float64) = (-(1/16))*exp(ϕ*im)*sqrt(165/(2*pi))*(1 - 14*cos(θ)^2 + 21*cos(θ)^4)*sin(θ)
sphharm52(θ::Float64, ϕ::Float64) = (1/16)*exp(2*ϕ*im)*sqrt(1155/(2*pi))*cos(θ)*(1 + 3*cos(2*θ))*sin(θ)^2
sphharm53(θ::Float64, ϕ::Float64) = (-(1/64))*exp(3*ϕ*im)*sqrt(385/pi)*(7 + 9*cos(2*θ))*sin(θ)^3
sphharm54(θ::Float64, ϕ::Float64) = (3/16)*exp(4*ϕ*im)*sqrt(385/(2*pi))*cos(θ)*sin(θ)^4
sphharm55(θ::Float64, ϕ::Float64) = (-(3/32))*exp(5*ϕ*im)*sqrt(77/pi)*sin(θ)^5
sphharm6m6(θ::Float64, ϕ::Float64) = ((1/64)*sqrt(3003/pi)*sin(θ)^6)/exp(6*ϕ*im)
sphharm6m5(θ::Float64, ϕ::Float64) = ((3/32)*sqrt(1001/pi)*cos(θ)*sin(θ)^5)/exp(5*ϕ*im)
sphharm6m4(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(91/(2*pi))*(9 + 11*cos(2*θ))*sin(θ)^4)/exp(4*ϕ*im)
sphharm6m3(θ::Float64, ϕ::Float64) = ((1/64)*sqrt(1365/pi)*cos(θ)*(5 + 11*cos(2*θ))*sin(θ)^3)/exp(3*ϕ*im)
sphharm6m2(θ::Float64, ϕ::Float64) = ((1/64)*sqrt(1365/pi)*(1 - 18*cos(θ)^2 +33*cos(θ)^4)*sin(θ)^2)/exp(2*ϕ*im)
sphharm6m1(θ::Float64, ϕ::Float64) = ((1/16)*sqrt(273/(2*pi))*cos(θ)*(5 - 30*cos(θ)^2 +33*cos(θ)^4)*sin(θ))/exp(ϕ*im)
sphharm60(θ::Float64, ϕ::Float64) = complex((1/32)*sqrt(13/pi)*(-5 + 105*cos(θ)^2 -315*cos(θ)^4 + 231*cos(θ)^6))
sphharm61(θ::Float64, ϕ::Float64) = (-(1/16))*exp(ϕ*im)*sqrt(273/(2*pi))*cos(θ)*(5 - 30*cos(θ)^2 + 33*cos(θ)^4)*sin(θ)
sphharm62(θ::Float64, ϕ::Float64) = (1/64)*exp(2*ϕ*im)*sqrt(1365/pi)*(1 - 18*cos(θ)^2 +33*cos(θ)^4)*sin(θ)^2
sphharm63(θ::Float64, ϕ::Float64) = (-(1/64))*exp(3*ϕ*im)*sqrt(1365/pi)*cos(θ)*(5 + 11*cos(2*θ))*sin(θ)^3
sphharm64(θ::Float64, ϕ::Float64) = (3/64)*exp(4*ϕ*im)*sqrt(91/(2*pi))*(9 + 11*cos(2*θ))*sin(θ)^4
sphharm65(θ::Float64, ϕ::Float64) = (-(3/32))*exp(5*ϕ*im)*sqrt(1001/pi)*cos(θ)*sin(θ)^5
sphharm66(θ::Float64, ϕ::Float64) = (1/64)*exp(6*ϕ*im)*sqrt(3003/pi)*sin(θ)^6
sphharm7m7(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(715/(2*pi))*sin(θ)^7)/exp(7*ϕ*im)
sphharm7m6(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(5005/pi)*cos(θ)*sin(θ)^6)/exp(6*ϕ*im)
sphharm7m5(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(385/(2*pi))*(-1 + 13*cos(θ)^2)*sin(θ)^5)/exp(5*ϕ*im)
sphharm7m4(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(385/(2*pi))*cos(θ)*(7 + 13*cos(2*θ))*sin(θ)^4)/exp(4*ϕ*im)
sphharm7m3(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(35/(2*pi))*(3 - 66*cos(θ)^2 +143*cos(θ)^4)*sin(θ)^3)/exp(3*ϕ*im)
sphharm7m2(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(35/pi)*cos(θ)*(15 - 110*cos(θ)^2 +143*cos(θ)^4)*sin(θ)^2)/exp(2*ϕ*im)
sphharm7m1(θ::Float64, ϕ::Float64) = ((1/64)*sqrt(105/(2*pi))*(-5 + 135*cos(θ)^2 -495*cos(θ)^4 + 429*cos(θ)^6)*sin(θ))/exp(ϕ*im)
sphharm70(θ::Float64, ϕ::Float64) = complex((1/32)*sqrt(15/pi)*cos(θ)*(-35 + 315*cos(θ)^2 -693*cos(θ)^4 + 429*cos(θ)^6))
sphharm71(θ::Float64, ϕ::Float64) = (-(1/64))*exp(ϕ*im)*sqrt(105/(2*pi))*(-5 + 135*cos(θ)^2 - 495*cos(θ)^4 + 429*cos(θ)^6)*sin(θ)
sphharm72(θ::Float64, ϕ::Float64) = (3/64)*exp(2*ϕ*im)*sqrt(35/pi)*cos(θ)*(15 - 110*cos(θ)^2 + 143*cos(θ)^4)*sin(θ)^2
sphharm73(θ::Float64, ϕ::Float64) = (-(3/64))*exp(3*ϕ*im)*sqrt(35/(2*pi))*(3 - 66*cos(θ)^2 + 143*cos(θ)^4)*sin(θ)^3
sphharm74(θ::Float64, ϕ::Float64) = (3/64)*exp(4*ϕ*im)*sqrt(385/(2*pi))*cos(θ)*(7 + 13*cos(2*θ))*sin(θ)^4
sphharm75(θ::Float64, ϕ::Float64) = (-(3/64))*exp(5*ϕ*im)*sqrt(385/(2*pi))*(-1 + 13*cos(θ)^2)*sin(θ)^5
sphharm76(θ::Float64, ϕ::Float64) = (3/64)*exp(6*ϕ*im)*sqrt(5005/pi)*cos(θ)*sin(θ)^6
sphharm77(θ::Float64, ϕ::Float64) = (-(3/64))*exp(7*ϕ*im)*sqrt(715/(2*pi))*sin(θ)^7
sphharm8m8(θ::Float64, ϕ::Float64) = ((3/256)*sqrt(12155/(2*pi))*sin(θ)^8)/exp(8*ϕ*im)
sphharm8m7(θ::Float64, ϕ::Float64) = ((3/64)*sqrt(12155/(2*pi))*cos(θ)*sin(θ)^7)/exp(7*ϕ*im)
sphharm8m6(θ::Float64, ϕ::Float64) = ((1/256)*sqrt(7293/pi)*(13 + 15*cos(2*θ))*sin(θ)^6)/exp(6*ϕ*im)
sphharm8m5(θ::Float64, ϕ::Float64) = ((3/128)*sqrt(17017/(2*pi))*cos(θ)*(3 + 5*cos(2*θ))*sin(θ)^5)/exp(5*ϕ*im)
sphharm8m4(θ::Float64, ϕ::Float64) = ((3/128)*sqrt(1309/(2*pi))*(1 - 26*cos(θ)^2 +65*cos(θ)^4)*sin(θ)^4)/exp(4*ϕ*im)
sphharm8m3(θ::Float64, ϕ::Float64) = ((1/64)*sqrt(19635/(2*pi))*cos(θ)*(3 - 26*cos(θ)^2 + 39*cos(θ)^4)*sin(θ)^3)/exp(3*ϕ*im)
sphharm8m2(θ::Float64, ϕ::Float64) = ((3/128)*sqrt(595/pi)*(-1 + 33*cos(θ)^2 - 143*cos(θ)^4 + 143*cos(θ)^6)*sin(θ)^2)/exp(2*ϕ*im)
sphharm8m1(θ::Float64, ϕ::Float64) = (1/4096)*((3*sqrt(17/(2*pi))*(178 + 869*cos(2*θ) +286*cos(4*θ) + 715*cos(6*θ))*sin(2*θ))/exp(ϕ*im))
sphharm80(θ::Float64, ϕ::Float64) = complex((1/256)*sqrt(17/pi)*(35 - 1260*cos(θ)^2 + 6930*cos(θ)^4 -12012*cos(θ)^6 + 6435*cos(θ)^8))
sphharm81(θ::Float64, ϕ::Float64) = -((1/4096)*(3*exp(ϕ*im)*sqrt(17/(2*pi))*(178 + 869*cos(2*θ) + 286*cos(4*θ) +715*cos(6*θ))*sin(2*θ)))
sphharm82(θ::Float64, ϕ::Float64) = (3/128)*exp(2*ϕ*im)*sqrt(595/pi)*(-1 + 33*cos(θ)^2 -143*cos(θ)^4 + 143*cos(θ)^6)*sin(θ)^2
sphharm83(θ::Float64, ϕ::Float64) = (-(1/64))*exp(3*ϕ*im)*sqrt(19635/(2*pi))*cos(θ)*(3 - 26*cos(θ)^2 + 39*cos(θ)^4)*sin(θ)^3
sphharm84(θ::Float64, ϕ::Float64) = (3/128)*exp(4*ϕ*im)*sqrt(1309/(2*pi))*(1 - 26*cos(θ)^2 + 65*cos(θ)^4)*sin(θ)^4
sphharm85(θ::Float64, ϕ::Float64) = (-(3/128))*exp(5*ϕ*im)*sqrt(17017/(2*pi))*cos(θ)*(3 + 5*cos(2*θ))*sin(θ)^5
sphharm86(θ::Float64, ϕ::Float64) = (1/256)*exp(6*ϕ*im)*sqrt(7293/pi)*(13 + 15*cos(2*θ))*sin(θ)^6
sphharm87(θ::Float64, ϕ::Float64) = (-(3/64))*exp(7*ϕ*im)*sqrt(12155/(2*pi))*cos(θ)*sin(θ)^7
sphharm88(θ::Float64, ϕ::Float64) = (3/256)*exp(8*ϕ*im)*sqrt(12155/(2*pi))*sin(θ)^8

realsphharm00(θ::Float64, ϕ::Float64) = 1/(2*sqrt(pi))
realsphharm1m1(θ::Float64, ϕ::Float64) = (1/2)*sqrt(3/pi)*sin(ϕ)*sin(θ)
realsphharm10(θ::Float64, ϕ::Float64) = (1/2)*sqrt(3/pi)*cos(θ)
realsphharm11(θ::Float64, ϕ::Float64) = (1/2)*sqrt(3/pi)*cos(ϕ)*sin(θ)
realsphharm2m2(θ::Float64, ϕ::Float64) = (-(1/4))*sqrt(15/pi)*sin(2*ϕ)*sin(θ)^2
realsphharm2m1(θ::Float64, ϕ::Float64) = (1/2)*sqrt(15/pi)*cos(θ)*sin(ϕ)*sin(θ)
realsphharm20(θ::Float64, ϕ::Float64) = (1/8)*sqrt(5/pi)*(1 + 3*cos(2*θ))
realsphharm21(θ::Float64, ϕ::Float64) = (1/2)*sqrt(15/pi)*cos(ϕ)*cos(θ)*sin(θ)
realsphharm22(θ::Float64, ϕ::Float64) = (1/4)*sqrt(15/pi)*cos(2*ϕ)*sin(θ)^2
realsphharm3m3(θ::Float64, ϕ::Float64) = (1/4)*sqrt(35/(2*pi))*sin(3*ϕ)*sin(θ)^3
realsphharm3m2(θ::Float64, ϕ::Float64) = (-(1/4))*sqrt(105/pi)*cos(θ)*sin(2*ϕ)*sin(θ)^2
realsphharm3m1(θ::Float64, ϕ::Float64) = (1/16)*sqrt(21/(2*pi))*sin(ϕ)*(sin(θ) + 5*sin(3*θ))
realsphharm30(θ::Float64, ϕ::Float64) = (1/8)*sqrt(7/pi)*cos(θ)*(-1 + 5*cos(2*θ))
realsphharm31(θ::Float64, ϕ::Float64) = (1/16)*sqrt(21/(2*pi))*cos(ϕ)*(sin(θ) + 5*sin(3*θ))
realsphharm32(θ::Float64, ϕ::Float64) = (1/4)*sqrt(105/pi)*cos(2*ϕ)*cos(θ)*sin(θ)^2
realsphharm33(θ::Float64, ϕ::Float64) = (1/4)*sqrt(35/(2*pi))*cos(3*ϕ)*sin(θ)^3
realsphharm4m4(θ::Float64, ϕ::Float64) = (-(3/16))*sqrt(35/pi)*sin(4*ϕ)*sin(θ)^4
realsphharm4m3(θ::Float64, ϕ::Float64) = (3/4)*sqrt(35/(2*pi))*cos(θ)*sin(3*ϕ)*sin(θ)^3
realsphharm4m2(θ::Float64, ϕ::Float64) = (3/8)*sqrt(5/pi)*(1 - 7*cos(θ)^2)*sin(2*ϕ)*sin(θ)^2
realsphharm4m1(θ::Float64, ϕ::Float64) = (3/32)*sqrt(5/(2*pi))*sin(ϕ)*(2*sin(2*θ) + 7*sin(4*θ))
realsphharm40(θ::Float64, ϕ::Float64) = (3*(9 + 20*cos(2*θ) + 35*cos(4*θ)))/(128*sqrt(pi))
realsphharm41(θ::Float64, ϕ::Float64) = (3/32)*sqrt(5/(2*pi))*cos(ϕ)*(2*sin(2*θ) + 7*sin(4*θ))
realsphharm42(θ::Float64, ϕ::Float64) = (3/16)*sqrt(5/pi)*cos(2*ϕ)*(5 + 7*cos(2*θ))*sin(θ)^2
realsphharm43(θ::Float64, ϕ::Float64) = (3/4)*sqrt(35/(2*pi))*cos(3*ϕ)*cos(θ)*sin(θ)^3
realsphharm44(θ::Float64, ϕ::Float64) = (3/16)*sqrt(35/pi)*cos(4*ϕ)*sin(θ)^4
realsphharm5m5(θ::Float64, ϕ::Float64) = (3/16)*sqrt(77/(2*pi))*sin(5*ϕ)*sin(θ)^5
realsphharm5m4(θ::Float64, ϕ::Float64) = (-(3/16))*sqrt(385/pi)*cos(θ)*sin(4*ϕ)*sin(θ)^4
realsphharm5m3(θ::Float64, ϕ::Float64) = (1/32)*sqrt(385/(2*pi))*(7 + 9*cos(2*θ))*sin(3*ϕ)*sin(θ)^3
realsphharm5m2(θ::Float64, ϕ::Float64) = (-(1/32))*sqrt(1155/pi)*(5*cos(θ) + 3*cos(3*θ))*sin(2*ϕ)*sin(θ)^2
realsphharm5m1(θ::Float64, ϕ::Float64) = (1/256)*sqrt(165/pi)*sin(ϕ)*(2*sin(θ) + 7*(sin(3*θ) + 3*sin(5*θ)))
realsphharm50(θ::Float64, ϕ::Float64) = (1/256)*sqrt(11/pi)*(30*cos(θ) + 35*cos(3*θ) + 63*cos(5*θ))
realsphharm51(θ::Float64, ϕ::Float64) = (1/256)*sqrt(165/pi)*cos(ϕ)*(2*sin(θ) + 7*(sin(3*θ) + 3*sin(5*θ)))
realsphharm52(θ::Float64, ϕ::Float64) = (1/16)*sqrt(1155/pi)*cos(2*ϕ)*cos(θ)*(1 + 3*cos(2*θ))*sin(θ)^2
realsphharm53(θ::Float64, ϕ::Float64) = (1/32)*sqrt(385/(2*pi))*cos(3*ϕ)*(7 + 9*cos(2*θ))*sin(θ)^3
realsphharm54(θ::Float64, ϕ::Float64) = (3/16)*sqrt(385/pi)*cos(4*ϕ)*cos(θ)*sin(θ)^4
realsphharm55(θ::Float64, ϕ::Float64) = (3/16)*sqrt(77/(2*pi))*cos(5*ϕ)*sin(θ)^5
realsphharm6m6(θ::Float64, ϕ::Float64) = (-(1/32))*sqrt(3003/(2*pi))*sin(6*ϕ)*sin(θ)^6
realsphharm6m5(θ::Float64, ϕ::Float64) = (3/16)*sqrt(1001/(2*pi))*cos(θ)*sin(5*ϕ)*sin(θ)^5
realsphharm6m4(θ::Float64, ϕ::Float64) = (-(3/64))*sqrt(91/pi)*(9 + 11*cos(2*θ))*sin(4*ϕ)*sin(θ)^4
realsphharm6m3(θ::Float64, ϕ::Float64) = (1/64)*sqrt(1365/(2*pi))*(21*cos(θ) + 11*cos(3*θ))*sin(3*ϕ)*sin(θ)^3
realsphharm6m2(θ::Float64, ϕ::Float64) = (-(1/256))*sqrt(1365/(2*pi))*(35 + 60*cos(2*θ) + 33*cos(4*θ))*sin(2*ϕ)*sin(θ)^2
realsphharm6m1(θ::Float64, ϕ::Float64) = (1/512)*sqrt(273/pi)*sin(ϕ)*(5*sin(2*θ) + 12*sin(4*θ) + 33*sin(6*θ))
realsphharm60(θ::Float64, ϕ::Float64) = (sqrt(13/pi)*(50 + 105*cos(2*θ) + 126*cos(4*θ) + 231*cos(6*θ)))/1024
realsphharm61(θ::Float64, ϕ::Float64) = (1/512)*sqrt(273/pi)*cos(ϕ)*(5*sin(2*θ) + 12*sin(4*θ) + 33*sin(6*θ))
realsphharm62(θ::Float64, ϕ::Float64) = (1/256)*sqrt(1365/(2*pi))*cos(2*ϕ)*(35 + 60*cos(2*θ) + 33*cos(4*θ))*sin(θ)^2
realsphharm63(θ::Float64, ϕ::Float64) = (1/64)*sqrt(1365/(2*pi))*cos(3*ϕ)*(21*cos(θ) + 11*cos(3*θ))*sin(θ)^3
realsphharm64(θ::Float64, ϕ::Float64) = (3/64)*sqrt(91/pi)*cos(4*ϕ)*(9 + 11*cos(2*θ))*sin(θ)^4
realsphharm65(θ::Float64, ϕ::Float64) = (3/16)*sqrt(1001/(2*pi))*cos(5*ϕ)*cos(θ)*sin(θ)^5
realsphharm66(θ::Float64, ϕ::Float64) = (1/32)*sqrt(3003/(2*pi))*cos(6*ϕ)*sin(θ)^6
realsphharm7m7(θ::Float64, ϕ::Float64) = (3/64)*sqrt(715/pi)*sin(7*ϕ)*sin(θ)^7
realsphharm7m6(θ::Float64, ϕ::Float64) = (-(3/32))*sqrt(5005/(2*pi))*cos(θ)*sin(6*ϕ)*sin(θ)^6
realsphharm7m5(θ::Float64, ϕ::Float64) = (3/128)*sqrt(385/pi)*(11 + 13*cos(2*θ))*sin(5*ϕ)*sin(θ)^5
realsphharm7m4(θ::Float64, ϕ::Float64) = (-(3/64))*sqrt(385/pi)*cos(θ)*(7 + 13*cos(2*θ))*sin(4*ϕ)*sin(θ)^4
realsphharm7m3(θ::Float64, ϕ::Float64) = (3/512)*sqrt(35/pi)*(189 + 308*cos(2*θ) + 143*cos(4*θ))*sin(3*ϕ)*sin(θ)^3
realsphharm7m2(θ::Float64, ϕ::Float64) = (-(3/512))*sqrt(35/(2*pi))*(350*cos(θ) + 275*cos(3*θ) + 143*cos(5*θ))*sin(2*ϕ)*sin(θ)^2
realsphharm7m1(θ::Float64, ϕ::Float64) = (sqrt(105/pi)*sin(ϕ)*(25*sin(θ) + 81*sin(3*θ) + 165*sin(5*θ) + 429*sin(7*θ)))/4096
realsphharm70(θ::Float64, ϕ::Float64) = (sqrt(15/pi)*(175*cos(θ) + 189*cos(3*θ) + 231*cos(5*θ) + 429*cos(7*θ)))/2048
realsphharm71(θ::Float64, ϕ::Float64) = (sqrt(105/pi)*cos(ϕ)*(25*sin(θ) + 81*sin(3*θ) + 165*sin(5*θ) + 429*sin(7*θ)))/4096
realsphharm72(θ::Float64, ϕ::Float64) = (3/512)*sqrt(35/(2*pi))*cos(2*ϕ)*(350*cos(θ) + 275*cos(3*θ) + 143*cos(5*θ))*sin(θ)^2
realsphharm73(θ::Float64, ϕ::Float64) = (3/512)*sqrt(35/pi)*cos(3*ϕ)*(189 + 308*cos(2*θ) + 143*cos(4*θ))*sin(θ)^3
realsphharm74(θ::Float64, ϕ::Float64) = (3/128)*sqrt(385/pi)*cos(4*ϕ)*(27*cos(θ) + 13*cos(3*θ))*sin(θ)^4
realsphharm75(θ::Float64, ϕ::Float64) = (3/128)*sqrt(385/pi)*cos(5*ϕ)*(11 + 13*cos(2*θ))*sin(θ)^5
realsphharm76(θ::Float64, ϕ::Float64) = (3/32)*sqrt(5005/(2*pi))*cos(6*ϕ)*cos(θ)*sin(θ)^6
realsphharm77(θ::Float64, ϕ::Float64) = (3/64)*sqrt(715/pi)*cos(7*ϕ)*sin(θ)^7
realsphharm8m8(θ::Float64, ϕ::Float64) = (-(3/256))*sqrt(12155/pi)*sin(8*ϕ)*sin(θ)^8
realsphharm8m7(θ::Float64, ϕ::Float64) = (3/64)*sqrt(12155/pi)*cos(θ)*sin(7*ϕ)*sin(θ)^7
realsphharm8m6(θ::Float64, ϕ::Float64) = (-(1/128))*sqrt(7293/(2*pi))*(13 + 15*cos(2*θ))*sin(6*ϕ)*sin(θ)^6
realsphharm8m5(θ::Float64, ϕ::Float64) = (3/256)*sqrt(17017/pi)*(11*cos(θ) + 5*cos(3*θ))*sin(5*ϕ)*sin(θ)^5
realsphharm8m4(θ::Float64, ϕ::Float64) = -((3*sqrt(1309/pi)*(99 + 156*cos(2*θ) + 65*cos(4*θ))*sin(4*ϕ)*sin(θ)^4)/1024)
realsphharm8m3(θ::Float64, ϕ::Float64) = (sqrt(19635/pi)*(126*cos(θ) + 91*cos(3*θ) + 39*cos(5*θ))*sin(3*ϕ)*sin(θ)^3)/1024
realsphharm8m2(θ::Float64, ϕ::Float64) = -((3*sqrt(595/(2*pi))*(210 + 385*cos(2*θ) + 286*cos(4*θ) + 143*cos(6*θ))*sin(2*ϕ)*sin(θ)^2)/2048)
realsphharm8m1(θ::Float64, ϕ::Float64) = (3*sqrt(17/pi)*sin(ϕ)*(70*sin(2*θ) + 154*sin(4*θ) + 286*sin(6*θ) + 715*sin(8*θ)))/8192
realsphharm80(θ::Float64, ϕ::Float64) = (sqrt(17/pi)*(1225 + 2520*cos(2*θ) + 2772*cos(4*θ) + 3432*cos(6*θ) + 6435*cos(8*θ)))/32768
realsphharm81(θ::Float64, ϕ::Float64) = (3*sqrt(17/pi)*cos(ϕ)*(70*sin(2*θ) + 154*sin(4*θ) + 286*sin(6*θ) + 715*sin(8*θ)))/8192
realsphharm82(θ::Float64, ϕ::Float64) = (3*sqrt(595/(2*pi))*cos(2*ϕ)*(210 + 385*cos(2*θ) + 286*cos(4*θ) + 143*cos(6*θ))*sin(θ)^2)/2048
realsphharm83(θ::Float64, ϕ::Float64) = (sqrt(19635/pi)*cos(3*ϕ)*(126*cos(θ) + 91*cos(3*θ) + 39*cos(5*θ))*sin(θ)^3)/1024
realsphharm84(θ::Float64, ϕ::Float64) = (3*sqrt(1309/pi)*cos(4*ϕ)*(99 + 156*cos(2*θ) + 65*cos(4*θ))*sin(θ)^4)/1024
realsphharm85(θ::Float64, ϕ::Float64) = (3/256)*sqrt(17017/pi)*cos(5*ϕ)*(11*cos(θ) + 5*cos(3*θ))*sin(θ)^5
realsphharm86(θ::Float64, ϕ::Float64) = (1/128)*sqrt(7293/(2*pi))*cos(6*ϕ)*(13 + 15*cos(2*θ))*sin(θ)^6
realsphharm87(θ::Float64, ϕ::Float64) = (3/64)*sqrt(12155/pi)*cos(7*ϕ)*cos(θ)*sin(θ)^7
realsphharm88(θ::Float64, ϕ::Float64) = (3/256)*sqrt(12155/pi)*cos(8*ϕ)*sin(θ)^8

i_lm(l::Int, m::Int) = l*(l+1)+1+m

sphharmtab = [
    sphharm00,
    sphharm1m1,
    sphharm10,
    sphharm11,
    sphharm2m2,
    sphharm2m1,
    sphharm20,
    sphharm21,
    sphharm22,
    sphharm3m3,
    sphharm3m2,
    sphharm3m1,
    sphharm30,
    sphharm31,
    sphharm32,
    sphharm33,
    sphharm4m4,
    sphharm4m3,
    sphharm4m2,
    sphharm4m1,
    sphharm40,
    sphharm41,
    sphharm42,
    sphharm43,
    sphharm44,
    sphharm5m5,
    sphharm5m4,
    sphharm5m3,
    sphharm5m2,
    sphharm5m1,
    sphharm50,
    sphharm51,
    sphharm52,
    sphharm53,
    sphharm54,
    sphharm55,
    sphharm6m6,
    sphharm6m5,
    sphharm6m4,
    sphharm6m3,
    sphharm6m2,
    sphharm6m1,
    sphharm60,
    sphharm61,
    sphharm62,
    sphharm63,
    sphharm64,
    sphharm65,
    sphharm66,
    sphharm7m7,
    sphharm7m6,
    sphharm7m5,
    sphharm7m4,
    sphharm7m3,
    sphharm7m2,
    sphharm7m1,
    sphharm70,
    sphharm71,
    sphharm72,
    sphharm73,
    sphharm74,
    sphharm75,
    sphharm76,
    sphharm77,
    sphharm8m8,
    sphharm8m7,
    sphharm8m6,
    sphharm8m5,
    sphharm8m4,
    sphharm8m3,
    sphharm8m2,
    sphharm8m1,
    sphharm80,
    sphharm81,
    sphharm82,
    sphharm83,
    sphharm84,
    sphharm85,
    sphharm86,
    sphharm87,
    sphharm88
]

realsphharmtab = [
    realsphharm00,
    realsphharm1m1,
    realsphharm10,
    realsphharm11,
    realsphharm2m2,
    realsphharm2m1,
    realsphharm20,
    realsphharm21,
    realsphharm22,
    realsphharm3m3,
    realsphharm3m2,
    realsphharm3m1,
    realsphharm30,
    realsphharm31,
    realsphharm32,
    realsphharm33,
    realsphharm4m4,
    realsphharm4m3,
    realsphharm4m2,
    realsphharm4m1,
    realsphharm40,
    realsphharm41,
    realsphharm42,
    realsphharm43,
    realsphharm44,
    realsphharm5m5,
    realsphharm5m4,
    realsphharm5m3,
    realsphharm5m2,
    realsphharm5m1,
    realsphharm50,
    realsphharm51,
    realsphharm52,
    realsphharm53,
    realsphharm54,
    realsphharm55,
    realsphharm6m6,
    realsphharm6m5,
    realsphharm6m4,
    realsphharm6m3,
    realsphharm6m2,
    realsphharm6m1,
    realsphharm60,
    realsphharm61,
    realsphharm62,
    realsphharm63,
    realsphharm64,
    realsphharm65,
    realsphharm66,
    realsphharm7m7,
    realsphharm7m6,
    realsphharm7m5,
    realsphharm7m4,
    realsphharm7m3,
    realsphharm7m2,
    realsphharm7m1,
    realsphharm70,
    realsphharm71,
    realsphharm72,
    realsphharm73,
    realsphharm74,
    realsphharm75,
    realsphharm76,
    realsphharm77,
    realsphharm8m8,
    realsphharm8m7,
    realsphharm8m6,
    realsphharm8m5,
    realsphharm8m4,
    realsphharm8m3,
    realsphharm8m2,
    realsphharm8m1,
    realsphharm80,
    realsphharm81,
    realsphharm82,
    realsphharm83,
    realsphharm84,
    realsphharm85,
    realsphharm86,
    realsphharm87,
    realsphharm88
]


#! /usr/bin/env julia

using DifferentialEquations

b, c = 0.25, 5.0
pend2 = @ode_def Pendulum begin
  dθ = ω
  dω = (-b * ω) - (c * sin(θ))
end

prob = ODEProblem(pend, y0, (0.0, 10.0))
sol = solve(prob)   # solve on [0,10]
t, y = sol.t, hcat(sol.u...)'

using Winston
plot(t, y[:, 1], t, y[:, 2])
title("2D Differential Equation")
savefig("figures/Pendulum_solution.png")


abstract type AbstractRelicDensitySolver end

# ==========================================================================
# ==== ODE Solving with Radau ==============================================
# ==========================================================================

struct ODE
    alg::OrdinaryDiffEqAlgorithm
    xstart::Float64
    xend::Float64
    reltol::Float64
    abstol::Float64
end

function ODE(;
    alg = RadauIIA5(),
    xstart = 1,
    xend = 1e4,
    reltol = 1e-20,
    abstol = 1e-20,
)
    ODE(alg, xstart, xend, reltol, abstol)
end

function boltzmann!(dw, w, model::AbstractDMModel, logx)
    mx = dm_mass(model)
    x = exp(logx)
    T = mx / x
    s = sm_entropy_density(T)
    we = weq(T, mx, dm_dof(model), dm_spin2(model))
    pf = -sqrt(π / 45) * M_PLANK * sm_sqrt_gstar(T) * T
    sv = dm_thermal_cross_section(x, model)

    dw[1] = pf * sv * (exp(w[1]) - exp(2we - w[1]))
end

function boltzmann_jac!(J, w, model::AbstractDMModel, logx)
    mx = dm_mass(model)
    x = exp(logx)
    T = mx / x
    s = sm_entropy_density(T)
    we = weq(T, mx, dm_dof(model), dm_spin2(model))
    pf = -sqrt(π / 45) * M_PLANK * sm_sqrt_gstar(T) * T
    sv = dm_thermal_cross_section(x, model)

    J[1, 1] = pf * sv * (exp(w[1]) + exp(2we - w[1]))
end

function relic_density(model::AbstractDMModel, ode::ODE)
    mx = dm_mass(model)
    logxspan = (log(ode.xstart), log(ode.xend))
    u0 = [weq(mx / ode.xstart, mx, 2.0, 1)]
    f = ODEFunction(boltzmann!, jac = boltzmann_jac!)
    prob = ODEProblem(f, u0, logxspan, model)
    sol = solve(prob, alg = ode.alg, reltol = ode.reltol, abstol = ode.abstol)
    exp(sol.u[end][1]) * mx * S_TODAY / RHO_CRIT
end

# ==========================================================================
# ==== Morrison-Patel-Ulbricht =============================================
# ==========================================================================

struct MPU <: AbstractRelicDensitySolver
    α::Float64
end

function MPU(; α = 0.0)
    MPU(α)
end

function compute_f(x::Real, model::AbstractDMModel)
    mx = dm_mass(model)
    T = mx / x

    sqrt(π / 45) * mx * M_PLANK / x^2 *
    sm_sqrt_gstar(T) *
    dm_thermal_cross_section(x, model)
end

function compute_q(x::Real, model::AbstractDMModel)
    mx = dm_mass(model)
    T = mx / x
    sv = dm_thermal_cross_section(x, model)
    ye = yeq(T, mx, dm_dof(model), dm_spin2(model))
    sqrt(π / 45) * mx * M_PLANK / x^2 * sm_sqrt_gstar(T) * sv * ye
end

function compute_dq(x::Real, model::AbstractDMModel)
    first(ForwardDiff.partials(compute_f(
        ForwardDiff.Dual{Float64}(x, 1.0),
        model,
    )))
end

function compute_p(x::Real, model::AbstractDMModel)
    mx = dm_mass(model)
    T = mx / x
    sv = dm_thermal_cross_section(x, model)
    f(xx) = compute_f(xx, model)

    val = f(ForwardDiff.Dual{ForwardDiff.Dual{Float64}}(
        ForwardDiff.Dual{Float64}(x, 1.0),
        ForwardDiff.Dual{Float64}(1.0, 0.0),
    ))

    f = ForwardDiff.value(ForwardDiff.value(val))
    df = first(ForwardDiff.partials(ForwardDiff.value(val)))
    d2f = first(ForwardDiff.partials(first(ForwardDiff.partials(val))))

    3 * (df / f)^2 / 4 - (d2f / f) / 2
end

function compute_xf(model::AbstractDMModel)
    f(xf) = compute_q(xf, model) - 0.41309880210998464
    fzero(f, 0.1, 100.0)
end

function compute_δx(α, model::AbstractDMModel)
    α == zero(typeof(α)) && return one(typeof(α))
    (
        log(gamma((2 + 3α) / (2 + 2α)) / gamma((2 + α) / (2 + 2α))) / α +
        (1.5772156649015329 + log(1 + α)) / (1 + α)
    )
end

function relic_density(model::AbstractDMModel, mpu::MPU)
    mx = dm_mass(model)
    xf = compute_xf(model)
    δx = compute_δx(mpu.α, model)
    function f(x)
        try
            return compute_f(x, model)
        catch
            return 0.0
        end
    end
    mx * S_TODAY / RHO_CRIT / quadgk(f, xf - δx, Inf)[1]
end


# This is a test of the MRIReco package. In the future I plan to change the reconstruction step
# to use this. I will need to change the raw data output to a ISMRMRDFile.

using MRIReco

f = ISMRMRDFile("/home/ccp/Downloads/knee_3dFSE_slice170.h5")
acqData = AcquisitionData(f);

recoParams = Dict{Symbol, Any}()
recoParams[:reco] = "direct"
recoParams[:reconSize] = (320,320) # this size is also contained in acqData.encodingSize
img = reconstruction(acqData, recoParams)
# #samplingDensity.jl, v \ u uses too much memory

# ##
img = sqrt.(sum(img.^2,dims=5))
using Plots
heatmap(reverse( abs.(img[:,:,1,1,1]),dims=1), c=:viridis)

# multiCoil
# sampling
redFac= 4.0
acqDataSub = sample_kspace(acqData,redFac,"poisson",calsize=30,profiles=false);
# show sampling pattern
msk = zeros(acqDataSub.encodingSize[1],acqDataSub.encodingSize[2])
msk[acqDataSub.subsampleIndices[1]] .= 1
# #sampling
# heatmap(reverse( msk,dims=1), c=:grays)

#ESPIRIT
# coil sensitivities
smaps = espirit(acqData,(6,6),30,eigThresh_1=0.035,eigThresh_2=0.98)
# SENSE reconstruction
params = Dict{Symbol, Any}()
params[:reco] = "multiCoil"
params[:reconSize] = (320,320) #previously it says :shape
params[:senseMaps] = smaps
params[:solver] = "cgnr"
params[:regularization] = "L2"
params[:λ] = 1e-4
params[:iterations] = 1
params[:normalizeReg] = true

img_reco = reconstruction(acqDataSub, params)

## show image
heatmap(reverse( abs.(img_reco[:,:,1,1,1]),dims=1), c=:viridis)

##
using MRIsim

imcoils = []
Nc = numChannels(data)
for coil = 1:Nc
    raw = reshape(data.kdata[1][:,coil], sz[1:2]...)
    img = MRIsim.ifftc(raw)
    append!(imcoils, [img])
end
imgss = sqrt.( sum([abs.(imcoils[i]).^2 for i=1:Nc]) ) #sum of squares

##

l = PlotlyJS.Layout(;height=500, width=500, colorbar=false)
p = PlotlyJS.heatmap(z = abs.(imgss), colorscale="Greys")
PlotlyJS.plot(p,l)


#=

  Secret Santa problem in Julia ConstraintSolver.jl

  From Ruby Quiz Secret Santa
  http://www.rubyquiz.com/quiz2.html
  """
  Honoring a long standing tradition started by my wife's dad, my friends
  all play a Secret Santa game around Christmas time. We draw names and
  spend a week sneaking that person gifts and clues to our identity. On the
  last night of the game, we get together, have dinner, share stories, and,
  most importantly, try to guess who our Secret Santa was. It's a crazily
  fun way to enjoy each other's company during the holidays.

  To choose Santas, we use to draw names out of a hat. This system was
  tedious, prone to many "Wait, I got myself..." problems. This year, we
  made a change to the rules that further complicated picking and we knew
  the hat draw would not stand up to the challenge. Naturally, to solve
  this problem, I scripted the process. Since that turned out to be more
  interesting than I had expected, I decided to share.

  This weeks Ruby Quiz is to implement a Secret Santa selection script.

  Your script will be fed a list of names on STDIN.
  ...
  Your script should then choose a Secret Santa for every name in the list.
  Obviously, a person cannot be their own Secret Santa. In addition, my friends
  no longer allow people in the same family to be Santas for each other and your
  script should take this into account.
  """

  Comment: Well, this model skips the file input and mail parts. We
           assume that the friends are identified with a number from 1..n,
           and the families is identified with a number 1..num_families.

           And we also do some random family configurations below.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
using Distributions
const CS = ConstraintSolver
include("constraints_utils.jl")

function secret_santa(family,print_solutions=true,all_solutions=false)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = length(family)
    @variable(model, 1 <= x[1:n] <= n, Int)

    # Everyone gives and receives a Secret Santa
    @constraint(model, x in CS.AllDifferent())

    # Can't be one own's Secret Santa
    for i in 1:n
        @constraint(model, x[i] != i)
    end

    # No Secret Santa to a person in the same family
    @variable(model,1 <= fxi[1:n] <= n, Int)
    for i in 1:n
        my_element(model, x[i],family,fxi[i])
        @constraint(model,family[i] != fxi[i])
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        if all_solutions
            println("num_sols:$num_sols\n")
        end

        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                xx = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$xx")

            end
        end
        if all_solutions
            println("num_sols:$num_sols\n")
        end
    else
        println("status:$status")
    end

    return status
end

#
# Testing
#

# The families
secret_santa_problems =  Dict(
:1 =>  [1,1,2,2, 3, 4,4], # Ruby Quiz example. 672 solutions
:2 =>  [1,1,1,1, 2, 3,3,3,3,3, 4,4], # From my Picat model
)
print_solutions=false
all_solutions=true
# Number of solutions for the simple instance
@time secret_santa(secret_santa_problems[:1],print_solutions,all_solutions)

# Generate random families

function random_families(n=10)
    d = round(Int,n/2)
    rand(DiscreteUniform(1,d),n)|>sort
end

print_solutions=true
all_solutions=false

# Test larger instances
for n in [i*10 for i in 1:10]
    println("\nn:$(n)")
    family = random_families(n)
    println("family:$family")
    c = 0
    timeout = false
    while c < 10
        @time status = secret_santa(family,print_solutions,all_solutions)
        c += 1
        if status == MOI.TIME_LIMIT
            timeout = true
        end
        if status == MOI.OPTIMAL || timeout || c >= 10
            break
        end
    end
    if timeout
        break
    end
end


# Experimental threads example. WIP!

using CmdStan
using Statistics
using Distributions

ProjDir = @__DIR__
cd(ProjDir) #do

bernoullimodel = "
data {
  int<lower=1> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  y ~ bernoulli(theta);
}
";

#isdir(ProjDir * "/tmp") && rm(ProjDir * "/tmp", recursive=true)
tmpdir = ProjDir * "/tmp"

p1 = 24 # p1 is the number of models to fit

n = 100;
p = range(0, stop=1, length=p1)
observeddata = [Dict("N" => n, "y" => Int.(rand(Bernoulli(p[i]), n))) for i in 1:p1]

sm = [Stanmodel(name="bernoulli_m$i", model=bernoullimodel;
  output_format=:namedtuple, tmpdir=tmpdir) for i in 1:p1]

println("\nThreads loop\n")

estimates = Vector(undef, p1)
Threads.@threads for i in 1:p1
    rc, samples, cnames = stan(sm[i], observeddata[i]; summary=false);
    if rc == 0
      estimates[i] = [mean(reshape(samples.theta, 4000)), std(reshape(samples.theta, 4000))]
    end
end

estimates |> display

#end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    BAT.MultimodalStudentT <: Distribution{Multivariate,Continuous}

*Experimental feature, not part of stable public API.*

The Multimodal Student-t Distribution (It's defined similarly to multimodal Cauchy
defined in [Caldwell et al.](https://arxiv.org/abs/1808.08051)).

Assumes two bimodal peaks, each in its own dimension.

Constructors:

* ```BAT.MultimodalStudentT(; μ::Real=1, σ::Float64=0.2, ν::Integer=1, n::Integer=4)```

Arguments:

* `μ::Real`: The location parameter used for the two bimodal peaks.
* `σ::Float64`: The scale parameter shared among all components.
* `ν::Int`: The degrees of freedom.
* `n::Int`: The number of dimensions.

Fields:

$(TYPEDFIELDS)

!!! note

    All fields of `MultimodalStudentT` are considered internal and subject to
    change without deprecation.
"""
struct MultimodalStudentT{M<:MixtureModel, P<:Product} <: Distribution{Multivariate,Continuous}
    bimodals::M
    σ::Float64
    n::Int
    dist::P
end

function MultimodalStudentT(;μ::Real=1., σ::Float64=0.2, ν::Int=1, n::Integer=4)
    @argcheck n > 1 "Minimum number of dimensions for MultimodalCauchy is 2"
    mixture_model = MixtureModel([LocationScale(-μ, σ, TDist(ν)), LocationScale(μ, σ, TDist(ν))])
    dist = _construct_dist(mixture_model, σ, ν, n)
    MultimodalStudentT(mixture_model, σ, n, dist)
end

Base.size(d::MultimodalStudentT) = size(d.dist)
Base.length(d::MultimodalStudentT) = length(d.dist)
Base.eltype(d::MultimodalStudentT) = eltype(d.dist)

Statistics.mean(d::MultimodalStudentT) = Distributions.mean(d.dist)
Statistics.var(d::MultimodalStudentT) = Distributions.var(d.dist)
Statistics.cov(d::MultimodalStudentT) = Distributions.cov(d.dist)

function StatsBase.params(d::MultimodalStudentT)#Check this
    (
        vcat(d.bimodals.components[1].μ, d.bimodals.components[2].μ, zeros(d.n-2)),
        [d.σ for i in 1:d.n],
        d.n
    )
end

function Distributions._logpdf(d::MultimodalStudentT, x::AbstractArray)
    convert(eltype(x), Distributions._logpdf(d.dist, x))::eltype(x)
end

function Distributions._rand!(rng::AbstractRNG, d::MultimodalStudentT, x::AbstractVector)
    Distributions._rand!(rng, d.dist, x)
end

function _construct_dist(mixture_model::MixtureModel, σ::Real, ν::Int, n::Integer)
    vector_of_dists = vcat(mixture_model, mixture_model, [LocationScale(0, σ, TDist(ν)) for i in 3:n])
    dist = product_distribution(vector_of_dists)
    return dist
end


module states

import Base: ==, +, -, *, /
import ..bases

using ..bases

export StateVector, Bra, Ket,
       tensor, dagger,
       basis_bra, basis_ket, basis,
       normalize, normalize!

"""
Abstract base class for Bra and Ket.

The state vector class stores the coefficients of an abstract state
in respect to a certain basis. These coefficients are stored in the
``data`` attribute and the basis is defined in the ``basis``
attribute.
"""
abstract StateVector

type Bra <: StateVector
    basis::Basis
    data::Vector{Complex{Float64}}
    Bra(b::Basis, data) = length(b) == length(data) ? new(b, data) : throw(DimensionMismatch())
end

type Ket <: StateVector
    basis::Basis
    data::Vector{Complex{Float64}}
    Ket(b::Basis, data) = length(b) == length(data) ? new(b, data) : throw(DimensionMismatch())
end

Bra(b::Basis) = Bra(b, zeros(Complex, length(b)))
Ket(b::Basis) = Ket(b, zeros(Complex, length(b)))

Base.eltype(x::StateVector) = Complex128
Base.zero{T<:StateVector}(x::T) = T(x.basis)

=={T<:StateVector}(x::T, y::T) = (x.basis == y.basis) && (x.data == y.data)

# Arithmetic operations
*(a::Bra, b::Ket) = (check_multiplicable(a.basis, b.basis); sum(a.data.*b.data))
*{T<:StateVector}(a::Number, b::T) = T(b.basis, complex(a)*b.data)
*{T<:StateVector}(a::T, b::Number) = T(a.basis, complex(b)*a.data)

/{T<:StateVector}(a::T, b::Number) = T(a.basis, a.data/complex(b))

+{T<:StateVector}(a::T, b::T) = (a.basis==b.basis ? T(a.basis, a.data+b.data) : throw(IncompatibleBases()))

-{T<:StateVector}(a::T, b::T) = (a.basis==b.basis ? T(a.basis, a.data-b.data) : throw(IncompatibleBases()))

"""
Tensor product of given bras or kets.
"""
bases.tensor{T<:StateVector}(a::T, b::T) = T(tensor(a.basis, b.basis), kron(a.data, b.data))
bases.tensor{T<:StateVector}(states::T...) = reduce(tensor, states)

"""
Hermitian conjugate of the given state vector.
"""
dagger(x::Bra) = Ket(x.basis, conj(x.data))
dagger(x::Ket) = Bra(x.basis, conj(x.data))

# Normalization functions
"""
Norm of the given state vector.
"""
Base.norm(x::StateVector, p=2) = norm(x.data, p)

"""
Normalized copy of the given state vector.
"""
normalize(x::StateVector, p=2) = x/norm(x, p)

"""
Normalize the given state vector.
"""
function normalize!(x::StateVector, p=2)
    u = 1./norm(x, p)
    for i=1:length(x.data)
        x.data[i]*=u
    end
    return x
end

# Creation of basis states.
function basis_vector(shape::Vector{Int}, indices::Vector{Int})
    x = zeros(Complex, shape...)
    x[indices] = Complex(1.)
    reshape(x, prod(shape))
end

basis_bra(b::Basis, indices::Array{Int}) = Bra(b, basis_vector(b.shape, indices))
basis_bra(b::Basis, index::Int) = basis_bra(b, [index])
basis_ket(b::Basis, indices::Array{Int}) = Ket(b, basis_vector(b.shape, indices))
basis_ket(b::Basis, index::Int) = basis_ket(b, [index])

end # module


# Simple tests to make sure Roots works with other
# packages. This is not part of runtests, as no expectation
# that CI should run these

using Roots
using Compat.Test
using Unitful
using SymEngine

@testset "Test composability with other packages" begin

    orders = [Order0(), Order1(), Order2(), Order5(),
              Order8(), Order16()]

    # Unitful
    @testset "find zero(s) with Unitful" begin

        s = u"s"; m = u"m"
        g = 9.8*m/s^2
        v0 = 10m/s
        y0 = 16m
        y(t) = -g*t^2 + v0*t + y0

        for order in orders
            @test find_zero(y, 1.8s, order) ≈  1.886053370668014s
        end

        for M in [Roots.Bisection(), Roots.A42(), Roots.AlefeldPotraShi()]
            @test find_zero(y, (1.8s, 1.9s), M) ≈ 1.886053370668014s
        end

        xrts = find_zeros(y, 0s, 10s)
        @test length(xrts) == 1
        @test xrts[1] ≈  1.886053370668014s

    end

    # SymEngine
    @testset "find zero(s) with SymEngine" begin

        m = s = 1.0
        g = 9.8*m/s^2
        v0 = 10m/s
        y0 = 16m
        y(t) = -g*t^2 + v0*t + y0

        @vars x
        for order in orders
            @test find_zero(y(x), 1.8, order) ≈  1.8860533706680143
        end

        for M in [Roots.Bisection(), Roots.A42(), Roots.AlefeldPotraShi()]
            @test find_zero(y(x), (1.8, 1.9), M) ≈ 1.8860533706680143
        end

        xrts = find_zeros(y(x), 0, 10)
        @test length(xrts) == 1
        @test xrts[1] ≈  1.8860533706680143
    end

end


using ArgParse
using Cambrian
using CartesianGeneticProgramming
using Dates
using IICGP
using Random
using UnicodePlots
using LinearAlgebra

import Cambrian.mutate

out(plt) = println(IOContext(stdout, :color=>true), plt)
default_resdir = joinpath(dirname(@__DIR__), "results")
default_cfgdir = joinpath(dirname(@__DIR__), "cfg")

s = ArgParseSettings()
@add_arg_table! s begin
    "--cfg"
    help = "configuration script"
    default = joinpath(default_cfgdir, "nsga2_example.yaml")
    "--seed"
    help = "random seed"
    default = 0
    "--out"
    help = "output directory"
    arg_type = String
    default = default_resdir
end

args = parse_args(ARGS, s)
const cfg_path = args["cfg"]
const seed = args["seed"]
const resdir = args["out"]
const cfg = get_config(cfg_path)

# Specific functions for this example of learning two functions
x = [i for i in 0:0.05:1]
const max_fitness = [0.0, 0.0]
const min_fitness = - length(x) * ones(cfg.d_fitness)
f1(x) = 0.2 * (2.0 + cos(2.0*π*x))
f2(x) = 0.25 * (2.0 + cos(4*π*x))
y1 = f1.(x)
y2 = f2.(x)
out(lineplot([f1, f2], 0, 1, border=:dotted))
function display_paretto(e::NSGA2Evo)
    o1 = [ind.fitness[1] for ind in e.population]
    o2 = [ind.fitness[2] for ind in e.population]
    out(scatterplot(o1, o2, title = "Paretto front"))#, xlim=[0,1], ylim=[0,1]))
end

# User-defined fitness function (normalized)
function my_fitness(ind::NSGA2Ind, seed::Int64, game::Game) # game is a placeholder
    cgp_ind = CGPInd(cfg, ind.chromosome)
    y_hat = zeros(length(y1))
    @inbounds for i in eachindex(y1)
        y_hat[i] = CartesianGeneticProgramming.process(cgp_ind, [x[i]])[1]
    end
    o1 = - norm(y1 - y_hat)
    o2 = - norm(y2 - y_hat)
    ([o1, o2] .- min_fitness) ./ (max_fitness .- min_fitness)
end

# User-defined mutation function
function mutate(ind::NSGA2Ind)
    cgpind = CGPInd(cfg, ind.chromosome)
    child = goldman_mutate(cfg, cgpind)
    NSGA2Ind(cfg, child.chromosome)
end

# User-defined population initialization function
function my_init(cfg::NamedTuple)
    [NSGA2Ind(cfg, CGPInd(cfg).chromosome) for _ in 1:cfg.n_population]
end

e = NSGA2Evo(cfg, resdir, my_fitness, my_init, "boxing") # boxing is a placeholder

init_backup(cfg.id, resdir, cfg_path)
#run!(evo)
for i in 1:e.config.n_gen
    e.gen += 1
    if e.gen > 1
        populate(e)
    end
    evaluate(e)
    display_paretto(e)
    generation(e, max_fitness, min_fitness)
    #=if ((e.config.log_gen > 0) && mod(e.gen, e.config.log_gen) == 0)
        log_gen(e, max_fitness, min_fitness)
    end
    if ((e.config.save_gen > 0) && mod(e.gen, e.config.save_gen) == 0)
        save_gen(e)
    end=#
end

# Display one of the elite (Pareto efficient) individuals
ind = CGPInd(cfg, e.population[1].chromosome)
cgp_f(x) = CartesianGeneticProgramming.process(ind, [x])[1]
out(lineplot([f1, f2, cgp_f], 0, 1, border=:dotted))


# using SymPy

# x, y = @vars x y
# u = x*(1-x)*y*(1-y)
# v = x*(1-x)*y*(1-y)^2

# p = x*(1-x)*y*(1-y)

# ux = diff(u,x)
# uy = diff(u,y)
# vx = diff(v,x)
# vy = diff(v,y)
# px = diff(p,x)
# py = diff(p,y)
# f = u*ux + v*uy - (diff(ux,x)+diff(uy,y)) + px
# g = u*vx + v*vy - (diff(vx,x)+diff(vy,y)) + py
# h = diff(u, x) + diff(v, y)
# println(replace(replace(sympy.julia_code(simplify(f)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(g)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(h)), ".^"=>"^"), ".*"=>"*"))

function u_exact(x,y)
    x*(1-x)*y*(1-y)
end

function v_exact(x,y)
    x*(1-x)*y*(1-y)
end

function p_exact(x,y)
    x*(1-x)*y*(1-y)
end

function t_exact(x,y)
    x*(1-x)*y*(1-y)
end

function k_exact(x,y)
    1.0
end

function ffunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - y) + 0.02*x*(1 - x) + y*(1 - x)*(1 - y) + 0.02*y*(1 - y)
end

function gfunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - x) + x*(1 - x)*(1 - y) + 0.02*x*(1 - x) + 0.02*y*(1 - y)
end

function hfunc_(x,y)
    -x*y*(1 - x) - x*y*(1 - y) + x*(1 - x)*(1 - y) + y*(1 - x)*(1 - y)
end

function heat_source_func(x, y)
    x^2*y*(x - 1)^2*(y - 1)*(2*y - 1) + x*y^2*(x - 1)*(2*x - 1)*(y - 1)^2 - 2.0*x*(x - 1) - 2.0*y*(y - 1)
end

function q1_func_(t, x, y)
    -1.0*x*y*(1 - x)*(1 - y)
end

function q2_func_(t, x, y)
    -1.0*x*y*(1 - x)*(1 - y)
end

using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

m = 50
n = 50
h = 1/n
nu = 0.01
NT_transport = 500
Δt = 1/NT_transport
κ1 = 1.0
κ2 = 1.0

NT = 5    # number of iterations for Newton's method

bd = bcnode("all", m, n, h)
bd = [bd; bd .+ (m+1)*(n+1);
    2*(m+1)*(n+1)+(n-1)*m+1:2*(m+1)*(n+1)+(n-1)*m+2;
    (2*(m+1)*(n+1)+m*n) .+ bd]
# only apply Dirichlet to velocity; set left bottom two points to zero to fix rank deficient problem for pressure

## STEP 1: NAVIER STOKES EQUATIONS
F1 = compute_fem_source_term1(eval_f_on_gauss_pts(ffunc_, m, n, h), m, n, h)
F2 = compute_fem_source_term1(eval_f_on_gauss_pts(gfunc_, m, n, h), m, n, h)
H = h^2*eval_f_on_fvm_pts(hfunc_, m, n, h)
B = constant(compute_interaction_matrix(m, n, h))

# compute F
Laplace = nu * constant(compute_fem_laplace_matrix1(m, n, h))
heat_source = eval_f_on_gauss_pts(heat_source_func, m, n, h)
heat_source = constant(compute_fem_source_term1(heat_source, m, n, h))
# The temperature term is
# T_xx + T_yy + q = u T_x + v T_y

function compute_residual(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
        S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
        S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
        S[2(m+1)*(n+1)+m*n+1:end]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    interaction = compute_interaction_term(p, m, n, h) # julia kernel needed
    f1 = compute_fem_source_term1(ugauss.*ux, m, n, h)
    f2 = compute_fem_source_term1(vgauss.*uy, m, n, h)
    f3 = -interaction[1:(m+1)*(n+1)]
    f4 = Laplace*u
    f5 = -F1
    F = f1 + f2 + f3 + f4 + f5

    g1 = compute_fem_source_term1(ugauss.*vx, m, n, h)
    g2 = compute_fem_source_term1(vgauss.*vy, m, n, h)
    g3 = -interaction[(m+1)*(n+1)+1:end]
    g4 = Laplace*v
    g5 = -F2
    G = g1 + g2 + g3 + g4 + g5

    H0 = -B * [u;v] + H

    # T_xx + T_yy = u T_x + v T_y - heat_source
    # ugauss = zeros_like(ugauss)
    # vgauss = zeros_like(vgauss)
    A = constant(compute_fem_laplace_matrix1(m, n, h))
    T0 = A * T + compute_fem_advection_matrix1(ugauss,vgauss, m, n, h) * T - heat_source
    R = [F;G;H0;T0]
    return R
end

function compute_jacobian(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
                 S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
                 S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
                 S[2(m+1)*(n+1)+m*n+1:end]

    graduv = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = graduv[:,1,1], graduv[:,1,2], graduv[:,2,1], graduv[:,2,2]

    M1 = constant(compute_fem_mass_matrix1(ux, m, n, h))
    M2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h)) # a julia kernel needed
    M3 = Laplace
    Fu = M1 + M2 + M3

    Fv = constant(compute_fem_mass_matrix1(uy, m, n, h))

    N1 = constant(compute_fem_mass_matrix1(vy, m, n, h))
    N2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
    N3 = Laplace
    Gv = N1 + N2 + N3

    Gu = constant(compute_fem_mass_matrix1(vx, m, n, h))

    M = constant(compute_fem_laplace_matrix1(m, n, h)) +
         compute_fem_advection_matrix1(ugauss,vgauss, m, n, h)

    gradT = eval_grad_on_gauss_pts1(T, m, n, h)
    Tx, Ty = gradT[:,1], gradT[:,2]
    DU_TX = constant(compute_fem_mass_matrix1(Tx, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)
    DV_TY = constant(compute_fem_mass_matrix1(Ty, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)

    J0 = [Fu Fv                         # Jacobian of [u; v]
          Gu Gv]

    J1 = [J0 -B'                        # Jacobian of [u; v; p]
          -B spdiag(zeros(size(B,1)))]

    N = 2*(m+1)*(n+1) + m*n                                     # length of [u; v; p]
    J = [J1 SparseTensor(spzeros(N,(m+1)*(n+1)))                # Jacobian of [u; v; p; T]
        DU_TX DV_TY SparseTensor(spzeros((m+1)*(n+1), m*n)) M]
end

function solve_steady_cavityflow_one_step(S)
    residual = compute_residual(S)
    J = compute_jacobian(S)

    J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
    residual = scatter_update(residual, bd, zeros(length(bd)))    # residual[bd] .= 0.0 in Tensorflow syntax

    d = J\residual
    residual_norm = norm(residual)
    op = tf.print("residual norm", residual_norm)
    d = bind(d, op)
    S_new = S - d
    return S_new
end

function condition(i, S_arr)
    i <= NT + 1
end

function body(i, S_arr)
    S = read(S_arr, i-1)
    op = tf.print("i=",i)
    i = bind(i, op)
    S_new = solve_steady_cavityflow_one_step(S)
    S_arr = write(S_arr, i, S_new)
    return i+1, S_arr
end

xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
u0 = @. u_exact(x,y)
v0 = @. v_exact(x,y)
T0 = @. t_exact(x,y)

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p0 = @. p_exact(x,y)

S_arr = TensorArray(NT+1)
S_arr = write(S_arr, 1, [u0; v0; p0; T0])
# S_arr = write(S_arr, 1, zeros(m*n+3*(m+1)*(n+1)))

i = constant(2, dtype=Int32)

_, S = while_loop(condition, body, [i, S_arr])
S = set_shape(stack(S), (NT+1, 2*(m+1)*(n+1)+m*n+(m+1)*(n+1)))

u = S[NT+1, 1:(m+1)*(n+1)]
v = S[NT+1, (m+1)*(n+1)+1:2*(m+1)*(n+1)]
p = S[NT+1, 2*(m+1)*(n+1)+1:2*(m+1)*(n+1)+m*n]
T = S[NT+1, 2*(m+1)*(n+1)+m*n+1:end]

## STEP 2: TRANSPORT EQUATION
# pre-compute source term
Q1 = zeros(NT_transport, (m+1)*(n+1))
Q2 = zeros(NT_transport, (m+1)*(n+1))
for i = 1:NT_transport
    t = i*Δt
    Q1[i, :] = eval_f_on_fem_pts((x,y)->q1_func_(t, x, y), m, n, h)
    Q2[i, :] = eval_f_on_fem_pts((x,y)->q2_func_(t, x, y), m, n, h)
end
Q1 = constant(Q1)
Q2 = constant(Q2)

function solve_transport_step(w1, w2, q1, q2)
    w1 = (1/Δt * w1 + κ1 * u + q1) / (1/Δt + κ1)
    w2 = (1/Δt * w2 + κ2 * v + q2) / (1/Δt + κ2)
    w1, w2
end

function transport_condition(i, w1_arr, w2_arr)
    i <= NT_transport
end

function transport_body(i, w1_arr, w2_arr)
    w1 = read(w1_arr, i)
    w2 = read(w2_arr, i)
    q1 = Q1[i]
    q2 = Q2[i]
    # op = tf.print("transport equation, step: ", i)
    # i = bind(i, op)
    w1, w2 = solve_transport_step(w1, w2, q1, q2)
    i+1, write(w1_arr, i+1, w1), write(w2_arr, i+1, w2)
end

i = constant(1, dtype = Int32)
w1_arr = TensorArray(NT_transport+1)
w2_arr = TensorArray(NT_transport+1)
w1_0 = eval_f_on_fem_pts((x,y)->x*y*(1-x)*(1-y), m, n, h)
w2_0 = eval_f_on_fem_pts((x,y)->x*y*(1-x)*(1-y), m, n, h)
w1_arr = write(w1_arr, 1, w1_0)
w2_arr = write(w2_arr, 1, w2_0)
_, w1, w2 = while_loop(transport_condition, transport_body, [i, w1_arr, w2_arr])
w1, w2 = stack(w1), stack(w2)
w1 = set_shape(w1, (NT_transport+1, (m+1)*(n+1)))
w2 = set_shape(w2, (NT_transport+1, (m+1)*(n+1)))

sess = Session(); init(sess)
W1, W2, u_, v_, p_ = run(sess, [w1, w2, u, v, p])

figure(figsize=(15,5))
subplot(131)
visualize_scalar_on_fem_points(W1[end,:], m, n, h)
title("Computed droplet x-velocity")
subplot(132)
visualize_scalar_on_fem_points(eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Exact droplet x-velocity")
subplot(133)
visualize_scalar_on_fem_points(W1[end,:]-eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Difference")

figure(figsize=(15,5))
subplot(131)
visualize_scalar_on_fem_points(W2[end,:], m, n, h)
title("Computed droplet y-velocity")
subplot(132)
visualize_scalar_on_fem_points(eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Exact droplet y-velocity")
subplot(133)
visualize_scalar_on_fem_points(W2[end,:]-eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Difference")
tight_layout()

# # S = output
# # # out_v = output[:, 1:2*(m+1)*(n+1)]
# # # out_p = output[:, 2*(m+1)*(n+1)+1:end]

matwrite("covid_data.mat", Dict(
        "W1"=>W1, "W2"=>W2))


using Plots

function mainAlex(NN=20,TT=3650,pb=0.01,cb=1,cr=4)

    AA = dynamicCommunicators(NN,TT,pb,cb,cr)
    degOut = totalOut(AA,NN)
    Cbroad = dynamicCentrality(AA,NN,TT)
    scatterPlot(AA,size(AA)[1],degOut,Cbroad)

   # AA2 = ringNet()
   # degOut = totalOut(AA2,21)
   # Cbroad = dynamicCentrality(AA2,21,3)
   # scatterPlot(AA2,size(AA2)[1],degOut,Cbroad)
end

function scatterPlot(AA,NN,degOut,Cbroad)
    labels = [string(x) for x in 1:NN]
    scatter(degOut,Cbroad,markersize=1,series_annotations =labels,grid=false)
    xlabel!("total deg out")
    ylabel!("broadcast centrality")
    title!(string("node num=",NN))

end

function ringNet()
    AA2 = zeros(21,21,3)
    AA2[1,7,1] =1;AA2[7,1,1] =1;AA2[2,7,1] =1;AA2[7,2,1] =1;
    AA2[21,20,1] =1;AA2[20,21,1] =1;AA2[21,19,1] =1;AA2[19,21,1] =1;
    AA2[2,8,2]=1;AA2[8,2,2]=1;AA2[3,8,2]=1;AA2[8,3,2]=1;
    AA2[20,18,2]=1;AA2[18,20,2]=1;AA2[20,17,2]=1;AA2[17,20,2]=1;
    AA2[19,16,2]=1;AA2[16,19,2]=1;
    AA2[4,10,3]=1;AA2[10,4,3]=1;AA2[5,10,3]=1;AA2[10,5,3]=1;
    AA2[18,15,3]=1;AA2[15,18,3]=1;AA2[18,14,3]=1;AA2[14,18,3]=1;
    AA2[17,14,3]=1;AA2[14,17,3]=1;AA2[17,13,3]=1;AA2[13,17,3]=1;
    AA2[16,13,3]=1;AA2[13,16,3]=1;AA2[16,12,3]=1;AA2[12,16,3]=1;
    return AA2
end

function dynamicCommunicators(NN,TT,pb,cb,cr)
    imp = [x^2 for x in 1:NN]
    #imp = 2*ones(1,NN)
    #imp[NN] = 100
    println(imp)
    println("that was imp vector now normal")
    AA = zeros(NN,NN,TT)
    #INIT DATA>>
    for ii in 1:NN
        if(rand(1)[1] <= pb)
            tmp = deleteat!(collect(1:NN),ii)
            destinationN = tmp[rand(1:end)]
            AA[ii,destinationN,1] = 1
        end

    end
    #DATA LOOP: 1st put basal edges into <present> / 2nd put response from past to now
    for tt in 2:(TT)
        #println(string("tt=",tt))
        #BASAL LOOP
        for ii in 1:NN
            if(rand(1)[1] <= pb)
                tmp = deleteat!(collect(1:NN),ii)
                destinationN = tmp[rand(1:end)]
                AA[ii,destinationN,tt] = 1#FIRE NOW
            end

        end
        #RESPONSE LOOP
        for ii in 1:NN
            msgsTo_ii = AA[:,ii,tt-1]#RESPOND TO PREVIOUS STATE (tt-1)
            if(sum(AA[:,ii,tt-1]) >= 1)

                totalImportance = sum([msgsTo_ii[x]*imp[x] for x in 1:NN])
                #println(msgsTo_ii)
                #println(imp)
                #println(sum([msgsTo_ii[x]*imp[x] for x in 1:NN]))

                r_nNumerator = totalImportance
                r_nDenominator = 1 + (findmax(imp)[1] * sum(AA[:,ii,tt-1]))
                r_next = r_nNumerator / r_nDenominator
                println(r_next)
                if(r_next >=  rand(1)[1])#if so generate cr links

                    tmp = deleteat!(collect(1:NN),ii)
                    destinationNodes = tmp[randperm(length(tmp))[1:cr]]
                    AA[ii,destinationNodes,tt] = 1#FIRE NOW
                    println("fire!")

               end
            end
        end

    end

    return AA
end

function totalOut(AA,NN)
    coltmp = zeros(NN,1)
    outDegVec = sum(sum(AA,3),2)
    coltmp[:] = outDegVec[:]
    return coltmp
end

function dynamicCentrality(AA,NN,TT)
    println(":")
    daynorm = zeros(TT,1);
    for ii in 1:TT
        Atemp = AA[:,:,ii]
        daynorm[ii] = findmax(abs(eigvals(Atemp)))[1]
    end
    alpha = 0.9*(1/findmax(daynorm)[1]);#90percent of the max
    println(alpha)
    #alpha = 0.15;
    BB = eye(NN,NN);
    for kk = TT:-1:1
        Aday = AA[:,:,kk]
        #println(Aday)
        BBtmp = inv(eye(NN,NN) - alpha*Aday)#\BB
        #BBtmp = (eye(NN,NN) - alpha*Aday) \ BB
        BBtmp = abs(BBtmp)
        #BB = BBtmp / norm(BBtmp)
        BB = (BB*BBtmp) / norm(BB)
    end
    #sum across columns
    Cbroad = sum(BB',2)
    return Cbroad
end


using Statistics

@testsuite "statistics" AT->begin
    @testset "std" begin
        @test compare(std, AT, rand(10))
        @test compare(std, AT, rand(10,1,2))
        @test compare(std, AT, rand(10,1,2); corrected=true)
        @test compare(std, AT, rand(10,1,2); dims=1)
    end

    @testset "var" begin
        @test compare(var, AT, rand(10))
        @test compare(var, AT, rand(10,1,2))
        @test compare(var, AT, rand(10,1,2); corrected=true)
        @test compare(var, AT, rand(10,1,2); dims=1)
        @test compare(var, AT, rand(10,1,2); dims=[1])
        @test compare(var, AT, rand(10,1,2); dims=(1,))
        @test compare(var, AT, rand(10,1,2); dims=[2,3])
        @test compare(var, AT, rand(10,1,2); dims=(2,3))
    end

    @testset "mean" begin
        @test compare(mean, AT, rand(2,2))
        @test compare(mean, AT, rand(2,2); dims=2)
        @test compare(mean, AT, rand(2,2,2); dims=[1,3])
        @test compare(x->mean(sin, x), AT, rand(2,2))
        @test compare(x->mean(sin, x; dims=2), AT, rand(2,2))
        @test compare(x->mean(sin, x; dims=[1,3]), AT, rand(2,2,2))
    end

    @testset "cov" begin
        s = 100
        @test compare(cov, AT, rand(s))
        @test compare(cov, AT, rand(Complex{Float64}, s))
        @test compare(cov, AT, rand(s, 2))
        @test compare(cov, AT, rand(Complex{Float64}, s, 2))
        @test compare(cov, AT, rand(s, 2); dims=2)
        @test compare(cov, AT, rand(Complex{Float64}, s, 2); dims=2)
        @test compare(cov, AT, rand(1:100, s))
    end

    @testset "cor" begin
        s = 100
        @test compare(cor, AT, rand(s))
        @test compare(cor, AT, rand(Complex{Float64}, s))
        @test compare(cor, AT, rand(s, 2))
        @test compare(cor, AT, rand(Complex{Float64}, s, 2))
        @test compare(cor, AT, rand(s, 2); dims=2)
        @test compare(cor, AT, rand(Complex{Float64}, s, 2); dims=2)
        @test compare(cor, AT, rand(1:100, s))
    end
end


# Parts of this code were taken / derived from Graphs.jl. See LICENSE for
# licensing details.

# The Bellman Ford algorithm for single-source shortest path

###################################################################
#
#   The type that capsulates the state of Bellman Ford algorithm
#
###################################################################
using Base.Threads

struct NegativeCycleError <: Exception end

# AbstractPathState is defined in core
"""
    BellmanFordState{T, U}

An `AbstractPathState` designed for Bellman-Ford shortest-paths calculations.
"""
struct BellmanFordState{T<:Real, U<:Integer} <: AbstractPathState
    parents::Vector{U}
    dists::Vector{T}
end

"""
    bellman_ford_shortest_paths(g, s, distmx=weights(g))
    bellman_ford_shortest_paths(g, ss, distmx=weights(g))

Compute shortest paths between a source `s` (or list of sources `ss`) and all
other nodes in graph `g` using the [Bellman-Ford algorithm](http://en.wikipedia.org/wiki/Bellman–Ford_algorithm).
Return a [`LightGraphs.BellmanFordState`](@ref) with relevant traversal information.
"""
function bellman_ford_shortest_paths(
    graph::AbstractGraph{U},
    sources::AbstractVector{<:Integer},
    distmx::AbstractMatrix{T}=weights(graph)
    ) where T<:Real where U<:Integer

    active = Set{U}(sources)
    sizehint!(active, nv(graph))
    dists = fill(typemax(T), nv(graph))
    parents = zeros(U, nv(graph))
    dists[sources] .= 0
    no_changes = false

    for i in one(U):nv(graph)
        no_changes = true
        new_active = Set{U}()
        for u in active
            for v in outneighbors(graph, u)
                relax_dist = distmx[u, v] + dists[u]
                if dists[v] > relax_dist
                    dists[v] = relax_dist
                    parents[v] = u
                    no_changes = false
                    push!(new_active, v)
                end
            end
        end
        if no_changes
            break
        end
        active = new_active
    end
    no_changes || throw(NegativeCycleError())
    return BellmanFordState(parents, dists)
end

bellman_ford_shortest_paths(
    graph::AbstractGraph{U},
    v::Integer,
    distmx::AbstractMatrix{T} = weights(graph);
    ) where T<:Real where U<:Integer = bellman_ford_shortest_paths(graph, [v], distmx)

has_negative_edge_cycle(g::AbstractGraph) = false

function has_negative_edge_cycle(
    g::AbstractGraph{U},
    distmx::AbstractMatrix{T}
    ) where T<:Real where U<:Integer
    try
        bellman_ford_shortest_paths(g, vertices(g), distmx)
    catch e
        isa(e, NegativeCycleError) && return true
    end
    return false
end

"""
    enumerate_paths(state[, vs])

Given a path state `state` of type `AbstractPathState`, return a
vector (indexed by vertex) of the paths between the source vertex used to
compute the path state and a single destination vertex, a list of destination
vertices, or the entire graph. For multiple destination vertices, each
path is represented by a vector of vertices on the path between the source and
the destination. Nonexistent paths will be indicated by an empty vector. For
single destinations, the path is represented by a single vector of vertices,
and will be length 0 if the path does not exist.

### Implementation Notes
For Floyd-Warshall path states, please note that the output is a bit different,
since this algorithm calculates all shortest paths for all pairs of vertices:
`enumerate_paths(state)` will return a vector (indexed by source vertex) of
vectors (indexed by destination vertex) of paths. `enumerate_paths(state, v)`
will return a vector (indexed by destination vertex) of paths from source `v`
to all other vertices. In addition, `enumerate_paths(state, v, d)` will return
a vector representing the path from vertex `v` to vertex `d`.
"""
function enumerate_paths(state::AbstractPathState, vs::Vector{T}) where T<:Integer
    parents = state.parents

    num_vs = length(vs)
    all_paths = Vector{Vector{T}}(undef, num_vs)
    for i = 1:num_vs
        all_paths[i] = Vector{T}()
        index = vs[i]
        if parents[index] != 0 || parents[index] == index
            while parents[index] != 0
                push!(all_paths[i], index)
                index = parents[index]
            end
            push!(all_paths[i], index)
            reverse!(all_paths[i])
        end
    end
    all_paths
end

enumerate_paths(state::AbstractPathState, v) = enumerate_paths(state, [v])[1]
enumerate_paths(state::AbstractPathState) = enumerate_paths(state, [1:length(state.parents);])


### Zero variance (ZV) Monte Carlo estimators
### Reference:
### Mira A, Solgi R and Imparato D
### Zero Variance Markov Chain Monte Carlo for Bayesian Estimators
### Statistics and Computing, 2013, 23 (5), pp 653-662

## Zero variance (ZV) Monte Carlo estimators linear polynomials

function lzv(chain::Vector{N}, grad::Vector{N}) where N<:Real
  z = -0.5*grad
  augmentedcov = cov([z chain])
  a = -augmentedcov[1, 2]/augmentedcov[1, 1]
  return chain+z*a, a
end

function lzv(chain::Matrix{N}, grad::Matrix{N}) where N<:Real
  npars = size(chain, 2)

  augmentedcov = Array(N, npars+1, npars+1, npars)
  precision = Array(N, npars, npars, npars)
  sigma = Array(N, npars, npars)
  a = Array(N, npars, npars)

  z = -0.5*grad

  for i in 1:npars
     augmentedcov[:, :, i] = cov([z chain[:, i]])
     precision[:, :, i] = inv(augmentedcov[1:npars, 1:npars, i])
     sigma[:, i] = augmentedcov[1:npars, npars+1, i]
     a[:, i] = -precision[:, :, i]*sigma[:, i]
  end

  return chain+z*a, a
end

lzv(s::ParameterNState{Continuous, Univariate}) = lzv(s.value, s.gradlogtarget)

lzv(s::ParameterNState{Continuous, Multivariate}) = lzv(transpose(s.value), transpose(s.gradlogtarget))

## Zero variance (ZV) Monte Carlo estimators quadratic polynomials

function qzv(chain::Vector{N}, grad::Vector{N}) where N<:Real
  z = -0.5*grad
  qz = [z 2*z.*chain-1]
  augmentedcov = cov([qz chain])
  a = -inv(augmentedcov[1:2, 1:2])*augmentedcov[1:2, 3]
  return chain+qz*a, a
end

function qzv(chain::Matrix{N}, grad::Matrix{N}) where N<:Real
  nsamples, npars = size(chain)
  k = convert(Integer, npars*(npars+3)/2)
  l = 2*npars+1

  qz = Array(N, nsamples, k)
  augmentedcov = Array(N, k+1, k+1, npars)
  precision = Array(N, k, k, npars)
  sigma = Array(N, k, npars)
  a = Array(N, k, npars)

  z = -grad/2

  qz[:, 1:npars] = z
  qz[:, (npars+1):(2*npars)] = 2*z.*chain-1
  for i in 1:(npars-1)
    for j = (i+1):npars
      qz[:, l] = chain[:, i].*z[:, j]+chain[:, j].*z[:, i]
      l += 1
    end
  end

  for i in 1:npars
    augmentedcov[:, :, i] = cov([qz chain[:, i]])
    precision[:, :, i] = inv(augmentedcov[1:k, 1:k, i])
    sigma[:, i] = augmentedcov[1:k, k+1, i]
    a[:, i] = -precision[:, :, i]*sigma[:, i]
  end

  return chain+qz*a, a
end

qzv(s::ParameterNState{Continuous, Univariate}) = qzv(s.value, s.gradlogtarget)

qzv(s::ParameterNState{Continuous, Multivariate}) = qzv(transpose(s.value), transpose(s.gradlogtarget))


include("../utils.jl")
include("group_utils.jl")

@testset "Semidirect product group" begin
    M1 = TranslationGroup(2)
    A = TranslationAction(M1, M1)
    G = SemidirectProductGroup(M1, M1, A)
    @test G === GroupManifold(
        TranslationGroup(2) × TranslationGroup(2),
        Manifolds.SemidirectProductOperation(A),
    )
    @test repr(G) == "SemidirectProductGroup($(M1), $(M1), $(A))"
    @test repr(G.op) == "SemidirectProductOperation($(A))"
    M = base_manifold(G)
    @test M === TranslationGroup(2) × TranslationGroup(2)

    ts1 = Vector{Float64}.([1:2, 2:3, 3:4])
    ts2 = Vector{Float64}.([1:2, 2:3, 3:4]) .* 10
    tuple_pts = [zip(ts1, ts2)...]
    tuple_v = ([-1.0, 1.0], [2.0, 3.0])

    shape_se = Manifolds.ShapeSpecification(Manifolds.ArrayReshaper(), M.manifolds...)
    pts = [Manifolds.prod_point(shape_se, tp...) for tp in tuple_pts]
    v_pts = [Manifolds.prod_point(shape_se, tuple_v...)]

    X = log(G, pts[1], pts[1])
    Y = zero_tangent_vector(G, pts[1])
    Z = Manifolds.allocate_result(G, zero_tangent_vector, pts[1])
    Z = zero_tangent_vector!(M, Z, pts[1])
    @test norm(G, pts[1], X) ≈ 0
    @test norm(G, pts[1], Y) ≈ 0
    @test norm(G, pts[1], Z) ≈ 0

    e = Identity(G, (zeros(2), zeros(2)))
    @test inv(G, e) === e
    @test identity(G, e) === e
    @test identity!(G, e, e) === e

    @test compose(G, e, pts[1]) == pts[1]
    @test compose(G, pts[1], e) == pts[1]
    @test compose(G, e, e) === e
end


@testset "Triangle" begin
    # Make some points between (0,0,0) and (1,1,1) and put in a point cloud
    N = 1000
    pos = rand(SVector{3,Float32}, N)
    triangle = Triangle(SVector(0.0, 0.0), SVector(1.0, 0.0), SVector(0.0, 1.0))

    search_correct = true
    for i = 1:length(pos)
        v = pos[i]
        expected = v[1] + v[2] < 1.0
        actual = v ∈ triangle
        search_correct &= actual == expected
    end
    @test search_correct

    triangle = Triangle(SVector(0.5, 0.5), SVector(0.5, 1.5), SVector(1.5, 1.0))
    @test SVector(0.0, 0.0, 0.0) ∉ triangle
    @test SVector(0.0, 0.0, 1.0) ∉ triangle
    @test SVector(0.0, 0.0, -1.0) ∉ triangle
    @test SVector(1.0, 1.0, 0.0) ∈ triangle
    @test SVector(1.0, 1.0, -1.0) ∈ triangle
    @test SVector(1.0, 1.0, 1.0) ∈ triangle

    @test boundingbox(triangle) == BoundingBox(0.5, 0.5, -Inf, 1.5, 1.5, Inf)

    @test area(triangle) == 0.5
end


#=
A type to solve the career / job choice model due to Derek Neal.

@author : Spencer Lyon <spencer.lyon@nyu.edu>

@date: 2014-08-05

References
----------

http://quant-econ.net/jl/career.html

[Neal1999] Neal, D. (1999). The Complexity of Job Mobility among Young Men,
Journal of Labor Economics, 17(2), 237-261.
=#

"""
Career/job choice model fo Derek Neal (1999)

##### Fields

- `beta::Real` : Discount factor in (0, 1)
- `N::Int` : Number of possible realizations of both epsilon and theta
- `B::Real` : upper bound for both epsilon and theta
- `theta::AbstractVector` : A grid of values on [0, B]
- `epsilon::AbstractVector` : A grid of values on [0, B]
- `F_probs::AbstractVector` : The pdf of each value associated with of F
- `G_probs::AbstractVector` : The pdf of each value associated with of G
- `F_mean::Real` : The mean of the distribution F
- `G_mean::Real` : The mean of the distribution G

"""
type CareerWorkerProblem
    beta::Real
    N::Int
    B::Real
    theta::AbstractVector
    epsilon::AbstractVector
    F_probs::AbstractVector
    G_probs::AbstractVector
    F_mean::Real
    G_mean::Real
end

"""
Constructor with default values for `CareerWorkerProblem`

##### Arguments

- `beta::Real(0.95)` : Discount factor in (0, 1)
- `B::Real(5.0)` : upper bound for both epsilon and theta
- `N::Real(50)` : Number of possible realizations of both epsilon and theta
- `F_a::Real(1), F_b::Real(1)` : Parameters of the distribution F
- `G_a::Real(1), G_b::Real(1)` : Parameters of the distribution F

##### Notes

There is also a version of this function that accepts keyword arguments for
each parameter
"""
# use key word argument
function CareerWorkerProblem(;beta::Real=0.95, B::Real=5.0, N::Real=50,
                             F_a::Real=1, F_b::Real=1, G_a::Real=1,
                             G_b::Real=1)
    theta = linspace(0, B, N)
    epsilon = copy(theta)
    F_probs::Vector{Float64} = pdf(BetaBinomial(N-1, F_a, F_b))
    G_probs::Vector{Float64} = pdf(BetaBinomial(N-1, G_a, G_b))
    F_mean = sum(theta .* F_probs)
    G_mean = sum(epsilon .* G_probs)
    CareerWorkerProblem(beta, N, B, theta, epsilon, F_probs, G_probs,
                        F_mean, G_mean)
end

"""
Apply the Bellman operator for a given model and initial value.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output
- `;ret_policy::Bool(false)`: Toggles return of value or policy functions

##### Returns

None, `out` is updated in place. If `ret_policy == true` out is filled with the
policy function, otherwise the value function is stored in `out`.

"""
function update_bellman!(cp::CareerWorkerProblem, v::Array, out::Array;
                           ret_policy=false)
    # new life. This is a function of the distribution parameters and is
    # always constant. No need to recompute it in the loop
    v3 = (cp.G_mean + cp.F_mean + cp.beta .*
          cp.F_probs' * v * cp.G_probs)[1]  # don't need 1 element array

    for j=1:cp.N
        for i=1:cp.N
            # stay put
            v1 = cp.theta[i] + cp.epsilon[j] + cp.beta * v[i, j]

            # new job
            v2 = (cp.theta[i] .+ cp.G_mean .+ cp.beta .*
                  v[i, :]'*cp.G_probs)[1]  # don't need a single element array

            if ret_policy
                if v1 > max(v2, v3)
                    action = 1
                elseif v2 > max(v1, v3)
                    action = 2
                else
                    action = 3
                end
                out[i, j] = action
            else
                out[i, j] = max(v1, v2, v3)
            end
        end
    end
end

function update_bellman(cp::CareerWorkerProblem, v::Array; ret_policy=false)
    out = similar(v)
    update_bellman!(cp, v, out, ret_policy=ret_policy)
    return out
end

"""
Extract the greedy policy (policy function) of the model.

##### Arguments

- `cp::CareerWorkerProblem` : Instance of `CareerWorkerProblem`
- `v::Matrix`: Current guess for the value function
- `out::Matrix` : Storage for output

##### Returns

None, `out` is updated in place to hold the policy function

"""
function get_greedy!(cp::CareerWorkerProblem, v::Array, out::Array)
    update_bellman!(cp, v, out, ret_policy=true)
end

function get_greedy(cp::CareerWorkerProblem, v::Array)
    update_bellman(cp, v, ret_policy=true)
end


using Statistics
using NCDatasets

using ArgParse
using JSON

include("LinearRegression.jl")

correlation = (x1, x2) -> x1' * x2 / (sum(x1.^2)*sum(x2.^2)).^0.5

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--data-file"
            help = "Atm data file. New variable will be appended."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

        "--output-file"
            help = "Output file."
            arg_type = String
            required = true

        "--beg-year"
            help = "Year of begin."
            arg_type = Int64
            required = true

        "--end-year"
            help = "Year of end."
            arg_type = Int64
            required = true

    end

    return parse_args(ARGS, s)
end

parsed = parse_commandline()
print(json(parsed, 4))

output_file = parsed["output-file"]

Dataset(parsed["data-file"], "r") do ds

    beg_t = (parsed["beg-year"] - 1) * 12 + 1
    end_t = (parsed["end-year"] - 1) * 12 + 12
    rng = (:,:,beg_t:end_t)

    global total_precip = replace(ds["PRECC"][rng...] + ds["PRECL"][rng...], missing=>NaN)
    global (Nx, Ny, Nt) = size(total_precip)

    if mod(Nt, 12) != 0
        ErrorException("Time record is not multiple of 12") |> throw
    end

    global nyears = Int64(Nt / 12)
end

Dataset(parsed["domain-file"], "r") do ds
    global mask = replace(ds["mask"], missing=>NaN)
end

# M   = Mean
# VAR = Variance

# MM   = Monthly Mean
# MVAR = Monthly Variance

# ZM   = Zonal Mean
# ZVAR = Zonal Variance

# ZMM   = Zonal Monthly Mean
# ZMVAR = Zonal Monthly Variance

total_precip_M     = zeros(Float64, Nx, Ny)
total_precip_VAR   = zeros(Float64, Nx, Ny)

total_precip_MM     = zeros(Float64, Nx, Ny, 12)
total_precip_MVAR   = zeros(Float64, Nx, Ny, 12)

total_precip_ZM     = zeros(Float64, Ny)
total_precip_ZVAR   = zeros(Float64, Ny)

total_precip_ZMM    = zeros(Float64, Ny, 12)
total_precip_ZMVAR  = zeros(Float64, Ny, 12)

x = collect(Float64, 1:Nt)
for i=1:Nx, j=1:Ny

    d = view(total_precip, i, j, :)

    total_precip_M[i, j]   = mean(d)
    total_precip_VAR[i, j] = std(d)

    d_wrap = reshape(d, 12, :)
    total_precip_MM[i, j, :]   = mean( d_wrap, dims=2 )[:, 1]
    total_precip_MVAR[i, j, :] = std( d_wrap, dims=2 )[:, 1]

end

for j=1:Ny

    d = mean(view(total_precip, :, j, :), dims=1)[1, :]

    total_precip_ZM[j]   = mean(d)
    total_precip_ZVAR[j] = std(d)

    d_wrap = reshape(d, 12, :)
    total_precip_ZMM[j, :]   = mean( d_wrap, dims=2 )[:, 1]
    total_precip_ZMVAR[j, :] = std(  d_wrap, dims=2 )[:, 1]

end

Dataset(output_file, "c") do ds

    defDim(ds, "months", 12)
    defDim(ds, "time", Inf)
    defDim(ds, "Nx", Nx)
    defDim(ds, "Ny", Ny)

    for (varname, vardata, vardim, attrib) in [
        ("total_precip_M",     total_precip_M,     ("Nx", "Ny",), Dict()),
        ("total_precip_VAR",   total_precip_VAR,   ("Nx", "Ny",), Dict()),
        ("total_precip_MM",    total_precip_MM,    ("Nx", "Ny", "months"), Dict()),
        ("total_precip_MVAR",  total_precip_MVAR,  ("Nx", "Ny", "months"), Dict()),
        ("total_precip_ZM",    total_precip_ZM,    ("Ny",), Dict()),
        ("total_precip_ZVAR",  total_precip_ZVAR,  ("Ny",), Dict()),
        ("total_precip_ZMM",   total_precip_ZMM,   ("Ny", "months"), Dict()),
        ("total_precip_ZMVAR", total_precip_ZMVAR, ("Ny", "months"), Dict()),
    ]

        if ! haskey(ds, varname)
            var = defVar(ds, varname, Float64, vardim)
            var.attrib["_FillValue"] = 1e20
        end

        var = ds[varname]

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end
end


push!(LOAD_PATH, "../LF_common/")
using m_LF3d
include("../LF_common/sparse_LF3d.jl")
include("../LF_common/prec_mkl_ilu0.jl")
include("../LF_common/apply_prec_ilu0.jl")
include("build_nabla2_x.jl")
include("build_nabla2_y.jl")
include("build_nabla2_z.jl")

const ⊗ = kron

function test_main( NN::Array{Int64} )

  println("NN = ", NN)

  Nx = NN[1]
  Ny = NN[2]
  Nz = NN[3]
  Npoints = Nx*Ny*Nz

  AA = [0.0, 0.0, 0.0]
  BB = [16.0, 16.0, 16.0]

  # Initialize LF
  LF = init_LF3d_c( NN, AA, BB, verbose=true )

  nabla2_x = build_nabla2_x(LF)
  nabla2_y = build_nabla2_y(LF)
  nabla2_z = build_nabla2_z(LF)

  # Using Kronecker product
  ∇2 = get_Laplacian3d_kron(LF)

  #
  srand(123)
  v = rand(Npoints)

  # Test multiplication operation
  #@time Lv_1 = nabla2_x*v + nabla2_y*v + nabla2_z*v
  #@time Lv_2 = ∇2*v

  #@printf("Difference: %18.10e\n", sum(abs(Lv_1 - Lv_2))/Npoints)

  # Test the preconditioner

  prec_x = prec_mkl_ilu0( nabla2_x )
  prec_y = prec_mkl_ilu0( nabla2_y )
  prec_z = prec_mkl_ilu0( nabla2_z )

  pvx = apply_prec_ilu0( prec_x, v )
  pvy = apply_prec_ilu0( prec_y, v )
  pvz = apply_prec_ilu0( prec_z, v )
  Pv_1 = pvx + pvy + pvz
  #
  prec_∇2 = prec_mkl_ilu0( ∇2 )
  Pv_2 = apply_prec_ilu0(prec_∇2, v)

  println(sum(v))

  println(sum(pvx))
  println(sum(pvy))
  println(sum(pvz))

  println(sum(Pv_1))
  println(sum(Pv_2))

end

test_main( [3, 3, 3])


using SparseArrays, Graphs

A = sprand(Bool, 20, 20, 0.5)
for i in 1:20; A[i,i] = 0; end
A = A + A'
dropzeros!(A)
g = IndexedGraph(A)

@testset "undirected graph" begin
    @testset "basics" begin
        @test !is_directed(g)
        @test length(collect(edges(g))) == ne(g)
        i = 3
        es = inedges(g, i)
        neigs = neighbors(g, i)
        @test all(
            ((src(e), dst(e)) == (i, j)) || ((src(e), dst(e)) == (j, i))
            for (e, j) in zip(es, neigs)
        )
    end

    @testset "edge indexing" begin
        @test all( e == get_edge(g, src(e), dst(e)) for e in edges(g) )
        @test all( e == get_edge(g, idx(e)) for e in edges(g) )
    end

    @testset "construct from SimpleGraph" begin
        sg = SimpleGraph(A)
        ig = IndexedGraph(sg)
        @test adjacency_matrix(sg) == adjacency_matrix(ig)
    end
end


Pkg.add(["TimeSeries", "MarketData", "Plots", "Pipe"])
using Pipe: @pipe
using TimeSeries, MarketData, Plots
# include("src/Portfolio.jl")
using Portfolio

function strategy_tearsheet(weights::Weights, returns::Returns)
    strategy_returns = run_weights(returns, weights, 0.0)
    equity_curves = plot(TimeSeries.merge(equity_curve(returns), equity_curve(strategy_returns)))
    plot(
        equity_curves,
        plot(weights),
        size = (1200, 1200),
        layout = grid(2, 1, heights = [0.8, 0.2])
    )
end

@time dell = DELL[:AdjClose] |> percentchange

# Target mean volatility, 30 day lookback
weight_vol_target_mean = @pipe strategy_vol_target_mean(dell, 30) |> lowpass(1.0).(_)
strategy_tearsheet(weight_vol_target_mean, dell)

# Ensemble lookback 20:40, still target the mean volatility
ensembled_weight_vol_target = @pipe ensemble_parameters(strategy_vol_target_mean, dell, 20:40) |> lowpass(1.0).(_)
strategy_tearsheet(ensembled_weight_vol_target, dell)

# If 200 day momentum is positive, 100% long, otherwise zero allocation
weight_momentum_above_0 = strategy_positive_mom_binary(dell, 200)
strategy_tearsheet(weight_momentum_above_0, dell)

# Ensemble 50:200 day momentum
ensembled_weight_momentum_above_0 = ensemble_parameters(strategy_positive_mom_binary, dell, 50:200)
strategy_tearsheet(ensembled_weight_momentum_above_0, dell)

# Merge the ensembled mom / vol targeting weights
all_weights_ensembled = merge_weights([ensembled_weight_vol_target, ensembled_weight_momentum_above_0])
strategy_tearsheet(all_weights_ensembled, dell)


#Test library imports
using Test

#Arithmetic library imports
using IntervalArithmetic
using IntervalContractors

#Preamble
setprecision(53)
setprecision(Interval, Float64)

# Using approximate checks for validaty update later?
import Base.isapprox
isapprox(x::Interval,y::Interval) = isapprox(x.lo,y.lo,atol=1E-4) && isapprox(x.hi,y.hi,atol=1E-4)

@testset "sinh_rev_test" begin
    @test isapprox(sinh_rev(∅, -∞..∞)[2], ∅)
    @test isapprox(sinh_rev(Interval(-10.0, -1.0), -∞..∞)[2], Interval(-2.99823, -0.881373))
    @test isapprox(sinh_rev(Interval(0.0, Inf), -∞..∞)[2], Interval(0.0, ∞))
    @test isapprox(sinh_rev(Interval(0.0, 1.0), -∞..∞)[2], Interval(0, 0.881374))
    @test isapprox(sinh_rev(Interval(-0.5, 1.0), -∞..∞)[2], Interval(-0.481212, 0.881374))
    @test isapprox(sinh_rev(Interval(-1000.0, 1.0), -∞..∞)[2], Interval(-7.60091, 0.881374))
    @test isapprox(sinh_rev(Interval(0.0, 25.0), -∞..∞)[2], Interval(0.0, 3.91243))
    @test isapprox(sinh_rev(Interval(-1.0, 25.0), -∞..∞)[2], Interval(-0.881374, 3.91243))
end

@testset "cosh_rev_test" begin
    @test isapprox(cosh_rev(∅, -∞..∞)[2], ∅)
    @test isapprox(cosh_rev(Interval(-10.0, -1.0), -∞..∞)[2], ∅)
    @test isapprox(cosh_rev(Interval(0.0, Inf), -∞..∞)[2], Interval(0.0, ∞))
    @test isapprox(cosh_rev(Interval(0.0, 1.0), -∞..∞)[2], Interval(0, 0))
    @test isapprox(cosh_rev(Interval(-0.5, 1.0), -∞..∞)[2], Interval(0, 0))
    @test isapprox(cosh_rev(Interval(-1000.0, 1.0), -∞..∞)[2], Interval(0, 0))
    @test isapprox(cosh_rev(Interval(0.0, 25.0), -∞..∞)[2], Interval(0, 3.91163))
    @test isapprox(cosh_rev(Interval(-1.0, 25.0), -∞..∞)[2], Interval(0, 3.91163))
end

@testset "tanh_rev_test" begin
    @test tanh_rev(∅, -∞..∞)[2] == ∅
    @test tanh_rev(Interval(-10.0, -1.0), -∞..∞)[2] == ∅
    @test isapprox(tanh_rev(Interval(0.0, Inf), -∞..∞)[2], Interval(0.0, ∞))
    @test isapprox(tanh_rev(Interval(0.0, 1.0), -∞..∞)[2], Interval(0.0, ∞))
    @test isapprox(tanh_rev(Interval(-0.5, 1.0), -∞..∞)[2], Interval(-0.549307, ∞))
    @test isapprox(tanh_rev(Interval(-1000.0, 1.0), -∞..∞)[2], Interval(-∞, ∞))
    @test isapprox(tanh_rev(Interval(0.0, 25.0), -∞..∞)[2], Interval(0, ∞))
    @test isapprox(tanh_rev(Interval(-1.0, 25.0), -∞..∞)[2], Interval(-∞, ∞))
end


#=

In this file we simulate a simple rocket landing

Units in kg, m, s

=#

# RUN:
# ] activate TrajOptSOCPs

using TrajOptSOCPs
using LinearAlgebra, SparseArrays

# If you don't want plotting comment the include and pyplot statements.
# Then, sent runplots = false
include("../plotting/plotTrajectory.jl")
include("../plotting/plotConstraintViolation.jl")
include("../plotting/plotObjective.jl")
include("../plotting/batchPlots.jl")
pyplot()
runplots = true
saveplots = false

println("\n--------------------------------------------")
println("          Setting Up Rocket Landing ")
println("--------------------------------------------")

# Based on the Falcon 9
# 549,054 kg (Mass)
# 282 s (Specific Impulse)
# select "y" as the vertical
mass = 1
isp = 1
grav = [0; -9.81]
deltaTime = 1
rocket = rocket_simple(mass, isp, grav, deltaTime)

# in m
# The Karman Line (100 km)
const rocketStart = [2.0; 20.0; 0.0; -18.0]
const rocketEnd = [0.0; 0.0; 0.0; 0.0]#[-5.0; 0.0; 0.0; 0.0]

uHover = mass * grav

# Number of time steps to discretize the trajectory
const NSteps = 60
# Initialize the trajectory with a line
initTraj = initializeTraj(rocketStart, rocketEnd, uHover, uHover, NSteps)

# Use a Linear Quadratic Regulator as the cost function
const lqrQMat = 0.0001 * Diagonal(I, size(rocketStart, 1))
const lqrRMat = 0.0025 * Diagonal(I, Int64(size(rocketStart, 1) / 2))
costFun = makeLQR_TrajReferenced(lqrQMat, lqrRMat, NSteps, initTraj)

# Create the Dynamics Constraints
const ADyn, BDyn = rocketDynamicsFull(rocket, rocketStart, rocketEnd, NSteps)
dynConstraint = AL_AffineEquality(ADyn, BDyn)
lambdaInit = -1 * ones(size(BDyn))

# Create the Ground Constraints
groundConstraint = makeGroundConstraint(NSteps, size(grav, 1), size(grav, 1))
groundLambda = zeros(size(groundConstraint.b, 1))

# Create the Max Thrust Constraints
thrustMax = 20.0
maxThrustConstraint = makeMaxThrustConstraint(NSteps, size(grav, 1), thrustMax)
maxThrustLambda = zeros(size(maxThrustConstraint.indicatorList))

# Create the constraint manager

# cMRocket = constraintManager_Base([dynConstraint], [lambdaInit])
# cMRocket = constraintManager_Dynamics([], [], dynConstraint, lambdaInit)
# cMRocket = constraintManager_Dynamics([groundConstraint], [groundLambda],
#                                         dynConstraint, lambdaInit)
cMRocket = constraintManager_Dynamics([groundConstraint, maxThrustConstraint],
                                      [groundLambda, maxThrustLambda],
                                      dynConstraint, lambdaInit)

# Initialize the primal-dual vector
initTrajPD = [initTraj; lambdaInit]

# Initialize the Augmented Lagrangian penalty on the constraints
penaltyStart = 1.0

# Test that the evaluations work
println("\n--------------------------------------------")
println(" Testing constraints are inputted correctly ")
println("--------------------------------------------")
println("Starting constraint violation: ")
println([evalConstraints(cMRocket, initTraj, penaltyStart)])
println("Starting gradient of constraints: ")
println(size(evalGradConstraints(cMRocket, initTraj, penaltyStart)))
println("Starting hessian of constraints: ")
println(size(evalHessConstraints(cMRocket, initTraj)))

# Equiped with the constraint term and the objective term, I now build the
# Augmented Lagrangian
alRocket = augLag(costFun, cMRocket, penaltyStart)

# Test that the evaluations work
println("--------------------------------------------")
println("           Testing AL is Functional         ")
println("--------------------------------------------")
println("Evaluating augmented lagrangian: ")
println([evalAL(alRocket, initTrajPD)])
println("Evaluating gradient of augmented lagrangian: ")
println(size(evalGradAL(alRocket, initTrajPD)))
println("Evaluating hessian of augmented lagrangian: ")
println(size(evalHessAl(alRocket, initTrajPD)))

# Next we select resonable solver parameters
currSolveParams = solverParams(0.1, 0.5,
                                8, 10,
                                10^-4,
                                10, 10^6,
                                2.5, 2, 0.2, 0.2, 0.4)
solParamPrint(currSolveParams)
println()
println()

println("--------------------------------------------")
println("             Beginning Solve                ")
println("--------------------------------------------")

# Solve the Trajectory Optimization problem
useALMethod = true

if useALMethod
    # Use an augmented lagrangian method
    trajLambdaSolved, resArr = ALPrimalNewtonMain(initTrajPD, alRocket,
                                            currSolveParams)
    trajStatesAllPD = [parsePrimalDualVec(trajL, size(initTraj, 1))
                                                for trajL in trajLambdaSolved]
    trajStatesAll = [trajPDThis.primals for trajPDThis in trajStatesAllPD]
else
    # Solve only the objective and linear dynamics
    costQ = costFun.lqr_qp.QR_Full
    costP = - (initTraj' * costQ)'
    trajLambdaSolvedTrue = solveQP_AffineEq(costQ, costP, ADyn, BDyn)
    trajStateLastTrue = parsePrimalDualVec(trajLambdaSolvedTrue,
                                                        size(initTraj, 1))
end

# Blocked so that it can be run independently after the fact
if runplots

    hDual = heatmap(hcat(trajStatesAllPD[end].duals))
    yflip!()
    ylabel!("Column")
    title!("Dual Vector for Dynamics Constraints")
    display(hDual)

    # Get the parsed list of trajectories
    nDim = size(grav, 1)
    pltTraj, pltCV, pltCV2, pltObj, plts, pltv, pltu =
                                    batchPlot(trajStatesAll, cMRocket, nDim)
end

# Blocked so that it can be run independently after the fact
if runplots && saveplots
    header = "MaxThrustNew_" * string(currSolveParams.maxOuterIters) *
             "Outer_" * string(currSolveParams.maxNewtonSteps) * "Newton" *
             string(Int64(rocketStart[4])) * "Vel" * "_"
    saveBulk(pltTraj, pltCV, pltCV2, pltObj, plts, pltv, pltu, header)
end


@testset "Broadcast" begin
    for T in (Float32, Float64)
        v = rand(Vec{3, T})
        A = rand(Mat{2, 3, T})
        vs = [rand(Vec{3, T}) for _ in 1:4]
        As = [rand(Mat{2, 3, T}) for _ in 1:4]
        # with array (tensors behave like scalar)
        @test (@inferred(v .+ vs))::Vector{Vec{3, T}} ≈ map(y -> v + y, vs)
        @test (@inferred(v .- vs))::Vector{Vec{3, T}} ≈ map(y -> v - y, vs)
        @test (@inferred(v .+ vs .- vs .+ v))::Vector{Vec{3, T}} ≈ map(y -> v + y - y + v, vs)
        @test (@inferred(broadcast(+, v, vs, vs, v)))::Vector{Vec{3, T}} ≈ map(y -> v + y + y + v, vs)
        @test (vs .= v)::Vector{Vec{3, T}} ≈ map(y -> v, vs)
        @test (Array(v) .= v)::Vector{T} ≈ v
        @test (@inferred(A .+ As))::Vector{Mat{2, 3, T, 6}} ≈ map(y -> A + y, As)
        @test (@inferred(A .- As))::Vector{Mat{2, 3, T, 6}} ≈ map(y -> A - y, As)
        @test (@inferred(A .+ As .- As .+ A))::Vector{Mat{2, 3, T, 6}} ≈ map(y -> A + y - y + A, As)
        @test (@inferred(broadcast(+, A, As, As, A)))::Vector{Mat{2, 3, T, 6}} ≈ map(y -> A + y + y + A, As)
        @test (As .= A)::Vector{Mat{2, 3, T, 6}} ≈ map(y -> A, As)
        @test (Array(A) .= A)::Matrix{T} ≈ A
        # with scalar
        @test (@inferred(v .+ 1))::Vec{3, T} ≈ map(y -> y + 1, v)
        @test (@inferred(A .+ 1))::Mat{2, 3, T} ≈ map(y -> y + 1, A)
        @test (@inferred(v .+ 1 .+ 2 .+ v))::Vec{3, T} ≈ map(y -> y + 1 + 2 + y, v)
        @test (@inferred(A .+ 1 .+ 2 .+ A))::Mat{2, 3, T} ≈ map(y -> y + 1 + 2 + y, A)
        # with tuple
        @test (@inferred(v .+ (1,2,3)))::Vec{3, T} ≈ map(+, v, (1,2,3))
        @test (@inferred(v .- (1,2,3)))::Vec{3, T} ≈ map(-, v, (1,2,3))
        # others
        @test (@inferred(broadcast(sqrt, v)))::Vec{3, T} ≈ map(sqrt, v)
        @test (@inferred(broadcast(sqrt, A)))::Mat{2, 3, T} ≈ map(sqrt, A)
    end
end


"Returns the time-scale for advection on a regular grid across a single grid cell
 for ShallowWaterModel."

import Oceananigans.Utils: cell_advection_timescale

function shallow_water_cell_advection_timescale(uh, vh, h, grid)
    uhmax = maximum(abs, uh)
    vhmax = maximum(abs, vh)
    hmin  = minimum(abs,  h)

    return min(grid.Δx / uhmax, grid.Δy / vhmax) * hmin
end

cell_advection_timescale(model::ShallowWaterModel) =
    shallow_water_cell_advection_timescale(
        model.solution.uh.data.parent,
        model.solution.vh.data.parent,
        model.solution.h.data.parent,
        model.grid
        )


import Dolo

model = Dolo.yaml_import("rbc_iid.yaml")

# define
include("tmp_module.jl")
import temp

# get grid for endogenous
gg = model.options.grid
gg.a
gg.orders             # this is imported from the yaml file
grid = CartesianGrid(gg.a, gg.b, gg.orders) # temporary compatibility

# now look at exogenous
# everything below should work for any type of process (normal, markov chain, ar1, ...)

ptype = "MC"
if ptype == "IID"
    process = MvNormal(0.01)
elseif ptype == "MC"
    values = [ -0.01  0.01 ]'
    transitions = [
        0.9 0.1;
        0.1 0.9
    ]
    process = DiscreteMarkovProcess(transitions, values)
end

dprocess = discretize(process)

model.calibration.flat[:beta]

model.calibration[:parameters]
#beta=model.calibration.flat[:beta]
β = 0.5

drv_answer = value_iteration(model, dprocess, process, grid, β, true, 300)

println(drv_answer)
using Gadfly
kvec = linspace(8,12,100)

drv_answer.grid
yvec = [drv_answer(1,[k])[1] for k in kvec]
plot(x=kvec, y=yvec, Geom.line)


using LinearAlgebra, MixedModels, StableRNGs, Test, SparseArrays

include("modelcache.jl")

@testset "femat" begin
    trm = MixedModels.FeMat(hcat(ones(30), repeat(0:9, outer = 3)), ["(Intercept)", "U"])
    piv = trm.piv
    ipiv = invperm(piv)
    @test size(trm) == (30, 2)
    @test length(trm) == 60
    @test size(trm') == (2, 30)
    @test eltype(trm) == Float64
    @test trm.x === trm.wtx
    prd = trm'trm
    @test typeof(prd) == Matrix{Float64}
    @test prd == [30.0 135.0; 135.0 855.0][piv, piv]
    wts = rand(StableRNG(123454321), 30)
    MixedModels.reweight!(trm, wts)
    @test mul!(prd, trm', trm)[ipiv[1], ipiv[1]] ≈ sum(abs2, wts)

    # empty fixed effects
    trm = MixedModels.FeMat(ones(10,0), String[])
    @test size(trm) == (10, 0)
    @test length(trm) == 0
    @test size(trm') == (0, 10)
    @test eltype(trm) == Float64
    @test trm.rank == 0
end

@testset "fematSparse" begin

    @testset "sparse and dense yield same fit" begin
        # deepcopy because we're going to modify
        m = deepcopy(last(models(:insteval)))
        # this is kinda sparse:
        # julia> mean(first(m.feterms).x)
        # 0.10040140325753434

        fe = first(m.feterms)
        X =  MixedModels.FeMat(SparseMatrixCSC(fe.x), fe.cnames)
        @test typeof(X.x) <: SparseMatrixCSC
        @test typeof(X.wtx) <: SparseMatrixCSC
        @test X.rank == 28
        @test X.cnames == fe.cnames

        dense_θ = copy(m.θ)
        dense_feval = m.optsum.feval
        m.feterms[1] = X
        refit!(m)

        @test dense_θ ≈ m.θ
    end

    @testset "rank defiency in sparse FeMat" begin
        trm = MixedModels.FeMat(SparseMatrixCSC(hcat(ones(30),
                                                     repeat(0:9, outer = 3),
                                                     2repeat(0:9, outer = 3))),
                                ["(Intercept)", "U", "V"])
        piv = trm.piv
        ipiv = invperm(piv)
        @test_broken rank(trm) == 2
        @test size(trm) == (30, 3)
        @test length(trm) == 90
        @test size(trm') == (3, 30)
        @test eltype(trm) == Float64
        @test trm.x === trm.wtx
        prd = trm'trm
        @test typeof(prd) == typeof(prd)
        @test prd == [30.0 135.0 270.0; 135.0 855.0 1710.0; 270.0 1710.0 3420][piv, piv]
        wts = rand(StableRNG(123454321), 30)
        MixedModels.reweight!(trm, wts)
        @test mul!(prd, trm', trm)[ipiv[1], ipiv[1]] ≈ sum(abs2, wts)
    end

end


#Written By Johnathan Bizzano
module Error
    using ..Functions
    using ..Utils
    using ..Data

    @Fun(Func, Float64, AbstractDataSet, fun::Any)
    @Fun(WrappedFunc, Float64)

    export Meanabserr, Meansqrerr, Meanerr, None

    diff(fun, entry) = fun(entry.inputs)[1] - entry.outputs[1]

    None() = Func((dataset, fun) -> 0.0)

    function Meanabserr()
        return Func(function (dataset, fun)
                sum = 0.0
                for item in dataset
                    sum += abs(diff(fun, item))
                end
                sum / length(dataset)
            end)
    end

    function Meansqrerr()
        return Func(function (dataset, fun)
                sum = 0.0
                for item in dataset
                    sum += diff(fun, item) ^ 2
                end
                sqrt(sum / length(dataset))
            end)
    end

    function Meanerr()
        return Func(function (dataset, fun)
                sum = 0.0
                for item in dataset
                    sum += diff(fun, item)
                end
                sum / length(dataset)
            end)
    end
end


"""
    plot(path::RegularizationPath, args...; <keyword arguments>)

Plots a `RegularizationPath` fitted with the `Lasso` package.

The minimum AICc segment is represented by a solid vertical line and the CVmin
and CV1se cross-validation selected segments by dashed vertical lines.

By default it shows nonzero coefficients at the AICc in color and the rest grayed out.

LassoPlot uses Plots.jl, so you can set any supported backend before plotting,
and add any features of the plot afterwards.

# Example:
```julia
    using Lasso, LassoPath, Plots
    path = fit(LassoPath, X, y, dist, link)
    plot(path)
```
# Arguments
- `args...` additional arguments passed to Plots.plot()

# Keywords
- `x=:segment` one of (:segment, :λ, :logλ)
- `varnames=nothing` specify variable names
- `select=:AICc` Selection criteria in (:AIC, :CVmin, :CV1se) for which coefficients
    will be shown in color. The rest are grayed out.
- `showselectors=[:AICc,:CVmin,:CV1se]` shown vertical lines
- `selectedvars=[]` Subset of the variables to present, or empty vector for all
- `kwargs...` additional keyword arguments passed along to fit(GammaLassoPath,...)
"""
function Plots.plot(path::RegularizationPath, args...;
    x=:segment, varnames=nothing, selectedvars=[], select=:AICc, showselectors=[:AICc,:CVmin,:CV1se], nCVfolds=10)
    β=coef(path)
    if hasintercept(path)
        β = β[2:end,:]
    end

    (p,nλ)=size(β)

    if varnames==nothing
        varnames=[Symbol("x$i") for i=1:p]
    end

    indata=DataFrame()
    if x==:λ
        indata[x]=path.λ
    elseif x==:logλ
        indata[x]=log.(path.λ)
    else
        x=:segment
        indata[x]=1:nλ
    end
    outdata = deepcopy(indata)

    # automatic selectors
    # xintercept = Float64[]
    dashed_vlines=Float64[]
    solid_vlines=Float64[]

    if select == :AICc || :AICc in showselectors
        minAICcix=minAICc(path)
        if select == :AICc
            push!(solid_vlines,indata[minAICcix,x])
        else
            push!(dashed_vlines,indata[minAICcix,x])
        end
    end

    if select == :CVmin || :CVmin in showselectors
        gen = Kfold(length(path.m.rr.y),nCVfolds)
        segCVmin = cross_validate_path(path;gen=gen,select=:CVmin)
        if select == :CVmin
            push!(solid_vlines,indata[segCVmin,x])
        else
            push!(dashed_vlines,indata[segCVmin,x])
        end
    end

    if select == :CV1se || :CV1se in showselectors
        gen = Kfold(length(path.m.rr.y),nCVfolds)
        segCV1se = cross_validate_path(path;gen=gen,select=:CV1se)
        if select == :CV1se
            push!(solid_vlines,indata[segCV1se,x])
        else
            push!(dashed_vlines,indata[segCV1se,x])
        end
    end

    if length(selectedvars) == 0
        if select == :all
            selectedvars = 1:p
        elseif select == :AICc
            selectedvars = findall(!iszero, β[:,minAICcix])
        elseif select == :CVmin
            selectedvars = findall(!iszero, β[:,segCVmin])
        elseif select == :CV1se
            selectedvars = findall(!iszero, β[:,segCV1se])
        else
            error("unknown selector $select")
        end
    end

    # colored paths
    for j in selectedvars
        indata[varnames[j]]=Vector(β[j,:])
    end

    # grayed out paths
    for j in setdiff(1:p,selectedvars)
        outdata[varnames[j]]=Vector(β[j,:])
    end

    inmdframe=melt(indata,x)
    outmdframe=melt(outdata,x)
    rename!(inmdframe,:value=>:coefficients)
    rename!(outmdframe,:value=>:coefficients)
    inmdframe = inmdframe[convert(BitArray,map(b->!isnan(b),inmdframe[:coefficients])),:]
    outmdframe = outmdframe[convert(BitArray,map(b->!isnan(b),outmdframe[:coefficients])),:]

    p = plot(xlabel=string(x), ylabel="Coefficient", args...)
    if size(inmdframe,1) > 0
      @df inmdframe plot!(cols(x), :coefficients, group=:variable)
    end
    if size(outmdframe,1) > 0
      @df outmdframe plot!(cols(x), :coefficients, group=:variable, palette=:grays)
    end
    if length(dashed_vlines) > 0
        vline!(dashed_vlines, line = (:dash, 0.5, 2, :black), label="")
    end
    if length(solid_vlines) > 0
        vline!(solid_vlines, line = (:solid, 0.5, 2, :black), label="")
    end

    p
end