text:
# Plotting functionalities from Matplotlib.

using PyPlot
using REPL.TerminalMenus # Command line UI

export plot, pcolormesh, pcolormeshslice, vdfslice, streamplot, quiver, plotmesh
export pui

@static if matplotlib.__version__ ≥ "3.3"
   matplotlib.rc("image", cmap="turbo") # set default colormap
end

@static if matplotlib.__version__ < "3.5"
   matplotlib.rc("pcolor", shading="nearest") # newer version default "auto"
end

matplotlib.rc("font", size=14)
matplotlib.rc("xtick", labelsize=10)
matplotlib.rc("ytick", labelsize=10)

"""
    plot(meta, var, ax=nothing; comp=0, kwargs)

Plot `var` from `meta` of 1D VLSV data. If `ax===nothing`, plot on the current active axes.
The keyword `comp` in (0, 1, 2, 3) is used to specify the component of a vector.
"""
function PyPlot.plot(meta::MetaVLSV, var, ax=nothing; comp=0, kwargs...)
   ndims(meta) == 1 || error("plot only accepts 1D data!")

   datafull = readvariable(meta, var)

   if comp == 0
      data = datafull
   else
      data = datafull[comp,:]
   end

   x = LinRange(meta.coordmin[1], meta.coordmax[1], meta.ncells[1])

   isnothing(ax) && (ax = plt.gca())

   ax.plot(x, data; kwargs...)
end

"""
    streamplot(meta, var, ax=nothing; comp="xy", axisunit=EARTH, kwargs...)

Wrapper over Matplotlib's streamplot function.
The keyword arguments can be any valid Matplotlib arguments into streamplot.

# Optional arguments
- `comp`: a subset of "xyz" in any order.
- `axisunit`: chosen from `EARTH` and `SI`.
"""
function PyPlot.streamplot(meta::MetaVLSV, var::AbstractString, ax=nothing;
   comp="xy", axisunit=EARTH, kwargs...)

   X, Y, v1, v2 = set_vector(meta, var, comp, axisunit)

   isnothing(ax) && (ax = plt.gca())

   streamplot(X, Y, v1, v2; kwargs...)
end

"""
    quiver(meta, var, ax=nothing; comp="xy", axisunit=EARTH, stride=10, kwargs...)

Wrapper over Matplotlib's quiver function. If `ax===nothing`, plot on the current active
axes.
The keyword arguments can be any valid Matplotlib arguments into quiver.

# Optional arguments
- `comp`: a subset of "xyz" in any order.
- `axisunit`: chosen from `EARTH` and `SI`.
- `stride::Integer`: arrow strides in number of cells.
"""
function PyPlot.quiver(meta::MetaVLSV, var::AbstractString, ax=nothing;
   comp="xy", axisunit::AxisUnit=EARTH, stride::Integer=10, kwargs...)

   X, Y, v1, v2 = set_vector(meta, var, comp, axisunit)

   isnothing(ax) && (ax = plt.gca())

   Xq,  Yq  = X[1:stride:end, 1:stride:end],  Y[1:stride:end, 1:stride:end]
   v1q, v2q = v1[1:stride:end, 1:stride:end], v2[1:stride:end, 1:stride:end]

   ax.quiver(Xq, Yq, v1q, v2q; kwargs...)
end

function set_vector(meta::MetaVLSV, var, comp, axisunit::AxisUnit)
   (;ncells, coordmin, coordmax) = meta
   if occursin("x", comp)
      v1_ = 1
      if occursin("y", comp)
         v2_ = 2
         sizes = (ncells[1], ncells[2])
         plotrange = (coordmin[1], coordmax[1], coordmin[2], coordmax[2])
      else
         v2_ = 3
         sizes = (ncells[1], ncells[3])
         plotrange = (coordmin[1], coordmax[1], coordmin[3], coordmax[3])
      end
   else
      v1_, v2_ = 2, 3
      sizes = (ncells[2], ncells[3])
      plotrange = (coordmin[2], coordmax[2], coordmin[3], coordmax[3])
   end

   data = readvariable(meta, var)

   if !startswith(var, "fg_") # vlasov grid
      @assert ndims(data) == 2 && size(data,1) == 3 "Vector data required!"
      data = reshape(data, 3, sizes[1], sizes[2])
   end

   v1 = data[v1_,:,:]
   v2 = data[v2_,:,:]

   x, y = get_axis(axisunit, plotrange, sizes)

   # meshgrid: note the array ordering difference between Julia and Python!
   X = [i for _ in y, i in x]
   Y = [j for j in y, _ in x]

   X, Y, v1', v2'
end

"""
    pcolormesh(meta::MetaVLSV, var::AbstractString, ax=nothing;
       comp=0, axisunit=EARTH, colorscale=Linear, vmin=-Inf32, vmax=Inf32,
       addcolorbar=true, extent=[-Inf32, Inf32, -Inf32, Inf32], kwargs...)

Plot a variable using pseudocolor from 2D VLSV data.
If `ax` is provided, then it will plot on that axes.
If 3D or AMR grid detected, it will pass arguments to [`pcolormeshslice`](@ref).

# Optional arguments
- `comp::Tuple{Int64, Symbol}`: the component of a vector, chosen from
`:mag, :x, :y, :z, 0, 1, 2, 3`.
- `axisunit::AxisUnit`: the unit of axis ∈ `EARTH, SI`.
- `colorscale::ColorScale`: `Linear`, `Log`, or `SymLog`.
- `vmin::Float`: minimum data range. Set to maximum of data if not specified.
- `vmax::Float`: maximum data range. Set to minimum of data if not specified.
- `addcolorbar::Bool`: whether to add a colorbar to the colormesh.
- `extent::Vector`: extent of axis ranges for plotting in the same unit as `axisunit`.
- `normal::Symbol`: normal direction for slice of 3D data, `:x`, `:y`, `:z`.
- `origin::Float`: origin of plane slice of 3D data.

`pcolormesh(meta, var)`

`pcolormesh(meta, var, axisunit=SI)`

`pcolormesh(data, var, colorscale=Log, extent=[0,1,0,2])`
"""
function PyPlot.pcolormesh(meta::MetaVLSV, var::AbstractString, ax=nothing;
   comp=0, axisunit::AxisUnit=EARTH, colorscale::ColorScale=Linear, addcolorbar=true,
   vmin=-Inf32, vmax=Inf32, extent=[-Inf32, Inf32, -Inf32, Inf32], kwargs...)

   if ndims(meta) == 3 || meta.maxamr > 0
      # check if origin and normal exist in kwargs
      normal = haskey(kwargs, :normal) ? kwargs.data.normal : :y
      origin = haskey(kwargs, :origin) ? kwargs.data.origin : 0.0
      kwargs = Base.structdiff(values(kwargs), (normal = normal, origin = origin))

      pArgs = set_args(meta, var, axisunit; normal, origin)
      data = prep2dslice(meta, var, normal, comp, pArgs)'
   else
      pArgs = set_args(meta, var, axisunit)
      data = prep2d(meta, var, comp)'
   end

   x1, x2 = get_axis(pArgs)

   if var in ("fg_b", "fg_e", "vg_b_vol", "vg_e_vol") || endswith(var, "vg_v")
      _fillinnerBC!(data, data)
   end

   norm, ticks = set_colorbar(colorscale, vmin, vmax, data)

   range1, range2 =
      if axisunit == EARTH
         searchsortedfirst(x1, extent[1]):searchsortedlast(x1, extent[2]),
         searchsortedfirst(x2, extent[3]):searchsortedlast(x2, extent[4])
      else
         searchsortedfirst(x1, extent[1]):searchsortedlast(x1, extent[2]),
         searchsortedfirst(x2, extent[3]):searchsortedlast(x2, extent[4])
      end

   isnothing(ax) && (ax = plt.gca())

   if colorscale != SymLog
      c = ax.pcolormesh(x1[range1], x2[range2], data[range2, range1];
         norm, kwargs...)
   else
      c = ax.pcolormesh(x1[range1], x2[range2], data[range2, range1];
         norm, cmap=matplotlib.cm.RdBu_r, kwargs...)
   end

   set_plot(c, ax, pArgs, ticks, addcolorbar)

   return c
end

