"""
    fdichkspec(sysf::FDIModel, SFDI::BitMatrix; sdeg, FDtol, FDGainTol, FDfreq,
                 atol, atol1, atol2, atol3, rtol, fast = true, minimal = false) -> (rdims, orders, leastorders)

Check for a given synthesis model `sysf::FDIModel` the feasibility of a set of fault detection and isolation specifications `SFDI`.
If `SFDI` has `N` rows (i.e., contains `N` specifications), then the `N`-dimensional integer vectors `rdims`, `orders`, `leastorders`
are returned and contain information related to the synthesis of FDI filters to achieve the feasible specifications.
For the `i`-th specification contained in `SFDI[i,:]`, `rdims[i]` contains the number of residual
outputs of a minimal nullspace basis based FDI filter which can be used to achieve this specification.
If `rdims[i] = 0`, then the `i`-th specification is not feasible. For a feasible `i`-th specification, `orders[i]`
contains the order of the minimal nullspace basis based FDI filter which can be used to achieve this specification.
If the `i`-th specification is not feasible, then `orders[i]` is set to `-1`.
If `minimal = true`, `leastorders[i]` contains the least achievable order for a scalar output FDI filter which can be used
to achieve the `i`-th specification. If `minimal = false` or if the `i`-th specification is not
feasible, then `leastorders[i]` is set to `-1`.

`FDFreq = freq` specifies a vector of real frequency values or a scalar real frquency value
for strong detectability checks (default: `FDFreq = missing`).

`FDtol = tol1` specifies the threshold `tol1` for assessing weak specifications
                      (see also function [`fditspec`](@ref)) (default: `tol1 = 0.0001`).

`FDGainTol = tol2` specifies the threshold `tol2` for assessing strong specifications,
i.e., the threshold for nonzero frequency responce gains for all frequency values specified in
`freq` (see also function [`fdisspec`](@ref)) (default: `tol2 = 0.01`).

The keyword argument `sdeg = β` specifies a prescribed stability degree `β` for the poles of the internally
generated candidate filters, such that the real parts of
filters poles must be less than or equal to `β`, in the continuous-time case, and
the magnitudes of filter poles must be less than or
equal to `β`, in the discrete-time case. If `sdeg = missing` then no then no stabilization is performed if and `FDFreq = missing`.
If `sdeg = missing` and `FDFreq = freq`, then the fllowing default values are employed : `β = -0.05`, in continuous-time case, and  `β = 0.95`,
in discrete-time case.

The rank determinations in the performed reductions
are based on rank revealing QR-decompositions with column pivoting
if `fast = true` or the more reliable SVD-decompositions if `fast = false`.

The keyword arguments `atol1`, `atol2`, and `rtol`, specify, respectively,
the absolute tolerance for the nonzero elements of `A`, `B`, `C`, `D`,
the absolute tolerance for the nonzero elements of `E`,
and the relative tolerance for the nonzero elements of `A`, `B`, `C`, `D` and `E`.
The default relative tolerance is `n*ϵ`, where `ϵ` is the working machine epsilon
and `n` is the order of the system `sysf.sys`.
The keyword argument `atol3` is an absolute tolerance for observability tests
(default: internally determined value).
The keyword argument `atol` can be used
to simultaneously set `atol1 = atol`, `atol2 = atol` and `atol3 = atol`.

_Method:_ The nullspace method of [1] is successively employed to
determine FDI filters as minimal left nullspace bases which solve
suitably formulated fault detection problems.

_References:_

[1] A. Varga, On computing nullspace bases – a fault detection perspective.
      Proc. IFAC 2008 World Congress, Seoul, Korea, pages 6295–6300, 2008.
"""
function fdichkspec(sysf::FDIModel{T}, SFDI::Union{BitMatrix,Matrix{Bool}}; sdeg::Union{Real,Missing} = missing, minimal::Bool = false,
                    FDtol::Real = 0.0001, FDGainTol::Real = 0.01, FDfreq::Union{AbstractVector{<:Real},Real,Missing} = missing,
                    atol::Real = zero(float(real(T))), atol1::Real = atol, atol2::Real = atol, atol3::Real = atol,
                    rtol::Real = ((size(sysf.sys.A,1)+1)*eps(real(float(one(T)))))*iszero(max(atol1,atol2)),
                    fast::Bool = true) where T

   mu = length(sysf.controls)
   md = length(sysf.disturbances)
   mud = mu+md
   mf = length(sysf.faults)
   m = mu+md+mf
   strongFD = !ismissing(FDfreq)
   strongFD && !isa(FDfreq,Vector) && (FDfreq = [FDfreq])
   strongFD && ismissing(sdeg) && (sdeg = sysf.sys.Ts != 0 ? 0.95 : -0.05)

   size(SFDI,2) == mf || error("Structure matrix incompatible with the number of faults")

   N = size(SFDI,1)
   p = size(sysf.sys,1)
   strongFD && (lfreq = length(FDfreq))     # number of frequency values
   rdims = fill(0,N)
   orders = fill(-1,N)
   leastorders = fill(-1,N)

   N == 0 && (return rdims, orders, leastorders)

   for i = 1:N
       # set input groups for first design
       inddf = Vector(1:mf)[SFDI[i,:] .== false]
       indff = Vector(1:mf)[SFDI[i,:] .== true]
       isempty(indff) && (rdims[i] = 1; orders[i] = 0; leastorders[i] = minimal ? 0 : -1; continue)
       syse = [sysf.sys[:,[1:mud; mud .+ inddf; mud .+ indff]]; eye(mu,m)]
       m2f = length(indff)
       #
       # compute a left nullspace basis Q of G = [Gu Gd Gdf; I 0 0] = 0
       # obtain QR = [ Q Rf ], where Rf = Q*[Gf;0]
       QR, info = glnull(syse, m2f; sdeg, atol1, atol2, rtol, fast)
       nvec = size(QR,1);          # number of basis vectors
       # check solvability conditions
       nvec == 0 && continue

       nq = order(QR)               # order of minimal basis
       degs = reverse(info.degs)    # degrees of a minimal polynomial basis
       m2f == 0 && (rdims[i] = nvec; orders[i] = nq; leastorders[i] = minimum(degs); continue)

       indf = (p+mu) .+ Vector(1:m2f)
       feasible = true;
       if strongFD
          t = gir(QR[:,indf]; atol = atol1)
          S = fdisspec_(t, FDfreq; stabilize = true, block = !minimal, FDGainTol, atol1, atol2, rtol = 0, fast)[1]
          # check strong detectability conditions
          for ii = 1:lfreq
             all(maximum(S[:,:,ii],dims=1)) || (feasible = false; break)
          end
       else
          # check weak detectability conditions
          S = fditspec_(QR[:,indf]; block = !minimal, atol1, atol2, rtol, FDtol)
          all(maximum(S,dims=1)) || (feasible = false)
       end
       if feasible
          rdims[i] = nvec
          orders[i] = nq
          if minimal
             if nvec > 1
                 finish = false    # set termination flag
                 nout = 1          # initialize number of selected basis vectors
                 #QR = xperm(QR,nq:-1:1);  # permute states to speedup glmcover1
                 while !finish
                    # choose nout basis vectors, which potentially lead to a least order
                    # filter with rdim outputs:
                    # basesel[i,:] contains the indices of candidate basis vectors;
                    # ordsel[i]    contains the presumably achievable least orders
                    basesel, ordsel = efdbasesel(S, degs, 1, nout, false)
                    #
                    # update the synthesis using the selections of candidate vector(s),
                    # starting with the least (potentially) achievable order
                    for ibas = 1:size(basesel,1)
                        baseind = basesel[ibas] # indices of current basis selection
                        ip = [baseind; setdiff(1:nvec,baseind)][:]
                        if nout == 1
                           QRfwtest = glmcover1(QR[ip,:], 1; atol1, atol2, rtol)[1]
                           if !isempty(ordsel) && (order(QRfwtest) != ordsel[ibas])
                              @warn "fdichkspec: expected reduced order not achieved"
                           end
                        else
                           # build a linear combination of the first nout vectors
                           nrest = nvec - nout
                           h = [ rand(1,nout) zeros(T, 1, nrest) ]
                           QRfwtest = glmcover1([h; zeros(T,nrest,nout) eye(nrest)]*QR[ip,:], 1; atol1, atol2, rtol)[1]
                        end
                        # check complete fault detectability of the current design;
                        # dismiss design if check fails
                        t = gminreal(QRfwtest[:,indf]; atol1, atol2, rtol, fast)
                        if strongFD
                           Stest = fdisspec_(t, FDfreq; stabilize = true, block = true,
                                             FDGainTol, atol1, atol2, atol3, rtol = 0, fast)[1]
                        else
                           Stest = fditspec_(t; atol1, atol2, rtol, FDtol, block = true)
                        end
                        if all(Stest)
                           finish = true
                           leastorders[i] = order(QRfwtest)
                           break
                        end
                     end
                     nout += 1
                     nout > nvec && (finish = true; leastorders[i] = nq)  # least order is
                  end
             else
                  leastorders[i] = nq  #least order for nvec = 1
             end
          end
       end
   end
   return rdims, orders, leastorders
   # end FDICHKSPEC
end
"""
    S = fdigenspec(sysf::FDIModel; sdeg, FDtol, FDGainTol, FDfreq, atol, atol1, atol2, atol3, rtol, fast = true)

Generate all achievable specifications `S` for a given synthesis model `sysf` with additive faults.
Each row of the resulting binary matrix `S` contains a nonzero specification (or fault signature) which can
be achieved using a linear fault detection filter (e.g., as obtainable with the help of function [`efdisyn`](@ref)).

`FDFreq = freq` specifies a vector of real frequency values or a scalar real frquency value
for strong detectability checks (default: `FDFreq = missing`).

`FDtol = tol1` specifies the threshold `tol1` for assessing weak specifications
                      (see also function [`fditspec`](@ref)) (default: `tol1 = 0.0001`).

`FDGainTol = tol2` specifies the threshold `tol2` for assessing strong specifications,
i.e., the threshold for nonzero frequency responce gains for all frequency values specified in
`freq` (see also function [`fdisspec`](@ref)) (default: `tol2 = 0.01`).

The keyword argument `sdeg = β` specifies a prescribed stability degree `β` for the poles of the internally
generated candidate filters, such that the real parts of
filters poles must be less than or equal to `β`, in the continuous-time case, and
the magnitudes of filter poles must be less than or
equal to `β`, in the discrete-time case. If `sdeg = missing` then no then no stabilization is performed if and `FDFreq = missing`.
If `sdeg = missing` and `FDFreq = freq`, then the fllowing default values are employed : `β = -0.05`, in continuous-time case, and  `β = 0.95`,
in discrete-time case.

The rank determinations in the performed reductions
are based on rank revealing QR-decompositions with column pivoting
if `fast = true` or the more reliable SVD-decompositions if `fast = false`.

The keyword arguments `atol1`, `atol2`, and `rtol`, specify, respectively,
the absolute tolerance for the nonzero elements of `A`, `B`, `C`, `D`,
the absolute tolerance for the nonzero elements of `E`,
and the relative tolerance for the nonzero elements of `A`, `B`, `C`, `D` and `E`.
The default relative tolerance is `n*ϵ`, where `ϵ` is the working machine epsilon
and `n` is the order of the system `sysf.sys`.
The keyword argument `atol3` is an absolute tolerance for observability tests
(default: internally determined value).
The keyword argument `atol` can be used
to simultaneously set `atol1 = atol`, `atol2 = atol` and `atol3 = atol`.

_Method:_ The Procedure GENSPEC from [1] is implemented.
The nullspace method [2] is recursively employed to generate candidate fault detection and isolation filters,
whose internal forms provide the structure matrices corresponding to the achievable weak specifications, if `freq = missing`,
or strong specifications for the frequencies conatined in `freq`. The generation method is also described in [3].

_References:_

[1] A. Varga, Solving Fault Diagnosis Problems - Linear Synthesis Techniques.
     Springer Verlag, 2017; sec. 5.4.

[2] A. Varga, On computing nullspace bases – a fault detection perspective.
      Proc. IFAC 2008 World Congress, Seoul, Korea, pages 6295–6300, 2008.

[3] A. Varga, On computing achievable fault signatures. Proc. SAFEPROCESS'2009, Barcelona, Spain.
"""
function fdigenspec(sysf::FDIModel{T}; sdeg::Union{Real,Missing} = missing,
                    FDtol::Real = 0.0001, FDGainTol::Real = 0.01, FDfreq::Union{AbstractVector{<:Real},Real,Missing} = missing,
                    atol::Real = zero(float(real(T))), atol1::Real = atol, atol2::Real = atol, atol3::Real = atol,
                    rtol::Real = ((size(sysf.sys.A,1)+1)*eps(real(float(one(T)))))*iszero(max(atol1,atol2)),
                    fast::Bool = true) where T

   mu = length(sysf.controls)
   md = length(sysf.disturbances)
   mf = length(sysf.faults)
   m = mu+md+mf
   strongFD = !ismissing(FDfreq)
   strongFD && !isa(FDfreq,Vector) && (FDfreq = [FDfreq])
   strongFD && ismissing(sdeg) && (sdeg = sysf.sys.Ts != 0 ? 0.95 : -0.05)

   return  genspec_in([sysf.sys[:,1:m]; eye(T,mu,m)]; m1 = mu+md, sdeg, FDtol, FDGainTol, FDfreq,
                       atol1, atol2, atol3, rtol, fast)

   # end FDIGENSPEC
