text:
using Distributions
using Random
using StaticArrays
using Test
using UnPack
using LinearAlgebra
using ArgParse
using Printf

using ClimateMachine
using ClimateMachine.Atmos # add include("AtmosCanopyModel.jl")
using ClimateMachine.Orientations
using ClimateMachine.ConfigTypes
using ClimateMachine.Diagnostics
using ClimateMachine.GenericCallbacks
using ClimateMachine.Mesh.Grids
using ClimateMachine.Mesh.Filters
using Thermodynamics.TemperatureProfiles # ClimateMachine.TemperatureProfiles is deprecated!
#using ClimateMachine.SystemSolvers
using ClimateMachine.ODESolvers
#using Thermodynamics, instead use
using ClimateMachine.Thermodynamics
using ClimateMachine.Thermodynamics: relative_humidity, air_density, internal_energy, total_energy, PhaseDry, PhaseEquil_pθq
using ClimateMachine.TurbulenceClosures
using ClimateMachine.TurbulenceConvection
using ClimateMachine.VariableTemplates
#import ClimateMachine.Atmos: atmos_source!
using ClimateMachine.Writers

using CLIMAParameters
using CLIMAParameters.Planet: MSLP, R_d, day, grav, cp_d, cv_d, planet_radius, LH_v0
using CLIMAParameters.Atmos.SubgridScale: C_smag

using ClimateMachine.BalanceLaws
import ClimateMachine.BalanceLaws: source, prognostic_vars
using ClimateMachine.Atmos: altitude

struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()
const n_tracer = Int(1)

struct CanopyAerodynamics{FT} <: TendencyDef{Source} # <: AbstractTendencyType, under BalanceLaw.
    "drag coefficient"
    c_d::FT
    "canopy height [m]"
    h_c::FT
    "Leaf Area Index LAI"
    LAI::FT
    #"Leaf Area Density LAD"
    #LAD::FT
    Δz::FT
end
prognostic_vars(::CanopyAerodynamics) = (Momentum(), SGSTKE(),)

function source(::Momentum, c::CanopyAerodynamics, m ,args)
    # currently only regular meshing is possible.
    # args include below:
    #_args = (; state, aux, t, direction, diffusive)
    #args = merge(_args, (precomputed = precompute(bl, _args, tend),))

    @unpack state, aux = args
    @unpack c_d, h_c, LAI = c
    LAD = LAI/h_c

    z = altitude(m, aux)
    if z <= h_c
        ρu = state.ρu
        ρ = state.ρ
        u_mag = norm(ρu/ρ)
        # see eq (9-10) in Patton et al 2016
        return -c_d * LAD * ρu * u_mag #SVector{3, FT}
    else
        FT = eltype(state)
        return SVector{3, FT}(0,0,0)
    end
end

function source(::SGSTKE, c::CanopyAerodynamics, m ,args)
    # currently only regular meshing is possible.
    @unpack state, aux = args
    @unpack c_d, h_c, LAI = c

    #turbulence_model(m)
    LAD = LAI/h_c
    FT = eltype(state)
    z = altitude(m, aux)
    if z <= h_c
        ρu = state.ρu
        ρ = state.ρ
        u_mag = norm(ρu/ρ)
        ρe = state.sgstke.ρe_SGS
        # see eq (9-10) in Patton et al 2016
        return -FT(8/3) * c_d * LAD * ρe * u_mag
    else
        return FT(0)
    end
end

struct CanopyAtmoInteraction{FT} <: TendencyDef{Source} # <: AbstractTendencyType, under BalanceLaw.
    "canopy height [m]"
    h_c::FT
    "Leaf Area Index LAI"
    LAI::FT
    #"Leaf Area Density LAD"
    #LAD::FT
    Δz::FT
    "vegetation canopy conductance"
    G_veg::FT #[m/s]
    "Gross Primary Production"
    GPP::FT
end

prognostic_vars(::CanopyAtmoInteraction) = (TotalMoisture(), Tracers{n_tracer}(),)
# Later it should include energy exchange due to Latent & sensible heat exchange

function source(::TotalMoisture, c::CanopyAtmoInteraction, m, args)
    @unpack state, aux = args
    @unpack h_c, LAI, Δz, G_veg = c
    @unpack ts = args.precomputed
    #@unpcak = node

    # incomplete! need to account for the gas moisture content...
    # First assuming T_l = T_air
    LAD = LAI/h_c
    FT = eltype(state)
    z = altitude(m, aux)
    if z <= h_c
        VPD = (FT(1) - relative_humidity(ts))*air_density(ts) # in density unit [kg/m3]
        return G_veg*LAD*VPD
    else
        return FT(0)
    end
end

function source(::Tracers{N}, c::CanopyAtmoInteraction, m, args) where {N}
    @unpack state, aux = args
    @unpack h_c, LAI, Δz, GPP = c
    #SPAC_Flux = spac_model(m)
    #node = SPAC_Flux.node
    #flux = SPAC_Flux.flux
    ##@unpcak = node
    LAD = LAI / h_c
    z = altitude(m, aux)
    FT = eltype(state)
    if z <= h_c
        return SVector{N, FT}(- GPP * LAD)
    else
        return SVector{N, FT}(0)
    end
end

function init_canopy_dry!(problem, bl, state, aux, localgeo, time)
    # here the first argument, `problem` is AtmosProblem, plugged in `init_state_prognostic`
    # ClimateMachine.Atmos.atmos_problem_init_state_auxiliary
    # orientation = FlatOrientation()
    FT = eltype(state)
    param_set = parameter_set(bl)
    (x, y, z) = localgeo.coord

    # These constants are those used by Stevens et al. (2005)
    #PhaseDry(param_set, e_int, state.e_int=19284.033384, ρ=1.160988, q_tot=1.000000, T = 404.512170, maxiter=3, tol=71.75060)
    # Patton et al 2015
    #qref = FT(0.0)  #FT(9.0e-3)

    R_gas = FT(R_d(param_set)) # for dry air
    c_v::FT = cv_d(param_set)
    c_p::FT = cp_d(param_set)
    p0::FT = MSLP(param_set)
    _grav = FT(grav(param_set))
    _LH_v0 = LH_v0(param_set) # Latent heat of vaporization

    # Use this if wish to follow the reference profile stated in the AtmosModel
    #ref_state = reference_state(bl)
    #θ, p = ref_state.virtual_temperature_profile(param_set, z)

    # Below is artifical profile, following Patton et al. 2016
    # Potential Temperature profile θ=300K for z<2h_c, and apply lapse rate of 3K/km.
    Γ = FT(3e-3) # [K/m], artificial lapse rate
    T_surf = FT(300) # This should be removed later
    h_c = FT(20) # Look up z_canopy in the main()
    if z <= 2*h_c
        θ = T_surf
    else
        θ = T_surf - Γ*(z - 2*h_c)
    end

    # initial velocity: 2-D geostrophic flow
    # MU case in Patton et al 2015
    ugeo = FT(5)
    vgeo = FT(0)
    u, v, w = ugeo, vgeo, FT(0)

    # Density, Temperature, ρ
    # Moisture: DryModel
    π_exner = FT(1) - _grav / (c_p * θ) * z
    ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas)
    T = θ * π_exner
    e_int = internal_energy(param_set, T)
    ts = PhaseDry(param_set, e_int, ρ)
    # Below are common
    e_kin = FT(0.5) * FT((u^2 + v^2 + w^2))
    e_pot = gravitational_potential(bl.orientation, aux)
    ρe_tot = ρ * total_energy(e_kin, e_pot, ts)

    state.ρ = ρ
    state.ρu = SVector(ρ * u, ρ * v, ρ * w)
    state.energy.ρe = ρe_tot
    if !(sgstke_model(bl) isa NoSGStke)
        state.sgstke.ρe_SGS = ( FT(0.2) + FT(rand(Uniform(-0.01, 0.01))) ) * e_kin
    end
    if !(tracer_model(bl) isa NoTracers)
        # tracer we would like to track is CO2--in ppm unit.
        # intial field of 340 ppm
        ρχ = SVector{n_tracer, FT}(FT(3.4e-4)*ρ)
        state.tracers.ρχ = ρχ
    end

    # initialize spac? do we need this?
    return nothing
end

function init_canopy_moist!(problem, bl, state, aux, localgeo, time)
    # here the first argument, `problem` is AtmosProblem, plugged in `init_state_prognostic`
    # ClimateMachine.Atmos.atmos_problem_init_state_auxiliary
    # orientation = FlatOrientation()
    FT = eltype(state)
    param_set = parameter_set(bl)
    (x, y, z) = localgeo.coord

    # These constants are those used by Stevens et al. (2005)
    #PhaseDry(param_set, e_int, state.e_int=19284.033384, ρ=1.160988, q_tot=1.000000, T = 404.512170, maxiter=3, tol=71.75060)
    # Patton et al 2015
    #qref = FT(0.0)  #FT(9.0e-3)

    q_tot = FT(1e-3) # kg/kg
    q_liq = FT(0)   # Assuming only water vapor state
    q_ice = FT(0)
    q_pt_sfc = PhasePartition(q_tot)
    Rm_sfc = gas_constant_air(param_set, q_pt_sfc)
    c_v::FT = cv_d(param_set)
    c_p::FT = cp_d(param_set)
    p0::FT = MSLP(param_set)
    _grav = FT(grav(param_set))
    _LH_v0 = LH_v0(param_set) # Latent heat of vaporization

    # Use this if wish to follow the reference profile stated in the AtmosModel
    #ref_state = reference_state(bl)
    #θ, p = ref_state.virtual_temperature_profile(param_set, z)

    # Below is artifical profile, following Patton et al. 2016
    # Potential Temperature profile θ=300K for z<2h_c, and apply lapse rate of 3K/km.
    Γ = FT(3e-3) # [K/m], artificial lapse rate
    T_surf = 300 # This should be removed later
    h_c = FT(20) # Look up z_canopy in the main()
    if z <= 2*h_c
        θ = T_surf
    else
        θ = T_surf - Γ*(z - 2*h_c)
    end

    # initial velocity: 2-D geostrophic flow
    # MU case in Patton et al 2015
    ugeo = FT(5)
    vgeo = FT(0)
    u, v, w = ugeo, vgeo, FT(0)

    # Density, Temperature, ρ
    # Moisture: EquilMoist
    H = Rm_sfc * T_surf / _grav
    p = p0 * exp(-z / H) # hydrostatic ressure
    θ_liq_ice = θ - (_LH_v0/c_p)*q_liq  # liquid-ice potential temperature, θ_liq = θ for q_liq = 0
    ts = PhaseEquil_pθq(param_set, p,θ_liq_ice, q_tot)
    ρ = air_density(ts)
    # Below are common
    e_kin = FT(0.5) * FT((u^2 + v^2 + w^2))
    e_pot = gravitational_potential(bl.orientation, aux)
    ρe_tot = ρ * total_energy(e_kin, e_pot, ts)

    state.ρ = ρ
    state.ρu = SVector(ρ * u, ρ * v, ρ * w)
    state.energy.ρe = ρe_tot
    # activate only for non-dry model!
    state.moisture.ρq_tot = ρ * q_tot
    if !(sgstke_model(bl) isa NoSGStke)
        state.sgstke.ρe_SGS =  ( FT(0.2) + FT(rand(Uniform(-0.01, 0.01))) ) * e_kin
    end
    if !(tracer_model(bl) isa NoTracers)
        # tracer we would like to track is CO2--in ppm unit.
        # intial field of 340 ppm
        ρχ = SVector{n_tracer, FT}(FT(3.4e-4)*ρ)
        state.tracers.ρχ = ρχ
    end

    # SPAC will be initialized separately
    return nothing
