module LargeMatrixEquations

export kpik, rksm, lp_lradi, lp_para

"""
    `kpik(A,B,E=1;<keyword arguments>)`

Julia code for K-PIK (Krylov-plus-inverted-Krylov)

Based on kpik.m avalible from V. Simoncini's website
(http://www.dm.unibo.it/~simoncin/software.html)
and includes many of the same comments and much of the same description.

Approximately solve

     A X E + E X A' + BB' = 0

by means of the extended Krylov subspace method

ARGUMENTS

    'A' : coeff matrix, A < 0

    'B' : factor of rhs,   nxk matrix with k << n

NAMED ARGUMENTS

    'E' : coeff matrix, spd, Defult: 1

    'LE' : lower triang factor of coeff matrix, Defult: 1 Note: This is an
    optional argument, if not provided it will be set to cholfact(E)[:L]

    'm' : max space dimension, Defult: 100

    'tol' : stopping tolerance based on the backwards error,
    with stopping criterion

      ||LE\\A X LE  + LE' X A'/LE'-LE\\BB'/LE'||
    -----------------------------------------------   < tol
    ||LE\\BB'/LE'|| + ||E^{-1}|| ||A|| ||LE'X LE ||

    computed in a cheap manner. Defult: 1e-9

Output

    'Z' : solution factor s.t.  X = Z Z'

    'er2' : history of scaled residual, as above

Comments

    1. The projected solution is computed at each iteration
    As an alternative, a periodic computation could be considered.

    2. This code performs a factorization of A. As an alternative,
    iterative solves could be considered.

    3. To solve

     A X + X A' + BB' = 0

     Use kpik(A,B) as E is set to 1 by Defult

If you use this code, please cite the following article:

    V. Simoncini
    A new iterative method for solving large-scale Lyapunov matrix equations,
    SIAM J.  Scient. Computing, v.29, n.3 (2007), pp. 1268-1288.
"""
function kpik(A,B,E=1;LE=1,m=100,tol=1e-9,tolY=1e-12,infoV=true)

  @assert(isdefined(:vecnorm),"Your julia version is too old. vecnorm() not defined")

  infoV && tic()

  #### Check if we solve the general lyapunov equation and user did not provide LE
  if E != 1 && LE == 1
    #### If this is the case, calculate LE
    if issparse(E)
      LE = try
            sparse(cholfact(E,perm=1:size(E,1))[:L])
           catch
            error("E must be SPD")
          end
    else
      LE = try
            cholfact(E)[:L]
          catch
            error("E must be SPD")
          end
    end
  end

  const rhs=LE\B
  const nrmb=vecnorm(rhs)^2
  const nrma=vecnorm(A)
  const sqrt2=sqrt(2)
  er2=zeros(m,1)

  const n,sh=size(rhs)

  Y=[]
  odds=[]
  er2=[]

  if E == 1
    const condestE=1
    const singE=condestE/vecnorm(eye(n))
  elseif (vecnorm(E-speye(n))>1e-14)
    const condestE=cond(full(E))
    const singE=condestE/vecnorm(E)
  else
    const singE=1
  end

  if norm(A-A',1)<1e-14
     if issparse(A)
       cfA = cholfact(-A,perm=1:size(A,1))
       const LA = sparse(cfA[:L])
       const UA = -LA'
     else
       const UA = chol(-A)
       const LA = -UA'
     end
     infoV && println("A sym. Completed Chol factorization\n")
     const k_max =2
   else
     luA=lufact(full(A),Val{false})
     const LA = luA[:L]
     const UA = luA[:U]

     infoV && println("A nonsym. Completed LU factorization\n")
     const k_max = m
   end

   const s=2*sh
   #rhs1=LE'*(UA\(LA\(LE*rhs)))

   # Orthogonalize [B,A^{-1}B] with an economy-size QR
   rr = [ rhs LE'*(UA\(LA\(LE*rhs))) ]

   # Julia qr decomposition is always "economy size"
   U,beta=qr(rr)
   U = U[1:n,1:s]

   const ibeta=inv(beta[1:s,1:s])
   beta = beta[1:sh,1:sh]
   const beta2=beta*beta'
   H=zeros((m+1)*s,m*s)
   T=zeros((m+1)*s,m*s)
   L=zeros((m+1)*s,m*s)
   local js, j, rho
   for j=1:m
     jms=(j-1)*s+1
     j1s=(j+1)*s
     js=j*s
     js1=js+1
     jsh=(j-1)*s+sh

     # Expand the basis
     # multiply by A
     Up = zeros(n,s)
     Up[1:n,1:sh] = LE\(A*(LE'\U[:,jms:jsh]))
     # solve with A

     Up[1:n,sh+1:s] = LE'*(UA\(LA\(LE*U[1:n,jsh+1:js])))

     # orthogonalize the new basis block wrt all the previous ones by modified gram
     for l=1:2
        k_min=max(1,j-k_max)
        for kk=k_min:j
            k1=(kk-1)*s+1
            k2=kk*s
            coef= U[1:n,k1:k2]'*Up
            H[k1:k2,jms:js] = H[k1:k2,jms:js]+ coef
            Up = Up - U[:,k1:k2]*coef
        end
      end

      if j<=m
        Up,H[js1:j1s,jms:js] = qr(Up)
        hinv=inv(H[js1:j1s,jms:js])
      end

      ###############################################################
      # Recover the columns of T=U'*A*U (projection of A onto the space) from
      # the colums of H.
      # REMARK: we need T as coefficient matrix of the projected problem.
      Iden=speye(js+s)

      if (j==1)
        L[1:j*s+sh,(j-1)*sh+1:j*sh] = [H[1:s+sh,1:sh]/ibeta[1:sh,1:sh] eye(s+sh,sh)/ibeta[1:sh,1:sh]]*ibeta[1:s,sh+1:s]
      else
        L[1:j*s+s,(j-1)*sh+1:j*sh] = L[1:j*s+s,(j-1)*sh+1:j*sh] + H[1:j*s+s,jms:jms-1+sh]*rho
      end

      odds = [odds; jms:(jms-1+sh)]   # store the odd block columns
      evens = 1:js
      flag = trues(size(evens))
      flag[odds] = false
      evens = evens[flag]
      T[1:js+s,odds]=H[1:js+s,odds]   #odd columns

      T[1:js+sh,evens]=L[1:js+sh,1:j*sh]   #even columns
      L[1:j*s+s,j*sh+1:(j+1)*sh] = ( Iden[1:j*s+s,(js-sh+1):js]- T[1:js+s,1:js]*H[1:js,js-sh+1:js])*hinv[sh+1:s,sh+1:s]
      rho = hinv[1:sh,1:sh]\hinv[1:sh,sh+1:s]

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

      # Solve the projected problem by Bartels-Stewart
      # Do "type lyap" from command window if interested

      Y = lyap((T[1:js,1:js]),eye(j*s,sh)*beta2*eye(j*s,sh)')

      # safeguard to preserve symmetry
      Y = (Y+Y')/2

      # Compute the residual norm. See the article by Valeria

      cc = [H[js1:j1s,js-s+1:js-sh] L[js1:j1s,(j-1)*sh+1:j*sh]]

      nrmx = vecnorm(Y)

      er2=[er2;sqrt2*vecnorm(cc*Y[js-s+1:js,:])/(nrmb+singE*nrma*nrmx)]

      infoV && println("KPIK It: $j -- Current Backwards Error: $(er2[j])")

      (er2[j]<tol) ? break : U = [U Up]

    end

    # Done
    # reduce solution rank if needed
    sY,uY=eig(Y)
    id=sortperm(sY)
    sY=sort(sY)

    sY=flipdim(sY,1)
    uY=uY[:,id[end:-1:1]]
    is = 0
    for ii in 1:size(sY)[1]
      if abs(sY[ii])>tolY
        is = is+1
      end
    end

    Y0 = uY[:,1:is]*diagm(sqrt(sY[1:is]))
    Z = LE'\(U[1:n,1:js]*Y0)
    er2=er2[1:j]

    if infoV
      println("its  Back. Error            space dim. CPU Time")
      println("$j    $(er2[j])  $js          $(toq())")
    end

    return Z, er2

end

"""
  `rksm(A,B,E=1;LE=1,s1=NaN,emax=NaN,m=100,tol=1e-9,ch=true,tolY=1e-12,infoV=true)`

Based on rksm.m on Valeria Simoncini's website
(http://www.dm.unibo.it/~simoncin/software.html)
and includes many of the same comments and much of the same description

Approximately Solve

    A X E + E X A' + BB' = 0

by the Rational Krylov subspace method
(Galerkin condition onto the Rational Krylov subspace)
This code performs system solves with (A-s E)

Input:

    A, E  coeff. matrices. A<0,  E is spd

    LE   Cholesky lower factor of E

    B     rhs factor

    m       max space dimension allowed

    tol     stopping tolerance based on the backwards error,
            with stopping criterion

                ||LE\A X LE  + LE' X A'/LE'-LE\BB'/LE'||
            ----------------------------------------------  < tol
            ||LE\BB'/LE'|| + ||E^{-1}|| ||A|| ||LE'X LE ||

            computed in a cheap manner

    s1,smax   estimates for real spectral interval associated with
              field of values of (A,E)

    ch      ch=true  complex poles  ch=false real poles

    tolY    truncation tolerance for final solution, e.g., tolY=1e-12

    infoV   If true, print out backwards error every iteration and a summary
            once the program finishes.

Output:

    Z    factor of approximate solution  X = Z Z'
         If ch=true,  Z may be complex

    nrmrestot  history of residuals

Hints:

    1. Before the call permute entries of A, E and B so as to
    limit fill-in in the system solves

    2. Provide "comfortable" (loose bounds) estimates s1, emax

When using this code, please cite the following reference:

    V. Druskin and V. Simoncini,
    Adaptive rational Krylov subspaces for large-scale dynamical systems
    Systems & Control Letters, 60 (2011), pp. 546-560.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
#FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
#COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
#IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
#CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
function rksm(A,B,E=1;LE=1,s1=NaN,emax=NaN,m=100,tol=1e-9,ch=false,tolY=1e-12,
                                                          infoV=true)

@assert(isdefined(:vecnorm),"Your julia version is too old. vecnorm() not defined")

infoV && tic()

const symm = norm(A-E-(A-E)',1) < 1e-14
# If not symmetric, we may have complex poles.
# Therefore we must use complex types
symm ? (const typ = Float64) : (const typ = Complex{Float64})

const n=size(A,1)
const p=size(B,2)

E==1 ? ee = eye(n) : ee = E
# If user does not give smallest generalized eigenvalue, calculate it
if (isnan(s1))
  (symm) && (s1 = eigmax(inv(full(ee))*full(-A)))
  (!symm) && (s1 = maximum(abs(eigvals(inv(full(ee))*full(-A)))))
end

# If user does not give largest generalized eigenvalue, calculate it
if (isnan(emax))
  (symm) && (emax = eigmin(inv(full(ee))*full(-A)))
  (!symm) && (emax = minimum(abs(eigvals(inv(full(ee))*full(-A)))))
end

if (E != 1 && LE == 1)
  if issparse(E)
    LE = try
          LE = sparse(cholfact(E,perm=1:size(E,1))[:L])
        catch
          error("E must be Symmetric Positive Definite")
        end
  else
    LE = try
          cholfact(E)[:L]
        catch
          error("E must be Symmetric Positive Definite")
        end
 end
end

const Iden=eye(p)
const O=0*Iden
const uno=ones(1,p)
const Lres = convert(Array{typ,2},LE\B)

V,irr=qr(Lres)
rr=inv(irr)
const nrmb=vecnorm(inv(rr))^2
const beta=V'*Lres
const beta2=beta*beta'
s=typ[]

VV=zeros(n,p*(m+2))
(typ == Float64) && (VV=VV*1.)
(typ == Complex{Float64}) && (VV=VV*1.*im)
VV[1:n,1:p]=V

H=zeros(p*(m+2),p*(m+1))
(typ == Float64) && (H=H*1.)
(typ == Complex{Float64}) && (H=H*1.*im)

nrmrestot=[]
const nrma=vecnorm(A)

if E == 1
  const condestE=1
  const singE=condestE/vecnorm(eye(n))
elseif (vecnorm(E-speye(n))>1e-14)
  const condestE=cond(full(E))
  const singE=condestE/vecnorm(full(E))
else
  const singE=1
end

newAv=LE\(A*(LE'\V))

K=V'*newAv
push!(s,s1)
eH=eig(K)
eHpoints = sort([s1,emax])
snew=newpolei(eHpoints,eH[1],s1*uno',typ)

if real(snew)<0
   snew=-real(snew)+im*imag(snew)
end
push!(s,snew)

# additional steps
cmplxflag=false

i=0

local Y, j1s
while i < m
  local i1, j1, js
  i=i+1

  paired=0
  while (paired==0)

    i1=i+1
    w=LE*V
    #E==1 ? wrk = (A-snew*diagm(ones(n),0))\w : wrk = (A-snew*E)\w
    wrk = (A-snew*speye(n)*E)\w
    wrk= LE'*wrk

    # Gram-Schmidt step
    jms=(i-1)*p+1
    j1s=(i+1)*p
    js=i*p
    js1=js+1
    for it=1:2
      for kk=1:i
        k1=(kk-1)*p+1
        k2=kk*p
        gamma=VV[1:n,k1:k2]'*wrk
        H[k1:k2,jms:js] = H[k1:k2,jms:js]+ gamma
        wrk = wrk - VV[:,k1:k2]*gamma
      end
    end
    V, H[js1:j1s,jms:js]=qr(wrk)
    if (cmplxflag)
      snew=conj(snew)
      push!(s,snew)
      cmplxflag=false
      newAv=LE\(A*(LE'\V))
      g = VV[1:n,1:js]'*newAv
      g1 = g
      g2 = V'*(LE\(A*(LE'\VV[1:n,1:js])))
      g3 = V'*(LE\(A*(LE'\V)))
      K = [K g1; g2 g3]
      VV[1:n,js+1:j1s]=V
      i=i+1
    else
      paired=1
    end
  end

  ih1=i1
  ih=i
  newAv=LE\(A*(LE'\V))
  g = VV[1:n,1:js]'*newAv

  if (symm)
    K=(K+K')/2
  end

  rhs2=speye(ih*p,p)*beta2*speye(ih*p,p)'
  Y = lyap(K,rhs2)
  nrmx = vecnorm(Y)

  # computed residual   (exact, in exact arithmetic)
  u1=newAv-VV[1:n,1:js]*g
  d=-VV[1:n,1:js]*(Y*(H[1:ih*p,1:ih*p]'\[zeros(p*(ih-1),p);Iden])*H[p*ih+1:p*ih1,p*ih-p+1:p*ih]')

  U=[-V*s[end]  d u1 ]
  extra,rr=qr(full(U))

  nrmres=vecnorm(rr*([O Iden O; Iden O Iden; O Iden O ])*rr')/(nrmb+singE*nrma*nrmx)
  push!(nrmrestot, nrmres)

  infoV && println("RKSM It: $i -- Current backwards error: $nrmres")
  (nrmres<tol) && break
  # New poles and zeros
  eH=eig(K)[1]
  symm ? sort!(eH) : eH[sortperm(abs(eH))]

  eHorig=eH

   if (ch)
     # Complex poles. Compute set for next complex pole of r_m

      if (countnz(imag(eH))>0 && length(eH)>2) # Roots lambdas come from convex hull too
        eH=full([eH;-emax])
        eH=convhull(eH)
        ieH=length(eH)
        missing=ih*p-ieH
        while missing>0                        # include enough points from the border
          neweH=(eH[1:ieH-1]+eH[2:ieH])/2
          missing=ih*p-length(eH)
          eH=[eH;neweH]
        end
        eHpoints=-eH
        eH=eHorig
      else     # if all real eigs, no convex hull possible
        eHpoints = sort([s1; emax.';-real(eH)])
      end

   else   # Real poles s from real set. Compute complex roots of r_m via Ritz convex hull
       if (countnz(imag(eH))>0 && length(eH)>2)    # Roots lambdas come from convex hull too
         eH=full([eH;-s1;-emax.'])
         eH=convhull(eH)
         ieH=length(eH)
         missing=ih*p-ieH
         while missing>0 # include enough points from the border
           neweH=(eH[1:ieH-1]+eH[2:ieH])/2
           eH=[eH;neweH]
           missing=ih*p-length(eH)
         end
         eH=eH[1:ih*p]
       end
        eHpoints = sort([s1; emax.';-real(eH)])
        eH=eHorig
   end

   gs=kron(s[2:i+1].',uno)'

   snew = newpolei(eHpoints,eH,gs,typ)
   (real(snew)<0) && (snew=-real(snew)+im*imag(snew)); #safeguard strategy

   # If pole is complex, include its conjugate
   (imag(snew) !=0) && (cmplxflag=true)

   push!(s,snew)

   g1 = g
   g2 = V'*(LE\(A*(LE'\VV[1:n,1:js])))
   g3 = V'*(LE\(A*(LE'\V)))
   K = [K g1; g2 g3]
   VV[1:n,js+1:j1s]=V

end

# Done
# Reduce rank of solution, if needed
sY,uY=eig(Y)
symm ? id = sortperm(sY) : id = sortperm(abs(sY))

sY=sY[id]
sY=flipdim(sY,1)
uY=uY[:,id[end:-1:1]]
is = 0
for ii in 1:size(sY)[1]
  if abs(sY[ii])>tolY
    is = is+1
  end
end

Y0 = uY[:,1:is]*diagm(sqrt(sY[1:is]))
Z = LE'\(VV[:,1:size(Y0,1)]*Y0)

if infoV
  RKStotal_time=toq()
  println("Space dim $j1s  Solution rank $is  time $RKStotal_time")
end

return Z,nrmrestot

end

#################################################
#
#  ratfun and newpolei function definitions. These
#  are auxiliary routines for rksm based on Simoncini's
#  corresponding routines.
#
##################################################

function ratfun(x,eH,s,typ)
r = Array(typ,length(x))

for j=1:length(x)
  xj = x[j]*ones(length(s))
  r[j]=abs(prod( (xj-s)./(xj-eH) ))
end
return r
end

##################################
function newpolei(eHpoints,eH,s,typ)
snew=Array(typ,length(eHpoints)-1)

for j=1:length(eHpoints)-1

    sval = linspace(real(eHpoints[j]),real(eHpoints[j+1]),200)
    (typ == Complex{Float64}) && (sval = sval + im*linspace(imag(eHpoints[j]),imag(eHpoints[j+1]),200))

    jx = indmax(abs(ratfun(sval,eH,s,typ)))

    snew[j]=sval[jx]
end
jx=indmax(abs(ratfun(snew,eH,s,typ)))
snew=snew[jx]
return snew
end

"""
  `convexhull(pnts)`
Based on code from (https://github.com/intdxdt/convexhull.jl)
Function to compute 2-D convex hull
"""
function convhull(pnts)
    const T = eltype(pnts) #get point type
    const N = length(pnts) #number of pnts
    # sort the points lexicographically.
    # copy points into mutable container

    #trivial case 0 or 1 point
    length(pnts) <= 1 && (return pnts)

    #sort points lexicographically
    sort!(pnts, lt=lt2d)

    #function to orient boundry using 2D cross product of OA and OB vectors,
    #i.e. z-component of their 3D cross product.
    #Returns a positive value, if OAB makes a counter-clockwise turn,
    #negative for clockwise turn, and zero if the points are collinear.
    orient(b, pnt) = (real(b[end]-b[end-1]) * imag(pnt - b[end-1])) -
                      (imag(b[end-1] - b[end-1]) * real(pnt - b[end-1]))

    #build lower hull
    lower = T[]
    for i in 1:N
        pnt = pnts[i]
        while length(lower) >= 2 && orient(lower, pnt) <= 0
            pop!(lower)
        end
        push!(lower, pnt)
    end

    #build upper hull
    upper = T[]
    for i in N:-1:1
        pnt = pnts[i]
        while length(upper) >= 2 && orient(upper, pnt) <= 0
            pop!(upper)
        end
        push!(upper, pnt)
    end

    # omit end points :
    #   - reapeated at the beginning of lower & upper
    length(upper) > 0 && pop!(upper)
    length(lower) > 0 && pop!(lower)

    # two reapeated points
    if (length(upper)==1 && length(lower)==1 && isequal(upper, lower))
      pop!(upper)
    end

    # concat lower and upper hull.
    hull =[lower; upper]

    #close ring
    length(hull) > 1 && push!(hull, hull[1])

    return hull
end

#less than comparator
function lt2d(a, b)
    dx, dy =  real(a - b), imag(a - b)
    dx != 0 &&  (return <(dx, 0))
    return <(dy, 0)
end

############################################################
#
#     Please note: The following functions lp_lradi, lp_nrmu, lp_para,
#                  lp_arn, lp_mnmx are based on the functions
#                  with the same name in LYAPACK 1.0 (lp_arn
#                  is based on lp_arn_m.m and lp_arn_p.m). The
#                  structure of these files is closely followed
#                  and even includes many of the same comments.
#
#############################################################

"""
Based on lp_lradi.m from LYAPACK 1.0

Low Rank Cholesky Factor Alternating Direction Implicit (LRCF-ADI) Method for
solving the stable Lyapunov equation:

  For tp = :B

    F * X + X * F' = -B * B',

  for tp = :C

    F' * X + X * F = -B' * B,

where F = A-Bf*Kf'.

The routine works in two modes (depending on the choice of zk)

    For zk = :Z, this routine deliveres a low rank factor Z,
    such that Z * Z' approximates X

    for zk = :K, it only generates the product K_out = Z * Z'* K
    without forming Z itself.

Input:

    A         matrix A
    B         matrix B (n-x-m if tp=:B or m-x-n if tp=:C,
              where  m << n !)
    p         l-vector with ADI shift parameters. Complex parameters must
              appear as conjugate complex pairs in consecutive order.
              If more than length(p) iterations need to be done, the
              parameters p[i] are reused in a cyclic manner. The user can calculate
              these values using lp_para.
    Bf        "feedback" matrix Bf
              Named argument with defult value [], so if not provided by user,
              the code solves
              A * X + X * A' = -B * B' or A' * X + X * A = -B'* B
              depending on choice of tp.
    Kf        "feedback" matrix Kf
              amed argument with defult value [], so if not provided by user,
              the code solves
              A * X + X * A' = -B * B' or A' * X + X * A = -B' * B
              depending on choice of tp.
    tp        (= :B or :C) named argument that determines the type of Lyapunov equation. Defult value is :B
    rc        (= :R or :C) named argument that determines whether the low rank factor Z
              must be real (:R) or if complex factors are allowed (:C)
              If p contains complex parameters, then the low rank factor
              Z is complex, too, although Z * Z' is real. A real low rank
              factor is determined from the complex data, if rc = :R.
              However, this requires some additional computation.
              If zk = :K, rc is ignored. Defult value is :C.
    K         named argument n-x-r matrix (where r should be small: r << n !).
              Defult value is [].
    max_it    stopping criterion: maximal number of LRCP-ADI steps
              Set to [] or +Inf (default) to avoid this criterion. Named argument
              With defult value of 100.
    min_res   stopping criterion: minimal relative residual norm. The
              iteration is stopped when res(i+1) <= min_res (See Remarks).
              Set min_res = [] or 0 (default) to avoid this criterion. Note:
              If min_res<=0 and with_rs=:N, the (often expensive) calculation
              of the residual norm is avoided, but, of course, res is not
              provided on exit. Named argument with defult value of 1e-6.
    with_rs   (= :S or :N) if with_rs = "S", the iteration is stopped,
              when the routine detects a stagnation of the residual norms,
              which is most likely the case, when roundoff errors rather
              than the approximation error start to dominante the residual
              norm. This implies that the residual norms are computed (which
              can be expensive). This criterion works quite well in practice,
              but is not absolutely sure. Use with_rs = "S" only if you
              want to compute the Lyapunov solution as accurate as possible
              (for a given machine precision).
              If with_rs = :N, this criterion is not used. Named argument with
              defult value :S.
    min_in    stopping criterion: This value limits the "minimal increase"
              in the matrix Z by the "new" columns . The iteration is
              terminated as soon as

                || Z_nc ||F^2 < min_in * || Z_new ||F^2

              holds for a certain number of consecutive iteration steps.
              Here, Z_nc are the currently computed "new" colums, which
              appended to the old iterate Z_old deliver the new iterate
              Z_new  = [ Z_old  Z_nc ]. Set min_in = 0 to avoid it.
              Default value is eps, the machine precision. min_in = []
              has the same effect. Note that this is an "adaptive"
              stopping criterion which does not require the
              (often expensive) computation of the residual norm. Named argument
              with defult value 0.
    infoV      Bool: Return infoVrmation given during the
              iteration. Named argument with default value true.

Output:

    Z         Z * Z' is a low rank approximation of X
              (Note that Z can be complex if rc=:C!)
    flag      the criterion which had stopped the iteration:
               = "I": max_it,
               = "R": min_res,
               = "S": with_rs,
               = "N": min_in,
    res       the relative residual norms attained in the course of
              the iterations (res[i+1] is the norm after the i-th step
              of the iteration!). See note in min_res.

Remarks:

    1. The eigenvalues of F must have negative real parts.

    2. The values in res correspond to the following "relative" norms

      tp = :B

        res(i+1) = ||F * Z_i * Z_i' + Z_i * Z_i' * F' + B * B'||F/|| B * B'||F

      tp = :C

        res(i+1) = ||F' * Z_i * Z_i' + Z_i * Z_i' * F + B' * B||F/||B' * B||F

    3. Note that all stopping criteria are checked only after a step
    with a real parameter or a "double step" with a pair of conjugate
    complex parameters. This ensures that Z*Z' is real, even if Z is
    not.

References:

    The algorithm is a slight modifivation of that proposed in:

    1. J.Li, F.Wang, and J.White.
    An efficient Lyapunov equation based approach for generating
    reduced-order models of interconnect.
    Proceedings of the 36th IEEE/ACM Design Automation Conference,
    New Orleans, LA, 1999.

    Another (though more expensive) low rank algorithm is proposed in:

    2. T. Penzl.
    A cyclic low rank Smith method for large sparse Lyapunov equations.
    To appear in SIAM Sci. Comp.

    See also:

    3. P. Benner, J. Li, and T. Penzl
    Numerical solution of large Lyapunov equations, Riccati equations,
    and linear-quadratic optimal control problems.
    In preparation.

    4. T. Penzl.
    LYAPACK (Users' Guide - Version 1.0).
    1999.

Internal remarks:

    Input data not completely checked!

    The procedure to generate real factors in case of complex parameters is
    different from that in [3]!

    The matrices SMi (i = 1:length(p)) for the "Sherman-Morrison trick"
    (only used if Bf and Kf are nonzero) are computed a priori, which
    is good in view of computation if parameters p(i) are used cyclically,
    but may be sometimes bad in view of memory demand, in particular, when
    length(p) is large.

    The stopping criterion related to the input parameter with_rs
    corresponds to the stagnation of the residual curve caused by
    round-off errors. Its performance depends on the constants stcf and
    min_rs. The iteration is stopped, as soon as

    (ra-rb) * (i-stcf+1) / ((r(1)-ra) * stcf) < min_rs

    and r(1)-ra > 0 and i >= stcf hold for stcf consecutive iteration
    steps. Here res(i+1) is the residual norm after the i-th LRCF-ADI step,
    r(i+1) = log(res(i+1)), ra = min(r(1:i-stcf+1)), rb = min(r(i-stcf+2:i+1)).

    stcf is also the number of consecutive steps, for which the criterion
    w.r.t. min_in must be fulfilled.
"""
function lp_lradi(A,B;p=[],Bf=[],Kf=[],K=[],max_it=100,tp=:B,rc=:C,
  min_res=1e-6,with_rs=:S,min_in=0,infoV=true)

tic()
isempty(p) && (p=lp_para(A))
const stcf = 10
const min_rs = .1

(tp!=:B && tp!=:C) && error("tp must be either :B or :C.")
(rc!=:R && rc!=:C) && error("rc must be either :R or :C.")

const tpB = (tp == :B)

const with_min_rs = (with_rs==:S)
const with_norm = (min_res>0)||with_min_rs
const make_real = (rc==:R)

const with_min_in = min_in>0

const with_BK = !isempty(Bf)

const l = length(p)

if tpB
  const n,m = size(B)
else
  const m,n = size(B)
end

const Ide = speye(size(A,1))
LP_L = Array(Complex{Float64},size(A)...,l)
LP_U = Array(Complex{Float64},size(A)...,l)
for i = 1:l
  LP_L[:,:,i],LP_U[:,:,i],~=lu(full(A+p[i]*Ide),Val{false})
end

if with_BK
  SM=Array(Complex{Float64},size(Bf)...,l)
  Im = speye(size(Bf,2))
  if tpB
    # SMi = TM*inv(I-Kf'*TM),
    # where TM = inv(F+p(i)*I)*Bf
    # (These are the columns of the LR terms for the
    # rank corrections of the "inverse"
    # in the Sherman-Morrison formulae.)

    for i = 1:l
      TM = LP_U[:,:,i]\(LP_L[:,:,i]\Bf)
      SM[:,:,i] = TM/(Im-Kf'*TM)
    end
  else  # tp==:C
    # SMi = TM*inv(I-Bf'*TM),
    # where TM = inv(F.'+p(i)*I)*Kf
    # (These are the columns of the LR terms for the
    # rank corrections of the "inverse"
    # in the Sherman-Morrison formulae.)
    for i = 1:l
      TM = LP_L[:,:,i].'\(LP_U[:,:,i].'\Kf)
      SM[:,:,i] = TM/(Im-Bf'*TM)
    end
  end
end

# Initialize QR factorization for norm computation
if with_norm
  res0,nrmQ,nrmR,nrmbs = lp_nrmu(A,B,Bf,Kf,tpB,[],[],[],[])
  res = [1.]
  res_log = [log(res0)]
  # Vector containing log of residual norms
  # corresponds to r(:) in prolog.
end

flag = "I"
if with_min_in
  nrmF_Z_2 = 0
  nrmF_rec = +Inf*ones(stcf,1)
  # Current squared Frobenius norm of Z
  # Records the values of
  # ||Z_nc||_F^2 / ||Z_new||_F^2 (see prolog)
  # for the last stcf iteration steps.
end

i_p = 1;                   # Pointer to i-th entry in p(:)
is_compl = imag(p[1])!=0;  # is_compl = (current parameter is complex.)
is_first = true;           # is_first = (current parameter is the first
                           #            of a pair, PROVIDED THAT is_compl.)

local V, Z
for i = 1:max_it       # The iteration itself

  if i==1

    if tpB
      # V = last columns of Cholesky factor Z
      # Initialize:
      # V := sqrt(-2*real(p[1]))*inv(F+p[1]*I)*B

      V=LP_U[:,:,1]\(LP_L[:,:,1]\B)
      with_BK && (V = V+SM[:,:,1]*(Kf'*V))
      V = sqrt(-2*real(p[1]))*V

    else #( tp = :C )
      # Initialize:
      # V := sqrt(-2*real(p[1]))*inv(F.'+p[1]*I)*B'

      V = LP_L[:,:,1].'\(LP_U[:,:,1].'\B')
      with_BK && (V = V+SM[:,:,1]*(Bf'*V))
      V = sqrt(-2*real(p[1]))*V
    end

    Z = V;             # Note:  Z*Z' = current ADI iterate

  else # i > 1

    p_old = p[i_p]

    i_p = i_p+1
    i_p>l && (i_p = 1)    # update current parameter index

    if is_compl && is_first
      is_first = false
      if i_p==1
        error("Parameter sequence ends in the 'middle' of a complex pair!")
      end
      if p[i_p]!=conj(p_old)
        error("Parameters p[i] must be either real or pairs of conjugate complex numbers.")
      end
    else
      is_compl = imag(p[i_p])!=0
      is_first = true
    end

    if tpB
      # Evaluate
      #   V := sqrt(real(p[i_p])/real(p_old))*...
      #        (V-(p[i_p]+conj(p_old))*inv(F+p[i_p]*I)*V)

      TM = LP_U[:,:,i_p]\(LP_L[:,:,i_p]\V)
      with_BK && (TM = TM+SM[:,:,i_p]*(Kf'*TM))
      TM = V-(p[i_p]+conj(p_old))*TM
      V = sqrt(real(p[i_p])/real(p_old))*TM

    else   #  tp==:C

      # Evaluate
      #   V := sqrt(real(p[i_p])/real(p_old))*...
      #        (V-(p[i_p]+conj(p_old))* inv(F.'+p[i_p]*I)*V)

      TM = LP_L[:,:,i_p].'\(LP_U[:,:,i_p].'\V)
      with_BK && (TM = TM+SM[:,:,i_p]*(Bf'*TM))
      TM = V-(p[i_p]+conj(p_old))*TM
      V = sqrt(real(p[i_p])/real(p_old))*TM
    end

    # Form new iterate Z.

    !is_compl && (V = real(V))

    Z = [Z V]

    # Make last 2*m columns real.
    if make_real && is_compl && !is_first
      for j = (i-1)*m+1:i*m
        U1,S1,V1 = svd([real(Z[:,j-m]) real(Z[:,j]) imag(Z[:,j-m]) imag(Z[:,j])])
        S1 = diagm(S1)
        U2,S2,V2 = svd(V1[1:2,1:2]'*V1[1:2,1:2]+V1[3:4,1:2]'*V1[3:4,1:2],thin=false)
        TMP = U1[:,1:2]*S1[1:2,1:2]*U2*diagm(sqrt(S2))
        Z[:,j-m] = TMP[:,1]
        Z[:,j] = TMP[:,2]
      end

    end

  end

  if with_norm
    # Compute residual norm

    resnrm,nrmQ,nrmR,nrmbs = lp_nrmu( A, B, Bf, Kf, tpB, V, nrmQ, nrmR, nrmbs)
    push!(res_log, log(resnrm))
    akt_res = resnrm/res0
    push!(res, akt_res)

    infoV && println("LRCF-ADI step $i -- norm. residual = $akt_res")

    # After pair of complex parameters or
    # a real parameter, check stopping criteria
    # based on residual norm.
    if  !(is_compl && is_first)

      if akt_res <= min_res
        flag = "R"
        break
      end

      if with_min_rs && i>=2*stcf
        ra = minimum(res_log[1:i-stcf+1])
        rb = minimum(res_log[i-stcf+2:i+1])
        if res_log[1]-ra > 0 && (ra-rb)*(i-stcf+1)/((res_log[1]-ra)*stcf) < min_rs
          flag = "S"
          break
        end
      end

    end

  end
  # Check stopping criteria based on increase in ||Z_i||_F.
  if with_min_in
    nrmF_V_2 = sum(sum(abs(V).*abs(V)));    # Note the "abs"; V is complex.
    nrmF_Z_2 = nrmF_Z_2 + nrmF_V_2
    nrmF_rec[1:stcf-1] = nrmF_rec[2:stcf]
    nrmF_rec[stcf] = nrmF_V_2/nrmF_Z_2
    if !(is_compl && is_first) && i>stcf && all(nrmF_rec .< min_in)
      flag = "N"
      break
    end
  end
end

infoV && println("Computational Time:  ", toq())

Z, flag, res
end

"""
Based on lp_nrmu.m from LYAPACK 1.0

Helper function for lp_lradi

Using updated QR factorizations, this routine computes efficiently
a sequence of norms which correspond to either of the following types:

for tp = :B

    nrm = || F*Z_i*Z_i' + Z_i*Z_i'* F' + B*B' ||F,

for tp = :C

    nrm = || F' * Z_i * Z_i' + Z_i * Z_i' * F + B' * B ||F.

Here, F = A-Bf*Kf'.

The matrices Z_i must have much more rows than columns and they
have to obey the recursion

    Z_i = [ Z_{i-1}  V ].

Calling sequence:

    nrm, nrmQ, nrmR, nrmbs = ...
    lp_nrmu(  A, B, Bf, Kf, tpB, V, nrmQ, nrmR, nrmbs )

Input:

    tpB        (= true or false) the type of the norm

    Bf        real matrix Bf
              Set Bf = [] if not existing or zero!

    Kf        real matrix Kf
              Set Kf = [] if not existing or zero!

    B         n-x-m or m-x-n matrix B (must be real)

    V         n-x-r matrix V (may be complex)

    nrmQ,

    nrmR,

    nrmbk     variables for internal use (they contain the data of
              the updated QR factorization).

Output:

    nrm       the current value of the Frobenius norm nrm_i.

    nrmQ,

    nrmR,

    nrmbk     variables for internal use (they must be output
              parameters).

Remarks:

    Using this routine for computing the norms within an iteration
    can save a lot of computation compared to repeatedly calling
    'lp_nrm'.

    This routine must be used as follows:

    1. Before the iteration starts, it must be invoked with the
       parameters V = nrmQ = nrmR = nrmbl = []. This initializes the
       variables nrmQ, nrmR, nrmbl.

    2. The routine must be called in each step of the iteration:

Example:

    ...

    nrm,nrmQ,nrmR,nrmbs = lp_nrmu(A,B,[],[],tpB,[],[],[],[])

    ...

    Z = zeros(n,0)

      ...

      for i = 1:100

        ...

        V = ...

        ...

        Z = [ Z V ];    # iteration:  Z_i = [ Z_i-1 V ]

        ...

        nrm,nrmQ,nrmR,nrmbs = lp_nrmu(A,B,[],[],tpB,V,nrmQ,nrmR,nrmbs)

        ...

      end
      ...
"""
function lp_nrmu( A, B, Bf, Kf, tpB, V, nrmQ, nrmR, nrmbs )

const with_BK = length(Bf)>0

const n= size(A,1);                    # Get system order.

if isempty(V)
  # The routine is called for the first time (before the iteration
  # 'outside' is started (i.e. V = zeros(n,0) or []!)) in order to
  # initialize the QR factorization nrmQ*nrmR correctly.

  if tpB
    nrmbs = [size(B,2)]
    nrmQ,nrmR = qr(B)
  else
    nrmbs = [size(B,1)]
    nrmQ,nrmR = qr(B')
  end
  nrm = vecnorm(nrmR*nrmR')

else
  # The routine is called during the iteration 'outside'.
  # The QR factorization is updated and the norm is computed.

  # Update of the QR factorization.
  push!(nrmbs,size(V,2))
  lw = size(nrmQ,2)
  lz = size(V,2)

  if tpB
    TM = A*V
    with_BK && (TM = TM-Bf*(Kf'*V))
    Z = [ TM V ]
  else
    TM = A'*V
    with_BK && (TM = TM-Kf*(Bf'*V))
    Z = [ TM V ]
  end

  for j = 1:2*lz
    a = Array(Complex{Float64},lw,1)
    #a = complex(zeros(lw,1))
    t = Z[:,j]
    for k = 1:lw
      u = nrmQ[:,k]
      alpha = dot(u,t)
      t = t-alpha*u
      a[k] = alpha[1]
    end
    beta = norm(t)
    nrmQ = [nrmQ t/beta]
    #nrmQ = push!(nrmQ, t/beta...)
    nrmR = [ nrmR a ; zeros(1,lw) beta ]
    lw = lw+1
  end

  # Computation of  nrmR * [permutation matrix] * nrmR'
  RT = copy(nrmR)
  ie2 = nrmbs[1]
  for j = 2:length(nrmbs)
    ia1 = ie2+1
    ie1 = ie2+nrmbs[j]
    ia2 = ie1+1
    ie2 = ie1+nrmbs[j]
    TMP = RT[:,ia1:ie1]
    RT[:,ia1:ie1] = RT[:,ia2:ie2]
    RT[:,ia2:ie2] = TMP
  end
  nrm = vecnorm(nrmR*RT')
end
nrm, nrmQ, nrmR, nrmbs
end

"""
    `lp_para(A,Bf,Kf,l0,kp,km,b0)`

Based on lp_para.m from LYAPACK 1.0

Estimation of suboptimal ADI shift parameters for the matrix

F = A - Bf * Kf'.

Calling sequence:

p = lp_para(A,Bf,Kf,l0,kp,km,b0)

Input:

    Bf        matrix Bf
              Set Bf = [] if not existing or zero!
    Kf        matrix Kf
              Set Kf = [] if not existing or zero!
    l0        desired number of shift parameters (kp+km > 2*l0)
              (The algorithm delivers either l0 or l0+1 parameters!)
    kp, km    numbers of Arnoldi steps w.r.t. F and inv(F),
              respectively (kp, km << order of A)
    b0         Arnoldi start vector (optional; chosen at random if not
              provided).

Output:

    p         an l0- or l0+1-vector of suboptimal ADI parameters

Remarks:

    Typical values are l0 = 10..40, kp = 20..80, km = 10..40.
    The harder the problem is the large values are necessary.
    Larger values mostly result in a faster convergence, but also in a
    larger memory requirement.
    However, for "well-conditioned" problems small values of l0 can
    lead to the optimal performance.

References:

    1. T. Penzl.
    LYAPACK (Users' Guide - Version 1.0).
    1999.

Input data not completely checked!
"""
function lp_para(A,Bf=[],Kf=[],l0=NaN,kp=NaN,km=NaN,b0=[])

err_code = 0
n=size(A,1); # Get system order.

if isnan(l0) && isnan(kp) && isnan(km)
  if n > 51
    kp=50;
    km=25;
    l0=20;
  elseif n > 15
    kp = n - 5
    km = Int(round(kp/2))
    l0 = km-2
  else
    kp = n-2
    km = Int(round(kp/2))
    l0 = km-1
  end
else
  kp >= n && error("kp must be smaller than n!")
  km >= n && error("km must be smaller than n!")
  2*l0 >= kp+km && error("2*l0 must be smaller than kp+km!")
end

isempty(b0) && (b0 = randn(n,1))
b0 = (1/norm(b0))*b0

rwp = []
rwm = []
rw = []
Hp = []
Hm = []

if kp > 0
  Hp,V = lp_arn(A,Bf,Kf,kp,b0,:p)
  rwp = eigvals(Hp[1:kp,1:kp]);                  # =: R_+
  rw = rwp
end

if km > 0
  Hm,V = lp_arn(A,Bf,Kf,km,b0,:m)
  rwm = ones(km,1)./eigvals(Hm[1:km,1:km]);      # =: 1 / R_-
  push!(rw,rwm...);                           # =: R
end

if any(real(rw) .>= zeros(size(rw)))
  err_code = 1
  println("These are the Ritz values computed by the Arnoldi process w.r.t. F: $rwp
  These are the Ritz values computed by the Arnoldi process w.r.t. inv(F): $rwm

  ####################################################################
  WARNING in 'lp_para': NON-STABLE RITZ VALUES DETECTED!!!

  This is quite a serious problem, that can be caused by
  (i)   non-stable matrices F (Be sure that F is stable. ADI like
        methods only work for stable or antistable problems. If your
        Lyapunov equation is antistable, multiply it by -1.)
  (ii)  matrices F that are stable but have an indefinite symmetric
        part (This is THE weak point of this algorithm. Try to work
        with the 'reduced' Ritz values, i.e., the unstable values are
        simply removed. This is not an elegant measure but it may work.
        However, the convergence of ADI can be poor. This measure is
        taken automatically. Another measure might be to enlarge the
        values of kp or km, and run the program again.
  (iii) matrices F with a negative definite, but ill-conditioned
        symmetric part (This is quite unlikely. The problem is
        caused by round-off errors).

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

  NOTE: The unstable Ritz values will be ignored in the further computation!!! ")
  rw0 = copy(rw)
  rw = []
  for j = 1:length(rw0)
    if real(rw0[j])<0
      push!(rw,rw0[j])
    end
  end
end

p = lp_mnmx(rw,l0)

p
end

"""
Based on lp_arn_m.m and lp_arn_p.m from LYAPACK 1.0 and includes many of the
same comments much of the same description.

helper function for lp_para

Arnoldi method w.r.t. inv(F), where F = A-Bf*Kf'.

Calling sequence:

    H,V = lp_arn_m(A,Bf,Kf,k,r)

Input:

    A         Matrix A

    Bf        matrix Bf
              Set Bf = [] if not existing or zero!

    Kf        matrix Kf

              Set Kf = [] if not existing or zero!

    k         number of Arnoldi steps (usually k << n, where
              n is the order of the system)

    r         initial n-vector
              (optional - chosen by random, if []).

    pm        %TODO write description

Output:

    H         matrix H ((k+1)-x-k matrix, upper Hessenberg)

    V         matrix V (n-x-(k+1) matrix, orthogonal columns).

Method:

    The ("inverse") Arnoldi method produces matrices V and H such that

    V(:,1) in span{r},
    V' * V = eye(k+1),
      inv(F) * V(:,1:k) = V * H.

Remark:

    This implementation does not check for (near-)breakdown!

"""
function lp_arn(A,Bf,Kf,k,r,pm)

const with_BK = !isempty(Bf)

const n = size(A,1)                 # Get system order.
k >= n-1 && error("k must be smaller than the order of A!")
isempty(r) && (r = randn(n,1))

V = zeros(n,k+1)
H = zeros(k+1,k)

V[:,1] = (1.0/norm(r))*r

beta = 0

LP_L,LP_U,~ = lu(full(A),Val{false})

if with_BK && pm==:m
  # SM = inv(F)*Bf*inv(I-Kf'*inv(F)*Bf)
  # (This is the main part of the term needed for the low
  # rank correction in the Sherman-Morrison formula.)
  Im = speye(size(Bf,2))
  TM = LP_U\(LP_L\Bf)
  SM = TM/(Im-Kf'*TM)
end

for j = 1:k

  if j > 1
    H[j,j-1] = beta
    V[:,j] = (1.0/beta)*r
  end

  if (pm == :m)
    w = LP_U\(LP_L\V[:,j])
    with_BK && (w = w + SM*(Kf'*w)); # LR correction by SM formula
  else
    w = A*V[:,j]
    with_BK && (w = w-Bf*(Kf'*V[:,j]))
  end
  r = w

  for i = 1:j
    H[i,j] = dot(V[:,i],w)
    r = r-H[i,j]*V[:,i]
  end

  beta = norm(r)
  H[j+1,j] = beta

end

V[:,k+1] = (1.0/beta)*r

H,V
end

"""
Based on lp_mnmx.m from LYAPACK 1.0 and includes many of the same comments and
much of the same description.

Suboptimal solution of the ADI minimax problem. The delivered parameter
set is closed under complex conjugation.

Calling sequence:

    p = lp_mnmx(rw,l0)

Input:

    rw        a vector containing numbers in the open left half plane, which
              approximate the spectrum of the corresponding matrix, e.g.,
              a set of Ritz values. The set must be closed w.r.t. complex
              conjugation

    l0        desired number of shift parameters (length(rw) >= l0)
              (The algorithm delivers either l0 or l0+1 parameters!).

Output:

    p         an l0- or l0+1-vector of suboptimal ADI parameters

"""
function lp_mnmx(rw,l0)

if length(rw)<l0
  error("length(rw) must be at least l0.")
end

max_rr = +Inf;                       # Choose initial parameter (pair)
p0 = 0
for i = 1:length(rw)
  max_r,~ = lp_s(rw[i],rw)
  if max_r < max_rr
    p0 = rw[i]
    max_rr = max_r
  end
end

imag(p0)!=0 ? (p = [ p0; conj(p0) ]) : (p = [p0])

max_r,i = lp_s(p,rw);         # Choose further parameters.

while size(p,1) < l0

  p0 = rw[i]
  imag(p0)!=0 ? push!(p,p0,conj(p0)) : push!(p,p0)

  max_r,i = lp_s(p,rw)

end

p
end

"""
Based on lp_s.m from LYAPACK 1.0 and includes much of the same description

Helper function for lp_para

Computation of the maximal magnitude of the rational ADI function over
a discrete subset of the left complex half plane.

Calling sequence:

    max_r,ind = lp_s(p,set)

Input:

    p        vector of ADI parameters

    set      vector representing the discrete set.

Output:

    max_r    maximal magnitude of the rational ADI function over set

    ind      index - maximum is attained for set(ind).
"""
function lp_s(p,set)

max_r = -1
ind = 0

for i = 1:length(set)

  x = set[i]

  rr = 1
  for j = 1:length(p)

    rr = rr*abs(p[j]-x)/abs(p[j]+x)

  end

  if rr > max_r

    max_r = rr
    ind = i

  end

end
max_r,ind
end

# end module matrixEqs
end


"""
Convenience function to return the energy levels and
degeneracies of a linear molecule.

B is the rotational constant, D is the centrifugal
distortion term, and J_max is the maximum value
to simulate up to. Assumes that B and D are the same units.
"""
function make_linear_molecule(B, D, J_max, T=10.)
    J = 1:J_max
    E = @. B * J * (J + 1) - D * J^2 * (J + 1)^2
    G = @. 2 * (J + 1)
    levels = BaseLevel.(E, G)
    # use the boltzmann weighting as the intensity
    qrot = Q.(levels, T)
    intensity = qrot ./ sum(qrot)
    # this calculates the transition energies
    ν = diff(E)
    transitions = Transitions()
    @inbounds for i in eachindex(ν)
        trans = Transition(ν[i], 0., qrot[i+1], levels[i], levels[i+1], 0., 0.)
        push!(transitions, trans)
    end
    return levels, transitions
end

function add_noise!(I::Vector{<:Real}, σ::Real)
    target_σ = max(I) / σ
    I .+= rand(Normal(0., target_σ,), length(I))
end


using EconPDEs, Distributions

# This script implements a generalization of the model presented
# in the Handbook of Macro Vol. 2 chapter
# "Macro, Money, and Finance: A Continuous Time Approach"
# by Brunnermeier and Sannikov.

# The model includes agent-specific risk aversion, EIS,
# discount rates, depreciation, consumption good productivity,
# and volatility of capital holdings.

# macro instantiate_parameters(m)
#     for name in fieldnames(m)
#         eval(:$(name) = m.$(name))
#     end
# end

mutable struct BrunnermeierSannikov2016Model
    # Utility Function
    γ₁::Float64
    γ₂::Float64
    ϵ₁::Float64
    ϵ₂::Float64
    ρ₁::Float64
    ρ₂::Float64

    # Technology
    a₁::Float64
    a₂::Float64
    σ₁::Float64
    σ₂::Float64
    δ₁::Float64
    δ₂::Float64
    A₁::Float64
    A₂::Float64
    B₁::Float64
    B₂::Float64

    # "Skin in the game" constraint
    χ_::Float64

    # Transfers
    τ::Float64
end

function BrunnermeierSannikov2016Model(; γ₁ = 5., γ₂ = 5., ϵ₁ = 2.,
                                       ϵ₂ = 2., ρ₁ = .05, ρ₂ = .05,
                                       a₁ = 1, a₂ = .7, σ₁ = 0.03,
                                       σ₂ = 0.03, δ₁ = 0.05, δ₂ = 0.05,
                                       A₁ = 53.2, A₂ = 53.2,
                                       B₁ = 0., B₂ = 0.,
                                       χ_ = 1., τ = .01)
    BrunnermeierSannikov2016Model(γ₁, γ₂, ϵ₁, ϵ₂,
                                  ρ₁, ρ₂, a₁, a₂,
                                  σ₁, σ₂, δ₁, δ₂, A₁, A₂, B₁, B₂, χ_, τ)
end

function initialize_stategrid(m::BrunnermeierSannikov2016Model; η_n = 80)
  OrderedDict(:η => range(0.01, stop = 0.99, length = η_n))
end

function initialize_y(m::BrunnermeierSannikov2016Model, stategrid::OrderedDict)
  x = fill(1.0, length(stategrid[:η]))
  OrderedDict(:ξ₁ => x,
              :ξ₂ => x)
end

function (m::BrunnermeierSannikov2016Model)(state::NamedTuple, y::NamedTuple)
    γ₁ = m.γ₁; γ₂ = m.γ₂; ϵ₁ = m.ϵ₁; ϵ₂ = m.ϵ₂;
    ρ₁ = m.ρ₁; ρ₂ = m.ρ₂; a₁ = m.a₁; a₂ = m.a₂;
    σ₁ = m.σ₁; σ₂ = m.σ₂; δ₁ = m.δ₁; δ₂ = m.δ₂;
    A₁ = m.A₁; A₂ = m.A₂; B₁ = m.B₁; B₂ = m.B₂;
    χ_ = m.χ_; τ = m.τ;

    η = state.η
    # ξ₁, ξ₁η, ξ₁ηη, ξ₂, ξ₂η, ξ₂ηη, q, qη, qηη = y.ξ₁, y.ξ₁η, y.ξ₁ηη, y.ξ₂, y.ξ₂η, y.ξ₂ηη, y.q, y.qη, y.qηη
    ξ₁, ξ₁η, ξ₁ηη, ξ₂, ξ₂η, ξ₂ηη = y.ξ₁, y.ξ₁η, y.ξ₁ηη, y.ξ₂, y.ξ₂η, y.ξ₂ηη

    # Determine capital allocations
    ψ₁ = 1.
    ψ₂ = 0.
    q = (a₁ * ψ₁ + a₂ * ψ₂) / (ρ₁ * η + ρ₂ * (1 - η))

    # Internal investment function
    #     Φᵢ(ιᵢ) = (Aᵢ / Bᵢ) * log(ιᵢ / Bᵢ - 1) - δᵢ
    # with optimality condition
    #     Φ'ᵢ(ιᵢ) = 1 / q
    # μ_K1 = (A₁ / B₁) * log(A₁ * q) - δ₁
    # μ_K2 = (A₂ / B₂) * log(A₂ * q) - δ₂
    # μ_K1 = (q - B₁) / (2 * A₁) - δ₁
    # μ_K2 = (q - B₂) / (2 * A₂) - δ₂
    # μ_K  = ψ₁ * μ_K1 + ψ₂ * μ_K2
    μ_K = 0.02
    μ_K1 = 0.02
    μ_K2 = 0.02
    # ι₁   = (A₁ * q - 1) / B₁
    # ι₂   = (A₂ * q - 1) / B₂
    # ι₁ = A₁ * ((q - B₁) / (2 * A₁))^2 + B₁ * ((q - B₁) / (2 * A₁))
    # ι₂ = A₂ * ((q - B₂) / (2 * A₂))^2 + B₂ * ((q - B₂) / (2 * A₂))
    # ι    = ψ₁ * ι₁ + ψ₂ * ι₂
    ι = 0.
    ι₁ = 0.

    # Compute volatilities
    # σ_q  = qη * (ψ₁ - η) * (ψ₁ * σ₁ + ψ₂ * σ₂) / (q - qη * (ψ₁ - η))
    σ_q = 0.
    σ_η  = ψ₁ / η * (σ₁ + σ_q) - σ_q - (ψ₁ * σ₁ + ψ₂ * σ₂)
    σ_ξ₁ = ξ₁η / ξ₁ * σ_η * η
    σ_ξ₂ = ξ₂η / ξ₂ * σ_η * η
    # σ_ξ₁ = σ_v₁ - σ_η - σ_q
    # σ_ξ₂ = σ_v₁ + η / (1 - η) * σ_η - σ_q

    # Compute risk premia
    ς_K₁  = ψ₁ / η * (σ₁ + σ_q)^2
    ς_ξ₁ = (σ₁ + σ_q) * σ_ξ₁

    # Consumption market-clearing
    c₁_per_n₁ = (ϵ₁ == 1.) ? ρ₁ : ξ₁^(1 - ϵ₁)
    c₂_per_n₂ = (ϵ₂ == 1.) ? ρ₂ : ξ₂^(1 - ϵ₂)
    # ∂q_∂t     = q * (c₁_per_n₁ * η + c₂_per_n₂ * (1 - η)) - (a₁ * ψ₁ + a₂ * ψ₂ - ι)
    ∂q_∂t = 0.

    # Compute drift
    μ_η = (a₁ - ι₁) / q - c₁_per_n₁ - τ +
        (ψ₁ / η - 1) * (γ₁ * ς_K₁ + (γ₁ - 1) * ς_ξ₁) -
        (ψ₁ / η * (σ₁ + σ_q) - (ψ₁ * σ₁ + ψ₂ * σ₂ + σ_q)) * (ψ₁ * σ₁ + ψ₂ * σ₂ + σ_q)

    # Compute interest rate
    # μ_q = qη / q * μ_η * η + (1 / 2) * qηη / q * (η * σ_η)^2
    μ_q = 0.
    E_returns = (a₁ - ι₁) / q + μ_q + μ_K1 + σ₁ * σ_q
    r_f = E_returns - γ₁ * ς_K₁ - (γ₁ - 1) * ς_ξ₁

    # Evolution of marginal value of agents' net worth
    μ_ξ₁ = ξ₁η / ξ₁ * μ_η * η + (1 / 2) * ξ₁ηη / ξ₁ * (η * σ_η)^2 # Drift from Ito's lemma
    μ_ξ₂ = ξ₂η / ξ₂ * μ_η * η + (1 / 2) * ξ₂ηη / ξ₂ * (η * σ_η)^2
    c₁_term = (ϵ₁ == 1.) ? 1 / (1 - 1 / ϵ₁) * (ξ₁^(1 - ϵ₁) - ρ₁) : ρ₁ * log(ρ₁ / ξ₁)
    c₂_term = (ϵ₂ == 1.) ? 1 / (1 - 1 / ϵ₂) * (ξ₂^(1 - ϵ₂) - ρ₂) : ρ₂ * log(ρ₂ / ξ₂)

    ∂ξ₁_∂t = ξ₁ * (c₁_term + μ_ξ₁ + r_f -
        c₁_per_n₁ - τ + γ₁ / 2 * ((ψ₁ / η * (σ₁ + σ_q))^2 - σ_ξ₁^2))
    ∂ξ₂_∂t = 0. #ξ₂ * (c₂_term + μ_ξ₂ + r_f -
        # c₂_per_n₂ + τ * η / (1 - η) - γ₂ / 2 * σ_ξ₂^2)

 return (∂ξ₁_∂t, ∂ξ₂_∂t), (μ_η), (μ_η = μ_η, q = q, ξ₁ = ξ₁, ξ₂ = ξ₂,
                                         ψ₁ = ψ₁, ψ₂ = ψ₂,
                                         μ_K = μ_K, ι = ι, σ_q = σ_q,
                                         σ_η = σ_η, σ_ξ₁ = σ_ξ₁, σ_ξ₂ = σ_ξ₂,
                                         ς_K₁ = ς_K₁, ς_ξ₁ = ς_ξ₁)
end

m = BrunnermeierSannikov2016Model(ϵ₁ = 1., ϵ₂ = 1., γ₁ = 1., γ₂ = 1.)
stategrid = initialize_stategrid(m)
y0 = initialize_y(m, stategrid)
y, result, distance = pdesolve(m, stategrid, y0)
# y, result, distance = pdesolve(m, stategrid, y0; is_algebraic = OrderedDict(:v₁ => false, :v₂ => false, :q => true))


function dden_complete(ygrid, W, eta, Z, mus, sig2; i, j)
  K = size(W, 2)
  L0 = size(eta[0], 3)
  L1 = size(eta[1], 3)
  L = Dict(0 => L0, 1 => L1)

  dden = [begin
            si = sqrt(sig2[i])
            dd = 0.0
            for k in 1:K
              ddk = 0.0
              z = Z[j, k]
              for ell in 1:L[z]
                mu_zl = mus[z][ell]
                eta_zijl = eta[z][i, j, ell]
                ddk += eta_zijl * pdf(Normal(mu_zl, si), y)
              end
              dd += W[i, k] * ddk
            end
            dd
          end for y in ygrid]

  return dden
end

function dden_complete(ygrid, W, eta, Z, mus, sig2)
  I, J = size(eta[0])[1:end-1]

  dden = [dden_complete(ygrid, W, eta, Z, mus, sig2, i=i, j=j)
          for i in 1:I, j in 1:J]

  return dden
end


#=
check a sufficient condition for pointwise membership of vector valued
polynomials in the epinormone/epinormeucl cone
=#

struct NormConePoly{T <: Real} <: ExampleInstanceJuMP{T}
    polys_name::Symbol
    deg::Int
    is_feas::Bool # whether model should be primal-dual feasible; only for testing
    use_L2::Bool # check for membership in epigraph of L2, otherwise use L1
end

function build(inst::NormConePoly{T}) where {T <: Float64}
    halfdeg = div(inst.deg + 1, 2)
    (U, pts, Ps) = PolyUtils.interpolate(
        PolyUtils.FreeDomain{T}(1), halfdeg)
    vals = normconepoly_data[inst.polys_name].(pts)
    l = length(vals[1])
    if inst.use_L2
        cone = Hypatia.WSOSInterpEpiNormEuclCone{T}(l, U, Ps)
    else
        cone = Hypatia.WSOSInterpEpiNormOneCone{T}(l, U, Ps)
    end

    model = JuMP.Model()
    JuMP.@constraint(model, [v[i] for i in 1:l for v in vals] in cone)

    return model
end

function test_extra(inst::NormConePoly{T}, model::JuMP.Model) where T
    stat = JuMP.termination_status(model)
    @test stat == (inst.is_feas ? MOI.OPTIMAL : MOI.INFEASIBLE)
end

normconepoly_data = Dict(
    :polys1 => (x -> [x^2 + 2, x]),
    :polys2 => (x -> [2x^2 + 2, x, x]),
    :polys3 => (x -> [x^2 + 2, x, x]),
    :polys4 => (x -> [2 * x^4 + 8 * x^2 + 4, x + 2 + (x + 1)^2, x]),
    :polys5 => (x -> [x, x^2 + x]),
    :polys6 => (x -> [x, x + 1]),
    :polys7 => (x -> [x^2, x]),
    :polys8 => (x -> [x + 2, x]),
    :polys9 => (x -> [x - 1, x, x]),
    )


# Simplifications go here

# Basic maps

composedmap1(m1::IdentityMap, m2) = m2
composedmap1(m1::IdentityMap{T}, m2::Map{T}) where {T} = m2
composedmap1(m1::IdentityMap{T}, m2::Map{S}) where {S,T} = convert(Map{T}, m2)

composedmap2(m1, m2::IdentityMap) = m1
composedmap2(m1::Map{T}, m2::IdentityMap{T}) where {T} = m1
composedmap2(m1::Map{S}, m2::IdentityMap{T}) where {S,T} = convert(Map{T}, m1)

composedmap2(m1, m2::ConstantMap) = m2
composedmap2(m1::Map{T}, m2::ConstantMap) where {T} = ConstantMap{T}(constant(m2))
composedmap1(m1::ConstantMap{T}, m2) where {T} = ConstantMap{T}(m2(constant(m1)))

composedmap2(m1, m2::ZeroMap) = m2
composedmap2(m1::Map{T}, m2::ZeroMap{S,U}) where {S,T,U} = ZeroMap{T,U}()
composedmap1(m1::ConstantMap{T}, m2::ZeroMap{S,U}) where {S,T,U} = ZeroMap{T,U}()

multiply_map1(m1::ZeroMap, m2) = m1
multiply_map2(m1, m2::ZeroMap) = m2
multiply_map2(m1::ConstantMap{T}, m2::ConstantMap{S}) where {S,T} =
    ConstantMap{promote_type(S,T)}(constant(m1)*constant(m2))

sum_map1(m1::ZeroMap, m2) = m2
sum_map2(m1, m2::ZeroMap) = m1
sum_map2(m1::ConstantMap{T}, m2::ConstantMap{S}) where {S,T} =
    ConstantMap{promote_type(S,T)}(constant(m1)+constant(m2))

## Affine maps

composedmap(m1::AbstractAffineMap, m2::AbstractAffineMap) = affine_composition(m1, m2)

"""
Compute the affine map that represents map2 after map1, that is:
`y = a2*(a1*x+b1)+b2 = a2*a1*x + a2*b1 + b2`.
"""
affine_composition(map1::AbstractAffineMap, map2::AbstractAffineMap) =
    AffineMap(matrix(map2) * matrix(map1), matrix(map2)*vector(map1) + vector(map2))

affine_composition(map1::AffineMap, map2::AffineMap) =
    AffineMap(unsafe_matrix(map2) * unsafe_matrix(map1), unsafe_matrix(map2)*unsafe_vector(map1) + unsafe_vector(map2))

affine_composition(map1::LinearMap, map2::LinearMap) =
    LinearMap(unsafe_matrix(map2) * unsafe_matrix(map1))

affine_composition(map1::LinearMap, map2::AffineMap) =
    AffineMap(unsafe_matrix(map2) * unsafe_matrix(map1), unsafe_vector(map2))

affine_composition(map1::AffineMap, map2::LinearMap) =
    AffineMap(unsafe_matrix(map2) * unsafe_matrix(map1), unsafe_matrix(map2)*unsafe_vector(map1))

affine_composition(map1::Translation, map2::Translation) =
    Translation(unsafe_vector(map2) + unsafe_vector(map1))

affine_composition(map1::Translation, map2::LinearMap) =
    AffineMap(unsafe_matrix(map2), unsafe_matrix(map2)*unsafe_vector(map1))

affine_composition(map1::LinearMap, map2::Translation) =
    AffineMap(unsafe_matrix(map1), unsafe_vector(map2))

# The sum of two affine maps is again an affine map
mapsum(map1::AbstractAffineMap, map2::AbstractAffineMap) =
    AffineMap(matrix(map1)+matrix(map2), vector(map1)+vector(map2))

==(m1::ProductMap, m2::IdentityMap) = all(map(isidentity, components(m1)))
==(m1::IdentityMap, m2::ProductMap) = m2 == m1


export MultiHomTotalDegreeSystem

"""
    MultiHomogeneousTotalDegreeSystem(polynomials, vars) <: AbstractSystem

Create a multi-homogeneous total degree system as described in

An efficient start system for multi-homogeneous polynomial continuation,
Wampler, C.W. Numer. Math. (1993) 66: 517. https://doi.org/10.1007/BF01385710
"""
struct MultiHomTotalDegreeSystem{T} <: AbstractSystem
    D::Matrix{Int}
    C::Matrix{Float64}
    scaling_factors::Vector{T}
end

function MultiHomTotalDegreeSystem(D, C)
    MultiHomTotalDegreeSystem(D, C, ones(size(D, 2)))
end

struct MultiHomTotalDegreeSystemCache{M,T} <: AbstractSystemCache
    B::Matrix{T}
    Ĝ::Matrix{T}
    R::Matrix{T}
    S::Matrix{T}
    ranges::NTuple{M,UnitRange{Int}}
    homvars::NTuple{M,Int}
end

Base.size(F::MultiHomTotalDegreeSystem) = (size(F.D, 2), sum(size(F.D)))
Base.length(F::MultiHomTotalDegreeSystem) = size(F.D, 2)

function cache(F::MultiHomTotalDegreeSystem, x::PVector{TX,M}) where {TX,M}
    T = typeof(F.C[1, 1] * x[1] + F.C[1, 1] * x[1])
    n = length(F)
    B, Ĝ = zeros(T, M, n), zeros(T, M, n)
    R, S = zeros(T, M, n), zeros(T, M, n)
    ranges_homvars = ProjectiveVectors.dimension_indices_homvars(x)
    ranges = first.(ranges_homvars)
    homvars = last.(ranges_homvars)
    MultiHomTotalDegreeSystemCache(B, Ĝ, R, S, ranges, homvars)
end

function evaluate!(
    u,
    F::MultiHomTotalDegreeSystem,
    z::ProjectiveVectors.PVector,
    cache::MultiHomTotalDegreeSystemCache{M},
) where {M}
    D, C = F.D, F.C
    B, ranges, homvars = cache.B, cache.ranges, cache.homvars

    n = size(B, 2)

    @boundscheck checkbounds(u, 1:n)

    # Compute all bᵢⱼ and store in B
    # Since B is column major we store bᵢⱼ in B[j, i]
    @inbounds for i = 1:n
        for j = 1:M
            if D[j, i] != 0
                bᵢⱼ = -zero(B[j, i])
                for k in ranges[j]
                    bᵢⱼ += C[k, i] * z[k]
                end
                B[j, i] = bᵢⱼ
            end
        end
    end

    @inbounds for i = 1:n
        gᵢ = one(eltype(u))
        for j = 1:M
            dᵢⱼ, bᵢⱼ = D[j, i], B[j, i]
            if dᵢⱼ != 0
                gᵢ *= bᵢⱼ^dᵢⱼ - z[homvars[j]]^dᵢⱼ
            end
        end
        u[i] = gᵢ * F.scaling_factors[i]
    end

    u
end
function evaluate(F::MultiHomTotalDegreeSystem, x, cache::MultiHomTotalDegreeSystemCache)
    evaluate!(similar(x, size(F, 1)), F, x, cache)
end

function jacobian!(
    U,
    F::MultiHomTotalDegreeSystem,
    z,
    cache::MultiHomTotalDegreeSystemCache{M},
) where {M}
    evaluate_and_jacobian!(nothing, U, F, z, cache)
    U
end
function jacobian(F::MultiHomTotalDegreeSystem, x, cache::MultiHomTotalDegreeSystemCache)
    jacobian!(similar(x, size(F)), F, x, cache)
end

function evaluate_and_jacobian!(
    u,
    U,
    F::MultiHomTotalDegreeSystem,
    z,
    cache::MultiHomTotalDegreeSystemCache{M},
) where {M}
    n, N = size(F)

    @boundscheck checkbounds(U, 1:n, 1:N)
    if u !== nothing
        @boundscheck checkbounds(u, 1:n)
    end

    D, C = F.D, F.C
    B, Ĝ, R, S, ranges, homvars =
        cache.B, cache.Ĝ, cache.R, cache.S, cache.ranges, cache.homvars

    n = size(B, 2)

    @inbounds for i = 1:n
        # Compute all bᵢⱼ and store in B
        # Since B is column major we store bᵢⱼ in B[j, i]
        for j = 1:M
            if D[j, i] != 0
                bᵢⱼ = -zero(B[j, i])
                for k in ranges[j]
                    bᵢⱼ += C[k, i] * z[k]
                end
                B[j, i] = bᵢⱼ
            end
        end

        # Compute all ĝᵢⱼ and store in Ĝ
        # Since Ĝ is column major we store ĝᵢⱼ in Ĝ[j, i]
        for j = 1:M
            dᵢⱼ, bᵢⱼ = D[j, i], B[j, i]
            if dᵢⱼ == 0
                Ĝ[j, i] = one(eltype(Ĝ))
            else
                Ĝ[j, i] = bᵢⱼ^dᵢⱼ - z[homvars[j]]^dᵢⱼ
            end
        end

        # Accumulate subproducts forward
        R[1, i] = rᵢⱼ_prev = Ĝ[1, i]
        for j = 2:M
            if D[j, i] != 0 # otherwise Ĝ[j, i] = 1
                R[j, i] = rᵢⱼ_prev = rᵢⱼ_prev * Ĝ[j, i]
            else
                R[j, i] = rᵢⱼ_prev
            end
            if u !== nothing
                u[i] = rᵢⱼ_prev * F.scaling_factors[i]
            end
        end

        # Accumulate subproducts backward
        S[M, i] = sᵢⱼ_prev = Ĝ[M, i]
        for j = M-1:-1:1
            if D[j, i] != 0 # otherwise Ĝ[j, i] = 1
                S[j, i] = sᵢⱼ_prev = sᵢⱼ_prev * Ĝ[j, i]
            else
                S[j, i] = sᵢⱼ_prev
            end
        end

        # Compute partial derivatives
        for j = 1:M
            dᵢⱼ = D[j, i]
            for k in ranges[j]
                c = C[k, i]
                if iszero(c) || iszero(dᵢⱼ)
                    U[i, k] = zero(eltype(U))
                else
                    if dᵢⱼ == 1
                        u_ik = c
                    else
                        u_ik = dᵢⱼ * B[j, i]^(dᵢⱼ - 1) * c
                    end
                    if j > 1
                        u_ik *= R[j-1, i]
                    end
                    if j < M
                        u_ik *= S[j+1, i]
                    end
                    U[i, k] = u_ik * F.scaling_factors[i]
                end
            end

            k = homvars[j]
            if iszero(dᵢⱼ)
                U[i, k] = zero(eltype(U))
            else
                if dᵢⱼ == 1
                    u_ik = -dᵢⱼ
                else
                    u_ik = -dᵢⱼ * z[k]^(dᵢⱼ - 1)
                end
                if j > 1
                    u_ik *= R[j-1, i]
                end
                if j < M
                    u_ik *= S[j+1, i]
                end
                U[i, k] = u_ik * F.scaling_factors[i]
            end
        end
    end
    nothing
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -6.103907844236108 <= q <= 8.378675476959124)

                     add_NL_constraint(m, :(log(1 + exp(-0.8297607126794757 + -0.5064858067335005*log(1 + exp(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2]))) + -0.8619658571743893*log(1 + exp(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2]))) + 0.8415580445654567*log(1 + exp(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))))) + log(1 + exp(-0.0809505831099564 + 0.6185835236357913*log(1 + exp(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2]))) + -0.5288114683363174*log(1 + exp(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2]))) + -0.42070009030061817*log(1 + exp(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))))) + log(1 + exp(0.16011867292674786 + -0.14898607305663658*log(1 + exp(-0.6188655598844357 + -0.38627203797486676*$(x[1]) + 0.4878933200630442*$(x[2]))) + -0.8220978033005912*log(1 + exp(-0.7914291135758744 + 0.8682104207735972*$(x[1]) + 0.3684602934879275*$(x[2]))) + -0.8518875556726564*log(1 + exp(-0.26405078449033637 + -0.6452849815457378*$(x[1]) + 0.9586530591287103*$(x[2]))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


"""
    rankinequalities(m::Model; max_ranks::Int = 5, suppress_warnings::Bool = true, U::Int = 10000, diag::Bool = false, active_constr::Bool = false, real_recourse_dim::Int = -1, conic_δ::Bool = false)
Returns ranking data in the form `Vector{Dict}` where each dictionary corresponds to a particular rank level and
contains the optimal flexibility_index, active constraint indexes, and flexibility model. The function will iteratively
solve copies of the flexibility model via `solve` where the prior active constraints are turned off in order rank the most limiting
constraints. The user can specify the maximum number of rank levels and the flexibility index problem will be repeatedly
solved until that maximum is acheived or the problem becomes unbounded, which occurs first.

**Arguments**
- `m::Model` The flexibility model.

**Keyword Arguments**
- `max_ranks::Int = 5` The maximum number of rank levels. 2
- `suppress_warnings::Bool = false` Indicates if solver warnings should be suppressed.
- `U::Union{Int, Float64} = 10000` The slack variable upper bound.
- `diag::Bool = false` Indicates whether or not to diagnonalize ellipsoidal uncertainty set (this is only active when an ellipsoidal set is used).
- `active_constr::Bool = false` Indicates if the optional active constraint should be used which enforces how many inequalities are active at the solution, this must be set to `true` for systems without control variables and/or contain state variables.
- `real_recourse_dim::Int = -1` The actual number of recourse variables in case state variables are included as recourse variables. This is mandatory if `active_constr = true` and no state variables are provided.
- `conic_δ::Bool = false` This should be set to `true` if a conic solver is used such as Pajarito.jl.

```julia
julia> rankinequalities(m, active_constr = true)
2-element Array{Dict,1}:
 Dict{String,Any}(Pair{String,Any}("flexibility_index", 3.59938),Pair{String,Any}("model", Feasibility problem with:
 * 6 linear constraints
 * 6 variables
Solver is Pavito),Pair{String,Any}("active_constraints", [3, 6]))
 Dict{String,Any}(Pair{String,Any}("flexibility_index", 9.58983),Pair{String,Any}("model", Feasibility problem with:
 * 6 linear constraints
 * 6 variables
Solver is Pavito),Pair{String,Any}("active_constraints", [1, 5]))
```
"""
function rankinequalities(m::Model; max_ranks::Int = 5, suppress_warnings::Bool = true, U::Int = 10000, diag::Bool = false, active_constr::Bool = false, real_recourse_dim::Int = -1, conic_δ::Bool = false)
    # Check the max_ranks
    if max_ranks < 1
        error("max_ranks set to $max_ranks, but must be an integer greater than 0.")
    end

    # Initialize data
    m_copy = deepcopy(m)
    ranked_data = Vector{Dict}(undef, max_ranks)
    ranked_constrs = []
    inactives = []
    counter = 0

    # Get the inequality information
    constr_bounds = JuMP.prepConstrBounds(m_copy)
    inequal_inds = findall(constr_bounds[1] .!= constr_bounds[2])

    # Iteratively solve the flexibility index problem and extract data
    for i = 1:max_ranks
        status = solve(m_copy, suppress_warnings = suppress_warnings, U = U, diag = diag, active_constr = active_constr, real_recourse_dim = real_recourse_dim, conic_δ = conic_δ, inactives = inactives)
        if status == :Optimal
            ranked_constrs = [findall(inequal_inds .== getactiveconstraints(m_copy)[j])[1] for j = 1:length(getactiveconstraints(m_copy))]
            inactives = unique([inactives; ranked_constrs])
            ranked_data[i] = Dict("flexibility_index" => getflexibilityindex(m_copy), "active_constraints" => getactiveconstraints(m_copy), "model" => m_copy)
            counter += 1
        else
            break
        end
        if length(inactives) == length(inequal_inds)
            break
        end
    end

    # Return the results
    return ranked_data[1:counter]
end

"""
    getconfidencelevel(m::Model)
Return the confidence level provided that the flexibility model `m` has been solved with an
ellipsoidal uncertainty set. This is equivalent to the quantile associated with the optimized
uncertainty set. Note that this assumes a multivariate Gaussian distribution with mean
`FlexibilityData.RVmeans` and covariance `FlexibilityData.covariance`.

**Arguments**
- `m::Model` The flexibility model.

```julia
julia> getconfidencelevel(m)
0.5370703369769008
```
"""
function getconfidencelevel(m::Model)
    flex_data = getflexibilitydata(m)
    if flex_data.uncertainty_set.name == :Ellipsoid && flex_data.flexibility_index != nothing
        return cdf(Chisq(flex_data.numRVs), flex_data.flexibility_index)
    elseif flex_data.flexibility_index == nothing
        error("Cannot compute the confidence level without first solving the flexibility model.")
    else
        error("Invalid uncertainty set, must solve the flexibility model with an ellipsoidal set.")
    end
end

"""
    findstochasticflexibility(m::Model; num_pts::Int = 10000, toler::Number = 1e-5, solver = Clp.ClpSolver(), only_positive::Bool = false, use_vulnerability_model::Bool = false, use_flexibility_index::Bool = false, seed::Int = -1)
Returns the estimated stochastic flexibility index that is evaluated via Monte Carlo sampling. At default
this estimation is carried out by evaluating the feasibility of each Monte Carlo sample. The samples are
generating from a multivariate Gaussian distribution with mean `FlexibilityData.RVmeans` and covariance
`FlexibilityData.covariance`. The vulnerability model also tests the feasibility of the samples, but does
so in one large optimization problem instead of evaluating each sample indiviually. The optimized flexibility
index can also be used to reduce the number of samples that need to be evaluated.

**Arguments**
- `m::Model` The flexibility model.

**Keyword Arguments**
- `num_pts::Int = 10000` Number of Monte Carlo samples.
- `toler::Number = 1e-5` The feasibility check tolerance.
- `solver = Clp.ClpSolver()` The solver, any LP or NLP solver should work.
- `only_positive::Bool = false` Indicates if only positive samples should be used.
- `use_vulnerability_model::Bool = false` Indicates if the vulnerability model should be used.
- `use_flexibility_index::Bool = false` Indicates if the optimal flexibility index should be used.
- `seed::Int = -1` Random seed for sample collection, any negative value will turn off the random seed.

```julia
julia> findstochasticflexibility(m)
0.9687

julia> findstochasticflexibility(m, use_vulnerability_model = true)
0.9705

julia> findstochasticflexibility(m, num_pts = 5000, use_flexibility_index = true)
0.973
```
"""
function findstochasticflexibility(m::Model; num_pts::Int = 10000, toler::Number = 1e-5, solver = Clp.ClpSolver(), only_positive::Bool = false, use_vulnerability_model::Bool = false, use_flexibility_index::Bool = false, seed::Int = -1)
    # Parse data
    flex_data = getflexibilitydata(m)

    if length(flex_data.covariance) == 0
        error("The covariance matrix needs to be set.")
    end

    # Load in the neceassary data
    input_dict = MakeInputDict(m)
    θ_nom = input_dict["theta_nom"]

    # Parse the necessary data dimensions
    n_z = input_dict["dims"][1]
    n_θ = input_dict["dims"][2]
    n_x = input_dict["dims"][3]
    n_f = input_dict["dims"][4]
    n_h = input_dict["dims"][5]

    # Setup a multivariate normal distribution and Monte Carlo samples
    if seed >= 0
        Random.seed!(seed)
    end
    d = MvNormal(Vector{Float64}(flex_data.RVmeans), Matrix{Float64}(getcovariance(m)))
    if only_positive
        samples = Array{Float64}(undef, flex_data.numRVs, num_pts)
        not_done = true
        counter = 1
        while not_done
            sample = rand(d, 1)
            if all(sample .>= 0)
                samples[:, counter] = sample
                counter += 1
            end
            if counter > num_pts
                not_done = false
            end
        end
    else
        samples = rand(d, num_pts)
    end

    # Determine which points are outside of the set and check for F
    F = getflexibilityindex(m)
    if use_flexibility_index && F == nothing
        @warn "Flexibility index hasn't yet been computed. Setting use_flexibility_index = false."
        use_flexibility_index = false
    elseif use_flexibility_index && flex_data.uncertainty_set.name == :Ellipsoid
        inv_covar = inv(getcovariance(m))
        outside_set = [(samples[:, k] - θ_nom)' * inv_covar * (samples[:, k] - θ_nom) > F for k = 1:num_pts]
    elseif use_flexibility_index && flex_data.uncertainty_set.name == :PNorm
        outside_set = [norm(samples[:, k] - θ_nom, flex_data.uncertainty_set.p) > F for k = 1:num_pts]
    elseif use_flexibility_index && flex_data.uncertainty_set.name == :Hyperbox
        outside_set = Vector{Bool}(undef, num_pts)
        for k = 1:num_pts
            all_inside = true
            for i = 1:flex_data.numRVs
                if samples[i, k] < θ_nom[i] - flex_data.uncertainty_set.neg_dev[i] * F || samples[i, k] > θ_nom[i] + flex_data.uncertainty_set.pos_dev[i] * F
                    all_inside = false
                end
            end
            outside_set[k] = !all_inside
        end
    end

    # Determine the points to test
    if use_flexibility_index
        test_pts = samples[:, outside_set]
        num_test_pts = size(test_pts)[2]
    else
        num_test_pts = num_pts
        test_pts = samples
    end

    # Stop if all the test points are feasible.
    if num_test_pts == 0
        return 1.0
    end

    # Setup the model to be solved
    m_solve = Model(solver = solver)

    if use_vulnerability_model
        # Initialize variables
        @variable(m_solve, θ[1:n_θ, 1:num_test_pts])
        @variable(m_solve, u[1:num_test_pts] >= 0)
        if n_z != 0
            @variable(m_solve, z[1:n_z, 1:num_test_pts])
        end
        if n_x != 0
            @variable(m_solve, x[1:n_x, 1:num_test_pts])
        end

        # Set the objective
        @objective(m_solve, Min, 1 / num_test_pts * sum(u[k] for k = 1:num_test_pts))

        # Add the constraints
        exprs = AddSystemExpressions(m_solve, input_dict, num_test_pts)
        @constraint(m_solve, f[j = 1:n_f, k = 1:num_test_pts], exprs[1][j, k] <= u[k])
        if n_h != 0
            @constraint(m_solve, h[i = 1:n_h, k = 1:num_test_pts], exprs[2][i, k] == 0)
        end
    else
        # Initialize variables
        @variable(m_solve, θ[1:n_θ])
        @variable(m_solve, u)
        if n_z != 0
            @variable(m_solve, z[1:n_z])
        end
        if n_x != 0
            @variable(m_solve, x[1:n_x])
        end

        # Set the objective
        @objective(m_solve, Min, u)

        # Add the constraints
        exprs = AddSystemExpressions(m_solve, input_dict)
        @constraint(m_solve, f[j = 1:n_f], exprs[1][j] <= u)
        if n_h != 0
            @constraint(m_solve, h[i = 1:n_h], exprs[2][i] == 0)
        end
    end

    # Solve the model(s) and compute the SF
    if use_vulnerability_model
        for i = 1:n_θ
            for k = 1:num_test_pts
                JuMP.fix(θ[i, k], test_pts[i, k])
            end
        end
        status = solve(m_solve)
        if status == :Optimal
            return 1 - sum(getvalue(u) .>= toler) / num_pts
        else
            error("Problem not solved to optimality, cannot estimate SF.")
        end
    else
        infeasible_counter = 0
        feasible_results = zeros(Bool, num_test_pts)
        for k = 1:num_test_pts
            for i = 1:n_θ
                JuMP.fix(θ[i], test_pts[i, k])
            end
            status = solve(m_solve)
            if status == :Optimal
                if getvalue(u) <= toler
                    feasible_results[k] = true
                end
            else
                infeasible_counter += 1
            end
        end
        if infeasible_counter != 0
            @warn "Not all scenario subproblems not solved to optimality, estmation of SF might not be correct."
        end
        if use_flexibility_index
            return (sum(feasible_results) + num_pts - num_test_pts) / (num_pts - infeasible_counter)
        else
            return sum(feasible_results) / (num_pts - infeasible_counter)
        end
    end
end


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 44c12b60-97e4-11eb-23f7-eb2bf0847809
using Plots, SSpline, SpectralVis

# ╔═╡ d5fbb184-9eee-41db-8180-1cf5c5b42db7
plotlyjs()

# ╔═╡ 26ffc9f9-e1ec-4187-9eb6-4b0c3fc04223
begin
	x=[10,20,30,40,50]
	y=[19,30,30,15,10]
end

# ╔═╡ 1bb1f1e1-b366-4b31-873b-e4e78d1dca43
marsspec=mars_sky_opportunity()

# ╔═╡ 4f6b10bb-810d-49c4-ab9e-97c935e94504
marsspline=extrap(spline3(marsspec.λ,marsspec.l,:natural),380:830,:linear)

# ╔═╡ e2292549-9f32-427d-aa5b-1b816374caa1
marssky= interp(marsspline,380:1:830)

# ╔═╡ bc9a854d-4fc5-47da-8f09-1e21afe70522
begin
	scatter(marsspec.λ,marsspec.l,markersize=3,yrange=0.0:1.0)
	plot!(marssky[1],marssky[2])
end

# ╔═╡ ae0407ff-283a-4bf8-81ea-f0ac027aab8f
begin
	ms=LSpec(marssky[1],marssky[2])
	mr=RSpec(marssky[1],marssky[2])
end

# ╔═╡ fceaba8b-3c12-4883-8140-0930a75a590e
begin
	two=ms*mr
	three=ms*mr*mr*mr*mr*mr*mr*mr*mr
	plot(ms.λ,ms.l)
	plot!(two.λ,two.l)
	plot!(three.λ,three.l)
end

# ╔═╡ 208ed350-6c4e-40b2-9b3a-0a821b18a778

# ╔═╡ c1485749-eb3b-4f51-8814-416a39e7f5e0
begin
	lin=spline1(x,y)
	linint=interp(lin,10:.5:50)
	cub=spline3(x,y,:natural)
	cubint=interp(cub,10:.5:50)
	cubext=extrap(cub,0:60,:linear)
	extint=interp(cubext,0:.5:60)
end

# ╔═╡ b1299c43-c5a7-434e-8543-1715e23dc33f
begin
	scatter(x,y,color=:black,markersize=3,label="points")
	plot!(linint[1],linint[2],color=:grey,style=:dash,label="linear")
	plot!(cubint[1],cubint[2],color=:cyan,label="cubic nat.")
	plot!(extint[1],extint[2],color=:orange,linestyle=:dash,label="lin. extrap.")
end

# ╔═╡ f536734c-b7c0-4589-9be6-f4abea07faa8
function extrapo(Spl,xrange,extr::Symbol)
    envx=collect(xrange)
    xs=0.0 # start of extrapolatin range
    xe=0.0 # end of extrapolation range
    envx[1] < Spl.x[1] ? xs=envx[1] : xs=Spl.x[1]
    envx[1] < Spl.x[1] ? xs=envx[1] : xs=Spl.x[1]
    envx[end] > Spl.x[end] ? xe=envx[end] : xe=Spl.x[end]

    x=copy(Spl.x)
    a=copy(Spl.a)
	nf=nextfloat(Float64(x[end]))

    if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
        b=copy(Spl.b)
    end
    if typeof(Spl) in (QuadraticSpline, CubicSpline)
        c=copy(Spl.c)
    end
    if typeof(Spl) == CubicSpline
        d=copy(Spl.d)
    end
	if extr == :zero
		if xs < x[1]
			pushfirst!(x,xs)
			pushfirst!(a, 0.0)
			if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
				pushfirst!(b,0.0)
			end
			if typeof(Spl) in (QuadraticSpline,CubicSpline)
				pushfirst!(c,0.0)
			end
			if typeof(Spl) == CubicSpline
				pushfirst!(d,0.0)
			end
		end
		if xe ≥ x[end]
			if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
				x[end]=nf
				push!(x,xe)
				push!(a,0.0)
				push!(b,0.0)
			end
			if typeof(Spl) in (QuadraticSpline,CubicSpline)
				push!(c,0.0)
			end
			if typeof(Spl) == CubicSpline
				push!(d,0.0)
			end
		end
	elseif extr == :boundary
		de=x[end]-x[end-1]
		if xs < x[1]
			pushfirst!(x,xs)
			pushfirst!(a, a[1])
			if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
				pushfirst!(b,0.0)
			end
			if typeof(Spl) in (QuadraticSpline,CubicSpline)
				pushfirst!(c,0.0)
			end
			if typeof(Spl) == CubicSpline
				pushfirst!(d,0.0)
			end
		end
		if xe ≥ x[end]
			if typeof(Spl) == LinearSpline
				push!(a,Spl.b[end]*de+Spl.a[end])
				push!(b,0.0)
				x[end]=nf
			end
			if typeof(Spl) == QuadraticSpline
				push!(a,c[end]*de^2 + b[end]*de + a[end])
				push!(b,0.0)
				push!(c,0.0)
				x[end]=nf
			end
			if typeof(Spl) == CubicSpline
				push!(a,d[end]*de^3 + c[end]*de^2 + b[end]*de + a[end])
				push!(b,0.0)
				push!(c,0.0)
				push!(d,0.0)
				x[end]=nf
			end
		end
		push!(x,xe)
	elseif extr == :linear
		if xs < x[1]
			pushfirst!(x,xs)
			#
			pushfirst!(a,a[1]+b[1]*(x[1]-x[2]))
			if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
				pushfirst!(b,b[1])
			end
			if typeof(Spl) in (QuadraticSpline,CubicSpline)
				pushfirst!(c,0.0)
			end
			if typeof(Spl) == CubicSpline
				pushfirst!(d,0.0)
			end
		end
		if xe ≥ x[end]
			if typeof(Spl) == LinearSpline
				push!(a,a[end]+b[end]*(x[end]-x[end-1]))
				push!(b,b[end])
			end
			if typeof(Spl) == QuadraticSpline
				push!(c,0.0)
			end
			if typeof(Spl) == CubicSpline
				de=x[end]-x[end-1]
				push!(a,d[end]*de^3 + c[end]*de^2 + b[end]*de + a[end])
				push!(b,3*d[end]*de^2 + 2*c[end]*de + b[end])
				push!(c,0.0)
				push!(d,0.0)
				x[end]=nf
			end
		end
		push!(x,xe)
    elseif extr == :quadratic
        if xs < x[1]
            pushfirst!(x,xs)
            if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
                pushfirst!(b,x[1]*((a[3]-a[2])+x[2]*(a[1]-b[3])+x[3]*(a[2]-a[1]))/((x[1]-x[2])*(x[1]-x[3])*(x[2]-x[3])))
                pushfirst!(a,a[1]/(b[1]*(x[1]-xs)))
            end
            if typeof(Spl) in (QuadraticSpline, CubicSpline)
                pushfirst!(c,0.0)
            end
            if typeof(Spl) == CubicSpline
                pushfirst!(d,0.0)
            end
        end
        if xe ≥ x[end]
            push!(x,xe)
            if typeof(Spl) in (LinearSpline, QuadraticSpline, CubicSpline)
                push!(b,x[end-2]*((a[end]-a[end-1])+x[end-1]*(a[end-2]-b[end])+x[end]*(a[end-1]-a[end-2]))/((x[end-2]-x[end-1])*(x[end-2]-x[end])*(x[end-1]-x[end])))
                push!(a,a[end]/(b[end]*(x[end]-xs)))
            end
            if typeof(Spl) in (QuadraticSpline, CubicSpline)
                push!(c,0.0)
            end
            if typeof(Spl) == CubicSpline
                push!(d,0.0)
            end
        end
    else error("Extrapolation type $extr does not exist.")
    end
    if typeof(Spl) == NearestNeighborSpline
        NearestNeighborSpline(a,x)
    elseif typeof(Spl) == LinearSpline
        LinearSpline(b,a,x)
    elseif typeof(Spl) == QuadraticSpline
        QuadraticSpline(c,b,a,x)
    else
        CubicSpline(d,c,b,a,x)
    end
end

# ╔═╡ Cell order:
# ╠═44c12b60-97e4-11eb-23f7-eb2bf0847809
# ╠═d5fbb184-9eee-41db-8180-1cf5c5b42db7
# ╠═26ffc9f9-e1ec-4187-9eb6-4b0c3fc04223
# ╠═1bb1f1e1-b366-4b31-873b-e4e78d1dca43
# ╠═4f6b10bb-810d-49c4-ab9e-97c935e94504
# ╠═e2292549-9f32-427d-aa5b-1b816374caa1
# ╠═bc9a854d-4fc5-47da-8f09-1e21afe70522
# ╠═ae0407ff-283a-4bf8-81ea-f0ac027aab8f
# ╠═fceaba8b-3c12-4883-8140-0930a75a590e
# ╠═208ed350-6c4e-40b2-9b3a-0a821b18a778
# ╠═c1485749-eb3b-4f51-8814-416a39e7f5e0
# ╠═b1299c43-c5a7-434e-8543-1715e23dc33f
# ╠═f536734c-b7c0-4589-9be6-f4abea07faa8


using OrdinaryDiffEq, Test
function lorenz(du,u,p,t)
 du[1] = 10.0(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end
u0 = [1.0;0.0;0.0]
tspan = (0.0,1.0)
prob = ODEProblem(lorenz,u0,tspan)
sol = solve(prob,Tsit5())
@test_broken @inferred solve(prob,Tsit5())

function f(du,u,p,t)
  du[1] = p.a
  du[2] = p.b
end

const alg = Tsit5()

function solve_ode(f::F, p::P) where {F,P}

  tspan = (0., 1.0)
  Δt = tspan[2] - tspan[1]
  dt = 1/252
  nodes = Int(ceil(Δt / dt) + 1)
  t = T = [tspan[1] + (i - 1) * dt for i = 1:nodes]

  # if I do not set {true}, prob type Any...
  prob = ODEProblem{true}(f, [0., 0.], tspan, p)
  # prob = ODEProblem(f, [0., 0.], tspan, p)

  prob_func = (prob, i, repeat) -> begin
    remake(prob, tspan = (T[i + 1], t[1]))
  end

  # ensemble problem
  odes = EnsembleProblem(prob, prob_func = prob_func)

  sol = OrdinaryDiffEq.solve(
    odes, OrdinaryDiffEq.Tsit5(), OrdinaryDiffEq.EnsembleThreads(),
    trajectories = nodes - 1, saveat = -dt
  )

  return sol
end
@test_broken @inferred solve_ode(f, (a = 1, b = 1))

using StochasticDiffEq, Test
u0=1/2
ff(u,p,t) = u
gg(u,p,t) = u
dt = 1//2^(4)
tspan = (0.0,1.0)
prob = SDEProblem(ff,gg,u0,(0.0,1.0))
sol = solve(prob,EM(),dt=dt)
@test_broken @inferred solve(prob,EM(),dt=dt)


#=
Here are two approaches to deal with convolutions in the inverse fitting approach.
Taking the convolution of the IRF and the model appears to be computationally expensive.
Each iteration in the least squares fit has to compare the new fit with the data after convolving the model w/IRF.

Here I am trying to take advantage of the fact that the fft(IRF) is fixed for each step in the iteration and there is no reason to repeatedly calculate it.
The challenge is that the conv function from the FFTW library is very optimzed as it is a wrapper of a C/C++ library that is highly optimized.
Even still, there is an opportunity to plan out the FFT beforehand using the fact that the u and v in conv(u, v) are of the same length for our problem.
This allows to calculate the linear operator by the FFT needed to compute fft for a given set of dimensions.
In other words we can calculate a plan_fft (call this p) operator and then calculate the fft by u*p and v*p. If u is the IRF this value is constant for each iteration in the least squares.

=#

model(t, p) = p[1] * exp.(-p[2] *t)

tdata = 1:0.01:10

ydata = model(tdata, [1.0, 2.0]) + randn(length(tdata))

IRF = [Vector(1:500); Vector(500:-1:100)]

ydata = conv(ydata, IRF)

function convFIT(tdata, p, IRF)

    Rt = model(tdata, p)

    Rt = conv(Rt, IRF)
end

function fit(tdata, ydata, IRF)
    curve_fit((tdata, p) -> convFIT(tdata, p, IRF), tdata, ydata, [0.5, 0.5])
end

function convFIT1(tdata, p, pl, IRF_fft)
    Rt = model(tdata, p)
    dims = 2*length(Rt) -1

    Rt_fft = convert(Array{Complex{Float64},1}, [Rt ; zeros(Float64, dims-length(Rt))])
    r = similar(Rt_fft)
    i = similar(Rt_fft)
    c = similar(Rt_fft)

    mul!(r, pl, Rt_fft)
    mul!(i, pl, IRF_fft)

   ldiv!(c, pl, (r.*i))
   return real(c)
end

function fit2(tdata, ydata, IRF)

    dims = 2*length(IRF) -1

    IRF_fft = convert(Array{Complex{Float64},1}, [IRF ; zeros(Float64, dims-length(IRF))])

    pl = plan_fft(IRF_fft)

    curve_fit((tdata, p) -> convFIT1(tdata, p, pl, IRF_fft), tdata, ydata, [0.5, 0.5])
end

#=

############# Run these commands after running previous lines ################
julia> @benchmark f = fit(tdata,ydata,IRF)
BenchmarkTools.Trial:
  memory estimate:  132.05 MiB
  allocs estimate:  837915
  --------------
  minimum time:     792.865 ms (0.64% GC)
  median time:      801.574 ms (0.64% GC)
  mean time:        801.359 ms (0.64% GC)
  maximum time:     807.842 ms (0.63% GC)
  --------------
  samples:          7
  evals/sample:     1

 julia> @benchmark f2 = fit2(tdata,ydata,IRF)
  BenchmarkTools.Trial:
    memory estimate:  255.31 MiB
    allocs estimate:  1578872
    --------------
    minimum time:     212.654 ms (4.39% GC)
    median time:      214.032 ms (4.38% GC)
    mean time:        218.092 ms (4.41% GC)
    maximum time:     242.318 ms (3.91% GC)
    --------------
    samples:          23
    evals/sample:     1
  =#

  su = size(IRF)
  outsize = 2 .* su .-1
 # nffts = nextfastfft(outsize)

  upad = similar(IRF, outsize)
  vpad = similar(IRF, outsize)
  fill!(upad, zero(eltype(upad)))
  fill!(vpad, zero(eltype(vpad)))

  copyto!(upad, IRF)

  p! = plan_fft!(convert(Array{Complex{Float64}}, upad))
  upad = p! * upad

function conv_DT(t, data::input_data, upad, p!, vpad)

	RtDT = Array{Float64}(undef, length(0:0.01:10))
    RtDT = DT_model(0:0.01:10, [0.01,10.0], 1., 1., 1.)

    copyto!(vpad, RtDT)

    vpad = p! * vpad
    vpad .*= upad
    real(ifft!(vpad))

	Rt = Rt./maximum(Rt)

	tidx = findfirst(x -> x == t[1], data.t)

	convDT = convDT[tidx:tidx+length(t)-1]
	return log.(convDT)
end

function pad_fft(u, v)

    su = size(u)
    sv = size(v)
    outsize = su .+ sv .-1
    nffts = nextfastfft(outsize)

    upad = similar(u, nffts)
    vpad = similar(v, nffts)
    fill!(upad, zero(eltype(upad)))
    fill!(vpad, zero(eltype(vpad)))

    copyto!(upad, u)
    copyto!(vpad, v)

    p! = plan_fft!(convert(Array{Complex{Float64}}, upad))
    upad = p! * upad

function testing(upad, vpad, p!)

    vpad = p! * vpad
    vpad .*= upad
    real(ifft!(vpad))
end

outsize = 2 * length(data.IRF) .- 1
nffts = nextfastfft(outsize)

upad = zeros(Float64, nffts)
vpad = zeros(Float64, nffts)
convpad = zeros(Float64, nffts)

copyto!(upad, data.IRF)
pl = plan_rfft(upad)

upadRfft = zeros(Complex{Float64}, 1 .+ div(length(upad), 2, RoundDown))
upadRfft = pl * upad

RtDT = Array{Float64}(undef, length(data.IRF))

vpadRfft = zeros(Complex{Float64}, 1 .+ div(length(upad), 2, RoundDown))
u_vpadRfft = zeros(Complex{Float64}, 1 .+ div(length(upad), 2, RoundDown))

function conv_DT(t, β::Array{Float64,1}, data::input_data,
    vpad, RtDT, pl, vpadRfft, upadRfft, u_vpadRfft,
    convpad
    )

RtDT = DT_model(data.t, β, data.ρ, data.nmed, data.ndet)
copyto!(vpad, RtDT)

vpadRfft = pl * vpad
u_vpadRfft = upadRfft .* vpadRfft
convpad = irfft(u_vpadRfft, length(upad))
convpad = convpad./maximum(convpad)

tidx = findfirst(x -> x == t[1], data.t)

convpad = convpad[tidx:tidx+length(t)-1]
return log.(convpad)
end

#=@benchmark conv_DT(data.t[ind1:ind2], [0.1,10.], data, vpad, RtDT, pl, vpadRfft, upadRfft, u_vpadRfft, convpad)
BenchmarkTools.Trial:
memory estimate:  401.22 KiB
allocs estimate:  134
--------------
minimum time:     194.223 μs (0.00% GC)
median time:      199.783 μs (0.00% GC)
mean time:        241.692 μs (2.87% GC)
maximum time:     5.035 ms (0.00% GC)
--------------
samples:          10000
evals/sample:     1

old version
@benchmark conv_DT(data.t[ind1:ind2], [0.1,10.], data)
BenchmarkTools.Trial:
  memory estimate:  628.81 KiB
  allocs estimate:  184
  --------------
  minimum time:     296.018 μs (0.00% GC)
  median time:      308.628 μs (0.00% GC)
  mean time:        423.330 μs (2.68% GC)
  maximum time:     7.641 ms (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1
=#

function conv_DT(t, β::Array{Float64,1}, data::input_data,
    vpad, RtDT, pl, vpadRfft, upadRfft, u_vpadRfft,
    convpad
    )

RtDT = DT_model(data.t, β, data.ρ, data.nmed, data.ndet)
copyto!(vpad, RtDT)

vpadRfft = pl * vpad
u_vpadRfft = upadRfft .* vpadRfft
convpad = irfft(u_vpadRfft, length(upad))
convpad = convpad./maximum(convpad)

tidx = findfirst(x -> x == t[1], data.t)

convpad = convpad[tidx:tidx+length(t)-1]
return log.(convpad)
end

function getfit(input_data, model_params; alpha=0.1)

    ## get FFT padding and plans
    outsize = 2 * length(data.IRF) .- 1
    nffts = nextfastfft(outsize)

    upad = zeros(Float64, nffts)
    vpad = zeros(Float64, nffts)
    convpad = zeros(Float64, nffts)

    copyto!(upad, data.IRF)
    pl = plan_rfft(upad)

    upadRfft = zeros(Complex{Float64}, 1 .+ div(length(upad), 2, RoundDown))
    upadRfft = pl * upad

    RtDT = Array{Float64}(undef, length(data.IRF))

    vpadRfft = zeros(Complex{Float64}, 1 .+ div(length(upad), 2, RoundDown))
    u_vpadRfft = zeros(Complex{Float64}, 1 .+ div(length(upad), 2, RoundDown))

    #get fitting window
    counts = input_data.DTOF
    maxvalue, maxindex = findmax(counts)

    ind1 = findfirst(x -> x > maxvalue*model_params.risefactor, counts)
    ind2 =  findlast(x -> x >maxvalue*model_params.tailfactor, counts)

    weights = 1 ./ input_data.DTOFerrors

    fit = curve_fit((t, β) -> conv_DT(t, β, input_data, vpad, RtDT, pl, vpadRfft, upadRfft, u_vpadRfft, convpad), input_data.t[ind1:ind2], log.(input_data.DTOF[ind1:ind2]),
    model_params.initparams, lower=model_params.lb, upper=model_params.ub)

    perrors = estimate_errors(fit,alpha)
    marginerror = margin_error(fit,alpha)

    chisq = sum(fit.resid.^2)
    chisqdist = Distributions.Chisq(2)
    probability = Distributions.ccdf(chisqdist,chisq)
    return fitresult(input_data.t[ind1:ind2],log.(input_data.DTOF[ind1:ind2]),input_data.DTOFerrors, DT_model, fit.resid,fit.param,marginerror,perrors,2,chisq,probability)
end


#!/usr/bin/env julia

export pe1

pe1(n::Integer=1000)=union(3:3:(n-1),5:5:(n-1))|>sum

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe1()|>print
end


"""
Provides functionality for computing the Cuthill-McKee graph permutation on a sparse matrix in CSC format.
"""
module CuthillMcKee

import SparseArrays: SparseMatrixCSC

export rcmpermute, symrcm

"""
	 rcmpermute(A)
Computes and applies the reverse Cuthill-McKee permutation to the structurally-symmetric sparse matrix `A`. No checks are performed to ensure that the matrix is indeed structurally symmetric.
"""
function rcmpermute(A::SparseMatrixCSC)
	adj = coladjacency(A)
	sdeg = sortperm(coldegrees(A))
	p = symrcm(adj, sdeg)
	return A[p,p]
end

"""
	 symrcm(A,[ rev, inv, warnunconnected])
Computes the Cuthill-McKee permutation of the structurally-symmetric sparse matrix `A`. The default is to compute the reverse permutation, i.e. `rev` defaults to `true`. Optionally, when `inv` is set to `true`, the permutation is inverted. Per default the warning displaying the number of unconnected regions in the graph is disabled, to enable, set `warnunconnected` to `true`. No checks are performed to ensure that the matrix is indeed structurally symmetric.
"""
function symrcm(A::SparseMatrixCSC, args...)
	adj = coladjacency(A)
	sdeg = sortperm(coldegrees(A))
	return symrcm(adj, sdeg, args...)
end

"""
	 symrcm(adjac, sdegr,[ rev, inv, warnunconnected])
Computes the Cuthill-McKee permutation of the graph given as the adjacency list `adjac`. The adjacency list is required to be of type `Vector{Vector{Int}}` such that `adjac[i]` is a vector of the adjacent vertex numbers to vertex `i` in order of ascending vertex degree. The additional required argument `sdegr` is an array of vertex numbers in order of ascending vertex degree. The default is to compute the reverse permutation, i.e. `rev` defaults to `true`. Optionally, when `inv` is set to `true`, the permutation is inverted. Per default the warning displaying the number of unconnected regions in the graph is disabled, to enable, set `warnunconnected` to `true`.
"""
function symrcm(adjac::Array{Array{Int64,1}}, sdegr::Array{Int64,1}, rev=true, invert=false, warnunconnected=false)
	# Based on the description of the RCM algorithm by Ciprian Zavoianu.
	# See: http://ciprian-zavoianu.blogspot.com/2009/01/project-bandwidth-reduction.html
	nverts = length(adjac)
	checkbounds(sdegr,nverts)
	sd = copy(sdegr)
	R = Array{Int}(undef, 0)
	Q = Array{Int}(undef, 0)
	P = 0
	C = 0
	unconnected = 0
	inserted = falses(nverts)
	while !isempty(sd)
		found = false
		while !isempty(sd)
			P = popfirst!(sd)
			if !inserted[P]
				push!(R, P)
				inserted[P] = true
				found = true
				break
			end
		end
		if found
			empty!(Q)
			append!(Q, adjac[P])
			while !isempty(Q)
				while !isempty(Q)
					C = popfirst!(Q)
					if !inserted[C]
						push!(R, C)
						inserted[C] = true
						append!(Q, adjac[C])
						break
					end
				end
			end
		end
		if length(R) == nverts
			break
		else
			unconnected += 1
		end
	end
	if unconnected > 0 && warnunconnected
		warn("There are $(unconnected) unconnected regions.")
	end
	if rev
		reverse!(R)
	end
	if invert
		L = zeros(Int64,nverts)
		for i = 1:nverts
			L[R[i]] = i
		end
		return L
	else
		return R
	end
end

function coldegrees(A::SparseMatrixCSC)
	# assumes A is structurally symmetric
	cptr = A.colptr
	ncols = length(cptr)-1
	degr = zeros(Int,ncols)
	@inbounds for j = 1:ncols
		degr[j] = coldeg(A, j)
	end
	return degr
end

function coldeg(A::SparseMatrixCSC, j::Int)
	# assumes A is structurally symmetric
	cptr = A.colptr
	return cptr[j+1]-cptr[j]
end

function coladjacency(A::SparseMatrixCSC)
	# assumes A is structurally symmetric
	cptr = A.colptr
	rval = A.rowval
	ncols = length(cptr)-1
	adjac = Vector{Vector{Int}}(undef, ncols)
	sbyf = let A = A
		j->coldeg(A, j)
	end
	for j = 1:ncols
		strt = cptr[j]
		jdeg = coldeg(A, j)
		jadj = Vector{Int}(undef, jdeg)
		for i = 1:jdeg
			jadj[i] = rval[strt+i-1]
		end
		sort!(jadj, by=sbyf)
		adjac[j] = jadj
	end
	return adjac
end

end


@test super(BasicKalmanFilter) == LinearKalmanFilter

s = State([1.,0],[0.5 0; 0 1e-10])

@test typeof(s) == State{Float64}
@test s.x == [1,0]
@test s.p == [0.5 0; 0 1e-10]

o = convert(Observation,2.0)

@test typeof(o) == Observation{Float64}
@test o.y == [2]

x0 = State([1.,0],[0.5 0; 0 1e-10])
a = [1 0.01; 0.01*(-1) 1]
g = [0 0; 0 1]
q = [1e-10 0; 0 1e-10]
f = LinearModel(a,g,q)
h = [1 0]
r = [0.5]'
z = LinearObservationModel(h,r)

dt = 0.01
t = [0:dt:10]
x = zeros(length(t))
m = cos(t)

kf0 = BasicKalmanFilter(x0,f,z)
kf = kf0
x[1] = kf.x.x[1]

y = map(t->Observation(m[t]+kf.z.r[1]*randn(1)),1:length(t))

kf1 = predictupdate(kf,y[1])
predictupdate!(kf,y[1])
@test kf1.x == kf.x

for i in 2:length(t)-1
    kf2 = predict(kf)
    kf = update(kf2,y[i])
    x[i+1] = kf.x.x[1]
end

@test predictupdate(kf0,y[1]).x == update(predict(kf0),y[1]).x


#! Transformations of the meshes. Uses:
#! LinearAlgebra.UniformScaling for scaling
#! Rotation.jl for rotation
#! CoordinateTransformations.jl for translation
# Voir https://stackoverflow.com/questions/10546320/remove-rotation-from-a-4x4-homogeneous-transformation-matrix
# pour extraire la rotation et la translation depuis la matrice 4*4, puis transformer
# cette matrice 4*4 en CoordinateTransformations.jl transformations et après on ne gèrera
# que ces transformations là.

"""
    transform_mesh!(node::Node, transformation)

Add a new `CoordinateTransformations.jl` transformation to the node geometry
`transformation` field. The transformation is composed with the previous
transformation if any.

`transformation` must be a `CoordinateTransformations.jl` transformation.

It is also possible to invert a transformation using `inv` from
`CoordinateTransformations.jl`.

# Examples

```julia
using PlantGeom, MultiScaleTreeGraph, GLMakie, Rotations, CoordinateTransformations

file = joinpath(dirname(dirname(pathof(PlantGeom))), "test", "files", "simple_plant.opf")
opf = read_opf(file)

# Visualize the mesh as is:
viz(opf)

# Copy the OPF, and translate the whole plant by 15 in the y direction (this is in cm, the mesh comes from XPlo):
opf2 = deepcopy(opf)
transform!(opf2, x -> transform_mesh!(x, Translation(0, 15, 0)))
viz!(opf2) # Visualize it again in the same figure

# Same but rotate the whole plant around the X axis:
opf3 = deepcopy(opf)
transform!(opf3, x -> transform_mesh!(x, LinearMap(RotX(0.3))))
# NB: we use Rotations.jl's RotX here. Input in radian, use rad2deg and deg2rad if needed.
viz!(opf3)

# Same but rotate only the second leaf around the Z axis:
opf4 = deepcopy(opf)
# Build the meshes from the reference meshes (need it because we want the coordinates of the parent):
transform!(opf4, refmesh_to_mesh!)

# Get the second leaf in the OPF:
leaf_node = get_node(opf4, 8)

# Get the parent node (internode) Z coordinates:
parent_zmax = zmax(leaf_node.parent)

# Define a rotation of the mesh around the Z axis defined by the parent node max Z:
transformation = recenter(LinearMap(RotZ(1.0)), Point3(0.0, 0.0, parent_zmax))

# Update the transformation matrix of the leaf and its mesh:
transform_mesh!(leaf_node, transformation)

# Plot the result:
viz(opf)
viz!(opf4)
```
"""
function transform_mesh!(node::MultiScaleTreeGraph.Node, transformation)
    if node[:geometry] !== nothing
        node[:geometry].transformation = transformation ∘ node[:geometry].transformation
        # If the node has a computed mesh, update it:
        if node[:geometry].mesh !== nothing
            refmesh_to_mesh!(node)
        end
    end
end


#=
Given a string, find the longest palindromic contiguous substring. If there are more than one with the maximum length, return any one.

For example, the longest palindromic substring of "aabcdcb" is "bcdcb". The longest palindromic substring of "bananas" is "anana".
=#

function longestpalindromicsubstring(s::String)
    longestpalindrome = ""
    for index in eachindex(s)
        current_palindrome = longestpalindromecenter(s, index)
        if length(current_palindrome) >= length(longestpalindrome)
            longestpalindrome = current_palindrome
        end
    end
    return longestpalindrome
end

function longestpalindromecenter(s::String, index::Int)
    longestpalindrome = s[index:index]
    i = 1
    if index != 1 && index != length(s)
        while index-i >= 1 && index+i <= length(s) && ispalindrome(s[index-i:index+i])
            longestpalindrome = s[index-i:index+i]
            i = i + 1
        end
    end
    i = i-1
    if index+i < length(s)
        if ispalindrome(s[index-i:index+i+1])
            longestpalindrome = s[index-i:index+i+1]
        end
    end
    if index-i > 1
        if ispalindrome(s[index-i-1:index+i])
            longestpalindrome = s[index-i-1:index+i]
        end
    end
    return longestpalindrome
end

function ispalindrome(s::String)
    if isodd(length(s))
        half = floor(Int8, length(s)/2)
        is_palindrome = s[1:half] == reverse(s[half+2:end])
    else
        half = floor(Int8, length(s)/2)
        is_palindrome = s[1:half] == reverse(s[half+1:end])
    end
    return is_palindrome
end