# Copyright (c) 2020: Matthew Wilhelm & Matthew Stuber.
# This work is licensed under the Creative Commons Attribution-NonCommercial-
# ShareAlike 4.0 International License. To view a copy of this license, visit
# http://creativecommons.org/licenses/by-nc-sa/4.0/ or send a letter to Creative
# Commons, PO Box 1866, Mountain View, CA 94042, USA.
#############################################################################
# Dynamic Bounds - pODEs Discrete
# A package for discretize and relax methods for bounding pODEs.
# See https://github.com/PSORLab/DynamicBoundspODEsDiscrete.jl
#############################################################################
# src/DiscretizeRelax/utilities/relax.jl
# Defines the relax! routine for the DiscretizeRelax integrator.
#############################################################################

function DBB.relax!(d::DiscretizeRelax{M,T,S,F,K,X,NY}) where {M <: AbstractStateContractor, T <: Number, S <: Real, F, K, X, NY}

    fill!(d.time, 0.0)

    storage_len = length(d.storage)
    for i=1:(1000-storage_len)
        push!(d.storage, zeros(T, d.nx))
        push!(d.storage_apriori, zeros(T, d.nx))
        push!(d.time, 0.0)
    end

    # reset relax!
    for i = 1:length(d.storage)
        fill!(d.storage[i], zero(T))
    end
    for i = 1:length(d.storage_apriori)
        fill!(d.storage_apriori[i], zero(T))
    end
    empty!(d.relax_t_dict_indx)
    empty!(d.relax_t_dict_flt)
    d.step_result.time = 0.0

    set_P!(d)         # Functor set P and P - p values for calculations
    compute_X0!(d)     # Compute initial condition values

    # Get initial time and integration direction
    t = d.tspan[1]
    tmax = d.tspan[2]
    sign_tstep = copysign(1, tmax - t)
    d.time[1] = t
    d.contractor_result.times[1] = t

    # Computes maximum step size to take (either hit support or max time)
    support_indx = 1
    next_support = Inf
    tsupports = d.tsupports
    if !isempty(tsupports)
        if (tsupports[1] == 0.0)
            support_indx += 1
        end
        next_support = tsupports[support_indx]
    end

    # initialize QR type storage
    set_Δ!(d.contractor_result.Δ, d.storage)
    fill!(d.contractor_result.A_Q[1], 0.0)
    fill!(d.contractor_result.A_inv[1], 0.0)
    for i = 1:d.nx
        d.contractor_result.A_Q[1][i,i] = 1.0
        d.contractor_result.A_inv[1][i,i] = 1.0
    end

    d.exist_result.status_flag = RELAXATION_NOT_CALLED

    # Begin integration loop
    hlast = 0.0
    d.step_count = 0
    is_adaptive = d.exist_result.hj <= 0.0
    d.exist_result.hj = !is_adaptive ? d.exist_result.hj : 0.01*(tmax - d.contractor_result.times[1])
    d.exist_result.predicted_hj = d.exist_result.hj
    stored_value_count = length(d.relax_t_dict_indx)
    for step_number = 2:(d.step_limit+2)
        if (sign_tstep*d.step_result.time <= sign_tstep*tmax) &&
           (d.exist_result.predicted_hj != 0.0)

            delT = abs(sign_tstep*d.step_result.time - sign_tstep*tmax)

            # max step size is min of predicted, when next support point occurs,
            # or the last time step in the span
            tv = d.step_result.time
            d.exist_result.hj = min(d.exist_result.hj, next_support - tv, tmax - tv)
            hj_limit = min(next_support - tv, tmax - tv)
            d.exist_result.hj_max = tmax - tv
            is_support_pnt = (hj_limit == next_support - tv)

            d.contractor_result.steps[1] = d.exist_result.hj
            d.contractor_result.step_count = step_number

            # perform step size calculation and update bound information
            single_step!(d.exist_result, d.contractor_result, d.step_params,
                         d.step_result, d.method_f!, step_number-1, hj_limit,
                         delT, d.constant_state_bounds, d.polyhedral_constraint)

            # unpack storage
            if step_number - 1 > length(d.time)
                push!(d.storage, copy(d.contractor_result.X_computed))
                push!(d.storage_apriori, copy(d.exist_result.Xj_apriori))
                push!(d.time, d.contractor_result.times[1])
            else
                copyto!(d.storage[step_number], d.contractor_result.X_computed)
                copyto!(d.storage_apriori[step_number], d.exist_result.Xj_apriori)
                d.time[step_number] = d.step_result.time
                d.contractor_result.times[1] = d.step_result.time
            end
            if is_support_pnt
                stored_value_count += 1
                d.relax_t_dict_indx[stored_value_count] = step_number
                d.relax_t_dict_flt[next_support] = step_number
                support_indx += 1
                if support_indx <= length(tsupports)
                    next_support = tsupports[support_indx]
                else
                    next_support = Inf
                end
            end

            # throw error if limit exceeded
            if d.exist_result.status_flag !== RELAXATION_NOT_CALLED
                d.error_code = d.exist_result.status_flag
                break
            end
            d.step_count += 1
        else
            break
        end
    end

    if (d.step_count > d.step_limit) && (sign_tstep*d.time[d.step_count + 1] < sign_tstep*tmax)
        d.error_code = LIMIT_EXCEEDED
    end

    # cut out any unnecessary array elements
    resize!(d.storage, d.step_count + 1)
    resize!(d.storage_apriori, d.step_count + 1)
    resize!(d.time, d.step_count + 1)

    if d.error_code === RELAXATION_NOT_CALLED
        d.error_code = COMPLETED
    end

    return nothing
end


@testset "Constraint Stamp" begin

    # CStamp
    N = 10
    p = 4
    s0 = CStamp(:v, :col, 1, 2, 3)
    @test valid(s0, N, p)
    s0 = CStamp(:v, :col, 1, 1, 3)
    @test !valid(s0, N, p)
    s0 = CStamp(:v, :col, 1, 3, 3)
    @test valid(s0, N, p)
    s0 = CStamp(:v, :col, 1, 5, 3)
    @test !valid(s0, N, p)
    s0 = CStamp(:v, :col, 0, 3, 3)
    @test !valid(s0, N, p)
    s0 = CStamp(:v, :col, 1, 3, 1)
    @test !valid(s0, N, p)
    s0 = CStamp(:v, :col, 3, 1, 3)
    @test !valid(s0, N, p)
    s0 = CStamp(:v, :col, 1, 2, 11)
    @test !valid(s0, N, p)

    s0 = CStamp(:h, :col, 1, 3, 3)
    @test valid(s0, N, p)
    s0 = CStamp(:h, :col, 2, 2, 3)
    @test !valid(s0, N, p)
    s0 = CStamp(:h, :col, 3, 1, 3)
    @test valid(s0, N, p)
    s0 = CStamp(:h, :col, 3, 1, 1)
    @test !valid(s0, N, p)
    s0 = CStamp(:h, :col, 3, 1, 11)
    @test !valid(s0, N, p)
    s0 = CStamp(:h, :col, 5, 1, 10)
    @test !valid(s0, N, p)
    s0 = CStamp(:h, :col, 4, 1, 10)
    @test valid(s0, N, p)

    s1 = CStamp()
    @test s1 != s0
    @test !(s1 == s0)

    s2 = stampify(:v,:col,1,2,10)
    stampify!(s1,:v,:col,1,2,10)
    @test s2 == s1
    @test isequal(s2, s1)
    @test hash(s2) == hash(s1)

end


using RecipesBase

@recipe function plot(f_ts::TimeSeries, ii::Vararg)
    tmp = get_coords(f_ts,ii...)
    f   = f_ts[ii...]

    # we want to select only those that are *not* Numbers (ie, Arrays or Ranges)
    idx = .!isa.(tmp, Number)
    idx = [idx...]
    coord = [tmp...]

    xx = coord[idx]

    xx_f = [xx..., f]

    match_dimensions := true

    (xx_f...,)
end


using Collatz
using Test

@testset "Collatz.jl" begin
    @test collatz(0) == 0
    @test collatz(1) == 0
    @test collatz(9) == 19
    @test collatz(97) == 118
    @test collatz(871) == 178
    @test collatz(6171) == 261
    @test collatz(77031) == 350
    @test collatz(837799) == 524
    @test collatz(8400511) == 685
    @test collatz(63728127) == 949
    @test collatz(670617279) == 986
    @test collatz(9780657630) == 1132
    @test collatz(75128138247) == 1228
    @test collatz(989345275647) == 1348
    @test collatz(Int128(7887663552367)) == 1563
    @test collatz(80867137596217) == 1662
    @test collatz(942488749153153) == 1862
    @test collatz(Int128(7579309213675935)) == 1958
    @test collatz(Int128(93571393692802302)) == 2091
end


# ---------------------------------------------------------------------------- #
#
#   mesher.jl
#
#   Abstract mesher type
#   This allows for writing an n-dimensional version of mesher methods
#
#   λυτέος
#   Fall 2017
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

"""
    Mesher

Mesher abstract type:
Overarching abstract type for meshers (BAMG, FEFLOA).
"""
abstract type Mesher

end

include("mesherBAMG.jl")
include("mesherFEFLOA.jl")

include("metricBAMG.jl")
include("metricFEFLOA.jl")


using DiffEqJump, DiffEqBase, SafeTestsets

