using DataFrames, CSV, Distributions

# Test Data
    # source_data = rand(Normal(), 1000, 5)
    # FileName = "sip_export"

"""
    dfnames(headings)

Adds ":" in front of all items in an array. This format is required by the DataFrame packages when giving names to columns.
"""
function dfnames(headings)
h_list = []
    for i = 1:size([headings...],1)
        name = ":" * headings[i]
        println(name)
        push!(h_list, name)

    end
    return h_list
end

"""
    sip2csv(FileName, source_df, s_name="SLURP", s_origin="Julia Language")

- In order to export a SIP Library from Julia, you simply need to have a DataFrame.
- Only specify header fields if your data does not contain any.
- Also make to specify the full filename, including extension
"""
function sip2csv(FileName, source_df, s_name="SLURP", s_origin="Julia Language")

    csvfilename = "$FileName.csv"
    csvfile = open(csvfilename, "w") #Create CSV csvfile
    numSamples = size(source_df, 1)
    colCount = size(source_df, 2)
    df_names = SIPNames(source_df)

    #colwise methods
    meta_mean = lineprep([[mean(col) for col = eachcol(source_df)]])
    meta_median = lineprep([[median(col) for col = eachcol(source_df)]])
    meta_min = lineprep([[minimum(col) for col = eachcol(source_df)]])
    meta_max = lineprep([[maximum(col) for col = eachcol(source_df)]])

    #Load MetaData Table
    metafilename = "$s_name"*"_meta.csv"

    if isfile(metafilename) == true
        meta_df = CSV.read(metafilename, datarow=4)
        AttributeCount = size(meta_df, 1)
        SipAttrsEnd = 23 + AttributeCount
        SipTlc = SipAttrsEnd + 4

        #metadata = convert(Matrix, meta_df)
        # Setup SIP Header
        write(csvfile,"CSV,,,,,,Always have something in A1,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        Control,,SheetName,N/A,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,FilePath,$FileName.csv,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,SlurpAttrs,C14:C16,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,SipAttrs,C19:C$SipAttrsEnd,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,SipTlc,A$SipTlc,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,ClearFirst,TRUE,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,numSamples,$numSamples,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        SLURP,,name,$s_name,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,count,$colCount,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,origin,$s_origin,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        SIP,,name,$df_names
        ,,mean,$meta_mean
        ,,median,$meta_median
        ,,min,$meta_min
        ,,max,$meta_max,\n")
        close(csvfile)

        CSV.write("$FileName.csv", meta_df, append=true)
        csvfile = open(csvfilename, "a")
        write(csvfile,",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n")

        close(csvfile)

        CSV.write("$FileName.csv", source_df, append=true, writeheader=true)

    elseif isfile(metafilename) == false

        # Setup SIP Header
        write(csvfile,"CSV,,,,,,Always have something in A1,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        Control,,SheetName,N/A,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,FilePath,$FileName.csv,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,SlurpAttrs,C14:C16,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,SipAttrs,C19:C23,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,SipTlc,A27,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,ClearFirst,TRUE,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,numSamples,$numSamples,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        SLURP,,name,$s_name,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,count,$colCount,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,origin,$s_origin,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        SIP,,name,$df_names
        ,,mean,$meta_mean
        ,,median,$meta_median
        ,,min,$meta_min
        ,,max,$meta_max
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
        ,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,\n")
        close(csvfile)
        CSV.write("$FileName.csv", source_df, append=true, writeheader=true)
    end

end

"""
    SIPNames(source_df)

Preps a vector of names for the CSV export of SIPs by removing unnecessary characters in front of all items in an array.
"""
function SIPNames(source_df)
    df_names=names(source_df)
    df_names=string([df_names])
    df_names = replace(df_names, "Array{Symbol,1}[[:" =>"")
    df_names = replace(df_names, "]]" =>"")
    df_names = replace(df_names, ":" =>"")
    return df_names
end
"""
    lineprep(line_array)

When you want to strip "[]" from Arrays for export to CSV

"""
function lineprep(line_array)
    val = string([line_array])
    val = replace(val, "Array{Array{Float64,1},1}" =>"")
    val = replace(val, "[" =>"")
    val = replace(val, "]" =>"")
    return val
end

#Loading a SIPloadsip
"""
    importxlsip(FileName, source="")

This function allows to import SIPs in CSV format from Excel using the SIP 2.0 Standard (Stochastic Information Packets, Savage[2009]) to build unified simulation models.
Based on the CSV package, the function cleans up meta data and cleans out redundant columns. You can also set a *source string* that will be included in your SIP DataFrame as a dimension.
"""
function importxlsip(FileName, source="")
println("Processing SIP Data...\n")
    open(FileName) do file
        linecounter = 1
            for ln in eachline(file)
                if occursin("SIP,,name,", ln) == true
                    println("header_row: $linecounter")
                    global header_row = linecounter
                    linecounter +=1
                elseif occursin("SIPs,,,", ln) == true
                    println("data_row: $linecounter")
                    global start_row = linecounter
                else linecounter +=1
                end
            end
    end
    println("")
    println("SIPs are acceccisble via Dataframe. \n")
    println("header=$header_row, datarow=$start_row")
    SLURP = CSV.read(FileName, header=header_row, datarow=start_row)
    select!(SLURP, Not([:SIP, :Column2, :name])) #replaces deletecols!
    #SLURP[!, :source] .= source # new syntax. now uses .functions to broadcast
    println("")
    print(describe(SLURP))
    println("\n")
    return SLURP
end

"""
    importsip(FileName, source="Not specified by user")

This function allows to import SIPs in CSV format from Julia using the SIP 2.0 Standard (Stochastic Information Packets, Savage[2009]) to build unified simulation models.
Based on the CSV package, the function cleans up meta data and cleans out redundant columns. You can also set a *source string* that will be included in your SIP DataFrame as a dimension.
"""
function importsip(FileName, source="Not specified by user")
    open(FileName) do file
        global linecounter = 1
            for ln in eachline(file)
            linecounter +=1
            end

            if linecounter < 1000
                rounding_digit = -2
            else rounding_digit = -3
            end

        Trial_Count = round(linecounter, digits=rounding_digit)
        global start_row = Int64(linecounter - Trial_Count)
        global header_row = start_row-1
        end
    println("Header Row: $header_row and Data_Row: $start_row")
    SLURP = CSV.read(FileName, header=header_row, datarow=start_row)
    #SLURP[!, :source] .= source # new syntax. now uses .functions to broadcast
    println("")
    print(describe(SLURP))
    println("\n")
return SLURP
end

"""
    genmeta(source_df, s_name="SLURPName")

SIPs using the SIP 2.0 Standard (Stochastic Information Packets, Savage[2009]) require that Meta Data be available and maintained. This function creates a file template to add MetaData to your SIP Library. This is a seperate file must accompany the DataFrame to generate the SIP Library correctly. If ommited, the file will export the SIP Library without any metadata.
"""
function genmeta(source_df, s_name="SLURPName")
    #File Operations
    csvfilename = string("$s_name","_meta.csv")
    #csvfilename = "$FileName.csv"
    csvfile = open(csvfilename, "w") #Create CSV csvfile
    ColHeadings = SIPNames(source_df)

    #CSV Body
    write(csvfile,"This file is used to store SIP level MetaData and dimensions for SIP Library : $s_name,,
    ,,,
    ,,SIP Attribute, $ColHeadings,,
    ,,,\n")

#Close File
    close(csvfile)
    CSV.write("$csvfilename.csv")
end


import GraphUtils: get_graph, get_vtx

export
    AbstractGraphState,
    GraphState,
    # get_vtx,
    get_t,
    AbstractGraphAction,
    GraphAction,
    get_e,
    get_dt

abstract type AbstractGraphState end
@with_kw_noshow struct GraphState <: AbstractGraphState
    vtx::Int        = -1 # vertex of graph
    t::Int          = -1
end
get_vtx(s::AbstractGraphState)  = s.vtx
get_t(s::AbstractGraphState)    = s.t
is_valid(s::AbstractGraphState) = get_vtx(s) >= 1
Base.string(s::AbstractGraphState) = "(v=$(get_vtx(s)),t=$(get_t(s)))"
Base.show(io::IO,s::AbstractGraphState) = print(io,typeof(s),string(s)) #"(v=$(get_vtx(s)),t=$(get_t(s)))"
# allow paths to start at non-zero start times
get_start_index(path::Path{S,A,C}) where {S<:AbstractGraphState,A,C} = get_t(get_initial_state(path))

# AbstractGraphAction
export AbstractGraphAction

abstract type AbstractGraphAction end
@with_kw_noshow struct GraphAction <: AbstractGraphAction
    e::Edge{Int}    = Edge(-1,-1)
    dt::Int         = 1
end
get_e(a::AbstractGraphAction)   = a.e
get_dt(a::AbstractGraphAction)  = a.dt
Base.reverse(a::AbstractGraphAction) = typeof(a)(a,e=reverse(a.e))
Base.string(a::AbstractGraphAction) = "(e=$(get_e(a).src) → $(get_e(a).dst))"
Base.show(io::IO,a::AbstractGraphAction) = print(io,typeof(a),string(a)) #"(v=$(get_vtx(s)),t=$(get_t(s)))"

export
    GraphEnv

"""
    GraphEnv{S,A,C} <: AbstractLowLevelEnv{S,A,C}

An abstract environment type whose concrete subtypes comply with a standard
interface.
"""
abstract type GraphEnv{S,A,C} <: AbstractLowLevelEnv{S,A,C} end

export
    # get_graph,
    get_cost_model,
    get_agent_id,
    get_constraints,
    get_goal,
    get_heuristic_model

"""
    GraphUtils.get_graph(env::GraphEnv)

Must be implemented for all concrete subtypes of `GraphEnv`
"""
get_graph(env::GraphEnv)            = env.graph
get_cost_model(env::GraphEnv)       = env.cost_model
get_agent_id(env::GraphEnv)         = env.agent_idx
get_constraints(env::GraphEnv)      = env.constraints
get_goal(env::GraphEnv)             = env.goal
get_heuristic_model(env::GraphEnv)  = env.heuristic
function GraphUtils.get_distance(env::GraphEnv,s1::AbstractGraphState,s2::AbstractGraphState)
    get_distance(get_graph(env),get_vtx(s1),get_vtx(s2))
end

set_agent_idx!(env::E,i) where {E<:GraphEnv} = E(env,agent_idx=i)

cost_type(env::GraphEnv)            = cost_type(get_cost_model(env))

is_valid(env::GraphEnv,s::AbstractGraphState) = has_vertex(get_graph(env),get_vtx(s))
is_valid(env::GraphEnv,a::AbstractGraphAction) = has_edge(get_graph(env),get_e(a))
# get_possible_actions(env::GraphEnv,s) = map(v->GraphAction(e=Edge(get_vtx(s),v)),outneighbors(get_graph(env),get_vtx(s)))
get_next_state(s::AbstractGraphState,a::AbstractGraphAction) = GraphState(get_e(a).dst,get_t(s)+get_dt(a))
# get_next_state(env::GraphEnv,s,a) = get_next_state(s,a)
wait(s::AbstractGraphState) = GraphAction(e=Edge(get_vtx(s),get_vtx(s)))
# wait(env::GraphEnv,s) = GraphAction(e=Edge(get_vtx(s),get_vtx(s)))
function get_transition_cost(c::TravelTime,env::GraphEnv,s,a,sp)
    return cost_type(c)(get_dt(a))
end
function get_conflict_value(c::C,agent_id::Int,env::GraphEnv,s::AbstractGraphState,t=get_t(s)) where {C<:SoftConflictTable}
    idx,t = serialize_jointly(env,s,t)
    get_conflict_value(c,agent_id,idx,t)
end
function get_conflict_value(c::C,agent_id::Int,env::GraphEnv,a::AbstractGraphAction,t) where {C<:SoftConflictTable}
    idx,t = serialize_jointly(env,a,t)
    get_conflict_value(c,agent_id,idx,t)
end
function get_transition_cost(c::C,env::GraphEnv,s,a,sp) where {C<:ConflictCostModel}
    state_conflict_value = get_conflict_value(c, get_agent_id(env), get_vtx(sp), get_t(sp))
    edge_conflict_value = min(
        get_conflict_value(c,get_agent_id(env), get_vtx(sp), get_t(s)),
        get_conflict_value(c,get_agent_id(env), get_vtx(s), get_t(sp)),
        )
    if edge_conflict_value > 0
        # println("Possible Edge Conflict")
        edge_conflict_value = 0
        for (i,p) in enumerate(c.table.paths)
            if (get_planned_vtx(c.table, i, get_t(s)) == get_vtx(sp)) && (get_planned_vtx(c.table, i, get_t(sp)) == get_vtx(s))
                edge_conflict_value += 1
            end
        end
    end
    return state_conflict_value + edge_conflict_value
end
function get_transition_cost(c::ConflictCostModel{S},env::GraphEnv,s,a,sp) where {S<:SoftConflictTable}
    state_conflict_value = get_conflict_value(c.table,get_agent_id(env),env,sp)
    action_conflict_value = get_conflict_value(c.table,get_agent_id(env),env,a,get_t(sp))
    return state_conflict_value + action_conflict_value
end
function get_transition_cost(c::TravelDistance,env::GraphEnv,s,a,sp)
    return (get_vtx(s) == get_vtx(sp)) ? 0.0 : 1.0
end
compute_path_cost(c::TravelTime,env::GraphEnv,path,i) = cost_type(c)(get_t(get_final_state(path)))
function compute_path_cost(c::TravelDistance,env::GraphEnv,path,i)
    cost_type(c)(sum(map(
        n->get_e(get_a(n)).src != get_e(get_a(n)).dst,
        path.path_nodes)))
end
# get_heuristic_cost(env::GraphEnv,s) = get_heuristic_cost(env,get_heuristic_model(env),s)
get_heuristic_cost(env::GraphEnv,s) = get_heuristic_cost(get_heuristic_model(env),env,s)
# function get_heuristic_cost(env::GraphEnv,h::H,s) where {H<:Union{PerfectHeuristic,DefaultPerfectHeuristic}}
function get_heuristic_cost(h::H,env::GraphEnv,s) where {H<:Union{PerfectHeuristic,DefaultPerfectHeuristic}}
    get_heuristic_cost(h, get_vtx(get_goal(env)), get_vtx(s))
end
# function get_heuristic_cost(env::GraphEnv,h::EnvDistanceHeuristic,s)
function get_heuristic_cost(h::EnvDistanceHeuristic,env::GraphEnv,s)
    get_distance(env, s, get_goal(env))
end
# function get_heuristic_cost(env::GraphEnv,h::H,s) where {E<:GraphEnv, H<:ConflictTableHeuristic}
function get_heuristic_cost(h::H,env::GraphEnv,s) where {E<:GraphEnv, H<:ConflictTableHeuristic}
    get_heuristic_cost(h, get_agent_id(env), get_vtx(s), get_t(s))
end
# states_match
states_match(s1::AbstractGraphState,s2::AbstractGraphState) = (get_vtx(s1) == get_vtx(s2))
states_match(env::GraphEnv,s1,s2) = (get_vtx(s1) == get_vtx(s2))
################################################################################
######################## Low-Level (Independent) Search ########################
################################################################################
num_states(env::GraphEnv)                           = nv(get_graph(env))
num_actions(env::GraphEnv)                          = num_states(env)^2 # NOT actually true, but necessary for O(1) serialization
state_space_trait(env::GraphEnv)                    = DiscreteSpace()
action_space_trait(env::GraphEnv)                   = DiscreteSpace()
serialize(env::GraphEnv,s::AbstractGraphState,t=get_t(s))            = get_vtx(s), t
# NOTE the return type of deserialize doesn't need to match the state/action type of env, since it is only used for constraint checking
deserialize(env::GraphEnv,s::AbstractGraphState,idx::Int,t=-1) = GraphState(vtx=idx,t=t), t
function serialize(env::GraphEnv,a::AbstractGraphAction,t=-1)
    (get_e(a).src-1)*num_states(env)+get_e(a).dst, t
end
function deserialize(env::GraphEnv,s::AbstractGraphAction,idx::Int,t=get_t(s))
    GraphAction(e = Edge(
        div(idx-1,num_states(env))+1,
        mod(idx-1,num_states(env))+1)
        ), t
end
serialize_jointly(env::GraphEnv,s::AbstractGraphState,t=get_t(s)) = serialize(env,s,t)
function serialize_jointly(env::GraphEnv,a::A,t=-1) where {A<:AbstractGraphAction}
    a_ = A(a,e=Edge(sort([get_e(a).src, get_e(a).dst])...))
    idx, t = serialize(env,a_,t)
    return idx+num_states(env), t
end

"""
    create_reservations(env::GraphEnv,s,a,sp,t=-1)

Generates three reservations as follows
```
t0 = get_t(s)
tF = get_t(sp)
t_mid = (t0+tF)/2
reservations = [
    ResourceReservation{Int}(s_idx,get_agent_id(env),   (t0,    t_mid)),
    ResourceReservation{Int}(a_idx,get_agent_id(env),   (t0,    tF)),
    ResourceReservation{Int}(sp_idx,get_agent_id(env),  (t_mid, tF)),
]
```
In this way, the reservations for one path node will not interfere with those
for the next path node.
"""
function create_reservations(env::GraphEnv,s,a,sp,t=-1)
    s_idx,t0    = serialize(env,s)
    a_idx,_     = serialize_jointly(env,a) # serialize action to not overlap with state serialization space?
    sp_idx,tF   = serialize(env,sp)
    t_mid = (t0+tF)/2
    reservations = [
        ResourceReservation{Float64}(s_idx,get_agent_id(env),   (t0,    t_mid)),
        ResourceReservation{Float64}(a_idx,get_agent_id(env),   (t0,    tF)),
        ResourceReservation{Float64}(sp_idx,get_agent_id(env),  (t_mid, tF)),
    ]
end
function create_reservations(env::GraphEnv,n::PathNode,t=-1)
    create_reservations(env,get_s(n),get_a(n),get_sp(n),t)
end
# is_goal
function is_goal(env::GraphEnv,s)
    if get_t(s) >= get_t(get_goal(env))
        if states_match(s, get_goal(env)) || !is_valid(get_goal(env))
            return true
        end
    end
    return false
end
function violates_constraints(env::GraphEnv, s, a, sp)
    t = get_t(sp)
    if has_constraint(env,get_constraints(env),
        state_constraint(get_agent_id(get_constraints(env)),PathNode(s,a,sp),t)
        )
        return true
    elseif has_constraint(env,get_constraints(env),
        action_constraint(get_agent_id(get_constraints(env)),PathNode(s,a,sp),t)
        )
        return true
    end
    return false
end

################################################################################
###################### Conflict-Based Search (High-Level) ######################
################################################################################
function detect_state_conflict(n1::N,n2::N) where {S<:AbstractGraphState,A<:AbstractGraphAction,N<:PathNode{S,A}}
    if get_vtx(n1.sp) == get_vtx(n2.sp) && get_t(n1.sp) == get_t(n2.sp)
        return true
    end
    return false
end
function detect_action_conflict(n1::N,n2::N) where {S<:AbstractGraphState,A<:AbstractGraphAction,N<:PathNode{S,A}}
    if (get_e(n1.a).src == get_e(n2.a).dst) && (get_e(n1.a).dst == get_e(n2.a).src) && (get_t(n1.sp) == get_t(n2.sp))
        return true
    end
    return false
end

################################################################################
############################### HELPER FUNCTIONS ###############################
################################################################################
""" Helper for displaying Paths """
function convert_to_vertex_lists(path::Path{S,A,C}) where {S<:AbstractGraphState,A<:AbstractGraphAction,C}
    vtx_list = [get_vtx(n.sp) for n in path.path_nodes]
    # if length(path) > 0
    #     vtx_list = [get_s(get_path_node(path,1)).vtx, vtx_list...]
    # else
        vtx_list = [get_vtx(get_initial_state(path)), vtx_list...]
    # end
    vtx_list
end


"""
This function reads the RTS-GMLC timeseries data and parses it into the format
compatible with AgentSimulation
"""

function read_rts(data_dir::String,
                  test_system_dir::String,
                  base_dir::String,
                  annual_growth::AxisArrays.AxisArray{Float64, 2},
                  start_year::Int64,
                  n_rep_days::Int64)

    test_system_load = DataFrames.DataFrame(CSV.File(joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "Load", "DAY_AHEAD_regional_Load.csv")))

    base_year = test_system_load[1, "Year"]

    test_sys_num_hours = DataFrames.nrow(test_system_load)
    if test_sys_num_hours >= 8760
        test_sys_hour_weight = ones(test_sys_num_hours)
    else
        test_sys_hour_weight = ones(test_sys_num_hours) * 8760 / test_sys_num_hours
    end

    zone_numbers = names(test_system_load)[5:end]
    zones = ["zone_$(i)" for i in zone_numbers]

    test_system_reserve_up = test_system_load[:, 1:4]
    test_system_reserve_down = test_system_load[:, 1:4]

    for zone in zone_numbers
        reserve_up_df = DataFrames.DataFrame(CSV.File(joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "Reserves", "DAY_AHEAD_regional_Reg_Up_R$(zone).csv")))

        test_system_reserve_up[:, zone] = reserve_up_df[:, "Reg_Up_R$(zone)"]

        reserve_down_df = DataFrames.DataFrame(CSV.File(joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "Reserves", "DAY_AHEAD_regional_Reg_Down_R$(zone).csv")))
        test_system_reserve_down[:, zone] = zeros(DataFrames.nrow(test_system_reserve_down))
        for d in 1:Int(DataFrames.nrow(test_system_reserve_down)/24)
            for h in 1:24
                test_system_reserve_down[(d - 1) * 24 + h, zone] = reserve_down_df[d, Symbol(h)]
            end
        end
    end

    scaled_test_sys_load = deepcopy(test_system_load)
    scaled_test_sys_reserve_up = deepcopy(test_system_reserve_up)
    scaled_test_sys_reserve_down = deepcopy(test_system_reserve_down)

    scaled_test_sys_load[:, "Year"] = fill(start_year, DataFrames.nrow(scaled_test_sys_load))
    remove_leap_day!(scaled_test_sys_load, start_year)

    scaled_test_sys_reserve_up[:, "Year"] = fill(start_year, DataFrames.nrow(scaled_test_sys_reserve_up))
    remove_leap_day!(scaled_test_sys_reserve_up, start_year)

    scaled_test_sys_reserve_down[:, "Year"] = fill(start_year, DataFrames.nrow(scaled_test_sys_reserve_down))
    remove_leap_day!(scaled_test_sys_reserve_down, start_year)

    for y in 1:(start_year - base_year)
        for zone in zone_numbers
            scaled_test_sys_load[:, "$(zone)"] =  scaled_test_sys_load[:, "$(zone)"] * (1 + annual_growth[y, Symbol("load_zone_$(zone)")])
            scaled_test_sys_reserve_up[:, "$(zone)"] =  scaled_test_sys_reserve_up[:, "$(zone)"] * (1 + annual_growth[y, Symbol("load_zone_$(zone)")])
            scaled_test_sys_reserve_down[:, "$(zone)"] =  scaled_test_sys_reserve_down[:, "$(zone)"] * (1 + annual_growth[y, Symbol("load_zone_$(zone)")])
        end
    end

    net_load_df = scaled_test_sys_load[:, 1:4]

    for zone in zone_numbers
        net_load_df[:, "load_zone_$(zone)"] = scaled_test_sys_load[:, zone]
    end

    existing_generator_data = DataFrames.DataFrame(CSV.File(joinpath(test_system_dir, "RTS_Data", "SourceData", "gen.csv")))

    wind_timeseries_file = joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "WIND", "DAY_AHEAD_wind.csv")
    if isfile(wind_timeseries_file)
        wind_timeseries_data = DataFrames.DataFrame(CSV.File(wind_timeseries_file))
        remove_leap_day!(wind_timeseries_data, start_year)
    end

    pv_timeseries_file = joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "PV", "DAY_AHEAD_pv.csv")
    if isfile(pv_timeseries_file)
        pv_timeseries_data = DataFrames.DataFrame(CSV.File(pv_timeseries_file))
        remove_leap_day!(pv_timeseries_data, start_year)
    end

    rtpv_timeseries_file = joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "RTPV", "DAY_AHEAD_rtpv.csv")
    if isfile(rtpv_timeseries_file)
        rtpv_timeseries_data = DataFrames.DataFrame(CSV.File(rtpv_timeseries_file))
        remove_leap_day!(rtpv_timeseries_data, start_year)
    end

    hydro_timeseries_file = joinpath(test_system_dir, "RTS_Data", "timeseries_data_files", "Hydro", "DAY_AHEAD_hydro.csv")
    if isfile(hydro_timeseries_file)
        hydro_timeseries_data = DataFrames.DataFrame(CSV.File(hydro_timeseries_file))
        remove_leap_day!(hydro_timeseries_data, start_year)
    end

    gen_availability_df = scaled_test_sys_load[:, 1:4]

    for i in 1:DataFrames.nrow(existing_generator_data)
        if existing_generator_data[i, "Unit Type"] == "WIND"
            gen_availability_df[:, existing_generator_data[i, "GEN UID"]] = wind_timeseries_data[:, existing_generator_data[i, "GEN UID"]] / existing_generator_data[i, "PMax MW"]
            net_load_df[:, existing_generator_data[i, "GEN UID"]] = wind_timeseries_data[:, existing_generator_data[i, "GEN UID"]]
        elseif existing_generator_data[i, "Unit Type"] == "PV"
            gen_availability_df[:, existing_generator_data[i, "GEN UID"]] = pv_timeseries_data[:, existing_generator_data[i, "GEN UID"]] / existing_generator_data[i, "PMax MW"]
            net_load_df[:, existing_generator_data[i, "GEN UID"]] = pv_timeseries_data[:, existing_generator_data[i, "GEN UID"]]
        elseif existing_generator_data[i, "Unit Type"] == "RTPV"
            gen_availability_df[:, existing_generator_data[i, "GEN UID"]] = rtpv_timeseries_data[:, existing_generator_data[i, "GEN UID"]] / existing_generator_data[i, "PMax MW"]
        elseif existing_generator_data[i, "Unit Type"] == "HYDRO"
            gen_availability_df[:, existing_generator_data[i, "GEN UID"]] = hydro_timeseries_data[:, existing_generator_data[i, "GEN UID"]] / existing_generator_data[i, "PMax MW"]
        else
            gen_availability_df[:, existing_generator_data[i, "GEN UID"]] = ones(DataFrames.nrow(gen_availability_df))
        end
    end

    write_data(joinpath(data_dir, "timeseries_data_files", "Availability"), "DAY_AHEAD_availability.csv", gen_availability_df)
    write_data(joinpath(data_dir, "timeseries_data_files", "Net Load Data"), "load_n_vg_data.csv", net_load_df)

    representative_days = find_representative_days(data_dir, base_dir, n_rep_days)
    rep_load_data = filter(row -> in(Dates.Date(row[:Year], row[:Month], row[:Day]), keys(representative_days)), scaled_test_sys_load)

    rep_hour_weight = zeros(DataFrames.nrow(rep_load_data))

    for i in 1:length(rep_hour_weight)
        rep_hour_weight[i] = representative_days[Dates.Date(rep_load_data[i, :Year], rep_load_data[i, :Month], rep_load_data[i, :Day])]
    end

    write_data(joinpath(data_dir, "timeseries_data_files", "Load"), "load_0.csv", scaled_test_sys_load)
    write_data(joinpath(data_dir, "timeseries_data_files", "Load"), "rep_load_0.csv", rep_load_data)

    system_peak_load = maximum(sum(scaled_test_sys_load[:, zone] for zone in zone_numbers))

    rep_reserve_up_data = filter(row -> in(Dates.Date(row[:Year], row[:Month], row[:Day]), keys(representative_days)), scaled_test_sys_reserve_up)
    write_data(joinpath(data_dir, "timeseries_data_files", "Reserves"), "reserve_up_0.csv", scaled_test_sys_reserve_up)
    write_data(joinpath(data_dir, "timeseries_data_files", "Reserves"), "rep_reserve_up_0.csv", rep_reserve_up_data)

    rep_reserve_down_data = filter(row -> in(Dates.Date(row[:Year], row[:Month], row[:Day]), keys(representative_days)), scaled_test_sys_reserve_down)
    write_data(joinpath(data_dir, "timeseries_data_files", "Reserves"), "reserve_down_0.csv", scaled_test_sys_reserve_down)
    write_data(joinpath(data_dir, "timeseries_data_files", "Reserves"), "rep_reserve_down_0.csv", rep_reserve_down_data)

    reserve_params_df = read_data(joinpath(test_system_dir, "RTS_Data", "SourceData", "reserves.csv"))

    reserve_up_eligible_categories = String[]
    reserve_down_eligible_categories = String[]

    for zone in zone_numbers
        push!(reserve_up_eligible_categories, reserve_params_df[findfirst(x-> x == "Reg_Up_R$(zone)", reserve_params_df[:, "Reserve Product"]), "Eligible Device SubCategories"])
        push!(reserve_down_eligible_categories, reserve_params_df[findfirst(x-> x == "Reg_Down_R$(zone)", reserve_params_df[:, "Reserve Product"]), "Eligible Device SubCategories"])
    end

    reserve_up_mkt_df = read_data(joinpath(data_dir, "markets_data", "reserve_up_mkt_param.csv"))
    reserve_down_mkt_df = read_data(joinpath(data_dir, "markets_data", "reserve_down_mkt_param.csv"))

    reserve_up_mkt_df[!, "eligible categories"] = reserve_up_eligible_categories
    reserve_down_mkt_df[!, "eligible categories"] = reserve_down_eligible_categories

    write_data(joinpath(data_dir, "markets_data"), "reserve_up_mkt_param.csv", reserve_up_mkt_df)
    write_data(joinpath(data_dir, "markets_data"), "reserve_down_mkt_param.csv", reserve_down_mkt_df)

    # Create zonal lines

    zonal_lines = ZonalLine[]

    branches = read_data(joinpath(test_system_dir, "RTS_Data", "SourceData", "branch.csv"))
    dc_branches = read_data(joinpath(test_system_dir, "RTS_Data", "SourceData", "dc_branch.csv"))

    for b in 1:DataFrames.nrow(branches)
        from_bus = "$(branches[b, "From Bus"])"
        from_zone = "zone_$(first(from_bus, 1))"

        to_bus = "$(branches[b, "To Bus"])"
        to_zone = "zone_$(first(to_bus, 1))"

        similar_line = filter(l -> (in(from_zone, [get_from_zone(l), get_to_zone(l)]) && in(to_zone, [get_from_zone(l), get_to_zone(l)])), zonal_lines)

        if length(similar_line) < 1
            push!(zonal_lines, ZonalLine(branches[b, "UID"], from_zone, to_zone, branches[b, "Cont Rating"]))
        else
            set_active_power_limit!(similar_line[1], get_active_power_limit(similar_line[1]) + branches[b, "Cont Rating"])
        end
    end

    for b in 1:DataFrames.nrow(dc_branches)
        from_bus = "$(dc_branches[b, "From Bus"])"
        from_zone = "zone_$(first(from_bus, 1))"

        to_bus = "$(dc_branches[b, "To Bus"])"
        to_zone = "zone_$(first(to_bus, 1))"

        similar_line = filter(l -> (in(from_zone, [get_from_zone(l), get_to_zone(l)]) && in(to_zone, [get_from_zone(l), get_to_zone(l)])), zonal_lines)

        if length(similar_line) < 1
            push!(zonal_lines, ZonalLine(dc_branches[b, "UID"], from_zone, to_zone, dc_branches[b, "MW Load"]))
        else
            set_active_power_limit!(similar_line[1], get_active_power_limit(similar_line[1]) + dc_branches[b, "MW Load"])
        end
    end

    return zones, representative_days, rep_hour_weight, system_peak_load, test_sys_hour_weight, zonal_lines
