using System;
using System.Diagnostics;
using System.Numerics;

namespace FluidSim.Core
{
    public class PipeSystem
    {
        // ---------- Master arrays ----------
        private float[] _rho, _rhou, _E, _Y;
        private readonly float[] _area;
        private readonly float[] _dx;
        private readonly int[] _pipeStart;
        private readonly int[] _pipeEnd;
        private readonly int _totalCells;          // original cell count (visible)
        private readonly int _allCells;            // total allocated (padded to Vector<float>.Count)
        private readonly int _pipeCount;

        // Derived state – _p is kept for visualization, _c is gone
        private float[] _p;

        // Flux arrays (size = _allCells + 1)
        private float[] _fluxM, _fluxP, _fluxE, _fluxY;

        // Damping and relaxation (computed on‑the‑fly only if used)
        private float[] _dampingFactors;
        private float[] _relaxFactors;
        private bool _applyDamping;
        private bool _applyRelax;

        // Ghost buffer
        private readonly GhostBuffer _ghost;

        // Wall mask – precomputed once
        private readonly bool[] _isWallFace;

        // ---------- Physical constants ----------
        private const float Gamma = 1.4f;
        private const float Gm1 = 0.4f;
        private const float Gm1Inv = 1f / Gm1;          // 2.5
        private const float GammaOverGm1 = Gamma / Gm1; // 3.5
        private float _coeffBase;
        private float _relaxRate;
        private float _ambientPressure = 101325f;
        private float _ambientEnergyRef;

        public float DampingMultiplier
        {
            set
            {
                _coeffBase = 0.1f * value;
                _applyDamping = _coeffBase != 0f;
            }
        }
        public float EnergyRelaxationRate
        {
            set
            {
                _relaxRate = value;
                _applyRelax = _relaxRate != 0f;
            }
        }
        public float AmbientPressure
        {
            set
            {
                _ambientPressure = value;
                _ambientEnergyRef = value * Gm1Inv;
            }
        }

        // ---------- Profiling ----------
        public bool EnableProfiling { get; set; }
        private long _profFluxTicks;
        private long _profUpdateTicks;
        private long _profCallCount;

        // ---------- Construction ----------
        public PipeSystem(int totalCells, int[] pipeStart, int[] pipeEnd,
                          float[] area, float[] dx,
                          float initialRho, float initialU, float initialP)
        {
            _pipeStart = pipeStart;
            _pipeEnd = pipeEnd;
            _pipeCount = pipeStart.Length;
            _totalCells = totalCells;
            _area = area;
            _dx = dx;

            // Pad to SIMD width so all vectorized loops cover the whole data
            int vecSize = Vector<float>.Count;
            _allCells = totalCells % vecSize == 0 ? totalCells : totalCells + vecSize - (totalCells % vecSize);

            _rho = new float[_allCells];
            _rhou = new float[_allCells];
            _E = new float[_allCells];
            _Y = new float[_allCells];
            _p = new float[_allCells];          // pressure for drawing
            int faceCount = _allCells + 1;
            _fluxM = new float[faceCount];
            _fluxP = new float[faceCount];
            _fluxE = new float[faceCount];
            _fluxY = new float[faceCount];

            _dampingFactors = new float[_allCells];
            _relaxFactors = new float[_allCells];
            _applyDamping = _coeffBase != 0f;
            _applyRelax = _relaxRate != 0f;

            _ghost = new GhostBuffer(_pipeCount);
            _ambientEnergyRef = initialP * Gm1Inv;

            // Pre‑compute wall face flags: each face that sits between two different pipes is a wall
            _isWallFace = new bool[faceCount];
            for (int f = 1; f < _totalCells; f++)
            {
                for (int p = 0; p < _pipeCount; p++)
                {
                    if (f == _pipeEnd[p] && f < _totalCells)
                    {
                        _isWallFace[f] = true;
                        break;
                    }
                }
            }

            // Initialize uniform state
            float initE = initialP / (Gm1 * initialRho);
            float rhoE = initialRho * initE + 0.5f * initialRho * initialU * initialU;
            for (int i = 0; i < totalCells; i++)
            {
                _rho[i] = initialRho;
                _rhou[i] = initialRho * initialU;
                _E[i] = rhoE;
                _Y[i] = 1f;
            }
        }

