export
    UAVActions,
    UAV,
    get_action,
    belief_state_from_slave,
    get_polar_state,
    get_belief_state,
    get_other_uav_id,
    apply_sensor_error

#=
This file defines the uav type, which uses different strategies to
make decisions based on the current state of the system.
=#

# type to hold the actions and signals emitted from uavs
type UAVActions
    actions::Vector{Float64}
    # signals sent by the uavs to coordinate next actions
    signals::Vector{Symbol}
    function UAVActions(num_uavs::Int64)
        actions = Array(Float64, num_uavs)
        signals = [:no_signal, :no_signal]
        return new(actions, signals)
    end
end

# type for uav
type UAV
    # id gives index into tables for this uav
    uav_id::Int64
    # actions
    actions::Vector{Float64}
    num_actions::Int64
    joint_actions::Matrix{Float64}
    num_joint_actions::Int64
    # tables
    policy_table::PolicyTable
    coordination_table::CoordinationTable
    # type of policy
    is_policy_joint::Bool
    # strategy
    strategy::Function
    action_selection_method::String
    # sensor accuracy information
    # we now need to pass the random number generator to the uav:
    rng::MersenneTwister
    # gaussian sensors, mean = 0, std dev these values
    sensor_range_std::Float64
    sensor_theta_std::Float64
    sensor_psi_std::Float64
    sensor_velocity_std::Float64
    use_sensor_noise::Bool
    function UAV(uav_id::Int64, actions::Vector{Float64}, joint_actions::Matrix{Float64},
            policy_table::PolicyTable, coordination_table::CoordinationTable, is_policy_joint::Bool,
            strategy::Function, action_selection_method::String, rng::MersenneTwister,
            sensor_range_std::Float64, sensor_theta_std::Float64, sensor_psi_std::Float64,
            sensor_velocity_std::Float64, use_sensor_noise::Bool)
        return new(uav_id, actions, length(actions), joint_actions, size(joint_actions, 2),
                policy_table, coordination_table, is_policy_joint, strategy, action_selection_method,
                rng, sensor_range_std, sensor_theta_std, sensor_psi_std, sensor_velocity_std, use_sensor_noise)
    end
end

#=
Description:
Delegates to the strategy of the uav to decide which action to take

Parameters:
- uav: the uav deciding which action to take
- state: state object which contains all variable information
    needed by the uav to decide which action to take

Return Value:
- action: float value that is the action of this uav
- signal: float value indiciating coordination action
- utilities: the utilities used to derive the actions
- belief_state: the belief_state used to index the policy table
- different_action: whether or not the coordination table prevented
    the taking of a different action from the greedy case and
    the coordinationt table action was coc. This is used to
    wake the pilot.
=#
function get_action(uav::UAV, state::State)

    # compute belief state
    belief_state = get_belief_state(uav, state)

    # get the baseline utilities from the policy table
    utilities = get_utilities(uav.policy_table, belief_state)

    # delegate to the strategy to decide which action to take
    # as well as whether to signal the other plane
    action, signals, different_action = uav.strategy(uav, state, utilities)

    return action, signals, utilities, belief_state, different_action
end

#=
Description:
Converts the state from the master point of view to the slave point of view
to select the correct action.
=#
function belief_state_from_slave(belief_state::Array{Float64,1})
    v0 = belief_state[4]
    v1 = belief_state[5]
    slave_belief_state = belief_state
    # [1] stays the same
    # [2] changes - requires calc
    # [3] changes - requires calc
    # [4] becomes [5] - no calculation
    slave_belief_state[4] = v1
    # [5] becomes [4] - no calculation
    slave_belief_state[5] = v0
    # [6] stays the same

    # solve for x and y coordinates
    range = belief_state[1]
    theta = belief_state[2] # rad
    psi = belief_state[3] # rad
    p1 = [0;0]
    p2 = [range*cos(theta);range*sin(theta)]

    # move the point 1 and rotate
    p1 = p1-p2
    # rotate point 1
    rot = [cos(psi) sin(psi);-sin(psi) cos(psi)]
    p1 = rot*p1

    # calculate new theta
    thetaNew = atan2(p1[2],p1[1]) # rad
    while thetaNew > pi
        thetaNew = thetaNew - 2*pi
    end
    while thetaNew < -pi
        thetaNew = thetaNew + 2*pi
    end

    # calculate new psi
    psiNew = 2*pi-psi # rad
    # ensure new psi is between 0 and 2*pi
    while psiNew < 0
        psiNew = psiNew+2*pi
    end

    slave_belief_state[2] = thetaNew # rad
    slave_belief_state[3] = psiNew # rad

    return slave_belief_state
end

#=
Description:
Converts the state to polar format and check that it is in bounds of the policy

Parameters:
- uav: the uav
- state: state object to convert to polar

Return Value:
- polar_state: the state converted to polar coordinates, which consists of
    [xr, yr, pr, vown, vint, resp, resp] to [rho, theta, psi, vown, vint, resp, resp]
=#
function get_polar_state(uav::UAV, state::State)
    # delegate to state for conversion
    if uav.is_policy_joint # joint policy
        polar_state = to_polar(state)
    else
        polar_state = to_polar_non_joint(state)
    end

    # get bounds of grid
    max_r = maximum(uav.policy_table.grid.cutPoints[1])
    min_t = minimum(uav.policy_table.grid.cutPoints[2])
    max_t = maximum(uav.policy_table.grid.cutPoints[2])

    # check that the new polar state is inside of the bounds
    # and if not return the terminal state
    r = polar_state[1]
    theta = polar_state[2]
    if r > max_r || theta < min_t || theta > max_t
        polar_state = TERMINAL_POLAR_STATE
    end

    return polar_state
end

#=
Description:
Populates a belief state of the uav

Parameters:
- uav: the uav, the belief state of which is being calculated
- state: the state converted to belief state

Return Value:
- belief_state: probability distribution over possible states
=#
function get_belief_state(uav::UAV, state::State)
    # initialize belief as sparse vector over all possible states
    belief_state = spzeros(uav.policy_table.num_states, 1)

    # convert state to polar
    polar_state = get_polar_state(uav, state)

    if !uav.is_policy_joint
        if uav.uav_id != 1
            if polar_state != TERMINAL_POLAR_STATE
                polar_state = belief_state_from_slave(polar_state)
            end
        end
    end

    # add sensor error if using noisy sensors
    if uav.use_sensor_noise
        sensor_error_state = apply_sensor_error(uav, polar_state)
    else
        sensor_error_state = polar_state
    end

    # populate belief_state
    if polar_state == TERMINAL_POLAR_STATE
        belief_state[end] = 1.0
    else
        indices, weights = interpolants(uav.policy_table.grid, sensor_error_state)
        belief_state[indices] = weights
    end

    return belief_state
end

#=
Description:
Returns the id value of the other uav

Parameters:
- uav: the current uav

Return Value:
- id: int value of the other uavs id
=#
function get_other_uav_id(uav::UAV)
    if uav.uav_id == 1
        return 2
    else
        return 1
    end
end

function apply_sensor_error(uav::UAV, polar_state::Vector{Float64})
    # copy polar state because it might be used
    # elsewhere in the above function and it's not
    # obvious whether we might want to use the
    # original or a noisy version above
    sensor_error_state = deepcopy(polar_state)

    # range
    sensor_error_state[1] = (sensor_error_state[1]
        + randn(uav.rng) * uav.sensor_range_std)
    sensor_error_state[2] = (sensor_error_state[2]
        + randn(uav.rng) * uav.sensor_theta_std)
    sensor_error_state[3] = (sensor_error_state[3]
        + randn(uav.rng) * uav.sensor_psi_std)

    # apply noise to other planes velocity
    if uav.uav_id == 1
        other_uav_velocity_idx = 5
    else
        other_uav_velocity_idx = 4
    end
    sensor_error_state[other_uav_velocity_idx] = (
        sensor_error_state[other_uav_velocity_idx]
        + randn(uav.rng) * uav.sensor_velocity_std)

    return sensor_error_state
end


# A custom geometry type to test that we can render arbitrary primitives
# by decomposing them into simple meshes. This replaces the previous test
# which did the same thing using a Polyhedron from Polyhedra.jl. The Polyhedra
# test was removed because it required too many external dependencies just to
# verify a simple interface.
struct CustomGeometry <: GeometryPrimitive{3, Float64}
end

GeometryTypes.isdecomposable(::Type{<:Face}, ::CustomGeometry) = true
function GeometryTypes.decompose(::Type{F}, c::CustomGeometry) where {F <: Face}
    [convert(F, Face(1,2,3))]
end
GeometryTypes.isdecomposable(::Type{<:Point}, ::CustomGeometry) = true
function GeometryTypes.decompose(::Type{P}, c::CustomGeometry) where {P <: Point}
    convert(Vector{P}, [Point(0., 0, 0), Point(0., 1, 0), Point(0., 0, 1)])
end

@testset "Custom decomposable geometry" begin
    Visualizer(CustomGeometry())
end


# Author: Youngjun Kim, youngjun@stanford.edu
# Date: 08/19/2014
#
# Originally written in MATLAB by Mykel Kochenderfer, mykel@stanford.edu
# Converted by Youngjun Kim, youngjun@stanford.edu

