using Compat: @compat
using ITensors: ITensor, array, inds, itensor, @disable_warn_order
using ITensorMPS: TimeDependentSum, to_vec, OpSum, MPO, MPS, contract, inner, random_mps, siteinds, tdvp
using LinearAlgebra: norm
using KrylovKit: exponentiate
using OrdinaryDiffEqTsit5
using OrdinaryDiffEqTsit5: ODEProblem, Tsit5, solve 
using Random: Random

function heisenberg(n; J=1.0, J2=0.0)
  ℋ = OpSum()
  if !iszero(J)
    for j in 1:(n - 1)
      ℋ += J / 2, "S+", j, "S-", j + 1
      ℋ += J / 2, "S-", j, "S+", j + 1
      ℋ += J, "Sz", j, "Sz", j + 1
    end
  end
  if !iszero(J2)
    for j in 1:(n - 2)
      ℋ += J2 / 2, "S+", j, "S-", j + 2
      ℋ += J2 / 2, "S-", j, "S+", j + 2
      ℋ += J2, "Sz", j, "Sz", j + 2
    end
  end
  return ℋ
end

function ode_updater(operator, init; internal_kwargs, alg=Tsit5(), kwargs...)
  @compat (; current_time, time_step) = (; current_time=zero(Bool), internal_kwargs...)
  time_span = typeof(time_step).((current_time, current_time + time_step))
  init_vec, to_itensor = to_vec(init)
  f(init::ITensor, p, t) = operator(t)(init)
  f(init_vec::AbstractArray, p, t) = to_vec(f(to_itensor(init_vec), p, t))[1]
  prob = ODEProblem(f, init_vec, time_span)
  sol = solve(prob, alg; kwargs...)
  state_vec = sol.u[end]
  return to_itensor(state_vec), (;)
end

function krylov_updater(operator, init; internal_kwargs, kwargs...)
  @compat (; current_time, time_step) = (; current_time=zero(Bool), internal_kwargs...)
  state, info = exponentiate(operator(current_time), time_step, init; kwargs...)
  return state, (; info)
end

function main(; eltype=Float64, device=identity)
  Random.seed!(1234)

  # Time dependent Hamiltonian is:
  # H(t) = H₁(t) + H₂(t) + …
  #      = f₁(t) H₁(0) + f₂(t) H₂(0) + …
  #      = cos(ω₁t) H₁(0) + cos(ω₂t) H₂(0) + …

  # Number of sites
  n = 6

  # How much information to output from TDVP
  # Set to 2 to get information about each bond/site
  # evolution, and 3 to get information about the
  # updater.
  outputlevel = 3

  # Frequency of time dependent terms
  ω₁ = one(eltype) / 10
  ω₂ = one(eltype) / 5

  # Nearest and next-nearest neighbor
  # Heisenberg couplings.
  J₁ = one(eltype)
  J₂ = one(eltype)

  time_step = one(eltype) / 10
  time_stop = one(eltype)

  # nsite-update TDVP
  nsite = 2

  # Starting state bond/link dimension.
  # A product state starting state can
  # cause issues for TDVP without
  # subspace expansion.
  start_linkdim = 4

  # TDVP truncation parameters
  maxdim = 100
  cutoff = √(eps(eltype))

  tol = 10 * eps(eltype)

  @show n
  @show ω₁, ω₂
  @show J₁, J₂
  @show maxdim, cutoff, nsite
  @show start_linkdim
  @show time_step, time_stop

  ω⃗ = (ω₁, ω₂)
  f⃗ = map(ω -> (t -> cos(ω * t)), ω⃗)

  # H₀ = H(0) = H₁(0) + H₂(0) + …
  ℋ₁₀ = heisenberg(n; J=J₁, J2=zero(eltype))
  ℋ₂₀ = heisenberg(n; J=zero(eltype), J2=J₂)
  ℋ⃗₀ = (ℋ₁₀, ℋ₂₀)

  s = siteinds("S=1/2", n)

  H⃗₀ = map(ℋ₀ -> device(MPO(eltype, ℋ₀, s)), ℋ⃗₀)

  # Initial state, ψ₀ = ψ(0)
  # Initialize as complex since that is what OrdinaryDiffEq.jl/DifferentialEquations.jl
  # expects.
  ψ₀ = device(
    complex.(random_mps(eltype, s, j -> isodd(j) ? "↑" : "↓"; linkdims=start_linkdim))
  )

  @show norm(ψ₀)

  println()
  println("#"^100)
  println("Running TDVP with ODE updater")
  println("#"^100)
  println()

  ψₜ_ode = tdvp(
    -im * TimeDependentSum(f⃗, H⃗₀),
    time_stop,
    ψ₀;
    updater=ode_updater,
    updater_kwargs=(; reltol=tol, abstol=tol),
    time_step,
    maxdim,
    cutoff,
    nsite,
    outputlevel,
    noise=1e-8,
  )

  println()
  println("Finished running TDVP with ODE updater")
  println()

  println()
  println("#"^100)
  println("Running TDVP with Krylov updater")
  println("#"^100)
  println()

  ψₜ_krylov = tdvp(
    -im * TimeDependentSum(f⃗, H⃗₀),
    time_stop,
    ψ₀;
    updater=krylov_updater,
    updater_kwargs=(; tol, eager=true),
    time_step,
    cutoff,
    nsite,
    outputlevel,
  )

  println()
  println("Finished running TDVP with Krylov updater")
  println()

  println()
  println("#"^100)
  println("Running full state evolution with ODE updater")
  println("#"^100)
  println()

  @disable_warn_order begin
    ψₜ_full, _ = ode_updater(
      -im * TimeDependentSum(f⃗, contract.(H⃗₀)),
      contract(ψ₀);
      internal_kwargs=(; time_step=time_stop, outputlevel),
      reltol=tol,
      abstol=tol,
    )
  end

  println()
  println("Finished full state evolution with ODE updater")
  println()

  @show norm(ψₜ_ode)
  @show norm(ψₜ_krylov)
  @show norm(ψₜ_full)

  @show 1 - abs(inner(contract(ψₜ_ode), ψₜ_full))
  @show 1 - abs(inner(contract(ψₜ_krylov), ψₜ_full))
  return nothing
