text:
using BlossomV
using Metis
using SparseArrays
using LightGraphs: add_edge!
using SimpleWeightedGraphs
using MetaGraphs
using ..Utils

export learn_vtree

#############
# Metis top down method
#############

# Add edge weights to Metis.jl
using Metis: idx_t, ishermitian

struct WeightedGraph
    nvtxs::idx_t
    xadj::Vector{idx_t}
    adjncy::Vector{idx_t}
    adjwgt::Vector{idx_t} # edge weights
    WeightedGraph(nvtxs, xadj, adjncy, adjwgt) = new(nvtxs, xadj, adjncy, adjwgt)
end

function my_graph(G::SparseMatrixCSC; check_hermitian=true)
    if check_hermitian
        ishermitian(G) || throw(ArgumentError("matrix must be Hermitian"))
    end
    N = size(G, 1)
    xadj = Vector{idx_t}(undef, N+1)
    xadj[1] = 1
    adjncy = Vector{idx_t}(undef, nnz(G))
    adjncy_i = 0
    adjwgt = Vector{idx_t}(undef, nnz(G))
    @inbounds for j in 1:N
        n_rows = 0
        for k in G.colptr[j] : (G.colptr[j+1] - 1)
            i = G.rowval[k]
            if i != j # don't include diagonal elements
                n_rows += 1
                adjncy_i += 1
                adjncy[adjncy_i] = i
                adjwgt[adjncy_i] = G[i, j]
            end
        end
        xadj[j+1] = xadj[j] + n_rows
    end
    resize!(adjncy, adjncy_i)
    resize!(adjwgt, adjncy_i)
    return WeightedGraph(idx_t(N), xadj, adjncy, adjwgt)
end

function my_partition(G::WeightedGraph, nparts::Integer)
    part = Vector{Metis.idx_t}(undef, G.nvtxs)
    edgecut = fill(idx_t(0), 1)
    # if alg === :RECURSIVE
        Metis.METIS_PartGraphRecursive(G.nvtxs, idx_t(1), G.xadj, G.adjncy, C_NULL, C_NULL, G.adjwgt,
                                 idx_t(nparts), C_NULL, C_NULL, Metis.options, edgecut, part)
    # elseif alg === :KWAY
    #     Metis.METIS_PartGraphKway(G.nvtxs, idx_t(1), G.xadj, G.adjncy, C_NULL, C_NULL, G.adjwgt,
    #                         idx_t(nparts), C_NULL, C_NULL, Metis.options, edgecut, part)
    # else
    #     throw(ArgumentError("unknown algorithm $(repr(alg))"))
    # end
    return part
end

my_partition(G, nparts) = my_partition(my_graph(G), nparts)

"Metis top down method"
function metis_top_down(data::DataFrame;α)
    δINT = 999999
    MIN_INT = 1
    MAX_INT = δINT + MIN_INT

    weight=ones(Float64, num_examples(data))
    (_, mi) = mutual_information(data, weight; α=α)
    vars = Var.(collect(1:num_features(data)))
    info = to_long_mi(mi, MIN_INT, MAX_INT)

    function f(leafs::Vector{PlainVtreeLeafNode})::Tuple{Vector{PlainVtreeLeafNode}, Vector{PlainVtreeLeafNode}}
        var2leaf = Dict([(variable(x),x) for x in leafs])
        vertices = sort(variable.(leafs))
        sub_context = info[vertices, vertices]
        len = length(vertices)
        for i in 1 : len
            sub_context[i, i] = 0
        end
        g = convert(SparseMatrixCSC, sub_context)
        partition = my_partition(my_graph(g), 2)

        subsets = (Vector{PlainVtreeLeafNode}(), Vector{PlainVtreeLeafNode}())
        for (index, p) in enumerate(partition)
            push!(subsets[p], var2leaf[vertices[index]])
        end

        return subsets
    end
    return f
end

#############
# Blossom bottom up method
#############

"Blossom bottom up method, vars are not used"
function blossom_bottom_up(data::DataFrame;α)
    weight = ones(Float64, num_examples(data))
    (_, mi) = mutual_information(data, weight; α)
    vars = Var.(collect(1:num_features(data)))
    info = round.(Int64, 1000001 .+ to_long_mi(mi, -1, -1000000))

    function f(leaf::Vector{<:Vtree})
        variable_sets = collect.(variables.(leaf))

        # even number of nodes, use blossomv alg
        function blossom_bottom_up_even!(variable_sets)::Tuple{Vector{Tuple{Var, Var}}, Int64}
            # 1. calculate pMI
            pMI = set_mutual_information(info, variable_sets)
            pMI = round.(Int64, pMI)

            # 2. solve by blossomv alg
            len = length(variable_sets)
            m = Matching(len)
            for i in 1 : len, j in i + 1 : len
                add_edge(m, i - 1, j - 1, pMI[i, j]) # blossomv index start from 0
            end

            solve(m)
            all_matches = Set{Tuple{Var, Var}}()
            for v in 1 : len
                push!(all_matches, order_asc(v, get_match(m, v - 1) + 1))
            end

            # 3. calculate scores, map index to var
            all_matches = Vector(collect(all_matches))
            score = 0

            for i in 1 : length(all_matches)
                (x, y) = all_matches[i]
                score += pMI[x, y]
            end

            return (all_matches, score)
        end

        # odd number of nodes, try every 2 combinations
        function blossom_bottom_up_odd!(variable_sets)
            # try all len - 1 conditions, find best score(minimun cost)
            (best_matches, best_score) = (nothing, typemax(Int64))
            len = length(variable_sets)
            for index in 1 : len
                indices = [collect(1:index-1);collect(index+1:len)]
                (matches, score) = blossom_bottom_up_even!(variable_sets[indices])
                if score < best_score
                    (best_matches, best_score) = ([[(indices[l], indices[r]) for (l,r) in matches];[index]], score)
                end
            end
            return (best_matches, best_score)
        end

        if length(variable_sets) % 2 == 0
            (matches, score) = blossom_bottom_up_even!(variable_sets)
        else
            (matches, score) = blossom_bottom_up_odd!(variable_sets)
        end

        pairs = []
        for x in matches
            if x isa Tuple
                push!(pairs, (leaf[x[1]], leaf[x[2]]))
            else
                push!(pairs, leaf[x])
            end
        end
        return pairs
    end
    return f
end

function learn_vtree(data::DataFrame; α=0.0, alg=:bottomup)
    if alg==:topdown
        PlainVtree(num_features(data), :topdown; f=metis_top_down(data;α))
    elseif alg==:bottomup
        PlainVtree(num_features(data), :bottomup; f=blossom_bottom_up(data;α))
    elseif alg==:clt
        clt = learn_chow_liu_tree(data)
        learn_vtree_from_clt(clt, vtree_mode="balanced")
    else
        error("Vtree learner $(alg) not supported.")
    end
end


text:
# functions for building various standard types of graphs

export Complete, Path, Cycle, RandomGraph, RandomRegular
export RandomTree, code_to_tree
export Grid, Wheel, Cube, BuckyBall
export Petersen, Kneser, Paley, Knight

"""
`Complete(n)` returns a complete graph with `n` vertices `1:n`.

`Complete(n,m)` returns a complete bipartite graph with `n` vertices
in one part and `m` vertices in the other.

`Complete([n1,n2,...,nt])` returns a complete multipartite graph with
parts of size `n1`, `n2`, ..., `nt`.
"""
# Create a complete graph
function Complete(n::Int)
    G = IntGraph(n)

    for k=1:n-1
        for j=k+1:n
            add!(G,j,k)
        end
    end
    return G
end

# Create a complete bipartite graph
function Complete(n::Int, m::Int)
    G = IntGraph(n+m)
    for u=1:n
        for v=n+1:n+m
            add!(G,u,v)
        end
    end
    return G
end

# Create the complete multipartite graph with given part sizes
function Complete(parts::Array{Int,1})
    # check all part sizes are positive
    for p in parts
        if p < 1
            error("All part sizes must be positive")
        end
    end

    n = sum(parts)
    G = IntGraph(n)

    np = length(parts)
    if np < 2
        return G
    end

    # create table of part ranges
    ranges = Array(Int,np,2)
    ranges[1,1] = 1
    ranges[1,2] = parts[1]
    for k = 2:np
        ranges[k,1] = ranges[k-1,2] + 1
        ranges[k,2] = ranges[k,1] + parts[k] - 1
    end

    # Add all edges between all parts
    for i=1:np-1
        for j=i+1:np
            for u=ranges[i,1]:ranges[i,2]
                for v=ranges[j,1]:ranges[j,2]
                    add!(G,u,v)
                end
            end
        end
    end

    return G