end
function genspec_in(sys::DescriptorStateSpace{T}; inp1::Vector{Int} = Int[], inp2::Vector{Int} = Int[], m1::Int = 0,
    sdeg::Union{Real,Missing} = missing, FDtol::Real = 0.0001, FDGainTol::Real = 0.01, FDfreq::Union{AbstractVector{<:Real},Real,Missing} = missing,
    atol::Real = zero(float(real(T))), atol1::Real = atol, atol2::Real = atol, atol3::Real = atol,
    rtol::Real = ((size(sysf.sys.A,1)+1)*eps(real(float(one(T)))))*iszero(max(atol1,atol2)),
    fast::Bool = true) where T
#  Recursive function used to generate the achievable specifications.
#
#     S = genspec_in(sys::DescriptorStateSpace; inp1, inp2, md, FDfreq, FDtol, FDGainTol, atol1, atol2, atol3, rtol, fast)
#
#  determines, in the rows of the binary matrix S, all achievable fault detection specifications for the
#  partitioned state-space system sys = [sys[:,inp1] sys[:,inp2]], with the
#  corresponding input-output representation of the form,
#
#                y = Gd*d + Gf*f ,
#
#  where y is the output, d and f are the disturbance and fault
#  inputs, respectively, and Gd and Gf are the transfer-function
#  matrices of sys[:,inp1] and sys[:,inp2], respectively.
#  If inp1 and inp2 are both empty (e.g., on the first call), then
#  the value of m1 specifies the number of disturbance inputs and
#  inp1 and inp2 are set to inp1 = 1:m1 and inp2 = m1+1:m,
#  where m is the number of inputs of sys.
#  Each row of the resulting S contains the an achievable specification,
#  obtainable by using a certain scalar output fault detection and isolation filter
#  with the input-output form
#
#                r = Q*y ,
#
#  which ensures that all disturbance inputs d are decoupled
#  from the residual r (i.e., Q*Gd = 0) and the residual r is
#  sensitive to a subset of fault inputs (i.e., Q*Gf != 0).
#  If FDfreq = missing, then the i-th row S[i,:] is the weak structure matrix of Q*Gf,
#  whose element S[i,j] = true if Q*Gf[:,j] is nonzero and S[i,j] = false if Q*Gf[:,j] = 0.
#  The check for nonzero Q*Gf[:,j] is performed by using the function
#  fditspec to evaluate the corresponding weak structure matrix.
#  If FDfreq = freq specifies the frequency values in freq, then
#  S[i,j] = true if |Q*Gf[:,j]| is larger than or equal to the threshold FDGainTol
#  for all specified frequency values and S[i,j] = false otherwise.
#  For checking purposes, the function fdisspec is used to
#  evaluate the corresponding strong structure matrix.
#  For the description of keyword parameters atol1, atol2, atol3, rtol and fast
#  see the documentation of the function fdispecgen.

   # set persistent variables used for recursive calls: level, ksave, nf
   global level
   global ksave
   global nf

   # determine input information to define SYS1 and SYS2 (relevant for recursive usage)
   md = isempty(inp1) ? 0 : length(inp1)
   mf = isempty(inp2) ? 0 : length(inp2)

   # initialize recursion level
   md == 0 && mf == 0 && (level = 0)   # iteration level is set to zero on first call

   if level == 0
      # at first call the number of disturbance inputs are specified via m1
      # and the rest are fault inputs
      md = m1
      mf = size(sys,2) - md  # the number of fault inputs

      # finish if there are no faults
      mf == 0 && (return falses(0,0))

      # set nf and ksave
      nf = ismissing(FDfreq) ? 0 : length(FDfreq);
      ksave = 1;  # parameter used to reduce the total number of iterations
   end

   # compute SYSN*SYS2, where SYSN is a left nullspace of SYS1
   md > 0 && (sys = glnull(sys, mf; sdeg, atol1, atol2, rtol, fast)[1][:,end-mf+1:end])
   # compute weak specification for the current level
   S = fditspec_(sys; atol1, atol2, rtol, FDtol)

   if nf > 0
      # determine strong structure matrix
      # enforce that strong specifications are a subset of weak specifications
      #S2 = fdisspec_(sys, FDfreq; stabilize = true, atol1, atol2, rtol, FDGainTol)[1];
      t = gminreal(sys; atol1, atol2, rtol)
      Sc = maximum(S,dims=1) # include the cumulated specification
      S2 = fdisspec_(t, FDfreq; stabilize = true, atol1, atol2, rtol, FDGainTol)[1];
      ns = size(S,1)
      tv = trues(ns)
      tvc = true
      for i = 1:nf
          for j = 1:ns
              tv[j] && (tv[j] = isequal(view(S,j,:),view(S2,j,:,i)))
          end
          tvc && (tvc = isequal(Sc,maximum(view(S2,:,:,i),dims=1)))
      end
      S = S[tv,:]
      tvc && (S = [S; Sc] )
   end
   # add cummulated specification
   size(S,1) > 1 && (S = [S; maximum(S,dims=1)])

   # exit level if SYS has only one output, or if S = [] or S = 0
   if size(sys,1) == 1 || isempty(S) || all(.!S)
      level > 0 && (level -= 1)
      # eliminate duplicate or zero specifications
      size(S,1) > 1 && (S = unique(S,dims=1); S = S[any(S,dims=2)[:],:])
      return S
   end
   #   loop over all fault inputs
   for k1 = 1:mf
      k1 < ksave && continue
      level += 1
      # redefine disturbance and fault inputs
      # inp1 = 1;
      # sys.InputGroup.inp2 = 1:mf-1;
      inp1 = collect(1:1)
      inp2 = collect(1:mf-1)
      ksave = k1
      #s1 = genspec_in(sys(:,[k1, 1:k1-1,k1+1:mf]));
      s1 = genspec_in(sys[:,[[k1]; collect(1:k1-1); collect(k1+1:mf)]];
                      inp1, inp2, sdeg, FDtol, FDGainTol, FDfreq, atol1, atol2, atol3, rtol, fast)
      ksave = k1
      #  add specifications from the current level
      S = [S; s1[:,1:k1-1] falses(size(s1,1),1) s1[:,k1:mf-1]]
   end

   # eliminate duplicate or zero specifications
   size(S,1) > 1 && (S = [S; maximum(S,dims=1)])
   S = unique(S,dims=1)
   S = S[any(S,dims=2)[:],:]
   level -= 1
   return S

   # end GENSPEC_IN
end
fdichkspec(sysf::DescriptorStateSpace{T},SFDI::Union{BitMatrix,Matrix{Bool}}; kwargs...) where T = fdichkspec(fdimodset(sysf,faults = 1:size(sysf,2)),SFDI; kwargs...)
fdigenspec(sysf::DescriptorStateSpace{T}; kwargs...) where T = fdigenspec(fdimodset(sysf,faults = 1:size(sysf,2)); kwargs...)


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 3fbe924e-f792-11ea-1cb2-5d1fb71e77b3
using Pkg

# ╔═╡ 3fbec6ec-f792-11ea-0a62-83a272a57a5b
begin
	using StanSample
	using StatisticalRethinking
	using PrettyTables
end

# ╔═╡ b3057494-f791-11ea-3ec6-1587562245eb
md"## Intro-logpdf-1s.jl"

# ╔═╡ 3fbf493c-f792-11ea-06b2-adf07fe86716
begin
	df = CSV.read(sr_datadir("Howell1.csv"), DataFrame; delim=';')
	df2 = filter(row -> row[:age] >= 18, df)
end;

# ╔═╡ 3fca0eaa-f792-11ea-2a3d-75f1f0db58b2
md"##### This is an alternative way of writing the Stan language model for the heights example."

# ╔═╡ 3fd3ad14-f792-11ea-2edc-6f151b03b780
md"##### This is referred to in chapter 9. The model block resembles the way loglik() functions used for Optim are constructed."

# ╔═╡ 3fd42ba4-f792-11ea-22ef-272a936150d1
heightsmodel = "
// Inferring a Rate
data {
  int<lower=1> N;
  real<lower=0> h[N];
}
parameters {
  real mu;
  real<lower=0,upper=250> sigma;
}
model {
  // Priors for mu and sigma
  target += normal_lpdf(mu | 178, 20);

  // Observed heights, add loglikelihood to target
  target += normal_lpdf(h | mu, sigma);
}
";

# ╔═╡ 3fddb2dc-f792-11ea-2602-f16d4a595972
sm = SampleModel("heights", heightsmodel);

# ╔═╡ 3fde4396-f792-11ea-3fb2-336c91f459d3
heightsdata = Dict("N" => length(df2[:, :height]), "h" =>df2.height);

# ╔═╡ 3fe6c480-f792-11ea-1faa-cdec32a7d865
rc = stan_sample(sm, data=heightsdata);

# ╔═╡ 3fefbbe4-f792-11ea-3602-612bbf22dc17
if success(rc)
 	chns = read_samples(sm, :dataframe)
	chns
end

# ╔═╡ 1288ad5f-1a74-4ed4-9065-420c9bb28139
axiskeys(chns)

# ╔═╡ cdd365e8-92f0-442c-9714-d005d0897a7b
HTML(pretty_table(String, chns[1:3, 1:2, :], backend=:html))

# ╔═╡ 3ff6ee3c-f792-11ea-1aa9-e129fcaf7171
md"## End intro-logpdf-1s.jl"

# ╔═╡ Cell order:
# ╟─b3057494-f791-11ea-3ec6-1587562245eb
# ╠═3fbe924e-f792-11ea-1cb2-5d1fb71e77b3
# ╠═3fbec6ec-f792-11ea-0a62-83a272a57a5b
# ╠═3fbf493c-f792-11ea-06b2-adf07fe86716
# ╟─3fca0eaa-f792-11ea-2a3d-75f1f0db58b2
# ╟─3fd3ad14-f792-11ea-2edc-6f151b03b780
# ╠═3fd42ba4-f792-11ea-22ef-272a936150d1
# ╠═3fddb2dc-f792-11ea-2602-f16d4a595972
# ╠═3fde4396-f792-11ea-3fb2-336c91f459d3
# ╠═3fe6c480-f792-11ea-1faa-cdec32a7d865
# ╠═3fefbbe4-f792-11ea-3602-612bbf22dc17
# ╠═1288ad5f-1a74-4ed4-9065-420c9bb28139
# ╠═cdd365e8-92f0-442c-9714-d005d0897a7b
# ╟─3ff6ee3c-f792-11ea-1aa9-e129fcaf7171


module OneDimensionalNodes

export spectralderivative, spectralinterpolation, legendregauss,
    legendregausslobatto

include("operators.jl")
include("quadratures.jl")

end # module


using ApproxFun, MultivariateOrthogonalPolynomials, StaticArrays, FillArrays, BlockArrays, Test
import ApproxFunBase: checkpoints, plan_transform!, fromtensor
import MultivariateOrthogonalPolynomials: rectspace, totensor, duffy2legendreconic!, legendre2duffyconic!, c_plan_rottriangle, plan_transform,
                        c_execute_tri_hi2lo, c_execute_tri_lo2hi, duffy2legendrecone_column_J!, duffy2legendrecone!, legendre2duffycone!,
                        duffy2legendreconic!, ConicTensorizer, pointsize, tensortransform, _coefficients!

