text:
mutable struct NoCrashRewardModel <: AbstractMLRewardModel
    cost_dangerous_brake::Float64 # POSITIVE NUMBER
    reward_in_target_lane::Float64 # POSITIVE NUMBER

    brake_penalty_thresh::Float64 # (POSITIVE NUMBER) if the deceleration is greater than this cost_dangerous_brake will be accured
    target_lane::Int
end

#XXX temporary
NoCrashRewardModel() = NoCrashRewardModel(100.,10.,2.5,4)
lambda(rm::NoCrashRewardModel) = rm.cost_dangerous_brake/rm.reward_in_target_lane

mutable struct NoCrashIDMMOBILModel{G<:BehaviorGenerator} <: AbstractMLDynamicsModel
    nb_cars::Int
    phys_param::PhysicalParam

    # behaviors::Vector{BehaviorModel}
    # behavior_probabilities::Weights
    behaviors::G

    adjustment_acceleration::Float64
    lane_change_rate::Float64 # in LANES PER SECOND

    p_appear::Float64 # probability of a new car appearing if the maximum number are not on the road
    appear_clearance::Float64 # minimum clearance for a car to appear

    vel_sigma::Float64 # std of new car speed about v0

    lane_terminate::Bool # if true, terminate the simulation when the car has reached the desired lane
    brake_terminate_thresh::Float64 # terminate simulation if braking is above this (always positive)
    max_dist::Float64 # terminate simulation if distance becomes greater than this
    speed_terminate_thresh::Float64 # terminate simulation if any speed is below this
end

function NoCrashIDMMOBILModel(nb_cars::Int,
                              pp::PhysicalParam;
                              vel_sigma = 0.5,
                              lane_terminate=false,
                              p_appear=0.5,
                              brake_terminate_thresh=Inf,
                              max_dist=Inf,
                              speed_terminate_thresh=-Inf,
                              behaviors=DiscreteBehaviorSet(IDMMOBILBehavior[IDMMOBILBehavior(x[1],x[2],x[3],idx) for (idx,x) in
                                                 enumerate(Iterators.product(["cautious","normal","aggressive"],
                                                        [pp.v_slow+0.5;pp.v_med;pp.v_fast],
                                                        [pp.l_car]))], Weights(ones(9)))
                              )

    return NoCrashIDMMOBILModel(
        nb_cars,
        pp,
        behaviors,
        1., # adjustment accel
        1.0/(2.0*pp.dt), # lane change rate
        p_appear, # p_appear
        35.0, # appear_clearance
        vel_sigma, # vel_sigma
        lane_terminate,
        brake_terminate_thresh,
        max_dist,
        speed_terminate_thresh
    )
end

function NoCrashIDMMOBILModel(old::NoCrashIDMMOBILModel, behaviors::BehaviorGenerator)
    return NoCrashIDMMOBILModel(old.nb_cars, old.phys_param, behaviors, old.adjustment_acceleration, old.lane_change_rate, old.p_appear, old.appear_clearance, old.vel_sigma, old.lane_terminate, old.brake_terminate_thresh, old.max_dist, old.speed_terminate_thresh)
end

const NoCrashMDP{R<:AbstractMLRewardModel, G} =  MLMDP{MLState, MLAction, NoCrashIDMMOBILModel{G}, R}
const NoCrashPOMDP{R<:AbstractMLRewardModel, G} =  MLPOMDP{MLState, MLAction, MLObs, NoCrashIDMMOBILModel{G}, R}

const NoCrashProblem{R<:AbstractMLRewardModel,G} =  Union{NoCrashMDP{R,G}, NoCrashPOMDP{R,G}}

gen_type(::Type{MLPOMDP{S,A,O,NoCrashIDMMOBILModel{G},R}}) where {S,A,O,G,R} = G

# create_action(::NoCrashProblem) = MLAction()

# action space = {a in {accelerate,maintain,decelerate}x{left_lane_change,maintain,right_lane_change} | a is safe} U {brake}
struct NoCrashActionSpace
    NORMAL_ACTIONS::Vector{MLAction} # all the actions except brake
    acceptable::IntSet
    brake::MLAction # this action will be EITHER braking at half the dangerous brake threshold OR the braking necessary to prevent a collision at all time in the future
end

const NB_NORMAL_ACTIONS = 9

function NoCrashActionSpace(mdp::NoCrashProblem)
    accels = (-mdp.dmodel.adjustment_acceleration, 0.0, mdp.dmodel.adjustment_acceleration)
    lane_changes = (-mdp.dmodel.lane_change_rate, 0.0, mdp.dmodel.lane_change_rate)
    NORMAL_ACTIONS = MLAction[MLAction(a,l) for (a,l) in Iterators.product(accels, lane_changes)] # this should be in the problem
    return NoCrashActionSpace(NORMAL_ACTIONS, IntSet(), MLAction()) # note: brake will be calculated later based on the state
end

function actions(mdp::NoCrashProblem, s::Union{MLState, MLPhysicalState})
    as = NoCrashActionSpace(mdp)
    acceptable = IntSet()
    for i in 1:NB_NORMAL_ACTIONS
        a = as.NORMAL_ACTIONS[i]
        ego_y = s.cars[1].y
        # prevent going off the road
        if ego_y == 1. && a.lane_change < 0. || ego_y == mdp.dmodel.phys_param.nb_lanes && a.lane_change > 0.0
            continue
        end
        # prevent running into the person in front or to the side
        if is_safe(mdp, s, as.NORMAL_ACTIONS[i])
            push!(acceptable, i)
        end
    end
    brake_acc = calc_brake_acc(mdp, s)
    brake = MLAction(brake_acc, 0)
    return NoCrashActionSpace(as.NORMAL_ACTIONS, acceptable, brake)
end

actions(mdp::NoCrashProblem, b::AbstractParticleBelief) = actions(mdp, MLPhysicalState(first(particles(b))))

calc_brake_acc(mdp::NoCrashProblem, s::Union{MLState, MLPhysicalState}) = min(max_safe_acc(mdp,s), -mdp.rmodel.brake_penalty_thresh/2.0)
# calc_brake_acc(mdp::NoCrashProblem{TargetLaneReward}, s::Union{MLState, MLPhysicalState}) = min(max_safe_acc(mdp, s), min(-mdp.dmodel.phys_param.brake_limit/2, -mdp.dmodel.brake_terminate_thresh/2))

iterator(as::NoCrashActionSpace) = as
Base.start(as::NoCrashActionSpace) = 1
function Base.next(as::NoCrashActionSpace, state::Integer)
    while !(state in as.acceptable)
        if state > NB_NORMAL_ACTIONS
            return (as.brake, state+1)
        end
        state += 1
    end
    return (as.NORMAL_ACTIONS[state], state+1)
end
Base.done(as::NoCrashActionSpace, state::Integer) = state > NB_NORMAL_ACTIONS+1
Base.length(as::NoCrashActionSpace) = length(as.acceptable) + 1

function rand(rng::AbstractRNG, as::NoCrashActionSpace, a::MLAction=MLAction())
    nb_acts = length(as.acceptable)+1
    index = rand(rng,1:nb_acts)
    for (i,a) in enumerate(as) # is this efficient at all ?
        if i == index
            return a
        end
    end
end

"""
Calculate the maximum safe acceleration that will allow the car to avoid a collision if the car in front slams on its brakes
"""
function max_safe_acc(mdp::NoCrashProblem, s::Union{MLState,MLObs}, lane_change::Float64=0.0)
    dt = mdp.dmodel.phys_param.dt
    v_min = mdp.dmodel.phys_param.v_min
    l_car = mdp.dmodel.phys_param.l_car
    bp = mdp.dmodel.phys_param.brake_limit
    ego = s.cars[1]

    car_in_front = 0
    smallest_gap = Inf
    # find car immediately in front
    if length(s.cars) > 1
        for i in 2:length(s.cars)#nb_cars
            if occupation_overlap(s.cars[i].y, 0.0, ego.y, lane_change) # occupying same lane
                gap = s.cars[i].x - ego.x - l_car
                if gap >= -l_car && gap < smallest_gap
                    car_in_front = i
                    smallest_gap = gap
                end
            end
        end
        # calculate necessary acceleration
        if car_in_front == 0
            return Inf
        else
            n_brake_acc = nullable_max_safe_acc(smallest_gap, ego.vel, s.cars[car_in_front].vel, bp, dt)
            return get(n_brake_acc, -mdp.dmodel.phys_param.brake_limit)
        end
    end
    return Inf
end

"""
Return the maximum acceleration that the car behind can have on this step so that it won't hit the car in front if it slams on its brakes
"""
function max_safe_acc(gap, v_behind, v_ahead, braking_limit, dt)
    bp = braking_limit
    v = v_behind
    vo = v_ahead
    g = gap
    # VVV see mathematica notebook
    return - (bp*dt + 2.*v - sqrt(8.*g*bp + bp^2*dt^2 - 4.*bp*dt*v + 4.*vo^2)) / (2.*dt)
end

"""
Return max_safe_acc or an empty Nullable if the discriminant is negative.
"""
function nullable_max_safe_acc(gap, v_behind, v_ahead, braking_limit, dt)
    bp = braking_limit
    v = v_behind
    vo = v_ahead
    g = gap
    # VVV see mathematica notebook
    discriminant = 8.*g*bp + bp^2*dt^2 - 4.*bp*dt*v + 4.*vo^2
    if discriminant >= 0.0
        return Nullable{Float64}(- (bp*dt + 2.*v - sqrt(discriminant)) / (2.*dt))
    else
        return Nullable{Float64}()
    end
end

"""
Test whether, if the ego vehicle takes action a, it will always be able to slow down fast enough if the car in front slams on his brakes and won't pull in front of another car so close they can't stop
"""
function is_safe(mdp::NoCrashProblem, s::Union{MLState,MLObs}, a::MLAction)
    dt = mdp.dmodel.phys_param.dt
    if a.acc >= max_safe_acc(mdp, s, a.lane_change)
        return false
    end
    # check whether we will go into anyone else's lane so close that they might hit us or we might run into them
    if isinteger(s.cars[1].y) && a.lane_change != 0.0
        l_car = mdp.dmodel.phys_param.l_car
        for i in 2:length(s.cars)
            car = s.cars[i]
            ego = s.cars[1]
            if car.x < ego.x + l_car && occupation_overlap(ego.y, a.lane_change, car.y, 0.0)  # ego is in front of car
                # New definition of safe - the car behind can brake at max braking to avoid the ego if the ego
                # slams on his brakes
                # XXX IS THIS RIGHT??
                # I think I need a better definition of "safe" here
                gap = ego.x - car.x - l_car
                if gap <= 0.0
                    return false
                end
                n_braking_acc = nullable_max_safe_acc(gap, car.vel, ego.vel, mdp.dmodel.phys_param.brake_limit, dt)

                # if isnull(n_braking_acc) || get(n_braking_acc) < -mdp.dmodel.phys_param.brake_limit
                if isnull(n_braking_acc) || get(n_braking_acc) < max_accel(mdp.dmodel.behaviors)
                    return false
                end
            end
        end
    end
    return true
end

#XXX temp
create_state(p::NoCrashProblem) = MLState(0.0, 0.0, Vector{CarState}(p.dmodel.nb_cars), nothing)
create_observation(pomdp::NoCrashPOMDP) = MLObs(0.0, 0.0, Vector{CarStateObs}(pomdp.dmodel.nb_cars), nothing)