end


#=
    DenseKKTSystem
=#

"""
    DenseKKTSystem{T, VT, MT} <: AbstractReducedKKTSystem{T, MT}

Implement [`AbstractReducedKKTSystem`](@ref) with dense matrices.

Requires a dense linear solver to be factorized (otherwise an error is returned).

"""
struct DenseKKTSystem{T, VT, MT} <: AbstractReducedKKTSystem{T, MT}
    hess::MT
    jac::MT
    pr_diag::VT
    du_diag::VT
    diag_hess::VT
    # KKT system
    aug_com::MT
    # Info
    ind_ineq::Vector{Int}
    ind_fixed::Vector{Int}
    jacobian_scaling::VT
    # Buffers
    etc::Dict{Symbol, Any}
end

function DenseKKTSystem{T, VT, MT}(n, m, ind_ineq, ind_fixed) where {T, VT, MT}
    ns = length(ind_ineq)
    hess = MT(undef, n, n)
    jac = MT(undef, m, n+ns)
    pr_diag = VT(undef, n+ns)
    du_diag = VT(undef, m)
    diag_hess = VT(undef, n)

    # If the the problem is unconstrained, then KKT system is directly equal
    # to the Hessian (+ some regularization terms)
    aug_com = if (m == 0)
        hess
    else
        MT(undef, n+ns+m, n+ns+m)
    end

    jacobian_scaling = VT(undef, m)

    # Init!
    fill!(aug_com, zero(T))
    fill!(hess,    zero(T))
    fill!(jac,     zero(T))
    fill!(pr_diag, zero(T))
    fill!(du_diag, zero(T))
    fill!(diag_hess, zero(T))
    fill!(jacobian_scaling, one(T))

    return DenseKKTSystem{T, VT, MT}(
        hess, jac, pr_diag, du_diag, diag_hess, aug_com,
        ind_ineq, ind_fixed, jacobian_scaling, Dict{Symbol, Any}(),
    )
