using Decapodes
using DiagrammaticEquations
using CombinatorialSpaces
using GeometryBasics
using MLStyle
using ComponentArrays
using OrdinaryDiffEq
using CairoMakie
import CairoMakie: wireframe, mesh, Figure, Axis
import Decapodes:nparts, mul!
using Zygote
using SciMLSensitivity
using DataFrames
using Chairmarks

Point2D = Point2{Float64}
Point3D = Point3{Float64}

rect = triangulated_grid(100, 100, 2, 2, Point3D)
d_rect = EmbeddedDeltaDualComplex2D{Bool,Float64,Point3D}(rect)
subdivide_duals!(d_rect, Circumcenter())

#fig = Figure()
#ax = CairoMakie.Axis(fig[1, 1])
#wf = wireframe!(ax, rect)
#fig

Heat = @decapode begin
    U::Form0

    k::Constant
    c::Constant
    ρ::Constant

    ∂ₜ(U) == k/(c*ρ) * Δ(U)
end


decapode_code = gensim(Heat, preallocate=true)
file = open("/home/jadonclugston/Documents/Work/dev/DecapodeCalibrateDemos/HeatEquation/Heat_alloc.jl", "w")
write(file, string("decapode_f = ", decapode_code))
close(file)
include("/home/jadonclugston/Documents/Work/dev/DecapodeCalibrateDemos/HeatEquation/Heat_alloc.jl")

#fₘ  = evalsim(Heat, preallocate=true)

fₘ = decapode_f(d_rect, nothing)

U = map(d_rect[:point]) do (x, _)
    return x
end



#fig = Figure()
#ax = CairoMakie.Axis(fig[1, 1])
#msh = mesh!(ax, rect, color=U, colormap=:jet)
#fig


# Aluminum Pure https://ncfs.ucf.edu/burn_db/Thermal_Properties/material_thermal.html
# k = 237, ρ = 2702, Cp = 0.903, ϵ = 0.03
u₀ = ComponentArray(U=U)

constants_and_parameters = ComponentArray(k = 100, c = 1, ρ = 1)

tₑ = 20.0

@info("Precompiling Solver")
prob = ODEProblem(fₘ, u₀, (0, 1e-4), constants_and_parameters)
soln = solve(prob, Tsit5())
soln.retcode != :Unstable || error("Solver was not stable")
@info("Solving")
data_prob = ODEProblem(fₘ, u₀, (0, tₑ), constants_and_parameters)
soln = solve(data_prob, Tsit5())
@info("Done")


#fig = Figure()
#ax = CairoMakie.Axis(fig[1, 1])
#msh = mesh!(ax, rect, color=soln(tₑ).U, colormap=:jet)
#fig

reference_dat = last(soln).U

function loss(u) #only compares last time step
    newp = ComponentArray(k=u[1], c=1, ρ=1)
    prob = remake(data_prob, p=newp)
    sol = solve(prob, Tsit5(), sensealg = GaussAdjoint(autojacvec = EnzymeVJP()))
    current_dat = last(sol).U
    sum(abs2, reference_dat .- current_dat)
end

Zygote.gradient(loss, 99.0)

ERROR: MethodError: no method matching augmented_primal(::EnzymeCore.EnzymeRules.RevConfigWidth{…}, ::EnzymeCore.Const{…}, ::Type{…}, ::EnzymeCore.Duplicated{…}, ::EnzymeCore.Duplicated{…}, ::EnzymeCore.Duplicated{…}, ::EnzymeCore.Const{…}, ::EnzymeCore.Const{…})

Closest candidates are:
  augmented_primal(::EnzymeCore.EnzymeRules.RevConfig, ::EnzymeCore.Const{typeof(mul!)}, ::Type{RT}, ::EnzymeCore.Annotation{<:StridedVecOrMat}, ::EnzymeCore.Const{<:Union{SparseArrays.AbstractSparseMatrixCSC{Tv, Ti}, SubArray{Tv, 2, <:SparseArrays.AbstractSparseMatrixCSC{Tv, Ti}, Tuple{Base.Slice{Base.OneTo{Int64}}, I}} where I<:AbstractUnitRange} where {Tv, Ti}}, ::EnzymeCore.Annotation{<:StridedVecOrMat}, ::EnzymeCore.Annotation{<:Number}, ::EnzymeCore.Annotation{<:Number}) where RT
   @ Enzyme ~/.julia/packages/Enzyme/Vjlrr/src/internal_rules.jl:732
  augmented_primal(::Any, ::EnzymeCore.Const{typeof(QuadGK.quadgk)}, ::Type{RT}, ::Any, ::EnzymeCore.Annotation{T}...; kws...) where {RT, T}
   @ QuadGKEnzymeExt ~/.julia/packages/QuadGK/BjmU0/ext/QuadGKEnzymeExt.jl:6
  augmented_primal(::Any, ::EnzymeCore.Const{typeof(NNlib._dropout!)}, ::Type{RT}, ::Any, ::OutType, ::Any, ::Any, ::Any) where {OutType, RT}
   @ NNlibEnzymeCoreExt ~/.julia/packages/NNlib/CkJqS/ext/NNlibEnzymeCoreExt/NNlibEnzymeCoreExt.jl:318
  ...

