using System.Numerics;

namespace FluidSimV2
{
    public class Pipe1D
    {
        private float _length;
        private int _cells;
        private float _dt;
        private float _dx;

        private float[] _rho; //rho (density) 
        private float[] _rho_u; //rho * u (momentum)
        private float[] _E; // total energy

        private float[] _u;
        private float[] _P;
        private float[] _c;

        float[] _flux_rho;
        float[] _flux_rho_u;
        float[] _flux_E;

        public Pipe1D(int sampleRate, float length)
        {
            float c = 345; // m/s

            _length = length;

            _dt = 1f / sampleRate; // no sub stepping yet, only at planning step
            _dx = c / sampleRate; // approximate for cells needed
            _cells = (int)(_length / _dx) + 2; //2 ghost cells, always 0 and n-1

            _rho = new float[_cells];
            _rho_u = new float[_cells];
            _E = new float[_cells];

            _u = new float[_cells];
            _P = new float[_cells];
            _c = new float[_cells];

            // +1 for cell faces
            _flux_rho = new float[_cells + 1];
            _flux_rho_u = new float[_cells + 1];
            _flux_E = new float[_cells + 1];

            InitializeAmbient();
        }

        public void InitializeAmbient(float ambientPressure = 101325f)
        {
            float R = 287.05f;
            float T = 293.15f; // 20 celsius
            float ambientRho = ambientPressure / (R * T);
            float ambientEnergy = ambientPressure / 0.4f;

            for(int i = 0; i < _cells; i++)
            {
                _rho[i] = ambientRho;
                _rho_u[i] = 0f;
                _E[i] = ambientEnergy;
            }
        }

        public void Step()
        {
            UpdateDerivedValues();

            for (int i = 0; i <= _cells; i++){
                //var stateL = GetStateL(i); // Returns rho, rho_u, energy, u, P, c
                //var stateR = GetStateR(i);

                // Lax-Friedrichs
                //var flux = CalculateFlux(stateL, stateR);
                
                //flux_rho[i] = flux.rho;
                //flux_rho_u[i] = flux.rho_u;
                //flux_energy[i] = flux.energy;
            }

            UpdateCells(_dt, _dx);

            // if boundary with a 0D volume, update it from flux
        }

        public void UpdateDerivedValues()
        {
            Vector<float> gammaMinusOne = new Vector<float>(1.4f - 1.0f);
            Vector<float> gamma = new Vector<float>(1.4f);
            Vector<float> half = new Vector<float>(0.5f);

            int vectorSize = Vector<float>.Count;
            int i = 0;

            /* using vector for SIMD
            Process as many full vector‑width blocks as possible.
            The condition i <= _cells - vectorSize ensures the block starting at i
            covers indices i .. i+vectorSize-1, all strictly less than _cells.
            Example: vectorSize=8, _cells=70 → last SIMD i=56 covers 56…63.
            After that i becomes 64 > 62, so the scalar loop handles 64…69.
            */
            for (; i <= _cells - vectorSize; i += vectorSize)
            {
                // load old states as vectors
                Vector<float> v_rho = new Vector<float>(_rho, i);
                Vector<float> v_rho_u = new Vector<float>(_rho_u, i);
                Vector<float> v_E = new Vector<float>(_E, i);

                // u = (rho_u) / rho
                Vector<float> v_u = v_rho_u / v_rho;
                v_u.CopyTo(_u, i);

                // P = (gamma-1) * (E - 0.5 * rho * u^2)
                // P = (gamma - 1) * internal energy
                // Equivalent to total energy (E) - kinetic energy (0.5 * rho * u^2)
                Vector<float> v_P = gammaMinusOne * (v_E - (half * v_rho * v_u * v_u));
                v_P.CopyTo(_P, i);


                // c = sqrt(gamma * P / rho)
                Vector.SquareRoot(gamma * v_P / v_rho).CopyTo(_c, i);
            }

            // handle rest of cells (simd can only handle multiples of vectorSize, see comment above)
            for (; i < _cells; i++)
            {
                float rho = _rho[i];
                float rho_u = _rho_u[i];
                float E = _E[i];

                _u[i] = rho_u / rho;
                
                float kinE = 0.5f * rho * (_u[i] * _u[i]);
                _P[i] = 0.4f * (E - kinE);

                _c[i] = MathF.Sqrt(1.4f * _P[i] / rho);
            }

        }

        public void UpdateCells(float dt, float dx)
        {
            float dt_dx = dt / dx;

            Vector<float> v_dt_dx = new Vector<float>(dt_dx);

            int vectorSize = Vector<float>.Count;
            int i = 0;

            // using vector for SIMD
            for (; i <= _rho.Length - vectorSize; i += vectorSize)
            {
                Vector<float> v_rho = new Vector<float>(_rho, i);
                Vector<float> v_rho_u = new Vector<float>(_rho_u, i);
                Vector<float> v_E = new Vector<float>(_E, i);

                Vector<float> v_fluxL_rho = new Vector<float>(_flux_rho, i);
                Vector<float> v_fluxR_rho = new Vector<float>(_flux_rho, i + 1);
                Vector<float> v_fluxL_rho_u = new Vector<float>(_flux_rho_u, i);
                Vector<float> v_fluxR_rho_u = new Vector<float>(_flux_rho_u, i + 1);
                Vector<float> v_fluxL_E = new Vector<float>(_flux_E, i);
                Vector<float> v_fluxR_E = new Vector<float>(_flux_E, i + 1);

                (v_rho + v_dt_dx * (v_fluxL_rho - v_fluxR_rho)).CopyTo(_rho, i);
                (v_rho_u + v_dt_dx * (v_fluxL_rho_u - v_fluxR_rho_u)).CopyTo(_rho_u, i);
                (v_E + v_dt_dx * (v_fluxL_E - v_fluxR_E)).CopyTo(_E, i);
            }

            for (; i < _cells; i++)
            {
                int left = i;
                int right = i+1;

                float rho = _rho[i];
                float rho_u = _rho_u[i];
                float E = _E[i];

                _rho[i] = rho + dt_dx * (_flux_rho[left] - _flux_rho[right]);
                _rho_u[i] = rho_u + dt_dx * (_flux_rho_u[left] - _flux_rho_u[right]);
                _E[i] = E + dt_dx * (_flux_E[left] - _flux_E[right]);
            }
        }

    }
}