function generate_s(mdp::NoCrashProblem, s::MLState, a::MLAction, rng::AbstractRNG)

    @if_debug dbg_rng = copy(rng)

    sp::MLState=create_state(mdp)

    try
        pp = mdp.dmodel.phys_param
        dt = pp.dt
        nb_cars = length(s.cars)
        resize!(sp.cars, nb_cars)
        sp.terminal = s.terminal
        sp.t = s.t + dt

        ## Calculate deltas ##
        #====================#

        dxs = Vector{Float64}(nb_cars)
        dvs = Vector{Float64}(nb_cars)
        dys = Vector{Float64}(nb_cars)
        lcs = Vector{Float64}(nb_cars)

        # agent
        dvs[1] = a.acc*dt
        dxs[1] = s.cars[1].vel*dt + a.acc*dt^2/2.
        lcs[1] = a.lane_change
        dys[1] = a.lane_change*dt

        for i in 2:nb_cars
            neighborhood = get_neighborhood(pp, s, i)

            behavior = s.cars[i].behavior

            acc = gen_accel(behavior, mdp.dmodel, s, neighborhood, i, rng)
            dvs[i] = dt*acc
            dxs[i] = (s.cars[i].vel + dvs[i]/2.)*dt

            lcs[i] = gen_lane_change(behavior, mdp.dmodel, s, neighborhood, i, rng)
            dys[i] = lcs[i] * dt
        end

        ## Consistency checking ##
        #========================#

        # first prevent lane changes into each other
        changers = IntSet()
        for i in 1:nb_cars
            if lcs[i] != 0
                push!(changers, i)
            end
        end
        sorted_changers = sort!(collect(changers), by=i->s.cars[i].x, rev=true) # this might be slow because anonymous functions are slow
        # from front to back

        if length(sorted_changers) >= 2 #something to compare
            # iterate through pairs
            iter_state = start(sorted_changers)
            j, iter_state = next(sorted_changers, iter_state)
            while !done(sorted_changers, iter_state)
                i = j
                j, iter_state = next(sorted_changers, iter_state)
                car_i = s.cars[i] # front
                car_j = s.cars[j] # back

                # check if they are both starting to change lanes on this step
                if isinteger(car_i.y) && isinteger(car_j.y)

                    # check if they are near each other lanewise
                    if abs(car_i.y - car_j.y) <= 2.0

                        # make sure there is a conflict longitudinally
                        # if car_i.x - car_j.x <= pp.l_car || car_i.x + dxs[i] - car_j.x + dxs[j] <= pp.l_car
                        # if car_i.x - car_j.x <= mdp.dmodel.appear_clearance # made more conservative on 8/19
                        # if car_i.x - car_j.x <= get_idm_s_star(car_j.behavior.p_idm, car_j.vel, car_j.vel-car_i.vel) # upgraded to sstar on 8/19
                        ivp = car_i.vel + dt*dvs[i]
                        jvp = car_j.vel + dt*dvs[j]
                        ixp = car_i.x + dt*(car_i.vel + ivp)/2.0
                        jxp = car_j.x + dt*(car_j.vel + jvp)/2.0
                        n_max_acc_p = nullable_max_safe_acc(ixp-jxp-pp.l_car, jvp, ivp, pp.brake_limit,dt)
                        if ixp - jxp <= pp.l_car || car_i.x - car_j.x <= pp.l_car || isnull(n_max_acc_p) || get(n_max_acc_p) < -pp.brake_limit

                            # check if they are moving towards each other
                            # if dys[i]*dys[j] < 0.0 && abs(car_i.y+dys[i] - car_j.y+dys[j]) < 2.0
                            if true # prevent lockstepping 8/19 (doesn't prevent it that well)

                                # make j stay in his lane
                                dys[j] = 0.0
                                lcs[j] = 0.0
                            end
                        end
                    end
                end
            end
        end

        # second, prevent cars hitting each other due to noise
        if mdp.throw
            # sorted = sort!(collect(1:length(s.cars)), by=i->s.cars[i].x, rev=true)
            sorted = sortperm(collect(c.x for c in s.cars), rev=true)

            if length(sorted) >= 2 #something to compare
                # iterate through pairs
                iter_state = start(sorted)
                j, iter_state = next(sorted, iter_state)
                while !done(sorted, iter_state)
                    i = j
                    j, iter_state = next(sorted, iter_state)
                    if j == 1
                        continue # don't check for the ego since the ego does not have noise
                    end
                    car_i = s.cars[i]
                    car_j = s.cars[j]

                    # check if they overlap longitudinally
                    if car_j.x + dxs[j] > car_i.x + dxs[i] - pp.l_car

                        # check if they will be in the same lane
                        if occupation_overlap(car_i.y + dys[i], car_j.y + dys[j])
                            # warn and nudge behind
                            if mdp.throw
                                @show car_j.x
                                @show car_i.x
                                @show car_j.x + dxs[j]
                                @show car_i.x + dxs[i]
                                @show pp.l_car
                            end
                            @if_debug begin
                                println("Conflict because of noise: front:$i, back:$j")
                                # Gallium.@enter generate_s(mdp, s, a, dbg_rng)
                                fname = tempname()*".jld"
                                println("saving debug args to $fname")
                                JLD.@save(fname, mdp, s, a, dbg_rng)
                            end
                            if i == 1
                                if mdp.throw
                                    error("Car nudged because of crash (ego in front).")
                                end
                            else
                                # warn("Car nudged because noise would cause a crash.")
                                if mdp.throw
                                    error("Car nudged because of crash.")
                                end
                            end
                            dxs[j] = car_i.x + dxs[i] - car_j.x - 1.01*pp.l_car
                            dvs[j] = 2.0*(dxs[j]/dt - car_j.vel)
                        end
                    end
                end
            end
        end

        ## Dynamics and Exits ##
        #======================#

        exits = IntSet()
        for i in 1:nb_cars
            car = s.cars[i]
            xp = car.x + (dxs[i] - dxs[1])
            yp = car.y + dys[i]
            # velp = max(min(car.vel + dvs[i],pp.v_max), pp.v_min)
            velp = max(car.vel + dvs[i], 0.0) # removed speed limits on 8/13
            # note lane change is updated above

            if velp < mdp.dmodel.speed_terminate_thresh
                sp.terminal = Nullable{Symbol}(:speed)
            end
            if dvs[i]/dt < -mdp.dmodel.brake_terminate_thresh
                sp.terminal = Nullable{Symbol}(:brake)
            end

            # check if a lane was crossed and snap back to it
            if isinteger(car.y)
                # prevent a multi-lane change in a single timestep
                if abs(yp-car.y) > 1.
                    yp = car.y + sign(dys[i])
                end
            else # car.y is not an integer
                if floor(yp) >= ceil(car.y)
                    yp = ceil(car.y)
                end
                if ceil(yp) <= floor(car.y)
                    yp = floor(car.y)
                end
            end

            # if yp < 1.0 || yp > pp.nb_lanes
            #     @show i
            #     @show yp
            #     println("mdp = $mdp")
            #     println("s = $s")
            #     println("a = $a")
            # end
            @assert yp >= 1.0 && yp <= pp.nb_lanes

            if xp < 0.0 || xp >= pp.lane_length
                push!(exits, i)
            else
                sp.cars[i] = CarState(xp, yp, velp, lcs[i], car.behavior, s.cars[i].id)
            end
        end

        sp.x = s.x + dxs[1]

        next_id = maximum([c.id for c in s.cars]) + 1

        deleteat!(sp.cars, exits)
        nb_cars -= length(exits)

        ## Generate new cars ##
        #=====================#

        if nb_cars < mdp.dmodel.nb_cars && rand(rng) <= mdp.dmodel.p_appear

            behavior = rand(rng, mdp.dmodel.behaviors)
            vel = typical_velocity(behavior) + randn(rng)*mdp.dmodel.vel_sigma

            clearances = Vector{Float64}(pp.nb_lanes)
            fill!(clearances, Inf)
            closest_cars = Vector{Int}(pp.nb_lanes)
            fill!(closest_cars, 0)
            sstar_margins = Vector{Float64}(pp.nb_lanes)
            if vel > sp.cars[1].vel
                # put at back
                # sstar is the sstar of the new guy
                for i in 1:length(sp.cars)
                    lowlane, highlane = occupation_lanes(sp.cars[i].y, 0.0)
                    back = sp.cars[i].x - pp.l_car
                    if back < clearances[lowlane]
                        clearances[lowlane] = back
                        closest_cars[lowlane] = i
                    end
                    if back < clearances[highlane]
                        clearances[highlane] = back
                        closest_cars[highlane] = i
                    end
                end
                for j in 1:pp.nb_lanes
                    other = closest_cars[j]
                    if other == 0
                        sstar = 0.0
                    else
                        sstar = get_idm_s_star(behavior.p_idm::IDMParam, vel, vel-sp.cars[other].vel::Float64)
                    end
                    sstar_margins[j] = clearances[j] - sstar
                end
            else
                for i in 1:length(sp.cars)
                    lowlane, highlane = occupation_lanes(sp.cars[i].y, 0.0)
                    front = pp.lane_length - (sp.cars[i].x + pp.l_car) # l_car is half the length of the old car plus half the length of the new one
                    if front < clearances[lowlane]
                        clearances[lowlane] = front
                        closest_cars[lowlane] = i
                    end
                    if front < clearances[highlane]
                        clearances[highlane] = front
                        closest_cars[highlane] = i
                    end
                end
                for j in 1:pp.nb_lanes
                    other = closest_cars[j]
                    if other == 0
                        sstar = 0
                    else
                        sstar = get_idm_s_star(sp.cars[other].behavior.p_idm,
                                               sp.cars[other].vel,
                                               sp.cars[other].vel-vel)
                    end
                    sstar_margins[j] = clearances[j] - sstar
                end
            end

            margin, lane = findmax(sstar_margins)

            if margin > 0.0
                if vel > sp.cars[1].vel
                    # at back
                    push!(sp.cars, CarState(0.0, lane, vel, 0.0, behavior, next_id))
                else
                    push!(sp.cars, CarState(pp.lane_length, lane, vel, 0.0, behavior, next_id))
                end
            end
        end

        if sp.x > mdp.dmodel.max_dist
            sp.terminal = Nullable{Any}(:distance)
        elseif mdp.dmodel.lane_terminate && sp.cars[1].y == mdp.rmodel.target_lane
            sp.terminal = Nullable{Any}(:lane)
        end

        @assert sp.cars[1].x == s.cars[1].x # ego should not move
        @if_debug begin
            if any(any(isnan.(c.behavior.p_idm)) for c in sp.cars)
                warn("NaN in idm.")
            end
        end
    catch ex
        if mdp.throw
            rethrow(ex)
        else
            sp = deepcopy(s)
            sp.terminal = Nullable{Any}(ex)
            return sp
        end
    end

    return sp
end

function reward(mdp::NoCrashProblem{NoCrashRewardModel}, s::MLState, ::MLAction, sp::MLState)
    r = 0.0
    if sp.cars[1].y == mdp.rmodel.target_lane
        r += mdp.rmodel.reward_in_target_lane
    end
    nb_brakes = detect_braking(mdp, s, sp)
    r -= mdp.rmodel.cost_dangerous_brake*nb_brakes
    return r
end

"""
Return the number of braking actions that occured during this state transition
"""
function detect_braking(mdp::NoCrashProblem, s::MLState, sp::MLState, threshold::Float64)
    nb_brakes = 0
    nb_leaving = 0
    dt = mdp.dmodel.phys_param.dt
    for (i,c) in enumerate(s.cars)
        if length(sp.cars) >= i-nb_leaving
            cp = sp.cars[i-nb_leaving]
        else
            break
        end
        if cp.id != c.id
            nb_leaving += 1
            continue
        else
            acc = (cp.vel-c.vel)/dt
            if acc < -threshold
                nb_brakes += 1
            end
        end
    end
    # @assert nb_leaving <= 5 # sanity check - can remove this if it is violated as long as it doesn't happen all the time
    return nb_brakes
end

detect_braking(mdp::NoCrashProblem, s::MLState, sp::MLState) = detect_braking(mdp, s, sp, mdp.rmodel.brake_penalty_thresh)

"""
Return the ids of cars that brake during this state transition
"""
function braking_ids(mdp::NoCrashProblem, s::MLState, sp::MLState, threshold=mdp.rmodel.brake_penalty_thresh)
    braking = Int[]
    nb_leaving = 0
    dt = mdp.dmodel.phys_param.dt
    for (i,c) in enumerate(s.cars)
        if length(sp.cars) >= i-nb_leaving
            cp = sp.cars[i-nb_leaving]
        else
            break
        end
        if cp.id != c.id
            nb_leaving += 1
            continue
        else
            acc = (cp.vel-c.vel)/dt
            if acc < -threshold
                push!(braking, c.id)
            end
        end
    end
    # @assert nb_leaving <= 5 # sanity check - can remove this if it is violated as long as it doesn't happen all the time
    return braking
end

"""
Return the maximum braking (a positive number) by any car during the transition between s an sp
"""
function max_braking(mdp::NoCrashProblem, s::MLState, sp::MLState)
    nb_leaving = 0
    dt = mdp.dmodel.phys_param.dt
    min_acc = 0.0
    for (i,c) in enumerate(s.cars)
        if length(sp.cars) >= i-nb_leaving
            cp = sp.cars[i-nb_leaving]
        else
            break
        end
        if cp.id != c.id
            nb_leaving += 1
            continue
        else
            acc = (cp.vel-c.vel)/dt
            if acc < min_acc
                min_acc = acc
            end
        end
    end
    # @assert nb_leaving <= 5 # sanity check - can remove this if it is violated as long as it doesn't happen all the time
    return -min_acc
end

"""
Assign behaviors to a given physical state.
"""
function initial_state(mdp::NoCrashProblem, ps::MLPhysicalState, rng::AbstractRNG)
    s = MLState(ps, Vector{CarState}(length(ps.cars)))
    s.cars[1] = CarState(ps.cars[1], NORMAL)
    for i in 2:length(s.cars)
        behavior = rand(rng, mdp.dmodel.behaviors)
        s.cars[i] = CarState(ps.cars[i], behavior)
    end
    s.x = 0.0
    s.t = 0.0
    return s
end

function initial_state(p::NoCrashProblem, rng::AbstractRNG=Base.GLOBAL_RNG)
    @if_debug println("debugging")
    mdp = NoCrashMDP{typeof(p.rmodel), typeof(p.dmodel.behaviors)}(p.dmodel, p.rmodel, p.discount, p.throw) # make sure an MDP
    return relaxed_initial_state(mdp, 200, rng)
end

function generate_o(mdp::NoCrashProblem, s::MLState, a::MLAction, sp::MLState)
    return MLObs(sp)
end

function generate_sor(pomdp::NoCrashPOMDP, s::MLState, a::MLAction, rng::AbstractRNG)
    sp, r = generate_sr(pomdp, s, a, rng)
    o = generate_o(pomdp, s, a, sp)
    return sp, o, r
end

@if_debug function generate_sori(pomdp::NoCrashPOMDP, s::MLState, a::MLAction, rng::AbstractRNG)
    rngcp = copy(rng)
    sp, o, r = generate_sor(pomdp, s, a, rng)
    return sp, o, r, Dict(:rng=>rngcp)
end

@if_debug function generate_sri(mdp::NoCrashProblem, s, a, rng::AbstractRNG)
    rngcp = copy(rng)
    sp, r = generate_sr(mdp, s, a, rng)
    return sp, r, Dict(:rng=>rngcp)
end

discount(mdp::Union{MLMDP,MLPOMDP}) = mdp.discount
isterminal(mdp::Union{MLMDP,MLPOMDP},s::MLState) = !isnull(s.terminal)


text:
## basic task functions and TLS

show(io::IO, t::Task) = print(io, "Task ($(t.state)) @0x$(hex(unsigned(pointer_from_objref(t)), WORD_SIZE>>2))")

macro task(ex)
    :(Task(()->$(esc(ex))))
end

current_task() = ccall(:jl_get_current_task, Any, ())::Task
istaskdone(t::Task) = ((t.state == :done) | (t.state == :failed))
istaskstarted(t::Task) = isdefined(t, :last)

# yield to a task, throwing an exception in it
function throwto(t::Task, exc)
    t.exception = exc
    yieldto(t)
