text:
"""
`module Constraints`

TODO: write documentation
"""
module Constraints

using JuLIP: Dofs, AbstractConstraint, AbstractAtoms,
             mat, vecs, JVecs, JVecsF, JMatF, JMat,
             set_positions!, set_cell!, virial, defm, set_defm!,
             forces, momenta, set_momenta!,
             constraint, rnn, AbstractCalculator, calculator

# temporary hack - we may need to rethink this
using JuLIP.Potentials: hessian_pos

import JuLIP: position_dofs, project!, set_position_dofs!, positions, gradient,
              energy, momentum_dofs, set_momentum_dofs!, dofs, project!,
              set_dofs!, positions, gradient, energy, hessian

export FixedCell, VariableCell, ExpVariableCell, FixedCell2D, atomdofs

function zeros_free{T}(n::Integer, x::Vector{T}, free::Vector{Int})
   z = zeros(T, n)
   z[free] = x
   return z
end

function insert_free!{T}(p::Array{T}, x::Vector{T}, free::Vector{Int})
   p[free] = x
   return p
end

# a helper function to get a valid positions array from a dof-vector
positions{TI<:Integer}(at::AbstractAtoms, ifree::AbstractVector{TI}, dofs::Dofs) =
      insert_free!(positions(at) |> mat, dofs, ifree) |> vecs

# ========================================================================
#          FIXED CELL IMPLEMENTATION
# ========================================================================

"""
`FixedCell`: no constraints are placed on the motion of atoms, but the
cell shape is fixed

Constructor:
```julia
FixedCell(at::AbstractAtoms; free=..., clamp=..., mask=...)
```
Set at most one of the kwargs:
* no kwarg: all atoms are free
* `free` : list of free atom indices (not dof indices)
* `clamp` : list of clamped atom indices (not dof indices)
* `mask` : 3 x N Bool array to specify individual coordinates to be clamped
"""
type FixedCell <: AbstractConstraint
   ifree::Vector{Int}
end

function analyze_mask(at, free, clamp, mask)
   if length(find((free != nothing, clamp != nothing, mask != nothing))) > 1
      error("FixedCell: only one of `free`, `clamp`, `mask` may be provided")
   elseif all( (free == nothing, clamp == nothing, mask == nothing) )
      # in this case (default) all atoms are free
      return collect(1:3*length(at))
   end
   # determine free dof indices
   Nat = length(at)
   if clamp != nothing
      # revert to setting free
      free = setdiff(1:Nat, clamp)
   end
   if free != nothing
      # revert to setting mask
      mask = Matrix{Bool}(3, Nat)
      fill!(mask, false)
      mask[:, free] = true
   end
   return find(mask[:])
end

FixedCell(at::AbstractAtoms; clamp = nothing, free=nothing, mask=nothing) =
   FixedCell(analyze_mask(at, free, clamp, mask))

position_dofs(at::AbstractAtoms, cons::FixedCell) = mat(positions(at))[cons.ifree]

set_position_dofs!(at::AbstractAtoms, cons::FixedCell, x::Dofs) =
      set_positions!(at, positions(at, cons.ifree, x))

momentum_dofs(at::AbstractAtoms, cons::FixedCell) = mat(momenta(at))[cons.ifree]

set_momentum_dofs!(at::AbstractAtoms, cons::FixedCell, p::Dofs) =
      set_momenta!(at, zeros_free(3 * length(at), p, cons.ifree) |> vecs)

project!(at::AbstractAtoms, cons::FixedCell) = at

# TODO: this is a temporary hack, and I think we need to
#       figure out how to do this for more general constraints
#       maybe not too terrible
#       but why is there a !???
project!(cons::FixedCell, A::SparseMatrixCSC) = A[cons.ifree, cons.ifree]

gradient(at::AbstractAtoms, cons::FixedCell) =
               scale!(mat(forces(at))[cons.ifree], -1.0)

energy(at::AbstractAtoms, cons::FixedCell) = energy(at)

# >>> TODO: move to abstractions.jl ??
hessian(at::AbstractAtoms, cons::FixedCell) =
      project!(cons, _pos_to_alldof(hessian_pos(calculator(at), at), at))

function _pos_to_alldof(Hpos::SparseMatrixCSC, at::AbstractAtoms)
   I, J, Z = Int[], Int[], Float64[]
   for C in (I, J, Z); sizehint!(C, 9 * nnz(Hpos)); end
   Nat = length(at)
   @assert Nat == size(Hpos, 2)
   # TODO: this findnz creates an extra copy of all data, which we should avoid
   for (iat, jat, zat) in zip(findnz(Hpos)...)
      for a = 1:3, b = 1:3
         push!(I, 3 * (iat-1) + a)
         push!(J, 3 * (jat-1) + b)
         push!(Z, zat[a,b])
      end
   end
   return sparse(I, J, Z, 3*Nat, 3*Nat)
end

# ===================
#   In-Plane Motion
# ===================

FixedCell2D(at::AbstractAtoms; kwargs...) = InPlaneFixedCell(at; kwargs...)

"""
preliminary implementation of a Constraint, restricting a
simulation to 2D in-plane motion
"""
function InPlaneFixedCell(at::AbstractAtoms; clamp = Int[], free = nothing)
   if free == nothing
      free = setdiff(1:length(at), clamp)
   end
   mask = fill(false, (3, length(at)))
   mask[:, free] = true
   mask[3, :] = false
   return FixedCell(at, mask = mask)
end

# =====================
#   Anti-Plane Motion
# =====================

"""
preliminary implementation of a Constraint, restricting a
simulation to 2D in-plane motion
"""
function AntiPlaneFixedCell(at::AbstractAtoms; free = Int[])
   mask = fill(false, (3, length(at)))
   mask[:, free] = true
   mask[1:2, :] = false
   return FixedCell(at, mask = mask)
end

# ========================================================================
#          VARIABLE CELL IMPLEMENTATION
# ========================================================================

"""
`VariableCell`: both atom positions and cell shape are free;

**WARNING:** before manipulating the dof-vectors returned by a `VariableCell`
constraint, read *meaning of dofs* instructions at bottom of help text!

Constructor:
```julia
VariableCell(at::AbstractAtoms; free=..., clamp=..., mask=..., fixvolume=false)
```
Set at most one of the kwargs:
* no kwarg: all atoms are free
* `free` : list of free atom indices (not dof indices)
* `clamp` : list of clamped atom indices (not dof indices)
* `mask` : 3 x N Bool array to specify individual coordinates to be clamped

### Meaning of dofs

On call to the constructor, `VariableCell` stored positions and deformation
`X0, F0`, dofs are understood *relative* to this "initial configuration".

`dofs(at, cons::VariableCell)` returns a vector that represents a pair
`(Y, F1)` of a displacement and a deformation matrix. These are to be understood
*relative* to the reference `X0, F0` stored in `cons` as follows:
* `F = F1`   (the cell is then `F'`)
* `X = [F1 * (F0 \ y)  for y in Y)]`

One aspect of this definition is that clamped atom positions still change via
`F`.
"""
type VariableCell <: AbstractConstraint
   ifree::Vector{Int}
   X0::JVecsF
   F0::JMatF
   pressure::Float64
   fixvolume::Bool
   volume::Float64    # this is meaningless if `fixvolume == false`
end

function VariableCell(at::AbstractAtoms;
               free=nothing, clamp=nothing, mask=nothing,
               pressure = 0.0, fixvolume=false)
   if pressure != 0.0 && fixvolume
      warning("the pressure setting will be ignores when `fixvolume==true`")
   end
   return VariableCell( analyze_mask(at, free, clamp, mask),
                        positions(at), defm(at),
                        pressure, fixvolume, det(defm(at)) )
end

# reverse map:
#   F -> F
#   X[n] = F * F^{-1} X0[n]

function position_dofs(at::AbstractAtoms, cons::VariableCell)
   X = positions(at)
   F = defm(at)
   A = cons.F0 * inv(F)
   U = [A * x for x in X]   # switch to broadcast!
   return [mat(U)[cons.ifree]; Matrix(F)[:]]
end

posdofs(x) = x[1:end-9]
celldofs(x) = x[end-8:end]

function set_position_dofs!(at::AbstractAtoms, cons::VariableCell, x::Dofs)
   F = JMatF(celldofs(x))
   A = F * inv(cons.F0)
   Y = copy(cons.X0)
   mat(Y)[cons.ifree] = posdofs(x)
   for n = 1:length(Y)
      Y[n] = A * Y[n]
   end
   set_positions!(at, Y)
   set_defm!(at, F)
   return at
end

# for a variation x^t_i = (F+tU) F_0^{-1} (u_i + t v_i)
#       ~ U F^{-1} F F0^{-1} u_i + F F0^{-1} v_i
# we get
#      dE/dt |_{t=0} = U : (S F^{-T}) - < (F * inv(F0))' * frc, v>
#
# this is nice because there is no contribution from the stress to
# the positions component of the gradient

vol(at::AbstractAtoms) = det(defm(at))

vol_d(at::AbstractAtoms) = vol(at) * inv(defm(at))'

# function vol_dd(at::AbstractAtoms)
#    hdetI = zeros(3,3,3,3)
#    h = 0.1
#    for i = 1:3, j = 1:3
#       Ih = eye(3); Ih[i,j] += h
#       hdetI[:,:,i,j] = (ddet(Ih) - ddetI) / h
#    end
#    round(Int, reshape(hdetI, 9, 9))
# end

function gradient(at::AbstractAtoms, cons::VariableCell)
   F = defm(at)
   A = F * inv(cons.F0)
   G = forces(at)
   for n = 1:length(G)
      G[n] = - A' * G[n]
   end
   S = - virial(at) * inv(F)'        # ∂E / ∂F
   S += cons.pressure * vol_d(at)     # applied stress
   return [ mat(G)[cons.ifree]; Array(S)[:] ]
end

energy(at::AbstractAtoms, cons::VariableCell) =
         energy(at) + cons.pressure * det(defm(at))

# TODO: fix this once we implement the volume constraint ??????
project!(at::AbstractAtoms, cons::VariableCell) = at

# TODO: fix the abstraction for projecting a preconditioner;
#       this will actually need to do quite a bit more in the future
# project!(cons::FixedCell, A::SparseMatrixCSC) = A[cons.ifree, cons.ifree]

# TODO: CONTINUE WITH EXPCELL IMPLEMENTATION
# include("expcell.jl")

# convenience function to return DoFs associated with a particular atom
atomdofs(a::AbstractAtoms, I::Integer) = findin(constraint(a).ifree, 3*I-2:3*I)

end # module


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct WhatIfPropertyChange <: SwaggerModel
    path::Any # spec type: Union{ Nothing, String } # spec name: path
    propertyChangeType::Any # spec type: Union{ Nothing, String } # spec name: propertyChangeType
    before::Any # spec type: Union{ Nothing, Any } # spec name: before
    after::Any # spec type: Union{ Nothing, Any } # spec name: after
    children::Any # spec type: Union{ Nothing, Vector{WhatIfPropertyChange} } # spec name: children

    function WhatIfPropertyChange(;path=nothing, propertyChangeType=nothing, before=nothing, after=nothing, children=nothing)
        o = new()
        validate_property(WhatIfPropertyChange, Symbol("path"), path)
        setfield!(o, Symbol("path"), path)
        validate_property(WhatIfPropertyChange, Symbol("propertyChangeType"), propertyChangeType)
        setfield!(o, Symbol("propertyChangeType"), propertyChangeType)
        validate_property(WhatIfPropertyChange, Symbol("before"), before)
        setfield!(o, Symbol("before"), before)
        validate_property(WhatIfPropertyChange, Symbol("after"), after)
        setfield!(o, Symbol("after"), after)
        validate_property(WhatIfPropertyChange, Symbol("children"), children)
        setfield!(o, Symbol("children"), children)
        o
    end