end

function DenseKKTSystem{T, VT, MT}(nlp::AbstractNLPModel, info_constraints=get_index_constraints(nlp)) where {T, VT, MT}
    return DenseKKTSystem{T, VT, MT}(
        get_nvar(nlp), get_ncon(nlp), info_constraints.ind_ineq, info_constraints.ind_fixed
    )
end

is_reduced(::DenseKKTSystem) = true
num_variables(kkt::DenseKKTSystem) = length(kkt.pr_diag)

# Special getters for Jacobian
function get_jacobian(kkt::DenseKKTSystem)
    n = size(kkt.hess, 1)
    ns = length(kkt.ind_ineq)
    return view(kkt.jac, :, 1:n)
end

get_raw_jacobian(kkt::DenseKKTSystem) = kkt.jac

nnz_jacobian(kkt::DenseKKTSystem) = length(kkt.jac)

function diag_add!(dest::AbstractMatrix, d1::AbstractVector, d2::AbstractVector)
    n = length(d1)
    @inbounds for i in 1:n
        dest[i, i] = d1[i] + d2[i]
    end
end

function _build_dense_kkt_system!(dest, hess, jac, pr_diag, du_diag, diag_hess, n, m, ns)
    # Transfer Hessian
    for i in 1:n, j in 1:i
        if i == j
            dest[i, i] = pr_diag[i] + diag_hess[i]
        else
            dest[i, j] = hess[i, j]
            dest[j, i] = hess[j, i]
        end
    end
    # Transfer slack diagonal
    for i in 1:ns
        dest[i+n, i+n] = pr_diag[i+n]
    end
    # Transfer Jacobian
    for i in 1:m, j in 1:(n+ns)
        dest[i + n + ns, j] = jac[i, j]
        dest[j, i + n + ns] = jac[i, j]
    end
    # Transfer dual regularization
    for i in 1:m
        dest[i + n + ns, i + n + ns] = du_diag[i]
    end
