using Pkg;Pkg.activate("../probprog/workspace/");Pkg.instantiate()

using Random
Random.seed!(87613) # Set random seed

using LinearAlgebra
using PyPlot
using ForneyLab

include("environment_2d.jl") # Include environmental dynamics
include("helpers_2d.jl") # Include helper functions for plotting
include("agent_2d.jl") # Load agent's internal beliefs over external dynamics
;

# Internal model perameters
Gamma   = 100.0*diageye(3) # Transition precision
Phi     = 10.0*diageye(3) # Observation precision
Upsilon = Diagonal([0.25*pi, 1.0].^2) # Control prior variance
Sigma   = Diagonal([1.0, 30.0, 30.0]) # Goal prior variance
;

T = 10 # Lookahead

# Build internal model
fg = FactorGraph()

o   = Vector{Variable}(undef, T) # Observed states
s_g = Vector{Variable}(undef, T) # Noiseless internal states
s   = Vector{Variable}(undef, T) # Noisy internal states
u   = Vector{Variable}(undef, T) # Control states

@RV s_t_min ~ GaussianMeanVariance(placeholder(:m_s_t_min, dims=(3,)),
                                   placeholder(:V_s_t_min, dims=(3,3))) # Prior internal state
u_t = placeholder(:u_t, dims=(2,))
@RV u[1]   ~ GaussianMeanVariance(u_t, tiny*diageye(2))
@RV s_g[1] ~ Nonlinear{Unscented}(s_t_min, u[1], g=g,g_inv=[g_inv_in1,g_inv_in2])
@RV s[1]   ~ GaussianMeanPrecision(s_g[1], Gamma)
@RV o[1]   ~ GaussianMeanPrecision(s[1], Phi)
placeholder(o[1], :o_t, dims=(3,))

s_k_min = s[1]
for k=2:T
    @RV u[k]   ~ GaussianMeanVariance(zeros(2), Upsilon)
    @RV s_g[k] ~ Nonlinear{Unscented}(s_k_min, u[k], g=g, g_inv=[g_inv_in1,g_inv_in2]) # Noiseless transition model
    @RV s[k]   ~ GaussianMeanPrecision(s_g[k], Gamma) # State transition model including noise
    @RV o[k]   ~ GaussianMeanPrecision(s[k], Phi) # Observation model
    GaussianMeanVariance(o[k], 
                         placeholder(:m_o, var_id=:m_o_*k, index=k-1, dims=(3,)),
                         placeholder(:V_o, var_id=:V_o_*k, index=k-1, dims=(3,3))) # Goal prior
    s_k_min = s[k]
end
;

# Schedule message passing algorithm
q = PosteriorFactorization(fg)
algo = messagePassingAlgorithm(u[2]) # Infer internal states
source = algorithmSourceCode(algo)
eval(Meta.parse(source)) # Loads the step!() function for inference
inspectSnippet(source) # Inspect a snippet of the algorithm code
;

# Initial state
ϕ_0 = pi # Initial orientation
x_0 = 2.0 # Initial coordinates
y_0 = 3.0
;

N = 20 # Total simulation time

(execute, observe)  = initializeWorld() # Let there be a world
(infer, act, slide) = initializeAgent() # Let there be an agent

# Step through action-perception loop
u_hat = Vector{Vector{Float64}}(undef, N) # Actions
o_hat = Vector{Vector{Float64}}(undef, N) # Observations
for t=1:N
    u_hat[t] = act() # Evoke an action from the agent
               execute(u_hat[t]) # The action influences hidden external states
    o_hat[t] = observe() # Observe the current environmental outcome (update p)
               infer(u_hat[t], o_hat[t]) # Infer beliefs from current model state (update q)
               slide() # Prepare for next iteration
end
;

# Plot active inference results
plotTrajectory(u_hat, o_hat)
;