text:
#!/usr/bin/env julia

module XboxCarSim

const MAX_LONGITUDINAL_ACC = 5.0
const MAX_LONGITUDINAL_BRAKE = -8.0
const MAX_CURVATURE = 1/20    # 1/(turning radius)
const SIM_RATE = 60

using PyCall
using RobotOS
using StaticArrays
using Rotations
using DifferentialDynamicsModels
using SimpleCarModels

@pyimport tf2_ros
@pyimport rospy
@rosimport safe_traffic_weaving.msg: XYThV
@rosimport sensor_msgs.msg: Joy
@rosimport std_msgs.msg: Bool, Header
@rosimport geometry_msgs.msg: Transform, TransformStamped, PoseStamped, TwistStamped, AccelStamped
@rosimport visualization_msgs.msg: Marker, MarkerArray
rostypegen()
import safe_traffic_weaving.msg: XYThV
import sensor_msgs.msg: Joy
import std_msgs.msg: BoolMsg, Header
import geometry_msgs.msg: Transform, TransformStamped, PoseStamped, TwistStamped, AccelStamped
import visualization_msgs.msg: Marker, MarkerArray

include("utils/utils.jl")

const terminate_pub = Publisher{BoolMsg}("/terminate", queue_size=10)
const reset_pub = Publisher{BoolMsg}("/reset", queue_size=1)
const reset_sim_pub = Publisher{BoolMsg}("/reset_sim", queue_size=1)
const xbox_car_pose_pub = Publisher{PoseStamped}("/xbox_car/pose", queue_size=10)    # w.r.t. /world frame
const xbox_car_vel_pub = Publisher{TwistStamped}("/xbox_car/vel", queue_size=10)     # w.r.t. /world frame
const xbox_car_acc_pub = Publisher{AccelStamped}("/xbox_car/acc", queue_size=10)     # w.r.t. /world frame
const xbox_car_xythv_pub = Publisher{XYThV}("/xbox_car/xythv", queue_size=10)
const xbox_car_marker_pub = Publisher{MarkerArray}("/xbox_car/viz", queue_size=10)
const xbox_car_init_marker_pub = Publisher{MarkerArray}("/xbox_car/init_viz", queue_size=10)
const tf_broadcaster = tf2_ros.TransformBroadcaster()
const tfBuffer = tf2_ros.Buffer()
const listener = tf2_ros.TransformListener(tfBuffer)

const time_scale_factor = RobotOS.get_param("time_scale_factor", 1.0)
const x_scale_factor = RobotOS.get_param("x_scale_factor", 1.0)

@show time_scale_factor
@show x_scale_factor

mutable struct SimulatorState
    terminate::Bool
    x::SE2vState{Float64}
    u::AccelerationCurvatureControl{Float64}
    x0::SE2vState{Float64}
    v0::Float64
    x0_offset_mode::Bool

    prev_joystick_time::Float64
end

function initial_state_callback(msg::XYThV, sim_state::SimulatorState)
    sim_state.x0 = SE2vState(XYThV.x, XYThV.y, XYThV.th, XYThV.v)
end

function x1_vel_callback(msg::TwistStamped, sim_state::SimulatorState)
    sim_state.v0 = norm([msg.twist.linear.x, msg.twist.linear.y])
end

function joystick_callback(msg::Joy, sim_state::SimulatorState)
    curr_joystick_time = Float64(RobotOS.get_rostime())
    dt = sim_state.prev_joystick_time == 0 ? 0.0 : curr_joystick_time - sim_state.prev_joystick_time
    sim_state.prev_joystick_time = curr_joystick_time

    if msg.buttons[1] == 1 && msg.buttons[2] == 1    # A and B (green and red)
        publish(terminate_pub, BoolMsg(true))
        sim_state.terminate = true
    end
    if msg.buttons[3] == 1                           # X (blue)
        sim_state.x0_offset_mode = !sim_state.x0_offset_mode
    end
    if msg.buttons[4] == 1                           # Y (yellow)
        sim_state.x = sim_state.x0
        if sim_state.x0_offset_mode
            try
                x1_tf = tfBuffer[:lookup_transform]("world", "xbox_car_init", rospy.Time())
                x = x1_tf[:transform][:translation][:x]
                y = x1_tf[:transform][:translation][:y]
                θ = RotZYX(Quat(x1_tf[:transform][:rotation][:w],
                                x1_tf[:transform][:rotation][:x],
                                x1_tf[:transform][:rotation][:y],
                                x1_tf[:transform][:rotation][:z])).theta1
                v = sim_state.v0 + sim_state.x0.v
                sim_state.x = SE2vState(x, y, θ, v)
            end
        end
        sim_state.u = zeros(sim_state.u)
    end
    if msg.buttons[5] == 1                           # left bumper
        RobotOS.set_param("robot_start_lane", "R")
        RobotOS.set_param("human_start_lane", "L")
    end
    if msg.buttons[6] == 1                           # right bumper
        RobotOS.set_param("robot_start_lane", "L")
        RobotOS.set_param("human_start_lane", "R")
    end
    if msg.buttons[7] == 1                           # select
        for i in 1:5
            rate = Rate(100)
            publish(reset_sim_pub, BoolMsg(true))
            rossleep(rate)
        end
    end
    if msg.buttons[8] == 1                           # start
        publish(reset_pub, BoolMsg(true))
    end
    if sim_state.x0_offset_mode
        if norm(msg.axes[4:5]) > .1
            sim_state.x0 = SE2vState(sim_state.x0.x + msg.axes[5]*5*dt,
                                     sim_state.x0.y + msg.axes[4]*5*dt,
                                     sim_state.x0.θ,
                                     sim_state.x0.v)
        end
        if abs(msg.axes[8]) > .1
            sim_state.x0 = SE2vState(sim_state.x0.x,
                                     sim_state.x0.y,
                                     sim_state.x0.θ,
                                     sim_state.x0.v + .05*msg.axes[8])
        end
    end
    accel_frac = (1 - msg.axes[6])/2
    brake_frac = (1 - msg.axes[3])/2
    sim_state.u = AccelerationCurvatureControl((MAX_LONGITUDINAL_ACC*accel_frac + MAX_LONGITUDINAL_BRAKE*brake_frac)*x_scale_factor*time_scale_factor^2,
                                               MAX_CURVATURE*msg.axes[1]/x_scale_factor)
end

init_node("xbox_car", anonymous=true)

const sim_state = SimulatorState(false, zeros(SE2vState{Float64}), zeros(AccelerationCurvatureControl{Float64}), zeros(SE2vState{Float64}), 0.0, true, 0.0)
const x0_sub = Subscriber{XYThV}("/xbox_car/set_init", initial_state_callback, (sim_state,), queue_size=10)
const X1_sub = Subscriber{TwistStamped}("/x1/vel", x1_vel_callback, (sim_state,), queue_size=10)
const joystick_sub = Subscriber{Joy}("/joy", joystick_callback, (sim_state,), queue_size=10)

function fill_2D_transform!(t::TransformStamped, x, y, θ, stamp = RobotOS.now())
    t.header.stamp = stamp
    t.transform.translation.x = x
    t.transform.translation.y = y
    q = Quat(RotZYX(θ, 0, 0))    # tf.transformations[:quaternion_from_euler](0.0, 0.0, θ)
    t.transform.rotation.x = q.x
    t.transform.rotation.y = q.y
    t.transform.rotation.z = q.z
    t.transform.rotation.w = q.w
    t
end

function run(x0 = zeros(SE2vState{Float64}), marker_color="lime")
    rate = Rate(SIM_RATE)

    world_header = Header()
    world_header.frame_id = "world"
    xbox_car_wrt_world = TransformStamped(world_header, "xbox_car", Transform())
    init_wrt_xbox_car = TransformStamped(Header(), "xbox_car_init", Transform())

    pose_msg, vel_msg, acc_msg = PoseStamped(), TwistStamped(), AccelStamped()
    xythv_msg = XYThV()
    pose_msg.header.frame_id = vel_msg.header.frame_id = acc_msg.header.frame_id = "world"
    # pose_msg.pose.orientation.w = 1    # pose is at the origin in the /xbox_car frame
    xbox_car_marker = MarkerArray([car_marker("xbox_car", marker_color),
                                   text_marker("xbox_car", "white", 5.0)])
    init_marker = MarkerArray([car_marker("xbox_car_init", marker_color, .2),
                               text_marker("xbox_car_init", "white", 5.0)])

    sim_state.x = sim_state.x0 = x0
    sim_state.terminate = false
    while !sim_state.terminate
        timestamp = RobotOS.now()
        sim_state.x = propagate(SimpleCarDynamics{1,0}(), sim_state.x, StepControl(1/SIM_RATE, sim_state.u))
        fill_2D_transform!(xbox_car_wrt_world, sim_state.x.x, sim_state.x.y, sim_state.x.θ, timestamp)
        fill_2D_transform!(init_wrt_xbox_car, sim_state.x0.x, sim_state.x0.y, sim_state.x0.θ, timestamp)
        init_wrt_xbox_car.header.frame_id = sim_state.x0_offset_mode ? "x1" : "world"

        pose_msg.header.stamp = vel_msg.header.stamp = acc_msg.header.stamp = timestamp
        xbox_car_marker.markers[2].text = "$(@sprintf("%.2f", sim_state.x.v))"
        init_marker.markers[2].text = "$(@sprintf("%.2f", sim_state.v0))"
        if sim_state.x0_offset_mode
            init_marker.markers[2].text = init_marker.markers[2].text*(sim_state.x0.v >= 0 ? "+" : "-")*"$(@sprintf("%.2f", abs(sim_state.x0.v)))"
        end
        # init_marker.markers[2].text = "$(@sprintf("%.2f", sim_state.x0_offset_mode ? sim_state.v0 + sim_state.x0.v : sim_state.x0.v))"
        # init_marker.markers[2].text = "$(@sprintf("%.2f", sim_state.x0.v))"
        # vel_msg.twist.linear.x = sim_state.x.v
        # acc_msg.accel.linear.x = sim_state.u.a
        # acc_msg.accel.linear.y = sim_state.u.κ*sim_state.x.v^2

        p, o = pose_msg.pose.position, pose_msg.pose.orientation
        p.x, p.y = sim_state.x.x, sim_state.x.y
        q =  Quat(RotZYX(sim_state.x.θ, 0, 0))    # tf.transformations[:quaternion_from_euler](0.0, 0.0, sim_state.x.θ)
        o.x, o.y, o.z, o.w = q.x, q.y, q.z, q.w
        v = vel_msg.twist.linear
        v.x, v.y = sim_state.x.v*cos(sim_state.x.θ), sim_state.x.v*sin(sim_state.x.θ)
        a = acc_msg.accel.linear
        a.x = sim_state.u.a*cos(sim_state.x.θ) - sim_state.u.κ*sim_state.x.v^2*sin(sim_state.x.θ)
        a.y = sim_state.u.a*sin(sim_state.x.θ) + sim_state.u.κ*sim_state.x.v^2*cos(sim_state.x.θ)
        xythv_msg.x, xythv_msg.y, xythv_msg.th, xythv_msg.v =
            sim_state.x.x, sim_state.x.y, sim_state.x.θ, sim_state.x.v

        tf_broadcaster[:sendTransform](xbox_car_wrt_world)
        tf_broadcaster[:sendTransform](init_wrt_xbox_car)
        publish(xbox_car_pose_pub, pose_msg)
        publish(xbox_car_vel_pub, vel_msg)
        publish(xbox_car_acc_pub, acc_msg)
        publish(xbox_car_xythv_pub, xythv_msg)
        publish(xbox_car_marker_pub, xbox_car_marker)
        publish(xbox_car_init_marker_pub, init_marker)
        rossleep(rate)
    end
