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

immutable LaneTag
    segment :: Int # segment id
    lane    :: Int # index in segment.lanes of this lane
end
Base.show(io::IO, tag::LaneTag) = @printf(io, "LaneTag(%d, %d)", tag.segment, tag.lane)
hash(id::LaneTag, h::UInt=zero(UInt)) = hash(id.segment, hash(id.lane, h))
@compat Base.:(==)(a::LaneTag, b::LaneTag) = a.segment == b.segment && a.lane == b.lane
const NULL_LANETAG = LaneTag(0,0)

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

immutable RoadIndex
    ind::CurveIndex
    tag::LaneTag
end
const NULL_ROADINDEX = RoadIndex(CurveIndex(-1,NaN), LaneTag(-1,-1))

Base.show(io::IO, r::RoadIndex) = @printf(io, "RoadIndex({%d, %3.f}, {%d, %d})", r.ind.i, r.ind.t, r.tag.segment, r.tag.lane)
Base.write(io::IO, r::RoadIndex) = @printf(io, "%d %.6f %d %d", r.ind.i, r.ind.t, r.tag.segment, r.tag.lane)

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

immutable LaneConnection
    downstream::Bool # if true, mylane → target, else target → mylane
    mylane::CurveIndex
    target::RoadIndex
end

Base.show(io::IO, c::LaneConnection) = print(io, "LaneConnection(", c.downstream ? "D" : "U", ", ", c.mylane, ", ", c.target)
function Base.write(io::IO, c::LaneConnection)
    @printf(io, "%s (%d %.6f) ", c.downstream ? "D" : "U", c.mylane.i, c.mylane.t)
    write(io, c.target)
end
function Base.parse(::Type{LaneConnection}, line::AbstractString)
    cleanedline = replace(line, r"(\(|\))", "")
    tokens = split(cleanedline, ' ')

    @assert(tokens[1] == "D" || tokens[1] == "U")
    downstream = tokens[1] == "D"
    mylane = CurveIndex(parse(Int, tokens[2]), parse(Float64, tokens[3]))
    target = RoadIndex(
                CurveIndex(parse(Int, tokens[4]), parse(Float64, tokens[5])),
                LaneTag(parse(Int, tokens[6]), parse(Int, tokens[7]))
            )
    LaneConnection(downstream, mylane, target)
end

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

const DEFAULT_LANE_WIDTH = 3.0 # [m]

type Lane
    tag            :: LaneTag
    curve          :: Curve
    width          :: Float64 # [m]
    speed_limit    :: SpeedLimit
    boundary_left  :: LaneBoundary
    boundary_right :: LaneBoundary
    exits          :: Vector{LaneConnection} # list of exits; put the primary exit (at end of lane) first
    entrances      :: Vector{LaneConnection} # list of entrances; put the primary entrance (at start of lane) first

    function Lane(
        tag::LaneTag,
        curve::Curve;
        width::Float64 = DEFAULT_LANE_WIDTH,
        speed_limit::SpeedLimit = DEFAULT_SPEED_LIMIT,
        boundary_left::LaneBoundary = NULL_BOUNDARY,
        boundary_right::LaneBoundary = NULL_BOUNDARY,
        exits::Vector{LaneConnection} = LaneConnection[],
        entrances::Vector{LaneConnection} = LaneConnection[],
        next::RoadIndex=NULL_ROADINDEX,
        prev::RoadIndex=NULL_ROADINDEX,
        )

        retval = new()
        retval.tag = tag
        retval.curve = curve
        retval.width = width
        retval.speed_limit = speed_limit
        retval.boundary_left = boundary_left
        retval.boundary_right = boundary_right
        retval.exits = exits
        retval.entrances = entrances

        if next != NULL_ROADINDEX
            unshift!(retval.exits, LaneConnection(true, curveindex_end(retval.curve), next))
        end
        if prev != NULL_ROADINDEX
            unshift!(retval.entrances, LaneConnection(false, CURVEINDEX_START, prev))
        end

        retval
    end
end

has_next(lane::Lane) = !isempty(lane.exits) && lane.exits[1].mylane == curveindex_end(lane.curve)
has_prev(lane::Lane) = !isempty(lane.entrances) && lane.entrances[1].mylane == CURVEINDEX_START

is_in_exits(lane::Lane, target::LaneTag) = findfirst(lc->lc.target.tag == target, lane.exits) != 0
is_in_entrances(lane::Lane, target::LaneTag) = findfirst(lc->lc.target.tag == target, lane.entrances) != 0

function connect!(source::Lane, dest::Lane)
    # place these at the front

    cindS = curveindex_end(source.curve)
    cindD = CURVEINDEX_START

    unshift!(source.exits,   LaneConnection(true,  cindS, RoadIndex(cindD, dest.tag)))
    unshift!(dest.entrances, LaneConnection(false, cindD, RoadIndex(cindS, source.tag)))
    (source, dest)
end
function connect!(source::Lane, cindS::CurveIndex, dest::Lane, cindD::CurveIndex)
    push!(source.exits,   LaneConnection(true,  cindS, RoadIndex(cindD, dest.tag)))
    push!(dest.entrances, LaneConnection(false, cindD, RoadIndex(cindS, source.tag)))
    (source, dest)
end

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

type RoadSegment
    # a list of lanes forming a single road with a common direction
    id::Int
    lanes::Vector{Lane} # lanes are stored right to left
    RoadSegment(id::Int, lanes::Vector{Lane}=Lane[]) = new(id, lanes)
end

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

type Roadway
    segments::Vector{RoadSegment}
    Roadway(segments::Vector{RoadSegment}=RoadSegment[]) = new(segments)
end

function Base.write(io::IO, roadway::Roadway)
    # writes to a text file

    println(io, "ROADWAY")
    println(io, length(roadway.segments)) # number of segments
    for seg in roadway.segments
        println(io, seg.id)
        println(io, "\t", length(seg.lanes)) # number of lanes
        for (i,lane) in enumerate(seg.lanes)
            @assert(lane.tag.lane == i)
            @printf(io, "\t%d\n", i)
            @printf(io, "\t\t%.3f\n", lane.width)
            @printf(io, "\t\t%.3f %.3f\n", lane.speed_limit.lo, lane.speed_limit.hi)
            println(io, "\t\t", lane.boundary_left.style, " ", lane.boundary_left.color)
            println(io, "\t\t", lane.boundary_right.style, " ", lane.boundary_right.color)
            println(io, "\t\t", length(lane.exits) + length(lane.entrances))
            for conn in lane.exits
                print(io, "\t\t\t"); write(io, conn); print(io, "\n")
            end
            for conn in lane.entrances
                print(io, "\t\t\t"); write(io, conn); print(io, "\n")
            end
            println(io, "\t\t", length(lane.curve))
            for pt in lane.curve
                @printf(io, "\t\t\t(%.4f %.4f %.6f) %.4f %.8f %.8f\n", pt.pos.x, pt.pos.y, pt.pos.θ, pt.s, pt.k, pt.kd)
            end
        end
    end
end
function Base.read(io::IO, ::Type{Roadway})
    lines = readlines(io)
    line_index = 1
    if contains(lines[line_index], "ROADWAY")
        line_index += 1
    end

    function advance!()
        line = strip(lines[line_index])
        line_index += 1
        line
    end

    nsegs = parse(Int, advance!())
    roadway = Roadway(Array{RoadSegment}(nsegs))
    for i_seg in 1:nsegs
        segid = parse(Int, advance!())
        nlanes = parse(Int, advance!())
        seg = RoadSegment(segid, Array{Lane}(nlanes))
        for i_lane in 1:nlanes
            @assert(i_lane == parse(Int, advance!()))
            tag = LaneTag(segid, i_lane)
            width = parse(Float64, advance!())

            tokens = split(advance!(), ' ')
            speed_limit = SpeedLimit(parse(Float64, tokens[1]), parse(Float64, tokens[2]))

            tokens = split(advance!(), ' ')
            boundary_left = LaneBoundary(Symbol(tokens[1]), Symbol(tokens[2]))

            tokens = split(advance!(), ' ')
            boundary_right = LaneBoundary(Symbol(tokens[1]), Symbol(tokens[2]))

            exits = LaneConnection[]
            entrances = LaneConnection[]
            n_conns = parse(Int, advance!())
            for i_conn in 1:n_conns
                conn = parse(LaneConnection, advance!())
                conn.downstream ? push!(exits, conn) : push!(entrances, conn)
            end

            npts = parse(Int, advance!())
            curve = Array{CurvePt}(npts)
            for i_pt in 1:npts
                line = advance!()
                cleanedline = replace(line, r"(\(|\))", "")
                tokens = split(cleanedline, ' ')
                x = parse(Float64, tokens[1])
                y = parse(Float64, tokens[2])
                θ = parse(Float64, tokens[3])
                s = parse(Float64, tokens[4])
                k = parse(Float64, tokens[5])
                kd = parse(Float64, tokens[6])
                curve[i_pt] = CurvePt(VecSE2(x,y,θ), s, k, kd)
            end

            seg.lanes[i_lane] = Lane(tag, curve, width=width, speed_limit=speed_limit,
                                     boundary_left=boundary_left,
                                     boundary_right=boundary_right,
                                     entrances=entrances, exits=exits)
        end
        roadway.segments[i_seg] = seg
    end

    roadway
end

"""
    lane[ind::CurveIndex, roadway::Roadway]
Accessor for lanes based on a CurveIndex.
Note that we extend the definition of a CurveIndex,
previously ind.i ∈ [1, length(curve)-1], to:

    ind.i ∈ [0, length(curve)]

where 1 ≤ ind.i ≤ length(curve)-1 is as before, but if the index
is on the section between two lanes, we use:

    ind.i = length(curve), ind.t ∈ [0,1] for the region between curve[end] → next
    ind.i = 0,             ind.t ∈ [0,1] for the region between prev → curve[1]
"""
function Base.getindex(lane::Lane, ind::CurveIndex, roadway::Roadway)
    if ind.i == 0
        pt_lo = prev_lane_point(lane, roadway)
        pt_hi = lane.curve[1]
        s_gap = abs(pt_hi.pos - pt_lo.pos)
        pt_lo = CurvePt(pt_lo.pos, -s_gap, pt_lo.k, pt_lo.kd)
        lerp(pt_lo, pt_hi, ind.t)
    elseif ind.i < length(lane.curve)
        lane.curve[ind]
    else
        pt_hi = next_lane_point(lane, roadway)
        pt_lo = lane.curve[end]
        s_gap = abs(pt_hi.pos - pt_lo.pos)
        pt_hi = CurvePt(pt_hi.pos, pt_lo.s + s_gap, pt_hi.k, pt_hi.kd)
        lerp( pt_lo, pt_hi, ind.t)
    end
end
function Base.getindex(roadway::Roadway, segid::Int)
    for seg in roadway.segments
        if seg.id == segid
            return seg
        end
    end
    error("Could not find segid $segid in roadway")
end
function Base.getindex(roadway::Roadway, tag::LaneTag)
    seg = roadway[tag.segment]
    seg.lanes[tag.lane]
end

is_between_segments_lo(ind::CurveIndex) = ind.i == 0
is_between_segments_hi(ind::CurveIndex, curve::Curve) = ind.i == length(curve)
is_between_segments(ind::CurveIndex, curve::Curve) = is_between_segments_lo(ind) || is_between_segments_hi(ind, curve)

next_lane(lane::Lane, roadway::Roadway) = roadway[lane.exits[1].target.tag]
prev_lane(lane::Lane, roadway::Roadway) = roadway[lane.entrances[1].target.tag]
next_lane_point(lane::Lane, roadway::Roadway) = roadway[lane.exits[1].target]
prev_lane_point(lane::Lane, roadway::Roadway) = roadway[lane.entrances[1].target]

function has_segment(roadway::Roadway, segid::Int)
    for seg in roadway.segments
        if seg.id == segid
            return true
        end
    end
    false
end
function has_lanetag(roadway::Roadway, tag::LaneTag)
    if !has_segment(roadway, tag.segment)
        return false
    end
    seg = roadway[tag.segment]
    1 ≤ tag.lane ≤ length(seg.lanes)
end

immutable RoadProjection
    curveproj::CurveProjection
    tag::LaneTag
end

function get_closest_perpendicular_point_between_points(A::VecSE2, B::VecSE2, Q::VecSE2;
    tolerance::Float64 = 0.01, # acceptable error in perpendicular component
    max_iter::Int = 50, # maximum number of iterations
    )

    # CONDITIONS: a < b, either f(a) < 0 and f(b) > 0 or f(a) > 0 and f(b) < 0
    # OUTPUT: value which differs from a root of f(x)=0 by less than TOL

    a = 0.0
    b = 1.0

    f_a = inertial2body(Q, A).x
    f_b = inertial2body(Q, B).x

    if sign(f_a) == sign(f_b) # both are wrong - use the old way
        t = get_lerp_time_unclamped(A, B, Q)
        t = clamp(t, 0.0, 1.0)
        return (t, lerp(A,B,t))
    end

    iter = 1
    while iter ≤ max_iter
        c = (a+b)/2 # new midpoint
        footpoint = lerp(A, B, c)
        f_c = inertial2body(Q, footpoint).x
        if abs(f_c) < tolerance # solution found
            return (c, footpoint)
        end
        if sign(f_c) == sign(f_a)
            a, f_a = c, f_c
        else
            b = c
        end
        iter += 1
    end

    # Maximum number of iterations passed
    # This will occur when we project with a point that is not actually in the range,
    # and we converge towards one edge

    if a == 0.0
        return (0.0, A)
    elseif b == 1.0
        return (1.0, B)
    else
        warn("get_closest_perpendicular_point_between_points - should not happen")
        c = (a+b)/2 # should not happen
        return (c, lerp(A,B,c))
    end
end

