using ShowItLikeYouBuildIt
using Base.Test
push!(LOAD_PATH, @__DIR__)
using ZArrays
pop!(LOAD_PATH)

const showit_old = VERSION < v"0.7.0-DEV.1790"
const eltype_string = showit_old ? "element type" : "eltype"

if showit_old
    printdims(io, dims) = print(io, dims)
else
    printdims(io, dims) = _printdims(io, dims...)
    _printdims(io, d1, d...) = (print(io, d1, ", "); _printdims(io, d...))
    _printdims(io, d1) = print(io, d1)
end
function printdims(dims)
    io = IOBuffer()
    printdims(io, dims)
    String(take!(io))
end

if showit_old
    for T in (Float64, Bool, Int8, UInt16, Symbol, String, Union{})
        @test type_complexity(T) == 1
    end
    @test type_complexity(Union{Float32,Int8}) == 2
    @test type_complexity(Array{Int,2}) == 3
    @test type_complexity(Array{Complex{Float32},2}) == 4
    @test type_complexity(Array{Array{Complex{Float32},1},2}) == 6
end

@test ShowItLikeYouBuildIt.dimstring(()) == "0-dimensional"

@test sprint(showarg, 3.0) == "::Float64"
@test sprint(showarg, Float64) == "::Type{Float64}"

# Display of objects

# Note that these next tests alter how SubArrays, ReshapedArrays, and
# PermutedDimsArrays are displayed; it's probably best to run this
# only in a "throwaway" julia session.

function ShowItLikeYouBuildIt.showarg(io::IO, v::SubArray)
    print(io, "view(")
    showarg(io, parent(v))
    print(io, ", ", join(v.indexes, ", "))
    print(io, ')')
end

function ShowItLikeYouBuildIt.showarg(io::IO, A::PermutedDimsArray{T,N,perm}) where {T,N,perm}
    print(io, "PermutedDimsArray(")
    showarg(io, parent(A))
    print(io, ", ", perm, ')')
end

function ShowItLikeYouBuildIt.showarg(io::IO, A::Base.ReshapedArray)
    print(io, "reshape(")
    showarg(io, parent(A))
    print(io, ", ")
    printdims(io, A.dims)
    print(io, ')')
end

Base.summary(A::SubArray) = summary_build(A)
Base.summary(A::Base.PermutedDimsArrays.PermutedDimsArray) = summary_build(A)
Base.summary(A::Base.ReshapedArray) = summary_build(A)

a = rand(3,5,7)
v = view(a, :, 3, 2:5)
view_name = showit_old ? "view(::Array{Float64,3}, $(v.indexes[1]), 3, 2:5)" :
                         "view(::Array{Float64,3}, :, 3, 2:5)"
@test summary(v) == "3×4 $view_name with $eltype_string Float64"

c = reshape(v, 4, 3)
str = summary(c)
dimstr = printdims((4, 3))
@test str == "4×3 reshape($view_name, $dimstr) with $eltype_string Float64"

a = reshape(1:24, 3, 4, 2)
b = PermutedDimsArray(a, (2,3,1))
str = summary(b)
intstr = string("Int", Sys.WORD_SIZE)
dimstr = printdims((3, 4, 2))
@test str == "4×2×3 PermutedDimsArray(reshape(::UnitRange{$intstr}, $dimstr), $((2, 3, 1))) with $eltype_string $intstr"

o = ZArray(rand(3,5))
vo = view(o, 0:2:2, :)
if showit_old
    @test summary(vo) == "Base.OneTo(2)×ZArrays.ZeroRange(5) view(::ZArrays.ZArray{Float64,2}, 0:2:2, $(vo.indexes[end])) with $eltype_string Float64"
else
    @test summary(vo) == "view(::ZArrays.ZArray{Float64,2}, 0:2:2, :) with eltype Float64 with indices Base.OneTo(2)×ZArrays.ZeroRange(5)"
end

if showit_old
    # Adding cthresh is deprecated
    Base.summary(A::SubArray) = summary_build(A,1000)
    @test summary(v) == "3×4 SubArray{Float64,2,Array{Float64,3},Tuple{$(typeof(v.indexes[1])),$Int,UnitRange{$Int}},false}"
end

nothing


module FMMTrees

#using VectorBackedLists
#using AbstractTrees

#include("lists.jl")
include("trees.jl")
include("simpletrees.jl")
include("pointerbasedtrees.jl")
#include("blocktrees.jl")
#include("hmatrix.jl")
#include("cluster_tree.jl")
#include("aca.jl")

include("octree.jl")
include("leveltree.jl")
include("blocktree.jl")

end # module


module EventsWrapper

using JavaCall

const Action = @jimport org.hipparchus.ode.events.Action
const BracketedRealFieldUnivariateSolver = @jimport org.hipparchus.analysis.solvers.BracketedRealFieldUnivariateSolver
const BracketedUnivariateSolver = @jimport org.hipparchus.analysis.solvers.BracketedUnivariateSolver
const Class = @jimport java.lang.Class
const Enum = @jimport java.lang.Enum
const EventFilter = @jimport org.hipparchus.ode.events.EventFilter
const EventState = @jimport org.hipparchus.ode.events.EventState
const EventState_EventOccurrence = @jimport org.hipparchus.ode.events.EventState$EventOccurrence
const FieldEventState = @jimport org.hipparchus.ode.events.FieldEventState
const FieldEventState_EventOccurrence = @jimport org.hipparchus.ode.events.FieldEventState$EventOccurrence
const FieldODEEventHandler = @jimport org.hipparchus.ode.events.FieldODEEventHandler
const FieldODEState = @jimport org.hipparchus.ode.FieldODEState
const FieldODEStateAndDerivative = @jimport org.hipparchus.ode.FieldODEStateAndDerivative
const FieldODEStateInterpolator = @jimport org.hipparchus.ode.sampling.FieldODEStateInterpolator
const FilterType = @jimport org.hipparchus.ode.events.FilterType
const JString = @jimport java.lang.String
const ODEEventHandler = @jimport org.hipparchus.ode.events.ODEEventHandler
const ODEState = @jimport org.hipparchus.ode.ODEState
const ODEStateAndDerivative = @jimport org.hipparchus.ode.ODEStateAndDerivative
const ODEStateInterpolator = @jimport org.hipparchus.ode.sampling.ODEStateInterpolator
const Object = @jimport java.lang.Object
const Optional = @jimport java.util.Optional
const RealFieldElement = @jimport org.hipparchus.RealFieldElement
const Transformer = @jimport org.hipparchus.ode.events.Transformer

export Action
export EventFilter
export EventState
export FieldEventState
export FieldODEEventHandler
export FilterType
export ODEEventHandler
export Transformer
export compare_to
export describe_constable
export do_event
export equals
export evaluate_step
export event_occurred
export g
export get_class
export get_convergence
export get_declaring_class
export get_event_handler
export get_event_time
export get_max_check_interval
export get_max_iteration_count
export hash_code
export init
export name
export notify
export notify_all
export ordinal
export reinitialize_begin
export reset_state
export to_string
export try_advance
export value_of
export values
export wait

include("action.jl")
include("event_filter.jl")
include("event_state.jl")
include("field_event_state.jl")
include("field_ode_event_handler.jl")
include("filter_type.jl")
include("ode_event_handler.jl")
include("transformer.jl")

end


#
#  OutputData.jl
#  CollagenPolymer
#
#  Created by Christopher Revell on 30/03/2020.
#
#

module OutputData

# Import Julia packages
using DelimitedFiles
using LinearAlgebra

# Import local program modules

function outputData(pos,outfile)#,t,tMax)#,nMonomers,nDomains,σ)

    writedlm(outfile,pos,", ")
    flush(outfile)

    # Measure monomer lengths
    # for i=0:nMonomers-1
    #     l = 0
    #     for j=1:nDomains-1
    #         dx = pos[i*nDomains+j+1,:]-pos[i*nDomains+j,:]
    #         l += sqrt(dot(dx,dx))
    #     end
    #     println("$(i+1) => $(l+σ), nDomains=$nDomains")
    # end

    return nothing

end

export outputData

end


println("Playing Penney's Game Against the computer.")

if randbool()
    mach = autobet()
    println(@sprintf "The computer bet first, chosing %s." pgdecode(mach))
    println("Now you can bet.")
    human = pgencode(humanbet())
else
    println("You bet first.")
    human = pgencode(humanbet())
    mach = autobet(human)
end
print(@sprintf "You bet %s " pgdecode(human))
println(@sprintf "and the computer bet %s." pgdecode(mach))


# Copyright 2020 Google LLC

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
@enum State floor empty occupied

function Base.show(io::Core.IO, state::State)
    if state == empty
        print(io, 'L')
    elseif state == occupied
        print(io, '#')
    else
        print(io, '.')
    end
end

function Base.show(io::Core.IO, state::AbstractArray{State, 2})
    height, width = size(state)
    for i in 1:height
        for j in 1:width
            show(io, state[i, j])
        end
        print(io, '\n')
    end
end

function parsechar(c)::State
    if c == '.'
        return floor
    elseif c == 'L'
        return empty
    elseif c == '#'
        return occupied
    end
    throw(ErrorException("Unrecognized character"))
end

function adjacents(state::AbstractArray{State, 2}, i, j)
    height, width = size(state)
    [(-1, -1), (-1, 0), (-1, +1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)] |>
        dirs->map(c->(x=i+c[1], y=j+c[2]), dirs) |>
        coords->filter(c -> !(c.x < 1 || c.y < 1 || c.x > height || c.y > width), coords) |>
        coords->filter(c -> state[c.x, c.y] != floor, coords) |>
        coords->map(c->(c.x, c.y), coords)
end

function nextstate(state::State, neighbors::AbstractArray{State}, threshold_to_leave::Int)
    if state == empty
        if count(s->s==occupied, neighbors) == 0
            return occupied
        end
    elseif state == occupied
        if count(s->s==occupied, neighbors) >= threshold_to_leave
            return empty
        end
    end
    return state
end

