"""
cohorts as 3 matrices: year_b - initial year, year_e - last year,  x - abundance, y - catch, w - weight (obsolete)
"""
mutable struct  cohorts_t
     year_b
    year_e
    x
     y
    w

end

"""
detail cohort representation
tmax - number of seasons
jmax - number of ages
yb - physical year assciated with cohort (when age=1)
jb - youngest age
l - length
jt - trusted age
xv,yv,wv - vectors with abundunce, catch, weight
"""
mutable struct cohort_v
    tmax
    jmax
    yb
    jb
    l
    jt
     xv::Array{Float64}
     yv::Array{Float64}
     wv::Array{Float64}
 end

"""
detail  representation of given cohorts
tmax - number of seasons
jmax - number of ages
cs -  vector of cohort of mutable struct cohort_v
"""
mutable struct cohorts_v
       tmax
    jmax
    cs::Array{cohort_v}
 end

"""
detail  representation of given cohorts
tmax - number of seasons
jmax - number of ages
cs -  vector of cohort of mutable struct cohort_v
"""

mutable struct cohorts_v1
    year_b
    tmax
    jmax
    cs::Array{cohort_v}
 end

"""
detail cohort representation
no- the number
yb - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
jbeg - initial age
len - length
jtrust - trusted age
xv,yv,wv - vectors with abundunce, catch, weight
"""
mutable struct cohort_v1
    no
    yb
    tmax
    jmax
    jbeg
    len
    jtrust
     xv::Array{Float64}
     yv::Array{Float64}
     wv::Array{Float64}
 end

"""
cohorts as 3 tables (=cohorts_m)
yb - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
x,y,w- tables with abundunce, catch, weight
"""
mutable struct  cohorts_t1
     year_b
   tmax
    jmax
    x
     y
    w
end

"""cohorts as cs: cohort_v1?
year_b - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
"""
mutable struct cohorts_v2
    year_b
    tmax
    jmax
    cs
end

"""
cohorts related matrix
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
mtx: x,y,w,m,f, etc.
"""
mutable struct csmtx
    year_b
    tmax
    jmax
    mtx
end

"""
detail mtx representation of a cohort as sequence of vecs
no- the number
yb - physical year associated with cohort (when age=1)
jbeg - youngest age
len - length
jtrust - trusted age
vec
"""
mutable struct cvec
    no
    yb
    jbeg
    len
    jtrust
    vec::Array{Float64}

 end

"""
cohort as vecs
"""
mutable struct csvecs
    year_b
    tmax
    jmax
    vecs
end

"""
cohorts_t1->cohorts_v2

cohorts_t1:
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
x,y,w -  matrices  with abundunce, catch, weight

cohorts_v2:
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
cs - cohorts of mutable struct cohort_v1

cohort_v1:
no- the number
yb - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
jbeg - initial age
len - length
jtrust - trusted age
xv,yv,wv- vectors with abundunce, catch, weight
tmax+jmax-1 -- number of related cohorts
"""
function cst2csv2(cs)
   x=cs.x
    y=cs.y
    w=cs.w
    tmax=cs.tmax
    jmax=cs.jmax
    year1=cs.year_b
    year0=year1-jmax+1
                #println(year0)
                #println(year1)

#    csv1=Array{cohort_v1}(tmax+jmax-1)
csv1=Array{cohort_v1}(undef,tmax+jmax-1)

    #csv1=Array{Any}(undef,tmax+jmax-1)

    for year in year0:year1-1
        i=year-year0+1
                #println(year)
                #println(i)

        csv1[i]=cohort_v1(i,year,tmax,jmax,
        jmax-i+1,i,jmax-i+1,
        map(j->x[j,jmax-i+j],1:i),
        map(j->y[j,jmax-i+j],1:i),
        map(j->w[j,jmax-i+j],1:i))
                #println(csv1[i])
    end

    year2=year1+tmax-jmax
                #println(year2)

    if tmax>=jmax
        for year in year1:year2
            i=year-year1+1
               #println(year)
               #println(i)
               #println(i-year0+year1)
            csv1[i-year0+year1]=cohort_v1(i-year0+year1,year,tmax,jmax,
            1,jmax,1,
            map(j->x[i+j-1,j],1:jmax),
            map(j->y[i+j-1,j],1:jmax),
            map(j->w[i+j-1,j],1:jmax))
               #println( csv1[i-year0+year1])
        end
    end

    year3=year2+jmax-1
                #println(year3)

    for year in year2+1:year3
            i=year-year1+1
                #println(year)
                #println(i)
                #println(i-year1+year2)
        csv1[i-year1+year2+1]=cohort_v1(i-year1+year2+1,year,tmax,jmax,
        1,tmax-i+1,1,
        map(j->x[i+j-1,j],1:tmax-i+1),
        map(j->y[i+j-1,j],1:tmax-i+1),
        map(j->w[i+j-1,j],1:tmax-i+1))
                #println( csv1[i-year0+year1])
    end

    cohorts_v2(year1,tmax,jmax,csv1)
end

"""
cohorts_v2->cohirts_t1

cohorts_t1:
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
x,y,w -  matrices  with abundunce, catch, weight

cohorts_v2:
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
cs - cohorts of mutable struct cohort_v1

cohort_v1:
no- the number
yb - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
jbeg - initial age
len - length
jtrust - trusted age
xv,yv,wv- vectors with abundunce, catch, weight
"""
function csv2cst(csv)

    year1=csv.year_b
    tmax=csv.tmax
    jmax=csv.jmax
    cs=csv.cs

    x=fill(1.0,tmax,jmax)
    y=fill(1.0,tmax,jmax)
    w=fill(1.0,tmax,jmax)

    ncs=tmax+jmax-2
    year0=year1-jmax+1
    year2=year1+tmax-jmax
    year3=year2+jmax-1

    #println(year0)
    #println(year1)
    #println(year2)
    #println(year3)

   #1
    for year in year0:year1-1
        i=year-year0+1
        #println(year)
        #println(i)

    cv1=cs[i]
       no=cv1.no
    yb=cv1.yb
    jbeg=cv1.jbeg
    len=cv1.len
        jtrust=cv1.jtrust
        xv=cv1.xv
    yv=cv1.yv
    wv=cv1.wv

     #println(' ',i,' ',no,' ',yb,' ',jbeg,' ',len)

        for j in 1:i
            #println(j,' ',jmax-i+j)
            x[j,jmax-i+j]=xv[j]
            y[j,jmax-i+j]=yv[j]
            w[j,jmax-i+j]=wv[j]
        end
    end
  #2
     if tmax>=jmax
        for year in year1:year2
            i=year-year0+1
            i1=year-year1+1
                 #println(year)
                 #println(i)

            cv1=cs[i]
       no=cv1.no
    yb=cv1.yb
    jbeg=cv1.jbeg
    len=cv1.len
        jtrust=cv1.jtrust
        xv=cv1.xv
    yv=cv1.yv
    wv=cv1.wv

               #println(i,' ',no,' ',yb,' ',jbeg,' ',len)

            for j in 1:jmax
                #println(i+j-1,' ',j)
            x[i1+j-1,j]=xv[j]
            y[i1+j-1,j]=yv[j]
            w[i1+j-1,j]=wv[j]

        end
    end
        end

  #3
  for year in year2+1:year3
            i=year-year0+1
        i1=year-year1+1

                #println(year)
                #println(i)

        cv1=cs[i]
       no=cv1.no
    yb=cv1.yb
    jbeg=cv1.jbeg
    len=cv1.len
        jtrust=cv1.jtrust
        xv=cv1.xv
    yv=cv1.yv
    wv=cv1.wv

             #println(i,' ',no,' ',yb,' ',jbeg,' ',len)

        for j in 1:tmax-i1+1
             #println(i1+j-1,' ',j)

        x[i1+j-1,j]=xv[j]
            y[i1+j-1,j]=yv[j]
            w[i1+j-1,j]=wv[j]
        end
    end

    cohorts_t1(year1,tmax,jmax,x,y,w)
end

"""
cohorts related matrix (of mutable struct csmtx: x,y,w,m,f,z,g)->sequence of vectors (of mutable struct csvecs)

csrm:cohorts related matrix (of mutable struct csrm)
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
x,y,w, m,f,g -  cohorts related matrices

cvec: detail cohort representation as vec
no- the number in the sequence
yb - physical year associated with cohort (when age=1)
jbeg - initial age
len - length
jtrust - trusted age
vec

csvecs: cohorts related mtx as a sequence of cvec
no- the number
yb - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
jbeg - initial age
len - length
jtrust - trusted age
xv,yv,wv- vectors with abundunce, catch, weight
tmax+jmax-1 -- number of related cohorts
"""
function csmtx2csvs(csrmtx)
   year1=csrmtx.year_b
    tmax=csrmtx.tmax
    jmax=csrmtx.jmax
    mtx=csrmtx.mtx

    year0=year1-jmax+1
      year2=year1+tmax-jmax
     year3=year2+jmax-1
    #println(year0)
    #println(year1)
    #println(year2)
    #println(year3)

    vecs=Array{cvec,1}(undef,tmax+jmax-1)

#    vecs=Array{Any}(undef,tmax+jmax-1)

    for year in year0:year1-1
        i=year-year0+1
        #println(year)
        #println(i)

        vecs[i]=cvec(i,year,
        jmax-i+1,i,argmax(map(j->mtx[j,jmax-i+j],1:i)),
        map(j->mtx[j,jmax-i+j],1:i))
        #println(vecs[i])
    end

    if tmax>=jmax
        for year in year1:year2
            i=year-year1+1
             #println(year)
             #println(i)
        #println(i-year0+year1)
            vecs[i-year0+year1]=cvec(i-year0+year1,year,
            1,jmax, argmax(map(j->mtx[i+j-1,j],1:jmax)),
            map(j->mtx[i+j-1,j],1:jmax))
            #println( vecs[i-year0+year1])
        end
    end

   #println(3)

    for year in year2+1:year3
            i=year-year1+1
            #println(year)
            #println(i)
            #println(i-year1+year2)
        vecs[i-year1+year2+1]=cvec(i-year1+year2+1,year,
        1,tmax-i+1,argmax(map(j->mtx[i+j-1,j],1:tmax-i+1)),
        map(j->mtx[i+j-1,j],1:tmax-i+1))
        #println( vecs[i-year1+year2+1])
    end

    csvecs(year1,tmax,jmax,vecs)
end

"""
sequence of vectors (of mutable struct csvecs)->cohorts related matrix (of mutable struct csrm) (x,y,w,m,f,z,g)

csrm:cohorts related matrix (of mutable struct csrm)
year_b - initial physical year
tmax - number of seasons
jmax - number of ages
x,y,w, m,f,g -  cohorts related matrices

cvec: detail cohort representation as vec
no- the number in the sequence
yb - physical year associated with cohort (when age=1)
jbeg - initial age
len - length
jtrust - trusted age
vec

csvecs: cohorts related mtx as a sequence of cvec
no- the number
yb - physical year associated with cohort (when age=1)
tmax - number of seasons
jmax - number of ages
jbeg - initial age
len - length
jtrust - trusted age
xv,yv,wv- vectors with abundunce, catch, weight
tmax+jmax-2 -- number of related cohorts
"""
function csvs2csmtx(csvs)

    year1=csvs.year_b
    tmax=csvs.tmax
    jmax=csvs.jmax
    vs=csvs.vecs

    mtx=fill(1.0,tmax,jmax)

    ncs=tmax+jmax-1
    year0=year1-jmax+1
    year2=year1+tmax-jmax
    year3=year2+jmax-1

    #println(year0)
    #println(year1)
    #println(year2)
    #println(year3)

   #1
    #println(1)
    for year in year0:year1-1
        i=year-year0+1
        #println(year)
        #println(i)

    cv1=vs[i]
    no=cv1.no
    yb=cv1.yb
    jbeg=cv1.jbeg
    len=cv1.len
    jtrust=cv1.jtrust
    xv=cv1.vec

     #println(' ',i,' ',no,' ',yb,' ',jbeg,' ',len)

        for j in 1:i
            #println(j,' ',jmax-i+j)
            mtx[j,jmax-i+j]=xv[j]

        end
    end
  #2
    #println(2)
     if tmax>=jmax
        for year in year1:year2
            i=year-year0+1
            i1=year-year1+1
            #println(year)
            #println(i)

    cv1=vs[i]
    no=cv1.no
    yb=cv1.yb
    jbeg=cv1.jbeg
    len=cv1.len
    jtrust=cv1.jtrust
    xv=cv1.vec

           #println(i,' ',no,' ',yb,' ',jbeg,' ',len)

            for j in 1:jmax
           # println(i+j-1,' ',j)
            mtx[i1+j-1,j]=xv[j]

        end
    end
        end

  #3
  #println(3)
    for year in year2+1:year3
            i=year-year0+1
        i1=year-year1+1

            #println(year)
            #println(i)

    cv1=vs[i]
    no=cv1.no
    yb=cv1.yb
    jbeg=cv1.jbeg
    len=cv1.len
    jtrust=cv1.jtrust
     xv=cv1.vec

            #println(i,' ',no,' ',yb,' ',jbeg,' ',len)

        for j in 1:tmax-i1+1
        # println(i1+j-1,' ',j)

        mtx[i1+j-1,j]=xv[j]

        end
    end

    csmtx(year1,tmax,jmax,mtx)
