"""
Module Elements --- elements of free modules.

A  `ModuleElt{K,V}`  represents  an  element  of  a free module where basis
elements  are of type  `K` and coefficients  of type `V`.  Usually you want
objects  of type `V` to be elements of  a ring, but it could also be useful
if  they just belong to  an abelian group. This  is similar to the SageMath
CombinatorialFreeModule.  You can  also see  them as  `SparseVector`s where
keys can be type `K` instead of integers.

This  basic  data  structure  is  used  in  my  packages  as  an  efficient
representation   at  many   places.  For   example,  the   `Monomial`  type
representing multivariate monomials is a `ModuleElt{Symbol,Int}`:

`x^2y^-3 ` is represented by `ModuleElt(:x=>2,:y=>-3)`

And  multivariate polynomials are  represented by a `ModuleElt{Monomial,C}`
where `C` is the type of the coefficients:

`x*y-z^2` is represented by ``ModuleElt(x*y=>1,z^2=>-1)

`ModuleElts`  are  also  used  for  cyclotomics, CycPols, elements of Hecke
algebras, etc…

A  `ModuleElt{K,V}` is essentially a  list of `Pairs{K,V}`. The constructor
takes  as argument a list of pairs, or a variable number of pair arguments,
or a generator of pairs.

We provide two implementations:

  - `HModuleElt`, an implementation by `Dict`s

This  requires  that  the  type  `K`  is  hashable.  It  is  a  very simple
implementation  since the interface of the type  is close to that of dicts;
the  only difference is weeding  out keys which have  a zero cofficient ---
which  is necessary since for testing equality of module elements one needs
a canonical form for each element.

-  a faster implementation  `ModuleElt` is obtained  by keeping the list of
pairs  sorted by key. This demands that the type `K` has a `isless` method.
This  implementation is two  to four times  faster than the  `Dict` one and
requires half the memory.

Both  implementations  have  the  same  methods,  which are mostly the same
methods  as a  `Dict` (`haskey`,  `getindex`, `setindex`, `keys`, `values`.
`pairs`,  `first`,  `iterate`,  `length`,  `eltype`), with some exceptions.
Adding  elements is implemented  as `merge(+,...)` which  is a variation on
`merge`  for  `Dict`s  where  keys  with  zero  value are deleted after the
operation (here `+` can be replaced by any operation `op` with the property
that `op(0,x)=op(x,0)=x`).

A module element can also be negated, or multiplied or divided (`/`or `//`)
by  some element (acting on coefficients)  if the method is defined between
type `V` and that element; there are also `zero` and `iszero` methods.

`ModuleElt`s  have  methods  `cmp`  and  `isless` which `HModuleElt`s don't
have. There is also `ModuleElts.merge2` which does the same as merge but is
valid  for more  general operations  (I use  it with  `min` and `max` which
implement `gcd` and `lcm` for `Monomial`s and `CycPol`s).

Here  is an example where basis elements are `Symbol`s and coefficients are
`Int`. As you can see in the examples, at the REPL (or in Jupyter or Pluto,
when  `IO`  has  the  `:limit`  attribute)  the  `show`  method  shows  the
coefficients  (bracketed  if  necessary,  which  is  when  they  have inner
occurences  of `+-*/`), followed by showing  the basis elements. The `repr`
method gives a representation which can be read back in julia:

```julia-repl
julia> a=ModuleElt(:xy=>1,:yx=>-1)
:xy-:yx

julia> repr(a)
"ModuleElt([:xy => 1, :yx => -1])"

julia> ModuleElt([:xy=>1//2,:yx=>-1])
(1//2):xy+(-1//1):yx
```

Setting  the  `IO`  property  `:showbasis`  to  a  custom printing function
changes how the basis elements are printed.

```julia-rep1
julia> show(IOContext(stdout,:showbasis=>(io,s)->string("<",s,">")),a)
3<xy>+2<yx>
```
We illustrate basic operations on `ModuleElt`s:

```julia-repl
julia> a-a
0

julia> a*99
99:xy-99:yx

julia> a//2
(1//2):xy+(-1//2):yx

julia> a/2
0.5:xy-0.5:yx

julia> a+ModuleElt(:yx=>1)
:xy

julia> a[:xy] # indexing by a basis element finds the coefficient
1

julia> a[:xx] # the coefficient of an absent basis element is zero.
0

julia> haskey(a,:xx)
false

julia> first(a)
:xy => 1

julia> collect(a)
2-element Vector{Pair{Symbol, Int64}}:
 :xy => 1
 :yx => -1

julia> collect(keys(a))
2-element Vector{Symbol}:
 :xy
 :yx

julia> collect(values(a))
2-element Vector{Int64}:
  1
 -1

julia> length(a)
2

julia> eltype(a)
Pair{Symbol, Int64}
```

In both implementations the constructor normalizes the constructed element,
removing zero coefficients and merging duplicate basis elements, adding the
corresponding   coefficients  (and   sorting  the   basis  in  the  default
implementation).  If  you  know  this  normalisation is unnecessary, to get
maximum  speed you can disable this  by giving the keyword `check=false` to
the constructor.

```julia-repl
julia> a=ModuleElt(:yy=>1, :yx=>2, :xy=>3, :yy=>-1;check=false)
:yy+2:yx+3:xy-:yy

julia> a=ModuleElt(:yy=>1, :yx=>2, :xy=>3, :yy=>-1)
3:xy+2:yx
```

Adding  or subtracting `ModuleElt`s does promotion  on the type of the keys
and the coefficients if needed:

```julia-repl
julia> a+ModuleElt([:z=>1.0])
3.0:xy+2.0:yx+1.0:z
```
"""
module ModuleElts

export ModuleElt, HModuleElt # data structure

#------------- implementation with Dicts ----------------------
struct HModuleElt{K,V}
  d::Dict{K,V} # the keys K should be hashable
  function HModuleElt(d::Dict{K,V};check::Bool=true)where {K,V}
    if check
      for (k,v) in d if iszero(v) delete!(d,k) end end
    end
    new{K,V}(d)
  end
end

function HModuleElt(x::Vector{Pair{K,V}};check=true) where{K,V}
  if !check || length(x)<=1 return HModuleElt(Dict(x)) end
  res=Dict{K,V}()
  for (k,v) in x
    if haskey(res,k) res[k]+=v
    else res[k]=v
    end
  end
  HModuleElt(res)
end

# forwarded methods
Base.haskey(x::HModuleElt,y...)=haskey(x.d,y...)
Base.keys(x::HModuleElt)=keys(x.d)
Base.values(x::HModuleElt)=values(x.d)

Base.getindex(x::HModuleElt{K,V},i) where{K,V}=haskey(x,i) ?  x.d[i] : zero(V)
Base.setindex!(x::HModuleElt{K,V},v,i) where{K,V}=
  if iszero(v) delete!(x.d,v) else x.d[i]=v end

# Valid for ops such that op(0,x)=x otherwise the result is wrong.
Base.merge(op::Function,a::HModuleElt,b::HModuleElt)=HModuleElt(merge(op,a.d,b.d))

@inline Base.:-(a::HModuleElt,b::HModuleElt)=a+(-b) # since merge2 not available
#-------------- faster implementation -------------------------------------
"""
`ModuleElt{K,V}`  has a  similar interface  to `Dict{K,V}`,  but instead of
assuming that `K` is hashable, it assumes that `K` is sortable. It also has
the advantage that ModuleElts are naturally sortable.

The  only  field,  a  `Vector{Pair{K,V}}`,  is  kept  sorted  by  `K`;  the
constructor  by  default  checks  sorting,  adds  values with same key, and
suppresses  keys  with  zero  value.  This  can  be bypassed by the keyword
`check=false`.
"""
struct ModuleElt{K,V}
  d::Vector{Pair{K,V}}
  function ModuleElt(d::AbstractVector{Pair{K,V}};check::Bool=true)where {K,V}
    if check normalize!(d) end
    new{K,V}(d)
  end
end

function normalize!(d::Vector{<:Pair})
  if isempty(d) return end
  if !issorted(d,by=first) sort!(d,by=first) end
  i=1
  ki,vi=d[i]
@inbounds for j in 2:length(d)
    if first(d[j])==ki
      vi+=last(d[j])
      d[i]=ki=>vi
    else
      if !iszero(vi) i+=1 end
      ki,vi=d[i]=d[j]
    end
  end
  resize!(d,iszero(vi) ? i-1 : i)
end

@inline Base.cmp(x::ModuleElt,y::ModuleElt)=cmp(x.d,y.d)
@inline Base.isless(x::ModuleElt,y::ModuleElt)=cmp(x,y)==-1

"""
`merge(op::Function,a::ModuleElt,b::ModuleElt)`

is  like merge(op,a,b) for Dicts, except that keys with value 0 are deleted
after the operation is done.

The  code is only  valid for `op`s  such that `op(0,x)=op(x,0)=x` otherwise
the result is wrong.
"""
function Base.merge(op::Function,a::ModuleElt,b::ModuleElt)
  (a,b)=promote(a,b)
  la=length(a.d)
  lb=length(b.d)
  res=similar(a.d,la+lb)
  ai=bi=1
  ri=0
@inbounds while ai<=la || bi<=lb
    if     ai>la res[ri+=1]=b.d[bi]; bi+=1
    elseif bi>lb res[ri+=1]=a.d[ai]; ai+=1
    else c=cmp(first(a.d[ai]),first(b.d[bi]))
      if     c>0 res[ri+=1]=b.d[bi]; bi+=1
      elseif c<0 res[ri+=1]=a.d[ai]; ai+=1
      else s=op(last(a.d[ai]),last(b.d[bi]))
        if !iszero(s) res[ri+=1]=first(a.d[ai])=>s end
        ai+=1; bi+=1
      end
    end
  end
  ModuleElt(resize!(res,ri);check=false)
end

"""
`merge2(op::Function,a::ModuleElt,b::ModuleElt)`

does  `op` between coefficients of  the same basis element  in `a` and `b`.
This  version works for  general ops (not  necessarily commutative or which
need not satisfy op(0,x)=op(x,0)=x), but has too much overhead currently to
replace  `merge` for  + or  other ops  such that  op(0,x)==op(x,0)=x. It is
useful for max or min which do lcm and gcd of `Monomial`s or `CycPol`s.
"""
function merge2(op::Function,a::ModuleElt,b::ModuleElt)
  (a,b)=promote(a,b)
  la=length(a.d)
  lb=length(b.d)
  res=similar(a.d,la+lb)
  ai=bi=1
  ri=0
@inbounds while ai<=la || bi<=lb
    if     ai>la
      be=b.d[bi]
      s=op(zero(last(be)),last(be))
      if !iszero(s) res[ri+=1]=first(be)=>s end
      bi+=1
    elseif bi>lb
      ae=a.d[ai]
      s=op(last(ae),zero(last(ae)))
      if !iszero(s) res[ri+=1]=first(ae)=>s end
      ai+=1
    else c=cmp(first(a.d[ai]),first(b.d[bi]))
      if     c>0
        be=b.d[bi]
        s=op(zero(last(be)),last(be))
        if !iszero(s) res[ri+=1]=first(be)=>s end
        bi+=1
      elseif c<0
        ae=a.d[ai]
        s=op(last(ae),zero(last(ae)))
        if !iszero(s) res[ri+=1]=first(ae)=>s end
        ai+=1
      else s=op(last(a.d[ai]),last(b.d[bi]))
        if !iszero(s) res[ri+=1]=first(a.d[ai])=>s end
        ai+=1; bi+=1
      end
    end
  end
  ModuleElt(resize!(res,ri);check=false)
end

Base.:-(a::ModuleElt,b::ModuleElt)=merge2(-,a,b) # faster than a+(-b)

"""
`getindex(x::ModuleElt,key)`  returns  the  value  associated  to the given
`key` in `x`. It returns zero if the key does not occur in `x`.
"""
function Base.getindex(x::ModuleElt{K,V},i) where {K,V}
  r=searchsortedfirst(x.d,i=>zero(V);by=first)
  if r>length(x.d) || first(x.d[r])!=i return zero(V) end
  last(x.d[r])
end

"""
`setindex!(x::ModuleElt,v,key)`  sets `v`  as the  value associated  to the
given `key` in `x`. Setting a value to zero or for a new key is expensive.
"""
function Base.setindex!(x::ModuleElt{K,V},v,key) where {K,V}
  r=searchsortedfirst(x.d,key=>zero(V);by=first)
  if r>length(x.d) || first(x.d[r])!=key
    if !iszero(v) splice!(x.d,r:r-1,[key=>v]) end
  elseif iszero(v) deleteat!(x.d,r)
  else x.d[r]=key=>v
  end
  v
end

function Base.haskey(x::ModuleElt{K,V},i) where {K,V}
  r=searchsortedfirst(x.d,i=>zero(V);by=first)
  r<=length(x.d) && first(x.d[r])==i
end

@inline Base.pairs(x::ModuleElt)=x.d
Base.keys(x::ModuleElt)=(first(p) for p in x.d)
Base.values(x::ModuleElt)=(last(p) for p in x.d)

# copied from Util.jl in order to have a self-contained file.
# Tries to determine which coefficients should be bracketed for unambiguous
# display
function format_coefficient(c::String)
  if c=="1" ""
  elseif c=="-1" "-"
  elseif match(r"^[-+]?([^-+*/]|√-|{-)*(\(.*\))?$",c)!==nothing c
  else "("*c*")"
  end
