text:
using Plots
# using IndexedTables
gr()

mutable struct XY
    x::Real
    y::Real
    XY(x::T where {T <: Real}, y::T where {T <: Real}) = new(x,y)
end

mutable struct Particles
    x::Real
    y::Real
    index::Int
    best::XY
    velocity::XY

    function Particles(w::T1, h::T2, index::T3) where {T1 <: Real, T2 <: Real, T3 <: Int}
        x = getRandomInt(0.1*w, 0.9*w)
        y = h/2.0
        best = XY(x,y)
        velocity = XY(getRandomInt(-5, 5), getRandomInt(0, 5))
        new(x, y, index, best, velocity)
    end
end # struct Particles

struct Bags_size
    w::Real
    h::Real
    x::Real
    y::Real
    sides::NamedTuple

    function Bags_size(w::T1, h::T2, x::T3, y::T4) where {T1 <: Real, T2 <: Real, T3 <: Real, T4 <: Real}
        temp1 = Set(x .+ [0,w,w,0])
        temp2 = Set(y .+ [0,0,h,h])
        sides = (left = minimum(temp1), rigth = maximum(temp1), bottom = minimum(temp2), top = maximum(temp2))
        new(w,h,x,y,sides)
    end
end

mutable struct Weights
    inertiaWeight::Real
    personalWeight::Real
    swarmWeight::Real
    Weights(i, p, s) = new(i, p, s)
end

"""
    getRandomInt(min,max)

generate a random number between min and max
"""
function getRandomInt(min,max)
    return floor(rand() * (max - min + 1)) + min
end # function getRandomInt

rectangle(w, h, x, y) = Shape(x .+ [0,w,w,0], y .+ [0,0,h,h])

"""
    draw(particles::Particles, canvas::Plots.Plot)

this function draw the bag and particles
"""
function draw(particles::Particles, canvas::Plots.Plot)
    scatter!(canvas,[particles.x],[particles.y],marker = :rect)
end # function draw

"""
    init(width::Real, height::Real)

this function create a plot canvas and place the particles and the
center of the canvas
"""
function init(width::T1, height::T2, number::Int) where {T1 <: Real, T2 <: Real}
    bag_size= Bags_size(width/3,height/3,0,0)
    particles = []
    for ind in 1:number
        push!(particles,Particles(bag_size.w, bag_size.h, ind))
    end
    canvas = plot(size = (width,height), legend = false)
    xlims!(canvas,(-width/2,width/2))
    ylims!(canvas,(-height/2,height/2))
    bag = rectangle(bag_size.w,bag_size.h,bag_size.x,bag_size.y)
    plot!(canvas, bag, opacity = 0.2)
    for ind in 1:number
        draw(particles[ind], canvas)
    end
    return (particles, bag_size, canvas)
end # function init

# """
#     update_all!(particles::Array{Particles}, canvas::Plots.Plot, width::Real, height::Real, bag_size::Bags_size)
#
# this function update all the particles location using the 'move!' funciton
# """
# function update_all!(particles::Array{Particles}, canvas::Plots.Plot, width::T1, height::T2, bag_size::Bags_size) where {T1 <: Real, T2 <: Real}
#     for ind in 1:length(particles)
#         move!(particles,ind)
#         canvas.series_list[ind + 1][:x] = particles[ind].x
#         canvas.series_list[ind + 1][:y] = particles[ind].y
#     end
#     # if !in_bag(particles, bag_size.sides)
#     #     init()
#     # end
# end # function update_all!
#
#
# """
#     update_single!(particles::Array{Particles}, index::Int, canvas::Plots.Plot, width::Real, height::Real, bag_size::Bags_size)
#
# this function update a single particle location using the 'move!' funciton
# """
# function update_single!(particles::Array{Particles}, index::T1, canvas::Plots.Plot, width::T2, height::T3, bag_size::Bags_size) where {T1 <: Int, T2 <: Real, T3 <: Real}
#     move!(particles,index)
#     canvas.series_list[index + 1][:x] = particles[index].x
#     canvas.series_list[index + 1][:y] = particles[index].y
#     if !in_bag(particles, index, bag_size.sides)
#         particles[index].step = 0
#     end
# end # function update_single!

"""
    move_in_range(v::T1, max::T2, particle::Particles, property::Symbol) where {T1 <: Real, T2 <: Real}

this funcion keep the particle with in the bags sides
"""
function move_in_range(v::T1, max::T2, particle::Particles, property::Symbol) where {T1 <: Real, T2 <: Real}
    val = getfield(particle, property) + v
    if val < 0
        setfield!(particle,property,0)
    elseif val > max
        setfield!(particle,property,max)
    else
        setfield!(particle,property,val)
        setfield!(particle.velocity,property,v)
    end
end # function move_in_range

"""
    move!(particles::Array{Particles}, weights::Weights, height::T1, width::T2, bestGlobal::XY) where {T1 <: Real, T2 <: Real, T3 <: Real}

this function add a random shift in the particles position
"""
function move!(particles::Array{Particles}, weights::Weights, height::T1, width::T2, bestGlobal::XY) where {T1 <: Real, T2 <: Real, T3 <: Real}
    for particle in particles
        r1 = getRandomInt(0,5)
        r2 = getRandomInt(0,5)
        vy =    weights.inertiaWeight * particle.velocity.y +
                weights.personalWeight * r1 * (particle.best.y - particle.y) +
                weights.swarmWeight * r2 * (bestGlobal.y - particle.y)
        vx =    weights.inertiaWeight * particle.velocity.x +
                weights.personalWeight * r1 * (particle.best.x - particle.x) +
                weights.swarmWeight * r2 * (bestGlobal.x - particle.x)
        move_in_range(vy,height,particle,:y)
        move_in_range(vx,width,particle,:x)
    end
    return particles
end # function move!

# """
#     nudge(particles::Array{Particles})
#
# documentation
# """
# function nudge(particles::Array{Particles})
#     N = length(particles)
#     iszero(N) && return 0
#     sum = [0.0,0.0]
#     for particle in particles
#         sum[1] += particle.x
#         sum[2] += particle.y
#     end
#     return (sum[1]/N,sum[2]/N)
# end # function nudge

# """
#     in_bag(particles::Array{Particles}, bag_sides::NamedTuple)::Bool
#
# this function has two methods
#     1. check if all the particles are in the bag
#     2. check if a specific particle is in the bag
# """
# function in_bag(particles::Array{Particles}, bag_sides::NamedTuple)::Bool
#     res = fill(false,(length(particles),1))
#     for (ind,particle) in enumerate(particles)
#         if (bag_sides.left < particle.x < bag_sides.rigth) && (bag_sides.bottom < particle.y < bag_sides.top)
#             res[ind] = true
#         end
#     end
#     return all(res)
# end # function in_bag

# function in_bag(particles::Array{Particles}, index::T, bag_sides::NamedTuple)::Bool  where {T <: Int}
#     res = false
#     if (bag_sides.left < particles[index].x < bag_sides.rigth) && (bag_sides.bottom < particles[index].y < bag_sides.top)
#         res = true
#     end
#     return res
# end # function in_bag

# """
#     addparticle(particles::Array{Particles},canvas::Plots.Plot)
#
# this function add a new particles at the center of the canvas
# """
# function addparticle(particles::Array{Particles}, canvas::Plots.Plot)
#     index = length(particles) + 1
#     push!(particles,Particles(index))
#     draw(particles[end], canvas)
# end # function addparticle

# """
#     euclidean_distance(p1,p2)
#
# this function calculate the euclidean distance between two particles
# """
# function euclidean_distance(p1::Particles, p2::Particles)
#     return sqrt((p1.x - p2.x)^2 + (p1.y - p2.y)^2)
# end # function euclidean_distance

# """
#     knn(particles::Array{Particles}, index::Int)
#
# this function finds the 'k' nearest neighbors
# """
# function knn(particles::Array{Particles}, index::T, k::T) where {T <: Int}
#     tested_particle = particles[index]
#     dist = []
#     ind = []
#     for particle in particles
#         if particle.index != index
#             push!(dist,euclidean_distance(tested_particle,particle))
#             push!(ind,particle.index)
#         end
#     end
#     m = [dist ind]
#     return iszero(length(m)) ? [] : m[sortperm(m[:,1]),:][1:k,2]
# end # function knn

"""
    updateBest(particles::Array{Particles}, bestGlobal::XY)

this function find the best position based on all the particles
"""
function updateBest(particles::Array{Particles}, bestGlobal::XY)
    y = particles[1].y
    x = particles[1].x
    for particle in particles
        if particle.y > y
            y = particle.y
            x = particle.x
        end
    end
    bestGlobal = XY(x,y)
end # function updateBest

"""
    pso(particles::Array{Particles}, epoch::Int, bestGlobal::XY, height::T2, width::T3) where {T1 <: Real, T2 <: Real, T3 <: Real}

documentation
"""
function pso(particles::Array{Particles}, epoch::Int, bestGlobal::XY, height::T2, width::T3) where {T1 <: Real, T2 <: Real, T3 <: Real}
    epoch += 1
    weights = Weights(0.9,0.5,0.5)
    particle_size = 4
    move!(particles, weights, height - particle_size, width - particle_size, bestGlobal)
    # draw
    bestGlobal = updateBest(particles, bestGlobal)
    cb(timer) = pso(particles, epoch, bestGlobal, height, width)
    t = Timer(cb,0, interval = 1)
    if epoch >= 40
        close(t)
    end
end # function pso

# """
#     main_follow_the_best(width::Real = 600, height::Real = 600)
#
# this function run the main loop and move a single particle at a time
# """
# function main_follow_the_best(width::T1 = 600, height::T2 = 600; number::Int = 20) where {T1 <: Real, T2 <: Real}
#     println("To add a new particles press 'n'")
#     println("To add a quit press 'q'")
#     (particles, bag_size, canvas) = init(width, height, number)
#     current_p = 0
#     @async begin
#         cb(timer) = (current_p = (current_p  % length(particles)) + 1;
#                      if particles[current_p].step > 0
#                          update_single!(particles, current_p, canvas, width, height, bag_size)
#                      end;
#                      display(canvas))
#         global t = Timer(cb,0, interval = 1)
#     end
#     while true
#         kb_input = readline()
#         if lowercase(kb_input) == "n"
#             addparticle(particles,canvas)
#             kb_input = ""
#         end
#         if lowercase(kb_input) == "q"
#             close(t)
#             break
#         end
#     end
# end # function main_follow_the_best
#
#
# """
#     create_animate()
#
# documentation
# """
# function create_animate(args)
#     anim = @animate for ii in 1:15
#         Timer(update_all!(particles, canvas, width, height, bag_size),0, interval = 1)
#         display(canvas)
#     end
#     gif(anim, joinpath(@__DIR__,"..","animations\\anim_fps15.gif"), fps = 15)
# end # function


text:
using DataStructures
using Compat
using Compat: @warn
using IterableTables: getiterator
using DataFrames

abstract type AbstractBranchTree{RT, N, B, LI, ND} <:
    AbstractTree{OneTree, RT, String, N, B}
end

import Phylo.API: _leafinfotype
_leafinfotype(::Type{<: AbstractBranchTree{RT, N, B, LI, ND}}) where
    {RT, N, B, LI, ND} = LI

import Phylo.API: _noderecordtype
_noderecordtype(::Type{<: AbstractBranchTree{RT, N, B, LI, ND}}) where
    {RT, N, B, LI, ND} = ND

import Phylo.API: _branchrecordtype
_branchrecordtype(::Type{<: AbstractBranchTree}) = Nothing

import Phylo.API: _getnodes
_getnodes(bt::AbstractBranchTree) = values(bt.nodes)

import Phylo.API: _getnodenames
_getnodenames(bt::AbstractBranchTree) = keys(bt.nodes)

import Phylo.API: _getbranches
_getbranches(bt::AbstractBranchTree) = values(bt.branches)

import Phylo.API: _getbranchnames
_getbranchnames(bt::AbstractBranchTree) = keys(bt.branches)

import Phylo.API: _getleafinfo
_getleafinfo(bt::AbstractBranchTree) = bt.leafinfos
function _getleafinfo(bt::AbstractBranchTree, leafname)
    return [info for info in getiterator(bt.leafinfos) if info[1] == leafname]
end

import Phylo.API: _setleafinfo!
function _setleafinfo!(bt::AbstractBranchTree, info)
    bt.leafinfos = info
end

import Phylo.API: _resetleaves!
function _resetleaves!(bt::AbstractBranchTree)
    bt.leafinfos = empty!(bt.leafinfos)
    return bt
end

import Phylo.API: _getnoderecord
function _getnoderecord(bt::AbstractBranchTree, name)
    return bt.noderecords[name]
end

import Phylo.API: _setnoderecord!
function _setnoderecord!(bt::AbstractBranchTree, name, value)
    bt.noderecords[name] = value
end

import Phylo.API: _hasnode
_hasnode(tree::AbstractBranchTree, name::String) = haskey(tree.nodes, name)

import Phylo.API: _getnode
_getnode(tree::AbstractBranchTree, name::String) = tree.nodes[name]

import Phylo.API: _createnode!
function _createnode!(tree::AbstractBranchTree{RT, N, B, LI, ND},
                      name::String, data::ND = ND()) where {RT, N, B, LI, ND}
    _hasnode(tree, name) && error("Node $name already present in tree")
    tree.nodes[name] = N(name)
    _setnoderecord!(tree, name, data)
    return name
end

import Phylo.API: _deletenode!
function _deletenode!(tree::AbstractBranchTree, node::AbstractNode)
    name = getnodename(tree, node)
    if _hasinbound(tree, node)
        deletebranch!(tree, _getinbound(tree, node))
    end
    for b in _getoutbounds(tree, node)
        deletebranch!(tree, b)
    end
    delete!(tree.nodes, name)
    delete!(tree.noderecords, name)
    return name
