text:
using Pkg
Pkg.activate(".")

using PowerDynamics
using OrdinaryDiffEq
using CSV
using Plots
plotly()

##

dir = @__DIR__

## set per unit MV
const base_power = 1E6 # 1MW
const base_voltage = 20E3 # 20kV
const base_current = base_power / base_voltage # 50A
const base_admittance = base_power / base_voltage^2 # 0.0025Ω^-1
const ω = 2 * π * 50.0 # 314.1593rad/s
# per unit HV
const base_voltage_HV = 110E3 # 110kV
const base_admittance_HV = base_power / base_voltage_HV^2 # 8.264462809917356e-5
const base_current_HV = base_power / base_voltage_HV

## define OLTC

include("$dir/OLTC.jl")

## define DGUnit

import Base: @__doc__
using PowerDynamics: @DynamicNode
import PowerDynamics:
    construct_vertex, dimension, symbolsof, showdefinition, AbstractNode
using LinearAlgebra: Diagonal
using NetworkDynamics: ODEVertex

function Limit(x, l, u)
    if x < l
        return l
    elseif x > u
        return u
    else
        return x
    end
end

begin
    @__doc__ struct StaticGeneratorObs <: AbstractNode
        K_pll
        S_ref
        K_FRT
        I_max
        Vref
        Vdead
        Y_n
    end
    StaticGeneratorObs(; K_pll, S_ref, K_FRT, I_max, Vref, Vdead, Y_n) =
        StaticGeneratorObs(K_pll, S_ref, K_FRT, I_max, Vref, Vdead, Y_n)
    function construct_vertex(par::StaticGeneratorObs)
        K_pll = par.K_pll
        S_ref = par.S_ref
        K_FRT = par.K_FRT
        I_max = par.I_max
        Vref = par.Vref
        Vdead = par.Vdead
        Y_n = par.Y_n
        function rhs!(dx, x, e_s, e_d, p, t)
            u = complex(x[1], x[2])
            # current from incoming/outgoing edges, nodal shunt, PQ load background
            i = total_current(e_s, e_d) + Y_n(t) * u
            Vamp = abs(u)

            # this is the only dynamic variable
            θ = x[3]

            # to keep track of the internal currents for debugging ...
            i_d_r = x[4] # behind FRT
            i_q_r = x[5] # behind FRT
            i_d_r_b = x[6] # before FRT
            i_q_r_b = x[7] # before FRT
            iqp = x[8]

            # PLL
            dθ = -K_pll * Vamp * sin(θ - angle(u))

            # FRT
            # additional reactive current during fault
            if Vamp > Vref + Vdead
                iqplus = K_FRT * (Vamp - (Vref + Vdead)) # pos for over_voltage
            elseif Vamp < Vref - Vdead
                iqplus = K_FRT * (Vamp - (Vref - Vdead)) # neg for under_voltage
            else
                iqplus = 0.0
            end

            # current set point
            I_r = conj(S_ref) / Vamp

            # save values for plotting
            di_d_r_b = i_d_r_b - real(I_r)
            di_q_r_b = i_q_r_b - imag(I_r)
            diqp = iqp - iqplus

            # FRT limiter
            if Vamp > Vref + Vdead || Vamp < Vref - Vdead
                iq_fault = Limit(imag(I_r) + iqplus, -I_max, I_max)
                id_lim = sqrt(I_max^2 - iq_fault^2)
                id_fault = Limit(real(I_r), 0.0, id_lim) #-id_lim, id_lim)
                I_r_limited = complex(id_fault, iq_fault)
            else
                id_normal = Limit(real(I_r), 0.0, I_max) #-I_max, I_max)
                iq_lim = sqrt(I_max^2 - id_normal^2)
                iq_normal = Limit(imag(I_r), -iq_lim, iq_lim)
                I_r_limited = complex(id_normal, iq_normal)
            end

            # save values for plotting
            di_d_r = i_d_r - real(I_r_limited)
            di_q_r = i_q_r - imag(I_r_limited)

            du = i - exp(θ * 1im) * I_r_limited # = 0 constraint

            try
                dx[1] = real(du)
                dx[2] = imag(du)
                dx[3] = dθ
                dx[4] = di_d_r
                dx[5] = di_q_r
                dx[6] = di_d_r_b
                dx[7] = di_q_r_b
                dx[8] = diqp
                return nothing
            catch e
                if typeof(e) === UndefVarError
                    throw(NodeDynamicsError("you need to provide $(e.var)"))
                else
                    throw(e)
                end
            end
        end
        ODEVertex(
            f! = rhs!,
            dim = 8,
            mass_matrix = Diagonal([0, 0, 1, 0, 0, 0, 0, 0]),
            sym = Symbol[
                :u_r,
                :u_i,
                :θ,
                :i_d_r,
                :i_q_r,
                :i_d_r_b,
                :i_q_r_b,
                :iqp,
            ],
        )
    end
    symbolsof(::StaticGeneratorObs) = begin
        [:u_r, :u_i, :θ, :i_d_r, :i_q_r, :i_d_r_b, :i_q_r_b, :iqp]
    end
    dimension(::StaticGeneratorObs) = begin
        8
    end
end

## setup system

# define events
t_fault = 0.5 # onset of node short circuit
t_duration = 0.15 # duration of node short circuit

Ron = 2.5 # Ω
Znsc = complex(Ron, 0.0) * base_admittance #  fixed by Emocosy
nsc_node = 2 # sc on low-voltage side

## setup system
SL = SlackAlgebraic(; U = complex(1.0, 0.0))

DG = StaticGeneratorObs(;
    K_pll = 1632.993, #Hz/pu
    S_ref = complex(0.5, 0.5), # pu
    K_FRT = 2.0,
    I_max = 1.0, # pu
    Vref = 1.0, # pu
    Vdead = 0.1, #pu
    Y_n = t ->(t_fault < t < t_fault + t_duration ? 1 / Znsc : 0.0),
)

T = OLTC(
    from = 1,
    to = 2,
    Uos = 110E3 / base_voltage_HV, # Bemessungsspannung Oberspannungsseite in kV
    Uus = 20E3 / base_voltage,# Bemessungsspannung Unterspannungsseite in kV
    k = 0, # Kennzahl der Schaltgruppe, Bsp: YD5-> 5
    ssp = 0.0, # Stufenschalterposition
    stufe_us = 0.625, # Stufung pro Stufenschalterstellung in %
    Sr = 25E6 / base_power, #Bemessungsscheinleistung in MVA
    uk = 12.0, # Kurzschlussspannung in %
    Pvk = 25E3 / base_power, # Kupferverluste in MW
    Pvl = 0.0, # Eisenverluste in kW
    iLeer = 0.0, # Leerlaufstrom in %
)

pg = PowerGrid([SL, DG], [T])

## load data from powerfactory for comparison

powerfactory =
    CSV.File("$dir/FRT_Vac_id_iq.csv", header = true, normalizenames = true)

## find operating point

# set all initial voltages to slack value
op = find_operationpoint(pg)

## compare power flow
@show op[2, :v] # 1.00241 pu
@show op[2, :φ] .|> rad2deg # 0.14 deg

@show op[1, :s] * 1e3 # -499.98 kW, -497.61 kvar
@show op[2, :s] * 1e3 # 500.0 kW, 500.0 kvar

# verketteter Strom / sqrt(3)
@show op[1, :iabs] * base_current_HV * 1e-3 / sqrt(3) # 0.004 kA
@show op[2, :iabs] * base_current * 1e-3 / sqrt(3) # 0.020 kA

@show op[2, :i_d_r] # 0.4988 p.u.
@show op[2, :i_q_r] # -0.4988 p.u.

# since FRT is not active, the currents should be identical
@assert op[2, :i_d_r_b] == op[2, :i_d_r]
@assert op[2, :i_q_r_b] == op[2, :i_q_r]

# iq_plus is zero
@assert op[2, :iqp] == 0

## simulate short circuit

ode = rhs(pg)
problem = ODEProblem(ode, op.vec, (0.0, 1.0))
_sol = solve(
    problem,
    Rodas4(),
    force_dtmin = true,
    tstops = [t_fault, t_fault + t_duration],
)

# experimental version:
# include("$dir/short_circuit.jl")
#
# nsc = NodeShortCircuit(;
#     node_number = nsc_node,
#     Y = 1/Znsc,
#     tspan_fault = (t_fault, t_fault + t_duration),
# )
#
# sol = simulate(nsc, op, (0., 1.))
sol = PowerGridSolution(_sol, pg)

## plot result

scatter(powerfactory.b_tnow_in_s, powerfactory.s_Vac_ist, label = "PF", ms = 2)
plot!(sol, 2:2, :v, label = "PD")
xlims!(0.49, 0.66)
xlabel!("t [s]")
ylabel!("V_ac [pu]")

scatter(powerfactory.b_tnow_in_s, powerfactory.s_id_ref, label = "PF", ms = 2)
plot!(sol, 2, :i_d_r, label = "PD")
xlims!(0.49, 0.66)
xlabel!("t [s]")
ylabel!("i_d_ref [pu]")

scatter(powerfactory.b_tnow_in_s, powerfactory.s_iq_ref, label = "PF", ms = 2)
plot!(sol, 2, :i_q_r, label = "PD")
xlims!(0.49, 0.66)
xlabel!("t [s]")
ylabel!("i_q_ref [pu]")

# check FRT function
scatter(
    powerfactory.b_tnow_in_s,
    powerfactory.s_iq_ref,
    label = "i_q_ref PF",
    ms = 2,
)
plot!(sol, 2, :i_q_r, label = "FRT+limit")
plot!(sol, 2, :i_q_r_b, label = "before FRT")
plot!(sol, 2, :iqp, label = "iqplus")
xlims!(0.49, 0.66)
xlabel!("t [s]")
ylabel!("q currents [pu]")

scatter(powerfactory.b_tnow_in_s, powerfactory.s_id_ref, label = "PF", ms = 2)
plot!(sol, 2, :i_d_r, label = "FRT+limit")
plot!(sol, 2, :i_d_r_b, label = "before FRT")
xlims!(0.49, 0.66)
xlabel!("t [s]")
ylabel!("d currents [pu]")

## alternative model to compare short circuit power provided by slack bus

# PQ = PQAlgebraic(; S=complex(0., 0.), Y_n = 0.)
# pg = PowerGrid([SL, PQ], [T])

## to use the NodeShortCircuit, T must be converted to a PiModelLine

# # For ssp=0, replace it with PiModelLine
# # assume t_km = 1, t_mk real
# Y = PiModel(T) #./ admittance_m
# ue = sqrt(-Y[2, 2] / Y[1, 1])
# y = Y[1, 2] / ue
# ys = -Y[1, 1] - y
#
# L = PiModelLine(;
#     from = T.from,
#     to = T.to,
#     y = y,
#     y_shunt_km = ys,
#     y_shunt_mk = ys,
# )
# @assert iszero(T.ssp) && all(PiModel(L) .== PiModel(T))

max_stars_repo_name: luap-pik/PowerDynIC

max_issues_repo_name: luap-pik/PowerDynIC

max_forks_repo_name: luap-pik/PowerDynIC


text:
include("state/individualstate.jl")
include("state/runningstats.jl")

# the mutable part of the simulation
mutable struct SimState
  rng::MersenneTwister
  time::TimePoint
  queue::EventQueue
  individuals::Vector{IndividualState}
  forest::InfectionForest
  stats::RunningStats

  SimState(rng::AbstractRNG, num_individuals::Integer) = num_individuals <= 0 || num_individuals > typemax(UInt32) ? error("number of individuals must be positive and smaller than $(typemax(UInt32))") :
    new(
      rng,

      0.0,
      EventQueue(),

      fill(IndividualState(), num_individuals),

      InfectionForest(num_individuals),

      RunningStats()
    )
end

SimState(num_individuals::Integer; seed::Integer=0) = SimState(MersenneTwister(seed), num_individuals)

function show(io::IO, state::SimState)
  print(io, "Simulation state for ", length(state.individuals), " individuals")
end

function reset!(state::SimState, rng::AbstractRNG)
  if isa(rng, AbstractRNG)
    state.rng = rng
  end
  state.time=0
  empty!(state.queue)
  reset!(state.forest)
  fill!(state.individuals, IndividualState())
  state.stats = RunningStats()
  state
end

time(state) = state.time
numindividuals(state::SimState) = length(state.individuals)

reset!(state::SimState) = reset!(state::SimState, state.rng)
reset!(state::SimState, seed::Integer) = reset!(state, MersenneTwister(seed))

individualstate(state::SimState, person_id::Integer) = state.individuals[person_id]
health(state::SimState, person_id::Integer)::HealthState = state.individuals[person_id].health
freedom(state::SimState, person_id::Integer)::FreedomState = state.individuals[person_id].freedom

quarantine_level(state::SimState, person_id::Integer) = state.individuals[person_id].quarantine_level
isquarantined(state::SimState, person_id::Integer)::Bool = quarantine_level(state, person_id) != 0
detected(state::SimState, person_id::Integer)::DetectionStatus = state.individuals[person_id].detected
isdetected(state::SimState, person_id::Integer)::Bool = (Detected == detected(state, person_id))

subjecthealth(state::SimState, event::Event)::HealthState = health(state, subject(event))
subjectfreedom(state::SimState, event::Event)::FreedomState = freedom(state, subject(event))

sourcehealth(state::SimState, event::Event)::HealthState = health(state, source(event))
sourcefreedom(state::SimState, event::Event)::FreedomState = freedom(state, source(event))

#forwardinfections(state::SimState, person_id::Integer) = inclusive(state.infections, searchequalrange(state.infections, person_id)...) |> values
forwardinfections(state::SimState, person_id::Integer) = forwardinfections(state.forest, person_id)
backwardinfection(state::SimState, person_id::Integer) = backwardinfection(state.forest, person_id)

function sethealth!(state::SimState, person_id::Integer, new_health::HealthState)
  orig = state.individuals[person_id]
  @assert orig.health <= new_health
  state.individuals[person_id] = @set orig.health = new_health
  nothing
end

function setfreedom!(state::SimState, person_id::Integer, new_freedom::FreedomState)
  orig = state.individuals[person_id]
  state.individuals[person_id] = @set orig.freedom = new_freedom
  nothing
end

function setdetected!(state::SimState, person_id::Integer, new_detected::DetectionStatus)
  orig = state.individuals[person_id]
  @assert orig.detected <= new_detected
  state.individuals[person_id] = @set orig.detected = new_detected
  nothing
end

function quarantine_advance!(state::SimState, person_id::Integer, adv_val::Integer)
  orig = state.individuals[person_id]
  if adv_val >= 0
    state.individuals[person_id] = @set orig.quarantine_level = orig.quarantine_level+adv_val
  else
    state.individuals[person_id] = @set orig.quarantine_level = orig.quarantine_level-(-adv_val)
  end
  nothing
end

