# KeldyshED.jl
#
# Copyright (C) 2019-2024 Igor Krivenko <igor.s.krivenko@gmail.com>
# Copyright (C) 2015 P. Seth, I. Krivenko, M. Ferrero and O. Parcollet
#
# KeldyshED.jl is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software
# Foundation, either version 3 of the License, or (at your option) any later
# version.
#
# KeldyshED.jl is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# KeldyshED.jl. If not, see <http://www.gnu.org/licenses/>.
#
# Authors: Hugo U.R. Strand, Igor Krivenko

using DocStringExtensions

using Keldysh

using .Hilbert: factorized_basis_map

"""
$(TYPEDSIGNATURES)

Compute the partition function ``Z = \\mathrm{Tr}[e^{-\\beta\\hat H}]``
at an inverse temperature ``\\beta``.
"""
function partition_function(ed::EDCore, β::Real)::Real
  return sum(sum(exp.(-β * es.eigenvalues)) for es in ed.eigensystems)
end

"""
$(TYPEDSIGNATURES)

Compute the equilibrium density matrix ``\\hat\\rho = e^{-\\beta\\hat H} / Z`` at an
inverse temperature ``\\beta`` represented in the eigenbasis of the Hamiltonian ``\\hat H``.
The density matrix is returned as a vector of diagonal blocks corresponding to invariant
subspaces of ``\\hat H``.
"""
function density_matrix(ed::EDCore, β::Real)
  z = partition_function(ed, β)
  return [diagm(exp.(-β * es.eigenvalues) / z) for es in ed.eigensystems]
end

const EvolutionOperator = Vector{GenericTimeGF{ComplexF64, false}}

"""
$(TYPEDSIGNATURES)

Compute the evolution operator

```math
    \\hat S(t, t') = \\exp\\left(-i \\int_{t'}^{t} \\hat H d\\bar t\\right)
```
on a given contour time grid and represented in the eigenbasis of the Hamiltonian
``\\hat H``.
The operator is returned as a vector of diagonal blocks (matrix-valued Green's function
containers) corresponding to invariant subspaces of ``\\hat H``.
"""
function evolution_operator(ed::EDCore,
                            grid::AbstractTimeGrid)::EvolutionOperator
  [GenericTimeGF(grid, length(es.eigenvalues)) do t1, t2
      diagm(exp.(-im * es.eigenvalues * (t1.bpoint.val - t2.bpoint.val)))
   end
   for es in ed.eigensystems]
end

#########################
# Basis transformations #
#########################

"""
$(TYPEDSIGNATURES)

Transform a block-diagonal evolution operator `S` written in the eigenbasis of the system
into the Fock state basis.
"""
function tofockbasis(S::Vector{GF}, ed::EDCore) where {GF <: AbstractTimeGF}
  S_fock = GF[]
  for (s, es) in zip(S, ed.eigensystems)
    U = es.unitary_matrix
    push!(S_fock, similar(s))
    for t1 in S_fock[end].grid, t2 in S_fock[end].grid
      S_fock[end][t1, t2] = U * s[t1, t2] * adjoint(U)
    end
  end
  return S_fock
end

"""
$(TYPEDSIGNATURES)

Transform a block-diagonal evolution operator `S` written in the Fock state basis into
the eigenbasis of the system.
"""
function toeigenbasis(S::Vector{GF}, ed::EDCore) where {GF <: AbstractTimeGF}
  S_eigen = GF[]
  for (s, es) in zip(S, ed.eigensystems)
    U = es.unitary_matrix
    push!(S_eigen, similar(s))
    for t1 in S_eigen[end].grid, t2 in S_eigen[end].grid
      S_eigen[end][t1, t2] = adjoint(U) * s[t1, t2] * U
    end
  end
  return S_eigen
end

#################
# Partial trace #
#################

"""
$(TYPEDSIGNATURES)

Compute a partial trace of a block-diagonal matrix `M`.

The matrix is expected to be written in the Fock state basis. The resulting reduced matrix
acts in a [`Hilbert space`](@ref KeldyshED.Hilbert.FullHilbertSpace) spanned by all
fermionic Fock states generated by the
[`set of indices`](@ref KeldyshED.Hilbert.SetOfIndices) `target_soi`.
"""
function partial_trace(M::Vector{Matrix{T}},
                       ed::EDCore,
                       target_soi::SetOfIndices) where T

  target_hs = FullHilbertSpace(target_soi)
  fmap = factorized_basis_map(ed.full_hs, target_hs)

  n = length(target_hs)
  M_target = zeros(T, n, n)

  for (Mss, hss) in zip(M, ed.subspaces)
    for (i1, fs1) in pairs(hss), (i2, fs2) in pairs(hss)
      i_target1, i_env1 = fmap[getstateindex(ed.full_hs, fs1)]
      i_target2, i_env2 = fmap[getstateindex(ed.full_hs, fs2)]

      i_env1 != i_env2 && continue
      M_target[i_target1, i_target2] += Mss[i1, i2]
    end
  end

  M_target