end

function task_local_storage()
    t = current_task()
    if is(t.storage, nothing)
        t.storage = ObjectIdDict()
    end
    (t.storage)::ObjectIdDict
end
task_local_storage(key) = task_local_storage()[key]
task_local_storage(key, val) = (task_local_storage()[key] = val)

function task_local_storage(body::Function, key, val)
    tls = task_local_storage()
    hadkey = haskey(tls,key)
    old = get(tls,key,nothing)
    tls[key] = val
    try body()
    finally
        hadkey ? (tls[key] = old) : delete!(tls,key)
    end
end

# NOTE: you can only wait for scheduled tasks
function wait(t::Task)
    if !istaskdone(t)
        if is(t.donenotify, nothing)
            t.donenotify = Condition()
        end
    end
    while !istaskdone(t)
        wait(t.donenotify)
    end
    if t.state == :failed
        throw(t.exception)
    end
    return t.result
end

# runtime system hook called when a task finishes
function task_done_hook(t::Task)
    err = (t.state == :failed)
    result = t.result
    nexttask = t.last

    q = t.consumers

    #### un-optimized version
    #isa(q,Condition) && notify(q, result, error=err)
    if isa(q,Task)
        nexttask = q
        nexttask.state = :runnable
    elseif isa(q,Condition) && !isempty(q.waitq)
        notify(q, result, error=err)
    end

    t.consumers = nothing

    isa(t.donenotify,Condition) && notify(t.donenotify, result, error=err)

    if nexttask.state == :runnable
        if err
            nexttask.exception = result
        end
        yieldto(nexttask, result)
    else
        wait()
    end
end

## produce, consume, and task iteration

function produce(v)
    #### un-optimized version
    #q = current_task().consumers
    #t = shift!(q.waitq)
    #empty = isempty(q.waitq)
    ct = current_task()
    local empty, t, q
    while true
        q = ct.consumers
        if isa(q,Task)
            t = q
            ct.consumers = nothing
            empty = true
            break
        elseif isa(q,Condition) && !isempty(q.waitq)
            t = shift!(q.waitq)
            empty = isempty(q.waitq)
            break
        end
        wait()
    end

    t.state = :runnable
    if empty
        if isempty(Workqueue)
            yieldto(t, v)
        else
            schedule_and_wait(t, v)
        end
        while true
            # wait until there are more consumers
            q = ct.consumers
            if isa(q,Task)
                return q.result
            elseif isa(q,Condition) && !isempty(q.waitq)
                return q.waitq[1].result
            end
            wait()
        end
    else
        schedule(t, v)
        return q.waitq[1].result
    end
end
produce(v...) = produce(v)

function consume(P::Task, values...)
    if istaskdone(P)
        return wait(P)
    end

    ct = current_task()
    ct.result = length(values)==1 ? values[1] : values

    #### un-optimized version
    #if P.consumers === nothing
    #    P.consumers = Condition()
    #end
    #push!(P.consumers.waitq, ct)
    # optimized version that avoids the queue for 1 consumer
    if P.consumers === nothing || (isa(P.consumers,Condition)&&isempty(P.consumers.waitq))
        P.consumers = ct
    else
        if isa(P.consumers, Task)
            t = P.consumers
            P.consumers = Condition()
            push!(P.consumers.waitq, t)
        end
        push!(P.consumers.waitq, ct)
    end
    ct.state = :waiting

    schedule_and_wait(P)
end

start(t::Task) = nothing
function done(t::Task, val)
    t.result = consume(t)
    istaskdone(t)
end
next(t::Task, val) = (t.result, nothing)

isempty(::Task) = error("isempty not defined for Tasks")

## condition variables

type Condition
    waitq::Vector{Any}

    Condition() = new({})
end

function wait(c::Condition)
    ct = current_task()

    ct.state = :waiting
    push!(c.waitq, ct)

    try
        return wait()
    catch
        filter!(x->x!==ct, c.waitq)
        if ct.state == :waiting
            ct.state = :runnable
        end
        rethrow()
    end
end

function notify(c::Condition, arg::ANY=nothing; all=true, error=false)
    if all
        for t in c.waitq
            schedule(t, arg, error=error)
        end
        empty!(c.waitq)
    elseif !isempty(c.waitq)
        t = shift!(c.waitq)
        schedule(t, arg, error=error)
    end
    nothing
end

notify1(c::Condition, arg=nothing) = notify(c, arg, all=false)

notify_error(c::Condition, err) = notify(c, err, error=true)
notify1_error(c::Condition, err) = notify(c, err, error=true, all=false)

## scheduler and work queue

global const Workqueue = Any[]

function enq_work(t::Task)
    ccall(:uv_stop,Void,(Ptr{Void},),eventloop())
    push!(Workqueue, t)
    t.state = :queued
    t
end

# schedule an expression to run asynchronously, with minimal ceremony
macro schedule(expr)
    expr = localize_vars(:(()->($expr)), false)
    :(enq_work(Task($(esc(expr)))))
end

schedule(t::Task) = enq_work(t)

function schedule(t::Task, arg; error=false)
    # schedule a task to be (re)started with the given value or exception
    if error
        t.exception = arg
    else
        t.result = arg
    end
    enq_work(t)
end

# fast version of schedule(t,v);wait()
function schedule_and_wait(t, v=nothing)
    if isempty(Workqueue)
        if t.state == :runnable
            return yieldto(t, v)
        end
    else
        if t.state == :runnable
            t.result = v
            push!(Workqueue, t)
            t.state = :queued
        end
    end
    wait()
end

yield() = (enq_work(current_task()); wait())

function wait()
    while true
        if isempty(Workqueue)
            c = process_events(true)
            if c==0 && eventloop()!=C_NULL && isempty(Workqueue)
                # if there are no active handles and no runnable tasks, just
                # wait for signals.
                pause()
            end
        else
            t = shift!(Workqueue)
            arg = t.result
            t.result = nothing
            t.state = :runnable
            result = yieldto(t, arg)
            process_events(false)
            # return when we come out of the queue
            return result
        end
    end
    assert(false)
end

function pause()
    @unix_only    ccall(:pause, Void, ())
    @windows_only ccall(:Sleep,stdcall, Void, (Uint32,), 0xffffffff)
end

## dynamically-scoped waiting for multiple items

sync_begin() = task_local_storage(:SPAWNS, ({}, get(task_local_storage(), :SPAWNS, ())))

function sync_end()
    spawns = get(task_local_storage(), :SPAWNS, ())
    if is(spawns,())
        error("sync_end() without sync_begin()")
    end
    refs = spawns[1]
    task_local_storage(:SPAWNS, spawns[2])
    for r in refs
        wait(r)
    end
end

macro sync(block)
    quote
        sync_begin()
        v = $(esc(block))
        sync_end()
        v
    end
end

function sync_add(r)
    spawns = get(task_local_storage(), :SPAWNS, ())
    if !is(spawns,())
        push!(spawns[1], r)
    end
    r
end

function async_run_thunk(thunk)
    t = Task(thunk)
    sync_add(t)
    enq_work(t)
    t
end

macro async(expr)
    expr = localize_vars(:(()->($expr)), false)
    :(async_run_thunk($(esc(expr))))
end

ext: jl

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia


text:
using CVRPLIB
using Test

@testset "CVRPLIB.jl" begin
    # Write your tests here.
    instances = [
        "E-n13-k4", "E-n101-k14", "X-n242-k48", "A-n80-k10", "F-n135-k7",
        "M-n200-k17", "P-n50-k8", "CMT11", "tai150b", "Golden_10",
        "Li_28", "X-n157-k13", "X-n502-k39"
        # , "Brussels2", "Ghent2" (out of memory in CI test for Windows)
    ]
    for inst in instances
        cvrp, vrp_file, sol_file = readCVRPLIB(inst)
        @testset "$inst" begin
            @test isfile(vrp_file)
            @test isfile(sol_file)
            cvrp2 = readCVRP(vrp_file)
            @test size(cvrp2.weights) == size(cvrp.weights)
        end
    end

    @testset "duration constraints" begin
        cvrp, _, _ = readCVRPLIB("CMT6")
        @test cvrp.distance == 200.0
        @test cvrp.service_time == 10.0

        cvrp, _, _ = readCVRPLIB("CMT7")
        @test cvrp.distance == 160.0
        @test cvrp.service_time == 10.0

    end

    @testset "add_dummy" begin
        @testset "P-n50-k8" begin
            n = 50
            cvrp1, _, _ = readCVRPLIB("P-n50-k8")
            @test size(cvrp1.weights) == (n, n)
            @test cvrp1.dimension == n
            @test length(cvrp1.demand) == n
            @test size(cvrp1.coordinates) == (n, 2)
            @test cvrp1.dummy == cvrp1.depot

            cvrp2, _, _ = readCVRPLIB("P-n50-k8", add_dummy=true)
            @test size(cvrp2.weights) == (n + 1, n + 1)
            @test cvrp2.dimension == n
            @test length(cvrp2.demand) == n + 1
            @test size(cvrp2.coordinates) == (n + 1, 2)
            @test cvrp2.dummy == n + 1
        end

        @testset "E-n13-k4" begin
            n = 13
            cvrp1, _, _ = readCVRPLIB("E-n13-k4")
            @test size(cvrp1.weights) == (n, n)
            @test cvrp1.dimension == n
            @test length(cvrp1.demand) == n
            @test cvrp1.dummy == cvrp1.depot

            cvrp2, _, _ = readCVRPLIB("E-n13-k4", add_dummy=true)
            @test size(cvrp2.weights) == (n + 1, n + 1)
            @test cvrp2.dimension == n
            @test length(cvrp2.demand) == n + 1
            @test cvrp2.dummy == n + 1
        end
    end

    wrong_instances = [
        "Daejeon", "E-n999-k999"
    ]
    for inst in wrong_instances
        @test_broken readCVRPLIB(inst)
    end

    write_instances = [
        "P-n16-k8", "E-n13-k4", "CMT6", "CMT7"
    ]
    @testset "write_CVRP" begin
        for inst in write_instances
            @testset "$inst" begin
                cvrp, vrp_file, sol_file = readCVRPLIB(inst)
                name, cvrp_file = write_cvrp(cvrp)
                @test isfile(cvrp_file)

                cvrp2 = readCVRP(cvrp_file)
                @test cvrp.dimension == cvrp2.dimension
                @test cvrp.capacity == cvrp2.capacity
                @test cvrp.distance == cvrp2.distance
                @test cvrp.service_time == cvrp2.service_time

                delete_cvrp(cvrp_file)
                @test !isfile(cvrp_file)
            end
        end
    end
    @testset "write_CVRP: add_dummy" begin
        for inst in write_instances
            @testset "$inst" begin
                cvrp, vrp_file, sol_file = readCVRPLIB(inst, add_dummy=true)
                name, cvrp_file = write_cvrp(cvrp)
                @test isfile(cvrp_file)

                cvrp2 = readCVRP(cvrp_file)
                @test cvrp.dimension == cvrp2.dimension
                @test cvrp.capacity == cvrp2.capacity
                @test cvrp.distance == cvrp2.distance
                @test cvrp.service_time == cvrp2.service_time

                delete_cvrp(cvrp_file)
                @test !isfile(cvrp_file)
            end
        end
    end
end


text:
using Test
using LinearAlgebra, StaticArrays

import ClimaCore:
    ClimaCore,
    slab,
    Spaces,
    Domains,
    Meshes,
    Geometry,
    Topologies,
    Spaces,
    Fields,
    Operators
import ClimaCore.Geometry: ⊗

import Logging
import TerminalLoggers
Logging.global_logger(TerminalLoggers.TerminalLogger())

# set up function space

function hvspace_3D(
    xlim = (-π, π),
    ylim = (-π, π),
    zlim = (0, 4π),
    xelem = 4,
    yelem = 4,
    zelem = 16,
    npoly = 4,
)
    FT = Float64

    xdomain = Domains.IntervalDomain(
        Geometry.XPoint{FT}(xlim[1]),
        Geometry.XPoint{FT}(xlim[2]),
        periodic = true,
    )
    ydomain = Domains.IntervalDomain(
        Geometry.YPoint{FT}(ylim[1]),
        Geometry.YPoint{FT}(ylim[2]),
        periodic = true,
    )

    horzdomain = Domains.RectangleDomain(xdomain, ydomain)
    horzmesh = Meshes.RectilinearMesh(horzdomain, xelem, yelem)
    horztopology = Topologies.Topology2D(horzmesh)

    zdomain = Domains.IntervalDomain(
        Geometry.ZPoint{FT}(zlim[1]),
        Geometry.ZPoint{FT}(zlim[2]);
        boundary_names = (:bottom, :top),
    )
    vertmesh = Meshes.IntervalMesh(zdomain, nelems = zelem)
    vert_center_space = Spaces.CenterFiniteDifferenceSpace(vertmesh)

    quad = Spaces.Quadratures.GLL{npoly + 1}()
    horzspace = Spaces.SpectralElementSpace2D(horztopology, quad)

    hv_center_space =
        Spaces.ExtrudedFiniteDifferenceSpace(horzspace, vert_center_space)
    hv_face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(hv_center_space)
    return (hv_center_space, hv_face_space)
end

# set up 3D domain - doubly periodic box
hv_center_space, hv_face_space = hvspace_3D((-500, 500), (-500, 500), (0, 1000))

const MSLP = 1e5 # mean sea level pressure
const grav = 9.8 # gravitational constant
const R_d = 287.058 # R dry (gas constant / mol mass dry air)
const γ = 1.4 # heat capacity ratio
const C_p = R_d * γ / (γ - 1) # heat capacity at constant pressure
const C_v = R_d / (γ - 1) # heat capacity at constant volume
const R_m = R_d # moist R, assumed to be dry

function pressure(ρθ)
    if ρθ >= 0
        return MSLP * (R_d * ρθ / MSLP)^γ
    else
        return NaN
    end
