using OpenTK.Mathematics;

namespace fluidsC_;

public class FluidSimulation3D
{
    private readonly int _sizeX;
    private readonly int _sizeY;
    private readonly int _sizeZ;

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

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

    // Obstacle field
    private bool[,,] _obstacles;

    public FluidSimulation3D(int sizeX, int sizeY, int sizeZ)
    {
        _sizeX = sizeX;
        _sizeY = sizeY;
        _sizeZ = sizeZ;

        // Initialize arrays in constructor
        _density = new float[_sizeX, _sizeY, _sizeZ];
        _densityOld = new float[_sizeX, _sizeY, _sizeZ];
        _vx = new float[_sizeX, _sizeY, _sizeZ];
        _vy = new float[_sizeX, _sizeY, _sizeZ];
        _vz = new float[_sizeX, _sizeY, _sizeZ];
        _vxOld = new float[_sizeX, _sizeY, _sizeZ];
        _vyOld = new float[_sizeX, _sizeY, _sizeZ];
        _vzOld = new float[_sizeX, _sizeY, _sizeZ];
        _obstacles = new bool[_sizeX, _sizeY, _sizeZ];

        CreateDefaultObstacles();
    }

    private void CreateDefaultObstacles()
    {
        // Create some default obstacles
        int centerX = _sizeX / 2;
        int centerZ = _sizeZ / 2;

        // Create a sphere obstacle in the center
        int sphereRadius = Math.Min(_sizeX, Math.Min(_sizeY, _sizeZ)) / 6;
        for (int x = 0; x < _sizeX; x++)
        {
            for (int y = 0; y < _sizeY; y++)
            {
                for (int z = 0; z < _sizeZ; z++)
                {
                    float dx = x - centerX;
                    float dy = y - _sizeY / 3;
                    float dz = z - centerZ;
                    float dist = MathF.Sqrt(dx * dx + dy * dy + dz * dz);

                    if (dist <= sphereRadius)
                    {
                        _obstacles[x, y, z] = true;
                    }

                    // Create walls
                    if (x == 0 || x == _sizeX - 1 || z == 0 || z == _sizeZ - 1)
                    {
                        _obstacles[x, y, z] = true;
                    }

                    // Create floor
                    if (y == 0)
                    {
                        _obstacles[x, y, z] = true;
                    }
                }
            }
        }
    }

    public void AddDensity(int x, int y, int z, float amount)
    {
        if (IsInBounds(x, y, z) && !_obstacles[x, y, z])
        {
            _density[x, y, z] += amount;
        }
    }

    public void AddVelocity(int x, int y, int z, float amountX, float amountY, float amountZ)
    {
        if (IsInBounds(x, y, z) && !_obstacles[x, y, z])
        {
            _vx[x, y, z] += amountX;
            _vy[x, y, z] += amountY;
            _vz[x, y, z] += amountZ;
        }
    }

    public void AddObstacle(int x, int y, int z, bool obstacle)
    {
        if (IsInBounds(x, y, z))
        {
            _obstacles[x, y, z] = obstacle;
        }
    }

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

        // Fixed Project3D calls - use correct parameters
        Project3D(_vxOld, _vyOld, _vzOld, _vxOld, _vyOld, _vzOld);

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

        // Fixed Project3D calls - use correct parameters
        Project3D(_vx, _vy, _vz, _vx, _vy, _vz);

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

    private void Diffuse(int b, float[,,] x, float[,,] x0, float diff, float dt)
    {
        float a = dt * diff * (_sizeX - 2) * (_sizeY - 2) * (_sizeZ - 2);
        LinearSolve(b, x, x0, a, 1 + 6 * 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++)
                {
                    for (int l = 1; l < _sizeZ - 1; l++)
                    {
                        if (!_obstacles[i, j, l])
                        {
                            x[i, j, l] = (x0[i, j, l] + a * (
                                x[i - 1, j, l] + x[i + 1, j, l] +
                                x[i, j - 1, l] + x[i, j + 1, l] +
                                x[i, j, l - 1] + x[i, j, l + 1])) / c;
                        }
                    }
                }
            }
            SetBoundary(b, x);
        }
    }

    private void Project3D(float[,,] velocX, float[,,] velocY, float[,,] velocZ, float[,,] p, float[,,] div, float[,,] temp)
    {
        for (int i = 1; i < _sizeX - 1; i++)
        {
            for (int j = 1; j < _sizeY - 1; j++)
            {
                for (int l = 1; l < _sizeZ - 1; l++)
                {
                    if (!_obstacles[i, j, l])
                    {
                        div[i, j, l] = -0.5f * (
                            velocX[i + 1, j, l] - velocX[i - 1, j, l] +
                            velocY[i, j + 1, l] - velocY[i, j - 1, l] +
                            velocZ[i, j, l + 1] - velocZ[i, j, l - 1]
                        ) / _sizeX;
                        p[i, j, l] = 0;
                    }
                }
            }
        }

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

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

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

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

        for (int i = 1; i < _sizeX - 1; i++)
        {
            for (int j = 1; j < _sizeY - 1; j++)
            {
                for (int l = 1; l < _sizeZ - 1; l++)
                {
                    if (!_obstacles[i, j, l])
                    {
                        float x = i - dt0 * velocX[i, j, l];
                        float y = j - dt0 * velocY[i, j, l];
                        float z = l - dt0 * velocZ[i, j, l];

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

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

                        float s1 = x - i0;
                        float s0 = 1 - s1;
                        float t1 = y - j0;
                        float t0 = 1 - t1;
                        float u1 = z - l0;
                        float u0 = 1 - u1;

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

        SetBoundary(b, d);
    }

    private void SetBoundary(int b, float[,,] x)
    {
        // Set boundary conditions with obstacles
        for (int i = 1; i < _sizeX - 1; i++)
        {
            for (int j = 1; j < _sizeY - 1; j++)
            {
                for (int l = 1; l < _sizeZ - 1; l++)
                {
                    if (_obstacles[i, j, l])
                    {
                        // Reflect velocity at obstacles
                        if (_obstacles[i - 1, j, l]) x[i, j, l] = b == 1 ? -x[i + 1, j, l] : x[i + 1, j, l];
                        if (_obstacles[i + 1, j, l]) x[i, j, l] = b == 1 ? -x[i - 1, j, l] : x[i - 1, j, l];
                        if (_obstacles[i, j - 1, l]) x[i, j, l] = b == 2 ? -x[i, j + 1, l] : x[i, j + 1, l];
                        if (_obstacles[i, j + 1, l]) x[i, j, l] = b == 2 ? -x[i, j - 1, l] : x[i, j - 1, l];
                        if (_obstacles[i, j, l - 1]) x[i, j, l] = b == 3 ? -x[i, j, l + 1] : x[i, j, l + 1];
                        if (_obstacles[i, j, l + 1]) x[i, j, l] = b == 3 ? -x[i, j, l - 1] : x[i, j, l - 1];
                    }
                }
            }
        }
    }

    private bool IsInBounds(int x, int y, int z)
    {
        return x >= 0 && x < _sizeX && y >= 0 && y < _sizeY && z >= 0 && z < _sizeZ;
    }

    public float GetDensity(int x, int y, int z)
    {
        if (IsInBounds(x, y, z))
            return _density[x, y, z];
        return 0;
    }

    public bool GetObstacle(int x, int y, int z)
    {
        if (IsInBounds(x, y, z))
            return _obstacles[x, y, z];
        return false;
    }

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

    public (int X, int Y, int Z) GetSize() => (_sizeX, _sizeY, _sizeZ);
}