import ReverseDiff
using ReverseDiff: SpecialInstruction, istracked, increment_deriv!, deriv, TrackedArray
using ReverseDiff: track, InstructionTape, TrackedReal, seed!, unseed!, reverse_pass!, record!
import ReverseDiff: special_reverse_exec!

mutable struct GFBackpropTraceState
    trace::GFTrace
    score::TrackedReal
    tape::InstructionTape
    visitor::AddressVisitor
    params::Dict{Symbol,Any}
    selection::AddressSet
    tracked_choices::HomogenousTrie{Any,TrackedReal}
    value_trie::DynamicAssignment
    gradient_trie::DynamicAssignment
end

function GFBackpropTraceState(trace, selection, params, tape)
    score = track(0., tape)
    visitor = AddressVisitor()
    tracked_choices = HomogenousTrie{Any,TrackedReal}()
    value_trie = DynamicAssignment()
    gradient_trie = DynamicAssignment()
    GFBackpropTraceState(trace, score, tape, visitor, params,
                       selection, tracked_choices, value_trie, gradient_trie)
end

get_args_change(state::GFBackpropTraceState) = nothing
get_addr_change(state::GFBackpropTraceState, addr) = nothing
set_ret_change!(state::GFBackpropTraceState, value) = begin end

function fill_gradient_trie!(gradient_trie::DynamicAssignment,
                             tracked_trie::HomogenousTrie{Any,TrackedReal})
    for (key, tracked) in get_leaf_nodes(tracked_trie)
        set_leaf_node!(gradient_trie, key, deriv(tracked))
    end
    # NOTE: there should be no address collision between these primitive
    # choices and the generator invocations, as enforced by the visitor
    for (key, node) in get_internal_nodes(tracked_trie)
        @assert !has_leaf_node(gradient_trie, key) && !has_internal_node(gradient_trie, key)
        gradient_trie_node = DynamicAssignment()
        fill_gradient_trie!(gradient_trie_node, node)
        set_internal_node!(gradient_trie, key, gradient_trie_node)
    end
end

function fill_value_trie!(value_trie::DynamicAssignment,
                          tracked_trie::HomogenousTrie{Any,TrackedReal})
    for (key, tracked) in get_leaf_nodes(tracked_trie)
        set_leaf_node!(value_trie, key, ReverseDiff.value(tracked))
    end
    # NOTE: there should be no address collision between these primitive
    # choices and the generator invocations, as enforced by the visitor
    for (key, node) in get_internal_nodes(tracked_trie)
        @assert !has_leaf_node(value_trie, key) && !has_internal_node(value_trie, key)
        value_trie_node = DynamicAssignment()
        fill_value_trie!(value_trie_node, node)
        set_internal_node!(value_trie, key, value_trie_node)
    end
end

function addr(state::GFBackpropTraceState, dist::Distribution{T}, args, addr) where {T}
    visit!(state.visitor, addr)
    call::CallRecord = get_primitive_call(state.trace, addr)
    retval::T = call.retval
    if has_internal_node(state.selection, addr)
        error("Got internal node but expected leaf node in selection at $addr")
    end
    if has_leaf_node(state.selection, addr)
        tracked_retval = track(retval, state.tape)
        set_leaf_node!(state.tracked_choices, addr, tracked_retval)
        score_tracked = logpdf(dist, tracked_retval, args...)
        state.score += score_tracked
        return tracked_retval
    else
        state.score += logpdf(dist, retval, args...)
        return retval
    end
end

struct BackpropTraceRecord
    generator::Generator
    subtrace::Any
    selection::AddressSet
    value_trie::DynamicAssignment
    gradient_trie::DynamicAssignment
    addr::Any
end

function addr(state::GFBackpropTraceState, gen::Generator{T}, args, addr, args_change) where {T}
    visit!(state.visitor, addr)
    if has_leaf_node(state.selection, addr)
        error("Cannot select a whole subtrace, tried to select $addr")
    end
    trace = get_subtrace(state.trace, addr)
    call::CallRecord = get_call_record(trace) # use the return value recorded in the trace
    retval::T = call.retval
    if accepts_output_grad(gen)
        retval_maybe_tracked = track(retval, state.tape)
        if !istracked(retval_maybe_tracked)
            error("Could not track return value at address $addr on AD tape.")
        end
    else
        retval_maybe_tracked = retval
    end
    # some of the args may be tracked (see special_reverse_exec!)
    # note: we still need to run backprop_params on gen, even if it does not
    # accept an output gradient, because it may make random choices.
    if has_internal_node(state.selection, addr)
        selection = get_internal_node(state.selection, addr)
    else
        selection = EmptyAddressSet()
    end
    record = BackpropTraceRecord(gen, trace, selection, state.value_trie,
                                 state.gradient_trie, addr)
    record!(state.tape, SpecialInstruction, record, (args...,), retval_maybe_tracked)
    retval_maybe_tracked
end

function backprop_trace(gf::GenFunction, trace::GFTrace, selection::AddressSet, retval_grad)
    tape = InstructionTape()
    state = GFBackpropTraceState(trace, selection, gf.params, tape)
    call = get_call_record(trace)
    args = call.args
    args_maybe_tracked = (map(maybe_track, args, gf.has_argument_grads, fill(tape, length(args)))...,)
    retval_maybe_tracked = exec(gf, state, args_maybe_tracked)
    if gf.accepts_output_grad
        if !istracked(retval_maybe_tracked)
            error("Output of $gf accepts gradient but return value was not tracked on AD tape.")
        end
        if retval_grad != nothing
            # it is not an error if retval_grad is nothing, it means the gradient is zero?
            deriv!(retval_maybe_tracked, retval_grad)
        end
    end
    seed!(state.score)
    reverse_pass!(tape)

    # fill trace gradient with gradients with respect to primitive random choices
    fill_gradient_trie!(state.gradient_trie, state.tracked_choices)
    fill_value_trie!(state.value_trie, state.tracked_choices)

    # return gradients with respect to inputs
    # NOTE: if a value isn't tracked the gradient is nothing
    input_grads::Tuple = (map((arg, has_grad) -> has_grad ? deriv(arg) : nothing,
                             args_maybe_tracked, gf.has_argument_grads)...,)

    (input_grads, state.value_trie, state.gradient_trie)
end

@noinline function special_reverse_exec!(instruction::SpecialInstruction{BackpropTraceRecord})
    record = instruction.func
    gen = record.generator
    args_maybe_tracked = instruction.input
    retval_maybe_tracked = instruction.output
    if accepts_output_grad(gen)
        @assert istracked(retval_maybe_tracked)
        retval_grad = deriv(retval_maybe_tracked)
    else
        retval_grad = nothing
    end
    (arg_grads, value_trie, grad_trie) = backprop_trace(
        gen, record.subtrace, record.selection, retval_grad)

    @assert !has_internal_node(record.gradient_trie, record.addr)
    @assert !has_leaf_node(record.gradient_trie, record.addr)
    set_internal_node!(record.gradient_trie, record.addr, grad_trie)
    set_internal_node!(record.value_trie, record.addr, value_trie)

    for (arg, grad, has_grad) in zip(args_maybe_tracked, arg_grads, has_argument_grads(gen))
        if has_grad && istracked(arg)
            increment_deriv!(arg, grad)
        end
    end
    nothing
end


module CineFiles

using ColorTypes: Color, N0f8, N6f10, N4f12, N0f16, gray, Gray, AbstractGray
using ColorVectorSpace
using LRUCache
using StaticArrays: SizedMatrix, SizedVector
using FixedPointNumbers: Normed, reinterpret, nbitsfrac

import Base: eltype, length, size, getindex, firstindex, lastindex, iterate, read, read!
export CineFile

include("LUT.jl")

abstract type RawFrame end
abstract type BinaryData end

rawcounts(v::AbstractGray) = reinterpret(gray(v))

struct Levels{T}
    white::T
    black::T
end

function Levels{T}(W::S, B::S) where {T<:AbstractGray, S<:Int32}
    scale = rawcounts(T(1))
    Levels{T}(T(W / scale), T(B / scale))
end

struct Unpacked{T} <: RawFrame
    width::Int
    height::Int
    tmp::Array{T,2}
    unpacked::Array{T,2}
    levels::Levels
    Unpacked{T}(width, height, whiteL, blackL) where {T<:Color} = new{T}(
        width,
        height,
        Array{T,2}(undef, width, height),
        Array{T,2}(undef, width, height),
        Levels{T}(whiteL, blackL),
    )
end

struct Packed{T,S} <: RawFrame
    width::Int
    height::Int
    tmp::Array{S,2}
    unpacked::Array{T,2}
    pack::Vector{UInt8}
    levels::Levels
end

function Packed{T,S}(width, height, whiteL, blackL) where {T<:Gray,S<:Gray}
    true_bits = nbitsfrac(eltype(T)) * width * height
    true_bytes = ceil(Int, true_bits / 8)
    Packed{T,S}(
        width,
        height,
        Array{S,2}(undef, width, height),
        Array{T,2}(undef, width, height),
        Vector{UInt8}(undef, true_bytes),
        Levels{T}(whiteL, blackL),
    )
end

function Packed{T}(width, height, whiteL, blackL) where {T<:Gray}
    true_bits = nbitsfrac(eltype(T)) * width * height
    true_bytes = ceil(Int, true_bits / 8)
    Packed{T,T}(
        width,
        height,
        Array{T,2}(undef, width, height),
        Array{T,2}(undef, width, height),
        Vector{UInt8}(undef, true_bytes),
        Levels{T}(whiteL, blackL),
    )
end

struct Time64 <: BinaryData
    fractions::UInt32
    seconds::UInt32
end

struct WBgain <: BinaryData
    R::Float32
    B::Float32
end

struct IMfilter <: BinaryData
    Dim::Int32
    Shifts::Int32
    Bias::Int32
    Coef::SizedMatrix{5,5,Int32,2,Matrix{Int32}}
end

struct Rect <: BinaryData
    Left::Int32
    Top::Int32
    Right::Int32
    Bottom::Int32
end

struct TimeCode <: BinaryData
    framesU::UInt8
    framesT::UInt8
    dropFrameFlag::UInt8
    colorFrameFlag::UInt8
    secondsU::UInt8
    secondsT::UInt8
    flag1::UInt8
    minutesU::UInt8
    minutesT::UInt8
    flag2::UInt8
    hoursU::UInt8
    hoursT::UInt8
    flag3::UInt8
    flag4::UInt8
    userBitData::UInt32
end

