########################################################
####                sfmodel_spec()                  ####
########################################################

"""
sfmodel_spec(<keyword arguments>)

Provide specifications of the stochastic frontier model, including the type of model
and names of variables or matrix used in estimating the model. Two ways to
specify: Method 1: Use DataFrame as input, and Method 2: use matrix as input.

# Method 1 (DataFrame input)
Variables come from a DataFrame, and column names of the Dataframe are used in
the variable input. With this method,
equations are identified by macros but not functions (e.g., `@depvar()` but
not `depvar()`).

## Arguments of Method 1
- `sfdist(::Vararg)`: the distribution assumption of the one-sided stochastic
  variable (aka inefficiency term) of the model;
  possible choices include `truncated` (or `trun`, `t`), `half` (or `h`),
  `exponential` (or `expo`, `e`), and `trun_scale` (or `trun_scaling`, `ts`).
- `sftype(::Vararg)`: whether the model is a `production` (or `prod`) frontier
  or a `cost` frontier.
- `sfpanel(::Vararg)`: the type of panel model. Choices include `TFE_WH2010`
  (true fixed effect model of Wang and Ho 2010 JE), `TFE_CSW2014` (true fixed
  model of Chen, Schmidt, and Wang 2014 JE),  `TRE` (true random effect model
  of Greene 2004), `TimeDecay` (time decay model of Battese and Coelli 1992).
- `@depvar(::Vararg)`: the dependent variable from a DataFrame.
- `@frontier(::Vararg)`: a list of variables, separated by commas, in the frontier function.
- `@μ(::Vararg)` or `@mu(::Vararg)`: a list of variable, separated by comma,
  in the linear function of μ. (`sftype(trun)` only).
- `@σᵥ²(::Vararg)` or `@sigma_v_2(::Vararg)`: a list of variable, separated by comma, in the σᵥ²
  equation.
- `@σᵤ²(::Vararg)` or `@sigma_u_2(::Vararg)`: a list of variable, separated by comma, in the σᵤ²
  equation.
- `@σₐ²(::Vararg)` or `@sigma_a_2(::Vararg)`: a list of variable, separated by comma, in the σₐ²
  equation. `sfpanel(TRE)` only.
- `@gamma(::Vararg)`: a list of variables, separated by commas, in the gamma
  equation. `sfpanel(TimeDecay)` only.
- `@timevar(::Vararg)`: the variable containing the time period information.
  Panel data model only.
- `@idvar(::Vararg)`: the variable identifying each individual. Panel data
  model only.
- message::Bool: Whether printing (=true) or not (=false, the default) the
  confirmation message "A dictionary from sfmodel_spec() is generated."
  on the screen after `sfmodel_spec()` is successfully executed.

# Method 2 (matrix/vector input)
Data of the variables are
provided by individual matrices or vectors, and names of the mat/vec are used
in the equations. With this method, equations are identified by functions but
not macros (e.g., `depvar()` but not `@depvar()`). Note that if, for instance,
the name of `depvar` or `σᵤ²` has been used elsewhere in the program, using
these names to read in mat/vec will cause name conflict (`MethodError: objects
of type ... are not callable`). The workaround is to fully qualify the function
names, e.g., `SFrontiers.depvar`, `SFrontiers.σᵤ²`, etc. Or, use the alias (if
available), e.g., `sigma_u_2` instead of `σᵤ²`.

## Arguments of Method 2
- `sfdist(::Vararg)`: the distribution assumption on the inefficiency term;
  possible choices include `truncated` (or `trun`, `t`), `half` (or `h`),
  `exponential` (or `expo`, `e`), and `trun_scale` (or `trun_scaling`, `ts`).
- `sftype(::Vararg)`: whether the model is a `production` (or `prod`) frontier
  or a `cost` frontier.
- `sfpanel(::Vararg)`: the type of panel model. Choices include `TFE_WH2010`
  (true fixed effect model of Wang and Ho 2010 JE), `TFE_CSW2014` (true fixed
  model of Chen, Schmidt, and Wang 2014 JE),  `TRE` (true random effect model
  of Greene 2004), `TimeDecay` (time decay model of Battese and Coelli 1992).
- `depvar(::Matrix)`: Matrix or vector of the dependent variable.
- `frontier(::Matrix)`: matrix or vector for frontier function.
- `μ(::Matrix)` or `mu(::Matrix)`: matrix or vector for the (linear) μ equation (`trun` type only).
- `σᵤ²(::Matrix)` or `sigma_u_2(::Matrix)`: matrix or vector for the σᵤ² equation.
- `σₐ²(::Matrix)` or `sigma_a_2(::Matrix)`: matrix or vector for the σₐ² equation.
- message::Bool: Whether printing (=true) or not (=false, the default) the
  confirmation message "A dictionary from sfmodel_spec() is generated."
  on the screen after `sfmodel_spec()` is successfully executed.

# Examples
```julia-repl
sfmodel_spec(sftype(prod), sfdist(trun),
             @depvar(output),
             @frontier(land, , labor, bull, year, _cons),
             @μ(age, school, year, _cons),
             @σᵤ²(age, school, year, _cons),
             @σᵥ²(_cons));

sfmodel_spec(sfpanel(TRE), sftype(prod), sfdist(half),
             @timevar(yr), @idvar(id),
             @depvar(y),
             @frontier(x1, x2, _cons),
             @σₐ²(_cons),
             @σᵤ²(_cons),
             @σᵥ²(_cons),
             message = false);
```
"""
function sfmodel_spec(arg::Vararg; message::Bool=false)

    global _dicM
           _dicM = Dict{Symbol, Any}()  # nullify and initiate new dictionaries when a new model is specified

         #* -- creates default values ---

        for k in (:panel, :timevar, :idvar, :dist, :type, :depvar, :frontier, :μ, :hscale, :gamma, :σᵤ², :σᵥ², :σₐ², :hasDF, :transfer, :misc) # take out :data, :η, :λ, :τ, in this revision
            _dicM[k] = nothing
        end

        #* -- replace the defaults with user's values ---

        for d in :($(arg))
            _dicM[d[1]] = d[2]
        end

        #* ==== Method 2, matrix input (such as in simulations), create a DataFrame

           _dicM[:hasDF]    = true
           _dicM[:transfer] = false

        if typeof(_dicM[:depvar]) != Array{Symbol,1} # not a DataFrame

           _dicM[:hasDF] = false

           isa(_dicM[:depvar][1], Vector) || isa(_dicM[:depvar][1], Matrix) || throw("
           `depvar()` has to be a Vector or Matrix (e.g., Array{Float64, 1} or Array{Float64, 2}).
           Check with `isa(your_thing, Matrix)` or `isa(your_thing, Vector)`.
           Try `convert()`, `reshape()`, `Matrix()`, or something similar.")

           comDF = _dicM[:depvar][1]  # create the first data column of comDF
           varname = [:depvar]
           _dicM[:depvar] = [:depvar]

           for k in (:timevar, :idvar, :frontier, :μ, :hscale, :gamma, :σₐ², :σᵤ², :σᵥ²)
               if _dicM[k] !== nothing # if not nothing, must be Array
                  isa(_dicM[k], Vector) || isa(_dicM[k][1], Vector) || isa(_dicM[k][1], Matrix) || throw("
                     `k` has to be a Vector or Matrix (e.g., Array{Float64, 1} or Array{Float64, 2}).
                     Check with `isa(your_thing, Matrix)` or `isa(your_thing, Vector)`.
                     To convert, try `convert()`, `reshape()`, `Matrix()`, or something similar.")

                  (isa(_dicM[k], Vector)  && length(_dicM[k][1]) == 1) ?  _dicM[k] = [_dicM[k]] : nothing # ugly fix for pure vector input

                  @views comDF = hcat(comDF, _dicM[k][1]) # combine the data
                  aa = Symbol[]
                  for i in 1:size(_dicM[k][1], 2)
                      push!(aa, Symbol(String(k)*"_var$(i)")) # create name for the data
                  end
                  varname = vcat(varname, aa) # combine the dataname
                  _dicM[k] = aa
               end
           end # for k in (...)

           comDF = DataFrame(comDF, varname)
           _dicM[:sdf] = comDF

        end # if typeof(...)

        #* -- check the model identifier and the model type ---

        (_dicM[:dist] !== nothing) || throw("You need to specify dist().")
        (_dicM[:type] !== nothing) || throw("You need to specify type().")

        #* --- get the model identifier -------

        s = uppercase(String(_dicM[:dist][1])[1:1])

        global tagD
        if _dicM[:panel] === nothing  # not panel
            if (s == "T")
                tagD = Dict{Symbol, Type{Trun}}()
                tagD[:modelid] = Trun
            elseif (s == "H")
                tagD = Dict{Symbol, Type{Half}}()
                tagD[:modelid] = Half
            elseif (s == "E")
                tagD = Dict{Symbol, Type{Expo}}()
                tagD[:modelid] = Expo
            elseif (s == "S")
                tagD = Dict{Symbol, Type{Trun_Scale}}()
                tagD[:modelid] = Trun_Scale
            end
        elseif (_dicM[:panel] == [:TFE_WH2010]) && (s=="T")
            tagD = Dict{Symbol, Type{PFEWHT}}()
            tagD[:modelid] = PFEWHT
        elseif (_dicM[:panel] == [:TFE_WH2010]) && (s == "H")  # panel and dist=half
            tagD = Dict{Symbol, Type{PFEWHH}}()
            tagD[:modelid] = PFEWHH
        elseif (_dicM[:panel] == [:TFE_WH2010])
            throw("The panel TFE_WH2010 model can only have `sfdist(trun)` or `sfdist(half)`.")
        elseif (_dicM[:panel] == [:TFE_CSW2014]) && (s=="H")
            tagD = Dict{Symbol, Type{PFECSWH}}()
            tagD[:modelid] = PFECSWH
        elseif (_dicM[:panel] == [:TFE_CSW2014])
            throw("The panel TFE_CSW2014 model can only have `sfdist(half)`.")
        elseif (_dicM[:panel] == [:TRE]) && (s=="H")
            tagD = Dict{Symbol, Type{PTREH}}()
            tagD[:modelid] = PTREH
        elseif (_dicM[:panel] == [:TRE]) && (s=="T")
            tagD = Dict{Symbol, Type{PTRET}}()
            tagD[:modelid] = PTRET
        elseif (_dicM[:panel] == [:TRE])
            throw("The panel TRE model can only have `sfdist(half)` or `sfdist(trun)`.")
        elseif (_dicM[:panel] == [:TimeDecay]) && (s == "T")
            tagD = Dict{Symbol, Type{PanDecay}}()
            tagD[:modelid] = PanDecay
        elseif (_dicM[:panel] == [:TimeDecay])
            throw("The panel time-decay model can only have `sfdist(trun)`.")
        elseif (_dicM[:panel] == [:Kumbhakar1990]) && (s == "T")
            tagD = Dict{Symbol, Type{PanKumb90}}()
            tagD[:modelid] = PanKumb90
        elseif (_dicM[:panel] == [:Kumbhakar1990])
            throw("The panel Kumbhakar 1990 model can only have `sfdist(trun)`.")
        else
            throw("The `sfpanel()` and/or `sfdist()` are not specified correctly.")
        end

        #* ---- check if the model has the correct syntax ---

        SFrontiers.checksyn(tagD[:modelid])

        #* ----- make return -----------
        if message
          printstyled("A dictionary from sfmodel_spec() is generated.\n"; color = :green)
        end
        return _dicM # for debugging purpose

end  # end of sfmodel_spec()

########################################################
####                sfmodel_init()                  ####
########################################################
"""
    sfmodel_init(<keyword arguments>)

Provide initial values for the stochastic frontier model estimation. The
values could be a vector or scalars. It creates a global dictionary `_dicINI`. Optional.

# Arguments
- `all_init(::Union{Vector, Real})`: initial values of all the parameters in the model
- `frontier(::Union{Vector, Real})`: initial values of parameters in
  the `frontier()` function
- `μ(::Union{Vector, Real})` or `mu(::Union{Vector, Real})`: initial values of
  parameters in the `μ` function
- `hscale(::Union{Vector, Real})`: initial values of parameters in the `hscale()` function
- `gamma(::Union{Vector, Real})`: initial values of parameters in the `gamma()` function
- `σᵤ²(::Union{Vector, Real})` or `sigma_u_2(::Union{Vector, Real})`: initial values of parameters in the
   `σᵤ²` function
- `σᵥ²(::Union{Vector, Real})` or `sigma_v_2(::Union{Vector, Real})`: initial values of parameters in the
   `σᵥ²` function
- `σₐ²(::Union{Vector, Real})` or `sigma_a_2(::Union{Vector, Real})`: initial
  values of parameters in the `σₐ²` function
- message::Bool: Whether printing (=true) or not (=false, the default) the
  confirmation message "A dictionary from sfmodel_init() is generated."
  on the screen after `sfmodel_init()` is successfully executed.

# Remarks
- Equations do not have to follow specific orders.
- `sfmodel_init(...)` is optional but is highly recommended. If it is not
  specified or is specified as an empty set, default values are used.
- It is not necessary to specify a complete set of equations. A partial list
  or even empty lists are acceptable. Default values will be substituted for the
  missing equations.
- The generated `_dicINI` is inheritable in the sense that an exiting
  `_dicINI` (from the previous run of the same or a different model, for
  example) will be used if the current model does not have its own
  `sfmodel_init(...)`. This design has advantages in a simulations study where
  `sfmodel_init(...)` needs to be specified only once.

# Examples
```julia-repl
b_ini = ones(2)*0.2
sfmodel_init( # frontier(bb),             # may skip and use default
             μ(b_ini),                    # may use a vector
             σᵤ²(-0.1, -0.1),
             σᵥ²(-0.1) )

sfmodel_init(all_init(0.1, 0.2, 0.5, 0.0, -0.1, -0.1, -0.1),
             message = false)
```
"""
function sfmodel_init(arg::Vararg; message::Bool =false) # create a dictionary of inital vectors

   global _dicINI
          _dicINI = Dict{Symbol, Any}()

    for d in :($(arg))
        _dicINI[d[1]] = d[2]
    end

    #* If has the key, creates the alias key with the same value.

    !(haskey(_dicINI, :μ))      || (_dicINI[:eqz] = _dicINI[:μ])
    !(haskey(_dicINI, :hscale)) || (_dicINI[:eqq] = _dicINI[:hscale])
    !(haskey(_dicINI, :gamma))  || (_dicINI[:eqq] = _dicINI[:gamma])
    !(haskey(_dicINI, :σₐ²))    || (_dicINI[:eqw] = _dicINI[:σₐ²])
    !(haskey(_dicINI, :σᵤ²))    || (_dicINI[:eqw] = _dicINI[:σᵤ²])
    !(haskey(_dicINI, :σᵥ²))    || (_dicINI[:eqv] = _dicINI[:σᵥ²])

    if message
      printstyled("A dictionary from sfmodel_init() is generated.\n"; color = :green)
    end
    return _dicINI # for debugging purpose
end

########################################################
####                sfmodel_opt()                   ####
########################################################
"""
    sfmodel_opt(<keyword arguments>)

Provide options to the optimization algorithms for the maiximum likelihood
estimation. It creates a global dictionary `_dicOPT`. Optional. The `Optim`
package is used for the optimization, and a subset of
`Optim`'s keywords are directly accessible from this API.

# Arguments
- `warmstart_solver(algorithm)`: The algorithm used in the first-stage ("warmstart")
  optimization process, which serves the purpose of improving upon the initial
  values for the second-stage ("main") estimation. The default is
  `NelderMead()`. Others include `SimulatedAnnealing()`, `SAMIN()`, `ParticleSwarm()`,
  `ConjugateGradient()`, `GradientDescent()`, `BFGS()`, `LBFGS()`,
  `Newton()`, `NewtonTrustRegion()`, and `IPNewton()`. See
  http://julianlsolvers.github.io/Optim.jl/stable/ for details.
  Non-gradient based algorithms are recommended for the warmstart solver.
- `warmstart_maxIT(::Int64)`: The iteration limit for the warmstart. Default
  is 100.
- `main_solver(algorithm)`: The algorithm used in the main opimization process.
  The default is `Newton()`. Others include `SimulatedAnnealing()`, `SAMIN()`, `ParticleSwarm()`,
  `ConjugateGradient()`, `GradientDescent()`, `BFGS()`, `LBFGS()`,
  `NewtonTrustRegion()`, and `IPNewton()`. See
  http://julianlsolvers.github.io/Optim.jl/stable/ for details.
- `main_maxIT(::Int64)`: The iteration limit for the main estimation. Default
  is 2000.
- `tolerance(::Float64)`: The convergence criterion ("tolerance") based on the
  absolute value of gradients. Default is 1.0e-8. For non-gradient algorithms,
  it controls the main convergence tolerance, which is solver specific.
  See `Optim`'s `g_tol` option for more information.
- `verbose(::Bool)`: Print on screen (`true`, the default) the information of
  the model and the optimization results.
- `banner(::Bool)`: Print on screen (`true`, the default) a banner to serve as
  a visual indicator of the start of the estimation.
- `ineff_index(::Bool)`: Whether to compute the Jondrow et al. (1982)
  inefficiency index and the Battese and Coelli (1988) efficiency index. The
  defauis `true`.
- `marginal(::Bool)`: Whether to compute the marginal effects of the exogenous
  determinants of inefficiency (if any).
- `table_format()`: The format to print the coefficient tables on the screen:
  `text` (default), `html`, or `latex`. A wrapper of `PrettyTables.jl`'s
  `backend` option.
- message::Bool: Whether printing (=true) or not (=false, the default) the
  confirmation message "A dictionary from sfmodel_opt() is generated."
  on the screen after `sfmodel_opt()` is successfully executed.

# Remarks
- `sfmodel_opt(...)` is optional. It can be omitted entirely, or specifying
  only a partial list of the keywords.
- If any of the keywords are missing, default values are used.
- If warmstart is not needed, you need to give empty keyword values to
  warmstart related keys. E.g., either `warmstart_solver()` or
  `warmstart_maxIT()`, or both. Omitting the keyword entirely (i.e., not
  writing down `warmstart_solver` or `warmstart_maxIT`) will not skip the
  warmstart, but will reinstate the default.
- Users do not need to provide gradient or Hessian functions even if
  gradient-based optimization algorithms are used. The package uses automatic
  differentiation (https://en.wikipedia.org/wiki/Automatic_differentiation) to
  compute the derivatives. It is not numerical finite differentiation. It is
  fast and as accurate as the symbolic differentiation.
- The `_dicOPT` is inheritable in the sense that an exiting `_dicOPT` (from
  the previous run of the same or a different model, for example) will be used
  if the current model does not have its own `sfmodel_opt(...)`. This design
  has advantages in simulation studies where `sfmodel_opt(...)` needs to be
  specified only once.

# Examples
```julia-repl
sfmodel_opt(warmstart_solver(NelderMead()),
            warmstart_maxIT(200),
            main_solver(Newton()),
            main_maxIT(2000),
            tolerance(1e-8),
            message = false)
```
"""
function sfmodel_opt(arg::Vararg; message::Bool=false) # create a dictionary of maximization options

    global _dicOPT
           _dicOPT = Dict{Symbol, Any}()

    #* -- creates the default ---

    _dicOPT[:warmstart_solver] = :(NelderMead())
    _dicOPT[:warmstart_maxIT]  =  100
    _dicOPT[:main_solver]      = :(Newton())
    _dicOPT[:main_maxIT]       =  2000
    _dicOPT[:tolerance]        =  1.0e-8
    _dicOPT[:verbose]          =  true
    _dicOPT[:banner]           =  true
    _dicOPT[:ineff_index]      =  true
    _dicOPT[:marginal]         =  true
    _dicOPT[:table_format]     = :(text)

    #* -- replace the defaults with the user's value ---

    for d in :($(arg))
        _dicOPT[d[1]] = d[2]
    end

    #* ---- error checking --

    if (_dicOPT[:main_solver] === nothing) || (_dicOPT[:main_maxIT] === nothing) || (_dicOPT[:tolerance] === nothing)
         throw("You cannot give empty keyword values to `main_solver()`, `main_maxIT()`, or `tolerance()`. If you want to use the default, you may do so by dropping (not emptying keyword values) the keywords.")
    end

    if message
      printstyled("A dictionary from sfmodel_opt() is generated.\n"; color = :green)
    end
    return _dicOPT # for debugging purpose

end    # end of sfmodel_opt()

########################################################
###                 sfmodel_fit()                   ####
########################################################
"""
    sfmodel_fit(<keyword arguments>)

Maximum likelihood estimation of the stochastic frontier model specified
in `sfmodel_spec(...)`. Estimate the model parameters, calculate Jondrow et al.
(1982) inefficiency index and Battese and Coelli (1988) efficiency index,
compute marginal effects of inefficiency determinants (if any).
Return a dictionary with results.

# Arguments
- `useData(::DataFrame)`: The DataFrame used with the Method 1 of
  `sfmodel_spec(...)`. If use Method 2 of `sfmodel_spec(...)` (viz., data
  is supplied by individual matrices), do not need this keyword argument.

# Remarks
- Use `Optim.jl` to carry out the estimation.
- Users do not need to provide gradient or Hessian functions even if
  gradient-based optimization algorithms are used. The package uses automatic
  differentiation (https://en.wikipedia.org/wiki/Automatic_differentiation) to
  compute the derivatives. AD is not numerical finite differentiation. AD is
  fast and as accurate as the symbolic differentiation.

# Examples
```julia-repl
sfmodel_fit(useData(df))    # Method 1
sfmodel_fit()               # Method 2
```
"""
function sfmodel_fit()
     # For Method 2 of `sfmodel_spec()`.

    !(_dicM[:hasDF]) || throw("Need to specify DataFrame in `sfmodel_fit()`.")

    _dicM[:transfer] = true
    sfmodel_fit(_dicM[:sdf])

end

function sfmodel_fit(sfdat::DataFrame) #, D1::Dict = _dicM, D2::Dict = _dicINI, D3::Dict = _dicOPT)

    (_dicM[:hasDF] || _dicM[:transfer])  || throw("You provided matrix in `sfmodel_spec()` so you cannot specify a DataFrame in `sfmodel_fit()`. Leave it blank.")

   #* for simulation, add a flag
   redflag::Bool = 0

  #* ###### Check if the OPT dictionary exists #####

      @isdefined(_dicINI) || sfmodel_init()  # if not exist, create one with default values
      @isdefined(_dicOPT) || sfmodel_opt()

     if _dicOPT[:banner]
        printstyled("\n###------------------------------------###\n"; color=:yellow)
        printstyled("###  Estimating SF models using Julia  ###\n"; color=:yellow)
        printstyled("###------------------------------------###\n\n"; color=:yellow)
      end

  #* ##### Get variables from dataset #######

     # pos: (begx, endx, begz, endz, ...); variables' positions in the parameter vector.
     # num: (nofobs, nofx, ..., nofpara); number of variables in each equation
     # eqvec: ("frontier"=2, "μ"=6,...); named tuple of equation names and equation position in the table
     # eqvec2: (xeq=(1,3), zeq=(4,5),...); named tuple of equation and parameter positions, for sfmodel_predict
     # varlist: ("x1", "x2",...); variable names for making table

     (minfo1, minfo2, pos, num, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar,
      vvar,         rowIDT, varlist) = SFrontiers.getvar(tagD[:modelid], sfdat)

  #* ### print preliminary information ########

    if _dicOPT[:verbose]

      printstyled("*********************************\n "; color=:cyan)
      printstyled("      Model Specification:\n"; color=:cyan);
      printstyled("*********************************\n"; color=:cyan)

      print("Model type: "); printstyled(minfo1; color=:yellow); println();println()
      printstyled(minfo2; color=:yellow); println()
    end

  #* ##### Get the type parameter #######

     _porc::Int64 = 1

     if (_dicM[:type] == [:cost])
         _porc = -1
     end

  #* ########## Process initial value dictionary  #####
     #* --- Get OLS results and other auxiliary values. --- #

     noffixed = 0  # number of fixed effect parameter
     if _dicM[:panel] == [:TFE_WH2010] || _dicM[:panel] == [:TFE_CSW2014]  # the fixed effects and differenced off
        noffixed = size(rowIDT,1)
     end

     β0     = xvar \ yvar;  # OLS estiamte, uses a pivoted QR factorization;
     resid  = yvar - xvar*β0
     sse    = sum((resid).^2)
     ssd    = sqrt(sse/(size(resid,1)-(num.nofx + noffixed ))) # sample standard deviation; σ² = (1/(N-K))* Σ ϵ^2
     ll_ols = sum(normlogpdf.(0, ssd, resid)) # ols log-likelihood
     sk_ols = sum((resid).^3) / ((ssd^3)*(size(resid,1))) # skewnewss of ols residuals

     #* --- Create the dictionary -----------

     if (:all_init in keys(_dicINI))
         sf_init = _dicINI[:all_init]
     else
         #*  Create ini vectors from user's values; if none, use the default.--- #
         b_ini  = get(_dicINI, :frontier, β0)
         d1_ini = get(_dicINI, :eqz, ones(num.nofz) * 0.1)
         t_ini  = get(_dicINI, :eqq, ones(num.nofq) * 0.1)
         d2_ini = get(_dicINI, :eqw, ones(num.nofw) * 0.1)
         g_ini  = get(_dicINI, :eqv, ones(num.nofv) * 0.1)

         #*  Make it Array{Float64,1}; otherwise Array{Float64,2}. ---#
         #*       Could also use sf_init[:,1]. *#
         sf_init = vcat(b_ini, d1_ini, t_ini, d2_ini, g_ini)
         sf_init = vec(sf_init)
     end # if :all_init

  #* ############ Misc.  ################
     # --- check if the number of initial values is correct
        (length(sf_init) == num.nofpara) ||  throw("The number of initial values does not match the number of parameters to be estimated. Make sure the number of init values in sfmodel_init() matches the number of variabls in sfmodel_spec().")

     # --- Make sure there is no numerical issue arising from int vs. Float64.
        sf_init = convert(Array{Float64,1}, sf_init)

  #* ############# process optimization dictionary  #######

         if (_dicOPT[:warmstart_solver] === nothing) || (_dicOPT[:warmstart_maxIT] === nothing)
             do_warmstart_search = 0
         else
             do_warmstart_search = 1
             sf_ini_algo  = eval(_dicOPT[:warmstart_solver])  # warmstart search algorithms
             sf_ini_maxit = _dicOPT[:warmstart_maxIT]         # warmstart search iter limit
         end

     # ---- main maximization algorithm -----
         sf_algo  = eval(_dicOPT[:main_solver])    # main algorithm
         sf_maxit = _dicOPT[:main_maxIT]
         sf_tol   = _dicOPT[:tolerance]
         sf_table = _dicOPT[:table_format]

  #* ########  Start the Estimation  ##########

    #* ----- Define the problem's Hessian -----#

     _Hessian = TwiceDifferentiable(rho -> SFrontiers.LL_T(tagD[:modelid],
                           yvar, xvar, zvar, qvar, wvar, vvar,
                           _porc, num.nofobs, pos, rho,
                                   rowIDT, _dicM[:misc]),
                     sf_init;
                    autodiff = :forward);

    #* ---- Make placeholders for dictionary recording purposes *#

    sf_init_1st_dic  = 0
    sf_init_2nd_dic  = 0
    sf_ini_algo_dic  = nothing
    sf_ini_maxit_dic = 0
    sf_total_iter    = 0

    _run = 1  # a counter; use the -if- instead of -for- to avoid using global variables

    if (do_warmstart_search == 1) && (_run == 1)

        if _dicOPT[:verbose]
            printstyled("The warmstart run...\n\n"; color = :green)
        end

        sf_init_1st_dic  = copy(sf_init) # for dict recording
        sf_ini_algo_dic  = sf_ini_algo
        sf_ini_maxit_dic = copy(sf_ini_maxit)

        # @time
               mfun = optimize(_Hessian,
                               sf_init,         # initial values
                               sf_ini_algo,
                               Optim.Options(g_tol = sf_tol,
                                             iterations  = sf_ini_maxit,
                                             store_trace = true,
                                             show_trace  = false))

        sf_total_iter += Optim.iterations(mfun) # for later use

        sf_init = Optim.minimizer(mfun)  # save as initials for the next run
        _run    = 2                      # modify the flag

        if _dicOPT[:verbose]
            println()
            print("$mfun \n")
            print("The warmstart results are:\n"); printstyled(Optim.minimizer(mfun); color=:yellow); println("\n")
        end

   end  # if  (do_warmstart_search == 1) && (_run == 1)

   if (do_warmstart_search == 0 ) || (_run == 2) # either no warmstart run and go straight here, or the 2nd run

       sf_init_2nd_dic = copy(sf_init) # for dict recording

       if _dicOPT[:verbose]
           println()
           printstyled("Starting the optimization run...\n\n" ; color = :green)
       end

       # @time
              mfun = optimize(_Hessian,
                              sf_init,       # initial values
                              sf_algo,       # different from search run
                              Optim.Options(g_tol = sf_tol,
                                            iterations  = sf_maxit, # different from search run
                                            store_trace = true,
                                            show_trace  = false))
       sf_total_iter += Optim.iterations(mfun)

       if _dicOPT[:verbose]
             println()
             print("$mfun \n")
             print("The resulting coefficient vector is:\n"); printstyled(Optim.minimizer(mfun); color=:yellow); println("\n")
       end

       if isnan(Optim.g_residual(mfun)) || (Optim.g_residual(mfun) > 0.1)
            redflag = 1
            printstyled("Note that the estimation may not have converged properly. The gradients are problematic (too large, > 0.1, or others).\n\n", color = :red)
       end

       if Optim.iteration_limit_reached(mfun)
             redflag = 1
             printstyled("Caution: The number of iterations reached the limit.\n\n"; color= :red)
       end

   end     # if (do_warmstart_search == 0 )....

  #* ###### Post-estimation process ###############

      _coevec            = Optim.minimizer(mfun)  # coef. vec.
      numerical_hessian  = hessian!(_Hessian, _coevec)  # Hessain

     #* ------ Check if the matrix is invertible. ----

     var_cov_matrix = try
                         inv(numerical_hessian)
                      catch err
                         redflag = 1
                         checkCollinear(tagD[:modelid], xvar, zvar, qvar, wvar, vvar) # check if it is b/c of multi-collinearity in the data
                         throw("The Hessian matrix is not invertible, indicating the model does not converge properly. The estimation is abort.")
                      end

          #* In some cases the matrix is invertible but the resulting diagonal
          #*    elements are negative. Check.

          if !all( diag(var_cov_matrix) .> 0 ) # not all are positive
               redflag = 1
               printstyled("Some of the diagonal elements of the var-cov matrix are non-positive, indicating problems in the convergence. The estimation is abort.\n\n"; color = :red)
               checkCollinear(tagD[:modelid], xvar, zvar, qvar, wvar, vvar) # check if it is b/c of multi-collinearity in the data
          end

     #* ------- JLMS and BC index -------------------

     if _dicOPT[:ineff_index]
        @views (_jlms, _bc) = jlmsbc(tagD[:modelid], _porc, pos, _coevec,
                                     yvar, xvar, zvar, qvar, wvar, vvar,         rowIDT)
        _jlmsM = mean(_jlms) # sum(_jlms)/length(_jlms)
        _bcM   = mean(_bc)   # sum(_bc)/length(_bc)
     else
        _jlms  = nothing
        _bc    = nothing
        _jlmsM = nothing
        _bcM   = nothing
     end

     #* ---- marginal effect on E(u) --------------

     if _dicOPT[:marginal]
        margeff, margMinfo = get_marg(tagD[:modelid], pos, num, _coevec, zvar, qvar, wvar)
     else
        margeff, margMinfo = nothing, ()
     end

     #* ------- Make Table ------------------

     stddev  = sqrt.(diag(var_cov_matrix)) # standard error
     t_stats = _coevec ./ stddev          # t statistics
     p_value = zeros(num.nofpara)   # p values
     ci_low  = zeros(num.nofpara) # confidence interval
     ci_upp  = zeros(num.nofpara)
     tt      = cquantile(Normal(0,1), 0.025)

     for i = 1:num.nofpara
         @views p_value[i,1] = pvalue(TDist(num.nofobs - num.nofpara), t_stats[i,1]; tail=:both)
         @views ci_low[i,1] = _coevec[i,1] - tt*stddev[i,1]
         @views ci_upp[i,1] = _coevec[i,1] + tt*stddev[i,1]
     end

       #* Build the table columns *#

    table = zeros(num.nofpara, 7)  # 7 columns in the table
    table[:,2] = _coevec   # estiamted coefficients
    table[:,3] = stddev    # std deviation
    table[:,4] = t_stats   # t statistic
    table[:,5] = p_value   # p value
    table[:,6] = ci_low
    table[:,7] = ci_upp
    table      = [" " "Coef." "Std. Err." "z" "P>|z|" "95%CI_l" "95%CI_u"; table]  # add to top of the table

       #*  creating a column of function names

    table[:, 1] .= ""
    for i in 1:length(eqvec)
        @views j = eqvec[i]
        @views table[j,1] = keys(eqvec)[i]
    end

       #*  Add the column of variable names

    table = hcat(varlist, table)                      # combine the variable names column (hcat, horizontal concatenate; see also vcat)
    table[:,1], table[:,2] = table[:,2], table[:,1]   # swap the first name column and the function column
    table[1,2] = "Var."

     # * ------ Print Results ----------- *#

     if _dicOPT[:verbose]

         printstyled("*********************************\n "; color=:cyan)
         printstyled("      Estimation Results:\n"; color=:cyan);
         printstyled("*********************************\n"; color=:cyan)

         print("Model type: "); printstyled(minfo1; color=:yellow); println()
         print("Number of observations: "); printstyled(num.nofobs; color=:yellow); println()
         print("Number of total iterations: "); printstyled(sf_total_iter; color=:yellow); println()
         if Optim.converged(mfun)
             print("Converged successfully: "); printstyled(Optim.converged(mfun); color=:yellow); println()
         elseif Optim.converged(mfun) == false
             print("Converged successfully: "); printstyled(Optim.converged(mfun); color=:red); println()
             redflag = 1
         end
         print("Log-likelihood value: "); printstyled(round(-1*Optim.minimum(mfun); digits=5); color=:yellow); println()
         println()

         pretty_table(table[2:end,:],    # could print the whole table as is, but this prettier
                      header=["", "Var.", "Coef.", "Std.Err.", "z", "P>|z|",
                              "95%CI_l", "95%CI_u"],
                      formatters = ft_printf("%5.4f", 3:8),
                      compact_printing = true,
                      backend = Val(sf_table))
         println()

         # *----- Auxiliary Table, log parameters to original scales --------

         auxtable = Array{Any}(undef,2,3)
         rn = 0 # row index

         if size(wvar,2) == 1 # single variable
            varstd = sqrt(sum((wvar .- sum(wvar)/length(wvar)).^2)/length(wvar))
            if varstd  <= 1e-7 # constant, assuming =1 anyway
                rn += 1
                auxtable[rn, 1] = :σᵤ²
                auxtable[rn, 2] = exp(_coevec[pos.begw])
                auxtable[rn, 3] = exp(_coevec[pos.begw])*stddev[pos.begw]
            end
         end

         if size(vvar,2) == 1 # single variable
            varstd = sqrt(sum((vvar .- sum(vvar)/length(vvar)).^2)/length(vvar))
            if varstd  <= 1e-7 # constant
                rn += 1
                auxtable[rn, 1] = :σᵥ²
                auxtable[rn, 2] = exp(_coevec[pos.begv])
                auxtable[rn, 3] = exp(_coevec[pos.begv])*stddev[pos.begv]
            end
         end

         if rn >= 1  # table is non-empty
             println("Convert the constant log-parameter to its original scale, e.g., σ² = exp(log_σ²):")
             pretty_table(auxtable[1:rn,:],
                          header=["", "Coef.", "Std.Err."],
                          formatters = ft_printf("%5.4f", 2:3),
                          compact_printing = true,
                          backend = Val(sf_table))

             print("\nTable format: "); printstyled("$(sf_table)"; color=:yellow); println(". Use sfmodel_opt() to choose between text, html, and latex.")
             println()
         end

         printstyled("***** Additional Information *********\n"; color=:cyan)

         print("* OLS (frontier-only) log-likelihood: "); printstyled(round(ll_ols; digits=5); color=:yellow); println("")
         print("* Skewness of OLS residuals: "); printstyled(round(sk_ols; digits=5); color=:yellow); println("")
         if _dicOPT[:ineff_index]
            print("* The sample mean of the JLMS inefficiency index: "); printstyled(round(_jlmsM; digits=5); color=:yellow); println("")
            print("* The sample mean of the BC efficiency index: "); printstyled(round(_bcM; digits=5); color=:yellow); println("\n")
         end
         if length(margMinfo) >= 1
            print("* The sample mean of inefficiency determinants' marginal effects on E(u): " ); printstyled(margMinfo; color=:yellow); println("")
            println("* Marginal effects of the inefficiency determinants at the observational level are saved in the return. See the follows.\n")
         end

         println("* Use `name.list` to see saved results (keys and values) where `name` is the return specified in `name = sfmodel_fit(..)`. Values may be retrieved using the keys. For instance:")
         println("   ** `name.loglikelihood`: the log-likelihood value of the model;")
         println("   ** `name.jlms`: Jondrow et al. (1982) inefficiency index;")
         println("   ** `name.bc`: Battese and Coelli (1988) efficiency index;")
         println("   ** `name.marginal`: a DataFrame with variables' (if any) marginal effects on E(u).")
         println("* Use `keys(name)` to see available keys.")

         printstyled("**************************************\n\n\n"; color=:cyan)

     end  # if_verbose

  #* ########### create a dictionary and make a tuple for return ########### *#

      _dicRES = OrderedDict{Symbol, Any}()
      _dicRES[:converged]          = Optim.converged(mfun)
      _dicRES[:iter_limit_reached] = Optim.iteration_limit_reached(mfun)
      _dicRES[:_______________] = "___________________"  #33
      _dicRES[:n_observations]  = num.nofobs
      _dicRES[:loglikelihood]   = -Optim.minimum(mfun)
      _dicRES[:table]           = [table][1]
      _dicRES[:coeff]           = _coevec
      _dicRES[:std_err]         = stddev
      _dicRES[:var_cov_mat]     = [var_cov_matrix][1]
      _dicRES[:jlms]            = _jlms
      _dicRES[:bc]              = _bc
      _dicRES[:OLS_loglikelihood] = ll_ols
      _dicRES[:OLS_resid_skew]    = sk_ols
      _dicRES[:marginal]      = margeff
      _dicRES[:marginal_mean] = margMinfo
      _dicRES[:_____________] = "___________________"  #31
      _dicRES[:model]         = minfo1
#     _dicRES[:data]          = "$sfdat"
      _dicRES[:depvar]        = _dicM[:depvar]
      _dicRES[:frontier]      = _dicM[:frontier]
      _dicRES[:μ]             = _dicM[:μ]
      _dicRES[:hscale]        = _dicM[:hscale]
      _dicRES[:gamma]         = _dicM[:gamma]
      _dicRES[:σₐ²]           = _dicM[:σₐ²]
      _dicRES[:σᵤ²]           = _dicM[:σᵤ²]
      _dicRES[:σᵥ²]           = _dicM[:σᵥ²]
      _dicRES[:log_σₐ²]       = _dicM[:σₐ²]
      _dicRES[:log_σᵤ²]       = _dicM[:σᵤ²]
      _dicRES[:log_σᵥ²]       = _dicM[:σᵥ²]
      _dicRES[:type]          = _dicM[:type]
      _dicRES[:dist]          = _dicM[:dist]
      _dicRES[:PorC]          = _porc
      _dicRES[:timevar]       = _dicM[:timevar]
      _dicRES[:idvar]         = _dicM[:idvar] # for bootstrap marginal effect
      _dicRES[:table_format]  = _dicOPT[:table_format]
      _dicRES[:modelid]       = tagD[:modelid]
      _dicRES[:verbose]       = _dicOPT[:verbose]
      _dicRES[:hasDF]         = _dicM[:hasDF]
      _dicRES[:transfer]      = _dicM[:transfer]

    for i in 1:length(eqvec2)
        _dicRES[keys(eqvec2)[i]] = _coevec[eqvec2[i]]
    end

      _dicRES[:________________]  = "___________________" #34
      _dicRES[:Hessian]           = [numerical_hessian][1]
      _dicRES[:gradient_norm]     = Optim.g_residual(mfun)
    # _dicRES[:trace]             = Optim.trace(mfun)     # comment out because not very informative and size could be large
      _dicRES[:actual_iterations] = Optim.iterations(mfun)
      _dicRES[:______________] = "______________________" #32
      _dicRES[:warmstart_solver] = sf_ini_algo_dic
      _dicRES[:warmstart_ini]    = sf_init_1st_dic
      _dicRES[:warmstart_maxIT]  = sf_ini_maxit_dic
      _dicRES[:main_solver]      = sf_algo
      _dicRES[:main_ini]         = sf_init_2nd_dic
      _dicRES[:main_maxIT]       = sf_maxit
      _dicRES[:tolerance]        = sf_tol
      _dicRES[:eqpo]             = eqvec2

      _dicRES[:redflag]          = redflag

     #* ----- Delete optional keys that have value nothing,

         for k in (:μ, :hscale,  :gamma)
             if _dicRES[k] === nothing
                delete!(_dicRES, k)
             end
         end

     #* ----- Create a NamedTuple from the dic as the final output;
     #* -----     put the dic in the tuple.

         _ntRES = NamedTuple{Tuple(keys(_dicRES))}(values(_dicRES))
         _ntRES = (; _ntRES..., list    = _dicRES)

     #* ---- Create a gloal dictionary for sf_predict ----

        global _eqncoe
        _eqncoe = Dict{Symbol, Vector}()  # nullify and initiate new dictionaries when a new model is specified

        for i in 1:length(eqvec2)
            _eqncoe[keys(eqvec)[i]]  = _coevec[eqvec2[i]] # for sf_predict
        end

  #* ############  make returns  ############ *#

      return _ntRES

end # sfmodel_fit

#############################################################
###  bootstrapping std.err. of mean marginal effect   #######
#############################################################
"""
    sfmodel_boot_marginal(<keyword arguments>)

Bootstrap standard errors and obtain bias-corrected (BC) confidence intervals
for the mean marginal effects of inefficiency determinants. In default, return
a ``K x 2``` matrix of standard errors (1st column) and confidence intervals
(tuples, 2nd column), where ``K`` is the number of exogenous inefficiency
determinants. With `getBootData=true`, return two matrices: the first is the
same as in the default return, and the second is the ``R x K`` bootstrapped
data.

See also the help file on `sfmodel_CI()`.

# Arguments
- result=<returned result>: The returned result from `sfmodel_fit()`.
- data=<dataset>: The DataFrame dataset containing the model's data.
  Same as the one used in `sfmodel_fit()`. If the data was supplied by matrix
  (instead of DataFrame; i.e., the Method 2 of `sfmodel_spec()`), this option
  should be skipped.
- R::Integer=<number>: The number of bootstrapped samples. The default is 500.
- level::Real=<number>: The significance level (default=0.05) of the bias-corrected
  confidence intervals. If `level`>0.5, it is automatically transformed to
  `1-level`, such that `level=0.05` and `level=0.95` both return 95%
  confidene intervals at the 5% significance level.
- seed::Integer=<number>: A postive integer used to seed the random
  number generator (rng) for resampling, which ensures reproducibility.
  This rng is not global and is only effective in this function. If not
  specified, the global random number generator is used, and the bootstrap
  result may change (slightly) between different runs.
- iter::Integer=<number>: The maximum number of iterations for each bootstrapped
  sample. If the number is larger than 0, it overwrites the `main_maxIT`
  specified in `sfmodel_opt()` which is the default.
- getBootData::Bool=false: Whether to return the bootstrapped data which is
  ``R x K`` where K is the number of exogenous determinants of inefficiency.
- every::Integer=10: Print bootstrapping progress for every `every` samples.

# Remarks
- Bootstrap samples are with replacement. For panel data, it samples
  cross-sectional units with replacement.
- In the MLE estimation, estimated
  coefficients from the main result is used as initial values. There is no
  `warmstart`. The `main_solver`, `main_maxIT`, and `tolerance` specified in
  `sfmodel_opt()` are used as default, but the value of `main_maxIT` may be
  replaced by the `iter` option.

# Examples
```julia
julia> std_ci = sfmodel_boot_marginal(result=res, data=df, R=250, seed=123)
bootstrap in progress..10..20..30..40..50..60..70..80..90..100..110..120..130..140..150..160..170..180..190..200..210..220..230..240..250..Done!

┌────────┬──────────────────────┬─────────────────┬──────────────────────┐
│        │ mean of the marginal │ std.err. of the │       bias-corrected │
│        │       effect on E(u) │     mean effect │    95.0%  conf. int. │
├────────┼──────────────────────┼─────────────────┼──────────────────────┤
│    age │             -0.00264 │         0.00225 │   (-0.00734, 0.0016) │
│ school │             -0.01197 │         0.01765 │    (-0.048, 0.01224) │
│     yr │             -0.02650 │         0.01221 │ (-0.05257, -0.00447) │
└────────┴──────────────────────┴─────────────────┴──────────────────────┘

3×2 Matrix{Any}:
 0.00224786  (-0.00734, 0.0016)
 0.0176472   (-0.048, 0.01224)
 0.012213    (-0.05257, -0.00447)

julia> std_ci, bsdata = sfmodel_boot_marginal(result=res, data=df, R=250, seed=123, getBootData=true);

 (output omitted)

julia> bsdata
250×3 adjoint(::Matrix{Real}) with eltype Real:
-0.000493033  -0.00840051   -0.0339088
-0.00510175   -0.0128184    -0.0147759
-0.0024987    -0.00971428   -0.00989262
-0.00157346   -0.0265515    -0.0140001
-0.00352179   -0.00670365   -0.0246122
-0.00375162   -0.0070496    -0.0306034
-0.00153094   -0.0154201    -0.0367731
0.000149329   0.00672036   -0.0461389
⋮
-0.00373306   -0.0108364    -0.00871898
-0.00170254   -0.0393002    -0.0500638
0.000686169   0.00241594   -0.018542
0.000258745   0.000183392  -0.039621
-0.00408104   -0.014574     -0.024126
-0.00417206   -0.0192443    -0.0406959
0.00266017   -0.0396552    -0.0359759
```
"""
function sfmodel_boot_marginal(; result::Any=nothing,  data::Any=nothing,
                                 R::Integer=500,
                                 level::Real=0.05,
                                 mymisc=nothing,
                                 iter::Integer=-1,
                                 getBootData::Bool=false,
                                 seed::Integer=-1,
                                 every::Integer=10)

  if data === nothing
      if result.transfer == true
          data = _dicM[:sdf]
      else
         throw("Need the `data=<DataFrame name>` option.")
      end
  elseif typeof(data) != DataFrame
      throw("`data=` needs to be a DataFrame.")
  end

   # getvar still use info from _dicM

  if result.marginal_mean==NamedTuple() || result.marginal_mean === nothing
      printstyled("The model does not have exogenous determinants of inefficiency, ", #=
               =# "or you didn't compute the marginal effect in the main estimation. Abort the bootstrap.\n"; color = :red)
      throw("No statistics to bootstrap.")
  end

  ((level > 0.0) && (level < 1.0)) || throw("The significance level (`level`) should be between 0 and 1.")

  if level > 0.5
     level = 1-level  # 0.95 -> 0.05
  end

     # In the following lines, the integer part had been taken care of in Type.
  (seed == -1) || ( seed > 0) || throw("`seed` needs to be a positive integer.")
  (iter == -1) || ( iter > 0) || throw("`iter` needs to be a positive integer.")
  (R > 0) || throw("`R` needs to be a positive integer.")

#* ##### Get parameters #######

          _porc = result.PorC
        sf_init = result.coeff
        sf_algo = eval(result.main_solver)
       sf_maxit = result.main_maxIT
         sf_tol = result.tolerance
       sf_table = result.table_format

       if iter > 0
           sf_maxit = iter
       end

#* ########  parepare data  ##########

  (minfo1, minfo2, pos, num, eqvec, eqvec2, yvar, xvar, zvar, qvar, wvar,
   vvar,         rowIDT, varlist) = SFrontiers.getvar(result.modelid, data)

  #----- make an orginal copy; also make later code easier --------

   yvar0 = deepcopy(yvar); xvar0 = deepcopy(xvar); zvar0 = deepcopy(zvar);
   qvar0 = deepcopy(qvar); wvar0 = deepcopy(wvar); vvar0 = deepcopy(vvar);
   rowIDT0 = deepcopy(rowIDT)

#* ####### begin looping #############

 sim_res = Array{Real}(undef, ncol(result.marginal), R)

 if seed > 0
    rng = MersenneTwister(seed)
 end

 for i in 1:R

    @label start1

   # ----- sampling data for cross-sectional data --------------

   if result.idvar === nothing  # i.e., cross-sectional

     if seed == -1
        select_row = sample( 1:nrow(data), nrow(data); replace=true)  # require StatsBase.jl
     else
        select_row = sample(rng, 1:nrow(data), nrow(data); replace=true)  # require StatsBase.jl
     end

     yvar =  yvar0[select_row, :]
     xvar =  xvar0[select_row, :]
     (zvar0 == ()) ||  (zvar =  zvar0[select_row, :])
     (qvar0 == ()) ||  (qvar =  qvar0[select_row, :])
     (wvar0 == ()) ||  (wvar =  wvar0[select_row, :])
     (vvar0 == ()) ||  (vvar =  vvar0[select_row, :])

  end

   # ------ sampling panel data ------

   if result.idvar !== nothing  # i.e., panel data

      if seed == -1
        select_row = sample(     1:size(rowIDT0,1), size(rowIDT0,1); replace=true)  # require StatsBase.jl
      else
        select_row = sample(rng, 1:size(rowIDT0,1), size(rowIDT0,1); replace=true)  # require StatsBase.jl
      end

      yvar = mySample(rowIDT0, yvar0, select_row)

      xvar = mySample(rowIDT0, xvar0, select_row)
      (zvar0 == ()) ||  (zvar = mySample(rowIDT0, zvar0, select_row))
      (qvar0 == ()) ||  (qvar = mySample(rowIDT0, qvar0, select_row))
      (wvar0 == ()) ||  (wvar = mySample(rowIDT0, wvar0, select_row))
      (vvar0 == ()) ||  (vvar = mySample(rowIDT0, vvar0, select_row))

       newidvar = mySample2!(rowIDT0, ones(size(yvar,1)), select_row) # a trick to get new id; note not yvar0
       rowIDT = get_rowIDT(vec(newidvar))
       # yvar = vec(yvar)

    end

 #*  --- other small things -----

   nofobs1 = size(yvar,1)

  #* ----- Define the problem's Hessian -----#

   _Hessian = TwiceDifferentiable(rho -> SFrontiers.LL_T(result.modelid,
                         yvar, xvar, zvar, qvar, wvar, vvar,
                         _porc, nofobs1, pos, rho,
                                 rowIDT, mymisc),
                          sf_init;
                          autodiff = :forward);

  #* ---- estimate ----------------- *#

             mfun = try
                       optimize(_Hessian,
                                sf_init,       # initial values
                                sf_algo,       # different from search run
                                Optim.Options(g_tol = sf_tol,
                                              iterations  = sf_maxit, # different from search run
                                              store_trace = false,
                                              show_trace  = false))
                    catch err
                       @goto start1
                    end

#* ###### check if valid ###############

  if (Optim.iteration_limit_reached(mfun) ) ||
     (isnan(Optim.g_residual(mfun)) ) ||
     (Optim.g_residual(mfun) > 1e-1)
         @goto start1
  end

    _coevec            = Optim.minimizer(mfun)  # coef. vec.
    numerical_hessian  = hessian!(_Hessian, _coevec)  # Hessain

   #* ------ Check if the matrix is invertible. ----

   var_cov_matrix = try
                       inv(numerical_hessian)
                    catch err
                       @goto start1
                    end

        if !all( diag(var_cov_matrix) .> 0 ) # not all are positive
              @goto start1
        end

   #* ---- marginal effect on E(u) --------------

   margeff, margMinfo = get_marg(result.modelid, pos, num, _coevec, zvar, qvar, wvar)

   theM = mean.(eachcol(margeff))  # mean marginal effect

   if sum(isnan.(theM)) != 0  # in this run some of the element is NaN
      @goto start1
   end

   sim_res[:, i] = theM

   if result.verbose
      if i == 1
         printstyled("bootstrap in progress.."; color = :cyan)
      end
      if i%every == 0   # print for every 10 replications
          print("$(i)..")
      end
      if i == R
          printstyled("Done!\n\n"; color = :cyan)
      end
   end

end   # for i=1:R

 sim_res = sim_res'  # nofobs x K; ith row = ith replication, column = statistics

 #*###### compute statistics for the mean marginal effect #####*#

   theMean = mean.(eachcol(result.marginal))  # Kx1 vector
   theSTD  = sqrt.( sum((sim_res .- theMean').^2, dims=1) ./(R-1)) # 1 x K

   ci_mat = sfmodel_CI(bootdata=sim_res, observed=theMean, level=level, verbose=false);

   if result.verbose

      table = Array{Any}(undef, length(theMean), 3+1)
      vname = names(result.marginal)

      for j in 1:length(theMean)
          table[j, 1] = string(vname[j])[6:end]
      end

      table[:,2], table[:,3] = theMean, theSTD
      table[:,4] = ci_mat

       mylevel = 100*(1-level)

      table = [" " "mean of the marginal" "std.err. of the"  "bias-corrected";
               " " "effect on E(u)"       "mean effect"      "$(mylevel)%  conf. int.";
               table]

            pretty_table(table,
                         noheader= true,
                         body_hlines = [2],
                         formatters = ft_printf("%0.5f", 2:4),
                         compact_printing = true,
                         backend = Val(sf_table))
            println()

   end

   if getBootData == false   # most likely
      return hcat(theSTD', ci_mat) # K x 1
   else
      return hcat(theSTD', ci_mat), sim_res
   end

end # sfmodel_boot_marginal

# -----------------------------------------------#
# --- utilities for sfmodel_boot_marginal -------#
# -----------------------------------------------#

function mySample(rowIDT, data, select)
  temp1 = Array{Float64}(undef, 0, size(data,2))
  @inbounds for i=1:size(rowIDT,1)
       @views   cc = data[rowIDT[select[i],1], :]
                temp1 = vcat(temp1, cc)
            end
  return temp1
end

function mySample2!(rowIDT, data, select)
  temp1 = Array{Float64}(undef, 0, size(data,2))
  @inbounds for i=1:size(rowIDT,1)
       @views   cc = data[rowIDT[select[i],1], :]
                cc[:,1] .= i     # caution, it changes data
                temp1 = vcat(temp1, cc)
            end
  return temp1
end

"""
function sfmodel_CI(; bootdata::Any=nothing, observed::Union{Vector, Real, Tuple, NamedTuple}=nothing, level::Real=0.05, verbose::Bool=true)

    sfmodel_CI(<keyword arguments>)

A general purpose (not specific to stochastic frontier models) function for
obtaining bias-correctred (BC) confidence intervals from bootstrapped data.
Return a ``K x 1`` matrix of confidence intervals in the form of tuples, where
``K`` is the number of bootstrap statistics.

See also the help file on `sfmodel_boot_marginal()`.

# Arguments
- bootdata::Array=<data>: The bootstrapped data of size ``R x K``,
  where ``R`` is the number of bootstrap samples (replications) and ``K`` is
  the number of statistics. An example is the bootstrapped data from
  `sfmodel_boot_marginal(, ... getBootData=true)`.
- observed::Union{Vector, Real, Tuple, NamedTuple}=<a vector of numbers>: The
  observed values of the statistics to which the confidence intervals are to
  be calculated. The length of `observed` should be equal to ``K``. It could take
  the form of a single value (if ``K=1``), a vector, a tuple, or a NamedTuple.
- level::Real=<number>: The significance level (default=0.05) of the bias-corrected
  confidence intervals. If `level`>0.5, it is automatically transformed to
  `1-level`, such that `level=0.05` and `level=0.95` both return 95%
  confidene intervals at the 5% significance level.
- verbose::Bool=true: Print the result.

# Examples
```julia
julia> myans = sfmodel_fit(useData(df));

 (output omitted)

julia> std_ci, bsdata = sfmodel_boot_marginal(result=myans, data=df, R=250, seed=123, getBootData=true);

 (output omitted)

julia> sfmodel_CI(bootdata=bsdata, observed=myans.marginal_mean, level=0.10)

Bias-Corrected 90.0% Confidence Interval:

3×1 Matrix{Any}:
 (-0.00655, 0.0009)
 (-0.04064, 0.0086)
 (-0.04663, -0.00872)

julia> # manually input observed values

julia> sfmodel_CI(bootdata=bsdata, observed=(-0.00264, -0.01197, -0.0265), level=0.10)

Bias-Corrected 90.0% Confidence Interval:

3×1 Matrix{Any}:
 (-0.00655, 0.0009)
 (-0.04064, 0.0086)
 (-0.04663, -0.00872)
```
"""
function sfmodel_CI(; bootdata::Any=nothing, observed::Union{Vector, Real, Tuple, NamedTuple}=nothing, level::Real=0.05, verbose::Bool=true)

    # bias-corrected (but not accelerated) confidence interval
    # For the "accelerated" factor, need to estimate the SF model
    #    for every jack-knifed sample, which is expensive.

    if (observed isa NamedTuple)
        observed = values(observed)
    end

   ((level > 0.0) && (level < 1.0)) || throw("The significance level (`level`) should be between 0 and 1.")

   if level > 0.50
      level = 1-level  # 0.95 -> 0.05
   end

    nofobs, nofK = size(bootdata)  # number of statistics

    (nofK == length(observed)) || throw("The number of statistics (`observed`) does not fit the number of columns of bootstrapped data.")

    ci = Array{Any}(undef, nofK, 1)

    z1 = quantile(Normal(), level/2)
    z2 = quantile(Normal(), 1 - level/2)  #! why z1 != z2?

    for i in 1:nofK
        @views data = bootdata[:,i]

        count = sum(data .< observed[i])
        z0 = quantile(Normal(), count/nofobs) # bias corrected factor

        alpha1 = cdf(Normal(), z0 + ((z0 + z1) ))
        alpha2 = cdf(Normal(), z0 + ((z0 + z2) ))

        order_data = sort(data)

        ci[i,1] = (  round(order_data[Int(ceil(nofobs*alpha1))], digits=5),    round(order_data[Int(ceil(nofobs*alpha2))], digits=5)  )
    end

    if verbose == true
       mylevel = 100*(1-level)
       println("\nBias-Corrected $(mylevel)% Confidence Interval:\n")
    end

    return ci
end

########################################################
####            catching type error                 ####
########################################################

function sfmodel_init(arg::Vector)
    throw(ArgumentError("The initial values in sfmodel_init() are specified incorrectly. They must be supplied in macros such as all_init(0.1, 0.1)."))
end

function sfmodel_init(args::Number)
    throw(ArgumentError("The initial values in sfmodel_init() are specified incorrectly. They must be supplied in macros such as all_init(0.1, 0.1)."))
end

function sfmodel_fit(sfdat::Any)
    throw(ArgumentError("The dataset specified in sfmodel_fit() must be a DataFrame."))
end


using SafeTestsets, Test
using Surrogates: sample, SobolSample

@safetestset "AbstractGPSurrogate" begin

using Surrogates
using SurrogatesAbstractGPs
using AbstractGPs
using Zygote

@testset "1D -> 1D" begin
    lb = 0.0
    ub = 3.0
    f = x -> log(x)*exp(x);
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    agp1D = AbstractGPSurrogate(x,y, gp=GP(SqExponentialKernel()), Σy=0.05)
    x_new = 2.5
    y_actual = f.(x_new)
    y_predicted = agp1D([x_new])
    @test isapprox(y_predicted, y_actual, atol=0.1)
end

@testset "add points 1D" begin
    lb = 0.0
    ub = 3.0
    f = x -> x^2;
    x_points = sample(5,lb,ub,SobolSample())
    y_points = f.(x_points)
    agp1D = AbstractGPSurrogate([x_points[1]],[y_points[1]], gp=GP(SqExponentialKernel()), Σy=0.05)
    x_new = 2.5
    y_actual = f.(x_new)
    for i in 2:length(x_points)
        add_point!(agp1D, x_points[i], y_points[i]);
    end
    y_predicted = agp1D([x_new]);
    @test isapprox(y_predicted, y_actual, atol=0.1)
end

@testset "2D -> 1D" begin
    lb = [0.0; 0.0]
    ub = [2.0; 2.0]
    log_exp_f = x -> log(x[1])*exp(x[2])
    x = sample(50,lb,ub,SobolSample())
    y = log_exp_f.(x)
    agp_2D = AbstractGPSurrogate(x, y)
    x_new_2D = (2.0, 1.0)
    y_actual = log_exp_f(x_new_2D)
    y_predicted = agp_2D(x_new_2D)
    @test isapprox(y_predicted, y_actual, atol=0.1)
end

@testset "add points 2D" begin
    lb = [0.0;0.0]
    ub = [2.0;2.0]
    sphere = x -> x[1]^2 + x[2]^2
    x = sample(20, lb, ub, SobolSample())
    y = sphere.(x)
    agp_2D = AbstractGPSurrogate([x[1]],[y[1]])
    logpdf_vals = []
    push!(logpdf_vals, logpdf_surrogate(agp_2D))
    for i in 2:length(x)
        add_point!(agp_2D,x[i],y[i])
        push!(logpdf_vals, logpdf_surrogate(agp_2D))
    end
    @test first(logpdf_vals) < last(logpdf_vals) #as more points are added log marginal posterior predictive probability increases
end

@testset "check ND prediction" begin
    lb = [-1.0;-1.0;-1.0]
    ub = [1.0;1.0;1.0]
    f = x ->hypot(x...)
    x = sample(25,lb,ub,SobolSample())
    y = f.(x)
    agpND = AbstractGPSurrogate(x, y, gp = GP(SqExponentialKernel()), Σy = 0.05)
    x_new = (-0.8,0.8,0.8)
    @test agpND(x_new) ≈ f(x_new) atol=0.2
end

@testset "Optimization 1D" begin
    objective_function = x -> 2*x+1
    lb = 0.0
    ub = 6.0
    x = [2.0,4.0,6.0]
    y = [5.0,9.0,13.0]
    p = 2
    a = 2
    b = 6
    my_k_EI1 = AbstractGPSurrogate(x,y)
    surrogate_optimize(objective_function,EI(),a,b,my_k_EI1,UniformSample(),maxiters=200,num_new_samples=155)
end

@testset "Optimization ND" begin
    objective_function_ND = z -> 3*hypot(z...)+1
    x = [(1.2,3.0),(3.0,3.5),(5.2,5.7)]
    y = objective_function_ND.(x)
    theta = [2.0,2.0]
    lb = [1.0,1.0]
    ub = [6.0,6.0]
    my_k_E1N = AbstractGPSurrogate(x,y)
    surrogate_optimize(objective_function_ND,EI(),lb,ub,my_k_E1N,UniformSample())
end

@testset "check working of logpdf_surrogate 1D" begin
    lb = 0.0
    ub = 3.0
    f = x -> log(x)*exp(x);
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    agp1D = AbstractGPSurrogate(x,y, gp=GP(SqExponentialKernel()), Σy=0.05)
    logpdf_surrogate(agp1D)
end

@testset "check working of logpdf_surrogate ND" begin
    lb = [0.0; 0.0]
    ub = [2.0; 2.0]
    f = x -> log(x[1])*exp(x[2]);
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    agpND = AbstractGPSurrogate(x,y, gp=GP(SqExponentialKernel()), Σy=0.05)
    logpdf_surrogate(agpND)
end

lb = 0.0
ub = 3.0
n = 10
x = sample(n,lb,ub,SobolSample())
f = x -> x^2
y = f.(x)

#AbstractGP 1D
@testset "AbstractGP 1D" begin
    agp1D = AbstractGPSurrogate(x,y, gp=GP(SqExponentialKernel()), Σy=0.05)
    g = x -> agp1D'(x)
    g([2.0])
end

lb = [0.0,0.0]
ub = [10.0,10.0]
n = 5
x = sample(n,lb,ub,SobolSample())
f = x -> x[1]*x[2]
y = f.(x)

# AbstractGP ND
@testset "AbstractGPSurrogate ND" begin
    my_agp = AbstractGPSurrogate(x,y, gp=GP(SqExponentialKernel()), Σy=0.05)
    g = x ->Zygote.gradient(my_agp, x)
    #g([(2.0,5.0)])
    g((2.0,5.0))
 end

end


export Polynomial, print

"""
    Polynomial(poly::SortedDict{Exponent, Number}, degree::Degree)

Define a mathematical polynomial, that is a linear combinaison of product of
`Variable` and conjugated `Variable`.

### Attributes
- `poly` : a dictionary associating an `Exponent` to a a number.
- `degree`: a `Degree` object indicating the global degree of the polynomial :
the maximum of the explicit exponent degrees and of the conjugated exponent
degrees.

### Exemple
```julia
julia > a, b = Variable("a", Complex), Variable("b", Complex)
julia > p = a*b*conj(b) + (2+3im) * b^3
julia > p.poly = SortedDict(a*b*conj(b)=>1, b^3=>2+3im)
julia > p.degree = (3,1)
```
"""
struct Polynomial <: AbstractPolynomial
    poly::SortedDict{Exponent, Number}
    degree::Degree

    function Polynomial(poly::SortedDict{Exponent, Number})
        degexpl, degconj = 0, 0
        for (expo, λ) in poly
            if λ!=0
                degexpl = max(degexpl, expo.degree.explvar)
                degconj = max(degconj, expo.degree.conjvar)
            else
                delete!(poly, expo)
            end
        end
        return new(poly, Degree(degexpl, degconj))
    end
end

#############################
## Constructors
#############################
Polynomial() = Polynomial(SortedDict{Exponent, Number}())

#############################
## Equality, hash
#############################
function ==(p1::Polynomial, p2::Polynomial)
  if (length(p1) != length(p2)) || (p1.degree != p2.degree)
    return false
  end

  state1 = start(p1)
  state2 = start(p2)
  while !done(p1, state1) && !done(p2, state2)
    (i1, state1) = next(p1, state1)
    (i2, state2) = next(p2, state2)

    if !(i1 == i2)
      return false
    end
  end
  return true
end
!=(pol1::Polynomial, pol2::Polynomial) = !(pol1 == pol2)

#############################
## Print
#############################
function Base.print(io::IO, P::Polynomial)
  poly = P.poly
  i = length(poly)
  un = Exponent()
  for (expo, λ) in poly

    if expo == un
      print(io, "$λ")
    else
      if λ != 1
        print(io, "(", λ, ")*")
      end
      print(io, expo)
    end
    if i > 1
      print(io, " + ")
    end
    i -= 1
  end
end


# Author: Lukasz Bratos

module Blocksys

import SparseArrays
using LinearAlgebra

export readMatrix, readVector, solveGauss, solveGaussWithChoice, lu, solveLu, gaussWithChoice, solveLuWithChoice, printXToFile, printBToFile, computeRightSideVector

"""
Funkcja wczytuje macierz z pliku

Dane wejściowe:
    file - ścieżka pliku wejściowego
Dane wyjściowe:
    krotka (A, n, l), gdzie:
        A - macierz rzadka w formacie SparseMatrixCSC
        n - rozmiar macierzy
        l - rozmiar bloku
"""
function readMatrix(file::String)
    open(file) do f
        firstLine = split(readline(f))
        n = parse(Int64, firstLine[1])
        l = parse(Int64, firstLine[2])

        rows = Int64[]
        columns = Int64[]
        values = Float64[]

        for line in eachline(f)
            line = split(line)
            push!(rows, parse(Int64, line[1]))
            push!(columns, parse(Int64, line[2]))
            push!(values, parse(Float64, line[3]))
        end

        A = SparseArrays.sparse(rows, columns, values)
        return (A, n, l)
    end
end

"""
Funkcja wczytująca wektor z pliku

Dane wejściowe:
    file - ścieżka pliku wejściowego
Dane wyjściowe:
    vector - wczytany wektor
"""
function readVector(file::String)::Vector{Float64}
    open(file) do f
        n = parse(Int64, readline(f))
        vector = []

        for line in eachline(f)
            push!(vector, parse(Float64, line))
        end

        return vector
    end
end

"""
Funkcja zapisująca wektor do pliku

Dane wejściowe:
    file - ścieżka do pliku
    x - wektor, który chcemy zapisać
"""
function printXToFile(file::String, x::Vector{Float64})
    open(file, "w") do f
        foreach(a->write(f, string(a), "\n"), x)
    end
end

"""
Funkcja zapisująca wektor do pliku wraz z błędem względnym

Dane wejściowe:
    file - ścieżka do pliku
    x - wektor, który chcemy zapisać
"""
function printBToFile(file::String, x::Vector{Float64})
    open(file, "w") do f
        y = ones(Float64, length(x))
        write(f, string(norm(y - x) / norm(x)), "\n")
        foreach(a->write(f, string(a), "\n"), x)
    end
end

"""
Funkcja wyliczająca wektor prawych stron na podstawie macierzy A

Dane wejściowe:
    A - macierz rzadka w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
Dane wyjściowe:
    b - wektor prawych stron
"""
function computeRightSideVector(A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64)
    x = ones(Float64, n)
    b = zeros(Float64, n)

    for i in 1:n
        for j in max(1, i - (2 + l)):min(n, i + l)
            b[i] += A[i, j]
        end
    end

    return b
end

"""
Funkcja aplikująca metodę eliminacji Gaussa na macierzy A

Dane wejściowe:
    A - macierz rzadka w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
    b - wektor prawych stron
Dane wyjściowe:
    krotka (A, b), gdzie:
        A - zmodyfikowana macierz
        b - zmodyfikowany wektor prawych stron
"""
function gauss(A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64, b::Vector{Float64})
    for k in 1:(n - 1)
        for i in (k + 1):min(n, k + l + 1)
            z = A[i, k] / A[k, k]
            A[i, k] = 0.0

            for j in (k + 1):min(n, k + l)
                A[i, j] -= z * A[k, j]
            end

            b[i] -= z * b[k]
        end
    end

    return (A, b)
end

"""
Funkcja rozwiązująca układ równań Ax = b metodą eliminacji Gaussa

Dane wejściowe:
    A - macierz rzadka współczynników w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
    b - wektor prawych stron
Dane wyjściowe:
    x - wektor wyników
"""
function solveGauss(_A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64, _b::Vector{Float64})::Vector{Float64}
    A, b = gauss(_A, n, l, _b)

    x = zeros(Float64, n)

    for i in n:-1:1
        sum = 0.0

        for j in (i + 1):min(n, i + l)
            sum += A[i, j] * x[j]
        end

        x[i] = (b[i] - sum) / A[i, i]
    end

    return x
end

"""
Funkcja aplikująca metodę eliminacji Gaussa (z częściowym wyborem elementu głównego) na macierzy A

Dane wejściowe:
    A - macierz rzadka w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
    b - wektor prawych stron
Dane wyjściowe:
    krotka (A, b, perm), gdzie:
        A - zmodyfikowana macierz
        b - zmodyfikowany wektor prawych stron
        perm - wektor permutacji
"""
function gaussWithChoice(A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64, b::Vector{Float64})
    perm = [1:n;]

    for k in 1:(n - 1)
        max = 0.0
        maxIdx = 0

        for i in k:min(n, k + l + 1)
            if abs(A[perm[i], k]) > max
                max = abs(A[perm[i], k])
                maxIdx = i
            end
        end

        perm[maxIdx], perm[k] = perm[k], perm[maxIdx]

        for i in (k + 1):min(n, k + l + 1)
            z = A[perm[i], k] / A[perm[k], k]
            A[perm[i], k] = 0.0

            for j in (k + 1):min(n, k + 2 * l)
                A[perm[i], j] -= z * A[perm[k], j]
            end

            b[perm[i]] -= z * b[perm[k]]
        end
    end

    return (A, b, perm)
end

"""
Funkcja rozwiązująca układ równań Ax = b metodą eliminacji Gaussa (z częściowym wyborem elementu głównego)

Dane wejściowe:
    A - macierz rzadka współczynników w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
    b - wektor prawych stron
Dane wyjściowe:
    x - wektor wyników
"""
function solveGaussWithChoice(_A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64, _b::Vector{Float64})
    A, b, perm = gaussWithChoice(_A, n, l, _b)

    x = zeros(Float64, n)

    for i in n:-1:1
        sum = 0.0

        for j in (i + 1):min(n, i + 2 * l)
            sum += A[perm[i], j] * x[j]
        end

        x[i] = (b[perm[i]] - sum) / A[perm[i], i]
    end

    return x
end

"""
Funkcja dokonująca rozkładu LU macierzy A metodą eliminacji Gaussa

Dane wejściowe:
    A - macierz rzadka w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
Dane wyjściowe:
    A - macierz U
    L - macierz L
"""
function lu(A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64)
    # A is our U
    L = SparseArrays.spzeros(n, n)

    for k in 1:(n - 1)
        L[k ,k] = 1.0

        for i in (k + 1):min(n, k + l + 1)
            z = A[i, k] / A[k, k]
            L[i, k] = z
            A[i, k] = 0.0
            for j in (k + 1):min(n, k + 2 * l)
                A[i, j] -= z * A[k, j]
            end
        end
    end
    L[n,n] = 1.0

    return (A, L)
end

"""
Funkcja rozwiązująca układ równań Ax = b korzystając z rozkładu LU

Dane wejściowe:
    U - macierz U w formacie SparseMatrixCSC
    L - macierz L w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
    b - wektor prawych stron
Dane wyjściowe:
    x - wektor wyników
"""
function solveLu(U::SparseArrays.SparseMatrixCSC{Float64,Int64}, L::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64, b::Vector{Float64})::Vector{Float64}
    x = zeros(Float64, n)

    for k in 1:(n - 1)
        for i in (k + 1):min(n, k + l + 1)
            b[i] -= L[i, k] * b[k]
        end
    end

    for i in n:-1:1
        sum = 0.0

        for j in (i + 1):min(n, i + l)
            sum += U[i, j] * x[j]
        end

        x[i] = (b[i] - sum) / U[i, i]
    end

    return x
end

"""
Funkcja dokonująca rozkładu LU macierzy A metodą eliminacji Gaussa (z częściowym wyborem elementu głównego)

Dane wejściowe:
    A - macierz rzadka w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
Dane wyjściowe:
    krotka (A, perm), gdzie:
        A - macierz U
        L - macierz L
        perm - wektor permutacji
"""
function luWithChoice(A::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64)
    # A is our U
    perm = [1:n;]

    L = SparseArrays.spzeros(n, n)

    for k in 1:(n - 1)

        max = 0.0
        maxIdx = 0

        for i in k:min(n, k + l + 1)
            if abs(A[perm[i], k]) > max
                max = abs(A[perm[i], k])
                maxIdx = i
            end
        end

        perm[maxIdx], perm[k] = perm[k], perm[maxIdx]

        for i in (k + 1):min(n, k + l + 1)
            z = A[perm[i], k] / A[perm[k], k]

            L[perm[i], k] = z
            A[perm[i], k] = 0.0

            for j in (k + 1):min(n, k + 2 * l)
                A[perm[i], j] -= z * A[perm[k], j]
            end
        end
    end

    return (A, L, perm)
end

"""
Funkcja rozwiązująca układ równań Ax = b korzystając z rozkładu LU (z częściowym wyborem elementu głównego)

Dane wejściowe:
    A - macierz rozkładu LU w formacie SparseMatrixCSC
    n - rozmiar macierzy
    l - rozmiar bloku
    b - wektor prawych stron
    perm - wektor permutacji
Dane wyjściowe:
    x - wektor wyników
"""
function solveLuWithChoice(U::SparseArrays.SparseMatrixCSC{Float64,Int64}, L::SparseArrays.SparseMatrixCSC{Float64,Int64}, n::Int64, l::Int64, b::Vector{Float64}, perm::Vector{Int64})
    x = zeros(Float64, n)

    for k in 1:(n - 1)
        for i in (k + 1):min(n, k + l + 1)
            b[perm[i]] -= L[perm[i], k] * b[perm[k]]
        end
    end

    for i in n:-1:1
        sum = 0.0

        for j in (i + 1):min(n, i + 2 * l)
            sum += U[perm[i], j] * x[j]
        end

        x[i] = (b[perm[i]] - sum) / U[perm[i], i]
    end

    return x
end

end


module myJuliaUtils

using Random
using LinearAlgebra
using Statistics
using Distributions
using StatsBase: counts, Histogram, fit

export

# extfunc.jl
lag0, eye, vech, vechRshp, vec2sym, vec2ltri, vec2ltriW1, stabcheck, stabcheckC, preparexy, sumsqr, nanmean,
wish, iwish, cholPSD, regMat2PD!, randnPSD, genPSDmatStrict, genPDmat, genPSDmat, quantileArr, companionf,
ismyapprox, normpdf, acf, decVCV, ols1, VARols1, inbetween, getmultdiag!,transf1To,transf1Back,split_data,standardize,unstandardize, showprog,
# evalF.jl
useN2fit, useMvN2fit, pspit, pscrps,
# data_manipulation.jl
dataTransf, detrend,
# generate.jl
genLPM, genEndoMC, simMC, genMC

include("extfunc.jl");
include("evalF.jl");
include("data_manipulation.jl");
include("generate.jl");

end # module


function transformation_matrices(Nx::Int32, Ny::Int32)
    # calculates the interpolation and restriction matrices given a certain resolution in each direction
    # Nx = grid.Nx
    # Ny = grid.Ny
    coarsexres = Int32(Nx/2)
    coarseyres = Int32(Ny/2)

    xcoarse = zeros(Float64, coarsexres)
    ycoarse = zeros(Float64, coarseyres)  # using the reduction in number of points from the paper
    # way to define these with each new iteration? Another overarching loop with the number of iterations?
    # could do a while loop potentially, if trying to get under a threshold

    # define the interpolation matrix, then the restriction matrix as its transpose

    # only working in 1D so far
    Interpx = zeros(Nx, coarsexres)
    # Restrx1D = zeros(coarsexres, Nx)

    for i=1:Nx
        for j=1:coarsexres
            # sort of brute force based on the form of the interpolation matrix given in the MIT site
            Interpx[2(j-1)+1, j] = .5
            Interpx[2(j-1)+2, j] = 1
            if 2(j-1)+3 <= i
                Interpx[2(j-1)+3, j] = .5
            end
            # may not need outer loop if indexing only depends on j
        end
    end

    Restrx1D = .5*transpose(Interpx) # definition of restriction matrix from the paper

    # moving to 2D?
    # defining values based on 2D interpolation scheme in the paper
    """for i=1:Nx
        for j=1:coarsexres
            if isodd(i)
                if isodd(j)
                    Interpx[i,j] = .25
                elseif iseven(j)
                    Interpx[i,j] = .5
                end
            elseif iseven(i)
                if isodd(j)
                    Interpx[i,j] = .5
                elseif iseven(j)
                    Interpx[i,j] = 1
                end
            end
        end
    end"""

    # Restrx = .25*transpose(Interpx) # definition of restriction matrix
    Restr2d = kron(Restrx1D, Restrx1D) # assuming we're using the 1D matrix
    Interp2d = 4*transpose(Restr2d)
    # same Interpolation/Restriction matrices in x and y?

    return Restr2d, Interp2d, coarsexres, coarseyres
end

# not sure if fine_to_coarse and coarse_to_fine are needed, but they could be useful so I'll keep them in
function fine_to_coarse(Nx::Int32, Ny::Int32, x::Vector{Float64}, y::Vector{Float64})
    # restricts values from fine grid to fit them onto coarse grid

    Restriction, Interpolation, newxres, newyres = transformation_matrices(Nx, Ny)

    # calculating the coarse resolution for reducing the number of grid points by 2 in each dimension, assuming odd number of points in each dimension
    coarsexres = Int32(Nx/2)
    coarseyres = Int32(Ny/2)

    # initializing coarse grid values in each dimension
    xcoarse = zeros(Float64, coarsexres)
    ycoarse = zeros(Float64, coarseyres)

    # assuming dimensionality of both are the same?
    xcoarse[:] = Restriction * x
    ycoarse[:] = Restriction * y

    return xcoarse, ycoarse
end

function fine_to_coarse(Nx::Int32, Ny::Int32, x::Vector{Float64}, y::Vector{Float64})
    # projects vector values from the fine grid to the coarse grid using interpolation

    Restriction, Interpolation, newxres, newyres = transformation_matrices(Nx, Ny)

    # calculating the coarse resolution for reducing the number of grid points by 2 in each dimension, assuming odd number of points in each dimension
    coarsexres = ceil((Nx-1)/2)
    coarseyres = ceil((Ny-1)/2)

    # initializing coarse grid values in each dimension
    xcoarse = zeros(Float64, coarsexres)
    ycoarse = zeros(Float64, coarseyres)

    # assuming dimensionality of both are the same?
    xcoarse[:] = Interpolation * x
    ycoarse[:] = Interpolation * y

    return xcoarse, ycoarse
end

function project_to_coarse(A::SparseMatrixCSC, Nx::Int32, Ny::Int32, R::Matrix{Float64}, I::Matrix{Float64})
    # projects the matrix from the matrix equation Ax = b onto the coarse grid

    # acquiring restriction and interpolation matrices
    # R, I, xres, yres = transformation_matrices(Nx, Ny)

    # using formula for projection of matrix onto coarse grid
    Acoarse = R * A * I

    return Acoarse
end

function correction_matrix(A::SparseMatrixCSC, Nx::Int32, Ny::Int32)
    # calculates the error correction matrix s

    # acquiring restriction and interpolation matrices
    R, I, xres, yres = transformation_matrices(Nx, Ny)

    # carrying out explicit calculation of s
    S = I * inv(R * A * I) * R * A

    return S
end

# NOTE THAT THE x I AM USING HERE IS THE ACTUAL SOLUTION AND NOT THE GUESSED/KNOWN SOLUTION
# Also the x above is x as in the coordinate, below the x is the solution to the matrix equation Ax = b

function residual(A::SparseMatrixCSC, x::Vector{Float64}, b::Vector{Float64})
    # calculates the residual associated with the vector equation we're trying to solve

    # definition of residual on fine grid
    rfine = b - A * x

    return rfine
end

function residual(A::Matrix{Float64}, x::Vector{Float64}, b::Vector{Float64})
    # calculates the residual associated with the vector equation we're trying to solve

    # definition of residual on fine grid
    rfine = b - A * x

    return rfine
end

function restricted_residual(A::SparseMatrixCSC, x::Vector{Float64}, b::Vector{Float64}, Nx::Int32, Ny::Int32)
    # calculates the residual on the coarse grid

    # getting restriction/interpolation matrices on the given grid
    R, I, xres, yres = transformation_matrices(Nx::Int32, Ny::Int32)

    # obtaining residual on fine grid
    rfine = residual(A, x, b)

    # calculation of restriction residual
    rcoarse = R * rfine

    return rcoarse
end

function multigrid_error(e::Vector{Float64}, A::SparseMatrixCSC, Nx::Int32, Ny::Int32)
    # given the error e = x(guessed) - x(calculated), this gives the multigrid error correction

    # calculation of correction matrix S that maps from coarse to fine grid
    S = correction_matrix(A, Nx, Ny)

    # definition of multigrid correction given true fine grid error
    E = S * e

    return E
end

# Need to solve Acoarse = Ecoarse * rcoarse, then interpolate back and add Efine to x(calculated)

# may not need y vector - check later if issue doesn't get resolved
function multigrid_step(grid::Grid, A::SparseMatrixCSC, xcalc::Vector{Float64}, xguess::Vector{Float64}, b::Vector{Float64})
    # steps through one step of multigrid from fine to coarse and back according to the steps in the MIT page

    Nx = grid.Nx
    Ny = grid.Ny

    R1, I1, Nx1, Ny1 = transformation_matrices(Nx, Ny)
    R2, I2, Nx2, Ny2 = transformation_matrices(Nx1, Ny1)
    R3, I3, Nx3, Ny3 = transformation_matrices(Nx2, Ny2)

    rcoarse1 = restricted_residual(A, xcalc, b, Nx, Ny)
    Acoarse1 = project_to_coarse(A, Nx, Ny, R1, I1)
    xcalc1 = R1 * xcalc
    b1 = R1 * b

    rcoarse2 = restricted_residual(Acoarse1, xcalc1, b1, Nx1, Ny1)
    Acoarse2 = project_to_coarse(Acoarse1, Nx1, Ny1, R2, I2)
    xcalc2 = R2 * xcalc
    b2 = R2 * b

    rcoarse3 = restricted_residual(Acoarse2, xcalc2, b2, Nx2, Ny2)
    Acoarse3 = project_to_coarse(Acoarse2, Nx2, Ny2, R3, I3)
    xcalc3 = R3 * xcalc
    b3 = R3 * b

    # actual calculation - use a vcycle for example?
    # iterate at each step, carry over E into next grid and use it as guess for iteration on that grid
    # potentially start with exact solution at lowest grid level
    # for example
    Ecoarse3 = Acoarse3 \ rcoarse3
    Ecoarse2 = I3 * Ecoarse3
    Ecoarse2_1 = p_jacobi(Acoarse2, Ecoarse2, rcoarse2, 1, 1, 1, .1)
    Ecoarse1 = I2 * Ecoarse2_1
    Ecoarse1_1 = p_jacobi(Acoarse1, Ecoarse1, rcoarse1, 1, 1, 1, .1)
    Efine = I1 * Ecoarse1_1

    xcalc = xcalc + Efine

    return xcalc # may want to return more depending on the information we want
end

function v_cycle(grid::Grid, A::SparseMatrixCSC, xcalc::Vector{Float64}, b::Vector{Float64})
    # steps through one step of multigrid from fine to coarse and back according to the steps in the MIT page

    Nx = grid.Nx
    Ny = grid.Ny

    R1, I1, Nx1, Ny1 = transformation_matrices(Nx, Ny)
    R2, I2, Nx2, Ny2 = transformation_matrices(Nx1, Ny1)
    R3, I3, Nx3, Ny3 = transformation_matrices(Nx2, Ny2)

    rcoarse1 = restricted_residual(A, xcalc, b, Nx, Ny)
    Acoarse1 = project_to_coarse(A, Nx, Ny, R1, I1)
    xcalc1 = R1 * xcalc
    b1 = R1 * b

    rcoarse2 = restricted_residual(Acoarse1, xcalc1, b1, Nx1, Ny1)
    Acoarse2 = project_to_coarse(Acoarse1, Nx1, Ny1, R2, I2)
    xcalc2 = R2 * xcalc1
    b2 = R2 * b1

    rcoarse3 = restricted_residual(Acoarse2, xcalc2, b2, Nx2, Ny2)
    Acoarse3 = project_to_coarse(Acoarse2, Nx2, Ny2, R3, I3)
    xcalc3 = R3 * xcalc2
    b3 = R3 * b2

    Einitial = zeros(Float64, size(xcalc))
    Ecoarse1 = R1 * Einitial
    Ecoarse1_1 = jacobi_smooth(Acoarse1, Ecoarse1, rcoarse1, 1, 1, 1, 1, .1)
    Ecoarse2 = R2 * Ecoarse1_1
    Ecoarse2_1 = jacobi_smooth(Acoarse2, Ecoarse2, rcoarse2, 1, 1, 1, 1, .1)
    Ecoarse3 = R3 * Ecoarse2_1
    Ecoarse3_1 = jacobi_smooth(Acoarse3, Ecoarse3, rcoarse3, 1, 1, 1, 1, .1)
    Ecoarse2_2 = I3 * Ecoarse3_1
    Ecoarse2_3 = jacobi_smooth(Acoarse2, Ecoarse2_2, rcoarse2, 1, 1, 1, 1, .1)
    Ecoarse1_2 = I2 * Ecoarse2_3
    Ecoarse1_3 = jacobi_smooth(Acoarse1, Ecoarse1_2, rcoarse1, 1, 1, 1, 1, .1)
    Efine = I1 * Ecoarse1_3

    xcalc = xcalc + Efine

    return xcalc # may want to return more depending on the information we want
end

function v_cycle_general(grids::Diffresgrids, A::SparseMatrixCSC, x::Vector{Float64}, b::Vector{Float64})
    # doing a v-cycle in the formatting of the rest of the code (no restricted grids)

    # loading the restriction/interpolation matrices into concise variables
    R1 = grids.R1
    R2 = grids.R2
    R3 = grids.R3
    I1 = grids.I1
    I2 = grids.I2
    I3 = grids.I3
    Nx = grids.Nx
    Ny = grids.Ny
    Nx1 = grids.Nx1
    Ny1 = grids.Ny1
    Nx2 = grids.Nx2
    Ny2 = grids.Ny2
    Nx3 = grids.Nx3
    Ny3 = grids.Ny3

    # calculating the form of the matrix/residual at different resolutions
    """- error is calculated through smoothing (using given error and residual)
    - residual is calculated using error and previous residual as well as spacing
    - residual is restricted
    - error is set to 0
    IF NOT AT LOWEST RESOLUTION
    - run cycle again using the new error/residual to get error
    IF AT LOWEST RESOLUTION
    - error is given by smoothing function at lowest resolution"""

    Acoarse1 = sparse(project_to_coarse(A, Nx, Ny, R1, I1))
    x1 = R1 * x
    b1 = R1 * b
    r = residual(Acoarse1, x1, b1)
    Einitial1 = zeros(Float64, size(x1))
    Ecoarse1 = jacobi_smooth(Acoarse1, Einitial1, r, 1., 1, 1, 1)
    rcoarse1 = Acoarse1 * Ecoarse1
    rcoarse2 = R2 * rcoarse1

    Acoarse2 = sparse(project_to_coarse(Acoarse1, Nx1, Ny1, R2, I2))
    x2 = R2 * x1
    b2 = R2 * b1
    Einitial2 = zeros(Float64, size(x2))
    Ecoarse2 = jacobi_smooth(Acoarse2, Einitial2, rcoarse2, 1., 1, 1, 1)
    rcoarse2_1 = Acoarse2 * Ecoarse2
    rcoarse3 = R3 * rcoarse2_1

    Acoarse3 = sparse(project_to_coarse(Acoarse2, Nx2, Ny2, R3, I3))
    x3 = R3 * x2
    b3 = R3 * b2
    Einitial3 = zeros(Float64, size(x3))
    Ecoarse3 = jacobi_smooth(Acoarse3, Einitial3, rcoarse3, 1., 1, 1, 1) # maybe want to do exact solution at lowest res?
    rcoarse3_1 = Acoarse3 * Ecoarse3
    rcoarse2_2 = I3 * rcoarse3_1

    Ecoarse2_1 = jacobi_smooth(Acoarse2, Einitial2, rcoarse2_2, 1., 1, 1, 1)
    rcoarse2_3 = Acoarse2 * Ecoarse2_1
    rcoarse1_2 = I2 * rcoarse2_3

    Ecoarse1_1 = jacobi_smooth(Acoarse1, Einitial1, rcoarse1_2, 1., 1, 1, 1)
    rcoarse1_3 = Acoarse1 * Ecoarse1_1 # these two calculations may not be necessary
    rfine = I1 * rcoarse1_3
    Efine = Ecoarse1_1

    # actual v-cycle steps
    """Einitial = zeros({Float64}, size(x))
    Ecoarse1 = R1 * Einitial
    Ecoarse1_1 = jacobi_smooth(Acoarse1, Ecoarse1, rcoarse1, 1, 1, 1, 1, .1)
    Ecoarse2 = R2 * Ecoarse1_1
    Ecoarse2_1 = jacobi_smooth(Acoarse2, Ecoarse2, rcoarse2, 1, 1, 1, 1, .1)
    Ecoarse3 = R3 * Ecoarse2_1
    Ecoarse3_1 = jacobi_smooth(Acoarse3, Ecoarse3, rcoarse3, 1, 1, 1, 1, .1)
    Ecoarse2_2 = I3 * Ecoarse3_1
    Ecoarse2_3 = jacobi_smooth(Acoarse2, Ecoarse2_2, rcoarse2, 1, 1, 1, 1, .1)
    Ecoarse1_2 = I2 * Ecoarse2_3
    Ecoarse1_3 = jacobi_smooth(Acoarse1, Ecoarse1_2, rcoarse1, 1, 1, 1, 1, .1)
    Efine = I1 * Ecoarse1_3"""

    x = x + Efine

    return x, Efine
end


module IPMeasures

using StatsBase
using LinearAlgebra
using Distances
using CUDA

using Distributions
using DistributionsAD
using ConditionalDists

const MetricOrFun = Union{PreMetric,Function}
const CuMatrix = CuArray{<:Float32,2}

"""
    samplecolumns(x::AbstractMatrix, n::Int)

Sample n columns from a matrix. Returns x if the matrix has less than n columns.
"""
function samplecolumns(x::AbstractMatrix, n::Int)
    (size(x,2) > n) ? x[:,sample(1:size(x,2), n, replace=false)] : x
end

"""
    split2(x::AbstractMatrix)

Splits a matrix into halves.
"""
function split2(x::AbstractMatrix)
	n = size(x,2)
	x[:,1:div(n,2)], x[:,div(n,2)+1:end]
end

"""
	pairwisel2(x,y)
	pairwisel2(x)

Calculates pairwise squared euclidean distances of the columns of `x` and `y`
or `x` and `x`. The dispatches for CuArrays are necessary until
https://github.com/JuliaStats/Distances.jl/pull/142 is merged.
"""
pairwisel2(x::Matrix, y::Matrix) = pairwise(SqEuclidean(), x, y, dims=2)
pairwisel2(x::CuMatrix, y::CuMatrix) =
    -2 .* x' * y .+ sum(x.^2, dims=1)' .+ sum(y.^2,dims=1)
pairwisel2(x::AbstractMatrix) = pairwisel2(x,x)

"""
	pairwisel2(x,y)
	pairwisel2(x)

Calculates pairwise cosine distances of the columns of `x` and `y` or
`x` and `x`. The dispatches for CuArrays are necessary until
https://github.com/JuliaStats/Distances.jl/pull/142 is merged.
"""
pairwisecos(x::Matrix, y::Matrix) = pairwise(CosineDist(), x, y, dims=2)
pairwisecos(x::CuMatrix, y::CuMatrix) = 1 .- (x' * y / norm(x) / norm(y))
pairwisecos(x::AbstractMatrix) = pairwisecos(x,x)

include("kernels.jl")
include("mmd.jl")
include("kl_divergence.jl")

"""
	null_distribution(k::AbstractKernel, x, n, l = div(size(x,2),2))

Estimates the null distribution of samples `x` from `n` random draws of subsets of size `l`
"""
null_distribution(k::AbstractKernel, x, n, l = div(size(x,2),2)) =
    _null_distribution(k, pairwisel2(x,x), n, l )

function _null_distribution(k, d, n, l)
	ns = size(d,2)
	map(1:n) do _
		idx = sample(1:ns, l, replace = false)
		cidx = setdiff(1:ns, idx)
		cidx = sample(cidx, min(l, length(cidx)), replace = false)
		mmdfromdist(k, d, idx, cidx)
	end
end

end # module


@testset "8 Special functions" begin
    include("8.1 Error functions.jl")
    include("8.10 Formal derivatives.jl")
    include("8.2 Fresnel integral functions.jl")
    include("8.3 Exponential integral functions.jl")
    include("8.4 Trig integral functions.jl")
    include("8.5 Hyperbolic integral functions.jl")
    include("8.6 Gamma functions.jl")
    include("8.7 Zeta function.jl")
    include("8.8 Polylogarithm function.jl")
    include("8.9 Product logarithm function.jl")
end


using ModelingToolkit
using OrdinaryDiffEq, DiffEqCallbacks
using BenchmarkTools
using LinearAlgebra
using DataFrames

@parameters t, I[1:3, 1:3], Iinv[1:3, 1:3], g[1:3], m, α
@variables q[1:4](t) x[1:3](t) ω[1:3](t) v[1:3](t)
@derivatives D'~t

function transition_matrix(q)
    q0, q1, q2, q3 = q
    P = [
        2*(q0^2 + q1^2) - 1 2*(q1*q2 - q0*q3)   2*(q1*q3 + q2*q0)
        2*(q1*q2 + q0*q3)   2*(q0^2 + q2^2) - 1 2*(q2*q3 - q1*q0)
        2*(q1*q3 - q2*q0)   2*(q2*q3 + q1*q0)   2*(q0^2 + q3^2) - 1
    ]
end

function quaternion_derivative(q, ω)
    a, b, c = ω
    Mₛ = [
        0 -a -b -c
        a  0  c -b
        b -c  0  a
        c  b -a  0
    ]
    1/2 * Mₛ * q
end

function dynamics(q, x, ω, v)
    P = transition_matrix(q)
    vcat(
        D.(q) .~ quaternion_derivative(q, ω),
        D.(x) .~ P * v, # world velocity
        D.(ω) .~ Iinv * cross(I * ω, ω),
        D.(v) .~ 1/m * (m * P' * g - α .* v) # body acceleration
    )
end

function postprocess(u) # dynamics_f does not work
    q, x, ω, v = u[1:4, end], u[5:7, end], u[8:10, end], u[11:13, end]
    dynamics_f(q, x, ω, v)
end

get_values(eqs) = getproperty.(getproperty.(eqs, :rhs), :value)
function extract_result(u)
    syms = [:q1, :q2, :q3, :q4, :x, :y, :z, :ωx, :ωy, :ωz, :vx, :vy, :vz]
    df = DataFrame(repeat([Float64], 13), syms)
    for i ∈ 1:size(u)[2]
        push!(df, Dict(syms .=> u[:, i]))
    end
    df
end

function quaternion_norm_callback(resid, u, p, t)
    resid[1] = @views norm(u[1:4])^2 - 1
    resid[2:13] .= 0
end

quaternion_manifold_proj = ManifoldProjection(quaternion_norm_callback)

u0 = vcat(
    q .=> [1., 0, 0, 0],
    x .=> 0.,
    ω .=> [0.1, 0.15, 0.2],
    v .=> [1., 0, 0],
)

p = vcat(
    (I .=> Diagonal([0.25, 0.2, 0.1]))...,
    # (Iinv .=> inv(I))...
    (Iinv .=> Diagonal([4, 5., 10.]))...,
    g .=> [0., 0., -9.81],
    m => 3.,
    α => 0.01
)

tspan = (0.0, 10000.0)
sys = ODESystem(dynamics(q, x, ω, v))
prob = ODEProblem(sys, u0, tspan, p)

u = solve(prob, Vern7(), save_everystep=false, callback=quaternion_manifold_proj)
df = extract_result(u)

# dynamics_f = build_function(dynamics(q, x, ω, v), q, x, ω, v; expression=Val{false})
# postprocess(u)


struct CMMatrix{T, C} <: AbstractMatrix{T} # M_r
    cc::Vector{C} # vector of conjugacy classes to fix the order
    r::Int # the index of conjugacy class
    m::Matrix{T} # cache of class coefficients

    function CMMatrix(cc::A, r::Int, T::Type=Int) where {C, A<:AbstractVector{C}}
        M = -ones(T, length(cc), length(cc))
        new{T, C}(cc, r, M)
    end
end

Base.size(M::CMMatrix) = size(M.m)
Base.IndexStyle(::Type{<:CMMatrix}) = IndexCartesian()

function Base.getindex(M::CMMatrix, s::Integer, t::Integer)
    if isone(-M.m[s,t])
        out = one(first(first(M.cc)))

        M.m[s,:] .= 0
        r = M.r

        for g in M.cc[r]
            for h in M.cc[s]
                out = PermutationGroups.mul!(out, g, h)
                for t in 1:size(M, 2)
                    if out == first(M.cc[t])
                        M.m[s, t] += 1
                        break
                    end
                end
            end
        end

    end
    return M.m[s,t]
end


using Base.Iterators, LinearAlgebra, Optim, Plotly;
import DelimitedFiles.readdlm;

NDimensions = Int64(4);
#Define the single qubit measurements, R is the state with "-i".
H = Float64[1, 0];
V = Float64[0, 1];
statevars = [H, V];

Id = Matrix{Float64}(I, 2, 2);
Hx = [1 1; 1 -1]/sqrt(2);
Hy = [1 1; im -im]/sqrt(2);
Trafovars = [Id, Hx, Hy];

Trafos = [kron(l,k) for k in Trafovars, l in Trafovars];
StateArray = [kron(l,k) for k in statevars, l in statevars];
StateMeasured = [];

for a in 1:length(Trafos)
    for b in 1:NDimensions
        push!(StateMeasured,Trafos[a]*StateArray[b])
    end
end

D2Array = readdlm("D:\\myJulia\\two_qubit_tomo_measurement_2.txt", skipstart=2)[1:end,34:37];
Data = vec(transpose(D2Array));
SigmaData = sqrt.(Data);

function BadnessPolynomial!(t)
    pM = [t[1] 0 0 0;
        t[5] + im*t[6]  t[2] 0 0;
        t[11] + im*t[12] t[7] + im*t[8] t[3] 0;
        t[15] + im*t[16] t[13] + im*t[14] t[9] + im*t[10] t[4]];

    pMd = conj(transpose(pM));

    global GeneralDM = pMd*pM;
    Prediction(DM, State) = sum(conj(transpose(State))*DM*State);

    BD = sum((Prediction(GeneralDM, StateMeasured[i]) - Data[i])^2/Prediction(GeneralDM, StateMeasured[i]) for i=1:length(StateMeasured));
    return real(BD)
end

Bell = (kron(H,V)-kron(V,H))/sqrt(2);
IdealDM = reshape(kron(Bell, conj(Bell)),(NDimensions,NDimensions))

t0 = vec(IdealDM);
@time res = optimize(BadnessPolynomial!, t0, NelderMead(), Optim.Options(g_tol = 10e-10, iterations = 20000));
#@time res = optimize(BadnessPolynomial!, t0, LBFGS(), Optim.Options(iterations = 20000)) #Requires a function and gradient (will be approximated if omitted)
#inner_optimizer = GradientDescent()
#lower = fill(-1.0, NDimensions);
#upper = fill(1.0, NDimensions);
#@time res = optimize(BadnessPolynomial!, lower, upper, t0, Fminbox(inner_optimizer), Optim.Options(iterations = 50000)) #Box constrained minimization
temp = Optim.minimizer(res);

MaxLikDM = GeneralDM/tr(GeneralDM)

Fidelity = tr(sqrt(sqrt(MaxLikDM)*IdealDM*sqrt(MaxLikDM)))^2;
Purity = tr(MaxLikDM*MaxLikDM);
LinearEntropy = NDimensions/(NDimensions - 1)*(1 - Purity);

function plotme3d(dta)
    x = [];
    y = [];
    re_z = [];
    im_z = [];
    for i = 1:size(dta)[1]
        for j = 1:size(dta)[1]
            push!(x, i);
            push!(y, j);
            push!(re_z, real(dta[i,j]));
            push!(im_z, imag(dta[i,j]));
        end
    end

    trace1 = scatter3d(x = x, y = y, z = re_z, mode = "markers", marker = attr(size = 5), name = "Real part");
    trace2 = scatter3d(x = x, y = y, z = im_z, mode = "markers", marker = attr(size = 5), name = "Imaginary part");
    layout = Layout(xaxis = attr(tickvals=[1,2,3,4], ticktext = ["HH", "HV", "VH", "VV "]),
            yaxis = attr(tickvals=[1,2,3,4], ticktext = ["HH", "HV", "VH", "VV "]));
    #plot([trace1, trace2])
    p1 = Plotly.plot(trace1, layout);
    p2 = Plotly.plot(trace2, layout);
    [p1, p2]
end

plotme3d(MaxLikDM)


# Albrecher, H., Mayer, P., & Schoutens, W. T. J. (2007)
# The Little Heston Trap
# Wilmott Magazine, January issue, 83-92.

function analyticinterval(m::Heston)
  # an approximation of the analytic interval is the
  # area between the roots of p(x) = ax^2+bx+c
  a = (1.0 - m.ρ^2) * m.σ^2
  b = 2.0 * m.ρ * m.σ * m.κ - m.σ^2
  c = -m.κ^2
  A = -0.5 * b / a
  B = sqrt(0.25 * b^2 / a^2 - c / a)
  A - B, A + B
end

function logmgf(m::Heston, x0, T, u)
  v0 = x0[2]
  σ2 = m.σ * m.σ
  a0 = m.ρ * m.σ * u - m.κ
  a1 = sqrt(a0 * a0 + σ2 * u * (1 - u))
  a2 = m.ρ * m.σ * u - m.κ
  a3 = -(a2 + a1)
  a4 = a3 / (a1 - a2)
  a5 = exp(-a1 * T)
  a6 = 1 - a4 * a5
  ((a3 * T - 2 * log(a6 / (1 - a4))) * m.θ * m.κ + v0 * a3 * (1 - a5) / a6) / σ2
end


using Test
using QuantumLattices.Mathematics.Combinatorics

@testset "Combinations" begin
    @test Combinations{0}(1:3)|>collect==[()]
    @test Combinations{1}(1:3)|>collect==[(1,),(2,),(3,)]
    @test Combinations{2}(1:3)|>collect==[(1,2),(1,3),(2,3)]
    @test Combinations{3}(1:3)|>collect==[(1,2,3)]
end

@testset "DulCombinations" begin
    @test DulCombinations{0}(1:3)|>collect==[()]
    @test DulCombinations{1}(1:3)|>collect==[(1,),(2,),(3,)]
    @test DulCombinations{2}(1:3)|>collect==[(1,1),(1,2),(1,3),(2,2),(2,3),(3,3)]
    @test DulCombinations{3}(1:3)|>collect==[(1,1,1),(1,1,2),(1,1,3),(1,2,2),(1,2,3),(1,3,3),(2,2,2),(2,2,3),(2,3,3),(3,3,3)]
end

@testset "Permutations" begin
    @test Permutations{0}(1:3)|>collect==[()]
    @test Permutations{1}(1:3)|>collect==[(1,),(2,),(3,)]
    @test Permutations{2}(1:3)|>collect==[(1,2),(1,3),(2,1),(2,3),(3,1),(3,2)]
    @test Permutations{3}(1:3)|>collect==[(1,2,3),(1,3,2),(2,1,3),(2,3,1),(3,1,2),(3,2,1)]
end

@testset "DulPermutations" begin
    @test DulPermutations{0}(1:3)|>collect==[()]
    @test DulPermutations{1}(1:3)|>collect==[(1,),(2,),(3,)]
    @test DulPermutations{2}(1:3)|>collect==[(1,1),(1,2),(1,3),(2,1),(2,2),(2,3),(3,1),(3,2),(3,3)]
    @test DulPermutations{3}(1:3)|>collect==[(i,j,k) for i=1:3 for j=1:3 for k=1:3]
end