text:
# Protocol implementations for:
# - Client - Namenode
# - Client - Datanode
# ref: https://wiki.apache.org/hadoop/HadoopRpc

const HRPC_HEADER               = "hrpc"
const HRPC_VERSION              = 0x09
const HRPC_SERVICE_CLASS        = 0x00

const HRPC_AUTH_PROTOCOL_NONE   = 0x00
const HRPC_AUTH_PROTOCOL_SASL   = 0xdf # -33

const HRPC_AUTH_METHOD_SIMPLE   = 0x80
const HRPC_AUTH_METHOD_TOKEN    = 0x82

const HRPC_PROTOBUFF_TYPE       = RpcKindProto.RPC_PROTOCOL_BUFFER
const HRPC_FINAL_PACKET         = RpcRequestHeaderProto_OperationProto.RPC_FINAL_PACKET

const HRPC_CALL_ID_SASL         = Int32(-33)
const HRPC_CALL_ID_CONNCTX      = Int32(-3)
const HRPC_CALL_ID_NORMAL       = Int32(0)  # Post the handshake it cycles sequentially between 1:typemax(Int32)

const HRPC_PROTOCOLS = Dict(
    ClientNamenodeProtocolBlockingStub              => Dict(:id => "org.apache.hadoop.hdfs.protocol.ClientProtocol",            :ver => UInt64(1), :name => "HDFSClient"),
    ApplicationClientProtocolServiceBlockingStub    => Dict(:id => "org.apache.hadoop.yarn.api.ApplicationClientProtocolPB",    :ver => UInt64(1), :name => "YarnClient"),
    ApplicationMasterProtocolServiceBlockingStub    => Dict(:id => "org.apache.hadoop.yarn.api.ApplicationMasterProtocolPB",    :ver => UInt64(1), :name => "YarnAppMaster"),
    ContainerManagementProtocolServiceBlockingStub  => Dict(:id => "org.apache.hadoop.yarn.api.ContainerManagementProtocolPB",  :ver => UInt64(1), :name => "YarnNMClient")
)

@doc doc"""
HadoopRpcException is thrown on Rpc interaction errors either with namenode or datanode.
Field `status` contains error code (enum) if received from the connected entity or just ERROR (1) to indicate failure.
Though HadoopRpcException is used while communicating with both namenodes and datanodes, SUCCESS and ERROR are coded
with the same values in both cases. Other specific error codes need knowledge of the context to interpret.
""" ->
type HadoopRpcException <: Exception
    status::Int32
    message::AbstractString
    function HadoopRpcException(header::RpcResponseHeaderProto)
        new(header.status, header.errorMsg)
    end
    function HadoopRpcException(msg::AbstractString)
        new(RpcResponseHeaderProto_RpcStatusProto.ERROR, msg)
    end
    function HadoopRpcException(status::Int32, message::AbstractString)
        new(status, message)
    end
end

# Utility methods
const crc32 = crc(CRC_32)
const crc32c = crc(CRC_32_C)

immutable ChkSum{T} end
typealias ChkSumCRC32   ChkSum{ChecksumTypeProto.CHECKSUM_CRC32}
typealias ChkSumCRC32C  ChkSum{ChecksumTypeProto.CHECKSUM_CRC32C}

chksum(::Type{ChkSumCRC32}, c_data::Vector{UInt8}) = crc32(c_data)
chksum(::Type{ChkSumCRC32C}, c_data::Vector{UInt8}) = crc32c(c_data)
chksum{T}(::Type{ChkSum{T}}, c_data::Vector{UInt8}) = throw(HadoopRpcException("Unknown CRC type $T"))
isvalid_chksum(::Type{ChkSumCRC32}) = true
isvalid_chksum(::Type{ChkSumCRC32C}) = true
isvalid_chksum{T}(::Type{ChkSum{T}}) = false

function _len_uleb{T <: Integer}(x::T)
    nw = 1
    while ((x >>>= 7) != 0)
        nw += 1
    end
    nw
end

function buffer_size_delimited(channelbuff::IOBuffer, obj)
    iob = IOBuffer()
    writeproto(iob, obj)

    data = takebuf_array(iob)
    len = write_bytes(channelbuff, data)
    @logmsg("$(typeof(obj)) -> $data")
    @logmsg("$(typeof(obj)) -> buffer. len $len")
    len
end

function send_buffered(buff::IOBuffer, sock::TCPSocket, delimited::Bool)
    data = takebuf_array(buff)
    len::UInt32 = 0
    if delimited
        datalen = UInt32(length(data))
        len = write(sock, hton(datalen))
    end
    len += write(sock, data)
    @logmsg("buffer -> sock. len $len")
    len
end

function recv_blockop(sock::TCPSocket)
    @logmsg("recv block read message")
    data_bytes = read_bytes(sock)
    @logmsg("block_resp <- sock. len $(length(data_bytes))")

    block_resp = BlockOpResponseProto()
    readproto(IOBuffer(data_bytes), block_resp)

    (block_resp.status == Status.SUCCESS) || throw(HadoopRpcException(block_resp.status, "Error in block operation: $(block_resp.message)"))
    return block_resp
end

#
# Hadoop RPC Channel
type HadoopRpcController <: ProtoRpcController
    debug::Bool
end

type HadoopRpcChannel <: ProtoRpcChannel
    host::AbstractString
    port::Integer
    protocol_attribs::Dict
    call_id::Int32          # see: RpcRequestHeaderProto.callId
    sent_call_id::Int32     # set to the last call id sent for verification purpose
    clnt_id::AbstractString # string(Base.Random.uuid4())
    ugi::UserGroupInformation
    iob::IOBuffer
    sock::Nullable{TCPSocket}

    function HadoopRpcChannel(host::AbstractString, port::Integer, ugi::UserGroupInformation, protocol::DataType)
        protocol = HRPC_PROTOCOLS[protocol]
        call_id = HRPC_CALL_ID_CONNCTX
        clnt_id = string(Base.Random.uuid4())[1:16]
        new(host, port, protocol, call_id, call_id, clnt_id, ugi, IOBuffer(), Nullable{TCPSocket}())
    end
end

function show(io::IO, ch::HadoopRpcChannel)
    user = username(ch.ugi)
    user_spec = isempty(user) ? user : "$(user)@"
    println(io, "$(user_spec)$(ch.host):$(ch.port)/")
    println(io, "    id: $(ch.clnt_id)")
    println(io, "    connected: $(isconnected(ch))")
    nothing
end

isconnected(channel::HadoopRpcChannel) = !isnull(channel.sock) && isopen(get(channel.sock))
begin_send(channel::HadoopRpcChannel) = Base.truncate(channel.iob, 0)
send_buffered(channel::HadoopRpcChannel, delimited::Bool) = send_buffered(channel.iob, get(channel.sock), delimited::Bool)

function next_call_id(channel::HadoopRpcChannel)
    id = channel.sent_call_id = channel.call_id

    if id == HRPC_CALL_ID_SASL
        channel.call_id = HRPC_CALL_ID_CONNCTX
    elseif id == HRPC_CALL_ID_CONNCTX
        channel.call_id = HRPC_CALL_ID_NORMAL
    elseif id < typemax(Int32)
        channel.call_id = id + Int32(1)
    else
        channel.call_id = HRPC_CALL_ID_NORMAL
    end
    id
end

function buffer_handshake(channel::HadoopRpcChannel, authprotocol::UInt8)
    write(channel.iob, [HRPC_HEADER.data, HRPC_VERSION, HRPC_SERVICE_CLASS, authprotocol;])
end

function buffer_connctx(channel::HadoopRpcChannel)
    protocol = channel.protocol_attribs[:id]
    connctx = protobuild(IpcConnectionContextProto, Dict(:userInfo => channel.ugi.userinfo, :protocol => protocol))

    buffer_size_delimited(channel.iob, connctx)
end

function buffer_reqhdr(channel::HadoopRpcChannel, call_id::Int32)
    hdr = protobuild(RpcRequestHeaderProto, Dict(:rpcKind => HRPC_PROTOBUFF_TYPE,
                :rpcOp => HRPC_FINAL_PACKET,
                :callId => call_id,
                #:retryCount => -1,
                :clientId => channel.clnt_id.data))

    buffer_size_delimited(channel.iob, hdr)
end

buffer_rpc_reqhdr(channel::HadoopRpcChannel) = buffer_reqhdr(channel, next_call_id(channel))
buffer_conctx_reqhdr(channel::HadoopRpcChannel) = (channel.sent_call_id = channel.call_id = HRPC_CALL_ID_CONNCTX; buffer_reqhdr(channel, HRPC_CALL_ID_CONNCTX))

function buffer_method_reqhdr(channel::HadoopRpcChannel, method::MethodDescriptor)
    protocol = channel.protocol_attribs[:id]
    protocol_ver = channel.protocol_attribs[:ver]
    hdr = protobuild(RequestHeaderProto, Dict(:methodName => method.name,
                :declaringClassProtocolName => protocol,
                :clientProtocolVersion => protocol_ver))

    buffer_size_delimited(channel.iob, hdr)
end

# Connect to the hadoop server and complete the handshake
function connect(channel::HadoopRpcChannel)
    # open connection
    @logmsg("connecting to HadoopRpcChannel $(channel.host):$(channel.port)")
    try
        sock = connect(channel.host, channel.port)
        channel.sock = Nullable{TCPSocket}(sock)

        # negotiate sasl authentication if ugi has appropriate tokens
        if !conditional_sasl_auth(channel)
            # else do a simple handshake
            begin_send(channel)
            buffer_handshake(channel, HRPC_AUTH_PROTOCOL_NONE)
            send_buffered(channel, false)
        end

        begin_send(channel)
        buffer_conctx_reqhdr(channel)
        buffer_connctx(channel)
        send_buffered(channel, true)
        @logmsg("connected to HadoopRpcChannel $(channel.host):$(channel.port)")
        channel.call_id = HRPC_CALL_ID_NORMAL
    catch ex
        @logmsg("exception connecting to $(channel.host):$(channel.port): $ex")
        disconnect(channel)
        rethrow(ex)
    end
    nothing
end

# Disconnect and reset the channel instance.
# Client Id is also reset and so this is in essence an entirely new channel.
function disconnect(channel::HadoopRpcChannel)
    try
        isconnected(channel) && close(get(channel.sock))
    catch ex
        @logmsg("exception while closing channel socket $ex")
    end
    channel.sock = Nullable{TCPSocket}()
    channel.call_id = HRPC_CALL_ID_CONNCTX
    channel.clnt_id = string(Base.Random.uuid4())[1:16]
end

function send_rpc_message(channel::HadoopRpcChannel, controller::HadoopRpcController, method::MethodDescriptor, params)
    isconnected(channel) || connect(channel)

    try
        @logmsg("send rpc message. method: $(method.name)")
        begin_send(channel)
        buffer_rpc_reqhdr(channel)
        buffer_method_reqhdr(channel, method)
        buffer_size_delimited(channel.iob, params)
        send_buffered(channel, true)
    catch ex
        @logmsg("exception sending to $(channel.host):$(channel.port): $ex")
        disconnect(channel)
        rethrow(ex)
    end
    nothing
end

function recv_rpc_message(channel::HadoopRpcChannel, resp)
    try
        @logmsg("recv rpc message")
        msg_len = ntoh(read(get(channel.sock), UInt32))
        hdr_bytes = read_bytes(get(channel.sock))
        @logmsg("hdr <- sock. len $(length(hdr_bytes))")

        resp_hdr = RpcResponseHeaderProto()
        readproto(IOBuffer(hdr_bytes), resp_hdr)

        (resp_hdr.callId == reinterpret(UInt32, channel.sent_call_id)) || throw(HadoopRpcException("unknown callid. received:$(resp_hdr.callId) sent:$(channel.sent_call_id). status: $(resp_hdr.status)"))
        (resp_hdr.status == RpcResponseHeaderProto_RpcStatusProto.SUCCESS) || throw(HadoopRpcException(resp_hdr))

        if resp_hdr.status == RpcResponseHeaderProto_RpcStatusProto.SUCCESS
            hdr_len = UInt32(length(hdr_bytes))
            hdr_len += _len_uleb(hdr_len)
            if msg_len > hdr_len
                data_bytes = read_bytes(get(channel.sock))
                @logmsg("data <- sock. len $(length(data_bytes))")
                data_len = msg_len - hdr_len - _len_uleb(length(data_bytes))
                (length(data_bytes) == data_len) || throw(HadoopRpcException("unexpected response data length. expected:$(data_len) read:$(length(data_bytes))"))
                readproto(IOBuffer(data_bytes), resp)
            end
        end
    catch ex
        @logmsg("exception receiving from $(channel.host):$(channel.port): $ex")
        disconnect(channel)
        rethrow(ex)
    end
    nothing
end

function call_method(channel::HadoopRpcChannel, method::MethodDescriptor, controller::HadoopRpcController, params)
    @logmsg("call_method $(method.name)")
    send_rpc_message(channel, controller, method, params)
    resp_type = get_response_type(method)
    resp = resp_type()
    recv_rpc_message(channel, resp)
    resp
end

#
# Hadoop Data Channel. Base type for communicating with data nodes.

# Opcodes
const HDATA_VERSION             = Int16(28)
const HDATA_WRITE_BLOCK         = Int8(80)
const HDATA_READ_BLOCK          = Int8(81)
const HDATA_READ_METADATA       = Int8(82)
const HDATA_REPLACE_BLOCK       = Int8(83)
const HDATA_COPY_BLOCK          = Int8(84)
const HDATA_BLOCK_CHECKSUM      = Int8(85)
const HDATA_TRANSFER_BLOCK      = Int8(86)

@doc doc"""
HadoopDataChannel is the connection to a datanode.
Also associated with it is a buffer to help in message preparation.
""" ->
type HadoopDataChannel
    host::AbstractString
    port::Integer
    iob::IOBuffer
    sock::Nullable{TCPSocket}

    HadoopDataChannel(host::AbstractString, port::Integer) = new(host, port, IOBuffer(), Nullable{TCPSocket}())