struct SetupHeader <: BinaryData
    FrameRate16::UInt16
    Shutter16::UInt16
    PostTrigger16::UInt16
    FrameDelay16::UInt16
    AspectRatio::UInt16
    Res7::UInt16
    Res8::UInt16
    Res9::UInt8
    Res10::UInt8
    Res11::UInt8
    TrigFrame::UInt8
    Res12::UInt8
    DescriptionOld::SizedVector{121,Int8}
    Mark::UInt16
    Length::UInt16
    Res13::UInt16
    SigOption::UInt16
    BinChannels::Int16
    SamplesPerImage::UInt8
    BinName::SizedMatrix{8,11,Int8}
    AnaOption::UInt16
    AnaChannels::Int16
    Res6::UInt8
    AnaBoard::UInt8
    ChOption::SizedVector{8,Int16}
    AnaGain::SizedVector{8,Float32}
    AnaUnit::SizedMatrix{8,6,Int8}
    AnaName::SizedMatrix{8,11,Int8}
    lFirstImage::Int32
    dwImageCount::UInt32
    nQFactor::Int16
    wCineFileType::UInt16
    szCinePath::SizedMatrix{4,65,Int8}
    Res14::UInt16
    Res15::UInt8
    Res16::UInt8
    Res17::UInt16
    Res18::Float64
    Res19::Float64
    Res20::UInt16
    Res1::Int32
    Res2::Int32
    Res3::Int32
    ImWidth::UInt16
    ImHeight::UInt16
    EDRShutter16::UInt16
    Serial::UInt32
    Saturation::Int32
    Res5::UInt8
    AutoExposure::UInt32
    bFlipH::Int32
    bFlipV::Int32
    Grid::UInt32
    FrameRate::UInt32
    Shutter::UInt32
    EDRShutter::UInt32
    PostTrigger::UInt32
    FrameDelay::UInt32
    bEnableColor::Int32
    CameraVersion::UInt32
    FirmwareVersion::UInt32
    SoftwareVersion::UInt32
    RecordingTimeZone::Int32
    CFA::UInt32
    Bright::Int32
    Contrast::Int32
    Gamma::Int32
    Res21::UInt32
    AutoExpLevel::UInt32
    AutoExpSpeed::UInt32
    AutoExpRect::Rect
    WBGain::SizedVector{4,WBgain}
    Rotate::Int32
    WBView::WBgain
    RealBPP::UInt32
    Conv8Min::UInt32
    Conv8Max::UInt32
    FilterCode::Int32
    FilterParam::Int32
    UF::IMfilter
    BlackCalSVer::UInt32
    WhiteCalSVer::UInt32
    GrayCalSVer::UInt32
    bStampTime::Int32
    SoundDest::UInt32
    FRPSteps::UInt32
    FRPImgNr::SizedVector{16,Int32}
    FRPRate::SizedVector{16,UInt32}
    FRPExp::SizedVector{16,UInt32}
    MCCnt::Int32
    MCPercent::SizedVector{64,Float32}
    CICalib::UInt32
    CalibWidth::UInt32
    CalibHeight::UInt32
    CalibRate::UInt32
    CalibExp::UInt32
    CalibEDR::UInt32
    CalibTemp::UInt32
    HeadSerial::SizedVector{4,UInt32}
    RangeCode::UInt32
    RangeSize::UInt32
    Decimation::UInt32
    MasterSerial::UInt32
    Sensor::UInt32
    ShutterNs::UInt32
    EDRShutterNs::UInt32
    FrameDelayNs::UInt32
    ImPosXAcq::UInt32
    ImPosYAcq::UInt32
    ImWidthAcq::UInt32
    ImHeightAcq::UInt32
    Description::SizedVector{4096,Int8}
    RisingEdge::Int32
    FilterTime::UInt32
    LongReady::Int32
    ShutterOff::Int32
    Res4::SizedVector{16,UInt8}
    bMetaWB::Int32
    Hue::Int32
    BlackLevel::Int32
    WhiteLevel::Int32
    LensDescription::SizedVector{256,Int8}
    LensAperture::Float32
    LensFocusDistance::Float32
    LensFocalLength::Float32
    fOffset::Float32
    fGain::Float32
    fSaturation::Float32
    fHue::Float32
    fGamma::Float32
    fGammaR::Float32
    fGammaB::Float32
    fFlare::Float32
    fPedestalR::Float32
    fPedestalG::Float32
    fPedestalB::Float32
    fChroma::Float32
    ToneLabel::SizedVector{256,Int8}
    TonePoints::Int32
    fTone::SizedVector{64,Float32}
    UserMatrixLabel::SizedVector{256,Int8}
    EnableMatrices::Int32
    fUserMatrix::SizedVector{9,Float32}
    EnableCrop::Int32
    CropRect::Rect
    EnableResample::Int32
    ResampleWidth::UInt32
    ResampleHeight::UInt32
    fGain16_8::Float32
    FRPShape::SizedVector{16,UInt32}
    TrigTC::TimeCode
    fPbRate::Float32
    fTcRate::Float32
    CineName::SizedVector{256,UInt8}
end

struct CineFileHeader <: BinaryData
    Type::UInt16
    HeaderSize::UInt16
    Compression::UInt16
    Version::UInt16
    FirstMovieImage::Int32
    TotalImageCount::UInt32
    FirstImageNum::Int32
    ImageCount::UInt32
    OffImageHeader::UInt32
    OffSetup::UInt32
    OffImageOffsets::UInt32
    TriggerTime::Time64
end

struct BitmapInfoHeader <: BinaryData
    Size::UInt32
    Width::Int32
    Height::Int32
    Planes::UInt16
    BitCount::UInt16
    Compression::UInt32
    SizeImage::UInt32
    XPelsPerMeter::Int32
    YPelsPerMeter::Int32
    ClrUsed::UInt32
    ClrImportant::UInt32
    FPS::UInt16
end

struct CineHeader{T}
    cine::CineFileHeader
    bitmap::BitmapInfoHeader
    setup::SetupHeader
    imglocs::Vector{Int64}
    imgoffset::Int
    dt::Vector{Float64}
    raw::RawFrame
end

struct CineFile{T}
    path::String
    header::CineHeader{T}
    data::LRU{Int,Array{T,2}}
end

function unpack!(pack::Vector{UInt8}, unpacked::Array{Gray{N6f10}})
    pack16 = UInt16.(pack)
    indx = CartesianIndices(unpacked)[:, end:-1:1]
    unpacked[indx[1:4:end]] =
        reinterpret.(N6f10, (pack16[1:5:end] .<< 2) .| (pack16[2:5:end] .>> 6))
    unpacked[indx[2:4:end]] =
        reinterpret.(
            N6f10,
            ((pack16[2:5:end] .& 0b00111111) .<< 4) .| (pack16[3:5:end] .>> 4),
        )
    unpacked[indx[3:4:end]] =
        reinterpret.(
            N6f10,
            ((pack16[3:5:end] .& 0b00001111) .<< 6) .| (pack16[4:5:end] .>> 2),
        )
    unpacked[indx[4:4:end]] =
        reinterpret.(N6f10, ((pack16[4:5:end] .& 0b00000011) .<< 8) .| pack16[5:5:end])
end

function unpack!(pack::Vector{UInt8}, unpacked::Array{Gray{N4f12}})
    @warn "unpacking 12 bit images untested"
    pack16 = UInt16.(pack)
    indx = CartesianIndices(unpacked)[:, end:-1:1]
    unpacked[indx[1:2:end]] =
        reinterpret.(N6f10, (pack16[1:3:end] .<< 4) .| (pack16[2:3:end] .>> 4))
    unpacked[indx[2:2:end]] =
        reinterpret.(N6f10, ((pack16[2:3:end] .& 0b00001111) .<< 8) .| (pack16[3:3:end]))
end

function linearize!(raw_data::RawFrame)
    Wlevel = raw_data.levels.white
    Blevel = raw_data.levels.black
    raw_data.tmp .= raw_data.unpacked
    raw_data.tmp[raw_data.tmp.>Wlevel] .= Wlevel
    raw_data.tmp[raw_data.tmp.<Blevel] .= Blevel
    raw_data.tmp .= (raw_data.tmp .- Blevel) ./ (Wlevel - Blevel)
end

function linearize!(raw_data::Packed{Gray{N6f10}})
    Blevel, Wlevel = Gray{N4f12}.((64, 4064) ./ (2^12))
    raw_data.tmp .= lookup.(raw_data.unpacked, Ref(CINE_LUT))
    raw_data.tmp[raw_data.tmp.>Wlevel] .= Wlevel
    raw_data.tmp[raw_data.tmp.<Blevel] .= Blevel
    raw_data.tmp .= (raw_data.tmp .- Blevel) ./ (Wlevel - Blevel)
end

function CineHeader(fname)
    open(fname) do f
        read(f, UInt16) == UInt(18755) || error(basename(fname), " is not a .cine file")
        seek(f, 0)
        cine = read(f, CineFileHeader)
        seek(f, cine.OffImageHeader)
        bitmap = read(f, BitmapInfoHeader)

        seek(f, cine.OffSetup)
        setup = read(f, SetupHeader)

        bitsized = (bitmap.BitCount == 8) || (bitmap.BitCount == 24) ? UInt8 : UInt16

        if (bitmap.BitCount == 8) || (bitmap.BitCount == 16)
            bittype = Gray{Normed{bitsized,Int(setup.RealBPP)}}
        elseif (bitmap.BitCount == 24) || (bitmap.BitCount == 36)
            bittype = BGR{Normed{bitsized,Int(setup.RealBPP)}}
        end

        cine.ImageCount > 0 || error("no images exist in file")

        seek(f, cine.OffImageOffsets)
        imglocs = read!(f, Array{Int64}(undef, cine.ImageCount))

        seekstart(f)
        imgoffset = 1
        while read(f, UInt16) != 1002
            imgoffset += 1
        end

        dt = zeros(cine.ImageCount)
        skip(f, 2)
        for i in eachindex(dt)
            fracstart = read(f, UInt32)
            secstart = read(f, UInt32)
            dt[i] =
                (secstart - cine.TriggerTime.seconds) +
                ((fracstart - cine.TriggerTime.fractions) / 2^32)
        end

        # Packing options
        if bitmap.Compression == 0
            raw = Unpacked{bittype}(
                bitmap.Width,
                bitmap.Height,
                setup.WhiteLevel,
                setup.BlackLevel,
            )
            return CineHeader{bittype}(cine, bitmap, setup, imglocs, imgoffset, dt, raw)
        elseif bitmap.Compression == 256
            raw = Packed{bittype,Gray{N6f10}}(
                bitmap.Width,
                bitmap.Height,
                setup.WhiteLevel,
                setup.BlackLevel,
            )
            return CineHeader{Gray{N4f12}}(cine, bitmap, setup, imglocs, imgoffset, dt, raw)
        elseif bitmap.Compression == 1024
            raw = Packed{bittype,Gray{N4f12}}(
                bitmap.Width,
                bitmap.Height,
                setup.WhiteLevel,
                setup.BlackLevel,
            )
            return CineHeader{bittype}(cine, bitmap, setup, imglocs, imgoffset, dt, raw)
        end
    end
end

function readframe!(f::IO, frame, h, frameidx)
    1 <= frameidx <= h.cine.ImageCount || BoundsError(h.dt, frameidx)
    seek(f, h.imglocs[frameidx])
    skip(f, read(f, UInt32) - 4)
    read!(f, h.raw)
    linearize!(h.raw)
    # non-allocating version of rotl90
    ind1, ind2 = axes(h.raw.tmp)
    n = first(ind2) + last(ind2)
    for i in axes(h.raw.tmp, 1), j in ind2
        frame[n-j, i] = h.raw.tmp[i, j]
    end
    return frame
end

function read!(f::IOStream, frame::Packed)
    read!(f, frame.pack)
    unpack!(frame.pack, frame.unpacked)
end

read!(f::IOStream, frame::Unpacked) = read!(f, frame.unpacked)
# TODO: implement with lower compilation overhead
read(f::IOStream, S::Type{T}) where {T<:BinaryData} = S(read.(Ref(f), S.types)...)

readframe!(filename, frame, h, frameidx) =
    open(f -> readframe!(f, frame, h, frameidx), filename)
readframe(f, h, frameidx) =
    readframe!(f, similar(h.raw.tmp, size(h.raw.tmp, 2), size(h.raw.tmp, 1)), h, frameidx)

"""
    CineFile(filepath, cachelimit=0.25)

Load header information and create a frame cache for a Phantom .cine file.
    `cachelimit` sets the maximum size of the frame cache as a fraction of
    your system's free RAM. Indexing and iteration is supported for CineFiles.

"""
function CineFile(filepath, cachelimit = 0.25)
    header = CineHeader(filepath)
    framesize = Base.summarysize(header.raw.tmp)
    maxcachedframes = ceil(Int, cachelimit * Sys.free_memory() / framesize)
    data = LRU{Int,Array{eltype(header),2}}(maxsize = maxcachedframes)
    return CineFile(filepath, header, data)
end

Base.eltype(h::CineHeader{T}) where {T} = T
Base.length(cf::CineFile) = length(cf.header.dt)
Base.eltype(cf::CineFile) = typeof(cf.header.raw.tmp)
Base.size(cf::CineFile) = (length(cf), size(cf.header.raw.tmp, 2), size(cf.header.raw.tmp, 1))

function Base.getindex(cf::CineFile, idx::Int)
    get!(cf.data, idx) do
        readframe(cf.path, cf.header, idx)
    end
end

Base.firstindex(cf::CineFile) = 1
Base.lastindex(cf::CineFile) = length(cf)
Base.getindex(cf::CineFile, I) = [cf[i] for i in I]
Base.iterate(cf::CineFile, state = 1) = state > length(cf) ? nothing : (cf[state], state + 1)

# precompile as the final step of the module definition:
if ccall(:jl_generating_output, Cint, ()) == 1   # if we're precompiling the package
    let
        precompile(read, (IOStream, SetupHeader))

        # CineFile(joinpath(dirname(@__DIR__), "test/data/8bpp.cine"))[1]
    end
end

end