"""
    proj(posG::VecSE2, lane::Lane, roadway::Roadway)
Return the RoadProjection for projecting posG onto the lane.
This will automatically project to the next or prev curve as appropriate.
"""
function Vec.proj(posG::VecSE2, lane::Lane, roadway::Roadway;
    move_along_curves::Bool = true, # if false, will only project to lane.curve
    )
    curveproj = proj(posG, lane.curve)
    rettag = lane.tag

    if curveproj.ind == CurveIndex(1,0.0) && has_prev(lane)
        pt_lo = prev_lane_point(lane, roadway)
        pt_hi = lane.curve[1]

        t = get_lerp_time_unclamped(pt_lo, pt_hi, posG)
        if t ≤ 0.0 && move_along_curves
            return proj(posG, prev_lane(lane, roadway), roadway)
        elseif t < 1.0 # for t == 1.0 we use the actual end of the lane
            @assert(!move_along_curves || 0.0 ≤ t < 1.0)

            # t was computed assuming a constant angle
            # this is not valid for the large distances and angle disparities between lanes
            # thus we now use a bisection search to find the appropriate location

            t, footpoint = get_closest_perpendicular_point_between_points(pt_lo.pos, pt_hi.pos, posG)

            ind = CurveIndex(0, t)
            curveproj = get_curve_projection(posG, footpoint, ind)
        end
    elseif curveproj.ind == curveindex_end(lane.curve) && has_next(lane)
        pt_lo = lane.curve[end]
        pt_hi = next_lane_point(lane, roadway)

        t = get_lerp_time_unclamped(pt_lo, pt_hi, posG)

        if t ≥ 1.0 && move_along_curves
             # for t == 1.0 we use the actual start of the lane
            return proj(posG, next_lane(lane, roadway), roadway)
        elseif t ≥ 0.0
            @assert(!move_along_curves || 0.0 ≤ t ≤ 1.0)

            # t was computed assuming a constant angle
            # this is not valid for the large distances and angle disparities between lanes
            # thus we now use a bisection search to find the appropriate location

            t, footpoint = get_closest_perpendicular_point_between_points(pt_lo.pos, pt_hi.pos, posG)

            ind = CurveIndex(length(lane.curve), t)
            curveproj = get_curve_projection(posG, footpoint, ind)
        end
    end
    RoadProjection(curveproj, rettag)
end

"""
    proj(posG::VecSE2, seg::RoadSegment, roadway::Roadway)
Return the RoadProjection for projecting posG onto the segment.
Tries all of the lanes and gets the closest one
"""
function Vec.proj(posG::VecSE2, seg::RoadSegment, roadway::Roadway)

    best_dist2 = Inf
    best_proj = RoadProjection(CurveProjection(CurveIndex(-1,-1), NaN, NaN), NULL_LANETAG)

    for lane in seg.lanes
        roadproj = proj(posG, lane, roadway)
        footpoint = roadway[roadproj.tag][roadproj.curveproj.ind, roadway]
        dist2 = abs2(posG - footpoint.pos)
        if dist2 < best_dist2
            best_dist2 = dist2
            best_proj = roadproj
        end
    end

    best_proj
end

"""
    proj(posG::VecSE2, seg::RoadSegment, roadway::Roadway)
Return the RoadProjection for projecting posG onto the roadway.
Tries all of the lanes and gets the closest one
"""
function Vec.proj(posG::VecSE2, roadway::Roadway)

    best_dist2 = Inf
    best_proj = RoadProjection(CurveProjection(CurveIndex(-1,-1), NaN, NaN), NULL_LANETAG)

    for seg in roadway.segments
        for lane in seg.lanes
            roadproj = proj(posG, lane, roadway, move_along_curves=false)
            targetlane = roadway[roadproj.tag]
            footpoint = targetlane[roadproj.curveproj.ind, roadway]
            dist2 = abs2(posG - footpoint.pos)
            if dist2 < best_dist2
                best_dist2 = dist2
                best_proj = roadproj
            end
        end
    end

    best_proj
end

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

RoadIndex(roadproj::RoadProjection) = RoadIndex(roadproj.curveproj.ind, roadproj.tag)

function Base.getindex(roadway::Roadway, roadind::RoadIndex)
    lane = roadway[roadind.tag]
    lane[roadind.ind, roadway]
end

"""
    move_along(roadind::RoadIndex, road::Roadway, Δs::Float64)
Return the RoadIndex at ind's s position + Δs
"""
function move_along(roadind::RoadIndex, roadway::Roadway, Δs::Float64, depth::Int=0)

    lane = roadway[roadind.tag]
    curvept = lane[roadind.ind, roadway]

    if curvept.s + Δs < 0.0
        if has_prev(lane)
            pt_lo = prev_lane_point(lane, roadway)
            pt_hi = lane.curve[1]
            s_gap = abs(pt_hi.pos - pt_lo.pos)

            if curvept.s + Δs < -s_gap
                lane_prev = prev_lane(lane, roadway)
                curveind = curveindex_end(lane_prev.curve)
                roadind = RoadIndex(curveind, lane_prev.tag)
                return move_along(roadind, roadway, Δs + curvept.s + s_gap, depth+1)
            else # in the gap between lanes
                t = (s_gap + curvept.s + Δs) / s_gap
                curveind = CurveIndex(0, t)
                RoadIndex(curveind, lane.tag)
            end

        else # no prev lane, return the beginning of this one
            curveind = CurveIndex(1, 0.0)
            return RoadIndex(curveind, roadind.tag)
        end
    elseif curvept.s + Δs > lane.curve[end].s
        if has_next(lane)
            pt_lo = lane.curve[end]
            pt_hi = next_lane_point(lane, roadway)
            s_gap = abs(pt_hi.pos - pt_lo.pos)

            if curvept.s + Δs ≥ pt_lo.s + s_gap # extends beyond the gap
                curveind = lane.exits[1].target.ind
                roadind = RoadIndex(curveind, lane.exits[1].target.tag)
                return move_along(roadind, roadway, Δs - (lane.curve[end].s + s_gap - curvept.s))
            else # in the gap between lanes
                t = (Δs - (lane.curve[end].s - curvept.s)) / s_gap
                curveind = CurveIndex(0, t)
                RoadIndex(curveind, lane.exits[1].target.tag)
            end
        else # no next lane, return the end of this lane
            curveind = curveindex_end(lane.curve)
            return RoadIndex(curveind, roadind.tag)
        end
    else
        if roadind.ind.i == 0
            ind = get_curve_index(CurveIndex(1,0.0), lane.curve, curvept.s+Δs)
        elseif roadind.ind.i == length(lane.curve)
            ind = get_curve_index(curveindex_end(lane.curve), lane.curve, curvept.s+Δs)
        else
            ind = get_curve_index(roadind.ind, lane.curve, Δs)
        end
        RoadIndex(ind, roadind.tag)
    end
end

n_lanes_right(lane::Lane, roadway::Roadway) = lane.tag.lane - 1
function n_lanes_left(lane::Lane, roadway::Roadway)
    seg = roadway[lane.tag.segment]
    length(seg.lanes) - lane.tag.lane
end

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

function _fit_curve(
    pts::AbstractMatrix{Float64}, # 2×n
    desired_distance_between_samples::Real;
    max_iterations::Int=50,
    epsilon::Float64=1e-4,
    n_intervals_in_arclen::Int=100,
    )

    @assert(size(pts, 1) == 2)

    spline_coeffs = fit_cubic_spline(pts)

    L = calc_curve_length(spline_coeffs[1], spline_coeffs[2], n_intervals_per_segment=n_intervals_in_arclen)
    n = round(Int, L/desired_distance_between_samples, RoundNearestTiesUp)+1

    s_arr = collect(linspace(0.0,L,n))
    t_arr = calc_curve_param_given_arclen(spline_coeffs[1], spline_coeffs[2], s_arr,
        curve_length=L, max_iterations=max_iterations, epsilon=epsilon, n_intervals_in_arclen=n_intervals_in_arclen)

    x_arr = sample_spline(spline_coeffs[1], t_arr)
    y_arr = sample_spline(spline_coeffs[2], t_arr)
    θ_arr = sample_spline_theta(spline_coeffs[1], spline_coeffs[2], t_arr)

    κ_arr = sample_spline_curvature(spline_coeffs[1], spline_coeffs[2], t_arr)
    κd_arr = sample_spline_derivative_of_curvature(spline_coeffs[1], spline_coeffs[2], t_arr)

    @assert(!any(s->isnan(s), s_arr))
    @assert(!any(s->isnan(s), x_arr))
    @assert(!any(s->isnan(s), y_arr))
    @assert(!any(s->isnan(s), θ_arr))

    curve = Array{CurvePt}(n)
    for i in 1 : n
        pos = VecSE2(x_arr[i], y_arr[i], θ_arr[i])
        curve[i] = CurvePt(pos, s_arr[i], κ_arr[i], κd_arr[i])
    end
    curve
end

"""
    read_dxf(io::IO, ::Type{Roadway})
Return a Roadway generated from a DXF file

    Layers with names such as seg001 will contain LWPOLYLINEs.
    Each LWPOLYLINE corresponds to a lane centerline, which together
    are all neighbored.
"""
function read_dxf(io::IO, ::Type{Roadway};
    dist_threshold_lane_connect::Float64 = 0.25, # [m]
    desired_distance_between_curve_samples::Float64 = 1.0; # [m]
    )

    lines = readlines(io)

    i = findfirst(lines, "ENTITIES\n")
    i != 0 || error("ENTITIES section not found")

    ###################################################
    # Pull pts for each lane
    lane_pts_dict = Dict{LaneTag, Vector{VecE2}}()

    i = findnext(lines, "LWPOLYLINE\n", i)
    while i != 0
        i = findnext(lines, "  8\n", i)
        if i != 0 # segment identifier found in LWPOLYLINE

            if ismatch(r"(?<=seg)(\d*)", lines[i+1])
                segid = parse(Int, match(r"(?<=seg)(\d*)", lines[i+1]).match)

                    i = findnext(lines, "AcDbPolyline\n", i)
                i != 0 || error("AcDbPolyline not found in LWPOLYLINE!")
                i = findnext(lines, " 90\n", i)
                i != 0 || error("Number of vertices not found in AcDbPolyline!")

                N = parse(Int, lines[i+1])
                N > 0 || error("Empty line segment!")

                pts = Array{VecE2}(N)

                i = findnext(lines, " 10\n", i)
                i != 0 || error("First point not found in AcDbPolyline!")

                for j in 1 : N
                    x = parse(Float64, lines[i+1])
                    y = parse(Float64, lines[i+3])
                    i += 4
                    pts[j] = VecE2(x,y)
                end

                laneid = 1
                for tag in keys(lane_pts_dict)
                    if tag.segment == segid
                        laneid += 1
                    end
                end
                lane_pts_dict[LaneTag(segid, laneid)] = pts
            end

            i = findnext(lines, "LWPOLYLINE\n", i)
        end
    end

    ###################################################
    # Shift pts to connect to previous / next pts
    lane_next_dict = Dict{LaneTag, Tuple{VecE2, LaneTag}}()
    lane_prev_dict = Dict{LaneTag, Tuple{VecE2, LaneTag}}()

    for (tag, pts) in lane_pts_dict
        # see if can connect to next
        best_tag = NULL_LANETAG
        best_ind = -1
        best_sq_dist = dist_threshold_lane_connect
        for (tag2, pts2) in lane_pts_dict
            if tag2.segment != tag.segment
                for (ind,pt) in enumerate(pts2)
                    sq_dist = abs2(pt - pts[end])
                    if sq_dist < best_sq_dist
                        best_sq_dist = sq_dist
                        best_ind = ind
                        best_tag = tag2
                    end
                end
            end
        end
        if best_tag != NULL_LANETAG
            # remove our last pt and set next to pt to their pt
            pop!(pts)
            lane_next_dict[tag] = (lane_pts_dict[best_tag][best_ind], best_tag)
            if best_ind == 1 # set connect prev as well
                lane_prev_dict[best_tag] = (pts[end], tag)
            end
        end
    end
    for (tag, pts) in lane_pts_dict
        # see if can connect to prev
        if !haskey(lane_prev_dict, tag)
            best_tag = NULL_LANETAG
            best_ind = -1
            best_sq_dist = dist_threshold_lane_connect
            for (tag2, pts2) in lane_pts_dict
                if tag2.segment != tag.segment
                    for (ind,pt) in enumerate(pts2)
                        sq_dist = abs2(pt - pts[1])
                        if sq_dist < best_sq_dist
                            best_sq_dist = sq_dist
                            best_ind = ind
                            best_tag = tag2
                        end
                    end
                end
            end
            if best_tag != NULL_LANETAG
                # connect 'em
                shift!(pts)
                lane_prev_dict[tag] = (lane_pts_dict[best_tag][best_ind], best_tag)
            end
        end
    end

    ###################################################
    # Build the roadway
    retval = Roadway()
    for (tag, pts) in lane_pts_dict
        if !has_segment(retval, tag.segment)
            push!(retval.segments, RoadSegment(tag.segment))
        end
    end

    lane_new_dict = Dict{LaneTag, LaneTag}() # old -> new tag
    for seg in retval.segments

        # pull lanetags for this seg
        lanetags = LaneTag[]
        for tag in keys(lane_pts_dict)
            if tag.segment == seg.id
                push!(lanetags, tag)
            end
        end

        # sort the lanes such that the rightmost lane is lane 1
        # do this by taking the first lane,
        # then project each lane's midpoint to the perpendicular at the midpoint

        @assert(!isempty(lanetags))
        proj_positions = Array{Float64}(length(lanetags))

        first_lane_pts = lane_pts_dict[lanetags[1]]
        n = length(first_lane_pts)
        lo = first_lane_pts[div(n,2)]
        hi = first_lane_pts[div(n,2)+1]
        midpt_orig = (lo + hi)/2
        dir = polar(1.0, atan2(hi - lo) + π/2) # direction perpendicular (left) of lane

        for (i,tag) in enumerate(lanetags)
            pts = lane_pts_dict[tag]
            n = length(pts)
            midpt = (pts[div(n,2)] + pts[div(n,2)+1])/2
            proj_positions[i] = proj(midpt - midpt_orig, dir, Float64)
        end

        for (i,j) in enumerate(sortperm(proj_positions))

            tag = lanetags[j]
            boundary_left = i == length(proj_positions) ? LaneBoundary(:solid, :white) : LaneBoundary(:broken, :white)
            boundary_right = i == 1 ? LaneBoundary(:solid, :white) : LaneBoundary(:broken, :white)

            pts = lane_pts_dict[tag]
            pt_matrix = Array{Float64}(2, length(pts))
            for (k,P) in enumerate(pts)
                pt_matrix[1,k] = P.x
                pt_matrix[2,k] = P.y
            end

            println("fitting curve ", length(pts), "  "); tic()
            curve = _fit_curve(pt_matrix, desired_distance_between_curve_samples)
            toc()

            tag_new = LaneTag(seg.id, length(seg.lanes)+1)
            lane = Lane(tag_new, curve,
                        boundary_left = boundary_left,
                        boundary_right = boundary_right)
            push!(seg.lanes, lane)
            lane_new_dict[tag] = tag_new
        end
    end

    ###################################################
    # Connect the lanes
    for (tag_old, tup) in lane_next_dict
        next_pt, next_tag_old = tup
        lane = retval[lane_new_dict[tag_old]]
        next_tag_new = lane_new_dict[next_tag_old]
        dest = retval[next_tag_new]
        roadproj = proj(VecSE2(next_pt, 0.0), dest, retval)

        # println("connecting $(lane.tag) to $(dest.tag)")

        cindS = curveindex_end(lane.curve)
        cindD = roadproj.curveproj.ind

        if cindD == CURVEINDEX_START # a standard connection
            connect!(lane, dest)
            # remove any similar connection from lane_prev_dict
            if haskey(lane_prev_dict, next_tag_old) && lane_prev_dict[next_tag_old][2] == tag_old
                delete!(lane_prev_dict, next_tag_old)
            end
        else
            # otherwise connect as before
            unshift!(lane.exits,  LaneConnection(true,  cindS, RoadIndex(cindD, dest.tag)))
            push!(dest.entrances, LaneConnection(false, cindD, RoadIndex(cindS, lane.tag)))
        end
    end
    for (tag_old, tup) in lane_prev_dict
        prev_pt, prev_tag_old = tup
        lane = retval[lane_new_dict[tag_old]]
        prev_tag_new = lane_new_dict[prev_tag_old]
        prev = retval[prev_tag_new]
        roadproj = proj(VecSE2(prev_pt, 0.0), prev, retval)

        # println("connecting $(lane.tag) from $(prev.tag)")

        cindS = roadproj.curveproj.ind
        cindD = CURVEINDEX_START

        if cindS == curveindex_end(prev) # a standard connection
            @assert(!has_prev(prev))
            connect!(prev, lane)
        else
            # a standard connection
            push!(prev.exits,  LaneConnection(true,  cindS, RoadIndex(cindD, lane.tag)))
            unshift!(lane.entrances, LaneConnection(false, cindD, RoadIndex(cindS, prev.tag)))
        end
    end

    retval
