local_dir = "/home/awise/Desktop/RvSpectML"
#local_dir = "/storage/work/afw5465/RvSpectML"

using Pkg
 Pkg.activate(joinpath(local_dir,"alex_sandbox/"))
 using Query, DataFrames, DataStructures, CSV, JLD, Statistics, StatsBase, Plots, JLD2, FileIO, ArgParse
 using RvSpectMLBase, EchelleInstruments, EchelleCCFs, RvSpectML, Scalpels

 include(joinpath(local_dir,"data_paths.jl"))
 include(joinpath(local_dir,"RvLineList/inputs/param.jl"))
 include(joinpath(local_dir,"RvLineList/src/masks.jl"))

 s = ArgParseSettings()

 @add_arg_table s begin
    "--norm_type"
       help = "normalization type. Options: :raw, :continuum, :blaze"
       arg_type = Symbol
    "--EXPRES_excalibur_only"
       help = "whether or not to only include excalibur wavelengths in expres data - could be replaced with orders_to_use param here"
       arg_type = Bool
   "--mask_type"
      help = "mask shape. Options: :tophat, :gaussian, :supergaussian, :halfcos"
      arg_type = Symbol
   "--mask_scale_factor"
      help = "approximate number of pixels in the width scale of the mask"
      arg_type = Float64
   "--overlap_cutoff"
      help = "distance between lines required for them to be categorized as a blend, expressed as a fraction of the speed of light x 10^-5"
      arg_type = Int
   "--quant"
      help = "quantile for stability of fit params for empirical masks"
      arg_type = String
   "--nbin"
      help = "number of mask subdivisions to make"
      arg_type = Int
   "--bin_n"
      help = "which subdivision from nbin to use - one is first as this is a julia index"
      arg_type = Int
 end

 parsed_args = parse_args(s)

if parsed_args["norm_type"] != nothing
   println("Found command-line arg norm_type. Overwriting param file definition of this arg.")
   global norm_type = parsed_args["norm_type"]
end

if parsed_args["EXPRES_excalibur_only"] != nothing
   println("Found command-line arg EXPRES_excalibur_only. Overwriting param file definition of this arg.")
   global EXPRES_excalibur_only = parsed_args["EXPRES_excalibur_only"]
end

if parsed_args["mask_type"] != nothing
   println("Found command-line arg mask_type. Overwriting param file definition of this arg.")
   global mask_type = parsed_args["mask_type"]
end

if parsed_args["mask_scale_factor"] != nothing
   println("Found command-line arg mask_scale_factor. Overwriting param file definition of this arg.")
   global mask_scale_factor = parsed_args["mask_scale_factor"]
end

if parsed_args["overlap_cutoff"] != nothing
   println("Found command-line arg overlap_cutoff. Overwriting param file definition of this arg.")
   global overlap_cutoff = parsed_args["overlap_cutoff"]
end

if parsed_args["quant"] != nothing
   println("Found command-line arg quant. Overwriting param file definition of this arg.")
   global quant = parsed_args["quant"]
end

if parsed_args["nbin"] != nothing
   println("Found command-line arg nbin. Overwriting param file definition of this arg.")
   global nbin = parsed_args["nbin"]
end

if parsed_args["bin_n"] != nothing
   println("Found command-line arg bin_n. Overwriting param file definition of this arg.")
   global bin_n = parsed_args["bin_n"]
end

function hasOneToOneMatch(X::A1,Y::A2) where { T1<:Real, A1<:AbstractArray{T1,1}, T2<:Real, A2<:AbstractArray{T2,1}}
   xTrue = zeros(Bool,size(X,1))
   for y in 1:length(Y)
      n_matches = 0
      match_i = -1
      for x in 1:length(X)
         if (abs(X[x]-Y[y])/X[x])<1e-5
            n_matches += 1
            match_i = x
         end
      end
      if n_matches == 1
         xTrue[match_i]=true
      end
   end
   return xTrue
end

#norm_type_results = Dict()

ccfs = Dict()

v_grids = Dict()

pipeline_plans = Dict()

rvs = Dict()

times = Dict()

targs = ["101501", "26965", "34411", "10700"]

for targ in targs

   GC.gc()

   global target = targ

   maskWavelengths = string("Reiners",target) #keyword for mask wavelengths - should be e.g. Reiners or Reiners101501
   badLineFilter = string("HD",target)

   linelist_names = SortedDict(Dict())

   push!(linelist_names, "empirical_mask"=>string(target,"_good_lines_fit_quant=0.",quant,"_frac=",min_frac_converged,"_fmt=espresso.mas"))

   VALDmaskStr = string("VALD", depthPercentile ? "" : "_nonDP", "_species=", iron1Only, "_depthcutoff=", depth_cutoff, "_overlapcutoff=", overlap_cutoff, "e-05_allowBlends=", allowBlends, "_badLineFilter=", badLineFilter, "_rejectTelluricSlope=", rejectTelluricSlope, "_waves=", maskWavelengths, "_nbin=", 1, "_binParam=", binParam, "_n=", 0, ".mas")
   push!(linelist_names, "vald_mask"=>VALDmaskStr)

   push!(linelist_names, "espresso"=>"G2.espresso.mas")

   global fits_target_str = target
   verbose=true
   #path_to_EXPRES_data = joinpath("/home/awise/data/",target)
   path_to_EXPRES_data = joinpath(expres_data_path,target)
   default_paths_to_search = [local_dir]
   pipeline_plan = PipelinePlan()
   dont_make_plot!(pipeline_plan, :movie)

   reset_all_needs!(pipeline_plan)
   if need_to(pipeline_plan,:read_spectra)
      if verbose println("# Finding what data files are avaliable.")  end
      global df_files = EXPRES.make_manifest(path_to_EXPRES_data, fits_target_str)

      if verbose println("# Reading in customized parameters from param.jl.")  end
      include(joinpath(local_dir,"param.jl"))

      if verbose println("# Reading in ", size(df_files_use,1), " FITS files.")  end
      @time all_spectra = map(row->EXPRES.read_data(row,store_min_data=true, store_tellurics=true, normalization=norm_type), eachrow(df_files_use))
      #for i in 1:length(all_spectra)
      #  all_spectra[i].metadata[:bjd] = all_spectra[i].metadata[:BARYMJD]
      #end
      #RvSpectMLBase.discard_blaze(all_spectra)
      RvSpectMLBase.discard_continuum(all_spectra)
      dont_need_to!(pipeline_plan,:read_spectra)
   end

   # ccf_mid_velocity = 0.0 #since the empirical mask is at v=0, and we shifted the VALD masks to be at v=0 too, we must reset this global variable to zero.

   pixels_to_use = map(ord->min_col_default(all_spectra[1].inst,ord):max_col_default(all_spectra[1].inst,ord),orders_to_use)
   order_list_timeseries = extract_orders(all_spectra,pipeline_plan, orders_to_use=orders_to_use);
   n_spectra = length(order_list_timeseries)
   n_orders = length(orders_to_use)

   ### NEW RUN: empirical mask
   # ccf_mid_velocity = 0.0 #since the empirical mask is at v=0, and we shifted the VALD masks to be at v=0 too, we must reset this global variable to zero.

   ll = string("empirical_mask")
   #ll = "espresso"
   ll_v = string("vald_mask")

   line_width_50 = 7000.

   need_to!(pipeline_plan,:read_line_list)
   need_to!(pipeline_plan,:clean_line_list_tellurics)
   empty!(pipeline_plan.cache)
   println(linelist_names[ll])
   linelist_path = joinpath(local_dir,"RvLineList","inputs","empirical_masks",linelist_names[ll])
   mask_ll = prepare_line_list(linelist_path, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=0.0,
      Δv_to_avoid_tellurics = 1*RvSpectMLBase.max_bc+5*line_width_50+2*default_ccf_mask_v_width(EXPRES1D()),  orders_to_use=orders_to_use, verbose=true,
      convert_air_to_vacuum=false)

   need_to!(pipeline_plan,:read_line_list)
   need_to!(pipeline_plan,:clean_line_list_tellurics)
   empty!(pipeline_plan.cache)
   println(linelist_names[ll_v])
   linelist_path = joinpath(local_dir,"RvLineList","outputs","masks",linelist_names[ll_v])
   mask_ll_v = prepare_line_list(linelist_path, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=0.0,
      Δv_to_avoid_tellurics = 1*RvSpectMLBase.max_bc+5*line_width_50+2*default_ccf_mask_v_width(EXPRES1D()),  orders_to_use=orders_to_use, verbose=true,
      convert_air_to_vacuum=false)

   nf = length(mask_ll.lambda)
   # hm = zeros(Bool,nf)
   # for i in 1:nf
   #   hm[i] = hasOneMatch(mask_ll.lambda[i],mask_ll_v.lambda)
      #hm[i] = hasOneMatch(linelists[ll].lambda[i]*(1.0-5e3/3e8),linelists[ll_v].lambda)
   hm = hasOneToOneMatch(mask_ll.lambda,mask_ll_v.lambda)
   #end
   mask_empirical_total = mask_ll[hm,:]

   masks_empirical = binMask(mask_empirical_total,nbin)
   mask_empirical = masks_empirical[bin_n]

   println("# mask_scale_factor = ", mask_scale_factor, ", frac_width_to_fit =  ", fwtf, ". ")
   @time (ccfs_empirical, v_grid) = ccf_total(order_list_timeseries, mask_empirical, pipeline_plan,  mask_scale_factor=mask_scale_factor, range_no_mask_change=5*line_width_50, ccf_mid_velocity=0.0, v_step=100,
      #v_max=RvSpectMLBase.max_bc,
      v_max= 0*RvSpectMLBase.max_bc+5*line_width_50, #+2*line_width_50,
      calc_ccf_var=false, recalc=true, mask_type = mask_type)
   push!(ccfs, string("ccfs_empirical_",target) => ccfs_empirical)
   push!(v_grids, string("v_grid_empirical_",target) => v_grid)
   push!(pipeline_plans, string("pipeline_plan_empirical_",target) => pipeline_plan)
   #if nbin == 1 #if nbin > 1, rv fits sometimes fail
   alg_fit_rv =  EchelleCCFs.MeasureRvFromCCFGaussian(frac_of_width_to_fit=fwtf)
   rvs_ccf_empirical = RvSpectML.calc_rvs_from_ccf_total(ccfs_empirical, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true, bin_nightly=true, alg_fit_rv=alg_fit_rv)
   push!(rvs, string("rvs_empirical_",target) => rvs_ccf_empirical)
   #end

   k="espresso"
   need_to!(pipeline_plan,:read_line_list)
   need_to!(pipeline_plan,:clean_line_list_tellurics)
   empty!(pipeline_plan.cache)
   println(linelist_names[k])
   linelist_path = joinpath("alex_sandbox","masks",linelist_names[k])
   line_list_df = prepare_line_list(linelist_path, all_spectra, pipeline_plan, v_center_to_avoid_tellurics=ccf_mid_velocity,
   Δv_to_avoid_tellurics = 2*RvSpectMLBase.max_bc+3*lsf_width, verbose=true)

   masks_espresso = binMask(line_list_df,nbin)
   mask_espresso = masks_espresso[bin_n]

   println("# mask_scale_factor = ", mask_scale_factor, ", frac_width_to_fit =  ", fwtf, ". ")
   @time (ccfs_espresso, v_grid) = ccf_total(order_list_timeseries, mask_espresso, pipeline_plan,  mask_scale_factor=mask_scale_factor, range_no_mask_change=5*line_width_50, ccf_mid_velocity=ccf_mid_velocity, v_step=100,
      #v_max=RvSpectMLBase.max_bc,
      v_max= 0*RvSpectMLBase.max_bc+5*line_width_50, #+2*line_width_50,
      calc_ccf_var=false, recalc=true, mask_type = mask_type)
   push!(ccfs, string("ccfs_espresso_",target) => ccfs_espresso)
   push!(v_grids, string("v_grid_espresso_",target) => v_grid)
   push!(pipeline_plans, string("pipeline_plan_espresso_",target) => pipeline_plan)
   #if nbin == 1 #if nbin > 1, rv fits sometimes fail
   alg_fit_rv =  EchelleCCFs.MeasureRvFromCCFGaussian(frac_of_width_to_fit=fwtf)
   rvs_ccf_espresso = RvSpectML.calc_rvs_from_ccf_total(ccfs_espresso, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true, bin_nightly=true, alg_fit_rv=alg_fit_rv)
   push!(rvs, string("rvs_espresso_",target) => rvs_ccf_espresso)
   #end

   times[target] = order_list_timeseries.times

