using GPUArraysCore
using Oceananigans.Grids: stretched_dimensions
using Oceananigans.Grids: XDirection, YDirection, ZDirection
using Oceananigans.Operators: Δxᶠᵃᵃ, Δyᵃᶠᵃ, Δzᵃᵃᶠ

using Oceananigans.Solvers: BatchedTridiagonalSolver,
                            stretched_direction,
                            tridiagonal_direction,
                            dimension,
                            HomogeneousNeumannFormulation,
                            ZTridiagonalSolver,
                            YTridiagonalSolver,
                            XTridiagonalSolver,
                            compute_main_diagonal!,
                            compute_lower_diagonal!

import Oceananigans.Solvers: compute_preconditioner_rhs!

struct DistributedFourierTridiagonalPoissonSolver{G, L, B, P, R, S, β}
    plan :: P
    global_grid :: G
    local_grid :: L
    batched_tridiagonal_solver :: B
    source_term :: R
    storage :: S
    buffer :: β
end

# Useful aliases for dispatch...
const XStretchedDistributedSolver = DistributedFourierTridiagonalPoissonSolver{<:Any, <:Any, <:XTridiagonalSolver}
const YStretchedDistributedSolver = DistributedFourierTridiagonalPoissonSolver{<:Any, <:Any, <:YTridiagonalSolver}
const ZStretchedDistributedSolver = DistributedFourierTridiagonalPoissonSolver{<:Any, <:Any, <:ZTridiagonalSolver}

architecture(solver::DistributedFourierTridiagonalPoissonSolver) =
    architecture(solver.global_grid)

@inline Δξᶠ(i, grid, ::Val{1}) = Δxᶠᵃᵃ(i, 1, 1, grid)
@inline Δξᶠ(j, grid, ::Val{2}) = Δyᵃᶠᵃ(1, j, 1, grid)
@inline Δξᶠ(k, grid, ::Val{3}) = Δzᵃᵃᶠ(1, 1, k, grid)