end


## Generative Merging MDP model

const HARD_BRAKE = 1
const RELEASE = 7
const WRAP_AROUND_TOL = 2.0 # when vehicles reach the end of the main lane - WRAP_AROUND_TOL they are spawned at the beginning

"""
    AugScene
Driving scene augmented with information about the ego vehicle
"""
struct AugScene
    scene::Scene
    ego_info::NamedTuple{(:acc,), Tuple{Float64}}
end

# Dummy type since the action is given by the policy
mutable struct EgoDriver{A} <: DriverModel{A}
    a::A
end

Base.rand(model::EgoDriver) = model.a
AutomotiveDrivingModels.observe!(m::EgoDriver, s::EntityFrame{S,D,I}, roadway::R, egoid::Int64) where {S,D,I,R} = m

"""
    GenerativeMergingMDP

A simulation environment for a highway merging scenario. Implemented using POMDPs.jl

# Parameters
    - `env::MergingEnvironment = MergingEnvironment(main_lane_angle = 0.0, merge_lane_angle = pi/7)`
    - `n_cars_main::Int64 = 1`
    - `n_cars_merge::Int64 = 1`
    - `n_agents::Int64 = n_cars_main + n_cars_merge`
    - `max_cars::Int64 = 16`
    - `min_cars::Int64 = 0`
    - `car_def::VehicleDef = VehicleDef()`
    - `dt::Float64 = 0.5 # time step`
    - `jerk_levels::SVector{5, Float64} = SVector(-1, -0.5, 0, 0.5, 1.0)`
    - `accel_levels::SVector{6, Float64} = SVector(-4.0, -2.0, -1.0, 0.0, 1.0, 2.0)`
    - `max_deceleration::Float64 = -4.0`
    - `max_acceleration::Float64 = 3.5`
    - `comfortable_acceleration::Float64 = 2.0`
    - `discount_factor::Float64 = 0.95`
    - `ego_idm::IntelligentDriverModel = IntelligentDriverModel(σ=0.0, v_des=env.main_lane_vmax)`
    - `default_driver_model::DriverModel{LaneFollowingAccel} = IntelligentDriverModel(v_des=env.main_lane_vmax)`
    - `observe_cooperation::Bool = false`
    - `observe_speed::Bool = true`
    - `traffic_speed::Symbol = :mixed`
    - `random_n_cars::Bool = false`
    - `driver_type::Symbol = :random` can be :binary, or :aggressive, or :cooperative
    - `max_burn_in::Int64 = 20`
    - `min_burn_in::Int64 = 10`
    - `initial_ego_velocity::Float64 = 10.0`
    - `initial_velocity::Float64 = 5.0`
    - `initial_velocity_std::Float64 = 1.0`
    - `main_lane_slots::LinRange{Float64} = LinRange(0.0, env.main_lane_length + env.after_merge_length, max_cars)`
    - `collision_cost::Float64 = -1.0`
    - `goal_reward::Float64 = 1.0`
    - `hard_brake_cost::Float64 = 0.0`
"""
@with_kw mutable struct GenerativeMergingMDP <: MDP{AugScene, Int64}
    env::MergingEnvironment = MergingEnvironment(main_lane_angle = 0.0, merge_lane_angle = pi/7)
    n_cars_main::Int64 = 1
    n_cars_merge::Int64 = 1
    n_agents::Int64 = n_cars_main + n_cars_merge
    max_cars::Int64 = 16
    min_cars::Int64 = 0
    car_def::VehicleDef = VehicleDef()
    dt::Float64 = 0.5 # time step
    jerk_levels::SVector{5, Float64} = SVector(-1, -0.5, 0, 0.5, 1.0)
    accel_levels::SVector{6, Float64} = SVector(-4.0, -2.0, -1.0, 0.0, 1.0, 2.0)
    max_deceleration::Float64 = -4.0
    max_acceleration::Float64 = 3.5
    comfortable_acceleration::Float64 = 2.0
    discount_factor::Float64 = 0.95
    ego_idm::IntelligentDriverModel = IntelligentDriverModel(σ=0.0, v_des=env.main_lane_vmax)
    default_driver_model::DriverModel{LaneFollowingAccel} = IntelligentDriverModel(v_des=env.main_lane_vmax)
    observe_cooperation::Bool = false
    observe_speed::Bool = true
    # initial state params
    traffic_speed::Symbol = :mixed
    random_n_cars::Bool = false
    driver_type::Symbol = :random
    max_burn_in::Int64 = 20
    min_burn_in::Int64 = 10
    initial_ego_velocity::Float64 = 10.0
    initial_velocity::Float64 = 5.0
    initial_velocity_std::Float64 = 1.0
    main_lane_slots::LinRange{Float64} = LinRange(0.0,
                                            env.main_lane_length + env.after_merge_length,
                                            max_cars)
    # reward params
    collision_cost::Float64 = -1.0
    goal_reward::Float64 = 1.0
    hard_brake_cost::Float64 = 0.0

    # internal states
    driver_models::Dict{Int64, DriverModel} = Dict{Int64, DriverModel}(EGO_ID=>EgoDriver(LaneFollowingAccel(0.0)))
end

POMDPs.discount(mdp::GenerativeMergingMDP) = mdp.discount_factor
POMDPs.actions(mdp::GenerativeMergingMDP) = 1:7
POMDPs.n_actions(mdp::GenerativeMergingMDP) = 7
POMDPs.actionindex(mdp::GenerativeMergingMDP, a::Int64) = a

function POMDPs.initialstate(mdp::GenerativeMergingMDP, rng::AbstractRNG)
    s0 = Scene()
    if mdp.random_n_cars
        mdp.n_cars_main = rand(rng, mdp.min_cars:mdp.max_cars)
    end
    mdp.driver_models = Dict{Int64, DriverModel}(EGO_ID=>EgoDriver(LaneFollowingAccel(0.0)))
    start_positions = sample(rng, mdp.main_lane_slots, mdp.n_cars_main, replace=false)

    start_velocities = mdp.initial_velocity .+ mdp.initial_velocity_std*randn(rng, mdp.n_cars_main)
    ego = initial_merge_car_state(mdp, rng, EGO_ID)
    ego_acc_0 = 0.0
    scene = s0
    for i=EGO_ID+1:EGO_ID+mdp.n_cars_main
        veh_state = vehicle_state(start_positions[i - EGO_ID], main_lane(mdp.env),
        start_velocities[i - EGO_ID], mdp.env.roadway)
        veh = Vehicle(veh_state, mdp.car_def, i)
        push!(s0, veh)
        mdp.driver_models[i] = CooperativeIDM()
        if mdp.traffic_speed == :mixed
            v_des = sample(rng, [4.0, 5., 6.0], Weights([0.2, 0.3, 0.5]))
        elseif mdp.traffic_speed == :fast
            v_des = 15.0
        end
        set_desired_speed!(mdp.driver_models[i], v_des)
        if mdp.driver_type == :random
            mdp.driver_models[i].c = rand(rng, 0:0.01:1) # change cooperativity
        elseif mdp.driver_type == :binary
            mdp.driver_models[i].c = sample(rng, [0,1], Weights([0.9, 0.1])) # change cooperativity
        elseif mdp.driver_type == :aggressive
            mdp.driver_models[i].c = 0.0
        elseif mdp.driver_type == :cooperative
            mdp.driver_models[i].c = 1.0
        end
    end
    # burn in
    acts = Vector{LaneFollowingAccel}(undef, length(scene))
    burn_in =rand(rng, mdp.min_burn_in:mdp.max_burn_in)
    scene = s0
    for t=1:burn_in
        get_actions!(acts, scene, mdp.env.roadway, mdp.driver_models)
        tick!(scene, mdp.env.roadway, acts, mdp.dt, true)
        for (i, veh) in enumerate(scene)
            scene[i] = wrap_around(mdp.env, scene[i])
        end
    end
    s = Scene()
    push!(s, ego)
    for veh in scene
        push!(s, veh)
    end
    return AugScene(s, (acc=ego_acc_0,))
end

function POMDPs.reward(mdp::GenerativeMergingMDP, s::AugScene, a::Int64, sp::AugScene)
    egop = get_by_id(sp.scene, EGO_ID)
    r = 0.0
    if reachgoal(mdp, egop)
       r += mdp.goal_reward
    elseif collision_checker(sp.scene, EGO_ID)
        r += mdp.collision_cost
    end
    if caused_hard_brake(mdp, sp.scene)
        r += mdp.hard_brake_cost
    end
    return r
end

function POMDPs.reward(pomdp::FullyObservablePOMDP{AugScene,Int64}, s::AugScene, a::Int64, sp::AugScene)
    return reward(pomdp.mdp, s, a, sp)
end

function POMDPs.isterminal(mdp::GenerativeMergingMDP, s::AugScene)
    return collision_checker(s.scene, EGO_ID) || reachgoal(mdp, get_by_id(s.scene, EGO_ID))
end

function POMDPs.gen(mdp::GenerativeMergingMDP, s::AugScene, a::Int64, rng::AbstractRNG)
    scene = deepcopy(s.scene)
    mdp.driver_models[EGO_ID].a = action_map(mdp, s.ego_info.acc, a)
    ego_acc = mdp.driver_models[EGO_ID].a.a
    acts = Vector{LaneFollowingAccel}(undef, length(scene))

    # call driver models
    for i=EGO_ID+1:EGO_ID+mdp.n_cars_main
        mdp.driver_models[i].other_acc = s.ego_info.acc
    end
    get_actions!(acts, scene, mdp.env.roadway, mdp.driver_models)

    # update scene
    tick!(scene, mdp.env.roadway, acts, mdp.dt, true)

    # clamp speed
    for (i, veh) in enumerate(scene)
        # scene[i] = clamp_speed(mdp.env, veh)
        scene[i] = wrap_around(mdp.env, scene[i])
    end

    return (sp=AugScene(scene, (acc=ego_acc,)), )
end