end

GC.gc()

save(joinpath(local_dir,"RvLineList","outputs",string("param_study_results","_norm_type=",norm_type,
   "_EXPRES_excalibur_only=",EXPRES_excalibur_only,"_mask_type=",mask_type,"_mask_scale_factor=",mask_scale_factor,
   "_overlap_cutoff=",overlap_cutoff,"_quant=",quant,"_nbin=",nbin,"_bin_n=",bin_n,"_orders_to_use=",join(orders_to_use),".jld")),
   "ccfs",ccfs,"v_grids",v_grids,"pipeline_plans",pipeline_plans,"rvs",rvs,"times",times,"orders_to_use",orders_to_use)

"""
sample job submission scripts:

for norm_type in raw blaze continuum; do for EXPRES_excalibur_only in true false; do for mask_type in gaussian tophat; do for mask_scale_factor in 1 2 3 4 5 6 7 8; do for overlap_cutoff in 2; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff expres_mask_param_study.pbs; done; done; done; done; done

for norm_type in raw continuum; do for EXPRES_excalibur_only in true; do for mask_type in gaussian; do for mask_scale_factor in 0.5 1 1.5; do for overlap_cutoff in 1 2 3 4 6; do for quant in 50 70 90; do for nbin in 1; do for bin_n in 0; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff,quant=$quant,nbin=$nbin,n=$n expres_mask_param_study.pbs; done; done; done; done; done; done; done; done

for norm_type in raw continuum; do for EXPRES_excalibur_only in true; do for mask_type in gaussian; do for mask_scale_factor in 1 4 8; do for overlap_cutoff in 1 6; do for quant in 90; do for nbin in 8; do for bin_n in 1 2 3 4 5 6 7 8; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff,quant=$quant,nbin=$nbin,n=$n expres_mask_param_study.pbs; done; done; done; done; done; done; done; done

for i in $(seq 1 50); do for norm_type in raw continuum; do for EXPRES_excalibur_only in true; do for mask_type in gaussian; do for mask_scale_factor in 1; do for overlap_cutoff in 1; do for quant in 90; do for nbin in 1; do for bin_n in 1; do qsub -v norm_type=$norm_type,EXPRES_excalibur_only=$EXPRES_excalibur_only,mask_type=$mask_type,mask_scale_factor=$mask_scale_factor,overlap_cutoff=$overlap_cutoff,quant=$quant,nbin=$nbin,n=$n expres_mask_param_study.pbs; done; done; done; done; done; done; done; done; done

"""


is_directed(::AdjacencyMap{D}) where {D} = D
is_directed(::Type{<:AdjacencyMap{D}}) where {D} = D

###
### has_edge
###
has_edge(g::AdjacencyMap, e::AbstractEdge) = has_edge(g, src(e), dst(e))
function has_edge(g::AdjMap, i::Integer, j::Integer)
    if i < j
        if j <= nvertices(g)
            unsafe_has_edge(getcolptr(g), rowvals(g), i, j)
        else
            return false
        end
    else ## j <= i
        if i <= nvertices(g)
            unsafe_has_edge(getcolptr(g), rowvals(g), j, i)
        else
            return false
        end
    end
end
function has_edge(g::AdjDiMap, i::Integer, j::Integer)
    N = nvertices(g)
    if !((1 <= i <= N) & (1 <= j <= N))
        return false
    else
        unsafe_has_edge(getcolptr(g), rowvals(g), i, j)
    end
end
function unsafe_has_edge(cptr, rvals, i, j)
    coljlastk = @inbounds(cptr[j + 1] - 1)
    searchk = searchsortedfirst(rvals, i, @inbounds(cptr[j]), coljlastk, Base.Order.Forward)
    return searchk <= coljlastk && @inbounds(rvals[searchk]) == i
end

###
### add_edge!
###
add_edge!(g::AdjacencyMap, i::Int, j::Int) = add_edge!(g, i, j, no_metadata)
function add_edge!(g::AdjMap, i::Int, j::Int, val) where {T,M}
    N = nvertices(g)
    if !((1 <= i <= N) & (1 <= j <= N))
        return false
    elseif i < j
        unsafe_add_edge!(g, i, j, val)
    else
        unsafe_add_edge!(g, j, i, val)
    end
end
function add_edge!(g::AdjDiMap, i::Int, j::Int, val) where {T,M}
    N = nvertices(g)
    if !((1 <= i <= N) & (1 <= j <= N))
        return false
    else
        return unsafe_add_edge!(g, i, j, val)
    end
end
function unsafe_add_edge!(g, i, j, val)
    coljlastk = @inbounds(getcolptr(g)[j + 1] - 1)
    searchk = searchsortedfirst(rowvals(g), i, @inbounds(getcolptr(g)[j]), coljlastk, Base.Order.Forward)
    if searchk <= coljlastk && rowvals(g)[searchk] == i
        _setindex!(metadata(g), val, searchk)
        return true
    else
        N = nvertices(g)
        nz = getcolptr(g)[N + 1]
        # if nnz(A) < length(rowval/nzval): no need to grow rowval and preserve values
        _insert!(rowvals(g), searchk, i, nz)
        _insert!(metadata(g), searchk, val, nz)
        @simd for i in (j + 1):(N + 1)
            @inbounds getcolptr(g)[i] += 1
        end
    end
    return true
end
# insert item at position pos, shifting only from pos+1 to nz
_insert!(m, pos, ::NoMetadata, nz) = error()
_insert!(::NoMetadata, pos, ::NoMetadata, nz) = nothing
function _insert!(m, pos, val, nz)
    if nz > length(m)
        insert!(m, pos, val)
    else # nz < length(v)
        Base.unsafe_copyto!(m, pos + 1, m, pos, nz - pos)
        @inbounds(setindex!(m, val, pos))
        m
    end
end
_insert!(m::ListOfMeta, pos, val, nz) = _insert!(parent(m), pos, val, nz)

_setindex!(nzvals, val, i) = @inbounds(setindex!(nzvals, val, i))
_setindex!(m::ListOfMeta, val, i) = _setindex!(parent(m), val, i)
_setindex!(::NoMetadata, ::NoMetadata, i) = nothing
_setindex!(::NoMetadata, val, i) = nothing
_setindex!(nzvals, ::NoMetadata, i) = nothing

###
### rem_edge!
###
rem_edge!(g::AdjDiMap, i::Integer, j::Integer) = unsafe_rem_edge!(g, i, j)
function rem_edge!(g::AdjMap, i::Integer, j::Integer)
    if i < j
        return unsafe_rem_edge!(g, i, j)
    else
        return unsafe_rem_edge!(g, j, i)
    end
end
function unsafe_rem_edge!(g::AdjacencyMap, i::Integer, j::Integer)
    N = nvertices(g)
    if !((1 <= i <= N) && (1 <= j <= N))
        return false
    else
        cptr = getcolptr(g)
        rval = rowvals(g)
        coljlastk = @inbounds(cptr[j + 1]) - 1
        searchk = searchsortedfirst(rval, i, @inbounds(cptr[j]), coljlastk, Base.Order.Forward)
        if searchk > coljlastk || @inbounds(rval[searchk]) != i
            return false
        else
            Awritepos = 1
            oldAcolptrAj = 1
            @inbounds for Aj in OneTo(N)
                for Ak in oldAcolptrAj:(cptr[Aj + 1]-1)
                    if Ak != searchk && Awritepos != Ak
                        # If this element should be kept, rewrite in new position
                        _move_value!(rval, Awritepos, Ak)
                        _move_value!(metadata(g), Awritepos, Ak)
                    end
                    Awritepos += 1
                end
                oldAcolptrAj = Acolptr[Aj + 1]
                Acolptr[Aj + 1] = Awritepos
            end

            # Trim A's storage if necessary
            Annz = Acolptr[end] - 1
            resize!(rval, Annz)
            resize!(m, Annz)
            return true
        end
    end
end

_move_value!(m, d, s) = @inbounds(setindex!(m, m[s], d))
_move_value!(::NoMetadata, d, s) = nothing

function SparseArrays.fkeep!(A::AdjacencyMap, f)
    An = size(A, 2)
    Acolptr = getcolptr(A)
    Arowval = rowvals(A)
    Anzval = metadata(A)

    # Sweep through columns, rewriting kept elements in their new positions
    # and updating the column pointers accordingly as we go.
    Awritepos = 1
    oldAcolptrAj = 1
    @inbounds for Aj in 1:An
        for Ak in oldAcolptrAj:(Acolptr[Aj+1]-1)
            Ai = Arowval[Ak]
            Ax = Anzval[Ak]
            # If this element should be kept, rewrite in new position
            if f(Ai, Aj, Ax)
                if Awritepos != Ak
                    _move_value!(rval, Awritepos, Ak)
                    _move_value!(Anzval, Awritepos, Ak)
                end
                Awritepos += 1
            end
        end
        oldAcolptrAj = Acolptr[Aj+1]
        Acolptr[Aj+1] = Awritepos
    end

    # Trim A's storage if necessary
    Annz = Acolptr[end] - 1
    resize!(Arowval, Annz)
    resize!(Anzval, Annz)

    return A
end

###
### add_vertices!
###
function add_vertices!(g::AdjacencyMap, n::Integer)
    if can_change_size(g)
        return _unsafe_add_vertices!(g, max(0, n))
    else
        return false
    end
end

function _unsafe_add_vertices!(g::AdjacencyMap, n)
    cptr = getcolptr(g)
    c = last(cptr)
    for _ in OneTo(n)
        push!(cptr, c)
    end
    return true
end

function add_vertex!(g::AdjacencyMap)
    cptr = getcolptr(g)
    push!(cptr, last(cptr))
    return true
end

###
### add_vertices!
###
# TODO LightGraphs.rem_vertex!(m::AdjMap, n::Integer)
# TODO LightGraphs.rem_vertices!(m::AdjMap, n::AbstractVector)
# TODO check this a bunch
function rem_vertex!(g::AdjacencyMap, n::Integer)
    fkeep!(g, (i, j, x) -> (i != n && j != n))
    deleteat!(getcolptr(g), n)
end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
"""

use Scenario struct to create reopt.jl model inputs
"""

