export forw2spot, spot2forw

## capital market -----------------------------------------------
"""
`getyield(τ, stock::Stock)`

calculates the yield of `stock` during year `τ`
"""
function getyield(τ, stock::Stock)
  if τ > 1
    return stock.x[τ] / stock.x[τ - 1] - 1
  elseif τ == 1
    return stock.x[τ] / stock.x_0 - 1
  else
    return stock.yield_0
  end
end

"""
`getyield(τ, rfr::RiskFreeRate)`

calculates the yield of `rfr` during year `τ`
"""
function getyield(τ, rfr::RiskFreeRate)
  if τ >= 1
    return rfr.x[τ]
  else
    return rfr.yield_0
  end
end

"""
`forw2spot(f::Vector{Float64})`

calculates the spot rate `s` rate from the forward rate `f`:

  `(1+f[1])(1+f[2])...(1+f[n]) = (1+s[n])^n`
"""
forw2spot(f::Vector{Float64}) =
  cumprod(1 .+ f) .^ (1 ./ collect(1:length(f))) .- 1

"""
`spot2forw(s::Vector{Float64})`

calculates the forward rate `f` rate from the spot rate `s`:

  `(1 + s[n-1])^(n-1) * (1+f[n]) = (1+s[n])^n`
"""
function spot2forw(s::Vector{Float64})
  f = zeros(Float64, length(s))
  for 𝑛 ∈ length(s):-1:2
    f[𝑛] = (1 + s[𝑛])^𝑛 / (1 + s[𝑛-1])^(𝑛 - 1) -1
  end
  f[1] = s[1]
  return f
end

## investments --------------------------------------------------
"""
`project!(τ::Int, mv_boy::Float64, invest::Invest)`

project `invest::Invest` one year for a given
initial market value.

**Changed**:  `invest`
"""
function project!(τ::Int, mv_boy::Float64, invest::Invest)
  invest.mv[τ] = (1 + getyield(τ, invest.proc)) * mv_boy
end

"""
`project!(τ::Int, mv_bop_total::Float64, ig::InvestGroup)`

project `ig::InvestGroup` one year for a given
initial market value.

**Changed**:  `ig`
"""
function project!(τ::Int, mv_bop_total::Float64, ig::InvestGroup)
  mv_bop = ig.alloc.total[τ] * mv_bop_total
  ig.mv[τ] = 0
  for (𝑖, 𝑖𝑔_𝑖𝑛𝑣𝑒𝑠𝑡) ∈ enumerate(ig.investments)
    project!(τ, ig.alloc.all[τ, 𝑖] * mv_bop, 𝑖𝑔_𝑖𝑛𝑣𝑒𝑠𝑡)
    ig.mv[τ] += 𝑖𝑔_𝑖𝑛𝑣𝑒𝑠𝑡.mv[τ]
  end
  ig.cost.total[τ] =
    ig.cost.abs[τ] * ig.cost.cum_infl_abs[τ] +
    mv_bop * ig.cost.rel[τ] * ig.cost.cum_infl_rel[τ]
end

"""
`alloc!(τ, cap_mkt::CapMkt, invs::InvPort)`

Dynamically re-allocate investments `invs::InvPort` at the
beginning of year `τ`

**Changed**:  `invs`
"""
function alloc!(τ, cap_mkt::CapMkt, invs::InvPort)
  if τ > 1
    invs.igs[:IGStock].alloc.total[τ] =
      0.5 * (1 - exp( -max(0, dynstateavg(τ, cap_mkt))))
    invs.igs[:IGCash].alloc.total[τ] =
      1 - invs.igs[:IGStock].alloc.total[τ]
    ## we leave the allocations within each group unchanged:
    for 𝑠𝑦𝑚𝑏 ∈ [:IGCash, :IGStock]
      for 𝑖 = 1:size(invs.igs[𝑠𝑦𝑚𝑏].alloc.all, 2)
        invs.igs[𝑠𝑦𝑚𝑏].alloc.all[τ, 𝑖] =
          invs.igs[𝑠𝑦𝑚𝑏].alloc.all[τ-1, 𝑖]
      end
    end
  end
end

"""
`project!(τ::Int, mv_boy::Float64, invs::InvPort)`

Project `invs::InvPort`  one year for a given
initial market value.

**Changed**:  `invs`
"""
function project!(τ::Int, mv_boy::Float64, invs::InvPort)
  invs.mv[τ] = 0.0
  invs.cost[τ] = 0.0
  invs.mv_boy[τ] = mv_boy
  for 𝑖𝑔 ∈ values(invs.igs)
    project!(τ, mv_boy, 𝑖𝑔)
    invs.mv[τ] += 𝑖𝑔.mv[τ]
    invs.cost[τ] += 𝑖𝑔.cost.total[τ]
  end
  invs.yield[τ] = invs.mv[τ] / mv_boy - 1
end

## insurance liabilities  ---------------------------------------
"""
`premium(ins_sum, rfr, prob, β, λ)`

Calculate the premium of a product
"""
function premium(ins_sum, rfr, prob, β, λ)
  lx_boy = [1; cumprod(prob[:px])[1:end-1]]
  v_eoy = 1 ./ cumprod(1 .+ rfr)
  v_boy = [1; v_eoy[1:end-1]]
  num =
    sum(lx_boy .* ins_sum .*
        (v_boy .* λ[:boy] .* λ[:cum_infl] ./ (1 .+ λ[:infl]) .+
           v_eoy .* (λ[:eoy] .* λ[:cum_infl] .+
                       prob[:px] .* β[:px] .+
                       prob[:qx] .* β[:qx])
         ))
  denom =
    sum(lx_boy .* β[:prem] .*
        (v_boy - v_eoy .* prob[:sx] .* β[:sx]))
  return num / denom
end

"""
`tpgrec(τ, tpg, rfr, prob, β, λ)`

One year backward recursion formula for technical provisions
of guaranteed benefits
"""
function tpgrec(τ, tpg, rfr, prob, β, λ)
  disc = 1/(1 + rfr[τ + 1])
  -β[τ + 1, :prem] +
    λ[τ + 1, :boy] *
    λ[τ + 1, :cum_infl] / (1 + λ[τ + 1, :infl]) +
    disc * (λ[τ + 1, :eoy] * λ[τ + 1, :cum_infl] +
              prob[τ + 1, :qx] * β[τ + 1, :qx] +
              prob[τ + 1, :sx] * β[τ + 1, :sx] +
              prob[τ + 1, :px] * (β[τ + 1, :px] + tpg))
end

"""
`tpg(τ, rfr, prob, β, λ)`

Technical provisions of guaranteed benefits
at the end of year `τ`
"""
function tpg(τ, rfr, prob, β, λ)
  dur = nrow(β)
  res = 0.0
  if τ >= dur
    return 0.0
  else
    for 𝑠 ∈ (dur-1):-1:τ
      res = tpgrec(𝑠, res, rfr, prob, β, λ)
    end
    return res
  end
end

"""
`tpg(τ, rfr, mp)`

Best estimate guaranteed technical provisions
of a model point `mp` at the end of year `τ`
"""
tpg(τ, rfr, mp) = tpg(τ, rfr, mp.prob, mp.β, mp.λ)

"""
`tpgfixed(τ, rfr, fixed_cost_gc::Vector)`

Technical provisions at the end of year `τ`
for future fixed (going concern) costs
"""
function tpgfixed(τ, rfr, fixed_cost_gc::Vector)
  dur = length(fixed_cost_gc)
  if dur < τ + 1
    return 0.0
  else
    disc = cumprod(1 ./ (1 .+ rfr[(τ + 1):dur]))
    return fixed_cost_gc[(τ + 1):dur] ⋅ disc
  end
end

## other liabilities --------------------------------------------
"""
`pv(τ::Int, cap_mkt::CapMkt, debt::Debt)`

Present value of `debt` at the end of year `τ`.
Any servicing of this debt during year `τ` has occured beforehand
"""
function pv(τ::Int, cap_mkt::CapMkt, debt::Debt)
  ## calculate pv at the end of the year after servicing debt
  if (debt.τ_init > τ) | (debt.τ_mat <= τ)
    return 0.0
  else
    p_v = debt.nominal
    for 𝑠 ∈ (debt.τ_mat - 1) : -1 : τ
      p_v = (debt.coupon + p_v) / (1 + cap_mkt.rfr.x[𝑠 + 1])
    end
    return p_v
  end
end

"""
`pv(τ::Int, cap_mkt::CapMkt, l_other::LiabOther)`

present value of other liabilities `l_other` at the end of year
`τ`. Any servicing of debt within this portfolio during year `τ`
has occured beforehand
"""
function pv(τ::Int, cap_mkt::CapMkt, l_other::LiabOther)
  p_v = 0.0
  for 𝑑𝑒𝑏𝑡 ∈ l_other.subord
    p_v += pv(τ, cap_mkt, 𝑑𝑒𝑏𝑡)
  end
  return p_v
end

"""
`paycoupon(τ::Int, debt::Debt)`

Coupon payment for `debt` at the end of year `τ`
"""
paycoupon(τ::Int, debt::Debt) =
  (debt.τ_init <= τ <= debt.τ_mat ? debt.coupon : 0.0)

"""
`paycoupon(τ::Int, l_other::LiabOther)`

Coupon payment for all debts within the portfolio of
other liabilities `l_other` at the end of year `τ`
"""
function paycoupon(τ::Int, l_other::LiabOther)
  pay = 0.0
  for 𝑑𝑒𝑏𝑡 ∈ l_other.subord
    pay += paycoupon(τ, 𝑑𝑒𝑏𝑡)
  end
  return pay
end

"""
`payprincipal(τ::Int, debt::Debt)`

Payment of the principal of `debt` at the end of year `τ`,
if the debt matures at this point in time
"""
payprincipal(τ::Int, debt::Debt) =
  (τ == debt.τ_mat ? debt.nominal : 0.0)

"""
`payprincipal(τ::Int, l_other::LiabOther)`

Payment of the total principal of all debts within the
portfolio of other liabilities `l_other`, which mature at
the end of year `τ`
"""
function payprincipal(τ::Int, l_other::LiabOther)
  pay = 0.0
  for 𝑑𝑒𝑏𝑡 ∈ l_other.subord
    pay += payprincipal(τ, 𝑑𝑒𝑏𝑡)
  end
  return pay
end

"""
`getloan(τ::Int, debt::Debt)`

Get the nominal of `debt` at time `τ`, if it has been
taken out at this point in time, otherwise return 0.0
"""
getloan(τ::Int, debt::Debt) =
  (τ == debt.τ_init ? debt.nominal : 0.0)

"""
`getloan(τ::Int, l_other::LiabOther)`

Get the total nominal of all debts within the
portfolio of other liabilities `l_other`,
which are taken out at the beginning of year `τ`
"""
function getloan(τ::Int, l_other::LiabOther)
  nominal = 0.0
  for 𝑑𝑒𝑏𝑡 ∈ l_other.subord
    nominal += getloan(τ, 𝑑𝑒𝑏𝑡)
  end
  return nominal
end

"""
`goingconcern(debts::Vector{Debt}, Δgc::Vector{Float64})`

Calculates a vector of debts from an existing vector of debts
according to the going concern assumption. The total initial debt
is the same, but the new debt vectors mature earlier so that the
total nominal decreases according to the going concern factors.
We do not input the factors `gc` directly but their year on year
differences `Δgc`.
"""
function goingconcern(debts::Vector{Debt}, Δgc::Vector{Float64})
  new_debt_vec = Array{Debt}(undef, 0)
  for 𝑑𝑒𝑏𝑡 ∈ debts
    if 𝑑𝑒𝑏𝑡.nominal > 0.0
      τ_init = max(1, 𝑑𝑒𝑏𝑡.τ_init)
      diff_nom = -Δgc * 𝑑𝑒𝑏𝑡.nominal
      for 𝜏 ∈ τ_init:𝑑𝑒𝑏𝑡.τ_mat
        t = 𝑑𝑒𝑏𝑡.t_init + 𝜏 - 𝑑𝑒𝑏𝑡.τ_init
        push!(new_debt_vec,
              Debt(𝑑𝑒𝑏𝑡.name,
                   𝑑𝑒𝑏𝑡.t_init,
                   t,
                   𝑑𝑒𝑏𝑡.τ_init,
                   𝜏,
                   diff_nom[𝜏],
                   𝑑𝑒𝑏𝑡.coupon * diff_nom[𝜏] / 𝑑𝑒𝑏𝑡.nominal))
      end
    end
  end
  return(new_debt_vec)
end

"""
`goingconcern!(l_other::LiabOther, Δgc::Vector{Float64})`

Transforms a portfolio of other liabilities `l_other::LiabOther`
according to the going concern assumption. We do not input the
factors `gc` directly but their year on year differences `Δgc`.

**Changed**: `l_other`
"""
function goingconcern!(l_other::LiabOther, Δgc::Vector{Float64})
  l_other.subord = goingconcern(l_other.subord, Δgc)
end

## dynamics -----------------------------------------------------
"""
`dynstate(τ, cap_mkt::CapMkt)`

indicator for the state of the economy at the end of year `τ`
"""
dynstate(τ, cap_mkt::CapMkt) =
  getyield(τ, cap_mkt.stock) / max(getyield(τ, cap_mkt.rfr), eps()) - 1

"""
`dynstateavg(τ, cap_mkt::CapMkt)`

Two year average of the indicator for the state of the economy
at the end of year `τ`.
"""
dynstateavg(τ, cap_mkt::CapMkt) =
  0.5 * (getyield(τ - 1, cap_mkt.stock) /
           max(getyield(τ-1, cap_mkt.rfr), eps())
         + getyield(τ, cap_mkt.stock) /
           max(getyield(τ, cap_mkt.rfr), eps())) - 1

"""
`bonusrate(yield_eoy, rfr_price, bonus_factor)`

Helper function for dynamic bonus rate declaration
"""
bonusrate(yield_eoy, rfr_price, bonus_factor) =
  max(bonus_factor * (yield_eoy - rfr_price), 0.0)

"""
`bonusrate(τ, yield_eoy, mp::ModelPoint, dyn)`

Dynamic bonus rate declaration for a model point at the end
of year `τ`
"""
bonusrate(τ, yield_eoy, mp::ModelPoint, dyn) =
  bonusrate(yield_eoy,
            ( τ == 0 ? mp.rfr_price_0 : mp.rfr_price[τ]),
            dyn.bonus_factor)

"""
`biquotient(τ, yield_eoy, cap_mkt, invs, mp, dyn)`

Indicator for bonus rate expectation at the end of year `τ`
"""
function biquotient(τ, yield_eoy, cap_mkt, invs, mp, dyn)
  if τ ≤ mp.dur
    ind_bonus =
      getyield(τ, cap_mkt.stock) /
      max(0.0,
          bonusrate(τ - 1, yield_eoy, mp, dyn) + mp.rfr_price[τ])
    ind_bonus_hypo =
      getyield(0, cap_mkt.stock) /
      max(eps(), mp.bonus_rate_hypo + mp.rfr_price_0)
    return ind_bonus / ind_bonus_hypo
  else
    return 0.0
  end