approximator_GNN = GeometricFlux.GraphConv(64 => 64, Flux.leakyrelu)
target_approximator_GNN = GeometricFlux.GraphConv(64 => 64, Flux.leakyrelu)
gnnlayers = 1
approximator_model = SeaPearl.CPNN(
    graphChain = Flux.Chain(
        GeometricFlux.GraphConv(3 => 64, Flux.leakyrelu),
        [approximator_GNN for i = 1:gnnlayers]...
    ),
    nodeChain = Flux.Chain(
        Flux.Dense(64, 32, Flux.leakyrelu),
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 16, Flux.leakyrelu),
    ),
    outputChain = Flux.Dense(16, 4),
) |> cpu
target_approximator_model = SeaPearl.CPNN(
    graphChain = Flux.Chain(
        GeometricFlux.GraphConv(3 => 64, Flux.leakyrelu),
        [target_approximator_GNN for i = 1:gnnlayers]...
    ),
    nodeChain = Flux.Chain(
        Flux.Dense(64, 32, Flux.leakyrelu),
        Flux.Dense(32, 32, Flux.leakyrelu),
        Flux.Dense(32, 16, Flux.leakyrelu),
    ),
    outputChain = Flux.Dense(16, 4),
) |> cpu

agent = RL.Agent(
    policy = RL.QBasedPolicy(
        learner = RL.DQNLearner(
            approximator = RL.NeuralNetworkApproximator(
                model = approximator_model,
                optimizer = ADAM(0.001f0)
            ),
            target_approximator = RL.NeuralNetworkApproximator(
                model = target_approximator_model,
                optimizer = ADAM(0.001f0)
            ),
            loss_func = Flux.Losses.huber_loss,
            stack_size = nothing,
            γ = 0.99f0,
            batch_size = 32,
            update_horizon = 1,
            min_replay_history = 1,
            update_freq = 1,
            target_update_freq = 100,
        ),
        explorer = RL.EpsilonGreedyExplorer(
            ϵ_stable = 0.01,
            kind = :exp,
            ϵ_init = 1.0,
            warmup_steps = 0,
            decay_steps = 500,
            step = 1,
            is_break_tie = false,
            #is_training = true,
        )
    ),
    trajectory = RL.CircularArraySLARTTrajectory(
        capacity = 500,
        state = SeaPearl.DefaultTrajectoryState[] => (),
        action = Int => (),
        legal_actions_mask = Vector{Bool} => (4, ),
    )
)

@testset "learning.jl" begin
    include("environment.jl")
    include("rewards/rewards.jl")
    include("utils.jl")
    include("learnedheuristic.jl")

    LearnedHeristicBasicConstructor = SeaPearl.LearnedHeuristic(agent)
    @test isa(LearnedHeristicBasicConstructor,SeaPearl.LearnedHeuristic{SeaPearl.DefaultStateRepresentation{SeaPearl.DefaultFeaturization, SeaPearl.DefaultTrajectoryState}, SeaPearl.DefaultReward, SeaPearl.FixedOutput})
    @test isnothing(LearnedHeristicBasicConstructor.fitted_problem)
    @test isnothing(LearnedHeristicBasicConstructor.fitted_strategy)
    @test isnothing(LearnedHeristicBasicConstructor.action_space)
    @test isnothing(LearnedHeristicBasicConstructor.current_state)
    @test isnothing(LearnedHeristicBasicConstructor.reward)
    @test isnothing(LearnedHeristicBasicConstructor.search_metrics)

end


# rna.jl
#
# Fit G state models (generalized telegraph models) to RNA abundance data
# For single cell RNA (scRNA) technical noise is included as a yieldfactor
# single molecule FISH (smFISH) is treated as loss less
#

"""
    fit_rna(nchains::Int,gene::String,cell::String,fittedparam::Vector,fixedeffects::Tuple,datacond,G::Int,maxtime::Float64,infolder::String,resultfolder::String,datafolder,fish::Bool,runcycle::Bool,inlabel,label,nsets::Int,cv=0.,transient::Bool=false,samplesteps::Int=100000,warmupsteps=20000,annealsteps=100000,temp=1.,tempanneal=100.,root = "/home/carsonc/scrna/",yieldprior = 0.05,ejectprior = 1.0)
    fit_rna(nchains::Int,data::AbstractRNAData,gene::String,cell::String,fittedparam::Vector,fixedeffects::Tuple,datacond,G::Int,maxtime::Float64,infolder::String,resultfolder::String,datafolder,fish::Bool,runcycle,inlabel,label,nsets,cv=0.,transient::Bool=false,samplesteps::Int=100000,warmupsteps=20000,annealsteps=100000,temp=1.,tempanneal=100.,root = "/home/carsonc/scrna/",yieldprior = 0.05,ejectprior = 1.0)

Fit steady state or transient GM model to RNA data for a single gene, write the result (through function finalize), and return nothing.

# Arguments
- `nchains`: number of MCMC chains
- `gene`: gene name
- `cell`: cell type
- `datacond`: condition, if more than one condition use vector of strings e.g. ["DMSO","AUXIN"]
- `maxtime`: float maximum time for entire run
- `infolder`: folder pointing to results used as initial conditions
- `resultfolder`: folder where results go
- `datafolder`: folder for data, string or array of strings
- `inlabel`: name of input files (not including gene name but including condition)
- `label`: = name of output files
- `nsets`: int number of rate sets
- `runcycle`: if true, cycle through all parameters sequentially in MCMC
- `cyclesteps`: int number of steps per cycle
- `samplesteps`: int number of samples
- `warmupsteps`: int number of warmup steps
- `annealsteps`: in number of annealing steps
- `temp`: MCMC temperature
- `tempanneal`: starting temperature for annealing
- `root`: root folder of data and Results folders
- `fittedparam`: vector of rate indices,  indices of parameters to be fit (input as string of ints separated by "-")
- `fixedeffects`: (tuple of vectors of rate indices) string indicating which rate is fixed, e.g. "eject"
- `data`: data structure

"""

function fit_rna(nchains::Int,gene::String,cell::String,fittedparam::Vector,fixedeffects::Tuple,datacond,G::Int,maxtime::Float64,infolder::String,resultfolder::String,datafolder,fish::Bool,runcycle::Bool,inlabel,label,nsets::Int,cv=0.,transient::Bool=false,samplesteps::Int=100000,warmupsteps=20000,annealsteps=100000,temp=1.,tempanneal=100.,root = ".",yieldprior=0.05)
    gene = check_genename(gene,"[")
    datafolder = joinpath("data",datafolder)
    if occursin("-",datafolder)
        datafolder = string.(split(datafolder,"-"))
    end
    if occursin("-",datacond)
        datacond = string.(split(datacond,"-"))
    end
    if transient
        data = data_rna(gene,datacond,datafolder,fish,label,root,["T0","T30","T120"],[0.,30.,120.])
    else
        data = data_rna(gene,datacond,datafolder,fish,label,root)
    end
    fit_rna(nchains,data,gene,cell,fittedparam,fixedeffects,datacond,G,maxtime,infolder,resultfolder,datafolder,fish,runcycle,inlabel,label,nsets,cv,transient,samplesteps,warmupsteps,annealsteps,temp,tempanneal,root,yieldprior)
end
function fit_rna(nchains::Int,data::AbstractRNAData,gene::String,cell::String,fittedparam::Vector,fixedeffects::Tuple,datacond,G::Int,maxtime::Float64,infolder::String,resultfolder::String,datafolder,fish::Bool,runcycle,inlabel,label,nsets,cv=0.,transient::Bool=false,samplesteps::Int=100000,warmupsteps=20000,annealsteps=100000,temp=1.,tempanneal=100.,root = ".",yieldprior=0.05)
    println(now())
    printinfo(gene,G,datacond,datafolder,infolder,resultfolder,maxtime)
    println("size of histogram: ",data.nRNA)

    resultfolder = joinpath("results",resultfolder)
    infolder = joinpath("results",infolder)

    if fish
        yieldprior = 1.
    end

    model = model_rna(gene,cell,G,cv,fittedparam,fixedeffects,inlabel,infolder,nsets,root,data,yieldprior)
    options = MHOptions(samplesteps,0,warmupsteps,annealsteps,maxtime*.9,temp,tempanneal)
    if runcycle > 0
        model = cycle(nchains,fish,fixedeffects,model,data,options)
    end
    print_ll(data,model)
    fit,stats,measures = run_mh(data,model,options,nchains);
    finalize(data,model,fit,stats,measures,temp,resultfolder,root)
    println(now())
    nothing
end

"""
cycle(nchains,fish,fixedeffects,model,data,options)

run_mh by cycling through individual parameters sequentially
"""
function cycle(nchains,fish,fixedeffects,model,data,options)
    maxtime = options.maxtime/10
    options = MHOptions(100,0,0,0,maxtime,options.temp,options.tempanneal)
    print_ll(data,model,"pre-cycle ll: ")
    t0 = time()
    nsets = length(data.nRNA)
    r = model.rates
    nalleles = model.nalleles
    G = model.G
    fittedparam = model.fittedparam
    cv = 0.02
    rateprior = model.rateprior
    while (time() - t0 < maxtime)
        for i in eachindex(fittedparam)
            model = model_rna(r,[rateprior[i]],G,nalleles,cv,[fittedparam[i]],fixedeffects,0)
            fit,_,_ = run_mh(data,model,options,nchains);
            r = get_rates(fit.parml,model)
        end
    end
    return model_rna(r,rateprior,G,nalleles,cv,fittedparam,fixedeffects,0)
end
"""
check_genename(gene,p1)

Check genename for p1
if p1 = "[" change to "("
(since swarm cannot parse "(")

"""
function check_genename(gene,p1)
    if occursin(p1,gene)
        if p1 == "["
            gene = replace(gene,"[" => "(")
            gene = replace(gene,"]" => ")")
        elseif p1 == "("
            gene = replace(gene,"(" => "]")
            gene = replace(gene,")" => "]")
        end
    end
    return gene
end

"""
print_ll(param,data,model,message="initial ll:")

compute and print initial loglikelihood
"""
function print_ll(param,data,model,message)
    ll,_ = loglikelihood(param,data,model)
    println(message,ll)
end
function print_ll(data,model,message="initial ll: ")
    ll,_ = loglikelihood(get_param(model),data,model)
    println(message,ll)
end

"""
printinfo(gene,G,cond,datafolder,infolder,resultfolder,maxtime)

print out run information
"""
function printinfo(gene,G,cond,datafolder,infolder,resultfolder,maxtime)
    println("Gene: ",gene," G: ",G," Treatment:  ",cond)
    println("data: ",datafolder)
    println("in: ", infolder," out: ",resultfolder)
    println("maxtime: ",maxtime)
end

"""
finalize(data,model,fit,stats,waic,temp,resultfolder,root)

write out run results and print out final loglikelihood and deviance
"""
function finalize(data,model,fit,stats,measures,temp,resultfolder,root)
    writefile = joinpath(root,resultfolder)
    writeall(writefile,fit,stats,measures,data,temp,model)
    println("final max ll: ",fit.llml)
    print_ll(vec(stats.meanparam),data,model,"mean ll: ")
    println("Mean fitted rates: ",stats.meanparam[:,1])
    println("Acceptance: ",fit.accept,"/",fit.total)
    println("Deviance: ",deviance(fit,data,model))
    println("rhat: ",maximum(measures.rhat))
end

#Prepare data structures
"""
    data_rna(gene::String,cond::String,datafolder,fish,label,root)
    data_rna(path,time,gene,time)
    data_rna(path,time,gene)

Load data structure
"""

function data_rna(gene::String,cond::String,datafolder::String,fish::Bool,label,root)
    if cond == "null"
        cond = ""
    end
    datafile = fish ? FISHpath(gene,cond,datafolder,root) : scRNApath(gene,cond,datafolder,root)
    data_rna(datafile,label,gene,fish)
end
function data_rna(gene::String,cond::Array,datafolder::String,fish::Bool,label,root)
    datafile = Array{String,1}(undef,length(cond))
    for i in eachindex(cond)
        datafile[i] = fish ? FISHpath(gene,cond[i],datafolder,root) : scRNApath(gene,cond[i],datafolder,root)
    end
    data_rna(datafile,label,gene,fish)
end
function data_rna(gene::String,cond::Array,datafolder::Array,fish::Bool,label,root)
    datafile = Array{String,1}(undef,length(cond))
    for i in eachindex(cond)
        datafile[i] = fish ? FISHpath(gene,cond[i],datafolder,root) : scRNApath(gene,cond[i],datafolder,root)
    end
    println(datafile)
    data_rna(datafile,label,gene,fish)
end
function data_rna(gene::String,cond::String,datafolder::String,fish::Bool,label,root,sets::Vector,time::Vector)
    if cond == "null"
        cond = ""
    end
    datafile =[]
    for set in sets
        folder = joinpath(datafolder,set)
        path = fish ? FISHpath(gene,cond,datafolder,root) : scRNApath(gene,cond,datafolder,root)
        datafile = vcat(datafile,path)
    end
    data_rna(datafile,label,times,gene,false)