end

#-------------- methods which have same code in both implementations-------
# we assume that converting the keys from K to K1 does not change hashing
# for the Dict implementation or sorting for the other implementation
function Base.convert(::Type{ModuleElt{K,V}},a::ModuleElt{K1,V1}) where {K,K1,V,V1}
  if K==K1
    if V==V1 a
    elseif iszero(a) zero(ModuleElt{K,V})
    else ModuleElt(k=>convert(V,v) for (k,v) in a.d;check=false)
    end
  else
    if iszero(a) zero(ModuleElt{K,V})
    elseif V==V1  ModuleElt(convert(K,k)=>v for (k,v) in a.d;check=false)
    else ModuleElt(convert(K,k)=>convert(V,v) for (k,v) in a.d;check=false)
    end
  end
end

# I get a bug when putting this routine in the next loop. So here duplicated
function Base.convert(::Type{HModuleElt{K,V}},a::HModuleElt{K1,V1}) where {K,K1,V,V1}
  if K==K1
    if V==V1 a
    elseif iszero(a) zero(HModuleElt{K,V})
    else HModuleElt(k=>convert(V,v) for (k,v) in a.d;check=false)
    end
  else
    if iszero(a) zero(HModuleElt{K,V})
    elseif V==V1  HModuleElt(convert(K,k)=>v for (k,v) in a.d;check=false)
    else HModuleElt(convert(K,k)=>convert(V,v) for (k,v) in a.d;check=false)
    end
  end
end

for M in (:HModuleElt, :ModuleElt)
  @eval begin
@inline $M(x::Pair...;u...)=$M(collect(x);u...)
@inline $M(x::Base.Generator;u...)=$M(collect(x);u...)

@inline Base.:+(a::$M,b::$M)=merge(+,a,b)
Base.:-(a::$M)=iszero(a) ? a : $M(k=>-v for(k,v) in a;check=false)

# multiply module element by scalar
function Base.:*(a::$M{K,V},b)where {K,V}
  if iszero(b) || iszero(a)
    return zero($M{K,promote_type(V,typeof(b))})
  end
  let b=b
    $M(k=>v*b for (k,v) in a;check=false)
  end
end

@inline Base.iszero(x::$M)=isempty(x.d)
Base.zero(x::$M)=$M(empty(x.d))
Base.zero(::Type{$M{K,V}}) where{K,V}=$M(Pair{K,V}[])
# forwarded methods
@inline Base.:(==)(a::$M,b::$M)=a.d==b.d
@inline Base.first(x::$M)=first(x.d)
@inline Base.iterate(x::$M,y...)=iterate(x.d,y...)
@inline Base.length(x::$M)=length(x.d)
@inline Base.eltype(x::$M)=eltype(x.d)
@inline Base.hash(x::$M, h::UInt)=hash(x.d,h)

function Base.promote_rule(a::Type{$M{K1,V1}},
                           b::Type{$M{K2,V2}})where {K1,K2,V1,V2}
  $M{promote_type(K1,K2),promote_type(V1,V2)}
end

function Base.show(io::IO,m::$M)
  if !(get(io,:limit,false) || get(io,:TeX,false))
    print(io,Symbol($M),"(",m.d,")")
    return
  end
  if iszero(m) print(io,"0"); return end
  showbasis=get(io,:showbasis,nothing)
  if isnothing(showbasis)
    showbasis=(io,x)->repr(x;context=io)
  end
  start=true
  res=""
  for (k,v) in m
    v=repr(v;context=io)
    k=showbasis(io,k)
    if !isempty(k) v=format_coefficient(v) end
    if (isempty(v) || v[1]!='-') && !start v="+"*v end
    res*=v*k
    if start start=false end
  end
  if res=="" res="1" end
  print(io,res)
end

function Base.merge(f,m::$M{K,V},b;check=true)where {K,V}
  let b=b, f=f
    $M(k=>f(v,b) for (k,v) in m;check)
  end
end

function Base.merge(f,m::$M;check=true)
  $M(k=>f(v) for (k,v) in m;check)
end

end
end