end

# Create a path graph on n vertices
"""
`Path(n)` creates a path graph with `n` vertices named `1:n`.

`Path(array)` creates a path graph with vertices `array[1]`,
`array[2]`, etc.
"""

function Path(n::Int)
    G = IntGraph(n)
    for v = 1:n-1
        add!(G,v,v+1)
    end
    return G
end

# Create a path graph from a list of vertices
function Path{T}(verts::Array{T})
    G = SimpleGraph{T}()
    n = length(verts)

    if n==1
        add!(G,verts[1])
    end
    for k = 1:n-1
        add!(G,verts[k],verts[k+1])
    end
    return G
end

# Create a cycle graph on n vertices
function Cycle(n::Int)
    if n<3
        error("Cycle requires 3 or more vertices")
    end
    G = Path(n)
    add!(G,1,n)
    return G
end

# Create the wheel graph on n vertices: a cycle on n-1 vertices plus
# an additional vertex adjacent to all the vertices on the wheel.
"""
`Wheel(n)` creates a wheel graph with `n` vertices. That is, a cycle
with `n-1` vertices `1:(n-1)` all adjacent to a common single vertex,
`n`.
"""
function Wheel(n::Int)
    if n < 4
        error("Wheel graphs must have at least 4 vertices")
    end
    G = Cycle(n-1)
    for k=1:n-1
        add!(G,k,n)
    end
    return G
end

# Create a grid graph
"""
`Grid(n,m)` creates an `n`-by-`m` grid graph. For other grids, we
suggest `Path(n1)*Path(n2)*Path(n3)` optionally wrapped in
`relabel`. See also: `Cube`.
"""
function Grid(n::Int, m::Int)
    G = SimpleGraph{Tuple{Int,Int}}()

    # add the vertices
    for u=1:n
        for v=1:m
            add!(G,(u,v))
        end
    end

    #horizontal edges
    for u=1:n
        for v=1:m-1
            add!(G,(u,v),(u,v+1))
        end
    end

    # vertical edges
    for v=1:m
        for u=1:n-1
            add!(G,(u,v),(u+1,v))
        end
    end
    return G
end

# Create an Erdos-Renyi random graph
"""
`RandomGraph(n,p=0.5)` creates an Erdos-Renyi random graph with `n`
vertices and edge probability `p`.
"""
function RandomGraph(n::Int, p::Real=0.5)
    G = IntGraph(n)

    # guess the size of the edge set to preallocate storage
    m = round(Int,n*n*p)+1

    # generate the edges
    for v=1:n-1
        for w=v+1:n
            if (rand() < p)
                add!(G,v,w)
            end
        end
    end
    return G
end

# Generate a random tree on vertex set 1:n. All n^(n-2) trees are
# equally likely.

"""
`RandomTree(n)` creates a random tree on `n` vertices each with
probability `1/n^(n-2)`.
"""
function RandomTree(n::Int)
    if n<0   # but we allow n==0 to give empty graph
        error("Number of vertices cannot be negative")
    end

    if n<2
        return IntGraph(n)
    end

    code = [ mod(rand(Int),n)+1 for _ in 1:n-2 ]
    return code_to_tree(code)
end

# This is a helper function for RandomTree that converts a Prufer code
# to a tree. No checks are done on the array fed into this function.
function code_to_tree(code::Array{Int,1})
    n = length(code)+2
    G = IntGraph(n)
    degree = ones(Int,n)  # initially all 1s

    #every time a vertex appears in code[], up its degree by 1
    for c in code
        degree[c]+=1
    end

    for u in code
        for v in 1:n
            if degree[v]==1
                add!(G,u,v)
                degree[u] -= 1
                degree[v] -= 1
                break
            end
        end
    end

    last = find(degree)
    add!(G,last[1],last[2])

    return G
end

# Create the Cube graph with 2^n vertices
"""
`Cube(n)` creates the `n`-dimensional cube graph. This graph has `2^n`
vertices named by all possible length-`n` strings of 0s and 1s. Two
vertices are adjacent iff they differ in exactly one position.
"""
function Cube(n::Integer=3)
    G = StringGraph()
    for u=0:2^n-1
        for shift=0:n-1
            v = (1<<shift) $ u
            add!(G,bin(u,n), bin(v,n))
        end
    end
    return G
end

# Create the BuckyBall graph

"""
`BuckyBall()` returns the Bucky ball graph.
"""
function BuckyBall()
    G = IntGraph()
    edges = [(1,3), (1,49), (1,60), (2,4), (2,10), (2,59),
	     (3,4), (3,37), (4,18), (5,7), (5,9), (5,13),
	     (6,8), (6,10), (6,17), (7,8), (7,21), (8,22),
	     (9,10), (9,57), (11,12), (11,13), (11,21), (12,28),
	     (12,48), (13,14), (14,47), (14,55), (15,16), (15,17),
	     (15,22), (16,26), (16,42), (17,18), (18,41), (19,20),
	     (19,21), (19,27), (20,22), (20,25), (23,24), (23,32),
	     (23,35), (24,26), (24,39), (25,26), (25,31), (27,28),
	     (27,31), (28,30), (29,30), (29,32), (29,36), (30,45),
	     (31,32), (33,35), (33,40), (33,51), (34,36), (34,46),
	     (34,52), (35,36), (37,38), (37,41), (38,40), (38,53),
	     (39,40), (39,42), (41,42), (43,44), (43,47), (43,56),
	     (44,46), (44,54), (45,46), (45,48), (47,48), (49,50),
	     (49,53), (50,54), (50,58), (51,52), (51,53), (52,54),
	     (55,56), (55,57), (56,58), (57,59), (58,60), (59,60),
	     ]
    for e in edges
        add!(G,e[1],e[2])
    end
    return G
end

# The Kneser graph Kneser(n,k) has C(n,k) vertices that are the
# k-element subsets of 1:n in which two vertices are adjacent if (as
# sets) they are disjoint. The Petersen graph is Kneser(5,2).

"""
`Kneser(n,m)` creates the Kneser graph whose vertices are all the
`m`-element subsets of `1:n` in which two vertices are adjacent iff
they are disjoint.
"""
function Kneser(n::Int,k::Int)
    A = collect(1:n)
    vtcs = [Set(v) for v in subsets(A,k)]
    G = SimpleGraph{Set{Int}}()

    for v in vtcs
        add!(G,v)
    end

    n = length(vtcs)
    for i=1:n-1
        u = vtcs[i]
        for j=i+1:n
            v = vtcs[j]
            if length(intersect(u,v))==0
                add!(G,u,v)
            end
        end
    end

    return G
end

# Create the Petersen graph.
"""
`Petersen()` returns the Petersen graph. The vertices are labeled as
the 2-element subsets of `1:5`. Wrap in `relabel` to have vertices
named `1:10`. See also: `Kneser`.
"""

Petersen() = Kneser(5,2)

# Create Paley graphs

"""
`Paley(p)` creates the Paley graph with `p` vertices named
`0:(p-1)`. Here `p` must be a prime with `p%4==1`. Vertices `u` and
`v` are adjacent iff `u-v` is a quadratic residue (perfect square)
modulo `p`.
"""
function Paley(p::Int)
    if mod(p,4) != 1 || ~isprime(p)
        error("p must be a prime congruent to 1 mod 4")
    end

    # Quadratic residues mod p
    qrlist = unique( [ mod(k*k,p) for k=1:p ] )

    G = IntGraph()
    for u = 0:p-1
        for k in qrlist
            v = mod(u+k,p)
            add!(G,u,v)
        end
    end
    return G
end

# Called by RandomRegular ... one step
function RandomRegularBuilder(n::Int, d::Int)
    # caller has already checked the values of n,d are legit
    vlist = randperm(n*d)
    G = IntGraph(n*d)
    for v=1:2:n*d
        add!(G,vlist[v], vlist[v+1])
    end

    for v = n:-1:1
        mushlist = collect( d*(v-1)+1 : v*d )
        for k=d-1:-1:1
            contract!(G,mushlist[k],mushlist[k+1])
        end
    end
    return relabel(G)
end