@time begin
  @time @safetestset "Constant Rate Tests" begin include("constant_rate.jl") end
  @time @safetestset "Variable Rate Tests" begin include("variable_rate.jl") end
  @time @safetestset "Monte Carlo Tests" begin include("monte_carlo_test.jl") end
  @time @safetestset "Split Coupled Tests" begin include("splitcoupled.jl") end
  @time @safetestset "SSA Tests" begin include("ssa_tests.jl") end
  @time @safetestset "Tau Leaping Tests" begin include("regular_jumps.jl") end
  @time @safetestset "Simple SSA Callback Test" begin include("ssa_callback_test.jl") end
  @time @safetestset "SIR Discrete Callback Test" begin include("sir_model.jl") end
  @time @safetestset "Linear Reaction SSA Test" begin include("linearreaction_test.jl") end
  @time @safetestset "Mass Action Jump Tests; Gene Expr Model" begin include("geneexpr_test.jl") end
  @time @safetestset "Mass Action Jump Tests; Nonlinear Rx Model" begin include("bimolerx_test.jl") end
  @time @safetestset "Mass Action Jump Tests; Special Cases" begin include("degenerate_rx_cases.jl") end
  @time @safetestset "Composition-Rejection Table Tests" begin include("table_test.jl") end
  @time @safetestset "Extinction test" begin include("extinction_test.jl") end
  @time @safetestset "Saveat Regression test" begin include("saveat_regression.jl") end
  @time @safetestset "Ensemble Uniqueness test" begin include("ensemble_uniqueness.jl") end
  @time @safetestset "Thread Safety test" begin include("thread_safety.jl") end
  @time @safetestset "A + B <--> C" begin include("reversible_binding.jl") end
  @time @safetestset "Remake tests" begin include("remake_test.jl") end
  @time @safetestset "Long time accuracy test" begin include("longtimes_test.jl") end
  @time @safetestset "Spatial utilities" begin include("spatial/utils_test.jl") end
  @time @safetestset "Spatial A + B <--> C" begin include("spatial/ABC.jl") end
  @time @safetestset "Pure diffusion" begin include("spatial/diffusion.jl") end
end


# This file contains a number of additions to BlockArrays.jl. These are completely
# independent of Stheno.jl, and will (hopefully) move over to BlockArrays.jl at some point.

using BlockArrays: cumulsizes, blocksizes, blocksizes, BlockSizes

import Base: +, *, size, getindex, eltype, copy, \, vec, getproperty, zero
import LinearAlgebra: UpperTriangular, LowerTriangular, logdet, Symmetric, transpose,
    adjoint, AdjOrTrans, AdjOrTransAbsMat, cholesky!, logdet, ldiv!, mul!, logabsdet
import BlockArrays: BlockArray, BlockVector, BlockMatrix, BlockVecOrMat, getblock,
    blocksize, setblock!, nblocks, getblock!

# Do some character saving.
const BV{T} = BlockVector{T}
const BM{T} = BlockMatrix{T}
const ABV{T} = AbstractBlockVector{T}
const ABM{T} = AbstractBlockMatrix{T}

const BlockCholesky{T, V} = Cholesky{T, <:BlockMatrix{T, V}}

const BlockArrayAdjoint = NamedTuple{(:blocks, :block_sizes)}

####################################### Various util #######################################

import BlockArrays: BlockVector, BlockMatrix, blocksizes, cumulsizes, AbstractBlockSizes
export BlockVector, BlockMatrix, blocksizes, blocklengths

@adjoint function _BlockArray(
    blocks::R,
    block_sizes::BS,
) where {T, N, R<:AbstractArray{<:AbstractArray{T,N},N}, BS<:AbstractBlockSizes{N}}
    back(Δ::NamedTuple{(:blocks, :block_sizes)}) = (Δ.blocks, nothing)
    back(Δ::BlockArray) = (Δ.blocks, nothing)
    back(Δ::AbstractArray) = (BlockArray(Δ, block_sizes).blocks, nothing)
    return _BlockArray(blocks, block_sizes), back
end

@adjoint Vector(x::BlockVector) = Vector(x), Δ::Vector->(BlockArray(Δ, blocksizes(x)),)
@adjoint function Matrix(X::BlockMatrix)
    back(Δ::Matrix) = (BlockArray(Δ, blocksizes(X)),)
    back(Δ::UpperTriangular{T, <:Matrix{T}} where {T}) = back(Matrix(Δ))
    return Matrix(X), back
end

# Tell Zygote to ignore some stuff that doesn't involve gradient info.
_get_block_sizes(xs::Vector) = (length.(xs),)
_get_block_sizes(Xs) = (size.(Xs[:, 1], 1), size.(Xs[1, :], 2))
@nograd _get_block_sizes

"""
    blocksizes(X::AbstractArray, d::Int)

Get a vector containing the block sizes over the `d`th dimension of `X`.
"""
blocksizes(X::AbstractArray, d::Int) = diff(cumulsizes(X, d))
@nograd blocksizes

"""
    BlockVector(xs::AbstractVector{<:AbstractVector})

Construct a `BlockVector` from a collection of `AbstractVector`s.
"""
BlockVector(xs::AV{<:AV}) = _BlockArray(xs, _get_block_sizes(xs)...)

"""
    BlockMatrix(Xs::Matrix{<:AbstractVecOrMat})

Construct a `BlockMatrix` from a matrix of `AbstractVecOrMat`s.
"""
BlockMatrix(Xs::Matrix{<:AbstractVecOrMat}) = _BlockArray(Xs, _get_block_sizes(Xs)...)

"""
    are_conformal(A::BlockVecOrMat, B::BlockVecOrMat)

Test whether two block matrices (or vectors) are conformal. This criterion is stricter than
that for general matrices / vectors as we additionally require that each block be conformal
with block of the other matrix with which it will be multiplied. This ensures that the
result is itself straightforwardly representable as `BlockVecOrMat`.
"""
are_conformal(A::AVM, B::AVM) = cumulsizes(A, 2) == cumulsizes(B, 1)

# BlockMatrix(Xs::Vector{<:AbstractVecOrMat}) = BlockMatrix(reshape(Xs, length(Xs), 1))
# BlockMatrix(x::AbstractVecOrMat) = BlockMatrix([x])

# """
#     BlockMatrix(xs::Vector{<:AM}, P::Int, Q::Int)

# Construct a block matrix with `P` rows and `Q` columns of blocks.
# """
# BlockMatrix(xs::Vector{<:AM}, P::Int, Q::Int) = BlockMatrix(reshape(xs, P, Q))

# zero(x::AbstractBlockVector) = BlockVector([zero(getblock(x, n)) for n in 1:nblocks(x, 1)])
# function zero(X::AbstractBlockMatrix)
#     blocks = [zero(getblock(X, p, q)) for p in 1:nblocks(X, 1), q in 1:nblocks(X, 2)]
#     return BlockMatrix(blocks)
# end

# # Convert from a BlockVector to a single-column BlockMatrix
# function _to_block_matrix(x::BlockVector)
#     return _BlockArray(reshape(reshape.(x.blocks, :, 1), :, 1), blocksizes(x, 1), [1])
# end
# @adjoint function _to_block_matrix(x::BlockVector)
#     X = _to_block_matrix(x)
#     back(Δ::BlockMatrix) = (_to_block_vector(Δ),)
#     back(Δ::AbstractMatrix) = back(BlockArray(Δ), blocksizes(X)...)
#     return X, back
# end

# # Convert a single-column BlockMatrix to a BlockVector.
# function _to_block_vector(X::BlockMatrix)
#     @assert size(X, 2) == 1
#     return _BlockArray(reshape(reshape.(X.blocks, :), :), blocksizes(X, 1))
# end
# @adjoint function _to_block_vector(X::BlockMatrix)
#     x = _to_block_vector(X)
#     back(Δ::BlockVector) = (_to_block_matrix(Δ),)
#     back(Δ::AbstractVector) = back(BlockArray(Δ), blocksizes(x)...)
#     return x, back
# end

# #
# # Dense BlockArrays
# #

# *(a::BlockVector, b::BlockMatrix) = _to_block_matrix(a) * b

# function *(A::BlockMatrix{T}, x::BlockVector{V}) where {T, V}
#     Ps = blocksizes(A, 1)
#     y = _BlockArray([Vector{promote_type(T, V)}(undef, P) for P in Ps], Ps)
#     return mul!(y, A, x)
# end

# @adjoint *(A::BlockMatrix, x::BlockVector) = A * x, Δ::BlockVector -> (Δ * x', A' * Δ)

# function mul!(y::BlockVector, A::BlockMatrix, x::BlockVector)
#     @assert are_conformal(A, x) && are_conformal(A', y)
#     for r in 1:nblocks(A, 1)
#         mul!(y[Block(r)], A[Block(r, 1)], x[Block(1)])
#         for c in 2:nblocks(A, 2)
#             y[Block(r)] = y[Block(r)] + A[Block(r, c)] * x[Block(c)]
#         end
#     end
#     return y
# end

# function *(A::BlockMatrix{T}, B::BlockMatrix{V}) where {T, V}
#     Ps, Qs = blocksizes(A, 1), blocksizes(B, 2)
#     C = _BlockArray([Matrix{promote_type(T, V)}(undef, P, Q) for P in Ps, Q in Qs], Ps, Qs)
#     return mul!(C, A, B)
# end

# @adjoint *(A::BlockMatrix, B::BlockMatrix) = A * B, Δ::BlockMatrix->(Δ * B', A' * Δ)

# function mul!(C::BlockMatrix, A::BlockMatrix, B::BlockMatrix)
#     @assert are_conformal(A, B)
#     @assert cumulsizes(A, 1) == cumulsizes(C, 1)
#     @assert cumulsizes(C, 2) == cumulsizes(B, 2)

#     P, Q, R = nblocks(A, 1), nblocks(A, 2), nblocks(B, 2)
#     for p in 1:P, r in 1:R
#         mul!(C[Block(p, r)], A[Block(p, 1)], B[Block(1, r)])
#         for q in 2:Q
#             C[Block(p, r)] = C[Block(p, r)] + A[Block(p, q)] * B[Block(q, r)]
#         end
#     end
#     return C
# end

# function cholesky_and_adjoint(A::BlockMatrix{T, V}) where {T, V}
#     U = BlockMatrix{T, V}(undef_blocks, blocksizes(A, 1), blocksizes(A, 2))
#     D = BlockMatrix{T, V}(undef_blocks, blocksizes(A, 1), blocksizes(A, 2))
#     backs = Vector{Any}(undef, nblocks(A, 2))

#     # Do an initial pass to fill each of the blocks with zeros. This is cheap.
#     for q in 1:nblocks(U, 2), p in 1:nblocks(U, 1)
#         U[Block(p, q)] = zeros{T}(blocksize(A, (p, q))...)
#     end

