"""
`PSSFSS` is a software package for the analysis of polarization and frequency selective surfaces (PSSs and FSSs).
The user specifies the geometry to be analyzed as a `Vector` containing two or more dielectric [`Layer`](@ref)s
and zero or more [`Sheet`](@ref) objects denoting the PSS/FSS surfaces.  After also specifying the scan angles or
unit cell incremental phasings, frequencies to be considered, and optionally some output parameters to be written
to CSV file(s),
the user then invokes the [`analyze`](@refs) function to perform the analysis.  FSS/PSS triangulations
can be conveniently visualized using the `plot` command of the `Plots` package.

"""
module PSSFSS

if isdefined(Base, :Experimental) && isdefined(Base.Experimental, Symbol("@optlevel"))
    @eval Base.Experimental.@optlevel 3
end

using Reexport
using PkgVersion
using InteractiveUtils: versioninfo
using Dates: now
using DelimitedFiles: writedlm
using Printf: @sprintf
using LinearAlgebra: ×, norm, ⋅, factorize, lu!, ldiv!
using StaticArrays: StaticArrays, SVector, SArray, @SVector
using Unitful: ustrip, @u_str
using Logging: with_logger
using ProgressMeter

include("Constants.jl")
include("Log.jl")
include("PSSFSSLen.jl")
include("Rings.jl")
include("Layers.jl")
include("Sheets.jl")
include("Meshsub.jl")
include("Elements.jl")
include("RWG.jl")
include("PGF.jl")
include("Zint.jl")
include("FillZY.jl")
include("GSMs.jl")
include("Modes.jl")
include("Outputs.jl")

using .Rings
@reexport using .Sheets: Sheet, RWGSheet, read_sheet_data
using .RWG: setup_rwg, rwgbfft!, RWGData
using .GSMs: GSM, cascade, cascade!, gsm_electric_gblock, gsm_magnetic_gblock,
             gsm_slab_interface, translate_gsm!, choose_gblocks, Gblock, pecgsm, pmcgsm
using .FillZY: fillz, filly
using .Modes: zhatcross, choose_layer_modes!, setup_modes!
using .Constants: twopi, c₀, tdigits, dbmin
using .Log: pssfss_logger, @logfile
@reexport using .PSSFSSLen
@reexport using .Layers: Layer
@reexport using .Elements: rectstrip, diagstrip, polyring, meander, loadedcross,
                           jerusalemcross, pecsheet, pmcsheet
@reexport using .Outputs: @outputs, extract_result_file, extract_result
using .Outputs: Result, append_result_data

export analyze

Base.isfile(f::Base.DevNull) = false
Base.open(f::Base.DevNull,::AbstractString) = f

"""
    analyze(strata::Vector, flist, steering; outlist=[], logfile="pssfss.log", resultfile="pssfss.res", showprogress::Bool=true)

Analyze a full FSS/PSS structure over a range of frequencies and steering angles/phasings.
Generate output files as specified in `outlist`.

## Positional Arguments
- `strata`:  A vector of `Layer` and `Sheet` objects. The first and last entries must be of type `Layer`.

- `flist`: An iterable containing the analysis frequencies in GHz.

- `steering`: A length 2 `NamedTuple` containing as keys the steering parameter labels and as values
  the iterables that define the values of steering parameters to be analyzed.

    - one of {`:phi` ,`:ϕ`} and one of {`:theta`, `:θ`}, or

    - one of {`:psi1` ,`:ψ₁`} and one of {`:psi2`, `:ψ₂`}.

  All steering parameters are input in degrees.

  The program will analyze while iterating over a triple loop over the two steering
  parameters and frequency, with frequency in the innermost loop (i.e. varying the fastest).
  The steering parameter listed first will be in the outermost loop and will therefore
  vary most slowly.

## Keyword Arguments

- `outlist`:  A matrix with 2 columns.  The first column in each row is a string
  containing the name of the CSV file to write the output to.  The second entry in
  each row is a tuple generated by the `@outputs` macro of the `Outputs` module. The
  contents of the specified file(s) will be updated as the program completes each analysis
  frequency.

- `logfile`:  A string containing the name of the log file to which timing and other
  information about the run is written. Defaults to `"pssfss.log"`.
  If this file already exists, it will be overwritten.

- `resultfile`:  A string containing the name of the results file. Defaults to `pssfss.res`.
  If this file already exists, it will be overwritten.  It is a binary
  file that contains information (including the generalized scattering matrix) from
  the analysis performed for each scan condition and frequency. The result file can be
  post-processed to produce similar or additional outputs that were requested at run time
  via the `outlist` argument.

- `showprogress`: If true (default), then show progress bar during execution.

"""
function analyze(strata::Vector, flist, steering; outlist=[], logfile="pssfss.log",
    resultfile="pssfss.res", showprogress::Bool=true)
    layers = Layer[deepcopy(s) for s in strata if s isa Layer]
    sheets = RWGSheet[s for s in strata if s isa Sheet]
    islayer = map(x -> x isa Layer, strata)
    issheet = map(x -> x isa RWGSheet, strata)
    nl = length(layers)
    nj = nl - 1
    ns = length(sheets)
    sint = cumsum(islayer)[issheet] # sint[k] contains dielectric interface number of k'th sheet
    junc = zeros(Int, nj)
    junc[sint] = 1:ns #  junc[i] is the sheet number present at interface i, or 0 if no sheet is there
    freqstemp = float.(collect(flist))
    if length(freqstemp) < 2
        freqs = Float64[freqstemp]
    else
        freqs::Vector{Float64} = freqstemp
    end

    stkeys::Tuple{Symbol, Symbol} = keys(steering)
    stvaluestemp = [float.(collect(s)) for s in steering]
    stvalues = Vector{Float64}[]
    for (i,s) in pairs(stvaluestemp)
        if length(s) < 2
            push!(stvalues, Float64[s])
        else
            push!(stvalues, s)
        end
    end

    with_logger(pssfss_logger(logfile)) do
        _analyze(layers, sheets, junc, freqs, stkeys, stvalues;
                 outlist, resultfile, showprogress)
    end
end # function

"""
function _analyze(layers, sheets, junc, freqs, stkeys, stvalues;
    outlist=Any[], resultfile="pssfss.res", showprogress::Bool=true)

## Positional Arguments
- `layers`:  A vector of `Layer` objects.

- `sheets`:  A vector of `RWGSheet` objects.

- `junc`: `junc[k]` contains the sheet number present at layer interface `k` or `0` if no sheet
  is present there.

- `freqs`: A vector containing the analysis frequencies in GHz.

- `stkeys`: A length 2 `Tuple` containing as the steering angle labels as `Symbols` Either

    - one of {`:phi` ,`:ϕ`} and one of {`:theta`, `:θ`}, or

    - one of {`:psi1` ,`:ψ₁`} and one of {`:psi2`, `:ψ₂`}.

  The angular steering parameters are input in degrees, while the incremental phase shift
  parameters are input in radians.

  The program will analyze while iterating over a triple loop over the two steering
  parameters and frequency, with frequency in the innermost loop (i.e. varying the fastest).
  The steering parameter listed first will be in the outermost loop and will therefore
  vary most slowly.

## Keyword Arguments

- `outlist`:  A matrix with 2 columns.  The first column in each row is a string
  containing the name of the CSV file to write the output to.  The second entry in
  each row is a tuple generated by the `@outputs` macro of the `Outputs` module. The
  contents of the specified file(s) will be updated as the program completes each analysis
  frequency.

- `resultfile`:  A string containing the name of the results file. Defaults to `pssfss.res`.
  If this file already exists, it will be overwritten.  It is a binary
  file that contains information (including the generalized scattering matrix) from
  the analysis performed for each scan condition and frequency. The result file can be
  post-processed to produce similar or additional outputs that were requested at run time
  via the `outlist` argument.

- `showprogress`: If true, use ProgressMeter to show execution progress
"""
function _analyze(layers, sheets, junc, freqs, stkeys, stvalues;
    outlist=[], resultfile="pssfss.res", showprogress::Bool=true)
    showprogress && println("Beginning PSSFSS Analysis")
    ncount = 0 # Number of analyses performed
    ntotal = length(freqs) * length(stvalues[1]) * length(stvalues[2])
    showprogress && (progress = Progress(ntotal,1))
    showprogress && update!(progress, 0)
    isfile(resultfile) && rm(resultfile)
    date, clock = split(string(now()),'T')
    pssfssv = PkgVersion.Version(PSSFSS)
    ss = "Environment"
    io=IOBuffer(); versioninfo(io); juliainfo = String(take!(io)) * ss
    i = findfirst(ss, juliainfo)
    juliainfo = juliainfo[1:first(i)-1]
    juliainfo = juliainfo * "  Threads.nthreads() = $(Threads.nthreads())\n"
    @logfile "\n\nStarting PSSFSS $(pssfssv) analysis on $(date) at $(clock)\n$(juliainfo)\n\n"
    check_inputs(layers, sheets, junc, freqs, stkeys, stvalues, outlist)
    k0min, k0max = twopi*1e9/c₀ .* extrema(freqs)
    gbls = choose_gblocks(layers, sheets, junc, k0min)
    gsm_save = Vector{GSM}(undef, length(gbls)) # Storage for reusable GSMs
    choose_layer_modes!(layers, sheets, junc, gbls, k0max, dbmin)
    gbldup = get_gbldup(gbls, layers, sheets, junc)
    usi = unique_indices(sheets)
    rwgdat = RWGData[setup_rwg(sheet) for sheet in sheets]
    report_layers_sheets(layers, sheets, junc, rwgdat, usi)
    uvec::Vector{Float64} = map(sheets) do sh  # Green's function smoothing factors
        sh.style == "NULL" && (return 0.0)
        ufactor = 0.5 * ustrip(Float64, sh.units, 1u"m")
        ufactor * max(norm(sh.β₁), norm(sh.β₂))
    end
    # Begin analysis loops over steering angles and frequency
    firstoutput = true
    results = Result[]
    for stout in stvalues[1], stin in stvalues[2]
        steer = getsttuple(stkeys, stout, stin)
        if keys(steer)[1] == :ψ₁
            ψ₁, ψ₂ = deg2rad.([steer...]) # radians
            upm::Float64 = ustrip(Float64, sheets[1].units, 1u"m")
            β₁, β₂ = sheets[1].β₁*upm, sheets[1].β₂*upm
            β⃗₀₀ = (ψ₁*β₁ + ψ₂*β₂) / twopi # Eq. (2.13b)
        else
            θ, ϕ = steer # degrees
            st = sind(θ)
            sp, cp = sincosd(ϕ)
        end
        @logfile "Beginning $(steer)"
        for fghz in freqs
            @logfile "  $(fghz) GHz"
            t_freq = time()
            k0 = twopi*fghz*1e9/c₀
            if keys(steer)[1] == :θ
                k1 = k0 * sqrt(real(layers[1].ϵᵣ * layers[1].μᵣ))
                β⃗₀₀ = @SVector([k1*st*cp, k1*st*sp])
            end
            setup_modes!.(layers, k0, Ref(β⃗₀₀))
            if !(angle(layers[begin].γ[1]) ≈ angle(layers[end].γ[1]) ≈ π/2)
                @logfile "  Skipping $(fghz) GHz due to cutoff principal modes in ambient medium"
                continue
            end
            # Initialize overall GSM and propagate it through layer 1's width:
            n1 = length(layers[1].P)
            gsma = GSM(n1,n1)
            gsmc = GSM(n1,n1)
            cascade!(gsma, layers[1])
            for (ig,gbl) in pairs(gbls) # Walk through the Gblocks
                i1 = first(gbl.rng) # Index of layer to left of Gblock
                i2 = 1 + last(gbl.rng) # Index of layer to right of Gblock
                i_junc = gbl.j # junction where FSS is located, or 0 if no sheet
                i_sheet = i_junc == 0 ? 0 : junc[i_junc]
                if i_sheet ≠ 0
                    if gbldup[ig] > 0
                        gsmb = deepcopy(gsm_save[gbldup[ig]]) # Use previously calculated GSM
                    else
                        region = @view layers[i1:i2]
                        sheet = sheets[i_sheet]
                        s = gbl.j - i1 + 1 # sheet interface location within `region`
                        if sheet.class == 'J'
                            gsmb = calculate_jtype_gsm(region, sheet, uvec[i_sheet],
                                                       rwgdat[i_sheet], s, k0, β⃗₀₀, usi[i_sheet])
                        elseif sheet.class == 'M'
                            gsmb = calculate_mtype_gsm(region, sheet, uvec[i_sheet],
                                                       rwgdat[i_sheet], s, k0, β⃗₀₀, usi[i_sheet])
                        elseif sheet.class == 'E'
                            gsmb = pecgsm(length(region[1].P), length(region[end].P))
                        elseif sheet.class == 'H'
                            gsmb = pmcgsm(length(region[1].P), length(region[end].P))
                        else
                            error("Illegal sheet class: $(sheet.class)")
                        end

                        gbldup[ig] < 0 && (gsm_save[ig] = deepcopy(gsmb))
                    end
                    # Apply translations if requested:
                    sheet = sheets[i_sheet]
                    if sheet.dx ≠ 0 || sheet.dy ≠ 0
                        upm = ustrip(Float64, sheet.units, 1u"m")
                        dx = sheet.dx / upm
                        dy = sheet.dy / upm
                        translate_gsm!(gsmb, dx, dy, first(region), last(region))
                    end

                else # no sheet
                    @assert i2 - i1 == 1
                    gsmb = gsm_slab_interface(layers[i1], layers[i2], k0)
                end
                gsmc = cascade(gsma, gsmb)
                cascade!(gsmc, layers[i2])
                gsma = gsmc
            end # Gblock loop
            t_freq = round(time()-t_freq, digits=tdigits)
            @logfile "  $(t_freq) seconds total at $(fghz) GHz"

            result = Result(gsmc, steer, β⃗₀₀, fghz, layers[1].ϵᵣ, layers[1].μᵣ,
            layers[1].β₁, layers[1].β₂, layers[end].ϵᵣ, layers[end].μᵣ,
            layers[end].β₁, layers[end].β₂)
            push!(results, result)
            ncount += 1
            # Write to output files
            append_result_data(resultfile,string(ncount),result)
            for row in eachrow(outlist)
                if firstoutput
                    open(row[1], "w") do io
                        writedlm(io, permutedims([r.label for r in row[2]]), ',')
                    end
                end
                open(row[1], "a") do io
                    writedlm(io, permutedims([r(result) for r in row[2]]), ',')
                end
            end
            firstoutput = false
            showprogress && next!(progress) # Bump progress meter
        end # Frequency loop
    end # steering angle loop

    date, clock = split(string(now()),'T')
    @logfile "\n\n PSSFSS analysis exiting on $(date) at $(clock)\n\n"
    return results