#   LOAD_SCRIPTS(FILENAME) returns a structure holding the encounter
#   scripts stored in the specified file.
#
#   SCRIPTS FILE:
#   The scripts file contains a set of scripted encounters. Each
#   encounter is defined by a set of encounter scripts associated with a
#   fixed number of aircraft. The file is organized as follows:
#
#   [Header]
#   uint32 (number of encounters)
#   uint32 (number of aircraft)
#       [Encounter 1]
#           [Initial]
#               [Aircraft 1]
#               double (initial airspeed in ft/s)
#               double (initial north position in ft)
#               double (initial east position in ft)
#               double (initial altitude in ft)
#               double (initial heading angle in radians)
#               double (initial flight path angle in radians)
#               double (initial roll angle in radians)
#               double (initial airspeed acceleration in ft/s^2)
#               ...
#               [Aircraft n]
#               double (initial airspeed in ft/s)
#               double (initial north position in ft)
#               double (initial east position in ft)
#               double (initial altitude in ft)
#               double (initial heading angle in radians)
#               double (initial flight path angle in radians)
#               double (initial roll angle in radians)
#               double (initial airspeed acceleration in ft/s^2)
#           [Updates]
#               [Aircraft 1]
#               uint8 (number of updates)
#                   [Update 1]
#                   double (time in s)
#                   double (commanded vertical rate in ft/s)
#                   double (commanded turn rate in rad/s)
#                   double (commanded airspeed acceleration in ft/s^2)
#                   ...
#                   [Update m]
#                   double (time in s)
#                   double (commanded vertical rate in ft/s)
#                   double (commanded turn rate in rad/s)
#                   double (commanded airspeed acceleration in ft/s^2)
#               ...
#               [Aircraft n]
#                   ...
#       ...
#       [Encounter k]
#           ...
#
#   SCRIPTS STRUCTURE:
#   The waypoints structure is an m x n structure matrix, where m is the
#   number of aircraft and n is the number of encounters. This structure
#   matrix has two fields: initial and update. The initial field is an 8
#   element array specifying the airspeed, north position, east position,
#   altitude, heading, flight path angle, roll angle, and airspeed
#   acceleration. The update field is a 4 x n matrix, where n is the number
#   of updates. The rows correspond to the time, vertical rate, turn rate
#   and airspeed acceleration.
#   The script also takes an additional parameter through varargin:  the
#   number of encounters to load, 'limit'.  load_encounters will load the
#   minimum of this value and the number of encounters specified in the
#   encounter file.  If specified, this impacts 'n' in the above paragraph.

function load_encounters(filename, initial_dim, update_dim, num_update_type; limit = 0)

    fio = open(filename, "r")

    num_encounters = read(fio, UInt32)
    num_ac = read(fio, UInt32)

    if limit != 0
        num_encounters = min(num_encounters, limit)
    end

    encounters = Array(Any, num_ac, num_encounters)

    for i = 1:num_encounters
        for j = 1:num_ac
            encounters[j, i] = Dict()
        end

        for j = 1:num_ac
            encounters[j, i]["initial"] = read(fio, Float64, initial_dim)
        end

        for j = 1:num_ac
            num_update = read(fio, num_update_type)

            encounters[j, i]["update"] = reshape(read(fio, Float64, update_dim * num_update), Int64(update_dim), Int64(num_update))
        end
    end

    close(fio)

    return encounters, num_ac, num_encounters
end

function load_scripts(filename; limit = 0)

    scripts, num_ac, num_enc = load_encounters(filename, 8, 4, UInt8; limit = limit)
end

if false
    scripts, num_ac, num_enc = load_scripts("encounters_0.dat", limit = 2)

    println(num_ac)
    println(num_enc)
    dump(scripts)
end


using KUparser, KUnet, HDF5, JLD, ArgParse, Compat, CUDArt
VERSION < v"0.4-" && eval(Expr(:using,:Dates))

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "datafiles"
        help = "JLD files for input: first one will be used for training"
        nargs = '+'
        required = true
        "--out"
        help = "JLD file for output trained network"
        default = ""
        "--net"
        help = "Initialize net from this file if specified"
        default = ""
        "--machinefile"
        help = "Need this to restart workers to avoid pmap leak"
        default = ""
        # "--nogpu"  # always use gpu
        # help = "Do not use gpu"
        # action = :store_true
        # "--parser"  # always use bparser
        # help = "Parsing algorithm to use: g(reedy)parser, b(eam)parser, or o(racle)parser for static training"
        # default = "oparser"
        "--arctype"
        help = "Move set to use: ArcEager{R1,13}, ArcHybrid{R1,13}"
        default = "ArcEager13"
        "--feats"
        help = "Features to use from KUparser.Flist"
        default = "tacl13hybrid"
        "--hidden"
        help = "One or more hidden layer sizes"
        nargs = '+'
        arg_type = Int
        default = [20000]
        "--epochs"
        help = "Minimum number of epochs to train"
        arg_type = Int
        default = 100
        # "--pepochs"
        # help = "Number of epochs between parsing"
        # arg_type = Int
        # default = 1
        # "--ncpu"
        # help = "Number of workers for multithreaded parsing"
        # arg_type = Int
        # default = 12
        "--pbatch"
        help = "Minibatch size for parsing"
        arg_type = Int
        default = 100
        "--tbatch"
        help = "Minibatch size for training"
        arg_type = Int
        default = 128
        "--nbeam"
        help = "Beam size for bparser"
        arg_type = Int
        default = 64
        "--adagrad"
        help = "If nonzero apply adagrad using arg as epsilon"
        arg_type = Float32
        nargs = '+'
        default = [1f-8]
        "--dropout"
        help = "Dropout probability"
        arg_type = Float32
        nargs = '+'
        default = [0.2f0, 0.7f0]
        "--l1reg"
        help = "L1 regularization parameter"
        arg_type = Float32
        nargs = '+'
        "--l2reg"
        help = "L2 regularization parameter"
        arg_type = Float32
        nargs = '+'
        "--learningRate"
        help = "Learning rate"
        arg_type = Float32
        nargs = '+'
        default = [2f-2]
        "--maxnorm"
        help = "If nonzero upper limit on weight matrix row norms"
        arg_type = Float32
        nargs = '+'
        "--momentum"
        help = "Momentum"
        arg_type = Float32
        nargs = '+'
        "--nesterov"
        help = "Apply nesterov's accelerated gradient"
        arg_type = Float32
        nargs = '+'
        "--seed"
        help = "Random number seed (-1 leaves the default seed)"
        arg_type = Int
        default = -1
    end
    args = parse_args(s)
    for (arg,val) in args
        print("$arg:$val ")
    end
    println("")
    return args
end

function pbatch(c::Corpus, n::Int)
    dist = Corpus[]
    for s1=1:n:length(c)
        s2=min(length(c), s1+n-1)
        push!(dist, c[s1:s2])
    end
    return dist
end

function evalpxy(pxy, data)
    p = map(a->a[1], pxy)
    evalp(p, data)
end

function evalp(p, d)
    pa = vcat(p...)
    da = vcat(d...)
    evalparse(pa, da)
end

function initdata(args)
    pt = eval(parse(args["arctype"]))
    data = Any[]
    ndeps = 0
    feats = eval(parse("Flist."*args["feats"]))
    net = nothing
    deprel = nothing
    for f in args["datafiles"]
        @date d=load(f)
        push!(data, pbatch(d["corpus"], args["pbatch"]))
        if deprel == nothing
            deprel = d["deprel"]
            ndeps = length(deprel)
        else
            @assert deprel == d["deprel"]
        end
    end
    if isempty(args["net"])
        s1 = data[1][1][1]
        p1 = pt(wcnt(s1),length(deprel))
        xrows=KUparser.flen(p1, s1, feats)
        yrows=p1.nmove
        net=newnet(relu, [xrows; args["hidden"]; yrows]...)
        net[end].f=KUnet.logp
        for k in [fieldnames(UpdateParam); :dropout]
            haskey(args, string(k)) || continue
            v = args[string(k)]
            if isempty(v)
                continue
            elseif length(v)==1
                setparam!(net, k, v[1])
            else
                @assert length(v)==length(net) "$k should have 1 or $(length(net)) elements"
                for i=1:length(v)
                    setparam!(net[i], k, v[i])
                end
            end
        end
        # Initialize net using oparse on first corpus
        @date (p,x,y) = oparse(pt, data[1], ndeps, feats)
        @date train(net, x, y; batch=args["tbatch"], loss=KUnet.logploss, shuffle=true)
    else
        net = newnet(args["net"])
    end
    return (pt, data, ndeps, feats, net)
end

function restart_all_workers(wlist)
    @date Base.terminate_all_workers()
    @date gc()
    @show Base.PGRP.refs
    @date addprocs(wlist)
    @date @everywhere eval(Expr(:using,:KUnet))
    @date @everywhere eval(Expr(:using,:KUparser))
end

function myparse{T<:Parser}(pt::Type{T}, ca::AbstractArray, ndeps::Integer, feats::Fvec, net::Net, args::Dict; trn::Bool=false)
    @date tnet = testnet(net)
    if trn
        @date pxy = pmap(ca) do c
            bparse(pt, c, ndeps, feats, copy(tnet,:gpu), args["nbeam"], args["pbatch"]; xy=true)
        end
        @show e = evalpxy(pxy, ca); flush(STDOUT)
        @date for (p,x,y) in pxy
            train(net, x, y; batch=args["tbatch"], loss=KUnet.logploss, shuffle=true)
        end
    else
        @date p = pmap(ca) do c
            bparse(pt, c, ndeps, feats, copy(tnet,:gpu), args["nbeam"], args["pbatch"])
        end
        @show e = evalp(p, ca); flush(STDOUT)
    end
    return e[1]
end