#     # Fill out the upper triangle with the Cholesky.
#     for j in 1:nblocks(A, 2)

#         # Update off-diagonals.
#         for i in 1:j-1
#             # U[Block(i, j)] = copy(A[Block(i, j)]
#             D[Block(i, j)] = copy(A[Block(i, j)])
#             for k in 1:i-1
#                 D[Block(i, j)] -= U[Block(k, i)]' * U[Block(k, j)]
#             end
#             U[Block(i, j)] = U[Block(i, i)]' \ D[Block(i, j)]
#             # ldiv!(UpperTriangular(U[Block(i, i)])', U[Block(i, j)])
#         end

#         # Update diagonal.
#         D[Block(j, j)] = copy(A[Block(j, j)])
#         for k in 1:j-1
#             D[Block(j, j)] -= U[Block(k, j)]' * U[Block(k, j)]
#         end
#         U_tmp, back = Zygote.forward(A->cholesky(A).U, D[Block(j, j)])
#         U[Block(j, j)] = U_tmp
#         backs[j] = back
#         # U[Block(j, j)] = cholesky(U[Block(j, j)]).U
#     end

#     _to_block(Ū::AbstractMatrix) = BlockArray(copy(Ū), blocksizes(A))
#     _to_block(Ū::AbstractBlockMatrix) = copy(Ū)
#     function _to_block(Ū::BlockUpperTriangular)
#         return copy(Ū.data)
#     end

#     return Cholesky(U, :U, 0), function(Δ::NamedTuple)

#         Ū = _to_block(Δ.factors)
#         Ā = BlockMatrix{T, V}(undef_blocks, blocksizes(A, 1), blocksizes(A, 1))

#         for j in reverse(1:nblocks(A, 2))

#             Ā[Block(j, j)] = first(backs[j](Ū[Block(j, j)]))
#             for k in reverse(1:j-1)
#                 Ū[Block(k, j)] -= U[Block(k, j)] * (Ā[Block(j, j)] + Ā[Block(j, j)]')
#             end

#             for i in reverse(1:j-1)
#                 Ā[Block(i, j)] = U[Block(i, i)] \ Ū[Block(i, j)]
#                 Ū[Block(i, i)] -= U[Block(i, j)] * Ā[Block(i, j)]'
#                 for k in reverse(1:i-1)
#                     Ū[Block(k, i)] -= U[Block(k, j)] * Ā[Block(i, j)]'
#                     Ū[Block(k, j)] -= U[Block(k, i)] * Ā[Block(i, j)]
#                 end
#             end
#         end

#         # Zero-out the lower triangle.
#         for q in 1:nblocks(A, 1)
#             for p in q+1:nblocks(A, 2)
#                 Ā[Block(p, q)] = zeros(size(A[Block(p, q)]))
#             end
#         end
#         return (Ā,)
#     end
# end

# cholesky(A::BlockMatrix{<:Real}) = first(cholesky_and_adjoint(A))
# @adjoint cholesky(A::BlockMatrix{<:Real}) = cholesky_and_adjoint(A)

# cholesky(A::Symmetric{T, <:BlockMatrix{T}} where {T<:Real}) = cholesky(A.data)

# # A slightly strange util function that shouldn't ever be used outside of `logdet`.
# reduce_diag(f, A::Matrix{<:Real}) = sum(f, view(A, diagind(A)))

# function reduce_diag(f, A::BlockMatrix{<:Real})
#     return sum([reduce_diag(f, getblock(A, n, n)) for n in 1:nblocks(A, 1)])
# end

# logdet(C::BlockCholesky{<:Real}) = 2 * reduce_diag(log, C.factors)

# @adjoint function logdet(C::BlockCholesky{<:Real})
#     return logdet(C), function(Δ::Real)
#         function update_diag!(X::Matrix, A::Matrix)
#             X[diagind(X)] .= (2 * Δ) ./ A[diagind(A)]
#             return X
#         end
#         function update_diag!(X::BlockMatrix, A::BlockMatrix)
#             for n in 1:nblocks(A)[1]
#                 update_diag!(getblock(X, n, n), getblock(A, n, n))
#             end
#             return X
#         end
#         factors = update_diag!(zero(C.factors), C.factors)
#         return ((factors=factors, uplo=nothing, info=nothing),)
#     end
# end

# ldiv!(C::BlockCholesky{<:Real}, x::BlockVector) = ldiv!(C.U, ldiv!(C.U', x))
# ldiv!(C::BlockCholesky{<:Real}, X::BlockMatrix) = ldiv!(C.U, ldiv!(C.U', X))

# \(C::BlockCholesky{<:Real}, x::BlockVector) = ldiv!(C, copy(x))
# \(C::BlockCholesky{<:Real}, X::BlockMatrix) = ldiv!(C, copy(X))

# #
# # Adjoint and Transpose
# #

# function Base.adjoint(X::BlockMatrix)
#     return _BlockArray(
#         permutedims(collect(adjoint.(X.blocks))),
#         blocksizes(X, 2),
#         blocksizes(X, 1),
#     )
# end
# function Base.adjoint(x::BlockVector)
#     return _BlockArray(
#         permutedims(collect(adjoint.(x.blocks))),
#         [1],
#         blocksizes(x, 1),
#     )
# end

# function Base.transpose(X::BlockMatrix)
#     return _BlockArray(
#         permutedims(collect(transpose.(X.blocks))),
#         blocksizes(X, 2),
#         blocksizes(X, 1),
#     )
# end
# function Base.transpose(x::BlockVector)
#     return _BlockArray(
#         permutedims(collect(transpose.(x.blocks))),
#         [1],
#         blocksizes(x, 1),
#     )
# end

####################################### Multiplication #####################################

# function +(u::UniformScaling, X::AbstractBlockMatrix)
#     @assert cumulsizes(X, 1) == cumulsizes(X, 2)
#     Y = copy(X)
#     for p in 1:nblocks(Y, 1)
#         setblock!(Y, getblock(Y, p, p) + u, p, p)
#     end
#     return Y
# end
# +(u::UniformScaling, X::Symmetric{T, <:ABM{T}} where T) = Symmetric(u + unbox(X))
# function +(X::AbstractBlockMatrix, u::UniformScaling)
#     @assert cumulsizes(X, 1) == cumulsizes(X, 2)
#     Y = copy(X)
#     for p in 1:nblocks(Y, 1)
#         setblock!(Y, u + getblock(Y, p, p), p, p)
#     end
#     return Y
# end
# +(X::Symmetric{T, <:ABM{T}} where T, u::UniformScaling) = Symmetric(unbox(X) + u)

# # Define addition and subtraction for compatible block matrices and vectors.
# import Base: +, -
# for foo in [:+, :-]
#     @eval function $foo(A::BV{T}, B::BV{T}) where T
#         @assert blocksizes(A) == blocksizes(B)
#         C = similar(A)
#         for p in 1:nblocks(C, 1)
#             setblock!(C, $foo(getblock(A, p), getblock(B, p)), p)
#         end
#         return C
#     end
#     @eval function $foo(A::BM{T}, B::BM{T}) where T
#         @assert blocksizes(A) == blocksizes(B)
#         C = similar(A)
#         for q in 1:nblocks(C, 2), p in 1:nblocks(C, 1)
#             setblock!(C, $foo(getblock(A, p, q), getblock(B, p, q)), p, q)
#         end
#         return C
#     end
# end

# #################################### Broadcasting ##########################################
# # Override the usual broadcasting machinery for BlockArrays. This is a pretty disgusting
# # hack. At the time of writing it, I was also writing my first year report, so was short of
# # time. This is generally an open problem for AbstractBlockArrays which really does need to
# # be resolved at some point.

# # Very specific `broadcast` method for particular use case. Needs to be generalised.
# function broadcasted(f, A::AbstractBlockVector, b::Real)
#     return BlockVector([broadcast(f, getblock(A, p), b) for p in 1:nblocks(A, 1)])
# end
# function broadcasted(f, A::AbstractBlockMatrix, b::Real)
#     return BlockMatrix([broadcast(f, getblock(A, p, q), b)
#         for p in 1:nblocks(A, 1), q in 1:nblocks(A, 2)])
# end


module ShallowWaters

    export run_model, Parameter, ∂x, ∂y, Ix, Iy, ∇²

    using NetCDF, Parameters, Printf, Dates, Interpolations

    include("default_parameters.jl")
    include("grid.jl")
    include("constants.jl")
    include("forcing.jl")
    include("model_setup.jl")
    include("initial_conditions.jl")
    include("preallocate.jl")

    include("time_integration.jl")
    include("ghost_points.jl")
    include("rhs.jl")
    include("gradients.jl")
    include("interpolations.jl")
    include("PV_advection.jl")
    include("continuity.jl")
    include("bottom_drag.jl")
    include("diffusion.jl")
    include("tracer_advection.jl")

    include("feedback.jl")
    include("output.jl")
    include("run_model.jl")

end


using Oceananigans.Fields: FieldSlicer

#####
##### JLD2OutputWriter tests
#####

function jld2_field_output(model)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u = (x, y, z) -> rand(),
                v = (x, y, z) -> rand(),
                w = (x, y, z) -> rand(),
                T = 0,
                S = 0)

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)

    simulation.output_writers[:velocities] = JLD2OutputWriter(model, model.velocities,
                                                                    schedule = TimeInterval(1),
                                                                        dir = ".",
                                                                     prefix = "test",
                                                                      force = true)

    u₀ = CUDA.@allowscalar data(model.velocities.u)[3, 3, 3]
    v₀ = CUDA.@allowscalar data(model.velocities.v)[3, 3, 3]
    w₀ = CUDA.@allowscalar data(model.velocities.w)[3, 3, 3]

    run!(simulation)

    file = jldopen("test.jld2")

    # Data is saved without halos by default
    u₁ = file["timeseries/u/0"][3, 3, 3]
    v₁ = file["timeseries/v/0"][3, 3, 3]
    w₁ = file["timeseries/w/0"][3, 3, 3]

    close(file)

    rm("test.jld2")

    FT = typeof(u₁)

    return FT(u₀) == u₁ && FT(v₀) == v₁ && FT(w₀) == w₁
