text:
module SixDOF

using LinearAlgebra: norm, cross

export Control, MassProp, Reference
export AbstractAeroModel, AbstractPropulsionModel, AbstractInertialModel,
    AbstractAtmosphereModel, AbstractController
export StabilityDeriv, MotorPropBatteryDataFit, UniformGravitationalField,
    ConstantAtmosphere, ConstantController
export CO, COUNTER, COCOUNTER
export sixdof!

# ------ General Structs -------

"""
    State(x, y, z, phi, theta, psi, u, v, w, p, q, r)

State of the aircraft: positions in inertial frame, euler angles,
velocities in body frame, angular velocities in body frame.
"""
struct State{TF}
    x::TF  # position (inertial frame)
    y::TF
    z::TF
    phi::TF  # orientation, euler angles
    theta::TF
    psi::TF
    u::TF  # velocity (body frame)
    v::TF
    w::TF
    p::TF  # angular velocity (body frame)
    q::TF
    r::TF
end

"""
    Control(de, dr, da, df, throttle)

Define the control settings: delta elevator, delta rudder, delta aileron,
delta flaps, and throttle.
"""
struct Control{TF}
    de::TF  # elevator
    dr::TF  # rudder
    da::TF  # aileron
    df::TF  # rudder
    throttle::TF
end

"""
    MassProp(m, Ixx, Iyy, Izz, Ixz, Ixy, Iyz)

Mass and moments of inertia in the body frame.
Ixx = int(y^2 + z^2, dm)
Ixz = int(xz, dm)

Most aircraft are symmetric about y and so there is a convenience method
to specify only the nonzero components.
MassProp(m, Ixx, Iyy, Izz, Ixz)
"""
struct MassProp{TF}
    m::TF
    Ixx::TF
    Iyy::TF
    Izz::TF
    Ixz::TF
    Ixy::TF
    Iyz::TF
end

# most aircraft are symmetric in y
MassProp(m, Ixx, Iyy, Izz, Ixz) = MassProp(m, Ixx, Iyy, Izz, Ixz, zero(Ixx), zero(Ixx))

"""
    Reference(S, b, c)

The reference area, span, and chord used in the aerodynamic computations.
"""
struct Reference{TF}
    S::TF  # area
    b::TF  # span
    c::TF  # chord
end

# ----------------------------------------------

# --------------- Interfaces ---------------

abstract type AbstractAtmosphereModel end

"""
    wind(model::AbstractAtmosphereModel, state)

Compute wind velocities.

**Returns**
- Wi: wind velocities in inertial frame
- Wb: gust velocities in body frame (just a convenience to allow some velocities in body frame)
"""
function wind(model::AbstractAtmosphereModel, state)
    @warn "wind function not implemented for AbstractAtmosphereModel"
    Wi = [0.0, 0.0, 0.0]
    Wb = [0.0, 0.0, 0.0]
    return Wi, Wb
end

"""
    properties(model::AbstractAtmosphereModel, state)

Compute atmospheric density and the speed of sound.
"""
function properties(model::AbstractAtmosphereModel, state)
    @warn "properties function not implemented for AbstractAtmosphereModel"
    rho = 1.225  # sea-level properties
    asound = 340.3
    return rho, asound
end

"""
    gravity(model::AbstractAtmosphereModel, state)

Compute the local acceleration of gravity.
"""
function gravity(model::AbstractAtmosphereModel, state)
    @warn "gravity function not implemented for AbstractAtmosphereModel"
    g = 9.81
    return g
end

# ----

abstract type AbstractAeroModel end

"""
    aeroforces(model::AbstractAeroModel, atm::AbstractAtmosphereModel, state::State, control::Control, mp::MassProp, ref::Reference)

Compute the aerodynamic forces and moments in the body frame.
return F, M
"""
function aeroforces(model::AbstractAeroModel, atm, state, control, mp, ref)
    @warn "aeroforces function not implemented for AbstractAeroModel"
    # forces and moments in body frame
    F = [0.0, 0.0, 0.0]
    M = [0.0, 0.0, 0.0]
    return F, M
end

# ----

abstract type AbstractPropulsionModel end

"""
    propulsionforces(model::AbstractPropulsionModel, atm::AbstractAtmosphereModel, state::State, control::Control, mp::MassProp, ref::Reference)

Compute the propulsive forces and moments in the body frame.
return F, M
"""
function propulsionforces(model::AbstractPropulsionModel, atm, state, control, mp, ref)
    @warn "propulsionforces function not implemented for AbstractPropulsionModel"
    # forces and moments in body frame
    F = [0.0, 0.0, 0.0]
    M = [0.0, 0.0, 0.0]
    return F, M
end

# ----

abstract type AbstractInertialModel end

"""
    gravityforces(model::AbstractInertialModel, atm::AbstractAtmosphereModel, state::State, control::Control, mp::MassProp, ref::Reference)

Compute the gravitational forces and moments in the body frame.
return F, M
"""
function gravityforces(model::AbstractInertialModel, atm, state, control, mp, ref)
    @warn "gravityforces function not implemented for AbstractInertialModel"
    # forces and moments in body frame
    F = [0.0, 0.0, 0.0]
    M = [0.0, 0.0, 0.0]
    return F, M
end

# ----

abstract type AbstractController end

"""
    setcontrol(controller::AbstractController, time, atm::AbstractAtmosphereModel, state::State, lastcontrol::Control, mp::MassProp, ref::Reference)

Compute control state for next time step given current state
return control::Control
"""
function setcontrol(controller::AbstractController, time, atm, state, mp, ref)
    @warn "setcontrol function not implemented for AbstractController"
    control = Control(0.0, 0.0, 0.0, 0.0, 0.0)
    return control
end

# -----------------------------

# ------------- helper functions (private) --------------

"""
    inertialtobody(state)

Construct a rotation matrix from inertial frame to body frame

The assumed order of rotation is
1) psi radians about the z axis,
2) theta radians about the y axis,
3) phi radians about the x axis.

This is an orthogonal transformation so its inverse is its transpose.
"""
function inertialtobody(state)

    R = Array{eltype(state.phi)}(undef, 3, 3)

    cphi, ctht, cpsi = cos.([state.phi, state.theta, state.psi])
    sphi, stht, spsi = sin.([state.phi, state.theta, state.psi])

    R[1, 1] = ctht*cpsi
    R[1, 2] = ctht*spsi
    R[1, 3] = -stht

    R[2, 1] = sphi*stht*cpsi - cphi*spsi
    R[2, 2] = sphi*stht*spsi + cphi*cpsi
    R[2, 3] = sphi*ctht

    R[3, 1] = cphi*stht*cpsi + sphi*spsi
    R[3, 2] = cphi*stht*spsi - sphi*cpsi
    R[3, 3] = cphi*ctht

    return R

end

"""
    windtobody(alpha, beta)

Rotation matrix from wind frame to body frame.
- alpha: angle of attack
- beta: sideslip angle
"""
function windtobody(alpha, beta)

    ca, cb = cos.([alpha, beta])
    sa, sb = sin.([alpha, beta])

    Rwb = [ca*cb  -ca*sb  -sa;
           sb      cb     0.0;
           sa*cb  -sa*sb  ca]

    return Rwb
end

"""
    stabilitytobody(alpha, beta)

Rotation matrix from stability frame to body frame.
- alpha: angle of attack
"""
function stabilitytobody(alpha)

    ca = cos(alpha)
    sa = sin(alpha)

    Rsb = [ca  0.0  -sa;
           0.0 1.0  0.0;
           sa  0.0  ca]

    return Rsb
end

"""
    windaxes(atm::AbstractAtmosphereModel, state)

Compute relative velocity in wind axes (airspeed, aoa, sideslip)
"""
function windaxes(atm, state)

    # velocity vectors
    Vb = [state.u, state.v, state.w]
    Wi, Wb = wind(atm, state)

    Rib = inertialtobody(state)

    # relative wind
    Vrel = Vb - (Rib*Wi + Wb)

    # airspeed
    Va = norm(Vrel)

    # angle of attack
    alpha = atan(Vrel[3], Vrel[1])

    # sideslip
    beta = asin(Vrel[2] / Va)

    return Va, alpha, beta

end

# ----------------------------------------------------

# ------- Some Default Interface Implementations -----

"""
    StabilityDeriv(CL0, CLalpha, CLq, CLM, CLdf, CLde, alphas,
        CD0, U0, exp_Re, e, Mcc, CDdf, CDde, CDda, CDdr,
        CYbeta, CYp, CYr, CYda, CYdr, Clbeta,
        Clp, Clr, Clda, Cldr,
        Cm0, Cmalpha, Cmq, CmM, Cmdf, Cmde,
        Cnbeta, Cnp, Cnr, Cnda, Cndr)

Stability derivatives of the aircraft.  Most are self explanatory if you are
familiar with stability derivatives (e.g., CLalpha is dCL/dalpha or the lift curve slope).
Some less familiar ones include
- M: Mach number
- alphas: the angle of attack for stall
- U0: the speed for the reference Reynolds number CD0 was computed at
- exp_Re: the coefficient in the denominator of the skin friction coefficient (0.5 laminar, 0.2 turbulent)
- e: Oswald efficiency factor
- Mcc: crest critical Mach number (when compressibility drag rise starts)

"""
struct StabilityDeriv{TF} <: AbstractAeroModel
    CL0::TF
    CLalpha::TF
    CLq::TF
    CLM::TF
    CLdf::TF
    CLde::TF
    alphas::TF  # TODO: should probably do in terms of CLmax

    CD0::TF
    U0::TF  # velocity corresponding to Reynolds number of CD0  (TODO: rethink this)
    exp_Re::TF  # exponent for Reynolds number scaling. typical values: exp_Re = 0.5 laminar, 0.2 turbulent
    e::TF  # Oswald efficiency factor
    Mcc::TF  # crest-critical Mach number when compressibility drag rise starts (quartic)
    CDdf::TF
    CDde::TF
    CDda::TF
    CDdr::TF

    CYbeta::TF
    CYp::TF
    CYr::TF
    CYda::TF
    CYdr::TF

    Clbeta::TF
    Clp::TF
    Clr::TF
    Clda::TF
    Cldr::TF

    Cm0::TF
    Cmalpha::TF
    Cmq::TF
    CmM::TF
    Cmdf::TF
    Cmde::TF

    Cnbeta::TF
    Cnp::TF
    Cnr::TF
    Cnda::TF
    Cndr::TF
end