end # type WhatIfPropertyChange

const _property_map_WhatIfPropertyChange = Dict{Symbol,Symbol}(Symbol("path")=>Symbol("path"), Symbol("propertyChangeType")=>Symbol("propertyChangeType"), Symbol("before")=>Symbol("before"), Symbol("after")=>Symbol("after"), Symbol("children")=>Symbol("children"))
const _property_types_WhatIfPropertyChange = Dict{Symbol,String}(Symbol("path")=>"String", Symbol("propertyChangeType")=>"String", Symbol("before")=>"Any", Symbol("after")=>"Any", Symbol("children")=>"Vector{WhatIfPropertyChange}")
Base.propertynames(::Type{ WhatIfPropertyChange }) = collect(keys(_property_map_WhatIfPropertyChange))
Swagger.property_type(::Type{ WhatIfPropertyChange }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_WhatIfPropertyChange[name]))}
Swagger.field_name(::Type{ WhatIfPropertyChange }, property_name::Symbol) =  _property_map_WhatIfPropertyChange[property_name]

const _allowed_WhatIfPropertyChange_propertyChangeType = ["Create", "Delete", "Modify", "Array"]

function check_required(o::WhatIfPropertyChange)
    (getproperty(o, Symbol("path")) === nothing) && (return false)
    true
end

function validate_property(::Type{ WhatIfPropertyChange }, name::Symbol, val)
    if name === Symbol("propertyChangeType")
        Swagger.validate_param(name, "WhatIfPropertyChange", :enum, val, _allowed_WhatIfPropertyChange_propertyChangeType)
    end
end

ext: jl


text:
###################################################################################################
# BIOINFORMATICS FUNCTIONS
###################################################################################################
"""
    `get_align_strand(PAIRED_END,FLAG1,FLAG2)`

    Function that returns the strand of methylation call based on Bismark's logic. In single-end mode,
    OT and CTOT reads will both receive a FLAG of 0 while OB and CTOB both receive a FLAG of 16. In
    paired end mode:

                Read 1       Read 2
    OT:         99          147
    OB:         83          163
    CTOT:      147           99
    CTOB:      163           83

    This table was extracted from

        https://github.com/FelixKrueger/Bismark/issues/151

    # Examples
    ```julia-repl
    julia> CpelTdm.get_align_strand(true,UInt16(99),UInt16(147))
    "OT"
    ```
"""
function get_align_strand(pe::Bool,flag1::UInt16,flag2::UInt16)::String

    # Treat SE and PE separately
    if !pe
        if flag1==0
            s="OT"
        elseif flag1==16
            s="OB"
        else
            print_log("SE: was expecting FLAG 0 or and encountered $(flag1) instead.")
            print_log("Exiting julia ...")
            exit(1)
        end
    else
        if flag1==99 && flag2==147
            s="OT"
        elseif flag1==83 && flag2==163
            s="OB"
        elseif flag1==147 && flag2==99
            s="CTOT"
        elseif flag1==163 && flag2==83
            s="CTOB"
        else
            print_log("PE: unexpected flag combination. Expected 99/147, 147/99, 83/163, 163/83.")
            print_log("Exiting julia ...")
            exit(1)
        end
    end

    # Return strand
    return s

end # end get_align_strand
"""
    `order_bams(PAIRED_END,RECORDS)`

    Function that returns an AlignTemp object with R1 as the first record in the forward strand and
    with the methylation call strand taken from `get_align_strand`.

    # Examples
    ```julia-repl
    julia> CpelTdm.order_bams(true,RECORDS)
    ```
"""
function order_bams(pe::Bool,records::Vector{BAM.Record})::AlignTemp

    # Check which record appears first in forward strand
    if BAM.position(records[1])<=BAM.position(records[2])
        s = get_align_strand(pe, BAM.flag(records[1]), BAM.flag(records[2]))
        return AlignTemp(s,records[1],records[2])
    else
        s = get_align_strand(pe, BAM.flag(records[2]), BAM.flag(records[1]))
        return AlignTemp(s,records[2],records[1])
    end

end
"""
    `clean_records(PAIRED_END,RECORDS)`

    Function that takes a set of records and returns an AllAlignTemps object that contains all the
    properly aligned templates as an array of AlignTemp, which contains information about the
    methylation call strand as well as the relevant BAM records. In the PE case, R1 corresponds to
    the BAM record that appears before in the forward strand.

    # Examples
    ```julia-repl
    julia> CpelTdm.clean_records(true,RECORDS)
    ```
"""
function clean_records(pe::Bool,records::Vector{BAM.Record})::AllAlignTemps

    # Initialize struct
    out = AllAlignTemps(pe,[])

    # Consider PE vs. SE
    if pe
        # Handle PE case, first retrieve unique template names
        temp_names = unique([BAM.tempname(x) for x in records])
        for temp_name in temp_names
            # Get records with same template name
            temp_recs = filter(x->BAM.tempname(x)==temp_name,records)

            # There should only be two records with the same template name
            length(temp_recs)==2 && push!(out.templates,order_bams(pe,temp_recs))
        end
    else
        # Handle SE case
        out.templates = [AlignTemp(get_align_strand(pe,BAM.flag(x),UInt16(0)),x,BAM.Record()) for
            x in records]
    end

    # Return struct
    return out

end # end clean_records
"""
    `try_olaps(READER,CHR,WINDOW)`

    Function that tries to find BAM records overlaping with `CHR` at positions `WINDOW`.

    # Examples
    ```julia-repl
    julia> CpelTdm.try_olaps(reader,chr,win)
    ```
"""
function try_olaps(reader::BAM.Reader,chr::String,win::Vector{Int64})::Vector{BAM.Record}

    # NOTE: known bug that has to do w/ BioAlignments when close to end?
    records_olap =
    try
        collect(eachoverlap(reader, chr, win[1]:win[2]))
    catch
        Array{BAM.Record,1}()
    end

    return records_olap

end # try_olaps
"""
    `read_bam(BAM_PATH,CHR,FEAT_ST,FEAT_END,CPG_POS,CHR_SIZE,PE,TRIM)`

    Function that reads in BAM file in `BAM_PATH` and returns methylation vectors for those records
    that overlap with (1-based) genomic coordinates `chr:chrSt-chrEnd` at `cpg_pos`. A trimming given
    by TRIM=(Rf_5p,Rf_3p,Rr_5p,Rr_3p) is applied to the reads. The information was taken from:

        XM: meth calls (https://github.com/FelixKrueger/Bismark/tree/master/Docs)
        XS: uniqueness (http://bowtie-bio.sourceforge.net/bowtie2/manual.shtml)

    For info on OT, CTOT, OB, CTOB nomenclature see

        http://software.broadinstitute.org/software/igv/book/export/html/37.

    # Examples
    ```julia-repl
    julia> CpelTdm.read_bam(BAM_PATH,"chr1",30,80,[40,60],1000,false,(0,0,0,0))
    ```
"""
function read_bam(bam::String,chr::String,roi_st::Int64,roi_end::Int64,cpg_pos::Vector{Int64},
                  chr_size::Int64,pe::Bool,trim::NTuple{4,Union{Int64,Float64}})::Array{Vector{Int64},1}

    # Number of CpG sites is determined by that in the region
    N = length(cpg_pos)

    # Expand window in PE case to include pair even if it's outside the window
    exp_win = pe ? [max(1,roi_st-75),min(chr_size,roi_end+75)] : [roi_st,roi_end]

    # Get records overlapping window.
    reader = open(BAM.Reader,bam,index=bam*".bai")
    records_olap = try_olaps(reader,chr,exp_win)
    length(records_olap)>0 || return Vector{Int64}[]
    close(reader)

    # Relevant flags in BAM file (both Bismark and Arioc)
    filter!(x-> (BAM.ismapped(x)) && (haskey(x,"XM")) && (!haskey(x,"XS")) && (BAM.flag(x) in
           FLAGS_ALLOWED) && (BAM.mappingquality(x)>THRESH_MAPQ),records_olap)

    # Clean records & organize them
    records_org = clean_records(pe,records_olap)

    # Loop over organized records
    xobs = Vector{Int64}[]
    for record in records_org.templates
        # Get methylation call and offset (depends on strand where call is made)
        if !(records_org.paired)
            # Obtain methylation call from single end ()
            meth_call = record.R1[:"XM"]
            meth_call = SubString(meth_call,(1+trim[1]):(length(meth_call)-trim[2]))
            OFFSET = record.strand in ["OT","CTOT"] ? 1 : 2
            OFFSET -= BAM.position(record.R1)+trim[1]
        else
            # Obtain methylation call
            R1_call = record.R1[:"XM"]
            R2_call = record.R2[:"XM"]
            R1_call = SubString(R1_call,(1+trim[1]):(length(R1_call)-trim[2]))
            R2_call = SubString(R2_call,(1+trim[4]):(length(R2_call)-trim[3]))
            dist_st = abs(BAM.position(record.R2)+trim[4] - BAM.position(record.R1)-trim[1])
            meth_call = R1_call * "."^max(0,dist_st-length(R1_call))
            meth_call *= R2_call[max(1,(length(R1_call)+1-dist_st)):end]
            OFFSET = record.strand in ["OT","CTOT"] ? 1 : 2
            OFFSET -= BAM.position(record.R1)+trim[1]
        end

        # Cross positions of CpG sites if template contains CpG sites
        obs_cpgs = [x.offset for x in eachmatch(r"[zZ]",meth_call)] .- OFFSET
        length(obs_cpgs)>0 || continue

        # If overlapping CpG sites then store and add to xobs
        olap_cpgs = findall(x-> x in obs_cpgs, cpg_pos)
        length(olap_cpgs)>0 || continue
        x = zeros(Int64,N)
        obs_cpgs = obs_cpgs[findall(x-> x in cpg_pos,obs_cpgs)] .+ OFFSET
        x[olap_cpgs] = reduce(replace,["Z"=>1,"z"=>-1],init=split(meth_call[obs_cpgs],""))

        # Add to set of observations
        push!(xobs,x)

    end # end loop over templates sequenced

    # Return
    return xobs

end # end read_bam
"""
    `split_bed_record(bed_rec,max_size_anal_reg)`

    Function that splits regions of interest in BED into analysis regions.

    # Examples
    ```julia-repl
    julia> str_j_rec = "chr1\t1000\t2001\t."
    julia> using BED; bed_rec = BED.Record(str_j_rec);
    julia> CpelTdm.split_bed_record(bed_rec,500)
    2-element Array{BED.Record,1}:
    BED.Record:
        chromosome: chr1
                start: 1000
                end: 1500
                name: .
    BED.Record:
        chromosome: chr1
                start: 1500
                end: 2001
                name: .
    ```
"""
function split_bed_record(bed_rec::BED.Record,max_size_anal_reg::Union{Int64,Float64})::Vector{BED.Record}

    ## BED files are:
    ## - 0-based coordinates
    ## - chromstart is inclusive
    ## - chromend is non-inclusive
    ## When read in by BED.jl the start becomes one-based.

    # Get size of record
    size_rec = BED.chromend(bed_rec) - BED.chromstart(bed_rec) + 1

    # Initialize
    recs = Vector{BED.Record}()

    # Divide into analysis regions
    chr = BED.chrom(bed_rec)
    n_anal_reg = Int(ceil(size_rec/max_size_anal_reg))
    size_anal_reg = Int(floor(size_rec/n_anal_reg))

    # Create n_anal_reg analysis regions
    @inbounds for j=1:n_anal_reg

        # Get coordinates
        chr_st = BED.chromstart(bed_rec) + (j-1) * size_anal_reg - 1
        chr_end = j==n_anal_reg ? BED.chromend(bed_rec) : chr_st + size_anal_reg + 1

        # Create string with BED record
        str_j_rec = "$(chr)\t$(chr_st)\t$(chr_end)\t."

        # Push record
        push!(recs,BED.Record(str_j_rec))

    end

    # Return vector
    return recs