end

function build_kkt!(kkt::DenseKKTSystem{T, VT, MT}) where {T, VT, MT}
    n = size(kkt.hess, 1)
    m = size(kkt.jac, 1)
    ns = length(kkt.ind_ineq)
    if m == 0 # If problem is unconstrained, just need to update the diagonal
        diag_add!(kkt.aug_com, kkt.diag_hess, kkt.pr_diag)
    else # otherwise, we update the full matrix
        _build_dense_kkt_system!(kkt.aug_com, kkt.hess, kkt.jac, kkt.pr_diag, kkt.du_diag, kkt.diag_hess, n, m, ns)
    end
    treat_fixed_variable!(kkt)
end

function compress_jacobian!(kkt::DenseKKTSystem{T, VT, MT}) where {T, VT, MT}
    m = size(kkt.jac, 1)
    n = size(kkt.hess, 1)
    # Add slack indexes
    for i in kkt.ind_ineq
        kkt.jac[i, i+n] = -one(T)
    end
    # Scale
    kkt.jac .*= kkt.jacobian_scaling
    return
end

function compress_hessian!(kkt::DenseKKTSystem)
    # Transfer diagonal term for future regularization
    diag!(kkt.diag_hess, kkt.hess)
end

function mul!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector)
    mul!(y, kkt.aug_com, x)
