"""
A wrapper for the interpolation methods of DataInterpolations.jl.

$(SIGNATURES)

Wraps the methods in such a way that they are callable as `f(u,t)` to
create and return an interpolation of `u` over `t`.
The first argument of the constructor always defines the interpolation method,
all following arguments will be used in the interpolation.

The additional keyword `crop = false` indicates to discard the first and last element of the time series. 

# Example

```julia
# Create the wrapper struct
itp_method = InterpolationMethod(QuadraticSpline)
# Create a callable interpolation
itp = itp_method(u,t)
# Return u[2]
itp(t[2])
```
"""
struct InterpolationMethod{T} <: AbstractInterpolationMethod
  itp::T
  args

  function InterpolationMethod(itp, args...)
    return new{typeof(itp)}(itp, args)
  end

end

(x::InterpolationMethod)(u, t) = x.itp(u,t,x.args...)

# TODO Wrap all types
# Wrap the common itps
InterpolationMethod() = InterpolationMethod(LinearInterpolation)



# Taken from DiffEqFlux
# https://github.com/SciML/DiffEqFlux.jl/blob/master/src/collocation.jl
# On 3-11-2021

struct EpanechnikovKernel <: CollocationKernel end
struct UniformKernel <: CollocationKernel end
struct TriangularKernel <: CollocationKernel end
struct QuarticKernel <: CollocationKernel end
struct TriweightKernel <: CollocationKernel end
struct TricubeKernel <: CollocationKernel end
struct GaussianKernel <: CollocationKernel end
struct CosineKernel <: CollocationKernel end
struct LogisticKernel <: CollocationKernel end
struct SigmoidKernel <: CollocationKernel end
struct SilvermanKernel <: CollocationKernel end

function calckernel(::EpanechnikovKernel,t)
    if abs(t) > 1
        return 0
    else
        return 0.75*(1-t^2)
    end
end

function calckernel(::UniformKernel,t)
    if abs(t) > 1
        return 0
    else
        return 0.5
    end
end

function calckernel(::TriangularKernel,t)
    if abs(t) > 1
        return 0
    else
        return (1-abs(t))
    end
end

function calckernel(::QuarticKernel,t)
  if abs(t)>0
    return 0
  else
    return (15*(1-t^2)^2)/16
  end
end

function calckernel(::TriweightKernel,t)
  if abs(t)>0
    return 0
  else
    return (35*(1-t^2)^3)/32
  end
end

function calckernel(::TricubeKernel,t)
  if abs(t)>0
    return 0
  else
    return (70*(1-abs(t)^3)^3)/80
  end
end

function calckernel(::GaussianKernel,t)
  exp(-0.5*t^2)/(sqrt(2*π))
end

function calckernel(::CosineKernel,t)
  if abs(t)>0
    return 0
  else
    return (π*cos(π*t/2))/4
  end
end

function calckernel(::LogisticKernel,t)
  1/(exp(t)+2+exp(-t))
end

function calckernel(::SigmoidKernel,t)
  2/(π*(exp(t)+exp(-t)))
end

function calckernel(::SilvermanKernel,t)
  sin(abs(t)/2+π/4)*0.5*exp(-abs(t)/sqrt(2))
end

function construct_t1(t,tpoints)
    hcat(ones(eltype(tpoints),length(tpoints)),tpoints.-t)
end

function construct_t2(t,tpoints)
  hcat(ones(eltype(tpoints),length(tpoints)),tpoints.-t,(tpoints.-t).^2)
end

function construct_w(t,tpoints,h,kernel)
    W = @. calckernel((kernel,),(tpoints-t)/h)/h
    Diagonal(W)
end


"""
$(SIGNATURES)

Unified interface for collocation techniques. The input can either be
a `CollocationKernel` (see list below) or a wrapped `InterpolationMethod` from
[DataInterpolations.jl](https://github.com/PumasAI/DataInterpolations.jl).

Computes a non-parametrically smoothed estimate of `u'` and `u`
given the `data`, where each column is a snapshot of the timeseries at
`tpoints[i]`.

# Examples
```julia
u′,u,t = collocate_data(data,tpoints,kernel=SigmoidKernel())
u′,u,t = collocate_data(data,tpoints,tpoints_sample,interp,args...)
u′,u,t = collocate_data(data,tpoints,interp)
```

# Collocation Kernels
See [this paper](https://www.ncbi.nlm.nih.gov/pmc/articles/PMC2631937/) for more information.
+ EpanechnikovKernel
+ UniformKernel
+ TriangularKernel
+ QuarticKernel
+ TriweightKernel
+ TricubeKernel
+ GaussianKernel
+ CosineKernel
+ LogisticKernel
+ SigmoidKernel
+ SilvermanKernel

# Interpolation Methods
See [DataInterpolations.jl](https://github.com/PumasAI/DataInterpolations.jl) for more information.
+ ConstantInterpolation
+ LinearInterpolation
+ QuadraticInterpolation
+ LagrangeInterpolation
+ QuadraticSpline
+ CubicSpline
+ BSplineInterpolation
+ BSplineApprox
+ Curvefit
"""
function collocate_data(data,tpoints,kernel=TriangularKernel(); crop = false, kwargs...)
  _one = oneunit(first(data))
  _zero = zero(first(data))
  e1 = [_one;_zero]
  e2 = [_zero;_one;_zero]
  n = length(tpoints)
  h = (n^(-1/5))*(n^(-3/35))*((log(n))^(-1/16))

  Wd = similar(data, n, size(data,1))
  WT1 = similar(data, n, 2)
  WT2 = similar(data, n, 3)
  x = map(tpoints) do _t
      T1 = construct_t1(_t,tpoints)
      T2 = construct_t2(_t,tpoints)
      W = construct_w(_t,tpoints,h,kernel)
      mul!(Wd,W,data')
      mul!(WT1,W,T1)
      mul!(WT2,W,T2)
      (e2'*((T2'*WT2)\T2'))*Wd,(e1'*((T1'*WT1)\T1'))*Wd
  end
  estimated_derivative = reduce(hcat,transpose.(first.(x)))
  estimated_solution = reduce(hcat,transpose.(last.(x)))
  crop && return estimated_derivative[:,2:end-1], estimated_derivative[:, 2:end-1], tpoints[2:end-1]
  estimated_derivative,estimated_solution, tpoints
end


# Adapted to dispatch on InterpolationMethod
collocate_data(data, tpoints, interp::InterpolationMethod; kwargs...) = collocate_data(data, tpoints, tpoints, interp; kwargs...)

function collocate_data(data::AbstractVector,tpoints::AbstractVector,tpoints_sample::AbstractVector,
                        interp::InterpolationMethod; kwargs...)
  u, du, tpoints = collocate_data(reshape(data, 1, :),tpoints,tpoints_sample,interp; kwargs...)
  return du[1, :], u[1, :] , tpoints
end

# Adapted to dispatch on InterpolationMethod
function collocate_data(data::AbstractMatrix{T},tpoints::AbstractVector{T},
                        tpoints_sample::AbstractVector{T},interp::InterpolationMethod; crop = false, kwargs...) where T

  u = zeros(T,size(data, 1),length(tpoints_sample))
  du = zeros(T,size(data, 1),length(tpoints_sample))
  for d1 in 1:size(data,1)
    interpolation = interp(data[d1,:],tpoints)
    u[d1,:] .= interpolation.(tpoints_sample)
    du[d1,:] .= DataInterpolations.derivative.((interpolation,), tpoints_sample)
  end
  crop && return du[:, 2:end-1], u[:, 2:end-1], tpoints[2:end-1]
  return du, u, tpoints
end
