#include "zeroerr/assert.h"
#include "zeroerr/unittest.h"

#include "fluids/FluidMath.h"
#include "fluids/FluidSdf.h"
#include "fluids/FluidSimulation.h"
#include "fluids/Fluids.h"
#include "graphics/Graphics.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <fstream>
#include <utility>
#include <vector>

using namespace eve::fluids;

namespace {

bool tryInitHeadlessGfx() {
    auto* gfx = eve::graphics::Graphics::create();
    if (!gfx) return false;
    gfx->initHeadless(320, 240);
    return true;
}

/** @brief Subdivided icosahedron sphere mesh (positions + triangle indices). */
std::pair<std::vector<glm::vec3>, std::vector<uint32_t>> makeIcosphere(float radius, int subdivisions) {
    const float            t     = (1.f + std::sqrt(5.f)) / 2.f;
    std::vector<glm::vec3> verts = {
        glm::normalize(glm::vec3(-1.f, t, 0.f)),  glm::normalize(glm::vec3(1.f, t, 0.f)),
        glm::normalize(glm::vec3(-1.f, -t, 0.f)), glm::normalize(glm::vec3(1.f, -t, 0.f)),
        glm::normalize(glm::vec3(0.f, -1.f, t)),  glm::normalize(glm::vec3(0.f, 1.f, t)),
        glm::normalize(glm::vec3(0.f, -1.f, -t)), glm::normalize(glm::vec3(0.f, 1.f, -t)),
        glm::normalize(glm::vec3(t, 0.f, -1.f)),  glm::normalize(glm::vec3(t, 0.f, 1.f)),
        glm::normalize(glm::vec3(-t, 0.f, -1.f)), glm::normalize(glm::vec3(-t, 0.f, 1.f)),
    };
    std::vector<uint32_t> tris = {
        0, 11, 5, 0, 5, 1, 0, 1, 7, 0, 7, 10, 0, 10, 11, 1, 5, 9, 5, 11, 4,  11, 10, 2,  10, 7, 6, 7, 1, 8,
        3, 9,  4, 3, 4, 2, 3, 2, 6, 3, 6, 8,  3, 8,  9,  4, 9, 5, 2, 4,  11, 6,  2,  10, 8,  6, 7, 9, 8, 1,
    };
    for (int s = 0; s < subdivisions; ++s) {
        std::vector<uint32_t> nextTris;
        nextTris.reserve(tris.size() * 4);
        for (size_t i = 0; i < tris.size(); i += 3) {
            const uint32_t a = tris[i], b = tris[i + 1], c = tris[i + 2];
            const uint32_t ab = uint32_t(verts.size());
            const uint32_t bc = uint32_t(verts.size()) + 1;
            const uint32_t ca = uint32_t(verts.size()) + 2;
            verts.push_back(glm::normalize((verts[a] + verts[b]) * 0.5f));
            verts.push_back(glm::normalize((verts[b] + verts[c]) * 0.5f));
            verts.push_back(glm::normalize((verts[c] + verts[a]) * 0.5f));
            nextTris.insert(nextTris.end(), {a, ab, ca, b, bc, ab, c, ca, bc, ab, bc, ca});
        }
        tris.swap(nextTris);
    }
    for (glm::vec3& v : verts) v *= radius;
    return {std::move(verts), std::move(tris)};
}

float distToCenter(const glm::vec3& p) { return glm::length(p); }

}  // namespace