function quarantine_cancel!(state::SimState, person_id::Integer)
  orig = state.individuals[person_id]
  state.individuals[person_id] = @set orig.quarantine_level = 0
  nothing
end

setsubjecthealth!(state::SimState, event::Event, health::HealthState) = sethealth!(state, subject(event), health)
setsubjectfreedom!(state::SimState, event::Event, freedom::FreedomState) = setfreedom!(state, subject(event), freedom)

registerinfection!(state::SimState, infection::Event) = push!(state.forest, infection)

function initialfeed!(state::SimState, num_initial_infections::Integer)
  @assert 0 == time(state)

  N = length(state.individuals)
  individuals = 1:N

  for _ in 1:num_initial_infections
    person_id = sample(state.rng, individuals)
    event = Event(Val(OutsideInfectionEvent), 0.0, person_id)
    push!(state.queue, event)
  end

end

saveparams(dict, state::SimState, prefix::AbstractString="") = saveparams(dict, state.forest, prefix)

ext: jl


text:
# Example 2   : Quadrotor path planning

# Constraints :
# 				Position in box
# 				Velocity in ball
# 				Thrust in ball ∩ cone
#				Thrust in halfspace (z component lower bound)
# 				Thrust rate in box

# problem definition
module prb

	using LinearAlgebra, StaticArrays

	const nx = 6												# state dim
	const nu = 3												# imput dim
	const N = 30												# horizon length
	const Δ = 0.25 												# sampling time
	const tvec = [(j-1)*Δ for j in 1:N]   						# time vector (t=0:N-1)

	# solution method settings
		# common
		const kmax = 2000										# max iterations for main algorithms (and outer loop)
		const scale_Hg = true

		# JuMP
		const ϵ_pd = 1e-9										# solver precision
		const ϵ_gap = 1e-9										# solver precision

		# PIPG
		const ϵ_pipg = 1e-4										# termination precision
		const restrt_idx = 300									# index for restarting

		# ADMM
		const α_admm = 1.0										# step size
		const ϵ_nest_admm = 1e-10								# Nesterov's method termination precision
		const kmax_nest = 100									# max iterations for Nesterov's method

		# Chambolle and Pock
		const β⁰ = 2.0											# initial step size
		const γ⁰ = 1.0											# initial step size

	# system parameters
	const nxb2 = Int64(nx/2)									# position veloity dim
	const α_tl = tand( 45 )										# tangent of max thrust pointing angle

	const m = 0.35												# quadrotor mass
	const gₑ = 9.806	 										# acceleration due to gravity at Earth sea-level (m s⁻²)
	const pmax = 3.0											# max position ∞ norm
	const vmax = 5.0											# max velocity 2 norm
	const umax = 5.0					 						# max input (forcing) 2 norm
	const umin = 2.0											# min z component for input
	const dumax = 3.0										 	# max input rate ∞ norm

	# const p1 = SVector{3}([-10.0,-10.0,0.0])  				# center of obstacle #1
	# const p2 = SVector{3}([8.0,-8.0,0.0])	  					# center of obstacle #2
	# const p3 = SVector{3}([10.0,10.0,0.0])	  				# center of obstacle #3

	# const r1 = 7.6										  	# radius of obstacle #1
	# const r2 = 9.0						      			    # radius of obstacle #2
	# const r3 = 7.6                            			    # radius of obstacle #3

	const p1 = SVector{3}([-1.5,-1.5,0.0])  					# center of obstacle #1
	const p2 = SVector{3}([1.2,-1.2,0.0])	  					# center of obstacle #2
	const p3 = SVector{3}([1.5,1.5,0.0])	  					# center of obstacle #3

	const r1 = 0.8										  	  	# radius of obstacle #1
	const r2 = 1.2						      					# radius of obstacle #2
	const r3 = 0.8                            					# radius of obstacle #3

	# reference trajectory
	const x0 = SVector{nx}(cat(p1 .-
		[0.0,1.2*r1,0.0],[0.0,0.0,0.0],dims=1))					# initial condition (t=0)
	const xf = SVector{nx}(cat(p3 .+
		[1.2*r3,0.0,0.0],[0.0,0.0,0.0],dims=1))
	const xref = [SVector{nx}((1-t/(N-1)) .* x0 .+
		(t/(N-1)) .* xf) for t in 1:N-1]		        		# state reference (t=1:N-1)
	const uref = [SVector{nu}(zeros(nu)) for _ in 1:N-1]	    # input reference (t=0:N-2)

	# dynamics (t=0:N-2)
	const A = [SMatrix{nx,nx}([1.0 0.0 0.0 Δ   0.0 0.0
							   0.0 1.0 0.0 0.0 Δ   0.0
							   0.0 0.0 1.0 0.0 0.0 Δ
							   0.0 0.0 0.0 1.0 0.0 0.0
							   0.0 0.0 0.0 0.0 1.0 0.0
							   0.0 0.0 0.0 0.0 0.0 1.0]) for _ in 1:N-1]
	const B = [SMatrix{nx,nu}([0.5*(Δ^2)/m  0.0          0.0
							   0.0          0.5*(Δ^2)/m  0.0
							   0.0          0.0          0.5*(Δ^2)/m
							   Δ/m          0.0          0.0
							   0.0          Δ/m          0.0
							   0.0          0.0          Δ/m]) for _ in 1:N-1]
	const e = [SVector{nx}([0.0, 0.0, -0.5*(Δ^2)*gₑ, 0.0, 0.0, -Δ*gₑ]) for _ in 1:N-1]

	# state and control penalty (LQR weights)

	# const q_tmp = [1.0,3.0]
	# const r_tmp = 5.0

	const q_tmp = [1.0,2.5]
	const r_tmp = 0.5

	const Q = [SMatrix{nx,nx}(Array(Diagonal(cat(q_tmp[1] .* ones(nxb2), q_tmp[2] .* ones(nxb2), dims=1)))) for _ in 1:N-1] # t=1:N-1
	const R = [SMatrix{nu,nu}(r_tmp .* Array(Diagonal(ones(nu)))) for _ in 1:N-1] 											# t=0:N-2

end

# auxiliary definitions for the the problem
module prb2
	using ..prb
	using LinearAlgebra, StaticArrays

	const nx = copy(prb.nx)
	const nu = copy(prb.nu)
	const nxb2 = copy(prb.nxb2)
	const N = copy(prb.N)

	# obstacle avoidance
	# linearization
	# initial and final point for obstacle linearization reference
	const x̄0 = SVector{nx}(cat(prb.p1 .- [0.0,prb.r1,0.0],[0.0,0.0,0.0],dims=1))
	const x̄f = SVector{nx}(cat(prb.p3 .+ [prb.r3,0.0,0.0],[0.0,0.0,0.0],dims=1))
	# obstructed straight-line reference for obstacle linearization
	const x̄_obs = [SVector{nx}( (1-t/(N-1)) .* x̄0 .+ (t/(N-1)) .* x̄f ) for t=1:(N-1)]

	# move state to nearest obstacle boundary if penetrated
	function path_unobstructed(z)
		# globals used p1, p2, p3
		yy = zeros(prb.nx)
		yy .= z

		tol_fac = 1.0

		nrm_y_p1 = norm(yy[1:3]-prb.p1)
		if nrm_y_p1 ≤ tol_fac*prb.r1
			yy[1:3] .= (yy[1:3] .- prb.p1) .* (tol_fac*prb.r1/nrm_y_p1) .+ prb.p1
		end

		nrm_y_p2 = norm(yy[1:3] .- prb.p2)
		if nrm_y_p2 ≤ tol_fac*prb.r2
			yy[1:3] .= (yy[1:3] .- prb.p2) .* (tol_fac*prb.r2/nrm_y_p2) .+ prb.p2
		end

		nrm_y_p3 = norm(yy[1:3]-prb.p3)
		if nrm_y_p3 ≤ tol_fac*prb.r3
			yy[1:3] .= (yy[1:3] .- prb.p3) .* (tol_fac*prb.r3/nrm_y_p3) .+ prb.p3
		end
		return yy
	end

	# unobstructed straight-line reference for obstacle linearization
	const x̄_unobs = [SVector{nx}(path_unobstructed(x̄_obs[t])) for t=1:(N-1)]

	# terms in the obstacle avoidance constraint linearization
	#  - ϕᵢ⋅x + θᵢ ≈ norm(x[1:3]-pᵢ) - rᵢ ≥ 0
	const θ₁ = SVector{N-1}([norm(x̄_unobs[t][1:3] .- prb.p1) - prb.r1 + (-1 * norm(x̄_unobs[t][1:3])^2 + prb.p1⋅x̄_unobs[t][1:3]) / norm(x̄_unobs[t][1:3] .- prb.p1) for t=1:(N-1)])
	const θ₂ = SVector{N-1}([norm(x̄_unobs[t][1:3] .- prb.p2) - prb.r2 + (-1 * norm(x̄_unobs[t][1:3])^2 + prb.p2⋅x̄_unobs[t][1:3]) / norm(x̄_unobs[t][1:3] .- prb.p2) for t=1:(N-1)])
	const θ₃ = SVector{N-1}([norm(x̄_unobs[t][1:3] .- prb.p3) - prb.r3 + (-1 * norm(x̄_unobs[t][1:3])^2 + prb.p3⋅x̄_unobs[t][1:3]) / norm(x̄_unobs[t][1:3] .- prb.p3) for t=1:(N-1)])

	const ϕ₁ = [SVector{nxb2}(-1 .* (x̄_unobs[t][1:3] .- prb.p1) ./ norm(x̄_unobs[t][1:3] .- prb.p1) ) for t=1:N-1]
	const ϕ₂ = [SVector{nxb2}(-1 .* (x̄_unobs[t][1:3] .- prb.p2) ./ norm(x̄_unobs[t][1:3] .- prb.p2) ) for t=1:N-1]
	const ϕ₃ = [SVector{nxb2}(-1 .* (x̄_unobs[t][1:3] .- prb.p3) ./ norm(x̄_unobs[t][1:3] .- prb.p3) ) for t=1:N-1]
end

include("proj_funcs.jl") 					# projection library

# assemble quantities for vectorized PIPG+ implementation
module asm

	using ..prb
	using ..prb2
	using ..proj 											# projection onto simple sets
	using LinearAlgebra, StaticArrays, JuMP

	const scale_Hg = prb.scale_Hg
	const nH = 2*prb.nu*(prb.N-2)+(prb.N-1)+3*(prb.N-1)	    # no. of inequality constraints

	function construct_H(nx,nu,N,A,B,Inx,Inu)
		nxnu = nx+nu
		nxb2 = Int64(nx/2)
		H = zeros(nx*(N-1) + nH, nxnu*(N-1))
		H[1:nx,1:nxnu] .= [-B[1] Inx]
		for j=1:N-2
			H[j*nx+1:(j+1)*nx, nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+2*nx+nu] .= [-A[j] -B[j] Inx]

			H[nx*(N-1)+(j-1)*nu+1:nx*(N-1)+j*nu, (j-1)*nxnu+1:(j-1)*nxnu+nu+nxnu] .= [Inu zeros(nu,nx) -Inu]

			H[nx*(N-1)+nu*(N-2)+(j-1)*nu+1:nx*(N-1)+nu*(N-2)+j*nu, (j-1)*nxnu+1:(j-1)*nxnu+nu+nxnu] .= [-Inu zeros(nu,nx) Inu]
		end
		for j=1:N-1
			H[nx*(N-1)+2*nu*(N-2)+j,(j-1)*nxnu+1:(j-1)*nxnu+nu] .= [0.0,0.0,1.0]

			H[nx*(N-1)+2*nu*(N-2)+(N-1)+j,nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nxb2] .= -1 .* prb2.ϕ₁[j]

			H[nx*(N-1)+2*nu*(N-2)+2*(N-1)+j,nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nxb2] .= -1 .* prb2.ϕ₂[j]

			H[nx*(N-1)+2*nu*(N-2)+3*(N-1)+j,nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nxb2] .= -1 .* prb2.ϕ₃[j]
		end

		norm_row_H = zeros(nx*(N-1)+nH)
		for j=1:nx*(N-1)+nH
			norm_row_H[j] = norm(H[j,1:end],Inf)
		end

		if scale_Hg
			print("Scaling rows of H\n")
			for j=1:length(norm_row_H)
				H[j,1:end] .= H[j,1:end] ./ norm_row_H[j]
			end
		end

		return H,norm_row_H
	end

	function construct_G(nx,nu,N,Q,R)
		nxnu = nx+nu
		G = zeros(nxnu*(N-1),nxnu*(N-1))
		for j=1:N-1
			G[(j-1)*nxnu+1:j*nxnu,(j-1)*nxnu+1:j*nxnu] .= [R[j] 		zeros(nu,nx)
														   zeros(nx,nu) Q[j]]
		end
		return G
	end

	H_scl,norm_row_H = construct_H(prb.nx,prb.nu,prb.N,prb.A,prb.B,Diagonal(ones(prb.nx)),Diagonal(ones(prb.nu)))
	const H = deepcopy(H_scl)
	const HT = transpose(H)

	const zref = cat([cat(prb.uref[t],prb.xref[t],dims=1) for t=1:prb.N-1]...,dims=1)
	const h = cat([cat(-prb.R[t]*prb.uref[t], -prb.Q[t]*prb.xref[t], dims=1) for t=1:prb.N-1]...,dims=1)

	if scale_Hg
		print("Scaling components of g\n")
		const scl_val_Hg = norm_row_H
	else
		const scl_val_Hg = 1.0
	end

	const g = cat(prb.A[1]*prb.x0 .+ prb.e[1],cat([prb.e[t] for t=2:prb.N-1]...,dims=1), -prb.dumax .* ones(2*prb.nu*(prb.N-2)),
			prb.umin .* ones(prb.N-1),-1 .* prb2.θ₁...,-1 .* prb2.θ₂...,-1 .* prb2.θ₃...,dims=1) ./ scl_val_Hg

	@assert length(h) == (prb.nx+prb.nu)*(prb.N-1) "Length of h is incorrect."
	@assert length(g) == 2*prb.nu*(prb.N-2)+(prb.N-1)+3*(prb.N-1)+prb.nx*(prb.N-1) "Length of g is incorrect."

	const Heq = [H -1 .* Array(Diagonal(ones(prb.nx*(prb.N-1) + nH)))]
	const HeqT = transpose(Heq)

	const G = construct_G(prb.nx,prb.nu,prb.N,prb.Q,prb.R)

	const svd_HTH = svd(HT*H)
	const svd_HeqTHeq = svd(HeqT*Heq)

	const svd_G = svd(G)

	const σ = max(svd_HTH.S...)
	const σ_eq = max(svd_HeqTHeq.S...)
	const λ = max(svd_G.S...)
	const μ = min(svd_G.S...)

	function construct_xu(z::Array{Float64,1})
		return cat([prb.x0],[z[nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nx] for j=1:N-1],dims=1), [z[(j-1)*nxnu+1:(j-1)*nxnu+nu] for j=1:N-1]
	end

	# useful constants
	const nx = copy(prb.nx)
	const nu = copy(prb.nu)
	const nxnu = nx+nu
	const nxb2 = Int64(nx/2)
	const N = copy(prb.N)

	const κ_x = MVector{nx}(randn(nx))
	const κ_xb2 = MVector{nxb2}(randn(nxb2))
	const κ_u = MVector{nu}(randn(nu))

	# projection
	function project_U!(u)
		κ_u .= u
		γ_u = randn()
		proj.soc_ball!(u,κ_u,prb.α_tl,prb.umax,nu,γ_u)
	end
	function project_X!(x)
		κ_x .= x
		proj.box!(κ_xb2,κ_x[1:nxb2],-prb.pmax,prb.pmax,nxb2)
		x[1:nxb2] .= κ_xb2
		proj.ball!(κ_xb2,κ_x[nxb2+1:nx],prb.vmax,nxb2)
		x[nxb2+1:nx] .= κ_xb2
	end

	# set state and input constraints in JuMP
	function set_constraint_jump_vectorized!(model,zz)
		# acceleration, position and velocity
		@constraint(model,[j=1:N-1],[prb.umax,zz[(j-1)*nxnu+1:(j-1)*nxnu+nu]...] in SecondOrderCone())
		@constraint(model,[j=1:N-1],[prb.α_tl*zz[(j-1)*nxnu+nu],zz[(j-1)*nxnu+1:(j-1)*nxnu+nu-1]...] in SecondOrderCone())
		@constraint(model,[j=1:N-1],-prb.pmax .≤ zz[nu+(j-1)*nxnu+1:nu+(j-1)*nxnu+nxb2] .≤  prb.pmax)
		@constraint(model,[j=1:N-1],[prb.vmax,zz[nu+(j-1)*nxnu+nxb2+1:nu+(j-1)*nxnu+nx]...] in SecondOrderCone())
	end

	# set state and input constraints (everything other than dynamics and initial condition) in JuMP (decomposed version)
	function set_constraint_jump_decomposed!(model,x,u)

		@constraint(model,[t=1:N-1], -prb.pmax .≤ x[1:nxb2,t+1] .≤ prb.pmax)
		@constraint(model,[t=1:N-1], [prb.vmax,x[nxb2+1:nx,t+1]...] in SecondOrderCone())
		@constraint(model,[t=1:N-1], [prb.umax,u[1:nu,t]...] in SecondOrderCone())
		@constraint(model,[t=1:N-1], [prb.α_tl*u[nu,t],u[1:nu-1,t]...] in SecondOrderCone())

		# linear inequality constraints
		@constraint(model,[t=1:N-2], -prb.dumax .≤ u[:,t+1] .- u[:,t] .≤ prb.dumax)
		@constraint(model,[t=1:N-1], u[3,t] ≥ prb.umin)
		@constraint(model,[t=1:N-1], prb2.ϕ₁[t]⋅x[1:nxb2,t+1] - prb2.θ₁[t] ≤ 0)
		@constraint(model,[t=1:N-1], prb2.ϕ₂[t]⋅x[1:nxb2,t+1] - prb2.θ₂[t] ≤ 0)
		@constraint(model,[t=1:N-1], prb2.ϕ₃[t]⋅x[1:nxb2,t+1] - prb2.θ₃[t] ≤ 0)
	end