"""
    set_colorbar(colorscale::ColorScale, v1, v2, data=[1.0]) -> norm, ticks

Set colorbar norm and ticks in a given range `v1` to `v2` for `data` in `colorscale`.
Matplotlib's Colormap Normalization Section presents various kinds of normlizations beyond
linear, logarithmic and symmetric logarithmic, like centered, discrete, and two slope norm.
For fine-grain control, it is suggested to use the norm methods from `matplotlib.colors`.
For instance,
```
julia> norm = matplotlib.colors.CenteredNorm(); # after matplotlib v3.4
julia> norm = matplotlib.colors.BoundaryNorm(boundaries=[0, 1], ncolors=2);
julia> norm = matplotlib.colors.TwoSlopeNorm(vmin=-500., vcenter=0, vmax=4000);
```
There are also various options for the ticks available in `matplotlib.ticker`.
"""
function set_colorbar(colorscale::ColorScale, v1, v2, data=[1.0f0])
   vmin, vmax = set_lim(Float32(v1), Float32(v2), data, colorscale)
   if colorscale == Linear
      levels = matplotlib.ticker.MaxNLocator(nbins=255).tick_values(vmin, vmax)
      norm = matplotlib.colors.BoundaryNorm(levels, ncolors=256, clip=true)
      ticks = matplotlib.ticker.LinearLocator(numticks=9)
   elseif colorscale == Log # logarithmic
      norm = matplotlib.colors.LogNorm(;vmin, vmax)
      ticks = matplotlib.ticker.LogLocator(base=10, subs=collect(0:9))
   else # symmetric log
      linthresh = 1.0
      logstep = 1

      logthresh = floor(Int, log10(linthresh))
      minlog = ceil(Int, log10(-vmin))
      maxlog = ceil(Int, log10(vmax))

      norm = matplotlib.colors.SymLogNorm(;linthresh, linscale=0.03, vmin, vmax, base=10)
      ticks = [ [-(10.0^x) for x in minlog:-logstep:logthresh]..., 0.0,
         [10.0^x for x in logthresh+1:logstep:maxlog]..., ]
   end

   norm, ticks
end

"Configure customized plot."
function set_plot(c, ax, pArgs::PlotArgs, ticks, addcolorbar)
   (;str_title, strx, stry, cb_title) = pArgs

   if addcolorbar
      cb = colorbar(c; ax, ticks, fraction=0.04, pad=0.02)
      !isempty(cb_title) && cb.ax.set_ylabel(cb_title)
      cb.ax.tick_params(direction="in")
   end

   ax.set_title(str_title, fontweight="bold")
   ax.set_xlabel(strx)
   ax.set_ylabel(stry)
   ax.set_aspect("equal")

   # Set border line widths
   for loc in ("left", "bottom", "right", "top")
      edge = get(ax.spines, loc, nothing)
      edge.set_linewidth(2.0)
   end

   ax.xaxis.set_tick_params(width=2.0, length=3)
   ax.yaxis.set_tick_params(width=2.0, length=3)
   return
end

"""
    vdfslice(meta, location, ax=nothing; kwargs...)

Plot the 2D slice cut of phase space distribution function at `location` within velocity
range `limits`. If `ax===nothing`, plot on the current active axes.
# Optional arguments
- `unit::AxisUnit`: location unit in `SI`, `EARTH`.
- `unitv::String`: velocity unit in ("km/s", "m/s").
- `limits::Vector{Real}`: velocity space range given in [xmin, xmax, ymin, ymax].
- `slicetype`: symbol for choosing the slice type from `:xy`, `:xz`, `:yz`,
`:bperp`, `:bpar1`, `:bpar2`.
- `center`: symbol for setting the reference frame from `:bulk`, `:peak`.
- `vslicethick`: setting the velocity space slice thickness in the normal direction. If set
to 0, the whole distribution along the normal direction is projected onto a plane. Currently
this is only meaningful when `center` is set such that a range near the bulk/peak normal
velocity is selected!
- `vmin, vmax`: minimum and maximum VDF values for plotting.
- `weight::Symbol`: choosing distribution weights from phase space density or particle flux
between `:particle` and `:flux`.
- `flimit`: minimum VDF threshold for plotting.
- `kwargs...`: any valid keyword argument for hist2d.
"""
function vdfslice(meta::MetaVLSV, location, ax=nothing;
   limits=[-Inf32, Inf32, -Inf32, Inf32], verbose=false, species="proton",
   vmin=-Inf32, vmax=Inf32, unit=SI, unitv="km/s", slicetype=:default, vslicethick=0.0,
   center=:nothing, weight=:particle, flimit=-1.0,
   kwargs...)

   v1, v2, r1, r2, weights, strx, stry, str_title =
      prep_vdf(meta, location;
         species, unit, unitv, slicetype, vslicethick, center, weight, flimit, verbose)

   isinf(vmin) && (vmin = minimum(weights))
   isinf(vmax) && (vmax = maximum(weights))

   verbose && @info "Active f range is $vmin, $vmax"

   isnothing(ax) && (ax = plt.gca())

   norm = matplotlib.colors.LogNorm(vmin, vmax)

   h = ax.hist2d(v1, v2; bins=(r1, r2), weights, norm, shading="flat")

   ax.set_title(str_title, fontweight="bold")
   ax.set_xlabel(strx)
   ax.set_ylabel(stry)
   ax.set_aspect("equal")
   ax.grid(color="grey", linestyle="-")
   ax.tick_params(direction="in")

   cb = colorbar(h[4]; ax, fraction=0.04, pad=0.02)
   cb.ax.tick_params(which="both", direction="in")
   cb_title = cb.ax.set_ylabel("f(v)")

   if slicetype in (:xy, :xz, :yz)
      # Draw vector of magnetic field direction
   end

   h[4] # h[1] is 2D data, h[2] is x axis, h[3] is y axis
end

"""
    plotmesh(meta; projection="3d", origin=0.0, marker="+", kwargs...)

Plot mesh cell centers from axis view `projection`. `projection` should be either "3d", "x",
"y" or "z". `origin` is center of projection plane in the normal direction.
"""
function plotmesh(meta::MetaVLSV, ax=nothing; projection="3d", origin=0.0, marker="+",
   kwargs...)
   (;coordmin, coordmax, cellid) = meta
   if projection in ("x", "y", "z")
      dirp, dir1, dir2 =
         if projection == "x"
            1, 2, 3
         elseif projection == "y"
            2, 1, 3
         else # "z"
            3, 1, 2
         end
      sliceoffset = origin - coordmin[dirp]
      ids, _ = getslicecell(meta, sliceoffset, dirp, coordmin[dirp], coordmax[dirp])
   else # 3D
      ids = cellid
   end

   centers = [zeros(SVector{3, Float32}) for _ in ids]
   for (i, id) in enumerate(ids)
      @inbounds centers[i] = getcellcoordinates(meta, id)
   end

   isnothing(ax) && (ax = plt.gca())

   if projection in ("x", "y", "z")
      x1 = getindex.(centers, dir1)
      x2 = getindex.(centers, dir2)
      s = ax.scatter(x1, x2; marker, kwargs...)
   else
      x1 = getindex.(centers, 1)
      x2 = getindex.(centers, 2)
      x3 = getindex.(centers, 3)
      s = ax.scatter(x1, x2, x3; marker, kwargs...)
   end
   s
end

"""
    pui(meta::MetaVLSV)

Quick plotting via command line interactive selections.
"""
function pui(meta::MetaVLSV; suppress_output=false)

   menu = RadioMenu(meta.variable; charset=:ascii)

   println("Choose variable to plot:")
   var_ = request(menu; suppress_output)

   var_ == -1 && error("Variable selection canceled.")

   comp = 0

   if meta.variable[var_] in ("fg_b", "fg_e", "vg_b_vol", "vg_e_vol") ||
      endswith(meta.variable[var_], "vg_v") ||
      occursin("vg_ptensor", meta.variable[var_])

      menu = RadioMenu(["x","y","z","mag"]; charset=:ascii)

      comp = request("Choose vector component to plot:", menu; suppress_output)

      comp == -1 && error("Vector component selection canceled.")

      comp == 4 && (comp = 0)
   end

   if ndims(meta) != 1
      pcolormesh(meta, meta.variable[var_]; comp)
   else # 1D
      plot(meta, meta.variable[var_]; comp)
   end

   return
end

"Direct file plotting."
pui(file::AbstractString) = file |> load |> plot_ui


text:
using ToyRockets
using Test

@testset "ToyRockets.jl" begin
    # Write your tests here.
end

ext: jl


text:
# this doesn't need to be exact, used for setting the dict size
const N_OPCODES = 256

instruction(c::CPU, m::Memory) = deref(Π(c.PC), m)
instruction(cs::Chipset) = instruction(cs.cpu, cs.ram)