Stacktrace:
  [1] custom_rule_method_error
    @ ~/.julia/packages/Enzyme/Vjlrr/src/rules/customrules.jl:452 [inlined]
  [2] mul!
    @ ~/.julia/juliaup/julia-1.10.4 0.x64.linux.gnu/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:237 [inlined]
  [3] f
    @ ~/Documents/Work/dev/DecapodeCalibrateDemos/HeatEquation/Heat_alloc.jl:45
  [4] ODEFunction
    @ ~/.julia/packages/SciMLBase/tEuIM/src/scimlfunctions.jl:2336 [inlined]
  [5] #139
    @ ~/.julia/packages/SciMLSensitivity/HRhwU/src/adjoint_common.jl:479 [inlined]
  [6] diffejulia__139_22024_inner_1wrap
    @ ~/.julia/packages/SciMLSensitivity/HRhwU/src/adjoint_common.jl:0
  [7] macro expansion
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:8839 [inlined]
  [8] enzyme_call
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:8405 [inlined]
  [9] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/Vjlrr/src/compiler.jl:8178 [inlined]
 [10] autodiff
    @ ~/.julia/packages/Enzyme/Vjlrr/src/Enzyme.jl:491 [inlined]
 [11] _vecjacobian!(dλ::Vector{…}, y::ComponentVector{…}, λ::Vector{…}, p::ComponentVector{…}, t::Float64, S::SciMLSensitivity.ODEGaussAdjointSensitivityFunction{…}, isautojacvec::EnzymeVJP, dgrad::Nothing, dy::Nothing, W::Nothing)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/HRhwU/src/derivative_wrappers.jl:710
 [12] #vecjacobian!#18
    @ ~/.julia/packages/SciMLSensitivity/HRhwU/src/derivative_wrappers.jl:232 [inlined]
 [13] vecjacobian!
    @ ~/.julia/packages/SciMLSensitivity/HRhwU/src/derivative_wrappers.jl:229 [inlined]
 [14] (::SciMLSensitivity.ODEGaussAdjointSensitivityFunction{…})(du::Vector{…}, u::Vector{…}, p::ComponentVector{…}, t::Float64)
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/HRhwU/src/gauss_adjoint.jl:102
 [15] ODEFunction
    @ ~/.julia/packages/SciMLBase/tEuIM/src/scimlfunctions.jl:2336 [inlined]
 [16] initialize!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, cache::OrdinaryDiffEqTsit5.Tsit5Cache{…})
    @ OrdinaryDiffEqTsit5 ~/.julia/packages/OrdinaryDiffEqTsit5/DHYtz/src/tsit_perform_step.jl:175
 [17] __init(prob::ODEProblem{…}, alg::Tsit5{…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{…}; saveat::Vector{…}, tstops::Vector{…}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Bool, callback::CallbackSet{…}, dense::Bool, calck::Bool, dt::Float64, dtmin::Float64, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{…}, abstol::Float64, reltol::Float64, qmin::Rational{…}, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Rational{…}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(LinearAlgebra.opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), progress_id::Symbol, userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEqCore.DefaultInit, kwargs::@Kwargs{})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/HwWWN/src/solve.jl:525
 [18] __init (repeats 5 times)
    @ ~/.julia/packages/OrdinaryDiffEqCore/HwWWN/src/solve.jl:11 [inlined]
 [19] #__solve#75
    @ ~/.julia/packages/OrdinaryDiffEqCore/HwWWN/src/solve.jl:6 [inlined]
 [20] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/HwWWN/src/solve.jl:1 [inlined]
 [21] solve_call(_prob::ODEProblem{…}, args::Tsit5{…}; merge_callbacks::Bool, kwargshandle::Nothing, kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:612
 [22] solve_call
    @ ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:569 [inlined]
 [23] #solve_up#53
    @ ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:1092 [inlined]
 [24] solve_up
    @ ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:1078 [inlined]
 [25] #solve#51
    @ ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:1015 [inlined]
 [26] _adjoint_sensitivities(sol::ODESolution{…}, sensealg::GaussAdjoint{…}, alg::Tsit5{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, checkpoints::Vector{…}, corfunc_analytical::Bool, callback::Nothing, kwargs::@Kwargs{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/HRhwU/src/gauss_adjoint.jl:580
 [27] _adjoint_sensitivities
    @ ~/.julia/packages/SciMLSensitivity/HRhwU/src/gauss_adjoint.jl:533 [inlined]
 [28] #adjoint_sensitivities#63
    @ ~/.julia/packages/SciMLSensitivity/HRhwU/src/sensitivity_interface.jl:401 [inlined]
 [29] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#313"{…})(Δ::ODESolution{…})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/HRhwU/src/concrete_solve.jl:627
 [30] ZBack
    @ ~/.julia/packages/Zygote/Tt5Gx/src/compiler/chainrules.jl:212 [inlined]
 [31] (::Zygote.var"#294#295"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/lib/lib.jl:206
 [32] (::Zygote.var"#2169#back#296"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [33] #solve#51
    @ ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:1015 [inlined]
 [34] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [35] #294
    @ ~/.julia/packages/Zygote/Tt5Gx/src/lib/lib.jl:206 [inlined]
 [36] (::Zygote.var"#2169#back#296"{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/ZygoteRules/M4xmc/src/adjoint.jl:72
 [37] solve
    @ ~/.julia/packages/DiffEqBase/uqSeD/src/solve.jl:1005 [inlined]
 [38] (::Zygote.Pullback{…})(Δ::ODESolution{…})
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [39] loss
    @ ~/Documents/Work/dev/DecapodeCalibrateDemos/HeatEquation/HeatEquation.jl:89 [inlined]
 [40] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface2.jl:0
 [41] (::Zygote.var"#78#79"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl:91
 [42] gradient(f::Function, args::Float64)
    @ Zygote ~/.julia/packages/Zygote/Tt5Gx/src/compiler/interface.jl:148
 [43] top-level scope
    @ ~/Documents/Work/dev/DecapodeCalibrateDemos/HeatEquation/HeatEquation.jl:94
Some type information was truncated. Use `show(err)` to see complete types.