end

function jld2_sliced_field_output(model, outputs=model.velocities)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u = (x, y, z) -> rand(),
                v = (x, y, z) -> rand(),
                w = (x, y, z) -> rand())

    simulation = Simulation(model, Δt=1.0, stop_iteration=1)

    simulation.output_writers[:velocities] =
        JLD2OutputWriter(model, outputs,
                             schedule = TimeInterval(1),
                         field_slicer = FieldSlicer(i=1:2, j=1:2:4, k=:),
                                  dir = ".",
                               prefix = "test",
                                force = true)

    run!(simulation)

    file = jldopen("test.jld2")

    u₁ = file["timeseries/u/0"]
    v₁ = file["timeseries/v/0"]
    w₁ = file["timeseries/w/0"]

    close(file)

    rm("test.jld2")

    return size(u₁) == (2, 2, 4) && size(v₁) == (2, 2, 4) && size(w₁) == (2, 2, 5)
end

function run_jld2_file_splitting_tests(arch)
    model = NonhydrostaticModel(architecture=arch, grid=RegularRectilinearGrid(size=(16, 16, 16), extent=(1, 1, 1)))
    simulation = Simulation(model, Δt=1, stop_iteration=10)

    function fake_bc_init(file, model)
        file["boundary_conditions/fake"] = π
    end

    ow = JLD2OutputWriter(model, (u=model.velocities.u,); dir=".", prefix="test", schedule=IterationInterval(1),
                          init=fake_bc_init, including=[:grid],
                          field_slicer=nothing, array_type=Array{Float64},
                          max_filesize=200KiB, force=true)

    push!(simulation.output_writers, ow)

    # 531 KiB of output will be written which should get split into 3 files.
    run!(simulation)

    # Test that files has been split according to size as expected.
    @test filesize("test_part1.jld2") > 200KiB
    @test filesize("test_part2.jld2") > 200KiB
    @test filesize("test_part3.jld2") < 200KiB

    for n in string.(1:3)
        filename = "test_part$n.jld2"
        jldopen(filename, "r") do file
            # Test to make sure all files contain structs from `including`.
            @test file["grid/Nx"] == 16

            # Test to make sure all files contain info from `init` function.
            @test file["boundary_conditions/fake"] == π
        end

        # Leave test directory clean.
        rm(filename)
    end
end

function jld2_time_averaging_of_horizontal_averages(model)

    model.clock.iteration = 0
    model.clock.time = 0.0

    set!(model, u = (x, y, z) -> 1,
                v = (x, y, z) -> 2,
                w = (x, y, z) -> 0,
                T = (x, y, z) -> 4)

    u, v, w = model.velocities
    T, S = model.tracers

    simulation = Simulation(model, Δt=1.0, stop_iteration=5)

    u, v, w = model.velocities
    T, S = model.tracers

    average_fluxes = (wu = AveragedField(w * u, dims=(1, 2)),
                      uv = AveragedField(u * v, dims=(1, 2)),
                      wT = AveragedField(w * T, dims=(1, 2)))

    simulation.output_writers[:fluxes] = JLD2OutputWriter(model, average_fluxes,
                                                          schedule = AveragedTimeInterval(4, window=2),
                                                               dir = ".",
                                                            prefix = "test",
                                                             force = true)

    run!(simulation)

    file = jldopen("test.jld2")

    # Data is saved without halos by default
    wu = file["timeseries/wu/4"][1, 1, 3]
    uv = file["timeseries/uv/4"][1, 1, 3]
    wT = file["timeseries/wT/4"][1, 1, 3]

    close(file)

    rm("test.jld2")

    FT = eltype(model.grid)

    return wu == zero(FT) && wT == zero(FT) && uv == FT(2)
end

for arch in archs
    # Some tests can reuse this same grid and model.
    topo =(Periodic, Periodic, Bounded)
    grid = RegularRectilinearGrid(topology=topo, size=(4, 4, 4), extent=(1, 1, 1))
    model = NonhydrostaticModel(architecture=arch, grid=grid)

    @testset "JLD2 output writer [$(typeof(arch))]" begin
        @info "  Testing JLD2 output writer [$(typeof(arch))]..."

        @test jld2_field_output(model)
        @test jld2_sliced_field_output(model)
        @test jld2_sliced_field_output(model, (u = model -> model.velocities.u,
                                               v = model -> model.velocities.v,
                                               w = model -> model.velocities.w))

        run_jld2_file_splitting_tests(arch)

        @test jld2_time_averaging_of_horizontal_averages(model)
    end
end


# Implementation taken from p316 RL Textbook and incorporating TB
Base.@kwdef mutable struct ETB{O, T<:AbstractTraceUpdate} <: LearningUpdate
    lambda::Float64
    opt::O
    hdpi::HordeDPI
    clip_threshold::Float64
    normalize_interest::Bool
    trace::T = AccumulatingTraces()
    e::IdDict = IdDict()
    followon::IdDict = IdDict()
    prev_discounts::IdDict = IdDict()
    emphasis_logging::Any = zeros(4) #I'm lazy and this is used for tracking hard to recalculate values
    rho_logging::Any = zeros(4) #I'm lazy and this is used for tracking hard to recalculate values
end

function get_demon_parameters(lu::ETB, learner, demons, obs, state, action, next_obs, next_state, next_action, env_reward)
    # C, next_discounts, _ = get(demons, obs, action, next_obs, next_action)
    C, next_discounts, _ = get(demons; state_t = obs, action_t = action, state_tp1 = next_obs, action_tp1 = next_action, reward = env_reward)
    target_pis = get_demon_pis(demons, learner.num_actions, state, obs)
    next_target_pis = get_demon_pis(demons, learner.num_actions, next_state, next_obs)
    C, next_discounts, target_pis, next_target_pis
end

function update!(lu::ETB,
                 learner::QLearner{M, LU},
                 demons,
                 obs,
                 next_obs,
                 state,
                 action,
                 next_state,
                 next_action,
                 is_terminal,
                 behaviour_pi_func,
                 env_reward) where {M<:AbstractMatrix, LU<:ETB}

    weights = learner.model
    λ = lu.lambda

    behaviour_pis = behaviour_pi_func(state, obs)
    C, next_discounts, target_pis, next_target_pis =
        get_demon_parameters(lu,
                             learner,
                             demons,
                             obs,
                             state,
                             action,
                             next_obs,
                             next_state,
                             next_action,
                             env_reward)

    discounts = get!(()->zero(next_discounts), lu.prev_discounts, learner)::typeof(next_discounts)
    e = get!(()->zero(weights), lu.e, weights)::typeof(weights)
    followon = get!(()->zeros(learner.num_actions), lu.followon, weights)::typeof(zeros(learner.num_actions))

    #Broadcast the policy and pseudotermination of each demon across the actions
    inds = get_action_inds(action, learner.num_actions, learner.num_demons)
    state_action_row_ind = inds

    # interest = get_interest(learner, lu, obs, action)

    # using dpi for interest instead
    interest = lu.hdpi(obs, action)
    if (lu.normalize_interest)
        interest = [ceil(i) for i in interest]
    end

    # getting IS ratio
    if (behaviour_pis[action] == 0)
        ρ = zeros(size(target_pis[:, action]))
    else
        ρ = target_pis[:, action] / behaviour_pis[action]
    end

    # accumulating followon
    followon[:] .+= interest

    # Get Emphasis vector - size is # demons
    # Correcting with ρ in the beginning since it is the action-value emphasis rather than the state-value emphasis
    emphasis = ρ .* (λ * behaviour_pis[action] * interest + (1 - λ * behaviour_pis[action]) * followon)

    # I'm lazy, here we're just logging the emphasis and rho as variables.
    lu.emphasis_logging = emphasis
    lu.rho_logging = ρ

    clamp!(emphasis, 0, lu.clip_threshold)

    # Update eligibility trace
    update_trace!(lu.trace,
                  e,
                  state,
                  λ,
                  repeat(discounts, inner = learner.num_actions),
                  repeat(target_pis[:,action], inner = learner.num_actions),
                  inds;
                  emphasis=emphasis)

    # decaying followon
    followon[:] .*= ρ .* next_discounts

    pred = learner(next_state)
    Qs = reshape(pred, (learner.num_actions, learner.num_demons))'

    # Target Pi is num_demons  x num_actions
    backup_est_per_demon = vec(sum((Qs .* next_target_pis ), dims = 2))
    target = C + next_discounts .* backup_est_per_demon

     # TD error per demon is the td error on Q
    td_err = target - (weights * state)[inds]
    td_err_across_demons = repeat(td_err, inner=learner.num_actions)
    # update discounts
    discounts .= next_discounts

    if lu.opt isa Auto
        next_state_action_row_ind = get_action_inds(next_action, learner.num_actions, learner.num_demons)
        state_discount = zero(e)
        state_discount[state_action_row_ind,:] .+= state'
        state_discount[next_state_action_row_ind,:] .-= next_discounts * next_state'
        abs_phi = abs.(e)
        update!(lu.opt, weights, e, td_err_across_demons, abs_phi .* max.(state_discount, abs_phi), learner.num_demons, learner.num_actions)
    else
        Flux.Optimise.update!(lu.opt, weights,  -(e .* td_err_across_demons))
    end
end