end

function jtprod!(y::AbstractVector, kkt::DenseKKTSystem, x::AbstractVector)
    mul!(y, kkt.jac', x)
end

function set_jacobian_scaling!(kkt::DenseKKTSystem, constraint_scaling::AbstractVector)
    copyto!(kkt.jacobian_scaling, constraint_scaling)
end


using .Glimpse
const Gl = Glimpse
using ColorSchemes
using .Glimpse: RGBf0

struct WfInterface <: Gl.System end

function Gl.update(::WfInterface, m::AbstractLedger)
    wfman = Gl.singleton(m, WfManager)
    curvis = wfman.entities[wfman.current]

    gui_func = () -> begin
        Gl.Gui.SetNextWindowPos((650, 200), Gl.Gui.ImGuiCond_FirstUseEver)
        Gl.Gui.SetNextWindowSize((550, 680), Gl.Gui.ImGuiCond_FirstUseEver)
        Gl.Gui.Begin("Wannier Function Manager")
        Gl.Gui.SetWindowFontScale(3.3f0)
        curid = Int32(wfman.current)
        Gl.Gui.@c Gl.Gui.InputInt("Id", &curid)
        if 0 < curid <= length(wfman.entities)
            wfman.current = Int(curid)
        elseif curid <= 0
            wfman.current = length(wfman.entities)
        else
            wfman.current = 1
        end
        Gl.Gui.@c Gl.Gui.InputDouble("dt", &wfman.dt,0.01,0.01,"%.3f")
        if wfman.dt <= 0.0
            wfman.dt += 0.01
        end
        if !wfman.iterating
            if Gl.Gui.Button("Iterate")
                wfman.iterating = true
            end
        else
            if Gl.Gui.Button("Stop Iterating")
                wfman.iterating = false
            end
        end
        Gl.Gui.End()
    end
    push!(Gl.singleton(m, Gl.GuiFuncs).funcs, gui_func)

    if wfman.iterating
        wfman.current_t += m[Gl.TimingData][1].dtime
        if wfman.current_t >= wfman.dt
            wfman.current += 1
            if wfman.current >= length(wfman.entities)
                wfman.current= 1
            end
            wfman.current_t = 0.0
        end
    end

    vis = m[Gl.Visible]
    for e in wfman.entities
        vis[e] = Gl.Visible(e == curvis)
    end
end

@component @with_kw mutable struct WfManager
    entities::Vector{Entity}
    wfuncs::Vector{<:WannierFunction}
    current::Int = 1
    dt::Float64 = 0.5
    current_t::Float64 = 0.0
    iterating::Bool = false
end

"""
    visualize_wfuncs(wfuncs::Vector{<:WannierFunction}, str::Structure;
                     iso_ratio = 1/4,
                     alpha     = 0.6,
                     material  = Gl.Material(),
                     phase_channel = Up())

Visualize the wannierfunctions in a Diorama.
The isosurface is determined from the charge, where the coloring signifies the phase of the wavefunction at that point in space.
The `iso_ratio` will be used to determine the isosurface values as the ratio w.r.t the maximum value.
The `phase_channel` denotes whether the phase of the spin up or down channel should be shown in the case of spinor Wannierfunctions.
"""
function visualize_wfuncs(wfuncs::Vector{<:WannierFunction}, str::Structure; kwargs...)
    dio = Diorama(background = RGBAf0(60/255, 60/255, 60/255, 1.0f0))
    visualize_wfuncs!(dio, wfuncs, str; kwargs...)
    return dio
end

function visualize_wfuncs!(dio::Diorama, wfuncs::Vector{<:WannierFunction}, str::Structure;
                          iso_ratio = 1/4,
                          alpha = 0.6,
                          material = Gl.Material(),
                          phase_channel = Up())

    DFControl.Display.add_structure!(dio, str)

    phase_id = length(wfuncs[1].values[1]) > 1 ? (phase_channel == Up() ? 1 : 2) : 1
    colfunc(x) = RGBf0(get(ColorSchemes.rainbow, (angle(x[phase_id]) .+ π)/2π))

    wfentities = fill(Entity(0), length(wfuncs))
    grid = Gl.Grid([Point3f0(w...) for w in wfuncs[1].points])

    for (i, w) in enumerate(wfuncs)
        d = Float32.(norm.(w.values))
        geom = Gl.DensityGeometry(d, iso_ratio * maximum(d))
        color = Gl.DensityColor(colfunc.(w.values))

        wfentities[i] =  Entity(dio, Gl.Spatial(), grid, geom, color, material, Gl.Alpha(alpha))

    end

    man = Entity(dio, WfManager(entities=wfentities, wfuncs=wfuncs))
    push!(dio, :simulation, WfInterface())
    update(WfInterface(), dio)
    return dio
end


function test_vKITTI_Conv(model,predict,w; batchsize=32,saveImages=true,
    testVariations=["30-deg-left","30-deg-right","clone","fog","rain"],
    testMode="even",
    experiment="First_Second"
    )
    @time begin
        printlnFlush("Loading weights from $(w)")
        w = loadNetwork(w)
        totalloss = totalcorrect = totalcount = 0
        for world in worlds
            for variation in testVariations
                x,y,colors = loadData_vKITTI([world],[variation],model,testMode)
                batchCount = Int64(ceil(size(y,4)/batchsize))
                varloss = imcounter = varcorrect = varcount = 0
                for batch=1:batchCount
                    printUpdate(@sprintf("Batch: %5d/%-5d",batch,batchCount))
                    #Calculate current indices
                    startIndex = 1+(batch-1)*batchsize
                    endIndex = min(batch*batchsize,size(y,4))
                    #Minibatch
                    batchx = KnetArray(x[:,:,:,startIndex:endIndex])
                    batchy = y[:,:,:,startIndex:endIndex]
                    #Accuracy
                    ypred = Array(predict(w,batchx))
                    varloss += pixelwiseSoftloss(ypred,batchy)
                    testAcc = pixelwiseAccuracy(ypred,batchy)
                    varcorrect += testAcc[1]
                    varcount += testAcc[2]
                    #Image conversion
                    if saveImages
                        for im=1:size(ypred,4)
                            savePrediction(ypred[:,:,:,im],colors,outputPath(@sprintf("vKITTI/%s/Conv/%s/%s/%s/%05d.png",experiment,model,world,variation,imcounter)))
                            imcounter += 1
                        end
                    end
                end
                printlnFlush(@sprintf("World: %s\tVariation: %s\tLoss: %0.10f\tAccuracy: %0.4f%%",world,variation,varloss/batchCount,varcorrect/varcount*100))
                totalloss += varloss/batchCount
                totalcorrect += varcorrect
                totalcount += varcount
            end
        end
        printlnFlush(@sprintf("Total Loss: %0.10f\tTotal Accuracy: %0.4f%%",totalloss,totalcorrect/totalcount*100))
    end
end

function test_vKITTI_RNN(model,predict,w; saveImages=true,
    testVariations=["30-deg-left","30-deg-right","clone","fog","rain"],
    testMode="even",
    experiment="First_Second"
    )
    @time begin
        printlnFlush("Loading weights from $(w)")
        @time w = loadNetwork(w)
        totalloss = totalcorrect = totalcount = 0
        for world in worlds
            for variation in testVariations
                varloss = imcounter = varcorrect = varcount = 0
                x,y,colors = loadData_vKITTI([world],[variation],model,testMode)
                batchCount = size(y,4)
                h = KnetArray(zeros(Float32,28,28,512,1))
                for t=1:batchCount
                    printUpdate(@sprintf("Time: %5d/%-5d",t,batchCount))
                    #Accuracy
                    ypred = Array(predict(w,h,KnetArray(x[:,:,:,t:t])))
                    varloss += pixelwiseSoftloss(ypred,y[:,:,:,t:t])
                    testAcc = pixelwiseAccuracy(ypred,y[:,:,:,t:t])
                    varcorrect += testAcc[1]
                    varcount += testAcc[2]
                    #Image conversion
                    if saveImages
                        savePrediction(ypred[:,:,:,1],colors,outputPath(@sprintf("vKITTI/%s/RNN/%s/%s/%s/%05d.png",experiment,model,world,variation,imcounter)))
                        imcounter += 1
                    end
                end
                printlnFlush(@sprintf("World: %s\tVariation: %s\tLoss: %0.10f\tAccuracy: %0.4f%%",world,variation,varloss/batchCount,varcorrect/varcount*100))
                totalloss += varloss/batchCount
                totalcorrect += varcorrect
                totalcount += varcount
            end
        end
        printlnFlush(@sprintf("Total Loss: %0.10f\tTotal Accuracy: %0.4f%%",totalloss,totalcorrect/totalcount*100))
    end
end

function test_vKITTI_LSTM(model,predict,w; saveImages=true,
    testVariations=["30-deg-left","30-deg-right","clone","fog","rain"],
    testMode="even",
    experiment="First_Second"
    )
    @time begin
        printlnFlush("Loading weights from $(w)")
        @time w = loadNetwork(w)
        totalloss = totalcorrect = totalcount = 0
        for world in worlds
            for variation in testVariations
                varloss = imcounter = varcorrect = varcount = 0
                x,y,colors = loadData_vKITTI([world],[variation],model,testMode)
                batchCount = size(y,4)
                h = KnetArray(zeros(Float32,28,28,512,1))
                c = KnetArray(zeros(Float32,28,28,512,1))
                for t=1:batchCount
                    printUpdate(@sprintf("Time: %5d/%-5d",t,batchCount))
                    #Accuracy
                    ypred = Array(predict(w,h,c,KnetArray(x[:,:,:,t:t])))
                    varloss += pixelwiseSoftloss(ypred,y[:,:,:,t:t])
                    testAcc = pixelwiseAccuracy(ypred,y[:,:,:,t:t])
                    varcorrect += testAcc[1]
                    varcount += testAcc[2]
                    #Image conversion
                    if saveImages
                        savePrediction(ypred[:,:,:,1],colors,outputPath(@sprintf("vKITTI/%s/LSTM/%s/%s/%s/%05d.png",experiment,model,world,variation,imcounter)))
                        imcounter += 1
                    end
                end
                printlnFlush(@sprintf("World: %s\tVariation: %s\tLoss: %0.10f\tAccuracy: %0.4f%%",world,variation,varloss/batchCount,varcorrect/varcount*100))
                totalloss += varloss/batchCount
                totalcorrect += varcorrect
                totalcount += varcount
            end
        end
        printlnFlush(@sprintf("Total Loss: %0.10f\tTotal Accuracy: %0.4f%%",totalloss,totalcorrect/totalcount*100))
    end
end

function test_vKITTI_GRU(model,predict,w; saveImages=true,
    testVariations=["30-deg-left","30-deg-right","clone","fog","rain"],
    testMode="even",
    experiment="First_Second"
    )
    @time begin
        printlnFlush("Loading weights from $(w)")
        @time w = loadNetwork(w)
        totalloss = totalcorrect = totalcount = 0
        for world in worlds
            for variation in testVariations
                varloss = imcounter = varcorrect = varcount = 0
                x,y,colors = loadData_vKITTI([world],[variation],model,testMode)
                batchCount = size(y,4)
                h = KnetArray(zeros(Float32,28,28,512,1))
                for t=1:batchCount
                    printUpdate(@sprintf("Time: %5d/%-5d",t,batchCount))
                    #Accuracy
                    ypred = Array(predict(w,h,KnetArray(x[:,:,:,t:t])))
                    varloss += pixelwiseSoftloss(ypred,y[:,:,:,t:t])
                    testAcc = pixelwiseAccuracy(ypred,y[:,:,:,t:t])
                    varcorrect += testAcc[1]
                    varcount += testAcc[2]
                    #Image conversion
                    if saveImages
                        savePrediction(ypred[:,:,:,1],colors,outputPath(@sprintf("vKITTI/%s/GRU/%s/%s/%s/%05d.png",experiment,model,world,variation,imcounter)))
                        imcounter += 1
                    end
                end
                printlnFlush(@sprintf("World: %s\tVariation: %s\tLoss: %0.10f\tAccuracy: %0.4f%%",world,variation,varloss/batchCount,varcorrect/varcount*100))
                totalloss += varloss/batchCount
                totalcorrect += varcorrect
                totalcount += varcount
            end
        end
        printlnFlush(@sprintf("Total Loss: %0.10f\tTotal Accuracy: %0.4f%%",totalloss,totalcorrect/totalcount*100))
    end
end

#E1 : Degrees
#E2 : Half Train/Test

#   Conv
function run_vKITTI_Conv_RGBD_E2_test()
    test_vKITTI_Conv("RGBD_VGG",vKITTI_Conv_RGBD,"vKITTI/First_Second/Conv_RGBD.jld"; batchsize=32,saveImages=true,
        testVariations=["30-deg-left","30-deg-right","15-deg-left","15-deg-right","clone","fog","morning","overcast","rain","sunset"],
        testMode="secondHalf",
        experiment="First_Second")
end

#   LSTM
function run_vKITTI_LSTM_RGBD_E2_test()
    test_vKITTI_LSTM("RGBD_VGG",vKITTI_LSTM_RGBD,"vKITTI/First_Second/LSTM_RGBD.jld"; saveImages=true,
        testVariations=["30-deg-left","30-deg-right","15-deg-left","15-deg-right","clone","fog","morning","overcast","rain","sunset"],
        testMode="secondHalf",
        experiment="First_Second")
end

export run_vKITTI_Conv_RGBD_E2_test
export run_vKITTI_LSTM_RGBD_E2_test


struct LCOSY
    # Metadata
    t1          # Indirect time axis
    echo_time   # Minimum echo time (first T1 increment)

    # Indices into underlying acqusition list
    ref_scans::Vector{Int}
    lcosy_scans::Matrix{Int}  # num_averages × length(t1)

    # Raw acquisitions
    acquisitions

    cycle_length
end

standard_metadata(l::LCOSY) = standard_metadata(l.acquisitions)

"""
An L-COSY experiment with (`num_averages` × `nsamp_t1`) acquisitions, possibly
with reference scans included.

The L-COSY pulse sequence is described in:

* Thomas, M. Albert, et al. "Localized two‐dimensional shift correlated MR
  spectroscopy of human brain." Magnetic Resonance in Medicine 46.1 (2001): 58-67.
* Thomas, M. Albert, et al. "Evaluation of two‐dimensional L‐COSY and JPRESS
  using a 3 T MRI scanner: from phantoms to human brain in vivo." NMR in
  Biomedicine 16.5 (2003): 245-251.
"""
function LCOSY(
  t1, echo_time, ref_scans::Vector{Int}, lcosy_scans::Matrix{Int}, acquisitions
)
    LCOSY(
        t1, echo_time, ref_scans, lcosy_scans, acquisitions, count_cycles(acquisitions)
    )
end

function Base.show(io::IO, lcosy::LCOSY)
    println(io, """
                LCOSY experiment:
                  t1 = $(lcosy.t1)
                  size(lcosy_scans) = $(size(lcosy.lcosy_scans))
                  length(ref_scans) = $(length(lcosy.ref_scans))""")
    show(io, lcosy.acquisitions)
end

"""
    extract_fids(lcosy::LCOSY; downsample=1)

Extract raw fid from LCOSY experiment and apply channel combination
"""
function extract_fids(lcosy::LCOSY; downsample=1)
    acqs = lcosy.acquisitions
    # Use reference scans if possible to figure out channel combination
    # weights; if not, use the actual data (which has worse SNR because it's
    # probably water-suppressed)
    combiner_inds = lcosy.ref_scans
    if isempty(combiner_inds)
        # TODO: Put this logic in pca_channel_combiner ?
        combiner_inds = vec(lcosy.lcosy_scans)
    end
    combiner = pca_channel_combiner(sampledata(acqs, i) for i in combiner_inds)

    # Hmm. Calling sampledata to initialize this is kinda ugly...
    fid1 = combiner(sampledata(acqs, lcosy.lcosy_scans[1,1], downsample=downsample))
    t2 = AxisArrays.axes(fid1, Axis{:time}).val

    num_averages = size(lcosy.lcosy_scans, 1)
    phase_cycle = mod1.(1:num_averages, lcosy.cycle_length)

    nsamp_t1 = size(lcosy.lcosy_scans, 2)
    signal = AxisArray(zeros(eltype(fid1), length(t2), num_averages, nsamp_t1),
                       Axis{:time2}(t2),
                       Axis{:phase_cycle}(phase_cycle),
                       Axis{:time1}(lcosy.t1))
    for i=1:nsamp_t1
        for j=1:num_averages
            fid = combiner(sampledata(acqs, lcosy.lcosy_scans[j,i], downsample=downsample))
            signal[:,j,i] = fid
        end
    end
    signal
end

"""
    simple_averaging(spectro_expt)

Simple channel combination and averaging for spectroscopic data acquisition.

Note that this uses a simple mean to combine acquisitions across the phase
cycling dimension; it does no frequency alignment or other calibration.
"""
function simple_averaging(lcosy::LCOSY; downsample=1)
    fids = extract_fids(lcosy, downsample=downsample)
    simple_averaging(fids)
end

"""
    spectrum(lcosy::LCOSY,
             win1=t->sinebell(t, pow=2),
             win2=t->sinebell(t, skew=0.3, pow=2),
             t1pad=4)

Compute spectrum from lcosy data with "standard" L-COSY processing parameters
as have traditionally been used by Mountford et. al.
   * Simple averaging
   * T2: skewed sine bell squared window, skew=0.3
   * T1: Sine bell squared window, 4x zero padded
"""
function spectrum(lcosy::LCOSY;
                  win1=t->sinebell(t, pow=2),
                  win2=t->sinebell(t, skew=0.3, pow=2),
                  t1pad=4)
    signal = simple_averaging(lcosy)
    # Apply sine bell squared windows to signal, as in TRI Felix workflow
    apply_window!(signal, Axis{:time2}, win2)
    apply_window!(signal, Axis{:time1}, win1)
    signal = zeropad(signal, Axis{:time1}, t1pad)
    spec = spectrum(signal)

    # TODO: Flip freq1 axis.
    # TODO: Figure out why we need this...
    #freq1 = AxisArrays.axes(s, Axis{:freq1})
    #freq2 = AxisArrays.axes(s, Axis{:freq2})

    spec
end


using BQCESubroutine
using BQCESubroutine.Utilities
using LinearAlgebra
using YaoLocations: YaoLocations, plain, Locations, CtrlLocations, merge_locations
using Test
using BenchmarkTools
using YaoArrayRegister
using LoopVectorization
using PaddedMatrices
using MatchCore
using MLStyle
using BQCESubroutine: BitContext, BitRoutine, multi_broutine2x2!
using BQCESubroutine: diag_kernel, perm_kernel, generic_kernel, threaded_subspace_loop
using BQCESubroutine: codegen_broutine2x2_multi, codegen_broutine, hoist_entries!, codegen_basic_broutine
using BQCESubroutine: threaded_basic_broutine!, basic_broutine!, threaded_multi_broutine2x2!

st = rand(ComplexF64, 1 << 10);
locs = Locations((9, 10))
threaded_basic_broutine!(copy(st), Val(:PSWAP), locs, 2.0) ≈ basic_broutine!(copy(st), Val(:PSWAP), locs, 2.0)
broutine!(copy(st), Val(:Z), locs) ≈ broutine!(copy(st), BQCESubroutine.BZ, locs)
threaded_basic_broutine!(copy(st), Val(:Z), locs) ≈ broutine!(copy(st), BQCESubroutine.BZ, locs)
threaded_basic_broutine!(copy(st), Val(:Z), locs) ≈ basic_broutine!(copy(st), Val(:Z), locs)

threaded_multi_broutine2x2!(copy(st), Val(:Z), locs) ≈ broutine!(copy(st), BQCESubroutine.BZ, locs)
threaded_multi_broutine2x2!(copy(st), Val(:Z), locs) ≈ multi_broutine2x2!(copy(st), Val(:Z), locs)

@code_lowered broutine!(st, Val(:Z), Locations(24))
BQCESubroutine.set_nthreads(12)
@benchmark broutine!(st, Val(:Z), Locations(24)) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark threaded_basic_broutine!(st, Val(:Z), Locations(25)) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark basic_broutine!(st, Val(:Z), Locations(25)) setup=(st=rand(ComplexF64, 1 << 25))

BQCESubroutine.disable_threads()
@benchmark broutine!(st, Val(:Z), Locations(25)) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark instruct!(st, Val(:Z), (25, )) setup=(st=rand(ComplexF64, 1 << 25))

@benchmark threaded_basic_broutine!(st, Val(:Z), Locations(25)) setup=(st=rand(ComplexF64, 1 << 25))

@benchmark threaded_multi_broutine2x2!(st, Val(:Z), Locations((23, 24))) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark multi_broutine2x2!(st, Val(:Z), Locations((23, 24))) setup=(st=rand(ComplexF64, 1 << 25))
@benchmark instruct!(st, Val(:Z), (23, 24)) setup=(st=rand(ComplexF64, 1 << 25))

@broutine X = [0 1;1 0]

threaded_basic_broutine!(copy(st), Val(:X), locs) ≈ broutine!(copy(st), BX, locs)
multi_broutine2x2!(copy(st), Val(:X), locs)≈ broutine!(copy(st), BX, locs)
threaded_multi_broutine2x2!(copy(st), Val(:X), locs) ≈ multi_broutine2x2!(copy(st), Val(:X), locs)

ctx = BitContext(;threading=true, expand_sz = 0)
brt = BitRoutine(Main, :(Z = [1 0;0 -1]))
codegen_basic_broutine(ctx, brt)|>flatten_block

function BQCESubroutine.threaded_basic_broutine!(var"##st#297"::AbstractVector{T} where T, ::Val{:Z}, var"##locs#298"::Locations)
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:162 =#
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:258 =#
    nothing
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:259 =#
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:296 =#
    var"##plain_locs#299" = (YaoLocations.plain)(var"##locs#298")
    var"##nqubits#300" = (BQCESubroutine.Utilities.log2i)(size(var"##st#297", 1))
    var"##nlocs_needed#301" = (BQCESubroutine.Utilities.log2i)(BQCESubroutine.BQCES_nthreads[] - 1) + 1
    var"##step_1_l#305" = 1 << (var"##plain_locs#299"[1] - 1)
    var"##step_1_h#306" = 1 << var"##plain_locs#299"[1]
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:297 =#
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:298 =#
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:561 =# @fastmath #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:560 =# @inbounds(begin
                if var"##nlocs_needed#301" ≤ var"##nqubits#300" - var"##plain_locs#299"[1]
                    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:455 =#
                    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:414 =#
                    if var"##step_1_l#305" == 1
                        #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:415 =#
                        Threads.@threads for var"##idx#302_1" = 0:var"##step_1_h#306":size(var"##st#297", 1) - var"##step_1_h#306"
                                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:312 =#
                                nothing
                                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:313 =#
                                var"##I#307_2" = ((var"##idx#302_1" + 0) + var"##step_1_l#305") + 1
                                var"##T#308_2" = -(var"##st#297"[var"##I#307_2"])
                                var"##st#297"[var"##I#307_2"] = var"##T#308_2"
                            end
                        #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:416 =#
                        return var"##st#297"
                    end
                    Threads.@threads for var"##idx#302_1" = 0:var"##step_1_h#306":size(var"##st#297", 1) - var"##step_1_h#306"
                            for var"##idx#302_2" = var"##idx#302_1":1:(var"##idx#302_1" + var"##step_1_l#305") - 1
                                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:426 =#
                                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:312 =#
                                nothing
                                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:313 =#
                                var"##I#309_2" = ((var"##idx#302_2" + 0) + var"##step_1_l#305") + 1
                                var"##T#310_2" = -(var"##st#297"[var"##I#309_2"])
                                var"##st#297"[var"##I#309_2"] = var"##T#310_2"
                            end
                        end
                    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:456 =#
                    return var"##st#297"
                end
                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:478 =#
                var"##m#304" = (var"##nqubits#300" - var"##nlocs_needed#301") - 1
                #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:479 =#
                Threads.@threads for var"##base#303" = (BQCESubroutine.Utilities.bsubspace)(var"##nqubits#300", (1:var"##m#304"..., var"##plain_locs#299"...))
                        #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:480 =#
                        for var"##idx#302_1" = 0:1 << var"##m#304" - 1
                            #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:481 =#
                            #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:312 =#
                            nothing
                            #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:313 =#
                            var"##I#311_2" = ((var"##base#303" + var"##idx#302_1") + var"##step_1_l#305") + 1
                            var"##T#312_2" = -(var"##st#297"[var"##I#311_2"])
                            var"##st#297"[var"##I#311_2"] = var"##T#312_2"
                        end
                    end
                return var"##st#297"
            end)
    #= /home/roger/julia_code/BQCESubroutine/src/codegen/broutine.jl:163 =#
    return var"##st#297"
