using OpenTK.Mathematics;

namespace fluidsC_;

public class FluidSimulation
{
    private readonly int _sizeX;
    private readonly int _sizeY;

    // Fluid fields
    private float[,] _density;
    private float[,] _densityOld;

    // Velocity fields
    private float[,] _vx;
    private float[,] _vy;
    private float[,] _vxOld;
    private float[,] _vyOld;

    public FluidSimulation(int sizeX, int sizeY)
    {
        _sizeX = sizeX;
        _sizeY = sizeY;

        // Initialize arrays in constructor to fix nullable warnings
        _density = new float[_sizeX, _sizeY];
        _densityOld = new float[_sizeX, _sizeY];
        _vx = new float[_sizeX, _sizeY];
        _vy = new float[_sizeX, _sizeY];
        _vxOld = new float[_sizeX, _sizeY];
        _vyOld = new float[_sizeX, _sizeY];
    }

    public void AddDensity(int x, int y, float amount)
    {
        if (x >= 0 && x < _sizeX && y >= 0 && y < _sizeY)
        {
            _density[x, y] += amount;
        }
    }

    public void AddVelocity(int x, int y, float amountX, float amountY)
    {
        if (x >= 0 && x < _sizeX && y >= 0 && y < _sizeY)
        {
            _vx[x, y] += amountX;
            _vy[x, y] += amountY;
        }
    }

    public void Step(float dt, float diffusion, float viscosity)
    {
        Diffuse(1, _vxOld, _vx, viscosity, dt);
        Diffuse(2, _vyOld, _vy, viscosity, dt);

        Project(_vxOld, _vyOld, _vx, _vy);

        Advect(1, _vx, _vxOld, _vxOld, _vyOld, dt);
        Advect(2, _vy, _vyOld, _vxOld, _vyOld, dt);

        Project(_vx, _vy, _vxOld, _vyOld);

        Diffuse(0, _densityOld, _density, diffusion, dt);
        Advect(0, _density, _densityOld, _vx, _vy, dt);
    }

    private void Diffuse(int b, float[,] x, float[,] x0, float diff, float dt)
    {
        float a = dt * diff * (_sizeX - 2) * (_sizeY - 2);
        LinearSolve(b, x, x0, a, 1 + 4 * a);
    }

    private void LinearSolve(int b, float[,] x, float[,] x0, float a, float c)
    {
        for (int k = 0; k < 20; k++)
        {
            for (int i = 1; i < _sizeX - 1; i++)
            {
                for (int j = 1; j < _sizeY - 1; j++)
                {
                    x[i, j] = (x0[i, j] + a * (x[i - 1, j] + x[i + 1, j] + x[i, j - 1] + x[i, j + 1])) / c;
                }
            }
            SetBoundary(b, x);
        }
    }

    private void Project(float[,] velocX, float[,] velocY, float[,] p, float[,] div)
    {
        for (int i = 1; i < _sizeX - 1; i++)
        {
            for (int j = 1; j < _sizeY - 1; j++)
            {
                div[i, j] = -0.5f * (
                    velocX[i + 1, j] - velocX[i - 1, j] +
                    velocY[i, j + 1] - velocY[i, j - 1]
                ) / _sizeX;
                p[i, j] = 0;
            }
        }

        SetBoundary(0, div);
        SetBoundary(0, p);
        LinearSolve(0, p, div, 1, 4);

        for (int i = 1; i < _sizeX - 1; i++)
        {
            for (int j = 1; j < _sizeY - 1; j++)
            {
                velocX[i, j] -= 0.5f * (p[i + 1, j] - p[i - 1, j]) * _sizeX;
                velocY[i, j] -= 0.5f * (p[i, j + 1] - p[i, j - 1]) * _sizeY;
            }
        }

        SetBoundary(1, velocX);
        SetBoundary(2, velocY);
    }

    private void Advect(int b, float[,] d, float[,] d0, float[,] velocX, float[,] velocY, float dt)
    {
        float dt0 = dt * (_sizeX - 2);

        for (int i = 1; i < _sizeX - 1; i++)
        {
            for (int j = 1; j < _sizeY - 1; j++)
            {
                float x = i - dt0 * velocX[i, j];
                float y = j - dt0 * velocY[i, j];

                x = Math.Clamp(x, 0.5f, _sizeX - 1.5f);
                y = Math.Clamp(y, 0.5f, _sizeY - 1.5f);

                int i0 = (int)Math.Floor(x);
                int i1 = i0 + 1;
                int j0 = (int)Math.Floor(y);
                int j1 = j0 + 1;

                float s1 = x - i0;
                float s0 = 1 - s1;
                float t1 = y - j0;
                float t0 = 1 - t1;

                d[i, j] = s0 * (t0 * d0[i0, j0] + t1 * d0[i0, j1]) +
                          s1 * (t0 * d0[i1, j0] + t1 * d0[i1, j1]);
            }
        }

        SetBoundary(b, d);
    }

    private void SetBoundary(int b, float[,] x)
    {
        for (int i = 1; i < _sizeX - 1; i++)
        {
            x[i, 0] = b == 2 ? -x[i, 1] : x[i, 1];
            x[i, _sizeY - 1] = b == 2 ? -x[i, _sizeY - 2] : x[i, _sizeY - 2];
        }

        for (int j = 1; j < _sizeY - 1; j++)
        {
            x[0, j] = b == 1 ? -x[1, j] : x[1, j];
            x[_sizeX - 1, j] = b == 1 ? -x[_sizeX - 2, j] : x[_sizeX - 2, j];
        }

        x[0, 0] = 0.5f * (x[1, 0] + x[0, 1]);
        x[0, _sizeY - 1] = 0.5f * (x[1, _sizeY - 1] + x[0, _sizeY - 2]);
        x[_sizeX - 1, 0] = 0.5f * (x[_sizeX - 2, 0] + x[_sizeX - 1, 1]);
        x[_sizeX - 1, _sizeY - 1] = 0.5f * (x[_sizeX - 2, _sizeY - 1] + x[_sizeX - 1, _sizeY - 2]);
    }

    public float GetDensity(int x, int y)
    {
        if (x >= 0 && x < _sizeX && y >= 0 && y < _sizeY)
            return _density[x, y];
        return 0;
    }

    public void FadeDensity(float amount)
    {
        for (int i = 0; i < _sizeX; i++)
        {
            for (int j = 0; j < _sizeY; j++)
            {
                _density[i, j] = Math.Max(0, _density[i, j] - amount);
            }
        }
    }
}