struct REoptInputs <: AbstractInputs
    techs::Array{String, 1}
    pvtechs::Array{String, 1}
    gentechs::Array{String,1}
    elec_techs::Array{String, 1}
    techs_no_turndown::Array{String, 1}
    min_sizes::DenseAxisArray{Float64, 1}  # (techs)
    max_sizes::DenseAxisArray{Float64, 1}  # (techs)
    existing_sizes::DenseAxisArray{Float64, 1}  # (techs)
    cap_cost_slope::DenseAxisArray{Float64, 1}  # (techs)
    om_cost_per_kw::DenseAxisArray{Float64, 1}  # (techs)
    max_grid_export_kwh::Float64
    elec_load::ElectricLoad
    time_steps::UnitRange
    time_steps_with_grid::Array{Int, 1}
    time_steps_without_grid::Array{Int, 1}
    hours_per_timestep::Float64
    months::UnitRange
    production_factor::DenseAxisArray{Float64, 2}  # (techs, time_steps)
    levelization_factor::DenseAxisArray{Float64, 1}  # (techs)
    VoLL::Array{R, 1} where R<:Real #default set to 1 US dollar per kwh
    pwf_e::Float64
    pwf_om::Float64
    two_party_factor::Float64
    owner_tax_pct::Float64
    offtaker_tax_pct::Float64
    microgrid_premium_pct::Float64
    pvlocations::Array{Symbol, 1}
    maxsize_pv_locations::DenseAxisArray{Float64, 1}  # indexed on pvlocations
    pv_to_location::DenseAxisArray{Int, 2}  # (pv_techs, pvlocations)
    etariff::ElectricTariff
    ratchets::UnitRange
    techs_by_exportbin::DenseAxisArray{Array{String,1}}  # indexed on [:NEM, :WHL, :CUR]
    storage::Storage
    generator::Generator
    elecutil::ElectricUtility
    min_resil_timesteps::Int
    mg_tech_sizes_equal_grid_sizes::Bool
    node::Int
    export_bins_by_tech::Dict
end

"""
    REoptInputs(fp::String)

Use `fp` to load in JSON scenario:
```
function REoptInputs(fp::String)
    s = Scenario(JSON.parsefile(fp))
    REoptInputs(s)
end
```
Useful if you want to manually modify REoptInputs before solving the model.
"""
function REoptInputs(fp::String)
    s = Scenario(JSON.parsefile(fp))
    REoptInputs(s)
end

function REoptInputs(s::Scenario)

    time_steps = 1:length(s.electric_load.loads_kw)
    hours_per_timestep = 8760.0 / length(s.electric_load.loads_kw)
    techs, pvtechs, gentechs, pv_to_location, maxsize_pv_locations, pvlocations, production_factor,
        max_sizes, min_sizes, existing_sizes, cap_cost_slope, om_cost_per_kw  = setup_tech_inputs(s)
    elec_techs = techs  # only modeling electric loads/techs so far
    techs_no_turndown = pvtechs

    max_grid_export_kwh = sum(s.electric_load.loads_kw)

    months = 1:length(s.electric_tariff.monthly_demand_rates)

    techs_by_exportbin = DenseAxisArray([ techs, techs, techs], s.electric_tariff.export_bins)
    # TODO account for which techs have access to export bins (when we add more techs than PV)

    levelization_factor, pwf_e, pwf_om, two_party_factor = setup_present_worth_factors(s, techs, pvtechs)
    # the following hardcoded value for levelization_factor matches the public REopt API value
    # for test_with_cplex (test_time_of_export_rate) and makes the test values match.
    # the REopt code herein uses the Desktop method for levelization_factor, which is more accurate
    # (Desktop has non-linear degradation vs. linear degradation in API)
    # levelization_factor = DenseAxisArray([0.9539], techs)
    # levelization_factor = DenseAxisArray([0.9539, 1.0], techs)  # w/generator
    time_steps_with_grid, time_steps_without_grid, = setup_electric_utility_inputs(s)

    if any(pv.existing_kw > 0 for pv in s.pvs)
        adjust_load_profile(s, production_factor)
    end

    export_bins_by_tech = Dict{String, Array{Symbol, 1}}()
    for t in elec_techs
        export_bins_by_tech[t] = s.electric_tariff.export_bins
    end
    # TODO implement export bins by tech (rather than assuming that all techs share the export_bins)

    REoptInputs(
        techs,
        pvtechs,
        gentechs,
        elec_techs,
        techs_no_turndown,
        min_sizes,
        max_sizes,
        existing_sizes,
        cap_cost_slope,
        om_cost_per_kw,
        max_grid_export_kwh,
        s.electric_load,
        time_steps,
        time_steps_with_grid,
        time_steps_without_grid,
        hours_per_timestep,
        months,
        production_factor,
        levelization_factor,
        typeof(s.financial.VoLL) <: Array{<:Real, 1} ? s.financial.VoLL : fill(s.financial.VoLL, length(time_steps)),
        pwf_e,
        pwf_om,
        two_party_factor,
        s.financial.owner_tax_pct,
        s.financial.offtaker_tax_pct,
        s.financial.microgrid_premium_pct,
        pvlocations,
        maxsize_pv_locations,
        pv_to_location,
        s.electric_tariff,
        1:length(s.electric_tariff.tou_demand_ratchet_timesteps),  # ratchets
        techs_by_exportbin,
        s.storage,
        s.generator,
        s.electric_utility,
        s.site.min_resil_timesteps,
        s.site.mg_tech_sizes_equal_grid_sizes,
        s.site.node,
        export_bins_by_tech
    )
end

function setup_tech_inputs(s::Scenario)

    pvtechs = String[pv.name for pv in s.pvs]
    if length(Base.Set(pvtechs)) != length(pvtechs)
        error("PV names must be unique, got $(pvtechs)")
    end

    techs = copy(pvtechs)
    gentechs = String[]
    if s.generator.max_kw > 0
        push!(techs, "Generator")
        push!(gentechs, "Generator")
    end

    time_steps = 1:length(s.electric_load.loads_kw)

    # REoptInputs indexed on techs:
    max_sizes = DenseAxisArray{Float64}(undef, techs)
    min_sizes = DenseAxisArray{Float64}(undef, techs)
    existing_sizes = DenseAxisArray{Float64}(undef, techs)
    cap_cost_slope = DenseAxisArray{Float64}(undef, techs)
    om_cost_per_kw = DenseAxisArray{Float64}(undef, techs)
    production_factor = DenseAxisArray{Float64}(undef, techs, time_steps)

    # PV specific arrays
    pvlocations = [:roof, :ground, :both]
    pv_to_location = DenseAxisArray{Int}(undef, pvtechs, pvlocations)
    maxsize_pv_locations = DenseAxisArray([1.0e5, 1.0e5, 1.0e5], pvlocations)
    # default to large max size per location. Max size by roof, ground, both

    if !isempty(pvtechs)
        setup_pv_inputs(s, max_sizes, min_sizes, existing_sizes, cap_cost_slope, om_cost_per_kw, production_factor,
                        pvlocations, pv_to_location, maxsize_pv_locations)
    end

    if "Generator" in techs
        setup_gen_inputs(s, max_sizes, min_sizes, existing_sizes, cap_cost_slope, om_cost_per_kw, production_factor)
    end

    return techs, pvtechs, gentechs, pv_to_location, maxsize_pv_locations, pvlocations, production_factor,
    max_sizes, min_sizes, existing_sizes, cap_cost_slope, om_cost_per_kw
end

function setup_pv_inputs(s::Scenario, max_sizes, min_sizes,
    existing_sizes, cap_cost_slope, om_cost_per_kw, production_factor,
    pvlocations, pv_to_location, maxsize_pv_locations)

    time_steps = 1:length(s.electric_load.loads_kw)

    pv_roof_limited, pv_ground_limited, pv_space_limited = false, false, false
    roof_existing_pv_kw, ground_existing_pv_kw, both_existing_pv_kw = 0.0, 0.0, 0.0
    roof_max_kw, land_max_kw = 1.0e5, 1.0e5

    for pv in s.pvs
        production_factor[pv.name, :] = prodfactor(pv, s.site.latitude, s.site.longitude)
        for location in pvlocations
            if pv.location == location
                pv_to_location[pv.name, location] = 1
            else
                pv_to_location[pv.name, location] = 0
            end
        end

        beyond_existing_kw = pv.max_kw
        if pv.location == "both"
            both_existing_pv_kw += pv.existing_kw
            if !(s.site.roof_squarefeet === nothing) && !(s.site.land_acres === nothing)
                # don"t restrict unless both land_area and roof_area specified,
                # otherwise one of them is "unlimited"
                roof_max_kw = s.site.roof_squarefeet * pv.kw_per_square_foot
                land_max_kw = s.site.land_acres / pv.acres_per_kw
                beyond_existing_kw = min(roof_max_kw + land_max_kw, beyond_existing_kw)
                pv_space_limited = true
            end
        elseif pv.location == "roof"
            roof_existing_pv_kw += pv.existing_kw
            if !(s.site.roof_squarefeet === nothing)
                roof_max_kw = s.site.roof_squarefeet * pv.kw_per_square_foot
                beyond_existing_kw = min(roof_max_kw, beyond_existing_kw)
                pv_roof_limited = true
            end

        elseif pv.location == "ground"
            ground_existing_pv_kw += pv.existing_kw
            if !(s.site.land_acres === nothing)
                land_max_kw = s.site.land_acres / pv.acres_per_kw
                beyond_existing_kw = min(land_max_kw, beyond_existing_kw)
                pv_ground_limited = true
            end
        end

        existing_sizes[pv.name] = pv.existing_kw
        min_sizes[pv.name] = pv.existing_kw + pv.min_kw
        max_sizes[pv.name] = pv.existing_kw + beyond_existing_kw

        cap_cost_slope[pv.name] = effective_cost(;
            itc_basis=pv.cost_per_kw,
            replacement_cost=0.0,
            replacement_year=s.financial.analysis_years,
            discount_rate=s.financial.owner_discount_pct,
            tax_rate=s.financial.owner_tax_pct,
            itc=pv.total_itc_pct,
            macrs_schedule = pv.macrs_option_years == 7 ? s.financial.macrs_seven_year : s.financial.macrs_five_year,
            macrs_bonus_pct=pv.macrs_bonus_pct,
            macrs_itc_reduction = pv.macrs_itc_reduction,
            rebate_per_kw = pv.total_rebate_per_kw
        )

        om_cost_per_kw[pv.name] = pv.om_cost_per_kw
    end

    if pv_roof_limited
        maxsize_pv_locations[:roof] = float(roof_existing_pv_kw + roof_max_kw)
    end
    if pv_ground_limited
        maxsize_pv_locations[:ground] = float(ground_existing_pv_kw + land_max_kw)
    end
    if pv_space_limited
        maxsize_pv_locations[:both] = float(both_existing_pv_kw + roof_max_kw + land_max_kw)
    end

    return nothing
end

function setup_gen_inputs(s::Scenario, max_sizes, min_sizes, existing_sizes,
    cap_cost_slope, om_cost_per_kw, production_factor)

    max_sizes["Generator"] = s.generator.max_kw
    min_sizes["Generator"] = s.generator.existing_kw + s.generator.min_kw
    existing_sizes["Generator"] = s.generator.existing_kw
    cap_cost_slope["Generator"] = s.generator.cost_per_kw
    om_cost_per_kw["Generator"] = s.generator.om_cost_per_kw
    production_factor["Generator", :] = prodfactor(s.generator)
    return nothing
end

function setup_present_worth_factors(s::Scenario, techs::Array{String, 1}, pvtechs::Array{String, 1})

    lvl_factor = DenseAxisArray{Float64}(undef, techs)
    for (i, tech) in enumerate(pvtechs)
        lvl_factor[tech] = levelization_factor(
            s.financial.analysis_years,
            s.financial.elec_cost_escalation_pct,
            s.financial.offtaker_discount_pct,
            s.pvs[i].degradation_pct  # TODO generalize for any tech (not just pvs)
        )
    end
    if "Generator" in techs
        lvl_factor["Generator"] = 1
    end

    pwf_e = annuity(
        s.financial.analysis_years,
        s.financial.elec_cost_escalation_pct,
        s.financial.offtaker_discount_pct
    )

    pwf_om = annuity(
        s.financial.analysis_years,
        s.financial.om_cost_escalation_pct,
        s.financial.owner_discount_pct
    )

    if s.financial.two_party_ownership
        pwf_offtaker = annuity(s.financial.analysis_years, 0.0, s.financial.offtaker_discount_pct)
        pwf_owner = annuity(s.financial.analysis_years, 0.0, s.financial.owner_discount_pct)
        two_party_factor = (pwf_offtaker * (1 - s.financial.offtaker_tax_pct)) /
                           (pwf_owner * (1 - s.financial.owner_tax_pct))
    else
        two_party_factor = 1.0
    end

    return lvl_factor, pwf_e, pwf_om, two_party_factor