end

"""
`δsx(τ, cap_mkt, invs, mp, dyn)`

Dynamic lapse probability factor to adjust the initial estimate
"""
function δsx(τ, cap_mkt, invs, mp, dyn)
  yield_eoy =
    invs.igs[:IGStock].alloc.total[τ] * getyield(τ, cap_mkt.stock) +
    invs.igs[:IGCash].alloc.total[τ] * getyield(τ, cap_mkt.rfr)
  if τ - 1 > mp.t_start
    bi_quot =  biquotient(τ, yield_eoy, cap_mkt, invs, mp, dyn)
  else
    bi_quot = 1
  end
  state_quot = dynstate(τ, cap_mkt) / dynstate(0, cap_mkt)
  δ_SX = 1.0
  if state_quot < 0.5
    δ_SX += 0.15
  elseif  state_quot > 2.0
    δ_SX -= 0.15
  end
  δ_SX += 0.25 * min(4.0, max(0.0, bi_quot - 1.2))
  return δ_SX
end

"""
`freesurp(dyn, invest_pre, liab)`

Helper function for the free surplus calculation
"""
freesurp(dyn, invest_pre, liab) =
 max(0, invest_pre - (1 + dyn.quota_surp) * liab)

"""
`freesurp(τ, proj::Projection, dyn)`

Free surplus for the dynamic dividend declaration
"""
freesurp(τ, proj::Projection, dyn) =
  if τ == 1
      freesurp(dyn,
               proj.val_0[1, :invest],
               proj.val_0[1, :tpg] + proj.val_0[1, :l_other])
  else
      freesurp(dyn,
               proj.val[τ-1, :invest],
               proj.val[τ-1, :tpg] + proj.val[τ-1, :l_other])
  end

"""
`update!(τ, proj::Projection, dyn::Dynamic)`

Update dynamic parameters
"""
function update!(τ, proj::Projection, dyn::Dynamic)
  if τ == 1
    dyn.free_surp_boy[τ] =
      freesurp(dyn,
               proj.val_0[1, :invest],
               proj.val_0[1, :tpg] + proj.val_0[1, :l_other])
  else
    dyn.free_surp_boy[τ] =
      freesurp(dyn,
               proj.val[τ-1, :invest],
               proj.val[τ-1, :tpg] + proj.val[τ-1, :l_other])
  end
end

## cashflow projection  -----------------------------------------
"""
`val0!(cap_mkt::CapMkt, invs::InvPort, liabs::LiabIns,
  l_other::LiabOther, proj::Projection)`

Valuation at time `t_0`

**Changed:** `proj::Projection`
"""
function val0!(cap_mkt::CapMkt,
               invs::InvPort,
               liabs::LiabIns,
               l_other::LiabOther,
               proj::Projection)
  proj.val_0[1, :invest] = invs.mv_0
  for 𝑚𝑝 ∈ liabs.mps
    if 0 <= 𝑚𝑝.dur
      proj.val_0[1, :tpg] += tpg(0, cap_mkt.rfr.x, 𝑚𝑝)
    end
  end
  proj.val_0[1, :l_other] = pv(0, cap_mkt, l_other)
  proj.val_0[1, :surplus] =
    proj.val_0[1, :invest] -
    proj.val_0[1, :tpg] -
    proj.val_0[1, :l_other]
end

"""
`projectboy!(τ, proj::Projection, liabs::LiabIns)`

Project one year, update values at the beginning of the year `τ`

**Changed:** `proj::Projection`, `liabs::LiabIns  (liabs.mp)`
"""
function projectboy!(τ, proj::Projection, liabs::LiabIns)
  proj.cf[τ, :prem] = 0.0
  proj.cf[τ, :λ_boy] = 0.0
  for 𝑚𝑝 ∈ liabs.mps
    if τ <= 𝑚𝑝.dur
      𝑚𝑝.lx_boy[τ] = (τ == 1 ? 1 : 𝑚𝑝.lx_boy_next)
      proj.cf[τ, :prem] += 𝑚𝑝.lx_boy[τ] * 𝑚𝑝.β[τ, :prem]
      proj.cf[τ, :λ_boy] -=
        𝑚𝑝.lx_boy[τ] * 𝑚𝑝.λ[τ, :boy] *
        𝑚𝑝.λ[τ, :cum_infl] / (1 + 𝑚𝑝.λ[τ, :infl])
    end
  end
end

"""
`projecteoy!(τ, cap_mkt::CapMkt, invs::InvPort, liabs::LiabIns,
  dyn::Dynamic, proj::Projection)`

Project one year, update values at the end of the year `τ`

**Changed:** `proj::Projection`, `liabs::LiabIns  (liabs.mp)`
"""
function projecteoy!(τ,
                     cap_mkt::CapMkt,
                     invs::InvPort,
                     liabs::LiabIns,
                     dyn::Dynamic,
                     proj::Projection)
  tpg_price_positive = 0.0
  for 𝑚𝑝 ∈ liabs.mps
    if τ <= 𝑚𝑝.dur
      tpg_price_positive +=
        𝑚𝑝.lx_boy[τ] *
        max(0, (τ == 1 ? 𝑚𝑝.tpg_price_0 : 𝑚𝑝.tpg_price[τ-1]))
    end
  end
  for 𝑚𝑝 ∈ liabs.mps
    if τ <= 𝑚𝑝.dur
      prob = deepcopy(𝑚𝑝.prob)
      prob[:,:sx] *=
        δsx(τ, cap_mkt, invs, 𝑚𝑝, dyn)
      prob[:,:px] = 1 .- prob[:,:qx] - prob[:,:sx]
      𝑚𝑝.lx_boy_next = 𝑚𝑝.lx_boy[τ] * prob[τ, :px]
      for 𝑤𝑥 ∈ [:qx, :sx, :px]
        proj.cf[τ, 𝑤𝑥] -=
          𝑚𝑝.lx_boy[τ] * prob[τ, 𝑤𝑥] * 𝑚𝑝.β[τ, 𝑤𝑥]
      end
      proj.cf[τ, :λ_eoy] -=
        𝑚𝑝.lx_boy[τ] * 𝑚𝑝.λ[τ, :eoy] * 𝑚𝑝.λ[τ, :cum_infl]
      proj.val[τ, :tpg] +=
        𝑚𝑝.lx_boy[τ] * prob[τ, :px] *
        tpg(τ, cap_mkt.rfr.x, prob, 𝑚𝑝.β, 𝑚𝑝.λ)
    end
  end
  proj.cf[τ, :Δtpg] =
    -(proj.val[τ, :tpg] -
        (τ == 1 ? proj.val_0[1, :tpg] : proj.val[τ - 1, :tpg]))

end

"""
`project!(τ, cap_mkt::CapMkt, invs::InvPort, dyn::Dynamic,
  proj::Projection)`

Project one year, investment results from the year `τ`

**Changed:** `proj::Projection`, `invs::InvPort`
"""
function project!(τ,
                  cap_mkt::CapMkt,
                  invs::InvPort,
                  dyn::Dynamic,
                  proj::Projection)
  mv_boy =
    (τ == 1 ? proj.val_0[1, :invest] : proj.val[τ - 1, :invest])
  mv_boy += proj.cf[τ, :prem] + proj.cf[τ, :λ_boy]
  alloc!(τ, cap_mkt, invs)
  project!(τ, mv_boy, invs)
  proj.cf[τ, :invest] = invs.mv[τ] - mv_boy
end

"""
`bonus!(τ, invs::InvPort, liabs::LiabIns, dyn::Dynamic,
  proj, surp_pre_profit_tax_bonus)`

Bonus at the end of year `τ`

**Changed:** `proj::Projection`
"""
function bonus!(τ,
                invs::InvPort,
                liabs::LiabIns,
                dyn::Dynamic,
                proj,
                surp_pre_profit_tax_bonus)
  for 𝑚𝑝 ∈ liabs.mps
    if τ <= 𝑚𝑝.dur
      proj.cf[τ, :bonus] -=
        min(surp_pre_profit_tax_bonus,
            𝑚𝑝.lx_boy[τ] *
              bonusrate(τ, invs.yield[τ], 𝑚𝑝, dyn) *
              max(0, (τ == 1 ?
                        𝑚𝑝.tpg_price_0 :
                        𝑚𝑝.tpg_price[τ-1])))
    end
  end
end

"""
`investpredivid(τ, invs::InvPort, proj::Projection)`

Market value of assets before payment of dividends
"""
function investpredivid(τ, invs::InvPort, proj::Projection)
  invs.mv_boy[τ] +
    sum(convert(Array,
                proj.cf[τ,
                        [:invest, :qx, :sx, :px, :λ_eoy, :bonus,
                         :l_other, :tax, :gc]]))
end

"""
`project!(τ, cap_mkt::CapMkt, invs::InvPort, liabs::LiabIns,
  liab_other::LiabOther, dyn::Dynamic, proj::Projection)`

Project one year

**Changed:** `proj::Projection`, `invs::InvPort`, `dyn::Dynamic`
"""
function project!(τ,
                  cap_mkt::CapMkt,
                  invs::InvPort,
                  liabs::LiabIns,
                  liab_other::LiabOther,
                  dyn::Dynamic,
                  proj::Projection)
  projectboy!(τ, proj, liabs)
  proj.cf[τ, :new_debt] = -getloan(τ, liab_other)
  project!(τ, cap_mkt, invs, dyn, proj)
  update!(τ, proj, dyn)
  proj.cf[τ, :λ_eoy] = -invs.cost[τ]
  proj.cf[τ, :l_other] = -paycoupon(τ, liab_other)
  proj.cf[τ, :l_other] -= payprincipal(τ, liab_other)
  proj.val[τ, :l_other] = pv(τ, cap_mkt, liab_other)
  projecteoy!(τ, cap_mkt, invs, liabs, dyn, proj)

  proj.cf[τ,:tax] = 0.0
  proj.cf[τ,:bonus] = 0.0
  surp_pre_profit_tax_bonus =
    max(0,
        investpredivid(τ, invs, proj) -
          proj.val[τ, :tpg] -
          proj.val[τ, :l_other]  )
  bonus!(τ, invs, liabs, dyn, proj, surp_pre_profit_tax_bonus)
  proj.cf[τ, :profit] =
    sum(convert(Array, proj.cf[τ, [:prem, :invest,
                                   :qx, :sx, :px, :λ_boy, :λ_eoy,
                                   :Δtpg, :bonus, :l_other]]))
  tax = proj.tax_rate * proj.cf[τ, :profit] ## could be negative
  tax_credit_pre =
    (τ == 1 ? proj.tax_credit_0 : proj.tax_credit[τ - 1])
  proj.cf[τ, :tax] = -max(0, tax - tax_credit_pre)
  proj.tax_credit[τ] =  ## no new tax credit generated
  tax_credit_pre - tax - proj.cf[τ, :tax]

  proj.val[τ, :invest] = investpredivid(τ, invs, proj)
  proj.cf[τ, :divid] =
    -freesurp(dyn,
              proj.val[τ, :invest],
              proj.val[τ, :tpg] + proj.val[τ, :l_other])
  proj.val[τ, :invest] +=  proj.cf[τ, :divid]

  proj.val[τ, :surplus] =
    proj.val[τ, :invest] -
    proj.val[τ, :tpg] -
    proj.val[τ, :l_other]
end

"""
`pvprev(rfr, cf, pv) = (cf + pv) /  (1 + rfr)`

Recursive step in generic present value calculation
"""
pvprev(rfr, cf, pv) = (cf + pv) /  (1 + rfr)

"""
`pvvec(rfr::Vector{Float64}, cf)``

Generic present value calculation, where `rfr` denotes the risk
free rate for each year and `cf`. The length of the vector `rfr`
may not be smaller than the length of the cashflow.
Each payment occurs at the end of the corresponding year.

Output: A vector of the same length as `cf` contaning the present
value at the end of each year. The last component is zero.
"""
function pvvec(rfr::Vector{Float64}, cf)
  T = length(cf)
  val = zeros(Float64, T)
  for 𝑡 ∈ reverse(collect(1:(T-1))) # [T-1:-1:1] #r
    val[𝑡] = pvprev(rfr[𝑡 + 1], val[𝑡 + 1], cf[𝑡 + 1])
  end
  return val
end

"""
`valbonus!(rfr::Vector{Float64}, proj::Projection) `

Provisions for future bonus payments (at each time `τ`).
Needs to be called after the projection is completed

**Changed:** proj::Projection
"""
function valbonus!(rfr::Vector{Float64},
                   proj::Projection)
  proj.val[:bonus] = pvvec(rfr,  -proj.cf[:bonus])
  proj.val_0[:bonus] =
    pvprev(rfr[1], -proj.cf[1, :bonus], proj.val[1, :bonus])
end

"""
`valcostprov!(rfr::Vector{Float64}, invs::InvPort,
  proj::Projection)`

Provisions for future costs (at each time `τ`). Absolute costs
(including absolute investment costs from *all* investments) and
relative investment costs for provisions are considered.
It is assumed that provisions are backed by cash investments
Needs to be called after the projection is completed

**Changed:** proj::Projection
"""
function valcostprov!(rfr::Vector{Float64},
                      invs::InvPort,
                      proj::Projection)
  cash_cost = deepcopy(invs.igs[:IGCash].cost)
  proj.cf[1, :cost_prov] =
    proj.fixed_cost_gc[1] +
    sum(convert(Array,
                proj.val_0[1, [:tpg, :bonus, :l_other]])) *
    cash_cost.cum_infl_rel[1] * cash_cost.rel[1]
  for 𝑡 ∈ 2:proj.dur
    proj.cf[𝑡, :cost_prov] =
      proj.fixed_cost_gc[𝑡] +
      sum(convert(Array,
                  proj.val[𝑡 - 1, [:tpg, :bonus, :l_other]])) *
      cash_cost.cum_infl_rel[𝑡] * cash_cost.rel[𝑡]
  end
  proj.val[:cost_prov] =
    pvvec(rfr - cash_cost.rel,  proj.cf[:cost_prov])
  proj.val_0[:cost_prov] = pvprev(rfr[1] - cash_cost.rel[1],
                                  proj.cf[1, :cost_prov],
                                  proj.val[1, :cost_prov])
end


"""
Calculates the SVD deomposition of a matrix by employing
a divide and conquer approach.
Returns a `SVD` factorization object.
"""
gesdd(x) = gesdd!(copy(x))
"""
Same as `gesdd` but saves space by overwriting the input matrix.
"""
function gesdd!(x)
  u,s,vt = LinearAlgebra.LAPACK.gesdd!('A', x)
  SVD(u,s,vt)
end

"""
Calculates the SVD deomposition of a matrix.
Returns a `SVD` factorization object.
"""
gesvd(x) = gesvd!(copy(x))
"""
Same as `gesvd` but saves space by overwriting the input matrix.
"""
function gesvd!(x)
  u,s,vt = LinearAlgebra.LAPACK.gesvd!('A', 'A', x)
  SVD(u,s,vt)