function rss()
    statm = split(readall("/proc/$(getpid())/statm"))
    pages = int(statm[2])
    mb = pages >> 8
end

function main()
    args = parse_commandline()
    args["seed"] >= 0 && (srand(args["seed"]); KUnet.gpuseed(args["seed"]))
    @date (pt, data, ndeps, feats, net) = initdata(args)
    @show (pt, ndeps, feats)
    @show map(size, data)
    @show net
    accuracy = Array(Float32, length(data))
    (bestscore,bestepoch,epoch)=(0,0,0)
    @show wlist = (isempty(args["machinefile"]) ? [] : split(readall(args["machinefile"])))

    while true
        @show epoch += 1; flush(STDOUT)
        for i=1:length(data)
            @show rss(); flush(STDOUT)
            @date isempty(wlist) || restart_all_workers(wlist)
            @show rss(); flush(STDOUT)
            @date gc()
            @show rss(); flush(STDOUT)
            @date accuracy[i] = myparse(pt, data[i], ndeps, feats, net, args; trn=(i==1))
            @show rss(); flush(STDOUT)
            @date gc()
            @show rss(); flush(STDOUT)
        end
        println("DATA:\t$epoch\t"*join(accuracy, '\t')); flush(STDOUT)

        # We look at accuracy[2] (dev score) by default,
        # accuracy[1] (trn score) if there is no dev.

        score = (length(accuracy) > 1 ? accuracy[2] : accuracy[1])

        # Update best score and save best net

        if (score > bestscore)
            @show (bestscore,bestepoch)=(score,epoch); flush(STDOUT);
            !isempty(args["out"]) && save(args["out"], net)
        end

        # Quit if we have the minimum number of epochs and have
        # not made progress in the last half.

        (epoch >= args["epochs"]) && (epoch >= 2*bestepoch) && break

    end # while true
end # main

main()


using ..Transformers: batchedmul, batched_triu!

# Bert specific initializers

FakeTHLinear(cfg::AbstractHGFConfig, args...; kwargs...) = FakeTHLinear(typeof(cfg), cfg, args...; kwargs...)
function FakeTHLinear(::Type{HGFBertConfig}, config, hi, ho; bias=true)
  weight = randn(Float32, ho, hi) .* config.initializer_range
  if bias
    bias = zeros(Float32, ho)
  else
    bias = nothing
  end
  FakeTHLinear(weight, bias)
end

FakeTHEmbedding(cfg::AbstractHGFConfig, args...; kwargs...) = FakeTHEmbedding(typeof(cfg), cfg, args...; kwargs...)
function FakeTHEmbedding(::Type{HGFBertConfig}, config, num, dims; pad_idx=nothing)
  weight = randn(Float32, dims, num) .* config.initializer_range
  if !isnothing(pad_idx)
    real_pad_idx = pad_idx+1
  else
    real_pad_idx = 0
  end
  FakeTHEmbedding(real_pad_idx, weight)
end

FakeTHLayerNorm(cfg::AbstractHGFConfig, args...; kwargs...) = FakeTHLayerNorm(typeof(cfg), cfg, args...; kwargs...)
function FakeTHLayerNorm(::Type{HGFBertConfig}, config, dims; eps::Float32=1e-05)
  weight = ones(Float32, dims)
  bias = zeros(Float32, dims)
  FakeTHLayerNorm(eps, weight, bias)
end

# HGF Bert compounts

# embedding

struct HGFBertEmbeddings{
  W<:FakeTHEmbedding,
  P<:FakeTHEmbedding,
  T<:FakeTHEmbedding,
  L<:FakeTHLayerNorm
} <: THModule
  LayerNorm::L
  word_embeddings::W
  position_embeddings::P
  token_type_embeddings::T
end

@functor HGFBertEmbeddings

