# RAPTOR: Round-Based Public Transit Optimized Routing
# Described in Delling, D., Pajor, T., & Werneck, R. (2012). Round-Based Public Transit Routing. http://research.microsoft.com/pubs/156567/raptor_alenex.pdf

using Dates

const MAX_TIME = 2147483646
const XFER_ROUTE = -1
const INT_MISSING = -2
const ORIGIN = -3
const BOARD_SLACK_SECONDS = 60
const EMPTY_SET = BitSet()

struct StopAndTime
    stop::Int64
    time::Int32
end

struct RaptorRequest
    origins::Vector{StopAndTime}
    max_rides::Int64
    date::Date
    walk_speed_meters_per_second::Float64
end

struct RaptorResult
    times_at_stops_each_round::Array{Int32, 2}
    prev_stop::Array{Int64, 2}
    prev_route::Array{Int64, 2}
    prev_boardtime::Array{Int32, 2}
end

function empty_no_resize!(s::BitSet)
    for set_bit in s
        delete!(s, set_bit)
    end
end

function raptor(net::TransitNetwork, req::RaptorRequest)
    nstops = length(net.stops)
    # * 2 for transfer rounds, + 1 for origin times, -1 for skipped final transfer round
    nrounds = req.max_rides * 2 + 1 - 1

    # these get allocated here, and the core RAPTOR algorithm should have 0 allocations
    times_at_stops::Array{Int32,2} = fill(MAX_TIME, (nrounds, nstops))
    prev_stop::Array{Int64,2} = fill(INT_MISSING, (nrounds, nstops))
    prev_route::Array{Int64,2} = fill(INT_MISSING, (nrounds, nstops))
    prev_boardtime::Array{Int32,2} = fill(INT_MISSING, (nrounds, nstops))
    # set bit 0 so that offset is forced to zero and there aren't allocations later
    prev_touched_stops::BitSet = BitSet([0])
    touched_stops::BitSet = BitSet([0])
    sizehint!(prev_touched_stops, nstops)
    sizehint!(touched_stops, nstops)

    # unset the 0 set bits that were to force offset
    delete!(prev_touched_stops, 0)
    delete!(touched_stops, 0)

    # initialize times at stops
    for sat in req.origins
        times_at_stops[1, sat.stop] = sat.time
        push!(prev_touched_stops, sat.stop)
    end

    @assert prev_touched_stops.offset == 0
    @assert touched_stops.offset == 0

    # get which service idxes are running
    services_running = BitSet(map(t -> t[1], filter(t -> is_service_running(t[2], req.date), collect(enumerate(net.services)))))

    @info "$(length(services_running)) services running on requested date"

    # ideally this would have no allocations, although it does have a few due to empty!ing and push!ing to the bitsets - would be nice to have
    # a bounded bitset implementation that did not dynamically resize.
    @time run_raptor!(net, times_at_stops, prev_stop, prev_route, prev_boardtime, req, services_running, prev_touched_stops, touched_stops)

    return RaptorResult(
        times_at_stops,
        prev_stop,
        prev_route,
        prev_boardtime
    )
end

function run_raptor!(net::TransitNetwork, times_at_stops::Array{Int32, 2}, prev_stop::Array{Int64, 2},
    prev_route::Array{Int64, 2}, prev_boardtime::Array{Int32, 2}, req::RaptorRequest, services_running::BitSet, prev_touched_stops::BitSet, touched_stops::BitSet)
    for round in 1:req.max_rides
        # where the results of this round will be recorded
        target = round * 2

        for stop in prev_touched_stops
            # find all patterns that touch this stop
            for patidx in net.patterns_for_stop[stop]
                tp = net.patterns[patidx]

                if !in(tp.service, services_running)
                    continue
                end

                # TODO handle loop trips
                stoppos = INT_MISSING

                # not using a vectorized findfirst to avoid allocations
                for (i, tpstop) in enumerate(tp.stops)
                    if tpstop == stop
                        stoppos = i
                        break
                    end
                end

                @assert stoppos != INT_MISSING

                # find the trip that departs at or after the earliest possible departure
                earliest_departure::Int32 = times_at_stops[target - 1, stop] + BOARD_SLACK_SECONDS
                best_departure::Int32 = MAX_TIME
                best_trip_idx::Int64 = INT_MISSING

                for tripidx in net.trips_for_pattern[patidx]
                    trip = net.trips[tripidx]
                    time_at_stop = trip.stop_times[stoppos].departure_time
                    if (time_at_stop >= earliest_departure && time_at_stop < best_departure)
                        best_trip_idx = tripidx
                        best_departure = time_at_stop
                        # pre-sorting trips by departure time would allow us to break here, but I think we found in R5 that it didn't really help
                    end
                end

                if best_trip_idx == INT_MISSING
                    continue  # no possible trip to board
                end

                best_trip = net.trips[best_trip_idx]

                for stidx in stoppos + 1:length(best_trip.stop_times)
                    stop_time = best_trip.stop_times[stidx]
                    if stop_time.arrival_time < times_at_stops[target, stop_time.stop]
                        # we have found a new fastest way to get to this stop!
                        times_at_stops[target, stop_time.stop] = stop_time.arrival_time
                        prev_stop[target, stop_time.stop] = stop
                        prev_route[target, stop_time.stop] = best_trip.route
                        prev_boardtime[target, stop_time.stop] = best_trip.stop_times[stoppos].departure_time
                        push!(touched_stops, stop_time.stop)
                    end
                end
            end
        end # loop over prev_touched_stop

        #@info "round $round found $(length(touched_stops)) stops accessible by transit"

        # clear prev_touched_stops and reuse as next_touched_stops, avoid allocation
        empty!(prev_touched_stops)
        next_touched_stops = prev_touched_stops

        # do transfers, but skip after last iteration
        if round < req.max_rides
            times_at_stops[target + 1, :] = times_at_stops[target, :]
            # leave other things missing if there were no transfers
            for stop in touched_stops
                push!(next_touched_stops, stop)  # this stop was touched by this round

                for xfer in net.transfers[stop]
                    xfer_walk_time = Base.round(xfer.distance_meters / req.walk_speed_meters_per_second)
                    pre_xfer_time = times_at_stops[target, stop]
                    time_after_xfer = pre_xfer_time + xfer_walk_time
                    if time_after_xfer < times_at_stops[target + 1, xfer.target_stop]
                        # transferring to this stop is optimal!
                        times_at_stops[target + 1, xfer.target_stop] = time_after_xfer
                        prev_stop[target + 1, xfer.target_stop] = stop
                        prev_route[target + 1, xfer.target_stop] = XFER_ROUTE
                        prev_boardtime[target + 1, stop] = pre_xfer_time
                        push!(next_touched_stops, xfer.target_stop)
                    end
                end
            end
        end

        # prepare for next iteration
        prev_touched_stops = next_touched_stops
        empty!(touched_stops)
    end
end


for (func, typ) in [(:clblasStrmm, cl.CL_float),
                    (:clblasDtrmm, cl.CL_double),
                    (:clblasCtrmm, CL_float2),
                    (:clblasZtrmm, CL_double2)]

    @eval @blasfun $func(order::clblasOrder, side::clblasSide,
                              uplo::clblasUplo, transA::clblasTranspose,
                              diag::clblasDiag,
                              M::Csize_t, N::Csize_t,
                              alpha::$typ,
                              A::cl.CL_mem, offA::Csize_t, lda::Csize_t,
                              B::cl.CL_mem, offB::Csize_t, ldb::Csize_t,
                              n_queues::cl.CL_uint,
                              queues::Ptr{cl.CL_command_queue},
                              n_events_in_wait_list::cl.CL_uint,
                              event_wait_list::Ptr{cl.CL_event},
                              events::Ptr{cl.CL_event})

    @eval @blasfun2 $func(order::clblasOrder, side::clblasSide,
                               uplo::clblasUplo, transA::clblasTranspose,
                               diag::clblasDiag,
                               M::Csize_t, N::Csize_t,
                               alpha::$typ,
                               A::cl.CL_mem, offA::Csize_t, lda::Csize_t,
                               B::cl.CL_mem, offB::Csize_t, ldb::Csize_t)

end