# this doesn't do any checking
bytes2arg(bytes::AbstractVector{UInt8}) = bincat(bytes[1], bytes[2])

# this gets the arguments to the op given the pointer to the instruction and nbytes total
deref_opargs(cs::Chipset, ptr::Π, nbytes::Integer) = deref(ptr+0x01, cs.ram, nbytes-1)

# this has been tested and so far yields inferior performance
function derefview_opargs(cs::Chipset, ptr::Π, nbytes::Integer)
    derefview(ptr+0x01, cs.ram, nbytes-1)
end

# note that now we very much depend on the instruction pointer being set correctly
# at the time the instructions execute
@inline opargs(cs::Chipset, n::Integer) = cs.ram[cs.cpu.PC+n]

#==============================================================================================
    <execution>
==============================================================================================#

# TODO figure out how ops handle instruction pointer, think ok now?

# this doesn't do the sleeping
function op!(cs::Chipset)
    exe!, nbytes, ncycles = OPCODES[deref(Π(cs.cpu.PC), cs.ram)+1]
    exe!(cs)
    cs.cpu.PC += nbytes
    cs.clock += ncycles
    nbytes, ncycles
end

# TODO eventually this should employ timing
function run!(cs::Chipset)
    while !cs.cpu.flags.B
        op!(cs)
        println(cs)
    end
end

#==============================================================================================
    </execution>
==============================================================================================#

#==============================================================================================
    <opcodes>
==============================================================================================#
# this is an (efficient) function pointer type for the instructions
# note this works even though ops take AbstractVector
const OpFunc = FunctionWrapper{UInt8,Tuple{Chipset}}

# for the time being opcode 0x00 would throw an error
const OPCODES = Vector{Tuple{OpFunc,Int,Int}}(undef, N_OPCODES)

# this is for the assembler
# TODO NEEDS TO SEE MODES!!! FIX!
const ASSEMBLY_DICT = Dict{Tuple{String,Symbol},UInt8}(); sizehint!(ASSEMBLY_DICT, N_OPCODES)

# format is
#    mode, opcode, nbytes, ncycles, (increment)
#
# the last arg is optional and tells it whether or not to increment the instruction pointer

# we treat "Accumulator" mode as "Implicit"

@opdef lda! begin
    Immediate, 0xa9, 2, 2
    ZeroPage,  0xa5, 2, 3
    ZeroPageX, 0xb5, 2, 4
    Absolute,  0xad, 3, 4
    AbsoluteX, 0xbd, 3, 4  # TODO should automatically take care of page crossings
    AbsoluteY, 0xb9, 3, 4
    IndirectX, 0xa1, 2, 6
    IndirectY, 0xb1, 2, 5
end

@opdef ldx! begin
    Immediate, 0xa2, 2, 2
    ZeroPage,  0xa6, 2, 3
    ZeroPageY, 0xb6, 2, 4
    Absolute,  0xae, 3, 4
    AbsoluteY, 0xbe, 3, 4
end

@opdef ldy! begin
    Immediate, 0xa0, 2, 2
    ZeroPage,  0xa4, 2, 3
    ZeroPageX, 0xb4, 2, 4
    Absolute,  0xac, 3, 4
    AbsoluteX, 0xbc, 3, 4
end

@opdef sta! begin
    ZeroPage,  0x85, 2, 3
    ZeroPageX, 0x95, 2, 4
    Absolute,  0x8d, 3, 4
    AbsoluteX, 0x9d, 3, 5
    AbsoluteY, 0x99, 3, 5
    IndirectX, 0x81, 2, 6
    IndirectX, 0x91, 2, 6
end

@opdef stx! begin
    ZeroPage,  0x86, 2, 3
    ZeroPageY, 0x96, 2, 4
    Absolute,  0x8e, 3, 4
end

@opdef sty! begin
    ZeroPage,  0x84, 2, 3
    ZeroPageX, 0x94, 2, 4
    Absolute,  0x8c, 2, 4
end

@opdef tax! begin
    Implicit,  0xaa, 1, 2
end

@opdef tay! begin
    Implicit,  0xa8, 1, 2
end

@opdef txa! begin
    Implicit,  0x8a, 1, 2
end

@opdef tya! begin
    Implicit,  0x98, 1, 2
end

@opdef tsx! begin
    Implicit,  0xba, 1, 2
end

@opdef txs! begin
    Implicit,  0x9a, 1, 2
end

@opdef pha! begin
    Implicit,  0x48, 1, 2
end

@opdef php! begin
    Implicit,  0x08, 1, 3
end

@opdef pla! begin
    Implicit,  0x68, 1, 4
end

@opdef plp! begin
    Implicit,  0x28, 1, 4
end

@opdef and! begin
    Immediate, 0x29, 2, 2
    ZeroPage,  0x25, 2, 3
    ZeroPageX, 0x35, 2, 4
    Absolute,  0x2d, 3, 4
    AbsoluteX, 0x3d, 3, 4
    AbsoluteY, 0x39, 3, 4
    IndirectX, 0x21, 2, 6
    IndirectY, 0x31, 2, 5
end

@opdef eor! begin
    Immediate, 0x49, 2, 2
    ZeroPage,  0x45, 2, 3
    ZeroPageX, 0x55, 2, 4
    Absolute,  0x4d, 3, 4
    AbsoluteX, 0x5d, 3, 4
    AbsoluteY, 0x59, 3, 4
    IndirectX, 0x41, 2, 6
    IndirectY, 0x51, 2, 5
end

@opdef ora! begin
    Immediate, 0x09, 2, 2
    ZeroPage,  0x05, 2, 3
    ZeroPageX, 0x15, 2, 4
    Absolute,  0x0d, 3, 4
    AbsoluteX, 0x1d, 3, 4
    AbsoluteY, 0x19, 3, 4
    IndirectX, 0x01, 2, 6
    IndirectY, 0x11, 2, 5
end

@opdef bit! begin
    ZeroPage,  0x24, 2, 3
    Absolute,  0x2c, 3, 4
end

@opdef adc! begin
    Immediate, 0x69, 2, 2
    ZeroPage,  0x65, 2, 3
    ZeroPageX, 0x75, 2, 4
    Absolute,  0x6d, 3, 4
    AbsoluteX, 0x7d, 3, 4
    AbsoluteY, 0x79, 3, 4
    IndirectX, 0x61, 2, 6
    IndirectY, 0x71, 2, 5
end

@opdef sbc! begin
    Immediate, 0xe9, 2, 2
    ZeroPage,  0xe5, 2, 3
    ZeroPageX, 0xf5, 2, 4
    Absolute,  0xed, 3, 4
    AbsoluteX, 0xfd, 3, 4
    AbsoluteY, 0xf9, 3, 4
    IndirectX, 0xe1, 2, 6
    IndirectY, 0xf1, 2, 5
end

@opdef cmp! begin
    Immediate, 0xc9, 2, 2
    ZeroPage,  0xc5, 2, 3
    ZeroPageX, 0xd5, 2, 4
    Absolute,  0xcd, 3, 4
    AbsoluteX, 0xdd, 3, 4
    AbsoluteY, 0xd9, 3, 4
    IndirectX, 0xc1, 2, 6
    IndirectY, 0xd1, 2, 5
end

@opdef cpx! begin
    Immediate, 0xe0, 2, 2
    ZeroPage,  0xe4, 2, 3
    Absolute,  0xec, 3, 4
end

@opdef cpy! begin
    Immediate, 0xc0, 2, 2
    ZeroPage,  0xc4, 2, 3
    Absolute,  0xcc, 3, 4
end

@opdef inc! begin
    ZeroPage,  0xe6, 2, 5
    ZeroPageX, 0xf6, 2, 6
    Absolute,  0xee, 3, 6
    AbsoluteX, 0xfe, 3, 7
end

@opdef inx! begin
    Implicit,  0xe8, 1, 2
end

@opdef iny! begin
    Implicit,  0xc8, 1, 2
end

@opdef dec! begin
    ZeroPage,  0xc6, 2, 5
    ZeroPageX, 0xd6, 2, 6
    Absolute,  0xce, 3, 6
    AbsoluteX, 0xde, 3, 7
end

@opdef dex! begin
    Implicit,  0xca, 1, 2
end

@opdef dey! begin
    Implicit,  0x88, 1, 2
end

@opdef asl! begin
    Implicit,  0x0a, 1, 2
    ZeroPage,  0x06, 2, 5
    ZeroPageX, 0x16, 2, 6
    Absolute,  0x0e, 3, 6
    AbsoluteX, 0x1e, 3, 7