        // ---------- Ghost setters (for BoundarySystem) ----------
        public void SetGhostLeft(int pipeIndex, float rho, float u, float p, float y)
            => _ghost.Set(pipeIndex, true, rho, u, p, y);
        public void SetGhostRight(int pipeIndex, float rho, float u, float p, float y)
            => _ghost.Set(pipeIndex, false, rho, u, p, y);

        // ---------- Public read methods ----------
        public int TotalCells => _totalCells;
        public int PipeCount => _pipeCount;
        public int GetPipeStart(int pipeIdx) => _pipeStart[pipeIdx];
        public int GetPipeEnd(int pipeIdx) => _pipeEnd[pipeIdx];
        public float GetCellPressure(int i) => _p[i];
        public float GetCellDensity(int i) => _rho[i];
        public float GetCellVelocity(int i)
        {
            float rho = _rho[i];
            return rho > 1e-12f ? _rhou[i] / rho : 0f;
        }
        public float GetCellAirFraction(int i) => _Y[i];

        public (float rho, float u, float p) GetInteriorStateLeft(int pipeIdx)
        {
            int i = _pipeStart[pipeIdx];
            float rho = _rho[i];
            float rhou = _rhou[i];
            float u = rhou / MathF.Max(rho, 1e-12f);
            float p = Gm1 * (_E[i] - 0.5f * rhou * u);
            return (rho, u, p);
        }
        public (float rho, float u, float p) GetInteriorStateRight(int pipeIdx)
        {
            int i = _pipeEnd[pipeIdx] - 1;
            float rho = _rho[i];
            float rhou = _rhou[i];
            float u = rhou / MathF.Max(rho, 1e-12f);
            float p = Gm1 * (_E[i] - 0.5f * rhou * u);
            return (rho, u, p);
        }
        public float GetInteriorAirFractionLeft(int pipeIdx) => _Y[_pipeStart[pipeIdx]];
        public float GetInteriorAirFractionRight(int pipeIdx) => _Y[_pipeEnd[pipeIdx] - 1];

        public void SetCellState(int i, float rho, float u, float p, float y = 1f)
        {
            if (i < 0 || i >= _totalCells) return;
            _rho[i] = rho;
            _rhou[i] = rho * u;
            _E[i] = p * Gm1Inv + 0.5f * rho * u * u;
            _Y[i] = y;
        }

        // ---------- Main step ----------
        public void SimulateStep(float dt)
        {
            long t0 = 0, t1 = 0;
            if (EnableProfiling)
            {
                _profCallCount++;
                t0 = Stopwatch.GetTimestamp();
            }

            ComputeFluxes(dt);

            if (EnableProfiling)
            {
                t1 = Stopwatch.GetTimestamp();
                _profFluxTicks += (t1 - t0);
                t0 = t1;
            }

            UpdateCells(dt);

            if (EnableProfiling)
            {
                t1 = Stopwatch.GetTimestamp();
                _profUpdateTicks += (t1 - t0);
            }
        }

