export generate_positions, initialize_mesh, start, getstats, plotgraph

to_mW(dBm) = 10^(dBm / 10)

const ble_37_channel_wavelength = 0.12491352
const part_of_path_loss_for_37 = (ble_37_channel_wavelength / (4 * π))^2
const message_length = 312
const scanner_sensitivity = to_mW(-95)

function randexclude(range::UnitRange{T}, exclude::T) where T
    while true
        x = rand(range)
        x != exclude && return x
    end
end

function agent_step!(node::Node, model::AgentBasedModel)
    # produce a packet and store it in the queue for later
    if rand() < model.packet_emit_rate / model.n
        packet_seq = length(model.packets) + 1
        packet = Packet(packet_seq, node.id, randexclude(1:model.n, node.id), model.ttl, model.tick, 0, false)
        push!(model.packets, packet)

        # (pomdp): keeping track of nodes that will receive this packet
        model.packet_xs[packet.seq] = Int[]
        push!(node.pocket, (model.tick + rand(1:node.t_back_off_delay), packet_seq, model.ttl))
    end

    if node.state == :scanning
        node.event_start > model.tick && return
        node.channel = div(abs(model.tick - node.event_start), node.t_scan_interval) + 37

        # no packets acquired, repeat scanning
        if node.channel > 39
            node.channel = 37
            node.event_start = model.tick
        end

        # there is no more job for sinks here
        if node.role == :relay && node.packet_seq != 0 && !(node.packet_seq in node.received_packets)
            push!(node.received_packets, node.packet_seq)

            # our own packets are treated the same way as those generated by not us
            if node.packet_ttl > 1
                # put a packet in the pocket
                push!(node.pocket, (model.tick + rand(1:node.t_back_off_delay), node.packet_seq, node.packet_ttl - 1))
            end
        end
    end

    # wait until it's time to emit the first packet from the queue
    if !isempty(node.pocket) && model.tick >= node.pocket[1][1]
        node.state = :advertising
        node.event_start = model.tick
    end

    if node.state == :advertising
        node.event_start > model.tick && return
        if model.tick == node.event_start
            _, node.packet_seq, node.packet_ttl = popfirst!(node.pocket)
        end

        d, r = divrem(abs(model.tick - node.event_start), node.t_interpdu)

        # transmitting happens once for each channel
        node.transmitting = iszero(r)
        node.channel = d + 37

        # the end of one advertising cycle
        if node.channel > 39
            # bonus retransmissions for the original/retx packet
            if node.n_transmit_left > 0
                push!(node.pocket, (model.tick + node.t_back_off_delay, node.packet_seq, node.packet_ttl))
            end

            node.packet_seq = 0
            node.transmitting = false
            node.state = :scanning
            node.event_start = model.tick + 1
            node.n_transmit_left = node.n_retx_transmit_count
            node.channel = 37
        end
    end
end

# without sqrt, and taking into account overlapping devices
distance(a::Tuple{Int, Int}, b::Tuple{Int, Int}) = (a[1] - b[1]) ^ 2 + (a[2] - b[2]) ^ 2 + 1e-3

function calc_rssi!(agents, P_tx::Vector{Float64})
    n = length(agents)
    rssi_map = Matrix{Float64}(undef, n, n)

    L_p = Matrix{Float64}(undef, n, n)

    for i in 1:n
        for j in 1:(i - 1)
            L_p[i, j] = L_p[j, i]
        end

        for j in (i + 1):n
            L_p[i, j] = 1 / distance(agents[i].pos, agents[j].pos)
        end
    end

    for i in 1:n
        for j in 1:n
            rssi_map[i, j] = P_tx[j] * L_p[i, j] * part_of_path_loss_for_37 # dB
        end

        rssi_map[i, i] = 0
    end

    rssi_map
end