"""
    DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid, planner_flag=FFTW.PATIENT;
                                               tridiagonal_formulation=nothing)

Return an FFT-based solver for the Poisson equation evaluated on a `local_grid` that has a non-uniform
spacing in exactly one direction (i.e. either in x, y or z)

```math
∇²φ = b
```

for `Distributed` architectures.

Supported configurations
========================

In the following, `Nx`, `Ny`, and `Nz` are the number of grid points of the **global** grid
in the `x`, `y`, and `z` directions, while `Rx`, `Ry`, and `Rz` are the number of ranks in the
`x`, `y`, and `z` directions, respectively. Furthermore, 'pencil' decomposition refers to a domain
decomposed in two different directions (i.e., with `Rx != 1` and `Ry != 1`), while 'slab' decomposition
refers to a domain decomposed only in one direction, (i.e., with either `Rx == 1` or `Ry == 1`).
Additionally, `storage` indicates the `TransposableField` used for storing intermediate results;
see [`TransposableField`](@ref).

1. Three dimensional configurations with 'pencil' decompositions in ``(x, y)``
where `Ny ≥ Rx` and `Ny % Rx = 0`, and `Nz ≥ Ry` and `Nz % Ry = 0`.

2. Two dimensional configurations decomposed in ``x`` where `Ny ≥ Rx` and `Ny % Rx = 0`

!!! warning "Unsupported decompositions"
    _Any_ configuration decomposed in ``z`` direction is _not_ supported.
    Furthermore, any ``(x, y)`` decompositions other than the configurations mentioned above are also _not_ supported.

Algorithm for pencil decompositions
============================================

For pencil decompositions (useful for three-dimensional problems),
there are two forward transforms, two backward transforms, one tri-diagonal matrix inversion
and a variable number of transpositions that require MPI communication, dependent on the
stretched direction:

- a stretching in the x-direction requires four transpositions
- a stretching in the y-direction requires six transpositions
- a stretching in the z-direction requires eight transpositions

!!! note "Computational cost"
    Because of the additional transpositions, a stretching in the x-direction
    is computationally cheaper than a stretching in the y-direction, and the latter
    is cheaper than a stretching in the z-direction

In our implementation we require `Nz ≥ Ry` and `Nx ≥ Ry` with the additional constraint
that `Nz % Ry = 0` and `Ny % Rx = 0`.

x - stretched algorithm
========================

1. `storage.zfield`, partitioned over ``(x, y)`` is initialized with the `rhs`.
2. Transform along ``z``.
3. Transpose `storage.zfield` + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
4. Transform along ``y``.
5. Transpose `storage.yfield` + communicate to `storage.xfield` partitioned into `(Rx, Ry)` processes in ``(y, z)``.
6. Solve the tri-diagonal linear system in the ``x`` direction.

Steps 5 -> 1 are reversed to obtain the result in physical space stored in `storage.zfield`
partitioned over ``(x, y)``.

y - stretched algorithm
========================

1. `storage.zfield`, partitioned over ``(x, y)`` is initialized with the `rhs`.
2. Transform along ``z``.
3. Transpose `storage.zfield` + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
4. Transpose `storage.yfield` + communicate to `storage.xfield` partitioned into `(Rx, Ry)` processes in ``(y, z)``.
5. Transform along ``x``.
6. Transpose `storage.xfield` + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
7. Solve the tri-diagonal linear system in the ``y`` direction.

Steps 6 -> 1 are reversed to obtain the result in physical space stored in `storage.zfield`
partitioned over ``(x, y)``.

z - stretched algorithm
========================

1. `storage.zfield`, partitioned over ``(x, y)`` is initialized with the `rhs`.
2. Transpose `storage.zfield` + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
3. Transform along ``y``.
4. Transpose `storage.yfield` + communicate to `storage.xfield` partitioned into `(Rx, Ry)` processes in ``(y, z)``.
5. Transform along ``x``.
6. Transpose `storage.xfield` + communicate to `storage.yfield` partitioned into `(Rx, Ry)` processes in ``(x, z)``.
7. Transpose `storage.yfield` + communicate to `storage.zfield` partitioned into `(Rx, Ry)` processes in ``(x, y)``.
8. Solve the tri-diagonal linear system in the ``z`` direction.

Steps 7 -> 1 are reversed to obtain the result in physical space stored in `storage.zfield`
partitioned over ``(x, y)``.

Algorithm for slab decompositions
=============================

The 'slab' decomposition works in the same manner while skipping the transposes that
are not required. For example if the domain is decomposed in ``x``, step 4. and 6. in the above algorithm
are skipped (and the associated reversed step in the backward transform)

Restrictions
============

1. Pencil decompositions:
    - `Ny ≥ Rx` and `Ny % Rx = 0`
    - `Nz ≥ Ry` and `Nz % Ry = 0`
    - If the ``z`` direction is `Periodic`, also the ``y`` and the ``x`` directions must be `Periodic`
    - If the ``y`` direction is `Periodic`, also the ``x`` direction must be `Periodic`

2. Slab decomposition:
    - Same as for two-dimensional decompositions with `Rx` (or `Ry`) equal to one

"""
function DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid, planner_flag=FFTW.PATIENT; tridiagonal_formulation=nothing)

    validate_poisson_solver_distributed_grid(global_grid)
    validate_poisson_solver_configuration(global_grid, local_grid)

    # Try to guess what direction should be tridiagonal
    if isnothing(tridiagonal_formulation)
        tridiagonal_dir = global_grid isa XYZRegularRG ? ZDirection() : stretched_direction(global_grid)
        tridiagonal_formulation = HomogeneousNeumannFormulation(tridiagonal_dir)
    else
        tridiagonal_dir = tridiagonal_direction(tridiagonal_formulation)
    end

    tridiagonal_dim = dimension(tridiagonal_dir)

    topology(global_grid, tridiagonal_dim) != Bounded &&
        error("`DistributedFourierTridiagonalPoissonSolver` requires that the stretched direction (dimension $tridiagonal_dim) is `Bounded`.")

    T = Complex{eltype(local_grid)}
    child_arch = child_architecture(local_grid)
    storage_field = CenterField(local_grid)
    storage = TransposableField(storage_field, T)

    TX, TY, TZ = topology(global_grid)
    tx, ty, tz = TX(), TY(), TZ()
    λx = dropdims(poisson_eigenvalues(global_grid, global_grid.Nx, global_grid.Lx, 1, tx), dims=(2, 3))
    λy = dropdims(poisson_eigenvalues(global_grid, global_grid.Ny, global_grid.Ly, 2, ty), dims=(1, 3))
    λz = dropdims(poisson_eigenvalues(global_grid, global_grid.Nz, global_grid.Lz, 3, tz), dims=(1, 2))

    Rx, Ry, _ = local_grid.architecture.ranks
    slab_x = Ry == 1 && Rx > 1  # slab decomposition in x only

    if tridiagonal_dim == 1
        arch = architecture(storage.xfield.grid)
        grid = storage.xfield.grid
        λ1 = partition_coordinate(λy, size(grid, 2), arch, 2)
        λ2 = partition_coordinate(λz, size(grid, 3), arch, 3)
    elseif tridiagonal_dim == 2
        arch = architecture(storage.yfield.grid)
        grid = storage.yfield.grid
        λ1 = partition_coordinate(λx, size(grid, 1), arch, 1)
        λ2 = partition_coordinate(λz, size(grid, 3), arch, 3)
    elseif tridiagonal_dim == 3 && slab_x
        # Optimization: for slab-x decomposition, solve tridiagonal in x-local space.
        # In x-local, z is fully local (all Nz points), so the tridiag solve works directly.
        # This avoids 2 MPI transposes (x→y + y→z before solve, and z→y + y→x after solve).
        arch = architecture(storage.xfield.grid)
        grid = storage.xfield.grid
        λ1 = on_architecture(child_arch, λx)  # x is local: all eigenvalues available
        λ2 = partition_coordinate(λy, size(grid, 2), arch, 2)  # y is partitioned
    elseif tridiagonal_dim == 3
        arch = architecture(storage.zfield.grid)
        grid = storage.zfield.grid
        λ1 = partition_coordinate(λx, size(grid, 1), arch, 1)
        λ2 = partition_coordinate(λy, size(grid, 2), arch, 2)
    end

    λ1 = on_architecture(child_arch, λ1)
    λ2 = on_architecture(child_arch, λ2)

    # Lower and upper diagonals are the same
    main_diagonal = zeros(grid, size(grid)...)

    Nd = size(grid, tridiagonal_dim) - 1
    lower_diagonal = zeros(grid, Nd)
    upper_diagonal = lower_diagonal

    compute_main_diagonal!(main_diagonal, tridiagonal_formulation, grid, λ1, λ2)
    Nd > 0 && compute_lower_diagonal!(lower_diagonal, tridiagonal_formulation, grid)

    # Set up batched tridiagonal solver
    btsolver = BatchedTridiagonalSolver(grid; lower_diagonal, upper_diagonal,
                                        diagonal = main_diagonal,
                                        tridiagonal_direction = tridiagonal_dir)

    # We need to permute indices to apply bounded transforms on the GPU (r2r or r2c with twiddling)
    x_buffer_needed = child_arch isa GPU && TX == Bounded
    z_buffer_needed = child_arch isa GPU && TZ == Bounded

    # Y-buffer is needed on GPU for permutedims before/after the y-FFT.
    # All GPU y-FFTs now use dim-1 layout with permutedims (both Periodic and Bounded).
    y_buffer_needed = child_arch isa GPU

    buffer_x = x_buffer_needed ? on_architecture(child_arch, zeros(T, size(storage.xfield)...)) : nothing
    buffer_y = y_buffer_needed ? on_architecture(child_arch, zeros(T, size(storage.yfield)...)) : nothing
    buffer_z = z_buffer_needed ? on_architecture(child_arch, zeros(T, size(storage.zfield)...)) : nothing

    buffer = if tridiagonal_dim == 1
        (; y = buffer_y, z = buffer_z)
    elseif tridiagonal_dim == 2
        (; x = buffer_x, z = buffer_z)
    elseif tridiagonal_dim == 3
        (; x = buffer_x, y = buffer_y)
    end

    plan = plan_distributed_transforms(global_grid, storage, planner_flag)

    if tridiagonal_dim == 1
        forward  = (y! = plan.forward.y!,  z! = plan.forward.z!)
        backward = (y! = plan.backward.y!, z! = plan.backward.z!)
    elseif tridiagonal_dim == 2
        forward  = (x! = plan.forward.x!,  z! = plan.forward.z!)
        backward = (x! = plan.backward.x!, z! = plan.backward.z!)
    elseif tridiagonal_dim == 3
        forward  = (x! = plan.forward.x!,  y! = plan.forward.y!)
        backward = (x! = plan.backward.x!, y! = plan.backward.y!)
    end

    plan = (; forward, backward)

    # Storage space for right hand side of Poisson equation
    T = complex(eltype(grid))
    source_term = zeros(T, size(grid)...)
    source_term = on_architecture(arch, source_term)

    return DistributedFourierTridiagonalPoissonSolver(plan, global_grid, local_grid, btsolver, source_term, storage, buffer)