end
function _deletenode!(tree::AbstractBranchTree, name::String)
    node = getnode(tree, name)
    if _hasinbound(tree, node)
        deletebranch!(tree, _getinbound(tree, node))
    end
    for b in _getoutbounds(tree, node)
        deletebranch!(tree, b)
    end
    delete!(tree.nodes, name)
    delete!(tree.noderecords, name)
    return name
end

import Phylo.API: _hasbranch
_hasbranch(tree::AbstractBranchTree, name::Int) = haskey(tree.branches, name)

import Phylo.API: _getbranch
_getbranch(tree::AbstractBranchTree, name::Int) = tree.branches[name]

import Phylo.API: _createbranch!
function _createbranch!(tree::AbstractBranchTree{RT}, source, destination,
                        length::Float64, name::Int,
                        data::Nothing = nothing) where RT
    # Add the new branch
    branch = Branch{RT}(name, _getnodename(tree, source),
                        _getnodename(tree, destination), length)
    tree.branches[name] = branch

    # Update the associated source and destination nodes
    _addoutbound!(tree, _getnode(tree, source), branch)
    _addinbound!(tree, _getnode(tree, destination), branch)

    # Return updated tree
    return name
end

import Phylo.API: _deletebranch!
function _deletebranch!(tree::AbstractBranchTree, name::Int)
    # Find the branch
    branch = _getbranch(tree, name)
    # Remove branch reference from its source node
    _removeoutbound!(tree, _getnode(tree, _src(tree, branch)), branch)
    # Remove branch reference from its destination node
    _removeinbound!(tree, _getnode(tree, _dst(tree, branch)), branch)
    # Remove branch itself
    delete!(tree.branches, name)
    # Return the branch name
    return name
end
function _deletebranch!(tree::AbstractBranchTree, branch::Branch)
    # Find the branch
    name = _getbranchname(tree, branch)
    # Remove branch reference from its source node
    _removeoutbound!(tree, _getnode(tree, _src(tree, branch)), branch)
    # Remove branch reference from its destination node
    _removeinbound!(tree, _getnode(tree, _dst(tree, branch)), branch)
    # Remove branch itself
    delete!(tree.branches, name)
    # Return the branch name
    return branch
end

import Phylo.API: _validate
function _validate(tree::TREE) where {TREE <: AbstractBranchTree}
    if _leafinfotype(TREE) != Nothing && length(getiterator(tree.leafinfos)) > 0
        if Set(info[1] for info in getiterator(tree.leafinfos)) !=
            Set(_getleafnames(tree))
            @warn "LeafInfo names do not match actual leaves of tree"
            return false
        end
    end

    if Set(keys(tree.noderecords)) != Set(_getnodenames(tree))
        @warn "Node names do not match node records of tree"
        return false
    end
    return true
end

import Phylo.API: _hasrootheight
function _hasrootheight(tree::AbstractBranchTree)
    return !isnan(tree.rootheight)
end

import Phylo.API: _getrootheight
function _getrootheight(tree::AbstractBranchTree)
    return tree.rootheight
end

import Phylo.API: _setrootheight!
function _setrootheight!(tree::AbstractBranchTree, height::Float64)
    tree.rootheight = height
    return height
end

import Phylo.API: _clearrootheight!
function _clearrootheight!(tree::AbstractBranchTree)
    tree.rootheight = NaN
end

"""
    BinaryTree

Binary phylogenetic tree object with known leaves and per node data
"""
mutable struct BinaryTree{RT <: Rooted, LI, ND} <:
    AbstractBranchTree{RT, BinaryNode{RT, String, Branch{RT, String}},
                       Branch{RT, String}, LI, ND}
    nodes::OrderedDict{String, BinaryNode{RT, String, Branch{RT, String}}}
    branches::Dict{Int, Branch{RT, String}}
    leafinfos::LI
    noderecords::OrderedDict{String, ND}
    rootheight::Float64
end

function BinaryTree(lt::BinaryTree{RT, LI, ND};
                    copyinfo=false) where {RT, LI, ND}
    validate(lt) || error("Tree to copy is not valid")
    leafnames = getleafnames(lt)
    # Leaf records are conserved across trees, as could be invariant?
    leafinfos = copyinfo ? deepcopy(lt.leafinfos) : lt.leafinfos
    nodes = OrderedDict(leaf => BinaryNode{RT, String, Branch{RT, String}}(leaf)
                        for leaf in leafnames)
    branches = Dict{Int, Branch{RT, String}}()
    noderecords = OrderedDict(leaf => ND() for leaf in leafnames)
    return BinaryTree{RT, LI, ND}(nodes, branches, leafinfos, noderecords,
                                  lt.rootheight)
end

function BinaryTree{RT, LI, ND}(leafnames::Vector{String},
                                treetype::Type{BinaryTree{RT, LI, ND}} =
                                BinaryTree{RT, LI, ND};
                                rootheight::Float64 = NaN) where {RT, LI, ND}
    nodes = OrderedDict(leaf => BinaryNode{RT, String, Branch{RT, String}}(leaf)
                        for leaf in leafnames)
    leafinfos = LI()
    noderecords = OrderedDict(leaf => ND() for leaf in leafnames)
    return BinaryTree{RT, LI, ND}(nodes, Dict{Int, Branch{RT, String}}(),
                                  leafinfos, noderecords, rootheight)
end

function BinaryTree{RT, LI, ND}(numleaves::Int = 0,
                                treetype::Type{BinaryTree{RT, LI, ND}} =
                                BinaryTree{LI, ND};
                                rootheight::Float64 = NaN) where {RT, LI, ND}
    leafnames = ["Leaf $num" for num in Base.OneTo(numleaves)]
    return BinaryTree{RT, LI, ND}(leafnames, treetype; rootheight = rootheight)
end

function BinaryTree{LI, ND}(leafinfos::LI;
                            rootheight::Float64 = NaN) where {LI, ND}
    leafnames = unique(info[1] for info in getiterator(leafinfos))
    nodes = OrderedDict(leaf => BinaryNode{RT, String, Branch{RT, String}}(leaf)
                        for leaf in leafnames)
    branches = Dict{Int, Branch{RT, String}}()
    noderecords = OrderedDict(leaf => ND() for leaf in leafnames)
    return BinaryTree{RT, LI, ND}(nodes, branches, leafinfos,
                                  noderecords, rootheight)
end

BinaryTree{RT}(leafinfos::LI; rootheight::Float64 = NaN) where {RT, LI} =
    BinaryTree{RT, LI, Dict{String, Any}}(leafinfos; rootheight = rootheight)

"""
    NamedBinaryTree

Binary phylogenetic tree object with known leaves
"""
const NamedBinaryTree = BinaryTree{ManyRoots, DataFrame, Dict{String, Any}}

"""
    PolytomousTree

Phylogenetic tree object with polytomous branching, and known leaves
and per node data
"""
mutable struct PolytomousTree{RT <: Rooted, LI, ND} <:
    AbstractBranchTree{RT, Node{RT, String, Branch{RT, String}},
                       Branch{RT, String}, LI, ND}
    nodes::OrderedDict{String, Node{RT, String, Branch{RT, String}}}
    branches::Dict{Int, Branch{RT, String}}
    leafinfos::LI
    noderecords::OrderedDict{String, ND}
    rootheight::Float64
end

function PolytomousTree(lt::PolytomousTree{RT, LI, ND};
                        copyinfo=false) where {RT, LI, ND}
    validate(lt) || error("Tree to copy is not valid")
    leafnames = getleafnames(lt)
    # Leaf records may be conserved across trees, as could be invariant?
    leafinfos = copyinfo ? deepcopy(lt.leafinfos) : lt.leafinfos
    nodes = OrderedDict(leaf => Node{RT, String, Branch{RT, String}}(leaf)
                        for leaf in leafnames)
    branches = Dict{Int, Branch{String}}()
    noderecords = OrderedDict(leaf => ND() for leaf in leafnames)
    return PolytomousTree{RT, LI, ND}(nodes, branches, leafinfos,
                                      noderecords, lt.rootheight)
end

function PolytomousTree{RT, LI, ND}(leafnames::Vector{String},
                                    treetype::Type{PolytomousTree{RT, LI, ND}} =
                                        PolytomousTree{RT, LI, ND};
                                    rootheight::Float64 = NaN) where
    {RT, LI, ND}
    nodes = OrderedDict(leaf => Node{RT, String, Branch{RT, String}}(leaf)
                        for leaf in leafnames)
    leafinfos = LI()
    noderecords = OrderedDict(leaf => ND() for leaf in leafnames)
    return PolytomousTree{RT, LI, ND}(nodes,
                                      OrderedDict{Int,
                                                  Branch{RT, String}}(),
                                      leafinfos, noderecords, rootheight)
end

function PolytomousTree{RT, LI, ND}(numleaves::Int = 0,
                                    treetype::Type{PolytomousTree{RT, LI, ND}} =
                                        PolytomousTree{RT, LI, ND};
                                    rootheight::Float64 = NaN) where
    {RT, LI, ND}
    leafnames = ["Leaf $num"  for num in Base.OneTo(numleaves)]
    return PolytomousTree{RT, LI, ND}(leafnames, treetype;
                                      rootheight = rootheight)
end

function PolytomousTree{RT, LI, ND}(leafinfos::LI,
                                    treetype::Type{PolytomousTree{RT, LI, ND}} =
                                        PolytomousTree{RT, LI, ND};
                                    rootheight::Float64 = NaN) where
    {RT, LI, ND}
    leafnames = unique(info[1] for info in getiterator(leafinfos))
    nodes = OrderedDict(leaf => Node{RT, String, Branch{RT, String}}(leaf)
                        for leaf in leafnames)
    branches = Dict{Int, Branch{String}}()
    noderecords = OrderedDict(leaf => ND() for leaf in leafnames)
    return PolytomousTree{RT, LI, ND}(nodes, branches, leafinfos,
                                      noderecords, rootheight)
end

PolytomousTree{RT}(leafinfos::LI; rootheight::Float64 = NaN) where {RT, LI} =
    PolytomousTree{RT, LI, Dict{String, Any}}(leafinfos;
                                              rootheight = rootheight)

"""
    NamedPolytomousTree

Polytomous phylogenetic tree object with known leaves
"""
const NamedTree = NamedPolytomousTree =
    PolytomousTree{ManyRoots, DataFrame, Dict{String, Any}}

ext: jl


text:
module IHT

using Distances: euclidean, chebyshev, sqeuclidean
using PLINK
using RegressionTools
#using StatsBase
using StatsFuns: logistic, logit, softplus
using DataFrames
using Gadfly

### idea from Julia Hoffimann Mendes to conditionally load OpenCL module
# only load if Julia can find OpenCL module
# otherwise warn and set "cl" variable to Void
# will load GPU code based on value of "cl"
try
    using OpenCL
catch e
    warn("IHT.jl cannot find an OpenCL library and will not load GPU functions correctly.")
    global cl = nothing
end

# since PLINK module also loads OpenCL, the previous test can fail
# check that PLINK actually loaded GPU code
try
    PLINK.PlinkGPUVariables
catch e
    warn("PLINK.jl could not find an OpenCL library, so IHT.jl will not load GPU functions correctly.")
    global cl = nothing
end

# used for pretty printing of IHTResults, IHTCrossvalidationResults
import Base.show

# used to plot MSEs v. models from IHTCrossvalidationResults
import Gadfly.plot

export L0_reg
export L0_log
export iht_path
export iht_path_log
export cv_iht
export cv_get_folds
export cv_log
export MendelIHT

# IHT will only work on single/double precision floats!
const Float = Union{Float64,Float32}

include("common.jl")
#if cl != nothing
try
    include("gpu.jl") # conditional load of GPU code
catch e
    warn("IHT.jl failed to load GPU functions!")
end
include("gwas.jl")
include("cv.jl")
include("hardthreshold.jl")
include("aiht.jl")
include("log.jl")
include("MendelIHT.jl")

end # end module IHT


