fd.jl

module FDM

using StaticArrays
using SparseArrays
using OffsetArrays
using LinearMaps
using IterativeSolvers

include("filtering.jl")

"Finite Differences"
module FD

using OffsetArrays

const forward = OffsetVector([-1, 1], -1)
const backward = OffsetVector([-1, 1], -2)
const central = OffsetVector([-1, 0, 1], -2)
function dirfd(fd, ndims, dir)
  dir <= ndims || throw(ArgumentError("need dir <= ndims!"))
  dfd = OffsetArray(zeros(ntuple(i -> i == dir ? length(fd) : 1, ndims)), ntuple(i -> i == dir ? fd.offsets[1] : -1, ndims))
  # Workaround: LinearIndices(::OffsetVector) are not 1-based, so copyto! chockes
  copyto!(dfd, fd.parent)
end

end


using .Filtering
using .FD

"""
    _fdkern(fd, ndims, dir)

Create from `fd` a finite difference kernel usable in `ndims` dimension
oriented in dimension `dir`.
"""
# TODO: make it typestable on ndims
function _fdkern(fd::OffsetArray{T,1,Array{T,1}}, ndims::Integer, dir::Integer) where T
  1 <= dir <= ndims || throw(ArgumentError("dimesion $(dir) out of range 1:$(ndims)"))
  dfd = OffsetArray{T}(undef, Tuple(i == dir ? axes(fd, 1) : (0:0) for i in 1:ndims))
  # copyto! from OffsetArray fails, broadcasting for differing axes fails too
  # workaround: copyto! from parent
  copyto!(dfd, fd.parent)
end


# Grid


abstract type AbstractGrid{d} <: AbstractArray{Float64, d} end

Base.ndims(grid::AbstractGrid{d}) where d = d
Base.axes(grid::AbstractGrid) = Base.OneTo.(size(grid))
Base.length(grid::AbstractGrid) = prod(size(grid))
Base.IndexStyle(::Type{<:AbstractGrid}) = IndexCartesian()

# aka `StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`
const FDStepRange = typeof(0:0.1:1)

struct Grid{d} <: AbstractGrid{d}
  domain::NTuple{d, FDStepRange}
end
Grid(domain::FDStepRange...) = Grid(domain)
Base.show(io::IO, grid::Grid) =
    print("Grid from $(first.(grid.domain)) to $(last.(grid.domain)) ",
          "with spacing $(step.(grid.domain))")

Base.size(grid::Grid) = ntuple(i -> grid.domain[i].len, ndims(grid))

Base.iterate(grid::Grid) = iterate(Iterators.product(grid.domain...))
Base.iterate(grid::Grid, state) = iterate(Iterators.product(grid.domain...), state)

Base.getindex(grid::Grid, idx...) = SVector(getindex.(grid.domain, idx))
Base.getindex(grid::Grid, I::CartesianIndex) = SVector(getindex.(grid.domain, Tuple(I)))


# GridView

struct GridView{d} <: AbstractGrid{d}
  grid::AbstractGrid{d}
  indices::NTuple{d, UnitRange{Int}}
end
GridView(grid, indices::UnitRange{Int}...) = GridView(grid, indices)

Base.size(grid::GridView) = length.(grid.indices)

# FIXME: this is terrible!
Base.iterate(grid::GridView) = iterate(grid.grid[I] for I in CartesianIndices(grid.indices))
Base.iterate(grid::GridView, state) = iterate([grid.grid[I] for I in CartesianIndices(grid.indices)], state)

Base.getindex(grid::GridView, idx...) = getindex(grid.grid, (first.(grid.indices) .- 1 .+ Base.to_indices(grid, idx))...)

# MetaGrid


struct MetaGrid{d} <: AbstractGrid{d}
  grid::Grid{d}
  indices::Array{NTuple{d, UnitRange{Int}}, d}
  overlap::Int
end

