# Convenience type to allow dispatch on solution objects that were created by Trixi.jl
#
# This is a union of a Trixi.jl-specific SciMLBase.ODESolution and of Trixi.jl's own
# TimeIntegratorSolution.
#
#! format: off
const TrixiODESolution = Union{ODESolution{T, N, uType, uType2, DType, tType, rateType, discType, P} where
    {T, N, uType, uType2, DType, tType, rateType, discType, P<:ODEProblem{uType_, tType_, isinplace, P_, F_} where
     {uType_, tType_, isinplace, P_<:AbstractSemidiscretization, F_}}, TimeIntegratorSolution}
#! format: on

# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin
#! format: noindent

# This file holds plotting types which can be used for both Plots.jl and Makie.jl.

# This abstract type is used to derive PlotData types of different dimensions; but still allows to share some functions for them.
abstract type AbstractPlotData{NDIMS} end

# Define additional methods for convenience.
# These are defined for AbstractPlotData, so they can be used for all types of plot data.
Base.firstindex(pd::AbstractPlotData) = first(pd.variable_names)
Base.lastindex(pd::AbstractPlotData) = last(pd.variable_names)
Base.length(pd::AbstractPlotData) = length(pd.variable_names)
Base.size(pd::AbstractPlotData) = (length(pd),)
Base.keys(pd::AbstractPlotData) = tuple(pd.variable_names...)

function Base.iterate(pd::AbstractPlotData, state = 1)
    if state > length(pd)
        return nothing
    else
        return (pd.variable_names[state] => pd[pd.variable_names[state]], state + 1)
    end
end

"""
    Base.getindex(pd::AbstractPlotData, variable_name)

Extract a single variable `variable_name` from `pd` for plotting with `Plots.plot`.
"""
function Base.getindex(pd::AbstractPlotData, variable_name)
    variable_id = findfirst(isequal(variable_name), pd.variable_names)

    if isnothing(variable_id)
        throw(KeyError(variable_name))
    end

    return PlotDataSeries(pd, variable_id)
end

Base.eltype(pd::AbstractPlotData) = Pair{String, PlotDataSeries{typeof(pd)}}

"""
    PlotData2D

Holds all relevant data for creating 2D plots of multiple solution variables and to visualize the
mesh.
"""
struct PlotData2DCartesian{Coordinates, Data, VariableNames, Vertices} <:
       AbstractPlotData{2}
    x::Coordinates
    y::Coordinates
    data::Data
    variable_names::VariableNames
    mesh_vertices_x::Vertices
    mesh_vertices_y::Vertices
    orientation_x::Int
    orientation_y::Int
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pd::PlotData2DCartesian)
    @nospecialize pd # reduce precompilation time

    print(io, "PlotData2DCartesian{",
          typeof(pd.x), ",",
          typeof(pd.data), ",",
          typeof(pd.variable_names), ",",
          typeof(pd.mesh_vertices_x),
          "}(<x>, <y>, <data>, <variable_names>, <mesh_vertices_x>, <mesh_vertices_y>)")
    return nothing
end

# holds plotting information for UnstructuredMesh2D and DGMulti-compatible meshes
struct PlotData2DTriangulated{DataType, NodeType, FaceNodeType, FaceDataType,
                              VariableNames, PlottingTriangulation} <:
       AbstractPlotData{2}
    x::NodeType # physical nodal coordinates, size (num_plotting_nodes x num_elements)
    y::NodeType
    data::DataType
    t::PlottingTriangulation
    x_face::FaceNodeType
    y_face::FaceNodeType
    face_data::FaceDataType
    variable_names::VariableNames
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pd::PlotData2DTriangulated)
    @nospecialize pd # reduce precompilation time

    print(io, "PlotData2DTriangulated{",
          typeof(pd.x), ", ",
          typeof(pd.data), ", ",
          typeof(pd.x_face), ", ",
          typeof(pd.face_data), ", ",
          typeof(pd.variable_names),
          "}(<x>, <y>, <data>, <plot_triangulation>, <x_face>, <y_face>, <face_data>, <variable_names>)")
    return nothing
end

"""
    PlotData1D

Holds all relevant data for creating 1D plots of multiple solution variables and to visualize the
mesh.
"""
struct PlotData1D{Coordinates, Data, VariableNames, Vertices} <: AbstractPlotData{1}
    x::Coordinates
    data::Data
    variable_names::VariableNames
    mesh_vertices_x::Vertices
    orientation_x::Integer
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pd::PlotData1D)
    print(io, "PlotData1D{",
          typeof(pd.x), ",",
          typeof(pd.data), ",",
          typeof(pd.variable_names), ",",
          typeof(pd.mesh_vertices_x),
          "}(<x>, <data>, <variable_names>, <mesh_vertices_x>)")
    return nothing
end