function calculate_par!(rssi_neighbourhood::Vector{Float64}, total_rssi::Float64)
    total_par = 0.0

    @inbounds for idx in eachindex(rssi_neighbourhood)
        rssi_neighbourhood[idx] == 0 && continue

        SINR = rssi_neighbourhood[idx] / (total_rssi - rssi_neighbourhood[idx])
        BER = 0.5 * erfc(sqrt(SINR / 2))
        PAR = (1 - BER)^message_length

        rssi_neighbourhood[idx] = PAR
        total_par += PAR
    end

    total_par
end

function select_rssi!(rssi_neighbourhood::Vector{Float64}, rssi_map::SubArray{Float64}, transmitters::CircularBuffer{Int})
    total_rssi = 0.0

    @inbounds for idx in transmitters
        rssi = rssi_map[idx]

        total_rssi += rssi
        if rssi > scanner_sensitivity
            rssi_neighbourhood[idx] = rssi
        end
    end

    total_rssi
end

function model_step!(model::AgentBasedModel)
    transmitters_count = 0

    for agent in allagents(model)
        if agent.transmitting
            push!(model.transmitters[agent.channel - 36], agent.id)
            transmitters_count += 1
        end
    end

    # (pomdp): clear previous reward nominees
    empty!(model.reward_plate)

    if iszero(transmitters_count)
        model.tick += 1
        return
    end

    rssi_neighbourhood = zeros(model.n)
    count_successes = zeros(model.n)

    for dst in allagents(model)
        dst.state == :scanning || continue

        # select RSSI of only transmitting devices reachable to here
        # begs the question whether RSSI can be calculated on the fly
        fill!(rssi_neighbourhood, 0)

        txs = model.transmitters[dst.channel - 36]
        total_rssi = select_rssi!(rssi_neighbourhood, view(model.rssi_map, :, dst.id), txs)

        iszero(total_rssi) && continue

        total_par = calculate_par!(rssi_neighbourhood, total_rssi)

        # do any of this packets pass?
        rand() <= total_par || continue

        tx_id = sample(1:length(txs), Weights(rssi_neighbourhood[1:length(txs)]))
        source_id = txs[tx_id]

        # @show model[source_id]
        packet = model.packets[model[source_id].packet_seq]
        count_successes[source_id] = 1

        if packet.dst == dst.id
            if !packet.done
                # since the packet is delivered, acknowledge everyone participating
                append!(model.reward_plate, model.packet_xs[packet.seq])

                # marking the packet as delivered
                packet.done = true
                packet.time_done = model.tick
            end
        else
            # register a touch upon this not delivered packet
            !packet.done && push!(model.packet_xs[packet.seq], dst.id)

            # copy packet's reference and source's ttl
            dst.packet_seq = packet.seq
            dst.packet_ttl = model[source_id].packet_ttl
            # dst.packet_rssi = neighbours_agents[src_i].packet_ttl
        end

        fill!(rssi_neighbourhood, 0.0)
    end

    if transmitters_count > 0
        model.packets_lost += transmitters_count - sum(count_successes)
        model.transmitters_count += transmitters_count
    end

    # reset transmitter signaling array
    for idx = 1:3
        empty!(model.transmitters[idx])
    end

    model.tick += 1
end

function plotgraph(model::AgentBasedModel)
    g = SimpleGraph(nagents(model))

    for (src_i, src) in enumerate(allagents(model)), (dst_i, dst) in enumerate(allagents(model))
        src.role == :relay || continue
        dst.role != :source || continue

        if model.rssi_map[src_i, dst_i] > model.scanner_sensitivity
            add_edge!(g, src.id, dst.id)
        end
    end

    colors = Dict(:relay => "orange", :sink => "darkblue", :source => "transparent")

    agents = sort(collect(allagents(model)), by=a -> a.id)
    filter!(a -> a.role != :source, agents)

    xs = map(a -> a.pos[1], agents)
    ys = map(a -> a.pos[2], agents)
    cs = map(a -> colors[a.role], agents)

    gplot(g, xs, ys; nodefillc = cs)
end