end # function

"""
    calculate_jtype_gsm(layers, sheet::RWGSheet, u::Real, rwgdat::RWGData, s::Int, k0, k⃗inc, is::Int) -> gsm

Compute the generalized scattering matrix for a sheet of class `'J'`.

### Input Arguments

- `layers`: An iterable of `Layer` instances containing the layers for the `Gblock`
    associated with the `sheet` under consideration.
- `sheet`:  A sheet of class `'J'` for which the GSM is desired.
- `u`: The Green's function smoothing parameter for the `sheet`.
- `rwgdat`: The `RWGData` object associated with the `sheet` argument.
- `s`: The interface number within `layers` at which the sheet is located.
- `k0`: The free-space wavenumber in radians/meter.
- `k⃗inc`: A 2-vector containing the incident field wave vector x and y components. Note
    that by the phase match condition this vector is the same for all layers in the entire FSS
    structure.
- `is`: The sheet index for the `sheet` argument within the global list of sheets.

### Return Value

- `gsm::GSM`  The full GSM for the GBlock including incident fields and scattered fields due
    to currents induced on the sheet surface.
"""
function calculate_jtype_gsm(layers::AbstractVector{Layer}, sheet::RWGSheet, u::Float64,
                                 rwgdat::RWGData, s::Int, k0::Float64, k⃗inc::SVector{2, Float64}, is::Int)
    one_meter = ustrip(Float64, sheet.units, 1u"m")
    area = norm(sheet.s₁ × sheet.s₂) / one_meter^2 # Unit cell area (m^2).
    nmodesmax = max(length(layers[begin].P), length(layers[end].P))
    nbf = size(rwgdat.bfe,2) # Number of basis functions
    if isempty(rwgdat.bfftstore)
        rwgdat.bfftstore = zeros(SArray{Tuple{2},ComplexF64,1,2}, (nbf, 2, nmodesmax))
    end
    bfftstore = rwgdat.bfftstore
    # Compute area correction factors for the mode normalization constants of
    # the two end regions:
    tempvec = zeros(2)
    for (i,l) in enumerate(@view layers[[begin,end]])
      area_i = twopi * twopi / norm(l.β₁ × l.β₂)
      tempvec[i] = √(area_i / area)
    end
    acf = SVector(tempvec[1], tempvec[2])

    # Set up the partial GSM due to incident field
    (gsm, tlgfvi, vincs) = gsm_electric_gblock(layers, s, k0)
    if sheet.style == "NULL"
        return gsm
    end

    # Calculate the scattered field partial scattering matrix of the
    # FSS sheet at this junction. Then add it to GSM already computed
    # for the dielectric discontinuity...

    #  Fill the interaction matrix for the current sheet:
    t_temp = time()
    ψ₁ = k⃗inc ⋅ sheet.s₁ / one_meter
    ψ₂ = k⃗inc ⋅ sheet.s₂ / one_meter
    @logfile "    Beginning matrix fill for sheet $(is)"
    zmat = fillz(k0,u,layers,s,ψ₁,ψ₂,sheet,rwgdat)
    t_fill = round(time() - t_temp, digits=tdigits)
    @logfile "      $(t_fill) seconds total matrix fill time for sheet $(is)"
    # Factor the matrix:
    t_temp = time()
    zmatf = lu!(zmat)
    t_factor = round(time() - t_temp, digits=tdigits)
    @logfile "      $(t_factor) seconds to factor matrix for sheet $(is)"
    t_temp = time()
    # Compute and store the basis function Fourier transforms:
    i_ft = 0
    for (sr,l) in enumerate(@view layers[[begin,end]]) # loop over possible source regions
        for qp = 1:length(l.P)
            kvec = l.β[qp]
            # If desired F.T. has already been computed, then copy it.
            if qp > 1 && kvec ≈ l.β[qp-1]
                bfftstore[:,sr,qp] = bfftstore[:,sr,qp-1]
                continue
            end
            if sr == 2 && length(layers[1].β) ≥ qp && kvec ≈ layers[1].β[qp]
                bfftstore[:,2,qp] = bfftstore[:,1,qp]
                continue
            end
            bfft = @view bfftstore[:,sr,qp]
            rwgbfft!(bfft, rwgdat, sheet, kvec, ψ₁, ψ₂) # Otherwise, compute from scratch
            i_ft += 1
        end
    end
    t_fft = round(time() - t_temp, digits=tdigits)
    #@logfile "      $(t_fft) seconds for basis function Fourier transforms at $(i_ft) points"
    nsolve = 0
    t_extract = 0.0
    i_extract = 0
    t_solve = 0.0
    imat = rwgdat.rhs
    for (sr,ls) in enumerate(@view layers[[begin,end]]) # Loop over source regions
        for qp in 1:length(ls.P) # Loop over srce reg modes
            # Incident field for source layer in absence of the FSS sheet:
            sourcevec = vincs[qp,sr] * ls.c[qp] * acf[sr] * ls.tvec[qp]
            # Compute generalized voltage vector:
            imat .= (b ⋅ sourcevec for b in bfftstore[:,sr,qp]) # Eq. (7.39)
            # Solve the matrix equation
            t_solve1 = time()
            ldiv!(zmatf, imat)
            t_solve2 = time()
            t_solve += t_solve2 - t_solve1
            nsolve += 1
            t_extract1 = time()
            for (or,lo) in enumerate(@view layers[[begin,end]]) # Loop over obs. regions
                smat = gsm[or,sr]
                for q in 1:length(lo.P)  # Loop obs. regn modes
                    # Extract partial scattering parameter due to scattered fields...
                    FTJ = sum((imat[n] * bfftstore[n,or,q] for n in 1:nbf)) # FT of total current
                    smat[q,qp] -= (lo.tvec[q] ⋅ FTJ) * (tlgfvi[q,or] /
                                                (lo.c[q] * acf[or] * area)) # Eq (6.18)
                    i_extract += 1
                end
            end
            t_extract2 = time()
            t_extract = t_extract + (t_extract2 - t_extract1)
        end
    end
    #@logfile "      $(round(t_extract,digits=tdigits)) seconds to extract $(i_extract) GSM entries"
    return gsm
end

