text:
"""
   immutable struct contain the model parameters for FDFD
*  data_format : Float32 or Float64
*  nz          : depth of model
*  nx          : length of model
*  npml        : number of PML layers
*  free_surface: true -> no obsorbtion on the surface
                 false-> PML on the surface
*  Nz          : depth of model (include PML padding)
*  Nx          : length of model (include PML padding)
*  ntop        : number of padded PML layers on the top
*  spt2wfd     : auxillary index vector for computing wavefield from snapshot
*  h           : grid spatial size
*  rho         : 2D density model
*  vel         : 2D P-wave velocity model
"""
struct FdParams{Ti<:Int64, Tv<:AbstractFloat}

    data_format  :: DataType
    nz           :: Ti
    nx           :: Ti
    npml         :: Ti
    apml         :: Tv
    free_surface :: Bool
    Nz           :: Ti
    Nx           :: Ti
    ntop         :: Ti
    spt2wfd      :: Vector{Ti}
    h            :: Tv
    rho          :: Matrix{Tv}
    vel          :: Matrix{Tv}
end

"""
   Overloading the show function for FDParams
"""
function show(io::IO, params::FdParams)

    # size of model
    @printf("nz = %4d, nx = %4d, npml = %4d\n", params.nz, params.nx, params.npml)
    @printf("h  = %4.1f\n", params.h)

    # top boundary condition
    if params.free_surface
       @printf("top boundary conditon = %20s\n", "free surface")
    else
       @printf("top boundary conditon = %20s\n", "PML obsorbing")
    end

    # data presicion (Float32, Float64)
    @printf("data_format=%16s\n", params.data_format)

    return nothing
end

"""
   get the auxillary index vectors which facilitate extracting wavefield from snapshot
"""
function get_spt2wfd_index(nz::Ti, nx::Ti, npml::Ti, free_surface::Bool) where {Ti<:Int64}

    # size of wavefield
    len = nz * nx

    # allocate 3 vectors of integer
    index = zeros(Int64, len)

    # size of padded model
    if free_surface
       Nz = nz + npml
       ntop = 0
    else
       Nz = nz + npml * 2
       ntop = npml
    end
    Nx = nx + npml * 2

    # map snapshot to wavefield
    icount = 0
    for i2 = 1 : nx
        tmp= (i2 + npml - 1)*Nz + ntop

        for i1 = 1 : nz
            icount = icount + 1
            index[icount] = tmp + i1
        end
    end

    return index
end

"""
   create the parameter structure for frequency domain finite difference method
"""
function FdParams(rho, vel, free_surface::Bool, h;
         data_format=Float32, npml=20, apml=900.)

    # get the size of the model
    (nz, nx) = size(rho)
    if (nz, nx) != size(vel)
       error("model size mismatch")
    end

    # unify the data format
    rho = convert(Matrix{data_format}, rho)
    vel = convert(Matrix{data_format}, vel)
    h   = convert(data_format, h)
    apml= convert(data_format, apml)

    # size of model after padding PML layers
    if free_surface
       Nz   = nz +     npml
       ntop = 0
    else
       Nz   = nz + 2 * npml
       ntop = npml
    end
    Nx = nx + 2 * npml

    # compute auxillary index vectors
    idx_spt2wfd = get_spt2wfd_index(nz, nx, npml, free_surface)

    # call the default struct constructor
    return FdParams(data_format, nz, nx, npml, apml, free_surface, Nz, Nx, ntop, idx_spt2wfd, h, rho, vel)
end

"""
   initialize the pml damping profile in z- and x-direction
"""
function DampingProfile(Nz::Ti, Nx::Ti, npml::Ti, apml::Tv,
         free_surface::Bool, h::Tv, omega::Float64) where {Ti<:Int64, Tv<:AbstractFloat}

    # convert the data format of radian frequency
    apml  = convert(Float64, apml )
    h     = convert(Float64, h    )

    # get the damping profile in Z direction
    if free_surface
       (zdamp_integer, zdamp_stagger) = get_fsdamping_profile(Nz, npml, apml, h, omega)
    else
       (zdamp_integer, zdamp_stagger) = get_damping_profile(Nz, npml, apml, h, omega)
    end

    # get the damping profile in X direction
    (xdamp_integer, xdamp_stagger) = get_damping_profile(Nx, npml, apml, h, omega)

    # # data format convert
    # zdamp_integer = convert(Vector{Complex{params.data_format}}, zdamp_integer)
    # zdamp_stagger = convert(Vector{Complex{params.data_format}}, zdamp_stagger)
    # xdamp_integer = convert(Vector{Complex{params.data_format}}, xdamp_integer)
    # xdamp_stagger = convert(Vector{Complex{params.data_format}}, xdamp_stagger)

    # group into a struct
    return zdamp_integer, zdamp_stagger, xdamp_integer, xdamp_stagger
end

"""
   set up damping profile for PML boundary conditions, n is the model size include pml,
npml is the number of PML layers, omega is radian frequency.
"""
function get_damping_profile(n::Ti, npml::Ti, apml::Tv, h::Tv, omega::Tv) where {Ti<:Int64, Tv<:Float64}

    # constants
    pi_over_2 = 0.5 * pi   # pi/2

    # damping profile of PML
    damp_integer = ones(Complex{Float64}, n+2)
    damp_stagger = ones(Complex{Float64}, n+2)

    # PML width
    pml_width = npml * h

    # model width include PML layers
    model_width = (n - 1) * h

    # both side of damp_integer and one side of damp_stagger
    for i = 1 : npml

        # current coordinate of x
        x_integer = (i-1) * h
        x_stagger = (i-1+0.5) * h

        # distance of current grid to the edge of computation area
        dis_integer = pml_width - x_integer
        dis_stagger = pml_width - x_stagger

        # pml coefficients
        gamma_integer = apml * (1.0 - cos(dis_integer / pml_width * pi_over_2))
        gamma_stagger = apml * (1.0 - cos(dis_stagger / pml_width * pi_over_2))

        # dampping profile (one extra point) 1/(1-iγ/w)
        damp_integer[i+1] = 1.0 / (1.0 - im * gamma_integer / omega)
        damp_stagger[i+1] = 1.0 / (1.0 - im * gamma_stagger / omega)

        # symmetrical damp_integer
        damp_integer[n-i+2] = damp_integer[i+1]
    end

    # one extra point facilitate building Helmhotz operator
    damp_integer[1]   = damp_integer[2]
    damp_integer[n+2] = damp_integer[n+1]

    # right side of damp_stagger
    for i = 1 : npml + 1

        # the coordinate of the stagger grid (from right to left)
        x_stagger = model_width + 0.5 * h - (i-1)*h

        # distance of current grid to the edge of computation area
        dis_stagger = x_stagger - (model_width - pml_width)

        # pml coefficients
        gamma_stagger = apml * (1.0 - cos(dis_stagger / pml_width * pi_over_2))

        # damping profile
        damp_stagger[n-i+2] = 1.0 / (1.0 - im * gamma_stagger / omega)
    end

    damp_stagger[1]   = damp_stagger[2]
    damp_stagger[n+2] = damp_stagger[n+1]

    return damp_integer, damp_stagger
end

"""
   get the damping profile for vertical direction when free surface boundary
condition is implemented for the top boundary.
"""
function get_fsdamping_profile(Nz::Ti, npml::Ti, apml::Tv, h::Tv, omega::Tv) where {Ti<:Int64, Tv<:Float64}

    # constants
    pi_over_2 = 0.5 * pi   # pi/2

    # damping profile of PML
    damp_integer = ones(Complex{Float64}, Nz+2)
    damp_stagger = ones(Complex{Float64}, Nz+2)

    # PML width
    pml_width = npml * h

    # model width include PML layers
    model_width = (Nz - 1) * h

    # the bottom side
    for i = 1 : npml + 1

        # the coordinate of the stagger grid (from right to left)
        x_integer = model_width           - (i-1)*h
        x_stagger = model_width + 0.5 * h - (i-1)*h

        # distance of current grid to the edge of computation area
        dis_integer = x_integer - (model_width - pml_width)
        dis_stagger = x_stagger - (model_width - pml_width)

        # pml coefficients
        gamma_integer = apml * (1.0 - cos(dis_integer / pml_width * pi_over_2))
        gamma_stagger = apml * (1.0 - cos(dis_stagger / pml_width * pi_over_2))

        # damping profile
        damp_integer[Nz-i+2] = 1.0 / (1.0 - im * gamma_integer / omega)
        damp_stagger[Nz-i+2] = 1.0 / (1.0 - im * gamma_stagger / omega)
    end

    damp_integer[Nz+2] = damp_integer[Nz+1]
    damp_stagger[Nz+2] = damp_stagger[Nz+1]

    return damp_integer, damp_stagger
end

"""
   compute complex bulk modulus which support visco-acoustic wave-equation, padding
one extra layer. The bulk modulus is padded with one more outer layer.
"""
function vel2bulk(rho::Matrix{Tv}, vel::Matrix{Tv}) where {Tv<:Float64}

    # size of model
    (nz, nx)    = size(rho)

    # allocate space for bulk modulus
    bulk = zeros(Float64, nz+2, nx+2)
    for ix = 1 : nx
        for iz = 1 : nz
            bulk[iz+1,ix+1] = rho[iz,ix] * vel[iz,ix] * vel[iz,ix]
        end
    end

    # visco-acoustic case
    # bulk = zeros(Complex{data_format}, nz+2, nx+2)
    # for ix = 1 : nx
    #     for iz = 1 : nz
    #         vc = vel[iz,ix] * (1.0 - 0.5*im / q[iz,ix])
    #         mu[iz+1,ix+1] = rho[iz,ix] * vc * vc
    #     end
    # end

    # left and right
    for iz = 1 : nz
        bulk[iz+1, 1]    = bulk[iz+1, 2]
        bulk[iz+1, nx+2] = bulk[iz+1, nx+1]
    end

    # upper and bottom
    for ix = 1 : nx
        bulk[1, ix+1]    = bulk[2, ix+1]
        bulk[nz+2, ix+1] = bulk[nz+1,ix+1]
    end

    # four corners
    bulk[1, 1]       = bulk[2, 2]
    bulk[1, nx+2]    = bulk[2, nx+1]
    bulk[nz+2, 1   ] = bulk[nz+1, 2]
    bulk[nz+2, nx+2] = bulk[nz+1, nx+1]

    return bulk
end

"""
   compute buoyancy from density, the computed buoyancy is paded with one extra layer.
"""
function rho2buoyancy(rho::Matrix{Tv}) where {Tv <: Float64}

    # size of the model
    (nz, nx) = size(rho)

    # allocate space for buoyancy
    buoyancy = zeros(Float64, nz+2, nx+2)
    for ix = 1 : nx
        for iz = 1 : nz
            if rho[iz, ix] != 0.0
               buoyancy[iz+1, ix+1] = 1.0 / rho[iz,ix]
            else
               error("density is zero")
            end
        end
    end

    # left and right layer
    for iz = 1 : nz
        buoyancy[iz+1, 1]    = buoyancy[iz+1, 2]
        buoyancy[iz+1, nx+2] = buoyancy[iz+1, nx+1]
    end

    # top and bottom layer
    for ix = 1 : nx
        buoyancy[1, ix+1]    = buoyancy[2, ix+1]
        buoyancy[nz+2, ix+1] = buoyancy[nz+1, ix+1]
    end

    # four corners
    buoyancy[1, 1]       = buoyancy[2, 2   ]
    buoyancy[1, nx+2]    = buoyancy[2, nx+1]
    buoyancy[nz+2, 1]    = buoyancy[nz+1, 2]
    buoyancy[nz+2, nx+2] = buoyancy[nz+1, nx+1]

    return buoyancy