for (T, N) ∈ ((:MPO, 4), (:MPS, 3))
    CuT = Symbol(:Cu, T)
    AT = Symbol(:Abstract, T)
    @eval begin
        export $CuT, $T

        struct $CuT{T <: Number} <: $AT{T}
            tensors::Vector{CuArray{T, $N}}
        end

        $CuT(::Type{T}, L::Int) where {T} = $CuT(Vector{CuArray{T, $N}}(undef, L))
        $CuT(L::Int) = $CuT(Float32, L)

        Base.copy(a::$CuT) = $CuT(copy(a.tensors))

        $T(A::$CuT) = $T([Array(arr) for arr in A.tensors])
        $CuT(A::$T) = $CuT([cu(arr) for arr in A.tensors])
    end
end
@inline CuMPS(A::AbstractArray) = CuMPS(cu(A), :right)
@inline CuMPS(A::AbstractArray, s::Symbol, args...) = CuMPS(cu(A), Val(s), typemax(Int), args...)
@inline CuMPS(A::AbstractArray, s::Symbol, Dcut::Int, args...) = CuMPS(cu(A), Val(s), Dcut, args...)
@inline CuMPS(A::AbstractArray, ::Val{:right}, Dcut::Int, args...) = _left_sweep_SVD(CuMPS, cu(A), Dcut, args...)
@inline CuMPS(A::AbstractArray, ::Val{:left}, Dcut::Int, args...) = _right_sweep_SVD(CuMPS, cu(A), Dcut, args...)

function CUDA.randn(::Type{CuMPS{T}}, L::Int, D::Int, d::Int) where {T}
    ψ = CuMPS(L)
    ψ[1] = CUDA.randn(T, 1, d, D)
    for i ∈ 2:(L-1)
        ψ[i] = CUDA.randn(T, D, d, D)
    end
    ψ[end] = CUDA.randn(T, D, d, 1)
    ψ
end

CUDA.randn(::Type{CuMPS}, args...) = CUDA.randn(CuMPS{Float32}, args...)

function CUDA.randn(::Type{CuMPO{T}}, L::Int, D::Int, d::Int) where {T}
    ψ = CUDA.randn(CuMPS{T}, L, D, d^2)
    CuMPO(ψ)
end

CUDA.randn(::Type{CuMPO}, args...) = CUDA.randn(CuMPO{Float32}, args...)

CuMPS(states::Vector{CuVector{T}}) where {T <: Number} = CuMPS(
    [reshape(copy(v), (1, length(v), 1)) for v ∈ states]
)

function is_left_normalized(ψ::CuMPS)
    for i ∈ eachindex(ψ)
        A = ψ[i]
        DD = size(A, 3)

        @cutensor Id[x, y] := conj(A[α, σ, x]) * A[α, σ, y] order = (α, σ)
        if norm(cu(I(DD)) - Id) > 1e-6 return false end
    end
    true
end

function is_right_normalized(ϕ::CuMPS)
    for i ∈ eachindex(ϕ)
        B = ϕ[i]
        DD = size(B, 1)

        @cutensor Id[x, y] := B[x, σ, α] * conj(B[y, σ, α]) order = (α, σ)
        if norm(cu(I(DD)) - Id) > 1e-6 return false end
    end
    true
end

function _truncate_svd(F::CuSVD, Dcut::Int)
    U, Σ, V = F
    c = min(Dcut, size(U, 2))
    U = U[:, 1:c]
    Σ = Σ[1:c]
    V = V[:, 1:c]
    U, Σ, V
end

function _truncate_qr(F::CuQR, Dcut::Int)
    Q, R = F
    Q = CuMatrix(Q)
    c = min(Dcut, size(Q, 2))
    Q = Q[:, 1:c]
    R = R[1:c, :]
    _qr_fix!(Q, R)
end

LinearAlgebra.svd(M::CuMatrix, Dcut::Int, args...) = _truncate_svd(CUDA.svd(M), Dcut)


####################################
#
# Extension of SLiDE model to include
#       dynamics and perfect foresight
#
####################################

using SLiDE
using CSV
using JuMP
using Complementarity
using DataFrames

############
# LOAD DATA
############

#SLiDE data needs to be built or point to pre-existing build directory
#can pass a name (d, set) = build_data("name_of_build_directory")
!(@isdefined(d_in) && @isdefined(set_in)) && ((d_in, set_in) = build("state_model"))
d = copy(d_in)
set = copy(set_in)

#Define range of years in time horizon
# !!!! or load from CSV or YAML
bmkyr = 2016
endyr = 2018
years = bmkyr:endyr
#years = [2017, 2016, 2019, 2018]

#Define years and associated flags
years = ensurearray(sort(years))
yrl = years[length(years)]
yrf = years[1]
islast = Dict(years[k] => (years[k] == yrl ? 1 : 0) for k in keys(years))
isfirst = Dict(years[k] => (years[k] == yrf ? 1 : 0) for k in keys(years))
yrint = Dict(years[k+1] => years[k+1]-years[k] for k in 1:(length(years)-1))

#Load slide data and time horizon to produce model data and appropriate time-indexed subsets
(sld, set, idx) = model_input(years, d, set)

########## Model ##########
cge = MCPModel();

##############
# PARAMETERS
##############

#benchmark values
#benchmark values
@NLparameter(cge, ys0[r in set[:r], s in set[:s], g in set[:g]] == sld[:ys0][r,s,g]); # Sectoral supply
@NLparameter(cge, id0[r in set[:r], s in set[:s], g in set[:g]] == sld[:id0][r,s,g]); # Intermediate demand
@NLparameter(cge, ld0[r in set[:r], s in set[:s]] == sld[:ld0][r,s]); # Labor demand
@NLparameter(cge, kd0[r in set[:r], s in set[:s]] == sld[:kd0][r,s]); # Capital demand
@NLparameter(cge, ty0[r in set[:r], s in set[:s]] == sld[:ty0][r,s]); # Production tax (benchmark)
@NLparameter(cge, m0[r in set[:r], g in set[:g]] == sld[:m0][r,g]); # Imports
@NLparameter(cge, x0[r in set[:r], g in set[:g]] == sld[:x0][r,g]); # Exports of goods and services
@NLparameter(cge, rx0[r in set[:r], g in set[:g]] == sld[:rx0][r,g]); # Re-exports of goods and services
@NLparameter(cge, md0[r in set[:r], m in set[:m], g in set[:g]] == sld[:md0][r,m,g]); # Total margin demand
@NLparameter(cge, nm0[r in set[:r], g in set[:g], m in set[:m]] == sld[:nm0][r,g,m]); # Margin demand from national market
@NLparameter(cge, dm0[r in set[:r], g in set[:g], m in set[:m]] == sld[:dm0][r,g,m]); # Margin supply from local market
@NLparameter(cge, s0[r in set[:r], g in set[:g]] == sld[:s0][r,g]); # Aggregate supply
@NLparameter(cge, a0[r in set[:r], g in set[:g]] == sld[:a0][r,g]); # Armington supply
@NLparameter(cge, ta0[r in set[:r], g in set[:g]] == sld[:ta0][r,g]); # Tax net of subsidy rate on intermediate demand
@NLparameter(cge, tm0[r in set[:r], g in set[:g]] == sld[:tm0][r,g]); # Import tariff
@NLparameter(cge, cd0[r in set[:r], g in set[:g]] == sld[:cd0][r,g]); # Final demand
@NLparameter(cge, c0[r in set[:r]] == sld[:c0][r]); # Aggregate final demand
@NLparameter(cge, yh0[r in set[:r], g in set[:g]] == sld[:yh0][r,g]); #Household production
@NLparameter(cge, bopdef0[r in set[:r]] == sld[:bopdef0][r]); #Balance of payments
@NLparameter(cge, hhadj[r in set[:r]] == sld[:hhadj][r]); # Household adjustment
@NLparameter(cge, g0[r in set[:r], g in set[:g]] == sld[:g0][r,g]); # Government demand
@NLparameter(cge, i0[r in set[:r], g in set[:g]] == sld[:i0][r,g]); # Investment demand
@NLparameter(cge, xn0[r in set[:r], g in set[:g]] == sld[:xn0][r,g]); # Regional supply to national market
@NLparameter(cge, xd0[r in set[:r], g in set[:g]] == sld[:xd0][r,g]); # Regional supply to local market
@NLparameter(cge, dd0[r in set[:r], g in set[:g]] == sld[:dd0][r,g]); # Regional demand to local market
@NLparameter(cge, nd0[r in set[:r], g in set[:g]] == sld[:nd0][r,g]); # Regional demand to national market