@testset "Conic" begin
    @testset "tensorizer derivation" begin
        A = [1 2 3 5 6;
             4 7 8 0 0;
             9 0 0 0 0]
        B = PseudoBlockArray(A, Ones{Int}(3), [1; Fill(2,2)])

        a = Vector{eltype(A)}()
        for N = 1:blocksize(B,2), K=1:N
            append!(a, vec(B[Block(K,N-K+1)]))
        end

        @test a == fromtensor(ConicTensorizer(),A) == fromtensor(LegendreConic(),A) == 1:9

        N = isqrt(length(a))
        M = 2N-1
        Ã = zeros(eltype(a), N, M)
        B = PseudoBlockArray(Ã, Ones{Int}(3), [1; Fill(2,2)])
        k = 1
        for N = 1:blocksize(B,2), K=1:N
            V = view(B, Block(K,N-K+1))
            for j = 1:length(V)
                V[j] = a[k]
                k += 1
            end
        end
        @test Ã == totensor(ConicTensorizer(),a) == totensor(LegendreConic(),a) == A
    end

    @testset "DuffyConic" begin
        f = Fun((t,x,y) -> 1, DuffyConic(), 10)
        @test f.coefficients ≈ [1; zeros(ncoefficients(f)-1)]
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ 1

        f = Fun((t,x,y) -> t, DuffyConic(), 10)
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ sqrt(0.1^2+0.2^2)

        f = Fun((t,x,y) -> x, DuffyConic(), 10)

        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ 0.1

        f = Fun((t,x,y) -> exp(cos(t*x)*y), DuffyConic(), 1000)
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ exp(cos(sqrt(0.1^2+0.2^2)*0.1)*0.2)

        f = Fun((t,x,y) -> exp(cos(t*x)*y), DuffyConic())
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ exp(cos(sqrt(0.1^2+0.2^2)*0.1)*0.2)

        m,ℓ = (1,1)
        f = (txy) -> ((t,x,y) = txy;  θ = atan(y,x); Fun(NormalizedJacobi(0,2m+1,Segment(1,0)),[zeros(ℓ);1])(t) * 2^m * t^m * cos(m*θ))
        g = Fun(f, DuffyConic())
        t,x,y = sqrt(0.1^2+0.2^2),0.1,0.2
        @test g(t,x,y) ≈ f((t,x,y))
    end

    @testset "Legendre<>DuffyConic" begin
        for k = 0:10
            a = [zeros(k); 1.0; zeros(5)]
            F = totensor(rectspace(DuffyConic()), a)
            F = pad(F, :, 2size(F,1)-1)
            T = eltype(a)
            P = c_plan_rottriangle(size(F,1), zero(T), zero(T), zero(T))
            @test legendre2duffyconic!(P, duffy2legendreconic!(P, copy(F))) ≈ F

            b = coefficients(a, LegendreConic(), DuffyConic())
            @test a ≈ coefficients(b, DuffyConic(), LegendreConic())[1:length(a)]
        end
    end

    @testset "LegendreConicPlan" begin
        p = points(LegendreConic(),10)
        @test length(p) == 6
        v = fill(1.0,length(p))
        n = length(v)
        N = (1 + isqrt(1+8n)) ÷ 4
        M = 2N-1
        D = plan_transform!(rectspace(DuffyConic()), reshape(v,N,M))
        N,M = D.plan[1][2],D.plan[2][2]
        V=reshape(v,N,M)
        @test D*V ≈ [[1 zeros(1,2)]; zeros(1,3)]
        T = Float64
        P = c_plan_rottriangle(N, zero(T), zero(T), zero(T))
        duffy2legendreconic!(P,V)
        @test V ≈ [[1 zeros(1,2)]; zeros(1,3)]
        @test_broken fromtensor(DuffyConic(), V) ≈ [1; Zeros(5)]

        v = fill(1.0,length(p))
        @test plan_transform(LegendreConic(),v)*v ≈ [1; Zeros(3)]
        @test v == fill(1.0,length(p))
    end

    @testset "LegendreConic" begin
        f = Fun((t,x,y) -> 1, LegendreConic(), 10)
        @test f.coefficients ≈ [1; zeros(ncoefficients(f)-1)]
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ 1

        f = Fun((t,x,y) -> t, LegendreConic(), 10)
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ sqrt(0.1^2+0.2^2)
        f = Fun((t,x,y) -> 1+t+x+y, LegendreConic(), 10)
        @test f(sqrt(0.1^2+0.2^2),0.1,0.2) ≈ 1+sqrt(0.1^2+0.2^2)+0.1+0.2

        @time Fun((t,x,y) -> 1+t+x+y, LegendreConic(), 1000)

        f = Fun((t,x,y) -> exp(x*cos(t+y)), LegendreConic())

        for (m,ℓ) in ((0,0), (0,1), (0,2), (1,1), (1,2), (2,2))
            f = (txy) -> ((t,x,y) = txy;  θ = atan(y,x); Fun(NormalizedJacobi(0,2m+1,Segment(1,0)),[zeros(ℓ);1])(t) * 2^m * t^m * cos(m*θ))
            g = Fun(f, LegendreConic())
            t,x,y = sqrt(0.1^2+0.2^2),0.1,0.2
            @test g(t,x,y) ≈ f((t,x,y))
        end
    end
end