end

ClimateMachine.init()

function canopy_model(
    ::Type{FT},
    c_aero::CanopyAerodynamics{FT},
    c_atmo::CanopyAtmoInteraction{FT};
    config_type = AtmosLESConfigType,
    turbconv = NoTurbConv(),
    moisture_model = "dry", #"dry", "equilibrium"
    sgstke = NoSGStke(),
    ) where {FT}
    #@unpack  = node;

    if n_tracer > 0
        δ_χ = @SVector [FT(1.6e-5) for ii in 1:n_tracer] # Diffusion coefficient for CO2
        tracers = NTracers{n_tracer, FT}(δ_χ)
    else
        tracers = NoTracers()
    end

    T_surf = FT(300)  # from Patton et al 2016
    T_min_ref = FT(0)
    #T_virt = FT(300)
    temp_profile_ref = DryAdiabaticProfile{FT}(param_set, T_surf, T_min_ref) #or use IsothermalProfile(param_set, T_virt)
    # Note that IsothermalProfile = DecayingTemperatureProfile{FT}(param_set, T_virt, T_virt)
    ref_state = HydrostaticState(temp_profile_ref)

    C_drag = FT(0.0011)
    #_C_smag = FT(C_smag(param_set))
    LHF = FT(50) # Latent heat flux
    SHF = FT(15) # Sensible heat flux
    moisture_flux = LHF / FT(LH_v0(param_set))

    # Assemble source components
    source_sgstke = (ShearProduction(), BuoyancyProduction(), Dissipation(),)
    source_default = (Gravity(), c_aero, source_sgstke...)
    if moisture_model == "dry"
        source = source_default
        moisture = DryModel()
        moisture_bcs = ()
        init_state_prognostic = init_canopy_dry!
        #tracer_bcs = ()
    elseif moisture_model == "equilibrium"
        source = (source_default..., c_atmo)
        moisture = EquilMoist(; maxiter = 5, tolerance = FT(0.1))
        moisture_bcs = (; moisture = PrescribedMoistureFlux((state, aux, t) -> moisture_flux)) # No-flux from the ground, modify `LHF`
        init_state_prognostic = init_canopy_moist!
        #tracer_bcs = (tracer = ImpermeableTracer())
    #elseif moisture_model == "nonequilibrium"
    #    source = (source_default..., CreateClouds())
    #    moisture = NonEquilMoist()
    else
        @warn @sprintf(
            """
%s: unrecognized moisture_model in source terms, using the defaults""",
            moisture_model,
        )
        source = source_default
        moisture_bcs = ()
        #tracer_bcs = ()
    end

    physics = AtmosPhysics{FT}(
        param_set;
        ref_state = ref_state,
        turbulence = Deardorff(), #SmagorinskyLilly(_C_smag),
        moisture = moisture,
        tracers = tracers,
        turbconv = turbconv,
        sgstke = sgstke,
    )

    problem = AtmosProblem(
        boundaryconditions = (
            AtmosBC(
                physics;
                momentum = Impenetrable(DragLaw(
                    (state, aux, t, normPu) -> C_drag,
                )),#default is Impenetrable{FreeSlip}
                ## Latent + Sensible heat flux from the ground?
                energy = PrescribedEnergyFlux((state, aux, t) -> LHF + SHF),
                moisture_bcs...,
                #tracer_bcs...,
                turbconv = turbconv_bcs(turbconv)[1],
                #sgstke = NoSGStkeFlux(),
            ),
            AtmosBC(physics; turbconv = turbconv_bcs(turbconv)[2]),          #it's not wrong, dont' worry!
            # default setting of AtmosBC: { Insulating, ClimateMachine.TurbulenceConvection.NoTurbConvBC}
        ),
        init_state_prognostic = init_state_prognostic, # you can either put it here or in `model`
        # ClimateMachine.Atmos.atmos_problem_init_state_auxiliary
    )

    #struct AtmosModel{FT, PH, PR, O, S, DC} <: BalanceLaw
    model = AtmosModel{FT}(
        config_type,                            # Flow in a box, requires the AtmosLESConfigType
        physics;
        problem = problem,
        source = source,
        # default settings: hyperdiffusion = NoHyperDiffusion(), radiation = NoRadiation(), precipitation = NoPrecipitation(), moisture = EquilMoist{Float64}(3, 0.1)
        # we can set hyperdiffusion by (τ_hyper = FT(4 * 3600); # hyperdiffusion time scale hyperdiffusion_model = DryBiharmonic(FT(4 * 3600));)
    )

    return model
end

function main()
    # add a command line argument to specify the kind of surface flux
    # TODO: this will move to the future namelist functionality
    canopy_args = ArgParseSettings(autofix_names = true)
    add_arg_group!(canopy_args, "ATMO-LAND")
    @add_arg_table! canopy_args begin
        "--surface-flux"
        help = "specify surface flux for energy and moisture"
        metavar = "prescribed|bulk"
        arg_type = String
        default = "prescribed"
    end

    cl_args =
        ClimateMachine.init(parse_clargs = true, custom_clargs = canopy_args)
    ClimateMachine.init(parse_clargs = true) # every 5 simulation minutes
    surface_flux = cl_args["surface_flux"]

    FT = Float64

    # DG polynomial order
    N = 4

    # Domain resolution and size
    Δh = FT(10)
    Δv = FT(5)
    resolution = (Δh, Δh, Δv)

    xmax = FT(100)
    ymax = FT(100)
    zmax = FT(100)

    t0 = FT(0)
    timeend = FT(600)
    CFLmax = FT(0.1)     # use this for single-rate explicit LSRK144

    # constants for struct Canopy
    c_d = FT(0.2)
    LAI = FT(2.0)       # Later improve this to LAD profile
    Δz = resolution[3] # This approach only works for uniform gridding
    G_veg = FT(0.2)
    h_c = FT(20)
    GPP = FT(1)

    ode_solver_type = ClimateMachine.ExplicitSolverType(
        solver_method = LSRK144NiegemannDiehlBusch,)

    model = canopy_model(
        FT,
        CanopyAerodynamics(c_d, h_c, LAI, Δz),
        CanopyAtmoInteraction(h_c, LAI, Δz, G_veg, GPP);
	moisture_model = "equilibrium",
	sgstke = SGStkeModel(),
        )
    ics = model.problem.init_state_prognostic

    # need to tweak the configuration to add land module flux simulation
    driver_config = ClimateMachine.AtmosLESConfiguration(
        "CanopyFlow",
        N,
        resolution,
        xmax,
        ymax,
        zmax,
        param_set,
        ics,
        model = model,
    )
    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config,
        ode_solver_type = ode_solver_type,
        init_on_cpu = true,
        Courant_number = CFLmax,
    )

    # The `user_callbacks` are passed to the ODE solver as callback functions;
    # see [`solve!`](@ref ODESolvers.solve!).

    # GenericCallbacks.init!(callbacks, solver, Q, param, t0)
    #=
    cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do
        Filters.apply!(
            solver_config.Q,
            ("moisture.ρq_tot", "tracers.ρχ",),
            solver_config.dg.grid,
            TMARFilter(),
        )
        nothing
    end
    =#
    check_cons = (
        ClimateMachine.ConservationCheck("ρ", "1smins", FT(0.0001)),
        ClimateMachine.ConservationCheck("energy.ρe", "1smins", FT(0.0025)),
    )

    result = ClimateMachine.invoke!(
        solver_config;
        check_cons = check_cons,
        #user_callbacks = (cbtmarfilter, ),
        check_euclidean_distance = true,
    )

end

main()


text:
using LinearAlgebra.BLAS
using OpenBLAS32
using OpenBLAS32_jll
using Test

@testset "OpenBLAS32 test suite" begin
    libs = BLAS.get_config().loaded_libs
    idx = findfirst(l -> l.libname == OpenBLAS32_jll.libopenblas_path, libs)
    @test idx !== nothing
    @test libs[idx].interface === :lp64
end

ext: jl


text:
#==============================================================================#
# System
#==============================================================================#

"""
	@enum Ensemble::UInt8 Canonical Grand

Statistical ensemble for a field-theoretic system.
"""
@enum Ensemble::UInt8 begin
	Canonical
    Grand
end

"""
	mutable struct FieldSystem <: AbstractSystem

A system containing data for a field-theoretic simulation,
represented as a set of species, field grids, and a corresponding unit cell.
Contains all species, field, density, and interaction information.
"""
mutable struct FieldSystem <: AbstractSystem
	# Grid and cell
	dims          :: NTuple{3,Int}
	ensemble      :: Ensemble
	cell          :: UnitCell

	# MDE and FFT
	fftplan       :: FFTHolder
	solver        :: PseudoSpectralSolver

	# Field grids for each monomer type
	monomers      :: Dict{Int,Monomer}
	monomer_fracs :: Dict{Int,Float64}
	fields        :: Dict{Int,FieldGrid{Float64}}
	potentials    :: Dict{Int,FieldGrid{Float64}}
	residuals     :: Dict{Int,FieldGrid{Float64}}
	density       :: Dict{Int,FieldGrid{Float64}}

	# Interactions and species
	#constraints  :: Vector{AbstractConstraint}
	interactions  :: Vector{AbstractInteraction}
	species       :: Vector{AbstractSpecies}
	phi_species   :: Vector{Float64}
	mu_species    :: Vector{Float64}

	# Constraints
	compress      :: Compressibility
end

