####################################
# The AbstractBlockArray interface #
####################################

"""
    abstract AbstractBlockArray{T, N} <: AbstractArray{T, N}

The abstract type that represents a blocked array. Types that implement
the `AbstractBlockArray` interface should subtype from this type.

** Typealiases **

* `AbstractBlockMatrix{T}` -> `AbstractBlockArray{T, 2}`

* `AbstractBlockVector{T}` -> `AbstractBlockArray{T, 1}`

* `AbstractBlockVecOrMat{T}` -> `Union{AbstractBlockMatrix{T}, AbstractBlockVector{T}}`
"""
abstract type AbstractBlockArray{T, N} <: LayoutArray{T, N} end
const AbstractBlockMatrix{T} = AbstractBlockArray{T, 2}
const AbstractBlockVector{T} = AbstractBlockArray{T, 1}
const AbstractBlockVecOrMat{T} = Union{AbstractBlockMatrix{T}, AbstractBlockVector{T}}

block2string(b, s) = string(join(map(string,b), '×'), "-blocked ", Base.dims2string(s))
_block_summary(a) = string(block2string(blocksize(a), size(a)), " ", typeof(a))
Base.summary(a::AbstractBlockArray) = _block_summary(a)
_show_typeof(io, a) = show(io, typeof(a))
function _block_summary(io, a)
    print(io, block2string(blocksize(a), size(a)))
    print(io, ' ')
    _show_typeof(io, a)
end
Base.summary(io::IO, a::AbstractBlockArray) = _block_summary(io, a)

# avoid to_shape which complicates axes
Base.similar(a::AbstractBlockArray{T}) where {T}                             = similar(a, T)
Base.similar(a::AbstractBlockArray, ::Type{T}) where {T}                     = similar(a, T, axes(a))
Base.similar(a::AbstractBlockArray{T}, dims::Tuple) where {T}                = similar(a, T, dims)

# If all we know is size, just return an Array which conforms to BlockArray interface
Base.similar(::Type{<:AbstractBlockArray{T,N}}, dims::Dims) where {T,N} = similar(Array{T,N}, dims)

Base.IndexStyle(::Type{<:AbstractBlockArray}) = IndexCartesian()

# need to overload axes to return BlockAxis
@inline size(block_array::AbstractBlockArray) = map(length, axes(block_array))
@inline axes(block_array::AbstractBlockArray) = throw(error("axes for ", typeof(block_array), " is not implemented"))

"""
    BlockBoundsError([A], [inds...])

Thrown when a block indexing operation into a block array, `A`, tried to access an out-of-bounds block, `inds`.
"""
struct BlockBoundsError <: Exception
    a::Any
    i::Any
    BlockBoundsError() = new()
    BlockBoundsError(a::AbstractArray) = new(a)
    BlockBoundsError(a::AbstractArray, @nospecialize(i)) = new(a,i)
end

BlockBoundsError(a::AbstractArray, I::Block) = BlockBoundsError(a, I.n)

function Base.showerror(io::IO, ex::BlockBoundsError)
    print(io, "BlockBoundsError")
    if isdefined(ex, :a)
        print(io, ": attempt to access ")
            print(io, summary(ex.a))
        if isdefined(ex, :i)
            print(io, " at block index [")
            join(io, ex.i, ',')
            print(io, ']')
        end
    end
end

"""
    blockcheckbounds(A, inds...)

Throw a `BlockBoundsError` if the specified block indexes are not in bounds for the given block array.
Subtypes of `AbstractBlockArray` should
specialize this method if they need to provide custom block bounds checking behaviors.

```jldoctest; setup = quote using BlockArrays end
julia> A = BlockArray(rand(2,3), [1,1], [2,1]);

julia> blockcheckbounds(A, 3, 2)
ERROR: BlockBoundsError: attempt to access 2×2-blocked 2×3 BlockMatrix{Float64, Matrix{Matrix{Float64}}, Tuple{BlockedUnitRange{Vector{Int64}}, BlockedUnitRange{Vector{Int64}}}} at block index [3,2]
[...]
```
"""
@inline function blockcheckbounds(A::AbstractArray{T, N}, i::Vararg{Integer, N}) where {T,N}
    if blockcheckbounds(Bool, A, i...)
        return
    else
        throw(BlockBoundsError(A, i))
    end
end

@inline function blockcheckbounds(::Type{Bool}, A::AbstractArray{T, N}, i::Vararg{Integer, N}) where {T,N}
    n = blockaxes(A)
    k = 0
    for idx in 1:N # using enumerate here will allocate
        k += 1
        @inbounds _i = i[idx]
        Block(_i) in n[k] || return false
    end
    for idx = N+1:length(i)
        isone(i[idx]) || return false
    end
    return true
end

blockcheckbounds(A::AbstractArray{T, N}, i::Block{N}) where {T,N} = blockcheckbounds(A, i.n...)
blockcheckbounds(A::AbstractArray{T, N}, i::Vararg{Block{1},N}) where {T,N} = blockcheckbounds(A, Int.(i)...)
blockcheckbounds(A::AbstractVector{T}, i::Block{1}) where {T,N} = blockcheckbounds(A, Int(i))

@propagate_inbounds Base.setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Block{N}) where {T,N} =
    setindex!(block_arr, v, Block.(block.n)...)
@inline @propagate_inbounds function Base.setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Vararg{Block{1}, N}) where {T,N}
    blockcheckbounds(block_arr, block...)
    dest = view(block_arr, block...)
    size(dest) == size(v) || throw(DimensionMismatch(string("tried to assign $(size(v)) array to $(size(dest)) block")))
    copyto!(dest, v)
    block_arr
end

@inline @propagate_inbounds Base.setindex!(block_arr::AbstractBlockArray{T,N}, v, blockindex::BlockIndex{N}) where {T,N} =
    view(block_arr, block(blockindex))[blockindex.α...] = v
@inline @propagate_inbounds Base.setindex!(block_arr::AbstractBlockVector{T}, v, blockindex::BlockIndex{1}) where {T} =
    view(block_arr, block(blockindex))[blockindex.α...] = v
@inline @propagate_inbounds Base.setindex!(block_arr::AbstractBlockArray{T,N}, v, blockindex::Vararg{BlockIndex{1},N}) where {T,N} =
    block_arr[BlockIndex(blockindex)] = v

viewblock(block_arr, block) = Base.invoke(view, Tuple{AbstractArray, Any}, block_arr, block)
@inline Base.view(block_arr::AbstractBlockArray{<:Any,N}, block::Block{N}) where N = viewblock(block_arr, block)
@inline Base.view(block_arr::AbstractBlockArray{<:Any,N}, block::Block{1}) where N = view(block_arr, Block(Base._ind2sub(map(Base.OneTo, blocksize(block_arr)), Int(block))...))
@inline Base.view(block_arr::AbstractBlockVector, block::Block{1}) = viewblock(block_arr, block)
@inline @propagate_inbounds Base.view(block_arr::AbstractBlockArray, block::Block{1}...) = view(block_arr, Block(block))

"""
    eachblock(A::AbstractBlockArray)

Create a generator that iterates over each block of an `AbstractBlockArray`
returning views.

```jldoctest; setup = quote using BlockArrays end
julia> v = Array(reshape(1:6, (2, 3)))
2×3 Matrix{Int64}:
 1  3  5
 2  4  6

julia> A = BlockArray(v, [1,1], [2,1])
2×2-blocked 2×3 BlockMatrix{Int64}:
 1  3  │  5
 ──────┼───
 2  4  │  6

julia> Matrix.(collect(eachblock(A)))
2×2 Matrix{Matrix{Int64}}:
 [1 3]  [5;;]
 [2 4]  [6;;]

julia> sum.(eachblock(A))
2×2 Matrix{Int64}:
 4  5
 6  6
```
"""
function eachblock(A::AbstractBlockArray)
    # blockinds = CartesianIndices(blocksize(A))
    blockinds = CartesianIndices(axes.(blocklasts.(axes(A)),1))
    (view(A, Block(Tuple(I))) for I in blockinds)
end

# Use memory laout for sub-blocks
@inline Base.getindex(A::AbstractMatrix, kr::Colon, jr::Block{1}) = ArrayLayouts.layout_getindex(A, kr, jr)
@inline Base.getindex(A::AbstractMatrix, kr::Block{1}, jr::Colon) = ArrayLayouts.layout_getindex(A, kr, jr)
@inline Base.getindex(A::AbstractMatrix, kr::Block{1}, jr::AbstractVector) = ArrayLayouts.layout_getindex(A, kr, jr)
@inline Base.getindex(A::AbstractArray{T,N}, kr::Block{1}, jrs...) where {T,N} = ArrayLayouts.layout_getindex(A, kr, jrs...)
@inline Base.getindex(A::AbstractArray{T,N}, block::Block{N}) where {T,N} = ArrayLayouts.layout_getindex(A, block)
@inline Base.getindex(A::AbstractMatrix, kr::AbstractVector, jr::Block) = ArrayLayouts.layout_getindex(A, kr, jr)
@inline Base.getindex(A::AbstractMatrix, kr::BlockRange{1}, jr::BlockRange{1}) = ArrayLayouts.layout_getindex(A, kr, jr)
@inline Base.getindex(A::LayoutMatrix, kr::BlockRange{1}, jr::BlockRange{1}) = ArrayLayouts.layout_getindex(A, kr, jr)
for Typ in (:AbstractTriangular, :Adjoint, :Transpose, :Symmetric, :Hermitian)
    @eval @inline Base.getindex(A::$Typ{<:Any,<:LayoutMatrix}, kr::BlockRange{1}, jr::BlockRange{1}) = ArrayLayouts.layout_getindex(A, kr, jr)
end

###
# permutedims
#
# just use transpose for now
###

Base.permutedims(A::AbstractBlockVector{<:Number}) = transpose(A)
Base.permutedims(A::AbstractBlockMatrix{<:Number}) = transpose(A)


"""
    bioclim(layer::Integer; left=nothing, right=nothing, bottom=nothing, top=nothing)

Download and prepare bioclim layers from the CHELSA database, and returns
them as an array of `SimpleSDMPredictor`s. Layers are called by their number,
from 1 to 19. The list of available layers is given in a table below.

The keyword argument is `path`, which refers to the path where the function
will look for the geotiff files.

Note that these files are *large* due the fine resolution of the data, and for
this reason this function will return the *integer* version of the layers. Also
note that the bioclim data are only available for the V1 of CHELSA, and are not
from the V2.

It is recommended to *keep* the content of the `path` folder, as it will
eliminate the need to download the tiff files (which are quite large). For
example, calling `bioclim(1:19)` will download and everything, and future
calls will be much faster.

| Variable | Description                                                |
| ------   | ------                                                     |
| 1        | Annual Mean Temperature                                    |
| 2        | Mean Diurnal Range (Mean of monthly (max temp - min temp)) |
| 3        | Isothermality (BIO2/BIO7) (* 100)                          |
| 4        | Temperature Seasonality (standard deviation *100)          |
| 5        | Max Temperature of Warmest Month                           |
| 6        | Min Temperature of Coldest Month                           |
| 7        | Temperature Annual Range (BIO5-BIO6)                       |
| 8        | Mean Temperature of Wettest Quarter                        |
| 9        | Mean Temperature of Driest Quarter                         |
| 10       | Mean Temperature of Warmest Quarter                        |
| 11       | Mean Temperature of Coldest Quarter                        |
| 12       | Annual Precipitation                                       |
| 13       | Precipitation of Wettest Month                             |
| 14       | Precipitation of Driest Month                              |
| 15       | Precipitation Seasonality (Coefficient of Variation)       |
| 16       | Precipitation of Wettest Quarter                           |
| 17       | Precipitation of Driest Quarter                            |
| 18       | Precipitation of Warmest Quarter                           |
| 19       | Precipitation of Coldest Quarter                           |

"""
function bioclim(layer::Integer; left=nothing, right=nothing, bottom=nothing, top=nothing)
    return raster(SimpleSDMPredictor, BioClim(), layer=layer, left=left, right=right, bottom=bottom, top=top)
end

bioclim(layers::Vector{T}; args...) where {T <: Integer} = [bioclim(l; args...) for l in layers]
bioclim(layers::UnitRange{T}; args...) where {T <: Integer} = [bioclim(l; args...) for l in layers]


function run_pandamodels_vstab(json_path)
    pm = load_pm_from_json(json_path)
    active_powermodels_silence!(pm)
    pm = check_powermodels_data!(pm)
    model = get_model(pm["pm_model"])
    solver = get_solver(pm)

    result = _run_vstab(
        pm,
        model,
        solver,
        setting = Dict("output" => Dict("branch_flows" => true)),
        ext = extract_params!(pm),
    )
    return result
end

function run_pandamodels_multi_vstab(json_path)
    pm = load_pm_from_json(json_path)
    active_powermodels_silence!(pm)
    pm = check_powermodels_data!(pm)
    model = get_model(pm["pm_model"])
    solver = get_solver(pm)
    mn = set_pq_values_from_timeseries(pm)

    result = _run_multi_vstab(
        mn,
        model,
        solver,
        setting = Dict("output" => Dict("branch_flows" => true)),
        ext = extract_params!(pm),
    )
    return result
end

function run_pandamodels_qflex(json_path)
    pm = load_pm_from_json(json_path)
    active_powermodels_silence!(pm)
    pm = check_powermodels_data!(pm)
    model = get_model(pm["pm_model"])
    solver = get_solver(pm)

    result = _run_qflex(
        pm,
        model,
        solver,
        setting = Dict("output" => Dict("branch_flows" => true)),
        ext = extract_params!(pm),
    )
    return result
end

function run_pandamodels_multi_qflex(json_path)
    pm = load_pm_from_json(json_path)
    active_powermodels_silence!(pm)
    pm = check_powermodels_data!(pm)
    model = get_model(pm["pm_model"])
    solver = get_solver(pm)
    mn = set_pq_values_from_timeseries(pm)

    result = _run_multi_qflex(
        mn,
        model,
        solver,
        setting = Dict("output" => Dict("branch_flows" => true)),
        ext = extract_params!(pm),
    )
    return result