end

#
# Construct a similar GF object with a different number of orbitals
#

function Base.similar(G::T, norb::Int) where T <: GenericTimeGF
  GenericTimeGF(eltype(T), G.grid, norb, is_scalar(T))
end
function Base.similar(G::T, norb::Int) where T <: FullTimeGF
  FullTimeGF(eltype(T), G.grid, norb, G.ξ, is_scalar(T))
end
function Base.similar(G::T, norb::Int) where T <: TimeInvariantFullTimeGF
  TimeInvariantFullTimeGF(eltype(T), G.grid, norb, G.ξ, is_scalar(T))
end
function Base.similar(G::T, norb::Int) where T <: KeldyshTimeGF
  KeldyshTimeGF(eltype(T), G.grid, norb, G.ξ, is_scalar(T))
end
function Base.similar(G::T, norb::Int) where T <: TimeInvariantKeldyshTimeGF
  TimeInvariantKeldyshTimeGF(eltype(T), G.grid, norb, G.ξ, is_scalar(T))
end
function Base.similar(G::T, norb::Int) where T <: ImaginaryTimeGF
  ImaginaryTimeGF(eltype(T), G.grid, norb, G.ξ, is_scalar(T))
end

"""
$(TYPEDSIGNATURES)

Compute a partial trace of a block-diagonal evolution operator `S` written in
the Fock state basis.

The resulting reduced evolution operator acts in a
[`Hilbert space`](@ref KeldyshED.Hilbert.FullHilbertSpace) spanned by
all fermionic Fock states generated by the
[`set of indices`](@ref KeldyshED.Hilbert.SetOfIndices) `target_soi`.
"""
function partial_trace(S::Vector{GF},
                       ed::EDCore,
                       target_soi::SetOfIndices) where {GF <: AbstractTimeGF}
  target_hs = FullHilbertSpace(target_soi)
  fmap = factorized_basis_map(ed.full_hs, target_hs)

  S_target = similar(S[1], length(target_hs))

  for (Sss, hss) in zip(S, ed.subspaces)
    for (i1, fs1) in pairs(hss), (i2, fs2) in pairs(hss)
      i_target1, i_env1 = fmap[getstateindex(ed.full_hs, fs1)]
      i_target2, i_env2 = fmap[getstateindex(ed.full_hs, fs2)]

      i_env1 != i_env2 && continue

      for t1 in S_target.grid, t2 in S_target.grid
        S_target[i_target1, i_target2, t1, t2] += Sss[i1, i2, t1, t2]
      end
    end
  end

  S_target
end

"""
$(TYPEDSIGNATURES)

Compute a reduced equilibrium density matrix

```math
\\hat\\rho_{target} = \\mathrm{Tr}_{env}[e^{-\\beta\\hat H}] / Z
```
at an inverse temperature ``\\beta``.

The Hamiltonian ``\\hat H`` acts in a direct product of
[`Hilbert spaces`](@ref KeldyshED.Hilbert.FullHilbertSpace)
``H_{target} \\otimes H_{env}``, while the resulting reduced matrix acts
in ``H_{target}`` spanned by all fermionic Fock states generated by the
[`set of indices`](@ref KeldyshED.Hilbert.SetOfIndices) `target_soi`.
"""
function reduced_density_matrix(ed::EDCore,
                                target_soi::SetOfIndices,
                                β::Real)
  ρ = tofockbasis(density_matrix(ed, β), ed)
  return partial_trace(ρ, ed, target_soi)
end

"""
$(TYPEDSIGNATURES)

Compute a reduced evolution operator

```math
\\hat S_{target}(t, t') =
\\mathrm{Tr}_{env}\\left[\\exp\\left(-i \\int_{t'}^{t} \\hat H d\\bar t\\right)\\right]
```

on a given contour time grid.

The Hamiltonian ``\\hat H`` acts in a direct product of
[`Hilbert spaces`](@ref KeldyshED.Hilbert.FullHilbertSpace)
``H_{target} \\otimes H_{env}``, while the resulting reduced evolution operator
acts in ``H_{target}`` spanned by all fermionic Fock states generated by
the [`set of indices`](@ref KeldyshED.Hilbert.SetOfIndices) `target_soi`.
"""
function reduced_evolution_operator(ed::EDCore,
                                    target_soi::SetOfIndices,
                                    grid::AbstractTimeGrid)
  S = tofockbasis(evolution_operator(ed, grid), ed)
  return partial_trace(S, ed, target_soi)
end