end

"""
   get the buoyancy at stagger grid by averaging the buoyancy at reference grid
"""
function buoyancy_average!(bu::OffsetArray, buoyancy::Matrix{Tv}, i::Ti, j::Ti) where {Ti<:Int64, Tv<:Float64}

    # some constants
    one_over_two  =  1.0 / 2.0
    one_over_four =  1.0 / 4.0

    # first column
    bu[-1,-1] = one_over_four * (buoyancy[i-1,j-1] + buoyancy[i,j-1] + buoyancy[i-1,j] + buoyancy[i,j])
    bu[0 ,-1] = one_over_two  * (buoyancy[i,j-1]   + buoyancy[i,j]                                    )
    bu[1 ,-1] = one_over_four * (buoyancy[i,j-1] + buoyancy[i+1,j-1] + buoyancy[i,j] + buoyancy[i+1,j])

    # second column
    bu[-1,0] = one_over_two * (buoyancy[i,j] + buoyancy[i-1,j])
    bu[0 ,0] =                 buoyancy[i,j]
    bu[1 ,0] = one_over_two * (buoyancy[i,j] + buoyancy[i+1,j])

    # third column
    bu[-1,1] = one_over_four * (buoyancy[i-1,j] + buoyancy[i,j] + buoyancy[i-1,j+1] + buoyancy[i,j+1])
    bu[0 ,1] = one_over_two  * (buoyancy[i,j] + buoyancy[i,j+1]                                      )
    bu[1 ,1] = one_over_four * (buoyancy[i,j] + buoyancy[i+1,j] + buoyancy[i,j+1] + buoyancy[i+1,j+1])

    return nothing
end

"""
   the kernal of building Helmhotz operator, output the LU factorization of Helmhotz operator
"""
function get_helmholtz_LU(params::FdParams, omega)

    # convert radian frequency to Float64
    omega = convert(Float64, omega)

    # constants for mass acceleration averaging
    w1 = 0.6287326  # c
    w2 = 0.0928166  # d
    w3 = (1.0 - w1 - 4.0*w2) / 4.0 # (1-c-4d)/4

    # weights for stagger and rotated grid
    alpha = 0.5617365     # α
    beta  = 1.0 - alpha   # 1-α

    # omega square and one over h square
    omega2       = omega * omega
    one_over_h2  = 1.0  / (params.h * params.h)
    one_over_4h2 = 0.25 / (params.h * params.h)

    # model parameter padding to accomdate PML layers
    vel = model_padding(params.vel, params.npml, params.free_surface)
    rho = model_padding(params.rho, params.npml, params.free_surface)

    # convert the model parameter to Float64
    if params.data_format == Float32
       rho = convert(Matrix{Float64}, rho)
       vel = convert(Matrix{Float64}, vel)
    end

    # add support for visco-acoustic in the future
    # qfactor = model_padding(params.qfactor, params.npml, params.free_surface)

    # compute bulk modulus
    k        = vel2bulk(rho, vel)
    buoyancy = rho2buoyancy(rho)

    # compute the damping profile
    (zdamp_integer, zdamp_stagger, xdamp_integer, xdamp_stagger) = DampingProfile(params.Nz, params.Nx,
                                                                   params.npml, params.apml, params.free_surface, params.h, omega)

    # count the number of non-zero elements
    nnz = 0
    if params.free_surface
       nnz = nnz + 1 + 1 + 4 * 2              # four corners
       nnz = nnz + params.Nx - 2              # upper
       nnz = nnz + (params.Nx-2) * 6          # bottom

    else
       nnz = nnz + 4*4                        # four corners
       nnz = nnz + (params.Nx-2)*6*2          # upper and bottom
    end

    nnz = nnz + (params.Nz-2)*6*2             # left and right
    nnz = nnz + (params.Nz-2)*(params.Nx-2)*9 # inner part

    # allocate space for sparse Helmhotz matrix in coordinate format
    row_idx = zeros(Int64, nnz) # each row has 9 non-zeros elements
    col_idx = zeros(Int64, nnz)
    nzval   = zeros(Complex{Float64}, nnz)
    icount  = 0

    # buoyancy at stagger grid
    b = OffsetArray{Float64}(undef, -1:1, -1:1)

    for i2 = 1 : params.Nx
        for i1 = 1 : params.Nz

            # the corresponding index for buoyancy and bulk modulus as one extra more layer is padded
            i = i1 + 1
            j = i2 + 1

            # buoyancy at stagger grid
            buoyancy_average!(b, buoyancy, i, j)

            # damping parameter in z direction
            d1o = zdamp_integer[i]
            d1p = zdamp_stagger[i]
            d1m = zdamp_stagger[i-1]

            # damping parameter in x direction
            d2o = xdamp_integer[j]
            d2p = xdamp_stagger[j]
            d2m = xdamp_stagger[j-1]

            # row index
            row_val = (i2-1) * params.Nz + i1

            # node (0, 0)
            col_val = (i2-1  ) * params.Nz + i1
            icount  = icount + 1
            row_idx[icount] = row_val
            col_idx[icount] = col_val
            if (i1 == 1) && params.free_surface
               nzval[icount] = -omega2 / k[i,j]
            else
               nzval[icount] = (-w1 * omega2 / k[i,j]
                                +alpha * one_over_h2  * ( d1o * (b[1,0]*d1p + b[-1,0]*d1m)
                                                         +d2o * (b[0,1]*d2p + b[0,-1]*d2m))
                                +beta  * one_over_4h2 * ( d1o * (b[1,1]*d1p + b[-1,-1]*d1m + b[1,-1]*d1p + b[-1,1]*d1m)
                                                         +d2o * (b[1,1]*d2p + b[-1,-1]*d2m + b[-1,1]*d2p + b[1,-1]*d2m)))
            end

            # node(-1,-1)
            if (i1 > 1) && (i2 > 1)
               col_val = (i2-1-1) * params.Nz + i1-1
               icount  = icount + 1
               row_idx[icount] = row_val
               col_idx[icount] = col_val
               nzval[icount] = (-w3   * omega2 / k[i-1,j-1]
                                -beta * one_over_4h2 * (b[-1,-1]*d1o*d1m + b[-1,-1]*d2o*d2m))
            end

            # node(0,-1)
            if (i2 > 1)
               if (i1 != 1) || !params.free_surface
                  col_val = (i2-1-1) * params.Nz + i1
                  icount  = icount + 1
                  row_idx[icount] = row_val
                  col_idx[icount] = col_val
                  nzval[icount]   = (-w2 * omega2 / k[i,j-1]
                                     -alpha  * one_over_h2 * b[0,-1] * d2o * d2m
                                     -beta   * one_over_4h2* ( d2o*(b[-1,-1]*d2m + b[1,-1]*d2m)
                                                              -d1o*(b[-1,-1]*d1m + b[1,-1]*d1p)))
               end
            end

            # node(1,-1)
            if (i1 < params.Nz) && (i2 > 1)
               if (i1 != 1) || !params.free_surface
                  col_val = (i2-1-1) * params.Nz + i1+1
                  icount  = icount + 1
                  row_idx[icount] = row_val
                  col_idx[icount] = col_val
                  nzval[icount]   = (-w3   * omega2 / k[i+1,j-1]
                                     -beta * one_over_4h2 * (b[1,-1]*d1o*d1p + b[1,-1]*d2o*d2m))
               end
            end

            # node(-1, 0)
            if (i1 > 1)
               col_val = (i2-1  ) * params.Nz + i1-1
               icount = icount + 1
               row_idx[icount] = row_val
               col_idx[icount] = col_val
               nzval[icount]   = (-w2 * omega2 / k[i-1,j]
                                  -alpha * one_over_h2  * b[-1,0] * d1o * d1m
                                  -beta  * one_over_4h2 * ( d1o*(b[-1,-1]*d1m + b[-1,1]*d1m)
                                                           -d2o*(b[-1,-1]*d2m + b[-1,1]*d2p)))
            end

            # node(1, 0)
            if (i1 < params.Nz)
               if (i1 != 1) || !params.free_surface
                  col_val = (i2-1  ) * params.Nz + i1+1
                  icount  = icount + 1
                  row_idx[icount] = row_val
                  col_idx[icount] = col_val
                  nzval[icount]   = (-w2 * omega2 / k[i+1,j]
                                     -alpha * one_over_h2  * b[1,0] * d1o * d1p
                                     -beta  * one_over_4h2 * ( d1o*(b[1,1]*d1p + b[1,-1]*d1p)
                                                              -d2o*(b[1,1]*d2p + b[1,-1]*d2m)))
               end
            end

            # node(-1, 1)
            if (i1 > 1) && (i2 < params.Nx)
               col_val = (i2-1+1) * params.Nz + i1-1
               icount  = icount + 1
               row_idx[icount] = row_val
               col_idx[icount] = col_val
               nzval[icount]   = (-w3 * omega2 / k[i-1,j+1]
                                  -beta * one_over_4h2 * (b[-1,1]*d1o*d1m + b[-1,1]*d2o*d2p))
            end

            # node (0, 1)
            if (i2 < params.Nx)
               if (i1 != 1) || !params.free_surface
                  col_val = (i2-1+1) * params.Nz + i1
                  icount  = icount + 1
                  row_idx[icount] = row_val
                  col_idx[icount] = col_val
                  nzval[icount]   = (-w2 * omega2 / k[i,j+1]
                                     -alpha * one_over_h2  * b[0,1]*d2o*d2p
                                     -beta  * one_over_4h2 * ( d2o*(b[-1,1]*d2p + b[1,1]*d2p)
                                                              -d1o*(b[1,1]*d1p + b[-1,1]*d1m)))
               end
            end

            # node(1, 1)
            if (i1 < params.Nz) && (i2 < params.Nx)
               if (i1 != 1) || !params.free_surface
                  col_val = (i2-1+1) * params.Nz + i1+1
                  icount  = icount + 1
                  row_idx[icount] = row_val
                  col_idx[icount] = col_val
                  nzval[icount]   = (-w3 * omega2 / k[i+1,j+1]
                                     -beta * one_over_4h2 * (b[1,1]*d1o*d1p + b[1,1]*d2o*d2p))
               end
            end

        end # i1
    end # i2

    # check the size
    if nnz != icount
       error("number of non-zero element is wrong")
    end

    # save the sparse matrix in CSC format
    H = sparse(row_idx, col_idx, nzval)

    # return the LU factorization of Helmhotz operator
    return lu(H)
end

"""
   generate source for frequency domain finite difference
"""
function Source(sz, sx, dt, params::FdParams;
                location_flag="index", ot=0.0, amp=1.0, fdom=20.0, hl=128, type_flag="ricker", p=Vector{Float32}(undef,0))

    # source location given index
    if location_flag == "index"
       isz = round(Int64, sz)
       isx = round(Int64, sx)

    # source location given as distance
    elseif location_flag == "distance"
       isz = round(Int64, sz/params.h) + 1
       isx = round(Int64, sx/params.h) + 1
    else
       error("wrong specification of source location")
    end

    # error checking
    if isz > params.nz || isx > params.nx || isz < 1 || isx < 1
       error("source located outside of modeling area")
    end

    # can't put on the free-surface
    if params.free_surface && isz == 1
       error("can't inject source on the surface")
    end

    # the auxillary index mapping the location of source to snapshot or wavefield
    src2spt = (isx+params.npml-1) * params.Nz + isz + params.ntop
    src2wfd = (isx            -1) * params.nz + isz

    # non-user provided source wavelet
    if length(p) == 0
       if type_flag == "ricker"
          p = ricker(fdom, dt)
          p = amp * p
       elseif type_flag == "miniphase"
          p = ricker(fdom, dt)
          p = convert2miniphase(p)
          p = amp * p
       elseif type_flag == "sinc"
          fc = params.data_format(fdom)
          hl = round(Int64, hl)
          p  = tapered_sinc(fc, hl, dt)
          p  = amp * p
       end
    end

    # make the data type consistent
    p = convert(Vector{params.data_format}, p)
    nt= length(p)

    # the index time range of source wavelet
    it_min = floor(Int64, ot/dt) + 1
    it_max = it_min + nt - 1

    # call the default constructor
    return Source(isz, isx, src2spt, src2wfd, dt, it_min, it_max, p)