end


# module for drawing goal flag

using Plots
pyplot()

mutable struct Goal
  pose
  radius
  value

  # init
  function Goal(x, y; radius=0.3, value=0.0)
    self = new()
    self.pose = [x, y]
    self.radius = radius
    self.value = value
    return self
  end
end

function inside(self::Goal, pose)
  if self.radius > sqrt((self.pose[1]-pose[1])^2 + (self.pose[2]-pose[2])^2)
    return true
  else
    return false
  end
end

function draw!(self::Goal)
  scatter!([self.pose[1]+0.1], [self.pose[2]+0.5],
           markershape=:rtriangle,
           markercolor=:red,
           markersize=9)
  plot!([self.pose[1], self.pose[1]], [self.pose[2], self.pose[2]+0.6],
        color=:black)
end


push!(LOAD_PATH, joinpath(expanduser("~/repos/InsarTimeseries.jl/src/")))
# import InsarTimeseries
# using InsarTimeseries: PHASE_TO_CM, STACK_FLAT_SHIFTED_DSET, STACK_FLAT_DSET, STACK_DSET, CC_FILENAME, UNW_FILENAME
import Polynomials
import Glob
using HDF5
using PyCall
using Printf: @printf
try
    using MapImages
catch
end

# gps = InsarTimeseries.gps
latlon = pyimport("apertools.latlon")
gps = pyimport("apertools.gps")

reloadpython(modname) = pyimport("importlib")."reload"(modname)
reloadpython(gps)

LOS_MAP_FILE = "los_enu.tif"

# PATH 78
station_name_list78 = [
    "NMHB",
    "TXAD",
    "TXBG",
    "TXBL",
    "TXCE",
    "TXFS",
    "TXKM",
    "TXL2",
    "TXMC",
    "TXMH",
    "TXOE",
    "TXOZ",
    "TXS3",
    "TXSO",
]
#
# PATH 85
station_name_list85 = ["TXKM", "TXMH", "TXFS", "NMHB", "TXAD", "TXS3"]
# Take out TXAL cuz it's not in crop
# BAD station reasons: MDO1 (nothing 2014-2017), TXVH (started 2018), TXPC (ended 2017)

all_stations = sort(unique([station_name_list78; station_name_list85]))

# Functions for error evaluation
rms(arr::AbstractArray{T,1}) where {T <: Any} = sqrt(mean(filter(!isnan, arr.^2)))
# rms(arr::AbstractArray{T, 2}) where {T <: Any} = [rms(arr[i, :]) for i in 1:size(arr, 1)]
maxabs(x) = maximum(abs.(filter(!isnan, x)))

# Function to take a velocity file and calculate the GPS errors at one station
function get_gps_error(
    insar_fname::String,
    station_name::String;
    dset="velos",
    window=5,
    ref_station=nothing,
    verbose=false,
    shift=0.0,
    negate=false,
)
    # If it's a file with LOS flipped:
    flip_sign = negate ? -1 : 1

    # insar derived solution in a small patch
    slope_insar_mm_yr =
        flip_sign * _get_val_at_station(insar_fname, station_name, dset=dset, window=window) + shift
    slope_insar_mm_yr -= isnothing(ref_station) ? 0 :
        (
        flip_sign * _get_val_at_station(insar_fname, ref_station, dset=dset, window=window) + shift
    )

    start_date, end_date = _get_date_range(insar_fname, dset)
    slope_gps_mm_yr = solve_gps_ts(
        station_name,
        ref_station,
        start_date=start_date,
        end_date=end_date,
        los_map_file=joinpath(dirname(insar_fname), LOS_MAP_FILE),
    )

    if verbose
        @show station_name
        println("GPS velocity (mm / year): $slope_gps_mm_yr")
        println("InSAR linear velocity (mm / year): $slope_insar_mm_yr")
        println("==="^10)
    end

    return slope_insar_mm_yr - slope_gps_mm_yr
end
get_gps_error(fname::String, station_list::Array{String}; kwargs...) =
    get_gps_error.(fname, station_list; kwargs...)
# ! note: dot broadcasting is same as [get_gps_error.(fname, stat; kwargs...) for stat in station_list]

_get_date_range(fname::AbstractString) = extrema(Sario.load_geolist_from_h5(fname))
_get_date_range(fname::AbstractString, dset::AbstractString) =
    extrema(Sario.load_geolist_from_h5(fname, dset))

function _get_station_rowcol(station_name, directory=".")
    demrsc = Sario.load(joinpath(directory, "dem.rsc"))
    return map(x -> convert(Int, x), MapImages.station_rowcol(station_name, demrsc))
end

"""Load the insar data in a small patch around the pixel"""
function _get_val_at_station(
    insar_fname,
    station_name;
    dset="velos",
    window=5,
    kwargs...,
)
    # Note: swapping row and col due to julia/hdf5 transposes
    row, col = _get_station_rowcol(station_name, dirname(insar_fname))
    halfwin = div(window, 2)
    try
        patch =
            h5read(insar_fname, dset, (col - halfwin:col + halfwin, row - halfwin:row + halfwin))
        return mean(patch)
    catch e
        println(e)
        return NaN
    end
end

get_station_values(fname, station_list::Array{String}; kwargs...) =
    [_get_val_at_station(fname, stat; kwargs...) for stat in station_list]

"""Given a list of errors from insar-gps, find the const to add
to the insar solution image to minimize these errors
(converts the gps from relative to absolute)"""
function minimize_errors(error_list, search_range=-5:0.1:5)
    best_rms, best_maxabs = 100, 100
    c_rms, c_maxabs = 0, 0
    for c in search_range
        cur_rms = rms(c .+ error_list)
        cur_maxabs = maxabs(c .+ error_list)
        if cur_rms < best_rms
            c_rms = c
            best_rms = cur_rms
        end

        if cur_maxabs < best_maxabs
            c_maxabs = c
            best_maxabs = cur_maxabs
        end
    end
    return c_rms, c_maxabs, best_rms, best_maxabs
end

function print_errors(
    errors78::AbstractArray{<:AbstractFloat},
    errors85::AbstractArray{<:AbstractFloat},
    station_name_list::AbstractArray{<:AbstractString},
)
    println("Name | error 78 | error 85 ")
    println("============================")
    for (err78, err85, name) in zip(errors78, errors85, station_name_list)
        # println("$name RESULTS:")
        @printf("%s |  %.2f | %.2f \n", name, err78, err85)
    end
    println("==============")
    println("RMS 78 | Max-abs 78 | RMS 85 | Max-abs 85")
    @printf(
        "%.2f |  %.2f | %.2f |  %.2f  \n",
        rms(errors78),
        maxabs(errors78),
        rms(errors85),
        maxabs(errors85)
    )
end

function compare_solutions_with_gps(geolist, intlist, unw_vals, station_name, linear=true)
    B = prepB(geolist, intlist, constant_velocity=linear)

    l1_diffs = Array{Float32,1}(undef, length(geolist))
    lstsq_diffs = similar(l1_diffs)

    slope_gps_mm_yr = solve_gps_ts(station_name, nothing)

    # First, solve with a dummy date so nothing is removed
    base_l1, base_lstsq = solve_without(Date(2000, 1, 1), intlist, unw_vals, B)
    base_l1_error = base_l1 - slope_gps_mm_yr
    base_lstsq_error = base_lstsq - slope_gps_mm_yr
    println("Base L1 error for $station_name = $base_l1_error")

    for (idx, d) in enumerate(geolist)
        l1, lstsq = solve_without(d, intlist, unw_vals, B)
        l1d = l1 - slope_gps_mm_yr
        lstsqd = lstsq - slope_gps_mm_yr

        # Note: if base is larger, the diff will be positive (an improvement)
        l1_diffs[idx] = abs(base_l1_error) - abs(l1d)
        lstsq_diffs[idx] = abs(base_lstsq_error) - abs(lstsqd)
    end
    return l1_diffs, lstsq_diffs, base_l1_error
end

############################
# GPS FUNCTIONS
############################
function get_gps_los(
    station_name;
    los_map_file=LOS_MAP_FILE,
    reference_station=nothing,
    start_date=Date(2014, 11, 1),
    end_date=Date(2019, 1, 1),
    days_smooth=0,
)
    dts, gps_los_data = gps.load_gps_los_data(
        days_smooth=days_smooth,
        los_map_file=los_map_file,
        station_name=station_name,
        start_date=start_date,
        end_date=end_date,
        zero_mean=true,
        reference_station=reference_station,
    )

    return [convert(Date, d) for d in dts], gps_los_data
end

function get_gps_enu(station_name)
    dts, enu_df = gps.load_station_enu(
        station_name,
        start_date=Date(2014, 11, 1),
        end_date=Date(2019, 1, 1),
        zero_mean=true,
    )

    # Convert from PyObjects to Arrays
    dts = [convert(Date, d) for d in dts]
    east = [r for r in enu_df.east]
    north = [r for r in enu_df.north]
    up = [r for r in enu_df.up]
    return dts, east, north, up
end

"""Find the linear fit of MM per year of the gps station"""
function solve_gps_ts(
    station_name,
    reference_station=nothing;
    start_date=Date(2014, 11, 1),
    end_date=Date(2019, 1, 1),
    los_map_file=LOS_MAP_FILE,
)
    dts, gps_los_data = get_gps_los(
        station_name,
        reference_station=reference_station,
        start_date=start_date,
        end_date=end_date,
        los_map_file=los_map_file,
    )
    # If we wanna compare with GPS subtracted too, do this:
    # dts, gps_los_data = get_gps_los(station_name, reference_station=reference_station)

    # Convert to "days since start" for line fitting
    gps_poly = fit_line(dts, gps_los_data)
    slope = length(gps_poly) == 2 ? Polynomials.coeffs(gps_poly)[2] :
        Polynomials.coeffs(gps_poly)[1]
    # offset, slope = Polynomials.coeffs(gps_poly)
    slope_gps_mm_yr = 365 * 10 * slope
    return slope_gps_mm_yr
end

function fit_line(dts, data)
    day_nums = InsarTimeseries._get_day_nums(dts)
    p = Polynomials.polyfit(day_nums, data, 1)
    # p(day_nums[end])
    return p
end

function shift_from(src::HDF5Dataset, shift::Real)
    tmp = src .+ shift
    tmp[src .== 0] .= 0
    return tmp
end

function shift_dset(shift, fname, src::String="velos/1", dest="velos_shifted/1")
    h5open(fname, "r+") do f
        out = shift_from(f[src], shift)
        f[dest] = out
    end
    h5writeattr(fname, dest, h5readattr(fname, src))
    return
end

function scale_dset!(scale, fname, src::String="velos/1")
    a = h5readattr(fname, src)
    h5open(fname, "cw") do f
        tmp = scale * read(f, src)
        o_delete(f, src)
        f[src] = tmp
    end
    h5writeattr(fname, src, a)
    return
end

function get_shift(fname, src::String="velos/1", dest="velos_shifted/1")
    h5open(fname, "r+") do f
        return f[dest][div(end, 2), div(end, 2)] - f[src][div(end, 2), div(end, 2)]
    end
    end

function save_as_unw(fname, dset="velos/1")
    amp = abs.(Sario.load(Glob.glob("*.int")[1]))
    img = Sario.load(fname, dset_name=dset)
    outname = replace(fname, ".h5" => ".unw")
    Sario.save(outname, cat(amp, Float32.(img ./ 10), dims=3))
end


    ######### getting all together in Housedorff
    # julia -g 2 C:\GitHub\GitHub\NuclearMedEval\test\distanceMetr\HousdorffTest\integrationHousedorff.jl
    using Revise, Parameters, Logging, Test
    using CUDA
    includet("C:\\GitHub\\GitHub\\NuclearMedEval\\test\\includeAllUseFullForTest.jl")
    using ..CUDAGpuUtils ,..IterationUtils,..ReductionUtils , ..MemoryUtils,..CUDAAtomicUtils
    using ..MainLoopKernel,..PrepareArrtoBool,..MetadataAnalyzePass,..MetaDataUtils,..WorkQueueUtils,..ProcessMainDataVerB,..HFUtils,..ResultListUtils, ..Housdorff

    CUDA.allowscalar(true)

    mainArrDims= (120,120,120);
    mainArrCPU= zeros(UInt8,mainArrDims);
    refArrCPU = zeros(UInt8,mainArrDims);
    ##### we will create two planes 20 units apart from each
    mainArrCPU[10:50,10:50,10].= 1;
    refArrCPU[10:50,10:50,80].= 1;
    # mainArrCPU[10:12,10:12,10].= 1;
    # refArrCPU[10:12,10:12,30].= 1;

    goldGPU = CuArray(mainArrCPU);
    segmGPU= CuArray(mainArrCPU);

    robustnessPercent= 0.95
    numberToLooFor=1
    boolKernelArgs, mainKernelArgs,threadsBoolKern,blocksBoolKern ,threadsMainKern,blocksMainKern ,shmemSizeBool,shmemSizeMain=    preparehousedorfKernel(goldGPU,segmGPU,robustnessPercent,numberToLooFor)
    mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength=boolKernelArgs

        gold3d = CuArray(mainArrCPU)
        segm3d= CuArray(refArrCPU)
        numberToLooFor =UInt8(1)
        robustnessPercent = 0.9

    # globalIterationNumb= Housdorff.getHousedorffDistance(goldGPU,segmGPU,boolKernelArgs,mainKernelArgs,threadsBoolKern,blocksBoolKern ,threadsMainKern,blocksMainKern,shmemSizeBool,shmemSizeMain)
    # globalIterationNumb[1]

        boolKernelArgs, mainKernelArgs,threadsBoolKern,blocksBoolKern ,threadsMainKern,blocksMainKern ,shmemSizeBool,shmemSizeMain=    preparehousedorfKernel(gold3d,segm3d,robustnessPercent,numberToLooFor)
        mainArrDims,dataBdim,metaData,metaDataDims,reducedGoldA,reducedSegmA,loopXinPlane,loopYinPlane,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp ,numberToLooFor,inBlockLoopXZIterWithPadding,shmemblockDataLoop,shmemblockDataLenght,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength=boolKernelArgs

        golddd = CuArray(mainArrCPU);
        segmmm= CuArray(refArrCPU);

    @cuda threads=threadsBoolKern blocks=blocksBoolKern shmem=shmemSizeBool  cooperative=true boolKernelLoad(golddd,segmmm,boolKernelArgs...)

    Int64(sum(mainArrCPU))
    Int64(fn[1])
    Int64(fp[1])
    Int64(minxRes[1])
    Int64(maxxRes[1])
    Int64( minyRes[1])
    Int64(maxyRes[1])
    Int64( minzRes[1])
    Int64(maxzRes[1])
