using Interpolations
using Plots

# Parameters
T = 10.;
β = 0.02;
bᵐᵃˣ = 0.07;
σ = 0.3;
κ = 0.25;
γ = 3.;
m = 0.;
λ = 1.;
c = 0.1;

# Meshing parameter
ρ⁻¹ = 2;

# Timesteps
P = 32 * ρ⁻¹;
dt = T/P;

# Spatial grid (2M+1 nodes)
xᵐᵃˣ = 1;
M = 8 * ρ⁻¹;
x = -xᵐᵃˣ:(1/M):xᵐᵃˣ;

# Control grid (2M₀+1 nodes)
M₀ = 2 * ρ⁻¹;
Bρ = -bᵐᵃˣ:(bᵐᵃˣ/M₀):bᵐᵃˣ;

# Normal random variables
C = 10^3 * (ρ⁻¹)^2;
ψ = randn(C);
ϕ = [ψ;-ψ]

# Solution and buffer
uⁿ = SharedArray(Float64, length(x))
u = zeros(length(x));

tic();
for n = (P - 1):-1:0
    discount = exp(-β * n * dt);
    interp = interpolate((x,), u, Gridded(Linear()));
    @sync @parallel for j = 1:length(x)
        uⁿ[j] = maximum(u - discount * (λ * abs(x - x[j]) + c));
        for b = Bρ
            X = x[j] - κ * b * dt + σ * ϕ * √dt;
            X = max(min(X, xᵐᵃˣ), -xᵐᵃˣ);
            f = -discount * ((x[j] - m)^2 + γ * b^2);
            uⁿ[j] = max(uⁿ[j], mean(interp[X]) + f * dt);
        end
    end
    u = copy(uⁿ);
end
toc();

# Solution
interp = interpolate((x,), u, Gridded(Linear()));
println("u(t=0,x=0) = $(interp[m])");