end

"""
   get multiple sources which are put into a vector
"""
function get_multi_sources(sz::Vector, sx::Vector, dt, params::FdParams;
                           location_flag="index", ot=0.0, amp=1.0, fdom=20.0, hl=128, type_flag="ricker", p=Vector{Float32}(undef,0))

    # number of source
    ns = length(sz)
    if length(sx) != ns
       error("length of source coordinate mismatch")
    end

    # source location given index
    isz = zeros(Int64, ns)
    isx = zeros(Int64, ns)
    if location_flag == "index"
       for i = 1 : ns
           isz[i] = round(Int64, sz[i])
           isx[i] = round(Int64, sx[i])
       end

    # source location given as distance
    elseif location_flag == "distance"
       isz[i] = round(Int64, sz[i]/params.h) + 1
       isx[i] = round(Int64, sx[i]/params.h) + 1
    else
       error("wrong specification of source location")
    end

    # error checking
    for i = 1 : ns
        if isz[i] > params.nz || isx[i] > params.nx
           error("source located outside of modeling area")
        end
    end

    # starting time
    if typeof(ot) <: Real  # all sources starting at same time
       ot = ot * ones(ns)
    end
    if length(ot) != ns
       error("length of starting time ot mismatch")
    end

    # amplitude for each source
    if typeof(amp) <: Real
       amp = amp * ones(ns)
    end
    if length(amp) != ns
       error("length of amplitude is wrong")
    end

    # dominant frequency for source wavelet
    if typeof(fdom) <: Real
       fdom = fdom * ones(ns)
    end
    if length(fdom) != ns
       error("length of dominant frequency is wrong")
    end

    # multiple source function are put into a vector
    wavelet = Vector{Vector{params.data_format}}(undef, ns)
    if length(p) == 0  # user don't provided wavelet
       if typeof(type_flag) == String               # all sources share same wavelet
          wavelet_type = Vector{String}(undef, ns)
          wavelet_type[:] .= type_flag
          type_flag    = wavelet_type

       elseif typeof(type_flag) == Vector{String}   # use different wavelet
          if length(type_flag) != ns
             error("length of type_flag mismatch")
          end
       end

       # generate the specific source wavelet
       for i = 1 : ns
           if type_flag[i] == "ricker"
              tmp = amp[i] * ricker(fdom[i], dt)
              wavelet[i] = convert(Vector{params.data_format}, tmp)

           elseif type_flag[i] == "miniphase"
              tmp = ricker(fdom[i], dt)
              tmp = amp[i] * convert2miniphase(tmp)
              wavelet[i] = convert(Vector{params.data_format}, tmp)

           elseif type_flag[i] == "sinc"
              hl   = round(Int64, hl)
              tmp  = amp[i] * tapered_sinc(fdom[i], hl, dt)
              wavelet[i] = convert(Vector{params.data_format}, tmp)
           end
       end

    elseif eltype(p) <: AbstractFloat   # all sources share same user-given wavelet
       for i = 1 : ns
           wavelet[i] = convert(Vector{params.data_format}, p)
       end

    else                                # sources have different user-given wavelet
       for i = 1 : ns
           wavelet[i] = convert(Vector{params.data_format}, p[i])
       end
    end

    # allocate memory for vector of sources
    srcs = Vector{Source}(undef, ns)
    for i = 1 : ns
        srcs[i] = Source(isz[i], isx[i], dt, params; p=wavelet[i], ot=ot[i])
    end
    return srcs
end

"""
   get the forward modeling wavefield via frequency domain finite difference
"""
function get_wavefield_FDFD(src::Source, params::FdParams;
         flower=0.0, fupper=60.0, tmax=2.0, print_flag=false, consistent_flag=false)

    # length of wavefield
    N = params.nz * params.nx

    # number of time samples
    nt = floor(Int64, tmax / src.dt) + 1
    if nt < src.it_max
       error("tmax isn't big enough")
    end

    # allocate space for 3D wavefield [nz*nx, nt]
    pressure = zeros(Complex{Float64}, N, nt)

    # frequency sampling interval
    df = 1.0 / (src.dt * nt)

    # lower frequency index bound
    iw_lower = floor(Int64, flower / df) + 1
    if iw_lower < 2
       iw_lower = 2
    end

    # upper frequency index bound
    iw_upper = ceil(Int64, fupper / df) + 1
    if iw_upper > floor(Int64, nt / 2) + 1
       iw_upper = floor(Int64, nt / 2) + 1
    end

    # allocate space for source time function
    wlet = zeros(nt)
    copyto!(wlet, src.it_min, src.p)

    # transform to frequency domain
    fwave = fft(wlet)

    # allocate space for left and right hand side
    b = zeros(Complex{Float64}, params.Nz * params.Nx)
    u = zeros(Complex{Float64}, params.Nz * params.Nx)

    # loop over all frequency slice
    for iw = iw_lower : iw_upper

        # radian frequency
        omega = 2.0 * pi * (iw-1) * df

        # build damping profile
        H = get_helmholtz_LU(params, omega)

        # two ways to inject source
        if consistent_flag
           bulk = params.rho[src.isz, src.isx] * params.vel[src.isz, src.isx] * params.vel[src.isz, src.isx]
           b[src.src2spt] = im * omega * fwave[iw] / bulk
        else
           b[src.src2spt] = fwave[iw]
        end

        # solve implicit wave equation
        ldiv!(u, H, b)

        # save the pressure component
        for i = 1 : N
            tmp = u[params.spt2wfd[i]]
            pressure[i, iw] = tmp
            pressure[i, nt-iw+2] = conj(tmp)
        end

        # print the frequency slice
        if print_flag
           println("finished $iw")
        end
    end

    # transform back to time domain
    pressure = real(ifft(pressure, 2))

    # return a 3D cube
    return reshape(pressure, params.nz, params.nx, nt)
end

"""
   constructor for recordings for frequency domain finite difference
"""
function Recordings(rz::Vector, rx::Vector, dt, tmax,
         params::FdParams; location_flag="index")

    # number of receivers
    nr  = length(rz)
    if length(rx) != nr
       error("length of receiver coordinate vector doesn't match")
    end

    irz = zeros(Int64, nr)
    irx = zeros(Int64, nr)

    # location are provided as index
    if location_flag == "index"
       for i = 1 : nr
           irz[i] = round(Int64, rz[i])
           irx[i] = round(Int64, rx[i])
       end

    # location are given as distance
    elseif location_flag == "distance"
       for i = 1 : nr
           irz[i] = round(Int64, rz[i]/params.h) + 1
           irx[i] = round(Int64, rx[i]/params.h) + 1
       end
    else
       error("wrong specification of receiver location")
    end

    # error checking
    for i = 1 : nr
        if irz[i] > params.nz || irx[i] > params.nx || irz[i] < 1 || irx[i] < 1
           error("receiver located outside of modeling area")
        end
    end

    # can't put receivers on the free surface
    if params.free_surface
       for i = 1 : nr
           if irz[i] == 1
              error("can't put receiver at free surface, move it deeper")
           end
       end
    end

    # convert dt format
    dt = convert(params.data_format, dt)
    nt = floor(Int64, tmax/dt) + 1

    # the auxillary vector mapping snapshot to recordings
    spt2rec = zeros(Int64, nr)
    for i = 1 : nr
        spt2rec[i] = (irx[i]+params.npml-1) * params.Nz + irz[i] + params.ntop
    end

    # initalize empty recordings
    rec= Recordings(nt, nr, dt, irz, irx, spt2rec,
         zeros(params.data_format, nt, nr))

    return rec
end

"""
   get the recordings via frequency domain finite difference
"""
function get_recordings!(rec::Recordings, src::Source, params::FdParams;
         flower=0.0, fupper=60.0, print_flag=false, consistent_flag=false) where {Ti <: Int64}

    # number of time samples
    nt = rec.nt
    if nt < src.it_max
       error("simulation length is less than the length of source wavelet")
    end

    # check the time sample interval
    if rec.dt != src.dt
       error("time sample interval doesn't match")
    end

    # frequency interval
    df   = 1.0 / (src.dt * nt)

    # lower frequency index bound
    iw_lower = floor(Int64, flower / df) + 1
    if iw_lower < 2
       iw_lower = 2
    end

    # upper frequency index bound
    iw_upper = ceil(Int64, fupper / df) + 1
    if iw_upper > floor(Int64, nt / 2) + 1
       iw_upper = floor(Int64, nt / 2) + 1
    end

    # allocate space for source time function
    wlet = zeros(nt)
    copyto!(wlet, src.it_min, src.p)

    # transform to frequency domain
    fwave = fft(wlet)

    # allocate space for left and right hand side
    b = zeros(Complex{Float64}, params.Nz * params.Nx)
    u = zeros(Complex{Float64}, params.Nz * params.Nx)

    # the pressure component
    p = zeros(Complex{Float64}, nt, rec.nr)

    # loop over all frequency slice
    for iw = iw_lower : iw_upper

        # radian frequency
        omega = 2.0 * pi * (iw-1) * df

        # build damping profile
        H = get_helmholtz_LU(params, omega)

        # two ways to inject source
        if consistent_flag
           bulk = params.rho[src.isz, src.isx] * params.vel[src.isz, src.isx] * params.vel[src.isz, src.isx]
           b[src.src2spt] = im * omega * fwave[iw] / bulk
        else
           b[src.src2spt] = fwave[iw]
        end

        # solve implicit equation
        ldiv!(u, H, b)

        # save the pressure component
        for i = 1 : rec.nr
            idx = rec.spt2rec[i]
            p[iw,i]     = u[idx]
            p[nt-iw+2,i]= conj(u[idx])
        end

        # print the frequency slice
        if print_flag
           println("finished $iw")
        end
    end

    # transform back to time domain
    ifft!(p, 1)
    p = real(p)

    # copy to recordings
    copyto!(rec.p, 1, p)

    return nothing
end