size(reducedGoldA)
size(reducedSegmA)
size(metaData)
    #now time to get data structures dependent on bool kernel like for example loading subsections of meta data, creating work queue ...
    #some arrays needs to be instantiated only after we know the number of the false and true positives
 referenceArrs,dilatationArrs, mainArrDims,dataBdim,numberToLooFor,metaDataDims,metaDataNull,iterThrougWarNumb,robustnessPercent,shmemSumLengthMaxDiv4,globalFpResOffsetCounter,globalFnResOffsetCounter    ,globalIterationNumber,globalCurrentFnCount,globalCurrentFpCount    ,globalIterationNumb,workQueue,workQueueCounter,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed    ,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop    ,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength, fn,fp,resList,inBlockLoopXZIterWithPadding,paddingStore,shmemblockDataLenght,shmemblockDataLoop= mainKernelArgs
        size(metaData)

    metaData,reducedGoldA  ,reducedSegmA ,paddingStore,resList,workQueue,workQueueCounter= getBigGPUForHousedorffAfterBoolKernel(metaData,minxRes,maxxRes,minyRes,maxyRes,minzRes,maxzRes,fn,fp,reducedGoldA,reducedSegmA,dataBdim)

    reducedSegmA[:,:,3]

    metaDataDims = size(metaData)
    metaDataLength= metaDataDims[1]*metaDataDims[2]*metaDataDims[3]
    loopMeta= fld(metaDataLength,blocksMainKern )
    loopWarpMeta= cld(metaDataLength,(blocksMainKern*threadsMainKern[1] ))

    dilatationArrs= (reducedGoldA,reducedSegmA)
    #reverse order as when dilatating gold we want to establish do we cover segm voxels
    referenceArrs=(golddd,segmmm )
    shmemblockDataLenght = threadsMainKern[1]*threadsMainKern[2]*2
    shmemblockDataLoop = fld(shmemblockDataLenght,threadsMainKern[1]*threadsMainKern[2])

using  BenchmarkTools
    BenchmarkTools.DEFAULT_PARAMETERS.samples = 2
    BenchmarkTools.DEFAULT_PARAMETERS.seconds =600
    BenchmarkTools.DEFAULT_PARAMETERS.gcsample = true

    #main calculations

    # @benchmark CUDA.@sync @cuda threads=threadsMainKern blocks=blocksMainKern shmem=shmemSizeMain cooperative=true mainKernelLoad( referenceArrs,dilatationArrs, mainArrDims,dataBdim
    @benchmark CUDA.@sync @cuda threads=threadsMainKern blocks=blocksMainKern shmem=shmemSizeMain cooperative=true mainKernelLoad( referenceArrs,dilatationArrs, mainArrDims,dataBdim
    ,numberToLooFor,metaDataDims,metaData,iterThrougWarNumb,robustnessPercent
    ,shmemSumLengthMaxDiv4,globalFpResOffsetCounter,globalFnResOffsetCounter
    ,globalIterationNumber,globalCurrentFnCount,globalCurrentFpCount
    ,globalIterationNumb,workQueue,workQueueCounter
    ,loopAXFixed,loopBXfixed,loopAYFixed,loopBYfixed,loopAZFixed,loopBZfixed
    ,loopdataDimMainX,loopdataDimMainY,loopdataDimMainZ,inBlockLoopX,inBlockLoopY
    ,inBlockLoopZ,metaDataLength,loopMeta,loopWarpMeta,clearIterResShmemLoop
    ,clearIterSourceShmemLoop,resShmemTotalLength,sourceShmemTotalLength, fn,fp,resList,inBlockLoopXZIterWithPadding,paddingStore,shmemblockDataLenght,shmemblockDataLoop)

    globalIterationNumb[1]

    Int64(workQueueCounter[1])
    Int64.(workQueue)

#     (threadIdxX())+(innerWarpNumb+21)*33
#     32+(16+21)*33
# dataBdim[1]* dataBdim[2]*2

resList

for i in 1:16
    # println(" i $(i)  $((Int64(metaData[2,2,2, (getOldCountersBeg()+1)])))")
    println(" i $(i)  $((Int64(metaData[2,2,2, (getNewCountersBeg()+1)])))")
end
mainArrDims


module ReactiveMPModelsTemplateTest

using Test, InteractiveUtils
using Rocket, ReactiveMP, GraphPPL, Distributions
using BenchmarkTools, Random, Plots, Dates, LinearAlgebra, StableRNGs

# Please use StableRNGs for random number generators

## Model definition
## -------------------------------------------- ##
@model function template_model(n, x0, c_, P_)
    error("define model")
end
## -------------------------------------------- ##
## Inference definition
## -------------------------------------------- ##
function template_inference(data, x0, c, P)
    error("define inference")
end

@testset "Model template" begin

    @testset "Use case template" begin
        ## -------------------------------------------- ##
        ## Data creation
        ## -------------------------------------------- ##
        _
        ## -------------------------------------------- ##
        ## Inference execution
        -
        ## -------------------------------------------- ##
        ## Test inference results
        _
        ## -------------------------------------------- ##
        ## Form debug output
        base_output = joinpath(pwd(), "_output", "models")
        mkpath(base_output)
        timestamp        = Dates.format(now(), "dd-mm-yyyy-HH-MM")
        plot_output      = joinpath(base_output, "template_model_plot_$(timestamp)_v$(VERSION).png")
        benchmark_output = joinpath(base_output, "template_model_benchmark_$(timestamp)_v$(VERSION).txt")
        ## -------------------------------------------- ##
        ## Create output plots

        savefig(p, plot_output)
        ## -------------------------------------------- ##
        ## Create output benchmarks
        benchmark = @benchmark 1 + 1#
        open(benchmark_output, "w") do io
            show(io, MIME("text/plain"), benchmark)
            versioninfo(io)
        end
        ## -------------------------------------------- ##
    end

end

end


module Lattices

using ..LatticeVars
import ..LatticeVars.init!, ..LatticeVars.get_num_sites_visited
using ..Points
using Distributions
using JDistributions: BernoulliBool

export AbstractLattice, Lattice, AbstractLatticeTypes, Traps, TrapsLattice
export SitePercolationLattice, SitePercolation

abstract type AbstractLattice end

abstract type AbstractLatticeTypes end
struct Traps <: AbstractLatticeTypes
end

"""
    Lattice{N, T} <: AbstractLattice

Type representing variables on the sites of an `N`-dimensional lattice.
`T` is a `Tuple` type of variables. These variables are meant to be a mixture of
types `Distribution` (annealed variables) and `LatticeVars` (quenched variables).
"""
struct Lattice{N, T, Opts} <: AbstractLattice
    vars::T
    Lattice{N}(vars::Tuple; lattice_type = Traps) where N = new{N, typeof(vars), lattice_type}(vars)
    function Lattice{N}(vars...; lattice_type = Traps) where N
        vartup = (vars...,)
        return new{N, typeof(vartup), lattice_type}(vartup)
    end
end

Lattice(args...; kwargs...) = Lattice{1}(args...; kwargs...)

const TrapsLattice = Lattice{N, T, Traps} where {N, T}

struct SitePercolation <: AbstractLatticeTypes
end

const SitePercolationLattice = Lattice{N, T, SitePercolation} where {N, T}

function SitePercolationLattice{N}(p) where {N}
    return Lattice{N}(LatticeVar{N}(BernoulliBool(p)); lattice_type = SitePercolation)
end

SitePercolationLattice(p) = SitePercolationLattice{1}(p)

_getindex(dist::Distribution, inds...) = rand(dist)
_getindex(lattice_var::AbstractLatticeVar, inds...) = lattice_var[inds...]
Base.getindex(lattice::Lattice, i::Integer) = lattice.vars[i]
Base.getindex(lattice::Lattice, i::Integer, inds::Integer...) = _getindex(lattice[i], inds)
Base.getindex(lattice::Lattice, i::Integer, p::Point) = _getindex(lattice[i], get_coords(p))
init!(lattice::Lattice) = foreach(x -> init!(x), lattice.vars)
get_vars(lattice::Lattice) = lattice.vars
get_num_sites_visited(lattice::Lattice) = get_num_sites_visited(get_vars(lattice)[1])

function Base.show(io::IO, lattice::Lattice{N}) where N
    println(io, N, "-dimensional Lattice")
    for v in get_vars(lattice)
        println(io, v)
    end
end

end


using DrWatson
@quickactivate
using HierarchicalAD
using Flux, CUDA
using BSON, FileIO
using ValueHistories, DataFrames
using Plots
using ArgParse
using IterTools

using HierarchicalAD: construct_conv_classifier, train_conv_classifier!, create_input_data
using HierarchicalAD: auc_val, classifier_score

arg_table = ArgParseSettings(;autofix_names=true)
@add_arg_table arg_table begin
	"--modeldir"
		default = "models/contaminated_training/run2/"
		help = "dir where the models are saved"
	"--model-id"
		default = nothing
		help = "id of the FVLAE model whose architecture is to be used for construction fo the classifier"
	"--model-ind"
		default = nothing
		help = "index of the FVLAE model whose architecture is to be used for construction fo the classifier"
end
args = parse_args(arg_table)
@unpack modeldir, model_id, model_ind = args

# get the data
X, y = HierarchicalAD.load_shapes2D();

# set device id
CUDA.device!(2)

# setup paths
modeldir = datadir(modeldir)
plotpath = datadir("plots/baseline_comparison/conv_classifier")
outpath = datadir("baseline_comparison/conv_classifier")
mkpath(outpath)
mkpath(plotpath)

# what models are we gonna use
inmfs = readdir(modeldir)
if !isnothing(model_ind)
	mfs = [inmfs[Meta.parse(model_ind)]]
elseif !isnothing(model_id)
	mfs = filter(x->occursin(model_id, x), mfs)
else
	mfs = inmfs # iterate over all
end

mfs = ["activation=relu_channels=8-8-16-32-64_data=shapes2D_gamma=50.0_lambda=0.0_latent_dim=4_model=fvlae_model_id=20210817174044364_xdist=bernoulli.bson"]

# some params
nns = floor.(Int, 10.0 .^ range(1, 5, length=10))
#nns = floor.(Int, 10.0 .^ range(1, 2, length=2))
nas = floor.(Int, 10.0 .^ range(1, 4, length=4))
#nas = floor.(Int, 10.0 .^ range(1, 2, length=2))
ntst = 25000
label_vec = [
	y[!,:shape] .== 2, y[!,:shape] .!= 2,
	y[!,:scale] .< 0.75, y[!,:scale] .> 0.75,
	y[!,:posX] .> 0.5, y[!,:posY] .> 0.5,
	y[!,:normalized_orientation] .< 0.1,
	.!((y[!,:shape] .== 1) .| ((y[!,:posX] .> 0.5) .& (y[!,:posY] .> 0.5)))
	]
label_desc = [
	"ovals normal", "ovals anomalous",
	"small normal", "large normal",
	"right normal", "bottom normal",
	"rotated anomalous",
	"squares or right bottom anomalous"
	]
conv_classifier_params = (lr = 0.001f0, λ = 0f0, batchsize = 64, nepochs = 1000, patience = 1, verb = true)

function train_and_evaluate_classifier(mf)
	model_id = split(split(mf, "model_id=")[2], "_")[1]
	autoencoder_data = load(joinpath(modeldir, mf))

	for seed = 1:10
		aucs_per_split = []
		scores_per_split = []
		tst_ys_per_split = []
		for (labels, desc) in zip(label_vec, label_desc)
			aucs_per_samples = []
			scores_per_samples = []
			tst_ys_per_samples = []
			for (nn, na) in product(nns, nas)
				@info "Training classifier using $(desc), seed=$seed."
				@info "# normal training/validation = $nn, # anomalous training/validation = $na, # test = $(ntst*2)"
				try
					global tr_nX, (tr_X, tr_y), (val_X, val_y), (tst_X, tst_y) = create_input_data(X, y, labels, nn, nn, na, na,
					    ntst; seed=seed);
				catch e
					if isa(e, BoundsError)
						@info "Not enough data, skipping training..."
						push!(aucs_per_samples, NaN)
						push!(scores_per_samples, [NaN])
						push!(tst_ys_per_samples, [NaN])
						continue
					else
						rethrow(e)
					end
				end
				classifier = construct_conv_classifier(
					autoencoder_data[:experiment_args].kernelsizes,
					autoencoder_data[:experiment_args].activation,
					autoencoder_data[:experiment_args].stride,
					autoencoder_data[:experiment_args].channels,
					size(X))

				classifier, history, opt = train_conv_classifier!(classifier, tr_X, tr_y, val_X, val_y; conv_classifier_params...);
				scores = cpu(classifier_score(classifier, tst_X, 128))
				auc = auc_val(tst_y, scores)
				@info "Finished training, test AUC = $(auc)"
				push!(aucs_per_samples, auc)
				push!(scores_per_samples, scores)
				push!(tst_ys_per_samples, tst_y)
			end
			push!(aucs_per_split, aucs_per_samples)
			push!(scores_per_split, scores_per_samples)
			push!(tst_ys_per_split, tst_ys_per_samples)
		end
		savef = joinpath(outpath, "classifier_output_model_id=$(model_id)_seed=$(seed).bson")
		outdata = Dict(
			:nns => nns, :nas => nas, :ntst => ntst,
			:model_id => model_id,
			:conv_classifier_params => conv_classifier_params,
			:aucs => aucs_per_split,
			:scores => scores_per_split,
			:tst_ys => tst_ys_per_split,
			:label_desc => label_desc
			)
		save(savef, outdata)
		@info "Saved to $savef"
	end