#counterfactual taxes
@NLparameter(cge, ty[r in set[:r], s in set[:s]] == sld[:ty0][r,s]); # Output tax
@NLparameter(cge, ta[r in set[:r], g in set[:g]] == sld[:ta0][r,g]); # Armington tax
@NLparameter(cge, tm[r in set[:r], g in set[:g]] == sld[:tm0][r,g]); # Import tariff

# benchmark value share parameters
@NLparameter(cge, alpha_kl[r in set[:r], s in set[:s]] == ensurefinite(value(ld0[r,s]) / (value(ld0[r,s]) + value(kd0[r,s]))));
@NLparameter(cge, alpha_x[r in set[:r], g in set[:g]] == ensurefinite((value(x0[r,g]) - value(rx0[r,g])) / value(s0[r,g])));
@NLparameter(cge, alpha_d[r in set[:r], g in set[:g]] == ensurefinite(value(xd0[r,g]) / value(s0[r,g])));
@NLparameter(cge, alpha_n[r in set[:r], g in set[:g]] == ensurefinite(value(xn0[r,g]) / value(s0[r,g])));
@NLparameter(cge, theta_n[r in set[:r], g in set[:g]] == ensurefinite(value(nd0[r,g]) / (value(nd0[r,g]) - value(dd0[r,g]))));
@NLparameter(cge, theta_m[r in set[:r], g in set[:g]] == ensurefinite((1+value(tm0[r,g])) * value(m0[r,g]) / (value(nd0[r,g]) + value(dd0[r,g]) + (1 + value(tm0[r,g])) * value(m0[r,g]))));

#Substitution and transformation elasticities
# !!!! Currently assigned to global variable in definitions.jl
@NLparameter(cge, es_va[r in set[:r], s in set[:s]] == SUB_ELAST[:va]); # value-added nest - substitution elasticity
@NLparameter(cge, es_y[r in set[:r], s in set[:s]]  == SUB_ELAST[:y]); # Top-level Y nest (VA,M) - substitution elasticity
@NLparameter(cge, es_m[r in set[:r], s in set[:s]]  == SUB_ELAST[:m]); # Materials nest - substitution elasticity
@NLparameter(cge, et_x[r in set[:r], g in set[:g]]    == TRANS_ELAST[:x]); # Disposition, distribute regional supply to local, national, export - transformation elasticity
@NLparameter(cge, es_a[r in set[:r], g in set[:g]]    == SUB_ELAST[:a]); # Top-level A nest for aggregate demand (set[:m], set[:g]) - substitution elasticity
@NLparameter(cge, es_mar[r in set[:r], g in set[:g]]  == SUB_ELAST[:mar]); # Margin supply - substitution elasticity
@NLparameter(cge, es_d[r in set[:r], g in set[:g]]    == SUB_ELAST[:d]); # Domestic demand aggregation nest (intranational) - substitution elasticity
@NLparameter(cge, es_f[r in set[:r], g in set[:g]]    == SUB_ELAST[:f]); # Domestic and foreign demand aggregation nest (international) - substitution elasticity

# Temporal/Dynamic modifications
@NLparameter(cge, ir == 0.05); # Interest rate
@NLparameter(cge, gr == 0.01); # Growth rate
@NLparameter(cge, dr == 0.02); # Depreciation rate

@NLparameter(cge, pvm[yr in years] == (1/(1+value(ir)))^(yr-yrf)); # Reference price path - Present value multiplier
@NLparameter(cge, qvm[yr in years] == (1 + value(gr))^(yr-yrf)); # Reference quantity path multiplier

@NLparameter(cge, rk0 == value(ir) + value(dr)); # Initial return to capital

# !!!! Not used, will likely be used in updated model
# and intertemporal choice
# Initial capital earnings vk0
@NLparameter(cge, vk0[r in set[:r], s in set[:s]] == value(kd0[r,s]));
# Initial capital stock k0
@NLparameter(cge, k0[r in set[:r], s in set[:s]] == value(vk0[r,s])/value(rk0));
# Initial investment inv0
@NLparameter(cge, inv0[r in set[:r], s in set[:s]] == (value(dr)+value(gr))*value(k0[r,s]));
# !!!!

##################
# -- VARIABLES --
##################

# Set lower bound - small value that acts as a lower limit to variable values
# default is zero
lo = MODEL_LOWER_BOUND

# !!!! Define variables separately or in blocks
# sectors
@variable(cge, Y[(yr,r,s) in set[:Y]] >= lo, start = value(qvm[yr])); # Output
@variable(cge, X[(yr,r,g) in set[:X]] >= lo, start = value(qvm[yr])); # Exports
@variable(cge, A[(yr,r,g) in set[:A]] >= lo, start = value(qvm[yr])); # Armington
@variable(cge, C[yr in years, r in set[:r]] >= lo, start = value(qvm[yr])); # Consumption
@variable(cge, MS[yr in years, r in set[:r], m in set[:m]] >= lo, start = value(qvm[yr])); # Margin supply

@variable(cge, K[(yr,r,s) in set[:PK]] >= lo, start = value(qvm[yr])*value(kd0[r,s])); # Capital
@variable(cge, I[(yr,r,s) in set[:PK]] >= lo, start = value(qvm[yr])*(value(dr)+value(gr))*value(kd0[r,s])); # Investment
@variable(cge, TK[(r,s) in set[:PKT]] >= lo, start = value(kd0[r,s]) * value(qvm[yrl])*(1+value(gr))); # Terminal Capital

# !!!! Future updates to capital stock/investment
# @variable(cge, K[(yr,r,s) in set[:PK]] >= lo, start = value(qvm[yr])*value(k0[r,s])); # Capital
# @variable(cge, I[(yr,r,s) in set[:PK]] >= lo, start = value(qvm[yr])*value(inv0[r,s])); # Investment
# @variable(cge, TK[(r,s) in set[:PKT]] >= lo, start = value(k0[r,s]) * value(qvm[yr]) * (1 + value(gr))); # Terminal Capital

# commodities
@variable(cge, PA[(yr,r,g) in set[:PA]] >= lo, start = value(pvm[yr])); # Regional market (input)
@variable(cge, PY[(yr,r,g) in set[:PY]] >= lo, start = value(pvm[yr])); # Regional market (output)
@variable(cge, PD[(yr,r,g) in set[:PD]] >= lo, start = value(pvm[yr])); # Local market price
@variable(cge, PN[yr in years, g in set[:g]] >= lo, start = value(pvm[yr])); # National market
@variable(cge, PL[yr in years, r in set[:r]] >= lo, start = value(pvm[yr])); # Wage rate
@variable(cge, PM[yr in years, r in set[:r], m in set[:m]] >= lo, start = value(pvm[yr])); # Margin price
@variable(cge, PC[yr in years, r in set[:r]] >= lo, start = value(pvm[yr])); # Consumer price index
@variable(cge, PFX[yr in years] >= lo, start = value(pvm[yr])); # Foreign exchange

@variable(cge, PK[(yr,r,s) in set[:PK]] >= lo, start = value(pvm[yr]) * (1 + value(ir))); # Price of capital
@variable(cge, RK[(yr,r,s) in set[:PK]] >= lo, start = value(pvm[yr])*value(rk0)); # Capital rental rate
@variable(cge, PKT[(r,s) in set[:PKT]] >= lo, start = value(pvm[yrl])); # Terminal capital price
#@variable(cge, PKT[(r,s) in set[:PKT]] >= lo, start = start_value(PK[yrl,r,s])/(1+value(ir)); # Terminal capital price

# consumer
# !!!! Updates needed for Intertemporal Consumption
@variable(cge,RA[yr in years, r in set[:r]]>=lo,start = value(qvm[yr])*value(pvm[yr])*value(c0[r])); #Representative Agent

###############################
# -- PLACEHOLDER VARIABLES --
###############################

#cobb-douglas function for value added (VA)
@NLexpression(cge,CVA[yr in years,r in set[:r],s in set[:s]],
  PL[yr,r]^alpha_kl[r,s] * ((haskey(RK.lookup[1], (yr,r,s)) ? RK[(yr,r,s)] : 1.0)/rk0)^ (1-alpha_kl[r,s]) );