end

Φ(z) = grav * z

# Reference: https://journals.ametsoc.org/view/journals/mwre/140/4/mwr-d-10-05073.1.xml, Section 5a
function init_dry_rising_bubble_3d(x, y, z)
    x_c = 0.0
    y_c = 0.0
    z_c = 350.0
    r_c = 250.0
    θ_b = 300.0
    θ_c = 0.5
    cp_d = C_p
    cv_d = C_v
    p_0 = MSLP
    g = grav

    # auxiliary quantities
    r = sqrt((x - x_c)^2 + (z - z_c)^2)
    θ_p = r < r_c ? 0.5 * θ_c * (1.0 + cospi(r / r_c)) : 0.0 # potential temperature perturbation

    θ = θ_b + θ_p # potential temperature
    π_exn = 1.0 - g * z / cp_d / θ # exner function
    T = π_exn * θ # temperature
    p = p_0 * π_exn^(cp_d / R_d) # pressure
    ρ = p / R_d / T # density
    ρθ = ρ * θ # potential temperature density

    # Horizontal momentum defined on cell centers
    return (ρ = ρ, ρθ = ρθ, ρuₕ = ρ * Geometry.UVVector(0.0, 0.0))
end

# initial conditions
coords = Fields.coordinate_field(hv_center_space)
face_coords = Fields.coordinate_field(hv_face_space)

Yc = map(coords) do coord
    bubble = init_dry_rising_bubble_3d(coord.x, coord.y, coord.z)
    bubble
end

# Vertical momentum defined on cell faces
ρw = map(face_coords) do coord
    Geometry.WVector(0.0)
end;

Y = Fields.FieldVector(Yc = Yc, ρw = ρw)

function energy(Yc, ρu, z)
    ρ = Yc.ρ
    ρθ = Yc.ρθ
    u = ρu / ρ
    kinetic = ρ * norm(u)^2 / 2
    potential = z * grav * ρ
    internal = C_v * pressure(ρθ) / R_d
    return kinetic + potential + internal
end
function combine_momentum(ρuₕ, ρw)
    Geometry.transform(Geometry.UVWAxis(), ρuₕ) +
    Geometry.transform(Geometry.UVWAxis(), ρw)
end
function center_momentum(Y)
    If2c = Operators.InterpolateF2C()
    combine_momentum.(Y.Yc.ρuₕ, If2c.(Y.ρw))
end
function total_energy(Y)
    ρ = Y.Yc.ρ
    ρu = center_momentum(Y)
    ρθ = Y.Yc.ρθ
    z = Fields.coordinate_field(axes(ρ)).z
    sum(energy.(Yc, ρu, z))
end

energy_0 = total_energy(Y)
mass_0 = sum(Yc.ρ) # Computes ∫ρ∂Ω such that quadrature weighting is accounted for.
theta_0 = sum(Yc.ρθ)

function rhs!(dY, Y, _, t)
    ρw = Y.ρw
    Yc = Y.Yc
    dYc = dY.Yc
    dρw = dY.ρw
    ρ = Yc.ρ
    ρuₕ = Yc.ρuₕ
    ρθ = Yc.ρθ
    dρθ = dYc.ρθ
    dρuₕ = dYc.ρuₕ
    dρ = dYc.ρ

    # spectral horizontal operators
    hdiv = Operators.Divergence()
    hgrad = Operators.Gradient()
    hwdiv = Operators.WeakDivergence()
    hwgrad = Operators.WeakGradient()

    # vertical FD operators with BC's
    vdivf2c = Operators.DivergenceF2C(
        bottom = Operators.SetValue(Geometry.WVector(0.0)),
        top = Operators.SetValue(Geometry.WVector(0.0)),
    )
    vvdivc2f = Operators.DivergenceC2F(
        bottom = Operators.SetDivergence(Geometry.WVector(0.0)),
        top = Operators.SetDivergence(Geometry.WVector(0.0)),
    )
    uvdivf2c = Operators.DivergenceF2C(
        bottom = Operators.SetValue(
            Geometry.WVector(0.0) ⊗ Geometry.UVVector(0.0, 0.0),
        ),
        top = Operators.SetValue(
            Geometry.WVector(0.0) ⊗ Geometry.UVVector(0.0, 0.0),
        ),
    )
    If = Operators.InterpolateC2F(
        bottom = Operators.Extrapolate(),
        top = Operators.Extrapolate(),
    )
    Ic = Operators.InterpolateF2C()
    ∂ = Operators.DivergenceF2C(
        bottom = Operators.SetValue(Geometry.WVector(0.0)),
        top = Operators.SetValue(Geometry.WVector(0.0)),
    )
    ∂f = Operators.GradientC2F()
    ∂c = Operators.GradientF2C()
    B = Operators.SetBoundaryOperator(
        bottom = Operators.SetValue(Geometry.WVector(0.0)),
        top = Operators.SetValue(Geometry.WVector(0.0)),
    )

    fcc = Operators.FluxCorrectionC2C(
        bottom = Operators.Extrapolate(),
        top = Operators.Extrapolate(),
    )
    fcf = Operators.FluxCorrectionF2F(
        bottom = Operators.Extrapolate(),
        top = Operators.Extrapolate(),
    )

    uₕ = @. ρuₕ / ρ
    w = @. ρw / If(ρ)
    wc = @. Ic(ρw) / ρ
    p = @. pressure(ρθ)
    θ = @. ρθ / ρ
    Yfρ = @. If(ρ)

    ### HYPERVISCOSITY
    # 1) compute hyperviscosity coefficients
    @. dρθ = hdiv(hgrad(θ))
    @. dρuₕ = hdiv(hgrad(uₕ))
    @. dρw = hdiv(hgrad(w))
    Spaces.weighted_dss!(dYc)
    Spaces.weighted_dss!(dρw)

    κ₄ = 100.0 # m^4/s
    @. dρθ = -κ₄ * hdiv(ρ * hgrad(dρθ))
    @. dρuₕ = -κ₄ * hdiv(ρ * hgrad(dρuₕ))
    @. dρw = -κ₄ * hdiv(Yfρ * hgrad(dρw))

    # density
    @. dρ = -∂(ρw)
    @. dρ -= hdiv(ρuₕ)

    # potential temperature
    @. dρθ += -(∂(ρw * If(ρθ / ρ)))
    @. dρθ -= hdiv(uₕ * ρθ)

    # Horizontal momentum
    @. dρuₕ += -uvdivf2c(ρw ⊗ If(uₕ))
    Ih = Ref(
        Geometry.Axis2Tensor(
            (Geometry.UVAxis(), Geometry.UVAxis()),
            @SMatrix [1.0 0.0; 0.0 1.0]
        ),
    )
    @. dρuₕ -= hdiv(ρuₕ ⊗ uₕ + p * Ih)

    # vertical momentum
    z = coords.z
    @. dρw += B(
        Geometry.transform(Geometry.WAxis(), -(∂f(p)) - If(ρ) * ∂f(Φ(z))) -
        vvdivc2f(Ic(ρw ⊗ w)),
    )
    uₕf = @. If(ρuₕ / ρ) # requires boundary conditions
    @. dρw -= hdiv(uₕf ⊗ ρw)

    ### UPWIND FLUX CORRECTION
    upwind_correction = true
    if upwind_correction
        @. dρ += fcc(w, ρ)
        @. dρθ += fcc(w, ρθ)
        @. dρuₕ += fcc(w, ρuₕ)
        @. dρw += fcf(wc, ρw)
    end

    ### DIFFUSION
    κ₂ = 0.0 # m^2/s
    #  1a) horizontal div of horizontal grad of horiz momentun
    @. dρuₕ += hdiv(κ₂ * (ρ * hgrad(ρuₕ / ρ)))
    #  1b) vertical div of vertical grad of horiz momentun
    @. dρuₕ += uvdivf2c(κ₂ * (Yfρ * ∂f(ρuₕ / ρ)))

    #  1c) horizontal div of horizontal grad of vert momentum
    @. dρw += hdiv(κ₂ * (Yfρ * hgrad(ρw / Yfρ)))
    #  1d) vertical div of vertical grad of vert momentun
    @. dρw += vvdivc2f(κ₂ * (Yc.ρ * ∂c(ρw / Yfρ)))

    #  2a) horizontal div of horizontal grad of potential temperature
    @. dρθ += hdiv(κ₂ * (Yc.ρ * hgrad(ρθ / ρ)))
    #  2b) vertical div of vertial grad of potential temperature
    @. dρθ += ∂(κ₂ * (Yfρ * ∂f(ρθ / ρ)))

    Spaces.weighted_dss!(dYc)
    Spaces.weighted_dss!(dρw)
    return dY
end

dYdt = similar(Y);
rhs!(dYdt, Y, nothing, 0.0);

# run!
using OrdinaryDiffEq
Δt = 0.05
prob = ODEProblem(rhs!, Y, (0.0, 1.0))

haskey(ENV, "CI_PERF_SKIP_RUN") && exit() # for performance analysis

sol = @timev solve(
    prob,
    SSPRK33(),
    dt = Δt,
    saveat = 50.0,
    progress = true,
    progress_message = (dt, u, p, t) -> t,
);

ENV["GKSwstype"] = "nul"
import Plots
Plots.GRBackend()

dirname = "bubble_3d"
path = joinpath(@__DIR__, "output", dirname)
mkpath(path)

# post-processing
Es = [total_energy(u) for u in sol.u]
Mass = [sum(u.Yc.ρ) for u in sol.u]
Theta = [sum(u.Yc.ρθ) for u in sol.u]

Plots.png(
    Plots.plot((Es .- energy_0) ./ energy_0),
    joinpath(path, "energy.png"),
)
Plots.png(Plots.plot((Mass .- mass_0) ./ mass_0), joinpath(path, "mass.png"))
Plots.png(
    Plots.plot((Theta .- theta_0) ./ theta_0),
    joinpath(path, "rho_theta.png"),
)

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig(
    relpath(joinpath(path, "energy.png"), joinpath(@__DIR__, "../..")),
    "Total Energy",
)
linkfig(
    relpath(joinpath(path, "rho_theta.png"), joinpath(@__DIR__, "../..")),
    "Potential Temperature",
)
linkfig(
    relpath(joinpath(path, "mass.png"), joinpath(@__DIR__, "../..")),
    "Mass",
)


text:
# Copyright (c) 2014: SCS.jl contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

module SCS

import MathOptInterface
import Requires
import SCS_jll
import SparseArrays

function __init__()
    global indirect = SCS_jll.libscsindir
    global direct = SCS_jll.libscsdir
    Requires.@require(
        CUDA_jll = "e9e359dc-d701-5aa8-82ae-09bbf812ea83",
        begin
            import SCS_GPU_jll
            global gpuindirect = SCS_GPU_jll.libscsgpuindir
            push!(available_solvers, GpuIndirectSolver)
        end
    )
    return
end

include("c_wrapper.jl")
include("linear_solvers/direct.jl")
include("linear_solvers/indirect.jl")
include("linear_solvers/gpu_indirect.jl")
include("MOI_wrapper/MOI_wrapper.jl")

const available_solvers = [DirectSolver, IndirectSolver]

export scs_solve

end


text:
abstract type Executor end
struct SequentialEx{K} <: Executor
    kwargs::K
end
struct ThreadedEx{K} <: Executor
    kwargs::K
end
struct DistributedEx{K} <: Executor
    kwargs::K
end
struct PreferParallel{K} <: Executor
    kwargs::K
end

const _SIMD_OPT_DOCS = """
- `simd`: If `true` or `:ivdep`, enable SIMD using `Base.@simd`.  If
  `:ivdep`, use `@simd ivdep for ... end` variant.  Read Julia manual
  of `Base.@simd` to understand when it is appropriate to use this
  option.  For example, `simd = :ivdep` _must not_ be used with stateful
  transducer like [`Scan`](@ref).  If `false` (default), `Base.@simd` is not
  used.
"""

"""
    SequentialEx(; simd)

Sequential fold executor.  It can be passed to APIs from packages such as
Folds.jl and FLoops.jl to run the algorithm sequentially.

See also: [`foldxl`](@ref), [`ThreadedEx`](@ref) and [`DistributedEx`](@ref).

# Keyword Arguments
$_SIMD_OPT_DOCS

# Examples
```jldoctest
julia> using Folds

julia> Folds.sum(1:3, SequentialEx())
6
```
"""
SequentialEx

const _THREADED_EX_OPTS_DOCS = """
- `basesize::Integer = amount(reducible) ÷ nthreads()`: A size of
  chunk in `reducible` that is processed by each worker.  A smaller
  size may be required when:
    * computation time for processing each item fluctuates a lot
    * computation can be terminated by [`reduced`](@ref) or
      transducers using it, such as [`ReduceIf`](@ref)
- `stoppable::Bool`: [This option usually does not have to be set
  manually.]  The threaded fold executed in the "stoppable"
  mode used for optimizing reduction with [`reduced`](@ref) has a
  slight overhead if `reduced` is not used.  This mode can be disabled
  by passing `stoppable = false`.  It is usually automatically
  detected and set appropriately.  Note that this option is purely for
  optimization and does not affect the result value.
- `nestlevel::Union{Integer,Val}`: Specify how many inner `Cat`
  (flatten) transducers to be multi-threaded (using [`TCat`](@ref)).
  It must be a positive integer, `Val` of positive integer, or
  `Val(:inf)`.  `Val(:inf)` means to use multi-threading for all `Cat`
  transducers.  Note that `Cat` transducer should be statically known.
  That is to say, the fold implementation sees two `Cat`s in `... |> Map(f) |>
  Cat() |> Cat()` but only one `Cat` in `... |> Map(x -> f(x) |> Cat()) |>
  Cat()` even though they are semantically identical.
"""