end

for mf in mfs
	train_and_evaluate_classifier(mf)
end


# Raytracer.jl
# Raytracing for the generation of photorealistic images in Julia
# Copyright (c) 2021 Samuele Colombo, Paolo Galli

# HitRecord for store the result of an intersection between a Ray and a Shape

"""
    HitRecord

A struct representing the result of an intersection between a [`Ray`](@ref) and a [`Shape`](@ref).

# Fields

- `world_point::Point`: a [`Point`](@ref) representing the world coordinates of the hit point.
- `normal::Normal`: a [`Normal`](@ref) representing the orientation of the normal to the surface where the hit happened.
- `surface_point::Vec2D`: a [`Vec2D`](@ref) representing the position of the hit point on the surface of the object.
- `t::Float32`: distance from the origin of the ray where the hit happened.
- `ray::Ray`: a [`Ray`](@ref) representing the the ray that hit the surface.
- `material::Material`: a [`Material`](@ref) representing the material of the point where the hit happened.
"""
struct HitRecord
    world_point::Point
    normal::Normal
    surface_point::Vec2D
    t::Float32
    ray::Ray
    material::Material
end

"""
    HitOrMiss

Alias for `Union{HitRecord, Nothing}`.
"""
const HitOrMiss = Union{HitRecord, Nothing}

function show(io::IO, ::MIME"text/plain", hr::T) where {T <: HitRecord}
    print(io, T)
    fns = fieldnames(T)
    n = maximum(fns .|> String .|> length)
    for fieldname ∈ fns
        println(io)
        print(io, " ↳ ", rpad(fieldname, n), " = ", getfield(hr, fieldname))
    end
end

"""
    ≈(hr1::HitRecord, hr2::HitRecord)

Check if two [`HitRecord`](@ref) represent the same hit event or not.
"""
function (≈)(hr1::HitRecord, hr2::HitRecord)
    sp1, sp2 = hr1.surface_point, hr2.surface_point
    hr1.world_point     ≈  hr2.world_point     &&
    hr1.normal          ≈  hr2.normal          &&
    findall(isnan, sp1) == findall(isnan, sp2) &&
    filter(!isnan, sp1) ≈  filter(!isnan, sp2) &&
    hr1.t               ≈  hr2.t               &&
    hr1.ray             ≈  hr2.ray
end

function Base.isless(hr1::HitRecord, hr2::HitRecord)
    hr1.t/norm(hr1.ray.dir) < hr2.t/norm(hr2.ray.dir)
end

function Base.isequal(hr1::HitRecord, hr2::HitRecord)
    isequal(hr1.t/norm(hr1.ray.dir), hr2.t/norm(hr2.ray.dir))
end

function Base.:(==)(hr1::HitRecord, hr2::HitRecord)
    hr1.t/norm(hr1.ray.dir) == hr2.t/norm(hr2.ray.dir)
end


rand_seis_unit() = rand() > 0.5 ? "m" : "m/s"
rand_seis_cc() = rand(rand() > 0.2 ? zne : nvc)

const cDict = Dict{Char, Function}(
  'H' => i->rand_seis_cc(),
  'L' => i->rand_seis_cc(),
  'N' => i->rand_seis_cc(),
  'A' => i->rand() > 0.5 ? 'N' : 'E',
  'B' => i->'_',
  'D' => i->rand(oidfhu),
  'F' => i->rand(zne),
  'G' => i->rand(nvc),
  'I' => i->rand(oid),
  'J' => i->rand(nvc),
  'K' => i->rand(oid),
  'M' => i->rand(nvc),
  'O' => i->'_',
  'P' => i->rand(zne),
  'Q' => i->'_',
  'R' => i->'_',
  'S' => i->rand(zne),
  'T' => i->'Z',
  'U' => i->'_',
  'V' => i->'_',
  'W' => i->rand() > 0.5 ? 'S' : 'D',
  'Z' => i->rand(icfo),
  )

const uDict = Dict{Char, Function}(
    'H' => i->rand_seis_unit(),
    'N' => i->"m/s2",
    'L' => i->rand_seis_unit(),
    'A' => i->"rad",
    'B' => i->"m",
    'D' => i->"Pa",
    'F' => i->"T",
    'G' => i->"m/s2",
    'I' => i->"%",
    'J' => i->rand(junits),
    'K' => i->rand() > 0.5 ? "Cel" : "K",
    'M' => i->"m",
    'O' => i->"m/s",
    'P' => i->rand_seis_unit(),
    'Q' => i->"V",
    'R' => i->rand_seis_unit(),
    'S' => i->"m/m",
    'T' => i->"m",
    'U' => i->"%{cloud_cover}",
    'V' => i->"m3/m3",
    'W' => i->i=='S' ? "m/s" : "{direction_vector}",
    'Z' => i->rand_seis_unit(),
  )

  """
     (cha, u) = iccodes_and_units(b::Char, s::Bool)

Using band code `b`, generate quasi-sane random instrument code `i` and channel code `c`, returning channel string `cha` = `b`*`i`*`c` and unit string `u`. If `s=true`, use only seismic intrument codes.
"""
function iccodes_and_units(b::Char, s::Bool)
  i = rand(s ? hln : iclist)
  c = cDict[i](b)
  u = uDict[i](c)
  return string(b, i, c), u
end


using JuMP, Gurobi, Graphs, JLD, DataFrames
include("DataImport.jl")
include("TypeInfo.jl")
include("contin_data.jl")
include("load24_data.jl")
include("PHsub_ini.jl")
include("PHsub_iter.jl")
include("DelayMatrix.jl")

#### Optimization Parameters Control Area ####
gap_ini = 1e-3
gap = 1e-3
rho_zero = 70.962 # 1MWh/300kW => 124.1835  #1MWh/150kW => 70.962
delayMode = 1
##############################################
########## Simulation Case Setting ##########
N_line = 14
N_time = 24
N_storage = 3
N_bus = 15
N_scen_normal = 1
N_scen_contin = 1
N_scen = N_scen_normal + N_scen_contin
N_event = 1
##############################################
#### Progressive Hedging Penalty Setting ####
rho_x = rho_zero * ones(N_storage)
rho_z = (rho_zero / 1 )* ones(N_storage,N_bus)
rho = PenaltyParameter(rho_z,rho_x)
##############################################
filename_Node = "Node.csv"
filename_Generator = "Generator.csv"
filename_Line = "Line.csv"

filename_Storage = "Storage1MWhx3.csv"

filename_Station = "Station.csv"
filename_SMP = "DataMiner-Export_2018-01-10-121427.csv"
filename_Load24 = "load24data-PJM.csv"

load24 = load24_data(filename_Load24)

contin = contin_data(N_line,N_time,N_scen_normal,N_scen_contin,N_event)
scen_data = ScenarioData(contin,load24)

S_N = Int64[]
S_E = Int64[]
for s = 1:N_scen
    if sum(contin,1:2)[s] == 0
        push!(S_N,s)
    elseif sum(contin,1:2)[s] > 0
        push!(S_E,s)
    end
end

Prob = Array{Float64}(N_scen,1) # Scenario Probability
Prob[S_N] = 0.9 / length(S_N)
Prob[S_E] = 0.1 / length(S_E)

status_ini = Array{String}(N_scen,1)
obj_value_ini = Array{Float64}(N_scen,1)
simulation_time_ini = Array{Float64}(N_scen,1)
ESlocation_ini = Array{Int64}(N_storage,N_time,N_scen)
SoC_ini = Array{Float64}(N_storage, N_time, N_scen)
u_ini = Array{Float64}(N_storage, N_bus, N_time, N_scen)
z_ini = Array{Float64}(N_storage, N_bus, N_scen)
v_ini = Array{Float64}(N_bus, N_time, N_scen)
a_ini = Array{Float64}(N_bus, N_time, N_scen)
pg_ini = Array{Float64}(N_bus, N_time, N_scen)
qg_ini = Array{Float64}(N_bus, N_time, N_scen)
IC_ini = Array{Float64}(N_scen,1)
OC_ini = Array{Float64}(N_scen,1)
EC_ini = Array{Float64}(N_scen,1)
OCinter_ini = Array{Float64}(N_scen,1)
OCgen_ini = Array{Float64}(N_scen,1)
OCloadshed_ini = Array{Float64}(N_scen,1)
OCstorage_ini = Array{Float64}(N_scen,1)
x_ini = Array{Float64}(N_storage, N_scen)
sigma_d_ini = Array{Float64}(N_bus, N_time, N_scen)
sigma_l_ini = Array{Float64}(N_line, N_time, N_scen)

#

for s = 1:N_scen
    if s in S_N
        contin_temp = contin[:,:,s]
        scen_data_temp = ScenarioData(contin_temp,load24[s,:])
    elseif s in S_E
        contin_temp = contin[:,:,s]
        scen_data_temp = ScenarioData(contin_temp,mean(load24[S_N,:],1))
    end
    status_ini[s],obj_value_ini[s],simulation_time_ini[s],ESlocation_ini[:,:,s],SoC_ini[:,:,s],u_ini[:,:,:,s],z_ini[:,:,s],v_ini[:,:,s],a_ini[:,:,s],pg_ini[:,:,s],qg_ini[:,:,s],IC_ini[s],OC_ini[s],EC_ini[s],OCinter_ini[s],OCgen_ini[s],OCloadshed_ini[s],OCstorage_ini[s],x_ini[:,s],sigma_d_ini[:,:,s],sigma_l_ini[:,:,s]  = PHsub_ini(scen_data_temp, gap_ini, delayMode, filename_Node, filename_Generator, filename_Line, filename_Storage, filename_Station, filename_SMP)
end

@save "PH_ini.jld"
#@load "PH_ini.jld"

#############
status = Array{String}(N_scen,1)
obj_value = Array{Float64}(N_scen,1)
simulation_time = Array{Float64}(N_scen,1)
ESlocation = Array{Int64}(N_storage,N_time,N_scen)
SoC = Array{Float64}(N_storage, N_time, N_scen)
u = Array{Float64}(N_storage, N_bus, N_time, N_scen)
z = Array{Float64}(N_storage, N_bus, N_scen)
v = Array{Float64}(N_bus, N_time, N_scen)
a = Array{Float64}(N_bus, N_time, N_scen)
pg = Array{Float64}(N_bus, N_time, N_scen)
qg = Array{Float64}(N_bus, N_time, N_scen)
IC = Array{Float64}(N_scen,1)
OC = Array{Float64}(N_scen,1)
EC = Array{Float64}(N_scen,1)
OCinter = Array{Float64}(N_scen,1)
OCgen = Array{Float64}(N_scen,1)
OCloadshed = Array{Float64}(N_scen,1)
OCstorage = Array{Float64}(N_scen,1)
x = Array{Float64}(N_storage, N_scen)
sigma_d = Array{Float64}(N_bus, N_time, N_scen)
sigma_l = Array{Float64}(N_line, N_time, N_scen)
#########################
N_iter = 40
#########################
status_iter = Array{String}(N_scen, N_iter)
obj_value_iter = Array{Float64}(N_scen, N_iter)
simulation_time_iter = Array{Float64}(N_scen, N_iter)
ESlocation_iter = Array{Int64}(N_storage,N_time,N_scen, N_iter)
SoC_iter = Array{Float64}(N_storage, N_time, N_scen, N_iter)
u_iter = Array{Float64}(N_storage, N_bus, N_time, N_scen, N_iter)
z_iter = Array{Float64}(N_storage, N_bus, N_scen, N_iter)
v_iter = Array{Float64}(N_bus, N_time, N_scen, N_iter)
a_iter = Array{Float64}(N_bus, N_time, N_scen, N_iter)
pg_iter = Array{Float64}(N_bus, N_time, N_scen, N_iter)
qg_iter = Array{Float64}(N_bus, N_time, N_scen, N_iter)
IC_iter = Array{Float64}(N_scen, N_iter)
OC_iter = Array{Float64}(N_scen, N_iter)
EC_iter = Array{Float64}(N_scen, N_iter)
OCinter_iter = Array{Float64}(N_scen, N_iter)
OCgen_iter = Array{Float64}(N_scen, N_iter)
OCloadshed_iter = Array{Float64}(N_scen, N_iter)
OCstorage_iter = Array{Float64}(N_scen, N_iter)
x_iter = Array{Float64}(N_storage, N_scen, N_iter)
sigma_d_iter = Array{Float64}(N_bus, N_time, N_scen, N_iter)
sigma_l_iter = Array{Float64}(N_line, N_time, N_scen, N_iter)
#########################
penalty1 = Array{Float64}(N_scen,1)
penalty2 = Array{Float64}(N_scen,1)
z_old = Array{Int64}(N_storage, N_bus, N_scen)
x_old = Array{Int64}(N_storage, N_scen)
w_old = PHmultiplier(zeros(N_storage, N_bus, N_scen),zeros(N_storage, N_scen))
w_new = PHmultiplier(zeros(N_storage, N_bus, N_scen),zeros(N_storage, N_scen))

#####################
convergence = Array{Float64}(N_iter,1)
convergence[:] = 777
penalty1_iter = Array{Float64}(N_scen, N_iter)
penalty2_iter = Array{Float64}(N_scen, N_iter)
z_bar_iter = Array{Float64}(N_storage, N_bus, N_iter)
x_bar_iter = Array{Float64}(N_storage, N_iter)
w_iter = PHmultiplier(zeros(N_storage, N_bus, N_scen, N_iter),zeros(N_storage, N_scen, N_iter))