end

# solve! requires that `b` in `A x = b` (the right hand side)
# is copied in the solver storage
# See: Models/NonhydrostaticModels/solve_for_pressure.jl
function solve!(x, solver::ZStretchedDistributedSolver)
    storage = solver.storage

    # Detect slab-x optimization: if tridiag solver grid matches xfield grid,
    # we solve in x-local space with only 2 MPI transposes instead of 4.
    if storage isa SlabYFields
        return _slab_x_solve!(x, solver)
    else
        return _general_z_solve!(x, solver)
    end
end

# Optimized solve for slab-x (Ry=1) Z-stretched: only 2 MPI transposes.
# In slab-x, z is always fully local. After forward FFTs and one transpose
# to x-local space, we can solve the tridiagonal directly without transposing back.
function _slab_x_solve!(x, solver::ZStretchedDistributedSolver)
    arch    = architecture(solver)
    storage = solver.storage
    buffer  = solver.buffer

    # Forward: y-FFT (local) → transpose to x-local → x-FFT
    solver.plan.forward.y!(parent(storage.yfield), buffer.y)
    transpose_y_to_x!(storage)
    solver.plan.forward.x!(parent(storage.xfield), buffer.x)

    # Tridiagonal solve in x-local space (z is fully local)
    parent(solver.source_term) .= parent(storage.xfield)
    solve!(storage.xfield, solver.batched_tridiagonal_solver, solver.source_term)

    # Backward: x-IFFT → transpose to y-local → y-IFFT
    solver.plan.backward.x!(parent(storage.xfield), buffer.x)
    transpose_x_to_y!(storage)
    solver.plan.backward.y!(parent(storage.yfield), buffer.y)

    # Copy the real component (yfield aliases zfield for slab-x)
    launch!(arch, solver.local_grid, :xyz,
            _copy_real_component!, x, parent(storage.zfield))

    return x
