# The `NetworkLoss` methods below dispatch on GML's own architectures, which is what keeps them out
# of piracy. A `NetworkLoss(::NeuralNetwork)` forwarding to `architecture(nn)` would be piracy on
# both counts, and belongs in `AbstractNeuralNetworks` itself.

@doc raw"""
    TransformerLoss(seq_length, prediction_window)

Make an instance of the transformer loss.

This should be used together with a neural network of type [`TransformerIntegrator`](@ref).

# Example

`TransformerLoss` applies a neural network to an input and compares it to the `output` via an ``L_2`` norm:

```jldoctest
using GeometricMachineLearning
using LinearAlgebra: norm
import Random

const d = 2
const seq_length = 3
const prediction_window = 2

Random.seed!(123)
arch = StandardTransformerIntegrator(d)
nn = NeuralNetwork(arch)

input_mat =  [1. 2. 3.; 4. 5. 6.]
output_mat = [1. 2.; 3. 4.]
loss = TransformerLoss(seq_length, prediction_window)

# start of prediction
const sop = seq_length - prediction_window + 1
loss(nn, input_mat, output_mat) ≈ norm(output_mat - nn(input_mat)[:, sop:end]) / norm(output_mat)

# output

true
```

So `TransformerLoss` simply does:

```math
    \mathtt{loss}(\mathcal{NN}, \mathtt{input}, \mathtt{output}) = || \mathcal{NN}(\mathtt{input})[(\mathtt{sl} - \mathtt{pw} + 1):\mathtt{end}] - \mathtt{output} || / || \mathtt{output} ||,
```
where ``||\cdot||`` is the ``L_2`` norm.

# Parameters

The `prediction_window` specifies how many time steps are predicted into the future.
It defaults to the value specified for `seq_length`.
"""
struct TransformerLoss <: NetworkLoss
    seq_length::Int
    prediction_window::Int
end

TransformerLoss(seq_length::Int) = TransformerLoss(seq_length, seq_length)

# The loss a `Batch{:Transformer}` implies: `NetworkLoss(arch, batch)` reads the sequence length and
# the prediction window off the batch instead of taking the two numbers.
function TransformerLoss(batch::Batch{:Transformer})
    TransformerLoss(batch.seq_length, batch.prediction_window)
end

# This crops the output array of the neural network so that it conforms with the output it should be compared to. This is needed for the transformer loss.
function crop_array_for_transformer_loss(nn_output::AT,
        output::BT) where {T, T2, AT <: AbstractArray{T, 3}, BT <: AbstractArray{T2, 3}}
    @view nn_output[
        axes(output, 1), axes(output, 2) .+ size(nn_output, 2) .- size(output, 2),
        axes(output, 3)]
end

function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
        ps::NetworkParameters, input::AT,
        output::AT) where {T, AT <: AbstractArray{T, 3}}
    input_dim, input_seq_length = size(input)
    output_dim, output_prediction_window = size(output)
    @assert input_dim == output_dim
    @assert input_seq_length == loss.seq_length
    @assert output_prediction_window == loss.prediction_window

    predicted_output_uncropped = model(input, ps)
    predicted_output_cropped = crop_array_for_transformer_loss(
        predicted_output_uncropped, output)
    _compute_loss(predicted_output_cropped, output)
end

function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
        ps::NetworkParameters, input::AT,
        output::AT) where {T, AT <: AbstractArray{T, 2}}
    loss(model, ps, reshape(input, size(input)..., 1), reshape(output, size(output)..., 1))
end

function (loss::TransformerLoss)(model::Union{Chain, AbstractExplicitLayer},
        ps::NetworkParameters,
        input::NamedTupleOfArrays{<:Any, (:q, :p)},
        output::NamedTupleOfArrays{<:Any, (:q, :p)})
    loss(model, ps, vcat(input.q, input.p), vcat(output.q, output.p))
end

@doc raw"""
    AutoEncoderLoss()

Make an instance of `AutoEncoderLoss`.

This loss should always be used together with a neural network of type [`AutoEncoder`](@ref) (and it is also the default for training such a network).

# Example

`AutoEncoderLoss` applies a neural network to an input and compares it to the input via an ``L_2`` norm:

```jldoctest
using GeometricMachineLearning
using LinearAlgebra: norm
import Random
Random.seed!(123)

const N = 4
const n = 1
arch = SymplecticAutoencoder(2*N, 2*n)
nn = NeuralNetwork(arch)

input_vec =  [1., 2., 3., 4., 5., 6., 7., 8.]
loss = AutoEncoderLoss()

loss(nn, input_vec) ≈ norm(input_vec - nn(input_vec)) / norm(input_vec)

# output

true
```

So `AutoEncoderLoss` simply does:

```math
    \mathtt{loss}(\mathcal{NN}, \mathtt{input}) = || \mathcal{NN}(\mathtt{input}) - \mathtt{input} || / || \mathtt{input} ||,
```
where ``||\cdot||`` is the ``L_2`` norm.

# Parameters

This loss does not have any parameters. The four-argument method reconstructs its input and
ignores the target it is given, so an autoencoder is trained on the input alone.
"""
struct AutoEncoderLoss <: NetworkLoss end