# Auxiliary data structure for visualizing a single variable
struct PlotDataSeries{PD <: AbstractPlotData}
    plot_data::PD
    variable_id::Int
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pds::PlotDataSeries)
    @nospecialize pds # reduce precompilation time

    print(io, "PlotDataSeries{", typeof(pds.plot_data), "}(<plot_data>, ",
          pds.variable_id, ")")
    return nothing
end

# Generic PlotMesh wrapper type.
struct PlotMesh{PD <: AbstractPlotData}
    plot_data::PD
end

# Show only a truncated output for convenience (the full data does not make sense)
function Base.show(io::IO, pm::PlotMesh)
    @nospecialize pm # reduce precompilation time

    print(io, "PlotMesh{", typeof(pm.plot_data), "}(<plot_data>)")
    return nothing
end

"""
    getmesh(pd::AbstractPlotData)

Extract grid lines from `pd` for plotting with `Plots.plot`.
"""
getmesh(pd::AbstractPlotData) = PlotMesh(pd)

"""
    PlotData2D(u, semi [or mesh, equations, solver, cache];
               solution_variables=nothing,
               grid_lines=true, max_supported_level=11, nvisnodes=nothing,
               slice=:xy, point=(0.0, 0.0, 0.0))

Create a new `PlotData2D` object that can be used for visualizing 2D/3D DGSEM solution data array
`u` with `Plots.jl`. All relevant geometrical information is extracted from the semidiscretization
`semi`. By default, the primitive variables (if existent) or the conservative variables (otherwise)
from the solution are used for plotting. This can be changed by passing an appropriate conversion
function to `solution_variables`.

For coupled semidiscretizations, i.e., `semi isa` [`SemidiscretizationCoupled`](@ref) a vector of
`PlotData2D` objects is returned, one for each semidiscretization which is part of the
coupled semidiscretization.

If `grid_lines` is `true`, also extract grid vertices for visualizing the mesh. The output
resolution is indirectly set via `max_supported_level`: all data is interpolated to
`2^max_supported_level` uniformly distributed points in each spatial direction, also setting the
maximum allowed refinement level in the solution. `nvisnodes` specifies the number of visualization
nodes to be used. If it is `nothing`, twice the number of solution DG nodes are used for
visualization, and if set to `0`, exactly the number of nodes in the DG elements are used.

When visualizing data from a three-dimensional simulation, a 2D slice is extracted for plotting.
`slice` specifies the plane that is being sliced and may be `:xy`, `:xz`, or `:yz`.
The slice position is specified by a `point` that lies on it, which defaults to `(0.0, 0.0, 0.0)`.
Both of these values are ignored when visualizing 2D data.

# Examples
```julia
julia> using Trixi, Plots

julia> trixi_include(default_example())
[...]

julia> pd = PlotData2D(sol)
PlotData2D(...)

julia> plot(pd) # To plot all available variables

julia> plot(pd["scalar"]) # To plot only a single variable

julia> plot!(getmesh(pd)) # To add grid lines to the plot
```
"""
function PlotData2D(u_ode, semi; kwargs...)
    return PlotData2D(wrap_array_native(u_ode, semi),
                      mesh_equations_solver_cache(semi)...;
                      kwargs...)
end

function PlotData2D(u_ode, semi::SemidiscretizationCoupled; kwargs...)
    plot_data_array = []
    @unpack semis = semi

    foreach_enumerate(semis) do (i, semi_)
        u_loc = get_system_u_ode(u_ode, i, semi)
        u_loc_wrapped = wrap_array_native(u_loc, semi_)

        return push!(plot_data_array,
                     PlotData2D(u_loc_wrapped,
                                mesh_equations_solver_cache(semi_)...;
                                kwargs...))
    end

    return plot_data_array
end

# Redirect `PlotDataTriangulated2D` constructor.
function PlotData2DTriangulated(u_ode, semi; kwargs...)
    return PlotData2DTriangulated(wrap_array_native(u_ode, semi),
                                  mesh_equations_solver_cache(semi)...;
                                  kwargs...)
end

# Create a PlotData2DCartesian object for TreeMeshes on default.
function PlotData2D(u, mesh::TreeMesh, equations, solver, cache; kwargs...)
    return PlotData2DCartesian(u, mesh::TreeMesh, equations, solver, cache; kwargs...)
end

# Create a PlotData2DTriangulated object for any type of mesh other than the TreeMesh.
function PlotData2D(u, mesh, equations, solver, cache; kwargs...)
    return PlotData2DTriangulated(u, mesh, equations, solver, cache; kwargs...)
end