"""
A simple (mostly) linear aerodynamics model
"""
function aeroforces(sd::StabilityDeriv, atm, state, control, ref, mp)

    # airspeed, angle of attack, sideslip
    Va, alpha, beta = windaxes(atm, state)

    # Mach number and dynamic pressure
    rho, asound = properties(atm, state)
    Mach = Va / asound
    qdyn = 0.5 * rho * Va^2

    # rename for convenience
    p = state.p
    q = state.q
    r = state.r
    de = control.de
    df = control.df
    dr = control.dr
    da = control.da

    # lift
    CL = sd.CL0 + sd.CLalpha*alpha + sd.CLq*q *ref.c/(2*Va) + sd.CLM*Mach
        + sd.CLdf*df + sd.CLde*de

    em = exp(-50*(alpha - sd.alphas))
    ep = exp(50*(alpha + sd.alphas))
    sigma = (1 + em + ep)/((1 + em)*(1 + ep))
    CL = (1- sigma)*CL + sigma * 2 * sign(alpha)*sin(alpha)^2*cos(alpha)

    # drag
    CDp = sd.CD0*(Va/sd.U0)^sd.exp_Re
    CDi = CL^2/(pi*(ref.b^2/ref.S)*sd.e)
    CDc = (Mach < sd.Mcc) ? 0.0 : 20*(Mach - sd.Mcc)^4

    CD = CDp + CDi + CDc + abs(sd.CDdf*df) + abs(sd.CDde*de) + abs(sd.CDda*da) + abs(sd.CDdr*dr)

    # side force
    CY = sd.CYbeta*beta + (sd.CYp*p + sd.CYr*r)*ref.b/(2*Va) + sd.CYda*da + sd.CYdr*dr

    # rolling moment
    Cl = sd.Clbeta*beta + (sd.Clp*p + sd.Clr*r)*ref.b/(2*Va) + sd.Clda*da + sd.Cldr*dr

    # pitching moment
    Cm = sd.Cm0 + sd.Cmalpha*alpha + sd.Cmq*q * ref.c/(2*Va) + sd.CmM*Mach + sd.Cmdf*df + sd.Cmde*de

    # yawing moment
    Cn = sd.Cnbeta*beta + (sd.Cnp*p + sd.Cnr*r)*ref.b/(2*Va) + sd.Cnda*da + sd.Cndr*dr

    # transfer forces from wind to body axes
    Rwb = windtobody(alpha, beta)

    F = Rwb*[-CD, CY, -CL] * qdyn * ref.S

    M = Rwb*[Cl*ref.b, Cm*ref.c, Cn*ref.b] * qdyn * ref.S

    return F, M
end

@enum PropType CO=1 COUNTER=-1 COCOUNTER=0

"""
    MotorPropBatteryDataFit(CT2, CT1, CT0, CQ2, CQ1, CQ0, D, num, type,
        R, Kv, i0, voltage)

**Inputs**
- CT2, CT1, CT0: quadratic fit to propeller thrust coefficient of form: CT = CT2*J2 + CT1*J + CT0
- CQ2, CQ1, CQ0: quadratic fit to propeller torque coefficient of form: CQ = CQ2*J2 + CQ1*J + CQ0
- D: propeller diameter
- num: number of propellers
- type: CO (torques add), COUNTER (torques add but with minus sign), COCOUNTER (no torque, they cancel out)
- R: motor resistance
- Kv: motor Kv
- i0: motor no-load current
- voltage: battery voltage
"""
struct MotorPropBatteryDataFit{TF, TI, PropType} <: AbstractPropulsionModel
    # CT = CT2*J2 + CT1*J + CT0
    # CQ = CQ2*J2 + CQ1*J + CQ0
    CT2::TF  # prop data fit
    CT1::TF
    CT0::TF
    CQ2::TF
    CQ1::TF
    CQ0::TF
    D::TF  # prop diameter
    num::TI
    type::PropType

    R::TF  # motor resistance
    Kv::TF  # motor Kv
    i0::TF  # motor no-load current

    voltage::TF  # battery voltage
end

function propulsionforces(prop::MotorPropBatteryDataFit, atm, state, control, ref, mp)

    # airspeed, angle of attack, sideslip
    Va, _, _ = windaxes(atm, state)

    # density
    rho, _ = properties(atm, state)

    D = prop.D

    # determine torque for motor/prop match (quadratic equation)
    a = rho*D^5/(2*pi)^2 * prop.CQ0
    b = rho*D^4/(2*pi)*Va * prop.CQ1 + 1.0/(prop.R*prop.Kv)
    c = rho*D^3*Va^2 * prop.CQ2 - control.throttle*prop.voltage/(prop.R*prop.Kv) + prop.i0/prop.Kv
    Omega = (-b + sqrt(b^2 - 4*a*c))/(2*a)

    # advance ratio
    n = Omega/(2*pi)
    J = Va/(n*D)

    # thrust and torque
    CT = prop.CT0 + prop.CT1*J + prop.CT2*J^2
    CQ = prop.CQ0 + prop.CQ1*J + prop.CQ2*J^2

    T = prop.num * CT * rho * n^2 * D^4
    Q = prop.num * CQ * rho * n^2 * D^5 * Int(prop.type)

    return [T, 0, 0], [Q, 0, 0]
end

"""
    UniformGravitationalField()

Assumes center of mass and center of gravity are coincident.
"""
struct UniformGravitationalField <: AbstractInertialModel end

function gravityforces(model::UniformGravitationalField, atm, state, control, ref, mp)

    W = mp.m * gravity(atm, state)
    ct, cp = cos.([state.theta, state.phi])
    st, sp = sin.([state.theta, state.phi])

    Fg = W*[-st, ct*sp, ct*cp]
    Mg = [zero(W), zero(W), zero(W)]  # no gravitational moment

    return Fg, Mg
end

"""
    ConstantAtmosphere(Wi, Wb, rho, asound, g)

Constant atmospheric properties.
"""
struct ConstantAtmosphere{TF, TV<:AbstractVector{TF}} <: AbstractAtmosphereModel
    Wi::TV
    Wb::TV
    rho::TF
    asound::TF
    g::TF
end

function wind(atm::ConstantAtmosphere, state)
    return atm.Wi, atm.Wb
end

function properties(atm::ConstantAtmosphere, state)
    return atm.rho, atm.asound
end

function gravity(atm::ConstantAtmosphere, state)
    return atm.g
end

"""
    ConstantController(de, dr, da, df, throttle)

Just a dummy controller that outputs constant control outputs the whole time.
"""
struct ConstantController{TF} <: AbstractController
    de::TF
    dr::TF
    da::TF
    df::TF
    throttle::TF
end

function setcontrol(controller::ConstantController, time, atm, state, mp, ref)
    return Control(controller.de, controller.dr, controller.da, controller.df, controller.throttle)
end
# --------------------------------------------------------

# ------------- main functions (public) --------------

"""
    sixdof!(ds, s, params, time)

dynamic and kinematic ODEs.  Follows format used in DifferentialEquations package.
- s = x, y, z, phi, theta, psi, u, v, w, p, q, r (same order as State)
- params = control, massproperties, reference, aeromodel, propmodel, inertialmodel, atmmodel
"""
function sixdof!(ds, s, params, time)

    x, y, z, phi, theta, psi, u, v, w, p, q, r = s
    mp, ref, aeromodel, propmodel, inertialmodel, atmmodel, controller = params

    # ---- controller -------
    state = State(s...)
    control = setcontrol(controller, time, atmmodel, state, mp, ref)
    # -----------------------

    # --------- forces and moments ---------
    # aerodynamics
    Fa, Ma = aeroforces(aeromodel, atmmodel, state, control, ref, mp)

    # propulsion
    Fp, Mp = propulsionforces(propmodel, atmmodel, state, control, ref, mp)

    # weight
    Fg, Mg = gravityforces(inertialmodel, atmmodel, state, control, ref, mp)

    # total forces and moments
    F = Fa + Fp + Fg
    M = Ma + Mp + Mg

    # --------------------------------------

    # ----- derivative of state --------
    Vb = [u, v, w]
    omegab = [p, q, r]

    # linear kinematics
    Rib = inertialtobody(state)
    rdot = Rib' * Vb

    # angular kinematics
    phidot = p + (q*sin(phi) + r*cos(phi))*tan(theta)
    thetadot = q*cos(phi) - r*sin(phi)
    psidot = (q*sin(phi) + r*cos(phi))/cos(theta)

    # linear dynamics
    vdot = F/mp.m - cross(omegab, Vb)

    # angular dynamics
    I = [mp.Ixx -mp.Ixy -mp.Ixz;
         -mp.Iyz mp.Iyy -mp.Iyz;
         -mp.Ixz -mp.Iyz mp.Izz]
    omegadot = I \ (M - cross(omegab, I*omegab))

    # -------------------------

    # TODO: if we need more efficiency we can avoid allocating and then assigning.
    ds[1:3] = rdot
    ds[4] = phidot
    ds[5] = thetadot
    ds[6] = psidot
    ds[7:9] = vdot
    ds[10:12] = omegadot
end

end # module


text:
# IupDialog: Example
#   Creates a dialog showing an icon, the "DEFAULTESC" attribute and a simple menu

module dialog1_

using IUP

export
	dialog1

include("../src/libiup_h.jl")

#------------------------example html/examples/C/dialog1.c ----------------------------------
function dialog1()
	IupOpen()	#Initializes IUP

	# Creating dialog's icon
	img = icon_img()
	icon = IupImage(32, 32, img)
	IupSetAttribute (icon, "1", "255 255 255")
	IupSetAttribute (icon, "2", "000 000 000")
	IupSetAttribute (icon, "3", "226 226 226")
	IupSetAttribute (icon, "4", "128 128 128")
	IupSetAttribute (icon, "5", "192 192 192")
	IupSetAttribute (icon, "6", "000 000 255")
	IupSetHandle ("icon", icon)

	# Creating dialog's content
	quit_bt = IupButton("Quit")
	IupSetCallback(quit_bt, "ACTION", cfunction(quit_cb, Int, (),))
	IupSetHandle("quit", quit_bt)

	# Creating dialog's menu
	options = IupMenu(IupItem("Exit","quit_cb"))
	submenu = IupSubmenu("File", options)
	menu = IupMenu(submenu)
	IupSetHandle("menu", menu)

	# Creating main dialog
	dialog = IupDialog(IupVbox(quit_bt))
	IupSetAttribute(dialog, "TITLE", "IupDialog")
	IupSetAttribute(dialog, "MENU", "menu")
	IupSetAttribute(dialog, "CURSOR", "CROSS")
	IupSetAttribute(dialog, "ICON", "icon")
	IupSetAttribute(dialog, "DEFAULTESC", "quit")

	IupShowXY(dialog,IUP_CENTER,IUP_CENTER)

	IupMainLoop()                               # Initializes IUP main loop
	IupClose()                                  # And close it when ready