function FieldSystem(dims::NTuple{3,<:Integer}, cell::UnitCell;
	monomers = Monomer[], ensemble::Ensemble = Canonical,
	compress = Compressibility(0.0), mde::Symbol = :RK2, nthreads::Integer = -1)
	# Create FFT and MDE solver
	fftplan = FFTHolder(dims; nthreads = nthreads)
	solver = PseudoSpectralSolver(dims; method = mde)

	sdims = Tuple(sort(collect(dims), rev = true))
	sys = FieldSystem(sdims, ensemble, cell, fftplan, solver,
		Dict(), Dict(), Dict(), Dict(), Dict(), Dict(), [], [], [], [],
		compress
	)

	# Allocate to system size
	setup!(cell, sys)
	setup!(compress, sys)

	# Add initial monomers
	for mon in monomers
		add_monomer!(sys, mon)
	end

	return sys
end

FieldSystem(nx::Integer, ny::Integer, nz::Integer, cell::UnitCell; kwargs...) = FieldSystem((nx, ny, nz), cell; kwargs...)

#==============================================================================#

function Base.show(io::IO, sys::FieldSystem)
	@printf(io, "FieldSystem(dims = %s, ensemble = %s, %d monomers, %d species)",
		sys.dims, sys.ensemble, nmonomers(sys), nspecies(sys))
end

volume(sys::FieldSystem) = sys.cell.volume

ngrid(sys::FieldSystem) = prod(sys.dims)

nspecies(sys::FieldSystem) = length(sys.species)

nmonomers(sys::FieldSystem) = length(sys.monomers)

hasmonomer(sys::FieldSystem, mid::Integer) = haskey(sys.monomers, mid)

#==============================================================================#

"""
	add_cell!(sys, cell)

Provide a new cell for the system.
"""
function add_cell!(sys::FieldSystem, cell::UnitCell)
	sys.cell = cell
	setup!(cell, sys)
	return nothing
end

"""
	add_monomer!(sys, mon)

Add a `Monomer` to the `FieldSystem`.
"""
function add_monomer!(sys::FieldSystem, mon::Monomer)
	if hasmonomer(sys, mon.id)
		@warn "Replacing monomer with id = $(mon.id) in FieldSystem."
	end

	# Add to system, add new fields
	sys.monomers[mon.id] = mon
	sys.monomer_fracs[mon.id] = 0.0

	sys.fields[mon.id] = zeros(Float64, sys.dims)
	sys.potentials[mon.id] = zeros(Float64, sys.dims)
	sys.residuals[mon.id] = zeros(Float64, sys.dims)
	sys.density[mon.id] = zeros(Float64, sys.dims)

	return nothing
end

"""
	add_species!(sys, species, phi/mu)

Add an `AbstractSpecies` to the `FieldSystem`.
For `Canonical` ensemble systems, the third argument is treated as the species volume fraction.
For `Grand` ensemble systems, this argument is the species chemical potential.
"""
function add_species!(sys::FieldSystem, species::AbstractSpecies, val::Real)
	# Check to make sure species has valid monomers
	for mid in species.mids
		if !hasmonomer(sys, mid)
			error("Species contains monomer with id = $(mid) not present in system.")
		end
	end

	setup!(species, sys)
	push!(sys.species, species)

	# Add value based on ensemble of simulation
	if sys.ensemble == Canonical
		if val <= 0.0
			error("Bulk volume fraction must be 0 < phi <= 1.")
		end
		push!(sys.phi_species, val)
		push!(sys.mu_species, 0.0)
	elseif sys.ensemble == Grand
		push!(sys.phi_species, 0.0)
		push!(sys.mu_species, val)
	end

	return nothing
end

"""
	add_interaction!(sys, itx)

Add an `AbstractInteraction` to the `FieldSystem`.
"""
function add_interaction!(sys::FieldSystem, itx::AbstractInteraction)
	# Check to make sure has valid monomers
	for mid in itx.mids
        if !hasmonomer(sys, mid)
            error("Interaction contains monomer with id = $(mid) not present in system.")
        end
    end

	setup!(itx, sys)
	push!(sys.interactions, itx)
	return nothing
end

"""
	validate(sys)

Determine if the system is properly initialized before a calculation.
Checks if the following requirements are met:

* All fields are initialized with non-zero mean or standard deviation
* All monomer types are found in at least one species
* Bulk fractions of all monomers sums to unity (canonical ensemble)
* Bulk charge fractions of all monomers sums to zero
"""
function validate(sys::FieldSystem)
	# Check if fields are initialized
	fields_init = false
	for omega in values(sys.fields)
		if mean(omega) != 0.0 || std(omega) != 0.0
			fields_init = true
			break
		end
	end
	if !fields_init
		error("Fields must be initialized before SCFT simulation.")
	end

	# Check bulk volume fraction constraint
	if sys.ensemble == Canonical
		if !(sum(sys.phi_species) ≈ 1.0)
			error("Species volume fractions do not sum to unity.")
		end
	end

	# Check bulk charge neutrality
	bulk_charge = 0.0
	monomer_fractions!(sys)
	for (mid, mon) in sys.monomers
		bulk_charge += sys.monomer_fracs[mid] * mon.charge / mon.vol
	end
	if !(bulk_charge ≈ 0.0)
		error("Bulk system not charge neutral.")
	end

	# Check if any monomer types do not have associated species
	for mid in keys(sys.monomers)
		has_species = false
		for species in sys.species
			if hasmonomer(species, mid)
				has_species = true
				break
			end
		end
		if !has_species
			error("Monomer with ID = $(mid) not found in any species.")
		end
	end

	return nothing
end

#==============================================================================#
# Methods for system density, potentials, and constraints
#==============================================================================#

"""
	muphi!(sys)

Update species volume fractions and chemical potentials.
"""
function muphi!(sys::FieldSystem)
	if sys.ensemble == Canonical
		for (isp, species) in enumerate(sys.species)
			sys.mu_species[isp] = log(sys.phi_species[isp] / species.Q)
		end
	elseif sys.ensemble == Grand
		for (isp, species) in enumerate(sys.species)
			sys.phi_species[isp] = species.Q * exp(sys.mu_species[isp])
		end
	end
	return nothing
end

"""
	density!(sys)

Compute density fields for all monomers in the system.
"""
function density!(sys::FieldSystem)
	# Zero density fields
	for rho in values(sys.density)
		rho .= 0.0
	end

	# Loop over all species, update density on block/type level
	for species in sys.species
		density!(species)
	end

	# Update all bulk fractions/chemical potentials
	muphi!(sys)

	# Pool the updated density by type at the system level
	for (isp, species) in enumerate(sys.species)
		phi = sys.phi_species[isp]
		for mid in species.mids
			@. sys.density[mid] += phi * species.density[mid]
		end
	end

	return nothing
end

"""
	potentials!(sys)

Compute differentials of the excess interactions with respect to each density field.
"""
function potentials!(sys::FieldSystem)
	# Zero current potentials
	for pot in values(sys.potentials)
		pot .= 0.0
	end

	# Add excess interactions
	for (mid, pot) in sys.potentials
		for itx in sys.interactions
			potential!(itx, mid, pot)
		end
	end

	return nothing
end

"""
    residuals!(sys)

Compute the gradient of the system Hamiltonian with respect to each density field.
"""
function residuals!(sys::FieldSystem)
	# Run updates on all field grids
	density!(sys)
	potentials!(sys)

	# Update compressibility field
	update!(sys.compress)
	eta = sys.compress.field

    # Make residuals from fields and calculated potentials/eta field
	for (mid, res) in sys.residuals
		@. res = sys.potentials[mid] + eta - sys.fields[mid]
	end

	# Mean-subtract the system fields and residuals
	if sys.ensemble == Canonical
		for mid in keys(sys.fields)
			sys.fields[mid] .-= mean(sys.fields[mid])
			sys.residuals[mid] .-= mean(sys.residuals[mid])
		end
	end

	return nothing
end

"""
	uniform_fields!(sys)

Set the value of all chemical potential fields equal to the summation of interaction potentials,
corresponding to the mean-field limit with vanishing constraint fields.
"""
function uniform_fields!(sys::FieldSystem)
	potentials!(sys)
	for (mid, omega) in sys.fields
		@. omega = sys.potentials[mid]
		if sys.ensemble == Canonical;
			omega .-= mean(omega);
		end
	end
	return nothing
end

"""
	monomer_fractions!(sys)

Update the bulk fractions for each monomer present in the system.
"""
function monomer_fractions!(sys::FieldSystem)
	bulk = sys.monomer_fracs
	for mid in keys(bulk); bulk[mid] = 0.0; end

	for (isp, species) in enumerate(sys.species)
		phi = sys.phi_species[isp]
		for mid in species.mids
			bulk[mid] += phi * monomer_fraction(species, mid)
		end
	end
	return bulk
end

#==============================================================================#

"""
	scfstress(sys)

Compute the free energy variation of the system with respect to each cell parameter.
"""
function scfstress(sys::FieldSystem)
	stress = zeros(nparams(sys.cell))
	muphi!(sys)

	# Add stress from each contributing species
	for (isp, species) in enumerate(sys.species)
		phi = sys.phi_species[isp]
		stress .-= phi * scfstress(species)
	end

	return stress
end

"""
	free_energy(sys)

Return the intensive Helmholtz free energy for the interacting, inhomogeneous system.
"""
function free_energy(sys::FieldSystem)
	# Add translational entropy terms
	ftrans = 0.0
	for (isp, species) in enumerate(sys.species)
		phi = sys.phi_species[isp]
		mu = sys.phi_species[isp]
		if phi > 1e-8
			ftrans += phi * (mu - 1.0) / species.Nref
		end
	end

	# Add field/density pair
	ffield = 0.0
	for (mid, omega) in sys.fields
		rho = sys.density[mid]
		@simd for i in eachindex(rho)
			@inbounds ffield -= rho[i] * omega[i]
		end
	end
	ffield /= ngrid(sys)

	# Add excess interaction terms
	fexc = 0.0
	for itx in sys.interactions
		fexc += energy(itx)
	end

	# All pieces together
	fhelm = ftrans + ffield + fexc

	return fhelm
end

"""
	free_energy_bulk(sys)

Return the intensive Helmholtz free energy for a system at the specified bulk composition.
"""
function free_energy_bulk(sys::FieldSystem)
	fhelm = 0.0

	# Add translational entropy terms
	for (isp, species) in enumerate(sys.species)
		phi = sys.phi_species[isp]
		if phi > 1e-8
			fhelm += phi * (log(phi) - 1.0) / species.Nref
		end
	end

	# Add the excess interaction energy
	monomer_fractions!(sys)
	for itx in sys.interactions
		fhelm += energy_bulk(itx)
	end

	return fhelm
end

"""
    field_error(sys)

Compute the weighted error based on the current field residuals.

Follows the definition in the Matsen (2009) Anderson mixing papers:
``
{\\rm err} = [∑_{α,i} {\\rm res}_{α,i}^2 / ∑_{α,i} ω_{α,i}^2]^{1/2}
``
"""
function field_error(sys::FieldSystem)
    # Residual weights
    rsum = 0.0
    for (mid, res) in sys.residuals
        @simd for i in eachindex(res)
            @inbounds rsum += res[i]^2
        end
    end
    # Field weights
    wsum = eps(Float64)
    for (mid, omega) in sys.fields
        @simd for i in eachindex(omega)
            @inbounds wsum += omega[i]^2
        end
    end
    return sqrt(rsum) / sqrt(wsum)