text:
"""
   immutable struct contain the model parameters
*  data_format : Float32 or Float64
*  nz          : depth of model
*  nx          : length of model
*  npml        : number of PML layers
*  free_surface: true -> no obsorbtion on the surface
                 false-> PML on the surface
*  Nz          : depth of model (include PML padding)
*  Nx          : length of model (include PML padding)
*  ntop        : number of padded PML layers on the top
*  spt2wfd     : auxillary index vector for computing wavefield from snapshot
*  spt2bnd     : auxillary index vector for computing wavefield boundaries from snapshot
*  bnd2wfd     : auxillary index vector for inserting wavefield boundaries to wavefield
*  dz          : vertical grid size
*  dx          : horizontal grid size
*  dt          : time step size
*  tmax        : maximum modelling length
*  nt          : number of time steps
*  rho         : 2D density model
*  vel         : 2D P-wave velocity model
*  MvzBvz      : FD stencil for updating v_z from v_z, vector length of Nz
*  MvzBp       : FD stencil for updating v_z from p  , vector length of Nz * Nx
*  dpdz        : sparse matrix computing dp/dz

*  MvxBvx      : FD stencil for updating v_x from v_x, vector length of Nx
*  MvxBp       : FD stencil for updating v_x from p  , vector length of Nz * Nx
*  dpdx        : sparse matrix computing dp/dx

*  MpzBpz      : FD stencil for updating p_z from p_z, vector length of Nz
*  MpzBvz      : FD stencil for updating p_z from v_z, vector length of Nz * Nx
*  dvdz        : sparse matrix computing dvz/dz

*  MpxBpx      : FD stencil for updating p_x from p_x, vector length of Nx
*  MpxBvx      : FD stencil for updating p_x from v_x, vector length of Nz * Nx
*  dvdx        : sparse matrix computing dvx/dx

*  RvzBp       : FD stencil for updating v_z from p, vector length of nz * nx
*  rpdz        : sparse matrix computing dp/dz

*  RvxBp       : FD stencil for updating v_x from p, vector length of nz * nx
*  rpdx        : sparse matrix computing dp/dx

*  RpzBvz      : FD stencil for updating p_z from v_z, vector length of nz * nx
*  rvdz        : sparse matrix computing dv/dz

*  RpxBvx      : FD stencil for updating p_x from v_x, vector length of nz * nx
*  rvdx        : sparse matrix computing dv/dx
"""
struct TdParams{Ti<:Int64, Tv<:AbstractFloat}

    data_format  :: DataType
    nz           :: Ti
    nx           :: Ti
    npml         :: Ti
    free_surface :: Bool
    Nz           :: Ti
    Nx           :: Ti
    ntop         :: Ti
    spt2wfd      :: Vector{Ti}
    spt2bnd      :: Vector{Ti}
    wfd2bnd      :: Vector{Ti}
    order        :: Ti
    dz           :: Tv
    dx           :: Tv
    dt           :: Tv
    tmax         :: Tv
    nt           :: Ti
    rho          :: Matrix{Tv}
    vel          :: Matrix{Tv}
    MvzBvz       :: Vector{Tv}
    MvzBp        :: Vector{Tv}
    RvzBp        :: Vector{Tv}
    dpdz         :: SparseMatrixCSC{Tv,Ti}
    rpdz         :: SparseMatrixCSC{Tv,Ti}
    MvxBvx       :: Vector{Tv}
    MvxBp        :: Vector{Tv}
    RvxBp        :: Vector{Tv}
    dpdx         :: SparseMatrixCSC{Tv,Ti}
    rpdx         :: SparseMatrixCSC{Tv,Ti}
    MpzBpz       :: Vector{Tv}
    MpzBvz       :: Vector{Tv}
    dvdz         :: SparseMatrixCSC{Tv,Ti}
    rvdz         :: SparseMatrixCSC{Tv,Ti}
    MpxBpx       :: Vector{Tv}
    MpxBvx       :: Vector{Tv}
    RpBv         :: Vector{Tv}
    dvdx         :: SparseMatrixCSC{Tv,Ti}
    rvdx         :: SparseMatrixCSC{Tv,Ti}
end

"""
   Overloading the show function for TdParams
"""
function show(io::IO, params::TdParams)

    # size of model
    @printf("nz = %4d, nx = %4d, npml = %4d\n", params.nz, params.nx, params.npml)
    @printf("dz = %4.1f, dx = %4.1f\n", params.dz, params.dx)
    @printf("time step size = %4.4f, max simulation time = %4.1f, number of step = %5d\n", params.dt, params.tmax, params.nt)

    # top boundary condition
    if params.free_surface
       @printf("top boundary conditon = %20s\n", "free surface")
    else
       @printf("top boundary conditon = %20s\n", "PML obsorbing")
    end

    # data presicion (Float32, Float64)
    @printf("data_format=%16s\n", params.data_format)

    # compute the memory cost of time domain parameters
    esize = sizeof(params)
    esize+= 172 + 8*3 + 1 + 8*3
    esize+= sizeof(params.spt2wfd) + sizeof(params.spt2bnd) + sizeof(params.wfd2bnd)
    esize+= 8*2 + sizeof(params.data_format) * 4
    esize+= sizeof(params.rho) + sizeof(params.vel)
    esize+= sizeof(params.MvzBvz)*2 + sizeof(params.MvxBvx)*2 + sizeof(params.MvzBp)*4 + sizeof(params.RvzBp)*3
    esize+= (sizeof(params.dpdz) + sizeof(params.dpdz.m)*2 + sizeof(params.dpdz.rowval) + sizeof(params.dpdz.colptr) + sizeof(params.dpdz.nzval)) * 2
    esize+= (sizeof(params.dpdx) + sizeof(params.dpdx.m)*2 + sizeof(params.dpdx.rowval) + sizeof(params.dpdx.colptr) + sizeof(params.dpdx.nzval)) * 2
    esize+= (sizeof(params.rpdz) + sizeof(params.rpdz.m)*2 + sizeof(params.rpdz.rowval) + sizeof(params.rpdz.colptr) + sizeof(params.rpdz.nzval)) * 2
    esize+= (sizeof(params.rpdz) + sizeof(params.rpdx.m)*2 + sizeof(params.rpdx.rowval) + sizeof(params.rpdx.colptr) + sizeof(params.rpdx.nzval)) * 2

    # convert to mega byte
    esize = esize / 1024^2
    @printf("esize of TdParams=%6.3f M\n", esize)

    return nothing
end

"""
   Padding PML obsorbing layers to model parameters
"""
function  model_padding(m::Matrix{Tv}, npml::Ti, free_surface::Bool) where {Ti<:Int64, Tv<:AbstractFloat}

    # model size
    (nz, nx) = size(m)

    # model size after padding
    Nx = nx + 2 * npml
    if free_surface
       Nz   = nz + npml
       ntop = 0
    else
       Nz   = nz + npml * 2
       ntop = npml
    end

    # allocate memory for the padded model
    mpad = zeros(eltype(m), Nz, Nx)

    # central part
    for ix = 1 : nx
        for iz = 1 : nz
            mpad[iz+ntop, ix+npml] = m[iz, ix]
        end
    end

    # left and right sides
    for ix = 1 : npml
        for iz = 1 : nz
            mpad[iz+ntop,ix        ] = m[iz,1 ]
            mpad[iz+ntop,ix+npml+nx] = m[iz,nx]
        end
    end

    # upper and bottom sides
    for ix = 1 : nx
        for iz = 1 : npml

            if free_surface == false              #  top
               mpad[iz,ix+npml] = m[1,ix]
            end
            mpad[iz+nz+ntop,ix+npml] = m[nz,ix]   # bottom
        end
    end

    # corners
    for ix = 1 : npml
        for iz = 1 : npml
            if free_surface == false
               mpad[iz, ix        ] = m[1, 1 ]    # upper-left
               mpad[iz, ix+npml+nx] = m[1, nx]    # upper-right
            end
            mpad[iz+ntop+nz, ix        ] = m[nz, 1 ]   # bottom-left
            mpad[iz+ntop+nz, ix+npml+nx] = m[nz, nx]   # bottom-right
        end
    end

    return mpad
end

"""
   Smoothing the 2D physical model with Hanning or Gaussian window whose half length is L
"""
function model_smooth(par::Matrix{Tv}, L::Ti; filter_flag="hanning") where {Ti<:Int64, Tv<:AbstractFloat}

    data_format = eltype(par)

    # size of smoother
    n = 2*L + 1

    # padding model before smoothing to handle boundary part
    par1 = model_padding(par, L, false)

    # convert the type of smoother
    if filter_flag == "hanning"
       s = convert(Vector{data_format}, hanning(n))
    elseif filter_flag == "gaussian"
       s = convert(Vector{data_format}, gaussian(n, 0.5))
    else
       error("wronge filter, only supper hanning and gaussian smoother")
    end
    s = s / sum(s)

    # smoothing via 2D convolution
    par1 = conv(s,s,par1)

    # truncate to original size
    par1 = par1[2*L+1:end-2*L, 2*L+1:end-2*L]

    return par1
end

"""
compute buoyancy from density term
"""
function density2buoyancy(rho::Matrix{Tv}) where {Tv<:AbstractFloat}

    data_format = eltype(rho)

    # dimensions of the model
    (nz, nx) = size(rho)

    # allocate space for buoyancy
    buoyancy = zeros(data_format, nz, nx)

    # constant
    c_one = one(data_format)

    # compute element-wise buoyancy
    for i2 = 1 : nx
        for i1 = 1 : nz
            buoyancy[i1,i2] = c_one / rho[i1,i2]
        end
    end

    return buoyancy
end

"""
   compute lambda from velocity and density
"""
function vp2lambda(rho::Matrix{Tv}, vp::Matrix{Tv}) where {Tv<:AbstractFloat}

    (m , n ) = size(rho)
    (m1, n1) = size(vp )
    if m1 != m || n1 != n
       error("size mismatch")
    end

    lambda = zeros(eltype(vp), m, n)
    for j = 1 : n
        for i = 1 : m
            lambda[i,j] = rho[i,j] * vp[i,j] * vp[i,j]
        end
    end
    return lambda
end

"""
   get the auxillary index vectors which facilitate extracting wavefield from snapshot,
saving wavefield boundaries and insert the boundaries for wavefield reconstruction.
"""
function get_index(nz::Ti, nx::Ti, npml::Ti, free_surface::Bool; order=2) where {Ti<:Int64}

    # size of wavefield
    len1 = nz * nx

    # length of boundary values need to save
    len2 = 2*order*nz + 2*order*(nx-2*order)

    # allocate 3 vectors of integer
    index1 = zeros(Int64, len1)
    index2 = zeros(Int64, len2)
    index3 = zeros(Int64, len2)

    # size of padded model
    if free_surface
       Nz = nz + npml
       ntop = 0
    else
       Nz = nz + npml * 2
       ntop = npml
    end
    Nx = nx + npml * 2

    # map snapshot to wavefield
    icount = 0
    for i2 = 1 : nx
        tmp= (i2 + npml - 1)*Nz + ntop

        for i1 = 1 : nz
            icount = icount + 1
            index1[icount] = tmp + i1
        end
    end

    # map snapshot to wavefield boundary
    # left side
    for i = 1 : order
        ilower = (npml+i-1) * Nz + ntop + 1
        iupper = ilower + nz - 1
        idx1   = (i-1)*nz + 1
        idx2   = idx1 +nz - 1
        index2[idx1:idx2] = collect(ilower : iupper)
    end

    # top side and bottom side
    icount = order * nz
    for ix = npml+order+1 : npml+nx-order
        tmp = (ix-1) * Nz + ntop

        for iz = 1 : order
            icount = icount + 1
            index2[icount] = tmp + iz
            icount = icount + 1
            index2[icount] = tmp + nz - iz + 1
        end
    end

    # right side
    for i = 1 : order
        ilower = (npml+nx-order+i-1) * Nz + ntop + 1
        iupper = ilower + nz - 1
        index2[icount+1: icount+nz] = collect(ilower : iupper)
        icount = icount + nz
    end

    # map wavefield to the boundary of Wavefield
    # left side
    index3[1:order*nz] = collect(1:order*nz)

    # top and bottom part in the middle
    icount = order * nz
    for ix = order+1 : nx-order
        tmp = (ix-1) * nz

        for iz = 1 : order
            icount = icount + 1
            index3[icount] = tmp + iz
            icount = icount + 1
            index3[icount] = tmp + nz - iz + 1
        end
    end

    # right side
    index3[end-order*nz+1:end] = collect((nx-order)*nz+1 : nz*nx)

    return index1, index2, index3

end

"""
   compute the finite-difference coefficients based on Tylor series expansion
"""
function finite_difference_coefficients(M::Integer)

    if M < 1
       error("M must be bigger or equal to 1")
    elseif M == 1
       c = ones(1)
       return c
    elseif M >= 2
       c = zeros(M)
       for i = 1 : M
           left  = (-1)^(i+1) / (2*i-1)
           right = 1.0
           for j = 1 : M
               if j != i
                  right = right * abs((2*j-1)^2 / ((2*i-1)^2 - (2*j-1)^2))
               end
           end
           c[i] = left * right
       end
       return c
    end
end

# finite difference coefficients estimated by minimizing LS of dispersion equation
const fdc = [[0.1129042e1, -0.4301412e-1],
             [0.1185991e1, -0.7249965e-1, 0.6301572e-2],
             [0.1217990e1, -0.9382142e-1, 0.1507536e-1, -0.1700324e-2],
             [0.1236607e1, -0.1082265   , 0.2343440e-1, -0.5033546e-2, 0.6817483e-3],
             [0.1247662e1, -0.1175538   , 0.2997970e-1, -0.8719078e-2, 0.2215897e-2, -0.3462075e-3],
             [0.1254799e1, -0.1238928   , 0.3494371e-1, -0.1208897e-1, 0.4132531e-2, -0.1197110e-2, 0.2122227e-3],
             [0.1259312e1, -0.1280347   , 0.3841945e-1, -0.1473229e-1, 0.5924913e-2, -0.2248618e-2, 0.7179226e-3, -0.1400855e-3],
             [0.1262502e1, -0.1310244   , 0.4103928e-1, -0.1686807e-1, 0.7530520e-2, -0.3345071e-2, 0.1380367e-2, -0.4808410e-3, 0.1023759e-3],
             [0.1264748e1, -0.1331606   , 0.4296909e-1, -0.1851897e-1, 0.8861071e-2, -0.4347073e-2, 0.2076101e-2, -0.9164925e-3, 0.3437446e-3, -0.7874250e-4]];

"""
   test the stable condition of finite difference method
"""
function is_stable(vel::Matrix{Tv}, dz::Tv, dx::Tv, dt::Tv) where {Tv<:AbstractFloat}

    vmax = maximum(vel)
    beta = vmax * sqrt(1/(dz^2) + 1/(dx^2))
    tt   = 1.0 / beta

    if beta * dt >= 1.0
       println("maximum acceptable time step: $tt")
       println("current dt is $dt")
       error("decrease dt or increase dz or dx")
    end

    return true
end

