module CellularFluid.Grid where

import           RIO

import qualified RIO.Vector                            as V
import qualified RIO.Vector.Unsafe                     as V'

import           Data.Vector.Parallel

import qualified CellularFluid.Grid.FD.FDSimpleStepper as FD
import           CellularFluid.Grid.Types

-- | Simulates one cell for one timestep
stepCell :: SimulationCFG -> AdjCells -> Cell -> Cell
stepCell cfg adjs cell = step cell
  where
    phy = cfgPhysics cfg
    step (FluidD p) =
      let width = phyGridSize phy
          dt = cfgTimeStep cfg
          mu = phyFDMu phy
          rho = phyFDRho phy
          nu = mu / rho
          p' = FD.step width dt nu p adjs
       in FluidD p'
    step x = x

-- | Gets cell at (r, c)
--
-- If out of bounds, cell is Wall
cellAt :: Grid -> Int -> Int -> Cell
cellAt g r c
  | r < 0 || c < 0 || r >= h || c >= w = Wall
  | otherwise = cs `idx` (w * r + c) -- We guarentee not out of bounds
        -- idx = parIndex'
  where
    idx = V'.unsafeIndex
    cs = gridCells g
    w = gridWidth g
    h = gridHeight g

-- | Get vector of adjacent cells
adjCells :: Grid -> Vector (Cell, AdjCells)
adjCells g = parIMapV f cs -- (V.imap f cs)
  where
    cs = gridCells g
    w = gridWidth g
    cell = cellAt g
    f i x =
      let (r, c) = i `divMod` w
          acN = cell (r - 1) c
          acNW = cell (r - 1) (c - 1)
          acW = cell r (c - 1)
          acSW = cell (r + 1) (c - 1)
          acS = cell (r + 1) c
          acSE = cell (r + 1) (c + 1)
          acE = cell r (c + 1)
          acNE = cell (r - 1) (c + 1)
          adjEdge = V.fromList [acN, acW, acE, acS]
          adjCorner = V.fromList [acNW, acNE, acSW, acSE]
          adjs = Adj {..}
       in (x, adjs)

-- | Steps grid once
stepGrid :: SimulationCFG -> Grid -> Grid
stepGrid cfg grid = grid {gridCells = cells'}
  where
    cells' = step `parMapV` adjs
    adjs = adjCells grid
    step (cell, adjcells) = stepCell cfg adjcells cell

-- | Checks whether the grid is sane.
--
-- Currently verifies grid size is consistent with (width * height)
gridIsSane :: Grid -> Bool
gridIsSane grid =
  (V.length $ gridCells grid) == (gridHeight grid * gridWidth grid)