end

function connect(channel::HadoopDataChannel)
    @logmsg("connecting to HadoopDataChannel $(channel.host):$(channel.port)")
    try
        sock = connect(channel.host, channel.port)
        channel.sock = Nullable(sock)
    catch ex
        @logmsg("exception connecting to HadoopDataChannel $(channel.host):$(channel.port): $ex")
        rethrow(ex)
    end
    @logmsg("connected to HadoopDataChannel $(channel.host):$(channel.port)")
    nothing
end

function disconnect(channel::HadoopDataChannel)
    try
        isconnected(channel) && close(get(channel.sock))
    catch ex
        @logmsg("exception while closing HadoopDataChannel socket $ex")
    end
    channel.sock = Nullable{TCPSocket}()
    nothing
end

isconnected(channel::HadoopDataChannel) = !isnull(channel.sock) && isopen(get(channel.sock))
begin_send(channel::HadoopDataChannel) = Base.truncate(channel.iob, 0)
send_buffered(channel::HadoopDataChannel, delimited::Bool) = send_buffered(channel.iob, get(channel.sock), delimited::Bool)

function buffer_opcode(channel::HadoopDataChannel, opcode::Int8)
    hver = UInt16(HDATA_VERSION)
    len = write(channel.iob, hton(hver))
    len += write(channel.iob, opcode)
end

function _select_node_for_block(block::LocatedBlockProto)
    node_infos = DatanodeInfoProto[]
    for loc in block.locs
        (loc.adminState == DatanodeInfoProto_AdminState.NORMAL) && push!(node_infos, loc)
    end
    isempty(node_infos) && (return Nullable{DatanodeInfoProto}())
    # TODO: algo to select best node
    return Nullable(node_infos[1])
end

@doc doc"""
HadoopDataChannelPool is a connection pool that holds connections to datanodes.
Connections are deemed stale after `keepalivesecs`.
""" ->
type HadoopDataChannelPool
    free::Dict{AbstractString,Vector}
    keepalivesecs::UInt64

    HadoopDataChannelPool(keepalivesecs::Integer) = new(Dict{AbstractString,Vector}(), keepalivesecs)
end

const _dcpool = HadoopDataChannelPool(30)
function _get(pool::HadoopDataChannelPool, host::AbstractString, port::Integer)
    dnid = "$host:$port"
    free = (dnid in keys(pool.free)) ? pool.free[dnid] : []

    timediff = _dcpool.keepalivesecs
    while !isempty(free) && (timediff >= _dcpool.keepalivesecs)
        channel,lastusetime = shift!(free)
        timediff = round(UInt64, time() - lastusetime)
    end

    (timediff < _dcpool.keepalivesecs) || (channel = HadoopDataChannel(host, port))
    @logmsg("return channel: $channel connected: $(isconnected(channel))")
    channel
end

function _put(pool::HadoopDataChannelPool, channel::HadoopDataChannel, reuse::Bool)
    isconnected(channel) || return
    @logmsg("keeping channel: $channel connected: $(isconnected(channel))")
    if !reuse
        try
            disconnect(channel)
        finally
            return
        end
    end
    dnid = "$(channel.host):$(channel.port)"
    free = dnid in keys(pool.free) ? pool.free[dnid] : (pool.free[dnid] = [])
    for (ch,tm) in free
        (object_id(channel) == object_id(ch)) && return
    end
    push!(free, (channel,time()))
end

#
# Begin Block Reader implementation to read data

@doc doc"""
HDFSBlockReader reads one block of data from a datanode.
""" ->
type HDFSBlockReader
    channel::HadoopDataChannel

    # block and region assigned
    block::LocatedBlockProto
    offset::UInt64
    len::UInt64

    # protocol data
    block_op_resp::Nullable{BlockOpResponseProto}
    total_read::UInt64

    pkt_hdr::Nullable{PacketHeaderProto}
    packet_len::UInt64
    packet_read::UInt64

    chunk_len::UInt64
    chunk_count::UInt64
    chunks_read::UInt64
    chunk::Array{UInt8,1}

    checksums::Array{UInt32,1}
    chk_crc::Bool

    initiated::Bool

    function HDFSBlockReader(host::AbstractString, port::Integer, block::LocatedBlockProto, offset::UInt64, len::UInt64, chk_crc::Bool)
        channel = _get(_dcpool, host, port)
        @logmsg("creating block reader for offset $offset at $host:$port for length $len")
        new(channel, block, offset, len,
            Nullable{BlockOpResponseProto}(), 0,
            Nullable{PacketHeaderProto}(), 0, 0,
            0, 0, 0, UInt8[],
            UInt32[], chk_crc, false)
    end
end

isconnected(reader::HDFSBlockReader) = isconnected(reader.channel)
function disconnect(reader::HDFSBlockReader, reuse::Bool)
    isconnected(reader.channel) && _put(_dcpool, reader.channel, reuse)

    reader.block_op_resp = Nullable{BlockOpResponseProto}()
    reader.total_read = 0
    reader.pkt_hdr = Nullable{PacketHeaderProto}()
    reader.packet_len = 0
    reader.packet_read = 0
    reader.chunk_len = 0
    reader.chunk_count = 0
    reader.chunks_read = 0
    reader.chunk = UInt8[]
    reader.checksums = UInt32[]
    reader.initiated = false
    nothing
end

function buffer_readblock(reader::HDFSBlockReader)
    channel = reader.channel
    block = reader.block
    offset = reader.offset
    len = reader.len

    token = TokenProto()
    for fld in (:identifier, :password, :kind, :service)
        set_field!(token, fld, get_field(block.blockToken, fld))
    end

    exblock = ExtendedBlockProto()
    for fld in (:poolId, :blockId, :generationStamp)
        set_field!(exblock, fld, get_field(block.b, fld))
    end

    basehdr = protobuild(BaseHeaderProto, Dict(:block => exblock, :token => token))
    hdr = protobuild(ClientOperationHeaderProto, Dict(:baseHeader => basehdr, :clientName => ELLY_CLIENTNAME))
    readblock = protobuild(OpReadBlockProto, Dict(:offset => offset, :len => len, :header => hdr))
    @logmsg("sending block read message for offset $offset len $len")

    buffer_size_delimited(channel.iob, readblock)
end

function buffer_client_read_status(reader::HDFSBlockReader, status::Int32)
    channel = reader.channel
    read_status = protobuild(ClientReadStatusProto, Dict(:status => status))
    buffer_size_delimited(channel.iob, read_status)
end

function send_block_read(reader::HDFSBlockReader)
    channel = reader.channel
    isconnected(channel) || connect(channel)

    try
        @logmsg("send block read message")
        begin_send(channel)
        buffer_opcode(channel, HDATA_READ_BLOCK)
        buffer_readblock(reader)
        send_buffered(channel, false)
    catch ex
        @logmsg("exception sending to $(channel.host):$(channel.port): $ex")
        disconnect(reader, false)
        rethrow(ex)
    end
    nothing
end

function send_read_status(reader::HDFSBlockReader, status::Int32=Status.SUCCESS)
    channel = reader.channel
    try
        @logmsg("send read status $status")
        buffer_client_read_status(reader, status)
        send_buffered(channel, false)
    catch ex
        @logmsg("exception sending to $(channel.host):$(channel.port): $ex")
        disconnect(reader, false)
        rethrow(ex)
    end
    nothing
end

function recv_blockop(reader::HDFSBlockReader)
    channel = reader.channel
    try
        block_resp = recv_blockop(get(channel.sock))
        checksum_type = block_resp.readOpChecksumInfo.checksum._type
        isvalid_chksum(ChkSum{checksum_type}) || throw(HadoopRpcException("Unknown checksum type $checksum_type"))
        reader.block_op_resp = Nullable(block_resp)
    catch ex
        @logmsg("exception receiving from $(channel.host):$(channel.port): $ex")
        disconnect(reader, false)
        rethrow(ex)
    end
    nothing
end

function recv_packet_hdr(reader::HDFSBlockReader)
    channel = reader.channel
    try
        @logmsg("recv block packet message")
        sock = get(channel.sock)

        pkt_len = ntoh(read(sock, UInt32))
        hdr_len = ntoh(read(sock, UInt16))
        hdr_bytes = Array(UInt8, hdr_len)
        read!(sock, hdr_bytes)
        @logmsg("pkt_hdr <- sock. len $(hdr_len) (pkt_len: $pkt_len)")

        pkt_hdr = PacketHeaderProto()
        readproto(IOBuffer(hdr_bytes), pkt_hdr)

        data_len = pkt_hdr.dataLen
        block_op_resp = get(reader.block_op_resp)
        reader.chunk_len = block_op_resp.readOpChecksumInfo.checksum.bytesPerChecksum
        reader.chunk = Array(UInt8, reader.chunk_len)
        reader.chunk_count = div(data_len + reader.chunk_len - 1, reader.chunk_len)    # chunk_len-1 to take care of chunks with partial data
        reader.chunks_read = 0
        @logmsg("received packet with $(reader.chunk_count) chunks of $(reader.chunk_len) bytes each in $data_len bytes of data")
        @logmsg("last packet: $(pkt_hdr.lastPacketInBlock)")

        checksums = Array(UInt32, reader.chunk_count)
        read!(sock, checksums)
        @logmsg("checksums <- sock. len $(sizeof(checksums))")
        for idx in 1:length(checksums)
            checksums[idx] = ntoh(checksums[idx])
        end
        reader.packet_len = pkt_len
        #reader.packet_read = 4 + 2 + hdr_len + sizeof(checksums)
        reader.packet_read = 4 + sizeof(checksums)
        reader.pkt_hdr = Nullable(pkt_hdr)
        reader.checksums = checksums
        @logmsg("current read position pkt $(reader.packet_read)/$(reader.packet_len), block $(reader.total_read)/$(reader.len)")
        nothing
    catch ex
        @logmsg("exception receiving from $(channel.host):$(channel.port): $ex")
        disconnect(reader, false)
        rethrow(ex)
    end
end

eofpacket(reader::HDFSBlockReader) = (reader.packet_read >= reader.packet_len)
eof(reader::HDFSBlockReader) = (reader.total_read >= reader.len)

function verify_pkt_checksums(reader::HDFSBlockReader, buff::Vector{UInt8})
    data_len = length(buff)
    @logmsg("verifying packet crc data_len: $data_len")
    (data_len > 0) || return

    block_op_resp = get(reader.block_op_resp)
    checksum_type = block_op_resp.readOpChecksumInfo.checksum._type
    C = ChkSum{checksum_type}
    if !isvalid_chksum(C)
        disconnect(reader, false)
        throw(HadoopRpcException("Unknown CRC type $checksum_type"))
    end
    verify_pkt_checksums(reader, buff, C)
end

function verify_pkt_checksums{T}(reader::HDFSBlockReader, buff::Vector{UInt8}, C::Type{ChkSum{T}})
    chunk_len = reader.chunk_len
    offset = 1
    checksums = reader.checksums
    for idx in 1:(length(checksums)-1)
        cs = chksum(C, unsafe_wrap(Array, pointer(buff, offset), chunk_len))
        (cs == checksums[idx]) || throw(HadoopRpcException("Checksum mismatch at chunk $(idx). Expected $(checksums[idx]), got $(cs)"))
        offset += chunk_len
    end

    cs = chksum(unsafe_wrap(Array, pointer(buff, offset), data_len-(offset-1)))
    (cs == checksums[end]) || throw(HadoopRpcException("Checksum mismatch at last chunk $(length(checksums)). Expected $(checksums[end]), got $(cs)"))
    nothing
end

function initiate(reader::HDFSBlockReader; retry=true)
    try
        isconnected(reader.channel) || connect(reader.channel)
        send_block_read(reader)
        recv_blockop(reader)
        reader.initiated = true
        return
    catch ex
        if retry
            @logmsg("retrying block reader initiation")
            disconnect(reader.channel)
            connect(reader.channel)
            return initiate(reader; retry=false)
        else
            rethrow()
        end
    end
end

@doc doc"""
Read one packet into `inbuff` starting from `offset`.
If `inbuff` has insufficient space, returns the minimum additional space required in `inbuff` to read the packet as a negative number.
Otherwise, returns the number of bytes available in `inbuff` after reading the packet.
""" ->
function read_packet!(reader::HDFSBlockReader, inbuff::Vector{UInt8}, offset::UInt64)
    if !reader.initiated
        # initiate the stream
        initiate(reader)
        recv_packet_hdr(reader)
    elseif eofpacket(reader)
        # recv the next packet
        recv_packet_hdr(reader)
    end

    packet_remaining = reader.packet_len - reader.packet_read
    excess = Int64(length(inbuff)+1-offset) - Int64(packet_remaining)
    (excess >= 0) || return excess

    buff = unsafe_wrap(Array, pointer(inbuff, offset), packet_remaining)

    channel = reader.channel
    sock = get(channel.sock)

    read!(sock, buff)
    reader.packet_read += packet_remaining
    reader.total_read += packet_remaining
    @logmsg("recv $(packet_remaining) packet_read $(reader.packet_read), total_read $(reader.total_read)")

    # verify crc only if required
    reader.chk_crc && verify_pkt_checksums(reader, buff)

    # send the success block read status
    if eof(reader)
        @logmsg("recv last empty packet")
        recv_packet_hdr(reader)
        send_read_status(reader)
    end
    excess
end

#
# Begin Block Writer implementation for writing / appending data

@doc doc"""
PipelinedPacket holds a packet of data while it is sent to datanodes and acknowledged.
The data bytes are reset once acks are received successfully to reduce memory usage.
""" ->
type PipelinedPacket
    seqno::Int64
    hdr::PacketHeaderProto
    bytes::Vector{UInt8}
    checksums::Vector{UInt32}
    acks::Vector{Int32}