"""
    calculate_mtype_gsm(layers, sheet::RWGSheet, u::Real, rwgdat::RWGData, s::Int, k⃗inc, is::Int) -> gsm

Compute the generalized scattering matrix for a sheet of class `'M'`.

### Input Arguments

- `layers`: An iterable of `Layer` instances containing the layers for the `Gblock`
    associated with the `sheet` under consideration.
- `sheet`:  A sheet of class `'M'` for which the GSM is desired.
- `u`: The Green's function smoothing parameter for the `sheet`.
- `rwgdat`: The `RWGData` object associated with the `sheet` argument.
- `s`: The interface number within `layers` at which the sheet is located.
- `k0`: The free-space wavenumber in radians/meter.
- `k⃗inc`: A 2-vector containing the incident field wave vector x and y components. Note
    that by the phase match condition this vector is the same for all layers in the entire FSS
    structure.
- `is`: The sheet index for the `sheet` argument within the global list of sheets.

### Return Value

- `gsm::GSM`  The full GSM for the GBlock including incident fields and scattered fields due
    to magnetic currents induced in the gaps on the sheet surface.
"""
function calculate_mtype_gsm(layers::AbstractVector{Layer}, sheet::RWGSheet, u::Float64,
                                 rwgdat::RWGData, s::Int, k0::Float64, k⃗inc::SVector{2, Float64}, is::Int)
    one_meter = ustrip(Float64, sheet.units, 1u"m")
    area = norm(sheet.s₁ × sheet.s₂) / one_meter^2 # Unit cell area (m^2).
    nmodesmax = max(length(layers[begin].P), length(layers[end].P))
    nbf = size(rwgdat.bfe,2) # Number of basis functions
    if isempty(rwgdat.bfftstore)
        rwgdat.bfftstore = zeros(SArray{Tuple{2},ComplexF64,1,2}, (nbf, 2, nmodesmax))
    end
    bfftstore = rwgdat.bfftstore
    # Compute area correction factors for the mode normalization constants of
    # the two end regions:
    tempvec = zeros(2)
    for (i,l) in enumerate(@view layers[[begin,end]])
      area_i = twopi * twopi / norm(l.β₁ × l.β₂)
      tempvec[i] = √(area_i / area)
    end
    acf = SVector(tempvec[1], tempvec[2])

    # Set up the partial GSM due to incident field
    (gsm, tlgfiv, iincs) = gsm_magnetic_gblock(layers, s, k0)
    if sheet.style == "NULL"
        return gsm
    end

    # Calculate the scattered field partial scattering matrix of the
    # FSS sheet at this junction. Then add it to GSM already computed
    # for the dielectric discontinuity...

    #  Fill the interaction matrix for the current sheet:
    t_temp = time()
    ψ₁ = k⃗inc ⋅ sheet.s₁ / one_meter
    ψ₂ = k⃗inc ⋅ sheet.s₂ / one_meter
    @logfile "    Beginning matrix fill for sheet $(is)"
    ymat = filly(k0,u,layers,s,ψ₁,ψ₂,sheet,rwgdat)
    t_fill = round(time() - t_temp, digits=tdigits)
    @logfile "      $(t_fill) seconds total matrix fill time for sheet $(is)"
    # Factor the matrix:
    t_temp = time()
    ymatf = lu!(ymat)
    t_factor = round(time() - t_temp, digits=tdigits)
    @logfile "      $(t_factor) seconds to factor matrix for sheet $(is)"
    t_temp = time()
    # Compute and store the basis function Fourier transforms:
    i_ft = 0
    for (sr,l) in enumerate(@view layers[[begin,end]]) # loop over possible source regions
        for qp = 1:length(l.P)
            kvec = l.β[qp]
            # If desired F.T. has already been computed, then copy it.
            if qp > 1 && kvec ≈ l.β[qp-1]
                bfftstore[:,sr,qp] = bfftstore[:,sr,qp-1]
                continue
            end
            if sr == 2 && length(layers[1].β) ≥ qp && kvec ≈ layers[1].β[qp]
                bfftstore[:,2,qp] = bfftstore[:,1,qp]
                continue
            end
            bfft = @view bfftstore[:,sr,qp]
            rwgbfft!(bfft, rwgdat, sheet, kvec, ψ₁, ψ₂) # Otherwise, compute from scratch
            i_ft += 1
        end
    end
    t_fft = round(time() - t_temp, digits=tdigits)
    @logfile "      $(t_fft) seconds for basis function Fourier transforms at $(i_ft) points"
    nsolve = 0
    t_extract = 0.0
    t_solve = 0.0
    vmat = rwgdat.rhs
    i_extract = 0
    σ = -1
    for (sr,ls) in enumerate(@view layers[[begin,end]]) # Loop over source regions
        σ *= -1 # 1 for sr == 1, and -1 for sr == 2
        for qp in 1:length(ls.P) # Loop over srce reg modes
            # Incident field for Region sr (Eq. (7.64))
            sourcevec = iincs[qp,sr] * ls.c[qp] * ls.Y[qp] * zhatcross(ls.tvec[qp])
            # Compute generalized current vector:
            vmat .= (b ⋅ sourcevec for b in bfftstore[:,sr,qp]) # Eq. (7.64)
            # Solve the matrix equation
            t_solve1 = time()
            ldiv!(ymatf, vmat)
            t_solve2 = time()
            t_solve += t_solve2 - t_solve1
            nsolve += 1
            t_extract1 = time()
            for (or,lo) in enumerate(@view layers[[begin,end]]) # Loop over obs. regions
                smat = gsm[or,sr]
                for q in 1:length(lo.P)  # Loop obs. regn. modes
                    # Extract partial scattering parameter due to scattered fields...
                    FTM = sum((vmat[n] * bfftstore[n,or,q] for n in 1:nbf)) # FT of total mag. current
                    smat[q,qp] += (zhatcross(lo.tvec[q]) ⋅ FTM) *
                                        (σ * tlgfiv[q,or] * lo.c[q]) # Eq. (6.37)
                    i_extract += 1
                end
            end
            t_extract2 = time()
            t_extract += t_extract2 - t_extract1
        end
    end
    @logfile "      $(t_extract) seconds to extract $(i_extract) GSM entries"
    return gsm
end

"""
    getsttuple(stkeys::Tuple{Symbol,Symbol}, stout::Float64, stin::Float64) -> NamedTuple

Return a named tuple either of the form `(θ = θ, ϕ = ϕ)` or `(ψ₁ = ψ₁, ψ₂ = ψ₂)` that serves
to define the current steering situation. Actually, the input field names can be spelled out in
English as `:theta`, `:phi`, `:psi1`, and `:psi2`.

### Arguments

- `stkeys`: A 2-tuple containing steering parameters as `Symbols`, either (`:ψ₁` and `ψ₂`) or (`:θ` and `:ϕ`)
  (or their spelled-out English versions as detailed above), either of which
  could be listed in either order.  The order is significant in that the first member of the pair
  defines the outer steering loop.

- `stout` and `stin`: These are the current values of the outer and inner steering variables,
  respectively.
"""
function getsttuple(stkeys::Tuple{Symbol, Symbol}, stout::Float64, stin::Float64)
    if stkeys[1] ∈ (:phi, :ϕ, :Phi, :PHI, :Φ)
        return (θ = stin, ϕ = stout)
    elseif stkeys[2] ∈ (:phi, :ϕ, :Phi, :PHI, :Φ)
        return (θ = stout, ϕ = stin)
    elseif keys(steering)[1] ∈ (:psi1, :Psi1, :PSI1, :ψ₁, :Ψ₁, :ψ1, :Ψ1)
        return (ψ₁ = stout, ψ₂ = stin)
    else
        return (ψ₁ = stin, ψ₂ = stout)
    end
end

function check_inputs(layers, sheets, junc, freqs, stkeys, stvalues, outlist)
    # Check that input and output media are lossless and the same
    imag(first(layers).ϵᵣ) == imag(last(layers).ϵᵣ) ==
    imag(first(layers).μᵣ) == imag(last(layers).μᵣ) == 0 ||
    error("First and last layers must be lossless")
    (first(layers).ϵᵣ == last(layers).ϵᵣ) &&
    (first(layers).μᵣ == last(layers).μᵣ) ||
    error("First and last layers must have identical electrical parameters")

    # todo:

    # Check that ψ₁ and ψ₂ are not specified when there are no nonnull sheets

    return
end

"""
    unique_indices(v::Vector)

Return a vector `ui` of the same length as `v`.
`ui[k]` contains the smallest `i` such that 'i ≤ k` and `ui[i] === ui[k]`
"""
function unique_indices(v::Vector)
    n = length(v)
    ui = collect(1:n)
    for io in 1:n-1
        ui[io] ≠ io && continue
        for it in io+1:n
            ui[it] ≠ it && continue
            v[io] === v[it]  && (ui[it] = io)
        end
    end
    return ui
end

"""
    get_gbldup(gbls::Vector{Gblock}, layers::Vector{Layer}, sheets::Vector{RWGSheet}, junc::Vector{Int})
    -> (gbldup, junc)

Return `gbldup::Vector{Int}` of the same length as `gbls`.
`gbldup[k]` contains `0` for an ordinary Gblock.  `gbldup[k] == -1` means that
`gbls[k]` is the first occurence of a repeated Gblock and that its GSM should be
saved for reuse.  `gbldup[k] == i` where `0<i<k` means that `gbls[k]` is identical
to `gbls[i]` and they can both use the same GSM.
Two `Gblock`s are considered identical if they

1. Contain identical (`===`) `Sheet` objects at the same location within the block.
2. Comprise the same number of dielectric layers with identical widths analyzed
   electrical characteristics.
3. Are embedded within similar adjacent dielectric layers, having identical electrical
   properties and numbers of modes.

`junc::Vector{Int}` is has length `length(Layers)-1`. `junc[i]`` is the sheet number
present at dielectric interface `i`, or `0` if no sheet is present there.
"""
function get_gbldup(gbls::Vector{Gblock}, layers::Vector{Layer}, sheets::Vector{RWGSheet}, junc::Vector{Int})
    gbldup = zeros(Int, length(gbls))
    for (g1,gbl1) in pairs(gbls)
        (gbldup[g1] ≠ 0 || gbl1.j == 0) && continue
        j1, rng1 = gbl1.j, gbl1.rng
        n1 = length(layers[first(rng1)].P) # modes at left side
        n2 = length(layers[last(rng1)+1].P) # modes at right side
        # Examine succeeding Gblocks to see if they match:
        for g2 in g1+1:length(gbls)
            gbl2 = gbls[g2]
            rng2 = gbl2.rng
            j2 = gbl2.j
            (gbldup[g2] ≠ 0 || j2 == 0) && continue
            sheets[junc[j1]] === sheets[junc[j2]] || continue
            last(rng1)-j1 ≠ last(rng2)-j2 && continue
            # Check that layers within blocks gbl1 and gbl2 are identical:
            length(rng1) ≠ length(rng2) && continue
            n1 ≠ length(layers[first(rng2)].P) && continue
            n2 ≠ length(layers[1+last(rng2)].P) && continue
            all(zip(rng1,rng2)) do (i1,i2)
                i1 == first(rng1) && (return true)
                layers[i1] == layers[i2]
            end || continue
            # If we made it to here, the two Gblocks are identical:
            gbldup[g1] = -1  # Indicate that GSM of Gblock g1 is to be saved
            gbldup[g2] = g1  # GSM of Gblock g2 is obtained from saved GSM of block g1
        end
    end
    return gbldup
end # function