        // ---------- Flux computation: fuses primitive calculation and flux evaluation ----------
        private void ComputeFluxes(float dt)
        {
            float fm, fp, fe;
            int vecSize = Vector<float>.Count;

            // ---- 1. Left ghost boundaries ----
            for (int p = 0; p < _pipeCount; p++)
            {
                int idx = _pipeStart[p];
                int ghostIdx = p * 2;
                float rL = _ghost.Rho[ghostIdx];
                float uL = _ghost.U[ghostIdx];
                float pL = _ghost.P[ghostIdx];
                float YL = _ghost.Y[ghostIdx];
                float cL = MathF.Sqrt(Gamma * pL / MathF.Max(rL, 1e-12f));

                float rR = _rho[idx], rhouR = _rhou[idx];
                float invRhoR = MathF.ReciprocalEstimate(MathF.Max(rR, 1e-12f));
                float uR = rhouR * invRhoR;
                float pR = Gm1 * (_E[idx] - 0.5f * rhouR * uR);
                float cR = MathF.Sqrt(Gamma * pR * invRhoR);
                float YR = _Y[idx];

                // store pressure for cell idx
                _p[idx] = pR;

                LaxFlux(rL, uL, pL, cL, rR, uR, pR, cR, out fm, out fp, out fe);
                _fluxM[idx] = fm; _fluxP[idx] = fp; _fluxE[idx] = fe;

                float alpha = MathF.Max(MathF.Abs(uL) + cL, MathF.Abs(uR) + cR);
                ScalarFlux(rL, uL, YL, rR, uR, YR, alpha, out float fy);
                _fluxY[idx] = fy;
            }

            // ---- 2. Right ghost boundaries ----
            for (int p = 0; p < _pipeCount; p++)
            {
                int idx = _pipeEnd[p] - 1;
                int face = idx + 1;
                int ghostIdx = p * 2 + 1;
                float rR = _ghost.Rho[ghostIdx];
                float uR = _ghost.U[ghostIdx];
                float pR = _ghost.P[ghostIdx];
                float YR = _ghost.Y[ghostIdx];
                float cR = MathF.Sqrt(Gamma * pR / MathF.Max(rR, 1e-12f));

                float rL = _rho[idx], rhouL = _rhou[idx];
                float invRhoL = MathF.ReciprocalEstimate(MathF.Max(rL, 1e-12f));
                float uL = rhouL * invRhoL;
                float pL = Gm1 * (_E[idx] - 0.5f * rhouL * uL);
                float cL = MathF.Sqrt(Gamma * pL * invRhoL);
                float YL = _Y[idx];

                // store pressure for cell idx
                _p[idx] = pL;

                LaxFlux(rL, uL, pL, cL, rR, uR, pR, cR, out fm, out fp, out fe);
                _fluxM[face] = fm; _fluxP[face] = fp; _fluxE[face] = fe;

                float alpha = MathF.Max(MathF.Abs(uL) + cL, MathF.Abs(uR) + cR);
                ScalarFlux(rL, uL, YL, rR, uR, YR, alpha, out float fy);
                _fluxY[face] = fy;
            }

            // ---- 3. Interior faces – vectorised SIMD ----
            for (int face = 1; face < _totalCells; face++)
            {
                // Handle walls (rare) with scalar code
                if (_isWallFace[face])
                {
                    int iL = face - 1;
                    float rL = _rho[iL], rhouL = _rhou[iL];
                    float invRhoL = MathF.ReciprocalEstimate(MathF.Max(rL, 1e-12f));
                    float uL = rhouL * invRhoL;
                    float pL = Gm1 * (_E[iL] - 0.5f * rhouL * uL);
                    float cL = MathF.Sqrt(Gamma * pL * invRhoL);
                    _p[iL] = pL;

                    LaxFlux(rL, uL, pL, cL, rL, -uL, pL, cL, out fm, out fp, out fe);
                    _fluxM[face] = fm; _fluxP[face] = fp; _fluxE[face] = fe;
                    _fluxY[face] = 0f;
                    continue;
                }

                // If the next vecSize faces contain a wall, fall back to scalar for this block
                if (face + vecSize - 1 < _totalCells)
                {
                    bool hasWall = false;
                    for (int f = face; f < face + vecSize; f++)
                        if (_isWallFace[f]) { hasWall = true; break; }

                    if (!hasWall)
                    {
                        // --- Vectorised block ---
                        var rhoL = new Vector<float>(_rho, face - 1);
                        var rhouL = new Vector<float>(_rhou, face - 1);
                        var EL = new Vector<float>(_E, face - 1);
                        var YL = new Vector<float>(_Y, face - 1);
                        var rhoR = new Vector<float>(_rho, face);
                        var rhouR = new Vector<float>(_rhou, face);
                        var ER = new Vector<float>(_E, face);
                        var YR = new Vector<float>(_Y, face);

                        var invRhoL = Vector<float>.One / Vector.Max(rhoL, new Vector<float>(1e-12f));
                        var invRhoR = Vector<float>.One / Vector.Max(rhoR, new Vector<float>(1e-12f));
                        var uL = rhouL * invRhoL;
                        var uR = rhouR * invRhoR;
                        var kinL = 0.5f * rhouL * uL;
                        var kinR = 0.5f * rhouR * uR;
                        var pL = Gm1 * (EL - kinL);
                        var pR = Gm1 * (ER - kinR);
                        var cL = Vector.SquareRoot(Gamma * pL * invRhoL);
                        var cR = Vector.SquareRoot(Gamma * pR * invRhoR);

                        // Store pressures for visualisation (left cell of each face)
                        pL.CopyTo(_p, face - 1);

                        // Lax‑Friedrichs fluxes
                        var ELs = pL * Gm1Inv * invRhoL + 0.5f * uL * uL; // energy per mass
                        var ERs = pR * Gm1Inv * invRhoR + 0.5f * uR * uR;

                        var FmL = rhoL * uL;
                        var FpL = rhoL * uL * uL + pL;
                        var FeL = (rhoL * ELs + pL) * uL;

                        var FmR = rhoR * uR;
                        var FpR = rhoR * uR * uR + pR;
                        var FeR = (rhoR * ERs + pR) * uR;

                        var absUL = Vector.Abs(uL);
                        var absUR = Vector.Abs(uR);
                        var alpha = Vector.Max(absUL + cL, absUR + cR);

                        var fmVec = 0.5f * (FmL + FmR) - 0.5f * alpha * (rhoR - rhoL);
                        var fpVec = 0.5f * (FpL + FpR) - 0.5f * alpha * (rhouR - rhouL);
                        var feVec = 0.5f * (FeL + FeR) - 0.5f * alpha * (rhoR * ERs - rhoL * ELs);

                        var fyL = FmL * YL;
                        var fyR = FmR * YR;
                        var fyVec = 0.5f * (fyL + fyR) - 0.5f * alpha * (rhoR * YR - rhoL * YL);

                        fmVec.CopyTo(_fluxM, face);
                        fpVec.CopyTo(_fluxP, face);
                        feVec.CopyTo(_fluxE, face);
                        fyVec.CopyTo(_fluxY, face);

                        face += vecSize - 1; // loop increment will add 1, so we advance vecSize faces
                        continue;
                    }
                }

                // --- Scalar interior face (fallback) ---
                {
                    int iLf = face - 1, iRf = face;
                    float rLf = _rho[iLf], rhouLf = _rhou[iLf];
                    float invRhoLf = MathF.ReciprocalEstimate(MathF.Max(rLf, 1e-12f));
                    float uLf = rhouLf * invRhoLf;
                    float pLf = Gm1 * (_E[iLf] - 0.5f * rhouLf * uLf);
                    float cLf = MathF.Sqrt(Gamma * pLf * invRhoLf);
                    float YLf = _Y[iLf];
                    _p[iLf] = pLf;

                    float rRf = _rho[iRf], rhouRf = _rhou[iRf];
                    float invRhoRf = MathF.ReciprocalEstimate(MathF.Max(rRf, 1e-12f));
                    float uRf = rhouRf * invRhoRf;
                    float pRf = Gm1 * (_E[iRf] - 0.5f * rhouRf * uRf);
                    float cRf = MathF.Sqrt(Gamma * pRf * invRhoRf);
                    float YRf = _Y[iRf];

                    LaxFlux(rLf, uLf, pLf, cLf, rRf, uRf, pRf, cRf, out fm, out fp, out fe);
                    _fluxM[face] = fm; _fluxP[face] = fp; _fluxE[face] = fe;

                    float alpha = MathF.Max(MathF.Abs(uLf) + cLf, MathF.Abs(uRf) + cRf);
                    ScalarFlux(rLf, uLf, YLf, rRf, uRf, YRf, alpha, out float fy);
                    _fluxY[face] = fy;
                }
            }

            // If damping/relaxation are active, compute the factors here (re-uses _dampingFactors/_relaxFactors arrays,
            // but we no longer have a separate precompute pass). We compute them on demand in UpdateCells anyway?
            // Actually UpdateCells multiplies by these factors; we can compute them there if needed.
        }

