"""
    FESetModule

Module for finite element sets.
"""
module FESetModule

__precompile__(true)

import Base.count
import Base.cat
import LinearAlgebra: mul!, Transpose
At_mul_B!(C, A, B) = mul!(C, Transpose(A), B)
import LinearAlgebra: norm, cross
using ..FTypesModule: FInt, FFlt, FCplxFlt, FFltVec, FIntVec, FFltMat, FIntMat, FMat, FVec, FDataDict

"""
    AbstractFESet{NODESPERELEM}

Abstract type of a finite element set. Parameterized with the number of of the nodes per element.
"""
abstract type AbstractFESet{NODESPERELEM} end
"""
    AbstractFESet0Manifold{NODESPERELEM} <: FESet{NODESPERELEM}

Abstract type of a finite element set for 0-dimensional manifolds (points).
Parameterized with the number of of the nodes per element.
"""
abstract type AbstractFESet0Manifold{NODESPERELEM} <: AbstractFESet{NODESPERELEM} end
"""
    AbstractFESet1Manifold{NODESPERELEM} <: FESet{NODESPERELEM}

Abstract type of a finite element set for 1-dimensional manifolds (curves).
Parameterized with the number of of the nodes per element.
"""
abstract type AbstractFESet1Manifold{NODESPERELEM} <: AbstractFESet{NODESPERELEM} end
"""
    AbstractFESet2Manifold{NODESPERELEM} <: FESet{NODESPERELEM}

Abstract type of a finite element set for 2-dimensional manifolds (surfaces).
Parameterized with the number of of the nodes per element.
"""
abstract type AbstractFESet2Manifold{NODESPERELEM} <: AbstractFESet{NODESPERELEM} end
"""
    AbstractFESet3Manifold{NODESPERELEM} <: FESet{NODESPERELEM}

Abstract type of a finite element set for 3-dimensional manifolds (solids).
Parameterized with the number of of the nodes per element.
"""
abstract type AbstractFESet3Manifold{NODESPERELEM} <: AbstractFESet{NODESPERELEM} end

"""
    add_FESet_fields(NODESPERELEM)

Generate standard fields for the finite element set.
"""
macro add_FESet_fields(NODESPERELEM)
    return esc(:(
    conn::Array{NTuple{$NODESPERELEM, FInt}, 1};
    label::FIntVec; )
    )
end
# show(macroexpand(:(@add_FESet_fields 6)))

"""
    define_FESet(NAME, MANIFOLD, NODESPERELEM)

Define the concrete type for a finite element set.
"""
macro define_FESet(NAME, MANIFOLD, NODESPERELEM)
    return esc(:(
        mutable struct $NAME <: $MANIFOLD{$NODESPERELEM}
            @add_FESet_fields $NODESPERELEM
            function $NAME(conn::FIntMat)
                self = new(NTuple{$NODESPERELEM, FInt}[], FInt[])
                self = fromarray!(self, conn)
                setlabel!(self, 0)
                return self
            end
        end
    ))
end
# show(macroexpand(:(@define_FESet FESetT10 AbstractFESet3Manifold 10)))

"""
    nodesperelem(fes::AbstractFESet{NODESPERELEM}) where {NODESPERELEM}

Provide the number of nodes per element.
"""
nodesperelem(fes::AbstractFESet{NODESPERELEM}) where {NODESPERELEM} = NODESPERELEM

"""
    nodesperelem(::Type{T}) where {NODESPERELEM, T<:AbstractFESet{NODESPERELEM}}

Provide the number of nodes per element for a given type.
"""
nodesperelem(::Type{T}) where {NODESPERELEM, T<:AbstractFESet{NODESPERELEM}} = NODESPERELEM

"""
    manifdim(me)

Get the manifold dimension.
"""
manifdim(me::AbstractFESet0Manifold{NODESPERELEM}) where {NODESPERELEM} = 0
manifdim(me::AbstractFESet1Manifold{NODESPERELEM}) where {NODESPERELEM} = 1
manifdim(me::AbstractFESet2Manifold{NODESPERELEM}) where {NODESPERELEM} = 2
manifdim(me::AbstractFESet3Manifold{NODESPERELEM}) where {NODESPERELEM} = 3

"""
    count(self::T)::FInt where {T<:AbstractFESet}

Get the number of individual connectivities in the FE set.
"""
count(self::T) where {T<:AbstractFESet} = length(self.conn)

"""
    fromarray!(self::AbstractFESet{NODESPERELEM}, conn::FIntMat) where {NODESPERELEM}

Set  the connectivity from an integer array.
"""
function fromarray!(self::AbstractFESet{NODESPERELEM}, conn::FIntMat) where {NODESPERELEM}
    @assert size(conn, 2) == NODESPERELEM
    self.conn = fill(tuple(fill(0, NODESPERELEM)...), size(conn, 1));
    for i = 1:length(self.conn)
        self.conn[i] = ntuple(y -> conn[i, y], NODESPERELEM);
    end
    return self
end

"""
    connasarray(self::AbstractFESet{NODESPERELEM}) where {NODESPERELEM}

Return the connectivity as an array.

Return the connectivity as an integer array (matrix), where the number of rows
matches the number of connectivities in the set.
"""
function connasarray(self::AbstractFESet{NODESPERELEM}) where {NODESPERELEM}
    conn = fill(zero(FInt), length(self.conn), NODESPERELEM)
    for i = 1:size(conn, 1)
        for j = 1:NODESPERELEM
            conn[i, j] = self.conn[i][j]
        end
    end
    return conn
end

"""
    boundaryconn(self::T) where {T<:AbstractFESet}

Get boundary connectivity.
"""
function boundaryconn(self::T) where {T<:AbstractFESet}
    return privboundaryconn(self)
end

"""
    boundaryfe(self::T) where {T<:AbstractFESet}

Return the constructor of the type of the boundary finite element.
"""
function boundaryfe(self::T) where {T<:AbstractFESet}
    return privboundaryfe(self)
end

"""
    bfun(self::T,  param_coords::FFltVec)::FFltMat where {T<:AbstractFESet}

Compute the values of the basis functions.

Compute the values of the basis functions at a given parametric coordinate.
One basis function per row.
"""
function bfun(self::T,  param_coords::FFltVec)::FFltMat where {T<:AbstractFESet}
    return privbfun(self, param_coords)
end

"""
    bfundpar(self::T,  param_coords::FFltVec)::FFltMat where {T<:AbstractFESet}

Compute the values of the basis function gradients.

Compute the values of the basis function gradients with respect to the
parametric coordinates at a given parametric coordinate. One basis function
gradients per row.
"""
function bfundpar(self::T,  param_coords::FFltVec)::FFltMat where {T<:AbstractFESet}
    return privbfundpar(self, param_coords)
end

"""
    setlabel!(self::T, val::FInt) where {T<:AbstractFESet}

Set the label of the entire finite elements set.

All elements are labeled with this number.
"""
function setlabel!(self::T, val::FInt) where {T<:AbstractFESet}
    self.label = zeros(FInt, size(self.conn, 1));
    fill!(self.label, val)
    return self
end

"""
    setlabel!(self::T, val::FIntVec) where {T<:AbstractFESet}

Set the labels of individual elements.
"""
function setlabel!(self::T, val::FIntVec) where {T<:AbstractFESet}
    #    Set the label of this set.
    @assert size(self.conn, 1)==length(val) "Must get one  label per finite element connectivity"
    self.label=zeros(FInt, size(self.conn, 1));
    copyto!(self.label, val);
    return self
end

"""
    subset(self::T, L::FIntVec) where {T<:AbstractFESet}

Extract a subset of the finite elements from the given finite element set.
"""
function subset(self::T, L::FIntVec) where {T<:AbstractFESet}
    result = deepcopy(self)
    result.conn = deepcopy(self.conn[L])
    result.label = deepcopy(self.label[L])
    return  result
end

"""
    cat(self::T,  other::T) where {T<:AbstractFESet}

Concatenate the connectivities of two FE sets.
"""
function cat(self::T,  other::T) where {T<:AbstractFESet}
    @assert nodesperelem(self) == nodesperelem(other)
    result =deepcopy(self)
    result.conn = vcat(self.conn, other.conn);
    setlabel!(result, vcat(self.label, other.label))
    return result
end

"""
    updateconn!(self::T, newids::FIntVec) where {T<:AbstractFESet}

Update the connectivity after the IDs of nodes changed.

`newids`= new node IDs. Note that indexes in the conn array "point"
_into_ the  `newids` array. After the connectivity was updated
this will no longer be true!
"""
function updateconn!(self::T, newids::FIntVec) where {T<:AbstractFESet}
    conn = connasarray(self)
    for i=1:size(conn, 1)
        for j=1:size(conn, 2)
            conn[i, j]=newids[conn[i, j]];
        end
    end
    return fromarray!(self, conn)
end

"""
    inparametric(self::AbstractFESet, param_coords::FFltVec)

Are given parametric coordinates inside the element parametric domain?

Return a Boolean: is the point inside, true or false?
"""
function inparametric(self::T, param_coords::FFltVec; tolerance = 0.0) where {T<:AbstractFESet}
    return privinparametric(self, param_coords, tolerance)
end

"""
    map2parametric(self::T, x::FFltMat, pt::FFltVec;
        tolerance = 0.001, maxiter =5) where {T<:AbstractFESet}

Map a spatial location to parametric coordinates.

- `x`=array of spatial coordinates of the nodes, size(x) = nbfuns x dim,
- `c`= spatial location
- `tolerance` = tolerance in parametric coordinates; default is 0.001.

# Return
- `success` = Boolean flag, true if successful, false otherwise.
- `pc` = Returns a row array of parametric coordinates if the solution was
  successful, otherwise NaN are returned.
"""
function map2parametric(self::T, x::FFltMat, pt::FFltVec; tolerance = 0.001, maxiter =5) where {T<:AbstractFESet}
    sdim = size(x, 2); # number of space dimensions
    mdim = manifdim(self); # manifold dimension of the element
    ppc = zeros(mdim);
    pc = deepcopy(ppc);
    J = [i==j ? one(FFlt) : zero(FFlt) for i=1:sdim, j=1:mdim] # Jac. matx: "identity"
    loc = fill(zero(FFlt), (1, sdim)) # point location buffer
    success = true
    for i = 1:maxiter
        gradNparams = privbfundpar(self, pc);
        At_mul_B!(J, x, gradNparams); # calculate the Jacobian matrix
        N = privbfun(self, pc);
        At_mul_B!(loc, N, x);# Iterated point location
        pc[:] = ppc .- (J\(vec(loc) - pt))
        if (norm(pc-ppc) < tolerance)
            return pc, success;
        end
        ppc[:] = pc[:];
    end
    success = false
    return pc, success;
end

"""
    centroidparametric(self::T) where {T<:AbstractFESet}

Return the parametric coordinates  of the centroid of the element.
"""
function centroidparametric(self::T) where {T<:AbstractFESet}
    return privcentroidparametric(self)
end

"""
    Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet0Manifold}

Evaluate the point Jacobian.

- `J` = Jacobian matrix, columns are tangent to parametric coordinates curves.
"""
function Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet0Manifold}
    return 1.0::FFlt;
end

"""
    Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet1Manifold}

Evaluate the curve Jacobian.

- `J` = Jacobian matrix, columns are tangent to parametric coordinates curves.
"""
function Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet1Manifold}
    sdim,  ntan = size(J);
    @assert ntan == 1 "Expected number of tangent vectors: 1"
    return norm(J)::FFlt;
end

"""
    gradN!(self::AbstractFESet1Manifold, gradN::FFltMat, gradNparams::FFltMat,
      redJ::FFltMat)

Compute the gradient of the basis functions with the respect to
the "reduced" spatial coordinates.

- `gradN`= output,  matrix of gradients,  one per row
- `gradNparams`= matrix of gradients with respect to parametric coordinates, one per row
- `redJ`= reduced Jacobian matrix `redJ=transpose(Rm)*J`
"""
function gradN!(self::AbstractFESet1Manifold, gradN::FFltMat, gradNparams::FFltMat, redJ::FFltMat)
    for r=1:size(gradN, 1)
        gradN[r, 1]=  gradNparams[r, 1]/redJ[1, 1]
    end
end

"""
    Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet2Manifold}

Evaluate the curve Jacobian.

- `J` = Jacobian matrix, columns are tangent to parametric coordinates curves.
"""
function Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet2Manifold}
    sdim,  ntan = size(J);
    @assert ntan == 2 "Expected number of tangent vectors: 2"
    if sdim == ntan
        @inbounds Jac = (J[1, 1]*J[2, 2] - J[2, 1]*J[1, 2])
        return Jac::FFlt;# is det(J);% Compute the Jacobian
    else
        return norm(cross(J[:, 1], J[:, 2]))::FFlt;
    end
end

"""
    gradN!(self::AbstractFESet2Manifold, gradN::FFltMat, gradNparams::FFltMat,
      redJ::FFltMat)

Compute the gradient of the basis functions with the respect to
the "reduced" spatial coordinates.

- `gradN`= output,  matrix of gradients,  one per row
- `gradNparams`= matrix of gradients with respect to parametric coordinates,
  one per row
- `redJ`= reduced Jacobian matrix `redJ=transpose(Rm)*J`
"""
function gradN!(self::AbstractFESet2Manifold, gradN::FFltMat, gradNparams::FFltMat, redJ::FFltMat)
    # This is the unrolled version that avoids allocation of a 3 x 3 matrix
    invdet=1.0/(redJ[1, 1]*redJ[2, 2] - redJ[1, 2]*redJ[2, 1]);
    invredJ11 =  (redJ[2, 2])*invdet;
    invredJ12 = -(redJ[1, 2])*invdet;
    invredJ21 = -(redJ[2, 1])*invdet;
    invredJ22 =  (redJ[1, 1])*invdet;
    @assert size(gradN, 1)==size(gradNparams, 1)
    @inbounds for r=1:size(gradN, 1)
        gradN[r, 1]= gradNparams[r, 1]*invredJ11 +gradNparams[r, 2]*invredJ21;
        gradN[r, 2]= gradNparams[r, 1]*invredJ12 +gradNparams[r, 2]*invredJ22;
    end