end

@spawn spin()

end # module

# if !isinteractive()
    using ArgParse
# end

if !isinteractive()
    using XboxCarSim

    function parse_commandline()
        s = ArgParseSettings()
        @add_arg_table s begin
            "--x0"
                arg_type = Float64
                default = 0.0
            "--y0"
                arg_type = Float64
                default = 0.0
            "--th0"
                arg_type = Float64
                default = 0.0
            "--v0"
                arg_type = Float64
                default = 0.0
            "--color"
                default = "lime"
        end
        parse_args(s)
    end
    parsed_args = parse_commandline()

    XboxCarSim.run(XboxCarSim.SE2vState(parsed_args["x0"],
                                        parsed_args["y0"],
                                        parsed_args["th0"],
                                        parsed_args["v0"]),
                   parsed_args["color"])
end


text:
using Cbc, Test, MathOptInterface

const MOI = MathOptInterface
const MOIT = MOI.Test
const MOIU = MOI.Utilities

@MOIU.model(ModelForCachingOptimizer,
    (MOI.ZeroOne, MOI.Integer),
    (MOI.EqualTo, MOI.GreaterThan, MOI.LessThan, MOI.Interval),
    (),
    (),
    (MOI.SingleVariable,),
    (MOI.ScalarAffineFunction,),
    (),
    ()
)

const OPTIMIZER = MOIU.CachingOptimizer(
    MOIU.UniversalFallback(ModelForCachingOptimizer{Float64}()),
    Cbc.Optimizer(logLevel = 0)
)

const CONFIG = MOIT.TestConfig(duals = false, infeas_certificates = false)

@testset "basic_constraint_tests" begin
    MOIT.basic_constraint_tests(OPTIMIZER, CONFIG)
end

@testset "Unit Tests" begin
    MOIT.unittest(OPTIMIZER, CONFIG, [
        "solve_affine_deletion_edge_cases",  # VectorAffineFunction
        "solve_duplicate_terms_vector_affine",  # VectorAffineFunction
        "solve_qcp_edge_cases", "solve_qp_edge_cases"  # No quadratics
    ])
end

@testset "ModelLike tests" begin
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "COIN Branch-and-Cut (Cbc)"
    @testset "default_objective_test" begin
         MOIT.default_objective_test(OPTIMIZER)
     end
     @testset "default_status_test" begin
         MOIT.default_status_test(OPTIMIZER)
     end
    @testset "nametest" begin
        MOIT.nametest(OPTIMIZER)
    end
    @testset "validtest" begin
        MOIT.validtest(OPTIMIZER)
    end
    @testset "emptytest" begin
        # Requires VectorOfVariables
        # MOIT.emptytest(OPTIMIZER)
    end
    @testset "orderedindicestest" begin
        MOIT.orderedindicestest(OPTIMIZER)
    end
    @testset "copytest" begin
        # Requires VectorOfVariables
        # MOIT.copytest(OPTIMIZER, MOIU.CachingOptimizer(
        #     ModelForCachingOptimizer{Float64}(),
        #     Cbc.Optimizer()
        # ))
    end
end

@testset "contlineartest" begin
    MOIT.contlineartest(OPTIMIZER, CONFIG, [
        "linear7", "linear15"  # VectorAffineFunction
    ])
end

@testset "intlineartest" begin
    MOIT.intlineartest(OPTIMIZER, CONFIG, [
        "int2"  # Requires Special-Ordered-Sets
    ])
end

@testset "Test params" begin
    knapsack_model = ModelForCachingOptimizer{Float64}()
    MOIU.loadfromstring!(knapsack_model, """
        variables: x, y
        maxobjective: x + y
        c1: x in ZeroOne()
        c2: y in ZeroOne()
        c3: x + y <= 1.0
    """)
    model = Cbc.Optimizer(maxNodes = 0, presolve = "off", cuts = "off",
                          heur = "off", logLevel = 0)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.NODE_LIMIT
    # We also check that options are not destroyed on `empty!`.
    MOI.empty!(model)
    MOI.copy_to(model, knapsack_model)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.NODE_LIMIT
end

ext: jl


text:
# Kalman filter for the problem

# See Optimal Estimation of Dynamic Systems by Crassidis and Junkins p.256

"""
LightDark2DKalman

A simple Kalman filter for a LightDark2D problem. This does not work very well at localizing.

Fields:
    pomdp::LightDark2D

    obs_noise_offset::Float64
        Makes the filter more conservative by calculating the observation
        at a point this fraction of a standard deviation away from the mean in
        the direction that will increase it.
"""
immutable LightDark2DKalman <: Updater
    pomdp::LightDark2D
    obs_noise_offset::Float64
end
LightDark2DKalman(pomdp::LightDark2D) = LightDark2DKalman(pomdp, 0.0)

function update(kf::LightDark2DKalman, b::SymmetricNormal2, a::Vec2, o::Vec2)
    x = b.mean
    s = b.std

    # propagation
    xpm = x + a
    # so = obs_std(kf.pomdp, x[1]) # VERY sketchy
    # (b.std stays the same in propagation)

    # update
    if xpm[1] > kf.pomdp.min_noise_loc
        so = obs_std(kf.pomdp, xpm[1]+kf.obs_noise_offset*s)
    else
        so = obs_std(kf.pomdp, xpm[1]-kf.obs_noise_offset*s)
    end
    k = s^2/(s^2+so^2)
    xp = xpm + k*(o-x)
    sp = sqrt(1-k)*s
    return SymmetricNormal2(xp, sp)
end


text:
module Movement

using LinearAlgebra: norm, cross
using DifferentialEquations
using FLOWUnsteady

#using ResumableFunctions
#using SimJulia

using Dierckx

export CMassProp, Reference
export AbstractAeroModel, AbstractPropulsionModel, AbstractInertialModel,
    AbstractAtmosphereModel
export UniformGravitationalField,
    ConstantAtmosphere
export CO, COUNTER, COCOUNTER
export equations

# Aliases
uns = FLOWUnsteady
vlm = FLOWUnsteady.vlm

# ------ 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

"""
    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))

"""
function necessary to evaluate the mass properties of the structure: in input
    are given the mass in kg and number of rotors, the radius of center mass in m,
    the distance between the center mass and the rotors. For output is given the
    total mass, the moment of inertia respect to the x,y and z axis
"""
function massprop(m,Mmotor,l,rotors,Rc)
    mtot=m+enumerate(rotors)*Mmotor
    Ixx= (2/5)*m*Rc^2+2*l^2*Mmotor
    Iyy= Ixx
    Izz= (2/5)*m*Rc^2+4*l^2*Mmotor
    return mtot,Ixx,Iyy,Izz
end

"""
    Reference(S, b, c)

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

"""
function used to elaborate the referenc surface, that is equal to the sum of the
    each rotor surface
"""
function reference(Rtip,rotors)
    Sref=Rtip^2*pi*enumerate(rotors)
    return Sref
end

"""
    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

"""
    ConstantAtmosphere(Wi, Wb, rho, asound, g)
    TV defines vectors, TF defines a constant number

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

"""

Computates the Thrust and Torque of each rotor in time
"""

function run (tinit, tfinal)

   for (t in seq(tinit, tfinal, by=tstep))
  	 for rotor in rotors
#	 	 prfrmnc, gammas = calc_distributedloads(self, Vinf, RPM, rho; t=t,
 #                                       include_comps=include_comps,
#                                        return_performance=return_performance,
#                                        Vref=Vref, sound_spd=sound_spd,
#                                        Uinds=Uinds,
#                                        _lookuptable=_lookuptable, _Vinds=_Vinds,
#                                        tiploss_correction=tiploss_correction,
#                                        AR_to_360extrap=AR_to_360extrap)
#     	  	 CT[rotor][t] = prfrmc[1][2]
#     	   	 CQ[rotor][t] = prfrmc[1][3]
#		T[rotor][t] = CT[rotor]*(rho*n^2*D^4)
#		Q[rotor][t] = CQ[rotor]*(rho*n^2*D^5)
		T[rotor][t], Q[rotor][t] = calc_thrust_torque(Rotor)
     	   	end

   end
   return T, Q
end

"""

Computates the spline of the Thrust and Torque of each rotor in time
"""
function spline(t, T, torque)
i=0
   for (t in seq(tinit, tfinal, by=tstep))
   x[i]=t
   i++
   end
 for rotor in rotors
y[rotor] = T[rotor]
z[rotor] = Q[rotor]
splT[rotor] = Spline1D(x, y[rotor])
splQ[rotor] = Spline1D(x, z[rotor])
end
return splT, splQ
end

"""
equations used in order to obtain the second derivatives of the state of the vehicle,
the forces vector contains the force produced by each rotor and the weight of the
whole structure
"""
function equations(atm, MassProp, State, l, rotors, splQ, splT)

    ct, cp, cpsi = cos.([state.theta, state.phi, state.psi])
    st, sp, spsi = sin.([state.theta, state.phi, state.psi])

    theta2=((-splT[1]-splT[2]+splT[3]+splT[4])*l)/
    .Ixx
    phi2=((-splT[1]+splT[2]+splT[3]-splT[4])*l)/massprop.Iyy
    psi2=(splQ[1]-splQ[2]+splQ[3]-splQ[4])/massprop.Izz
    x2=sum(splT[i]*(sp*spsi+cp*cpsi*st) for i in 1:enumerate(rotors))/massprop.m
    y2=sum(splT[i]*(spsi*st*cp-cpsi*sp) for i in 1:enumerate(rotors))/massprop.m
    z2=sum(splT[i]*(cp*ct)-atm.g for i in 1:enumerate(rotors))/massprop.m
    acc=[theta2, phi2, psi2, x2, y2, z2]
    return acc
end

#function integrationv(acc,state)
 #   Pkg.add("SymPy")
#    using SymPy
#    theta1=integrate(acc[1], t)+state.q
#    phi1=integrate(acc[2],t)+state.p
#    psi1=integrate(acc[3],t)+state.r
#    x1=integrate(acc[4],t)+state.u
 #   y1=integrate(acc[5],t)+state.v
 #   z1=integrate(acc[6],t)+state.w
 #   vel=[theta1, phi1, psi1, x1, y1, z1]
#    return vel
#end

#function integrationp(vel,state)
 #   theta=integrate(vel[1],t)+state.theta
 #   phi=integrate(vel[2],t)+state.phi
#    psi=integrate(vel[3],t)+state.psi
#    x=integrate(vel[4],t)+state.x
#    y=integrate(vel[5],t)+state.y
#    z= integrate(vel[6],t)+state.z
#    pos=[theta,phi,psi,x,y,z]
#    return pos
#end

#function simulate
#	sim = Simulation()
#	for t in 1:100
#  	run(sim, t)
#  	forces[i][t], torque[i][t] = calcforce (tspan, rotors)
#	end
#end

function integrate(massprop,State, splQ, splT, l,n_motor, atm, tspan, s0, acc, Reference)
p = massprop,state, splQ, splT,l, enumerate(rotors), atm, Reference
prob = DifferentialEquations.ODEProblem(equations, s0, tspan, p)
sol = DifferentialEquations.solve(prob)

return sol

end

function integrate2(massprop,state,splQ, splT, l,n_motor, atm, tspan, s0, acc, Reference)
p = massprop,state, splQ, splT,l, enumerate(rotors), atm, Reference
prob = DifferentialEquations.SecondOrderODEProblem(equations, u0, s0, tspan, p)
sol = DifferentialEquations.solve(prob)
nothing # hide