end

"""
generates xv1,yv1 of csvecs mutable struct with use of xv::csvecs,yv::csvecs,wv::csvecs,zv::csvecs,gv::csvecs
first makes a deep copy, then modifies vectors
"""
function gvecs(xv,yvs,wv,zv,gv)
    xv1=deepcopy(xv)
    yv1=deepcopy(yv)

    year1=xv.year_b
    tmax=xv.tmax
    jmax=xv.jmax
    xvecs=xv.vecs
    yvecs=yv.vecs
    wvecs=yv.vecs
    zvecs=zv.vecs
    gvecs=gv.vecs

    x1vecs=xv1.vecs
    y1vecs=yv1.vecs

    #println(xv1)

    # println( mvecs)
 #1
    for i in 2:jmax-1
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
        if jtr >1
        for i1 in 1:jtr-1
        #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(sum(zi[i1:jtr-1]))
        end
        end
        if jtr+1<len
        for i1 in jtr+1:len
          #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(-sum(zi[jtr+1:i1]))
                y1vecs[i].vec[i1]= wi[i1]*x1vecs[i].vec[i1]*gi[i1]
        end
        end

        #println(x1vecs[i])

    end
   #2
   for i in jmax:tmax
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
        if jtr >1
        for i1 in 1:jtr-1
         #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(sum(zi[i1:jtr-1]))
        end
        end
        if jtr+1<len
        for i1 in jtr+1:len
          #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(-sum(zi[jtr+1:i1]))
                 y1vecs[i].vec[i1]= wi[i1]*x1vecs[i].vec[i1]*gi[i1]
        end
        end

        #println(x1vecs[i])

    end
    #3
    for i in tmax+1:tmax+jmax-2
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
        if jtr >1
        for i1 in 1:jtr-1
         #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(sum(zi[i1:jtr-1]))
        end
        end
        if jtr+1<len
        for i1 in jtr+1:len
          #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(-sum(zi[jtr+1:i1]))
                 y1vecs[i].vec[i1]= wi[i1]*x1vecs[i].vec[i1]*gi[i1]
        end
        end

        #println(x1vecs[i])
        #println(y1vecs[i])

    end

    return(xv1,yv1)

end

"""
generates xv1,yv1 of csvecs mutable struct with use of xv::csvecs,yv::csvecs,wv::csvecs,zv::csvecs,gv::csvecs
first makes a deep copy, then modifies vectors
"""
function gvecs1(xv,yv,wv,zv,gv)
    xv1=deepcopy(xv)
    yv1=deepcopy(yv)

    year1=xv.year_b
    tmax=xv.tmax
    jmax=xv.jmax
    xvecs=xv.vecs
    yvecs=yv.vecs
    wvecs=yv.vecs
    zvecs=zv.vecs
    gvecs=gv.vecs

    x1vecs=xv1.vecs
    y1vecs=yv1.vecs

    #println(xv1)

    # println( mvecs)
 #1
    for i in 2:jmax-1
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
        if jtr >1
        for i1 in 1:jtr-1
        #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(sum(zi[i1:jtr-1]))

        end
        end
        if jtr+1<len
        for i1 in jtr+1:len
          #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(-sum(zi[jtr+1:i1]))

        end
        end

        #println(x1vecs[i])

    end

   #2
   for i in jmax:tmax
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
        if jtr >1
        for i1 in 1:jtr-1
         #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(sum(zi[i1:jtr-1]))

        end
        end
        if jtr+1<len
        for i1 in jtr+1:len
          #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(-sum(zi[jtr+1:i1]))

        end
        end

        #println(x1vecs[i])

    end
    #3
    for i in tmax+1:tmax+jmax-2
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
        if jtr >1
        for i1 in 1:jtr-1
         #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(sum(zi[i1:jtr-1]))

        end
        end
        if jtr+1<len
        for i1 in jtr+1:len
          #println(i1)
            x1vecs[i].vec[i1]=xijtr*exp(-sum(zi[jtr+1:i1]))

        end
        end

        #println(x1vecs[i])
        #println(y1vecs[i])

    end

    ######

     for i in 1:jmax-1
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)

        for i1 in 1:len
        #println(i1)
           y1vecs[i].vec[i1]= wi[i1]*x1vecs[i].vec[i1]*gi[i1]

        end
                end

   #2
   for i in jmax:tmax
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

        #println(xijtr)
       for i1 in 1:len
        #println(i1)
           y1vecs[i].vec[i1]= wi[i1]*x1vecs[i].vec[i1]*gi[i1]

        end
                end

    #3
    for i in tmax+1:tmax+jmax-1
        xvi=xvecs[i]
        yvi=yvecs[i]
        wvi=wvecs[i]
        zvi=zvecs[i]
        gvi=gvecs[i]

        x1vi=x1vecs[i]
        y1vi=y1vecs[i]
      #  println(mvi)

         no=xvi.no
        yb=xvi.yb
        jbeg=xvi.jbeg
         len=xvi.len
        jtr=xvi.jtrust
        xi=xvi.vec

        yi=yvi.vec
        wi=wvi.vec
        zi=zvi.vec
        gi=gvi.vec

        x1i=x1vi.vec
        y1i=y1vi.vec

        #println(jbeg,' ',jtr)

        xijtr=xi[jtr]

        x1vecs[i].vec[jtr]=xijtr

       for i1 in 1:len
        #println(i1)
           y1vecs[i].vec[i1]= wi[i1]*x1vecs[i].vec[i1]*gi[i1]

        end
                end

    return(xv1,yv1)

end

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

"""
look how to evaluate closeness of data and model!!!!
as L1+L2
no prints out
"""
function Lp(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)

    function g(hh,mm,ss,ff)
   ss*ff/(hh*mm+ss*ff)*(1-exp(-hh*mm-ss*ff))
    end

    local  x=cst.x

      local  y=cst.y
     local   w=cst.w

     (tmax,jmax)=size(x)

   #println(tmax," ",jmax)
   # println(a1m,a2m,b1m,b2m,cm,d1m,d2m,as,bs)
   #println(htj)
   #println(ft)
     local  L1=0.0

    for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+abs(log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t])^p
           #println(t," ",j," ",L1t)
        end
        L1=L1+L1t

    end

   #println(L1)

     local   L2=0.0
    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+abs(log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t])))^p

           #println(t," ",j," ",L2t)
        end;
        L2=L2+L2t;

    end
    #println(L2)
  # println(L1+L2)

    if L1+L2==Inf 10000.0 else L1+L2 end
end

"""
L as L1+L2 with dispersions for the normal distribution, p=2
no prints out
"""
function Lmine2(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)

    function g(hh,mm,ss,ff)
   ss*ff/(hh*mm+ss*ff)*(1-exp(-hh*mm-ss*ff))
    end

    local  x=cst.x

      local  y=cst.y
     local   w=cst.w

     (tmax,jmax)=size(x)

   #println(tmax," ",jmax)
   # println(a1m,a2m,b1m,b2m,cm,d1m,d2m,as,bs)
   #println(htj)
   #println(ft)
     local  L1=0.0
    local  L11=0.0
    local mean1

    for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+
                    log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t]
           #println(t," ",j," ",L1t)
        end
      L1=L1+L1t
     end

        mean1=L1/((tmax-1)*(jmax-1))

     for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+
                    (log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t]-mean1)^2
           #println(t," ",j," ",L1t)
        end
        L11=L11+L1t
      end

   #println(L1)

     local   L2=0.0
        local  L22=0.0
        local mean2

    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+
                log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t]))

           #println(t," ",j," ",L2t)
        end;
        L2=L2+L2t;

    end

     mean2=L2/(tmax*jmax)

    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+
                    (log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t]))-mean2)^2

           #println(t," ",j," ",L2t)
        end;
        L22=L22+L2t;

    end

    #println(L2)
  # println(L1+L2)
    local result
    result=lambda *L11/((tmax-1)*(jmax-1)-1)+(1. - lambda)*L22/(tmax*jmax)

    if result==Inf 10000.0 else result end
end

"""
L as max
no prints out
"""
function Lmine2max(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)

    function g(hh,mm,ss,ff)
   ss*ff/(hh*mm+ss*ff)*(1-exp(-hh*mm-ss*ff))
    end

    local  x=cst.x

      local  y=cst.y
     local   w=cst.w

     (tmax,jmax)=size(x)

   #println(tmax," ",jmax)
   # println(a1m,a2m,b1m,b2m,cm,d1m,d2m,as,bs)
   #println(htj)
   #println(ft)
     local  L1=0.0
    local  L11=0.0
    local mean1

    for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+
                    log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t]
           #println(t," ",j," ",L1t)
        end
      L1=L1+L1t
     end

        mean1=L1/((tmax-1)*(jmax-1))

     for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+
                    (log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t]-mean1)^2
           #println(t," ",j," ",L1t)
        end
        L11=L11+L1t
      end

   #println(L1)

     local   L2=0.0
        local  L22=0.0
        local mean2

    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+
                log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t]))

           #println(t," ",j," ",L2t)
        end;
        L2=L2+L2t;

    end

     mean2=L2/(tmax*jmax)

    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+
                    (log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t]))-mean2)^2

           #println(t," ",j," ",L2t)
        end;
        L22=L22+L2t;

    end

    #println(L2)
  # println(L1+L2)
    local result
    global lambda

    result=max(lambda*L11/((tmax-1)*(jmax-1)-1),(1. - lambda)*L22/(tmax*jmax))

    if result==Inf 10000.0 else result end
end

"""
look how to evaluate closeness of data and model!!!!
as L1+L2
"""
function Lp12(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)

    function g(hh,mm,ss,ff)
   ss*ff/(hh*mm+ss*ff)*(1-exp(-hh*mm-ss*ff))
    end

    local  x=cst.x

      local  y=cst.y
     local   w=cst.w

     (tmax,jmax)=size(x)

   #println(tmax," ",jmax)
   # println(a1m,a2m,b1m,b2m,cm,d1m,d2m,as,bs)
   #println(htj)
   #println(ft)
     local  L1=0.0

    for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+abs(log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t])^p
           #println(t," ",j," ",L1t)
        end
        L1=L1+L1t

    end

   println(L1)

     local   L2=0.0
    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+abs(log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t])))^p

           #println(t," ",j," ",L2t)
        end;
        L2=L2+L2t;

    end
    println(L2)
   println(L1+L2)

    if L1+L2==Inf 10000.0 else L1+L2 end
end

function Lp1(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)

    function g(hh,mm,ss,ff)
   ss*ff/(hh*mm+ss*ff)*(1-exp(-hh*mm-ss*ff))
    end

    local  x=cst.x

      local  y=cst.y
     local   w=cst.w

     (tmax,jmax)=size(x)

   #println(tmax," ",jmax)
   # println(a1m,a2m,b1m,b2m,cm,d1m,d2m,as,bs)
   #println(htj)
   #println(ft)
     local  L1=0.0

    for t=1:tmax-1
        L1t=0.0
        for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1t=L1t+abs(log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t])^p
           #println(t," ",j," ",L1t)
        end
        L1=L1+L1t

    end

   println(L1)

     local   L2=0.0
    for t=1:tmax
        L2t=0.0
        for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2t=L2t+abs(log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t])))^p

           #println(t," ",j," ",L2t)
        end;
        L2=L2+L2t;

    end
    println(L2)
   println(L1+L2)

    if L1+L2==Inf 10000.0 else L1+L2 end
end