end

"""
    Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet3Manifold}

Evaluate the volume Jacobian.

`J` = Jacobian matrix, columns are tangent to parametric coordinates curves.
"""
function Jacobian(self::T, J::FFltMat)::FFlt where {T<:AbstractFESet3Manifold}
    sdim,  ntan = size(J);
    @assert (ntan == 3) && (sdim == 3) "Expected number of tangent vectors: 3"
    #Jac = det(J);# Compute the Jacobian
    # The unrolled version
    return (+J[1, 1]*(J[2, 2]*J[3, 3]-J[3, 2]*J[2, 3])
    -J[1, 2]*(J[2, 1]*J[3, 3]-J[2, 3]*J[3, 1])
    +J[1, 3]*(J[2, 1]*J[3, 2]-J[2, 2]*J[3, 1]) )::FFlt;
end

"""
    gradN!(self::AbstractFESet3Manifold, gradN::FFltMat, gradNparams::FFltMat,
      redJ::FFltMat)

Compute the gradient of the basis functions with the respect to
the "reduced" spatial coordinates.

- `gradN`= output,  matrix of gradients,  one per row
- `gradNparams`= matrix of gradients with respect to parametric coordinates,
  one per row
- `redJ`= reduced Jacobian matrix `redJ=transpose(Rm)*J`
"""
function gradN!(self::AbstractFESet3Manifold, gradN::FFltMat, gradNparams::FFltMat, redJ::FFltMat)
    invdet = 1.0 / ( +redJ[1, 1]*(redJ[2, 2]*redJ[3, 3]-redJ[3, 2]*redJ[2, 3])
                    -redJ[1, 2]*(redJ[2, 1]*redJ[3, 3]-redJ[2, 3]*redJ[3, 1])
                    +redJ[1, 3]*(redJ[2, 1]*redJ[3, 2]-redJ[2, 2]*redJ[3, 1]) );
    # This is the unrolled version that avoids allocation of a 3 x 3 matrix
    invredJ11 =  (redJ[2, 2]*redJ[3, 3]-redJ[3, 2]*redJ[2, 3])*invdet;
    invredJ12 = -(redJ[1, 2]*redJ[3, 3]-redJ[1, 3]*redJ[3, 2])*invdet;
    invredJ13 =  (redJ[1, 2]*redJ[2, 3]-redJ[1, 3]*redJ[2, 2])*invdet;
    invredJ21 = -(redJ[2, 1]*redJ[3, 3]-redJ[2, 3]*redJ[3, 1])*invdet;
    invredJ22 =  (redJ[1, 1]*redJ[3, 3]-redJ[1, 3]*redJ[3, 1])*invdet;
    invredJ23 = -(redJ[1, 1]*redJ[2, 3]-redJ[2, 1]*redJ[1, 3])*invdet;
    invredJ31 =  (redJ[2, 1]*redJ[3, 2]-redJ[3, 1]*redJ[2, 2])*invdet;
    invredJ32 = -(redJ[1, 1]*redJ[3, 2]-redJ[3, 1]*redJ[1, 2])*invdet;
    invredJ33 =  (redJ[1, 1]*redJ[2, 2]-redJ[2, 1]*redJ[1, 2])*invdet;
    @assert size(gradN, 1)==size(gradNparams, 1)
    @inbounds for r=1:size(gradN, 1)
        gradN[r, 1]= gradNparams[r, 1]*invredJ11 +gradNparams[r, 2]*invredJ21 +gradNparams[r, 3]*invredJ31;
        gradN[r, 2]= gradNparams[r, 1]*invredJ12 +gradNparams[r, 2]*invredJ22 +gradNparams[r, 3]*invredJ32;
        gradN[r, 3]= gradNparams[r, 1]*invredJ13 +gradNparams[r, 2]*invredJ23 +gradNparams[r, 3]*invredJ33;
    end
end

################################################################################
################################################################################

"""
    FESetP1

Type for sets of point-like of finite elements.
"""
@define_FESet FESetP1 AbstractFESet0Manifold 1

privbfun(self::FESetP1,  param_coords::FFltVec) = reshape([1.0], 1, 1) # make sure this is a matrix
privbfundpar(self::FESetP1,  param_coords::FFltVec) = zeros(1, 0)

function privinparametric(self::FESetP1, param_coords::FFltVec, tolerance::FFlt)
    return  param_coords[1] == 0.0;
end

function privcentroidparametric(self::FESetP1)
    return vec([0.0])
end

function map2parametric(self::FESetP1, x::FFltMat, pt::FFltVec;
    tolerance = 0.001, maxiter =5)
    success = false; pc = [0.0]
    if norm(vec(x) - pt) < tolerance
        success = true
    end
    return pc, success;
end

################################################################################
################################################################################

"""
    FESetL2

Type for sets of curve-like finite elements with two nodes.
"""
@define_FESet FESetL2 AbstractFESet1Manifold 2

privbfun(self::FESetL2,  param_coords::FFltVec) = reshape([(1. - param_coords[1]); (1. + param_coords[1])] / 2.0, 2, 1) # make sure this is a matrix
privbfundpar(self::FESetL2,  param_coords::FFltVec) = reshape([-1.0; +1.0]/2.0, 2, 1)

function privboundaryconn(self::FESetL2)
    conn = connasarray(self)
    return [conn[:, 1]; conn[:, 2]];
end

function privboundaryfe(self::FESetL2)
    return FESetP1;
end

function privinparametric(self::FESetL2, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetL2)
    return vec([0.0])
end

################################################################################
################################################################################

"""
    FESetL3

Type for sets of curve-like of finite elements with three nodes.
"""
@define_FESet FESetL3 AbstractFESet1Manifold 3

function privbfun(self::FESetL3,  param_coords::FFltVec)
    xi=param_coords[1];
    val = [(xi-1)*xi/2;  (xi+1)*xi/2;  -(xi+1)*(xi-1)];
    return reshape(val, 3, 1) # make sure this is a matrix
end

function privbfundpar(self::FESetL3,  param_coords::FFltVec)
    xi=param_coords[1];
    val = [(xi-1/2); (xi+1/2); -2*xi];
    return reshape(val, 3, 1) # make sure this is a matrix
end

function privboundaryconn(self::FESetL3)
    conn = connasarray(self)
    return [conn[:, 1]; conn[:, 2]];
end

function privboundaryfe(self::FESetL3)
    return FESetP1;
end

function privinparametric(self::FESetL3, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetL3)
    return vec([0.0])
end

################################################################################
################################################################################

"""
    FESetT3

Type for sets of surface-like triangular finite elements with three nodes.
"""
@define_FESet FESetT3 AbstractFESet2Manifold 3