end

function quit_cb()
	return IUP_CLOSE
end

# defines icon's image
function icon_img()
	img = uint8([
	1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,
	1,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,3,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,6,6,6,6,6,6,5,5,6,6,5,5,6,6,5,5,6,6,6,6,6,6,6,5,5,5,4,2,
	1,3,5,5,6,6,6,6,6,6,5,5,6,6,5,5,6,6,5,5,6,6,6,6,6,6,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,6,5,5,6,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,6,5,5,6,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,6,6,6,6,6,6,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,6,6,6,6,6,5,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,6,6,5,5,5,5,6,6,5,5,6,6,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,6,6,6,6,6,6,5,5,6,6,6,6,6,6,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,6,6,6,6,6,6,5,5,5,6,6,6,6,5,5,5,6,6,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,4,2,
	1,3,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,2,
	1,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2
	])
end

end 	# module

ext: jl


text:
using TuringModels
using Test

@testset "TuringModels" begin

include("../scripts/00/basic_example.jl")

end


text:
# ------------------------------------------------------------------------------
# Dirichlet BDDC example
# ------------------------------------------------------------------------------

using LFAToolkit
using LinearAlgebra

# setup
mesh = Mesh2D(1.0, 1.0)
p = 3

# diffusion operator
diffusion = GalleryOperator("diffusion", p + 1, p + 1, mesh)

# Dirichlet BDDC preconditioner
bddc = DirichletBDDC(diffusion)

# compute operator symbols
A = computesymbols(bddc, [0.2], [π, π])
eigenvalues = real(eigvals(A))

# ------------------------------------------------------------------------------


text:
# This file is part of Intersections.jl
export umerge

"""
    _remove_empty!(P, L)

Inplace removal of empty lists
"""
function _remove_empty!(P, L)
    j = 1
    n = length(P)
    @inbounds for i in 1:n
        # @show (i, j, P[i] <= length(L[i]), P[i], length(L[i]))
        if P[i] <= length(L[i])
            if i != j
                P[j] = P[i]
                L[j] = L[i]
            end
            j += 1
        end
    end
    # @show j n P L

    while length(P) >= j
        pop!(P); pop!(L)
    end
end

"""
    _sort!(P, L)

Adaptive bubble sort, efficient than other approaches because we expect a few sets and almost sorted
"""
function _sort!(P::Vector{IType}, L) where {IType<:Integer}
    s = 1
    n = length(P)

    while s > 0
        s = 0
        for i in 1:n-1
            @inbounds if _get_key(L[i], P[i]) > _get_key(L[i+1], P[i+1])
                _swap_items!(P, i, i+1)
                _swap_items!(L, i, i+1)
                s += 1
            end
        end
    end
end

"""
    umerge(L, output=eltype(L[1])[], P=nothing; t=1)
    umerge(onmatch::Function, L, P=nothing; t::Int=1)

Merges posting lists in `L` and saves the union in `output`.
The method accepts a callback function `onmatch` that is called whenever an object occurs in at least `t` posting lists.

# Arguments:
- `L`: The array of posting lists, the array can be destroyed in the process.
- `P`: The array of current positions in posting lists, i.e., initial state as an array of ones of size ``|L|``.

# Keyword arguments
- `t`: A positive treshold (number of occurrences) to push or apply the callback function, i.e., ``t=1`` means for union.

# About the callback function

The callback signature is `onmatch(L, P, m)` where:

- `L` is an array of postings lists, the same of the input but maybe sorted and perhaps some lists could be removed by the process.
- `P` a list of indices of the current position of `L`, e.g., `L[i][P[i]]` to obtain the current position of i-th list.
- `m` the number of occurences of the current element pointed by `P`. Note that all `L[i][P[i]]` are the same for `1 ≤ i ≤ m` (it is an alignment position).

"""
function umerge(L_, output=eltype(L_[1])[], P=nothing; t::Int=1)
    umerge(L_, P; t) do L, P, m
        @inbounds m >= t && push!(output, _get_key(L[1], P[1]))
    end

    output
end

function umerge(onmatch::Function, L, P=nothing; t::Int=1)
    P = P === nothing ? ones(Int32, length(L)) : P
    _sort!(P, L)  # sort!(L, by=first)
    usize = 0

    @inbounds while true
        _remove_empty!(P, L)
        t > length(L) && break
        n = length(P)
        n == 0 && break
        n > 1 && _sort!(P, L)
        val = _get_key(L[1], P[1])
        m = 1
        @inbounds for i in 2:n
            if _get_key(L[i], P[i]) == val
                m += 1
            else
                break
            end
        end

        if m >= t
            onmatch(L, P, m)
            usize += 1
        end

        for i in 1:m
            P[i] += 1
        end
    end

    usize
end


text:
function convm(x, p)
    """Generates a convolution matrix

    Usage: X = convm(x,p)
    Given a vector x of length N, an N+p-1 by p convolution matrix is
    generated of the following form:
              |  x(0)  0      0     ...      0    |
              |  x(1) x(0)    0     ...      0    |
              |  x(2) x(1)   x(0)   ...      0    |
         X =  |   .    .      .              .    |
              |   .    .      .              .    |
              |   .    .      .              .    |
              |  x(N) x(N-1) x(N-2) ...  x(N-p+1) |
              |   0   x(N)   x(N-1) ...  x(N-p+2) |
              |   .    .      .              .    |
              |   .    .      .              .    |
              |   0    0      0     ...    x(N)   |

    That is, x is assumed to be causal, and zero-valued after N.
    """
    N = length(x) + 2 * p - 2
    xpad = [zeros(p - 1); x; zeros(p - 1)]
    X = zeros(length(x) + p - 1, p)
    # Construct X column by column
    for i = 0:p-1
        X[:, i+1] = xpad[p-i:N-i]
    end

    return X
end


text:
@testset "ExternalModel" begin
    sourcedir = tempdir()
    sourcefiles = ["in.txt"]
    extrafiles = []

    numberformats = Dict(:x => FormatSpec(".8e"), :* => FormatSpec(".8e"))

    workdir = joinpath(tempdir(), "external-model-test")

    r = Extractor(
        base -> begin
            map(x -> parse(Float64, x), readlines(joinpath(base, "out.txt")))[1]
        end, :r
    )

    binary = ""
    if Sys.iswindows()
        binary = joinpath(pwd(), "solvers/bin/radius.exe")
    elseif Sys.isapple()
        binary = joinpath(pwd(), "solvers/bin/radius-mac")
    else
        binary = joinpath(pwd(), "solvers/bin/radius")
    end

    opensees = Solver(binary, "", "in.txt")

    ext = ExternalModel(
        sourcedir, sourcefiles, extrafiles, numberformats, workdir, [r], opensees
    )

    open(joinpath(sourcedir, "in.txt"), "w") do input
        println(input, "{{{ :x }}}")
        println(input, "{{{ :y }}}")
    end

    x = RandomVariable(Uniform(0, 1), :x)
    y = RandomVariable(Uniform(0, 1), :y)

    df = sample([x, y], 1)

    evaluate!(ext, df)

    @test isapprox(df.r, sqrt.(df.x .^ 2 + df.y .^ 2))
end

max_stars_repo_name: FriesischScott/julia-uq

max_issues_repo_name: FriesischScott/julia-uq

max_forks_repo_name: FriesischScott/julia-uq


text:
# const TRAVIS_SKIP = VERSION.minor != 4 && !isnothing(get(ENV, "TRAVIS_BRANCH", nothing))

function test_fixed_size(A, At, B, Bt, Aa, Aat, Ba, Bat)::NTuple{4,Float64}
  t0 = @elapsed(@test Aa * Ba ≈ A * B)
  t1 = @elapsed(@test Aa * Bat ≈ A * Bt)
  t2 = @elapsed(@test Aat * Ba ≈ At * B)
  t3 = @elapsed(@test Aat * Bat ≈ At * Bt)
  t0, t1, t2, t3
end

function test_fixed_size(M, K, N)
  A = @StrideArray rand(M, K)
  B = @StrideArray rand(K, N)
  At = (@StrideArray rand(K, M))'
  Bt = (@StrideArray rand(N, K))'
  Aa = Array(A)
  Ba = Array(B)
  Aat = Array(At)
  Bat = Array(Bt)
  t0 = test_fixed_size(A, At, B, Bt, Aa, Aat, Ba, Bat)
  t1 = test_fixed_size(
    StrideArrays.make_dynamic(A),
    StrideArrays.make_dynamic(At),
    B,
    Bt,
    Aa,
    Aat,
    Ba,
    Bat,
  )
  t2 = test_fixed_size(
    A,
    At,
    StrideArrays.make_dynamic(B),
    StrideArrays.make_dynamic(Bt),
    Aa,
    Aat,
    Ba,
    Bat,
  )
  t3 = test_fixed_size(
    StrideArrays.make_dynamic(A),
    StrideArrays.make_dynamic(At),
    StrideArrays.make_dynamic(B),
    StrideArrays.make_dynamic(Bt),
    Aa,
    Aat,
    Ba,
    Bat,
  )
  gflops = let gflop = 2e-9M * K * N
    map((t0, t1, t2, t3)) do t
      gflop ./ t
    end
  end
  @show (M, K, N), gflops
  nothing
end

test_fixed_size(M) = test_fixed_size(M, M, M)

# gmul(A, B) = LinearAlgebra.generic_matmatmul('N','N', A, B)