end

function setup_electric_utility_inputs(s::AbstractScenario)
    if s.electric_utility.outage_end_timestep > 0 &&
            s.electric_utility.outage_end_timestep > s.electric_utility.outage_start_timestep
        time_steps_without_grid = Int[i for i in range(s.electric_utility.outage_start_timestep,
                                                    stop=s.electric_utility.outage_end_timestep)]
        if s.electric_utility.outage_start_timestep > 1
            time_steps_with_grid = append!(
                Int[i for i in range(1, stop=s.electric_utility.outage_start_timestep - 1)],
                Int[i for i in range(s.electric_utility.outage_end_timestep + 1,
                                     stop=length(s.electric_load.loads_kw))]
            )
        else
            time_steps_with_grid = Int[i for i in range(s.electric_utility.outage_end_timestep + 1,
                                       stop=length(s.electric_load.loads_kw))]
        end
    else
        time_steps_without_grid = Int[]
        time_steps_with_grid = Int[i for i in range(1, stop=length(s.electric_load.loads_kw))]
    end
    return time_steps_with_grid, time_steps_without_grid
end

function adjust_load_profile(s::Scenario, production_factor::DenseAxisArray)
    if s.electric_load.loads_kw_is_net
        for pv in s.pvs if pv.existing_kw > 0
            s.electric_load.loads_kw .+= pv.existing_kw * production_factor[pv.name, :].data
        end end
    end

    if s.electric_load.critical_loads_kw_is_net
        for pv in s.pvs if pv.existing_kw > 0
            s.electric_load.critical_loads_kw .+= pv.existing_kw * production_factor[pv.name, :].data
        end end
    end
end


function exclude(a::String,b::String)
    a2 = deepcopy(a)
    for idx in 1:length(b)
        i = b[idx:idx]
        if occursin(i,a2)
            a2 = replace(a2,i=>"")
        end
    end
    a2
end
function getdims(iC,iA,iB,A,B)
    sza = size(A)
    szb = size(B)
    szc = []
    for i in 1:length(iA)
        if occursin(iA[i:i],iC)
            push!(szc,sza[i])
        end
    end
    for i in 1:length(iB)
        if occursin(iB[i:i],iC)
            push!(szc,szb[i])
        end
    end
    szc
end
function contract(A,B,iA,iB)
    # figure out size of C, iC
    contract(A,B,1.0,1.0,iA,iB)
end
function contract(A,B,sA,sB,iA,iB)
    iC = exclude(iA,iB)*exclude(iB,iA)
    C = zeros(eltype(A),getdims(iC,iA,iB,A,B)...)
    contract!(C,A,B,1.0,sA,sB,iC,iA,iB)
    C
end

function contract!(C,A,B,iC,iA,iB)
    contract!(C,A,B,1.0,1.0,1.0,iC,iA,iB)
end

function contract!(C,A,B,sC::T,sA::T,sB::T,iC,iA,iB) where T <: AbstractFloat
    contract!(C,A,B,sC,sA,sB,iC,iA,iB,
              Fermi.ComputeEnvironment.interconnect,Fermi.ComputeEnvironment.communicator,
              Fermi.ComputeEnvironment.accelerator,Fermi.ComputeEnvironment.contractor)
end

function contract!(C,A,B, #tensors
                   sC,sA,sB, #scaling factors
                   iC,iA,iB, #indices
                   ic::I where I <: Fermi.Environments.AbstractInterconnect,
                   comm::Fermi.Environments.NoCommunicator,
                   acc::Fermi.Environments.NoAccelerator,
                   contractor::Fermi.Environments.TTGT)
    # transpose
    # transpose
    # GEMM
    # transpose
end

function contract!(C::T1,
                   A::T2,
                   B::T3,  #tensors
                   sC,sA,sB, #scaling factors
                   iC,iA,iB, #indices
                   ic::I where I <: Fermi.Environments.AbstractInterconnect,
                   comm::Fermi.Environments.NoCommunicator,
                   acc::Fermi.Environments.NoAccelerator,
                   contractor::Fermi.Environments.TBLIS) where { T1 <: Union{Array,Fermi.MemTensor},
                                                                 T2 <: Union{Array,Fermi.MemTensor},
                                                                 T3 <: Union{Array,Fermi.MemTensor}}

    Ap = Fermi.data(A)
    Bp = Fermi.data(B)
    Cp = Fermi.data(C)
    _A = TBLIS.TTensor{eltype(Ap)}(Ap,sA)
    _B = TBLIS.TTensor{eltype(Bp)}(Bp,sB)
    _C = TBLIS.TTensor{eltype(Cp)}(Cp,sC)
    TBLIS.mul!(_C,_A,_B,iA,iB,iC)
end

function contract!(C,A,B, #tensors
                   sC,sA,sB, #scaling factors
                   iC,iA,iB, #indices
                   ic::Fermi.Environments.MPI_IC,
                   comm::Fermi.Environments.MPI_World,
                   acc::Fermi.Environments.AMD_GPU,
                   contractor::Fermi.Environments.TTGT)
    # Super special MPI aware GPU code specialized for AMD products
end


# Atomic opterations

const atomic_ops = ((:add!, :+), (:sub!, :-), (:min!, :min), (:max!, :max),
                    (:and!, :&), (:or!, :|), (:xor!, :xor))

# Methods for writing to a WritableGridData grid from . These are
# associative and commutative so that write order does not affect the result.
for (f, op) in atomic_ops
    @eval begin
        @propagate_inbounds ($f)(d::AbstractSimData, x, I...) = ($f)(first(d), x, I...)
        @propagate_inbounds ($f)(d::WritableGridData{<:Any,R}, x, I...) where R = ($f)(d, proc(d), x, I...)
        @propagate_inbounds function ($f)(d::WritableGridData{<:Any,R}, ::Processor, x, I...) where R
            @boundscheck checkbounds(dest(d), I...)
            @inbounds _setdeststatus!(d, x, I...)
            @inbounds dest(d)[I...] = ($op)(dest(d)[I...], x)
        end
        @propagate_inbounds function ($f)(
            d::WritableGridData{<:Any,R}, proc::Union{ThreadedCPU,CPUGPU}, x, I...
        ) where R
            lock(proc)
            ($f)(d, SingleCPU(), x, I...)
            unlock(proc)
        end
    end
end


struct SIGPLAN end
textwidth(::Type{SIGPLAN}) = 7.00137inch
linewidth(::Type{SIGPLAN}) = 3.33461inch
defaultsm(::Type{SIGPLAN}) = 8

scale(::Type{Plots.PyPlotBackend}) = 4.629 * 0.6485500820772881
scale(::Type{Plots.PGFPlotsBackend}) = 4.629  * 0.6485500820772881

dpi(::Type{Plots.PGFPlotsBackend}) = 300
dpi(::Type{Plots.PyPlotBackend}) = 300


"""
Set of Julia functions to assist with interfacing between the Julia MPC and the
Python neural network model (mats).
"""

using PyCall
push!(PyVector(pyimport("sys")."path"), "../scripts")
python_utils = pyimport("python_utils")

struct Node
    x               ::Vector
    y               ::Vector
    θ               ::Vector
    v               ::Vector
    first_timestep  ::Int64
    last_timestep   ::Int64
    type            ::String
    id              ::Int64
end

struct Scene{T}
    # information about agents in scene
    robot        ::Node
    nodes        ::Vector{Node}

    # possibly useful metadata
    dt           ::T
    timesteps    ::Int64
    node_ids     ::Vector{String}
    x_offset     ::Float64
    y_offset     ::Float64
end

mutable struct PredictionSettings
    mats
    hyperparams
    env
    num_modes
end

mutable struct Obstacle
    ps        ::Vector
    active    ::Bool
end

function Scene(env, scene_num::Int64)
    # information about agents in scene
    robot = nothing
    non_robot_nodes = Vector{Node}()
    x_offset = env.scenes[scene_num].x_min
    y_offset = env.scenes[scene_num].y_min
    for node in env.scenes[scene_num].nodes
        if node.is_robot
            x = node.data.data[:, 1] .+ x_offset
            y = node.data.data[:, 2] .+ y_offset
            θ = node.data.data[:, 9]
            v = node.data.data[:, 11]
            first_timestep = node.first_timestep
            last_timestep = node.last_timestep
            type = node.type.name
            id = node.id
            robot = Node(x, y, θ, v, first_timestep, last_timestep, type, id)
        else
            x = node.data.data[:, 1] .+ x_offset
            y = node.data.data[:, 2] .+ y_offset
            θ = []
            v = []
            first_timestep = node.first_timestep
            last_timestep = node.last_timestep
            type = node.type.name
            id = node.id
            push!(non_robot_nodes, Node(x, y, θ, v, first_timestep, last_timestep, type, id))
        end
    end

    # possibly useful metadata
    dt = env.scenes[scene_num].dt
    timesteps = env.scenes[scene_num].timesteps
    node_ids = [string(node.id) for node in non_robot_nodes]

    Scene(robot, non_robot_nodes, dt, timesteps, node_ids, x_offset, y_offset)
end

function load_model(model_path, env; ts=100)
    mats, hyperparams =
        pycall(python_utils.load_model, NTuple{2, PyObject}, model_path, env, ts);

    mats.set_environment(env)
    mats.set_annealing_params()
    python_utils.calculate_scene_graphs(env, hyperparams)
    mats, hyperparams
end

function predicted_dynamics(pred_settings, scene_num, timestep)
    prediction_info, dynamics_dict = python_utils.predict(pred_settings.mats,
                                                          pred_settings.hyperparams,
                                                          pred_settings.env.scenes[scene_num],
                                                          timestep-1,
                                                          pred_settings.num_modes)

    pred_horizon = pred_settings.hyperparams.get("prediction_horizon")
    num_modes = pred_settings.num_modes

    Aps = [[dynamics_dict["A"][ts, mode, :, :] for ts in 1:pred_horizon] for mode in 1:num_modes]
    Bps = [[dynamics_dict["B"][ts, mode, :, :] for ts in 1:pred_horizon] for mode in 1:num_modes]
    gps = [[dynamics_dict["affine_terms"][ts, mode, :, :] for ts in 1:pred_horizon] for mode in 1:num_modes]

    state_dim = 4 #TODO: state dimension is hardcoded and assumed uniform
    q0 = zeros((length(prediction_info) + 1) * state_dim) # the plus one accounts for robot state
    ordered_node_ids = Array{String, 1}(undef, length(prediction_info))
    for (k, v) in prediction_info
        node_idx = v["node_idx"]
        state_vector_idx = node_idx * state_dim + 1
        q0[state_vector_idx:state_vector_idx+state_dim-1] = v["current_state"]
        ordered_node_ids[v["node_idx"]] = split(k, "/")[2]
    end

    #TODO: hard coded indices
    q_robot = [env.scenes[scene_num].robot.data.data[timestep, 1], # x
               env.scenes[scene_num].robot.data.data[timestep, 2], # y
               env.scenes[scene_num].robot.data.data[timestep, 9], # heading
               env.scenes[scene_num].robot.data.data[timestep, 11]] # v
    q0[1:4] = q_robot

    Aps, Bps, gps, q0, ordered_node_ids