end
"""
    `gen_anal_reg_file(bed,max_size_anal_reg)`

    Function that dvides regions of interest in BED into analysis regions.

    # Examples
    ```julia-repl using CpelTdm
    julia> max_size_anal_reg = 500;
    julia> bed_out = CpelTdm.get_bed_out_filename(bed);
    julia> CpelTdm.gen_anal_reg_file(bed,bed_out,max_size_anal_reg)
    ```
"""
function gen_anal_reg_file(bed::String,bed_out::String,max_size_anal_reg::Union{Int64,Float64})::Nothing

    # Check if BED file already exists
    isfile(bed_out) && (print_log("Found existing CpelTdm BED file ..."); return nothing)

    # Read new BED file
    reader = open(BED.Reader,bed)
    record = BED.Record()
    anal_regs = Vector{BED.Record}()
    while !eof(reader)
        # Get record
        read!(reader,record)
        # Append new analysis regions
        append!(anal_regs,split_bed_record(record,max_size_anal_reg))
    end
    close(reader)

    # Write new BED file
    output = open(bed_out,"a")
    writer = BED.Writer(output)
    while length(anal_regs)>0
        write(writer,popfirst!(anal_regs))
    end
    close(writer)

    # Return nothing
    return nothing

end
"""
    `get_anal_reg_chr(bed,chr)`

    Function that returns analyses regions in chr.

    # Examples
    ```julia-repl
    julia> CpelTdm.get_anal_reg_chr(bed,"chr1")
    ```
"""
function get_anal_reg_chr(bed::String,chr::String)::Vector{BED.Record}

    # Read new BED file
    reader = open(BED.Reader,bed)
    anal_regs = Vector{BED.Record}()
    while !eof(reader)
        # Get record
        # Note: when reading converts to a 1-based and start and end are inclusive
        record = BED.Record()
        read!(reader,record)
        # Keep if in chromosome
        BED.chrom(record)==chr && push!(anal_regs,record)
    end
    close(reader)

    # Return nothing
    return anal_regs

end
"""
    `get_cpg_pos!(ROI,FASTA)`

    Function that stores vector of CpG site positions in ROI.

    # Examples
    ```julia-repl
    julia> CpelTdm.get_cpg_pos!(roi,fasta)
    ```
"""
function get_cpg_pos!(roi_data::RoiData,fasta::String)::Nothing

    # Get sequence
    fasta_reader = open(FASTA.Reader,fasta,index=fasta*".fai")
    fa_record = fasta_reader[roi_data.chr]
    close(fasta_reader)

    # Get position CpG sites
    roi_seq = convert(String,FASTA.sequence(fa_record,roi_data.chrst:roi_data.chrend))
    cpg_pos = map(x->getfield(x,:offset),eachmatch(r"CG",roi_seq)) .+ roi_data.chrst .- 1

    # Set cpg_pos
    roi_data.cpg_pos = cpg_pos

    # Return nothing
    return nothing

end
"""
    `get_ns!(ROI,MAX_SUBREGION_SIZE)`

    Function that stores vector [N1,...,NK] given the position of the CpG sites & maximum subregion size.

    # Examples
    ```julia-repl
    julia> roi_data = CpelTdm.RoiData(10,10);
    julia> roi_data.chr = "chr22"; roi_data.chrst = 10; roi_data.chrend = 400;
    julia> roi_data.cpg_pos = [100,200,300,350];
    julia> max_size_subreg = 200;
    julia> CpelTdm.get_ns!(roi_data,max_size_subreg)
    julia> roi_data.n_vec
    2-element Array{Int64,1}:
    2
    2
    ```
"""
function get_ns!(roi_data::RoiData,max_size_subreg::Union{Int64,Float64})::Nothing

    # Check if need to partition
    if (roi_data.chrend-roi_data.chrst)<=max_size_subreg
        roi_data.n_vec = [length(roi_data.cpg_pos)]
        return nothing
    end

    # Get K of model
    k = ceil((roi_data.chrend-roi_data.chrst+1)/max_size_subreg)

    # Get delimiters
    delimiters = vcat(collect(roi_data.chrst:(roi_data.chrend-roi_data.chrst+1)/k:roi_data.chrend),roi_data.chrend)

    # Partition CpG sites
    ids = [findfirst(y->y>=x,delimiters)-1 for x in roi_data.cpg_pos]

    # Set n vector
    roi_data.n_vec = [sum(ids.==i) for i in unique(ids)]

    # Return nothing
    return nothing

end
"""
    `get_obs_per_cpg(XOBS)`

    Function that returns the number of observations per CpG site as a vector.

    # Examples
    ```julia-repl
    julia> CpelTdm.get_obs_per_cpg(xobs)
    ```
"""
function get_obs_per_cpg(xobs::Array{Vector{Int64},1})::Vector{Int64}

    # Get observations per CpG site
    out = length(xobs)>0 ? vcat(sum(abs.(hcat(xobs...)),dims=2)...) : [0]

    # Return
    return out

end
"""
    `mean_cov(XOBS)`

    Function returns the average coverage per CpG given some observations `XOBS`.

    # Examples
    ```julia-repl
    julia> xobs=[[1,-1] for i=1:10]; append!(xobs,[[1,0] for i=1:10]);
    julia> CpelTdm.mean_cov(xobs)
    15.0
    ```
"""
function mean_cov(xobs::Array{Vector{Int64},1})::Float64

    # Return 0 if no observations
    return length(xobs)>0 ? norm(hcat(xobs...),1)/length(xobs[1]) : 0.0

end
"""
    `pmap_anal_roi(ROI,CHR,CHR_SIZE,BAMS1,BAMS2,FASTA,CONFIG)`

    Function that estimates θ and computes MML/NME for all samples, and does permutation test for a given ROI.

    # Examples
    ```julia-repl
    julia> CpelTdm.pmap_anal_roi(roi,chr,chr_size,bams1,bams2,fasta,config)
    ```
"""
function pmap_anal_roi(roi::BED.Record,chr::String,chr_size::Int64,bams1::Vector{String},bams2::Vector{String},
                       fasta::String,config::CpeltdmConfig)::RoiData

    # Print ROI for testing
    # println(roi)

    # Get sample size
    s1 = length(bams1)
    s2 = length(bams2)

    # Empty output
    roi_data = RoiData(s1,s2)

    # Store chromosomal info
    roi_data.chr = BED.chrom(roi)
    roi_data.chrst = BED.chromstart(roi)
    roi_data.chrend = BED.chromend(roi)

    # Get CpG sites
    get_cpg_pos!(roi_data,fasta)
    length(roi_data.cpg_pos)>0 || return roi_data

    # Get vector of Ns
    get_ns!(roi_data,config.max_size_subreg)

    ## Estimate θs

    # Loop over group 1 samples
    # print_log("First group ...")
    @inbounds for i=1:s1

        # Get BAM file
        bam = bams1[i]

        # Get vectors from BAM overlapping region
        xobs = read_bam(bam,chr,roi_data.chrst,roi_data.chrend,roi_data.cpg_pos,chr_size,config.pe,config.trim)
        obs_per_cpg = get_obs_per_cpg(xobs)
        (config.min_cov<=mean_cov(xobs)<=400) && (sum(obs_per_cpg.==0)<=1.0/5.0*sum(roi_data.n_vec)) || continue

        # Estimate parameters of CPEL models
        θhat = est_theta_sa(roi_data.n_vec,xobs)
        config.bound_check && check_boundary(θhat) && continue
        roi_data.θ1s[i] = θhat

        # Estimate intermediate quantities
        ∇logZ = get_∇logZ(roi_data.n_vec,θhat)

        # Compute output
        roi_data.mmls1[i] = comp_mml(roi_data.n_vec,∇logZ)
        roi_data.nmes1[i] = comp_nme(roi_data.n_vec,θhat,∇logZ)

        # Set i-th sample done
        roi_data.analyzed1[i] = true

    end

    # Loop over group 2 samples
    # print_log("Second group ...")
    @inbounds for i=1:s2

        # Get BAM file
        bam = bams2[i]

        # Get vectors from BAM overlapping region
        xobs = read_bam(bam,chr,roi_data.chrst,roi_data.chrend,roi_data.cpg_pos,chr_size,config.pe,config.trim)
        obs_per_cpg = get_obs_per_cpg(xobs)
        (config.min_cov<=mean_cov(xobs)<=400) && (sum(obs_per_cpg.==0)<=1.0/5.0*sum(roi_data.n_vec)) || continue

        # Estimate parameters of CPEL models
        θhat = est_theta_sa(roi_data.n_vec,xobs)
        config.bound_check && check_boundary(θhat) && continue
        roi_data.θ2s[i] = θhat

        # Estimate intermediate quantities
        ∇logZ = get_∇logZ(roi_data.n_vec,θhat)

        # Compute output
        roi_data.mmls2[i] = comp_mml(roi_data.n_vec,∇logZ)
        roi_data.nmes2[i] = comp_nme(roi_data.n_vec,θhat,∇logZ)

        # Set i-th sample done
        roi_data.analyzed2[i] = true

    end

    ## Differential analysis
    # print_log("Differential analysis ...")
    if config.matched

        # Find pairs with data. TODO: issue with -1 power
        ind = roi_data.analyzed1 .& roi_data.analyzed2
        (0.5^(sum(ind)-1))<0.05 || return roi_data

        # Get sample pairs with data
        θ1s = roi_data.θ1s[ind]
        θ2s = roi_data.θ2s[ind]

        # Perform matched differential analysis
        tmml_test,tnme_test,tpdm_test = mat_tests(roi_data.n_vec,θ1s,θ2s)

    else

        # Check if sufficient data
        s1 = sum(roi_data.analyzed1)
        s2 = sum(roi_data.analyzed2)
        1.0/binomial(s1+s2,s1)<0.05 || return roi_data

        # Get samples with data
        θ1s = roi_data.θ1s[roi_data.analyzed1]
        θ2s = roi_data.θ2s[roi_data.analyzed2]

        # Perform unmatched differential analysis
        tmml_test,tnme_test,tpdm_test = unmat_tests(roi_data.n_vec,θ1s,θ2s)

    end

    # Set differential analysis done
    roi_data.diff_anal_done = true

    # Store results
    roi_data.mml_test = tmml_test
    roi_data.nme_test = tnme_test
    roi_data.pdm_test = tpdm_test

    # Return
    return roi_data

end
"""
    `anal_bed_file(BAMS1,BAMS2,BED,FASTA,OUT_PATHS,CONFIG)`

    Function that estimates θ for all samples, computes statistics, and does a matched/unmatched
    permutation test at each ROI.

    # Examples
    ```julia-repl
    julia> CpelTdm.anal_bed_file(bams1,bams2,bed,fasta,out_paths,config)
    ```
"""
function anal_bed_file(bams1::Vector{String},bams2::Vector{String},bed::String,fasta::String,
                       out_paths::NTuple{3,Vector{String}},config::CpeltdmConfig)::Nothing

    # Find chromosomes
    fasta_reader = open(FASTA.Reader,fasta,index=fasta*".fai")
    chr_sizes = fasta_reader.index.lengths
    chr_names = fasta_reader.index.names
    close(fasta_reader)

    # Loop over chromosomes
    for chr in chr_names

        # Print info
        print_log("Processing chr: $(chr) ...")

        # Get windows pertaining to current chromosome
        print_log("Obtaining corresponding ROIs ...")
        roi_chr = get_anal_reg_chr(bed,chr)
        chr_size = chr_sizes[findfirst(x->x==chr,chr_names)]
        length(roi_chr)>0 || continue

        # Process regions of interest in chromosome
        print_log("Performing differential analysis on ROIs ...")
        out_pmap = pmap(x->pmap_anal_roi(x,chr,chr_size,bams1,bams2,fasta,config),roi_chr)

        # Add last to respective bedGraph file
        length(out_pmap)>0 && write_output(out_pmap,out_paths)

    end

    # Correct for multiple hypothesis in differential analysis
    print_log("Performing multiple hypothesis correction ...")
    mult_hyp_corr(out_paths[3])

    # Return nothing
    return nothing