#######
z_old = z_ini
x_old = x_ini
z_bar = sum( Prob[s] * z_old[:,:,s] for s in 1:N_scen)
x_bar = sum( Prob[s] * x_old[:,s] for s in 1:N_scen)

for kk = 1:N_iter
    println("###iteration begin###")
    for s = 1:N_scen
        w_new.z[:,:,s] = w_old.z[:,:,s] + rho.z .* (z_old[:,:,s] - z_bar[:,:] )
        w_new.x[:,s] = w_old.x[:,s] + rho.x .* (x_old[:,s] - x_bar[:] )
        w_new_temp = PHmultiplier(w_new.z[:,:,s],w_new.x[:,s])
        if s in S_N
            contin_temp = contin[:,:,s]
            scen_data_temp = ScenarioData(contin_temp,load24[s,:])
        elseif s in S_E
            contin_temp = contin[:,:,s]
            scen_data_temp = ScenarioData(contin_temp,mean(load24[S_N,:],1))
        end
        status[s],obj_value[s],penalty1[s],penalty2[s],simulation_time[s],ESlocation[:,:,s],SoC[:,:,s],u[:,:,:,s],z[:,:,s],v[:,:,s],a[:,:,s],pg[:,:,s],qg[:,:,s],IC[s],OC[s],EC[s],OCinter[s],OCgen[s],OCloadshed[s],OCstorage[s],x[:,s],sigma_d[:,:,s],sigma_l[:,:,s]  = PHsub_iter(scen_data_temp, gap, delayMode, z_bar, w_new_temp, rho, filename_Node, filename_Generator, filename_Line, filename_Storage, filename_Station, filename_SMP)
    end
    ###
    z_old = z
    x_old = x
    z_bar = sum( Prob[s] * z_old[:,:,s] for s in 1:N_scen)
    x_bar = sum( Prob[s] * x_old[:,s] for s in 1:N_scen)

    penalty1_iter[:,kk] = penalty1[:]
    penalty2_iter[:,kk] = penalty2[:]
    z_bar_iter[:,:,kk] = z_bar[:,:]
    x_bar_iter[:,kk] = x_bar[:]
    w_iter.z[:,:,:,kk] = w_new.z[:,:,:]
    w_iter.x[:,:,kk] = w_new.x[:,:]

    status_iter[:,kk] = status[:]
    obj_value_iter[:,kk] = obj_value[:]
    simulation_time_iter[:,kk] = simulation_time[:]
    ESlocation_iter[:,:,:,kk] = ESlocation[:,:,:]
    SoC_iter[:,:,:,kk] = SoC[:,:,:]
    u_iter[:,:,:,:,kk] = u[:,:,:,:]
    z_iter[:,:,:,kk] = z[:,:,:]
    v_iter[:,:,:,kk] = v[:,:,:]
    a_iter[:,:,:,kk] = a[:,:,:]
    pg_iter[:,:,:,kk] = pg[:,:,:]
    qg_iter[:,:,:,kk] = qg[:,:,:]
    IC_iter[:,kk] = IC[:]
    OC_iter[:,kk] = OC[:]
    EC_iter[:,kk] = EC[:]
    OCinter_iter[:,kk] = OCinter[:]
    OCgen_iter[:,kk] = OCgen[:]
    OCloadshed_iter[:,kk] = OCloadshed[:]
    OCstorage_iter[:,kk] = OCstorage[:]
    x_iter[:,:,kk] = x[:,:]
    sigma_d_iter[:,:,:,kk] = sigma_d[:,:,:]
    sigma_l_iter[:,:,:,kk] = sigma_l[:,:,:]
    ###

    w_old = w_new
    z_temp = Array{Int64}(N_storage , N_bus, N_scen)
    z_temp[:]=0
    for s = 1:N_scen
        z_temp[:,:,s] = z[:,:,s][:]
    end
    temp_conv = sum( Prob[s] * norm( z_temp[:,:,s][:] - z_bar[:,:][:] ) for s in 1:N_scen)
    obj_value_global_temp = sum(Prob[s]*obj_value[s] for s=1:N_scen)
    println("--------------------")
    println("iteration #:")
    display(kk)
    println()
    println("convergence:")
    display(temp_conv)
    println()
    println("ESlocation:")
    display(ESlocation)
    println()
    println("Objective values:")
    display(obj_value)
    println()
    println("Global objective value:")
    display(obj_value_global_temp)
    println()
    println("Simulation time:")
    display(simulation_time)
    println()
    convergence[kk] = temp_conv
    println("###iteration end###")
    println("--------------------")
end
###########
obj_value_global = Array{Float64}(N_iter,1)
obj_value_global = sum( Prob[s] * obj_value_iter[s,:] for s = 1:N_scen)
simulation_time_series = sum(simulation_time_iter)
simulation_time_parallel=sum(maximum(simulation_time_iter,1))
simulation_time_longest = maximum(simulation_time_iter)

@save "result.jld"


"""
Test for AVR model : EXST1 available in PSS/e
This case study defines a four bus system with an infinite bus in 1,
a GENSAL in bus 2 and a constant impedance load in bus 3
The GENSAL machine has the EXST1 and a HYGOV.
The disturbance is the outage of one line between buses 1 and 4
"""
##################################################
############### LOAD DATA ########################
##################################################

include(joinpath(TEST_FILES_DIR, "data_tests/test39.jl"))

##################################################
############### SOLVE PROBLEM ####################
##################################################

csv_file = joinpath(TEST_FILES_DIR, "benchmarks/psse//EXST1/results_PSSe.csv")

@testset "Test 39 EXST1 ResidualModel" begin
    path = joinpath(pwd(), "test-psse-exst1")
    !isdir(path) && mkdir(path)
    try
        # Define Simulation Problem
        sim = Simulation(
            ResidualModel,
            sys, #system
            path,
            (0.0, 20.0), #time span
            BranchTrip(1.0, Line, "BUS1-BUS4-i_1"), #Type of Fault
        )

        # Test Initial Condition
        diff_val = [0.0]
        res = get_init_values_for_comparison(sim)
        for (k, v) in test39_x0_init
            diff_val[1] += LinearAlgebra.norm(res[k] - v)
        end

        @test diff_val[1] < 1e-3

        # Obtain small signal results for initial conditions
        small_sig = small_signal_analysis(sim)
        eigs = small_sig.eigenvalues
        @test small_sig.stable

        # Test Eigenvalues
        @test LinearAlgebra.norm(eigs - test39_eigvals) < 1e-3

        # Solve problem
        @test execute!(sim, IDA(), dtmax = 0.005, saveat = 0.005) ==
              PSID.SIMULATION_FINALIZED
        results = read_results(sim)

        # Obtain results

        t_psid, v2_psid = get_voltage_magnitude_series(results, 2)
        _, v3_psid = get_voltage_magnitude_series(results, 3)
        _, ω_psid = get_state_series(results, ("generator-2-1", :ω))

        # Obtain PSSE results
        M = get_csv_data(csv_file)
        t_psse = M[:, 1]
        v1_psse = M[:, 2]
        v2_psse = M[:, 3]
        v3_psse = M[:, 4]
        v4_psse = M[:, 5]
        ω_psse = M[:, 6] .+ 1.0
        efd_psse = M[:, 7]

        # Test Transient Simulation Results

        @test LinearAlgebra.norm(t_psid - round.(t_psse, digits = 3)) == 0.0
        @test LinearAlgebra.norm(v2_psid - v2_psse, Inf) <= 1e-3
        @test LinearAlgebra.norm(v3_psid - v3_psse, Inf) <= 1e-3
        @test LinearAlgebra.norm(ω_psid - ω_psse, Inf) <= 1e-3

    finally
        @info("removing test files")
        rm(path, force = true, recursive = true)
    end
end


"""
    ProgressiveHedgingSolver(qpsolver::AbstractMathProgSolver; <keyword arguments>)

Return a progressive-hedging algorithm object specified. Supply `qpsolver`, a MathProgBase solver capable of solving quadratic problems.

The following penalty parameter update procedures are available
- `Fixed`:  Fixed penalty (default) ?Fixed for parameter descriptions.
- `adaptive`: Adaptive penalty update ?Adaptive for parameter descriptions.

The following execution policies are available
- `Serial`:  Classical progressive-hedging (default)
- `Synchronous`: Classical progressive-hedging run in parallel ?Synchronous for parameter descriptions.
- `Asynchronous`: Asynchronous progressive-hedging ?Asynchronous for parameter descriptions.

...
# Arguments
- `qpsolver::AbstractMathProgSolver`: MathProgBase solver capable of solving quadratic programs.
- `penalty::AbstractPenalizer = Fixed()`: Specify penalty update procedure (Fixed, Adaptive)
- `execution::AbstractExecuter = Serial`: Specify how algorithm should be executed (Serial, Synchronous, Asynchronous). Distributed variants requires worker cores.
- <keyword arguments>: Algorithm specific parameters, consult individual docstrings (see above list) for list of possible arguments and default values.
...

## Examples

The following solves a stochastic program `sp` created in `StochasticPrograms.jl` using the progressive-hedging algorithm with Ipopt as an `qpsolver`.

```jldoctest
julia> solve(sp,solver=ProgressiveHedgingSolver(IpoptSolver(print_level=0)))
Progressive Hedging Time: 0:00:06 (1315 iterations)
  Objective:  -855.8332803469432
  δ:          9.436947935542464e-7
:Optimal
```
"""
mutable struct ProgressiveHedgingSolver{S <: QPSolver,
                                        E <: Execution,
                                        P <: AbstractPenalizer} <: AbstractStructuredSolver
    qpsolver::S
    execution::E
    penalty::P
    crash::CrashMethod
    parameters::Dict{Symbol,Any}

    function ProgressiveHedgingSolver(qpsolver::QPSolver;
                                      execution::Execution = Serial(),
                                      penalty::AbstractPenalizer = Fixed(),
                                      crash::CrashMethod = Crash.None(), kwargs...)
        S = typeof(qpsolver)
        E = typeof(execution)
        P = typeof(penalty)
        return new{S, E, P}(qpsolver,
                            execution,
                            penalty,
                            crash,
                            Dict{Symbol,Any}(kwargs))
    end
end

function StructuredModel(stochasticprogram::StochasticProgram, solver::ProgressiveHedgingSolver)
    x₀ = solver.crash(stochasticprogram, solver.qpsolver)
    return ProgressiveHedging(stochasticprogram, solver.qpsolver, solver.execution, solver.penalty; solver.parameters...)
end

function add_params!(solver::ProgressiveHedgingSolver; kwargs...)
    push!(solver.parameters, kwargs...)
    for (k,v) in kwargs
        if k ∈ [:qpsolver, :execution, :penalty, :crash]
            setfield!(solver, k, v)
            delete!(solver.parameters, k)
        end
    end
    return nothing
end

function internal_solver(solver::ProgressiveHedgingSolver)
    return get_solver(solver.qpsolver)
end

function optimize_structured!(ph::AbstractProgressiveHedgingSolver)
    return ph()
end

function fill_solution!(stochasticprogram::StochasticProgram, ph::AbstractProgressiveHedgingSolver)
    # First stage
    first_stage = StochasticPrograms.get_stage_one(stochasticprogram)
    nrows, ncols = first_stage_dims(stochasticprogram)
    StochasticPrograms.set_decision!(stochasticprogram, ph.ξ)
    # stochasticprogram.redCosts = try
    #     getreducedcosts(ph.mastersolver.lqmodel)[1:ncols]
    # catch
    #     fill(NaN, ncols)
    # end
    # stochasticprogram.linconstrDuals = try
    #     getconstrduals(ph.mastersolver.lqmodel)[1:nrows]
    # catch
    #     fill(NaN, nrows)
    # end
    # Second stage
    fill_submodels!(ph, scenarioproblems(stochasticprogram))
end

function solverstr(solver::ProgressiveHedgingSolver)
    return "Progressive-hedging solver under $(str(solver.execution)) and $(str(solver.penalty))"
end


# pathToHDF5="D:\\dataSets\\forMainHDF5\\smallLiverDataSet.hdf5"

# using Revise,CUDA
# includet("D:\\projects\\vsCode\\JuliaMedPipeB\\tests\\includeAll.jl")
# using Main.GaussianPure, Main.HDF5saveUtils,Main.visualizationFromHdf5, Main.distinctColorsSaved
# using MedEye3d
# import MedEye3d.ForDisplayStructs.TextureSpec
# using Distributions
# using Clustering
# using IrrationalConstants
# using ParallelStencil
# # using PythonCall
# # using CondaPkg
# # CondaPkg.add.("monai");
# # CondaPkg.resolve()

# patienGroupName="0"
# z=7# how big is the area from which we collect data to construct probability distributions
# klusterNumb = 5# number of clusters - number of probability distributions we will use

# #******************for display

# #1) open HDF5 file and define additional arrays needed for our algorithm
# fid = openHDF5(pathToHDF5)
# listOfColorUsed= falses(18)

# #manual Modification array
# manualModif = TextureSpec{UInt8}(# choosing number type manually to reduce memory usage
#     name = "manualModif",
#     color =getSomeColor(listOfColorUsed)# automatically choosing some contrasting color
#     ,minAndMaxValue= UInt8.([0,1]) #important to keep the same number type as chosen at the bagining
#     ,isEditable = true ) # we will be able to manually modify this array in a viewer

# algoVisualization = TextureSpec{Float32}(
#     name = "algoOutput",
#     # we point out that we will supply multiple colors
#     isContinuusMask=true,
#     colorSet = [getSomeColor(listOfColorUsed),getSomeColor(listOfColorUsed)]
#     ,minAndMaxValue= Float32.([0,1])# values between 0 and 1 as this represent probabilities
#    )