"""
    extract_features(mdp::GenerativeMergingMDP, s::AugScene)
extract a feature vector from AugScene
"""
function extract_features(mdp::GenerativeMergingMDP, s::AugScene)
    # @show keys(mdp.driver_models)
    scene = s.scene
    env = mdp.env
    features = -3*ones(15)
    ego_ind = findfirst(EGO_ID, scene)
    ego = scene[ego_ind]

    # distance to merge point
    s_ego = dist_to_merge(env, ego)
    v_ego = ego.state.v
    features[1] = s_ego
    features[2] = v_ego
    features[3] = s.ego_info.acc
    # get neighbors
    fore, merge, fore_main, rear_main = get_neighbors(env, scene, EGO_ID)

    # front neighbor
    features[6] = 0.0
    if fore.ind != nothing
        fore_id = scene[fore.ind].id
        v_oth = scene[fore.ind].state.v
        headway = fore.Δs
        features[4] = headway
        if mdp.observe_speed
            features[5] = v_oth
        else
            features[5] = 0.0
        end
        if mdp.observe_cooperation
            features[6] = mdp.driver_models[fore_id].c
        else
            features[6] = 0.
        end
    end

    # two closest cars in main lane
    features[9] = 0.0
    if fore_main.ind != nothing
        fore_main_id = scene[fore_main.ind].id
        v_oth_main_fore = scene[fore_main.ind].state.v
        headway_main_fore = fore_main.Δs
        features[7] = headway_main_fore
        if mdp.observe_speed
            features[8] = v_oth_main_fore
        else
            features[8] = 0.0
        end
        if mdp.observe_cooperation
            features[9] = mdp.driver_models[fore_main_id].c
        else
            features[9] = 0.
        end
    end
    features[12] = 0.0
    if rear_main.ind != nothing
        rear_main_id = scene[rear_main.ind].id
        v_oth_main_rear = scene[rear_main.ind].state.v
        headway_main_rear = rear_main.Δs
        features[10] = headway_main_rear
        if mdp.observe_speed
            features[11] = v_oth_main_rear
        else
            features[11] = 0.0
        end
        if mdp.observe_cooperation
            features[12] = mdp.driver_models[rear_main_id].c
        else
            features[12] = 0.
        end
    end

    # rear car to the merge point
    features[15] = 0.0
    if merge.ind != nothing
        merge_id = scene[merge.ind].id
        v_oth = scene[merge.ind].state.v
        headway = merge.Δs
        features[13] = headway
        if mdp.observe_speed
            features[14] = v_oth
        else
            features[14] = 0.0
        end
        if mdp.observe_cooperation
            features[15] = mdp.driver_models[merge_id].c
        else
            features[15] = 0.
        end
    end
    return features
end

"""
    normalize_features!(mdp::GenerativeMergingMDP, features::Vector{Float64})
normalize a feature vector extracted from a scene
"""
function normalize_features!(mdp::GenerativeMergingMDP, features::Vector{Float64})
    features[1] /=  mdp.env.main_lane_length
    features[2] /=  mdp.env.main_lane_vmax
    features[3] /=   mdp.max_deceleration
    features[4] /= mdp.env.main_lane_length
    features[5] /=  mdp.env.main_lane_vmax
    features[6] /= 1.
    features[7] /= mdp.env.main_lane_length
    features[8] /=   mdp.env.main_lane_vmax
    features[9] /= 1.
    features[10] /= mdp.env.main_lane_length
    features[11] /=  mdp.env.main_lane_vmax
    features[12] /= 1.
    features[13] /= mdp.env.main_lane_length
    features[14] /=  mdp.env.main_lane_vmax
    features[15] /= 1.

    return features
end

"""
     unnormalize_features!(mdp::GenerativeMergingMDP, features::Vector{Float64})
rescale feature vector
"""
function unnormalize_features!(mdp::GenerativeMergingMDP, features::Vector{Float64})
    features[1]  *=  mdp.env.main_lane_length
    features[2]  *=  mdp.env.main_lane_vmax
    features[3]  *=   mdp.max_deceleration
    features[4]  *= mdp.env.main_lane_length
    features[5]  *=  mdp.env.main_lane_vmax
    features[6]  *= 1.
    features[7]  *= mdp.env.main_lane_length
    features[8]  *=   mdp.env.main_lane_vmax
    features[9]  *= 1.
    features[10] *= mdp.env.main_lane_length
    features[11] *=  mdp.env.main_lane_vmax
    features[12] *= 1.
    features[13] *= mdp.env.main_lane_length
    features[14] *=  mdp.env.main_lane_vmax
    features[15] *= 1.

    return features
end

function POMDPs.convert_s(::Type{V}, s::AugScene, mdp::GenerativeMergingMDP) where V<:AbstractArray
    feature_vec = extract_features(mdp, s)
    return normalize_features!(mdp, feature_vec)
end

function POMDPs.convert_s(::Type{AugScene}, o::V, mdp::GenerativeMergingMDP) where V<:AbstractArray
    feature_vec = deepcopy(o)
    unnormalize_features!(mdp, feature_vec)
    s_ego, v_ego, s_front, v_front, s_fm, v_fm, s_rm, v_rm, s_m, v_m, a_ego = feature_vec[1:11]

    # reconstruct the scene
    scene = Scene()
    if s_ego < 0.0
        lane_ego = merge_lane(mdp.env)
        s_ego = get_end(lane_ego) + s_ego
    else
        lane_ego = main_lane(mdp.env)
        s_ego = mdp.env.roadway[mdp.env.merge_index].s + s_ego
    end
    ego = Vehicle(vehicle_state(s_ego, lane_ego, v_ego, mdp.env.roadway), VehicleDef(), EGO_ID)
    push!(scene, ego)

    # front neighbor
    if lane_ego == main_lane(mdp.env)
        s_front = ego.state.posF.s + s_front
    else
        s_front = mdp.env.main_lane_length + s_front
    end
    front = Vehicle(vehicle_state(s_front, main_lane(mdp.env), v_front, mdp.env.roadway), VehicleDef(), EGO_ID+1)
    push!(scene, front)

    # front neighbor to projection
    proj_lane = main_lane(mdp.env)
    main_lane_proj = proj(ego.state.posG, proj_lane, mdp.env.roadway)
    s_main = proj_lane[main_lane_proj.curveproj.ind, mdp.env.roadway].s
    s_fm = s_main + s_fm
    frontmain = Vehicle(vehicle_state(s_fm, main_lane(mdp.env), v_fm, mdp.env.roadway), VehicleDef(), EGO_ID+2)
    push!(scene, frontmain)

    s_rm = s_main - s_rm
    rearmain = Vehicle(vehicle_state(s_rm, main_lane(mdp.env), v_rm, mdp.env.roadway), VehicleDef(), EGO_ID+3)
    push!(scene, rearmain)

    if lane_ego == main_lane(mdp.env)
        s_m = ego.state.posF.s - s_m
    else
        s_m = mdp.env.main_lane_length - s_m
    end
    merge = Vehicle(vehicle_state(s_m, main_lane(mdp.env), v_m, mdp.env.roadway), VehicleDef(), EGO_ID+4)
    push!(scene, merge)
    return AugScene(scene, (acc=a_ego,))
end

## helpers

"""
    initial_merge_car_state(mdp::GenerativeMergingMDP, rng::AbstractRNG, id::Int64)
returns a Vehicle, at the initial state of the merging car.
"""
function initial_merge_car_state(mdp::GenerativeMergingMDP, rng::AbstractRNG, id::Int64)
    v0 = mdp.initial_velocity + mdp.initial_velocity_std*randn(rng)
    v0 = mdp.initial_ego_velocity
    veh_state = vehicle_state(0.0, merge_lane(mdp.env), v0, mdp.env.roadway)
    return Vehicle(veh_state, mdp.car_def, id)
end

"""
    reset_main_car_state(mdp::GenerativeMergingMDP, veh::Vehicle)
initialize a car at the beginning of the main lane
"""
function reset_main_car_state(mdp::GenerativeMergingMDP, veh::Vehicle)
    v0 = mdp.initial_velocity + mdp.initial_velocity_std*randn(rng)
    veh_state = vehicle_state(0.0, main_lane(mdp.env), v0, mdp.env.roadway)
    return Vehicle(veh_state, mdp.car_def, veh.id)
end

"""
    reachgoal(mdp::GenerativeMergingMDP, ego::Vehicle)
return true if `ego` reached the goal position
"""
function reachgoal(mdp::GenerativeMergingMDP, ego::Vehicle)
    lane = get_lane(mdp.env.roadway, ego)
    s = ego.state.posF.s
    return lane.tag == main_lane(mdp.env).tag && s >= get_end(lane)
end

"""
    caused_hard_brake(mdp::GenerativeMergingMDP, scene::Scene)
returns true if the ego vehicle caused its rear neighbor to hard brake
"""
function caused_hard_brake(mdp::GenerativeMergingMDP, scene::Scene)
    ego_ind = findfirst(EGO_ID, scene)
    fore_res = get_neighbor_rear_along_lane(scene, ego_ind, mdp.env.roadway)
    if fore_res.ind == nothing
        return false
    else
        return mdp.driver_models[fore_res.ind].a <= mdp.driver_models[fore_res.ind].idm.d_max
    end
end

"""
    action_map(mdp::GenerativeMergingMDP, acc::Float64, a::Int64)
maps integer to a LaneFollowingAccel
"""
function action_map(mdp::GenerativeMergingMDP, acc::Float64, a::Int64)
    if a == HARD_BRAKE
        return LaneFollowingAccel(mdp.max_deceleration)
    elseif a == RELEASE
        return LaneFollowingAccel(0.0)
    else
        return LaneFollowingAccel(clamp(acc + mdp.jerk_levels[a-1], mdp.max_deceleration, mdp.max_acceleration))
    end
end

"""
    vehicle_state(s::Float64, lane::Lane, v::Float64, roadway::Roadway)
convenient constructor for VehicleState
"""
function vehicle_state(s::Float64, lane::Lane, v::Float64, roadway::Roadway)
    posF = Frenet(lane, s)
    return VehicleState(posF, roadway, v)
end

"""
    wrap_around(env::MergingEnvironment, veh::Vehicle)
respawn vehicle at the beginning of the main lane
"""
function wrap_around(env::MergingEnvironment, veh::Vehicle)
    lane = get_lane(env.roadway, veh)
    s_end = get_end(lane)
    s = veh.state.posF.s
    if s >= s_end - WRAP_AROUND_TOL && lane == main_lane(env) && veh.id != EGO_ID
        veh_state = vehicle_state(0.0, main_lane(env), veh.state.v, env.roadway)
        return Vehicle(veh_state,veh.def, veh.id)
    end
    return veh
end

"""
    clamp_speed(env::MergingEnvironment, veh::Vehicle)
clamp the speed of `veh` between 0. and main lane vmax
"""
function clamp_speed(env::MergingEnvironment, veh::Vehicle)
    v = clamp(veh.state.v, 0.0, env.main_lane_vmax)
    vehstate = VehicleState(veh.state.posG, veh.state.posF, v)
    return Vehicle(vehstate, veh.def, veh.id)
end

"""
    spread_out_initialization(mdp::GenerativeMergingMDP, rng::AbstractRNG)
spread out vehicles on the main lane
"""
function spread_out_initialization(mdp::GenerativeMergingMDP, rng::AbstractRNG)
    start_positions = zeros(mdp.n_cars_main)
    start_positions[1] = rand(rng, mdp.main_lane_slots)
    gap_length = div(mdp.env.main_lane_length + mdp.env.after_merge_length, mdp.n_cars_main)
    main_roadway = StraightRoadway(mdp.env.main_lane_length + mdp.env.after_merge_length)
    for i=2:mdp.n_cars_main
        start_positions[i] = mod_position_to_roadway(start_positions[i-1] + gap_length, main_roadway)
    end
    return start_positions
end

"""
    global_features(mdp::GenerativeMergingMDP, s::AugScene)
extract a vector with the states of all the vehicles in AugScene
"""
function global_features(mdp::GenerativeMergingMDP, s::AugScene)
    n_features = 2*mdp.n_cars_main + 3 + mdp.n_cars_main
    features = zeros(n_features)
    ego = get_by_id(s.scene, EGO_ID)
    features[1] = dist_to_merge(mdp.env, ego)
    features[2] = ego.state.v
    features[3] = s.ego_info.acc
    @assert s.scene[1].id == EGO_ID
    for i=2:length(s.scene)
        veh = s.scene[i]
        features[3*i-2] = veh.state.posF.s
        features[3*i-1] = veh.state.v
        obs_c = 0.5
        if mdp.observe_cooperation
            obs_c = mdp.driver_models[i].c
        end
        features[3*i] = obs_c
    end
    return features
end

function normalize_global_features!(mdp::GenerativeMergingMDP, features::Vector{Float64})
    features[1] /= mdp.env.main_lane_length
    features[2] /= mdp.env.main_lane_vmax
    features[3] /= mdp.max_deceleration
    for i=2:mdp.n_cars_main+1
        features[3*i - 2] /= mdp.env.main_lane_length
        features[3*i - 1] /= mdp.env.main_lane_vmax
    end
    return features
end

function unnormalize_global_features!(mdp::GenerativeMergingMDP, features::Vector{Float64})
    features[1] *= mdp.env.main_lane_length
    features[2] *= mdp.env.main_lane_vmax
    features[3] *= mdp.max_deceleration
    for i=2:mdp.n_cars_main+1
        features[3*i-2] *= mdp.env.main_lane_length
        features[3*i-1] *= mdp.env.main_lane_vmax
    end
    return features
end


"""
    LMPTools

Convenience and utility functions for LongwaveModePropagator.jl.
"""
module LMPTools

using PyCall, Conda
using Rotations, StaticArrays
using HDF5, Dates
using GeographicLib, SatelliteToolbox
using LongwaveModePropagator

project_path(parts...) = normpath(@__DIR__, "..", parts...)

const GROUND_DATA = h5read(project_path("data", "conductivity_data.h5"), "/")

const TRANSMITTER = Dict(
    :NAA => Transmitter("NAA", 44.646394, -67.281069, 24e3),
    :NAU => Transmitter("NAU", 18.39875, -67.177433, 40.75e3),
    :NML => Transmitter("NML", 46.366014, -98.335544, 25.2e3),
    :NLK => Transmitter("NLK", 48.203611, -121.917056, 24.8e3),
    :NPM => Transmitter("NPM", 21.420, -158.151, 21.4e3)
)

const CHAOS = PyNULL()
const CHAOS_MODEL = PyNULL()

export TRANSMITTER
export range
export get_ground, get_groundcode, get_epsilon, get_sigma, groundsegments
export igrf, chaos, load_CHAOS_matfile
export zenithangle, isday, ferguson, flatlinearterminator, smoothterminator, fourierperturbation