"""
   the constructor for TdParams
"""
function TdParams(rho, vel, free_surface::Bool, dz, dx, dt, tmax;
         data_format=Float32, fd_flag="taylor", order=2, npml=20, apml=900.) where {Ti<:Int64}

    # get the size of the model
    (nz, nx) = size(rho)
    if (nz, nx) != size(vel)
       error("model size mismatch")
    end

    # unify the data format
    rho = convert(Matrix{data_format}, rho)
    vel = convert(Matrix{data_format}, vel)
    dz  = convert(data_format, dz)
    dx  = convert(data_format, dx)
    dt  = convert(data_format, dt)
    tmax= convert(data_format, tmax)

    # test the stable condition
    is_stable(vel, dz, dx, dt)

    # compute finite difference coefficients
    if fd_flag == "taylor"
       fd_coefficients = finite_difference_coefficients(order)
    elseif fd_flag == "ls"
       fd_coefficients = fdc[order-1]
    else
       error("non-supported method for computing FD coefficients")
    end
    fd_coefficients = convert(Vector{data_format}, fd_coefficients)
    order = length(fd_coefficients)

    # size of model after padding PML layers
    if free_surface
       Nz   = nz +     npml
       ntop = 0
    else
       Nz   = nz + 2 * npml
       ntop = npml
    end
    Nx = nx + 2 * npml

    # number of time stepping
    nt = round(Int64, tmax/dt) + 1

    # compute auxillary index vectors
    (spt2wfd, spt2bnd, wfd2bnd) = get_index(nz, nx, npml, free_surface; order=order)

    # Z direction pml coefficients
    if free_surface
       (gammaz_integer, gammaz_stagger) = get_gamma_profile(Nz, npml, dz; apml=apml, oneside_flag=true)
    else
       (gammaz_integer, gammaz_stagger) = get_gamma_profile(Nz, npml, dz; apml=apml, oneside_flag=false)
    end

    # X direction pml coefficients
    (gammax_integer, gammax_stagger) = get_gamma_profile(Nx, npml, dx; apml=apml, oneside_flag=false)

    # data format convert
    gammaz_integer = convert(Vector{data_format}, gammaz_integer)
    gammaz_stagger = convert(Vector{data_format}, gammaz_stagger)
    gammax_integer = convert(Vector{data_format}, gammax_integer)
    gammax_stagger = convert(Vector{data_format}, gammax_stagger)

    # compute buoyancy and bulk modulus
    buoy0 = density2buoyancy(rho)
    bulk0 = vp2lambda(rho, vel)

    # implement free surface condition
    if free_surface
       for j = 1 : nx
           bulk0[1,j] = data_format(0.0)
       end
    end

    # padding the model parameters to accomdate PML
    buoy1 = model_padding(buoy0, npml, free_surface)
    bulk1 = model_padding(bulk0, npml, free_surface)

    # FD stencil with PML boundaries
    (MvzBvz, MvzBp , dpdz) = Mvz(gammaz_stagger, buoy1, dz, dt, fd_coefficients)
    (MvxBvx, MvxBp , dpdx) = Mvx(gammax_stagger, buoy1, dx, dt, fd_coefficients)
    (MpzBpz, MpzBvz, dvdz) = Mpz(gammaz_integer, bulk1, dz, dt, fd_coefficients)
    (MpxBpx, MpxBvx, dvdx) = Mpx(gammax_integer, bulk1, dx, dt, fd_coefficients)

    # FD stencil with rigid boundaries
    (RvzBp , rpdz) = Rvz(buoy0, dz, dt, fd_coefficients)
    (RvxBp , rpdx) = Rvx(buoy0, dx, dt, fd_coefficients)
    (RpBv, rvdz, rvdx) = Rpzx(bulk0, dz, dx, dt, fd_coefficients)

    # call the default struct constructor
    return TdParams(data_format, nz, nx, npml, free_surface, Nz, Nx, ntop,
                       spt2wfd, spt2bnd, wfd2bnd, order,
                       dz, dx, dt, tmax, nt, rho, vel,
                       MvzBvz, MvzBp , RvzBp , dpdz, rpdz,
                       MvxBvx, MvxBp , RvxBp , dpdx, rpdx,
                       MpzBpz, MpzBvz,         dvdz, rvdz,
                       MpxBpx, MpxBvx, RpBv  , dvdx, rvdx)
end

# """
#    immutable struct contain the model parameters
# *  data_format : Float32 or Float64
# *  nz          : depth of model
# *  nx          : length of model
# *  npml        : number of PML layers
# *  free_surface: true -> no obsorbtion on the surface
#                  false-> PML on the surface
# *  Nz          : depth of model (include PML padding)
# *  Nx          : length of model (include PML padding)
# *  ntop        : number of padded PML layers on the top
# *  spt2wfd     : auxillary index vector for computing wavefield from snapshot
# *  spt2bnd     : auxillary index vector for computing wavefield boundaries from snapshot
# *  bnd2wfd     : auxillary index vector for inserting wavefield boundaries to wavefield
# *  dz          : vertical grid size
# *  dx          : horizontal grid size
# *  dt          : time step size
# *  tmax        : maximum modelling length
# *  nt          : number of time steps
# *  rho         : 2D density model
# *  vel         : 2D P-wave velocity model
# *  fd_coefficients : finite difference coefficients
# *  order       : the order of precision, will decide the number of layers required for wavefield reconstruction
# """
# struct TdParams{Ti<:Int64, Tv<:AbstractFloat}
#
#     data_format  :: DataType
#     nz           :: Ti
#     nx           :: Ti
#     npml         :: Ti
#     free_surface :: Bool
#     Nz           :: Ti
#     Nx           :: Ti
#     ntop         :: Ti
#     spt2wfd      :: Vector{Ti}
#     spt2bnd      :: Vector{Ti}
#     bnd2wfd      :: Vector{Ti}
#     dz           :: Tv
#     dx           :: Tv
#     dt           :: Tv
#     tmax         :: Tv
#     nt           :: Ti
#     rho          :: Matrix{Tv}
#     vel          :: Matrix{Tv}
#     fd_coefficients :: Vector{Tv}
# end
#
# """
#    Overloading the show function for TdParams
# """
# function show(io::IO, params::TdParams)
#
#     # size of model
#     @printf("nz = %4d, nx = %4d, npml = %4d\n", params.nz, params.nx, params.npml)
#     @printf("dz = %4.1f, dx = %4.1f\n", params.dz, params.dx)
#     @printf("time step size = %4.4f, max simulation time = %4.1f, number of step = %5d\n", params.dt, params.tmax, params.nt)
#
#     # top boundary condition
#     if params.free_surface
#        @printf("top boundary conditon = %20s\n", "free surface")
#     else
#        @printf("top boundary conditon = %20s\n", "PML obsorbing")
#     end
#
#     # data presicion (Float32, Float64)
#     @printf("data_format=%16s\n", params.data_format)
#
#     return nothing
# end
#
# """
#    Padding PML obsorbing layers to model parameters
# """
# function  model_padding(m::Matrix{Tv}, npml::Ti, free_surface::Bool) where {Ti<:Int64, Tv<:AbstractFloat}
#
#     # model size
#     (nz, nx) = size(m)
#
#     # model size after padding
#     Nx = nx + 2 * npml
#     if free_surface
#        Nz   = nz + npml
#        ntop = 0
#     else
#        Nz   = nz + npml * 2
#        ntop = npml
#     end
#
#     # allocate memory for the padded model
#     mpad = zeros(eltype(m), Nz, Nx)
#
#     # central part
#     for ix = 1 : nx
#         for iz = 1 : nz
#             mpad[iz+ntop, ix+npml] = m[iz, ix]
#         end
#     end
#
#     # left and right sides
#     for ix = 1 : npml
#         for iz = 1 : nz
#             mpad[iz+ntop,ix        ] = m[iz,1 ]
#             mpad[iz+ntop,ix+npml+nx] = m[iz,nx]
#         end
#     end
#
#     # upper and bottom sides
#     for ix = 1 : nx
#         for iz = 1 : npml
#
#             if free_surface == false              #  top
#                mpad[iz,ix+npml] = m[1,ix]
#             end
#             mpad[iz+nz+ntop,ix+npml] = m[nz,ix]   # bottom
#         end
#     end
#
#     # corners
#     for ix = 1 : npml
#         for iz = 1 : npml
#             if free_surface == false
#                mpad[iz, ix        ] = m[1, 1 ]    # upper-left
#                mpad[iz, ix+npml+nx] = m[1, nx]    # upper-right
#             end
#             mpad[iz+ntop+nz, ix        ] = m[nz, 1 ]   # bottom-left
#             mpad[iz+ntop+nz, ix+npml+nx] = m[nz, nx]   # bottom-right
#         end
#     end
#
#     return mpad
# end
#
# """
#    Smoothing the 2D physical model with Hanning or Gaussian window whose half length is L
# """
# function model_smooth(par::Matrix{Tv}, L::Ti; filter_flag="hanning") where {Ti<:Int64, Tv<:AbstractFloat}
#
#     data_format = eltype(par)
#
#     # size of smoother
#     n = 2*L + 1
#
#     # padding model before smoothing to handle boundary part
#     par1 = model_padding(par, L, false)
#
#     # convert the type of smoother
#     if filter_flag == "hanning"
#        s = convert(Vector{data_format}, hanning(n))
#     elseif filter_flag == "gaussian"
#        s = convert(Vector{data_format}, gaussian(n, 0.5))
#     else
#        error("wronge filter, only supper hanning and gaussian smoother")
#     end
#     s = s / sum(s)
#
#     # smoothing via 2D convolution
#     par1 = conv2(s,s,par1)
#
#     # truncate to original size
#     par1 = par1[2*L+1:end-2*L, 2*L+1:end-2*L]
#
#     return par1
# end
#
# """
#    get the auxillary index vectors which facilitate extracting wavefield from snapshot,
# saving wavefield boundaries and insert the boundaries for wavefield reconstruction.
# """
# function get_index(nz::Ti, nx::Ti, npml::Ti, free_surface::Bool; order=2) where {Ti<:Int64}
#
#     # size of wavefield
#     len1 = nz * nx
#
#     # length of boundary values need to save
#     len2 = 2*order*nz + 2*order*(nx-2*order)
#
#     # allocate 3 vectors of integer
#     index1 = zeros(Int64, len1)
#     index2 = zeros(Int64, len2)
#     index3 = zeros(Int64, len2)
#
#     # size of padded model
#     if free_surface
#        Nz = nz + npml
#        ntop = 0
#     else
#        Nz = nz + npml * 2
#        ntop = npml
#     end
#     Nx = nx + npml * 2
#
#     # map snapshot to wavefield
#     icount = 0
#     for i2 = 1 : nx
#         tmp= (i2 + npml - 1)*Nz + ntop
#
#         for i1 = 1 : nz
#             icount = icount + 1
#             index1[icount] = tmp + i1
#         end
#     end
#
#     # map snapshot to wavefield boundary
#     # left side
#     for i = 1 : order
#         ilower = (npml+i-1) * Nz + ntop + 1
#         iupper = ilower + nz - 1
#         idx1   = (i-1)*nz + 1
#         idx2   = idx1 +nz - 1
#         index2[idx1:idx2] = collect(ilower : iupper)
#     end
#
#     # top side and bottom side
#     icount = order * nz
#     for ix = npml+order+1 : npml+nx-order
#         tmp = (ix-1) * Nz + ntop
#
#         for iz = 1 : order
#             icount = icount + 1
#             index2[icount] = tmp + iz
#             icount = icount + 1
#             index2[icount] = tmp + nz - iz + 1
#         end
#     end
#
#     # right side
#     for i = 1 : order
#         ilower = (npml+nx-order+i-1) * Nz + ntop + 1
#         iupper = ilower + nz - 1
#         index2[icount+1: icount+nz] = collect(ilower : iupper)
#         icount = icount + nz
#     end
#
#     # map wavefield to the boundary of Wavefield
#     # left side
#     index3[1:order*nz] = collect(1:order*nz)
#
#     # top and bottom part in the middle
#     icount = order * nz
#     for ix = order+1 : nx-order
#         tmp = (ix-1) * nz
#
#         for iz = 1 : order
#             icount = icount + 1
#             index3[icount] = tmp + iz
#             icount = icount + 1
#             index3[icount] = tmp + nz - iz + 1
#         end
#     end
#
#     # right side
#     index3[end-order*nz+1:end] = collect((nx-order)*nz+1 : nz*nx)
#
#     return index1, index2, index3
#
# end
#
# """
#    compute the finite-difference coefficients based on Tylor series expansion
# """
# function finite_difference_coefficients(M::Integer)
#
#     if M < 1
#        error("M must be bigger or equal to 1")
#     elseif M == 1
#        c = ones(1)
#        return c
#     elseif M >= 2
#        c = zeros(M)
#        for i = 1 : M
#            left  = (-1)^(i+1) / (2*i-1)
#            right = 1.0
#            for j = 1 : M
#                if j != i
#                   right = right * abs((2*j-1)^2 / ((2*i-1)^2 - (2*j-1)^2))
#                end
#            end
#            c[i] = left * right
#        end
#        return c
#     end
# end
#
# # finite difference coefficients estimated by minimizing LS of dispersion equation
# const fdc = [[0.1129042e1, -0.4301412e-1],
#              [0.1185991e1, -0.7249965e-1, 0.6301572e-2],
#              [0.1217990e1, -0.9382142e-1, 0.1507536e-1, -0.1700324e-2],
#              [0.1236607e1, -0.1082265   , 0.2343440e-1, -0.5033546e-2, 0.6817483e-3],
#              [0.1247662e1, -0.1175538   , 0.2997970e-1, -0.8719078e-2, 0.2215897e-2, -0.3462075e-3],
#              [0.1254799e1, -0.1238928   , 0.3494371e-1, -0.1208897e-1, 0.4132531e-2, -0.1197110e-2, 0.2122227e-3],
#              [0.1259312e1, -0.1280347   , 0.3841945e-1, -0.1473229e-1, 0.5924913e-2, -0.2248618e-2, 0.7179226e-3, -0.1400855e-3],
#              [0.1262502e1, -0.1310244   , 0.4103928e-1, -0.1686807e-1, 0.7530520e-2, -0.3345071e-2, 0.1380367e-2, -0.4808410e-3, 0.1023759e-3],
#              [0.1264748e1, -0.1331606   , 0.4296909e-1, -0.1851897e-1, 0.8861071e-2, -0.4347073e-2, 0.2076101e-2, -0.9164925e-3, 0.3437446e-3, -0.7874250e-4]];
#
# """
#    test the stable condition of finite difference method
# """
# function is_stable(vel::Matrix{Tv}, dz::Tv, dx::Tv, dt::Tv) where {Tv<:AbstractFloat}
#
#     vmax = maximum(vel)
#     beta = vmax * sqrt(1/(dz^2) + 1/(dx^2))
#     tt   = 1.0 / beta
#
#     if beta * dt >= 1.0
#        println("maximum acceptable time step: $tt")
#        println("current dt is $dt")
#        error("decrease dt or increase dz or dx")
#     end
#
#     return true
# end
#
# """
#    the constructor for TdParams
# """
# function TdParams(rho, vel, npml::Ti, free_surface::Bool, dz, dx, dt, tmax;
#          data_format=Float32, fd_flag="taylor", order=2) where {Ti<:Int64}
#
#     # get the size of the model
#     (nz, nx) = size(rho)
#     if (nz, nx) != size(vel)
#        error("model size mismatch")
#     end
#
#     # unify the data format
#     rho = convert(Matrix{data_format}, rho)
#     vel = convert(Matrix{data_format}, vel)
#     dz  = convert(data_format, dz)
#     dx  = convert(data_format, dx)
#     dt  = convert(data_format, dt)
#     tmax= convert(data_format, tmax)
#
#     # test the stable condition
#     is_stable(vel, dz, dx, dt)
#
#     # compute finite difference coefficients
#     if fd_flag == "taylor"
#        fd_coefficients = finite_difference_coefficients(order)
#     elseif fd_flag == "ls"
#        fd_coefficients = fdc[order-1]
#     else
#        error("non-supported method for computing FD coefficients")
#     end
#     fd_coefficients = convert(Vector{data_format}, fd_coefficients)
#     order = length(fd_coefficients)
#
#     # size of model after padding PML layers
#     if free_surface
#        Nz   = nz +     npml
#        ntop = 0
#     else
#        Nz   = nz + 2 * npml
#        ntop = npml
#     end
#     Nx = nx + 2 * npml
#
#     # number of time stepping
#     nt = round(Int64, tmax/dt) + 1
#
#     # compute auxillary index vectors
#     (spt2wfd, spt2bnd, bnd2wfd) = get_index(nz, nx, npml, free_surface; order=order)
#
#     # call the default struct constructor
#     return TdParams(data_format, nz, nx, npml, free_surface, Nz, Nx, ntop,
#            spt2wfd, spt2bnd, bnd2wfd, dz, dx, dt, tmax, nt, rho, vel, fd_coefficients)
#
# end