#     addTextSpecs=Vector{TextureSpec}(undef,2)
#     addTextSpecs[1]=manualModif
#     addTextSpecs[2]=algoVisualization

# #2) primary display of chosen image and annotating couple points with a liver
# mainScrollDat= loadFromHdf5Prim(fid,patienGroupName,addTextSpecs,listOfColorUsed)
# #3) save manual modifications to HDF5
# saveManualModif(fid,patienGroupName , mainScrollDat)
# #4) filtering out from the manually modified array all set pixels and get constants needed for later evaluation of gaussian PDF

# manualModif= getArrByName("manualModif" ,mainScrollDat)
# image=  getArrByName("image" ,mainScrollDat)
# algoOutput= getArrByName("algoOutput" ,mainScrollDat)

# mainArrSize= size(image)

# #****************** constructing probability distributions

# ##coordinates of manually set points
# coordsss= GaussianPure.getCoordinatesOfMarkings(eltype(image),eltype(manualModif),  manualModif, image) |>
#     (seedsCoords) ->GaussianPure.getPatchAroundMarks(seedsCoords,z ) |>
#     (patchCoords) ->GaussianPure.allNeededCoord(patchCoords,z )

# #getting patch statistics - mean and covariance
# patchStats = GaussianPure.calculatePatchStatistics(eltype(image),Float64, coordsss, image)

# #separate distribution for each marked point
# distribs = map(patchStat-> fit(MvNormal, reduce(hcat,(patchStat)))  , patchStats  )

# #in order to reduce computational complexity  we will reduce the number of used distributions using kl divergence

# #we are comparing all distributions
# klDivs =map(outerDist->    map(dist->kldivergence( outerDist  ,dist), distribs  ), distribs  )
# klDivsInMatrix = reduce(hcat,(klDivs))
# #clustering with kmeans
# R = kmeans(klDivsInMatrix, klusterNumb; maxiter=200, display=:iter)

# #now identify indexes for some example distributions from each cluster
# indicies = zeros(Int64,klusterNumb )
# a = assignments(R) # get the assignments of points to clusters
# for i in 1:klusterNumb
#     for j in 1:length(distribs)
#         if(a[j] == i)
#             indicies[i]=j
#         end
#     end
# end
# indicies

# #ditributions from diffrent clusters
# chosenDistribs = map(ind->distribs[ind] ,indicies)

# #************************ applying  probability distributions to image

# ### getting constants from distributions

# """
# calculate log normalization constant from distribution
# """
# function mvnormal_c0(d::AbstractMvNormal)
#     ldcd = logdetcov(d)
#     return - (length(d) * oftype(ldcd, log2π) + ldcd) / 2
# end

# """
# get constants needed for applying probability distributions
#  return vector   1) logConst 2) mu1 3) mu2 4) invcov00 5)invcov01 6)invcov10 7)invcov11
# """
# function getDistrConstants(exampleDistr)
#     c0= mvnormal_c0(exampleDistr)
#     invCov= inv(exampleDistr.Σ)
#     return [c0,exampleDistr.μ[1],exampleDistr.μ[2],invCov[1,1],invCov[1,2],invCov[2,1],invCov[2,2]  ]
# end#getDistrConstants

# # creating matrix from constants
# allConstants = map(distr-> getDistrConstants(distr)  , chosenDistribs) |>
#                (vectOfvects)-> reduce(hcat, vectOfvects)

# #### defining CUDA kernel

# """
# utility macro to iterate in given range around given voxel
# """
# macro iterAround(ex   )
#     return esc(quote
#         for xAdd in -r:r
#             x= (threadIdx().x+ ((blockIdx().x -1)*CUDA.blockDim_x()))+xAdd
#             if(x>0 && x<=mainArrSize[1])
#                 for yAdd in -r:r
#                     y= (threadIdx().y+ ((blockIdx().y -1)*CUDA.blockDim_y()))+yAdd
#                     if(y>0 && y<=mainArrSize[2])
#                         for zAdd in -r:r
#                             z= (threadIdx().z+ ((blockIdx().z -1)*CUDA.blockDim_z()))+zAdd
#                             if(z>0 && z<=mainArrSize[3])
#                                 if((abs(xAdd)+abs(yAdd)+abs(zAdd)) <=r)
#                                     $ex
#                                 end
#                             end
#                         end
#                     end
#                 end
#             end
#         end
#     end)
# end

# """
# con - matrix of precalculated constants
# image - main image here computer tomography image
# mainArrSize - dimensions of image
# output - where we want to save the calculations
# r - size of the evaluated patch
# klusterNumb- number of clusters - number of probability distributions we will use
# """
# function applyGaussKernel(con,image,mainArrSize,output, r::Int,klusterNumb::Int)
#     for probDist in 1:klusterNumb
#         summ=0.0
#         sumCentered=0.0
#         lenn= UInt8(0)
#         #get mean
#         @iterAround begin
#             lenn=lenn+1
#             summ+=image[x,y,z]
#         end
#         summ=summ/lenn
#         #get standard deviation
#         @iterAround sumCentered+= ((image[x,y,z]-summ )^2)

#         #here we have standard deviation
#         sumCentered= sqrt(sumCentered/(lenn-1))
#         #centering - subtracting means...
#         summ=summ-con[2,probDist]
#         sumCentered=sumCentered-con[3,probDist]
#         #saving output
#         x= (threadIdx().x+ ((blockIdx().x -1)*CUDA.blockDim_x()))
#         y= (threadIdx().y+ ((blockIdx().y -1)*CUDA.blockDim_y()))
#         z= (threadIdx().z+ ((blockIdx().z -1)*CUDA.blockDim_z()))
#         if(x>0 && x<=mainArrSize[1] && y>0 && y<=mainArrSize[2] &&z>0 && z<=mainArrSize[3] )
#             output[x,y,z]=  max(exp(con[1,probDist]-( ((summ*con[4,probDist]+sumCentered*con[6,probDist])*summ+(summ*con[5,probDist]+sumCentered*con[7,probDist])*sumCentered)/2 ) ),output[x,y,z]  )
#         end
#     end#for
#     return
# end#main kernel

# # for simplicity not using the occupancy API - in production one rather should
# threads=(8,4,8)
# blocks = (cld(mainArrSize[1],threads[1]), cld(mainArrSize[2],threads[2])  , cld(mainArrSize[3],threads[3]))
# using CUDA
# algoOutputGPU=CuArray(algoOutput)
# imageGPU=CuArray(image)
# conGPU = CuArray(allConstants)
# @cuda threads=threads blocks=blocks applyGaussKernel(conGPU,imageGPU,mainArrSize,algoOutputGPU, 5,klusterNumb)
# #@cuda threads=threads blocks=blocks applyGaussKernel(conGPU,imageGPU,mainArrSize,algoOutputGPU, z,klusterNumb)
# copyto!(algoOutput,algoOutputGPU)
# sum(algoOutput)# just to check is anythink copied

# #copy and divide by max so will be easier to visualize
# algoOutputB= getArrByName("algoOutput" ,mainScrollDat)
# maxEl = maximum(algoOutputGPU)
# algoOutputB[:,:,:]=algoOutput./maxEl

# ### just to show how slow it would be to achieve the same on CPU

# # ## single thread
# # function getMaxProb(point)
# #     coords= getCartesianAroundPoint(point,z)
# #     xxx=getSampleMeanAndStd( Float64,Float64, coords , image  )
# #     return maximum(map(dist-> Distributions.pdf(dist, xxx),chosenDistribs))
# # end

# # output = map(getMaxProb, CartesianIndices(image))
# # maximum(output)
# # algoOutput[:,:,:]=output./maximum(output)

# # ## multithread
# # cartss = CartesianIndices(image)
# # Threads.@threads for i = 1:length(image)
# #     algoOutput[i] = getMaxProb(cartss[1])
# # end

# saveMaskbyName(fid,patienGroupName , mainScrollDat, "algoOutput")

# #******************************************************** relaxation labelling

# ########5)
# """
# in case of the relaxation labelling for 2D case algorithm basically can be simplified to
# iteratively look into the surrounding elements and increase value of a voxel given voxels around
# had high enough value and decrese otherwise

# adapted from https://discourse.julialang.org/t/3d-medical-app-stencil/64019/4
# """

# const USE_GPU = true
# using ParallelStencil
# using ParallelStencil.FiniteDifferences3D

# @init_parallel_stencil(CUDA, Float64, 3);

# #cutoff set manually to rate
# @parallel_indices (ix,iy,iz) function relaxationLabellKern(In, rate)
#     # 7-point Neuman stencil
#     if (ix>1 && iy>1 && iz>1 &&      ix<(size(In,1))&& iy<(size(In,2)) && iz<(size(In,3)))
#         In[ix,iy,iz] = ( (In[ix-1,iy  ,iz  ] >rate)+
#                           (In[ix-1,iy  ,iz  ]>rate)+ (In[ix+1,iy  ,iz  ]>rate) +
#                           (In[ix  ,iy-1,iz  ]>rate) + (In[ix  ,iy+1,iz  ]>rate) +
#                           (In[ix  ,iy  ,iz-1]>rate) + (In[ix  ,iy  ,iz+1]>rate) )/7.0

#     end
#     return
# end

# @views function relaxLabels(In, iterNumb,rate)
#     # Calculation
#     for i in 1:iterNumb
#         innerRate=rate + ((i/iterNumb)/3)
#         @parallel relaxationLabellKern(In,innerRate)
#     end#for
#     return
# end

# rate=0.15
# relaxLabels(algoOutputGPU,50,rate)

# copyto!(algoOutput,algoOutputGPU)
# Int(round(sum(algoOutput)))# just to check is anythink copied  #85162

# #copy and divide by max so will be easier to visualize
# algoOutputB= getArrByName("algoOutput" ,mainScrollDat)
# algoOutputB[:,:,:]=algoOutput

# #relaxLabels(algoOutputGPU,10)

# ###########7) displaying performance metrics

# # first we need to define the cutoff  over which we will decide that probability indicates that it is truly a liver
# #####simple  tresholding
# function tresholdingKernel(mainArrSize,output)

#         x= (threadIdx().x+ ((blockIdx().x -1)*CUDA.blockDim_x()))
#         y= (threadIdx().y+ ((blockIdx().y -1)*CUDA.blockDim_y()))
#         z= (threadIdx().z+ ((blockIdx().z -1)*CUDA.blockDim_z()))
#         if(x>0 && x<=mainArrSize[1] && y>0 && y<=mainArrSize[2] &&z>0 && z<=mainArrSize[3] )
#             output[x,y,z]=  (output[x,y,z]>0.5)
#         end
#     return
# end#main kernel

# # for simplicity not using the occupancy API - in production one rather should
# threads=(8,8,8)
# blocks = (cld(mainArrSize[1],threads[1]), cld(mainArrSize[2],threads[2])  , cld(mainArrSize[3],threads[3]))
# @cuda threads=threads blocks=blocks tresholdingKernel(mainArrSize,algoOutputGPU)

# using MedEval3D
# using MedEval3D.BasicStructs
# using MedEval3D.MainAbstractions
# conf= ConfigurtationStruct(md=true, dice=true)
# numberToLookFor = 1.0
# liverGold= getArrByName("liver" ,mainScrollDat)

# preparedDict=MedEval3D.MainAbstractions.prepareMetrics(conf)
# calculateAndDisplay(preparedDict,mainScrollDat, conf, numberToLookFor,CuArray(liverGold),algoOutputGPU )

# mainScrollDat.mainTextToDisp

# copyto!(algoOutput,algoOutputGPU)
# Int(round(sum(algoOutput)))# just to check is anythink copied  #85162
# #copy and divide by max so will be easier to visualize
# algoOutputB= getArrByName("algoOutput" ,mainScrollDat)
# algoOutputB[:,:,:]=algoOutput

# saveManualModif(fid,patienGroupName , mainScrollDat)

# close(fid)


import Adapt

using Oceananigans: short_show
using Oceananigans.Utils: user_function_arguments
using Oceananigans.Operators: assumed_field_location, index_and_interp_dependencies
using Oceananigans.Fields: show_location
using Oceananigans.Utils: tupleit

"""
    ContinuousForcing{X, Y, Z, P, F, D, I}

A callable object that implements a "continuous form" forcing function
on a field at the location `X, Y, Z` with optional parameters.
"""
struct ContinuousForcing{X, Y, Z, P, F, D, I, ℑ}
                          func :: F
                    parameters :: P
            field_dependencies :: D
    field_dependencies_indices :: I
     field_dependencies_interp :: ℑ

    # Non-public "temporary" constructor that stores func, parameters, and field_dependencies
    # for later regularization
    function ContinuousForcing(func, parameters, field_dependencies)
        field_dependencies = tupleit(field_dependencies)

        return new{Nothing, Nothing, Nothing,
                   typeof(parameters),
                   typeof(func),
                   typeof(field_dependencies),
                   Nothing,
                   Nothing}(func, parameters, field_dependencies, nothing, nothing)
    end

    # Non-public "final" constructor.
    function ContinuousForcing{X, Y, Z}(func, parameters=nothing, field_dependencies=(),
                                        field_dependencies_indices=(), field_dependencies_interp=()) where {X, Y, Z}
        return new{X, Y, Z,
                   typeof(parameters),
                   typeof(func),
                   typeof(field_dependencies),
                   typeof(field_dependencies_indices),
                   typeof(field_dependencies_interp)}(func, parameters, field_dependencies,
                                                      field_dependencies_indices, field_dependencies_interp)
    end
end