function report_layers_sheets(layers, sheets, junc, rwgdat, usi)
    @logfile "Dielectric layer information... \n"
    @logfile " Layer  Width  units  epsr   tandel   mur  mtandel modes  beta1x  beta1y  beta2x  beta2y"
    @logfile " ----- ------------- ------- ------ ------- ------ ----- ------- ------- ------- -------"
    for (jl,l) in pairs(layers)
        eps = real(l.ϵᵣ)
        tandel = -imag(l.ϵᵣ) / eps
        mu = real(l.μᵣ)
        mtandel = -imag(l.μᵣ) / mu
        nmode = length(l.P)
        units = string(unit(l.user_width))
        units == "inch" && (units = "in")
        uw_unitless = ustrip(l.user_width)
        str = @sprintf(" %5i %9.4f %3s %7.2f %6.4f %7.2f %6.4f %5i %7.1f %7.1f %7.1f %7.1f",
              jl,uw_unitless, units, eps, tandel, mu, mtandel, nmode, l.β₁[1], l.β₁[2],
              l.β₂[1], l.β₂[2])
        @logfile "$str"
        if jl < length(layers) && junc[jl] ≠ 0
            js = junc[jl]
            s = sheets[js]
            om = ustrip(Float64, s.units, 1.0u"m")
            str = @sprintf(
                " ==================  Sheet %3i  ======================== %7.1f %7.1f %7.1f %7.1f",
                usi[js], s.β₁[1] * om, s.β₁[2] * om, s.β₂[1] * om, s.β₂[2] * om)
            @logfile "$str"
        end
    end
    sint = findall(junc .≠ 0) # sint[k] contains dielectric interface number of k'th sheet

    @logfile "\n\n\nPSS/FSS sheet information...\n"
    @logfile "Sheet  Loc         Style      Rot  J/M Faces Edges Nodes Unknowns  NUFP"
    @logfile "-----  ---  ---------------- ----- --- ----- ----- ----- -------- ------"
    for (js,s) in pairs(sheets)
        str = @sprintf("%4i   %3i  %16s %5.1f  %1s  %5i %5i %5i  %6i %7i",
             usi[js], sint[js], s.style, s.rot, s.class, size(s.fe,2), length(s.e1),
             length(s.ρ), size(rwgdat[js].bfe,2), length(rwgdat[js].ufp2fp))
        @logfile "$str"
    end
    @logfile "\n\n"
    nothing
end

end # module


lattice = load_lattice(joinpath(@__DIR__, "GaAs20.json"))
atoms   = load_atoms(  joinpath(@__DIR__, "GaAs20.json"))
atoms   = attach_pseudos(atoms, Ga="hgh/pbe/ga-q3.hgh", As="hgh/pbe/as-q5.hgh")

cases = standard_cases()
cases["adaptive_03"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.3, α_min=0.05), tol=1e-12)
cases["adaptive_04"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.4, α_min=0.05), tol=1e-12)
cases["adaptive_05"] = (damping=DFTK.AdaptiveDamping(;α_trial_min=0.5, α_min=0.05), tol=1e-12)

function load_GaAs_20_potential(basis)
    guessfile = joinpath(@__DIR__, "GaAs20_initV.jld2")
    if !isfile(guessfile)
        ρ0 = guess_density(basis)
        _, ham = DFTK.energy_hamiltonian(basis, nothing, nothing; ρ=ρ0)
        V = DFTK.total_local_potential(ham)
        V .+= 3.0rand(Float64, size(V)...)

        jldopen(guessfile, "w") do file
            file["V"] = V
        end
    end

    jldopen(guessfile, "r") do file
        file["V"]
    end
end

# Load guess from file
config[@__DIR__] = (
    directory        = @__DIR__,
    lattice          = lattice,
    atoms            = atoms,
    magnetic_moments = [],
    #
    kwargs_model     = (smearing=Smearing.None(), temperature=0.0),
    Ecut             = 20,
    kgrid            = [4, 4, 1],
    supersampling    = 1.5,
    potential        = load_GaAs_20_potential,
    kwargs_scf       = (; ),
    mixing           = SimpleMixing(),
    cases            = cases,
)


include("autotuning_utils.jl")
include("zygote_rules.jl")
include("forwarddiff_rules.jl")

function update_hyperparameters!(m::GP, state, x, y)
    μ₀ = pr_mean(m.f) # Get prior means
    k = kernel(m.f) # Get kernels
    if !isnothing(opt(m.f))
        if ADBACKEND[] == :Zygote
            hp_state = state.hyperopt_state
            # Compute gradients
            Δμ₀, Δk = Zygote.gradient(μ₀, k) do μ₀, k # Compute gradients for the whole model
                ELBO(m, x, y, μ₀, k)
            end

            # Optimize prior mean
            hp_state = state.hyperopt_state
            if !isnothing(Δμ₀)
                hp_state = update!(μ₀, Δμ₀, hp_state)
            end

            # Optimize kernel parameters
            if isnothing(Δk)
                @warn "Kernel gradients are equal to zero" maxlog = 1
            else
                state_k = update_kernel!(opt(m.f), kernel(m.f), Δk, hp_state)
                hp_state = merge(hp_state, (; state_k))
            end
            state = merge(state, (; hyperopt_state=hp_state))
        elseif ADBACKEND[] == :ForwardDiff
            θ, re = destructure((μ₀, k))
            Δ = ForwardDiff.gradient(θ) do θ
                ELBO(m, re(θ)...)
            end
        end
    end
    # end
    return state
end

# @traitfn function update_hyperparameters!(
#     m::TGP,
# ) where {TGP <: AbstractGPModel; IsFull{TGP}}
#     update_hyperparameters!.(m.f, [xview(m)])
#     setHPupdated!(m.inference, true)
# end

@traitfn function update_hyperparameters!(
    m::TGP, state, x, y
) where {TGP <: AbstractGPModel; IsFull{TGP}}
    if any((!) ∘ isnothing ∘ opt, m.f) # Check there is a least one optimiser
        hp_state = state.hyperopt_state
        μ₀ = pr_means(m) # Get prior means
        ks = kernels(m) # Get kernels
        if ADBACKEND[] == :Zygote
            Δμ₀, Δk = Zygote.gradient(μ₀, ks) do μ₀, ks # Compute gradients for the whole model
                ELBO(m, x, y, μ₀, ks, state)
            end
            # Optimize prior mean
            if !isnothing(Δμ₀)
                hp_state = update!.(μ₀, Δμ₀, hp_state)
            end
            if isnothing(Δk)
                @warn "Kernel gradients are equal to zero" maxlog = 1
            else
                hp_state = map(m.f, Δk, hp_state) do gp, Δ, hp_st
                    if isnothing(opt(gp))
                        return hp_st
                    else
                        state_k = update_kernel!(opt(gp), kernel(gp), Δ, hp_st.state_k)
                        return merge(hp_st, (; state_k))
                    end
                end
            end
        elseif ADBACKEND[] == :ForwardDiff
            θ, re = destructure((μ₀, ks))
            Δ = ForwardDiff.gradient(θ) do θ
                ELBO(m, re(θ)...)
            end
        end
        state = merge(state, (; hyperopt_state=hp_state))
    end
    return state
end

@traitfn function update_hyperparameters!(
    m::TGP, state, x, y
) where {TGP <: AbstractGPModel; !IsFull{TGP}}
    # Check that here is least one optimiser
    if any((!) ∘ isnothing ∘ opt, m.f) || any((!) ∘ isnothing ∘ Zopt, m.f)
        hp_state = state.hyperopt_state
        μ₀ = pr_means(m)
        ks = kernels(m)
        Zs = Zviews(m)
        if ADBACKEND[] == :Zygote
            Δμ₀, Δk, ΔZ = Zygote.gradient(μ₀, ks, Zs) do μ₀, ks, Zs
                ELBO(m, x, y, μ₀, ks, Zs, state)
            end
            # Optimize prior mean
            if !isnothing(Δμ₀)
                hp_state = update!.(μ₀, Δμ₀, hp_state)
            end

            # Optimize kernel parameters
            if isnothing(Δk)
                @warn "Kernel gradients are equal to zero" maxlog = 1
            else
                hp_state = map(m.f, Δk, hp_state) do gp, Δ, hp_st
                    if !isnothing(opt(gp))
                        state_k = update_kernel!(opt(gp), kernel(gp), Δ, hp_st.state_k)
                        return merge(hp_st, (; state_k))
                    else
                        return hp_st
                    end
                end
            end

            # Optimize inducing point locations
            if isnothing(ΔZ)
                @warn "Inducing point locations gradients are equal to zero" maxlog = 1
            else
                hp_state = map(m.f, ΔZ, hp_state) do gp, Δ, hp_st
                    if !isnothing(Zopt(gp))
                        state_Z = update_Z!(Zopt(gp), Zview(gp), Δ, hp_st.state_Z)
                        return merge(hp_st, (; state_Z))
                    else
                        return hp_st
                    end
                end
            end
        elseif ADBACKEND[] == :ForwardDiff
            θ, re = destructure((μ₀, ks, Zs))
            Δ = ForwardDiff.gradient(θ) do θ
                ELBO(m, re(θ)...)
            end
        end
        state = merge(state, (; hyperopt_state=hp_state))
    end
    return state
end

## Update all hyperparameters for the full batch GP models ##
function update_hyperparameters!(gp::AbstractLatent, X::AbstractVector)
    if !isnothing(gp.opt)
        f_l, f_μ₀ = hyperparameter_gradient_function(gp, X)
        ad_backend = K_ADBACKEND[] == :auto ? ADBACKEND[] : K_ADBACKEND[]
        Δμ₀ = f_μ₀()
        Δk = if ad_backend == :forward
            ∇L_ρ_forward(f_l, gp, X)
        elseif ad_backend == :zygote
            ∇L_ρ_zygote(f_l, gp, X)
        end
        apply_Δk!(gp.opt, kernel(gp), Δk) # Apply gradients to the kernel parameters
        apply_gradients_mean_prior!(pr_mean(gp), Δμ₀, X)
    end
end