"""
    MetaGrid(domain::NTuple{d, FDStepRange}, pnum::NTuple{d, Int}, overlap::Int)

Create a grid on `domain`, partitioned along dimensions according to `pnum`
respecting `overlap`.
"""
function MetaGrid(domain::NTuple{d, FDStepRange}, pnum::NTuple{d, Int}, overlap::Int) where d
  grid = Grid(domain)
  tsize = size(grid) .+ (pnum .- 1) .* overlap

  psize = tsize .÷ pnum
  osize = tsize .- pnum .* psize

  overhang(I, j) = I[j] == pnum[j] ? osize[j] : 0

  indices = Array{NTuple{d, UnitRange{Int}}, d}(undef, pnum)
  for I in CartesianIndices(pnum)
    indices[I] = ntuple(j -> ((I[j] - 1) * psize[j] - (I[j] - 1) * overlap + 1) :
                             (      I[j] * psize[j] - (I[j] - 1) * overlap + overhang(I, j)), d)
  end
  MetaGrid{d}(grid, indices, overlap)
end

Base.size(grid::MetaGrid) = size(grid.grid)

Base.iterate(grid::MetaGrid) = iterate(grid.grid)
Base.iterate(grid::MetaGrid, state) = iterate(grid.grid, state)

Base.getindex(grid::MetaGrid, idx...) = getindex(grid.grid, idx...)

indices(grid::MetaGrid, idx...) = getindex(grid.indices, idx...)

Base.view(grid::MetaGrid, idx...) = GridView(grid, indices(grid, idx...))

parts(grid::MetaGrid) = CartesianIndices(grid.indices)


## GridFunction

abstract type AbstractGridFunction{G,T,d,n} <: AbstractArray{T,1} end

rangedim(f::AbstractGridFunction{G,T,d,n}) where {G,T,d,n} = n
domaindim(f::AbstractGridFunction{G,T,d}) where {G,T,d} = d

Base.similar(f::F, ::Type{S}, dims::Int...) where {F<:AbstractGridFunction, S} =
    GridFunction(f.grid, Array{S}(undef, length(f)))


"""
    GridFunction{G<:AbstractGrid,T,d,n,D}

`n`-dimensional function values of type `T` on a Grid `G` of dimension `d`.

`D == d + 1`
"""
struct GridFunction{G<:AbstractGrid,T,d,n,A,D} <: AbstractGridFunction{G,T,d,n}
  grid::G
  data::A
  function GridFunction{G,T,d,n,A,D}(grid::G, data::A) where {G,T,d,n,D,A<:AbstractArray{T,D}}
    D == d + 1 || throw(ArgumentError)
    d == ndims(grid) || throw(ArgumentError)
    (size(grid)..., n) == size(data) || throw(ArgumentError("data does not match given size"))
    new(grid, data)
  end
end
# TODO: check size(data)
#GridFunction(grid::G, data::AbstractArray{T,D}) where {d,G<:AbstractGrid{d},T,D} =
#    GridFunction{G,T,d,size(data, d+1),typeof(data),d+1}(grid, data)
function GridFunction(grid::G, data::AbstractArray, ::Val{n}) where {d,G<:AbstractGrid{d}, n}
  if length(data) ÷ length(grid) == 1 && n > 1
    data = repeat(data, inner=ntuple(i -> i == d+1 ? n : 1, d+1))
  else
    # TODO: check size(data)
    data = reshape(data, size(grid)..., n)
  end
  length(data) ÷ length(grid) == n || throw(ArgumentError("dimensions don't match"))
  GridFunction{G,eltype(data),d,n,typeof(data),d+1}(grid, data)
end
function GridFunction(grid::G, data::AbstractArray) where {d,G<:AbstractGrid{d}}
  GridFunction(grid, data, Val(length(data) ÷ length(grid)))
end
function GridFunction(f::Function, n::Int, grid::G) where {G<:AbstractGrid}
  data = Array{Float64,ndims(grid)+1}(undef, size(grid)..., n)
  for I in CartesianIndices(grid)
    data[Tuple(I)..., :] .= f(grid[I])
  end
  GridFunction(grid, data)