function privbfun(self::FESetT3,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 3-node triangle.
    return reshape([(1 - param_coords[1] - param_coords[2]); param_coords[1]; param_coords[2]], 3, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetT3,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    return  [-1. -1.;  +1.  0.;  0. +1.];
end

function privboundaryconn(self::FESetT3)
    conn = connasarray(self)
    return [conn[:, 1:2]; conn[:, 2:3]; conn[:, [3, 1]]];
end

function privboundaryfe(self::FESetT3)
    # Get  the constructor of the class of the  boundary finite element.
    return FESetL2;
end

function privinparametric(self::FESetT3, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords)) &&
        (-3.0 * tolerance <= sum(param_coords) <= 1.0 + 3.0 * tolerance);
end

function privcentroidparametric(self::FESetT3)
    return vec([1.0/3 1.0/3])
end

################################################################################
################################################################################

"""
    FESetQ4

Type for sets of surface-like quadrilateral finite elements with four nodes.
"""
@define_FESet FESetQ4 AbstractFESet2Manifold 4

function privbfun(self::FESetQ4,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 4-node quadrilateral.
    val = [0.25 * (1. - param_coords[1]) * (1. - param_coords[2]);
           0.25 * (1. + param_coords[1]) * (1. - param_coords[2]);
           0.25 * (1. + param_coords[1]) * (1. + param_coords[2]);
           0.25 * (1. - param_coords[1]) * (1. + param_coords[2])];
    return reshape(val, 4, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetQ4,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    val = [-(1. - param_coords[2])*0.25 -(1. - param_coords[1])*0.25;
            (1. - param_coords[2])*0.25 -(1. + param_coords[1])*0.25;
            (1. + param_coords[2])*0.25 (1. + param_coords[1])*0.25;
           -(1. + param_coords[2])*0.25 (1. - param_coords[1])*0.25];
    return val
end

function privboundaryconn(self::FESetQ4)
    conn = connasarray(self)
    return [conn[:, 1:2]; conn[:, 2:3]; conn[:, 3:4]; conn[:, [4, 1]]];
end

function privboundaryfe(self::FESetQ4)
    # Get  the constructor of the class of the  boundary finite element.
    return FESetL2;
end

function privinparametric(self::FESetQ4, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetQ4)
    return vec([0.0 0.0])
end

################################################################################
################################################################################

"""
    FESetQ9

Type for sets of surface-like quadrilateral finite elements with nine nodes.
"""
@define_FESet FESetQ9 AbstractFESet2Manifold 9

function privbfun(self::FESetQ9,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 9-node quadrilateral.
    xi=param_coords[1];
    xis = [(xi-1)*xi/2;  (xi+1)*xi/2;  -(xi+1)*(xi-1)];
    eta=param_coords[2];
    etas = [(eta-1)*eta/2;  (eta+1)*eta/2;  -(eta+1)*(eta-1)];
    xisetas=(xis*etas');
    val = xisetas[[1     2     5     4     3     8     6     7     9]];
    return reshape(val, 9, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetQ9,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    xi=param_coords[1];
    xis = [(xi-1)*xi/2;  (xi+1)*xi/2;  -(xi+1)*(xi-1)];
    dxis = [(xi-1/2); (xi+1/2); -2*xi];
    eta=param_coords[2];
    etas = [(eta-1)*eta/2;  (eta+1)*eta/2;  -(eta+1)*(eta-1)];
    detas=[(eta-1/2); (eta+1/2); -2*eta];
    dxisetas=(dxis*etas');
    xisdetas=(xis*detas');
    ix=[1     2     5     4     3     8     6     7     9]
    val =  [dxisetas[ix]'             xisdetas[ix]'];
    return reshape(val, 9, 2); # Make sure this is a matrix
end

function privboundaryconn(self::FESetQ9)
    conn = connasarray(self)
    return [conn[:, [1, 2, 5]]; conn[:, [2, 3, 6]]; conn[:, [3, 4, 7]]; conn[:, [4, 1, 8]]];
end

function privboundaryfe(self::FESetQ9)
    return FESetL3;
end

function privinparametric(self::FESetQ9, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetQ9)
    return vec([0.0 0.0])
end

################################################################################
################################################################################

"""
    FESetQ8

Type for sets of surface-like quadrilateral finite elements with eight nodes.
"""
@define_FESet FESetQ8 AbstractFESet2Manifold 8

function privbfun(self::FESetQ8,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 4-node quadrilateral.
    xim    = (-1 + param_coords[1]);
    etam   = (-1 + param_coords[2]);
    xip    = (1 + param_coords[1]);
    etap   = (1 + param_coords[2]);
    val = [ -1.0/4*xim*etam*(1+param_coords[1]+param_coords[2]);
           1.0/4*xip*etam*(1-param_coords[1]+param_coords[2]);
           -1.0/4*xip*etap*(1-param_coords[1]-param_coords[2]);
           1.0/4*xim*etap*(1+param_coords[1]-param_coords[2]);
           1.0/2*xim*xip*etam;
           -1.0/2*etam*etap*xip;
           -1.0/2*xim*xip*etap;
           1.0/2*etam*etap*xim];
    return reshape(val, 8, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetQ8,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    xi =param_coords[1];
    eta =param_coords[2];
    val=zeros(FFlt, 8, 2)
    val[:, 1]= [1.0/4*(1-eta)*(1+xi+eta)-1.0/4*(1-xi)*(1-eta);
               -1.0/4*(1-eta)*(1-xi+eta)+1.0/4*(1+xi)*(1-eta);
               -1.0/4*(1+eta)*(1-xi-eta)+1.0/4*(1+xi)*(1+eta);
               1.0/4*(1+eta)*(1+xi-eta)-1.0/4*(1-xi)*(1+eta);
               -1.0/2*(1+xi)*(1-eta)+1.0/2*(1-xi)*(1-eta);
               1.0/2*(1-eta)*(1+eta);
               -1.0/2*(1+xi)*(1+eta)+1.0/2*(1-xi)*(1+eta);
               -1.0/2*(1-eta)*(1+eta)];
    val[:, 2] = [1.0/4*(1-xi)*(1+xi+eta)-1.0/4*(1-xi)*(1-eta);
                1.0/4*(1+xi)*(1-xi+eta)-1.0/4*(1+xi)*(1-eta);
                -1.0/4*(1+xi)*(1-xi-eta)+1.0/4*(1+xi)*(1+eta);
                -1.0/4*(1-xi)*(1+xi-eta)+1.0/4*(1-xi)*(1+eta);
                -1.0/2*(1-xi)*(1+xi);
                -1.0/2*(1+xi)*(1+eta)+1.0/2*(1+xi)*(1-eta);
                1.0/2*(1-xi)*(1+xi);
                -1.0/2*(1-xi)*(1+eta)+1.0/2*(1-xi)*(1-eta)];
    return  reshape(val, 8, 2); # Make sure this is a matrix
end

function privboundaryconn(self::FESetQ8)
    conn = connasarray(self)
    return [conn[:, [1, 2, 5]]; conn[:, [2, 3, 6]]; conn[:, [3, 4, 7]]; conn[:, [4, 1, 8]]];
end

function privboundaryfe(self::FESetQ8)
    return FESetL3;
end

function privinparametric(self::FESetQ8, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetQ8)
    return vec([0.0 0.0])
end

################################################################################
################################################################################

"""
    FESetT6

Type for sets of surface-like triangular finite elements with six nodes.
"""
@define_FESet FESetT6 AbstractFESet2Manifold 6

function privbfun(self::FESetT6,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 4-node quadrilateral.
    r=param_coords[1];
    s=param_coords[2];
    t = 1. - r - s;
    val = [t * (t + t - 1);
           r * (r + r - 1);
           s * (s + s - 1);
           4 * r * t;
           4 * r * s;
           4 * s * t];
    return reshape(val, 6, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetT6,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    r =param_coords[1];
    s =param_coords[2];
    t = 1. - r - s;
    val = [-3+4*r+4*s  -3+4*r+4*s;
           4*r-1  0.0;
           0.0 4*s-1;
           4-8*r-4*s  -4*r;
           4*s  4*r;
           -4*s  4-4*r-8*s];
    return  reshape(val, 6, 2); # Make sure this is a matrix
end

function privboundaryconn(self::FESetT6)
    conn = connasarray(self)
    return [conn[:, [1,  2,  4]]; conn[:, [2,  3,  5]]; conn[:, [3,  1,  6]]];
end

function privboundaryfe(self::FESetT6)
    # Get  the constructor of the class of the  boundary finite element.
    return FESetL3;
end

function privinparametric(self::FESetT6, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords)) &&
      (-3.0 * tolerance <= sum(param_coords) <= 1.0 + 3.0 * tolerance);
end

function privcentroidparametric(self::FESetT6)
    return vec([1.0/3 1.0/3])
end

################################################################################
################################################################################

"""
    FESetH8

Type for sets of volume-like hexahedral finite elements with eight nodes.
"""
@define_FESet FESetH8 AbstractFESet3Manifold 8

function privbfun(self::FESetH8,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 8-node hexahedron.
    one_minus_xi    = (1.0 - param_coords[1]);
    one_minus_eta   = (1.0 - param_coords[2]);
    one_minus_theta = (1.0 - param_coords[3]);
    one_plus_xi     = (1.0 + param_coords[1]);
    one_plus_eta    = (1.0 + param_coords[2]);
    one_plus_theta  = (1.0 + param_coords[3]);
    val = [one_minus_xi*one_minus_eta*one_minus_theta;
           one_plus_xi*one_minus_eta*one_minus_theta;
           one_plus_xi*one_plus_eta*one_minus_theta;
           one_minus_xi*one_plus_eta*one_minus_theta;
           one_minus_xi*one_minus_eta*one_plus_theta;
           one_plus_xi*one_minus_eta*one_plus_theta;
           one_plus_xi*one_plus_eta*one_plus_theta;
           one_minus_xi*one_plus_eta*one_plus_theta] / 8.0;
    return reshape(val, 8, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetH8,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
     omxi    = (1.0 - param_coords[1]);
    ometa   = (1.0 - param_coords[2]);
    omtheta = (1.0 - param_coords[3]);
    opxi     = (1.0 + param_coords[1]);
    opeta    = (1.0 + param_coords[2]);
    optheta  = (1.0 + param_coords[3]);
   val = [-ometa*omtheta  ometa*omtheta  opeta*omtheta  -opeta*omtheta  -ometa*optheta  ometa*optheta  opeta*optheta   -opeta*optheta;
          -omxi*omtheta  -opxi*omtheta   opxi*omtheta   omxi*omtheta    -omxi*optheta   -opxi*optheta  opxi*optheta   omxi*optheta;
          -omxi*ometa    -opxi*ometa     -opxi*opeta    -omxi*opeta      omxi*ometa     opxi*ometa     opxi*opeta    omxi*opeta]'/8.0;
    return val
end

function privboundaryconn(self::FESetH8)
    conn = connasarray(self)
    return  [conn[:, [1,  4,  3,  2]];
             conn[:, [1,  2,  6,  5]];
             conn[:, [2,  3,  7,  6]];
             conn[:, [3,  4,  8,  7]];
             conn[:, [4,  1,  5,  8]];
             conn[:, [6,  7,  8,  5]]];
end

function privboundaryfe(self::FESetH8)
    # Get  the constructor of the class of the  boundary finite element.
    return FESetQ4;
end

function privinparametric(self::FESetH8, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetH8)
    return vec([0.0 0.0 0.0])
end

################################################################################
################################################################################

"""
    FESetH20

Type for sets of volume-like hexahedral finite elements with 20 nodes.
"""
@define_FESet FESetH20 AbstractFESet3Manifold 20

function privbfun(self::FESetH20,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 20-node hexahedron.
    xim    = (-1 + param_coords[1]);
    etam   = (-1 + param_coords[2]);
    zetam  = (-1 + param_coords[3]);
    xip    = (1 + param_coords[1]);
    etap   = (1 + param_coords[2]);
    zetap  = (1 + param_coords[3]);
    val = [ 1.0/8*xim*etam*zetam*(2+param_coords[1]+param_coords[2]+param_coords[3]);
           -1.0/8*xip*etam*zetam*(2-param_coords[1]+param_coords[2]+param_coords[3]);
           1.0/8*xip*etap*zetam*(2-param_coords[1]-param_coords[2]+param_coords[3]);
           -1.0/8*xim*etap*zetam*(2+param_coords[1]-param_coords[2]+param_coords[3]);
           1.0/8*xim*etam*zetap*(-2-param_coords[1]-param_coords[2]+param_coords[3]);
           -1.0/8*xip*etam*zetap*(-2+param_coords[1]-param_coords[2]+param_coords[3]);
           1.0/8*xip*etap*zetap*(-2+param_coords[1]+param_coords[2]+param_coords[3]);
           -1.0/8*xim*etap*zetap*(-2-param_coords[1]+param_coords[2]+param_coords[3]);
           -1.0/4*xim*xip*etam*zetam;
           1.0/4*etam*etap*xip*zetam;
           1.0/4*xim*xip*etap*zetam;
           -1.0/4*etam*etap*xim*zetam;
           1.0/4*xim*xip*etam*zetap;
           -1.0/4*etam*etap*xip*zetap;
           -1.0/4*xim*xip*etap*zetap;
           1.0/4*etam*etap*xim*zetap;
           -1.0/4*zetam*zetap*xim*etam;
           1.0/4*zetam*zetap*xip*etam;
           -1.0/4*zetam*zetap*xip*etap;
           1.0/4*zetam*zetap*xim*etap];
    return reshape(val, 20, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetH20,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    xim    = -(-1 + param_coords[1]);
    etam   = -(-1 + param_coords[2]);
    zetam  = -(-1 + param_coords[3]);
    xip    = (1 + param_coords[1]);
    etap   = (1 + param_coords[2]);
    zetap  = (1 + param_coords[3]);
    twoppp =(2+param_coords[1]+param_coords[2]+param_coords[3]);
    twompp =(2-param_coords[1]+param_coords[2]+param_coords[3]);
    twopmp =(2+param_coords[1]-param_coords[2]+param_coords[3]);
    twoppm =(2+param_coords[1]+param_coords[2]-param_coords[3]);
    twommp =(2-param_coords[1]-param_coords[2]+param_coords[3]);
    twopmm =(2+param_coords[1]-param_coords[2]-param_coords[3]);
    twompm =(2-param_coords[1]+param_coords[2]-param_coords[3]);
    twommm =(2-param_coords[1]-param_coords[2]-param_coords[3]);
    val =zeros(FFlt, 20, 3)
   val[:, 1]= [1.0/8*etam*zetam*twoppp-1.0/8*xim*etam*zetam;
           -1.0/8*etam*zetam*twompp+1.0/8*xip*etam*zetam;
           -1.0/8*etap*zetam*twommp+1.0/8*xip*etap*zetam;
           1.0/8*etap*zetam*twopmp-1.0/8*xim*etap*zetam;
           1.0/8*etam*zetap*twoppm-1.0/8*xim*etam*zetap;
           -1.0/8*etam*zetap*twompm+1.0/8*xip*etam*zetap;
           -1.0/8*etap*zetap*twommm+1.0/8*xip*etap*zetap;
           1.0/8*etap*zetap*twopmm-1.0/8*xim*etap*zetap;
           -1.0/4*xip*etam*zetam+1.0/4*xim*etam*zetam;
           1.0/4*etam*etap*zetam;
           -1.0/4*xip*etap*zetam+1.0/4*xim*etap*zetam;
           -1.0/4*etam*etap*zetam;
           -1.0/4*xip*etam*zetap+1.0/4*xim*etam*zetap;
           1.0/4*etam*etap*zetap;
           -1.0/4*xip*etap*zetap+1.0/4*xim*etap*zetap;
           -1.0/4*etam*etap*zetap;
           -1.0/4*zetam*zetap*etam;
           1.0/4*zetam*zetap*etam;
           1.0/4*zetam*zetap*etap;
           -1.0/4*zetam*zetap*etap];
    val[:, 2]= [1.0/8*xim*zetam*twoppp-1.0/8*xim*etam*zetam;
             1.0/8*xip*zetam*twompp-1.0/8*xip*etam*zetam;
             -1.0/8*xip*zetam*twommp+1.0/8*xip*etap*zetam;
             -1.0/8*xim*zetam*twopmp+1.0/8*xim*etap*zetam;
             1.0/8*xim*zetap*twoppm-1.0/8*xim*etam*zetap;
             1.0/8*xip*zetap*twompm-1.0/8*xip*etam*zetap;
             -1.0/8*xip*zetap*twommm+1.0/8*xip*etap*zetap;
             -1.0/8*xim*zetap*twopmm+1.0/8*xim*etap*zetap;
             -1.0/4*xim*xip*zetam;
             -1.0/4*xip*etap*zetam+1.0/4*xip*etam*zetam;
             1.0/4*xim*xip*zetam;
             -1.0/4*xim*etap*zetam+1.0/4*xim*etam*zetam;
             -1.0/4*xim*xip*zetap;
             -1.0/4*xip*etap*zetap+1.0/4*xip*etam*zetap;
             1.0/4*xim*xip*zetap;
             -1.0/4*xim*etap*zetap+1.0/4*xim*etam*zetap;
             -1.0/4*zetam*zetap*xim;
             -1.0/4*zetam*zetap*xip;
             1.0/4*zetam*zetap*xip;
             1.0/4*zetam*zetap*xim];
    val[:, 3]= [1.0/8*xim*etam*twoppp-1.0/8*xim*etam*zetam;
             1.0/8*xip*etam*twompp-1.0/8*xip*etam*zetam;
             1.0/8*xip*etap*twommp-1.0/8*xip*etap*zetam;
             1.0/8*xim*etap*twopmp-1.0/8*xim*etap*zetam;
             -1.0/8*xim*etam*twoppm+1.0/8*xim*etam*zetap;
             -1.0/8*xip*etam*twompm+1.0/8*xip*etam*zetap;
             -1.0/8*xip*etap*twommm+1.0/8*xip*etap*zetap;
             -1.0/8*xim*etap*twopmm+1.0/8*xim*etap*zetap;
             -1.0/4*xim*xip*etam;
             -1.0/4*etam*etap*xip;
             -1.0/4*xim*xip*etap;
             -1.0/4*etam*etap*xim;
             1.0/4*xim*xip*etam;
             1.0/4*etam*etap*xip;
             1.0/4*xim*xip*etap;
             1.0/4*etam*etap*xim;
             -1.0/4*xim*etam*zetap+1.0/4*xim*etam*zetam;
             -1.0/4*xip*etam*zetap+1.0/4*xip*etam*zetam;
             -1.0/4*xip*etap*zetap+1.0/4*xip*etap*zetam;
             -1.0/4*xim*etap*zetap+1.0/4*xim*etap*zetam];
    return reshape(val, 20, 3);
end

function privboundaryconn(self::FESetH20)
    conn = connasarray(self)
    return  [conn[:, [1,  4,  3,  2,  12,  11,  10,  9]];
             conn[:, [1,  2,  6,  5,  9,  18,  13,  17]];
             conn[:, [2,  3,  7,  6,  10,  19,  14,  18]];
             conn[:, [3,  4,  8,  7,   11,  20,  15,  19]];
             conn[:, [4,  1,  5,  8,  12,  17,  16,  20]];
             conn[:, [6,  7,  8,  5,  14,  15,  16,  13]]];
end

function privboundaryfe(self::FESetH20)
    # Get  the constructor of the class of the  boundary finite element.
    return FESetQ8;
end

function privinparametric(self::FESetH20, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetH20)
    return vec([0.0 0.0 0.0])
end

################################################################################
################################################################################

"""
    FESetH27

Type for sets of volume-like hexahedral finite elements with 27 nodes.
"""
@define_FESet FESetH27 AbstractFESet3Manifold 27

function privbfun(self::FESetH27,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 8-node hexahedron.
    xi=param_coords[1];
    eta=param_coords[2];
    zet=param_coords[3];
    x1 =(xi-1);
    y1=(eta-1);
    z1 =(zet-1);
    x2 =(xi+1);
    y2=(eta+1);
    z2 =(zet+1);
    val =[1.0/8.0*z1*zet*x1*xi*y1*eta
          1.0/8.0*z1*zet*x2*xi*y1*eta
          1.0/8.0*z1*zet*x2*xi*y2*eta
          1.0/8.0*z1*zet*x1*xi*y2*eta
          1.0/8.0*z2*zet*x1*xi*y1*eta
          1.0/8.0*z2*zet*x2*xi*y1*eta
          1.0/8.0*z2*zet*x2*xi*y2*eta
          1.0/8.0*z2*zet*x1*xi*y2*eta
          1.0/4.0*z1*zet*(-x2)*x1*y1*eta
          1.0/4.0*z1*zet*x2*xi*(-y2)*y1
          1.0/4.0*z1*zet*(-x2)*x1*y2*eta
          1.0/4.0*z1*zet*x1*xi*(-y2)*y1
          1.0/4.0*z2*zet*(-x2)*x1*y1*eta
          1.0/4.0*z2*zet*x2*xi*(-y2)*y1
          1.0/4.0*z2*zet*(-x2)*x1*y2*eta
          1.0/4.0*z2*zet*x1*xi*(-y2)*y1
          1.0/4.0*(-z2)*z1*x1*xi*y1*eta
          1.0/4.0*(-z2)*z1*x2*xi*y1*eta
          1.0/4.0*(-z2)*z1*x2*xi*y2*eta
          1.0/4.0*(-z2)*z1*x1*xi*y2*eta
          1.0/2.0*z1*zet*(-x2)*x1*(-y2)*y1
          1.0/2.0*(-z2)*z1*(-x2)*x1*y1*eta
          1.0/2.0*(-z2)*z1*x2*xi*(-y2)*y1
          1.0/2.0*(-z2)*z1*(-x2)*x1*y2*eta
          1.0/2.0*(-z2)*z1*x1*xi*(-y2)*y1
          1.0/2.0*z2*zet*(-x2)*x1*(-y2)*y1
          (-z2)*z1*(-x2)*x1*(-y2)*y1];
    return reshape(val, 27, 1); # Make sure this is a matrix
end

function privbfundpar(self::FESetH27,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    xi=param_coords[1];
    eta=param_coords[2];
    zet=param_coords[3];
    x1 =(xi-1.0/2.0);
    x2 =(xi+1.0/2.0);
    x3 =(xi-1.0);
    x4 =(xi+1.0);
    z1 =(zet-1.0);
    z2 =(zet-1.0/2.0);
    z3 =(zet+1.0);
    z4 =(zet+1.0/2.0);
    y1 =(eta-1.0);
    y2 =(eta-1.0/2.0);
    y3 =(eta+1.0);
    y4 = (eta+1.0/2.0);
    val = [      1.0/4.0*z1*zet*x1*y1*eta        1.0/4.0*z1*zet*x3*xi*y2        1.0/4.0*z2*x3*xi*y1*eta;
                  1.0/4.0*z1*zet*x2*y1*eta        1.0/4.0*z1*zet*x4*xi*y2        1.0/4.0*z2*x4*xi*y1*eta;
                  1.0/4.0*z1*zet*x2*y3*eta        1.0/4.0*z1*zet*x4*xi*y4        1.0/4.0*z2*x4*xi*y3*eta;
                  1.0/4.0*z1*zet*x1*y3*eta        1.0/4.0*z1*zet*x3*xi*y4        1.0/4.0*z2*x3*xi*y3*eta;
                  1.0/4.0*z3*zet*x1*y1*eta        1.0/4.0*z3*zet*x3*xi*y2        1.0/4.0*z4*x3*xi*y1*eta;
                  1.0/4.0*z3*zet*x2*y1*eta        1.0/4.0*z3*zet*x4*xi*y2        1.0/4.0*z4*x4*xi*y1*eta;
                  1.0/4.0*z3*zet*x2*y3*eta        1.0/4.0*z3*zet*x4*xi*y4        1.0/4.0*z4*x4*xi*y3*eta;
                  1.0/4.0*z3*zet*x1*y3*eta        1.0/4.0*z3*zet*x3*xi*y4        1.0/4.0*z4*x3*xi*y3*eta;
                       -1.0/2.0*z1*zet*xi*y1*eta   1.0/2.0*z1*zet*(-x4)*x3*y2   1.0/2.0*z2*(-x4)*x3*y1*eta;
             1.0/2.0*z1*zet*x2*(-y3)*y1             -1.0/2.0*z1*zet*x4*xi*eta   1.0/2.0*z2*x4*xi*(-y3)*y1;
                       -1.0/2.0*z1*zet*xi*y3*eta   1.0/2.0*z1*zet*(-x4)*x3*y4   1.0/2.0*z2*(-x4)*x3*y3*eta;
             1.0/2.0*z1*zet*x1*(-y3)*y1             -1.0/2.0*z1*zet*x3*xi*eta   1.0/2.0*z2*x3*xi*(-y3)*y1;
                       -1.0/2.0*z3*zet*xi*y1*eta   1.0/2.0*z3*zet*(-x4)*x3*y2   1.0/2.0*z4*(-x4)*x3*y1*eta;
             1.0/2.0*z3*zet*x2*(-y3)*y1             -1.0/2.0*z3*zet*x4*xi*eta   1.0/2.0*z4*x4*xi*(-y3)*y1;
                       -1.0/2.0*z3*zet*xi*y3*eta   1.0/2.0*z3*zet*(-x4)*x3*y4   1.0/2.0*z4*(-x4)*x3*y3*eta;
             1.0/2.0*z3*zet*x1*(-y3)*y1             -1.0/2.0*z3*zet*x3*xi*eta   1.0/2.0*z4*x3*xi*(-y3)*y1;
             1.0/2.0*(-z3)*z1*x1*y1*eta            1.0/2.0*(-z3)*z1*x3*xi*y2             -1.0/2.0*zet*x3*xi*y1*eta;
             1.0/2.0*(-z3)*z1*x2*y1*eta            1.0/2.0*(-z3)*z1*x4*xi*y2             -1.0/2.0*zet*x4*xi*y1*eta;
             1.0/2.0*(-z3)*z1*x2*y3*eta            1.0/2.0*(-z3)*z1*x4*xi*y4             -1.0/2.0*zet*x4*xi*y3*eta;
             1.0/2.0*(-z3)*z1*x1*y3*eta            1.0/2.0*(-z3)*z1*x3*xi*y4             -1.0/2.0*zet*x3*xi*y3*eta;
                      -z1*zet*xi*(-y3)*y1            -z1*zet*(-x4)*x3*eta  z2*(-x4)*x3*(-y3)*y1;
                      -(-z3)*z1*xi*y1*eta  (-z3)*z1*(-x4)*x3*y2            -zet*(-x4)*x3*y1*eta;
            (-z3)*z1*x2*(-y3)*y1            -(-z3)*z1*x4*xi*eta            -zet*x4*xi*(-y3)*y1;
                      -(-z3)*z1*xi*y3*eta  (-z3)*z1*(-x4)*x3*y4            -zet*(-x4)*x3*y3*eta;
            (-z3)*z1*x1*(-y3)*y1            -(-z3)*z1*x3*xi*eta            -zet*x3*xi*(-y3)*y1;
                      -z3*zet*xi*(-y3)*y1            -z3*zet*(-x4)*x3*eta  z4*(-x4)*x3*(-y3)*y1;
               -2.0*(-z3)*z1*xi*(-y3)*y1     -2.0*(-z3)*z1*(-x4)*x3*eta     -2.0*zet*(-x4)*x3*(-y3)*y1;
           ];
    return reshape(val, 27, 3);
end

function privboundaryconn(self::FESetH27)
    conn = connasarray(self)
    return  [conn[:, [1,  4,  3,  2,  12,  11,  10,  9,  21]];
             conn[:, [1,  2,  6,  5,  9,  18,  13,  17,  22]];
             conn[:, [2,  3,  7,  6,  10,  19,  14,  18,  23]];
             conn[:, [3,  4,  8,  7,   11,  20,  15,  19,  24]];
             conn[:, [4,  1,  5,  8,  12,  17,  16,  20,  25]];
             conn[:, [6,  7,  8,  5,  14,  15,  16,  13,  26]]];
end

function privboundaryfe(self::FESetH27)
    return FESetQ9;
end

function privinparametric(self::FESetH27, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-1.0 - tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords));
end

function privcentroidparametric(self::FESetH27)
    return vec([0.0 0.0 0.0])
end

################################################################################
################################################################################

"""
    FESetT4

Type for sets of volume-like tetrahedral finite elements with four nodes.
"""
@define_FESet FESetT4 AbstractFESet3Manifold 4

function privbfun(self::FESetT4,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 3-node triangle.
     val = [(1 - param_coords[1] - param_coords[2] - param_coords[3]);
            param_coords[1];
            param_coords[2];
            param_coords[3]];
    return reshape(val, 4, 1)
end

function privbfundpar(self::FESetT4,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    val = [-1. -1. -1.;
           +1.  0.  0.;
           0. +1.  0.;
           0.  0. +1.];
    return reshape(val, 4, 3)
end

function privboundaryconn(self::FESetT4)
    conn = connasarray(self)
    return [conn[:, [1,  3,  2]]; conn[:, [1,  2,  4]]; conn[:, [2,  3,  4]]; conn[:, [1,  4,  3]]];
end

function privboundaryfe(self::FESetT4)
    return FESetT3;
end

function privinparametric(self::FESetT4, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords)) &&
      (-3.0 * tolerance <= sum(param_coords) <= 1.0 + 3.0 * tolerance);
end

function privcentroidparametric(self::FESetT4)
    return vec([1.0/4 1.0/4 1.0/4])
end

################################################################################
################################################################################

"""
    FESetT10

Type for sets of volume-like tetrahedral finite elements with 10 nodes.
"""
@define_FESet FESetT10 AbstractFESet3Manifold 10

function privbfun(self::FESetT10,  param_coords::FFltVec)
    # Evaluate the basis function matrix for an 3-node triangle.
    r = param_coords[1];
    s = param_coords[2];
    t = param_coords[3];
    val = [(1-r-s-t) * (2*(1-r-s-t)-1);
           r * (2*r-1);
           s * (2*s-1);
           t * (2*t-1);
           4*(1-r-s-t)*r;
           4*r*s;
           4*s*(1-r-s-t);
           4*(1-r-s-t)*t;
           4*r*t;
           4*s*t;
           ];
    return reshape(val, 10, 1)
end

function privbfundpar(self::FESetT10,  param_coords::FFltVec)
    # Evaluate the derivatives of the basis function matrix.
    r = param_coords[1];
    s = param_coords[2];
    t = param_coords[3];
            val = [-3+4*r+4*s+4*t   4*r-1 0   0  -8*r+4-4*s-4*t  4*s -4*s  -4*t  4*t  0;
                -3+4*r+4*s+4*t  0  4*s-1  0   -4*r   4*r 4-4*r-8*s-4*t -4*t  0 4*t;
                -3+4*r+4*s+4*t  0  0  4*t-1   -4*r   0  -4*s -8*t+4-4*r-4*s  4*r  4*s
                ]';
    return reshape(val, 10, 3)
end

function privboundaryconn(self::FESetT10)
    conn = connasarray(self)
    return [conn[:, [1,  3,  2,  7,  6,  5]];
            conn[:, [1,  2,  4,  5,  9,  8]];
            conn[:, [2,  3,  4,  6,  10,  9]];
            conn[:, [3,  1,  4,  7,  8,  10]]];
end

function privboundaryfe(self::FESetT10)
    return FESetT6;
end

function privinparametric(self::FESetT10, param_coords::FFltVec, tolerance::FFlt)
    s = findall(v -> (v >= (-tolerance)) && (v <= (+1.0 + tolerance)), param_coords);
    return  (length(s) == length(param_coords)) &&
      (-3.0 * tolerance <= sum(param_coords) <= 1.0 + 3.0 * tolerance);
end

function privcentroidparametric(self::FESetT10)
    return vec([1.0/4 1.0/4 1.0/4])
end

################################################################################
################################################################################

end


module PhaseRec
using FFTW
using LinearAlgebra
using Statistics
using Logging
using Base.Iterators

include("reconstruction.jl")
export phaserec, two_point

end # module


"""
   Traits for the NEID spectrograph (data products from DRP v0.6).
   https://neid.psu.edu/
Author: Eric Ford and collaborators
Created: August 2020
"""

""" Delegates loading of code specifying types essential to the package.  """

import RvSpectMLBase: min_order, max_order, min_pixel_in_order, max_pixel_in_order, min_pixel, max_pixel

#= Pre-ship
min_order(::NEID2D) = 1
max_order(::NEID2D) = 90
=#
# At KPNO
min_order(::NEID2D) = 1
#max_order(::NEID2D) = 126 # DRS 0.5
max_order(::NEID2D) = 122 # DRS 0.6
min_pixel_in_order(inst::NEID2D) = 1
max_pixel_in_order(inst::NEID2D) = 9216

min_pixel(::NEID1D) = 1
max_pixel(::NEID1D) = (max_order(NEID1D())-min_order(NEID1D())+1)*9216 # TODO: Update once know size of NEID's 1d extracted spectra

bad_col_ranges(inst::NEID2D, ord::Int) = [439:449, 1934:1943, 6714:6714]

import RvSpectMLBase: orders_to_use_default, min_col_default, max_col_default

#= Pre-ship
#orders_to_use_default(inst::NEID2D) = 1:52   # Avoiding redder orders due to tellurics
orders_to_use_default(inst::NEID2D) = 1:71   # Avoiding 72 because of NaNs in solar data
min_col_default(::NEID2D, ord::Integer) = 451              # Avoiding smaller columns due to NaNs
#min_col_default(::NEID2D) = 2000              # Avoiding smaller columns due to lower flux and distortions
max_col_default(::NEID2D, ord::Integer) = 9216 - (min_col_default(NEID2D(),ord)-1)   # Avoiding larger columns for symmetry
=#
# At KPNO
#orders_to_use_default(::NEID2D) = 55:113 # everything plaussibly usable as of DRS 0.5
#orders_to_use_default(::NEID2D) = 56:111 # everything plaussibly usable as of DRS 0.6
#orders_to_use_default(::NEID2D) = 60:90  # relatively safe
orders_to_use_default(::NEID2D) = 56:98   # avoid worst of tellurics
function min_col_default(::NEID2D, ord::Integer)
#    return 1445 +500 # DRS 0.6, avoiding NaN in cols 1435-1444
    if ord == 55
        return 787
    else
        return 500
    end
end
#max_col_default(::NEID2D, ord::Integer) = 8429  # DRS 0.5
#max_col_default(::NEID2D, ord::Integer) = 6214  # DRS 0.6, avoiding NaN in col 6215
max_col_default(::NEID2D, ord::Integer) = 9215  # DRS 0.6

import RvSpectMLBase: get_pixel_range
function get_pixel_range(inst::NEID2D, ord::Integer)
    minc = max(min_col_default(inst, ord)) #, min_col_excalibur(inst,order), min_col_nonnan(inst,order))
    maxc = min(max_col_default(inst, ord)) #, max_col_excalibur(inst,order), max_col_nonnan(inst,order))
    return minc:maxc
end

import RvSpectMLBase: metadata_symbols_default, metadata_strings_default
#= Pre-ship headers
#metadata_symbols_default(::AnyNEID) = Symbol[:bjd, :target, :ssbz]
#metadata_strings_default(::AnyNEID) = String["OBSJD", "SKY-OBJ", "SSBZ000"]
=#
# Headers at KPNO
metadata_symbols_default(::AnyNEID) = Symbol[:bjd, :target, :exptime, :airmass, :ssbz]
metadata_strings_default(::AnyNEID) = String["OBSJD", "OBJECT", "EXPTIME", "AIRMASS",  "SSBZ100"]

import RvSpectMLBase: default_ccf_mask_v_width
default_ccf_mask_v_width(::AnyNEID) = 620.953

import RvSpectMLBase: get_inst_module
get_inst_module(::AnyNEID) = NEID

import RvSpectMLBase: get_λ_range
function get_λ_range(data::CLT) where { T1<:Real, T2<:Real, T3<:Real, A1<:AbstractArray{T1,2}, A2<:AbstractArray{T2,2}, A3<:AbstractArray{T3,2},
                                       IT<:AnyNEID, CLT<:Spectra2DBasic{T1,T2,T3,A1,A2,A3,IT} }
   (λmin, λmax) = NaNMath.extrema(data.λ)
   return (min=λmin, max=λmax)
end

#=
# Pre-ship
default_λmin = 3950.0  # Based on HD solar data from PSU, should generalize
default_λmax = 9500.0  #
=#
# At KPNO
#default_λmin = 5040.04045429369  # Based on one FITS file, should generalize
#default_λmax = 9810.539033567005 # Based on one FITS file, should generalize
default_λmin = 5145.0
default_λmax = 9822.0


include("../../__lib__/std/util/src/util_module.jl")
include("../../__lib__/std/collections/gridded/src/gridded_module.jl")

module compgraph
using ..util

function contour_level(xs, ys, z, level)
    #   vertex    edges
	##  1-----2   .-1-.
    ##  |     |   |   |
    ##  |     |   4   2
    ##  |     |   |   |
    ##  4-----3   .-3-.

    Z = z .- level
    mask = Z .>= 0
    mark = mark_cells(mask)

    # show(stdout, MIME("text/plain"), mark)

    accum_array = zeros(2, 0)

    (n_row, n_col) = size(mark)

    # top search
    step_dir   = down::StepDirections
    first           = true
    start_edge      = 1
    rr              = 1
    for cc = 1:n_col
        id = mark[rr, cc]
        if id > 0 && (  id == 2 || id == 6 || id == 14 ||  id == 1 || id == 9 || id == 13 )
            (mark, accum_array) = __search_path!(xs, ys, Z, mark, accum_array, rr, cc, start_edge, first, step_dir, level)
        end
    end

    # bottom search
    step_dir   = up::StepDirections
    first           = true
    start_edge      = 3
    rr              = n_row
    for cc = 1:n_col
        id = mark[rr, cc]
        if id > 0 && ( id == 8 || id == 9 || id == 11 || id == 4 || id == 6 || id == 7)
            (mark, accum_array) = __search_path!(xs, ys, Z, mark, accum_array, rr, cc, start_edge, first, step_dir, level)
        end
    end

    # left search
    step_dir   = right::StepDirections
    first           = true
    start_edge      = 4
    cc              = 1
    for rr = 1:n_row
        id = mark[rr, cc]
        if id > 0 && ( id == 1 || id == 3 || id == 7 || id == 8 || id == 12 || id == 14)
            (mark, accum_array) = __search_path!(xs, ys, Z, mark, accum_array, rr, cc, start_edge, first, step_dir, level)
        end
    end

    # right search
    step_dir   = left::StepDirections
    first           = true
    start_edge      = 2
    cc              = n_col
    for rr = 1:n_row
        id = mark[rr, cc]
        if id > 0 &&  (id == 4 || id == 12 || id == 13 || id == 2 || id == 3 || id == 11)
            (mark, accum_array) = __search_path!(xs, ys, Z, mark, accum_array, rr, cc, start_edge, first, step_dir, level)
        end
    end

    # interior search
    step_dir = none::StepDirections
    first = true
    start_edge = 1
    for cc = 1:n_col-1
        for rr = 1:n_row-1
            id = mark[rr, cc]
            if id > 0
                (mark, accum_array) = __search_path!(xs, ys, Z, mark, accum_array, rr, cc, start_edge, first, step_dir, level)
            end
        end
    end

    return (accum_array, mark)
end

@enum StepDirections begin
    none
    up
    right
    down
    left
end

function __search_path!(xs, ys, Z, mark, accum_array, r, c, init_start_edge, first, step_dir::StepDirections, level)

    (n_row, n_col) = size(mark)
    start_idx = size(accum_array, 2) + 1

    # step_dir = none::StepDirections
    prev_step_dir = none::StepDirections # overrite for now

    stop_edge = Int8(127)
    while r >= 1 && c >= 1 && r <= n_row && c <= n_col && mark[r, c] > 0

        px = ( xs[c], xs[c+1], xs[c+1], xs[c] )
        py = ( ys[r], ys[r],   ys[r+1], ys[r+1] )
        pz = ( Z[r, c], Z[r, c+1], Z[r+1, c+1], Z[r+1, c] )

        # find start edge
        (start_edge, stop_edge) = find_start_and_stop_edge(mark, r, c)
        # if first
        #     start_edge = init_start_edge
        # end

        #   vertex    edges
        ##  1-----2   .-1-.
        ##  |     |   |   |
        ##  |     |   4   2
        ##  |     |   |   |
        ##  4-----3   .-3-.

        # switch the start and stop edges
        if step_dir == none::StepDirections
            # do nothing
        elseif step_dir == up::StepDirections
            if start_edge != 3
                start_edge, stop_edge = stop_edge, start_edge
            end
        elseif step_dir == down::StepDirections
            if start_edge != 1
                start_edge, stop_edge = stop_edge, start_edge
            end
        elseif step_dir == left::StepDirections
            if start_edge != 2
                start_edge, stop_edge = stop_edge, start_edge
            end
        elseif step_dir == right::StepDirections
            if start_edge != 4
                start_edge, stop_edge = stop_edge, start_edge
            end
        end

        if start_edge == 127
            # unsuccessful search, break up
            break
        end

        if first

            pt = (start_edge, util.modulo_index(start_edge+1,4))
            (ct_x, ct_y) = lerp_interp(px, py, pz, pt)

            accum_array = add_point!(accum_array, level, 0)
            accum_array = add_point!(accum_array, ct_x, ct_y)
            first = false
        end

        pt = (stop_edge, util.modulo_index(stop_edge+1,4))
        (ct_x, ct_y) = lerp_interp(px, py, pz, pt)
        accum_array = add_point!(accum_array, ct_x, ct_y)

        mark[r, c] -= mark[r, c]

        prev_step_dir = step_dir
        if (stop_edge == 1)
            r -= 1;
            step_dir = up::StepDirections;
        elseif (stop_edge == 2)
            c += 1;
            step_dir = right::StepDirections;
        elseif (stop_edge == 3)
            r += 1;
            step_dir = down::StepDirections;
        elseif (stop_edge == 4)
            c -= 1;
            step_dir = left::StepDirections;
        end

    end # while

    end_idx = size(accum_array, 2)
    if start_idx >= 1
        accum_array[2,start_idx] = end_idx - start_idx
    end

    return (mark, accum_array)
end

function add_point!(accum_array, x, y)
    accum_array = hcat(accum_array, [x, y])
    return accum_array
end

function mark_cells(mask)
    (n_row, n_col) = size(mask)

    mark = zeros(Int8, n_row-1, n_col-1)

    for jj = 1:n_col-1
        for ii = 1:n_row-1

            f1 = mask[ii, jj]
            f2 = mask[ii, jj+1]
            f3 = mask[ii+1, jj+1]
            f4 = mask[ii+1, jj]

            mark[ii, jj] = f1 + f2 * 2 + f3 * 4 + f4 * 8

        end
    end

    mark[ mark .>= 15] .= 0

    return mark

end

function lerp_interp(px, py, pz, pt)

    tmp = pz[pt[2]] / pz[pt[1]];
    tmp = abs(tmp);

    if !isnan(tmp)
        lamda = 1 / (1 + tmp)
        ct_x = px[pt[1]] * (1-lamda) + px[pt[2]] * lamda
        ct_y = py[pt[1]] * (1-lamda) + py[pt[2]] * lamda
    else
        lamda = 0.5;
        ct_x = px[pt[1]] * (1-lamda) + px[pt[2]] * lamda
        ct_y = py[pt[1]] * (1-lamda) + py[pt[2]] * lamda
    end

    return (ct_x, ct_y)
end

function find_start_and_stop_edge(mark, r, c)
    id = mark[r, c]

    start_edge = Int8(127)
    stop_edge = Int8(127)

    if id == 0 || id == 15
        # impossible condition
        return (start_edge, stop_edge)
    end

    if id == 1 || id == 3 || id == 7
        start_edge = 4
    elseif id == 2 || id == 6 || id == 14
        start_edge = 1
    elseif id == 4 || id == 12 || id == 13
        start_edge = 2
    elseif id == 8 || id == 9 || id == 11
        start_edge = 3
    end

    if id == 1 || id == 9 || id == 13
        stop_edge = 1
    elseif id == 2 || id == 3 || id == 11
        stop_edge = 2
    elseif id == 4 || id == 6 || id == 7
        stop_edge = 3
    elseif id == 8 || id == 12 || id == 14
        stop_edge = 4
    end

    if id == 5 || id == 10
        start_edge = -1
        stop_edge = -1
    end

    start_edge = Int8(start_edge)
    stop_edge = Int8(stop_edge)

    return (start_edge, stop_edge)

end

end

module ctest
using ..util
using ..compgraph
using ..gridded
using PyPlot
PyPlot.pygui(true)

Nx = 10 # 10;
Ny = 8 # 8;

x_min = -1.0;
x_max = 2.0;

y_min = -2.0;
y_max = 3.0;

xs = collect( range(x_min, x_max, Nx) );
ys = collect( range(y_min, y_max, Ny) );

(xm, ym) = util.mesh_grid(xs, ys);

xm = xm'
ym = ym'
# zm = xm .+ ym;

# zm = xm.^2 .+ ym.^2

gmap = gridded.peaks()
zm = gmap.z'
xm = gmap.x
ym = gmap.y

xs = xm
ys = ym

level = 1.5;

# show(stdout, MIME("text/plain"), zm)
# println("")

(accum_array, mark) = compgraph.contour_level(xs, ys, zm, level)

PyPlot.figure()
PyPlot.grid()
PyPlot.contour(xm, ym, zm, [level], linewidths = 5, colors = :red)

PyPlot.plot( accum_array[1,2:end], accum_array[2,2:end], marker = :., markersize = 10 )

end


using Colors
using FileIO
using GLMakie
using GeoJSON
using GeoMakie.GeoInterface
using XLSX

inzidenz_colors = [
    RGB(244/255, 225/255, 96/255),
    RGB(224/255, 164/255, 74/255),
    RGB(220/255, 117/255, 56/255),
    RGB(195/255, 75/255, 64/255),
    RGB(117/255, 47/255, 46/255)
]

read_kreise() = GeoJSON.read(read("landkreise.geojson"))

function color_for_inzidenz(x)
    if x <= 25
        inzidenz_colors[1]
    elseif x <= 50
        inzidenz_colors[2]
    elseif x <= 100
        inzidenz_colors[3]
    elseif x <= 250
        inzidenz_colors[4]
    elseif x <= 500
        inzidenz_colors[5]
    else
        RGB(0, 0, 0)
    end
end

function read_inzidenzen(rki_excel_file)
    xf = XLSX.readxlsx(rki_excel_file)
    rows = xf["LK_7-Tage-Inzidenz (fixiert)"][:][6:end, [3, end]]

    rows_with_content = filter(collect(eachrow(rows))) do row
        !ismissing(row[1])
    end
    Dict(map(rows_with_content) do row
         @sprintf("%05i", row[1]) => row[2]
    end)
end

function colors_from_rki(kreise, rki_excel_file)
    inz_dict = read_inzidenzen(rki_excel_file)
    map(kreise.features) do kreis
        ags = kreis.properties["AGS"]

        try
            color_for_inzidenz(inz_dict[ags])
        catch
            RGB(1, 1, 1) # white for not known data
        end
    end
end

function inzidenz_legend(scene::Scene, date)
    u25 = PolyElement(color = inzidenz_colors[1], strokecolor=:black)
    u50 = PolyElement(color = inzidenz_colors[2], strokecolor=:black)
    u100 = PolyElement(color = inzidenz_colors[3], strokecolor=:black)
    u250 = PolyElement(color = inzidenz_colors[4], strokecolor=:black)
    u500 = PolyElement(color = inzidenz_colors[5], strokecolor=:black)
    above = PolyElement(color = RGB(0, 0, 0), strokecolor=:black)

    leg = Legend(scene,
                 [u25, u50, u100, u250, u500, above],
                 ["<=25", "25<x<=50", "50<x<=100", "100<x<=250", "250<x<=500", "x>500"],
                 "7T-Inzidenz, $date",
                 patchsize=(35,35),
                 orientation=:vertical,
                 tellwidth=false,
                 tellheight=false,
                 halign=:center,
                 margin=(350, 0, 350, 0),
                )
end

function create_plot(kreise, colors, date)
    fig = Figure(resolution=(1500, 1920), title="11. November 2021")
    ax = fig[1,1] = Axis(fig)

    poly!(ax, kreise, color=colors, strokewidth=1, strokecolor=:white, )
    inzidenz_legend(fig.scene, date)

    display(fig)
    return fig
end

function rki_plot(rki_excel_file, date)
    kreise = read_kreise()
    colors = colors_from_rki(kreise, rki_excel_file)

    create_plot(kreise, colors, date)
end


#TODO: Add doc string
@with_kw struct ArrowCar{A<:AbstractArray{Float64}, C<:Colorant} <: Renderable
    pos::A         = SVector(0.0, 0.0)
    angle::Float64 = 0.0
    length::Float64 = 4.8
    width::Float64 = 1.8
    color::C       = COLOR_CAR_OTHER
    text::String   = "" # some debugging text to print by the car
    id::Int        = 0
end

ArrowCar(pos::AbstractArray, angle::Float64=0.0; length = 4.8, width = 1.8,  color=COLOR_CAR_OTHER, text="", id=0) = ArrowCar(pos, angle, length, width, color, text, id)
ArrowCar(x::Real, y::Real, angle::Float64=0.0; length = 4.8, width = 1.8,  color=COLOR_CAR_OTHER, text="", id=0) = ArrowCar(SVector(x, y), angle, length, width, color, text, id)

function render!(rm::RenderModel, c::ArrowCar)
    x = c.pos[1]
    y = c.pos[2]
    add_instruction!(rm, render_vehicle, (x, y, c.angle, c.length, c.width, c.color))
    add_instruction!(rm, render_text, (c.text, x, y-c.width/2 - 2.0, 10, colorant"white"))
    return rm
end

id(ac::ArrowCar) = ac.id
pos(ac::ArrowCar) = ac.pos


SRC_DIR = joinpath(Base.source_dir(), "..", "..", "..", "src")
OUT_DIR = joinpath(Base.source_dir(), "..", "..", "..", "output11")
mkpath(OUT_DIR)

#include(joinpath(SRC_DIR, "BridgeSDEInference.jl"))
#using Main.BridgeSDEInference
include(joinpath(SRC_DIR, "BridgeSDEInference_for_tests.jl"))

using StaticArrays, LinearAlgebra, GaussianDistributions
using Distributions
using Random
using Bridge
using CSV, DataFrames
DIR = "auxiliary"
include(joinpath(SRC_DIR, DIR, "data_simulation_fns.jl"))
include(joinpath(SRC_DIR, DIR, "utility_functions.jl"))
include(joinpath(SRC_DIR, DIR, "plotting_fns.jl"))

#=============================================================================
                Routines for setting up the MCMC sampler
=============================================================================#
function _prepare_setup(updt_order, ρ=0.96, num_mcmc_steps=4*10^3,
                        thin_path=10^0, save_path_every=1*10^3)
    model_setup = DiffusionSetup(P˟, P̃, PartObs())
    set_observations!(model_setup, Ls, Σs, obs_vals,
                      obs_time)
    set_imputation_grid!(model_setup, 1/100)
    set_x0_prior!(model_setup,
                  GsnStartingPt(x0, SArray{Tuple{4,4}, Float64}(I)), # prior over starting point
                  x0)
    set_auxiliary!(model_setup; skip_for_save=thin_path,
                   adaptive_prop=Adaptation(x0, fill(0.0, 50), fill(200, 50), 1))
    initialise!(eltype(x0), model_setup, Vern7(), false, NoChangePt())

    mcmc_setup = MCMCSetup(
        Imputation(NoBlocking(), ρ, Vern7()),
        ParamUpdate(MetropolisHastingsUpdt(), [1], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [1]))),
        ParamUpdate(MetropolisHastingsUpdt(), [2], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [2]))),
        ParamUpdate(MetropolisHastingsUpdt(), [3], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [3]))),
        ParamUpdate(MetropolisHastingsUpdt(), [4], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [4]))),
        ParamUpdate(MetropolisHastingsUpdt(), [5], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [5]))),
        #ParamUpdate(MetropolisHastingsUpdt(), [6], fill(0.0, 8),
        #            UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
        #            UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [6]))),
        ParamUpdate(MetropolisHastingsUpdt(), [7], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [7]))),
        ParamUpdate(MetropolisHastingsUpdt(), [8], fill(0.0, 8),
                    UniformRandomWalk(0.5, true), ExpUnif(-7.0, 2.0),
                    UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [8]))),

    )
    schedule = MCMCSchedule(num_mcmc_steps, updt_order,
                            (save=save_path_every, verbose=10^2, warm_up=0,
                             readjust=(x->x%100==0), fuse=(x->x==2000)) )
    mcmc_setup, schedule, model_setup
end

#==============================================================================
                        Auxiliary routines for saving
==============================================================================#

function save_paths(tt, paths, filename)
    d = length(paths[1][1])
    xx = copy(tt)
    for i in 1:length(paths)
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_marginals(tt, paths, filename, indices)
    d = length(paths[1][1])
    xx = copy(tt[indices])
    N = length(paths)
    stride = div(N, 100)
    for i in 1:length(paths)
        if i % stride == 0
            print(div(i, 100), "% done...\n")
        end
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i][indices]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx'))
end

function save_param_chain(chain, filename)
    d = length(chain[1])
    xx = [c[1] for c in chain]
    for i in 2:d
        xx = hcat(xx, [c[i] for c in chain])
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_obs(tt, paths, filename)
    d = length(paths[1][1])
    xx = copy(tt)
    for i in 1:length(paths)
        for j in 1:d
            xx = hcat(xx, [p[j] for p in paths[i]])
        end
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

function save_history(to_save, filename)
    n, k = length(to_save), length(to_save[1])
    xx = zeros(n, k)
    for i in 1:n
        xx[i,1:end] .= to_save[i]
    end
    CSV.write(joinpath(OUT_DIR, filename), DataFrame(xx))
end

#================================================================================
                            Generate the process
==============================================================================#

_obs = open(joinpath(OUT_DIR, "../output/prokaryote_custom.dat")) do f
    [map(x->parse(Float64, x), split(l, ' '))[[2,3]] for (i,l) in enumerate(eachline(f)) if i != 1]
end

#_obs = open(joinpath(OUT_DIR, "prokaryote_custom.dat")) do f
#    [map(x->parse(Float64, x), split(l, ' '))[2] for (i,l) in enumerate(eachline(f)) if i != 1]
#end
#obs_vals = vcat([0.0],[o for o in _obs])

_obs_time = collect(range(1.0, length(_obs), step=1))

#obs_vals = vcat([0.0],[o[2] for o in _obs])
obs_vals = vcat([[0.0]],[(i % 8 == 0 ? o : [o[2]]) for (i,o) in enumerate(_obs)])
obs_time = vcat(0.0, _obs_time)

# Permutation of observations
Random.seed!(5)
obs_RNA = [o[1] for (i,o) in enumerate(obs_vals) if i % 8 == 1][2:end]
obs_RNA_permuted = obs_RNA[randperm(length(obs_RNA))]
obs_vals = [((i % 8 == 1 && i != 1) ? [obs_RNA_permuted[div(i,8)], o[2]] : o) for (i,o) in enumerate(obs_vals)]

# let's start from the true values that generated the data
θ_init = [0.05, 0.05, 0.05, 0.05, 0.05, 0.5, 0.05, 0.05]#[0.1, 0.7, 0.35, 0.2, 0.1, 0.9, 0.3, 0.1]#
K = 10.0
P˟ = Prokaryote(θ_init..., K)
x0 = ℝ{4}([8.0,8.0,8.0,5.0])

auxFlag = Val{:custom}()
start_v = @SVector[5.0, 5.0, 5.0, 5.0]
P̃ = [ProkaryoteAux(θ_init..., K, t₀, u, T, v, auxFlag, start_v) for (t₀, T, u, v)
     in zip(obs_time[1:end-1], obs_time[2:end], obs_vals[1:end-1], obs_vals[2:end])]

Σ = @SMatrix[4.0]
L = @SMatrix[0.0 1.0 2.0 0.0]
#Ls = [L for _ in P̃]
#Σs = [Σ for _ in P̃]

Σextra = @SMatrix[1.0 0.0; 0.0 4.0]
Lextra = @SMatrix[1.0 0.0 0.0 0.0; 0.0 1.0 2.0 0.0]
Ls = [(i % 8 == 0 ? Lextra : L) for i in 1:length(P̃)]
Σs = [(i % 8 == 0 ? Σextra : Σ) for i in 1:length(P̃)]

#=
model_setup = DiffusionSetup(P˟, P̃, PartObs())
set_observations!(model_setup, [L for _ in P̃], [Σ for _ in P̃], obs_vals,
                  obs_time)
set_imputation_grid!(model_setup, 1/100)
set_x0_prior!(model_setup,
              GsnStartingPt(x0, SArray{Tuple{4,4}, Float64}(I)), # prior over starting point
              x0)
set_auxiliary!(model_setup; skip_for_save=1,
               adaptive_prop=Adaptation(x0, fill(0.0, 10), fill(100, 10), 1))
initialise!(eltype(x0), model_setup, Vern7(), false, NoChangePt())
model_setup

mcmc_setup = MCMCSetup(
    Imputation(NoBlocking(), 0.95, Vern7()),
    ParamUpdate(MetropolisHastingsUpdt(), [1], fill(0.0, 8),
                UniformRandomWalk(0.5, true), ImproperPosPrior(),
                UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [1]))),
    ParamUpdate(MetropolisHastingsUpdt(), [2], fill(0.0, 8),
                UniformRandomWalk(0.5, true), ImproperPosPrior(),
                UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [2]))),
    ParamUpdate(MetropolisHastingsUpdt(), [3], fill(0.0, 8),
                UniformRandomWalk(0.5, true), ImproperPosPrior(),
                UpdtAuxiliary(Vern7(), check_if_rec->falseompute_ODEs(P̃, [3]))),
    ParamUpdate(MetropolisHastingsUpdt(), [7], fill(0.0, 8),
                UniformRandomWalk(0.5, true), ImproperPosPrior(),
                UpdtAuxiliary(Vern7(), check_if_recompute_ODEs(P̃, [7]))))