end

module plotter
using LinearAlgebra
using Plots, LaTeXStrings
gr()
# This ensures that a legend entry is not created by default
default(lab="",markersize=2,markerstrokewidth=0.1,xtickfontsize=12,ytickfontsize=12,
	ztickfontsize=12,legendfontsize=9)
using ..prb
using ..prb2

function trajectory2D(x1,u1,x2,u2,slver=:ecos,pltflg=true)

	y_unscl = cat([prb.x0],prb.xref,dims=1)

	# Position

	# obstacle regions
	x = prb.p1[1] .+ collect(-prb.r1:(2*prb.r1/(100-1)):prb.r1)
	y = prb.p1[2] .+ sqrt.( abs.(prb.r1^2 .- (x .- prb.p1[1]) .^ 2) )
	p1 = plot(x,y,color=:black,line=:solid)
	y = prb.p1[2] .- sqrt.( abs.(prb.r1^2 .- (x .- prb.p1[1]) .^ 2) )
	plot!(x,y,color=:black,line=:solid)

	x = prb.p2[1] .+ collect(-prb.r2:(2*prb.r2/(100-1)):prb.r2)
	y = prb.p2[2] .+ sqrt.( abs.(prb.r2^2 .- (x .- prb.p2[1]) .^ 2) )
	plot!(x,y,color=:black,line=:solid)
	y = prb.p2[2] .- sqrt.( abs.(prb.r2^2 .- (x .- prb.p2[1]) .^ 2) )
	plot!(x,y,color=:black,line=:solid)

	x = prb.p3[1] .+ collect(-prb.r3:(2*prb.r3/(100-1)):prb.r3)
	y = prb.p3[2] .+ sqrt.( abs.(prb.r3^2 .- (x .- prb.p3[1]) .^ 2) )
	plot!(x,y,color=:black,line=:solid)
	y = prb.p3[2] .- sqrt.( abs.(prb.r3^2 .- (x .- prb.p3[1]) .^ 2) )
	plot!(x,y,color=:black,line=:solid)

	plot!([prb.x0[1]],[prb.x0[2]],color=:orange,marker=:circle,linealpha=0.01,lab=L"x_0[1:3]")
	plot!([prb.xf[1]],[prb.xf[2]],color=:orange,marker=:square,linealpha=0.01,lab=L"x_f[1:3]")

	plot!([y_unscl[t][1] for t=1:prb.N],[y_unscl[t][2] for t=1:prb.N],color=:red,line=:solid,lw=0.7,lab="Cost Ref.")
	plot!([prb2.x̄_unobs[t][1] for t=1:prb.N-1],[prb2.x̄_unobs[t][2] for t=1:prb.N-1],color=:magenta,line=:solid,lw=1.3,lab="Obstcl. Ref.")
	# lim_set = [-max(prb.r1,prb.r2,prb.r3),max(prb.r1,prb.r2,prb.r3)]
	lim_set = [-1.0*prb.pmax,prb.pmax*1.0]
	plot!(xlim=lim_set,ylim=lim_set)

	if pltflg
		plot!([x1[t][1] for t=1:prb.N],[x1[t][2] for t=1:prb.N],color=:blue,line=:solid,lab="PIPG")
		plot!([x2[t][1] for t=1:prb.N],[x2[t][2] for t=1:prb.N],color=:cyan,line=:dash,lab=string("JuMP - ",String(slver)))
	end

	plot!(size=[500,500],legend=:topleft,title="Position")

	vmax = prb.vmax
	umax = prb.umax
	umin = prb.umin

	# Velocity
	p2 = plot(prb.tvec,vmax .* ones(prb.N),color=:black,line=:solid)

	if pltflg
		plot!(prb.tvec,[norm(x1[t][4:6]) for t=1:prb.N],color=:blue,line=:solid,lab="PIPG")
		plot!(prb.tvec,[norm(x2[t][4:6]) for t=1:prb.N],color=:cyan,line=:dash,lab=string("JuMP - ",String(slver)))
	end

	plot!(size=[500,500],legend=:none,title="Velocity")

	# Acceleration

	p3 = plot(prb.tvec[1:prb.N-1],umax .* ones(prb.N-1),color=:black,line=:solid,label=L"u_{max}")
	plot!(prb.tvec[1:prb.N-1],umin .* ones(prb.N-1),color=:black,line=:solid,label=L"u_{min}")

	if pltflg
		plot!(prb.tvec[1:prb.N-1],[norm(u1[t]) for t=1:prb.N-1],color=:blue,line=:solid,lab="PIPG")
		plot!(prb.tvec[1:prb.N-1],[norm(u2[t]) for t=1:prb.N-1],color=:cyan,line=:dash,lab=string("JuMP - ",String(slver)))
	end

	plot!(size=[500,500],legend=:topright,title="Acceleration",ylim=[0.8*umin,1.1*umax])

	λ = @layout [a b c]
	display(plot(p1,p2,p3,layout=λ,size=[999,333]))

	if pltflg
		p0 = plot(prb.tvec,[x1[t][3] for t=1:prb.N],color=:orange,line=:solid,lab="PIPG",title="Altitude",size=[400,400])
		plot!(prb.tvec,[x2[t][3] for t=1:prb.N],color=:green,line=:dash,lab=string("JuMP - ",String(slver)))
		display(p0)
	end
end

end

max_stars_repo_name: purnanandelango/pipg-demo

max_issues_repo_name: UW-ACL/pipg-demo

max_forks_repo_name: UW-ACL/pipg-demo


text:
module MPV2
using LinearAlgebra
function back_subst(R,b,xrem)
 n = length(b)
 x = zeros(n)
 for i=n:-1:1
   if abs(R[i,i]) > 1.e-15
     x[i] = (b[i] - R[i,i+1:n]'*x[i+1:n]) / R[i,i]
   else x[i] = xrem[i]
   end
   end
 return x
end;

function CheckConvergency(p,m,k, b, Nji, F, ex, xx, xs, StandardPressure)
# check MB
  epsmb = 0
  for j in 1:m
    x=0
    for i in 1:k
      x+=ex[i]*Nji[i,j]
    end
  x=abs(b[j]-x)
  if x > epsmb epsmb = x end
  end
  SumMolarFractions=0
  SumG=0
  px=(1.e6/StandardPressure)*p
  for i in 1:k
    x=0
    for j in 1:m
    x+=Nji[i,j]*xx[j]
    end
  SumMolarFractions+=exp(F[i]+x)/px
  SumG+=(log(px*ex[i]/xs)-F[i])*ex[i]
  end
  SumL=0
  for j in 1:m
    SumL+=b[j]*xx[j]
  end
  epsg=abs(SumL-SumG) # sum(bj*Lamj) - sum(Gi(T,pi))
  return epsmb, epsg, SumMolarFractions, SumG
end

function MicPV(T,p,v, m,k, b, Nji, F)
      Parameters=zeros(4); EquilibriumComposition=zeros(k); bb = zeros(m+1)
      Rgas = 8.314462618
      StandardPressure = 101325   #  =1.e5, if 1 bar
      global xx, xs
sum_b=0;
for i in 1:m sum_b+=b[i] end
for i in 1:m b[i] = b[i]/sum_b end
v=v/sum_b
    if (m == 0) | (k == 0) ier = " (m == 0) | (k == 0)"; return ier, Parameters, EquilibriumComposition,0 end
      m1=m+1; m2=m+2; iter=0; ier="";
      x=zeros(k); ex=zeros(k); str=zeros(m2); xrem=zeros(m1)
        for i in 1:k
          x[i] = -15.0
          ex[i]= exp(x[i]) end
        y=0.0
while true
      eps = 0.0; S=zeros(m1,m1)
      for j in 1:m bb[j]=b[j] end
      if v <= 0.0
	S[m1,m1]=(1.e6/StandardPressure)*p*exp(-y)
 	bb[m1]=S[m1,m1]*(1.0+y)
      else
        S[m1,m1]=1.0
        bb[m1]=log((Rgas/StandardPressure)*T/v) end
      for i in 1:k
        str[m1]=-ex[i]
        str[m2]=(x[i]-F[i]-1.0)*ex[i]
        for j in 1:m
          str[j]=Nji[i,j]*ex[i]; end
       if v<=0.0
        for j in 1:m1
          S[m1,j]=S[m1,j]+str[j] end;
       bb[m1]=bb[m1]+str[m2]
       end
       for jj in 1:m # do
         for j in 1:m1 # do
           S[jj,j] = S[jj,j]+str[j]*Nji[i,jj]; end;
        bb[jj] = bb[jj]+str[m2]*Nji[i,jj]
        end
      end;
# find the SLE solution
      z=qr(S)
      yy=z.Q'*bb
      xx = back_subst(z.R,yy, xrem)
# modify the results of iteraion
      w = 0.5
      if (abs(y) > 1.0) w = 0.5*abs(y);  end
      if (abs(xx[m1]-y) < w) w=abs(xx[m1]-y);  end
      if (xx[m1] < y) w = -w;  end
      if (abs(y)>1e-10) & (abs(w/y) > eps) eps=abs(w/y); end
      y = y+w
      xs = 0.0;
      for i in 1:k
        u=F[i]-xx[m1]
        for j in 1:m
          u=u+Nji[i,j]*xx[j]; end
        w=3.0
        if u>=x[i]
          w=1.0
          if x[i]<-7.0 w=5.0 end
        end
        if abs(u-x[i]) < w w=abs(u-x[i]); end
        if u < x[i] w=-w; end
        if (x[i] > -80) & (abs(x[i])>1e-10) & (abs(w/x[i])>eps) eps=abs(w/x[i]); end
        x[i]=x[i]+w;
        ex[i]=exp(x[i]);
        xs=xs+ex[i]
        end;
# remember the result of current iteration
      xrem = copy(xx)
      iter=iter+1;
    eps <= 1.0e-6 && break
    iter > 200 && break
    end; # while true
    if iter > 200 (ier="after 200 iterations eps="*string(eps)) end
# Prepare the results for the export
      if v <= 0.0
        Parameters[1]=p;
      else
        Parameters[1]=1.e-6*Rgas*T*xs/v;
      end;
# check up the convergence
    epsmb, epsg, SumMolarFractions, SumG =CheckConvergency(Parameters[1],m,k, b, Nji, F, ex, xx, xs,StandardPressure)
      Parameters[2]=T;
 println("max relative error = ", eps)
 println("material balance error = ", epsmb)
 println("Gibbs energy error = ", epsg)
 println("sum of molar fractions = ", SumMolarFractions)
      if v <= 0.0
        Parameters[3]=(Rgas/StandardPressure)*T*sum_b/exp(y);
      else
        Parameters[3]=v*sum_b
      end;
      Parameters[4]=xs*sum_b;
      for i in 1:k EquilibriumComposition[i]=ex[i]*sum_b; end
      for i in 1:m b[i]=b[i]*sum_b; end
      return ier, Parameters, EquilibriumComposition, SumG*sum_b
  end
end

max_stars_repo_name: gvbelov/Chemical-Equilibrium

max_issues_repo_name: gvbelov/Chemical-Equilibrium

max_forks_repo_name: gvbelov/Chemical-Equilibrium


text:
# Tests formulations in prob/security-stage
@testset "test security-stage" begin

function contingency_stage_data(network)
    network = deepcopy(network)

    network["cont_label"] = "testing"
    network["delta"] = 0.0

    for (a,gens) in network["area_gens"]
        network["response_gens"] = gens
        break
    end

    for (i,branch) in network["branch"]
        g, b = PowerModels.calc_branch_y(branch)
        tr, ti = PowerModels.calc_branch_t(branch)
        branch["g"] = g
        branch["b"] = b
        branch["tr"] = tr
        branch["ti"] = ti
    end

    bus_gens = gens_by_bus(network)

    for (i,bus) in network["bus"]
        bus["vm_base"] = bus["vm"]
        bus["vm_start"] = bus["vm"]
        bus["va_start"] = bus["va"]
        bus["vm_fixed"] = length(bus_gens[i]) != 0
    end

    for (i,gen) in network["gen"]
        gen["pg_base"] = gen["pg"]
        gen["pg_start"] = gen["pg"]
        gen["qg_start"] = gen["qg"]
        gen["pg_fixed"] = false
        gen["qg_fixed"] = false
    end

    return network
end

solution2_lines = [31,600]
@testset "write_solution2 - $(i)" for (i,network) in enumerate(networks)
    network = contingency_stage_data(network)

    bus_gens = gens_by_bus(network)
    cont = network["gen_contingencies"][1]

    cont_gen = network["gen"]["$(cont.idx)"]
    cont_gen["contingency"] = true
    cont_gen["gen_status"] = 0
    pg_lost = cont_gen["pg"]

    gen_bus = network["bus"]["$(cont_gen["gen_bus"])"]
    if length(bus_gens["$(gen_bus["index"])"]) == 1
        gen_bus["vm_fixed"] = false
    end

    result = run_fixpoint_pf_bqv!(network, pg_lost, nlp_solver)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)

    cont_sol = result["solution"]
    cont_sol["label"] = cont.label
    cont_sol["feasible"] = (result["termination_status"] == LOCALLY_SOLVED)
    cont_sol["cont_type"] = "gen"
    cont_sol["cont_comp_id"] = cont.idx

    cont_sol["gen"]["$(cont.idx)"] = Dict("pg" => 0.0, "qg" => 0.0)
    cont_sol["delta"] = 0.0

    write_solution2(network, [cont_sol], solution_file="solution2-tmp.txt")

    open("solution2-tmp.txt", "r") do file
        lines = countlines(file)
        @test lines == solution2_lines[i]
    end
    rm("solution2-tmp.txt")