end
function GridFunction(grid::AbstractGrid, ::UndefInitializer, ::Val{n} = Val(1)) where n
  data = Array{Float64}(undef, size(grid)..., 1)
  GridFunction{typeof(grid), eltype(data), ndims(grid), n, typeof(data), ndims(grid) + 1}(grid, data)
end

#Base.show(io::IO, grid::GridFunction) =
#    print("GridFunction on $(length(grid)) grid points")

Base.size(f::GridFunction) = (length(f.data),)
#Base.axes(f::GridFunction) = (Base.OneTo(length(f.data)),)
Base.getindex(f::GridFunction, idx::Int) = getindex(f.data, idx)
Base.setindex!(f::GridFunction, v, idx::Int) = setindex!(f.data, v, idx)

## TODO: using this prints #undef for display(f)
#Base.getindex(f::GridFunction, idx...) = getindex(f.data, idx...)
#Base.setindex!(f::GridFunction, v, idx...) = setindex!(f.data, v, idx...)

#Base.IndexStyle(::Type{<:GridFunction}) = IndexLinear()

#component(f::GridFunction, k::Integer) = view(reshape(f.data,

Base.view(f::AbstractGridFunction{G}, idx...) where G<:MetaGrid =
    GridFunction(GridView(f.grid, indices(f.grid, idx...)), view(f.data, indices(f.grid, idx...)..., :))

grid(f::GridFunction) = f.grid


Base.:*(a::Number, b::GridFunction) = GridFunction(b.grid, a * b.data, Val(rangedim(b)))



import Colors: Gray
import Plots: plot

# TODO: add clamp option
my_plot(img::AbstractGridFunction{G,T,d,1}) where {G,T,d} =
    plot(reshape(Gray.((img .- minimum(img)) ./ (maximum(img) - minimum(img))), size(img.grid)), show=true)

## Operations

function fdfilter(f::GridFunction{G,T,d,n}, kern) where {G,T,d,n}
  A = zeros(axes(f.data))
  for k = 1:n
    vidx = ntuple(i -> i <= d ? Colon() : k, d + 1)
    imfilter!(view(A, vidx...), f.data[vidx...], kern)
  end
  GridFunction(f.grid, A)
end

derivate(f::GridFunction, dir::Integer; fd = FD.forward) = fdfilter(f, _fdkern(fd, ndims(f.grid), dir))


# currently gradient and divergence are specific implementations (as needed for
# chambolles algorithm), other choices may be reasonable

function gradient(f::GridFunction{G,T,d,1}; fd=FD.forward) where {G,T,d}
  #A = Array{T,d+1}(undef, (size(f.grid)..., d))
  A = zeros(size(f.grid)..., d)
  for k = 1:d
    vidx = ntuple(i -> i <= d ? Colon() : k, d + 1)
    vidx2 = ntuple(i -> i <= d ? Colon() : 1, d + 1)
    imfilter!(view(A, vidx...), f.data[vidx2...], bidx -> trim_kernel(_fdkern(fd, ndims(f.grid), k), bidx))
    # FIXME: hackish, only for forward fd
    selectdim(selectdim(A, d+1, k), k, size(f.grid, k)) .= 0
  end
  GridFunction(f.grid, A)
end

function divergence(f::GridFunction{G,T,d}; fd=FD.backward) where {d,G<:AbstractGrid{d},T}
  A = zeros(size(f.grid)..., 1)
  for k = 1:d
    vidx = ntuple(i -> i <= d ? Colon() : 1, d + 1)
    vidx2 = ntuple(i -> i <= d ? Colon() : k, d + 1)
    imfilter!(view(A, vidx...), f.data[vidx2...], bidx -> trim_kernel(_fdkern(fd, ndims(f.grid), k), bidx))
  end
  GridFunction(f.grid, A)
end

function pw_norm(f::GridFunction)
  d = domaindim(f)
  GridFunction(f.grid, sqrt.(sum(f.data.^2, dims=d+1)), Val(1))