# Create a PlotData2DCartesian for a TreeMesh.
function PlotData2DCartesian(u, mesh::TreeMesh, equations, solver, cache;
                             solution_variables = nothing,
                             grid_lines = true, max_supported_level = 11,
                             nvisnodes = nothing,
                             slice = :xy, point = (0.0, 0.0, 0.0))
    @assert ndims(mesh) in (2, 3) "unsupported number of dimensions $ndims (must be 2 or 3)"
    solution_variables_ = digest_solution_variables(equations, solution_variables)

    # Extract mesh info
    center_level_0 = mesh.tree.center_level_0
    length_level_0 = mesh.tree.length_level_0
    leaf_cell_ids = leaf_cells(mesh.tree)
    coordinates = mesh.tree.coordinates[:, leaf_cell_ids]
    levels = mesh.tree.levels[leaf_cell_ids]

    unstructured_data = get_unstructured_data(u, solution_variables_, mesh, equations,
                                              solver, cache)
    x, y, data, mesh_vertices_x, mesh_vertices_y = get_data_2d(center_level_0,
                                                               length_level_0,
                                                               leaf_cell_ids,
                                                               coordinates, levels,
                                                               ndims(mesh),
                                                               unstructured_data,
                                                               nnodes(solver),
                                                               grid_lines,
                                                               max_supported_level,
                                                               nvisnodes,
                                                               slice, point)
    variable_names = SVector(varnames(solution_variables_, equations))

    orientation_x, orientation_y = _get_orientations(mesh, slice)

    return PlotData2DCartesian(x, y, data, variable_names, mesh_vertices_x,
                               mesh_vertices_y,
                               orientation_x, orientation_y)
end

"""
    PlotData2D(sol; kwargs...)

Create a `PlotData2D` object from a solution object created by either `OrdinaryDiffEq.solve!` (which
returns a `SciMLBase.ODESolution`) or Trixi.jl's own `solve!` (which returns a
`TimeIntegratorSolution`).
"""
function PlotData2D(sol::TrixiODESolution; kwargs...)
    return PlotData2D(sol.u[end], sol.prob.p; kwargs...)
end

# Also redirect when using PlotData2DTriangulate.
function PlotData2DTriangulated(sol::TrixiODESolution; kwargs...)
    return PlotData2DTriangulated(sol.u[end], sol.prob.p; kwargs...)
end

# If `u` is an `Array{<:SVectors}` and not a `StructArray`, convert it to a `StructArray` first.
function PlotData2D(u::Array{<:SVector}, mesh, equations, dg::DGMulti, cache;
                    solution_variables = nothing, nvisnodes = 2 * nnodes(dg))
    nvars = length(first(u))
    u_structarray = StructArray{eltype(u)}(ntuple(_ -> zeros(eltype(first(u)), size(u)),
                                                  nvars))
    for (i, u_i) in enumerate(u)
        u_structarray[i] = u_i
    end

    # re-dispatch to PlotData2D with mesh, equations, dg, cache arguments
    return PlotData2D(u_structarray, mesh, equations, dg, cache;
                      solution_variables = solution_variables, nvisnodes = nvisnodes)
end

# constructor which returns an `PlotData2DTriangulated` object.
function PlotData2D(u::StructArray, mesh, equations, dg::DGMulti, cache;
                    solution_variables = nothing, nvisnodes = 2 * nnodes(dg))
    rd = dg.basis
    md = mesh.md

    # Vp = the interpolation matrix from nodal points to plotting points
    @unpack Vp = rd
    interpolate_to_plotting_points!(out, x) = mul!(out, Vp, x)

    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names = SVector(varnames(solution_variables_, equations))

    if Vp isa UniformScaling
        num_plotting_points = size(u, 1)
    else
        num_plotting_points = size(Vp, 1)
    end
    nvars = nvariables(equations)
    uEltype = eltype(first(u))
    u_plot = StructArray{SVector{nvars, uEltype}}(ntuple(_ -> zeros(uEltype,
                                                                    num_plotting_points,
                                                                    md.num_elements),
                                                         nvars))

    for e in eachelement(mesh, dg, cache)
        # interpolate solution to plotting nodes element-by-element
        StructArrays.foreachfield(interpolate_to_plotting_points!, view(u_plot, :, e),
                                  view(u, :, e))

        # transform nodal values of the solution according to `solution_variables`
        transform_to_solution_variables!(view(u_plot, :, e), solution_variables_,
                                         equations)
    end

    # interpolate nodal coordinates to plotting points
    x_plot, y_plot = map(x -> Vp * x, md.xyz) # md.xyz is a tuple of arrays containing nodal coordinates

    # construct a triangulation of the reference plotting nodes
    t = reference_plotting_triangulation(rd.rstp) # rd.rstp = reference coordinates of plotting points

    x_face, y_face, face_data = mesh_plotting_wireframe(u, mesh, equations, dg, cache;
                                                        nvisnodes = nvisnodes)

    return PlotData2DTriangulated(x_plot, y_plot, u_plot, t, x_face, y_face, face_data,
                                  variable_names)
end

# One can also call the `PlotData2DTriangulated` constructor directly for `DGMulti`
function PlotData2DTriangulated(u, mesh, equations, dg::DGMulti, cache;
                                solution_variables = nothing,
                                nvisnodes = 2 * nnodes(dg))
    return PlotData2D(u, mesh, equations, dg, cache;
                      solution_variables = solution_variables,
                      nvisnodes = nvisnodes)
end