schedule = MCMCSchedule(1*10^4, [[1,2,3,4,5]],
                        (save=10^2, verbose=10^2, warm_up=0,
                         readjust=(x->x%100==0), fuse=(x->false)) )
out = mcmc(mcmc_setup, schedule, model_setup)
=#
setup = _prepare_setup([[1,2,3,4,5,6,7,8]], 0.9, 1*10^5, 1, 10^3)
Random.seed!(6)
out, elapsed = @timeit mcmc(setup...)
plot_chains(out[2]; truth=[0.1, 0.7, 0.35, 0.2, 0.1, 0.9, 0.3, 0.1])
plot_paths(out[1], out[2], setup[2])

save_paths(out[1].time, out[1].paths, "prokaryote_paths.csv")
save_param_chain(out[2].θ_chain, "prokaryote_chain.csv")
save_history(out[2].updates[1].accpt_history, "prokaryote_accpt_hist_1.csv")
save_history(out[2].updates[2].accpt_history, "prokaryote_accpt_hist_2.csv")
save_history(out[2].updates[3].accpt_history, "prokaryote_accpt_hist_3.csv")
save_history(out[2].updates[4].accpt_history, "prokaryote_accpt_hist_4.csv")
save_history(out[2].updates[1].ρs_history, "prokaryote_rho_hist_1.csv")