end

function my_norm(f::GridFunction)
  d = domaindim(f)
  GridFunction(f.grid, sqrt.(sum(f.data.^2, dims=d+1)), Val(rangedim(f)))
end

function pwmul(f::GridFunction{G,T,d,1}, g::GridFunction{G,T,d}) where {G,T,d}
  # TODO: maybe comprehension?
  data = similar(g.data)
  for I in CartesianIndices(f.grid)
    data[Tuple(I)..., :] = f.data[Tuple(I)..., 1] * g.data[Tuple(I)..., :]
  end
  GridFunction(f.grid, data, Val(rangedim(g)))
end




"""
    neuman_kernel(kernel, bidx)

Modify `kernel` to be usable on a boundary `bixd` in a neumannn boundary setting.

Ghost nodes are eliminated through central finite differences. Currently only
the basic 5-star Laplace is supported.
"""
function neumann_kernel(kernel, bidx)
  out = copy(kernel)
  for (i, bi) in enumerate(bidx)
    colrange = (i == j ? Colon() : 0 for j in 1:ndims(kernel))
    out[colrange...] .+= bi * FD.forward
  end
  out /= 2 ^ count(!iszero, bidx)
  trim_kernel(out, bidx)
end

function divergence(grid::Grid{d}, field) where d
  field = reshape(field, (size(grid)..., ndims(grid)))
  size(field, d+1) == d || throw(ArgumentError("need a d-vectorfield on a d-grid"))
  dst = zeros(size(field))
  for k in 1:d
    kern = FD.dirfd(FD.backward, d, k)
    imfilter!(selectdim(dst, d+1, k), selectdim(field, d+1, k), bidx -> trim_kernel(kern, bidx))
  end
  out = dropdims(sum(dst, dims=d+1), dims=d+1)
  out[:]
end

function gradient(grid::Grid, img::AbstractArray)
  img = reshape(img, size(grid))
  d = ndims(img)
  dst = zeros(size(img)..., d)
  for k in 1:d
    kern = FD.dirfd(FD.forward, d, k)
    imfilter!(selectdim(dst, d+1, k), img, bidx -> trim_kernel(kern, bidx))
    # FIXME: hack
    selectdim(selectdim(dst, d+1, k), k, size(grid, k)) .= 0
  end
  reshape(dst, length(grid) * ndims(grid))
end

function my_norm(grid::Grid, img::AbstractArray)
  img = reshape(img, (size(grid)..., ndims(grid)))
  d = ndims(grid)

  out = sqrt.(sum(img.^2, dims=d+1))
  out = repeat(out, inner=ntuple(i -> i == d+1 ? size(img, d+1) : 1, d+1))

  out[:]
end


function apply(kern, flatimg, size)
  img = reshape(flatimg, size)
  out = imfilter(kern, img)
  reshape(out, prod(size))
end

function solve(pde, shape::NTuple{d, Integer}) where d
  ndims(pde.stencil) == d || throw("dimension mismatch")
  length(pde.f) == prod(shape) || throw("bounds mismatch")

  A = LinearMap{Float64}(u -> apply(pde.stencil, u, shape), prod(shape), issymmetric=true, isposdef=true)

  u = zeros(prod(shape))
  cg!(u, A, pde.f)
  return u
end

end