function __init__()
    cp = try
        pyimport("chaosmagpy")
    catch
        # Install chaosmagpy
        Conda.add(["numpy", "scipy", "pandas", "cython", "cartopy", "matplotlib", "h5py"])

        Conda.pip_interop(true)
        Conda.pip("install", ["cdflib", "hdf5storage", "lxml"])
        Conda.pip("install", "chaosmagpy")

        pyimport("chaosmagpy")
    end

    chaosfile = newestchaos()
    @info "Loading $chaosfile"

    model = cp.load_CHAOS_matfile(joinpath(project_path("data"), chaosfile))
    copy!(CHAOS, cp)
    copy!(CHAOS_MODEL, model)
end

function newestchaos()
    datafiles = readdir(project_path("data"); sort=true)
    newestfile = ""
    for f in datafiles
        if occursin("CHAOS", f) && f > newestfile
            newestfile = f
        end
    end
    return newestfile
end

"""
    load_CHAOS_matfile(file)

Load a CHAOS model coefficient MATLAB `.mat` file for use by [`chaos`](@ref).

These files can be found at https://www.spacecenter.dk/files/magnetic-models/CHAOS-7/.

By default, LMPTools loads $(newestchaos()).
"""
function load_CHAOS_matfile(file)
    cp = pyimport("chaosmagpy")
    @info "Loading $file"
    model = cp.load_CHAOS_matfile(joinpath(project_path("data"), file))
    copy!(CHAOS, cp)
    copy!(CHAOS_MODEL, model)

    return nothing
end

###
# Extend library functions

function GeographicLib.GeodesicLine(tx::Transmitter, rx::Receiver)
    return GeodesicLine(tx.longitude, tx.latitude; lon2=rx.longitude, lat2=rx.latitude)
end

"""
    range(tx::Transmitter, rx::Receiver)

Return the range (great-circle distance, but on the WGS-84 ellipsoid) in meters between
`tx` and `rx`.
"""
function Base.range(tx::Transmitter, rx::Receiver)
    return inverse(tx.longitude, tx.latitude, rx.longitude, rx.latitude).dist
end

###
# Ground

function searchsortednearest(a, x)
    @assert issorted(a) "`a` must be sorted!"

    idx = searchsortedfirst(a, x)

    idx == 1 && return idx
    idx > length(a) && return length(a)
    a[idx] == x && return idx
    abs(a[idx]-x) < abs(a[idx-1]-x) ? (return idx) : (return idx-1)
end

"""
    get_ground(param, lat, lon)

Return the ground `param` "codemap", "epsilonmap", or "sigmamap" at position `lat`, `lon`
in degrees.

Convenience functions are available for the most common fields with arguments `lat`, `lon`
only:

    * [`get_groundcode`](@ref)
    * [`get_sigma`](@ref)
    * [`get_epsilon`](@ref)
"""
function get_ground(param, lat, lon)
    latidx = searchsortednearest(GROUND_DATA["lat"], lat)
    lonidx = searchsortednearest(GROUND_DATA["lon"], lon)

    return GROUND_DATA[param][latidx, lonidx]
end
get_groundcode(lat, lon) = get_ground("codemap", lat, lon)
get_sigma(lat, lon) = get_ground("sigmamap", lat, lon)
get_epsilon(lat, lon) = get_ground("epsilonmap", lat, lon)

"""
    groundsegments(tx::Transmitter, rx::Receiver; resolution=20e3, require_permittivity_change=false)

Return a `Vector` of `Ground`'s and a `Vector` of distances at which the ground changes
between `tx` and `rx` using path `resolution` in meters.

If `require_permittivity_change == true`, then if the relative permittivity doesn't change
we don't consider it a ground change. This helps reduce the number of short small changes
that occur during a long path which have the same permittivity but a small change in
conductivity that has a small influence on the electric field in the EIWG but comes a the
cost of additional (potentially nearly identical) segments in the waveguide model.
"""
function groundsegments(tx::Transmitter, rx::Receiver; resolution=20e3, require_permittivity_change=false)

    line = GeodesicLine(tx, rx)
    pts = waypoints(line, dist=resolution)

    lat, lon, dist = first(pts).lat, first(pts).lon, first(pts).dist
    lastground = GROUND[get_groundcode(lat, lon)]

    # initialize
    grounds = [lastground]
    distances = [dist]  # == 0.0

    for pt in pts
        lat, lon, dist = pt.lat, pt.lon, pt.dist
        ground = GROUND[get_groundcode(lat, lon)]

        if require_permittivity_change
            if ground.ϵᵣ != lastground.ϵᵣ
                push!(grounds, ground)
                push!(distances, dist)
                lastground = ground
            end
        elseif ground != lastground
            push!(grounds, ground)
            push!(distances, dist)
            lastground = ground
        end
    end

    return grounds, distances
end

###
# Magnetic field

"""
    igrf(geoaz, lat, lon, year; alt=60e3)

Return a `BField` from IGRF-13 for position (`lat`, `lon`)°  at fractional
`year`. `geoaz`° is the geodetic azimuth of the path from transmitter to receiver.
By default, the magnetic field at an altitude of 60,000 meters is returned,
but this can be overridden with the `alt` keyword argument.
"""
function igrf(geoaz, lat, lon, year; alt=60e3)
    n, e, v, t = igrf13syn(0, year, 1, alt/1000, 90-lat, mod(lon, 360))  # in nT!, @60km altitude

    # Rotate the nev frame to the propagation path xyz frame
    # negate az to correct rotation direction for downward pointing v
    R = RotXZ(π, -deg2rad(geoaz))
    Rd = R*SVector(n,e,v)

    return BField(t*1e-9, Rd[1]/t, Rd[2]/t, Rd[3]/t)
end

"""
    igrf(tx::Transmitter, rx::Receiver, year, dists; alt=60e3)

Return a `Vector{BField}` at each distance in `dists` in meters along the path from `tx`
to `rx` in fractional `year`.

If applying `igrf` to a single distance, it is recommended to use the `geoaz`, `lat`, `lon`
form of [`igrf`](@ref).
"""
function igrf(tx::Transmitter, rx::Receiver, year, dists; alt=60e3)
    line = GeodesicLine(tx, rx)
    geoaz = inverse(tx.longitude, tx.latitude, rx.longitude, rx.latitude).azi

    bfields = Vector{BField}(undef, length(dists))
    for (i, d) in enumerate(dists)
        lo, la, _ = forward(line, d)
        bfields[i] = igrf(geoaz, la, lo, year; alt=alt)
    end
    return bfields
end

"""
    chaos(geoaz, lat::Real, lon:Real, year; alt=60e3)

Return a `BField` from CHAOS-7 using internal and external sources for position
(`lat`, `lon`)°  at fractional `year`.

`geoaz`° is the geodetic azimuth of the path from transmitter to receiver.
By default, the magnetic field at an altitude of 60,000 meters is returned,
but this can be overridden with the `alt` keyword argument.

# References

[^1]: C. C. Finlay et al., “The CHAOS-7 geomagnetic field model and observed changes in the
    South Atlantic Anomaly,” Earth Planets Space, vol. 72, no. 1, Art. no. 1, Dec. 2020,
    doi: 10.1186/s40623-020-01252-9.
"""
function chaos(geoaz, lat::Number, lon::Number, year; alt=60e3)
    Re = 6371.2
    r = Re + alt/1000
    t = CHAOS.data_utils.dyear_to_mjd(year)  # MJD2000

    theta = 90 - lat  # geocentric co-lat (deg)
    phi = lon  # geocentric lon (deg)

    Br, Bt, Bp = CHAOS_MODEL(t, r, theta, phi)
    u, s, e = only(Br), only(Bt), only(Bp)

    # Rotate the "use" frame to the propagation path xyz frame
    # negate az to correct rotation direction for downward pointing v
    R = RotXZ(π, -deg2rad(geoaz))
    Rd = R*SVector(-s,e,-u)

    mag = hypot(u, s, e)
    return BField(mag*1e-9, Rd[1]/mag, Rd[2]/mag, Rd[3]/mag)
end

"""
    chaos(tx::Transmitter, rx::Receiver, year, dists; alt=60e3)

Return a `Vector{BField}` at each distance in `dists` in meters along the path from `tx`
to `rx` in fractional `year`.

If applying `chaos` to a single distance, it is recommended to use the `geoaz`, `lat`, `lon`
form of [`chaos`](@ref).
"""
function chaos(tx::Transmitter, rx::Receiver, year, dists; alt=60e3)
    line = GeodesicLine(tx, rx)
    geoaz = inverse(tx.longitude, tx.latitude, rx.longitude, rx.latitude).azi

    lats = Vector{Float64}(undef, length(dists))
    lons = similar(lats)
    for (i, d) in enumerate(dists)
        lo, la, _ = forward(line, d)
        lats[i] = la
        lons[i] = lo
    end

    bfields = chaos(geoaz, lats, lons, year; alt)

    return bfields
end

"""
    chaos(geoaz, lats, lons, year; alt=60e3)

Return a `Vector{BField}` at each latitude-longitude pair formed from each element of `lats`
and `lons` in fractional `year`.
"""
function chaos(geoaz, lats, lons, year; alt=60e3)
    length(lats) == length(lons) || throw(ArgumentError("length(lats) must equal length(lons)"))

    Re = 6371.2
    r = Re + alt/1000
    t = CHAOS.data_utils.dyear_to_mjd(year)  # MJD2000

    theta = 90 .- lats  # geocentric co-lat (deg)
    phi = lons  # geocentric lon (deg)

    u, s, e = CHAOS_MODEL(t, r, theta, phi)

    # Rotate the "use" frame to the propagation path xyz frame
    # negate az to correct rotation direction for downward pointing v
    R = RotXZ(π, -deg2rad(geoaz))

    bfields = Vector{BField}(undef, length(lats))
    for i in eachindex(bfields)
        Rd = R*SVector(-s[i],e[i],-u[i])

        mag = hypot(u[i], s[i], e[i])
        bfields[i] = BField(mag*1e-9, Rd[1]/mag, Rd[2]/mag, Rd[3]/mag)
    end

    return bfields
end

###
# Sun position

"""
    zenithangle(lat, lon, y, m, d, h, Δτ=nothing)

Return the solar zenith angle in degrees for year `y`, month `m`, day `d`, and hour `h` in
universal time (UT1). `Δτ` is the difference between terrestrial time and universal time.
The position in `lat`, `lon` is in degrees.

If `Δτ` is `nothing`, it will be computed by a linear extrapolation. Although this is not
accurate because `Δτ` is irregular, an error up to 30 seconds does not significantly affect
the result.

This function uses the calculation-optimized formulation of Algorithm 5 from

    Grena, R., “Five new algorithms for the computation of sun position from 2010 to 2110,”
    Solar Energy, 86, 2012, pp. 1323-1337. doi: 10.1016/j.solener.2012.01.024

which has a maximum error of 0.0027° (~10 arcsec) between year 2010 and 2110.

Refraction is not considered in this implementation.
"""
function zenithangle(lat, lon, yr::Int, mo::Int, dy::Int, hr::Real, Δτ=nothing)
    θ, ϕ = deg2rad(lon), deg2rad(lat)

    # Time scale computation
    if mo <= 2
        mo += 12
        yr -= 1
    end

    t = trunc(Int, 365.25*(yr - 2000)) + trunc(Int, 30.6001*(mo + 1)) -
        trunc(Int, 0.01*yr) + dy + 0.0416667*hr - 21958

    if isnothing(Δτ)
        Δτ = 96.4 + 0.00158*t
    end

    te = t + 1.1574e-5*Δτ

    # Heliocentric ecliptic longitude L
    ωa = 0.0172019715
    a = SVector(3.33024e-2, 3.512e-4, 5.2e-6)
    b = SVector(-2.0582e-3, -4.07e-5, -9e-7)
    s1, c1 = sincos(ωa*te)
    s2, c2 = 2*s1*c1, (c1 + s1)*(c1 - s1)
    s3, c3 = s2*c1 + c2*s1, c2*c1 - s2*s1
    s, c = SVector(s1, s2, s2), SVector(c1, c2, c3)

    β = 2.92e-5
    dβ = -8.23e-5
    ω = SVector(1.49e-3, 4.31e-3, 1.076e-2, 1.575e-2, 2.152e-2, 3.152e-2, 2.1277e-1)
    d = SVector(1.27e-5, 1.21e-5, 2.33e-5, 3.49e-5, 2.67e-5, 1.28e-5, 3.14e-5)
    φ = SVector(-2.337, 3.065, -1.533, -2.358, 0.074, 1.547, -0.488)

    L = 1.7527901 + 1.7202792159e-2*te
    for k = 1:3
        L += a[k]*s[k] + b[k]*c[k]
    end
    L += dβ*s1*sin(β*te)
    for i = 1:7
        L += d[i]*sin(ω[i]*te + φ[i])
    end

    # Nutation correction
    ωn = 9.282e-4
    ν = ωn*te - 0.8
    sν, cν = sincos(ν)
    Δλ = 8.34e-5*sν
    λ = L + π + Δλ
    ϵ = 4.089567e-1 - 6.19e-9*te + 4.46e-5*cν
    sλ, cλ = sincos(λ)
    sϵ, cϵ = sincos(ϵ)

    α = atan(sλ*cϵ, cλ)
    α < 0 && (α += 2π)  # [-π/2, π/2] → [0, 2π]

    δ = asin(sλ*sϵ)
    H = 1.7528311 + 6.300388099*t + θ - α + 0.92*Δλ
    H = mod2pi(H + π) - π  # hour angle

    # Calculate zenith and azimuth angle with parallax correction for topocentric coords
    sϕ, cϕ = sincos(ϕ)
    sδ, cδ = sincos(δ)
    cH = cos(H)
    # sH, cH = sincos(H)

    se0 = sϕ*sδ + cϕ*cδ*cH
    ep = asin(se0) - 4.26e-5*sqrt(1 - se0^2)

    # Γ = atan(sH, cH*sϕ - sδ*cϕ/cδ)  # azimuth

    # Refraction correction to elevation
    # Δre = 0.08422*P/(273 + T)/tan(ep + 0.003138/(ep + 0.08919))

    z = π/2 - ep  # - Δre  # zenith angle

    return rad2deg(z)