"""
    generate_positions(; dims = (200, 200), n = 100) → positions::Array{Tuple{Int,Int},1}

Generates random xy-coordinates for `n` nodes with `dims`
"""
generate_positions(; dims=(100, 100), n=64) = [(rand(1:dims[1]), rand(1:dims[2])) for _ in 1:n]

"""
    initialize_mesh(positions::Vector{Tuple{Int, Int}}, roles::Vector{Int}) → mesh::AgentBasedModel

Create a mesh by specifying roles of each node by `roles`=[0, 1, 0...] where 1 means that the node is a relay, sink otherwise, and including `positions`=[(Int, Int)...] as nodes' positions
"""
function initialize_mesh(positions::Vector{NTuple{2, Int}}, roles::AbstractVector)
    length(roles) == length(positions) || throw(ArgumentError("Both arguments must be of equal length"))
    n = length(roles)

    properties = Dict(
        :n => n,
        :tick => 0,
        :packet_emit_rate => 50 / 1000,
        :transmitters => [CircularBuffer{Int}(n) for _ = 1:3],
        :packets_lost => 0,
        :transmitters_count => 0,
        :ttl => 5,
        :rssi_map => Matrix{Float64}(undef, n, n),
        :tx_powers => Vector{Float64}(undef, n),
        :scanner_sensitivity => to_mW(-95),
        :shadow_d => LogNormal(0, 1),
        :multipath_d => Rayleigh(4),
        :wifi_noise_d => Normal((to_mW(-125) - to_mW(-135)) / 2 + to_mW(-135), (to_mW(-125) - to_mW(-135)) / 2),
        :gaussian_noise => Normal((to_mW(-125) - to_mW(-135)) / 2 + to_mW(-135), (to_mW(-125) - to_mW(-135)) / 2),
        :packets => Packet[],
        # packet.seq => nodes that have touched that packet
        :packet_xs => Dict{Int, Vector{Int}}(),
        # plate for the actors who contributed to packet's successful delivery
        :reward_plate => Int[]
    )

    space = GridSpace((maximum(first.(positions)), maximum(last.(positions))))
    model = ABM(Node, space; properties, warn = false)

    for i in 1:model.n
        node = Node(id = i; pos = positions[i], role = roles[i] == 1 ? :relay : :sink)

        add_agent_pos!(node, model)

        model.tx_powers[node.id] = node.tx_power
        model.tx_powers[i] = node.tx_power
    end

    model.rssi_map = calc_rssi!(collect(allagents(model)), model.tx_powers)
    model
end

function getstats(model::AgentBasedModel)
    # filter packets which were too recently produced, or produced in an unreachable position
    filter!(packet -> packet.time_start < model.tick - 500, model.packets)

    # worst case in terms of PDR to a single device
    worstnode = argmax(map(id -> count(p -> p.dst == id && p.done == false, model.packets), 1:nagents(model)))
    deprived = filter(p -> p.dst == worstnode, model.packets)
    delivered = filter(p -> p.done == true, model.packets)
    delays = map(p -> p.time_done - p.time_start, delivered)

    # probing centrality measure
    devices = collect(allagents(model))
    filter!(n -> n.role != :source, devices)

    g = SimpleDiGraph(length(devices))

    for (src_i, src) in enumerate(devices), (dst_i, dst) in enumerate(devices)
        src.role == :relay || continue

        if model.rssi_map[src_i, dst_i] > model.scanner_sensitivity
            add_edge!(g, src.id, dst.id)
        end
    end

    centrality = indegree_centrality(g, normalize=false) ./ length(devices) |> mean

    return ( pdr = length(delivered) / length(model.packets),
             worstpdr = count(p -> p.done == true, deprived) / length(deprived),
             delay = mean(delays),
             centrality = centrality,
             packetloss = model.packets_lost / model.transmitters_count )
end

"""
    start(model::AgentBasedModel, minutes = 1) → (received, produced)

Start an experiment from `model` for the given number of minutes.
"""
function start(model::AgentBasedModel; minutes = 1)
    steps = minutes * 60 * 1000

    run!(model, agent_step!, model_step!, steps)

    getstats(model)
end