TEST_CASE("fluids.math.kernels") {
    const float h = 0.2f;
    // Poly6 is positive inside the support and zero at/outside the cutoff.
    CHECK(fluidPoly6(0.001f, h) > 0.f);
    CHECK(std::fabs(fluidPoly6(h * h, h)) < 1e-9f);
    CHECK(std::fabs(fluidPoly6((h + 0.1f) * (h + 0.1f), h)) < 1e-12f);
    // Poly6 is radially decreasing (denser near the center).
    CHECK(fluidPoly6(0.001f, h) > fluidPoly6(0.01f, h));

    // With dx = p_i - p_j, the radially decreasing kernel gradient points toward j.
    const glm::vec3 dx(0.1f, 0.f, 0.f);
    const glm::vec3 g = fluidSpikyGrad(dx, h);
    CHECK(g.x < 0.f);
    CHECK(std::fabs(g.y) < 1e-9f);
    CHECK(glm::length(fluidSpikyGrad(glm::vec3(h, 0.f, 0.f), h)) < 1e-6f);
    CHECK(glm::length(fluidSpikyGrad(glm::vec3(0.f), h)) < 1e-9f);

    // Viscosity Laplacian is positive inside and zero at cutoff.
    CHECK(fluidViscLaplacian(0.05f, h) > 0.f);
    CHECK(std::fabs(fluidViscLaplacian(h, h)) < 1e-6f);
}

TEST_CASE("fluids.sdf.sphere") {
    const glm::vec3 center(0.f, 1.f, 0.f);
    const float     radius = 1.f;
    const MeshSdf   sdf    = MeshSdf::makeSphere(center, radius, glm::ivec3(32));
    REQUIRE(sdf.voxelCount() == 32 * 32 * 32);

    // Center is deep inside (distance ≈ -radius).
    CHECK(std::fabs(sdf.sample(center) + radius) < 0.2f);
    // Point on the surface reads ~0.
    CHECK(std::fabs(sdf.sample(center + glm::vec3(radius, 0.f, 0.f))) < 0.2f);
    // Far outside reads positive.
    CHECK(sdf.sample(center + glm::vec3(2.f * radius, 0.f, 0.f)) > 0.5f);
    CHECK(std::fabs(sdf.sample(center + glm::vec3(2.f * radius, 0.f, 0.f)) - radius) < 0.2f);

    // Gradient points outward along +Y at the top.
    const glm::vec3 g = sdf.gradient(center + glm::vec3(0.f, radius, 0.f));
    CHECK(glm::dot(g, glm::vec3(0.f, 1.f, 0.f)) > 0.f);
    CHECK(glm::distance(g, sdf.sampleWithGradient(center + glm::vec3(0.f, radius, 0.f)).gradient) < 1e-6f);
}

TEST_CASE("fluids.sdf.triangleMesh") {
    auto [verts, tris] = makeIcosphere(1.f, 2);
    REQUIRE(verts.size() >= 100);
    REQUIRE((tris.size() % 3) == 0);
    const MeshSdf mesh     = MeshSdf::makeFromTriangles(verts, tris, glm::ivec3(32));
    const MeshSdf analytic = MeshSdf::makeSphere(glm::vec3(0.f), 1.f, glm::ivec3(32));

    // The voxelized mesh should match the analytic sphere within ~2 voxels.
    const float                  tol    = 2.f * analytic.cellSize;
    const std::vector<glm::vec3> probes = {
        glm::vec3(0.f, 1.2f, 0.f), glm::vec3(1.1f, 0.f, 0.f),  glm::vec3(0.2f, -1.15f, 0.f),
        glm::vec3(0.f, 0.f, 0.f),  glm::vec3(0.9f, 0.6f, 0.f),
    };
    for (const glm::vec3& p : probes) {
        CHECK(std::fabs(mesh.sample(p) - analytic.sample(p)) < tol);
    }
    // Inside the mesh is negative.
    CHECK(mesh.sample(glm::vec3(0.f)) < 0.f);
}