end

"""
    zenithangle(lat, lon, dt::DateTime, Δτ=nothing)

Return the solar zenith angle in degrees for universal `dt` with `Δτ` difference in seconds
between terrestrial time and universal time. The position `lat`, `lon` should be in degrees.
"""
function zenithangle(lat, lon, dt::DateTime, Δτ=nothing)
    y, m, d = yearmonthday(dt)
    h = hour(dt) + minute(dt)/60 + second(dt)/3600 + millisecond(dt)/3600_000

    return zenithangle(lat, lon, y, m, d, h, Δτ)
end

"""
    isday(sza)

Return `true` if solar zenith angle `sza` is less than 98°, otherwise return `false`.
"""
function isday(sza)
    abs(sza) < 98 ? true : false
end

"""
    ferguson(lat, sza, month_num::Int) → (h′, β)
    ferguson(lat, sza, dt::DateTime)

Ferguson ionosphere for geographic latitude `lat` in degrees, solar zenith angle `sza` in
degrees, and month number `month_num` or `DateTime` `dt`.

If the month term of the ionosphere fit is not desired, specify `month_num = 0`.

The higher order terms in [^Ferguson1980] (sunspot number and magnetic absorption) are ignored.

# References

[Ferguson1980]: J. A. Ferguson, "Ionospheric profiles for predicting nighttime VLF/LF
    propagation," NOSC/TR-530, 1980.
"""
function ferguson(lat, sza, month_num::Integer)
    lat = deg2rad(lat)
    sza = deg2rad(sza)

    cos_lat = cos(lat)
    cos_sza = cos(sza)

    month_scalar = cospi(2*(month_num-0.5)/12)

    hprime = 74.37 - 8.097*cos_sza + 5.779*cos_lat - 1.213*month_scalar
    beta = 0.3849 - 0.1658*cos_sza - 0.08584*cos_lat + 0.1296*month_scalar

    return hprime, beta
end
ferguson(lat, sza, dt::DateTime) = ferguson(lat, sza, month(dt))

"""
    flatlinearterminator(sza, hp_day=74, hp_night=86, b_day=0.3, b_night=0.5) → (h′, β)

A simple terminator model with flat day and night ionospheres and a linear transition
between `sza_min` and `sza_max`. Solar zenith angle `sza` is in degrees.
"""
function flatlinearterminator(sza, hp_day=74, hp_night=86, b_day=0.3, b_night=0.5;
    sza_min=90, sza_max=100)

    if sza_min < sza < sza_max
        # Linear trend between day and night
        m_hprime = (hp_night - hp_day) / (sza_max - sza_min)
        m_beta = (b_night - b_day) / (sza_max - sza_min)

        hprime = m_hprime*(sza - sza_min) + hp_day
        beta = m_beta*(sza - sza_min) + b_day
    elseif sza < sza_min
        hprime = hp_day
        beta = b_day
    else  # sza > sza_max
        hprime = hp_night
        beta = b_night
    end

    return hprime, beta
end

"""
    flatlinearterminator(sza::AbstractArray; hp_day=74, hp_night=86, b_day=0.3, b_night=0.5) → (h′, β)

Given an array of `sza`, return the corresponding arrays of h′s and βs.
"""
function flatlinearterminator(sza::AbstractArray, hp_day=74, hp_night=86, b_day=0.3, b_night=0.5;
    sza_min=90, sza_max=100)

    # Linear trend between day and night
    m_hprime = (hp_night - hp_day) / (sza_max - sza_min)
    m_beta = (b_night - b_day) / (sza_max - sza_min)

    hprimes = similar(sza, Float64)
    betas = similar(hprimes)
    for i in eachindex(sza)
        if sza_min < sza[i] < sza_max
            hprimes[i] = m_hprime*(sza[i] - sza_min) + hp_day
            betas[i] = m_beta*(sza[i] - sza_min) + b_day
        elseif sza[i] < sza_min
            hprimes[i] = hp_day
            betas[i] = b_day
        else  # sza > sza_max
            hprimes[i] = hp_night
            betas[i] = b_night
        end
    end

    return hprimes, betas
end

"""
    smoothterminator(sza, hp_day=74, hp_night=86, b_day=0.3, b_night=0.5;
        sza_min=90, sza_max=100, steepness=1) → (h′, β)

A smooth terminator model using a logistic curve from day to night ionospheres as a function
of solar zenith angle `sza`.

The `steepness` parameter `1` with the default values of `sza_min` and `sza_max`
approximately asymptotes to the day and night values by `sza_min` and `sza_max` and is
centered on their midpoint.
"""
function smoothterminator(sza, hp_day=74, hp_night=86, b_day=0.3, b_night=0.5;
    sza_min=90, sza_max=100, steepness=1)

    hp_night > hp_day || throw(ArgumentError("`hp_night` should be greater than `hp_day`."))
    b_night > b_day || throw(ArgumentError("`b_night` should be greater than `b_day`."))
    sza_max > sza_min || throw(ArgumentError("`sza_max` should be greater than `sza_min`."))

    hprime_range = hp_night - hp_day
    beta_range = b_night - b_day
    sza_mid = (sza_max - sza_min)/2 + sza_min

    # logistic curve
    hprime = hprime_range/(1 + exp(-steepness*(sza - sza_mid))) + hp_day
    beta = beta_range/(1 + exp(-steepness*(sza - sza_mid))) + b_day

    return hprime, beta
end

"""
    fourierperturbation(sza, coeff; a=deg2rad(45)/2)

Return the perturbation produced by the Fourier series:
```
pert = coeff[1] + coeff[2]*cos(sza*π/a) + coeff[3]*cos(sza*2π/a) + coeff[4]*sin(sza*π/a) + coeff[5]*sin(sza*2π/a)
```
for solar zenith angle `sza` in degrees.

This can be used to produce simulated ionospheres in conjunction with [`ferguson`](@ref).

# Examples
```julia-repl
julia> lat, lon = 45, -105
julia> dt = DateTime(2020, 1, 15);
julia> sza = zenithangle(lat, lon, dt);
julia> h, b = ferguson(lat, sza, dt);

julia> coeffH = (1.345, 0.668, -0.177, -0.248, 0.096);
julia> hpert = fourierperturbation(sza, coeffH)
1.7991484480844766

julia> h + hpert
79.5769471288202
```
"""
function fourierperturbation(sza, coeff; a=deg2rad(45)/2)
    sza = deg2rad(sza)

    pert = coeff[1] +
        coeff[2]*cospi(sza/a) +
        coeff[4]*sinpi(sza/a) +
        coeff[3]*cospi(sza*2/a) +
        coeff[5]*sinpi(sza*2/a)

    return pert
end

end  # module


using Recombinase, Plots
##
args, kwargs = Recombinase.series2D(
    Recombinase.discrete(FlippingModel.survival),
    b,
    error = :MouseID,
    Recombinase.Group(color = :Protocol, linestyle = :Wall),
    select = (:Omissions_plus_one, :Leave),
    ribbon = true
    )
p1 = plot(args...;
    kwargs...,
    legend = false,
    grid = false,
    thickness_scaling = 2.5,
    tick_orientation = :out,
    size = (809,500),
    linewidth=2,
    fillalpha= 0,
    xlims = (0,16),
    xticks = 1:2:16)
##
savefig(p1,"/Volumes/GoogleDrive/My Drive/Reports for Zach/NewProtocols/Fitting/Survival_from_data.pdf")
##
figures = "/Users/dariosarra/Documents/Lab/Mainen/Presentations/Walton_visit/figures/unprocessed_plots"
savefig(p1,joinpath(figures,"New_data_survival.pdf"))
##
Simulated[!, :Leave] .= true
args, kwargs = series2D(
    Recombinase.discrete(FlippingModel.survival),
    Simulated,
    error = :MouseID,
    Recombinase.Group(color = :Protocol, linestyle = :Wall),
    select = (:simulation, :Leave),
    ribbon = true
    )
p2 = plot!(p1,args...;kwargs...,
    legend = false,
    fillalpha = 0.25,linealpha=0)
##
savefig(p1,"/Volumes/GoogleDrive/My Drive/Reports for Zach/NewProtocols/Fitting/Survival_from_fit.pdf")
##
args, kwargs = series2D(
    Recombinase.discrete(FlippingModel.survival),
    b,
    error = :MouseID,
    Recombinase.Group(color = :Protocol, linestyle = :Wall),
    select = (:Omissions_plus_one, :Leave),
    ribbon = true
    )
p2 = plot!(p1,args...;kwargs...,
    legend = false
    ,fillalpha = 0,linewidth=2)
##
savefig(p1,"/Volumes/GoogleDrive/My Drive/Reports for Zach/NewProtocols/Fitting/Survival_fit.pdf")
##
figures = "/Users/dariosarra/Documents/Lab/Mainen/Presentations/Walton_visit/figures/unprocessed_plots"
savefig(p1,joinpath(figures,"New_data_survival_fit.pdf"))
###
s = filter(streaks) do row
    row.Correct_start &&
    row.Streak > 10 &&
    3 < row.Streak_within_Block< 23 &&
    row.ReverseStreak > 20 &&
    (ismissing(row.Pre_Interpoke) || row.Pre_Interpoke < 1) &&
    row.Protocol_Day > 1 &&
    row.Day > 20191215
end
s[!,:AftrePlusOne] = s.AfterLast .+ 1
args, kwargs = series2D(
    Recombinase.discrete(Recombinase.density),
    s,
    error = :MouseID,
    Recombinase.Group(color = :Protocol, linestyle = :Wall),
    select = (:AftrePlusOne),
    ribbon = true
    )

p3 = plot!(p1,args...;kwargs...,legend = false,fillalpha = 0,linewidth=2)

plot!(p3,args...;kwargs...,)


__precompile__()
"""
An outline of the API for ACTRModels is presented below. In the REPL, use `?` to see documentation,
    i.e.,

````julia
? print_memory

to see documentation for `print_memory`.
````
# API
## Types
- `ACTR`
- `Declarative`
- `Goal`
- `Imaginal`
- `Chunk`
- `Visual`
- `VisualLocation`
- `Motor`
- `Procedural`
- `Rule`
## Functions
- `get_chunks`
- `get_iconic_memory`
- `get_visicon`
- `get_buffer`
- `set_buffer!`
- `update_lags!`
- `update_recent!`
- `update_chunk!`
- `modify!`
- `add_chunk!`
- `retrieval_prob`
- `retrieval_probs`
- `retrieve`
- `compute_activation!`
- `get_parm`
- `spreading_activation`
- `match`
- `compute_RT`
- `retrieval_request`
- `get_subset`
- `first_chunk`
- `posterior_predictive`
- `find_index`
- `find_indices`
- `import_printing`
- `print_chunk`
- `print_memory`
-  `run`
- `get_mean_activations`
"""
module ACTRModels
    using Reexport
    @reexport using Distributions, Parameters, Random, StatsBase, StatsFuns
    using  ConcreteStructs, PrettyTables
    import Distributions: pdf, logpdf
    import SequentialSamplingModels: LNR
    import Base: rand, match
    export AbstractACTR, ACTR, Declarative, Imaginal, Chunk, BufferState, Mod
    export Goal, Visual, Motor, VisualLocation, Procedural, Rule
    export AbstractVisualObject, VisualObject, Parms, AbstractParms
    export default_penalty, LNR, reduce_data, get_buffer, set_buffer!
    export get_chunks, update_lags!, update_recent!, update_chunk!, modify!, add_chunk!
    export retrieval_prob, retrieval_probs, retrieve, compute_activation!, get_parm
    export spreading_activation!, match, compute_RT, retrieval_request, get_subset
    export first_chunk, posterior_predictive, find_index, find_indices, get_mean_activations
    export get_visicon, get_iconic_memory, get_rules
    export get_time, set_time!, add_time!, reset_time!, rnd_time, get_name
    export blend_chunks, blended_activation

    include("Structs.jl")
    include("Utilities/Utilities.jl")
    include("MemoryFunctions.jl")
end


using Revise
import LPDM, D3Trees, LightDarkLPDM, execute2d
# execute2d.execute(solver=:lpdm)
steps = 1
# execute2d.execute(output=2, steps = steps, reward_func = :quadratic)
# execute2d.execute(vis=Int64[1], solv_mode=:lpdm, reward_func=:quadratic, action_space_type=:adaptive, output=1, steps = steps)
execute2d.execute(s0=Vec2(3/2*π,3/2*π), solver_mode = :lpdm, action_mode = :standard, obs_mode = :continuous, steps=steps, output=3)
# execute2d.execute(vis=Int64[], solv_mode=:lpdm_bv, reward_func=:quadratic, action_space_type=:bv, output=0, steps = steps)
# execute2d.execute(vis=Int64[], solv_mode=:despot, reward_func=:quadratic, action_space_type=:small, output=0, steps = steps)
# execute2d.batch_execute(n=50, debug=0, reward_func=:quadratic)
# execute2d.batch_execute(n=50, debug=0, reward_func=:quadratic)


function causality(source::Vector{<:UVT1}, target::Vector{<:UVT2}, test::CT) where {
        UVT1<:AbstractUncertainValue,
        UVT2<:AbstractUncertainValue,
        CT<:CausalityTest}
    causality(resample.(source), resample.(target), test)
end

function causality(source, target::Vector{<:UVT}, test::CT) where {
        UVT<:AbstractUncertainValue,
        CT<:CausalityTest}
    causality(source, resample.(target), test)
end

function causality(source::Vector{<:UVT}, target, test::CT) where {
        UVT<:AbstractUncertainValue,
        CT<:CausalityTest}
    causality(resample.(source), target, test)
end