function Lp2(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)

    function g(hh,mm,ss,ff)
   ss*ff/(hh*mm+ss*ff)*(1-exp(-hh*mm-ss*ff))
    end

    local  x=cst.x

      local  y=cst.y
     local   w=cst.w

     (tmax,jmax)=size(x)

   #println(tmax," ",jmax)
   # println(a1m,a2m,b1m,b2m,cm,d1m,d2m,as,bs)
   #println(htj)
   #println(ft)
     local  L1=0.0

    for t=1:tmax-1
                for j=1:jmax-1
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)
            L1=L1+abs(log(x[t,j]/x[t+1,j+1])-htj[t,j]*mj-sj*ft[t])^p
           #println(t," ",j," ",L1t)
        end

    end

   println(L1)

     local   L2=0.0
    for t=1:tmax
              for j= 1:jmax
            sj= sigma(j,as,bs,jmax)
            mj=ushape(j,a1m,a2m,b1m,b2m,cm,d1m,d2m,jmax)

            L2=L2+abs(log(y[t,j])-log(kw*w[t,j]*kx*x[t,j]*g(htj[t,j],mj,sj,ft[t])))^p

           #println(t," ",j," ",L2t)
        end;

    end
    println(L2)
   println(L1+L2)

    if L1+L2==Inf 10000.0 else L1+L2 end
end

"""
generates x,y with M,F,W,kx,kw,x1j,xt1
"""
function gm1(tmax,jmax,M,F,W,kx,kw,x1j,xt1)
##0
    x=fill(0.0,tmax,jmax);

    for j in [1:jmax;] x[1,j]=x1j[j] end

    for t in 1:tmax x[t,1]=xt1[t] end

 ##1
        for t in [1:tmax-jmax+1;]
           for i in [0:jmax-2;]
            #println(t+i+1," ",i+2,"<-", t+i," ",i+1)
            x[t+i+1,i+1+1]=x[t+i,i+1]*exp(-M[t+i,i+1]-F[t+i,i+1])
           end
        end

 ##2
    for t in [tmax-jmax+2:tmax-1;]
           for i in [0:tmax-t-1;]
            #println(t+i+1," ",i+2,"<-", t+i," ",i+1)
            x[t+i+1,i+1+1]=x[t+i,i+1]*exp(-M[t+i,i+1]-F[t+i,i+1])
           end
        end
   ##3

    for j in [2:jmax;]
           for i in [0:jmax-j-1;]
           #println(i+1+1,"! ",j+i+1,"<-", i+1," ",j+i)
            x[i+1+1,j+i+1]=x[i+1,j+i]*exp(-M[i+1,j+i]-F[i+1,j+i])
           end
        end

 ##0
     y=fill(0.0,tmax,jmax);

##1
   for t in 1:tmax
        for j in 1:jmax
            #println(t+i," ",i+1)
            Z=M[t,j]+F[t,j]
            y[t,j]=
            (kw*W[t,j])*(kx*x[t,j])*F[t,j]/Z*(1- exp(-Z))
        end
    end

    return x,y
end

"""
returns ovsyannikov's weight for j
"""
function ww_osp(t,j)
   wv=[0.012 0.052 0.11 0.179 0.253 0.323 0.483 0.677 0.908 1.116 1.297 1.541 1.758 1.971 1.997 2.358 2.475];
    return wv[j]
end

"""
returns total mortality z as sum of natural m and fishing f mortality matrices
"""
function gz(m::Array{Float64,2},f::Array{Float64,2},tmax,jmax)
    z=Array{Float64}(undef,tmax,jmax)
    z=m+f
    return(z)
end

"""
generates g matrix with m and f, tmax,jmax
"""
function gg(m::Array{Float64,2},f::Array{Float64,2},tmax,jmax)
    z=Array{Float64}(undef,tmax,jmax)
    g=Array{Float64}(undef,tmax,jmax)
    z=gz(m,f,tmax,jmax)
    for t in 1:tmax
        for j in 1:jmax
            g[t,j]=f[t,j]/z[t,j]*(1-exp(-z[t,j]))
        end
    end
    return(g)
end

"""
evaluates criteria
"""
function Lpv(cst,a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,as,bs,ft,kx,kw,p)
    local  x=cst.x

      local  y=cst.y
     local   w=cst.w
   local year1= cst.year_b

     (tmax,jmax)=size(x)

     M0=gmtj(a1m,a2m,b1m,b2m,cm,d1m,d2m,htj,tmax,jmax);
    F0=gftj(as,bs,ft,tmax,jmax);
    Z0=M0+F0;
    G0=gg(M0,F0,tmax,jmax);

    xv=csmtx2csvs(csmtx(year1,tmax,jmax,x));
yv=csmtx2csvs(csmtx(year1,tmax,jmax,y));
    mv=csmtx2csvs(csmtx(year1,tmax,jmax,M0));
fv=csmtx2csvs(csmtx(year1,tmax,jmax,F0));
zv=csmtx2csvs(csmtx(year1,tmax,jmax,Z0));
gv=csmtx2csvs(csmtx(year1,tmax,jmax,G0));
wv=csmtx2csvs(csmtx(year1,tmax,jmax,w));
 xv0,yv0=gvecs(xv,yv,wv,zv,gv);

   x11=csvs2csmtx(xv0);
y11=csvs2csmtx(yv0);

    L=0;
for i in 1:tmax+jmax-1
    lenxv0=xv0.vecs[i].len;
   vecxv0=xv0.vecs[i].vec;
    vecyv0=yv0.vecs[i].vec;
    vecxv=xv.vecs[i].vec;
    vecyv=yv.vecs[i].vec;
           for n in 1:lenxv0
        L=L+abs(log(vecxv0[n])-log(vecxv[n]))^p+abs(log(vecyv0[n])-log(vecyv[n]))^p
    end

end

    if L==Inf 10000.0 else L end
end

"""
optimizes Lpv
"""
function optimizeLp_v(x,tmax,jmax,a,lb,lb1,xtol,xtol1,bc4,ctol,is1,is2,is3,pop,pop1,maxe)

   global count = 0
opt = Opt(ALGS[a], 4*tmax+9)
#l_opt= Opt(ALGS[a1], 4*tmax+9)
    lower_bounds!(opt, fill(lb,4*tmax+9))
    #lower_bounds!(l_opt, fill(lb1,4*tmax+9))
    xtol_rel!(opt,xtol)
    #xtol_rel!(l_opt,xtol1)

min_objective!(opt, buragofunc_v)
    inequality_constraint!(opt, (x,g) -> buragoconstraint1(x,jmax), ctol)
    inequality_constraint!(opt, (x,g) -> buragoconstraint2(x,jmax), ctol)
    inequality_constraint!(opt, (x,g) -> buragoconstraint3(x,jmax),ctol )
    inequality_constraint!(opt, (x,g) -> buragoconstraint4(x,bc4), ctol)
#for i in 1:3*tmax
#    inequality_constraint!(opt, (x,g) -> buragoconstraint5(x,i,0.1), 1e-8)
#     inequality_constraint!(opt, (x,g) -> buragoconstraint6(x,i,5.), 1e-8)
#end

#local_optimizer!(opt,l_opt)

is=fill(0.4,4*tmax+9)
    is[1:3*tmax]=fill(is1,3*tmax)#h
    is[3*tmax+1:4*tmax]=fill(is2,tmax)#f
    is[4*tmax+1:4*tmax+9]=fill(is3,9)#pm&pf

initial_step!(opt,is)
population!(opt,pop)
#population!(l_opt,pop1)
NLopt.srand_time()
    maxeval!(opt, maxe)

(minf,minx,ret) = NLopt.optimize(opt,x)

end

"""
returns Lpv
"""
function buragofunc_v(x::Vector, grad::Vector)

    global cst
    (tmax,jmax)=size(cst.x)

    local Hl=fill(1.0,tmax,jmax)
    local ftl=fill(1.,tmax)

    global kx,kw
    global count

    if length(grad) > 0
        grad[1] = 0
        grad[2] = 0.5/sqrt(x[2])
    end

    count::Int += 1

for t in 1:tmax
        Hl[t,1:3]=x[3*(t-1)+1:3*t]
    end

    local a1ml,a2ml,b1ml,b2ml,cml,d1ml,d2ml,asl,bsl
    ftl=x[3*tmax+1:4*tmax]
    a1ml,a2ml,b1ml,b2ml,cml,d1ml,d2ml=x[4*tmax+1:4*tmax+7]
    asl,bsl=x[4*tmax+8:4*tmax+9]

    #println(cst,psl,Hl,ftl)
    return Lpv(cst,a1ml,a2ml,b1ml,b2ml,cml,d1ml,d2ml,Hl,asl,bsl,ftl,kx,kw,1)

    #println("f_$count($x)=$y")

end


module DiffEqSensitivity

using DiffEqBase, ForwardDiff, Tracker, FiniteDiff, Statistics
using DiffEqCallbacks, QuadGK, RecursiveArrayTools, LinearAlgebra
using DiffEqOperators
using Adapt
using LinearSolve
using Parameters: @unpack
using StochasticDiffEq
import DiffEqNoiseProcess
import RandomNumbers: Xorshifts
using Random
import ZygoteRules, Zygote, ReverseDiff
import ArrayInterfaceCore, ArrayInterfaceTracker
import Enzyme
import GPUArrays

using Cassette, DiffRules
using Core: CodeInfo, SlotNumber, SSAValue, ReturnNode, GotoIfNot
using EllipsisNotation

using Markdown

using Reexport
import ChainRulesCore: @thunk, NoTangent, @not_implemented
abstract type SensitivityFunction end
abstract type TransformedFunction end

include("hasbranching.jl")
include("sensitivity_algorithms.jl")
include("derivative_wrappers.jl")
include("sensitivity_interface.jl")
include("forward_sensitivity.jl")
include("adjoint_common.jl")
include("lss.jl")
include("nilss.jl")
include("nilsas.jl")
include("backsolve_adjoint.jl")
include("interpolating_adjoint.jl")
include("quadrature_adjoint.jl")
include("callback_tracking.jl")
include("concrete_solve.jl")
include("second_order.jl")
include("steadystate_adjoint.jl")
include("sde_tools.jl")

# AD Extensions
include("reversediff.jl")
include("tracker.jl")
include("zygote.jl")

export extract_local_sensitivities

export ODEForwardSensitivityFunction, ODEForwardSensitivityProblem, SensitivityFunction,
       ODEAdjointSensitivityProblem, ODEAdjointProblem, AdjointSensitivityIntegrand,
       SDEAdjointProblem, RODEAdjointProblem, SensitivityAlg,
       adjoint_sensitivities, adjoint_sensitivities_u0,
       ForwardLSSProblem, AdjointLSSProblem,
       NILSSProblem, NILSASProblem,
       shadow_forward, shadow_adjoint

export BacksolveAdjoint, QuadratureAdjoint, InterpolatingAdjoint,
       TrackerAdjoint, ZygoteAdjoint, ReverseDiffAdjoint,
       ForwardSensitivity, ForwardDiffSensitivity,
       ForwardDiffOverAdjoint,
       SteadyStateAdjoint,
       ForwardLSS, AdjointLSS, NILSS, NILSAS

export second_order_sensitivities, second_order_sensitivity_product

export TrackerVJP, ZygoteVJP, EnzymeVJP, ReverseDiffVJP

export StochasticTransformedFunction

end # module


### A Pluto.jl notebook ###
# v0.12.20

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ a6c0b440-5688-11eb-0635-39912417c1f9
begin
	using NeXLDatabase, Gadfly, DataFrames, PlutoUI
md"""
## Plot k-ratios
This script makes it easy to plot some or all of the k-ratios associated with a single element over different measurement sets.
  1) Select an element
  2) Specify a minimum k-ratio to consider (eliminates materials with trace quantities of the element)
  3) Select one or more sets of measurements to display.  Measurements are as identified by material, standard material and beam energy.
"""
end

# ╔═╡ 3742ed50-1852-11eb-01e3-5b27d6c85250
db=openNeXLDatabase("C:\\Users\\nritchie\\Documents\\DrWatson\\K-Ratio Project\\data\\exp_pro\\kratio.db")

# ╔═╡ e793f370-1852-11eb-24b6-05fa57c3a3c6
md"""
Select an element or elements: $(@bind els MultiSelect(map(el->symbol(el),elements[1:94]),default=["Al"]))
"""

# ╔═╡ 78885a50-1859-11eb-0b6c-3368b717d267
md"""
Specify the minimum k-ratio: $(@bind mink NumberField(0.00 : 0.01 : 1.0; default=0.1))
"""

# ╔═╡ a51b0a40-1921-11eb-2c1c-ab31e48e7846
allkrs = mapreduce(el->findall(db, DBKRatio, elm=parse(Element,el), mink=mink), append!, els)

# ╔═╡ 22f61d10-185a-11eb-2c53-85f7ef0fb78a
md"""
Select which set of measurements to display:

$(@bind krsnames MultiSelect(sort(unique(map(kr->NeXLDatabase.krname(kr), allkrs)))))
"""