"""
`RandomRegular(n,d)` creates a random `d`-regular graph on `n`
vertices. This can take a while especially if the arguments are
large. Call with an optional third argument to activate verbose
progress reports: `RandomRegular(n,p,true)`.
"""
function RandomRegular(n::Int, d::Int, verbose::Bool=false)
    # sanity checks
    if n<1 || d<1 || (n*d)%2==1
        error("n,d must be positive integers and n*d even")
    end
    if verbose
        println("Trying to build ", d, "-regular graph on ",
                n, " vertices")
        count::Int = 0
    end

    while true
        if verbose
            count += 1
            println("Attempt ", count)
            tic()
        end
        g = RandomRegularBuilder(n,d)
        if verbose
            toc();
        end
        dlist = deg(g)
        if dlist[1] == dlist[n]
            if verbose
                println("Success")
            end
            return g
        end
        if verbose
            println("Failed; trying again")
        end
    end
end

"""
`Knight(r::Int=8,c::Int=8)` creates a Knight's Moves graph on a
`r`-by-`c` grid. That is, the vertices of this graph are the squares
of an `r`-by-`c` chess board. Two vertices are adjacent if a Knight
can go from one of these squares to the other in a single move.
"""
function Knight(r::Int=8,c::Int=8)
    vtcs = collect(product(1:r,1:c))
    G = SimpleGraph{Tuple{Int64,Int64}}()
    for v in vtcs
        add!(G,v)
    end

    for v in vtcs
        for w in vtcs
            xv = collect(v)
            xw = collect(w)
            z = sort(map(abs,xv - xw))
            if z==[1,2]
                add!(G,v,w)
            end
        end
    end

    return G
end

ext: jl


text:
using Sounds

lines, cols = displaysize(stdout)

clear() = join(repeat([""], cols), "\n") |> print
ft(t) = string(round(Int, 1e-6t)) * "s"

file = joinpath(@__DIR__, "piano.wav")
sound = Sound(file)
d = duration(sound)

function progress(p::Int, N = 100)
    str = ""
    for n in 1:N
        str *= n > p ?  "░" : "▓"
    end
    return str
end

interval = 0.5

try
    play!(sound)
    while isplaying(sound)
        t = timepos(sound)
        pre = ft(t) * " / " * ft(d)
        n = cols - length(pre) - 2
        p = floor(Int, n * t / d)
        clear()
        print(pre, " ", progress(p, n))
        sleep(interval)
    end
finally
    stop!(sound)
end


text:
@model function form_model(y :: Vector{Float64})
    N = length(y)
    s = Vector(undef, N)

    start ~ Uniform(-1, 1)

    s[1] ~ Normal(start, 0.05)
    y[1] ~ Normal(s[1], 0.5)

    for i ∈ 2:N
        s[i] ~ Normal(s[i-1], 0.05)
        y[i] ~ Normal(s[i], 0.5)
    end
end

function run_form_analysis(games :: Dict{String, Vector{Game}})
    results = Dict{String, Pair}()
    teams = [team for team ∈ keys(games)]
    pbar = Progress(length(teams))
    Threads.@threads for team ∈ teams
        N = 100
        g = NUTS()
        samples = sample(form_model(discounted_score.(games[team])), g, N);

        s = @pipe describe(group(samples, :s))[1][:, :mean] |> vec
        s_σ = @pipe describe(group(samples, :s))[1][:, :std] |> vec

        results[team] = Pair(s, s_σ)
        next!(pbar)
    end
    return results
end

max_stars_repo_name: shalinkpatel/football-form

max_issues_repo_name: shalinkpatel/football-form

max_forks_repo_name: shalinkpatel/football-form


text:
@testset "Scalar product" begin
    @test 1v1 ⦿ 2v3 == 0
    @test 1v1 ⦿ 2v1 == 2
    @test ((1v1 + 2v2) ⦿ (1v1 + 2v2)) == (triplet(SIGNATURE) == (1, 3, 0) ? -3 : 5)
end


text:
"""
    Base.show(io::IO, a::FlexibilityExpr)
Extend `Base.show` to print flexibility expressions.
"""
function Base.show(io::IO, a::FlexibilityExpr)
    # Setup empty case
    if length(a.vars) == 0
        return print(io, a.constant)
    end

    # Otherwise get data and parse it
    m = a.vars[1].m
    flex_data = getflexibilitydata(m)
    vars = a.vars
    names = []
    coeffs = []
    for i = 1:length(vars)
        push!(coeffs, a.coeffs[i])
        var = vars[i]
        if isa(var, RecourseVariable)
            push!(names, flex_data.recourse_names[var.idx])
        elseif isa(var, RandomVariable)
            push!(names, flex_data.RVnames[var.idx])
        end
    end

    # Print the expressions to the io
    strs = ["$(JuMP.aff_str(JuMP.REPLMode,coeffs[i], true))*$(names[i])" for i in 1:length(a.vars)]
    print(io,string(join(strs," + "), " + ", JuMP.aff_str(JuMP.REPLMode, a.constant, true)))
end

"""
    JuMP.addconstraint(m::Model, constr::FlexibilityConstraint)
Extend the `JuMP.addconstraint` function to handle `FlexibilityConstraint` types.
"""
function JuMP.addconstraint(m::Model, constr::FlexibilityConstraint)
    # Get the data
    flex_data = getflexibilitydata(m)
    push!(flex_data.flexibility_constraints, constr)

    # Add constraint to model here using the flexibility constraint information
    expr = constr.flex_expr
    sense = constr.sense

    #NOTE Look into why these coeffcients get conflated as expressions
    flex_coeffs = [expr.coeffs[i].constant for i = 1:length(expr.coeffs)]
    flex_vars = expr.vars

    # Parse the variable information
    flex_vars_jump = []
    for var in flex_vars
        if isa(var, RecourseVariable)
            jump_var = Variable(m, flex_data.recourse_cols[linearindex(var)])
            push!(flex_vars_jump, jump_var)
        elseif isa(var, RandomVariable)
            jump_var = Variable(m, flex_data.RVcols[linearindex(var)])
            push!(flex_vars_jump, jump_var)
        else
            error("Variable type $(typeof(var)) not recognized for constructing constraint")
        end
    end

    # Get non flexibility variable information
    state_coeffs = expr.constant.coeffs
    state_vars = expr.constant.vars
    constant = expr.constant.constant #should be the right hand side of the constraint

    # Add constraints to the model using the anonymous variables
    all_coeffs = vcat(flex_coeffs,state_coeffs)
    all_vars = vcat(flex_vars_jump,state_vars)
    if sense == :(<=)
        con_reference = @constraint(m, sum(all_coeffs[i]*all_vars[i] for i = 1:length(all_coeffs)) + constant <= 0)
    elseif sense == :(>=)
        con_reference = @constraint(m, sum(all_coeffs[i]*all_vars[i] for i = 1:length(all_coeffs)) + constant >= 0)
    elseif sense == :(==)
        con_reference = @constraint(m, sum(all_coeffs[i]*all_vars[i] for i = 1:length(all_coeffs)) + constant == 0)
    else
        error("Constraint sense $(sense) not recognized")
    end

    return ConstraintRef{JuMP.Model,FlexibilityConstraint}(m, length(flex_data.flexibility_constraints))
end

"""
    JuMP.show(io::IO,c::FlexibilityConstraint)
Extend the `JuMP.show` function to handle `FlexibilityConstraint` types.
"""
function JuMP.show(io::IO,c::FlexibilityConstraint)
    s = "$(string(c.flex_expr)) $(c.sense) 0"
    return print(io,s)
end


text:
include("ITI_additional.jl")