function update!(lu::ETB,
                 learner::Union{SRLearner,GPI},
                 demons,
                 obs,
                 next_obs,
                 state::SparseVector,
                 action,
                 next_state,
                 next_action,
                 is_terminal,
                 behaviour_pi_func,
                 env_reward)

    ψ = learner.ψ
    w = learner.r_w

    C, next_discounts, target_pis, next_target_pis =
        get_demon_parameters(lu,
                             learner,
                             demons,
                             obs,
                             state,
                             action,
                             next_obs,
                             next_state,
                             next_action,
                             env_reward)

    discounts = get!(()->zero(next_discounts), lu.prev_discounts, learner)::typeof(next_discounts)
    e_nz = get!(()->Int[], lu.e, learner)::Array{Int, 1}
    e_w_nz = get!(()->Int[], lu.e, learner)::Array{Int, 1}
    e_ψ = get!(()->zero(ψ), lu.e, ψ)::typeof(ψ)
    e_w = get!(()->zero(w), lu.e, w)::typeof(w)
    λ = lu.lambda

    next_active_state_action = get_active_action_state_vector(next_state, next_action,length(next_state), learner.num_actions)
    active_state_action = get_active_action_state_vector(state, action, length(state), learner.num_actions)

    projected_state = learner.feature_projector(obs, action, next_obs)

    (reward_C, SF_C) = C[1:learner.num_tasks] , C[learner.num_tasks + 1:end]
    (reward_discounts, SF_discounts) = discounts[1:learner.num_tasks], discounts[learner.num_tasks+1:end]
    (reward_next_discounts, SF_next_discounts) = next_discounts[1:learner.num_tasks], next_discounts[learner.num_tasks+1:end]
    (reward_target_pis, SF_target_pis) = target_pis[1:learner.num_tasks,:], target_pis[learner.num_tasks+1:end,:]
    (reward_next_target_pis, SF_next_target_pis) = next_target_pis[1:learner.num_tasks,:], next_target_pis[learner.num_tasks+1:end, :]

    # Followon
    followon = get!(()->zeros(learner.num_demons), lu.followon, weights)::typeof(zeros(learner.num_demons))

    # Setting interest to constant ones for now for each demon
    # interest = get_interest(learner, lu, obs, action)

    demon_interest = lu.hdpi(obs, action)

    if (lu.normalize_interest)
        demon_interest = [ceil(i) for i in demon_interest]
    end

    num_repeat = convert(Integer, (learner.num_demons - learner.num_tasks) / learner.num_tasks)
    if (rem(learner.num_demons - learner.num_tasks, learner.num_tasks) != 0)
        println("The number to repeat for SF demons aren't even with the number of tasks for SR. Something might be wrong :(")
    end

    SF_interest = repeat(demon_interest, inner=num_repeat)
    reward_interest = demon_interest

    interest = vcat(reward_interest, SF_interest)

    # getting IS ratio
    behaviour_pis = behaviour_pi_func(state, obs)
    if (behaviour_pis[action] == 0)
        ρ = zeros(size(target_pis[:, action]))
    else
        ρ = target_pis[:, action] / behaviour_pis[action]
    end
    reward_ρ, SF_ρ = ρ[1:learner.num_tasks], ρ[learner.num_tasks+1:end]

    # accumulating followon
    followon[:] .+= interest

    # Get Emphasis vector - size is # demons
    # Correcting with ρ in the beginning since it is the action-value emphasis rather than the state-value emphasis
    emphasis = ρ .* (λ * behaviour_pis[action] * interest + (1 - λ * behaviour_pis[action]) * followon)

    # I'm lazy, here we're just logging the emphasis and rho as variables.
    lu.emphasis_logging = emphasis
    lu.rho_logging = ρ

    clamp!(emphasis, 0, lu.clip_threshold)
    reward_emphasis, SF_emphasis = emphasis[1:learner.num_tasks], emphasis[learner.num_tasks+1:end]

    # Update Traces: See update_utils.jl
    update_trace!(lu.trace, e_ψ, active_state_action, λ, SF_discounts, SF_target_pis[:, action]; emphasis=SF_emphasis)
    update_trace!(lu.trace, e_w, projected_state, λ, reward_discounts, reward_target_pis[:, action]; emphasis=reward_emphasis)
    # e_nz = e_nz ∪ active_state_action.nzind
    # e_w_nz = e_w_nz ∪ projected_state_action.nzind
    if λ == 0.0
        e_nz = active_state_action.nzind
        e_w_nz = projected_state.nzind
    else
        e_nz = e_nz ∪ active_state_action.nzind
        e_w_nz = e_w_nz ∪ projected_state.nzind
    end

    # decaying followon
    followon[:] .*= ρ .* next_discounts

    pred = ψ * next_active_state_action
    reward_feature_backup = zeros(length(SF_C))
    for a in 1:learner.num_actions
        next_possible_active_state_action = get_active_action_state_vector(next_state, a, length(next_state), learner.num_actions)
        reward_feature_backup += SF_next_target_pis[:,a] .* (ψ * next_possible_active_state_action)
    end

    target = SF_C + SF_next_discounts .* reward_feature_backup
    td_err = target - ψ * active_state_action

    pred_err = reward_C - w * projected_state

    # This should always be true as this is immediate next step prediction which is equivalent to having discounts of 0 for all states
    @assert sum(reward_discounts) == 0
    # TD err is applied across rows

    if lu.opt isa Auto
        # next_state_action_row_ind = get_action_inds(next_action, learner.num_actions, learner.num_demons)
        state_discount = -SF_next_discounts * next_active_state_action'
        state_discount .+= active_state_action'
        abs_ϕ_ψ = if λ == 0.0
            abs.(repeat(active_state_action, outer=(1, length(td_err)))')
        else
            abs.(e_ψ)
        end
        z = abs_ϕ_ψ .* max.(abs_ϕ_ψ, state_discount)
        update!(lu.opt, ψ, e_ψ, td_err, z,  learner.num_demons - learner.num_tasks, 1)

        # reward state discount is always  0 so no need for next_state_action.
        # state_discount_r = -reward_next_discounts * projected_next_state_action'
        state_discount_r = projected_state'
        abs_ϕ_w = if λ == 0.0
            abs.(repeat(projected_state, outer=(1, length(pred_err)))')
        else
            abs.(e_w)
        end
        z_r = abs_ϕ_w .* max.(abs_ϕ_w, state_discount_r)
        update!(lu.opt, w, e_w, pred_err, z_r, learner.num_tasks, 1)
    elseif lu.opt isa Flux.Descent
        α = lu.opt.eta
        if λ == 0.0
            ψ[:, active_state_action.nzind] .+= (α  * td_err) * active_state_action.nzval'
            w .= w .+ α * pred_err * projected_state'
        else
            ψ[:, e_nz] .+= (α  * td_err) .* e_ψ[:, e_nz]
            w[:, e_w_nz] .+= (α  * pred_err) .* e_w[:, e_w_nz]
        end
    else
        update!(lu.opt, ψ, - e_ψ .* td_err)
        update!(lu.opt, w, - e_w .* pred_err)
    end
    discounts .= next_discounts
end

function zero_eligibility_traces!(lu::ETB)
    for (k, v) ∈ lu.e
        if eltype(v) <: Integer
            lu.e[k] = Int[]
        elseif eltype(v) <: Vector #Used for LSTD
            lu.e[k] = zero.(v)
        else
            v .= 0
        end
    end

end


abstract type Perturbation end

mutable struct ThreePhaseFault <: Perturbation
    time::Float64
    Ybus::SparseMatrixCSC{Complex{Float64}, Int64}
end

function get_affect(::PSY.System, pert::ThreePhaseFault)
    return (integrator) -> begin
        integrator.p.Ybus = pert.Ybus
    end
end

mutable struct ControlReferenceChange <: Perturbation
    time::Float64
    device::PSY.DynamicInjection
    signal_index::Int64
    ref_value::Float64
end

function get_affect(system::PSY.System, pert::ControlReferenceChange)
    device = PSY.get_component(typeof(pert.device), system, PSY.get_name(pert.device))
    pert.device = device
    return (integrator) -> begin
        control_ref = PSY.get_ext(device)[CONTROL_REFS]
        return control_ref[pert.signal_index] = pert.ref_value
    end
end


using Knet

include("initializers.jl")

"""
A classic fully connected layer.
    w   : Weight matrix
    b   : Bias vector
    f   : activation function
    p   : dropout ratio
"""
mutable struct Dense w; b; f; p; end

"""
A custom constructor for Dense
    i       : Dimension of the input
    o       : Dimension of the input
    f       : Activation function
    pdrop   : Dropout ratio
"""
function Dense(i::Int, o::Int; init=xavier, f=nothing, pdrop=0, dtype=Array{Float32}, bias=true)
    b = convert(dtype, zeros(o, 1))
    if bias b = Param(b) end
    return Dense(Param(convert(dtype, init(o, i))), b, f, pdrop)
end

# Dense Layer Forward Pass
function (d::Dense)(x; train=true)
    if train
        x_val = d.w * mat(dropout(x,d.p)) .+ d.b
        if d.f !== nothing x_val = d.f.(x_val) end
        return x_val
    else
        x_val = value(d.w) * mat(x) .+ value(d.b)
        if d.f !== nothing x_val = d.f.(x_val) end
        return x_val
    end
end

"""
A convolutional layer with 2D kernel dimensions.
    w   : Weight matrix
    b   : Bias vector
    f   : activation function
    p   : dropout ratio
"""
mutable struct Conv2D w; b; f; p; padding; stride; dilation; end

"""
A custom constructor for Conv2D
    w1          : First dimension of the kernel
    w2          : Second dimension of the kernel
    input_dim   : Channel dimension of the input
    output_dim  : Filter size
    f           : Activation function
    pdrop       : Dropout ratio
    padding     : Padding size
    stride      : Stride size
    dilation    : Dilation size
"""
function Conv2D(
    w1::Int, w2::Int, input_dim::Int, output_dim::Int; init=kaiming_uniform, f=nothing,
    pdrop=0, padding=0, stride=1, dilation=1, bias=true, dtype=Array{Float32}
    )
    b = convert(dtype, zeros(1, 1, output_dim, 1))
    if bias b = Param(b) end

    return Conv2D(
        Param(convert(dtype, init(w1, w2, input_dim, output_dim))),
        b, f, pdrop, padding, stride, dilation
        )
end

# Conv2D Layer Forward Pass
function (c::Conv2D)(x; train=true)
    # apply dropout only in training
    if train
        if c.p > 0
            x = dropout(x, c.p)
        end
        x_val = conv4(c.w, x, padding=c.padding, stride=c.stride, dilation=c.dilation) .+ c.b
    else
        x_val = conv4(value(c.w), x, padding=c.padding, stride=c.stride, dilation=c.dilation) .+ value(c.b)
    end
    if c.f === nothing return x_val
    else return c.f.(x_val)
    end
end

"""
Batch Normalization Layer
- works both for convolutional and fully connected layers
"""
mutable struct BatchNorm bn_moments; bn_params; end

function BatchNorm(;channels=1, momentum=0.1, dtype=Array{Float32})
    return BatchNorm(
        bnmoments(momentum=momentum),
        Param(convert(dtype, bnparams(Float32, channels)))
    )
end

function (bn::BatchNorm)(x; train=true)
    return batchnorm(x, bn.bn_moments, bn.bn_params; training=train, eps=1e-5)
end


# Phases in a 1D Earth

"""
    Phase(model, name, delta, depth, time, dtdd, inc, takeoff)
    Phase(..., pierce, distance, radius)

Construct a `Phase`, which represents a single event-station path and one single phase
arrival.

`Phase` objects have the following fields:

- `model`: Name of model used to calculate phase properties
- `name`: Name of phase
- `delta`: Epicentral distance in °
- `depth`: Depth of event below surface in km
- `time`: Travel time in s
- `dtdd`: Horizontal slowness (ray parameter) in s/°
- `inc`: Incidence angle at receiver, measured from downwards in °
- `takeoff`: Takeoff angle at source, measured from downwards in °

The following fields are filled if the ray path has been calculated via `path`:

- `distance`: `Vector` of distances in ° along the path
- `radius`: `Vector` of radii in km along the path
"""
struct Phase{T} <: AbstractPhase
    model::String
    name::String
    delta::T
    depth::T
    time::T
    dtdd::T
    inc::T
    takeoff::T
    pierce::Vector{T}
    distance::Vector{T}
    radius::Vector{T}
end
Phase(model, name, delta, depth, time, dtdd, inc, takeoff) =
    Phase{TauPyFloat}(model, name, delta, depth, time, dtdd, inc, takeoff, [], [], [])
function Phase(p::AbstractPhase, pierce::AbstractVector, distance::AbstractVector, radius::AbstractVector)
    length(pierce) == length(distance) == length(radius) ||
        throw(ArgumentError("pierce, distance and radius must be the same length"))
    Phase{TauPyFloat}(p.model, p.name, p.delta, p.depth, p.time, p.dtdd, p.inc, p.takeoff, pierce, distance, radius)
end

"""
    path(depth, distance, phase="ttall"; model="$DEFAULT_MODEL", cache=true) -> p::Vector{Phase}

Create a set of `Phase`s `p` which contain the computed ray path for a set of `phase`s
from an event `depth` km deep and at `distance`° away.
Optionally specify the `model` (one of: $(AVAILABLE_MODELS)).

If no arrivals are found for the geometry, an empty `Vector{Phase}` is returned.

Set `cache` to `false` to always recompute the ray path and ignore the module cache.
"""
function path(depth, distance, phase="ttall"; model=DEFAULT_MODEL, cache=true)
    phase = phase isa AbstractString ? [phase] : phase
    _call_taup(:get_ray_paths, model, depth, distance, phase; cache=cache)
end

"""
    travel_time(depth, distance, phase="ttall"; model="$DEFAULT_MODEL", cache=true) -> p::Vector{Phase}

Return a `Vector` of `Phase`s, given an event `depth` km deep and `distance`°
away.  Optionally specify a `phase` name; otherwise all arrivals are returned.
Optionally specify the model (one of: $(AVAILABLE_MODELS)).

If no arrivals are found for the geometry, an empty `Vector{Phase}` is returned.

Set `cache` to `false` to always recompute the travel time and ignore the module cache.
"""
function travel_time(depth, distance, phase="ttall"; model=DEFAULT_MODEL, cache=true)
    phase = phase isa AbstractString ? [phase] : phase
    _call_taup(:get_travel_times, model, depth, distance, phase; cache=cache)
end

"""Helper function which takes `obspy.taup.Arrivals` and returns `Phase`s.
Arguments should be the same as those to osbpy.taup calls, prefixed
by `model`."""
function _phases_from_arrivals(arr, model, depth, distance, phase)
    p = Vector{Phase{TauPyFloat}}()
    for a in arr
        name = a.name
        delta = a.distance
        depth = a.source_depth
        time = a.time
        dtdd = deg2rad(a.ray_param)
        inc = a.incident_angle
        takeoff = a.takeoff_angle
        path = a.path
        distance, radius = if path === nothing
            TauPyFloat[], TauPyFloat[]
        else
            pathlist = path.tolist()
            rad2deg.([pp[3] for pp in pathlist]), RADIUS[model] .- [pp[4] for pp in pathlist]
        end
        push!(p, Phase{TauPyFloat}(model, name, delta, depth, time, dtdd, inc, takeoff,
              [], distance, radius))
    end
    p
end


module iPredict
include("TableProcessing.jl")
include("FindKm.jl")
include("Simulate.jl")
include("Plots.jl")
include("pValues.jl")
import Cairo
import Fontconfig
import DataFrames
using Distributions
using Statistics
using DelimitedFiles

function InvestigateIsoformChoice(countsMatrixPath::String, GeneIsoTablePath::String,
     NumIsoformsToSimulate::Int64,
     filteringThreshold::Int64)

     global pFN = 0.04
     global pFP = 0.01

     #Read in counts matrix
     println("Data wrangling commencing!")
     println("Reading in counts matrix.") #Quick
     countsMatrix = ReadCountsMatrix(countsMatrixPath)

     #Read in gene-isoform table
     println("Reading in gene-isoform relationships")
     GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

     #Make sure gene-isoform table matches counts matrix
     if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
         throw(ArgumentError("Your counts matrix doesn't match your list of genes
         and isoforms - there is something wrong with your input files."))
     end

     #Add gene information to isoforms table
     println("Adding Gene information to countsMatrix")
     countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

     #Clean data for finding Km
     println("Cleaning data for finding Km")
     MMdf = CleanDataForMM(countsMatrix)

     #Find Km
     println("Calculating Km....")
     MMParams = FindKm(MMdf)
     Km = MMParams[1]
     sigma = MMParams[2]

     #Extract genes with filteringThreshold expressed isoforms
     println("Converting counts to dict")
     GenesWithNIsoforms = KeepGenesWithNIsoforms(countsMatrixWithGenes,
     filteringThreshold)
     workingDict = GenesWithNIsoforms[1]
     unexpr_dict = GenesWithNIsoforms[2]
     MaxNumIsos = GenesWithNIsoforms[3]

     #Get ranking and dropout info for each gene
     rankingDict = Dict()
     pObsDict = Dict()
     pDropoutDict = Dict()
     genesList = []

     for (key, value) in workingDict
         rankingDict[key] = FindRank(value)
         pObsDict[key] = FindPObs(value)
         pDropoutDict[key] = FindPDropout(value, Km)
         push!(genesList, key)
     end

     df = FindMeanUnifObsProbabilities(NumIsoformsToSimulate,
     genesList, pObsDict, pDropoutDict, rankingDict)

     return df
end

#Function that performs global level predictions
function globalPredict(countsMatrixPath::String, GeneIsoTablePath::String,
     NumSimulations::Int64, output::String, NumIsoformsToSimulate::Int64,
     filteringThreshold::Int64, test::Bool, alpha, beta, pFNeg::Float64,
     pFPos::Float64, error_model::String, choice_model::String, path::String)

    #Set probability(FP|N) and probability(FN|P)
    if error_model == "none"
        global pFN = pFNeg
        global pFP = pFPos
    end

    #Read in counts matrix
    println("Data wrangling commencing!")
    println("Reading in counts matrix.") #Quick
    countsMatrix = ReadCountsMatrix(countsMatrixPath)

    #Read in gene-isoform table
    println("Reading in gene-isoform relationships")
    GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

    #Make sure gene-isoform table matches counts matrix
    if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
        throw(ArgumentError("Your counts matrix doesn't match your list of genes
        and isoforms - there is something wrong with your input files."))
    end

    #Add gene information to isoforms table
    println("Adding Gene information to countsMatrix")
    countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

    #Only find Km if no Beta distribution parameters provided
    if isnothing(alpha) || isnothing(beta)

        #Clean data for finding Km
        println("Cleaning data for finding Km")
        MMdf = CleanDataForMM(countsMatrix)

        #Find Km
        println("Calculating Km....")
        MMParams = FindKm(MMdf)
        Km = MMParams[1]
        sigma = MMParams[2]
    end

    #Extract genes with filteringThreshold expressed isoforms
    println("Converting counts to dict")
    GenesWithNIsoforms = KeepGenesWithNIsoforms(countsMatrixWithGenes,
    filteringThreshold)
    workingDict = GenesWithNIsoforms[1]
    unexpr_dict = GenesWithNIsoforms[2]
    MaxNumIsos = GenesWithNIsoforms[3]

    #Get ranking and dropout info for each gene
    rankingDict = Dict()
    pObsDict = Dict()
    pDropoutDict = Dict()
    genesList = []

    #How we calculate dropouts depends on whether or not beta parameters are provided
    if isnothing(alpha) || isnothing(beta)
        for (key, value) in workingDict
            rankingDict[key] = FindRank(value)
            pObsDict[key] = FindPObs(value)
            pDropoutDict[key] = FindPDropout(value, Km)
            push!(genesList, key)
        end
    else
        for (key, value) in workingDict
            rankingDict[key] = FindRank(value)
            pObsDict[key] = FindPObs(value)
            pDropoutDict[key] = rand(Beta(alpha, beta), length(value))
            push!(genesList, key)
        end
    end

    plot_arr = []
    log_plot_arr = []
    mean_arr = []
    median_arr = []
    std_arr = []

    println("Beginning simulations!")

    for i in 1:NumIsoformsToSimulate

        #I pre-allocate arr in the hope it will make my code faster
        arr = Array{Float64}(undef, NumSimulations * length(genesList))
        arr_index = 1

        overlapArr = Array{Float64}(undef, NumSimulations * length(genesList))
        overlap_index = 1

        pDropoutArr = Array{Float64}(undef, NumSimulations * length(genesList) *
        (ncol(countsMatrix) - 1) * i)
        drop_index = 1

        #Need to add a for loop with num simulations here
        for j in 1:NumSimulations

            #@show arr_index

            #Get an array of mean isoforms per gene per cell
            simulation_out = globalArraySimulation(1, genesList, rankingDict, pObsDict,
            pDropoutDict, ncol(countsMatrix) - 1, i, filteringThreshold,
            unexpr_dict, MaxNumIsos, error_model, choice_model)

            arr[arr_index:(arr_index + length(genesList) - 1)] = simulation_out[1]
            overlapArr[overlap_index:(overlap_index + length(genesList) - 1)] = simulation_out[2]
            pDropoutArr[drop_index:(drop_index + length(simulation_out[3]) - 1)] = simulation_out[3]

            #Increment arr_index
            arr_index += length(genesList)
            overlap_index += length(genesList)
            drop_index += length(simulation_out[3])
        end

        #Save histogram of data
        if i == 1
            p1 = makeNormalisedHistogram(arr, string(i) * " Isoform")
        else
            p1 = makeNormalisedHistogram(arr, string(i) * " Isoforms")
        end
        push!(plot_arr, p1[1])
        push!(log_plot_arr, p1[2])
        push!(mean_arr, mean(arr))
        push!(std_arr, std(arr))
        push!(median_arr, median(arr))

        #Write to file if required
        if path != "none"
            writedlm(path * "_" * string(i) * "_isoforms_array.txt", arr)
            writedlm(path * "_" * string(i) * "_isoforms_overlap_array.txt", overlapArr)
            writedlm(path * "_" * string(i) * "_isoforms_pdrop_array.txt", pDropoutArr)
        end
    end

    #Get an array and a dict of mean isoforms per gene per cell in real data
    println("Finding means in real data.")
    findRealMeans = MeanNumIsoformsPerGenePerCellReal(workingDict)
    realMeansArr = findRealMeans[1]

    #Write to file if required
    if path != "none"
        writedlm(path * "_" * "real_isoforms_array.txt", realMeansArr)
    end

    #Plot real data
    real_plot = makeNormalisedHistogramReal(realMeansArr, "Real")
    push!(plot_arr, real_plot[1])
    push!(log_plot_arr, real_plot[2])

    #Arrange plots and write to output
    graph = plot(plot_arr..., layout=(length(plot_arr),1),legend=false, size = (300, 600), bottom_margin = 1.0mm)
    savefig(graph, output)
    graph = plot(log_plot_arr..., layout=(length(plot_arr),1),legend=false, size = (300, 600))
    savefig(graph, output * "_log.png")

    return (mean_arr, std_arr, median_arr)

end

#Function that makes negative control model
function NegativeControl(countsMatrixPath::String, GeneIsoTablePath::String,
     NumSimulations::Int64, output::String, NumIsoformsToSimulate::Int64,
     filteringThreshold::Int64, path::String)

     #Read in counts matrix
     println("Data wrangling commencing!")
     println("Reading in counts matrix.") #Quick
     countsMatrix = ReadCountsMatrix(countsMatrixPath)

     #Read in gene-isoform table
     println("Reading in gene-isoform relationships")
     GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

     #Make sure gene-isoform table matches counts matrix
     if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
         throw(ArgumentError("Your counts matrix doesn't match your list of genes
         and isoforms - there is something wrong with your input files."))
     end

     #Add gene information to isoforms table
     println("Adding Gene information to countsMatrix")
     countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

     println("Converting counts to dict")
     GenesWithNIsoforms = KeepGenesWithNIsoforms(countsMatrixWithGenes,
     filteringThreshold)
     workingDict = GenesWithNIsoforms[1]
     unexpr_dict = GenesWithNIsoforms[2]
     MaxNumIsos = GenesWithNIsoforms[3]

     #Get ranking and dropout info for each gene
     rankingDict = Dict()
     pObsDict = Dict()
     pDropoutDict = Dict()
     genesList = []

     #Fill arrays
     for (key, value) in workingDict
         rankingDict[key] = FindRank(value)
         pObsDict[key] = FindPObs(value)

         push!(genesList, key)
     end

     println("Beginning simulations!")

     for i in 1:NumIsoformsToSimulate

         #I pre-allocate arr in the hope it will make my code faster
         arr = Array{Float64}(undef, NumSimulations * length(genesList))
         arr_index = 1

         #Need to add a for loop with num simulations here
         for j in 1:NumSimulations

             #@show arr_index

             #Get an array of mean isoforms per gene per cell
             arr[arr_index:(arr_index + length(genesList) - 1)] =
             negControlSimulation(1, genesList, rankingDict, pObsDict,
             ncol(countsMatrix) - 1, i, filteringThreshold, pDropoutDict,
             MaxNumIsos)

             #Increment arr_index
             arr_index += length(genesList)
         end

         #Write to file if required
         if path != "none"
             writedlm(path * "_" * string(i) * "_isoforms_array.txt", arr)
         end
     end
     return true
end

#Function that performs global level predictions
function NegativeControlDrop(countsMatrixPath::String, GeneIsoTablePath::String,
     NumSimulations::Int64, output::String, NumIsoformsToSimulate::Int64,
     filteringThreshold::Int64, test::Bool, pFNeg::Float64,
     pFPos::Float64, error_model::String, choice_model::String, path::String)

    #Set probability(FP|N) and probability(FN|P)
    if error_model == "none"
        global pFN = pFNeg
        global pFP = pFPos
    end

    #Read in counts matrix
    println("Data wrangling commencing!")
    println("Reading in counts matrix.") #Quick
    countsMatrix = ReadCountsMatrix(countsMatrixPath)

    #Read in gene-isoform table
    println("Reading in gene-isoform relationships")
    GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

    #Make sure gene-isoform table matches counts matrix
    if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
        throw(ArgumentError("Your counts matrix doesn't match your list of genes
        and isoforms - there is something wrong with your input files."))
    end

    #Add gene information to isoforms table
    println("Adding Gene information to countsMatrix")
    countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

    #Extract genes with filteringThreshold expressed isoforms
    println("Converting counts to dict")
    GenesWithNIsoforms = KeepGenesWithNIsoforms(countsMatrixWithGenes,
    filteringThreshold)
    workingDict = GenesWithNIsoforms[1]
    unexpr_dict = GenesWithNIsoforms[2]
    MaxNumIsos = GenesWithNIsoforms[3]

    #Get ranking and dropout info for each gene
    rankingDict = Dict()
    pObsDict = Dict()
    pDropoutDict = Dict()
    genesList = []

    for (key, value) in workingDict
        rankingDict[key] = FindRank(value)
        pObsDict[key] = FindPObs(value)
        pDropoutDict[key] = Float64[0.0, 0.0, 0.0, 0.0]
        push!(genesList, key)
    end

    plot_arr = []
    log_plot_arr = []
    mean_arr = []
    median_arr = []
    std_arr = []

    println("Beginning simulations!")

    for i in 1:NumIsoformsToSimulate

        #I pre-allocate arr in the hope it will make my code faster
        arr = Array{Float64}(undef, NumSimulations * length(genesList))
        arr_index = 1

        overlapArr = Array{Float64}(undef, NumSimulations * length(genesList))
        overlap_index = 1

        pDropoutArr = Array{Float64}(undef, NumSimulations * length(genesList) *
        (ncol(countsMatrix) - 1) * i)
        drop_index = 1

        #Need to add a for loop with num simulations here
        for j in 1:NumSimulations

            #@show arr_index

            #Get an array of mean isoforms per gene per cell
            simulation_out = globalArraySimulation(1, genesList, rankingDict, pObsDict,
            pDropoutDict, ncol(countsMatrix) - 1, i, filteringThreshold,
            unexpr_dict, MaxNumIsos, error_model, choice_model)

            arr[arr_index:(arr_index + length(genesList) - 1)] = simulation_out[1]
            overlapArr[overlap_index:(overlap_index + length(genesList) - 1)] = simulation_out[2]
            pDropoutArr[drop_index:(drop_index + length(simulation_out[3]) - 1)] = simulation_out[3]

            #Increment arr_index
            arr_index += length(genesList)
            overlap_index += length(genesList)
            drop_index += length(simulation_out[3])
        end

        #Save histogram of data
        if i == 1
            p1 = makeNormalisedHistogram(arr, string(i) * " Isoform")
        else
            p1 = makeNormalisedHistogram(arr, string(i) * " Isoforms")
        end
        push!(plot_arr, p1[1])
        push!(log_plot_arr, p1[2])
        push!(mean_arr, mean(arr))
        push!(std_arr, std(arr))
        push!(median_arr, median(arr))

        #Write to file if required
        if path != "none"
            writedlm(path * "_" * string(i) * "_isoforms_array.txt", arr)
            writedlm(path * "_" * string(i) * "_isoforms_overlap_array.txt", overlapArr)
            writedlm(path * "_" * string(i) * "_isoforms_pdrop_array.txt", pDropoutArr)
        end
    end

    #Get an array and a dict of mean isoforms per gene per cell in real data
    println("Finding means in real data.")
    findRealMeans = MeanNumIsoformsPerGenePerCellReal(workingDict)
    realMeansArr = findRealMeans[1]

    #Write to file if required
    if path != "none"
        writedlm(path * "_" * "real_isoforms_array.txt", realMeansArr)
    end

    #Plot real data
    real_plot = makeNormalisedHistogramReal(realMeansArr, "Real")
    push!(plot_arr, real_plot[1])
    push!(log_plot_arr, real_plot[2])

    #Arrange plots and write to output
    graph = plot(plot_arr..., layout=(length(plot_arr),1),legend=false, size = (300, 600), bottom_margin = 1.0mm)
    savefig(graph, output)
    graph = plot(log_plot_arr..., layout=(length(plot_arr),1),legend=false, size = (300, 600))
    savefig(graph, output * "_log.png")

    return (mean_arr, std_arr, median_arr)

end

function oneGenePrediction(countsMatrixPath::String, GeneIsoTablePath::String,
    gene::String, NumSimulations::Int64, output::String, test::Bool)

    global pFN = 0.04
    global pFP = 0.01

    #Read in counts matrix
    println("Data wrangling commencing!")
    println("Reading in counts matrix.")
    countsMatrix = ReadCountsMatrix(countsMatrixPath)

    #Read in gene-isoform table
    println("Reading in gene-isoform relationships")
    GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

    #Make sure gene-isoform table matches counts matrix
    if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
        throw(ArgumentError("Your counts matrix doesn't match your list of genes
        and isoforms - there is something wrong with your input files."))
    end

    #Add gene information to isoforms table
    println("Adding Gene information to countsMatrix")
    countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

    #Make a dictionary just for the expressed isoforms in gene - no point
    #simulating unexpressed isoforms.
    tmp = filter(r -> any(occursin.([gene], r.Genes)), countsMatrixWithGenes)
    NumIsos = nrow(tmp)  #Get number of isoforms (incl. unexpressed) for gene of interest
    countsDict = ConvertExpressedIsosToDict(tmp)

    #Clean data for finding Km
    println("Cleaning data for finding Km")
    MMdf = CleanDataForMM(countsMatrix)

    #Find Km
    println("Calculating Km....")
    MMParams = FindKm(MMdf)
    Km = MMParams[1]
    sigma = MMParams[2]

    println("Preparing for simulations...")

    #Find ranking info
    cumulativeFreqs = GetIsoformChoiceProbabilitiesGene(length(countsDict[gene]))

    meansDict = Dict()

    #For each gene - bit unnecessary since just 1 gene, but convenient way of
    #getting value.
    for (key, value) in countsDict

        #Rank isoforms
        rankArr::Array{Int64, 1} = FindRank(value)

        #Find p(dropouts)
        dropoutArr::Array{Float64, 1} = FindPDropout(value, Km)

        for i in 1:length(countsDict[gene])

            meansArr = []

            for j in 1:NumSimulations
                mean = Simulate(cumulativeFreqs, rankArr, dropoutArr,
                ncol(countsMatrix) - 1, i, NumIsos, test)
                push!(meansArr, mean)
            end

            meansDict[i] = meansArr

        end
    end

    #Make a plot and write it to file
    OneGeneHistPlot(meansDict, tmp, output)

    #Find and return p-values
    mean = MeanIsoPerGenePerCell(tmp)
    return FindpValues(meansDict, mean, 40) #40 is an arbitrary number of required unique values

end

#Function that performs global level predictions
function rankingPredict(countsMatrixPath::String, GeneIsoTablePath::String,
     NumSimulations::Int64, output::String, NumIsoformsToSimulate::Int64,
     filteringThreshold::Int64, test::Bool, alpha, beta, pFNeg::Float64,
     pFPos::Float64, error_model::String, choice_model::String, path::String)

    #Set probability(FP|N) and probability(FN|P)
    if error_model == "none"
        global pFN = pFNeg
        global pFP = pFPos
    end

    #Read in counts matrix
    println("Data wrangling commencing!")
    println("Reading in counts matrix.") #Quick
    countsMatrix = ReadCountsMatrix(countsMatrixPath)

    #Read in gene-isoform table
    println("Reading in gene-isoform relationships")
    GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

    #Make sure gene-isoform table matches counts matrix
    if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
        throw(ArgumentError("Your counts matrix doesn't match your list of genes
        and isoforms - there is something wrong with your input files."))
    end

    #Add gene information to isoforms table
    println("Adding Gene information to countsMatrix")
    countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

    #Only find Km if no Beta distribution parameters provided
    if isnothing(alpha) || isnothing(beta)

        #Clean data for finding Km
        println("Cleaning data for finding Km")
        MMdf = CleanDataForMM(countsMatrix)

        #Find Km
        println("Calculating Km....")
        MMParams = FindKm(MMdf)
        Km = MMParams[1]
        sigma = MMParams[2]
    end

    #Extract genes with filteringThreshold expressed isoforms
    println("Converting counts to dict")
    GenesWithNIsoforms = KeepGenesWithNIsoforms(countsMatrixWithGenes,
    filteringThreshold)
    workingDict = GenesWithNIsoforms[1]
    unexpr_dict = GenesWithNIsoforms[2]
    MaxNumIsos = GenesWithNIsoforms[3]

    #Get ranking and dropout info for each gene
    rankingDict = Dict()
    pObsDict = Dict()
    pDropoutDict = Dict()
    genesList = []

    #How we calculate dropouts depends on whether or not beta parameters are provided
    if isnothing(alpha) || isnothing(beta)
        for (key, value) in workingDict
            rankingDict[key] = FindRank(value)
            pObsDict[key] = FindPObs(value)
            pDropoutDict[key] = FindPDropout(value, Km)
            push!(genesList, key)
        end
    else
        for (key, value) in workingDict
            rankingDict[key] = FindRank(value)
            pObsDict[key] = FindPObs(value)
            pDropoutDict[key] = rand(Beta(alpha, beta), length(value))
            push!(genesList, key)
        end
    end

    plot_arr = []
    log_plot_arr = []
    mean_arr = []
    median_arr = []
    std_arr = []

    println("Beginning simulations!")

    for i in 4:NumIsoformsToSimulate #Only interested in 4 isoform sims

        #I pre-allocate arr in the hope it will make my code faster
        arr = []
        arr_index = 1

        std_arr = []
        std_arr_index = 1

        #Need to add a for loop with num simulations here
        for j in 1:NumSimulations

            #@show arr_index

            #Get an array of mean isoforms per gene per cell
            simulation_out = rankingSimulation(1, genesList, rankingDict, pObsDict,
            pDropoutDict, ncol(countsMatrix) - 1, i, filteringThreshold,
            unexpr_dict, MaxNumIsos, error_model, choice_model)

            #@show simulation_out

            append!(arr, simulation_out)
            #append!(std_arr, simulation_out[2])

            #arr[arr_index:(arr_index + length(simulation_out[1]) - 1)] = simulation_out[1]
            #Increment arr_index
            #arr_index += length(genesList)

            #std_arr[std_arr_index:(std_arr_index + length(simulation_out[2]) - 1)] = simulation_out[2]
            #Increment arr_index
            #std_arr_index += length(genesList)

        end

        if path != "none"
            writedlm(path * "_spearmans_array.txt", arr)
            writedlm(path * "_std_array.txt", std_arr)
        end
    end

    return ("Simulations complete")

end

#Function that performs global level predictions
function rankingOriginal(countsMatrixPath::String, GeneIsoTablePath::String,
     NumSimulations::Int64, output::String, NumIsoformsToSimulate::Int64,
     filteringThreshold::Int64, test::Bool, alpha, beta, pFNeg::Float64,
     pFPos::Float64, error_model::String, choice_model::String, path::String)

    #Set probability(FP|N) and probability(FN|P)
    if error_model == "none"
        global pFN = pFNeg
        global pFP = pFPos
    end

    #Read in counts matrix
    println("Data wrangling commencing!")
    println("Reading in counts matrix.") #Quick
    countsMatrix = ReadCountsMatrix(countsMatrixPath)

    #Read in gene-isoform table
    println("Reading in gene-isoform relationships")
    GeneIsoformRelationships = ReadGeneIsoformRelationships(GeneIsoTablePath)

    #Make sure gene-isoform table matches counts matrix
    if CheckGenesMatchIsoforms(countsMatrix, GeneIsoformRelationships) == false
        throw(ArgumentError("Your counts matrix doesn't match your list of genes
        and isoforms - there is something wrong with your input files."))
    end

    #Add gene information to isoforms table
    println("Adding Gene information to countsMatrix")
    countsMatrixWithGenes = AddGeneInfo(countsMatrix, GeneIsoformRelationships)

    #Only find Km if no Beta distribution parameters provided
    if isnothing(alpha) || isnothing(beta)

        #Clean data for finding Km
        println("Cleaning data for finding Km")
        MMdf = CleanDataForMM(countsMatrix)

        #Find Km
        println("Calculating Km....")
        MMParams = FindKm(MMdf)
        Km = MMParams[1]
        sigma = MMParams[2]
    end

    #Extract genes with filteringThreshold expressed isoforms
    println("Converting counts to dict")
    GenesWithNIsoforms = KeepGenesWithNIsoforms(countsMatrixWithGenes,
    filteringThreshold)
    workingDict = GenesWithNIsoforms[1]
    unexpr_dict = GenesWithNIsoforms[2]
    MaxNumIsos = GenesWithNIsoforms[3]

    #Get ranking and dropout info for each gene
    rankingDict = Dict()
    pObsDict = Dict()
    pDropoutDict = Dict()
    genesList = []
    spearmans_arr = []
    pDropout_Rank_cor = []
    std_arr = []

    for (key, value) in workingDict
        Arr1 = FindRank(value)
        Arr2 = FindRankbyNoCells(value)
        Arr3 = ordinalrank(FindPDropout(value, Km))
        push!(std_arr, FindStd(value))

        push!(spearmans_arr,
        corspearman(convert(Array{Float64, 1}, Arr1),
        convert(Array{Float64, 1}, Arr2)))

        push!(pDropout_Rank_cor,
        corspearman(convert(Array{Float64, 1}, Arr1),
        convert(Array{Float64, 1}, Arr3)))

    end

    writedlm(path * "_spearmans_array.txt", spearmans_arr)
    writedlm(path * "_pdropout_rank_cor_array.txt", pDropout_Rank_cor)
    writedlm(path * "_std.txt", std_arr)

end

export globalPredict
export oneGenePrediction
export NegativeControl
export NegativeControlDrop
export rankingPredict
export rankingOriginal

end #end of module

#GPU structure of kernel
#find p(dropout)
#for i in 1:10000
#simulation process
#save average
#return an array of averages