# ╔═╡ 2a8bd5f0-185b-11eb-0b96-8592774730a0
begin
	krs = filter(kr->NeXLDatabase.krname(kr) in krsnames, allkrs)
	if length(krs)>0
		plot(krs) |> SVG(joinpath(homedir(),"Desktop","k-ratio plot.svg"),6inch,4inch)
		plot(krs)
	else
		md"No k-ratios have been selected"
	end
end

# ╔═╡ cb991c00-1888-11eb-02fb-b33a153473dc
asa(DataFrame, DBCampaign[ read(db, DBCampaign, kr.campaign) for kr in krs])

# ╔═╡ 6185fa20-19d9-11eb-04d3-3132221ca798
begin
	if length(krs)>0
		plot(krs, style=:XY)
	else
		md"No k-ratios have been selected"
	end
end

# ╔═╡ Cell order:
# ╟─a6c0b440-5688-11eb-0635-39912417c1f9
# ╟─3742ed50-1852-11eb-01e3-5b27d6c85250
# ╟─e793f370-1852-11eb-24b6-05fa57c3a3c6
# ╟─78885a50-1859-11eb-0b6c-3368b717d267
# ╟─a51b0a40-1921-11eb-2c1c-ab31e48e7846
# ╟─22f61d10-185a-11eb-2c53-85f7ef0fb78a
# ╠═2a8bd5f0-185b-11eb-0b96-8592774730a0
# ╟─cb991c00-1888-11eb-02fb-b33a153473dc
# ╟─6185fa20-19d9-11eb-04d3-3132221ca798


"""
    log_embeddings(logger::TBLogger, name::AbstractString, mat::AbstractMatrix; metadata, metadata_header, img_labels, step=step(logger))

Log embedding data to tensorboard and visualize in 3-D or 2-D with PCA, t-SNE or UMAP.
- mat: 2-D `Matrix` of data, with rows representing the samples and columns representing the features
- metadata: Array of labels for each sample. Each element across 1st dimenstion will be converted to `string`
- metadata_header: 1-D Array. Useful when samples have multiple labels. size should be same as size of each row in `metadata`
- img_labels: TBImages object representing image labels for each sample.
    - each number of images (N) must be equal to number of samples in `mat`
    - each image must be a square (H == W)
    - the value √N * W must be less than or equal to 8192 because of tensorboard restrictions.
"""
function log_embeddings(logger::TBLogger, name::AbstractString, mat::AbstractMatrix; metadata=nothing, metadata_header=nothing, img_labels=nothing, step=nothing)
    ndims(mat) == 2 || throw(DimensionMismatch("Embedding matrix must be 2-Dimensional. ("*string(ndims(mat))*" ≠ 2)"))
    matrix_path = joinpath(logger.logdir, repr(step), name)
    mkpath(matrix_path)
    if metadata != nothing
        size(metadata, 1) == size(mat, 1) || throw(ErrorException("#labels must be equal to #samples. ("*string(length(metadata))*" ≠ "*string(size(mat, 1))*")"))
        if metadata_header == nothing
            metadata = [string(x) for x in metadata]
        else
            length(metadata_header) == size(metadata, 2) || throw(ErrorException("length of header must be equal to the number of columns in metadata. ("*string(length(metadata_header))*" ≠ "*string(size(metadata, 2))*")"))
            metadata = [join(metadata_header, '\t'); [join(metadata[i, :], '\t') for i in 1:size(metadata, 1)]]
        end
        write_metadata(metadata, matrix_path)
    end
    if img_labels != nothing
        img_labels = convert_to_NCHW(content(img_labels), img_labels.format)
        size(img_labels, 1) == size(mat, 1) || throw(ErrorException("#label images must be equal to #samples. ("*string(size(img_labels, 1))*" ≠ "*string(size(mat, 1))*")"))
        size(img_labels, 3) == size(img_labels, 4) || throw(ErrorException("#label images must be square"))
        write_sprite(img_labels, matrix_path)
    end
    write_matrix(mat, matrix_path)
    write_pbtext(name, logger.logdir, matrix_path, metadata, img_labels, step)
end

function write_matrix(mat::AbstractMatrix, matrix_path)
    matrix_path = joinpath(matrix_path, "tensor.tsv")
    mat = convert(Array{Float64,2}, mat)
    open(matrix_path, "w") do file
        for i in 1:size(mat, 1)
            write(file, join(mat[i, :], '\t'))
            write(file, '\n')
        end
    end
end

function write_metadata(metadata::AbstractArray, matrix_path)
    matrix_path = joinpath(matrix_path, "metadata.tsv")
    open(matrix_path, "w") do file
        for x in metadata
            write(file, string(x)*'\n')
        end
    end
end

function write_sprite(img_labels::AbstractArray, matrix_path)
    n, _, _, w = size(img_labels)
    sqrt(n)*w <= 8192 || throw(ErrorException("the value √N * W must be less than or equal to 8192 because of tensorboard restrictions"))
    total_pixels = size(img_labels, 1)*size(img_labels, 3)*size(img_labels, 4)
    pixels_one_side = sqrt(total_pixels)
    number_of_images_per_row = Integer(ceil(pixels_one_side/size(img_labels, 4)))
    arranged_img_CHW = make_grid_of_images(img_labels, number_of_images_per_row)
    sprite_size = size(arranged_img_CHW, 3)
    arranged_augment_square_CHW = zeros((3, sprite_size, sprite_size))
    arranged_augment_square_CHW[:, 1:size(arranged_img_CHW, 2), :] = arranged_img_CHW
    sprite_path = joinpath(matrix_path, "sprite.png")
    open(sprite_path; write=true) do io
        save(Stream{format"PNG"}(io), colorview(RGB, arranged_augment_square_CHW))
    end
end

function make_grid_of_images(img_labels::AbstractArray, ncols::Integer)
    if size(img_labels, 2) == 1
        img_labels = cat(img_labels, img_labels, img_labels, dims = 2)
    end
    n, c, h, w = size(img_labels)
    ncols = min(ncols, n)
    nrows = Integer(ceil(n/ncols))
    grid = zeros(c, h*nrows, w*ncols)
    i = 0
    for y in 1:nrows
        for x in 1:ncols
            if i >= n
                break
            end
            grid[:, (y-1)*h+1:y*h, (x-1)*w+1:x*w] = img_labels[i+1, :, :, :]
            i = i+1
        end
    end
    grid
end

function write_pbtext(name::AbstractString, path, matrix_path, metadata, img_labels, step)
    metadata_path = joinpath(matrix_path, "metadata.tsv")
    img_labels_path = joinpath(matrix_path, "sprite.png")
    matrix_path = joinpath(matrix_path, "tensor.tsv")
    path = joinpath(path, "projector_config.pbtxt")
    isfile(path) || write(path, "") # workaround https://github.com/JuliaCloud/AWSS3.jl/issues/173
    open(path, "a") do file
        write(file, "embeddings {\n")
        write(file, "tensor_name: \""*name*":"*repr(step)*"\"\n")
        write(file, string("tensor_path: \"", matrix_path, "\"\n"))
        if metadata != nothing
            write(file, string("metadata_path: \"", metadata_path, "\"\n"))
        end
        if img_labels != nothing
            write(file, "sprite {\n")
            write(file, string("image_path: \"", img_labels_path, "\"\n"))
            write(file, "single_image_dim: "*string(size(img_labels, 4))*"\n")
            write(file, "single_image_dim: "*string(size(img_labels, 3))*"\n")
            write(file, "}\n")
        end
        write(file, "}\n")
    end
    return
end


module SamplingUtils

md(m, σ) = Dict("m"=>m, "σ"=>σ)

p_rng=Dict("x"=>50
    ,"ηA"=>0
    ,"ηF"=>0
    ,"muA"=>0
    ,"muF"=>0
    ,"s2"=> md(1.0, 1.5)
    ,"kR5"=>1.0
    ,"kr2"=>1.0
    ,"a15"=>md(1.0,1.5)
    ,"a25"=>md(0.6,1.5)
    ,"a35"=>md(0.3,1.5)
    ,"a5"=>md(2.0,0.5)
    ,"b15"=>md(0.5,1.5)
    ,"b25"=>md(1.0,1.5)
    ,"b35"=>md(2.0,1.5)
    ,"b5"=>md(1.0,0.5)
    ,"a18"=>md(1.0,1.5)
    ,"a28"=>md(0.8,1.5)
    ,"a38"=>md(0.6,1.5)
    ,"a48"=>md(0.3,1.5)
    ,"a8"=>md(2.0,0.5)
    ,"b18"=>md(0.5,1.5)
    ,"b28"=>md(1.0,1.5)
    ,"b38"=>md(2.0,1.5)
    ,"b48"=>md(3.0,1.5)
    ,"b8"=>md(1.0,0.5)
    ,"l05"=>md(2.0, 1.0)
    ,"l15"=>0.010
    ,"l25"=>0.010
    ,"l35"=>0.010
    ,"kP5"=>1.0
    ,"K"=>1000.0
    ,"maxA"=>2.3
    ,"maxF"=>2.1
    ,"s9"=> md(1.0, 1.5)
    ,"kR8"=>1.0
    ,"kR8"=>1.0
    ,"kr9"=>1.0
    ,"l08"=>md(2.0, 1.0)
    ,"l18"=>0.010 * 1
    ,"l28"=>0.010 * 1
    ,"l38"=>0.010 * 1
    ,"l48"=>0.010 * 1
    ,"kP8"=>1.0 * 1
    ,"gA"=>1.0
    ,"gF"=>1.0
    ,"gc"=>10.0
    ,"s05"=>md(0.01,1.5)
    ,"s5"=> md(1.0*10, 1.5)
    ,"s08"=>md(0.01,1.5)
    ,"s8"=> md(1.0*10, 1.5)
    ,"K5A"=>md(1.0,0.5)
    ,"n5A"=>[6.0, 6.00]
    ,"K58"=>md(0.1, 1.5)
    ,"n58"=>[6.0, 6.00]
    ,"K85"=>md(0.1, 1.5)
    ,"n85"=>[6.0, 6.00]
    ,"K8F"=>md(1.0,0.5)
    ,"n8F"=>[6.0, 6.00]
    ,"K9F"=>0.99 * 6.00
    ,"n9F"=>[6.0, 6.00]
    ,"K2A"=>md(1.0, 0.5)
    ,"n2A"=>[6.0, 6.00]
    ,"K25"=>md(0.1, 1.5)
    ,"n25"=>[6.0, 6.00]
    ,"K98"=>md(0.1, 1.5)
    ,"n98"=>[6.0, 6.00]
    ,"y"=>1)

end


#=
# spin10.jl - This module demonstrates grape on the spin system in julia.
=#
module Foo

# using BenchmarkTools
# using CUDAnative
# using CuArrays
# using DiffEqSensitivity
# using DifferentialEquations
# using ForwardDiff
using HDF5
using LinearAlgebra
# using Polynomials
using Printf
using StaticArrays
using Zygote

EXPERIMENT_META = "spin"
EXPERIMENT_NAME = "spin0"
WDIR = ENV["QC_PATH"]
SAVE_PATH = joinpath(WDIR, "out/$EXPERIMENT_META/$EXPERIMENT_NAME")

function generate_save_file_path(save_file_name, save_path)
    # Ensure the path exists.
    mkpath(save_path)

    # Create a save file name based on the one given; ensure it will
    # not conflict with others in the directory.
    max_numeric_prefix = -1
    for (_, _, files) in walkdir(save_path)
        for file_name in files
            if occursin("_$save_file_name.h5", file_name)
                max_numeric_prefix = max(parse(Int, split(file_name, "_")[1]))
            end
        end
    end

    save_file_name = "_$save_file_name.h5"
    save_file_name = @sprintf("%05d%s", max_numeric_prefix + 1, save_file_name)

    return joinpath(save_path, save_file_name)
end

mutable struct Reporter
    cost :: Float64
    costs :: Array{Float64, 1}
end

mutable struct PState
    control_count :: Int64
    control_eval_count :: Int64
    control_eval_times :: AbstractArray{Float64}
    cost_multipliers :: Array{Float64, 1}
    dt :: Float64
    evolution_time :: Float64
    initial_astate :: Any
    lagrange_multipliers :: Array{Float64, 1}
end

# We need an isomorphism from complex numbers to real numbers.
# TODO: It may be more memory efficient to store the augmented
# matrix as [mat_r ; mat_i] and define special multiplication
# operations.
"""
Promote a complex vector on the Hilbert space to
a real operator on the augmented Hilbert space.
"""
function augment_vec(vec)
    return [real(vec); imag(vec)]
end

