# # Problems and the system

#md # ```@meta
#md # CurrentModule = ImmersedLayers
#md # ```

#=
In specific problems that we wish to solve with immersed layers, there may
be other data and operators that we would like to cache. We do this with
an *extra cache*, which the user can define, along with a problem type associated
with this cache. The basic cache and the extra cache are generated and associated
together in a *system*.

There are a few basic ingredients to do this:
* **Create a problem type**, using the macro `@ilmproblem`. This mostly just serves as a means of dispatching correctly, but also holds the grid and bodies so they can be passed along when the caches are constructed.
* **Create an extra cache type**, making it a subtype of `AbstractExtraILMCache`. This can hold pretty much anything you want it to.
* **Extend the function `prob_cache(prob,base_cache)`** to serve as a constructor for your extra cache, when your problem type is passed in.

Optionally, you can also extend the function `solve` in order to perform the steps
of your algorithm. However, generically, you can just use pass in the system structure,
which holds the basic ILM cache and your extra cache, into any function.
=#

#=
## Example of problem and system use
We will demonstrate the use of problems and systems with the example
given in [A Dirichlet Poisson problem](@ref). Here, we will
assemble the various additional data structures and operators used to
solve this problem into an extra cache. We will also create a problem
type called `DirichletPoissonProblem`, which we make a subtype of
`AbstractScalarILMProblem`.
=#

using ImmersedLayers
using Plots
using UnPack


#=
### Create your problem type
For this job, we have a macro
=#
@ilmproblem DirichletPoisson scalar

#=
This generates a type `DirichletPoissonProblem` for dispatch, and an associated
constructor that we'll use later.
=#

#=
### Create your extra cache
Here, we'd like this extra cache to hold the Schur complement and the
filtering matrices, as well as some cache variables.
=#
struct DirichletPoissonCache{SMT,CMT,ST,FT} <: ImmersedLayers.AbstractExtraILMCache
   S :: SMT
   C :: CMT
   fb :: ST
   fstar :: FT
end

#=
### Extend the `prob_cache` function
We need this to construct our extra cache
=#
function ImmersedLayers.prob_cache(prob::DirichletPoissonProblem,base_cache::BasicILMCache)
    S = create_RTLinvR(base_cache)
    C = create_surface_filter(base_cache)
    fb = zeros_surface(base_cache)
    fstar = zeros_grid(base_cache)
    DirichletPoissonCache(S,C,fb,fstar)
end

#=
### Extend the `solve` function
Here, we actually do the work of the algorithm, making use of all of the
operators and data structures that we have cached for efficiency.
The example below takes in some surface Dirichlet data `fbplus`,
and returns the solutions `f` and `s` (filtered).
=#
function ImmersedLayers.solve(fbplus,prob::DirichletPoissonProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack S, C, fb, fstar = extra_cache

    f = zeros_grid(base_cache)
    s = zeros_surface(base_cache)

    surface_divergence!(fstar,fbplus,base_cache)
    fb .= 0.5*fbplus

    inverse_laplacian!(fstar,base_cache)

    interpolate!(s,fstar,base_cache)
    s .= fb - s
    s .= -(S\s);

    regularize!(f,s,base_cache)
    inverse_laplacian!(f,base_cache)
    f .+= fstar;

    return f, C^6*s
end

#=
### Set up the grid, shape, and cache
We do this just as we did in [Immersed layer caches](@ref), but
now we don't create a cache, since it will be done internally.
=#
Δx = 0.01
Lx = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx)
RadC = 1.0
Δs = 1.4*cellsize(g)
body = Circle(RadC,Δs)

#=
### Do the work
We do this in three steps:
- Create the problem instance. This is where we use the constructor generated by `@ilmproblem`
- Call `__init` to create the caches, assembled into a system
- Call `solve` to solve the problem.
Also, note that pretty much any function that accepts `base_cache`
as an argument also accepts `sys`.
=#
prob = DirichletPoissonProblem(g,body,scaling=GridScaling)
sys = ImmersedLayers.__init(prob)

pts = points(sys)
f, s = solve(pts.u,prob,sys)

plot(f,sys)
#-
plot(s)