#demand for labor in VA
@NLexpression(cge,AL[yr in years, r in set[:r], s in set[:s]], ld0[r,s] * CVA[yr,r,s] / PL[yr,r] );

#demand for capital in VA
@NLexpression(cge,AK[yr in years, r in set[:r],s in set[:s]],
  kd0[r,s] * CVA[yr,r,s] / ((haskey(RK.lookup[1], (yr,r,s)) ? RK[(yr,r,s)] : 1.0)/rk0));

#CES function for output demand - including
# exports (absent of re-exports) times the price for foreign exchange,
# region's supply to national market times the national market price
# regional supply to local market times domestic price

@NLexpression(cge,RX[yr in years,r in set[:r],g in set[:g]],
  (alpha_x[r,g]*PFX[yr]^(1 + et_x[r,g])+alpha_n[r,g]*PN[yr,g]^(1 + et_x[r,g])+alpha_d[r,g]*(haskey(PD.lookup[1], (yr,r,g)) ? PD[(yr,r,g)] : 1.0)^(1 + et_x[r,g]))^(1/(1 + et_x[r,g])) );

#demand for exports via demand function
@NLexpression(cge,AX[yr in years,r in set[:r],g in set[:g]], (x0[r,g] - rx0[r,g])*(PFX[yr]/RX[yr,r,g])^et_x[r,g] );

#demand for contribution to national market
@NLexpression(cge,AN[yr in years,r in set[:r],g in set[:g]], xn0[r,g]*(PN[yr,g]/(RX[yr,r,g]))^et_x[r,g] );

#demand for regionals supply to local market
@NLexpression(cge,AD[yr in years,r in set[:r],g in set[:g]],
  xd0[r,g] * ((haskey(PD.lookup[1], (yr,r,g)) ? PD[(yr,r,g)] : 1.0) / (RX[yr,r,g]))^et_x[r,g] );

# CES function for tradeoff between national and domestic market
@NLexpression(cge,CDN[yr in years,r in set[:r],g in set[:g]],
(theta_n[r,g]*PN[yr,g]^(1-es_d[r,g])+(1-theta_n[r,g])*(haskey(PD.lookup[1], (yr,r,g)) ? PD[(yr,r,g)] : 1.0)^(1-es_d[r,g]))^(1/(1-es_d[r,g])) );

# CES function for tradeoff between domestic consumption and foreign exports
# recall tm in the import tariff thus tm / tm0 is the relative change in import tariff rates
@NLexpression(cge,CDM[yr in years,r in set[:r],g in set[:g]],
((1-theta_m[r,g])*CDN[yr,r,g]^(1-es_f[r,g])+theta_m[r,g]*(PFX[yr]*(1+tm[r,g])/(1+tm0[r,g]))^(1-es_f[r,g]))^(1/(1-es_f[r,g])) );

# set[:r] demand from the national market <- note nesting of CDN in CDM
@NLexpression(cge,DN[yr in years,r in set[:r],g in set[:g]],
  nd0[r,g]*(CDN[yr,r,g]/PN[yr,g])^es_d[r,g]*(CDM[yr,r,g]/CDN[yr,r,g])^es_f[r,g] );

# region demand from local market <- note nesting of CDN in CDM
@NLexpression(cge,DD[yr in years,r in set[:r],g in set[:g]],
  dd0[r,g]*(CDN[yr,r,g]/(haskey(PD.lookup[1], (yr,r,g)) ? PD[(yr,r,g)] : 1.0))^es_d[r,g]*(CDM[yr,r,g]/CDN[yr,r,g])^es_f[r,g] );

# import demand
@NLexpression(cge,MD[yr in years,r in set[:r],g in set[:g]],
  m0[r,g]*(CDM[yr,r,g]*(1+tm[r,g])/(PFX[yr]*(1+tm0[r,g])))^es_f[r,g] );

# final demand
@NLexpression(cge,CD[yr in years,r in set[:r],g in set[:g]],
  cd0[r,g]*PC[yr,r] / (haskey(PA.lookup[1], (yr,r,g)) ? PA[(yr,r,g)] : 1.0) );

###############################
# -- Zero Profit Conditions --
###############################

@mapping(cge,profit_y[(yr,r,s) in set[:Y]],
# cost of intermediate demand
    sum((haskey(PA.lookup[1], (yr,r,g)) ? PA[(yr,r,g)] : 1.0) * id0[r,g,s] for g in set[:g])
# cost of labor inputs
    + PL[yr,r] * AL[yr,r,s]
# cost of capital inputs
    + ((haskey(RK.lookup[1], (yr,r,s)) ? RK[(yr,r,s)] : 1.0)/rk0)* AK[yr,r,s]
    -
# revenue from sectoral supply (take note of r/s/g indices on ys0)
    sum((haskey(PY.lookup[1], (yr,r,g)) ? PY[(yr,r,g)] : 1.0)  * ys0[r,s,g] for g in set[:g]) * (1-ty[r,s])
);

@mapping(cge,profit_x[(yr,r,g) in set[:X]],
# output 'cost' from aggregate supply
    (haskey(PY.lookup[1], (yr,r,g)) ? PY[(yr,r,g)] : 1.0) * s0[r,g]
    - (
# revenues from foreign exchange
        PFX[yr] * AX[yr,r,g]
# revenues from national market
        + PN[yr,g] * AN[yr,r,g]
# revenues from domestic market
        + (haskey(PD.lookup[1], (yr,r,g)) ? PD[(yr,r,g)] : 1.0) * AD[yr,r,g]
    )
);

@mapping(cge,profit_a[(yr,r,g) in set[:A]],
# costs from national market
    PN[yr,g] * DN[yr,r,g]
# costs from domestic market
    + (haskey(PD.lookup[1], (yr,r,g)) ? PD[(yr,r,g)] : 1.0) * DD[yr,r,g]
# costs from imports, including import tariff
    + PFX[yr] * (1+tm[r,g]) * MD[yr,r,g]
# costs of margin demand
    + sum(PM[yr,r,m] * md0[r,m,g] for m in set[:m])
    - (
# revenues from regional market based on armington supply
        (haskey(PA.lookup[1], (yr,r,g)) ? PA[(yr,r,g)] : 1.0) * (1-ta[r,g]) * a0[r,g]
# revenues from re-exports
        + PFX[yr] * rx0[r,g]
    )
);

@mapping(cge, profit_c[yr in years,r in set[:r]],
# costs of inputs - computed as final demand times regional market prices
    sum((haskey(PA.lookup[1], (yr,r,g)) ? PA[(yr,r,g)] : 1.0) * CD[yr,r,g] for g in set[:g])
    -
# revenues/benefit computed as CPI * reference consumption
    PC[yr,r] * c0[r]
);

@mapping(cge,profit_k[(yr,r,s) in set[:PK]],
    (haskey(PK.lookup[1], (yr,r,s)) ? PK[(yr,r,s)] : 1.0)
    - (
        (haskey(RK.lookup[1], (yr,r,s)) ? RK[(yr,r,s)] : 1.0)
        + (1-dr) * (yr!=yrl ? (haskey(PK.lookup[1], (yr+1,r,s)) ? PK[(yr+1,r,s)] : 1.0) : (haskey(PKT.lookup[1], (r,s)) ? PKT[(r,s)] : 1.0))
    )
);

@mapping(cge,profit_i[(yr,r,s) in set[:PK]],
    (haskey(PY.lookup[1], (yr,r,s)) ? PY[(yr,r,s)] : 1.0)
    -
    (yr!=yrl ? (haskey(PK.lookup[1], (yr+1,r,s)) ? PK[(yr+1,r,s)] : 1.0) : (haskey(PKT.lookup[1], (r,s)) ? PKT[(r,s)] : 1.0))
);

@mapping(cge,profit_ms[yr in years,r in set[:r],m in set[:m]],
# provision of set[:m] to national market
    sum(PN[yr,gm]   * nm0[r,gm,m] for gm in set[:gm])
# provision of set[:m] to domestic market
    + sum((haskey(PD.lookup[1], (yr,r,gm)) ? PD[(yr,r,gm)] : 1.0) * dm0[r,gm,m] for gm in set[:gm])
    -
# total margin demand
    PM[yr,r,m] * sum(md0[r,m,gm] for gm in set[:gm])
);