end


using Test
using JuLIP

data = joinpath(dirname(pathof(JuLIP)), "..", "data") * "/"

@testset "EAM with ASE" begin
    using ASE
    import Random

    test_pots = joinpath(datadep"JuLIP_testpots", "JuLIP_data") * "/"
    alloy = test_pots * "PdAgH_HybridPd3Ag.eam.alloy"
    eam_fs = test_pots * "Fe-P.eam.fs"
    Ni = test_pots * "Ni.eam.fs"

    @test EAM(alloy) isa EAM{T, 1} where {T}
    @test EAM(eam_fs) isa EAM{T, 2} where {T}

    eam = pyimport("ase.calculators.eam")
    ase_calc = ASECalculator(eam.EAM(potential=Ni))
    # make sure we get a perfect fit ...
    julip_calc = EAM(Ni)

    atoms = bulk(:Ni) * 3
    Random.seed!(0)
    rattle!(atoms, 0.1)
    E_jl = energy(julip_calc, atoms)
    E_py = energy(ase_calc, atoms)
    # ... but even then we the evaluation codes aren't the
    #     same so we only get ca 1e-6 to 1e-7 match.
    @test abs(E_jl - E_py) < 1e-5

    # Test again with unordered species in parameter file
    ase_calc = ASECalculator(eam.EAM(potential=alloy))
    julip_calc = EAM(alloy)

    atoms = bulk(:Pd) * 3
    atoms.Z[1:3:end] .= AtomicNumber(:Ag)
    atoms.Z[1:5:end] .= AtomicNumber(:H)
    rattle!(atoms, 0.1)
    E_jl = energy(julip_calc, atoms)
    E_py = energy(ase_calc, atoms)
    @test E_jl ≈ E_py rtol=1e-6