## Update all hyperparameters for the sparse variational GP models ##
function update_hyperparameters!(
    gp::AbstractLatent,
    X::AbstractVector,
    ∇E_μ::AbstractVector,
    ∇E_Σ::AbstractVector,
    i::AbstractInference,
    vi_opt::InferenceOptimizer,
)
    Δμ₀, Δk = if !isnothing(gp.opt)
        f_ρ, f_Z, f_μ₀ = hyperparameter_gradient_function(gp)
        ad_backend = K_ADBACKEND[] == :auto ? ADBACKEND[] : K_ADBACKEND[]
        Δμ₀ = f_μ₀()
        Δk = if ad_backend == :forward
            ∇L_ρ_forward(f_ρ, gp, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        elseif ad_backend == :zygote
            ∇L_ρ_zygote(f_ρ, gp, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        end
        (Δμ₀, Δk)
    else
        nothing, nothing
    end
    if !isnothing(Zopt(gp))
        ad_backend = Z_ADBACKEND[] == :auto ? ADBACKEND[] : Z_ADBACKEND[]
        Z_grads = if ad_backend == :forward
            Z_gradient_forward(gp, f_Z, X, ∇E_μ, ∇E_Σ, i, vi_opt) #Compute the gradient given the inducing points location
        elseif ad_backend == :zygote
            Z_gradient_zygote(gp, f_Z, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        end
        update_Z!(opt(gp.Z), gp.Z, Z_grads) #Apply the gradients on the location
    end
    if !all([isnothing(Δk), isnothing(Δμ₀)])
        apply_Δk!(gp.opt, kernel(gp), Δk) # Apply gradients to the kernel parameters
        apply_gradients_mean_prior!(pr_mean(gp), Δμ₀, X)
    end
end

function update_hyperparameters!(
    gp::AbstractLatent,
    l::AbstractLikelihood,
    i::AbstractInference,
    X::AbstractVector,
    Y::AbstractVector,
)
    Δμ₀, Δk = if !isnothing(gp.opt)
        ad_backend = K_ADBACKEND[] == :auto ? ADBACKEND[] : K_ADBACKEND[]
        (Δμ₀, Δk) = if ad_backend == :forward
            ∇L_ρ_forward(f_ρ, gp, X, ∇E_μ, ∇E_Σ, i, vi_opt)
        elseif ad_backend == :zygote
            ∇L_ρ_zygote(gp, l, i, X, Y)
        end
    else
        nothing, nothing
    end
    if !isnothing(opt(gp.Z))
        ad_backend = Z_ADBACKEND[] == :auto ? ADBACKEND[] : Z_ADBACKEND[]
        global Z_grads = if ad_backend == :forward
            Z_gradient_forward(gp, f_Z, X, ∇E_μ, ∇E_Σ, i, vi_opt) #Compute the gradient given the inducing points location
        elseif ad_backend == :zygote
            Z_gradient_zygote(gp, l, i, X, Y)
        end
        update!(opt(gp.Z), gp.Z.Z, Z_grads) #Apply the gradients on the location
    end
    if !all([isnothing(Δk), isnothing(Δμ₀)])
        apply!(kernel(gp), Δk, gp.opt) # Apply gradients to the kernel parameters
        update!(pr_mean(gp), Δμ₀, X)
    end
end

## Return the derivative of the KL divergence between the posterior and the GP prior ##
function hyperparameter_KL_gradient(J::AbstractMatrix, A::AbstractMatrix)
    return trace_ABt(J, A) / 2
end

function hyperparameter_gradient_function(gp::LatentGP{T}, ::AbstractVector) where {T}
    A = (inv(cov(gp)) - mean(gp) * transpose(mean(gp))) # μ = inv(K+σ²)*(y-μ₀)
    return (function (Jnn)
        return -hyperparameter_KL_gradient(Jnn, A)
    end, function ()
        return -mean(gp)
    end)
end

## Return functions computing gradients of the ELBO given the kernel hyperparameters for a non-sparse model ##
function hyperparameter_gradient_function(
    gp::VarLatent{T}, X::AbstractVector
) where {T<:Real}
    μ₀ = pr_mean(gp, X)
    A =
        (I - pr_cov(gp) \ (cov(gp) + (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))) /
        pr_cov(gp)
    return (
        function (Jnn)
            return -hyperparameter_KL_gradient(Jnn, A)
        end,
        function ()
            return pr_cov(gp) \ (mean(gp) - μ₀)
        end,
    )
end

## Return functions computing gradients of the ELBO given the kernel hyperparameters for a non-sparse latent GP ##
function hyperparameter_gradient_function(gp::SparseVarLatent{T}) where {T<:Real}
    μ₀ = pr_mean(gp, gp.Z)
    A =
        (I(dim(gp)) - pr_cov(gp) \ (cov(gp) + (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))) /
        cov(gp)
    κΣ = gp.κ * cov(gp)
    return (
        function (Jmm, Jnm, Jnn, ∇E_μ, ∇E_Σ, i, opt)
            return (
                hyperparameter_expec_gradient(gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, Jnn) -
                hyperparameter_KL_gradient(Jmm, A)
            )
        end,
        function (Jmm, Jnm, ∇E_μ, ∇E_Σ, i, opt)
            return hyperparameter_expec_gradient(
                gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, zero(gp.K̃)
            ) - hyperparameter_KL_gradient(Jmm, A)
        end,
        function ()
            return pr_cov(gp) \ (mean(gp) - μ₀)
        end,
    )
end

function hyperparameter_gradient_function(
    gp::TVarLatent{T}, X::AbstractVector
) where {T<:Real}
    μ₀ = pr_mean(gp, X)
    A =
        (
            I(dim(gp)) .-
            (pr_cov(gp)) \ (cov(gp) .+ (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))
        ) / (pr_cov(gp))
    return (
        function (Jnn)
            return -hyperparameter_KL_gradient(prior(gp).χ * Jnn, A)
        end,
        function ()
            return (pr_cov(gp)) \ (mean(gp) - μ₀)
        end,
    )
end

function hyperparameter_gradient_function(gp::OnlineVarLatent{T}) where {T<:Real}
    μ₀ = pr_mean(gp, gp.Z)
    A =
        (I(dim(gp)) - pr_cov(gp) \ (cov(gp) + (mean(gp) - μ₀) * transpose(mean(gp) - μ₀))) /
        pr_cov(gp)
    κΣ = gp.κ * cov(gp)
    κₐΣ = gp.κₐ * cov(gp)
    return (
        function (Jmm, Jnm, Jnn, Jab, Jaa, ∇E_μ, ∇E_Σ, i, opt)
            ∇E = hyperparameter_expec_gradient(gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, Jnn)
            ∇KLₐ = hyperparameter_online_gradient(gp, κₐΣ, Jmm, Jab, Jaa)
            ∇KL = hyperparameter_KL_gradient(Jmm, A)
            return ∇E + ∇KLₐ - ∇KL
        end, # Function gradient given kernel parameters
        function (Jmm, Jnm, Jab, ∇E_μ, ∇E_Σ, i, opt)
            return hyperparameter_expec_gradient(
                gp, ∇E_μ, ∇E_Σ, i, opt, κΣ, Jmm, Jnm, zero(gp.K̃)
            ) + hyperparameter_online_gradient(
                gp, κₐΣ, Jmm, Jab, zeros(T, length(gp.Zₐ), length(gp.Zₐ))
            ) - hyperparameter_KL_gradient(Jmm, A)
        end, # Function gradient given inducing points locations
        function ()
            return -(pr_cov(gp) \ (μ₀ - mean(gp)))
        end,
    ) # Function gradient given mean prior
end

## Gradient with respect to hyperparameter for analytical VI ##
function hyperparameter_expec_gradient(
    gp::AbstractLatent,
    ∇E_μ::AbstractVector{<:Real},
    ∇E_Σ::AbstractVector{<:Real},
    i::AnalyticVI,
    opt::AVIOptimizer,
    κΣ::AbstractMatrix{<:Real},
    Jmm::AbstractMatrix{<:Real},
    Jnm::AbstractMatrix{<:Real},
    Jnn::AbstractVector{<:Real},
)
    ι = (Jnm - gp.κ * Jmm) / pr_cov(gp)
    J̃ = Jnn - (diag_ABt(ι, gp.Knm) + diag_ABt(gp.κ, Jnm))
    dμ = dot(∇E_μ, ι * mean(gp))
    dΣ = -dot(∇E_Σ, J̃)
    dΣ += -dot(∇E_Σ, 2.0 * (diag_ABt(ι, κΣ)))
    dΣ += -dot(∇E_Σ, 2.0 * (ι * mean(gp)) .* (gp.κ * mean(gp)))
    return ρ(i) * (dμ + dΣ)
end

## Gradient with respect to hyperparameters for numerical VI ##
function hyperparameter_expec_gradient(
    gp::AbstractLatent,
    ∇E_μ::AbstractVector{<:Real},
    ∇E_Σ::AbstractVector{<:Real},
    i::NumericalVI,
    opt::NVIOptimizer,
    κΣ::AbstractMatrix{<:Real},
    Jmm::AbstractMatrix{<:Real},
    Jnm::AbstractMatrix{<:Real},
    Jnn::AbstractVector{<:Real},
)
    ι = (Jnm - gp.κ * Jmm) / pr_cov(gp)
    J̃ = Jnn - (diag_ABt(ι, gp.Knm) + diag_ABt(gp.κ, Jnm))
    dμ = dot(∇E_μ, ι * mean(gp))
    dΣ = dot(∇E_Σ, J̃ + 2.0 * diag_ABt(ι, κΣ))
    return ρ(i) * (dμ + dΣ)
end

function hyperparameter_online_gradient(
    gp::AbstractLatent,
    κₐΣ::Matrix{<:Real},
    Jmm::AbstractMatrix{<:Real},
    Jab::AbstractMatrix{<:Real},
    Jaa::AbstractMatrix{<:Real},
)
    ιₐ = (Jab - gp.κₐ * Jmm) / pr_cov(gp)
    trace_term =
        -sum(
            trace_ABt.(
                [gp.invDₐ],
                [
                    Jaa,
                    2 * ιₐ * transpose(κₐΣ),
                    -ιₐ * transpose(gp.Kab),
                    -gp.κₐ * transpose(Jab),
                ],
            ),
        ) / 2
    term_1 = dot(gp.prevη₁, ιₐ * mean(gp))
    term_2 = -dot(ιₐ * mean(gp), gp.invDₐ * gp.κₐ * mean(gp))
    return trace_term + term_1 + term_2
end


function define_weights_voltages(strategy)
    # Create dictionary of spatial weights for different equipment types / areas
    spatial_weights = Dict{String, Any}()

    # Weights for AC OHL
    spatial_weights["ac_ohl"] =  Dict{String, Any}()
    spatial_weights["ac_ohl"]["nature"] = 2.5
    spatial_weights["ac_ohl"]["urban"] = 40
    spatial_weights["ac_ohl"]["mountain"] = 10
    spatial_weights["ac_ohl"]["sea"] = 40
    spatial_weights["ac_ohl"]["Natura2000"] = 40
    spatial_weights["ac_ohl"]["grid"] = 1
    spatial_weights["ac_ohl"]["agricultural"] = 1
    spatial_weights["ac_ohl"]["roads"] = 1
    spatial_weights["ac_ohl"]["railroads"] = 1

    if strategy == "OHL_on_existing_corridors"
        spatial_weights["ac_ohl"]["agricultural"] = 40
        spatial_weights["ac_ohl"]["nature"] = 40
        spatial_weights["ac_ohl"]["mountain"] = 40
    end

    if strategy == "cables_only"
        spatial_weights["ac_ohl"]["agricultural"] = 40
        spatial_weights["ac_ohl"]["nature"] = 40
        spatial_weights["ac_ohl"]["mountain"] = 40
        spatial_weights["ac_ohl"]["grid"] = 40
        spatial_weights["ac_ohl"]["roads"] = 40
        spatial_weights["ac_ohl"]["railroads"] = 40
    end

    # Weights for AC UGC
    spatial_weights["ac_ugc"] =  Dict{String, Any}()
    spatial_weights["ac_ugc"]["nature"] = 1.5
    spatial_weights["ac_ugc"]["urban"] = 2.5
    spatial_weights["ac_ugc"]["mountain"] = 4.5
    spatial_weights["ac_ugc"]["sea"] = 0.75
    spatial_weights["ac_ugc"]["Natura2000"] = 40
    spatial_weights["ac_ugc"]["grid"] = 1
    spatial_weights["ac_ugc"]["agricultural"] = 1
    spatial_weights["ac_ugc"]["roads"] = 1
    spatial_weights["ac_ugc"]["railroads"] = 1

    # Weights for DC OHL
    spatial_weights["dc_ohl"] =  Dict{String, Any}()
    spatial_weights["dc_ohl"]["nature"] = 2.5
    spatial_weights["dc_ohl"]["urban"] = 40
    spatial_weights["dc_ohl"]["mountain"] = 10
    spatial_weights["dc_ohl"]["sea"] = 40
    spatial_weights["dc_ohl"]["Natura2000"] = 10
    spatial_weights["dc_ohl"]["grid"] = 1
    spatial_weights["dc_ohl"]["agricultural"] = 1
    spatial_weights["dc_ohl"]["roads"] = 1
    spatial_weights["dc_ohl"]["railroads"] = 1

    if strategy == "OHL_on_existing_corridors"
        spatial_weights["dc_ohl"]["agricultural"] = 40
        spatial_weights["dc_ohl"]["nature"] = 40
        spatial_weights["dc_ohl"]["mountain"] = 40
    end

    if strategy == "cables_only"
        spatial_weights["dc_ohl"]["agricultural"] = 40
        spatial_weights["dc_ohl"]["nature"] = 40
        spatial_weights["dc_ohl"]["mountain"] = 40
        spatial_weights["dc_ohl"]["grid"] = 40
        spatial_weights["dc_ohl"]["roads"] = 40
        spatial_weights["dc_ohl"]["railroads"] = 40
    end

    # Weights for AC UGC
    spatial_weights["dc_ugc"] =  Dict{String, Any}()
    spatial_weights["dc_ugc"]["nature"] = 1.5
    spatial_weights["dc_ugc"]["urban"] = 2.5
    spatial_weights["dc_ugc"]["mountain"] = 4.5
    spatial_weights["dc_ugc"]["sea"] = 0.75
    spatial_weights["dc_ugc"]["Natura2000"] = 10
    spatial_weights["dc_ugc"]["grid"] = 1
    spatial_weights["dc_ugc"]["agricultural"] = 1
    spatial_weights["dc_ugc"]["roads"] = 1
    spatial_weights["dc_ugc"]["railroads"] = 1

    # Define transmission voltage levels
    voltages = Dict{String, Any}()
    voltages["ac_ohl"] = 400
    voltages["dc_ohl"] = 400
    voltages["ac_ugc"] = 400
    voltages["dc_ugc"] = 320

    #Minimum resolution
    resolution = Dict{String, Any}()
    resolution["delta_min"] = 4
    resolution["weight_min"] = Dict{String, Any}()
    resolution["weight_min"]["ac_ohl"] = 1
    resolution["weight_min"]["ac_ugc"] = 1
    resolution["weight_min"]["dc_ohl"] = 1
    resolution["weight_min"]["dc_ugc"] = 1

    impedances = Dict{String,Any}()
    impedances["ac"] = Dict{String,Any}()
    impedances["ac"]["ohl"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["16"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["16"]["r"] = 1.874
    impedances["ac"]["ohl"]["16"]["x"] = 0.347
    impedances["ac"]["ohl"]["16"]["c"] = 13.06

    impedances["ac"]["ohl"]["25"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["25"]["r"] = 1.118
    impedances["ac"]["ohl"]["25"]["x"] = 0.334
    impedances["ac"]["ohl"]["25"]["c"] = 13.06

    impedances["ac"]["ohl"]["50"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["50"]["r"] = 0.574
    impedances["ac"]["ohl"]["50"]["x"] = 0.312
    impedances["ac"]["ohl"]["50"]["c"] = 13.06

    impedances["ac"]["ohl"]["70"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["70"]["r"] = 0.404
    impedances["ac"]["ohl"]["70"]["x"] = 0.302
    impedances["ac"]["ohl"]["70"]["c"] = 13.06

    impedances["ac"]["ohl"]["95"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["95"]["r"] = 0.300
    impedances["ac"]["ohl"]["95"]["x"] = 0.291
    impedances["ac"]["ohl"]["95"]["c"] = 13.06

    impedances["ac"]["ohl"]["120"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["120"]["r"] = 0.252
    impedances["ac"]["ohl"]["120"]["x"] = 0.284
    impedances["ac"]["ohl"]["120"]["c"] = 13.06

    impedances["ac"]["ohl"]["150"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["150"]["r"] = 0.190
    impedances["ac"]["ohl"]["150"]["x"] = 0.279
    impedances["ac"]["ohl"]["150"]["c"] = 13.06

    impedances["ac"]["ohl"]["170"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["170"]["r"] = 0.165
    impedances["ac"]["ohl"]["170"]["x"] = 0.276
    impedances["ac"]["ohl"]["170"]["c"] = 13.06

    impedances["ac"]["ohl"]["185"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["185"]["r"] = 0.154
    impedances["ac"]["ohl"]["185"]["x"] = 0.274
    impedances["ac"]["ohl"]["185"]["c"] = 13.06

    impedances["ac"]["ohl"]["210"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["210"]["r"] = 0.133
    impedances["ac"]["ohl"]["210"]["x"] = 0.270
    impedances["ac"]["ohl"]["210"]["c"] = 13.06

    impedances["ac"]["ohl"]["230"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["230"]["r"] = 0.122
    impedances["ac"]["ohl"]["230"]["x"] = 0.268
    impedances["ac"]["ohl"]["230"]["c"] = 13.06

    impedances["ac"]["ohl"]["240"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["240"]["r"] = 0.116
    impedances["ac"]["ohl"]["240"]["x"] = 0.267
    impedances["ac"]["ohl"]["240"]["c"] = 13.06

    impedances["ac"]["ohl"]["265"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["265"]["r"] = 0.107
    impedances["ac"]["ohl"]["265"]["x"] = 0.264
    impedances["ac"]["ohl"]["265"]["c"] = 13.06

    impedances["ac"]["ohl"]["300"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["300"]["r"] = 0.093
    impedances["ac"]["ohl"]["300"]["x"] = 0.261
    impedances["ac"]["ohl"]["300"]["c"] = 13.06

    impedances["ac"]["ohl"]["340"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["340"]["r"] = 0.083
    impedances["ac"]["ohl"]["340"]["x"] = 0.258
    impedances["ac"]["ohl"]["340"]["c"] = 13.06

    impedances["ac"]["ohl"]["380"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["380"]["r"] = 0.074
    impedances["ac"]["ohl"]["380"]["x"] = 0.255
    impedances["ac"]["ohl"]["380"]["c"] = 13.06

    impedances["ac"]["ohl"]["435"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["435"]["r"] = 0.065
    impedances["ac"]["ohl"]["435"]["x"] = 0.252
    impedances["ac"]["ohl"]["435"]["c"] = 13.06

    impedances["ac"]["ohl"]["450"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["450"]["r"] = 0.063
    impedances["ac"]["ohl"]["450"]["x"] = 0.251
    impedances["ac"]["ohl"]["450"]["c"] = 13.06

    impedances["ac"]["ohl"]["490"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["490"]["r"] = 0.058
    impedances["ac"]["ohl"]["490"]["x"] = 0.249
    impedances["ac"]["ohl"]["490"]["c"] = 13.06

    impedances["ac"]["ohl"]["550"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["550"]["r"] = 0.051
    impedances["ac"]["ohl"]["550"]["x"] = 0.246
    impedances["ac"]["ohl"]["550"]["c"] = 13.06

    impedances["ac"]["ohl"]["560"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["560"]["r"] = 0.050
    impedances["ac"]["ohl"]["560"]["x"] = 0.26
    impedances["ac"]["ohl"]["560"]["c"] = 13.06

    impedances["ac"]["ohl"]["680"] = Dict{String,Any}()
    impedances["ac"]["ohl"]["680"]["r"] = 0.042
    impedances["ac"]["ohl"]["680"]["x"] = 0.241
    impedances["ac"]["ohl"]["680"]["c"] = 13.06

    impedances["ac"]["ugc"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["300"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["300"]["r"] = 0.069
    impedances["ac"]["ugc"]["300"]["x"] = 0.110
    impedances["ac"]["ugc"]["300"]["c"] = 375

    impedances["ac"]["ugc"]["400"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["400"]["r"] = 0.057
    impedances["ac"]["ugc"]["400"]["x"] = 0.130
    impedances["ac"]["ugc"]["400"]["c"] = 170

    impedances["ac"]["ugc"]["500"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["500"]["r"] = 0.047
    impedances["ac"]["ugc"]["500"]["x"] = 0.125
    impedances["ac"]["ugc"]["500"]["c"] = 180

    impedances["ac"]["ugc"]["630"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["630"]["r"] = 0.038
    impedances["ac"]["ugc"]["630"]["x"] = 0.122
    impedances["ac"]["ugc"]["630"]["c"] = 200

    impedances["ac"]["ugc"]["800"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["800"]["r"] = 0.027
    impedances["ac"]["ugc"]["800"]["x"] = 0.102
    impedances["ac"]["ugc"]["800"]["c"] = 402

    impedances["ac"]["ugc"]["1200"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["1200"]["r"] = 0.0151
    impedances["ac"]["ugc"]["1200"]["x"] = 0.116
    impedances["ac"]["ugc"]["1200"]["c"] = 210

    impedances["ac"]["ugc"]["1600"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["1600"]["r"] = 0.0113
    impedances["ac"]["ugc"]["1600"]["x"] = 0.111
    impedances["ac"]["ugc"]["1600"]["c"] = 231

    impedances["ac"]["ugc"]["2000"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["2000"]["r"] = 0.010
    impedances["ac"]["ugc"]["2000"]["x"] = 0.1099
    impedances["ac"]["ugc"]["2000"]["c"] = 240

    impedances["ac"]["ugc"]["2500"] = Dict{String,Any}()
    impedances["ac"]["ugc"]["2500"]["r"] = 0.0095
    impedances["ac"]["ugc"]["2500"]["x"] = 0.098
    impedances["ac"]["ugc"]["2500"]["c"] = 260

    impedances["dc"] = Dict{String,Any}()
    impedances["dc"]["ohl"] = Dict{String,Any}()
    for (cross_section, cs) in impedances["ac"]["ohl"]
        impedances["dc"]["ohl"][cross_section] = Dict{String,Any}()
        impedances["dc"]["ohl"][cross_section]["r"] = cs["r"]
        impedances["dc"]["ohl"][cross_section]["x"] = 0
        impedances["dc"]["ohl"][cross_section]["c"] = 0
    end

    impedances["dc"]["ugc"] = Dict{String,Any}()
    impedances["dc"]["ugc"]["95"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["95"]["r"] = 0.247
    impedances["dc"]["ugc"]["95"]["x"] = 0
    impedances["dc"]["ugc"]["95"]["c"] = 0

    impedances["dc"]["ugc"]["120"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["120"]["r"] = 0.196
    impedances["dc"]["ugc"]["120"]["x"] = 0
    impedances["dc"]["ugc"]["120"]["c"] = 0

    impedances["dc"]["ugc"]["150"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["150"]["r"] = 0.16
    impedances["dc"]["ugc"]["150"]["x"] = 0
    impedances["dc"]["ugc"]["150"]["c"] = 0

    impedances["dc"]["ugc"]["185"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["185"]["r"] = 0.128
    impedances["dc"]["ugc"]["185"]["x"] = 0
    impedances["dc"]["ugc"]["185"]["c"] = 0

    impedances["dc"]["ugc"]["240"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["240"]["r"] = 0.097
    impedances["dc"]["ugc"]["240"]["x"] = 0
    impedances["dc"]["ugc"]["240"]["c"] = 0

    impedances["dc"]["ugc"]["300"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["300"]["r"] = 0.08
    impedances["dc"]["ugc"]["300"]["x"] = 0
    impedances["dc"]["ugc"]["300"]["c"] = 0

    impedances["dc"]["ugc"]["400"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["400"]["r"] = 0.062
    impedances["dc"]["ugc"]["400"]["x"] = 0
    impedances["dc"]["ugc"]["400"]["c"] = 0

    impedances["dc"]["ugc"]["500"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["500"]["r"] = 0.0508
    impedances["dc"]["ugc"]["500"]["x"] = 0
    impedances["dc"]["ugc"]["500"]["c"] = 0

    impedances["dc"]["ugc"]["630"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["630"]["r"] = 0.039
    impedances["dc"]["ugc"]["630"]["x"] = 0
    impedances["dc"]["ugc"]["630"]["c"] = 0

    impedances["dc"]["ugc"]["800"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["800"]["r"] = 0.032
    impedances["dc"]["ugc"]["800"]["x"] = 0
    impedances["dc"]["ugc"]["800"]["c"] = 0

    impedances["dc"]["ugc"]["1000"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1000"]["r"] = 0.024
    impedances["dc"]["ugc"]["1000"]["x"] = 0
    impedances["dc"]["ugc"]["1000"]["c"] = 0

    impedances["dc"]["ugc"]["1200"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1200"]["r"] = 0.021
    impedances["dc"]["ugc"]["1200"]["x"] = 0
    impedances["dc"]["ugc"]["1200"]["c"] = 0

    impedances["dc"]["ugc"]["1400"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1400"]["r"] = 0.019
    impedances["dc"]["ugc"]["1400"]["x"] = 0
    impedances["dc"]["ugc"]["1400"]["c"] = 0

    impedances["dc"]["ugc"]["1600"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1600"]["r"] = 0.017
    impedances["dc"]["ugc"]["1600"]["x"] = 0
    impedances["dc"]["ugc"]["1600"]["c"] = 0

    impedances["dc"]["ugc"]["1800"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["1800"]["r"] = 0.0155
    impedances["dc"]["ugc"]["1800"]["x"] = 0
    impedances["dc"]["ugc"]["1800"]["c"] = 0

    impedances["dc"]["ugc"]["2000"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2000"]["r"] = 0.014
    impedances["dc"]["ugc"]["2000"]["x"] = 0
    impedances["dc"]["ugc"]["2000"]["c"] = 0

    impedances["dc"]["ugc"]["2200"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2200"]["r"] = 0.013
    impedances["dc"]["ugc"]["2200"]["x"] = 0
    impedances["dc"]["ugc"]["2200"]["c"] = 0

    impedances["dc"]["ugc"]["2400"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2400"]["r"] = 0.012
    impedances["dc"]["ugc"]["2400"]["x"] = 0
    impedances["dc"]["ugc"]["2400"]["c"] = 0

    impedances["dc"]["ugc"]["2600"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2600"]["r"] = 0.010
    impedances["dc"]["ugc"]["2600"]["x"] = 0
    impedances["dc"]["ugc"]["2600"]["c"] = 0

    impedances["dc"]["ugc"]["2800"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["2800"]["r"] = 0.09
    impedances["dc"]["ugc"]["2800"]["x"] = 0
    impedances["dc"]["ugc"]["2800"]["c"] = 0

    impedances["dc"]["ugc"]["3000"] = Dict{String, Any}()
    impedances["dc"]["ugc"]["3000"]["r"] = 0.075
    impedances["dc"]["ugc"]["3000"]["x"] = 0
    impedances["dc"]["ugc"]["3000"]["c"] = 0

    return spatial_weights, voltages, resolution, impedances
end


using YaoArrayRegister
using YaoArrayRegister: matvec

export ControlBlock, control, cnot, cz

struct ControlBlock{N,BT<:AbstractBlock,C,M} <: AbstractContainer{BT,N}
    ctrl_locs::NTuple{C,Int}
    ctrl_config::NTuple{C,Int}
    content::BT
    locs::NTuple{M,Int}
    function ControlBlock{N,BT,C,M}(
        ctrl_locs,
        ctrl_config,
        block,
        locs,
    ) where {N,C,M,BT<:AbstractBlock}
        @assert_locs_safe N (ctrl_locs..., locs...)
        @assert nqubits(block) == M "number of locations doesn't match the size of block"
        @assert block isa AbstractBlock "expect a block, got $(typeof(block))"
        new{N,BT,C,M}(ctrl_locs, ctrl_config, block, locs)
    end
end

"""
    decode_sign(ctrls...)

Decode signs into control sequence on control or inversed control.
"""
decode_sign(ctrls::Int...) = decode_sign(ctrls)
decode_sign(ctrls::NTuple{N,Int}) where {N} =
    tuple(ctrls .|> abs, ctrls .|> sign .|> (x -> (1 + x) ÷ 2))

function ControlBlock{N}(
    ctrl_locs::NTuple{C},
    ctrl_config::NTuple{C},
    block::BT,
    locs::NTuple{K},
) where {N,M,C,K,BT<:AbstractBlock{M}}
    M == K || throw(DimensionMismatch("block position not maching its size!"))
    return ControlBlock{N,BT,C,M}(ctrl_locs, ctrl_config, block, locs)
end

function ControlBlock{N}(
    ctrl_locs::NTuple{C},
    ctrl_config::NTuple{C},
    block,
    locs::NTuple{K},
) where {N,M,C,K}
    error("expect a block, got $(typeof(block))")
end

# control bit configs are 1 by default, it use sign to encode control bit code
ControlBlock{N}(ctrl_locs::NTuple{C}, block::AbstractBlock, locs::NTuple) where {N,C} =
    ControlBlock{N}(decode_sign(ctrl_locs)..., block, locs)
ControlBlock{N}(ctrl_locs::NTuple{C}, block::Function, locs::NTuple) where {N,C} =
    ControlBlock{N}(decode_sign(ctrl_locs)..., parse_block(length(locs), block), locs)
ControlBlock{N}(ctrl_locs::NTuple{C}, block, locs::NTuple) where {N,C} =
    ControlBlock{N}(decode_sign(ctrl_locs)..., block, locs) # trigger error

# use pair to represent block under control in a compact way
ControlBlock{N}(ctrl_locs::NTuple{C}, target::Pair) where {N,C} =
    ControlBlock{N}(ctrl_locs, target.second, (target.first...,))

"""
    control(n, ctrl_locs, target)

Return a [`ControlBlock`](@ref) with number of active qubits `n` and control locs
`ctrl_locs`, and control target in `Pair`.

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control(4, (1, 2), 3=>X)
nqubits: 4
control(1, 2)
└─ (3,) X

julia> control(4, 1, 3=>X)
nqubits: 4
control(1)
└─ (3,) X
```
"""
control(total::Int, ctrl_locs, target::Pair) = ControlBlock{total}(Tuple(ctrl_locs), target)
control(total::Int, control_location::Int, target::Pair) = control(total, (control_location,), target)

"""
    control(ctrl_locs, target) -> f(n)

Return a lambda that takes the number of total active qubits as input. See also
[`control`](@ref).

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control((2, 3), 1=>X)
(n -> control(n, (2, 3), 1 => X))

julia> control(2, 1=>X)
(n -> control(n, 2, 1 => X))
```
"""
control(ctrl_locs, target::Pair) = @λ(n -> control(n, ctrl_locs, target))
control(control_location::Int, target::Pair) = @λ(n -> control(n, control_location, target))

"""
    control(target) -> f(ctrl_locs)

Return a lambda that takes a `Tuple` of control qubits locs as input. See also
[`control`](@ref).

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control(1=>X)
(ctrl_locs -> control(ctrl_locs, 1 => X))

julia> control((2, 3) => YaoBlocks.ConstGate.CNOT)
(ctrl_locs -> control(ctrl_locs, (2, 3) => CNOT))
```
"""
control(target::Pair) = @λ(ctrl_locs -> control(ctrl_locs, target))

"""
    control(ctrl_locs::Int...) -> f(target)

Return a lambda that takes a `Pair` of control target as input.
See also [`control`](@ref).

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> control(1, 2)
(target -> control((1, 2), target))
```
"""
control(ctrl_locs::Int...) = @λ(target -> control(ctrl_locs, target))

"""
    cnot([n, ]ctrl_locs, location)

Return a speical [`ControlBlock`](@ref), aka CNOT gate with number of active qubits
`n` and locs of control qubits `ctrl_locs`, and `location` of `X` gate.

# Example

```jldoctest; setup=:(using YaoBlocks)
julia> cnot(3, (2, 3), 1)
nqubits: 3
control(2, 3)
└─ (1,) X

julia> cnot(2, 1)
(n -> cnot(n, 2, 1))
```
"""
cnot(total::Int, ctrl_locs, locs::Int) = control(total, ctrl_locs, locs => X)
cnot(ctrl_locs, loc::Int) = @λ(n -> cnot(n, ctrl_locs, loc))

"""
    cz([n, ]ctrl_locs, location)

Return a speical [`ControlBlock`](@ref), aka CZ gate with number of active qubits
`n` and locs of control qubits `ctrl_locs`, and `location` of `Z` gate. See also
[`cnot`](@ref).
"""
cz(total::Int, ctrl_locs, locs::Int) = control(total, ctrl_locs, locs => Z)
cz(ctrl_locs, loc::Int) = @λ(n -> cz(n, ctrl_locs, loc))

mat(::Type{T}, c::ControlBlock{N,BT,C}) where {T,N,BT,C} =
    cunmat(N, c.ctrl_locs, c.ctrl_config, mat(T, c.content), c.locs)

function _apply!(r::AbstractRegister, c::ControlBlock)
    instruct!(r, mat_matchreg(r, c.content), c.locs, c.ctrl_locs, c.ctrl_config)
    return r
end

# specialization
for G in [:X, :Y, :Z, :S, :T, :Sdag, :Tdag]
    GT = Expr(:(.), :ConstGate, QuoteNode(Symbol(G, :Gate)))

    @eval function _apply!(r::AbstractRegister, c::ControlBlock{N,<:$GT}) where {N}
        instruct!(r, Val($(QuoteNode(G))), c.locs, c.ctrl_locs, c.ctrl_config)
        return r
    end
end

PropertyTrait(::ControlBlock) = PreserveAll()

occupied_locs(c::ControlBlock) = (c.ctrl_locs..., map(x -> c.locs[x], occupied_locs(c.content))...)
chsubblocks(pb::ControlBlock{N}, blk::AbstractBlock) where {N} =
    ControlBlock{N}(pb.ctrl_locs, pb.ctrl_config, blk, pb.locs)

# NOTE: ControlBlock will forward parameters directly without loop
cache_key(ctrl::ControlBlock) = cache_key(ctrl.content)

function Base.:(==)(lhs::ControlBlock{N,BT,C,M}, rhs::ControlBlock{N,BT,C,M}) where {BT,N,C,M}
    return (lhs.ctrl_locs == rhs.ctrl_locs) && (lhs.content == rhs.content) && (lhs.locs == rhs.locs)
end

Base.adjoint(blk::ControlBlock{N}) where {N} =
    ControlBlock{N}(blk.ctrl_locs, blk.ctrl_config, adjoint(blk.content), blk.locs)

# NOTE: we only copy one hierachy (shallow copy) for each block
function Base.copy(ctrl::ControlBlock{N,BT,C,M}) where {BT,N,C,M}
    return ControlBlock{N,BT,C,M}(ctrl.ctrl_locs, ctrl.ctrl_config, ctrl.content, ctrl.locs)
end

function YaoBase.iscommute(x::ControlBlock{N}, y::ControlBlock{N}) where {N}
    if x.locs == y.locs
        return iscommute(x.content, y.content)
    elseif !any(l -> l in y.ctrl_locs, x.locs) && !any(l -> l in x.ctrl_locs, y.locs)
        return true
    else
        return iscommute_fallback(x, y)
    end
end


# Sequence-form integrated cyber air defense, G = (N, Σ, g, C)
using Pkg; Pkg.activate(pwd())
using Revise, LinearAlgebra, FileIO, JLD2, ProgressMeter
using JuMP, Clp
using Clp: ClpCInterface

# load functions
dir = pwd()
include(joinpath(dir, "src\\ops_utility.jl"))
include(joinpath(dir, "src\\ops_build.jl"))
include(joinpath(dir, "src\\ops_methods.jl"))

#############
# best response
#############

# fix player 2's realization plan as uniformly random
# how to make better realization plan? Maybe pick a few pure strategies?
# greedy plan: defender assigns cyber assets to IADS with most covered value
function build_rp2(gos::GameOptSet)
    ns2, ns2_stage, na_stage = gos.ns2, gos.ns2_stage, gos.na_stage
    rp2 = zeros(Float64, ns2)
    ns_cumsum = cumsum(ns2_stage)
    p0 = 1.0
    p3 = p0 / na_stage[3]
    p6 = p3 / na_stage[6]
    rp2[1] = p0
    rp2[(ns_cumsum[1] + 1):ns_cumsum[2]] .= p3
    rp2[(ns_cumsum[2] + 1):ns_cumsum[3]] .= p6
    return rp2
end

function build_rp1(gos::GameOptSet)
    ns1, ns1_stage, na_stage = gos.ns1, gos.ns1_stage, gos.na_stage
    rp = zeros(Float64, ns1)
    ns_cumsum = cumsum(ns1_stage)
    p0 = 1.0
    p5 = p0 / na_stage[5]
    rp[1] = p0
    rp[(ns_cumsum[1] + 1):ns_cumsum[2]] .= p5
    return rp
end

function lp_best_response(gos::GameOptSet, rfixed::Vector{Float64}; fixedplayer::Int64 = 2)
    ns1, ns2, ni1, ni2, reward = gos.ns1, gos.ns2, gos.ni1, gos.ni2, gos.reward
    seqI1, seqI2, nextseqI1, nextseqI2 = gos.seqI1, gos.seqI2, gos.nextseqI1, gos.nextseqI2
    ind_nz = findall(!iszero, gos.reward)
    nsi, nso = fixedplayer == 1 ? (ns2, ns1) : (ns1, ns2)
    nii, nio = fixedplayer == 1 ? (ni2, ni1) : (ni1, ni2)
    seqIi, seqIo = fixedplayer == 1 ? (seqI2, seqI1) : (seqI1, seqI2)
    nextseqIi, nextseqIo = fixedplayer == 1 ? (nextseqI2, nextseqI1) : (nextseqI1, nextseqI2)
    println("Building model...")
    m = Model(Clp.Optimizer)
    @variable(m, r[1:nsi] >= 0)
    if fixedplayer == 1 # reward is player 1's reward(seq1, seq2) so we can't switch the indexes
        @objective(m, Max, sum(-1 * gos.reward[I] * rfixed[I[1]] * r[I[2]] for I in ind_nz))
    else
        @objective(m, Max, sum(gos.reward[I] * r[I[1]] * rfixed[I[2]] for I in ind_nz))
    end
    @constraint(m, r[1] == 1.0)
    # @progress for i in 1:nii
    for i in 1:nii
        @constraint(m, sum(r[si] for si in nextseqIi[i]) == r[seqIi[i]])
    end
    println("Solving model...")
    @show m
    status = solve(m)
    return status, getobjectivevalue(m), getvalue(r)
end

#############
# Nash Equilibrium
#############

# TODO: make building the objective function faster in lp_nash
function lp_nash(gos::GameOptSet; timelimit = 5 * 60)
    ns2, ni1, ni2, reward = gos.ns2, gos.ni1, gos.ni2, gos.reward
    ns1_stage, na_stage = gos.ns1_stage, gos.na_stage
    seqI2, nextseqI2 = gos.seqI2, gos.nextseqI2
    println("Building model...")
    mne = Model(solver = ClpSolver(MaximumSeconds = timelimit))
    @variable(mne, rp2[1:ns2] >= 0)
    @variable(mne, v1[1:ni1])
    @variable(mne, v0)
    @objective(mne, Min, v0)
    @constraint(mne, v0 - sum(v1[j] for j in 1:ni1) >= sum(reward[1, s2] * rp2[s2] for s2 in 1:ns2))
    nvconstraint = 1
    # @progress for s1 in (1 + 1):(1 + ns1_stage[2])
    for s1 in (1 + 1):(1 + ns1_stage[2])
        iind = ceil(Int, (s1 - 1) / na_stage[5]) # Σ1 is in blocks na_stage[5] actions for each infoset
        @constraint(mne, v1[iind] >= sum(reward[s1, s2] * rp2[s2] for s2 in 1:ns2))
        nvconstraint += 1
    end
    @constraint(mne, rp2[1] == 1)
    # @progress for i in 1:ni2
    for i in 1:ni2
        @constraint(mne, sum(rp2[s2] for s2 in nextseqI2[i]) == rp2[seqI2[i]])
    end
    @show mne
    println("Solving model...")
    status, t_solve = @timed solve(mne)
    println("obj value: $(getobjectivevalue(mne)), status: $status")
    rp1 = ClpCInterface.dual_row_solution(getrawsolver(mne))[1:nvconstraint] # player 1 strategy is the dual of player 1 values for each sequence constraint
    return status, getobjectivevalue(mne), rp1, getvalue(rp2), t_solve
end

#############
# Robust Best Response
#############

function lp_robust_br(gos::GameOptSet, rp2lb, rp2ub; timelimit = 5 * 60)
    ns2, ni1, ni2, reward = gos.ns2, gos.ni1, gos.ni2, gos.reward
    ns1_stage, na_stage = gos.ns1_stage, gos.na_stage
    seqI2, nextseqI2 = gos.seqI2, gos.nextseqI2
    println("Building model...")
    mr = Model(solver = ClpSolver(MaximumSeconds = timelimit))
    @variable(mr, rp2[1:ns2])
    @variable(mr, v1[1:ni1])
    @variable(mr, v0)
    @objective(mr, Min, v0)
    @constraint(mr, v0 - sum(v1[j] for j in 1:ni1) >= sum(reward[1, s2] * rp2[s2] for s2 in 1:ns2))
    # @progress for s1 in (1 + 1):(1 + ns1_stage[2])
    nvconstraint = 1
    for s1 in (1 + 1):(1 + ns1_stage[2])
        iind = ceil(Int, (s1 - 1) / na_stage[5]) # Σ1 is in blocks na_stage[5] actions for each infoset
        @constraint(mr, v1[iind] >= sum(reward[s1, s2] * rp2[s2] for s2 in 1:ns2))
        nvconstraint += 1
    end
    @constraint(mr, rp2[1] == 1)
    # @progress for i in 1:ni2
    for i in 1:ni2
        @constraint(mr, sum(rp2[s2] for s2 in nextseqI2[i]) == rp2[seqI2[i]])
    end
    @constraint(mr, rp2 .>= rp2lb)
    @constraint(mr, rp2 .<= rp2ub)
    @show mr
    println("Solving model...")
    status, t_solve = @timed solve(mr)
    rp1 = ClpCInterface.dual_row_solution(getrawsolver(mr))[1:nvconstraint] # player 1 strategy is the dual of player 1 values for each sequence constraint
    return status, getobjectivevalue(mr), rp1, getvalue(rp2), t_solve
end


#!/usr/bin/julia

# Sidef concept expressed in Julia

import Base.sqrt
import Base./;

mutable struct S_Number
    value::Number;

    function S_Number(x::Number)
        this = new();
        this.value = x;
        return this;
    end
end

# S_number / S_Number
function /(x::S_Number, y::S_Number)
    S_Number(x.value / y.value);
end

# int(S_Number)
function int(x::S_Number)
    S_Number(round(Int, x.value));
end

# sqrt(S_Number)
function sqrt(x::S_Number)
    S_Number(sqrt(x.value));
end

n1 = S_Number(42);
n2 = S_Number(2);

println(n1);
println(n2);

println(n1 / n2);

method = Symbol("/");
println(eval(:($method(n1, n2))));

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

d = Symbol("/");
x = S_Number(12);
y = S_Number(4);
println("$x / $y == ", eval(:($d(x, y))));

code = "25.3->int->sqrt";

if (m = match(r"\G(\d+(?:\.\d+)?)", code)) != nothing
    str = m.captures[1];
    n = S_Number(parse(Float64, str));
    pos = length(str);

    while (m = match(r"\G(->|\.)", code, pos)) != nothing
        pos += length(m.captures[1]);
        if (m = match(r"\G(\w+)", code, pos)) != nothing
            cap = m.captures[1];
            pos += length(cap);
            method = Symbol(cap);
            print(n, "->", cap);
            n = eval(:($method(n)));
            println(" == ", n);
        end
    end

end

abstract type Animal
end

struct Horse <: Animal
  name::String
end

run(::Animal) = println("Run, animal, run!")
sleep(a::Animal) = println("It's late; ", a.name, " is going to bed...")

function //(z::Horse)
    println("Run horsie, run!")
    println("Horsie's name is ", z.name, ".")
end

foo = Horse("Hashy");
//(foo);

sleep(Horse("Hashy"))