end
"""
    `cpel_tdm(BAMS1,BAMS2,BED,FASTA,OUTDIR,PREFIX)`

    Function to call to run CpelTdm analysis.

        - BAMS1: BAM files associated to group 1 (require index file).
        - BAMS2: BAM files associated to group 2 (require index file).
        - BED: BED file containing regions of interest.
        - FASTA: reference genome (requires index file).
        - OUTDIR: output directory to store results.
        - PREFIX: prefix of output differential files.

    # Examples
    ```julia-repl
    julia> CpelTdm.cpel_tdm(bams1,bams2,bed,fasta,outdir,prefix)
    ```
"""
function cpel_tdm(bams1::Vector{String},bams2::Vector{String},bed::String,fasta::String,outdir::String,prefix::String;
                  pe::Bool=false,max_size_anal_reg::Union{Int64,Float64}=1000,max_size_subreg::Union{Int64,Float64}=250,
                  min_cov::Union{Int64,Float64}=5,matched::Bool=false,trim::NTuple{4,Union{Int64,Float64}}=(0,0,0,0),
                  bound_check::Bool=false)::Nothing

    # Print initialization of juliASM
    print_log("Starting CpelTdm ...")

    ## IO
    print_log("Checking IO ...")

    # Get paths
    out_paths = get_paths(bams1,bams2,fasta,bed,outdir,prefix)

    # Create output folder if it doesn't exist
    isdir(outdir) || mkdir(outdir)

    ## Configure run
    print_log("Configuring CpelTdm ...")
    matched ? print_log("Matched comparison ...") : print_log("Unmatched comparison ...")
    config = CpeltdmConfig(pe,min_cov,matched,bound_check,trim,max_size_subreg,max_size_anal_reg)

    # Check same number of samles is matched
    matched && (length(bams1)!=length(bams2)) && sleep_exit("Unmatched samples found. Exiting ...")

    ## Define analysis regions
    print_log("Defining analysis regions from BED file ...")

    # Divide BED file into analysis regions
    bed_out = get_bed_out_filename(bed)
    gen_anal_reg_file(bed,bed_out,max_size_anal_reg)

    ## Inference

    # Estimate θ, compute MML & NME, and do test for each ROI
    print_log("Running differential analysis ...")
    anal_bed_file(bams1,bams2,bed_out,fasta,out_paths,config)

    ## Done

    # Print done
    print_log("Done.")

    # Return
    return nothing

end


text:
module TestTools

using Test

using BioSequences,
      BioTools.BLAST

import BioCore.Testing:
    get_bio_fmt_specimens

fmtdir = get_bio_fmt_specimens()

if !Sys.iswindows()  # temporarily disable the BLAST tests on Windows (issue: #197)

@testset "BLAST+ blastn" begin
    na1 = dna"""
    CGGACCAGACGGACACAGGGAGAAGCTAGTTTCTTTCATGTGATTGANAT
    NATGACTCTACTCCTAAAAGGGAAAAANCAATATCCTTGTTTACAGAAGA
    GAAACAAACAAGCCCCACTCAGCTCAGTCACAGGAGAGAN
    """
    na2 = dna"""
    CGGAGCCAGCGAGCATATGCTGCATGAGGACCTTTCTATCTTACATTATG
    GCTGGGAATCTTACTCTTTCATCTGATACCTTGTTCAGATTTCAAAATAG
    TTGTAGCCTTATCCTGGTTTTACAGATGTGAAACTTTCAA
    """
    fna = joinpath(fmtdir, "FASTA", "f002.fasta")
    nucldb = joinpath(fmtdir, "BLASTDB", "f002")
    nuclresults = joinpath(fmtdir, "BLASTDB", "f002.xml")

    @test typeof(blastn(na1, na2)) == Array{BLASTResult, 1}
    @test typeof(blastn(na1, [na1, na2])) == Array{BLASTResult, 1}
    @test typeof(blastn([na1, na2], [na1, na2])) == Array{BLASTResult, 1}
    @test typeof(blastn(na1, nucldb, db=true)) == Array{BLASTResult, 1}
    @test typeof(blastn(na1, fna)) == Array{BLASTResult, 1}
    @test typeof(blastn(fna, nucldb, db=true)) == Array{BLASTResult, 1}
    @test typeof(blastn([na1, na2], nucldb, db=true)) == Array{BLASTResult, 1}
    @test typeof(blastn([na1, na2], fna)) == Array{BLASTResult, 1}
    @test typeof(blastn(fna, [na1, na2])) == Array{BLASTResult, 1}
end

@testset "BLAST+ blastp" begin
    aa1 = aa"""
    MWATLPLLCAGAWLLGVPVCGAAELSVNSLEKFHFKSWMSKHRKTYSTEE
    YHHRLQTFASNWRKINAHNNGNHTFKMALNQFSDMSFAEIKHKYLWSEPQ
    NCSATKSNYLRGTGPYPPSVDWRKKGNFVSPVKNQGACGS
    """
    aa2 = aa"""
    MWTALPLLCAGAWLLSAGATAELTVNAIEKFHFTSWMKQHQKTYSSREYS
    HRLQVFANNWRKIQAHNQRNHTFKMGLNQFSDMSFAEIKHKYLWSEPQNC
    SATKSNYLRGTGPYPSSMDWRKKGNVVSPVKNQGACGSCW
    """
    faa = joinpath(fmtdir, "FASTA", "cysprot.fasta")
    protdb = joinpath(fmtdir, "BLASTDB", "cysprot")
    protresults = joinpath(fmtdir, "BLASTDB", "cysprot.xml")

    @test typeof(blastp(aa1, aa2)) == Array{BLASTResult, 1}
    @test typeof(blastp(aa1, [aa1, aa2])) == Array{BLASTResult, 1}
    @test typeof(blastp([aa1, aa2], [aa1, aa2])) == Array{BLASTResult, 1}
    @test typeof(blastp(aa1, protdb, db=true)) == Array{BLASTResult, 1}
    @test typeof(blastp(aa1, faa)) == Array{BLASTResult, 1}
    @test typeof(blastp(faa, protdb, db=true)) == Array{BLASTResult, 1}
    @test typeof(blastp([aa1, aa2], protdb, db=true)) == Array{BLASTResult, 1}
    @test typeof(blastp([aa1, aa2], faa)) == Array{BLASTResult, 1}
    @test typeof(blastp(faa, [aa1, aa2])) == Array{BLASTResult, 1}
end

end  # if !is_windows()

end # TestTools