end

@opdef lsr! begin
    Implicit,  0x4a, 1, 2
    ZeroPage,  0x46, 2, 5
    ZeroPageX, 0x56, 2, 6
    Absolute,  0x4e, 3, 6
    AbsoluteX, 0x5e, 3, 7
end

@opdef rol! begin
    Implicit,  0x2a, 1, 2
    ZeroPage,  0x26, 2, 5
    ZeroPageX, 0x36, 2, 6
    Absolute,  0x2e, 3, 6
    AbsoluteX, 0x3e, 3, 7
end

@opdef ror! begin
    Implicit,  0x6a, 1, 2
    ZeroPage,  0x66, 2, 5
    ZeroPageX, 0x76, 2, 6
    Absolute,  0x6e, 3, 6
    AbsoluteX, 0x7e, 3, 7
end

# note that these get 0's because they don't increment the instruction pointer
@opdef jmp! begin
    Immediate, 0x4c, 0, 3
    Absolute,  0x6c, 0, 5
end

@opdef jsr! begin
    Immediate, 0x20, 0, 6
end

# TODO not sure if implicit works with cpu and ram arguments
@opdef rts! begin
    Implicit,  0x60, 0, 6
end

# TODO check if instruction pointer is handled correctly in branches
@opdef bcc! begin
    Immediate, 0x90, 2, 2
end

@opdef bcs! begin
    Immediate, 0xb0, 2, 2
end

@opdef beq! begin
    Immediate, 0xf0, 2, 2
end

@opdef bmi! begin
    Immediate, 0x30, 2, 2
end

@opdef bne! begin
    Immediate, 0xd0, 2, 2
end

@opdef bpl! begin
    Immediate, 0x10, 2, 2
end

@opdef bvc! begin
    Immediate, 0x50, 2, 2
end

@opdef bvs! begin
    Immediate, 0x70, 2, 2
end

@opdef clc! begin
    Implicit,  0x18, 1, 2
end

@opdef cld! begin
    Implicit,  0xd8, 1, 2
end

@opdef cli! begin
    Implicit,  0x58, 1, 2
end

@opdef clv! begin
    Implicit,  0xb8, 1, 2
end

@opdef sec! begin
    Implicit,  0x38, 1, 2
end

@opdef sed! begin
    Implicit,  0xf8, 1, 2
end

@opdef sei! begin
    Implicit,  0x78, 1, 2
end

@opdef brk! begin
    Implicit,  0x00, 1, 7
end

@opdef nop! begin
    Implicit,  0xea, 1, 2
end

@opdef rti! begin
    Implicit,  0x40, 1, 6
end
#==============================================================================================
    </opcodes>
==============================================================================================#


text:
CONTEXTS = (Training(), Validation(), Inference())

"""
    checkmethod(method, sample, model; device = identity)
    checkmethod(method; device = identity)

Check if `method` conforms to the `DLPipelines.jl` interfaces.
`sample` and `model` are used for testing. If you have implemented the testing
interface and don't supply these as arguments, `mocksample(method)` and
`mockmodel(method)` will be used.

Checks *core* and *interpretation* interfaces.

"""
function checkmethod(method, sample, model; kwargs...)
    checkmethod_core(method, sample, model; kwargs...)
    checkmethod_interpretation(method, sample, model; kwargs...)
end

"""
    checkmethod_core(method, sample, model; device = identity)
    checkmethod_core(method; device = identity)

Check if `method` conforms to the [core interface](docs/interfaces/core.md).
`sample` and `model` are used for testing. If you have implemented the testing
interface and don't supply these as arguments, `mocksample(method)` and
`mockmodel(method)` will be used.
"""
function checkmethod_core(
        method;
        model = mockmodel(method),
        sample = mocksample(method),
        devicefn = identity)
    @testset "Core interface" begin
        @testset "`encode`" begin
            for context in CONTEXTS
                @test_nowarn encode(method, context, sample)
            end
        end
        @testset "Model compatibility" begin
            for context in CONTEXTS
                x, y = encode(method, context, sample)
                @test_nowarn ŷ = _predictx(method, model, x, devicefn)
            end
        end
        @testset "`decodeŷ" begin
            for context in CONTEXTS
                x, y = encode(method, context, sample)
                ŷ = _predictx(method, model, x, devicefn)
                @test_nowarn decodeŷ(method, context, ŷ)
            end
        end
    end
end

"""
    checkmethod_interpretation(method, sample, model; device = identity)
    checkmethod_interpretation(method; device = identity)

Check if `method` conforms to the [core interface](docs/interfaces/core.md).
`sample` and `model` are used for testing. If you have implemented the testing
interface and don't supply these as arguments, `mocksample(method)` and
`mockmodel(method)` will be used.
"""
function checkmethod_interpretation(
        method;
        model = mockmodel(method),
        sample = mocksample(method),
        devicefn = identity)

    @testset "Interpretation interface" begin
        context = Training()
        x, y = encode(method, context, sample)
        ŷ = _predictx(method, model, x, devicefn)
        target = decodeŷ(method, context, ŷ)
        @testset "`interpretsample`" begin
            @test_nowarn interpretsample(method, sample)
        end
        @testset "`interpretx`" begin
            @test_nowarn interpretx(method, x)
        end
        @testset "`interprety`" begin
            @test_nowarn interprety(method, y)
        end
        @testset "`interpretŷ`" begin
            @test_nowarn interpretŷ(method, ŷ)
        end
        @testset "`interprettarget`" begin
            @test_nowarn interprettarget(method, target)
        end
    end

end


text:
module PatientBlockResults

using DataFrames
using LinearAlgebra
using Dates
using Distributions
using Memoize

function results_all(
		sent::Array{<:Real,4},
		beds::Array{<:Real,2},
		initial::Array{<:Real,2},
		discharged::Array{<:Real,3},
		admitted::Array{<:Real,3},
		los::Array{<:Any,1},
		adj_matrix::BitArray{2},
		transfer_graph::BitArray{2},
		bed_types::Array{Int,1},
		locations::Array{String,1},
		start_date::Date,
)
	B = size(beds,1)
	G, N, T = size(admitted)

	active = results_active_patients(sent, beds, initial, discharged, admitted, los, adj_matrix, transfer_graph, bed_types)

	summary = results_summary(active, sent, locations)
	complete = results_complete(active, sent, admitted, beds, locations, start_date)

	sent_matrix_table = results_sentmatrix_table(sent, locations)
	sent_matrix_vis = results_sentmatrix_vis(sent, initial, admitted, locations)

	sent_to = Dict(locations[i] => locations[row] for (i,row) in enumerate(eachrow(sum(sent, dims=[1,4])[1,:,:,1] .> 0)))

	return (
		summary_table=summary,
		complete_table=complete,
		sent_matrix_table=sent_matrix_table,
		sent_matrix_vis=sent_matrix_vis,
		sent_to=sent_to,
		total_overflow=sum(active.overflow_total),
		total_overflow_nosent=sum(active.overflow_null_total),
		active...,
	)
end

function results_summary(
		active,
		sent::Array{<:Real,4},
		locations::Array{String,1},
)
	summary = DataFrame(
		location=locations,
		total_sent=sum(sent, dims=[1,3,4])[:],
		total_received=sum(sent, dims=[1,2,4])[:],
		overflow=sum(active.overflow_total, dims=2)[:],
	)
	return summary
end