end

@testset "pf soft rect - $(i)" for (i,network) in enumerate(networks)
    network = contingency_stage_data(network)

    result = run_pf_bqv_acr(network, nlp_solver)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)
end

@testset "pf soft rect fixedpoint - $(i)" for (i,network) in enumerate(networks)
    network = contingency_stage_data(network)
    pg_lost = 0.0

    result = run_fixpoint_pf_bqv!(network, pg_lost, nlp_solver, iteration_limit=10)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)
end

@testset "pf soft rect fixedpoint - infeasible" begin
    network = contingency_stage_data(network_infeasible)
    pg_lost = 0.0

    result = run_fixpoint_pf_bqv!(network, pg_lost, nlp_solver, iteration_limit=10)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)
end

@testset "pf fixed nbf rect2 - $(i)" for (i,network) in enumerate(networks)
    network = contingency_stage_data(network)

    result = run_pf_fixed_acr(network, nlp_solver)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)
end

@testset "pf fixed nbf rect2 ds- $(i)" for (i,network) in enumerate(networks)
    network = contingency_stage_data(network)

    for (i,gen) in network["gen"]
        if gen["index"] in network["response_gens"]
            gen["pg_fixed"] = true
        end
    end

    result = run_pf_fixed_bp_slack_acr(network, nlp_solver)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)
end

@testset "pf v2_3 fixpoint - $(i)" for (i,network) in enumerate(networks)
    network = contingency_stage_data(network)
    pg_lost = 0.0

    result = run_fixpoint_pf_pvpq!(network, pg_lost, nlp_solver, iteration_limit=5)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
    @test isapprox(result["objective"], 0.0; atol = 1e0)
end

@testset "pf v2_3 fixpoint - infeasible" begin
    network = contingency_stage_data(network_infeasible)
    pg_lost = 0.0

    result = run_fixpoint_pf_pvpq!(network, pg_lost, nlp_solver, iteration_limit=5)

    @test isapprox(result["termination_status"], LOCALLY_SOLVED)
end

end # close test group


text:
module Utils

using DSP
using Misfits
using LinearAlgebra
using Test
using FFTW
using AxisArrays

for file in ["freq", "taper", "adjtest"]
	fn=joinpath(@__DIR__, string(file,".jl"))
	include(fn)
end

end


text:
# Compute a few eigenvalues of the dielectric matrix (q=0,ω=0) iteratively

using DFTK
using Plots
using KrylovKit
using Printf

# Calculation parameters
kgrid = [1, 1, 1]
Ecut = 5

# Silicon lattice
a = 10.26
lattice = a / 2 .* [[0 1 1.]; [1 0 1.]; [1 1 0.]]
Si = ElementPsp(:Si, psp=load_psp("hgh/lda/Si-q4"))
atoms = [Si => [ones(3)/8, -ones(3)/8]]

# change the symmetry to compute the dielectric operator with and without symmetries
model = model_LDA(lattice, atoms, symmetry=:off)
basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)
scfres = self_consistent_field(basis, tol=1e-14)

# Apply ε = 1 - χ0 (vc + fxc)
function eps_fun(dρ)
    dρ = reshape(dρ, size(scfres.ρ.real))
    dρ = from_real(basis, dρ)
    dv = apply_kernel(basis, dρ; ρ=scfres.ρ)
    χdv = apply_χ0(scfres.ham, scfres.ψ, scfres.εF, scfres.eigenvalues, dv)
    vec((dρ - χdv).real)
end

# A straightfoward Arnoldi eigensolver that diagonalizes the matrix at each step
# This is more efficient than Arpack when `f` is very expensive
println("Starting Arnoldi ...")
function arnoldi(f, x0; howmany=5, tol=1e-4, maxiter=30, n_print=howmany)
    for (V, B, r, nr, b) in ArnoldiIterator(f, x0)
        # A * V = V * B + r * b'
        V = hcat(V...)
        AV = V*B + r*b'

        ew, ev = eigen(B, sortby=real)
        Vr = V*ev
        AVr = AV*ev
        R = AVr - Vr * Diagonal(ew)

        N = size(V, 2)
        normr = [norm(r) for r in eachcol(R)]

        println("#--- $N ---#")
        println("idcs      evals     residnorms")
        inds = unique(append!(collect(1:min(N, n_print)), max(1, N-n_print):N))
        for i in inds
            @printf "% 3i  %10.6g  %10.6g\n" i real(ew[i]) normr[i]
        end
        any(imag.(ew[inds]) .> 1e-5) && println("Warn: Suppressed imaginary part.")
        println()

        is_converged = (N ≥ howmany && all(normr[1:howmany] .< tol)
                        && all(normr[end-howmany:end] .< tol))
        if is_converged || (N ≥ maxiter)
            return (λ=ew, X=Vr, AX=AVr, residual_norms=normr)
        end
    end
end
arnoldi(eps_fun, vec(randn(size(scfres.ρ.real))); howmany=5)


text:
# ---
# jupyter:
#   jupytext:
#     formats: jl:light,ipynb
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.5.1
#   kernelspec:
#     display_name: Julia 1.4.2
#     language: julia
#     name: julia-1.4
# ---

# + [markdown] run_control={"frozen": false, "read_only": false}
# # Molecule Objects in  Psi4
#
# This tutorial provides an overview on creating and manipulating molecule objects in  Psi4, illustrated with an example parameterization of the Lennard-Jones potential for Helium dimer.
#
# In order to use  Psi4  within a Julia environment, we may install PyCall
# executing `using Pkg; Pkg.add("PyCall")`. We point PyCall to the Python installation that has  Psi4  installed, that is `ENV["PYTHON"] = ~/miniconda3/envs/p4env/bin/python`, rebuild PyCall, `using Pkg; Pkg.build("PyCall")` and restart Julia. The next time we start Julia we may import  Psi4  with pyimport:
#
# (check [here](https://github.com/dgasmith/psi4numpy/blob/master/README.md) for more details)

