#=

Author: Shunsuke Hori

=#

using LS

using QuantEcon
using NLopt
using NLsolve
using Dierckx

Para=LS.Para
Planners_Allocation_Sequential = LS.Planners_Allocation_Sequential
"""
Bellman equation for the continuation of the Lucas-Stokey Problem
"""
mutable struct BellmanEquation{TP <: Para, TI <: Integer, TR <: Real}
    para::TP
    S::TI
    xbar::Array{TR}
    time_0::Bool
    z0::Array{Array}
    cFB::Vector{TR}
    nFB::Vector{TR}
    xFB::Vector{TR}
    zFB::Vector{Vector{TR}}
end

"""
Compute the planner's allocation by solving Bellman
equation.
"""
struct Planners_Allocation_Bellman{TP <: Para, TI <: Integer,
                                   TVg <: AbstractVector, TT <: Tuple}
    para::TP
    mc::MarkovChain
    S::TI
    T::BellmanEquation
    mugrid::TVg
    xgrid::TVg
    Vf::Array
    policies::TT
end

"""
Initializes the type from the calibration Para
"""
function Planners_Allocation_Bellman(para::Para, mugrid::AbstractArray)
    G = para.G
    S = size(para.Pi, 1) # number of states
    mc=MarkovChain(para.Pi)
    #now find the first best allocation
    Vf, policies, T, xgrid = solve_time1_bellman(para, mugrid)
    T.time_0 = true #Bellman equation now solves time 0 problem
    return Planners_Allocation_Bellman(para, mc, S, T, mugrid, xgrid, Vf, policies)
end

"""
Solve the time  1 Bellman equation for calibration Para and initial grid mugrid
"""
function solve_time1_bellman{TR <: Real}(para::Para{TR}, mugrid::AbstractArray)
    Pi = para.Pi
    S = size(para.Pi, 1)

    #First get initial fit from lucas stockey solution.
    #Need to change things to be ex_ante
    PP = LS.Planners_Allocation_Sequential(para)

    function incomplete_allocation(PP::LS.Planners_Allocation_Sequential,
                                    mu_::AbstractFloat, s_::Integer)
        c, n, x, V = LS.time1_value(PP, mu_)
        return c,n,dot(Pi[s_, :], x), dot(Pi[s_, :], V)
    end
    cf = Array{Function}(S, S)
    nf = Array{Function}(S, S)
    xprimef = Array{Function}(S, S)
    Vf = Vector{Function}(S)
    # cf = Array{LinInterp}(S, S)
    # nf = Array{LinInterp}(S, S)
    # xprimef = Array{LinInterp}(S, S)
    # Vf = Vector{LinInterp}(S)
    xgrid = Array{TR}(S, length(mugrid))
    for s_ in 1:S
        c = Array{TR}(length(mugrid), S)
        n = Array{TR}(length(mugrid), S)
        x = Array{TR}(length(mugrid))
        V = Array{TR}(length(mugrid))
        for (i_mu, mu) in enumerate(mugrid)
            c[i_mu, :], n[i_mu, :], x[i_mu], V[i_mu] =
                incomplete_allocation(PP, mu, s_)
        end
        xprimes = repmat(x, 1, S)
        xgrid[s_, :] = x
        for sprime = 1:S
            splc = Spline1D(x[end:-1:1], c[:, sprime][end:-1:1], k=3)
            spln = Spline1D(x[end:-1:1], n[:, sprime][end:-1:1], k=3)
            splx = Spline1D(x[end:-1:1], xprimes[:, sprime][end:-1:1], k=3)
            cf[s_, sprime] = y -> splc(y)
            nf[s_, sprime] = y -> spln(y)
            xprimef[s_, sprime] = y -> splx(y)
            # cf[s_, sprime] = LinInterp(x[end:-1:1], c[:, sprime][end:-1:1])
            # nf[s_, sprime] = LinInterp(x[end:-1:1], n[:, sprime][end:-1:1])
            # xprimef[s_, sprime] = LinInterp(x[end:-1:1], xprimes[:, sprime][end:-1:1])
        end
        splV = Spline1D(x[end:-1:1], V[end:-1:1], k=3)
        Vf[s_] = y -> splV(y)
        # Vf[s_] = LinInterp(x[end:-1:1], V[end:-1:1])
    end

    policies = [cf, nf, xprimef]

    #create xgrid
    xbar = [maximum(minimum(xgrid)), minimum(maximum(xgrid))]
    xgrid = linspace(xbar[1], xbar[2], length(mugrid))

    #Now iterate on Bellman equation
    T = BellmanEquation(para, xgrid, policies)
    diff = 1.0
    while diff > 1e-6
        PF = (i_x, x, s) -> get_policies_time1(T, i_x, x, s, Vf, xbar)
        Vfnew, policies = fit_policy_function(T, PF, xgrid)

        diff = 0.0
        for s=1:S
            diff = max(diff,maximum(abs,
                       (Vf[s].(xgrid)-Vfnew[s].(xgrid))./Vf[s].(xgrid)))
        end

        println("diff = $diff")
        Vf = copy(Vfnew)
    end

    return Vf, policies, T, xgrid