"""
   get the recordings via frequency domain finite difference for the simultaneous source case
"""
function get_recordings!(rec::Recordings, srcs::Vector{Source}, params::FdParams;
         flower=0.0, fupper=60.0, print_flag=false, consistent_flag=false) where {Ti <: Int64}

    # number of sources
    ns = length(srcs)

    # number of time samples
    nt = rec.nt
    for i = 1 : ns
        if nt < srcs[i].it_max
           error("simulation length is less than the length of source wavelet")
        end

        # check the time sample interval
        if rec.dt != srcs[i].dt
           error("time sample interval doesn't match")
        end
    end

    # frequency interval
    df   = 1.0 / (srcs[1].dt * nt)

    # lower frequency index bound
    iw_lower = floor(Int64, flower / df) + 1
    if iw_lower < 2
       iw_lower = 2
    end

    # upper frequency index bound
    iw_upper = ceil(Int64, fupper / df) + 1
    if iw_upper > floor(Int64, nt / 2) + 1
       iw_upper = floor(Int64, nt / 2) + 1
    end

    # convert the source wavelet to frequency domain
    fwave = Vector{Vector{Complex{Float64}}}(undef, ns)
    wlet  = zeros(nt)
    for i = 1 : ns
        copyto!(wlet, srcs[i].it_min, srcs[i].p)
        fwave[i] = fft(wlet)
        fill!(wlet, 0.0)
    end

    # allocate space for left and right hand side
    b = zeros(Complex{Float64}, params.Nz * params.Nx)
    u = zeros(Complex{Float64}, params.Nz * params.Nx)

    # the pressure component
    p = zeros(Complex{Float64}, nt, rec.nr)

    # loop over all frequency slice
    for iw = iw_lower : iw_upper

        # radian frequency
        omega = 2.0 * pi * (iw-1) * df

        # build damping profile
        H = get_helmholtz_LU(params, omega)

        for i = 1 : ns
            # two ways to inject source
            if consistent_flag
               bulk = params.rho[srcs[i].isz, srcs[i].isx] * params.vel[srcs[i].isz, srcs[i].isx] * params.vel[srcs[i].isz, srcs[i].isx]
               b[srcs[i].src2spt] += im * omega * fwave[i][iw] / bulk
            else
               b[srcs[i].src2spt] += fwave[i][iw]
            end
        end

        # solve implicit equation
        ldiv!(u, H, b)

        # clean b to prepare for next frequency slice
        for i = 1 : ns
            b[srcs[i].src2spt] = zero(Complex{Float64})
        end

        # save the pressure component
        for i = 1 : rec.nr
            idx = rec.spt2rec[i]
            p[iw,i]     = u[idx]
            p[nt-iw+2,i]= conj(u[idx])
        end

        # print the frequency slice
        if print_flag
           println("finished $iw")
        end
    end

    # transform back to time domain
    ifft!(p, 1)
    p = real(p)

    # copy to recordings
    copyto!(rec.p, 1, p)

    return nothing
end

# vel = model_padding(params.vel, params.npml, params.free_surface)
# rho = model_padding(params.rho, params.npml, params.free_surface)
# k        = vel2bulk(rho, vel)
# buoyancy = rho2buoyancy(rho)

# """
#    convert sparse matrix from coordinate format to CSC format
# """
# function COOCSC(m::Ti, n::Ti,
#          irow::Vector{Ti}, icol::Vector{Ti}, nzval::Vector{Tv}) where {Ti<:Int64, Tv<:Number}
#
#   etype = eltype(nzval)
#   nnz   = length(nzval)
#
#   colptr = zeros(Int64, n+1)
#   rowidx = zeros(Int64, nnz)
#   nzval1 = zeros(etype, nnz)
#
#   # Determine the column length
#   for i = 1 : nnz
#       colptr[icol[i]] = colptr[icol[i]] + 1
#   end
#
#   # starting position of the element in each column
#   k = 1
#   for j = 1 : n + 1
#       k0 = colptr[j]
#       colptr[j] = k
#       k = k + k0
#   end
#
#   # fill in the output
#   for k = 1 : nnz
#       i = irow[k]
#       j = icol[k]
#       x = nzval[k]
#       idx = colptr[j]
#       nzval1[idx] = x
#       rowidx[idx] = i
#       colptr[j  ] = idx + 1
#   end
#
#   # shift back colptr
#   for j = n : -1 : 1
#       colptr[j+1] = colptr[j]
#   end
#   colptr[1] = 1
#
#   A = spzeros(m, n)
#   A.rowval = rowidx
#   A.colptr = colptr
#   A.nzval  = nzval1
#
#   return A
#
# end
#
# """
#    For sparse matrix in coordinate form, remove the repeated ones
# """
# function remove_repeated_samples(m::Ti, n::Ti,
#          irow::Ti, icol::Ti, nzval::Tv) where {Ti<:Int64, Tv<:Number}
#
#   etype = eltype(nzval)
#   nnz   = length(nzval)
#
#   colptr = zeros(Int64, n+1)
#   rowidx = zeros(Int64, nnz)
#   nzval1 = zeros(etype, nnz)
#
#   # Determine the column length
#   for i = 1 : nnz
#       colptr[icol[i]] = colptr[icol[i]] + 1
#   end
#
#   # starting position of each column
#   k = 1
#   for j = 1 : N + 1
#       k0 = colptr[j]
#       colptr[j] = k
#       k = k + k0
#   end
#
#   # fill in the output
#   for k = 1 : nnz
#       i = irow[k]
#       j = icol[k]
#       x = nzval[k]
#       idx = colptr[j]
#       nzval1[idx] = x
#       rowidx[idx] = i
#       colptr[j  ] = idx + 1
#   end
#
#   # shift back colptr
#   for j = N : -1 : 1
#       colptr[j+1] = colptr[j]
#   end
#   colptr[1] = 1
#
#   A = spzeros(m, n)
#   A.rowval = rowidx
#   A.colptr = colptr
#   A.nzval  = nzval1
#
#   return A
#
# end


text:
module ContinuousWavelets
using Wavelets
using Interpolations
using AbstractFFTs
using FFTW
using LinearAlgebra
using SpecialFunctions

import Wavelets.wavelet
import Wavelets.WT.name, Wavelets.WT.class,
   Wavelets.WT.vanishingmoments

export ContWave, CWT, cwt, icwt

# Boundaries
export WaveletBoundary, PerBoundary, ZPBoundary, NullBoundary, SymBoundary,
    Periodic, DEFAULT_BOUNDARY, padded, NaivePer, SymBound
# waveletTypes (note there are also export statements in the for loops of waveletTypes.jl)
export Morlet, Paul, Dog, ContOrtho, Morse, name, vanishingmoments, ContOrtho
export cHaar, cBeyl, cVaid, morl, morse
# averaging types
export Average, Dirac, Father, NoAve
# inverse types
export DualFrames, NaiveDelta, PenroseDelta, InverseType
# CWT constructors
export wavelet, waveletType
# general utilities
export qmf, computeWavelets, getNWavelets, mother, father, getMeanFreq, isAnalytic, getDualCoverage

@doc """
    ContWave{Boundary,T}
The abstract type encompassing the various types of wavelets implemented in
the package. The abstract type has parameters `Boundary<:WaveletBoundary` and
`T<:Number`, which gives the element output type. Each has both a constructor,
and a default predefined entry. These
are:
- `Morlet`: A complex approximately analytic wavelet that is just a frequency
    domain Gaussian with mean subtracted. `Morlet(σ::T) where T<: Real`. `σ`
    gives the frequency domain variance of the mother Wavelet. As `σ` goes to
    zero, all of the information becomes spatial. Default is `morl` which has
    σ=2π.

    ``\\psi\\hat(\\omega) \\propto \\textrm{e}^{-\\frac{\\mu^2}{2}}\\big(\\textrm{e}^{-(\\mu - \\omega)^2} -\\textrm{e}^{\\frac{\\omega^2-\\mu^2}{2}}\\big)``

- `Paul{N}`: A complex analytic wavelet, also known as Cauchy wavelets. `pauln` for n in `1:20` e.g. `paul5`

    ``\\psi\\hat(\\omega) \\propto \\chi_{\\omega \\geq 0} \\omega^N\\textrm{e}^{-\\omega}``
- `Dog{N}`: Derivative of a Gaussian, where N is the number of
    derivatives. `dogn` for `n` in `0:6`. The Sombrero/mexican hat/Marr wavelet
    is `n=2`.

    ``\\psi\\hat(\\omega) \\propto \\omega^N\\textrm{e}^{-\\frac{\\omega^2}{2}}``
- `ContOrtho{OWT}`. OWT is some orthogonal wavelet of type `OrthoWaveletClass`
    from [Wavelets.jl](https://github.com/JuliaDSP/Wavelets.jl). This uses an
    explicit construction of the mother wavelet for these orthogonal wavelets
    to do a continuous transform. Constructed via `ContOrtho(o::W)` where `o`
    is from Wavelets.jl. Alternatively, you can get them directly as
    `ContOrtho` objects via:
   + `cHaar` Haar Wavelets
   + `cBeyl` Beylkin Wavelets
   + `cVaid` Vaidyanthan Wavelets
   + `cDbn` Daubhechies Wavelets. n ranges from `1:Inf`
   + `cCoifn` Coiflets. n ranges from `2:2:8`
   + `cSymn` Symlets. n ranges from `4:10`
   + `cBattn` Battle-Lemarie wavelets. n ranges from `2:2:6`
"""
abstract type ContWave{Boundary,T} end # equivalent to ContinuousWavelet in Wavelets.jl
# BOUNDARY TYPES

"""
the abstract type for the various types of boundaries
"""
abstract type WaveletBoundary end
# periodic (default)
"""
    PerBoundary() <: WaveletBoundary
standard periodic boundary assumption made by the fft. Aliases of `NaivePer`
and `Periodic`.
"""
struct PerBoundary <: WaveletBoundary end
# zero padding
"""
    ZPBoundary() <: WaveletBoundary
zero pads the signal before doing an fft, rounding up to the nearest power of
two. Alias of `padded`.
"""
struct ZPBoundary <: WaveletBoundary end
# constant padding
# struct CPBoundary <: WaveletBoundary end
# symmetric boundary (as in the DCTII)
"""
    SymBoundary() <: WaveletBoundary
symmetric boundary, as in the DCT type II. Repeats the edge value, which
alleviates derivative discontinuities. Alias of `DEFAULT_BOUNDARY` and
`SymBound`.
"""
struct SymBoundary <: WaveletBoundary end
# and so on...

const Periodic = PerBoundary()
const DEFAULT_BOUNDARY = SymBoundary()
const padded = ZPBoundary()
const NaivePer = PerBoundary()
const SymBound = SymBoundary()

abstract type ContWaveClass end # equivalent to ContinuousWaveletClass in Wavelets.jl
include("waveletTypes.jl")
include("CWTConstruction.jl")
include("utils.jl")
include("sanityChecks.jl")
include("createWavelets.jl")
include("apply.jl")
end

ext: jl


text:
export  PlaneStress
mutable struct PlaneStress
    H::Array{Float64}
    E::Float64
    ν::Float64
    ρ::Float64
    σ0::Array{Float64} # stress at last time step
    σ0_::Array{Float64} # σ0 to be updated in `commitHistory`
    ε0::Array{Float64}
    ε0_::Array{Float64}
end

function PlaneStress(prop::Dict{String, Any})
    E = prop["E"]; ν = prop["nu"]; ρ = prop["rho"]
    H = zeros(3,3)

    H[1,1] = E/(1. -ν*ν)
    H[1,2] = H[1,1]*ν
    H[2,1] = H[1,2]
    H[2,2] = H[1,1]
    H[3,3] = E/(2.0*(1.0+ν))
    # @show H
    # # ! remove it!!!
    # H = [ 4.09627e12  2.88053e11  -1.77231
    #     2.88053e11  8.29483e11   0.616748
    # -1.77231     0.616748     3.43662e11]
    σ0 = zeros(3); σ0_ = zeros(3); ε0 = zeros(3); ε0_ = zeros(3)
    PlaneStress(H, E, ν, ρ, σ0, σ0_,ε0,ε0_)
end

function getStress(self::PlaneStress, strain::Array{Float64}, Dstrain::Array{Float64}, Δt::Float64 = 0.0)
    sigma = self.H * strain
    self.σ0_ = copy(sigma)
    self.ε0_ = copy(strain)
    return sigma, self.H
end

function getTangent(self::PlaneStress)
    self.H
end

function commitHistory(self::PlaneStress)
    self.σ0 = self.σ0_
    self.ε0 = self.ε0_