#module FD
#
#using OffsetArrays
#include("filtering.jl")
#
#
#"FD-Kernels"
#module Kernels
#
#using OffsetArrays
#
#const forward = OffsetVector([-1, 1], 0:1)
#const backward = OffsetVector([-1, 1], -1:0)
#const central = OffsetVector([-1, 0, 1], -1:1)
#
#end
#
#using .Filtering
#using .Kernels
#
#
#
### Implementation
#
#
### Grid
#
## aka `StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}`
#const FDStepRange = typeof(0:0.1:1)
#
#struct Grid{d}
#  domain::NTuple{d, FDStepRange}
#end
#
#Grid(domain::FDStepRange...) = Grid(domain)
#Grid(ndims::Integer, range::FDStepRange) = Grid(ntuple(i -> range, ndims))
#
#Base.show(io::IO, grid::Grid) =
#    print("FDGrid from $(first.(grid.domain)) to $(last.(grid.domain)) ",
#          "with spacing $(step.(grid.domain))")
#
## this is not type-stable?
##Base.size(grid::Grid) = getproperty.(grid.domain, :len)
#Base.size(grid::Grid) = ntuple(i -> grid.domain[i].len, ndims(grid))
#Base.getindex(grid::Grid, idx...) = SVector(getindex.(grid.domain, idx))
#
#Base.ndims(grid::Grid{d}) where d = d
#Base.axes(grid::Grid) = Base.OneTo.(size(grid))
#Base.length(grid::Grid) = prod(size(grid))
#
## TODO: there might be a better way to wrap that product iterator
#Base.iterate(grid::Grid) = iterate(Iterators.product(grid.domain...))
#Base.iterate(grid::Grid, state) = iterate(Iterators.product(grid.domain...), state)
#
#
### GridFunction
#
#struct GridFunction{T,d} <: AbstractArray{T,1}
#  grid::Grid{d}
#  data::Array{T,1}
#end
#
#Base.show(io::IO, grid::GridFunction) =
#    print("GridFunction on $(length(grid)) grid points")
#
#Base.size(f::GridFunction) = size(f.data)
#Base.getindex(f::GridFunction, idx...) = getindex(f.data, idx...)
#Base.setindex!(f::GridFunction, v, idx...) = setindex!(f.data, v, idx...)
#
#Base.IndexStyle(::Type{<:GridFunction}) = IndexLinear()
#Base.similar(f::GridFunction, ::Type{S}, size::Int...) where S =
#    GridFunction(f.grid, Array{S}(undef, size))
#
### Operations
#
#function fdfilter(f::GridFunction, kern)
#  A = imfilter(reshape(f.data, size(f.grid)), kern)
#  GridFunction(f.grid, A[:])
#end
#
#derivate(f::GridFunction, dir::Integer; fd = FD.forward) = fdfilter(f, _fdkern(fd, ndims(f.grid), dir))
#
##gradient(f::GridFunction) = GridFunction(f.grid,
#
#
#"""
#    neuman_kernel(kernel, bidx)
#
#Modify `kernel` to be usable on a boundary `bixd` in a neumannn boundary setting.
#
#Ghost nodes are eliminated through central finite differences. Currently only
#the basic 5-star Laplace is supported.
#"""
#function neumann_kernel(kernel, bidx)
#  out = copy(kernel)
#  for (i, bi) in enumerate(bidx)
#    colrange = (i == j ? Colon() : 0 for j in 1:ndims(kernel))
#    out[colrange...] .+= bi * FD.forward
#  end
#  out /= 2 ^ count(!iszero, bidx)
#  trim_kernel(out, bidx)
#end
#
#function divergence(grid::Grid{d}, field) where d
#  field = reshape(field, (size(grid)..., ndims(grid)))
#  size(field, d+1) == d || throw(ArgumentError("need a d-vectorfield on a d-grid"))
#  dst = zeros(size(field))
#  for k in 1:d
#    kern = FD.dirfd(FD.backward, d, k)
#    imfilter!(selectdim(dst, d+1, k), selectdim(field, d+1, k), kern)
#  end
#  out = dropdims(sum(dst, dims=d+1), dims=d+1)
#  reshape(out, length(grid))
#end
#
#function gradient(grid::Grid, img::AbstractArray)
#  img = reshape(img, size(grid))
#  d = ndims(img)
#  dst = zeros(size(img)..., d)
#  for k in 1:d
#    kern = FD.dirfd(FD.forward, d, k)
#    imfilter!(selectdim(dst, d+1, k), img, kern)
#  end
#  reshape(dst, length(grid) * ndims(grid))
#end
#
#end