end


text:
using Test
using TestProject
using ComoniconOptions
using ComoniconBuilder
using ComoniconBuilder: install, install_script, install_completion,
    cmd_script, contains_path, contains_fpath, write_path, write_fpath, detect_rcfile,
    install_env_path

home_dir = mktempdir()
test_dir = pkgdir(ComoniconBuilder, "test")
usr_dir = joinpath(test_dir, "usr")

options = ComoniconOptions.Comonicon(
    name="test",
    install=ComoniconOptions.Install(
        path=usr_dir,
        completion=false,
    ),
    sysimg=nothing,
)

@testset "installation" begin
    install(TestProject, options)
    @test isfile(joinpath(usr_dir, "bin", "test"))
    @test isfile(joinpath(usr_dir, "bin", "test.jl"))
    rm(usr_dir; recursive=true)
end

options = ComoniconOptions.Comonicon(
    name="test",
    install=ComoniconOptions.Install(
        path=usr_dir,
    ),
    sysimg=ComoniconOptions.SysImg(precompile=ComoniconOptions.Precompile(["deps/precompile_execution.jl"], []))
)

# empty!(ARGS)
# push!(ARGS, "sysimg")
# install(TestProject, options)

if get(ENV, "TEST_BUILD_SYSIMG", false) == "true"
    @testset "System Image building" begin
        options_sys = ComoniconOptions.Comonicon(
            name="test",
            install=ComoniconOptions.Install(
                path=usr_dir,
            ),
        sysimg=ComoniconOptions.SysImg(precompile=ComoniconOptions.Precompile(["precompile_execution.jl"], []))
    )
        outpath = ComoniconBuilder.pkgdir(TestProject, "deps/lib/libtest.so")

        isfile(outpath) && rm(outpath)
        ComoniconBuilder.build_sysimg(TestProject, options_sys)
        # Test if the new system image is present
        @test isfile(outpath) && isnothing(rm(outpath))
    end
end

if get(ENV, "TEST_BUILD_APP", false) == "true"
    @testset "App building" begin
        options_app = ComoniconOptions.Comonicon(
            name="test",
            install=ComoniconOptions.Install(
                path=usr_dir,
            ),
        application=ComoniconOptions.Application(
            precompile=ComoniconOptions.Precompile(["precompile_execution.jl"], []),
            incremental=true,
            filter_stdlibs=false,
            )
    )
        outpath = ComoniconBuilder.pkgdir(TestProject, "build/test/bin/test")

        isfile(outpath) && rm(outpath)
        ComoniconBuilder.build_application(TestProject, options_app)
        # Test if the app has been built
        @test isfile(outpath) && isnothing(rm(outpath))
    end
end


text:
function solution_point(instance_path::String)
    solution1 = readdlm(joinpath(instance_path,"solution1.txt"), ',')
    solution2 = readdlm(joinpath(instance_path,"solution2.txt"), ',')
    basecase_generator_solution = read_solution1(solution1)
    contingency_generator_solution, bus_solution, delta_solution, line_flow_solution = read_solution2(solution2)
    power_data= getpowerdata(instance_path)
    index = get_bus_index(power_data)
    participation_factors = create_participation_factors_dict(instance_path, index)
    global_point = create_global_point(solution1, solution2, participation_factors, index)
    # Add delta solution
    delta_point = create_delta_solution(solution2)
    global_point = merge(global_point, delta_point)
    return global_point
end

function create_voltagesolutionpoint(bus_solution, index)
    point = Point()
    for i in 1:size(bus_solution,1)
        num_bus = index[bus_solution[i,2]]
        value = bus_solution[i,3]*exp(im*bus_solution[i,4]*pi/180)
        if bus_solution[i,1] == 0
            add_coord!(point, Variable(variable_name("VOLT", bus_name(num_bus), "",basecase_scenario_name()), Complex), value)
        else
            id_contingency = bus_solution[i,1]
            scenario_name = scenarioname(id_contingency)
            add_coord!(point, Variable(variable_name("VOLT", bus_name(num_bus), "", scenario_name), Complex), value)
        end
    end
    return point
end

function gen_solution_basecase(basecase_generator_solution,index)
    point = Point()
    for i in 1:size(basecase_generator_solution,1)
        num_bus = index[basecase_generator_solution[i,1]]
        gen_id = basecase_generator_solution[i,2]
        if typeof(gen_id)!=Int64
            gen_id = matchall(r"\d+" ,gen_id)[1]
        else
            gen_id = Int(gen_id)
        end
        Pg = basecase_generator_solution[i,3]
        Qg = basecase_generator_solution[i,4]
        add_coord!(point, Variable(variable_name("Sgen", bus_name(num_bus), generator_name(gen_id), basecase_scenario_name()), Complex), Pg + im*Qg)
    end
    return point
end

function gen_solution_contingencies(basecase_generator_solution,contingency_generator_solution,delta_solution, participation_factors, index)
    point = Point()

    ##reactive power
    for i in 1:size(contingency_generator_solution,1)
        id_contingency = contingency_generator_solution[i,1]
        scenario_name = scenarioname(id_contingency)
        gen_id = contingency_generator_solution[i,4]
        num_bus = index[Int(contingency_generator_solution[i,3])]
        if typeof(gen_id)!=Int64
            gen_id = matchall(r"\d+" ,gen_id)[1]
        else
            gen_id = Int(gen_id)
        end
        Qg = contingency_generator_solution[i,5]
        add_coord!(point, Variable(variable_name("Sgen", bus_name(num_bus), generator_name(gen_id), scenario_name), Complex), im*Qg)
    end

    #active power
    for l in 1:size(delta_solution,1)
        id_contingency = delta_solution[l,1]
        scenario_name = scenarioname(id_contingency)
        for i in 1:size(basecase_generator_solution,1)
            num_bus = index[Int(basecase_generator_solution[i,1])]
            gen_id = basecase_generator_solution[i,2]
            if typeof(gen_id)!=Int64
                gen_id = matchall(r"\d+" ,gen_id)[1]
            else
                gen_id = Int(gen_id)
            end
            Pg = basecase_generator_solution[i,3]

            var = Variable(variable_name("Sgen",bus_name(num_bus), generator_name(gen_id), scenario_name), Complex)

            add_coord!(point, var, Pg + delta_solution[l,2] * participation_factors[bus_name(num_bus)][generator_name(gen_id)])
        end
    end
    return point
end

function create_participation_factors_dict(instance_path, index)
    path_folder = instance_path
    if !("generator.csv" ∈ readdir(path_folder))
        error("missing file : generator.csv must be in $path_folder")
    end
    generator_data = readdlm(joinpath(path_folder,"generator.csv"),',',skipstart=1)

    participation_factors = Dict{String,Dict{String,Float64}}()
    for line in 1:size(generator_data,1)
        term_id = generator_data[line,3]
        if term_id == 9
            num_bus = index[Int(generator_data[line,1])]
            gen_id = generator_data[line,2]
            if typeof(gen_id)!=Float64
                gen_id = matchall(r"\d+" ,gen_id)[1]
            else
                gen_id = Int(gen_id)
            end
            value =  generator_data[line,4]
            if !haskey(participation_factors,bus_name(num_bus))
                participation_factors[bus_name(num_bus)] = Dict{String,Float64}()
            end
            participation_factors[bus_name(num_bus)][generator_name(gen_id)] = value
        end
    end
    return participation_factors
end

function create_delta_solution(solution2)
    point = Point()
    contingency_generator_solution, bus_solution, delta_solution, line_flow_solution = read_solution2(solution2)
    for i in 1:size(delta_solution,1)
        scenario = scenarioname(delta_solution[i,1])
        add_coord!(point, Variable(get_delta_varname(scenario),Real), delta_solution[i,2])
    end
    return point
end

function create_global_point(solution1,solution2, participation_factors,index)
    basecase_generator_solution = read_solution1(solution1)
    contingency_generator_solution, bus_solution, delta_solution, line_flow_solution = read_solution2(solution2)
    voltage_point = create_voltagesolutionpoint(bus_solution,index)
    genpoint =  gen_solution_basecase(basecase_generator_solution,index)

    genpoint_contingencies = gen_solution_contingencies(basecase_generator_solution,contingency_generator_solution,delta_solution, participation_factors,index)

    return merge(voltage_point, genpoint, genpoint_contingencies)
end

function nb_active_constraints_in_scenario_cc_Qgen(filenames, epsilon::Float64)
    filenames_path = joinpath("instances","GOC", filenames[1], filenames[2])
    solution1 = readdlm(joinpath(filenames_path,"solution1.txt"), ',')
    solution2 = readdlm(joinpath(filenames_path,"solution2.txt"), ',')
    gen_solution_basecase = read_solution1(solution1)
    ~, bus_solution, ~, ~ = read_solution2(solution2)
    power_data = getpowerdata(filenames)
    index = get_bus_index(power_data)
    #bus_id = Int(generator_data[line,1])
    generators = collect(gen_solution_basecase[:,1])
    generators = [ index[gen] for gen in generators]
    #println(length(generators))
    module_v_basecase = Dict{Int64, Float64}()
    module_v_scenarios = Dict{String,Dict{Int64, Float64}}()
    for i in 1:size(bus_solution,1)
        num_bus = index[bus_solution[i,2]]
            if bus_solution[i,1] == 0
                module_v_basecase[num_bus] = bus_solution[i,3]
            else
                id_contingency = bus_solution[i,1]
                scenario_name = scenarioname(id_contingency)
                if !haskey(module_v_scenarios, scenario_name)
                    module_v_scenarios[scenario_name] = Dict{Int64, Float64}()
                end
                module_v_scenarios[scenario_name][num_bus] = bus_solution[i,3]
            end
    end
    nb = 0
    for (scenario, dict_modules) in module_v_scenarios
        for num_bus in generators
            if abs(dict_modules[num_bus] - module_v_basecase[num_bus]) > epsilon
                nb += 1
            end
        end
    end

    return nb
end
# filenames = ["Phase_0_IEEE14_1Scenario", "scenario_1"]
# filenames = ["Phase_0_RTS96", "scenario_10"]
# nb_active_constraints_in_scenario_cc_Qgen(filenames, 1e-4)
# folder = "Phase_0_IEEE14"
# scenario = "scenario_10"
# epsilon = 1e-4
# get_nb_cc_reactive_power_active(folder, scenario, epsilon)

##########################################
#READ
#########################################

