module ClimaCorePlots

import RecipesBase
import TriplotBase

import ClimaCore:
    ClimaCore,
    DataLayouts,
    Fields,
    Geometry,
    Meshes,
    Operators,
    Topologies,
    Spaces


RecipesBase.@recipe function f(space::Spaces.SpectralElementSpace2D;)
    quad = Spaces.quadrature_style(space)
    quad_name = Base.typename(typeof(quad)).name
    dof = Spaces.Quadratures.degrees_of_freedom(quad)

    topology = Spaces.topology(space)
    mesh = topology.mesh
    @assert mesh isa Meshes.RectilinearMesh "plotting only defined for RectilinearMesh"

    n1 = Meshes.nelements(mesh.intervalmesh1)
    n2 = Meshes.nelements(mesh.intervalmesh2)

    coord_field = Fields.coordinate_field(space)
    x1coord = vec(parent(coord_field)[:, :, 1, :])
    x2coord = vec(parent(coord_field)[:, :, 2, :])

    coord_symbols = propertynames(coord_field)

    seriestype := :scatter
    title --> "$n1 × $n2 $quad_name{$dof} element space"
    xguide --> "$(coord_symbols[1])"
    yguide --> "$(coord_symbols[2])"
    marker --> :cross
    markersize --> 1
    seriescolor --> :blue
    legend --> false

    (x1coord, x2coord)
end

RecipesBase.@recipe function f(space::Spaces.ExtrudedFiniteDifferenceSpace)
    coord_field = Fields.coordinate_field(space)
    data = Fields.field_values(coord_field)
    Ni, Nj, _, Nv, Nh = size(data)

    #TODO: assumes VIFH layout
    @assert Nj == 1 "plotting only defined for 1D extruded fields"

    hspace = space.horizontal_space

    quad = Spaces.quadrature_style(hspace)
    quad_name = Base.typename(typeof(quad)).name
    dof = Spaces.Quadratures.degrees_of_freedom(quad)

    coord_symbols = propertynames(coord_field)
    hcoord = vec(parent(coord_field)[:, :, 1, :])
    vcoord = vec(parent(coord_field)[:, :, 2, :])

    stagger = space.staggering isa Spaces.CellCenter ? :center : :face

    seriestype := :scatter
    title --> "$Nh $quad_name{$dof} element × $Nv level $stagger space"
    xguide --> "$(coord_symbols[1])"
    yguide --> "$(coord_symbols[2])"
    marker --> :cross
    markersize --> 1
    seriescolor --> :blue
    legend --> false

    (hcoord, vcoord)
end

RecipesBase.@recipe function f(field::Fields.FiniteDifferenceField)
    # unwrap the data to plot
    space = axes(field)
    coord_field = Fields.coordinate_field(space)

    xdata = parent(field)[:, 1]
    ydata = parent(Spaces.coordinates_data(space))[:, 1]

    coord_symbols = propertynames(coord_field)

    # set the plot attributes
    xguide --> "value"
    yguide --> (
        if field isa Spaces.FaceFiniteDifferenceSpace
            "$(coord_symbols[1]) faces"
        else
            "$(coord_symbols[1]) centers"
        end
    )

    # fix the ylim to the column space (domain)
    ylims --> extrema(ydata)
    linewidth --> 2
    legend --> false

    (xdata, ydata)
end