@testset "MatrixMultiply" begin
  r = 2:12
  for M ∈ r, K ∈ r, N ∈ r
    test_fixed_size(M, K, N)
  end
  r = 13:33
  for M ∈ r
    test_fixed_size(M)
  end
  M = K = N = 80
  A = @StrideArray rand(M, K)
  B = @StrideArray rand(K, N)
  C = StrideArray{Float64}(undef, (StaticInt(M), StaticInt(N)))
  M, K, N = 23, 37, 19
  @views begin
    Av = A[1:M, 1:K]
    Bv = B[1:K, 1:N]
    Cv = C[1:M, 1:N]
    Avsl = A[StaticInt(1):M, StaticInt(1):K]
    Bvsl = B[StaticInt(1):K, StaticInt(1):N]
    Cvsl = C[StaticInt(1):M, StaticInt(1):N]
    Avsr = A[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(K)]
    Bvsr = B[StaticInt(1):StaticInt(K), StaticInt(1):StaticInt(N)]
    Cvsr = C[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(N)]
  end
  Creference = Array(Av) * Array(Bvsl)
  time = @elapsed mul!(Cv, Av, Bv)
  @test Creference ≈ Cv
  @show M, K, N, time
  time = @elapsed mul!(Cvsl, Avsl, Bvsl)
  @test Creference ≈ Cv
  @show M, K, N, time
  time = @elapsed mul!(Cvsr, Avsr, Bvsr)
  @test Creference ≈ Cv
  @show M, K, N, time

  Av = A[1:M, 1:K]
  Bv = B[1:K, 1:N]
  Cv = C[1:M, 1:N]
  Avsl = A[StaticInt(1):M, StaticInt(1):K]
  Bvsl = B[StaticInt(1):K, StaticInt(1):N]
  Cvsl = C[StaticInt(1):M, StaticInt(1):N]
  Avsr = A[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(K)]
  Bvsr = B[StaticInt(1):StaticInt(K), StaticInt(1):StaticInt(N)]
  Cvsr = C[StaticInt(1):StaticInt(M), StaticInt(1):StaticInt(N)]
  Creference = Array(Av) * Array(Bvsl)
  time = @elapsed mul!(Cv, Av, Bv)
  @test Creference ≈ Cv
  @show M, K, N, time
  time = @elapsed mul!(Cvsl, Avsl, Bvsl)
  @test Creference ≈ Cv
  @show M, K, N, time
  time = @elapsed mul!(Cvsr, Avsr, Bvsr)
  @test Creference ≈ Cv
  @show M, K, N, time
end


text:
# ------------------------------------------------------------------
# Copyright (c) 2018, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    StructuredGridData(data, X, Y, Z, ...)

Data spatially distributed on a structured grid where points are
georeferenced by coordinates `X`, `Y`, `Z`, ...

The `data` argument is a dictionary mapping variable names to Julia
arrays with the actual data.

## Examples

A very popular structured grid data format is NetCDF. Given 2D arrays
`LAT` and `LON` with coordinates and arrays with climate data `precipitation`,
`temperature`, the following code can be used to construct a structured grid:

```julia
julia> data = Dict(:precipitation => precipitation, :temperature => temperature)
julia> StructuredGridData(data, LAT, LON)
```
"""
struct StructuredGridData{T<:Real,N} <: AbstractSpatialData{T,N}
  data::Dict{Symbol,<:AbstractArray}
  coords::Matrix{T}

  # state fields
  dims::Dims{N}

  function StructuredGridData{T,N}(data, coords) where {N,T<:Real}
    sizes = [size(array) for array in values(data)]
    @assert length(unique(sizes)) == 1 "data dimensions must be the same for all variables"
    @assert prod(sizes[1]) == size(coords, 2) "data and coordinates arrays must have the same number of indices"

    new(data, coords, sizes[1])
  end
end

function StructuredGridData(data::Dict{Symbol,<:AbstractArray},
                            coordarrays::Vararg{<:AbstractArray{T},N}) where {N,T<:Real}
  sizes = [size(array) for array in coordarrays]
  @assert length(unique(sizes)) == 1 "coordinates arrays must have the same dimensions"

  coords = Matrix{T}(N, prod(sizes[1]))
  for (i, array) in enumerate(coordarrays)
    coords[i,:] .= array[:]
  end

  StructuredGridData{T,N}(data, coords)
end

coordinates(geodata::StructuredGridData{T,N}) where {N,T<:Real} = Dict(Symbol("x$i") => T for i=1:N)

variables(geodata::StructuredGridData) = Dict(var => eltype(array) for (var,array) in geodata.data)

npoints(geodata::StructuredGridData) = size(geodata.coords, 2)

function coordinates!(buff::AbstractVector{T}, geodata::StructuredGridData{T,N}, ind::Int) where {N,T<:Real}
  for i in 1:N
    @inbounds buff[i] = geodata.coords[i,ind]
  end
end

value(geodata::StructuredGridData, ind::Int, var::Symbol) = geodata.data[var][ind]

function Base.view(geodata::StructuredGridData{T,N}, inds::AbstractVector{Int}) where {N,T<:Real}
  data = Dict(var => view(array, inds) for (var,array) in geodata.data)
  coords = view(geodata.coords, :, inds)

  PointSetData(data, coords)
end

# ------------
# IO methods
# ------------
function Base.show(io::IO, geodata::StructuredGridData{T,N}) where {N,T<:Real}
  dims = join(geodata.dims, "×")
  print(io, "$dims StructuredGridData{$T,$N}")
end

function Base.show(io::IO, ::MIME"text/plain", geodata::StructuredGridData{T,N}) where {N,T<:Real}
  println(io, geodata)
  println(io, "  variables")
  varlines = ["    └─$var ($(eltype(array)))" for (var, array) in geodata.data]
  print(io, join(varlines, "\n"))
end


text:
for n = 1:2
    println(n)
    for n = 1:5
        println("  ", n)
    end
end

max_stars_repo_name: gabriel-doriath-dohler/pjulia

max_issues_repo_name: gabriel-doriath-dohler/pjulia

max_forks_repo_name: gabriel-doriath-dohler/pjulia


text:
type Z_data
    data::Array{Float64}
    Omega::Array{Float64}
end

type Dat
    y::Array{Float64}
    X::Array{Float64}
    K::Integer
    Z::Array{Float64}
    z_P::Array{Integer}
    sigma2_true::Float64
    beta_true::Array{Float64}
    gamma_true::Array{Integer}
    tau_true::Array{Float64}
end

"""
    generateZ(p, n, signal)

Generate geneteic data by Gaussian graphical model.

# Arguments
- `p::Integer`: dimension of each pathway.
- `n::Integer`: sample size.
- `signal::Bool`: if to generate functional pathway.

# Examples
```julia-repl
julia> Z =  generateZ(p, n, signal)
```
"""
function generateZ(p::Integer, n::Integer, signal::Bool)
	if signal
        Omega = matrixdepot("toeplitz", [1; .5; .25; fill(0, p - 3)]);
        # make Omega invertible
        Omega = fixMatrix(Omega, 2.0)
        @assert all(eig(Omega)[1] .> 0)
        # generate data
        Cov_True = inv(Omega);
        # make it symmetric
        for i = 1:(size(Cov_True, 1) - 1)
            for j = (i + 1): size(Cov_True, 1)
                Cov_True[j, i] = Cov_True[i, j]
            end
        end
        data = rand(MvNormal(Cov_True), n)'
	else
        Cov_True = eye(p) * 1/10000
        data = rand(MvNormal(Cov_True), n)'
        Omega = inv(Cov_True)
	end
	return Z_data(data, Omega)
end

"""
    generateData(p, n, signal)

Generate geneteic data and clinical data.

# Arguments
- `K::Integer`: dimension of each pathway.
- `p_beta::Integer`: number of clinical covariates.
- `n::Bool`: sample size.
- `df::Integer`: scale parameter.
- `num_sig::Bool`: number of significant pathways to include.

# Examples
```julia-repl
julia> generateData(K, p_beta, n, df, num_sig)
```
"""

function generateData(K::Integer, p_beta::Integer, n::Integer, df::Float64, num_sig::Integer)
    # set the true parameters
    sigma2_true = 1.0;
    p_tau = fill(20, K);
    beta_true = rand(Uniform(0, 10), p_beta);
    gamma_true = [fill(0, K - num_sig); fill(1, num_sig)]
    tau_true = gamma_true * df
    signal_true = gamma_true .== 1
    # generate Z(list)
    Z = []
    for k in 1:K
        Z = push!(Z, generateZ(p_tau[k], n, signal_true[k]))
    end
    # X is correlated with first 2 gens in the last pathway
    X = 3 * cos(Z[K].data[:, 1:p_beta]) .+ rand(Normal(), n)
    # generate random effects h
    H_true = zeros(n, K)
    for k in 1:K
        if gamma_true[k] != 0
            H_true[:, k] = rand(MvNormal(tau_true[k] * calcSigma(Z[k].data, Z[k].Omega)), 1)
        end
    end
    # give pathway membership and overwrite Z with z
    z_P = Int64[]
    z = zeros(n, sum(p_tau))
    for k in 1:K
        z_P = append!(z_P, fill(k, p_tau[k]))
        z[:, find(z_P .== k)] =  Z[k].data
    end
    # combine fixed and random effects
    y = X * beta_true + sum(H_true, 2) + rand(Normal(), n)
	return Dat(y, X, K, z, z_P, sigma2_true, beta_true, gamma_true, tau_true)
end


text:
"""
Provides the [`doctest`](@ref) function that makes sure that the `jldoctest` code blocks
in the documents and docstrings run and are up to date.
"""
module DocTests
using DocStringExtensions

import ..Documenter:
    Documenter,
    Documents,
    Expanders,
    Utilities

import Markdown, REPL

# Julia code block testing.
# -------------------------

"""
$(SIGNATURES)