end

#TODO: hard coded indices
#TODO: check direction of accel vector so there's nothing weird happening with accel norm
function get_recorded_robot_controls(pred_settings, scene_num, timestep)
    ω = pred_settings.env.scenes[scene_num].robot.data.data[timestep:end, 10]
    a = pred_settings.env.scenes[scene_num].robot.data.data[timestep:end, 12]
    [ω a]
end

#TODO: hardcoded indices
function predict_future_states(q0, us, As, Bs, gs, mode)
    pred_horizon = pred_settings.hyperparams.get("prediction_horizon")
    state_dim = size(As[1][1])[1]
    q_pred = zeros(state_dim, pred_horizon+1)
    q_pred[:, 1] = q0

    for k in 1:pred_horizon
        q_pred[:, k+1] =  gs[mode][k] + As[mode][k] * q_pred[:, k] + Bs[mode][k] * us[k, :]
    end

    q_pred
end

function init_node_obstacles!(scene, vals)
    for node_id in scene.node_ids
        vals.p_obs[node_id] = Obstacle([[[0., 0.] for i = 1:vals.obs_horizon]
            for j = 1:vals.n_modes], 0)
    end
end

#TODO: hardcoded indices
function update_obstacles_from_predictions!(q_pred, node_ids, vals, scene; iter=nothing)
    #####
    obstacle_heading = Dict()
    for (obs_id, obs) in vals.p_obs
        obstacle_heading[obs_id] = []
    end
    #####

    for mode = 1:vals.n_modes
        # node_idx = 2 # start at 2 to avoid robot state
        for (node_id, obs) in vals.p_obs
            if node_id in node_ids
                node_idx = findall(node_ids .== node_id)[1] + 1 #TODO: fix this convoluted mess (+1 from julia vs. python indexing)
                if length(findall(node_ids .== node_id)) > 1
                    println("Check update_obstacles_from_predictions!().")
                end
                row_idx = (node_idx - 1) * 4 + 1
                vals.p_obs[node_id].ps[mode] = [[q_pred[mode][row_idx, k] + scene.x_offset, q_pred[mode][row_idx+1, k] + scene.y_offset] for k = 1:vals.obs_horizon]
                vals.p_obs[node_id].active = 1

                #####
                push!(obstacle_heading[node_id], q_pred[mode][row_idx+2, :])
                #####

                node_idx += 1
            else
                vals.p_obs[node_id].active = 0
            end
        end
    end

    ####
    filename = ""
    if iter == nothing
        filename = "data/obstacle_heading_iteration_1"
    else
        filename = "data/obstacle_heading_iteration_"*string(iter)
    end
    np.save(filename, obstacle_heading)
    ####

end

#TODO: hardcoded indices
function extract_positions_from_state_pred!(q_pred, node_ids, vals)
    num_nodes = length(node_ids)
    horizon = size(q_pred)[2]

    for node_idx = 1:num_nodes
        row_idx = (node_idx - 1) * 4 + 1
        ps_node = []
        for k = 1:horizon
            push!(ps_node, [q_pred[row_idx, k], q_pred[row_idx+1, k]])
        end
        vals.p_obs[node_id] = ps_node
    end
end

"""
Sets positions for static obstacles. The variable p_obs is a list of the
static obstacle position vectors for each "mode"
(i.e., [[x_1, y_1], [x_2, y_2], ... , [x_n_modes, y_n_modes]]).
"""
function add_static_obstacle!(obstacles, vals)
    @assert length(obstacles) <= vals.n_modes
    @assert length(vals.p_obs) < vals.n_obs

    idx = length(vals.p_obs) + 1
    n_obs_modes = length(obstacles)
    vals.p_obs[idx] = [[obstacles[i] for j = 1:vals.obs_horizon] for i = 1:n_obs_modes]
    if n_obs_modes < vals.n_modes
        append!(vals.p_obs[idx], [[[0, 0] for i = 1:vals.obs_horizon]
            for j = 1:(vals.n_modes - n_obs_modes)])
    end
end


using Test

@testset "Windows" begin
    ap33 = AP33Model()
    @test repr(ap33) == "Moxtek AP3.3 model"
    ap5 = AP5Model()
    @test repr(ap5) == "Moxtek AP5 model"
    ap33t = AP33Tabulation()
    @test repr(ap33t) == "Moxtek AP3.3"
    ap5t = AP5Tabulation()
    @test repr(ap5t) == "Moxtek AP5"
    be = Beryllium()
    @test repr(be) == "5.0 μm Be window"
    ac1 = AmptekC1()
    @test repr(ac1) == "Amptek C1"
    ac2 = AmptekC2()
    @test repr(ac2) == "Amptek C2"
    now = NoWindow() # 100% transmission
    @test repr(now) == "No window"

    @test transmission(now, 300.0) == 1.0
    @test isapprox(transmission(ap33, 300.0), 0.197, atol = 0.001)
    @test isapprox(transmission(ap33t, 300.0), 0.342, atol = 0.001)
    @test isapprox(transmission(ap5t, 300.0), 0.132, atol = 0.001)
    @test isapprox(transmission(ap5, 300.0), 0.200, atol = 0.001)
    @test isapprox(transmission(be, 300.0), 3.6e-7, atol = 0.1e-7)
    @test isapprox(transmission(ac1, 300.0), 0.0702, atol = 0.001)
    @test isapprox(transmission(ac2, 300.0), 0.5198, atol = 0.001)

    @test transmission(now, 3000.0) == 1.0
    @test isapprox(transmission(ap33, 3000.0), 0.760, atol = 0.001)
    @test isapprox(transmission(ap33t, 3000.0), 0.772, atol = 0.001)
    @test isapprox(transmission(ap5t, 3000.0), 0.772, atol = 0.001)
    @test isapprox(transmission(ap5, 3000.0), 0.773, atol = 0.001)
    @test isapprox(transmission(be, 3000.0), 0.982, atol = 0.001)
    @test isapprox(transmission(ac1, 3000.0), 0.865, atol = 0.001)
    @test isapprox(transmission(ac2, 3000.0), 0.928, atol = 0.001)

    @test isapprox(transmission(ap33, 30000.0), 0.978, atol = 0.001)
    @test isapprox(transmission(ap33t, 30000.0), 0.993, atol = 0.001)
    @test isapprox(transmission(ap5t, 30000.0), 0.999, atol = 0.001)
    @test isapprox(transmission(ap5, 30000.0), 0.999, atol = 0.001)
    @test isapprox(transmission(be, 30000.0), 0.9999, atol = 0.001)
    @test isapprox(transmission(ac1, 30000.0), 0.9999, atol = 0.0001)
    @test isapprox(transmission(ac2, 30000.0), 0.9999, atol = 0.0001)
end


# Contents:
# Sprite type
# Board type

type Sprite
    y::Int
    x::Int
    is_alive::Bool

    Sprite(ypos,xpos) = new(ypos,xpos,true)
end

function move!(s::Sprite, direction::ASCIIString)
    contains(direction, "n") && (s.y -= 1)
    contains(direction, "s") && (s.y += 1)
    contains(direction, "e") && (s.x += 1)
    contains(direction, "w") && (s.x -= 1)
end

function teleport!(s::Sprite, h, w)
    s.y = rand(1:h)
    s.x = rand(1:w)
end

skip(s::Sprite) = nothing

type Board
    height::Int
    width::Int
    init_num_robots::Int
    live_robots::Int
    wait_mode::Bool
    robots_when_waiting::Int

    #=
    The matrix representations is a pair of 3D matrices (1 4D matrix)
    in which each "slice" shows
    the existence of a sprite, robot, or scrap pile at each coordinate
    e.g. matrix_rep[1,2,5,3] will be 1 iff there is a scrap pile at
    coordinate 2,5 in the first 3D matrix.
    Likewise, matrix_rep[2,2,5,1] and [2,2,5,2] denote the existence
    of a sprite or robot at [2,5] in the second 3D matrix.
    After mutations are finished, switch_active_board is called.
    =#

    matrix_rep::Array{Int,4}
    active::Int
    inactive::Int

    sprite::Sprite

    # Constructors
    Board(h=24,w=60,nrobots=10) = begin
        this = new( h , w , nrobots , nrobots , false , 0 , zeros(Int,2,h,w,3) , 1 , 2 )

        #random initialization of sprite and robots
        rand_coords = sample(1:h*w ,nrobots+1, replace = false)
        this.sprite = Sprite((rand_coords[1]-1) ÷ w + 1, (rand_coords[1]-1) ÷ h + 1)
        this.matrix_rep[this.active,this.sprite.y,this.sprite.x,1] = 1
        for i = 2:length(rand_coords)
            slice(this.matrix_rep,this.active,:,:,2)[rand_coords[i]] = 1
        end
        return this
    end

    Board(nrobots::Int) = Board(24,60,nrobots)
end

has_robots(b::Board) = b.live_robots != 0

function robots_chase_sprite!(s::Sprite,old_robot_field::AbstractArray{Int,2}, new_robot_field::AbstractArray{Int,2})
    height, width = size(old_robot_field)
    for y = 1:height, x = 1:width
        if old_robot_field[y,x] == 1
            old_robot_field[y,x] = 0
            new_robot_field[towards(y,s.y) , towards(x,s.x)] += 1
        end
    end
end

function scrap_robots!(robot_field::AbstractArray{Int,2},scrap_field::AbstractArray{Int,2})
    for i = 1:length(robot_field)
        if scrap_field[i] == 1
            robot_field[i] = 0
        elseif robot_field[i] > 1
            robot_field[i] = 0
            scrap_field[i] = 1
        end
    end
end

function copy_scrap_field!(b::Board)
    b.matrix_rep[b.active,:,:,3] = b.matrix_rep[b.inactive,:,:,3]
end

function process_robot_turn!(b::Board)
    old_robot_field = slice(b.matrix_rep,b.active,:,:,2)
    new_robot_field = slice(b.matrix_rep,b.inactive,:,:,2)
    old_scrap_field = slice(b.matrix_rep,b.active,:,:,3)
    new_scrap_field = slice(b.matrix_rep,b.inactive,:,:,3)

    new_scrap_field[:] = old_scrap_field[:]

    robots_chase_sprite!(b.sprite, old_robot_field, new_robot_field)

    scrap_robots!(new_robot_field,new_scrap_field)

    b.live_robots = sum(b.matrix_rep[b.inactive,:,:,2])
end

function switch_active_board!(b::Board)
    b.active, b.inactive = b.inactive, b.active
end

unset_sprite_pos!(b::Board) = b.matrix_rep[b.active,b.sprite.y,b.sprite.x,1] = 0
set_sprite_pos!(b::Board) = b.matrix_rep[b.inactive,b.sprite.y,b.sprite.x,1] = 1

function move_sprite!(b::Board, direction::ASCIIString)
    unset_sprite_pos!(b)
    move!(b.sprite, direction)
    set_sprite_pos!(b)
end

function teleport_sprite!(b::Board)
    unset_sprite_pos!(b)
    teleport!(b.sprite,b.height,b.width)
    set_sprite_pos!(b)
end

function skip_sprite(b::Board) # equivalent to nothing for now
    unset_sprite_pos!(b)
    skip(b.sprite)
    set_sprite_pos!(b)
end

robot_on_square(b::Board,y,x) = b.matrix_rep[b.inactive,y,x,2] == 1
scrap_on_square(b::Board,y,x) = b.matrix_rep[b.inactive,y,x,3] == 1
is_sprite_on_robot(b::Board) = robot_on_square(b,b.sprite.y,b.sprite.x)
is_sprite_on_scrap(b::Board) = scrap_on_square(b,b.sprite.y,b.sprite.x)

