using Parameters
using StaticArrays
using HCubature
using NLsolve
include("ForwardDiff.jl")
include("LinearAlgebra.jl/src/LinearAlgebra.jl")

""" Cubic Anderson model parameters (units of t_c)"""
@with_kw mutable struct PAM
    hyb0::Float64 = 0.  # Hybridization strength
    t_f::Float64 = 0.   # F-electron hopping
    ε⁰f::Float64 = 0.   # F-electron ground state e
    μ::Float64 = 0.     # Chemical potential
    β::Float64 = Inf    # Inverse temperature

    b::Real = NaN       # Saddle-point condensed slave boson
    λ::Real = NaN       # Saddle-point lagrange multiplier
end

function Hamiltonian(k::SVector, pam::PAM)
    # σ = [[0 1; 1 0], [0 -1im; 1im 0], [1 0; 0 -1]]
    σ = [[0 1; 1 0]]
    c = kron([1 0; 0 0], eye(2)) * (-2.0 * sum(cos, k) - pam.μ)
    f = kron([0 0; 0 1], eye(2)) * (-2.0 * pam.t_f * sum(cos, k) * pam.b^2 + pam.ε⁰f + pam.λ - pam.μ)
    cf = kron([0 1; 0 0], sin.(k)' * σ) * (pam.hyb0 * pam.b)
    return Hermitian(c + f + cf, :U)
end

function FreeEnergy(x::AbstractArray, pam::PAM)
    const integration(f) = hcubature(f, [-π],[π]; atol=1e-3, rtol=1e-3, maxevals=100)[1] / (2π)^1
    const lower_bands(k::SVector) = first(LinearAlgebra.EigenSelfAdjoint.eig!(Hamiltonian(k,pam)))[1:2]

    pam.b, pam.λ = x

    return pam.λ * (pam.b^2 - 1) + integration(q -> sum(lower_bands(q)))
end

PAM(hyb0=1.0, t_f=-0.1, ε⁰f=-2.5)
nlsolve(x -> ForwardDiff.gradient(x0 -> FreeEnergy(x0, pam), x), [0.1, -pam.ε⁰f]; inplace=false)