"""
Complex conjugate an augmented vector.
"""
function conj_vec(vec)
    (len,) = size(vec)
    imag_indices = Int(len / 2) + 1:len
    vec[imag_indices] = -vec[imag_indices]
    return vec
end

"""
Promote a complex operator on the Hilbert space to a real
operator on the augmented Hilbert space.
"""
function augment_mat(mat_r, mat_i)
    return [mat_r -mat_i ; mat_i mat_r]
end

"""
Complex conjugate an augmented matrix.
"""
function conj_mat(mat)
    (len, len) = size(mat)
    len_by_2 = Int(len / 2)
    lo = 1:len_by_2
    hi = len_by_2 + 1:len
    mat[lo, hi] = -mat[lo, hi]
    mat[hi, lo] = -mat[hi, lo]
    return mat
end

"""
Multiply an augmented matrix by the scalar negative imaginary unit (-i).

This is faster than multiplying by a matrix but
cannot be used in the differentiable code
because Zygote does not like the indexing.
"""
function neg_i_mat(mat)
    (len, len) = size(mat)
    len_by_2 = Int(len / 2)
    lo = 1:len_by_2
    hi = len_by_2 + 1:len
    mat_r = mat[lo, lo]
    mat_i = mat[hi, lo]
    return [mat_i mat_r ; -mat_r mat_i]
end

"""
Construct an augmented matrix that represents the conjugate transpose of
a vector.
"""
function conjugate_transpose_vec(vec)
    vec_r_t = real(vec)'
    vec_i_t = imag(vec)'
    return [vec_r_t vec_i_t ; -vec_i_t vec_r_t]
end

# Define experimental constants.
OMEGA = 2 * pi * 1e-2
MAX_CONTROL_NORM_0 = 2 * pi * 3e-1
HAMILTONIAN_ARGS = [
    OMEGA,
]
MAX_CONTROL_NORMS = [
    MAX_CONTROL_NORM_0,
]

# Define the system.
NEG_I = SA_F64[0   0  1  0 ;
               0   0  0  1 ;
               -1  0  0  0 ;
               0  -1  0  0 ;]
SIGMA_X = SA_F64[0 1 0 0;
                 1 0 0 0;
                 0 0 0 1;
                 0 0 1 0]
SIGMA_Z = SA_F64[1   0  0   0;
                 0  -1  0   0;
                 0   0  1   0;
                 0   0  0  -1]
H_S = SIGMA_Z / 2
NEG_I_H_S = NEG_I * H_S
OMEGA_NEG_I_H_S = OMEGA * NEG_I_H_S
H_C1 = SIGMA_X / 2
NEG_I_H_C1 = NEG_I * H_C1

# Define the optimization.
EVOLUTION_TIME = 120
COMPLEX_CONTROLS = false
CONTROL_COUNT = 1
CONTROL_EVAL_COUNT = SYSTEM_EVAL_COUNT = Int(EVOLUTION_TIME) + 1
CONTROL_EVAL_TIMES = SVector{CONTROL_EVAL_COUNT, Float64}(Array(range(0., stop=EVOLUTION_TIME, length=CONTROL_EVAL_COUNT)))
DT = 1e-2
GRAB_CONTROLS = false
if GRAB_CONTROLS
    initial_controls_ = h5open(joinpath(SAVE_PATH, "00007_spin10.h5")) do save_file
        save_file["controls"][200, :, :]
    end
    initial_controls_ = SMatrix{CONTROL_EVAL_COUNT, CONTROL_COUNT, Float64}(initial_controls_)
else
    initial_controls_ = @SMatrix ones(CONTROL_EVAL_COUNT, CONTROL_COUNT)
end
INITIAL_CONTROLS = initial_controls_

# Define the problem.
INITIAL_STATE = augment_vec(SA[1.; 0])
ZEROS_INITIAL_STATE = augment_vec(SA[0; 0])
TARGET_STATE_DAGGER = conjugate_transpose_vec(SA[0; 1.])
INITIAL_ASTATE = [
    INITIAL_STATE; # state
    # ZEROS_INITIAL_STATE; # dstate_dw
    # ZEROS_INITIAL_STATE; # d2state_dw2
]
(STATE_SIZE,) = size(INITIAL_STATE)
(ASTATE_SIZE,) = size(INITIAL_ASTATE)
STATE_INDICES = 1:STATE_SIZE
DSTATE_DW_INDICES = STATE_SIZE + 1:STATE_SIZE * 2
D2STATE_DW2_INDICES = STATE_SIZE * 2 + 1:STATE_SIZE * 3

HILBERT_DIM = Int(STATE_SIZE / 2)

# Define Misc
LEARNING_RATE = 1e-3
ITERATION_COUNT = 200
INITIAL_COSTS = [0., 0]
INITIAL_COST_MULTIPLIERS = [1., 1.]
INITIAL_LAGRANGE_MULTIPLIERS = [0., 0]
COST_MULTIPLIER_STEP = 5.
SAVE = false

"""
"""
function rk3_step(rhs_, x, u, t, dt, us)
    k1 = rhs_(x,             u, t       , us) * dt
    k2 = rhs_(x + k1/2,      u, t + dt/2, us) * dt
    k3 = rhs_(x - k1 + 2*k2, u, t + dt  , us) * dt
    return x + (k1 + 4 * k2 + k3) / 6
end

# Attempt at in-place was slower than static array version.
# function rk3_step(rhs_, x, u, t, dt, dx, ks)
#     rhs_(dx, x,             u, t       )
#     ks[1,:] = dx * dt
#     rhs_(dx, x + ks[1,:]/2,      u, t + dt/2)
#     ks[2,:] = dx * dt
#     rhs_(dx, x - ks[1,:] + 2*ks[2,:], u, t + dt  )
#     ks[3, :] = dx * dt
#     return x + (ks[1,:] + 4 * ks[2,:] + ks[3,:]) / 6
# end

"""
"""
function rk4_step(rhs_, x, u, t, dt, us)
    k1 = rhs_(x,        u, t       , us) * dt
	k2 = rhs_(x + k1/2, u, t + dt/2, us) * dt
	k3 = rhs_(x + k2/2, u, t + dt/2, us) * dt
	k4 = rhs_(x + k3,   u, t + dt  , us) * dt
	return x + (k1 + 2 * k2 + 2 * k3 + k4) / 6
end

"""
Evaluate the polynomial using horner's method.
"""
function horner(coeffs, x)
    (len,) = size(coeffs)
    if len == 1
        val = coeffs[1]
    else
        val = coeffs[len]
        for i in len:-1:2
            val = coeffs[i - 1] + x * val
        end
    end
    return val
end

"""
Interpolate at x3 given (x1, y1), (x2, y2)
"""
function interpolate_linear_points(x1, x2, x3, y1, y2)
    return y1 + (((y2 - y1) / (x2 - x1)) * (x3 - x1))
end

"""
Interpolate the controls at a given time.
"""
function interpolate(controls, control_eval_times, time)
    (len,) = size(control_eval_times)
    if time < control_eval_times[1]
        t1 = control_eval_times[1]
        t2 = control_eval_times[2]
        c1 = controls[1]
        c2 = controls[2]
    elseif time >= control_eval_times[len]
        t1 = control_eval_times[len]
        t2 = control_eval_times[len - 1]
        c1 = controls[len]
        c2 = controls[len - 1]
    else
        hi = Int(findall(t -> t > time, control_eval_times)[1][1])
        t1 = control_eval_times[hi]
        t2 = control_eval_times[hi - 1]
        c1 = controls[hi]
        c2 = controls[hi - 1]
    end

    controls_ = interpolate_linear_points(t1, t2, time, c1, c2)
    return controls_
end

"""
The right-hand-side of the Schroedinger equation for the dynamics
we consider.
"""
function rhs(astate, controls, time, control_step)
    # Unpack augmented state.
    state = astate[STATE_INDICES]
    # dstate_dw = astate[DSTATE_DW_INDICES]
    # d2state_dw2 = astate[D2STATE_DW2_INDICES]
    # less gc time but more allocations = slower in the no gradient regime
    # state = SVector{4, Float64}(astate[STATE_INDICES])
    # dstate_dw = SVector{4, Float64}(astate[DSTATE_DW_INDICES])
    # d2state_dw2 = SVector{4, Float64}(astate[D2STATE_DW2_INDICES])

    # Compute
    # linear interpolation
    # controls_ = interpolate_linear_points(CONTROL_EVAL_TIMES[control_step],
    #                                       CONTROL_EVAL_TIMES[control_step + 1],
    #                                       time,
    #                                       controls[control_step],
    #                                       controls[control_step + 1])

    # zero-order-hold
    controls_ = controls[control_step]

    neg_i_hamiltonian = OMEGA * NEG_I_H_S + controls_[1] * NEG_I_H_C1

    delta_state = neg_i_hamiltonian * state
    # delta_dstate_dw = neg_i_H_S * state + neg_i_hamiltonian * dstate_dw
    # delta_d2state_dw2 = neg_i_H_S * dstate_dw + neg_i_hamiltonian * d2state_dw2

    delta_astate = [
        delta_state;
        # delta_dstate_dw;
        # delta_d2state_dw2
    ]

    return delta_astate
end

"""
Infidelity captures the difference between
the final state and the target final state.
"""
function infidelity(final_astate)
    final_state = final_astate[STATE_INDICES]
    inner_product = TARGET_STATE_DAGGER * final_state
    fidelity = inner_product[1]^2 + inner_product[2]^2
    infidelity_ = 1 - fidelity
    return infidelity_
end

"""
Infidelity Robustness captures the sensitivity of the difference between
the final state and the target final state to the parameter OMEGA.
"""
function infidelity_robustness(final_astate)
    d2final_state_dw2 = final_astate[D2STATE_DW2_INDICES]
    inner_product = TARGET_STATE_DAGGER * d2final_state_dw2
    d2j_dw2 = 2 * sqrt(inner_product[1]^2 + inner_product[2]^2)
    return d2j_dw2 * 1e-3
end

function augment_cost(cost, multiplier, lagrange_multiplier)
    # return cost * lagrange_multiplier + (multiplier / 2) * cost ^ 2
    return cost
end

"""
Evolve a state and compute associated costs.
"""
function evolve(controls, pstate, reporter)
    # Evolve the state.
    t = 0
    dt = pstate.dt
    N = pstate.evolution_time / dt
    final_astate = pstate.initial_astate
    control_step = 1
    for i=1:N
        final_astate = rk3_step(rhs, final_astate, controls, t, dt, control_step)
        t += dt
        if t > pstate.control_eval_times[control_step + 1]
            control_step += 1
        end
    end
    final_state = final_astate[STATE_INDICES]
    # println(final_state)

    # Compute costs.
    infidelity_cost = infidelity(final_astate)
    # infidelity_robustness_cost = infidelity_robustness(final_astate)
    augmented_cost = (
        augment_cost(infidelity_cost, pstate.cost_multipliers[1], pstate.lagrange_multipliers[1])
        # + augment_cost(infidelity_robustness_cost, pstate.cost_multipliers[2], pstate.lagrange_multipliers[2])
    )

    # Report.
    reporter.costs = [
        infidelity_cost
        # infidelity_robustness_cost
    ]
    reporter.cost = augmented_cost

    return augmented_cost
end

function main()
    # Grab relevant information.
    control_count = CONTROL_COUNT
    control_eval_count = CONTROL_EVAL_COUNT
    controls = INITIAL_CONTROLS
    cost_multipliers = INITIAL_COST_MULTIPLIERS
    dt = DT
    evolution_time = EVOLUTION_TIME
    control_eval_times = Array(range(0., stop=evolution_time, length=control_eval_count))
    initial_astate = INITIAL_ASTATE
    iteration_count = ITERATION_COUNT
    lagrange_multipliers = INITIAL_LAGRANGE_MULTIPLIERS
    learning_rate = LEARNING_RATE
    save = SAVE
    # Construct the program state.
    pstate = PState(
        control_count, control_eval_count, control_eval_times,
        cost_multipliers, dt, evolution_time,
        initial_astate, lagrange_multipliers,
    )
    reporter = Reporter(0, INITIAL_COSTS)
    # Create a wrapper around the cost function. Zygote takes gradients with repsect to all arguments.
    evolve_(controls_) = evolve(controls_, pstate, reporter)
    # Generate save file path
    if save
        save_file_path = generate_save_file_path(EXPERIMENT_NAME, SAVE_PATH)
        println("saving this optimization to $save_file_path")
        h5write(save_file_path, "controls", zeros(iteration_count, control_eval_count, control_count))
    end

    # AL loop
    # for i = 1:5
    #     # iLQR loop
    #     println("costs")
    #     println(reporter.costs)
    #     println("cost_multipliers")
    #     println(pstate.cost_multipliers)
    #     println("lagrange_multipliers")
    #     println(pstate.lagrange_multipliers)

    for iteration = 1:iteration_count
        (dcontrols,) = Zygote.gradient(evolve_, controls)
        dcontrols_norm = norm(dcontrols)
        controls = controls - dcontrols * learning_rate
        @printf("%05d %f %f\n", iteration, reporter.cost, dcontrols_norm,)
        if save
            h5open(save_file_path, "r+") do save_file
                save_file["controls"][iteration, :, :] = controls
            end
        end
    end

    #     pstate.cost_multipliers = pstate.cost_multipliers * COST_MULTIPLIER_STEP
    #     pstate.lagrange_multipliers = (
    #         pstate.lagrange_multipliers + pstate.cost_multipliers .* reporter.costs
    #     )
    # end

    # evolve_(controls)
    # cfg = ForwardDiff.GradientConfig(evolve_, controls)
    # out = zeros(11, control_eval_count, control_count)
    # ForwardDiff.gradient!(out[11, :, :], evolve_, controls, cfg)
    # Zygote.gradient(evolve_, controls)
    # function test()
    #     for i=1:10
    #         # evolve_(controls)
    #         # ForwardDiff.gradient!(out[i, :, :], evolve_, controls, cfg)
    #         Zygote.gradient(evolve_, controls)
    #     end
    #     return
    # end

    # @time test()