text:
########################################################
####       process data for estimation              ####
########################################################

function get_rowIDT(ivar)

    # Create a matrix of panel information.

    N  = length(unique(ivar)) # N=number of panels
    id = Array{Any}(undef,N,2)
    id[:,1] = unique(ivar)    # list of id with no repetition
    @inbounds for i = 1:N
        @views id[i,2]= sum(ivar.== id[i,1])    # number of periods for the panel
    end
    @views Tᵢ = id[:,2]

    rowID =  Vector{Vector}()
    @inbounds for i=1:N
        @views cc = findall(x-> x == id[i,1], ivar) # row index of i'th firm
        push!(rowID, cc) # put the id info in the vector; faster than using UnitRange
    end

    rowIDT = hcat(rowID, Tᵢ)

    return rowIDT # (Nx2): col_1 is panel's row info; col_2 is panel's number of periods
end

#=  replaced by sf_demean()
function get_INVtrM(rowIDT)

    # Create the within-transformation matrix for panels.
    # The matrix is panel-specific.

    N = size(rowIDT,1)
    INVtrM = Vector{Matrix{Float64}}()  # store the inverted matrix
    @inbounds for i=1:N
        @views trM = Matrix(I, rowIDT[i, 2], rowIDT[i, 2]) .- (1/rowIDT[i,2]) # transformation matrix
      # push!(INVtrM, pinv(trM)) # put the inverted matrix in the vector; note, trM is an idempotent matrix, and so the inverse is itself
        push!(INVtrM, trM) # put the inverted matrix in the vector; note, trM is an idempotent matrix, and so the inverse is itself
      end
    return INVtrM # (Nx1): ith row is the transformation matrix for the ith panel
end
=#

function sf_demean(a::AbstractArray)  # subtract mean from columns
    return a .- mean(a, dims=1)
end

########################################################
####        process variables for estimation        ####
########################################################

#? ----- truncated normal -----------------

function getvar(::Type{Trun}, dat::DataFrame)

    yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
    xvar = dat[:, _dicM[:frontier]]
    zvar = dat[:, _dicM[:μ]]
    wvar = dat[:, _dicM[:σᵤ²]]
    vvar = dat[:, _dicM[:σᵥ²]]

    #* --- model info printout ---------

    modelinfo1 = "normal and truncated-normal"

    if _dicM[:μ] == _dicM[:σᵤ²] && length(_dicM[:μ]) > 1
      modelinfo1 = "the non-monotonic model of Wang (2002, JPA), normal and truncated-normal"
    elseif _dicM[:μ] == _dicM[:σᵤ²] && length(_dicM[:μ]) == 1
      modelinfo1 = "the normal and truncated-normal model (Stevenson 1980, JoE)"
    elseif length(_dicM[:μ]) > 1 && length(_dicM[:σᵤ²] ) == 1
      modelinfo1 = "the normal and truncated-normal model of Battese and Coelli (1995, JoE)"
    else
      modelinfo1 = "the normal and truncated-normal model"
    end

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢ" below is changed to "+ uᵢ".

     $(_dicM[:depvar][1]) = frontier($(_dicM[:frontier])) + vᵢ - uᵢ,

     where vᵢ ∼ N(0, σᵥ²),
                σᵥ² = exp(log_σᵥ²)
                    = exp($(_dicM[:σᵥ²]));
           uᵢ ∼ N⁺(μ, σᵤ²),
                μ = $(_dicM[:μ]),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

    #* --- retrieve and generate important parameters -----

    #*  number of obs and number of variables
    nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list
    nofobs  = nrow(dat)
    nofx    = size(xvar,2)  # nofx: number of x vars
    nofz    = size(zvar,2)
    nofw    = size(wvar,2)
    nofv    = size(vvar,2)
    nofpara = nofx + nofz + nofw + nofv
    nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
              nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

    #* positions of the variables/parameters
    begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0
    begx = 1
    endx = (nofx)
    begz = (nofx) + 1
    endz = (nofx) + (nofz )
    begw = (nofx) + (nofz ) + 1
    endw = (nofx) + (nofz ) + (nofw )
    begv = (nofx) + (nofz ) + (nofw ) + 1
    endv = nofpara
    posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
              begq=begq, endq=endq, begw=begw, endw=endw,
              begv=begv, endv=endv)

    #* create equation names and mark positions for making tables
    eqvec = (frontier = begx + 1,
                    μ = begz + 1,
              log_σᵤ² = begw + 1,
              log_σᵥ² = begv + 1)

    #* create equation names and mark positions
    eqvec2 = (coeff_frontier = (begx:endx),
              coeff_μ        = (begz:endz),
              coeff_log_σᵤ²  = (begw:endw),
              coeff_log_σᵥ²  = (begv:endv))

    #* retrieve variable names for making tables
    xnames  = names(xvar)
    znames  = names(zvar)
    wnames  = names(wvar)
    vnames  = names(vvar)
    varlist = vcat(" ", xnames, znames, wnames, vnames)

    #* Converting the dataframe to matrix in order to do computation
        # Matrix(.) converts DataFrame to Matrix, but DataFrame(load(.)) would
        # have Union{Missing, Float64} type, which causes problems in
        # `ForwardDiff` in `SFmarginal.jl`.
    yvar = convert(Array{Float64}, Matrix(yvar))
    xvar = convert(Array{Float64}, Matrix(xvar))
    zvar = convert(Array{Float64}, Matrix(zvar))
    wvar = convert(Array{Float64}, Matrix(wvar))
    vvar = convert(Array{Float64}, Matrix(vvar))

    qvar = ()

    dum2 = ()

    return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar, vvar,       dum2, varlist

  end

#? --------- half-normal ----------------

function getvar(::Type{Half}, dat::DataFrame)

    yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
    xvar = dat[:, _dicM[:frontier]]
  # zvar = dat[:, _dicM[:μ]]
    wvar = dat[:, _dicM[:σᵤ²]]
    vvar = dat[:, _dicM[:σᵥ²]]

    #* --- model info printout ---------

    modelinfo1 = "normal and half-normal"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢ" below is changed to "+ uᵢ".

     $(_dicM[:depvar][1]) = frontier($(_dicM[:frontier])) + vᵢ - uᵢ,

     where vᵢ ∼ N(0, σᵥ²),
                σᵥ² = exp(log_σᵥ²)
                    = exp($(_dicM[:σᵥ²]));
           uᵢ ∼ N⁺(0, σᵤ²),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

    #* --- retrieve and generate important parameters -----

    #*  number of obs and number of variables
    nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

    nofobs  = nrow(dat)
    nofx    = size(xvar,2)  # nofx: number of x vars
  # nofz    = size(zvar,2)
    nofw    = size(wvar,2)
    nofv    = size(vvar,2)
    nofpara = nofx + nofz + nofw + nofv

    nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
              nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

    #* positions of the variables/parameters
    begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

    begx = 1
    endx = (nofx)
 #  begz = (nofx) + 1
 #  endz = (nofx) + (nofz )
    begw = (nofx) + (nofz ) + 1
    endw = (nofx) + (nofz ) + (nofw )
    begv = (nofx) + (nofz ) + (nofw ) + 1
    endv = nofpara

    posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
              begq=begq, endq=endq, begw=begw, endw=endw,
              begv=begv, endv=endv)

    #* create equation names and mark positions for making tables
    eqvec = (frontier = begx + 1,
           #        μ = begz + 1,
              log_σᵤ² = begw + 1,
              log_σᵥ² = begv + 1)

    #* create equation names and mark positions
    eqvec2 = (coeff_frontier = (begx:endx),
              coeff_log_σᵤ²  = (begw:endw),
              coeff_log_σᵥ²  = (begv:endv))

    #* retrieve variable names for making tables
    xnames = names(xvar)
  # znames = names(zvar)
    wnames = names(wvar)
    vnames = names(vvar)
    varlist= vcat(" ", xnames, # znames,
                   wnames, vnames)

    #* Converting the dataframe to matrix in order to do computation
    yvar = convert(Array{Float64}, Matrix(yvar))
    xvar = convert(Array{Float64}, Matrix(xvar))
  # zvar = convert(Array{Float64}, Matrix(zvar))
    wvar = convert(Array{Float64}, Matrix(wvar))
    vvar = convert(Array{Float64}, Matrix(vvar))

    zvar = ()
    qvar = ()
    dum2 = ()

    return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar, vvar,       dum2, varlist

end

#? ---------- Exponential (same as Half, only change equation names) ----------------

function getvar(::Type{Expo}, dat::DataFrame)

    yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
    xvar = dat[:, _dicM[:frontier]]
  # zvar = dat[:, _dicM[:μ]]
    wvar = dat[:, _dicM[:σᵤ²]]
    vvar = dat[:, _dicM[:σᵥ²]]

    #* --- model info printout ---------

    modelinfo1 = "normal and exponential"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢ" below is changed to "+ uᵢ".

     $(_dicM[:depvar][1]) = frontier($(_dicM[:frontier])) + vᵢ - uᵢ,

     where vᵢ ∼ N(0, σᵥ²),
                σᵥ² = exp(log_σᵥ²)
                    = exp($(_dicM[:σᵥ²]));
           uᵢ ∼ Exp(1/σᵤ) = (1/σᵤ)e^(-(1/σᵤ)U),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

    #* --- retrieve and generate important parameters -----

    #*   number of obs and number of variables
    nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

    nofobs  = nrow(dat)
    nofx    = size(xvar,2)  # nofx: number of x vars
  # nofz    = size(zvar,2)
    nofw    = size(wvar,2)
    nofv    = size(vvar,2)
    nofpara = nofx + nofz + nofw + nofv

    nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
              nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

    #* positions of the variables/parameters
    begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

    begx = 1
    endx = (nofx)
  # begz = (nofx) + 1
  # endz = (nofx) + (nofz )
    begw = (nofx) + (nofz ) + 1
    endw = (nofx) + (nofz ) + (nofw )
    begv = (nofx) + (nofz ) + (nofw ) + 1
    endv = nofpara

    posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
              begq=begq, endq=endq, begw=begw, endw=endw,
              begv=begv, endv=endv)

    #* create equation names and mark positions for making tables
    eqvec = (frontier = begx + 1,
           #        μ = begz + 1,
              log_σᵤ² = begw + 1,  # old: λ
              log_σᵥ² = begv + 1)

    #* create equation names and mark positions
    eqvec2 = (coeff_frontier = (begx:endx),
              coeff_log_σᵤ²  = (begw:endw),
              coeff_log_σᵥ²  = (begv:endv))

    #* retrieve variable names for making tables
    xnames  = names(xvar)
  # znames  = names(zvar)
    wnames  = names(wvar)
    vnames  = names(vvar)
    varlist = vcat(" ", xnames, # znames,
                   wnames, vnames)

    #* Converting the dataframe to matrix in order to do computation

    yvar  = convert(Array{Float64}, Matrix(yvar))
    xvar  = convert(Array{Float64}, Matrix(xvar))
  # zvar  = Matrix(zvar)
    wvar  = convert(Array{Float64}, Matrix(wvar))
    vvar  = convert(Array{Float64}, Matrix(vvar))

    zvar = ()
    qvar = ()

    dum2 = ()

    return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar, vvar,       dum2, varlist