###################################
# -- Market Clearing Conditions --
###################################

@mapping(cge,market_pa[(yr,r,g) in set[:PA]],
# absorption or supply
    (haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * a0[r,g]
    - (
# government demand (exogenous)
        g0[r,g]*qvm[yr]
# demand for investment (exogenous)
        + i0[r,g]*qvm[yr]
# final demand
        + C[yr,r] * CD[yr,r,g]
# intermediate demand
        + sum((haskey(Y.lookup[1], (yr,r,s)) ? Y[(yr,r,s)] : 1.0) * id0[r,g,s] for s in set[:s] if (yr,r,s) in set[:Y])
    )
)

@mapping(cge,market_py[(yr,r,g) in set[:PY]],
# sectoral supply
    sum((haskey(Y.lookup[1], (yr,r,s)) ? Y[(yr,r,s)] : 1.0) *ys0[r,s,g] for s in set[:s])
# household production (exogenous)
    + yh0[r,g]*qvm[yr]
    -
# aggregate supply (akin to market demand)
    (haskey(X.lookup[1], (yr,r,g)) ? X[(yr,r,g)] : 1.0) * s0[r,g]
);

@mapping(cge,market_pd[(yr,r,g) in set[:PD]],
# aggregate supply
    (haskey(X.lookup[1], (yr,r,g)) ? X[(yr,r,g)] : 1.0)  * AD[yr,r,g]
    - (
# demand for local market
        (haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * DD[yr,r,g]
# margin supply from local market
        + sum(MS[yr,r,m] * dm0[r,g,m] for m in set[:m] if (g in set[:gm] ) )
    )
);

@mapping(cge,market_pn[yr in years,g in set[:g]],
# supply to the national market
    sum((haskey(X.lookup[1], (yr,r,g)) ? X[(yr,r,g)] : 1.0)  * AN[yr,r,g] for r in set[:r])
    - (
# demand from the national market
        sum((haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * DN[yr,r,g] for r in set[:r])
# market supply to the national market
        + sum(MS[yr,r,m] * nm0[r,g,m] for r in set[:r] for m in set[:m] if (g in set[:gm]) )
    )
);

@mapping(cge,market_pl[yr in years,r in set[:r]],
# supply of labor
    sum(ld0[r,s]*qvm[yr] for s in set[:s])
    -
# demand for labor in all set[:s]
    sum((haskey(Y.lookup[1], (yr,r,s)) ? Y[(yr,r,s)] : 1.0) * AL[yr,r,s] for s in set[:s])
);

@mapping(cge,market_pk[(yr,r,s) in set[:PK]],
# if first year, initial capital
# else investment plus previous year's decayed capital
    (yr==yrf ? kd0[r,s] : (haskey(I.lookup[1], (yr-1,r,s)) ? I[(yr-1,r,s)] : 0.0))
    +(1-dr) * (yr>yrf ? (haskey(K.lookup[1], (yr-1,r,s)) ? K[(yr-1,r,s)] : 0.0) : 0.0)
    -
# capital in year [yr]
    (haskey(K.lookup[1], (yr,r,s)) ? K[(yr,r,s)] : 0.0)
);

@mapping(cge,market_rk[(yr,r,s) in set[:PK]],
    (haskey(K.lookup[1], (yr,r,s)) ? K[(yr,r,s)] : 0.0)
    -
    (haskey(Y.lookup[1], (yr,r,s)) ? Y[(yr,r,s)] : 1.0) * AK[yr,r,s]
);

@mapping(cge,market_pkt[(r,s) in set[:PKT]],
    (1-dr)*(haskey(K.lookup[1], (yrl,r,s)) ? K[(yrl,r,s)] : 0.0)
    + (haskey(I.lookup[1], (yrl,r,s)) ? I[(yrl,r,s)] : 0.0)
    -
    (haskey(TK.lookup[1], (r,s)) ? TK[(r,s)] : 0.0)
);

@mapping(cge,termk[(r,s) in set[:PKT]],
    (haskey(I.lookup[1], (yrl,r,s)) ? I[(yrl,r,s)] : 0.0)
    / ((haskey(I.lookup[1], (yrl-1,r,s)) ? I[(yrl-1,r,s)] : 0.0))
    -
    (haskey(Y.lookup[1], (yrl,r,s)) ? Y[(yrl,r,s)] : 1.0)
    / (haskey(Y.lookup[1], (yrl-1,r,s)) ? Y[(yrl-1,r,s)] : 1.0)
);

@mapping(cge,market_pm[yr in years,r in set[:r], m in set[:m]],
# margin supply
    MS[yr,r,m] * sum(md0[r,m,gm] for gm in set[:gm])
    -
# margin demand
    sum((haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * md0[r,m,g] for g in set[:g])
);

@mapping(cge,market_pfx[yr in years],
# balance of payments (exogenous)
    sum(bopdef0[r] for r in set[:r])*qvm[yr]
# supply of exports
    + sum((haskey(X.lookup[1], (yr,r,g)) ? X[(yr,r,g)] : 1.0)  * AX[yr,r,g] for r in set[:r] for g in set[:g])
# supply of re-exports
    + sum((haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * rx0[r,g] for r in set[:r] for g in set[:g] if (yr,r,g) in set[:A])
    -
# import demand
    sum((haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * MD[yr,r,g] for r in set[:r] for g in set[:g] if (yr,r,g) in set[:A])
);

@mapping(cge,market_pc[yr in years,r in set[:r]],
# a period's final demand
    C[yr,r] * c0[r]
    -
# consumption / utiltiy
    RA[yr,r] / PC[yr,r]
);

# !!!! Updates needed for Intertemporal Consumption
@mapping(cge,income_ra[yr in years,r in set[:r]],
# consumption/utility
    RA[yr,r]
    -
    (
# labor income
        PL[yr,r] * sum(ld0[r,s] for s in set[:s])*qvm[yr]
# provision of household supply
        + sum( (haskey(PY.lookup[1], (yr,r,g)) ? PY[(yr,r,g)] : 1.0) * yh0[r,g]*qvm[yr] for g in set[:g])
# revenue or costs of foreign exchange including household adjustment
        + PFX[yr] * (bopdef0[r] + hhadj[r])*qvm[yr]
# government and investment provision
        - sum((haskey(PA.lookup[1], (yr,r,g)) ? PA[(yr,r,g)] : 1.0) * (g0[r,g] + i0[r,g])*qvm[yr] for g in set[:g])
# import taxes - assumes lumpsum recycling
        + sum((haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * MD[yr,r,g] * PFX[yr] * tm[r,g] for g in set[:g] if (yr,r,g) in set[:A])
# taxes on intermediate demand - assumes lumpsum recycling
        + sum((haskey(A.lookup[1], (yr,r,g)) ? A[(yr,r,g)] : 1.0) * a0[r,g]*(haskey(PA.lookup[1], (yr,r,g)) ? PA[(yr,r,g)] : 1.0)*ta[r,g] for g in set[:g] if (yr,r,g) in set[:A])
# production taxes - assumes lumpsum recycling
        + sum( pvm[yr]*(haskey(Y.lookup[1], (yr,r,s)) ? Y[(yr,r,s)] : 1.0) * ys0[r,s,g] * ty[r,s] for s in set[:s], g in set[:g])
# capital income
        + (1-islast[yr]) * sum((haskey(PK.lookup[1], (yr,r,s)) ? PK[(yr,r,s)] : 1.0) * (haskey(K.lookup[1], (yr,r,s)) ? K[(yr,r,s)] : 0.0) for s in set[:s]) / (1+ir)
        + (islast[yr]) * sum((haskey(PKT.lookup[1], (r,s)) ? PKT[(r,s)] : 1.0) * (haskey(TK.lookup[1], (r,s)) ? TK[(r,s)] : 0.0) for s in set[:s]) / (1+gr)
    )
);

####################################
# -- Complementarity Conditions --
####################################

# define complementarity conditions
# note the pattern of ZPC -> primal variable  &  MCC -> dual variable (price)
@complementarity(cge,profit_y,Y);
@complementarity(cge,profit_x,X);
@complementarity(cge,profit_a,A);
@complementarity(cge,profit_c,C);
@complementarity(cge,profit_ms,MS);
@complementarity(cge,market_pa,PA);
@complementarity(cge,market_py,PY);
@complementarity(cge,market_pd,PD);
@complementarity(cge,market_pn,PN);
@complementarity(cge,market_pl,PL);
@complementarity(cge,market_pm,PM);
@complementarity(cge,market_pc,PC);
@complementarity(cge,market_pfx,PFX);
@complementarity(cge,income_ra,RA);
@complementarity(cge,market_pk,PK);
@complementarity(cge,market_rk,RK);
@complementarity(cge,market_pkt,PKT)
@complementarity(cge,termk,TK)
@complementarity(cge,profit_k,K)
@complementarity(cge,profit_i,I)

####################
# -- Model Solve --
####################

#set up the options for the path solver
#PATHSolver.options(convergence_tolerance=1e-6, output=:yes, time_limit=3600)
PATHSolver.options(convergence_tolerance=1e-6, output=:yes, time_limit=3600, cumulative_iteration_limit=0)

# solve the model
status = solveMCP(cge)


## Automatically generated using the following and then hand-edited to
## remove duplicates and coerce constants to the appropriate type.

## using Clang.wrap_c
## context = wrap_c.init(; output_file="superlu_h.jl",
##                       header_library=x->"/usr/include/superlu",
##                       header_wrapped=(x,y)->contains(y, "superlu"),
##                       clang_diagnostics=true,
##                       clang_args=["-v"])
## context.options.wrap_structs = true
## wrap_c.wrap_c_headers(context, ["/usr/include/superlu/slu_ddefs.h"])

const EMPTY = int32(-1)
const FALSE = int32(0)
const TRUE = int32(1)
const NO_MEMTYPE = int32(4)

const NODROP = (0x0000)
const DROP_BASIC = (0x0001)
const DROP_PROWS = (0x0002)
const DROP_COLUMN = (0x0004)
const DROP_AREA = (0x0008)
const DROP_SECONDARY = (0x000E)
const DROP_DYNAMIC = (0x0010)
const DROP_INTERP = (0x0100)
const MILU_ALPHA = (1.0e-2)

typealias int_t Cint

# begin enum Stype_t
typealias Stype_t Cint   # technically, Uint32 but printing a Cint is less verbose
const SLU_NC = int32(0)
const SLU_NCP = int32(1)
const SLU_NR = int32(2)
const SLU_SC = int32(3)
const SLU_SCP = int32(4)
const SLU_SR = int32(5)
const SLU_DN = int32(6)
const SLU_NR_loc = int32(7)
# end enum Stype_t

# begin enum Dtype_t
typealias Dtype_t Cint
const SLU_S = int32(0)
const SLU_D = int32(1)
const SLU_C = int32(2)
const SLU_Z = int32(3)
# end enum Dtype_t

# begin enum Mtype_t
typealias Mtype_t Cint
const SLU_GE = int32(0)
const SLU_TRLU = int32(1)
const SLU_TRUU = int32(2)
const SLU_TRL = int32(3)
const SLU_TRU = int32(4)
const SLU_SYL = int32(5)
const SLU_SYU = int32(6)
const SLU_HEL = int32(7)
const SLU_HEU = int32(8)
# end enum Mtype_t

type SuperMatrix
    Stype::Stype_t
    Dtype::Dtype_t
    Mtype::Mtype_t
    nrow::int_t
    ncol::int_t
    Store::Ptr{None}
end

type NCformat
    nnz::int_t
    nzval::Ptr{None}
    rowind::Ptr{int_t}
    colptr::Ptr{int_t}
end

type NRformat
    nnz::int_t
    nzval::Ptr{None}
    colind::Ptr{int_t}
    rowptr::Ptr{int_t}
end

type SCformat
    nnz::int_t
    nsuper::int_t
    nzval::Ptr{None}
    nzval_colptr::Ptr{int_t}
    rowind::Ptr{int_t}
    rowind_colptr::Ptr{int_t}
    col_to_sup::Ptr{int_t}
    sup_to_col::Ptr{int_t}
end
type SCPformat
    nnz::int_t
    nsuper::int_t
    nzval::Ptr{None}
    nzval_colbeg::Ptr{int_t}
    nzval_colend::Ptr{int_t}
    rowind::Ptr{int_t}
    rowind_colbeg::Ptr{int_t}
    rowind_colend::Ptr{int_t}
    col_to_sup::Ptr{int_t}
    sup_to_colbeg::Ptr{int_t}
    sup_to_colend::Ptr{int_t}
end
type NCPformat
    nnz::int_t
    nzval::Ptr{None}
    rowind::Ptr{int_t}
    colbeg::Ptr{int_t}
    colend::Ptr{int_t}
end
type DNformat
    lda::int_t
    nzval::Ptr{None}
end
type NRformat_loc
    nnz_loc::int_t
    m_loc::int_t
    fst_row::int_t
    nzval::Ptr{None}
    rowptr::Ptr{int_t}
    colind::Ptr{int_t}
end

# begin enum yes_no_t
typealias yes_no_t Cint
const NO = int32(0)
const YES = int32(1)
# end enum yes_no_t

# begin enum fact_t
typealias fact_t Cint
const DOFACT = int32(0)
const SamePattern = int32(1)
const SamePattern_SameRowPerm = int32(2)
const FACTORED = int32(3)
# end enum fact_t

# begin enum rowperm_t
typealias rowperm_t Cint
const NOROWPERM = int32(0)
const LargeDiag = int32(1)
const MY_PERMR = int32(2)
# end enum rowperm_t

# begin enum colperm_t
typealias colperm_t Cint
const NATURAL = int32(0)
const MMD_ATA = int32(1)
const MMD_AT_PLUS_A = int32(2)
const COLAMD = int32(3)
const METIS_AT_PLUS_A = int32(4)
const PARMETIS = int32(5)
const ZOLTAN = int32(6)
const MY_PERMC = int32(7)
# end enum colperm_t

# begin enum trans_t
typealias trans_t Cint
const NOTRANS = int32(0)
const TRANS = int32(1)
const CONJ = int32(2)
# end enum trans_t

# begin enum DiagScale_t
typealias DiagScale_t Cint
const NOEQUIL = int32(0)
const ROW = int32(1)
const COL = int32(2)
const BOTH = int32(3)
# end enum DiagScale_t

# begin enum IterRefine_t
typealias IterRefine_t Cint
const NOREFINE = int32(0)
const SLU_SINGLE = int32(1)
const SLU_DOUBLE = int32(2)
const SLU_EXTRA = int32(3)
# end enum IterRefine_t

# begin enum MemType
typealias MemType Cint
const LUSUP = int32(0)
const UCOL = int32(1)
const LSUB = int32(2)
const USUB = int32(3)
const LLVL = int32(4)
const ULVL = int32(5)
# end enum MemType

# begin enum stack_end_t
typealias stack_end_t Cint
const HEAD = int32(0)
const TAIL = int32(1)
# end enum stack_end_t

# begin enum LU_space_t
typealias LU_space_t Cint
const SYSTEM = int32(0)
const USER = int32(1)
# end enum LU_space_t

# begin enum norm_t
typealias norm_t Cint
const ONE_NORM = int32(0)
const TWO_NORM = int32(1)
const INF_NORM = int32(2)
# end enum norm_t

# begin enum milu_t
typealias milu_t Cint
const SILU = int32(0)
const SMILU_1 = int32(1)
const SMILU_2 = int32(2)
const SMILU_3 = int32(3)
# end enum milu_t

# begin enum PhaseType - these are names of 1-based positions in SuperLUStat_t vectors.
typealias PhaseType Int
const COLPERM = 1
const ROWPERM = 2
const RELAX = 3
const ETREE = 4
const EQUIL = 5
const SYMBFAC = 6
const DIST = 7
const FACT = 8
const COMM = 9
const SOL_COMM = 10
const RCOND = 11
const SOLVE = 12
const REFINE = 13
const TRSV = 14
const GEMV = 15
const FERR = 16
const NPHASES = 16 # number of phases in SuperLUStat_t vectors.
# end enum PhaseType

typealias flops_t Cfloat
typealias Logical Cuchar
type superlu_options_t
    Fact::fact_t
    Equil::yes_no_t
    ColPerm::colperm_t
    Trans::trans_t
    IterRefine::IterRefine_t
    DiagPivotThresh::Cdouble
    SymmetricMode::yes_no_t
    PivotGrowth::yes_no_t
    ConditionNumber::yes_no_t
    RowPerm::rowperm_t
    ILU_DropRule::Cint
    ILU_DropTol::Cdouble
    ILU_FillFactor::Cdouble
    ILU_Norm::norm_t
    ILU_FillTol::Cdouble
    ILU_MILU::milu_t
    ILU_MILU_Dim::Cdouble
    ParSymbFact::yes_no_t
    ReplaceTinyPivot::yes_no_t
    SolveInitialized::yes_no_t
    RefineInitialized::yes_no_t
    PrintStat::yes_no_t
    nnzL::Cint
    nnzU::Cint
    num_lookaheads::Cint
    lookahead_etree::yes_no_t
    SymPattern::yes_no_t
end
type e_node
    size::Cint
    mem::Ptr{None}
end
type ExpHeader
    size::Cint
    mem::Ptr{None}
end
type LU_stack_t
    size::Cint
    used::Cint
    top1::Cint
    top2::Cint
    array::Ptr{None}
end
type SuperLUStat_t
    panel_histo::Ptr{Cint}
    utime::Ptr{Cdouble}
    ops::Ptr{flops_t}
    TinyPivots::Cint
    RefineSteps::Cint
    expansions::Cint
end
type mem_usage_t
    for_lu::Cfloat
    total_needed::Cfloat
end
type GlobalLU_t
    xsup::Ptr{Cint}
    supno::Ptr{Cint}
    lsub::Ptr{Cint}
    xlsub::Ptr{Cint}
    lusup::Ptr{Cdouble}
    xlusup::Ptr{Cint}
    ucol::Ptr{Cdouble}
    usub::Ptr{Cint}
    xusub::Ptr{Cint}
    nzlmax::Cint
    nzumax::Cint
    nzlumax::Cint
    n::Cint
    MemModel::LU_space_t
    num_expansions::Cint
    expanders::Ptr{ExpHeader}
    stack::LU_stack_t
end

const i1 = one(Cint)


export time_modeling

# Setup time-domain linear or nonlinear foward and adjoint modeling and interface to OPESCI/devito
function time_modeling(model_full::Model, srcGeometry, srcData, recGeometry, recData, dm, srcnum::Int64, op::Char, mode::Int64, options)
    pm = load_pymodel()

    # Load full geometry for out-of-core geometry containers
    typeof(recGeometry) == GeometryOOC && (recGeometry = Geometry(recGeometry))
    typeof(srcGeometry) == GeometryOOC && (srcGeometry = Geometry(srcGeometry))
    length(model_full.n) == 3 ? dims = [3,2,1] : dims = [2,1]   # model dimensions for Python are (z,y,x) and (z,x)

    # limit model to area with sources/receivers
    if options.limit_m == true
        model = deepcopy(model_full)
        if op=='J' && mode==1
            model,dm = limit_model_to_receiver_area(srcGeometry,recGeometry,model,options.buffer_size;pert=dm)
        else
            model = limit_model_to_receiver_area(srcGeometry,recGeometry,model,options.buffer_size)
        end
    else
        model = model_full
    end

    # Set up Python model structure
    if op=='J' && mode == 1
        modelPy = pm."Model"(origin=(0.,0.,0.), spacing=model.d, shape=model.n, vp=process_physical_parameter(sqrt.(1f0./model.m), dims), nbpml=model.nb,
            rho=process_physical_parameter(model.rho, dims), dm=process_physical_parameter(reshape(dm,model.n), dims), space_order=options.space_order)
    else
        modelPy = pm."Model"(origin=(0.,0.,0.), spacing=model.d, shape=model.n, vp=process_physical_parameter(sqrt.(1f0./model.m), dims), nbpml=model.nb,
            rho=process_physical_parameter(model.rho, dims), space_order=options.space_order)
    end

    # Load shot record if stored on disk
    if recData != nothing
        if typeof(recData[1]) == SegyIO.SeisCon
            recDataCell = Array{Any}(undef, 1); recDataCell[1] = convert(Array{Float32,2},recData[1][1].data); recData = recDataCell
        elseif typeof(recData[1]) == String
            recData = load(recData[1])."d".data
        end
    end

    # Remove receivers outside the modeling domain (otherwise leads to segmentation faults)
    if mode==1 && recGeometry != nothing
        recGeometry = remove_out_of_bounds_receivers(recGeometry, model)
    elseif mode==-1 && recGeometry != nothing
        recGeometry, recData = remove_out_of_bounds_receivers(recGeometry, recData, model)
    end

    # Devito interface
    argout = devito_interface(modelPy, model.o, srcGeometry, srcData, recGeometry, recData, dm, options)

    # Extend gradient back to original model size
    if op=='J' && mode==-1 && options.limit_m==true
        argout = vec(extend_gradient(model_full, model, reshape(argout, model.n)))
    end

    return argout
end

# Function instance without options
time_modeling(model::Model, srcGeometry, srcData, recGeometry, recData, perturbation, srcnum::Int64, op::Char, mode::Int64) =
    time_modeling(model, srcGeometry, srcData, recGeometry, recData, perturbation, srcnum, op, mode, Options())

######################################################################################################################################################

# d_obs = Pr*F*Ps'*q
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Geometry, srcData::Array, recGeometry::Geometry, recData::Nothing, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    qIn = time_resample(srcData[1],srcGeometry,dtComp)[1]
    ntComp = size(qIn,1)
    ntRec = Int(trunc(recGeometry.t[1]/dtComp + 1))

    # Set up coordinates with devito dimensions
    src_coords = setup_grid(srcGeometry, modelPy.shape, origin)
    rec_coords = setup_grid(recGeometry, modelPy.shape, origin)

    # Devito call
    dOut = get(pycall(ac."forward_modeling", PyObject, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(qIn))), PyReverseDims(copy(transpose(rec_coords))),
                  space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface), 0)
    ntRec > ntComp && (dOut = [dOut zeros(size(dOut,1), ntRec - ntComp)])
    dOut = time_resample(dOut,dtComp,recGeometry)

    # Output shot record as judiVector
    if options.save_data_to_disk
        container = write_shot_record(srcGeometry,srcData,recGeometry,dOut,options)
        return judiVector(container)
    else
        return judiVector(recGeometry,dOut)
    end
end

# q_ad = Ps*F'*Pr'*d_obs
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Geometry, srcData::Nothing, recGeometry::Geometry, recData::Array, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    dIn = time_resample(recData[1],recGeometry,dtComp)[1]
    ntComp = size(dIn,1)
    ntSrc = Int(trunc(srcGeometry.t[1]/dtComp + 1))

    # Set up coordinates with devito dimensions
    src_coords = setup_grid(srcGeometry, modelPy.shape, origin)
    rec_coords = setup_grid(recGeometry, modelPy.shape, origin)

    # Devito call
    qOut = pycall(ac."adjoint_modeling", Array{Float32,2}, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(rec_coords))), PyReverseDims(copy(transpose(dIn))),
                  space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface)
    ntSrc > ntComp && (qOut = [qOut zeros(size(qOut), ntSrc - ntComp)])
    qOut = time_resample(qOut,dtComp,srcGeometry)

    # Output adjoint data as judiVector
    return judiVector(srcGeometry,qOut)
end

# u = F*Ps'*q
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Geometry, srcData::Array, recGeometry::Nothing, recData::Nothing, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    qIn = time_resample(srcData[1],srcGeometry,dtComp)[1]
    ntComp = size(qIn,1)

    # Set up coordinates with devito dimensions
    src_coords = setup_grid(srcGeometry, modelPy.shape, origin)

    # Devito call
    u = pycall(ac."forward_modeling", PyObject, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(qIn))), nothing,
               space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface)

    # Output forward wavefield as judiWavefield
    options.save_wavefield_to_disk && (u = dump_wavefield(u))
    return judiWavefield(Info(prod(modelPy.shape), 1, ntComp), dtComp, u)
end

# v = F'*Pr'*d_obs
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Nothing, srcData::Nothing, recGeometry::Geometry, recData::Array, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    dIn = time_resample(recData[1],recGeometry,dtComp)[1]
    ntComp = size(dIn,1)

    # Set up coordinates with devito dimensions
    rec_coords = setup_grid(recGeometry, modelPy.shape, origin)

    # Devito call
    v = pycall(ac."adjoint_modeling", PyObject, modelPy, nothing, PyReverseDims(copy(transpose(rec_coords))), PyReverseDims(copy(transpose(dIn))),
               space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface)

    # Output adjoint wavefield as judiWavefield
    options.save_wavefield_to_disk && (v = dump_wavefield(v))
    return judiWavefield(Info(prod(modelPy.shape), 1, ntComp), dtComp, v)
end

# d_obs = Pr*F*u
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Nothing, srcData::Array, recGeometry::Geometry, recData::Nothing, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    ntRec = Int(trunc(recGeometry.t[1]/dtComp + 1))

    # Set up coordinates with devito dimensions
    rec_coords = setup_grid(recGeometry, modelPy.shape, origin)

    # Devito call
    dOut = get(pycall(ac."forward_modeling", PyObject, modelPy, nothing, srcData[1], PyReverseDims(copy(transpose(rec_coords))),
                  space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface), 0)
    #ntRec > ntComp && (dOut = [dOut zeros(size(dOut,1), ntRec - ntComp)])
    dOut = time_resample(dOut,dtComp,recGeometry)

    # Output shot record as judiVector
    if options.save_data_to_disk
        throw("Writing shot record to SEG-Y file not supported for modeling with wavefield as right-hand-side.")
    else
        return judiVector(recGeometry,dOut)
    end
end

# q_ad = Ps*F'*v
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Geometry, srcData::Nothing, recGeometry::Nothing, recData::Array, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    ntSrc = Int(trunc(srcGeometry.t[1]/dtComp + 1))

    # Set up coordinates with devito dimensions
    src_coords = setup_grid(srcGeometry, modelPy.shape, origin)

    # Devito call
    qOut = pycall(ac."adjoint_modeling", Array{Float32,2}, modelPy, PyReverseDims(copy(transpose(src_coords))), nothing, recData[1],
                  space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface)
    #ntSrc > ntComp && (qOut = [qOut zeros(size(qOut), ntSrc - ntComp)])
    qOut = time_resample(qOut,dtComp,srcGeometry)

    # Output adjoint data as judiVector
    return judiVector(srcGeometry,qOut)
end

# u_out = F*u_in
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Nothing, srcData::Array, recGeometry::Nothing, recData::Nothing, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    ntComp = srcData[1].shape[1]

    # Devito call
    u = pycall(ac."forward_modeling", PyObject, modelPy, nothing, srcData[1], nothing,
                  space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface)

    # Output forward wavefield as judiWavefield
    options.save_wavefield_to_disk && (u = dump_wavefield(u))
    return judiWavefield(Info(prod(modelPy.shape), 1, ntComp), dtComp, u)
end

# v_out = F'*v_in
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Nothing, srcData::Nothing, recGeometry::Nothing, recData::Array, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    ntComp = recData[1].shape[1]

    # Devito call
    v = pycall(ac."adjoint_modeling", PyObject, modelPy, nothing, nothing, recData[1],
                  space_order=options.space_order, nb=modelPy.nbpml, free_surface=options.free_surface)

    # Output adjoint wavefield as judiWavefield
    options.save_wavefield_to_disk && (v = dump_wavefield(v))
    return judiWavefield(Info(prod(modelPy.shape), 1, ntComp), dtComp, v)
end

# d_lin = J*dm
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Geometry, srcData::Array, recGeometry::Geometry, recData::Nothing, dm::Array, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    tmaxRec = recGeometry.t[1]
    qIn = time_resample(srcData[1],srcGeometry,dtComp)[1]
    ntComp = size(qIn,1)
    ntRec = Int(trunc(tmaxRec/dtComp + 1))

    # Set up coordinates with devito dimensions
    #origin = get_origin(modelPy)
    src_coords = setup_grid(srcGeometry, modelPy.shape, origin)
    rec_coords = setup_grid(recGeometry, modelPy.shape, origin)

    # Devito call
    dOut = pycall(ac."forward_born", Array{Float32,2}, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(qIn))), PyReverseDims(copy(transpose(rec_coords))),
                  space_order=options.space_order, nb=modelPy.nbpml, isic=options.isic)
    ntRec > ntComp && (dOut = [dOut zeros(size(dOut,1), ntRec - ntComp)])
    dOut = time_resample(dOut,dtComp,recGeometry)

    # Output linearized shot records as judiVector
    if options.save_data_to_disk
        container = write_shot_record(srcGeometry,srcData,recGeometry,dOut,options)
        return judiVector(container)
    else
        return judiVector(recGeometry,dOut)
    end