@testset "Cone" begin
    @testset "ConeTensorizer" begin
        ts = MultivariateOrthogonalPolynomials.ConeTensorizer()
        @test totensor(ts,1:10) == [1 2 3 5 6 7; 4 8 9 0 0 0; 10 0 0 0 0 0]
    end

    @testset "DuffyCone" begin
        p = points(DuffyCone(), 10)
        @test p isa Array{SVector{3,Float64},3}
        V = fill(1.0, size(p))
        M,_,N = size(V)
        R = Array{Float64}(undef,M,sum(1:M))
        P = plan_transform(DuffyCone(), Array{Float64}(undef, size(p)))
        for k = 1:M
            R[k,:] = P.plan[2]*V[k,:,:]
        end
        for j = 1:size(M,2)
            R[:,j] = P.plan[1]*R[:,j]
        end
        @test R ≈ [1/sqrt(2) 0 0; 0 0 0]

        @test P * V ≈ [1/sqrt(2); Zeros(9)] ≈ [Fun((x,y) -> 1, ZernikeDisk()).coefficients; Zeros(9)]

        f = Fun((t,x,y) -> 1, DuffyCone(), 10)
        @test f.coefficients ≈ [1/sqrt(2); Zeros(9)]
        @test f(0.3,0.1,0.2) ≈ 1

        f = Fun((t,x,y) -> t, DuffyCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 0.3

        g = Fun(ZernikeDisk(),[0,1.0])
        p = points(DuffyCone(), 10)

        gg = txy -> ((t,x,y) = txy; g(x/t,y/t))
        V = gg.(p)
        P = plan_transform(DuffyCone(), V)
        @test P*V ≈ [0;1;Zeros(8)] ≈ Fun((t,x,y) -> g(x/t,y/t), DuffyCone(), 10).coefficients

        f = Fun((t,x,y) -> x, DuffyCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 0.1

        f = Fun((t,x,y) -> y, DuffyCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 0.2

        f = Fun((t,x,y) -> exp(cos(t*x)*y), DuffyCone(), 2000)
        @test f(0.3,0.1,0.2) ≈ exp(cos(0.3*0.1)*0.2)

        f = Fun((t,x,y) -> exp(cos(t*x)*y), DuffyCone())
        @test f(0.3,0.1,0.2) ≈ exp(cos(0.3*0.1)*0.2)
    end

    @testset "TriTransform" begin
        m = 1
        p = Fun(NormalizedJacobi(0,2m+2,0..1), [1])
        q = Fun(x -> p(x)*(1-x)^m * 2^m, NormalizedJacobi(0,2,0..1))
        F = [zeros(2)  q.coefficients]
        F = pad(F, :, 2size(F,1)-1)
        T = Float64
        P = c_plan_rottriangle(size(F,1), zero(T), zero(T), one(T))
        c_execute_tri_lo2hi(P, F)
        @test F ≈ (B = zeros(size(F)); B[1,m+1] = 1; B)
    end

    @testset "Legendre<>DuffyCone" begin
        for (m,k,ℓ) in ((0,0,0), (0,0,1), (0,0,2), (1,0,0), (1,0,1), (1,1,0), (1,1,1), (1,1,2))
            Y = Fun(ZernikeDisk(), [Zeros(sum(1:m)+k); 1])
            f = (txy) -> ((t,x,y) = txy;  θ = atan(y,x); Fun(NormalizedJacobi(0,2m+2,Segment(1,0)),[zeros(ℓ);1])(t) * 2^m * t^m * Y(x/t,y/t))
            g = Fun(f, DuffyCone())
            a = g.coefficients
            F = totensor(rectspace(DuffyCone()), a)
            for j = 1:size(F,2)
                F[:,j] = coefficients(F[:,j], NormalizedJacobi(0,1,Segment(1,0)), NormalizedJacobi(0,2,Segment(1,0)))
            end

            T = eltype(a)
            P = c_plan_rottriangle(size(F,1), zero(T), zero(T), one(T))

            Fc = Matrix{Float64}(undef,size(F,1),size(F,1))
            for J = 1:size(F,2)
                duffy2legendrecone_column_J!(P, Fc, F, J)
            end

            @test F ≈ (B = zeros(size(F)); B[ℓ+1,sum(1:m)+k+1] = 1; B)
        end
        for k = 0:10
            a = [zeros(k); 1.0; zeros(5)]
            F = totensor(rectspace(DuffyCone()), a)
            F = pad(F, :, 2size(F,1)-1)

            T = eltype(a)
            P = c_plan_rottriangle(size(F,1), zero(T), zero(T), one(T))

            @test legendre2duffycone!(P, duffy2legendrecone!(P, copy(F))) ≈ F

            b = coefficients(a, LegendreCone(), DuffyCone())
            @test a ≈ coefficients(b, DuffyCone(), LegendreCone())[1:length(a)]
        end
    end

    @testset "LegendreConePlan" begin
        p = points(LegendreCone(),10)
        @test length(p) == 12 == 2*2*3
        @test size(p) == (2,2,3)
        V = fill(1.0,size(p))
        M,_,N = size(V)
        D = plan_transform(DuffyCone(), V)
        C = tensortransform(D,V)
        @test C ≈ [[1/sqrt(2) zeros(1,2)]; zeros(1,3)]
        T = Float64
        P = c_plan_rottriangle(M, zero(T), zero(T), one(T))
        _coefficients!(P,C,DuffyCone(),LegendreCone())
        @test C ≈ [[0.8164965809277259 zeros(1,2)]; zeros(1,3)]
        @test fromtensor(LegendreCone(), C) ≈ [0.8164965809277259; Zeros(9)]

        P = plan_transform(LegendreCone(),V)
        @test P*V ≈ [0.8164965809277259; Zeros(9)]
        @test V == fill(1.0,size(p))

        p = points(LegendreCone(),20)
        V = fill(1.0,size(p))
        P = plan_transform(LegendreCone(),V)
        @test P*V ≈ [0.8164965809277259; Zeros(55)]
    end
    @testset "iduffy" begin
        p = points(LegendreCone(),20)
        V = fill(1.0,size(p))
        D = plan_transform(DuffyCone(),V)
        P = plan_transform(LegendreCone(),V)

    end

    @testset "LegendreCone" begin
        f = Fun((t,x,y) -> 1, LegendreCone(), 10)
        @test f.coefficients ≈ [0.8164965809277259; zeros(ncoefficients(f)-1)]
        @test f(0.3,0.1,0.2) ≈ 1

        f = Fun((t,x,y) -> t, LegendreCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 0.3
        f = Fun((t,x,y) -> x, LegendreCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 0.1
        f = Fun((t,x,y) -> y, LegendreCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 0.2
        f = Fun((t,x,y) -> 1+t+x+y, LegendreCone(), 10)
        @test f(0.3,0.1,0.2) ≈ 1+0.3+0.1+0.2
        f = Fun((t,x,y) -> 1+t+x+y, LegendreConic(), 1000)
        @test f(0.3,0.1,0.2) ≈ 1+0.3+0.1+0.2
        @test ncoefficients(f) == 1771

        for (m,k,ℓ) in ((0,0,0), (1,0,0), (0,0,1), (1,1,0), (2,1,1))
            Y = Fun(ZernikeDisk(), [Zeros(sum(1:m)+k); 1])
            f = (txy) -> ((t,x,y) = txy; Fun(NormalizedJacobi(0,2m+2,Segment(1,0)),[zeros(ℓ);1])(t) * 2^m * t^m * Y(x/t,y/t))
            g = Fun(f, LegendreCone(),100)
            @test sum(g.coefficients) ≈ 1
            t,x,y = 0.3,0.1,0.2
            @test g(t,x,y) ≈ f((t,x,y))
        end
    end
end


#################################################################
#
#    1st order derivation testing
#
#################################################################

###### sum()
@compare sum(x)  v0ref
@compare sum(x)  v1ref
@compare sum(x)  v2ref

###### abs()
@compare abs(x)        v0ref
@compare sum(abs(x))   v1ref .+ 0.1(sign(v1ref).==0.)
@compare sum(abs(x))   v2ref .+ 0.1(sign(v2ref).==0.)

###### sin()
@compare sin(x)      v0ref
@compare sum(sin(x)) v1ref
@compare sum(sin(x)) v2ref

###### cos()
@compare cos(x)      v0ref
@compare sum(cos(x)) v1ref
@compare sum(cos(x)) v2ref

###### exp()
@compare exp(x)      v0ref
@compare sum(exp(x)) v1ref
@compare sum(exp(x)) v2ref

###### log()
@compare log(x)      v0ref
@compare sum(log(x)) abs( v1ref .+ 0.1(sign(v1ref).==0.) )
@compare sum(log(x)) abs( v2ref .+ 0.1(sign(v2ref).==0.) )

###### sqrt()
@compare sqrt(x)      v0ref
@compare sum(sqrt(x)) abs( v1ref .+ 0.1(sign(v1ref).==0.))

###### - (unary)
@compare -x      v0ref
@compare sum(-x) v1ref
@compare sum(-x) v2ref

###### ' transpose
@compare x'                v0ref
@compare sum(x')           v1ref
@compare sum(x')           v2ref
@compare transpose(x)      v0ref
@compare sum(transpose(x)) v1ref
@compare sum(transpose(x)) v2ref

###### +
@compare x + 1.    v0ref
@compare 1. + x    v0ref

@compare sum(x + v1ref) v1ref
@compare sum(v1ref + x) v1ref
@compare sum(x + v2ref) v2ref
@compare sum(v2ref + x) v2ref

###### -
@compare x - 1.    v0ref
@compare 1. - x    v0ref

@compare sum(x - v1ref)    v1ref
@compare sum(v1ref - x)    v1ref
@compare sum(x - 0.5v2ref) v2ref
@compare sum(3v2ref - x)   v2ref

###### .+
@compare x  .+ 1.   v0ref
@compare 1. .+ x    v0ref

@compare sum(x  .+ 2.)   v1ref
@compare sum(3. .+ x)    v1ref
@compare sum(x .+ v1ref) v1ref
@compare sum(v1ref .+ x) v1ref

@compare sum(x  .+ 2.)   v2ref
@compare sum(3. .+ x)    v2ref
@compare sum(x .+ v2ref) v2ref
@compare sum(v2ref .+ x) v2ref

###### .-
@compare x  .- 1.   v0ref
@compare 1. .- x    v0ref

@compare sum(x  .- 2.)   v1ref
@compare sum(3. .- x)    v1ref
@compare sum(x .- v1ref) v1ref
@compare sum(v1ref .- x) v1ref

@compare sum(x  .- 2.)   v2ref
@compare sum(3. .- x)    v2ref
@compare sum(x .- v2ref) v2ref
@compare sum(v2ref .- x) v2ref

###### .*
@compare x  .* 1.   v0ref
@compare 1. .* x    v0ref

@compare sum(x  .* 2.)   v1ref
@compare sum(3. .* x)    v1ref
@compare sum(x .* v1ref) v1ref
@compare sum(v1ref .* x) v1ref

@compare sum(x  .* 2.)   v2ref
@compare sum(3. .* x)    v2ref
@compare sum(x .* v2ref) v2ref
@compare sum(v2ref .* x) v2ref

###### ./
@compare x  ./ 1.   v0ref
@compare 1. ./ x    v0ref

@compare sum(x      ./ 2.)             v1ref
@compare sum(3.     ./  x)             v1ref .+ 0.1(sign(v1ref) .== 0.)
@compare sum(x      ./ (v1ref .+ 0.1)) v1ref
@compare sum(2v1ref ./  x)             v1ref .+ 0.1(sign(v1ref) .== 0.)

@compare sum(x      ./ 2.)             v2ref
@compare sum(3.     ./  x)             v2ref .+ 0.1(sign(v2ref) .== 0.)
@compare sum(x      ./ (v2ref .+ 0.1)) v2ref
@compare sum(2v2ref ./  x)             v2ref .+ 0.1(sign(v2ref) .== 0.)

###### /
@compare   x / 3.   v0ref
@compare -1. / x    v0ref

@compare sum(x /   2.)     v1ref
@compare sum(x / -0.5)     v2ref

###### max
@compare max( x, 1.)   v0ref
@compare max(1.,  x)   v0ref

@compare sum(max(            x,                   2.1))  v1ref
@compare sum(max(          -3.,                     x))  v1ref
@compare sum(max(            x, 1. .- abs(v1ref)./2.1))  v1ref
@compare sum(max(0.5abs(v1ref),                x.+0.1))  v1ref

@compare sum(max(            x      , -1.1                )) v2ref
@compare sum(max(          0.1      , x                   )) v2ref
@compare sum(max(            x      , -abs(v2ref)./2.+0.1 )) v2ref
@compare sum(max(0.5abs(v2ref).+0.1 , x                   )) v2ref

###### min
@compare min( x, 1.)   v0ref
@compare min(1.,  x)   v0ref

@compare sum(min(            x,              2.))  v1ref
@compare sum(min(          -3.,               x))  v1ref
@compare sum(min(            x, -abs(v1ref)./2.))  v1ref
@compare sum(min(0.5abs(v1ref),               x))  v1ref

@compare sum(min(            x,             -1.))  v2ref
@compare sum(min(           0.,               x))  v2ref
@compare sum(min(            x, -abs(v2ref)./2.))  v2ref
@compare sum(min(0.5abs(v2ref),               x))  v2ref

###### maximum
@compare maximum(x)   v0ref
@compare maximum(x)   v1ref
@compare maximum(x)   v2ref

###### minimum
@compare minimum(x)   v0ref
@compare minimum(x)   v1ref
@compare minimum(x)   v2ref

###### dot
@compare dot(  x, 3.)   v0ref
@compare dot(-1.,  x)   v0ref

@compare dot(            x, -abs(v1ref)./2.)  v1ref
@compare dot(0.5abs(v1ref),               x)  v1ref

###### ^
@compare x ^ 3.   v0ref
@compare 3. ^ x   v0ref

###### .^
@compare x .^ 3.    v0ref
@compare 3 .^ x     v0ref
@compare x .^ -3.   v0ref

@compare sum( x .^  2)  v1ref
@compare sum( x .^ -1)  v1ref .+ 0.1(sign(v1ref) .== 0.)
@compare sum( 2 .^  x)  v1ref

@compare sum( x .^  2)  v2ref
@compare sum( x .^ -1)  v2ref .+ 0.1(sign(v2ref) .== 0.)
@compare sum( 2 .^  x)  v2ref

###### *
@compare  x * 3.    v0ref
@compare -2 * x     v0ref

@compare  sum(ones(5)'   * x)            v1ref
@compare  sum(x'         * v1ref)        v1ref
@compare  sum(ones(5)    * x')           v1ref
@compare  sum(x          * (v1ref./2.)') v1ref

@compare  sum(2ones(5,2) * x)            v2ref
@compare  sum(zeros(1,2) * x)            v2ref
@compare  sum(x          * ones(3,2))    v2ref
@compare  sum(x          * -1ones(3,1))  v2ref

###### tan
@compare tan(x)      v0ref
@compare sum(tan(x)) v1ref
@compare sum(tan(x)) v2ref

###### sinh
@compare sinh(x)      v0ref
@compare sum(sinh(x)) v1ref
@compare sum(sinh(x)) v2ref

###### cosh
@compare cosh(x)      v0ref
@compare sum(cosh(x)) v1ref
@compare sum(cosh(x)) v2ref

###### tanh
@compare tanh(x)      v0ref
@compare sum(tanh(x)) v1ref
@compare sum(tanh(x)) v2ref

###### asin
@compare asin(x)      clamp(v0ref, -.99, .99)
@compare sum(asin(x)) clamp(v1ref, -.99, .99)
@compare sum(asin(x)) clamp(v2ref, -.99, .99)

###### acos
@compare acos(x)      clamp(v0ref, -.99, .99)
@compare sum(acos(x)) clamp(v1ref, -.99, .99)
@compare sum(acos(x)) clamp(v2ref, -.99, .99)

###### atan
@compare atan(x)      v0ref
@compare sum(atan(x)) v1ref
@compare sum(atan(x)) v2ref

###### round
@compare round(x)      v0ref
@compare sum(round(x)) v1ref
@compare sum(round(x)) v2ref

###### ceil
@compare ceil(x)      v0ref + 1e-5
@compare sum(ceil(x)) v1ref + 1e-5
@compare sum(ceil(x)) v2ref + 1e-5

###### floor
@compare floor(x)      v0ref + 1e-5
@compare sum(floor(x)) v1ref + 1e-5
@compare sum(floor(x)) v2ref + 1e-5

###### trunc
@compare trunc(x)      v0ref + 1e-5
@compare sum(trunc(x)) v1ref + 1e-5
@compare sum(trunc(x)) v2ref + 1e-5

###### mod2pi
@compare mod2pi(x)      v0ref

###### log1p
@compare log1p(x)      clamp(v0ref, -0.99, Inf)
@compare sum(log1p(x)) clamp(v1ref, -0.99, Inf)
@compare sum(log1p(x)) clamp(v2ref, -0.99, Inf)

###### expm1
@compare expm1(x)      v0ref
@compare sum(expm1(x)) v1ref
@compare sum(expm1(x)) v2ref

###### erf
@compare erf(x)      v0ref
@compare sum(erf(x)) v1ref
@compare sum(erf(x)) v2ref

###### erfc
@compare erfc(x)      v0ref
@compare sum(erfc(x)) v1ref
@compare sum(erfc(x)) v2ref

###### gamma
@compare gamma(x)      v0ref + 0.1 # to stay away from negative integers
@compare sum(gamma(x)) v1ref + 0.1
@compare sum(gamma(x)) v2ref + 0.1

###### lgamma
@compare lgamma(x)      v0ref + 0.1
@compare sum(lgamma(x)) v1ref + 0.1
@compare sum(lgamma(x)) v2ref + 0.1

###### beta
@compare beta(x, 2.)             v0ref + 0.1
@compare sum(beta(x, x .* 0.2))  v1ref + 0.1
@compare sum(beta(x, x))         v2ref + 0.1

###### lbeta
@compare lbeta(x, 2.)             v0ref + 0.1
@compare sum(lbeta(x, x .* 0.2))  v1ref + 0.1
@compare sum(lbeta(x, x))         v2ref + 0.1


# To run:
# include("DI_design.jl")

# DESCRIPTION
# Trims and linearizes the aircraft model at incremental flight speeds, and
# generates the gains of the dynamic inversion controller. Always run after
# modifying the aircraft parameters in the F16_constants.jl file.
#
# ------------------------------------------------------------------------------

# include packages
#using DifferentialEquations
using LinearAlgebra
using SparseArrays
using ControlSystems
using MAT
using Interpolations
using ApproxFun
using Printf
#ENV["MPLBACKEND"]="tkagg"
using Plots
using LaTeXStrings
using JLD
#using PyPlot
#using GR

# include functions
include("atan2.jl")
include("linearize.jl")
include("trimmer.jl")
#include("finp_F16_loop.jl")
include("THRUST.jl")
include("DAMP.jl")
include("DNDR.jl")
include("DNDA.jl")
include("DLDR.jl")
include("DLDA.jl")
include("CN.jl")
include("CL.jl")
include("CM.jl")
include("CZ.jl")
include("CY.jl")
include("CX.jl")
include("RTAU.jl")
include("TGEAR.jl")
include("PDOT.jl")
include("ADC.jl")
include("fix.jl")
include("F16.jl")
# load F-16 constants
include("F16_constants.jl")

# vector of trim speeds [kts]
speed_vec=range(200,stop=800,length=7)
# initialize DI matrices
CATAB = zeros(3,3,length(speed_vec))
CBinvTAB = zeros(3,3,length(speed_vec))
x0_mat = zeros(NSTATES,length(speed_vec))
u0_mat = zeros(NCTRLS,length(speed_vec))
# trim at different speeds
for iv=1:length(speed_vec)
    # set velocity [ft/s]
    speed=speed_vec[iv]*1.688
    # vertical flight path angle [deg]
    gamma=0.0
    # horizontal flight path angle [deg]
    chi=0.0
    # set initial altitude [ft]
    alt=2000.0
    # set initial North-South and East-West position [ft]
    p_north=0.0
    p_east=0.0
    # turn rate [deg/s]
    turnrate=0.0
    # pull-up rate [deg/s]
    pullup=0.0
    # roll rate [deg/s]
    rollrate=0.0
    # initial guess for state and control vectors
    phi_est=turnrate*D2R*speed/32.17
    pndot=speed*cos(gamma*D2R)*cos(chi*D2R)
    pedot=speed*cos(gamma*D2R)*sin(chi*D2R)
    hdot=speed*sin(gamma*D2R)
    x0=[speed; 0.0; 0.0; phi_est; 0; chi*D2R; 0.0; 0.0; 0.0; p_north; p_east; alt;
        0.0]
    u0=[0.2; 0; 0; 0]
    targ_des=[0; 0; 0; rollrate*D2R; pullup*D2R; turnrate*D2R; 0; 0; 0; pndot;
              pedot; hdot; 0; 0]
    # trim aircraft
    x0, u0, iter = trimmer(F16,x0,u0,targ_des)
    # store trim state and controls
    x0_mat[:,iv] = x0
    u0_mat[:,iv] = u0
    # linearized aircraft dynamics
    A, Bmat, C, D = linearize(F16,x0,u0)
    # dynamic inverse matrices
    C_DI=[1 0 0
          0 1 0
          0 0 1]
    CATAB[:,:,iv]=C_DI*A[7:9,7:9]
    CBinvTAB[:,:,iv]=inv(C_DI*Bmat[7:9,[2,3,4]])
end

# command filters time constants
tau_p=1/10
tau_q=1/10
tau_r=1/4
# washout filter for controls
tau_u=5.0;
# disturbance rejection frequency and damping and integrator poles
wnroll_d=10.0
dmproll_d=1.0
wnpitch_d=10.0
dmppitch_d=1.0
wnyaw_d=4.;
dmpyaw_d=1.;
# feedback gains
kp_roll=2*wnroll_d*dmproll_d
ki_roll=wnroll_d^2
kp_pitch=2*wnpitch_d*dmppitch_d
ki_pitch=wnpitch_d^2
kp_yaw=2*dmpyaw_d*wnyaw_d
ki_yaw=wnyaw_d^2;
# pilot gains
k_a=(290*pi/180)/21.5
k_e=(9-1)*G/25 ./(500*1.688)
k_r=1/30
# save DI matrices and gains into a file
save("DI.jld","CATAB",CATAB,"CBinvTAB",CBinvTAB,"kp_roll",kp_roll,"ki_roll",
    ki_roll,"kp_pitch",kp_pitch,"ki_pitch",ki_pitch,"kp_yaw",kp_yaw,"ki_yaw",
    ki_yaw,"k_a",k_a,"k_e",k_e,"k_r",k_r,"tau_p",tau_p,
    "tau_q",tau_q,"tau_r",tau_r,"tau_u",tau_u)

# plot trim pitch attitude and controls
gr(size=(800,600))
p10=plot(speed_vec,x0_mat[5,:]*180/pi,label="",line=(3, :royalblue, :solid), marker = (:circ, :royalblue, 8), markerstrokecolor = :royalblue)
xaxis!((200, 800), 200:50:800)
yaxis!("θ [deg]",yguidefontsize=14)
p11=plot(speed_vec,u0_mat[1,:]*100,label="",line=(3, :royalblue, :solid), marker = (:circ, :royalblue, 8), markerstrokecolor = :royalblue)
xaxis!((200, 800), 200:50:800)
yaxis!(L"\delta_t [\%]",yguidefontsize=14)
p12=plot(speed_vec,u0_mat[3,:],label="",line=(3, :royalblue, :solid), marker = (:circ, :royalblue, 8), markerstrokecolor = :royalblue)
xaxis!((200, 800), 200:50:800)
yaxis!(L"\delta_e [\mathrm{deg}]",yguidefontsize=14)
p13=plot(p10, p12, p11, layout=(3,1), leftmargin=3Plots.mm)
display(p13)
savefig(".\\Plots\\trim_F16.svg")


function solve(case, D, X)
    earliest = D
    for i in Iterators.reverse(X)
        earliest = div(earliest, i) * i
    end

    println("Case #$(case): $(earliest)")
end

function main()
    T = parse(Int, readline(stdin))
    for case in 1:T
        _, D = parse.(Int, split(readline(stdin)))
        X = parse.(Int, split(readline(stdin)))
        solve(case, D, X)
    end
end

isinteractive() || main()


"""
Make a lattice with cubic structure.

# Arguments

 - **box** (Array) 3D Box to be filled with spheres of radius r. e.g. [[0,1],[1,2],[0,4]]
 - **r** (Number) Radius of the sphere

# Optional keyword arguments
 - **noiseRatio** (Number) Noise ratio to add to the center of the cells. Default 0.
 - **holesRatio** (Number) Probability of holes in lattice. Default 0.

# Returns
[posX,posY,posZ] Where pos are arrays with the positions of the cells for each component.
"""
function latticeCubic(box::Array{Array{Float64,1},1},r::Number;noiseRatio::Number=0.,holesRatio::Number=0)

    d = length(box)
    if !(d in [1,2,3])
        error("Incompatible number of dimensions ", d)
    end

    #Make first dimension
    lineX = Array(box[1][1]:2*r:box[1][2])

    #Make second dimension
    if d > 1
        lineY = fill(box[2][1]+r,length(lineX))
        nY = ceil(Int,(box[2][2]-box[2][1])/(2*r))
        areaX = Float64[]
        areaY = Float64[]
        dx = 2*r
        dy = 2*r
        for i in 0:nY-1
            append!(areaX,lineX)
            append!(areaY,lineY.+dy*i)
        end
    end

    #Make third dimension
    if d == 3
        areaZ = fill(box[2][1]+r,length(areaX))
        nZ = ceil(Int,(box[3][2]-box[3][1])/(2*r))
        volumeX = Float64[]
        volumeY = Float64[]
        volumeZ = Float64[]
        dy = 2*r
        dz = 2*r
        for i in 0:nZ-1
            append!(volumeX,areaX)
            append!(volumeY,areaY)
            append!(volumeZ,areaZ.+dz*i)
        end
    end

    #Add noise
    if noiseRatio > 0
        dist = Normal(0,r*noiseRatio)
        if d == 1
            n = length(lineX)
            lineX += rand(dist,n);
        elseif d == 2
            n = length(areaX)
            areaX += rand(dist,n);
            areaY += rand(dist,n);
        else
            n = length(volumeX)
            volumeX += rand(dist,n)
            volumeY += rand(dist,n);
            volumeZ += rand(dist,n)
        end
    end

    #Add holes
    if holesRatio > 0
        if d == 1
            l = rand(length(lineX))
            l = l .> holesRatio

            lineX = lineX[l]
        elseif d == 2
            l = rand(length(areaX))
            l = l .> holesRatio

            areaX = areaX[l]
            areaY = areaY[l]
        else
            l = rand(length(volumeX))
            l = l .> holesRatio

            volumeX = volumeX[l]
            volumeY = volumeY[l]
            volumeZ = volumeZ[l]
        end
    end

    #Return
    if d == 1
        return lineX
    elseif d == 2
        return areaX, areaY
    else
        return volumeX, volumeY, volumeZ
    end

end


using Test
using Unfold

firbase = firbasis((-1, 1), 10)

# test optional call
@test Unfold.kernel(firbase) == Unfold.kernel(firbasis((-1, 1), 10))

@test typeof(Unfold.name(firbase)) <: String
# test basics of basisfunction
@test length(collect(Unfold.colnames(firbase))) == 21
@test unique(Unfold.kernel(firbase)(1)) == [1.0, 0.0]

# test length consistency
@test length(Unfold.colnames(firbase)) ==
      length(Unfold.times(firbase)) ==
      size(Unfold.kernel(firbase)(3.1))[2]

# testing the non-sampling rate samples
@test Unfold.kernel(firbase)(0.5)[1:3, 1:3] == [0.5 0.0 0.0; 0.5 0.5 0.0; 0.0 0.5 0.5]

boldbase = hrfbasis(2.0, name = "test")

@test Unfold.name(boldbase) == "test"
@test Unfold.kernel(boldbase)(0) == Unfold.kernel(boldbase)(1)
@test Unfold.kernel(boldbase)(0.1) != Unfold.kernel(boldbase)(1) # testing fractional kernel generation
@test findmax(Unfold.kernel(boldbase)(0.3))[2][1] == 4


@testset "Spectra" begin

@testset "NoLight" begin

    N = NoLight()

    @test N == GRAYTR.nolight
    @test 2.0 .* N == N
    @test N .* 2.0 == N
    @test N .* N == N
    @test N .+ N == N

end

@testset "SampledSpectrum" begin

    S1 = SampledSpectrum(300, 800, [0.0, 1.0, 0.0])
    S2 = SampledSpectrum(300, 800, [0.5, 0.5, 0.5])

    @test S1.values ≈ [0.0, 1.0, 0.0]
    @test S2.values ≈ [0.5, 0.5, 0.5]

    # SampledSpectrum with itself

    S3 = S1 .+ S2
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.5, 1.5, 0.5]

    S3 = S1 .* S2
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.0, 0.5, 0.0]

    S3 = @. S2 * (S2 + S1)
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.25, 0.75, 0.25]

    # SampledSpectrum with numbers

    S3 = S1 .* 2.0
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.0, 2.0, 0.0]

    S3 = S1 ./ 2.0
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.0, 0.5, 0.0]

    S3 = S1 .+ 2.0
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [2.0, 3.0, 2.0]

    S3 = 2.0 .* S1
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.0, 2.0, 0.0]

    S3 = 2.0 .+ S1
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [2.0, 3.0, 2.0]

    # SampledSpectrum with NoLight

    N = NoLight()

    S3 = N .* (S1 .+ S2)
    @test typeof(S3) == NoLight
    @test S3 == nolight

    S3 = N .+ (S1 .+ S2)
    @test typeof(S3) == SampledSpectrum
    @test S3.values ≈ [0.5, 1.5, 0.5]

end

end # main testset


function makesingleplot(vec1::AbstractVector, vec2::AbstractVector,
                                        E::Int, tau_s::Int, tau_p::Int,
                                        var1name::ASCIIString, var2name::ASCIIString;
                                        lib_start=1, xmin=false, xmax=false, ymin=false,
                                        ymax=false, nboots=0, plot=true,
                                        libsizemin=0, libsizemax=0,
                                        npred=0, pred_start=0)
    """
    vec1: Time series 1
    vec2: Time series 2
    libsizemin: Minimum library size
    ibsizemax: Maximum library size
    E: System dimensionality
    tau_s: Lag for manifold reconstruction
    tau_p: Lag for causal effect
    npred: Number of points to predict
    pred_start: Start for prediction
    var1name: Name of first variable
    var2name: Name of second variable
    lib_start: Start of library
    """
    @pyimport matplotlib.pyplot as plt
    shadowmat_dict_vec1, distmat_dict_vec1 = precalc_manif_dists(E:E, tau_s:tau_s, vec1)
    shadowmat_dict_vec2, distmat_dict_vec2 = precalc_manif_dists(E:E, tau_s:tau_s, vec2)

    librange12, yval_12 = calcCCM(vec1, vec2,
                                                        shadowmat_dict_vec1, distmat_dict_vec1,
                                                        E, tau_s, tau_p; libsizemin=libsizemin, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start,
                                                        lib_start=lib_start, nboots=nboots
                                                    )

    librange21, yval_21 = calcCCM(vec2, vec1,
                                                        shadowmat_dict_vec2, distmat_dict_vec2,
                                                        E, tau_s, tau_p; libsizemin=libsizemin, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start,
                                                        lib_start=lib_start, nboots=nboots
                                                    )

    plt.plot(librange12, yval_12, label = "$var2name influences $var1name?")
    plt.plot(librange21, yval_21, label = "$var1name influences $var2name?")
    plt.xlabel("L")
    plt.ylabel("\$\\rho_{ccm}\$")
    plt.legend(loc="lower right")
    ax = plt.gca()
    ax[:xaxis][:set_ticks_position]("bottom")
    ax[:yaxis][:set_ticks_position]("left")
    if xmin != false
        ax[:set_xlim]((xmin, xmax))
    end
    if ymin != false
        ax[:set_ylim]((ymin, ymax))
    end
    if plot
        plt.show()
    else
        plt.close()
    end
end

function makeoptimizationplots(vec1::AbstractVector, vec2::AbstractVector, E_vals::AbstractVector,
                                                    tau_s_vals::AbstractVector, tau_p_vals::AbstractVector,
                                                    var1name::ASCIIString, var2name::ASCIIString;
                                                    nreps=5, b_offset=1, ncols=28, left_E=false, left_tau_p=false,
                                                    right_E=false, right_tau_p=false, lagunit=1, unit=false,
                                                    imfont="medium", nboots=0, legend=true,
                                                    lib_start=1, plot=true, libsizemin=0, libsizemax=0,
                                                    pred_start=0, npred=0, show_tau_s=true)

    println("Calculating manifolds")
    shadowmat_dict_vec1, distmat_dict_vec1 = precalc_manif_dists(E_vals, tau_s_vals, vec1)
    shadowmat_dict_vec2, distmat_dict_vec2 = precalc_manif_dists(E_vals, tau_s_vals, vec2)

    println("\nCoordDescent opt1")
    res12 = CoordDescentOpt(vec1, vec2,
                                                shadowmat_dict_vec1, distmat_dict_vec1,
                                                E_vals, tau_s_vals, tau_p_vals;
                                                libsizemin=libsizemin, libsizemax=libsizemax,
                                                npred=npred, pred_start=pred_start,
                                                nreps=nreps, nboots=nboots)
    println("\nCoordDescent opt2")
    res21 =CoordDescentOpt(vec2, vec1,
                                                shadowmat_dict_vec2, distmat_dict_vec2,
                                                E_vals, tau_s_vals, tau_p_vals;
                                                libsizemin=libsizemin, libsizemax=libsizemax,
                                                npred=npred, pred_start=pred_start,
                                                nreps=nreps, nboots=nboots)

    libsizemin_12 = max(res12["E"] + b_offset + 1, 10)
    libsizemin_21 = max(res21["E"] + b_offset + 1, 10)

    println("\nstarting calcCCM1")
    librange12, yval_12 = calcCCM(vec1, vec2,
                                                        shadowmat_dict_vec1, distmat_dict_vec1,
                                                        res12["E"], res12["tau_s"], res12["tau_p"];
                                                        libsizemin=libsizemin_12, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start,
                                                        lib_start=1, nboots=nboots)
    println("starting calcCCM2")
    librange21, yval_21 = calcCCM(vec2, vec1,
                                                        shadowmat_dict_vec2, distmat_dict_vec2,
                                                        res12["E"], res12["tau_s"], res12["tau_p"];
                                                        libsizemin=libsizemin_21, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start,
                                                        lib_start=1, nboots=nboots)

    @pyimport matplotlib.pyplot as plt
    fig = plt.figure()
    label1 = "\n$(ucfirst(var2name)) influences $var1name?"
    label2 = "\n$(ucfirst(var1name)) influences $var2name?"

    rho1   = @sprintf("%.2f", res12["rho"])
    rho2   = @sprintf("%.2f", res21["rho"])

    if show_tau_s
        statsuff1 = "$(res12["tau_s"])"
        statsuff2 = "$(res21["tau_s"])"
    else
        statsuff1 = ""
        statsuff2 = ""
    end
    stats1 = "   E = $(res12["E"]), \$\\tau_{p}\$ = $(int(res12["tau_p"] * lagunit))" * statsuff1
    stats2 = "   E = $(res21["E"]), \$\\tau_{p}\$ = $(int(res21["tau_p"] * lagunit))" * statsuff2

    ax1 = plt.subplot2grid((2,ncols), (0,0), rowspan=2, colspan=ifloor(ncols/3))

    ax1[:plot](librange12, yval_12, label="$label1\n$stats1")
    ax1[:plot](librange21, yval_21, label="$label2\n$stats2")

    ax1[:yaxis][:set_ticks_position]("left")
    ax1[:xaxis][:set_ticks_position]("bottom")

    ax1[:set_ylabel]("\$\\rho_{ccm}\$", fontsize="x-large")
    ax1[:set_xlabel]("L", fontsize="large", labelpad=12)

    if legend == true
        ax1[:legend](loc=4, fontsize="x-small")
    end

    mat12, Es12, taus12 = get_E_taupcurves(vec1, vec2,
                                                                        shadowmat_dict_vec1, distmat_dict_vec1,
                                                                        libsizemax, res12["E"], res12["tau_s"], res12["tau_p"],
                                                                        E_vals, tau_p_vals, npred, pred_start;
                                                                        left_E=left_E, left_tau_p=left_tau_p,
                                                                        right_E=right_E, right_tau_p=right_tau_p, nboots=nboots
                                                                    )

    mat21, Es21, taus21 = get_E_taupcurves(vec2, vec1,
                                                                        shadowmat_dict_vec2, distmat_dict_vec2,
                                                                        libsizemax, res21["E"], res21["tau_s"], res21["tau_p"],
                                                                        E_vals, tau_p_vals, npred, pred_start;
                                                                        left_E=left_E, left_tau_p=left_tau_p,
                                                                        right_E=right_E, right_tau_p=right_tau_p, nboots=nboots
                                                                    )

    catmat = vcat(mat12, mat21)
    if minimum(catmat) < 0
        vmin = 0
        llab = "\$\\leq 0\$"
    else
        vmin = minimum(catmat)
        llab = @sprintf("%.2f", vmin)
    end

    vmax = 1  # maximum correlation

    if unit != false
        xlabel = "\$\\tau_p\$\ ($unit)"
    else
        xlabel =  "\$\\tau_p\$"
    end

    ax2 = plt.subplot2grid((2,ncols), (0,ifloor(ncols/3)), colspan=2*ifloor(ncols/3))
    imax12 = plotheatmap(ax2, mat12, vmin, vmax, label1, taus12, Es12, imfont, lagunit, xlabel)

    ax3 = plt.subplot2grid((2,ncols), (1,ifloor(ncols/3)), colspan=2*ifloor(ncols/3))
    imax21 = plotheatmap(ax3, mat21, vmin, vmax, label2, taus21, Es21, imfont, lagunit, xlabel)

    ax4 = plt.subplot2grid((2,ncols), (0,ncols-1), rowspan=2)  # axis for colorbar

    cbar = fig[:colorbar](imax21, cax=ax4, ticks=[vmin, vmax], format = "%.2f")
    cbar[:ax][:set_yticklabels]([llab, "\$1\$"], fontsize="large")
    cbar[:set_label]("\$max\\ \\rho_{ccm}\$", fontsize="large", labelpad=2)
    cbar[:ax][:get_yaxis]()[:labelpad]=0

    if plot
        plt.show()
    else
        plt.close()
    end
end

function get_totest(val::Int, numtotest::Int, minval::Int, maxval::Int)
    if maxval - minval - 1 < numtotest
        println("Maxval and minval for E_taup curves are inconsistent with number to test")
        println("Fixing maxval")
        maxval = minval + numtotest - 1
    end

    left::Int  = val-iceil(numtotest/2)
    right::Int = val+ifloor(numtotest/2)-1

    if left < minval
        right += (minval-left)
        left   = minval
    end

    if right >  maxval
        left -= right%maxval
        right = maxval
    end

    if right-left+1 != numtotest
        error("start and stop not equal to npred")
    elseif left < minval
        println("left: $left, right: $right, Max: $maxval, min: $minval.")
        error("Left value less than minval")
    end

    return left, right
end

function get_E_taupcurves(source_series::AbstractVector, target_series::AbstractVector, shadowmat_dict::Dict,
                                            distmat_dict::Dict, libsizemax::Int, E::Int, tau_s::Int, tau_p::Int,
                                            E_vals::AbstractVector, tau_p_vals::AbstractVector, npred::Int, pred_start::Int;
                                            nlag=10, b_offset=1, num_E=6, Estart=2, num_tau_p=14, tau_p_start=0,
                                            left_E=false, left_tau_p=false, right_E=false, right_tau_p=false, nboots=0)

    num_E = min(maximum(E_vals) - Estart + 1, num_E)
    num_tau_p = min(maximum(tau_p_vals-tau_p_start+1), num_tau_p)

    if (typeof(left_E) == Bool) && (left_E == false)
        left_E, right_E = get_totest(E, num_E, Estart, maximum(E_vals))
    end

    if (typeof(left_tau_p) == Bool) && (left_tau_p == false)
        left_tau_p, right_tau_p = get_totest(tau_p, num_tau_p, tau_p_start, maximum(tau_p_vals))
    end

    res = fill(NaN, num_E, num_tau_p)

    for (count_E, Eval) in enumerate(left_E:right_E)
        println("\nCalculating CCM max for E: $Eval")
        for (count_tau_p, tau_p_val) in enumerate(left_tau_p:right_tau_p)
            # libsizemin will be ignored. libsizemax will be reduced if it is too large
            println("\ttau_p: $tau_p_val")

            librange, rhos = calcCCM(source_series, target_series,
                                                        shadowmat_dict, distmat_dict,
                                                        Eval, tau_s, tau_p_val;
                                                        libsizemin=libsizemax - nlag, libsizemax=libsizemax,
                                                        npred=npred, pred_start=pred_start,
                                                        quick=true, nboots=nboots)
            rhos = rhos[!isnan(rhos)]

            if length(rhos) > 0
                rho = median(rhos)
            else
                rho = NaN
            end

            res[count_E, count_tau_p] = rho
        end
    end
    return res, left_E:right_E, left_tau_p:right_tau_p
end

function plotheatmap(ax, mat, vmin, vmax, label, taus, Es, imfont, lagunit, xlabel)
    imax = ax[:imshow](mat, vmin=vmin, vmax=vmax, origin="lower")
    ax[:set_title](label, fontsize="medium")

    ax[:xaxis][:set_ticks_position]("bottom")
    ax[:yaxis][:set_ticks_position]("right")

    xticks = 0:(length(taus)-1)
    ax[:set_xticks](xticks[1:2:end]) # take every other tick
    ax[:set_yticks](0:(length(Es)-1))

    ax[:set_xticklabels](
                                    format_ticks(
                                                        (taus * lagunit)[1:2:end]
                                                        ), fontsize=imfont)
    ax[:set_yticklabels](Es, fontsize=imfont)

    ax[:set_xlim]((0, length(taus)-1))
    ax[:set_ylim]((0, length(Es)-1))
    ax[:set_xlabel](xlabel, fontsize=imfont)
    ax[:set_ylabel]("E", fontsize=imfont, labelpad = 12)
    return imax
end

function format_ticks(xticks)

    try
        xticks = convert(Vector{Float64}, xticks) # float range to float vector (otherwise subtraction fails)
    catch
        return xticks # This may fail if you have an int range
    end

    xticks_int = fill(0, length(xticks))
    try
        xticks_int = convert(Vector{Int}, xticks)  # If values are too far from ints, this raises an InexactError
    catch
        return xticks
    end

    variance = var(xticks)
    if all(abs(xticks - xticks_int) .< (variance * 1e-12))  # possibly more stringent than InexactError if values are small
        return xticks_int
    else
        return xticks
    end
end


@doc raw"""
Tableau for variational Lobatto IIIA method with s stages

Uses Lobatto IIIA for the coefficients $a_{ij}$ and its symplectic conjugate, Lobatto IIIB, for the coefficients $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIA(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIA$(s)"), 2s-2, TableauLobattoIIIA(T,s), (-1)^(s+1), get_lobatto_nullvector(s))
end

TableauVPLobattoIIIA(s) = TableauVPLobattoIIIA(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIIB method with s stages

Uses Lobatto IIIB for the coefficients $a_{ij}$ and its symplectic conjugate, Lobatto IIIA, for the coefficients $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIB(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIB$(s)"), 2s-2, TableauLobattoIIIB(T,s), (-1)^(s+1), get_lobatto_nullvector(s))
end

TableauVPLobattoIIIB(s) = TableauVPLobattoIIIB(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIIC method with s stages

Uses Lobatto IIIC for the coefficients $a_{ij}$ and its symplectic conjugate, Lobatto IIIC̄, for the coefficients $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIC(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIC$(s)"), 2s-2, TableauLobattoIIIC(T,s), (-1)^(s+1))
end

TableauVPLobattoIIIC(s) = TableauVPLobattoIIIC(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIIC̄ method with s stages

Uses Lobatto IIIC̄ for the coefficients $a_{ij}$ and its symplectic conjugate, Lobatto IIIC, for the coefficients $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIC̄(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIC̄$(s)"), 2s-2, TableauLobattoIIIC̄(T,s), (-1)^(s+1))
end

TableauVPLobattoIIIC̄(s) = TableauVPLobattoIIIC̄(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIID method with s stages

Uses Lobatto IIID for both coefficients $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIID(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIID$(s)"), 2s-s, TableauLobattoIIID(T,s), (-1)^s)
end

TableauVPLobattoIIID(s) = TableauVPLobattoIIID(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIIE method with s stages

Uses Lobatto IIIE for both coefficients $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIE(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIE$(s)"), 2, TableauLobattoIIIE(T,s), (-1)^s)
end

TableauVPLobattoIIIE(s) = TableauVPLobattoIIIE(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIIF method with s stages

Uses Lobatto IIIF for the coefficients $a_{ij}$ and its symplectic conjugate for the coefficients $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIF(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIF$(s)"), 2, TableauLobattoIIIF(T,s), (-1)^s)
end

TableauVPLobattoIIIF(s) = TableauVPLobattoIIIF(Float64,s)

@doc raw"""
Tableau for variational Lobatto IIIG method with s stages

Uses Lobatto IIIG for both coefficients $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIG(::Type{T}, s) where {T}
    TableauVPRK(Symbol("LobattoIIIG$(s)"), 2, TableauLobattoIIIG(T,s), (-1)^s)
end

TableauVPLobattoIIIG(s) = TableauVPLobattoIIIG(Float64,s)

@doc raw"""
Tableau for variational Gauss-Legendre method with s stages

Uses Gauss coefficients for both $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPGLRK(::Type{T}, s) where {T}
    TableauVPRK(Symbol("GLRK", s), 2s, TableauGauss(T, s), (-1)^s)
end

TableauVPGLRK(s) = TableauVPGLRK(Float64,s)

@doc raw"""
Tableau for variational symmetric Runge-Kutta method with 3 stages

Uses SRK3 coefficients for both $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPSRK3(T=Float64)
    srk = TableauSRK3(T)
    TableauVPRK(:SRK3, srk.o, srk, (-1)^3)
end

@doc raw"""
Tableau for an implicit partitioned Lobatto IIIA Runge-Kutta method with s stages

Uses Lobatto IIIA for both coefficients $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIAIIIA(::Type{T}, s) where {T}
    lob = TableauLobattoIIIA(T,s)
    TableauVPRK(Symbol("LobattoIIIAIIIA$s"), 2s-2, lob, lob, (-1)^(s+1))
end

TableauVPLobattoIIIAIIIA(s) = TableauVPLobattoIIIAIIIA(Float64,s)

@doc raw"""
Tableau for an implicit partitioned Lobatto IIIB Runge-Kutta method with s stages

Uses Lobatto IIIB for both coefficients $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPLobattoIIIBIIIB(::Type{T}, s) where {T}
    lob = TableauLobattoIIIB(T,s)
    TableauVPRK(Symbol("LobattoIIIBIIIB$s"), 2s-2, lob, lob, (-1)^(s+1))
end

TableauVPLobattoIIIBIIIB(s) = TableauVPLobattoIIIBIIIB(Float64,s)

@doc raw"""
Tableau for an implicit partitioned Radau IIA Runge-Kutta method with s stages

Uses Radau IIA for both coefficients $a_{ij}$ and $\bar{a}_{ij}$.
"""
function TableauVPRadauIIAIIA(::Type{T}, s) where {T}
    rad = TableauRadauIIA(T,s)
    TableauVPRK(Symbol("RadauIIAIIA$s"), 2s-1, rad, rad, (-1)^(s+1))
end

TableauVPRadauIIAIIA(s) = TableauVPRadauIIAIIA(Float64,s)


# This file defines staionary arm models

"""
    StationaryArmBase

`StationaryArmBase` is the concrete type for all the stationary arm models. Any
subtype should implement ```pull!()``` method
"""
abstract type StationaryArmBase <: ArmBase end

# Default method for interface
pull!( arm::StationaryArmBase )     = nothing
tick!( arm::StationaryArmBase )     = nothing
reset!( arm::StationaryArmBase )    = nothing
mean( arm::StationaryArmBase )      = Distributions.mean( arm.armDist )

"""
    Bernoulli( p )
`Bernoulli` creates a Bernoulli arm with success probability `p`.
"""
mutable struct Bernoulli <: StationaryArmBase
    armDist::Distributions.Bernoulli

    function Bernoulli( p::Real )
        new( Distributions.Bernoulli(p) )
    end
end

function pull!( arm::Bernoulli )
    return Distributions.rand(arm.armDist)
end

"""
    Beta( α, β )

`Beta` creates a Beta arm with parameters `α` and `β`.
"""
mutable struct Beta <: StationaryArmBase
    armDist::Distributions.Beta

    function Beta( α::Real, β::Real )
        new( Distributions.Beta( α, β ) )
    end
end

function pull!( arm::Beta )
    return Distributions.rand( arm.armDist )
end

"""
    Exponetial( θ )

`Expoenential` creates a Expoenentially distributed arm with mean reward `θ`
"""
mutable struct Exponential <: StationaryArmBase
    armDist::Distributions.Exponential

    function Exponential( λ::Real ) # λ is the mean parameter
        new( Distributions.Exponential(λ) )
    end
end

function pull!( arm::Exponential )
    return Distributions.rand( arm.armDist )
end

"""
    Normal( μ, σ )

Creates an arm with reward distributed as Normal random varaible with mean `μ`
and standard deviation `σ`.
"""
mutable struct Normal <: StationaryArmBase
    armDist::Distributions.Normal

    function Normal( μ::Real, σ::Real )
        new(
            Distributions.Normal( μ, σ )
        )
    end
end

function pull!( arm::Normal )
    return rand( arm.armDist )
end


# Efficient implementation a' * Q * b that avoid unnecessary type promotion as
# well as unnecessary allocation
function quad_form(a::AbstractVector{<:Real},
                   Q::SymMatrix{<:Real},
                   b::AbstractVector{<:Real})
    n = length(a)
    @assert n == LinearAlgebra.checksquare(Q)
    @assert n == length(b)
    out = zero(typeof(zero(eltype(a)) * zero(eltype(Q)) * zero(eltype(b))))
    k = 0
    for j in 1:n
        α = a[j]
        β = b[j]
        for i in 1:(j-1)
            k += 1
            out += a[i] * Q.Q[k] * β
            out += α * Q.Q[k] * b[i]
        end
        k += 1
        out += a[j] * Q.Q[k] * b[j]
    end
    return out
end
function quad_form(a::AbstractVector,
                   Q::AbstractMatrix,
                   b::AbstractVector)
    (n, m) = size(Q)
    n == length(a) || throw(DimensionMismatch())
    m == length(b) || throw(DimensionMismatch())
    U = MA.promote_operation(*, eltype(a), eltype(Q), eltype(b))
    out = zero(MA.promote_operation(+, U, U))
    for j in 1:m
        for i in 1:n
            out = MA.add_mul!!(out, a[i], Q[i, j], b[j])
        end
    end
    return out
end
# Same as quad_form(a, Q, a)
function quad_form(Q::Symmetric,
                   a::AbstractVector)
    return quad_form(a, Q, a)
end
function quad_form(Q::Symmetric{JuMP.VariableRef}, a::AbstractVector{<:Real})
    return quad_form(a, Q, a)
end

# We have x' Q x and we want y' Q y where y is obtained by substituting
# vars for A * new_vars in x. We want to compute the matrix M such that
# y = M z where z is the vector of monomials of degree d of new_vars
# Then we will have y' Q y = z' M' Q M z
struct GramTransformation{T, MT <: MP.AbstractMonomial,
                          MVT <: AbstractVector{MT}}
    M::Matrix{T}
    monos::MVT
end

function apply_transformation(p::SumOfSquares.GramMatrix,
                              t::GramTransformation) where T
    new_n = length(t.monos)
    new_Q = [quad_form(t.M[:, i], p.Q, t.M[:, j]) for j in 1:new_n for i in 1:j]
    return GramMatrix(SymMatrix(new_Q, new_n), t.monos)
end

function transformation(old_monos, A::AbstractMatrix, new_vars, d)
    new_monos = monomials(new_vars, d)
    new_n = length(new_monos)
    M = zeros(eltype(A), length(old_monos), new_n)
    mapped_vars = A * new_vars
    # Cache the result of mapped_vars[i]^n
    powers = [Union{Nothing, eltype(mapped_vars)}[mapped_vars[i]]
              for i in eachindex(mapped_vars)]
    # Compute mapped_vars[i]^n by "Power by Squaring" and cache it in `powers`
    function _power(i, n)
        @assert n > 0
        while n > length(powers[i])
            push!(powers[i], nothing)
        end
        if powers[i][n] === nothing
            if isodd(n)
                powers[i][n] = mapped_vars[i] * _power(i, n - 1)
            else
                p_2 = _power(i, div(n, 2))
                powers[i][n] = p_2 * p_2
            end
        end
        return powers[i][n]::eltype(mapped_vars)
    end
    function _map(mono::MP.AbstractMonomial)
        exps = MP.exponents(mono)
        length(exps) == length(mapped_vars) || throw(ArgumentError("A monomial have less variables than `new_vars`"))
        cur = one(eltype(mapped_vars))
        for i in eachindex(exps)
            if exps[i] > 0
                cur *= _power(i, exps[i])
            end
        end
        return cur
    end
    for i in eachindex(old_monos)
        y = _map(old_monos[i])
        j = 1
        for term in MP.terms(y)
            mono = MP.monomial(term)
            while j <= length(new_monos) && mono < new_monos[j]
                j += 1
            end
            M[i, j] = MP.coefficient(term)
        end
    end
    return GramTransformation(M, new_monos)
end

# computes p ∘ A or more precisely p(variables(p) => A * new_vars)
function apply_matrix(p::SumOfSquares.GramMatrix{T, <:MonomialBasis},
                      A::AbstractMatrix, new_vars, d) where T
    return apply_transformation(p, transformation(p.basis.monomials, A, new_vars, d))
end

# computes A # μ or more precisely p(variables(p) => A * new_vars)

function psd_constraint(Q::Symmetric)
    n = LinearAlgebra.checksquare(Q)
    q = [Q[i, j] for j in 1:n for i in 1:j]
    # For n == 0, it will create no constraint, for n == 1, it will simply
    # be a Nonnegatives constraint and for n == 2 it will be a rotated SOC.
    set = SumOfSquares.matrix_cone(MOI.PositiveSemidefiniteConeTriangle, n)
    return PolyJuMP.bridgeable(JuMP.build_constraint(error, q, set),
                               JuMP.moi_function_type(typeof(q)), typeof(set))
end

function psd_constraint(model, Q::Symmetric)
    return JuMP.add_constraint(model, psd_constraint(Q))
end


import ExperimentalAnalysis
using Base.Test
using DataFrames
using GLM
# write your own tests here

df = DataFrame(a = randn(10),b = randn(10),c = randn(10),d = randn(10))
df[2,2] = NA

@test ExperimentalAnalysis.perform_example_analysis() == nothing
ExperimentalAnalysis.scattermatrix(randn(10,3))
ExperimentalAnalysis.scattermatrix(df)
ExperimentalAnalysis.scattermatrix(df, reglines=true)
ExperimentalAnalysis.scattermatrix(df, f ~ a + b)
ExperimentalAnalysis.scattermatrix(df, c ~ a + b)
ExperimentalAnalysis.scattermatrix(randn(10,3),["x1", "x2", "x3"])

modela = lm(b ~ a + c, df)
modelb = lm(d ~ a + c, df)
ExperimentalAnalysis.scattermatrix([modela, modelb])
ExperimentalAnalysis.modelheatmap([modela, modelb])
ExperimentalAnalysis.modelheatmap(["a", "b"], modela, modelb)
ExperimentalAnalysis.modelheatmap(["a", "b"], [modela, modelb])


using LinearAlgebra, ForwardDiff, Distributions
using Convex, SCS, ECOS
using IterativeSolvers

# problem setup
"Convex.jl"
P = [1.0 0.0 0.0;
     0.0 1.0 0.0]
D = [P; -P]
v = [1.0; 1.0; 0.0]
p = 4
c = D * v
y = -1.0e-8
β = Variable(p)
prob = minimize(c' * β)
prob.constraints += β' * ones(p) <= y
prob.constraints += β >= 0.0

@time solve!(prob, ECOS.Optimizer)

@show prob.status
@show prob.constraints
@show D' * β.value
# @show prob.constraints[1].dual
# @show prob.constraints[2].dual
A = [-1.0 * Diagonal(ones(p)); ones(1,p)]
b = [zeros(p); y]
m = size(A, 1)
k = p + m + 1

Q = Array([zeros(p, p) A' c;
           -A zeros(m, m) b;
           -c' -b'      0.0])

function κ(z)
    # z_proj = zero(z)
    # z_proj[1:N] = z[1:N]
    # z_proj[N .+ (1:p)] = max.(0.0, z[N .+ (1:p)])

    z_proj = max.(0.0, z)
    return z_proj
    # return z
end

function Pc(z)
    z_proj = zero(z)
    z_proj[1:p] = z[1:p]
    z_proj[p .+ (1:m)] = κ(z[p .+ (1:m)])
    z_proj[p + m + 1] = max(0.0, z[p + m + 1])

    return z_proj
end

function F(z)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2 * k .+ (1:k)]

    [(I + Q) * ũ - (u + v);
     u - Pc(ũ - v);
     ũ - u]
end

function Ju(z)
    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2k .+ (1:k)]

    JP = zeros(m, m)
    dif = ũ[p .+ (1:m)] - v[p .+ (1:m)]
    for i = 1:m
        if dif[i] >= 0.0
            JP[i, i] = 1.0
        else
            JP[i, i] = 0.0
        end
    end

    if z[k] - z[2k] >= 0.0
        ℓ = 1.0
    else
        ℓ = 0.0
    end

    [-I zeros(p, m) zeros(p) I zeros(p, m) zeros(p) I zeros(p, m) zeros(p);
     zeros(m, p) -JP zeros(m) zeros(m, p) I zeros(m) zeros(m, p) JP zeros(m);
     [zeros(1, p) zeros(1, m) -ℓ zeros(1, p) zeros(1, m) 1.0 zeros(1, p) zeros(1, m) ℓ]]
end

function J(z)
    [(I + Q) -I -I
     Ju(z);
     I -I zeros(k, k)]
end

ũ = zeros(k)
u = zeros(k)
v = zeros(k)
ũ[end] = 1.0
u[end] = 1.0
v[end] = 1.0

z = [ũ; u; v]
z = rand(3k)
F(z)
rank(Ju(z))
J(z)
rank(ForwardDiff.jacobian(F, z))
norm(J(z) - ForwardDiff.jacobian(F, z), Inf)
gmres(J(z), -F(z))
# (J(z)' * J(z)) \ (J(z)' * F(z))
# Δ = zero(z)
# Δ[k] = 1.0
# Δ[2k] = 1.0
# Δ[3k] = 1.0
# gmres!(Δ, J(z), -1.0 * F(z), abstol = 1.0)

function solve()
    ũ = zeros(k)
    u = zeros(k)
    v = zeros(k)
    ũ[end] = 1.0
    u[end] = 1.0
    v[end] = 1.0

    z = [ũ; u; v]
    # z = rand(3k)
    # Δ = zero(z)

    extra_iters = 0

    for i = 1:500
        _F = F(z)
        _J = J(z)
        Δ = gmres(_J, 1.0 * _F, abstol = 1.0e-12, maxiter = i + extra_iters)

        iter = 0
        α = 1.0
        while norm(F(z - α * Δ))^2.0 >= (1.0 - 0.001 * α) * norm(_F)^2.0 && α > 1.0e-4
            α = 0.5 * α
            # println("   α = $α")
            iter += 1
            if iter > 100
                @error "line search fail"
                # return z
                x = z[k .+ (1:p)]
                y = z[k + p .+ (1:m)]
                τ = z[k + k]
                κ = z[3k]
                println("τ = $τ")
                println("κ = $κ")
                return x ./ τ, z
            end
        end

        if α <= 1.0e-4
            extra_iters += 1
        end

        println("iter ($i) - norm: $(norm(F(z)))")

        z .-= α * Δ
    end

    # return z
    x = z[k .+ (1:p)]
    y = z[k + p .+ (1:m)]
    τ = z[k + k]
    κ = z[3k]
    println("τ = $τ")
    println("κ = $κ")
    return x ./ τ, z
end

x_sol, z_sol = solve()

x_sol
b - A * x_sol
norm(x_sol - β.value)

function Fθ(z, θ)
    A = reshape(θ[1:m * p], m, p)
    b = θ[m * p .+ (1:m)]
    c = θ[m * p + m .+ (1:p)]

    Q = Array([zeros(eltype(θ), p, p) A' c;
               -A zeros(eltype(θ), m, m) b;
               -c' -b'      0.0])

    ũ = z[1:k]
    u = z[k .+ (1:k)]
    v = z[2 * k .+ (1:k)]

    [(I + Q) * ũ - (u + v);
     u - Pc(ũ - v);
     ũ - u]
end
θ[m * p + m]
θ = [vec(A); b; c]
_F(y) = Fθ(z_sol, y)
_F(θ)
ForwardDiff.jacobian(_F, θ)
rank(J(z_sol))
J(z_sol)
((J(z_sol)' * J(z_sol)) \ (J(z_sol)' * ForwardDiff.jacobian(_F, θ)))[k .+ (1:p), m * p + (m-1) .+ (1:1+p)] ./ z_sol[2k]
# (J(z_sol) \ ForwardDiff.jacobian(_F, θ))[k .+ (1:p), m * p + (m-1) .+ (1:1+p)] ./ z_sol[2k]
# minimum(ForwardDiff.jacobian(_F, θ)[:, m * p + (m-1) .+ (1:1+p)])
#
# maximum(J(z_sol))
# minimum(J(z_sol))
#
# pinv(J(z_sol))
# F(z_sol)


# Definition of the composite type "Resolution"

using Printf

"""
    struct Resolution

`Resolution` objects are needed to perform a number of pixel-related
functions, e.g., convert a direction into a pixel number and vice
versa.

The fields of a `Resolution` object are the following:

- `nside`: the NSIDE parameter
- `nsideTimesTwo`: 2 * NSIDE
- `nsideTimesFour`: 4 * NSIDE
- `numOfPixels`: number of pixels in the map
- `order`: order of the map
- `pixelsPerFace`: number of pixels in each Healpix face
- `ncap`
- `fact2`
- `fact1`
"""
struct Resolution
    nside::Int
    nsideTimesTwo::Int
    nsideTimesFour::Int
    numOfPixels::Int

    order::Int
    pixelsPerFace::Int
    ncap::Int
    fact2::Float64
    fact1::Float64
end

# Customize printing
Base.show(io::IO, r::Resolution) = @printf(io, "Healpix resolution(NSIDE = %d)", r.nside)

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

"""
    Resolution(nside) -> Resolution

Create a `Resolution` object, given a value for `NSIDE`.
"""
function Resolution(nside)
    (1 ≤ nside ≤ NSIDE_MAX) || throw(DomainError(nside))
    # The expression (nside & (nside - 1)) == 0 is a quick check for
    # detecting if nside is a power of two or not
    (nside & (nside - 1) == 0) || throw(DomainError())

    order          = ilog2(nside)
    pixelsPerFace  = nside^2
    numOfPixels    = 12pixelsPerFace
    ncap           = 2 * nside * (nside - 1)
    fact2          = 4 / numOfPixels
    fact1          = 2 * nside * fact2

    result = Resolution(nside,
                        2nside,
                        4nside,
                        numOfPixels,
                        order,
                        pixelsPerFace,
                        ncap,
                        fact2,
                        fact1)

end

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

"""
    nsideok(nside::Integer) -> Bool

Check whether `nside` is a valid `NSIDE` parameter.
"""
nsideok(nside::Integer) = (nside > 0) && ((nside) & (nside - 1) == 0)

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

"""
    nside2npix(nside::Integer) -> Integer

Return the number of pixels for a Healpix map with the specified
`NSIDE` value. If `NSIDE` is not an integer power of two, the function
throws a `DomainError` exception.
"""
function nside2npix(nside::Integer)
    (nside > 0) || throw(DomainError(nside, "`NSIDE` is not a positive number"))

    nsidelog2 = round(Int, log2(nside))
    (2^nsidelog2 == nside) || throw(DomainError(nside, "`NSIDE` is not an integer power of two"))

    12(nside^2)
end

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

"""
    npix2nside(npix::Integer) -> Integer

Given the number of pixels in a Healpix map, return the `NSIDE`
resolution parameter. If the number is invalid, throw a `DomainError`
exception.
"""
function npix2nside(npix::Integer)
    (npix % 12 == 0) || throw(DomainError(npix, "Invalid number of pixels"))

    square_root = sqrt(npix / 12)
    (square_root^2 == npix / 12) || throw(DomainError(npix, "Invalid number of pixels"))

    convert(Int, round(square_root))
end

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

"""
    nside2pixarea(nside::Integer) -> Real

Return the solid angle of a pixel in a map with the specified `NSIDE` parameter.
The result is expressed in steradians.
"""
nside2pixarea(nside::Integer) = 4π / nside2npix(nside)

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

"""
    nside2resol(nside::Integer) -> Real

Return the approximate resolution of a map with the specified `NSIDE`. The
resolution is expressed in radians, and it is the square root of the pixel
size.
"""
nside2resol(nside::Integer) = sqrt(nside2pixarea(nside))

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


using RobustAndOptimalControl, ControlSystems, MonteCarloMeasurements

d = δr()
@test d.val == 0
@test d.radius == 1

d2 = d+d
@test d2.val == 0
@test d2.radius == 2
@test d2.name === d.name

d0 = d - d
@test d0.val == 0
@test d0.radius == 0
@test d2.name === d.name

d_ = δr()
d0_ = d - d_
@test d0_ isa UncertainSS

sys = uss(d)
@test sys.nx == 0
@test sys.D == [0 1; d.radius d.val] # lft representation of uncertain scalar

D = RobustAndOptimalControl.Δ(2, δc)
@test rand(D, 100) isa Matrix{Complex{Particles{Float64, 100}}}

# Example from Mu tools user guide
d = 2.4 + 0.4δr()
sys = uss(d)
@test sys.D == [0 1; 0.4 2.4]
isys = inv(sys)
@test isys.D ≈ [-1/6 1/2.4; -1/6 1/2.4] # Mu user guide seems wrong here

P = uss([δc(), δc()])
@test iszero(P.D11)
@test P.D21 == I
@test P.D12 == I
@test iszero(P.D22)
# W = tf(1.0, [1, 1])
# W*P

δel = δr(2, 1, :δel)
η = δr(6, 1, :η)
ρ = δr(-1, 1, :ρ)

s1 = uss(δel)
s2 = uss(η)

@test (1/η).D ≈ [
    -0.1667    0.1667
    -0.1667    0.1667
] atol=1e-3

@test (δel*η).D ≈ [
    0     1     6
    0     0     1
    1     2    12
] atol = 1e-3

@test (δel*η + ρ).D == [
    0     1     0     6
    0     0     0     1
    0     0     0     1
    1     2     1    11
]

@test (δel/η).D ≈ [
    0   -0.1667    0.1667
    0   -0.1667    0.1667
1.0000   -0.3333    0.3333
] atol=1e-3

@test svdvals([η+3+δel δel/η].D)[1] ≈ 11.1896 rtol=1e-2

A = [[3+δel+η δel/η]; [7+ρ ρ+δel*η]]
@test svdvals(A.D)[1] ≈ 15.3545 rtol=0.01

# lft(A.D, A.Δ) not yet supported
# A.Δ
# A2 = lft(ss(A.sys), ss(Diagonal(A.Δ)), :u)
# @test A2 ≈ A

δel_mcm = rand(δel, 100)
η_mcm = rand(η, 100)
ρ_mcm = rand(ρ, 100)

A_mcm = [[3+δel_mcm+η_mcm δel_mcm/η_mcm]; [7+ρ_mcm ρ_mcm+δel_mcm*η_mcm]]

As = rand(A, 100)
@test all(As.D .≈ A_mcm)

##
for i = 1:2, j = 1:2, k = 1:2
    d1 = δss(i,j)
    d2 = δss(j,k)
    p = d1*d2
    @test p.ny == d1.ny
    @test p.nu == d2.nu
    @test size(p.M) == (j+k, i+j)
end

##
@test uss(I(2)).M.D == I
@test sum(δss(2, 2).D) == 4
@test sum(δss(4, 4).D) == 8

##
H = δss(2, 3)
W = tf([1, .1],[.1, 1]) .* I(2)
WH = W*H
@test WH.D == [zeros(3,2) I(3); 10I(2) zeros(2,3)]

##

w = 2π .* exp10.(LinRange(-2, 2, 500))
W = makeweight(0.40,15,3)
Pn = tf(1, [1/60, 1]) |> ss
d = δss(1,1)

Gu = rand(d, 100)
@test Gu.nx == d.Δ[1].nx
@test Gu.ny == d.Δ[1].ny
@test Gu.nu == d.Δ[1].nu

@test d.ny == 1
@test d.nu == 1
@test d.nz == 1
@test d.nw == 1

@test d.zinds == 1:1
@test d.winds == 1:1
@test d.uinds == 2:2
@test d.yinds == 2:2

temp = (W*d)
@test temp.nu == temp.ny == 1
@test temp.nz == temp.nw == 1

temp = I(1) + W*d
@test temp.nu == temp.ny == 1
@test temp.nz == temp.nw == 1

@test length(d.Δ) == 1

Pd = Pn*(I(1) + W*d)

usyss = sminreal(system_mapping(Pd))
@test !iszero(usyss.B)
@test !iszero(usyss.C)
@test (usyss.C*usyss.B)[] == 60
@test usyss == Pn

Pp = rand(Pd, 50)
@test Pp.nx == 1+1+2 # == nom, W, 2 sample unc

if isinteractive()
    bodeplot(Pp, w, ylims=(1e-1, 1e1))
    bodeplot!(Pn, w)
end

Pn = ssrand(3,4,5)
@test δss(4,4, bound=0.2).D == [0I(4) sqrt(0.2)*I(4); sqrt(0.2)*I(4) 0I(4)]
mu = I(4) + δss(4,4, bound=0.2)

@test mu.D == [0I(4) sqrt(0.2)*I(4); sqrt(0.2)*I(4) I(4)]
@test size(Pn) == size(convert(UncertainSS, Pn))
@test convert(UncertainSS, Pn).D22 == Pn.D

P = Pn*mu
Pn2 = system_mapping(P)
@test Pn2 == Pn
@test size(P) == (7,8)
@test size(P.M) == (4,4)

## this time mimo real
delta = uss([δr(), δr()])
a = 1
P = ss([0 a; -a -1], I(2), [1 a; 0 1], 0)* (I(2) + delta)
K = ss(I(2))

G = lft(P, -K)
hn = norm(G, Inf)

w = exp10.(LinRange(-5, 1, 100))
M = freqresp(G.M, w)
M = permutedims(M, (2,3,1))
# mu = mussv_DG(M)
# maximum(mu)
# # maximum(structured_singular_value(M))
# @test 1/maximum(mu) ≈ √(2) atol=0.01

## Test block structure
d = δr()
@test RobustAndOptimalControl.makeblock([d]) == [-1, 0]
@test RobustAndOptimalControl.makeblock([d,d]) == [-2, 0]

d = δc()
@test RobustAndOptimalControl.makeblock([d]) == [1,0]
@test RobustAndOptimalControl.makeblock([d,d]) == [2, 0]
@test RobustAndOptimalControl.makeblock(δss(2,3).Δ) == [2, 3]

Δ = [δr(0,1,:qkarne), δr(0,1,:chuck_norris)]
@test RobustAndOptimalControl.block_structure(Δ)[1] == [[-1, 0], [-1, 0]]

Δ = [δr(0,1,:qkarne), δc(0,1,:chuck_norris)]
@test RobustAndOptimalControl.block_structure(Δ)[1] == [[1, 0], [-1, 0]] # The two elements should have been sorted and 1,-1 are therefore switched

Δ = [δr(0,1,:a), δss(2,3).Δ[], δr(0,1,:a), δc(0,1,:chuck_norris), δr(0,1,:a)]

@test RobustAndOptimalControl.block_structure(Δ)[1] == [
    [2, 3], # The gensym name will come first
    [-3, 0],
    [1, 0],
]

perm = RobustAndOptimalControl.block_structure(Δ)[2]
@test perm == [2,1,3,5,4]

P = partition(ssrand(8, 7, 2), 6, 7)
P = UncertainSS(P, Δ)

blocks, M = RobustAndOptimalControl.blocksort(P)
@test blocks == [
    [2, 3], # The gensym name will come first
    [-3, 0],
    [1, 0],
]

@test M[7,6] == P.M[6, 5] # test that the permutation was correctly applied, only test for the complex permutation
@test M[1:3,1:2] == P.M[2:4, 2:3]# test that the permutation was correctly applied, only test for the matrix block

## Test structured_singular_value with different kinds of blocks
w = exp10.(LinRange(-2, 2, 500))
delta = δss(1,1)
P = (tf(1,[1, .2, 1])) * (1+0.2*delta)
s = tf("s")
K = ss(1 + 2/s + 0.9s/(0.1s+1))
Gcl = lft(P, -K)

# These should all be equivalent
# Full block complex uncertainty gives μ = hinfnorm
@test norm(Gcl, Inf) ≈ 1.1861236982687908 rtol=1e-3
mu = structured_singular_value(Gcl, w)
@test 1/norm(mu, Inf) ≈ 1/1.1861236982687908 atol=0.005
@test RobustAndOptimalControl.robstab(Gcl) ≈ 0.84 atol=0.005
@test 1/structured_singular_value(Gcl) ≈ 0.84 atol=0.005

## same as above but with scalar instead of 1×1 system
w = exp10.(LinRange(-2, 2, 500))
delta = δc()
P = (tf(1,[1, .2, 1])) * (1+0.2*delta)
s = tf("s")
K = ss(1 + 2/s + 0.9s/(0.1s+1))
Gcl = lft(P, -K)

# These should all be equivalent
# Full block complex uncertainty gives μ = hinfnorm
@test norm(Gcl, Inf) ≈ 1.1861236982687908 rtol=1e-3
mu = structured_singular_value(Gcl, w)
@test 1/norm(mu, Inf) ≈ 1/1.1861236982687908 atol=0.005
@test RobustAndOptimalControl.robstab(Gcl) ≈ 0.84 atol=0.005
@test 1/structured_singular_value(Gcl) ≈ 0.84 atol=0.005

## this time mimo complex
delta = uss([δc(), δc()])
P = ss([0 1; 0 0], I(2), [1 0], 0) * (I(2) + delta)
# diagonal input uncertainty

K = ss([1;1])
G = lft(P, -K)
hn = norm(G, Inf)
@test robstab(G) ≈ 0.4824 atol=0.001
dm = diskmargin(system_mapping(P), K, 0, w)
dmm = argmin(dm->dm.margin, dm.simultaneous_input)
@test dmm.margin ≈ 0.5335 atol = 0.001


function plane_strain_transformation_stress(lambda, mu, theta)
    pt = (lambda + 2mu / 3) * theta
    transfstress = [pt, pt, 0.0]
    return transfstress
end

function bulk_transformation_rhs(
    transfstress,
    basis,
    quad,
    dim,
    jac,
    detjac,
    vectosymmconverter,
)

    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    rhs = zeros(ndofs)

    for (p, w) in quad
        grad = transform_gradient(gradient(basis, p), jac)
        NK = sum([
            make_row_matrix(vectosymmconverter[k], grad[:, k]) for k = 1:dim
        ])
        rhs .+= NK' * transfstress * detjac * w
    end
    return rhs
end

function coherent_interface_transformation_linear_form(
    transfstress,
    basis,
    quad,
    normals,
    dim,
    scalearea,
    vectosymmconverter,
)

    numqp = length(quad)
    @assert size(normals) == (dim, numqp)
    @assert length(scalearea) == numqp

    nf = number_of_basis_functions(basis)
    totalndofs = dim * nf
    rhs = zeros(totalndofs)

    for qpidx = 1:numqp
        p, w = quad[qpidx]
        vals = basis(p)
        NI = interpolation_matrix(vals, dim)

        normal = normals[:, qpidx]
        NK = sum([normal[k] * vectosymmconverter[k]' for k = 1:dim])

        rhs .+= NI' * NK * transfstress * scalearea[qpidx] * w
    end
    return rhs
end

function component_interface_transformation_linear_form(
    transfstress,
    basis,
    quad,
    components,
    normals,
    dim,
    scalearea,
    vectosymmconverter,
)

    numqp = length(quad)
    @assert size(components) == (dim, numqp)
    @assert size(normals) == (dim, numqp)
    @assert length(scalearea) == numqp
    nf = number_of_basis_functions(basis)
    ndofs = dim * nf
    rhs = zeros(ndofs)

    for qpidx = 1:numqp
        p, w = quad[qpidx]

        vals = basis(p)
        normal = normals[:, qpidx]
        component = components[:, qpidx]

        projector = component * component'
        NK = sum([normal[k] * vectosymmconverter[k]' for k = 1:dim])
        NI = interpolation_matrix(vals, dim)

        rhs .+= NI' * projector * NK * transfstress * scalearea[qpidx] * w
    end
    return rhs
end