function scrap_sprite!(b::Board)
    (is_sprite_on_scrap(b) || is_sprite_on_robot(b)) && (b.sprite.is_alive = false)
end

function enter_wait_mode!(b::Board)
    b.wait_mode = true
    b.robots_when_waiting = b.live_robots
end

function is_inbounds(b::Board, y ,x)
    (y < 1 || y > b.height) && return false
    (x < 1 || x > b.width) && return false
    return true
end

function is_valid(m::Char, b::Board)
    y = b.sprite.y
    x = b.sprite.x
    if m == 'h'
        checky, checkx = y,x-1
    elseif m == 'j'
        checky, checkx = y+1,x
    elseif m == 'k'
        checky, checkx = y-1,x
    elseif m == 'l'
        checky, checkx = y,x+1
    elseif m == 'y'
        checky, checkx = y-1,x-1
    elseif m == 'u'
        checky, checkx = y-1,x+1
    elseif m == 'b'
        checky, checkx = y+1,x-1
    elseif m == 'n'
        checky, checkx = y+1,x+1
    elseif m == ' '
        return !robot_in_dist_one(b,y,x)
    elseif m == 't'
        return true
    elseif m == 'w'
        return true
    else
        return false
    end

    #=if is_inbounds(b,checky, checkx) && !robot_in_dist_one(b,checky, checkx) && !scrap_on_square(b,checky, checkx)=#
    #=else=#
        #=@show is_inbounds(b,checky, checkx)=#
        #=@show !robot_in_dist_one(b,checky, checkx)=#
        #=@show !scrap_on_square(b,checky, checkx)=#
    #=end=#
    return is_inbounds(b,checky, checkx) &&
    !robot_in_dist_one(b,checky, checkx) &&
    !scrap_on_square(b,checky, checkx)
end

function robot_in_dist_one(b::Board, y::Int,x::Int)
    for i = -1:1, j = -1:1
        if is_inbounds(b,y+i,x+j)
            if b.matrix_rep[b.active,y+i,x+j,2] == 1
                return true
            end
        end
    end
    return false
end

function get_score(b::Board)
    retVal = 10 * (b.init_num_robots - b.live_robots)
    if b.live_robots == 0 && b.sprite.is_alive && b.wait_mode
        retVal += b.robots_when_waiting - b.live_robots
    end
    retVal
end


type TextureAtlas
    rectangle_packer::RectanglePacker
    mapping         ::Dict{Any, Int} # styled glyph to index in sprite_attributes
    index           ::Int
    images          ::Texture{Float16, 2}
    attributes      ::Vector{Vec4f0}
    scale           ::Vector{Vec2f0}
    extent          ::Vector{FontExtent{Float64}}

    function TextureAtlas(initial_size=(4096, 4096))

        images = Texture(fill(Float16(0.0), initial_size...), minfilter=:linear, magfilter=:linear)

        new(
            RectanglePacker(SimpleRectangle(0, 0, initial_size...)),
            Dict{Any, Int}(),
            1,
            images,
            Vec4f0[],
            Vec2f0[],
            FontExtent{Float64}[]
        )
    end
end

begin #basically a singleton for the textureatlas
const local TEXTURE_ATLAS = TextureAtlas[]
reset_texture_atlas!() = empty!(TEXTURE_ATLAS)
function get_texture_atlas()
    not_initilized = isempty(TEXTURE_ATLAS)
    if not_initilized
        global DEFAULT_FONT_FACE = newface(assetpath("fonts", "DejaVuSans.ttf"))
        alternatives = [
            "DejaVuSans.ttf",
            "NotoSansCJKkr-Regular.otf",
            "NotoSansCuneiform-Regular.ttf",
            "NotoSansSymbols-Regular.ttf",
            "FiraMono-Medium.ttf"
        ]
        global ALTERNATIVE_FONTS = map(alternatives) do font
            newface(assetpath("fonts", font))
        end
        atlas = push!(TEXTURE_ATLAS, TextureAtlas())[] # initialize only on demand
        for c in '\u0000':'\u00ff' #make sure all ascii is mapped linearly
            insert_glyph!(atlas, c, DEFAULT_FONT_FACE)
        end
        return atlas
    else
        return TEXTURE_ATLAS[]
    end
end
end

function glyph_index!(atlas::TextureAtlas, c::Char, font)
    if FT_Get_Char_Index(font[], c) == 0
        for afont in ALTERNATIVE_FONTS
            if FT_Get_Char_Index(afont[], c) != 0
                font = afont
            end
        end
    end
    if c < '\u00ff' && font == DEFAULT_FONT_FACE # characters up to '\u00ff'(255), are directly mapped for default font
        Int(c)+1
    else #others must be looked up, since they're inserted when used first
        return insert_glyph!(atlas, c, font)
    end
end

glyph_scale!(c::Char) = glyph_scale!(get_texture_atlas(), c, DEFAULT_FONT_FACE)
glyph_uv_width!(c::Char) = glyph_uv_width!(get_texture_atlas(), c, DEFAULT_FONT_FACE)

function glyph_uv_width!(atlas::TextureAtlas, c::Char, font)
    atlas.attributes[glyph_index!(atlas, c, font)]
end
function glyph_scale!(atlas::TextureAtlas, c::Char, font)
    atlas.scale[glyph_index!(atlas, c, font)]
end
function glyph_extent!(atlas::TextureAtlas, c::Char, font)
    atlas.extent[glyph_index!(atlas, c, font)]
end

function bearing(extent)
     Point2f0(extent.horizontal_bearing[1], -(extent.scale[2]-extent.horizontal_bearing[2]))
end
function glyph_bearing!{T}(atlas::TextureAtlas, c::Char, font, scale::T)
    T(bearing(atlas.extent[glyph_index!(atlas, c, font)])) .* scale
end
function glyph_advance!{T}(atlas::TextureAtlas, c::Char, font, scale::T)
    T(atlas.extent[glyph_index!(atlas, c, font)].advance) .* scale
end

insert_glyph!(atlas::TextureAtlas, glyph::Char, font) = get!(atlas.mapping, (glyph, font)) do
    uv, rect, extent, width_nopadd = render(atlas, glyph, font)
    tex_size       = Vec2f0(size(atlas.images))
    uv_start       = Vec2f0(uv.x, uv.y)
    uv_width       = Vec2f0(uv.w, uv.h)
    real_heightpx  = width_nopadd[2]
    halfpadding    = (uv_width - width_nopadd) / 2f0
    real_start     = uv_start + halfpadding # include padding
    relative_start = real_start ./ tex_size # use normalized texture coordinates
    relative_width = (real_start+width_nopadd) ./ tex_size

    uv_offset_width = Vec4f0(relative_start..., relative_width...)
    i               = atlas.index
    push!(atlas.attributes, uv_offset_width)
    push!(atlas.scale, Vec2f0(width_nopadd))
    push!(atlas.extent, extent)
    atlas.index = i+1
    return i
end

function sdistancefield(img, restrict_steps=2)
    w, h = size(img)
    w1, h1 = w, h

    halfpad = 24*(2^restrict_steps) # padd so that after restrict it comes out as roughly 48 pixel
    w, h = w+2halfpad, h+2halfpad #pad this, to avoid cuttoffs
    in_or_out = Bool[begin
        x, y = i-halfpad, j-halfpad
        if checkbounds(Bool, img, x,y)
            img[x,y] >= 1.0
        else
            false
        end
    end for i=1:w, j=1:h]

    sd = sdf(in_or_out)
    for i=1:restrict_steps
        w1, h1 = Images.restrict_size(w1), Images.restrict_size(h1)
        sd = Images.restrict(sd) #downsample
    end
    sz = Vec2f0(size(img))
    maxlen = norm(sz)
    sw, sh = size(sd)

    Float16[clamp(sd[i,j]/maxlen, -1, 1) for i=1:sw, j=1:sh], Vec2f0(w1, h1), (2^restrict_steps)
end

function GLAbstraction.render(atlas::TextureAtlas, glyph::Char, font)
    #select_font_face(cc, font)
    if glyph == '\n' # don't render  newline
        glyph = ' '
    end
    bitmap, extent = renderface(font, glyph, (164, 164))
    restrict_steps=2
    sd, width_nopadd, scaling_factor = sdistancefield(bitmap, restrict_steps)
    extent = extent ./ Vec2f0(2^restrict_steps)
    rect = SimpleRectangle(0, 0, size(sd)...)
    uv   = push!(atlas.rectangle_packer, rect) #find out where to place the rectangle
    uv == nothing && error("texture atlas is too small. Resizing not implemented yet. Please file an issue at GLVisualize if you encounter this") #TODO resize surface
    atlas.images[uv.area] = sd
    uv.area, rect, extent, width_nopadd
end


# vinicius- creation of the system directory
# Legend of the systems: 020 = 0.5, 0.35 = 1.0, 0.50 = 1.5, 0.65 = 2.0, 0.80 = 2.5, 0.95 = 3.0

  include("input_gen.jl")
" Example :: EMIMDCABF4 - The system will contain the ions EMIM (two mols), DCA and BF4 (each one with one mole)"
  function folder(name::String,MM1,MM2)
    concs  = ["0.50", "1.00", "1.50", "2.00", "2.50", "3.00"]
    replic = ["00", "01", "02", "03", "04", "05", "06", "07", "08", "09", "10", "11","12", "13", "14", "15", "16", "17", "18", "19", "20"]
     gmds = ["gmd_an1.inp","gmd_an2.inp","gmd_sol.inp", "gmd_ca1.inp"]

    println("entrou")

    main_dir = pwd()
    for c in concs
      for r in replic

        work_dir = joinpath(main_dir,name,c,r)
        cil = parse(Float64,c)
        cd(work_dir)
        conc(MM1,"$(name[1:4])","$(name[5:7])",MM2,"EMIM","$(name[8:10])",cil)

        for file in gmds
          w = open(file,"r+")
          new = open("gmdx_$(file[5:7]).inp","w")

          # change names
          for line in eachline(w)
            if occursin("AN1",line)
              if "$(name[5:7])" == "DCA"
                println(new,replace(line, "AN1" => "NC", count=1))
              else
                println(new,replace(line, "AN1" => "$(name[8:10])", count=1))
              end
            elseif occursin("dist",line)
              println(new,replace(line, "dist" => "20", count=1))
            elseif occursin("AN2",line)
              if "$(name[8:10])" == "DCA"
                println(new,replace(line, "AN2" => "NC", count=1))
              else
                println(new,replace(line, "AN2" => "$(name[8:10])", count=1))
              end
            elseif  occursin("CA1",line)
              println(new,replace(line, "CA1" => "$(name[1:3])", count=1))
            else
              println(new,line)
            end

          end
          close(new)
          close(w)

          rm(file)

        end

        cd(main_dir)
      end
    end
  end

  folder("EMIMDCABF4",177.21, 197.97)


# Source/receiver geometry structure
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

export Geometry, compareGeometry, GeometryIC, GeometryOOC, get_nsrc, n_samples

abstract type Geometry{T} end

function rotateXY!(x::Array{<:Union{Float32,Float64},1}, y::Array{<:Union{Float32,Float64},1}, ϕ::Number)
    for i in 1:length(x)
        ynew = sin(ϕ)*x[i] + cos(ϕ)*y[i]
        x[i] = cos(ϕ)*x[i] - sin(ϕ)*y[i]
        y[i] = ynew
    end
end

const CoordT = Union{Array{T, 1}, Array{Array{T, 1}, 1}} where T<:Number
(::Type{CoordT})(x::Vector{Any}) = rebuild_maybe_jld(x)