end
function data_rna(path,name,time,gene::String,fish::Bool)
    len,h = histograms_rna(path,gene,fish)
    TransientRNAData(name,gene,len,time,h)
end
function data_rna(path,name,gene::String,fish::Bool)
    len,h = histograms_rna(path,gene,fish)
    RNAData(name,gene,len,h)
end
function data_rna(path,name,gene::String,fish::Array{Bool,1})
    len,h = histograms_rna(path,gene,fish)
    RNAMixedData(name,gene,len,fish,h)
end
"""
histograms_rna(path,gene)

prepare mRNA histograms for gene given data folder path
"""
function histograms_rna(path::Array,gene::String,fish::Bool)
    n = length(path)
    h = Array{Array,1}(undef,n)
    lengths = Array{Int,1}(undef,n)
    for i in eachindex(path)
        lengths[i], h[i] = histograms_rna(path[i],gene,fish)
    end
    return lengths,h
end
function histograms_rna(path::String,gene::String,fish::Bool)
    if fish
        h = read_fish(path,gene,.98)
    else
        h = read_scrna(path,.99)
        if length(h) == 0
            throw("data not found")
        end
    end
    return length(h),h
end
function histograms_rna(path::Array,gene::String,fish::Array{Bool,1})
    n = length(path)
    h = Array{Array,1}(undef,n)
    lengths = Array{Int,1}(undef,n)
    for i in eachindex(path)
        lengths[i], h[i] = histograms_rna(path[i],gene,fish[i])
    end
    return lengths,h
end

# Prepare model structures
"""
model_rna(gene,cell,G,fish,fittedparam,fixedeffects,inlabel,infolder,nsets,root,data,verbose=true)
model_rna(r::Vector,d::Vector,G::Int,nalleles::Int,propcv,fittedparam,fixedeffects,fish::Bool,method=0)
model_rna(r,G,nalleles,nsets,propcv,fittedparam,decayprior,yieldprior,method)
model_rna(r,G,nalleles,nsets,propcv,fittedparam,decayprior,method)
model_rna(r,G,nalleles,nsets,propcv,fittedparam,decayprior,noisepriors,method)

make model structure

"""
function model_rna(gene::String,cell::String,G::Int,cv,fittedparam,fixedeffects,inlabel,infolder,nsets,root,data,yieldprior,verbose=true)
    decayrate = get_decay(gene,cell,root)
    nalleles = alleles(gene,cell,root)
    if verbose
        println("alleles: ",nalleles)
        if decayrate < 0
            throw("decayrate < 0")
        else
            println("decay rate: ",decayrate)
        end
    end
    if cv <= 0
        cv = getcv(gene,G,nalleles,fittedparam,inlabel,infolder,root,verbose)
    end
    if G == 1
        ejectrate = mean_histogram(data.histRNA) * decayrate/nalleles * yieldprior
    else
        ejectrate = yieldprior
    end
    r = getr(gene,G,nalleles,decayrate,ejectrate,inlabel,infolder,nsets,root,verbose)
    model = model_rna(r,G,nalleles,nsets,cv,fittedparam,fixedeffects,decayrate,ejectrate)
    return model
end
# function model_rna2(gene::String,cell::String,G::Int,fish::Bool,cv,fittedparam,fixedeffects,inlabel,infolder,nsets,root,data,yield,ejectrate,verbose=true)
#     decayrate = get_decay(gene,cell,root)
#     nalleles = alleles(gene,cell,root)
#     if verbose
#         println("alleles: ",nalleles)
#         if decayrate < 0
#             throw("decayrate < 0")
#         else
#             println("decay rate: ",decayrate)
#         end
#     end
#     if cv <= 0
#         cv = getcv(gene,G,nalleles,fittedparam,inlabel,infolder,root,verbose)
#     end
#     if G == 1
#         ejectrate = mean_histogram(data.histRNA) * decayrate/nalleles
#     end
#     r = getr2(gene,G,nalleles,decayrate,ejectrate,yield,inlabel,infolder,nsets,root,fish,verbose)
#     model = model_rna2(r,G,nalleles,nsets,cv,fittedparam,fixedeffects,decayrate,ejectrate,yield,fish)
#     return model
# end

function model_rna(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam,fixedeffects,decayprior,ejectprior)
    d = prior_rna(r,G,nsets,fittedparam,decayprior,ejectprior)
    model_rna(r::Vector,d,G::Int,nalleles,propcv,fittedparam,fixedeffects)
end
function model_rna(r::Vector,d,G::Int,nalleles,propcv,fittedparam,fixedeffects,method=0)
    if length(fixedeffects) > 0
        model = GMfixedeffectsmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,fixedeffects,method)
    else
        model = GMmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,method)
    end
    return model
end
# function model_rna2(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam,fixedeffects,decayprior,ejectprior,yieldprior,fish::Bool,method=0)
#     if fish
#         d = prior_rna(r,G,nsets,fittedparam,decayprior,ejectprior)
#     else
#         d = prior_rna2(r,G,nsets,fittedparam,decayprior,ejectprior,yieldprior)
#     end
#     model_rna2(r,d,G,nalleles,propcv,fittedparam,fixedeffects,fish,method)
# end
# function model_rna2(r::Vector,d,G::Int,nalleles,propcv,fittedparam,fixedeffects,fish,method)
#         if length(fixedeffects) > 0
#             model = GMfixedeffectsmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,fixedeffects,method)
#         else
#             model = GMmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,method)
#         end
#     return model
# end

function model_rna(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam::Array,decayprior::Float64,method::Int)
    # propcv = proposal_cv_rna(propcv,fittedparam)
    d = prior_rna(r,G,nsets,fittedparam,decayprior,LogNormal)
    GMmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,method)
end
function model_rna(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam::Array,decayprior::Float64,yieldprior::Float64,method::Int)
    # propcv = proposal_cv_rna(propcv,fittedparam)
    d = prior_rna(r,G,nsets,fittedparam,decayprior,yieldprior)
    GMlossmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,method)
end
function model_rna(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam::Array,fixedeffects::Tuple,decayprior::Float64,method::Int)
    # propcv = proposal_cv_rna(propcv,fittedparam)
    d = prior_rna(r,G,nsets,fittedparam,decayprior)
    GMfixedeffectsodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,fixedeffects,method)
end
function model_rna(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam::Array,fixedeffects::Tuple,decayprior::Float64,yieldprior::Float64,method::Int)
    # propcv = proposal_cv_rna(propcv,fittedparam)
    d = prior_rna(r,G,nsets,fittedparam,decayprior,yieldprior)
    GMfixedeffectslossmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,fixedeffects,method)
end
function model_rna(r::Vector,G::Int,nalleles::Int,propcv,fittedparam::Array,decayprior::Float64,noisepriors::Array,method::Int)
    # propcv = proposal_cv_rna(propcv,fittedparam)
    d = prior_rna(r,G::Int,1,fittedparam,decayprior,noisepriors)
    if method == 1
        GMrescaledmodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,method)
    else
        GMmultimodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),typeof(method)}(G,nalleles,r,d,propcv,fittedparam,method)
    end
end
function model_delay_rna(r::Vector,G::Int,nalleles::Int,nsets::Int,propcv,fittedparam::Array,decayprior,delayprior)
    # propcv = proposal_cv_rna(propcv,fittedparam)
    d = prior_rna(r,G,nsets,fittedparam,decayprior,delayprior)
    GMdelaymodel{typeof(r),typeof(d),typeof(propcv),typeof(fittedparam),Int64}(G,nalleles,r,d,propcv,fittedparam,1)
end

"""
transient_rna(nchains,gene::String,fittedparam,cond::Vector,G::Int,maxtime::Float64,infolder::String,resultfolder,datafolder,inlabel,label,nsets,runcycle::Bool=false,samplesteps::Int=40000,warmupsteps=20000,annealsteps=100000,temp=1.,tempanneal=100.,root = "/home/carsonc/scrna/")

make structures for transient model fit
"""
function transient_rna(nchains,gene::String,cell,fittedparam,cond::Vector,G::Int,maxtime::Float64,infolder::String,resultfolder,datafolder,inlabel,label,nsets,runcycle::Bool=false,samplesteps::Int=40000,warmupsteps=20000,annealsteps=100000,temp=1.,tempanneal=100.,root = "/home/carsonc/scrna/")
    data = make_data(gene,cond,G,datafolder,label,nsets,root)
    model = make_model(gene,cell,G,fittedparam,inlabel,infolder,nsets,root,data)
    param,_ = initial_proposal(model)
    return param, data, model
end

"""
getr(gene,G,nalleles,decayrate,ejectrate,inlabel,infolder,nsets::Int,root,verbose)

"""
function getr(gene,G,nalleles,decayrate,ejectrate,inlabel,infolder,nsets::Int,root,verbose)
    r = getr(gene,G,nalleles,inlabel,infolder,root,verbose)
    if ~isnothing(r)
        if length(r) == 2*G*nsets + 1
            for n in nsets
                r[2*G*n-1] *= clamp(r[2*G*nsets + 1],eps(Float64),1-eps(Float64))
            end
            r = r[1:2*G*nsets]
        end
        if length(r) == 2*G*nsets
            if verbose
                println("init rates: ",r)
            end
            return r
        end
    end
    println("No r")
    setr(G,nsets,decayrate,ejectrate)
end
# function getr2(gene,G,nalleles,decayrate,ejectrate,yield,inlabel,infolder,nsets::Int,root,fish::Bool,verbose)
#     r = getr(gene,G,nalleles,inlabel,infolder,root,verbose)
#     if ~isnothing(r)
#         if length(r) == 2*G*nsets + 1
#             for i in 1:nsets
#                 r[2*G*i-1] *= clamp(r[2*G*nsets + 1],eps(Float64),1-eps(Float64))
#             end
#             r = r[1:2*G*nsets]
#         end
#         if length(r) == 2*G*nsets
#             if verbose
#                 println(r)
#             end
#             return r
#         end
#     end
#     println("No r")
#     setr2(G,nsets,decayrate,ejectrate,yield,fish)
# end
function getr(gene,G,nalleles,inlabel,infolder,root,verbose)
    ratefile = path_Gmodel("rates",gene,G,nalleles,inlabel,infolder,root)
    if verbose
        println("rate file: ",ratefile)
    end
    if isfile(ratefile)
        r = readrates(ratefile,2)
    else
        return nothing
    end
end
function getcv(gene,G,nalleles,fittedparam,inlabel,infolder,root,verbose = true)
    paramfile = path_Gmodel("param-stats",gene,G,nalleles,inlabel,infolder,root)
    if isfile(paramfile)
        cv = read_covlogparam(paramfile)
        cv = float.(cv)
        if ~ isposdef(cv) || size(cv)[1] != length(fittedparam)
            cv = .02
        end
    else
        cv = .02
    end
    if verbose
        println("cv: ",cv)
    end
    return cv
end
"""
    get_decay(gene::String,cell::String,root::String,col::Int=2)
    get_decay(gene::String,path::String,col::Int)

    Get decay rate for gene and cell

"""
function get_decay(gene::String,cell::String,root::String,col::Int=2)
    path = get_file(root,"data/halflives",cell,"csv")
    if isnothing(path)
        println(gene," has no decay time")
        return -1.
    else
        get_decay(gene,path,col)
    end
end
function get_decay(gene::String,path::String,col::Int)
    a = nothing
    in = readdlm(path,',')
    ind = findfirst(in[:,1] .== gene)
    if ~isnothing(ind)
        a = in[ind,col]
    end
    get_decay(a,gene)
end
function get_decay(a,gene::String)
    if typeof(a) <: Number
        return get_decay(a)
    else
        println(gene," has no decay time")
        return -1.
    end
end
get_decay(a::Float64) = log(2)/a/60.

"""
    alleles(gene::String,cell::String,root::String,col::Int=3)
    alleles(gene::String,path::String,col::Int=3)

    Get allele number for gene and cell
"""
function alleles(gene::String,cell::String,root::String,col::Int=3)
    path = get_file(root,"data/alleles",cell,"csv")
    if isnothing(path)
        return 2
    else
        alleles(gene,path,col)
    end
end

function alleles(gene::String,path::String,col::Int=3)
    a = nothing
    in,h = readdlm(path,',',header=true)
    ind = findfirst(in[:,1] .== gene)
    if isnothing(ind)
        return 2
    else
        a = in[ind,col]
        return isnothing(a) ? 2 : Int(a)
    end