# + run_control={"frozen": false, "read_only": false}
using PyCall: pyimport
psi4 = pyimport("psi4")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Unlike in  Psi4  input files, defining a molecule in Julia is done by passing the molecular coordinates as a triple-quoted string to the [`psi4.geometry()`](http://psicode.org/psi4manual/master/api/psi4.driver.geometry.html#psi4.driver.geometry "API Details") function:

# + run_control={"frozen": false, "read_only": false}
he = psi4.geometry("""
He
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Here, not only does the variable `he` refer to the helium molecule, but also an instance of the [`psi4.core.Molecule`](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule "Go to API")
# class in  Psi4; this will be discussed in more detail later.  For a more
# complicated system than an isolated atom, the coordinates can be given in Cartesian or Z-Matrix formats:

# + run_control={"frozen": false, "read_only": false}
h2o = psi4.geometry("""
O
H 1 0.96
H 1 0.96 2 104.5
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Information like the molecular charge, multiplicity, and units are assumed to be 0, 1, and Angstroms, respectively, if not specified within the molecule definition.  This is done by adding one or more [molecule keywords](http://psicode.org/psi4manual/master/psithonmol.html#molecule-keywords "Go to Documentation") to the geometry string used to define the molecule.  Additionally,  Psi4  can detect molecular symmetry, or it can be specified manually.  For example, to define a doublet water cation in $C_{2V}$ symmetry using lengths in Bohr,

# + run_control={"frozen": false, "read_only": false}
doublet_h2o_cation = psi4.geometry("""
1 2
O
H 1 1.814
H 1 1.814 2 104.5

units bohr
symmetry c2v
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# where the line `1 2` defines the charge and multiplicity, respectively.  For systems of non-bonded fragments, the coordinates of each fragment are separated by a double-hyphen `"--"`; this allows for one fragment to be defined with Cartesian and another to be defined with Z-Matrix. For example, the hydronium-benzene complex can be defined with:

# + run_control={"frozen": false, "read_only": false}
hydronium_benzene = psi4.geometry("""
0 1
C          0.710500000000    -0.794637665924    -1.230622098778
C          1.421000000000    -0.794637665924     0.000000000000
C          0.710500000000    -0.794637665924     1.230622098778
C         -0.710500000000    -0.794637665924     1.230622098778
H          1.254500000000    -0.794637665924    -2.172857738095
H         -1.254500000000    -0.794637665924     2.172857738095
C         -0.710500000000    -0.794637665924    -1.230622098778
C         -1.421000000000    -0.794637665924     0.000000000000
H          2.509000000000    -0.794637665924     0.000000000000
H          1.254500000000    -0.794637665924     2.172857738095
H         -1.254500000000    -0.794637665924    -2.172857738095
H         -2.509000000000    -0.794637665924     0.000000000000
--
1 1
X  1  CC  3  30   2  A2
O  13 R   1  90   2  90
H  14 OH  13 TDA  1  0
H  14 OH  15 TDA  13 A1
H  14 OH  15 TDA  13 -A1

CC    = 1.421
CH    = 1.088
A1    = 120.0
A2    = 180.0
OH    = 1.05
R     = 4.0
units angstrom
""")

# + [markdown] run_control={"frozen": false, "read_only": false}
# For non-bonded fragments, the charge and multiplicity should be given explicitly for each fragment.  If not, the
# charge and multiplicity given (or inferred) for the first fragment is assumed to be the same for all fragments.  In
# addition to defining the coordinates outright, we have used variables within the geometry specification strings to
# define bond lengths, angles, and dihedrals in the molecule.  Similarly, we could define the X, Y, or Z Cartesian
# coordinate for any atom in our molecule.
#
# In order to define these variables after the molecule is built, as opposed to within the geometry specification
# itself, there are several ways to do so; one of which will be illustrated in the Lennard-Jones potential example
# below.
#
# When a Psi4 molecule is first built using ``psi4.geometry()``, it is in an unfinished state, as a user may wish to
# tweak the molecule. This can be solved by calling [``psi4.Molecule.update_geometry()``](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.update_geometry "Go to API"). This will update the molecule and restore sanity
# to chemistry.

# +
h2cch2 = psi4.geometry("""
H
C 1 HC
H 2 HC 1 A1
C 2 CC 3 A1 1 D1
H 4 HC 2 A1 1 D1
H 4 HC 2 A1 1 D2

HC = 1.08
CC = 1.4
A1 = 120.0
D1 = 180.0
D2 = 0.0
""")

println("Ethene has $(h2cch2.natom()) atoms")

h2cch2.update_geometry()
println("Ethene has $(h2cch2.natom()) atoms")
# -

# Finally, one can obtain useful information from a molecule by invoking one of several [`psi4.core.Molecule`](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule "Go to Documentation") class methods on the molecule of interest.  For example, if we were interested in verifying that our doublet water cation from above is, in fact, a doublet, we could invoke
# ~~~python
# doublet_h2o_cation.multiplicity()
# ~~~
# Below, some useful Molecule class methods are tabulated; please refer to the documentation for more details.
#
# | Method | Description |
# |--------|-------------|
# | [center_of_mass()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.center_of_mass "Go to Documentation") | Computes center of mass of molecule |
# | [molecular_charge()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.molecular_charge "Go to Documentation") | Gets the molecular charge |
# | [multiplicity()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.multiplicity "Go to Documentation") | Gets the total multiplicity |
# | [nuclear_repulsion_energy()](http://psicode.org/psi4manual/master/psi4api.html#psi4.core.Molecule.nuclear_repulsion_energy "Go to Documentation") | Computes the nuclear repulsion energy of the molecule |

# + [markdown] run_control={"frozen": false, "read_only": false}
# ## Example: Fitting Lennard-Jones Parameters from Potential Energy Scan
#
# In this example, we will compute and fit a potential energy curve for the Helium dimer.  To begin with, let's create a string representation for our He dimer in Z-Matrix format, with the variable `**R**` representing the distance between the He atoms. The stars surrounding the ``R`` aren't any special syntax, just a convenient marker for future substitution.

# + run_control={"frozen": false, "read_only": false}
# Define He Dimer
he_dimer = """
He
--
He 1 **R**
"""

# + [markdown] run_control={"frozen": false, "read_only": false}
# Now we can build a series of dimers with the He atoms at different separations, and compute the energy at each point:

# + run_control={"frozen": false, "read_only": false}
distances = [2.875, 3.0, 3.125, 3.25, 3.375, 3.5, 3.75, 4.0, 4.5, 5.0, 6.0, 7.0]
energies = []
for d in distances
    # Build a new molecule at each separation
    mol = psi4.geometry(replace(he_dimer,"**R**" => string(d)))

    # Compute the Counterpoise-Corrected interaction energy
    en = psi4.energy("MP2/aug-cc-pVDZ", molecule=mol, bsse_type="cp")

    # Place in a reasonable unit, Wavenumbers in this case
    en *= 219474.6

    # Append the value to our list
    push!(energies,en)
end

println("Finished computing the potential!")

# + [markdown] run_control={"frozen": false, "read_only": false}
# Next, we can use the [NumPy](http://www.numpy.org/) library to fit a curve to these points along the potential scan.  In this case, we will fit a Lennard-Jones potential.

# + run_control={"frozen": false, "read_only": false}
np   = pyimport("numpy")

# + run_control={"frozen": false, "read_only": false}
# Fit data in least-squares way to a -12, -6 polynomial
powers = [-12, -6]
x = np.power(np.reshape(np.array(distances),(-1, 1)), powers)
coeffs = np.linalg.lstsq(x, energies)[1]

# Build list of points
fpoints = np.reshape(np.linspace(2, 7, 50),(-1, 1))
fdata = np.power(fpoints, powers)

fit_energies = np.dot(fdata, coeffs)

# + [markdown] run_control={"frozen": false, "read_only": false}
# To visualize our results, we can use the [Matplotlib](http://matplotlib.org/) library.  Since we're working in a Jupyter notebook, we can also use the `%matplotlib inline` "magic" command so that the graphs will show up in the notebook itself (check [here](https://ipython.org/ipython-doc/3/interactive/magics.html) for a comprehensive list of magic commands).

# + run_control={"frozen": false, "read_only": false}
# %matplotlib inline
import PyPlot
plt = PyPlot

# + run_control={"frozen": false, "read_only": false}
plt.xlim((2, 7))  # X limits
plt.ylim((-7, 2))  # Y limits
plt.scatter(distances, energies)  # Scatter plot of the distances/energies
plt.plot(fpoints, fit_energies)  # Fit data
plt.plot([0,10], [0,0], "k-")  # Make a line at 0
# -


text:
using SimpleSolvers
using SimpleSolvers: initialize!, solver_step!
using Test

struct NonlinearSolverTest{T} <: NonlinearSolver end

test_solver = NonlinearSolverTest{Float64}()

@test_throws ErrorException config(test_solver)
@test_throws ErrorException status(test_solver)
@test_throws ErrorException initialize!(test_solver, rand(3))
@test_throws ErrorException solver_step!(test_solver)

function F!(f, x)
    f .= x.^2
end

function J!(g, x)
    g .= 0
    for i in eachindex(x)
        g[i,i] = 2x[i]
    end
end

for (Solver, kwarguments) in (
                (NewtonSolver, (linesearch = Static(),)),
                (NewtonSolver, (linesearch = Backtracking(),)),
                (NewtonSolver, (linesearch = Quadratic(),)),
                (NewtonSolver, (linesearch = Bisection(),)),
                (QuasiNewtonSolver, (linesearch = Static(),)),
                (QuasiNewtonSolver, (linesearch = Backtracking(),)),
                (QuasiNewtonSolver, (linesearch = Quadratic(),)),
                (QuasiNewtonSolver, (linesearch = Bisection(),)),
                (NLsolveNewton, NamedTuple()),
            )

    n = 1
    x = ones(n)
    y = zero(x)
    nl = Solver(x, y, F!; kwarguments...)

    @test config(nl) == nl.config
    @test status(nl) == nl.status

    solve!(x, nl)
    # println(status(nl))
    for x in nl.x
        @test x ≈ 0 atol=1E-7
    end

    x = ones(n)
    nl = Solver(x, y, F!; J! = J!, kwarguments...)
    solve!(x, nl)
    # println(status(nl))
    for x in nl.x
        @test x ≈ 0 atol=1E-7
    end
end


text:
"""
Newton-Raphson step of the steady state solver.
"""

# This line intentionally left blank.

"""
    NRData

A data structure that holds the necessary buffers and internal data for
performing a step of the Newton-Raphson algorithm

Internal use. Do not call directly.

"""
struct NRData
    sd::SolverData
    ss::SteadyStateData
    v_buffer::Vector{Float64}
    r_buffer::Vector{Float64}
    updated::BitArray{1}
    linesearch::Bool
end

@inline NRData(model::Model; kwargs...) = NRData(model, SolverData(model); kwargs...)
@inline NRData(model::Model, sd::SolverData; linesearch::Bool = false) = NRData(sd, model.sstate,
                    Array{Float64}(undef, sd.nvars), Array{Float64}(undef, sd.neqns),
                    trues(sd.nvars), linesearch)

"""
    step_nr!(x, dx, resid, J, nr::NRData; verbose=false)

Attempt a Newton-Raphson step.
The `nr` structure and the `dx` vector would be updated accordingly.
Vectors `x`, `resid` and the matrix `J` are read-only inputs.

!!! warning
    Internal function, do not call directly.
"""
function step_nr!(xx::AbstractArray{Float64,1}, dx::AbstractArray{Float64,1},
                    rr::AbstractArray{Float64,1}, jj::AbstractArray{Float64,2},
                    nr::NRData; verbose::Bool = false)
    ff = qr(jj, Val(true))
    rj = rank(ff.R)
    # nr.r_buffer .= ff.Q' * rr
    # dx[ff.p[1:rj]] = ff.R[1:rj,1:rj] \ nr.r_buffer[1:rj]
    nr.updated[ff.p[1:rj]] .= true
    nr.updated[ff.p[rj + 1:end]] .= false
    dx .= ff \ rr
    if verbose && any(.!nr.updated)
        problem_vars = join((ModelBaseEcon.ss_symbol(nr.ss, vi) for vi in findall(nr.sd.solve_var)[.!nr.updated]), ", ")
        @warn "Unable to update $(problem_vars)."
    end
    nf = norm(rr)
    if nr.linesearch
        # the Armijo rule: C.T.Kelly, Iterative Methods for Linear and Nonlinear Equations, ch.8.1, p.137
        λ = 1.0
        α = 1e-4
        σ = 0.5
        while λ > 0.00001
            nr.v_buffer .= xx - λ .* dx
            nrb2 = 0.0
            try
                global_SS_R!(nr.r_buffer, nr.v_buffer, nr.sd)
                nrb2 = norm(nr.r_buffer)
            catch e
                nrb2 = Inf
            end
            if nrb2 < (1.0 - α * λ) * nf
                if verbose && λ < 1.0
                    @info "Linesearch success with λ = $λ."
                end
                dx .*= λ
                break
            end
            λ = σ * λ
        end
        if verbose && λ <= 0.00001
            @warn "Linesearch failed."
        end
    end
    return nothing
end
@assert precompile(step_nr!, (Array{Float64,1}, Array{Float64,1}, Array{Float64,1}, Array{Float64,2}, NRData))


text:
module Plugins

end


text:
#!/usr/bin/env julia
#=
#
# at Package Manager
# * add Fire
#
# Help: julia ${path_to_code}/basic-cli.jl --help
=#

using Fire

"""
This will come in help for 'greet'
julia \${path_to_code}/basic-cli.jl greet ABK
"""
@main function greet(name)
  println("Hey, $name")
end

"""add all provided numbers"""
@main function add(num::Integer...)
  println("Adding: $num")
  println(sum(num))
end

@main function sysinfo(entity)
  lowercase_entity = lowercase(entity)
  if lowercase_entity == "cpu"
    Sys.cpu_summary()
  elseif lowercase_entity == "mem"
    println("Total Mem: $(Sys.total_memory() / (2*1024)) | Free Mem: $(Sys.free_memory() / (2*1024))")
  else
    println("[error] available options: cpu | mem")
  end
end


text:
function EventEnablingPredicateFilter(arg0::EventDetector, arg1::EnablingPredicate)
    return EventEnablingPredicateFilter((EventDetector, EnablingPredicate), arg0, arg1)
end

function g(obj::EventEnablingPredicateFilter, arg0::SpacecraftState)
    return jcall(obj, "g", jdouble, (SpacecraftState,), arg0)
end

function init(obj::EventEnablingPredicateFilter, arg0::SpacecraftState, arg1::AbsoluteDate)
    return jcall(obj, "init", void, (SpacecraftState, AbsoluteDate), arg0, arg1)
end


text:
#
#  Denoising a cyclic (S1-valued) Signal with first and second order TV.
#  using the Cyclic Proximal Point Algorithm
#
# This example was in a similar form part of the paper
#
# > Bergmann, R., Laus, F., Steidl, G., Weinmann, A.:
# > _Second Order Differences of Cyclic Data and Applications in Variational Denoising_
# > SIAM J. Imaging Sci., 7(4), 2916–2953, 2014.
# > doi: [10.1137/140969993](https://dx.doi.org/10.1137/140969993)
# > arXiv: [1405.5349](https://arxiv.org/abs/1405.5349)
#
using Manopt, Manifolds, Plots
#
# Settings
results_folder = joinpath(@__DIR__, "Signal_TV")
experiment_name = "S1_Signal_TV12_CPPA"
export_orig = true
export_result = true
!isdir(results_folder) && mkdir(results_folder)
#
# Parameters
n = 500
σ = 0.2
α = 0.5
β = 1.0
#
# Coloes
data_color = RGBA{Float64}(colorant"#BBBBBB")
result_color = RGBA{Float64}(colorant"#EE7733") # data Color: Tol Vibrant Orange
orig_color = RGBA{Float64}(colorant"#33BBEE") # tangent vector: Tol Vibrant Teal
#
# Manifolds and Data
M = Circle()
N = PowerManifold(M, n)
f = artificial_S1_signal(n)
xCompare = f
fn = exp.(Ref(M), f, random_tangent.(Ref(M), f, Val(:Gaussian), σ))
data = fn
t = range(0.0, 1.0; length=n)

if export_orig
    scene = scatter(
        t,
        f;
        markersize=2,
        markercolor=data_color,
        markerstrokecolor=data_color,
        lab="original",
    )
    scatter!(
        scene,
        t,
        fn;
        markersize=2,
        markercolor=orig_color,
        markerstrokecolor=orig_color,
        lab="noisy",
    )
    yticks!(
        [-π, -π / 2, 0, π / 2, π],
        [raw"$-\pi$", raw"$-\frac{\pi}{2}$", raw"$0$", raw"$\frac{\pi}{2}$", raw"$\pi$"],
    )
    png(scene, joinpath(results_folder, experiment_name * "-original.png"))
end
#
# Setup and Optimize
F = (N, x) -> costL2TVTV2(N, data, α, β, x)
proxes = (
    (N, λ, x) -> prox_distance(N, λ, data, x, 2),
    (N, λ, x) -> prox_TV(N, α * λ, x),
    (N, λ, x) -> prox_TV2(N, β * λ, x),
)

o = cyclic_proximal_point(
    N,
    F,
    proxes,
    data;
    λ=i -> π / (2 * i),
    debug=[
        :Iteration,
        " | ",
        DebugProximalParameter(),
        " | ",
        :Cost,
        " | ",
        :Change,
        "\n",
        1000,
        :Stop,
    ],
    record=[:Iteration, :Cost, :Change, :Iterate],
    return_options=true,
)
fR = get_solver_result(o)
r = get_record(o)
#
# Result
if export_result
    scene = scatter(
        t,
        f;
        markersize=2,
        markercolor=data_color,
        markerstrokecolor=data_color,
        lab="original",
    )
    scatter!(
        scene,
        t,
        fR;
        markersize=2,
        markercolor=result_color,
        markerstrokecolor=result_color,
        lab="reconstruction",
    )
    yticks!(
        [-π, -π / 2, 0, π / 2, π],
        [raw"$-\pi$", raw"$-\frac{\pi}{2}$", raw"$0$", raw"$\frac{\pi}{2}$", raw"$\pi$"],
    )
    png(scene, joinpath(results_folder, experiment_name * "-result.png"))
end

print("MSE (input):  ", 1 / n * distance(N, xCompare, data)^2, "\n")
print("MSE (result): ", 1 / n * distance(N, xCompare, fR)^2, "\n")


text:
function deferred_acceptance(apply_prefs::Array{Int64}, accept_prefs::Array{Int64}, caps=ones(Int64, length(accept_prefs[1, :])))
    num_of_apply = length(accept_prefs[:,1])-1
    num_of_accept = length(accept_prefs[1,:])

    ranking = zeros(Int64, num_of_apply+1, num_of_accept)

    da = zeros(Int64, num_of_apply+1, num_of_accept+1)

    num_of_da = zeros(Int64, num_of_accept)

    for i in 1:num_of_apply+1
        if i == num_of_apply+1
            search = 0
        else
            search = i
        end

        for j in 1:num_of_accept
            ranking[i, j] = findfirst(accept_prefs[:,j], search)
        end
    end

    rem = [1 for i in 1:length(apply_prefs[1,:])]
    depth = ones(Int64, length(accept_prefs[:,1])-1)

    b=1
    while true
        apply = findfirst(rem, 1)
        if apply == 0
            break
        end

        d = depth[apply]
        if d >= num_of_apply
            rem[apply] = 0
        else

            applied = apply_prefs[d, apply]

            if applied == 0
                rem[apply] = 0
            else
                my_rank = ranking[apply, applied]
                if num_of_da[applied] < caps[applied]
                    if my_rank < ranking[end, applied]
                        da[apply, applied] = 1
                        num_of_da[applied] += 1
                        rem[apply] = 0
                    else
                        depth[apply] += 1
                    end
                else
                    list_of_da = findn(da[:, applied])

                    da_rank = Array(Int64, length(list_of_da))

                    for i in 1:length(list_of_da)
                        da_rank[i] = ranking[list_of_da[i], applied]
                    end

                    worst_st_rank = maximum(da_rank)
                    worst_st = list_of_da[findfirst(da_rank, worst_st_rank)]

                    if my_rank < worst_st_rank
                        rem[apply] = 0
                        da[apply, applied] = 1
                        da[worst_st, applied] = 0
                        rem[worst_st] = 1
                        depth[worst_st] += 1
                    else
                        depth[apply] += 1
                    end

                end
            end
        end
    end

    accept_matched = Array(Int64, 0)
    apply_matched = zeros(Int64, num_of_apply)

    for i in 1:num_of_accept
        das = findn(da[:, i])
        append!(das, zeros(Int64, caps[i] - length(das)))
        append!(accept_matched, das)
        for d in das
            if d != 0
                apply_matched[d] = i
            end
        end

    end

    indptr = Array(Int64, num_of_accept+1)
    indptr[1] = 1
    for i in 1:num_of_accept
        indptr[i+1] = indptr[i] + caps[i]
    end
    return apply_matched, accept_matched, indptr
end

max_stars_repo_name: NlGG/Economics

max_issues_repo_name: NlGG/Economics

max_forks_repo_name: NlGG/Economics


text:
@test NewModule(:Main, nothing).name == :Main
@test NewTypeVar(:T)  != nothing
@test NewDataType(NewTypeName(:DT, MainModule), nothing, NewEmpty, true, false, false) != nothing


text:
"""
    SavedValues{tType<:Real, savevalType}

A struct used to save values of the time in `t::Vector{tType}` and
additional values in `saveval::Vector{savevalType}`.
"""
struct SavedValues{tType, savevalType}
    t::Vector{tType}
    saveval::Vector{savevalType}
end

"""
    SavedValues(tType::DataType, savevalType::DataType)

Return `SavedValues{tType, savevalType}` with empty storage vectors.
"""
function SavedValues(::Type{tType}, ::Type{savevalType}) where {tType,savevalType}
    SavedValues{tType, savevalType}(Vector{tType}(), Vector{savevalType}())
end

function Base.show(io::IO, saved_values::SavedValues)
    tType = eltype(saved_values.t)
    savevalType = eltype(saved_values.saveval)
    print(io, "SavedValues{tType=", tType, ", savevalType=", savevalType, "}",
                "\nt:\n", saved_values.t, "\nsaveval:\n", saved_values.saveval)
end

@recipe function plot(saved_values::SavedValues)
    DiffEqArray(saved_values.t, saved_values.saveval)
end

mutable struct SavingAffect{SaveFunc, tType, savevalType, saveatType, saveatCacheType}
    save_func::SaveFunc
    saved_values::SavedValues{tType, savevalType}
    saveat::saveatType
    saveat_cache::saveatCacheType
    save_everystep::Bool
    save_start::Bool
    save_end::Bool
    saveiter::Int
end

function (affect!::SavingAffect)(integrator,force_save = false)

    just_saved = false
    # see OrdinaryDiffEq.jl -> integrator_utils.jl, function savevalues!
    while !isempty(affect!.saveat) && integrator.tdir*top(affect!.saveat) <= integrator.tdir*integrator.t # Perform saveat
        affect!.saveiter += 1
        curt = pop!(affect!.saveat) # current time
        if curt != integrator.t # If <t, interpolate
            if typeof(integrator) <: ODEIntegrator
                # Expand lazy dense for interpolation
                DiffEqBase.addsteps!(integrator)
            end
            if !DiffEqBase.isinplace(integrator.sol.prob)
                curu = integrator(curt)
            else
                curu = first(get_tmp_cache(integrator))
                integrator(curu,curt) # inplace since save_func allocates
            end
            copyat_or_push!(affect!.saved_values.t, affect!.saveiter, curt)
            copyat_or_push!(affect!.saved_values.saveval, affect!.saveiter,
                            affect!.save_func(curu, curt, integrator),Val{false})
        else # ==t, just save
            just_saved = true
            copyat_or_push!(affect!.saved_values.t, affect!.saveiter, integrator.t)
            copyat_or_push!(affect!.saved_values.saveval, affect!.saveiter, affect!.save_func(integrator.u, integrator.t, integrator),Val{false})
        end
    end
    if !just_saved &&
        affect!.save_everystep || force_save ||
        (affect!.save_end && integrator.t == integrator.sol.prob.tspan[end])

        affect!.saveiter += 1
        copyat_or_push!(affect!.saved_values.t, affect!.saveiter, integrator.t)
        copyat_or_push!(affect!.saved_values.saveval, affect!.saveiter, affect!.save_func(integrator.u, integrator.t, integrator),Val{false})
    end
    u_modified!(integrator, false)
end

function saving_initialize(cb, u, t, integrator)
    if cb.affect!.saveiter != 0
        if integrator.tdir > 0
            cb.affect!.saveat = BinaryMinHeap(cb.affect!.saveat_cache)
        else
            cb.affect!.saveat = BinaryMaxHeap(cb.affect!.saveat_cache)
        end
        cb.affect!.saveiter = 0
    end
    cb.affect!.save_start && cb.affect!(integrator)
end

"""
    SavingCallback(save_func, saved_values::SavedValues;
                    saveat=Vector{eltype(saved_values.t)}(),
                    save_everystep=isempty(saveat),
                    save_start = save_everystep || isempty(saveat) || saveat isa Number,
                    save_end = save_everystep || isempty(saveat) || saveat isa Number,
                    tdir=1)

A `DiscreteCallback` applied after every step, saving the time `t` and the value
of `save_func(u, t, integrator)` in `saved_values`.

If `save_everystep`, every step of the integrator is saved.
If `saveat` is specified, the values are saved at the given times, using
interpolation if necessary.
If the time `tdir` direction is not positive, i.e. `tspan[1] > tspan[2]`,
`tdir = -1` has to be specified.
"""
function SavingCallback(save_func, saved_values::SavedValues;
                        saveat=Vector{eltype(saved_values.t)}(),
                        save_everystep=isempty(saveat),
                        save_start = save_everystep || isempty(saveat) || saveat isa Number,
                        save_end = save_everystep || isempty(saveat) || saveat isa Number,
                        tdir=1)
    # saveat conversions, see OrdinaryDiffEq.jl -> integrators/type.jl
    saveat_vec = collect(saveat)
    if tdir > 0
        saveat_internal = BinaryMinHeap(saveat_vec)
    else
        saveat_internal = BinaryMaxHeap(saveat_vec)
    end
    affect! = SavingAffect(save_func, saved_values, saveat_internal, saveat_vec, save_everystep, save_start, save_end, 0)
    condtion = (u, t, integrator) -> true
    DiscreteCallback(condtion, affect!;
                     initialize = saving_initialize,
                     save_positions=(false,false))
end

export SavingCallback, SavedValues


text:
# Here in this file I will try to translate the example file v2 provided by
# MAGMA official docs


text:
mutable struct JenksResult
    n::Int                          # number of classes / intervals
    ndata::Int                      # number of data points
    breaks::Array{Int,1}            # break indices in data array
    bounds::Array{Float64,1}        # lower bound of each class
    centres::Array{Float64,1}       # centres of intervals
    n_in_class::Array{Int,1}        # number of data points per class
    ARE::Float64                    # Average rounding error (L1 norm)
    GVF::Float64                    # Goodness of variance fit
    AREhistory::Array{Float64,1}    # Rounding error for each iteration
    GVFhistory::Array{Float64,1}    # GVF for each iteration
    errornorm::Int                  # 1 for L1, 2 for L2
    maxiter::Int                    # number of iterations
    flux::Float64                   # initial fraction of class size that is moved to adjacent class
    fluxes::Array{Float64,1}        # flux for each boundary between classes
    dt::Float64                     # time for all iterations (excl initialisation)
end

function JenksResult(n::Int,data::Array{T,1};
                    maxiter::Int=0,
                    errornorm::Int=1,
                    flux::Float64=0.1) where {T<:AbstractFloat}

    ndata = length(data)

    # classification results
    breaks = Array{Int,1}(undef,n+1)
    centres = Array{Float64,1}(undef,n)
    bounds = Array{Float64,1}(undef,n+1)
    n_in_class = Array{Int,1}(undef,n)

    # error measures
    ARE = 0.0
    GVF = 0.0
    AREhistory = Array{Float64,1}(undef,maxiter+1)
    GVFhistory = Array{Float64,1}(undef,maxiter+1)

    # flux adjustment
    fluxes = fill(flux,n-1)

    # time measures
    dt = 0.0

    JenksResult(n,ndata,breaks,bounds,centres,n_in_class,
                ARE,GVF,AREhistory,GVFhistory,
                errornorm,maxiter,flux,fluxes,dt)
end

"""Calculate the class centres from the class break indices. Assumes data in X to be sorted."""
function Breaks2Centres!(JR::JenksResult,X::Array{T,1}) where {T<:AbstractFloat}
    for i in 1:JR.n
        JR.centres[i] = (X[JR.breaks[i]] + X[JR.breaks[i+1]-1])/2
    end
end

"""Calculate the class sizes (number of data points per class/interval) from break indices."""
function Breaks2ClassSize!(JR::JenksResult)
    for i in 1:JR.n
        JR.n_in_class[i] = JR.breaks[i+1] - JR.breaks[i]
    end
end

"""Calculate the lower bounds from break indices. Assumes data in X to be sorted."""
function Breaks2Bounds!(JR::JenksResult,X::Array{T,1}) where {T<:AbstractFloat}
    for i in 1:JR.n
        JR.bounds[i] = X[JR.breaks[i]]
    end
    JR.bounds[end] = X[end]
end


text:
using DynamicGrids, Setfield, Test, LinearAlgebra, StaticArrays
import DynamicGrids: SimData, Extent, WritableGridData,
       radius, dest, hoodsize, neighborhoodkey, _buffer

@testset "neighbors" begin
    init = [0 0 0 1 1 1
            1 0 1 1 0 1
            0 1 1 1 1 1
            0 1 0 0 1 0
            0 0 0 0 1 1
            0 1 0 1 1 0]

    buf1 = [0 0 0
            0 1 0
            0 0 0]
    buf2 = [1 1 1
            1 0 1
            1 1 1]
    buf3 = [1 1 1
            0 0 1
            0 0 1]

    @testset "Moore" begin
        moore = Moore{1}(init[1:3, 1:3])
        @test _buffer(moore) == init[1:3, 1:3]
        @test hoodsize(moore) == 3
        @test moore[2, 2] == 0
        @test length(moore) == 8
        @test eltype(moore) == Int
        @test neighbors(moore) isa Base.Generator
        @test collect(neighbors(moore)) == [0, 1, 0, 0, 1, 0, 1, 1]
        @test sum(moore) == sum(neighbors(moore)) == 4
        @test Tuple(offsets(moore)) == ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0),
                                        (0, 1), (1, -1), (1, 0), (1, 1))

        moore1 = @set moore._buffer = buf1
        @test moore1._buffer == buf1
        moore2 = DynamicGrids._setbuffer(moore, buf2)
        @test moore2._buffer == buf2

        @test sum(Moore{1}(buf1)) == 0
        @test sum(Moore{1}(buf2)) == 8
        @test sum(Moore{1}(buf3)) == 5
    end
    @testset "Window" begin
        @test Window{1}() == Window(1) == Window(zeros(3, 3))
        window = Window{1}(init[1:3, 1:3])
        @test _buffer(window) == init[1:3, 1:3]
        @test hoodsize(window) == 3
        @test window[2, 2] == 0
        @test length(window) == 9
        @test eltype(window) == Int
        @test neighbors(window) isa Array
        @test neighbors(window) == _buffer(window)
        @test sum(window) == sum(neighbors(window)) == 4
        @test Tuple(offsets(window)) == ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0),
                                        (0, 1), (1, -1), (1, 0), (1, 1))

        window1 = @set window._buffer = buf1
        @test window1._buffer == buf1
        window2 = DynamicGrids._setbuffer(window, buf2)
        @test window2._buffer == buf2

        @test sum(Window{1}(buf1)) == 1
        @test sum(Window{1}(buf2)) == 8
        @test sum(Window{1}(buf3)) == 5
    end
    @testset "VonNeumann/Positional" begin
        vonneumann = VonNeumann(1, init[1:3, 1:3])
        @test offsets(vonneumann) == ((0, -1), (-1, 0), (1, 0), (0, 1))
        @test _buffer(vonneumann) == init[1:3, 1:3]
        @test hoodsize(vonneumann) == 3
        @test vonneumann[2, 1] == 1
        @test length(vonneumann) == 4
        @test eltype(vonneumann) == Int
        @test neighbors(vonneumann) isa Base.Generator
        @test collect(neighbors(vonneumann)) == [1, 0, 1, 1]
        @test sum(neighbors(vonneumann)) == 3
        vonneumann2 = VonNeumann(2)
        @test offsets(vonneumann2) ==
           ((0, -2), (-1, -1), (0, -1), (1, -1),
             (-2 , 0), (-1, 0), (1, 0), (2, 0),
             (-1, 1), (0, 1), (1, 1), (0, 2))

        vonneumann1 = @set vonneumann._buffer = buf2
        @test vonneumann1._buffer == buf2
        vonneumann2 = DynamicGrids._setbuffer(vonneumann, buf3)
        @test vonneumann2._buffer == buf3

        @test sum(VonNeumann(1, buf1)) == 0
        @test sum(VonNeumann(1, buf2)) == 4
        @test sum(VonNeumann(1, buf3)) == 2
    end
    @testset "Layered/Positional" begin
        buf = [0 1 0 0 1
               0 0 1 0 0
               0 0 0 1 1
               0 0 1 0 1
               1 0 1 0 1]
        custom1 = Positional(((-1,-1), (2,-2), (2,2), (-1,2), (0,0)), buf)
        custom2 = Positional(((-1,-1), (0,-1), (1,-1), (2,-1), (0,0)), buf)
        layered = LayeredPositional(
            (Positional((-1,1), (-2,2)), Positional((1,2), (2,2))), buf)

        @test neighbors(custom1) isa Base.Generator
        @test collect(neighbors(custom1)) == [0, 1, 1, 0, 0]

        @test sum(custom1) == 2
        @test sum(custom2) == 0
        @test sum(layered) == (1, 2)
        @test offsets(layered) == (((-1, 1), (-2, 2)), ((1, 2), (2, 2)))

        layered1 = @set layered._buffer = 2buf
        @test layered1._buffer == 2buf
        layered2 = DynamicGrids._setbuffer(layered, 3buf)
        @test layered2._buffer == 3buf
    end
    @testset "neighbors works on rule" begin
        rule = Life(;neighborhood=Moore{1}([0 1 1; 0 0 0; 1 1 1]))
        @test sum(neighbors(rule)) == 5
    end

end

@testset "Kernel" begin
    buf = reshape(1:9, 3, 3)
    @testset "Window" begin
        mat = zeros(3, 3)
        @test Kernel(mat) == Kernel(Window(1), mat)
        @test_throws ArgumentError Kernel(Window(2), mat)
        k = Kernel(Window{1,9,typeof(buf)}(buf), SMatrix{3,3}(reshape(1:9, 3, 3)))
        @test dot(k) == sum((1:9).^2)
        @test neighbors(k) == reshape(1:9, 3, 3)
        @test collect(offsets(k)) ==
            [(-1, -1) (0, -1) (1, -1)
             (-1, 0)  (0, 0)  (1, 0)
             (-1, 1)  (0, 1)  (1, 1)]
        @test collect(positions(k, (2, 2))) ==
             [(1, 1) (2, 1) (3, 1)
              (1, 2) (2, 2) (3, 2)
              (1, 3) (2, 3) (3, 3)]

    end
    @testset "Moore" begin
        k = Kernel(Moore{1,8,typeof(buf)}(buf), (1:4..., 6:9...))
        @test dot(k) == sum((1:4).^2) + sum((6:9).^2)
    end
    @testset "Positional" begin
        off = ((0,-1),(-1,0),(1,0),(0,1))
        hood = Positional{1,4,typeof(off),typeof(buf)}(off, buf)
        k = Kernel(hood, 1:4)
        @test dot(k) == 1 * 2 + 2 * 4 + 3 * 6 + 4 * 8
    end
end

struct TestNeighborhoodRule{R,W,N} <: NeighborhoodRule{R,W}
    neighborhood::N
end
DynamicGrids.applyrule(data, rule::TestNeighborhoodRule, state, index) = state

struct TestSetNeighborhoodRule{R,W,N} <: SetNeighborhoodRule{R,W}
    neighborhood::N
end
function DynamicGrids.applyrule!(
    data, rule::TestSetNeighborhoodRule{R,Tuple{W1,}}, state, index
) where {R,W1}
    add!(data[W1], state[1], index...)
end

@testset "neighborhood rules" begin
    ruleA = TestSetNeighborhoodRule{:a,:a}(Moore{3}())
    ruleB = TestSetNeighborhoodRule{Tuple{:b},Tuple{:b}}(Moore{2}())
    @test offsets(ruleA) isa Base.Generator
    @test positions(ruleA, (1, 1)) isa Base.Generator
    @test neighborhood(ruleA) == Moore{3}()
    @test neighborhood(ruleB) == Moore{2}()
    @test neighborhoodkey(ruleA) == :a
    @test neighborhoodkey(ruleB) == :b
    ruleA = TestNeighborhoodRule{:a,:a}(Moore{3}())
    ruleB = TestNeighborhoodRule{Tuple{:b},Tuple{:b}}(Moore{2}())
    @test neighbors(ruleA) isa Base.Generator
    @test offsets(ruleA) isa Base.Generator
    @test neighborhood(ruleA) == Moore{3}()
    @test neighborhood(ruleB) == Moore{2}()
    @test neighborhoodkey(ruleA) == :a
    @test neighborhoodkey(ruleB) == :b
    @test Tuple(offsets(ruleB)) ===
        ((-2, -2), (-2, -1), (-2, 0), (-2, 1), (-2, 2), (-1, -2), (-1, -1), (-1, 0),
         (-1, 1), (-1, 2), (0, -2), (0, -1), (0, 0), (0, 1), (0, 2), (1, -2), (1, -1),
         (1, 0), (1, 1), (1, 2), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2))
    @test Tuple(positions(ruleB, (10, 10))) ==
        ((8, 8), (8, 9), (8, 10), (8, 11), (8, 12), (9, 8), (9, 9), (9, 10), (9, 11),
         (9, 12), (10, 8), (10, 9), (10, 10), (10, 11), (10, 12), (11, 8), (11, 9),
         (11, 10), (11, 11), (11, 12), (12, 8), (12, 9), (12, 10), (12, 11), (12, 12))
end

@testset "radius" begin
    init = (a=[1.0 2.0], b=[10.0 11.0])
    ruleA = TestNeighborhoodRule{:a,:a}(Moore{3}())
    ruleB = TestSetNeighborhoodRule{Tuple{:b},Tuple{:b}}(Moore{2}())
    ruleset = Ruleset(ruleA, ruleB)
    @test radius(ruleA) == 3
    @test radius(ruleB) == 2
    @test radius(ruleset) == (a=3, b=2)
    @test radius(Ruleset()) == NamedTuple()

    output = ArrayOutput(init; tspan=1:3)
    sim!(output, ruleset)
    # TODO make sure 2 radii can coexist
end

@testset "Positional" begin
    buf = reshape(2:10, 3, 3)
    hood = Positional(((-1, -1), (1, 1)), buf)
    @test Tuple(neighbors(hood)) == (2, 10)
    @test offsets(hood) == ((-1, -1), (1, 1))
    @test Tuple(positions(hood, (2, 2))) == ((1, 1), (3, 3))
    @set! hood.offsets = ((-5, -5), (5, 5))
    @test offsets(hood) == hood.offsets == ((-5, -5), (5, 5))
end

@testset "LayeredPositional" begin
    lhood = LayeredPositional(
        Positional(((-1, -1), (1, 1)), ), Positional(((-2, -2), (2, 2)), )
    )
    @test offsets(lhood) == (((-1, -1), (1, 1)), ((-2, -2), (2, 2)))
    @test collect.(collect(positions(lhood, (1, 1)))) ==
        [[(0, 0), (2, 2)],
         [(-1, -1), (3, 3)]]

    buf = reshape(1:25, 5, 5)
    lhood_buf = DynamicGrids._setbuffer(lhood, buf)
    @test lhood_buf._buffer == lhood_buf.layers[1]._buffer ===
          lhood_buf.layers[2]._buffer === buf
    @test map(radius, lhood_buf.layers) == (2, 2)
    @test map(Tuple, neighbors(lhood_buf)) == ((7, 19), (1, 25))
end


text:
using Distributions, Turing, Plots, StatsPlots, MCMCChains, DataFrames

function BiExponential(theta1::Real, theta2::Real, p::Real)
    return MixtureModel(Exponential[
                             Exponential(theta1),
                             Exponential(theta2),
                            ], [p, 1 - p])
end

function generate_samples(theta1::Real, theta2::Real, p::Real, n::Int64)
    m = BiExponential(theta1, theta2, p)
    rand(m, n)
end

@model bi_exponential(samples::Array{Real, 1}) = begin
    theta1 ~ Gamma(1, 10)
    theta2 ~ truncated(Normal(), theta1, Inf)
    p ~ Beta(1, 1)
    m = BiExponential(theta1, theta2, p)
    for i in 1:length(samples)
        samples[i] ~ m
    end
end

function svf(x::Array{Real, 1}, digits::Int64)
    x = round.(x, digits=digits)
    unqx = unique(x) |> sort
    counts = map(x_ -> sum(x .== x_), unqx)
    DataFrame(x = unqx, count = counts, svr = 1. .- cumsum(counts ./ sum(counts)))
end

function predict(chains::Chains, x::Array{Real, 1})
    theta1 = mean(chains["theta1"])
    theta2 = mean(chains["theta2"])
    p = mean(chains["p"])
    becdf = cdf(BiExponential(theta1, theta2, p), x)
    return 1. .- becdf
end

samples = Array{Real, 1}(generate_samples(0.1, 5, 0.9, 2000))
chains = sample(bi_exponential(samples), NUTS(), MCMCThreads(), 1000, 4)

svf_samples = svf(samples, 2)
pred_samples = predict(chains, Array{Real, 1}(svf_samples.x))
scatter(svf_samples.x, log10.(svf_samples.svr), markersize=2)
plot!(svf_samples.x, log10.(pred_samples))
ylims!(-3, 0)


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule cryptopp_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("cryptopp")
JLLWrappers.@generate_main_file("cryptopp", UUID("9ced465f-3d53-56cd-9803-9d1649b5e965"))
end  # module cryptopp_jll


text:
using CompileBot

println("tests infer benchmark")

snoop_bench(BotConfig("TestPackage0", tmin =0.0))


text:
module GDAL

using PROJ_jll
using GDAL_jll
using CEnum
using NetworkOptions

const Ctm = Base.Libc.TmStruct

include("common.jl")

include("gdal_h.jl")
include("gdal_alg.jl")
include("gdal_vrt.jl")
include("gdal_utils.jl")
include("gdalwarper.jl")

include("ogr_api.jl")
include("ogr_srs_api.jl")
include("ogr_core.jl")

include("cpl_error.jl")
include("cpl_conv.jl")
include("cpl_minixml.jl")
include("cpl_progress.jl")
include("cpl_virtualmem.jl")
include("cpl_vsi.jl")

include("error.jl")

const GDALVERSION = Ref{VersionNumber}()
const GDAL_DATA = Ref{String}()
const PROJ_LIB = Ref{String}()

function __init__()
    # register custom error handler
    funcptr = @cfunction(gdaljl_errorhandler, Ptr{Cvoid}, (CPLErr, Cint, Cstring))
    cplseterrorhandler(funcptr)

    # get GDAL version number
    versionstring = gdalversioninfo("RELEASE_NAME")
    GDALVERSION[] = VersionNumber(versionstring)

    # set GDAL_DATA location, this overrides setting the environment variable
    GDAL_DATA[] = joinpath(GDAL_jll.artifact_dir, "share", "gdal")
    cplsetconfigoption("GDAL_DATA", GDAL_DATA[])

    # set path to CA certificates
    ca_path = @static if VERSION >= v"1.6"
        ca_roots()
    else
        ca_roots_path()
    end
    if ca_path !== nothing
        cplsetconfigoption("CURL_CA_BUNDLE", ca_path)
    end

    # set PROJ_LIB location, this overrides setting the environment variable
    PROJ_LIB[] = joinpath(PROJ_jll.artifact_dir, "share", "proj")
    osrsetprojsearchpaths([PROJ_LIB[]])
end

end # module


text:
function molecule(::Molecule"ethanol")
  return """
C	1.1615830	-0.4067550	0.0000000
C	0.0000000	0.5527180	0.0000000
O	-1.1871140	-0.2128600	0.0000000
H	-1.9324340	0.3838170	0.0000000
H	2.1028600	0.1358400	0.0000000
H	1.1223470	-1.0398290	0.8811340
H	1.1223470	-1.0398290	-0.8811340
H	0.0561470	1.1935530	0.8808960
H	0.0561470	1.1935530	-0.8808960
  """
end


text:
struct LexerConf
    tokens::Array{TerminalDef}
    ignore::Array{String}
    postlex  #Call the "process" method on this with the stream after lexing
    callbacks::Dict{String,Function}
    g_regex_flags::String
    skip_validation::Bool
end

LexerConf(tokens;ignore=(),postlex=nothing,callbacks=Dict(),g_regex_flags=0,skip_validation=false) = begin
    # convert Lark/Python integer flags back to characters
    new_flags = ""
    if (g_regex_flags & 256) > 0
        new_flags *='a'  #ascii
    end
    if (g_regex_flags & 2 ) > 0
        new_flags *= 'i'  #ignore case
    end
    if (g_regex_flags & 8 ) > 0
        new_flags *= 'm'   #multiline
    end
    if (g_regex_flags & 16) > 0
        new_flags *= 's'   #dotall
    end
    if (g_regex_flags & 64) > 0
        new_flags *= 'x'   #verbose
    end
    LexerConf(tokens,ignore,postlex,callbacks,new_flags,skip_validation)
end

reconfigure(l::LexerConf,tokens) = begin
    LexerConf(tokens,l.ignore,l.postlex,l.callbacks,l.g_regex_flags,l.skip_validation)
end

struct ParserConf
    rules::Array{Rule}
    callbacks::Dict{Rule,Function}
    start::Array{String} #The names of the starting rule
end


text:
@enum Int64 S I R

include("tree_helpers.jl")

# Note: I usually define a helper function to more easily add edges
mutable struct Person
    SIR
    ID
    is_infected::Bool
    neighbors::Vector{Person}
    Person(SIR, ID) = new(SIR, ID, false, [])
end

function infect(p::Person)
    if p.SIR == S
        p.is_infected = true
        p.SIR = I
        for neighbor in p.neighbors
            infect(neighbor)
        end
    else
        return
    end

end

function sir_test()
    people = [Person(S,i) for i = 1:6]

    # setting the 4th person to be recovered
    people[4].SIR = R

    # this is where the helper function would allow us to specify only 1 dir!
    people[1].neighbors = [people[2], people[3]]
    people[2].neighbors = [people[3]]
    people[3].neighbors = [people[4], people[5]]
    people[5].neighbors = [people[6]]

    infect(people[1])

    for person in people
        println("infection status: ", person.is_infected)
    end

    return people
end


text:
if VERSION >= v"0.7"
  using Test
else
  using Base.Test
  include("random.jl")
end
include("../src/define_constants.jl")
include("loglinspace.jl")
include("test_cel_vec.jl")
include("test_s2_gradient.jl")
include("test_transit_poly.jl")

include("test_IJv_derivative.jl")
#include("test_sn_jacobian2.jl")

skip_plots = false
include("test_transit_poly_gradient2.jl")

max_stars_repo_name: tagordon/limbdark

max_issues_repo_name: tagordon/limbdark

max_forks_repo_name: tagordon/limbdark


text:
using Test
using QuantumOptics, CollectiveSpins

@testset "squeezing" begin

cs = CollectiveSpins

function compare_squeezing(;phi=0., theta=0., N=1., χT=0., axis=[1.,0.,0.])
    state0 = cs.mpc.blochstate(phi, theta, N)
    state_squeezed = cs.mpc.squeeze(axis, χT, state0)
    Ψ₀ = CollectiveSpins.quantum.blochstate(phi,theta,N)
    ρ₀ = Ψ₀⊗dagger(Ψ₀)
    ρ_squeezed = cs.quantum.squeeze(axis, χT, ρ₀)
    return QuantumOptics.tracedistance(ρ_squeezed, cs.mpc.densityoperator(state_squeezed))
end

td = compare_squeezing(phi=0., theta=0., N=2, χT=2.5, axis=[1.,0,0])
@test td < 1e-5

td = compare_squeezing(phi=0.7, theta=1.34, N=2, χT=2.5, axis=[1.,3.,2.5])
@test td < 1e-5

td = compare_squeezing(phi=0.7, theta=1.34, N=5, χT=0.5, axis=[5.,2.,1.])
@test td < 0.1

td = compare_squeezing(phi=0., theta=0., N=5, χT=2.5, axis=[0.,0.,1.])
@test td < 1e-12

end # testset


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

function calc_error_norms(func, u, t, analyzer,
                          mesh::ParallelP4estMesh{3}, equations,
                          initial_condition, dg::DGSEM, cache, cache_analysis)
  @unpack vandermonde, weights = analyzer
  @unpack node_coordinates, inverse_jacobian = cache.elements
  @unpack u_local, u_tmp1, u_tmp2, x_local, x_tmp1, x_tmp2, jacobian_local, jacobian_tmp1, jacobian_tmp2 = cache_analysis

  # Set up data structures
  l2_error   = zero(func(get_node_vars(u, equations, dg, 1, 1, 1, 1), equations))
  linf_error = copy(l2_error)
  volume = zero(real(mesh))

  # Iterate over all elements for error calculations
  for element in eachelement(dg, cache)
    # Interpolate solution and node locations to analysis nodes
    multiply_dimensionwise!(u_local, vandermonde, view(u,                :, :, :, :, element), u_tmp1, u_tmp2)
    multiply_dimensionwise!(x_local, vandermonde, view(node_coordinates, :, :, :, :, element), x_tmp1, x_tmp2)
    multiply_scalar_dimensionwise!(jacobian_local, vandermonde, inv.(view(inverse_jacobian, :, :, :, element)), jacobian_tmp1, jacobian_tmp2)

    # Calculate errors at each analysis node
    @. jacobian_local = abs(jacobian_local)

    for k in eachnode(analyzer), j in eachnode(analyzer), i in eachnode(analyzer)
      u_exact = initial_condition(get_node_coords(x_local, equations, dg, i, j, k), t, equations)
      diff = func(u_exact, equations) - func(get_node_vars(u_local, equations, dg, i, j, k), equations)
      l2_error += diff.^2 * (weights[i] * weights[j] * weights[k] * jacobian_local[i, j, k])
      linf_error = @. max(linf_error, abs(diff))
      volume += weights[i] * weights[j] * weights[k] * jacobian_local[i, j, k]
    end
  end

  # Accumulate local results on root process
  global_l2_error = Vector(l2_error)
  global_linf_error = Vector(linf_error)
  MPI.Reduce!(global_l2_error, +, mpi_root(), mpi_comm())
  MPI.Reduce!(global_linf_error, max, mpi_root(), mpi_comm())
  total_volume = MPI.Reduce(volume, +, mpi_root(), mpi_comm())
  if mpi_isroot()
    l2_error   = convert(typeof(l2_error),   global_l2_error)
    linf_error = convert(typeof(linf_error), global_linf_error)
    # For L2 error, divide by total volume
    l2_error = @. sqrt(l2_error / total_volume)
  else
    l2_error   = convert(typeof(l2_error),   NaN * global_l2_error)
    linf_error = convert(typeof(linf_error), NaN * global_linf_error)
  end

  return l2_error, linf_error
end

function integrate_via_indices(func::Func, u,
                               mesh::ParallelP4estMesh{3}, equations,
                               dg::DGSEM, cache, args...; normalize=true) where {Func}
  @unpack weights = dg.basis

  # Initialize integral with zeros of the right shape
  integral = zero(func(u, 1, 1, 1, 1, equations, dg, args...))
  volume = zero(real(mesh))

  # Use quadrature to numerically integrate over entire domain
  for element in eachelement(dg, cache)
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
      volume_jacobian = abs(inv(cache.elements.inverse_jacobian[i, j, k, element]))
      integral += volume_jacobian * weights[i] * weights[j] * weights[k] * func(u, i, j, k, element, equations, dg, args...)
      volume += volume_jacobian * weights[i] * weights[j] * weights[k]
    end
  end

  global_integral = MPI.Reduce!(Ref(integral), +, mpi_root(), mpi_comm())
  total_volume = MPI.Reduce(volume, +, mpi_root(), mpi_comm())
  if mpi_isroot()
    integral = convert(typeof(integral), global_integral[])
  else
    integral = convert(typeof(integral), NaN * integral)
    total_volume = volume # non-root processes receive nothing from reduce -> overwrite
  end

  # Normalize with total volume
  if normalize
    integral = integral / total_volume
  end

  return integral
end

end # @muladd


text:
using JLD2
using FileIO

include("src/init.jl")

"""
This file contains the script to evaluate the trained controller on the real physical system
"""

params.t_horizon = 300;       #Time until end
params.l_traj    = length(0:params.dt:params.t_horizon); #Total timesteps per episode
#params.σ *= 0.5
params.n_traj = 10;
rms_all = init_rms(params.s_init)

#function for smoothing the data
function f_filt(x)
    l = length(x)
    tau = 0.01
    y = zeros(l)
    y[1] = x[1]

    for i = 2:l
        y[i] = (1-tau)*y[i-1] + tau *x[i]
    end

    return y
end

"""
1. Load the trained network:
"""

#SELECT THE TRAINING RESULTS TO BE EVALUATED FROM:
#-------------------------------------------------------------------------------
trained_data = load("Results/10_01_real_segway/train_data.jld")
#trained_data = load("Results/Videos/08_30_PPO_Ball/Actor_Critic_2018-09-04iteartion_2000.jld")
#trained_data = load("Results/Videos/09_27_ballbot_more/Actor_Critic_2018-09-27iteartion_330.jld")
#trained_data = load("Results/Videos/09_26_ballbot/Actor_Critic_2018-09-26iteartion_1080.jld")

trained_data0 = load("Results/11_07_longterm_ball/1_all.jld2")
trained_data1 = load("Results/11_07_longterm_ball/PPO1_1_1500.jld2")
trained_data2 = load("Results/11_07_longterm_ball/PPO2_1_1500.jld2")

#Rewrite the results
data0 = trained_data0["result"][2,:]
data1 = trained_data1["(θa, θv, all_traj, best_net, result, kl_res)"][5][5,:]
data2 = trained_data2["(θa, θv, all_traj, best_net, result, kl_res)"][5][6,:]

θa  = trained_data1["(θa, θv, all_traj, best_net, result, kl_res)"][1]
θa2 = trained_data2["(θa, θv, all_traj, best_net, result, kl_res)"][1]

θa  = trained_data["θa"]
θa2 = trained_data["θa"]
#-------------------------------------------------------------------------------

# Test the network for functionality:
res = zeros(2)
prec = zeros(2)
i_iter = 1

all_traj  = init_all_traj(params)
traj = all_traj.traj[1]
#(ĝ_a, ĝ_v, res[i_iter], prec[i_iter]) = parallel_episodes(params,θa,θv,all_traj,rms_all)
a_t_old = all_traj.traj[1].a_t
weights_old = zeros(length(net_to_vec(θa)),2)

#a_t_old ,weights_old = plot_stuff(params, all_traj, a_t_old,weights_old, i_iter,θa, θv, res, prec/params.n_traj, ĝ_a, ĝ_v)

"""
2. we need a task-trajectory, that the segway has to follow:
"""

using HDF5

c = h5open("Results/10_01_real_segway/data.h5", "r") do file
    read(file, "xu")
end

h5file = h5open("Results/10_01_real_segway/trajectory.h5", "r")
des_traj = read(h5file)["des_traj"] .* 0.5
close(h5file)

using HDF5

c = h5open("Results/10_01_real_segway/data.h5", "r") do file
    read(file, "xu")
end

t_start = 10

x_real = Float64.(c)[1:7,:]
u1     = Float64.(c)[8:9,:]
t_real = Float64.(c)[10,:]
len_real = length(t_real)

des_traj_real = zeros(length(x_real[1,:]),2)
iter_beginn = t_start * 100
iter_end    = iter_beginn + length(des_traj[:,1])
l_end = len_real - iter_end +1
des_traj_real[iter_beginn:iter_end-1,:] = des_traj
des_traj_real[iter_end:end,:] = des_traj[end,:]' .* ones(l_end,2)

l_des = length(des_traj_real[:,1])

traj = init_traj(params,length(des_traj_real[:,1]))
traj2 = init_traj(params,length(des_traj_real[:,1]))

"""
3. The System was trained on going to the center starting from point (1,0) with random rotation.
    => we need to rotate the koordinate system at every step
"""
function run_episode_continous!(params::Params,θa::Any, traj::Trajectory, rms::RunMeanStd, des_traj::Array{Float64,2})
    dt          = params.dt
    t_horizon   = params.t_horizon
    s_init      = params.s_init
    σ           = params.σ
    u_range     = params.u_range
    l_out       = length(u_range)

    #Initialize :
    s_t1 = copy(s_init)
    s_t2 = similar(s_t1)
    totalr = 0
    rand_vec = zeros(l_out)

    #Declare the right position here:
    des_x      = des_traj[1,1]
    des_y      = des_traj[1,2]

    l_traj      = length(des_traj[:,1])
    l_state     = length(s_init)

    s_fake = s_init
    s_t1   = zeros(l_state)
    s_t    = copy(s_t1)

    for i_t = 1:l_traj

        #Update input
        s_t .= s_t1
        s_fake = s_t

        #Recalculate the Input
        d_x      = des_traj[i_t,1] - s_t1[1]
        d_y      = des_traj[i_t,2] - s_t2[2]

        scale = (d_x^2 + d_y^2)
        part  = 0.3

        s_fake[1]    = -sign(d_x) * sqrt(scale - d_y^2) * part
        s_fake[2]    = -sign(d_y) * sqrt(scale - d_x^2) * part

        #s_fake += 0.000*randn(length(s_fake))

        s_fake[1]    = -sign(d_x) * 1*min(abs(d_x),part)
        s_fake[2]    = -sign(d_y) * 1*min(abs(d_y),part)

        #-=---------------------------------
        s_mod        = zeros(8)
        s_mod[1:2]   = s_fake[1:2]
        s_mod[3]     = sin(s_fake[3])
        s_mod[4]     = cos(s_fake[3])
        s_mod[5:end] = s_fake[4:end]
        #-------------------------------=#

        #Evaluate the Network Output
        out_t = mlp_nlayer_policy(θa,s_mod,params)

        #Sample:
        rand_vec = randn(l_out) .* 0.0

        #Get Action with stochastic policy:
        loc_a_t = sample_action(out_t, params.σ, rand_vec)
        #loc_a_t = out_t
        p_a_t   = get_π_prob(out_t,σ,loc_a_t)

        act_fac = 3
        #Interact with the Environment:
        s_t2, loc_r_t, terminal = environment_call(s_t1,act_fac*loc_a_t,params)

        #Save data in trajectory:
        traj.s_t[i_t,:]    .= s_t1
        traj.s_t2[i_t,:]   .= s_t2
        traj.a_t[i_t,:]    .= loc_a_t
        traj.p_a_t[i_t,:]   = p_a_t
        traj.r_t[i_t]       = loc_r_t*0.01

        traj.out_t[i_t,:]   = out_t
        traj.len            = i_t
        totalr             += loc_r_t

        #Override the state:
        s_t1 .= s_t2;

        #Check if the terminal condition is fullfilled and break
        if terminal == 1
                break
        end

    end
    return totalr / traj.len
end

run_episode_continous!(params,θa, traj, rms_all, des_traj_real)
run_episode_continous!(params,θa2, traj2, rms_all, des_traj_real)

"""
4. Plot the results of the received trajectory
    => we need to rotate the koordinate system at every step
"""

using HDF5

c = h5open("Results/10_01_real_segway/data.h5", "r") do file
    read(file, "xu")
end

t_start = 10

x_real = Float64.(c)[1:7,:]
u1     = Float64.(c)[8:9,:]
t_real = Float64.(c)[10,:]
len_real = length(t_real)

des_traj_real = zeros(length(x_real[1,:]),2)
iter_beginn = t_start * 100
iter_end    = iter_beginn + length(des_traj[:,1])
l_end = len_real - iter_end +1
des_traj_real[iter_beginn:iter_end-1,:] = des_traj
des_traj_real[iter_end:end,:] = des_traj[end,:]' .* ones(l_end,2)

f_size = 12

y1 = 0.0
y2 = 0
yd = 4

fig = figure("all_res")
clf()
PyPlot.axes([0.12, 0.45, 0.33, 0.5])
title("PPO-Result", fontsize = f_size)
grid(alpha = 0.25)
plot(t_real,des_traj_real[:,1] .+ y1, label="x_des", color="black", linestyle = "--")
plot(t_real,traj.s_t[:,1] .+ y1)
plot(t_real,des_traj_real[:,2] .+ y1, label="y_des", color="black", linestyle = "-.")
plot(t_real,traj.s_t[:,2] .+ y1)
ylim([-0.5, 5])
#xlabel("time (s)")
ylabel("position (m)", fontsize = f_size)

PyPlot.axes([0.12, 0.1, 0.33, 0.25])
grid(alpha = 0.25)
plot(0,0)
plot(0,0)
plot(t_real, traj.a_t[:,1] .+ y2)#, color = "red")
plot(t_real, traj.a_t[:,2] .+ y2)#, color = "blue")
#plot(t_real, traj.a_t[:,3] .+ y2)#, color = "blue")
ylim([-yd+y2, yd+y2])
xlabel("time (s)", fontsize = f_size)
ylabel("u (V)", fontsize = f_size)

#------------------------------------------
PyPlot.axes([0.12+0.43, 0.45, 0.33, 0.5])
title("PPO-soft Result", fontsize = f_size)
grid(alpha = 0.25)
plot(t_real,des_traj_real[:,1] .+ y1, label="x_des", color="black", linestyle = "--")
plot(t_real,traj2.s_t[:,1] .+ y1)
plot(t_real,des_traj_real[:,2] .+ y1, label="y_des", color="black", linestyle = "-.")
plot(t_real,traj2.s_t[:,2] .+ y1)
ylim([-0.5,5])
#xlabel("time (s)")
#ylabel("position (m)")

PyPlot.axes([0.12+0.43, 0.1, 0.33, 0.25])
grid(alpha = 0.25)
plot(0,0)
plot(0,0)
plot(t_real, traj2.a_t[:,1] .+ y2)#, color = "red")
plot(t_real, traj2.a_t[:,2] .+ y2)#, color = "blue")
#plot(t_real, traj2.a_t[:,3] .+ y2)#, color = "blue")
ylim([-yd+y2, yd+y2])
xlabel("time (s)", fontsize = f_size)
#ylabel("u (Nm)")

figure("legend_seg3")
clf()
axis("off")
plot(0,0, label="x_des", color="black", linestyle = "--")
plot(0,0, label="y_des", color="black", linestyle = "-.")
plot(0,0, label = "x_true")
plot(0,0, label = "y_true")
plot(0,0, label = "u_1")
plot(0,0, label = "u_2")
#plot(0,0, label = "u_3")
legend(ncol = 3)

figure("draw")
clf()
plot(des_traj_real[:,1], des_traj_real[:,2])
plot(traj2.s_t[:,1], traj2.s_t[:,2])
plot(traj.s_t[:,1], traj.s_t[:,2])

X  = traj.s_t'
X2 = traj2.s_t'
u  = traj.a_t'
u2 = traj.a_t'
t  = vec(0:0.01:((length(u[1,:])-1) * 0.01))
length(t)

#using Model_Vis
filenameit = "results_ballvsball.jld2"
#visualise_ballbott_compare(X,X2,u,u2,t, filename, des_traj_real; options= Dict())

#For Visualisation in Blender
using NPZ

randomis = 0.001

X_soll = zeros(12,9085)
X_soll[1:7,:]  = copy(X)
X_soll[1:2,:]  = des_traj_real'
X_soll[1,:]   += f_filt(randomis*randn(length(X_soll[1,:])))
X_soll[2,:]   += f_filt(randomis*randn(length(X_soll[1,:])))

X[1,:] += f_filt(randomis*randn(length(X_soll[1,:])))
X[2,:] += f_filt(randomis*randn(length(X_soll[1,:])))

plot(X_soll[:,1])

npzwrite("X_soll4.npy", X_soll)
npzwrite("simulation.npy", X)
npzwrite("experiment.npy", x_real)

npzread("X_soll2.npy")

PPOs = npzread("PPO-soft.npy")
PPO  = npzread("PPO.npy")

PPOs[1,:] += f_filt(randomis*randn(length(PPOs[1,:])))
PPOs[2,:] += f_filt(randomis*randn(length(PPOs[1,:])))

PPO[1,:] += f_filt(randomis*randn(length(PPO[1,:])))
PPO[2,:] += f_filt(randomis*randn(length(PPO[1,:])))

npzwrite("PPO-soft2.npy", PPO)
npzwrite("PPO2.npy", PPOs)

max_stars_repo_name: NikonPic/REINFORCEMENT-JULIA

max_issues_repo_name: NikonPic/REINFORCEMENT-JULIA

max_forks_repo_name: NikonPic/REINFORCEMENT-JULIA