function causality(source::UVT1, target::UVT2, test::CT) where {
        UVT1<:AbstractUncertainValueDataset,
        UVT2<:AbstractUncertainValueDataset,
        CT<:CausalityTest}
    causality(resample.(source), resample.(target), test)
end

function causality(source, target::UVT, test::CT) where {
        UVT<:AbstractUncertainValueDataset,
        CT<:CausalityTest}
    causality(source, resample.(target), test)
end

function causality(source::UVT, target, test::CT) where {
        UVT<:AbstractUncertainValueDataset,
        CT<:CausalityTest}
    causality(resample.(source), target, test)
end


using NeRCA
using ROy
using Sockets

calib = Calibration("scripts/NeRCA_00000043_20022019.detx")
calib = Calibration("/home/tgal/data/detx/NeRCA_-00000001_20171212.detx")

client = NeRCA.CHClient(ip"127.0.0.1", 5553, ["IO_EVT"])

try
    for message in client
        event = read(IOBuffer(message.data), NeRCA.DAQEvent)
        hits = event.snapshot_hits
        chits = NeRCA.calibrate(calib, hits)
        thits = event.triggered_hits
        cthits = NeRCA.calibrate(calib, thits)
        println(length(cthits))
    end
catch InterruptException
    println("ok")
end


#
# proximal ALM implementation
#
module ProxAL

using Ipopt, JuMP
using Printf, CatViews
using ExaPF
using ExaTron
using LinearAlgebra
using SparseArrays
using MPI

abstract type AbstractPrimalSolution end
abstract type AbstractDualSolution end
abstract type AbstractBackend end
abstract type AbstractBlocks end
mutable struct ProxALProblem
    opfBlockData::AbstractBlocks

    #---- iterate information ----
    x::AbstractPrimalSolution
    λ::AbstractDualSolution
    xprev::AbstractPrimalSolution
    λprev::AbstractDualSolution
    objvalue::Vector{Float64}
    lyapunov::Vector{Float64}
    maxviol_t::Vector{Float64}
    maxviol_c::Vector{Float64}
    maxviol_d::Vector{Float64}
    maxviol_t_actual::Vector{Float64}
    maxviol_c_actual::Vector{Float64}
    nlp_opt_sol::Array{Float64,2}
    nlp_soltime::Vector{Float64}
    wall_time_elapsed_actual::Float64
    wall_time_elapsed_ideal::Float64
    iter

    #---- other/static information ----
    blkLinIndex
end

include("Evaluators/Evaluators.jl")
include("params.jl")
include("OPF/opfdata.jl")
include("OPF/opfsolution.jl")
include("backends.jl")
include("blocks.jl")
include("OPF/opfmodel.jl")
include("communication.jl")
include("Evaluators/ProxALEvalutor.jl")
include("Evaluators/NonDecomposedModel.jl")

export ModelInfo, AlgParams
export ProxALEvaluator, NonDecomposedModel
export optimize!
export JuMPBackend, ExaPFBackend, ExaTronBackend

function update_primal_nlpvars(
    x::AbstractPrimalSolution,
    opfBlockData::AbstractBlocks,
    blk::Int,
    modelinfo::ModelInfo,
    algparams::AlgParams
)
    block = opfBlockData.blkIndex[blk]
    k = block[1]
    t = block[2]
    range_k = algparams.decompCtgs ? (k:k) : (1:(modelinfo.num_ctgs + 1))

    from = 1
    to = size(x.Pg,1)*length(range_k)
    @views x.Pg[:,range_k,t][:] .= opfBlockData.colValue[from:to,blk]

    from = 1+to
    to = to + size(x.Qg, 1)*length(range_k)
    @views x.Qg[:,range_k,t][:] .= opfBlockData.colValue[from:to,blk]

    from = 1+to
    to = to + size(x.Vm, 1)*length(range_k)
    @views x.Vm[:,range_k,t][:] .= opfBlockData.colValue[from:to,blk]

    from = 1+to
    to = to + size(x.Va, 1)*length(range_k)
    @views x.Va[:,range_k,t][:] .= opfBlockData.colValue[from:to,blk]

    from = 1+to
    to = to + length(range_k)
    @views x.ωt[range_k,t] .= opfBlockData.colValue[from:to,blk]

    from = 1+to
    to = to + size(x.St, 1)
    if !algparams.decompCtgs || k == 1
        @views x.St[:,t] .= opfBlockData.colValue[from:to,blk]
    end

    # Zt will be updated in update_primal_penalty
    from = 1+to
    to = to + size(x.Zt, 1)

    from = 1+to
    to = to + size(x.Sk, 1)*length(range_k)
    if !algparams.decompCtgs || k > 1
        @views x.Sk[:,range_k,t][:] .= opfBlockData.colValue[from:to,blk]
    end

    from = 1+to
    to = to + size(x.Zk, 1)*length(range_k)
    if !algparams.decompCtgs
        @views x.Zk[:,range_k,t][:] .= opfBlockData.colValue[from:to,blk]
    end

    return nothing
end

function update_primal_penalty(
    x::AbstractPrimalSolution,
    opfdata::OPFData,
    opfBlockData::AbstractBlocks,
    blk::Int,
    primal::AbstractPrimalSolution,
    dual::AbstractDualSolution,
    modelinfo::ModelInfo,
    algparams::AlgParams
)
    ngen = size(x.Pg, 1)
    block = opfBlockData.blkIndex[blk]
    k = block[1]
    t = block[2]

    if t > 1 && k == 1 && modelinfo.time_link_constr_type == :penalty
        β = [opfdata.generators[g].ramp_agc for g=1:ngen]
        @views x.Zt[:,t] .= (((algparams.ρ_t/32.0)*primal.Zt[:,t]) .- dual.ramping[:,t] .-
                                (algparams.ρ_t*(x.Pg[:,1,t-1] .- x.Pg[:,1,t] .+ x.St[:,t] .- β))
                            ) ./  max(algparams.zero, (algparams.ρ_t/32.0) + algparams.ρ_t + (modelinfo.obj_scale*algparams.θ_t))
    end
    if k > 1 && algparams.decompCtgs
        if modelinfo.ctgs_link_constr_type ∈ [:frequency_penalty, :preventive_penalty, :corrective_penalty]
            if modelinfo.ctgs_link_constr_type == :corrective_penalty
                β = [opfdata.generators[g].scen_agc for g=1:ngen]
            elseif modelinfo.ctgs_link_constr_type == :frequency_penalty
                β = [-opfdata.generators[g].alpha*x.ωt[k,t] for g=1:ngen]
                @assert norm(x.Sk[:,k,t]) <= algparams.zero
            else
                β = zeros(ngen)
                @assert norm(x.Sk[:,k,t]) <= algparams.zero
            end
            @views x.Zk[:,k,t] .=   (((algparams.ρ_c/32.0)*primal.Zk[:,k,t]) .- dual.ctgs[:,k,t] .-
                                        (algparams.ρ_c*(x.Pg[:,1,t] .- x.Pg[:,k,t] .+ x.Sk[:,k,t] .- β))
                                    ) ./  max(algparams.zero, (algparams.ρ_c/32.0) + algparams.ρ_c + (modelinfo.obj_scale*algparams.θ_c))
        end
    end

    return nothing
end

function update_dual_vars(
    λ::AbstractDualSolution,
    opfdata::OPFData,
    opfBlockData::AbstractBlocks,
    blk::Int,
    primal::AbstractPrimalSolution,
    modelinfo::ModelInfo,
    algparams::AlgParams
)
    block = opfBlockData.blkIndex[blk]
    k = block[1]
    t = block[2]

    maxviol_t, maxviol_c = 0.0, 0.0

    if t > 1 || (k > 1 && algparams.decompCtgs)
        d = Dict(:Pg => primal.Pg,
                 :ωt => primal.ωt,
                 :St => primal.St,
                 :Zt => primal.Zt,
                 :Sk => primal.Sk,
                 :Zk => primal.Zk)
    end

    if t > 1 && k == 1
        @assert modelinfo.time_link_constr_type == :penalty
        link_constr = compute_time_linking_constraints(d, opfdata, modelinfo, t)
        λ.ramping[:,t] += algparams.ρ_t*link_constr[:ramping][:]
        maxviol_t = maximum(abs.(link_constr[:ramping][:]))
    end
    if k > 1 && algparams.decompCtgs
        @assert modelinfo.ctgs_link_constr_type ∈ [:frequency_penalty, :preventive_penalty, :corrective_penalty]
        link_constr = compute_ctgs_linking_constraints(d, opfdata, modelinfo, k, t)
        λ.ctgs[:,k,t] += algparams.ρ_c*link_constr[:ctgs][:]
        maxviol_c = maximum(abs.(link_constr[:ctgs][:]))
    end

    return maxviol_t, maxviol_c
end

function ProxALProblem(
    opfdata::OPFData,
    rawdata::RawData,
    modelinfo::ModelInfo,
    algparams::AlgParams,
    backend::AbstractBackend,
    comm::Union{MPI.Comm,Nothing},
)
    # initial values
    x = OPFPrimalSolution(opfdata, modelinfo)
    λ = OPFDualSolution(opfdata, modelinfo)
    backend = if isa(backend, JuMPBackend)
        JuMPBlockBackend
    elseif isa(backend, ExaPFBackend)
        ExaBlockBackend
    elseif isa(backend, ExaTronBackend)
        TronBlockBackend
    end
    # NLP blocks
    blocks = OPFBlocks(
        opfdata, rawdata;
        modelinfo=modelinfo, algparams=algparams,
        backend=backend, comm
    )
    blkLinIndex = LinearIndices(blocks.blkIndex)
    for blk in blkLinIndex
        if is_my_work(blk, comm)
            if algparams.mode ∉ [:nondecomposed, :lyapunov_bound]
                init!(blocks.blkModel[blk], algparams)
            end
            blocks.colValue[:,blk] .= get_block_view(x, blocks.blkIndex[blk], modelinfo, algparams)
        end
    end

    iter = 0
    objvalue = []
    lyapunov = []
    maxviol_t = []
    maxviol_c = []
    maxviol_d = []
    maxviol_t_actual = []
    maxviol_c_actual = []
    nlp_opt_sol = Array{Float64, 2}(undef, blocks.colCount, blocks.blkCount)
    nlp_opt_sol .= blocks.colValue
    nlp_soltime = Vector{Float64}(undef, blocks.blkCount)
    wall_time_elapsed_actual = 0.0
    wall_time_elapsed_ideal = 0.0
    xprev = deepcopy(x)
    λprev = deepcopy(λ)

    return ProxALProblem(
        blocks,
        x,
        λ,
        xprev,
        λprev,
        objvalue,
        lyapunov,
        maxviol_t,
        maxviol_c,
        maxviol_d,
        maxviol_t_actual,
        maxviol_c_actual,
        nlp_opt_sol,
        nlp_soltime,
        wall_time_elapsed_actual,
        wall_time_elapsed_ideal,
        iter,
        blkLinIndex,
    )
end

function runinfo_update(
    runinfo::ProxALProblem,
    opfdata::OPFData,
    opfBlockData::AbstractBlocks,
    modelinfo::ModelInfo,
    algparams::AlgParams,
    comm::Union{MPI.Comm,Nothing}
)
    iter = runinfo.iter
    obj = 0.0
    for blk in runinfo.blkLinIndex
        if is_my_work(blk, comm)
            obj += compute_objective_function(runinfo.x, opfdata, opfBlockData, blk, modelinfo, algparams)
        end
    end
    obj = comm_sum(obj, comm)
    push!(runinfo.objvalue, obj)

    iter = runinfo.iter
    lyapunov = 0.0
    for blk in runinfo.blkLinIndex
        if is_my_work(blk, comm)
            lyapunov += compute_lyapunov_function(runinfo.x, runinfo.λ, opfdata, opfBlockData, blk, runinfo.xprev, modelinfo, algparams)
        end
    end
    lyapunov = comm_sum(lyapunov, comm)
    push!(runinfo.lyapunov, lyapunov)

    maxviol_t_actual = 0.0
    for blk in runinfo.blkLinIndex
        if is_my_work(blk, comm)
            lmaxviol_t_actual = compute_true_ramp_error(runinfo.x, opfdata, opfBlockData, blk, modelinfo)
            maxviol_t_actual = max(maxviol_t_actual, lmaxviol_t_actual)
        end
    end
    maxviol_t_actual = comm_max(maxviol_t_actual, comm)
    push!(runinfo.maxviol_t_actual, maxviol_t_actual)

    maxviol_c_actual = 0.0
    for blk in runinfo.blkLinIndex
        if is_my_work(blk, comm)
            lmaxviol_c_actual = compute_true_ctgs_error(runinfo.x, opfdata, opfBlockData, blk, modelinfo)
            maxviol_c_actual = max(maxviol_c_actual, lmaxviol_c_actual)
        end
    end
    maxviol_c_actual = comm_max(maxviol_c_actual, comm)
    push!(runinfo.maxviol_c_actual, maxviol_c_actual)

    maxviol_d = 0.0
    for blk in runinfo.blkLinIndex
        if is_my_work(blk, comm)
            lmaxviol_d = compute_dual_error(runinfo.x, runinfo.xprev, runinfo.λ, runinfo.λprev, opfdata, opfBlockData, blk, modelinfo, algparams)
            maxviol_d = max(maxviol_d, lmaxviol_d)
        end
    end
    maxviol_d = comm_max(maxviol_d, comm)
    push!(runinfo.maxviol_d, maxviol_d)
    if algparams.verbose > 0 && comm_rank(comm) == 0
        if iter == 1
            @printf("--------------------------------------------------------------------------------------------------------------------\n");
            @printf("iter ramp_err   ramp_err   ctgs_err   ctgs_err   dual_error lyapunov_f    rho_t    rho_c  theta_t  theta_c      tau \n");
            @printf("     (penalty)  (actual)   (penalty)  (actual)\n");
            @printf("--------------------------------------------------------------------------------------------------------------------\n");
        end
        @printf("%4d ", iter-1);
        @printf("%10.4e ", runinfo.maxviol_t[iter])
        @printf("%10.4e ", runinfo.maxviol_t_actual[iter])
        @printf("%10.4e ", runinfo.maxviol_c[iter])
        @printf("%10.4e ", runinfo.maxviol_c_actual[iter])
        @printf("%10.4e ", runinfo.maxviol_d[iter])
        @printf("%10.4e ", runinfo.lyapunov[iter])
        @printf("%8.2e ", algparams.ρ_t)
        @printf("%8.2e ", algparams.ρ_c)
        @printf("%8.2e ", algparams.θ_t)
        @printf("%8.2e ", algparams.θ_c)
        @printf("%8.2e ", algparams.τ)
        @printf("\n")
    end