# specializes the PlotData2D constructor to return an PlotData2DTriangulated for any type of mesh.
function PlotData2DTriangulated(u, mesh, equations, dg::DGSEM, cache;
                                solution_variables = nothing,
                                nvisnodes = 2 * polydeg(dg))
    @assert ndims(mesh)==2 "Input must be two-dimensional."

    n_nodes_2d = nnodes(dg)^ndims(mesh)
    n_elements = nelements(dg, cache)

    # build nodes on reference element (seems to be the right ordering)
    r, s = reference_node_coordinates_2d(dg)

    # reference plotting nodes
    if nvisnodes == 0 || nvisnodes === nothing
        nvisnodes = polydeg(dg) + 1
    end
    plotting_interp_matrix = plotting_interpolation_matrix(dg; nvisnodes = nvisnodes)

    # create triangulation for plotting nodes
    r_plot, s_plot = (x -> plotting_interp_matrix * x).((r, s)) # interpolate dg nodes to plotting nodes

    # construct a triangulation of the plotting nodes
    t = reference_plotting_triangulation((r_plot, s_plot))

    # extract x,y coordinates and solutions on each element
    uEltype = eltype(u)
    nvars = nvariables(equations)
    x = reshape(view(cache.elements.node_coordinates, 1, :, :, :), n_nodes_2d,
                n_elements)
    y = reshape(view(cache.elements.node_coordinates, 2, :, :, :), n_nodes_2d,
                n_elements)
    u_extracted = StructArray{SVector{nvars, uEltype}}(ntuple(_ -> similar(x,
                                                                           (n_nodes_2d,
                                                                            n_elements)),
                                                              nvars))
    for element in eachelement(dg, cache)
        sk = 1
        for j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, element)
            u_extracted[sk, element] = u_node
            sk += 1
        end
    end

    # interpolate to volume plotting points
    xplot, yplot = plotting_interp_matrix * x, plotting_interp_matrix * y
    uplot = StructArray{SVector{nvars, uEltype}}(map(x -> plotting_interp_matrix * x,
                                                     StructArrays.components(u_extracted)))

    xfp, yfp, ufp = mesh_plotting_wireframe(u_extracted, mesh, equations, dg, cache;
                                            nvisnodes = nvisnodes)

    # convert variables based on solution_variables mapping
    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names = SVector(varnames(solution_variables_, equations))

    transform_to_solution_variables!(uplot, solution_variables_, equations)
    transform_to_solution_variables!(ufp, solution_variables_, equations)

    return PlotData2DTriangulated(xplot, yplot, uplot, t, xfp, yfp, ufp, variable_names)
end

# Wrapper struct to indicate that an array represents a scalar data field. Used only for dispatch.
struct ScalarData{T}
    data::T
end

"""
    ScalarPlotData2D(u, semi::AbstractSemidiscretization; kwargs...)
    ScalarPlotData2D(function_to_visualize, u, semi::AbstractSemidiscretization; kwargs...)

Returns a `PlotData2DTriangulated` object which is used to visualize a single scalar field.
`u` should be an array whose entries correspond to values of the scalar field at nodal points.

The optional argument `function_to_visualize(u, equations)` should be a function which takes
in the conservative variables and equations as input and outputs a scalar variable to be visualized,
e.g., [`pressure`](@ref) or [`density`](@ref) for the compressible Euler equations.
"""
function ScalarPlotData2D(function_to_visualize::Func, u_ode,
                          semi::AbstractSemidiscretization;
                          kwargs...) where {Func}
    u = wrap_array(u_ode, semi)
    scalar_data = evaluate_scalar_function_at_nodes(function_to_visualize,
                                                    u,
                                                    mesh_equations_solver_cache(semi)...)
    return ScalarPlotData2D(scalar_data,
                            mesh_equations_solver_cache(semi)...; kwargs...)
end

function ScalarPlotData2D(scalar_data, semi::AbstractSemidiscretization; kwargs...)
    return ScalarPlotData2D(scalar_data, mesh_equations_solver_cache(semi)...;
                            kwargs...)
end

function evaluate_scalar_function_at_nodes(function_to_visualize, u, mesh, equations,
                                           dg::DGMulti, cache)
    # for DGMulti solvers, eltype(u) should be SVector{nvariables(equations)}, so
    # broadcasting `func_to_visualize` over the solution array will work.
    return function_to_visualize.(u, equations)
end

function evaluate_scalar_function_at_nodes(function_to_visualize, u,
                                           mesh::AbstractMesh{2}, equations,
                                           dg::Union{<:DGSEM, <:FDSBP}, cache)

    # `u` should be an array of size (nvariables, nnodes, nnodes, nelements)
    f = zeros(eltype(u), nnodes(dg), nnodes(dg), nelements(dg, cache))
    @threaded for element in eachelement(dg, cache)
        for j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, element)
            f[i, j, element] = function_to_visualize(u_node, equations)
        end
    end
    return f
end