return sol
end


text:
using NumericalMethodsInQuantumMechanics
using Documenter

makedocs(;
    modules=[NumericalMethodsInQuantumMechanics],
    authors="Qi Zhang <singularitti@outlook.com>",
    repo="https://github.com/singularitti/NumericalMethodsInQuantumMechanics.jl/blob/{commit}{path}#L{line}",
    sitename="NumericalMethodsInQuantumMechanics.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://singularitti.github.io/NumericalMethodsInQuantumMechanics.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "Manual" => Any[
            "EigenvalueProblems/NumerovMethod.md",
        ],
    ],
)

deploydocs(;
    repo="github.com/singularitti/NumericalMethodsInQuantumMechanics.jl",
)


text:
## Process measures

function integrate_measure(qs, integrand, domain, δ::DiracWeight, properties...)
    x = point(δ)
    if x ∈ domain
        I = integrand(x)
        E = zero(I)
    else
        I, E = zero_result(integrand)
    end
    I, E
end

function integrate_measure(qs, integrand, domain, μ::DiscreteWeight, properties...)
    if domain == covering(μ)
        I = sum(w*integrand(x) for (w,x) in zip(weights(μ), points(μ)))
        I, zero_error(I)
    else
        # Only sum over the elements that lie in the domain
        I, E = zero_result(integrand)
        x = points(μ)
        w = weights(μ)
        for i in 1:length(x)
            if x[i] ∈ domain
                I += w[i]*integrand(x[i])
            end
        end
        I, zero_error(I)
    end
end

# Lebesgue measures can pass through unaltered
integrate_measure(qs, integrand, domain, measure::LebesgueMeasure, properties...) =
    integrate_domain(qs, integrand, domain, measure, properties...)

process_measure(qs, integrand, domain, measure::LebesgueMeasure, properties...) =
    (integrand, domain, measure, properties...)

process_measure(qs, integrand, domain, measure::Measure, properties...) =
    process_measure_default(qs, integrand, domain, measure, properties...)

# By default we replace all measures by the Lebesgue on the space
process_measure_default(qs, integrand, domain, measure::Measure{T}, properties...) where {T} =
    ((t->unsafe_weightfun(measure, t)) * integrand, domain, Lebesgue{T}(), properties...)

function process_measure(qs, integrand, domain::MappedDomain, μ::MappedWeight, properties...)
    map1 = forward_map(μ)
    map2 = forward_map(domain)
    if map1 == map2
        # f -> f(m1(x))
        process_measure(qs, integrand ∘ map1, superdomain(domain), supermeasure(μ), properties...)
    else
        process_measure_default(qs, integrand, domain, μ, properties...)
    end
end

function process_measure(qs, integrand, domain, μ::MappedWeight, properties...)
    # f -> f(m(x))
    map1 = forward_map(μ)
    process_measure(qs, integrand ∘ map1, mapped_domain(map1, domain), supermeasure(μ), properties...)
end

# Truncate an infinite domain to a finite one for numerical evaluation of Hermite integrals
function process_measure(qs::AdaptiveStrategy, integrand, domain::FullSpace{T}, measure::HermiteWeight{T}, properties...) where {T}
    U = sqrt(-log(eps(T)))
    (hermite_weightfun * integrand, -U..U, Lebesgue{T}(), properties...)
end

# apply the cosine map for integrals with the ChebyshevT weight, to avoid the singularities
function process_measure(qs::AdaptiveStrategy, integrand, domain::ChebyshevInterval, measure::ChebyshevTWeight{T}, properties...) where {T}
    # Transformation is: f(t) -> pi*f(cos(pi*t))
    Tpi = convert(T, pi)
    map = t -> cos(Tpi*t)
    (Tpi * (integrand ∘ map), UnitInterval{T}(), LebesgueUnit{T}(), properties...)
end

# same as above, but on a subinterval
function process_measure(qs::AdaptiveStrategy, integrand, domain::AbstractInterval, measure::ChebyshevTWeight{T}, properties...) where {T}
    Tpi = convert(T, pi)
    map = t -> cos(Tpi*t)
    a, b = extrema(domain)
    a < -1.001 && throw(BoundsError(measure, a))
    b > 1.001 && throw(BoundsError(measure, b))
    # Set a and b to be within [-1,1] in order to avoid errors with acos below
    a = max(a, -1)
    b = min(b, 1)
    (Tpi * (integrand ∘ map), acos(b)/pi..acos(a)/pi, Lebesgue{T}(), properties...)
end

# apply the cosine map for integrals with the ChebyshevU weight as well
function process_measure(qs::AdaptiveStrategy, integrand, domain::ChebyshevInterval, measure::ChebyshevUWeight{T}, properties...) where {T}
    Tpi = convert(T, pi)
    prefactor = t -> Tpi * sin(Tpi*t)^2
    map = t -> cos(Tpi*t)
    (prefactor * (integrand ∘ map), UnitInterval{T}(), LebesgueUnit{T}(), properties...)
end

function process_measure(qs::AdaptiveStrategy, integrand, domain::AbstractInterval, measure::ChebyshevUWeight{T}, properties...) where {T}
    Tpi = convert(T, pi)
    prefactor = t -> Tpi * sin(Tpi*t)^2
    map = t -> cos(Tpi*t)
    a, b = extrema(domain)
    a < -1.001 && throw(BoundsError(measure, a))
    b > 1.001 && throw(BoundsError(measure, b))
    # Set a and b to be within [-1,1] in order to avoid errors with acos below
    a = max(a, -1)
    b = min(b, 1)
    (prefactor * (integrand ∘ map), acos(b)/pi..acos(a)/pi, Lebesgue{T}(), properties...)
end

function measure_fetch_transformations(qs, domain, weight, properties...)
    T = promote_type(numtype(domain),numtype(weight))
    I = TransformationLogIntegrand{T}()
    I_trans, domain_trans, weight_trans = process_measure(qs, I, domain, weight)
    I_trans.prefactor, I_trans.fun, domain_trans, weight_trans
end

# For product domains, process the measures dimension per dimension
function process_measure(qs::AdaptiveStrategy, integrand, domain::ProductDomain, μ::ProductWeight, properties...)
    if ncomponents(domain) == ncomponents(μ)
        if ncomponents(domain) == 2
            pre1, map1, domain1, μ1 = measure_fetch_transformations(qs, component(domain, 1), component(μ, 1), properties...)
            pre2, map2, domain2, μ2 = measure_fetch_transformations(qs, component(domain, 2), component(μ, 2), properties...)
            prefactor = t -> pre1(t[1])*pre2(t[2])
            map = t -> SA[map1(t[1]), map2(t[2])]
            prefactor * (integrand ∘ map), productdomain(domain1, domain2), productmeasure(μ1, μ2)
        elseif ncomponents(domain) == 3
            pre1, map1, domain1, μ1 = measure_fetch_transformations(qs, component(domain, 1), component(μ, 1), properties...)
            pre2, map2, domain2, μ2 = measure_fetch_transformations(qs, component(domain, 2), component(μ, 2), properties...)
            pre3, map3, domain3, μ3 = measure_fetch_transformations(qs, component(domain, 3), component(μ, 3), properties...)
            prefactor = t -> pre1(t[1])*pre2(t[2])*pre3(t[3])
            map = t -> SA[map1(t[1]), map2(t[2]), map3(t[3])]
            prefactor * (integrand ∘ map), productdomain(domain1, domain2, domain3), productmeasure(μ1, μ2, μ3)
        else
            process_measure_default(qs, integrand, domain, μ, properties...)
        end
    else
        process_measure_default(qs, integrand, domain, μ, properties...)
    end
end

# The "best" rule for certain measures becomes a Gauss rule
integrate_measure(qs::BestRule, integrand, domain::ChebyshevInterval{T}, μ::LegendreWeight, properties...) where {T} =
    integrate_measure(Q_GaussLegendre(T, qs.n), integrand, domain, μ, properties...)

integrate_measure(qs::BestRule, integrand, domain::ChebyshevInterval, μ::JacobiWeight{T}, properties...) where {T} =
    integrate_measure(Q_GaussJacobi(qs.n, μ.α, μ.β), integrand, domain, LegendreWeight{T}(), properties...)

integrate_measure(qs::BestRule, integrand, domain::HalfLine, μ::LaguerreWeight{T}, properties...) where {T} =
    integrate_measure(Q_GaussLaguerre(qs.n, μ.α), integrand, domain, Lebesgue{T}(), properties...)

integrate_measure(qs::BestRule, integrand, domain::FullSpace{T}, μ::HermiteWeight{T}, properties...) where {T} =
    integrate_measure(Q_GaussHermite(qs.n), integrand, domain, Lebesgue{T}(), properties...)