end

end


module LazyArrays

# Use README as the docstring of the module:
@doc read(joinpath(dirname(@__DIR__), "README.md"), String) LazyArrays

using Base, Base.Broadcast, LinearAlgebra, FillArrays, StaticArrays, ArrayLayouts, MatrixFactorizations, SparseArrays
import LinearAlgebra.BLAS

import Base: AbstractArray, AbstractMatrix, AbstractVector,
        ReinterpretArray, ReshapedArray, AbstractCartesianIndex, Slice,
             RangeIndex, BroadcastStyle, copyto!, length, broadcastable, axes,
             getindex, eltype, tail, IndexStyle, IndexLinear, getproperty,
             *, +, -, /, \, ==, isinf, isfinite, sign, angle, show, isless,
         fld, cld, div, min, max, minimum, maximum, mod,
         <, ≤, >, ≥, promote_rule, convert, copy,
         size, step, isempty, length, first, last, ndims,
         getindex, setindex!, setindex, intersect, @_inline_meta, inv,
         sort, sort!, issorted, sortperm, diff, accumulate, cumsum, sum, in, broadcast,
         eltype, parent, real, imag,
         conj, transpose, adjoint, permutedims, vec,
         exp, log, sqrt, cos, sin, tan, csc, sec, cot,
                   cosh, sinh, tanh, csch, sech, coth,
                   acos, asin, atan, acsc, asec, acot,
                   acosh, asinh, atanh, acsch, asech, acoth, (:),
         AbstractMatrix, AbstractArray, checkindex, unsafe_length, OneTo, one, zero,
        to_shape, _sub2ind, print_matrix, print_matrix_row, print_matrix_vdots,
      checkindex, Slice, @propagate_inbounds, @_propagate_inbounds_meta,
      _in_range, _range, _rangestyle, Ordered,
      ArithmeticWraps, floatrange, reverse, unitrange_last,
      AbstractArray, AbstractVector, axes, (:), _sub2ind_recurse, broadcast, promote_eltypeof,
      similar, @_gc_preserve_end, @_gc_preserve_begin,
      @nexprs, @ncall, @ntuple, tuple_type_tail,
      all, any, isbitsunion, issubset, replace_with_centered_mark, replace_in_print_matrix,
      unsafe_convert, strides, union, map, searchsortedfirst, searchsortedlast, searchsorted

import Base.Broadcast: BroadcastStyle, AbstractArrayStyle, Broadcasted, broadcasted,
                        combine_eltypes, DefaultArrayStyle, instantiate, materialize,
                        materialize!, eltypes

import LinearAlgebra: AbstractTriangular, AbstractQ, checksquare, pinv, fill!, tilebufsize, Abuf, Bbuf, Cbuf, dot, factorize, qr, lu, cholesky,
                        norm2, norm1, normInf, normp, normMinusInf, diag, det, logabsdet, tr, AdjOrTrans, triu, tril,
                        lmul!, rmul!

import LinearAlgebra.BLAS: BlasFloat, BlasReal, BlasComplex

import FillArrays: AbstractFill, getindex_value

import StaticArrays: StaticArrayStyle

import ArrayLayouts: MatMulVecAdd, MatMulMatAdd, MulAdd, Lmul, Rmul, Ldiv, Dot, Mul, _inv,
                        transposelayout, conjlayout, sublayout, triangularlayout, triangulardata,
                        reshapedlayout, diagonallayout, tridiagonallayout, symtridiagonallayout, bidiagonallayout, symmetriclayout, hermitianlayout,
                        adjointlayout, sub_materialize, mulreduce,
                        check_mul_axes, _mul_eltype, check_ldiv_axes, ldivaxes, colsupport, rowsupport,
                        _fill_lmul!, scalarone, scalarzero, fillzeros, zero!, layout_getindex, _copyto!,
                        AbstractQLayout, StridedLayout, layout_replace_in_print_matrix

import Base: require_one_based_indexing, oneto

export Mul, Applied, MulArray, MulVector, MulMatrix, InvMatrix, PInvMatrix,
        Hcat, Vcat, Kron, BroadcastArray, BroadcastMatrix, BroadcastVector, cache, Ldiv, Inv, PInv, Diff, Cumsum, Accumulate,
        applied, materialize, materialize!, ApplyArray, ApplyMatrix, ApplyVector, apply, @~, LazyArray,
        PaddedArray, PaddedVector, PaddedMatrix

include("lazyapplying.jl")
include("lazybroadcasting.jl")
include("linalg/linalg.jl")
include("cache.jl")
include("lazyconcat.jl")
include("lazysetoperations.jl")
include("lazyoperations.jl")
include("lazymacro.jl")

# support x^2
Base.to_power_type(x::LazyArray) = x

# Special broadcasting for BlockArrays.jl
map(::typeof(length), A::BroadcastArray{OneTo{Int},1,Type{OneTo}}) = A.args[1]
map(::typeof(length), A::BroadcastArray{<:Fill,1,Type{Fill},<:NTuple{2,Any}}) = A.args[2]
map(::typeof(length), A::BroadcastArray{<:Zeros,1,Type{Zeros}}) = A.args[1]
map(::typeof(length), A::BroadcastArray{<:Vcat,1,Type{Vcat}}) = broadcast(+,map.(length,A.args)...)
broadcasted(::LazyArrayStyle{1}, ::typeof(length), A::BroadcastArray{OneTo{Int},1,Type{OneTo}}) = A.args[1]
broadcasted(::LazyArrayStyle{1}, ::typeof(length), A::BroadcastArray{<:Fill,1,Type{Fill},<:NTuple{2,Any}}) = A.args[2]

end # module


# # Helper functions
to_point2(a::Vector{<: Number}) = Point2f(a[1], a[2])

"""
    geo2basic(obj::GeoInterface)

Converts any GeoInterface object to the equivalent GeometryBasics, which can be plotted by Makie.
"""
function geo2basic(poly::GeoInterface.AbstractPolygon)
    polygon_coordinates = GeoInterface.coordinates(poly)
    linestrings = map(x-> to_point2.(x), polygon_coordinates)
    return GeometryBasics.Polygon(linestrings[1], linestrings[2:end])
end

function geo2basic(poly::Vector{Vector{Vector{Float64}}})
    linestrings = map(x-> to_point2.(x), poly)
    return GeometryBasics.Polygon(linestrings[1], linestrings[2:end])
end

function geo2basic(linestring::Vector{Vector{Float64}})
    return to_point2.(linestring)
end

function geo2basic(point::GeoInterface.AbstractPoint)
    return to_point2(GeoInterface.coordinates(point))
end

function geo2basic(multi_poly::GeoInterface.AbstractMultiPolygon)
    polygons = GeoInterface.coordinates(multi_poly)
    return geo2basic.(polygons)
end

geo2basic(feature::GeoInterface.Feature) = geo2basic(GeoInterface.geometry(feature))
function geo2basic(feature::GeoInterface.AbstractLineString)
    return GeometryBasics.LineString(geo2basic(GeoInterface.coordinates(feature)))
end

to_multipoly(poly::Polygon) = GeometryBasics.MultiPolygon([poly])
to_multipoly(mp::MultiPolygon) = mp
to_multipoly(any) = GeometryBasics.MultiPolygon(any)

# Only converts polygons and multipolygons
function geo2basic(fc::GeoInterface.AbstractFeatureCollection)
    return map(geo2basic, GeoInterface.features(fc))
end

function convert_arguments(P::Type{<: Union{Poly, Mesh}}, geom::GeoInterface.AbstractGeometry)
    return convert_arguments(P, geo2basic(geom))
end

function convert_arguments(P::Type{<:Poly}, geom::GeoInterface.AbstractFeatureCollection)
    return convert_arguments(P, to_multipoly.(geo2basic(geom)))
end

# set the default plot type for Vectors of polygons,
# so that they are plotted using the most efficient method!
plottype(::Vector{<: GeoInterface.AbstractMultiPolygon}) = Mesh
plottype(::Vector{<: GeoInterface.AbstractPolygon}) = Mesh


"""
    train!(model::AbstractGPModel, X::AbstractMatrix, y::AbstractArray; obsdim = 1, iterations::Int=10,callback=nothing,conv=0)
    train!(model::AbstractGPModel, X::AbstractVector, y::AbstractArray; iterations::Int=20,callback=nothing,conv=0)

Function to train the given GP `model`.

## Keyword Arguments

there are options to change the number of max iterations,
- `iterations::Int` : Number of iterations (not necessarily epochs!)for training
- `callback::Function` : Callback function called at every iteration. Should be of type `function(model,iter) ...  end`
- `conv::Function` : Convergence function to be called every iteration, should return a scalar and take the same arguments as `callback`
"""
function train!(
    m::OnlineSVGP,
    X::AbstractMatrix,
    y::AbstractArray,
    state=nothing;
    iterations::Int=20,
    callback=nothing,
    convergence=nothing,
    obsdim::Int=1,
)
    return train!(
        m,
        KernelFunctions.vec_of_vecs(X; obsdim=obsdim),
        y,
        state;
        iterations=iterations,
        callback=callback,
    )
end

function train!(
    m::OnlineSVGP,
    X::AbstractVector,
    y::AbstractArray,
    state=nothing;
    iterations::Int=20,
    callback::Union{Nothing,Function}=nothing,
    conv::Union{Nothing,Function}=nothing,
)
    iterations > 0 || error("Number of iterations should be positive")
    X, _ = wrap_X(X)
    y = check_data!(y, likelihood(m))

    data = wrap_data(X, y) # Set the data in the model
    if is_stochastic(m)
        0 < batchsize(inference(m)) <= n_sample(data) || error(
            "The size of mini-batch $(batchsize(inference(m))) is incorrect (negative or bigger than number of samples), please set `batchsize` correctly in the inference object",
        )
        set_ρ!(model, n_sample(data) / batchsize(inference))
    else
        set_batchsize!(inference(m), n_sample(data))
    end

    if n_iter(m) == 0 # The first time data is seen, initialize all parameters
        init_online_model(m, X)
    else
        state = save_old_parameters!(m, state)
        updateZs!(m, X)
    end
    state = isnothing(state) ? init_state(m) : state

    # model.evol_conv = [] # Array to check on the evolution of convergence
    local_iter = 1
    conv = Inf

    while true # Loop until one condition is matched
        try # Allow for keyboard interruption without losing the model
            if local_iter == 1 # This is needed for dealing with change of sizes
                if n_iter(m) == 0 # First iteration Zₐ is empty
                    state = compute_kernel_matrices(m, state, X, true)
                else
                    state = compute_old_matrices(m, state, X)
                end
                local_vars = local_updates!(
                    state.local_vars,
                    likelihood(m),
                    y,
                    mean_f(m, state.kernel_matrices),
                    var_f(m, state.kernel_matrices),
                )
                # We compute the updates given current data but with
                # the previous inducing points to get directly
                # the right size for the variational parameters
                ∇E_μs = ∇E_μ(m, y, local_vars)
                ∇E_Σs = ∇E_Σ(m, y, local_vars)
                state = compute_kernel_matrices(m, state, X, true)
                natural_gradient!.(
                    m.f,
                    ∇E_μs,
                    ∇E_Σs,
                    ρ(m),
                    opt(inference(m)),
                    Zviews(m),
                    state.kernel_matrices,
                    state.opt_state,
                )
                state = global_update!(m, state)

            else
                state = update_parameters!(m, state, X, y) # Update all the variational parameters
            end
            set_trained!(m, true)
            if !isnothing(callback)
                callback(m, state, n_iter(m)) # Use a callback method if given by user
            end
            if (n_iter(m) % m.atfrequency == 0) && n_iter(m) >= 3
                state = update_hyperparameters!(m, state, X, y) # Update the hyperparameters
            end
            if verbose(m) > 2 || (verbose(m) > 1 && local_iter % 10 == 0)
                print("Iteration : $(n_iter(m)), ")
                print("ELBO is : $(objective(m, state, y))")
                print("\n")
                println("number of points : $(dim(m[1]))")
            end
            ### Print out informations about the convergence
            local_iter += 1
            m.inference.n_iter += 1
            local_iter <= iterations || break # Check if the number of maximum iterations has been reached
        # (iter < model.nEpochs && conv > model.ϵ) || break; # Check if any condition has been broken
        catch e
            if e isa InterruptException
                @warn "Training interrupted by user at iteration $local_iter"
                break
            else
                rethrow(e)
            end
        end
    end
    if verbose(m) > 0
        println(
            "Training ended after",
            local_iter,
            " iterations. Total number of iterations ",
            n_iter(model),
        )
    end
    state = compute_kernel_matrices(m, state, X, true) #Compute final version of the matrices for prediction
    set_trained!(m, true)
    return m, state
