vanVleck Recursion

[oameye]

https://github.com/oameye/vanVleck-recursion

[oameye]

using SymPy
using Base.deepcopy

struct Term
    rotating::Int
    factor::Sym
    freq_denom::Int
    term1::Union{Nothing, Term}
    term2::Union{Nothing, Term}
    footprint::String
    term_count::Int
end

function Term(rotating::Int=1)
    return Term(rotating, Sym(1), 0, nothing, nothing, string(rotating), 1)
end

struct Terms
    terms::Vector{Term}
end

function Terms(terms::Vector{Term})
    return Terms(terms)
end

function bracket(self::Term, term::Term, factor::Sym=Sym(1))
    term_r = Term(1)
    term_r.term1 = deepcopy(self)
    term_r.term2 = deepcopy(term)
    term_r.factor = factor * self.factor * term.factor
    term_r.footprint = "[" * self.footprint * "," * term.footprint * "]"
    term_r.term_count = self.term_count + term.term_count

    term_s = deepcopy(term_r)
    term_s.rotating = 0
    term_s.footprint *= "0"
    if self.rotating != term.rotating
        return Terms([term_r])
    end
    if self.rotating == 0
        return Terms([term_s])
    end
    return Terms([term_s, term_r])
end

function dot(self::Term, factor::Sym=Sym(1))
    if self.rotating == 0
        return Terms([])
    end
    new_term = deepcopy(self)
    new_term.factor *= factor
    new_term.freq_denom -= 1

    if new_term.freq_denom == 0
        new_term.footprint = new_term.footprint[1:end-2]
    else
        @warn "Unexpected derivative over an unintegrated term"
        new_term.footprint *= "/" * string(new_term.freq_denom)
    end
    return Terms([new_term])
end

function simplify(terms::Terms)
    # Placeholder for the simplify method
end

mutable struct Kamiltonian
    Ks::Dict{String, Terms}
end

function Kamiltonian()
    return Kamiltonian(Dict{String, Terms}())
end

function get_K(n::Int, k::Int, cls::Kamiltonian)
    key = string([n, k])
    if haskey(cls.Ks, key)
        return cls.Ks[key]
    end

    if n == 0 && k == 0
        throw(Exception("Base case of recursion not specified"))
    elseif n == 0 && k == 1
        cls.Ks[key] = dot(Term(1), Sym(1))
    elseif n != 0 && k == 1
        term1 = dot(Term(n + 1), Sym(1))
        term2 = bracket(Term(n), Term(0))
        cls.Ks[key] = Terms(term1.terms + term2.terms)
    elseif k > 1 && k <= n + 1
        terms = Terms([])
        for m in 0:(n-1)
            terms = Terms(terms.terms + bracket(Term(n - m), get_K(m, k - 1, cls), 1 / Sym(k)).terms)
        end
        cls.Ks[key] = terms
    else
        return Terms([])
    end
    simplify(cls.Ks[key])
    return cls.Ks[key]
end

function set_H(cls::Kamiltonian, H::Terms)
    cls.Ks[string([0, 0])] = H
end

struct Generator <: Terms
    Ss::Dict{Int, Terms}
end

function Generator()
    return Generator(Dict{Int, Terms}())
end

function get_S(cls::Generator, n::Int)
    if haskey(cls.Ss, n)
        return cls.Ss[n]
    end

    terms = Terms([])
    if n == 0
        return Terms([])
    elseif n == 1
        terms = Terms(terms.terms + bracket(Term(0), Term(0), Sym(-1)).terms)
    elseif n > 1
        terms = Terms(terms.terms + bracket(get_S(cls, n - 1), Term(0), Sym(-1)).terms)
        for k in 2:n
            terms = Terms(terms.terms + bracket(get_K(n - 1, k, cls), Term(0), Sym(-1)).terms)
        end
    end
    terms = integrate(terms)
    simplify(terms)
    cls.Ss[n] = terms
    return cls.Ss[n]
end

function K(n::Int, k::Int=-1, cls::Kamiltonian=Kamiltonian())
    if k != -1
        return get_K(n, k, cls)
    end
    Kn = Terms([])
    for ki in 0:(n + 1)
        Kn = Terms(Kn.terms + get_K(n, ki, cls).terms)
    end
    simplify(Kn)
    return Kn
end

function S(n::Int, cls::Generator=Generator())
    return get_S(cls, n)
end

Kamiltonian_inst = Kamiltonian()
set_H(Kamiltonian_inst, Terms([Term(0), Term(1)]))

k = S(3, Generator()).terms[1]
println(k)