############################################################################
### Types
############################################################################

struct RAMSymbolic{F1, F2, F3, A1, A2, A3, S1, S2, S3, V, V2, F4, A4, F5, A5, D1} <: SemImplySymbolic
    Σ_function::F1
    ∇Σ_function::F2
    ∇²Σ_function::F3
    Σ::A1
    ∇Σ::A2
    ∇²Σ::A3
    Σ_symbolic::S1
    ∇Σ_symbolic::S2
    ∇²Σ_symbolic::S3
    n_par::V
    ram_matrices::V2
    μ_function::F4
    μ::A4
    ∇μ_function::F5
    ∇μ::A5
    identifier::D1
end

############################################################################
### Constructors
############################################################################

function RAMSymbolic(;
        specification,
        loss_types = nothing,
        vech = false,
        gradient = true,
        hessian = false,
        kwargs...)

    if specification isa RAMMatrices
        ram_matrices = specification
        identifier = StructuralEquationModels.identifier(ram_matrices)
    elseif specification isa ParameterTable
        ram_matrices = RAMMatrices(specification)
        identifier = StructuralEquationModels.identifier(ram_matrices)
    else
        throw(ErrorException("The RAMSymbolic constructor does not know how to handle your specification object. 
        \n Please specify your model as either a ParameterTable or RAMMatrices."))
    end

    n_par = length(ram_matrices.parameters)
    n_var, n_nod = ram_matrices.size_F

    par = (Symbolics.@variables θ[1:n_par])[1]

    A = zeros(Num, n_nod, n_nod)
    S = zeros(Num, n_nod, n_nod)
    !isnothing(ram_matrices.M_ind) ? M = zeros(Num, n_nod) : M = nothing
    F = zeros(ram_matrices.size_F); F[CartesianIndex.(1:n_var, ram_matrices.F_ind)] .= 1.0

    set_RAMConstants!(A, S, M, ram_matrices.constants)
    fill_A_S_M(A, S, M, ram_matrices.A_ind, ram_matrices.S_ind, ram_matrices.M_ind, par)

    A, S, F = sparse(A), sparse(S), sparse(F)

    if !isnothing(loss_types)
        any(loss_types .<: SemWLS) ? vech = true : nothing
    end

    # Σ
    Σ_symbolic = get_Σ_symbolic_RAM(S, A, F; vech = vech)
    #print(Symbolics.build_function(Σ_symbolic)[2])
    Σ_function = Symbolics.build_function(Σ_symbolic, par, expression=Val{false})[2]
    Σ = zeros(size(Σ_symbolic))
    precompile(Σ_function, (typeof(Σ), Vector{Float64}))

    # ∇Σ
    if gradient
        ∇Σ_symbolic = Symbolics.sparsejacobian(vec(Σ_symbolic), [par...])
        ∇Σ_function = Symbolics.build_function(∇Σ_symbolic, par, expression=Val{false})[2]
        constr = findnz(∇Σ_symbolic)
        ∇Σ = sparse(constr[1], constr[2], fill(1.0, nnz(∇Σ_symbolic)), size(∇Σ_symbolic)...)
        precompile(∇Σ_function, (typeof(∇Σ), Vector{Float64}))
    else
        ∇Σ_symbolic = nothing
        ∇Σ_function = nothing
        ∇Σ = nothing
    end

    if hessian
        n_sig = length(Σ_symbolic)
        n_par = size(par, 1)
        ∇²Σ_symbolic_vec = [Symbolics.sparsehessian(σᵢ, [par...]) for σᵢ in vec(Σ_symbolic)]

        @variables J[1:n_sig]
        ∇²Σ_symbolic = zeros(Num, n_par, n_par)
        for i in 1:n_sig
            ∇²Σ_symbolic += J[i]*∇²Σ_symbolic_vec[i]
        end
    
        ∇²Σ_function = Symbolics.build_function(∇²Σ_symbolic, J, par, expression=Val{false})[2]
        ∇²Σ = zeros(n_par, n_par)
    else
        ∇²Σ_symbolic = nothing
        ∇²Σ_function = nothing
        ∇²Σ = nothing
    end

    # μ
    if !isnothing(M)
        μ_symbolic = get_μ_symbolic_RAM(M, A, F)
        μ_function = Symbolics.build_function(μ_symbolic, par, expression=Val{false})[2]
        μ = zeros(size(μ_symbolic))
        if gradient
            ∇μ_symbolic = Symbolics.jacobian(μ_symbolic, [par...])
            ∇μ_function = Symbolics.build_function(∇μ_symbolic, par, expression=Val{false})[2]
            ∇μ = zeros(size(F, 1), size(par, 1))
        else
            ∇μ_function = nothing
            ∇μ = nothing
        end
    else
        μ_function = nothing
        μ = nothing
        ∇μ_function = nothing
        ∇μ = nothing
    end

    return RAMSymbolic(
        Σ_function,
        ∇Σ_function,
        ∇²Σ_function,
        Σ,
        ∇Σ,
        ∇²Σ,
        Σ_symbolic,
        ∇Σ_symbolic,
        ∇²Σ_symbolic,
        n_par,
        ram_matrices,
        μ_function,
        μ,
        ∇μ_function,
        ∇μ,
        identifier
    )
end

############################################################################
### functors
############################################################################

function (imply::RAMSymbolic)(par, F, G, H, model)
    imply.Σ_function(imply.Σ, par)
    if G || H
        imply.∇Σ_function(imply.∇Σ, par)
    end
    if !isnothing(imply.μ)
        imply.μ_function(imply.μ, par)
        if G || H
            imply.∇μ_function(imply.∇μ, par)
        end
    end
end

############################################################################
### Recommended methods
############################################################################

identifier(imply::RAMSymbolic) = imply.identifier
n_par(imply::RAMSymbolic) = imply.n_par

function update_observed(imply::RAMSymbolic, observed::SemObs; kwargs...) 
    if n_man(observed) == size(imply.Σ)
        return imply
    else
        return RAMSymbolic(;observed = observed, kwargs...)
    end
end

############################################################################
### additional methods
############################################################################

Σ(imply::RAMSymbolic) = imply.Σ
∇Σ(imply::RAMSymbolic) = imply.∇Σ
∇²Σ(imply::RAMSymbolic) = imply.∇²Σ

μ(imply::RAMSymbolic) = imply.μ
∇μ(imply::RAMSymbolic) = imply.∇μ

Σ_function(imply::RAMSymbolic) = imply.Σ_function
∇Σ_function(imply::RAMSymbolic) = imply.∇Σ_function
∇²Σ_function(imply::RAMSymbolic) = imply.∇²Σ_function

############################################################################
### additional functions
############################################################################

function get_Σ_symbolic_RAM(S, A, F; vech = false)
    invia = neumann_series(A)
    Σ_symbolic = F*invia*S*permutedims(invia)*permutedims(F)
    Σ_symbolic = Array(Σ_symbolic)
    # Σ_symbolic = Symbolics.simplify.(Σ_symbolic)
    Threads.@threads for i in eachindex(Σ_symbolic)
        Σ_symbolic[i] = Symbolics.simplify(Σ_symbolic[i])
    end
    if vech Σ_symbolic = Σ_symbolic[tril(trues(size(F, 1), size(F, 1)))] end
    return Σ_symbolic
end

function get_μ_symbolic_RAM(M, A, F)
    invia = neumann_series(A)
    μ_symbolic = F*invia*M
    μ_symbolic = Array(μ_symbolic)
    Threads.@threads for i in eachindex(μ_symbolic)
        μ_symbolic[i] = Symbolics.simplify(μ_symbolic[i])
    end
    return μ_symbolic
end

############################################################################
### Pretty Printing
############################################################################

function Base.show(io::IO, struct_inst::RAMSymbolic)
    print_type_name(io, struct_inst)
    print_field_types(io, struct_inst)
end