        // ---------- Cell update (unchanged core, but skips relaxation/damping when not needed) ----------
        private void UpdateCells(float dt)
        {
            int vecSize = Vector<float>.Count;
            float dtRelax = -_relaxRate * dt;

            // Compute damping and relaxation factors if needed
            if (_applyDamping)
            {
                for (int i = 0; i < _totalCells; i++)
                {
                    float rho = _rho[i];
                    _dampingFactors[i] = rho > 1e-12f
                        ? MathF.Exp(-_coeffBase * dt / rho)
                        : 1f;
                }
            }
            if (_applyRelax)
            {
                var relaxVal = MathF.Exp(dtRelax);
                for (int i = 0; i < _totalCells; i++)
                    _relaxFactors[i] = relaxVal;
            }

            int iCell = 0;
            for (; iCell <= _totalCells - vecSize; iCell += vecSize)
            {
                var rhoOld = new Vector<float>(_rho, iCell);
                var rhouOld = new Vector<float>(_rhou, iCell);
                var EOld = new Vector<float>(_E, iCell);
                var YOld = new Vector<float>(_Y, iCell);

                var fluxM_L = new Vector<float>(_fluxM, iCell);
                var fluxP_L = new Vector<float>(_fluxP, iCell);
                var fluxE_L = new Vector<float>(_fluxE, iCell);
                var fluxY_L = new Vector<float>(_fluxY, iCell);

                var fluxM_R = new Vector<float>(_fluxM, iCell + 1);
                var fluxP_R = new Vector<float>(_fluxP, iCell + 1);
                var fluxE_R = new Vector<float>(_fluxE, iCell + 1);
                var fluxY_R = new Vector<float>(_fluxY, iCell + 1);

                var dtdx = new Vector<float>(dt) / new Vector<float>(_dx, iCell);

                var rhoNew = rhoOld - dtdx * (fluxM_R - fluxM_L);
                var rhouNew = rhouOld - dtdx * (fluxP_R - fluxP_L);
                var ENew = EOld - dtdx * (fluxE_R - fluxE_L);
                var rhoYOld = rhoOld * YOld;
                var rhoYNew = rhoYOld - dtdx * (fluxY_R - fluxY_L);

                if (_applyDamping)
                    rhouNew *= new Vector<float>(_dampingFactors, iCell);
                if (_applyRelax)
                {
                    var ambRef = new Vector<float>(_ambientEnergyRef);
                    var relax = new Vector<float>(_relaxFactors, iCell);
                    ENew = ambRef + (ENew - ambRef) * relax;
                }

                rhoNew = Vector.Max(rhoNew, new Vector<float>(1e-12f));
                var kinNew = 0.5f * rhouNew * rhouNew / rhoNew;
                var eMin = new Vector<float>(100f * Gm1Inv) + kinNew;
                ENew = Vector.Max(ENew, eMin);

                rhoNew.CopyTo(_rho, iCell);
                rhouNew.CopyTo(_rhou, iCell);
                ENew.CopyTo(_E, iCell);
                var yNew = rhoYNew / rhoNew;
                yNew = Vector.Min(Vector.Max(yNew, Vector<float>.Zero), Vector<float>.One);
                yNew.CopyTo(_Y, iCell);
            }

            // Scalar remainder (only a few cells)
            for (; iCell < _totalCells; iCell++)
            {
                float rhoOld = _rho[iCell], rhouOld = _rhou[iCell], EOld = _E[iCell], YOld = _Y[iCell];
                float fluxM_L = _fluxM[iCell], fluxP_L = _fluxP[iCell], fluxE_L = _fluxE[iCell], fluxY_L = _fluxY[iCell];
                float fluxM_R = _fluxM[iCell + 1], fluxP_R = _fluxP[iCell + 1], fluxE_R = _fluxE[iCell + 1], fluxY_R = _fluxY[iCell + 1];
                float dtdx = dt / _dx[iCell];

                float rhoNew = rhoOld - dtdx * (fluxM_R - fluxM_L);
                float rhouNew = rhouOld - dtdx * (fluxP_R - fluxP_L);
                float ENew = EOld - dtdx * (fluxE_R - fluxE_L);
                float rhoYOld = rhoOld * YOld;
                float rhoYNew = rhoYOld - dtdx * (fluxY_R - fluxY_L);

                if (_applyDamping) rhouNew *= _dampingFactors[iCell];
                if (_applyRelax) ENew = _ambientEnergyRef + (ENew - _ambientEnergyRef) * _relaxFactors[iCell];

                rhoNew = MathF.Max(rhoNew, 1e-12f);
                float kin = 0.5f * rhouNew * rhouNew / rhoNew;
                float eMin = 100f * Gm1Inv + kin;
                ENew = MathF.Max(ENew, eMin);

                _rho[iCell] = rhoNew;
                _rhou[iCell] = rhouNew;
                _E[iCell] = ENew;
                _Y[iCell] = Math.Clamp(rhoYNew / rhoNew, 0f, 1f);
            }
        }