Traverses the document tree and tries to run each Julia code block encountered. Will abort
the document generation when an error is thrown. Use `doctest = false` keyword in
[`Documenter.makedocs`](@ref) to disable doctesting.
"""
function doctest(doc::Documents.Document)
    if doc.user.doctest === :fix || doc.user.doctest
        @debug "running doctests."
        for (src, page) in doc.internal.pages
            empty!(page.globals.meta)
            for element in page.elements
                page.globals.meta[:CurrentFile] = page.source
                Documents.walk(page.globals.meta, page.mapping[element]) do block
                    doctest(block, page.globals.meta, doc, page)
                end
            end
        end
    else
        @debug "skipped doctesting."
    end
end

function doctest(block::Markdown.Code, meta::Dict, doc::Documents.Document, page)
    lang = block.language
    if startswith(lang, "jldoctest")
        # Define new module or reuse an old one from this page if we have a named doctest.
        name = match(r"jldoctest[ ]?(.*)$", split(lang, ';', limit = 2)[1])[1]
        sym = isempty(name) ? gensym("doctest-") : Symbol("doctest-", name)
        sandbox = get!(() -> Expanders.get_new_sandbox(sym), page.globals.meta, sym)

        # Normalise line endings.
        block.code = replace(block.code, "\r\n" => "\n")

        # parse keyword arguments to doctest
        d = Dict()
        idx = findfirst(c -> c == ';', lang)
        if idx !== nothing
            kwargs = Meta.parse("($(lang[nextind(lang, idx):end]),)")
            for kwarg in kwargs.args
                if !(isa(kwarg, Expr) && kwarg.head === :(=) && isa(kwarg.args[1], Symbol))
                    file = meta[:CurrentFile]
                    lines = Utilities.find_block_in_file(block.code, file)
                    @warn("""
                        invalid syntax for doctest keyword arguments in $(Utilities.locrepr(file, lines))
                        Use ```jldoctest name; key1 = value1, key2 = value2

                        ```$(lang)
                        $(block.code)
                        ```
                        """)
                    return false
                end
                d[kwarg.args[1]] = Core.eval(sandbox, kwarg.args[2])
            end
        end
        meta[:LocalDocTestArguments] = d

        for expr in [get(meta, :DocTestSetup, []); get(meta[:LocalDocTestArguments], :setup, [])]
            Meta.isexpr(expr, :block) && (expr.head = :toplevel)
            try
                Core.eval(sandbox, expr)
            catch e
                push!(doc.internal.errors, :doctest)
                @error("could not evaluate expression from doctest setup.",
                    expression = expr, exception = e)
                return false
            end
        end
        if occursin(r"^julia> "m, block.code)
            eval_repl(block, sandbox, meta, doc, page)
        elseif occursin(r"^# output$"m, block.code)
            eval_script(block, sandbox, meta, doc, page)
        else
            push!(doc.internal.errors, :doctest)
            file = meta[:CurrentFile]
            lines = Utilities.find_block_in_file(block.code, file)
            @warn("""
                invalid doctest block in $(Utilities.locrepr(file, lines))
                Requires `julia> ` or `# output`

                ```$(lang)
                $(block.code)
                ```
                """)
        end
       delete!(meta, :LocalDocTestArguments)
    end
    false
end
doctest(block, meta::Dict, doc::Documents.Document, page) = true

function doctest(block::Markdown.MD, meta::Dict, doc::Documents.Document, page)
    haskey(block.meta, :path) && (meta[:CurrentFile] = block.meta[:path])
    return true
end

# Doctest evaluation.

mutable struct Result
    block  :: Markdown.Code # The entire code block that is being tested.
    input  :: String # Part of `block.code` representing the current input.
    output :: String # Part of `block.code` representing the current expected output.
    file   :: String # File in which the doctest is written. Either `.md` or `.jl`.
    value  :: Any        # The value returned when evaluating `input`.
    hide   :: Bool       # Semi-colon suppressing the output?
    stdout :: IOBuffer   # Redirected stdout/stderr gets sent here.
    bt     :: Vector     # Backtrace when an error is thrown.

    function Result(block, input, output, file)
        new(block, input, rstrip(output, '\n'), file, nothing, false, IOBuffer())
    end
end

function eval_repl(block, sandbox, meta::Dict, doc::Documents.Document, page)
    for (input, output) in repl_splitter(block.code)
        result = Result(block, input, output, meta[:CurrentFile])
        for (ex, str) in Utilities.parseblock(input, doc, page; keywords = false)
            # Input containing a semi-colon gets suppressed in the final output.
            result.hide = REPL.ends_with_semicolon(str)
            (value, success, backtrace, text) = Utilities.withoutput() do
                Core.eval(sandbox, ex)
            end
            Core.eval(sandbox, Expr(:global, Expr(:(=), :ans, QuoteNode(value))))
            result.value = value
            print(result.stdout, text)
            if !success
                result.bt = backtrace
            end
        end
        checkresult(sandbox, result, meta, doc)
    end
end

function eval_script(block, sandbox, meta::Dict, doc::Documents.Document, page)
    # TODO: decide whether to keep `# output` syntax for this. It's a bit ugly.
    #       Maybe use double blank lines, i.e.
    #
    #
    #       to mark `input`/`output` separation.
    input, output = split(block.code, "# output\n", limit = 2)
    input  = rstrip(input, '\n')
    output = lstrip(output, '\n')
    result = Result(block, input, output, meta[:CurrentFile])
    for (ex, str) in Utilities.parseblock(input, doc, page; keywords = false)
        (value, success, backtrace, text) = Utilities.withoutput() do
            Core.eval(sandbox, ex)
        end
        result.value = value
        print(result.stdout, text)
        if !success
            result.bt = backtrace
            break
        end
    end
    checkresult(sandbox, result, meta, doc)
end

function filter_doctests(strings::NTuple{2, AbstractString},
                         doc::Documents.Document, meta::Dict)
    meta_block_filters = get(meta, :DocTestFilters, [])
    meta_block_filters == nothing && meta_block_filters == []
    doctest_local_filters = get(meta[:LocalDocTestArguments], :filter, [])
    for r in [doc.user.doctestfilters; meta_block_filters; doctest_local_filters]
        if all(occursin.((r,), strings))
            strings = replace.(strings, (r => "",))
        end
    end
    return strings
end

# Regex used here to replace gensym'd module names could probably use improvements.
function checkresult(sandbox::Module, result::Result, meta::Dict, doc::Documents.Document)
    sandbox_name = nameof(sandbox)
    mod_regex = Regex("(Main\\.)?(Symbol\\(\"$(sandbox_name)\"\\)|$(sandbox_name))[,.]")
    mod_regex_nodot = Regex(("(Main\\.)?$(sandbox_name)"))
    outio = IOContext(result.stdout, :module => sandbox)
    if isdefined(result, :bt) # An error was thrown and we have a backtrace.
        # To avoid dealing with path/line number issues in backtraces we use `[...]` to
        # mark ignored output from an error message. Only the text prior to it is used to
        # test for doctest success/failure.
        head = replace(split(result.output, "\n[...]"; limit = 2)[1], mod_regex  => "")
        head = replace(head, mod_regex_nodot => "Main")
        str  = replace(error_to_string(outio, result.value, result.bt), mod_regex => "")
        str  = replace(str, mod_regex_nodot => "Main")

        str, head = filter_doctests((str, head), doc, meta)
        # Since checking for the prefix of an error won't catch the empty case we need
        # to check that manually with `isempty`.
        if isempty(head) || !startswith(str, head)
            if doc.user.doctest === :fix
                fix_doctest(result, str, doc)
            else
                report(result, str, doc)
                push!(doc.internal.errors, :doctest)
            end
        end
    else
        value = result.hide ? nothing : result.value # `;` hides output.
        output = replace(rstrip(sanitise(IOBuffer(result.output))), mod_regex => "")
        str = replace(result_to_string(outio, value), mod_regex => "")
        # Replace a standalone module name with `Main`.
        str = rstrip(replace(str, mod_regex_nodot => "Main"))
        filteredstr, filteredoutput = filter_doctests((str, output), doc, meta)
        if filteredstr != filteredoutput
            if doc.user.doctest === :fix
                fix_doctest(result, str, doc)
            else
                report(result, str, doc)
                push!(doc.internal.errors, :doctest)
            end
        end
    end
    return nothing
end

# Display doctesting results.

function result_to_string(buf, value)
    value === nothing || Base.invokelatest(show, IOContext(buf, :limit => true), MIME"text/plain"(), value)
    return sanitise(buf)
end

function error_to_string(buf, er, bt)
    # Remove unimportant backtrace info.
    index = findlast(ptr -> Base.ip_matches_func(ptr, :eval), bt)
    # Print a REPL-like error message.
    print(buf, "ERROR: ")
    Base.invokelatest(showerror, buf, er, index === nothing ? bt : bt[1:(index - 1)])
    return sanitise(buf)
end

# Strip trailing whitespace from each line and return resulting string
function sanitise(buffer)
    out = IOBuffer()
    for line in eachline(seekstart(Base.unwrapcontext(buffer)[1]))
        println(out, rstrip(line))
    end
    return rstrip(String(take!(out)), '\n')
end

import .Utilities.TextDiff

function report(result::Result, str, doc::Documents.Document)
    diff = TextDiff.Diff{TextDiff.Words}(result.output, rstrip(str))
    lines = Utilities.find_block_in_file(result.block.code, result.file)
    @error("""
        doctest failure in $(Utilities.locrepr(result.file, lines))

        ```$(result.block.language)
        $(result.block.code)
        ```

        Subexpression:

        $(result.input)

        Evaluated output:

        $(rstrip(str))

        Expected output:

        $(result.output)

        """, diff)
end

function fix_doctest(result::Result, str, doc::Documents.Document)
    code = result.block.code
    filename = Base.find_source_file(result.file)
    # read the file containing the code block
    content = read(filename, String)
    # output stream
    io = IOBuffer(sizehint = sizeof(content))
    # first look for the entire code block
    # make a regex of the code that matches leading whitespace
    rcode = "(\\h*)" * replace(Utilities.regex_escape(code), "\\n" => "\\n\\h*")
    r = Regex(rcode)
    codeidx = findfirst(r, content)
    if codeidx === nothing
        @warn "could not find code block in source file"
        return
    end
    # use the capture group to identify indentation
    indent = match(r, content).captures[1]
    # write everything up until the code block
    write(io, content[1:prevind(content, first(codeidx))])
    # next look for the particular input string in the given code block
    # make a regex of the input that matches leading whitespace (for multiline input)
    rinput = "\\h*" * replace(Utilities.regex_escape(result.input), "\\n" => "\\n\\h*")
    r = Regex(rinput)
    inputidx = findfirst(r, code)
    if inputidx === nothing
        @warn "could not find input line in code block"
        return
    end
    # construct the new code-snippet (without indent)
    # first part: everything up until the last index of the input string
    newcode = code[1:last(inputidx)]
    isempty(result.output) && (newcode *= '\n') # issue #772
    # second part: the rest, with the old output replaced with the new one
    newcode *= replace(code[nextind(code, last(inputidx)):end], result.output => str, count = 1)
    # replace internal code block with the non-indented new code, needed if we come back
    # looking to replace output in the same code block later
    result.block.code = newcode
    # write the new code snippet to the stream, with indent
    newcode = replace(newcode, r"^(.+)$"m => Base.SubstitutionString(indent * "\\1"))
    write(io, newcode)
    # write rest of the file
    write(io, content[nextind(content, last(codeidx)):end])
    # write to file
    write(filename, seekstart(io))
    return
end

# REPL doctest splitter.

const PROMPT_REGEX = r"^julia> (.*)$"
const SOURCE_REGEX = r"^       (.*)$"
const ANON_FUNC_DECLARATION = r"#[0-9]+ \(generic function with [0-9]+ method(s)?\)"

function repl_splitter(code)
    lines  = split(string(code, "\n"), '\n')
    input  = String[]
    output = String[]
    buffer = IOBuffer()
    while !isempty(lines)
        line = popfirst!(lines)
        # REPL code blocks may contain leading lines with comments. Drop them.
        # TODO: handle multiline comments?
        # ANON_FUNC_DECLARATION deals with `x->x` -> `#1 (generic function ....)` on 0.7
        # TODO: Remove this special case and just disallow lines with comments?
        startswith(line, '#') && !occursin(ANON_FUNC_DECLARATION, line) && continue
        prompt = match(PROMPT_REGEX, line)
        if prompt === nothing
            source = match(SOURCE_REGEX, line)
            if source === nothing
                savebuffer!(input, buffer)
                println(buffer, line)
                takeuntil!(PROMPT_REGEX, buffer, lines)
            else
                println(buffer, source[1])
            end
        else
            savebuffer!(output, buffer)
            println(buffer, prompt[1])
        end
    end
    savebuffer!(output, buffer)
    zip(input, output)
end

function savebuffer!(out, buf)
    n = bytesavailable(seekstart(buf))
    n > 0 ? push!(out, rstrip(String(take!(buf)))) : out
end

function takeuntil!(r, buf, lines)
    while !isempty(lines)
        line = lines[1]
        if !occursin(r, line)
            println(buf, popfirst!(lines))
        else
            break
        end
    end
end

end


text:
# constructors

# values array + index array + fields singleton
TS(v::V, t::T, f::F) where {V<:VALARR,T<:IDXARR,F<:FLDTYPE} = TS(collect(hcat(v)), collect(t), [Symbol(f)])

# values array + index array + fields array
TS(v::V, t::T, f::F) where {V<:VALARR,T<:IDXARR,F<:FLDARR} = TS(collect(hcat(v)), collect(t), Symbol.(f))

# values array + index array
TS(v::V, t::T) where {V<:VALARR,T<:IDXARR} = TS(collect(hcat(v)), collect(t), autocol(1:size(v,2)))

# values array + index singleton + fields any
TS(v::V, t::T, f) where {V<:VALARR,T<:IDXTYPE} = TS(collect(hcat(v)), [t], f)

# values singleton + index array + fields any
TS(v::V, t::T, f) where {V<:VALTYPE,T<:IDXARR} = TS([v], collect(t), f)

# values singleton + index singelton + fields any
TS(v::V, t::T, f) where {V<:VALTYPE,T<:IDXTYPE} = TS([v][:,:], [t], f)

# values singleton + index array
TS(v::V, t::T) where {V<:VALTYPE, T<:IDXARR} = TS([v], collect(t), [:A])

# values singleton + index singleton
TS(v::V, t::T) where {V<:VALTYPE, T<:IDXTYPE} = TS([v], [t], [:A])

# values array
TS(v::V) where {V<:VALARR} = TS(collect(v), autoidx(size(v,1)), autocol(1:size(v,2)))

# default constructor
TS() = TS(zeros((0,0)), Date[], Symbol[])

# copy constructor
TS(X::TS) = copy(X)

# named tuple constructor
function TS(named_tuple::NamedTuple, index_element::Int=0)::TS
    function find_index_col(fields)::Int
        regex = r"date|time|index"i
        @inbounds for (j, field) in enumerate(fields)
            if occursin(regex, String(field))
                return j
            end
        end
        return 0
    end
    if index_element == 0
        index_element = find_index_col(keys(named_tuple))
    end
    index = named_tuple[index_element]
    array = zeros(Float64, (length(index), length(named_tuple)-1))
    value_elements = setdiff(1:length(named_tuple), index_element)
    fields = [f for f in keys(named_tuple)[value_elements]]
    @inbounds for (array_col, tuple_element) in enumerate(value_elements)
        array[:,array_col] = VALTYPE.(named_tuple[tuple_element])
    end
    return TS(array, index, fields)
end

# conversions

# boolean -> numeric
convert(::Type{TS{Bool}}, x::TS{V}) where {V<:VALTYPE} = TS{Bool}(map(V, x.values), x.index, x.fields)

## boolean -> float
#convert(::Type{TS{Float64}}, x::TS{Bool}) = TS{Float64}(map(Float64, x.values), x.index, x.fields)
#
## boolean -> integer
#convert(::Type{TS{Int}}, x::TS{Bool}) = TS{Int}(map(Int, x.values), x.index, x.fields)


text:
using .DifferentialEquations
const DE = DifferentialEquations

# extend the solve API for initial-value problems
function _solve_ensemble(ivp::InitialValueProblem, args...;
                         trajectories=100,
                         trajectories_alg=DE.Tsit5(),
                         ensemble_alg=DE.EnsembleThreads(),
                         inplace=true,
                         kwargs...)

    # get problem's vector field
    if inplace
        field = inplace_field!(ivp)
    else
        field = outofplace_field(ivp)
    end

    # get time span
    dt = _get_tspan(args...; kwargs...)
    tspan = (inf(dt), sup(dt))

    # sample initial states
    X0 = initial_state(ivp)
    sampler = get(kwargs, :sampler, LazySets._default_sampler(X0))
    rng = get(kwargs, :rng, LazySets.GLOBAL_RNG)
    seed = get(kwargs, :seed, nothing)
    include_vertices = get(kwargs, :include_vertices, false)
    X0_samples = sample(X0, trajectories; sampler=sampler, rng=rng, seed=seed, include_vertices=include_vertices)

    # formulate ensemble ODE problem
    ensemble_prob = ODEProblem(field, first(X0_samples), tspan)
    _prob_func(prob, i, repeat) = remake(prob, u0 = X0_samples[i])

    # choose tolerances
    reltol = get(kwargs, :reltol, 1e-3)
    abstol = get(kwargs, :abstol, 1e-6)

    ensemble_prob = EnsembleProblem(ensemble_prob, prob_func = _prob_func)
    return DE.solve(ensemble_prob, trajectories_alg, ensemble_alg;
                    trajectories=trajectories, reltol=reltol, abstol=abstol)
end


text:
using Random: shuffle!

"""
Default Data Loaders for ULMFiT training for Sentiment Analysis
 - WikiText-103 corpus is to pre-train the Language model
 - IMDB movie review dataset - unsup data is used for fine-tuning Language Mode for Sentiment Analysis
 - IMDB movie review dataset - labelled data is used for training classifier for Sentiment Analysis