"""
Computes a global solution for a model via backward Improved Time Iteration. The algorithm is applied to the residuals of the arbitrage equations. The idea is to solve the system G(x) = 0 as a big nonlinear system in x, where the inverted Jacobian matrix is approximated by an infinite sum (Neumann series).

If the initial guess for the decision rule is not explicitly provided, the initial guess is provided by `ConstantDecisionRule`.
If the stochastic process for the model is not explicitly provided, the process is taken from the default provided by the model object, `model.exogenous`

# Arguments
* `model::NumericModel`: Model object that describes the current model environment.
* `dprocess`: The stochastic process associated with the exogenous variables in the model.
* `init_dr`: Initial guess for the decision rule.
* `maxbsteps` Maximum number of backsteps.
* `verbose` Set "true" if you would like to see the details of the infinite sum convergence.
* `smaxit` Maximum number of iterations to compute the Neumann series.
* `complementarities`
* `compute_radius`
* `trace` Record Iteration informations
# Returns
* `sol`: Improved Time Iteration results
"""
function improved_time_iteration(model::AbstractModel, dprocess::AbstractDiscretizedProcess,
                                 init_dr::AbstractDecisionRule, grid;
                                 maxbsteps::Int=10, verbose::Bool=true, verbose_jac::Bool=false,
                                 tol::Float64=1e-8, smaxit::Int=500, maxit::Int=1000,
                                 complementarities::Bool=true, compute_radius::Bool=false, trace::Bool=false,
                                 method=:gmres)

    parms = model.calibration[:parameters]

    n_m = max(n_nodes(dprocess), 1) # number of exo states today
    n_mt = n_inodes(dprocess,1)  # number of exo states tomorrow
    n_s = length(model.symbols[:states]) # number of endo states

    s = nodes(ListOfPoints, grid)
    N_s = length(s)
    n_x = size(model.calibration[:controls],1)

    x0 = [init_dr(i, s) for i=1:n_m]
    ddr = CachedDecisionRule(dprocess, grid, x0)
    ddr_filt = CachedDecisionRule(dprocess, grid, x0)
    set_values!(ddr,x0)

    steps = 0.5.^collect(0:maxbsteps)

    p = SVector(parms...)

    x = x0
    N = length(x[1])

    if complementarities == true
        x_lb = [controls_lb(model, node(Point,dprocess,i),s,p) for i=1:n_m]
        x_ub = [controls_ub(model, node(Point,dprocess,i),s,p) for i=1:n_m]
        BIG = 100000
        for i=1:n_m
          for n=1:N
            x_lb[i][n] = max.(x_lb[i][n],-BIG)
            x_ub[i][n] = min.(x_ub[i][n], BIG)
          end
        end
    end

    trace_data = []

    ######### Loop     for it in range(maxit):
    it=0
    it_invert=0

   err_0 = 1.0 #abs(maximum(res_init))
   err_2 = err_0

   lam0 = 0.0

   verbose && println(repeat("-", 120))
   verbose && println("N\tf_x\t\td_x\tTime_residuals\tTime_inversion\tTime_search\tLambda_0\tN_invert\tN_search\t")
   verbose && println(repeat("-", 120))

   while it <= maxit && err_0>tol

      it += 1

      t1 = time();

      # compute derivatives and residuals:
      # R_i: residuals
      # D_i: derivatives w.r.t. x
      # J_ij: derivatives w.r.t. ~x
      # S_ij: future states

      # set_values!(ddr,x)   # implicit in the next call

      _,J_ij,S_ij =   euler_residuals(model,s,x,ddr,dprocess,p,keep_J_S=true,set_dr=true)

      fun(u) = euler_residuals(model,s,u,ddr,dprocess,p,keep_J_S=false,set_dr=false)
      R_i, D_i = DiffFun(fun, x)

      if complementarities == true
          PhiPhi!(R_i,x,x_lb,x_ub,D_i,J_ij)
      end

      J_ij *= -1.0

      push!(trace_data, [deepcopy(R_i)])

      err_0 = maxabs((R_i))

      ####################
      # Invert Jacobians
      t2 = time();

      π_i, M_ij, S_ij = Dolo.preinvert!(R_i, D_i, J_ij, S_ij)

      if method==:gmres
        L = LinearThing(M_ij, S_ij, ddr_filt)
        v = cat([reshape(reinterpret(Float64, vec(e)), (n_x*N,)) for e in π_i]...; dims=1)
        n1 = L.counter
        w = gmres(L, v, verbose=false)
        it_invert = L.counter-n1
        ww = reshape(w,n_x,N,n_m)
        tt = [ww[:,:,i]  for i=1:n_m]
        tot = [reshape(reinterpret(SVector{n_x,Float64}, vec(t)), (N,)) for t in tt]
      else
        tot, it_invert, lam0, errors = invert_jac(π_i, M_ij, S_ij, ddr_filt; maxit=smaxit)
      end

      t3 = time();

      i_bckstps=0
      new_err=err_0
      new_x = x

      while new_err>=err_0 && i_bckstps<length(steps)
        i_bckstps +=1

        new_x = x-tot*steps[i_bckstps]

        new_res = euler_residuals(model,s,new_x,ddr,dprocess,p,keep_J_S=false,set_dr=true)

        if complementarities == true
            new_res = [PhiPhi0.(new_res[i],new_x[i],x_lb[i],x_ub[i]) for i=1:n_m]
        end

        new_err = maxabs(new_res)

      end

      err_2 = maxabs(tot)

      t4 = time();

      x = new_x
      verbose && @printf "%-6i% -10e% -17e% -15.4f% -15.4f% -15.5f% -17.3f%-17i%-5i\n" it  err_0  err_2  t2-t1 t3-t2 t4-t3 lam0 it_invert i_bckstps

   end

   verbose && println(repeat("-", 120))
   set_values!(ddr,x)

   if compute_radius == true
       lam, lam_max, lambdas = radius_jac(res,dres,jres,S_ij,ddr_filt)
   else
       lam = NaN
   end

   converged = err_0<tol

   return ImprovedTimeIterationResult(ddr.dr, it, err_0, err_2, converged, complementarities, tol, lam0, it_invert, 5.0, lam, trace_data)

end

function improved_time_iteration(model:: AbstractModel, dprocess::AbstractDiscretizedProcess,
                                 init_dr::AbstractDecisionRule;grid=Dict(), kwargs...)
    grid = get_grid(model, options=grid)
    return improved_time_iteration(model, dprocess, init_dr, grid;  kwargs...)
end

function improved_time_iteration(model, dprocess::AbstractDiscretizedProcess; grid=Dict(), kwargs...)

    init_dr = ConstantDecisionRule(model.calibration[:controls])
    return improved_time_iteration(model, dprocess, init_dr; grid=grid, kwargs...)
end

function  improved_time_iteration(model, init_dr; grid=Dict(), kwargs...)
    dprocess = discretize( model.exogenous )
    return improved_time_iteration(model, dprocess, init_dr; grid=grid, kwargs...)
end

function improved_time_iteration(model; grid=Dict(), kwargs...)
    dprocess = discretize( model.exogenous )
    init_dr = ConstantDecisionRule(model.calibration[:controls])
    return improved_time_iteration(model, dprocess, init_dr; grid=grid, kwargs...)
end


text:
module VennDiagrams

using Compose
using Color

function venn(xs::Union(AbstractArray,Set)...; proportional::Bool = true, labels=Union(Bool,Vector{String}), colors=Union(Bool,Vector{ColorValue},Vector{AlphaColorValue}))
    n = length(xs)

    if eltype(xs) != Set
        cols = map(Set, xs)
    else
        cols = xs
    end

    if proportional
        error("Not implemented yet!")
        #sizes = 0
    end

    return p
end

end # module


text:
@block SimonDanisch [layout] begin

    @cell "Layouting" [scatter, lines, surface, heatmap, vbox] begin
        p1 = scatter(rand(10), markersize = 1)
        p2 = lines(rand(10), rand(10))
        p3 = surface(0..1, 0..1, rand(100, 100))
        p4 = heatmap(rand(100, 100))
        x = 0:0.1:10
        p5 = lines(0:0.1:10, sin.(x))
        pscene = vbox(
            hbox(p1, p2),
            p3,
            hbox(p4, p5, sizes = [0.7, 0.3]),
            sizes = [0.2, 0.6, 0.2]
        )
    end

    @cell "Comparing contours, image, surfaces and heatmaps" [image, contour, surface, heatmap, vbox] begin
        N = 20
        x = LinRange(-0.3, 1, N)
        y = LinRange(-1, 0.5, N)
        z = x .* y'
        hbox(
            vbox(
                contour(x, y, z, levels = 20, linewidth =3),
                contour(x, y, z, levels = 0, linewidth = 0, fillrange = true),
                heatmap(x, y, z),
            ),
            vbox(
                image(x, y, z, colormap = :viridis),
                surface(x, y, fill(0f0, N, N), color = z, shading = false),
                image(-0.3..1, -1..0.5, AbstractPlotting.logo())
            )
        )
    end
end