        // ---------- Scalar flux helpers (used in boundaries and scalar fallback) ----------
        private static void LaxFlux(float rL, float uL, float pL, float cL,
                                    float rR, float uR, float pR, float cR,
                                    out float fm, out float fp, out float fe)
        {
            float EL = pL * Gm1Inv / rL + 0.5f * uL * uL;
            float ER = pR * Gm1Inv / rR + 0.5f * uR * uR;
            float FmL = rL * uL;
            float FpL = rL * uL * uL + pL;
            float FeL = (rL * EL + pL) * uL;
            float FmR = rR * uR;
            float FpR = rR * uR * uR + pR;
            float FeR = (rR * ER + pR) * uR;
            float alpha = MathF.Max(MathF.Abs(uL) + cL, MathF.Abs(uR) + cR);
            fm = 0.5f * (FmL + FmR) - 0.5f * alpha * (rR - rL);
            fp = 0.5f * (FpL + FpR) - 0.5f * alpha * (rR * uR - rL * uL);
            fe = 0.5f * (FeL + FeR) - 0.5f * alpha * (rR * ER - rL * EL);
        }

        private static void ScalarFlux(float rL, float uL, float YL,
                                       float rR, float uR, float YR,
                                       float alpha, out float fy)
        {
            float FyL = rL * uL * YL;
            float FyR = rR * uR * YR;
            fy = 0.5f * (FyL + FyR) - 0.5f * alpha * (rR * YR - rL * YL);
        }