"""

# WikiText-103 corpus loade

# This funciton given as an example of doing preprocessing for such corpus.
# In this IMDB corpus, it was noticed that after tokenization some tokens were mixture of
# some valid words and punctuations like ',', '.' etc which leads to increase in UNKNOWN token count in the corpus,
# to reduce the UNKNOWN count those tokens are splitted using `split_worda` and `put` functions.
function imdb_preprocess(doc::AbstractDocument)
    ## Edit here if any preprocessing step is needed ##
    function put(en, symbol)
        l = length(en)
        (l == 1) && return en
        for i=1:l-1
            insert!(en, i*2, string(symbol))
        end
        return en
    end
    function split_word(word, symbol)
        length(word) == 1 && return [word]
        return split(word, symbol)
    end
    text = text(doc)
    remove_corrupt_utf8!(text)
    remove_case!(text)
    prepare!(text, strip_html_tags)
    tokens = tokens(text)
    for symbol in [',', '.', '-', '/', "'s"]
        tokens = split_word.(tokens, symbol)
        temp = []
        for token in tokens
            try
                append!(temp, put(token, symbol))
            catch
                append!(temp, token)
            end
        end
        tokens = temp
    end
    deleteat!(tokens, findall(x -> isequal(x, "")||isequal(x, " "), tokens))
    return tokens
end

# Loads WikiText-103 corpus and output a Channel to give a mini-batch at each call
function load_wikitext_103(batchsize::Integer, bptt::Integer; type = "train")
    corpuspath = joinpath(datadep"WikiText-103", "wiki.$(type).tokens")
    corpus = read(open(corpuspath, "r"), String)
    corpus = tokenize(corpus)
    return Channel(x -> generator(x, corpus; batchsize = batchsize, bptt = bptt));
end

# IMDB Data loaders for Sentiment Analysis specifically
# IMDB data loader for fine-tuning Language Model
function imdb_fine_tune_data(batchsize::Integer, bptt::Integer, num_examples::Integer=50000)
    imdb_dataset = IMDB("train_unsup")
    dataset = []
    for path in imdb_dataset.filepaths   #extract data from the files in directory and put into channel
        open(path) do fileio
            cur_text = read(fileio, String)
            append!(dataset, imdb_preprocess(cur_text))
        end #open
    end #for
    return Channel(x -> generator(x, dataset; batchsize=batchsize, bptt=bptt))
end

# IMDB data loader for training classifier
function imdb_classifier_data(batchsize::Integer)
    filepaths = IMDB("train_neg").filepaths
    append!(filepaths, IMDB("train_pos").filepaths)
    [shuffle!(filepaths) for _=1:10]
    corpus = Corpus(FileDocument.(filepaths))
    Channel(csize=2) do docs
        n_batches = Int(floor(length(corpus)/batchsize))
        put!(docs, n_batches)
        for i=1:n_batches
            X, Y = [], []
            for j=1:batchsize
                path = filepaths[(i-1)*batchsize+j]
                cur_text = text(corpus[(i-1)*batchsize+j])
                tokens = imdb_preprocess(cur_text)
                push!(X, tokens)
                y = (parse(Int32, split(path, '_')[2][1:end-4]) > 5) ? Flux.onehotbatch([1], 1:2) : Flux.onehotbatch([2], 1:2)
                push!(Y, y)
            end#for
            X = pre_pad_sequences(X, "_pad_")
            put!(docs, [Flux.batch(X[k][j] for k=1:batchsize) for j=1:length(X[1])])
            put!(docs, cat(Y..., dims=2))
        end #for
    end #channel
end


text:
module Billard4Balls

using ModiaLang

import Modia3D
using  Modia3D.ModiaInterface

vmatBalls = Modia3D.VisualMaterial(color="Red", transparency=0.0)
vmatTable = Modia3D.VisualMaterial(color="DarkGreen", transparency=0.5)

diameter = 0.06
radius = diameter/2
distance_balls = sqrt(3)/2*diameter
dist = 0.0001
TableX = 2.24
TableY = 1.12
TableZ = 0.05
LyBox = 0.05
LzBox = 0.05
hz = 7*radius/5
hzz = hz/3

Table = Model(
    table = Object3D(parent=:world,
                     translation=[0.0, 0.0, -TableZ/2],
                     feature=Solid(shape=Box(lengthX=TableX, lengthY=TableY, lengthZ=TableZ),
                                   solidMaterial="BilliardTable",
                                   visualMaterial=vmatTable,
                                   collision=true,
                                   collisionSmoothingRadius=0.001))
)

Cushion = Model(
    cushion11 = Object3D(parent=:world,
                         translation=[TableX/2-LyBox/2, 0.0, hz/2],
                         feature=Solid(shape=Box(lengthX=LyBox, lengthY=TableY, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable)),
    cushion12 = Object3D(parent=:cushion11,
                         translation=[-LyBox/4, 0.0, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=1.5*LyBox, lengthY=TableY, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable,
                                       collision=true)),
    cushion31 = Object3D(parent=:world,
                         translation=[-TableX/2+LyBox/2, 0.0, hz/2],
                         feature=Solid(shape=Box(lengthX=LyBox, lengthY=TableY, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable)),
    cushion32 = Object3D(parent=:cushion31,
                         translation=[LyBox/4, 0.0, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=1.5*LyBox, lengthY=TableY, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable,
                                       collision=true)),
    cushion21 = Object3D(parent=:world,
                         translation=[0.0, TableY/2-LyBox/2, hz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=LyBox, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable)),
    cushion22 = Object3D(parent=:cushion21,
                         translation=[0.0, -LyBox/4, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=1.5*LyBox, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable,
                                       collision=true)),
    cushion41 = Object3D(parent=:world,
                         translation=[0.0, -TableY/2+LyBox/2, hz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=LyBox, lengthZ=hz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable)),
    cushion42 = Object3D(parent=:cushion41,
                         translation=[0.0, LyBox/4, hz/2+hzz/2],
                         feature=Solid(shape=Box(lengthX=TableX, lengthY=1.5*LyBox, lengthZ=hzz),
                                       solidMaterial="BilliardCushion",
                                       visualMaterial=vmatTable,
                                       collision=true))
)

Ball = Model(
    ball = Object3D(feature=Solid(shape=Sphere(diameter=diameter),
                                  solidMaterial="BilliardBall",
                                  visualMaterial=vmatBalls,
                                  collision=true))
)

Billard = Model(
    world = Object3D(feature=Scene(gravityField=UniformGravityField(g=9.81, n=[0, 0, -1]),
                                          enableContactDetection=true)),
    table = Table,
    cushion = Cushion,
    ball0 = Ball,
    joint0 = FreeMotion(obj1=:world, obj2=:(ball0.ball),
                        r=Var(init=[-0.8, -0.1, diameter/2]),
                        v=Var(init=[3.0, 0.1, 0.0]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball1 = Ball,
    joint1 = FreeMotion(obj1=:world, obj2=:(ball1.ball),
                        r=Var(init=[TableX/6, 0.0, diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball2 = Ball,
    joint2 = FreeMotion(obj1=:world, obj2=:(ball2.ball),
                        r=Var(init=[TableX/6+1*distance_balls+dist, 1/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0])),
    ball3 = Ball,
    joint3 = FreeMotion(obj1=:world, obj2=:(ball3.ball),
                        r=Var(init=[TableX/6+1*distance_balls+dist, -1/2*(diameter+dist), diameter/2]),
                        rot=Var(init=[pi/2, 0.0, 0.0]))
)

billard = @instantiateModel(buildModia3D(Billard), unitless=true, log=false, logStateSelection=false, logCode=false)

stopTime = 5.0
testTime = 2.5
tolerance = 1e-8
interval = 0.0001
requiredFinalStates = [-0.06367465272629756, 0.24238657622364665, 0.02999809758732255, -0.8055793547583174, 0.5049384414649515, -1.2942520100327887e-6, -41.98758798677127, 1.0302592492754234, -61.383962718912315, 7.327199670112295, 30.778026759406714, -1.8364952739701552, -0.15196276181074325, 0.179157995489763, 0.029997955278049118, -0.5749051726066565, 0.19607773653519803, -9.23765746465685e-7, 1.823582864167988, -0.11082435017660355, 18.026385421294833, -6.535777886518464, 3.361853105628269e-5, 19.163080641737785, 0.415730418200003, -0.25313065448897354, 0.02999753001826395, -0.004216104905417898, -0.49464670233658653, 3.901393071314717e-9, 11.54177678780837, 0.005646322377435987, 0.011097198279087365, 16.486263237425984, -0.25596215953403123, -0.026989059565206447, 0.15504140104701092, -0.18658321851036563, 0.029997948409715088, -0.297314286397064, -0.176662447164565, 4.779910373218372e-7, 0.9427627045555219, 0.6640319309374063, 12.884489822750403, -2.1107480248530672, -5.435105655586106, 9.944019158456165]
simulate!(billard, stopTime=testTime, tolerance=tolerance, interval=interval, log=true, logStates=true, logEvents=true, requiredFinalStates=requiredFinalStates)

@usingModiaPlot
plot(billard, ["joint0.r" "joint0.rot"; "joint0.v" "joint0.w"], figure=1)

end


text:
# It was proposed by Christian Goldbach that every odd composite number can be
# written as the sum of a prime and twice a square.
#
# 9 = 7 + 2×1^2
# 15 = 7 + 2×2^2
# 21 = 3 + 2×3^2
# 25 = 7 + 2×3^2
# 27 = 19 + 2×2^2
# 33 = 31 + 2×1^2
#
# It turns out that the conjecture was false.
#
# What is the smallest odd composite that cannot be written as the sum of a
# prime and twice a square?

using ProjectEulerSolutions

# Make a resonable list of twice squares and primes, then sum combinations of
# them and make a set that you compare against the set of odd composite
# numbers.  Would probably be faster to not calculate all the prime numbers
# and odd composite numbers upfront.
function p046solution(n::Integer=100)::Integer
    primearray = sieve_eratosthenes(n)
    primeset = Set(primearray)
    oddnumset = Set(9:2:n)
    oddcompositeset = setdiff(oddnumset, primeset)

    twicesquares = 2 .* (1:Integer(floor(0.75*sqrt(n)))).^2

    sums = Set()
    for ts in twicesquares
        union!(sums, Set(ts .+ primearray))
    end

    oddcomp = setdiff(oddcompositeset, sums)
    if length(oddcomp) > 0
        return minimum(oddcomp)
    else
        return -1
    end
end

p046 = Problems.Problem(p046solution)

Problems.benchmark(p046, 10_000)

max_stars_repo_name: gnujosh/julia-euler-project

max_issues_repo_name: gnujosh/julia-euler-project

max_forks_repo_name: gnujosh/julia-euler-project


text:
module CollectGemsMultiAgentUndirectedModule

import ..GridWorlds as GW
import Random
import ReinforcementLearningBase as RLBase

#####
##### game logic
#####

const NUM_ACTIONS = 4

mutable struct CollectGemsMultiAgentUndirected{R, RNG} <: GW.AbstractGridWorld
    tile_map::BitArray{3}
    agent_positions::Vector{CartesianIndex{2}}
    current_agent::Int
    reward::R
    rng::RNG
    done::Bool
    num_gem_init::Int
    num_gem_current::Int
    gem_reward::R
    gem_positions::Vector{CartesianIndex{2}}
end

function CollectGemsMultiAgentUndirected(; R = Float32, height = 8, width = 8, num_gem_init = floor(Int, sqrt(height * width)), num_agents = 4, rng = Random.GLOBAL_RNG)
    tile_map = falses(num_agents + 2, height, width)
    WALL = num_agents + 1
    GEM = num_agents + 2

    tile_map[WALL, 1, :] .= true
    tile_map[WALL, height, :] .= true
    tile_map[WALL, :, 1] .= true
    tile_map[WALL, :, width] .= true

    agent_positions = Array{CartesianIndex{2}}(undef, num_agents)
    for i in 1 : num_agents
        agent_position = GW.sample_empty_position(rng, tile_map)
        agent_positions[i] = agent_position
        tile_map[i, agent_position] = true
    end

    gem_positions = Array{CartesianIndex{2}}(undef, num_gem_init)
    for i in 1 : num_gem_init
        gem_position = GW.sample_empty_position(rng, tile_map)
        gem_positions[i] = gem_position
        tile_map[GEM, gem_position] = true
    end

    reward = zero(R)
    done = false
    gem_reward = one(R)
    num_gem_current = num_gem_init
    current_agent = 1

    env = CollectGemsMultiAgentUndirected(tile_map, agent_positions, current_agent, reward, rng, done, num_gem_init, num_gem_current, gem_reward, gem_positions)

    GW.reset!(env)

    return env
end

function GW.reset!(env::CollectGemsMultiAgentUndirected)
    tile_map = env.tile_map
    rng = env.rng
    agent_positions = env.agent_positions
    gem_positions = env.gem_positions
    num_agents = length(agent_positions)
    WALL = num_agents + 1
    GEM = num_agents + 2

    for i in 1 : num_agents
        tile_map[i, agent_positions[i]] = false
    end

    for i in 1 : env.num_gem_init
        tile_map[GEM, gem_positions[i]] = false
    end

    for i in 1 : num_agents
        agent_position = GW.sample_empty_position(rng, tile_map)
        agent_positions[i] = agent_position
        tile_map[i, agent_position] = true
    end

    for i in 1 : env.num_gem_init
        gem_position = GW.sample_empty_position(rng, tile_map)
        gem_positions[i] = gem_position
        tile_map[GEM, gem_position] = true
    end

    env.reward = zero(env.reward)
    env.done = false
    env.num_gem_current = env.num_gem_init

    return nothing
end

function GW.act!(env::CollectGemsMultiAgentUndirected, action)
    @assert action in Base.OneTo(NUM_ACTIONS) "Invalid action $(action). Action must be in Base.OneTo($(NUM_ACTIONS))"

    tile_map = env.tile_map
    agent_positions = env.agent_positions
    num_agents = length(agent_positions)
    current_agent = env.current_agent
    agent_position = agent_positions[current_agent]

    WALL = num_agents + 1
    GEM = num_agents + 2

    if action == 1
        new_agent_position = GW.move_up(agent_position)
    elseif action == 2
        new_agent_position = GW.move_down(agent_position)
    elseif action == 3
        new_agent_position = GW.move_left(agent_position)
    else
        new_agent_position = GW.move_right(agent_position)
    end

    if !any(@view tile_map[1 : num_agents + 1, new_agent_position]) # assuming WALL = num_agents + 1
        tile_map[current_agent, agent_position] = false
        agent_positions[current_agent] = new_agent_position
        tile_map[current_agent, new_agent_position] = true
    end

    agent_position = agent_positions[current_agent]

    env.reward = zero(env.reward)
    if tile_map[GEM, agent_position]
        tile_map[GEM, agent_position] = false
        env.num_gem_current -= 1
        env.reward = env.gem_reward
    end

    env.done = env.num_gem_current <= 0

    if current_agent == num_agents
        env.current_agent = 1
    else
        env.current_agent = current_agent + 1
    end

    return nothing
end

#####
##### miscellaneous
#####

GW.get_height(env::CollectGemsMultiAgentUndirected) = size(env.tile_map, 2)
GW.get_width(env::CollectGemsMultiAgentUndirected) = size(env.tile_map, 3)

GW.get_action_names(env::CollectGemsMultiAgentUndirected) = (:MOVE_UP, :MOVE_DOWN, :MOVE_LEFT, :MOVE_RIGHT)

function GW.get_object_names(env::CollectGemsMultiAgentUndirected)
    num_agents = length(env.agent_positions)
    object_names = Array{Symbol}(undef, num_agents + 2)
    for i in 1:num_agents
        object_names[i] = Symbol("AGENT", "$(i)")
    end
    object_names[end - 1] = :WALL
    object_names[end] = :GEM

    return object_names
end

function GW.get_pretty_tile_map(env::CollectGemsMultiAgentUndirected, position::CartesianIndex{2})
    tile_map = env.tile_map
    object = findfirst(@view tile_map[:, position])
    num_agents = size(tile_map, 1) - 2

    if isnothing(object)
        return "⋅"
    elseif object in 1 : num_agents
        return "$(object)"
    elseif object == num_agents + 1
        return "█"
    else
        return "♦"
    end
end

function GW.get_pretty_sub_tile_map(env::CollectGemsMultiAgentUndirected, window_size, position::CartesianIndex{2})
    tile_map = env.tile_map
    agent_positions = env.agent_positions
    agent_position = agent_positions[env.current_agent]
    num_agents = length(agent_positions)

    sub_tile_map = GW.get_sub_tile_map(tile_map, agent_position, window_size)

    object = findfirst(@view sub_tile_map[:, position])
    if isnothing(object)
        return "⋅"
    elseif object in 1 : num_agents
        return "$(object)"
    elseif object == num_agents + 1
        return "█"
    else
        return "♦"
    end
end

function Base.show(io::IO, ::MIME"text/plain", env::CollectGemsMultiAgentUndirected)
    str = "tile_map:\n"
    str = str * GW.get_pretty_tile_map(env)
    str = str * "\nsub_tile_map:\n"
    str = str * GW.get_pretty_sub_tile_map(env, GW.get_window_size(env))
    str = str * "\nreward: $(env.reward)\ndone: $(env.done)\ncurrent_agent = $(env.current_agent)"
    str = str * "\naction_names: $(GW.get_action_names(env))"
    str = str * "\nobject_names: $(GW.get_object_names(env))"
    print(io, str)
    return nothing
end

GW.get_action_keys(env::CollectGemsMultiAgentUndirected) = ('w', 's', 'a', 'd')

#####
##### RLBase API
#####

RLBase.StateStyle(env::GW.RLBaseEnv{E}) where {E <: CollectGemsMultiAgentUndirected} = RLBase.InternalState{Any}()
RLBase.state_space(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: CollectGemsMultiAgentUndirected} = nothing
RLBase.state(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: CollectGemsMultiAgentUndirected} = env.env.tile_map

RLBase.reset!(env::GW.RLBaseEnv{E}) where {E <: CollectGemsMultiAgentUndirected} = GW.reset!(env.env)

RLBase.action_space(env::GW.RLBaseEnv{E}) where {E <: CollectGemsMultiAgentUndirected} = Base.OneTo(NUM_ACTIONS)
(env::GW.RLBaseEnv{E})(action) where {E <: CollectGemsMultiAgentUndirected} = GW.act!(env.env, action)

RLBase.reward(env::GW.RLBaseEnv{E}) where {E <: CollectGemsMultiAgentUndirected} = env.env.reward
RLBase.is_terminated(env::GW.RLBaseEnv{E}) where {E <: CollectGemsMultiAgentUndirected} = env.env.done

end # module


text:
# # Comparing several models for water waves
#
#md # [`notebook`](@__NBVIEWER_ROOT_URL__notebooks/WWvsXX.ipynb)
#

using Distributed
using WaterWaves1D
include("../src/models/WaterWaves.jl")
include("../src/models/PseudoSpectral.jl")

rmprocs(workers())
addprocs(3)

@everywhere begin
    using Pkg
    Pkg.activate(".")
    using WaterWaves1D
    include("../src/models/WaterWaves.jl")
    include("../src/models/PseudoSpectral.jl")

end

using Test
using TimerOutputs

function run_simulation()

    reset_timer!()

    @everywhere begin

        #---- parameters
        param = ( μ  = .1,
       		  ϵ  = 1,
                  N  = 2^11,  # number of collocation points
                  L  = 10,    # size of the mesh (-L,L)
                  T  = 5,     # final time of computation
                  dt = 0.001, # timestep
        				);
        #---- initial data
        g(x) = exp.(-abs.(x).^4);
        z(x) = 0*exp.(-x.^2);
        init = Init(g,z);

        #---- models to compare
        models=[]
        push!(models,WaterWaves(param))
        push!(models,PseudoSpectral(param;order=2,dealias=1,lowpass=1/100))
        push!(models,PseudoSpectral(param;order=3,dealias=1,lowpass=1/100))

        function solve_problem(model)
            problem = Problem(model, init, param)
            WaterWaves1D.solve!(problem)
            last(problem.data.U)
        end

    end

    @timeit "Consecutive solving" begin
        p2 = []
        for model in models
            println(model.label)
            push!(p2, solve_problem(model))
        end
    end

    @timeit "Distributed" begin

        rc = RemoteChannel(()->Channel(3));
        pids = collect(workers())
        @sync for model in models
             @async put!(rc, (model.label,solve_problem(model)))
        end

        results = (take!(rc) for _ in 1:3)

        p1 = Dict( k => v for (k,v) in results)

        @show keys(p1)
    end

    print_timer()

    #---- Tests
    @testset "Final" begin
     	@test p1["water waves"] == p2[1]
     	@test p1["WW2"] == p2[2]
     	@test p1["WW3"] == p2[3]
    end

end

run_simulation()


text:
"""
a lot of evaluation metrics can be found in
https://pymia.readthedocs.io/en/latest/pymia.evaluation.metric.html?highlight=dice
"""

module MainOverlapTest

using Distances, Test, Revise
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\kernelEvolutions.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\utils\\CUDAGpuUtils.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\structs\\BasicStructs.jl")
includet("C:\\GitHub\\GitHub\\NuclearMedEval\\src\\overLap\\MainOverlap.jl")
using ..BasicStructs
using ..BasicPreds, ..CUDAGpuUtils,Cthulhu,BenchmarkTools , CUDA, ..MainOverlap

goldBoolGPU,segmBoolGPU,tp,tn,fp,fn, tpArr,tnArr,fpArr, fnArr, blockNum , nx,ny,nz ,tpTotalTrue,tnTotalTrue,fpTotalTrue, fnTotalTrue ,tpPerSliceTrue,  tnPerSliceTrue,fpPerSliceTrue,fnPerSliceTrue ,flattG, flattSeg ,FlattGoldGPU,FlattSegGPU,intermediateResTp,intermediateResFp,intermediateResFn = getSmallTestBools();

#tpTotalTrue,tnTotalTrue,fpTotalTrue, fnTotalTrue
@testset "jaccard test" begin
    dicee = dice(tpTotalTrue,fpTotalTrue, fnTotalTrue)
    x =  (dicee)/(2-dicee)
    jaccardOut = 1- evaluate(Jaccard(), flattG, flattSeg)
    @test jaccardOut ≈ jaccard(tpTotalTrue,fpTotalTrue, fnTotalTrue )
    @test jaccardOut ≈ x

end;

@testset "global consistency error test" begin
    local_refinement_error( flattG, flattSeg)
end;

end# MainOverlapTest


text:
import Base: ==

using .NapiTypes: NapiValue
abstract type NodeValue end
abstract type JsValue end
abstract type JsObjectType <: JsValue end

const ValueTypes = Union{NapiValue, NodeValue, JsValue}

Base.convert(::Type{T}, v::T) where T <: ValueTypes = v
Base.convert(::Type{NapiValue}, v::T) where T = NapiValue(v)
Base.convert(::Type{NapiValue}, v::ValueTypes) = NapiValue(v)
NapiValue(v) = napi_value(v)
napi_value(v::NapiValue) = v

==(a::ValueTypes, b::ValueTypes) = @with_scope try
    @napi_call napi_strict_equals(a::NapiValue, b::NapiValue)::Bool
catch _
    false
end

JsValue(v; this=nothing) = value(v; this=this)
value(v; this=nothing) = v
value(v::ValueTypes; this=nothing) = @with_scope value(NapiValue(v), this=this)
JsValue(::Type{T}, v) where T = value(T, v)
Base.convert(::Type{Any}, v::ValueTypes) = v
Base.convert(::Type{T}, v::ValueTypes) where T = value(T, v)
value(::Type{T}, v::T) where T = v
value(::Type{T}, v) where T = @with_scope value(T, NapiValue(v))
value(::Type{T}, v::NapiValue) where T = error("Unimplemented to convert a NapiValue to $T")

function value(napi_value::NapiValue; this=nothing)
    t = get_type(napi_value)
    if t ∈ (NapiTypes.napi_undefined, NapiTypes.napi_null)
        nothing
    elseif t == NapiTypes.napi_boolean
        value(Bool, napi_value)
    elseif t == NapiTypes.napi_number
        value(Float64, napi_value)
    elseif t == NapiTypes.napi_string
        value(String, napi_value; is_string = true)
    elseif t == NapiTypes.napi_symbol
        value(Symbol, napi_value)
    elseif t == NapiTypes.napi_object
        object_value(napi_value)
    elseif t == NapiTypes.napi_function
        value(JsFunction, napi_value, this=this)
    elseif t == NapiTypes.napi_external
        value(NodeExternal, napi_value)
    elseif t == NapiTypes.napi_bigint
        value(BigInt, napi_value)
    else
        nothing
    end
end

# For being called in C++ code.
@noinline napi_value_from_value(v) = napi_value(v)
@noinline napi_value_to_value(nv::Ptr{Nothing}) = value(convert(NapiValue, nv))

macro global_js_const(def, is_object=true)
    @assert def isa Expr && def.head == :(=) && length(def.args) == 2
    typ = is_object ? :NodeObject : :NodeValueTemp
    name, script = def.args
    # Just a temporary use of `ObjectReference`
    esc(quote
        const $name = $typ()
        if !haskey(ObjectReference, :global_init)
            ObjectReference[:global_init] = Ref(Tuple{
                Union{NodeObject, NodeValueTemp},
                String
            }[])
        end
        push!(ObjectReference[:global_init][], ($name, $script))
    end)
end

_initialize_globals() = @with_scope begin
    for (ref, script) in ObjectReference[:global_init][]
        nv = run_script(script, RESULT_RAW)
        if ref isa NodeObject
            v = @napi_call napi_create_reference(nv::NapiValue, 1::UInt32)::NapiRef
            setfield!(ref, :ref, v)
        else
            v = getfield(NodeValueTemp(nv), :tempname)
            setfield!(ref, :tempname, v)
        end
    end
    delete!(ObjectReference, :global_init)
end