end

# Update all variational parameters of the online sparse
# variational GP Model
function update_parameters!(model::OnlineSVGP, state, X, y)
    state = compute_kernel_matrices(model, state, X) #Recompute the matrices if necessary (always for the stochastic case, or when hyperparameters have been updated)
    state = variational_updates(model, state, y)
    return state
end

function updateZs!(m::OnlineSVGP, x)
    for gp in m.f
        gp.Z = InducingPoints.updateZ(gp.Z, gp.Zalg, x; kernel=kernel(gp))
        gp.post.dim = length(Zview(gp))
    end
    setHPupdated!(inference(m), true)
    return nothing
end

function save_old_parameters!(m::OnlineSVGP, state)
    opt_state =
        map(m.f, state.opt_state, state.kernel_matrices) do gp, opt_state, kernel_matrices
            previous_gp = save_old_gp!(gp, opt_state.previous_gp, kernel_matrices.K)
            merge(opt_state, (; previous_gp))
        end
    return merge(state, (; opt_state))
end

function save_old_gp!(gp::OnlineVarLatent{T}, previous_gp, K) where {T}
    gp.Zₐ = deepcopy(gp.Z)
    gp.Z = InducingPoints.remove_point(Random.GLOBAL_RNG, gp.Z, gp.Zalg, Matrix(K))# Matrix(pr_cov(gp)))
    invDₐ = Symmetric(-2.0 * nat2(gp) - inv(K)) # Compute Σ⁻¹ₐ - K⁻¹ₐ
    prevη₁ = nat1(gp)
    prev𝓛ₐ = (-logdet(cov(gp)) + logdet(K) - dot(mean(gp), nat1(gp))) / 2
    return merge(previous_gp, (; invDₐ, prevη₁, prev𝓛ₐ))
end

function init_online_model(m::OnlineSVGP{T}, x) where {T<:Real}
    m.f = ntuple(n_latent(m)) do i
        init_online_gp!(m.f[i], x)
    end
    return setHPupdated!(inference(m), false)
end

function init_online_gp!(gp::OnlineVarLatent{T}, x, jitt::T=T(jitt)) where {T}
    Z = InducingPoints.initZ(gp.Zalg, x; kernel=kernel(gp))
    k = length(Z)
    post = OnlineVarPosterior{T}(k)
    prior = GPPrior(kernel(gp), pr_mean(gp))
    return OnlineVarLatent(
        prior, post, Z, typeof(Z)([]), gp.Zalg, gp.Zupdated, gp.opt, gp.Zopt
    )
end

function compute_old_matrices(m::OnlineSVGP{T}, state, x) where {T}
    kernel_matrices = if haskey(state, :kernel_matrices)
        state.kernel_matrices
    else
        ntuple(x -> (;), n_latent(m))
    end
    kernel_matrices = map(m.f, kernel_matrices) do gp, k_mat
        return compute_old_matrices(gp, k_mat, x, T(jitt))
    end
    return merge(state, (; kernel_matrices))
end

function compute_old_matrices(gp::OnlineVarLatent, state, X::AbstractVector, jitt::Real)
    K = cholesky(kernelmatrix(kernel(gp), gp.Zₐ) + jitt * I)
    Knm = kernelmatrix(kernel(gp), X, gp.Zₐ)
    κ = Knm / K
    K̃ = kernelmatrix_diag(kernel(gp), X) .+ jitt - diag_ABt(κ, Knm)
    all(>(0), K̃) || error("K̃ has negative values")
    return merge(state, (; K, Knm, κ, K̃))
end


module Seleroute

using JuMP
using LightGraphs, MetaGraphs
using Compose, Colors, GraphPlot, Cairo, Fontconfig

using Printf, Statistics, LinearAlgebra, SparseArrays, DelimitedFiles, Random
import Base: summary, getindex
import LightGraphs: edges, vertices
# import JuMP: copy_model

const MOF = MOI.FileFormats

const AbstractSimpleGraph = LightGraphs.AbstractSimpleGraph

inedges(g, v) = [Edge(iv, v) for iv in inneighbors(g, v)]
outedges(g, v) = [Edge(v, ov) for ov in outneighbors(g, v)]

# https://github.com/JuliaOpt/JuMP.jl/pull/1982
function write_LP(io, model::Model)
    lp_model = MOF.Model(format = MOF.FORMAT_LP)
    MOI.copy_to(lp_model, backend(model))
    MOI.write_to_file(lp_model, io)
end

function write_MOF(io, model::Model)
    mof_model = MOF.Model(format = MOF.FORMAT_MOF)
    MOI.copy_to(mof_model, backend(model))
    MOI.write_to_file(mof_model, io)
end

# https://github.com/JuliaOpt/JuMP.jl/issues/2019
function Base.getindex(reference_map::ReferenceMap, container::Union{Containers.DenseAxisArray, Containers.SparseAxisArray})
    return getindex.(reference_map, container)
end

include("modeltype.jl")
include("data/parameters.jl")
include("data/model.jl")
include("data/solution.jl")
include("data/topology.jl")
include("export.jl")
include("basemodels/helpers.jl")
include("basemodels/knowncapacities.jl")
include("basemodels/knowncapacities_flow.jl")
include("basemodels/knowncapacities_path.jl")
include("basemodels/objectives.jl")
include("compute/compute.jl")
include("compute/metrics_evaluate.jl")
include("compute/certain/certain.jl")
include("compute/certain/certain_cg.jl")
include("compute/certain/mmf.jl")
include("compute/oblivious/iter.jl")
include("compute/oblivious/iter_pathcg.jl")
include("compute/oblivious/rr.jl")
include("compute/oblivious/rr_flow.jl")
include("compute/oblivious/rr_path.jl")
include("compute/oblivious/rr_pathcg.jl") # Must be included after iter_pathcg.jl due to shared helpers.
include("output/plots.jl")
include("output/reports.jl")
include("experiments/random.jl")

# Export all symbols. Code copied from JuMP.
const _EXCLUDE_SYMBOLS = [Symbol(@__MODULE__), :eval, :include]

for sym in names(@__MODULE__, all=true)
    sym_string = string(sym)
    if sym in _EXCLUDE_SYMBOLS || startswith(sym_string, "_")
        continue
    end
    if !(Base.isidentifier(sym) || (startswith(sym_string, "@") &&
         Base.isidentifier(sym_string[2:end])))
       continue
    end
    @eval export $sym
end

end


############################################################
# judiProjection #############################################
############################################################

# Authors: Philipp Witte (pwitte@eos.ubc.ca), Henryk Modzelewski (hmodzelewski@eos.ubc.ca)
# Date: January 2017

export judiProjection, judiProjectionException

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

# Type for linear operator representing  Pr*A(m)^-1*Ps,
# i.e. it includes source and receiver projections
struct judiProjection{DDT<:Number,RDT<:Number} <: joAbstractLinearOperator{DDT,RDT}
    name::String
    m::Integer
    n::Integer
    info::Info
    geometry::Geometry
end

mutable struct judiProjectionException <: Exception
    msg :: String
end

############################################################
## Constructor
"""
    judiProjection(info, geometry)

Projection operator for sources/receivers to restrict or inject data at specified locations.\\
`info` is an `Info` structure and `geometry` is a `Geometry` structure with either source or\\
receiver locations.

Examples
========

`F` is a modeling operator of type `judiModeling` and `q` is a seismic source of type `judiVector`:

    Pr = judiProjection(info, rec_geometry)
    Ps = judiProjection(info, q.geometry)

    dobs = Pr*F*Ps'*q
    qad = Ps*F'*Pr'*dobs

"""
function judiProjection(info::Info, geometry::Geometry; DDT::DataType=Float32, RDT::DataType=DDT)
    (DDT == Float32 && RDT == Float32) || throw(judiProjectionException("Domain and range types not supported"))
    m = n_samples(geometry, info)
    n = info.n * sum(info.nt)

    return judiProjection{Float32,Float32}("restriction operator",m,n,info,geometry)
end

############################################################
## overloaded Base functions

# conj(judiProjection)
conj(A::judiProjection{DDT,RDT}) where {DDT,RDT} =
    judiProjection{DDT,RDT}("conj("*A.name*")",A.m,A.n,A.info,A.geometry)

# transpose(judiProjection)
transpose(A::judiProjection{DDT,RDT}) where {DDT,RDT} =
    judiProjection{DDT,RDT}("injection operator",A.n,A.m,A.info,A.geometry)

adjoint(A::judiProjection{DDT,RDT}) where {DDT,RDT} =
    judiProjection{DDT,RDT}("injection operator",A.n,A.m,A.info,A.geometry)

############################################################
## overloaded Base *(...judiProjection...)

# *(judiProjection,judiVector)
function *(A::judiProjection{ADDT,ARDT},v::judiVector{vDT, AT}) where {ADDT,ARDT,vDT, AT}
    A.n == size(v,1) || throw(judiProjectionException("Shape mismatch: A:$(size(A)), v: $(size(v))"))
    compareGeometry(A.geometry,v.geometry) == true || throw(judiProjectionException("geometry mismatch"))
    jo_check_type_match(ADDT,vDT,join(["DDT for *(judiProjection,judiVector):",A.name,typeof(A),vDT]," / "))
    V = judiRHS(A.info,v.geometry,v.data)
    jo_check_type_match(ARDT,eltype(V),join(["RDT from *(judiProjection,judiVector):",A.name,typeof(A),eltype(V)]," / "))
    return V
end

# *(judiProjection,vec)
function *(A::judiProjection{ADDT,ARDT}, v::AbstractVector{vDT}) where {ADDT,ARDT,vDT}
    A.n == size(v,1) || throw(judiProjectionException("Shape mismatch: A:$(size(A)), v: $(size(v))"))
    jo_check_type_match(ADDT,vDT,join(["DDT for *(judiProjection,judiVector):",A.name,typeof(A),vDT]," / "))
    V = judiRHS(A.info,A.geometry, process_input_data(v, A.geometry, A.info))
    jo_check_type_match(ARDT,eltype(V),join(["RDT from *(judiProjection,judiVector):",A.name,typeof(A),eltype(V)]," / "))
    return V
end

# *(judiProjection, judiModeling)
function *(A::judiProjection{CDT,ARDT}, B::judiModeling{BDDT,CDT}) where {ARDT,BDDT,CDT}
    A.n == size(B,1) || throw(judiProjectionException("Shape mismatch: A:$(size(A)), B: $(size(B))"))
    compareInfo(A.info, B.info) == true || throw(judiProjectionException("info mismatch"))
    return judiPDE("judiProjection*judiModeling",B.info,B.model,A.geometry;options=B.options,DDT=CDT,RDT=ARDT)
end

function *(A::judiProjection{CDT,ARDT}, B::judiModelingAdjoint{BDDT,CDT}) where {ARDT,BDDT,CDT}
    A.n == size(B,1) || throw(judiProjectionException("Shape mismatch: A:$(size(A)), B: $(size(B))"))
    compareInfo(A.info, B.info) == true || throw(judiProjectionException("info mismatch"))
    return judiPDEadjoint("judiProjection*judiModelingAdjoint",B.info,B.model,A.geometry;options=B.options,DDT=CDT,RDT=ARDT)
end

############################################################
## Additional overloaded functions