function read_solution1(solution1::Array{Any,2})
    nb_lines = size(solution1,1)
    #description_line = Array{Any}()
    line_start = 0
    line_end = 0
    for l in 1:nb_lines
        if solution1[l,1] == "--generation dispatch"
            #description_line = solution1[l+1,1:4]
            line_start = l + 2
            break
        end
    end
    i = line_start
    while solution1[i,1] != "--end of generation dispatch"
        if i==nb_lines
            error("no --end of generation dispatch in input solution1")
        else
        i +=1
        end
    end
    line_end = i-1
    basecase_generator_solution = solution1[line_start:line_end,1:4]
    return basecase_generator_solution
end

#
# basecase_generator_solution_description = ["bus id", "unit id", "pg(MW)", "qg(MVar)"]
# contingency_generator_solution_description = ["conID", "genID", "busID", "unitID", "q(MW)"]
# contingency_bus_solution_description = ["contingency id", "bus id", "v(pu)", "theta(deg)"]
# contingency_delta_solution_description = ["contingency id", "Delta(MW)"]
# contingency_line_flow_solution_description = ["contingency id", "line id", "origin bus id", "destination bus id", "circuit id", "p_origin(MW)", "q_origin(MVar)", "p_destination(MW)", "q_destination(MVar)"]
#

function read_solution2(solution2::Array{Any,2})
    nblines = size(solution2,1)

    i = 1
    while solution2[i,1] != "--contingency generator"
        if i==nblines
            error("no --contingency generator in input solution2")
        else
        i +=1
        end
    end
    #println(solution2[i+1,1:5])
    line_start_ctg, line_end_ctg = findnumarray(i+2, solution2)

    i = line_end_ctg + 1
    while solution2[i,1] != "--bus"
        if i==nblines
            error("no --bus in input solution2")
        else
        i +=1
        end
    end
    #println(solution2[i+1,1:4])
    line_start_bus, line_end_bus = findnumarray(i+2, solution2)

    i = line_end_bus + 1
    while solution2[i,1] != "--Delta"
        if i==nblines
            error("no --Delta in input solution2")
        else
        i +=1
        end
    end
    #println(solution2[i+1,1:2])
    line_start_delta, line_end_delta = findnumarray(i+2, solution2)

    i = line_end_delta + 1
    while solution2[i,1] != "--line flow"
        if i==nblines
            error("no --line flow in input solution2")
        else
        i +=1
        end
    end
    #println(solution2[i+1,1:9])
    line_start_line, line_end_line = findnumarray(i+2, solution2)

    if line_end_line + 1 == nblines
        contingency_generator_solution = solution2[line_start_ctg:line_end_ctg,1:5]
        bus_solution = solution2[line_start_bus:line_end_bus,1:4]
        delta_solution = solution2[line_start_delta:line_end_delta,1:2]
        line_flow_solution = solution2[line_start_line:line_end_line,1:9]

        return contingency_generator_solution, bus_solution, delta_solution, line_flow_solution
    else
        error()
    end
end

function findnumarray(line_start, data)
    i = line_start
    name = data[i-2,1][3:end]
    nb_lines = size(data,1)
    while data[i,1]!="--end of "*name
        if i==nb_lines
            error("no --end of $name in input solution2")
        else
        i +=1
        end
    end
    line_end = i-1
    line_start, line_end
end

max_stars_repo_name: Gimaju/GOC

max_issues_repo_name: Gimaju/GOC

max_forks_repo_name: Gimaju/GOC


text:
# -*- encoding: utf-8 -*-
#
# The MIT License (MIT)
#
# Copyright © 2021 Matteo Foglieni and Riccardo Gervasoni
#

"""
    create_onb_from_z(normal::Union{Vec, Normal}) :: (Normal, Normal, Normal)

Return an orthonormal base of 3 `Normal`s with the z-axes
(i.e. the third `Normal` returned) parallel to the
input `Vec`/ `Normal`.

The implementation of this function is based on the paper
of Duff et al. (2017), which improved the already efficient
work made by Frisvad (2012).

## References

- Duff et al. (2017), ["Building an Orthonormal Basis,
  Revisited"](https://graphics.pixar.com/library/OrthonormalB/paper.pdf)

- Frisvad (2012), ["Building an Orthonormal Basis from a 3D
  Unit Vector Without Normalization"](https://backend.orbit.dtu.dk/ws/portalfiles/portal/126824972/onb_frisvad_jgt2012_v2.pdf)

See also: [`Vec`](@ref), [`Normal`](@ref)
"""
function create_onb_from_z(normal::Union{Vec, Normal})
    (typeof(normal) == Vec) && (normal = Normal(normal))

    s = copysign(1, normal.z)
    a = -1.0/(s + normal.z)
    b = normal.x * normal.y * a

    e1 = Normal(1.0 + s * normal.x^2 * a, s * b, -s * normal.x)
    e2 = Normal(b, s + normal.y^2 * a, -normal.y)

    return e1, e2, Normal(normal.x, normal.y, normal.z)
end


text:
using DimensionalData, Statistics, Test, Unitful, SparseArrays, Dates, LinearAlgebra

using Combinatorics: combinations

