"""
Represent any source (incident) wave

Subtypes may have a symmetry (such as [`PlaneSource`](@ref)) and will contain information about physical medium.
"""
abstract type AbstractSource{T,P} end

"""
    PlaneSource(medium::P, amplitude::SVector, direction::SVector)

Is a struct type which describes a plane-wave source that drives/forces the whole system. It has three fields: a physical `medium`, an `amplitude` of the source, and the direction the propagate in `direction`.

For any given angular frequency ω, the PlaneSource has the value ``e^{i ω/c \\mathbf v \\cdot \\mathbf x }`` at the point ``\\mathbf x``, where ``c`` is the medium wavespeed and ``\\mathbf v`` is the direction.
"""
struct PlaneSource{T<:AbstractFloat,Dim,FieldDim,P<:PhysicalMedium} <: AbstractSource{T,P}
    medium::P
    direction::SVector{Dim,T}
    position::SVector{Dim,T}
    amplitude::SVector{FieldDim,Complex{T}}
    # Check that P has same Dim and FieldDim
    function PlaneSource{T,Dim,FieldDim,P}(medium::P, direction::AbstractArray{T}, position::AbstractArray{T}, amplitude::AbstractArray{CT} where CT<:Union{T,Complex{T}}) where {T,Dim,FieldDim,P<:PhysicalMedium{T,Dim,FieldDim}}
        normw = sqrt(sum(direction .^2)) # note if direction is complex this is different from norm(direction)
        if !(normw ≈ one(T))
            @warn "The direction will be normalised so that sum(direction .^2) == 1.0"
        end

        if length(direction) != Dim || length(position) != Dim
            throw(DimensionMismatch("The spatial dimensions of the medium do not match dimensions of the direction or position vector."))
        elseif length(amplitude) != FieldDim
            throw(DimensionMismatch("The dimensions of the field in the medium (i.e. always 1 for acoustics) do not match the dimension of the amplitude vector."))
        else
            new{T,Dim,FieldDim,P}(medium,direction ./ normw, position, amplitude)
        end
    end
end

field(s::PlaneSource{T}) where T = function (x::AbstractArray{T}, ω::T)
        s.amplitude * exp(im * (ω / s.medium.c) * dot(x - s.position,s.direction))
    end

field(s::PlaneSource{T,Dim,1}) where {T, Dim} = function (x::AbstractArray{T}, ω::T)
        s.amplitude[1] * exp(im * (ω / s.medium.c) * dot(x - s.position,s.direction))
    end

field(s::PlaneSource{T}, x::AbstractArray{T}, ω::T) where T = field(s)(x,ω)

# Convenience constructor which does not require explicit types/parameters
function PlaneSource(medium::PhysicalMedium{T,Dim,FieldDim},
        direction::AbstractArray{T},
        position::AbstractArray{T} = zeros(T,Dim),
        amplitude::Union{CT,AbstractArray{CT}} where CT<:Union{T,Complex{T}} = SVector{FieldDim}(ones(Complex{T},FieldDim))
    ) where {T,Dim,FieldDim}

    if !(typeof(amplitude) <: AbstractArray)
        amplitude = [amplitude]
    end

    PlaneSource{T,Dim,FieldDim,typeof(medium)}(medium,direction,position,amplitude)
end

function PlaneSource(medium::PhysicalMedium{T,Dim,FieldDim};
        direction::AbstractArray{T} = [one(T); zeros(T,Dim-1)],
        position::AbstractArray{T} = zeros(T,Dim),
        amplitude::Union{CT,AbstractArray{CT}} where CT<:Union{T,Complex{T}} = SVector{FieldDim}(ones(Complex{T},FieldDim))
    ) where {T,Dim,FieldDim}

    PlaneSource(medium,direction,position,amplitude)
end

Symmetry(s::PlaneSource{T,Dim}) where {T,Dim} = (abs(dot(s.direction,azimuthalnormal(Dim))) == norm(s.direction)) ? PlanarAzimuthalSymmetry{Dim}() : PlanarSymmetry{Dim}()

"""
    RegularSource(medium::P, field::Function, coef::Function)

Is a struct type which describes the source field that drives/forces the whole system. It is also known as an incident wave. It has three fields `RegularSource.medium`, `RegularSource.field`, and `RegularSource.coef`.

The source field at the position 'x' and angular frequency 'ω' is given by
```julia-repl
x = [1.0,0.0]
ω = 1.0
RegularSource.field(x,ω)
```

The field `RegularSource.coef`
regular_basis_function(medium::Acoustic{T,2}, ω::T)
"""
struct RegularSource{T<:AbstractFloat,P<:PhysicalMedium,S<:AbstractSymmetry} <: AbstractSource{T,P}
    medium::P
    "Use: field(x,ω)"
    field::Function
    "Use: coefficients(n,x,ω)"
    coefficients::Function
    # Enforce that the Types are the same
    function RegularSource{T,P,S}(medium::P,field::Function,coefficients::Function) where {T,Dim,FieldDim,P<:PhysicalMedium{T,Dim,FieldDim},S<:AbstractSymmetry{Dim}}
        s = new{T,P,S}(medium,field,coefficients)
        self_test(s)
        return s
    end
end