"""
    ThreadedEx(; basesize, stoppable, nestlevel, simd)

Multi-threaded fold executor.  This is the default [^1] parallel executor used
by Folds.jl and FLoops.jl.

See also: [`foldxt`](@ref), [`SequentialEx`](@ref) and [`DistributedEx`](@ref).

[^1]: More specifically, Folds.jl and FLoops.jl uses [`PreferParallel`](@ref)
      which in turn defaults to `ThreadedEx`.

# Keyword Arguments
$_THREADED_EX_OPTS_DOCS
$_SIMD_OPT_DOCS

# Examples
```jldoctest
julia> using Folds

julia> Folds.sum(1:3, ThreadedEx(basesize = 1))
6
```
"""
ThreadedEx

const _DISTRIBUTED_EX_OPTS_DOCS = """
- `pool::AbstractWorkerPool`: Passed to `Distributed.remotecall`.
- `basesize::Integer = amount(array) ÷ nworkers()`: A size of chunk in
  `array` that is processed by each worker.  A smaller size may be
  required when computation time for processing each item can
  fluctuate a lot.
- `threads_basesize::Integer = basesize ÷ nthreads()`: A size of chunk
  in `array` that is processed by each task in each worker process.
  The default setting assumes that the number of threads used in all
  workers are the same.  For heterogeneous setup where each worker
  process has different number of threads, it may be required to use
  smaller `threads_basesize` _and_ `basesize` to get a good performance.
"""

"""
    DistributedEx(; pool, basesize, threads_basesize, simd)

Distributed fold executor.  It can be passed to APIs from packages such as
Folds.jl and FLoops.jl to run the algorithm sequentially.

See also: [`foldxd`](@ref), [`SequentialEx`](@ref) and [`ThreadedEx`](@ref).

# Keyword Arguments
$_DISTRIBUTED_EX_OPTS_DOCS
$_SIMD_OPT_DOCS

# Examples
```jldoctest
julia> using Folds

julia> Folds.sum(1:3, DistributedEx())
6
```
"""
DistributedEx

"""
    PreferParallel(; simd, basesize)

A "placeholder" executor that indicates preference to parallel execution.

This lets the input data collection decide preferred execution strategy
(e.g., `CUDAEx` for `CuArray` when FoldsCUDA.jl is available), **assuming
that the reducing function is associative**. The default executor is
[`ThreadedEx`](@ref). As an optional feature, some input data collections
support (e.g., `AbstractChannel`) automatically demoting the execution
strategy to [`SequentialEx`](@ref).  An error is thrown if the automatic
detection fails,
"""
PreferParallel

maybe_set_simd(exc::Executor, simd) =
    isdefined(exc.kwargs, :simd) ? exc : @set exc.kwargs = merge(exc.kwargs, (simd = simd,))

(::Type{Ex})(; kwargs...) where {Ex<:Executor} = Ex(values(kwargs))

transduce(xf, rf::RF, init, coll, exc::SequentialEx) where {RF} =
    transduce(xf, rf, init, coll; exc.kwargs...)
transduce(xf, rf, init, coll, exc::ThreadedEx) =
    transduce_assoc(xf, rf, init, coll; exc.kwargs...)
transduce(xf, rf, init, coll, exc::DistributedEx) =
    dtransduce(xf, rf, init, coll; exc.kwargs...)

fold(rf::RF, coll, exec::Executor; init = DefaultInit) where {RF} =
    unreduced(transduce(IdentityTransducer(), Completing(rf), init, coll, exec))
fold(rf::RF, coll; init = DefaultInit, kwargs...) where {RF} =
    fold(rf, coll, PreferParallel(; kwargs...); init = init)

function transduce(xf, rf, init, coll, exc::PreferParallel)
    xf0, coll0 = extract_transducer(coll)
    return transduce(xf ∘ xf0, rf, init, coll0, executor_for(coll0, exc))
end

fold(rf; kw...) = itr -> fold(rf, itr; kw...)

executor_from(::Type{Ex}, exc::PreferParallel) where {Ex<:Executor} = Ex(exc.kwargs)
executor_from(::Type{PreferParallel}, exc::PreferParallel) = ThreadedEx(exc.kwargs)
executor_from(::Type{SequentialEx}, exc::PreferParallel) =
    SequentialEx(; _seq_ex_kwargs(; exc.kwargs...)...)
_seq_ex_kwargs(; simd = Val(false), _...) = (simd = simd,)

executor_for(xs::Any, exc) = executor_from(executor_type(xs), exc)
# * `ThreadedEx` part sould be determined by the property of transducers (and,
#   ideally, reducing function).
# * We can extend the API of `executor_type` that returns a list of comptaible
#   executors rather than only one executor.

# `executor_type(::CustomContainer)` is the main (still hidden) API.
# TODO: should `executor_type` be renamed to `parallel_executor_type` if the
# default type is `PreferParallel`?
@inline executor_type(::Any) = PreferParallel

# Some "famous" containers known to be sequential-only:
@inline executor_type(::AbstractChannel) = SequentialEx
# @inline executor_type(::Distributed.RemoteChannel) = SequentialEx
@inline executor_type(::Iterators.Stateful) = SequentialEx

@inline executor_type(xs::Iterators.Pairs) =
    executor_promote_type(executor_type(keys(xs)), executor_type(values(xs)))

@inline executor_type(xs::Iterators.ProductIterator) =
    mapfoldlargs(executor_type, executor_promote_type, PreferParallel, xs.iterators...)

@inline executor_type(xs::_Zip) =
    mapfoldlargs(executor_type, executor_promote_type, PreferParallel, arguments(xs)...)

@inline executor_promote_type(::Type{T}, ::Type{T}) where {T} = T
@inline function executor_promote_type(::Type{T}, ::Type{S}) where {T,S}
    result =
        executor_promote_result(executor_promote_rule(T, S), executor_promote_rule(S, T))
    result === nothing && throw_promote_failure(T, S)
    return result
end

# `nothing` means promotion failure
@inline executor_promote_result(::Nothing, ::Nothing) = nothing
@inline executor_promote_result(::Type{T}, ::Type{S}) where {T,S} = nothing
@inline executor_promote_result(::Type{T}, ::Type{T}) where {T} = T
@inline executor_promote_result(::Nothing, ::Type{T}) where {T} = T
@inline executor_promote_result(::Type{T}, ::Nothing) where {T} = T

# Only the owner of the first argument should define the rule. Furthermore, a
# custom executor `CustomEx` should define
# `executor_promote_rule(::Type{SequentialEx}, ::Type{OtherEx}) = OtherEx` only
# if it knows how to handle all `xs` such that `executor_type(xs) = OtherEx`.
@inline executor_promote_rule(::Type{PreferParallel}, ::Type{T}) where {T} = T
@inline executor_promote_rule(::Type, ::Type) = nothing

@noinline throw_promote_failure(T, S) = error("failed to promote executors:\n$T\n$S")
# TODO: show information on which container "suggested" which executor?

function Base.show(io::IO, exc::Executor)
    T = typeof(exc.kwargs)
    if (
        fieldtype(typeof(exc), 1) !== typeof(exc.kwargs) ||
        any(i -> typeof(exc.kwargs[i]) != fieldtype(T, i), 1:nfields(exc.kwargs))
    )
        return invoke(show, Tuple{IO,Any}, io, exc)
    end
    Ex = getfield(parentmodule(typeof(exc)), nameof(typeof(exc)))
    print(io, Ex)
    if isempty(exc.kwargs)
        print(io, "()")
    elseif length(exc.kwargs) == 1
        print(io, '(', fieldname(T, 1), " = ", exc.kwargs[1], ')')
    else
        print(io, exc.kwargs)
    end
    return
end


text:
function ticker(id::String)
  res = HTTP.get(ENDPOINT * "products/$id/ticker")
  rawData = JSON.parse(String(res.body))
  DataFrame(rawData)
end


text:
using Printf

# Lots more info on meta-programming here:
# https://docs.julialang.org/en/v1/manual/metaprogramming/
# And more examples here:
# https://gist.github.com/MikeInnes/8299575

macro apply_expression(n, exp)
    quote
        # For 1 to n
        for i = 1:$(esc(n))
            # run the expression
            $(esc(exp))
        end
    end
end

macro do_while(exp)
    # Check that first item in expression is 'while'
    @assert exp.head == :while
    esc(quote
        $(exp.args[2])
        $exp
    end)
end

function macro_examples()
    println("--------------------------------------")
    println("Macro examples")

    println()

    println("We can call macros (note that macros must be preceeded by @)")
    @apply_expression(5, println("Hello world"))

    z = 0
    @do_while while z < 10
        z += 1
        println(z)
    end
end

max_stars_repo_name: James-P-D/JuliaDump

max_issues_repo_name: James-P-D/JuliaDump

max_forks_repo_name: James-P-D/JuliaDump


text:
# Petsc matrix-free implementation

"""
  Type for managing a Petsc matrix-free preconditioner.  Actual preconditioners
  can be built on top of this.  Because matrix-free preconditioners do not have
  a particular form, this type requires the preconditioner to implement the
  entire [`AbstractPC`](@ref) interface.  The benefit of using this type to
  build a preconditioner is that is automatically exposes the preconditioner
  to Petsc.

  The [`calcPC`](@ref) function the user defines must call [`setPCCtx`](@ref).

  **Public Fields**

   * none
"""
mutable struct PetscMatFreePC <: AbstractPetscMatFreePC
  pc::PC
  ctx  # Petsc PC ctx
  is_shared::Bool
  is_finalized::Bool

  nsetups::Int
  napplies::Int
  ntapplies::Int

  # MPI stuff
  comm::MPI.Comm
  myrank::Int
  commsize::Int
end

"""
  Outer constructor

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
"""
function PetscMatFreePC(mesh::AbstractMesh, sbp::AbstractOperator,
                    eqn::AbstractSolutionData, opts::Dict)

  pc = createPetscPC(mesh, sbp, eqn, opts)

  # register wrapper routines with Petsc
  fptr = cfunction(applyPC_wrapper, PetscErrorCode, (PC, PetscVec, PetscVec))
  PCShellSetApply(pc, fptr)

  fptr = cfunction(applyPCTranspose_wrapper, PetscErrorCode, (PC, PetscVec, PetscVec))
  PCShellSetApplyTranspose(pc, fptr)

  fptr = cfunction(calcPC_wrapper, PetscErrorCode, (PC,))
  PCShellSetSetUp(pc, fptr)

  ctx = ()
  is_shared = false
  is_finalized = false

  nsetups = 0
  napplies = 0
  ntapplies = 0

  comm = eqn.comm
  myrank = eqn.myrank
  commsize = eqn.commsize

  # only update the PC when explicitly requested
  PCSetReusePreconditioner(pc, PETSC_TRUE)

  return PetscMatFreePC(pc, ctx, is_shared, is_finalized, nsetups, napplies,
                        ntapplies, comm, myrank, commsize)
end

function free(pc::PetscMatFreePC)

  if !pc.is_finalized
    if pc.pc.pobj != C_NULL
      PetscDestroy(pc.pc)
      pc.pc.pobj = C_NULL
    end
  end

  pc.is_finalized = true

  return nothing
end

"""
  This function should be called by the users [`calcPC`](@ref) function.
  The arguments passed to this function are passed to [`applyPC`](@ref)
  during the next preconditioner application.
  See that function for a description of the arguments

  **Inputs**

   * pc:
   * mesh
   * sbp
   * eqn
   * opts
   * ctx_residual
   * t
"""
function setPCCtx(pc::AbstractPetscMatFreePC, mesh::AbstractMesh,
                  sbp::AbstractOperator,
                  eqn::AbstractSolutionData, opts::Dict, ctx_residual, t)
# users *must* call this function from within their calcPC function

  pc2 = getBasePC(pc)
  @assert typeof(pc2) <: PetscMatFreePC
  pc2.ctx = (mesh, sbp, eqn, opts, pc, ctx_residual, t)
  PCShellSetContext(pc2.pc, pointer_from_objref(pc2.ctx))

  return nothing
end

function calcPC_wrapper(_pc::PC)

  ctx_ptr = PCShellGetContext(_pc)
  @assert ctx_ptr != C_NULL
  ctx = unsafe_pointer_to_objref(ctx_ptr)
  checkPCCtx(ctx)
  mesh = ctx[1]
  sbp = ctx[2]
  eqn = ctx[3]
  opts = ctx[4]
  pc = ctx[5]  # the user PC
  ctx_residual = ctx[6]
  t = ctx[7]

  calcPC(pc, mesh, sbp, eqn, opts, ctx_residual, t)

  return PetscErrorCode(0)
end

# wrapper for Petsc
function applyPC_wrapper(_pc::PC, b::PetscVec, x::PetscVec)

  ctx = unsafe_pointer_to_objref(PCShellGetContext(_pc))
  checkPCCtx(ctx)
  mesh = ctx[1]
  sbp = ctx[2]
  eqn = ctx[3]
  opts = ctx[4]
  pc = ctx[5]  # the user PC
  ctx_residual = ctx[6]
  t = ctx[7]

  # get the local parts of the vectors
  btmp = VecGetArrayRead(b)
  xtmp = VecGetArray(x)

  applyPC(pc, mesh, sbp, eqn, opts, t, btmp, xtmp)

  VecRestoreArrayRead(b, btmp)
  VecRestoreArray(x, xtmp)

  return PetscErrorCode(0)
end

function applyPCTranspose_wrapper(_pc::PC, b::PetscVec, x::PetscVec)

  ctx_ptr = PCShellGetContext(_pc)
  @assert ctx_ptr != C_NULL
  ctx = unsafe_pointer_to_objref(ctx_ptr)
  checkPCCtx(ctx)
  mesh = ctx[1]
  sbp = ctx[2]
  eqn = ctx[3]
  opts = ctx[4]
  pc = ctx[5]  # the user PC
  ctx_residual = ctx[6]
  t = ctx[7]

  # get the local parts of the vectors
  btmp = VecGetArrayRead(b)
  xtmp = VecGetArray(x)

  applyPCTranspose(pc, mesh, sbp, eqn, opts, t, btmp, xtmp)

  VecRestoreArrayRead(b, btmp)
  VecRestoreArray(x, xtmp)

  return PetscErrorCode(0)