end


text:
(cnet::CachedNet)(σ::LUState) =
    logψ(cnet, σ)

logψ(cnet::CachedNet, σ::LUState{<:DoubleState}) = begin
    σ_r = row(state(σ))
    σ_c = col(state(σ))

    return logψ(cnet.net, cnet.cache, lut(σ),
                config(raw_state(σ_r)), config(raw_state(σ_c)),
                changes(σ_r), changes(σ_c))
end

logψ_and_∇logψ(cnet::CachedNet, σ::LUState) = begin
   # @warn "Inefficient calling logψ_and_∇logψ for cachedNet"
    ∇lnψ = grad_cache(cnet)
    lnψ, ∇lnψ = logψ_and_∇logψ!(∇lnψ, cnet, σ)
    return lnψ, ∇lnψ
end

logψ_and_∇logψ!(∇lnψ, cnet::CachedNet, σ::LUState) = begin
    σ_r = row(state(σ))
    σ_c = col(state(σ))

    lnψ = logψ_and_∇logψ!(∇lnψ, cnet.net, cnet.cache, lut(σ),
                          config(raw_state(σ_r)), config(raw_state(σ_c)),
                          changes(σ_r), changes(σ_c))
    return lnψ, ∇lnψ
end

logψ_Δ(cnet::CachedNet, σ::LUState{<:DoubleState}) = begin
    σ_r = row(state(σ))
    σ_c = col(state(σ))
    return Δ_logψ(cnet.net, cnet.cache, lut(σ),
                  config(raw_state(σ_r)), config(raw_state(σ_c)),
                  changes(σ_r), changes(σ_c))
end

logψ_Δ(cnet::CachedNet, σ::LUState{<:DoubleState}, changes_r, changes_c) = begin
    σ_r = row(state(σ))
    σ_c = col(state(σ))
    return Δ_logψ(cnet.net, cnet.cache, lut(σ),
                  config(raw_state(σ_r)), config(raw_state(σ_c)),
                  changes_r, changes_c)
end

function Δ_logψ_and_∇logψ(cnet::CachedNet, σ::LUState, args...)
    ∇lnψ = grad_cache(cnet)
    return Δ_logψ_and_∇logψ!(∇lnψ, cnet, σ, args...)
end

Δ_logψ_and_∇logψ!(∇lnψ, cnet::CachedNet, σ::LUState{<:DoubleState}) = begin
    σ_r  = row(state(σ))
    σ_c  = col(state(σ))
    Δlnψ = Δ_logψ_and_∇logψ!(∇lnψ,
                  cnet.net, cnet.cache, lut(σ),
                  config(raw_state(σ_r)), config(raw_state(σ_c)),
                  changes(σ_r), changes(σ_c))
    return Δlnψ, ∇lnψ
end

Δ_logψ_and_∇logψ!(∇lnψ, cnet::CachedNet, σ::LUState{<:DoubleState}, changes_r, changes_c) = begin
    σ_r  = row(state(σ))
    σ_c  = col(state(σ))
    Δlnψ = Δ_logψ_and_∇logψ!(∇lnψ,
                  cnet.net, cnet.cache, lut(σ),
                  config(raw_state(σ_r)), config(raw_state(σ_c)),
                  changes_r, changes_c)
    return Δlnψ, ∇lnψ
end

prepare_lut!(σ::LUState, cnet::CachedNet) =
    set_lookup!(lut(σ), cnet.net, cnet.cache, raw_config(state(σ))...)
    #set_lookup!(lut(σ), cnet.net, cnet.cache, config(state(σ))...)

apply_lut_updates!(σ::LUState{<:DoubleState}, cnet::CachedNet) = begin
    σ_r = row(state(σ))
    σ_c = col(state(σ))
    update_lookup!(lut(σ), cnet.net, cnet.cache,
        config(raw_state(σ_r)), config(raw_state(σ_c)),
        changes(σ_r), changes(σ_c))
end

update_lookup!(σ::LUState, cnet::CachedNet) = begin
    # TODO Check if this is correct. I could just rename this, no?
    update_lut!(σ, cnet)
end


text:
## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using Measures
using HDF5

##
function fink_filter_bank_fast(J, L;wd=1,pc=1)
    #plane coverage (default 1, full 2Pi 2)
    #width of the wavelets (default 1, wide 2)

    # -------- set parameters
    dθ   = pc*π/L
    wdθ  = wd*dθ
    nx   = 256
    dx   = nx/2-1

    # -------- allocate output array of zeros
    filt = zeros(256, 256, J, L)

    # -------- allocate theta and logr arrays
    logr = zeros(nx, nx)
    θ    = zeros(nx, nx)

    for l = 0:L-1
        θ_l = dθ*l

    # -------- allocate anggood BitArray
        anggood = falses(nx, nx)

    # -------- loop over pixels
        for x = 1:nx
            sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
            for y = 1:nx
                sy = mod(y+dx,nx)-dx -1
                θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                θ_good = abs(θ_pix-π) <= wdθ

            # If this is a pixel we might use, calculate log2(r)
                if θ_good
                    anggood[y, x] = θ_good
                    θ[y, x]       = θ_pix
                    r = sqrt(sx^2 + sy^2)
                    logr[y, x] = log2(max(1,r))
                end
            end
        end
        angmask = findall(anggood)
    # -------- compute the wavelet in the Fourier domain
    # -------- the angular factor is the same for all j
        F_angular = cos.((θ[angmask].-π).*L./(2 .*wd .*pc))

    # -------- loop over j for the radial part
        for j = 0:J-1
            jrad  = 7-j
            Δj    = abs.(logr[angmask].-jrad)
            rmask = (Δj .<= 1)

    # -------- radial part
            F_radial = cos.(Δj[rmask] .* (π/2))
            ind      = angmask[rmask]
            filt[ind,j+1,l+1] = F_radial .* F_angular[rmask]
        end
    end
    return filt
end

@time fink_filter_set = fink_filter_bank_fast(8,16,wd=1)

findall(test[:,:,6] .== maximum(test[:,:,6]))

theme(:juno)