function RegularSource(medium::P,field::Function,coefficients::Function; symmetry::AbstractSymmetry{Dim} = WithoutSymmetry{Dim}()) where {T,Dim,FieldDim,P<:PhysicalMedium{T,Dim,FieldDim}}
    return RegularSource{T,P,typeof(symmetry)}(medium,field,coefficients)
end

Symmetry(reg::RegularSource{T,P,S}) where {T,P,S} = S()

"""
    regular_spherical_coefficients(source::RegularSource)

return a function which can calculate the coefficients of a regular spherical wave basis.
"""
function regular_spherical_coefficients(source::RegularSource)
    source.coefficients
end

"""
Check that the source functions return the correct types
"""
function self_test(source::RegularSource{T,P}) where {P,T}
    ω = one(T)

    # choose rand postion, hopefully not the source position/origin
    x = SVector(rand(T.(0.1:0.1:1.0),spatial_dimension(P))...)

    # Check that the result of field has same dimension and type as PhysicalProperty field
    if field_dimension(P) == 1
        source.field(x,ω)::Union{Complex{T},SVector{field_dimension(P),Complex{T}}}
    else
        source.field(x,ω)::SVector{field_dimension(P),Complex{T}}
    end

    # Check that the result of field has the same dimension as field dimension of P
    if field_dimension(P) == 1
        source.coefficients(1,x,ω)::Union{Vector{Complex{T}},Vector{SVector{field_dimension(P),Complex{T}}}}
    # else # this else is not necessarily correct..
    #     source.coefficients(1,x,ω)::SVector{field_dimension(P),Complex{T}}
    end

    return true
end

field(s::RegularSource) = s.field
field(s::RegularSource{T},x::AbstractArray{T}, ω::T) where T = field(s)(x,ω)

function constant_source(medium::P, num::Complex{T} = zero(Float64) * im) where {P,T}
    return RegularSource{T,P}(medium, (x,ω) -> num, (order,x,ω) -> [num])
end

"""
    regular_spherical_source(PhysicalMedium,regular_coefficients; amplitude = one(T), position = zeros(T,Dim))

``c_n = ```regular_coefficients[n]`, ``x_o=```position`, and let ``v_n(kx)`` represent the regular spherical basis with wavenumber ``k`` at position ``x``. The above returns a `source` where `source.field(x,ω) =```\\sum_n c_n v_n(kx -k x_0)``
"""
function regular_spherical_source(medium::PhysicalMedium{T,Dim},regular_coefficients::AbstractVector{CT};
        symmetry::AbstractSymmetry{Dim} = WithoutSymmetry{Dim}(),
        amplitude::Union{T,Complex{T}} = one(T),
        position::AbstractArray{T} = SVector(zeros(T,Dim)...)) where {T,Dim,CT<:Union{T,Complex{T}}}

    coeff_order = basislength_to_basisorder(typeof(medium),length(regular_coefficients))

    function source_field(x,ω)
        vs = regular_basis_function(medium, ω)(coeff_order,x-position)
        return amplitude * sum(vs .* regular_coefficients)

    end

    function source_coef(order,centre,ω)
        k = ω / medium.c

        V = regular_translation_matrix(medium, max(order,coeff_order), ω, centre - position)

        len2 = basisorder_to_basislength(typeof(medium),order)
        len1 = length(regular_coefficients)

        return amplitude .* sum(regular_coefficients[n] .* V[n,1:len2] for n in 1:len1)
    end

    return RegularSource(medium, source_field, source_coef; symmetry = symmetry)
end

"""
    source_expand(RegularSource, centre; basis_order = 4)

Returns a function of `(x,ω)` which approximates the value of the source at `(x,ω)`. That is, the source is written in terms of a regular basis expansion centred at `centre`.
"""
function source_expand(source::AbstractSource{T}, centre::AbstractVector{T}; basis_order::Int = 4) where T

    # Convert to SVector for efficiency and consistency
    centre = SVector(centre...)

    return function (x::AbstractVector{T}, ω::T)
        vs = regular_basis_function(source.medium, ω)
        regular_coefficients = regular_spherical_coefficients(source)

        sum(regular_coefficients(basis_order,centre,ω) .* vs(basis_order, x - centre))
    end
end

import Base.(+)
function +(s1::RegularSource{T,P},s2::RegularSource{T,P})::RegularSource{T,P} where {P,T}
    if s1.medium != s2.medium
        error("Can not add sources from different physical mediums.")
    end
    field(x,ω) = s1.field(x,ω) + s2.field(x,ω)
    coef(n,centre,ω) = s1.coefficients(n,centre,ω) + s2.coefficients(n,centre,ω)

    sym1 = Symmetry(s1)
    sym2 = Symmetry(s2)
    S = typeof(Symmetry(sym1,sym2))

    RegularSource{T,P,S}(s1.medium,field,coef)
end

import Base.(*)
function *(a,s::RegularSource{T,P,S})::RegularSource{T,P,S} where {P,T,S}
    if typeof(one(Complex{T})*a) != Complex{T}
        error("Multiplying source by $a would cause source type to change, please explicitly cast $a to same type as RegularSource ($T)")
    end

    field(x,ω) = a * s.field(x,ω)
    coef(n,centre,ω) = a .* s.coefficients(n,centre,ω)

    RegularSource{T,P,S}(s.medium,field,coef)
end

*(s::RegularSource,a) = *(a,s)