end

#? --------- scaling property model ----------------

function getvar(::Type{Trun_Scale}, dat::DataFrame)

    yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
    xvar = dat[:, _dicM[:frontier]]
    zvar = dat[:, _dicM[:μ]]
    qvar = dat[:, _dicM[:hscale]]
    wvar = dat[:, _dicM[:σᵤ²]]
    vvar = dat[:, _dicM[:σᵥ²]]

    #* --- model info printout ---------

    modelinfo1 = "normal and truncated-normal with the scaling property (e.g., Wang and Schmidt (2002 JPA))"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢ" below is changed to "+ uᵢ".

     $(_dicM[:depvar][1]) = frontier($(_dicM[:frontier])) + vᵢ - uᵢ,

     where vᵢ ∼ N(0, σᵥ²),
                σᵥ² = exp(log_σᵥ²)
                    = exp($(_dicM[:σᵥ²]));
           uᵢ ∼ hscaleᵢ * N⁺(μ, σᵤ²),
                hscaleᵢ = exp($(_dicM[:hscale])),
                μ = $(_dicM[:μ]),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

    #* --- retrieve and generate important parameters -----

    #*   number of obs and number of variables
    nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

    nofobs  = nrow(dat)
    nofx    = size(xvar,2)  # nofx: number of x vars
    nofz    = size(zvar,2)
    nofq    = size(qvar,2)
    nofw    = size(wvar,2)
    nofv    = size(vvar,2)
    nofpara = nofx + nofz + nofq + nofw + nofv

    nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
              nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

    #* positions of the variables/parameters
    begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

    begx = 1
    endx = nofx
    begz = endx + 1
    endz = begz + nofz-1
    begq = endz + 1
    endq = begq + nofq-1
    begw = endq + 1
    endw = begw + nofw-1
    begv = endw + 1
    endv = nofpara

    posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
              begq=begq, endq=endq, begw=begw, endw=endw,
              begv=begv, endv=endv)

    #* create equation names and mark positions for making tables
    eqvec = (frontier = begx + 1,
                    μ = begz + 1,
           log_hscale = begq + 1,
              log_σᵤ² = begw + 1,
              log_σᵥ² = begv + 1)

    #* create equation names and mark positions
    eqvec2 = (coeff_frontier = (begx:endx),
                     coeff_μ = (begz:endz),
            coeff_log_hscale = (begq:endq),
               coeff_log_σᵤ² = (begw:endw),
               coeff_log_σᵥ² = (begv:endv))

    #* retrieve variable names for making tables
    xnames  = names(xvar)
    znames  = names(zvar)
    qnames  = names(qvar)
    wnames  = names(wvar)
    vnames  = names(vvar)
    varlist = vcat(" ", xnames, znames, qnames, wnames, vnames)

    #* Converting the dataframe to matrix in order to do computation
    yvar  = convert(Array{Float64}, Matrix(yvar))
    xvar  = convert(Array{Float64}, Matrix(xvar))
    zvar  = convert(Array{Float64}, Matrix(zvar))
    qvar  = convert(Array{Float64}, Matrix(qvar))
    wvar  = convert(Array{Float64}, Matrix(wvar))
    vvar  = convert(Array{Float64}, Matrix(vvar))

    #* various functions can and cannot contain a constant, check! ---- *#
    checkConst(zvar, :μ,      @requireConst(1))
    checkConst(qvar, :hscale, @requireConst(0))
    checkConst(wvar, :σᵤ²,    @requireConst(1))

    dum2 = ()

    return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar, vvar,       dum2, varlist

end

#? --------- panel FE Wang and Ho, half-normal ----------------

function getvar(::Type{PFEWHH}, dat::DataFrame)

  yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
  xvar = dat[:, _dicM[:frontier]]
# zvar = dat[:, _dicM[:μ]]
  qvar = dat[:, _dicM[:hscale]]
  wvar = dat[:, _dicM[:σᵤ²]]
  vvar = dat[:, _dicM[:σᵥ²]]

  tvar = dat[:, _dicM[:timevar]]
  ivar = dat[:, _dicM[:idvar]] # need the [1] otherwise causes problems down the road

    #* --- model info printout ---------

    modelinfo1 = "true fixed effect of Wang and Ho (2010 JE), normal and half-normal"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

     $(_dicM[:depvar][1]) = frontier(αᵢ + $(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

     where vᵢₜ ∼ N(0, σᵥ²),
                 σᵥ² = exp(log_σᵥ²)
                     = exp($(_dicM[:σᵥ²]));
           uᵢₜ ∼ hscaleᵢₜ * uᵢ,
                 hscaleᵢₜ = exp($(_dicM[:hscale])),
           uᵢ ∼ N⁺(0, σᵤ²),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

  #* --- retrieve and generate important parameters -----

  #*   number of obs and number of variables
  nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

  nofobs  = nrow(dat)
  nofx    = size(xvar,2)  # nofx: number of x vars
# nofz    = size(zvar,2)
  nofq    = size(qvar,2)
  nofw    = size(wvar,2)
  nofv    = size(vvar,2)
  nofpara = nofx + nofz + nofq + nofw + nofv

  nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
            nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

  #* positions of the variables/parameters
  begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

  begx = 1
  endx = nofx
# begz = endx + 1
# endz = begz + nofz -1
  begq = endx + 1
  endq = begq + nofq -1
  begw = endq + 1
  endw = begw + nofw -1
  begv = endw + 1
  endv = nofpara

  posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
            begq=begq, endq=endq, begw=begw, endw=endw,
            begv=begv, endv=endv)

  #* create equation names and mark positions for making tables
  eqvec = (frontier = begx + 1,
           #      μ = begz + 1,
         log_hscale = begq + 1,
            log_σᵤ² = begw + 1,
            log_σᵥ² = begv + 1)

  #* create equation names and mark positions
  eqvec2 = (coeff_frontier = (begx:endx),
          coeff_log_hscale = (begq:endq),
             coeff_log_σᵤ² = (begw:endw),
             coeff_log_σᵥ² = (begv:endv))

  #* retrieve variable names for making tables
  xnames  = names(xvar)
# znames  = names(zvar)
  qnames  = names(qvar)
  wnames  = names(wvar)
  vnames  = names(vvar)
  varlist = vcat(" ", xnames, # znames,
                      qnames, wnames, vnames)

  #* Converting the dataframe to matrix in order to do computation
  yvar = convert(Array{Float64}, Matrix(yvar))
  xvar = convert(Array{Float64}, Matrix(xvar))
# zvar = convert(Array{Float64}, Matrix(zvar))
  qvar = convert(Array{Float64}, Matrix(qvar))
  wvar = convert(Array{Float64}, Matrix(wvar))
  vvar = convert(Array{Float64}, Matrix(vvar))

  tvar = convert(Array{Float64}, Matrix(tvar))
  ivar = convert(Array{Float64}, Matrix(ivar))

  zvar = ()

  #* various functions can and cannot contain a constant, check! ---- *#
  checkConst(xvar, :frontier, @requireConst(0))
  checkConst(qvar, :hscale,   @requireConst(0))
# checkConst(zvar, :μ,        @requireConst(1))
  checkConst(wvar, :σᵤ²,      @requireConst(1))
  checkConst(vvar, :σᵥ²,      @requireConst(1))

  #* panel info and within transformation

  rowIDT = get_rowIDT(vec(ivar))   # (Nx2): col_1 is panel's row info; col_2 is panel's number of periods

  yxdata = hcat(yvar, xvar)
       D = zeros(nofobs, 1+nofx)  # pre-allocate the transformed dataset
       N = length(unique(vec(ivar)))   # N=number of panels

  for i=1:N  # within-transform the data
      @views D[rowIDT[i,1], :] = sf_demean(yxdata[rowIDT[i,1], :]) # INVtrM[i] * yxdata[rowIDT[i,1], :] # transform the data
  end

  ỹ = (D[:, 1])
  x̃ = (D[:, 2:end])

  return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, ỹ, x̃, zvar, qvar, wvar, vvar,         rowIDT, varlist

end

#?--------- panel FE Wang and Ho, truncated normal ----------------

function getvar(::Type{PFEWHT}, dat::DataFrame)

  yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
  xvar = dat[:, _dicM[:frontier]]
  zvar = dat[:, _dicM[:μ]]
  qvar = dat[:, _dicM[:hscale]]
  wvar = dat[:, _dicM[:σᵤ²]]
  vvar = dat[:, _dicM[:σᵥ²]]

  tvar = dat[:, _dicM[:timevar]]
  ivar = dat[:, _dicM[:idvar]]

  #* --- model info printout ---------
  modelinfo1 = "true fixed effect model of Wang and Ho (2010 JE), normal and truncated-normal"
  modelinfo2 = begin
   """
   * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

   $(_dicM[:depvar][1]) = frontier(αᵢ + $(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

   where vᵢₜ ∼ N(0, σᵥ²),
               σᵥ² = exp(log_σᵥ²)
                   = exp($(_dicM[:σᵥ²]));
         uᵢₜ ∼ hscaleᵢₜ * uᵢ,
               hscaleᵢₜ = exp($(_dicM[:hscale])),
         uᵢ ∼ N⁺(μ, σᵤ²),
              μ = $(_dicM[:μ])
              σᵤ² = exp(log_σᵤ²)
                  = exp($(_dicM[:σᵤ²]));
   """
  end

  #* --- retrieve and generate important parameters -----

  #*   number of obs and number of variables
  nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

  nofobs  = nrow(dat)
  nofx    = size(xvar,2)  # nofx: number of x vars
  nofz    = size(zvar,2)
  nofq    = size(qvar,2)
  nofw    = size(wvar,2)
  nofv    = size(vvar,2)
  nofpara = nofx + nofz + nofq + nofw + nofv

  nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
            nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

  #* positions of the variables/parameters
  begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

  begx = 1
  endx = nofx
  begz = endx + 1
  endz = begz + nofz-1
  begq = endz + 1
  endq = begq + nofq-1
  begw = endq + 1
  endw = begw + nofw-1
  begv = endw + 1
  endv = nofpara

  posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
            begq=begq, endq=endq, begw=begw, endw=endw,
            begv=begv, endv=endv)

  #* create equation names and mark positions for making tables
  eqvec = (frontier = begx + 1,
                  μ = begz + 1,
         log_hscale = begq + 1,
            log_σᵤ² = begw + 1,
            log_σᵥ² = begv + 1)

  #* create equation names and mark positions
  eqvec2 = (coeff_frontier = (begx:endx),
                   coeff_μ = (begz:endz),
          coeff_log_hscale = (begq:endq),
             coeff_log_σᵤ² = (begw:endw),
             coeff_log_σᵥ² = (begv:endv))

  #* retrieve variable names for making tables
  xnames  = names(xvar)
  znames  = names(zvar)
  qnames  = names(qvar)
  wnames  = names(wvar)
  vnames  = names(vvar)
  varlist = vcat(" ", xnames, znames, qnames, wnames, vnames)

  #* Converting the dataframe to matrix in order to do computation
  yvar  = convert(Array{Float64}, Matrix(yvar))
  xvar  = convert(Array{Float64}, Matrix(xvar))
  zvar  = convert(Array{Float64}, Matrix(zvar))
  qvar  = convert(Array{Float64}, Matrix(qvar))
  wvar  = convert(Array{Float64}, Matrix(wvar))
  vvar  = convert(Array{Float64}, Matrix(vvar))

  tvar  = convert(Array{Float64}, Matrix(tvar))
  ivar  = convert(Array{Float64}, Matrix(ivar))

  #* various functions can and cannot contain a constant, check! ---- *#
  checkConst(xvar, :frontier, @requireConst(0))
  checkConst(zvar, :μ,        @requireConst(1))
  checkConst(qvar, :hscale,   @requireConst(0))
  checkConst(wvar, :σᵤ²,      @requireConst(1))
  checkConst(vvar, :σᵥ²,      @requireConst(1))

  #* Within transformation

  rowIDT = get_rowIDT(vec(ivar))   # rowIDT (Nx2): col_1 is panel's row info; col_2 is panel's number of periods

  yxdata = hcat(yvar, xvar)
       D = zeros(nofobs, 1+nofx)      # pre-allocate the transformed dataset
       N = length(unique(vec(ivar))) # N=number of panels

  for i=1:N
      @views D[rowIDT[i,1], :] = sf_demean(yxdata[rowIDT[i,1], :]) # INVtrM[i] * yxdata[rowIDT[i,1], :] # transform the data
  end

  ỹ = (D[:, 1])
  x̃ = (D[:, 2:end])

  return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, ỹ, x̃, zvar, qvar, wvar, vvar,         rowIDT, varlist

end

#?--------- time decay model   ----------------

# function getvar(::Type{PanDecay}, dat::DataFrame)