end


import Clp:ClpSolver
import ECOS:ECOSSolver
import GLPK
import GLPKMathProgInterface:GLPKSolverLP
import Ipopt:IpoptSolver
import Tulip:TulipSolver

import CPLEX:CplexSolver
import Gurobi
import Gurobi:GurobiSolver
import Mosek:MosekSolver

global const GRBENV = Gurobi.Env()  # define this to avoid extra prints

# Open-source solvers
sCLP() = ClpSolver(
    LogLevel=0,
    PresolveType=1,
    SolveType=4,
    MaximumSeconds=TIME_LIMIT
);

sECOS() = ECOSSolver(
    verbose=0
)

sGLPK() = GLPKSolverLP(
#   presolve=true,  # no pre-solve for interior-point (apparently)
    method=:InteriorPoint,
    msg_lev=0,
    # TimeLimit=TIME_LIMIT  # GLPK will ignore time limit for IPM algo
);

sIPOPT() = IpoptSolver(
    print_level=0,
    max_cpu_time=Float64(TIME_LIMIT),
    mehrotra_algorithm="yes"
);

sTLP() = Tulip.TulipSolver(
    verbose=0,
    time_limit=TIME_LIMIT,
    barrier_iter_max=200,
    algo=1
);

# Commercial solvers
sCPX() = CplexSolver(
    CPX_PARAM_SCRIND=0,
    CPX_PARAM_PREIND=0,
    CPX_PARAM_THREADS=1,
    CPX_PARAM_LPMETHOD=4, # 0:default, 1:PS, 2:DS, 4:IPM, 6:Concurrent
    CPX_PARAM_SOLUTIONTYPE=2  # 0: default, 1: basic, 2:no crossover,
)

sGRB() = GurobiSolver(
    GRBENV,
    OutputFlag=0,
    Method=2,
    Presolve=0,
    Threads=1,
    Crossover=0
)

sMSK() = MosekSolver(
    MSK_IPAR_INTPNT_BASIS=0,
    MSK_IPAR_NUM_THREADS=1,
    MSK_IPAR_PRESOLVE_USE=0,
    MSK_IPAR_LOG=0,
    MSK_DPAR_OPTIMIZER_MAX_TIME=TIME_LIMIT
);

solver_name(::ClpSolver) = :CLP
solver_name(::ECOSSolver) = :ECOS
solver_name(::GLPKSolverLP) = :GLPK
solver_name(::IpoptSolver) = :IPOPT
solver_name(::TulipSolver) = :TLP
solver_name(::CplexSolver) = :CPX
solver_name(::GurobiSolver) = :GRB
solver_name(::MosekSolver) = :MSK