text:
@testset "Generic" begin
    N = 5

    # Generate random test quantities for specific types.
    ∇Arrays = Union{Type{∇Array}, Type{∇ArrayOrScalar}}

    trandn(rng::AbstractRNG, ::∇Arrays) = randn(rng, N, N)
    trandn(rng::AbstractRNG, ::Type{∇Scalar}) = randn(rng)
    trandn(rng::AbstractRNG, ::Type{<:Transpose}) = Transpose(randn(rng, N, N))
    trandn(rng::AbstractRNG, ::Type{<:Adjoint}) = Adjoint(randn(rng, N, N))

    trandn2(rng::AbstractRNG, ::∇Arrays) = randn(rng, N^2, N^2)
    trandn2(rng::AbstractRNG, ::Type{<:Transpose}) = Transpose(randn(rng, N^2, N^2))
    trandn2(rng::AbstractRNG, ::Type{<:Adjoint}) = Adjoint(randn(rng, N^2, N^2))

    trand(rng::AbstractRNG, ::∇Arrays) = rand(rng, N, N)
    trand(rng::AbstractRNG, ::Type{∇Scalar}) = rand(rng)
    trand(rng::AbstractRNG, ::Type{<:Transpose}) = Transpose(rand(rng, N, N))
    trand(rng::AbstractRNG, ::Type{<:Adjoint}) = Adjoint(rand(rng, N, N))

    @testset "Unary sensitivities" begin
        _ϵ, lb, ub = 3e-2, -3.0, 3.0
        unary_linalg_optimisations = [
            (-,          ∇Array,  (lb, ub)),
            (tr,         ∇Array,  (lb, ub)),
            (inv,        ∇Array,  (lb, ub)),
            (det,        ∇Array,  (_ϵ, ub)),
            (logdet,     ∇Array,  (_ϵ, ub)),
            (transpose,  ∇Array,  (lb, ub)),
            (adjoint,    ∇Scalar, (_ϵ, ub)),
            (adjoint,    ∇Array,  (lb, ub)),
            (norm,       ∇Array,  (lb, ub)),
            (norm,       ∇Scalar, (lb, ub)),
        ]

        rng = MersenneTwister(123)
        @testset "$f" for (f, T_In, bounds) in unary_linalg_optimisations
            for _ in 1:5
                Z = trand(rng, T_In) .* (bounds[2] .- bounds[1]) .+ bounds[1]
                X = Z'Z + 1e-6 * one(Z)
                Ȳ, V = f(X), trandn(rng, T_In)
                @test check_errs(f, Ȳ, X, 1e-1 .* V)
            end
        end
    end

    @testset "Binary sensitivities" begin
        rng = MersenneTwister(2)
        binary_linalg_optimisations = [
            (*, ∇Array, ∇Array,),
            (/, ∇Array, ∇Array,),
            (\, ∇Array, ∇Array,),
            (norm, ∇Array, ∇Scalar,),
            (norm, ∇Scalar, ∇Scalar,),
        ]
        @testset "$f" for (f, T_A, T_B) in binary_linalg_optimisations
            for _ in 1:5
                A, B, VA, VB = trandn.(Ref(rng), (T_A, T_B, T_A, T_B))
                @test check_errs(f, eval(f)(A, B), (A, B), (VA, VB))
            end
        end
    end

    @testset "kron" begin
        rng = MersenneTwister(3)
        for _ in 1:5
            A, B, VA, VB = trandn.(Ref(rng), (∇Array, ∇Array, ∇Array, ∇Array))
            @test check_errs(kron, kron(A, B), (A, B), (VA, VB))
        end
    end

    @testset "I" begin
        rng = MersenneTwister(4)
        for _ in 1:5
            A, VA, tI = randn(rng, 5, 5), randn(rng, 5, 5), 0.65 * I
            @test check_errs(X->X + tI, VA, A, 1e-1 * randn(rng, 5, 5))
            @test check_errs(X->tI + X, VA, A, 1e-1 * randn(rng, 5, 5))
        end
    end

    @testset "dot" begin
        rng = MersenneTwister(123456)
        for _ in 1:10
            x, y, vx, vy = randn.(Ref(rng), [5, 5, 5, 5])
            @test check_errs(LinearAlgebra.dot, LinearAlgebra.dot(x, y), (x, y), (vx, vy))
        end
    end

    @testset "copy" begin
        rng = MersenneTwister(12345)

        # Scalars (no-op)
        x = randn(rng)
        y = randn(rng)
        @test check_errs(copy, x, x, y)
        x_ = Leaf(Tape(), x)
        c = copy(x_)
        @test c isa Branch{Float64}
        @test getfield(c, :f) === Base.copy

        # Unwrapping adjoint/transposes
        X = randn(rng, 6, 6)'
        Y = randn(rng, 6, 6)
        @test check_errs(copy, X, copy(X), Y)
        X_ = Leaf(Tape(), X)
        C = copy(X_)
        @test C isa Branch{Matrix{Float64}}
        @test getfield(c, :f) === Base.copy
    end

    @testset "exp" begin
        rng = MersenneTwister(12345)
        n = 10
        @testset "Symmetric" begin
            symm!(X) = (X .= (X .+ X') ./ 2; X)
            X = symm!(randn(rng, n, n))
            VX = symm!(randn(rng, n, n))
            @test check_errs(exp, randn(rng, n, n), X, VX)
        end
        @testset "Asymmetric" begin
            A = randn(rng, n, n)
            VA = randn(rng, n, n)
            @test check_errs(exp, randn(rng, n, n), A, VA)
        end
    end
end


text:
using Base.Iterators
using LinearAlgebra

include("Helper.jl")

function binPacking(params, nodes, demands, costs, t)

	Q = params["Q"]
	cumsum = 0

	circuits = Vector{Int64}[]

	#circuit = Vector{Int64}()
	circuit = [0]

	for (i, demand) in enumerate(demands[:, t]) #Itérer sur les demandes au temps t

		cumsum += demand

		if cumsum > Q
			push!(circuits, circuit)

			#Reset
			cumsum = demand
			circuit = [0]
		end

		push!(circuit, i)

	end

	if length(circuit) > 1
		push!(circuits, circuit)
	end

	return circuits

end

#=
Retourne une liste de tournées qui vérifie la contrainte de poids sur les véhicules
Le nombre de tournée peut toutefois excéder le nombre de véhicules disponible.

Note: Version plus rapide mais qui prend un peu plus de mémoire
=#
function clark_wright(params,nodes,demands,costs,t)

	n = params["n"]
	Q = params["Q"]

	#Initialiser S
	S = [[0, i] for i in 1:n]
	whereIs = Dict(i => circuit for (i, circuit) in enumerate(S))

	#Calcul des s_{i,j} et mettre dans l'ordre décroissant
	s=[]
	for i in 1:n, j in 1:n
		if i != j
			push!(s, [ (i,j), costs[(0, i)] + costs[(0, j)] - costs[(i, j)] ])
		end
	end

	sort!(s, by = x->x[2], rev=true) #by = key pour sort, rev = reverse (ordre décroissant)

	#Construction des circuits (attention il peut en avoir plus de m (le nombre de camion disponible))
	for k in 1:length(s)

		i=s[k][1][1]
		j=s[k][1][2]

		circuit_i = whereIs[i]
		circuit_j = whereIs[j]

		#Si i et j ne sont pas dans le même circuit
		if circuit_j != circuit_i

			index_circuit_i = findfirst(item -> item == circuit_i, S)
			index_circuit_j = findfirst(item -> item == circuit_j, S)

			circuit_union = union(S[index_circuit_i], S[index_circuit_j])

			#Calcul demande du nouveau circuit circuit_union
			demandeUnion=0
			for l in circuit_union[2:end] #commence à 2 car le premier noeud du circuit est 0 (le dépot n'a pas de demande)
				demandeUnion += demands[l, t]
			end

			#Si la demande du nouveau circuit n'est pas trop pour un seul camion
			if demandeUnion <= Q

				deleteat!(S, sort!([index_circuit_i, index_circuit_j]))

				push!(S, circuit_union)

				for index in circuit_union
					push!(whereIs, index=>circuit_union)
				end

			end
		end

	end

	return S #la liste des circuits qui sont des tournées valides (qui respectent la contrainte de poids)
end

#=
Retourne une liste de tournées qui vérifie la contrainte de poids sur les véhicules
Le nombre de tournée peut toutefois excéder le nombre de véhicules disponible.

Note: Version plus lente mais qui prend un peu moins de mémoire
=#
function clark_wright_old(params,nodes,demands,costs,t)

	n = params["n"]
	Q = params["Q"]

	#Initialiser S
	S = [[0, i] for i in 1:n]

	#Calcul des s_{i,j} et mettre dans l'ordre décroissant
	s=[]
	for (i,j) in product(1:n,1:n)
		if i != j
			push!(s,[(i,j),costs[(0,i)]+costs[(0,j)]-costs[(i,j)]])
		end
	end
	sort!(s, by = x->x[2], rev=true) #by = key pour sort, rev = reverse (ordre décroissant)

	#Construction des circuits (attention il peut en avoir plus de m (le nombre de camion disponible))
	for k in 1:length(s)

		i=s[k][1][1]
		j=s[k][1][2]
		foundi=false
		foundj=false
		circuit_i=NaN
		circuit_j=NaN
		index_i=0
		index_j=0

		#trouver les circuits où sont i et j
		for (index, circuit) in enumerate(S)
			if i in circuit
				circuit_i=circuit
				foundi=true
				index_i = index
			end
			if j in circuit
				circuit_j=circuit
				foundj=true
				index_j = index
			end
			if foundi && foundj
				break
			end
		end

		#Si i et j ne sont pas dans le même circuit
		if circuit_j != circuit_i

			#création de l'union de circuit_i et circuit_j
			circuit_union=union(circuit_i, circuit_j)

			#Calcul demande du nouveau circuit circuit_union
			demandeUnion=0
			for l in 2:length(circuit_union) #commence à 2 car le premier noeud du circuit est 0 (le dépot n'a pas de demande)
				demandeUnion+=demands[circuit_union[l],t]
			end

			#Si la demande du nouveau circuit n'est pas trop pour un seul camion
			if demandeUnion<=Q
				#suppression des circuits circuit_i et circuit_j
				deleteat!(S, sort!([index_i, index_j]))

				push!(S,circuit_union)
			end
		end
	end

	return S #la liste des circuits qui sont des tournées valides (qui respectent la contrainte de poids)
end

function sectorielle(params, nodes, demands, costs, t, angle)

	parts = div(360, angle)
	angle = 360 / parts

	center = [nodes[0]["y"], nodes[0]["x"]]
	top = [1, 0]

	sectors = [Int[] for i in 1:parts]

	for (i, node) in enumerate(nodes[1:end])
		centredPoint = [node["y"], node["x"]] - center

		#Calcul du degré entre le vecteur directionnel [1, 0] et le vecteur centre vers le point
		degree = mod(rad2deg(atan(top...) - atan(centredPoint...)), 360)

		index = Int( div(degree, angle) ) % parts + 1

		push!(sectors[index], i)
	end

	Q = params["Q"]

	circuits = Vector{Int64}[]
	circuit = [0]
	cumsum = 0

	for sector in sectors
		for i in sector

			demand = demands[i, t]
			cumsum += demand

			if cumsum > Q
				push!(circuits, circuit)

				#Reset
				cumsum = demand
				circuit = [0]
			end

			push!(circuit, i)

		end
	end

	if length(circuit) > 1
		push!(circuits, circuit)
	end

	return circuits

end

#=
Reorder the circuit by putting the value (currtently at pos) at the first index
=#
function replace(circuit, pos)

	return vcat(circuit[pos:end], circuit[1:pos-1])

end

function TSPVoisins(circuit)

	n = length(circuit)

	voisins = Vector{Int64}[]

	for i in 1:n
		for j in i+1:n
			voisin = copy(circuit)
			voisin[i], voisin[j] = voisin[j], voisin[i]

			if i == 1
				voisin = replace(voisin, j)
			end
			push!(voisins, voisin)
		end
	end

	return voisins

end

function TSPBoost(circuits, params, costs)

	for (i, circuit) in enumerate(circuits)

		best = circuit
		bestCost = getCircuitCost(circuit, costs)

		while true

			voisins = TSPVoisins(best)
			noChange = true

			for voisin in voisins
				cost = getCircuitCost(voisin, costs)

				if cost < bestCost
					bestCost = cost
					best = voisin
					noChange = false
				end

			end

			if noChange
				break
			end

		end

		circuits[i] = best

	end

	return circuits

end

function SwapCircuitVoisins(circuits)

	voisins = Vector{Vector{Int64}}[]

	for (i, circuit) in enumerate(circuits)
		for (j, elem) in enumerate(circuit[2:end])

			j += 1

			for indexCircuit in 1:length(circuits)
				if i != indexCircuit
					voisin = deepcopy(circuits)

					deleteat!(voisin[i], j)
					push!(voisin[indexCircuit], elem)

					#If only zero remain in i
					if length(voisin[i]) == 1
						deleteat!(voisin, i)
					end

					push!(voisins, voisin)
				end
			end

		end
	end

	return voisins

end

#=

Métaheuristique faisant:
	- Ajout d'un circuit vide si possible

	Voisin:
		- Déplace chaque noeud de chaque circuit à la fin de chaque autre circuit
		- Suppression du circuit vide si il l'est encore
		- Appliquer la recherche local de TSP sur chaque circuit de chaque voisin

=#
function mixMetaheuristic(circuits, params, costs)

	###Initiliaze

	#Deep copy pour ne pas modifier circuits de base
	currentCircuits = deepcopy(circuits)
	currentCircuits = TSPBoost(currentCircuits, params, costs)

	#Noeud courant est le meilleur actuellement
	best = deepcopy(currentCircuits)
	bestCost = getCircuitsCost(best, costs)

	while true
		#Debug
		#println("Loop ", bestCost)

		###Creation des voisins

		#Ajouter un circuit vide si le nombre n'est pas déjà dépassé
		if length(best) < params["k"]
			push!(best, [0])
		end

		voisins = SwapCircuitVoisins(best)

		#Retirer les 0 les circuits contenant seulement 0 et de best (-1 car on rajouté 1 avant)
		if length(best) - 1 < params["k"]

			pop!(best)

			for voisin in voisins
				if length(voisin[end]) == 1
					pop!(voisin)
				end
			end

		end

		#Appliquer la recherche locale de TSP
		for voisin in voisins

			voisin = TSPBoost(voisin, params, costs)
		end

		###Fin de création des voisins

		noChange = true

		#Trouver le meilleur voisin
		for voisin in voisins
			cost = getCircuitsCost(voisin, costs)

			if cost < bestCost

				bestCost = cost
				best = voisin
				noChange = false

			end

		end

		if noChange
			break
		end

	end

	return best

end

#=
function sectorielle_old(params,nodes,demands,costs,t,angle) #ON SUPPOSE QUE 360 EST DIVISIBLE PAR ANGLE (sinon trop compliqué, flemme)
	distance_du_point_le_plus_eloigne=0
	origin=nodes[0]
	#Calcul de la distance du point le plus eloigné
	for node in nodes
		distance=sqrt((node["x"]-origin["x"])^2+(node["y"]-origin["y"])^2)
		if(distance_du_point_le_plus_eloigne<distance)
			distance_du_point_le_plus_eloigne=distance
		end
	end

	#Calcul distance point fictif a l'origine
		#La distance de ces points fictif à l'origine ne peut pas être égale
		#a la distance du point le plus eloigné car celui ci peut ne pas être aucun des triangles
		#C'est pourquoi que la distance des points fictifs à l'origine sera égal à
		#la somme de la distance du point le plus eloigné + la distance entre le (point D qui est le point au milieu
		# du segment entre deux points fictifs quelconquz) et le
		#(point E de fin du segment qui commence à l'origine, passant par D )

	distance_point_fictif_a_lorigine=distance_du_point_le_plus_eloigne*(2-cos(angle/2))

	#Création des points fictifs pour créer les secteurs en triangles
	angleNormalise=angle/360
	triangles=[]
	for i in 1:360/angle
		push!(triangles,[(float(origin["x"]),float(origin["y"]))])
	end
	k=1
	for i in 1:2*360/angle
		point=(origin["x"]+(1-(i-1)*angleNormalise)*distance_point_fictif_a_lorigine,origin["y"]+(i-1)*angleNormalise*distance_point_fictif_a_lorigine)
		push!(triangles[k],point)
		k+=Integer((i+1)%2)
	end

	#Création secteurs (il y en a 360/angle)
	secteurs=[]
	for i in 1:length(triangles)
		push!(secteurs,[origin])
	end
	initial=true
	for node in nodes
		if initial
			initial=false
			continue
		end
		k=1
		for (Ori,A,B) in triangles #detection de quel secteur est situé le point
			# AB=vecteurEntreDeuxPoints(A,B)
			# BA=vecteurEntreDeuxPoints(B,A)

			# OriA=vecteurEntreDeuxPoints(Ori,A)
			# AOri=vecteurEntreDeuxPoints(A,Ori)

			# OriB=vecteurEntreDeuxPoints(Ori,B)
			# BOri=vecteurEntreDeuxPoints(B,Ori)

			# ANode=vecteurEntreDeuxPoints(A,(node["x"],node["y"]))
			# BNode=vecteurEntreDeuxPoints(B,(node["x"],node["y"]))
			# OriNode=vecteurEntreDeuxPoints(C,(node["x"],node["y"]))

			# #print(AB,BA,OriA,AOri,OriB,BOri,ANode,BNode,OriNode)
			# if dot(cross(AB,ANode),cross(ANode,AOri))>=0 && dot(cross(BA,BNode),cross(BNode,BOri))>=0 && dot(cross(OriA,OriNode),cross(OriNode,OriB))>=0
			# 	push!(secteurs[k],node)
			# 	break
			# end

			cond1=(A[1]-node["x"])*(B[2]-node["y"])-(A[2]-node["y"])*(B[1]-node["x"])
			cond2=(B[1]-node["x"])*(Ori[2]-node["y"])-(B[2]-node["y"])*(Ori[1]-node["x"])
			cond3=(Ori[1]-node["x"])*(A[2]-node["y"])-(Ori[2]-node["y"])*(A[1]-node["x"])
			if((cond1>=0 && cond2>=0 && cond3>=0) || (cond1<0 && cond2<0 && cond3<0))
				push!(secteurs[k],node)
				break
			end
			k+=1
		end
	end

	ens_de_ens_de_circuits=[]
	for secteur in secteurs
		push!(ens_de_ens_de_circuits,clark_wright(params,secteur,demands,costs,t))
	end

	ens_final=[]
	for ens_de_circuits in ens_de_ens_de_circuits
		for circuit in ens_de_circuits
			push!(ens_final,circuit)
		end
	end
	return ens_final
end

=#


text:
using Test #src

# # Exercise 1
#
# Set `x` to be a 20-item random sample with replacement from the integers 1 to 15.

x = "Change this String to be the random sample"

# @assert length(x) == 20             # src
# @assert all(x -> 1 ≤ x ≤ 15, x)     # src

# # Exercise 2
#
# What is the probability of

max_stars_repo_name: fercarozzi/JuliaAcademyMaterials

max_issues_repo_name: fercarozzi/JuliaAcademyMaterials

max_forks_repo_name: fercarozzi/JuliaAcademyMaterials


text:
using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra, StableRNGs
import QuantumOpticsBase

@testset "initial state" begin

    Random.seed!(StableRNG(0), 1)

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(op1::AbstractSuperOperator, op2::AbstractSuperOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))

    mode1 = Mode(0.2, 1.0)
    mode2 = Mode(0.3, 2.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 2, [2.0, 200.0])
    mol2 = Molecule([mode2], 2, [3.0, 300.0])
    aggCore = AggregateCore([mol1, mol2])
    aggCore.coupling[2, 3] = 50
    aggCore.coupling[3, 2] = 50
    agg = setupAggregate(aggCore)
    aggTools = agg.tools
    aggOperators = agg.operators

    Ham_I = agg.operators.Ham_I
    Ham_0 = agg.operators.Ham_0
    Ham = agg.operators.Ham

    basis = agg.tools.basis
    indicesLen = agg.tools.bSize
    indices = agg.tools.indices
    indicesMap = agg.tools.indicesMap
    FCFact = agg.tools.FCfactors
    FCProd = agg.tools.FCproduct

    T = 300.0
    mu_array = [[1, 1, 2]]
    W0_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.25000143755694526 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.2499997124870238 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.2500002875090083 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.24999856244702248]
    W0 = thermal_state(T, mu_array, aggCore, aggTools, aggOperators; diagonalize = true)
    @test 1e-15 > D(W0_ref, W0.data)

    W0 = thermal_state(T, mu_array, agg; diagonalize = true)
    @test 1e-15 > D(W0_ref, W0.data)

    W01 = thermal_state(T, [[1, 2, 1]], aggCore, aggTools, aggOperators; diagonalize = true)
    W02 = thermal_state(T, [[1, 1, 2]], aggCore, aggTools, aggOperators; diagonalize = true)
    W0_composite_ref = 0.8 * W01 + 0.2 * W02
    normalize!(W0_composite_ref)
    W0 = thermal_state_composite(T, [0.0, 0.8, 0.2], aggCore, aggTools, aggOperators; diagonalize = true)
    @test 1e-15 > D(W0_composite_ref, W0)

    W0 = thermal_state_composite(T, [0.0, 0.8, 0.2], agg; diagonalize = true)
    @test 1e-15 > D(W0_composite_ref, W0)

    W0_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.25029983141333245 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.24993996472850133 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.25005986276494624 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.2497003410932199]
    W0 = thermal_state(T, mu_array, aggCore, aggTools, aggOperators; diagonalize = false, diagonal = true)
    @test 1e-15 > D(W0_ref, W0.data)

    W01 = thermal_state(T, [[1, 2, 1]], aggCore, aggTools, aggOperators; diagonalize = false, diagonal = true)
    W02 = thermal_state(T, [[1, 1, 2]], aggCore, aggTools, aggOperators; diagonalize = false, diagonal = true)
    W0_composite_ref = 0.8 * W01 + 0.2 * W02
    normalize!(W0_composite_ref)
    W0 = thermal_state_composite(
        T,
        [0.0, 0.8, 0.2],
        aggCore,
        aggTools,
        aggOperators;
        diagonalize = false,
        diagonal = true,
    )
    @test 1e-15 > D(W0_composite_ref, W0)

    W01 = thermal_state(T, [[1, 2, 1]], aggCore, aggTools, aggOperators; diagonalize = true)
    W02 = thermal_state(T, [[1, 1, 2]], aggCore, aggTools, aggOperators; diagonalize = true)
    W0_composite_ref = 0.8 * W01 + 0.2 * W02
    # normalize!(W0_composite_ref)
    W0 = thermal_state_composite(T, [0.0, 0.8, 0.2], aggCore, aggTools, aggOperators; diagonalize = true)
    @test 1e-15 > D(W0_composite_ref, W0)

    W0_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.25029983141333245 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.24993996472850133 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.25005986276494624 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.2497003410932199]
    W0 = thermal_state(T, mu_array, aggCore, aggTools, aggOperators; diagonalize = false, diagonal = true)
    @test 1e-15 > D(W0_ref, W0.data)

    W01 = thermal_state(T, [[1, 2, 1]], aggCore, aggTools, aggOperators; diagonalize = false, diagonal = true)
    W02 = thermal_state(T, [[1, 1, 2]], aggCore, aggTools, aggOperators; diagonalize = false, diagonal = true)
    W0_composite_ref = 0.8 * W01 + 0.2 * W02
    # normalize!(W0_composite_ref)
    W0 = thermal_state_composite(
        T,
        [0.0, 0.8, 0.2],
        aggCore,
        aggTools,
        aggOperators;
        diagonalize = false,
        diagonal = true,
    )
    @test 1e-15 > D(W0_composite_ref, W0)

    T = 0.5
    mu_array = [[1, 1, 2]]
    W0_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.7750674418239181 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.06464140010302448 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.1479518276234691 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.012339330449588293]
    W0 = thermal_state(T, mu_array, aggCore, aggTools, aggOperators; diagonalize = true)
    @test 1e-15 > D(W0_ref, W0.data)

    W0, rho0, W0_bath = ultrafast_laser_excitation(T, [0.0, 0.8, 0.2], agg; diagonalize = true)
    W0_ref = thermal_state_composite(
        T,
        [0.0, 0.8, 0.2],
        aggCore,
        aggTools,
        aggOperators;
        diagonalize = true
    )
    normalize!(W0_ref)
    W0_ref = DenseOperator(W0_ref.basis_l, W0_ref.basis_r, complex(W0_ref.data))

    rho0_ref = trace_bath(W0_ref, agg.core, agg.tools)
    rho0_ref = DenseOperator(rho0_ref.basis_l, rho0_ref.basis_r, complex(rho0_ref.data))

    W0_bath_ref = get_rho_bath(W0_ref, agg.core, agg.tools)
    W0_bath_ref = DenseOperator(W0_bath_ref.basis_l, W0_bath_ref.basis_r, complex(W0_bath_ref.data))
    @test 1e-15 > D(W0_ref, W0)
    @test 1e-15 > D(rho0_ref, rho0)
    @test 1e-15 > D(W0_bath_ref, W0_bath)