end

"""
Fits the policy functions
"""
function fit_policy_function{TF<:AbstractFloat}(
                    T::BellmanEquation,
                    PF::Function, xgrid::AbstractVector{TF})
    S = T.S
    # preallocation
    PFvec = Array{TF}(4S+1, length(xgrid))
    cf = Array{Function}(S, S)
    nf = Array{Function}(S, S)
    xprimef = Array{Function}(S, S)
    TTf = Array{Function}(S, S)
    Vf = Vector{Function}(S)
    # fit policy fuctions
    for s_ in 1:S
        for (i_x, x) in enumerate(xgrid)
            PFvec[:, i_x] = PF(i_x, x, s_)
        end
        splV = Spline1D(xgrid, PFvec[1,:], k=3)
        Vf[s_] = y -> splV(y)
        # Vf[s_] = LinInterp(xgrid, PFvec[1, :])
        for sprime=1:S
            splc = Spline1D(xgrid, PFvec[1+sprime,:], k=3)
            spln = Spline1D(xgrid, PFvec[1+S+sprime,:], k=3)
            splxprime = Spline1D(xgrid, PFvec[1+2S+sprime,:], k=3)
            splTT = Spline1D(xgrid, PFvec[1+3S+sprime,:], k=3)
            cf[s_,sprime] = y -> splc(y)
            nf[s_,sprime] = y -> spln(y)
            xprimef[s_,sprime] = y -> splxprime(y)
            TTf[s_,sprime] = y -> splTT(y)
        end
    end
    policies = (cf, nf, xprimef, TTf)
    return Vf, policies
end

"""
Computes Tau given c,n
"""
function Tau(pab::Planners_Allocation_Bellman,
            c::AbstractArray, n::AbstractArray)
    para = pab.para
    Uc, Un = para.Uc(c, n), para.Un(c, n)
    return 1+Un./(para.Theta .* Uc)
end

Tau(pab::Planners_Allocation_Bellman, c::Real, n::Real) =
    Tau(pab, [c], [n])

"""
Finds the optimal allocation given initial government debt B_ and state s_0
"""
function time0_allocation(pab::Planners_Allocation_Bellman,
                           B_::Real, s0::Integer)
    T, Vf = pab.T, pab.Vf
    xbar = T.xbar
    z0 = get_policies_time0(T, B_, s0, Vf, xbar)

    c0, n0, xprime0, T0 = z0[2], z0[3], z0[4], z0[5]
    return c0, n0, xprime0, T0
end

"""
Simulates planners policies for `T` periods
"""
function simulate{TF <: AbstractFloat}(pab::Planners_Allocation_Bellman,
                   B_::TF, s_0::Integer, T::Integer,
                   sHist::Vector=QuantEcon.simulate(pab.mc, T, init=s_0))
    para, mc, Vf, S = pab.para, pab.mc, pab.Vf, pab.S
    Pi, Uc = para.Pi, para.Uc
    cf, nf, xprimef, TTf = pab.policies

    cHist=Array{TF}(T)
    nHist=Array{TF}(T)
    Bhist=Array{TF}(T)
    xHist=Array{TF}(T)
    TauHist=Array{TF}(T)
    THist=Array{TF}(T)
    muHist=Array{TF}(T)

    #time0
    cHist[1], nHist[1], xHist[1], THist[1]  = time0_allocation(pab, B_, s_0)
    TauHist[1] = Tau(pab, cHist[1], nHist[1])[s_0]
    Bhist[1] = B_
    muHist[1] = Vf[s_0](xHist[1])

    #time 1 onward
    for t in 2:T
        s_, x, s = sHist[t-1], xHist[t-1], sHist[t]
        c = Array{TF}(S)
        n = Array{TF}(S)
        xprime = Array{TF}(S)
        TT = Array{TF}(S)
        for sprime=1:S
            c[sprime], n[sprime], xprime[sprime], TT[sprime] =
                cf[s_, sprime](x), nf[s_, sprime](x),
                xprimef[s_, sprime](x), TTf[s_, sprime](x)
        end

        Tau_val = Tau(pab, c, n)[s]
        u_c = Uc(c, n)
        Eu_c = dot(Pi[s_, :], u_c)

        muHist[t] = Vf[s](xprime[s])

        cHist[t], nHist[t], Bhist[t], TauHist[t] = c[s], n[s], x/Eu_c, Tau_val
        xHist[t], THist[t] = xprime[s], TT[s]
    end
    return cHist, nHist, Bhist, xHist, TauHist, THist, muHist, sHist