end

# dm = J'*d_lin
function devito_interface(modelPy::PyCall.PyObject, origin, srcGeometry::Geometry, srcData::Array, recGeometry::Geometry, recData::Array, dm::Nothing, options::Options)
    ac = load_acoustic_codegen()

    # Interpolate input data to computational grid
    dtComp = modelPy.critical_dt
    qIn = time_resample(srcData[1],srcGeometry,dtComp)[1]
    dIn = time_resample(recData[1],recGeometry,dtComp)[1]

    # Set up coordinates with devito dimensions
    src_coords = setup_grid(srcGeometry, modelPy.shape, origin)
    rec_coords = setup_grid(recGeometry, modelPy.shape, origin)

    if options.optimal_checkpointing == true
        op_F = pycall(ac."forward_modeling", PyObject, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(qIn))), PyReverseDims(copy(transpose(rec_coords))),
                      op_return=true, space_order=options.space_order, nb=modelPy.nbpml)
        grad = pycall(ac."adjoint_born", Array{Float32, length(modelPy.shape)}, modelPy, PyReverseDims(copy(transpose(rec_coords))), PyReverseDims(copy(transpose(dIn))), op_forward=op_F, space_order=options.space_order,
            nb=modelPy.nbpml, is_residual=true, isic=options.isic, n_checkpoints=options.num_checkpoints, maxmem=options.checkpoints_maxmem)
    elseif ~isempty(options.frequencies)    # gradient in frequency domain
        typeof(options.frequencies) == Array{Any,1} && (options.frequencies = options.frequencies[1])
        d_pred, uf_real, uf_imag = pycall(ac."forward_freq_modeling", PyObject, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(qIn))), PyReverseDims(copy(transpose(rec_coords))),
                                          options.frequencies, space_order=options.space_order, nb=modelPy.nbpml, factor=options.dft_subsampling_factor)
        grad = pycall(ac."adjoint_freq_born", Array{Float32, length(modelPy.shape)}, modelPy, PyReverseDims(copy(transpose(rec_coords))), PyReverseDims(copy(transpose(dIn))),
                      options.frequencies, uf_real, uf_imag, space_order=options.space_order, nb=modelPy.nbpml, isic=options.isic, factor=options.dft_subsampling_factor)
    else
        u0 = get(pycall(ac."forward_modeling", PyObject, modelPy, PyReverseDims(copy(transpose(src_coords))), PyReverseDims(copy(transpose(qIn))), PyReverseDims(copy(transpose(rec_coords))),
                    space_order=options.space_order, nb=modelPy.nbpml, save=true, tsub_factor=options.subsampling_factor), 1)
        grad = pycall(ac."adjoint_born", Array{Float32, length(modelPy.shape)}, modelPy, PyReverseDims(copy(transpose(rec_coords))), PyReverseDims(copy(transpose(dIn))), u=u0,
                      space_order=options.space_order, tsub_factor=options.subsampling_factor, nb=modelPy.nbpml, isic=options.isic)
    end

    # Remove PML and return gradient as Array
    grad = remove_padding(grad,modelPy.nbpml, true_adjoint=options.sum_padding)
    return vec(grad)
end


@testset "randomForest.jl" begin

    data_surv = CSO.dataset("cso_raw")
    data_trudeau = CSO.dataset("Montreal-Trudeau_15min_precipitations")
    data = innerjoin(data_surv, data_trudeau, on=:Date)
    dropmissing!(data, :Duration)
    train, test = train_test_year(data, 2020)

    varnames = [:d15min, :d1h, :d6h]

    model = RandomForestClassifier(n_trees=20, max_depth = 3)

    #The RMSD is well derived
    @test typeof(RForest(model, "Duration", varnames, train, test)) == Float64

    #Predictions are well derived
    @test typeof(RForest(model, "Duration", varnames, train, test, true, false)) == Vector{Float64}

end