# simplified model from Chakrawal21

function chak21_phys_system()
    #using ModelingToolkit, DifferentialEquations, Plots
    @parameters t ks km kd Y HS HB kmr ksm
    # t in hours
    # masses per soilmass in mol/g
    @variables s(t) b(t) cr_tot(t) r_tot(t) q(t) dec_s(t) tvr_b(t)
    @variables mm_s(t) u1(t) u2(t) r(t) r_gr(t) r_m(t) db(t)
    D = Differential(t)
    HG = HS - Y * HB # can take out of system because does not involve t
    eqs = [
        D(s) ~ -dec_s + tvr_b,
        db ~ Y * u1 - tvr_b, # store output
        D(b) ~ db,
        D(cr_tot) ~ r_tot,
        D(r) ~ Y * u1/b * (s/(s+kmr) - r),
        u1 ~ b * ks * r * s/(s + km),
        u2 ~ b * ksm * (1-r) * s/(s + km),
        dec_s ~ u1 + u2,
        tvr_b ~ kd*b,
        r_gr ~ (1-Y) * u1,
        r_m ~ u2,
        r_tot ~ r_gr + r_m,
        q ~ -HG * u1 - HS * u2,
        ]
    @named chak21_phys = ODESystem(eqs)
    chak21_phys_s = structural_simplify(chak21_phys)
    p_straw = Dict(
        :parms => [ks => 0.15, km => 2.62e-6, kd => 1.35e-2, Y => 0.72,
            HB => -492, HS => -469,
            kmr => 1e-9, # half-saturation low so that optimum r -> 1
            ksm => 0.15, # uptake for P (maint comp) equal that for U (growth)
            ],
        :x0 => [s => 4.17e-5, b => 1.14e-5, cr_tot => 0, r => 1]
    ) # see S5
    chak21_phys_s, chak21_phys, p_straw
end


# Validated from http://www.ams.jhu.edu/~castello/625.414/Handouts/FractionalProg.pdf

@testset "Simple example" begin
    lfp = LinearFractionalModel(with_optimizer(Clp.Optimizer))
    x1 = @variable(lfp, lower_bound=0, base_name="x1")
    x2 = @variable(lfp, lower_bound=0, upper_bound=6, base_name="x2")
    @constraint(lfp, -x1 + x2 <= 4)
    @constraint(lfp, 2x1 + x2 <= 14)

    LinearFractional.set_objective(lfp, MOI.MIN_SENSE,
              @expression(lfp, -2x1 + x2 + 2),
              @expression(lfp, x1 + 3x2 + 4))

    optimize!(lfp)
    @test termination_status(lfp) === MOI.OPTIMAL
    @test value(x1) ≈ 7.0
    @test value(x2) ≈ 0.0
end


# =====================================
# MNA 5
#
# system type: affine continuous system
# state dimension: 10913
# =====================================
using ReachabilityBenchmarks, MathematicalSystems, MAT, SparseArrays