end

main()

ERROR: UndefVarError: `phi` not defined in `ITensorMPS`
Suggestion: check for spelling errors or missing imports.
Stacktrace:
  [1] region_update!(nsite_val::Val{…}, reverse_step_val::Val{…}, reduced_operator::TimeDependentSum{…}, state::MPS, b::Int64; updater::typeof(ode_updater), updater_kwargs::@NamedTuple{…}, current_time::Float64, time_step::Float64, outputlevel::Int64, normalize::Bool, direction::Base.Order.ForwardOrdering, noise::Float64, which_decomp::Nothing, svd_alg::Nothing, cutoff::Float64, maxdim::Int64, mindim::Int64, maxtruncerr::Float64)
    @ ITensorMPS ~/ITensorMPS.jl/src/solvers/sweep_update.jl:417
  [2] region_update!(reduced_operator::TimeDependentSum{…}, state::MPS, b::Int64; updater::Function, updater_kwargs::@NamedTuple{…}, nsite::Int64, reverse_step::Bool, current_time::Float64, outputlevel::Int64, time_step::Float64, normalize::Bool, direction::Base.Order.ForwardOrdering, noise::Float64, which_decomp::Nothing, svd_alg::Nothing, cutoff::Float64, maxdim::Int64, mindim::Int64, maxtruncerr::Float64)
    @ ITensorMPS ~/ITensorMPS.jl/src/solvers/sweep_update.jl:187
  [3] sub_sweep_update(direction::Base.Order.ForwardOrdering, reduced_operator::TimeDependentSum{…}, state::MPS; updater::Function, updater_kwargs::@NamedTuple{…}, which_decomp::Nothing, svd_alg::Nothing, sweep::Int64, current_time::Float64, time_step::Float64, nsite::Int64, reverse_step::Bool, normalize::Bool, observer!::ITensorMPS.EmptyObserver, outputlevel::Int64, maxdim::Int64, mindim::Int64, cutoff::Float64, noise::Float64)
    @ ITensorMPS ~/ITensorMPS.jl/src/solvers/sweep_update.jl:107
  [4] sweep_update(order::ITensorMPS.TDVPOrder{…}, reduced_operator::TimeDependentSum{…}, state::MPS; current_time::Float64, time_step::Float64, kwargs::@Kwargs{…})
    @ ITensorMPS ~/ITensorMPS.jl/src/solvers/sweep_update.jl:25
  [5] macro expansion
    @ ~/ITensorMPS.jl/src/solvers/alternating_update.jl:69 [inlined]
  [6] macro expansion
    @ ./timing.jl:421 [inlined]
  [7] alternating_update(operator::TimeDependentSum{…}, init::MPS; updater::Function, updater_kwargs::@NamedTuple{…}, nsweeps::Int64, checkdone::Returns{…}, write_when_maxdim_exceeds::Nothing, nsite::Int64, reverse_step::Bool, time_start::Float64, time_step::Float64, order::Int64, observer!::ITensorMPS.EmptyObserver, sweep_observer!::ITensorMPS.EmptyObserver, outputlevel::Int64, normalize::Bool, maxdim::Int64, mindim::Int64, cutoff::Float64, noise::Float64)
    @ ITensorMPS ~/ITensorMPS.jl/src/solvers/alternating_update.jl:68
  [8] alternating_update
    @ ~/ITensorMPS.jl/src/solvers/alternating_update.jl:23 [inlined]
  [9] #tdvp#767
    @ ~/ITensorMPS.jl/src/solvers/tdvp.jl:86 [inlined]
 [10] main(; eltype::Type, device::typeof(identity))
    @ Main ~/ITensorMPS.jl/examples/solvers/03_tdvp_time_dependent.jl:105
 [11] main()
    @ Main ~/ITensorMPS.jl/examples/solvers/03_tdvp_time_dependent.jl:27
 [12] top-level scope
    @ ~/ITensorMPS.jl/examples/solvers/03_tdvp_time_dependent.jl:174
Some type information was truncated. Use `show(err)` to see complete types.