end

@doc doc"""
WriterPipeline holds all packets of data for a block and provides methods to maintain
their states. Pipeline `failed` status is set if any of the acks received is a failure.
Field `acked_bytes` contains the count of bytes successfully sent till an error is
encountered.
""" ->
type WriterPipeline
    pkt_prepared::Vector{PipelinedPacket}
    pkt_ackwait::Vector{PipelinedPacket}
    pkt_ackrcvd::Vector{PipelinedPacket}
    acked_bytes::UInt64
    failed::Bool
    triggers::Vector{Condition}
    flush::Bool

    WriterPipeline() = new(PipelinedPacket[], PipelinedPacket[], PipelinedPacket[], 0, false, [Condition(), Condition()], false)
end

function move(from::Vector{PipelinedPacket}, to::Vector{PipelinedPacket}, seqno::Int64)
    for idx in 1:length(from)
        pkt = from[idx]
        if pkt.seqno == seqno
            push!(to, splice!(from, idx))
            return pkt
        end
    end
    throw(HadoopRpcException("Packet seqno:$seqno not found in pipeline"))
end

# since these use condition variables that are edge triggered,
# take care to ensure atomicity of the methods (do not introduce anything that may yield)
wait_pkt(pipeline::WriterPipeline) = pipeline.flush || wait(pipeline.triggers[1])
wait_flushed(pipeline::WriterPipeline) = pipeline.failed || wait(pipeline.triggers[2])
trigger_pkt(pipeline::WriterPipeline) = notify(pipeline.triggers[1])
trigger_flushed(pipeline::WriterPipeline) = notify(pipeline.triggers[2])
function flush_and_wait(pipeline::WriterPipeline)
    pipeline.flush = true
    trigger_pkt(pipeline)
    wait_flushed(pipeline)
    pipeline.failed && throw(HadoopRpcException("Error flushing pipeline"))
    nothing
end

@inline function enqueue(pipeline::WriterPipeline, pkt::PipelinedPacket)
    push!(pipeline.pkt_prepared, pkt)
    nothing
end

function ackwait(pipeline::WriterPipeline, seqno::Int64)
    pkt = move(pipeline.pkt_prepared, pipeline.pkt_ackwait, seqno)
    pkt.bytes = UInt8[]
    pkt.checksums = UInt32[]
    nothing
end

function ackrcvd(pipeline::WriterPipeline, seqno::Int64, acks::Vector{Int32})
    pkt = move(pipeline.pkt_ackwait, pipeline.pkt_ackrcvd, seqno)
    pkt.acks = acks
    for ack in acks
        (ack == Status.SUCCESS) || (pipeline.failed = true)
    end
    if !pipeline.failed
        pipeline.acked_bytes += pkt.hdr.dataLen
    end
    nothing
end

function assert_no_failure(pipeline::WriterPipeline)
    if pipeline.failed
        throw(HadoopRpcException("Received failure ack from datanode. Pipeline proecssing has failed"))
    end
    nothing
end

pending(pipeline::WriterPipeline) = (length(pipeline.pkt_prepared) + length(pipeline.pkt_ackwait))

# TODO: write optimizations
# - schedule a block only after block is complete
# - schedule and manage multiple block streams to different nodes
# - select data node to write to based on whether it already has an active write to the datanode
# - maintain a list of IOBuffers and LocatedBlockProtos in HDFSFileWriters that are queued / being written to, with a max limit beyond which it needs to wait for blocks to be written.
# - close of HDFSFileWriter should wait for all blocks to be flushed
# - this can later be tweaked to incorporate node failover
@doc doc"""
HDFSBlockWriter writes one block worth of data to a datanode.

The `write` method returns number of bytes accepted for writing into this block, which may be less than
what was requested if block is full.

Packets are sent and acks received as and when data being written exceed packet size limit.

The `close` operation flushes remaining data in the block and waits for all pending acks.
""" ->
type HDFSBlockWriter
    channel::HadoopDataChannel
    server_defaults::FsServerDefaultsProto

    block::LocatedBlockProto
    source_node::DatanodeInfoProto
    buffer::IOBuffer

    total_written::UInt64

    pkt_seq::Int64
    pkt_pipeline::WriterPipeline
    initiated::Bool
    pipeline_task::Task

    function HDFSBlockWriter(block::LocatedBlockProto, defaults::FsServerDefaultsProto)
        nnode_info = _select_node_for_block(block)
        isnull(nnode_info) && throw(HadoopRpcException("Could not get a valid datanode to write"))

        node_info = get(nnode_info)
        node = node_info.id
        host = node.ipAddr
        port = node.xferPort

        channel = _get(_dcpool, host, port)
        @logmsg("creating block writer for offset $(block.offset) at $host:$port")
        writer = new(channel, defaults, block, node_info, PipeBuffer(), 0, 0, WriterPipeline(), false)
        writer.pipeline_task = @async process_pipeline(writer)
        writer
    end
end

isconnected(writer::HDFSBlockWriter) = isconnected(writer.channel)
function disconnect(writer::HDFSBlockWriter, reuse::Bool)
    if !istaskdone(writer.pipeline_task)
        reuse = false
        trigger_pkt(writer.pkt_pipeline)
    end
    isconnected(writer.channel) && _put(_dcpool, writer.channel, reuse)
    writer.initiated = false
    nothing
end

function connect(writer::HDFSBlockWriter)
    connect(writer.channel)
    nothing
end

function check_write_buffer(writer::HDFSBlockWriter)
    defaults = writer.server_defaults
    max_buffer_pkt = div(defaults.blockSize, defaults.writePacketSize) + 1
    max_buffer_bytes = defaults.writePacketSize * max_buffer_pkt
    buffered_bytes = nb_available(writer.buffer)
    if ((buffered_bytes >= max_buffer_bytes) || (pending(writer.pkt_pipeline) >= max_buffer_pkt))
        @logmsg("slow pipeline. pending bytes:$(nb_available(writer.buffer)), packets:$(pending(writer.pkt_pipeline)). waiting...")
        trigger_pkt(writer.pkt_pipeline)
        wait_pkt(writer.pkt_pipeline)
    elseif buffered_bytes >= defaults.writePacketSize
        @logmsg("triggering pkt_pipeline. buffered:$(nb_available(writer.buffer)), pending packets:$(pending(writer.pkt_pipeline))")
        trigger_pkt(writer.pkt_pipeline)
        yield()
    end
end

# write only enough till blockSize
_can_write(writer::HDFSBlockWriter, n::Int) = min(writer.server_defaults.blockSize - writer.total_written - nb_available(writer.buffer), n)

function write(writer::HDFSBlockWriter, x::UInt8)
    nbytes = _can_write(writer, 1)
    (nbytes > 0) && write(writer.buffer, x)
    check_write_buffer(writer)
    nbytes
end

function write(writer::HDFSBlockWriter, buff::Vector{UInt8})
    nbytes = _can_write(writer, length(buff))
    (nbytes > 0) && Base.write_sub(writer.buffer, buff, 1, nbytes)
    check_write_buffer(writer)
    nbytes
end

function flush(writer::HDFSBlockWriter)
    @logmsg("flush of block writer invoked")
    flush_and_wait(writer.pkt_pipeline)
end

function initiate(writer::HDFSBlockWriter; retry=true)
    try
        isconnected(writer.channel) || connect(writer.channel)
        send_block_write(writer)
        recv_blockop(writer)
        writer.initiated = true
        return
    catch ex
        if retry
            @logmsg("retrying block write initiation")
            disconnect(writer.channel)
            connect(writer.channel)
            return initiate(writer; retry=false)
        else
            rethrow()
        end
    end
end

function process_pipeline(writer::HDFSBlockWriter)
    @logmsg("process_pipeline: started")
    pipeline = writer.pkt_pipeline
    defaults = writer.server_defaults
    flushed = false
    failed = pipeline.failed

    loop_state = 0
    try
        readable_bytes = 0
        while !(failed || flushed)
            loop_state = 0
            if (nb_available(writer.buffer) < defaults.writePacketSize) &&
                isempty(pipeline.pkt_prepared) &&
                (isempty(pipeline.pkt_ackwait) || (readable_bytes == 0))
                wait_pkt(pipeline)
            end

            writer.initiated || initiate(writer)

            channel = writer.channel
            sock = get(channel.sock)
            readable_bytes = nb_available(sock)

            @logmsg("processing pipeline. pending $(nb_available(writer.buffer))bytes -> [$(length(pipeline.pkt_prepared)) -> $(length(pipeline.pkt_ackwait)) -> $(length(pipeline.pkt_ackrcvd))]pkts <- $(readable_bytes)bytes flush:$(flushed)")
            while nb_available(writer.buffer) >= defaults.writePacketSize
                prepare_packet(writer)
            end

            flushed = pipeline.flush

            if flushed
                while nb_available(writer.buffer) > 0
                    @logmsg("flushing remaining data packet...")
                    prepare_packet(writer)
                end
                @logmsg("preparing last empty packet...")
                prepare_packet(writer)
            end

            loop_state = 1
            @logmsg("writing packets...")
            while !isempty(pipeline.pkt_prepared)
                pkt = pipeline.pkt_prepared[1]
                write_packet(writer, pkt)
            end

            @logmsg("reading acks...")
            loop_state = 2
            while !failed && !isempty(pipeline.pkt_ackwait) && (flushed || (readable_bytes > 0))
                read_packet_ack(writer)
                failed = pipeline.failed
                readable_bytes = nb_available(sock)
            end
            @logmsg("processed pipeline. pending $(nb_available(writer.buffer))bytes -> [$(length(pipeline.pkt_prepared)) -> $(length(pipeline.pkt_ackwait)) -> $(length(pipeline.pkt_ackrcvd))]pkts <- $(readable_bytes)bytes flush:$(flushed)")
            trigger_pkt(pipeline)
        end
    catch ex
        channel = writer.channel
        extype = (loop_state == 1) ? "send" : (loop_state == 2) ? "read" : "unknown"
        @logmsg("pipeline for $(channel.host):$(channel.port) failed. stage $extype. $ex")
        disconnect(writer, false)
        failed = pipeline.failed = true
    end
    @logmsg("process_pipeline finished. failed: $(failed), flushed: $(flushed)")
    trigger_pkt(pipeline)
    trigger_flushed(pipeline)
end

function buffer_writeblock(writer::HDFSBlockWriter)
    channel = writer.channel
    block = writer.block
    defaults = writer.server_defaults

    token = TokenProto()
    for fld in (:identifier, :password, :kind, :service)
        set_field!(token, fld, get_field(block.blockToken, fld))
    end

    exblock = ExtendedBlockProto()
    for fld in (:poolId, :blockId, :generationStamp)
        set_field!(exblock, fld, get_field(block.b, fld))
    end
    set_field!(exblock, :numBytes, zero(UInt64))

    basehdr = protobuild(BaseHeaderProto, Dict(:block => exblock, :token => token))
    hdr = protobuild(ClientOperationHeaderProto, Dict(:baseHeader => basehdr, :clientName => ELLY_CLIENTNAME))
    chksum = protobuild(ChecksumProto, Dict(:_type => defaults.checksumType, :bytesPerChecksum => defaults.bytesPerChecksum))

    targets = DatanodeInfoProto[]
    for loc in block.locs
        (loc.id != writer.source_node.id) && push!(targets, loc)
    end

    writeblock = protobuild(OpWriteBlockProto, Dict(:header => hdr,
                    :targets => targets,
                    :source => writer.source_node,
                    :stage => OpWriteBlockProto_BlockConstructionStage.PIPELINE_SETUP_CREATE,
                    :pipelineSize => length(block.locs),
                    :minBytesRcvd => 0,
                    :maxBytesRcvd => 0,
                    :latestGenerationStamp => exblock.generationStamp,
                    :requestedChecksum => chksum))
    @logmsg("sending block write message for offset $(block.offset)")

    buffer_size_delimited(channel.iob, writeblock)
end

function send_block_write(writer::HDFSBlockWriter)
    isconnected(writer) || connect(writer)
    channel = writer.channel

    try
        @logmsg("send block write message")
        begin_send(channel)
        buffer_opcode(channel, HDATA_WRITE_BLOCK)
        buffer_writeblock(writer)
        send_buffered(channel, false)
    catch ex
        @logmsg("exception sending to $(channel.host):$(channel.port): $ex")
        disconnect(writer, false)
        rethrow(ex)
    end
    nothing
end

function recv_blockop(writer::HDFSBlockWriter)
    channel = writer.channel
    try
        recv_blockop(get(channel.sock))
    catch ex
        @logmsg("exception receiving from $(channel.host):$(channel.port): $ex")
        disconnect(writer, false)
        rethrow(ex)
    end
    nothing
end

populate_checksums(bytes::Vector{UInt8}, chunk_len::UInt32, checksums::Vector{UInt32}, checksum_type::Int32) = populate_checksums(bytes, chunk_len, checksums, ChkSum{checksum_type})
function populate_checksums{T}(bytes::Vector{UInt8}, chunk_len::UInt32, checksums::Vector{UInt32}, C::Type{ChkSum{T}})
    isvalid_chksum(C) || throw(HadoopRpcException("Unknown checksum type $T"))
    nbytes = length(bytes)
    nchunks = length(checksums)
    c_offset = 1
    for idx in 1:nchunks
        c_len = min(nbytes-c_offset+1, chunk_len)
        c_data = unsafe_wrap(Array, pointer(bytes, c_offset), c_len)
        checksums[idx] = hton(chksum(C, c_data))
        #@logmsg("chksum $(checksums[idx])")
        c_offset += c_len
    end
    nothing
end