## Copyright (c) 2015, Júlio Hoffimann Mendes <juliohm@stanford.edu>
##
## Permission to use, copy, modify, and/or distribute this software for any
## purpose with or without fee is hereby granted, provided that the above
## copyright notice and this permission notice appear in all copies.
##
## THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
## WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
## MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
## ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
## WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
## ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
## OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

function imfilter_cpu{T<:Real,K<:Real,N}(img::AbstractArray{T,N}, kern::AbstractArray{K,N})
  imfilter(img, centered(kern), Inner(), Algorithm.FFT())
end


#from PC-SAFT paper https://doi.org/10.1021/ie0003887
function fugacity_coeff_impl(mt::MultiVT,model::HelmholtzModel,v,t,x)
    Z = compressibility_factor_impl(mt,model,v,t,x)
    ar = z -> αR_impl(mt,model,one(v)/v,t,z)
    dadx = ForwardDiff.gradient(ar,x)
    sumdadx = dot(x,dadx)
    lnZ = log(Z)
    c0 = ar(x) + (Z-1) - lnZ - sumdadx
    return map(xi->xi+c0,dadx)
end

function flash_impl(mt::MultiPT,model,p,t,z0;threaded=true,moles=one(eltype(z0)))
    p = float(p)
    t = float(t)
    _0 = zero(p*t*first(z0))
    _1 = one(_0)
    logϕ0 = similar(z0)
    if threaded
        res_z0 = Threads.@spawn update_logϕ!!(model,logϕ0,$p,$t,z0,nothing,:undefined)
    end
    kmodel = WilsonK(model)
    k0 = kvalues_impl(mt,kmodel,p,t,z0)
    k = copy(k0)
    #@show k
    g_0 = flash_eval(k,z0,_0)
    g_1 = flash_eval(k,z0,_1)
    if g_0 <= 0  #bubble point assumption
        #println("bubble point")
        β = _0
        β0 = _0
        xil = copy(z0)
        xiv = flash_vapor(k,z0,_0)
    elseif g_1 >= 0 #dew point assumption
        #println("dew point")
        β = _1
        β0 = _1
        xil = flash_liquid(k,z0,_1)
        xiv = copy(z0)
    else #two phase assumption
        #println("flash point")
        β = flash_vfrac(RR(),k,z0)
        xil = flash_liquid(k,z0,β)
        xiv = flash_vapor(k,z0,β)
    end
    logϕl = similar(xil)
    logϕv = similar(xiv)
    xil = normalizefrac!!(xil)
    xiv = normalizefrac!!(xiv)

    vl,logϕl,vv,logϕv = update_logϕ_vl!!(model,p,t,z0,xil,xiv,logϕl,logϕv,nothing,nothing,threaded)
    lnk = logϕl - logϕv
    @! k .= exp.(lnk)

    stable_phase = false
    for i = 1:5
        res = refine_k!!(model,p,t,z0,xil,xiv,vl,vv,k,logϕl,logϕv,lnk,threaded)
        xil,xiv,vl,vv,k,logϕl,logϕv,lnk,β = res

        if !(_0 <= β <= _1)
            stable_phase = true
            break
        end
    end
    #if the process was threaded, the value is fetched here
    if !threaded
        vz0,logϕ0 = update_logϕ!!(model,logϕ0,p,t,z0,nothing,:undefined)
    else
        vz0,logϕ0 = fetch(res_z0)
    end

    if stable_phase == false
        tpdl = flash_tpd(xil, logϕl,z0,logϕ0)
        tpdv = flash_tpd(xiv, logϕv,z0,logϕ0)
        ΔG = (1 - β) * tpdl + β * tpdv
        if all(>(0),(ΔG,tpdv,tpdl))
            stable_phase = true # more detailed stability phase analisis
        end
    end

    if stable_phase == true
        mul_v = copy(xiv)
        mul_l = copy(xiv)
        initial_data = (model,p,t,z0)
        xil,xiv,vl,vv = trial_compositions!!(initial_data,xil,xiv,vl,vv,logϕl,logϕv,logϕ0,mul_v,mul_l,true;k0=k0)
        count = 1
        max_count = 10
        while count < max_count
            xil,xiv,vl,vv,tpdl,tpdv = trial_compositions!!(initial_data,xil,xiv,vl,vv,logϕl,logϕv,logϕ0,mul_v,mul_l,true)
            converged,stable_phase = trial_convergence(z0,xil,xiv,mul_l,mul_v,tpdl,tpdv,stable_phase)
            if converged == true
                break
            else
                count += 1
            end
        end
        if stable_phase == false
            for i in 1:3
                xil,xiv,vl,vv,tpdl,tpdv = trial_compositions!!(initial_data,xil,xiv,vl,vv,logϕl,logϕv,logϕ0,mul_v,mul_l,true)
            end
        end
    end

    if stable_phase == false
        vle_res = VLEResults(
            model = model
            ,vl = vl #liquid mol volume
            ,vv = vv #gas mol volume
            ,tv = t #temperature
            ,pv = p #pressure
            ,tl = t #temperature
            ,pl = p #pressure
            ,z0 = z0 #initial composition
            ,lnϕl = logϕl #efective log fugacity coef - liquid
            ,lnϕv = logϕv #efective log fugacity coef - gas
            ,nl = copy(xil) #moles of liquid
            ,nv = copy(xiv) #moles of gas
            ,xil = xil #mol fraction of liquid
            ,xiv = xiv) #mol fraction of gas
            #@show vle_res.xiv
            #@show vle_res.xil
            #@show vle_res.nv
            #@show vle_res.nl
            fg! = generate_ptflash_objective(vle_res,threaded=threaded)
            nv0 = zeros(length(xiv))
            nv0 .+= xiv
            nv0 .*= β
            opts = Optim.Options(g_tol = 1e-12,
                iterations = 10,
                store_trace = true,
                show_trace = true)
            optim_obj = OnceDifferentiable(only_fg!(fg!), nv0)
         try
            optimize(optim_obj,nv0, Optim.BFGS(linesearch = Optim.LineSearches.MoreThuente()))
         catch err
            @show vle_res
         finally
            return vtx_states(vle_res,moles)
         end
    end

    return state(t=t,v=vz0,xn = z0)