TEST_CASE("fluids.cpu.surfaceFlowDownhill") {
    FluidParams params;
    params.gravity   = glm::vec3(0.f, -9.8f, 0.f);
    params.viscosity = 0.02f;
    FluidSimulation sim(2048, params);
    sim.setSdf(MeshSdf::makeSphere(glm::vec3(0.f), 1.f, glm::ivec3(32)));
    const int added = sim.spawnDrop(glm::vec3(0.f, 1.5f, 0.f), 0.25f, 200);
    REQUIRE(added == 200);

    // Spawn projects particles onto the surface (nothing inside the sphere).
    for (int i = 0; i < sim.particleCount(); ++i) {
        CHECK(distToCenter(sim.particles()[size_t(i)].pos) >= 1.f - 0.02f);
    }

    float meanY0 = 0.f;
    for (int i = 0; i < sim.particleCount(); ++i) meanY0 += sim.particles()[size_t(i)].pos.y;
    meanY0 /= float(sim.particleCount());

    for (int s = 0; s < 150; ++s) sim.step(1.f / 60.f);

    float meanY1       = 0.f;
    bool  allOnSurface = true;
    for (int i = 0; i < sim.particleCount(); ++i) {
        meanY1 += sim.particles()[size_t(i)].pos.y;
        if (distToCenter(sim.particles()[size_t(i)].pos) < 1.f - 0.03f) allOnSurface = false;
    }
    meanY1 /= float(sim.particleCount());

    // Water runs down the sphere.
    CHECK(meanY1 < meanY0 - 0.1f);
    // Projection keeps every particle glued to the surface.
    CHECK(allOnSurface);
}

TEST_CASE("fluids.cpu.viscosityDampsRelativeMotion") {
    FluidParams params;
    params.gravity    = glm::vec3(0.f);
    params.viscosity  = 0.5f;
    params.damping    = 0.f;
    params.iterations = 1;
    FluidSimulation sim(16, params);
    sim.setSdf(MeshSdf::makePlane(-2.f, glm::ivec3(8), 2.f));

    const int added = sim.spawnDrop(glm::vec3(0.f, 0.f, 0.f), 0.02f, 2);
    REQUIRE(added == 2);
    REQUIRE(sim.particleCount() == 2);
    // Co-locate the pair with opposing velocities: XSPH pulls them together.
    sim.particles()[0].pos = glm::vec3(0.f);
    sim.particles()[1].pos = glm::vec3(0.03f, 0.f, 0.f);
    sim.particles()[0].vel = glm::vec3(5.f, 0.f, 0.f);
    sim.particles()[1].vel = glm::vec3(-5.f, 0.f, 0.f);

    const float rel0 = glm::length(sim.particles()[0].vel - sim.particles()[1].vel);
    sim.step(1.f / 60.f);
    const float rel1 = glm::length(sim.particles()[0].vel - sim.particles()[1].vel);
    CHECK(rel1 < rel0);
}

TEST_CASE("fluids.cpu.densityPositiveNearNeighbors") {
    FluidParams     params;
    FluidSimulation sim(64, params);
    sim.setSdf(MeshSdf::makePlane(-2.f, glm::ivec3(8), 2.f));
    sim.spawnDrop(glm::vec3(0.f), 0.12f, 32);
    sim.step(1.f / 60.f);
    REQUIRE(sim.particleCount() > 0);
    for (int i = 0; i < sim.particleCount(); ++i) {
        CHECK(sim.densities()[size_t(i)] > 0.f);
    }
}

TEST_CASE("fluids.module.create") {
    auto* mod = Fluids::create();
    REQUIRE(mod != nullptr);
    CHECK_EQ(mod->getName(), std::string("Fluids"));
    FluidSimulator* sim = mod->newSimulator(1024);
    REQUIRE(sim != nullptr);
    CHECK_EQ(sim->getMaxParticles(), 1024);
    CHECK_EQ(mod->getSimulatorCount(), 1);
    CHECK_EQ(sim->getParticleCount(), 0);
    FluidSurfaceRenderer* r = mod->newSurfaceRenderer(32, 32);
    REQUIRE(r != nullptr);
    CHECK_EQ(mod->getRendererCount(), 1);
    CHECK_EQ(r->getWidth(), 32);
    r->render(*sim);
    CHECK_EQ(r->getHeight(), 32);
}