end

"""
Initializes the class from the calibration Para
"""
function BellmanEquation{TF <: AbstractFloat}(para::Para{TF},
                                              xgrid::AbstractVector{TF},
                                              policies0::Array)
    S = size(para.Pi, 1) # number of states
    xbar = [minimum(xgrid), maximum(xgrid)]
    time_0 = false
    z0 = Array{Array}(length(xgrid), S)
    cf, nf, xprimef = policies0[1], policies0[2], policies0[3]
    for s in 1:S
        for (i_x, x) in enumerate(xgrid)
            cs=Array{TF}(S)
            ns=Array{TF}(S)
            xprimes=Array{TF}(S)
            for j=1:S
                cs[j], ns[j], xprimes[j] = cf[s, j](x), nf[s, j](x), xprimef[s, j](x)
            end
            z0[i_x, s] = vcat(cs, ns, xprimes, zeros(S))
        end
    end
    cFB, nFB, IFB, xFB, zFB = LS.find_first_best(para, S, 2)
    return BellmanEquation(para, S, xbar, time_0, z0, cFB, nFB, xFB, zFB)
end


"""
Finds the optimal policies
"""
function get_policies_time1(T::BellmanEquation,
        i_x::Integer, x::Real, s_::Integer, Vf::AbstractArray{Function}, xbar::AbstractVector)
    para, S = T.para, T.S
    beta, Theta, G, Pi = para.beta, para.Theta, para.G, para.Pi
    U,Uc,Un = para.U, para.Uc, para.Un

    S_possible = sum(Pi[s_, :].>0)
    sprimei_possible = find(Pi[s_, :].>0)

    function objf(z, grad)
        c, xprime = z[1:S_possible], z[S_possible+1:2S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        Vprime = [Vf[sprimei_possible[si]](xprime[si]) for si in 1:S_possible]
        return -dot(Pi[s_, sprimei_possible], U.(c, n)+beta*Vprime)
    end

    function cons(out, z, grad)
        c, xprime, TT =
            z[1:S_possible], z[S_possible+1:2S_possible], z[2S_possible+1:3S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        u_c = Uc.(c, n)
        Eu_c = dot(Pi[s_, sprimei_possible], u_c)
        out .= x*u_c/Eu_c - u_c.*(c-TT)-Un(c, n).*n - beta*xprime
    end
    function cons_no_trans(out, z, grad)
        c, xprime = z[1:S_possible], z[S_possible+1:2S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        u_c = Uc.(c, n)
        Eu_c = dot(Pi[s_, sprimei_possible], u_c)
        out .= x*u_c/Eu_c - u_c.*c-Un(c, n).*n - beta*xprime
    end

    if para.transfers == true
        lb = vcat(zeros(S_possible), ones(S_possible)*xbar[1], zeros(S_possible))
        if para.n_less_than_one == true
            ub = vcat(100*ones(S_possible)-G[sprimei_possible],
                      ones(S_possible)*xbar[2], 100*ones(S_possible))
        else
            ub = vcat(100*ones(S_possible),
                      ones(S_possible)*xbar[2],
                      100*ones(S_possible))
        end
        init = vcat(T.z0[i_x, s_][sprimei_possible],
                    T.z0[i_x, s_][2S+sprimei_possible],
                    T.z0[i_x, s_][3S+sprimei_possible])
        opt = Opt(:LN_COBYLA, 3S_possible)
        equality_constraint!(opt, cons, zeros(S_possible))
    else
        lb = vcat(zeros(S_possible), ones(S_possible)*xbar[1])
        if para.n_less_than_one == true
            ub = vcat(ones(S_possible)-G[sprimei_possible], ones(S_possible)*xbar[2])
        else
            ub = vcat(ones(S_possible), ones(S_possible)*xbar[2])
        end
        init = vcat(T.z0[i_x, s_][sprimei_possible],
                    T.z0[i_x, s_][2S+sprimei_possible])
        opt = Opt(:LN_COBYLA, 2S_possible)
        equality_constraint!(opt, cons_no_trans, zeros(S_possible))
    end
    init[init .> ub] = ub[init .> ub]
    init[init .< lb] = lb[init .< lb]

    min_objective!(opt, objf)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 10000000)
    maxtime!(opt, 10)
    ftol_rel!(opt, 1e-8)
    ftol_abs!(opt, 1e-8)

    (minf, minx, ret) = optimize(opt, init)

    if ret != :SUCCESS && ret != :ROUNDOFF_LIMITED && ret != :MAXEVAL_REACHED &&
         ret != :FTOL_REACHED && ret != :MAXTIME_REACHED
        error("optimization failed: ret = $ret")
    end

    T.z0[i_x, s_][sprimei_possible] = minx[1:S_possible]
    T.z0[i_x, s_][S+sprimei_possible] = minx[1:S_possible]+G[sprimei_possible]
    T.z0[i_x, s_][2S+sprimei_possible] = minx[S_possible+1:2S_possible]
    if para.transfers == true
        T.z0[i_x, s_][3S+sprimei_possible] = minx[2S_possible+1:3S_possible]
    else
        T.z0[i_x, s_][3S+sprimei_possible] = zeros(S)
    end

    return vcat(-minf, T.z0[i_x, s_])
end

"""
Finds the optimal policies
"""
function get_policies_time0(T::BellmanEquation,
        B_::Real, s0::Integer, Vf::AbstractArray{Function}, xbar::AbstractVector)
    para = T.para
    beta, Theta, G = para.beta, para.Theta, para.G
    U, Uc, Un = para.U, para.Uc, para.Un

    function objf(z, grad)
        c, xprime = z[1], z[2]
        n=(c+G[s0])/Theta[s0]
        return -(U(c, n)+beta*Vf[s0](xprime))
    end

    function cons(z,grad)
        c, xprime, TT = z[1], z[2], z[3]
        n=(c+G[s0])/Theta[s0]
        return -Uc(c, n)*(c-B_-TT)-Un(c, n)*n - beta*xprime
    end
    cons_no_trans(z, grad) = cons(vcat(z, 0), grad)

    if para.transfers == true
        lb = [0.0, xbar[1], 0.0]
        if para.n_less_than_one == true
            ub = [1-G[s0], xbar[2], 100]
        else
            ub = [100.0, xbar[2], 100.0]
        end
        init = vcat(T.zFB[s0][1], T.zFB[s0][3], T.zFB[s0][4])
        init = [0.95124922, -1.15926816,  0.0]
        opt = Opt(:LN_COBYLA, 3)
        equality_constraint!(opt, cons)
    else
        lb = [0.0, xbar[1]]
        if para.n_less_than_one == true
            ub = [1-G[s0], xbar[2]]
        else
            ub = [100, xbar[2]]
        end
        init = vcat(T.zFB[s0][1], T.zFB[s0][3])
        init = [0.95124922, -1.15926816]
        opt = Opt(:LN_COBYLA, 2)
        equality_constraint!(opt, cons_no_trans)
    end
    init[init.> ub] = ub[init.> ub]
    init[init.< lb] = lb[init.< lb]


    min_objective!(opt, objf)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 100000000)
    maxtime!(opt, 30)

    (minf, minx, ret) = optimize(opt, init)

    if ret != :SUCCESS && ret != :ROUNDOFF_LIMITED && ret != :MAXEVAL_REACHED && ret != :FTOL_REACHED
        error("optimization failed: ret = $ret")
    end

    if para.transfers == true
        return -minf, minx[1], minx[1]+G[s0], minx[2], minx[3]
    else
        return -minf, minx[1], minx[1]+G[s0], minx[2], 0
    end
end