end

function trial_compositions!!(initial_data,xil,xiv,vl,vv,logϕl,logϕv,logϕ0,mul_v,mul_l,threaded=true;k0=nothing)
    model,p,t,z0 = initial_data
    if k0 !== nothing

        @! xiv .= z0 .* k0
        @! xil .= z0 ./ k0

        xil = normalizefrac!!(xil)
        xiv = normalizefrac!!(xiv)
        vl = nothing
        vv = nothing
        tpdv = convert(typeof(p),Inf)
        tpdl = convert(typeof(p),Inf)
    else
        vl,logϕl,vv,logϕv = update_logϕ_vl!!(model,p,t,z0,xil,xiv,logϕl,logϕv,nothing,nothing,threaded)
        mul_v = exp.(logϕ0 .- logϕv)
        mul_l = exp.(logϕ0 .- logϕl)
        @! xiv += z0 .* mul_v - xiv
        @! xil += z0 .* mul_l - xil
        xil = normalizefrac!!(xil)
        xiv = normalizefrac!!(xiv)
        tpdl = flash_tpd(xil, logϕl,z0,logϕ0)
        tpdv = flash_tpd(xiv, logϕv,z0,logϕ0)
        @show tpdl
        @show tpdv

    end
    return xil,xiv,vl,vv,tpdl,tpdv