end

"""
prior_rna(r::Vector,G::Int,nsets::Int,propcv,fittedparam::Array,decayprior,yieldprior)
prior_rna(r::Vector,G::Int,nsets::Int,fittedparam::Array,decayprior::Float64,noisepriors::Array)

prepare prior distribution
r[mod(1:2*G,nsets)] = rates for each model set  (i.e. rates for each set are stacked)
r[2*G*nsets + 1] == yield factor (i.e remaining after loss due to technical noise)

prior for multithresholded smFISH data
r[1:2G] = model rates
r[2G+1] = additive noise mean
r[2G + 1 + 1:length(noisepriors)] = remaining fraction after thresholding (i.e. yield)
"""
# function prior_rna(r::Vector,G::Int,nsets::Int,fittedparam::Array,decayprior::Float64,ejectprior::Float64,f=LogNormal)
#         ind = 2*G * nsets + 1
#         if length(r) == ind
#             rm,rcv = setrate(G,nsets,decayprior,ejectprior,yieldprior)
#             if in(ind,fittedparam)
#                 return distributionBeta_array(rm[fittedparam],rcv[fittedparam],findfirst(ind.==fittedparam),f)
#             else
#                 return distribution_array(rm[fittedparam],rcv[fittedparam],f)
#             end
#         else
#             throw("rates have wrong length")
#         end
# end
# function prior_rna2(r::Vector,G::Int,nsets::Int,fittedparam::Array,decayprior::Float64,ejectprior::Float64,yieldprior::Float64,f=LogNormal)
#         ind = 2*G * nsets
#         if length(r) == ind
#             rm,rcv = setrate2(G,nsets,decayprior,ejectprior,yieldprior)
#             return distribution_array(rm[fittedparam],rcv[fittedparam],f)
#         else
#             throw("rates have wrong length")
#         end
# end
function prior_rna(r::Vector,G::Int,nsets::Int,fittedparam::Array,decayprior::Float64,ejectprior::Float64,f=LogNormal)
        if length(r) == 2*G*nsets
            rm,rcv = setrate(G,nsets,decayprior,ejectprior)
            return distribution_array(rm[fittedparam],rcv[fittedparam],f)
        else
            throw("rates have wrong length")
        end
end
function prior_rna(r::Vector,G::Int,nsets::Int,fittedparam::Array,decayprior::Float64,ejectprior::Float64,noisepriors::Array,f=LogNormal)
        if length(r) == 2*G * nsets + length(noisepriors)
            rm,rcv = setrate(G,nsets,decayprior,ejectprior,noisepriors)
            return distribution_array(rm[fittedparam],rcv[fittedparam],f)
        else
            throw("rates have wrong length")
        end
end

"""
    setrate(G::Int,nsets::Int,decayrate::Float64,ejectrate::Float64)
    setrate(G::Int,nsets::Int,decayrate,ejectrate)
    setrate(G::Int,nsets::Int,decayrate,ejectrate,yield)
    setrate(G::Int,nsets::Int,decayrate::Float64,ejectrate::Float64,noisepriors::Array)

    Set mean and cv of rates

"""

function setr(G,nsets,decayrate,ejectrate)
    r = setrate(G,nsets,decayrate,ejectrate)[1]
    println("init rates: ",r)
    return r
end
# function setr2(G,nsets,decayrate,ejectrate,yield,fish)
#     if fish
#         r = setrate(G,nsets,decayrate,ejectrate)[1]
#     else
#         r = setrate2(G,nsets,decayrate,ejectrate,yield)[1]
#     end
#     println(r)
#     return r
# end
function setrate(G::Int,nsets::Int,decayrate::Float64,ejectrate::Float64)
    rm = Vector{Float64}(undef,0)
    rc = similar(rm)
    for i in 1:nsets
        rm = vcat(rm,[repeat([0.01;0.1],(G-1));ejectrate;decayrate])
        rc = vcat(rc,[.4*ones(2*(G-1));1.;0.01])
    end
    return rm,rc
end
# function setrate(G::Int,nsets::Int,decayrate::Float64,ejectrate::Float64,yield::Float64)
#     rm,rcv = setrate(G,nsets,decayrate,ejectrate)
#     return [rm;yield],[rcv;.25]
# end
# function setrate2(G::Int,nsets::Int,decayrate::Float64,ejectprior::Float64,yieldprior::Float64)
#     rm = Vector{Float64}(undef,0)
#     rc = similar(rm)
#     for i in 1:nsets
#         rm = vcat(rm,[repeat([0.01;0.1],(G-1));ejectprior*yieldprior;decayrate])
#         rc = vcat(rc,[.4*ones(2*(G-1));1.14;0.01])
#     end
#     return rm,rc
# end
function setrate(G::Int,nsets::Int,decayrate::Float64,ejectrate::Float64,noisepriors::Array)
    rm,rc = setrate(G,nsets,decayrate,ejectrate)
    for nm in noisepriors
        rm = vcat(rm,nm)
        rc = vcat(rc,.25)
    end
    return rm,rc
end

"""
proposal_cv_rna(cv)
proposal_cv_rna(propcv,fittedparam)

set proposal cv to a vector or matrix
"""
proposal_cv_rna(cv) = typeof(cv) == Float64 ? propcv*ones(length(fittedparam)) : propcv

proposal_cv_rna(propcv,fittedparam) = typeof(propcv) == Float64 ? propcv*ones(length(fittedparam)) : propcv

"""
fittedparam_rna(G,nsets,loss)
fittedparam_rna(G,nsets)

select all parameters to be fitted except decay rates
yieldfactor is last parameter in loss models
"""
function fittedparam_rna(G,nsets,loss)
    fp = fittedparam_rna(G,nsets)
    if loss == 1
        return vcat(fp,nsets*2*G+1)
    else
        return fp
    end
end
function fittedparam_rna(G,nsets)
    nrates = 2*G  # total number of rates in GM model
    k = nrates - 1  # adjust all rate parameters except decay time
    fp = Array{Int,1}(undef,nsets*k)
    for j in 0:nsets-1
        for i in 1:k
            fp[k*j + i] = nrates*j + i
        end
    end
    return fp
end

"""
rescale_rate_rna(r,G,decayrate::Float64)

set new decay rate and rescale
transition rates such that steady state distribution is the same
"""
function rescale_rate_rna(r,G,decayrate::Float64)
    rnew = copy(r)
    if mod(length(r),2*G) == 0
        rnew *= decayrate/r[2*G]
    else
        stride = fld(length(r),fld(length(r),2*G))
        for i in 0:stride:length(r)-1
            rnew[i+1:i+2*G] *= decayrate/r[2*G]
        end
    end
    return rnew
end

# Read in data and construct histograms
"""
read_scrna(filename::String,yield::Float64=.99,nhistmax::Int=1000)
Construct mRNA count per cell histogram array of a gene
"""
function read_scrna(filename::String,threshold::Float64=.99,nhistmax::Int=500)
    if isfile(filename) && filesize(filename) > 0
        x = readdlm(filename)[:,1]
        x = truncate_histogram(x,threshold,nhistmax)
        if x == 0
            dataFISH = Array{Int,1}(undef,0)
        else
            dataFISH = x
        end
        return dataFISH
    else
        # println("data file not found")
        return Array{Int,1}(undef,0)
    end
end

"""
read_fish(path,gene,threshold)
Read in FISH data from 7timepoint type folders

"""
function read_fish(path::String,cond::String,threshold::Float64=.98,maxlength = 1800)
    xr = zeros(maxlength)
    lx = 0
    for (root,dirs,files) in walkdir(path)
        for file in files
            target = joinpath(root, file)
            if occursin(cond,target) && occursin("cellular",target)
                # println(target)
                x1 = readdlm(target)[:,1]
                x1 = truncate_histogram(x1,threshold,maxlength)
                lx = length(x1)
                # println(lx)
                xr[1:min(lx,maxlength)] += x1[1:min(lx,maxlength)]
            end
        end
    end
    return truncate_histogram(xr,1.0,maxlength)
end

function read_fish(path1::String,cond1::String,path2::String,cond2::String,threshold::Float64=.98)
    x1 = read_fish(path1,cond1,threshold)
    x2 = read_fish(path2,cond2,threshold)
    combine_histogram(x1,x2)
end

# functions to build paths to RNA/FISH data and results
"""
scRNApath(gene::String,cond::String,datapath::String,root::String)

generate path to scRNA data for given gene and condition cond
    data files must have format genename_cond.txt
"""
function scRNApath(gene::String,cond::String,datapath::String,root::String)
    datapath = joinpath(root,datapath)
    if cond == ""
        joinpath(datapath,gene * ".txt")
    else
        joinpath(datapath,gene * "_" * cond * ".txt")
    end
end
scRNApath(gene,cond,datapath) = joinpath(datapath,gene * "_" * cond * ".txt")

"""
FISHpath(gene,cond,datapath,root)

generate path to FISH data
"""
# FISHpath(gene,cond,datapath,root) = joinpath(joinpath(joinpath(root,datapath),gene),cond)
FISHpath(gene,cond,datapath,root) = joinpath(root,datapath,gene,cond)
FISHpath(gene,cond,datapath) = joinpath(datapath,gene,cond)

"""
ratepath_Gmodel(gene::String,cond::String,G::Int,nalleles::Int,label,folder,root)

"""
function ratepath_Gmodel(gene::String,cond::String,G::Int,nalleles::Int,label,folder,root)
    path_Gmodel("rates",gene,G,nalleles,label * "-" * cond,folder,root)
end

"""
path_Gmodel(type,gene::String,G::Int,nalleles::Int,label::String,folder,root)
"""
function path_Gmodel(type,gene::String,G::Int,nalleles::Int,label::String,folder,root)
    filelabel = label  * "_" * gene *  "_" * "$G" * "_" * "$nalleles" * ".txt"
    ratefile = type * "_" * filelabel
    joinpath(root,folder,ratefile)
end

"""
stats_rna(genes::Vector,conds,datapath,threshold=.99)
"""
function stats_rna(genes::Vector,conds,datapath,threshold=.99)
    g = Vector{String}(undef,0)
    z = Vector{Float64}(undef,0)
    h1 = Vector{Float64}(undef,0)
    h2 = Vector{Float64}(undef,0)
    for gene in genes
        t,m1,m2 = expression_rna(gene,conds,datapath,threshold)
        push!(g,gene)
        push!(z,t)
        push!(h1,m2)
        push!(h2,m1)
    end
    return g,z,h1,h2
end

"""
expression_rna(gene,conds::Vector,folder::String,threshold=.99)
"""
function expression_rna(gene,conds::Vector,folder::String,threshold=.99)
    h = Vector{Vector}(undef,2)
    for i in eachindex(conds)
        datapath = scRNApath(gene,conds[i],folder)
        h[i] = read_scrna(datapath,threshold)
    end
    if length(h[1]) > 0 && length(h[2]) > 0
        return log_2sample(h[1],h[2]), mean_histogram(h[1]), mean_histogram(h[2])
    else
        return 0,0,0
    end
end

"""
expression_rna(gene,conds::Vector,folder::Vector,threshold=.99)
"""
function expression_rna(gene,conds::Vector,folder::Vector,threshold=.99)
    h = Vector{Vector}(undef,2)
    for i in eachindex(conds)
        datapath = scRNApath(gene,conds[i],folder[i])
        h[i] = read_scrna(datapath,threshold)
    end
    if length(h[1]) > 0 && length(h[2]) > 0
        return log_2sample(h[1],h[2]), mean_histogram(h[1]), mean_histogram(h[2])
    else
        return 0,0,0
    end
end

"""
expression_rna(genes::Vector,cond::String,folder,threshold=.99)

"""
function expression_rna(genes::Vector,cond::String,folder,threshold=.99)
    h1 = Array{Any,2}(undef,0,2)
    for gene in genes
        datapath = scRNApath(gene,cond,folder)
        h = read_scrna(datapath,threshold)
        if length(h) > 0
            h1 = vcat(h1,[gene mean_histogram(h)])
        end
    end
    return h1
end

"""
rna_fish(gene,cond,fishfolder,rnafolder,yield,root)

output RNA histogram and downsampled FISH histogram with loss
"""
function rna_fish(gene,cond,fishfolder,rnafolder,yield,root)
    datarna = histograms_rna(scRNApath(gene,cond,rnafolder,root),gene,false)
    f=reduce_fish(gene,cond,datarna[1],fishfolder,yield,root)
    return datarna[2],f
end