text:
## Preloads
module DHC_2DUtils

    using FFTW
    using LinearAlgebra
    using Statistics

    export fink_filter_bank
    export fink_filter_list
    export fink_filter_hash
    export DHC_compute

    function fink_filter_bank(c, L; nx=256, wd=1, pc=1, shift=false)
        #c     - sets the scale sampling rate (1 is dyadic, 2 is half dyadic)
        #L     - number of angular bins (usually 8*pc or 16*pc)
        #wd    - width of the wavelets (default 1, wd=2 for a double covering)
        #pc    - plane coverage (default 1, full 2pi 2)
        #shift - shift in θ by 1/2 of the θ spacing

        # -------- set parameters
        dθ   = pc*π/L
        wdθ  = wd*dθ
        θ_sh = shift ? dθ/2 : 0.0
        dx   = nx/2-1
        norm = 1.0/sqrt(wd)

        im_scale = convert(Int8,log2(nx))
        # -------- number of bins in radial direction (size scales)
        J = (im_scale-2)*c

        # -------- allocate output array of zeros
        filt      = zeros(nx, nx, J*L+1)
        psi_index = zeros(Int32, J, L)
        theta     = zeros(Float64, L)
        j_value   = zeros(Float64, J)

        # -------- allocate theta and logr arrays
        θ    = zeros(nx, nx)
        logr = zeros(nx, nx)

        # -------- loop over l
        for l = 0:L-1
            θ_l        = dθ*l+θ_sh
            theta[l+1] = θ_l

        # -------- allocate anggood BitArray
            anggood = falses(nx, nx)

        # -------- loop over pixels
            for x = 1:nx
                sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
                for y = 1:nx
                    sy = mod(y+dx,nx)-dx -1
                    θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                    θ_good = abs(θ_pix-π) <= wdθ

                    # If this is a pixel we might use, calculate log2(r)
                    if θ_good
                        anggood[y, x] = θ_good
                        θ[y, x]       = θ_pix
                        r2            = sx^2 + sy^2
                        logr[y, x]    = 0.5*log2(max(1,r2))
                    end
                end
            end
            angmask = findall(anggood)
        # -------- compute the wavelet in the Fourier domain
        # -------- the angular factor is the same for all j
            F_angular = norm .* cos.((θ[angmask].-π).*(L/(2*wd*pc)))

        # -------- loop over j for the radial part
        #    for (j_ind, j) in enumerate(1/c:1/c:im_scale-2)
            for j_ind = 1:J
                j = j_ind/c
                j_value[j_ind] = j  # store for later
                jrad  = im_scale-j-1
                Δj    = abs.(logr[angmask].-jrad)
                rmask = (Δj .<= 1/c)

        # -------- radial part
                F_radial = cos.(Δj[rmask] .* (c*π/2))
                ind      = angmask[rmask]
        #      Let's have these be (J,L) if you reshape...
        #        f_ind    = (j_ind-1)*L+l+1
                f_ind    = j_ind + l*J
                filt[ind, f_ind] = F_radial .* F_angular[rmask]
                psi_index[j_ind,l+1] = f_ind
            end
        end

        # -------- phi contains power near k=0 not yet accounted for
        filter_power = (sum(filt.*filt, dims=3))[:,:,1]

        # -------- for plane half-covered (pc=1), add other half-plane
        if pc == 1
            filter_power .+= circshift(filter_power[end:-1:1,end:-1:1],(1,1))
        end

        # -------- compute power required to sum to 1.0
        i0 = round(Int16,nx/2-2)
        i1 = round(Int16,nx/2+4)
        center_power = 1.0 .- fftshift(filter_power)[i0:i1,i0:i1]
        zind = findall(center_power .< 1E-15)
        center_power[zind] .= 0.0  # set small numbers to zero
        phi_cen = zeros(nx, nx)
        phi_cen[i0:i1,i0:i1] = sqrt.(center_power)

        # -------- before adding ϕ to filter bank, renormalize ψ if pc=1
        if pc==1 filt .*= sqrt(2.0) end  # double power for half coverage

        # -------- add result to filter array
        phi_index  = J*L+1
        filt[:,:,phi_index] .= fftshift(phi_cen)

        # -------- metadata dictionary
        info=Dict()
        info["npix"]         = nx
        info["j_value"]      = j_value
        info["theta_value"]  = theta
        info["psi_index"]    = psi_index
        info["phi_index"]    = phi_index
        info["pc"]           = pc
        info["wd"]           = wd

        return filt, info
    end

    ## Make a list of non-zero pixels for the (Fourier plane) filters
    function fink_filter_list(filt)
        (ny,nx,Nf) = size(filt)

        # Allocate output arrays
        filtind = fill(CartesianIndex{2}[], Nf)
        filtval = fill(Float64[], Nf)

        # Loop over filters and record non-zero values
        for l=1:Nf
            f = @view filt[:,:,l]
            ind = findall(f .> 1E-13)
            val = f[ind]
            filtind[l] = ind
            filtval[l] = val
        end
        return [filtind, filtval]
    end

    function fink_filter_hash(c, L; nx=256, wd=1, pc=1, shift=false)
        # -------- compute the filter bank
        filt, hash = fink_filter_bank(c, L; nx=nx, wd=wd, pc=pc, shift=shift)

        # -------- list of non-zero pixels
        flist = fink_filter_list(filt)

        # -------- pack everything you need into the info structure
        hash["filt_index"] = flist[1]
        hash["filt_value"] = flist[2]
        return hash
    end

    function DHC_compute(image::Array{Float64,2}, filter_hash, filter_hash2)
        # Use 2 threads for FFT
        FFTW.set_num_threads(2)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate coeff arrays
        out_coeff = []
        S0  = zeros(Float64, 2)
        S1  = zeros(Float64, Nf)
        S20 = zeros(Float64, Nf, Nf)
        S12 = zeros(Float64, Nf, Nf)
        S2  = zeros(Float64, Nf, Nf)  # traditional 2nd order

        # allocate image arrays for internal use
        im_fdf_0_1 = zeros(Float64,           Nx, Ny, Nf)   # this must be zeroed!
        im_rd_0_1  = Array{Float64, 3}(undef, Nx, Ny, Nf)

        ## 0th Order
        S0[1]   = mean(image)
        norm_im = image.-S0[1]
        S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
        norm_im ./= sqrt(Nx*Ny*S0[2])

        append!(out_coeff,S0[:])

        ## 1st Order
        im_fd_0 = fft(norm_im)  # total power=1.0

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        f_ind2   = filter_hash2["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val2   = filter_hash2["filt_value"]

        zarr = zeros(ComplexF64, Nx, Ny)  # temporary array to fill with zvals

        # make a FFTW "plan" for an array of the given size and type
        P   = plan_ifft(im_fd_0)   # P is an operator, P*im is ifft(im)

        ## Main 1st Order and Precompute 2nd Order
        for f = 1:Nf
            S1tot = 0.0
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i
            # for (ind, val) in zip(f_i, f_v)   # this is slower!
            if length(f_i) > 0
                for i = 1:length(f_i)
                    ind       = f_i[i]
                    zval      = f_v[i] * im_fd_0[ind]
                    S1tot    += abs2(zval)
                    zarr[ind] = zval
                    im_fdf_0_1[ind,f] = abs(zval)
                end
                S1[f] = S1tot/(Nx*Ny)  # image power
                im_rd_0_1[:,:,f] .= abs2.(P*zarr)
                zarr[f_i] .= 0
            end
        end
        append!(out_coeff, S1[:])

        # we stored the abs()^2, so take sqrt (this is faster to do all at once)
        im_rd_0_1 .= sqrt.(im_rd_0_1)

        ## 2nd Order
        #Amat = reshape(im_fdf_0_1, Nx*Ny, Nf)
        #S12  = Amat' * Amat
        #Amat = reshape(im_rd_0_1, Nx*Ny, Nf)
        #S20  = Amat' * Amat

        #append!(out_coeff, S20)
        #append!(out_coeff, S12)

        ## Traditional second order
        for f1 = 1:Nf
            thisim = fft(im_rd_0_1[:,:,f1])  # Could try rfft here
            # println("  f1",f1,"  sum(fft):",sum(abs2.(thisim))/Nx^2, "  sum(im): ",sum(abs2.(im_rd_0_1[:,:,f1])))
            # Loop over f2 and do second-order convolution
            for f2 = 1:Nf
                f_i = f_ind2[f2]  # CartesianIndex list for filter
                f_v = f_val2[f2]  # Values for f_i
                # sum im^2 = sum(|fft|^2/npix)
                S2[f1,f2] = sum(abs2.(f_v .* thisim[f_i]))/(Nx*Ny)
            end
        end
        append!(out_coeff, S2)

        return out_coeff
    end

end # of module

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
using Random, LinearAlgebra, BlockArrays
using BlockArrays: cumulsizes, _BlockArray, BlockSizes

@timedtestset "dense" begin

    # Test construction of a BlockArray via _BlockArray
    @testset "_BlockArray" begin
        rng, P, Q, R = MersenneTwister(123456), 7, 13, 11
        @testset "Block Vector" begin
            x, blk_sizes = randn(rng, P), BlockSizes([3, 2, 2])
            ȳ = randn(rng, P)
            x_blks = BlockArray(x, blk_sizes).blocks
            ȳ_blks = BlockArray(ȳ, blk_sizes).blocks

            y, back = Zygote.forward(_BlockArray, x_blks, blk_sizes)
            @test y == _BlockArray(x_blks, blk_sizes)
            @test first(back(ȳ)) == ȳ_blks
            @test last(back(ȳ)) === nothing
            @test first(back(ȳ)) == first(back(BlockArray(ȳ, blk_sizes)))
            @test first(back(ȳ)) == first(back((blocks=ȳ_blks, block_sizes=nothing)))
        end
        @testset "Block Array 3" begin
            X, blk_sizes = randn(rng, P, Q, R), BlockSizes([3, 4], [7, 6], [2, 3, 6])
            Ȳ = randn(rng, P, Q, R)
            X_blks = BlockArray(X, blk_sizes).blocks
            Ȳ_blks = BlockArray(Ȳ, blk_sizes).blocks

            Y, back = Zygote.forward(_BlockArray, X_blks, blk_sizes)
            @test Y == _BlockArray(X_blks, blk_sizes)
            @test first(back(Ȳ)) == Ȳ_blks
            @test last(back(Ȳ)) === nothing
            @test first(back(Ȳ)) == first(back(BlockArray(Ȳ, blk_sizes)))
            @test first(back(Ȳ)) == first(back((blocks=Ȳ_blks, block_sizes=nothing)))
        end
    end

    @testset "Vector(::BlockVector)" begin
        rng, Ps = MersenneTwister(123456), [5, 6, 7]
        x = BlockArray(randn(rng, sum(Ps)), Ps)
        adjoint_test(Vector, randn(rng, sum(Ps)), x)
    end

    @testset "Matrix(::BlockMatrix)" begin
        rng, Ps, Qs = MersenneTwister(123456), [3, 4, 5], [6, 7, 8, 9]
        X = BlockArray(randn(rng, sum(Ps), sum(Qs)), Ps, Qs)
        adjoint_test(Matrix, randn(rng, sum(Ps), sum(Qs)), X)
    end

    # # Test construction of `BlockVector` from a vector of vectors. Also test copying.
    # @testset "BlockVector" begin
    #     rng, N, N′ = MersenneTwister(123456), 5, 6
    #     x, x′ = randn(rng, N), randn(rng, N′)
    #     x̂ = BlockVector([x, x′])
    #     @test length(x̂) == N + N′
    #     @test x̂ == vcat(x, x′)
    #     @test getblock(x̂, 1) == x && getblock(x̂, 2) == x′

    #     # Test BlockVector construction adjoint.
    #     adjoint_test(x->BlockVector([x]), BlockVector([randn(rng, N)]), x)
    #     adjoint_test(x->BlockVector([x]), randn(rng, N), x)

    #     adjoint_test(
    #         x->BlockVector([x, x′]),
    #         BlockVector([randn(rng, N), randn(rng, N′)]),
    #         x,
    #     )
    #     adjoint_test(x->BlockVector([x, x′]), randn(rng, N + N′), x)

    #     # Test BlockVector conversion to Vector adjoint.
    #     adjoint_test(Vector, randn(rng, N), BlockVector([x]))
    #     adjoint_test(Vector, randn(rng, N + N′), x̂)

    #     # zero of a BlockVector should be a BlockVector
    #     @test zero(x̂) isa BlockVector
    #     @test cumulsizes(x̂) == cumulsizes(zero(x̂))
    #     @test zero(x̂) == zero(Vector(x̂))
    # end

    # # Test construction of `BlockMatrix` from matrix of matrices. Also test copying.
    # @testset "BlockMatrix" begin
    #     rng, P1, P2, P3, Q1, Q2 = MersenneTwister(123456), 2, 3, 6, 4, 5
    #     X11, X12 = randn(rng, P1, Q1), randn(rng, P1, Q2)
    #     X21, X22 = randn(rng, P2, Q1), randn(rng, P2, Q2)
    #     X31, X32 = randn(rng, P3, Q1), randn(rng, P3, Q2)
    #     X = BlockMatrix(reshape([X11, X21, X31, X12, X22, X32], 3, 2))
    #     @test size(X, 1) == P1 + P2 + P3
    #     @test size(X, 2) == Q1 + Q2
    #     @test getblock(X, 1, 1) == X11
    #     @test getblock(X, 1, 2) == X12
    #     @test getblock(X, 2, 1) == X21
    #     @test getblock(X, 2, 2) == X22

    #     @test cumulsizes(X, 1) == cumsum([1, P1, P2, P3])
    #     @test cumulsizes(X, 2) == cumsum([1, Q1, Q2])

    #     @test blocksizes(X, 1) == [P1, P2, P3]
    #     @test blocksizes(X, 2) == [Q1, Q2]

    #     @test BlockMatrix([X11, X21, X31, X12, X22, X32], 3, 2) == X

    #     Y = BlockMatrix([X11, X21, X31])
    #     @test size(Y, 1) == P1 + P2 + P3
    #     @test size(Y, 2) == Q1
    #     @test getblock(Y, 1, 1) == X11
    #     @test getblock(Y, 2, 1) == X21
    #     @test cumulsizes(Y, 1) == cumsum([1, P1, P2, P3])
    #     @test cumulsizes(Y, 2) == cumsum([1, Q1])

    #     @test blocksizes(Y, 1) == [P1, P2, P3]
    #     @test blocksizes(Y, 2) == [Q1]

    #     P, Q = P1 + P2 + P3, Q1 + Q2
    #     adjoint_test(
    #         X11->BlockMatrix(reshape([X11, X21, X31, X12, X22, X32], 3, 2)),
    #         randn(rng, P, Q),
    #         X11,
    #     )
    #     adjoint_test(
    #         X11->BlockMatrix(reshape([X11, X21, X31, X12, X22, X32], 3, 2)),
    #         X,
    #         X11,
    #     )
    #     adjoint_test(X11->BlockMatrix([X11, X21]), randn(rng, P1 + P2, Q1), X11)
    #     adjoint_test(X11->BlockMatrix([X11, X21]), BlockMatrix([X11, X21]), X11)
    #     adjoint_test(X11->BlockMatrix([X11, X12], 1, 2), BlockMatrix([X11, X12], 1, 2), X11)

    #     # BlockMatrix conversion to Matrix.
    #     adjoint_test(Matrix, randn(rng, size(X)), X)

    #     # zero of a BlockMatrix should be a BlockMatrix.
    #     @test zero(X) isa BlockMatrix
    #     @test cumulsizes(X) == cumulsizes(zero(X))
    #     @test zero(X) == zero(Matrix(X))
    # end

    # @testset "dense, transpose, and adjoint mul" begin
    #     rng, Ps, Qs, Rs = MersenneTwister(123456), [2, 3], [4, 5], [6, 7]
    #     X = _BlockArray([randn(rng, P, Q) for P in Ps, Q in Qs], Ps, Qs)

    #     for foo in [transpose, adjoint]
    #         @test foo(X) isa BlockMatrix
    #         @test blocksizes(foo(X), 1) == blocksizes(X, 2)
    #         @test blocksizes(foo(X), 2) == blocksizes(X, 1)
    #         @test foo(Matrix(X)) == Matrix(foo(X))
    #         @test foo(foo(X)) == X
    #     end

    #     y = _BlockArray([randn(rng, Q) for Q in Qs], Qs)
    #     z = _BlockArray([randn(rng, P) for P in Ps], Ps)
    #     dense_BlockMatrix_BlockVector_mul_tests(rng, X, y)
    #     dense_BlockMatrix_BlockVector_mul_tests(rng, X', z)

    #     Y = _BlockArray([randn(rng, Q, R) for Q in Qs, R in Rs], Qs, Rs)

    #     Xt_blocks = Matrix.(transpose(X.blocks))
    #     Xt = _BlockArray(Xt_blocks, Qs, Ps)
    #     Yt_blocks = Matrix.(transpose(Y.blocks))
    #     Yt = _BlockArray(Yt_blocks, Rs, Qs)

    #     dense_BlockMatrix_BlockMatrix_mul_tests(rng, X, Y)
    #     for foo in [transpose, adjoint]
    #         dense_BlockMatrix_BlockMatrix_mul_tests(rng, foo(Xt), Y)
    #         dense_BlockMatrix_BlockMatrix_mul_tests(rng, X, foo(Yt))
    #         dense_BlockMatrix_BlockMatrix_mul_tests(rng, foo(Xt), foo(Yt))
    #     end
    # end

    # @testset "dense cholesky" begin
    #     rng, Ps = MersenneTwister(123456), [3, 4, 5]
    #     tmp = randn(rng, sum(Ps), sum(Ps))
    #     A_ = tmp * tmp' + 1e-9I
    #     A = BlockArray(A_, Ps, Ps)
    #     @assert A_ == A

    #     # Compute cholesky and compare.
    #     C_, C = cholesky(A_), cholesky(A)
    #     @test C isa Cholesky
    #     @test C_.U ≈ Matrix(C.U)
    #     @test C.U isa UpperTriangular{T, <:AbstractBlockMatrix{T}} where T

    #     function cholesky_gradients(C̄_factors, A_, A)

    #         # Compute cholesky adjoint with Matrix.
    #         C̄_dense = (factors=C̄_factors, uplo=nothing, info=nothing)
    #         _, back_dense = Zygote.forward(cholesky, A_)
    #         Ā_ = first(back_dense(C̄))

    #         # Compute cholesky adjoint with BlockMatrix.
    #         C̄_block = (factors=BlockArray(C̄_factors, blocksizes(A)...),)
    #         _, back_block = Zygote.forward(cholesky, A)
    #         Ā = first(back_block(C̄_block))

    #         # Ensure approximate agreement.
    #         @test Ā_ ≈ Ā
    #     end

    #     cholesky_gradients(randn(rng, size(A)), A_, A)
    #     cholesky_gradients(UpperTriangular(randn(rng, size(A))), A_, A)

    #     # Test `logdet`.
    #     @test logdet(C) ≈ logdet(C_)

    #     # Check that the forwards-pass agrees.
    #     @test Zygote.forward(logdet, C)[1] == logdet(C)

    #     # Check that reverse-pass agrees with dense revese-pass.
    #     @testset "logdet gradients" begin
    #         ȳ = randn(rng)

    #         # Compute adjoint with dense.
    #         _, back_dense = Zygote.forward(logdet, C_)
    #         C̄_ = back_dense(ȳ)

    #         # Compute adjoint with block
    #         _, back_block = Zygote.forward(logdet, C)
    #         C̄ = back_block(ȳ)

    #         # Check that both answers approximately agree.
    #         @test C̄_[1].factors ≈ C̄[1].factors
    #     end

    #     @testset "cholesky backsolving" begin
    #         x = BlockVector([randn(rng, P) for P in Ps])
    #         X = BlockMatrix([randn(rng, P, Q) for P in Ps, Q in Qs])
    #         @test Vector(ldiv!(C, copy(x))) ≈ ldiv!(C_, Vector(x))
    #         @test Matrix(ldiv!(C, copy(X))) ≈ ldiv!(C_, Matrix(X))
    #     end
    # end
end


text:
include("input.jl")

input = split(input, "\n")

gamma = ""
epsilon = ""

for bit in 1:length(input[1])
    num0 = 0
    num1 = 0
    for byte in input
        if byte[bit] == '1'
            num1+=1
        else
            num0+=1
        end
    end
    global gamma *= string(Int(num0>num1))
    global epsilon *= string(Int(num0<num1))

end

println(parse(Int,gamma,base=2)*parse(Int,epsilon,base=2))


text:
# https://adventofcode.com/2021/day/16
using AdventOfCode
hex2bin(x)=lpad.(string.(x |>hex2bytes, base=2),8,"0")|>join
input = hex2bin(readlines("2021/data/day_16.txt")|>first)
bin2Int(x)=parse(Int,x,base=2)
function readliteral!(str,head)
    head[]+=5
    if bin2Int(str[head[]-5])==0
        return(str[head[]-4:head[]-1])
    end
    join([str[head[]-4:head[]-1];readliteral!(str,head)])
end
function readpacket!(str,head=fill(1))
    ver=bin2Int(str[head[]:head[]+2])
    head[]+=3
    typ=bin2Int(str[head[]:head[]+2])
    head[]+=3
    if typ == 4 # literal
        return((ver=ver,typ=typ,content=readliteral!(str,head)|>bin2Int))
    end
    lentype=ifelse(str[head[]]=='0', 15, 11) # length type 15 or 11
    head[]+=1
    len=bin2Int(str[head[]:head[]+lentype-1])
    head[]+=lentype
    origin=head[]
    subpackets=[]
    while((lentype==15&&head[]<origin+len) || (lentype==11 &&length(subpackets)<len))
        push!(subpackets,readpacket!(str,head))
    end
    return((ver=ver,typ=typ,content=subpackets))
end
function sumver(packet)
    if packet.typ==4 return(packet.ver) end
    sum(sumver,(packet.content))+packet.ver
end
function evalpacket(packet)
    if packet.typ==4 return(packet.content) end
    if packet.typ==0 return(sum(evalpacket.(packet.content))) end
    if packet.typ==1 return(prod(evalpacket.(packet.content))) end
    if packet.typ==2 return(minimum(evalpacket.(packet.content))) end
    if packet.typ==3 return(maximum(evalpacket.(packet.content))) end
    if packet.typ==5 return(>(evalpacket.(packet.content)...)) end
    if packet.typ==6 return(<(evalpacket.(packet.content)...)|>Int) end
    if packet.typ==7 return(==(evalpacket.(packet.content)...)|>Int) end
end
function part_1(input)
    sumver(readpacket!(input))
end

function part_2(input)
    evalpacket(readpacket!(input))
end


text:
# functions for mating

# mating between two groups
"""
    pregnant_dams = mating!(mgroup::PTGroup, fgroup::PTGroup, progeny="fgroup"; ...)

Perform mating between sires from `mgroup` and dams from `fgroup`.
The progeny will be in a group defined by `progeny`; `"mgroup"` for the male (sire) group or `"fgroup"` for the female (dam) group.
The pregnant dams (stored in the population database) will be excluded from the mating plan.
The generation of progeny is defined as the largest generation number of parents (in `progeny`) + 1.
It returns a list of dams that have been pregnant in this mating opportinity.
The following options are available.

- `n`: `"all"` for mating as many females as possible in particular `method` and `plan`. Or, you can limit the number of females for mating.
- `prob`: `nothing` for the equal use of sires (default). If you want to choose a sire with a certain probability, provide `prob` with the probability for each sire, e.g. `[0.1,0.2,0.7]` for 3 sires.
- `calving`: `true` for mating and calving at the same time in this function (but still returning `pregnant_dams`). By default, it is `false` to keep pregnancy for dams (use `calving!` to clear the pregnancy status).
- `method`: mating strategy. The default is `"dairy_standard_ai"`, in which a cow has only one chance to mate one of the sires.
- `plan`: `"once_per_female"` for using each dam once (default), or `"once_per_male"` for using each sire once.
- `pmp`: proportion of progeny male (default=0.5).
- `year`: birth year of progeny.
"""
function mating!(mgroup::PTGroup, fgroup::PTGroup, progeny="fgroup"; prob::Union{Nothing,Vector{Float64}}=nothing,
   n::Union{Int,String}="all", year::Int=0, method::String="dairy_standard_ai", plan::String="once_per_female",
   pmp::Float64=0.5, calving::Bool=false, updateinb::Bool=true, verbose::Bool=false)

   if length(mgroup.sires)<1 || length(fgroup.dams)<1
      error("No sires or no dams")
   end
   if !isnothing(prob)
      if length(mgroup.sires) != length(prob)
         throw(ArgumentError("length mismatch: mgroup.sires and prob"))
      end
   end

   if progeny=="mgroup"
      group = mgroup
      maxGen = maximum( group.generation[findall(x->in(x,group.sires),group.id)] )
   elseif progeny=="fgroup"
      group = fgroup
      maxGen = maximum( group.generation[findall(x->in(x,group.dams),group.id)] )
   else
      error("invalid argument progeny=$(progeny)")
   end

   sgroup = mgroup.groupid
   dgroup = fgroup.groupid
   non_pregnant_dams = selectid([:pregnant] => x->x==false, fgroup, idlist=fgroup.dams)
   if typeof(n) == String
      if n=="all"
         ndams = length(non_pregnant_dams)
      else
         error("n: to be `all`, or integer value")
      end
   else
      ndams = min(n,length(non_pregnant_dams))
   end
   if verbose
      if length(non_pregnant_dams)<1
         @warn "There is no non-pregnant dams in this group (gr=$(group.gid))."
      else
         println("Mating gr$(mgroup.gid) x gr$(fgroup.gid); non-pregnant dams:$(length(non_pregnant_dams)) out of $(length(fgroup.dams))")
      end
   end
   if method=="dairy_standard_ai"
      pregnant_dams = _mating_dairy_standard_ai!(group, mgroup.sires, non_pregnant_dams, sgroup, dgroup, maxGen+1, pmp, method="random", plan=plan, year=year, n=ndams, prob=prob)
   else
      error("unknown method $(method) for mating")
   end

   # keep the pregnancy status
   if !calving
      fgroup.pop.df[pregnant_dams,:pregnant] .= true
   end

   # update inbreeding
   if updateinb
      update_inbreeding!(group.pop.df.sire, group.pop.df.dam, group.pop.df.inb)
   end

   return pregnant_dams
end

# standard dairy AI mating
#  - litter size = 1
#  - one cow receives one service; producing n progeny from n dams.
#  - no action if the cow has already been serviced in the same year.
function _mating_dairy_standard_ai!(group::PTGroup, sires::Vector{Int}, dams::Vector{Int}, sgroup::Int, dgroup::Int, gen::Int, pmp::Float64;
   method::String="random", plan::String="once_per_female", year::Int=0, n::Int=length(dams), prob::Union{Nothing,Vector{Float64}}=nothing)
   var_a = get_var_poly(group.pop.par)
   var_g = get_var_qtl(group.pop.par)
   var_g_sim = group.pop.par.var_g_sim
   nai = 0
   ndams = length(dams)
   # non-pregnant dam list
   damlist = copy(dams)
   # pregnant dams
   pregnant_dams_id = Dict{Int,Bool}()
   if method=="random"
      shuffle!(damlist)
      if plan=="once_per_male"
         @inbounds for s in sires
            if s>0
               if isempty(damlist)
                  break
               end
               d = pop!(damlist)
               pregnant_dams_id[d] = true
               male = rand_male_calf(pmp)
               generate_and_add_animal!(group, s, d, sgroup, dgroup, male, gen, year, var_a, var_g, var_g_sim)
               nai = nai + 1
            end
            if nai>=n
               break
            end
         end
      elseif plan=="once_per_female"
         @inbounds for d in damlist
            if d>0
               pregnant_dams_id[d] = true
               if isnothing(prob)
                  s = rand(sires)
               else
                  s = sample(sires,Weights(prob))
               end
               male = rand_male_calf(pmp)
               generate_and_add_animal!(group, s, d, sgroup, dgroup, male, gen, year, var_a, var_g, var_g_sim)
               nai = nai + 1
            end
            if nai>=n
               break
            end
         end
      else
         error("unknown mating plan: $(plan)")
      end
   else
      error("unknown mating method: $(method)")
   end
   return Int.(keys(pregnant_dams_id))
end

function rand_male_calf(pmp::Float64)
   if rand()<=pmp
      return true
   else
      return false
   end
end

function generate_and_add_animal!(group::PTGroup, s::Int, d::Int, sgroup::Int, dgroup::Int, male::Bool, gen::Int, year::Int, var_a::Float64, var_g::Float64, var_g_sim::Float64)
   pop = group.pop

   # pedigree-based (polygenic) breeding valur
   pa = (pop.df[s,:pbv] + pop.df[d,:pbv])/2
   ms = get_ms_deviation(pop.df.inb[s], pop.df.inb[d], var_a)
   pbv = pa + ms

   # QTL-based breeding value
   if pop.df[s,:genotyped] && pop.df[d,:genotyped]
      gs = read_qmsim_individual_hdf5(pop.map,pop.gfile,s)
      gd = read_qmsim_individual_hdf5(pop.map,pop.gfile,d)
      ga = mating(pop.map,gs,gd)
      qbv = sqrt(var_g/var_g_sim)*ga.tbv
      hassnp = true
   else
      qbv = sqrt(var_g)*randn()  # just a random number
      hassnp = false
   end

   # total breeding values
   tbv = pbv + qbv

   # for future use
   cg = zeros(Int,0)
   y = zeros(Float64,0)
   pe = zeros(Float64,0)
   e = zeros(Float64,0)
   animal = PTAnimal(cg,y,pe,e)

   # add to population
   add_new_animal!(pop, male, animal, year=year, sire=s, dam=d, siregroup=sgroup, damgroup=dgroup, inb=0.0, pbv=pbv, qbv=qbv, tbv=tbv, ebv=0.0, gebv=missing, genotyped=hassnp)
   if hassnp
      add_qmsim_individual_hdf5(pop.map,pop.gfile,ga)
   end

   # add to group
   group.n = group.n + 1
   push!(group.id, pop.maxAnimal)
   push!(group.generation, gen)

   return nothing
end

function get_ms_deviation(inbs::Float64, inbd::Float64, var_a::Float64, k::Float64=0.0)
   sd_ms = sqrt( (1-k) * ( 0.5*(1 - 0.5*(inbs+inbd)) ) * var_a )
   return sd_ms*randn()
end

function update_polygenic_ms!(pop::PTPopulation, id::Int, k::Float64=0.0)
   s = pop.df[i,:sire]
   d = pop.df[i,:dam]
   if s==0
      us = 0.0
      inbs = 0.0
   else
      us = pop.df[s,:pbv]
      inbs = pop.df[s,:inb]
   end
   if d==0
      ud = 0.0
      inbd = 0.0
   else
      ud = pop.df[d,:pbv]
      inbd = pop.df[d,:inb]
   end
   pa = (us + ud)/2
   ms = get_ms_deviation(inbs, inbd, var_a, k)
   pop.df[id,:pbv] = pa + ms
end

function update_polygenic_ms!(pop::PTPopulation, idlist::Vector{Int}, k::Float64=0.0)
   for id in idlist
      update_polygenic_ms!(pop, id, k)
   end
end

# calving
"""
    calving(pop::PTPopulation)
    calving(group::PTGroup)

Change the pregnancy status for females from "pregnant" to "non-pregnant".
"""
function calving!(pop::PTPopulation)
   pop.df[:,:pregnant] .= false
   return nothing
end
function calving!(group::PTGroup)
   group.pop.df[group.dams,:pregnant] .= false
   return nothing
end
function calving!(groups::Vector{PTGroup})
   for group in groups
      calving!(group)
   end
   return nothing
end


text:
module MemPool

using Serialization, Sockets, Random
import Serialization: serialize, deserialize
export DRef, FileRef, poolset, poolget, pooldelete, destroyonevict,
       movetodisk, copytodisk, savetodisk, mmwrite, mmread, cleanup,
       deletefromdisk

## Wrapping-unwrapping of payloads:

struct MMWrap # wrap an object to use custom
                 # memory-mapping/fast serializer
    x::Any
end

# Wrapping is an implementation detail unwrap_payload will be called by poolget
# on every fetched value this will give back the value as it was saved.

unwrap_payload(x) = x
unwrap_payload(x::Tuple) = map(unwrap_payload, x)
unwrap_payload(x::MMWrap) = unwrap_payload(x.x)

function serialize(io::AbstractSerializer, w::MMWrap)
    # ticket to mm land
    mmwrite(io, w.x)
end

function deserialize(io::AbstractSerializer, T::Type{MMWrap})
    MMWrap(mmread(T, io)) # gotta keep that wrapper on
end

# Unwrapping FileRef == reading the file.
# This allows communicating only the file name across processors,
# the receiver process will simply read from file while unwrapping
struct FileRef
    host::IPAddr
    file::String
    size::UInt
    force_pid::Ref{Union{Int, Nothing}}

    function FileRef(file, size; pid=nothing)
        new(host, file, size, Ref{Union{Int, Nothing}}(pid))
    end
end

unwrap_payload(f::FileRef) = unwrap_payload(open(deserialize, f.file, "r+"))

include("io.jl")
include("datastore.jl")

"""
`approx_size(d)`

Returns the size of `d` in bytes used for accounting in MemPool datastore.
"""
function approx_size(@nospecialize(d))
    Base.summarysize(d) # note: this is accurate but expensive
end

function approx_size(d::Union{Base.BitInteger, Float16, Float32, Float64})
    sizeof(d)
end

function approx_size(d::AbstractDict{K,V}) where {K,V}
    N = length(d)
    ksz = approx_size(K, N, keys(d))
    vsz = approx_size(V, N, values(d))
    ksz + vsz + 8*N
end

function approx_size(d::AbstractArray{T}) where T
    isempty(d) && return 0
    isbitstype(T) && return sizeof(d)
    approx_size(T, length(d), d)
end

function approx_size(T, L, d)
    fl = fixedlength(T)
    if fl > 0
        return L * fl
    elseif T === Any
        return L * 64  # approximation (override with a more specific method where exact calculation is needed)
    elseif isempty(d)
        return 0
    else
        return sum(approx_size(x) for x in d)
    end
end

function approx_size(xs::AbstractArray{String})
    # doesn't check for redundant references, but
    # really super fast in comparison to summarysize
    s = 0
    for x in xs
        s += sizeof(x)
    end
    s + 4 * length(xs)
end

function __init__()
    global session = "sess-" * randstring(5)
    try
        global host = getipaddr()
    catch err
        global host = Sockets.localhost
    end
    datastore_lock[] = ReentrantLock()
    id_counter[] = Atomic{Int}(0)
end

end # module


text:
import NaiveNASlib

@testset "Sugar" begin

    struct ScaleByTwo
        f
    end
    (s::ScaleByTwo)(x...) = 2s.f(x...)

    @testset "Create input" begin
        @test issame(inputvertex("input", 5), InputSizeVertex("input", 5))
    end

    @testset "Create immutable" begin
        v = immutablevertex(x -> x * [1 2 3; 4 5 6], 2, inputvertex("in", 3))
        @test nin(v) == [3]
        @test nout(v) == 2
        @test v([1 2]) == [9  12  15]
        v = immutablevertex(identity, 1, inputvertex("in", 1), mutation=IoSize, traitdecoration = t -> NamedTrait(t, "v"))
        @test name(v) == "v"
    end

    @testset "Create absorbing" begin
        v = absorbvertex(x -> x * [1 2 3; 4 5 6], 2, inputvertex("in", 3))
        @test nin(v) == [3]
        @test nout(v) == 2
        @test v([1 2]) == [9  12  15]
        v = absorbvertex(identity, 1, inputvertex("in", 1), mutation=IoSize, traitdecoration = t -> NamedTrait(t, "v"))
        @test name(v) == "v"
    end

    @testset "Create invariant" begin
        v = invariantvertex(x -> 2 .* x, inputvertex("in", 2))
        @test nin(v) == [2]
        @test nout(v) == 2
        @test v([1 2]) == [2 4]
        v = invariantvertex(identity, inputvertex("in", 1), mutation=IoSize, traitdecoration = t -> NamedTrait(t, "v"))
        @test name(v) == "v"
    end

    @testset "Create stacking" begin
        v = conc(inputvertex.(("in1", "in2"), (1,2))..., dims=1, traitdecoration = t -> NamedTrait(t, "v"), outwrap = ScaleByTwo)
        @test nin(v) == [1 ,2]
        @test nout(v) == 3
        @test v(1, [2, 3]) == [2,4,6] # Due to outwrap = ScaleByTwo
        @test name(v) == "v"
        @test typeof(v.base.base.computation) == ScaleByTwo
    end

    @testset "Create element wise" begin

        conf = traitconf(t -> NamedTrait(t, "v"))

        @testset "Create elemwise addition" begin
            v = conf >> inputvertex("in1", 2) + inputvertex("in2", 2) + inputvertex("in3" ,2)

            @test nin(v) == [2, 2, 2]
            @test nout(v) == 2
            @test v([1 2], [3 4], [5 6]) == [9 12]
            @test v([1, 2], [3, 4], [5, 6]) == [9, 12]
            @test v([2], [3, 4], [5, 6]) == [10, 12]
            @test name.(inputs(v)) == ["in1", "in2", "in3"]

            @test_throws DimensionMismatch inputvertex("in1", 2) + inputvertex("in2", 3)

            v = inputvertex("in1", 2) + inputvertex("in2", 2)
            @test v([1, 2], [3, 4]) == [4, 6]

            v = outwrapconf(f -> (x...) -> 2f((x...))) >> inputvertex("in1", 2) + inputvertex("in2", 2)
            @test v([1, 2], [3, 4]) == [8, 12]
        end

        @testset "Create elemwise multiplication" begin
            v = conf >> inputvertex("in1", 2) * inputvertex("in2", 2) * inputvertex("in3" ,2)

            @test nin(v) == [2, 2, 2]
            @test nout(v) == 2
            @test v([1 2], [3 4], [5 6]) == [15 48]
            @test v([1, 2], [3, 4], [5, 6]) == [15, 48]
            @test v([2], [3,4], [5,6]) == [30, 48]
            @test name.(inputs(v)) == ["in1", "in2", "in3"]

            @test_throws DimensionMismatch inputvertex("in1", 2) * inputvertex("in2", 3)

            v = inputvertex("in1", 2) * inputvertex("in2", 2)
            @test v([1, 2], [3, 4]) == [3, 8]
        end
        @testset "Create elemwise subtraction" begin
            v = conf >> inputvertex("in1", 2) - inputvertex("in2", 2)

            @test nin(v) == [2, 2]
            @test nout(v) == 2
            @test v([1 2], [3 4]) == [-2 -2]
            @test v([1, 2], [3, 4]) == [-2, -2]
            @test v([2], [3,4]) == [-1, -2]
            @test name.(inputs(v)) == ["in1", "in2"]

            @test_throws DimensionMismatch inputvertex("in1", 2) - inputvertex("in2", 3)

            v = inputvertex("in1", 2) - inputvertex("in2", 2)
            @test v([1, 2], [3, 4]) == [-2, -2]
        end

        @testset "Create unary subtraction" begin
            v = -(conf >> inputvertex("in1", 2))

            @test nin(v) == [nout(v)] == [2]
            @test v([1 2]) == [-1 -2]
            @test v([1, 2]) == [-1, -2]
            @test name.(inputs(v)) == ["in1"]

            v = -inputvertex("in1", 2)
            @test v([1, 2]) == [-1, -2]
        end

        @testset "Create elemwise division" begin
            v = conf >> inputvertex("in1", 2) / inputvertex("in2", 2)

            @test nin(v) == [2, 2]
            @test nout(v) == 2
            @test v([6 8], [2 4]) == [3 2]
            @test v([6, 8], [2, 4]) == [3, 2]
            @test v([12], [3, 4]) == [4, 3]
            @test name.(inputs(v)) == ["in1", "in2"]

            @test_throws DimensionMismatch inputvertex("in1", 2) / inputvertex("in2", 3)

            v = inputvertex("in1", 2) / inputvertex("in2", 2)
            @test v([6, 8], [2, 4]) == [3, 2]
        end

        @testset "Named elementwise" begin
            v = "eladd" >> inputvertex("in1", 2) + inputvertex("in2", 2)

            @test name(v) == "eladd"
            @test nin(v) == [2, 2]
            @test nout(v) == 2

            v = -("unsub" >> inputvertex("in1", 2))
            @test name(v) == "unsub"
            @test nin(v) == [2]
            @test nout(v) == 2
        end

        @testset "Outwrap shortcut" begin
            scale(f) = ScaleByTwo(f)
            v = scale >> inputvertex("in1", 2) + inputvertex("in2", 2)

            @test nin(v) == [2, 2]
            @test nout(v) == 2

            @test v([1, 2], [3, 4]) == [8, 12]
            @test typeof(v.base.base.computation) == ScaleByTwo
        end
    end
end


text:
# -*- coding: utf-8 -*-
# + {}
module WiringDiagrams

using Catlab.WiringDiagrams
using Catlab.Theories
using Catlab.Graphics
using Catlab.Graphics.Graphviz
using MacroTools
import MacroTools: postwalk, striplines

export wirenames, label!, odeTemplate, drawhom, canonical, label!

"""    drawhom(hom, name::String, format="svg")

draw a hom expression as a wiring diagram and store it as a file.
Defaults to SVG format. The filename is <name.format>.
"""
function drawhom(hom, name::String, format="svg")
    d = to_wiring_diagram(hom)
    g = to_graphviz(d, orientation=LeftToRight)
    t = Graphics.Graphviz.run_graphviz(g, format=format)
    write("$name.$format", t)
    return g
end

"""    canonical(Syntax::Module, hom)

canonicalizes a hom expression by converting it to a WiringDiagram and then back again.
The Syntax parameter is the Doctrine for the morphism such as FreeSymmetricMonoidalCategory.
"""
canonical(Syntax::Module, hom) = begin d = to_wiring_diagram(hom)
    to_hom_expr(Syntax, d)
end

wirenames(d::WiringDiagram) = foldr(union,
    map(box->union(input_ports(box), output_ports(box)),
        boxes(d)))

function label!(g::Graphviz.Graph, v::Vector{String})
    iter = Iterators.Stateful(v)
    map(enumerate(g.stmts)) do (i,s)
        if typeof(s) <: Edge
            g.stmts[i].attrs[:label] = popfirst!(iter)
        end
    end
    return g
end

# CODE TO CONVERT WIRING DIAGRAM TO ODE

function fluxes(d::WiringDiagram)
    # TODO design Multiple Dispatch Lens API
    # TODO use ModelingToolkit variables
    nb = nboxes(d)
    vars = wirenames(d)
    byvar = Dict{Symbol, Expr}()
    homnames = Vector{Symbol}()
    for var in vars
        byvar[var] = :(+())
    end
    map(enumerate(boxes(d))) do (i, box)
        invars = input_ports(box)
        outvars = output_ports(box)
        homname = box.value
        push!(homnames, homname)
        βᵢ = :(p.$(homname))
        ϕ =  :(*($βᵢ, $(invars...)))
        map(invars) do v
            push!(byvar[v].args, :(-$ϕ))
        end
        map(outvars) do v
            push!(byvar[v].args, :($ϕ))
        end
    end
    return byvar, vars, homnames
end

function oderhs(dudt::Dict{T, Expr}, vars::Vector{T}, homnames::Any) where {T}
    vardefs = map(enumerate(vars)) do (i, v)
        :($(v) = u[$(i)])
    end
    lines = map(enumerate(vars)) do (i, v)
        ϕ = dudt[v]
        #:(du.$(ϕ.first) = $(ϕ.second))
        :(du[$(i)] = $(ϕ))
    end
    fdef = quote
    function f(du, u, p, t)
        $(vardefs...)
        $(lines...)
        return du
    end
    end
    return fdef, vars, homnames
end

"""    oderhs(d::WiringDiagram)

convert a wiring diagram into a dynamical system described as julia Exprs.
We need to keep track of the names we give the variables and the parameters so those are the second and third arguments respectively.

Returns an expression that defines a function, a list of symbols representing the variables, and a list of symbols representing the parameter names.

see also fluxes(d::WiringDiagram) for implementation details
"""
oderhs(d::WiringDiagram) = oderhs(fluxes(d)...)

"""    odeTemplate(d::WiringDiagram)

create an expression that defines a code that solves the ODE.
Given just the wiring diagram, we don't know the paramters, initial conditions, or timedomain,
so they are passed in as arguments to the function we generate.

These parameters and initial conditions are destructured in the main function so you can
see what the code is expecting to receive by reading the generated output.

The structure of the timedomain is not implied by the wiring diagram so it is passed directly to the
ODEProblem constructor. Any keyword arguments you pass to `main()` are forwarded to `solve()`.

"""
function odeTemplate(d::WiringDiagram)
    f, vars, homnames = oderhs(d)
    params = Expr(:tuple, map(enumerate(homnames)) do (i, name)
        Expr(:(=), name, :(β[$i]))
            end...)
    vars = wirenames(d)
    initials = Expr(:vect, map(vars) do v
            :(i₀.$v)
            end...)
    out = quote
        using DifferentialEquations
        $(f.args[end])
        function main(β, i₀, tdomain; kwargs...)
            params = $params
            initials = $initials
            prob = ODEProblem(f, params, initials, tdomain)
            soln = solve(prob; kwargs...)
            return prob, soln
        end
    end
    out |> x -> postwalk(x) do x
        return striplines(x)
    end
end

# ADD CATEGORY THEORY REWRITE RULES FOR WIRING DIAGRAMS

end


text:
using SchumakerSpline
x = Array{Union{Missing,Float64}}(collect(0.0:0.01:2.0))
y = Array{Union{Missing,Float64}}(sqrt.(x))
y[5] = missing
x[2] = missing
s1 = Schumaker{Float64}(x, y)

# This should plot something over [0,2]
plt = plot(s1)
# This should plot something over [0,1] but only two points because it is input as an array and this is the exact array of plotpoints.
plt = plot(s1, [0.0,1.0])
# Now it is a tuple and so there will be grid_len intermediate points put it (200 by default)
plt = plot(s1, (0.0,1.0))
# And we can plot the derivatives too
plt = plot(s1, (0.0,1.0); derivs = true)
# This should add a second spline just below with same derivaitves
glt = plot(s1 -10.0, [0.5,1.0]; plot_options = (label = "shifted",), deriv_plot_options = (label = "shifted deriv1",), deriv2_plot_options = (label = "shifted deriv2",), plt = plt)
# This should only plot the first spline.
qlt = plot(s1 + 0.2, (0.0,1.0); plot_options = (label = "shifted",))

ss = Array{Schumaker,1}([s1, s1+0.2, s1*0.8+0.2])
plt  = plot(ss, [0.0,1.0])

plt = plot(ss)
plt2 = plot(0-2.0*s1 + 0.1, [0.0,1.0]; derivs = false, plot_options = (label = "new one",), plt = plt)

# I can't really test anything programatically here. Just have to eyeball it.
# So just giving a true so it all runs happily.
true


text:
using  Compat: isbitstype
import Base: write

"""
    write(f::FortranFile, items...)
    write(f::FortranFile, rec=N, items...)

Write a data record to a `FortranFile`. Each `item` should be a scalar
of a Fortran-compatible datatype (e.g. `Int32`, `Float64`, `FString{10}`),
or an array of such scalars. If no `item`s are given, an empty record is
written. Returns the number of bytes written, **not** including the space
taken up by the record markers.

For direct-access files, the number of the record to be written must be
specified with the `rec` keyword (N=1 for the first record).
"""
function write(f::FortranFile, items...)
   # how much data to write?
   towrite = sizeof_vars(items)
   record = Record(f, towrite)
   result = fwrite(record, items...)
   close(record)
   return result
end

function write(f::FortranFile{DirectAccess}, items...; rec::Integer=0)
   if rec==0
      fthrow("direct-access files require specifying the record to be written (use rec keyword argument)")
   end
   towrite = sizeof_vars(items)
   record = Record(f, rec, towrite)
   result = fwrite(record, items...)
   close(record)
   return result
end

function fwrite( rec::Record )
   return 0
end

function fwrite( rec::Record, vars... )
   written = sum( write_var(rec,var) for var in vars )
   return written
end

# workarounds for "does not support byte I/O"
function write_var( rec::Record, var::Int8 )
   write_var( rec, [var] )
end

function write_var( rec::Record, arr::Array{Int8,N} ) where {N}
   write(rec, arr)
end

# write scalars
function write_var( rec::Record, var::T ) where {T}
   write( rec, rec.convert.onwrite(var) )
end

# write arrays
function write_var( rec::Record, arr::Array{T,N} ) where {T,N}
   written = 0
   for x in arr
      written += write(rec, rec.convert.onwrite(x))
   end
   return written
end

# write strings: delegate to data field
write_var( rec::Record, var::FString ) = write_var(rec, var.data)
write_var( rec::Record, arr::Array{FString{L},N} ) where {L,N} = write_fstrings(rec, arr)
write_fstrings( rec::Record, arr::Array{FString{L},N} ) where {L,N} = sum( write_var(rec, var.data) for var in arr )

# specialized versions for no byte-order conversion
write_var( rec::RecordWithSubrecords{NOCONV}, arr::Array{T,N} ) where {T,N} = write(rec, arr)
write_var( rec::RecordWithSubrecords{NOCONV}, arr::Array{Int8,N} ) where {N} = write(rec, arr)
write_var( rec::RecordWithoutSubrecords{R,NOCONV}, arr::Array{T,N} ) where {T,N,R} = write(rec, arr)
write_var( rec::RecordWithoutSubrecords{R,NOCONV}, arr::Array{Int8,N} ) where {N,R} = write(rec, arr)

# resolve ambiguities
write_var( rec::RecordWithSubrecords{NOCONV}, arr::Array{FString{L},N} ) where {L,N} = write_fstrings(rec, arr)
write_var( rec::RecordWithoutSubrecords{R,NOCONV}, arr::Array{FString{L},N} ) where {L,N,R} = write_fstrings(rec, arr)

# check for type compatibility with Fortran
check_fortran_type(x::Array) = _check_fortran_type(eltype(x))
check_fortran_type(x)        = _check_fortran_type(typeof(x))
_check_fortran_type(::Type{FString{L}}) where L = true
_check_fortran_type(T::Type) = isbitstype(T)

function sizeof_var( var::T ) where {T}
   check_fortran_type(var) ||
      fthrow("cannot serialize datatype $T for Fortran")
   sizeof(var)
end

sizeof_vars(vars) = isempty(vars) ? 0 : sum( sizeof_var(var) for var in vars )