function mna5_model()
    file = matopen(@relpath "mna5.mat")

    # system matrix
    A = sparse(read(file, "A"))

    # affine term
    b = sparsevec(19:27, [fill(-0.1, 5); fill(-0.2, 4)], 10913)

    # affine continuous system
    S = @system(x' = Ax + b)

    return S
end


function beta(u::UGrid, v::VGrid, dt)
    umax = maximum(u)
    vmax = maximum(v)

    (abs(umax) + abs(vmax)) / min(u.grid.dx, u.grid.dy) * dt
end
beta(sim::Simulation) = beta(sim.u, sim.v, sim.dt)

function diagr(sim::Simulation)
    getnu(sim.params, sim.u.grid) * sim.dt / min(sim.u.grid.dx, sim.u.grid.dy)^2
end

function reynh(sim::Simulation)
    beta(sim) / diagr(sim)
end

function reynhvort(sim::Simulation)
    vort = VortGrid(sim.u, sim.v)
    maxvort = maximum(vort)
    maxvort * min(sim.u.grid.dx, sim.v.grid.dy)^2 / getnu(sim.params, sim.u.grid)
end


using Healpix
using PowerSpectra
using Plots

# retrieve the downgraded planck maps
nside = 256
map1 = 1e6 * PowerSpectra.planck256_map("100", "hm1", "I_STOKES")
map2 = 1e6 * PowerSpectra.planck256_map("100", "hm2", "I_STOKES")
plot(map1, clim=(-200,200))

mask1 = PowerSpectra.planck256_mask("100", "hm1", :T)
mask2 = PowerSpectra.planck256_mask("100", "hm2", :T)
plot(mask1)

# compute the cross-spectra of the masked maps
masked_alm_1 = map2alm(map1 * mask1)
masked_alm_2 = map2alm(map2 * mask2)
pCl = SpectralVector(alm2cl(masked_alm_1, masked_alm_2))

# compute the mode coupling matrix
M_TT = mcm(:TT, map2alm(mask1), map2alm(mask2))

# perform a linear solve to undo the effects of mode-coupling
Cl = M_TT \ pCl

# get lmax from nyquist frequency
lmax = nside2lmax(nside)

# get the planck instrumental beam and pixel window
Wl = PowerSpectra.planck_beam_Wl("100", "hm1", "100", "hm2", :TT, :TT; lmax=lmax)
pixwinT = SpectralVector(pixwin(nside)[1:(lmax+1)])

# plot our spectra
ell = eachindex(Cl)
prefactor = ell .* (ell .+ 1) ./ (2π)
plot(prefactor .*  Cl ./ (Wl .* pixwinT.^2), label="\$D_{\\ell}\$", xlim=(0,2nside) )

# compare it to theory
theory = PowerSpectra.planck_theory_Dl()  # returns a Dict of Dl indexed with :TT, :EE, ...
plot!(theory[:TT], label="theory TT")

##
using Healpix
using PowerSpectra
using Plots

nside = 256
m₁ = PowerSpectra.planck256_polmap("100", "hm1")
m₂ = PowerSpectra.planck256_polmap("100", "hm2")
maskT₁ = PowerSpectra.planck256_mask("100", "hm1", :T)
maskP₁ = PowerSpectra.planck256_mask("100", "hm1", :P)
maskT₂ = PowerSpectra.planck256_mask("100", "hm2", :T)
maskP₂ = PowerSpectra.planck256_mask("100", "hm2", :P)

# convert to μK
scale!(m₁, 1e6)
scale!(m₂, 1e6)

lmax = nside2lmax(256)
# utility function for doing TEB mode decoupling
Cl = master(m₁, maskT₁, maskP₁,
            m₂, maskT₂, maskP₂)
lmax = nside2lmax(256)
print(keys(Cl))
spec = :TE
Wl = PowerSpectra.planck_beam_Wl("100", "hm1", "100", "hm2", spec, spec; lmax=lmax)
pixwinT, pixwinP = pixwin(nside; pol=true)
pixwinT = SpectralVector(pixwinT[1:(lmax+1)])
pixwinP = SpectralVector(pixwinP[1:(lmax+1)])

ℓ = eachindex(Wl)
plot( (ℓ.^2 / (2π)) .*  Cl[spec] ./ (Wl .* pixwinT .* pixwinP),
    label="\$D_{\\ell}\$", xlim=(0,2nside) )
theory = PowerSpectra.planck_theory_Dl()
plot!(theory[spec], label="theory $(spec)")
##

spec = :EE
Wl = PowerSpectra.planck_beam_Wl("100", "hm1", "100", "hm2", spec, spec; lmax=lmax)
ℓ = eachindex(Wl)
plot( (ℓ.^2 / (2π)) .*  Cl[spec] ./ (Wl .* pixwinP.^2),
    label="\$D_{\\ell}\$", xlim=(0,2nside) )
theory = PowerSpectra.planck_theory_Dl()
plot!(theory[spec], label="theory $(spec)", ylim=(-10,50))
##


using CairoMakie, Makie.GeometryBasics
n = 20
function ngonShape(h, k, r, n)
    Polygon([Point2f(h .+ r*sin.(m*2π/n),k .+ r*cos.(m*2π/n)) for m in 1:n])
end
polysCentric = [ngonShape(0, 0, 3/i^1.5, i) for i in 3:n]
polysCircular = [ngonShape(√2/2*sin(θ), √2/2*cos(θ), 0.15/√idx, idx + 2)
    for (idx, θ) in enumerate(LinRange(0,2π*(1 -1/(n-2)), n-2))]
cmap = to_colormap(:linear_protanopic_deuteranopic_kbw_5_98_c40_n256)[3:end]

with_theme(theme_black()) do
    fig, ax, = poly(polysCentric; color = 1:n-2, colormap = cmap,
        axis = (;aspect = DataAspect()), figure = (;resolution = (600,400)))
    poly!(polysCircular; color = 1:n-2, colormap = cmap)
    hidedecorations!(ax; grid = false)
    hidespines!(ax)
    save("polygons.svg", fig)
    display(fig)
end


# light transport, e.g. reflection, refraction

"""Compute reflection vector for v (pointing to surface) and normal n⃗.

	See [diagram](https://raytracing.github.io/books/RayTracingInOneWeekend.html#metal/mirroredlightreflection)"""
@inline @fastmath reflect(v::Vec3{T}, n⃗::Vec3{T}) where T = v - (2v⋅n⃗)*n⃗

"""
	Args:
		refraction_ratio: incident refraction index divided by refraction index of
			hit surface. i.e. η/η′ in the figure above"""
@inline @fastmath function refract(dir::Vec3{T}, n⃗::Vec3{T}, refraction_ratio::T) where T
	cosθ = min(-dir ⋅ n⃗, one(T))
	r_out_perp = refraction_ratio * (dir + cosθ*n⃗)
	r_out_parallel = -√(abs(one(T)-squared_length(r_out_perp))) * n⃗
	normalize(r_out_perp + r_out_parallel)
end

@inline @fastmath function reflectance(cosθ, refraction_ratio)
	# Use Schlick's approximation for reflectance.
	# claforte: may be buggy? I'm getting black pixels in the Hollow Glass Sphere...
	r0 = (1-refraction_ratio) / (1+refraction_ratio)
	r0 = r0^2
	r0 + (1-r0)*((1-cosθ)^5)
end


const fp64max = floatmax(Float64)
const maxhyp_fp64max = 710.4758600739439
const minhyp_fp64max = 5.562684646268003e-309

function sinh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    abs(HI(x)) > maxhyp_fp64max && return DoubleFloat{T}(copysign(Inf, HI(x)))
    !isfinite(HI(x)) && return x
    iszero(x) && return zero(x)
    signbit(x) && return -sinh(-x)

    y = exp(x) - exp(-x)
    z = DoubleFloat{T}(y.hi*0.5, y.lo*0.5)
    return z
end

function cosh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    abs(HI(x)) > maxhyp_fp64max && return DoubleFloat{T}(Inf)
    !isfinite(x) && return abs(x)
    iszero(x) && return one(x)
    signbit(x) && return cosh(-x)

    y = exp(x) + exp(-x)
    z = DoubleFloat{T}(y.hi*0.5, y.lo*0.5)
    return z
end

function tanh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    abs(HI(x)) > maxhyp_fp64max && return copysign(one(DoubleFloat{T}), x)
    return sinh(x) / cosh(x)
end
#=
function tanh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    x < 0 && return -tanh(-x)
    iszero(x) && return zero(x)
    !isfinite(x) && return nan(typeof(x))
    epos = exp(x)
    eneg = exp(-x)
    z = (epos - eneg) / (epos + eneg)
    return z
end
=#

function csch(x::DoubleFloat{T}) where {T<:IEEEFloat}
    abs(HI(x)) > maxhyp_fp64max && return copysign(zero(DoubleFloat{T}), x)
    return inv(sinh(x))
end

function sech(x::DoubleFloat{T}) where {T<:IEEEFloat}
    abs(HI(x)) > maxhyp_fp64max && return zero(DoubleFloat{T})
    return inv(cosh(x))
end

function coth(x::DoubleFloat{T}) where {T<:IEEEFloat}
    abs(HI(x)) > maxhyp_fp64max && return copysign(one(DoubleFloat{T}), x)
    return cosh(x) / sinh(x)
end


using RegressionDiscontinuity
using Hypatia
using Roots, Distributions
using GLM, DataFrames
using Test

data = RDData(RegressionDiscontinuity.Lee08())

result = fit(
    ImbensWagerOptRD(
        B = 14.28,
        solver = Hypatia.Optimizer,
        num_buckets = 2000,
        variance = Homoskedastic(),
    ),
    data,
)

#test that the result is close to the R package result with the same settings
@test result.tau_est ≈ 0.05936146 atol = 0.001
@test result.se_est ≈ 0.01975009 atol = 0.0001
@test result.maxbias ≈ 0.01019907 atol = 0.0001

level = 0.95
rel_bias = result.maxbias / result.se_est
zz = fzero(
    z -> cdf(Normal(), rel_bias - z) + cdf(Normal(), -rel_bias - z) + level - 1,
    0,
    rel_bias - quantile(Normal(), (1 - level) / 2.1),
)

result_plusminus = result.tau_est - result.ci[1]
est_plusminus = zz * result.se_est

@test result_plusminus ≈ est_plusminus atol = 0.00001


@doc raw"""
    MultivariateBasicStructural(y::Matrix{Fl}, s::Int) where Fl

An implementation of a non-homogeneous seemingly unrelated time series equations for basic structural
state-space model consists of trend (local linear trend) and seasonal components.
It is defined by:
```math
\begin{gather*}
    \begin{aligned}
    y_{t} &=  \mu_{t} + \gamma_{t} + \varepsilon_{t} \quad &\varepsilon_{t} \sim \mathcal{N}(0, \Sigma_{\varepsilon})\\
    \mu_{t+1} &= \mu_{t} + \nu_{t} + \xi_{t} \quad &\xi_{t} \sim \mathcal{N}(0, \Sigma_{\xi})\\
    \nu_{t+1} &= \nu_{t} + \zeta_{t} \quad &\zeta_{t} \sim \mathcal{N}(0, \Sigma_{\zeta})\\
    \gamma_{t+1} &= -\sum_{j=1}^{s-1} \gamma_{t+1-j} + \omega_{t} \quad & \omega_{t} \sim \mathcal{N}(0, \Sigma_{\omega})\\
    \end{aligned}
\end{gather*}
```

# Example
```jldoctest
julia> model = MultivariateBasicStructural(rand(100, 2), 12)
MultivariateBasicStructural model
```

# References
 * Durbin, James, & Siem Jan Koopman. (2012). "Time Series Analysis by State Space Methods: Second Edition." Oxford University Press.
"""
mutable struct MultivariateBasicStructural <: StateSpaceModel
    hyperparameters::HyperParameters
    system::LinearMultivariateTimeInvariant
    seasonality::Int
    results::Results

    function MultivariateBasicStructural(y::Matrix{Fl}, s::Int) where Fl
        p = size(y, 2)
        Z = kron(Matrix{Fl}(I, p, p), [1 0 1 zeros(Fl, 1, s - 2)])
        T = kron(
            Matrix{Fl}(I, p, p),
            [
                1 1 zeros(Fl, 1, s - 1)
                0 1 zeros(Fl, 1, s - 1)
                0 0 -ones(Fl, 1, s - 1)
                zeros(Fl, s - 2, 2) Matrix{Fl}(I, s - 2, s - 2) zeros(Fl, s - 2)
            ],
        )

        R = kron(
            Matrix{Fl}(I, p, p),
            [
                Matrix{Fl}(I, 3, 3)
                zeros(Fl, s - 2, 3)
            ],
        )

        d = zeros(Fl, p)
        c = zeros(Fl, p * (s + 1))
        H = kron(one(Fl), Matrix{Fl}(I, p, p))
        Q = kron(ones(Fl, p, p), Matrix{Fl}(I, 3, 3))

        system = LinearMultivariateTimeInvariant{Fl}(y, Z, T, R, d, c, H, Q)

        names = handle_multivariate_basicstructural_names(p)
        hyperparameters = HyperParameters{Fl}(names)
        return new(hyperparameters, system, s, Results{Fl}())
    end
end

function handle_multivariate_basicstructural_names(p::Int)
    # generate \varepsilon  names
    greek_letters_for_states = ["ε", "ξ", "ζ", "ω"]
    names = String[]
    # Walk the lower triangle positions
    for letter in greek_letters_for_states, i in 1:p, j in 1:i
        if i == j
            push!(names, "sigma2_$(letter)$(i)")
        else
            push!(names, "sigma_$(letter)$(i)sigma_$(letter)$(j)")
        end
    end
    return names
end

function default_filter(model::MultivariateBasicStructural)
    Fl = typeof_model_elements(model)
    steadystate_tol = Fl(1e-5)
    a1 = zeros(Fl, num_states(model))
    P1 = Fl(1e6) .* Matrix{Fl}(I, num_states(model), num_states(model))
    return MultivariateKalmanFilter(
        size(model.system.y, 2), a1, P1, num_states(model), steadystate_tol
    )
end

function initial_hyperparameters!(model::MultivariateBasicStructural)
    Fl = typeof_model_elements(model)
    initial_hyperparameters = Dict{String,Fl}()
    for variable in get_names(model)
        if occursin("sigma2", variable)
            initial_hyperparameters[variable] = one(Fl)
        else
            initial_hyperparameters[variable] = zero(Fl)
        end
    end
    set_initial_hyperparameters!(model, initial_hyperparameters)
    return nothing
end

function constrain_hyperparameters!(model::MultivariateBasicStructural)
    Fl = typeof_model_elements(model)
    p = size(model.system.y, 2)
    # H
    H_unconstrained = zeros(Fl, p, p)
    current_name = 1
    for i in 1:p, j in 1:i
        H_unconstrained[i, j] = get_unconstrained_value(
            model, get_names(model)[current_name]
        )
        current_name += 1
    end
    H_constrained = H_unconstrained * H_unconstrained'
    # Q
    Q_unconstrained = zeros(Fl, 3p, 3p)
    sigma2_ξ = filter(x -> occursin("sigma2_ξ", x), get_names(model))
    sigma2_ξ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, 0)))
        if mod1(i, 3) == 1
            Q_unconstrained[pos] = get_unconstrained_value(
                model, sigma2_ξ[sigma2_ξ_counter]
            )
            sigma2_ξ_counter += 1
        end
    end
    sigma2_ζ = filter(x -> occursin("sigma2_ζ", x), get_names(model))
    sigma2_ζ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, 0)))
        if mod1(i, 3) == 2
            Q_unconstrained[pos] = get_unconstrained_value(
                model, sigma2_ζ[sigma2_ζ_counter]
            )
            sigma2_ζ_counter += 1
        end
    end
    sigma2_ω = filter(x -> occursin("sigma2_ω", x), get_names(model))
    sigma2_ω_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, 0)))
        if mod1(i, 3) == 3
            Q_unconstrained[pos] = get_unconstrained_value(
                model, sigma2_ω[sigma2_ω_counter]
            )
            sigma2_ω_counter += 1
        end
    end
    sigma_ξxsigma_ξx = filter(x -> occursin("sigma_ξ", x), get_names(model))
    sigma_ξxsigma_ξx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 1
            Q_unconstrained[pos] = get_unconstrained_value(
                model, sigma_ξxsigma_ξx[sigma_ξxsigma_ξx_counter]
            )
            sigma_ξxsigma_ξx_counter += 1
        end
    end
    sigma_ζxsigma_ζx = filter(x -> occursin("sigma_ζ", x), get_names(model))
    sigma_ζxsigma_ζx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 2
            Q_unconstrained[pos] = get_unconstrained_value(
                model, sigma_ζxsigma_ζx[sigma_ζxsigma_ζx_counter]
            )
            sigma_ζxsigma_ζx_counter += 1
        end
    end
    sigma_ωxsigma_ωx = filter(x -> occursin("sigma_ω", x), get_names(model))
    sigma_ωxsigma_ωx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 3
            Q_unconstrained[pos] = get_unconstrained_value(
                model, sigma_ωxsigma_ωx[sigma_ωxsigma_ωx_counter]
            )
            sigma_ωxsigma_ωx_counter += 1
        end
    end
    Q_constrained = Q_unconstrained * Q_unconstrained'

    # Updates
    # H
    current_name = 1
    for i in 1:p, j in 1:i
        update_constrained_value!(
            model, get_names(model)[current_name], H_constrained[i, j]
        )
        current_name += 1
    end
    # Q
    sigma2_ξ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 1
            update_constrained_value!(model, sigma2_ξ[sigma2_ξ_counter], Q_constrained[pos])
            sigma2_ξ_counter += 1
        end
    end
    sigma2_ζ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 2
            update_constrained_value!(model, sigma2_ζ[sigma2_ζ_counter], Q_constrained[pos])
            sigma2_ζ_counter += 1
        end
    end
    sigma2_ω_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 3
            update_constrained_value!(model, sigma2_ω[sigma2_ω_counter], Q_constrained[pos])
            sigma2_ω_counter += 1
        end
    end
    sigma_ξxsigma_ξx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 1
            update_constrained_value!(
                model, sigma_ξxsigma_ξx[sigma_ξxsigma_ξx_counter], Q_constrained[pos]
            )
            sigma_ξxsigma_ξx_counter += 1
        end
    end
    sigma_ζxsigma_ζx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 2
            update_constrained_value!(
                model, sigma_ζxsigma_ζx[sigma_ζxsigma_ζx_counter], Q_constrained[pos]
            )
            sigma_ζxsigma_ζx_counter += 1
        end
    end
    sigma_ωxsigma_ωx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 3
            update_constrained_value!(
                model, sigma_ωxsigma_ωx[sigma_ωxsigma_ωx_counter], Q_constrained[pos]
            )
            sigma_ωxsigma_ωx_counter += 1
        end
    end
    return model