@block JuliusKrumbiegel ["layout", "2d"] begin
    using MakieLayout

    @cell "Faceting" [faceting, grid] begin
        # layoutscene is a convenience function that creates a Scene and a GridLayout
        # that are already connected correctly and with Outside alignment
        scene, layout = layoutscene(30, resolution = (1200, 900))

        using ColorSchemes

        ncols = 4
        nrows = 4

        # create a grid of LAxis objects
        axes = [LAxis(scene) for i in 1:nrows, j in 1:ncols]
        # and place them into the layout
        layout[1:nrows, 1:ncols] = axes

        # link x and y axes of all LAxis objects
        linkxaxes!(axes...)
        linkyaxes!(axes...)

        lineplots = [lines!(axes[i, j], (1:0.1:8pi) .+ i, sin.(1:0.1:8pi) .+ j,
                color = get(ColorSchemes.rainbow, ((i - 1) * nrows + j) / (nrows * ncols)), linewidth = 4)
            for i in 1:nrows, j in 1:ncols]

        for i in 1:nrows, j in 1:ncols
            # remove unnecessary decorations in some of the facets, this will have an
            # effect on the layout as the freed up space will be used to make the axes
            # bigger
            i > 1 && (axes[i, j].titlevisible = false)
            j > 1 && (axes[i, j].ylabelvisible = false)
            j > 1 && (axes[i, j].yticklabelsvisible = false)
            j > 1 && (axes[i, j].yticksvisible = false)
            i < nrows && (axes[i, j].xticklabelsvisible = false)
            i < nrows && (axes[i, j].xticksvisible = false)
            i < nrows && (axes[i, j].xlabelvisible = false)
        end

        autolimits!(axes[1]) # hide

        legend = LLegend(scene, permutedims(lineplots, (2, 1)), ["Line $i" for i in 1:length(lineplots)],
            ncols = 2)
        # place a legend on the side by indexing into one column after the current last
        layout[:, end+1] = legend

        # index into the 0th row, thereby adding a new row into the layout and place
        # a text object across the first four columns as a super title
        layout[0, 1:4] = LText(scene, text="MakieLayout Facets", textsize=50)

        scene

    end

    @cell "ManualTicks" [layout, ticks] begin
        scene = Scene(resolution = (1000, 1000));
        campixel!(scene);

        maingl = GridLayout(scene, alignmode = Outside(30))

        la = maingl[1, 1] = LAxis(scene)
        la.attributes.yautolimitmargin = (0f0, 0.05f0)

        poly!(la, BBox(0, 1, 4, 0), color=:blue)
        poly!(la, BBox(1, 2, 7, 0), color=:red)
        poly!(la, BBox(2, 3, 1, 0), color=:green)

        la.attributes.xticks = ManualTicks([0.5, 1.5, 2.5], ["blue", "red", "green"])
        la.attributes.xlabel = "Color"
        la.attributes.ylabel = "Value"

        scene
    end

    @cell "Protrusion changes" [protrusions] begin
        using MakieLayout.Animations

        scene, layout = layoutscene(resolution = (600, 600))

        axes = [LAxis(scene) for i in 1:2, j in 1:2]
        layout[1:2, 1:2] = axes

        a_title = Animation([0, 2], [30.0, 50.0], sineio(n=2, yoyo=true, prewait=0.2))
        a_xlabel = Animation([2, 4], [20.0, 40.0], sineio(n=2, yoyo=true, prewait=0.2))
        a_ylabel = Animation([4, 6], [20.0, 40.0], sineio(n=2, yoyo=true, prewait=0.2))

        record(scene, @replace_with_a_path(mp4), 0:1/60:6, framerate = 60) do t
            axes[1, 1].titlesize = a_title(t)
            axes[1, 1].xlabelsize = a_xlabel(t)
            axes[1, 1].ylabelsize = a_ylabel(t)
        end
    end

    @cell "Hiding decorations" [decorations] begin
        scene = Scene(resolution = (600, 600), camera=campixel!)

        layout = GridLayout(
            scene, 2, 2, # we need to specify rows and columns so the gap sizes don't get lost
            addedcolgaps = Fixed(0),
            addedrowgaps = Fixed(0),
            alignmode = Outside(30)
        )

        axes = [LAxis(scene) for j in 1:2, i in 1:2]
        layout[1:2, 1:2] = axes

        re = record(scene, @replace_with_a_path(mp4), framerate=3) do io
            recordframe!(io)
            for ax in axes
                ax.titlevisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.xlabelvisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.ylabelvisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.xticklabelsvisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.yticklabelsvisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.xticksvisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.yticksvisible = false
                recordframe!(io)
            end
            for ax in axes
                ax.bottomspinevisible = false
                ax.leftspinevisible = false
                ax.topspinevisible = false
                ax.rightspinevisible = false
                recordframe!(io)
            end
        end
        return re
    end

    @cell "Axis aspects" [axis, aspect] begin
        using FileIO
        scene, layout = layoutscene(30, resolution = (1200, 900))

        axes = [LAxis(scene) for i in 1:2, j in 1:3]
        tightlimits!.(axes)
        layout[1:2, 1:3] = axes

        img = rotr90(MakieGallery.loadasset("cow.png"))

        for ax in axes
            image!(ax, img)
        end

        axes[1, 1].title = "Default"

        axes[1, 2].title = "DataAspect"
        axes[1, 2].aspect = DataAspect()

        axes[1, 3].title = "AxisAspect(418/348)"
        axes[1, 3].aspect = AxisAspect(418/348)

        axes[2, 1].title = "AxisAspect(1)"
        axes[2, 1].aspect = AxisAspect(1)

        axes[2, 2].title = "AxisAspect(2)"
        axes[2, 2].aspect = AxisAspect(2)

        axes[2, 3].title = "AxisAspect(0.5)"
        axes[2, 3].aspect = AxisAspect(0.5)

        scene
    end

    @cell "Linked axes" [axis, link] begin
        scene, layout = layoutscene()

        layout[1, 1:3] = axs = [LAxis(scene) for i in 1:3]
        linkxaxes!(axs[1:2]...)
        linkyaxes!(axs[2:3]...)

        axs[1].title = "x linked"
        axs[2].title = "x & y linked"
        axs[3].title = "y linked"

        for i in 1:3
            lines!(axs[i], 1:10, 1:10, color = "green")
            if i != 1
                lines!(axs[i], 1:10, 11:20, color = "blue")
            end
            if i != 3
                lines!(axs[i], 11:20, 1:10, color = "red")
            end
        end

        scene
    end

    @cell "Nested grids" [grid] begin
        scene, layout = layoutscene(30, resolution = (1200, 900))

        subgl_left = GridLayout()
        subgl_left[1:2, 1:2] = [LAxis(scene) for i in 1:2, j in 1:2]

        subgl_right = GridLayout()
        subgl_right[1:3, 1] = [LAxis(scene) for i in 1:3]

        layout[1, 1] = subgl_left
        layout[1, 2] = subgl_right

        scene
    end

    @cell "Grid alignment" [grid] begin
        scene, layout = layoutscene(30, resolution = (1200, 1200))

        layout[1, 1] = LAxis(scene, title="No grid layout")
        layout[2, 1] = LAxis(scene, title="No grid layout")
        layout[3, 1] = LAxis(scene, title="No grid layout")

        subgl_1 = layout[1, 2] = GridLayout(alignmode=Inside())
        subgl_2 = layout[2, 2] = GridLayout(alignmode=Outside())
        subgl_3 = layout[3, 2] = GridLayout(alignmode=Outside(50))

        subgl_1[1, 1] = LAxis(scene, title="Inside")
        subgl_2[1, 1] = LAxis(scene, title="Outside")
        subgl_3[1, 1] = LAxis(scene, title="Outside(50)")

        layout[1:3, 2] = [LRect(scene, color = :transparent, strokecolor = :red) for i in 1:3]

        scene
    end

    @cell "Spanned grid content" [grid] begin

        scene, layout = layoutscene(4, 4, 30, resolution = (1200, 1200))

        layout[1, 1:2] = LAxis(scene, title="[1, 1:2]")
        layout[2:4, 1:2] = LAxis(scene, title="[2:4, 1:2]")
        layout[:, 3] = LAxis(scene, title="[:, 3]")
        layout[1:3, end] = LAxis(scene, title="[1:3, end]")
        layout[end, end] = LAxis(scene, title="[end, end]")

        scene
    end

    @cell "Indexing outside grid" [grid] begin
        scene, layout = layoutscene(30, resolution = (1200, 1200))

        layout[1, 1] = LAxis(scene)
        for i in 1:3
            layout[:, end+1] = LAxis(scene)
            layout[end+1, :] = LAxis(scene)
        end

        layout[0, :] = LText(scene, text="Super Title", textsize=50)
        layout[end+1, :] = LText(scene, text="Sub Title", textsize=50)
        layout[2:end-1, 0] = LText(scene, text="Left Text", textsize=50,
        rotation=pi/2)
        layout[2:end-1, end+1] = LText(scene, text="Right Text", textsize=50,
        rotation=-pi/2)

        scene
    end

    @cell "Column and row sizes" [grid] begin
        scene = Scene(resolution = (1200, 900), camera=campixel!)

        layout = GridLayout(
        scene, 6, 6,
        colsizes = [Fixed(200), Relative(0.25), Auto(), Auto(), Auto(2), Auto()],
        rowsizes = [Auto(), Fixed(100), Relative(0.25), Aspect(2, 1), Auto(), Auto()],
        alignmode = Outside(30, 30, 30, 30))

        for i in 2:6, j in 1:5
            if i == 6 && j == 3
                layout[i, j] = LText(scene, text="My Size is Inferred")
            else
                layout[i, j] = LRect(scene)
            end
        end

        for (j, label) in enumerate(["Fixed(200)", "Relative(0.25)", "Auto()", "Auto()", "Auto(2)"])
            layout[1, j] = LText(scene, width = Auto(), text = label,
            tellwidth = false
            )
        end

        for (i, label) in enumerate(["Fixed(100)", "Relative(0.25)", "Aspect(2, 1)", "Auto()", "Auto()"])
            layout[i + 1, 6] = LText(scene, height = Auto(), text = label, tellwidth = false)
        end

        scene
    end

    @cell "Trimming grids" [grid] begin

        scene, layout = layoutscene(resolution = (600, 600))

        record(scene, @replace_with_a_path(mp4), framerate=1) do io

            ax1 = layout[1, 1] = LAxis(scene, title = "Axis 1")
            recordframe!(io)

            ax2 = layout[1, 2] = LAxis(scene, title = "Axis 2")
            recordframe!(io)

            layout[2, 1] = ax2
            recordframe!(io)

            trim!(layout)
            recordframe!(io)

            layout[2, 3:4] = ax1
            recordframe!(io)

            trim!(layout)
            recordframe!(io)
        end
    end

    @cell "3D scenes" ["3d"] begin
        scene, layout = layoutscene()
        makescene() = LScene(scene, camera = cam3d!, raw = false,
            scenekw = (backgroundcolor = RGBf0(0.9, 0.9, 0.9), clear = true))
        layout[1, 1]   = makescene()
        layout[1, 2]   = makescene()
        layout[2, 1:2] = makescene()
        layout[1:2, 3] = makescene()
        foreach(LScene, layout) do s
            scatter!(s, rand(100, 3));
        end
        scene
    end

    @cell "Window resizing" [record] begin
        using MakieLayout.Animations

        container_scene = Scene(camera = campixel!, resolution = (1200, 1200))

        t = Node(0.0)

        a_width = Animation([1.0, 7.0], [1200.0, 800.0], sineio(n=2, yoyo=true, postwait=0.5))
        a_height = Animation([2.5, 8.5], [1200.0, 800.0], sineio(n=2, yoyo=true, postwait=0.5))
        scene_area = lift(t) do t
            IRect(0, 0, a_width(t), a_height(t))
        end

        scene = Scene(container_scene, scene_area, camera = campixel!)

        rect = poly!(scene, scene_area,
            raw=true, color=RGBf0(0.97, 0.97, 0.97), strokecolor=:transparent, strokewidth=0)[end]

        outer_gl = GridLayout(scene, alignmode = Outside(30))

        inner_gl = outer_gl[1, 1] = GridLayout()

        ax1 = inner_gl[1, 1] = LAxis(scene, xautolimitmargin=(0, 0), yautolimitmargin=(0, 0))
        ax2 = inner_gl[1, 2] = LAxis(scene, xautolimitmargin=(0, 0), yautolimitmargin=(0, 0))
        ax3 = inner_gl[2, 1:2] = LAxis(scene, xautolimitmargin=(0, 0), yautolimitmargin=(0, 0))

        guigl = inner_gl[3, 1:2] = GridLayout()
        b1 = guigl[1, 1] = LButton(scene, label = "prev", width = Auto())
        sl = guigl[1, 2] = LSlider(scene, startvalue = 6, height = 40)
        b2 = guigl[1, 3] = LButton(scene, label = "next", width = Auto())

        data = randn(200, 200) .+ 3 .* sin.((1:200) ./ 20) .* sin.((1:200)' ./ 20)
        h1 = heatmap!(ax1, data)
        h2 = heatmap!(ax2, data, colormap = :blues)
        h3 = heatmap!(ax3, data, colormap = :heat)

        agl1 = gridnest!(inner_gl, 1, 1)
        agl1[1, 2] = LColorbar(scene, h1, width = 30, label = "normal bar")
        agl1[2, 1:2] = LSlider(scene, height = 20, startvalue = 4)
        agl1[3, 1:2] = LSlider(scene, height = 20, startvalue = 5)
        agl1[4, 1:2] = LSlider(scene, height = 20, startvalue = 6)

        agl2 = gridnest!(inner_gl, 1, 2)
        agl2[1, 2] = LColorbar(scene, h2, width = 30, height = Relative(0.66), label = "two thirds bar")
        agl2gl = agl2[2, :] = GridLayout()
        agl2gl[1, 1] = LButton(scene, label = "Run", height = Auto())
        agl2gl[1, 2] = LButton(scene, label = "Start")

        agl3 = gridnest!(inner_gl, 2, 1:2)
        agl3[:, 3] = LColorbar(scene, h3, width = 30, height=200, label = "fixed height bar")
        rowsize!(agl3, 1, Auto(1.0))

        inner_gl[0, :] = LText(scene, text = "MakieLayout", textsize = 50)

        record(container_scene, @replace_with_a_path(mp4), 0:1/30:9; framerate=30) do ti
        t[] = ti
        end

    end

    @cell "Aspect ratios stretching circles" [layout] begin
        using Animations

        # scene setup for animation

        container_scene = Scene(camera = campixel!, resolution = (1200, 1200))

        t = Node(0.0)

        a_width = Animation([1, 7], [1200.0, 800], sineio(n=2, yoyo=true, postwait=0.5))
        a_height = Animation([2.5, 8.5], [1200.0, 800], sineio(n=2, yoyo=true, postwait=0.5))

        scene_area = lift(t) do t
        IRect(0, 0, round(Int, a_width(t)), round(Int, a_height(t)))
        end

        scene = Scene(container_scene, scene_area, camera = campixel!)

        rect = poly!(scene, scene_area,
        raw=true, color=RGBf0(0.97, 0.97, 0.97), strokecolor=:transparent, strokewidth=0)[end]

        outer_layout = GridLayout(scene, alignmode = Outside(30))

        # example begins here

        layout = outer_layout[1, 1] = GridLayout()

        titles = ["aspect via layout", "axis aspect", "no aspect", "data aspect"]
        axs = layout[1:2, 1:2] = [LAxis(scene, title = t) for t in titles]

        # plot to all axes using an array comprehension
        [lines!(a, Circle(Point2f0(0, 0), 100f0)) for a in axs]

        rowsize!(layout, 1, Fixed(400))
        # force the layout cell [1, 1] to be square
        colsize!(layout, 1, Aspect(1, 1))

        axs[2].aspect = 1
        axs[4].autolimitaspect = 1

        rects = layout[1:2, 1:2] = [LRect(scene, color = (:black, 0.05),
        strokecolor = :transparent) for _ in 1:4]

        record(container_scene, @replace_with_a_path(mp4), 0:1/60:9; framerate=60) do ti
            t[] = ti
        end

    end

    @cell "Ellipse markers and arrows" [layout, "2d"] begin

        # here we plot some grouped coordinates, find their mean and variance which
        # we plot as ellipses, plot their actual means, and some additional details.

        using DataFrames, Distributions, Colors
        using MakieLayout
        import AbstractPlotting:px

        # Here we create `ngrp` groups of `n` coordinates. The mean coordinates of
        # the groups are distributed equally around the unit circle. The coordinates
        # are drawn from a normal distribution (with standard deviation σ).
        ngrp = 4
        groups = Symbol.(Iterators.take('a':'z', ngrp))
        n = 5
        σ = [0.2, 0.4]
        intended = (; Dict(g => Point2f0(Iterators.reverse(sincos(θ))...) for (g, θ) in zip(groups, range(0, step = 2π/ngrp, length = ngrp)))...)
        df = DataFrame(group = Symbol[], x = Point2f0[])
        for g in groups
            d = MvNormal(intended[g], σ)
            for _ in 1:n
                push!(df, (group = g, x = rand(d)))
            end
        end

        # assign colors to each group
        colors = Dict(zip(groups, distinguishable_colors(ngrp, [colorant"white", colorant"black"], dropseed = true)))

        # calculate the mean of each group, and the FWHM of a Gaussian fit to the
        # data. We later use the mean and FWHM to plot ellipses.
        ellipses = by(df, :group) do g
            n = length(g.x)
            X = Array{Float64}(undef, 2, n)
            for i in 1:n
                X[:,i] = g.x[i]
            end
            dis = fit(DiagNormal, X)
            radii = sqrt(2log(2))*sqrt.(var(dis)) # half the FWHM
            ellipse = (origin = Point2f0(mean(dis)), radius = Vec2f0(radii))
            (ellipse = ellipse, )
        end

        # some helper functions

        # return a vector of coordinates of an ellipse
        mydecompose(origin, radii) = [origin + radii .* Iterators.reverse(sincos(t)) for t in range(0, stop = 2π, length = 51)]
        # brighten a color
        brighten(c, p = 0.5) = weighted_color_mean(p, c, colorant"white")
        # darken a color
        darken(c, p = 0.5) = weighted_color_mean(p, c, colorant"black")

        scene, layout = layoutscene(0, fontsize = 10, resolution = (500,400));

        ax = layout[1,1] = LAxis(scene,
                                 xlabel = "X (cm)",
                                 ylabel = "Y (cm)",
                                 xticklabelsize = 8,
                                 yticklabelsize = 8,
                                 aspect = DataAspect()
                                )

        for r in eachrow(ellipses)
            c = r.ellipse
            xy = mydecompose(c...)
            poly!(ax, xy, color = brighten(colors[r.group], 0.25))
        end

        scatter!(ax, getfield.(ellipses.ellipse, :origin), color = :white, marker = '+', markersize = 10px)
        scatter!(ax, [zero(Point2f0)], color = :black, strokecolor = :black, marker = '⋆', strokewidth = 0.5, markersize = 15px)
        scatter!(ax, [intended[k] for k in keys(colors)], color = :white, strokecolor = collect(values(colors)), marker = '⋆', strokewidth = 0.5, markersize = 15px)

        for g in groupby(df, :group)
            scatter!(ax, g.x, color = RGBA(colors[g.group[1]], 0.75), marker = '●', markersize = 5px)
        end

        # Here, we manually construct a list of legend entries to be provided to the
        # LLegend constructor. This allows us a larger degree of control over how
        # the legend looks.
        polys = [PolyElement(color = colors[k], strokecolor = :transparent) for k in unique(df.group)]
        shapes = [MarkerElement(color = :black, marker = '⋆', strokecolor = :black, markerstrokewidth = 0.5, markersize = 15px),
                  MarkerElement(color = :white, marker = '⋆', strokecolor = :black, markerstrokewidth = 0.5, markersize = 15px),
                  MarkerElement(color = :black, marker = '●', strokecolor = :transparent, markersize = 5px),
                  [PolyElement(color = brighten(colorant"black", 0.75), strokecolor = :transparent, polypoints = mydecompose(Point2f0(0.5, 0.5), Vec2f0(0.75, 0.5))),
                   MarkerElement(color = :white, marker = '+', strokecolor = :transparent, markersize = 10px),
                  ]]

        leg = ([polys, shapes], [string.(unique(df.group)), ["center", "intended means", "coordinates", "μ ± FWHM"]], ["Groups", "Shapes"])
        layout[1, 2] = LLegend(scene, leg..., markersize = 10px, markerstrokewidth = 1, patchsize = (10, 10), rowgap = Fixed(0), titlegap = Fixed(5), groupgap = Fixed(10), titlehalign = :left, gridshalign = :left)

        # draw an arrow with some text
        textlayer = Scene(scene, ax.scene.px_area, camera = campixel!, raw = true)
        topright = lift(textlayer.px_area) do w
            xy = widths(w)
            xy .- 0.05max(xy...)
        end
        text!(textlayer, "arrow", position = topright, align = (:right, :top), textsize = 10, font = "noto sans")
        lines!(textlayer, @lift([$topright, $topright .- (30, 0)]))
        scatter!(textlayer, @lift([$topright]), marker = '►', markersize = 10px)

        scene

    end
end


text:
module SalsaML

using Salsa

const N_FEATURES = 1
const Sample = NTuple{N_FEATURES, Number}
const learning_rate = 0.0005

@declare_input num_samples(s)::Int
@declare_input sample(s, i::Int)::Sample
@declare_input response(s, i::Int)::Float64

init_lr!(s::Runtime) = set_num_samples!(s, 0)
function new_lr_runtime()
    s = Runtime()
    init_lr!(s)
    return s
end

function insert_training_pair!(s, sample::Sample, response::Float64)
    i = num_samples(s) + 1
    set_num_samples!(s, i)
    set_sample!(s, i, sample)
    set_response!(s, i, response)
end

function predict(s::Runtime, x::Sample)::Float64
    sum(x .* lr_learned_weights(s))
end

# Derived function with no inputs (essentially a constant)
@derived init_weights(s::Runtime) = Tuple(rand(Float64, N_FEATURES))

@derived function lr_learned_weights(s::Runtime)::NTuple{N_FEATURES,Float64}
    if num_samples(s) == 0
        error("Cannot learn a linear regression with no samples.")
    end
    weights = lr_learned_weights_unrolled(s, 0)
    for i in 1:100  # Stop after 100 iterations
        new_weights = lr_learned_weights_unrolled(s, i)
        if abs(sum(new_weights .- weights)) < 0.0001
            @info "stopped after $i iterations"
            return weights
        end
        weights = new_weights
    end
    @info "timed out in learn iterations"
    return weights
end
@derived function lr_learned_weights_unrolled(s::Runtime, iteration::Int)

    if iteration == 0
        return init_weights(s)
    else
        n = num_samples(s)
        weights = lr_learned_weights_unrolled(s, iteration-1)
        δweights_δcost = lr_δweights_δmse(s, iteration-1) .* learning_rate
        weights = weights .- δweights_δcost
        return weights
    end
end

@derived function lr_δweights_δmse(s::Runtime, iteration::Int)
    n = num_samples(s)
    (-2/n) * sum(
        sum(sample(s, i) .* (response(s, i) - lr_predicted_unrolled(s, i, iteration)))
        for i in 1:n
    )
end
@derived function lr_predicted_unrolled(s::Runtime, i::Int, iteration::Int)
    sum(sample(s, i) .* lr_learned_weights_unrolled(s, iteration))
end

# For logging purposes
@derived function lr_mse(s::Runtime)
    n = num_samples(s)
    1/n * sum(
        (response(s, i) - lr_predicted(s, i)) ^ 2
        for i in 1:num_samples(s)
    )
end

end


text:
using PowerDynamics: simulate, Perturbation, Inc, PowerGrid
using Distributions: Uniform
using Sundials
using Statistics
using NetworkDynamics
using LaTeXStrings
using FileIO, JLD2
using DelimitedFiles

include("../new node types/plotting.jl")
include("../new node types/ThirdOrderEq.jl")
include("../../src/Random-Pertubation-on-Constrained-Manifolds/RandPertOnConstraindManifolds.jl")
include("../../src/PDpatches.jl")

"""
Evaluates the result of a simulation after a pertubation.
Checks if conditions for a stable synchronous state where met.
Conditions implemented so far:
1. Is ω returning to zero? (Realized through checking the deviations from 0)
#2. Did one of the machines get too far away from a stable state? -> Survivability
Inputs:
       pg: Power grid, a graph containg nodes and lines
       result: Solution of a power grid after a pertubation
       nodesarray: Array containg the index of all nodes that contain varible
       endtime: Endtime of the simulation
Outputs:
       Returns 1 if the system fullfills the condition for a stable state
       Returns 0 if the condition was not met
"""
function StableState(result::PowerGridSolution,nodesarray,variable,endtime)
        #=
        maxdeviation = maximum(result(0:endtime,node,:ω))
        if abs(maxdeviation) > threshold, (5)   # this could be used for Survivability later on...
            return 0
        end
        =#
    for node in  1:length(powergrid.nodes) # runs over all nodes
        summation = 0
        for time in endtime-10:0.1:endtime
            if node in nodesarray  # checks only nodes which have ω as a symbol
                summation =+ abs.(result(time,node,:ω))
            end
            if summation > 100  # checks if ω returns to 0
                                #this is just an arbitrary value that i have choosen so far
                return 0
            end
            if isapprox(result(time,node, :v) , 0, atol=1e-2) # checks for a short cirtuit
                return 0
            end
        end
    end
    return 1
end

function sim(pg::PowerGrid, x0::State, timespan, force_dtmin::Bool, rpg)
    problem = ODEProblem(rpg,x0.vec,timespan)
    # sets the kwarg force_dtmin to avoid the dt < dtmin error
    if force_dtmin == true
        solution = solve(problem, Rodas5(autodiff=false), abstol=1e-8,reltol=1e-6,tspan=Inf, force_dtmin = true)
    else
        solution = solve(problem, Rodas5(autodiff=false), abstol=1e-8,reltol=1e-6,tspan=Inf)
    end
    PowerGridSolution(solution, pg)
end

function dae_sim(rpg_ode, tspan, x0::State,pg)
        function dae_form(res, du, u, p, t)
            du_temp = similar(du)
            rpg_ode.f(du_temp, u, p, t)
            @. res = du - du_temp
        end
        u0 = x0.vec
        du0 = similar(u0)
        p = nothing
        diff_vars = (diag(rpg_ode.mass_matrix) .== 1)
        rpg_dae = DAEFunction{true, true}(dae_form, syms = rpg_ode.syms)
        dae = DAEProblem(rpg_dae, du0, u0, tspan, p; differential_vars=diff_vars)
        sol = solve(dae, DABDF2(autodiff=false), initializealg=BrownFullBasicInit(), force_dtmin = true, matiters= 1e7,save_everystep=false)
        #sol = solve(dae, IDA(linear_solver=:GMRES))
        return PowerGridSolution(sol,pg)
end

"""
Removes all nodes from the from the calculation procedure that do not contain
variable as a symbol.
Returns a list containg the indexes of all the remaning nodes.
Inputs:
    pg: Power grid, a graph containg nodes and lines
Outputs:
    nodesarray: Array containg the index of all nodes with the symbol varibale
"""
function RemoveNodes(pg::PowerGrid, variable::Symbol)
    nodesarray = findall(variable.∈ symbolsof.(pg.nodes))
    return nodesarray
end

"""
Produces a plot of the calulated basin stability against the nodes index.
Inputs:
    nodesarray: Array containg the index of all nodes that contain varible
    basinstability: Array containing the basin stability of each node
    standarddev: Array of the standard deviations of the basin stability
"""
function PlotBasinStability(nodesarray, basinstability, standarddev; labtext = "Basin Stability")
    plot(
        nodesarray,
        basinstability,
        xlabel = "Node Index",
        ylabel = "Basin Stability",
        marker = (:circle, 8, "red"),
        line = (:path, 2,"gray"),
        label = labtext,
        grid = false,
        show = true
    )
end

function PlotBasinStability!(nodesarray, basinstability, standarddev; labtext = "Basin Stability")
    plot!(
        nodesarray,
        basinstability,
        xlabel = "Node Index",
        ylabel = "Basin Stability",
        marker = (:circle, 8, "blue"),
        line = (:path, 2,"gray"),
        label = labtext,
        grid = false,
        show = true
    )
end

"""
Calculates the single node basin stability of the nodes in a power grid using
PowerDynamics.jl.
The calculations were performed on the basis of the three following papers:
1. Menck, P., Heitzig, J., Marwan, N. et al. How basin stability complements
   the linear-stability paradigm. Nature Phys 9, 89–92 (2013)
2. Menck, P., Heitzig, J., Kurths, J. et al. How dead ends undermine
   power grid stability. Nat Commun 5, 3969 (2014).
3. C Mitra, A Choudhary, S Sinha, J Kurths, and R V Donnerk,
   Multiple-node basin stability in complex dynamical networks
   Physical Review E, (2017)

To calculate the basin stability "numtries" random pertubations are
drawn according to the interval.
The behaviour of the power grid after the pertubation is then calculated.
In the function StableState the number of times the system returns
to the stable synchronous state is counted.
If numtries -> ∞: Then BasinStability ≈ stablecounter / numtries.
Inputs:
    pg: Power grid, a graph containg nodes and lines
    endtime: Endtime of the simulation
    variable: Internal variable of the node types to be perturbed
    numtries: Number of random pertubations within interval to be evaluated
    interval: Lower and Upper Bound for the pertubation
Outputs:
    basinstability: Array conating the basinstability of each node
    standarddev: Array of the standard deviations of the basin stability
"""
function BasinStability(pg::PowerGrid, endtime::Int, variable::Symbol,
                        numtries::Int, interval_v::Vector{Float64}, interval_θ::Vector{Float64};
                        ProjectionMethod = "Optim")
    @assert interval_v[2] >= interval_v[1] "Upper bound should be bigger than lower bound!"
    num_errors = 0
    operationpoint = find_steady_state(pg)
    nodesarray = RemoveNodes(pg,variable)
    basinstability = []
    standarddev =  []
    P, Q = SimplePowerFlow(pg, operationpoint)
    rpg = rhs(pg)

    state_file = File(format"JLD2","State_File.jld2")
    save(state_file, "operationpoint", operationpoint)
    for node in nodesarray
        stablecounter = Array{Int64}(undef,numtries)
        # do some parralelization\ multithreading here to speed up calculation?
        for tries = 1:numtries
            println("NODE:",node," TRY:",tries)
            PerturbedState, dz = RandPertWithConstrains(powergrid, operationpoint,node, interval_v, interval_θ, Q[node], ProjectionMethod = ProjectionMethod)
            column_name = string(node) * string(tries)
            save(state_file, column_name, PerturbedState)
            λ, stable = check_eigenvalues(powergrid, PerturbedState)
            #=
            if stable != true
                println("Positive eigenvalue detected. The system might be unstable.")
                try
                    result = sim(pg, PerturbedState, (0.0, endtime),true, rpg)
                    stablecounter[tries] = StableState(result, nodesarray, variable, endtime, interval)
                    if stablecounter[tries] == 0
                        display(plot_res(result, pg, node))
                        fn = "node$(node)_try$(tries)"
                        png(fn)
                    end
                catch err
                    if isa(err, GridSolutionError)
                        stablecounter[tries] = 0 # counted as an unstable solution
                        num_errors += 1
                    end
                end
            else
            =#
            result = sim(pg, PerturbedState, (0.0, endtime),false, rpg)
            stablecounter[tries] = StableState(result,nodesarray,variable,endtime)
            #if stablecounter[tries] == 0
            display(plot_res(result, pg, node))
                #fn = "node$(node)_try$(tries)"
                #png(fn)
            #end
        end
        append!(basinstability, mean(stablecounter))
        append!(standarddev, std(stablecounter)) # is this even a useful measure here?
    end
    println("Percentage of failed Simualtions:  ",100 * num_errors/(numtries * length(nodesarray)))
    display(PlotBasinStability(nodesarray, basinstability, standarddev))
    png("BasinStability_Plot")
    writedlm("bs.txt", basinstability)
    return basinstability, standarddev
end

function BasinStability(pg::PowerGrid, endtime::Float64, variable::Symbol,
                        numtries::Int, interval_v::Vector{Float64}, interval_θ::Vector{Float64};
                        ProjectionMethod = "Optim",
                        dae)
    @assert interval_v[2] >= interval_v[1] "Upper bound should be bigger than lower bound!"
    num_errors = 0
    operationpoint = find_steady_state(pg)
    nodesarray = RemoveNodes(pg,variable)
    basinstability = []
    standarddev =  []
    P, Q = SimplePowerFlow(pg, operationpoint)
    rpg = rhs(pg)
    state_file = File(format"JLD2","State_File.jld2")
    save(state_file, "operationpoint", operationpoint)
    for node in nodesarray
        stablecounter = Array{Int64}(undef,numtries)
        # do some parralelization\ multithreading here to speed up calculation?
        for tries = 1:numtries
            println("NODE:",node," TRY:",tries)
            PerturbedState, dz = RandPertWithConstrains(powergrid, operationpoint,node, interval_v, interval_θ, Q[node], ProjectionMethod = ProjectionMethod)
            column_name = string(node) * string(tries)
            save(state_file, column_name, PerturbedState)
            result = dae_sim(rpg, (0.0, endtime), PerturbedState,pg)
            stablecounter[tries] = StableState(result,nodesarray,variable,endtime)
            display(plot_res(result, pg, node))
            fn = "node$(node)_try$(tries)"
            png(fn)
        end
        append!(basinstability, mean(stablecounter))
        append!(standarddev, std(stablecounter)) # is this even a useful measure here?
    end
    println("Percentage of failed Simualtions:  ",100 * num_errors/(numtries * length(nodesarray)))
    display(PlotBasinStability(nodesarray, basinstability, standarddev))
    png("BasinStability_Plot")
    writedlm("bs.txt", basinstability)
    return basinstability, standarddev
end

max_stars_repo_name: luap-pik/PowerDynIC

max_issues_repo_name: luap-pik/PowerDynIC

max_forks_repo_name: luap-pik/PowerDynIC