function results_complete(
		active,
		sent::Array{<:Real,4},
		admitted::Array{<:Real,3},
		beds::Array{<:Real,2},
		locations::Array{String,1},
		start_date::Date,
)
	G, N, _, T = size(sent)
	B = size(beds, 1)
	sent_total = sum(sent, dims=1)[1,:,:,:]

	outcomes = DataFrame()
	for (i,s) in enumerate(locations)
		single_state_outcome = DataFrame(
			location=fill(s, T),
			date=start_date .+ Day.(0:T-1),

			sent_total=sum(sent_total[i,:,:], dims=1)[:],
			received_total=sum(sent_total[:,i,:], dims=1)[:],
			new_patients_total=sum(admitted[:,i,:], dims=1)[:],
			capacity_total=fill(sum(beds[:,i]), T),
			active_patients_total=sum(active.active_bybedtype[:,i,:], dims=1)[:],
			active_patients_nosent_total=sum(active.active_null_bybedtype[:,i,:], dims=1)[:],
			overflow_total=sum(active.overflow_bybedtype[:,i,:], dims=1)[:],
			overflow_nosent_total=sum(active.overflow_null_bybedtype[:,i,:], dims=1)[:],
			load_total=active.load_total[i,:],
			load_nosent_total=active.load_null_total[i,:],

			sent_icu=sum(sent[1,i,:,:], dims=1)[:],
			received_icu=sum(sent[1,:,i,:], dims=1)[:],
			new_patients_icu=admitted[1,i,:],
			capacity_icu=fill(beds[1,i], T),
			active_patients_icu=active.active_bybedtype[1,i,:],
			active_patients_nosent_icu=active.active_null_bybedtype[1,i,:],
			overflow_icu=active.overflow_bybedtype[1,i,:],
			overflow_nosent_icu=active.overflow_null_bybedtype[1,i,:],
			load_icu=active.load_bybedtype[1,i,:],
			load_nosent_icu=active.load_null_bybedtype[1,i,:],

			sent_regular=sum(sent[2,i,:,:], dims=1)[:],
			received_regular=sum(sent[2,:,i,:], dims=1)[:],
			new_patients_regular=admitted[1,i,:],
			capacity_regular=fill(beds[2,i], T),
			active_patients_regular=active.active_bybedtype[2,i,:],
			active_patients_nosent_regular=active.active_null_bybedtype[2,i,:],
			overflow_regular=active.overflow_bybedtype[2,i,:],
			overflow_nosent_regular=active.overflow_null_bybedtype[2,i,:],
			load_regular=active.load_bybedtype[2,i,:],
			load_nosent_regular=active.load_null_bybedtype[2,i,:],

			sent_to=[collect(zip(locations[sent_total[i,:,t] .> 0], sent_total[i,sent_total[i,:,t].>0,t])) for t in 1:T],
			sent_from=[collect(zip(locations[sent_total[:,i,t] .> 0], sent_total[sent_total[:,i,t].>0,i,t])) for t in 1:T],
		)
		outcomes = vcat(outcomes, single_state_outcome)
	end

	return outcomes
end

function results_sentmatrix_table(sent::Array{<:Real,4}, locations::Array{String,1})
	sent_matrix = DataFrame(sum(sent, dims=[1,4])[1,:,:,1])
	rename!(sent_matrix, Symbol.(locations))
	insertcols!(sent_matrix, 1, :state => locations)
	return sent_matrix
end

function results_sentmatrix_vis(
		sent::Array{<:Real,4},
		initial_patients::Array{<:Real,2},
		admitted_patients::Array{<:Real,3},
		locations::Array{String,1}
)
	sent_total = sum(sent, dims=1)[1,:,:,:]
	selfedges = sum(initial_patients, dims=1)[:] + sum(admitted_patients, dims=[1,3])[:] - sum(sent_total, dims=[2,3])[:]
	sent_vis_matrix = sum(sent_total, dims=3)[:,:,1] + diagm(selfedges)
	sent_vis_matrix = DataFrame(sent_vis_matrix)
	rename!(sent_vis_matrix, Symbol.(locations))
	return sent_vis_matrix
end

function results_active_patients(
		sent::Array{<:Real,4},
		beds::Array{<:Real,2},
		initial::Array{<:Real,2},
		discharged::Array{<:Real,3},
		admitted::Array{<:Real,3},
		los::Array{<:Any,1},
		adj_matrix::BitArray{2},
		transfer_graph::BitArray{2},
		bed_types::Array{Int,1},
)
	B = size(beds,1)
	G, N, T = size(admitted)

	F = zeros(Float64, G, T)
	f = zeros(Float64, G, T)
	for g in 1:G
		if isa(los[g], Int)
			F[g,:] = vcat(ones(Int, los[g]), zeros(Int, T-los[g]))
			f[g,los[g]+1] = 1.0
		elseif isa(los[g], Distribution)
			F[g,:] = 1.0 .- cdf.(los[g], 0:T-1)
			f[g,:] = pdf.(los[g], 0:T-1)
		else
			error("Invalid length of stay distribution")
		end
		f[g,:] = f[g,:] / sum(f[g,:])
	end

	groups_bybedtype = [sort(findall(x -> x == b, bed_types)) for b in 1:B]

	@memoize q(g,i,t) = (
		admitted[g,i,t] + tr(g,i,t) + sum(sent[g,j,i,t] - sent[g,i,j,t] for j in 1:N)
	)
	@memoize l(g,i,t) = (
		discharged[g,i,t] + sum(f[g,t-t₁+1] * q(g,i,t₁) for t₁ in 1:t)
	)
	@memoize tr(g,i,t) = (
		any(transfer_graph[:,g]) ?
			sum(l(g₁,i,t) for g₁ in findall(transfer_graph[:,g]))
			: 0
	)
	active_bygroup = [(
		initial[g,i] + sum(sent[g,i,:,t]) + sum((q(g,i,t₁) - l(g,i,t₁)) for t₁ in 1:t)
	) for g in 1:G, i in 1:N, t in 1:T]

	@memoize q_null(g,i,t) = (
		admitted[g,i,t] + tr_null(g,i,t)
	)
	@memoize l_null(g,i,t) = (
		discharged[g,i,t] + sum(f[g,t-t₁+1] * q_null(g,i,t₁) for t₁ in 1:t)
	)
	@memoize tr_null(g,i,t) = (
		any(transfer_graph[:,g]) ?
			sum(l_null(g₁,i,t) for g₁ in findall(transfer_graph[:,g]))
			: 0
	)
	active_null_bygroup = [(
		initial[g,i] + sum((q_null(g,i,t₁) - l_null(g,i,t₁)) for t₁ in 1:t)
	) for g in 1:G, i in 1:N, t in 1:T]

	active_bybedtype = [sum(active_bygroup[g,i,t] for g in groups_bybedtype[b]) for b in 1:B, i in 1:N, t in 1:T]
	active_null_bybedtype = [sum(active_null_bygroup[g,i,t] for g in groups_bybedtype[b]) for b in 1:B, i in 1:N, t in 1:T]

	active_total = [sum(active_bygroup[:,i,t]) for i in 1:N, t in 1:T]
	active_null_total = [sum(active_null_bygroup[:,i,t]) for i in 1:N, t in 1:T]

	overflow_bybedtype = [max(0, active_bybedtype[b,i,t] - beds[b,i]) for b in 1:B, i in 1:N, t in 1:T]
	overflow_total = [sum(overflow_bybedtype[:,i,t]) for i in 1:N, t in 1:T]
	load_bybedtype = [active_bybedtype[b,i,t] / beds[b,i] for b in 1:B, i in 1:N, t in 1:T]
	load_total = [sum(active_bybedtype[:,i,t]) / sum(beds[:,i]) for i in 1:N, t in 1:T]

	overflow_null_bybedtype = [max(0, active_null_bybedtype[b,i,t] - beds[b,i]) for b in 1:B, i in 1:N, t in 1:T]
	overflow_null_total = [sum(overflow_null_bybedtype[:,i,t]) for i in 1:N, t in 1:T]
	load_null_bybedtype = [active_null_bybedtype[b,i,t] / beds[b,i] for b in 1:B, i in 1:N, t in 1:T]
	load_null_total = [sum(active_null_bybedtype[:,i,t]) / sum(beds[:,i]) for i in 1:N, t in 1:T]

	return (
		active_bygroup          = active_bygroup,
		active_null_bygroup     = active_null_bygroup,
		active_bybedtype        = active_bybedtype,
		active_null_bybedtype   = active_null_bybedtype,
		active_total            = active_total,
		active_null_total       = active_null_total,
		overflow_bybedtype      = overflow_bybedtype,
		overflow_null_bybedtype = overflow_null_bybedtype,
		overflow_total          = overflow_total,
		overflow_null_total     = overflow_null_total,
		load_bybedtype          = load_bybedtype,
		load_null_bybedtype     = load_null_bybedtype,
		load_total              = load_total,
		load_null_total         = load_null_total,
	)
end

end

max_stars_repo_name: flixpar/covid-resource-allocation

max_issues_repo_name: flixpar/covid-resource-allocation

max_forks_repo_name: flixpar/covid-resource-allocation


text:
using ML_Hopf
using LinearAlgebra
using Statistics
using MAT
using DiffEqFlux
using PGFPlotsX
using LaTeXStrings
using DifferentialEquations

#= Use to format code in BlueStyle
using JuliaFormatter
format("./src/phys_ex.jl", BlueStyle())
format("./src/num_ex.jl", BlueStyle())
format("./src/vdp_ex.jl", BlueStyle())
format("./src/ML_Hopf.jl", BlueStyle())
=#