end

function unconstrain_hyperparameters!(model::MultivariateBasicStructural)
    Fl = typeof_model_elements(model)
    p = size(model.system.y, 2)
    # H
    H_constrained = zeros(Fl, p, p)
    current_name = 1
    for i in 1:p, j in 1:i
        H_constrained[i, j] = get_constrained_value(model, get_names(model)[current_name])
        current_name += 1
    end
    H_unconstrained = cholesky(Symmetric(H_constrained, :L)).L
    # Q
    Q_constrained = zeros(Fl, 3p, 3p)
    sigma2_ξ = filter(x -> occursin("sigma2_ξ", x), get_names(model))
    sigma2_ξ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 1
            Q_constrained[pos] = get_constrained_value(model, sigma2_ξ[sigma2_ξ_counter])
            sigma2_ξ_counter += 1
        end
    end
    sigma2_ζ = filter(x -> occursin("sigma2_ζ", x), get_names(model))
    sigma2_ζ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 2
            Q_constrained[pos] = get_constrained_value(model, sigma2_ζ[sigma2_ζ_counter])
            sigma2_ζ_counter += 1
        end
    end
    sigma2_ω = filter(x -> occursin("sigma2_ω", x), get_names(model))
    sigma2_ω_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 3
            Q_constrained[pos] = get_constrained_value(model, sigma2_ω[sigma2_ω_counter])
            sigma2_ω_counter += 1
        end
    end
    sigma_ξxsigma_ξx = filter(x -> occursin("sigma_ξ", x), get_names(model))
    sigma_ξxsigma_ξx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 1
            Q_constrained[pos] = get_constrained_value(
                model, sigma_ξxsigma_ξx[sigma_ξxsigma_ξx_counter]
            )
            sigma_ξxsigma_ξx_counter += 1
        end
    end
    sigma_ζxsigma_ζx = filter(x -> occursin("sigma_ζ", x), get_names(model))
    sigma_ζxsigma_ζx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 2
            Q_constrained[pos] = get_constrained_value(
                model, sigma_ζxsigma_ζx[sigma_ζxsigma_ζx_counter]
            )
            sigma_ζxsigma_ζx_counter += 1
        end
    end
    sigma_ωxsigma_ωx = filter(x -> occursin("sigma_ω", x), get_names(model))
    sigma_ωxsigma_ωx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 3
            Q_constrained[pos] = get_constrained_value(
                model, sigma_ωxsigma_ωx[sigma_ωxsigma_ωx_counter]
            )
            sigma_ωxsigma_ωx_counter += 1
        end
    end
    Q_unconstrained = cholesky(Symmetric(Q_constrained, :L)).L

    # Updates
    # H
    current_name = 1
    for i in 1:p, j in 1:i
        update_unconstrained_value!(
            model, get_names(model)[current_name], H_unconstrained[i, j]
        )
        current_name += 1
    end
    # Q
    sigma2_ξ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, 0)))
        if mod1(i, 3) == 1
            update_unconstrained_value!(
                model, sigma2_ξ[sigma2_ξ_counter], Q_unconstrained[pos]
            )
            sigma2_ξ_counter += 1
        end
    end
    sigma2_ζ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, 0)))
        if mod1(i, 3) == 2
            update_unconstrained_value!(
                model, sigma2_ζ[sigma2_ζ_counter], Q_unconstrained[pos]
            )
            sigma2_ζ_counter += 1
        end
    end
    sigma2_ω_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, 0)))
        if mod1(i, 3) == 3
            update_unconstrained_value!(
                model, sigma2_ω[sigma2_ω_counter], Q_unconstrained[pos]
            )
            sigma2_ω_counter += 1
        end
    end
    sigma_ξxsigma_ξx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 1
            update_unconstrained_value!(
                model, sigma_ξxsigma_ξx[sigma_ξxsigma_ξx_counter], Q_unconstrained[pos]
            )
            sigma_ξxsigma_ξx_counter += 1
        end
    end
    sigma_ζxsigma_ζx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 2
            update_unconstrained_value!(
                model, sigma_ζxsigma_ζx[sigma_ζxsigma_ζx_counter], Q_unconstrained[pos]
            )
            sigma_ζxsigma_ζx_counter += 1
        end
    end
    sigma_ωxsigma_ωx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_unconstrained, -3)))
        if mod1(i, 3) == 3
            update_unconstrained_value!(
                model, sigma_ωxsigma_ωx[sigma_ωxsigma_ωx_counter], Q_unconstrained[pos]
            )
            sigma_ωxsigma_ωx_counter += 1
        end
    end
    return model
end

function fill_model_system!(model::MultivariateBasicStructural)
    Fl = typeof_model_elements(model)
    p = size(model.system.y, 2)
    H_constrained = zeros(Fl, p, p)
    current_name = 1
    for i in 1:p, j in 1:i
        H_constrained[i, j] = get_constrained_value(model, get_names(model)[current_name])
        current_name += 1
    end
    model.system.H = Matrix(Symmetric(H_constrained, :L))
    # Q
    Q_constrained = zeros(Fl, 3p, 3p)
    sigma2_ξ = filter(x -> occursin("sigma2_ξ", x), get_names(model))
    sigma2_ξ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 1
            Q_constrained[pos] = get_constrained_value(model, sigma2_ξ[sigma2_ξ_counter])
            sigma2_ξ_counter += 1
        end
    end
    sigma2_ζ = filter(x -> occursin("sigma2_ζ", x), get_names(model))
    sigma2_ζ_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 2
            Q_constrained[pos] = get_constrained_value(model, sigma2_ζ[sigma2_ζ_counter])
            sigma2_ζ_counter += 1
        end
    end
    sigma2_ω = filter(x -> occursin("sigma2_ω", x), get_names(model))
    sigma2_ω_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, 0)))
        if mod1(i, 3) == 3
            Q_constrained[pos] = get_constrained_value(model, sigma2_ω[sigma2_ω_counter])
            sigma2_ω_counter += 1
        end
    end
    sigma_ξxsigma_ξx = filter(x -> occursin("sigma_ξ", x), get_names(model))
    sigma_ξxsigma_ξx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 1
            Q_constrained[pos] = get_constrained_value(
                model, sigma_ξxsigma_ξx[sigma_ξxsigma_ξx_counter]
            )
            sigma_ξxsigma_ξx_counter += 1
        end
    end
    sigma_ζxsigma_ζx = filter(x -> occursin("sigma_ζ", x), get_names(model))
    sigma_ζxsigma_ζx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 2
            Q_constrained[pos] = get_constrained_value(
                model, sigma_ζxsigma_ζx[sigma_ζxsigma_ζx_counter]
            )
            sigma_ζxsigma_ζx_counter += 1
        end
    end
    sigma_ωxsigma_ωx = filter(x -> occursin("sigma_ω", x), get_names(model))
    sigma_ωxsigma_ωx_counter = 1
    for (i, pos) in enumerate(collect(diagind(Q_constrained, -3)))
        if mod1(i, 3) == 3
            Q_constrained[pos] = get_constrained_value(
                model, sigma_ωxsigma_ωx[sigma_ωxsigma_ωx_counter]
            )
            sigma_ωxsigma_ωx_counter += 1
        end
    end
    model.system.Q = Matrix(Symmetric(Q_constrained, :L))
    return model