TEST_CASE("fluids.cpu.pbfSeparatesCompressedParticles") {
    FluidParams params;
    params.gravity       = glm::vec3(0.f);
    params.viscosity     = 0.02f;
    params.damping       = 0.2f;
    params.pbfIterations = 4;
    params.iterations    = 1;
    FluidSimulation sim(64, params);
    sim.setSdf(MeshSdf::makePlane(-2.f, glm::ivec3(8), 2.f));
    REQUIRE(sim.spawnDrop(glm::vec3(0.f), 0.03f, 12) == 12);

    auto meanPairDist = [&]() {
        double sum = 0.0;
        int    n   = 0;
        for (int i = 0; i < sim.particleCount(); ++i) {
            for (int j = i + 1; j < sim.particleCount(); ++j) {
                sum += glm::length(sim.particles()[size_t(i)].pos - sim.particles()[size_t(j)].pos);
                ++n;
            }
        }
        return float(sum / double(std::max(n, 1)));
    };

    const float before = meanPairDist();
    for (int s = 0; s < 40; ++s) sim.step(1.f / 60.f);
    const float after = meanPairDist();
    // A compressed cloud must relax toward rest density (particles spread).
    CHECK(after > before * 1.2f);
}

TEST_CASE("fluids.cpu.cohesionClumpsScatteredParticles") {
    FluidParams params;
    params.gravity       = glm::vec3(0.f);
    params.viscosity     = 0.08f;
    params.cohesion      = 0.5f;
    params.pbfIterations = 0;  // isolate the cohesion force
    params.damping       = 0.05f;
    FluidSimulation sim(64, params);
    sim.setSdf(MeshSdf::makePlane(-2.f, glm::ivec3(8), 2.f));
    REQUIRE(sim.spawnDrop(glm::vec3(0.f), 0.15f, 12) == 12);

    auto meanToCentroid = [&]() {
        glm::vec3 c(0.f);
        for (int i = 0; i < sim.particleCount(); ++i) c += sim.particles()[size_t(i)].pos;
        c /= float(sim.particleCount());
        double sum = 0.0;
        for (int i = 0; i < sim.particleCount(); ++i) sum += glm::length(sim.particles()[size_t(i)].pos - c);
        return float(sum / double(sim.particleCount()));
    };

    const float before = meanToCentroid();
    for (int s = 0; s < 90; ++s) sim.step(1.f / 60.f);
    const float after = meanToCentroid();
    // Surface tension (cohesion) pulls the scattered film into a tighter drop.
    CHECK(after < before * 0.75f);
}

TEST_CASE("fluids.cpu.adhesionKeepsDropOnUpsideDownSurface") {
    FluidParams sticky;
    sticky.gravity       = glm::vec3(0.f, -9.8f, 0.f);
    sticky.viscosity     = 0.05f;
    sticky.adhesion      = 1.5f;
    sticky.pbfIterations = 2;
    FluidSimulation sim(128, sticky);
    sim.setSdf(MeshSdf::makeSphere(glm::vec3(0.f), 1.f, glm::ivec3(32)));
    REQUIRE(sim.spawnDrop(glm::vec3(0.f, -1.45f, 0.f), 0.2f, 64) == 64);
    for (int s = 0; s < 90; ++s) sim.step(1.f / 60.f);

    float meanR = 0.f;
    for (int i = 0; i < sim.particleCount(); ++i) meanR += distToCenter(sim.particles()[size_t(i)].pos);
    meanR /= float(sim.particleCount());
    // Adhesion holds the drop on the underside of the sphere.
    CHECK(meanR < 1.08f);
}