end # testset

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
using ArgParse
s = ArgParseSettings()
@add_arg_table s begin
    "--part2"
    help = "Do part 2"
    action = :store_true
    "--test"
    help = "Run test data"
    action = :store_true
end
parsed_args = parse_args(ARGS, s)

test_suffix = parsed_args["test"] ? "_test" : ""
input = readlines("input/day_15$test_suffix.txt")
input_x = length(input[1])
input_y = length(input)

function risk(cave)
    (sx, sy) = size(cave)
    risks = fill(typemax(UInt16), sx, sy)
    risks[1,1] = 0
    function update_risk(x, y)
        codb = cave[x, y]
        risk = risks[x, y]
        path_risk = risk - codb
        if x > 1 && risks[x - 1, y] < path_risk
            path_risk = risks[x - 1, y]
        end
        if y > 1 && risks[x, y - 1] < path_risk
            path_risk = risks[x, y - 1 ]
        end
        if x < sx && risks[x + 1, y] < path_risk
            path_risk = risks[x+1, y]
        end
        if y < sy && risks[x, y+1] < path_risk
            path_risk = risks[x, y + 1]
        end
        path_risk += codb
        if path_risk < risk
            risks[x, y] = path_risk
            return true
        end
        false
    end

    while true
        updated = false
        for x in 1: sx
            for y in 1: sy
                if update_risk(x, y)
                    updated = true
                end
            end
        end
        if !updated
            break;
        end
    end
    risks[sx, sy]