RecipesBase.@recipe function f(
    field::Fields.SpectralElementField2D;
    interpolate = 10,
)
    @assert interpolate ≥ 1 "number of element quadrature points for uniform interpolation must be ≥ 1"

    # compute the interpolated data to plot
    space = axes(field)
    topology = Spaces.topology(space)
    mesh = topology.mesh

    Nu = interpolate
    coord_field = Fields.coordinate_field(space)

    M_coords = Operators.matrix_interpolate(coord_field, Nu)
    M = Operators.matrix_interpolate(field, Nu)

    domain = Meshes.domain(mesh)
    x1min = Geometry.component(domain.interval1.coord_min, 1)
    x2min = Geometry.component(domain.interval2.coord_min, 1)
    x1max = Geometry.component(domain.interval1.coord_max, 1)
    x2max = Geometry.component(domain.interval2.coord_max, 1)

    # our interpolated field is transposed
    x1coord = [Geometry.component(pt, 1) for pt in M_coords[:, 1]]
    x2coord = [Geometry.component(pt, 2) for pt in M_coords[1, :]]

    coord_symbols = propertynames(coord_field)

    # set the plot attributes
    seriestype := :heatmap

    xguide --> "$(coord_symbols[1])"
    yguide --> "$(coord_symbols[2])"
    seriescolor --> :balance

    (x1coord, x2coord, M')
end

function _unfolded_pannel_matrix(field, interpolate)
    space = axes(field)
    topology = Spaces.topology(space)
    mesh = topology.mesh
    nelem = Meshes.nelements(mesh)
    panel_size = mesh.ne

    quad_from = Spaces.quadrature_style(space)
    quad_to = Spaces.Quadratures.Uniform{interpolate}()
    Imat = Spaces.Quadratures.interpolation_matrix(Float64, quad_to, quad_from)

    dof = interpolate

    pannel_range(i) =
        ((panel_size * dof) * (i - 1) + 1):((panel_size * dof) * i)

    # construct a matrix to fill in the rotated / flipped pannel data
    unfolded_panels =
        fill(NaN, ((panel_size * dof) * 3, (panel_size * dof) * 4))

    # temporary pannels as we have to rotate / flip some and not all operators are in place
    # TODO: inefficient memory wise, but good enough for now
    panels = [fill(NaN, (panel_size * dof, panel_size * dof)) for _ in 1:6]

    interpolated_data =
        DataLayouts.IJFH{Float64, interpolate}(Array{Float64}, nelem)
    field_data = Fields.field_values(field)

    Operators.tensor_product!(interpolated_data, field_data, Imat)

    # element index ordering defined by a specific layout
    eidx = 1
    for panel_idx in 1:6
        panel_data = panels[panel_idx]
        # elements are ordered along fastest axis defined in sphere box mesh
        for ex2 in 1:panel_size, ex1 in 1:panel_size
            # compute the nodal extent index range for this element
            x1_nodal_range = (dof * (ex1 - 1) + 1):(dof * ex1)
            x2_nodal_range = (dof * (ex2 - 1) + 1):(dof * ex2)
            # transpose the data as our plotting axis order is
            # reverse nodal element order (x1 axis varies fastest)
            data_element = permutedims(parent(interpolated_data)[:, :, 1, eidx])
            panel_data[x2_nodal_range, x1_nodal_range] = data_element
            eidx += 1
        end
    end

    # (px, py, rot) for each panel
    # px, py are the locations of the panel
    # rot is number of clockwise rotations (0:3)
    # https://extranet.gfdl.noaa.gov/~atw/ferret/cubed_sphere/
    # Equatorial strip, poles connected through the Americas
    panel_locations =
        [(4, 2, 0), (1, 2, 0), (3, 3, 2), (2, 2, 1), (3, 2, 1), (3, 1, 1)]
    for (i, (px, py, rot)) in enumerate(panel_locations)
        unfolded_panels[pannel_range(py), pannel_range(px)] .= if rot == 0
            panels[i]
        elseif rot == 1
            reverse(transpose(panels[i]), dims = 1)
        elseif rot == 2
            reverse(panels[i], dims = (1, 2))
        else
            reverse(transpose(panels[i]), dims = 2)
        end
    end
    #=
    unfolded_panels[pannel_range(1), pannel_range(2)] =
        reverse(panels[5], dims = 1)
    unfolded_panels[pannel_range(2), pannel_range(1)] =
        reverse(panels[4], dims = 2)
    unfolded_panels[pannel_range(2), pannel_range(2)] = transpose(panels[1])
    unfolded_panels[pannel_range(2), pannel_range(3)] = transpose(panels[2])
    unfolded_panels[pannel_range(2), pannel_range(4)] =
        reverse(panels[6], dims = 2)
    unfolded_panels[pannel_range(3), pannel_range(2)] =
        reverse(panels[3], dims = 2)
    =#
    return unfolded_panels
end

RecipesBase.@recipe function f(
    field::Fields.CubedSphereSpectralElementField2D;
    interpolate = 10,
)
    @assert interpolate ≥ 1 "number of element quadrature points for uniform interpolation must be ≥ 1"

    unfolded_panels = _unfolded_pannel_matrix(field, interpolate)

    # construct the title for info about the field space
    space = axes(field)
    nelem = Topologies.nlocalelems(Spaces.topology(space))
    quad_from = Spaces.quadrature_style(space)
    dof_in = Spaces.Quadratures.degrees_of_freedom(quad_from)
    quad_from_name = Base.typename(typeof(quad_from)).name

    # set the plot attributes
    seriestype := :heatmap
    title --> "$nelem $quad_from_name{$dof_in} element space"
    xguide --> "panel x1"
    yguide --> "panel x2"
    seriescolor --> :balance

    (unfolded_panels)
end

RecipesBase.@recipe function f(
    field::Fields.Field{<:Any, S};
    level = nothing,
    hinterpolate = 10,
) where {
    S <: Spaces.ExtrudedFiniteDifferenceSpace{
        <:Any,
        <:Spaces.CubedSphereSpectralElementSpace2D,
    },
}
    @assert hinterpolate ≥ 1 "number of element quadrature points for uniform interpolation must be ≥ 1"

    space = axes(field)
    if level === nothing
        if space isa Spaces.FaceExtrudedFiniteDifferenceSpace
            level = Utilities.PlusHalf(0)
        else
            level = 1
        end
    end
    level_field = Fields.level(field, level)
    unfolded_panels = _unfolded_pannel_matrix(level_field, hinterpolate)

    # construct the title for info about the field space
    nlevel = Spaces.nlevels(space)
    nelem = Topologies.nlocalelems(Spaces.topology(space))
    quad_from = Spaces.quadrature_style(space)
    dof_in = Spaces.Quadratures.degrees_of_freedom(quad_from)
    quad_from_name = Base.typename(typeof(quad_from)).name

    # set the plot attributes
    seriestype := :heatmap
    title -->
    "level $level of $nlevel × $nelem $quad_from_name{$dof_in} element space"
    xguide --> "panel x1"
    yguide --> "panel x2"
    seriescolor --> :balance

    (unfolded_panels)
end

RecipesBase.@recipe function f(
    field::Fields.Field{<:Any, S};
    hinterpolate = 0,
    ncolors = 256,
) where {
    S <: Spaces.ExtrudedFiniteDifferenceSpace{
        <:Any,
        <:Spaces.SpectralElementSpace1D,
    },
}
    data = Fields.field_values(field)
    Ni, Nj, _, Nv, Nh = size(data)

    space = axes(field)
    @assert Nj == 1

    coord_symbols = propertynames(Fields.coordinate_field(space))
    hcoord_field = getproperty(Fields.coordinate_field(space), 1)
    vcoord_field = getproperty(Fields.coordinate_field(space), 2)

    if hinterpolate ≥ 1
        Nu = hinterpolate
        uquad = Spaces.Quadratures.ClosedUniform{Nu}()
        M_hcoord = Operators.matrix_interpolate(hcoord_field, uquad)
        M_vcoord = Operators.matrix_interpolate(vcoord_field, uquad)
        M_data = Operators.matrix_interpolate(field, uquad)

        hcoord_data = vec(M_hcoord)
        vcoord_data = vec(M_vcoord)
        data = vec(M_data)
        triangles = hcat(Spaces.triangles(Nv, Nu, Nh)...)'
    else
        hcoord_data = vec(parent(hcoord_field))
        vcoord_data = vec(parent(vcoord_field))
        data = vec(parent(data))
        triangles = hcat(Spaces.triangles(Nv, Ni, Nh)...)'
    end
    cmap = range(extrema(data)..., length = ncolors)

    z = TriplotBase.tripcolor(
        hcoord_data,
        vcoord_data,
        data,
        triangles,
        cmap;
        bg = NaN,
        # unique number of nodal element coords
        # (number of nodal values minus number of shared faces)
        px = (Ni * Nh) - (Nh - 1),
        py = Nv,
    )

    # set the plot attributes
    seriestype := :heatmap

    xguide --> "$(coord_symbols[1])"
    yguide --> "$(coord_symbols[2])"
    seriescolor --> :balance

    # some plots backends need coordinates in sorted order
    (sort(unique(hcoord_data)), sort(unique(vcoord_data)), z')
end

end # module