        // ---------- Profiling report ----------
        public string GetProfileReport()
        {
            if (!EnableProfiling || _profCallCount == 0)
                return "Pipe profiling disabled or no data.";

            double freq = Stopwatch.Frequency;
            long totalTicks = _profFluxTicks + _profUpdateTicks;
            if (totalTicks == 0) return "No pipe profile data collected.";

            double totalMs = totalTicks * 1000.0 / freq;
            double avgCallUs = totalMs * 1000.0 / _profCallCount;

            double fluxMs = _profFluxTicks * 1000.0 / freq;
            double updateMs = _profUpdateTicks * 1000.0 / freq;

            double fluxAvgUs = fluxMs * 1000.0 / _profCallCount;
            double updateAvgUs = updateMs * 1000.0 / _profCallCount;

            string report = $"  Pipe kernel (over {_profCallCount} calls, total {totalMs:F2} ms, avg {avgCallUs:F2} µs/call):\n";
            report += $"    Fluxes (incl. primitives): {fluxMs:F2} ms  ({_profFluxTicks * 100.0 / totalTicks:F1}%), avg {fluxAvgUs:F2} µs/call\n";
            report += $"    Update cells:              {updateMs:F2} ms  ({_profUpdateTicks * 100.0 / totalTicks:F1}%), avg {updateAvgUs:F2} µs/call\n";

            _profFluxTicks = 0;
            _profUpdateTicks = 0;
            _profCallCount = 0;

            return report;
        }
    }
}