end

function populate_input!(cave)
    for (y, l) in enumerate(input)
        for (x, v) in enumerate([parse(UInt8, c) for c in l])
            cave[x, y] = v
        end
    end
end

function part1()
    cave = Array{UInt8}(undef, input_x, input_y)
    populate_input!(cave)
    risk(cave)
end

function part2()
    cave = Array{UInt8}(undef, input_x * 5, input_y * 5)
    populate_input!(cave)
    for tile_x in 0:4
        offset_x = input_x * tile_x
        for tile_y in 0:4
            offset_y = input_y * tile_y
            if tile_x ==0 && tile_y == 0
                continue
            end
            for x in 1: input_x
                for y in 1: input_y
                    c = cave[x,y] + tile_x + tile_y
                    while c > 9
                        c -= 9
                    end
                    cave[offset_x + x, offset_y + y] = c
                end
            end
        end
    end
    risk(cave)
end

p = parsed_args["part2"] ? part2 : part1
answer = p()
println("The answer is: $answer")


text:
function intensity_by_state(mmpp::MMPP, m = nothing)
    λ = [intensity(emission(mmpp, s), m) for s = 1:nb_states(mmpp)]
    Λ = Diagonal(λ)
    return Λ
end

function renewal_density(mmpp::MMPP, y, m = nothing)
    Q = rate_matrix(mmpp)
    Λ = intensity_by_state(mmpp)
    Λm = intensity_by_state(mmpp, m)
    return exp((Q - Λ) * y) * Λm
end

function renewal_void_probability(mmpp::MMPP, y)
    Q = rate_matrix(mmpp)
    Λ = intensity_by_state(mmpp)
    return exp((Q - Λ) * y)
end

function stationary_distribution(mmpp::MMPP)
    Q = rate_matrix(mmpp)
    Λ = intensity_by_state(mmpp)
    P = inv(Λ - Q) * Λ
    return stationary_distribution(DiscreteMarkovChain(ones(nb_states(mmpp)), P))
end