# In-core geometry structure for seismic header information
mutable struct GeometryIC{T} <: Geometry{T}
    xloc::CoordT  # Array of receiver positions (fixed for all experiments)
    yloc::CoordT
    zloc::CoordT
    dt::Array{T,1}
    nt::Array{Integer,1}
    t::Array{T,1}
end

# Out-of-core geometry structure, contains look-up table instead of coordinates
mutable struct GeometryOOC{T} <: Geometry{T}
    container::Array{SegyIO.SeisCon,1}
    dt::Array{T,1}
    nt::Array{Integer,1}
    t::Array{T,1}
    nsamples::Array{Integer,1}
    key::String
    segy_depth_key::String

    # h5geo members
    h5geo::PyCall.PyObject
    crs::String
    lengthUnits::String
    temporalUnits::String
    model_origin_x::Float64
    model_origin_y::Float64
    model_orientation::Float64
end

######################## shapes easy access ################################
get_nsrc(g::GeometryIC) = length(g.xloc)
get_nsrc(g::GeometryOOC) = length(g.container)

n_samples(g::GeometryOOC, ::Info) = sum(g.nsamples)
n_samples(g::GeometryIC, info::Info) = sum([length(g.xloc[j])*g.nt[j] for j=1:info.nsrc])
n_samples(g::GeometryOOC, ::Integer) = sum(g.nsamples)
n_samples(g::GeometryIC, nsrc::Integer) = sum([length(g.xloc[j])*g.nt[j] for j=1:nsrc])

################################################ Constructors ####################################################################

"""
    Geometry
        xloc::Array{Array{T, 1},1}
        yloc::Array{Array{T, 1},1}
        zloc::Array{Array{T, 1},1}
        dt::Array{T,1}
        nt::Array{Integer,1}
        t::Array{T,1}

Geometry structure for seismic sources or receivers. Each field is a cell array, where individual cell entries\\
contain values or arrays with coordinates and sampling information for the corresponding shot position. The \\
first three entries are in meters and the last three entries in milliseconds.

Constructors
============

Only pass `dt` and `n` and automatically set `t`:

    Geometry(xloc, yloc, zloc; dt=[], nt=[])

Pass single array as coordinates/parameters for all `nsrc` experiments:

    Geometry(xloc, yloc, zloc, dt=[], nt=[], nsrc=1)

Create geometry structure for either source or receivers from a SegyIO.SeisBlock object.\\
`segy_depth_key` is the SegyIO keyword that contains the depth coordinate and `key` is \\
set to either `source` for source geometry or `receiver` for receiver geometry:

    Geometry(SeisBlock; key="source", segy_depth_key="")

Create geometry structure for from a SegyIO.SeisCon object (seismic data container):

    Geometry(SeisCon; key="source", segy_depth_key="")

Examples
========

(1) Set up receiver geometry for 2D experiment with 4 source locations and 80 fixed receivers:

    xrec = range(100,stop=900,length=80)
    yrec = range(0,stop=0,length=80)
    zrec = range(50,stop=50,length=80)
    dt = 4f0
    t = 1000f0

    rec_geometry = Geometry(xrec, yrec, zrec; dt=dt, t=t, nsrc=4)

(2) Set up corresponding source geometry (coordinates can be of type `linspace` or regular arrays):

    xsrc = [200,400,600,800]
    ysrc = [0,0,0,0]
    zsrc = [50,50,50,50]

    src_geometry = Geometry(xsrc, ysrc, zsrc; dt=dt, t=t, nsrc=4)

(3) Read source and receiver geometries from SEG-Y file:

    using SegyIO
    seis_block = segy_read("test_file.segy")
    rec_geometry = Geometry(seis_block; key="receiver", segy_depth_key="RecGroupElevation")
    src_geometry = Geometry(seis_block; key="source", segy_depth_key="SourceDepth")

Check the seis_block's header entries to findall out which keywords contain the depth coordinates.\\
The source depth keyword is either `SourceDepth` or `SourceSurfaceElevation`. The receiver depth \\
keyword is typically `RecGroupElevation`.

(4) Read source and receiver geometries from out-of-core SEG-Y files (for large data sets). Returns an out-of-core \\
geometry object `GeometryOOC` without the source/receiver coordinates, but a lookup table instead:

    using SegyIO
    seis_container = segy_scan("/path/to/data/directory","filenames",["GroupX","GroupY","RecGroupElevation","SourceDepth","dt"])
    rec_geometry = Geometry(seis_container; key="receiver", segy_depth_key="RecGroupElevation")
    src_geometry = Geometry(seis_container; key="source", segy_depth_key="SourceDepth")

"""
Geometry(xloc::CoordT, yloc::CoordT, zloc::CoordT, dt::Array{T,1}, nt::Array{Integer,1}, t::Array{T,1}) where T = GeometryIC{T}(xloc,yloc,zloc,dt,nt,t)

# Fallback constructors for non standard input types
function Geometry(xloc, yloc, zloc; dt=[], t=[], nsrc=nothing)
    if any(typeof(x) <: StepRangeLen for x=[xloc, yloc, zloc])
        args = [typeof(x) <: StepRangeLen ? collect(x) : x for x=[xloc, yloc, zloc]]
        isnothing(nsrc) && (return Geometry(args...; dt=dt, t=t))
        return Geometry(args...; dt=dt, t=t, nsrc=nsrc)
    end
    isnothing(nsrc) && (return Geometry(tof32(xloc), tof32(yloc), tof32(zloc); dt=dt, t=t))
    return Geometry(tof32(xloc), tof32(yloc), tof32(zloc); dt=dt, t=t, nsrc=nsrc)
end

# Constructor if nt is not passed
function Geometry(xloc::Array{Array{T, 1},1}, yloc::CoordT, zloc::Array{Array{T, 1},1};dt=[],t=[]) where T
    nsrc = length(xloc)
    # Check if single dt was passed
    dtCell = typeof(t) <: Real ? [T(dt) for j=1:nsrc] : T.(dt)
    # Check if single t was passed
    tCell = typeof(t) <: Real ? [T(t) for j=1:nsrc] : T.(t)

    # Calculate number of time steps
    ntCell = typeof(t) <: Real ? [floor(Int, t / dt) + 1 for j=1:nsrc] : floor.(Int, tCell ./ dtCell) .+ 1
    return GeometryIC{T}(xloc, yloc, zloc, dtCell, ntCell, tCell)
end

# Constructor if coordinates are not passed as a cell arrays
function Geometry(xloc::Array{T, 1}, yloc::CoordT, zloc::Array{T, 1}; dt=[], t=[], nsrc::Integer=1) where T
    xlocCell = [xloc for j=1:nsrc]
    ylocCell = [yloc for j=1:nsrc]
    zlocCell = [zloc for j=1:nsrc]
    dtCell = [T(dt) for j=1:nsrc]
    ntCell = [floor(Int, t/dt)+1 for j=1:nsrc]
    tCell = [T(t) for j=1:nsrc]
    return GeometryIC{T}(xlocCell, ylocCell, zlocCell, dtCell, ntCell, tCell)
end

################################################ Constructors from SEGY data  ####################################################

# Set up source geometry object from in-core data container
function Geometry(data::SegyIO.SeisBlock; key="source", segy_depth_key="")
    src = get_header(data,"FieldRecord")
    nsrc = length(unique(src))
    if key=="source"
        isempty(segy_depth_key) && (segy_depth_key="SourceSurfaceElevation")
        params = ["SourceX","SourceY",segy_depth_key]
        gt = Float32
    elseif key=="receiver"
        isempty(segy_depth_key) && (segy_depth_key="RecGroupElevation")
        params = ["GroupX","GroupY",segy_depth_key]
        gt = Array{Float32, 1}
    else
        throw("Specified keyword not supported")
    end
    xloc = Array{gt, 1}(undef, nsrc)
    yloc = Array{gt, 1}(undef, nsrc)
    zloc = Array{gt, 1}(undef, nsrc)
    dt = Array{Float32}(undef, nsrc)
    nt = Array{Integer}(undef, nsrc)
    t = Array{Float32}(undef, nsrc)

    xloc_full = get_header(data, params[1])
    yloc_full = get_header(data, params[2])
    zloc_full = get_header(data, params[3])
    dt_full = get_header(data, "dt")[1]
    nt_full = get_header(data, "ns")[1]

    for j=1:nsrc
        traces = findall(src .== unique(src)[j])
        if key=="source"    # assume same source location for all traces within one shot record
            xloc[j] = convert(gt, xloc_full[traces][1])
            yloc[j] = convert(gt,yloc_full[traces][1])
            zloc[j] = abs.(convert(gt,zloc_full[traces][1]))
        else
            xloc[j] = convert(gt, xloc_full[traces])
            yloc[j] = convert(gt, yloc_full[traces])
            zloc[j] = abs.(convert(gt, zloc_full[traces]))
        end
        dt[j] = Float32(dt_full/1f3)
        nt[j] = convert(Integer,nt_full)
        t[j] =  Float32((nt[j]-1)*dt[j])
    end

    if key == "source"
        xloc = convertToCell(xloc)
        yloc = convertToCell(yloc)
        zloc = convertToCell(zloc)
    end
    return GeometryIC{Float32}(xloc,yloc,zloc,dt,nt,t)
end

# Set up geometry summary from out-of-core data container
function Geometry(data::SegyIO.SeisCon; key="source", segy_depth_key="")

    if key=="source"
        isempty(segy_depth_key) && (segy_depth_key="SourceSurfaceElevation")
    elseif key=="receiver"
        isempty(segy_depth_key) && (segy_depth_key="RecGroupElevation")
    else
        throw("Specified keyword not supported")
    end

    # read either source or receiver geometry
    nsrc = length(data)
    container = Array{SegyIO.SeisCon}(undef, nsrc)
    dt = Array{Float32}(undef, nsrc)
    nt = Array{Integer}(undef, nsrc)
    t = Array{Float32}(undef, nsrc)
    nsamples = Array{Integer}(undef, nsrc)
    for j=1:nsrc
        container[j] = split(data,j)
        dt[j] = data.blocks[j].summary["dt"][1]/1f3
        nt[j] = data.ns
        t[j] = (nt[j]-1)*dt[j]
        key=="source" ? nsamples[j] = data.ns : nsamples[j] = Int((data.blocks[j].endbyte - data.blocks[j].startbyte)/(240 + data.ns*4)*data.ns)
    end
    return  GeometryOOC{Float32}(container,dt,nt,t,nsamples,key,segy_depth_key,nothing,"","","",0.0,0.0,0.0)
end

# Set up geometry summary from out-of-core data container passed as cell array
function Geometry(data::Array{SegyIO.SeisCon,1}; key="source", segy_depth_key="")

    if key=="source"
        isempty(segy_depth_key) && (segy_depth_key="SourceSurfaceElevation")
    elseif key=="receiver"
        isempty(segy_depth_key) && (segy_depth_key="RecGroupElevation")
    else
        throw("Specified keyword not supported")
    end

    nsrc = length(data)
    container = Array{SegyIO.SeisCon}(undef, nsrc)
    dt = Array{Float32}(undef, nsrc); nt = Array{Integer}(undef, nsrc); t = Array{Float32}(undef, nsrc)
    nsamples = Array{Integer}(undef, nsrc)
    for j=1:nsrc
        container[j] = data[j]
        dt[j] = data[j].blocks[1].summary["dt"][1]/1f3
        nt[j] = data[j].ns
        t[j] = (nt[j]-1)*dt[j]
        key=="source" ? nsamples[j] = data[j].ns : nsamples[j] = Int((data[j].blocks[1].endbyte - data[j].blocks[1].startbyte)/(240 + data[j].ns*4)*data[j].ns)
    end
    return  GeometryOOC{Float32}(container,dt,nt,t,nsamples,key,segy_depth_key,nothing,"","","",0.0,0.0,0.0)