end

# extra operations for SVD factorizations
"""
    fact_mult(A::SVD, B::SVD) -> SVD

Stabilized multiplication of two SVD decompositions.
Returns a `SVD` factorization object.
"""
function fact_mult(A::SVD, B::SVD)
    mat = A.Vt * B.U
    lmul!(Diagonal(A.S), mat)
    rmul!(mat, Diagonal(B.S))
    F = svd!(mat)
    SVD(A.U * F.U, F.S, F.Vt * B.Vt)
end

##############################################################
#
#                   SVD / USVt
#
##############################################################

"""
    svd_inv_one_plus(F::SVD) -> SVD

Stabilized calculation of [1 + USVt]^(-1):

  * Use one intermediate SVD decomposition.

Options for preallocation via keyword arguments:

  * `t = similar(F.U)`
  * `a = similar(F.U)`
  * `c = similar(F.U)`
"""
function svd_inv_one_plus(F::SVD;
                          svdalg = svd!,
                          t = similar(F.U),
                          a = similar(F.U),
                          c = similar(F.U),
                          internaluse=false)
  U, S, V = F
  mul!(t, U', V)
  t[diagind(t)] .+= S
  u, s, v = svdalg(t)
  mul!(a, V, v)
  s .= 1 ./ s
  mul!(c, u', U')
  if internaluse
    SVD(a, s, c)
  else
    SVD(copy(a), s, copy(c))
  end
end

"""
    inv_one_plus!(res, F::SVD) -> res

Stabilized calculation of [1 + USVt]^(-1):

  * Use one intermediate SVD decomposition.

Writes the result into `res`.

See `svd_inv_one_plus` for preallocation options.
"""
function inv_one_plus!(res, F::SVD; kwargs...)
  M = svd_inv_one_plus(F; internaluse=true, kwargs...)
  rmul!(M.U, Diagonal(M.S))
  mul!(res, M.U, M.Vt)
  res
end

"""
    inv_one_plus(F::SVD) -> AbstractMatrix

Stabilized calculation of [1 + USVt]^(-1):

  * Use one intermediate SVD decomposition.

See `svd_inv_one_plus` for preallocation options.
"""
function inv_one_plus(F::SVD; kwargs...)
  res = similar(F.U)
  inv_one_plus!(res, F; kwargs...)
  res
end

"""
    svd_inv_one_plus_loh(F::SVD) -> SVD

Stabilized calculation of [1 + USVt]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate SVD decompositions.

Options for preallocation via keyword arguments:

  * `l = similar(F.U)`
  * `r = similar(F.U)`
  * `Dp = similar(F.D)`
  * `Dm = similar(F.D)`
"""
function svd_inv_one_plus_loh(F::SVD;
                              svdalg = svd!,
                              Sp = similar(F.S),
                              Sm = similar(F.S),
                              l = similar(F.V),
                              r = similar(F.U))
  U, S, V = F

  copyto!(l, V)
  copyto!(r, U)

  Sp .= max.(S, 1)
  Sm .= min.(S, 1)

  Sp .\= 1 # Sp now Spinv

  rmul!(l, Diagonal(Sp))
  rmul!(r, Diagonal(Sm))

  M = svdalg(l + r)
  m = inv(M)
  lmul!(Diagonal(Sp), m)
  u, s, v = svdalg(m)
  mul!(m, V, u)

  SVD(m, s, v')
end

"""
  inv_one_plus_loh!(res, F::SVD) -> res

Stabilized calculation of [1 + USVt]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate SVD decompositions.

Writes the result into `res`.

See `svd_inv_one_plus_loh` for preallocation options.
"""
function inv_one_plus_loh!(res, F::SVD; kwargs...)
  M = svd_inv_one_plus_loh(F; kwargs...)
  rmul!(M.U, Diagonal(M.S))
  mul!(res, M.U, M.Vt)
  res
end

"""
  inv_one_plus_loh(F::SVD) -> AbstractMatrix

Stabilized calculation of [1 + USVt]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate SVD decompositions.

See `svd_inv_one_plus_loh` for preallocation options.
"""
function inv_one_plus_loh(F::SVD; kwargs...)
  res = similar(F.U)
  inv_one_plus_loh!(res, F; kwargs...)
end

"""
    svd_inv_sum(A::SVD, B::SVD) -> SVD

Stabilized calculation of [UaSaVta + UbSbVtb]^(-1):

  * Use one intermediate SVD decompositions.

Options for preallocations via keyword arguments:

  * `m1 = similar(A.U)`
  * `m2 = similar(A.U)`

"""
function svd_inv_sum(A::SVD, B::SVD; m1 = similar(A.U), m2 = similar(A.U), svdalg = svd!, internaluse = false)
  Ua, Sa, Va = A
  Ub, Sb, Vb = B

  mul!(m1, Va', Vb)
  lmul!(Diagonal(Sa), m1)

  mul!(m2, Ua', Ub)
  rmul!(m2, Diagonal(Sb))

  u,s,v = svdalg(m1 + m2)

  mul!(m1, Ua, u)
  mul!(m2, Vb, v)

  if internaluse
    SVD(m2, 1 ./ s, m1')
  else
    SVD(copy(m2), 1 ./ s, copy(m1'))
  end
end

"""
    inv_sum!(res, A::SVD, B::SVD) -> res

Stabilized calculation of [UaSaVta + UbSbVtb]^(-1):

  * Use one intermediate SVD decompositions.

Writes the result into `res`.

See `svd_inv_sum` for preallocation options.
"""
function inv_sum!(res, A::SVD, B::SVD; kwargs...)
  M = svd_inv_sum(A, B; internaluse = true, kwargs...)
  rmul!(M.U, Diagonal(M.S))
  mul!(res, M.U, M.Vt)
  res
end

"""
    inv_sum(A::SVD, B::SVD) -> res

Stabilized calculation of [UaSaVta + UbSbVtb]^(-1):

  * Use one intermediate SVD decompositions.

See `svd_inv_sum` for preallocation options.
"""
function inv_sum(A::SVD, B::SVD; kwargs...)
  res = similar(A.U)
  inv_sum!(res, A, B; kwargs...)
  res
end

"""
    svd_inv_sum_loh(A::SVD, B::SVD) -> SVD

Stabilized calculation of [UaSaVta + UbSbVtb]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate SVD decompositions.

Options for preallocations via keyword arguments:

  * None (yet)

"""
function svd_inv_sum_loh(A::SVD, B::SVD; svdalg = svd!)
    Ua, Sa, Va = A
    Ub, Sb, Vb = B

    d = length(Sa)

    Sap = max.(Sa,1.)
    Sam = min.(Sa,1.)
    Sbp = max.(Sb,1.)
    Sbm = min.(Sb,1.)

    mat1 = Va' * Vb
    for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] * Sam[j]/Sbp[k]
    end

    mat2 = adjoint(Ua) * Ub
    for j in 1:d, k in 1:d
        mat2[j,k]=mat2[j,k] * Sbm[k]/Sap[j]
    end

    mat1 = mat1 + mat2

    M = svdalg(mat1)
    mat1 = inv(M)

    for j in 1:d, k in 1:d
        mat1[j,k]=mat1[j,k] / Sbp[j] / Sap[k]
    end

    u, s, v = svdalg(mat1)
    mul!(mat1, Vb, u)
    mul!(mat2, v', adjoint(Ua))

    SVD(mat1, s, mat2)
end

"""
    inv_sum_loh!(res, A::SVD, B::SVD) -> res

Stabilized calculation of [UaSaVta + UbSbVtb]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate SVD decompositions.

Writes the result into `res`.

See `svd_inv_sum_loh` for preallocation options.
"""
function inv_sum_loh!(res, A::SVD, B::SVD; kwargs...)
  M = svd_inv_sum_loh(A, B; kwargs...)
  rmul!(M.U, Diagonal(M.S))
  mul!(res, M.U, M.Vt)
  res

end

"""
    inv_sum_loh(A::SVD, B::SVD) -> res

Stabilized calculation of [UaSaVta + UbSbVtb]^(-1):

  * Separate scales larger and smaller than unity
  * Use two intermediate SVD decompositions.

See `svd_inv_sum_loh` for preallocation options.
"""
function inv_sum_loh(A::SVD, B::SVD; kwargs...)
  res = similar(A.U)
  inv_sum_loh!(res, A, B; kwargs...)
  res
end


# Current sheet model.

"""
    getB_CS_harris(B₀, L)

Return the magnetic field at location `r` near a current sheet with magnetic strength `B₀`
and sheet length `L`.
"""
function getB_CS_harris(r, B₀, L)
   x = r[1]
   [0.0, 0.0, B₀ * tanh(x/L)]
end


using LinearAlgebra, Krylov, LinearOperators, SparseArrays, Stopping

#Krylov @kdot
macro kdot(n, x, y)
  return esc(:(Krylov.krylov_dot($n, $x, 1, $y, 1)))
end

using SolverTools, Logging

"""
Randomized coordinate descent

Sect. 3.7 in Gower, R. M., & Richtárik, P. (2015).
Randomized iterative methods for linear systems.
SIAM Journal on Matrix Analysis and Applications, 36(4), 1660-1690.

Using Stopping and a fake line search
"""
function StopRandomizedCD_LS(A             :: AbstractMatrix,
                             b             :: AbstractVector{T};
                             is_zero_start :: Bool = true,
                             x0            :: AbstractVector{T} = zeros(T,size(A,2)),
                             atol          :: AbstractFloat = 1e-7,
                             rtol          :: AbstractFloat = 1e-15,
                             max_iter      :: Int = size(A,2)^2,
                             verbose       :: Int = 100,
                             kwargs...) where T <: AbstractFloat

 stp = LAStopping(LinearSystem(A,b),
                  GenericState(x0, similar(b)),
                  max_cntrs = Stopping._init_max_counters_linear_operators(),
                  atol = atol, rtol = rtol, max_iter = max_iter,
                  tol_check = (atol, rtol, opt0)->(atol + rtol * opt0),
                  retol = false,
                  optimality_check = (pb :: LinearSystem, state :: GenericState{Vector{T},Vector{T}}) -> state.res,
                  kwargs...)

 return StopRandomizedCD_LS(stp, is_zero_start = is_zero_start, verbose = verbose, kwargs...)
end

function StopRandomizedCD_LS(stp           :: AbstractStopping;
                             is_zero_start :: Bool = true,
                             verbose       :: Int = 100,
                             kwargs...)
    state = stp.current_state

    A,b = stp.pb.A, stp.pb.b
    m,n = size(A)
    x  = state.x
    T  = eltype(x)

    state.res = is_zero_start ? b : b - A*x

    OK = start!(stp)

    #We create a fake line search
    substate = GenericState(x)
    #stop_remote = cheap_stop_remote_control()
    #meta = StoppingMeta() #the expansive part
    #list = VoidListStates()
    #stopping_user_struct = nothing
    #substp = GenericStopping(stp.pb.b, cheap_stop_remote_control(), substate, main_stp = stp, max_iter = 0)
    substp = GenericStopping(stp.pb.b, cheap_stop_remote_control(), substate, main_stp = stp)
    #substp = GenericStopping(stp.pb.b, substate)
    #substp = GenericStopping(stp.pb.b, meta, stop_remote, substate)

    #@info log_header([:iter, :nrm, :time], [Int, T, T],
    #                 hdr_override=Dict(:nrm=>"||Ax-b||"))
    #@info log_row(Any[0, state.current_score[1], state.current_time])

    while !OK

        #rand a number between 1 and n
        #224.662 ns (4 allocations: 79 bytes) - independent of the n
        i  = mod(stp.meta.nb_of_stop,n)+1#Int(floor(rand() * n) + 1)
        Ai = A[:,i]

        #ei = zeros(n); ei[i] = 1.0 #unit vector in R^n
        #xk  = Ai == 0 ? x0 : x0 - dot(Ai,res)/norm(Ai,2)^2 * ei
        Aires = @kdot(m, Ai, state.res)
        nAi   = @kdot(m, Ai, Ai)
        state.x[i] -= Aires/nAi

        state.res += Ai * Aires/nAi

        #reinitialize and "solve" the fake subproblem
        #reinit!(substp)
        #substp.current_state.x[i] = state.x[i] #reinit!(substp.current_state, state.x) #would reassign the whole vector x
        #substp.current_state.x = state.x
        #solve_fake_subproblem!(substp)

        OK = stop!(stp)

        if mod(stp.meta.nb_of_stop, verbose) == 0 #print every 20 iterations
    #     @info log_row(Any[stp.meta.nb_of_stop, state.current_score[1], state.current_time])
        end

    end

 return stp
end

function solve_fake_subproblem!(stp :: GenericStopping)
    start!(stp)
    return stop!(stp)
end


using Test
using LinearAlgebra
using QBase

@testset "./src/math/matrices.jl" begin

@testset "partial_trace()" begin
    @testset "invalid inputs" begin
        @test_throws DomainError partial_trace([1 0;0 1], [2,2], 1)
        @test_throws DomainError partial_trace([1 0 0 0 0;0 0 0 0 0;0 0 0 0 0;0 0 0 0 0], [2,2], 1)
    end

    @testset "diagonal density operator" begin
        ρ = (1/6)*Matrix(1I,6,6)

        @test partial_trace(ρ, [2,3],2) == [0.5 0; 0 0.5]
        @test partial_trace(ρ, [2,3],1) == (1/3)*I
        @test partial_trace(ρ, [3,2],2) == (1/3)*I
    end

    @testset "maximally entangled bell state" begin
        ρ = 0.5*[1 0 0 1; 0 0 0 0; 0 0 0 0; 1 0 0 1]

        @test partial_trace(ρ,[2,2],2) == [0.5 0; 0 0.5]
        @test partial_trace(ρ,[2,2],1) == [0.5 0; 0 0.5]
    end
end

@testset "computational_basis_vectors()" begin
    @test computational_basis_vectors(2) == [[1,0],[0,1]]
    @test computational_basis_vectors(5) == [
        [1,0,0,0,0],[0,1,0,0,0],[0,0,1,0,0],[0,0,0,1,0],[0,0,0,0,1]
    ]
end

@testset "commutes()" begin
    @test_throws DomainError commutes([1 0;0 1;1 1],[0 1;1 0])

    @test commutes([1 0;0 0],[0 0;0 1])
    @test !commutes([0 1;1 0],[1 0;0 -1])

    @test commutes.([[0 1;1 0],[1 0;0 -1]],[[0 im;-im 0],[0 im;-im 0]]) == [false,false]
end

@testset "is_hermitian()" begin
    @testset "invalid cases" begin
        @test !is_hermitian([1 0;1 1])
        @test !is_hermitian([1 1;-1 1])
    end
    @testset "valid cases" begin
        @test is_hermitian([1 0;0 1])
        @test is_hermitian([1 0;0 -1])
        @test is_hermitian([1 im; -im 1])
    end
    @testset "using atol" begin
        @test !is_hermitian([1 1;1+1e-6 1])
        @test is_hermitian([1 1;1+1e-6 1], atol=1e-5)

        @test !is_hermitian([1 1e-6 + 1im;-1im 1])
        @test is_hermitian([1 1e-6 + 1im;-1im 1], atol=1e-5)
    end
end

@testset "is_positive_semidefinite()" begin
    @testset "invalid cases" begin
        @test !is_positive_semidefinite([1 0 ; 0 -1])
        @test !is_positive_semidefinite([0 1;1 0])
        @test !is_positive_semidefinite([1 1;-1 1])
        @test !is_positive_semidefinite([1 0 0;0 1 0])
    end

    @testset "valid cases" begin
        @test is_positive_semidefinite([1 0; 0 1])
    end

    @testset "using atol" begin
        @test !is_positive_semidefinite([1 0;0 -1e-6])
        @test is_positive_semidefinite([1 0;0 -1e-6], atol=1e-6)
    end
end

@testset "is_orthonormal_basis()" begin
    @testset "valid cases" begin
        @test is_orthonormal_basis([[1,0],[0,1]])
        @test is_orthonormal_basis(Ket.([[1,0],[0,1]]))
        @test is_orthonormal_basis(computational_basis_vectors(5))
        @test is_orthonormal_basis([[1,1im]/sqrt(2),[1,-1im]/sqrt(2)])
    end

    @testset "invalid cases" begin
        @test !is_orthonormal_basis([[1,1]/2,[1,-1]/2])
        @test !is_orthonormal_basis([[0.5,0.5],[0.5,0.5]])
    end

    @testset "atol" begin
        @test !is_orthonormal_basis([[1,0.1],[0,1]])
        @test is_orthonormal_basis([[1,0.1],[0,1]],atol=0.2)
    end
end

@testset "n_product_id()" begin
    m0 = [1 0;0 0]
    m1 = [0 0;1 0]
    m2 = [0 0 0;0 0 0;0 1 0;0 0 0]
    m3 = [0 0 0 0;0 0 0 0;0 0 0 1]

    col_id = n_product_id([1,1],[2,2])
    row_id = n_product_id([1,2],[2,2])
    @test col_id == 1
    @test row_id == 2
    @test kron(m0,m1)[row_id,col_id] == 1

    col_id = n_product_id([1,1,2,4],[2,2,3,4])
    row_id = n_product_id([1,2,3,3],[2,2,4,3])
    @test col_id == 8
    @test row_id == 21
    @test kron(m0,m1,m2,m3)[row_id,col_id] == 1

    @testset "errors" begin
        @test_throws DomainError n_product_id([1,1,1],[2,2])
        @test_throws DomainError n_product_id([1,3],[2,2])
    end
end

@testset "is_complete()" begin
    @testset "valid cases" begin
        @test is_complete([[1 0;0  0],[0 0;0 1]])
        @test is_complete([[0.25 -0.25im;0.25im 0.25],[0.25 0.25im;-0.25im 0.25],[0.5 0;0 0.5]])
        @test is_complete(computational_basis_states(3))
        @test is_complete([[1 0;0.5 0],[0 0;0 1]], atol=0.51)
    end

    @testset "invalid cases" begin
        @test !is_complete([[1 0;0.5 0],[0 0;0 1]])
    end
end

end


"""
    eigsolve(A::AbstractMatrix, [howmany = 1, which = :LM, T = eltype(A)]; kwargs...)
    eigsolve(f, n::Int, [howmany = 1, which = :LM, T = Float64]; kwargs...)
    eigsolve(f, x₀, [howmany = 1, which = :LM]; kwargs...)
    eigsolve(f, x₀, howmany, which, algorithm)

Compute `howmany` eigenvalues from the linear map encoded in the matrix `A` or by the
function `f`. Return eigenvalues, eigenvectors and a `ConvergenceInfo` structure.

### Arguments:
The linear map can be an `AbstractMatrix` (dense or sparse) or a general function or callable
object. If an `AbstractMatrix` is used, a starting vector `x₀` does not need to be provided,
it is then chosen as `rand(T, size(A,1))`. If the linear map is encoded more generally as a
a callable function or method, the best approach is to provide an explicit starting guess `x₀`.
Note that `x₀` does not need to be of type `AbstractVector`, any type that behaves as a vector
and supports the required methods (see KrylovKit docs) is accepted. If instead of `x₀` an integer
`n` is specified, it is assumed that `x₀` is a regular vector and it is initialized to `rand(T,n)`,
where the default value of `T` is `Float64`, unless specified differently.

The next arguments are optional, but should typically be specified. `howmany` specifies how many
eigenvalues should be computed; `which` specifies which eigenvalues should be targetted. Valid
specifications of `which` are given by
*   `LM`: eigenvalues of largest magnitude
*   `LR`: eigenvalues with largest (most positive) real part
*   `SR`: eigenvalues with smallest (most negative) real part
*   `LI`: eigenvalues with largest (most positive) imaginary part, only if `T <: Complex`
*   `SI`: eigenvalues with smallest (most negative) imaginary part, only if `T <: Complex`
*   [`ClosestTo(λ)`](@ref): eigenvalues closest to some number `λ`
!!! note "Note about selecting `which` eigenvalues"
    Krylov methods work well for extremal eigenvalues, i.e. eigenvalues on the periphery of
    the spectrum of the linear map. Even with `ClosestTo`, no shift and invert is performed.
    This is useful if, e.g., you know the spectrum to be within the unit circle in the complex
    plane, and want to target the eigenvalues closest to the value `λ = 1`.

The argument `T` acts as a hint in which `Number` type the computation should be performed, but
is not restrictive. If the linear map automatically produces complex values, complex arithmetic
will be used even though `T<:Real` was specified.

### Return values:
The return value is always of the form `vals, vecs, info = eigsolve(...)` with
*   `vals`: a `Vector` containing the eigenvalues, of length at least `howmany`, but could be
    longer if more eigenvalues were converged at the same cost. Eigenvalues will be real if
    [`Lanczos`](@ref) was used and complex if [`Arnoldi`](@ref) was used (see below).
*   `vecs`: a `Vector` of corresponding eigenvectors, of the same length as `vals`. Note that
    eigenvectors are not returned as a matrix, as the linear map could act on any custom Julia
    type with vector like behavior, i.e. the elements of the list `vecs` are objects that are
    typically similar to the starting guess `x₀`, up to a possibly different `eltype`. In particular,
    for a general matrix (i.e. with `Arnoldi`) the eigenvectors are generally complex and are
    therefore always returned in a complex number format.
    When the linear map is a simple `AbstractMatrix`, `vecs` will be `Vector{Vector{<:Number}}`.
*   `info`: an object of type [`ConvergenceInfo`], which has the following fields
    -   `info.converged::Int`: indicates how many eigenvalues and eigenvectors were actually
        converged to the specified tolerance `tol` (see below under keyword arguments)
    -   `info.residual::Vector`: a list of the same length as `vals` containing the residuals
        `info.residual[i] = f(vecs[i]) - vals[i] * vecs[i]`
    -   `info.normres::Vector{<:Real}`: list of the same length as `vals` containing the norm
        of the residual `info.normres[i] = norm(info.residual[i])`
    -   `info.numops::Int`: number of times the linear map was applied, i.e. number of times
        `f` was called, or a vector was multiplied with `A`
    -   `info.numiter::Int`: number of times the Krylov subspace was restarted (see below)
!!! warning "Check for convergence"
    No warning is printed if not all requested eigenvalues were converged, so always check
    if `info.converged >= howmany`.

### Keyword arguments:
Keyword arguments and their default values are given by:
*   `tol::Real`: the requested accuracy (corresponding to the 2-norm of the residual for
    Schur vectors, not the eigenvectors). If you work in e.g. single precision (`Float32`),
    you should definitely change the default value.
*   `krylovdim::Integer`: the maximum dimension of the Krylov subspace that will be constructed.
    Note that the dimension of the vector space is not known or checked, e.g. `x₀` should not
    necessarily support the `Base.length` function. If you know the actual problem dimension
    is smaller than the default value, it is useful to reduce the value of `krylovdim`, though
    in principle this should be detected.
*   `maxiter::Integer`: the number of times the Krylov subspace can be rebuilt; see below for
    further details on the algorithms.
*   `orth::Orthogonalizer`: the orthogonalization method to be used, see [`Orthogonalizer`](@ref)
*   `issymmetric::Bool`: if the linear map is symmetric, only meaningful if `T<:Real`
*   `ishermitian::Bool`: if the linear map is hermitian
The default values are given by `tol = KrylovDefaults.tol`, `krylovdim = KrylovDefaults.krylovdim`,
`maxiter = KrylovDefaults.maxiter`, `orth = KrylovDefaults.orth`; see [`KrylovDefaults`](@ref) for details.

The default value for the last two parameters depends on the method. If an `AbstractMatrix`
is used, `issymmetric` and `ishermitian` are checked for that matrix, ortherwise the default
values are `issymmetric = false` and `ishermitian = T <: Real && issymmetric`.

### Algorithm
The last method, without default values and keyword arguments, is the one that is finally called,
and can also be used directly. Here, one specifies the algorithm explicitly as either [`Lanczos`](@ref),
for real symmetric or complex hermitian problems, or [`Arnoldi`](@ref), for general problems.
Note that these names refer to the process for building the Krylov subspace, but the actual
algorithm is an implementation of the Krylov-Schur algorithm, which can dynamically shrink and
grow the Krylov subspace, i.e. the restarts are so-called thick restarts where a part of the
current Krylov subspace is kept.

!!! note "Note about convergence"
    In case of a general problem, where the `Arnoldi` method is used, convergence of an eigenvalue
    is not based on the norm of the residual `norm(f(vecs[i]) - vals[i]*vecs[i])` for the eigenvectors
    but rather on the norm of the residual for the corresponding Schur vectors.

    See also [`schursolve`](@ref) if you want to use the partial Schur decomposition directly,
    or if you are not interested in computing the eigenvectors, and want to work in real arithmetic
    all the way true (if the linear map and starting guess are real).
"""
function eigsolve end

"""
    ClosestTo(λ)

A simple `struct` to be used in combination with [`eigsolve`](@ref) or [`schursolve`](@ref) to
indicate which eigenvalues need to be targetted, namely those closest to `λ`.
"""
struct ClosestTo{T}
    λ::T
end

const Selector = Union{ClosestTo, Symbol}

function eigsolve(A::AbstractMatrix, howmany::Int = 1, which::Selector = :LM, T::Type = eltype(A); kwargs...)
    eigsolve(A, rand(T, size(A,1)), howmany, which; kwargs...)
end
function eigsolve(f, n::Int, howmany::Int = 1, which::Selector = :LM, T::Type = Float64; kwargs...)
    eigsolve(f, rand(T, n), howmany, which; kwargs...)
end
function eigsolve(f, x₀, howmany::Int = 1, which::Selector = :LM; kwargs...)
    alg = eigselector(f, eltype(x₀); kwargs...)
    checkwhich(which) || error("Unknown eigenvalue selector: which = $which")
    if alg isa Lanczos
        if which == :LI || which == :SI || which isa ClosestTo{<:Complex}
            error("Eigenvalue selector which = $which invalid: real eigenvalues expected with Lanczos algorithm")
        end
    elseif eltype(x₀) <: Real
        if which == :LI || which == :SI || which isa ClosestTo{<:Complex}
            error("Eigenvalue selector which = $which invalid: work in complex arithmetic by providing a complex starting vector `x₀`")
        end
    end
    eigsolve(f, x₀, howmany, which, alg)
end

function eigselector(f, T::Type;
    issymmetric::Bool = false, ishermitian::Bool = T<:Real && issymmetric,
    krylovdim::Int = KrylovDefaults.krylovdim, maxiter::Int = KrylovDefaults.maxiter,
    tol::Real = KrylovDefaults.tol, orth::Orthogonalizer = KrylovDefaults.orth)
    if (T<:Real && issymmetric) || ishermitian
        return Lanczos(krylovdim = krylovdim, maxiter = maxiter, tol=tol, orth = orth)
    else
        return Arnoldi(krylovdim = krylovdim, maxiter = maxiter, tol=tol, orth = orth)
    end
end
function eigselector(A::AbstractMatrix, T::Type;
    issymmetric::Bool = issymmetric(A), ishermitian::Bool = ishermitian(A),
    krylovdim::Int = KrylovDefaults.krylovdim, maxiter::Int = KrylovDefaults.maxiter,
    tol::Real = KrylovDefaults.tol, orth::Orthogonalizer = KrylovDefaults.orth)
    if (T<:Real && issymmetric) || ishermitian
        return Lanczos(krylovdim = krylovdim, maxiter = maxiter, tol=tol, orth = orth)
    else
        return Arnoldi(krylovdim = krylovdim, maxiter = maxiter, tol=tol, orth = orth)
    end
end

checkwhich(::ClosestTo) = true
checkwhich(s::Symbol) = s in (:LM, :LR, :SR, :LI, :SI)

function eigsort(which::ClosestTo)
    by = x->abs(x-which.λ)
    rev = false
    return by, rev
end
function eigsort(which::Symbol)
    if which == :LM
        by = abs
        rev = true
    elseif which == :LR
        by = real
        rev = true
    elseif which == :SR
        by = real
        rev = false
    elseif which == :LI
        by = imag
        rev = true
    elseif which == :SI
        by = imag
        rev = false
    else
        error("invalid specification of which eigenvalues to target: which = $which")
    end
    return by, rev
end


using DiffEqSensitivity, ForwardDiff, Distributions, OrdinaryDiffEq,
      LinearAlgebra, Test

function fiip(du,u,p,t)
    du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
    du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
end
function g(sol)
    J = extract_local_sensitivities(sol,true)[2]
    det(J'*J)
end

u0 = [1.0,1.0]
p = [1.5,1.0,3.0,1.0]
prob = ODEForwardSensitivityProblem(fiip,u0,(0.0,10.0),p,saveat=0:10)
sol = solve(prob, Tsit5())
u0_dist = [Uniform(0.9,1.1), 1.0]
p_dist = [1.5, truncated(Normal(1.5,.1),1.1, 1.9),3.0,1.0]
u0_dist_extended = vcat(u0_dist,zeros(length(p)*length(u0)))

function fiip_expe_SciML_forw_sen_SciML()
    prob = ODEForwardSensitivityProblem(fiip,u0,(0.0,10.0),p,saveat=0:10)
    prob_func = function (prob, i, repeat)
        _prob = remake(prob, u0=[isa(ui,Distribution) ? rand(ui) : ui for ui in u0_dist], p=[isa(pj,Distribution) ? rand(pj) : pj for pj in p_dist])
        _prob
    end
    output_func = function (sol, i)
        (g(sol), false)
    end
    monte_prob = EnsembleProblem(prob;output_func=output_func,prob_func=prob_func)
    sol = solve(monte_prob,Tsit5(),EnsembleSerial(),trajectories=100_000)
    mean(sol.u)
end

@test fiip_expe_SciML_forw_sen_SciML() ≈ 3.56e6 rtol=4e-2


## Instrumental variables problems

dt = readcsv("iv.csv");

y = dt[:,1];
x = dt[:,2:4];
z = dt[:,5:end];

f(theta) = z.*(y-x*theta);

gmm2s = GMMEstimator(f, [.0, .0, .0])
estimate!(gmm2s)

Dsn(θ)    = -z'x;
Dws(θ, p) = -z'*Diagonal(p)*x;
Dsl(θ, λ) = -x.*(z*λ);
Hwsl(θ, p, λ) = zeros(3,3);

mds = MDEstimator(f, [.0, .0, .0], grad = (Dsn, Dws, Dsl, Hwsl));

## GMM using linear algebra

## Stata Coef
## -.0050209   .0708816   .0593661

## Stata vcov
## symmetric e(V)[3,3]
##             v2          v3          v4
##    .00478911
##   -.00272551   .01026473
##    .00163441  -.00201291   .00983372

##   x1 |  -.0050209   .0692034    -0.07   0.942     -.140657    .1306153
##   x2 |   .0708816    .101315     0.70   0.484    -.1276922    .2694554
##   x3 |   .0593661   .0991651     0.60   0.549     -.134994    .2537261

## Hansen J statistic (overidentification test of all instruments):         1.191
##                                                     Chi-sq(4) P-val =    0.8795

## srand(1)
## using MomentBasedEstimators
## using ModelsGenerators
## y,x,z = ModelsGenerators.sim_iv_d01(n = 10000, CP = 2000)
## xb    = randn(10000,2)
## x     = [x xb]
## z     = [z xb]

## f(theta) = z.*(y-x*theta);

## gmm = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0], initialW = z'z);
## md = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0]);

## MomentBasedEstimators.initialize!(gmm);
## MomentBasedEstimators.initialize!(md);

## estimate!(gmm);
## estimate!(md);

## MomentBasedEstimators.setparLB!(md, coef(md))
## MomentBasedEstimators.setparUB!(md, coef(md))
## estimate!(md);

## MomentBasedEstimators.setparLB!(md, [0.,0,0])
## MomentBasedEstimators.setparUB!(md, [0.,0,0])

## ## Methods

## # 1. Coefficient
## coef(gmm)
## coef(md)

## ## 2. Variances
## vcov(gmm)
## vcov(md)

## ## 3. Objectiv value

## srand(1)
## y,x,z=randiv(100)
## xb = randn(100,2)
## x = [x xb]
## z = [z xb]

## z=sparse(z)

## f(theta) = broadcast(*, z, (y-x*theta))

## g = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0])
## m = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0])

## MomentBasedEstimators.initialize!(g);
## MomentBasedEstimators.initialize!(m);

## @time estimate!(g);
## @time estimate!(m);

## srand(1)
## y,x,z=randiv(100)
## xb = randn(100,2)
## x = [x xb]
## z = [z xb]

## f(theta) = z.*(y-x*theta);

## g = MomentBasedEstimators.GMMEstimator(f, [.1, 0, 0])
## m = MomentBasedEstimators.MDEstimator(f, [.1, 0, 0])

## MomentBasedEstimators.initialize!(g)
## MomentBasedEstimators.initialize!(m)

## @time estimate!(g);
## @time estimate!(m);

## Before
## elapsed time: 0.129503039 seconds (69486712 bytes allocated, 40.36% gc time)


################################################################################
#                               Linear                                         #
################################################################################

mutable struct Linear
    W::AbstractMatrix{T} where T
    b::AbstractVector{T} where T
end

function Linear(in_features::Int, out_features::Int)
    k_sqrt = sqrt(1 / in_features)
    d = Uniform(-k_sqrt, k_sqrt)
    return Linear(rand(d, out_features, in_features), rand(d, out_features))
end
Linear(in_out::Pair{Int, Int}) = Linear(in_out[1], in_out[2])

function Base.show(io::IO, l::Linear)
    o, i = size(l.W)
    print(io, "Linear($i=>$o)")
end

(l::Linear)(x::Union{AbstractVector{T}, AbstractMatrix{T}}) where T = l.W * x .+ l.b

################################################################################
#                               Sequential                                     #
################################################################################

mutable struct Sequential
    seq::Tuple
end

Sequential(args...) = Sequential(args)

function Base.show(io::IO, s::Sequential)
    println(io, "Sequential(")
    for m in s.seq
        println(io, "  $m,")
    end
    print(io, ")")
end

function (s::Sequential)(x)
    for m in s.seq
        x = m(x)
    end
    return x
end

################################################################################
#                           Convolutions                                       #
################################################################################

# Although internall 1D, 2D and 3D convolutions use the same convNd method,
# we keep different wrappers to make API easier to understand

################################# Conv1D #######################################

"""
1D convolutional layer.

Input and output dimensions: (sequence_length, channels, batch_size)
"""
mutable struct Conv1d
    W::AbstractArray{T,3} where T
    b::Union{AbstractVector{T} where T, Nothing}
    stride
    padding
    dilation
end

function Conv1d(in_channels::Int, out_channels::Int, kernel_size::Int;
                stride=1, padding=0, dilation=1, bias=true)
    kernel_tuple = (kernel_size,)
    # init weights same as in https://pytorch.org/docs/stable/nn.html#conv2d
    k_sqrt = sqrt(1 / (in_channels * prod(kernel_tuple)))
    d = Uniform(-k_sqrt, k_sqrt)
    W = rand(d, kernel_tuple..., in_channels, out_channels)
    b = bias ? rand(d, out_channels) : nothing
    return Conv1d(W, b, stride, padding, dilation)
end
Conv1d(in_out::Pair{Int, Int}, k) = Conv1d(in_out[1], in_out[2], k)

function Base.show(io::IO, c::Conv1d)
    k1, i, o = size(c.W)
    print(io, "Conv1d($i=>$o, $(k1))")
end

function (c::Conv1d)(x::AbstractArray{T,3}) where T
    y = convNd(x, c.W; stride=c.stride, padding=c.padding, dilation=c.dilation)
    if c.b != nothing
        y = y .+ reshape(c.b, (1, length(c.b), 1))  # TODO: check for 1D and 3D!
    end
    return y
end

################################# Conv2D #######################################

mutable struct Conv2d
    W::AbstractArray{T,4} where T
    b::Union{AbstractVector{T} where T, Nothing}
    stride
    padding
    dilation
end

function Conv2d(in_channels::Int, out_channels::Int, kernel_size::Union{Int, NTuple{2, Int}};
                stride=1, padding=0, dilation=1, bias=true)
    kernel_tuple = kernel_size isa Tuple ? kernel_size : (kernel_size, kernel_size)
    # init weights same as in https://pytorch.org/docs/stable/nn.html#conv2d
    k_sqrt = sqrt(1 / (in_channels * prod(kernel_tuple)))
    d = Uniform(-k_sqrt, k_sqrt)
    W = rand(d, kernel_tuple..., in_channels, out_channels)
    b = bias ? rand(d, out_channels) : nothing
    return Conv2d(W, b, stride, padding, dilation)
end
Conv2d(in_out::Pair{Int, Int}, k) = Conv2d(in_out[1], in_out[2], k)

function Base.show(io::IO, c::Conv2d)
    k1, k2, i, o = size(c.W)
    print(io, "Conv2d($i=>$o, $(k1)x$(k2))")
end

function (c::Conv2d)(x::AbstractArray{T,4}) where T
    y = convNd(x, c.W; stride=c.stride, padding=c.padding, dilation=c.dilation)
    if c.b != nothing
        y = y .+ reshape(c.b, (1, 1, length(c.b), 1))
    end
    return y
end

################################# Conv3D #######################################

mutable struct Conv3d
    W::AbstractArray{T,5} where T
    b::Union{AbstractVector{T} where T, Nothing}
    stride
    padding
    dilation
end

function Conv3d(in_channels::Int, out_channels::Int, kernel_size::Union{Int, NTuple{2, Int}};
                stride=1, padding=0, dilation=1, bias=true)
    kernel_tuple = kernel_size isa Tuple ? kernel_size : (kernel_size, kernel_size, kernel_size)
    # init weights same as in https://pytorch.org/docs/stable/nn.html#conv2d
    k_sqrt = sqrt(1 / (in_channels * prod(kernel_tuple)))
    d = Uniform(-k_sqrt, k_sqrt)
    W = rand(d, kernel_tuple..., in_channels, out_channels)
    b = bias ? rand(d, out_channels) : nothing
    return Conv3d(W, b, stride, padding, dilation)
end
Conv3d(in_out::Pair{Int, Int}, k) = Conv3d(in_out[1], in_out[2], k)

function Base.show(io::IO, c::Conv3d)
    k1, k2, k3, i, o = size(c.W)
    print(io, "Conv3d($i=>$o, $(k1)x$(k2)x$(k3))")
end

function (c::Conv3d)(x::AbstractArray{T,5}) where T
    y = convNd(x, c.W; stride=c.stride, padding=c.padding, dilation=c.dilation)
    if c.b != nothing
        y = y .+ reshape(c.b, (1, 1, length(c.b), 1))
    end
    return y
end

################################################################################
#                           Loss Functions                                     #
################################################################################

mutable struct NLLLoss
end

Base.show(io::IO, loss::NLLLoss) = print(io, "NLLLoss()")

(loss::NLLLoss)(logp::AbstractMatrix, c::Union{AbstractMatrix{<:Real}, AbstractVector{<:Real}}) =
    nllloss(logp, c)

mutable struct CrossEntropyLoss
end

Base.show(io::IO, loss::CrossEntropyLoss) = print(io, "CrossEntropyLoss()")

(loss::CrossEntropyLoss)(x::AbstractMatrix, c::Union{AbstractMatrix{<:Real}, AbstractVector{<:Real}}) =
    crossentropyloss(x, c)

mutable struct MSELoss
end

Base.show(io::IO, loss::MSELoss) = print(io, "MSELoss()")

(loss::MSELoss)(x::AbstractArray, x_target::AbstractArray) = mseloss(x, x_target)

################################################################################
#                                 Common                                       #
################################################################################

function trainmode!(m, train::Bool=true)
    # by default, recursively call trainmode!() on all fields
    T = typeof(m)
    for fld in fieldnames(T)
        trainmode!(getfield(m, fld), train)
    end
end
testmode!(m) = trainmode!(m, false)


using StructJuMP

NS = 3;                        # number of scenarios
probability = [1/3, 1/3, 1/3]; # probability

CROPS  = 1:3 # set of crops (wheat, corn and sugar beets, resp.)
PURCH  = 1:2 # set of crops to purchase (wheat and corn, resp.)
SELL   = 1:4 # set of crops to sell (wheat, corn, sugar beets under 6K and those over 6K)

Cost     = [150., 230., 260.]    # cost of planting crops
Budget   = 500.                # budget capacity
Purchase = [238., 210.];        # purchase price
Sell     = [170., 150., 36., 10.] # selling price
Yield    = [3.0 3.6 24.0;
            2.5 3.0 20.0;
            2.0 2.4 16.0]
Minreq   = [200., 240., 0.]      # minimum crop requirement

function farmer_stochastic()
    m = StructuredModel(num_scenarios = NS)

    @variable(m, x[i=CROPS] >= 0, Int)
    @objective(m, Min, sum(Cost[i] * x[i] for i=CROPS))
    @constraint(m, sum(x[i] for i=CROPS) <= Budget)

    for s in 1:NS
        blk = StructuredModel(parent = m, id = s, prob = probability[s])
        @variable(blk, y[j=PURCH] >= 0)
        @variable(blk, w[k=SELL] >= 0)
        @objective(blk, Min,
            sum(Purchase[j] * y[j] for j=PURCH)
            - sum(Sell[k] * w[k] for k=SELL))
        @constraint(blk, const_minreq[j=PURCH], Yield[s,j] * x[j] + y[j] - w[j] >= Minreq[j])
        @constraint(blk, const_minreq_beets, Yield[s,3] * x[3] - w[3] - w[4] >= Minreq[3])
        @constraint(blk, const_aux, w[3] <= 6000)
    end

    return m
end

function farmer_blocks()
    m = StructuredModel(num_scenarios = NS)

    @variable(m, x[i=CROPS,s=1:NS] >= 0, Int)
    @variable(m, y[j=PURCH,s=1:NS] >= 0)
    @variable(m, w[k=SELL,s=1:NS] >= 0)
    @objective(m, Min,
        sum(probability[s] * Cost[i] * x[i,s] for i=CROPS for s=1:NS)
        + sum(probability[s] * Purchase[j] * y[j,s] for j=PURCH for s=1:NS)
        - sum(probability[s] * Sell[k] * w[k,s] for k=SELL for s=1:NS))
    @constraint(m, nonant[i=CROPS,s=2:NS], x[i,s-1] - x[i,s] == 0)

    for s in 1:NS
        blk = StructuredModel(parent = m, id = s)
        @objective(blk, Min, 0.)
        @constraint(blk, sum(x[i,s] for i=CROPS) <= Budget)
        @constraint(blk, const_minreq[j=PURCH], Yield[s,j] * x[j,s] + y[j,s] - w[j,s] >= Minreq[j])
        @constraint(blk, const_minreq_beets, Yield[s,3] * x[3,s] - w[3,s] - w[4,s] >= Minreq[3])
        @constraint(blk, const_aux, w[3,s] <= 6000)
    end
    return m
end


import YAML
config = YAML.load(open("config.yaml"))

using DiskJockey.constants
using DiskJockey.model
using DiskJockey.visibilities
using DiskJockey.image
using DiskJockey.gridding
using MultivariateStats
using HDF5

dv = DataVis(config["data_file"], true)[5]

# Conjugation is necessary for the SMA and ALMA
visibilities.conj!(dv)

qq = get_qq(dv)

# Find the indices that sort these in increasing order
ind = sortperm(qq)
qq = qq[ind]
uu = dv.uu[ind]
vv = dv.vv[ind]

# Points at which to create images
parameters = readdlm("parameters.dat")
npoints = size(parameters, 1)

fid = h5open("grid_VV.hdf5", "r")
# Load the sampled visibilities from HDF5
reals = read(fid, "real")
imags = read(fid, "imag")
VV = reals + imags .* im # Complex visibility

close(fid)

nvis = size(reals, 1)

println("Dimensions of input matrix ", size(reals))

# suppose Xtr and Xte are training and testing data matrix,
# with each observation in a column

# Need to do the reals and imaginaries separately, since MultivariateStats does not work w/ Complex128

mean_estimate = mean(reals, 2)
std_estimate = std(reals, 2)

norm_reals = (reals .- mean_estimate) ./ std_estimate
# norm_reals = reals .- mean_estimate

# train a PCA model
M = fit(PCA, norm_reals; mean=0, pratio=0.999)

dump(M)

eigen = projection(M)

println("projection", size(eigen))

ncomp = outdim(M)

# a (npoints, ncomp) matrix
weights = Array(Float64, (npoints, ncomp))

# Calculate the weights that lead to reconstruction of the visibilities
for i=1:npoints
  for j=1:ncomp
    weights[i,j] = sum(reals[:,i] .* eigen[:,j])
  end
end

# Now try reconstructing the spectra
reconstructed_reals = Array(Float64, (nvis, npoints))
for i=1:npoints
  reconstructed_reals[:,i] = std_estimate .* ((eigen * squeeze(weights[i,:], 1)) .+ mean_estimate)
  # reconstructed_reals[:,i] = (eigen * squeeze(weights[i,:], 1)) .+ mean_estimate
end

ms=0.1

# Compare the reconstructed spectra vs the reals.
import PyPlot.plt
using LaTeXStrings

fig, ax = plt[:subplots](nrows=3, figsize=(10, 6), sharex=true)

for i=1:3
  ax[1][:plot](qq, reals[:,i], ".", ms=ms)
  ax[2][:plot](qq, reconstructed_reals[:,i], ".", ms=ms)
  ax[3][:plot](qq, reals[:,i] - reconstructed_reals[:,i], ".", ms=ms)
end

ax[1][:set_xlabel](L"$qq\quad k\lambda$")
ax[1][:grid]()
ax[2][:grid]()
ax[3][:grid]()

plt[:savefig]("visibilities_reconstructed.png", dpi=600)

fig, ax = plt[:subplots](nrows=3, figsize=(10, 6), sharex=true)

ax[1][:plot](qq, mean_estimate, ".", ms=ms)
ax[2][:plot](qq, std_estimate, ".", ms=ms)

for i=1:size(eigen, 2)
  ax[3][:plot](qq, eigen[:,i], ".", ms=ms)
end

ax[3][:set_xlabel](L"$qq\quad k\lambda$")

ax[1][:grid]()
ax[2][:grid]()
ax[3][:grid]()

ax[1][:set_ylabel]("Mean")
ax[2][:set_ylabel]("Std")
ax[3][:set_ylabel]("Real")

plt[:savefig]("visibilities_PCA_real.png", dpi=600)


using DataFrames, Plots, MultivariateStats, GLM

# The code is taken from the original
# Bernanke Boivin and Eliasz paper
# FAVAR Bayesian structural factor augmented VAR
# -------------------------------------------------------------------------
# The model is:
#
#   FACTOR (OBSERVATION) EQUATION:
#     _      _           _      _
#    |  X(t)  |         |  F(t)  |
#    |        |  = L  x |        | + e(t)
#    |_ Y(t) _|         |_ Y(t) _|
#
#
#   VAR (STATE) EQUATION:
#     _      _            _        _
#    |  F(t)  |          |  F(t-1)  |
#    |        |  = PHI x |          |  +  u(t)
#    |_ Y(t) _|          |_ Y(t-1) _|
#
#
#  with u(t)~N(0,S_F), and e(t)~N(0,Sigma), with Sigma = diag{(sigma_1)^2 ,
#  ...., (sigma_N)^2}.
# -------------------------------------------------------------------------
#-----------------------------LOAD DATA------------------------------------
# Load quarterly data
# load data used to extract factors
xdata = readtable("/home/lbrugnol/Downloads/FAVAR/xdata.csv", header = true)
# load data on inflation, unemployment and interest rate
ydata = readtable("/home/lbrugnol/Downloads/FAVAR/ydata.csv", header = false)
# load transformation codes (see file transx.m)
# load the slow or fast moving variables codes (see Bernanke, Boivin and
# Eliasz, 2005, QJE)
tcode = readtable("/home/lbrugnol/Downloads/FAVAR/tcode.csv", header=false)
slowcode = readtable("/home/lbrugnol/Downloads/FAVAR/slowcode.csv", header = false)
# load the file with the dates of the data (quarters)
yearlab = readtable("/home/lbrugnol/Downloads/FAVAR/yearlab.csv", header = false)
# load the file with the names of the variables.
namesX = readtable("/home/lbrugnol/Downloads/FAVAR/namesX.csv", header = false)

# namesX.mat contains only the short codes of the data. For a complete
# description see the attached excel file.
stddata = [std(convert(Array,xdata),1) ones(1,3)]

const dRule = Dict{String,Integer}("none" => 1, "first_diff" => 2, "log_yoy"=> 3,
"log" => 4, "log_first_diff" => 5, "yoy" => 6)
const sFile_sd = "/home/lbrugnol/Downloads/FAVAR/xdata.csv"
const bHeader_sd = true
const sFile_tc = "/home/lbrugnol/Downloads/FAVAR/tcode.csv"
const bHeader_tc = false
# const sStart = "1999-01-31"

# Call starting dataset, transformation and create final dataset
sd = StartData(sFile_sd,bHeader_sd)
tc = TransRule(sFile_tc,bHeader_tc,dRule,sd)
fd = FinalData(sd,tc)

# Correct size after stationarity transformation
xdata = convert(Array,fd.dfData[3:end,:])
ydata = convert(Array,ydata[3:end,:])
yearlab = yearlab[3:end,:]

# Demean data (no intercept is estimated as in BBE (2005))
t1 = size(xdata,1)     # time series observations of xdata
t2 = size(ydata,1)     # time series dimension of ydata
stdffr = std(ydata[:,3],1)       # standard deviation of the Fed Funds Rate
xdata = xdata .- mean(xdata,1)
ydata = ydata .- mean(ydata,1)

# Define X and Y matrices
X = xdata    # X contains the 'xdata' which are used to extract factors.
Y = ydata  # Y contains inflation, unemployment and interest rate
# NamesXY has the short codes of X and Y
namesXY = [fd.sd.vNames; "Inflation"; "Unemployment"; "Fed_funds"]

# Number of observations and dimension of X and Y
T,M=size(Y) # T time series observations
N=size(X,2) # N series from which we extract factors
#----------------------------PRELIMINARIES---------------------------------
# Set some Gibbs - related preliminaries
const nrep = 15000  # Number of replications
const nburn = 5000   # Number of burn-in-draws
const nthin = 1   # Consider every thin-th draw (thin value)
const it_print = 100  #Print in the screen every "it_print"-th iteration
const constant = 1  # Set 0:no constant on the FAVAR equation, 1:include constant

# Number of factors & lags:
const K = 2               # Number of Factors
const p = K+M             # p is the dimensionality of [Factors, Y]
const plag = 2            # plag is number of lags in the VAR part
# ==============================| FACTOR EQUATION |========================
# Extract principal components from X (to be used as starting values for the
# MCMC estimation of the factors)
X_st = X./std(X,1)# First standardize data to extract PC
[F0,Lf]=extract(X_st,K)  # F0 are the factors, Lf are the loadings

##### UNDERSTAND WHY DIFFERENT FACTORS WRT MATLAB!!!!!!

function pca(data::Array{Float64}; n::Int64 = 1)
    # Need MultivariateStats
    mX = data.-mean(data,1)
    M = fit(PCA, mX, method = :cov)
    return projection(M)[:,1:n]
  end

F0 = pca(X_st,n = K)

#----If using PC estimates of the factors:
# Now rotate the factor space as in Bernanke, Boivin and Eliasz (2005)
slowindex = find(convert(Array,slowcode).==1)
xslow = X[:,slowindex]
Fslow0 = pca(xslow,n = K)

Ffast0 = Y[:,end]
t1,k1=size(Ffast0[:,:])
ff = ([ones(t1,1) Ffast0 Fslow0]'*[ones(t1,1) Ffast0 Fslow0])
fF = [ones(t1,1) Ffast0 Fslow0]'*F0
b = ff\fF # original paper ols svd
Fr0 = F0 - Ffast0*b[2:k1+1,:]

X=X_st
Y=Y./std(Y)

# #----If using MCMC estimates of the factors:
# # This part is explained in the PhD thesis of Eliasz (2005). Essentially,
# # we rotate the factors in a certain way which gives a loadings matrix L,
# # with the upper KxK block being the identity matrix. This structure of L
# # is the same structure we will have when using MCMC to estimate the
# # factors (because of the identification restrictions imposed on L).
# # regress X on F0 and Y, obtain loadings
# Lfy=olssvd(X(:,K+1:N),[F0 Y])'     # upper KxM block of Ly set to zero
# Lf=[Lf(1:K,:);Lfy(:,1:K)]
# Ly=[zeros(K,M);Lfy(:,K+1:K+M)]
# # transform factors and loadings for LE normalization
# [ql,rl]=qr(Lf')
# Lf=rl  # do not transpose yet, is upper triangular
# F0=F0*ql
# # need identity in the first K columns of Lf, call them A for now
# A=Lf(:,1:K)
# Lf=[eye(K),inv(A)*Lf(:,(K+1):N)]'
# Fr0=F0*A

# Put it all in state-space representation, write obs equ as XY=FY*L+e
XY=[X Y]    # the data X used to extract factors, plus Y (infl,unemp and interest)
FY=[Fr0 Y]  # the extracted factors, plus Y (infl,unemp and interest)

# Obtain L (the loadings matrix)
#----If using PC estimates of the factors:
L = (FY'*FY)\FY'*XY # original paper ols svd decom

# #----If using MCMC estimates of the factors:
# L=[Lf Ly;zeros(M,K),eye(M)]

# Obtain SIGMA (the error varaince in the factor equation)
e = XY - FY*L
SIGMA = e'*e./T
SIGMA = diagm(diag(SIGMA))

# ================================| VAR EQUATION |=========================
# Generate lagged FY matrix.
function mlag2(mX,p)
    T,N = size(mX)
    mXlag = zeros(T,N*p)
    for ii in 1:p
        mXlag[p+1:T,(N*(ii-1)+1):N*ii] = mX[p+1-ii:T-ii,1:N]
    end
    return mXlag
end

FY_lag = mlag2(FY,plag)

FY_lag = FY_lag[plag+1:T,:]
FY_temp = FY[plag+1:T,:]

# Get an initial estimate of PHI (FAVAR autoregressive coefficients) and
# S_F (FAVAR covariance matrix). These values will be used as initial
# values for the Gibbs sampler (but you can use 'zeros' instead, to declare
# 'ignorance' about the initial values
PHI = inv(FY_lag'*FY_lag)*(FY_lag'*FY_temp)
SSE = (FY_temp - FY_lag*PHI)'*(FY_temp - FY_lag*PHI)
S_F = SSE./(T-p)

# If plag is not 1, we need to write the state equation as a VAR(1) model,
# in order to be able to estimate the latent factors using the Kalman
# filter. In this case, PHI and S_F are the parameter matrices of the VAR(p)
# model, while PHI_mat and S_F_mat are the matrices of the VAR(1)
# transformed model (see the time series book of Luthkephol (2005), for
# example). If using PC estimates this is not necessary, but we will need
# the VAR(1) form anyway for the impulse responses, so do not comment this
# part.
S_F_mat = [S_F zeros(p,p*(plag-1));zeros(p*(plag-1),p*plag)]
PHI_mat = [PHI' ; eye(p*(plag-1)) zeros(p*(plag-1),p)]

#-------------------------------- PRIORS ----------------------------------
# #========= INITIAL CONDITION ON LATENT FACTORS:
# # F(0) ~ N(F_0_prmean, F_0_prvar)
# F_0_prmean = zeros(p*plag,1);
# F_0_prvar = 4*eye(p*plag)
#
# # variable indexnM indexes which draws of FY refer only to the factors F
# # (and not Y, since Y is assumed to be observed and known). See also
# # function kfgibbsnv.m
# indexnM = [ones(K,plag);zeros(M,plag)]
# indexnM = find(indexnM==1)

# PRIORS ON FACTOR EQUATION:
# Prior on loadings L_i ~ N(0, I), where i = 1,...,N
Li_prvar = 4*eye(p)
alpha = 0.01
beta = 0.01

# Prior on covariance SIGMA_i ~ iG(a, b), where i = 1,...,N
alpha = 0.01
beta = 0.01

# PRIORS ON VAR EQUATION:

# IMPULSE RESPONSES:
# Note that impulse response and related stuff involves a lot of storage
# and, hence, put istore=0 if you do not want them
istore = 1
if istore == 1
    # Impulse response horizon
    nhor = 21
    shock_init = diagm([zeros(p-1); 1./stdffr]) # in terms of standard deviation, identification is recursive
    imp = zeros(nrep,N+M,nhor)
    bigj = zeros(p,p*plag)
    bigj[1:p,1:p] = eye(p)
end
#-------------------------- END OF PRELIMINARIES --------------------------

# START SAMPLING
println("Number of iterations")

for irep = 1:nrep + nburn    #  GIBBS iterations starts here
    # Print iterations
    # if mod(irep,it_print) == 0
    #     disp(irep); toc
    # end

    #FACTOR (MEASUREMENT) EQUATION:
    # -----------------------------------------------------------------------------------------
    #   STEP I: Sample latent factors F_t using Carter and Kohn (1994) or
    #   use PC estimates
    # -----------------------------------------------------------------------------------------

    # #----If using MCMC estimates of the factors:
    # FY = kfgibbsnv(XY,F_0_prmean,F_0_prvar,L,SIGMA,PHI_mat,S_F_mat,M,indexnM)
    #
    # # Make sure the factors have mean zero
    # FY = FY - repmat(mean(FY),T,1)

    # Commenting the above, means that we are using the PC estimate of FY.

    # -----------------------------------------------------------------------------------------
    #   STEP II: Sample L and SIGMA
    # -----------------------------------------------------------------------------------------
    # Since the covariance matrix of the error (SIGMA) in this equation is
    # diagonal, we can estimate the parameters equation-by-equation
    for i = 1:N
        # Sample L from a Normal distribution. The upper KxK block of L is
        # the identity matrix, so we sample the rest N-K rows only:
        if i > K
            Li_postvar = inv(inv(Li_prvar) + inv(SIGMA[i,i])*FY'*FY)
            Li_postmean = Li_postvar*(inv(SIGMA[i,i])*FY'*X[:,i])
            Lidraw = Li_postmean' + randn(1,p)*chol(Hermitian(Li_postvar))
            L'[i,1:p] = Lidraw
        end

        ed = X[:,i] - FY*L'[i,:]

        # Sample SIGMA(i,i) from iGamma
		S_1 = alpha/2 + T/2
		S_2 = beta/2 + ((X[:,i] - FY*L'[i,:])'*(X[:,i] - FY*L'[i,:]))/2
        Sidraw = inv(rand(Gamma(S_1,1/S_2)))
        SIGMA[i,i] = Sidraw
    end

    #VAR (STATE) EQUATION:
    # Generate lagged FY and then correct sizes (we loose observations
    # because of taking lags, as with the simple VAR model)
    FY_lag = mlag2(FY,plag)
    FY_lag = FY_lag[plag+1:T,:]

    FY_temp = FY[plag+1:T,:]
    # Now we have everything we need to estimate PHI and S_F, as we did
    # with the simple BVAR models. You can use any of the priors we used in
    # the BVAR models, like the Diffuse, the Normal-Wishart, the Independent
    # Normal-Wishart, or the Minnesota (Littermann) prior. Here I use the
    # Diffuse prior, based on the OLS quantities (remember, at each Gibbs
    # iteration we treat the dependent variables (FY_temp) and the lags of
    # the dependent variables (FY_lag) as observed data.

    # -----------------------------------------------------------------------------------------
    #   STEP I: Sample autoregressive coefficients PHI
    # -----------------------------------------------------------------------------------------

    PHI_OLS = inv(FY_lag'*FY_lag)*(FY_lag'*FY_temp)
    phi_OLS = PHI_OLS(:)

    V_post = kron(S_F,inv(FY_lag'*FY_lag))
    phi_vec = phi_OLS + chol(V_post)'*randn(p*p*plag,1)
    PHI = reshape(phi_vec,p*plag,p)'

    # Now create draw of PHI_mat, the parameters in the VAR(1) transformation
    PHI_mat(1:p,:) = PHI

#     # truncate to ensure stationarity
#     while max(abs(eig(PHI_mat)))>0.999
#         phi_vec = phi_OLS + chol(V_post)'*randn(p*p*plag,1)
#         PHI = reshape(phi_vec,p*plag,p)'
#         PHI_mat(1:p,:) = PHI
#     end
    # -----------------------------------------------------------------------------------------
    #   STEP II: Sample covariance matrix S_F
    # -----------------------------------------------------------------------------------------
    SSE = (FY_temp - FY_lag*PHI_OLS)'*(FY_temp - FY_lag*PHI_OLS)
    S_F = inv(wish(inv(SSE),T-p))

    # Now create draw of S_F_mat, the parameter in the VAR(1) transformation
    S_F_mat(1:p,1:p) = S_F

    #--------------------------SAVE AFTER-BURN-IN DRAWS AND IMPULSE RESPONSES -----------------
    if irep > nburn && mod((irep-nburn),nthin)==0

        # Save here draws of the parameters to get means and variances

        #----------------- Impulse response analysis:
        if istore==1
            # Note that Htsd contains the
            # structural error cov matrix
            # Set up things in VAR(1) format as in Lutkepohl page 11
            biga = PHI_mat

            # st dev matrix for structural VAR
            shock = chol(S_F)'
            d = diag(diag(shock))
            shock = inv(d)*shock
            #shock = Hsd2*shock_init
            #shock=eye(p)

            #now get impulse responses for 1 through nhor future periods@
            impresp = zeros(p,p*nhor)
            impresp(1:p,1:p) = shock
            bigai = biga
            for j = 1:nhor-1
                impresp(:,j*p+1:(j+1)*p) = bigj*bigai*bigj'*shock
                bigai = bigai*biga
            end

            imp_m = zeros(p,nhor)
            jj=0
            for ij = 1:nhor
                jj = jj + p
                imp_m(:,ij) = impresp(:,jj)
            end
            imp(irep-nburn,:,:) = L*imp_m
       end #END the impulse response calculation section
    end # END saving after burn-in results
end #END main Gibbs loop (for irep = 1:nrep+nburn)
#=============================GIBBS SAMPLER ENDS HERE==================================

if istore == 1
    scale=repmat(stddata',[1 3 nhor])
    scale=permute(scale,[2 1 3])

    imp = permute(imp,[2 3 1])

    for i=1:N
        if tcode(i)==4
            imp(i,:,:) = exp(imp(i,:,:))-ones(1,nhor,nrep/nthin)
        elseif tcode(i)==5
            imp(i,:,:)=exp(cumsum(imp(i,:,:),2))-ones(1,nhor,nrep/nthin)
        end
    end
    imp = permute(imp,[3 1 2])

    # Set quantiles from the posterior density of the impulse responses
    qus = [.10, .5, .9]
    #if using PC estimates of the factors, use this
    impXY=squeeze(quantile(imp,qus))
#     #if using MCMC estimates of the factors, use this
#     impXY=squeeze(quantile(imp,qus))./scale

    #====================PLOTS
    #---Plot I: impulse responses of inflation, unemployment, interest
    figure
    set(0,'DefaultAxesColorOrder',[0 0 0],...
        'DefaultAxesLineStyleOrder','--|-|--')
    subplot(3,1,1)
    plot(1:nhor,squeeze(impXY(:,end-2,:)))
    hold
    plot(zeros(1,nhor),'-')
    title('Impulse response of inflation')
    xlim([1 nhor])
    set(gca,'XTick',0:3:nhor)
    subplot(3,1,2)
    plot(1:nhor,squeeze(impXY(:,end-1,:)))
    hold
    plot(zeros(1,nhor),'-')
    title('Impulse response of unemployment')
    xlim([1 nhor])
    set(gca,'XTick',0:3:nhor)
    subplot(3,1,3)
    plot(1:nhor,squeeze(impXY(:,end,:)))
    hold
    plot(zeros(1,nhor),'-')
    title('Impulse response of interest rate')
    xlim([1 nhor])
    set(gca,'XTick',0:3:nhor)
    #-------------------------

    #---Plot II: impulse responses of other variables
    figure
    set(0,'DefaultAxesColorOrder',[0 0 0],...
        'DefaultAxesLineStyleOrder','--|-|--')
    # I will plot only 12 out of the 115 variables here
    var_numbers = [2  9 10 28 42 46 77 91 92 108 109 111]
    # These variables have the following short codes:
    var_names = namesXY(var_numbers)

    for i=1:12
        subplot(4,3,i)
        plot(1:nhor,squeeze(impXY(:,var_numbers(i),:)))
        hold
        plot(zeros(1,nhor),'-')
        title(['Impulse response of ' var_names(i)])
        xlim([1 nhor])
        set(gca,'XTick',0:3:nhor)
    end
    #-------------------------

end

clc
toc # Stop timer and print total time


@testset "Simplify" begin
  Qx, x = PolynomialRing(FlintQQ, "x")
  K, a = NumberField(x^2 - 10, cached = false)
  L, mL = simplify(K)
  @test is_isomorphic(K, L)

  K, a = NumberField(x^2 - 100000000000000000000000, cached = false)
  L, mL = simplify(K)
  @test is_isomorphic(K, L)

  K, a = NumberField(x - 10, cached = false)
  L, mL = simplify(K)
  @test is_isomorphic(K, L)

  K,a = NumberField(x^4-100020001*x^2+100040006000400010, cached = false)
  L, mL = simplify(K)
  @test is_isomorphic(K, L)

  K, a = number_field(8*x^3 + 4*x^2 - 4*x - 1, cached = false)
  L, mL = simplify(K)
  @test Hecke.is_defining_polynomial_nice(L)
  L1, mL1 = simplify(K, canonical = true)
  @test Hecke.is_defining_polynomial_nice(L1)
  @test is_isomorphic(L1, L)

  f = -1//3*x^2 - x + 1
  K, a = number_field(f, cached = false)
  L, mL = simplify(K)
  @test is_isomorphic(K, L)

  f = x - 1
  K, a = number_field(f, cached = false)
  _, g = Hecke.polredabs(K)
  @test g == f
end


# This version shows how verifiability for proof of shuffle is added

using ShuffleProofs
import ShuffleProofs: PrimeGenerator, Enc, Dec, shuffle, ElGamal, ProtocolSpec, verify, Shuffle

sk = 123 # Only authorithy who do decryption would know

g = PrimeGenerator(3, 23)
options = [g, g^2, g^3]
pk = g^sk

verifier = ProtocolSpec(; g) # Some default parameters. In practice needs to be chosen carefully.

bbord = (;) # A representation of a buletinboard

######## Step 1. ########

bbord = let
    enc = Enc(pk, g)
    ciphertexts_in = ElGamal([enc(options[rand(1:3)], rand(1:10)) for i in 1:10])
    (; bbord..., ciphertexts_in)
end

######## Step 2. ########

bbord = let
    enc = Enc(pk, g)
    simulator = shuffle(bbord.ciphertexts_in, enc, verifier) # now returns simulator as verifier is added
    ciphertexts_out = simulator.proposition.𝐞′
    (; bbord..., ciphertexts_out, simulator.proof)
end

####### Step 3. ########

bbord = let
    dec = Dec(sk)
    votes = dec(bbord.ciphertexts_out)
    (; bbord..., votes)
end

######### As last step we can tally the votes ############

# Before counting we perform an audit that reencryption have happened honestly
proposition = Shuffle(g, pk, bbord.ciphertexts_in, bbord.ciphertexts_out)
@assert verify(proposition, bbord.proof, verifier)

for (i, o) in enumerate(options)
    s = 0
    for v in bbord.votes
        if v == o
            s += 1
        end
    end

    println("Choice $i received $s votes")
end


@testset "points.jl" begin
    p1 = rad2deg(Point(0.5, 1.0))
    p2 = rad2deg(Point(1.0, 0.5))
    p3 = rad2deg(Point(0.6, 1.0))  # above p1

    rs = RouteSection(p1,p2)

    @test midpoint(p1, p2).λ ≈ rad2deg(0.8106650896663075) atol = 0.0001
    @test midpoint(p1, p2).ϕ ≈ rad2deg(0.76484691145) atol = 0.0001
    @test midpoint(rs).λ ≈ rad2deg(0.8106650896663075) atol = 0.0001
    @test midpoint(rs).ϕ ≈ rad2deg(0.76484691145) atol = 0.0001

    @test intermediate_point(p1, p2, 0.5).λ ≈ rad2deg(0.8106650896663075) atol = 0.0001
    @test intermediate_point(p1, p2, 0.5).ϕ ≈ rad2deg(0.76484691145) atol = 0.0001
    @test intermediate_point(rs, 0.5).λ ≈ rad2deg(0.8106650896663075) atol = 0.0001
    @test intermediate_point(rs, 0.5).ϕ ≈ rad2deg(0.76484691145) atol = 0.0001

    dst = 3888000.0
    brg = 5.81412807071
    @test destination_point(p1, dst, rad2deg(brg)).ϕ ≈ rad2deg(1.0) atol = 0.01
    @test destination_point(p1, dst, rad2deg(brg)).λ ≈ rad2deg(0.5) atol = 0.01

    brg13 = 0.5
    brg23 = 1.5
    brg13_ambiguous = 2.61799387799
    @test isnan(intersection_point(p2, p1, rad2deg(brg13_ambiguous), rad2deg(brg23)).ϕ) == true

    @test intersection_point(p1, p2, rad2deg(brg13), rad2deg(brg23)).ϕ ≈ rad2deg(0.86121333) atol = 0.001
    @test intersection_point(p1, p2, rad2deg(brg13), rad2deg(brg23)).λ ≈ rad2deg(1.31431535) atol = 0.001

    p1ni = Point(50.0, 10.0)
    p2ni = Point(50.0, 30.0)
    p3ni = Point(40.0, 20.0)
    p4ni = Point(60.0, 20.0)

    @test intersection_point(p1ni, p2ni, p3ni, p4ni).ϕ ≈ 50.4313888888 atol = 0.0001
    @test intersection_point(p1ni, p2ni, p3ni, p4ni).λ ≈ 20.0 atol = 0.0001
    @test intersection_point(p1ni, p2ni, p3ni, p2ni).λ ≈ 30.0 atol = 0.0001

    p4ni2 = Point(49.0, 20.0)
    @test isnan(intersection_point(p1ni, p2ni, p3ni, p4ni2).λ)

    p1a = Point(10.0374, 0.0)
    p2a = Point(-90.0, 0.0)
    p3a = Point(0.0, 0.0)
    @test isinf(intersection_point(p3a, p2a, p1a, p3a).ϕ)

    p7 = Point(5.0, 0.0)
    p8 = Point(-5.0, 0.0)
    brg7 = 0.0
    brg8 = 0.0
    @test isinf(intersection_point(p7, p8, brg7, brg8).λ)

    #TODO Check if intersection_point also works when crossing the datum line

    #TODO Uncomment after network.jl has been added
    # ap = Airspace("test", [Point(50.0, 10.0), Point(50.0, 30.0),
    # Point(60.0, 30.0), Point(60.0, 10.0), Point(50.0, 10.0)])
    #
    # bx = box_spherical_polygon(ap.polygon)
    # @test bx[3].ϕ ≈ 60.378348124804496 atol = 0.00001
    # # @test isnan(intersection_point(Point(40.0, 20.0), Point(50.0, 20.0), ap).ϕ)

    @test opposite_point(p1).ϕ ≈ -28.6478897565 atol = 0.0001
    @test opposite_point(p1).λ ≈ -122.704220487 atol = 0.0001

    @test max_latitude(0.0, 45.0) ≈ 45.0 atol = 0.0001
    @test max_latitude(rad2deg(1.0), 90.0) ≈ rad2deg(1.0) atol = 0.0001
    @test max_latitude(rad2deg(1.0), 0.0) ≈ 90.0 atol = 0.0001
    @test max_latitude(rad2deg(-1.0), 0.0) ≈ 90.0 atol = 0.0001
    @test max_latitude(rad2deg(-1.0), 180.0) ≈ 90.0 atol = 0.0001

    p5 = Point(45.0, 45.0)
    @test closest_point_to_pole(p5, 45.0).ϕ ≈ 60.0 atol = 0.001
    @test closest_point_to_pole(p5, 45.0).λ ≈ 99.73561 atol = 0.001
    @test closest_point_to_pole(p5, 89.0).ϕ ≈ 45.00873 atol = 0.001
    @test closest_point_to_pole(p5, 89.0).λ ≈ 46.41407 atol = 0.001
    @test closest_point_to_pole(p5, 135.0).ϕ ≈ -60.0 atol = 0.001
    @test closest_point_to_pole(p5, 135.0).λ ≈ 170.2644 atol = 0.001
    @test closest_point_to_pole(p5, -135.0).ϕ ≈ -60.0 atol = 0.001
    @test closest_point_to_pole(p5, -135.0).λ ≈ -80.2644 atol = 0.001
    @test closest_point_to_pole(p5, -89.0).ϕ ≈ 45.00873 atol = 0.001
    @test closest_point_to_pole(p5, -89.0).λ ≈ 43.58593 atol = 0.001
    @test closest_point_to_pole(p5, -45.0).ϕ ≈ 60.0 atol = 0.001
    @test closest_point_to_pole(p5, -45.0).λ ≈ -9.73561 atol = 0.001

    p6 = Point(-45.0, 45.0)
    @test closest_point_to_pole(p6, 45.0).ϕ ≈ 60.0 atol = 0.001
    @test closest_point_to_pole(p6, 45.0).λ ≈ 170.2644 atol = 0.001
    @test closest_point_to_pole(p6, 91.0).ϕ ≈ -45.00873 atol = 0.001
    @test closest_point_to_pole(p6, 91.0).λ ≈ 46.41407 atol = 0.001
    @test closest_point_to_pole(p6, 135.0).ϕ ≈ -60.0 atol = 0.001
    @test closest_point_to_pole(p6, 135.0).λ ≈ 99.73561 atol = 0.001
    @test closest_point_to_pole(p6, -135.0).ϕ ≈ -60.0 atol = 0.001
    @test closest_point_to_pole(p6, -135.0).λ ≈ -9.73561 atol = 0.001
    @test closest_point_to_pole(p6, -91.0).ϕ ≈ -45.00873 atol = 0.001
    @test closest_point_to_pole(p6, -91.0).λ ≈ 43.58593 atol = 0.001
    @test closest_point_to_pole(p6, -45.0).ϕ ≈ 60.0 atol = 0.001
    @test closest_point_to_pole(p6, -45.0).λ ≈ -80.2644 atol = 0.001
end


using ErrorStateKalmanFilter
using Test
using LinearAlgebra
using Quaternions

@testset "types" begin
    x = collect(range(0, length=19, step=0.5))
    @test vec(NominalState(x)) == x
    x = collect(range(0, length=18, step=0.1))
    @test vec(ErrorState(x)) == x
end

@testset "ESKF" begin
    @test skew_matrix([1, 2, 3]) == [0 -3 2; 3 0 -1; -2 1 0]

    dt = 1e-3
    times = 0:dt:5

    nominal = NominalState(
        [0.0, 0, 0],
        [0.0, 0, 0],
        quat(1.0),
        zeros(3),
        zeros(3),
        [0.0, 0.0, -1.0]
    )

    error = ErrorState(zeros(18))

    error_cov = Matrix{Float64}(I, 18, 18)

    measurement = Measurement([0.0, 0, 0], [0.0, 0, 0])

    noise = Noise(
        2e-7  * (1.0/512),
        1e-10 * (1.0/512),
        1e-9,
        1e-9
    )

    new_nominal, new_error, new_error_cov = predict(nominal, error, error_cov, measurement, noise, dt)

    # error state is unchanged
    @test new_error == error
    # error covariance increases
    @test all(abs.(new_error_cov) .>= abs.(error_cov))
end


using DistributedFactorGraphs
using IncrementalInference
using Manifolds
using StaticArrays
using Test

@testset "Test TranslationGroup(2)" begin
## ======================================================================================
##
## ======================================================================================

Base.convert(::Type{<:Tuple}, M::TranslationGroup{Tuple{2},ℝ}) = (:Euclid, :Euclid)
Base.convert(::Type{<:Tuple}, ::IIF.InstanceType{TranslationGroup{Tuple{2},ℝ}})  = (:Euclid, :Euclid)

@defVariable TranslationGroup2 TranslationGroup(2) [0.0, 0.0]
getManifold(TranslationGroup2)
getPointType(TranslationGroup2)
getPointIdentity(TranslationGroup2)

fg = initfg()

v0 = addVariable!(fg, :x0, TranslationGroup2)
v1 = addVariable!(fg, :x1, TranslationGroup2)

mp = ManifoldPrior(TranslationGroup(2), SA[10., 20], MvNormal([1.0,1.0]))
p = addFactor!(fg, [:x0], mp;  graphinit=true)

doautoinit!(fg, :x0)

mf = ManifoldFactor(TranslationGroup(2), MvNormal([1.0, 2.0], [0.1,0.1]))
f = addFactor!(fg, [:x0, :x1], mf)

solveGraph!(fg)

end


module _EnofEn
export EnofEn
using StatsBase: countmap, Histogram, fit
    """
        EoE, AvEn, S2 = EnofEn(Sig)

    Returns the entropy of entropy (`EoE`), the average Shannon entropy (`AvEn`), and
    the number of levels (`S2`) across all windows estimated from the data sequence (`Sig`) using
    the default parameters:
    window length (samples) = 10, slices = 10, logarithm = natural,
    heartbeat interval range (xmin, xmax) = (min(Sig), max(Sig))

        EoE, AvEn, S2 = EnofEn(Sig::AbstractArray{T,1} where T<:Real; tau::Int=10, S::Int=10, Xrange::Tuple{Real,REal}, Logx::Real=exp(1))

    Returns the entropy of entropy (`EoE`) estimated from the data sequence
    (`Sig`)  using the specified 'keyword' arguments:

    # Arguments:
    `tau`    - Window length, an integer > 1 \n
    `S`      - Number of slices (s1,s2), a two-element tuple of integers > 2 \n
    `Xrange` - The min and max heartbeat interval, a two-element tuple where X[1] <= X[2]\n
    `Logx`   - Logarithm base, a positive scalar  \n

    # See also `SampEn`, `MSEn`, `ApEn`

    # References:
        [1] Chang Francis Hsu, et al.,
            "Entropy of entropy: Measurement of dynamical complexity for
            biological systems."
            Entropy
            19.10 (2017): 550.

    """
    function EnofEn(Sig::AbstractArray{T,1} where T<:Real; tau::Int=10, S::Int=10,
                        Xrange::Tuple{Real,Real}=(minimum(Sig),maximum(Sig)), Logx::Real=exp(1))

    N = size(Sig,1)
    (N > 10) ? nothing :  error("Sig:   must be a numeric vector")
    (tau > 1 && tau < length(Sig)) ? nothing : error("tau:   must be an integer > 1")
    (S > 1) ? nothing : error("S:  must be an integer > 1")
    (length(Xrange)==2 && (Xrange[1]<=Xrange[2])) ? nothing : error("Xrange:   must be a two-element numeric tuple where Xrange[1]<Xrange[2]")
    (Logx>0) ? nothing :  error("Logx:     must be a positive number > 0")

    Wn = Int(floor(N/tau))
    Wj = transpose(reshape(Sig[1:Wn*tau],tau,Wn))
    Yj = zeros(Wn)
    #Edges = collect(range(minimum(Sig),maximum(Sig),length=(S[1]+1)))
    Edges = collect(range(Xrange[1],Xrange[2],length=(S+1)))
    Edges[1] -= .1;  Edges[end] += .1
    for n = 1:Wn
        Temp = fit(Histogram,Wj[n,:],Edges).weights/tau
        Temp = Temp[Temp.>0]
        Yj[n] = -sum(Temp.*log.(Logx, Temp))
    end

    AvEn = sum(Yj)/Wn
    #Edges = collect(range(minimum(Yj),maximum(Yj),length=(S[2]+1)))
    #Edges[1] -= .1;  Edges[end] += .1
    #Pjl = fit(Histogram,Yj,Edges).weights/Wn
    #Pjl = Pjl[Pjl.>0]
    Pjl = collect(values(countmap(round.(Yj,digits=12))))./Wn
    S2 = length(Pjl)
    if round(sum(Pjl),digits=5) != 1
        @warn("Possible error estimating probabilities")
    end

    EoE = -sum(Pjl.*log.(Logx, Pjl))
    return EoE, AvEn, S2

    end
end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


_πACGT(mod::NASM) = 0.25
_πACGU(mod) = _πACGT(mod)
_πA(mod) = _πACGT(mod)
_πC(mod) = _πACGT(mod)
_πG(mod) = _πACGT(mod)
_πT(mod) = _πACGT(mod)
_π(mod::NASM) = SVector(_πA(mod), _πC(mod), _πG(mod), _πT(mod))
_πR(mod::NASM) = _πA(mod) + _πG(mod)
_πY(mod::NASM) = _πT(mod) + _πC(mod)

"""
Generate a Q matrix for a `NucleicAcidSubstitutionModel`, of the form:

\$Q = \\begin{bmatrix}
Q_{A, A} & Q_{A, C} & Q_{A, G} & Q_{A, T} \\\\
Q_{C, A} & Q_{C, C} & Q_{C, G} & Q_{C, T} \\\\
Q_{G, A} & Q_{G, C} & Q_{G, G} & Q_{G, T} \\\\
Q_{T, A} & Q_{T, C} & Q_{T, G} & Q_{T, T} \\end{bmatrix}\$

Call as either
1) `Q(model)`, or
2) `Q(model, bool)`
Form (2) scales the matrix so that ``_π(model) ⋅ -diag(Q(model)) = 1``
when bool=true. `Q(model, false)` is equivalent to `Q(model)`.
"""
@inline function Q(mod::NASM)
    α = _α(mod)
    β = _β(mod)
    γ = _γ(mod)
    δ = _δ(mod)
    ϵ = _ϵ(mod)
    η = _η(mod)
    πA = _πA(mod)
    πC = _πC(mod)
    πG = _πG(mod)
    πT = _πT(mod)

    return Qmatrix(-(δ * πC + η * πG + β * πT), δ * πA, η * πA, β * πA,
                   δ * πC, -(δ * πA + ϵ * πG + α * πT), ϵ * πC, α * πC,
                   η * πG, ϵ * πG, -(η * πA + ϵ * πC + γ * πT), γ * πG,
                   β * πT, α * πT, γ * πT, -(β * πA + α * πC + γ * πG))
end

function Q(mod::NASM, scale::Bool)
  if scale
    q = Q(mod)
    scale = 1/(-diag(q)' * _π(mod))
    return q * scale
  else
    Q(mod)
  end
end

scale_generic(mod::NASM) = 1/(-diag(Q(mod))' * _π(mod))

_scale(mod::NASM) = scale_generic(mod)

@inline function P_generic(mod::NASM, t::Float64)
    if t < 0.0
        error("t must be positive")
    end
    return exp(Q(mod) * t)
end

function P_generic(mod::NASM, t, scale::Bool)
    if scale
        P(mod, t * _scale(mod))
    else
        P(mod, t)
    end
end

function P_generic(mod::NASM, t::Array{Float64})
    if any(t .< 0.0)
        error("t must be positive")
    end

    try
        q = Q(mod)

        # Use symmetrical similar matrix if possible:
        # B is similar (has same eigenvalues) to Q, but is symmetrical if the model is reversible.
        # Eigenvalue decomposition is easier with symmetrical matrix
        # NB: computation of B described in Inferring Phylogenies, Felsenstein, p.206

        rootπ = sqrt.(_π(mod))
        b = diagm(0 => rootπ) * q * diagm(0 => 1.0 ./ rootπ)

        # If B is symmetrical, then do eigenvalue decomposition on B. The resulting
        # eigenvalues are the same as for Q, then Q's eigenvectors can be obtained by
        # a simple conversion.
        if ishermitian(round.(b, digits=12))
            eig_vals, r = eigen(b)
            eig_vecs = diagm(0 => 1.0 ./ rootπ) * r # eig_vecs are the eigenvectors of Q
            eig_vecs_inv = r' * diagm(0 => rootπ)

        # If B is not symmetrical, fall back to directly obtaining eigenvalues from Q
        else
            eig_vals, eig_vecs = eigen(Array(q))
            eig_vecs_inv = inv(eig_vecs)
        end

        return [SMatrix{size(q)...}(eig_vecs * diagm(0 => exp.(eig_vals .* i)) * eig_vecs_inv) for i in t]

    catch
        # Any errors, fall back to direct use of matrix exponential
        return [P_generic(mod, i) for i in t]

    end
end

"""
Generate a P matrix for a `NucleicAcidSubstitutionModel`, of the form:

\$P = \\begin{bmatrix}
P_{A, A} & P_{A, C} & P_{A, G} & P_{A, T} \\\\
P_{C, A} & P_{C, C} & P_{C, G} & P_{C, T} \\\\
P_{G, A} & P_{G, C} & P_{G, G} & P_{G, T} \\\\
P_{T, A} & P_{T, C} & P_{T, G} & P_{T, T} \\end{bmatrix}\$

for specified time

Call as either
1) `P(model, t)`, or
2) `P(model, t, bool)`
Form (2) obtains its probabilities from the scaled Q matrix if bool=true.
Branch lengths estimated from a scaled P matrix are in units of expected
number of substitutions per site. `P(model, t, false)` is equivalent to
`P(model, t)`.
"""
P(mod::NASM, t) = P_generic(mod, t)

P(mod::NASM, t, scale) = P_generic(mod, t, scale)


using ZonalFlow
using Test

@testset "QL/CE2 tests" begin
    @info "Testing QL/CE2 conformity..."
    lx = 4.0*Float64(pi);
    ly = 2.0*Float64(pi);
    Ω = 2.0*Float64(pi)

    Λ = 0
    for nx=4:4
        for ny=nx:4
            for θ in [0.0]
                for τ in [20.0]
                    for Ξ in [0.2]

                        @info "Parameters: θ = $θ, τ = $τ, Ξ = $Ξ"

                        β = 2.0*Ω*cos(θ*Float64(pi))
                        ζ0 = ic_pert_eqm(lx,ly,nx,ny,Ξ*Ω); # one ic for all

                        sol2 = gql(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,t_end=500.0,dt=0.001,ic=ζ0);
                        sol3 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,t_end=500.0,ic=ζ0,dt=0.001);
                        sol4 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,t_end=500.0,ic=ζ0,dt=0.001,poscheck=true);
                        @test sol2.u[end][:,1:Λ+1] ≈ sol3.u[end].x[1] atol = 1e-6
                        @test sol2.u[end][:,1:Λ+1] ≈ sol4.u[end].x[1] atol = 1e-6
                        P2,O2 = zonalenergy(lx,ly,nx,ny,sol2.u);
                        P3,O3 = zonalenergy(lx,ly,nx,ny,Λ,sol3.u);
                        P4,O4 = zonalenergy(lx,ly,nx,ny,Λ,sol4.u);
                        @test P2[end,:] ≈ P3[end,:] atol = 1e-6
                        @test P2[end,:] ≈ P4[end,:] atol = 1e-6

                    end
                end
            end
        end
    end
end

@testset "NL/GQL(M)/GCE2(M) tests" begin
    @info "Testing NL/GQL(M)/GCE2(M) conformity..."
    lx = 4.0*Float64(pi);
    ly = 2.0*Float64(pi);
    Ω = 2.0*Float64(pi)

    for nx=4:4
        for ny=nx:4
            for θ in [0.0]
                for τ in [20.0]
                    for Ξ in [0.2]

                        @info "Parameters: θ = $θ, τ = $τ, Ξ = $Ξ"

                        β = 2.0*Ω*cos(θ*Float64(pi))
                        ζ0 = ic_pert_eqm(lx,ly,nx,ny,Ξ*Ω); # one ic for all

                        sol1 = nl(lx,ly,nx,ny,Ξ*Ω,β,τ/Ω,t_end=500.0,dt=0.001,ic=ζ0);
                        P1,O1 = zonalenergy(lx,ly,nx,ny,sol1.u);

                        Λ = nx-1
                        sol2 = gql(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,t_end=500.0,dt=0.001,ic=ζ0);
                        sol3 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,t_end=500.0,ic=ζ0,dt=0.001);
                        sol4 = gce2(lx,ly,nx,ny,Λ,Ξ*Ω,β,τ/Ω,t_end=500.0,ic=ζ0,dt=0.001,poscheck=true);
                        @test sol1.u[end] ≈ sol2.u[end] atol = 1e-6
                        @test sol1.u[end] ≈ sol3.u[end].x[1] atol = 1e-6
                        @test sol1.u[end] ≈ sol4.u[end].x[1] atol = 1e-6
                        P2,O2 = zonalenergy(lx,ly,nx,ny,sol2.u);
                        P3,O3 = zonalenergy(lx,ly,nx,ny,Λ,sol3.u);
                        P4,O4 = zonalenergy(lx,ly,nx,ny,Λ,sol4.u);
                        @test P1[end,:] ≈ P2[end,:] atol = 1e-6
                        @test P1[end,:] ≈ P3[end,:] atol = 1e-6
                        @test P1[end,:] ≈ P4[end,:] atol = 1e-6

                    end
                end
            end
        end
    end
end