plot(Plots.heatmap(fftshift(fink_filter_set[:,:,2,1]),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

plot(Plots.heatmap(fftshift(real.(ifft(fink_filter_set[:,:,7,2]))),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:twilight,
    legend = :none,
    margin=0mm))

ifft(fink_filter_set[:,:,8,1])

nx = 256
c=2
im_scale = convert(Int8,log2(nx))
J = (im_scale-2)*c

println("start")
for i = 1/c:1/c:im_scale-2
    println(i)
end

## dev

function fink_filter_bank_fast(c, L; nx=256,wd=1,pc=1)
    #plane coverage (default 1, full 2Pi 2)
    #width of the wavelets (default 1, wide 2)
    #c sets the scale sampling rate (1 is dyadic, 2 is half dyadic)

    # -------- set parameters
    dθ   = pc*π/L
    wdθ  = wd*dθ
    dx   = nx/2-1

    im_scale = convert(Int8,log2(nx))
    J = (im_scale-2)*c

    # -------- allocate output array of zeros
    filt = zeros(nx, nx, J*L+1)

    # -------- allocate theta and logr arrays
    logr = zeros(nx, nx)
    θ    = zeros(nx, nx)

    # -------- allocate temp phi building zeros
    phi_b = zeros(nx, nx)

    for l = 0:L-1
        θ_l = dθ*l

    # -------- allocate anggood BitArray
        anggood = falses(nx, nx)

    # -------- loop over pixels
        for x = 1:nx
            sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
            for y = 1:nx
                sy = mod(y+dx,nx)-dx -1
                θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                θ_good = abs(θ_pix-π) <= wdθ

            # If this is a pixel we might use, calculate log2(r)
                if θ_good
                    anggood[y, x] = θ_good
                    θ[y, x]       = θ_pix
                    r = sqrt(sx^2 + sy^2)
                    logr[y, x] = log2(max(1,r))
                end
            end
        end
        angmask = findall(anggood)
    # -------- compute the wavelet in the Fourier domain
    # -------- the angular factor is the same for all j
        F_angular = cos.((θ[angmask].-π).*L./(2 .*wd .*pc))

    # -------- loop over j for the radial part
        for (j_ind, j) in enumerate(1/c:1/c:im_scale-2)
            jrad  = 7-j
            Δj    = abs.(logr[angmask].-jrad)
            rmask = (Δj .<= 1/c)

    # -------- radial part
            F_radial = cos.(Δj[rmask] .* (c*π/2))
            ind      = angmask[rmask]
            filt[ind,(j_ind-1)*L+l+1] = F_radial .* F_angular[rmask]
        end
    # -------- handle the phi case (jrad=0, j=7)
        for (j_ind, j) in enumerate(im_scale-2+1/c:1/c:im_scale-1)
            jrad  = 7-j
            Δj    = abs.(logr[angmask].-jrad)
            rmask = (Δj .<= 1/c)

            # -------- radial part
            F_radial = cos.(Δj[rmask] .* (c*π/2))
            ind      = angmask[rmask]
            phi_b[ind] .+= (F_radial .* F_angular[rmask]).^2
        end
    end

    # -------- normalize the phi case correctly
    if pc == 1
        phi_b .+= circshift(phi_b[end:-1:1,end:-1:1],(1,1))
        phi_b[1,1] /= 2
    end

    phi_b .= sqrt.(phi_b)

    filt[:,:,J*L+1] .= phi_b
    return filt
end

test = fink_filter_bank_fast(1,8,pc=1)
h5write("./Data/filter_bank_1_8_1.h5", "main/data", test)
test = fink_filter_bank_fast(1,16,pc=2)
h5write("./Data/filter_bank_1_16_2.h5", "main/data", test)
test = fink_filter_bank_fast(2,8,pc=1)
h5write("./Data/filter_bank_2_8_1.h5", "main/data", test)
test = fink_filter_bank_fast(2,16,pc=2)
h5write("./Data/filter_bank_2_16_2.h5", "main/data", test)

push!(LOAD_PATH, pwd())
using DHC_Utils

# Generate phi from an L=16 filter bank
donut = fink_filter_bank(8,8)
phi   = sqrt.(reshape(sum(donut[:,:,8,:].^2,dims=[3,4]),256,256))
phi_r = real(ifft(phi))
# check that the real-space function sums to 1 as expected.
print(sum(phi_r))

plot(Plots.heatmap(fftshift(phi_r),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

plot(Plots.heatmap(fftshift(phi),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

phi_shift = fftshift(phi)

plot(Plots.heatmap(phi_shift.+phi_shift[end:-1:1,end:-1:1],
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

plot(Plots.heatmap(fftshift(donut[:,:,8,2].^2),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

X = fftshift(donut[:,:,8,2].^2)

plot(Plots.heatmap(X[end:-1:1,end:-1:1],
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

# Generate phi from an L=16 filter bank
donut = fink_filter_bank(8,8)
phi   = sqrt.(reshape(sum(donut[:,:,8,:].^2,dims=[3,4]),256,256))
phi_r = real(ifft(phi))
# check that the real-space function sums to 1 as expected.
print(sum(phi_r))

plot(Plots.heatmap(fftshift(donut[:,:,8,8].^2),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

# Generate phi from an L=16 filter bank
donut = fink_filter_bank(8,8)
phi   = sqrt.(reshape(sum(donut[:,:,8,:].^2,dims=[3,4]),256,256))
phi_r = real(ifft(phi))
# check that the real-space function sums to 1 as expected.
print(sum(phi_r))

plot(Plots.heatmap(fftshift(phi),
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

phi_shift2 = fftshift(phi)
phi_shift3 = phi_shift2.+circshift(phi_shift2[end:-1:1,end:-1:1],(1,1))
phi_shift = fftshift(phi)

plot(Plots.heatmap(phi_shift.+phi_shift[end:-1:1,end:-1:1],
    xlims=(1,256),
    ylims=(1,256),
    aspectratio=1,
    axis=nothing,
    border=:none,
    c=:cubehelix,
    legend = :none,
    margin=0mm))

donut = fink_filter_bank(8,16)
phi   = sqrt.(reshape(sum(donut[:,:,8,:].^2,dims=[3,4]),256,256))
phi_r = real(ifft(phi))
# check that the real-space function sums to 1 as expected.
print(sum(phi_r))

## check run straight through...
push!(LOAD_PATH, pwd())
using DHC_2DUtils

fink_filter_set = fink_filter_list(fink_filter_bank(2,8,wd=1,pc=1))
DHC_compute(rand(256,256),fink_filter_set)

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
# Cross validation

## cross validation generators

abstract CrossValGenerator

# K-fold

immutable Kfold <: CrossValGenerator
    permseq::Vector{Int}
    k::Int
    coeff::Float64

    function Kfold(n::Int, k::Int)
        2 <= k <= n || error("The value of k must be in [2, length(a)].")
        new(randperm(n), k, n / k)
    end
end

length(c::Kfold) = c.k

immutable KfoldState
    i::Int      # the i-th of the subset
    s::Int      # starting index
    e::Int      # ending index
end

start(c::Kfold) = KfoldState(1, 1, @compat(round(Integer,c.coeff)))
next(c::Kfold, s::KfoldState) =
    (i = s.i+1; (setdiff(1:length(c.permseq), c.permseq[s.s:s.e]), KfoldState(i, s.e+1, @compat(round(Integer,c.coeff * i)))))
done(c::Kfold, s::KfoldState) = (s.i > c.k)

# Stratified K-fold

immutable StratifiedKfold <: CrossValGenerator
    n::Int                         #Total number of observations
    permseqs::Vector{Vector{Int}}  #Vectors of vectors of indexes for each stratum
    k::Int                         #Number of splits
    coeffs::Vector{Float64}        #About how many observations per strata are in a val set
    function StratifiedKfold(strata, k)
        2 <= k <= length(strata) || error("The value of k must be in [2, length(strata)].")
        strata_labels, permseqs = unique_inverse(strata)
        map(shuffle!, permseqs)
        coeffs = Float64[]
        for (stratum, permseq) in zip(strata_labels, permseqs)
            k <= length(permseq) || error("k is greater than the length of stratum $stratum")
            push!(coeffs, length(permseq) / k)
        end
        new(length(strata), permseqs, k, coeffs)
    end
end

length(c::StratifiedKfold) = c.k

start(c::StratifiedKfold) = 1
function next(c::StratifiedKfold, s::Int)
    r = Int[]
    for (permseq, coeff) in zip(c.permseqs, c.coeffs)
        a, b = @compat(round(Integer, [s-1, s] .* coeff))
        append!(r, view(permseq, a+1:b))
    end
    setdiff(1:c.n, r), s+1
end
done(c::StratifiedKfold, s::Int) = (s > c.k)

# LOOCV (Leave-one-out cross-validation)

function leave_one_out(n::Int, i::Int)
    @assert 1 <= i <= n
    x = Array(Int, n-1)
    for j = 1:i-1
        x[j] = j
    end
    for j = i+1:n
        x[j-1] = j
    end
    return x
end

immutable LOOCV <: CrossValGenerator
    n::Int
end

length(c::LOOCV) = c.n

start(c::LOOCV) = 1
next(c::LOOCV, s::Int) = (leave_one_out(c.n, s), s+1)
done(c::LOOCV, s::Int) = (s > c.n)

# Repeated random sub-sampling

immutable RandomSub <: CrossValGenerator
    n::Int    # total length
    sn::Int   # length of each subset
    k::Int    # number of subsets
end

length(c::RandomSub) = c.k

start(c::RandomSub) = 1
next(c::RandomSub, s::Int) = (sort!(sample(1:c.n, c.sn; replace=false)), s+1)
done(c::RandomSub, s::Int) = (s > c.k)

# Stratified repeated random sub-sampling

immutable StratifiedRandomSub <: CrossValGenerator
    idxs::Vector{Vector{Int}}    # total length
    sn::Int                      # length of subset
    sns::Vector{Int}             # num from each stratum for each subset
    k::Int                       # number of subsets
    function StratifiedRandomSub(strata, sn, k)
        n = length(strata)
        strata_labels, idxs = unique_inverse(strata)
        sn >= length(strata_labels) || error("sn is too small for all strata to be represented")
        lengths_ord = sortperm(map(length, idxs))
        sns = zeros(Int, n)
        remaining_n = n   # total in the strata we haven't seen yet
        remaining_sn = sn # total room in the subset that hasn't been "assinged" to a stratum yet
        #loop through strata from smallest to largest, making sure there is at least one idx
        #from each strata in each subset:
        for stratum_num in lengths_ord
            stratum_n = length(idxs[stratum_num])
            remaining_proportion = remaining_sn/remaining_n
            stratum_sn = max(@compat(round(Integer, remaining_proportion*stratum_n)), 1)
            remaining_n -= stratum_n
            remaining_sn -= stratum_sn
            sns[stratum_num] = stratum_sn
        end
        #@assert mapreduce(sum, +, sns) == sn
        new(idxs, sn, sns, k)
    end
end

length(c::StratifiedRandomSub) = c.k

start(c::StratifiedRandomSub) = 1
function next(c::StratifiedRandomSub, s::Int)
    idxs = Array(Int, 0)
    sizehint!(idxs, c.sn)
    for (stratum_sn, stratum_idxs) in zip(c.sns, c.idxs)
        append!(idxs, sample(stratum_idxs, stratum_sn, replace=false))
    end
    (sort!(idxs), s+1)
end
done(c::StratifiedRandomSub, s::Int) = (s > c.k)

## Cross validation algorithm
#
#  estfun: model estimation function
#
#          model = estfun(train_inds)
#
#          it takes as input the indices of
#          the samples for training, and returns
#          a trained model.
#
#  evalfun: model evaluation function
#
#           v = evalfun(model, test_inds)
#
#           it applies a trained model to a subset
#           of samples whose indices are given in
#           test_inds, and returns an overall score.
#
#  n:   the total number of samples
#
#  gen: an iterable object (e.g. an instance of CrossValGenerator)
#       where each element is a vector of indices
#
#  This function returns a tuple a vector of validation scores
#
function cross_validate(estfun::Function, evalfun::Function, n::Int, gen)
    best_model = nothing
    best_score = NaN
    best_inds = Int[]
    first = true

    scores = Float64[]
    for (i, train_inds) in enumerate(gen)
        test_inds = setdiff(1:n, train_inds)
        model = estfun(train_inds)
        score = evalfun(model, test_inds)
        push!(scores, score)
    end
    return scores
end

cross_validate(estfun::Function, evalfun::Function, n::Integer, gen) =
    cross_validate(estfun, evalfun, n, gen, Forward)


text:
### A Pluto.jl notebook ###
# v0.11.4

using Markdown
using InteractiveUtils

# ╔═╡ 57abe58c-e897-11ea-2806-33ea69204c99
md"""# Experimentation with multilevel intelligent tables
"""

# ╔═╡ 272d908e-e899-11ea-3fa0-035efee04065
md""" What if a cell is implemented as an object that is an instance of an heirarchical class structure with an inherited notion of value.

When an object's value is accessed for rendering, per the accessing client, the response reflects the class heirarchy:
1. CONSTANT: Starting from an entered value,
2. CALCULATION: A calculated value,
3. DISTRIBUTION: A value sampled from a prior distribution

In practical terms, if an entered value is unavailable, the calculation is accessed, failing which a sample from the prior is used, and so on."""

# ╔═╡ b937fb56-e897-11ea-3855-8d3facef5038
md"""## Cell #1
"""

# ╔═╡ d541fd92-e897-11ea-258b-b3c8b3f09d62
md"""## Cell #2
"""

# ╔═╡ dd09d7d4-e897-11ea-1088-217bcf66639f
md"""## Cell #3
"""

# ╔═╡ 6a9552ac-e897-11ea-21b5-4b27d1faacbe
using PlutoUI, Random

# ╔═╡ 96efdea6-e89e-11ea-1d03-ddb47316b341
function access_cell(c)
	return isnothing(c[1]) ?
				(isnothing(c[2]) ?
					c[3] :
				c[2]) :
			c[1]
end

# ╔═╡ 38f10e34-e898-11ea-0b2d-772ab12e1d67
c1 = [0.5,nothing,round(rand(),digits=1)]

# ╔═╡ 28180f0c-e89f-11ea-12c2-df5eaaa1046c
accessed_c1 = access_cell(c1)

# ╔═╡ dbd67dca-e897-11ea-105a-ddd2dcb4b8fe
c2 = [0.2,nothing,round(rand(),digits=1)]

# ╔═╡ cf0b45e0-e898-11ea-082b-8b51a5f60299
accessed_c2 = access_cell(c2)

# ╔═╡ e25c23a2-e897-11ea-366f-013c58886d25
c3 = [0.6,accessed_c1+accessed_c2,2*round(rand(),digits=1)]

# ╔═╡ e1fc59b0-e897-11ea-1b33-7fdd98bfcd47
accessed_c3 = access_cell(c3)

# ╔═╡ Cell order:
# ╟─6a9552ac-e897-11ea-21b5-4b27d1faacbe
# ╟─57abe58c-e897-11ea-2806-33ea69204c99
# ╟─272d908e-e899-11ea-3fa0-035efee04065
# ╠═96efdea6-e89e-11ea-1d03-ddb47316b341
# ╟─b937fb56-e897-11ea-3855-8d3facef5038
# ╠═38f10e34-e898-11ea-0b2d-772ab12e1d67
# ╟─28180f0c-e89f-11ea-12c2-df5eaaa1046c
# ╟─d541fd92-e897-11ea-258b-b3c8b3f09d62
# ╠═dbd67dca-e897-11ea-105a-ddd2dcb4b8fe
# ╟─cf0b45e0-e898-11ea-082b-8b51a5f60299
# ╟─dd09d7d4-e897-11ea-1088-217bcf66639f
# ╠═e25c23a2-e897-11ea-366f-013c58886d25
# ╟─e1fc59b0-e897-11ea-1b33-7fdd98bfcd47

max_stars_repo_name: kripa-experiments/causaltab

max_issues_repo_name: kripa-experiments/causaltab

max_forks_repo_name: kripa-experiments/causaltab


text:
function mo_from_ao!(I::IntegralHelper{T,E,O}, aoints::IntegralHelper{T,E,AtomicOrbitals}, entries...) where {T<:AbstractFloat,E<:AbstractERI,
                                                                                O<:AbstractRestrictedOrbitals}
    t = @elapsed begin
        output("Converting integrals:")
        output(replace("($T, $E, AtomicOrbitals) ⇒ ($T, $E, $O)", "Fermi.Orbitals."=>""))
        for entry in entries
            output("• {:6s}", entry, ending=" ")
            compute!(I, entry, aoints)
            output("✔")
        end
    end
    output("Done in {:5.5f} seconds.\n", t)
end

function mo_from_ao!(I::IntegralHelper{T1,E1,O}, aoints::IntegralHelper{T2,E2,AtomicOrbitals}, entries...) where {T1<:AbstractFloat,E1<:AbstractERI,T2<:AbstractFloat,
                                                                                E2<:AbstractERI,O<:AbstractRestrictedOrbitals}
    t = @elapsed begin
        output("Converting integrals:")
        output(replace("($T2, $E2, AtomicOrbitals) ⇒ ($T1, $E1, $O)", "Fermi.Orbitals."=>""))
        basis = I.orbitals.basis
        aoorbs = AtomicOrbitals(I.molecule, basis)
        newaoints = IntegralHelper{T1}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
        output("IntegralHelper have different ERI handlers. A new AO object will be used:", ending=" ")
        output("{}", typeof(newaoints))

        for entry in entries
            if occursin(r"F[dijab]{0,2}", entry)
                output("Fock matrix will be computed using original ERI")
                output("• {:6s}", entry, ending=" ")
                compute_F!(I, aoints)
                output("✔")
            end
        end
        entries = [entries...]
        filter!(i->!occursin(r"F[dijab]{0,2}",i), entries)

        for entry in entries
            output("• {:6s}", entry, ending=" ")
            compute!(I, entry, newaoints)
            output("✔")
        end
    end
    output("Done in {:5.5f} seconds.\n", t)
end

# function mo_from_ao!(I::IntegralHelper{T1,E1,O}, aoints::IntegralHelper{T2,E2,AtomicOrbitals}, entries...) where {T1<:AbstractFloat,T2<:AbstractFloat,
#                                                             E1<:AbstractERI,E2<:AbstractERI,O<:AbstractRestrictedOrbitals}
#     if T1 !== T2 || ((E2 <: AbstractDFERI) ⊻ (E1 <: AbstractDFERI))
#         output("!! AO Integrals are not the same type as the MO. New integrals will be computed.")
#         for entry in entries
#             if occursin(r"F[dijab]{0,2}", entry)
#                 output("Fock matrix will be computed using old ERI")
#                 compute_F!(I, aoints)
#             end
#             entries = [entries...]
#             filter!(i->i==entry, entries)
#         end
#         basis = I.orbitals.basis
#         aoorbs = AtomicOrbitals(I.molecule, basis)
#         aoints = IntegralHelper{T1}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
#     end
#     t = @elapsed begin
#         output("Computing MO Integrals...")
#         for entry in entries
#             compute!(I, entry, aoints)
#         end
#     end
#     output("Done in {:5.5f} seconds.", t)
# end

# function mo_from_ao!(I::IntegralHelper{T1,E1,O}, aoints::IntegralHelper{T2,SparseERI,AtomicOrbitals}, entries...) where {T1<:AbstractFloat,T2<:AbstractFloat,
#                                                             E1<:AbstractERI,O<:AbstractRestrictedOrbitals}
#     if T1 !== T2 || E1 <: AbstractDFERI
#         output("!! AO Integrals are not the same type as the MO. New integrals will be computed.")
#         for entry in entries
#             if occursin(r"F[dijab]{0,2}", entry)
#                 output("Fock matrix will be computed using old ERI")
#                 compute_F!(I, aoints)
#             end
#             entries = [entries...]
#             filter!(i->i==entry, entries)
#         end
#         basis = I.orbitals.basis
#         aoorbs = AtomicOrbitals(I.molecule, basis)
#         aoints = IntegralHelper{T1}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
#     end
#     t = @elapsed begin
#         output("Converting integrals:")
#         output(replace("($T2, SparseERI, AtomicOrbitals) ⇒ ($T1, $E1, $O)", "Fermi.Orbitals."=>""))
#         for entry in entries
#             compute!(I, entry, aoints)
#         end
#     end
#     output("Done in {:5.5f} seconds.\n", t)
# end

function compute!(I::IntegralHelper{T,Chonky,O}, entry::String, x...) where {T<: AbstractFloat,O<:AbstractRestrictedOrbitals}
    if entry == "S"
        compute_S!(I, x...)
    elseif entry == "T"
        compute_T!(I, x...)
    elseif entry == "V"
        compute_V!(I, x...)
    elseif entry == "ERI"
        compute_ERI!(I, x...)
    elseif entry == "OOOO"
        compute_OOOO!(I, x...)
    elseif entry == "OOOV"
        compute_OOOV!(I, x...)
    elseif entry == "OVOV"
        compute_OVOV!(I, x...)
    elseif entry == "OOVV"
        compute_OOVV!(I, x...)
    elseif entry == "OVVV"
        compute_OVVV!(I, x...)
    elseif entry == "VVVV"
        compute_VVVV!(I, x...)
    elseif occursin(r"F[dijab]{0,2}", entry)
        compute_F!(I, x...)
    else
        throw(FermiException("Invalid key for IntegralHelper: $(entry)."))
    end
end

function compute!(I::IntegralHelper{T,E,O}, entry::String, x...) where {T<: AbstractFloat, E<:AbstractDFERI, O<:AbstractRestrictedOrbitals}
    if entry == "S"
        compute_S!(I, x...)
    elseif entry == "T"
        compute_T!(I, x...)
    elseif entry == "V"
        compute_V!(I, x...)
    elseif entry == "ERI"
        compute_ERI!(I, x...)
    elseif entry == "BOO"
        compute_BOO!(I, x...)
    elseif entry == "BOV"
        compute_BOV!(I, x...)
    elseif entry == "BVV"
        compute_BVV!(I, x...)
    elseif entry == "OOOO"
        compute_OOOO!(I, x...)
    elseif entry == "OOOV"
        compute_OOOV!(I, x...)
    elseif entry == "OVOV"
        compute_OVOV!(I, x...)
    elseif entry == "OOVV"
        compute_OOVV!(I, x...)
    elseif entry == "OVVV"
        compute_OVVV!(I, x...)
    elseif entry == "VVVV"
        compute_VVVV!(I, x...)
    elseif occursin(r"F[dijab]{0,2}", entry)
        compute_F!(I, x...)
    else
        throw(FermiException("Invalid key for IntegralHelper: $(entry)."))
    end
end

function compute_ERI!(I::IntegralHelper{T, E, O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
        # Create AO integral object
        basis = I.orbitals.basis
        aoorbs = AtomicOrbitals(I.molecule, basis)
        aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
        compute_ERI!(I,aoints)
end

include("ROIntegrals/OneElectron.jl")
include("ROIntegrals/DFERI.jl")

function compute_OOOO!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_OOOO!(I,aoints)
end

function compute_OOOV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_OOOV!(I,aoints)
end

function compute_OOVV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_OOVV!(I,aoints)
end

function compute_OVOV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_OVOV!(I,aoints)
end

function compute_OVVV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_OVVV!(I,aoints)
end
function compute_VVVV!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    # Create AO integral object
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper{T}(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_VVVV!(I,aoints)
end

function compute_ERI!(I::IntegralHelper{T, Chonky, O}, aoints::IntegralHelper{T, Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}
    AOERI = aoints["ERI"]
    C = I.orbitals.C
    if eltype(I.orbitals.C) !== T
        C = T.(C)
    end
    @tensoropt MOERI[i,j,k,l] :=  AOERI[μ, ν, ρ, σ]*C[μ, i]*C[ν, j]*C[ρ, k]*C[σ, l]
    I["ERI"] = MOERI
end

function compute_OOOO!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    AOERI = aoints["ERI"]

    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    o = (1+core):ndocc
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Co = T.(Co)
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOOO[i,j,k,l] :=  AOERI[μ, ν, ρ, σ]*Co[μ, i]*Co[ν, j]*Co[ρ, k]*Co[σ, l]
    end
    I["OOOO"] = OOOO
end

function compute_OOOO!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = o
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOOO[i,j,k,l] :=  iνρσ[i, ν, ρ, σ]*Co[ν, j]*Co[ρ, k]*Co[σ, l]
    end
    I["OOOO"] = OOOO
end

function compute_OOOV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    AOERI = aoints["ERI"]

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOOV[i,j,k,a] :=  AOERI[μ, ν, ρ, σ]*Co[μ, i]*Co[ν, j]*Co[ρ, k]*Cv[σ, a]
    end
    I["OOOV"] = OOOV
end

function compute_OOOV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = o
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOOV[i,j,k,a] :=  iνρσ[i, ν, ρ, σ]*Co[ν, j]*Co[ρ, k]*Cv[σ, a]
    end
    I["OOOV"] = OOOV
end

function compute_OOVV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    AOERI = aoints["ERI"]

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOVV[i,j,a,b] :=  AOERI[μ, ν, ρ, σ]*Co[μ, i]*Co[ν, j]*Cv[ρ, a]*Cv[σ, b]
    end
    I["OOVV"] = OOVV
end

function compute_OOVV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = o
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OOVV[i,j,a,b] :=  iνρσ[i, ν, ρ, σ]*Co[ν, j]*Cv[ρ, a]*Cv[σ, b]
    end
    I["OOVV"] = OOVV
end

function compute_OVOV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    AOERI = aoints["ERI"]

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OVOV[i,a,j,b] :=  AOERI[μ, ν, ρ, σ]*Co[μ, i]*Cv[ν, a]*Co[ρ, j]*Cv[σ, b]
    end
    I["OVOV"] = OVOV
end

function compute_OVOV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = o
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
    end
end
    end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OVOV[i,a,j,b] :=  iνρσ[i, ν, ρ, σ]*Cv[ν, a]*Co[ρ, j]*Cv[σ, b]
    end
    I["OVOV"] = OVOV
end

function compute_OVVV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    AOERI = aoints["ERI"]

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OVVV[i,a,b,c] :=  AOERI[μ, ν, ρ, σ]*Co[μ, i]*Cv[ν, a]*Cv[ρ, b]*Cv[σ, c]
    end
    I["OVVV"] = OVVV
end

function compute_OVVV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    core = Options.get("drop_occ")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    Co = I.orbitals.C[:,o]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
        Co = T.(Co)
    end

    # Create a partially-contracted dense array from SparseERI
    iνρσ = zeros(T, length(o), nbf, nbf, nbf)
    @sync for i = o
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Co[μ,i]
            Vν = V*Co[ν,i]
            Vρ = V*Co[ρ,i]
            Vσ = V*Co[σ,i]

            if γab && γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ
                iνρσ[i,ρ,ν,μ] += Vσ

            elseif γab && γμν
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,σ,ν,μ] += Vρ

            elseif γab && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
                iνρσ[i,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,ν,σ,ρ] += Vμ
                iνρσ[i,μ,ρ,σ] += Vν
                iνρσ[i,μ,σ,ρ] += Vν
            elseif γab
                iνρσ[i,ν,ρ,σ] += Vμ
                iνρσ[i,σ,μ,ν] += Vρ
            else
                iνρσ[i,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end

    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        OVVV[i,a,b,c] := iνρσ[i, ν, ρ, σ]*Cv[ν, a]*Cv[ρ, b]*Cv[σ, c]
    end
    I["OVVV"] = OVVV
end

function compute_VVVV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,Chonky, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    AOERI = aoints["ERI"]

    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        VVVV[a,b,c,d] :=  AOERI[μ, ν, ρ, σ]*Cv[μ, a]*Cv[ν, b]*Cv[ρ, c]*Cv[σ, d]
    end
    I["VVVV"] = VVVV
end

function compute_VVVV!(I::IntegralHelper{T,Chonky,O}, aoints::IntegralHelper{T,SparseERI, AtomicOrbitals}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}

    eri_vals = aoints["ERI"].data
    idxs = aoints["ERI"].indexes

    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    v = (ndocc+1):(nbf - inac)
    Cv = I.orbitals.C[:,v]
    if eltype(I.orbitals.C) !== T
        Cv = T.(Cv)
    end

    # Create a partially-contracted dense array from SparseERI
    aνρσ = zeros(T, length(v), nbf, nbf, nbf)
    @sync for a = 1:length(v)
        Threads.@spawn begin
        @inbounds begin
        @fastmath begin
        for z = eachindex(eri_vals)
            V = eri_vals[z]
            μ,ν,ρ,σ = idxs[z] .+ 1
            γμν = μ !== ν
            γρσ = ρ !== σ
            γab = Fermi.index2(μ,ν) !== Fermi.index2(ρ, σ)
            Vμ = V*Cv[μ,a]
            Vν = V*Cv[ν,a]
            Vρ = V*Cv[ρ,a]
            Vσ = V*Cv[σ,a]

            if γab && γμν && γρσ
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,ν,σ,ρ] += Vμ
                aνρσ[a,μ,ρ,σ] += Vν
                aνρσ[a,μ,σ,ρ] += Vν
                aνρσ[a,σ,μ,ν] += Vρ
                aνρσ[a,σ,ν,μ] += Vρ
                aνρσ[a,ρ,μ,ν] += Vσ
                aνρσ[a,ρ,ν,μ] += Vσ

            elseif γab && γμν
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,μ,ρ,σ] += Vν
                aνρσ[a,σ,μ,ν] += Vρ
                aνρσ[a,σ,ν,μ] += Vρ

            elseif γab && γρσ
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,ν,σ,ρ] += Vμ
                aνρσ[a,σ,μ,ν] += Vρ
                aνρσ[a,ρ,μ,ν] += Vσ

            elseif γμν && γρσ
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,ν,σ,ρ] += Vμ
                aνρσ[a,μ,ρ,σ] += Vν
                aνρσ[a,μ,σ,ρ] += Vν
            elseif γab
                aνρσ[a,ν,ρ,σ] += Vμ
                aνρσ[a,σ,μ,ν] += Vρ
            else
                aνρσ[a,ν,ρ,σ] += Vμ
            end
        end
        end
        end
        end
    end
    @tensoropt (μ=>100x, ν=>100x, ρ=>100x, σ=>100x, i=>10x, j=>10x, k=>10x, l=>10x, a=>80x, b=>80x, c=>80x, d=>80) begin
        VVVV[a,b,c,d] :=  aνρσ[a, ν, ρ, σ]*Cv[ν, b]*Cv[ρ, c]*Cv[σ, d]
    end
    I["VVVV"] = VVVV
end

function compute_F!(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat, E<:AbstractERI, O<:AbstractRestrictedOrbitals}
    basis = I.orbitals.basis
    aoorbs = AtomicOrbitals(I.molecule, basis)
    aoints = IntegralHelper(molecule=I.molecule, orbitals=aoorbs, basis=basis, eri_type=I.eri_type)
    compute_F!(I, aoints)
end

function compute_F!(I::IntegralHelper{<:AbstractFloat,Chonky,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{<:AbstractFloat, Chonky, AtomicOrbitals})

    core = Options.get("drop_occ")
    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)

    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)

    Co = I.orbitals.C[:,1:ndocc]
    if eltype(I.orbitals.C) !== T
        Co = T.(Cv)
    end

    # Build Atomic Fock
    F = aoints["T"] + aoints["V"]
    ERI = aoints["ERI"]
    D = Co*transpose(Co)

    @tensor F[m,n] += 2*D[r,s]*ERI[m,n,r,s]
    @tensor F[m,n] -= D[r,s]*ERI[m,r,n,s]
    C = I.orbitals.C

    # Atomic Fock -> Molecular Fock
    @tensoropt Fmol[p,q] := C[μ,p]*F[μ,ν]*C[ν,q]

    Fd = FermiMDArray(diag(Fmol))
    I["Fd"] = Fd
    I["Fii"] = Fd[o]
    I["Faa"] = Fd[v]
    I["Fia"] = Fmol[o,v]
    I["Fij"] = Fmol[o,o] - diagm(Fd[o])
    I["Fab"] = Fmol[v,v] - diagm(Fd[v])
end

function compute_F!(I::IntegralHelper{T1,E1,<:AbstractRestrictedOrbitals}, aoints::IntegralHelper{T2,E2,AtomicOrbitals}) where {T1<:AbstractFloat,T2<:AbstractFloat,
                                                                                E1<:AbstractDFERI,E2<:AbstractDFERI}

    core = Options.get("drop_occ")
    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)

    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)

    Co = I.orbitals.C[:,1:ndocc]
    if eltype(I.orbitals.C) !== T1
        Co = T1.(Co)
    end

    # Build Atomic Fock
    F = aoints["T"] + aoints["V"]
    b = aoints["ERI"]
    D = Co*transpose(Co)

    @tensoropt F[m,n] += 2*D[r,s]*b[Q,m,n]*b[Q,r,s]
    @tensoropt F[m,n] -= D[r,s]*b[Q,m,r]*b[Q,n,s]
    C = I.orbitals.C

    # Atomic Fock -> Molecular Fock
    @tensoropt Fmol[p,q] := C[μ,p]*F[μ,ν]*C[ν,q]

    Fd = FermiMDArray(diag(Fmol))
    I["Fd"] = Fd
    I["Fii"] = Fd[o]
    I["Faa"] = Fd[v]
    I["Fia"] = Fmol[o,v]
    I["Fij"] = Fmol[o,o] - diagm(Fd[o])
    I["Fab"] = Fmol[v,v] - diagm(Fd[v])
end

function compute_F!(I::IntegralHelper{<:AbstractFloat,Chonky,RHFOrbitals}, aoints::IntegralHelper{<:AbstractFloat, SparseERI, AtomicOrbitals})
    compute_F!(I)
end

function compute_F!(I::IntegralHelper{<:AbstractFloat,RIFIT,RHFOrbitals}, aoints::IntegralHelper{<:AbstractFloat, JKFIT, AtomicOrbitals})
    compute_F!(I)
end

function compute_F!(I::IntegralHelper{T,E,RHFOrbitals}) where {T<:AbstractFloat, E<:AbstractERI}
    core = Options.get("drop_occ")
    inac = Options.get("drop_vir")
    ndocc = I.molecule.Nα
    nbf = size(I.orbitals.C,1)
    o = (1+core):ndocc
    v = (ndocc+1):(nbf - inac)
    No = length(o)
    Nv = length(v)

    I["Fd"] = T.(I.orbitals.eps)
    I["Fii"] = T.(I.orbitals.eps[o])
    I["Faa"] = T.(I.orbitals.eps[v])
    I["Fia"] = FermiMDzeros(T, No,Nv)
    I["Fij"] = FermiMDzeros(T, No,No)
    I["Fab"] = FermiMDzeros(T, Nv,Nv)
end

function compute_ref_energy(I::IntegralHelper{T,Chonky,O}) where {T<:AbstractFloat, O<:AbstractRestrictedOrbitals}
    ndocc = I.molecule.Nα
    o = 1:ndocc

    H = I["T"] + I["V"]

    OOOO = I["OOOO"]

    E0 = zero(T)
    for i in o
        E0 += 2*H[i,i]
        for j in o
            E0 += 2*OOOO[i,i,j,j] - OOOO[i,j,i,j]
        end
    end

    return E0
end

function compute_ref_energy(I::IntegralHelper{T,E,O}) where {T<:AbstractFloat,E<:AbstractDFERI,O<:AbstractRestrictedOrbitals}
    ndocc = I.molecule.Nα
    o = 1:ndocc

    H = I["T"] + I["V"]

    Boo = I["BOO"]

    E0 = zero(T)
    for i in o
        E0 += 2*H[i,i]
        for j in o
            E0 += 2*sum(Boo[:,i,i] .* Boo[:,j,j])
            E0 -= sum(Boo[:,i,j] .* Boo[:,i,j])
        end
    end

    return E0
end


text:
function updatevc!(system::CVSystem,n::Int64,terms::Vector{Int64})
    ivcflow = Float64(0);
    svcflow = Float64(0);

    # sum terminal/liver outflows
    for i = 1:length(terms)
        if system.branches.group[terms[i]] == "lower"
            if terms[i] != 12 && terms[i] != 13 && terms[i] != 14 && terms[i] != 15 && terms[i] != 21 # skip portal vein, hepatic vein inputs
                ivcflow += system.branches.term[terms[i]].Q[n+1,5];
            end
        else
            svcflow += system.branches.term[terms[i]].Q[n+1,5];
        end
    end
    ivcflow += system.liver.Q[n+1,4];

    # Euler forward update to vena cavae
    system.ivc.V[n+2] = (system.ivc.V[n+1] + system.solverparams.h*(
        ivcflow - system.ivc.Q[n+1]));
    system.ivc.Q[n+2] = ((1 - system.ivc.R/system.ivc.L*
        system.solverparams.h)*system.ivc.Q[n+1]+
        system.solverparams.h/system.ivc.L*(system.ivc.P[n+1]-
        system.heart.ra.P[n+1]));
    system.ivc.P[n+2] = (1/system.ivc.C*(system.ivc.V[n+2]-
        system.ivc.V0));
    system.svc.V[n+2] = (system.svc.V[n+1] + system.solverparams.h*(
        svcflow - system.svc.Q[n+1]));
    system.svc.Q[n+2] = ((1 - system.svc.R/system.svc.L*
        system.solverparams.h)*system.svc.Q[n+1]+
        system.solverparams.h/system.svc.L*(system.svc.P[n+1]-
        system.heart.ra.P[n+1]));
    system.svc.P[n+2] = (1/system.svc.C*(system.svc.V[n+2]-
        system.svc.V0));
end

max_stars_repo_name: dcanuto/CardioModeling

max_issues_repo_name: dcanuto/CardioModeling

max_forks_repo_name: dcanuto/CardioModeling


text:
using Test
using NSGAII
using Cambrian

include("./cornercases.jl")

@testset "crowdingDistanceAssignement" begin
    fitness1D=subPop1D()
    cfg=Cambrian.get_config("./test.yaml";n_population= 3,n_offsprings=2,d_fitness=1)
    fitness(x::FloatIndividual)=0
    dim1=NSGA2Evolution{FloatIndividual}(cfg,fitness)
    oldPop1D=copy(dim1.population)
    for i in 1:3
        copyto!(dim1.population[i].fitness,fitness1D[i])
    end
    crowdingDistanceAssignement!(dim1,dim1.population)
    @test length(oldPop1D)==length(dim1.population)
    for x in oldPop1D
        @test x in dim1.population
    end
    for x in dim1.population
        @test dim1.distance[objectid(x)]==0
    end

    fitness3D=subPop3D()
    cfg=Cambrian.get_config("./test.yaml";n_population= 5,n_offsprings=5,d_fitness=3)
    dim3=NSGA2Evolution{FloatIndividual}(cfg,fitness)
    oldPop3D=copy(dim3.population)
    for i in 1:5
        copyto!(dim3.population[i].fitness,fitness3D[i])
    end
    crowdingDistanceAssignement!(dim3,dim3.population)
    @test length(oldPop3D)==length(dim3.population)
    for x in oldPop3D
        @test x in dim3.population
    end
    max3D=maxPop(dim3.population)
    min3D=minPop(dim3.population)
    for x in max3D
        @test dim3.distance[objectid(x)]==Inf
    end
    for x in min3D
        @test dim3.distance[objectid(x)]==Inf
    end

    fitnessEQ=subPopEqual3D()
    cfg=Cambrian.get_config("./test.yaml";n_population= 8,n_offsprings=5,d_fitness=3)
    equal=NSGA2Evolution{FloatIndividual}(cfg,fitness)
    for i in 1:8
        copyto!(equal.population[i].fitness,fitnessEQ[i])
    end
    crowdingDistanceAssignement!(equal,equal.population)
    for x in equal.population
        @test equal.distance[objectid(x)]==0
    end
end