function forward_backward(mmpp::MMPP, h::TemporalHistory)
    S = nb_states(mmpp)
    π0 = stationary_distribution(mmpp)

    y = diff(vcat(min_time(h), event_times(h), max_time(h)))
    marks = event_marks(h)
    n = nb_events(h)

    Λ = intensity_by_state(mmpp)
    f = [renewal_density(mmpp, y[k], marks[k]) for k = 1:n]
    F̄ = renewal_void_probability(mmpp, y[n+1])

    L = OffsetVector{Vector{Float64}}(undef, 0:n+1)
    c = Vector{Float64}(undef, n + 1)
    L[0] = π0
    for k = 1:n+1
        L[k] = k < n + 1 ? L[k-1] * f[k] : L[n] * F̄
        c[k] = sum(L[k])
        L[k] ./= c[k]
    end

    R = Vector{Vector{Float64}}(undef, n + 2)
    R[n+2] = ones(S)
    for k = n+1:-1:1
        R[n+1] = k == n + 1 ? F̄ * R[n+2] : f[k] * R[k+1]
        R[k] ./= c[k]
    end

    C = Matrix{Float64}(undef, 2S, 2S)
    C[1:S, 1:S] = Q - Λ
    C[S+1:2S, 1:S] .= 0.0
    C[S+1:2S, S+1:2S] = Q - Λ

    ∫ = Vector{Matrix{Float64}}(undef, n + 1)
    for k = 1:n+1
        if k < n + 1
            Λm = intensity_by_state(mmpp, marks[k])
            C[1:S, S+1:2S] = Λm * R[k+1] * L[k-1]
        else
            C[1:S, S+1:2S] = R[k+1] * L[k-1]
        end
        ∫[k] = exp(C[k] * y[k])[1:S, S+1:2S]
    end

    Q = rate_matrix(mmpp)

    m̂ = Q .* sum(∫[k]' / c[k] for k = 1:n+1)
    D̂ = diag(m̂) ./ diag(Q)
    n̂ = [DefaultDict{M,Float64}(0.0) for s = 1:S]
    for k = 1:n
        local_n̂ = L[k]' .* R[k+1]
        for s = 1:S
            n̂[s][marks[k]] += local_n̂[s]
        end
    end

    logL = sum(log.(c))

    return m̂, n̂, D̂, logL
end

function ryden(mmpp::MMPP{M,Tr,Em}, h::TemporalHistory{M}; iterations) where {M,Tr,Em}
    logL_evolution = Float64[]
    for _ = 1:iterations
        m̂, n̂, D̂, logL = forward_backward(mmpp, h)
        push!(logL_evolution, logL)
        new_transitions = fit(Tr, m̂ = m̂, D̂ = D̂)
        new_emissions = [fit(Em, n̂ = n̂[s], D̂ = D̂[s]) for s = 1:nb_states(mmpp)]  # TODO: @view
        mmpp = MMPP(new_transitions, new_emissions)
    end
    return mmpp, logL_evolution
end


text:
# Julia wrapper for header: /usr/include/libavutil/log.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

export
    av_vlog,
    av_log_get_level,
    av_log_set_level,
    av_log_set_callback,
    av_log_default_callback,
    av_default_item_name,
    av_log_set_flags

function av_vlog(avcl,level::Integer,fmt,arg1)
    ccall((:av_vlog,libavutil),Void,(Ptr{Void},Cint,Ptr{UInt8},Ptr{__va_list_tag}),avcl,level,fmt,arg1)
end

function av_log_get_level()
    ccall((:av_log_get_level,libavutil),Cint,())
end

function av_log_set_level(arg1::Integer)
    ccall((:av_log_set_level,libavutil),Void,(Cint,),arg1)
end

function av_log_set_callback(arg1)
    ccall((:av_log_set_callback,libavutil),Void,(Ptr{Void},),arg1)
end

function av_log_default_callback(ptr,level::Integer,fmt,vl)
    ccall((:av_log_default_callback,libavutil),Void,(Ptr{Void},Cint,Ptr{UInt8},Ptr{__va_list_tag}),ptr,level,fmt,vl)
end

function av_default_item_name(ctx)
    ccall((:av_default_item_name,libavutil),Ptr{UInt8},(Ptr{Void},),ctx)
end

function av_log_set_flags(arg::Integer)
    ccall((:av_log_set_flags,libavutil),Void,(Cint,),arg)
end


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

module TrajSaveCommon

export extract_params!, extract_params, trajLoggedPlay

using AdaptiveStressTesting
using CPUTime
using RLESUtils, Obj2Dict, RunCases, Observers, Loggers, Obj2DataFrames
using DataFrames

using ..DefineSave
using ..SaveTypes
using ..DefineLog
using ..SaveHelpers

function extract_params!(paramObj, case::Case, key::AbstractString)
  #assumes format of key is: "key.field"

  for (k, v) in case
    k_ = split(k, '.')

    if length(k_) < 2
      warn("extract_params!::dot separator not found in $k")
    end

    if k_[1] == key
      if isdefined(paramObj, Symbol(k_[2]))
        setfield!(paramObj, Symbol(k_[2]), v)
      else
        warn("$(k_[2]) is not a member of type $(typeof(paramObj))")
      end
    end
  end

  return paramObj
end

function extract_params(case::Case, key::AbstractString)
    kwargs = Tuple{Symbol,Any}[]
    for (k, v) in case
        toks = split(k, '.')
        if length(toks) < 2
            warn("extract_params::dot separator not found in $k")
        end
        if toks[1] == key
            push!(kwargs, (toks[2], v))
        end
    end
    kwargs
end

"""
Replay trajectory instrumented with logs
"""
function trajLoggedPlay(ast::AdaptiveStressTest, reward, action_seq,
        compute_info,
        sim_params,
        ast_params::ASTParams;
        suppress_warn::Bool=false
        )

    sim = ast.sim

    #replay to get logs
    log = addObservers(ast.sim)
    replay_reward, action_seq2 = play_sequence(ast, action_seq)

    @notify_observer(sim.observer, "run_info", Any[replay_reward, sim.md_time, sim.hmd,
        sim.vmd, sim.label_as_nmac])

    #sanity check replay
    @assert action_seq2 == action_seq
    if !suppress_warn && !isapprox(replay_reward, reward)
        warn("trajLoggedPlay: replay reward is different than original reward. replay_reward=$(replay_reward) and reward=$reward")
    end

    #Save
    set_compute_info!(log, compute_info)
    set_sim_params!(log, sim_params)
    set_ast_params!(log, ast_params)
    set_action_seq!(log, action_seq)

    log
end

end #module

max_stars_repo_name: sisl/RLESCAS

max_issues_repo_name: sisl/RLESCAS

max_forks_repo_name: sisl/RLESCAS


text: println("Placeholder for a full test set.")


text:
@testset "Widgets" begin
    entry = Entry(text = "Foo")
    @test entry isa GtkEntry

    set!(entry, text = "Hello")
    @test getprop(entry, :text, String) == "Hello"

    set!(entry, text = "Ok")
    @test getprop(entry, :text, String) == "Ok"

    @test Label("Foo")              isa GtkLabel
    @test ProgressBar()             isa GtkProgressBar
    @test Scale(true, 0:10)         isa GtkScale
    @test Slider(0:10, init = 5)    isa GtkScale
    @test SpinButton(0:10)          isa GtkSpinButton
    @test Spinner()                 isa GtkSpinner
    @test Switch(true)              isa Gtk.GtkSwitch

    color_btn =  ColorButton()
    @test color_btn isa GtkColorButton

    set!(color_btn, rgba = colorant"#f1a")
    @test value(color_btn) == RGBA(colorant"#f1a")
end

@testset "Containers" begin
    win = Window("Window", 800, 600) do
        Box(:v) do
            Frame("Hello") do
                Paned(:h) do
                    # Expander(())
                    Box(:v) do
                        button"Hello",
                        SpinButton(1:100),
                        colorbtn"#1af"
                    end, # Box
                    Box(:v) do
                        entry"World!",
                        Slider(-10:10, @hexpand),
                        label"Ok"
                    end # Box
                end # Paned
            end # Frame
        end # Box
    end # Window

    showall(win)
    destroy(win)
end

@testset "Events" begin
    i = Ref{Int}(0)
    slider = Slider(1:50,       @on value changed () -> i[] += 1)
    spin = SpinButton(0:1000,   @on value changed () -> i[] += 1)
    entry = Entry(text = "Ok",  @on changed () -> i[] += 1)
    btn = Button("Click-me!",   @on clicked () -> println("Button was clicked!"))
    checkbox = CheckBox(        @on clicked () -> println("CheckBox was clicked!"))
    switch = Switch(true)
    cbtn = ColorButton()

    @test i[] == 3

    onevent(:state_set, switch) do
        # error test
        @test_throws UndefVarError println(ok)
    end

    set!(switch, active = false)

    onevent(:clicked, cbtn) do
        println("Color Button value changed!")
    end

    win = Window("Event test", 800, 600) do
        Box(:v, @margin(20), spacing = 10) do
            slider,
            spin,
            btn,
            entry,
            checkbox,
            switch,
            cbtn
        end
    end

    showall(win)
    destroy(win)
end

@testset "Prop Helpers" begin
    value = 10
    box = Box((), :h, @margin(value), @hexpand, @vexpand(false), @align(:center, :fill))

    for m in [:top, :right, :bottom, :left]
        @test getprop(box, "margin-$m", Int) == value
    end

    set!(box, @margin(value, 2value))

    @test getprop(box, "margin-top", Int) == value
    @test getprop(box, "margin-bottom", Int) == value

    @test getprop(box, "margin-left", Int) == 2value
    @test getprop(box, "margin-right", Int) == 2value

    @test getprop(box, :valign, UInt32) == GtkAlign.CENTER
    @test getprop(box, :halign, UInt32) == GtkAlign.FILL

    @test getprop(box, :hexpand, Bool) == true
    @test getprop(box, :vexpand, Bool) == false
end


text:
function kron_δ(i,j)
    if i == j
        return 1
    else
        return 0
    end
end

function nbm_matrix(g::SimpleGraph)
    edgeidmap = Dict{Edge, Int}()
    aristas = ne(g)
    m = 0
    for e in edges(g)
        m += 1
        edgeidmap[e] = m
	edgeidmap[reverse(e)] = m + aristas
    end

    B = spzeros(Int64, 2*aristas, 2*aristas)

    for (e,u) in edgeidmap
        i, j = src(e), dst(e)
	eles = neighbors(g,j)
	k = j
        for l in eles
          B[edgeidmap[k=>l],u] = (kron_δ(j,k)*(1-kron_δ(i,l)))
        end
    end
    return B, edgeidmap
end


text:
@inline function add_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(add_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? add_dbdb_db_nonfinite(x,y) : res
end

@inline function sub_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(sub_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? sub_dbdb_db_nonfinite(x,y) : res
end

@inline function mul_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(mul_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? mul_dbdb_db_nonfinite(x,y) : res
end

@inline function dvi_dbdb_db(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    res = DoubleFloat{T}(dvi_dddd_dd(HILO(x), HILO(y)))
    isnan(HI(res))  ? dvi_dbdb_db_nonfinite(x,y) : res
end

@inline function add_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = HI(x) + HI(y)
    return DoubleFloat(z, zero(T))
end

@inline function sub_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = HI(x) - HI(y)
    return DoubleFloat(z, zero(T))
end

@inline function mul_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = HI(x) * HI(y)
    return DoubleFloat(z, zero(T))
end

@inline function dvi_dbdb_db_nonfinite(x::DoubleFloat{T}, y::DoubleFloat{T}) where {T<:AbstractFloat}
    z = div(HI(x), HI(y))
    return DoubleFloat(z, zero(T))
end


text:
module NNN

using InplaceOps
using Base.LinAlg.BLAS

include("type.jl")
include("backprop.jl")
include("activation.jl")
include("filter.jl")
include("cost.jl")
include("optimize.jl")
include("util.jl")

export Layer,
       HiddenLayer,
       OutputLayer,
       RecurrentLayer,
       Network,
       Optimizer,
       cross_entropy_cost,
       rnn_cross_entropy_cost,
       quadratic_cost,
       rnn_quadratic_cost,
       report,
       train,
       gradient_clip!,
       Filter,
       Dropout

end


text:
module RemPiO2

import Base: TwicePrecision, significand_mask
import Base.Math: significand_bits, exponent_mask, exponent_bias

include("codywaite.jl")
include("paynehanek.jl")

"""
    highword(x)

returns the high word of x as a UInt32.
"""
@inline highword(x::UInt64) = unsafe_trunc(UInt32,x >> 32)
@inline highword(x::Float64) = unsafe_trunc(UInt32,highword(reinterpret(UInt64, x)))

"""
    rem_pio2(x::Float64)

returns the remainder of x modulo π/2 as a TwicePrecision number, along with a k
the specifies
"""

function rem_pio2(x::Float64)
    xh = highword(x)
    xhp = xh & 0x7fffffff # positive part of highword

    if xhp <= 0x3fe921fb #|x| ~<= pi/4, no need for reduction, just return input
        return Int(0), TwicePrecision(x, 0.0)
    end
    rem_pio2_kernel(x, xh, xhp)
end

"""
    rem_pio2_kernel(x, xh, xhp)

returns the remainder of x modulo π/2 as a TwicePrecision number, along with a k
such that.
"""
function rem_pio2_kernel(x)
    xh = highword(x)
    xhp = xh & 0x7fffffff # positive part of highword
    rem_pio2_kernel(x, xh, xhp)
end
function rem_pio2_kernel(x, xh, xhp)
    if xhp <= 0x400f6a7a #|x| ~<= 5pi/4, use Cody Waite with two constants
        if (xhp & 0xfffff) == 0x921fb # |x| ~= pi/2 or 2pi/2, use precise Cody Waite scheme
            return cody_waite_ext_pio2(x, xhp)
        end
        if xhp <= 0x4002d97c # |x| ~<= 3pi/4
            if x > 0
                return 1, cody_waite_2c_pio2(x, 1.0)
            else
                return -1, cody_waite_2c_pio2(x, -1.0)
            end
        else
            if x > 0
                return 2, cody_waite_2c_pio2(x, 2.0)
            else
                return -2, cody_waite_2c_pio2(x, -2.0)
            end
        end
    end

    if xhp <= 0x401c463b # |x| ~<= 9pi/4, use Cody Waite with two constants
        if (xhp <= 0x4015fdbc) # |x| ~<= 7pi/4
            if (xhp == 0x4012d97c) # |x| ~= 3pi/2, use precise Cody Waite scheme
                return cody_waite_ext_pio2(x, xhp)
            end
            if x > 0
                return 3, cody_waite_2c_pio2(x, 3.0)
            else
                return -3, cody_waite_2c_pio2(x, -3.0)
            end
        else
            if xhp == 0x401921fb # |x| ~= 4pi/2, use precise Cody Waite scheme
                return cody_waite_ext_pio2(x, xhp)
            end
            if x > 0
                return 4, cody_waite_2c_pio2(x, 4.0)
            else
                return -4, cody_waite_2c_pio2(x, -4.0)
            end
        end
    end
    if xhp<0x413921fb # |x| ~< 2^20*pi/2, use precise Cody Waite scheme
        return cody_waite_ext_pio2(x, xhp)
    end

    # if |x| >= 2^20*pi/2 switch to Payne Hanek
    return paynehanek(x)
end

end # module


text:
begin
    using Revise
    using BayesNegativeBinomial
    using BNPRegressionGGA2021
    using DataFrames
    using Distributions
    using FileIO
    using Gadfly
    using Random
    using RCall
    using Statistics
    using StatsBase
    using LinearAlgebra
    const BNP = BNPRegressionGGA2021
end

function foo(dy, x0, x1; filename, discrete = true, mcmcsize = 10000)
    # Simulate the responses
    y0 = @. rand(dy(x0))
    y1 = LinRange(-3, 3, 50) |> collect

    # Expand the grid
    R"""
    grid <- expand.grid(y1 = $y1, x1 = $x1)
    y1 <- grid$y1
    x1 <- grid$x1
    """
    @rget y1 x1

    # Expand x0-x1 using splines
    if discrete == true
        X0 = x0[:, :]
        X1 = x1[:, :]
    else
        x0min, x0max = extrema(x0)
        R"""
        X0 <- splines::bs($x0, df = 6, Boundary.knots = c($x0min, $x0max))
        X1 <- predict(X0, $x1)
        """
        @rget X0 X1
    end

    # Standardise responses and predictors
    if discrete == false
        mX0, sX0 = mean_and_std(X0, 1)
        for col in 1:size(X0, 2)
            if sX0[col] > 0
                X0[:, col] = (X0[:, col] .- mX0[col]) ./ sX0[col]
                X1[:, col] = (X1[:, col] .- mX0[col]) ./ sX0[col]
            end
        end
    end
    my0, sy0 = mean_and_std(y0, 1)
    y0 = (y0 .- my0) ./ sy0

    # Add a constant term to the design matrices
    X0 = [ones(size(X0, 1)) X0]
    X1 = [ones(size(X1, 1)) X1]

    # Generate the true (standardized) density over the grid
    f1 = pdf.(dy.(x1), my0 .+ y1 .* sy0) .* sy0

    # Fit the model
    smpl = BNP.DGSBPNormal(; y0, X0, y1, X1);
    fhchain, _ = BNP.sample!(smpl; mcmcsize);
    fh = mean(fhchain)

    # Organize the results as a DataFrame
    df = DataFrame(; y1, x1, f1, fh)
    df = stack(df, [:f1, :fh])

    # Plot the results
    R"""
    p <-
        $df |>
        ggplot2::ggplot(ggplot2::aes(x = y1, y = value, color = variable)) +
        ggplot2::geom_line() +
        ggplot2::facet_grid(
            cols = ggplot2::vars(x1),
            labeller = ggplot2::labeller(.cols = ggplot2::label_both)
        ) +
        ggplot2::theme_classic()
    ggplot2::ggsave(filename = $filename)
    """
    return y0, y1, X0, X1
end

# Example 1: No predictor (normal distribution)
begin
    dy(x) = Normal(0, 1)
    Random.seed!(1)
    x1 = [-1.0, 1.0]
    x0 = rand(x1, 500)
    filename = "figures/ex-normal-1.png"
    a = foo(dy, x0, x1; filename, mcmcsize = 4000)
end;

# Example 2: No predictor (mixture distribution)
begin
    dy(x) = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    Random.seed!(1)
    x1 = [-1.0, 1.0]
    x0 = rand(x1, 500)
    filename = "figures/ex-normal-2.png"
    a = foo(dy, x0, x1; filename, mcmcsize = 4000)
end;

# Example 3: One discrete predictor (normal distribution)
begin
    dy(x) = Normal(x, 1)
    Random.seed!(1)
    x1 = [-1.0, 1.0]
    x0 = rand(x1, 500)
    filename = "figures/ex-normal-3.png"
    a = foo(dy, x0, x1; filename, mcmcsize = 4000)
end;

# Example 4: One discrete predictor (mixture distribution)
begin
    dy(x) = MixtureModel(Normal, [(x + 1, 1.0), (x - 1, 0.25)], [0.5, 0.5])
    Random.seed!(1)
    x1 = [-1.0, 1.0]
    x0 = rand(x1, 500)
    filename = "figures/ex-normal-4.png"
    a = foo(dy, x0, x1; filename, mcmcsize = 10000)
end;

# Example 1: No predictor ------------------------------------------------------

function simulate_sample_01(N0, N1)
    dy = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    X00 = ones(N0, 1)
    X10 = ones(N1, 1)
    y00 = rand(dy, N0)
    y10 = LinRange(minimum(y00), maximum(y00), N1) |> collect
    y0, y1, X0, X1, my00, sy00 = standardize_data!(y00, y10, X00, X10)
    return y0, y1, X0, X1
end;

begin
    Random.seed!(1)
    N0, N1 = 1000, 50;
    rng = MersenneTwister(1);
    y0, y1, X0, X1 = simulate_sample_01(N0, N1);
    smpl = BNP.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNP.sample!(smpl; mcmcsize = 10000);
end;

f1 = mean(chainf);
plot(
    layer(x = y1, y = f1, Geom.line, color=["bnp"]),
    layer(x = y0, Geom.density, color=["kden"]),
    layer(x = y0, Geom.histogram(density = true), color=["hist"])
)

#========================================================#
# Example 2 - 1 discrete predictor (simple distribution) #
#========================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(x0, N0)]
    y0 = X0 * ones(2) + randn(N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 2)
    x1 = [x0[1] * ones(N1); x0[2] * ones(N1)]
    X1 = [ones(2 * N1) x1]
    return y0, X0, y1, X1
end;

begin
    Random.seed!(1)
    N0, N1 = 1000, 50;
    y0, X0, y1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
dy(x) = Normal(1 + x, 1);
plot(x = y0, color = X0[:, 2], Geom.density)
plot(
    layer(x = y1, y = fb, color = X1[:, 2], Geom.line),
    layer(x = y1, y = pdf.(dy.(X1[:, 2]), y1), color = X1[:, 2], Geom.line)
)
mean(chainβ)

#=========================================================#
# Example 3 - 1 discrete predictor (mixture distribution) #
#=========================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(x0, N0)]
    y0 = X0 * ones(2) + rand(dϵ, N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 2)
    x1 = [x0[1] * ones(N1); x0[2] * ones(N1)]
    X1 = [ones(2 * N1) x1]
    return y0, X0, y1, X1
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, X0, y1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
plot(x = y0, color = X0[:, 2], Geom.histogram(density = true, bincount = 50))
plot(x = y0, color = X0[:, 2], Geom.density)
plot(x = y1, y = fb, color = X1[:, 2], Geom.line)
mean(chainβ)

#=================================================#
# Example 4 - 2 predictors (mixture distribution) #
#=================================================#

# Simulate a sample
function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(1.0, 1.0), (-1.0, 0.25)], [0.5, 0.5])
    x0 = [-0.5, 0.5]
    X0 = [ones(N0) rand(x0, N0, 2)]
    y0 = X0 * [ones(2); 0] + rand(dϵ, N0)
    y1 = LinRange(minimum(y0), maximum(y0), N1) |> collect |> x -> repeat(x, 4)
    x1 = kron([x0[1], x0[2], x0[1], x0[2]], ones(N1))
    x2 = kron([x0[1], x0[1], x0[2], x0[2]], ones(N1))
    X1 = [ones(4 * N1) x1 x2]
    return y0, X0, y1, X1
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, X0, y1, X1 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

fb = mean(chainf);
X0concat = @. string(X0[:, 2]) * " " * string(X0[:, 3])
X1concat = @. string(X1[:, 2]) * " " * string(X1[:, 3])
plot(x = y0, color = X0concat, Geom.density)
plot(x = y1, y = fb, color = X1concat, Geom.line)
mean(chainβ)

#==========================================================#
# Example 5 - 1 continuous predictor (simple distribution) #
#==========================================================#

function simulate_sample(N0, N1)
    dϵ = Normal(0.0, 1.0)
    x0 = LinRange(-2, 2, N0) |> collect
    X0 = rcopy(R"cbind(1, splines::bs($x0, df = 6, Boundary.knots = c(-4, 4)))")
    y0 = x0 + rand(dϵ, N0)
    x_grid = LinRange(-2, 2, N1) |> collect
    y_grid = LinRange(-3, 3, N1) |> collect
    R"""
    grid <- expand.grid(y1 = $y_grid, x1 = $x_grid)
    y1 <- grid$y1
    x1 <- grid$x1
    """
    @rget y1 x1
    X1 = rcopy(R"cbind(1, predict(splines::bs($x0, df = 6, Boundary.knots = c(-4, 4)), $x1))")
    # Standardization
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0[:, 1] .= 1
    X1[:, 1] .= 1
    return y0, x0, X0, y1, x1, X1, mean_y0, std_y0
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, x0, X0, y1, x1, X1, mean_y0, std_y0 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

dy1(x) = Normal((x - mean_y0[1]) / std_y0[1], 1 / std_y0[1])
f0 = pdf.(dy1.(x1), y1);
fh = mean(chainf);
df = DataFrame(; y1, x1, f0, fh);

R"""
$df |>
    tidyr::pivot_longer(c(f0, fh)) |>
    ggplot2::ggplot(ggplot2::aes(x = x1, y = y1, z = value)) +
    ggplot2::geom_contour() +
    ggplot2::facet_grid(~ name)
ggplot2::ggsave("figures/ex-normal-5.png")
"""

# Save the estimated surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-f0) |>
    tidyr::pivot_wider(values_from = fh, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-5-surface-hat.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "fh"
)
dev.off()
"""

# Save the true surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-fh) |>
    tidyr::pivot_wider(values_from = f0, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-5-surface-true.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "f1"
)
dev.off()
"""

#===========================================================#
# Example 6 - 1 continuous predictor (mixture distribution) #
#===========================================================#

function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(1, 1), (-1, 0.25)], [0.5, 0.5])
    x0 = LinRange(-2, 2, N0) |> collect
    X0 = rcopy(R"cbind(1, splines::bs($x0, df = 8, Boundary.knots = c(-5, 5)))")
    y0 = x0 + rand(dϵ, N0)
    x_grid = LinRange(-2, 2, N1)
    y_grid = LinRange(-3, 3, N1)
    R"""
    grid <- expand.grid(y1 = $y_grid, x1 = $x_grid)
    y1 <- grid$y1
    x1 <- grid$x1
    """
    @rget y1 x1
    X1 = rcopy(R"cbind(1, predict(splines::bs($x0, df = 8, Boundary.knots = c(-5, 5)), $x1))")
    # Standardization
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0[:, 1] .= 1
    X1[:, 1] .= 1
    return y0, x0, X0, y1, x1, X1, mean_y0, std_y0
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, x0, X0, y1, x1, X1, mean_y0, std_y0 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

dy1(x) = MixtureModel(Normal[
        Normal((x - 1 - mean_y0[1]) / std_y0[1], 0.25 / std_y0[1]),
        Normal((x + 1 - mean_y0[1]) / std_y0[1], 1 / std_y0[1])
    ],
    [0.5, 0.5]
)
f0 = pdf.(dy1.(x1), y1);
fh = mean(chainf);
df = DataFrame(; y1, x1, f0, fh);

# Save the estimated surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-f0) |>
    tidyr::pivot_wider(values_from = fh, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-6-surface-hat.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "fh"
)
dev.off()
"""

# Save the true surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-fh) |>
    tidyr::pivot_wider(values_from = f0, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-6-surface-true.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "f1"
)
dev.off()
"""

R"""
png("figures/ex-normal-6-scatter.png")
plot($x0, $y0)
dev.off()
"""

#==================================#
# Example 5 (now without bsplines) #
#==================================#

function simulate_sample(N0, N1)
    dϵ = Normal(0.0, 1.0)
    x0 = LinRange(-2, 2, N0) |> collect
    X0 = x0[:, :]
    y0 = x0 + rand(dϵ, N0)
    x_grid = LinRange(-2, 2, N1) |> collect
    y_grid = LinRange(-3, 3, N1) |> collect
    R"""
    grid <- expand.grid(y1 = $y_grid, x1 = $x_grid)
    y1 <- grid$y1
    x1 <- grid$x1
    """
    @rget y1 x1
    X1 = x1[:, :]
    # Standardization
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0 = [ones(size(X0, 1)) X0]
    X1 = [ones(size(X1, 1)) X1]
    return y0, x0, X0, y1, x1, X1, mean_y0, std_y0
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, x0, X0, y1, x1, X1, mean_y0, std_y0 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, _ = BNP.sample!(smpl; mcmcsize = 10000);
end;

dy1(x) = Normal((x - mean_y0[1]) / std_y0[1], 1 / std_y0[1])
f0 = pdf.(dy1.(x1), y1);
fh = mean(chainf);
df = DataFrame(; y1, x1, f0, fh);

# Save the estimated surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-f0) |>
    tidyr::pivot_wider(values_from = fh, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-5-surface-no-bspline-hat.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "fh"
)
dev.off()
"""

# Save the true surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-fh) |>
    tidyr::pivot_wider(values_from = f0, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-5-surface-no-bspline-true.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "f1"
)
dev.off()
"""

#==============================#
# Example 6 (without bsplines) #
#==============================#

function simulate_sample(N0, N1)
    dϵ = MixtureModel(Normal, [(1, 1), (-1, 0.25)], [0.5, 0.5])
    x0 = LinRange(-2, 2, N0) |> collect
    X0 = x0[:, :]
    y0 = x0 + rand(dϵ, N0)
    x_grid = LinRange(-2, 2, N1) |> collect
    y_grid = LinRange(-3, 3, N1) |> collect
    R"""
    grid <- expand.grid(y1 = $y_grid, x1 = $x_grid)
    y1 <- grid$y1
    x1 <- grid$x1
    """
    @rget y1 x1
    X1 = x1[:, :]
    # Standardization
    mean_y0, std_y0 = mean_and_std(y0, 1)
    mean_X0, std_X0 = mean_and_std(X0, 1)
    y0 = (y0 .- mean_y0) ./ std_y0
    X0 = (X0 .- mean_X0) ./ std_X0
    X1 = (X1 .- mean_X0) ./ std_X0
    X0 = [ones(size(X0, 1)) X0]
    X1 = [ones(size(X1, 1)) X1]
    return y0, x0, X0, y1, x1, X1, mean_y0, std_y0
end

begin
    Random.seed!(1)
    N0, N1 = 1000, 50
    y0, x0, X0, y1, x1, X1, mean_y0, std_y0 = simulate_sample(N0, N1);
    smpl = BNPRegressionGGA2021.DGSBPNormal(; y0, X0, y1, X1);
    chainf, chainβ = BNPRegressionGGA2021.sample!(smpl; mcmcsize = 10000);
end;

dy1(x) = MixtureModel(Normal[
        Normal((x - 1 - mean_y0[1]) / std_y0[1], 0.25 / std_y0[1]),
        Normal((x + 1 - mean_y0[1]) / std_y0[1], 1 / std_y0[1])
    ],
    [0.5, 0.5]
)
f0 = pdf.(dy1.(x1), y1);
fh = mean(chainf);
df = DataFrame(; y1, x1, f0, fh);

# Save the estimated surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-f0) |>
    tidyr::pivot_wider(values_from = fh, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-6-no-bspline-surface-hat.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "fh"
)
dev.off()
"""

# Save the true surface
R"""
xv <- sort(unique($df[["x1"]]))
yv <- sort(unique($df[["y1"]]))
zv <- $df |>
    dplyr::select(-fh) |>
    tidyr::pivot_wider(values_from = f0, names_from = y1) |>
    dplyr::select(-x1) |>
    as.matrix()
png("figures/ex-normal-6-no-bspline-surface-true.png")
fig <- persp(
    x = xv, y = yv, z = zv, phi = 45, theta = 45,
    xlab = "x", ylab = "y", zlab = "f1"
)
dev.off()
"""

R"""
png("figures/ex-normal-6-no-bspline-scatter.png")
plot($x0, $y0)
dev.off()
"""