# Returns an `PlotData2DTriangulated` which is used to visualize a single scalar field
function ScalarPlotData2D(u, mesh, equations, dg::DGMulti, cache;
                          variable_name = nothing, nvisnodes = 2 * nnodes(dg))
    rd = dg.basis
    md = mesh.md

    # Vp = the interpolation matrix from nodal points to plotting points
    @unpack Vp = rd

    # interpolate nodal coordinates and solution field to plotting points
    x_plot, y_plot = map(x -> Vp * x, md.xyz) # md.xyz is a tuple of arrays containing nodal coordinates
    u_plot = Vp * u

    # construct a triangulation of the reference plotting nodes
    t = reference_plotting_triangulation(rd.rstp) # rd.rstp = reference coordinates of plotting points

    # Ignore face data when plotting `ScalarPlotData2D`, since mesh lines can be plotted using
    # existing functionality based on `PlotData2D(sol)`.
    x_face, y_face, face_data = mesh_plotting_wireframe(ScalarData(u), mesh, equations,
                                                        dg, cache;
                                                        nvisnodes = 2 * nnodes(dg))

    # wrap solution in ScalarData struct for recipe dispatch
    return PlotData2DTriangulated(x_plot, y_plot, ScalarData(u_plot), t,
                                  x_face, y_face, face_data, variable_name)
end

function ScalarPlotData2D(u, mesh, equations, dg::Union{<:DGSEM, <:FDSBP}, cache;
                          variable_name = nothing,
                          nvisnodes = 2 * nnodes(dg))
    n_nodes_2d = nnodes(dg)^ndims(mesh)
    n_elements = nelements(dg, cache)

    # build nodes on reference element (seems to be the right ordering)
    r, s = reference_node_coordinates_2d(dg)

    # reference plotting nodes
    if nvisnodes == 0 || nvisnodes === nothing
        nvisnodes = polydeg(dg) + 1
    end
    plotting_interp_matrix = plotting_interpolation_matrix(dg; nvisnodes = nvisnodes)

    # create triangulation for plotting nodes
    r_plot, s_plot = (x -> plotting_interp_matrix * x).((r, s)) # interpolate dg nodes to plotting nodes

    # construct a triangulation of the plotting nodes
    t = reference_plotting_triangulation((r_plot, s_plot))

    # extract x,y coordinates and reshape them into matrices of size (n_nodes_2d, n_elements)
    x = view(cache.elements.node_coordinates, 1, :, :, :)
    y = view(cache.elements.node_coordinates, 2, :, :, :)
    x, y = reshape.((x, y), n_nodes_2d, n_elements)

    # interpolate to volume plotting points by multiplying each column by `plotting_interp_matrix`
    x_plot, y_plot = plotting_interp_matrix * x, plotting_interp_matrix * y
    u_plot = plotting_interp_matrix * reshape(u, size(x))

    # Ignore face data when plotting `ScalarPlotData2D`, since mesh lines can be plotted using
    # existing functionality based on `PlotData2D(sol)`.
    x_face, y_face, face_data = mesh_plotting_wireframe(ScalarData(u), mesh, equations,
                                                        dg, cache;
                                                        nvisnodes = 2 * nnodes(dg))

    # wrap solution in ScalarData struct for recipe dispatch
    return PlotData2DTriangulated(x_plot, y_plot, ScalarData(u_plot), t,
                                  x_face, y_face, face_data, variable_name)
end

"""
    PlotData1D(u, semi [or mesh, equations, solver, cache];
               solution_variables=nothing, nvisnodes=nothing,
               reinterpolate=Trixi.default_reinterpolate(solver),
               slice=:x, point=(0.0, 0.0, 0.0), curve=nothing)

Create a new `PlotData1D` object that can be used for visualizing 1D DGSEM solution data array
`u` with `Plots.jl`. All relevant geometrical information is extracted from the
semidiscretization `semi`. By default, the primitive variables (if existent)
or the conservative variables (otherwise) from the solution are used for
plotting. This can be changed by passing an appropriate conversion function to
`solution_variables`, e.g., [`cons2cons`](@ref) or [`cons2prim`](@ref).

Alternatively, you can also pass a function `u` with signature `u(x, equations)`
returning a vector. In this case, the `solution_variables` are ignored. This is useful,
e.g., to visualize an analytical solution. The spatial variable `x` is a vector with one
element.

`nvisnodes` specifies the number of visualization nodes to be used. If it is `nothing`,
twice the number of solution DG nodes are used for visualization, and if set to `0`,
exactly the number of nodes as in the DG elements are used. The solution is interpolated to these
nodes for plotting. If `reinterpolate` is `false`, the solution is not interpolated to the `visnodes`,
i.e., the solution is plotted at the solution nodes. In this case `nvisnodes` is ignored. By default,
`reinterpolate` is set to `false` for `FDSBP` approximations and to `true` otherwise.

When visualizing data from a two-dimensional simulation, a 1D slice is extracted for plotting.
`slice` specifies the axis along which the slice is extracted and may be `:x`, or `:y`.
The slice position is specified by a `point` that lies on it, which defaults to `(0.0, 0.0)`.
Both of these values are ignored when visualizing 1D data.
This applies analogously to three-dimensional simulations, where `slice` may be `:xy`, `:xz`, or `:yz`.

Another way to visualize 2D/3D data is by creating a plot along a given curve.
This is done with the keyword argument `curve`. It can be set to a list of 2D/3D points
which define the curve. When using `curve` any other input from `slice` or `point` will be ignored.
"""
function PlotData1D(u_ode, semi; kwargs...)
    return PlotData1D(wrap_array_native(u_ode, semi),
                      mesh_equations_solver_cache(semi)...;
                      kwargs...)