"""
    ContinuousForcing(func; parameters=nothing, field_dependencies=())

Construct a "continuous form" forcing with optional `parameters` and optional
`field_dependencies` on other fields in a model.

If neither `parameters` nor `field_dependencies` are provided, then `func` must be
callable with the signature

    `func(x, y, z, t)`

where `x, y, z` are the east-west, north-south, and vertical spatial coordinates, and `t` is time.

If `field_dependencies` are provided, the signature of `func` must include them.
For example, if `field_dependencies=(:u, :S)` (and `parameters` are _not_ provided), then
`func` must be callable with the signature

    `func(x, y, z, t, u, S)`

where `u` is assumed to be the `u`-velocity component, and `S` is a tracer. Note that any field
which does not have the name `u`, `v`, or `w` is assumed to be a tracer and must be present
in `model.tracers`.

If `parameters` are provided, then the _last_ argument to `func` must be `parameters`.
For example, if `func` has no `field_dependencies` but does depend on `parameters`, then
it must be callable with the signature

    `func(x, y, z, t, parameters)`

With `field_dependencies=(:u, :v, :w, :c)` and `parameters`, then `func` must be
callable with the signature

    `func(x, y, z, t, u, v, w, c, parameters)`

"""
ContinuousForcing(func; parameters=nothing, field_dependencies=()) =
    ContinuousForcing(func, parameters, field_dependencies)

"""
    regularize_forcing(forcing::ContinuousForcing, field, field_name, model_field_names)

Regularize `forcing::ContinuousForcing` by determining the indices of `forcing.field_dependencies`
in `model_field_names`, and associated interpolation functions so `forcing` can be used during
time-stepping `IncompressibleModel`.
"""
function regularize_forcing(forcing::ContinuousForcing, field, field_name, model_field_names)

    X, Y, Z = location(field)

    indices, interps = index_and_interp_dependencies(X, Y, Z,
                                                     forcing.field_dependencies,
                                                     model_field_names)

    return ContinuousForcing{X, Y, Z}(forcing.func, forcing.parameters, forcing.field_dependencies,
                                      indices, interps)
end

#####
##### Functions for calling ContinuousForcing in a time-stepping kernel
#####

@inline function (forcing::ContinuousForcing{X, Y, Z, F})(i, j, k, grid, clock, model_fields) where {X, Y, Z, F}

    args = user_function_arguments(i, j, k, grid, model_fields, forcing.parameters, forcing)

    return @inbounds forcing.func(xnode(X, i, grid), ynode(Y, j, grid), znode(Z, k, grid), clock.time, args...)
end

"""Show the innards of a `ContinuousForcing` in the REPL."""
Base.show(io::IO, forcing::ContinuousForcing{X, Y, Z, P}) where {X, Y, Z, P} =
    print(io, "ContinuousForcing{$P} at ", show_location(X, Y, Z), '\n',
        "├── func: $(short_show(forcing.func))", '\n',
        "├── parameters: $(forcing.parameters)", '\n',
        "└── field dependencies: $(forcing.field_dependencies)")

"""Show the innards of an "non-regularized" `ContinuousForcing` in the REPL."""
Base.show(io::IO, forcing::ContinuousForcing{Nothing, Nothing, Nothing, P}) where P =
    print(io, "ContinuousForcing{$P}", '\n',
        "├── func: $(short_show(forcing.func))", '\n',
        "├── parameters: $(forcing.parameters)", '\n',
        "└── field dependencies: $(forcing.field_dependencies)")

Adapt.adapt_structure(to, forcing::ContinuousForcing{X, Y, Z}) where {X, Y, Z} =
    ContinuousForcing{X, Y, Z}(Adapt.adapt(to, forcing.func),
                               Adapt.adapt(to, forcing.parameters),
                               nothing,
                               Adapt.adapt(to, forcing.field_dependencies_indices),
                               Adapt.adapt(to, forcing.field_dependencies_interp))


module SequentialQueryStrategies

import SVDD
import MLLabelUtils
import Statistics: mean

import ..QueryStrategies:
    QueryStrategy,
    MissingLabelTypeException,

    get_query_objects,
    knn_indices,
    knn_mean_dist,
    leave_out_one_cv_kde,
    multi_kde,
    qs_score

include("sequential_qs_base.jl")

include("TestPQs.jl")
include("RandomPQs.jl")
include("RandomOutlierPQs.jl")
include("MinimumMarginPQs.jl")
include("ExpectedMinimumMarginPQs.jl")
include("ExpectedMaximumEntropyPQs.jl")
include("MinimumLossPQs.jl")
include("HighConfidencePQs.jl")
include("DecisionBoundaryPQs.jl")
include("NeighborhoodBasedPQs.jl")
include("BoundaryNeighborCombinationPQs.jl")

export
    PoolQs,
    SequentialPQs,
    DataBasedPQs,
    ModelBasedPQs,
    HybridPQs,

    # data-based query strategies
    TestPQs, RandomPQs, MinimumMarginPQs, ExpectedMinimumMarginPQs, ExpectedMaximumEntropyPQs,
    MinimumLossPQs,
    # model-based query strategies
    RandomOutlierPQs, HighConfidencePQs, DecisionBoundaryPQs,
    # hybrid query strategies
    NeighborhoodBasedPQs, BoundaryNeighborCombinationPQs
end


__precompile__(true)

"""
The module for [Random123 Family](@ref).

Provide 8 RNG types:

- [`Threefry2x`](@ref)
- [`Threefry4x`](@ref)
- [`Philox2x`](@ref)
- [`Philox4x`](@ref)
- [`AESNI1x`](@ref)
- [`AESNI4x`](@ref)
- [`ARS1x`](@ref)
- [`ARS4x`](@ref)
"""
module Random123

export R123_USE_AESNI, set_counter!
include("common.jl")

export Threefry2x, Threefry4x
include("threefry.jl")

export Philox2x, Philox4x
include("philox.jl")

if R123_USE_AESNI
    export AESNI1x, AESNI4x
    include("aesni.jl")

    export ARS1x, ARS4x
    include("ars.jl")
end

end


#=
This file is part of the replication code for: Hasenzagl, T., Pellegrino, F., Reichlin, L., & Ricco, G. (2020). A Model of the Fed's View on Inflation.
Please cite the paper if you are using any part of the code for academic work (including, but not limited to, conference and peer-reviewed papers).
=#

function mwg_run(θ_unb::Array{Float64,1}, par::ParSsm, h::Int64, par_ind::BoolParSsm, par_size::SizeParSsm,
                  prior_opt::PriorOpt, MIN::Array{Float64, 1}, MAX::Array{Float64, 1}, opt_transf::Array{Int64, 1},
                  nDraws::Int64, burnin::Int64, mwg_const::Float64, algorithm_name::String, t::Int64, end_oos::Int64,
                  Σ::Array{Float64, 2}, gibbs_step::Int64)

     # -----------------------------------------------------------------------------------------------------------------
     # Initialisation
     # -----------------------------------------------------------------------------------------------------------------

     # Pre-allocate memory
     k             = size(par.T)[1];
     n, m          = size(par.y);
     distr_α       = zeros(k, m, nDraws-burnin);
     distr_fcst    = zeros(h, n, nDraws-burnin);
     distr_par     = Array{Any}(undef, nDraws-burnin);
     chain_θ_unb   = zeros(par_size.θ, nDraws);
     chain_θ_bound = zeros(par_size.θ, nDraws);

     # Initial values
     θ_bound           = get_par_bound(θ_unb, MIN, MAX, opt_transf);
     par.logposterior  = -Inf;
     par_prop          = copy(par);
     apriori_rejection = [0.0];

     # Parallel computing: random jumps for the proposal distribution
     n_workers        = length(workers());
     nDraws_worker    = convert(Int64, floor(nDraws/n_workers));
     θ_unb_jumps_pmap = pmap(get_mwg_jump,
                             [Σ for i=1:n_workers],
                             [mwg_const for i=1:n_workers],
                             [nDraws_worker for i=1:n_workers]);

     θ_unb_jumps = reduce(hcat, θ_unb_jumps_pmap);

     # Additional constants
     workers_vec = repeat(collect(1:n_workers), convert(Int64, floor(100/n_workers)), 1);
     n_fields    = length(fieldnames(typeof(par)));
     n_iter      = convert(Int64, ceil(n_fields/n_workers));

     # -----------------------------------------------------------------------------------------------------------------
     # Start estimation
     # -----------------------------------------------------------------------------------------------------------------

     for draw=1:nDraws

          last_draw = draw-1;

          if draw > 1 && mod(last_draw, 1000) == 0

               last_logposterior = round(par.logposterior, digits=2);

               # Display the estimation status
               if t != 0 && end_oos != 0
                    print("OOS > Running the $t-th iteration (out of $end_oos) \n");
               end

               print("$algorithm_name > Executed $last_draw draws (out of $nDraws): \n");
               print("- Logposterior: $last_logposterior\n")
               get_progress(chain_θ_unb, 1, last_draw, 1);
               if last_draw > burnin+100
                    get_progress(chain_θ_unb, burnin+1, last_draw, 1, "Acceptance rate (post burn-in)");
               end
               print("\n");
          end

          # ------------------------------------------------------------------------------------------------------------
          # Block 1: Propose θ* | θ_{draw-1}
          # ------------------------------------------------------------------------------------------------------------

          θ_prop_unb   = θ_unb + θ_unb_jumps[:, draw];
          θ_prop_bound = get_par_bound(θ_prop_unb, MIN, MAX, opt_transf);

          set_par!(θ_prop_bound, θ_prop_unb, par_prop, opt_transf, MIN, MAX, par_ind, par_size, prior_opt, apriori_rejection);

          if apriori_rejection[1] == 0

               # Accept / Reject
               if rand(Uniform(0, 1)) < exp(par_prop.logposterior-par.logposterior)

                    # Accept the proposed state-space parameters
                    # - Copy the fields of par_prop in par (in parallel)
                    begin
                         par_list = Array{Any}(undef, n_fields);

                         for i=1:n_iter

                              # Set end for the sync loop
                              end_sync_loop = n_workers;
                              if i*n_workers > n_fields
                                   end_sync_loop = n_fields-(i-1)*n_iter;
                              end

                              @sync for j=1:end_sync_loop

                                   ind_par = j+(i-1)*n_workers;

                                   # @async remotecall_fetch
                                   @async par_list[ind_par] =
                                        remotecall_fetch(copy, workers_vec[j], getfield(par_prop, j+(i-1)*n_workers));
                              end
                         end
                    end

                    par = ParSsm(par_list...);

                    # Copy θ_unb and θ_bound in parallel
                    @sync begin
                         @async θ_unb   = remotecall_fetch(copy, workers_vec[1], θ_prop_unb);
                         @async θ_bound = remotecall_fetch(copy, workers_vec[2], θ_prop_bound);
                    end
               end
          end

          # Chains
          chain_θ_unb[:, draw]   = θ_unb;
          chain_θ_bound[:, draw] = θ_bound;

          # ------------------------------------------------------------------------------------------------------------
          # Block 2: Draw alpha (Gibbs step)
          # ------------------------------------------------------------------------------------------------------------

          if draw > burnin && gibbs_step == 1

               # Draw α conditional to the state-space parameters
               α_draw, _                  = kalman_diffuse!(par, 0, 1, 1);
               distr_α[:, :, draw-burnin] = α_draw;
               distr_par[draw-burnin]     = par;

               # Forecast
               density_fcst_draw             = (par.Z*α_draw)';
               density_fcst_draw             = density_fcst_draw[end-h+1:end, :];
               distr_fcst[:, :, draw-burnin] = density_fcst_draw;
          end
     end

     return chain_θ_unb, chain_θ_bound, distr_α, distr_fcst, par, distr_par;
end


###################################################
## routingviz.jl
## network and routing visualizer
###################################################

"""
    `RoutingViz`: Network visualizer that also shows routing information.
    If `P` is clicked, enters path visualization mode.
    The shortest path will be shown, starting from the selected node. click to fix the path.

"""
mutable struct RoutingViz <: NetworkVisualizer
    # Mandatory attributes
    network::Network
    window::RenderWindow
    view::View
    nodes::Vector{CircleShape}
    roads::Dict{Tuple{Int,Int},Line}
    nodeRadius::Float64
    colors::VizColors

    "basic visualizer"
    networkviz::NetworkViz
    "the routing information"
    routing::RoutingPaths
    "if path-mode is on"
    pathMode::Bool
    "path is frozen"
    pathFrozen::Bool
    "destination node (path destination)"
    destNode::Int
    "the path to show"
    path::Vector{Line}

    "contructor"
    function RoutingViz(r::RoutingPaths; colors::VizColors=RelativeSpeedColors(r))
        obj = new()
        obj.network  = r.network
        obj.routing = r
        obj.colors = colors

        obj.networkviz = NetworkViz(r.network; colors=colors)
        return obj
    end
end

function visualInit(v::RoutingViz)
    #give visuals to nodeinfo
    copyVisualData(v,v.networkviz)

    v.pathMode = false
    v.pathFrozen = false
    v.destNode = 1
end

function visualStartUpdate(v::RoutingViz,frameTime::Float64)
    if v.pathMode && !v.pathFrozen
        mouseCoord = pixel2coords(v.window, get_mousepos(v.window))
        nodeId = knn(v.networkviz.tree,[Float64(mouseCoord.x),-Float64(mouseCoord.y)],1)[1][1]

        if nodeId != v.destNode
            if v.destNode != v.networkviz.selectedNode
                # normal color to previous node
                set_fillcolor(v.nodes[v.destNode], nodeColor(v.colors, v.network.nodes[v.destNode]))
                # red for new node
                set_fillcolor(v.nodes[nodeId], SFML.red)
            end
            v.destNode = nodeId
            set_title(v.window, string("Routing path: ", v.networkviz.selectedNode, " => ", v.destNode))
            drawPath(v)
        end
    end
end

function visualEndUpdate(v::RoutingViz, frameTime::Float64)
    if v.pathMode #print the path on top of other roads
        for segment in v.path
            draw(v.window, segment)
        end
    end
end

function visualRedraw(v::RoutingViz)
    visualRedraw(v.networkviz)
    if v.pathMode
        # redraw the path
        drawPath(v)
        set_fillcolor(v.nodes[v.destNode], SFML.red)
    end