## Save data
nh = 30
l = 6000
vars = matread("./src/measured_data/CBC_stable_v14_9.mat")
uu = get(vars, "data", 1)
ind1 = 1;
ind2 = 4;
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
mu1 = mean(uu[ind1, 1:l])
mu2 = mean(uu[ind2, 1:l])
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = [[transpose(uu1); transpose(uu2)]]
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = c

vars = matread("./src/measured_data/CBC_stable_v15_6.mat")
uu = get(vars, "data", 1)
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = vcat(t_series, [[transpose(uu1); transpose(uu2)]])
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = hcat(AA, c)

vars = matread("./src/measured_data/CBC_stable_v16_5.mat")
uu = get(vars, "data", 1)
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = vcat(t_series, [[transpose(uu1); transpose(uu2)]])
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = hcat(AA, c)

vars = matread("./src/measured_data/CBC_stable_v17_3.mat")
uu = get(vars, "data", 1)
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = vcat(t_series, [[transpose(uu1); transpose(uu2)]])
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = hcat(AA, c)

vars = matread("./src/measured_data/CBC_unstable_v14_9.mat")
uu = get(vars, "data", 1)
ind1 = 1;
ind2 = 4;
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = vcat(t_series, [[transpose(uu1); transpose(uu2)]])
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = hcat(AA, c)

vars = matread("./src/measured_data/CBC_unstable_v15_6.mat")
uu = get(vars, "data", 1)
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = vcat(t_series, [[transpose(uu1); transpose(uu2)]])
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = hcat(AA, c)

vars = matread("./src/measured_data/CBC_unstable_v16_5.mat")
uu = get(vars, "data", 1)
uu1 = uu[ind1, 1:l]
uu2 = uu[ind2, 1:l]
uu1 = uu[ind1, 1:l] - mu1 * ones(l)
uu2 = uu[ind2, 1:l] - mu2 * ones(l)
t_series = vcat(t_series, [[transpose(uu1); transpose(uu2)]])
rr = sqrt.(uu1 .^ 2 + uu2 .^ 2)
tt = atan.(uu2, uu1)
c = ML_Hopf.LS_harmonics(rr, tt, 1, nh).coeff
AA = hcat(AA, c)
##
vel_l = 4
Vel = [14.9, 15.6, 16.5, 17.3]
Vel2 = [14.9, 15.6, 16.5]
θ_l = 300
θ = range(0, stop = 2π, length = θ_l)
coθ = cos.(θ)
siθ = sin.(θ)
## Linear transformation
rot = π * 0.1
R = [cos(rot) -sin(rot); sin(rot) cos(rot)]
θ = vec(1e-2 * R * [8.0 0.0; 0.0 1.7])
θ = vcat(θ, zeros(2))
scale_f_l = 1e1 # optimization works for scale_f_l>=50 for small scale_f_l optimization does not work.

hidden = 12
ann_l = FastChain(
    FastDense(2, hidden, tanh),
    FastDense(hidden, hidden, tanh),
    FastDense(hidden, 4),
)
θl = initial_params(ann_l)
scale_f2 = 1e2
θ = vcat(θ, θl)
pp = [17.95, 3.85]
θ = vcat(θ, pp)

## Nonlinear transformation
hidden = 11
ann = FastChain(
    FastDense(3, hidden, tanh),
    FastDense(hidden, hidden, tanh),
    FastDense(hidden, 2),
)
θn = initial_params(ann)
scale_f = 1e3

loss_lt_nn_(θ_t) =
    ML_Hopf.loss_lt_nn(θ_t, ann_l, scale_f_l, scale_f2, AA, θ_l, nh, Vel, Vel2, coθ, siθ)
res_l = DiffEqFlux.sciml_train(loss_lt_nn_, θ, ADAM(0.001), maxiters = 100)
U₀ = res_l.minimizer[end-1];
s_ = res_l.minimizer[end];
loss_lt_nn_2(θ_t) = ML_Hopf.loss_lt_nn2(
    θ_t,
    ann_l,
    scale_f_l,
    scale_f2,
    U₀,
    s_,
    AA,
    θ_l,
    nh,
    Vel,
    Vel2,
    coθ,
    siθ,
)
res_l = DiffEqFlux.sciml_train(
    loss_lt_nn_2,
    res_l.minimizer,
    BFGS(initial_stepnorm = 1e-4),
    maxiters = 300,
) # First, train the simple model
θ_ = res_l.minimizer
θn = vcat(θn, [U₀, s_])
loss_nt_(θ_t) = ML_Hopf.loss_nt2(
    θ_t,
    θ_,
    ann_l,
    scale_f,
    scale_f2,
    scale_f_l,
    ann,
    AA,
    nh,
    Vel,
    Vel2,
    θ_l,
    coθ,
    siθ,
)
res1 = DiffEqFlux.sciml_train(loss_nt_, θn, ADAM(0.001), maxiters = 100) # Train more complicated model

U₀ = res1.minimizer[end-1];
s_ = res1.minimizer[end];
θ_n = res1.minimizer
## Define parameters of the model as global variable
np1 = θ_n[end-1];
np2 = θ_n[end];
nf = ML_Hopf.nf_dis(np1, np2, Vel, Vel2, coθ, siθ)
vl = nf.v;
vl2 = nf.v2;
p1, p2, p3, p4, p5, p6 = θ_[1:6]
T = [p1 p3; p2 p4]
pn1 = θ_[7:end-2]
pn2 = θ_n[1:end-2]
##
sens = 18 / 100 # sensitivity of the lasor sensor
Ap = ML_Hopf.lt_pp_n2(
    θ_n,
    θ_,
    ann_l,
    scale_f,
    scale_f2,
    scale_f_l,
    ann,
    Vel,
    Vel2,
    θ_l,
    coθ,
    siθ,
)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -4e-2,
        xmax = 4e-2,
        ymax = 7e-2,
        ymin = -8e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[1][1, :] * sens, Ap[1][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", mark = "o", mark_size = "0.5pt"},
        Coordinates(t_series[1][1, :] * sens, t_series[1][2, :]),
    ),
    LegendEntry("Measured data"),
)
pgfsave("./Figures/exp/pp_s149.pdf", a)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -4e-2,
        xmax = 4e-2,
        ymax = 7e-2,
        ymin = -8e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[5][1, :] * sens, Ap[5][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", mark = "o", mark_size = "0.5pt"},
        Coordinates(t_series[5][1, :] * sens, t_series[5][2, :]),
    ),
    LegendEntry("Measured data"),
)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -4e-2,
        xmax = 4e-2,
        ymax = 7e-2,
        ymin = -8e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[3][1, :] * sens, Ap[3][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", mark = "o", mark_size = "0.5pt"},
        Coordinates(t_series[3][1, :] * sens, t_series[3][2, :]),
    ),
    LegendEntry("Measured data"),
)

a = @pgf Axis(
    {
        xlabel = L"$h$ (m)",
        ylabel = L"$\alpha$ (rad)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = -4e-2,
        xmax = 4e-2,
        ymax = 7e-2,
        ymin = -8e-2,
    },
    Plot({color = "red", no_marks}, Coordinates(Ap[7][1, :] * sens, Ap[7][2, :])),
    LegendEntry("Learnt model"),
    Plot(
        {color = "blue", mark = "o", mark_size = "0.5pt"},
        Coordinates(t_series[7][1, :] * sens, t_series[7][2, :]),
    ),
    LegendEntry("Measured data"),
)

## Plot bifurcation diagram
bd = ML_Hopf.lt_b_dia2(1, θ_n, θ_, ann_l, ann, scale_f_l, scale_f, scale_f2, coθ, siθ)
h = [maximum(t_series[i][1, :]) - minimum(t_series[i][1, :]) for i = 1:length(Vel)]
h2 = [maximum(t_series[i+4][1, :]) - minimum(t_series[i+4][1, :]) for i = 1:length(Vel2)]

a = @pgf Axis(
    {
        xlabel = "Wind speed (m/sec)",
        ylabel = "Heave amplitude (m)",
        legend_pos = "north west",
        height = "11cm",
        width = "15cm",
        ymin = 0,
        ymax = 9e-2,
        mark_options = {scale = 1.5},
    },
    Plot({color = "blue", only_marks}, Coordinates(Vel, h * sens)),
    LegendEntry("Measured data  (stable LCO)"),
    Plot({color = "red", only_marks, mark = "triangle*"}, Coordinates(Vel2, h2 * sens)),
    LegendEntry("Measured data (unstable LCO)"),
    Plot(
        {color = "blue", no_marks},
        Coordinates(vcat(reverse(bd.v), bd.v), vcat(reverse(bd.u * sens), bd.s * sens)),
    ),
    LegendEntry("Bifrucation based ML model"),
)