AbstractNeuralNetworks.NetworkLoss(::NeuralNetworkIntegrator) = FeedForwardLoss()
AbstractNeuralNetworks.NetworkLoss(::AutoEncoder) = AutoEncoderLoss()

function (loss::AutoEncoderLoss)(nn::NeuralNetwork, input::ArrayOrNamedTuple)
    loss(nn.model, params(nn), input)
end

function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer},
        ps::NetworkParameters, input::ArrayOrNamedTuple)
    FeedForwardLoss()(model, ps, input, input)
end

function (loss::AutoEncoderLoss)(model::Union{Chain, AbstractExplicitLayer},
        ps::NetworkParameters, input::ArrayOrNamedTuple, output::ArrayOrNamedTuple)
    loss(model, ps, input)
end

@doc raw"""
    ReducedLoss(encoder, decoder)

Make an instance of `ReducedLoss` based on an [`Encoder`](@ref) and a [`Decoder`](@ref).

This loss should be used together with a [`NeuralNetworkIntegrator`](@ref) or [`TransformerIntegrator`](@ref).

# Example

`ReducedLoss` applies the *encoder*, *integrator* and *decoder* neural networks in this order to an input and compares it to the `output` via an ``L_2`` norm:

```jldoctest
using GeometricMachineLearning
using LinearAlgebra: norm
import Random
Random.seed!(123)

const N = 4
const n = 1

Ψᵉ = NeuralNetwork(Chain(Dense(N, n), Dense(n, n))) |> encoder
Ψᵈ = NeuralNetwork(Chain(Dense(n, n), Dense(n, N))) |> decoder
transformer = NeuralNetwork(StandardTransformerIntegrator(n))

input_mat =  [1.  2.;  3.  4.;  5.  6.;  7.  8.]
output_mat = [9. 10.; 11. 12.; 13. 14.; 15. 16.]
loss = ReducedLoss(Ψᵉ, Ψᵈ)

output_prediction = Ψᵈ(transformer(Ψᵉ(input_mat)))
loss(transformer, input_mat, output_mat) ≈ norm(output_mat - output_prediction) / norm(output_mat)

# output

true
```

So the loss computes:

```math
\mathrm{loss}_{\mathcal{E}, \mathcal{D}}(\mathcal{NN}, \mathrm{input}, \mathrm{output}) = ||\mathcal{D}(\mathcal{NN}(\mathcal{E}(\mathrm{input}))) - \mathrm{output}||,
```
where ``\mathcal{E}`` is the [`Encoder`](@ref), ``\mathcal{D}`` is the [`Decoder`](@ref).
``\mathcal{NN}`` is the neural network we compute the loss of.
"""
struct ReducedLoss{ET <: NeuralNetwork{<:Encoder}, DT <: NeuralNetwork{<:Decoder}} <:
       NetworkLoss
    encoder::ET
    decoder::DT
end

"""
    ReducedLoss(autoencoder)

Make an instance of `ReducedLoss` based on a neural network of type [`AutoEncoder`](@ref).

Internally this does:
```julia
ReducedLoss(encoder(autoencoder), decoder(autoencoder))
```
so it calls the functions [`encoder`](@ref) and [`decoder`](@ref).
"""
function ReducedLoss(autoencoder::NeuralNetwork{<:AutoEncoder})
    ReducedLoss(encoder(autoencoder), decoder(autoencoder))
end

function (loss::ReducedLoss)(model::Chain, ps::NetworkParameters,
        input::ArrayOrNamedTuple, output::ArrayOrNamedTuple)
    _compute_loss(loss.decoder(model(loss.encoder(input), ps)), output)
end

@doc raw"""
    ParametricLoss()

Make an instance of the loss of a network whose forward pass takes the parameters ``\mu`` of the
system beside its input.

`ParametricLoss()(model, ps, input, output, μ)` applies the network with
[`apply_parametric`](@ref) and compares the result with `output` in the relative ``L_2`` norm:

```math
\mathrm{loss}(\mathcal{NN}, \mathrm{input}, \mathrm{output}, \mu) = ||\mathcal{NN}(\mathrm{input}, \mu) - \mathrm{output}|| / ||\mathrm{output}||.
```

This is `FeedForwardLoss` with the system parameters passed to the network. `μ` is one
`OptionalParameters` for one sample, or a vector of them, one per sample of a batch, as a
`ParametricDataLoader` draws it. `train!` on a `ParametricDataLoader` trains on this loss
when no `loss` is given.
"""
struct ParametricLoss <: NetworkLoss end

function (loss::ParametricLoss)(model::Chain, ps::NetworkParameters,
        input::ArrayOrNamedTuple, output::ArrayOrNamedTuple,
        μ::Union{OptionalParameters, AbstractVector})
    _compute_loss(apply_parametric(model, input, μ, ps), output)
end