function prepare_packet(writer::HDFSBlockWriter)
    defaults = writer.server_defaults

    #@logmsg("prepare block packet")
    bytes_in_packet = min(defaults.writePacketSize, nb_available(writer.buffer))

    last_pkt = (bytes_in_packet == 0)
    seq_no = Int64(writer.pkt_seq += 1)
    #@logmsg("packet seqno $seq_no with $(bytes_in_packet)/$(defaults.writePacketSize) bytes is last packet: $last_pkt")

    pkt_hdr = protobuild(PacketHeaderProto, Dict(:offsetInBlock => writer.total_written,
                    :seqno => seq_no,
                    :lastPacketInBlock => last_pkt,
                    :dataLen => bytes_in_packet))

    writer.total_written += bytes_in_packet

    if bytes_in_packet == nb_available(writer.buffer)
        bytes = takebuf_array(writer.buffer)
    else
        bytes = Array(UInt8, bytes_in_packet)
        read!(writer.buffer, bytes)
    end

    chunk_len = defaults.bytesPerChecksum
    chunk_count = div(bytes_in_packet + chunk_len - 1, chunk_len)
    checksums = Array(UInt32, chunk_count)
    populate_checksums(bytes, chunk_len, checksums, defaults.checksumType)

    enqueue(writer.pkt_pipeline, PipelinedPacket(seq_no, pkt_hdr, bytes, checksums, Int32[]))
    nothing
end

const hdr_iob = IOBuffer()
function write_packet(writer::HDFSBlockWriter, pkt::PipelinedPacket)
    #@logmsg("write block packet")
    channel = writer.channel

    sock = get(channel.sock)
    pkt_len = UInt32(4 + sizeof(pkt.checksums) + sizeof(pkt.bytes))

    seek(hdr_iob, 0)
    hdr_len = UInt16(writeproto(hdr_iob, pkt.hdr))

    write(sock, hton(pkt_len))
    write(sock, hton(hdr_len))
    unsafe_write(sock, pointer(hdr_iob.data), hdr_iob.size)
    write(sock, pkt.checksums)
    write(sock, pkt.bytes)

    ackwait(writer.pkt_pipeline, pkt.seqno)
    #@logmsg("sent packet to $(channel.host):$(channel.port). total_written:$(writer.total_written)")
    nothing
end

function read_packet_ack(writer::HDFSBlockWriter)
    #@logmsg("reading packet ack")
    channel = writer.channel

    sock = get(channel.sock)
    #@logmsg("recv packet ack message")
    data_bytes = read_bytes(sock)
    #@logmsg("ack <- sock. len $(length(data_bytes))")

    ack = PipelineAckProto()
    readproto(IOBuffer(data_bytes), ack)

    pipeline = writer.pkt_pipeline
    ackrcvd(pipeline, ack.seqno, ack.reply)

    exblk = writer.block.b
    set_field!(exblk, :numBytes, pipeline.acked_bytes)

    #@logmsg("received ack for seqno: $(ack.seqno), status: $(ack.reply) bytes acked: $(exblk.numBytes)")
    nothing
end

@doc doc"""
HadoopRpcProtocol binds a channel and controller with a service protocol implementation.
Used by actual service implementations.
""" ->
type HadoopRpcProtocol{T<:AbstractProtoServiceStub}
    channel::HadoopRpcChannel
    controller::HadoopRpcController
    stub::T

    function HadoopRpcProtocol(host::AbstractString, port::Integer, ugi::UserGroupInformation=UserGroupInformation())
        channel = HadoopRpcChannel(host, port, ugi, T)
        controller = HadoopRpcController(false)
        stub = T(channel)

        new(channel, controller, stub)
    end
end

function show{T}(io::IO, client::HadoopRpcProtocol{T})
    srvcname = (HRPC_PROTOCOLS[T])[:name]
    print(io, "$(srvcname): ")
    show(io, client.channel)
end

set_debug(protocol::HadoopRpcProtocol, debug::Bool) = (protocol.controller.debug = debug)