for M in (:HModuleElt, :ModuleElt), op in (:/,:(//))
  @eval begin
    Base.$op(m::$M,b)=merge($op,m,b)
  end
end
end


type SGDSolver <: RegressionSolver end

function solve(model::RegressionModel, method::RegERM, ::SGDSolver, X::AbstractMatrix, y::AbstractVector)
    function loss_grad(theta::Vector, i::Int)
        grad_model = gradient(model.f, X[i,:], theta)
        grad_loss = derivs(loss(method), values(model, X[i,:], theta), [y[i]])
        vec(broadcast(*, grad_loss',grad_model))
    end

    reg_grad(theta::Vector) = gradient(regularizer(method, theta)) / method.n
    grad(theta::Vector, i::Int) = loss_grad(theta, i) + reg_grad(theta)

    sgd(grad, method.n, model.theta)
end

# step size has to statisfy the conditions sum alpha^2 < inf and sum alpha = inf
# 1 / k and 1 / sqrt(k) are common choicses
sqrt_step_size(k::Int) = 1 / sqrt(k)
lin_step_size(k::Int) = 1 / k

function sgd{T}(grad::Function,
                n::Int,
                initial_x::Array{T};
                iterations::Int=1000,
                xtol::Real = 1e-8,
                step_size::Function = sqrt_step_size)

    # Maintain current state in x and previous state in x_previous
    x, x_previous = copy(initial_x), copy(initial_x)

    # Count the total number of iterations
    iteration = 0

    # Count the total number of updates
    k = 0

    # Iterate until convergence
    converged = false
    while !converged && iteration < iterations
        # Increment the number of steps we've had to perform
        iteration += 1

        # randomly iterate over gradients
        for i = randperm(n)
            k += 1

            # Update step size
            alpha = step_size(k)

            # Update current solution
            x -= alpha*grad(x, i)
        end

        converged = norm(x - x_previous) < xtol

        # Maintain a record of previous position
        copy!(x_previous, x)
    end
    x
end


using Plots, Statistics, LaTeXStrings

function Leveling(surface)
    maxlen = 0
    surf = surface
    for i in surf
        if length(i) > maxlen
            maxlen = length(i)
        end
    end
    for i in surf
        append!(i, [3 for i=1:maxlen - length(i)])
    end
    return surf
end

function Deposition_2DG(;len, tot_time, rate, color_step)
    surf = [[] for i=1:len]
    for n in 0:tot_time
        randsurf = rand(1:len,rate)
        for i in randsurf
            append!(surf[i], n%color_step)
        end
    end
    return hcat(Leveling(surf)...)
end

# theme(:dark)
# gr()

Parameter = Dict( :len => 200,
                    :tot_time => 100,
                        :rate => 10000,
                            :color_step => 10)

surf = Deposition_2DG(;Parameter...)

heatmap(surf, c = :solar, legend = false, title = L"Surface \ length = %$(Parameter[:len]), Total \ number \ of \ particles = 10^{%$(Int(log10(Parameter[:tot_time]*Parameter[:rate])))}")
savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\ComputationalPhysics-Fall2021\\ProblemSet2\\Figs\\Q2\\RBDGr.png")


# ----------------------------------------------------------------
# Radiative forcing from methane (Etminan et al. 2016 equations).
# ----------------------------------------------------------------

@defcomp rf_ch4_etminan begin

    a₃         = Parameter()             # Methane forcing constant.
    b₃         = Parameter()             # Methane forcing constant.
    N₂O_0      = Parameter()             # Initial (pre-industrial) nitrous oxide concentration (ppb).
    scale_CH₄  = Parameter()             # Scaling factor for uncertainty in methane radiative forcing.
    CH₄_0      = Parameter()             # Initial (pre-industrial) methane concentration (ppb).
    N₂O        = Parameter(index=[time]) # Atmospheric nitrous oxide concentration (ppb).
    CH₄        = Parameter(index=[time]) # Atmospheric methane concentration (ppb).

    rf_CH₄     = Variable(index=[time])  # Direct forcing from atmospheric methane concentrations (Wm⁻²).

    function run_timestep(p, v, d, t)

        if is_first(t)
            # Set initial radiative forcing to 0.
            v.rf_CH₄[t] = 0.0

        else
            # Create temporary averaging variables.
            M̄ = 0.5 * (p.CH₄[t] + p.CH₄_0)
            N̄ = 0.5 * (p.N₂O[t] + p.N₂O_0)

            # Direct methane radiative forcing.
            v.rf_CH₄[t] = ((p.a₃ * M̄ + p.b₃ * N̄ + 0.043) * (sqrt(p.CH₄[t]) - sqrt(p.CH₄_0))) * p.scale_CH₄
        end
    end
end


#
# Compute Polar Decomposition via SVD
#

immutable SVDAlg <: PolarAlg end

function solve!{T}(alg::SVDAlg,
                X::Matrix{T}, U::Matrix{T}, H::Matrix{T})

    F = svdfact(X, thin = true)
    PQt = Array(T, size(U))

    A_mul_Bt!(PQt, F[:U], F[:V])
    copy!(U, PQt)
    copy!(H, F[:V] * diagm(F[:S]) * F[:Vt])
    H = 0.5 * (H + H')

    return SVDResult{T}(U, H)
end


include("model.jl")

type cpaResults
  key
  ρ::Array
  keyGraphs::Array
  pointOfinterests::Array #the time instants at which the activity for the byte happened
  res_dir::String
end

function cpa(attack_round, model_type, cipher_type, trace_data)
  V = modelmaker(attack_round, model_type, cipher_type, trace_data)

  num_unit = trace_data.num_unit
  unit_size = trace_data.unit_size
  keymax = 2^unit_size -1
  num_trace = length(trace_data.traces[1,:])

  ρ = zeros(num_unit,keymax+1,num_trace)

  for i in 1:num_unit
    for keyGuess in 0:keymax
      for t in 1:num_trace
        ρ[i,keyGuess+1,t] = abs(cor(V[i,:,keyGuess+1], trace_data.traces[:,t]))
      end
    end
  end

  #now I have the final matrix, now to extract information
  #to recover key
  recovered_key = ""
  keyGraphs = Array{}[]
  poi = Int[]
  for i in 1:num_unit
    push!(keyGraphs, findmax(ρ[i,:,:],2)[1])
    predict_k = findmax(findmax(ρ[i,:,:],2)[1])[2]-1
    #now the correct key bytes has been found, calculating the instants of time
    #where the actual event occured for each byte
    push!(poi,findmax(ρ[i,predict_k+1,:])[2])
    recovered_key = join([recovered_key,hex(predict_k,div(unit_size,4))])
  end
  res = cpaResults(recovered_key, ρ, keyGraphs, poi, trace_data.res_dir)
end


# Sinusoidal Arm - Restless

type Sinusoidal <: BanditArmBase
    step::Int64         # Time step
    period::Int64       # Period of the sinusoidal wave
    isRestless::Bool    # Is this arm restless - True by default
    offset::Float64     # Initial pahse, random by default

    function Sinusoidal( _period::Int64 )
        new(
            0,
            _period,
            true,
            2*π*rand()
        )
    end

    function Sinusoidal( _period::Int64, _isRestless::Bool )
        new(
            0,
            _period,
            _isRestless,
            2*π*rand()
        )
    end

    function Sinusoidal( _period::Int64, _offset::Real )
        new(
            0,
            _period,
            true,
            _offset
        )
    end

    function Sinusoidal( _period::Int64, _isRestless::Bool, _offset::Real )
        new(
            0,
            _period,
            _isRestless,
            _offset
        )
    end
end

function pull!( arm::Sinusoidal )
    arm.step = arm.step + 1
    return 1/2 + 1/2 * sin( 2 * π * (arm.step-1)/arm.period + arm.offset )  # To limit rewards between 0 and 1
end

function tick!( arm::Sinusoidal )
    if arm.isRestless
        arm.step = arm.step + 1
    end
    nothing
end

function reset!( arm::Sinusoidal )
    arm.step = 0
    nothing
end


@testset "Element stiffness matrix" begin

    # 1x1 linear quadrilateral cell
    # Isoparameteric
    # 2nd order quadrature rule

    E = 1
    nu = 0.3
    k = [
        1 / 2 - nu / 6,
        1 / 8 + nu / 8,
        -1 / 4 - nu / 12,
        -1 / 8 + 3 * nu / 8,
        -1 / 4 + nu / 12,
        -1 / 8 - nu / 8,
        nu / 6,
        1 / 8 - 3 * nu / 8,
    ]
    Ke =
        E / (1 - nu^2) * [
            k[1] k[2] k[3] k[4] k[5] k[6] k[7] k[8]
            k[2] k[1] k[8] k[7] k[6] k[5] k[4] k[3]
            k[3] k[8] k[1] k[6] k[7] k[4] k[5] k[2]
            k[4] k[7] k[6] k[1] k[8] k[3] k[2] k[5]
            k[5] k[6] k[7] k[8] k[1] k[2] k[3] k[4]
            k[6] k[5] k[4] k[3] k[2] k[1] k[8] k[7]
            k[7] k[4] k[5] k[2] k[3] k[8] k[1] k[6]
            k[8] k[3] k[2] k[5] k[4] k[7] k[6] k[1]
        ]

    problem = HalfMBB((2, 2), (1.0, 1.0), 1.0, 0.3, 1.0)

    #@test ElementFEAInfo(problem).Kes[1] ≈ Ke
end


# Load modules
using Plots, LaTeXStrings
pyplot()
using ThinFilmsTools
using Optim

x = -10:0.01:10
len_x = length(x)

options = Optim.Options(
    g_abstol=1e-8, g_reltol=1e-8, iterations=10^5, show_trace=false, store_trace=true,
);

let
    # Gaussian 1 peak
    seed1 = [[0.0], [15.0, 0.0, sqrt(0.2)]]
    # We generate some data using the Gaussian model
    y_data = Utils.gaussian(x, seed1) .+ 0.5.*randn(len_x)
    sol = fit_curve_model(:gaussian, x, y_data, seed1)
    plot(x, [y_data, sol.ymodel], lab=["Data" "Model"])
    gui()
    # Gaussian 2 peaks
    seed2 = [[1.0], [15.0, -1.0, sqrt(0.2)], [10.0, 1.5, sqrt(0.1)]]
    # We generate some data using the Gaussian model
    y_data = Utils.gaussian(x, seed2) .+ 0.5.*randn(len_x)
    sol = fit_curve_model(:gaussian, x, y_data, seed2; options=options)
    plot(x, [y_data, sol.ymodel], lab=["Data" "Model"])
    gui()
end

let
    # Lorentzian 1 peak
    seed1 = [[0.0], [15.0, 0.0, 0.5]]
    # We generate some data using the Lorentzian model
    y_data = Utils.lorentzian(x, seed1) .+ 0.5.*randn(len_x)
    sol = fit_curve_model(:lorentzian, x, y_data, seed1; options=options)
    plot(x, [y_data, sol.ymodel], lab=["Data" "Model"])
    gui()
    # Lorentzian 2 peaks
    seed2 = [[0.0], [15.0, -1.0, 0.5], [30.0, 1.0, 0.5]]
    # We generate some data using the Lorentzian model
    y_data = Utils.lorentzian(x, seed2) .+ 0.5.*randn(len_x)
    sol = fit_curve_model(:lorentzian, x, y_data, seed2; options=options)
    plot(x, [y_data, sol.ymodel], lab=["Data" "Model"])
    gui()
end

let
    # Voigtian 1 peak
    seed1 = [[0.0], [18.0, 0.0, 0.0, 1.53]]
    # We generate some data using the Voigtian model
    y_data = Utils.voigtian(x, seed1) .+ randn(len_x)
    sol = fit_curve_model(:voigtian, x, y_data, seed1; options=options)
    plot(x, [y_data, sol.ymodel], lab=["Data" "Model"])
    gui()
    # Voigtian 2 peaks
    seed2 = [[0.0], [10.0, -1.6, 0.0, 1.53], [24.0, 5.0, 0.0, 1.53]]
    # We generate some data using the Voigtian model
    y_data = Utils.voigtian(x, seed2) .+ randn(len_x)
    sol = fit_curve_model(:voigtian, x, y_data, seed2; options=options)
    plot(x, [y_data, sol.ymodel], lab=["Data" "Model"])
    gui()
end


"""
Euler functions for entropy variables and more

"""

# 1D wavespeed for use in interface fluxes
function wavespeed(rho,rhou,E)
    cvel = (@. sqrt(γ*pfun(rho,rhou,E)/rho))
    return (@. sqrt(abs(rhou/rho)) + cvel)
end

Unorm(U) = sum((x->x.^2).(U))

"primitive pressure to conservative vars"
function primitive_to_conservative(rho,U,p)
    rhoU = (x->rho.*x).(U)
    unorm = Unorm(U)
    E = (@. p/(γ-1) + .5*rho*unorm)
    return (rho,rhoU...,E)
end

function primitive_to_conservative(rho,u,v,p)
    return primitive_to_conservative(rho,(u,v),p)
end
function primitive_to_conservative(rho,u,v,w,p)
    return primitive_to_conservative(rho,(u,v,w),p)
end

"inverse temperature (used in entropy conservative fluxes)"
function betafun(rho,rhoU,E)
    p = pfun(rho,rhoU,E)
    return (@. rho/(2*p))
end
function betafun(rho,rhou,rhov,E)
    return betafun(rho,(rhou,rhov),E)
end
function betafun(rho,rhou,rhov,rhow,E)
    return betafun(rho,(rhou,rhov,rhow),E)
end

"pressure as a function of ρ,u,v,E"
function pfun(rho,rhoU,E,rhounorm)
    return @. (γ-1)*(E-.5*rhounorm)
end
function pfun(rho,rhoU,E)
    rhounorm = Unorm(rhoU)./rho
    return pfun(rho,rhoU,E,rhounorm)
end

# hack
function pressure_fun(rho,rhou,rhov,E)
    rhounorm = @. (rhou^2+rhov^2)/rho
    return @. (γ-1)*(E-.5*rhounorm)
    #return pfun(rho,(rhou,rhov),E)
end

"specific energy as a function of conservative variables"
function rhoefun(rho,rhoU,E)
    rhoUnorm = Unorm(rhoU)
    return (@. E - .5*rhoUnorm/rho)
end

"Thermodynamic entropy as a function of conservative variables"
function sfun(rho,rhoU,E)
    rhoe = rhoefun(rho,rhoU,E)
    return (@. log((γ-1)*rhoe/(rho^γ)))
end

"Mathematical entropy"
function Sfun(rho,rhou,rhov,E)
    return -rho.*sfun(rho,(rhou,rhov),E)
end
function Sfun(rho,rhou,rhov,rhow,E)
    return -rho.*sfun(rho,(rhou,rhov,rhow),E)
end

"Entropy variables as functions of conservative vars"
function v_ufun(rho,rhoU,E)
    ρe = rhoefun(rho,rhoU,E)
    sU = sfun(rho,rhoU,E)
    v1 = (@. (-E + ρe*(γ + 1 - sU))/ρe)
    vU = [@. rhoUi/ρe for rhoUi in rhoU]
    vE = (@. (-rho)/ρe)
    return (v1,vU...,vE)
end
function v_ufun(rho,rhou,rhov,E)
    return v_ufun(rho,(rhou,rhov),E)
end
function v_ufun(rho,rhou,rhov,rhow,E)
    return v_ufun(rho,(rhou,rhov,rhow),E)
end

"entropy as function of entropy variables"
function s_vfun(v1,vU,vE)
    vUnorm = Unorm(vU)
    return (@. γ - v1 + vUnorm/(2*vE))
end

"specific energy as function of entropy variables"
function rhoe_vfun(v1,vU,vE)
    s = s_vfun(v1,vU,vE)
    return (@. ((γ-1)/((-vE)^γ))^(1/(γ-1)) * exp(-s/(γ-1)))
end

"Conservative vars as functions of entropy variables"
function u_vfun(v1,vU,vE)
    rhoeV = rhoe_vfun(v1,vU,vE)
    vUnorm = Unorm(vU)
    rho   = (@. rhoeV*(-vE))
    rhoU  = [@. rhoeV*(vUi) for vUi in vU]
    E     = (@. rhoeV*(1-vUnorm/(2*vE)))
    return (rho,rhoU...,E)
end
function u_vfun(v1,v2,v3,v4)
    return u_vfun(v1,(v2,v3),v4)
end
function u_vfun(v1,v2,v3,v4,v5)
    return u_vfun(v1,(v2,v3,v4),v5)
end


using TrainView

using DataFrames, CSV

# Some track and camera properties
    focal_length = 40e-3
    track_gauge = 1.435

    # horizontal curvature
    β = 0.0

    # vertical curvature
    α = 0.0

    left_track  = [[x, -track_gauge/2.0 + β*x^2, α*x^2] for x = LinRange(4.0,40.0,40)]
    right_track = [[x, track_gauge/2.0 + β*x^2, α*x^2] for x = LinRange(4.0,40.0,40)]

# The reference track image
    cameraposition_reference = [0.0,0.0,2.4];
    camera_reference = VideoCamera(cameraposition_reference;
        focal_length = focal_length
        # , ψθφ = [0.0,-5.0,0.0] .* (pi/180.0)
    )

    left_ref_uvs = camera_image(camera_reference, left_track);
    right_ref_uvs = camera_image(camera_reference, right_track);

# The disturbed track image
    cameraposition = cameraposition_reference + [0.0,-0.06,0.05];
    camera = VideoCamera(cameraposition;
        focal_length = focal_length
        , ψθφ = [1.05,-1.2,0.001] .* (pi/180.0)
    )

    uvs_left = camera_image(camera, left_track);
    uvs_right = camera_image(camera, right_track);

# Save the results
    data_ref = DataFrame(
        left_ref_u = [uv[1] for uv in left_ref_uvs],
        left_ref_v = [uv[2] for uv in left_ref_uvs],
        right_ref_u = [uv[1] for uv in right_ref_uvs],
        right_ref_v = [uv[2] for uv in right_ref_uvs]
    )
    data = DataFrame(
        left_u = [uv[1] for uv in uvs_left],
        left_v = [uv[2] for uv in uvs_left],
        right_u = [uv[1] for uv in uvs_right],
        right_v = [uv[2] for uv in uvs_right]
    )

# CSV.write("reference_image.csv", data_ref)
# CSV.write("disturbed_image.csv", data)

# Plot the results
using Plots

w = camera.opticalproperties.sensor_width;
h = camera.opticalproperties.sensor_height;

plot([-w/2.0,w/2.0,w/2.0,-w/2.0,-w/2.0],[-h/2.0,-h/2.0,h/2.0,h/2.0,-h/2.0], color=:yellow, lab = "")

w = camera_reference.opticalproperties.sensor_width;
h = camera_reference.opticalproperties.sensor_height;

plot!([-w/2.0,w/2.0,w/2.0,-w/2.0,-w/2.0],[-h/2.0,-h/2.0,h/2.0,h/2.0,-h/2.0], color=:red, lab = "")

scatter!([uv[1] for uv in left_ref_uvs],[uv[2] for uv in left_ref_uvs])
scatter!([uv[1] for uv in right_ref_uvs],[uv[2] for uv in right_ref_uvs])
scatter!([uv[1] for uv in left_uvs],[uv[2] for uv in left_uvs])
scatter!([uv[1] for uv in right_uvs],[uv[2] for uv in right_uvs])


# this was taken from a gist by luke stagner
# https://gist.github.com/floswald/61e201d64f1ee9a5f3b6
# there was no license file.

type PolyharmonicSpline
  dim::Int64
  order::Int64
  coeff::Vector{Float64}
  centers::Array{Float64,2}
  error::Float64
end

function polyharmonicK(r,K)
  if iseven(K)
    return r< 1 ? (r.^(K-1))*log(r.^r) : (r^K)*log(r)
  else
    return r^K
  end
end

function PolyharmonicSpline(K::Int64, centers::Array{Float64,2},values::Array{Float64}; s = 0.0)
  m,n = size(centers)
  m != length(values) && throw(DimensionMismatch())

  M = zeros(m,m)
  N = zeros(m,n+1)

  for i=1:m
    N[i,1] = 1
    N[i,2:end] = centers[i,:]
    for j=1:m
      M[i,j] = polyharmonicK(norm(centers[i,:] .- centers[j,:]),K)
    end
  end
  M = M .+ s.*eye(m)
  L = [[M N],[N' zeros(n+1,n+1)]]

  w = pinv(L)*[values,zeros(n+1)]

  ivalues = zeros(m)
  for i=1:m
    tmp = 0.0
    for j=1:m
      tmp = tmp + w[j]*polyharmonicK(norm(centers[i,:] .- centers[j,:]),K)
    end
    tmp = tmp + w[m+1]
    for j=2:n+1
      tmp = tmp + w[m+j]*centers[i,j-1]
    end
    ivalues[i] = tmp
  end
  error = norm(values .- ivalues)

  return PolyharmonicSpline(n,K,w,centers,error)
end

function PolyharmonicSpline(K::Int64, centers::Vector{Float64},values::Vector{Float64};s = 0.0)
  PolyharmonicSpline(K,centers'',values,s=s)
end

function interpolate(S::PolyharmonicSpline,x::Array{Float64,2})
  m,n = size(x)

  n != S.dim && throw(DimensionMismatch("$m != $(S.dim)"))

  interpolates = zeros(m)
  for i=1:m
    tmp = 0.0
    l = length(S.coeff)-(n+1)
    for j=1:l
      tmp = tmp + S.coeff[j]*polyharmonicK(norm(x[i,:] .- S.centers[j,:]),S.order)
    end
    tmp = tmp + S.coeff[l+1]
    for j=2:n+1
      tmp = tmp + S.coeff[l+j]*x[i,j-1]
    end
    interpolates[i] = tmp
  end
  return interpolates
end

function interpolate(S::PolyharmonicSpline,x::Vector{Float64})
  return interpolate(S,x'')
end

function interpolate(S::PolyharmonicSpline,x::Vector{Float64},y::Vector{Float64})
  return interpolate(S,[x y])
end

function interpolate(S::PolyharmonicSpline,x::Vector{Float64},y::Vector{Float64},z::Vector{Float64})
  return interpolate(S,[x y z])
end

function plotBasis()
  x = collect(0:0.01:1)
  d=[polyharmonicK(ix,i) for ix in x, i=collect(1:7)]
  PyPlot.plot(x,d)
end


@testset "Greedy" begin
    @polyvar x y z w
    f = x^2+y^2+1
    grid = Grid2D((-2, 2), 800)
    A = amoeba(f, grid=grid, alg=Greedy())
    @test any(A.data)

    I = imaginary_projection(f, grid=grid, alg=Greedy())
    @test any(I.data)

    f = x^2+y^2+z^2+1
    F = AmoebaFiber3D(f)
    grid = Grid3D(xlims=(-2, 2), ylims=(-2, 2), zlims=(-2, 2), res=(21, 21, 21))
    A3 = PolynomialAmoebas.amoeba(f, grid=grid, alg=Greedy())
    @test any(A3.data)

    I3 = imaginary_projection(f, grid=grid, alg=Greedy())
    @test any(I3.data)

    f = 1 + z * cis(0.8π) + z^3*w*cis(1.6π) + z*w^2*cis(0.4π) + w*cis(1.2π)
    C = coamoeba(f, alg=Greedy(), resolution=500)
    @test any(C.data)
end


#=
Be careful, these functions are not all coded to be used in a neural network, especially the functions that return 0 or 1.
They return a value and not an array.
However, you can easily modify them and adapt them to your project, as they remain very general.

Coded by Adrien I, 12/27/2020
=#

function sigmoid(z)
    return 1 ./ (1 .+ exp.(-z))
end

function sigmoid_derivative(z)
    return sigmoid(z) .* (1 .- sigmoid(z))
end

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

function tanh(z)
    return 2 ./ ((1 .+ exp.(-2 .* z)) .+ 1 )
end

function tanh_derivative(z)
    return 1 .- tanh(z) .^ 2
end

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

function relu(x)
    if x < 0
        return 0
    end
    if x >= 0
        return x
    end
end

function derivative_relu(x)
    if x < 0
        return 0
    end
    if x >= 0
        return 1
    end
end

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

function prelu(alpha, x)
    if x < 0
        return alpha .* x
    end
    if x >= 0
        return x
    end
end

function derivative_prelu(alpha, x)
    if x < 0
        return alpha
    end
    if x >= 0
        return 1
    end
end

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

function softmax(x)
    return exp.(x) ./ sum(exp.(x))
end

function derivative_softmax(x)
    return (exp.(x[1]) .* (exp.(x[2]) .+ exp.(x[3]))) ./ ((exp.(x[1]) .+ exp.(x[2]) .+ exp.(x[3])) .^ 2 )
end

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

function elu(alpha, x)
    if x < 0
        return alpha .* (exp.(x) .- 1 )
    end
    if x >= 0
        return x
    end
end

function derivative_elu((alpha, x))
    if x < 0
        return elu(alpha, x) .+ alpha
    end
    if x >= 0
        return 1
    end
end

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

function identity(x)
    return x
end

function derivative_identity(x)
    return 1
end

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

function binary(x)
    if x < 0
        return 0
    end
    if x >= 0
        return 1
    end
end

function derivative_binary(x)
    if x != 0
        return 0
    end
    if x == 0
        return nothing
    end
end


function bt03()
#
# Brugnano and Trigiante
#
    p = poly([im*ones(5);-im*ones(5);0.5im*ones(4);-0.5im*ones(4);0.75im;-0.75im])
    z = [im 5; -im  5; 0.5im  4; -0.5im  4; .75im  1; -0.75im  1];

    p, PolyZeros(z)
end


using DynamicHMC
using DiffWrappers
using ContinuousTransformations
using StochasticVolatility
import Distributions: Uniform, InverseGamma
using JLD

RNG = Base.Random.GLOBAL_RNG

ρ = 0.8
σ = 0.6
y = simulate_stochastic(ρ, σ, 10000)

# set up the model
model = StochasticVolatility(StochasticVolatility(y, Uniform(-1, 1), InverseGamma(1, 1), 10000)...)

# we start the estimation process from the true values
θ₀ = inverse(model.transformation, (ρ, σ))

# wrap for gradient calculations

fgw = ForwardGradientWrapper(model, θ₀)

# sample

sample, tuned_sampler = NUTS_tune_and_mcmc(RNG, fgw, 5000; q = θ₀)

# posterior analysis

posterior = variable_matrix(sample)

sample_ρ, sample_σ = map_by_row(model.transformation, posterior)

# save

save(StochasticVolatility.path("results", "results.jld"), Dict("sample" => sample, "posterior" => posterior, "sample_ρ" => sample_ρ, "sample_σ" => sample_σ, "ρ" => ρ, "σ" => σ))


function miyak02()
#
#  test polynomial suggested by Goedecker
#  square of Fibocacci polynomial
#
    p = poly([(1.1+1.1im)*ones(4); (3.2+2.3im)*[1;1]; 2.1+1.5im])
    p = reverse(conv(p.a, p.a))
    z = [[(1.1+1.1im);(3.2+2.3im);2.1+1.5im] 2*[4;2;1]]
    p, PolyZeros(z)
end


"""
    pmkstruct(P; CF1, grade=l, fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> KRInfo, iz, ip

Determine the Kronecker-structure and infinite pole-zero structure information of the polynomial matrix `P(λ)` and
return an `KRInfo` object and the multiplicities of infinite zeros and poles.
The computation of the Kronecker-structure employs strong linearizations of `P(λ)` in either
the first companion form, if `CF1 = true`, or the second companion form, if `CF1 = false`.
The effective grade `l` to be used for linearization can be specified via the keyword argument `grade` as
`grade = l`, where `l` must be chosen equal to or greater than the degree of `P(λ)`.
The default value used for `l` is `l = deg(P(λ))`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The information on the Kronecker-structure consists of the right Kronecker indices `rki`,
left Kronecker indices `lki`, infinite elementary divisors `id`, the
number of finite eigenvalues `nf` and normal rank `nrank` and can be obtained from `KRInfo` as
`KRInfo.rki`, `KRInfo.lki`, `KRInfo.id`, `KRInfo.nf` and `KRInfo.nrank`, respectively.
For more details, see  [`KRInfo`](@ref).

The determination of the Kronecker-structure information is performed by building a companion
form linearization `M-λN` of `P(λ)` and reducing the pencil `M-λN` to an
appropriate Kronecker-like form (KLF) which exhibits the number of finite eigenvalues,
the multiplicities of the infinite eigenvalues,
the left and right Kronecker indices and the normal rank.
The Kronecker-structure information and pole-zero stucture information on `P(λ)` are recovered from
the Kronecker-structure information on `M-λN` using the results of [1] and [2].

The right Kronecker indices are provided in the integer vector `rki`.
The number of elements of `rki` is the dimension of the right nullspace of the polynomial matrix `P(λ)`
and their sum is the least degree of a right polynomial nullspace basis.

The left Kronecker indices are provided in the integer vector `lki`.
The number of elements of `lki` is the dimension of the left nullspace of the polynomial matrix `P(λ)`
and their sum is the least degree of a left polynomial nullspace basis.

The multiplicities of infinite eigenvalues are provided in the integer vector `id`,
where each `i`-th element `id[i]` is the order of an infinite elementary divisor
(i.e., the multiplicity of an infinite eigenvalue).

The multiplicities of the infinite zeros of `P(λ)` are returned in `iz` and represent
the positive differences between the multiplicities of
the infinite eigenvalues of `P(λ)` and the effective grade `l` of `P(λ)` [2].

The multiplicities of the infinite poles of `P(λ)` are returned in `ip` and represent
the absolute values of the negative differences between the multiplicities of
the infinite eigenvalues of `P(λ)` and the effective grade `l` of `P(λ)` [2].

The reduction to the KLF is performed using orthogonal similarity transformations and involves rank decisions based
on rank revealing QR-decompositions with column pivoting, if `fast = true`, or, the more reliable,
SVD-decompositions, if `fast = false`.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] F. De Terán, F. M. Dopico, D. S. Mackey, Spectral equivalence of polynomial matrices and
the Index Sum Theorem, Linear Algebra and Its Applications, vol. 459, pp. 264-333, 2014.

[2] A. Varga, On computing the Kronecker structure of polynomial matrices using Julia, June 2020,
[arXiv:2006.06825](https://arxiv.org/pdf/2006.06825).
"""
function pmkstruct(P::AbstractArray{T,3}; CF1::Bool = size(P,1) <= size(P,2) ? false : true,
                   grade::Int = pmdeg(P), fast::Bool = false, atol::Real = zero(real(T)),
                   rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   if CF1
      M, N = pm2lpCF1(P, grade = grade)
      kinfo = pkstruct(M, N, fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
      kinfo.rki == [] || (grade > 1 && (kinfo.rki .-= grade-1))
      kinfo.nrank > 0 && grade > 1 && (kinfo.nrank -= size(P,2)*(grade-1))
   else
      M, N = pm2lpCF2(P, grade = grade)
      kinfo = pkstruct(M, N, fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
      kinfo.lki == [] || (grade > 1 && (kinfo.lki .-= grade-1))
      kinfo.nrank > 0 && grade > 1 && (kinfo.nrank -= size(P,1)*(grade-1))
   end
   # determine the infinite structural indices
   grade < 0 ? is = Int[] : (is = [zeros(Int,kinfo.nrank-length(kinfo.id)); kinfo.id] .- grade)
   iz = is[is.>0]
   ip = sort(abs.(is[is.<0]))

   return kinfo, iz, ip
end
pmkstruct(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
          pmkstruct(poly2pm(P); kwargs...)
"""
    pmeigvals(P; CF1, grade = l, fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, KRInfo)

Return the finite and infinite eigenvalues of the polynomial matrix `P(λ)` in `val`
and information on the Kronecker-structure of `P(λ)` in the `KRInfo` object.
The computation of eigenvalues and Kronecker-structure employs strong linearizations of `P(λ)` in either
the first companion form, if `CF1 = true`, or the second companion form, if `CF1 = false`.
The effective grade `l` to be used for linearization can be specified via the keyword argument `grade` as
`grade = l`, where `l` must be chosen equal to or greater than the degree of `P(λ)`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The information on the Kronecker-structure consists of the right Kronecker indices `rki`,
left Kronecker indices `lki`, infinite elementary divisors `id`, the
number of finite eigenvalues `nf` and normal rank `nrank` and can be obtained from `KRInfo` as
`KRInfo.rki`, `KRInfo.lki`, `KRInfo.id`, `KRInfo.nf` and `KRInfo.nrank`, respectively. For more details, see  [`KRInfo`](@ref).

The computation of the eigenvalues is performed by building the companion form based linearization `M-λN`
of the polynomial matrix `P(λ)` and then reducing the pencil `M-λN` to an appropriate Kronecker-like form (KLF)
which exhibits the number of finite eigenvalues, the multiplicities of the infinite eigenvalues,
the left and right Kronecker indices and the normal rank.
The left and right Kronecker indices of `P(λ)` are returned in `KRInfo.rki` and `KRInfo.rki`, respectively, and
their values are recovered from the left and right Kronecker indices of `M-λN` using the results of [1].
The multiplicities of the infinite eigenvalues of `P(λ)` of effective grade `l` is returned in `KRInfo.id`.
The number of finite eigenvalues in `val` is equal to the number of finite eigenvalues of `M-λN` (returned in `KRInfo.nf`),
while the number of infinite eigenvalues in `val` is the sum of multiplicites returned in `KRInfo.id`.

The reduction of `M-λN` to the KLF is performed using orthonal similarity transformations and involves rank decisions
based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] F. De Terán, F. M. Dopico, D. S. Mackey, Spectral equivalence of polynomial matrices and
the Index Sum Theorem, Linear Algebra and Its Applications, vol. 459, pp. 264-333, 2014.
"""
function pmeigvals(P::AbstractArray{T,3}; CF1::Bool = size(P,1) <= size(P,2) ? false : true,
                   grade::Int = pmdeg(P), fast::Bool = false, atol::Real = zero(real(T)),
                   rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   if CF1
      M, N = pm2lpCF1(P, grade = grade)
      val, kinfo = peigvals(M, N, fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
      kinfo.rki == [] || (grade > 1 && (kinfo.rki .-= grade-1))
      kinfo.nrank > 0 && grade > 1 && (kinfo.nrank -= size(P,2)*(grade-1))
   else
      M, N = pm2lpCF2(P, grade = grade)
      val, kinfo = peigvals(M, N, fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
      kinfo.lki == [] || (grade > 1 && (kinfo.lki .-= grade-1))
      kinfo.nrank > 0 && grade > 1 && (kinfo.nrank -= size(P,1)*(grade-1))
   end
   return val, kinfo
end
pmeigvals(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
          pmeigvals(poly2pm(P); kwargs...)
"""
    pmzeros(P; CF1, fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, iz, KRInfo)

Return the finite and infinite zeros of the polynomial matrix `P(λ)` in `val`,
the multiplicities of infinite zeros in `iz` and
information on the  Kronecker-structure of `P(λ)` in the `KRInfo` object.
The computation of zeros and Kronecker-structure employs strong linearizations of `P(λ)` in either
the first companion form, if `CF1 = true`, or the second companion form, if `CF1 = false`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The information on the Kronecker-structure consists of the right Kronecker indices `rki`,
left Kronecker indices `lki`, infinite elementary divisors `id`, the
number of finite eigenvalues `nf`, the normal rank `nrank`, and can be obtained from `KRInfo` as
`KRInfo.rki`, `KRInfo.lki`, `KRInfo.id`, `KRInfo.nf` and `KRInfo.nrank`, respectively.
For more details, see  [`KRInfo`](@ref).

The computation of the zeros is performed by building the companion form based linearization `M-λN`
of the polynomial matrix `P(λ)` and then reducing the pencil `M-λN` to an appropriate Kronecker-like form (KLF)
which exhibits information on its Kronecker structure.
The left and right Kronecker indices of `P(λ)` are returned in `KRInfo.rki` and `KRInfo.rki`, respectively, and
their values are recovered from the left and right Kronecker indices of `M-λN` using the results of [1].
The multiplicities of the infinite zeros of `P(λ)`, returned in `iz`, are the positive differences between the multiplicities of
the infinite eigenvalues of `M-λN` and the degree of `P(λ)` [2].
The number of finite zeros in `val` is equal to the number of finite eigenvalues of `M-λN` (returned in `KRInfo.nf`),
while the number of infinite zeros in `val` is the sum of multiplicites in `iz`.

The reduction of `M-λN` to the KLF is performed using orthonal similarity transformations and involves rank decisions based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] F. De Terán, F. M. Dopico, D. S. Mackey, Spectral equivalence of polynomial matrices and
the Index Sum Theorem, Linear Algebra and Its Applications, vol. 459, pp. 264-333, 2014.

[2] A. Varga, On computing the Kronecker structure of polynomial matrices using Julia, June 2020,
[arXiv:2006.06825](https://arxiv.org/pdf/2006.06825).
"""
function pmzeros(P::AbstractArray{T,3}; CF1::Bool = size(P,1) <= size(P,2) ? false : true,
                 fast = false, atol::Real = zero(real(T)),
                 rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   dm1 = pmdeg(P)-1
   if CF1
      M, N = pm2lpCF1(P)
      val, iz, kinfo = pzeros(M, N, fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
      kinfo.rki == [] || dm1 <= 0 || (kinfo.rki .-= dm1)
      kinfo.nrank > 0 && dm1 > 0 && (kinfo.nrank -= size(P,2)*dm1)
   else
      M, N = pm2lpCF2(P)
      val, iz, kinfo = pzeros(M, N, fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
      kinfo.lki == [] || dm1 <= 0 || (kinfo.lki .-= dm1)
      kinfo.nrank > 0 && dm1 > 0 && (kinfo.nrank -= size(P,1)*dm1)
   end
   if iz != [] && dm1 > 0
      iz = iz[iz .> dm1] .- (dm1)
      val = [val[1:kinfo.nf]; Inf*ones(real(T),sum(iz))]
   end
   return val, iz, kinfo
end
pmzeros(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       pmzeros(poly2pm(P); kwargs...)
"""
    pmpoles(P; CF1, fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, ip, id)

Return the finite and infinite poles of the polynomial matrix `P(λ)` in `val`,
the multiplicities of infinite poles in `ip` and the infinite elementary divisors of
`Q(λ) = [P(λ) I; I 0]` in `id`.
The computation of pole-structure employs strong linearizations of `Q(λ)` in either
the first companion form, if `CF1 = true`, or the second companion form, if `CF1 = false`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The information on the finite and infinite poles of the polynomial matrix `P(λ)` is obtained from the
Kronecker-structure information of the underlying linearization `Q(λ)` using the results of [1] and [2].
The computation of the poles is performed by building the companion form based linearization `M-λN`
of the extended polynomial matrix `Q(λ) = [P(λ) I; I 0]` (see [1]) and then reducing the pencil `M-λN` to an appropriate
Kronecker-like form (KLF) which exhibits information on its Kronecker structure.
Since `Q(λ)` is regular, `M-λN` has no left or right Kronecker indices.
The multiplicities of the infinite poles of `P(λ)`, returned in `ip`, are the positive differences between the multiplicities of
the infinite eigenvalues of `M-λN` and the degree of `P(λ)` [2].
The number of finite poles in `val` is equal to the number of finite eigenvalues of `M-λN` (returned in `KRInfo.nf`),
while the number of infinite poles in `val` is the sum of multiplicites in `ip`.
The multiplicities of the infinite eigenvalues of `Q(λ)` are returned in `id`.

The reduction of `M-λN` to the KLF is performed using orthonal similarity transformations and involves rank decisions
based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] F. De Terán, F. M. Dopico, D. S. Mackey, Spectral equivalence of polynomial matrices and
the Index Sum Theorem, Linear Algebra and Its Applications, vol. 459, pp. 264-333, 2014.

[2] A. Varga, On computing the Kronecker structure of polynomial matrices using Julia, June 2020,
[arXiv:2006.06825](https://arxiv.org/pdf/2006.06825).
"""
function pmpoles(P::AbstractArray{T,3}; kwargs...) where T
   p, m, k1 = size(P)
   Ptilde = zeros(T,p+m,p+m,k1);
   Ptilde[1:p,1:m,:] = P;
   Ptilde[1:p,m+1:m+p,1] = Matrix{T}(I,p,p);
   Ptilde[p+1:p+m,1:m,1] = Matrix{T}(I,m,m);
   val, ip, kinfo = pmzeros(Ptilde; kwargs...)
   return val, ip, kinfo.id
end
pmpoles(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       pmpoles(poly2pm(P); kwargs...)
"""
    pmzeros1(P; fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, iz, KRInfo)

Return the finite and infinite zeros of the polynomial matrix `P(λ)` in `val`,
the multiplicities of infinite zeros in `iz` and information on the
Kronecker-structure of the underlying strongly irreducible pencil based linearization of `P(λ)`
in the `KRInfo` object.

The information on the  Kronecker-structure of the underlying linearization consists of
the right Kronecker indices `rki` (the same as of `P(λ)`), left Kronecker indices `lki` (the same as of `P(λ)`),
infinite elementary divisors `id`, the number of finite eigenvalues `nf` (the same as of `P(λ)`) and the
normal rank 'nrank', and can be obtained from `KRInfo` as
`KRInfo.rki`, `KRInfo.lki`, `KRInfo.id`, `KRInfo.nf` and `KRInfo.nrank`, respectively.
For more details, see  [`KRInfo`](@ref).

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The computation of the zeros is performed by first building a strongly irreducible pencil based
linearization of `P(λ)` [1] of order `ν` as a structured system matrix pencil

              | A-λE | B-λF |
     M - λN = |------|------| ,
              | C-λG | D-λH |

such that `P(λ) = (C-λG)*inv(λE-A)*(B-λF)+D-λH`, and then reducing the pencil `M-λN` to an
appropriate Kronecker-like form (KLF) which exhibits information on its Kronecker structure.
The left and right Kronecker indices of `M-λN` and `P(λ)` are the same [2] and are
returned in `KRInfo.rki` and `KRInfo.rki`, respectively.
The multiplicities of the infinite zeros of `M-λN` and of `P(λ)` are the same [2] and are returned in `iz`.
The number of finite zeros in `val` is equal to the number of finite eigenvalues of `M-λN`
(returned in `KRInfo.nf`), while the number of infinite zeros in `val` is the sum of multiplicites in `iz`.
The normal rank of `P(λ)` can be evaluated as `KRInfo.nrank - ν`.

The reduction of `M-λN` to the KLF is performed using orthonal similarity transformations and involves rank decisions based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] F.M. Dopico, M.C. Quintana and P. Van Dooren, Linear system matrices of rational transfer functions, to appear in "Realization and Model Reduction of Dynamical Systems",
A Festschrift to honor the 70th birthday of Thanos Antoulas", Springer-Verlag. [arXiv:1903.05016](https://arxiv.org/pdf/1903.05016.pdf)

[2] G. Verghese, Comments on ‘Properties of the system matrix of a generalized state-space system’,
Int. J. Control, Vol.31(5) (1980) 1007–1009.
"""
function pmzeros1(P::AbstractArray{T,3}; fast = false, atol::Real = zero(real(T)),
                  rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   if size(P,1) <= size(P,2)
      A, E, B, F, C, G, D, H, = lpsminreal(pm2lps(P; obs = true)...; fast = fast,
                                           atol1 = atol, atol2 = atol, rtol = rtol, obs = false)
   else
      A, E, B, F, C, G, D, H, = lpsminreal(pm2lps(P; contr = true)...; fast = fast,
                                           atol1 = atol, atol2 = atol, rtol = rtol, contr = false)
   end
   val, iz, info = pzeros([A B; C D], [E F; G H]; fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
   info.nrank -= size(A,1)
   return val, iz, info
end
pmzeros1(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) where {T} =
       pmzeros1(poly2pm(P); kwargs...)
"""
    pmpoles1(P; fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, ip, id)

Return the infinite poles of the polynomial matrix `P(λ)` in `val`, the multiplicities of infinite poles in `ip` and
the infinite elementary divisors of the pole pencil `Sp(λ)` of the underlying strongly irreducible
pencil based linearization of  `Q(λ) := [P(λ) I; I 0]` in `id`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The information on the infinite poles of the polynomial matrix `P(λ)` is obtained from the
pole-structure information of the underlying pencil based linearization using the results of [1] and [2].
The determination of the pole-structure information is performed by building a strongly minimal
pencil based linearization `(A-λE,B-λF,C-λG,D-λH)` with `A-λE` a regular pencil,
satisfying `P(λ) = (C-λG)*inv(λE-A)*(B-λF)+D-λH` and reducing the pole system matrix pencil
`Sp(λ) = [A-λE -λF; -λG -λH]` to an appropriate Kronecker-like form (KLF) which exhibits
the multiplicities of the infinite eigenvalues (in excess with one to the multiplicities of infinite poles of `P(λ)`).
The multiplicities of the infinite zeros of `Sp(λ)` and of infinite poles of `P(λ)` are the same [2]
and are returned in `ip`. The number of infinite poles in `val` is equal to the sum of multiplicites in `ip`.
The multiplicities of the infinite eigenvalues of `Sp(λ)` are returned in `id`.

The reduction of `Sp(λ)` to the KLF is performed using orthonal similarity transformations and involves rank decisions
based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] F.M. Dopico, M.C. Quintana and P. Van Dooren, Linear system matrices of rational transfer functions, to appear in "Realization and Model Reduction of Dynamical Systems",
A Festschrift to honor the 70th birthday of Thanos Antoulas", Springer-Verlag. [arXiv:1903.05016](https://arxiv.org/pdf/1903.05016.pdf)

[2] G. Verghese, Comments on ‘Properties of the system matrix of a generalized state-space system’,
Int. J. Control, Vol.31(5) (1980) 1007–1009.
"""
function pmpoles1(P::AbstractArray{T,3}; fast = false, atol::Real = zero(real(T)),
                  rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   p = size(P,1)
   m = size(P,2)
   if p <= m
      A, E, B, F, C, G, D, H, = lpsminreal(pm2lps(P; obs = true)...; fast = fast,
                                           atol1 = atol, atol2 = atol, rtol = rtol, obs = false)
   else
      A, E, B, F, C, G, D, H, = lpsminreal(pm2lps(P; contr = true)...; fast = fast,
                                           atol1 = atol, atol2 = atol, rtol = rtol, contr = false)
   end
   n = size(A,1)
   M = [A zeros(T,n,p+m); zeros(T,p,n+m) I; zeros(T,m,n) I zeros(T,m,p)]
   N = [E F zeros(T,n,p); G H zeros(T,p,p); zeros(T,m,n+m+p)]
   val, iz, info = pzeros(M, N; fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
   return val, iz, info.id
end
pmpoles1(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       pmpoles1(poly2pm(P); kwargs...)
"""
      pmzeros2(P; fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, iz, KRInfo)

Return the finite and infinite zeros of the polynomial matrix `P(λ)` in `val`,
the multiplicities of infinite zeros in `iz` and information on the
Kronecker-structure of the underlying irreducible descriptor system based linearization of `P(λ)`
in the `KRInfo` object.

The information on the  Kronecker-structure of the underlying linearization consists of
the right Kronecker indices `rki` (the same as of `P(λ)`), left Kronecker indices `lki` (the same as of `P(λ)`),
infinite elementary divisors `id`, the number of finite eigenvalues `nf` (the same as of `P(λ)`) and the
normal rank 'nrank', and can be obtained from `KRInfo` as
`KRInfo.rki`, `KRInfo.lki`, `KRInfo.id`, `KRInfo.nf` and `KRInfo.nrank`, respectively.
For more details, see  [`KRInfo`](@ref).

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The computation of the zeros is performed by first building an irreducible descriptor system based
linearization of `P(λ)` [1] of order `ν` as a structured system matrix pencil

              | A-λE | B |
     M - λN = |------|---| ,
              |  C   | D |

such that `P(λ) = C*inv(λE-A)*B+D`, and then reducing the pencil `M-λN` to an
appropriate Kronecker-like form (KLF) which exhibits information on its Kronecker structure.
The left and right Kronecker indices of `M-λN` and `P(λ)` are the same [2] and are
returned in `KRInfo.rki` and `KRInfo.rki`, respectively.
The multiplicities of the infinite zeros of `M-λN` and of `P(λ)` are the same [2] and are returned in `iz`.
The number of finite zeros in `val` is equal to the number of finite eigenvalues of `M-λN`
(returned in `KRInfo.nf`), while the number of infinite zeros in `val` is the sum of multiplicites in `iz`.
The normal rank of `P(λ)` can be evaluated as `KRInfo.nrank - ν`.

The reduction of `M-λN` to the KLF is performed using orthonal similarity transformations and involves rank decisions based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] G. Verghese, B. Levy, and T. Kailath, Generalized state-space systems, IEEE Trans. Automat. Control,
26:811-831 (1981).

[2] G. Verghese, Comments on ‘Properties of the system matrix of a generalized state-space system’,
Int. J. Control, Vol.31(5) (1980) 1007–1009.
"""
function pmzeros2(P::AbstractArray{T,3}; fast = false, atol::Real = zero(real(T)),
                  rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   sys = pm2ls(P; contr = true, obs = true, fast = fast, atol = atol, rtol = rtol)
   val, iz, info = spzeros(sys...; fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
   info.nrank -= size(sys[1],1)
   return val, iz, info
end
pmzeros2(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       pmzeros2(poly2pm(P); kwargs...)
"""
    pmpoles2(P; fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> (val, ip, id)

Return the infinite poles of the polynomial matrix `P(λ)` in `val`,
the multiplicities of infinite poles in `ip` and the infinite elementary divisors of the pole pencil
of the underlying irreducible descriptor system based linearization of `P(λ)` in `id`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The information on the infinite poles of the polynomial matrix `P(λ)` is obtained from the
pole-structure information of the underlying linearization using the results of [1] and [2].
The determination of the pole-structure information is performed by building an irreducible
descriptor system realization `(A-λE,B,C,D)` with `A-λE` unimodular,
satisfying `R(λ) = C*inv(λE-A)*B+D` and reducing the pole pencil `A-λE`
to an appropriate Kronecker-like form (KLF) which exhibits
the multiplicities of the infinite eigenvalues (in excess with one to the multiplicities of infinite poles of `P(λ)`).
The multiplicities of the infinite zeros of `A-λE` and of infinite poles of `P(λ)` are the same [2]
and are returned in `ip`. The number of infinite poles in `val` is equal to the sum of multiplicites in `ip`.
The multiplicities of the infinite eigenvalues of `A-λE` are returned in `id`.

The reduction of `A-λE` to the KLF is performed using orthonal similarity transformations and involves rank decisions
based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.

[1] G. Verghese, B. Levy, and T. Kailath, Generalized state-space systems, IEEE Trans. Automat. Control,
26:811-831 (1981).

[2] G. Verghese, Comments on ‘Properties of the system matrix of a generalized state-space system’,
Int. J. Control, Vol.31(5) (1980) 1007–1009.
"""
function pmpoles2(P::AbstractArray{T,3}; fast = false, atol::Real = zero(real(T)),
                  rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   A, E, _ = pm2ls(P; contr = true, obs = true, fast = fast, atol = atol, rtol = rtol)
   val, ip, info = pzeros(A, E; fast = fast, atol1 = atol, atol2 = atol, rtol = rtol)
   return val, ip, info.id
end
pmpoles2(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       pmpoles2(poly2pm(P); kwargs...)
"""
    pmroots(P; fast = false, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> val

Compute the roots of the determinant of the regular polynomial matrix `P(λ)` in `val`.

`P(λ)` can be specified as a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`,
for which the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The roots of `det(P(λ))` are computed as the finite eigenvalues of a companion form linearization `M-λN`of `P(λ)`.
The finite eigenvalues are computed by reducing the pencil `M-λN` to an appropriate Kronecker-like form (KLF)
exhibiting the spliting of the infinite and finite eigenvalue structures of the pencil `M-λN`.
The reduction is performed using orthonal similarity transformations and involves rank decisions based on rank revealing QR-decompositions with column pivoting,
if `fast = true`, or, the more reliable, SVD-decompositions, if `fast = false`. For efficiency purposes, the reduction is only
partially performed, without accumulating the performed orthogonal transformations.

The keyword arguments `atol`  and `rtol` specify the absolute and relative tolerances for the nonzero
coefficients of `P(λ)`,  respectively.
The default relative tolerance is `n*ϵ`, where `n` is the size of the smallest dimension of `P(λ)`, and `ϵ` is the
machine epsilon of the element type of coefficients of `P(λ)`.
"""
function pmroots(P::AbstractArray{T,3}; fast = false, atol::Real = zero(real(T)),
                 rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   ispmregular(P,atol = atol, rtol = rtol) || error("The polynomial matrix is not regular")
   val, iz, _ = pmzeros(P, fast = fast, atol = atol, rtol = rtol)
   return val[1:end-sum(iz)]
end
pmroots(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
        pmroots(poly2pm(P); kwargs...)
"""
       pmrank(P; fastrank = true, atol = 0, rtol = atol > 0 ? 0 : n*ϵ)

Determine the normal rank of a polynomial matrix `P(λ)`.

`P(λ)` is a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`, for which
the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`). The normal rank of `P(λ)`
is the number of linearly independent rows or columns.

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

If `fastrank = true`, the rank is evaluated by counting how many singular values of `P(γ)` have magnitude greater
than `max(atol, rtol*σ₁)`, where `σ₁` is the largest singular value of `P(γ)` and `γ` is a randomly generated
complex value of magnitude equal to one.
If `fastrank = false`, first a structured linearization of `P(λ)` is built in the form `[A-λE B; C D]` with `A-λE`
an order `n` regular subpencil, and then its normal rank `nr` is determined. The normal rank of `P(λ)` is `nr - n`.

The keyword arguments `atol` and `rtol`, specify, respectively, the absolute and relative tolerance for the
nonzero coefficients of `P(λ)`. The default relative tolerance is `n*ϵ`, where `n` is the size of the
smallest dimension of `P(λ)` and `ϵ` is the machine epsilon of the element type of its coefficients.
"""
function pmrank(P::AbstractArray{T,3}; fastrank::Bool = true, atol::Real = zero(real(T)),
                rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   if fastrank
       return rank(pmeval(P,exp(rand()*im)), atol = atol, rtol = rtol)
   else
       p, m, k = size(P)
       d = pmdeg(P)
       if p < m
          M, N = pm2lpCF2(P)
          r = prank(M, N, atol1 = atol, atol2 = atol, rtol = rtol)
          return d < 2 ? r : r-p*(d-1)
       else
          M, N = pm2lpCF1(P)
          r = prank(M, N, atol1 = atol, atol2 = atol, rtol = rtol)
          return d < 2 ? r : r-m*(d-1)
       end
   end
end
pmrank(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       pmrank(poly2pm(P); kwargs...)
"""
    ispmregular(P; fastrank = true, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> Bool

Test whether the polynomial matrix `P(λ)` is regular (i.e., `P(λ)` is square and ``{\\small\\det(P(λ)) \\not\\equiv 0}``).
The underlying computational procedure checks that the normal rank of the square `P(λ)` is equal to its order.

`P(λ)` is a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`, for which
the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.
f
If `fastrank = true`, the rank is evaluated by counting how many singular values of `P(γ)` have magnitude greater
than `max(atol, rtol*σ₁)`, where `σ₁` is the largest singular value of `P(γ)` and `γ` is a randomly generated
complex value of magnitude equal to one.
If `fastrank = false`, first a linearization of `P(λ)` is built in a companion form `M-λN` of order `n`
and then its normal rank `nr` is determined. `P(λ)` is regular if `nr = n`.

The keyword arguments `atol` and `rtol`, specify, respectively, the absolute and relative tolerance for the
nonzero coefficients of `P(λ)`. The default relative tolerance is `n*ϵ`, where `n` is the size of the
smallest dimension of `P(λ)` and `ϵ` is the machine epsilon of the element type of its coefficients.
"""
function ispmregular(P::AbstractArray{T,3}; fastrank = true, atol::Real = zero(real(T)),
                     rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   p, m, = size(P)
   m == p || (return false)
   if fastrank
      return rank(pmeval(P,exp(rand()*im)), atol = atol, rtol = rtol) == m
   else
      p <= m ? MN = pm2lpCF1(P) : MN = pm2lpCF2(P)
      return isregular(MN..., atol1 = atol, atol2 = atol, rtol = rtol)
   end
end
ispmregular(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
       ispmregular(poly2pm(P); kwargs...)
"""
    ispmunimodular(P; fastrank = true, atol::Real = 0, rtol::Real = atol>0 ? 0 : n*ϵ) -> Bool

Test whether the polynomial matrix `P(λ)` is unimodular (i.e., `P(λ)` is square, regular and `det(P(λ)) == constant`).

The underlying computational procedure checks that `P(λ)` is square and its first companion form linearization
is regular and has no finite eigenvalues.

`P(λ)` is a grade `k` polynomial matrix of the form `P(λ) = P_1 + λ P_2 + ... + λ**k P_(k+1)`, for which
the coefficient matrices `P_i`, `i = 1, ..., k+1`, are stored in the 3-dimensional matrix `P`,
where `P[:,:,i]` contains the `i`-th coefficient matrix `P_i` (multiplying `λ**(i-1)`).

`P(λ)` can also be specified as a matrix, vector or scalar of elements of the `Polynomial` type
provided by the [Polynomials](https://github.com/JuliaMath/Polynomials.jl) package.

The keyword arguments `atol` and `rtol`, specify, respectively, the absolute and relative tolerance for the
nonzero coefficients of `P(λ)`. The default relative tolerance is `n*ϵ`, where `n` is the size of the
smallest dimension of `P(λ)` and `ϵ` is the machine epsilon of the element type of its coefficients.
"""
function ispmunimodular(P::AbstractArray{T,3}; atol::Real = zero(real(T)),
                        rtol::Real = (min(size(P)...)*eps(real(float(one(T)))))*iszero(atol)) where T
   p, m, = size(P)
   m == p || (return false)
   return isunimodular(pm2lpCF1(P)..., atol1 = atol, atol2 = atol, rtol = rtol)
end
ispmunimodular(P::Union{AbstractVecOrMat{<:Polynomial},Polynomial,Number,AbstractVecOrMat{<:Number}}; kwargs...) =
               ispmunimodular(poly2pm(P); kwargs...)


#=
The Computer Language Benchmarks Game
 https://salsa.debian.org/benchmarksgame-team/benchmarksgame/
 direct transliteration of the swift#3 program by Ralph Ganszky and Daniel Muellenborn:
 https://benchmarksgame-team.pages.debian.net/benchmarksgame/program/mandelbrot-swift-3.html
 modified for Julia 1.0 by Simon Danisch
=#
const zerov8 = ntuple(x-> 0f0, 8)

# Calculate mandelbrot set for one Vec8 into one byte
Base.@propagate_inbounds function mand8(cr, ci)
    Zr = zerov8
    Zi = zerov8
    Tr = zerov8
    Ti = zerov8
    t = zerov8
    for i in 0:49
        Zi = 2f0 .* Zr .* Zi .+ ci
        Zr = Tr .- Ti .+ cr
        Tr = Zr .* Zr
        Ti = Zi .* Zi
        t = Tr .+ Ti
        all(x-> x > 4f0, t) && break
    end
    byte = UInt8(0)
    t[1] <= 4f0 && (byte |= 0x80)
    t[2] <= 4f0 && (byte |= 0x40)
    t[3] <= 4f0 && (byte |= 0x20)
    t[4] <= 4f0 && (byte |= 0x10)
    t[5] <= 4f0 && (byte |= 0x08)
    t[6] <= 4f0 && (byte |= 0x04)
    t[7] <= 4f0 && (byte |= 0x02)
    t[8] <= 4f0 && (byte |= 0x01)
    return byte
end

function mandel_inner(rows, ci, y, N, xvals)
    @simd for x in 1:8:N
        @inbounds begin
            cr = ntuple(i-> xvals[x + (i - 1)], 8)
            rows[((y-1)*N÷8+(x-1)÷8) + 1] = mand8(cr, ci)
        end
    end
end

function mandelbrot(n = 200)
    inv_ = 2.0 / n
    N = n
    xvals = zeros(Float32, n)
    yvals = zeros(Float32, n)
    Threads.@threads for i in 0:(N-1)
        @inbounds xvals[i + 1] = i * inv_ - 1.5
        @inbounds yvals[i + 1] = i * inv_ - 1.0
    end
    rows = zeros(UInt8, n*N÷8)
    Threads.@threads for y in 1:N
        @inbounds ci = yvals[y]
        mandel_inner(rows, ci, y, N, xvals)
    end
    write(stdout, "P4\n$n $n\n")
    write(stdout, rows)
end

mandelbrot(parse(Int, ARGS[1]))


using Serialization

function julia_func(z)
    c = complex(-0.5, 0.61)
    z = z^2 + c
    return z
end

function gen_julia()
    im_width, im_height = 5000, 5000
    max_it = 300
    max_z2 = 100
    min_x, max_x = -1, 1
    min_y, max_y = -1, 1

    julia_set = zeros(im_width,im_height)
    for ix in 1:im_width
        for iy in 1:im_height
            real_part = ((max_x-min_x)/im_width)*ix + min_x
            im_part = ((max_y-min_y)/im_height)*iy + min_y
            z = complex(real_part, im_part)
            it = 0
            while(abs2(z) < max_z2 && it < max_it)
                z = julia_func(z)
                it += 1
            end
            ratio = it/max_it
            julia_set[ix,iy] = ratio
        end
    end
    return julia_set
end

function main()
    @time julia_set = gen_julia()
    serialize("julia_data",julia_set)
end

main()


# Code used to generate Semidefinite Relaxation of Mixed Integer Quadratic Objective problems with an accurate
# solution from MOSEK

# Original problem has the following format:
# min_x   ||Ax-b||_2^2
# s.t.    x ∈ Z^n

workspace()
include("../../OSSDP/Code/src/Solver.jl")

using JLD, JuMP, Mosek, OSSDP

# this function creates a matrix A that slices out the diagonal entries Xii of a vectorized square matrix x=vec(X)
function createDiagonalExtractor(n)
  A = spzeros(n,n^2)
  A[1,1] = 1
  for iii=2:n-1
    col = (iii-1)*(n+1)
    A[iii,col+1] = 1
  end
  A[n,n^2] = 1
  return A
end

function createLMIconstMatrix(n)
  # assume the decision vector is like this [x,vec(X)]
  # M = spzeros((n+1)*(2*n),n+n^2)

  # # first add parts that belong to x and assign correct position in LMI
  # for iii=1:n
  #   E = spzeros(n+1,n+1)
  #   E[iii,end] = 1
  #   E[end,iii] = 1
  #   M[:,iii] = vec(E)
  # end

  # ind = n
  # # assign correct position of elements of X in LMI constraint
  # for jjj=1:n, iii=1:n
  #   E = spzeros(n+1,n+1)
  #   E[iii,jjj] = 1
  #   ind+=1
  #   M[:,ind] = vec(E)
  # end
  A1 = [zeros(n^2,n) kron(speye(n),speye(n))]

  # takes care of vec([X;x']) part
  Aa = [A1[1:n,:];1 zeros(1,n+n^2-1)]
  for i=1:n-1
    Aa = [Aa;A1[i*n+1:(i+1)*n,:];[zeros(1,i) 1 zeros(1,n+n^2-1-i)]]
  end
  # add vec([x;0]) part
  for i=0:n-1
    Aa = [Aa;zeros(1,i) 1 zeros(1,n+n^2-1-i)]
  end
  Aa = [Aa;spzeros(1,n+n^2)]
  return Aa
end

rng = MersenneTwister(12345)
dirPath =  "../DataFiles/Julia/MIQO/"
!ispath(dirPath) && mkdir(dirPath)

nn = 25

 for iii =1:1:nn
  # choose size of problem
  np = rand(rng,5:15)
  mp = 2*np
  A = randn(rng,mp,np)
  P = A'*A
  #xc: continous relaxation solution
  xc = rand(rng,np)
  q = -P*xc

  # scale problem to get fcts = -q'Pinv*q = -1
  # s = -q'*(P\q)
  # q = q./-s
  # P = P./-s
  # put problem into solver format
  # min   x'Ix - 1*t
  # s.t.  Ax + s1 = 0
  #       x = s4
  #       -t <= c'*x <=t
  # s1 in 0, s2,s3 in R+, s4 in SDPcone

  n = np+np^2
  m = (np+1)^2+np

  D = createDiagonalExtractor(np)
  LMI = createLMIconstMatrix(np)
  Aa = [speye(np,np) -D;-LMI]
  ba = [spzeros(np,1);spzeros((np+1)^2-1,1);1]
  Pa = spzeros(np+np^2,np+np^2)
  qa = [2*q;vec(P)]
  ra = 0.

  # specify cone
  Kf = 0
  Kl = np
  Kq = []
  Ks = [(np+1)^2]

  # K = OSSDPTypes.Cone(Kf,Kl,Kq,Ks)
  # setOFF = OSSDPSettings(rho=0.1,sigma=1e-6,alpha=1.6,max_iter=5000,verbose=true,checkTermination=1,scaling = 0,eps_abs = 1e-5,eps_rel=1e-5,adaptive_rho=false)
  # res1,nothing = OSSDP.solve(Pa,qa,Aa,ba,K,setOFF);

  # model = Model()
  # @variable(model, 0 <= x[1:np] <= 1,Int)
  # @objective(model,Min, norm(A*x-b)^2)

  # status = JuMP.solve(model)

  # solve accurately once with mosek
  model = Model(solver=MosekSolver())
  @variable(model, X[1:np,1:np])
  @variable(model, x[1:np])
  @objective(model,Min, trace(P*X) + 2*q'*x)

  @SDconstraint(model, [X x;x' 1] >= 0)
  for i=1:np
    @constraint(model, X[i,i] >= x[i])
  end
  status = JuMP.solve(model)

  objTrue = getobjectivevalue(model)
  solTrue = getvalue(x)

  nr = "$(iii)"
  if iii < 10
    nr = "0$(iii)"
  end
  problemType = "MIQO Relaxation"
  problemName = "MIQO"*nr

  fn = "MIQO"*nr*".jld"
  JLD.save(dirPath*fn,"n",n,"m",m,"A",Aa,"b",ba,"P",Pa,"q",qa,"r",ra,"objTrue",objTrue,"solTrue",solTrue,"problemType",problemType,"problemName",problemName,"Kf",Kf,"Kl",Kl,"Kq",Kq,"Ks",Ks)
  println("$(iii)/$(nn) completed!")
end


function accuracy(y::AbstractVector{T}, ŷ::AbstractVector{S}) where {T,S}
    @assert length(y) == length(ŷ)
    return sum(y .== ŷ) / length(y)
end

function accuracy(y::AbstractVector{T}, ŷ::AbstractMatrix{S}) where {T,S}
    max_idxs = dropdims(mapslices(x -> findmax(x)[2], ŷ; dims=1), dims=1)
    return accuracy(y, max_idxs)
end

# true labels on rows, predicted labels on columns
function confusion_matrix(y::AbstractVector{T}, ŷ::AbstractVector{T}; norm=false) where T
    labels = sort(collect(Set(y)))
    index = Dict(l => i for (i, l) in enumerate(labels))
    num_classes = length(labels)
    mat = zeros(Int, num_classes, num_classes)
    for (t, p) in zip(y, ŷ)
        mat[index[t], index[p]] += 1
    end
    if norm
        mat = mat ./ sum(mat, dims=2)
    end
    return mat
end

function recall(y::AbstractVector{T}, ŷ::AbstractVector{T}) where T
    tn, fn, fp, tp = confusion_matrix(y, ŷ)
    return tp / (tp + fn)
end

function precision(y::AbstractVector{T}, ŷ::AbstractVector{T}) where T
    tn, fn, fp, tp = confusion_matrix(y, ŷ)
    return tp / (tp + fp)
end


using Plots
using CSV
using Printf
using DataFrames
using LaTeXStrings
using LsqFit            # Required for the statistics on the fitting process
using FrictionModels

# Read in a CSV file containing the data to fit.
#
# The file is structured as follows:
# Column 1: Time
# Column 2: Velocity
# Column 3: Normal force
# Column 4: Friction Force
println(pwd())
data = CSV.read("examples//test_data_1.csv", DataFrame)

# Construct an initial estimate for model parameters
init_guess = HyperbolicModel(
    0.15,
    2.0,
    20.0,
    1.0,
    0.1,
    0.1
)

# Fit the model
results = fit_model(
    init_guess,     # Initial Guess
    data[:,1],      # Time
    data[:,4],      # Friction Force
    data[:,3],      # Normal Force
    data[:,2]       # Velocity
)

# Solve the model
F = friction.(results.model, data[:,3], data[:,2])

# Plot the data
plt = plot(
    data[:,2], F,
    xlabel = L"v(t)",
    ylabel = L"F(t)",
    label = "Fitted Model",
    lw = 2,
    legend = :bottomright
)
plot!(
    plt,
    data[:,2], data[:,4],
    st = :scatter,
    label = "Measured Data"
)
display(plt)

# Print out the coefficients of the fitted model
@printf("Fitted Model Coefficients:\n")

# Print out statistitics related to the fit
sigma = stderror(results.fit)
confidence = confidence_interval(results.fit)

# Print out the results
@printf("friction_coefficient: %f\n\tError: %f\n\tConfidence Interval: (%f, %f)\n",
    results.model.friction_coefficient,
    sigma[1],
    confidence[1][1],
    confidence[1][2]
)
@printf("normalization_coefficient: %f\n\tError: %f\n\tConfidence Interval: (%f, %f)\n",
    results.model.normalization_coefficient,
    sigma[2],
    confidence[2][1],
    confidence[2][2]
)
@printf("dissipation_coefficient: %f\n\tError: %f\n\tConfidence Interval: (%f, %f)\n",
    results.model.dissipation_coefficient,
    sigma[3],
    confidence[3][1],
    confidence[3][2]
)
@printf("hysteresis_coefficient: %f\n\tError: %f\n\tConfidence Interval: (%f, %f)\n",
    results.model.hysteresis_coefficient,
    sigma[4],
    confidence[4][1],
    confidence[4][2]
)
@printf(
    "stribeck_velocity: %f\n\tError: %f\n\tConfidence Interval: (%f, %f)\n",
    results.model.stribeck_velocity,
    sigma[5],
    confidence[5][1],
    confidence[5][2]
)
@printf(
    "viscous_damping: %f\n\tError: %f\n\tConfidence Interval: (%f, %f)\n",
    results.model.viscous_damping,
    sigma[6],
    confidence[6][1],
    confidence[6][2]
)


using PyPlot

include("simulate_exp_gp2.jl")
nt = 10
t = collect(linspace(0,10,nt))
beta = complex(0.05,1.0)
#beta = 0.3+im*0.0
alpha = 1.0+im*0.0
#println(alpha)
#println(beta)
#beta = abs(randn())+im*randn()
#alpha = randn()+im*randn()
c1 = real(alpha)*real(beta)+imag(alpha)*imag(beta)
c2 = real(alpha)*real(beta)-imag(alpha)*imag(beta)
while c1 < 0 || c2 < 0
  beta = 0.01*abs(randn())+im*randn()
  alpha = randn()+im*randn()
  c1 = real(alpha)*real(beta)+imag(alpha)*imag(beta)
  c2 = real(alpha)*real(beta)-imag(alpha)*imag(beta)
end
println(c1," ",c2)
clf()
#for i=1:1
ndev = randn(nt)
data = simulate_exp_gp2(t,alpha,beta,ndev)
#datac = simulate_exp_gp2(t,conj(alpha),conj(beta),ndev)
#datac = simulate_exp_gp2(t,alpha,conj(beta),ndev)
# Now build matrix & Cholesky:
mat = zeros(nt,nt)
for i=1:nt
  for j=1:nt
    dt = abs(t[i]-t[j])
    mat[i,j]=exp(-real(beta)*dt)*(real(alpha)*cos(imag(beta)*dt)+imag(alpha)*sin(imag(beta)*dt))
  end
end
chmat=chol(mat)
sim=*(transpose(chmat),ndev)
# Now compute cholesky decomposition analytically:
chol_anal = zeros(Complex{Float64},nt,nt)
# Loop over columns (see 10/26/16 notes):
gam = 0.0+im*0.0
for i=1:nt
  if i > 1
    gam = sqrt(1.-exp(-2.0*beta*(t[i]-t[i-1])))
  else
    gam = 1.0
  end
# Loop over rows:
  for j=i:nt
    chol_anal[i,j]=sqrt(alpha)*gam*exp(-beta*(t[j]-t[i]))
  end
end
plot(t,0.5*real(data),color="red")
plot(t,sim,color="blue")
plot(t,sim-0.5*real(data),color="green")
#println(abs(alpha)," ",std(real(data)))
#end


function ∂(field, grid, direction)
    gp = length(grid)
    access = define_shifts(1:gp-1, :, direction)
    access_shift = define_shifts(2:gp, :, direction)
    Δf = field[access_shift...] - field[access...]
    Δg = grid[2:end] - grid[1:end-1]
    reshape_Δg = define_shifts(gp-1, 1, direction)
    Δg = reshape(Δg, reshape_Δg)
    return Δf ./ Δg
end

function ∫dz(field, grid; direction = 3)
    gp = length(grid)
    access = define_shifts(1:gp-1, :, direction)
    access_shift = define_shifts(2:gp, :, direction)
    f_avg = (field[access...] + field[access_shift...]) ./ 2
    Δg = grid[2:end] - grid[1:end-1]
    reshape_Δg = define_shifts(gp-1, 1, direction)
    Δg = reshape(Δg, reshape_Δg)
    ∫f = zeros(size(field[access...]))
    ∫f[:, :, 1] = Δg[1] * field[:, :, 1]
    for i in 2:length(Δg)
        ∫f[:, :, i] = Δg[i] * field[:, :, i] + ∫f[:, :, i-1]
    end
    return ∫f
end

function define_shifts(tmp, tmp2, direction)
    if direction==1
        return (tmp, tmp2, tmp2)
    elseif direction==2
        return (tmp2, tmp, tmp2)
    else
        return (tmp2, tmp2, tmp)
    end
    return nothing
end
#=
# this instead?
function appropriate_dims(n1, n2, N)
    y = ones(Int, n1)
    y[n2] = N
    return Tuple(y)
end
=#
function avg_other(a, direction)
    nx, ny, nz = size(a)
    if direction==1
        b = (a[:, 1:ny-1, :] + a[:, 2:ny, :]) .* 0.5
        c = (b[:, :, 1:nz-1] + b[:, :, 2:nz]) .* 0.5
        return c
    elseif direction==2
        b = (a[:, :, 1:nz-1] +a[:, :, 2:nz])  .* 0.5
        c = (b[1:nx-1, :, :] + b[2:nx, :, :]) .* 0.5
        return c
    else
        b = (a[1:nx-1, :, :] + a[2:nx, :, :]) .* 0.5
        c = (b[:, 1:ny-1, :] + b[:, 2:ny, :]) .* 0.5
        return c
    end
    return nothing
end

import Base: *

struct PartialDerivative{T}
    d::T
end
struct Gradient{T}
    g::T
end
struct ScalarField{T}
    data::T
end

function *(∇::PartialDerivative, ϕ::AbstractArray)
    tmp = ∂(ϕ, ∇.d[1], ∇.d[2])
    tmp2 = avg_other(tmp, ∇.d[2])
    return tmp2
end

function (p::PartialDerivative)(ϕ::AbstractArray)
    return *(p, ϕ)
end

*(∇::PartialDerivative, ϕ::ScalarField) = *(∇, ϕ.data)

function hydrostatic_pressure(b,x,y,z)
    p = ∫dz(b, z)
    ∂ᶻp = (b[:,:,1:end-1] + b[:,:,2:end]) ./ 2
    ∂ᶻp = avg_other(∂ᶻp, 3)
    Δx = reshape(x[1:end-1] - x[2:end], (length(x)-1, 1, 1))
    ∂ˣp = (p[1:end-1,:,:] - p[2:end,:,:]) ./ Δx
    ∂ˣp = (∂ˣp[:,1:end-1,:] + ∂ˣp[:,2:end,:]) ./ 2
    Δy = reshape(x[1:end-1] - x[2:end], (1,length(y)-1, 1))
    ∂ʸp = (p[:,1:end-1,:] - p[:,2:end,:]) ./ Δy
    ∂ʸp = (∂ʸp[1:end-1,:,:] + ∂ʸp[2:end,:,:]) ./ 2
    ∇p = [∂ˣp, ∂ʸp, ∂ᶻp]
    return ∇p
end

import LinearAlgebra: dot
function dot(ω::Array{Array{Float64,3},1}, ∇b::Array{Array{Float64,3},1})
    tmparray = copy(ω[1]) .* 0.0
    for i in eachindex(ω)
        tmparray += ω[i] .* ∇b[i]
    end
    return tmparray
end

## 2D stuff
function average1(b)
    return (b[1:end-1,:] + b[2:end, :]) * 0.5
end
function average2(b)
    return (b[:,1:end-1] + b[:, 2:end]) * 0.5
end

function Δ1(f)
    return f[2:end,:] - f[1:end-1,:]
end
function Δ2(f)
    return f[:,2:end] - f[:,1:end-1]
end


module CarterMethodBLCoords

using ..CarterBoyerLindquist
using DocStringExtensions

"""
    $(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.10):

```math
T = E \\left( r^2 + a^2 \\right) - L * a.
```
"""
@inline T(E, L, r, a) = E * (r^2 + a^2) - L * a

"""
    $(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.10), for the special case of a null-geodesic ``\\mu = 0``:

```math
V_r = T^2 - \\Delta \\left[ (L - a E)^2 + Q \\right]
```
"""
@inline Vr(E, L, M, Q, r, a) = T(E, L, r, a)^2 - Δ(M, r, a) * ((L - a * E)^2 + Q)

"""
    $(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.10), for the special case of a null-geodesic ``\\mu = 0``:

```math
V_\\theta =
    Q + \\cos^2 (\\theta) \\left[ a^2 E^2 - \\frac{L^2}{\\sin^2 (\\theta) } \\right].
```
"""
@inline Vθ(E, L, Q, a, θ) = Q + cos(θ)^2 * ((a * E)^2 - (L / sin(θ))^2)

"""
    $(TYPEDSIGNATURES)

The ``t`` compontent of the equation of motion for a photon around a black hole, multiplied
by ``\\Sigma``.

From Bardeen et al. (1972) eq. (2.9d):

```math
\\Sigma \\frac{\\text{d}t}{\\text{d}\\lambda} =
- a \\left( a E \\sin^2 \\theta - L \\right)
+ \\frac{\\left( r^2 + a^2 \\right) T}{\\Delta}.
```
"""
@inline function Σδt_δλ(E, L, M, r, a, θ)
    -a * (a * E * sin(θ)^2 - L) + (r^2 + a^2) * T(E, L, r, a) / Δ(M, r, a)
end

"""
    $(TYPEDSIGNATURES)

The ``r`` compontent of the equation of motion for a photon around a black hole, multiplied
by ``\\Sigma``.

Modified from Bardeen et al. (1972) eq. (2.9a):

```math
\\Sigma \\frac{\\text{d}r}{\\text{d}\\lambda} =
\\pm \\sqrt{\\lvert V_r \\rvert},
```

where, for implementation reason, the sign is always positive. Instead, the sign is applied
in [`δ`](@ref).
"""
@inline function Σδr_δλ(E, L, M, Q, r, a)
    V = Vr(E, L, M, Q, r, a)
    √abs(V)
end

"""
    $(TYPEDSIGNATURES)

The ``\\theta`` compontent of the equation of motion for a photon around a black hole,
multiplied by ``\\Sigma``.

Modified from Bardeen et al. (1972) eq. (2.9b):

```math
\\Sigma \\frac{\\text{d}\\theta}{\\text{d}\\lambda} =
\\pm \\sqrt{\\lvert V_\\theta \\rvert},
```

where, for implementation reason, the sign is always positive. Instead, the sign is applied
in [`δ`](@ref).
"""
@inline function Σδθ_δλ(E, L, Q, a, θ)
    V = Vθ(E, L, Q, a, θ)
    √abs(V)
end

"""
    $(TYPEDSIGNATURES)

The ``\\phi`` compontent of the equation of motion for a photon around a black hole,
multiplied by ``\\Sigma``.

From Bardeen et al. (1972) eq. (2.9c):

```math
\\Sigma \\frac{\\text{d}\\phi}{\\text{d}\\lambda} =
- \\frac{L}{\\sin^2 \\theta} - aE
+ \\frac{aT}{\\Delta}.
```
"""
@inline function Σδϕ_δλ(E, L, M, r, a, θ)
    (L / sin(θ)^2) - (a * E) + a * T(E, L, r, a) / Δ(M, r, a)
end

"""
$(TYPEDSIGNATURES)

From Fanton et al. (1997), eq. (76):

```math
S = 1 + \\alpha \\omega \\sin \\theta.
```
"""
S(θ, α, ω) = 1 + α * ω * sin(θ)

"""
$(TYPEDSIGNATURES)

Calculates and returns the observer's angles ``\\sin \\Theta`` and ``\\sin \\Phi``, where
the parameters in the function signature correspond to the Bardeen et al. (1972) paper.

From Fanton et al. (1997), eq. (74) and eq. (75):

```math
\\begin{align*}
\\sin \\Theta &=
\\frac{\\alpha \\Sigma}{\\sqrt{A}}
\\left\\{
        \\beta^2 + \\left( \\alpha + a \\sin \\theta \\right)^2
        + \\frac{
            A S^2 - \\left( r^2 + a^2 + a \\alpha \\sin \\theta \\right)
        }{\\Delta}
\\right\\}, \\\\
\\sin \\Phi &=
-\\frac{\\alpha \\Sigma \\sqrt{\\Delta}}{A S \\sin\\Theta}.
\\end{align*}
```
"""
function sinΦsinΨ(Σ₀, sinθ, A₀, Δ₀, S₀, r, a, α, β)
    # calc 1
    sinΦ =
        ((α * Σ₀) / √A₀) /
        sqrt(β^2 + (α + a * sinθ)^2 + (A₀ * S₀^2 - (r^2 + a^2 + a * α * sinθ)^2) / Δ₀)

    # calc 2
    sinΨ = -(α * Σ₀ * √Δ₀) / (S₀ * A₀ * sinΦ)

    # return
    (sinΦ, sinΨ)
end

function sinΦsinΨ(M, r, a, θ, α, β)
    # value cache
    Σ₀ = Σ(r, a, θ)
    sinθ = sin(θ)
    A₀ = A(M, r, a, θ)
    Δ₀ = Δ(M, r, a)
    ω = 2.0 * a * r / A₀
    S₀ = S(θ, α, ω)

    sinΦsinΨ(Σ₀, sinθ, A₀, Δ₀, S₀, r, a, α, β)
end

"""
$(TYPEDSIGNATURES)

Calculates conserved quantities
- angular momentum ``L``
- Carter parameter ``Q``

for a photon described with position described by `x` in a Kerr spacetime given by
`p`.

From Fanton et al. (1997), eq. (69):

```math
L = \\frac{\\Upsilon_1}{\\Upsilon_2},
```

where

```math
\\begin{align*}
\\Upsilon_1 &= \\sin \\theta \\sin \\Phi \\sin \\Theta, \\\\
\\Upsilon_2 &= \\frac{\\Sigma \\sqrt{\\Delta}}{A} + \\omega \\Upsilon_1,
\\omega = \\frac{2 a r}{A},
\\end{align*}
```

taken from eq. (72) and eq. (73).

From Fanton et al. (1997), eq. (70):

```math
Q =
\\frac{P^2}{\\Delta}
- \\left( \\lambda - a \\right)^2
- \\frac{\\Sigma^2}{A} \\left( \\frac{\\cos \\Phi}{\\Upsilon_2} \\right)^2,
```

and

```math
P = \\left( r^2 + a^2 - a L \\right),
```

taken from eq. (71).
"""
function LQ(M, r, a, θ, sinΦ, sinΨ)
    sinθ = sin(θ)
    Σ₀ = Σ(r, a, θ)
    A₀ = A(M, r, a, θ)
    Δ₀ = Δ(M, r, a)
    ω = 2.0 * a * r / A₀

    Υ₁ = sinθ * sinΦ * sinΨ
    Υ₂ = (Σ₀ * √Δ₀ / A₀) + ω * Υ₁

    Υ₁ = sinθ * sinΦ * sinΨ
    Υ₂ = (Σ₀ * √Δ₀ / A₀) + ω * Υ₁

    L = Υ₁ / Υ₂
    P = (r^2 + a^2 - (a * L))
    Q = (P^2 / Δ₀) - (L - a)^2 - ((Σ₀^2 / A₀) * (cos(asin(sinΨ)) / Υ₂)^2)
    L, Q
end

"""
$(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.3):

```math
\\Sigma = r^2 + a^2 \\cos^2( \\theta ).
```
"""
Σ(r, a, θ) = r^2 + (a * cos(θ))^2

"""
$(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.3):

```math
\\Delta = r^2 - 2 M r + a^2.
```
"""
Δ(M, r, a) = r^2 - 2 * M * r + a^2

"""
$(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.3):

```math
A = (r^2 + a^2)^2 - a^2 \\Delta \\sin^2 ( \\theta ).
```
"""
A(M, r, a, θ) = (r^2 + a^2)^2 - a^2 * Δ(M, r, a) * sin(θ)^2

"""
$(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.21):

```math
Z_1 =
1 + \\sqrt[3]{1 - \\frac{a^2}{M^2}}
\\left[
    \\sqrt[3]{1 + \\frac{a}{M}} + \\sqrt[3]{1 - \\frac{a}{M}}
\\right].
```
"""
Z₁(M, a) = 1 + ∛(1 - (a / M)^2) * (∛(1 + (a / M)) + ∛(1 - (a / M)))

"""
$(TYPEDSIGNATURES)

From Bardeen et al. (1972) eq. (2.21):

```math
Z_2 = \\sqrt{\\frac{3a^2}{M^2} + Z_1^2}.
```
"""
Z₂(M, a) = √(3(a / M)^2 + Z₁(M, a)^2)

end # module


using Test
using ConvDiffMIPDECO
using jInv.ForwardShare
using jInv.LinearSolvers
using KrylovMethods
using jInv.Mesh
using MUMPSjInv
using Statistics
using LinearAlgebra

domain = [0 1. 0 3. 0 2.]

N = [8;16;32;]
v = X -> [ones(size(X,1)) zeros(size(X,1)) zeros(size(X,1))]
bc = (:dir,:neu,:neu,:neu,:neu,:neu)

u  = X -> cos.(pi*X[:,1]).*cos.(pi*X[:,2]).*cos.(pi*X[:,3])
ux = X -> -pi*sin.(pi*X[:,1]).*cos.(pi*X[:,2]).*cos.(pi*X[:,3])
uy = X -> -pi*cos.(pi*X[:,1]).*sin.(pi*X[:,2]).*cos.(pi*X[:,3])
uz = X -> -pi*cos.(pi*X[:,1]).*cos.(pi*X[:,2]).*sin.(pi*X[:,3])
sig = .01;

f  = X->  3*sig*pi^2*u(X)+ux(X)

err = zeros(length(N))

times = zeros(length(N),3)
pFor=1;
for j=1:2
	for k=1:length(N)
		Ainvs = (getJuliaSolver(), getMUMPSsolver())

		nk  = [N[k];N[k];N[k]]
		Mk  = getRegularMesh(domain,nk)
		xck = getCellCenteredGrid(Mk)
		xnk = getNodalGrid(Mk)
		pFork = getConvDiffFEMParam(Mk,sig=sig,gd=u,bc=bc,Ainv=Ainvs[j])

		dobsk,pFork = getData(f(xck), pFork)
		pFork.Fields .-= mean(pFork.Fields)
		utruek = u(xnk)
		utruek .-= mean(utruek)

		 times[k,j]= @elapsed begin
			 err[k] = norm(pFork.Fields - utruek,Inf)/norm(utruek,Inf)
		 end
		println(err[k])
	end
	@test all(abs.(diff(log2.(err))).>1.4)
end