HGFBertEmbeddings(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertEmbeddings(typeof(cfg), cfg, args...; kwargs...)
function HGFBertEmbeddings(T::Type{HGFBertConfig}, config)
  layernorm = FakeTHLayerNorm(T, config, config.hidden_size; eps=config.layer_norm_eps)

  word_emb = FakeTHEmbedding(T, config, config.vocab_size,
                             config.hidden_size;
                             pad_idx=config.pad_token_id)
  posi_emb = FakeTHEmbedding(T, config, config.max_position_embeddings,
                             config.hidden_size)
  toke_emb = FakeTHEmbedding(T, config, config.type_vocab_size,
                             config.hidden_size)

  HGFBertEmbeddings(layernorm, word_emb, posi_emb, toke_emb)
end

_arange(x, len) = cumsum(fill!(similar(x, Int, len), one(Int)))
_arange(x, shift, len) = _arange(x, len) .+ Int(shift)

Flux.@nograd _arange

@inline get_word_emb(be::HGFBertEmbeddings, input_ids::AbstractArray{<:Integer}) = be.word_embeddings(input_ids)
@inline get_word_emb(be::HGFBertEmbeddings, input_embed::AbstractArray{T}) where T = input_embed

@inline get_position_emb(be::HGFBertEmbeddings, inputs_embeds, ::Nothing) = get_position_emb(be, inputs_embeds, _arange(inputs_embeds, size(inputs_embeds)[end-1]))
@inline get_position_emb(be::HGFBertEmbeddings, inputs_embeds, position_ids) = be.position_embeddings(position_ids)

@inline get_token_type_emb(be::HGFBertEmbeddings, inputs_embeds, ::Nothing) = get_token_type_emb(be, inputs_embeds, fill!(similar(inputs_embeds, Int, Base.tail(size(inputs_embeds))), one(Int)))
@inline get_token_type_emb(be::HGFBertEmbeddings, inputs_embeds, token_type_ids) = be.token_type_embeddings(token_type_ids)

function (be::HGFBertEmbeddings)(input,
                                 position_ids::Union{Nothing, AbstractArray{<:Integer}},
                                 token_type_ids::Union{Nothing, AbstractArray{<:Integer}})

  inputs_embeds = get_word_emb(be, input)
  position_embeds = get_position_emb(be, inputs_embeds, position_ids)
  token_type_embeds = get_token_type_emb(be, inputs_embeds, token_type_ids)
  return be(inputs_embeds, position_embeds, token_type_embeds)
end

function (be::HGFBertEmbeddings)(inputs_embeds::AbstractArray{T},
                                 position_embeds::AbstractArray{T},
                                 token_type_embeds::AbstractArray{T}) where T
  embeddings = (inputs_embeds .+ position_embeds) + token_type_embeds
  embeddings = be.LayerNorm(embeddings)
  return embeddings
end

(be::HGFBertEmbeddings)(input; position_ids=nothing, token_type_ids=nothing) = be(input, position_ids, token_type_ids)

# self attention part

struct HGFBertSelfAttention{
  Q<:FakeTHLinear,
  K<:FakeTHLinear,
  V<:FakeTHLinear
} <: THModule
  num_attention_heads::Int
  query::Q
  key::K
  value::V
end

Functors.functor(::Type{<:HGFBertSelfAttention}, sa) = (query=sa.query, key=sa.key, value=sa.value), y->HGFBertSelfAttention(sa.num_attention_heads, y...)

function _split_tranpose_for_scores(x, num_head)
  head_size = div(size(x, 1), num_head)
  split_size = (head_size, num_head, size(x, 2), size(x, 3))
  permute_order = (1, 3, 2, 4)

  return reshape(x, split_size) |> Base.Fix2(permutedims, permute_order)
end

function _compute_attention_scores(query_layer, key_layer, attention_mask::Union{Nothing, <:AbstractArray})
  attentions_scores = batchedmul(key_layer, query_layer; transA = true)
  attentions_scores = attentions_scores ./ convert(eltype(attentions_scores), sqrt(size(key_layer, 1)))

  !isnothing(attention_mask) &&
    (attentions_scores = attentions_scores .+ attention_mask)

  return attentions_scores
end

function _merge_transpose_for_output(x, num_head)
  permute_order = (1, 3, 2, 4)
  final_size = (:, size(x, 2), size(x, 4))

  return permutedims(x, permute_order) |>
    Base.Fix2(reshape, final_size)
end

function (sa::HGFBertSelfAttention)(hidden_states, attention_mask, output_attentions::Val)
  mixed_query_layer = sa.query(hidden_states)
  mixed_key_layer = sa.key(hidden_states)
  mixed_value_layer = sa.value(hidden_states)

  return sa(mixed_query_layer, mixed_key_layer, mixed_value_layer, attention_mask, output_attentions)
end

function (sa::HGFBertSelfAttention)(hidden_states, encoder_hidden_states, attention_mask, output_attentions::Val)
  mixed_query_layer = sa.query(hidden_states)
  mixed_key_layer = sa.key(encoder_hidden_states)
  mixed_value_layer = sa.value(encoder_hidden_states)

  return sa(mixed_query_layer, mixed_key_layer, mixed_value_layer, attention_mask, output_attentions)
end

function (sa::HGFBertSelfAttention)(mixed_query_layer, mixed_key_layer, mixed_value_layer,
                                    attention_mask,
                                    ::Val{output_attentions}) where output_attentions

  query_layer = _split_tranpose_for_scores(mixed_query_layer, sa.num_attention_heads)
  key_layer = _split_tranpose_for_scores(mixed_key_layer, sa.num_attention_heads)
  value_layer = _split_tranpose_for_scores(mixed_value_layer, sa.num_attention_heads)

  attentions_scores = _compute_attention_scores(query_layer, key_layer, attention_mask)
  attentions_probs = softmax(attentions_scores; dims=1)

  mixed_context_layer = batchedmul(value_layer, attentions_probs)
  context_layer = _merge_transpose_for_output(mixed_context_layer, sa.num_attention_heads)

  if output_attentions
    outputs = context_layer, attentions_probs
    return outputs
  else
    return context_layer
  end
end

HGFBertSelfAttention(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertSelfAttention(typeof(cfg), cfg, args...; kwargs...)
function HGFBertSelfAttention(T::Type{HGFBertConfig}, config)
  attention_head_size = config.hidden_size ÷ config.num_attention_heads
  all_head_size = config.num_attention_heads * attention_head_size

  query = FakeTHLinear(T, config, config.hidden_size, all_head_size)
  key   = FakeTHLinear(T, config, config.hidden_size, all_head_size)
  value = FakeTHLinear(T, config, config.hidden_size, all_head_size)

  HGFBertSelfAttention(config.num_attention_heads, query, key, value)
end

# self attention output part

struct HGFBertSelfOutput{
  L<:FakeTHLayerNorm,
  D<:FakeTHLinear
} <: THModule
  LayerNorm::L
  dense::D
end

@functor HGFBertSelfOutput

function (so::HGFBertSelfOutput)(hidden_states, input_tensor)
  hidden_states = so.dense(hidden_states)
  hidden_states = so.LayerNorm(hidden_states + input_tensor)
  return hidden_states
end

HGFBertSelfOutput(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertSelfOutput(typeof(cfg), cfg, args...; kwargs...)
function HGFBertSelfOutput(T::Type{HGFBertConfig}, config)
  layernorm = FakeTHLayerNorm(T, config, config.hidden_size; eps=config.layer_norm_eps)
  dense = FakeTHLinear(T, config, config.hidden_size, config.hidden_size)
  HGFBertSelfOutput(layernorm, dense)
end

# self attention

struct HGFBertAttention{
  S<:HGFBertSelfAttention,
  O<:HGFBertSelfOutput
} <: THModule
  self::S
  output::O
end

@functor HGFBertAttention

function (a::HGFBertAttention)(hidden_states,
                               attention_mask::Union{Nothing, <:AbstractArray},
                               _output_attentions::Val{output_attentions}) where output_attentions

  self_output = a.self(hidden_states, attention_mask, _output_attentions)

  if output_attentions
    output, attention_prob = self_output
    attention_output = a.output(output, hidden_states)
    return attention_output, attention_prob
  else
    attention_output = a.output(self_output, hidden_states)
    return attention_output
  end
end

HGFBertAttention(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertAttention(typeof(cfg), cfg, args...; kwargs...)
function HGFBertAttention(T::Type{HGFBertConfig}, config)
  self = HGFBertSelfAttention(T, config)
  output = HGFBertSelfOutput(T, config)
  HGFBertAttention(self, output)
end

# positionwise first dense

struct HGFBertIntermediate{F, D<:FakeTHLinear} <: THModule
  intermediate_act::F
  dense::D
end

Functors.functor(::Type{<:HGFBertIntermediate}, intermediate) = (dense = intermediate.dense,), y->HGFBertIntermediate(intermediate.intermediate_act, y...)

(i::HGFBertIntermediate)(hidden_states) = i.intermediate_act.(i.dense(hidden_states))

HGFBertIntermediate(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertIntermediate(typeof(cfg), cfg, args...; kwargs...)
function HGFBertIntermediate(T::Type{HGFBertConfig}, config)
  global ACT2FN
  act = ACT2FN[Symbol(config.hidden_act)]
  dense = FakeTHLinear(T, config, config.hidden_size, config.intermediate_size)
  HGFBertIntermediate(act, dense)
end

# positionwise second dense

struct HGFBertOutput{
  D<:FakeTHLinear,
  L<:FakeTHLayerNorm
} <: THModule
  dense::D
  LayerNorm::L
end

@functor HGFBertOutput

function (o::HGFBertOutput)(hidden_states, input_tensor)
  hidden_states = o.dense(hidden_states)
  hidden_states = o.LayerNorm(hidden_states + input_tensor)
  return hidden_states
end

HGFBertOutput(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertOutput(typeof(cfg), cfg, args...; kwargs...)
function HGFBertOutput(T::Type{HGFBertConfig}, config)
  dense = FakeTHLinear(T, config, config.intermediate_size, config.hidden_size)
  layernorm = FakeTHLayerNorm(T, config, config.hidden_size; eps=config.layer_norm_eps)
  HGFBertOutput(dense, layernorm)
end

# transformer layer

struct HGFBertLayer{DEC<:Union{Nothing, HGFBertAttention},
                    A<:HGFBertAttention,
                    I<:HGFBertIntermediate,
                    O<:HGFBertOutput
} <: THModule
  attention::A
  crossattention::DEC
  intermediate::I
  output::O
end

HGFBertLayer(a, i, o) = HGFBertLayer(a, nothing, i, o)

_is_decode(::HGFBertLayer{Nothing}) = false
_is_decode(::HGFBertLayer) = true

Functors.functor(::Type{<:HGFBertLayer}, layer) = (!_is_decode(layer) ?
    (attention = layer.attention, intermediate = layer.intermediate, output = layer.output) :
    (attention = layer.attention, crossattention = layer.crossattention, intermediate = layer.intermediate, output = layer.output)),
    y ->HGFBertLayer(y...)

function (l::HGFBertLayer{Nothing})(hidden_states, attention_mask::Union{Nothing, <:AbstractArray},
                                    _output_attentions::Val{output_attentions}) where output_attentions
  if output_attentions
    attention_output, attention_prob = l.attention(hidden_states, attention_mask, _output_attentions)
  else
    attention_output = l.attention(hidden_states, attention_mask, _output_attentions)
  end

  intermediate_output = l.intermediate(attention_output)
  layer_output = l.output(intermediate_output, attention_output)

  if output_attentions
    return layer_output, attention_prob
  else
    return layer_output
  end
end

function (l::HGFBertLayer)(hidden_states, attention_mask,
                           encoder_hidden_states, encoder_attention_mask,
                           _output_attentions::Val{output_attentions}) where output_attentions
  if output_attentions
    attention_output, attention_prob = l.attention(hidden_states, attention_mask, _output_attentions)
    attention_output, cross_attention_prob = l.crossattention(attention_output, encoder_attention_mask, _output_attentions)
  else
    attention_output = l.attention(hidden_states, attention_mask, _output_attentions)
    attention_output = l.crossattention(attention_output, encoder_hidden_states, encoder_attention_mask, _output_attentions)
  end

  intermediate_output = l.intermediate(attention_output)
  layer_output = l.output(intermediate_output, attention_output)

  if output_attentions
    return layer_output, attention_prob, cross_attention_prob
  else
    return layer_output
  end
end

HGFBertLayer(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertLayer(typeof(cfg), cfg, args...; kwargs...)
function HGFBertLayer(T::Type{HGFBertConfig}, config)
  attention = HGFBertAttention(T, config)
  crossattention = config.is_decoder ?
    HGFBertAttention(T, config) :
    nothing
  intermediate = HGFBertIntermediate(T, config)
  output = HGFBertOutput(T, config)
  HGFBertLayer(attention, crossattention, intermediate, output)
end

# stacked transformers

struct HGFBertEncoder{N, L<:FakeTHModuleList{N}} <: THModule
  layer::L
end

@functor HGFBertEncoder

(e::HGFBertEncoder)(hidden_states;
                    attention_mask = nothing,
                    output_attentions = false,
                    output_hidden_states = false
                    ) = e(hidden_states, attention_mask, Val(output_attentions), Val(output_hidden_states))

@generated function (e::HGFBertEncoder{N})(hidden_states, attention_mask,
                             _output_attentions::Val{output_attentions},
                             _output_hidden_states::Val{output_hidden_states}
                             ) where {N, output_attentions, output_hidden_states}
  if output_attentions
    all_attentions = Expr(:tuple)
  end

  if output_hidden_states
    all_hidden_states = Expr(:tuple, :hidden_states)
  end

  body = Expr[]

  for i = 1:N
    previous = i == 1 ? :hidden_states : Symbol(:hidden_states, i-1)
    current = Symbol(:hidden_states, i)
    if output_attentions
      current_attention = Symbol(:attention_, i)
      current_output = :($current, $current_attention)
    else
      current_output = current
    end

    expr = :($current_output = e.layer[$i]($previous, attention_mask, _output_attentions))
    push!(body, expr)

    if output_attentions
      push!(all_attentions.args, current_attention)
    end

    if output_hidden_states
      push!(all_hidden_states.args, current)
    end
  end

  current = Symbol(:hidden_states, N)

  if output_attentions
    push!(body, :(all_attentions = $all_attentions))
  else
    push!(body, :(all_attentions = nothing))
  end

  if output_hidden_states
    push!(body, :(all_hidden_states = $all_hidden_states))
  else
    push!(body, :(all_hidden_states = nothing))
  end

  return quote
    $(body...)

    return (
      last_hidden_state = $current,
      hidden_states = all_hidden_states,
      attentions = all_attentions
    )
  end
end

(e::HGFBertEncoder)(hidden_states, encoder_hidden_states;
                    attention_mask = nothing,
                    encoder_mask = nothing,
                    output_attentions = false,
                    output_hidden_states = false
                    ) = e(hidden_states, attention_mask,
                          encoder_hidden_states, encoder_mask,
                          Val(output_attentions), Val(output_hidden_states))

@generated function (e::HGFBertEncoder{N})(hidden_states, attention_mask,
                                           encoder_hidden_states, encoder_mask,
                                           _output_attentions::Val{output_attentions},
                                           _output_hidden_states::Val{output_hidden_states}
                                           ) where {N, output_attentions, output_hidden_states}
  if output_attentions
    all_attentions = Expr(:tuple)
  end

  if output_hidden_states
    all_hidden_states = Expr(:tuple, :hidden_states)
  end

  body = Expr[]

  for i = 1:N
    previous = i == 1 ? :hidden_states : Symbol(:hidden_states, i-1)
    current = Symbol(:hidden_states, i)
    if output_attentions
      current_attention = Symbol(:attention_, i)
      current_output = :($current, $current_attention)
    else
      current_output = current
    end

    expr = :($current_output = e.layer[$i]($previous, attention_mask,
                                           encoder_hidden_states, encoder_mask,
                                           _output_attentions))
    push!(body, expr)

    if output_attentions
      push!(all_attentions.args, current_attention)
    end

    if output_hidden_states
      push!(all_hidden_states.args, current)
    end
  end

  current = Symbol(:hidden_states, N)

  if output_attentions
    push!(body, :(all_attentions = $all_attentions))
  else
    push!(body, :(all_attentions = nothing))
  end

  if output_hidden_states
    push!(body, :(all_hidden_states = $all_hidden_states))
  else
    push!(body, :(all_hidden_states = nothing))
  end

  return quote
    $(body...)

    return (
      last_hidden_state = $current,
      hidden_states = all_hidden_states,
      attentions = all_attentions
    )
  end
end

HGFBertEncoder(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertEncoder(typeof(cfg), cfg, args...; kwargs...)
function HGFBertEncoder(T::Type{HGFBertConfig}, config)
  layer = FakeTHModuleList(
    [HGFBertLayer(T, config) for _ in 1:config.num_hidden_layers]
  )
  HGFBertEncoder(layer)
end

# classify token projection

struct HGFBertPooler{D<:FakeTHLinear} <: THModule
  dense::D
end

@functor HGFBertPooler

function (p::HGFBertPooler)(hidden_states)
  first_token_tensor = hidden_states[:, 1, :]
  pooled_output = tanh.(p.dense(first_token_tensor))
  return pooled_output
end

HGFBertPooler(cfg::AbstractHGFConfig, args...; kwargs...) = HGFBertPooler(typeof(cfg), cfg, args...; kwargs...)
function HGFBertPooler(T::Type{HGFBertConfig}, config)
  dense = FakeTHLinear(T, config, config.hidden_size, config.hidden_size)
  HGFBertPooler(dense)
end

# label prediction layer

struct HGFBertPredictionHeadTransform{
  F,
  D<:FakeTHLinear,
  L<:FakeTHLayerNorm
} <: THModule
  transform_act_fn::F
  dense::D
  LayerNorm::L
end

@functor HGFBertPredictionHeadTransform

function (pht::HGFBertPredictionHeadTransform)(hidden_states)
  hidden_states = pht.dense(hidden_states)
  hidden_states = pht.transform_act_fn.(hidden_states)
  hidden_states = pht.LayerNorm(hidden_states)
  return hidden_states
end

function HGFBertPredictionHeadTransform(config::HGFBertConfig)
  global ACT2FN
  dense = FakeTHLinear(config, config.hidden_size, config.hidden_size)
  act = ACT2FN[Symbol(config.hidden_act)]
  layernorm = FakeTHLayerNorm(config, config.hidden_size; eps=config.layer_norm_eps)

  return HGFBertPredictionHeadTransform(act, dense, layernorm)
end

# language model prediction layer

struct HGFBertLMPredictionHead{
  B<:AbstractArray,
  T<:HGFBertPredictionHeadTransform,
  D<:FakeTHLinear
} <: THModule
  transform::T
  decoder::D
  bias::B
end

@functor HGFBertLMPredictionHead

function (ph::HGFBertLMPredictionHead)(hidden_states)
  hidden_states = ph.transform(hidden_states)
  hidden_states = ph.decoder(hidden_states) .+ ph.bias
  return hidden_states
end

function HGFBertLMPredictionHead(config::HGFBertConfig; input_embedding=nothing)
  trans = HGFBertPredictionHeadTransform(config)
  if isnothing(input_embedding)
    decoder = FakeTHLinear(config, config.hidden_size, config.vocab_size; bias=false)
  else
    decoder = FakeTHLinear(Transpose(input_embedding), nothing)
  end
  bias = zeros(Float32, config.vocab_size)
  return HGFBertLMPredictionHead(trans, decoder, bias)
end

# language model prediction wrapper

struct HGFBertOnlyMLMHead{P<:HGFBertLMPredictionHead} <: THModule
  predictions::P
end

@functor HGFBertOnlyMLMHead

(h::HGFBertOnlyMLMHead)(sequence_output) = h.predictions(sequence_output)

function HGFBertOnlyMLMHead(config::HGFBertConfig; input_embedding=nothing)
  predictions = HGFBertLMPredictionHead(config; input_embedding=input_embedding)
  return HGFBertOnlyMLMHead(predictions)
end

# next sentence prediction layer

struct HGFBertOnlyNSPHead{S<:FakeTHLinear} <: THModule
  seq_relationship::S
end

@functor HGFBertOnlyNSPHead

(h::HGFBertOnlyNSPHead)(pooled_output) = h.seq_relationship(pooled_output)

function HGFBertOnlyNSPHead(config::HGFBertConfig)
  seq_relationship = FakeTHLinear(config, config.hidden_size, 2)
  return HGFBertOnlyNSPHead(seq_relationship)
end

# pretrain prediction layers

struct HGFBertPreTrainingHeads{
  P<:HGFBertLMPredictionHead,
  S<:FakeTHLinear
} <: THModule
  predictions::P
  seq_relationship::S
end

@functor HGFBertPreTrainingHeads

(pth::HGFBertPreTrainingHeads)(sequence_output, pooled_output) = pth.predictions(sequence_output), pth.seq_relationship(pooled_output)

function HGFBertPreTrainingHeads(config::HGFBertConfig; input_embedding=nothing)
  predictions = HGFBertLMPredictionHead(config; input_embedding=input_embedding)
  seq_relationship = FakeTHLinear(config, config.hidden_size, 2)
  return HGFBertPreTrainingHeads(predictions, seq_relationship)
end

# bert model without prediction
abstract type HGFBertPreTrainedModel <: HGFPreTrainedModel end

struct HGFBertModel{
  E<:HGFBertEmbeddings,
  T<:HGFBertEncoder,
  P<:HGFBertPooler
} <: HGFBertPreTrainedModel
  embeddings::E
  encoder::T
  pooler::P
end

@functor HGFBertModel

function maybe_prepare_mask(embedding_output, ::Nothing)
  mask_size = size(embedding_output) |> Base.tail
  fill!(similar(embedding_output, mask_size), 1) |> maybe_prepare_mask
end
maybe_prepare_mask(embedding_output, attention_mask) = maybe_prepare_mask(attention_mask)
maybe_prepare_mask(embedding_output, attention_mask::AbstractArray{T, 4}) where T = attention_mask

function maybe_prepare_mask(attention_mask::AbstractArray{T, 3}) where T
  seq_len1, seq_len2, batch_size = size(attention_mask)
  attention_mask = reshape(attention_mask, seq_len1, seq_len2, 1, batch_size)
  return attention_mask
end

function maybe_prepare_mask(attention_mask::AbstractMatrix)
  seq_len, batch_size = size(attention_mask)
  attention_mask = reshape(attention_mask, seq_len, 1, 1, batch_size)
  return attention_mask
end

function create_attention_mask(attention_mask::AbstractArray{T, 4}) where T
  return (one(T) .- attention_mask) .* convert(T, -10000)
end

function create_attention_mask(embedding_output, attention_mask)
  attention_mask = maybe_prepare_mask(embedding_output, attention_mask)
  create_attention_mask(attention_mask)
end

maybe_prepare_causal_mask(embedding_output, attention_mask::AbstractArray{T, 4}) where T = attention_mask
maybe_prepare_causal_mask(embedding_output, attention_mask::AbstractArray{T, 3}) where T = maybe_prepare_mask(embedding_output, attention_mask)
function maybe_prepare_causal_mask(embedding_output, attention_mask::Union{Nothing, AbstractMatrix})
  regular_mask = maybe_prepare_mask(embedding_output, attention_mask) # seq_len, 1, 1, b
  attention_mask = permutedims(regular_mask, (2,1,3,4)) .* regular_mask # seq_len, seq_len, 1, b
  return attention_mask
end

function create_causal_attention_mask(embedding_output, attention_mask)
  attention_mask = maybe_prepare_causal_mask(embedding_output, attention_mask)
  batched_triu!(attention_mask, 0)
  return create_attention_mask(attention_mask)
end

Flux.@nograd create_attention_mask
Flux.@nograd create_causal_attention_mask

(bm::HGFBertModel)(input; position_ids = nothing, token_type_ids = nothing,
                   attention_mask = nothing,
                   output_attentions = false,
                   output_hidden_states = false
                   ) = bm(input, position_ids, token_type_ids, attention_mask,
                          Val(output_attentions), Val(output_hidden_states))

function (bm::HGFBertModel)(
  input, position_ids, token_type_ids,
  attention_mask,
  _output_attentions::Val{output_attentions},
  _output_hidden_states::Val{output_hidden_states}
) where {output_attentions, output_hidden_states}

  embedding_output = bm.embeddings(input, position_ids, token_type_ids)
  attention_mask = create_attention_mask(embedding_output, attention_mask)

  encoder_outputs = bm.encoder(embedding_output, attention_mask,
                               _output_attentions, _output_hidden_states)

  sequence_output = encoder_outputs.last_hidden_state
  pooled_output = bm.pooler(sequence_output)
  return (
    last_hidden_state = sequence_output,
    pooler_output = pooled_output,
    hidden_states = encoder_outputs.hidden_states,
    attentions = encoder_outputs.attentions
  )
end

(bm::HGFBertModel)(input, encoder_hidden_states;
                   position_ids = nothing, token_type_ids = nothing,
                   attention_mask = nothing, encoder_attention_mask = nothing,
                   output_attentions = false,
                   output_hidden_states = false
                   ) = bm(input, encoder_hidden_states,
                          position_ids, token_type_ids,
                          attention_mask, encoder_attention_mask,
                          Val(output_attentions), Val(output_hidden_states))

function (bm::HGFBertModel)(
  input, encoder_hidden_states, position_ids, token_type_ids,
  attention_mask, encoder_attention_mask,
  _output_attentions::Val{output_attentions},
  _output_hidden_states::Val{output_hidden_states}
) where {output_attentions, output_hidden_states}

  embedding_output = bm.embeddings(input, position_ids, token_type_ids)
  attention_mask = create_causal_attention_mask(embedding_output, attention_mask)
  encoder_attention_mask = create_attention_mask(encoder_hidden_states, encoder_attention_mask)

  encoder_outputs = bm.encoder(embedding_output, attention_mask,
                               encoder_hidden_states, encoder_attention_mask,
                               _output_attentions, _output_hidden_states)

  sequence_output = encoder_outputs.last_hidden_state
  pooled_output = bm.pooler(sequence_output)

  return (
    last_hidden_state = sequence_output,
    pooler_output = pooled_output,
    hidden_states = encoder_outputs.hidden_states,
    attentions = encoder_outputs.attentions
  )
end

function HGFBertModel(config::HGFBertConfig)
  embeddings = HGFBertEmbeddings(config)
  encoder = HGFBertEncoder(config)
  pooler = HGFBertPooler(config)
  HGFBertModel(embeddings, encoder, pooler)
end

get_input_embedding(model::HGFBertModel) = model.embeddings.word_embeddings.weight

# bert models for different task

# pretrain

struct HGFBertForPreTraining{B<:HGFBertModel, C<:HGFBertPreTrainingHeads} <: HGFBertPreTrainedModel
  bert::B
  cls::C
end

@functor HGFBertForPreTraining

(self::HGFBertForPreTraining)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForPreTraining)(input, position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output, pooled_output = outputs
  prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
  total_loss = nothing

  return (
    loss = total_loss,
    prediction_logits = prediction_scores,
    seq_relationship_logits = seq_relationship_score,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForPreTraining)(input, labels, next_sentence_label;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, labels, next_sentence_label,
                                       position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForPreTraining)(input, labels, next_sentence_label,
                                       position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  prediction_scores = outputs.prediction_logits
  seq_relationship_score = outputs.seq_relationship_logits

  masked_lm_loss = Flux.logitcrossentropy(prediction_scores, labels)
  next_sentence_loss = Flux.logitcrossentropy(seq_relationship_score, next_sentence_label)
  total_loss = masked_lm_loss + next_sentence_loss

  return (
    loss = total_loss,
    prediction_logits = prediction_scores,
    seq_relationship_logits = seq_relationship_score,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForPreTraining(config::HGFBertConfig)
  bert = HGFBertModel(config)
  input_embedding = get_input_embedding(bert)
  cls = HGFBertPreTrainingHeads(config; input_embedding=input_embedding)
  return HGFBertForPreTraining(bert, cls)
end

# clm finetune

struct HGFBertLMHeadModel{B<:HGFBertModel, C<:HGFBertOnlyMLMHead} <: HGFBertPreTrainedModel
  bert::B
  cls::C
end

@functor HGFBertLMHeadModel

(self::HGFBertLMHeadModel)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertLMHeadModel)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  prediction_scores = self.cls(sequence_output)
  lm_loss = nothing

  return (
    loss = lm_loss,
    prediction_logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertLMHeadModel)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertLMHeadModel)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  prediction_scores = outputs.prediction_logits

  shifted_prediction_scores = prediction_scores[:, 1:end-1, :]
  shifted_labels = labels[:, 2:end]

  lm_loss = Flux.logitcrossentropy(shifted_prediction_scores, shifted_labels)

  return (
    loss = lm_loss,
    prediction_logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertLMHeadModel(config::HGFBertConfig)
  bert = HGFBertModel(config)
  input_embedding = get_input_embedding(bert)
  cls = HGFBertOnlyMLMHead(config; input_embedding=input_embedding)
  return HGFBertLMHeadModel(bert, cls)
end

# maked lm

struct HGFBertForMaskedLM{B<:HGFBertModel, C<:HGFBertOnlyMLMHead} <: HGFBertPreTrainedModel
  bert::B
  cls::C
end

@functor HGFBertForMaskedLM

(self::HGFBertForMaskedLM)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForMaskedLM)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  prediction_scores = self.cls(sequence_output)
  masked_lm_loss = nothing

  return (
    loss = masked_lm_loss,
    logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForMaskedLM)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForMaskedLM)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  prediction_scores = outputs.logits
  masked_lm_loss = Flux.logitcrossentropy(prediction_scores, labels)

  return (
    loss = masked_lm_loss,
    logits = prediction_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForMaskedLM(config::HGFBertConfig)
  bert = HGFBertModel(config)
  input_embedding = get_input_embedding(bert)
  cls = HGFBertOnlyMLMHead(config; input_embedding=input_embedding)
  return HGFBertForMaskedLM(bert, cls)
end

# next sentence

struct HGFBertForNextSentencePrediction{B<:HGFBertModel, C<:HGFBertOnlyNSPHead} <: HGFBertPreTrainedModel
  bert::B
  cls::C
end

@functor HGFBertForNextSentencePrediction

(self::HGFBertForNextSentencePrediction)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForNextSentencePrediction)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  pooled_output = outputs[2]
  seq_relationship_scores = self.cls(pooled_output)
  next_sentence_loss = nothing

  return (
    loss = next_sentence_loss,
    logits = seq_relationship_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForNextSentencePrediction)(input, next_sentence_label;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, next_sentence_label,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForNextSentencePrediction)(input, next_sentence_label,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  seq_relationship_scores = outputs.logits
  next_sentence_loss = Flux.logitcrossentropy(seq_relationship_scores, next_sentence_label)

  return (
    loss = next_sentence_loss,
    logits = seq_relationship_scores,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForNextSentencePrediction(config::HGFBertConfig)
  bert = HGFBertModel(config)
  cls = HGFBertOnlyNSPHead(config)
  return HGFBertForNextSentencePrediction(bert, cls)
end

# seq classify

struct HGFBertForSequenceClassification{B<:HGFBertModel, C<:FakeTHLinear} <: HGFBertPreTrainedModel
  bert::B
  classifier::C
end

@functor HGFBertForSequenceClassification

(self::HGFBertForSequenceClassification)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForSequenceClassification)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  pooled_output = outputs[2]
  logits = self.classifier(pooled_output)
  loss = nothing

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForSequenceClassification)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForSequenceClassification)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  logits = outputs.logits
  if size(logits, 1) == 1
    loss = Flux.mse(logits, labels)
  else
    loss = Flux.logitcrossentropy(logits, labels)
  end

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForSequenceClassification(config::HGFBertConfig)
  bert = HGFBertModel(config)
  classifier = FakeTHLinear(config, config.hidden_size, config.num_labels)
  HGFBertForSequenceClassification(bert, classifier)
end

# multiple choice

struct HGFBertForMultipleChoice{B<:HGFBertModel, C<:FakeTHLinear} <: HGFBertPreTrainedModel
  bert::B
  classifier::C
end

@functor HGFBertForMultipleChoice

(self::HGFBertForMultipleChoice)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForMultipleChoice)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                          ) where {output_attentions, output_hidden_states}

  num_choices = size(input, ndims(input)-1)
  flat_choice(x) = reshape(x, size(x)[1:end-2]..., :)
  flat_choice(::Nothing) = nothing
  flat_input = flat_choice(input)
  flat_position_ids = flat_choice(position_ids)
  flat_token_type_ids = flat_choice(token_type_ids)
  flat_attention_mask = flat_choice(attention_mask)

  outputs = self.bert(flat_input, flat_position_ids, flat_token_type_ids,
                      flat_attention_mask, _output_attentions, _output_hidden_states)

  pooled_output = outputs[2]
  logits = self.classifier(pooled_output)
  reshaped_logits = reshape(logits, num_choices, :)
  loss = nothing

  return (
    loss = loss,
    logits = reshaped_logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForMultipleChoice)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForMultipleChoice)(input, labels,
                                          position_ids, token_type_ids,
                                          attention_mask,
                                          _output_attentions::Val{output_attentions},
                                          _output_hidden_states::Val{output_hidden_states}
                                          ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                 _output_attentions, _output_hidden_states)

  logits = outputs.logits
  loss = Flux.logitcrossentropy(logits, labels)

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForMultipleChoice(config::HGFBertConfig)
  bert = HGFBertModel(config)
  classifier = FakeTHLinear(config, config.hidden_size, 1)
  HGFBertForMultipleChoice(bert, classifier)
end

# token classify

struct HGFBertForTokenClassification{B<:HGFBertModel, C<:FakeTHLinear} <: HGFBertPreTrainedModel
  bert::B
  classifier::C
end

@functor HGFBertForTokenClassification

(self::HGFBertForTokenClassification)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForTokenClassification)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  logits = self.classifier(sequence_output)
  loss = nothing

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForTokenClassification)(input, labels;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, labels,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForTokenClassification)(input, labels,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  logits = outputs.logits
  loss = Flux.logitcrossentropy(logits, labels)

  return (
    loss = loss,
    logits = logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForTokenClassification(config::HGFBertConfig)
  bert = HGFBertModel(config)
  classifier = FakeTHLinear(config, config.hidden_size, config.num_labels)
  HGFBertForTokenClassification(bert, classifier)
end

# qa

struct HGFBertForQuestionAnswering{B<:HGFBertModel, C<:FakeTHLinear} <: HGFBertPreTrainedModel
  bert::B
  qa_outputs::C
end

@functor HGFBertForQuestionAnswering

(self::HGFBertForQuestionAnswering)(input;
                              position_ids = nothing, token_type_ids = nothing,
                              attention_mask = nothing,
                              output_attentions = false,
                              output_hidden_states = false
                              ) = self(input, position_ids, token_type_ids,
                                       attention_mask,
                                       Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForQuestionAnswering)(input, position_ids, token_type_ids,
                                    attention_mask,
                                    _output_attentions::Val{output_attentions},
                                    _output_hidden_states::Val{output_hidden_states}
                                    ) where {output_attentions, output_hidden_states}
  outputs = self.bert(input, position_ids, token_type_ids,
                      attention_mask, _output_attentions, _output_hidden_states)
  sequence_output = outputs[1]
  logits = self.qa_outputs(sequence_output)
  start_logits = @view(logits[1, :])
  end_logits = @view(logits[2, :])
  total_loss = nothing

  return (
    loss = total_loss,
    start_logits = start_logits,
    end_logits = end_logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

(self::HGFBertForQuestionAnswering)(input, start_positions, end_positions;
                           position_ids = nothing, token_type_ids = nothing,
                           attention_mask = nothing,
                           output_attentions = false,
                           output_hidden_states = false
                           ) = self(input, start_positions, end_positions,
                                    position_ids, token_type_ids,
                                    attention_mask,
                                    Val(output_attentions), Val(output_hidden_states))

function (self::HGFBertForQuestionAnswering)(input, start_positions, end_positions,
                                    position_ids, token_type_ids,
                                       attention_mask,
                                       _output_attentions::Val{output_attentions},
                                       _output_hidden_states::Val{output_hidden_states}
                                       ) where {output_attentions, output_hidden_states}
  outputs = self(input, position_ids, token_type_ids, attention_mask,
                _output_attentions, _output_hidden_states)

  start_logits = outputs.start_logits
  end_logits = outputs.end_logits
  start_loss = Flux.logitcrossentropy(start_logits, start_positions)
  end_loss = Flux.logitcrossentropy(end_logits, end_positions)
  total_loss = (start_loss + end_loss) / 2

  return (
    loss = total_loss,
    start_logits = start_logits,
    end_logits = end_logits,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFBertForQuestionAnswering(config::HGFBertConfig)
  bert = HGFBertModel(config)
  classifier = FakeTHLinear(config, config.hidden_size, config.num_labels)
  HGFBertForQuestionAnswering(bert, classifier)
end

# load model utils

basemodelkey(::HGFBertPreTrainedModel) = :bert
basemodel(m::HGFBertPreTrainedModel) = getproperty(m, basemodelkey(m))
basemodel(m::HGFBertModel) = m

isbasemodel(m::HGFBertModel) = true
isbasemodel(m::HGFBertPreTrainedModel) = false

get_model_type(::Val{:bert}) = (
  :model => HGFBertModel,
  :forpretraining => HGFBertForPreTraining,
  :lmheadmodel => HGFBertLMHeadModel,
  :formaskedlm => HGFBertForMaskedLM,
  :fornextsentenceprediction => HGFBertForNextSentencePrediction,
  :forsequenceclassification => HGFBertForSequenceClassification,
  :formultiplechoice => HGFBertForMultipleChoice,
  :fortokenclassification => HGFBertForTokenClassification,
  :forquestionanswering => HGFBertForQuestionAnswering,
)


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

using .GeoStatsBase
import .GeoStatsBase: preprocess, solve_single

"""
    ImgQuilt(var₁=>param₁, var₂=>param₂, ...)

Image quilting simulation solver as described in Hoffimann et al. 2017.

## Parameters

### Required

* `TI`       - Training image
* `tilesize` - Tile size in x, y and z

### Optional

* `overlap`  - Overlap size in x, y and z (default to (1/6, 1/6, 1/6))
* `path`     - Simulation path (:raster (default), :dilation, or :random)
* `inactive` - Vector of inactive voxels (i.e. tuples (i,j,k)) in the grid
* `soft`     - A vector of `(data,dataTI)` pairs
* `tol`      - Initial relaxation tolerance in (0,1] (default to 0.1)

## Global parameters

### Optional

* `threads`      - Number of threads in FFT (default to number of physical CPU cores)
* `gpu`          - Whether to use the GPU or the CPU (default to false)
* `showprogress` - Whether to show or not the estimated time duration (default to false)
"""
@simsolver ImgQuilt begin
  @param TI
  @param tilesize
  @param overlap       = (1/6, 1/6, 1/6)
  @param path          = :raster
  @param inactive      = nothing
  @param soft          = []
  @param tol           = .1
  @global threads      = cpucores()
  @global gpu          = false
  @global showprogress = false
end

function preprocess(problem::SimulationProblem, solver::ImgQuilt)
  # retrieve problem info
  pdata = data(problem)
  pdomain = domain(problem)

  # sanity checks
  @assert pdomain isa RegularGrid "ImgQuilt requires RegularGrid domain"
  @assert ndims(pdomain) ∈ [2,3] "Number of dimensions must be 2 or 3"

  # result of preprocessing
  preproc = Dict{Symbol,Tuple}()

  for (var, V) in variables(problem)
    @assert var ∈ keys(solver.params) "Parameters for variable `$var` not found"

    # get user parameters
    varparams = solver.params[var]

    # add ghost dimension to simulation grid if necessary
    simsize = ndims(pdomain) == 2 ? (size(pdomain)..., 1) : size(pdomain)

    # create hard data object
    hdata = Dict{CartesianIndex{3},Real}()
    for (loc, datloc) in datamap(problem, var)
      push!(hdata, lin2cart(simsize, loc) => pdata[datloc,var])
    end

    # disable inactive voxels
    shape = Dict{CartesianIndex{3},Real}()
    if varparams.inactive ≠ nothing
      for icoords in varparams.inactive
        push!(shape, icoords => NaN)
      end
    end

    preproc[var] = (varparams, simsize, merge(hdata, shape))
  end

  preproc
end

function solve_single(problem::SimulationProblem, var::Symbol,
                      solver::ImgQuilt, preproc)
  # unpack preprocessed parameters
  par, simsize, hard = preproc[var]

  # run image quilting core function
  reals = iqsim(par.TI, par.tilesize, simsize;
                overlap=par.overlap, path=par.path,
                soft=par.soft, hard=hard, tol=par.tol,
                threads=solver.threads, gpu=solver.gpu,
                showprogress=solver.showprogress)

  # flatten result
  reals[1][:]
end


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl
=#

relaxed_tols = (default_tol_relax = 100,)
insts = OrderedDict()
insts["minimal"] = [
    # rootdet
    ((2, 2, MatNegGeom()),),
    ((2, 2, MatNegGeomEFExp()),),
    ((2, 2, MatNegGeomEFPow()),),
    # tr inv
    ((2, 2, MatInv()),),
    ((2, 2, MatInvEigOrd()),),
    ((2, 2, MatInvDirect()),),
    # tr neg2sqrt
    ((1, 2, MatNeg2Sqrt()),),
    ((1, 2, MatNeg2SqrtEigOrd()),),
    # tr negexp1 (domain not positive)
    ((1, 2, MatNegExp1()),),
    ((1, 2, MatNegExp1EigOrd()),),
    # tr power12conj (domain not positive)
    ((1, 2, MatPower12Conj(1.5)),),
    ((1, 2, MatPower12ConjEigOrd(1.5)),),
    ]
insts["fast"] = [
    ((1, 10, MatNegGeomEFExp()),),
    ((1, 15, MatNegGeom()),),
    ((2, 3, MatNegGeomEFPow()),),
    ((2, 3, MatInvEigOrd()),),
    ((2, 5, MatInvDirect()),),
    ((2, 6, MatInv()),),
    ((2, 6, MatNegEntropy()),),
    ((2, 7, MatNegLog()),),
    ((3, 2, MatNegEntropyEigOrd()),),
    ((3, 2, MatNegExp1()),),
    ((3, 2, MatNegExp1EigOrd()),),
    ((3, 2, MatPower12Conj(1.5)),),
    ((3, 2, MatPower12ConjEigOrd(1.5)),),
    ((3, 2, MatPower12EigOrd(1.5)),),
    ((3, 4, MatPower12(1.5)),),
    ((3, 4, MatNegGeom()),),
    ((7, 2, MatNegLog()),),
    ((7, 2, MatInv()),),
    ((7, 2, MatNegEntropy()),),
    ((7, 2, MatPower12(1.5)),),
    ]
insts["various"] = [
    ((2, 5, MatInvEigOrd()),),
    ((2, 5, MatNegLogEigOrd()),),
    ((2, 4, MatPower12EigOrd(1.5)), nothing, relaxed_tols),
    ((2, 4, MatNegEntropyEigOrd()),),
    ((2, 10, MatNegGeomEFExp()),),
    ((2, 8, MatNegGeomEFPow()),),
    ((2, 14, MatNegGeom()),),
    ((2, 12, MatInv()),),
    ((2, 10, MatNegLog()),),
    ((2, 8, MatNegEntropy()),),
    ((2, 6, MatPower12(1.5)),),
    ((3, 6, MatInvDirect()),),
    ((3, 5, MatNegLogDirect()),),
    ((3, 6, MatNegEntropy()),),
    ((3, 6, MatNegLog()),),
    ((3, 5, MatPower12(1.5)),),
    ((3, 5, MatNegExp1()),),
    ((3, 4, MatPower12Conj(1.5)),),
    ((5, 2, MatNegExp1EigOrd()), nothing, relaxed_tols),
    ((5, 2, MatPower12ConjEigOrd(1.5)),),
    ((5, 3, MatInvDirect()),),
    ((5, 3, MatNegLogDirect()),),
    ((6, 3, MatNegGeom()),),
    ((6, 3, MatInv()),),
    ]
return (CentralPolyMatJuMP, insts)


module Streamfall

using Statistics
using Graphs, MetaGraphs, Distributed, DataFrames

const MODPATH = @__DIR__

if Sys.iswindows()
    libext = ".dll"
elseif Sys.islinux()
    libext = ".so"
elseif Sys.isapple()
    libext = ".dynlib"
else
    throw(DomainError("Unsupported platform"))
end

# Can't use string, DLL location has to be a const
# (which makes sense but still, many hours wasted!)
# https://github.com/JuliaLang/julia/issues/29602
const IHACRES = joinpath(MODPATH, "../deps", "usr", "lib", "ihacres$(libext)")

"""@def macro

Inline code to avoid repetitious declarations.
"""
macro def(name, definition)
    return quote
        macro $(esc(name))()
            esc($(Expr(:quote, definition)))
        end
    end
end

@def add_preprefix begin
    if !isnothing(id_prefix)
        prefix = id_prefix * prefix
    end
end

include("Network.jl")
include("Node.jl")
include("Climate.jl")
include("IHACRESNode.jl")
include("IHACRESExpuhNode.jl")
include("GR4JNode.jl")
include("HyModNode.jl")
include("SYMHYDNode.jl")
include("DamNode.jl")
include("EnsembleNode.jl")

function timestep_value(timestep::Int, gauge_id::String, col_partial::String,
                        dataset::Union{DataFrame, Vector, Number, Nothing}=nothing)::Float64
    amount::Float64 = 0.0
    if !isnothing(dataset)
        if dataset isa Vector
            amount = dataset[ts]
        elseif dataset isa DataFrame
            # Extract data for time step based on partial match
            target_col = filter(x -> occursin(gauge_id, x)
                                & occursin(col_partial, x),
                                names(dataset))
            if !isempty(target_col)
                amount = checkbounds(Bool, dataset.Date, timestep) ? dataset[timestep, target_col][1] : 0.0
            end
        elseif dataset isa Number
            amount = dataset
        end
    end

    return amount
end

"""
    find_common_timeframe(timeseries::T...)

Find common time frame between time series.

Requires that all DataFrames have a "Date" column.
"""
function find_common_timeframe(timeseries::T...) where {T<:DataFrame}
    for t in timeseries
        @assert "Date" in names(t)
    end

    min_date = maximum([x.Date[1] for x in timeseries])
    max_date = minimum([x.Date[end] for x in timeseries])
    return (min_date, max_date)
end

"""
    align_time_frame(timeseries::T...)

Subset an arbitrary number of DataFrames to their shared period of time.

Returns subsetted copy of data in same order as input.

# Example
```julia-repl
julia> climate, streamflow = align_time_frame(climate, streamflow)
```
"""
function align_time_frame(timeseries::T...) where {T<:DataFrame}
    min_date, max_date = find_common_timeframe(timeseries...)
    modded = [t[min_date .<= t.Date .<= max_date, :] for t in timeseries]

    return modded
end

"""
    run_node!(sn::StreamfallNetwork, node_id::Int, climate::Climate, timestep::Int;
              extraction::Union{DataFrame, Nothing}=nothing,
              exchange::Union{DataFrame, Nothing}=nothing)

Run a model attached to a node for a given time step.
Recurses upstream as needed.

# Arguments
- `sn::StreamfallNetwork`
- `node_id::Int`
- `climate::Climate`
- `timestep::Int`
- `inflow::DataFrame` : Additional inflow to consider (in ML/timestep)
- `extraction::DataFrame` : Volume of water to be extracted (in ML/timestep)
- `exchange::DataFrame` : Volume of flux (in ML/timestep), where negative values are losses to the groundwater system
"""
function run_node!(sn::StreamfallNetwork, node_id::Int, climate::Climate, timestep::Int;
                   inflow::Union{DataFrame, Nothing}=nothing,
                   extraction::Union{DataFrame, Nothing}=nothing,
                   exchange::Union{DataFrame, Nothing}=nothing)
    ts = timestep

    sim_inflow = 0.0
    ins = inneighbors(sn.mg, node_id)

    if !isempty(ins)
        for i in ins
            # Get inflow from previous node
            res = run_node!(sn, i, climate, timestep;
                            inflow=inflow,
                            extraction=extraction,
                            exchange=exchange)
            if res isa Number
                sim_inflow += res
            elseif length(res) > 1
                # get outflow from (outflow, level)
                sim_inflow += res[1]
            end
        end
    end

    node = sn[node_id]
    ts_inflow = timestep_value(ts, node.name, "inflow", inflow)
    ts_inflow += sim_inflow

    run_func! = get_prop(sn, node_id, :nfunc)

    # Run for a time step, dependent on previous state
    run_func!(node, climate, ts; inflow=ts_inflow, extraction=extraction, exchange=exchange)
end

"""
    run_basin!(sn::StreamfallNetwork, climate::Climate; inflow=nothing, extraction=nothing, exchange=nothing)

Run scenario for an entire catchment/basin.
"""
function run_basin!(sn::StreamfallNetwork, climate::Climate;
                    inflow=nothing, extraction=nothing, exchange=nothing)
    _, outlets = find_inlets_and_outlets(sn)
    for outlet in outlets
        run_node!(sn, outlet, climate;
                  inflow=inflow, extraction=extraction, exchange=exchange)
    end
end

run_catchment! = run_basin!

"""
    run_node!(sn::StreamfallNetwork, node_id::Int, climate::Climate;
              extraction=nothing, exchange=nothing)::Nothing

Generic run method that runs a model attached to a given node for all time steps.
Recurses upstream as needed.

# Arguments
- `sn::StreamfallNetwork`
- `node_id::Int` : node to run in the network
- `climate::Climate` : Climate object holding rainfall and evaporation data (or temperature)
- `extraction::DataFrame` : water orders for each time step (defaults to nothing)
- `exchange::DataFrame` : exchange with groundwater system at each time step (defaults to nothing)
"""
function run_node!(sn::StreamfallNetwork, node_id::Int, climate::Climate;
                   inflow=nothing, extraction=nothing, exchange=nothing)
    timesteps = sim_length(climate)
    for ts in 1:timesteps
        run_node!(sn, node_id, climate, ts;
                  inflow=inflow, extraction=extraction, exchange=exchange)
    end
end

"""
    run_node!(node::NetworkNode, climate::Climate;
              inflow=nothing, extraction=nothing, exchange=nothing)

Run a specific node, and only that node, for all time steps.

# Arguments
- `node::NetworkNode` : Any Streamfall NetworkNode
- `climate` : Climate data
- `inflow::Union{DataFrame, Vector, Number}` : Inflow to node
- `extraction::Union{DataFrame, Vector, Number}` : Extractions from this subcatchment
- `exchange::Union{DataFrame, Vector, Number}` : Groundwater flux
"""
function run_node!(node::NetworkNode, climate::Climate; inflow=nothing, extraction=nothing, exchange=nothing)
    timesteps = sim_length(climate)

    node_name = node.name
    for ts in 1:timesteps
        rain, et = climate_values(node, climate, ts)
        ext = timestep_value(ts, node_name, "extraction", extraction)
        flux = timestep_value(ts, node_name, "exchange", exchange)
        in_flow = timestep_value(ts, node_name, "inflow", inflow)

        run_timestep!(node, rain, et, ts; inflow=in_flow, extraction=ext, exchange=flux)
    end

    return node.outflow, node.level
end

include("metrics.jl")
include("calibration.jl")

include("Analysis/Analysis.jl")
include("plotting.jl")

export @def

# Nodes
export NetworkNode, GenericNode, GenericDirectNode
export IHACRES, IHACRESNode, BilinearNode, ExpuhNode, DamNode, Climate
export create_node, GR4JNode, HyModNode, SimpleHyModNode, SYMHYDNode
export EnsembleNode, BaseEnsemble
export run_step!, run_timestep!

# Network
export find_inlets_and_outlets, inlets, outlets, create_network, create_node
export climate_values, get_node, get_node_id, get_prop, set_prop!
export param_info, update_params!, sim_length, reset!
export run_catchment!, run_basin!, run_node!, run_node_with_temp!
export calibrate!

# plotting methods
export quickplot, plot_network, save_figure

# data interface (climate)
export timesteps

# Analysis
export Analysis

end  # end module