text:
"""
Surprise Actor-critic
---------------------

Actor-Critic with a surprise-modulated learning rate for the actor.

See actor-critic.jl for some more implementational info.
See learner/actorcriticlearner.jl and learner/actorcriticmodulatedlearner.jl
for the core implementation.

Note 1
------
With particle filtering we need to run multiple times with different seeds.
As final LL we report the MEAN across seeds.
For signals/results we keep track of the ones of all seeds
and in the end report the one with the MAXIMUM LL.

Note 2
------
get_all_puzzles(), get_all_episodes_in_puzzle(), get_data_by_puzzle_episode()
are in rl_data_toolsforalgos.jl

"""
function getLL_ActorCritic_TDLambda_modulatedActor(data::Array{Int,2},
                                            alpha::Float64,
                                            surpriseprobability::Float64,
                                            stochasticity::Float64,
                                            eta::Float64,
                                            gamma::Float64,
                                            lambda_actor::Float64,
                                            lambda_critic::Float64,
                                            temp::Float64,
                                            nparticles::Union{Int, Float64};
                                            particlefilterseeds = [1,2,3],
                                            init_policy=nothing,
                                            init_Vs = nothing,
                                            use_replacing_trace=true,
                                            do_debug_print = false)
    nparticles = converttoInt(nparticles)
    data_stats, nr_states, nr_actions, nr_samples = get_data_stats_and_dimensions(data)
    if data_stats.nr_subjects != 1
        error("Invalid argument: getLL_ActorCritic_TDLambda_modulatedActor supports only one subject")
    end

    action_selection_probs_seeds = Array{Array{Float64, 1}, 1}(undef, length(particlefilterseeds))
    [action_selection_probs_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    LL_seeds = zeros(Float64, length(particlefilterseeds))

    Tsas_seeds = Array{Any, 1}(undef, length(particlefilterseeds))

    # signals:
    TD_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [TD_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    Surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [Surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    V_state_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [V_state_seeds[i_seed] = zeros(Float64, nr_samples + data_stats.nr_rewards) for i_seed in 1:length(particlefilterseeds)]

    γ_surprise_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [γ_surprise_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    SPE_seeds =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [SPE_seeds[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    finalV_record_seed = Array{Dict{Int,Any}, 1}(undef, length(particlefilterseeds))
    [finalV_record_seed[i_seed] = Dict{Int,Any}() for i_seed in 1:length(particlefilterseeds)]
    final_policyparameters_record_seed = Array{Dict{Int,Any}, 1}(undef, length(particlefilterseeds))
    [final_policyparameters_record_seed[i_seed] = Dict{Int,Any}() for i_seed in 1:length(particlefilterseeds)]

    PolicyParametersImbalance_record_seed  = Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyParametersImbalance_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]
    PolicyImbalance_record_seed = Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyImbalance_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    PolicyParametersDiffChosenVsUnchosen_record_seed =Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyParametersDiffChosenVsUnchosen_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    PolicyDiffChosenVsUnchosen_record_seed = Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyDiffChosenVsUnchosen_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    PolicyParameters_record_seed = Array{Array{Float64,1}, 1}(undef, length(particlefilterseeds))
    [PolicyParameters_record_seed[i_seed] = zeros(Float64, nr_samples) for i_seed in 1:length(particlefilterseeds)]

    for i_seed in 1:length(particlefilterseeds)
        current_puzzle_id = -1
        Testimate = []
        learner = []
        j = 1
        for i = 1:nr_samples
            if current_puzzle_id != data[i,Int(c_puzzle)]
                # Puzzle = independent block of the experiment.
                # The code was made for general purposes, in this experiment we
                # have only ONE puzzle.
                current_puzzle_id = data[i,Int(c_puzzle)]
                if do_debug_print
                    println("new puzzle: $current_puzzle_id")
                end

                # Particle Filter
                Testimate = TParticleFilterJump(nr_states = nr_states,
                                        nr_actions = nr_actions,
                                        nparticles = nparticles,
                                        surpriseprobability = surpriseprobability,
                                        stochasticity = stochasticity,
                                        seed = particlefilterseeds[i_seed])

                if init_policy == nothing # providing init_policy w/o init_Vs is considered a bug
                    learner = ActorCriticModulated(nr_states = nr_states,
                                        nr_actions = nr_actions,
                                        α_critic = eta,
                                        α_actor = alpha,
                                        γ = gamma,
                                        λ_critic = lambda_critic,
                                        λ_actor = lambda_actor,
                                        use_replacing_trace = use_replacing_trace,
                                        do_debug_print = do_debug_print)
                else
                    learner = ActorCriticModulated(nr_states = nr_states,
                                        nr_actions = nr_actions,
                                        α_critic = eta,
                                        α_actor = alpha,
                                        γ = gamma,
                                        λ_critic = lambda_critic,
                                        λ_actor = lambda_actor,
                                        policyparameters = copy(init_policy[current_puzzle_id]),
                                        V = copy(init_Vs[current_puzzle_id]),
                                        use_replacing_trace = use_replacing_trace,
                                        do_debug_print = do_debug_print)
                end
            end
            state = data[i, Int(c_state)]
            action = data[i, Int(c_action)]
            reward::Float64= data[i, Int(c_reward)]
            next_state = data[i, Int(c_nextState)]
            done = reward>0 ? true : false

            p_actions = getSoftmaxProbs(learner.policyparameters[state,:], temp, [0., 0.])

            unbiased_p_sa = p_actions[action]

            action_selection_probs_seeds[i_seed][i] = copy(unbiased_p_sa)
            LL_seeds[i_seed] += log(unbiased_p_sa)

            # Particle Filtering model learner
            updatet!(Testimate, state, action, next_state, done)
            γ_surprise_seeds[i_seed][i] = copy(Testimate.γ_surprise[1])
            Surprise_seeds[i_seed][i] = copy(Testimate.Sgm[1])
            Tsas_seeds[i_seed]= deepcopy(Testimate.Ps1a0s0)
            SPE_seeds[i_seed][i] = copy(Testimate.SPE[1])

            # Track V values (used as fmri regressor)
            V_state_seeds[i_seed][j] = learner.V[state]
            if done
                j+=1
                if j < (nr_samples + data_stats.nr_rewards)
                    V_state_seeds[i_seed][j] = learner.V[next_state] # Track V value of landing state in case it is the Goal
                end
            end
            j+=1

            PolicyParametersImbalance_record_seed[i_seed][i] = abs(learner.policyparameters[state, 1] - learner.policyparameters[state, 2])
            PolicyImbalance_record_seed[i_seed][i] = abs(p_actions[1] - p_actions[2])
            PolicyParametersDiffChosenVsUnchosen_record_seed[i_seed][i] = learner.policyparameters[state, action] - learner.policyparameters[state, mod(action,2)+1]
            PolicyDiffChosenVsUnchosen_record_seed[i_seed][i] = p_actions[action] - p_actions[mod(action,2)+1]
            PolicyParameters_record_seed[i_seed][i] = learner.policyparameters[state, action]

            # Actor-Critic
            update!(learner, state, action, next_state, reward, unbiased_p_sa, (1. - Testimate.γ_surprise[1]))
            TD_seeds[i_seed][i] = copy(learner.td_error[1])

            finalV_record_seed[i_seed][current_puzzle_id] = copy(learner.V)
            final_policyparameters_record_seed[i_seed][current_puzzle_id] = copy(learner.policyparameters)

        end
    end

    LL = get_meanLL_acrossseeds(LL_seeds)

    action_selection_probs_cat = hcat(action_selection_probs_seeds...)
    action_selection_probs = mean(action_selection_probs_cat, dims=2)[:]
    trial_by_trial_p_sa = (action_selection_probs)

    # For signals: return the ones corresponding to the Max LL across seeds
    LL_max, seed_idx = findmax(LL_seeds)

    TD_record = copy(TD_seeds[seed_idx])
    SPE_record = copy(SPE_seeds[seed_idx])
    trial_by_trial_p_sa_max = (action_selection_probs_seeds[seed_idx])

    trace_type = use_replacing_trace ? "ReplTrc" : "EligTrc"
    algo = "Actor-Critic (TD-Lambda, $trace_type) modulated Actor"
    param_v = [alpha, surpriseprobability, stochasticity, eta, gamma, lambda_actor, lambda_critic, temp, nparticles]
    param_str = "alpha(actor)=$(round(alpha,digits=4)),
                surpriseprobability=$(round(surpriseprobability,digits=4)),
                stochasticity=$(round(stochasticity,digits=4)),
                eta(critic)=$(round(eta,digits=4)),
                gamma=$(round(gamma,digits=4)),
                lambda_actor=$(round(lambda_actor,digits=4)),
                lambda_critic=$(round(lambda_critic,digits=4)),
                temp=$(round(temp,digits=4)),
                nparticles=$nparticles,
                init_policy=$(repr(init_policy)),
                init_Vs=$(repr(init_Vs)),
                use_replacing_trace=$use_replacing_trace"

    results_dict = Dict{String,Any}(
        "trialByTrial_Vs" => V_state_seeds[seed_idx],
        "final_V" => finalV_record_seed[seed_idx],
        "final_PolicyParameters" => final_policyparameters_record_seed[seed_idx],
        "final_Tsas" => Tsas_seeds[seed_idx],
        "Surprise" => Surprise_seeds[seed_idx],
        "gamma_surprise" => γ_surprise_seeds[seed_idx],
        "nparticles" => nparticles,
        "particlefilterseeds" => particlefilterseeds,
        "trial_by_trial_p_sa_max" => trial_by_trial_p_sa_max,
        "PolicyParametersImbalance" => PolicyParametersImbalance_record_seed[seed_idx],
        "PolicyImbalance" => PolicyImbalance_record_seed[seed_idx],
        "PolicyParametersDiffChosenVsUnchosen" => PolicyParametersDiffChosenVsUnchosen_record_seed[seed_idx],
        "PolicyDiffChosenVsUnchosen" => PolicyDiffChosenVsUnchosen_record_seed[seed_idx],
        "PolicyParameters" => PolicyParameters_record_seed[seed_idx],
        )

    signals = Step_by_step_signals(
    data_stats.all_subject_ids[1],
    TD_record,  # RPE
    SPE_record,  # no SPE
    trial_by_trial_p_sa,
    LL,
    data_stats,
    algo,
    param_v,
    param_str,
    99, # dummy value
    results_dict)
    return LL, signals
end

ext: jl


text:
@testset "Test the MathProgBase wrapper with linprog" begin
    include("mpb_linear.jl")
end

import MathProgBase
for T in [SCS.Direct, SCS.Indirect]
    @testset "MathProgBase $T" begin
        include(joinpath(dirname(dirname(pathof(MathProgBase))), "test", "conicinterface.jl"))
        coniclineartest(SCS.SCSSolver(linear_solver=T, eps=1e-6, verbose=0),
                        duals=true, tol=1e-5)
        conicSOCtest(SCS.SCSSolver(linear_solver=T, eps=1e-6, verbose=0),
                     duals=true, tol=1e-5)
        conicEXPtest(SCS.SCSSolver(linear_solver=T, eps=1e-6, verbose=0),
                     duals=true, tol=1e-5)
        conicSDPtest(SCS.SCSSolver(linear_solver=T, eps=1e-6, verbose=0),
                     duals=true, tol=1e-5)
    end
end


text:
using CellFreeModelGenerationKit

# path to test metabolism file -
path_to_test_metabolism_file = "./test/data/Test.vff"
path_to_output_dir = "./test/generated_code/model"

# need to build a model object -
build_result = build_julia_model_object(path_to_test_metabolism_file, path_to_output_dir)
if (isa(build_result.value, Exception) == true)
    @show build_result.value
end
julia_model_object = build_result.value

max_stars_repo_name: varnerlab/CellFreeModelGenerationKit

max_issues_repo_name: varnerlab/CellFreeModelGenerationKit

max_forks_repo_name: varnerlab/CellFreeModelGenerationKit


text:
export NDM

struct NDM{VT,MT} <: MatrixNeuralNetwork
    b_μ::VT
    h_μ::VT
    w_μ::MT
    u_μ::MT

    b_λ::VT
    h_λ::VT
    d_λ::VT
    w_λ::MT
    u_λ::MT
end
@treelike NDM

"""
    NDM([T=STD_REAL_PREC], N, αₕ, αₐ, [initW, initb, inita])

Constructs a Neural Density Matrix with numerical precision `T` (Defaults to
Float32), `N` input neurons, N⋅αₕ hidden neurons and N⋅αₐ ancillary neurons.
This network ensure that the density matrix is always positive definite.

The number of input neurons `N` must match the size of the lattice.

The initial parameters of the neurons are initialized with a rescaled normal
distribution of width 0.01 for the coupling matrix and 0.005 for the local
biases. The default initializers can be overriden by specifying

initW=(dims...)->rescaled_normal(T, 0.01, dims...),
initb=(dims...)->rescaled_normal(T, 0.005, dims...),
inita=(dims...)->rescaled_normal(T, 0.005, dims...))

Refs:
    https://arxiv.org/abs/1801.09684
    https://arxiv.org/abs/1902.10104
"""
NDM(args...) = NDM(STD_REAL_PREC, args...)
NDM(T::Type{<:Real}, in, αh, αa,
    initW=(dims...)->rescaled_normal(T, 0.01, dims...),
    initb=(dims...)->rescaled_normal(T, 0.005, dims...),
    inita=(dims...)->rescaled_normal(T, 0.005, dims...)) =
    NDM(inita(in),
        initb(convert(Int,αh*in)),
        initW(convert(Int,αh*in), in),
        initW(convert(Int,αa*in), in),
        inita(in),
        initb(convert(Int,αh*in)),
        initb(convert(Int,αa*in)),
        initW(convert(Int,αh*in), in),
        initW(convert(Int,αa*in), in))

input_type(net::NDM{VT,MT}) where {VT,MT} = eltype(VT)
weight_type(net::NDM) = input_type(net)
out_type(net::NDM{VT,MT}) where {VT,MT} = Complex{eltype(VT)}
input_shape(net::NDM) = (length(net.b_μ), length(net.b_μ))
random_input_state(net::NDM{VT,MT}) where {VT,MT} =
    (eltype(VT).([rand(0:1) for i=1:length(net.b_μ)]), eltype(VT).([rand(0:1) for i=1:length(net.b_μ)]))
is_analytic(net::NDM) = true

Base.show(io::IO, m::NDM) = print(io,
    "NDM($(eltype(m.b_μ)), n=$(length(m.b_μ)), αₕ=$(length(m.h_μ)/length(m.b_μ)), αₐ=$(length(m.d_λ)/length(m.b_μ)))")

Base.show(io::IO, ::MIME"text/plain", m::NDM) = print(io,
    "NDM($(eltype(m.b_μ)), n=$(length(m.b_μ)), α=$(length(m.h_μ)/length(m.b_μ)), αₐ=$(length(m.d_λ)/length(m.b_μ)))")

@inline (net::NDM)(σ::State) = net(config(σ)...)
@inline (net::NDM)(σ::Tuple) = net(σ...)
function (W::NDM)(σr, σc)
    T=eltype(W.u_λ)
    ∑logℒ_λ_σ = sum_autobatch(logℒ.(W.h_λ .+ W.w_λ*σr))
    ∑logℒ_μ_σ = sum_autobatch(logℒ.(W.h_μ .+ W.w_μ*σr))

    ∑logℒ_λ_σp = sum_autobatch(logℒ.(W.h_λ .+ W.w_λ*σc))
    ∑logℒ_μ_σp = sum_autobatch(logℒ.(W.h_μ .+ W.w_μ*σc))

    ∑σ = σr .+ σc
    Δσ = σr .- σc

    _Π = T(0.5)  .* W.u_λ*∑σ .+
           T(0.5)im .* W.u_μ*Δσ .+ W.d_λ

    Γ_λ = T(0.5)   * (∑logℒ_λ_σ + ∑logℒ_λ_σp + transpose(W.b_λ)*∑σ )
    Γ_μ = T(0.5)im * (∑logℒ_μ_σ - ∑logℒ_μ_σp + transpose(W.b_μ)*Δσ )
    Π  = sum_autobatch(logℒ.(_Π))

    lnψ = Γ_λ + Γ_μ + Π
    return lnψ
end

# Cached version
mutable struct NDMCache{T,VT,VCT} <: NNCache{NDM}
    θλ_σ::VT
    θμ_σ::VT
    θλ_σp::VT
    θμ_σp::VT

    θλ_σ_tmp::VT
    θμ_σ_tmp::VT
    θλ_σp_tmp::VT
    θμ_σp_tmp::VT

    σr::VT
    σc::VT
    ∑σ::VT
    Δσ::VT

    ∑logℒ_λ_σ::T
    ∑logℒ_μ_σ::T
    ∂logℒ_λ_σ::VT
    ∂logℒ_μ_σ::VT
    ∂logℒ_λ_σp::VT
    ∂logℒ_μ_σp::VT
    _Π::VCT
    _Π2::VCT
    _Π_tmp::VT
    ∂logℒ_Π::VCT

    σ_row_cache::VT
    i_σ_row_cache::Int

    valid::Bool # = false
end

cache(net::NDM) =
    NDMCache(similar(net.h_μ),
              similar(net.h_μ),
              similar(net.h_μ),
              similar(net.h_μ),

              similar(net.h_μ),
              similar(net.h_μ),
              similar(net.h_μ),
              similar(net.h_μ),

              similar(net.b_μ),
              similar(net.b_μ),
              similar(net.b_μ),
              similar(net.b_μ),

              zero(eltype(net.b_μ)), zero(eltype(net.b_μ)),
              similar(net.h_μ), similar(net.h_μ),
              similar(net.h_μ), similar(net.h_μ),
              similar(net.d_λ, complex(eltype(net.d_λ))),
              similar(net.d_λ, complex(eltype(net.d_λ))),
              similar(net.d_λ, eltype(net.d_λ)),
              similar(net.d_λ, complex(eltype(net.d_λ))),
                  #VT(T, length(net.h_μ)), zeros(T, length(net.h_μ)),
                  #VT(T, length(net.h_μ)), zeros(T, length(net.h_μ)),
                  #zeros(Complex{T}, length(net.d_λ)), zeros(T, length(net.d_λ)), zeros(Complex{T}, length(net.d_λ)),

              similar(net.b_μ),
              -1,

              false)

(net::NDM)(c::NDMCache, σ::State) = net(c, config(σ))
(net::NDM)(c::NDMCache, (σr, σc)::Tuple{AbstractArray,AbstractArray}) = net(c, σr, σc)
function (W::NDM)(c::NDMCache, σr, σc)
    ∑σ      = c.∑σ
    Δσ      = c.Δσ
    θλ_σ    = c.θλ_σ
    θμ_σ    = c.θμ_σ
    θλ_σp   = c.θλ_σp
    θμ_σp   = c.θμ_σp
    _Π      = c._Π
    _Π_tmp  = c._Π_tmp
    T       = eltype(c.θλ_σ)

    if !c.valid || c.σ_row_cache ≠ σr
        c.σ_row_cache .= σr
        c.valid = true

        # θλ_σ .= W.h_λ + W.w_λ*σr
        LinearAlgebra.BLAS.blascopy!(length(W.h_λ), W.h_λ, 1, θλ_σ, 1)
        LinearAlgebra.BLAS.gemv!('N', one(T), W.w_λ, σr, one(T), θλ_σ)

        # θμ_σ .= W.h_μ + W.w_μ*σr
        LinearAlgebra.BLAS.blascopy!(length(W.h_μ), W.h_μ, 1, θμ_σ, 1)
        LinearAlgebra.BLAS.gemv!('N', one(T), W.w_μ, σr, one(T), θμ_σ)

        # cache.∑logℒ_λ_σ = sum(logℒ.(θλ_σ, (NT(),)))
        c.θλ_σ_tmp .= logℒ.(θλ_σ)
        c.∑logℒ_λ_σ = sum(c.θλ_σ_tmp)

        # cache.∑logℒ_μ_σ = sum(logℒ.(θμ_σ, (NT(),)))
        c.θμ_σ_tmp .= logℒ.(θμ_σ)
        c.∑logℒ_μ_σ = sum(c.θμ_σ_tmp)

        c.∂logℒ_λ_σ .= ∂logℒ.(θλ_σ)
        c.∂logℒ_μ_σ .= ∂logℒ.(θμ_σ)
    end

    ∑logℒ_λ_σ = c.∑logℒ_λ_σ
    ∑logℒ_μ_σ = c.∑logℒ_μ_σ
    ∂logℒ_λ_σ = c.∂logℒ_λ_σ
    ∂logℒ_μ_σ = c.∂logℒ_μ_σ

    ∑σ .= σr .+ σc
    Δσ .= σr .- σc

    #θλ_σp .= W.h_λ + W.w_λ*σc
    #θμ_σp .= W.h_μ + W.w_μ*σc
    LinearAlgebra.BLAS.blascopy!(length(W.h_λ), W.h_λ, 1, θλ_σp, 1)
    LinearAlgebra.BLAS.gemv!('N', T(1.0), W.w_λ, σc, T(1.0), θλ_σp)
    LinearAlgebra.BLAS.blascopy!(length(W.h_μ), W.h_μ, 1, θμ_σp, 1)
    LinearAlgebra.BLAS.gemv!('N', T(1.0), W.w_μ, σc, T(1.0), θμ_σp)

    #∑logℒ_λ_σp = sum(logℒ.(θλ_σp, (NT(),)))
    c.θλ_σp_tmp .= logℒ.(θλ_σp)
    ∑logℒ_λ_σp = sum(c.θλ_σp_tmp)

    #∑logℒ_μ_σp = sum(logℒ.(θμ_σp, (NT(),)))
    c.θμ_σp_tmp .= logℒ.(θμ_σp)
    ∑logℒ_μ_σp = sum(c.θμ_σp_tmp)

    #_Π         = (T(0.5)  * W.u_λ*∑σ
    #               + T(0.5)im * W.u_μ*Δσ .+ W.d_λ)
    LinearAlgebra.BLAS.blascopy!(length(W.d_λ), W.d_λ, 1, _Π_tmp, 1)
    LinearAlgebra.BLAS.gemv!('N', T(0.5), W.u_λ, ∑σ, one(T), _Π_tmp)
    _Π .= _Π_tmp
    LinearAlgebra.BLAS.gemv!('N', T(0.5), W.u_μ, Δσ, T(0.0), _Π_tmp)
    _Π .+= T(1.0)im.* _Π_tmp

    Γ_λ = T(0.5) * (∑logℒ_λ_σ + ∑logℒ_λ_σp + ∑σ ⋅ W.b_λ)
    Γ_μ = T(0.5) * (∑logℒ_μ_σ - ∑logℒ_μ_σp + Δσ ⋅ W.b_μ)

    # Π = sum(logℒ.(_Π, (NT(),)))
    _Π .= logℒ.(_Π)
    Π   = sum(_Π)

    logψ = Γ_λ + T(1.0)im * Γ_μ + Π
    return logψ
end

function logψ_and_∇logψ!(∇logψ, W::NDM, c::NDMCache, σr,σc)
    ∑σ      = c.∑σ
    Δσ      = c.Δσ
    θλ_σ    = c.θλ_σ
    θμ_σ    = c.θμ_σ
    θλ_σp   = c.θλ_σp
    θμ_σp   = c.θμ_σp
    _Π      = c._Π
    _Π_tmp  = c._Π_tmp
    T       = eltype(θλ_σ)

    if !c.valid || c.σ_row_cache ≠ σr
        c.σ_row_cache .= σr
        c.valid = true

        # θλ_σ .= W.h_λ + W.w_λ*σr
        LinearAlgebra.BLAS.blascopy!(length(W.h_λ), W.h_λ, 1, θλ_σ, 1)
        LinearAlgebra.BLAS.gemv!('N', one(T), W.w_λ, σr, one(T), θλ_σ)

        # θμ_σ .= W.h_μ + W.w_μ*σr
        LinearAlgebra.BLAS.blascopy!(length(W.h_μ), W.h_μ, 1, θμ_σ, 1)
        LinearAlgebra.BLAS.gemv!('N', one(T), W.w_μ, σr, one(T), θμ_σ)

        # cache.∑logℒ_λ_σ = sum(logℒ.(θλ_σ, (NT(),)))
        c.θλ_σ_tmp .= logℒ.(θλ_σ)
        c.∑logℒ_λ_σ = sum(c.θλ_σ_tmp)

        # cache.∑logℒ_μ_σ = sum(logℒ.(θμ_σ, (NT(),)))
        c.θμ_σ_tmp .= logℒ.(θμ_σ)
        c.∑logℒ_μ_σ = sum(c.θμ_σ_tmp)

        c.∂logℒ_λ_σ .= ∂logℒ.(θλ_σ)
        c.∂logℒ_μ_σ .= ∂logℒ.(θμ_σ)
    end

    ∑logℒ_λ_σ = c.∑logℒ_λ_σ
    ∑logℒ_μ_σ = c.∑logℒ_μ_σ
    ∂logℒ_λ_σ = c.∂logℒ_λ_σ
    ∂logℒ_μ_σ = c.∂logℒ_μ_σ

    # --- Common terms with computation of ψ --- #
    ∑σ .= σr .+ σc
    Δσ .= σr .- σc

    #θλ_σp .= W.h_λ + W.w_λ*σc
    #θμ_σp .= W.h_μ + W.w_μ*σc
    LinearAlgebra.BLAS.blascopy!(length(W.h_λ), W.h_λ, 1, θλ_σp, 1)
    LinearAlgebra.BLAS.gemv!('N', T(1.0), W.w_λ, σc, T(1.0), θλ_σp)
    LinearAlgebra.BLAS.blascopy!(length(W.h_μ), W.h_μ, 1, θμ_σp, 1)
    LinearAlgebra.BLAS.gemv!('N', T(1.0), W.w_μ, σc, T(1.0), θμ_σp);

    #∑logℒ_λ_σp = sum(logℒ.(θλ_σp, (NT(),)))
    c.θλ_σp_tmp .= logℒ.(θλ_σp)
    ∑logℒ_λ_σp = sum(c.θλ_σp_tmp)

    #∑logℒ_μ_σp = sum(logℒ.(θμ_σp, (NT(),)))
    c.θμ_σp_tmp .= logℒ.(θμ_σp)
    ∑logℒ_μ_σp = sum(c.θμ_σp_tmp)

    #_Π         = (T(0.5)  * W.u_λ*∑σ
    #               + T(0.5)im * W.u_μ*Δσ .+ W.d_λ)
    LinearAlgebra.BLAS.blascopy!(length(W.d_λ), W.d_λ, 1, _Π_tmp, 1)
    LinearAlgebra.BLAS.gemv!('N', T(0.5), W.u_λ, ∑σ, T(1.0), _Π_tmp)
    _Π .= _Π_tmp
    LinearAlgebra.BLAS.gemv!('N', T(0.5), W.u_μ, Δσ, T(0.0), _Π_tmp)
    _Π .+= T(1.0)im.* _Π_tmp
    #@info "_Π diff " maximum(abs.(_Π - (T(0.5)  * transpose(transpose(∑σ)*W.u_λ) + T(0.5)im* transpose(transpose(Δσ)*W.u_μ) .+ W.d_λ)))

    # --- End common terms with computation of ψ --- #

    # Compute additional terms for derivatives
    ∂logℒ_λ_σp = c.∂logℒ_λ_σp; ∂logℒ_λ_σp .= ∂logℒ.(θλ_σp)
    ∂logℒ_μ_σp = c.∂logℒ_μ_σp; ∂logℒ_μ_σp .= ∂logℒ.(θμ_σp)
    ∂logℒ_Π    = c.∂logℒ_Π;    ∂logℒ_Π    .= ∂logℒ.(_Π)

    # Store the derivatives
    ∇logψ.b_λ .= T(0.5)   .* ∑σ
    ∇logψ.b_μ .= T(0.5)im .* Δσ

    ∇logψ.h_λ .= T(0.5)   .* (∂logℒ_λ_σ .+ ∂logℒ_λ_σp)
    ∇logψ.h_μ .= T(0.5)im .* (∂logℒ_μ_σ .- ∂logℒ_μ_σp)

    ∇logψ.w_λ .= T(0.5)   .* (∂logℒ_λ_σ.*transpose(σr) .+ ∂logℒ_λ_σp.*transpose(σc))
    ∇logψ.w_μ .= T(0.5)im .* (∂logℒ_μ_σ.*transpose(σr) .- ∂logℒ_μ_σp.*transpose(σc))

    ∇logψ.d_λ .= ∂logℒ_Π
    ∇logψ.u_λ .= T(0.5) .* ∂logℒ_Π .* transpose(∑σ)
    ∇logψ.u_μ .= T(0.5)im .*  ∂logℒ_Π .* transpose(Δσ)

    Γ_λ = T(0.5) * (∑logℒ_λ_σ + ∑logℒ_λ_σp + ∑σ ⋅ W.b_λ)
    Γ_μ = T(0.5) * (∑logℒ_μ_σ - ∑logℒ_μ_σp + Δσ ⋅ W.b_μ)
    _Π .= logℒ.(_Π)
    Π   = sum(_Π)
    logψ = Γ_λ + T(1.0)im * Γ_μ + Π

    return logψ
end

function vectorize_gradient(net::NDM{T}, gradient::NamedTuple) where T
    fnames_μ = [:b_μ, :h_μ, :w_μ, :u_μ]
    fnames_λ = [:b_λ, :h_λ, :d_λ, :w_λ, :u_λ]
    fnames_sets = [fnames_μ, fnames_λ]
    kvpairs    =Dict{Symbol, Any}()
    params_vecs = []

    for fnames=fnames_sets
        lens = Int[1]
        for f=fnames
            push!(lens, length(gradient[f]))
        end
        n_params   = sum(lens)-1
        indices    = cumsum(lens)

        params_vec = Vector{Complex{T}}(undef, n_params)
        push!(params_vecs, params_vec)
        i=1;
        for f=fnames
            @views datavec = params_vec[indices[i]:indices[i+1]-1]
            reshpd_params  = reshape(datavec, size(gradient[f]))
            reshpd_params .= gradient[f]
            push!(kvpairs, f=>reshpd_params)
            i = i+1
        end
    end
    newgrad = (;kvpairs...)
    return (newgrad, params_vecs)
end


text:
### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5b918be6-0003-11ec-1e1f-534faed7dfe5
begin
    using Pkg
    Pkg.activate("..")
    using Revise,
        Javis, Distributions, StatsBase, JUtils, Random, Colors, FileIO, PlutoUI, Plots
    import Luxor
end

# ╔═╡ 46fd9cf5-eae0-4dec-91df-ebd892cbf4d6
converter = Dict(["bernoulli" => Bernoulli, "pareto" => Pareto]);

# ╔═╡ b6e3b06f-a521-40fb-9d5a-9268302cd17b
md"""
Select the distribution $(@bind dist Select(collect(keys(converter)), default="bernoulli"))
"""

# ╔═╡ d3ed0756-5fdf-4eba-a8ed-7d944acc079d
md"""
Select the distribution Parameters $(@bind param NumberField(0.0:0.1:5.0, default=0.3))
"""

# ╔═╡ e8fef37b-e0a7-42f3-89a7-414d61092c6e
md"""
Check to start rendering gif$(@bind start CheckBox())
"""

# ╔═╡ e725b87f-fb23-472d-84d7-8e6a8dd01aae
param_dist = converter[dist](param)

# ╔═╡ 47bcc4fe-a44f-4715-8ed3-bbe00c02f30f
begin
    start && JUtils.CLT(
        param_dist,
        boundingbox = BoundingBox(O + (-250, -300), O + (250, 200)),
        pathname = "../output/central_limit_theorem.gif",
    )
end

# ╔═╡ Cell order:
# ╟─5b918be6-0003-11ec-1e1f-534faed7dfe5
# ╟─46fd9cf5-eae0-4dec-91df-ebd892cbf4d6
# ╟─b6e3b06f-a521-40fb-9d5a-9268302cd17b
# ╟─d3ed0756-5fdf-4eba-a8ed-7d944acc079d
# ╟─e8fef37b-e0a7-42f3-89a7-414d61092c6e
# ╟─e725b87f-fb23-472d-84d7-8e6a8dd01aae
# ╟─47bcc4fe-a44f-4715-8ed3-bbe00c02f30f


text:
module TestDataArray
    using Base.Test
    using DataArrays

    v = [1, 2, 3, 4]
    dv = DataArray(v, falses(size(v)))

    m = [1 2; 3 4]
    dm = DataArray(m, falses(size(m)))

    t = Array(Int, 2, 2, 2)
    t[1:2, 1:2, 1:2] = 1
    dt = DataArray(t, falses(size(t)))

    dv = DataArray(v)
    dv = DataArray(v, [false, false, false, false])

    dv = DataArray(Int, 2)
    dm = DataArray(Int, 2, 2)
    dt = DataArray(Int, 2, 2, 2)

    similar(dv)
    similar(dm)
    similar(dt)

    similar(dv, 2)
    similar(dm, 2, 2)
    similar(dt, 2, 2, 2)

    @test isequal(DataArray([NA, NA], [true, true]), DataArray([NA, NA], [false, false]))
    @test isequal(DataArray(Any[1, NA], [false, true]), DataArray(Any[1, NA], [false, false]))
    @test isequal(DataArray(Any[1, 2], [false, true]), DataArray(Any[1, NA], [false, false]))

    x = DataArray([9, 9, 8])
    y = DataArray([1, 9, 3, 2, 2])
    @assert append!(x, y) == [9, 9, 8, 1, 9, 3, 2, 2]

    x = DataArray([9, 9, 8])
    y = [1, 9, 3, 2, 2]
    @assert append!(x, y) == [9, 9, 8, 1, 9, 3, 2, 2]

    x = @data [1, 2, NA]
    y = @data [3, NA, 5]
    @test isequal(copy(x), x)
    @test isequal(copy!(y, x), x)

    x = @data [1, NA, -2, 1, NA, 4]
    @assert isequal(unique(x), @data [1, NA, -2, 4])
    @assert isequal(unique(reverse(x)), @data [4, NA, 1, -2])
    @assert isequal(unique(dropna(x)), @data [1, -2, 4])
    @assert isequal(unique(reverse(dropna(x))), @data [4, 1, -2])
    @assert isequal(levels(x), @data [1, -2, 4])
    @assert isequal(levels(reverse(x)), @data [4, 1, -2])

    # check case where only NA occurs in final position
    @assert isequal(unique(@data [1, 2, 1, NA]), @data [1, 2, NA])

    # Test copy!
    function nonbits(dv)
        ret = similar(dv, Integer)
        for i = 1:length(dv)
            if !isna(dv, i)
                ret[i] = dv[i]
            end
        end
        ret
    end
    set1 = Any[@data([1, NA, 3]),
               @data([NA, 5]), @data([1, 2, 3, 4, 5]), data(Int[]),
               @data([NA, 5, 3]), @data([1, 5, 3])]
    set2 = map(nonbits, set1)
    set3 = map(pdata, set1)

    for (dest, src, bigsrc, emptysrc, res1, res2) in Any[set1, set2, set3]

        @test isequal(copy!(copy(dest), src), res1)
        @test isequal(copy!(copy(dest), 1, src), res1)

        @test isequal(copy!(copy(dest), 2, src, 2), res2)
        @test isequal(copy!(copy(dest), 2, src, 2, 1), res2)

        @test isequal(copy!(copy(dest), 99, src, 99, 0), dest)

        @test isequal(copy!(copy(dest), 1, emptysrc), dest)
        @test_throws BoundsError copy!(dest, 1, emptysrc, 1)

        for idx in [0, 4]
            @test_throws BoundsError copy!(dest, idx, src)
            @test_throws BoundsError copy!(dest, idx, src, 1)
            @test_throws BoundsError copy!(dest, idx, src, 1, 1)
            @test_throws BoundsError copy!(dest, 1, src, idx)
            @test_throws BoundsError copy!(dest, 1, src, idx, 1)
        end

        if VERSION >= v"0.5.0-dev+4711"
            @test_throws ArgumentError copy!(dest, 1, src, 1, -1)
        else
            @test_throws BoundsError copy!(dest, 1, src, 1, -1)
        end

        @test_throws BoundsError copy!(dest, bigsrc)

        @test_throws BoundsError copy!(dest, 3, src)
        @test_throws BoundsError copy!(dest, 3, src, 1)
        @test_throws BoundsError copy!(dest, 3, src, 1, 2)
        @test_throws BoundsError copy!(dest, 1, src, 2, 2)
    end

end


text:
using BenchmarkTools, StaticArrays
include("shared.jl")

const oo = @SVector ones(2)
const zz = @SVector zeros(2)
const TESTING_PARAMS = (@SVector [1.25, 0.25, 0.5],
                        @SVector [1.25, 5.25, 0.5],
                        @SVector [1.25, 0.25, 0.75],
                        @SVector [1.25, 5.25, 0.75],
                        @SVector [1.25, 0.25, 1.0],
                        @SVector [1.25, 5.25, 1.0],
                        @SVector [1.25, 0.25, 1.75],
                        @SVector [1.25, 5.25, 1.75],
                        @SVector [1.25, 0.25, 3.0],
                        @SVector [1.25, 5.25, 3.0],
                        @SVector [1.25, 0.25, 4.75],
                        @SVector [1.25, 5.25, 4.75])

for pp in TESTING_PARAMS

  println("\n###")
  println("Parameters $pp")
  println("###\n")

  println("Fast finite diff, h=$h:")
  @btime matern_fdfast_d1($oo, $zz, $pp)
  @btime matern_fdfast_d2($oo, $zz, $pp)
  @btime matern_fdfast_d3($oo, $zz, $pp)

  println("Adaptive finite diff, order 2:")
  @btime matern_fd2_d1($oo, $zz, $pp)
  @btime matern_fd2_d2($oo, $zz, $pp)
  @btime matern_fd2_d3($oo, $zz, $pp)

  println("Adaptive finite diff, order 5:")
  @btime matern_fd5_d1($oo, $zz, $pp)
  @btime matern_fd5_d2($oo, $zz, $pp)
  @btime matern_fd5_d3($oo, $zz, $pp)

  println("Complex step, h=$ch:")
  try
    @btime matern_cstep_d1($oo, $zz, $pp)
    @btime matern_cstep_d2($oo, $zz, $pp)
    @btime matern_cstep_d3($oo, $zz, $pp)
  catch
    println("Failure/error. Probably at integer values.")
  end

  println("Autodiff:")
  @btime matern_ad_d1($oo, $zz, $pp)
  @btime matern_ad_d2($oo, $zz, $pp)
  @btime matern_ad_d3($oo, $zz, $pp)

end


text:
#AbstractTensor
Base.:(*)(α::Number, A::AbstractTensor) = apply(A, x -> @. x *= α)
Base.:(*)(A::AbstractTensor, α::Number) = α   * A
Base.:(/)(A::AbstractTensor, α::Number) = 1/α * A
Base.conj!(A::AbstractTensor) = apply!(A, conj!)
Base.conj(A::AbstractTensor)  = apply(A, conj!)
function LA.pinv(A::AbstractTensor, inds = (1,2))
    u, s, vd = tensorsvd(A, inds)
    apply!(s, x -> x .= LA.pinv(x))
    @TO.tensor A2[1,2] := vd'[-1,1] * s[-1,-2] * u'[2,-2]
    return A2
end

#DTensor
apply!(A::DTensor, fun!) = (fun!(A.array); return A)
apply(A::DTensor, fun!) = apply!(deepcopy(A), fun!)
Base.adjoint(a::DTensor) = conj(a)

#= Operations on Tensors =#
function apply!(A::DASTensor, fun!)
    foreach(fun!, values(A))
    return A
end
apply(A::DASTensor, fun!) = apply!(deepcopy(A), fun!)
function Base.adjoint(A::DASTensor)
    B = apply(A, conj!)
    B.io = inv(B.io)
    setcharge!(B, -charge(A))
    return B
end


text:
module _MSEn
export MSEn, EMD
using Statistics: std, mean, median, var
#using DataInterpolations: CubicSpline
using Dierckx: Spline1D
using Plots
#=function __init__()
 @warn("\n\n Methodx option IMF (Intrinisic Mode Function) is not stable.
 Random or highly aperiodic signals may not decompose fully.
 Access to the IMFs decomposed by the empirical mode decomposition (EMD) function
 can be found by calling EntropyHub.EMD(`Sig`,`MaxIMFs`).
 A stable EMD function will be included in future releases.\n\n")
end=#

    """
        MSx, CI = MSEn(Sig, Mobj)

    Returns a vector of multiscale entropy values `MSx` and the complexity
    index `CI` of the data sequence `Sig` using the parameters specified
    by the multiscale object `Mobj` over 3 temporal scales with coarse-
    graining (default).

        MSx, CI = MSEn(Sig::AbstractArray{T,1} where T<:Real, Mobj::NamedTuple; Scales::Int=3,
                             Methodx::String="coarse", RadNew::Int=0, Plotx::Bool=false)

    Returns a vector of multiscale entropy values `MSx` and the complexity
    index `CI` of the data sequence `Sig` using the parameters specified by
    the multiscale object `Mobj` and the following 'keyword' arguments:

    # Arguments:
    `Scales`   - Number of temporal scales, an integer > 1   (default: 3) \n
    `Method`   - Graining method, one of the following:
                 {`coarse`,`modified`,`imf`,`timeshift`} [default = `coarse`]
                 For further info on these graining procedures, see the EntropyHub guide.  \n
    `RadNew`   - Radius rescaling method, an integer in the range [1 4].
                 When the entropy specified by `Mobj` is `SampEn` or `ApEn`,
                 RadNew allows the radius threshold to be updated at each
                 time scale (Xt). If a radius value is specified by `Mobj` (`r`),
                 this becomes the rescaling coefficient, otherwise it is set
                 to 0.2 (default). The value of RadNew specifies one of the
                 following methods:\n
                 [1]    Standard Deviation          - r*std(Xt)\n
                 [2]    Variance                    - r*var(Xt) \n
                 [3]    Mean Absolute Deviation     - r*mean_ad(Xt) \n
                 [4]    Median Absolute Deviation   - r*med_ad(Xt)\n
    `Plotx`    - When Plotx == true, returns a plot of the entropy value at each
                 time scale (i.e. the multiscale entropy curve) [default: false]\n

   `For further info on these graining procedures see the EntropyHub guide.`

    # See also `MSobject`, `rMSEn`, `cMSEn`, `hMSEn`, `SampEn`, `ApEn`, `XMSEn`

    # References:
        [1] Madalena Costa, Ary Goldberger, and C-K. Peng,
                "Multiscale entropy analysis of complex physiologic time series."
                Physical review letters
                89.6 (2002): 068102.

        [2] Vadim V. Nikulin, and Tom Brismar,
                "Comment on “Multiscale entropy analysis of complex physiologic
                time series”."
                Physical review letters
                92.8 (2004): 089803.

        [3] Madalena Costa, Ary L. Goldberger, and C-K. Peng.
                "Costa, Goldberger, and Peng reply."
                Physical Review Letters
                92.8 (2004): 089804.

        [4] Madalena Costa, Ary L. Goldberger and C-K. Peng,
                "Multiscale entropy analysis of biological signals."
                Physical review E
                71.2 (2005): 021906.

        [5] Ranjit A. Thuraisingham and Georg A. Gottwald,
                "On multiscale entropy analysis for physiological data."
                Physica A: Statistical Mechanics and its Applications
                366 (2006): 323-332.

        [6] Meng Hu and Hualou Liang,
                "Intrinsic mode entropy based on multivariate empirical mode
                decomposition and its application to neural data analysis."
                Cognitive neurodynamics
                5.3 (2011): 277-284.

        [7] Anne Humeau-Heurtier
                "The multiscale entropy algorithm and its variants: A review."
                Entropy
                17.5 (2015): 3110-3123.

        [8] Jianbo Gao, et al.,
                "Multiscale entropy analysis of biological signals: a
                fundamental bi-scaling law."
                Frontiers in computational neuroscience
                9 (2015): 64.

        [9]  Paolo Castiglioni, et al.,
                "Multiscale Sample Entropy of cardiovascular signals: Does the
                choice between fixed-or varying-tolerance among scales
                influence its evaluation and interpretation?."
                Entropy
                19.11 (2017): 590.

        [10] Tuan D Pham,
                "Time-shift multiscale entropy analysis of physiological signals."
                Entropy
                19.6 (2017): 257.

        [11] Hamed Azami and Javier Escudero,
                "Coarse-graining approaches in univariate multiscale sample
                and dispersion entropy."
                Entropy 20.2 (2018): 138.

    """
    function MSEn(Sig::AbstractArray{T,1} where T<:Real, Mobj::NamedTuple; Scales::Int=3,
        Methodx::String="coarse", RadNew::Int=0, Plotx::Bool=false)

    (size(Sig,1)>10) ? nothing : error("Sig:   must be a numeric vector" )
    (length(Mobj) >= 1) ? nothing :  error("Mobj:    must be a multiscale entropy object created
                with the function EntropyHub.MSobject")
    (Scales>1) ? nothing : error("Scales:     must be an integer > 1")
    (lowercase(Methodx) in ["coarse","modified","imf","timeshift"]) ? nothing :
        error("Method:  must be one of the following string names -
                'coarse','modified','imf','timeshift'")
    (RadNew==0 || (RadNew in 1:4 && String(Symbol(Mobj.Func)) in ("SampEn","ApEn"))) ? nothing :
        error("RadNew:     must be 0, or an integer in range [1 4] with
                entropy function 'SampEn' or 'ApEn'")

    if lowercase(Methodx)=="imf"
        Sig,_  = EMD(Sig,Scales-1)
        sum(all(Sig.==0,dims=2))==0 ? nothing : Sig = Sig[all(Sig.!=0,dims=2),:]
        Scales >= size(Sig,1) ? nothing :
        @warn("Max number of IMF's decomposed from EMD is less than number of Scales.
            MSEn evaluated over $(size(Sig,1)) scales instead of $Scales.")
        Scales = size(Sig,1)
    end
    MSx = zeros(Scales)
    Args = NamedTuple{keys(Mobj)[2:end]}(Mobj)
    Func2 = getfield(_MSEn,Symbol(lowercase(Methodx)))
    if RadNew > 0
        if RadNew == 1
            Rnew = x -> std(x, corrected=false)
        elseif RadNew == 2
            Rnew = x -> var(x, corrected=false)
        elseif RadNew == 3
            Rnew = x -> mean(abs.(x .- mean(x)))
        elseif RadNew == 4
            Rnew = x -> median(abs.(x .- median(x)))
        end

        if haskey(Mobj,:r)
            Cx = Mobj.r
        else
            Cy = ("Standard Deviation","Variance","Mean Abs Deviation",
                    "Median Abs Deviation")
            @warn("No radius value provided in Mobj.
                Default set to 0.2*$(Cy[RadNew]) of each new time-series.")
            Cx = .2
        end
    end

    for T = 1:Scales
        print(". ")
        Temp = Func2(Sig,T)

        if lowercase(Methodx) == "timeshift"
            Tempx = zeros(T)
            for k = 1:T
                RadNew > 0 ? Args = (Args..., r=Cx*Rnew(Temp[k,:])) : nothing
                Tempy = Mobj.Func(Temp[k,:]; Args...)
                typeof(Tempy)<:Tuple ? Tempx[k] = Tempy[1][end] : Tempx[k] = Tempy[end]
                # Tempx[k] = Tempy[end]
            end
            Temp2 = mean(Tempx)
        else
            RadNew > 0 ? Args = (Args..., r=Cx*Rnew(Temp[:])) : nothing
            Tempx = Mobj.Func(Temp[:]; Args...)
            typeof(Tempx)<:Tuple ? Temp2 = Tempx[1][end] : Temp2 = Tempx[end]
            #Temp2 = Tempx[1][end]
        end

        MSx[T] = Temp2
    end
    CI = sum(MSx)
    print("\n")
    if any(isnan.(MSx))
        println("Some entropy values may be undefined.")
    end

    if Plotx

        p1 = plot(1:Scales, MSx, c=RGB(8/255, 63/255, 77/255), lw=3)
        scatter!(1:Scales, MSx, markersize=6, c=RGB(1, 0, 1),
        xlabel = "Scale Factor", ylabel = "Entropy Value",
        guidefont = font(12, "arial", RGB(7/255, 54/255, 66/255)),
        tickfontsize = 10, tickfontfamily="arial", legend=false,
        title = "Multiscale $(Mobj.Func) ($(titlecase(Methodx))-graining method)",
        plot_titlefontsize=16, plot_titlefontcolor=RGB(7/255, 54/255, 66/255)) #ylim=(0,maximum(MSx)+.2),
        display(p1)
    end

    return MSx, CI
    end

    function coarse(Z,sx)
        Ns = Int(floor(size(Z,1)/sx))
        Y = mean(reshape(Z[1:sx*Ns],sx,Ns),dims=1)
        return Y
    end
    function modified(Z,sx)
        Ns = size(Z,1) - sx + 1
        Y = zeros(Ns)
        for k = 1:Ns
            Y[k] = mean(Z[k:k+sx-1])
        end
        return Y
    end
    function imf(Z,sx)
        Y = sum(Z[1:sx,:],dims=1)
        return Y
    end
    function timeshift(Z,sx)
        Y =  reshape(Z[1:Int(sx*floor(length(Z)/sx))],
                (sx,Int(floor(length(Z)/sx))))
        return Y
    end

    function PkFind(X)
        Nx = length(X)
        Indx = zeros(Int,Nx);
        for n = 2:Nx-1
            if X[n-1]< X[n] > X[n+1]
                Indx[n] = n

            elseif X[n-1] < X[n] == X[n+1]
                k = 1
                Indx[n] = n
                while (n+k)<Nx && X[n] == X[n+k]
                    Indx[n+k] = n+k
                    k+=1
                end
                n+=k
            end
        end
        Indx = Indx[Indx.!==0]
    return Indx
    end

    function EMD(X, Scales::Int)
        Xt = copy(X); N = size(Xt,1); n=1; IMFs = zeros(Scales+1,N)
        MaxER = 20;  MinTN = 2;  #Xt .-= mean(Xt)
        r1 = Xt

        while n <= Scales
            r0 = Xt
            x = 0;
            Upx = PkFind(r0);  Lwx = PkFind(-r0)
            UpEnv = Spline1D(Upx,r0[Upx],k=3,bc="nearest")
            LwEnv = Spline1D(Lwx,r0[Lwx],k=3,bc="nearest")
            r1 = r0.- (UpEnv.(1:N) .+ LwEnv.(1:N))./2
            RT = (sum(r0.*r0) - sum(r1.*r1))/sum(r0.*r0)
            length(vcat(Upx,Lwx)) <= MinTN ? (LOG = "Decomposition hit minimal extrema criteria."; break) : nothing

            while x < 100 && RT > 0.2
                r0 = 1*r1
                Upx = PkFind(r0);  Lwx = PkFind(-r0)
                UpEnv = Spline1D(Upx,r0[Upx],k=3,bc="nearest")
                LwEnv = Spline1D(Lwx,r0[Lwx],k=3,bc="nearest")
                r1 = r0.- (UpEnv.(1:N) .+ LwEnv.(1:N))./2
                RT = (sum(r0.*r0) - sum(r1.*r1))/sum(r0.*r0)
                x += 1;
                10*log10(sqrt(sum(r0.*r0))/sqrt(sum(r1.*r1))) > MaxER ?
                (LOG = "Decomposition hit energy ratio criteria."; break) : nothing
            end

            IMFs[n,:] = r1
            Xt .-= r1
            IMFs[Scales+1,:] = r0 .+ mean(X)
            n+=1
        end
        LOG = "All went well :) "
        return IMFs, LOG
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


text:
using Compose, GraphPlot, Cairo, Gadfly

function drawTunnel(topology, a, T, Tf, edges, num_nodes, dir, algorithm)
    edgelabels = []
    graph = LightGraphs.SimpleDiGraph(num_nodes)
    ntunnels = length(T)
    nedges = length(edges)

    L = zeros(ntunnels, nedges)
    for t in 1:ntunnels
        for e in 1:nedges
            if in(e, T[t])
                L[t,e] = 1
            end
        end
    end
    for e in 1:size(edges,1)
        LightGraphs.add_edge!(graph, edges[e][1], edges[e][2])
        s = 0
        for f in 1:size(a,1)
            for t in 1:size(a,2)
                s += a[f,t] * L[Tf[f][t],e]
            end
        end
        s = round(s, digits=2)
        push!(edgelabels, s)
    end
    nodelabel = collect(1:num_nodes)
    Compose.draw(PNG("$(dir)/$(topology)/00_$(algorithm)_tunnel_graph.png", 1000, 1000), gplot(graph, edgelabelc="white", EDGELABELSIZE = 20.0, NODELABELSIZE=20.0, nodelabel=nodelabel, edgelabeldisty=0.5, edgelabel=edgelabels, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
end

function drawGraph(topology, IPTopo, topology_index, OpticalTopo, dir)
    graph = LightGraphs.SimpleDiGraph(length(IPTopo["nodes"]))
    for e in 1:length(IPTopo["links"])
        LightGraphs.add_edge!(graph, IPTopo["links"][e][1], IPTopo["links"][e][2])
        # println("$(IPTopo["links"][e][1])->$(IPTopo["links"][e][2])@$(IPTopo["capacity"][e])")
    end

    iplink_bw = []  # the sequence of labeling is internal decided, different from the sequence of input files
    for ss in 1:length(IPTopo["nodes"])
        for dd in 1:length(IPTopo["nodes"])
            if LightGraphs.has_edge(graph, ss, dd)
                for e in 1:length(IPTopo["links"])
                    if ss==IPTopo["links"][e][1] && dd==IPTopo["links"][e][2]
                        push!(iplink_bw, IPTopo["capacity"][e])
                    end
                end
            end
        end
    end

    nodelabel = collect(1:length(IPTopo["nodes"]))
    Compose.draw(PNG("$(dir)/$(topology)/00_IP_topo_graph.png", 1000, 1000), gplot(graph, edgelabelc="white", EDGELABELSIZE = 20.0, NODELABELSIZE=20.0, nodelabel=nodelabel, edgelabeldisty=0.5, edgelabel=iplink_bw, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    if !isfile("data/topology/$(topology)/IP_topo_$(topology_index)/IP_topo_graph.png")
        Compose.draw(PNG("data/topology/$(topology)/IP_topo_$(topology_index)/IP_topo_graph.png", 1000, 1000), gplot(graph, edgelabelc="white", EDGELABELSIZE = 20.0, NODELABELSIZE=20.0, nodelabel=nodelabel, edgelabeldisty=0.5, edgelabel=iplink_bw, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    end

    optical_graph = LightGraphs.SimpleDiGraph(length(OpticalTopo["nodes"]))
    for e in 1:length(OpticalTopo["links"])
        LightGraphs.add_edge!(optical_graph, OpticalTopo["links"][e][1], OpticalTopo["links"][e][2])
    end

    fiber_capacity = []  # the sequence of labeling is internal decided, different from the sequence of input files
    fiber_spectrum_availability = []
    for ss in 1:length(OpticalTopo["nodes"])
        for dd in 1:length(OpticalTopo["nodes"])
            if LightGraphs.has_edge(optical_graph, ss, dd)
                for e in 1:length(OpticalTopo["links"])
                    if ss==OpticalTopo["links"][e][1] && dd==OpticalTopo["links"][e][2]
                        fiber_available_wave_num = sum(OpticalTopo["capacityCode"][e,:])
                        push!(fiber_capacity, fiber_available_wave_num)
                        fiber_available_spectrum = []
                        for w in 1:length(OpticalTopo["capacityCode"][e,:])
                            if OpticalTopo["capacityCode"][e,w] == 1
                                push!(fiber_available_spectrum, w)
                            end
                        end
                        push!(fiber_spectrum_availability, fiber_available_spectrum)
                    end
                end
            end
        end
    end

    optical_nodelabel = collect(1:length(OpticalTopo["nodes"]))
    if !isfile("data/topology/$(topology)/optical_topo_graph.png")
        Compose.draw(PNG("data/topology/$(topology)/optical_topo_graph.png", 1000, 1000), gplot(optical_graph, edgelabelc="yellow", layout=spring_layout, NODELABELSIZE=20.0, nodelabel=optical_nodelabel, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    end
    Compose.draw(PNG("$(dir)/$(topology)/00_optical_topo_graph_residual_fiber_capacity.png", 1000, 1000), gplot(optical_graph, edgelabelc="yellow", layout=spring_layout, EDGELABELSIZE = 20.0, NODELABELSIZE=20.0, nodelabel=optical_nodelabel, edgelabeldisty=0.5, edgelabel=fiber_capacity, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    if !isfile("data/topology/$(topology)/IP_topo_$(topology_index)/optical_topo_graph_residual_fiber_capacity.png")
        Compose.draw(PNG("data/topology/$(topology)/IP_topo_$(topology_index)/optical_topo_graph_residual_fiber_capacity.png", 1000, 1000), gplot(optical_graph, edgelabelc="yellow", layout=spring_layout, EDGELABELSIZE = 20.0, NODELABELSIZE=20.0, nodelabel=optical_nodelabel, edgelabeldisty=0.5, edgelabel=fiber_capacity, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    end
    Compose.draw(PNG("$(dir)/$(topology)/00_optical_topo_graph_spectrum_availability.png", 2000, 2000), gplot(optical_graph, edgelabelc="yellow", layout=spring_layout, EDGELABELSIZE = 15.0, NODELABELSIZE=20.0, nodelabel=optical_nodelabel, edgelabeldisty=0.5, edgelabel=fiber_spectrum_availability, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    if !isfile("data/topology/$(topology)/IP_topo_$(topology_index)/optical_topo_graph_spectrum_availability.png")
        Compose.draw(PNG("data/topology/$(topology)/IP_topo_$(topology_index)/optical_topo_graph_spectrum_availability.png", 2000, 2000), gplot(optical_graph, edgelabelc="yellow", layout=spring_layout, EDGELABELSIZE = 15.0, NODELABELSIZE=20.0, nodelabel=optical_nodelabel, edgelabeldisty=0.5, edgelabel=fiber_spectrum_availability, EDGELINEWIDTH=3.0, arrowlengthfrac=.04))
    end
end

max_stars_repo_name: hipersys-team/arrow

max_issues_repo_name: arrow-bow/arrow

max_forks_repo_name: arrow-bow/arrow


text:
# Tests for tree structures

@testset "Tree structures" begin

    periods = 36

    ## Build a test FullCPIBase
    varbase = getzerobase(
        T_type = Float32,
        G = 10,
        T_periods = periods,
        startdate = Date(2001, 1)
    )

    # Codes with characters hierarchy defined by sequence (3, 4, 5, 7)
    # Try several combinations
    codes = [
        "_011101", # One division, single subgroup, single item
        "_011201", # More items per subgroup
        "_011202",
        "_021101", # Higher groups with only one Item
        "_022101",
        "_031101", # One division with single Item
        "_041101", # One division, one agrupation, several sub-agrupations and items
        "_041201",
        "_041202",
        "_041301",
    ]

    # Map some individual names for the Item objects
    item_names = "Item " .* string.('A':'J')

    # Create some group_codes
    group_codes = [
        "_01",
        "_02",
        "_03",
        "_04",
        "_011",
        "_021",
        "_022",
        "_031",
        "_041",
        "_0111",
        "_0112",
        "_0211",
        "_0221",
        "_0311",
        "_0411",
        "_0412",
        "_0413",
    ]

    # Map generic group codes according to length
    group_names = map(group_codes) do code
        l = length(code)
        l == 3 && return "Div." * code
        l == 4 && return "Agr." * code
        l == 5 && return "Subgr." * code
        return "Gen." * code
    end

    # Modify last columns to check for correct upward computations
    v = varbase.v
    # This should scale 150 to the division level for the last period
    v[:, 6] .= 100*((150.0/100)^(1/periods) - 1)
    # This should scale 120 to the subgroup level for the last period
    v[:, 9] .= 100*((120.0/100)^(1/periods) - 1)

    base = FullCPIBase(
        capitalize(varbase.v),
        varbase.v,
        varbase.w,
        varbase.dates,
        varbase.baseindex,
        codes,
        item_names
    )

    ## Build items and groups manually
    all_items = [Item(base.codes[i], base.names[i], base.w[i]) for i in 1:10]
    group = Group("_0", "Single group", all_items)
    @info "Test building items and groups manually"
    print_tree(group)

    @test length(children(group)) == 10
    @test abs(sum(c.weight for c in children(group)) - 100) <= 1e-2

    ## Try recursively building the tree
    @info "Test recursively building the CPI tree"
    cpi_test_tree = get_cpi_tree(
        base = base,
        group_names = group_names,
        group_codes = group_codes,
        characters = (3, 4, 5, 7)
    )
    print_tree(cpi_test_tree)

    # Structure tests
    @test cpi_test_tree["_0"] === cpi_test_tree # returns the same structure
    @test cpi_test_tree["_01"] isa Group
    @test length(children(cpi_test_tree)) == 4
    @test abs(cpi_test_tree.weight - 100) < 1e-2

    ## Tests to check the computations of compute_index
    @info "Tests to check the computations of compute_index"
    # Test upper level result
    @test all(compute_index(cpi_test_tree, base) .!= 100)
    # Test other levels

    # Division with single Item should scale the values from the bottom
    @test all(compute_index(cpi_test_tree["_03"], base) .≈ compute_index(cpi_test_tree["_031101"], base))
    @test last(compute_index(cpi_test_tree["_03"], base)) ≈ 150

    # These results should be the same as the bottom ones, as price changes were not modified
    @test all(compute_index(cpi_test_tree["_011"], base) .≈ 100)
    @test all(compute_index(cpi_test_tree["_0111"], base) .≈ 100)
    @test all(compute_index(cpi_test_tree["_0112"], base) .≈ 100)
    @test all(compute_index(cpi_test_tree["_011101"], base) .≈ 100)

    # Testing how values scale from Item to Group
    result = base.ipc[:, 8:9] * base.w[8:9] / sum(base.w[8:9])
    @test all(compute_index(cpi_test_tree["_0412"], base) .≈ result)

    # Testing how values scale from Group to upper Group. The rest of upward
    # categories work in the same way. It is recursive =)
    result = base.ipc[:, 7:end] * base.w[7:end] / sum(base.w[7:end])
    @test all(compute_index(cpi_test_tree["_041"], base) .≈ result)

    # This node does not exist
    @info "Test producing a warning for node not found"
    @test compute_index(cpi_test_tree["_011102"], base) === nothing

    ## Tests to check compute_index!
    @info "Tests to check compute_index!"

    # Create a cache dict
    d = Dict{String, Vector{Float32}}()

    # Compute all the nodes and save them in d
    idx = compute_index!(d, cpi_test_tree, base)
    # Check for the length of the
    @test length(d) == length(group_codes) + length(codes) + 1 # Groups + Items + upper level
    # Check that all group names were computed and cached
    @test all(key in keys(d) for key in group_codes)

    @info "Test producing a warning for node not found"
    @test compute_index(cpi_test_tree["_011102"], base) === nothing

    ## Test CPITree wrapper type
    @info "Tests with CPITree wrapper type"
    cpi_wrapper = CPITree(base, cpi_test_tree, group_names, group_codes)

    # Test indexing the top level code returns the same object
    @test cpi_wrapper["_0"] === cpi_wrapper

    # Test under-the-hood call to compute_index
    @test all(compute_index(cpi_wrapper, "_0") .== compute_index(cpi_test_tree["_0"], base))

    # Test computing an index not available
    @test compute_index(cpi_wrapper, "_011102") === nothing
    # Test indexing returns a proper CPITree
    @test hasproperty(cpi_wrapper["_01"], :group_names)
    @test hasproperty(cpi_wrapper["_01"], :group_codes)

    # Compute index with one argument
    @test all(compute_index(cpi_wrapper["_01"]) .== compute_index(cpi_wrapper, "_01"))
end


text:
mutable struct Mesures

    deb_mesure::Array{Int32,1} ;
    temps_mesure::Array{Int32,1} ;

    function Mesures(nb_pop::Int32)

        deb_mesure = zeros(Int32,nb_pop) ;
        temps_mesure = zeros(Int32,nb_pop) ;

        if(mesures.tracage)
            deb_mesure .= Trace.delai_trace ;
            temps_mesure .= Trace.jours_conf ;
        end

        return new(deb_mesure,temps_mesure)

    end

end


text:
dy = iy - ix
dz = iz - ix
@inbounds for i = ix:ix-1+n
    x[i] = $op( y[dy+i], z[dz+i] )
end