end

function trial_convergence(z0,xil,xiv,mul_l,mul_v,tpdl,tpdv,stable_phase)
    new_stable_phase = stable_phase
    _ones = ones(length(mul_l))

    #this checks for:
    # a: the composition aproaches z0
    # b: variation between iterations is very small

    if mul_v ≈ _ones
        converged = true
     elseif mul_l ≈ _ones
         converged = true
     elseif xiv ≈ z0
        converged = true
    elseif xil ≈ z0
        converged = true
    elseif tpdl < 0
        converged = true
        new_stable_phase = false
    elseif tpdv < 0
        converged = true
        new_stable_phase = false
    else
        converged = false
    end
    return converged,new_stable_phase
end

function refine_k!!(model,p,t,z0,xil,xiv,vl,vv,k,logϕl,logϕv,lnk,threaded)
    vl,logϕl,vv,logϕv = update_logϕ_vl!!(model,p,t,z0,xil,xiv,logϕl,logϕv,vl,vv,threaded)
    #@show xil,xiv,vl,vv
    @! lnk .= logϕl .- logϕv
    @! k .= exp.(lnk)
   k = remove_maximums!!(k)
    β = flash_vfrac(RR(),k,z0)
    #@show β
    xil = flash_liquid!!(xil,k,z0,β)
    xiv = flash_vapor!!(xiv,k,z0,β)

    xil = normalizefrac!!(xil)
    xiv = normalizefrac!!(xiv)
    return xil,xiv,vl,vv,k,logϕl,logϕv,lnk,β

end

#update both liquid and gas fugacity coefficients, considering threading
function update_logϕ_vl!!(model,p,t,z0,xil,xiv,logϕl,logϕv,vl=nothing,vv=nothing,threaded=true)
    if threaded
        res_l = Threads.@spawn update_logϕ!!(model,logϕl,$p,$t,xil,$vl,:liquid)
        res_v = Threads.@spawn update_logϕ!!(model,logϕv,$p,$t,xiv,$vv,:vapor)
         vl,logϕl = fetch(res_l)
         vv,logϕv = fetch(res_v)
     else
         vl,logϕl = update_logϕ!!(model,logϕl,p,t,xil,vl,:liquid)
         vv,logϕv = update_logϕ!!(model,logϕv,p,t,xiv,vv,:vapor)
     end
     #@show logϕl
     #@show logϕv
     #@! lnk .= logϕl .- logϕv
     #@! k .= exp.(lnk)
     #@show k
     return vl,logϕl,vv,logϕv

end
#with values of p,t,x; returns a new logϕ,v
function update_logϕ!!(model,logϕ,p,t,x,v=nothing,phase=:liquid)
    x = normalizefrac!!(x)
    _vtx = QuickStates.vtx()
    _ptx = QuickStates.ptx()
    _v =  v_zero(_ptx,model,p,t,x,v;phase=phase)
    @! logϕ = fugacity_coeff_impl(_vtx,model,_v,t,x)
    return _v,logϕ
end

function flash_tpd(xi, lnϕ,x0,lnϕ0)
    f(xiᵢ,lnϕᵢ,x0ᵢ,lnϕ0ᵢ) = xiᵢ*(log(xiᵢ) + lnϕᵢ-log(x0ᵢ)- lnϕ0ᵢ)
    return mapreduce(f,+,xi, lnϕ,x0,lnϕ0)
end

Base.@kwdef mutable struct VLEResults{MODEL,S,V}
    model::MODEL
    vl::S #liquid mol volume
    vv::S #gas mol volume
    tv::S #temperature, gas
    pv::S #pressure, gas
    tl::S #temperature,liquid
    pl::S #pressure,liquid
    z0::V #initial composition
    lnϕl::V #efective log fugacity coef - liquid
    lnϕv::V#efective log fugacity coef - gas
    nl::V #moles of liquid
    nv::V #moles of gas
    xil::V #mol fraction of liquid
    xiv::V #mol fraction of gas
end

function vtx_states(obj::VLEResults,moles=1.0)
    xv = obj.xiv
    xl = obj.xil
    z0 = obj.z0
    β = (z0[1] - xl[1])/(xv[1]-xl[1])
    nv = β*moles
    nl = (1-β)*moles
    stᵥ = state(mol_v=obj.vv,t=obj.tv,xn=obj.xiv,moles=nv,phase=:gas)
    stₗ = state(mol_v=obj.vl,t=obj.tl,xn=obj.xil,moles=nl,phase=:liquid)
    return (stₗ,stᵥ)
end

function vtn_states(obj::VLEResults)
    stᵥ = state(mol_v=obj.vv,t=obj.tv,n=obj.nv,phase=:gas)
    stₗ = state(mol_v=obj.vl,t=obj.tl,n=obj.nl,phase=:liquid)
    return (stₗ,stᵥ)
end

function _flash_obj(xlᵢ,xvᵢ,lnϕlᵢ,lnϕvᵢ,β)
    _0 = zero(β)
    l = ifelse(xlᵢ>_0,(1-β)*xlᵢ*(log(xlᵢ)+lnϕlᵢ),_0)
    v = ifelse(xvᵢ>_0,β*xvᵢ*(log(xvᵢ)+lnϕvᵢ),_0)
    return l+v
end

function generate_ptflash_objective(obj::VLEResults;threaded = true)
    #TODO: implement threading
    function flash_objective!(F, G, nv)
        t = obj.tv
        p = obj.pv
        β = sum(nv)
        obj.nv = nv
        #@show nv
        obj.nl =obj.z0 - obj.nv
        #corr = minimum(obj.nl)
        #@show obj.nv
        #obj.nl = obj.z0 - obj.nv - corr
        #@show obj.nl
        obj.nl = remove_minimums!!(obj.nl)
        obj.nv = remove_minimums!!(obj.nv)
        obj.xil = normalizefrac!!(obj.nl)
        obj.xiv = normalizefrac!!(obj.nv)

        #@show obj.xil
        obj.vl ,obj.lnϕl = update_logϕ!!(obj.model,obj.lnϕl,p,t,obj.xil,obj.vl,:liquid)
        obj.vv ,obj.lnϕv = update_logϕ!!(obj.model,obj.lnϕv,p,t,obj.xiv,obj.vv,:gas)
        #@show obj.lnϕl
        #@show obj.lnϕv
        if !(G === nothing)
            #@show G
            G .= zero(eltype(G))
            for i in 1:length(G)
                if obj.xiv[i] <= eps(obj.xiv[i])
                    G[i] += obj.lnϕv[i]  - log(obj.xil[i])
                elseif obj.xil[i] <= eps(obj.xil[i])
                    G[i] += log(obj.xiv[i])   - obj.lnϕl[i]
                else
                    G[i] += obj.lnϕv[i] + log(obj.xiv[i])   - obj.lnϕl[i] - log(obj.xil[i])
                #G .+= obj.lnϕv .+ log.(obj.xiv)   .- obj.lnϕl .- log.(obj.xil)
                end
            #@show G
            end
        end
        if !(F === nothing)
            fobjective = (xlᵢ,xvᵢ,lnϕlᵢ,lnϕvᵢ)->_flash_obj(xlᵢ,xvᵢ,lnϕlᵢ,lnϕvᵢ,β)
            F = mapreduce(fobjective,+,obj.xil,obj.xiv,obj.lnϕl,obj.lnϕv)
            #@show F
            return F
        end
    end
    return flash_objective!