function coords_visible(state::AbstractArray{State, 2}, i, j)
    height, width = size(state)
    neighbor_dirs = [(-1, -1), (-1, 0), (-1, +1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
    coords = Tuple{Int, Int}[]
    for dir ∈ neighbor_dirs
        x = i
        y = j
        while true
            x = x + dir[1]
            y = y + dir[2]
            if x < 1 || x > height || y < 1 || y > width
                break
            end
            if state[x, y] != floor
                push!(coords, (x, y))
                break
            end
        end
    end
    coords
end

function find_relevant_coords(state::AbstractArray{State, 2}, relevancy_function)
    height, width = size(state)
    relevant_coords = Dict{Tuple{Int, Int}, AbstractArray{Tuple{Int, Int}}}()
    for i ∈ 1:height
        for j ∈ 1:width
            relevant_coords[(i, j)] = relevancy_function(state, i, j)
        end
    end
    relevant_coords
end

function parseinput(lines)
    height = length(lines)
    width = length(lines[1])
    state = Array{State}(undef, (height, width))
    for i ∈ 1:height
        for j ∈ 1:width
            state[i, j] = parsechar(lines[i][j])
        end
    end
    state
end

function evolve(state::AbstractArray{State, 2}, relevant_coords, threshold_to_leave::Int)
    next = copy(state)
    for i ∈ 1:size(state)[1]
        for j ∈ 1:size(state)[2]
            neighbors = map(ind->state[ind[1], ind[2]], relevant_coords[(i, j)])
            next[i, j] = nextstate(state[i, j], neighbors, threshold_to_leave)
        end
    end
    next
end

function part1(filename="day11input.txt")
    state = parseinput(readlines(filename))
    relevant_coords = find_relevant_coords(state, adjacents)

    prev = nothing
    while prev != state
        prev = state
        state = evolve(state, relevant_coords, 4)
    end
    return count(s->s==occupied, state)
end

function part2(filename="day11input.txt")
    state = parseinput(readlines(filename))
    relevant_coords = find_relevant_coords(state, coords_visible)
    prev = nothing
    while prev != state
        prev = state
        state = evolve(state, relevant_coords, 5)
    end
    return count(s->s==occupied, state)
end


using simulationQP934

data = Data()
md1  = MDinput(dt=1)
md2  = MDinput(dt=0.1)
md3  = MDinput(dt=0.01)
md4  = MDinput(dt=0.001)

println("calculos para dt = 1")
md(data,md1)
println("  ")

println("calculos para dt = 0.1")
md(data,md2)
println("  ")

println("calculos para dt = 0.01")
md(data,md3)
println("  ")

println("calculos para dt = 0.001")
md(data,md4)
println("  ")


type ADMetadata
  valueSize::Tuple{Int64, Int64}
  opNode::Bool
  gradSize::Tuple{Int64, Int64}
end
type ADNodesMetadata
  id::String
  name::String
  metadata::ADMetadata
  group::Int64

  ADNodesMetadata(node::AD, id::String) = new(id, node.name, ADMetadata(size(node.value), node.opNode, size(node.grad)), node.opNode?1:0)
  ADNodesMetadata(node::AD, id::String, grp::Int64) = new(id, node.name, ADMetadata(size(node.value), node.opNode, size(node.grad)), grp)
end
type ADLinksMetadata
  source::String
  target::String
  value::Int64
end
type D3Data
  nodes::Array{ADNodesMetadata}
  links::Array{ADLinksMetadata}
end

# traverse the graph top-down
# convert AD to ADMetadata
function getGraphMetadata(graph::AD)
  uniqueNodes = Dict()
  nodes = Array(ADNodesMetadata,0)
  links = Array(ADLinksMetadata,0)

  bfs = [graph]
  while length(bfs) != 0
    current = pop!(bfs)
    currentId = string(pointer(current.value))
    if !haskey(uniqueNodes, currentId)
      push!(nodes, ADNodesMetadata(current, currentId))
      uniqueNodes[currentId] = true
    end
    currentPtr = string(pointer(current.value))
    numParents = length(current.parents)
    for i=1:numParents
      push!(links, ADLinksMetadata(currentPtr, string(pointer(current.parents[i].value)), current.level))
      push!(bfs, current.parents[i])
    end
  end

  return D3Data(nodes, links)
end

function d3ForcedDirectedGraphHTML(data::D3Data, supportsHTML::Bool)
  d3HTML =
    """
      <!DOCTYPE html>
      <meta charset="utf-8">
      <style>

        .links line {
          stroke: #999;
          stroke-opacity: 0.6;
        }

        .nodes text {
          pointer-events: none;
          font: 10px sans-serif;
        }

        text{
		      visibility: visible;
	      }

        .nodes circle {
          stroke: #fff;
          stroke-width: 1.5px;
        }

      </style>
      <svg width="960" height="600"></svg>
      <script type=text/javascript>
    """

  if supportsHTML
    # mbostock comment: https://github.com/mpld3/mpld3/pull/34
    d3HTML = d3HTML *
      """
        require.config({
          paths: {
            d3: '//d3js.org/d3.v4.min',
            mathjax: '//edx-static.s3.amazonaws.com/mathjax-MathJax-727332c/MathJax.js?config=TeX-MML-AM_HTMLorMML-full'
          },
          shim: {
            'd3': {
              exports: 'd3',
              init: function() {
                window.d3 = d3;
              }
            },
          }
        });
        require(["d3", "mathjax"], function(d3, mathjax) {
      """
  end

  d3HTML = d3HTML *
    """
      var svg = d3.select("svg"),
      width = +svg.attr("width"),
      height = +svg.attr("height");

      var color = d3.scaleOrdinal(d3.schemeCategory20);
      var simulation = d3.forceSimulation()
                          .force("link", d3.forceLink().id(function(d) { return d.id; }))
                          .force("charge", d3.forceManyBody())
                          .force("center", d3.forceCenter(width / 2, height / 2));

      var graph = $(json(data));
      var link = svg.append("g")
                    .attr("class", "links")
                    .selectAll("line")
                    .data(graph.links)
                    .enter().append("line")
                    .attr("stroke-width", function(d) { return Math.sqrt(d.value); });
      var node = svg.append("g")
                    .attr("class", "nodes")
                    .selectAll("circle")
                    .data(graph.nodes)
                    .enter()
                    .append("circle")
                    .attr("r", 5)
                    .attr("fill", function(d) { return color(d.group); })
                    .call(d3.drag()
                    .on("start", dragstarted)
                    .on("drag", dragged)
                    .on("end", dragended));

      node.append("title")
        .text(function(d) { return JSON.stringify(d.metadata); });
      node.append("text")
        .attr("dx", 12)
        .attr("dy", ".35em")
        .text(function(d) { return d.name });

      simulation
        .nodes(graph.nodes)
        .on("tick", ticked);

      simulation.force("link")
        .links(graph.links);

      setTimeout(textDisplay, 500);

  	  function textDisplay(){
  		    svg.append("g")
  		      .selectAll("text")
  		      .data(d3.selectAll("circle")._groups[0])
            .enter()
  		      .append("text")
  		      .attr("x", function(d,i){ return d.cx.animVal.value+5;})
  		      .attr("y", function(d,i){ return d.cy.animVal.value-5;})
  		      .text(function(d) { return d.__data__.name })
  		      .attr("style","")
  		      .attr("font-family", "sans-serif")
  		      .attr("font-size", "15px");
  	  }

      function ticked() {
        link
          .attr("x1", function(d) { return d.source.x; })
          .attr("y1", function(d) { return d.source.y; })
          .attr("x2", function(d) { return d.target.x; })
          .attr("y2", function(d) { return d.target.y; });
        node
          .attr("cx", function(d) { return d.x; })
          .attr("cy", function(d) { return d.y; });
      }

      function dragstarted(d) {
        if (!d3.event.active) simulation.alphaTarget(0.3).restart();
        d.fx = d.x;
        d.fy = d.y;
        svg.selectAll("text").remove()
	      textDisplay();
      }

      function dragged(d) {
        d.fx = d3.event.x;
        d.fy = d3.event.y;
        svg.selectAll("text").remove()
	      textDisplay();
      }

      function dragended(d) {
        if (!d3.event.active) simulation.alphaTarget(0);
        d.fx = null;
        d.fy = null;
        svg.selectAll("text").remove()
	      textDisplay();
      }
    """

  if supportsHTML
    d3HTML = d3HTML *
      """
        });
      """
  end

  d3HTML = d3HTML *
    """
      </script>
    """

  return d3HTML
end

# D3 operations on graph metadata
function getD3GraphMetadataHTML(graph::AD, supportsHTML::Bool)
  d3Data = getGraphMetadata(graph)
  d3HTML = d3ForcedDirectedGraphHTML(d3Data, supportsHTML)
  return d3HTML
end

# display in a new window (Electron) in case
# text/html mime type is not supported
function blink_show(graph::AD)
    w = Window()
    try
      d3 = joinpath(dirname(@__FILE__), "..", "deps", "d3.min.js")
      for file in (d3,)
        load!(w, file)
      end
    catch
      error("Could not load the file: d3.min.js from the deps directory.")
    end
    body = getD3GraphMetadataHTML(graph, false)
    body!(w, body)
end

# plot the graph
function plot(graph::AD)
  if displayable("text/html")
    display("text/html", getD3GraphMetadataHTML(graph, true))
  else
    #try
      blink_show(graph)
    #catch
    #  error("MIME type text/html is not available.")
    #end
  end
end


#= Testing data access patterns

The following methods will be evaluated:
 	_variable_
	varsubset
 	datasubset
	getx/getx!
	getindex
	setindex
=#

function t_data_access_patterns()

N = 10	# variables
M = 25  # samples

vectordata = rand(M)
matrixdata = rand(N,M)
labels = round.(2*rand(M))
multilabels = round.(2*rand(2,M))

A01 = j4pr.datacell(vectordata)
A02 = j4pr.datacell(vectordata, labels)
A03 = j4pr.datacell(vectordata, multilabels)
A04 = j4pr.datacell(matrixdata)
A05 = j4pr.datacell(matrixdata, labels)
A06 = j4pr.datacell(matrixdata, multilabels)
AV = [A01, A02, A03, A04, A05, A06]
AVv = [A01, A02, A03]
AVm = [A04, A05, A06]

# Data selection
# --------------

# Select observations
#println("Checking data selection...")
#print("\t|> Observations: datasubset - getindex - getx/getx!...")
for idx in [1,1:1,1:2,[1,2]]
	# Test for vector data
	for A in AVv
		Test.@test all(+j4pr.datasubset(A,idx) .== +A[idx] .== j4pr.getx(A)[idx] .== j4pr.getx!(A)[idx])
	end

	# Test for matrix data
	for A in AVm
		Test.@test all(+j4pr.datasubset(A,idx) .== +A[:,idx] .== j4pr.getx(A)[:,idx] .== j4pr.getx!(A)[:,idx])
	end
end
#println("PASSED")

# Select variables
#print("\t|> Variables: varsubset - _variable_ - getindex - getx/getx!...")
for idx in [1,1:1,1:2,[1,2]]
	# Test for vector data
	for A in AVv
		if idx isa Int 	# idx has to be 1 otherwise most of the variable access methods fail, as they should;
				# Note that there is no getindex method for A
			Test.@test all(+j4pr.varsubset(A,idx) .== j4pr._variable_(A,idx) .== +A .== j4pr.getx(A)[:,idx] .== j4pr.getx!(A)[:,idx])
		end
	end

	# Test for matrix data
	for A in AVm
		if idx isa Int
			# data cell indexing returns a matrix data if single row selected, result must be transposed as all the other methods use vectors
			Test.@test all(+j4pr.varsubset(A,idx) .== j4pr._variable_(A,idx) .== (+A[idx,:])' .== j4pr.getx(A)[idx,:] .== j4pr.getx!(A)[idx,:])
		else
			Test.@test all(+j4pr.varsubset(A,idx) .== j4pr._variable_(A,idx) .== +A[idx,:] .== j4pr.getx(A)[idx,:] .== j4pr.getx!(A)[idx,:])
		end
	end
end
#println("PASSED")

# Selecting observations + variables is done by applying successively one of the methods for variable selection and observation selection

# Data modification
# -----------------

# Modify observations
#println("Checking data modification...")
#print("\t|> Observations: datasubset - settindex - getx! and NOT getx...")
for idx in [1,1:1,1:2,[1,2]]
	# Test for vector data
	for A in AVv
		# datasubset approach
		Ac = deepcopy(A)
		newvals = idx isa Int ? rand() : rand(length(idx))
		j4pr.datasubset(Ac,idx)[:] = newvals;
		Test.@test all(+j4pr.datasubset(Ac,idx) .== newvals)

		# setindex! approach
		Ac = deepcopy(A)
		newvals = idx isa Int ? rand() : rand(length(idx))
		Ac[idx] = newvals;
		Test.@test all(+Ac[idx] .== newvals)

		# getx! approach
		Ac = deepcopy(A)
		newvals = idx isa Int ? rand() : rand(length(idx))
		j4pr.getx!(Ac)[idx] = newvals;
		Test.@test all(j4pr.getx!(Ac)[idx] .== newvals)

		# getx approach (MUST NOT MODIDY)
		Ac = deepcopy(A)
		newvals = idx isa Int ? rand() : rand(length(idx))
		j4pr.getx(Ac)[idx] = newvals;
		Test.@test all(j4pr.getx!(Ac)[idx] .!= newvals)
		Test.@test all(j4pr.getx!(Ac)[idx] .== j4pr.getx!(A)[idx])
	end

	# Test for matrix data
	for A in AVm
		# datasubset approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nvars(Ac), length(idx))
		j4pr.datasubset(Ac,idx)[:] = newvals;
		Test.@test all(+j4pr.datasubset(Ac,idx) .== newvals)

		# setindex! approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nvars(Ac), length(idx))
		Ac[:,idx] = newvals;
		Test.@test all(+Ac[:,idx] .== newvals)

		# getx! approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nvars(Ac), length(idx))
		j4pr.getx!(Ac)[:,idx] = newvals;
		Test.@test all(j4pr.getx!(Ac)[:,idx] .== newvals)

		# getx approach (MUST NOT MODIDY)
		Ac = deepcopy(A)
		newvals = rand(j4pr.nvars(Ac), length(idx))
		j4pr.getx(Ac)[:,idx] = newvals;
		Test.@test all(j4pr.getx!(Ac)[:,idx] .!= newvals)
		Test.@test all(j4pr.getx!(Ac)[:,idx] .== j4pr.getx!(A)[:,idx])
	end
end
#println("PASSED")

# Modify variables
#print("\t|> Variables: varsubset - _variable_ - settindex - getx! and NOT getx...")
for idx in [1,1:1,1:2,[1,2]]
	# Test for vector data
	for A in AVv
		if idx isa Int
		# varsubset approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nobs(A))
		j4pr.varsubset(Ac,idx)[:] = newvals;
		Test.@test all(+j4pr.varsubset(Ac,idx) .== newvals)

		# _variable_ approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nobs(A))
		j4pr._variable_(Ac,idx)[:] = newvals;
		Test.@test all(+j4pr._variable_(Ac,idx) .== newvals)

		# setindex! approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nobs(A))
		Ac[:] = newvals;
		Test.@test all(+Ac .== newvals)

		# getx! approach
		Ac = deepcopy(A)
		newvals = rand(j4pr.nobs(A))
		j4pr.getx!(Ac)[:,idx] = newvals;
		Test.@test all(j4pr.getx!(Ac)[:,idx] .== newvals)

		# getx approach (MUST NOT MODIDY)
		Ac = deepcopy(A)
		newvals = rand(j4pr.nobs(A))
		j4pr.getx(Ac)[:,idx] = newvals;
		Test.@test all(j4pr.getx!(Ac)[:,idx] .!= newvals)
		Test.@test all(j4pr.getx!(Ac)[:,idx] .== j4pr.getx!(A)[:,idx])
		end
	end
	# Test for matrix data
	for A in AVm
		# varsubset approach
		Ac = deepcopy(A)
		newvals = rand(size(j4pr.varsubset(Ac,idx)[:]))
		j4pr.varsubset(Ac,idx)[:] = newvals;
		Test.@test all(+j4pr.varsubset(Ac,idx) .== newvals)

		# _variable_ approach
		Ac = deepcopy(A)
		newvals = rand(size(j4pr._variable_(Ac,idx)))
		j4pr._variable_(Ac,idx)[:] = newvals;
		Test.@test all(+j4pr._variable_(Ac,idx) .== newvals)

		# setindex! approach
		Ac = deepcopy(A)
		newvals = rand(size(Ac[idx,:]))
		Ac[idx,:] = newvals;
		Test.@test all(+Ac[idx,:] .== newvals)

		# getx! approach
		Ac = deepcopy(A)
		newvals = rand(size(j4pr.getx!(Ac)[idx,:]))
		j4pr.getx!(Ac)[idx,:] = newvals;
		Test.@test all(j4pr.getx!(Ac)[idx,:] .== newvals)

		# getx approach (MUST NOT MODIDY)
		Ac = deepcopy(A)
		newvals = rand(size(j4pr.getx!(Ac)[idx,:]))
		j4pr.getx(Ac)[idx,:] = newvals;
		Test.@test all(j4pr.getx!(Ac)[idx,:] .!= newvals)
		Test.@test all(j4pr.getx!(Ac)[idx,:] .== j4pr.getx!(A)[idx,:])
	end
end
#println("PASSED")

end


module BlockNLPModels

using NLPModels
using LinearAlgebra
using SparseArrays

export AbstractBlockNLPModel
export BlockNLPModel
export DualizedNLPBlockModel
export AugmentedNLPBlockModel
export FullSpaceModel
export ProxAugmentedNLPBlockModel
export add_block,
    add_links,
    n_constraints,
    get_linking_matrix,
    get_linking_matrix_blocks,
    get_rhs_vector,
    update_dual!,
    update_primal!

"""
    AbstractBlockNLPModel
Abstract supertype for the definition of NLP problems with a block structure.
"""
abstract type AbstractBlockNLPModel end

"""
    AbstractLinearLinkConstraint
Abstract supertype for the definition of linear linking constraints
"""
abstract type AbstractLinearLinkConstraint end
"""
    AbstractNonLinearLinkConstraint
Abstract supertype for the definition of nonlinear linking constraints
"""
abstract type AbstractNonLinearLinkConstraint end

"""
    BlockNLPCounters()
Keeps a count of blocks, linking constraints, variables, and (block) constraints attached to the model.
"""
mutable struct BlockNLPCounters
    block_counter::Int
    link_counter::Int
    var_counter::Int
    con_counter::Int
    function BlockNLPCounters()
        return new(0, 0, 0, 0)
    end
end

"""
    AbstractBlockModel{T, S} <: AbstractNLPModel{T, S}
A data type to store block subproblems of the form
```math
\\begin{aligned}
  \\min_{x \\in \\mathbb{R^{m_i}}} \\quad & f_i(x_i) \\\\
  \\mathrm{subject \\, to} \\quad & c_{ij} (x_i) \\leq 0 \\quad \\forall j \\in \\mathcal{C}_i, \\\\
\\end{aligned}
where ``i`` is the index of the block.
"""
mutable struct AbstractBlockModel{T,S} <: AbstractNLPModel{T,S}
    meta::NLPModelMeta{T,S}
    counters::Counters
    problem_block::AbstractNLPModel
    idx::Int
    var_idx::UnitRange{Int} # Indices of variables wrt the full space model
    con_idx::UnitRange{Int} # Indices of constraints wrt the full space model
    linking_constraint_id::Vector{Int} # ID of linking constraints that connect this block with other blocks
end

"""
    LinearLinkConstraint <: AbstractLinearLinkConstraint
A data type to store linear link constraints of the form
```math
\\begin{aligned}
  \\sum\\limits_{i \\in \\mathcal{B}} A_i x_i = b
\\end{aligned}
"""
mutable struct LinearLinkConstraint <: AbstractLinearLinkConstraint
    linking_blocks::Vector{SparseMatrixCSC{Float64,Int}}
    idx::UnitRange{Int}
    link_map::Dict{Int,Vector{Int}} # constraint => blocks, i.e., which constraint connects which blocks
    rhs_vector::Vector{Float64}
end

"""
    BlockNLPModel <: AbstractBlockNLPModel
A data type designed to store nonlinear optimization models of the form
```math
\\begin{aligned}
  \\min_{x \\in \\mathbb{R^{m_i}}} \\quad & \\sum\\limits_{i \\in \\mathcal{B}} f_i(x_i) \\\\
  \\mathrm{subject \\, to} \\quad & c_{ij} (x_i) \\leq 0 \\quad \\forall i \\in \\mathcal{B}, j \\in \\mathcal{C}_i \\\\
  & \\sum\\limits_{i \\in \\mathcal{B}} A_i x_i = b,
\\end{aligned}
where ``\\mathcal{B}`` is the set of variable blocks.

---

    BlockNLPModel()

Initializes an empty `BlockNLPModel`.
"""
mutable struct BlockNLPModel <: AbstractBlockNLPModel
    problem_size::BlockNLPCounters
    blocks::Vector{AbstractBlockModel}
    linking_constraints::Vector{
        Union{AbstractLinearLinkConstraint,AbstractNonLinearLinkConstraint},
    }
    function BlockNLPModel()
        return new(
            BlockNLPCounters(),
            Vector{AbstractBlockModel}(),
            Vector{Union{AbstractLinearLinkConstraint,AbstractNonLinearLinkConstraint}}(),
        )
    end
end

"""
    add_block(
     block_nlp::AbstractBlockNLPModel,
     nlp::AbstractNLPModel
    )

Adds a block subproblem to a `BlockNLPModel`.
The `BlockNLPModel` needs to be initialized before this function can be called.

# Arguments

- `block_nlp::AbstractBlockNLPModel`: name of the BlockNLPModel to which a subproblem is to be added.
- `nlp::AbstractNLPModel`: the subproblem.
"""
function add_block(block_nlp::AbstractBlockNLPModel, nlp::AbstractNLPModel)
    block_nlp.problem_size.block_counter += 1

    temp_var_counter = block_nlp.problem_size.var_counter + 1
    block_nlp.problem_size.var_counter += nlp.meta.nvar
    var_idx = temp_var_counter:block_nlp.problem_size.var_counter

    temp_con_counter = block_nlp.problem_size.con_counter + 1
    block_nlp.problem_size.con_counter += nlp.meta.ncon
    con_idx = temp_con_counter:block_nlp.problem_size.con_counter

    push!(
        block_nlp.blocks,
        AbstractBlockModel(
            nlp.meta,
            Counters(),
            nlp,
            block_nlp.problem_size.block_counter,
            var_idx,
            con_idx,
            Vector{Int}(),
        ),
    )
end

"""
    add_links(
      block_nlp::AbstractBlockNLPModel,
      n_constraints::Int,
      links::Union{Dict{Int, Array{Float64, 2}}, Dict{Int, Vector{Float64}}},
      constants::Union{AbstractVector, Float64}
    )

Creates linear link constraint(s) between different subproblems of the BlockNLPModel.
This function must be called after subproblems have already been added.

# Arguments

- `block_nlp::AbstractBlockNLPModel`: The BlockNLPModel to which the constraint(s) is to be added.
- `n_constraints::Int`: Number of link constraints.
- `links::Union{Dict{Int, Array{Float64, 2}}, Dict{Int, Vector{Float64}}}`: coefficients of the block matrices that link different blocks,
- `constants::Union{AbstractVector, Float64}`: RHS vector
"""
function add_links(
    block_nlp::AbstractBlockNLPModel,
    n_constraints::Int,
    links::Union{Dict{Int,Array{Float64,2}},Dict{Int,Vector{Float64}}},
    constants::Union{AbstractVector,Float64},
)

    # Initialize linking_blocks with empty matrices
    linking_blocks = [
        spzeros(n_constraints, block_nlp.blocks[i].meta.nvar) for
        i = 1:block_nlp.problem_size.block_counter
    ]

    # check if everything is of appropriate size and add to linking_blocks
    @assert length(constants) == n_constraints
    if n_constraints == 1
        rhs_vector = [constants]
        for block_idx in keys(links)
            @assert length(links[block_idx]) == block_nlp.blocks[block_idx].meta.nvar
            linking_blocks[block_idx][1, :] = sparse(links[block_idx])
        end
    else
        rhs_vector = constants
        for block_idx in keys(links)
            @assert size(links[block_idx])[1] == n_constraints
            @assert size(links[block_idx])[2] == block_nlp.blocks[block_idx].meta.nvar
            linking_blocks[block_idx] = sparse(links[block_idx])
        end
    end

    # Prepare other information
    link_map = Dict{Int,Vector{Int}}()
    link_con_idx = Vector{Int}()

    # Parse the constraints one-by-one to collect this information
    temp_link_counter = block_nlp.problem_size.link_counter + 1
    for con = 1:n_constraints
        block_nlp.problem_size.link_counter += 1
        link_map[block_nlp.problem_size.link_counter] = Vector{Int}()
        for block_idx in keys(links)
            if linking_blocks[block_idx][con, :] !=
               spzeros(block_nlp.blocks[block_idx].meta.nvar)
                push!(
                    block_nlp.blocks[block_idx].linking_constraint_id,
                    block_nlp.problem_size.link_counter,
                )
                push!(link_map[block_nlp.problem_size.link_counter], block_idx)
            end
        end
    end
    link_con_idx = temp_link_counter:block_nlp.problem_size.link_counter
    lin_link_con = LinearLinkConstraint(linking_blocks, link_con_idx, link_map, rhs_vector)
    push!(block_nlp.linking_constraints, lin_link_con)
end

include("utils.jl")
include("full_space.jl")
include("dualized_block.jl")
include("augmented_block.jl")
include("prox_augmented_block.jl")
end # module


#
# This file is a part of MolecularGraph.jl
# Licensed under the MIT License http://opensource.org/licenses/MIT
#

export
    sssr, atom_sssrmem, bond_sssrmem,
    atom_sssrsizes, bond_sssrsizes,
    isringatom, isringbond,
    atom_sssrcount, bond_sssrcount,
    scaffoldmem, componentmem,
    atomsymbol, charge, multiplicity, bondorder,
    nodedegree, valence, lonepair, heavyatomcount,
    explicithcount, implicithcount, hcount, connectivity,
    ishdonor, hacceptorcount, ishacceptor, hdonorcount,
    isrotatable, rotatablecount,
    countatoms, molecularformula, empiricalformula,
    pielectron, hybridization,
    isaromaticring, isaromatic, isaromaticbond

# Molecular graph topology

sssr = mincyclenodes
atom_sssrmem = node_cyclemem
bond_sssrmem = edge_cyclemem
scaffoldmem = two_edge_membership
componentmem = connected_membership

function atom_sssrsizes(mol::UndirectedGraph)
    vec = [Set{Int}() for i in 1:nodecount(mol)]
    rings = sssr(mol)
    for (i, cs) in enumerate(atom_sssrmem(mol))
        for c in cs
            push!(vec[i], length(rings[c]))
        end
    end
    return vec
end

atom_sssrsizes(view::SubgraphView) = atom_sssrsizes(view.graph)

function bond_sssrsizes(mol::UndirectedGraph)
    vec = [Set{Int}() for i in 1:edgecount(mol)]
    rings = sssr(mol)
    for (i, cs) in enumerate(bond_sssrmem(mol))
        for c in cs
            push!(vec[i], length(rings[c]))
        end
    end
    return vec
end

bond_sssrsizes(view::SubgraphView) = bond_sssrsizes(view.graph)

# TODO: waiting for fix #28992
# @cache atom_isringmem(mol) = .!isempty.(node_cyclemem(mol))
isringatom(mol::UndirectedGraph) = [!i for i in isempty.(atom_sssrmem(mol))]
isringatom(view::SubgraphView) = isringatom(view.graph)

# @cache bond_isringmem(mol) = .!isempty.(node_cyclemem(mol))
isringbond(mol::UndirectedGraph) = [!i for i in isempty.(bond_sssrmem(mol))]
isringbond(view::SubgraphView) = isringbond(view.graph)

atom_sssrcount(mol::UndirectedGraph) = length.(atom_sssrmem(mol))
atom_sssrcount(view::SubgraphView) = atom_sssrcount(view.graph)

bond_sssrcount(mol::UndirectedGraph) = length.(bond_sssrmem(mol))
bond_sssrcount(view::SubgraphView) = bond_sssrcount(view.graph)

# Elemental properties

"""
    atomsymbol(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of atom symbols within the molecule that have ``n``
atoms.
"""
@cache atomsymbol(mol::GraphMol) = getproperty.(nodeattrs(mol), :symbol)
atomsymbol(view::SubgraphView) = atomsymbol(view.graph)

"""
    charge(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of atom charges within the molecule that have ``n``
atoms.
"""
@cache charge(mol::GraphMol) = getproperty.(nodeattrs(mol), :charge)
charge(view::SubgraphView) = charge(view.graph)

"""
    multiplicity(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of atom multiplicities within the molecule that
have ``n`` atoms (1: non-radical, 2: radical, 3: biradical).
"""
@cache multiplicity(mol::GraphMol) = getproperty.(nodeattrs(mol), :multiplicity)
multiplicity(view::SubgraphView) = multiplicity(view.graph)

"""
    bondorder(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of bond orders within the molecule that have ``n``
bonds.
"""
@cache bondorder(mol::GraphMol) = getproperty.(edgeattrs(mol), :order)
bondorder(view::SubgraphView) = bondorder(view.graph)

"""
    nodedegree(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of node degrees within the molecular graph that
have ``n`` atom nodes.

`nodedegree` values correspond to SMARTS `D` property.
"""
@cache nodedegree(mol::GraphMol) = [degree(mol, n) for n in 1:nodecount(mol)]
nodedegree(view::SubgraphView) = nodedegree(view.graph)

# Valence

"""
    lonepair(mol::GraphMol) -> Vector{Union{Int,Nothing}}

Return the size ``n`` vector of the number of lone pairs within the molecule that
have ``n`` atoms.

Note that implicit hydrogens are available for only organic atoms. The lonepair
value of inorganic atoms would be `nothing`. The result can take negative value if the atom has empty valence shells.
"""
function lonepair(mol::GraphMol)
    defs = Dict(
        :H => 0, :B => -1, :C => 0, :N => 1, :O => 2, :F => 3,
        :Si => 0, :P => 1, :S => 2, :Cl => 3,
        :As => 1, :Se => 2, :Br => 3, :I => 3
    )
    vec = Union{Int,Nothing}[]
    for (sym, chg) in zip(atomsymbol(mol), charge(mol))
        num = get(defs, sym, nothing)
        v = num === nothing ? nothing : num - chg
        push!(vec, v)
    end
    return vec
end

lonepair(view::SubgraphView) = lonepair(view.graph)

function apparentvalence(mol::GraphMol)
    vec = zeros(Int, nodecount(mol))
    bondorder_ = bondorder(mol)
    for i in 1:nodecount(mol)
        for inc in incidences(mol, i)
            vec[i] += bondorder_[inc]
        end
    end
    return vec
end

apparentvalence(view::SubgraphView) = apparentvalence(view.graph)

"""
    valence(mol::GraphMol) -> Vector{Union{Int,Nothing}}

Return the size ``n`` vector of intrinsic valences (with considering implicit
hydrogens) within the molecule that have ``n`` atoms.

`valence` values correspond to SMARTS `v` property. `valence` values of inorganic atoms would be `nothing`.
"""
@cache function valence(mol::GraphMol)
    vec = Union{Int,Nothing}[]
    atomsymbol_ = atomsymbol(mol)
    lonepair_ = lonepair(mol)
    for i in 1:nodecount(mol)
        if lonepair_[i] === nothing
            push!(vec, nothing)
        elseif atomsymbol_[i] === :H
            push!(vec, 1)
        else
            push!(vec, 4 - abs(lonepair_[i]))
        end
    end
    return vec
end

valence(view::SubgraphView) = valence(view.graph)

"""
    explicithcount(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of the number of adjacent explicit hydrogen nodes within the molecule that have ``n`` atoms.
"""
@cache function explicithcount(mol::GraphMol)
    vec = zeros(Int, nodecount(mol))
    atomsymbol_ = atomsymbol(mol)
    for i in 1:nodecount(mol)
        for adj in adjacencies(mol, i)
            if atomsymbol_[adj] == :H
                vec[i] += 1
            end
        end
    end
    return vec
end

explicithcount(view::SubgraphView) = explicithcount(view.graph)

"""
    implicithcount(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of the number of implicit hydrogens within
the molecule that have ``n`` atoms.
"""
@cache function implicithcount(mol::GraphMol)
    hcnt = (v, av) -> v === nothing ? 0 : max(0, v - av)
    return hcnt.(valence(mol), apparentvalence(mol))
end

implicithcount(view::SubgraphView) = implicithcount(view.graph)

"""
    heavyatomcount(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of the number of adjacent non-hydrogen atoms within
the molecule that have ``n`` atoms.
"""
heavyatomcount(mol::GraphMol) = nodedegree(mol) - explicithcount(mol)
heavyatomcount(view::SubgraphView) = heavyatomcount(view.graph)

"""
    hcount(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of the total number of hydrogens attached to the atom within the molecule that have ``n`` atoms.

`hcount` values correspond to SMARTS `H` property.
"""
@cache hcount(mol::GraphMol) = explicithcount(mol) + implicithcount(mol)
hcount(view::SubgraphView) = hcount(view.graph)

"""
    connectivity(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of the total number of adjacent atoms (including implicit hydrogens) within the molecule that have ``n`` atoms.

`connectivity` values correspond to SMARTS `X` property.
"""
@cache connectivity(mol::GraphMol) = nodedegree(mol) + implicithcount(mol)
connectivity(view::SubgraphView) = connectivity(view.graph)

# Hydrogen bond donor/acceptor

function ishacceptor(mol::GraphMol)
    ac = (sym, lp) -> lp === nothing ? false : sym in (:N, :O, :F) && lp > 0
    return ac.(atomsymbol(mol), lonepair(mol))
end

ishacceptor(view::SubgraphView) = ishacceptor(view.graph)

"""
    hacceptorcount(mol::GraphMol) -> Int

Return the number of hydrogen bond acceptors (N, O and F).
"""
hacceptorcount(mol::GraphMol) = reduce(+, ishacceptor(mol); init=0)

function ishdonor(mol::GraphMol)
    dc = (sym, h) -> sym in (:N, :O) && h > 0
    return dc.(atomsymbol(mol), hcount(mol))
end

ishdonor(view::SubgraphView) = ishdonor(view.graph)

"""
    hdonorcount(mol::GraphMol) -> Int

Return the number of hydrogen bond donors (O and N attached to hydrogens).
"""
hdonorcount(mol::GraphMol) = reduce(+, ishdonor(mol); init=0)

# Rotatable bonds

"""
    isrotatable(mol::GraphMol)

Return whether the bonds are rotatable or not.
"""
@cache function isrotatable(mol::GraphMol)
    nodedegree_ = nodedegree(mol)
    isringbond_ = isringbond(mol)
    bondorder_ = bondorder(mol)
    vec = Bool[]
    for (i, (u, v)) in enumerate(edgesiter(mol))
        rot = (!isringbond_[i] && bondorder_[i] == 1
            && nodedegree_[u] != 1 && nodedegree_[v] != 1)
        push!(vec, rot)
    end
    return vec
end

isrotatable(view::SubgraphView) = isrotatable(view.graph)

"""
    rotatablecount(mol::GraphMol) -> Int

Return the number of rotatable bonds.
"""
rotatablecount(mol::GraphMol) = reduce(+, isrotatable(mol); init=0)

# Composition

"""
    countatoms(mol::GraphMol) -> Dict{Symbol,Int}

Count the number of atoms and return symbol => count dict.
"""
function countatoms(mol::GraphMol)
    counter = Dict{Symbol,Int}()
    for sym in atomsymbol(mol)
        if !haskey(counter, sym)
            counter[sym] = 1
        else
            counter[sym] += 1
        end
    end
    hcnt = reduce(+, implicithcount(mol); init=0)
    if hcnt > 0
        if !haskey(counter, :H)
            counter[:H] = hcnt
        else
            counter[:H] += hcnt
        end
    end
    return counter
end

function writeformula(counter::Dict{Symbol,Int})
    strs = []
    if haskey(counter, :C)
        push!(strs, "C")
        c = pop!(counter, :C)
        c > 1 && push!(strs, string(c))
        if haskey(counter, :H)
            push!(strs, "H")
            h = pop!(counter, :H)
            h > 1 && push!(strs, string(h))
        end
    end
    for sym in sort(collect(keys(counter)))
        push!(strs, string(sym))
        cnt = counter[sym]
        cnt > 1 && push!(strs, string(cnt))
    end
    return join(strs)
end

"""
    molecularformula(mol::GraphMol) -> String

Return the molecular formula in Hill system.
"""
function molecularformula(mol::GraphMol)
    counter = countatoms(mol)
    return writeformula(counter)
end

"""
    empiricalformula(mol::GraphMol) -> String

Return the empirical formula in Hill system.
"""
function empiricalformula(mol::GraphMol)
    counter = countatoms(mol)
    if length(counter) > 1
        dv = reduce(gcd, values(counter))
        for k in keys(counter)
            counter[k] = div(counter[k], dv)
        end
    end
    return writeformula(counter)
end

# Hybridization

"""
    pielectron(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of the number of ``\\pi`` electrons (including implicit hydrogens) within the molecule that have ``n`` atoms.

The ``\\pi`` electron enumelation is based on the following simple rules.

- Any atom incident to a double bond -> +1
- Any atom incident to two double bond -> +2
- Any atom incident to a triple bond -> +1
- Uncharged N and O not incident to any multiple bonds and adjacent to atoms described above -> +2

These rules are applied for only typical organic atoms. The values for inorganic atoms will be 0.
"""
@cache function pielectron(mol::GraphMol)
    atomsymbol_ = atomsymbol(mol)
    charge_ = charge(mol)
    pie_ = apparentvalence(mol) - nodedegree(mol)
    vec = zeros(Int, nodecount(mol))
    for i in 1:nodecount(mol)
        vec[i] = pie_[i]
        atomsymbol_[i] in (:N, :O) || continue
        for adj in adjacencies(mol, i)
            if pie_[i] == 0 && pie_[adj] > 0 && charge_[i] == 0
                vec[i] = 2
                break
            end
        end
    end
    return vec
end

pielectron(view::SubgraphView) = pielectron(view.graph)

"""
    hybridization(mol::GraphMol) -> Vector{Int}

Return the size ``n`` vector of orbital hybridization symbols (`:sp3`, `:sp2`, `:sp`, `:none`) within the molecule that have ``n`` atoms.

These hybridizations are for only typical organic atoms. Inorganic atoms and other orbitals like s, sp3d and sp3d2 will be `:none`.
"""
@cache function hybridization(mol::GraphMol)
    atomsymbol_ = atomsymbol(mol)
    connectivity_ = connectivity(mol)
    lonepair_ = lonepair(mol)
    pielectron_ = pielectron(mol)
    organic_symbols = [:B, :C, :N, :O, :F, :Si, :P, :S, :Cl, :As, :Se, :Br, :I]
    vec = Symbol[]
    for i in 1:nodecount(mol)
        if atomsymbol_[i] in organic_symbols
            val = connectivity_[i] + lonepair_[i]
            if val == 4
                if (atomsymbol_[i] in [:N, :O] && pielectron_[i] == 2)
                    push!(vec, :sp2)  # adjacent to conjugated bonds
                else
                    push!(vec, :sp3)
                end
            elseif val == 3
                push!(vec, :sp2)
            elseif val == 2
                push!(vec, :sp)
            else
                push!(vec, :none)
            end
        else
            push!(vec, :none)
        end
    end
    return vec
end

hybridization(view::SubgraphView) = hybridization(view.graph)

# Aromatic rings

function satisfyhuckel(mol::GraphMol, ring)
    # TODO: need improvement
    cnt = 0
    atomsymbol_ = atomsymbol(mol)
    nodedegree_ = nodedegree(mol)
    pie_ = apparentvalence(mol) - nodedegree(mol)
    lonepair_ = lonepair(mol)
    carbonylO = findall(
        (atomsymbol_ .== :O) .* (nodedegree_ .== 1) .* (pie_ .== 1))
    carbonylC = Int[]
    for o in carbonylO
        c = iterate(adjacencies(mol, o))[1]
        if atomsymbol_[c] == :C
            push!(carbonylC, c)
        end
    end
    for r in ring
        if r in carbonylC
            continue
        elseif pie_[r] == 1
            cnt += 1
        elseif lonepair_[r] === nothing
            return false
        elseif lonepair_[r] > 0
            cnt += 2
        elseif lonepair_[r] < 0
            continue
        else
            return false
        end
    end
    return cnt % 4 == 2
end

@cache function isaromaticring(mol::GraphMol)
    vec = falses(circuitrank(mol))
    for (i, ring) in enumerate(sssr(mol))
        if satisfyhuckel(mol, ring)
            vec[i] = true
        elseif nodeattrtype(mol) === SmilesAtom # SMILES aromatic atom
            sub = nodesubgraph(mol, Set(ring))
            if all(nodeattr(mol, n).isaromatic for n in nodeset(sub))
                vec[i] = true
            end
        end
    end
    return vec
end

"""
    isaromatic(mol::GraphMol)

Return the vector whether the atom belongs to an aromatic ring.

Note that aromaticity described here means simplified binary descriptor
(aromatic or not) based on classical Huckel's rule. This is intended for use in
some kind of pharmaceutical researches. Non-classical aromaticity such as
Moebius aromaticity is not considered.
"""
@cache function isaromatic(mol::GraphMol)
    aromatic = falses(nodecount(mol))
    for ring in sssr(mol)[isaromaticring(mol)]
        sub = nodesubgraph(mol, Set(ring))
        for n in nodeset(sub)
            aromatic[n] = true
        end
    end
    return aromatic
end

isaromatic(view::SubgraphView) = isaromatic(view.graph)

@cache function isaromaticbond(mol::GraphMol)
    aromaticbond = falses(edgecount(mol))
    for ring in sssr(mol)[isaromaticring(mol)]
        sub = nodesubgraph(mol, Set(ring))
        for e in edgeset(sub)
            aromaticbond[e] = true
        end
    end
    return aromaticbond
end

isaromaticbond(view::SubgraphView) = isaromaticbond(view.graph)


using Knet: Knet, AutoGrad, Data, adam, adam!, progress!, minibatch, save, relu,
    gpu, KnetArray, load, sigm, goldensection, hyperband
using Statistics: mean
using DataFrames
using Plots
using Dates

PARAMS = Dict("batchsize" => 32,
              "pdrop" => 0.1,
              "lr" => 0.01,
              "hidden" => 100,
              "optimizer" => adam!
              )

mutable struct Model
    config::Config;
    name::String
    model
    savepath::String
    datapath::String
    vocabulary
    params
    extractor
    fdict
end

Model(config::Config; name="model", savedir=SAVEDIR, datadir="",
      voc=nothing, params=PARAMS, extractor=Dict(), fdict=Dict()) = Model(
          config, name, nothing, savedir, datadir, voc, params, extractor,
          fdict)

Model(inputs::Array{Tuple{String, String},1},
      outputs::Array{Tuple{String, String},1};
      name="model", savedir=SAVEDIR, datadir="", voc=nothing, params=PARAMS,
      extractor=Dict(), fdict=Dict()) = Model(Config(inputs, outputs), name=name,
                                savedir=SAVEDIR, datadir=datadir,
                                voc=voc, params=params,
                                extractor=extractor, fdict=fdict)

getfeatures(m::Model; ftype="all") = getfeatures(m.config; ftype=ftype)
getfnames(m::Model; ftype="all") = getfnames(m.config; ftype=ftype)
getftypes(m::Model; ftype="all") = getftypes(m.config; ftype=ftype)

function iscategorical(m::Model)
    in(CATEGORY, getftypes(m; ftype="output")) ||
        in(BINARYCATEGORY, getftypes(m; ftype="output"))
end

isimagemodel(m::Model) = in(IMAGE, getftypes(m; ftype="input"))
istextmodel(m::Model) = in(TEXT, getftypes(m; ftype="input"))

function savemodel(m::Model, savepath=nothing)
    savepath = savepath == nothing ? joinpath(SAVEDIR, "$(m.name).jld2") :
        savepath
    save(savepath,
         "config", m.config,
         "name", m.name,
         "model", m.model,
         "savepath", m.savepath,
         "params", m.params,
         "extractor", m.extractor,
         "fdict", m.fdict)
end

function loadmodel(m::Model, loadpath::String)
    m.config = load(loadpath, "config")
    m.name = load(loadpath, "name")
    m.model = load(loadpath, "model")
    m.savepath = load(loadpath, "savepath")
    m.params = load(loadpath, "params")
    m.extractor = load(loadpath, "extractor")
    m.fdict = load(loadpath, "fdict")

    return m
end

function loadmodel(loadpath::String)
    m = Model(Config())
    loadmodel(m, loadpath)
    m
end

function predictdata(m::Model, example)
    data = Dict(fname => [value] for (fname, value) in example)
    data = preprocess(m, data)
    x = preparedata(m, data; output=false)

    iscategorical(m) ? predict(m.model, x) : m.model(x)
end


#===================================
Define GNNGraph type as a subtype of LightGraphs' AbstractGraph.
For the core methods to be implemented by any AbstractGraph, see
https://juliagraphs.org/LightGraphs.jl/latest/types/#AbstractGraph-Type
https://juliagraphs.org/LightGraphs.jl/latest/developing/#Developing-Alternate-Graph-Types
=============================================#

const COO_T = Tuple{T, T, V} where {T <: AbstractVector, V}
const ADJLIST_T = AbstractVector{T} where T <: AbstractVector
const ADJMAT_T = AbstractMatrix
const SPARSE_T = AbstractSparseMatrix # subset of ADJMAT_T
const CUMAT_T = Union{AnyCuMatrix, CUDA.CUSPARSE.CuSparseMatrix}

"""
    GNNGraph(data; [graph_type, ndata, edata, gdata, num_nodes, graph_indicator, dir])
    GNNGraph(g::GNNGraph; [ndata, edata, gdata])

A type representing a graph structure and storing also
feature arrays associated to nodes, edges, and to the whole graph (global features).

A `GNNGraph` can be constructed out of different objects `data` expressing
the connections inside the graph. The internal representation type
is determined by `graph_type`.

When constructed from another `GNNGraph`, the internal graph representation
is preserved and shared. The node/edge/global features are transmitted
as well, unless explicitely changed though keyword arguments.

A `GNNGraph` can also represent multiple graphs batched togheter
(see [`Flux.batch`](@ref) or [`SparseArrays.blockdiag`](@ref)).
The field `g.graph_indicator` contains the graph membership
of each node.

A `GNNGraph` is a LightGraphs' `AbstractGraph`, therefore any functionality
from the LightGraphs' graph library can be used on it.

# Arguments

- `data`: Some data representing the graph topology. Possible type are
    - An adjacency matrix
    - An adjacency list.
    - A tuple containing the source and target vectors (COO representation)
    - A LightGraphs' graph.
- `graph_type`: A keyword argument that specifies
                the underlying representation used by the GNNGraph.
                Currently supported values are
    - `:coo`. Graph represented as a tuple `(source, target)`, such that the `k`-th edge
              connects the node `source[k]` to node `target[k]`.
              Optionally, also edge weights can be given: `(source, target, weights)`.
    - `:sparse`. A sparse adjacency matrix representation.
    - `:dense`. A dense adjacency matrix representation.
    Default `:coo`.
- `dir`. The assumed edge direction when given adjacency matrix or adjacency list input data `g`.
        Possible values are `:out` and `:in`. Default `:out`.
- `num_nodes`. The number of nodes. If not specified, inferred from `g`. Default `nothing`.
- `graph_indicator`. For batched graphs, a vector containing the graph assigment of each node. Default `nothing`.
- `ndata`: Node features. A named tuple of arrays whose last dimension has size num_nodes.
- `edata`: Edge features. A named tuple of arrays whose whose last dimension has size num_edges.
- `gdata`: Global features. A named tuple of arrays whose has size num_graphs.

# Usage.

```julia
using Flux, GraphNeuralNetworks

# Construct from adjacency list representation
data = [[2,3], [1,4,5], [1], [2,5], [2,4]]
g = GNNGraph(data)

# Number of nodes, edges, and batched graphs
g.num_nodes  # 5
g.num_edges  # 10
g.num_graphs # 1

# Same graph in COO representation
s = [1,1,2,2,2,3,4,4,5,5]
t = [2,3,1,4,5,3,2,5,2,4]
g = GNNGraph(s, t)

# From a LightGraphs' graph
g = GNNGraph(erdos_renyi(100, 20))

# Add 2 node feature arrays
g = GNNGraph(g, ndata = (x=rand(100, g.num_nodes), y=rand(g.num_nodes)))

# Add node features and edge features with default names `x` and `e`
g = GNNGraph(g, ndata = rand(100, g.num_nodes), edata = rand(16, g.num_edges))

g.ndata.x
g.ndata.e

# Send to gpu
g = g |> gpu

# Collect edges' source and target nodes.
# Both source and target are vectors of length num_edges
source, target = edge_index(g)
```
"""
struct GNNGraph{T<:Union{COO_T,ADJMAT_T}}
    graph::T
    num_nodes::Int
    num_edges::Int
    num_graphs::Int
    graph_indicator
    ndata::NamedTuple
    edata::NamedTuple
    gdata::NamedTuple
end

@functor GNNGraph

function GNNGraph(data;
                        num_nodes = nothing,
                        graph_indicator = nothing,
                        graph_type = :coo,
                        dir = :out,
                        ndata = (;),
                        edata = (;),
                        gdata = (;),
                        )

    @assert graph_type ∈ [:coo, :dense, :sparse] "Invalid graph_type $graph_type requested"
    @assert dir ∈ [:in, :out]

    if graph_type == :coo
        graph, num_nodes, num_edges = to_coo(data; num_nodes, dir)
    elseif graph_type == :dense
        graph, num_nodes, num_edges = to_dense(data; dir)
    elseif graph_type == :sparse
        graph, num_nodes, num_edges = to_sparse(data; dir)
    end

    num_graphs = !isnothing(graph_indicator) ? maximum(graph_indicator) : 1

    ndata = normalize_graphdata(ndata, default_name=:x, n=num_nodes)
    edata = normalize_graphdata(edata, default_name=:e, n=num_edges, duplicate_if_needed=true)
    gdata = normalize_graphdata(gdata, default_name=:u, n=num_graphs)

    GNNGraph(graph,
            num_nodes, num_edges, num_graphs,
            graph_indicator,
            ndata, edata, gdata)
end

# COO convenience constructors
GNNGraph(s::AbstractVector, t::AbstractVector, v = nothing; kws...) = GNNGraph((s, t, v); kws...)
GNNGraph((s, t)::NTuple{2}; kws...) = GNNGraph((s, t, nothing); kws...)

# GNNGraph(g::AbstractGraph; kws...) = GNNGraph(adjacency_matrix(g, dir=:out); kws...)

function GNNGraph(g::AbstractGraph; kws...)
    s = LightGraphs.src.(LightGraphs.edges(g))
    t = LightGraphs.dst.(LightGraphs.edges(g))
    if !LightGraphs.is_directed(g)
        # add reverse edges since GNNGraph are directed
        s, t = [s; t], [t; s]
    end
    GNNGraph((s, t); num_nodes=LightGraphs.nv(g), kws...)
end

function GNNGraph(g::GNNGraph; ndata=g.ndata, edata=g.edata, gdata=g.gdata)

    ndata = normalize_graphdata(ndata, default_name=:x, n=g.num_nodes)
    edata = normalize_graphdata(edata, default_name=:e, n=g.num_edges, duplicate_if_needed=true)
    gdata = normalize_graphdata(gdata, default_name=:u, n=g.num_graphs)

    GNNGraph(g.graph,
            g.num_nodes, g.num_edges, g.num_graphs,
            g.graph_indicator,
            ndata, edata, gdata)
end

function Base.show(io::IO, g::GNNGraph)
    println(io, "GNNGraph:
    num_nodes = $(g.num_nodes)
    num_edges = $(g.num_edges)
    num_graphs = $(g.num_graphs)")
    println(io, "    ndata:")
    for k in keys(g.ndata)
        println(io, "        $k => $(size(g.ndata[k]))")
    end
    println(io, "    edata:")
    for k in keys(g.edata)
        println(io, "        $k => $(size(g.edata[k]))")
    end
    println(io, "    gdata:")
    for k in keys(g.gdata)
        println(io, "        $k => $(size(g.gdata[k]))")
    end
end

"""
    edge_index(g::GNNGraph)

Return a tuple containing two vectors, respectively storing
the source and target nodes for each edges in `g`.

```julia
s, t = edge_index(g)
```
"""
edge_index(g::GNNGraph{<:COO_T}) = g.graph[1:2]

edge_index(g::GNNGraph{<:ADJMAT_T}) = to_coo(g.graph)[1][1:2]

edge_weight(g::GNNGraph{<:COO_T}) = g.graph[3]

LightGraphs.edges(g::GNNGraph) = zip(edge_index(g)...)

LightGraphs.edgetype(g::GNNGraph) = Tuple{Int, Int}

function LightGraphs.has_edge(g::GNNGraph{<:COO_T}, i::Integer, j::Integer)
    s, t = edge_index(g)
    return any((s .== i) .& (t .== j))
end

LightGraphs.has_edge(g::GNNGraph{<:ADJMAT_T}, i::Integer, j::Integer) = g.graph[i,j] != 0

LightGraphs.nv(g::GNNGraph) = g.num_nodes
LightGraphs.ne(g::GNNGraph) = g.num_edges
LightGraphs.has_vertex(g::GNNGraph, i::Int) = 1 <= i <= g.num_nodes
LightGraphs.vertices(g::GNNGraph) = 1:g.num_nodes

function LightGraphs.outneighbors(g::GNNGraph{<:COO_T}, i::Integer)
    s, t = edge_index(g)
    return t[s .== i]
end

function LightGraphs.outneighbors(g::GNNGraph{<:ADJMAT_T}, i::Integer)
    A = g.graph
    return findall(!=(0), A[i,:])
end

function LightGraphs.inneighbors(g::GNNGraph{<:COO_T}, i::Integer)
    s, t = edge_index(g)
    return s[t .== i]
end

function LightGraphs.inneighbors(g::GNNGraph{<:ADJMAT_T}, i::Integer)
    A = g.graph
    return findall(!=(0), A[:,i])
end

LightGraphs.is_directed(::GNNGraph) = true
LightGraphs.is_directed(::Type{GNNGraph}) = true

"""
    adjacency_list(g; dir=:out)

Return the adjacency list representation (a vector of vectors)
of the graph `g`.

Calling `a` the adjacency list, if `dir=:out`
`a[i]` will contain the neighbors of node `i` through
outgoing edges. If `dir=:in`, it will contain neighbors from
incoming edges instead.
"""
function adjacency_list(g::GNNGraph; dir=:out)
    @assert dir ∈ [:out, :in]
    fneighs = dir == :out ? outneighbors : inneighbors
    return [fneighs(g, i) for i in 1:g.num_nodes]
end

function LightGraphs.adjacency_matrix(g::GNNGraph{<:COO_T}, T::DataType=Int; dir=:out)
    if g.graph[1] isa CuVector
        A, n, m = to_dense(g.graph, T, num_nodes=g.num_nodes)
    else
        A, n, m = to_sparse(g.graph, T, num_nodes=g.num_nodes)
    end
    @assert size(A) == (n, n)
    return dir == :out ? A : A'
end

function LightGraphs.adjacency_matrix(g::GNNGraph{<:ADJMAT_T}, T::DataType=eltype(g.graph); dir=:out)
    @assert dir ∈ [:in, :out]
    A = g.graph
    A = T != eltype(A) ? T.(A) : A
    return dir == :out ? A : A'
end

function LightGraphs.degree(g::GNNGraph{<:COO_T}, T=Int; dir=:out)
    s, t = edge_index(g)
    degs = fill!(similar(s, T, g.num_nodes), 0)
    o = fill!(similar(s, Int, g.num_edges), 1)
    if dir ∈ [:out, :both]
        NNlib.scatter!(+, degs, o, s)
    end
    if dir ∈ [:in, :both]
        NNlib.scatter!(+, degs, o, t)
    end
    return degs
end

function LightGraphs.degree(g::GNNGraph{<:ADJMAT_T}, T=Int; dir=:out)
    @assert dir ∈ (:in, :out)
    A = adjacency_matrix(g, T)
    return dir == :out ? vec(sum(A, dims=2)) : vec(sum(A, dims=1))
end

function LightGraphs.laplacian_matrix(g::GNNGraph, T::DataType=Int; dir::Symbol=:out)
    A = adjacency_matrix(g, T; dir=dir)
    D = Diagonal(vec(sum(A; dims=2)))
    return D - A
end

"""
    normalized_laplacian(g, T=Float32; add_self_loops=false, dir=:out)

Normalized Laplacian matrix of graph `g`.

# Arguments

- `g`: A `GNNGraph`.
- `T`: result element type.
- `add_self_loops`: add self-loops while calculating the matrix.
- `dir`: the edge directionality considered (:out, :in, :both).
"""
function normalized_laplacian(g::GNNGraph, T::DataType=Float32;
                        add_self_loops::Bool=false, dir::Symbol=:out)
    Ã = normalized_adjacency(g, T; dir, add_self_loops)
    return I - Ã
end

function normalized_adjacency(g::GNNGraph, T::DataType=Float32;
                        add_self_loops::Bool=false, dir::Symbol=:out)
    A = adjacency_matrix(g, T; dir=dir)
    if add_self_loops
        A = A + I
    end
    degs = vec(sum(A; dims=2))
    inv_sqrtD = Diagonal(inv.(sqrt.(degs)))
    return inv_sqrtD * A * inv_sqrtD
end

@doc raw"""
    scaled_laplacian(g, T=Float32; dir=:out)

Scaled Laplacian matrix of graph `g`,
defined as ``\hat{L} = \frac{2}{\lambda_{max}} L - I`` where ``L`` is the normalized Laplacian matrix.

# Arguments

- `g`: A `GNNGraph`.
- `T`: result element type.
- `dir`: the edge directionality considered (:out, :in, :both).
"""
function scaled_laplacian(g::GNNGraph, T::DataType=Float32; dir=:out)
    L = normalized_laplacian(g, T)
    @assert issymmetric(L) "scaled_laplacian only works with symmetric matrices"
    λmax = _eigmax(L)
    return  2 / λmax * L - I
end

# _eigmax(A) = eigmax(Symmetric(A)) # Doesn't work on sparse arrays
function _eigmax(A)
    x0 = _rand_dense_vector(A)
    KrylovKit.eigsolve(Symmetric(A), x0, 1, :LR)[1][1] # also eigs(A, x0, nev, mode) available
end

_rand_dense_vector(A::AbstractMatrix{T}) where T = randn(float(T), size(A, 1))
_rand_dense_vector(A::CUMAT_T)= CUDA.randn(size(A, 1))

# Eigenvalues for cuarray don't seem to be well supported.
# https://github.com/JuliaGPU/CUDA.jl/issues/154
# https://discourse.julialang.org/t/cuda-eigenvalues-of-a-sparse-matrix/46851/5

"""
    add_self_loops(g::GNNGraph)

Return a graph with the same features as `g`
but also adding edges connecting the nodes to themselves.

Nodes with already existing
self-loops will obtain a second self-loop.
"""
function add_self_loops(g::GNNGraph{<:COO_T})
    s, t = edge_index(g)
    @assert g.edata === (;)
    @assert edge_weight(g) === nothing
    n = g.num_nodes
    nodes = convert(typeof(s), [1:n;])
    s = [s; nodes]
    t = [t; nodes]

    GNNGraph((s, t, nothing),
        g.num_nodes, length(s), g.num_graphs,
        g.graph_indicator,
        g.ndata, g.edata, g.gdata)
end

function add_self_loops(g::GNNGraph{<:ADJMAT_T})
    A = g.graph
    @assert g.edata === (;)
    A = A + I
    num_edges =  g.num_edges + g.num_nodes
    GNNGraph(A,
            g.num_nodes, num_edges, g.num_graphs,
            g.graph_indicator,
            g.ndata, g.edata, g.gdata)
end

function remove_self_loops(g::GNNGraph{<:COO_T})
    s, t = edge_index(g)
    # TODO remove these constraints
    @assert g.edata === (;)
    @assert edge_weight(g) === nothing

    mask_old_loops = s .!= t
    s = s[mask_old_loops]
    t = t[mask_old_loops]

    GNNGraph((s, t, nothing),
            g.num_nodes, length(s), g.num_graphs,
            g.graph_indicator,
            g.ndata, g.edata, g.gdata)
end

function _catgraphs(g1::GNNGraph{<:COO_T}, g2::GNNGraph{<:COO_T})
    s1, t1 = edge_index(g1)
    s2, t2 = edge_index(g2)
    nv1, nv2 = g1.num_nodes, g2.num_nodes
    s = vcat(s1, nv1 .+ s2)
    t = vcat(t1, nv1 .+ t2)
    w = cat_features(edge_weight(g1), edge_weight(g2))

    ind1 = isnothing(g1.graph_indicator) ? fill!(similar(s1, Int, nv1), 1) : g1.graph_indicator
    ind2 = isnothing(g2.graph_indicator) ? fill!(similar(s2, Int, nv2), 1) : g2.graph_indicator
    graph_indicator = vcat(ind1, g1.num_graphs .+ ind2)

    GNNGraph((s, t, w),
            nv1 + nv2, g1.num_edges + g2.num_edges, g1.num_graphs + g2.num_graphs,
            graph_indicator,
            cat_features(g1.ndata, g2.ndata),
            cat_features(g1.edata, g2.edata),
            cat_features(g1.gdata, g2.gdata))
end

### Cat public interfaces #############

"""
    blockdiag(xs::GNNGraph...)

Equivalent to [`Flux.batch`](@ref).
"""
function SparseArrays.blockdiag(g1::GNNGraph, gothers::GNNGraph...)
    g = g1
    for go in gothers
        g = _catgraphs(g, go)
    end
    return g
end

"""
    batch(xs::Vector{<:GNNGraph})

Batch together multiple `GNNGraph`s into a single one
containing the total number of nodes and edges of the original graphs.

Equivalent to [`SparseArrays.blockdiag`](@ref).
"""
Flux.batch(xs::Vector{<:GNNGraph}) = blockdiag(xs...)

### LearnBase compatibility
LearnBase.nobs(g::GNNGraph) = g.num_graphs
LearnBase.getobs(g::GNNGraph, i) = getgraph(g, i)[1]

# Flux's Dataloader compatibility. Related PR https://github.com/FluxML/Flux.jl/pull/1683
Flux.Data._nobs(g::GNNGraph) = g.num_graphs
Flux.Data._getobs(g::GNNGraph, i) = getgraph(g, i)[1]

#########################
Base.:(==)(g1::GNNGraph, g2::GNNGraph) = all(k -> getfield(g1,k)==getfield(g2,k), fieldnames(typeof(g1)))

"""
    getgraph(g::GNNGraph, i)

Return the getgraph of `g` induced by those nodes `v`
for which `g.graph_indicator[v] ∈ i`. In other words, it
extract the component graphs from a batched graph.

It also returns a vector `nodes` mapping the new nodes to the old ones.
The node `i` in the getgraph corresponds to the node `nodes[i]` in `g`.
"""
getgraph(g::GNNGraph, i::Int) = getgraph(g::GNNGraph{<:COO_T}, [i])

function getgraph(g::GNNGraph{<:COO_T}, i::AbstractVector{Int})
    if g.graph_indicator === nothing
        @assert i == [1]
        return g
    end

    node_mask = g.graph_indicator .∈ Ref(i)

    nodes = (1:g.num_nodes)[node_mask]
    nodemap = Dict(v => vnew for (vnew, v) in enumerate(nodes))

    graphmap = Dict(i => inew for (inew, i) in enumerate(i))
    graph_indicator = [graphmap[i] for i in g.graph_indicator[node_mask]]

    s, t, w = g.graph
    edge_mask = s .∈ Ref(nodes)
    s = [nodemap[i] for i in s[edge_mask]]
    t = [nodemap[i] for i in t[edge_mask]]
    w = isnothing(w) ? nothing : w[edge_mask]

    ndata = getobs(g.ndata, node_mask)
    edata = getobs(g.edata, edge_mask)
    gdata = getobs(g.gdata, i)

    num_nodes = length(graph_indicator)
    num_edges = length(s)
    num_graphs = length(i)

    gnew = GNNGraph((s,t,w),
                num_nodes, num_edges, num_graphs,
                graph_indicator,
                ndata, edata, gdata)
    return gnew, nodes
end

function node_features(g::GNNGraph)
    if isempty(g.ndata)
        return nothing
    elseif length(g.ndata) > 1
        @error "Multiple feature arrays, access directly through `g.ndata`"
    else
        return g.ndata[1]
    end
end

function edge_features(g::GNNGraph)
    if isempty(g.edata)
        return nothing
    elseif length(g.edata) > 1
        @error "Multiple feature arrays, access directly through `g.edata`"
    else
        return g.edata[1]
    end
end

function graph_features(g::GNNGraph)
    if isempty(g.gdata)
        return nothing
    elseif length(g.gdata) > 1
        @error "Multiple feature arrays, access directly through `g.gdata`"
    else
        return g.gdata[1]
    end
end

@non_differentiable normalized_laplacian(x...)
@non_differentiable normalized_adjacency(x...)
@non_differentiable scaled_laplacian(x...)
@non_differentiable adjacency_matrix(x...)
@non_differentiable adjacency_list(x...)
@non_differentiable degree(x...)
@non_differentiable add_self_loops(x...)     # TODO this is wrong, since g carries feature arrays, needs rrule
@non_differentiable remove_self_loops(x...)  # TODO this is wrong, since g carries feature arrays, needs rrule

# # delete when https://github.com/JuliaDiff/ChainRules.jl/pull/472 is merged
# function ChainRulesCore.rrule(::typeof(copy), x)
#     copy_pullback(ȳ) = (NoTangent(), ȳ)
#     return copy(x), copy_pullback
# end


using POMDPs
using POMDPModels
using POMDPModelTools
using StaticArrays
using JLD2, FileIO
using Random
using Logging
using Distributions
using GridInterpolations
using LocalFunctionApproximation
using LocalApproximationValueIteration
using JSON
using CMSSPs

global_logger(SimpleLogger(stderr, Logging.Error))

rng = MersenneTwister(3456)

param_files = ["./paramsets/scale-1.toml","./paramsets/simtime-1.toml","./paramsets/cost-1.toml"]

scale_file = param_files[1]
simtime_file = param_files[2]
cost_file = param_files[3]

inhor_fn = ARGS[1]
outhor_fn = ARGS[2]
infhor_fn = "./policies/dreamr-uf-params111.jld2"

ep_file_prefix = ARGS[3]
num_eps = parse(Int64, ARGS[4])
outfn = ARGS[5]

# Parse parameters
params = parse_params(scale_file=scale_file, simtime_file=simtime_file, cost_file=cost_file)

@info "Loading policies"
flight_policy = load_localapproxvi_policy_from_jld2(infhor_fn)
hopon_policy = load_modal_horizon_policy_localapprox(inhor_fn, outhor_fn)
ride_policy = DREAMRDeterministicPolicy(params)

# Create modal policies
modal_policies = Dict(FLIGHT => ModalFinInfHorPolicy(hopon_policy, flight_policy),
                      RIDE => ModalFinInfHorPolicy(ride_policy,flight_policy))

deltaT = convert(Int64, (params.time_params.MAX_REPLAN_TIMESTEP/params.time_params.MDP_TIMESTEP))

costs = Vector{Float64}(undef, 0)
steps = Vector{Int64}(undef, 0)
mode_switches = Vector{Int64}(undef, 0)
unsuccessful = Vector{Int64}(undef, 0)

for iter = 1:num_eps

    episode_dict = Dict

    ep_file = string(ep_file_prefix,"-",iter,".json")

    @show ep_file

    open(ep_file,"r") do f
        episode_dict = JSON.parse(f)
    end

    context_set = get_dreamr_episode_context(episode_dict["epochs"])

    start_pos = Point(episode_dict["start_pos"][1], episode_dict["start_pos"][2])
    goal_pos = Point(episode_dict["goal_pos"][1], episode_dict["goal_pos"][2])

    start_state = DREAMRStateType{MultiRotorUAVState}(FLIGHT, get_state_at_rest(MultiRotorUAVState, start_pos))

    # Create CMSSP
    cmssp = create_dreamr_cmssp(MultiRotorUAVState, MultiRotorUAVAction, context_set, goal_pos, params)

    # Create solver
    dreamr_solver = DREAMRSolverType{MultiRotorUAVState,typeof(rng)}(0, modal_policies,
                                    deltaT, [FLIGHT], start_state, DREAMRBookkeeping(), rng)

    (cost, timesteps, successful, num_mode_switches) = solve(dreamr_solver, cmssp)

    if successful
        energy_cost = cost - params.cost_params.TIME_COEFFICIENT*params.time_params.MDP_TIMESTEP*timesteps
        push!(costs, energy_cost)
        push!(steps, timesteps)
        push!(mode_switches, num_mode_switches)
    else
        push!(unsuccessful, iter)
    end
end

results_dict = Dict("costs"=>costs, "steps"=>steps, "mode_switches"=>mode_switches, "failed_iters"=>unsuccessful)

open(outfn, "w") do f
    JSON.print(f, results_dict, 2)
end


export CRRLearner

"""
    CRRLearner(;kwargs)

See paper: [Critic Regularized Regression](https://arxiv.org/abs/2006.15134).

# Keyword arguments

- `approximator`::[`ActorCritic`](@ref): used to get Q-values (Critic) and logits (Actor) of a state.
- `target_approximator`::[`ActorCritic`](@ref): similar to `approximator`, but used to estimate the target.
- `γ::Float32`, reward discount rate.
- `batch_size::Int=32`
- `policy_improvement_mode::Symbol=:exp`, type of the weight function f. Possible values: :binary/:exp.
- `ratio_upper_bound::Float32`, when `policy_improvement_mode` is ":exp", the value of the exp function is upper-bounded by this parameter.
- `beta::Float32`,  when `policy_improvement_mode` is ":exp", this is the denominator of the exp function.
- `advantage_estimator::Symbol=:mean`, type of the advantage estimate \\hat{A}. Possible values: :mean/:max.
- `m::Int=4`, when `continuous=true`, sample `m` action to estimate \\hat{A}.
- `update_freq::Int`: the frequency of updating the `approximator`.
- `update_step::Int=0`
- `target_update_freq::Int`: the frequency of syncing `target_approximator`.
- `continuous::Bool`: type of action space.
- `rng = Random.GLOBAL_RNG`
"""
mutable struct CRRLearner{
    Aq<:ActorCritic,
    At<:ActorCritic,
    R<:AbstractRNG,
} <: AbstractLearner
    approximator::Aq
    target_approximator::At
    γ::Float32
    batch_size::Int
    policy_improvement_mode::Symbol
    ratio_upper_bound::Float32
    beta::Float32
    advantage_estimator::Symbol
    m::Int
    update_freq::Int
    update_step::Int
    target_update_freq::Int
    continuous::Bool
    rng::R
    # for logging
    actor_loss::Float32
    critic_loss::Float32
end

function CRRLearner(;
    approximator::Aq,
    target_approximator::At,
    γ::Float32 = 0.99f0,
    batch_size::Int = 32,
    policy_improvement_mode::Symbol = :exp,
    ratio_upper_bound::Float32 = 20.0f0,
    beta::Float32 = 1.0f0,
    advantage_estimator::Symbol = :mean,
    m::Int = 4,
    update_freq::Int = 10,
    update_step::Int = 0,
    target_update_freq::Int = 100,
    continuous::Bool,
    rng = Random.GLOBAL_RNG,
) where {Aq<:ActorCritic, At<:ActorCritic}
    copyto!(approximator, target_approximator)
    CRRLearner(
        approximator,
        target_approximator,
        γ,
        batch_size,
        policy_improvement_mode,
        ratio_upper_bound,
        beta,
        advantage_estimator,
        m,
        update_freq,
        update_step,
        target_update_freq,
        continuous,
        rng,
        0.0f0,
        0.0f0,
    )
end

Flux.functor(x::CRRLearner) = (Q = x.approximator, Qₜ = x.target_approximator),
y -> begin
    x = @set x.approximator = y.Q
    x = @set x.target_approximator = y.Qₜ
    x
end

function (learner::CRRLearner)(env)
    s = state(env)
    s = Flux.unsqueeze(s, ndims(s) + 1)
    s = send_to_device(device(learner), s)
    if learner.continuous
        learner.approximator.actor(s; is_sampling=true) |> vec |> send_to_host
    else
        learner.approximator.actor(s) |> vec |> send_to_host
    end
end

function RLBase.update!(learner::CRRLearner, batch::NamedTuple)
    if learner.continuous
        continuous_update!(learner, batch)
    else
        discrete_update!(learner, batch)
    end
end

function continuous_update!(learner::CRRLearner, batch::NamedTuple)
    AC = learner.approximator
    target_AC = learner.target_approximator
    γ = learner.γ
    beta = learner.beta
    batch_size = learner.batch_size
    policy_improvement_mode = learner.policy_improvement_mode
    ratio_upper_bound = learner.ratio_upper_bound
    advantage_estimator = learner.advantage_estimator
    D = device(AC)

    s, a, r, t, s′ = (send_to_device(D, batch[x]) for x in SARTS)
    a = reshape(a, :, batch_size)
    r = reshape(r, :, batch_size)
    t = reshape(t, :, batch_size)

    target_a_t = target_AC.actor(s′; is_sampling=true)
    target_q_input = vcat(s′, target_a_t)
    expected_target_q = target_AC.critic(target_q_input)

    target = r .+ γ .* (1 .- t) .* expected_target_q

    q_t = Matrix{Float32}(undef, learner.m, batch_size)
    for i in 1:learner.m
        a_sample = AC.actor(s; is_sampling=true)
        q_t[i, :] = AC.critic(vcat(s, a_sample))
    end

    ps = Flux.params(AC)
    gs = gradient(ps) do
        # Critic loss
        qa_t = AC.critic(vcat(s, a))
        critic_loss = Flux.Losses.mse(qa_t, target)

        # Actor loss
        log_π = AC.actor.model(s, a)

        if advantage_estimator == :max
            advantage = qa_t .- maximum(q_t, dims=1)
        elseif advantage_estimator == :mean
            advantage = qa_t .- mean(q_t, dims=1)
        else
            error("Wrong parameter.")
        end

        if policy_improvement_mode == :binary
            actor_loss_coef = (advantage .> 0.0f0)
        elseif policy_improvement_mode == :exp
            actor_loss_coef = clamp.(exp.(advantage ./ beta), 0, ratio_upper_bound)
        else
            error("Wrong parameter.")
        end

        actor_loss = mean(-log_π .* actor_loss_coef)

        ignore() do
            learner.actor_loss = actor_loss
            learner.critic_loss = critic_loss
        end

        actor_loss + critic_loss
    end

    update!(AC, gs)
end

function discrete_update!(learner::CRRLearner, batch::NamedTuple)
    AC = learner.approximator
    target_AC = learner.target_approximator
    γ = learner.γ
    beta = learner.beta
    batch_size = learner.batch_size
    policy_improvement_mode = learner.policy_improvement_mode
    ratio_upper_bound = learner.ratio_upper_bound
    advantage_estimator = learner.advantage_estimator
    D = device(AC)

    s, a, r, t, s′ = (send_to_device(D, batch[x]) for x in SARTS)
    a = CartesianIndex.(a, 1:batch_size)
    r = reshape(r, :, batch_size)
    t = reshape(t, :, batch_size)

    target_a_t = softmax(target_AC.actor(s′))
    target_q_t = target_AC.critic(s′)
    expected_target_q = sum(target_a_t .* target_q_t, dims=1)

    target = r .+ γ .* (1 .- t) .* expected_target_q

    ps = Flux.params(AC)
    gs = gradient(ps) do
        # Critic loss
        q_t = AC.critic(s)
        qa_t = reshape(q_t[a], :, batch_size)
        critic_loss = Flux.Losses.mse(qa_t, target)

        # Actor loss
        a_t = softmax(AC.actor(s))

        if advantage_estimator == :max
            advantage = qa_t .- maximum(q_t, dims=1)
        elseif advantage_estimator == :mean
            advantage = qa_t .- mean(q_t, dims=1)
        else
            error("Wrong parameter.")
        end

        if policy_improvement_mode == :binary
            actor_loss_coef = (advantage .> 0.0f0)
        elseif policy_improvement_mode == :exp
            actor_loss_coef = clamp.(exp.(advantage ./ beta), 0, ratio_upper_bound)
        else
            error("Wrong parameter.")
        end

        actor_loss = mean(-log.(a_t[a]) .* actor_loss_coef)

        ignore() do
            learner.actor_loss = actor_loss
            learner.critic_loss = critic_loss
        end

        actor_loss + critic_loss
    end

    update!(AC, gs)
end


import Combinatorics
using Random
using DataFrames
import CSV
using Printf
#import PyPlot

using JuMP
#using GLPK
using Gurobi

using PyPlot, PyCall
#@pyimport mpl_toolkits.axes_grid1 as axes_grid1
axes_grid1 = pyimport("mpl_toolkits.axes_grid1")

include("DataLog.jl")
include("Agents/ScenarioTree.jl")

include("Agents/Parameters.jl")
include("Agents/PowerFlow.jl")
include("Agents/EnergyProducer.jl")

include("Agents/System.jl")

#include("RepresentativeDays_ABM/RepresentativeDays.jl")
include("RepresentativeDays_ABM/RepresentativeData.jl")

#GUROBI_ENV = Gurobi.Env()
#setparams!(GUROBI_ENV, OutputFlag = 0,Method = 1)

Random.seed!(12)

#Initialize global parameters
investmentModelSelection = :GEP # :EMCAS, :EMLab, :GEP
investmentSequence = :sequential
simulationSettings = SimulationSettings(day2hour,year2day,segmentNumbers,downpaymentRatio,representativeData,investmentSequence,investmentModelSelection)
systemSettings = SystemSettings(GenCoNumbers,totalHorizon,milestoneYear,interestRate,simulationSettings)

function InitSystem(systemSettings::SystemSettings,genCoHorizon::Real; kwargs...)
    V = Dict(kwargs)
    if haskey(V,:gapCoverPercentage) && haskey(V,:divisionTypeShare)
        gapCoverPercentage = V[:gapCoverPercentage]
        divisionTypeShare = V[:divisionTypeShare]
    end
##########################################################################################################
    #Initialize technology pool
    technologyPool = Array{Technology,1}(undef,0)
    initialTechnologyHistory = TechnologyHistory(systemSettings)

    for i in 1:length(technologyType)
        thisDesignProperty = DesignProperty(availability[i], constructionTimeStamp[i],constructionTime[i],deconstructionTime[i],efficiency[i],emission_ton_perMWh[i],lifeTime[i],unitCapacity[i])
        thisEconomicProperty = EconomicProperty(unitConstructionCost[i], unitDeconstructionCost[i], unitVOMCost[i],unitFOMCost[i],systemSettings)

        if technologyType[i] in [:Base,:Mid,:Peak]
            thisTechnology = ConventionalPlants(technologyType[i],thisDesignProperty, thisEconomicProperty,systemSettings)
            thisTechnology.UpdateCost(emissionPrice,fuelPrice[i])
        end

        #if technologyType[i] in [:OnShoreWind,:OffShoreWind,:Solar]
        #    thisTechnology = RenewablePlants(technologyType[i],thisDesignProperty, thisEconomicProperty, capacityFactorImproveRate,[:offline,:subsidized],systemSettings)
        #end

        thisTechnology.MultipleUnits(1,systemSettings)
        println(thisTechnology.economicProperties.constructionCost)
        push!(technologyPool, thisTechnology)
    end

    ##########################################################################################################
    initial_ExistingCapacity_TimeStamp = [11 11 11 11;
                                          11 11 11 11;
                                          11 11 11 11]
    initial_Capacity = Array{Technology,1}(undef,0)

    for pre_TimeStamp in 1:size(initial_ExistingCapacity_TimeStamp)[2]
        for techCount in 1:length(technologyPool)
            thisInitialPlant = deepcopy(technologyPool[techCount])
            for techNum in 1:initial_ExistingCapacity_TimeStamp[techCount,pre_TimeStamp]
                thisInitialPlant.designProperties.constructionTimeStamp = (pre_TimeStamp-4) * systemSettings.milestoneYear
                push!(initial_Capacity,deepcopy(thisInitialPlant))
            end
        end
    end

    ##########################################################################################################
    #Initialize GenCos and ownerships

    systemCapacityMix = Array{Array{Technology,1},1}(undef,0)
    systemHistory = SystemHistory()
    allGenerators = Array{Technology,1}(undef,0)
    allOwnerShips = Array{GenCoPortfolioOwnership,1}(undef,0)
    allGenCos = Array{GenCo,1}(undef,0)
    for i in 1:systemSettings.GenCoNumbers
        thisName = string("GenCo", i)
        thisTechnologies = Array{Technology,1}(undef,0)
        thisTechnologyPortfolio = TechnologyPortfolio(Array{Technology,1}(undef,0))
        thisCompetitorExpansion = Array{CapacityExpansionEdge,1}(undef,0)
        thisCompetitorDivision=Array{CapacityDivisionEdge,1}(undef,0)

        if haskey(V,:gapCoverPercentage) && haskey(V,:divisionTypeShare)
            for m in 1:length(gapCoverProbability[i])
                push!(thisCompetitorExpansion, CapacityExpansionEdge(gapCoverProbability[i][m],gapCoverPercentage[i][m]))
            end
            for k in 1:length(divisionProbability[i])
                push!(thisCompetitorDivision, CapacityDivisionEdge(divisionProbability[i][k],divisionTypeShare[i][k]))
            end
        end

        thisGenCo = GenCo(thisName, genCoHorizon,thisTechnologyPortfolio,deepcopy(technologyPool);competitorExpansion = thisCompetitorExpansion,competitorDivision=thisCompetitorDivision)
        thisOwnerShip = GenCoPortfolioOwnership(thisGenCo, thisTechnologyPortfolio)
        for thisType in technologyType
            thisGenCo.history.investmentYearByType[thisType] = zeros(totalHorizon)
        end
        push!(allOwnerShips,thisOwnerShip)
        push!(allGenCos,thisGenCo)
    end

    #Randomly distribute initial capacities
    for n in 1:length(initial_Capacity)
        luckyGenCo = allGenCos[rand(1:GenCoNumbers)]
        push!(luckyGenCo.technologyPortfolio.technologies, initial_Capacity[n])
        push!(allGenerators, initial_Capacity[n])
    end

    #push!(systemHistory.systemCapacityMix, deepcopy(allGenerators))
    println("Generator numbers: ", length(allGenerators))

    println("OwnerShips: ", length(allOwnerShips))

    ##########################################################################################################
    #Initialize consumers
    #priceCap = 200.0 #€/MWh
    #averageLoad = 10000 #MWh
    systemLoad = Load(priceCap,0.0,representativeData[:Load])

    systemState = SystemState(1,allGenerators)
    ##########################################################################################################
    #Initialize market operators
    MO = MarketOperator("MarketOperator")

    ##########################################################################################################
    #Initialize powerflows
    allPowerFlows = Array{PowerFlow,1}(undef,0)
    for generator in allGenerators
        #if typeof(generator) != VirtualPlants
        thisPowerFlow = PowerFlow(generator, MO, 0,0)
        push!(allPowerFlows, thisPowerFlow)
        #end
    end

    return MO,allGenCos,systemLoad,allPowerFlows,technologyPool,systemState,systemHistory
end


# Data structure of particle and basic functions are defined in this file.

mutable struct Ptype
  x::Float64
  p::Float64
  m::Float64
  rho::Float64
  U::Float64
  P::Float64

  x_old::Float64
  p_old::Float64
  U_old::Float64

  F1::Float64
  dU::Float64
  hsml::Float64
  c11::Float64

  alpha::Float64
  alpha_old::Float64
  dalpha::Float64

  dydh::Float64

  function Ptype(x, v, m, rho, P)
    p = m * v
    U = P / (gamma - 1) * m / rho
    x_old = 0
    p_old = 0
    U_old = 0
    F1 = 0
    dU = 0
    hsml = 0
    c11 = 0
    alpha = alpha_min
    alpha_old = 0
    dalpha = 0
    dydh = 0
    return new(
      x,
      p,
      m,
      rho,
      U,
      P,
      x_old,
      p_old,
      U_old,
      F1,
      dU,
      hsml,
      c11,
      alpha,
      alpha_old,
      dalpha,
      dydh,
    )
  end
end

function Z(i::Int)
  if volume_element == "mass"
    return sph[i].m
  elseif volume_element == "U"
    return sph[i].U
  end
end

function y(i::Int)
  if volume_element == "mass"
    return sph[i].rho
  elseif volume_element == "U"
    return sph[i].P / (gamma - 1)
  end
end

function v(i::Int)
  return sph[i].p / sph[i].m
end

function dist(i::Int, j::Int)
  rij = abs(sph[i].x - sph[j].x)
  rij = min(rij, lbox - rij)
  return rij
end

function eij(i::Int, j::Int)
  rij = abs(sph[i].x - sph[j].x)
  eij = (sph[i].x - sph[j].x) / rij

  if abs(sph[i].x - sph[j].x) > lbox/2
    eij *= -1
  end

  return eij
end

function vij(i::Int, j::Int)
  vij = v(i) - v(j)

  return vij
end

function cs(i)
  cs = sqrt(gamma * sph[i].P / sph[i].rho)
  return cs
end