## Train speed of the phase
function Inv_T_u(th0, vel)
    s_amp = sqrt(np2 / 2 + sqrt(np2^2 + 4 * (vel - np1)) / 2)
    ttl = Int(1e5)
    theta = range(0, 2π, length = ttl)
    uu = [transpose(s_amp * cos.(theta)); transpose(s_amp * sin.(theta))]
    dis = transpose([p5 * ones(ttl) p6 * ones(ttl)]) / scale_f_l
    Tu = ML_Hopf.Nt(uu, T, vel, dis, pn1, np1, pn2, ann_l, scale_f, scale_f2, ann)
    t0 = [abs(atan(Tu[2, i], Tu[1, i]) - th0) for i = 1:length(theta)]
    er = minimum(t0)
    return theta[argmin(t0)]
end

function Inv_T_uu(th0, vel)
    u_amp = sqrt(np2 / 2 + sqrt(np2^2 - 4 * (vel - np1)) / 2)
    ttl = Int(1e4)
    theta = range(0, 2π, length = ttl)
    uu = [transpose(u_amp * cos.(theta)); transpose(u_amp * sin.(theta))]
    dis = transpose([p5 * ones(ttl) p6 * ones(ttl)]) / scale_f_l
    Tu = ML_Hopf.Nt(uu, T, vel, dis, pn1, np1, pn2, ann_l, scale_f, scale_f2, ann)
    t0 = [abs(atan(Tu[2, i], Tu[1, i]) - th0) for i = 1:length(theta)]
    er = minimum(t0)
    return theta[argmin(t0)]
end

function dudt_ph(u, p, t) # speed of phase
    np1 = U₀
    np2 = s_
    θ = u[1]
    r = u[2]
    c = u[3]
    δ₀ = np1
    ν = (c - δ₀)
    ω₀ = p[1]
    uu = [r * cos(θ), r * sin(θ), ν]
    du₁ = ω₀ + ann3(uu, p[2:end])[1] / om_scale
    du₂ = 0
    du₃ = 0
    [du₁, du₂, du₃]
end

function predict_time_T(p) # Predict the time series of the model with computed initial conditions
    #vel_l=300
    A1 = [
        Array(
            concrete_solve(
                ODEProblem(dudt_ph, u_t0[i], (0, tl2), p),
                Tsit5(),
                u_t0[i],
                p,
                saveat = st,
                abstol = 1e-8,
                reltol = 1e-8,
                sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP()),
            ),
        ) for i = 1:length(Vel)
    ]
    uu = [
        transpose(
            hcat(
                A1[i][2, :] .* cos.(A1[i][1, :]),
                A1[i][2, :] .* sin.(A1[i][1, :]),
                A1[i][3, :],
            ),
        ) for i = 1:length(Vel)
    ]
    delU = zeros(2, spl)
    delU2 = -np1 * ones(1, spl)
    delU = vcat(delU, delU2)
    uu = [uu[i] + delU for i = 1:length(Vel)]
    vl = [uu[i][1:2, :] for i = 1:length(Vel)]
    vlT = [
        ML_Hopf.Nt(vl[i], T, Vel[i], dis, pn1, np1, pn2, ann_l, scale_f, scale_f2, ann)
        for i = 1:length(Vel)
    ]
    Pr = zeros(0, spl)
    for i = 1:length(Vel)
        theta = vlT[i][[1, 2], :]
        Pr = vcat(Pr, theta)
    end
    ##
    A1 = [
        Array(
            concrete_solve(
                ODEProblem(dudt_ph, uu_t0[i], (0, tl2), p),
                Tsit5(),
                uu_t0[i],
                p,
                saveat = st,
                abstol = 1e-8,
                reltol = 1e-8,
                sensealg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP()),
            ),
        ) for i = 1:length(Vel2)
    ]
    uu = [
        transpose(
            hcat(
                A1[i][2, :] .* cos.(A1[i][1, :]),
                A1[i][2, :] .* sin.(A1[i][1, :]),
                A1[i][3, :],
            ),
        ) for i = 1:length(Vel2)
    ]
    delU = zeros(2, spl)
    delU2 = -np1 * ones(1, spl)
    delU = vcat(delU, delU2)
    uu = [uu[i] + delU for i = 1:1:length(Vel2)]

    vl = [uu[i][1:2, :] for i = 1:length(Vel2)]
    vlT = [
        ML_Hopf.Nt(vl[i], T, Vel2[i], dis, pn1, np1, pn2, ann_l, scale_f, scale_f2, ann) for i = 1:length(Vel2)
    ]
    Pr2 = zeros(0, spl)
    for i = 1:length(Vel2)
        theta = vlT[i][[1, 2], :]
        Pr2 = vcat(Pr2, theta)
    end
    ##
    [Pr; Pr2]
end
##
tl2 = 1.0
st = 1e-3
spl = Int(tl2 / st + 1)
vl = 7
t_s = zeros(vl * 2, spl)
t_series2 = [t_series[j][:, 1:5:end] for j = 1:7]

for i = 1:vl
    t_s[[2 * (i - 1) + 1, 2 * (i - 1) + 2], :] = t_series2[i][:, 1:1001]
end
A3 = t_s
## Generate data and initial θ
np1 = U₀;
np2 = s_;
s_amp = [sqrt(np2 / 2 + sqrt(np2^2 + 4 * (Vel[i] - np1)) / 2) for i = 1:length(Vel)]
u_amp = [sqrt(np2 / 2 - sqrt(np2^2 + 4 * (Vel2[i] - np1)) / 2) for i = 1:length(Vel2)]
theta0 = [atan(t_series[i][2, 1], t_series[i][1, 1]) for i = 1:length(Vel)]
theta0u = [
    atan(t_series[i+length(Vel)][2, 1], t_series[i+length(Vel)][1, 1]) for
    i = 1:length(Vel2)
]
θ₀ = [Inv_T_u(theta0[i], Vel[i]) for i = 1:length(Vel)]
θ₀u = [Inv_T_uu(theta0u[i], Vel2[i]) for i = 1:length(Vel2)]
u_t0 = [[θ₀[i], s_amp[i], Vel[i]] for i = 1:length(Vel)]
uu_t0 = [[θ₀u[i], u_amp[i], Vel2[i]] for i = 1:length(Vel2)]

function loss_time_T(p) # Loss function for the time series
    pred = predict_time_T(p)
    sum(abs2, A3 .- pred) # + 1e-5*sum(sum.(abs, params(ann)))
end

om_scale = 0.3
tl2 = 1.0
st = 1e-3
spl = Int(tl2 / st + 1)
dis = transpose([p5 * ones(spl) p6 * ones(spl)]) / scale_f_l

vl = length(Vel)
hidden = 21
ann3 = FastChain(FastDense(3, hidden, tanh), FastDense(hidden, 1, tanh))
np = initial_params(ann3)
omega = 16.3
p = vcat(omega, np)

loss_time_T(p)
res_t = DiffEqFlux.sciml_train(loss_time_T, p, ADAM(0.01), maxiters = 200)
res_t = DiffEqFlux.sciml_train(
    loss_time_T,
    res_t.minimizer,
    BFGS(initial_stepnorm = 1e-3),
    maxiters = 10000,
)

p = res_t.minimizer

## Plot time-series

tv = range(0, tl2, length = spl) #Generate a time vector
a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$h$ (m)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 1,
        ymax = 5e-2,
        ymin = -5e-2,
        mark_options = {scale = 0.1},
    },
    Plot(
        {color = "red", no_marks},
        Coordinates(tv, sens * predict_time_T(p)[2*(1-1)+1, :]),
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, sens * t_series2[1][1, 1:1001])),
    LegendEntry("Measured data"),
)
#pgfsave("./Figures/exp/t_s149.pdf",a)

a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$h$ (m)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 1,
        ymax = 5e-2,
        ymin = -5e-2,
        mark_options = {scale = 0.1},
    },
    Plot(
        {color = "red", no_marks},
        Coordinates(tv, sens * predict_time_T(p)[2*(5-1)+1, :]),
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, sens * t_series2[5][1, 1:1001])),
    LegendEntry("Measured data"),
)

#pgfsave("./Figures/exp/t_u149.pdf",a)
##