TEST_CASE("fluids.cpu.yieldStressFreezesLowShearMotion") {
    auto run = [&](float yield) {
        FluidParams params;
        params.gravity       = glm::vec3(0.f);
        params.viscosity     = 0.01f;
        params.yieldStress   = yield;
        params.damping       = 0.f;
        params.pbfIterations = 0;
        FluidSimulation sim(64, params);
        sim.setSdf(MeshSdf::makePlane(-2.f, glm::ivec3(8), 2.f));
        REQUIRE(sim.spawnDrop(glm::vec3(0.f), 0.12f, 8) == 8);
        // Small deterministic velocities -> low shear rate.
        for (int i = 0; i < sim.particleCount(); ++i) {
            const float a                  = 0.02f;
            sim.particles()[size_t(i)].vel = glm::vec3(std::sin(float(i + 1)) * a, std::cos(float(i + 1) * 2.f) * a,
                                                       std::sin(float(i + 1) * 3.f) * a);
        }
        for (int s = 0; s < 30; ++s) sim.step(1.f / 60.f);
        glm::vec3 meanVelocity(0.f);
        for (int i = 0; i < sim.particleCount(); ++i) meanVelocity += sim.particles()[size_t(i)].vel;
        meanVelocity /= float(sim.particleCount());
        float ke = 0.f;
        for (int i = 0; i < sim.particleCount(); ++i) {
            const glm::vec3 relative = sim.particles()[size_t(i)].vel - meanVelocity;
            ke += 0.5f * glm::dot(relative, relative);
        }
        return ke;
    };

    const float keFree = run(0.f);
    const float keMud  = run(10.f);
    // Bingham yield stress turns low-shear flow into a quasi-solid.
    CHECK(keMud < keFree * 0.2f);
}

TEST_CASE("fluids.ssf.cpu.singleParticleDepthAndNormal") {
    FluidSurfaceParams params;
    // Odd dimensions place the particle center exactly on a pixel center.
    params.width            = 65;
    params.height           = 65;
    params.eye              = glm::vec3(0.f, 0.f, -3.f);
    params.particleRadius   = 0.08f;
    params.smoothIterations = 0;
    FluidSurfaceRenderer r(params, false);
    r.render(std::vector<glm::vec3>{glm::vec3(0.f)}, 0.08f);

    const int cx = 32;
    const int cy = 32;
    REQUIRE(std::fabs(r.depth()[size_t(cy * 65 + cx)] - (3.f - 0.08f)) < 1e-5f);
    const glm::vec3 n = r.normals()[size_t(cy * 65 + cx)];
    REQUIRE(n.z > 0.99f);
    REQUIRE(std::fabs(r.thickness()[size_t(cy * 65 + cx)] - 0.16f) < 1e-5f);
    REQUIRE(r.color()[size_t(cy * 65 + cx) * 4u + 3] > 0u);
}

TEST_CASE("fluids.ssf.cpu.waterVsMudShading") {
    auto renderMode = [](int mode) {
        FluidSurfaceParams params;
        params.width            = 64;
        params.height           = 64;
        params.eye              = glm::vec3(0.f, 0.f, -3.f);
        params.particleRadius   = 0.08f;
        params.smoothIterations = 0;
        params.mode             = mode;
        FluidSurfaceRenderer r(params, false);
        r.render(std::vector<glm::vec3>{glm::vec3(0.f)}, 0.08f);
        const size_t idx = size_t(32 * 64 + 32) * 4u;
        return std::array<uint8_t, 3>{r.color()[idx], r.color()[idx + 1], r.color()[idx + 2]};
    };

    const auto water = renderMode(0);
    const auto mud   = renderMode(1);
    // Water is blue-dominant; mud is brown-dominant (red > blue).
    CHECK(water[2] > water[0]);
    CHECK(mud[0] > mud[2]);
}