end

function PlotData1D(func::Function, semi; kwargs...)
    return PlotData1D(func,
                      mesh_equations_solver_cache(semi)...;
                      kwargs...)
end

function PlotData1D(u, mesh::TreeMesh, equations, solver, cache;
                    solution_variables = nothing, nvisnodes = nothing,
                    reinterpolate = default_reinterpolate(solver),
                    slice = :x, point = (0.0, 0.0, 0.0), curve = nothing,
                    variable_names = nothing)
    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names_ = digest_variable_names(solution_variables_, equations,
                                            variable_names)

    original_nodes = cache.elements.node_coordinates
    unstructured_data = get_unstructured_data(u, solution_variables_, mesh, equations,
                                              solver, cache)

    orientation_x = 0 # Set 'orientation' to zero on default.

    if ndims(mesh) == 1
        x, data, mesh_vertices_x = get_data_1d(original_nodes, unstructured_data,
                                               nvisnodes, reinterpolate)
        orientation_x = 1

        # Special care is required for first-order FV approximations since the nodes are the
        # cell centers and do not contain the boundaries
        n_nodes = size(unstructured_data, 1)
        if n_nodes == 1
            n_visnodes = length(x) ÷ nelements(solver, cache)
            if n_visnodes != 2
                throw(ArgumentError("This number of visualization nodes is currently not supported for finite volume approximations."))
            end
            left_boundary = mesh.tree.center_level_0[1] - mesh.tree.length_level_0 / 2
            dx_2 = zero(left_boundary)
            for i in 1:div(length(x), 2)
                # Adjust plot nodes so that they are at the boundaries of each element
                dx_2 = x[2 * i - 1] - left_boundary
                x[2 * i - 1] -= dx_2
                x[2 * i] += dx_2
                left_boundary = left_boundary + 2 * dx_2

                # Adjust mesh plot nodes
                mesh_vertices_x[i] -= dx_2
            end
            mesh_vertices_x[end] += dx_2
        end
    elseif ndims(mesh) == 2
        if curve !== nothing
            x, data, mesh_vertices_x = unstructured_2d_to_1d_curve(original_nodes,
                                                                   unstructured_data,
                                                                   nvisnodes, curve,
                                                                   mesh, solver, cache)
        else
            x, data, mesh_vertices_x = unstructured_2d_to_1d(original_nodes,
                                                             unstructured_data,
                                                             nvisnodes, reinterpolate,
                                                             slice, point)
        end
    else # ndims(mesh) == 3
        if curve !== nothing
            x, data, mesh_vertices_x = unstructured_3d_to_1d_curve(original_nodes,
                                                                   unstructured_data,
                                                                   nvisnodes, curve,
                                                                   mesh, solver, cache)
        else
            x, data, mesh_vertices_x = unstructured_3d_to_1d(original_nodes,
                                                             unstructured_data,
                                                             nvisnodes, reinterpolate,
                                                             slice, point)
        end
    end

    return PlotData1D(x, data, variable_names_, mesh_vertices_x,
                      orientation_x)
end