@testset "*" begin
    timespan = DateTime(2001):Month(1):DateTime(2001,12)
    A1 = DimArray(rand(12), (Ti(timespan),))
    A2 = DimArray(rand(12, 1), (Ti(timespan), X(10:10:10)))

    @test length.(dims(A1)) == size(A1)
    @test dims(parent(A1) * permutedims(A1)) isa Tuple{<:AnonDim,<:Ti}
    @test parent(A1) * permutedims(A1) == parent(A1) * permutedims(parent(A1))
    @test dims(permutedims(A1) * parent(A1)) isa Tuple{<:AnonDim}
    @test permutedims(A1) * parent(A1) == permutedims(parent(A1)) * parent(A1)

    @test length.(dims(permutedims(A1) * parent(A1))) == size(permutedims(parent(A1)) * parent(A1))
    @test length.(dims(permutedims(A1) * A1)) == size(permutedims(parent(A1)) * parent(A1))
    @test length.(dims(permutedims(parent(A1)) * A1)) == size(permutedims(parent(A1)) * parent(A1))

    @test length.(dims(parent(A1) * permutedims(A1))) == size(parent(A1) * permutedims(parent(A1)))
    @test length.(dims(A1 * permutedims(A1))) == size(parent(A1) * permutedims(parent(A1)))
    @test length.(dims(A1 * permutedims(parent(A1)))) == size(parent(A1) * permutedims(parent(A1)))

    @test length.(dims(A2)) == size(A2)
    @test length.(dims(A2')) == size(A2')

    sze1 = (12, 12)
    @test size(parent(A2) * parent(A2)') == sze1
    @test length.(dims(A2 * A2')) == sze1
    @test length.(dims(parent(A2) * A2')) == sze1
    @test length.(dims(A2 * parent(A2'))) == sze1
    sze2 = (1, 1)
    @test size(parent(A2') * parent(A2)) == sze2
    @test length.(dims(A2' * A2)) == sze2
    @test length.(dims(parent(A2') * A2)) == sze2
    @test length.(dims(A2' * parent(A2))) == sze2

    B1 = DimArray(rand(12, 6), (Ti(timespan), X(1:6)))
    B2 = DimArray(rand(8, 12), (Y(1:8), Ti(timespan)))
    b1 = DimArray(rand(12), Ti(timespan))

    # Test dimension propagation
    @test (B2 * B1) == (parent(B2) * parent(B1))
    @test dims(B2 * B1) isa Tuple{<:Y, <:X}
    @test length.(dims(B2 * B1)) == (8, 6)

    # Test where results have an empty dim
    true_result = (parent(b1)' * parent(B1))
    for flip in (adjoint, transpose, permutedims)
        result = flip(b1) * B1
        @test result ≈ true_result  # Permute dims is not exactly transpose
        @test dims(result) isa Tuple{<:AnonDim, <:X}
        @test length.(dims(result)) == (1, 6)
    end

    # Test flipped * flipped
    for (flip1, flip2) in combinations((adjoint, transpose, permutedims), 2)
        result = flip1(B1) * flip2(B2)
        @test result == flip1(parent(B1)) * flip2(parent(B2))
        @test dims(result) isa Tuple{<:X, <:Y}
        @test size(result) == (6, 8)
    end

end

@testset "some matmul ambiguity methods" begin
    special_types = (Adjoint, Transpose, Diagonal, Symmetric, Tridiagonal, SymTridiagonal, BitArray,)

    @testset "DimMatrix" begin
        da = DimArray(ones(5,5), (:a, :b))
        @testset "$T" for T in special_types
            x = T(ones(5,5))
            @test dims(x * da) isa Tuple{<:AnonDim, Dim{:b}}
            @test dims(da * x) isa Tuple{Dim{:a}, <:AnonDim}
            @test typeof(x * da) <: DimArray
            @test typeof(da * x) <: DimArray
            @test parent(da' * x) == parent(da)' * x
            @test parent(x * da) == x * parent(da)
        end
    end

    @testset "DimVector" begin
        dv = DimArray(ones(5), :vec)
        @testset "$T" for T in special_types
            x = T(ones(5,5))
            @test dims(x * dv) isa Tuple{<:AnonDim}
            @test dims(dv' * x) isa Tuple{<:AnonDim,<:AnonDim}
            @test typeof(x * dv) <: DimArray
            @test typeof(dv' * x) <: DimArray
            @test parent(dv' * x) == parent(dv)' * x
            @test parent(x * dv) == x * parent(dv)
        end
        @testset "$T" for T in (Adjoint, Transpose)
            x = T(1:5)
            @test x * dv === 15.0
            @test typeof(dv * x) <: DimArray
            @test dims(dv * x) isa Tuple{<:Dim{:vec},AnonDim}
            @test dv * x == vcat(x, x, x, x, x)
        end

    end

end


text:
import Base: @deprecate, broadcast
export sampler, sample

@deprecate broadcast(t::Tile) wrapbehavior(t)
@deprecate dropdownmenu(t::Tile, xs::AbstractArray) dropdownmenu(t, menu(xs))
@deprecate addinterpreter(x, y) intent(x, y)
@deprecate sample(x::Sampler, t) intent(x, t)
@deprecate plugsampler(s::Sampler, t::Tile) intent(s, t)
@deprecate constant(x, y) intent(constant(x), y)
@deprecate pairwith(x, y) intent(pairwith(x), y)


text:
"""
Code declaring Spectra1DBasic, Spectra1DBasic, and their abstract versions.

Author: Eric Ford
Created: August 2020
"""

#=
const min_orders_in_spectra = 1
const max_orders_in_spectra = 128
const min_pixels_in_order = 1
const max_pixels_in_order = 9128
const min_pixels_in_spectra = min_orders_in_spectra*min_pixels_in_order
const max_pixels_in_spectra = max_orders_in_spectra*max_pixels_in_order
min_usable_pixels_in_order = 128
=#

""" Abstract type for any Spectrum (or region of spectrum) """
abstract type AbstractSpectra end
""" Abstract type for any 1-d Spectrum (or region of spectrum) """
abstract type AbstractSpectra1D <: AbstractSpectra end
""" Abstract type for any 2-d Spectrum """
abstract type AbstractSpectra2D <: AbstractSpectra end

MetadataT = Dict{Symbol,Any}

""" Basic struct for Spectra1D (or region of specturm)
    Instruments can specialize their own if additional data is avaliable. """
struct Spectra1DBasic{T1<:Real,T2<:Real,T3<:Real,
    AA1<:AbstractArray{T1,1},AA2<:AbstractArray{T2,1},AA3<:AbstractArray{T3,1},
    InstT<:AbstractInstrument
    } <: AbstractSpectra1D
    λ::AA1
    flux::AA2
    var::AA3
    inst::InstT
    metadata::MetadataT #Dict{Symbol,Any}
end

""" Basic struct for Spectra2D (or region of specturm)
    Instruments can specialize their own if additional data is avaliable. """
struct Spectra2DBasic{T1<:Real,T2<:Real,T3<:Real,
     AA1<:AbstractArray{T1,2},AA2<:AbstractArray{T2,2},AA3<:AbstractArray{T3,2},
     InstT<:AbstractInstrument
    } <: AbstractSpectra2D
    λ::AA1
    flux::AA2
    var::AA3
    inst::InstT
    metadata::MetadataT # Dict{Symbol,Any}
end

""" Basic struct for Spectra2D (or region of specturm)
    Instruments can specialize their own if additional data is avaliable. """
struct Spectra2DExtended{T1<:Real,T2<:Real,T3<:Real,
     AA1<:AbstractArray{T1,2},AA2<:AbstractArray{T2,2},AA3<:AbstractArray{T3,2},
     InstT<:AbstractInstrument
    } <: AbstractSpectra2D
    λ::AA1
    λ_obs::AA1
    flux::AA2
    var::AA3
    inst::InstT
    metadata::MetadataT # Dict{Symbol,Any}
end

function Spectra1DBasic(λ::A1, flux::A2, var::A3, inst::InstT;
        metadata::Dict{Symbol,Any} = Dict{Symbol,Any}() ) where {
          T1<:Real, T2<:Real, T3<:Real, A1<:AbstractArray{T1,1}, A2<:AbstractArray{T2,1}, A3<:AbstractArray{T3,1},
          InstT<:AbstractInstrument }
    @assert size(λ) == size(flux)
    @assert size(λ) == size(var)
    @assert 1 <= length(λ) <= max_pixel(inst)-min_pixel(inst)+1
    Spectra1DBasic{eltype(λ),eltype(flux),eltype(var),typeof(λ),typeof(flux),typeof(var),typeof(inst)}(λ,flux,var,inst,metadata)
end

function Spectra2DBasic(λ::A1, flux::A2, var::A3, inst::InstT;
     metadata::MetadataT = MetadataT() ) where {
     T1<:Real, T2<:Real, T3<:Real, A1<:AbstractArray{T1,2}, A2<:AbstractArray{T2,2}, A3<:AbstractArray{T3,2},
     InstT<:AbstractInstrument  }
    @assert size(λ) == size(flux)
    @assert size(λ) == size(var)
    @assert 1 <= size(λ,1) <= max_pixel_in_order(inst)
    @assert 1 <= size(λ,2) <= max_order(inst)
    @assert 1 <= min_order(inst) <= max_order(inst)
    @assert 1 <= min_pixel_in_order(inst) <= max_pixel_in_order(inst)
    Spectra2DBasic{eltype(λ),eltype(flux),eltype(var),typeof(λ),typeof(flux),typeof(var),typeof(inst)}(λ,flux,var,inst,metadata)
end

function Spectra2DExtended(λ::A1, λ_obs::A1, flux::A2, var::A3, inst::InstT;
     metadata::MetadataT = MetadataT() ) where {
     T1<:Real, T2<:Real, T3<:Real, A1<:AbstractArray{T1,2}, A2<:AbstractArray{T2,2}, A3<:AbstractArray{T3,2},
     InstT<:AbstractInstrument  }
    @assert size(λ) == size(λ_obs) == size(flux) == size(var)
    @assert 1 <= size(λ,1) <= max_pixel_in_order(inst)-min_pixel_in_order(inst)+1
    @assert 1 <= size(λ,2) <= max_order(inst)-min_order(inst)+1
    Spectra2DExtended{eltype(λ),eltype(flux),eltype(var),typeof(λ),typeof(flux),typeof(var),typeof(inst)}(λ,λ_obs,flux,var,inst,metadata)
end


text:
using ChebyshevQuantum
using Test

@testset "test quad and derivative" begin

    function f(x)
        return exp(1.2*x)
    end

    c = ChebyshevQuantum.Interp.make_cheb(f , a=2, b = 3)

    exact = (exp(1.2*3) - exp(1.2*2))/1.2
    int = ChebyshevQuantum.Interp.integrateCC(c)
    int2 = sum(c)

    @test isapprox(exact, int, atol=1e-9)
    @test isapprox(exact, int2, atol=1e-9)

    x=2.2
    exact = 1.2 * exp(x*1.2)

    cprime = ChebyshevQuantum.Interp.derivative(c)
    cprime2 = c'
#    println(exact, " ", cprime(x))
    @test isapprox(exact, cprime(x), atol=1e-9)
    @test isapprox(exact, cprime2(x), atol=1e-9)

end


text:
module CauchyBorn

using JuLIP
using JuLIP.Potentials: StillingerWeber

import JuLIPMaterials.CLE: elastic_moduli

export WcbQuad

# "a fully equilibrated SW potential"
# function sw_eq()
#     T(σ, at) = trace(stress(StillingerWeber(σ=σ), at))
#     at = JuLIP.ASE.bulk("Si", pbc=true)
#     r0 = 2.09474
#     r1 = r0 - 0.1
#     s0, s1 = T(r0, at), T(r1, at)
#     while (abs(s1) > 1e-8) && abs(r0 - r1) > 1e-8
#         rnew = (r0 * s1 - r1 * s0) / (s1 - s0)
#         r0, r1 = r1, rnew
#         s0, s1 = s1, T(rnew, at)
#     end
# #     @show r1
#     return StillingerWeber(σ=r1)
# end

function DpWcb(F, p, at, calc)
    @assert length(at) == 2
    set_defm!(at, F)
    X = positions(at)
    X[2] = X[1] + p
    set_positions!(at, X)
    return -forces(calc, at)[2]
end

function DpDpWcb(at)
    calc = calculator(at)
    F0 = defm(at)
    p0 = positions(at)[2] |> Vector
    h = 1e-5
    DpDpW = zeros(3, 3)
    for i = 1:3
        p0[i] += h
        DpW1 = DpWcb(F0, JVecF(p0), at, calc)
        p0[i] -= 2*h
        DpW2 = DpWcb(F0, JVecF(p0), at, calc)
        p0[i] += h
        DpDpW[:, i] = (DpW1 - DpW2) / (2*h)
    end
    return 0.5 * (DpDpW + DpDpW')
end

type WcbQuad{TA, TF, TC}
    DpDpW::Matrix{Float64}
    DpDpW_inv::Matrix{Float64}
    at::TA
    calc::TC
    F0::TF
end

function WcbQuad(calc)
    # sw = sw_eq()
    at = bulk("Si", pbc=true)
    set_calculator!(at, calc)
    DpDpW = DpDpWcb(at)
    return WcbQuad(DpDpW, pinv(DpDpW), at, calc, defm(at))
end

# TODO: replace with get_shift
function (W::WcbQuad)(F)
    # TODO this is fishy - why is the initial position not reset?
    p0 = positions(W.at)[2]
    p1 = p0 - W.DpDpW_inv * DpWcb(F, p0, W.at, W.calc)
    p2 = p1 - W.DpDpW_inv * DpWcb(F, p1, W.at, W.calc)
    return p2
end

function DW_infp(W::WcbQuad, F)
    p = W(F)
    at = W.at
    set_defm!(at, F)
    X = positions(at)
    X[2] = p
    set_positions!(at, X)
    return stress(W.calc, at)
end

function elastic_moduli(W::WcbQuad)
   F0 = W.F0 |> Matrix
   Ih = Matrix(1.0*I, 3,3) # eye(3)
   h = eps()^(1/3)
   C = zeros(3,3,3,3)
   for i = 1:3, a = 1:3
      Ih[i,a] += h
      Sp = DW_infp(W, Ih * F0)
      Ih[i,a] -= 2*h
      Sm = DW_infp(W, Ih * F0)
      C[i, a, :, :] = (Sp - Sm) / (2*h)
      Ih[i,a] += h
   end
   # symmetrise it - major symmetries C_{iajb} = C_{jbia}
   for i = 1:3, a = 1:3, j=1:3, b=1:3
      t = 0.5 * (C[i,a,j,b] + C[j,b,i,a])
      C[i,a,j,b] = t
      C[j,b,i,a] = t
   end
   # minor symmetries - C_{iajb} = C_{iabj}
   for i = 1:3, a = 1:3, j=1:3, b=1:3
      t = 0.5 * (C[i,a,j,b] + C[i,a,b,j])
      C[i,a,j,b] = t
      C[i,a,b,j] = t
   end
   return C
end

end


text:
"""
    DeterministicFlow(n, method[, projection])

Deterministic log-homotopy flow discretized by `n` steps, where vector field is estimated using `method`.

### Optional arguments
- `projection`: a function to be applied after every step of the flow.
"""
struct DeterministicFlow{TG<:GainEstimationMethod, TP<:Function} <: ParticleFlowMethod
    n_steps::Int
    gain_method::TG
    projection::TP
end

DeterministicFlow(n, method) = DeterministicFlow(n, method, identity)

#####################
### BASIC METHODS ###
#####################

function Base.show(io::IO, ::MIME"text/plain", flow::DeterministicFlow)
    print(io, "Deterministic particle flow
    gain estimation method:                 ", flow.gain_method, "
    number of steps:                        ", flow.n_steps)
end

function Base.show(io::IO, flow::DeterministicFlow)
    print(io, "Deterministic particle flow with ", flow.n_steps, " steps and ", flow.gain_method)
end

######################
### MAIN ALGORITHM ###
######################

function flow!(ensemble::UnweightedParticleEnsemble{T}, log_likelihood::Function, method::DeterministicFlow) where T
    eq = PoissonEquation(log_likelihood, ensemble)
    if eq.H == zeros(T, size(eq.H)) #nothing to compute
        foreach(method.projection, eachcol(ensemble.positions))
        return ensemble.positions
    end

    n = method.n_steps
    for i in 1:n
        solve!(eq, method.gain_method)
        ensemble.positions .+= eq.gain[:,:,1]/n
        foreach(method.projection, eachcol(ensemble.positions))
        update!(eq, ensemble)
    end
    return ensemble.positions
end

max_stars_repo_name: simsurace/FeedbackParticleFilters

max_issues_repo_name: simsurace/FeedbackParticleFilters

max_forks_repo_name: simsurace/FeedbackParticleFilters


text:
module GlobalSensitivity

using Statistics, RecursiveArrayTools, LinearAlgebra
using QuasiMonteCarlo, ForwardDiff
using Parameters: @unpack, @with_kw
using FFTW, Distributions

abstract type GSAMethod end

include("morris_sensitivity.jl")
include("sobol_sensitivity.jl")
include("regression_sensitivity.jl")
include("DGSM_sensitivity.jl")
include("eFAST_sensitivity.jl")

export Sobol, Morris, gsa,
       SensitivityAlg, RegressionGSA, DGSM, eFAST

end # module


text:
####################################################
# Determine cell subgraph
####################################################

using LightGraphs
using CellularPotts
M = ModelParameters()
CPM = CellPotts(M)

g = CPM.graph.network
cellIdx = findall(isequal(1), CPM.graph.σ)

#This function attempts to find a subgraph of the a faster, however performance is about the same 😖
function induced_subgraph_mod(g::T, vlist::AbstractVector{U}) where T <: AbstractGraph where U <: Integer

    gSubRaw = view(g.fadjlist, vlist)
    fadjlist = filter.(x -> x ∈ vlist, gSubRaw)
    #d = Dict(vlist .=> eachindex(vlist))

    for list in fadjlist
        for (j,vertex) in enumerate(list)
            list[j] = searchsortedfirst(vlist,vertex) #assumes vlist is sorted
            #list[j] = d[vertex]
        end
    end

    return SimpleGraph( sum(length,fadjlist) ÷ 2, fadjlist)
end

induced_subgraph_mod(g,cellIdx)

####################################################
# Determine if a node is an Articulation point
####################################################

#Currently, articulation points are recalculated each time MHStep is called which is very slow/wasteful
#Need a better way to calculate to determine how articulation points change after a node is added removed


text:
module NeuronBuilder

using ModelingToolkit

const t = Num(Sym{Real}(:t))
const D = Differential(t)
export t, D

include("Species.jl")
export Species, SpeciesProperty, Ion, SpeciesDynamics, Voltage, Sodium, Potassium, Calcium, Proton, Leak, Reversal, Current, Conductance, mRNA
export shorthand_name

include("Components.jl")
export Component, Compartment, FlowChannel, Neuron, Synapse, EmptyConnection, PlasticityRule, PlasticisedChannel, Geometry
export get_name

include("FlowChannels.jl")
export sensed, actuated, currents, sensedvars, mrna, reversals, conductances, sensed_ions, instantiate_variables, instantiate_parameters, instantiate_hooks, get_sensor, get_actuator, voltage

include("PlasticityRules.jl")
export OLeary_reg, OLearyCalcRegulation

include("SpeciesDynamics.jl")
export BasicVoltageDynamics
export get_parameters, get_states, default_params, default_states

include("Geometries.jl")
export NoGeometry
export capacitance

include("helper_functions.jl")
export vardivide

include("assets/Liu/Liu.jl")
using .Liu
export Liu

include("assets/Prinz/Prinz.jl")
using .Prinz
export Prinz

include("Neurons.jl")
export BasicNeuron, BasicVoltageDynamics, EmptyNeuron

include("Synapses.jl")
export Chol, Glut, directed_synapse

include("Networks.jl")
export build_network, build_group, add_connection, add_all_connections

end


text:
include("constants.jl")
include("VelocityScale.jl")
include("holtslag_psim.jl")
include("holtslag_psis.jl")
include("holtslag_local.jl")

function holtslag2(θv,U,V,LH,SH,ustar,ρA,ZF)
# a vertical diffusion scheme based on Holtslag and Boville (1993)
# Includes enhenced diffusion below the PBL
#θv-virtual potential temperature

Rcrit=0.5;
b=8.5
eps=0.1;
minshear=0.003

Z=ZF[2:end];
DZ=diff(ZF,dims=1);
ZM=Z-DZ./2;
SZM=length(ZM);

#wθv-surface virtual heat flux, positive upward
wθv=SH./ρA./cpa.+humid_fac.*θv[1].*LH./ρA./Av;

#TF=(SH+LH)./ρA./cpa;
ustar=abs(ustar)
L=-θv[1].*ustar.^3 ./(karman.*gravity_mks.*wθv);
H=ZF[2:end]./L;

# surface velocity scale
out=VelocityScale.(0.1*Z,Z,L,ustar,θv[1],wθv)
ws=map(x->x[1], out)
wm=map(x->x[2], out)

θs=θv[1].+max.(b.*wθv./ws,0);

#calculate PBLH at high resolution
ha=(Rcrit.*θs.*(U.^2+V.^2)./(gravity_mks.*(θv-θs)))-Z;
ha[ha.>0].=Inf;
if all(isinf.(ha))
    if wθv>0
        id=SZM;
    else
        id=1;
    end
else
    ~,id=findmin(abs.(ha));
end
id=max(id,1);
h=abs(Z[id[1]]);
# n - levels inside PBL
n=(h.-Z).*((h.-Z).>0);
n[n.==0].=1e6;
~,i=findmin(n);
# frac - fraction of the upper PBL level that is indide the PBL
frac=(h-ZF[i])./DZ[i];

# nsl - levels inside surface layer
nsl=(0.1 .*h.-ZF).*((0.1 .*h.-ZF).>0);
nsl[nsl.==0].=1e6;
~,isl=findmin(nsl);

#ipbl - levels inside the PBL
#npbl - levels outside the PBL
ipbl=[1:i;];
if i==SZM
    npbl=[];
else
    npbl=i+1:SZM;
end

kM=zeros(SZM); kS=zeros(SZM);
wm=zeros(SZM); wt=zeros(SZM);

out=VelocityScale.(ZM[ipbl],h,L,ustar,θv[ipbl],wθv)

wt[ipbl]=map(x->x[1], out)
wm[ipbl]=map(x->x[2], out)

kS[ipbl]=karman.*wt[ipbl].*min.(ZM[ipbl],h).*(1 .- min.(ZM[ipbl]/h,1)).^2
kM[ipbl]=karman.*wm[ipbl].*min.(ZM[ipbl],h).*(1 .- min.(ZM[ipbl]/h,1)).^2

kM[i]=kM[i].*frac
kS[i]=kS[i].*frac

#vertical diffusion above the PBL
S=max.(abs.(diff([0; U])./diff(ZF)),minshear);
GRi=(gravity_mks./θv[1:end-1]).*diff(θv)./diff(ZM)./S[1:end-1].^2;

fm=[0; holtslag_local.(GRi)]

lc=30

i2=min(i+1,SZM)
kM[npbl]=kM[npbl]+lc.^2 .*S[npbl].*fm[npbl];
kS[npbl]=kM[npbl]+lc.^2 .*S[npbl].*fm[npbl];
println(h)
return kM,kS
#[i max(kM) max(kS) wθv max(Ts) U(1)]

end


text:
println("Building LiterateOrg.jl")

# Bootstrap package by extracting all package code to a Julia file and
# include it.

src_file = joinpath(dirname(@__FILE__), "..", "src", "LiterateOrg.org")

deps_dir = joinpath(dirname(@__FILE__), "..", "deps")
bootstrap_file = joinpath(deps_dir, "bootstrap.jl")

start_code_pat = r"[ ]*#\+begin_src[ ]+julia(.*)"
end_code_pat = r"[ ]*#\+end_src"

test_pat = r"^([*]+)[ ]+test"
sec_pat = r"^([*]+)"

println("Bootstrapping LiterateOrg.jl to $(bootstrap_file)")

code_mode = false
# We use this to keep track of which section level the last test
# section was seen at. Lower-level sections are parts of the test
# section, and should be skipped from the bootstrap file.
test_level = Inf

open(src_file) do infile
    open(bootstrap_file, "w") do outfile
        for line in readlines(infile)
            global code_mode, test_level
            lline = lowercase(line)
            if occursin(test_pat, lline) && test_level == Inf
                # If we've found a test section and have not seen one
                # before, count the number of stars to find out which
                # level we are on.
                m = match(test_pat, lline)
                test_level = length(m[1])
                continue
            elseif occursin(sec_pat, lline)
                # If we've found a normal section on a higher level
                # than the last seen test section, leave test mode.
                m = match(sec_pat, lline)
                if length(m[1]) < test_level
                    test_level = Inf
                end
                continue
            elseif occursin(start_code_pat, lline) && test_level == Inf
                # Copy code to the bootstrap file, as long as we're
                # not in test mode.
                code_mode = true
                continue
            elseif occursin(end_code_pat, lline)
                code_mode = false
                continue
            end
            code_mode && write(outfile, "$(line)\n")
        end
    end
end

println("Running bootstrap file")
include(bootstrap_file)
rm(bootstrap_file)

println("Tangling LiterateOrg.jl")

tangle_package(src_file, "LiterateOrg")


text:
using Main.AtBashM: atbash

while true

    println("\nText to shift [leave blank to break]: ")
    human_input = readline()
    human_input == "" ? break :
    println("\nAlphabet [leave blank for default -- default is lower case]: ")
    human_alphabet = readline()
    human_alphabet == "" ? human_alphabet = "abcdefghijklmnopqrstuvwxyz" :
    println()

    println("Decrypt:")
    println(atbash(human_input, human_alphabet))

end # while

max_stars_repo_name: Squalm/Cipher-Tools

max_issues_repo_name: Squalm/Cipher-Tools

max_forks_repo_name: Squalm/Cipher-Tools


text:
const jobs = Channel{Int}(32)
const results = Channel{Tuple}(32)

function do_work()
   for job_id in jobs
       #prev_time=Dates.Time()
       exec_time = rand()*job_id
       sleep(exec_time)# simulates elapsed time doing actual work typically performed externally.
       #Libc.systemsleep(exec_time)
       put!(results, (job_id, exec_time, "hi"))
       #time=Dates.Time()-prev_time
       println("$job_id with $(round(exec_time;digits=2)) seconds")
   end
end

function make_jobs(n)
   for i in 1:n
       put!(jobs, i)
   end
end

n = 12
schedule(@task make_jobs(n))

for i in 1:4 # start 4 tasks to process requests in parallel
    schedule(@task do_work())
end

@elapsed while n > 0 # print out results
  global n
   job_id, exec_time, msg = take!(results)
   println("$job_id with \"$msg\" finished in $(round(exec_time;digits=2)) seconds")
   n = n - 1
end

max_stars_repo_name: Gilga/JuliaOpenGL

max_issues_repo_name: Gilga/JuliaOpenGL

max_forks_repo_name: Gilga/JuliaOpenGL


text:
################
# DUMMY SOLVER
################
import GeoStatsBase: solvesingle

@simsolver DummySimSolver begin end
function solvesingle(problem::SimulationProblem, covars::NamedTuple,
                     solver::DummySimSolver, preproc)
  reals = map(covars.names) do var
    npts = npoints(domain(problem))
    V    = variables(problem)[var]
    real = vcat(fill(zero(V), npts÷2), fill(one(V), npts÷2))
    var => real
  end

  Dict(reals)
end

###################
# DUMMY ESTIMATOR
###################
import GeoStatsBase: fit, predict, status

struct DummyEstimator end
struct FittedDummyEstimator
  z
end

fit(estimator::DummyEstimator, X, z) = FittedDummyEstimator(z)
function predict(fitted::FittedDummyEstimator, xₒ)
  z  = fitted.z
  μ  = sum(z) / length(z)
  σ² = sum((v - μ)^2 for v in z) / length(z)
  μ, σ²
end
status(fitted::FittedDummyEstimator) = true


text:
abstract type AbstractContext <: AbstractElement end

struct Style <: AbstractElement
    context::Union{AbstractContext, Nothing}
    value::NamedTuple
end
Style(s::NamedTuple=NamedTuple()) = Style(nothing, s)
Style(c::AbstractContext) = Style(c, NamedTuple())
Style(s::Style) = s

function Base.show(io::IO, s::Style)
    print(io, "Style with entries ")
    show(io, keys(s.value))
end

style(args...; kwargs...) = Style(namedtuple(args...; kwargs...))

function Base.:*(s1::Union{Style, AbstractContext}, s2::Union{Style, AbstractContext})
    merge(Style(s1), Style(s2))
end

function Base.merge(s1::Style, s2::Style)
    c1, c2 = s1.context, s2.context
    c = c2 === nothing ? c1 : c2
    if c !== c1 && !isempty(s1.value)
        @warn "Changing context on a non empty style"
    end
    return _merge(c, s1, s2)
end
Base.pairs(s::Style) = _pairs(s.context, s)

# interface and fallbacks

# Must return a Style
_merge(c, s1::Style, s2::Style) = Style(c, merge(s1.value, s2.value))

# Must return a vector of NamedTuple => Style pairs
_pairs(c, s::Style) = [NamedTuple() => s]

## Dims context

struct DimsSelector{N} <: AbstractContext
    dims::NTuple{N, Int}
end
dims(args...) = DimsSelector(args)

Broadcast.broadcastable(d::DimsSelector) = fill(d)

Base.isless(s1::DimsSelector, s2::DimsSelector) = isless(s1.dims, s2.dims)

function aos(d::NamedTuple{names}) where names
    v = broadcast((args...) -> NamedTuple{names}(args), d...)
    return v isa NamedTuple ? fill(v) : v
end

function adjust(val::DimsSelector, c, nts)
    is = collect(val.dims)
    ax, cf = axes(nts)[is], Tuple(c)[is]
    return LinearIndices(ax)[cf...]
end

# make it a pair list
function _pairs(::DimsSelector{0}, s::Style)
    nts = aos(s.value)
    ps = map(CartesianIndices(nts)) do c
        nt = nts[c]
        dskeys = filter(keys(nt)) do key
            nt[key] isa DimsSelector
        end
        dsvalues = map(dskeys) do key
            val = nt[key]
            return adjust(val, c, nts)
        end
        ds = (; zip(dskeys, dsvalues)...)
        nds = Base.structdiff(nt, ds)
        return ds => Style(nds)
    end
    return vec(ps)
end

function _pairs(d::DimsSelector, s::Style)
    value = map(s.value) do v
        v isa AbstractArray ? mapslices(x -> [x], v; dims=d.dims) : v
    end
    return pairs(Style(dims(), value))
end

# data context: integers and symbols are columns

struct DataContext{NT, I<:AbstractVector{Int}} <: AbstractContext
    data::ColumnDict
    pkeys::NT
    perm::I
end

function DataContext(data)
    col = getcolumn(data, 1)
    DataContext(data, NamedTuple(), axes(col, 1))
end

data(x) = DataContext(ColumnDict(x))

iscategorical(v) = !(eltype(v) <: Number)

iswrappedcategorical(::Any) = false
iswrappedcategorical(v::AbstractArray{<:Any, 0}) = iscategorical(v[])

# unwrap categorical vectors from the 0-dim array that contains them
function unwrap_categorical(values)
    iter = (k => v[] for (k, v) in pairs(values) if iswrappedcategorical(v))
    return (; iter...)
end

function refine_perm(perm, pc, n)
    if n == length(pc)
        perm
    elseif n == 0
        sortperm(StructArray(pc))
    else
        x, y = pc[n], pc[n+1]
        refine_perm!(copy(perm), pc, n, x, y, 1, length(x))
    end
end

optimize_cols(v::NamedDimsArray) = refs(parent(v))
optimize_cols(v::Union{Tuple, NamedTuple}) = map(optimize_cols, v)

function _merge(c::DataContext, s1::Style, s2::Style)
    data, pkeys, perm = c.data, c.pkeys, c.perm
    nt = map(val -> extract_columns(data, val), s2.value)
    newpkeys = unwrap_categorical(keyword(nt))
    @assert isempty(intersect(keys(pkeys), keys(newpkeys)))
    allpkeys = merge(pkeys, newpkeys)
    # unwrap from NamedDimsArray to perform the sorting
    allperm = refine_perm(perm, optimize_cols(allpkeys), length(pkeys))
    ctx = DataContext(data, allpkeys, allperm)
    return Style(ctx, merge(s1.value, Base.structdiff(nt, newpkeys)))
end

function _pairs(c::DataContext, s::Style)
    data, pkeys, perm = c.data, c.pkeys, c.perm
    isempty(pkeys) && return pairs(Style(dims(), s.value))
    # unwrap for sorting computations
    itr = GroupPerm(StructArray(optimize_cols(pkeys)), perm)
    sa = StructArray(pkeys)
    nestedpairs = map(itr) do idxs
        i1 = perm[first(idxs)]
        # keep value categorical and preserve name by taking a mini slice (document this!)
        # consider using these levels directly in discrete scales
        k = map(col -> col[i1:i1], pkeys)
        cols = map(s.value) do val
            extract_views(val, perm[idxs])
        end
        [merge(k, p) => v for (p, v) in pairs(Style(dims(), cols))]
    end
    return reduce(vcat, nestedpairs)
end

extract_views(cols, idxs) = map(t -> view(t, idxs), cols)
extract_views(cols::DimsSelector, idxs) = cols

function extract_column(t, (nm, f)::Pair)
    colname = nm isa Symbol ? nm : columnnames(t)[nm]
    vals = f(getcolumn(t, nm))
    p = iscategorical(vals) ? categorical(vals) : vals
    return NamedDimsArray{(colname,)}(p)
end
extract_column(t, nm::Union{Symbol, Int}) = extract_column(t, nm => identity)

extract_columns(t, val::DimsSelector) = val
extract_columns(t, val::Union{Tuple, AbstractArray}) = map(v -> extract_column(t, v), val)
extract_columns(t, val) = fill(extract_column(t, val))


text:
using BenchmarkTools, Plots, Plots.PlotMeasures, LaTeXStrings
using BenchmarkTools: minimum, median

SUITE = BenchmarkGroup()
model = "CVDP21"
cases = ["μ=1.0", "μ=2.0"]
SUITE[model] = BenchmarkGroup()

include("vanderpol.jl")
validation = []

# ----------------------------------------
# Case 1: μ = 1
# ----------------------------------------

prob = vanderpolN2(μ=1.0);
alg = TMJets(abstol=1e-10, orderT=8, orderQ=1);

# warm-up run
sol_cvdp1 = solve(prob, T=7.0, alg=alg);
solz_cvdp1 = overapproximate(sol_cvdp1, Zonotope);

# verify that specification holds
property = (ρ(ey₁, solz_cvdp1) < 2.75) && (ρ(ey₂, solz_cvdp1) < 2.75)
push!(validation, Int(property))

# benchmark
SUITE[model][cases[1]] = @benchmarkable solve($prob, T=7.0, alg=$alg)

# ----------------------------------------
# Case 2: μ = 2
# ----------------------------------------

prob = vanderpolN2(μ=2.0);
alg = TMJets(abstol=1e-10, orderT=8, orderQ=1);

# warm-up run
sol_cvdp2 = solve(prob, T=8.0, alg=alg);
solz_cvdp2 = overapproximate(sol_cvdp2, Zonotope);

# verify that specification holds
property = (ρ(ey₁, solz_cvdp2) < 4.05) && (ρ(ey₂, solz_cvdp2) < 4.05)
push!(validation, Int(property))

# benchmark
SUITE[model][cases[2]] = @benchmarkable solve($prob, T=8.0, alg=$alg)

# ==============================================================================
# Execute benchmarks and save benchmark results
# ==============================================================================

# tune parameters
tune!(SUITE)

# run the benchmarks
results = run(SUITE, verbose=true)

# return the sample with the smallest time value in each test
println("minimum time for each benchmark:\n", minimum(results))

# return the median for each test
println("median time for each benchmark:\n", median(results))

# export runtimes
runtimes = Dict()
for (i, c) in enumerate(cases)
    t = median(results[model][c]).time * 1e-9
    runtimes[c] = t
end

for (i, c) in enumerate(cases)
    print(io, "JuliaReach, $model, $c, $(validation[i]), $(runtimes[c])\n")
end

# ==============================================================================
# Plot
# ==============================================================================

fig = Plots.plot()

fig = plot()
plot!(fig, solz_cvdp1,  vars=(1, 2), lw=0.0, alpha=1.0, color=:red)

plot!(fig, solz_cvdp2,  vars=(1, 2), lw=0.0, alpha=1.0, color=:blue,
    tickfont=font(30, "Times"), guidefontsize=45,
    xlab=L"x_{1}",
    ylab=L"y_1",
    xtick=[-2.0, 0.0, 2.0], ytick=[-4.0, -2.0, 0.0, 2.0, 4.0],
    xlims=(-2.5, 2.5), ylims=(-4.05, 4.05),
    bottom_margin=8mm, left_margin=2mm, right_margin=8mm, top_margin=3mm,
    size=(1000, 1000))

savefig(fig, "ARCH-COMP21-JuliaReach-VanDerPol.png")