function getvar(::Union{Type{PanDecay}, Type{PanKumb90}}, dat::DataFrame)

  yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
  xvar = dat[:, _dicM[:frontier]]
  zvar = dat[:, _dicM[:μ]]
  qvar = dat[:, _dicM[:gamma]]
  wvar = dat[:, _dicM[:σᵤ²]]
  vvar = dat[:, _dicM[:σᵥ²]]

  ivar = dat[:, _dicM[:idvar]] # need the [1] otherwise causes problems down the road

  #* --- model info printout ---------

 if _dicM[:panel][1]  == :TimeDecay

      modelinfo1 = "panel time-decay model of Battese and Coelli (1992) (see also Wang and Kumbhakar 2005)"
      modelinfo2 = begin
      """
      * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

      $(_dicM[:depvar][1]) = frontier($(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

      where vᵢₜ ∼ N(0, σᵥ²),
                  σᵥ² = exp(log_σᵥ²)
                      = exp($(_dicM[:σᵥ²]));
            uᵢₜ ∼ gammaₜ * uᵢ,
                  gammaₜ = exp($(_dicM[:gamma])),
            uᵢ ∼ N⁺(μ, σᵤ²),
                  μ = $(_dicM[:μ])
                  σᵤ² = exp(log_σᵤ²)
                      = exp($(_dicM[:σᵤ²]));
      """
      end

  end

  if _dicM[:panel][1]  == :Kumbhakar1990

    modelinfo1 = "panel Kumbhakar (1990) model"
    modelinfo2 = begin
    """
    * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

    $(_dicM[:depvar][1]) = frontier($(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

    where vᵢₜ ∼ N(0, σᵥ²),
                σᵥ² = exp(log_σᵥ²)
                    = exp($(_dicM[:σᵥ²]));
          uᵢₜ ∼ gammaₜ * uᵢ,
                gammaₜ = 2/(1 + exp($(_dicM[:gamma]))),
          uᵢ ∼ N⁺(μ, σᵤ²),
                μ = $(_dicM[:μ])
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
    """
    end

  end

  #* --- retrieve and generate important parameters -----

  #*   number of obs and number of variables
  nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

  nofobs  = nrow(dat)
  nofx    = size(xvar,2)  # nofx: number of x vars
  nofz    = size(zvar,2)
  nofq    = size(qvar,2)
  nofw    = size(wvar,2)
  nofv    = size(vvar,2)
  nofpara = nofx + nofz + nofq + nofw + nofv

  nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
            nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

  #* positions of the variables/parameters
  begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

  begx = 1
  endx = nofx
  begz = endx + 1
  endz = begz + nofz-1
  begq = endz + 1
  endq = begq + nofq-1
  begw = endq + 1
  endw = begw + nofw-1
  begv = endw + 1
  endv = nofpara

  posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
            begq=begq, endq=endq, begw=begw, endw=endw,
            begv=begv, endv=endv)

  #* create equation names and mark positions for making tables
  eqvec = (frontier = begx + 1,
                  μ = begz + 1,
          log_gamma = begq + 1,
            log_σᵤ² = begw + 1,
            log_σᵥ² = begv + 1)

  #* create equation names and mark positions
  eqvec2 = (coeff_frontier = (begx:endx),
                   coeff_μ = (begz:endz),
           coeff_log_gamma = (begq:endq),
             coeff_log_σᵤ² = (begw:endw),
             coeff_log_σᵥ² = (begv:endv))

  #* retrieve variable names for making tables
  xnames  = names(xvar)
  znames  = names(zvar)
  qnames  = names(qvar)
  wnames  = names(wvar)
  vnames  = names(vvar)
  varlist = vcat(" ", xnames, znames, qnames, wnames, vnames)

  #* Converting the dataframe to matrix in order to do computation
  yvar  = convert(Array{Float64}, Matrix(yvar))
  xvar  = convert(Array{Float64}, Matrix(xvar))
  zvar  = convert(Array{Float64}, Matrix(zvar))
  qvar  = convert(Array{Float64}, Matrix(qvar))
  wvar  = convert(Array{Float64}, Matrix(wvar))
  vvar  = convert(Array{Float64}, Matrix(vvar))
  ivar  = convert(Array{Float64}, Matrix(ivar))

  #* get panel related variables
  rowIDT = get_rowIDT(vec(ivar))   # (Nx2): col_1 is panel's row info; col_2 is panel's number of periods

  #* various functions should and/or should not contain a constant, check! ---- *#
  checkConst(zvar, :μ,      @requireConst(1), rowIDT[:,1])
  checkConst(qvar, :gamma,  @requireConst(0))
  checkConst(wvar, :σᵤ²,    @requireConst(1))
  checkConst(vvar, :σᵥ²,    @requireConst(1))

  return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2,  yvar, xvar, zvar, qvar, wvar, vvar,       rowIDT, varlist

end

#? --------- panel FE CSW 2014 JoE (CSN), half-normal ----------------

function getvar(::Type{PFECSWH}, dat::DataFrame)

  yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
  xvar = dat[:, _dicM[:frontier]]