end

end # module


module FastScanning

using Unitful, ImagineInterface, ImagineAnalyses, ImagineProcedures
using Images, CoordinateTransformations, ProgressMeter, FixedPointNumbers, IntervalSets
using Statistics, Distributed
using RegisterQD, RegisterMismatch
using Requires

import Unitful:s, Hz, μm, V

include("slicetiming_analysis.jl")
#include("bsearch.jl")
include("slicetiming_generate.jl")
include("experiments.jl")

function __init__()
    @require Imagine="28fce3f0-1dbc-11e9-10e3-e71e6ceb4e7c" include("run_experiments.jl")
end

export default_toffsets, pos_commands, get_cyc_pulses, pos_mon_lag_experiment, slicetiming_experiment

end # module


using .JSort
import Base: isless

struct GammaCascade
    gate::AbstractGate
    e::Vector{Float32}
    regions::Vector{Tuple{Float32, Float32}}
    keywords::Dict{Symbol, Any}
    #function GammaCascade(gate::Gate, e::AbstractVector, regions::AbstractVector)
        #new(gate, sort(e) |> reverse)
    #end
end
GammaCascade(gate::AbstractGate, e::AbstractVector, regions::AbstractVector) = GammaCascade(gate, e, regions, Dict(:align => Dict()))

function extend!(base::GammaCascade, others::AbstractArray{GammaCascade})
    lower = [c for c in sort(others) if c < base]
    for other in lower
        extend!(base, other)
    end
end

function extend!(base::GammaCascade, other::GammaCascade; width=3)
    from = center(base.gate)
    to = center(other.gate)
    for e in base.e
        if abs(from - e - to) < width
            append!(base.e, other.e)
            break
        end
    end
end

isless(c1::GammaCascade, c2::GammaCascade) = isless(c1.gate, c2.gate)
center(c::GammaCascade) = center(c.gate)


module DistributionsAD

using PDMats,
      LinearAlgebra,
      Distributions,
      Random,
      SpecialFunctions,
      StatsFuns,
      Compat,
      Requires,
      ZygoteRules,
      ChainRules,  # needed for `ChainRules.chol_blocked_rev`
      ChainRulesCore,
      FillArrays,
      Adapt

using SpecialFunctions: logabsgamma, digamma
using LinearAlgebra: copytri!, AbstractTriangular
using Distributions: AbstractMvLogNormal,
                     ContinuousMultivariateDistribution
using Base.Iterators: drop

import StatsBase
import StatsFuns: logsumexp,
                  binomlogpdf,
                  nbinomlogpdf,
                  poislogpdf,
                  nbetalogpdf
import Distributions: MvNormal,
                      MvLogNormal,
                      logpdf,
                      quantile,
                      PoissonBinomial,
                      Binomial,
                      BetaBinomial,
                      Erlang

export TuringScalMvNormal,
       TuringDiagMvNormal,
       TuringDenseMvNormal,
       TuringMvLogNormal,
       TuringPoissonBinomial,
       TuringWishart,
       TuringInverseWishart,
       arraydist,
       filldist

include("common.jl")
include("arraydist.jl")
include("filldist.jl")
include("univariate.jl")
include("multivariate.jl")
include("matrixvariate.jl")
include("flatten.jl")

include("chainrules.jl")
include("zygote.jl")

@init begin
    @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" begin
        using .ForwardDiff
        using .ForwardDiff: @define_binary_dual_op # Needed for `eval`ing diffrules here
        include("forwarddiff.jl")
    end

    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
        # ensures that we can load ForwardDiff without depending on it
        # (it is a dependency of ReverseDiff and therefore always available)
        @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" begin
            include("reversediff.jl")
        end
    end

    @require Tracker="9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" begin
        using DiffRules
        using SpecialFunctions
        using LinearAlgebra: AbstractTriangular
        using .Tracker: Tracker, TrackedReal, TrackedVector, TrackedMatrix,
                        TrackedArray, TrackedVecOrMat, track, @grad, data
        include("tracker.jl")
    end

    @require LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" begin
        using .LazyArrays: BroadcastArray, BroadcastVector, LazyArray

        const LazyVectorOfUnivariate{
            S<:ValueSupport,
            T<:UnivariateDistribution{S},
            Tdists<:BroadcastVector{T},
        } = VectorOfUnivariate{S,T,Tdists}

        function Distributions._logpdf(
            dist::LazyVectorOfUnivariate,
            x::AbstractVector{<:Real},
        )
            return sum(copy(logpdf.(dist.v, x)))
        end

        function Distributions.logpdf(
            dist::LazyVectorOfUnivariate,
            x::AbstractMatrix{<:Real},
        )
            size(x, 1) == length(dist) ||
                throw(DimensionMismatch("Inconsistent array dimensions."))
            return vec(sum(copy(logpdf.(dists, x)), dims = 1))
        end

        const LazyMatrixOfUnivariate{
            S<:ValueSupport,
            T<:UnivariateDistribution{S},
            Tdists<:BroadcastArray{T,2},
        } = MatrixOfUnivariate{S,T,Tdists}

        function Distributions._logpdf(
            dist::LazyMatrixOfUnivariate,
            x::AbstractMatrix{<:Real},
        )
            return sum(copy(logpdf.(dist.dists, x)))
        end

        lazyarray(f, x...) = LazyArray(Base.broadcasted(f, x...))
        export lazyarray
    end
end

end


#=
    NEMO: Next Energy Modeling system for Optimization.
    https://github.com/sei-international/NemoMod.jl

    Copyright © 2018: Stockholm Environment Institute U.S.

    File description: NemoMod module.
=#

module NemoMod
export calculatescenario, createnemodb, dropdefaultviews, dropresulttables, logmsg, setparamdefault

#= List of module global variables.
    • csdbpath - dbpath argument in last invocation of calculatescenario()
    • csquiet - quiet argument in last invocation of calculatescenario()
    • csjumpmodel - jumpmodel argument in last invocation of calculatescenario() (only set if a customconstraints include is performed)
=#

# BEGIN: Access other modules and code files.
using JuMP, SQLite, DataFrames, Distributed, Dates, ConfParser, MathOptInterface
using GLPK, Cbc  # Open-source solvers

# Proprietary solvers - enclosed in try blocks for users who aren't using NEMO installer
try
    using CPLEX
catch
    # Just continue
end

try
    using Gurobi
catch
    # Just continue
end

try
    using Mosek
catch
    # Just continue
end

# Do not load Xpress since it throws an unrecoverable error when loaded in a package on a machine without an Xpress license

include("db_structure.jl")  # Functions for manipulating structure of scenario databases
include("other_functions.jl")  # Core NEMO functions
include("scenario_calculation.jl")  # Functions for calculating a scenario with NEMO
# END: Access other modules and code files.

end  # module NemoMod


type GuidType
    guidance_mode::UInt
    num_waypts::UInt
    current_wypt::UInt
    WaypointList::Array{WGS_Pos,1}
    tgo::Float64
    miss::Float64
    miss_vec_ned::Array{Float64,1}
    eta_dive_cmd::Float64
    eta_turn_cmd::Float64
    alt_cmd::Float64
    gamma_cmd::Float64
    Mach_cmd::Float64
end
################################################################################
## Cooperative Formation
function GuidInit(wyptlist::Array{WGS_Pos,1}, num_waypoints::Float64)
    guidance_mode = 0;
    current_wypt = 1;
    tgo = 99999.0;
    miss = 9999.0;
    miss_vec_ned = zeros(3);
    eta_dive_cmd = 0.0;
    eta_turn_cmd = 0.0;
    alt_cmd = 0.0;
    gamma_cmd = 0.0;
    Mach_cmd = 0.0;

    Guidance_Data = GuidType(guidance_mode, num_waypoints, current_wypt, wyptlist,
                             tgo, miss, miss_vec_ned, eta_turn_cmd, eta_dive_cmd,
                             alt_cmd, gamma_cmd, Mach_cmd );

    return Guidance_Data;
end

################################################################################
## Cooperative Formation
function CoopFormation(nav_state::NavStates, guidance_data::GuidType, SS::Array{Float64,2},
    missile_id::UInt,  salvo_size::UInt, message_list::Array{CoopFormation,1})

    # Computing NED to Body DCM
    Tr_BN = nav_state.euler.Tr_bn;
    Tr_NE = nav_state.position.Tr_ne;
    rm_ecef = nav_state.position.rm_ecef;
    vm_ecef = nav_state.velocity.vm_ecef;
    nn = length(message_list);

    ## Clearing the summation variables
    sumvec = zeros(3);
    sum = 0;
    ii = missile_id;
    icoop_flag = true;

    ## Gravity Vector in NED frame
    tmpvec = zeros;
    tmpvec[3] = 1.0;
    gvec_LL[1] = 0.0;
    gvec_LL[2] = 0.0;
    gvec_LL[3] = go;

    rm_ned = Tr_NE*rm_ecef;
    vm_ned = Tr_NE*vm_ecef;

    ## Computing the Guidance Command
    jj = 0;

    for jj = 1:NN

        ## Determine Communication Topology
        if ( message_list[jj].time >= 0 )
            SS[jj] = 1;
        else
            SS[jj] = 0;
        end

        ## For right now, Missile 0 is the leader.
        if( jj == 1 )
            rm_lead_ecef = message_list[jj].rm_ecef;
            vm_lead_ecef = message_list[jj].vm_ecef;
            rm_lead_ned = Tr_NE*rm_lead_ecef;
            vm_lead_ned = Tr_NE*vm_lead_ecef;

            ## Unit Velocity Vector
            unit_vmlead_ned = vm_lead_ned/norm(vm_lead_ned);

            ## Gamma and Azimuth of Lead Vector
            gamma_lead = atan2(-unit_vmlead_ned[3], sqrt( unit_vmlead_ned[1]^2 + unit_vmlead_ned[2]^2 ));
            azimuth_lead = atan2( unit_vmlead_ned[1], unit_vmlead_ned[0] );

            ## NED to Velocity Frame Transformation Matrix
            Roll_lead = 0.0;
            Tr_VN_LEAD = TR_BN(Roll_lead, gamma_lead, azimuth_lead );
            Tr_NV_LEAD = Tr_VN_LEAD';

            #delta_vel[0] = 0.0;
            #delta_vel[1] = YY(jj) - YY(missile_number);
            #delta_vel[2] = 0.0;

            delta_vel = zeros(3);
            delta_vel[1] = 0.0;
            delta_vel[2] = 200.0*sin(0.20*guidtime + (ii-1)*2.0*MC::PI/(salvo_size - 1));
            delta_vel[3] = 200.0*cos(0.20*guidtime + (ii-1)*2.0*MC::PI/(salvo_size - 1));
            delta_ned = Tr_NV_LEAD*delta_vel;

            wn = 0.5;
            zz = 0.70;
            Kpf = wn*wn;
            Kvf = 2.0*wn*zz;
            sumvec = sumvec + SS[jj]*( Kpf*(rm_lead_ned - rm_ned - delta_ned ) + Kvf*(vm_lead_ned - vm_ned) );

        end

        jj = jj + 1;
    end

    ## Guidance Command
    acc_cmd_ned = sumvec - gvec_LL;
    acc_cmd_vel = Tr_VN*acc_cmd_ned;

    speed_accel_cmd = acc_cmd_vel[0];
    eta_turn_cmd = acc_cmd_vel[1];
    eta_dive_cmd = acc_cmd_vel[2];

    turn_lim = 3.0*go;
    speed_accel_cmd = Bound(speed_accel_cmd, -turn_lim, turn_lim );

    vcmd_coop = vcmd_coop + dt50*speed_accel_cmd;
    eta_turn_cmd = Bound(eta_turn_cmd, -turn_lim, turn_lim );
    eta_dive_cmd = Bound(eta_dive_cmd, -turn_lim, turn_lim );

end

################################################################################
## Altitude Command
function AltitudeControl(nav_state::NavStates, alt_cmd::Float64)

    ## Parameter Computation
    # Computing NED to Body DCM
    Tr_BN = nav_state.euler.Tr_bn;
    Tr_NE = nav_state.wgs_pos.Tr_ne;
    rm_ecef = nav_state.wgs_pos.rm_ecef;
    vm_ecef = nav_state.velocity.vm_ecef;
    vm_ned = Tr_NE*vm_ecef;

    gamma = atan2( -vm_ned[3], sqrt( vm_ned[1]^2 + vm_ned[2]^2 ) );
    alt = nav_state.wgs_pos.altitude;
    Vmag = nav_state.velocity.vmag;

    ## Altitude Control Gains
    Kalt = 0.210;

    ## Altitude Error
    alt_err = alt_cmd - alt;
    println("alt_err = ", alt_err)

    ## Gamma Command
    gammac = Kalt/Vmag*alt_err;

    gamlim = 10.0*d2r;
    gammac = Bound(gammac, -gamlim, gamlim);

    ## Gamma Control Gains
    wn = 1.2;
    zz = 0.70;
    Kvf = 2*zz*wn;
    Kpf = wn*wn;

    ## Eta Dive Command
    eta_dive_cmd = -Kpf*Vmag*(gammac - gamma);
    eta_dive_cmd = Bound(eta_dive_cmd, -1.5*go, 1.5*go );

    return eta_dive_cmd;
end