end

function reinstantiate(model::MultivariateBasicStructural, y::Matrix{Fl}) where Fl
    return MultivariateBasicStructural(y, model.seasonality)
end

has_exogenous(::MultivariateBasicStructural) = false


using TensorKit
using statmech_tm_solver
using TensorOperations
using LinearAlgebra

io = open("result_triangular_ising_fulldiag_obc.txt", "w")
close(io)
T = mpo_triangular_AF_ising()

BL = Tensor(zeros, ComplexF64, ℂ^2)
BL[1] = BL[2] = 1;
BL
BR = Tensor(zeros, ComplexF64, ℂ^2)
BR[1] = BR[2] = 1;
BR

T5 = ncon([BL', T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6]])
T5 = permute(T5, (1, 3, 5, 7, 9), (2, 4, 6, 8, 10))

T5

T6 = ncon([BL', T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7]])
T6 = permute(T6, (1, 3, 5, 7, 9, 11), (2, 4, 6, 8, 10, 12))

T7 = ncon([BL', T, T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7, -13, -14, 8], [8]])
T7 = permute(T7, (1, 3, 5, 7, 9, 11, 13), (2, 4, 6, 8, 10, 12, 14))

T8 = ncon([BL', T, T, T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7, -13, -14, 8], [8, -15, -16, 9], [9]])
T8 = permute(T8, (1, 3, 5, 7, 9, 11, 13, 15), (2, 4, 6, 8, 10, 12, 14, 16))

T9 = ncon([BL', T, T, T, T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7, -13, -14, 8], [8, -15, -16, 9], [9, -17, -18, 10], [10]])
T9 = permute(T9, (1, 3, 5, 7, 9, 11, 13, 15, 17), (2, 4, 6, 8, 10, 12, 14, 16, 18))

T10 = ncon([BL', T, T, T, T, T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7, -13, -14, 8], [8, -15, -16, 9], [9, -17, -18, 10], [10, -19, -20, 11], [11]])
T10 = permute(T10, (1, 3, 5, 7, 9, 11, 13, 15, 17, 19), (2, 4, 6, 8, 10, 12, 14, 16, 18, 20))

T11 = ncon([BL', T, T, T, T, T, T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7, -13, -14, 8], [8, -15, -16, 9], [9, -17, -18, 10], [10, -19, -20, 11], [11, -21, -22, 12], [12]])
T11 = permute(T11, (1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21), (2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22))

T12 = ncon([BL', T, T, T, T, T, T, T, T, T, T, T, T, BR], [[1], [1, -1, -2, 2], [2, -3, -4, 3], [3, -5, -6, 4], [4, -7, -8, 5], [5, -9, -10, 6], [6, -11, -12, 7], [7, -13, -14, 8], [8, -15, -16, 9], [9, -17, -18, 10], [10, -19, -20, 11], [11, -21, -22, 12], [12, -23, -24, 13], [13]])
T12 = permute(T12, (1, 3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23), (2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24))

for (n, Tn) in zip([5, 6, 7, 8, 9, 10, 11, 12], [T5, T6, T7, T8, T9, T10, T11, T12])
    w, _ = eig(Tn)
    w = diag(w.data)
    w_reals = real.(w)
    w_imags = imag.(w)

    io = open("result_triangular_ising_fulldiag_obc.txt", "a+")
    for (w_real, w_imag) in zip(w_reals, w_imags)
        println(io, "$(n)   $(w_real)   $(w_imag) ")
    end
    close(io)
end


using SpecialFunctions, QuadGK, ForwardDiff

##### Diffusion approximation for 2-layered medium with matched index of refraction #####

"""
    fluence_spatial_freq(z, s, μa, μsp, l)

Compute the spatial frequency components of the fluence for 2-layered model

# Arguments
- `z`: depth (z=0 at the surface, z=l at interface of two layers).
- `s`: spatial frequency component s^2 = sx^2 + sy^2
- `μa`: absorption coefficient of both layers [μa1, μa2]
- `μsp`: reduced scattering coefficient of both layers [μsp1, μsp2]
- `l`: thickness of first layer

# Examples
```jldoctest
julia> fluence_spatial_freq(0., 1,  [0.005,0.1], [10,100.], 1)
0.8647044203305079
```
"""
function fluence_spatial_freq(z, s, μa, μsp, l)

    #get numerical instabillity at large s
    D = @. 1/(3(μsp + μa))
    α = @. sqrt((D*s^2 + μa)/D)
    A = 0.493
    zb = 2*A*D[1]
    z0 = 1/(μsp[1] + μa[1])

    ## equation 11
    if z < z0

        # have to account for exp blowing up to inf
        if (α[1]*z0 - α[1]*z) > 15
            ϕ = exp(α[1]*(z0 - 2*l)) - exp(α[1]*(-z0 - 2*l - zb)) - exp(α[1]*(z0 - 2*l - 2*zb)) + exp(α[1]*(-z0 - 2*l - 4*zb))
            ϕ = ϕ*(D[1]*α[1] - D[2]*α[2])/(D[1]*α[1] + D[2]*α[2])
            ϕ = ϕ + exp(-α[1]*z0) - exp(-α[1]*(2*zb + z0))
            ϕ = ϕ/(2*D[1]*α[1])

        else

            ϕ = D[1]*α[1]*(1 + exp(-2*α[1]*(l - z))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l - z)))
            ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
            ϕ = ϕ*(1 - exp(-2*α[1]*(zb + z0)))/(2*D[1]*α[1])
            ϕ = ϕ*exp(α[1]*z0 - α[1]*z)
            ϕ = ϕ - ((exp(α[1]*(z0 - z))*(1 - exp(-2*α[1]*(z0 - z))))/(2*D[1]*α[1]))

        end

    ## equation 12
    elseif z > z0 && z < l
        ϕ = D[1]*α[1]*(1 + exp(-2*α[1]*(l - z))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l - z)))
        ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
        ϕ = ϕ*(1 + exp(-2*α[1]*(zb + z0)))/(2*D[1]*α[1])
        ϕ = ϕ*exp(α[1]*z0 - α[1]*z)

    #equation 13
    elseif z >= l

        ϕ = exp(α[1]*(z0 - l) + α[2]*(l - z))
        ϕ = ϕ*(1 - exp(-2*α[1]*(zb + z0)))
        ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) + D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
    end

    return ϕ
end

function fluence_spatial_complexfreq(z, s, μa, μsp, l, f, nmed)

    #get numerical instabillity at large s
    D = @. 1/(3(μsp + μa))
    ν = 29.9792345/nmed
    ω = 2*π*f
    α = @. sqrt((D*s^2 + μa + ω*im/ν)/D)
    A = 0.493
    zb = 2*A*D[1]
    z0 = 1/(μsp[1] + μa[1])

    ## equation 11
    if z < z0

        # have to account for exp blowing up to inf
        if abs2(α[1]*z0 - α[1]*z) > 15
            ϕ = exp(α[1]*(z0 - 2*l)) - exp(α[1]*(-z0 - 2*l - zb)) - exp(α[1]*(z0 - 2*l - 2*zb)) + exp(α[1]*(-z0 - 2*l - 4*zb))
            ϕ = ϕ*(D[1]*α[1] - D[2]*α[2])/(D[1]*α[1] + D[2]*α[2])
            ϕ = ϕ + exp(-α[1]*z0) - exp(-α[1]*(2*zb + z0))
            ϕ = ϕ/(2*D[1]*α[1])

        else

            ϕ = D[1]*α[1]*(1 + exp(-2*α[1]*(l - z))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l - z)))
            ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
            ϕ = ϕ*(1 - exp(-2*α[1]*(zb + z0)))/(2*D[1]*α[1])
            ϕ = ϕ*exp(α[1]*z0 - α[1]*z)
            ϕ = ϕ - ((exp(α[1]*(z0 - z))*(1 - exp(-2*α[1]*(z0 - z))))/(2*D[1]*α[1]))

        end

    ## equation 12
    elseif z > z0 && z < l
        ϕ = D[1]*α[1]*(1 + exp(-2*α[1]*(l - z))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l - z)))
        ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) +  D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
        ϕ = ϕ*(1 + exp(-2*α[1]*(zb + z0)))/(2*D[1]*α[1])
        ϕ = ϕ*exp(α[1]*z0 - α[1]*z)

    #equation 13
    elseif z >= l

        ϕ = exp(α[1]*(z0 - l) + α[2]*(l - z))
        ϕ = ϕ*(1 - exp(-2*α[1]*(zb + z0)))
        ϕ = ϕ/(D[1]*α[1]*(1 + exp(-2*α[1]*(l + zb))) + D[2]*α[2]*(1 - exp(-2*α[1]*(l + zb))))
    end

    return ϕ
end

function compute_complexintegral(ρ, z, s, μa, μsp, l, f, nmed)
    return fluence_spatial_complexfreq(z, s, μa, μsp, l, f, nmed)*s*besselj0(ρ*s)
end

function fluence_FD(ρ, z, μa, μsp, l, f, nmed)
    return (quadgk((s) -> compute_complexintegral(ρ, z, s, μa, μsp, l, f, nmed), 0, Inf, rtol=1e-6)[1])/2π
end

function fluence_TD(ρ, z, μa, μsp, l, nmed)
    fr = 0.19531 .* Vector(1:512)
	Rt = Array{Complex{Float64},1}(undef, length(fr))
    map!((f) -> fluence_FD(ρ, z, μa, μsp, l, f, nmed), Rt, fr)
end

# compute equation 14 integrand
function compute_integral(ρ, z, s, μa, μsp, l)
    return fluence_spatial_freq(z, s, μa, μsp, l)*s*besselj0(ρ*s)
end

# numerically integrate equation 14
function fluence_steadystate(ρ, z, μa, μsp, l)
    return (quadgk((s) -> compute_integral(ρ, z, s, μa, μsp, l), 0, Inf, rtol=1e-6)[1])/2π
 end

# Eqn 16 to compute reflectance (use auto diff)
function refl_steadystate(ρ, μa, μsp, l)
    f(z) = fluence_steadystate(ρ, z, μa, μsp, l)
    g(z) = ForwardDiff.derivative(z -> f(z), z)
    return 0.118*fluence_steadystate(ρ, 0, μa, μsp, l) + 0.306/(3*(μa[1] + μsp[1]))*g(0)
end

# Eqn 3 Kienle 97 for spatial fluence semi-infinite (use as comparison for l -> ∞)
function ss_compare(ρ, z, μa, μsp)
    D = 1/(3(μsp + μa))
    A = 0.493
    z0 = 1/(μsp + μa)
    zb = 2*A*D
    μe = sqrt(3*μa*(μa + μsp))

    ϕ = exp(-μe*sqrt((z - z0)^2 + ρ^2))/sqrt((z - z0)^2 + ρ^2)
    ϕ = ϕ - exp(-μe*sqrt((z + z0 + 2*zb)^2 + ρ^2))/sqrt((z + z0 + 2*zb)^2 + ρ^2)
    ϕ = ϕ/(4*π*D)
end

function refl_steadystate1(ρ, μa, μsp, l)
    f(z) = ss_compare(ρ, z, μa, μsp)
    g(z) = ForwardDiff.derivative(z -> f(z), z)
    return 0.118*ss_compare(ρ, 0, μa, μsp) + 0.306/(3*(μa[1] + μsp[1]))*g(0)
end

##reproduce Figure 2 Kienle 98 *seem to be slightly off by some factor
a = map((ρ) -> refl_steadystate(ρ, [0.02,0.01], [1.3,1.2], 2), 0.5:0.5:20)
b = map((ρ) -> refl_steadystate(ρ, [0.02,0.01], [1.3,0.7], 2), 0.5:0.5:20)

plot(0.5:0.5:20, a, yscale=:log10, lw = 2, label = "crosses", xlabel = "distance [mm]", ylabel = "reflectance")
plot!(0.5:0.5:20, b, yscale=:log10, lw = 2, label = "circles")

##reproduce Figure 3 Kienle 98
b = map((ρ) -> refl_steadystate(ρ, [0.005,0.01], [1.3,1.0], 6), 0.5:0.5:20)
a = map((ρ) -> refl_steadystate(ρ, [0.005,0.022], [1.3,1.0], 6), 0.5:0.5:20)

plot(0.5:0.5:20, a, yscale=:log10, lw = 2, label = "crosses", xlabel = "distance [mm]", ylabel = "reflectance")
plot!(0.5:0.5:20, b, yscale=:log10, lw = 2, label = "circles")

##compare to semi-infinite model from Kienle 97 as l gets large
a = map((ρ) -> fluence_steadystate(ρ, 0, [0.02,0.02], [1.3,1.3], 10), 0.5:0.5:20)
b = map((ρ) -> ss_compare(ρ, 0, 0.02, 1.3), 0.5:0.5:20)

plot(0.5:0.5:20, a, yscale=:log10, lw = 2, label = "2-layer", xlabel = "distance [mm]", ylabel = "reflectance")
plot!(0.5:0.5:20, b, yscale=:log10, lw = 2, label = "semi-inf", alpha = 0.7)


include("common.jl")

mmesh = Mesh(10, 10, 1/10, degree=1)
NT = 500
Δt = 2.0/NT
ηmax = 1
ηmin = 0.5

bdedge = bcedge((x1, y1, x2, y2)->(x1>1-1e-5 && x2>1-1e-5), mmesh)
bdnode = bcnode((x1, y1)->(y1<1e-5), mmesh)

invη = 50 * constant(ones(get_ngauss(mmesh)))
μ = 1.0
λ = 1.0

Forces = zeros(NT, 2mmesh.ndof)
T = eval_f_on_boundary_edge((x,y)->0.1, bdedge, mmesh)
T = [-T T]
rhs = compute_fem_traction_term(T, bdedge, mmesh)
for i = 1:NT
  Forces[i, :] = rhs
end
Forces = constant(Forces)

U, Sigma, Varepsilon = viscoelasticity_solver(invη, μ, λ, Forces, mmesh)

sess = Session(); init(sess)
Uval, Sigmaval, Varepsilonval = run(sess, [U, Sigma, Varepsilon])
p = visualize_displacement(Uval[1:5:end,:], mmesh)
saveanim(p, "space_u.gif")


using QuantumGate
using LinearAlgebra

# 初期状態
# Notation: |0> = [1, 0], |1> = [0, 1]
#
firstbit = [1; 0]
ψ = [1; 0]
ket = firstbit ⊗ ψ

# アダマールテストでHadを測定する回路
res = Hadamardgate(1, 2)*controlled_Ugate(1, 2, 2, Had)*Hadamardgate(1, 2)*ket

# 1つ目のqubitが0である確率は
# p = (1 + Re<ψ|H|ψ>)/2
# if |ψ> = |0>, p = (1 + 1/√2)/2
measured = (P⁰ ⊗ σ⁰) * res

dot(measured, measured)
(1 + 1/sqrt(2))/2

# 測定後の状態はHの固有状態である。
Had * measured[1:2] ≈ measured[1:2]

# 1つ目のqubitが1である確率は
# p = (1 - Re<ψ|H|ψ>)/2
# if |ψ> = |0>, p = (1 - 1/√2)/2
measured = (P¹ ⊗ σ⁰) * res

dot(measured, measured)
(1 - 1/sqrt(2))/2

# 測定後の状態はHの固有状態である。
Had * measured[3:4] ≈ -measured[3:4]


export polar,
       eigenHermitian,
       factorize

import LinearAlgebra.qr
function qr(A::ITensor,
            Linds...;
            kwargs...)
  tags::TagSet = get(kwargs,:tags,"Link,qr")
  Lis = commoninds(inds(A),IndexSet(Linds...))
  Ris = uniqueinds(inds(A),Lis)
  Lpos,Rpos = getperms(inds(A),Lis,Ris)
  QT,RT = qr(tensor(A),Lpos,Rpos)
  Q,R = ITensor(QT),ITensor(RT)
  q = commonindex(Q,R)
  settags!(Q,tags,q)
  settags!(R,tags,q)
  q = settags(q,tags)
  return Q,R,q
end

# TODO: allow custom tags in internal indices?
function polar(A::ITensor,
               Linds...;
               kwargs...)
  Lis = commoninds(inds(A),IndexSet(Linds...))
  Ris = uniqueinds(inds(A),Lis)
  Lpos,Rpos = getperms(inds(A),Lis,Ris)
  UT,PT = polar(tensor(A),Lpos,Rpos)
  U,P = ITensor(UT),ITensor(PT)
  u = commoninds(U,P)
  p = uniqueinds(P,U)
  replaceinds!(U,u,p')
  replaceinds!(P,u,p')
  return U,P,commoninds(U,P)
end

import LinearAlgebra.svd
"""
    svd(A::ITensor,
        leftind1::Index,
        leftind2::Index,
        ...
        ;kwargs...)

Singular value decomposition (SVD) of an ITensor A, computed
by treating the "left indices" provided collectively
as a row index, and the remaining "right indices" as a
column index (matricization of a tensor).

Whether the SVD performs a trunction depends on the keyword
arguments provided. The following keyword arguments are recognized:
* `maxdim` [Int]
* `mindim` [Int]
* `cutoff` [Float64]
* `absoluteCutoff` [Bool] Default value: false.
* `doRelCutoff` [Bool] Default value: true.
* `utags` [String] Default value: "Link,u".
* `vtags` [String] Default value: "Link,v".
* `fastSVD` [Bool] Defaut value: false.
"""
function svd(A::ITensor,
             Linds...;
             kwargs...)
  utags::TagSet = get(kwargs,:utags,"Link,u")
  vtags::TagSet = get(kwargs,:vtags,"Link,v")
  Lis = commoninds(inds(A),IndexSet(Linds...))
  Ris = uniqueinds(inds(A),Lis)
  Lpos,Rpos = getperms(inds(A),Lis,Ris)
  UT,ST,VT = svd(tensor(A),Lpos,Rpos;kwargs...)
  U,S,V = ITensor(UT),ITensor(ST),ITensor(VT)
  u₀ = commonindex(U,S)
  v₀ = commonindex(S,V)

  u = settags(u₀,utags)
  v = settags(u₀,vtags)

  U *= δ(dag(u₀),u)
  S = δ(dag(u₀),u)*S*δ(dag(v₀),v)
  V *= δ(dag(v₀),v)

  return U,S,V,u,v
end

function _factorize_center(A::ITensor,
                           Linds...;
                           kwargs...)
  tags::TagSet = get(kwargs,:tags,"Link,u")
  U,S,V = svd(A,Linds...;kwargs...)
  u = commonindex(U,S)
  v = commonindex(S,V)
  for ss = 1:dim(u)
    S[ss,ss] = sqrt(S[ss,ss])
  end
  FU = settags(U*S,tags,v)
  FV = settags(S*V,tags,u)
  return FU,FV,commonindex(FU,FV)
end

function _factorize_from_left_svd(A::ITensor,
                                  Linds...;
                                  kwargs...)
  tags::TagSet = get(kwargs,:tags,"Link,u")
  U,S,V = svd(A,Linds...;kwargs...)
  u = commonindex(U,S)
  FU = settags(U,tags,u)
  FV = settags(S*V,tags,u)
  return FU,FV,commonindex(FU,FV)
end

function _factorize_from_right_svd(A::ITensor,
                                   Linds...;
                                   kwargs...)
  tags::TagSet = get(kwargs,:tags,"Link,u")
  U,S,V = svd(A,Linds...;kwargs...)
  v = commonindex(S,V)
  FU = settags(U*S,tags,v)
  FV = settags(V,tags,v)
  return FU,FV,commonindex(FU,FV)
end

function _factorize_from_left_eigen(A::ITensor,
                                    Linds...;
                                    kwargs...)
  Lis = commoninds(inds(A),IndexSet(Linds...))
  A² = A*prime(dag(A),Lis)
  FU,D = eigenHermitian(A²,Lis,prime(Lis); ispossemidef=true,
                                           kwargs...)
  FV = dag(FU)*A
  return FU,FV,commonindex(FU,FV)
end

function _factorize_from_right_eigen(A::ITensor,
                                     Linds...;
                                     kwargs...)
  Ris = uniqueinds(inds(A),IndexSet(Linds...))
  A² = A*prime(dag(A),Ris)
  FV,D = eigenHermitian(A²,Ris,prime(Ris); ispossemidef=true,
                                           kwargs...)
  FU = A*dag(FV)
  return FU,FV,commonindex(FU,FV)
end

import LinearAlgebra.factorize
"""
factorize(A::ITensor, Linds...; dir = "center", which_factorization = "svd", cutoff = 0.0)

Do a low rank factorization of A either using an SVD or an eigendecomposition of A'A or AA'.
"""
function factorize(A::ITensor,
                   Linds...;
                   kwargs...)
  dir::String = get(kwargs,:dir,"center")
  if dir == "center"
    return _factorize_center(A,Linds...;kwargs...)
  end
  which_factorization::String = get(kwargs,:which_factorization,"svd")
  cutoff::Float64 = get(kwargs,:cutoff,0.0)
  use_eigen = false
  if which_factorization == "eigen" || (which_factorization == "automatic" && cutoff > 1e-12)
    use_eigen = true
  end
  if dir == "fromleft"
    if use_eigen
      return _factorize_from_left_eigen(A,Linds...;kwargs...)
    else
      return _factorize_from_left_svd(A,Linds...;kwargs...)
    end
  elseif dir == "fromright"
    if use_eigen
      return _factorize_from_right_eigen(A,Linds...;kwargs...)
    else
      return _factorize_from_right_svd(A,Linds...;kwargs...)
    end
  end
  throw(ArgumentError("In factorize, no dir = $dir supported. Use center, fromleft or fromright."))
end

function eigenHermitian(A::ITensor,
                        Linds=findinds(A,("",0)),
                        Rinds=prime(IndexSet(Linds));
                        kwargs...)
  tags::TagSet = get(kwargs,:tags,"Link,eigen")
  lefttags::TagSet = get(kwargs,:lefttags,tags)
  righttags::TagSet = get(kwargs,:righttags,prime(tags))
  Lis = commoninds(inds(A),IndexSet(Linds))
  Ris = uniqueinds(inds(A),Lis)
  Lpos,Rpos = getperms(inds(A),Lis,Ris)
  UT,DT = eigenHermitian(tensor(A),Lpos,Rpos;kwargs...)
  U,D = ITensor(UT),ITensor(DT)
  u = commonindex(U,D)
  settags!(U,lefttags,u)
  settags!(D,lefttags,u)
  u = settags(u,lefttags)
  v = uniqueindex(D,U)
  D *= δ(v,settags(u,righttags))
  return U,D,u,v
end

import LinearAlgebra.eigen
function eigen(A::ITensor,
               Linds=findinds(A,("",0)),
               Rinds=prime(IndexSet(Linds));
               kwargs...)
  tags::TagSet = get(kwargs,:tags,"Link,eigen")
  Lis = commoninds(inds(A),IndexSet(Linds))
  Ris = uniqueinds(inds(A),Lis)
  Lpos,Rpos = getperms(inds(A),Lis,Ris)
  UT,DT = eigen(tensor(A),Lpos,Rpos;kwargs...)
  U,D = ITensor(UT),ITensor(DT)
  u = commonindex(U,D)
  settags!(U,tags,u)
  settags!(D,tags,u)
  u = settags(u,tags)
  v = uniqueindex(D,U)
  D *= δ(v,u')
  return U,D,u
end


typeof(2) #> Int64
typeof(Int64) #> DataType
typeof(Complex{Int64}) #> DataType
typeof(Any) #> DataType
typeof(DataType) #> DataType

super(Int64) #> Signed
super(Signed) #> Integer
super(Integer) #> Real
super(Real) #> Number
super(Number) #> Any
super(Any) #> Any

subtypes(Integer) #> 5-element Array{Any,1}:
# BigInt
# Bool
# Char
# Signed
# Unsigned
subtypes(Signed) #> 5-element Array{Any,1}:
# Int128
# Int16
# Int32
# Int64
# Int8
subtypes(Int64) #> 0-element Array{Any,1}
typeof(None) #> UnionType

Bool <: Integer #> true
Bool <: Any #> true
Bool <: Char #> false
issubtype(Bool, Integer) #> true
issubtype(Float64, Integer) #> false

typeof(5) #> Int64
isa(5, Number) #> true


# A simple function to read user inputs
function input(prompt)
    println(prompt)
    uinput = readline();
    chomp(uinput)
    uinput = parse(Int,uinput)
end

"""
recodeimage(pathtoimage, n_singularvlas)

This function accepts path to a folder containing images (.png & .jpg) and
    number of singular values you want as arguments and generates a CSV File
    in the same folder with each row corresoponding to each images in the folder
    and each column corresoponding to it's respective singular values.
    USAGE:
    ```julia
    recodeimage(path,3)
    ```
    In the above example, we are asking to return a CSV File that has the first 3 singular
    values of the images in the folder.
"""
function recodeimage(pathtoimage)
    # The following code block looks for .jpg/.JPG/.png/.PNG files and create a list of them
    imagelist = glob("*.jpg", pathtoimage)
    append!(imagelist,  glob("*.JPG", pathtoimage))
    append!(imagelist,  glob("*.png", pathtoimage))
    append!(imagelist,  glob("*.PNG", pathtoimage))

    # Algorithm for scree plot ===================================================================================
    # creating a temporary array of arrays to hold the numerical values of the grayscaled images
    # temp[1] will have the grayscaled information of image 1, and so on...
    temp = Array{Array{Float64,2},1}(undef,length(imagelist))

    # @showprogress is a macro to print the progress of this loop when this function is run
        @showprogress "Loading Images " for i in 1:length(imagelist)
        temp[i] = Float64.(Gray.(imresize(load(imagelist[i]),128,128)))
        # TODO: delete the remaining commented code in this loop.
        # @pipe is a macro for chaining multiple tasks
        # the next two blocks of code takes the image, converts it into grayscale and compute
        # the singular values, then only the first n_singular values are stored in the recodedArray
        #img_singluar = @pipe X |> Float64.(Gray.(_)) |> svdvals(_)[1:n_singularvlas]'
        #recodedArray = vcat(recodedArray, img_singluar)
    end

    # stacking individual images to create a giant image matrix
    stackedX = vcat(temp...)
    println("Running Singular Value Decomposition...")
    U,S,V= svd(stackedX)
    S = S ./norm(S)
    screeplot(S)
    # ============================================================================================================

    # After examining the scree plot, the user decides the no. of singular values
    n_singularvlas = input("No. of Features (due to bug in the code that reads user inputs, you might have to enter the no twice, if the program didn't run first time)")

    # Initializing an empty array to store the n_singular values of the images
    recodedArray = Array{Float64}(undef, 0, n_singularvlas)
    # computing singular values using only the first n_singulars
    for imagearrays in temp
        img_singular = @pipe imagearrays |> svdvals(_)[1:n_singularvlas]'
        recodedArray = vcat(recodedArray, img_singular)
    end

    # Normalizing the singular values
    recodedArray =  eachcol(recodedArray) ./ norm.(eachcol(recodedArray))
    # writing the array as a .csv file
    filename = joinpath(pathtoimage, "image_recoded.csv")
    CSV.write(filename,  DataFrame(recodedArray), writeheader=true)
end

"""
recodeaudio(filepath,ncep)

This function reads audio files (.wav files) in the file path and generates an excel file (.csv) with
    each row corresoponding to a particular audio file and each column corresoponding to the mean
    cepstral coefficient. Only the first ncep coefficients are computed.

    USAGE:
    ```julia
    recodeaudio(path,3)
    ```
"""
function recodeaudio(filepath)

    # creating a list of .wav files in the folder and an empty array to store the mean cepstral coefficients
    audiolist = glob("*.wav",filepath)
    recodedArray = Array{Float64}(undef, 0, 26)

    @showprogress for files in audiolist
        x,fs = wavread(files)

        # mfcc is a function from the MFCC package and computes the Mel Frequency Cepstral Coefficients
        # mfcc returns 3 values: [1] a matrix of numcep columns with for each speech frame a row of MFCC coefficient,
        # [2] the power spectrum, and [3] a dictionary containing the information about the parameters
        # mfcc(x,fs, numcep = ncep)[1]...here the [1] specifies that we accesing only the cep coefficients matrix.
        # From the above step, the mean cep is calculated, then the first 13 deltas are extracted
        # from the signal and mean_cep and deltas are horizontally concatinated to form the feature vector
        deltavals = deltas(x)[1:13]'
        cep = collect(mfcc(x,fs)[1])
        features = [mean(cep,dims=1) deltavals] # mean(_,dims=1) returns the column means
        recodedArray = vcat(recodedArray, features)
    end
    replace!(recodedArray, Inf =>0)
    replace!(recodedArray, -Inf =>0)
    replace!(recodedArray, NaN =>0)
    U,S,V = svd(recodedArray)
    S = S ./norm(S)
    screeplot(S)
    # ============================================================================================================

    # After examining the scree plot, the user decides the no. of singular values
    n_singularvlas = input("No. of Features (due to bug in the code that reads user inputs, you might have to enter the no twice, if the program didn't run first time)")
    # Saving the recodedArray as a .csv file
    featurematrix = U[:,1:n_singularvlas]
    filename = joinpath(filepath, "audio_recoded.csv")

    # Normalizing values before writing onto excel sheet.
    #recodedArray = eachcol(recodedArray) ./ norm(eachcol(recodedArray))
    CSV.write(filename,  DataFrame(featurematrix), writeheader=true)
end

"""
recodetext(pathtotxt, n_singularvals)

This function reads PDF files in the pathtotxt and generates an excel file (.CSV) which contains the
    document by term matrix reconstructed using only the first n_singularvals. If the passed n_singularvals
    is a larger number than the singular values present, all singular values are used.

    USAGE:
    ```julia
    recodetext(pathtotxt,3)
    ```
"""
function recodetext(pathtotxt)

    # PDF parsing utilities for this function are provided by the Taro julia package
    # and Taro requires to initialize a java virtual machine. Taro.init() achieves that
    # Creating a list of PDF files in the path.
    files = glob("*.pdf", pathtotxt)

    # Initializing an empty array to store the text from each document
    docs = Array{StringDocument{String},1}[]

    # This block of code extracts the text from the PDFs and creates a corpus
    getTitle(t) = TextAnalysis.sentence_tokenize(Languages.English(), t)[1]
   @showprogress for i in 1:length(files)
        meta,txt = Taro.extract(files[i]);
        txt = replace(txt, '\n' => ' ')
        title = getTitle(txt)
        dm = TextAnalysis.DocumentMetadata(Languages.English(), title, "", meta["Creation-Date"] )
        doc = StringDocument(txt, dm)
        docs = vcat(docs, doc)
    end
    crps = Corpus(docs)

    # cleaning the corpus. getting rid of non-words, punctuations etc...
    prepare!(crps, strip_non_letters | strip_punctuation | strip_case | strip_stopwords)
    TextAnalysis.stem!(crps)
    update_lexicon!(crps)
    update_inverse_index!(crps)

    # The corpus is used to generates a document term matrix,
    # and we perform singluar value decomposition on this matrix.
    U,S,V = @pipe crps |> DocumentTermMatrix(_) |> dtm(_, :dense) |> svd(_)

    S = S ./norm(S)
    p1 = screeplot(S)
    display(p1)
    n_singularvlas = input("No. of Features (due to bug in the code that reads user inputs, you might have to enter the no twice, if the program didn't run first time)")
    # Reconstructing the document term matrix with just the first n_singularvals.
   # DTM = U[:,1:n_singularvals]*diagm(S[1:n_singularvals])*Vt[:,1:n_singularvals]'
    # saving the reconstructed document term matrix as a .csv file
    filename = joinpath(pathtotxt, "reconstructed_document_term_matrix.csv")
    CSV.write(filename,  DataFrame(U[:,1:n_singularvlas]), writeheader=true)
end

# ===================================
function recodeaudio2(filepath)

    # creating a list of .wav files in the folder and an empty array to store the mean cepstral coefficients
    audiolist = glob("*.wav",filepath)
    recodedArray = Array{Float64}(undef, 0, 26)

    @showprogress for files in audiolist
        x,fs = wavread(files)

        # mfcc is a function from the MFCC package and computes the Mel Frequency Cepstral Coefficients
        # mfcc returns 3 values: [1] a matrix of numcep columns with for each speech frame a row of MFCC coefficient,
        # [2] the power spectrum, and [3] a dictionary containing the information about the parameters
        # mfcc(x,fs, numcep = ncep)[1]...here the [1] specifies that we accesing only the cep coefficients matrix.
        # From the above step, the mean cep is calculated, then the first 13 deltas are extracted
        # from the signal and mean_cep and deltas are horizontally concatinated to form the feature vector
        deltavals = deltas(x)[1:13]'
        cep = collect(mfcc(x,fs)[1])
        features = [mean(cep[1],dims=1) deltavals] # mean(_,dims=1) returns the column means
        recodedArray = vcat(recodedArray, features)
    end
    replace!(recodedArray, Inf =>0)
    replace!(recodedArray, -Inf =>0)
    replace!(recodedArray, NaN =>0)
    U,S,V = svd(recodedArray)
    S = S ./norm(S)
    screeplot(S)
    # ============================================================================================================

    # After examining the scree plot, the user decides the no. of singular values
    n_singularvlas = input("No. of Features (due to bug in the code that reads user inputs, you might have to enter the no twice, if the program didn't run first time)")
    # Saving the recodedArray as a .csv file
    featurematrix = U[:,1:n_singularvlas]
    filename = joinpath(filepath, "audio_recoded.csv")

    # Normalizing values before writing onto excel sheet.
    #recodedArray = eachcol(recodedArray) ./ norm(eachcol(recodedArray))
    CSV.write(filename,  DataFrame(featurematrix), writeheader=true)
end


# Tests for the conditional mutual information function

entropy3 = get_entropy(arr3)
entropy13 = get_entropy(arr1, arr3)
entropy23 = get_entropy(arr2, arr3)
entropy123 = get_entropy(arr1, arr2, arr3)
cmi = entropy13 + entropy23 - entropy123 - entropy3

# Test conditional mutual information
@test get_conditional_mutual_information(arr1, arr2, arr3) ≈ cmi
println("Conditional mutual information passed.")

# Test conditional mutual information formula
@test apply_conditional_mutual_information_formula(entropy13, entropy23, entropy123, entropy3) ≈ cmi
println("Conditional mutual information formula passed.")


# Copyright 2017 Raytheon BBN Technologies
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the license at
#
#		http://www.apache.org/licenses/LICENSE-2.0

using LsqFit

## TODO: Devise a way to get specify a fit as a string...
## TODO: Fix types...

"""
FitResult type that stores the output of a fit.
"""
struct FitResult
    fit_params::Dict{String,Float64}
    sq_error::Float64
    Nσ::Float64
    errors::Dict{String,Float64}
    fit_curve::Function
    model_str::String
end

"""
`generic_fit(xpts, ypts, model, initial_guess, fit_params, model_string; yvars=[])`
Helper function to return a fit to a model.
# Arguments
* xpts: Independent fit variable.
* ypts: Dependent fit variable.
* model: Fit function(x, p) of data x and parameter vector p.
* initial_guess: Initial guess of fit parameters.
* fit_params::function(Array{Float64} => Dict{String, Float64}): Result dict of fit parameters.
* model_string::String: String indicating fit function.
* yvars: Variance of y points. Defaults to empty vector.
# Returns
* A FitResult object for the specified fit.
"""
function generic_fit(xpts, ypts, model, initial_guess, fit_params, model_string::String; yvars=[])
    @assert length(xpts) == length(ypts) "X and Y data length must match."
    #try
        if isempty(yvars)
          result = curve_fit(model, xpts, ypts, initial_guess)
          errors = margin_error(result)
          sq_error = sum(((model(xpts, result.param) - ypts) .^ 2))
        else
          @assert length(ypts) == length(yvars) "Y data and Y variance lengths must match."
          result = curve_fit(model, xpts, ypts, 1 ./ sqrt(yvars), initial_guess)
          errors = margin_error(result)
          sq_error = sum(((model(xpts, result.param) - ypts) .^ 2) ./ yvars)
        end
    #catch
    #    result =
    # Compute badness of fit:
    # Under the null hypothesis (that the model is valid and that the observations
    # do indeed have Gaussian statistics), the mean squared error is χ² distributed
    # with `dof` degrees of freedom. We can quantify badness-of-fit in terms of how
    # far the observed MSE is from the expected value, in units of σ = 2dof (the expected
    # standard deviation for the χ² distribution)
    dof = length(xpts)-length(initial_guess)
    Nσ = sq_error/sqrt(2dof) - dof/sqrt(2dof)
    return FitResult( fit_params(result.param),
              sq_error,
              Nσ,
              fit_params(errors),
              xpts->model(xpts,result.param),
              model_string)
end

""" `fit_t1(xpts, ypts, yvars=[])`
Fit to exponential T decay y = a exp(-t/T) + b.
"""
function fit_t1(xpts, ypts, yvars=[])
    T1_fit_dict(p) = Dict("a" => p[1], "T" => p[2], "b" => p[3])
    model(t, p) = p[1]*exp.(-t ./ p[2]) .+ p[3]
    p_guess = [ypts[1]-ypts[end], xpts[end]/3., ypts[end]]
    return generic_fit(xpts, ypts, model, p_guess, T1_fit_dict,
              "a * exp(-t/T) + b", yvars=yvars)
end

""" `fit_double_exp(xpts, ypts, yvars=[])`
Fit to double exponential decay y = a exp(-t/Ta) + b exp(-t/Tb) + c.
"""
function fit_double_exp(xpts, ypts, yvars=[])
    T1_fit_dict(p) = Dict("a" => p[1], "Ta" => p[2], "b" => p[3], "Tb" => p[4], "c" => p[5])
    model(t, p) = p[1]*exp.(-t ./ p[2]) .+ p[3]*exp.(-t ./ p[4]) .+ p[5]
    p_guess = [0.5*(ypts[1]-ypts[end]), xpts[end]/3., 0.5*(ypts[1]-ypts[end]), xpts[end]/3., ypts[end]]
    return generic_fit(xpts, ypts, model, p_guess, T1_fit_dict,
              "a * exp(-t/Ta) + b * exp(-t/Tb) + c", yvars=yvars)
end

""" `fit_line(xpts, ypts, yvars=[])`
Fit to linear model y = ax + b.
"""
function fit_line(xpts, ypts, yvars=[])
    model(t, p) = p[1]*t .+ p[2]
    p_guess = [-ypts[findmin(abs.(xpts))[2]] / xpts[findmin(abs.(ypts))[2]], ypts[findmin(abs.(xpts))[2]]]
    return generic_fit(xpts, ypts, model, p_guess,
    x->Dict("a" => x[1], "b"=>x[2]), "ax + b", yvars=yvars)
end

""" `fit_ramsey(xpts, ypts, yvars=[])`
Fit data to a Ramsey decay of the form a*exp(-t ./ T).*cos(2πf .* t + ϕ) + b
with optional weights for each point.
"""
function fit_ramsey(xpts, ypts, yvars=[])

    fit_dict1(p) = Dict("a"=>p[1], "T"=>p[2], "f"=>p[3], "ϕ"=>p[4], "b"=>p[5])
    model(t, p) = p[1]*exp.(-t ./ p[2]).*cos.(2π*p[3] .* t .+ p[4]) .+ p[5]
    # Use KT estimation to get a guess for the fit
    freqs,Ts,amps = KT_estimation(ypts .- mean(ypts), xpts[2]-xpts[1], 1)
    p_guess = [abs(amps[1]), Ts[1], freqs[1], angle(amps[1]), mean(ypts)];

    return generic_fit(xpts, ypts, model, p_guess,
              fit_dict1, "a*exp(-t/T).*cos(2πf .* t + ϕ) + b)", yvars=yvars)
end

""" `fit_twofreq_ramsey(xpts, ypts, yvars=[])`

Fit the Ramsey data to a one-frequency model a*exp(-t ./ T).*cos(2πf .* t + ϕ) + b
or a two-frequency model
a₁*exp(-t ./ T₁).*cos(2πf₁ .* t + ϕ₁) + a₂*exp(-t ./ T₂).*cos(2πf₂ .* t + ϕ₂) + b
optionally taking variances into account, and report which should be
used based on Akaike's Information Criterion (AIC). This assumes
Gaussian statistics for each observation.
"""
function fit_twofreq_ramsey(xpts, ypts, yvars=[])
    fit_dict2(p) = Dict("a₁"=>p[1], "T₁"=>p[2], "f₁"=>p[3], "ϕ₁"=>p[4],
    "a₂"=>p[5], "T₂"=>p[6], "f₂"=>p[7], "ϕ₂"=>p[8],
    "b"=>p[9])
    model2(t, p) = ( p[1] * exp.(-t ./ p[2]) .* cos.(2π * p[3] .*t .+ p[4]) .+
    p[5] * exp.(-t ./ p[6]) .* cos.(2π * p[7] .*t .+ p[8]) .+ p[9] )

    #Use KT estimation to get a guess for the fit
    freqs,Ts,amps = KT_estimation(ypts.-mean(ypts), xpts[2].-xpts[1], 2)
    phases = angle.(amps)
    amps = abs.(amps)
    p_guess = [1, Ts[1], freqs[1], phases[1], 1, Ts[2], freqs[2], phases[2], mean(ypts)]

    fitresult2 = generic_fit(xpts, ypts, model2, p_guess, fit_dict2,
    "a₁*exp(-t ./ T₁).*cos(2πf₁ .* t .+ ϕ₁) .+ a₂*exp(-t ./ T₂).*cos(2πf₂ .* t .+ ϕ₂) .+ b",
    yvars=yvars)

    fitresult1 = fit_ramsey(xpts, ypts, yvars)

    # AIC computation
    k1 = 5
    k2 = 9
    corr(k,n) = (k+1)*(k+1)/(n-k-2)
    aicc(e,k,n) = 2 * k + e + corr(k,n)

    aic = aicc(fitresult2.sq_error,k2,length(xpts)) - aicc(fitresult1.sq_error,k1,length(xpts))

    println(aic)
    return fitresult2 #(aic > 0) ? fitresult2 : fitresult1
end

"""`fit_sin(xpts, ypts, yvars)`
Fit to a sine a*sin(2πf t) + b
"""
function fit_sin(xpts, ypts, yvars=[])
    model(t, p) = p[1]*sin.(2π*p[2] .* t) .+ p[3]
    # Use KT estimation to get a guess for the fit
    freqs,Ts,amps = KT_estimation(ypts, xpts[2]-xpts[1], 2)
    idx = argmax(abs.(amps))
    p_guess = [abs(amps[idx]), Ts[idx], freqs[idx], mean(ypts)];
    return generic_fit(xpts, ypts, model, p_guess,
              x->Dict("a"=>x[1], "f"=>x[2], "b"=>x[3]),
              "a*sin(2πf t) + b", yvars=yvars)
end

## TODO: Change this to use generic_fit
"""`fit_photon_ramsey(xpts, ypts, params)`
Fit function in McClure et al., Phys. Rev. App. 2016. Params are:
1. cavity decay rate κ (MHz)
2. detuning Δ (MHz)
3. dispersive shift 2χ (MHz)
4. Ramsey decay time T₂* (μs)
5. exp(-t_meas/T₁) (μs)
6. initial qubit state (0/1)
"""
function fit_photon_ramsey(xpts, ypts, params)
    params[1:3]*=2*pi #convert to angular frequencies
    model_0(t, p) = (-imag.(exp.(-(1/params[4]+params[2]*1im).*t + (p[1]-p[2]*params[3]*(1-exp.(-((params[1] + params[3]*1im).*t)))/(params[1]+params[3]*1im))*1im)))
    function model(t, p)
        if params[6] == 1
            return params[5]*model_0(t, p) + (1-params[5])*model_0(t, [p[1]+pi; p[2:end]])
        else
            return model_0(t, p)
        end
    end
    p_guess = [0., 1.]
    result = curve_fit(model, xpts, ypts, p_guess)
    errors = margin_error(result)
    xfine = range(xpts[1], stop=xpts[end],length=1001)
    fit_curve = (xfine, model(xfine, result.param))
    return (result.param[2], errors[2], fit_curve)
end

"""
analyzeRB(ypts, seqlengths; purity=False)
analyzeRB Analyzes a randomized benchmarking experiment
seqlengths, example: [4,8,16,32,64,128,256,384,512,640,768,896,1024,1152,1280,1408,1536,1664]
if purity = true, fit for incoherent errors. See J.J. Wallman et al., New J. Phys. 17, 113020 (2015)
Assuming data of the form <Z>s1	, <Z>s2, ...., <Z>sN, <X>s1, ..., <X>sN, ..., <Y>s1, ..., <Y>sN
with seqlengths = [s1, s2, ..., sN]
"""
function analyzeRB(ypts, seqlengths; purity=false)
    # figure out how many sequences of each length we have
    num_repeats = length(ypts) ÷ length(seqlengths) ÷ (purity ? 3 : 1)
    xpts = repeat(seqlengths, inner=num_repeats)
    if purity
        # compute length of Bloch vector
        data = vec(sum(reshape(ypts, (length(ypts) ÷ 3),3).^2,  2))
    else
        # otherwise convert <Z> to prob of 0
        data = .5 * (1 + vec(ypts))
    end
    model(n, p) = p[1] * (1-p[2]).^n + p[3]
    fit = curve_fit(model, xpts-purity, data, [0.5, .01, 0.5]) #fit to ...^(n-1) for purity
    xfine = range(seqlengths[1],stop=seqlengths[end],length=1001)
    fit_curve = (xfine, model(xfine, fit.param))
    errors = margin_error(fit)
    if purity
        println("ϵ_inc = $(0.5*(1-sqrt(1-fit.param[2]))*100)")
    else
        println("ϵ = $(fit.param[2]/2*100)")
    end
    return (xpts, data, fit.param, fit_curve, errors)
end

"""`fit_RB(xpts, ypts, yvars=[])`
Fit to RB decay a*(1-b)^x + c
"""
function fit_RB(xpts, ypts, yvars=[])
    RB_fit_dict(p) = Dict("a" => p[1], "b" => p[2], "c" => p[3])
    model(n, p) = p[1] * (1-p[2]).^n + p[3]
    p_guess = [0.5, .01, 0.5]
    return generic_fit(xpts, ypts, model, p_guess, RB_fit_dict, "a * exp(-n/b) + c", yvars=yvars)
end

function calculate_residueCR(params, t, rvec)
    rx        = params[1]; # Rabi frequency along x
    ry        = params[2]; # Rabi frequency along y
    delta     = params[3]; # detuning
    x0        = params[4]; # initial x
    y0        = params[5]; # initial y
    z0        = params[6]; # initial z
    omega = sqrt(rx^2+ry^2+delta^2);
    x = (rx*delta*(-z0+x0*rx+y0*ry)+(rx*(z0*delta-y0*ry)+x0*(delta^2+ry^2)).*cos.(t*omega)+omega*(y0*delta+z0*ry).*sin.(t*omega))/omega^2;
    y = (ry*delta*(-z0+x0*rx+y0*ry)+(ry*(z0*delta-x0*rx)+y0*(delta^2+rx^2)).*cos.(t*omega)-omega*(x0*delta+z0*rx).*sin.(t*omega))/omega^2;
    z = (delta*(z0*delta-x0*rx-y0*ry)+(x0*delta*rx+y0*delta*ry+z0*(rx^2+ry^2)).*cos.(t*omega)+omega*(y0*rx-x0*ry).*sin.(omega*t))/omega^2;
    rmodel = [x,y,z];
    diffvec = rvec-rmodel;
    return sum(hcat(diffvec...).^2), rmodel
end

function fit_CR_hamiltonian(xpoints, datavec, params_estimate = [5,5,0,0,0,1])
    result = optimize(x-> calculate_residueCR(x, xpoints, datavec)[1], params_estimate)
    return Optim.minimizer(result)
end