# zvar = dat[:, _dicM[:μ]]
# qvar = dat[:, _dicM[:hscale]]
  wvar = dat[:, _dicM[:σᵤ²]]
  vvar = dat[:, _dicM[:σᵥ²]]

  tvar = dat[:, _dicM[:timevar]]
  ivar = dat[:, _dicM[:idvar]]
    #* --- model info printout ---------

    modelinfo1 = "true fixed effect of Chen, Schmidt, and Wang (2014 JE), normal and half-normal"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

     $(_dicM[:depvar][1]) = frontier(αᵢ + $(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

     where vᵢₜ ∼ N(0, σᵥ²),
                 σᵥ² = exp(log_σᵥ²)
                     = exp($(_dicM[:σᵥ²]));
           uᵢₜ ∼ N⁺(0, σᵤ²),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

  #* --- retrieve and generate important parameters -----

  #*   number of obs and number of variables
  nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

  nofobs  = nrow(dat)
  nofx    = size(xvar,2)  # nofx: number of x vars
# nofz    = size(zvar,2)
# nofq    = size(qvar,2)
  nofw    = size(wvar,2)
  nofv    = size(vvar,2)
  nofpara = nofx + nofz + nofq + nofw + nofv

  nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
            nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

  #* positions of the variables/parameters
  begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

  begx = 1
  endx = nofx
# begz = endx + 1
# endz = begz + nofz -1
# begq = endx + 1
# endq = begq + nofq -1
  begw = endx + 1
  endw = begw + nofw -1
  begv = endw + 1
  endv = nofpara

  posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
            begq=begq, endq=endq, begw=begw, endw=endw,
            begv=begv, endv=endv)

  #* create equation names and mark positions for making tables
  eqvec = (frontier = begx + 1,
           #      μ = begz + 1,
       # log_hscale = begq + 1,
            log_σᵤ² = begw + 1,
            log_σᵥ² = begv + 1)

  #* create equation names and mark positions
  eqvec2 = (coeff_frontier = (begx:endx),
             coeff_log_σᵤ² = (begw:endw),
             coeff_log_σᵥ² = (begv:endv))

  #* retrieve variable names for making tables
  xnames  = names(xvar)
# znames  = names(zvar)
# qnames  = names(qvar)
  wnames  = names(wvar)
  vnames  = names(vvar)
  varlist = vcat(" ", xnames, # znames,
                      # qnames,
                      wnames, vnames)

  #* Converting the dataframe to matrix in order to do computation
  yvar = convert(Array{Float64}, Matrix(yvar))
  xvar = convert(Array{Float64}, Matrix(xvar))
# zvar = convert(Array{Float64}, Matrix(zvar))
# qvar = convert(Array{Float64}, Matrix(qvar))
  wvar = convert(Array{Float64}, Matrix(wvar))
  vvar = convert(Array{Float64}, Matrix(vvar))

  tvar = convert(Array{Float64}, Matrix(tvar))
  ivar = convert(Array{Float64}, Matrix(ivar))

  zvar = ()
  qvar = ()

  #* various functions can and cannot contain a constant, check! ---- *#
  checkConst(xvar, :frontier, @requireConst(0))
# checkConst(qvar, :hscale,   @requireConst(0))
# checkConst(zvar, :μ,        @requireConst(1))
  checkConst(wvar, :σᵤ²,      @requireConst(1))
  checkConst(vvar, :σᵥ²,      @requireConst(1))

  #* panel info and within transformation

  rowIDT = get_rowIDT(vec(ivar))   # (Nx2): col_1 is panel's row info; col_2 is panel's number of periods

  yxdata = hcat(yvar, xvar)
       D = zeros(nofobs, 1+nofx)  # pre-allocate the transformed dataset
       N = length(unique(vec(ivar)))   # N=number of panels

  for i=1:N  # within-transform the data
      @views D[rowIDT[i,1], :] = sf_demean(yxdata[rowIDT[i,1], :]) # INVtrM[i] * yxdata[rowIDT[i,1], :] # transform the data
  end

  ỹ = (D[:, 1])
  x̃ = (D[:, 2:end])

  return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, ỹ, x̃, zvar, qvar, wvar, vvar,         rowIDT, varlist

end

#? --------- panel TRE (true random effect), half-normal ----------------

function getvar(::Type{PTREH}, dat::DataFrame)

  yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
  xvar = dat[:, _dicM[:frontier]]
# zvar = dat[:, _dicM[:μ]]
  qvar = dat[:, _dicM[:σₐ²]]
  wvar = dat[:, _dicM[:σᵤ²]]
  vvar = dat[:, _dicM[:σᵥ²]]

  tvar = dat[:, _dicM[:timevar]]
  ivar = dat[:, _dicM[:idvar]] # need the [1] otherwise causes problems down the road

    #* --- model info printout ---------

    modelinfo1 = "true random effect model of Greene (2005 JoE), normal and half-normal"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

     $(_dicM[:depvar][1]) = frontier(αᵢ + $(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

     where
            αᵢ ∼ N(0, σₐ²)
                 σₐ² = exp(log_σₐ²)
                     = exp($(_dicM[:σₐ²]));
           vᵢₜ ∼ N(0, σᵥ²),
                 σᵥ² = exp(log_σᵥ²)
                     = exp($(_dicM[:σᵥ²]));
           uᵢₜ ∼ N⁺(0, σᵤ²),
                σᵤ² = exp(log_σᵤ²)
                    = exp($(_dicM[:σᵤ²]));
     """
    end

  #* --- retrieve and generate important parameters -----

  #*   number of obs and number of variables
  nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

  nofobs  = nrow(dat)
  nofx    = size(xvar,2)  # nofx: number of x vars
# nofz    = size(zvar,2)
  nofq    = size(qvar,2)
  nofw    = size(wvar,2)
  nofv    = size(vvar,2)
  nofpara = nofx + nofz + nofq + nofw + nofv

  nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
            nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

  #* positions of the variables/parameters
  begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

  begx = 1
  endx = nofx
# begz = endx + 1
# endz = begz + nofz -1
  begq = endx + 1
  endq = begq + nofq -1
  begw = endq + 1
  endw = begw + nofw -1
  begv = endw + 1
  endv = nofpara

  posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
            begq=begq, endq=endq, begw=begw, endw=endw,
            begv=begv, endv=endv)

  #* create equation names and mark positions for making tables
  eqvec = (frontier = begx + 1,
           #      μ = begz + 1,
            log_σₐ² = begq + 1,
            log_σᵤ² = begw + 1,
            log_σᵥ² = begv + 1)

  #* create equation names and mark positions
  eqvec2 = (coeff_frontier = (begx:endx),
             coeff_log_σₐ² = (begq:endq),
             coeff_log_σᵤ² = (begw:endw),
             coeff_log_σᵥ² = (begv:endv))

  #* retrieve variable names for making tables
  xnames  = names(xvar)
# znames  = names(zvar)
  qnames  = names(qvar)
  wnames  = names(wvar)
  vnames  = names(vvar)
  varlist = vcat(" ", xnames, # znames,
                      qnames,
                      wnames, vnames)

  #* Converting the dataframe to matrix in order to do computation

  yvar = convert(Array{Float64}, Matrix(yvar))
  xvar = convert(Array{Float64}, Matrix(xvar))
  # zvar = convert(Array{Float64}, Matrix(zvar))
  qvar = convert(Array{Float64}, Matrix(qvar))
  wvar = convert(Array{Float64}, Matrix(wvar))
  vvar = convert(Array{Float64}, Matrix(vvar))

  tvar = convert(Array{Float64}, Matrix(tvar))
  ivar = convert(Array{Float64}, Matrix(ivar))

  zvar = ()

  #* various functions can and cannot contain a constant, check! ---- *#
# checkConst(xvar, :frontier, @requireConst(0))
# checkConst(qvar, :hscale,   @requireConst(0))
  checkConst(qvar, :σₐ²,      @requireConst(1))
  checkConst(wvar, :σᵤ²,      @requireConst(1))
  checkConst(vvar, :σᵥ²,      @requireConst(1))

  #* panel info and within transformation

  rowIDT = get_rowIDT(vec(ivar))   # (Nx2): col_1 is panel's row info; col_2 is panel's number of periods

  return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar, vvar,    rowIDT, varlist

end

#? --------- panel TRE (true random effect), truncated-normal ----------------

function getvar(::Type{PTRET}, dat::DataFrame)

  yvar = dat[:, _dicM[:depvar]]   # still a DataFrame
  xvar = dat[:, _dicM[:frontier]]
  zvar = dat[:, _dicM[:μ]]
  qvar = dat[:, _dicM[:σₐ²]]
  wvar = dat[:, _dicM[:σᵤ²]]
  vvar = dat[:, _dicM[:σᵥ²]]

  tvar = dat[:, _dicM[:timevar]]
  ivar = dat[:, _dicM[:idvar]] # need the [1] otherwise causes problems down the road

    #* --- model info printout ---------

    modelinfo1 = "true random effect model of Greene (2005 JoE), normal and truncated-normal"

    modelinfo2 = begin
     """
     * In the case of type(cost), "- uᵢₜ" below is changed to "+ uᵢₜ".

     $(_dicM[:depvar][1]) = frontier(αᵢ + $(_dicM[:frontier])) + vᵢₜ - uᵢₜ,

     where
            αᵢ ∼ N(0, σₐ²)
                 σₐ² = exp(log_σₐ²)
                     = exp($(_dicM[:σₐ²]));
           vᵢₜ ∼ N(0, σᵥ²),
                 σᵥ² = exp(log_σᵥ²)
                     = exp($(_dicM[:σᵥ²]));
           uᵢₜ ∼ N⁺(μ, σᵤ²),
                   μ = $(_dicM[:μ]),
                 σᵤ² = exp(log_σᵤ²)
                     = exp($(_dicM[:σᵤ²]));
     """
    end

  #* --- retrieve and generate important parameters -----

  #*   number of obs and number of variables
  nofx = nofz = nofq = nofw = nofv = 0  # to make a complete list

  nofobs  = nrow(dat)
  nofx    = size(xvar,2)  # nofx: number of x vars
  nofz    = size(zvar,2)
  nofq    = size(qvar,2)
  nofw    = size(wvar,2)
  nofv    = size(vvar,2)
  nofpara = nofx + nofz + nofq + nofw + nofv

  nofvar = (nofobs=nofobs, nofx=nofx, nofz=nofz, nofq=nofq,
            nofw=nofw, nofv=nofv, nofpara=nofpara, nofmarg = nofz+nofq+nofw)

  #* positions of the variables/parameters
  begx=endx=begz=endz=begq=endq=begw=endw=begv=endv=0

  begx = 1
  endx = nofx
  begz = endx + 1
  endz = begz + nofz -1
  begq = endz + 1
  endq = begq + nofq -1
  begw = endq + 1
  endw = begw + nofw -1
  begv = endw + 1
  endv = nofpara

  posvec = (begx=begx, endx=endx, begz=begz, endz=endz,
            begq=begq, endq=endq, begw=begw, endw=endw,
            begv=begv, endv=endv)

  #* create equation names and mark positions for making tables
  eqvec = (frontier = begx + 1,
                  μ = begz + 1,
            log_σₐ² = begq + 1,
            log_σᵤ² = begw + 1,
            log_σᵥ² = begv + 1)

  #* create equation names and mark positions
  eqvec2 = (coeff_frontier = (begx:endx),
             coeff_μ       = (begz:endz),
             coeff_log_σₐ² = (begq:endq),
             coeff_log_σᵤ² = (begw:endw),
             coeff_log_σᵥ² = (begv:endv))

  #* retrieve variable names for making tables
  xnames  = names(xvar)
  znames  = names(zvar)
  qnames  = names(qvar)
  wnames  = names(wvar)
  vnames  = names(vvar)
  varlist = vcat(" ", xnames,  znames,
                      qnames,
                      wnames, vnames)

  #* Converting the dataframe to matrix in order to do computation
  yvar = convert(Array{Float64}, Matrix(yvar))
  xvar = convert(Array{Float64}, Matrix(xvar))
  zvar = convert(Array{Float64}, Matrix(zvar))
  qvar = convert(Array{Float64}, Matrix(qvar))
  wvar = convert(Array{Float64}, Matrix(wvar))
  vvar = convert(Array{Float64}, Matrix(vvar))

  tvar = convert(Array{Float64}, Matrix(tvar))
  ivar = convert(Array{Float64}, Matrix(ivar))

  #* various functions can and cannot contain a constant, check! ---- *#
# checkConst(xvar, :frontier, @requireConst(0))
# checkConst(qvar, :hscale,   @requireConst(0))
  checkConst(qvar, :σₐ²,      @requireConst(1))
  checkConst(wvar, :σᵤ²,      @requireConst(1))
  checkConst(vvar, :σᵥ²,      @requireConst(1))

  #* panel info and within transformation

  rowIDT = get_rowIDT(vec(ivar))   # (Nx2): col_1 is panel's row info; col_2 is panel's number of periods

  return modelinfo1, modelinfo2, posvec, nofvar, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar, vvar,         rowIDT, varlist

end


text:
# push!(LOAD_PATH,"./src/")
using Documenter
using model1d1d
makedocs(sitename="DOC - 1d1d Scrap preheating model library")


text:
## solving by the L&S method
begin
    using JLD2, FileIO
    using LinearAlgebra
    using Interpolations
    using Roots
    using Plots
    using Optim
    pyplot()
    #CPROB=[1.0/3.0 2.0/3.0;1.0/4.0 3.0/4.0];
    #CPROB=[0.35 0.65;0.22 0.78];
    CPROB=[1.0/3.0 2.0/3.0;1.0/4.0 3.0/4.0];

    nos=size(CPROB)[1];
    [sum(CPROB[i,:]) for i=1:nos];

    mM=CPROB'-I;
    mM[end,:]=[1 for i=1:nos];
    rhs=[floor(i/nos) for i=1:nos];
    BSP=mM\rhs; #the list of steady-state probabilities

    hours = [0.0,1.0];
    wage=1.0;
    inc=wage*hours;
    ι=inc'*BSP; #payoff from the bond (= the aggregate output in the economy)

    #ρ=0.05
    ρ=-log(0.96)
    R=3.0;
end

function U(ξ::Float64,R::Float64)
    if ξ <= 0.
        -9e10
        else ((1/ξ^(R-1.0))/(1.0 - R))
    end
end

function II(aa::Int64,k::Int64,a::Int64,tbl2::Array{Int64,2})
    if aa==tbl2[k,a]
        return 1
    else
        return 0
    end
end

## the step on the grid created above

function LS_method(intr::Float64,accu_pow1::Int64,accu_pow2::Int64,nos::Int64,hours::Array{Float64,1},wage::Float64,β::Float64,grda::Array{Float64,1},Lg::Int64,CPROB::Array{Float64,2},R::Float64)
    local TBL,TBL1,TBL2,tbl,tbl1,tbl2,Λ0,Λ1
    currentU=[U((1+intr)*grda[i]+wage*hours[k]-grda[j],R) for k=1:nos, i=1:Lg, j=1:Lg];
    nextV=[U((1+intr)*grda[j]+wage*hours[k],R) for k=1:nos, j=1:Lg];

    TBL=[findmax(currentU[k,i,:]+β*(CPROB[k,:]'*nextV)') for k=1:nos, i=1:Lg];
    TBL1=first.(TBL);TBL2=last.(TBL);

    tbl=[findmax(currentU[k,i,:]+β*(CPROB[k,:]'*TBL1)') for k=1:nos, i=1:Lg];
    tbl1=first.(tbl);tbl2=last.(tbl);

    while maximum(abs.(TBL2-tbl2))>0||maximum(abs.(TBL1-tbl1))>1/(10^accu_pow1)
        TBL1=tbl1;
        TBL2=tbl2;
        tbl=[findmax(currentU[k,i,:]+β*(CPROB[k,:]'*TBL1)') for k=1:nos, i=1:Lg];
        tbl1=first.(tbl);tbl2=last.(tbl);
    end

    Λ0=[Float64(1/(nos*Lg)) for k=1:nos, i=1:Lg];

    Λ1=[sum([Λ0[k,i]*CPROB[k,kk]*II(ii,k,i,tbl2) for k=1:nos, i=1:Lg]) for kk=1:nos, ii=1:Lg];

    Λ0=Λ1;
    Λ1=[sum([Λ0[k,i]*CPROB[k,kk]*II(ii,k,i,tbl2) for k=1:nos, i=1:Lg]) for kk=1:nos, ii=1:Lg];
    while maximum(abs.(Λ1-Λ0))>1/10^accu_pow2
        Λ0=Λ1;
        Λ1=[sum([Λ0[k,i]*CPROB[k,kk]*II(ii,k,i,tbl2) for k=1:nos, i=1:Lg]) for kk=1:nos, ii=1:Lg];
    end
    return Λ1,tbl2
end


text:
export complex_embeddings, real_embeddings, evaluation_function

@doc Markdown.doc"""
    complex_embeddings(K::NumField; conjugates::Bool = true) -> Vector{NumFieldEmb}

Return the complex embeddings of $K$. If `conjugates` is `false`, only one
imaginary embedding per conjugated pairs is returned.

# Examples

```jldoctest
julia> K, a = quadratic_field(-3);

julia> complex_embeddings(K)
2-element Vector{Hecke.NumFieldEmbNfAbs}:
 Embedding corresponding to ≈ 0.00 + 1.73 * i
 Embedding corresponding to ≈ 0.00 - 1.73 * i

julia> complex_embeddings(K, conjugates = false)
1-element Vector{Hecke.NumFieldEmbNfAbs}:
 Embedding corresponding to ≈ 0.00 + 1.73 * i
```
"""
complex_embeddings(K::NumField)

@doc Markdown.doc"""
    real_embeddings(K::NumField) -> Vector{NumFieldEmb}

Return the real embeddings of $K$.

# Examples

```jldoctst
julia> K, a = quadratic_field(3);

julia> real_embeddings(K)
2-element Vector{Hecke.NumFieldEmbNfAbs}:
 Embedding corresponding to ≈ -1.73
 Embedding corresponding to ≈ 1.73
```
"""
function real_embeddings(K::NumField)
  res = get_attribute!(K, :real_embeddings) do
    r, s = signature(K)
    return complex_embeddings(K)[1:r]
  end::Vector{embedding_type(K)}
end

@doc Markdown.doc"""
    number_field(f::NumFieldEmb) -> NumField

Return the corresponding number field of the embedding $f$.

# Examples

```jldoctet
julia> K, a = quadratic_field(-3); e = complex_embeddings(K)[1];

julia> number_field(e)
Imaginary quadratic field defined by x^2 + 3
```
"""
number_field(::NumFieldEmb)

@doc Markdown.doc"""
    isreal(f::NumFieldEmb) -> Bool

Return `true` if the embedding is real.

# Examples

```jldoctest
julia> K, a = quadratic_field(3); e = complex_embeddings(K)[1];

julia> isreal(e)
true
```
"""
isreal(f::NumFieldEmb)

@doc Markdown.doc"""
    is_imaginary(f::NumFieldEmb) -> Bool

Returns `true` if the embedding is imaginary, that is, not real.

# Examples

```jldoctest
julia> K, a = quadratic_field(-3); e = complex_embeddings(K)[1];

julia> is_imaginary(e)
true
```
"""
is_imaginary(f::NumFieldEmb) = !isreal(f)

@doc Markdown.doc"""
    conj(f::NumFieldEmb) -> NumFieldEmb

Returns the conjugate embedding of `f`.

# Examples

```jldoctest
julia> K, a = quadratic_field(-3); e = complex_embeddings(K);

julia> conj(e[1]) == e[2]
true
```
"""
conj(f::NumFieldEmb)

# Restriction

@doc Markdown.doc"""
    restrict(f::NumFieldEmb, K::NumField)

Given an embedding $f$ of a number field $L$ and a number field $K$ appearing
as a base field of $L$, return the restriction of $f$ to $L$.

# Examples

```jldoctest
julia> K, a = quadratic_field(3);

julia> L, b = NumberField(polynomial(K, [1, 0, 1]), "b");

julia> e = complex_embeddings(L);

julia> restrict(e[1], K)
Embedding of
Real quadratic field defined by x^2 - 3
corresponding to ≈ -1.73
```
"""
restrict(f::NumFieldEmb, K::NumField)

@doc Markdown.doc"""
    restrict(f::NumFieldEmb, g::NumFieldMor)

Given an embedding $f$ of a number field $L$ and a morphism $g \colon K \to L$,
return the embedding $g \circ f$ of $K$.

This is the same as `g * f`.

# Examples

```jldoctest
julia> K, a = CyclotomicField(5, "a");

julia> k, ktoK = Hecke.subfield(K, [a + inv(a)]);

julia> e = complex_embeddings(K);

julia> restrict(e[1], ktoK)
Embedding of
Number field over Rational Field with defining polynomial x^2 + x - 1
corresponding to ≈ 0.62
```
"""
restrict(f::NumFieldEmb, K::NumFieldMor)

function Base.:(*)(f::NumFieldMor, e::NumFieldEmb)
  return restrict(e, f)
end

################################################################################
#
#  Extension
#
################################################################################

@doc Markdown.doc"""
    extend(e::NumFieldEmb, f::NumFieldMor)

Given an embedding $e$ of $k$ and a morphism $f \colon k \to K$, determine
all embedings of $K$ which restrict to $e$ along $f$.

# Example

```jldoctest
julia> K, a = CyclotomicField(5, "a");

julia> k, ktoK = Hecke.subfield(K, [a + inv(a)]);

julia> e = complex_embeddings(k)[1];

julia> extend(e, ktoK)
2-element Vector{Hecke.NumFieldEmbNfAbs}:
 Embedding corresponding to ≈ 0.31 + 0.95 * i
 Embedding corresponding to ≈ 0.31 - 0.95 * i
```
"""
function extend(e::NumFieldEmb, f::NumFieldMor)
  @req number_field(e) === domain(f) "Number field of embedding must be domain"
  emb = complex_embeddings(codomain(f))
  res = eltype(emb)[ E for E in emb if f * E == e ]
  @assert length(res) == div(absolute_degree(codomain(f)), absolute_degree(domain(f)))
  return res
end

################################################################################
#
#  Evaluation function
#
################################################################################

@doc Markdown.doc"""
    evaluation_function(e::NumFieldEmb, prec::Int) -> Function

Return the anonymous function `x -> e(x, prec)`.

# Examples

```jldoctest
julia> K, a = quadratic_field(-3);

julia> e = complex_embeddings(K)[1];

julia> fn = evaluation_function(e, 64);

julia> fn(a)
[+/- 3.99e-77] + [1.73205080756887729353 +/- 5.41e-21]*im
```
"""
function evaluation_function(e::NumFieldEmb, prec::Int = 32)
  return x -> e(x, prec)
end