function PlotData1D(u, mesh::TreeMesh1D, equations, solver::BlockFV, cache;
                    solution_variables = nothing, nvisnodes = nothing,
                    reinterpolate = default_reinterpolate(solver),
                    slice = :x, point = (0.0, 0.0, 0.0), curve = nothing,
                    variable_names = nothing)
    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names_ = digest_variable_names(solution_variables_, equations,
                                            variable_names)
    unstructured_data = get_unstructured_data(u, solution_variables_, mesh,
                                              equations, solver, cache)

    n_nodes = size(unstructured_data, 1)  # number of subcells per cell
    n_elements = nelements(solver, cache) # number of cells
    total_plot_points = 2 * n_nodes * n_elements # The total array size we need

    x = Vector{Float64}(undef, total_plot_points)
    data = Array{Float64}(undef, total_plot_points, length(variable_names_))
    mesh_vertices_x = Vector{Float64}(undef, n_elements + 1)

    left_boundary = mesh.tree.center_level_0[1] - mesh.tree.length_level_0 / 2

    if ndims(mesh) === 1
        orientation_x = 1

        for i in 1:n_elements #the loop over the cells
            element_width = 2.0 / cache.elements.inverse_jacobian[i]
            sub_cell_width = element_width / n_nodes

            for j in 1:n_nodes #the loop over the subscells
                idx_left = (i - 1) * (2 * n_nodes) + 2 * j - 1 #slot number in the list x for the left side of this subcell
                idx_right = (i - 1) * (2 * n_nodes) + 2 * j #slot number in the list x for the right side of this subcell

                x[idx_left] = left_boundary + (j - 1) * sub_cell_width #coordinate for the left edge of the subcell
                x[idx_right] = left_boundary + j * sub_cell_width #coordinate for the right edge of the subcell

                for v in 1:length(variable_names_) #loop over the values
                    data[idx_left, v] = unstructured_data[j, i, v] #value for quantity v at left edge of subcell
                    data[idx_right, v] = unstructured_data[j, i, v] #value for quantity v at right edge of subcell
                end
            end

            mesh_vertices_x[i] = left_boundary #write left boundary of the cell under consideration into the i'th entry of mesh_vertices_x 
            left_boundary += element_width #go to the left boundary of the next cell
        end
        mesh_vertices_x[end] = left_boundary #this is the right boundary. We do that by going out of the cell. In that case its the left boundary

    else
        error("BlockFV is not yet supported for 2D and 3D.")
    end
    return PlotData1D(x, data, variable_names_, mesh_vertices_x, orientation_x)
end

# PlotData2DTriangulated for BlockFV on P4estMesh{2}.
# Each FV cell becomes a flat-colored quad (split into 2 triangles) so that
# cell-to-cell jumps are visible rather than smoothed away by interpolation.
function PlotData2DTriangulated(u, mesh, equations, dg::BlockFV, cache;
                                solution_variables = nothing, kwargs...)
    @assert ndims(mesh)==2 "Input must be two-dimensional."

    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names = SVector(varnames(solution_variables_, equations))
    nvars = length(variable_names)
    uEltype = eltype(u)

    n = nnodes(dg)
    n_elements = nelements(dg, cache)
    n_cells = n^2 * n_elements

    corners_x, corners_y = calc_fv_cell_corners(mesh, dg, cache)

    x = Array{Float64}(undef, 4, n_cells)
    y = similar(x)
    data = StructArray{SVector{nvars, uEltype}}(ntuple(_ -> zeros(uEltype, 4, n_cells),
                                                       nvars))

    cell = 0
    for element in 1:n_elements
        for j in 1:n, i in 1:n
            cell += 1
            x[1, cell] = corners_x[i, j, element]
            x[2, cell] = corners_x[i + 1, j, element]
            x[3, cell] = corners_x[i + 1, j + 1, element]
            x[4, cell] = corners_x[i, j + 1, element]
            y[1, cell] = corners_y[i, j, element]
            y[2, cell] = corners_y[i + 1, j, element]
            y[3, cell] = corners_y[i + 1, j + 1, element]
            y[4, cell] = corners_y[i, j + 1, element]

            u_node = solution_variables_(get_node_vars(u, equations, dg, i, j, element),
                                         equations)
            data[1, cell] = u_node
            data[2, cell] = u_node
            data[3, cell] = u_node
            data[4, cell] = u_node
        end
    end

    # Reference triangulation: split each quad (corners 1-2-3-4) into 2 triangles.
    # The plotting recipe expands this to all cells.
    t = [1 2 3; 1 3 4]

    x_face, y_face = calc_fv_grid_wireframe(corners_x, corners_y)

    return PlotData2DTriangulated(x, y, data, t, x_face, y_face, nothing,
                                  variable_names)
end

# unwrap u if it is VectorOfArray
PlotData1D(u::VectorOfArray, mesh, equations, dg::DGMulti{1}, cache; kwargs...) = PlotData1D(parent(u),
                                                                                             mesh,
                                                                                             equations,
                                                                                             dg,
                                                                                             cache;
                                                                                             kwargs...)
PlotData2D(u::VectorOfArray, mesh, equations, dg::DGMulti{2}, cache; kwargs...) = PlotData2D(parent(u),
                                                                                             mesh,
                                                                                             equations,
                                                                                             dg,
                                                                                             cache;
                                                                                             kwargs...)