a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$h$ (m)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 1,
        ymax = 5e-2,
        ymin = -5e-2,
        mark_options = {scale = 0.1},
    },
    Plot(
        {color = "red", no_marks},
        Coordinates(tv, sens * predict_time_T(p)[2*(3-1)+1, :]),
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, sens * t_series2[3][1, 1:1001])),
    LegendEntry("Measured data"),
)
#pgfsave("./Figures/exp/t_s165.pdf",a)

a = @pgf Axis(
    {
        xlabel = "Time (sec)",
        ylabel = L"$h$ (m)",
        legend_pos = "north west",
        height = "9cm",
        width = "9cm",
        xmin = 0,
        xmax = 1,
        ymax = 5e-2,
        ymin = -5e-2,
        mark_options = {scale = 0.1},
    },
    Plot(
        {color = "red", no_marks},
        Coordinates(tv, sens * predict_time_T(p)[2*(7-1)+1, :]),
    ),
    LegendEntry("Learnt model"),
    Plot({color = "blue", no_marks}, Coordinates(tv, sens * t_series2[7][1, 1:1001])),
    LegendEntry("Measured data"),
)
#pgfsave("./Figures/exp/t_u165.pdf",a)


text:
"""
Evaluates an expression without compiling it.
Uses AST and symbol lookups. Only supports :call and :(=)
expressions at the moment.
Example:
tab = SymbolTable(:f => f, :x => x)
ex = :(f(x))
interpret(tab, ex)
"""
module Interpreter

export SymbolTable, interpret

const SymbolTable = Dict{Symbol,Any}

interpret(tab::SymbolTable, x::Any) = x
interpret(tab::SymbolTable, s::Symbol) = tab[s]

function interpret(tab::SymbolTable, ex::Expr)
    args = ex.args
    if ex.head == :call
        if ex.args[1] == Symbol(".&")
            return (interpret(tab, args[2]) .& interpret(tab, args[3]))
        elseif ex.args[1] == Symbol(".|")
            return (interpret(tab, args[2]) .| interpret(tab, args[3]))
        elseif ex.args[1] == Symbol(".==")
            return (interpret(tab, args[2]) .== interpret(tab, args[3]))
        elseif ex.args[1] == Symbol(".>=")
            return (interpret(tab, args[2]) .>= interpret(tab, args[3]))
        elseif ex.args[1] == Symbol(".<=")
            return (interpret(tab, args[2]) .<= interpret(tab, args[3]))
        else
            len = length(args)
            #unroll for performance and avoid excessive allocations
            if len == 1
                return tab[args[1]]()
            elseif len == 2
                return tab[args[1]](interpret(tab,args[2]))
            elseif len == 3
                return tab[args[1]](interpret(tab,args[2]), interpret(tab,args[3]))
            elseif len == 4
                return tab[args[1]](interpret(tab,args[2]), interpret(tab,args[3]), interpret(tab,args[4]))
            elseif len == 5
                return tab[args[1]](interpret(tab,args[2]), interpret(tab,args[3]), interpret(tab,args[4]),
                                       interpret(tab,args[5]))
            elseif len == 6
                return tab[args[1]](interpret(tab,args[2]), interpret(tab,args[3]), interpret(tab,args[4]),
                                       interpret(tab,args[5]), interpret(tab,args[6]))
            else
                return tab[args[1]](interpret.(Ref(tab),args[2:end])...)
            end
        end
    elseif ex.head == :(.)
        return Base.broadcast(Base.eval(args[1]), interpret(tab, args[2])...)
    elseif ex.head == :tuple
        return tuple(interpret.(Ref(tab), args)...)
    elseif ex.head == :vect
        return [interpret.(Ref(tab), args)...]
    elseif ex.head == :||
        return (interpret(tab, args[1]) || interpret(tab, args[2]))
    elseif ex.head == :&&
        return (interpret(tab, args[1]) && interpret(tab, args[2]))
    elseif ex.head == :(=)
        return (tab[args[1]] = interpret(tab, args[2])) #assignments made to symboltable
    elseif ex.head == :block
        result = nothing
        for x in args
            result = interpret(tab, x)
        end
        return result
    elseif ex.head == :if
        if interpret(tab, args[1])
            return interpret(tab, args[2])
        else
            return interpret(tab, args[3])
        end
    else
        error("Expression type not supported")
    end
end

### Raw interpret, no symbol table
function interpret(ex::Expr, M::Module=Main)
    result = if ex.head == :call
        call_func(M, ex.args...)
    elseif ex.head == :vect
        ex.args
    else
        Core.eval(M, ex)
    end
end
call_func(M::Module, f::Symbol) = getproperty(M,f)()
call_func(M::Module, f::Symbol, x1) = getproperty(M,f)f(x1)
call_func(M::Module, f::Symbol, x1, x2) = getproperty(M,f)(x1, x2)
call_func(M::Module, f::Symbol, x1, x2, x3) = getproperty(M,f)(x1, x2, x3)
call_func(M::Module, f::Symbol, x1, x2, x3, x4) = getproperty(M,f)(x1, x2, x3, x4)
call_func(M::Module, f::Symbol, x1, x2, x3, x4, x5) = getproperty(M,f)(x1, x2, x3, x4, x5)

end #module


text:
using JuMP

function objective_feasibility_problem(pm::AbstractPowerModel, x_vals::Dict{Int64,Float64}=Dict{Int64,Float64}(); nw::Int=pm.cnw)
    for l in ids(pm,pm.cnw,:branch)
        if !haskey(x_vals,l)
	        if l in pm.data["inactive_branches"]
		        x_vals[l]=1
	        else
		        x_vals[l]=0
	        end
        end
    end
    pd_br = var(pm,nw,:pd_br)
    pt_br = var(pm,nw,:pt_br)
    qd_br = var(pm,nw,:qd_br)
    qt_br = var(pm,nw,:qt_br)
    protected_arcs = filter(a->(a[1] in pm.data["protected_branches"]),ref(pm,nw,:arcs))
    inactive_arcs = filter(a->(a[1] in pm.data["inactive_branches"]),ref(pm,nw,:arcs))
    undecided_arcs = filter(a->!(a in protected_arcs || a in inactive_arcs),ref(pm,nw,:arcs))

    return JuMP.@objective(pm.model, Min,
	sum( (1-x_vals[a[1]])*(pd_br[a,0] + qd_br[a,0] + pd_br[a,1] + qd_br[a,1])
            + x_vals[a[1]]*(pt_br[a,0] + qt_br[a,0] + pt_br[a,1] + qt_br[a,1]) for a in undecided_arcs)
            + sum( (pd_br[a,0] + qd_br[a,0] + pd_br[a,1] + qd_br[a,1])  for a in protected_arcs)
            + sum( (pt_br[a,0] + qt_br[a,0] + pt_br[a,1] + qt_br[a,1])  for a in inactive_arcs)
    )
end

function objective_minmax_problem(pm::AbstractPowerModel; nw::Int=pm.cnw)
    K = pm.data["attacker_budget"]
    pd_br = var(pm,nw,:pd_br)
    pt_br = var(pm,nw,:pt_br)
    qd_br = var(pm,nw,:qd_br)
    qt_br = var(pm,nw,:qt_br)
    u_ord_aux = var(pm,nw,:u_ord_aux)
    u_K = var(pm,nw,:u_K)
    undecided_branches = filter(l->!(l in pm.data["protected_branches"] || l in pm.data["inactive_branches"]), ids(pm,nw,:branch))
    undecided_arcs = filter(a->(a[1] in undecided_branches),ref(pm,nw,:arcs))
    protected_arcs = filter(a->(a[1] in pm.data["protected_branches"]),ref(pm,nw,:arcs))
    attacked_arcs = filter(a->(a[1] in pm.data["inactive_branches"]),ref(pm,nw,:arcs))

    return JuMP.@objective(pm.model, Min, K*u_K + sum( u_ord_aux[l] for l in undecided_branches )
            + sum( pd_br[a,0] + qd_br[a,0] + pd_br[a,1] + qd_br[a,1] for a in undecided_arcs)
            + sum( pd_br[a,0] + qd_br[a,0] + pd_br[a,1] + qd_br[a,1] for a in protected_arcs)
            + sum( pt_br[a,0] + qt_br[a,0] + pt_br[a,1] + qt_br[a,1] for a in attacked_arcs)
    )
end


text:
n = 10 # rand(1:10)
@test matrixdepot("fiedler", n) == matrixdepot("fiedler", [1:n;])

A = matrixdepot("fiedler", n)

@test issymmetric(A)
for i = 1:n, j =1:n
    @test A[i,j] ≈ abs(i-j)
end
println("'fiedler' passed test...")