Test_ApproxFun

params=define_env()
   
   function ι(q)
        @unpack κ=params
        return((q-1.0)/κ)
    end

    function ℓ(C,k,q)
        @unpack κ,α,A=params
        return(k*((C/k+ι(q)+κ*(ι(q))^(2.0)/2.0)/A)^(1.0/(1.0-α)))
    end

    function w(C,k,q)
        @unpack γ,ψ =   params
        return (ℓ(C,k,q)^(ψ)*C^(γ))
    end


    function ν_k(C,k,q)
        @unpack α=params
        return((1.0/q)*(α/(1.0-α))*w(C,k,q)*(ℓ(C,k,q)/k))
    end
        
    function χ(C,k,q)
        @unpack α=params
        return((w(C,k,q)/(1.0-α))^(1.0-α)*(q*ν_k(C,k,q)/α)^(α))
    end

    function Y(C,k,q)
        @unpack α,A=params
        return(A*(k)^(α)*(ℓ(C,k,q))^(1-α))
    end

t=Fun(identity, 0..T)
N = (q,C,k,π,ρ,i) -> [q' - (q*(i-π+(ι(q)+κ*(ι(q))^(2.0)/2.0)/q-ι(q)-ν_k(C,k,q)));
                    C' -(σ*C*(i-π-ρ));
                    k' - ι(q)*k;
                    π' - (π*((1.0-σ)*(i-π)+σ*ρ)-((ϵ-1.0)/θ)*(χ(C,k,q)/χₙ-1.0));
                    ρ' -(-θᵨ*(ρ-ρ̄));
                    i' - (-θᵢ*(i-ϕ*π))]

    function SS()
        @unpack α,γ,σ,ϵ,θ,ϕ,ψ,ρ̄,θᵨ,θᵢ,κ,δ,A =   params

                    k_c=(α/(ρ̄))*((ϵ-1.0)/ϵ)*(1.0+θ/(ϵ-1.0)*ρ̄^(2)/(ϕ-1.0))
    
                    k_l=(A*k_c)^(1.0/(1.0-α))

            q_ss=1.0
            k_ss=(ρ̄*((1-α)/α)*(k_l)^(1+ψ)*(k_c)^(γ))^(1/(ψ+γ))
            C_ss=k_ss/k_c
            π_ss=ρ̄/(ϕ-1.0)
            ρ_ss=ρ̄
            i_ss=ϕ*π_ss
       return(π_ss=π_ss,
                C_ss=C_ss,
                q_ss=q_ss,
                k_ss=k_ss,
                ρ_ss=ρ_ss,
                i_ss=i_ss,
                ℓ_ss=k_ss/k_l,
                ι_ss=0.0)    
    end

    function bc1!(residual,u,p,t)
        @unpack q_ss,C_ss,k_ss,π_ss,ρ_ss,i_ss= SS()
        @unpack init_ρ  =   params
        residual[1] =   u[end][1]- q_ss
        residual[2] =   u[end][2]- C_ss
        residual[3] =   u[1][3]- k_ss
        residual[4] =   u[end][4]- π_ss
        residual[5] =   u[1][5]- init_ρ
        residual[6] =   u[1][6]- i_ss
    end

    @unpack q_ss,C_ss,k_ss,π_ss,ρ_ss,i_ss,ℓ_ss,ι_ss= SS()
    @unpack T,dt,init_ρ,α = params

    SS_vec = [q_ss,C_ss,k_ss,π_ss,ρ_ss,i_ss]

    u0    =   [q_ss,C_ss,k_ss,π_ss,init_ρ,i_ss]
    
@unpack σ,ϵ,θ,ϕ,ψ,ρ̄,θᵨ,θᵢ,κ,δ,A =   params
χₙ=A*(ϵ-1.0)/ϵ

t=Fun(0..T)
sys = (q,C,k,π,ρ,i,ι) -> [q' - (q*(i-π+((q-1.0)/κ+κ*(ι(q))^(2.0)/2.0)/q-ι(q)-ν_k(C,k,q)));
                    C' -(σ*C*(i-π-ρ));
                    k' - ι(q)*k;
                    π' - (π*((1.0-σ)*(i-π)+σ*ρ)-((ϵ-1.0)/θ)*(χ(C,k,q)/χₙ-1.0));
                    ρ' -(-θᵨ*(ρ-ρ̄));
                    i' - (-θᵢ*(i-ϕ*π));
                    ι-(q-1.0)/κ;
                    ι(T)-(q_ss-1.0)/κ
                    q(T)-q_ss;
                    C(T)-C_ss;
                    π(T)-π_ss;
                    k(0)-k_ss;
                    ρ(0)-init_ρ;
                    i(0)-i_ss]

u₀=[q_ss,C_ss,k_ss,π_ss,init_ρ,i_ss].*one(t)

q,C,k,π,ρ,i = newton(sys, u₀)

plot(k, label="k")

    function SS()
        @unpack σ,ϵ,θ,σ,ϕ,ψ,ρ̄,θᵨ,θᵢ =   params

        ρ_ss    =   ρ̄
        π_ss    =   ρ̄/(ϕ-1.0)
        i_ss    =   ρ̄*ϕ/(ϕ-1.0)
        x_ss    =   (1.0+θ/(ϵ-1.0)*π_ss*(σ*ρ̄+(1-σ)*(i_ss-π_ss)))^(1.0/(1.0/σ+ψ))

        return (ρ_ss=ρ_ss,π_ss=π_ss,i_ss=i_ss,x_ss=x_ss)
    end


    @unpack ρ_ss,π_ss,i_ss,x_ss= SS()
    @unpack α,γ,σ,ϵ,θ,ϕ,ψ,ρ̄,θᵨ,θᵢ,κ,δ,A =   params
    @unpack T,dt,init_ρ = params

    init    =   [x_ss,π_ss,i_ss,init_ρ]


t=Fun(identity, 0..T)

sys = (x,π,i,ρ) -> [i(0)-i_ss;
                    ρ(0)-init_ρ;
                    i(T)-i_ss;
                    ρ(0)-ρ_ss
                    x(T)-x_ss;
                    π(T)-π_ss;
                    x' - (σ*i*x - σ*π*x - σ*ρ*x);
                    π' - (-(ϵ-1.0)/θ*(x^(1.0\σ+ψ)-1.0)+(1.0-σ)*(i*π-π^2)+σ*π*ρ);
                    ρ' - (-θᵨ*ρ+θᵨ*ρ_ss);
                    i' - (-θᵢ*i+θᵢ*ϕ*π)]

x,π,i,ρ = newton(sys, [x_ss,π_ss,i_ss,ρ_ss].*one(t))

plot(x)