end

# General Z-stretched solve (pencil decomposition): 4+ MPI transposes.
function _general_z_solve!(x, solver::ZStretchedDistributedSolver)
    arch    = architecture(solver)
    storage = solver.storage
    buffer  = solver.buffer

    transpose_z_to_y!(storage)
    solver.plan.forward.y!(parent(storage.yfield), buffer.y)
    transpose_y_to_x!(storage)
    solver.plan.forward.x!(parent(storage.xfield), buffer.x)
    transpose_x_to_y!(storage)
    transpose_y_to_z!(storage)

    parent(solver.source_term) .= parent(storage.zfield)
    solve!(storage.zfield, solver.batched_tridiagonal_solver, solver.source_term)

    transpose_z_to_y!(storage)
    transpose_y_to_x!(storage)
    solver.plan.backward.x!(parent(storage.xfield), buffer.x)
    transpose_x_to_y!(storage)
    solver.plan.backward.y!(parent(storage.yfield), buffer.y)
    transpose_y_to_z!(storage)

    # Copy the real component of xc to x.
    launch!(arch, solver.local_grid, :xyz,
            _copy_real_component!, x, parent(storage.zfield))

    return x
end

function solve!(x, solver::YStretchedDistributedSolver)
    arch    = architecture(solver)
    storage = solver.storage
    buffer  = solver.buffer

    solver.plan.forward.z!(parent(storage.zfield), buffer.z)
    transpose_z_to_y!(storage) # copy data from storage.zfield to storage.yfield
    transpose_y_to_x!(storage) # copy data from storage.yfield to storage.xfield
    solver.plan.forward.x!(parent(storage.xfield), buffer.x)
    transpose_x_to_y!(storage) # copy data from storage.xfield to storage.yfield

    # copy results in the source term
    parent(solver.source_term) .= parent(storage.yfield)

    # Perform the implicit vertical solve here on storage.yfield...
    # Solve tridiagonal system of linear equations at every y-column.
    solve!(storage.yfield, solver.batched_tridiagonal_solver, solver.source_term)

    transpose_y_to_x!(storage) # copy data from storage.yfield to storage.xfield
    solver.plan.backward.x!(parent(storage.xfield), buffer.x)
    transpose_x_to_y!(storage) # copy data from storage.xfield to storage.yfield
    transpose_y_to_z!(storage) # copy data from storage.yfield to storage.zfield
    solver.plan.backward.z!(parent(storage.zfield), buffer.z)

    # Copy the real component of xc to x.
    launch!(arch, solver.local_grid, :xyz,
            _copy_real_component!, x, parent(storage.zfield))

    return x