#=
In the previous example, we supplied the surface data as an argument directly to
the `solve` function. However, sometimes we need to supply surface data with a little
more complexity (as in time-varying cases). For this purpose, we can create
functions that supply the boundary data when called, e.g.,
=#
get_fbplus(base_cache) = points(base_cache).u
get_fbminus(base_cache) = zeros_surface(base_cache)

#=
We can pass these along to the problem specification by various means.
It's at the user's discretion how to do it, and how to make use of them
in the solution. Here's an example, using a `Dict`
=#
bcdict = Dict("fbplus"=>get_fbplus,"fbminus"=>get_fbminus)

#=
We can also pass along physical parameters and forcing functions.
Let's use the latter feature in order to set some point sources at
a few locations. For this, we create a function that will take in
the base cache and the right-hand side data vector, and populate
this vector with the point sources. We use the `Regularize` function
in the `CartesianGrids.jl` package to immerse the point sources
into the grid. In the example shown here, two points: one at (-1.0,1.5)
and the other at (1.0,1.5), with respective strengths -1 and 1,
are placed.
=#
function rhs_func!(rhs,base_cache)
    X = VectorData([-1.0,1.0],[1.5,1.5])
    str = ScalarData([-1.0,1.0])
    reg = Regularize(X,cellsize(base_cache.g),I0=origin(base_cache.g),ddftype=CartesianGrids.M4prime)
    return reg(rhs,str)
end


#=
We redefine the `solve` function and use these functions in place of
the original argument:
=#
function ImmersedLayers.solve(prob::DirichletPoissonProblem,sys::ILMSystem)
    @unpack extra_cache, base_cache = sys
    @unpack bc, f_funcs, gdata_cache = base_cache
    @unpack S, C, fb, fstar = extra_cache

    f = zeros_grid(base_cache)
    s = zeros_surface(base_cache)

    ## Get the boundary data on each side of the interface
    fbplus = bc["fbplus"](base_cache)
    fbminus = bc["fbminus"](base_cache)

    ## Evaluate the forcing field
    f_funcs(gdata_cache,base_cache)

    ## Evaluate the right-hand side of Poisson equation
    surface_divergence!(fstar,fbplus-fbminus,base_cache)
    fstar .+= gdata_cache

    fb .= 0.5*(fbplus+fbminus)

    inverse_laplacian!(fstar,base_cache)

    interpolate!(s,fstar,base_cache)
    s .= fb - s
    s .= -(S\s);

    regularize!(f,s,base_cache)
    inverse_laplacian!(f,base_cache)
    f .+= fstar;

    return f, C^6*s
end

#=
Now we specify the problem, create the system, and solve it, as before,
but now supplying the boundary condition functions with the `bc` keyword:
=#
prob = DirichletPoissonProblem(g,body,scaling=GridScaling,bc=bcdict,f_funcs=rhs_func!)
sys = ImmersedLayers.__init(prob)
f, s = solve(prob,sys)
plot(f,sys)

#=
Same solution, of course. But suppose we wish to change the
the boundary conditions or source points? We can do it easily without regenerating the
cache and system, simply by redefining our bc and forcing functions, e.g.
to create an internal solution, with surface data equal to the $y$ coordinate,
=#
get_fbplus(base_cache) = zeros_surface(base_cache)
get_fbminus(base_cache) = points(base_cache).v
function rhs_func!(rhs,base_cache)
    X = VectorData([-0.2,0.2],[0.0,0.0])
    str = ScalarData([-1.0,1.0])
    reg = Regularize(X,cellsize(base_cache.g),I0=origin(base_cache.g),ddftype=CartesianGrids.M4prime)
    return reg(rhs,str)
end
#=
We can solve immediately without having to reconstruct the system, so
it's very fast.
=#
@time f, s = solve(prob,sys)
plot(f,sys)

#md # ## Problem types and functions
#md #
#md # ```@docs
#md # AbstractScalarILMProblem
#md # AbstractVectorILMProblem
#md # @ilmproblem
#md # BasicScalarILMProblem
#md # BasicVectorILMProblem
#md # prob_cache
#md # ```

#md # ## System types and functions
#md #
#md # ```@docs
#md # ILMSystem
#md # __init
#md # ```