end

function visualEvent(v::RoutingViz, event::Event)
    if get_type(event) == EventType.KEY_PRESSED && get_key(event).key_code == KeyCode.P
        if v.pathMode
            v.pathMode = false
            set_title(v.window, "")
            # normal color to previous node
            set_fillcolor(v.nodes[v.destNode], nodeColor(v.colors, v.network.nodes[v.destNode]))
        else
            v.pathMode = true
            v.pathFrozen = false
            v.destNode = v.networkviz.selectedNode
            v.path = []
        end
    elseif v.pathMode
        if get_type(event) == EventType.MOUSE_BUTTON_PRESSED && get_mousebutton(event).button == MouseButton.LEFT
            v.pathFrozen = !v.pathFrozen
        end
    else
        visualEvent(v.networkviz,event)
    end
end

"""
    `drawPath` create the path drawings
"""
function drawPath(v::RoutingViz)
    v.path = [Line(copy(v.roads[o, d].rect)) for (o,d) in getPathEdges(v.routing, v.networkviz.selectedNode, v.destNode)]
    for line in v.path
        set_thickness(line, get_thickness(line)*4.)
        set_fillcolor(line, SFML.Color(0, 0, 125))
    end
end


using DynamicRemoteCall
using Distributed

addprocs(3)

@everywhere function testfunc()
    r = rand()/10
    sleep(r)
    return r
end

results = zeros(100)

begin
    starttime = time()
    dynamicremotecall!(testfunc, results)
    endtime = time()
    (endtime-starttime)*nworkers()/sum(results)
end


using Setfield
replacein(x, oldnode, newnode) = x
replacein(x::Tuple, oldnode, newnode) = tuple([replacein(m, oldnode, newnode) for m in x]...)
replacein(x::NamedTuple, oldnode, newnode) = (;[k => replacein(x[k], oldnode, newnode) for k in keys(x)]...)

function replacein(x::T, oldnode, newnode) where {T<:Union{AbstractNode, AbstractMillModel}}
    x === oldnode && return(newnode)
    fields = map(f -> replacein(getproperty(x, f), oldnode, newnode), fieldnames(T))
    n = nameof(T)
    p = parentmodule(T)
    eval(:($p.$n))(fields...)
end

function replacein(x::LazyNode{N}, oldnode, newnode) where {N}
    x === oldnode && return(newnode)
    LazyNode{N}(replacein(x.data, oldnode, newnode))
end

function replacein(x::LazyModel{N}, oldnode, newnode) where {N}
    x === oldnode && return(newnode)
    LazyModel{N}(replacein(x.m, oldnode, newnode))
end

function findin(x, node)
    x === node && return(@lens _)
    return(nothing)
end

function findin(x::T, node) where {T<:Union{AbstractNode, AbstractMillModel}}
    x === node && return(@lens _)
    for k in fieldnames(T)
        l = findin(getproperty(x, k), node)
        if l != nothing
            lo = Setfield.PropertyLens{k}() ∘ l
            return(lo)
        end
    end
    return(nothing)
end

function findin(x::NamedTuple, node)
	x === node && return(@lens _)
    for k in keys(x)
    	l = findin(x[k], node)
    	if l != nothing
    		lo = Setfield.PropertyLens{k}() ∘ l
    		return(lo)
    	end
    end
    return(nothing)
end

function findin(x::Tuple, node)
    error("findin does not support Tuples due to restrinctions of Lens from Setfield.")
end

# function findin(x::T, node) where {T<:ArrayNode}
# 	x === node && return(@lens _)
# 	x.data === node && return(@lens _.data)
#     return(nothing)
# end

# function findin(x::LazyNode{N}, node) where {N}
# 	x === node && return(@lens _)
# 	x.data === node && return(@lens _.data)
#     return(nothing)
# end

# function findin(x::AbstractBagNode, node)
# 	x === node && return(@lens _)
# 	l = findin(x.data, node)
# 	if l != nothing
# 		return((@lens _.data) ∘  l)
# 	end
#     return(nothing)
# end


"""
    TermSet <: AbstractSet{AbstractTerm}

Wrapped `Set{AbstractTerm}` that specifies a collection of terms.
Commonly used methods for `Set` work in the same way for `TermSet`.

Compared with `StatsModels.TermOrTerms`,
it does not maintain order of terms
but is more suitable for dynamically constructed terms.
"""
struct TermSet <: AbstractSet{AbstractTerm}
    terms::Set{AbstractTerm}
    TermSet(set::Set{AbstractTerm}) = new(set)
end

"""
    TermSet([itr])
    TermSet(ts::Union{Int, Symbol, AbstractTerm}...)

Construct a [`TermSet`](@ref) from a collection of terms.
Instead of passing an iterable collection,
one may pass the terms as arguments directly.
In the latter case, any `Int` or `Symbol` will be converted to a `Term`.
See also [`termset`](@ref), which is an alias for the constructor.
"""
TermSet() = TermSet(Set{AbstractTerm}())
TermSet(ts) = TermSet(Set{AbstractTerm}(ts))
TermSet(ts::Union{Int, Symbol, AbstractTerm}...) =
    TermSet(Set{AbstractTerm}(t isa AbstractTerm ? t : term(t) for t in ts))

"""
    termset([itr])
    termset(ts::Union{Int, Symbol, AbstractTerm}...)

Construct a [`TermSet`](@ref) from a collection of terms.
Instead of passing an iterable collection,
one may pass the terms as arguments directly.
In the latter case, any `Int` or `Symbol` will be converted to a `Term`.
"""
termset() = TermSet(Set{AbstractTerm}())
termset(ts) = TermSet(Set{AbstractTerm}(ts))
termset(ts::Union{Int, Symbol, AbstractTerm}...) =
    TermSet(Set{AbstractTerm}(t isa AbstractTerm ? t : term(t) for t in ts))

Base.isempty(ts::TermSet) = isempty(ts.terms)
Base.length(ts::TermSet)  = length(ts.terms)
Base.in(x, ts::TermSet) = in(x, ts.terms)
Base.push!(ts::TermSet, x) = push!(ts.terms, x)
Base.pop!(ts::TermSet, x) = pop!(ts.terms, x)
Base.pop!(ts::TermSet, x, default) = pop!(ts.terms, x, default)
Base.delete!(ts::TermSet, x) = delete!(ts.terms, x)
Base.empty!(ts::TermSet) = empty!(ts.terms)

Base.eltype(::Type{TermSet}) = AbstractTerm
Base.iterate(ts::TermSet, i...) = iterate(ts.terms, i...)
Base.emptymutable(::TermSet, ::Type{<:AbstractTerm}) = termset()

const Terms{N} = NTuple{N, AbstractTerm} where N

"""
    eachterm(t)

Return an iterable collection of terms in `t`.
"""
eachterm(@nospecialize(t::AbstractTerm)) = (t,)
eachterm(@nospecialize(t::Terms)) = t

==(@nospecialize(a::Terms), @nospecialize(b::Terms)) =
    length(a)==length(b) && all(t->t in b, a)
==(@nospecialize(a::InteractionTerm), @nospecialize(b::InteractionTerm)) =
    a.terms==b.terms
==(@nospecialize(a::FormulaTerm), @nospecialize(b::FormulaTerm)) =
    a.lhs==b.lhs && a.rhs==b.rhs

"""
    TreatmentTerm{T<:AbstractTreatment} <: AbstractTerm

A term that contains specifications on treatment and parallel trends assumption.
See also [`treat`](@ref).

# Fields
- `sym::Symbol`: the column name of data representing treatment status.
- `tr::T`: a treatment type that specifies the causal parameters of interest.
- `pr::P`: a parallel type that specifies the parallel trends assumption.
"""
struct TreatmentTerm{T<:AbstractTreatment,P<:AbstractParallel} <: AbstractTerm
    sym::Symbol
    tr::T
    pr::P
end

"""
    treat(s::Symbol, t::AbstractTreatment, p::AbstractParallel)

Construct a [`TreatmentTerm`](@ref) with fields set by the arguments.
"""
treat(s::Symbol, t::AbstractTreatment, p::AbstractParallel) = TreatmentTerm(s, t, p)

"""
    treat(s::Term, t::FunctionTerm{FT}, p::FunctionTerm{FP})

Construct a [`TreatmentTerm`](@ref) with fields set by `s.sym`,
`FT.instance(t.args_parsed...)` and `FP.instance(p.args_parsed...)`.
This method is called by [`parse_treat(formula)`](@ref)
when the formula is constructed by `@formula`.
"""
treat(s::Term, t::FunctionTerm{FT}, p::FunctionTerm{FP}) where {FT,FP} =
    treat(s.sym, FT.instance(t.args_parsed...), FP.instance(p.args_parsed...))

"""
    hastreat(t)

Determine whether the term `t` contains an instance of [`TreatmentTerm`](@ref).
"""
hastreat(::TreatmentTerm) = true
hastreat(::FunctionTerm{typeof(treat)}) = true
hastreat(t::InteractionTerm) = any(hastreat(x) for x in t.terms)
hastreat(::AbstractTerm) = false
hastreat(@nospecialize(t::FormulaTerm)) = any(hastreat(x) for x in eachterm(t.rhs))

"""
    parse_treat(formula::FormulaTerm)

Extract terms related to treatment specifications from `formula`.

# Returns
- `TreatmentTerm`: the unique `TreatmentTerm` contained in the `formula`.
- `TermSet`: a set of terms that are interacted with the `TreatmentTerm`.
- `TermSet`: a set of remaining terms in `formula.rhs`.

Error will be raised if either existence or uniqueness of the `TreatmentTerm` is violated.
"""
function parse_treat(@nospecialize(formula::FormulaTerm))
    mettreat = false
    treatterm = nothing
    ints = TermSet()
    xterms = TermSet()
    for term in eachterm(formula.rhs)
        if term isa TreatmentTerm
            if !mettreat
                mettreat = true
                treatterm = term
            else
                throw(ArgumentError("cannot accept more than one TreatmentTerm"))
            end
        elseif term isa FunctionTerm{typeof(treat)}
            if !mettreat
                mettreat = true
                treatterm = treat(term.args_parsed...)
            else
                throw(ArgumentError("cannot accept more than one TreatmentTerm"))
            end
        elseif term isa InteractionTerm
            if hastreat(term)
                !mettreat ||
                    throw(ArgumentError("cannot accept more than one TreatmentTerm"))
                for t in term.terms
                    if t isa TreatmentTerm
                        if !mettreat
                            mettreat = true
                            treatterm = t
                        else
                            throw(ArgumentError("cannot accept more than one TreatmentTerm"))
                        end
                    elseif t isa FunctionTerm{typeof(treat)}
                        if !mettreat
                            mettreat = true
                            treatterm = treat(t.args_parsed...)
                        else
                            throw(ArgumentError("cannot accept more than one TreatmentTerm"))
                        end
                    else
                        push!(ints, t)
                    end
                end
            else
                push!(xterms, term)
            end
        else
            push!(xterms, term)
        end
    end
    mettreat || throw(ArgumentError("no TreatmentTerm is found"))
    return treatterm::TreatmentTerm, ints, xterms
end

isintercept(t::AbstractTerm) = t in (InterceptTerm{true}(), ConstantTerm(1))
isomitsintercept(t::AbstractTerm) =
    t in (InterceptTerm{false}(), ConstantTerm(0), ConstantTerm(-1))

"""
    parse_intercept(ts::TermSet)

Remove any `ConstantTerm` or `InterceptTerm`
and return Boolean values indicating whether terms explictly requiring
including/excluding the intercept exist.

This function is useful for obtaining a unique way of specifying the intercept
before going through the `schema`--`apply_schema` pipeline defined in `StatsModels`.
"""
function parse_intercept!(ts::TermSet)
    hasintercept = false
    hasomitsintercept = false
    for t in ts
        if isintercept(t)
            delete!(ts, t)
            hasintercept = true
        end
        if isomitsintercept(t)
            delete!(ts, t)
            hasomitsintercept = true
        end
    end
    return hasintercept, hasomitsintercept
end

schema(ts::TermSet, data, hints=nothing) =
    Schema(t=>concrete_term(t, VecColumnTable(data), hints) for t in ts)

concrete_term(t::Term, dt::VecColumnTable, hint) =
    concrete_term(t, getproperty(dt, t.sym), hint)
concrete_term(t::Term, dt::VecColumnTable, hints::Dict{Symbol}) =
    concrete_term(t, getproperty(dt, t.sym), get(hints, t.sym, nothing))
concrete_term(::Term, ::VecColumnTable, hint::AbstractTerm) = hint

termvars(ts::TermSet) = mapreduce(termvars, union, ts, init=Symbol[])


@testset "Trajectory" begin

    s = get_simulation("flamingo", "flat_2D_lc", "flat")
    ref_traj = deepcopy(ContactImplicitMPC.get_trajectory(s.model, s.env,
        joinpath(module_dir(), "src/dynamics/flamingo/gaits/gait_forward_36_4.jld2"),
        load_type = :split_traj_alt))

    N_repeat = 5
    N_sample = 1
    dupl_traj = repeat_ref_traj(ref_traj, N_repeat, idx_shift=[1])
    @test all(ContactImplicitMPC.tracking_error(ref_traj, dupl_traj, N_sample, idx_shift=[1]) .== 0.0)

end


"""
    scatterlines(xs, ys, [zs]; kwargs...)

Plots `lines` between sets of x and y coordinates provided,
as well as plotting those points using `scatter`.

## Attributes
$(ATTRIBUTES)
"""
@recipe(ScatterLines) do scene
    merge(default_theme(scene, Scatter), default_theme(scene, Lines))
end

function plot!(p::Combined{scatterlines, <:NTuple{N, Any}}) where N
    plot!(p, Lines, attributes_from(Lines, p), p[1:N]...)
    plot!(p, Scatter, attributes_from(Scatter, p), p[1:N]...)
end