"""
reduce_fish(gene,cond,nhist,fishfolder,yield,root)

sample fish histogram with loss (1-yield)
"""
function reduce_fish(gene,cond,nhist,fishfolder,yield,root)
    fish = histograms_rna(FISHpath(gene,cond,fishfolder,root),gene,true)
    technical_loss(fish[2],yield,nhist)
    fish[2]
end


mutable struct gpuCTPF <: GPUTopicModel
	K::Int
	M::Int
	V::Int
	U::Int
	N::Vector{Int}
	C::Vector{Int}
	R::Vector{Int}
	B::Int
	corp::Corpus
	batches::Vector{UnitRange{Int}}
	topics::VectorList{Int}
	scores::Matrix{Float32}
	libs::VectorList{Int}
	drecs::VectorList{Int}
	urecs::VectorList{Int}
	a::Float32
	b::Float32
	c::Float32
	d::Float32
	e::Float32
	f::Float32
	g::Float32
	h::Float32
	alef::Matrix{Float32}
	bet::Vector{Float32}
	gimel::VectorList{Float32}
	dalet::Vector{Float32}
	he::Matrix{Float32}
	vav::Vector{Float32}
	zayin::VectorList{Float32}
	het::Vector{Float32}
	phi::MatrixList{Float32}
	xi::MatrixList{Float32}
	newalef::Void
	newhe::Void
	Npsums::VectorList{Int}
	Jpsums::VectorList{Int}
	Rpsums::VectorList{Int}
	Ypsums::VectorList{Int}
	terms::VectorList{Int}
	counts::VectorList{Int}
	words::VectorList{Int}
	readers::VectorList{Int}
	ratings::VectorList{Int}
	views::VectorList{Int}
	device::cl.Device
	context::cl.Context
	queue::cl.CmdQueue
	Npsumsbuf::cl.Buffer{Int}
	Jpsumsbuf::cl.Buffer{Int}
	Rpsumsbuf::cl.Buffer{Int}
	Ypsumsbuf::cl.Buffer{Int}
	termsbuf::cl.Buffer{Int}
	countsbuf::cl.Buffer{Int}
	wordsbuf::cl.Buffer{Int}
	readersbuf::cl.Buffer{Int}
	ratingsbuf::cl.Buffer{Int}
	viewsbuf::cl.Buffer{Int}
	alefkern::cl.Kernel
	newalefkern::cl.Kernel
	betkern::cl.Kernel
	gimelkern::cl.Kernel
	daletkern::cl.Kernel
	hekern::cl.Kernel
	newhekern::cl.Kernel
	vavkern::cl.Kernel
	zayinkern::cl.Kernel
	hetkern::cl.Kernel
	phikern::cl.Kernel
	phinormkern::cl.Kernel
	xikern::cl.Kernel
	xinormkern::cl.Kernel
	alefbuf::cl.Buffer{Float32}
	newalefbuf::cl.Buffer{Float32}
	betbuf::cl.Buffer{Float32}
	gimelbuf::cl.Buffer{Float32}
	daletbuf::cl.Buffer{Float32}
	hebuf::cl.Buffer{Float32}
	newhebuf::cl.Buffer{Float32}
	vavbuf::cl.Buffer{Float32}
	zayinbuf::cl.Buffer{Float32}
	hetbuf::cl.Buffer{Float32}
	phibuf::cl.Buffer{Float32}
	xibuf::cl.Buffer{Float32}
	elbo::Float32
	newelbo::Float32

	function gpuCTPF(corp::Corpus, K::Integer, batchsize::Integer=length(corp), basemodel::Union{Void, BaseTopicModel}=nothing)
		@assert !isempty(corp)
		@assert all(ispositive.([K, batchsize]))
		checkcorp(corp)

		M, V, U = size(corp)
		N = [length(doc) for doc in corp]
		C = [size(doc) for doc in corp]
		R = [length(doc.readers) for doc in corp]

		batches = partition(1:M, batchsize)
		B = length(batches)

		@assert ispositive(U)
		@assert isequal(collect(1:U), sort(collect(keys(corp.users))))
		libs = [Int[] for _ in 1:U]
		for u in 1:U, d in 1:M
			u in corp[d].readers && push!(libs[u], d)
		end

		a, b, c, d, e, f, g, h = fill(0.1, 8)

		if isa(basemodel, Union{AbstractLDA, AbstractCTM})
			@assert isequal(size(basemodel.beta), (K, V))
			alef = exp.(basemodel.beta)
			topics = basemodel.topics
		elseif isa(basemodel, Union{AbstractfLDA, AbstractfCTM})
			@assert isequal(size(basemodel.fbeta), (K, V))
			alef = exp.(basemodel.fbeta)
			topics = basemodel.topics
		else
			alef = exp.(rand(Dirichlet(V, 1.0), K)' - 0.5)
			topics = [collect(1:V) for _ in 1:K]
		end

		bet = ones(K)
		gimel = [ones(K) for _ in 1:M]
		dalet = ones(K)
		he = ones(K, U)
		vav = ones(K)
		zayin = [ones(K) for _ in 1:M]
		het = ones(K)
		phi = [ones(K, N[d]) / K for d in batches[1]]
		xi = [ones(2K, R[d]) / 2K for d in batches[1]]

		model = new(K, M, V, U, N, C, R, B, copy(corp), batches, topics, zeros(M, U), libs, Vector[], Vector[], a, b, c, d, e, f, g, h, alef, bet, gimel, dalet, he, vav, zayin, het, phi, xi)
		fixmodel!(model, check=false)

		for (b, batch) in enumerate(batches)
			model.phi = [ones(K, N[d]) / K for d in batch]
			model.xi = [ones(2K, R[d]) / 2K for d in batch]
			updateNewELBO!(model, b)
		end
		model.phi = [ones(K, N[d]) / K for d in batches[1]]
		model.xi = [ones(2K, R[d]) / 2K for d in batches[1]]
		updateELBO!(model)
		return model
	end
end

gpuCTPF(corp::Corpus, K::Int, basemodel::Union{Void, BaseTopicModel}) = gpuCTPF(corp, K, length(corp), basemodel)

function Elogpya(model::gpuCTPF, d::Int, m::Int)
	x = 0
	readers, ratings = model.corp[d].readers, model.corp[d].ratings
	for (u, (re, ra)) in enumerate(zip(readers, ratings)), i in 1:model.K
		binom = Binomial(ra, model.xi[m][i,u])
		x += (ra * model.xi[m][i,u] * (digamma(model.gimel[d][i]) - log(model.dalet[i]) + digamma(model.he[i,re]) - log(model.vav[i]))
				- (model.gimel[d][i] / model.dalet[i]) * (model.he[i,re] / model.vav[i]) - sum([pdf(binom, y) * lgamma(y + 1) for y in 0:ra]))
	end
	return x
end

function Elogpyb(model::gpuCTPF, d::Int, m::Int)
	x = 0
	readers, ratings = model.corp[d].readers, model.corp[d].ratings
	for (u, (re, ra)) in enumerate(zip(readers, ratings)), i in 1:model.K
		binom = Binomial(ra, model.xi[m][model.K+i,u])
		x += (ra * model.xi[m][model.K+i,u] * (digamma(model.zayin[d][i]) - log(model.het[i]) + digamma(model.he[i,re]) - log(model.vav[i]))
				- (model.zayin[d][i] / model.het[i]) * (model.he[i,re] / model.vav[i]) - sum([pdf(binom, y) * lgamma(y + 1) for y in 0:ra]))
	end
	return x
end

function Elogpz(model::gpuCTPF, d::Int, m::Int)
	x = 0
	terms, counts = model.corp[d].terms, model.corp[d].counts
	for (n, (j, c)) in enumerate(zip(terms, counts)), i in 1:model.K
		binom = Binomial(c, model.phi[m][i,n])
		x += (c * model.phi[m][i,n] * (digamma(model.gimel[d][i]) - log(model.dalet[i]) + digamma(model.alef[i,j]) - log(model.bet[i]))
				- (model.gimel[d][i] / model.dalet[i]) * (model.alef[i,j] / model.bet[i]) - sum([pdf(binom, z) * lgamma(z + 1) for z in 0:c]))
	end
	return x
end

function Elogpbeta(model::gpuCTPF)
	x = model.V * model.K * (model.a * log(model.b) - lgamma(model.a))
	for j in 1:model.V, i in 1:model.K
		x += (model.a - 1) * (digamma(model.alef[i,j]) - log(model.bet[i])) - model.b * model.alef[i,j] / model.bet[i]
	end
	return x
end

function Elogptheta(model::gpuCTPF, d::Int)
	x = model.K * (model.c * log(model.d) - lgamma(model.c))
	for i in 1:model.K
		x += (model.c - 1) * (digamma(model.gimel[d][i]) - log(model.dalet[i])) - model.d * model.gimel[d][i] / model.dalet[i]
	end
	return x
end

function Elogpeta(model::gpuCTPF)
	x = model.U * model.K * (model.e * log(model.f) - lgamma(model.e))
	for u in 1:model.U, i in 1:model.K
		x += (model.e - 1) * (digamma(model.he[i,u]) - log(model.vav[i])) - model.f * model.he[i,u] / model.vav[i]
	end
	return x
end

function Elogpepsilon(model::gpuCTPF, d::Int)
	x = model.K * (model.g * log(model.h) - lgamma(model.g))
	for i in 1:model.K
		x += (model.g - 1) * (digamma(model.zayin[d][i]) - log(model.het[i])) - model.h * model.zayin[d][i] / model.het[i]
	end
	return x
end

function Elogqy(model::gpuCTPF, d::Int, m::Int)
	x = 0
	for (u, ra) in enumerate(model.corp[d].ratings)
		x -= entropy(Multinomial(ra, model.xi[m][:,u]))
	end
	return x
end

function Elogqz(model::gpuCTPF, d::Int, m::Int)
	x = 0
	for (n, c) in enumerate(model.corp[d].counts)
		x -= entropy(Multinomial(c, model.phi[m][:,n]))
	end
	return x
end

function Elogqbeta(model::gpuCTPF)
	x = 0
	for j in 1:model.V, i in 1:model.K
		x -= entropy(Gamma(model.alef[i,j], 1 / model.bet[i]))
	end
	return x
end

function Elogqtheta(model::gpuCTPF, d::Int)
	x = 0
	for i in 1:model.K
		x -= entropy(Gamma(model.gimel[d][i], 1 / model.dalet[i]))
	end
	return x
end

function Elogqeta(model::gpuCTPF)
	x = 0
	for u in 1:model.U, i in 1:model.K
		x -= entropy(Gamma(model.he[i,u], 1 / model.vav[i]))
	end
	return x
end

function Elogqepsilon(model::gpuCTPF, d::Int)
	x = 0
	for i in 1:model.K
		x -= entropy(Gamma(model.zayin[d][i], 1 / model.het[i]))
	end
	return x
end

function updateELBO!(model::gpuCTPF)
	model.elbo = model.newelbo + Elogpbeta(model) + Elogpeta(model) - Elogqbeta(model) - Elogqeta(model)
	model.newelbo = 0
	return model.elbo
end

function updateNewELBO!(model::gpuCTPF, b::Int)
	batch = model.batches[b]
	for (m, d) in enumerate(batch)
		model.newelbo += (Elogpya(model, d, m)
						+ Elogpyb(model, d, m)
						+ Elogpz(model, d, m)
						+ Elogptheta(model, d)
						+ Elogpepsilon(model, d)
						- Elogqy(model, d, m)
						- Elogqz(model, d, m)
						- Elogqtheta(model, d)
						- Elogqepsilon(model, d))
	end
end

const CTPF_ALEF_cpp =
"""
kernel void
updateAlef(long K,
			float a,
			global float *newalef,
			global float *alef)

			{
			long i = get_global_id(0);
			long j = get_global_id(1);

			alef[K * j + i] = newalef[K * j + i];
			newalef[K * j + i] = a;
			}
			"""

function updateAlef!(model::gpuCTPF)
	model.queue(model.alefkern, (model.K, model.V), nothing, model.K, model.a, model.newalefbuf, model.alefbuf)
end

const CTPF_NEWALEF_cpp =
"""
kernel void
updateNewalef(long K,
				const global long *Jpsums,
				const global long *counts,
				const global long *words,
				const global float *phi,
				global float *newalef)

				{
				long i = get_global_id(0);
				long j = get_global_id(1);

				float acc = 0.0f;

				for (long w=Jpsums[j]; w<Jpsums[j+1]; w++)
					acc += counts[words[w]] * phi[K * words[w] + i];

				newalef[K * j + i] += acc;
				}
				"""

function updateNewalef!(model::gpuCTPF)
	model.queue(model.newalefkern, (model.K, model.V), nothing, model.K, model.Jpsumsbuf, model.countsbuf, model.wordsbuf, model.phibuf, model.newalefbuf)
end

const CTPF_BET_cpp =
"""
kernel void
updateBet(long K,
			long M,
			float b,
			const global float *alef,
			const global float *gimel,
			const global float *dalet,
			global float *bet)

			{
			long i = get_global_id(0);

			float acc = 0.0f;

			for (long d=0; d<M; d++)
				acc += gimel[K * d + i];

			bet[i] = b + acc / dalet[i];
			}
			"""

function updateBet!(model::gpuCTPF)
	model.queue(model.betkern, model.K, nothing, model.K, model.M, model.b, model.alefbuf, model.gimelbuf, model.daletbuf, model.betbuf)
end

const CTPF_GIMEL_cpp =
"""
kernel void
updateGimel(long F,
			long K,
			float c,
			const global long *Npsums,
			const global long *Rpsums,
			const global long *counts,
			const global long *ratings,
			const global float *phi,
			const global float *xi,
			global float *gimel)

			{
			long i = get_global_id(0);
			long d = get_global_id(1);

			float accphi = 0.0f;
			float accxi = 0.0f;

			for (long n=Npsums[d]; n<Npsums[d+1]; n++)
				accphi += phi[K * n + i] * counts[n];

			for (long r=Rpsums[d]; r<Rpsums[d+1]; r++)
				accxi += xi[2 * K * r + i] * ratings[r];

			gimel[K * (F + d) + i] = c + accphi + accxi;
			}
			"""

function updateGimel!(model::gpuCTPF, b::Int)
	batch = model.batches[b]
	model.queue(model.gimelkern, (model.K, length(batch)), nothing, batch[1] - 1, model.K, model.c, model.Npsumsbuf, model.Rpsumsbuf, model.countsbuf, model.ratingsbuf, model.phibuf, model.xibuf, model.gimelbuf)
end

const CTPF_DALET_cpp =
"""
kernel void
updateDalet(long K,
			long V,
			long U,
			float d,
			const global float *alef,
			const global float *bet,
			const global float *he,
			const global float *vav,
			global float *dalet)

			{
			long i = get_global_id(0);

			float accalef = 0.0f;
			float acche = 0.0f;

			for (long j=0; j<V; j++)
				accalef += alef[K * j + i];

			for (long u=0; u<U; u++)
				acche += he[K * u + i];

			dalet[i] = d + accalef / bet[i] + acche / vav[i];
			}
			"""

function updateDalet!(model::gpuCTPF)
	model.queue(model.daletkern, model.K, nothing, model.K, model.V, model.U, model.d, model.alefbuf, model.betbuf, model.hebuf, model.vavbuf, model.daletbuf)
end

const CTPF_HE_cpp =
"""
kernel void
updateHe(long K,
			float e,
			global float *newhe,
			global float *he)

			{
			long i = get_global_id(0);
			long u = get_global_id(1);

			he[K * u + i] = newhe[K * u + i];
			newhe[K * u + i] = e;
			}
			"""

function updateHe!(model::gpuCTPF)
	model.queue(model.hekern, (model.K, model.U), nothing, model.K, model.e, model.newhebuf, model.hebuf)
end

const CTPF_NEWHE_cpp =
"""
kernel void
updateNewhe(long K,
			const global long *Ypsums,
			const global long *ratings,
			const global long *views,
			const global float *xi,
			global float *newhe)

			{
			long i = get_global_id(0);
			long u = get_global_id(1);

			float acc = 0.0f;

			for (long r=Ypsums[u]; r<Ypsums[u+1]; r++)
				acc += ratings[views[r]] * (xi[2 * K * views[r] + i] + xi[K * (2 * views[r] + 1) + i]);

			newhe[K * u + i] += acc;
			}
			"""

function updateNewhe!(model::gpuCTPF)
	model.queue(model.newhekern, (model.K, model.U), nothing, model.K, model.Ypsumsbuf, model.ratingsbuf, model.viewsbuf, model.xibuf, model.newhebuf)
end

const CTPF_VAV_cpp =
"""
kernel void
updateVav(long K,
			long M,
			float f,
			const global float *gimel,
			const global float *dalet,
			const global float *zayin,
			const global float *het,
			global float *vav)

			{
			long i = get_global_id(0);

			float accgimel = 0.0f;
			float acczayin = 0.0f;

			for (long d=0; d<M; d++)
			{
				accgimel += gimel[K * d + i];
				acczayin += zayin[K * d + i];
			}

			vav[i] = f + accgimel / dalet[i] + acczayin / het[i];
			}
			"""

function updateVav!(model::gpuCTPF)
	model.queue(model.vavkern, model.K, nothing, model.K, model.M, model.f, model.gimelbuf, model.daletbuf, model.zayinbuf, model.hetbuf, model.vavbuf)
end

const CTPF_ZAYIN_cpp =
"""
kernel void
updateZayin(long F,
			long K,
			float g,
			const global long *Rpsums,
			const global long *ratings,
			const global float *xi,
			global float *zayin)

			{
			long i = get_global_id(0);
			long d = get_global_id(1);

			float acc = 0.0f;

			for (long r=Rpsums[d]; r<Rpsums[d+1]; r++)
				acc += xi[K * (2 * r + 1) + i] * ratings[r];

			zayin[K * (F + d) + i] = g + acc;
			}
			"""

function updateZayin!(model::gpuCTPF, b::Int)
	batch = model.batches[b]
	model.queue(model.zayinkern, (model.K, length(batch)), nothing, batch[1] - 1, model.K, model.g, model.Rpsumsbuf, model.ratingsbuf, model.xibuf, model.zayinbuf)
end

const CTPF_HET_cpp =
"""
kernel void
updateHet(long K,
			long U,
			float h,
			const global float *he,
			const global float *vav,
			global float *het)

			{
			long i = get_global_id(0);

			float acc = 0.0f;

			for (long u=0; u<U; u++)
				acc += he[K * u + i];

			het[i] = h + acc / vav[i];
			}
			"""

function updateHet!(model::gpuCTPF)
	model.queue(model.hetkern, model.K, nothing, model.K, model.U, model.h, model.hebuf, model.vavbuf, model.hetbuf)
end

const CTPF_PHI_cpp =
"""
$(DIGAMMA_cpp)

kernel void
updatePhi(long F,
			long K,
			const global long *Npsums,
			const global long *terms,
			const global float *alef,
			const global float *bet,
			const global float *gimel,
			const global float *dalet,
			global float *phi)

			{

			long i = get_global_id(0);
			long d = get_global_id(1);

			float gdb = digamma(gimel[K * (F + d) + i]) - log(dalet[i]) - log(bet[i]);

			for (long n=Npsums[d]; n<Npsums[d+1]; n++)
				phi[K * n + i] = exp(gdb + digamma(alef[K * terms[n] + i]));
			}
			"""

const CTPF_PHI_NORM_cpp =
"""
kernel void
normalizePhi(long K,
				global float *phi)

				{
				long dn = get_global_id(0);

				float normalizer = 0.0f;

				for (long i=0; i<K; i++)
					normalizer += phi[K * dn + i];

				for (long i=0; i<K; i++)
					phi[K * dn + i] /= normalizer;
				}
				"""

function updatePhi!(model::gpuCTPF, b::Int)
	batch = model.batches[b]
	model.queue(model.phikern, (model.K, length(batch)), nothing, batch[1] - 1, model.K, model.Npsumsbuf, model.termsbuf, model.alefbuf, model.betbuf, model.gimelbuf, model.daletbuf, model.phibuf)
	model.queue(model.phinormkern, sum(model.N[batch]), nothing, model.K, model.phibuf)
end

const CTPF_XI_cpp =
"""
$(DIGAMMA_cpp)

kernel void
updateXi(long F,
			long K,
			const global long *Rpsums,
			const global long *readers,
			const global float *bet,
			const global float *gimel,
			const global float *dalet,
			const global float *he,
			const global float *vav,
			const global float *zayin,
			const global float *het,
			global float *xi)

			{
			long i = get_global_id(0);
			long d = get_global_id(1);

			float gdv = digamma(gimel[K * (F + d) + i]) - log(dalet[i]) - log(bet[i]);
			float zhv = digamma(zayin[K * (F + d) + i]) - log(het[i]) - log(vav[i]);

			for (long r=Rpsums[d]; r<Rpsums[d+1]; r++)
			{
				xi[2 * K * r + i] = exp(gdv + digamma(he[K * readers[r] + i]));
				xi[K * (2 * r + 1) + i] = exp(zhv + digamma(he[K * readers[r] + i]));
			}
			}
			"""

const CTPF_XI_NORM_cpp =
"""
kernel void
normalizeXi(long K,
			global float *xi)

				{
				long dn = get_global_id(0);

				float normalizer = 0.0f;

				for (long i=0; i<2*K; i++)
					normalizer += xi[2 * K * dn + i];

				for (long i=0; i<2*K; i++)
					xi[2 * K * dn + i] /= normalizer;
				}
				"""

function updateXi!(model::gpuCTPF, b::Int)
	batch = model.batches[b]
	model.queue(model.xikern, (model.K, length(batch)), nothing, batch[1] - 1, model.K, model.Rpsumsbuf, model.readersbuf, model.betbuf, model.gimelbuf, model.daletbuf, model.hebuf, model.vavbuf, model.zayinbuf, model.hetbuf, model.xibuf)
	model.queue(model.xinormkern, sum(model.R[batch]), nothing, model.K, model.xibuf)
end

function train!(model::gpuCTPF; iter::Int=150, tol::Real=1.0, viter::Int=10, vtol::Real=1/model.K^2, chkelbo::Int=1)
	@assert all(.!isnegative.([tol, vtol]))
	@assert all(ispositive.([iter, viter, chkelbo]))
	lowVRAM = model.B > 1

	for k in 1:iter
		chk = (k % chkelbo == 0)
		for b in 1:model.B
			for _ in 1:viter
				oldgimel = @host b model.gimelbuf
				updatePhi!(model, b)
				updateXi!(model, b)
				updateGimel!(model, b)
				updateZayin!(model, b)
				gimel = @host b model.gimelbuf
				if sum([norm(diff) for diff in oldgimel - gimel]) < length(model.batches[b]) * vtol
					break
				end
			end
			updateNewalef!(model)
			updateNewhe!(model)
			if chk
				updateHost!(model, b)
				updateNewELBO!(model, b)
			end
			lowVRAM && updateBuf!(model, b)
		end
		if checkELBO!(model, k, chk, tol)
			updateDalet!(model)
			updateHet!(model)
			updateAlef!(model)
			updateBet!(model)
			updateHe!(model)
			updateVav!(model)
			break
		end
		updateDalet!(model)
		updateHet!(model)
		updateAlef!(model)
		updateBet!(model)
		updateHe!(model)
		updateVav!(model)
	end
	updateHost!(model, 1)
	Ebeta = model.alef ./ model.bet
	model.topics = [reverse(sortperm(vec(Ebeta[i,:]))) for i in 1:model.K]

	Eeta = model.he ./ model.vav
	for d in 1:model.M
		Etheta = model.gimel[d] ./ model.dalet
		Eepsilon = model.zayin[d] ./ model.het
		model.scores[d,:] = sum(Eeta .* (Etheta + Eepsilon), 1)
	end

	model.drecs = Vector{Int}[]
	for d in 1:model.M
		nr = setdiff(keys(model.corp.users), model.corp[d].readers)
		push!(model.drecs, nr[reverse(sortperm(vec(model.scores[d,nr])))])
	end

	model.urecs = Vector{Int}[]
	for u in 1:model.U
		ur = filter(d -> !(u in model.corp[d].readers), collect(1:model.M))
		push!(model.urecs, ur[reverse(sortperm(model.scores[ur,u]))])
	end
	nothing
end


# first pass: read whole file, save in binary format, count IDs for
# collation, store AM spell strings in the order of appearance

using AMDB: data_path, all_data_files, data_file
using ByteParsers:
    Line, DateYYYYMMDD, PosInteger, ViewBytes, parsedtype, ByteVector, parsenext,
    isparsed, Skip
using DocStringExtensions: SIGNATURES
using CodecZlib: GzipDecompressorStream
using WallTimeProgress: WallTimeTracker, increment!

"""
A parser for just the columns we need for labor market transitions.
"""
const lineparser = Line(PosInteger(Int64), # PENR
                        Skip(),            # start, skip
                        Skip(),            # end, skip
                        Skip(),            # firm, skip
                        DateYYYYMMDD(),    # STICHTAG
                        ViewBytes())       # AM

"""
   $SIGNATURES

Read stream from `io` line by line, extract

1. person id,

2. sample date,

3. labor market status

then reclassify the last one with `aggregator`, then return the assembled
histories as a dictionary of person id to 12-element vectors.

The latter contain the histories of individuals. The value `missing` is used
when no observation is found in the sample.
"""
function read_aggregated_AMs(io::IO, aggregator::Associative{ByteVector, T}, missing::T,
                             lineparser = lineparser,
                             tracker = WallTimeTracker(10_000_000;
                                                       item_name = "line")) where T
    individuals = Dict{Int32, Vector{T}}()
    while !eof(io)
        line_content = readuntil(io, 0x0a)
        record = parsenext(lineparser, line_content, 1, UInt8(';'))
        if isparsed(record)
            (penr, stich, am) = unsafe_get(record)
            mon = Dates.month(stich) # just the month
            state = aggregator[am]   # look up from dictionary
            if !haskey(individuals, penr)
                record = individuals[penr] = fill(missing, 12)
            else
                record = individuals[penr]
            end
            record[mon] = state
            increment!(tracker) # display progress
        end
    end
    individuals
end

function read_aggregated_AMs(filename::String, args...)
    io_gz = open(filename, "r")
    io = GzipDecompressorStream(io_gz)
    individuals = read_aggregated_AMs(io, args...)
    close(io)
    close(io_gz)
    individuals
end

"""
    $SIGNATURES

Expand the shell mapping from a compact format of a dictionary of

```julia
classification => [keys...]
```

to a dictionary of

```julia
key => classification
```
"""
function expand_spell_mapping(spell_mapping::Dict{T, Vector{String}}) where T
    aggregator = Dict{ByteVector, T}()
    for (k, vs) in spell_mapping
        for v in vs
            aggregator[convert(ByteVector, v)] = k
        end
    end
    aggregator
end

# runtime code here

spell_mapping = Dict(# unemployed
                     Int8(1) => String["D2", "SR", "VM", "AS", "LS", "SC", "AL",
                                       "AO","66"],
                     # NA
                     Int8(0) => String["KD", "SV", "TO",
                                       # FIXME PLEASE PUT THESE SOMETHERE
                                       "AO", "AL", "SC", "VM", "66", "LS"],
                     # employed
                     Int8(2) => String["FBENG", "FBUN1", "FBKOM", "LE", "FBEB2",
                                       "FBBS1", "FBEB1", "FBSOL", "LFIBA", "FBKUA",
                                       "LFNRM", "FBEPU", "LFP30", "FBSOB", "SO",
                                       "FBES", "LW", "LFUBA", "SBSVA", "FBBS3",
                                       "FBUN3", "LFVOL", "FBUN2", "S1S2", "G1", "FU",
                                       "FBBP", "FBEB", "FBARB", "BE", "AA", "LFLST",
                                       "LFTEL", "LFVRL", "FD", "LFJAS", "FBGEB",
                                       "FBBEB", "FBES1", "FBBS2", "FBGBP"],
                     # other (inactive)
                     Int8(1) => String["SF", "KG", "AU", "W1", "TA", "MS", "BA",
                                       "ED", "LL", "AG", "AF", "EO", "KO", "RE",
                                       "MK", "PZ", "AM", "LF", "MP", "SG", "W2"])

aggregator = expand_spell_mapping(spell_mapping)
test_file = data_file(2000)
histories = read_aggregated_AMs(test_file, aggregator, Int8(-1))


#!/usr/bin/env julia

const dir = abspath( splitdir(@__FILE__)[1] )

tic()
println( STDERR, "EvoDuplexes.jl loading and compiling... " )

using ArgParse
using BufferedStreams
using Libz
using GenomicFeatures
using BioSequences
using DataFrames
using Gadfly
using Distances
using Clustering
using MultivariateStats
using StatsBase

unshift!( LOAD_PATH, dir * "/../src" )
using EvoDuplexes
using EvoDuplexes: score, score!

function parse_cmd()
   s = ArgParseSettings()
   @add_arg_table s begin
      "--tree"
       help     = "Phylogenetic tree with neutral branch lengths (in newick format)"
       arg_type = String
       required = true
      "--cons-regions"
       help     = "BED file containing conserved regions"
       arg_type = String
       required = false
      "--gene-regions"
       help     = "BED file containing gene units to allow long-range folding within"
       arg_type = String
       required = true
      "--maf"
       help     = "Directory with MAF files named by chromosome (chr1.maf.gz...)"
       arg_type = String
       required = true
       default  = "../maf"
      "--model-load"
       help     = "Load pre-trained IsolationForest models, .evt.jls"
       arg_type = String
       required = false
      "--model-data"
       help     = "Load training data from `.jlt` file, output `.evt.jls` file"
       arg_type = String
       required = false
      "--model-train"
       help     = "Train the IsolationForest models from scratch using input regions, output `.evt.jls` file"
       action   = :store_true
      "--output","-o"
       help     = "Output prefix"
       arg_type = String
       default  = "evoduplex"
      "--cons-branch"
       help     = "When conserved regions are given, use this branch length multiplier"
       arg_type = Float64
       default  = 0.33
      "--sig-ratio"
       help     = "Output the top fraction of outliers"
       arg_type = Float64
       default  = 0.05
      "--output-all"
       help     = "Output all folds regardless of prediction status"
       action   = :store_true
      "--max-distance"
       help     = "Set a limit on the maximum distance between a left/right arm of a duplex"
       arg_type = Int64
       default  = 2000
      "--max-duplex-length"
       help     = "Set the maximum duplex length that will be accessible from the suffix array"
       arg_type = Int64
       default  = 50
      "--max-bulges"
       help     = "Set the maximum number of bulges to allow for a duplex"
       arg_type = Int64
       default  = 3
      "--max-mismatches"
       help     = "Set the maximum number of mismatches to allow for a duplex"
       arg_type = Int64
       default  = 3
      "--max-deltag"
       help     = "Set the maximum allowable deltaG value"
       arg_type = Float64
       default  = -0.8
    end
   return parse_args(s)
end

function main()

   istrue( arg ) = (arg != nothing)

   args = parse_cmd()

   println(STDERR, " $( round( toq(), 6 ) ) seconds" )

   neutraltree = parsenewick(String(chomp(String(read(open(args["tree"]))))))
   pairtree = deepcopy(neutraltree)
   constree = deepcopy(neutraltree)

   extend_branches!( neutraltree, 1.0 )
   extend_branches!( constree,    args["cons-branch"] )
   extend_branches!( pairtree,    args["cons-branch"] )

   set_prob_mat!( neutraltree, GTR_SINGLE_Q )
   set_prob_mat!( constree,    GTR_SINGLE_Q )
   set_prob_mat!( pairtree,    GTR_PAIRED_Q )

   secbool   = false
   conserved = false
   if istrue( args["cons-regions"] )
      regs  = loadbed( args["cons-regions"], expandfirst=25, expandlast=25 )
      conserved = true
      if istrue( args["gene-regions"] )
         genes   = loadbed( args["gene-regions"] )#, expandfirst=2500, expandlast=2500 )
         secbool = true
      else
         error("No specific `--gene-regions` given! Can't fold entire chromosomes..")
      end
   elseif istrue( args["gene-regions"] )
      genes = loadbed( args["gene-regions"] )
      regs  = genes
   else
      error("No specific regions given! Can't fold entire chromosomes..")
   end

   train = false
   if istrue( args["model-load"] )
      println("Loading model from $(args["model-load"]) file...")
      model = open(deserialize, args["model-load"])
   elseif istrue( args["model-data"] )
      println("Training model from .jlt datafile $(args["model-data"])...")
      tabdata = open(readdlm, args["model-data"])
      model = DistanceForest( args["sig-ratio"] )
      train!( model, tabdata )
   elseif args["model-train"]
      train = true
   else
      error("No --model-file or --model-data given!, if training is desired use --model-train!")
   end

   datout = BufferedOutputStream(open("$(args["output"])-raw.jlt", "w"))
   bedout = open("$(args["output"]).bed", "w")
   duplexes = Vector{DuplexInterval}()

   for chr in keys(genes.trees)
      println("Reading $(args["maf"])/$chr.maf.gz")
      reader = ZlibInflateInputStream(open("$(args["maf"])/$chr.maf.gz"), bufsize=129000) |> MAFReader
      @time maf = readmaf!( reader, neutraltree.index, minlength=18, minspecies=2, minscore=-Inf, regionbool=true, regions=genes, secbool=secbool, secregions=regs )

      println(maf)

      for r in collect(genes)
         haskey( maf.trees, r.seqname ) || continue
         col = collect(intersect( maf.trees[r.seqname], r ))
         length(col) >= 1 || continue
         if r.strand == STRAND_NEG
            for i in col
               EvoDuplexes.reverse_complement!( i.metadata )
            end
         end
         println("Building DuplexArray for $(r.metadata)")
         rda = RNADuplexArray{DNAAlphabet{2},UInt16,UInt16}( col[1].metadata, neutraltree, args["max-duplex-length"] )
         for i in 2:length(col)
            push!( rda, col[i].metadata, neutraltree )
         end

         @time trav = collect(traverse( rda, 1:args["max-distance"],
                                         single=neutraltree,
                                         bulge_max=args["max-bulges"],
                                         mismatch_max=args["max-mismatches"],
                                         minfold=args["max-deltag"] ))

         println("Calculating structural conservation across the phylogeny...")
         for i in trav
            #const conscore = (score( i.duplex, constree, gapdenom=1.25 ) - score( i.duplex, neutraltree )) / npairs( i.duplex )
            if !conserved
               score!( i.duplex, neutraltree, pairtree )
               const cons = score( i.duplex ) / npairs( i.duplex )
            else
               score!( i.duplex, constree, pairtree )
               const cons = score( i.duplex ) / npairs( i.duplex )
            end
            const ent     = dinucleotide_entropy(i.duplex.duplex) / (npairs( i.duplex ) + 5)
            const meancov = covariance( i.duplex ) * 10
            const len     = npairs( i.duplex )
            const dist    = distance( i )
            const deltag  = energy( i.duplex )

            write( datout, string(dist) ); write( datout, '\t' )
            write( datout, string(cons) ); write( datout, '\t' )
            write( datout, string(deltag) ); write( datout, '\t' )
            write( datout, string(len) ); write( datout, '\t' )
            write( datout, string(ent) ); write( datout, '\t' )
            write( datout, string(meancov) ); write( datout, '\n' )
         end

         duplexes = vcat(duplexes, trav)

      end
   end

   full_char = [collect('0':'9'); collect('A':'Z')]

   close(datout)
   tabdata = open(readdlm, "$(args["output"])-raw.jlt")

   if train
      println("Training IsolationForest models...")
      model = DistanceForest( args["sig-ratio"] )
      train!( model, tabdata )
   end

   println("Identifying statistical anomalies...")
   if !args["output-all"]
      sig = EvoDuplexes.predict( model, tabdata )
   else
      sig = collect(1:length(duplexes))
   end

   println("Found $(length(sig)) significant duplexes..")

   for i in sig
      name = randcode(full_char)
      writebed( bedout, duplexes[i], name )
      s,b = brackets( duplexes[i].duplex )
      println("$name: $b")
      println("$name: $s")
   end

   modelhndl = open( args["output"], "w" )
   println("Serializing model to $( args["output"] ).evt.jls")
   serialize(modelhndl, model)
   close(modelhndl)

   close(bedout)
end

@time main()


module DepParsing

importall ..JukaiNLP
using Merlin
using Compat
using ProgressMeter
using TransitionParser: beamsearch, max_violation!, to_seq

include("my_adagrad.jl")
include("../iddict.jl")
include("depparser.jl")
include("token.jl")
include("arcstd.jl")
include("util.jl")
include("feedforward.jl")
include("structuredfeedforward.jl")
include("unlabeled_feedforward.jl")
include("struct_extras.jl")
include("accuracy.jl")
include("io.jl")
include("saver.jl")
include("extra/utils.jl")
include("extra/perceptron.jl")

export DepParser,
       Unlabeled,
       Labeled,
       FeedForward,
       UnlabeledFeedForward,
       LabelerFF,
       assign!,
       StructuredFeedForward,
       Perceptron,
       readconll,
       train!,
       decode,
       evaluate,
       toconll,
       initmodel!,
       MyAdaGrad

end