end

function solve!(x, solver::XStretchedDistributedSolver)
    arch    = architecture(solver)
    storage = solver.storage
    buffer  = solver.buffer

    # Apply forward transforms to b = first(solver.storage).
    solver.plan.forward.z!(parent(storage.zfield), buffer.z)
    transpose_z_to_y!(storage) # copy data from storage.zfield to storage.yfield
    solver.plan.forward.y!(parent(storage.yfield), buffer.y)
    transpose_y_to_x!(storage) # copy data from storage.yfield to storage.xfield

    # copy results in the source term
    parent(solver.source_term) .= parent(storage.xfield)

    # Perform the implicit vertical solve here on storage.xfield...
    # Solve tridiagonal system of linear equations at every x-column.
    solve!(storage.xfield, solver.batched_tridiagonal_solver, solver.source_term)

    transpose_x_to_y!(storage) # copy data from storage.xfield to storage.yfield
    solver.plan.backward.y!(parent(storage.yfield), buffer.y)
    transpose_y_to_z!(storage) # copy data from storage.yfield to storage.zfield
    solver.plan.backward.z!(parent(storage.zfield), buffer.z) # last backwards transform is in z

    # Copy the real component of xc to x.
    launch!(arch, solver.local_grid, :xyz,
            _copy_real_component!, x, parent(storage.zfield))

    return x
end

#####
##### Preconditioning support for ConjugateGradientPoissonSolver
#####

using Oceananigans.Operators: V⁻¹ᶜᶜᶜ

@kernel function distributed_fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::XDirection, grid, rhs)
    i, j, k = @index(Global, NTuple)
    @inbounds preconditioner_rhs[i, j, k] = rhs[i, j, k] * V⁻¹ᶜᶜᶜ(i, j, k, grid)
end

@kernel function distributed_fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::YDirection, grid, rhs)
    i, j, k = @index(Global, NTuple)
    @inbounds preconditioner_rhs[i, j, k] = rhs[i, j, k] * V⁻¹ᶜᶜᶜ(i, j, k, grid)
end

@kernel function distributed_fourier_tridiagonal_preconditioner_rhs!(preconditioner_rhs, ::ZDirection, grid, rhs)
    i, j, k = @index(Global, NTuple)
    @inbounds preconditioner_rhs[i, j, k] = rhs[i, j, k] * V⁻¹ᶜᶜᶜ(i, j, k, grid)
end

function compute_preconditioner_rhs!(solver::DistributedFourierTridiagonalPoissonSolver, rhs)
    grid = solver.local_grid
    arch = architecture(grid)
    tridiagonal_dir = solver.batched_tridiagonal_solver.tridiagonal_direction

    if tridiagonal_dir isa XDirection
        launch!(arch, grid, :xyz, distributed_fourier_tridiagonal_preconditioner_rhs!,
                solver.storage.xfield, tridiagonal_dir, grid, rhs)
    elseif tridiagonal_dir isa YDirection
        launch!(arch, grid, :xyz, distributed_fourier_tridiagonal_preconditioner_rhs!,
                solver.storage.yfield, tridiagonal_dir, grid, rhs)
    elseif tridiagonal_dir isa ZDirection
        launch!(arch, grid, :xyz, distributed_fourier_tridiagonal_preconditioner_rhs!,
                solver.storage.zfield, tridiagonal_dir, grid, rhs)
    end
    return nothing
end
