using QuantEcon

maxage = 5    # Maximum asset age
repcost = 75  # Replacement cost
beta = 0.9    # Discount factor
m = 2         # Number of actions; 1: keep, 2: replace;

S = [i for i in 1:maxage]
n = length(S);

# Reward array
R = Array{Float64}(n,m)
for i in 1:n
R[i, 1] = 50 - 2.5 * S[i] - 2.5 * S[i]^2
R[i, 2] = 50 - repcost
end

# Infeasible action
R[n,1] = -Inf;

R

# (Degenerate) transition probability array
Q = zeros(n,m,n)
for i in 1:n
    Q[i, 1, min(i+1, n)] = 1
    Q[i, 2, 1] = 1
end

Q

# Create a DiscreteDP
ddp = DiscreteDP(R, Q, beta);

# Solve the dynamic optimization problem (by policy iteration)
res = solve(ddp, PFI);

# Number of iterations
res.num_iter

# Optimal value function
res.v

# Optimal policy
res.sigma

# Transition probability matrix
res.mc.p

# Simulate the controlled Markov chain
initial_state_value = 1
nyrs = 12
spath = simulate(res.mc, nyrs+1, init=initial_state_value);

using PyPlot

import PyPlot.plt
fig, axes = plt[:subplots](1, 2, figsize=(12, 4))
axes[1][:plot]( [i for i in 1:n], res.v)
axes[1][:set_xlim](1, 5)
axes[1][:set_ylim](160, 220)
axes[1][:set_xticks](linspace(1, 5, 5))
axes[1][:set_xlabel]("Age of Machine")
axes[1][:set_ylabel]("Value")
axes[1][:set_title]("Optimal Value Function")

axes[2][:plot](spath)
axes[2][:set_xlim](0, nyrs)
axes[2][:set_ylim](1, 4)
axes[2][:set_yticks](linspace(1, 4, 4))
axes[2][:set_xlabel]("Year")
axes[2][:set_ylabel]("Age of Machine")
axes[2][:set_title]("Optimal State Path")