# Subsample Modeling operator
function subsample(P::judiProjection{ADDT,ARDT}, srcnum) where {ADDT,ARDT}
    geometry = subsample(P.geometry,srcnum)     # Geometry of subsampled data container
    info = Info(P.info.n, length(srcnum), P.info.nt[srcnum])
    return judiProjection(info, geometry;DDT=ADDT,RDT=ARDT)
end

getindex(P::judiProjection,a) = subsample(P,a)


using ..ChemistryFeaturization.Utils.ElementFeatureUtils
using DataFrames
using ..ChemistryFeaturization: elements
using ..ChemistryFeaturization.AbstractType: AbstractCodec, AbstractAtoms
using ..ChemistryFeaturization.Codec: OneHotOneCold

include("abstractfeatures.jl")

"""
    ElementFeatureDescriptor

Describe features associated with individual atoms that depend only upon their elemental identity

## Fields
- `name::String`: Name of the feature
- `encoder_decoder::AbstractCodec`: Codec defined which handles the feature's encoding and decoding logic
- `categorical::Bool`: flag for whether the feature is categorical or continuous-valued
- `lookup_table::DataFrame`: table containing values of feature for every encodable element
"""
struct ElementFeatureDescriptor{C<:AbstractCodec} <: AbstractAtomFeatureDescriptor
    name::String
    encoder_decoder::C
    categorical::Bool
    lookup_table::DataFrame
    function ElementFeatureDescriptor(
        feature_name::String,
        encoder_decoder::AbstractCodec,
        categorical::Bool,
        lookup_table = atom_data_df,
    )
        colnames = names(lookup_table)
        @assert feature_name in colnames && "Symbol" in colnames "Your lookup table must have a column called :Symbol and one with the same name as your feature to be usable!"

        lookup_table = lookup_table[:, ["Symbol", feature_name]]
        dropmissing!(lookup_table)

        new{typeof(encoder_decoder)}(
            feature_name,
            encoder_decoder,
            categorical,
            lookup_table,
        )
    end
end

"""
    ElementFeatureDescriptor(feature_name, lookup_table, nbins, logspaced, categorical)

Construct a feature object that encodes features associated with individual atoms that depend only upon their elemental identity.
If a Codec isn't explicity specified, [OneHotOneCold](@ref) with [default_efd_encode](@ref) and [default_efd_decode](@ref)
as the encoding and decoding functions respectively is the default choice.

## Required Arguments
- `name::String`: the name of the feature
- `lookup_table::DataFrame`: table containing values of feature for every encodable element (defaults to standard included data table, but must be included for custom features)

## Keyword Arguments
All will populate with default choices if not included.
- `nbins::Integer`: Number of bins to use for one-cold decoding of continuous-valued features
- `logspaced::Bool`: whether onehot-style bins should be logarithmically spaced or not
- `categorical::Bool`: flag for whether the feature is categorical or continuous-valued
"""
function ElementFeatureDescriptor(
    feature_name::String,
    lookup_table::DataFrame = atom_data_df;
    nbins::Integer = default_nbins,
    logspaced::Bool = default_log(feature_name, lookup_table),
    categorical::Bool = default_categorical(feature_name, lookup_table),
)
    lookup_table = lookup_table[:, ["Symbol", feature_name]]
    dropmissing!(lookup_table)

    bins = get_bins(
        feature_name,
        lookup_table;
        nbins = nbins,
        logspaced = logspaced,
        categorical = categorical,
    )

    ElementFeatureDescriptor(
        feature_name,
        OneHotOneCold(categorical, bins),
        categorical,
        lookup_table,
    )
end

# pretty printing, long version
function Base.show(io::IO, ::MIME"text/plain", efd::ElementFeatureDescriptor)
    st = "ElementFeature $(efd.name):\n   categorical: $(efd.categorical)\n   encoded length: $(output_shape(efd))"
    print(io, st)
end

encodable_elements(efd::ElementFeatureDescriptor) = efd.lookup_table[:, :Symbol]

function get_value(efd::ElementFeatureDescriptor, a::AbstractAtoms)
    @assert all([el in encodable_elements(efd) for el in elements(a)]) "Feature $(efd.name) cannot encode some element(s) in this structure!"

    feature_vals = efd.lookup_table[:, [:Symbol, Symbol(efd.name)]]
    map(
        el ->
            getproperty(feature_vals[feature_vals.Symbol.==el, :][1, :], Symbol(efd.name)),
        elements(a),
    )
end


"""
    points, lines, triangles, tetrahedra, domains=read_nastran(filename)

Read simplicial nastran mesh.

# Note
There is currently no sanity check for non-simplicial elements. These elements are just skipped while reading.
"""
function read_nastran(filename)
  number_of_points=0
  number_of_lines=0
  number_of_triangles=0
  number_of_tetrahedra=0
  name_tags=Dict()

  #first read to get number of points
  open(filename) do fid
    while !eof(fid)
      line=readline(fid)
      if line[1]=='$' #line is a comment
        #comments are freqeuntly used to leave more information
        #especially the domain tags
        #this is not standardized so it depends on the software used to write the file

        #ANSA-style tag comment
        #example:
        #$ANSA_NAME_COMMENT;2;PSHELL;TA_IN_Inlet;;NO;NO;NO;NO;
        if length(line)>=18 && line[2:18]=="ANSA_NAME_COMMENT"
          data=split(line[2:end],";")
          if data[3]=="PSOLID" || data[3]=="PSHELL"
            name_tags[data[2]]=data[4]
          end
        end

        #Centaur-style tag comment
        #example:
        #$HMNAME COMP                   1"Walls"
        if length(line)>=12 && line[2:12]=="HMNAME COMP"
          data=split(strip(line[13:end]),'"')
          name_tags[data[1]]=data[2]
        end

        continue
      end
      if length(line)<8 #line is too short probably it is the ENDDATA tag.
        continue
      end
      head=line[1:8]
      if head =="GRID    " || head == "GRID*   " || head[1:5] == "GRID," #TODO allow for whitespace between commas in free format
        number_of_points+=1
      elseif (head[1:6] == "CTRIA3") || (head[1:6] == "CTRIA6")
        number_of_triangles+=1
      elseif head[1:6] == "CTETRA"
        number_of_tetrahedra+=1
      end
    end
  end
  tri_idx=0
  tet_idx=0
  #TODO: sanity check for non-consecutive indexing
  points=Array{Float64}(undef,3,number_of_points)
  lines=Array{Array{UInt32,1}}(undef,number_of_lines)
  triangles=Array{Array{UInt32,1}}(undef, number_of_triangles)
  tetrahedra=Array{Array{UInt32,1}}(undef, number_of_tetrahedra)
  domains=Dict()
  surface_tags=Dict()
  volume_tags=Dict()
  #second read to get data
  open(filename) do fid
    while !eof(fid)
      line=readline(fid)
      if line[1]=='$' #line is a comment
        continue
      end
      # data = parse_nas_line(line)
      # if length(data)==0 #line is to short probaboly it is the ENDDATA tag.
      #   continue
      # end

      if length(line)<8 #line is too short probably it is the ENDDATA tag.
        continue
      end

      free=false #sentinel value for checking for free format
      if occursin(",",line)
        free=true
      end

      head=line[1:8]

      if head =="GRID    " #grid point short format
        data = parse_nas_line(line)
        idx=parse(UInt,data[2])
        points[1,idx]=parse_nas_number(data[4])
        points[2,idx]=parse_nas_number(data[5])
        points[3,idx]=parse_nas_number(data[6])

      elseif head =="GRID*   " #grid point long format
        data = parse_nas_line(line,format=:long)
        idx=parse(UInt,data[2])
        points[1,idx]=parse_nas_number(data[4])
        points[2,idx]=parse_nas_number(data[5])
        line=readline(fid) #long format has two lines
        data = parse_nas_line(line,format=:long)
        points[3,idx]=parse_nas_number(data[2])

      elseif head[1:5] =="GRID,"#grid point free format
        data = parse_nas_line(line,format=:free)
        idx=parse(UInt,data[2])
        points[1,idx]=parse_nas_number(data[4])
        points[2,idx]=parse_nas_number(data[5])
        points[3,idx]=parse_nas_number(data[6])

      elseif (head[1:6] == "CTRIA3") || (head[1:6] == "CTRIA6") #Triangle
        tri_idx+=1
        if free
          data = parse_nas_line(line,format=:free)
        else
          data = parse_nas_line(line)
        end
          #TODO: find nastran file specification and see whether there is
          #something like long format for stuff other than GRID
        idx=tri_idx
        dom=strip(data[3])
        if dom in keys(name_tags)
          dom=name_tags[dom]
        else
          dom="surf"*lpad(dom,4,"0")
        end
        if dom in keys(domains)
          append!(domains[dom]["simplices"],idx)
        else
          domains[dom]=Dict("dimension"=>2,"simplices"=>UInt32[idx])
        end
        tri=[parse(UInt32,data[4]),parse(UInt32,data[5]),parse(UInt32,data[6])]
        triangles[idx]=tri

      elseif head[1:6] == "CTETRA" #Tetrahedron
        tet_idx+=1
        if free
          data = parse_nas_line(line,format=:free)
        else
          data = parse_nas_line(line)
        end
        idx=tet_idx#parse(UInt,data[2])-number_of_triangles
        dom=strip(data[3])
        if dom in keys(name_tags)
          dom=name_tags[dom]
        else
          dom="vol"*lpad(dom,4,"0")
        end
        if dom in keys(domains)
          append!(domains[dom]["simplices"],idx)
        else
          domains[dom]=Dict("dimension"=>3,"simplices"=>UInt32[idx])
        end
        tet=[parse(UInt32,data[4]),parse(UInt32,data[5]),parse(UInt32,data[6]),parse(UInt32,data[7])]
        tetrahedra[idx]=tet
      end
    end
  end

  #sanity check when using higher order simplices (CTRIA6 or CTETRA with ten nodes)
  #Comsol uses this stuff -.-
  used_idx=sort(unique(vcat(tetrahedra...)))
  if used_idx[end]==length(used_idx)
    if length(triangles)!=0
      tri_max_idx=maximum(vcat(triangles...))
    else
      tri_max_idx=0
    end
    if tri_max_idx > used_idx[end]
      println("Warning: Nastran mesh needs manual reordering of point indeces. Not all points are actually used to define simplices.")
    else
      points=points[:,1:used_idx[end]]
    end
  else
    println("Warning: Nastran mesh needs manual reordering of point indeces. Not all points are actually used to define simplices.")
  end
  return points, lines, triangles, tetrahedra, domains
end

function parse_nas_line(txt;format=:short)
  #split line
  if format==:short
    number_of_entries=length(txt)÷8
    data=Array{String}(undef,number_of_entries)
    for idx=1:number_of_entries
      data[idx]=txt[(1+8*(idx-1)):(idx*8)]
    end
  elseif format==:long
    number_of_entries=(length(txt)-8)÷16+1
    data=Array{String}(undef,number_of_entries)
    data[1]=txt[1:8]
    for idx=2:number_of_entries
      data[idx]=txt[(1+16*(idx-1)-8):((idx*16)-8)]
    end
  elseif format==:free
    data=split(replace(txt," "=>""),",")
  end
  return data
end

function parse_nas_number(txt)
  # nastran numbers may or may not have sign symbols
  # and leading 0
  # also there might be lower or upper Case indicating
  # scientific notation or just a directly a sign symbol
  # This means, everything that is correctly parsed as
  # as a julia Float is a nastran float
  # but also something exotic like "-.314+7" is valid
  # nastran and would be the same as "-0.314E+7"
  # this routine correctly parses nastran numbers

  txt=strip(txt)
  E=false
  ins=0
  if !occursin("E",txt) && !occursin("e",txt)
    for (idx,c) in enumerate(txt)
      if c=="E"
        E=true
      end
      if idx!=1 && ( (!E && c=='-') || (!E && c=='+'))
        ins=idx
        break
      end
    end
    if ins >1
      txt=txt[1:(ins-1)]*"E"*txt[ins:end]
    end
  end
  return parse(Float64,txt)
end

# function relabel_nas_domains!(domains,filename)
#   open(filename) do fid
#     while !eof(fid)
#       line=readline(fid)
#       if line[1]=='#' #line is a comment
#         continue
#       end
#       data=split(line,",")
#       old_idx=strip(data[1])
#       new_idx=strip(data[2])
#       if old_idx in keys(domains)
#         domains[new_idx]=domains[old_idx]
#         delete!(domains,old_idx)
#       else
#         println("Warning: key `$old_idx` not in domains!")
#       end
#     end
#   end
#   return domains
# end