end

"""
  This function verifies the [`PetscMatFreePC`](@ref) ctx is valid, throwing
  an exception otherwise.

  Currently this only checks whether the members of the tuple have the right
  abstract supertypes

  **Inputs**

   * ctx: the ctx retrieved from the Petsc matrix-free preconditioner object
"""
function checkPCCtx(ctx)

  mesh = ctx[1]
  sbp = ctx[2]
  eqn = ctx[3]
  opts = ctx[4]
  pc = ctx[5]  # the user PC
#  pc2 = ctx[6]  # PetscMatFreePC
  ctx_residual = ctx[6]
  t = ctx[7]

  pc2 = getBasePC(pc)

  @assert typeof(mesh) <: AbstractMesh
  @assert typeof(sbp) <: AbstractOperator
  @assert typeof(eqn) <: AbstractSolutionData
  @assert typeof(opts) <: Dict
  @assert typeof(pc) <: AbstractPC
  @assert typeof(pc2) <: PetscMatFreePC
  @assert typeof(t) <: Number

  # verify equality with pc2.ctx ?

  return nothing
end

function getBasePC(pc::PetscMatFreePC)

  # this is the bottom of the recursion tree
  return pc
end


text:
module LinAlgTests

using ReverseDiff, ForwardDiff, Test, LinearAlgebra, Statistics

include(joinpath(dirname(@__FILE__), "../utils.jl"))

x, a, b = rand(3, 3), rand(3, 3), rand(3, 3)
tp = InstructionTape()

function test_arr2num(f, x, tp; ignore_tape_length = false)
    xt = track(copy(x), tp)
    y = f(x)

    # record
    yt = f(xt)
    if !ignore_tape_length
        @test length(tp) == 1
    end
    @test abs(yt - y) <= abs(y) * eps(typeof(y))

    # reverse
    ReverseDiff.seed!(yt)
    ReverseDiff.reverse_pass!(tp)
    test_approx(deriv(xt), ForwardDiff.gradient(f, x))

    # forward
    x2 = rand(eltype(x), size(x))
    ReverseDiff.value!(xt, x2)
    ReverseDiff.forward_pass!(tp)
    y = f(x2)
    @test abs(value(yt) - y) <= abs(y) * eps(typeof(y))
    ReverseDiff.value!(xt, x)

    empty!(tp)
end

function test_arr2arr(f, x, tp)
    xt = track(copy(x), tp)
    y = f(x)

    # record
    yt = f(xt)
    @test yt == y
    @test length(tp) == 1

    # reverse
    out = similar(y, (length(y), length(x)))
    ReverseDiff.seeded_reverse_pass!(out, yt, xt, tp)
    test_approx(out, ForwardDiff.jacobian(f, x))

    # forward
    x2 = rand(eltype(x), size(x))
    ReverseDiff.value!(xt, x2)
    ReverseDiff.forward_pass!(tp)
    @test value(yt) == f(x2)
    ReverseDiff.value!(xt, x)

    empty!(tp)
end

function test_arr2arr(f, a, b, tp)
    at, bt = track(copy(a), tp), track(copy(b), tp)
    c = f(a, b)

    ########################################

    # record
    ct = f(at, b)
    @test ct == c
    @test length(tp) == 1

    # reverse
    out = similar(c, (length(c), length(a)))
    ReverseDiff.seeded_reverse_pass!(out, ct, at, tp)
    test_approx(out, ForwardDiff.jacobian(x -> f(x, b), a))

    # forward
    a2 = rand(eltype(a), size(a))
    ReverseDiff.value!(at, a2)
    ReverseDiff.forward_pass!(tp)
    @test value(ct) == f(a2, b)

    ReverseDiff.value!(at, a)
    empty!(tp)

    ########################################

    # record
    ct = f(a, bt)
    @test ct == c
    @test length(tp) == 1

    # reverse
    out = similar(c, (length(c), length(b)))
    ReverseDiff.seeded_reverse_pass!(out, ct, bt, tp)
    test_approx(out, ForwardDiff.jacobian(x -> f(a, x), b))

    # forward
    b2 = rand(eltype(b), size(b))
    ReverseDiff.value!(bt, b2)
    ReverseDiff.forward_pass!(tp)
    @test value(ct) == f(a, b2)
    ReverseDiff.value!(bt, b)

    empty!(tp)

    ########################################

    # record
    ct = f(at, bt)
    @test ct == c
    @test length(tp) == 1

    # reverse
    out_a = similar(c, (length(c), length(a)))
    out_b = similar(c, (length(c), length(b)))
    ReverseDiff.seeded_reverse_pass!(out_a, ct, at, tp)
    ReverseDiff.seeded_reverse_pass!(out_b, ct, bt, tp)
    test_approx(out_a, ForwardDiff.jacobian(x -> f(x, b), a))
    test_approx(out_b, ForwardDiff.jacobian(x -> f(a, x), b))

    # forward
    a2, b2 = rand(eltype(a), size(a)), rand(eltype(b), size(b))
    ReverseDiff.value!(at, a2)
    ReverseDiff.value!(bt, b2)
    ReverseDiff.forward_pass!(tp)
    @test value(ct) == f(a2, b2)
    ReverseDiff.value!(at, a)
    ReverseDiff.value!(bt, b)

    empty!(tp)
end

function test_arr2arr_inplace(f!, f, c, a, b, tp)
    at, bt = track(copy(a), tp), track(copy(b), tp)
    f!(c, a, b)

    ########################################

    # record
    ct = track(c, eltype(c), ReverseDiff.NULL_TAPE)
    f!(ct, at, b)
    @test ct == c
    @test length(tp) == 1

    # reverse
    out = similar(c, (length(c), length(a)))
    ReverseDiff.seeded_reverse_pass!(out, ct, at, tp)
    test_approx(out, ForwardDiff.jacobian(x -> f(x, b), a))

    # forward
    a2 = rand(eltype(a), size(a))
    ReverseDiff.value!(at, a2)
    ReverseDiff.forward_pass!(tp)
    @test value(ct) == f(a2, b)
    ReverseDiff.value!(at, a)

    empty!(tp)

    ########################################

    # record
    ct = track(c, eltype(c), ReverseDiff.NULL_TAPE)
    f!(ct, a, bt)
    @test ct == c
    @test length(tp) == 1

    # reverse
    out = similar(c, (length(c), length(b)))
    ReverseDiff.seeded_reverse_pass!(out, ct, bt, tp)
    test_approx(out, ForwardDiff.jacobian(x -> f(a, x), b))

    # forward
    b2 = rand(eltype(b), size(b))
    ReverseDiff.value!(bt, b2)
    ReverseDiff.forward_pass!(tp)
    @test value(ct) == f(a, b2)
    ReverseDiff.value!(bt, b)

    empty!(tp)

    ########################################

    # record
    ct = track(c, eltype(c), ReverseDiff.NULL_TAPE)
    f!(ct, at, bt)
    @test ct == c
    @test length(tp) == 1

    # reverse
    out_a = similar(c, (length(c), length(a)))
    out_b = similar(c, (length(c), length(b)))
    ReverseDiff.seeded_reverse_pass!(out_a, ct, at, tp)
    ReverseDiff.seeded_reverse_pass!(out_b, ct, bt, tp)
    test_approx(out_a, ForwardDiff.jacobian(x -> f(x, b), a))
    test_approx(out_b, ForwardDiff.jacobian(x -> f(a, x), b))

    # forward
    a2, b2 = rand(eltype(a), size(a)), rand(eltype(b), size(b))
    ReverseDiff.value!(at, a2)
    ReverseDiff.value!(bt, b2)
    ReverseDiff.forward_pass!(tp)
    @test value(ct) == f(a2, b2)
    ReverseDiff.value!(at, a)
    ReverseDiff.value!(bt, b)

    empty!(tp)
end

for f in (
    sum,
    det,
    mean,
    y -> dot(vec(y), vec(y)),
    y -> vec(y)' * vec(y),
    y -> vec(y)' * ones(length(y)),
    y -> ones(length(y))' * vec(y),
)
    test_println("Array -> Number functions", f)
    test_arr2num(f, x, tp)
end

for f in (
    y -> vec(y)' * Matrix{Float64}(I, length(y), length(y)) * vec(y),
    y -> transpose(vec(y)) * vec(y),
    y -> transpose(vec(y)) * ones(length(y)),
    y -> transpose(ones(length(y))) * vec(y),
)
    test_println("Array -> Number functions", f)
    # TODO: transpose needs investigation; it is giving an unusually large tape length here
    test_arr2num(f, x, tp, ignore_tape_length=true)
end

for f in (-, inv)
    test_println("Array -> Array functions", f)
    test_arr2arr(f, x, tp)
end

for f in (x -> copy(transpose(x)), x -> copy(adjoint(x)))
    test_println("Array -> Array functions", f)
    test_arr2arr(f, x, tp)
end

for f in (+, -)
    test_println("(Array, Array) -> Array functions", f)
    test_arr2arr(f, a, b, tp)
end

test_println("*(A, B) functions", "*(a, b)")

test_arr2arr(*, a, b, tp)
test_arr2arr_inplace(mul!, *, x, a, b, tp)

for f in (transpose, adjoint)
    test_println("*(A, B) functions", string("*(", f, "(a), b)"))
    test_arr2arr(*, f(a), b, tp)
    test_arr2arr_inplace(mul!, *, x, f(a), b, tp)
    test_println("*(A, B) functions", string("*(a, ", f, "(b))"))
    test_arr2arr(*, a, f(b), tp)
    test_arr2arr_inplace(mul!, *, x, a, f(b), tp)
    test_println("*(A, B) functions", string("*(", f, "(a), ", f, "(b))"))
    test_arr2arr(*, f(a), f(b), tp)
    test_arr2arr_inplace(mul!, *, x, f(a), f(b), tp)
end

test_println("*(A, B) functions", "*(adjoint(a), transpose(b))")
test_arr2arr(*, adjoint(a), transpose(b), tp)
test_arr2arr_inplace(mul!, *, x, adjoint(a), transpose(b), tp)

test_println("*(A, B) functions", "*(transpose(a), adjoint(b))")
test_arr2arr(*, transpose(a), adjoint(b), tp)
test_arr2arr_inplace(mul!, *, x, transpose(a), adjoint(b), tp)

end # module


text:
include("./utils.jl")
include("../../phantoms/sperm_phantom.jl")
include("../snake_forward.jl")
include("../snake.jl")
include("../curve_utils.jl")
using IterTools
using LinearAlgebra

function count_parallel(projection, r; limit=30)
    #get the 'end points' of the projection, by getting the first and last value where value is greater than tail diameter, which is minimum
    projection_end1 = findfirst(p -> p > 2*r(0.0), projection)
    projection_end2 = findlast(p -> p > 2*r(0.0), projection)

    first = min(projection_end1, projection_end2)
    last = max(projection_end1, projection_end2)
    all_parallel = findall(p -> p <= 2.0*r(0.0), projection[first:last])
    neighbor_diff = all_parallel[2:end]-all_parallel[1:end-1]
    changes = count(c -> c > limit, neighbor_diff)+1
    return changes, all_parallel
end

images, tracks = get_sperm_phantom(301,grid_size=0.1)

r(s) = 1.0
detmin, detmax = -38.0, 38.0
bins = collect(detmin:0.125:detmax)

cwd = @__DIR__
savepath = normpath(joinpath(@__DIR__, "result"))
!isdir(savepath) && mkdir(savepath)

cd(savepath)
c = []
b = []
for (i,track) in enumerate(tracks)
    outline, normals = get_outline(track, r)
    projection = parallel_forward(outline, [π/2], bins)

    changes, a = count_parallel(projection[:,1], r, limit=30)
    curvature_changes, k, length_ok, changes, prime = could_be_sperm_tail(30.0, track, π/2, smoothness=0.1)
    plot(track[:,1], track[:,2], aspect_ratio=:equal, label=(@sprintf "%d_%f" changes maximum(abs.(prime))))
    push!(c, curvature_changes)
    push!(b, norm(prime))
    savefig(@sprintf "result_all_%d" i)
end

max_stars_repo_name: QPImaging/xfromProjections

max_issues_repo_name: QPImaging/xfromProjections

max_forks_repo_name: QPImaging/xfromProjections