function PlotData1D(u, mesh, equations, solver, cache;
                    solution_variables = nothing, nvisnodes = nothing,
                    reinterpolate = default_reinterpolate(solver),
                    slice = :x, point = (0.0, 0.0, 0.0), curve = nothing,
                    variable_names = nothing)
    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names_ = digest_variable_names(solution_variables_, equations,
                                            variable_names)

    original_nodes = cache.elements.node_coordinates

    orientation_x = 0 # Set 'orientation' to zero on default.

    if ndims(mesh) == 1
        unstructured_data = get_unstructured_data(u, solution_variables_,
                                                  mesh, equations, solver, cache)
        x, data, mesh_vertices_x = get_data_1d(original_nodes, unstructured_data,
                                               nvisnodes, reinterpolate)
        orientation_x = 1
    elseif ndims(mesh) == 2
        x, data, mesh_vertices_x = unstructured_2d_to_1d_curve(u, mesh, equations,
                                                               solver, cache,
                                                               curve, slice,
                                                               point, nvisnodes,
                                                               solution_variables_)
    else # ndims(mesh) == 3
        # Extract the information required to create a PlotData1D object.
        # If no curve is defined, create an axis curve.
        if curve === nothing
            curve = axis_curve(view(original_nodes, 1, :, :, :, :),
                               view(original_nodes, 2, :, :, :, :),
                               view(original_nodes, 3, :, :, :, :),
                               slice, point, nvisnodes)
        end

        # We need to loop through all the points and check in which element they are
        # located. A general implementation working for all mesh types has to perform
        # a naive loop through all nodes. However, the P4estMesh and the T8codeMesh
        # can make use of the efficient search functionality of p4est/t8code
        # to speed up the process. Thus, we pass the mesh, too.
        x, data, mesh_vertices_x = unstructured_3d_to_1d_curve(u, mesh, equations,
                                                               solver, cache,
                                                               curve,
                                                               solution_variables_)
    end

    return PlotData1D(x, data, variable_names_, mesh_vertices_x,
                      orientation_x)
end

function PlotData1D(func::Function, mesh, equations, dg::DGMulti{1}, cache;
                    solution_variables = nothing, variable_names = nothing)
    x = mesh.md.x
    u = func.(x, equations)

    return PlotData1D(u, mesh, equations, dg, cache;
                      solution_variables, variable_names)
end

# Specializes the `PlotData1D` constructor for one-dimensional `DGMulti` solvers.
function PlotData1D(u, mesh, equations, dg::DGMulti{1}, cache;
                    solution_variables = nothing, variable_names = nothing)
    solution_variables_ = digest_solution_variables(equations, solution_variables)
    variable_names_ = digest_variable_names(solution_variables_, equations,
                                            variable_names)

    orientation_x = 0 # Set 'orientation' to zero on default.

    if u isa StructArray
        # Convert conserved variables to the given `solution_variables` and set up
        # plotting coordinates
        # This uses a "structure of arrays"
        data = map(x -> vcat(dg.basis.Vp * x, fill(NaN, 1, size(u, 2))),
                   StructArrays.components(solution_variables_.(u, equations)))
        x = vcat(dg.basis.Vp * mesh.md.x, fill(NaN, 1, size(u, 2)))

        # Here, we ensure that `DGMulti` visualization uses the same data layout and format
        # as `TreeMesh`. This enables us to reuse existing plot recipes. In particular,
        # `hcat(data...)` creates a matrix of size `num_plotting_points` by `nvariables(equations)`,
        # with data on different elements separated by `NaNs`.
        x_plot = vec(x)
        data_plot = hcat(vec.(data)...)
    else
        # Convert conserved variables to the given `solution_variables` and set up
        # plotting coordinates
        # This uses an "array of structures"
        data_tmp = dg.basis.Vp * solution_variables_.(u, equations)
        data = vcat(data_tmp, fill(NaN * zero(eltype(data_tmp)), 1, size(u, 2)))
        x = vcat(dg.basis.Vp * mesh.md.x, fill(NaN, 1, size(u, 2)))

        # Same as above - we create `data_plot` as array of size `num_plotting_points`
        # by "number of plotting variables".
        x_plot = vec(x)
        data_ = reinterpret(reshape, eltype(eltype(data)), vec(data))
        # If there is only one solution variable, we need to add a singleton dimension
        if ndims(data_) == 1
            data_plot = reshape(data_, :, 1)
        else
            data_plot = permutedims(data_, (2, 1))
        end
    end

    return PlotData1D(x_plot, data_plot, variable_names_, mesh.md.VX, orientation_x)
end

"""
    PlotData1D(sol; kwargs...)

Create a `PlotData1D` object from a solution object created by either `OrdinaryDiffEq.solve!`
(which returns a `SciMLBase.ODESolution`) or Trixi.jl's own `solve!` (which returns a
`TimeIntegratorSolution`).
"""
function PlotData1D(sol::TrixiODESolution; kwargs...)
    return PlotData1D(sol.u[end], sol.prob.p; kwargs...)
end

function PlotData1D(time_series_callback::TimeSeriesCallback, point_id::Integer)
    @unpack time, variable_names, point_data = time_series_callback

    n_solution_variables = length(variable_names)
    data = Matrix{Float64}(undef, length(time), n_solution_variables)
    reshaped = reshape(point_data[point_id], n_solution_variables, length(time))
    for v in 1:n_solution_variables
        @views data[:, v] = reshaped[v, :]
    end

    mesh_vertices_x = Vector{Float64}(undef, 0)

    return PlotData1D(time, data, SVector(variable_names), mesh_vertices_x, 0)
end

function PlotData1D(cb::DiscreteCallback{<:Any, <:TimeSeriesCallback},
                    point_id::Integer)
    return PlotData1D(cb.affect!, point_id)
end
end # @muladd
