DecisionMakingPolicies

using Distributions
using DecisionMakingPolicies
import Random
using ComponentArrays

num_actions = 4
policy = StatelessSoftmax(num_actions, init_weights=(rng, dims...)->zeros(Float64, dims...))
rng = Random.default_rng()

ps, st = setup(rng, policy)  # initializes policy parameters and state of the policy. 
ps = ps |> ComponentArray  # map the parameters to vector representation for easy manipulation. 

d, st = policy(ps, st)  # returns a distribution over actions along with new state of the policy.
a = rand(d)    # samples an action from the distribution d
logp = logpdf(d, a) # uses the Distributions.jl computation of the logpdf

num_features = 4
num_actions = 3
policy = LinearSoftmax(num_features, num_actions)   # default initialization is zeros of type Float32
ps, st = setup(rng, policy)  
x = randn(Float32, num_features)

d, st = policy(x, ps, st)  # compute the conditional distributon over actions given the feature vector x
a = rand(d)  # sample the actions
logp = logpdf(d,a) # compute the log probability of a given x using Distributions.jl
logp, st = logpdf(policy, x, a, ps, st)  # log probability of a given x
logps, st = logpdf(policy, x, a, ps, st)  # compute the log probability for each action

logp, ψ, st = grad_logpdf(policy, x, a, ps, st)  # compute the partial derivative of the log probability of a given a with respect to the policy parameters

a, logp, ψ, st = sample_with_trace(policy, x, ps, st)  # sample an action, compute log probability of that action, and compute the partial derivative of with respect to the policy weights. This is faster than doing things individually. 

policy = LinearSoftmax(num_features, num_actions, buffered=true)
ps, st = setup(rng, policy)

x = randn(num_features)
d = policy(x, ps, st)
a = 1 
logp, ψ, st = grad_logpdf(policy, x, a, ps, st)
a, logp, ψ, st = sample_with_trace(policy, x, ps, st)

using Zygote

policy = LinearSoftmax(num_features, num_actions)
ps, st = setup(rng, policy)
x = rand(num_features)
a = 1
g = gradient(p->logpdf(policy(x, p, st)[1],a), ps)[1]  # gradient with respect to the policy parameters
# Going through the Distributions.jl interface is slow to compute the gradients. The faster was is:
g = gradient(p->logpdf(policy, x, a, ps, st)[1], ps)[1]

using DecisionMakingUtils
using BenchmarkEnvironments
import Lux

env = cartpole_finitetime()

# construct a Fourier basis function 
# note that this basis construction is an old style and will be migrated to the lux style soon. 
# However this shows a valid way to create an arbitrary basis function and incorporate it into the policy.
function make_linearbasis(env, order)
    X = env.X # observation ranges
    n = size(X,1)
    
    nrm = ZeroOneNormalization(X)
    nbuff = zeros(n)
    nrm = BufferedFunction(nrm, nbuff)  # optional preallocations of the normalization function (do not use buffers unless you want single sample)
    fb = FourierBasis(n, order, order, false)  # create Fourier basis with given order
    num_features = length(fb)
    fbuff = FourierBasisBuffer(fb)
    fb = BufferedFunction(fb, fbuff) # optional preallocation of the basis function
    ϕ(x) = fb(nrm(x))  # just some function that takes in an input x and outputs a vector of features
    return ϕ, num_features
end

ϕ, num_features = make_linearbasis(env, 3)
num_actions = length(env.A)
plinear = LinearSoftmax(num_features, num_actions, buffered=true)
policy = LinearPolicyWithBasis(ϕ, plinear)
ps, st = setup(rng, policy)

s,x = env.d0()
a, logp, ψ, st = sample_with_trace(policy, x, ps, st)

function nn_policy(env, num_layers, width)
    X = env.X 
    n = size(X,1)
    num_actions = length(env.A)
    
    act_fn = Lux.tanh  # can be whatever you want
    layers = []
    for i in 1:num_layers
        push!(layers,Lux.Dense(n => width, act_fn))
        n = width
    end
    
    model = Lux.Chain(layers...)
    linear_policy = LinearSoftmax(n, num_actions)
    policy = DeepPolicy(model, linear_policy)
    return policy
end

env = cartpole_finitetime()
policy = nn_policy(env, 2, 32)
ps, st = setup(rng, policy)

s, x = env.d0()
d, st = policy(x, ps, st)
a = rand(d)
g = gradient(p->logpdf(policy, x, a, p, st)[1], ps)[1]
logp, g2, st = grad_logpdf(policy, x, a, ps, st)

input_dim = 4
action_dim = 3
policy = LinearNormal(input_dim, action_dim; init_weights=(rng, dims...)->zeros(Float64, dims...), init_std=(rng, dims...)->zeros(Float64, dims...))
ps, st = setup(rng, policy)

batch_size = 10
X = randn(Float64, input_dim, batch_size)  # last dim should always be batch_size
A = randn(Float64, action_dim, batch_size) 

logps, st = logpdf(policy, X, A, ps, st)  # compute log probability of each action in the batch
println(size(logps))  # (batch_size,)

logps, ψ, st = grad_logpdf(policy, X, A, ps, st)  # compute log probability of each action in the batch and the gradient of the sum log probabilities with respect to the policy parameters. 
println(size(logps))  # (batch_size,)