end

# Load geometry from out-of-core Geometry container
function Geometry(geometry::GeometryOOC)
    nsrc = length(geometry.container)

    # read either source or receiver geometry
    if geometry.key=="source"
        params = ["SourceX","SourceY",geometry.segy_depth_key,"dt","ns"]
        gt = Float32
    elseif geometry.key=="receiver"
        params = ["GroupX","GroupY",geometry.segy_depth_key,"dt","ns"]
        gt = Array{Float32, 1}
    else
        throw("Specified keyword not supported")
    end
    xloc = Array{gt, 1}(undef, nsrc)
    yloc = Array{gt, 1}(undef, nsrc)
    zloc = Array{gt, 1}(undef, nsrc)
    dt = Array{Float32}(undef, nsrc); nt = Array{Integer}(undef, nsrc); t = Array{Float32}(undef, nsrc)

    for j=1:nsrc
        header = read_con_headers(geometry.container[j], params, 1)
        if geometry.key=="source"
            xloc[j] = convert(gt, get_header(header, params[1])[1])
            yloc[j] = convert(gt, get_header(header, params[2])[1])
            zloc[j] = abs.(convert(gt,get_header(header, params[3])[1]))
        else
            xloc[j] = convert(gt, get_header(header, params[1]))
            yloc[j] = convert(gt, get_header(header, params[2]))
            zloc[j] = abs.(convert(gt, get_header(header, params[3])))
        end
        if isnothing(geometry.h5geo)
            dt[j] = get_header(header, params[4])[1]/1f3
        else
            dt[j] = get_header(header, params[4])[1]
        end
        nt[j] = convert(Integer, get_header(header, params[5])[1])
        t[j] =  (nt[j]-1)*dt[j]
    end
    if geometry.key == "source"
        xloc = convertToCell(xloc)
        yloc = convertToCell(yloc)
        zloc = convertToCell(zloc)
    end

    if isnothing(geometry.h5geo)
        return GeometryIC(xloc,yloc,zloc,dt,nt,t)
    end

    dt = geometry.h5geo.sr.convertUnits(dt, geometry.temporalUnits, "ms")
    t = geometry.h5geo.sr.convertUnits(t, geometry.temporalUnits, "ms")
    if isnothing(dt) || length(dt) < 1 || isnothing(t) || length(t) < 1
        @error "Unable to convert temporal units"
    end

    for i in 1:length(xloc)
        xloc[i], yloc[i] = geometry.h5geo.sr.transformCoordFrom(
                            xloc[i], yloc[i], geometry.lengthUnits, geometry.crs)
        xloc[i] -= geometry.model_origin_x
        yloc[i] -= geometry.model_origin_y
        rotateXY!(xloc[i], yloc[i], geometry.model_orientation)
        xloc[i] += geometry.model_origin_x
        yloc[i] += geometry.model_origin_y
        xloc[i] = geometry.h5geo.sr.convertUnits(xloc[i], geometry.lengthUnits, "m")
        yloc[i] = geometry.h5geo.sr.convertUnits(yloc[i], geometry.lengthUnits, "m")
    end

    return GeometryIC(xloc,yloc,zloc,dt,nt,t)
end

Geometry(geometry::GeometryIC) = geometry
Geometry(::Nothing) = nothing

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

# Subsample in-core geometry structure
function subsample(geometry::GeometryIC,srcnum)
    if length(srcnum)==1
        srcnum = srcnum[1]
        geometry = Geometry(geometry.xloc[srcnum], geometry.yloc[srcnum], geometry.zloc[srcnum];
                            dt=geometry.dt[srcnum],t=geometry.t[srcnum],nsrc=1)
    else
        geometry = Geometry(geometry.xloc[srcnum], geometry.yloc[srcnum], geometry.zloc[srcnum],
                            geometry.dt[srcnum], geometry.nt[srcnum], geometry.t[srcnum])
    end
    return geometry
end

# Subsample out-of-core geometry structure
subsample(geometry::GeometryOOC, srcnum) = Geometry(geometry.container[srcnum]; key=geometry.key, segy_depth_key=geometry.segy_depth_key)

# Compare if geometries match
function compareGeometry(geometry_A::Geometry, geometry_B::Geometry)
    if isequal(geometry_A.xloc, geometry_B.xloc) && isequal(geometry_A.yloc, geometry_B.yloc) && isequal(geometry_A.zloc, geometry_B.zloc) &&
    isequal(geometry_A.dt, geometry_B.dt) && isequal(geometry_A.nt, geometry_B.nt)
        return true
    else
        return false
    end
end

isequal(geometry_A::Geometry, geometry_B::Geometry) = compareGeometry(geometry_A, geometry_B)

function compareGeometry(geometry_A::GeometryOOC, geometry_B::GeometryOOC)
    check = true
    for j=1:length(geometry_A.container)
        if ~isequal(geometry_A.container[j].blocks[1].summary["GroupX"], geometry_B.container[j].blocks[1].summary["GroupX"]) ||
        ~isequal(geometry_A.container[j].blocks[1].summary["GroupY"], geometry_B.container[j].blocks[1].summary["GroupY"]) ||
        ~isequal(geometry_A.container[j].blocks[1].summary["SourceX"], geometry_B.container[j].blocks[1].summary["SourceX"]) ||
        ~isequal(geometry_A.container[j].blocks[1].summary["SourceY"], geometry_B.container[j].blocks[1].summary["SourceY"]) ||
        ~isequal(geometry_A.container[j].blocks[1].summary["dt"], geometry_B.container[j].blocks[1].summary["dt"])
            check = false
        end
    end
    return check
end

isequal(geometry_A::GeometryOOC, geometry_B::GeometryOOC) = compareGeometry(geometry_A, geometry_B)

compareGeometry(geometry_A::GeometryOOC, geometry_B::Geometry) = true
compareGeometry(geometry_A::Geometry, geometry_B::GeometryOOC) = true

for G in [GeometryOOC, GeometryIC]
    @eval function push!(G1::$G, G2::$G)
        for k in fieldnames($G)
            pushfield!(getfield(G1, k), getfield(G2, k))
        end
    end
end

pushfield!(a::Array, b::Array) = append!(a, b)
pushfield!(a, b) = nothing


import .ChainRulesCore

if isdefined(ChainRulesCore, :ZeroTangent)
  const ChainRulesZero = ChainRulesCore.ZeroTangent
else
  const ChainRulesZero = ChainRulesCore.Zero
end

function ChainRulesCore.rrule(::typeof(tanh_fast), x)
    t = tanh_fast(x)
    ∂ = let t = t
        y -> (ChainRulesZero(), mul_fast(vfnmadd_fast(t, t, one(t)), y))
    end
    t, ∂
end
function ChainRulesCore.rrule(::typeof(sigmoid_fast), x)
    s = sigmoid_fast(x)
    ∂ = let s = s
        y -> (ChainRulesZero(), mul_fast(vfnmadd_fast(s, s, s), y))
    end
    s, ∂
end
function ChainRulesCore.rrule(::typeof(relu), v)
    z = zero(v)
    cmp = v < z
    r = ifelse(cmp, z, v)
    ∂ = let cmp = cmp
        y -> (ChainRulesZero(), ifelse(cmp, zero(y), y))
    end
    r, ∂
end

function ∂vmap_singlethread!(
    f::F, ∂y::Tuple{Vararg{DenseArray{T},A}}, y::DenseArray{T},
    args::Vararg{DenseArray{<:Base.HWReal},A}
) where {F,T <: Base.HWReal, A}
    N = length(y)
    ptry = VectorizationBase.zero_offsets(stridedpointer(y))
    ptrargs = VectorizationBase.zero_offsets.(stridedpointer.(args))
    ptr∂y = VectorizationBase.zero_offsets.(stridedpointer.(∂y))

    i = 0
    V = VectorizationBase.pick_vector_width(T)
    W = Int(V)
    st = VectorizationBase.static_sizeof(T)
    zero_index = MM{W}(StaticInt(0), st)
    while i < vsub_nsw(N, ((W << 2) - 1))
        index = VectorizationBase.Unroll{1,W,4,1,W,0x0000000000000000}((i,))
        v = f(init_dual(vload.(ptrargs, index))...)
        dual_store!(ptr∂y, ptry, v, index)
        i = vadd_nw(i, 4W)
    end
    while i < vsub_nsw(N, (W - 1))
        vᵣ = f(init_dual(vload.(ptrargs, ((MM{W}(i),),)))...)
        dual_store!(ptr∂y, ptry, vᵣ, (MM{W}(i),))
        i = vadd_nw(i, W)
    end
    if i < N
        m = mask(T, N & (W - 1))
        dual_store!(ptr∂y, ptry, f(init_dual(vload.(ptrargs, ((MM{W}(i),),), m))...), (MM{W}(i,),), m)
    end
    nothing
end

struct SIMDMapBack{K,T<:Tuple{Vararg{Any,K}}}
    jacs::T
end
@generated function (b::SIMDMapBack{K,T})(Δ::A) where {K,T,A}
    preloop = Expr(:block, :(jacs = b.jacs))
    loop_body = Expr(:block, :(Δᵢ = Δ[i]))
    ret = Expr(:tuple, ChainRulesZero(), ChainRulesZero())
    for k ∈1:K
        jₖ = Symbol(:j_, k)
        push!(preloop.args, :($jₖ = jacs[$k]))
        push!(loop_body.args, :($jₖ[i] *= Δᵢ))
        push!(ret.args, jₖ)
    end
    quote
        $preloop
        @turbo for i ∈ eachindex(Δ)
            $loop_body
        end
        $ret
    end
end

function ChainRulesCore.rrule(::typeof(vmap), f::F, args::Vararg{Any,K}) where {F,K}
    out = similar(first(args))
    jacs = map(similar, args)
    ∂vmap_singlethread!(f, jacs, out, args...)
    out, SIMDMapBack(jacs)
end


module UMICORE
using ..AGFFileReader: Glass, GlassID, AGF
export GASIR1, GASIR2

const GASIR1 = Glass(GlassID(AGF, 1774), 3, 2.50730798, 0.0761612483, 0.0116256625, -0.000129196515, -6.01958096e-8, -1.78603241e-10, 0.0, 0.0, 0.0, 0.0, 2.0, 14.0, 4.9751e-5, 1.5398e-7, 0.0, 0.0, 0.0, 0.0, 20.0, 0.0, -1.0, -1.0, 17.0, -1.0, 0, -1.0, [(2.0, 0.9998, 2.0), (4.0, 0.9968, 2.0), (6.0, 0.9954, 2.0), (8.0, 0.9941, 2.0), (10.0, 0.9926, 2.0), (12.0, 0.9567, 2.0), (14.0, 0.9275, 2.0)], 1.0, -1.0, -1.0, 0, 0.0, 0, 4.4, 0)
const GASIR2 = Glass(GlassID(AGF, 1775), 3, 2.60073644, 0.0917946407, 0.0122684939, -0.000165642849, -7.8013525e-8, -1.51225803e-10, 0.0, 0.0, 0.0, 0.0, 2.0, 14.0, 4.6e-5, 0.0, 0.0, 0.0, 0.0, 0.0, 20.0, 0.0, -1.0, -1.0, 16.0, -1.0, 2, -1.0, [(2.0, 0.999, 2.0), (4.0, 0.999, 2.0), (6.0, 0.995, 2.0), (8.0, 0.9934, 2.0), (10.0, 0.9984, 2.0), (12.0, 0.9823, 2.0), (14.0, 0.9724, 2.0)], 1.0, -1.0, -1.0, 0, 0.0, 0, 4.7, 0)
end # module