#pragma once
#include "Graphics/Buffer.h"
#include "Graphics/Descriptor.h"
#include "Graphics/Initializer.h"
#include "Graphics/RenderPass/RenderPass.h"
#include "Graphics/Resources.h"
#include "Math/Vector.h"
#include "Scene/Scene.h"

namespace Seele {
class SimulationComputePass : public RenderPass {
  public:
    SimulationComputePass(Gfx::PGraphics graphics, PScene scene);
    virtual ~SimulationComputePass() override;
    virtual void beginFrame(const Component::Camera& camera, const Seele::Component::Transform& transform) override;
    virtual void render() override;
    virtual void endFrame() override;
    virtual void publishOutputs() override;
    virtual void createRenderPass() override;

  private:
    UVector getDispatchSize(const UVector& gridSize, const UVector& threadGroupSize);
    void setBoundsPass(Gfx::PShaderBuffer buffer, int b, UVector gridSize);
    void linearSolvePass(Gfx::PShaderBuffer next, Seele::Gfx::PShaderBuffer current, Seele::Gfx::PShaderBuffer grid0, float a, float c, UVector gridSize);
    void divergencePass(Gfx::PShaderBuffer divergence, Seele::Gfx::PShaderBuffer pressure, Seele::Gfx::PShaderBuffer velocityX,
                        Seele::Gfx::PShaderBuffer velocityY, Seele::Gfx::PShaderBuffer velocityZ, UVector gridSize);
    void projectPass(Gfx::PShaderBuffer velocityX, Seele::Gfx::PShaderBuffer velocityY, Seele::Gfx::PShaderBuffer velocityZ,
                     Seele::Gfx::PShaderBuffer pressure, UVector gridSize);
    void advectPass(Gfx::PShaderBuffer quantity, Seele::Gfx::PShaderBuffer quantity0, Seele::Gfx::PShaderBuffer velocityX,
                    Seele::Gfx::PShaderBuffer velocityY, Seele::Gfx::PShaderBuffer velocityZ, float dt, UVector gridSize);
    void reinitializeLevelSetPass(Gfx::PShaderBuffer phi, Seele::Gfx::PShaderBuffer phi0, UVector gridSize);
    struct SetBounds {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        Gfx::OComputeShader shaderEdges;
        Gfx::PComputePipeline pipelineEdges;
        Gfx::OComputeShader shaderCorners;
        Gfx::PComputePipeline pipelineCorners;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } setBounds;
    struct LinearSolve {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } linearSolve;
    struct ComputeDivergence {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } computeDivergence;
    struct Project {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } project;
    struct Advect {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } advect;
    struct Reinitialize {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } reinitialize;
    struct ApplyForces {
        Gfx::OPipelineLayout pipelineLayout;
        Gfx::ODescriptorLayout descriptorLayout;
        Gfx::OComputeShader shader;
        Gfx::PComputePipeline pipeline;
        constexpr static UVector threadGroupSize = {32, 8, 1};
    } applyForces;
    void applyForcesPass(Gfx::PShaderBuffer velocityX, Gfx::PShaderBuffer velocityY, Gfx::PShaderBuffer velocityZ, Vector gravity, float dt, UVector gridSize);
    Gfx::ODescriptorLayout gridSizeLayout;
    Gfx::ODescriptorSet gridSizeSet;
    PScene scene;
};
DEFINE_REF(SimulationComputePass)
} // namespace Seele