TEST_CASE("fluids.gpu.surfaceFlow") {
    if (!tryInitHeadlessGfx()) return;
    FluidParams params;
    params.gravity   = glm::vec3(0.f, -9.8f, 0.f);
    params.viscosity = 0.02f;
    FluidSimulator sim(2048, params, true);
    sim.setSdfSphere(0.f, 0.f, 0.f, 1.f, 32);
    sim.spawnDrop(glm::vec3(0.f, 1.5f, 0.f), 0.25f, 200);
    sim.step(1.f / 60.f);

    if (!sim.usingGpu()) return;  // no Vulkan device / no shader compiler — skip

    std::vector<float> dens;
    sim.readDensities(dens);
    REQUIRE(dens.size() == size_t(sim.getParticleCount()));
    for (float d : dens) CHECK(d > 0.f);

    float                  meanY0 = 0.f;
    std::vector<glm::vec3> pos;
    sim.readPositions(pos);
    for (const glm::vec3& p : pos) meanY0 += p.y;
    meanY0 /= float(pos.size());

    for (int s = 0; s < 120; ++s) sim.step(1.f / 60.f);
    sim.readPositions(pos);

    float meanY1       = 0.f;
    bool  allOnSurface = true;
    float maxP         = 0.f;
    for (const glm::vec3& p : pos) {
        meanY1 += p.y;
        maxP = std::max(maxP, distToCenter(p));
        if (distToCenter(p) < 1.f - 0.05f) allOnSurface = false;
    }
    meanY1 /= float(pos.size());
    CHECK(meanY1 < meanY0 - 0.05f);
    CHECK(allOnSurface);
    // Particles must stay glued to the SDF; a broken upload would let the
    // projection push them straight off the model (regression guard).
    CHECK(maxP < 3.f);
}

TEST_CASE("fluids.gpu.pbfCohesionClusterRemainsBounded") {
    if (!tryInitHeadlessGfx()) return;
    FluidParams params;
    params.gravity       = glm::vec3(0.f, -9.8f, 0.f);
    params.viscosity     = 0.05f;
    params.cohesion      = 0.5f;
    params.adhesion      = 0.5f;
    params.pbfIterations = 2;
    FluidSimulator sim(1024, params, true);
    sim.setSdfSphere(0.f, 0.f, 0.f, 1.f, 32);
    sim.spawnDrop(glm::vec3(0.f, 1.45f, 0.f), 0.22f, 128);
    sim.step(1.f / 60.f);
    if (!sim.usingGpu()) return;

    auto clusterRadius = [&]() {
        std::vector<glm::vec3> pos;
        sim.readPositions(pos);
        glm::vec3 c(0.f);
        for (const glm::vec3& p : pos) c += p;
        c /= float(pos.size());
        double sum = 0.0;
        for (const glm::vec3& p : pos) sum += glm::length(p - c);
        return float(sum / double(pos.size()));
    };

    const float before = clusterRadius();
    for (int s = 0; s < 60; ++s) sim.step(1.f / 60.f);
    const float after = clusterRadius();
    // Gravity, density relaxation and cohesion compete on the curved surface.
    // The GPU contract is bounded retention; CPU coverage above separately proves
    // that cohesion contracts a force-isolated cloud.
    CHECK(after < before * 1.25f);

    std::vector<glm::vec3> pos;
    sim.readPositions(pos);
    bool allOnSurface = true;
    for (const glm::vec3& p : pos) {
        if (distToCenter(p) < 1.f - 0.06f) allOnSurface = false;
    }
    CHECK(allOnSurface);
}

TEST_CASE("fluids.ssf.gpu.pipeline") {
    if (!tryInitHeadlessGfx()) return;
    FluidSurfaceParams params;
    params.width            = 64;
    params.height           = 64;
    params.eye              = glm::vec3(0.f, 0.f, -3.f);
    params.particleRadius   = 0.08f;
    params.smoothIterations = 0;
    FluidSurfaceRenderer r(params, true);
    r.render(std::vector<glm::vec3>{glm::vec3(0.f)}, 0.08f);
    if (!r.usingGpu()) return;

    const int cx = 32;
    const int cy = 32;
    CHECK(std::fabs(r.depth()[size_t(cy * 64 + cx)] - 3.f) < 0.25f);
    CHECK(r.normals()[size_t(cy * 64 + cx)].z > 0.8f);
    CHECK(r.thickness()[size_t(cy * 64 + cx)] > 0.f);
    CHECK(r.color()[size_t(cy * 64 + cx) * 4u + 3] > 0u);

    r.writePpm("fluid_ssf_gpu.ppm");
    std::ifstream f("fluid_ssf_gpu.ppm", std::ios::binary | std::ios::ate);
    REQUIRE(f.good());
    const long size = long(f.tellg());
    f.close();
    CHECK(size > 100);
    std::remove("fluid_ssf_gpu.ppm");
}