text:
module MPSforQuantum
    export MPS, restore, OneQubitGate, TwoQubitGate, dstack, ddstack, init_rand_MPS
    export CX, inner_product, expectation, SVD_R, SVD_L, iterative_ground_state_search, j1j2_2D_Hamiltonian_MPO
    using LinearAlgebra
    using TensorOperations

    function SVD(C::Array{ComplexF64,2}, eps::Float64)
        #eps = 1e-2
        A = svd(C)
        filter!((x) -> x > eps, A.S)
        l = length(A.S)
        return A.U[:, 1:l], diagm(A.S), A.Vt[1:l, :]
        #return A.U[:, 1:l], diagm(A.S), transpose(A.Vt)[1:l, :]
    end

    function SVD(C::Array{ComplexF64,2}, D::Int64)
        A = svd(C)
        if length(A.S) > D
            return A.U[:, 1:D], diagm(A.S[1:D]), A.Vt[1:D, :]
            #return A.U[:, 1:D], diagm(A.S[1:D]), transpose(A.Vt)[1:D, :]
        else
            return A.U[:, :], diagm(A.S), A.Vt[:, :]
            #return A.U[:, :], diagm(A.S), transpose(A.Vt)[:, :]
        end
    end

    function SVD_L(C::Array{ComplexF64,2}, eps::Float64)
        U, S, Vt = SVD(C, eps)
        return U, S * Vt
    end

    function SVD_L(C::Array{ComplexF64,2}, D::Int64)
        U, S, Vt = SVD(C, D)
        return U, S * Vt
    end

    function SVD_R(C::Array{ComplexF64,2}, eps::Float64)
        U, S, Vt = SVD(C, eps)
        return U * S, Vt
    end

    function SVD_R(C::Array{ComplexF64,2}, D::Int64)
        U, S, Vt = SVD(C, D)
        return U * S, Vt
    end

    function MPS(C::Array{ComplexF64,1}, param, normalize::Char='l')
        d = 2
        N = Int64(log2(size(C)[1]))
        arrs = []
        r = 1
        if normalize == 'l'
            for i = 1:N-1
                C = reshape(C, d * r, d^(N - i))
                tmp_A, C = SVD_L(C, param)
                r = size(tmp_A, 2)
                col = convert(Int64, size(tmp_A, 1) / 2)
                push!(arrs, [tmp_A[1:col, :], tmp_A[(col+1):col*2, :]])
            end

            Ct = transpose(C)
            col2 = convert(Int64, size(C, 2) / 2)
            push!(arrs, [C[:, 1:col2], C[:, (col2+1):col2*2]])
            return arrs
        elseif normalize == 'r'
            for i = 1:N-1
                #C = reshape(C, d^(N - i), d * r)
                l = convert(Int64, size(C, 1) / 2)
                C = cat(C[1:l, :], C[l+1:l*2, :], dims=2)
                C, tmp_B = SVD_R(C, param)
                r = size(tmp_B, 1)
                col = convert(Int64, size(tmp_B, 2) / 2)
                push!(arrs, [tmp_B[:, 1:col], tmp_B[:, (col+1):col*2]])
            end

            Ct = transpose(C)
            col2 = convert(Int64, size(C, 1) / 2)
            push!(arrs, [C[1:col2, :], C[(col2+1):col2*2, :]])
            arrs_ = []
            for i = 1:N
                push!(arrs_, arrs[N - (i-1)])
            end
            return arrs_
        end
    end

    function OneQubitGate(arrs::Array{Any, 1}, O::Array{Complex{Float64},2}, n::Int64)
        arrs_ = similar(arrs[n])
        arrs__ = copy(arrs)
        arrs_[1] = arrs[n][1] * O[1, 1] + arrs[n][2] * O[1, 2]
        arrs_[2] = arrs[n][1] * O[2, 1] + arrs[n][2] * O[2, 2]
        arrs__[n] = arrs_
        return arrs__
    end

    function swap_ct(arr::Array{Complex{Float64},2})
        swap = [1 0 0 0; 0 0 1 0; 0 1 0 0; 0 0 0 1]
        _arr = copy(arr)
        _arr[:, 2] = arr[:, 3]
        _arr[:, 3] = arr[:, 2]
        return swap * _arr
    end

    function reduce(mps::Array{Any, 1}, t1::Int64, t2::Int64, param::Number)
        eps = 1e-3
        arr1 = cat(mps[t1][1], mps[t1][2], dims = 1)
        arr2 = cat(mps[t2][1], mps[t2][2], dims = 2)
        arr = arr1 * arr2

        # U, S, V = svd(arr)
        # filter!((x) -> x > eps, S)
        # sdim = length(S)
        U, S, V = SVD(arr, param)

        #U1 = U[:, 1:sdim]
        U1 = U
        U1 = [U1[1:convert(Int, size(U1)[1]/2), :], U1[convert(Int, size(U1)[1]/2)+1:size(U1)[1], :]]
        #U2 = diagm(S) * transpose(V)[1:sdim, :]
        U2 = S * V
        U2 = [U2[:, 1:convert(Int, size(U2)[2]/2)], U2[:, convert(Int, size(U2)[2]/2)+1:size(U2)[2]]]

        mps[t1] = U1
        mps[t2] = U2
        return mps
    end

    function TwoQubitGate(mps::Array{Any, 1}, arr::Array{Complex{Float64},2}, c::Int64, t::Int64, param::Number)
        eps = 1e-6 # constant cutoff for TwoQubitGate decomposition

        if c == t + 1
            arr = swap_ct(arr)
            (c, t) = (t, c)
        end

        arr_00 = transpose(reshape(arr[:, 1], 2, 2))
        arr_01 = transpose(reshape(arr[:, 2], 2, 2))
        arr_10 = transpose(reshape(arr[:, 3], 2, 2))
        arr_11 = transpose(reshape(arr[:, 4], 2, 2))

        tensor1 = cat(arr_00, arr_01, dims = 2)
        tensor2 = cat(arr_10, arr_11, dims = 2)
        tensor = cat(tensor1, tensor2, dims = 1)

        U, S, V = svd(tensor)
        filter!((x) -> x > eps, S)
        sdim = length(S)
        U1 = U[:, 1:sdim]
        U1 = [U1[1:convert(Int, size(U1)[1]/2), :], U1[convert(Int, size(U1)[1]/2)+1:size(U1)[1], :]]
        U2 = diagm(S) * transpose(V)[1:sdim, :]
        U2 = [U2[:, 1:convert(Int, size(U2)[2]/2)], U2[:, convert(Int, size(U2)[2]/2)+1:size(U2)[2]]]

        after_c = kron(U1[1], mps[c][1]) + kron(U1[2], mps[c][2])
        after_t = kron(U2[1], mps[t][1]) + kron(U2[2], mps[t][2])
        c_row = convert(Int, size(after_c)[1]/2)
        t_col = convert(Int, size(after_t)[2]/2)

        mps[t][1] = convert(Array{ComplexF64,2}, after_t[:, 1:t_col])
        mps[t][2] = convert(Array{ComplexF64,2}, after_t[:, t_col+1:size(after_t)[2]])
        mps[c][1] = convert(Array{ComplexF64,2}, after_c[1:c_row, :])
        mps[c][2] = convert(Array{ComplexF64,2}, after_c[c_row+1:size(after_c)[1], :])

        mps = reduce(mps, c, t, param)

        return mps
    end

    function dstack(operators)
        return cat(operators..., dims = 3)
    end

    function ddstack(operators)
        m = size(operators)[1]
        tmp2 = []
        for i=1:m
            push!(tmp2, cat(operators[i]..., dims=4))
        end
        tmp3 = Tuple(tmp2)
        mpo = cat(tmp3..., dims=3)
        return mpo
    end

    function inner_product(arrs1::Array{Any, 1}, arrs2::Array{Any, 1})
        N = Int64(size(arrs1)[1])
        ip = arrs1[1][1]' * arrs2[1][1] + arrs1[1][2]' * arrs2[1][2]
        for i = 2:N
            phys0 = arrs1[i][1]' * ip * arrs2[i][1]
            phys1 = arrs1[i][2]' * ip * arrs2[i][2]
            ip = phys0 + phys1
        end
        return ip
    end

    # function CX(arrs::Array{Any, 1}, n::Int64, t::Number)
    #     tmp = copy(arrs)
    #     arrs_ = [tmp[n + 1][1]*tmp[n + 2][1] tmp[n + 1][2]*tmp[n + 2][2]
    #             tmp[n + 1][1]*tmp[n + 2][2] tmp[n + 1][2]*tmp[n + 2][1]]
    #     arrs__ = MPSforQuantum.SVD_L(arrs_, t)
    #     col = convert(Int64, size(arrs__[1], 1) / 2)

    #     tmp[n + 1][1] = arrs__[1][1:col, :]
    #     tmp[n + 1][2] = arrs__[1][(col+1):(2*col), :]
    #     tmp[n + 2][1] = arrs__[2][:, 1:col]
    #     tmp[n + 2][2] = arrs__[2][:, (col+1):(2*col)]

    #     return tmp
    # end

    function pauli()
        pauliX = convert(Array{ComplexF64,2}, [0 1; 1 0])
        pauliY = convert(Array{ComplexF64,2}, [0 -1*im; 1*im 0])
        pauliZ = convert(Array{ComplexF64,2}, [1 0; 0 -1])
        pauliI = convert(Array{ComplexF64,2}, [1 0; 0 1])
        return (pauliX, pauliY, pauliZ, pauliI)
    end

    function pauli_zero()
        zero = convert(Array{ComplexF64,2}, [0 0; 0 0])
        return zero
    end

    function expectation(mps::Array{Any,1}, O::Array{Any,1})
        N_site = size(mps)[1]
        contracted_sites = []
        n_phys = 2 #ss
        for i=1:N_site
            if i==1
                # first step
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                a_len = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :])[1] # bb0
                arr_0 = zeros(ComplexF64, n_phys, b_len, a_len)
                arr = zeros(ComplexF64, a_len, b_len, a_len)
                @tensor begin
                    arr_0[sigma1, b, a_] = tmp_O[sigma1, sigma2, b] * tmp_mps[a_, sigma2]
                    arr[a, b, a_] = tmp_mps_dag[a, sigma1] * arr_0[sigma1, b, a_]
                end

            elseif i==N_site
                # last step
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                tmp_site = contracted_sites[i-1]
                a_len1 = size(mps[i-1][1])[2]
                b_len2 = size(O[i][1, 1, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len1, b_len2)
                arr_1 = zeros(ComplexF64, n_phys, a_len1)
                @tensor begin
                    arr_0[sigma2, a1, b2] = tmp_site[a1, b2, a1_0] * tmp_mps[a1_0, sigma2]
                    arr_1[sigma1, a1] = tmp_O[sigma1, sigma2, b2] * arr_0[sigma2, a1, b2]
                    arr = tmp_mps_dag[a1, sigma1] * arr_1[sigma1, a1]
                end

            else
                # Middle step
                tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                tmp_O = O[i]
                tmp_site = contracted_sites[i-1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len1, b_len2 = size(O[i][1, 1, :, :])
                arr_0 = zeros(ComplexF64, n_phys, a_len1, b_len1, a_len2)
                arr_1 = zeros(ComplexF64, n_phys, b_len2, a_len1, a_len2)
                arr = zeros(ComplexF64, a_len2, b_len2, a_len2)
                @tensor begin
                    arr_0[sigma2, a1, b1, a2] = tmp_site[a1, b1, a0_0] * tmp_mps[a0_0, a2, sigma2]
                    arr_1[sigma1, b2, a1, a2] = tmp_O[sigma1, sigma2, b1, b2] * arr_0[sigma2, a1, b1, a2]
                    arr[a2, b2, a2_] = tmp_mps_dag[a2, a1, sigma1] * arr_1[sigma1, b2, a1, a2_]
                end
            end
            push!(contracted_sites, arr)
        end
        return contracted_sites[N_site]
    end

    function R_expression(mps::Array{Any,1}, O::Array{Any,1}, t::Int64)
        # t: target site, sites t+1 to N are merged to R expression
        N_site = size(mps)[1]
        n_phys = 2
        contracted_sites = []
        for i = t+1:N_site
            if i==t+1
                if i==N_site
                    tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                    tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                    tmp_O = O[i]
                    a_len1 = size(mps[i-1][1])[2]
                    a_len2 = size(mps[i][1])[2]
                    b_len = size(O[i][1, 1, :])[1]
                    arr = zeros(ComplexF64, a_len1, b_len, a_len1)
                    @tensor begin
                        arr[a0, b0, a0_] = tmp_mps_dag[a0, sigma1] * tmp_O[sigma1, sigma2, b0] * tmp_mps[a0_, sigma2]
                    end
                    push!(contracted_sites, arr)
                else
                    tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                    tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                    tmp_O = O[i]
                    a_len1 = size(mps[i-1][1])[2]
                    a_len2 = size(mps[i][1])[2]
                    b_len = size(O[i][1, 1, :, :])[1]
                    arr = zeros(ComplexF64, a_len1, a_len2, b_len, b_len, a_len1, a_len2)
                    @tensor begin
                        arr[a0, a1, b0, b1, a0_, a1_] = tmp_mps_dag[a1, a0, sigma1] * tmp_O[sigma1, sigma2, b0, b1] * tmp_mps[a0_, a1_, sigma2]
                    end
                    push!(contracted_sites, arr)
                end

            elseif i==N_site
                # last step
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                tmp_site = contracted_sites[1]
                a_len0 = size(tmp_site)[1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len0, a_len1, b_len, b_len, a_len0)
                arr_1 = zeros(ComplexF64, n_phys, n_phys, a_len0, b_len, b_len, a_len0)
                arr = zeros(ComplexF64, a_len0, b_len, a_len0)
                @tensor begin
                    arr_0[sigma2, a0, a1, b0, b1, a0_] = tmp_site[a0, a1, b0, b1, a0_, a1_] * tmp_mps[a1_, sigma2]
                    arr_1[sigma1, sigma2, a0, b0, b1, a0_] = tmp_mps_dag[a1, sigma1] * arr_0[sigma2, a0, a1, b0, b1, a0_]
                    arr[a0, b0, a0_] = arr_1[sigma1, sigma2, a0, b0, b1, a0_] * tmp_O[sigma1, sigma2, b1]
                end
                contracted_sites[1] = arr
            else
                tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                tmp_O = O[i]
                tmp_site = contracted_sites[1]
                a_len0 = size(tmp_site)[1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len0, a_len1, b_len, b_len, a_len0, a_len2)
                arr_1 = zeros(ComplexF64, n_phys, n_phys, a_len0, a_len2, b_len, b_len, a_len0, a_len2)
                arr = zeros(ComplexF64, a_len0, a_len2, b_len, b_len, a_len0, a_len2)
                @tensor begin
                    arr_0[sigma2, a0, a1, b0, b1, a0_, a2_] = tmp_site[a0, a1, b0, b1, a0_, a1_] * tmp_mps[a1_, a2_, sigma2]
                    arr_1[sigma1, sigma2, a0, a2, b0, b1, a0_, a2_] = tmp_mps_dag[a2, a1, sigma1] * arr_0[sigma2, a0, a1, b0, b1, a0_, a2_]
                    arr[a0, a2, b0, b2, a0_, a2_] = arr_1[sigma1, sigma2, a0, a2, b0, b1, a0_, a2_] * tmp_O[sigma1, sigma2, b1, b2]
                end
                contracted_sites[1] = arr
            end
        end
        return contracted_sites[1]
    end

    function L_expression(mps::Array{Any,1}, O::Array{Any,1}, t::Int64)
        # t: target site, sites 1 to t-1 are merged to L expression
        N_site = size(mps)[1]
        n_phys = 2
        contracted_sites = []
        for i = 1:t-1
            if i==1
                tmp_mps = cat(mps[i][1][:], mps[i][2][:], dims=2)
                tmp_mps_dag = cat(mps[i][1]'[:], mps[i][2]'[:], dims=2)
                tmp_O = O[i]
                a_len = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :, :])[1]
                arr = zeros(ComplexF64, a_len, b_len, a_len)
                @tensor begin
                    arr[a0, b0, a0_] = tmp_mps_dag[a0, sigma1] * tmp_O[sigma1, sigma2, b0] * tmp_mps[a0_, sigma2]
                end
                push!(contracted_sites, arr)
            else
                tmp_mps = cat(mps[i][1], mps[i][2], dims=3)
                tmp_mps_dag = cat(mps[i][1]', mps[i][2]', dims=3)
                tmp_O = O[i]
                tmp_site = contracted_sites[1]
                a_len0 = size(tmp_site)[1]
                a_len1 = size(mps[i-1][1])[2]
                a_len2 = size(mps[i][1])[2]
                b_len = size(O[i][1, 1, :, :])[1]
                arr_0 = zeros(ComplexF64, n_phys, a_len1, b_len, a_len2)
                arr_1 = zeros(ComplexF64, n_phys, n_phys, a_len2, b_len, a_len2)
                arr = zeros(ComplexF64, a_len2, b_len, a_len2)
                @tensor begin
                    arr_0[sigma2, a1, b1, a2_] = tmp_site[a1, b1, a1_] * tmp_mps[a1_, a2_, sigma2]
                    arr_1[sigma1, sigma2, a2, b1, a2_] = tmp_mps_dag[a2, a1, sigma1] * arr_0[sigma2, a1, b1, a2_]
                    arr[a2, b2, a2_] = arr_1[sigma1, sigma2, a2, b1, a2_] * tmp_O[sigma1, sigma2, b1, b2]
                end
                contracted_sites[1] = arr
            end
        end
        return contracted_sites[1]
    end

    function left_norm_for_2_sites(mps::Array{Any,1}, t::Int64, D::Int64)
        # t: target, site to be left normalized
        mps_ = copy(mps)
        site_1 = cat(mps[t][1], mps[t][2], dims=1)
        site_2 = cat(mps[t+1][1], mps[t+1][2], dims=2)
        mixed_site = site_1 * site_2
        A, M = SVD_L(mixed_site, D)
        col = convert(Int64, size(A, 1) / 2)
        mps_[t] = [A[1:col, :], A[(col+1):col*2, :]]
        col2 = convert(Int64, size(M, 2) / 2)
        mps_[t+1] = [M[:, 1:col2], M[:, (col2+1):col2*2]]
        return mps_
    end

    function right_norm_for_2_sites(mps::Array{Any,1}, t::Int64, D::Int64)
        # t: target, site to be right normalized
        mps_ = copy(mps)
        site_1 = cat(mps[t-1][1], mps[t-1][2], dims=1)
        site_2 = cat(mps[t][1], mps[t][2], dims=2)
        mixed_site = site_1 * site_2
        M, B = SVD_R(mixed_site, D)
        col = convert(Int64, size(M, 1) / 2)
        mps_[t-1] = [M[1:col, :], M[(col+1):col*2, :]]
        col2 = convert(Int64, size(B, 2) / 2)
        mps_[t] = [B[:, 1:col2], B[:, (col2+1):col2*2]]
        return mps_
    end

    function left_most_site_update(mps_::Array{Any,1}, O::Array{Any,1}, t::Int64)
        R = R_expression(mps_, O, t)
        tmp_O = O[t]
        H = zeros(ComplexF64, size(R, 1), size(R, 3), size(tmp_O, 1), size(tmp_O, 2))
        @tensor begin
            H[sigma1, a0, sigma2, a0_] = tmp_O[sigma1, sigma2, b0] * R[a0, b0, a0_]
        end
        H_ = zeros(ComplexF64, size(tmp_O, 1) * size(R, 1), size(tmp_O, 2) * size(R, 3))
        for i=1:size(H, 1), j=1:size(H, 2), k=1:size(H, 3), l=1:size(H, 4)
            H_[(i-1)*size(H, 2) + j, (k-1)*size(H, 4) + l] = H[i,j,k,l]
        end
        v = eigvecs(H_)[:, 1]
        d = convert(Int64, size(v, 1) / 2)
        mps_[t][1][:] = v[1:d]
        mps_[t][2][:] = v[d+1:2*d]
        return mps_
    end

    function mid_site_update(mps_::Array{Any,1}, O::Array{Any,1}, t::Int64)
        R = R_expression(mps_, O, t)
        L = L_expression(mps_, O, t)
        tmp_O = O[t]
        H = zeros(ComplexF64, size(tmp_O, 1), size(L, 1), size(R, 1), size(tmp_O, 2), size(L, 3), size(R, 3))
        @tensor begin
            H[sigma1, a0, a1, sigma2, a0_, a1_] = L[a0, b0, a0_] * tmp_O[sigma1, sigma2, b0, b1] * R[a1, b1, a1_]
        end
        H_ = zeros(ComplexF64, size(tmp_O, 1) * size(L, 1) * size(R, 1), size(tmp_O, 2) * size(L, 3) * size(R, 3))
        for i=1:size(H, 1), j=1:size(H, 2), k=1:size(H, 3), l=1:size(H, 4), m=1:size(H, 5), n=1:size(H, 6)
            H_[(i-1)*size(H, 2)*size(H, 3) + (j-1)*size(H, 3) + k, (l-1)*size(H, 5)*size(H, 6) + (m-1)*size(H, 6) + n] = H[i,j,k,l,m,n]
        end
        v = eigvecs(H_)[:, 1]
        d = convert(Int64, size(v, 1) / 2)
        M_1 = transpose(reshape(v[1:d], size(transpose(mps_[t][1]))))
        M_2 = transpose(reshape(v[d+1:2*d], size(transpose(mps_[t][1]))))
        mps_[t][1][:, :] = M_1
        mps_[t][2][:, :] = M_2
        return mps_
    end

    function right_most_site_update(mps_::Array{Any,1}, O::Array{Any,1}, t::Int64)
        L = L_expression(mps_, O, t)
        tmp_O = O[t]
        H = zeros(ComplexF64, size(L, 1), size(L, 3), size(tmp_O, 1), size(tmp_O, 2))
        @tensor begin
            H[sigma1, a0, sigma2, a0_] = tmp_O[sigma1, sigma2, b0] * L[a0, b0, a0_]
        end
        H_ = zeros(ComplexF64, size(tmp_O, 1) * size(L, 1), size(tmp_O, 2) * size(L, 3))
        for i=1:size(H, 1), j=1:size(H, 2), k=1:size(H, 3), l=1:size(H, 4)
            H_[(i-1)*size(H, 2) + j, (k-1)*size(H, 4) + l] = H[i,j,k,l]
        end
        v = eigvecs(H_)[:, 1]
        d = convert(Int64, size(v, 1) / 2)
        mps_[t][1][:] = v[1:d]
        mps_[t][2][:] = v[d+1:2*d]
        return mps_
    end

    function iterative_ground_state_search(mps::Array{Any,1}, O::Array{Any,1}, D::Int64, ite::Int64)
        hist = []
        N = size(mps, 1)
        mps_ = copy(mps)
        push!(hist, expectation(mps_, O))
        for i = 1:ite
            for t = 1:N
                if t == 1
                    mps_ = left_most_site_update(mps_, O, t)
                    mps_ = left_norm_for_2_sites(mps_, t, D)
                elseif t == N
                    mps_ = right_most_site_update(mps_, O, t)
                else
                    mps_ = mid_site_update(mps_, O, t)
                    mps_ = left_norm_for_2_sites(mps_, t, D)
                end
                # push!(hist, expectation(mps_, O))
            end

            for t = N-1:-1:1
                mps_ = right_norm_for_2_sites(mps_, t+1, D)
                if t == 1
                    mps_ = left_most_site_update(mps_, O, t)
                else
                    mps_ = mid_site_update(mps_, O, t)
                end
                push!(hist, expectation(mps_, O))
            end
        end
        return (mps_, hist)
    end

    function init_rand_MPS(N::Int64, D::Int64, normalize::Char='l')
        mps = []
        mps2 = []
        half_N = convert(Int64, floor(N/2))
        for i in 1:half_N
            l = 2^(i-1)
            mps_1 = []
            mps_2 = []
            for j in 1:2
                arr1 = rand(ComplexF64, (min(l, D), min(l*2, D))) / min(l*2, D)
                arr2 = rand(ComplexF64, (min(l*2, D), min(l, D))) / min(l, D)
                push!(mps_1, arr1)
                push!(mps_2, arr2)
            end
            push!(mps, mps_1)
            push!(mps2, mps_2)
        end

        if N%2 == 1
            mps_ = []
            for j in 1:2
                arr = rand(ComplexF64, (min(2^half_N, D), min(2^half_N, D))) / min(2^half_N, D)
                push!(mps_, arr)
            end
            push!(mps, mps_)
        end

        for j in 1:half_N
            push!(mps, mps2[half_N - (j-1)])
        end

        for i in 1:N-1
            mps = left_norm_for_2_sites(mps, i, D)
        end
        norm = mps[N][1]' * mps[N][1] + mps[N][2]' * mps[N][2]
        mps[N][1] = mps[N][1]/sqrt(norm)
        mps[N][2] = mps[N][2]/sqrt(norm)
        if normalize == 'r'
            for i in N:-1:2
                mps = right_norm_for_2_sites(mps, i, D)
            end
        end
        return mps
    end

    function push_N(arr::Array{Any,1}, op::Array{Complex{Float64},2}, N::Int64)
        arr_ = copy(arr)
        for i in 1:N
            push!(arr_, op)
        end
        return arr_
    end

    function j1j2_connection(L::Int64)
        J1_term = Set()
        for i in 1:L
            for j in 1:(L-1)
                row = 1 + 4*(i-1)
                push!(J1_term, (row + j - 1, row + j))
                col = i
                push!(J1_term, (col + 4*(j - 1), col + 4 + 4*(j - 1)))
            end
        end

        J2_term = Set()
        for i in 1:(L-1)
            for j in 1:(L-1)
                row = 1 + 4*(i-1)
                push!(J2_term, (row + (j - 1), row + (j - 1) + 5))
                push!(J2_term, (row + j, row + j + 3))
            end
        end
        return (J1_term, J2_term)
    end

    function j1j2_2D_Hamiltonian_MPO(L::Int64, j1::Float64, j2::Float64)
        N = L^2
        max_distance = L + 1
        MPO_L = 3*(max_distance + 2)
        O = []
        first_MPO = []
        (pauliX, pauliY, pauliZ, pauliI) = pauli()
        zero = pauli_zero()
        ops = [pauliX, pauliY, pauliZ]
        (J1_term, J2_term) = j1j2_connection(L)
        for op in ops
            push!(first_MPO, zero)
            for k in 1:max_distance
                if (1, 1 + k) in J1_term
                    push!(first_MPO, op*j1)
                elseif (1, 1 + k) in J2_term
                    push!(first_MPO, op*j2)
                else
                    push!(first_MPO, zero)
                end
            end
            push!(first_MPO, zero)
        end
        first_MPO_tuple = Tuple(first_MPO)
        push!(O, dstack(first_MPO_tuple))

        for site in 1:N-2
            mid_MPO = []
            for i in 1:length(ops)
                for row in 1:(max_distance + 2)
                    tmp_row = []
                    for j in 1:(i-1) # fill zero
                        tmp_row = push_N(tmp_row, zero, max_distance + 2)
                    end

                    if row == 1
                        push!(tmp_row, pauliI)
                        tmp_row = push_N(tmp_row, zero, max_distance + 1)
                    elseif row == 2
                        push!(tmp_row, ops[i])
                        tmp_row = push_N(tmp_row, zero, max_distance + 1)
                    elseif row == (max_distance + 2)
                        push!(tmp_row, zero)
                        for k in 1:max_distance
                            if (site, site + k) in J1_term
                                push!(tmp_row, j1*ops[i])
                            elseif (site, site + k) in J1_term
                                push!(tmp_row, j2*ops[i])
                            else
                                push!(tmp_row, zero)
                            end
                        end
                        push!(tmp_row, zero)
                    else
                        tmp_row = push_N(tmp_row, zero, row - 2)
                        push!(tmp_row, pauliI)
                        tmp_row = push_N(tmp_row, zero, max_distance + 1 - (row - 2))
                    end
                    for j in 1:(3-i) # fill zero
                        tmp_row = push_N(tmp_row, zero, max_distance + 2)
                    end
                    tmp_row_tuple = Tuple(tmp_row)
                    push!(mid_MPO, tmp_row_tuple)
                end
            end
            push!(O, ddstack(mid_MPO))
        end

        last_MPO = []
        for op in ops
            push!(last_MPO, pauliI)
            push!(last_MPO, op)
            last_MPO = push_N(last_MPO, zero, max_distance)
        end
        last_MPO_tuple = Tuple(last_MPO)
        push!(O, dstack(last_MPO_tuple))
        return O
    end

end

max_stars_repo_name: ryuNagai/blueqat-julia

max_issues_repo_name: ryuNagai/blueqat-julia

max_forks_repo_name: ryuNagai/blueqat-julia