end

function equilibria(mt::MultiPT,model::HelmholtzModel,st::ThermodynamicState)
    t = temperature(FromState(),st)
    opts = options(FromState(),st)
    thr = get(opts,:threaded,true)
    p = pressure(FromState(),st)
    z0 = mol_fraction(FromState(),st,nothing,molecular_weight(model))
    _moles = moles(FromState(),st,u"mol",molecular_weight(model))
    return flash_impl(mt,model,p,t,z0,threaded=thr,moles=_moles)
end

#v = xiv*β
#l = z0 - v

#dfugv = I/β - 1/vt + dlnfugv/vt
#dfugl = I/(1-β) - 1/lt + dlnfugl/lt

#=

water-o2 (GERG2008)
logϕl = [-0.7705794332079066, 14.528869934800472]
logϕv = [-0.009360929108700455, 0.0031516545028221]
k = [0.46709692059495184, 2.034413807942005e6]

logϕl = [-0.7705794332079066, 14.528869934800472]
logϕv = [-0.009360929108700452, 0.0031516545028221037]
k = [0.46709692059495184, 2.034413807942005e6]

logϕl = [-0.7705816360739632, 14.539131948004911]
logϕv = [-0.007992505036874167, 0.0017762880364641798]
k = [0.4664571435209862, 2.0582273486877447e6]

logϕl = [-0.770581636063973, 14.539132013140865]
logϕv = [-0.007980392574881982, 0.0017656849290196598]
k = [0.4664514936154415, 2.0582493064751942e6]

logϕl = [-0.7705816360640014, 14.53913201314133]
logϕv = [-0.007980285727848792, 0.0017655915176444346]
k = [0.4664514437764727, 2.0582494987400582e6]

logϕl = [-0.7705816360640014, 14.53913201314133]
logϕv = [-0.007980284785338597, 0.0017655906936611884]
k = [0.4664514433368374, 2.05824950043602e6]

water-o2 (VdW)
logϕl = [2.297372915993078, 7.406207150097631]
logϕv = [-0.005105694210269807, -6.038786149844145e-5]
k = [9.998935228783932, 1646.2702241435495]

logϕl = [2.297372915993078, 7.406207150097631]
logϕv = [-0.005105694210269807, -6.038786149844145e-5]
k = [9.998935228783932, 1646.2702241435495]

logϕl = [2.297372728601342, 7.406704158470847]
logϕv = [-0.003828630603820233, -0.0005244515961122089]
k = [9.986172231292917, 1647.8531690781929]

logϕl = [2.297411298402489, 7.389040025797274]
logϕv = [-0.005202637502164627, 6.585518575033793e-5]
k = [10.000288433265684, 1618.0454227403882]

logϕl = [2.297412692566457, 7.388698667718173]
logϕv = [-0.005203259181312256, 6.678678750513239e-5]
k = [10.000308592298598, 1617.4916772651154]
=#


using Knet
using ArgParse
using JLD
using HDF5
using MAT
using JSON
using AutoGrad

include("lib/vocab.jl")
include("lib/base.jl")
include("lib/init.jl")
include("lib/model.jl")
include("lib/batch.jl")
include("lib/convnet.jl")
include("lib/train.jl")
include("lib/eval.jl")
include("lib/util.jl")
include("lib/data.jl")

function main(args)
    s = ArgParseSettings()
    s.description = string(
        "Show and Tell: A Neural Image Caption Generator",
        " Knet implementation by Ilker Kesen [ikesen16_at_ku.edu.tr], 2016.")

    @add_arg_table s begin
        # load/save files
        ("--images"; help="images JLD file")
        ("--captions"; help="captions zip file (shared by Karpathy)")
        ("--vocabfile"; help="vocabulary JLD file")
        ("--loadfile"; default=nothing; help="pretrained model file if any")
        ("--savefile"; default=nothing; help="model save file after train")
        ("--cnnfile"; help="pre-trained CNN MAT file")
        ("--extradata"; action=:store_true;
         help="use restval split for training")

        # model options
        ("--winit"; arg_type=Float32; default=Float32(0.08))
        ("--hidden"; arg_type=Int; default=512)
        ("--embed"; arg_type=Int; default=512)
        ("--convnet"; default="vgg16")
        ("--lastlayer"; default="relu7")
        ("--visual"; arg_type=Int; nargs='+'; default=[4096])
        ("--cnnmode"; arg_type=Int; default=1)

        # training options
        ("--nogpu"; action=:store_true)
        ("--epochs"; arg_type=Int; default=1)
        ("--batchsize"; arg_type=Int; default=200)
        ("--lr"; arg_type=Float32; default=Float32(0.001))
        ("--gclip"; arg_type=Float32; default=Float32(5.0))
        ("--seed"; arg_type=Int; default=-1; help="random seed")
        ("--gcheck"; arg_type=Int; default=0; help="gradient checking")
        ("--finetune"; action=:store_true; help="fine tune convnet")
        ("--optim"; default="Adam"; help="optimizer (Sgd|Adam|Adagrad)")
        ("--eps"; default=1e-6; help="epsilon for Adagrad optimizer")
        ("--adam"; action=:store_true; help="use adam optimizer")
        ("--decay"; arg_type=Float32; default=Float32(1.0); help="lr decay")
        ("--decayperiod"; arg_type=Int64; default=0)
        ("--fast"; action=:store_true; help="do not compute train loss")
        ("--saveperiod"; arg_type=Int; default=0)
        ("--newoptimizer"; action=:store_true)
        ("--evalmetric"; default="bleu")
        ("--beamsize"; arg_type=Int; default=1)
        ("--checkpoints"; arg_type=Int; default=1)
        ("--sortbylen"; action=:store_true)

        # dropout values
        ("--fc6drop"; arg_type=Float32; default=Float32(0.5))
        ("--fc7drop"; arg_type=Float32; default=Float32(0.5))
        ("--softdrop"; arg_type=Float32; default=Float32(0.5))
        ("--wembdrop"; arg_type=Float32; default=Float32(0.5))
        ("--vembdrop"; arg_type=Float32; default=Float32(0.5))
        ("--membdrop"; arg_type=Float32; default=Float32(0.5))
        ("--attdrop"; arg_type=Float32; default=Float32(0.5))
    end

    # parse args
    @printf("\nScript started. [%s]\n", now()); flush(STDOUT)
    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true); display(o); flush(STDOUT)

    # random seed
    s = o[:seed] > 0 ? srand(o[:seed]) : srand()
    @printf("\nrandom seed:\n")
    display(s.seed); println(); flush(STDOUT)

    # load vocabulary
    vocab = load(o[:vocabfile], "vocab")
    o[:vocabsize] = vocab.size
    println("Vocabulary loaded."); flush(STDOUT)

    # initialize state and weights
    o[:atype] = !o[:nogpu] ? KnetArray{Float32} : Array{Float32}
    prevscore = bestscore =
        o[:loadfile] == nothing ? Inf : load(o[:loadfile], "score")
    prevloss = bestloss =
        o[:loadfile] == nothing ? Inf : load(o[:loadfile], "lossval")
    w, srnn = get_weights(o)
    wcnn =  get_wcnn(o)
    if wcnn != nothing && o[:finetune]
        w["wcnn"] = wcnn
    end
    opts = get_opts(o, w)
    println("Model loaded."); flush(STDOUT)

    # get samples used during training process
    train, restval, val = get_entries(o[:captions], ["train", "restval", "val"])
    if o[:extradata]
        train = [train; restval]
    else
        val = [val; restval]
    end
    restval = 0

    # split samples into image/sentence pairs
    train = get_pairs(train)
    valid = get_pairs(val) # keep val for validation
    gc()
    const nsamples = length(train)
    const nbatches = div(nsamples, o[:batchsize])
    println("Data loaded."); flush(STDOUT)

    # gradient check
    # if o[:gcheck] > 0
    #     ids = shuffle([1:nsamples...])[1:o[:batchsize]]
    #     images, captions, masks = make_batch(o, train[ids], vocab)
    #     gradcheck(loss, w, copy(s), images, captions, masks;
    #               gcheck=o[:gcheck], verbose=true, atol=0.01)
    #     images
    #     gc()
    # end

    # checkpoints
    checkpoints = []

    # sort sequences
    if o[:sortbylen]
        sort!(train, by=i->length(i[2]), rev=true)
    end
    sort!(valid, by=i->length(i[2]), rev=true)
    offsets = collect(1:o[:batchsize]:nsamples+1)

    # training
    const saveperiod = o[:saveperiod] > 0 ? o[:saveperiod] : length(offsets)-1
    @printf("Training started (nsamples=%d, nbatches=%d, loss=%g, score=%g). [%s]\n",
            nsamples, nbatches, prevloss, prevscore, now()); flush(STDOUT)
    for epoch = 1:o[:epochs]
        t0 = now()

        # data split training
        losstrn = 0
        nwords  = 0
        if !o[:sortbylen]
            shuffle!(train)
            orders = [1:(length(offsets)-1)...]
        else
            orders = randperm(length(offsets)-1)
        end
        for (i,k) in enumerate(orders)
            iter = (epoch-1)*nbatches+i
            lower, upper = offsets[k:k+1]
            samples = train[lower:upper-1]
            images, x, y, batchsizes = make_batch(o, samples, vocab)
            this_loss, this_words = train!(
                w, srnn, images, x, y, batchsizes, opts, o)
            flush(STDOUT)
            images = 0; captions = 0; ans = 0; gc()
            losstrn += this_loss
            nwords  += this_words

            if iter % saveperiod == 0
                lossval = bulkloss(w,srnn,o,valid,vocab)
                @printf("\n(epoch/iter): %d/%d, loss: %g/%g [%s] ",
                        epoch, iter, losstrn/nwords, lossval, now())
                flush(STDOUT)
                # scores, bp, hlen, rlen =
                #     validate(w, val, vocab, o)
                # @printf("\nBLEU = %.1f/%.1f/%.1f/%.1f ",
                #         map(i->i*100,scores)...)
                # @printf("(BP=%g, ratio=%g, hyp_len=%d, ref_len=%d) [%s]\n",
                #         bp, hlen/rlen, hlen, rlen, now())
                # flush(STDOUT)
                # score = scores[end]
                score = lossval

                # learning rate decay
                decay!(o, opts, lossval, prevloss)
                prevscore = score
                prevloss  = lossval
                gc()

                # check and save best model
                score <= bestscore || continue

                path, ext = splitext(abspath(o[:savefile]))
                filename  = abspath(string(path, "-iter-", iter, ext))
                savemodel(o, w, opts, filename, score, lossval)
                @printf("Model saved to %s.\n", filename); flush(STDOUT)

                # keep track of checkpoints
                push!(checkpoints, (score, -lossval, filename))
                sort!(checkpoints, rev=true)
                if length(checkpoints) > o[:checkpoints]
                    _, _, worst = pop!(checkpoints)
                    rm(worst)
                end
                bestscore = checkpoints[end][1]
            end
        end # batches end

        t1 = now()
        elapsed = Int64(round(Float64(t1-t0)*0.001))
        @printf("\nepoch #%d finished. (time elapsed: %s)\n",
                epoch, pretty_time(elapsed))
        flush(STDOUT)
        gc()
    end # epoch end
end

function decay!(o, opts, lossval, prevloss)
    if o[:decay] < 1.0 && lossval > prevloss
        o[:lr] *= o[:decay]
        @printf("\nlr decay. new lr=%g\n", o[:lr]); flush(STDOUT)
        decay!(o[:lr], opts)
    end
end

function decay!(lr, opts::Dict)
    for k in keys(opts)
        decay!(lr, opts[k])
    end
end

function decay!(lr, opts::Array)
    for k in 1:length(opts)
        decay!(lr, opts[k])
    end
end

function decay!(lr, opt::Union{Knet.Adam,Knet.Sgd,Knet.Adagrad})
    opt.lr = lr
end

function savemodel(o, w, opts, filename, score, lossval)
    if o[:savefile] != nothing
        save(filename,
             "w", copy_weights(w),
             "opts", copy_opts(opts, w, true),
             "score", score,
             "lossval", lossval)
    end
end

function validate(w, data, vocab, o; metric=bleu, split="val")
    wcnn = o[:finetune] ? w[:wcnn] : nothing
    wdec = w
    hyp, ref = [], []
    s = initstate(o[:atype], o[:hidden], 1)

    hyp, ref = h5open(o[:images], "r") do f
        hyp, ref = [], []
        for entry in data
            image = read(f, entry["filename"])
            caption = generate(
                wdec, copy(s), image, vocab; beamsize=o[:beamsize])
            push!(hyp, caption)
            push!(ref, map(s->s["raw"], entry["sentences"]))
        end
        hyp, ref
    end

    return metric(hyp, ref)
end

!isinteractive() && !isdefined(Core.Main, :load_only) && main(ARGS)