Taylor.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}

module Math.Taylor where

-- base
import Prelude
  hiding ( Num(..), (/), fromRational )
import Control.Monad.ST
  ( ST, runST )
import Data.Kind
  ( Type, Constraint )
import Data.Ratio
  ( (%) )
import Data.Semigroup
  ( Product(..) )
import GHC.Exts
  ( Double(D#), IsList(..)
  , proxy#, fmaddDouble#
  )
import GHC.TypeNats
  ( Nat, KnownNat
  , natVal'
  , type (+), type (-)
  )

-- primitive
import qualified Data.Primitive.Array as Prim
import qualified Data.Primitive.PrimArray as Prim

-- brush-strokes
import Math.Interval
  ( 𝕀(𝕀), 𝕀ℝ
  , hull, addWithErr, subWithErr, prodWithErr, scaleInterval
  )
import Math.Linear
  ( ℝ, Fin (..), Representable (index) )
import Math.Ring
import Math.Root.Isolation.Core
  ( forEachCoord )

--------------------------------------------------------------------------------

-- | The shape of a tensor of the given order.
type Shape :: Nat -> Type
data Shape order where
  Nil  :: Shape 0
  (:.) :: Nat -> Shape ( n - 1 ) -> Shape n

type KnownShape :: forall {order}. Shape order -> Constraint
class KnownShape shape where
  knownShape :: [ Int ]
instance KnownShape Nil where
  knownShape = []
instance ( KnownNat n, KnownShape sh ) => KnownShape ( n :. sh ) where
  knownShape = fromIntegral ( natVal' @n proxy# ) : knownShape @sh

nbCoefficients :: forall sh. KnownShape sh => Int
nbCoefficients = getProduct $ foldMap ( Product . nbCoeffs ) ( knownShape @sh )
  where
    -- A degree d polynomial has (d+1) coefficients.
    nbCoeffs :: Int -> Int
    nbCoeffs deg = deg + 1

-- | A Taylor model in Bernstein form, consisting of a multivariate polynomial
-- given as a Bernstein tensor + an error term (interval).
--
-- Invariant: the interval always contains @0@ (although it is not necessarily
-- symmetric around @0@).
type BernsteinModel :: Shape order -> Type
data BernsteinModel degs
  = BernsteinModel
  { bernsteinCoefficients   :: !( Prim.PrimArray Double )
     -- Tensor coefficients, stored in row-major order.
  , bernsteinModelRemainder :: !𝕀
  }

instance KnownShape degs => AbelianGroup ( BernsteinModel degs ) where
  BernsteinModel c1 e1 + BernsteinModel c2 e2 =
    let (fpErr, cNew) = zipWithCoeffs addWithErr c1 c2
    in BernsteinModel cNew (e1 + e2 + fpErr)

  BernsteinModel c1 e1 - BernsteinModel c2 e2 =
    let (fpErr, cNew) = zipWithCoeffs subWithErr c1 c2
    in BernsteinModel cNew (e1 + e2 + fpErr)

  negate ( BernsteinModel coeffs err ) =
    -- Floating-point negation is exact, so no change in error.
    BernsteinModel ( Prim.mapPrimArray negate coeffs ) ( negate err )
  fromInteger i =
    case fromInteger i of
      𝕀 lo hi ->
        if lo == hi
        then
          BernsteinModel
            ( Prim.replicatePrimArray ( nbCoefficients @degs ) ( fromIntegral i ) )
            0
        else
          let !x = 0.5 * ( lo + hi )
          in
          BernsteinModel
            ( Prim.replicatePrimArray ( nbCoefficients @degs ) x )
            ( 𝕀 ( lo - x ) ( hi - x ) )

-- | Bernstein coefficient strides (assuming a row-major layout).
getStrides :: [ Int ] -> [ Int ]
getStrides shape =
  let degs = map (+1) shape
  in drop 1 $ scanr (*) 1 degs

choose :: Int -> Int -> Integer
choose n k
    | k < 0 || k > n   = 0
    | k == 0 || k == n = 1
    | k > n `div` 2    = choose n (n - k)
    | otherwise        = go 1 1
  where
    n', k' :: Integer
    n' = fromIntegral n
    k' = fromIntegral k
    go :: Integer -> Integer -> Integer
    go acc i
      | i > k' = acc
      | otherwise = go (acc * (n' - i + 1) `div` i) (i + 1)

-- | Compute the Bernstein weight \( \frac { { n \choose i } { m \choose j } }{ { {n + m} \choose { i + j } } } \).
bernsteinWeight :: Int -> Int -> Int -> Int -> 𝕀
bernsteinWeight n m i j =
  let
    -- Compute the exact rational using integer arithmetic,
    -- then use 'fromRational'. This avoids any accumulation of floating-point
    -- errors, at the cost of using large integers.
    num = choose n i * choose m j
    den = choose (n + m) (i + j)

  in
    -- NB: the 'fromRational' implementation for 𝕀 produces an enclosure of
    -- the exact rational number by 'Double's.
    fromRational $ num % den

-- | Helper function for computing Bernstein weights in the computation of
-- multiplication of two Bernstein models.
bernsteinWeights
  :: forall { order } ( degs1 :: Shape order ) ( degs2 :: Shape order )
  .  ( KnownNat order, KnownShape degs1, KnownShape degs2 )
  => ( Prim.Array ( Prim.PrimArray Double )
     , Prim.Array ( Prim.PrimArray 𝕀      )
     )
bernsteinWeights =
  let
    !dims1 = fromList $ knownShape @degs1
    !dims2 = fromList $ knownShape @degs2
    !order = fromIntegral $ natVal' @order proxy#
  in runST \ @s -> do
    marrBoxedVal <- Prim.newArray @( ST s ) @( Prim.PrimArray Double ) order undefined
    marrBoxedErr <- Prim.newArray @( ST s ) @( Prim.PrimArray 𝕀      ) order undefined

    forLoop 0 order \ d -> do
      let
        n = Prim.indexPrimArray dims1 d
        m = Prim.indexPrimArray dims2 d
        rowSize = m + 1
        totalSize = (n + 1) * rowSize

      marrWeightsVal <- Prim.newPrimArray @( ST s ) @Double totalSize
      marrWeightsErr <- Prim.newPrimArray @( ST s ) @𝕀      totalSize

      forLoop 0 ( n + 1 ) \ i -> do
        let
          !offsetI = i * rowSize
        forLoop 0 ( m + 1 ) \ j -> do
          let
            !( 𝕀 lo hi ) = bernsteinWeight n m i j
            !val = 0.5 * ( lo + hi )
            !err = 𝕀 ( lo - val ) ( hi - val )

          Prim.writePrimArray marrWeightsVal (offsetI + j) val
          Prim.writePrimArray marrWeightsErr (offsetI + j) err

      frozenVal <- Prim.unsafeFreezePrimArray marrWeightsVal
      frozenErr <- Prim.unsafeFreezePrimArray marrWeightsErr

      Prim.writeArray marrBoxedVal d frozenVal
      Prim.writeArray marrBoxedErr d frozenErr

    weightsVal <- Prim.unsafeFreezeArray marrBoxedVal
    weightsErr <- Prim.unsafeFreezeArray marrBoxedErr

    return ( weightsVal, weightsErr )

type ZipSum :: Shape order -> Shape order -> Shape order
type family ZipSum as bs where
  ZipSum Nil Nil = Nil
  ZipSum (a :. as) (b :. bs) = (a + b) :. ZipSum as bs

-- | Multiply two Bernstein models (without any degree reduction).
timesBernsteinModel
  :: forall { order } ( degs1 :: Shape order ) ( degs2 :: Shape order )
  .  ( KnownNat order
     , KnownShape degs1, KnownShape degs2
     , KnownShape ( ZipSum degs1 degs2 )
     )
  => BernsteinModel degs1
  -> BernsteinModel degs2
  -> BernsteinModel ( ZipSum degs1 degs2 )
timesBernsteinModel =
  let
    -- 1. Setup constants
    !env = ConvolutionEnv
      { envOrder      = fromIntegral $ natVal' @order proxy#
      , envDims1      = fromList $ knownShape @degs1
      , envDims2      = fromList $ knownShape @degs2
      , envDimsOut    = fromList $ knownShape @( ZipSum degs1 degs2 )
      , envStrides1   = fromList $ getStrides $ knownShape @degs1
      , envStrides2   = fromList $ getStrides $ knownShape @degs2
      , envStridesOut = fromList $ getStrides $ knownShape @( ZipSum degs1 degs2 )
      , envWeightsVal = fst $ bernsteinWeights @degs1 @degs2
      , envWeightsErr = snd $ bernsteinWeights @degs1 @degs2
      }
    !nbOut = nbCoefficients @( ZipSum degs1 degs2 )

  in
    \ (BernsteinModel coeffs1 err1) (BernsteinModel coeffs2 err2) ->
    runST $ do

    -- 2. Initialise output coefficient array & output FP-error array
    marrVal <- Prim.newPrimArray nbOut
    marrErr <- Prim.newPrimArray nbOut
    Prim.setPrimArray marrVal 0 nbOut 0
    Prim.setPrimArray marrErr 0 nbOut 0

    -- 3. Run the gathering weighted convolution
    runOutputWalker env coeffs1 coeffs2 marrVal marrErr

    -- 4. Compute the final error term and finish.
    resultCoeffs <- Prim.unsafeFreezePrimArray marrVal
    fpErrs       <- Prim.unsafeFreezePrimArray marrErr

    let
      !( 𝕀 min1 max1 ) = minMaxCoeffs coeffs1
      !( 𝕀 min2 max2 ) = minMaxCoeffs coeffs2
      !bound1 = max (abs min1) (abs max1)
      !bound2 = max (abs min2) (abs max2)
      !properErr = scaleInterval bound2 err1 + scaleInterval bound1 err2 + err1 * err2
      !finalErr  = properErr + hulls fpErrs

    return $ BernsteinModel resultCoeffs finalErr
{-# INLINEABLE timesBernsteinModel #-}

data ConvolutionEnv = ConvolutionEnv
  { envOrder      :: !Int
  , envDims1      :: !( Prim.PrimArray Int )
  , envDims2      :: !( Prim.PrimArray Int )
  , envDimsOut    :: !( Prim.PrimArray Int )
  , envStrides1   :: !( Prim.PrimArray Int )
  , envStrides2   :: !( Prim.PrimArray Int )
  , envStridesOut :: !( Prim.PrimArray Int )
  , envWeightsVal :: !( Prim.Array ( Prim.PrimArray Double ) )
  , envWeightsErr :: !( Prim.Array ( Prim.PrimArray 𝕀      ) )
  }

data FiberBuffer s = FiberBuffer
  { fbVal :: !( Prim.MutablePrimArray s Double )
  , fbErr :: !( Prim.MutablePrimArray s 𝕀 )
  }

-- | Lightweight grouping for the current traversal state.
data Cursor = Cursor
  { curOff1 :: !Int   -- ^ Current offset in input 1
  , curOff2 :: !Int   -- ^ Current offset in input 2
  , curWVal :: !Double -- ^ Accumulated Bernstein weight product
  , curWErr :: !𝕀      -- ^ Accumulated weight error
  }

runOutputWalker
  :: forall s
  .  ConvolutionEnv
  -> Prim.PrimArray Double          -- ^ Coeffs 1
  -> Prim.PrimArray Double          -- ^ Coeffs 2
  -> Prim.MutablePrimArray s Double -- ^ Output Coeffs
  -> Prim.MutablePrimArray s 𝕀      -- ^ Output Errors
  -> ST s ()
{-# INLINE runOutputWalker #-}
runOutputWalker env c1 c2 marrVal marrErr =
  if envOrder env == 0
  then do
    -- Trivial scalar codepath
    let !val1 = Prim.indexPrimArray c1 0
        !val2 = Prim.indexPrimArray c2 0
        !(# pVal, pErr #) = weightedProduct val1 val2 1 0
    oldVal <- Prim.readPrimArray marrVal 0
    oldErr <- Prim.readPrimArray marrErr 0
    let !(# newVal, sumErr #) = addWithErr oldVal pVal
        !newErr = oldErr + pErr + sumErr
    Prim.writePrimArray marrVal 0 newVal
    Prim.writePrimArray marrErr 0 newErr
  else do

    let lastDimLen = Prim.indexPrimArray (envDimsOut env) (envOrder env - 1)

    -- Allocate reusable buffers for the fiber dimension
    tmpVal <- Prim.newPrimArray lastDimLen
    tmpErr <- Prim.newPrimArray lastDimLen
    let fiberBuf = FiberBuffer tmpVal tmpErr

    -- Array to track output indices (k) for prefix dimensions
    kIndices <- Prim.newPrimArray (envOrder env)

    -- Iterate over every output fiber
    iterateOutputFibers env kIndices 0 0 $ \outputOffset -> do

      -- A. Reset the temporary accumulation buffer
      Prim.setPrimArray tmpVal 0 lastDimLen 0
      Prim.setPrimArray tmpErr 0 lastDimLen 0

      -- B. Accumulate all input contributions into the buffer
      let startCursor = Cursor 0 0 1 0
      accumulateFiberInputs env c1 c2 kIndices fiberBuf 0 startCursor

      -- C. Write the fiber into the main output array
      commitFiber marrVal marrErr fiberBuf outputOffset

-- | Iterates over all output indices for dimensions 0 to N-2.
-- Calls the 'action' callback when a fiber (dimension N-1) is identified.
iterateOutputFibers
  :: ConvolutionEnv
  -> Prim.MutablePrimArray s Int   -- ^ current output indices
  -> Int                           -- ^ current depth
  -> Int                           -- ^ current output offset
  -> ( Int -> ST s () )            -- ^ action to run on the fiber (receives offset)
  -> ST s ()
{-# INLINE iterateOutputFibers #-}
iterateOutputFibers env kIndices depth offOut action
  | depth == envOrder env - 1 = action offOut
  | otherwise = do
      let
        !outDeg = Prim.indexPrimArray (envDimsOut env) depth
        !stride = Prim.indexPrimArray (envStridesOut env) depth

      -- Loop over k for this dimension
      forLoop 0 ( outDeg + 1 ) \ k -> do
        Prim.writePrimArray kIndices depth k
        iterateOutputFibers env kIndices (depth + 1) (offOut + k * stride) action

-- | Finds all (i, j) pairs such that i + j = k for the prefix dimensions.
-- Updates the Cursor (offsets and weights) and recurses.
accumulateFiberInputs
  :: ConvolutionEnv
  -> Prim.PrimArray Double -- ^ argument 1 coefficients
  -> Prim.PrimArray Double -- ^ argument 2 coefficients
  -> Prim.MutablePrimArray s Int   -- ^ output indices
  -> FiberBuffer s
  -> Int                           -- ^ current depth
  -> Cursor                        -- ^ current accumulated state (offsets/weights)
  -> ST s ()
{-# INLINE accumulateFiberInputs #-}
accumulateFiberInputs env c1 c2 kIndices buf depth cur
  -- Base case: run the 1D convolution.
  | depth == envOrder env - 1 =
      weightedConvolve1D env c1 c2 buf cur

  -- Recursive Step: match i + j = k for this dimension
  | otherwise = do
      targetK <- Prim.readPrimArray kIndices depth

      let
        !n  = Prim.indexPrimArray (envDims1 env) depth
        !m  = Prim.indexPrimArray (envDims2 env) depth
        !s1 = Prim.indexPrimArray (envStrides1 env) depth
        !s2 = Prim.indexPrimArray (envStrides2 env) depth

          -- Bernstein weights for this dimension
        !wArrVal = Prim.indexArray (envWeightsVal env) depth
        !wArrErr = Prim.indexArray (envWeightsErr env) depth
        !wRowLen = m + 1

          -- Valid range for i
        !startI = max 0 (targetK - m)
        !endI   = min n targetK

      forLoop startI (endI + 1) \i -> do
        let
          !j = targetK - i
          !wIdx = (i * wRowLen) + j

          -- Fetch weight for this (i, j) pair
          !dimWVal = Prim.indexPrimArray wArrVal wIdx
          !dimWErr = Prim.indexPrimArray wArrErr wIdx

          -- Update accumulated weight via product: wNew = wOld * wDim
          !(# nextWVal, wProdErr #) = prodWithErr (curWVal cur) dimWVal
          !nextWErr = scaleInterval (curWVal cur) dimWErr
                    + scaleInterval dimWVal (curWErr cur)
                    + (curWErr cur) * dimWErr
                    + wProdErr

          -- Update offsets
          !nextCursor =
            Cursor
              { curOff1 = curOff1 cur + (i * s1)
              , curOff2 = curOff2 cur + (j * s2)
              , curWVal = nextWVal
              , curWErr = nextWErr
              }

        accumulateFiberInputs env c1 c2 kIndices buf (depth + 1) nextCursor

-- | Performs a weighted convolution for the innermost dimension and adds
--  the result into the FiberBuffer.
weightedConvolve1D
  :: ConvolutionEnv
  -> Prim.PrimArray Double
  -> Prim.PrimArray Double
  -> FiberBuffer s
  -> Cursor
  -> ST s ()
{-# INLINE weightedConvolve1D #-}
weightedConvolve1D env coeffs1 coeffs2 (FiberBuffer bufVal bufErr) (Cursor off1 off2 wVal wErr) = do
  let
    !lastDim = envOrder env - 1
    !n  = Prim.indexPrimArray (envDims1 env) lastDim
    !m  = Prim.indexPrimArray (envDims2 env) lastDim
    !s1 = Prim.indexPrimArray (envStrides1 env) lastDim
    !s2 = Prim.indexPrimArray (envStrides2 env) lastDim

    !wArrVal = Prim.indexArray (envWeightsVal env) lastDim
    !wArrErr = Prim.indexArray (envWeightsErr env) lastDim
    !wRowLen = m + 1
    !outDeg  = n + m

  -- Standard 1D convolution
  forLoop 0 (outDeg + 1) \k -> do
    let
      !startI = max 0 (k - m)
      !endI   = min n k

      -- Inner reduction loop: sum( c1[i]*c2[j]*w[i,j] )
      !(fiberVal, fiberErr) =
        foldLoop startI (endI + 1) (0, 0) \i (accVal, accErr) ->
          let
            !j = k - i

            !c1 = Prim.indexPrimArray coeffs1 (off1 + i * s1)
            !c2 = Prim.indexPrimArray coeffs2 (off2 + j * s2)

            !wIdx = i * wRowLen + j
            !dimWVal = Prim.indexPrimArray wArrVal wIdx
            !dimWErr = Prim.indexPrimArray wArrErr wIdx

            -- Compute term: (c1 * c2) * w
            !(# p, pErr #) = prodWithErr c1 c2
            !(# termVal, termProdErr #) = prodWithErr p dimWVal
            !termErr = scaleInterval p dimWErr
                     + scaleInterval dimWVal pErr
                     + pErr * dimWErr
                     + termProdErr

            !(# nextAcc, sumErr #) = addWithErr accVal termVal
            !nextAccErr = accErr + termErr + sumErr
          in (nextAcc, nextAccErr)

    -- Apply the prefix weight (wVal) accumulated from higher dimensions.
    -- result = fiber_sum * prefix_weight
    let !(# finalVal, finalProdErr #) = prodWithErr fiberVal wVal
        !finalErr = scaleInterval fiberVal wErr
                  + scaleInterval wVal fiberErr
                  + wErr * fiberErr
                  + finalProdErr

    -- Accumulate into temporary buffer
    oldBufVal <- Prim.readPrimArray bufVal k
    oldBufErr <- Prim.readPrimArray bufErr k
    let !(# newBufVal, bufSumErr #) = addWithErr oldBufVal finalVal
        !newBufErr = oldBufErr + finalErr + bufSumErr

    Prim.writePrimArray bufVal k newBufVal
    Prim.writePrimArray bufErr k newBufErr

commitFiber
  :: Prim.MutablePrimArray s Double
  -> Prim.MutablePrimArray s 𝕀
  -> FiberBuffer s
  -> Int
  -> ST s ()
{-# INLINE commitFiber #-}
commitFiber marrVal marrErr ( FiberBuffer bufVal bufErr ) outputOffset = do
  len <- Prim.getSizeofMutablePrimArray bufVal
  Prim.copyMutablePrimArray marrVal outputOffset bufVal 0 len
  Prim.copyMutablePrimArray marrErr outputOffset bufErr 0 len

weightedProduct
  :: Double  -- ^ coefficient from tensor 1
  -> Double  -- ^ coefficient from tensor 2
  -> Double  -- ^ accumulated weight
  -> 𝕀       -- ^ accumulated weight error
  -> (# Double, 𝕀 #)
{-# INLINE weightedProduct #-}
weightedProduct c1 c2 wVal wErr =
  let
    -- Compute ( c1 * c2 ) * w, keeping track of FP error.
    !(# p, pErr #)              = prodWithErr c1 c2
    !(# termVal, termProdErr #) = prodWithErr p wVal

    !termErr = scaleInterval p wErr
             + scaleInterval wVal pErr
             + pErr * wErr
             + termProdErr
  in
    (# termVal, termErr #)

fma :: Double -> Double -> Double -> Double
fma ( D# x ) ( D# y ) ( D# z ) = D# ( fmaddDouble# x y z )

boundingBox :: forall order ( degs :: Shape order ). BernsteinModel degs -> 𝕀
boundingBox ( BernsteinModel coeffs err ) =
  err + minMaxCoeffs coeffs

minMaxCoeffs :: Prim.PrimArray Double -> 𝕀
minMaxCoeffs arr = loop 0 $ 𝕀 (1/0) (-1/0)
  where
    len = Prim.sizeofPrimArray arr
    loop !i ival@( 𝕀 mn mx )
      | i >= len = ival
      | otherwise =
          let !x = Prim.indexPrimArray arr i
          in loop (i+1) $ 𝕀 (min mn x) (max mx x)

hulls :: Prim.PrimArray 𝕀 -> 𝕀
hulls arr = loop 0 $ 𝕀 (1/0) (-1/0)
  where
    len = Prim.sizeofPrimArray arr
    loop !i ival
      | i >= len = ival
      | otherwise =
          let !x = Prim.indexPrimArray arr i
          in loop (i+1) $ hull x ival


refineBernsteinModel
  :: forall { order } ( degs :: Shape order )
  .  ( KnownNat order
     , KnownShape degs
     , Representable 𝕀 ( 𝕀ℝ order )
     )
  => 𝕀ℝ order -- ^ The domains @[u_i, v_i]@ for each dimension
  -> BernsteinModel degs
  -> BernsteinModel degs
refineBernsteinModel =
  let
    !dims    = knownShape @degs
    !maxDeg  = if null dims then 0 else maximum dims
    !strides = getStrides dims
  in \ domains ( BernsteinModel coeffs inputErr ) -> runST \ @s -> do

    -- Allocate a scratch buffer for 'restrict1D'
    scratchBuf <- Prim.newPrimArray @( ST s ) @Double ( maxDeg + 1 )

    -- 1. Copy coefficients to a mutable array (we modify in place)
    let len = Prim.sizeofPrimArray coeffs
    marr <- Prim.newPrimArray len
    Prim.copyPrimArray marr 0 coeffs 0 len


    let
      -- Iterate over all 1D lines for a specific dimension 'd'
      -- This loops over all indices *except* dimension 'd'.
      applyToDim :: Fin order -> 𝕀 -> 𝕀 -> ST s 𝕀
      applyToDim ( Fin d_one_ixed ) ( 𝕀 u v ) currentAccErr = do
        let
          -- Fin uses 1-based indexing (not for any good reason), so account for that
          !d      = fromIntegral $ d_one_ixed - 1
          !n      = dims !! d
          !stride = strides !! d
          !dimSize = n + 1
          !block   = stride * dimSize

          loopBlock !base !acc
            | base >= len = return acc
            | otherwise = do
              let loopOffset !off !innerAcc
                    | off >= stride = return innerAcc
                    | otherwise = do
                        let !idx = base + off
                        e <- restrict1D scratchBuf ( Strided marr idx stride ) n (𝕀 u v)
                        loopOffset (off + 1) (innerAcc + e)
              errChunk <- loopOffset 0 0
              loopBlock (base + block) (acc + errChunk)

        loopBlock 0 currentAccErr

    -- 2. Apply refinement dimension by dimension
    totalRefinementErr <-
      forEachCoord @order 0 \ d accErr -> do
        let !dom = domains `index` d
        applyToDim d dom accErr

    finalCoeffs <- Prim.unsafeFreezePrimArray marr

    -- 3. Update error
    -- The accumulated floating point error from De Casteljau must be added.
    return $ BernsteinModel finalCoeffs (inputErr + totalRefinementErr)

-- | Restricts a 1D Bernstein polynomial from domain [0,1] to [u,v].
--
-- Returns the floating-point error incurred during the transformation.
restrict1D
  :: forall s
  .  Prim.MutablePrimArray s Double -- ^ scratch buffer (size >= degree + 1)
  -> Strided s -- coefficient buffer
  -> Int -- ^ degree
  -> 𝕀
  -> ST s 𝕀
restrict1D buf coeffs deg ( 𝕀 u v ) = do

  -- Use the passed-in scratch buffer for the computaiton.
  let !count = deg + 1
  forLoop 0 count \ k -> do
    val <- readStrided coeffs k
    Prim.writePrimArray buf k val

  -- Do right Casteljau on our local buffer, [0,1] --> [u,0].
  err1 <-
    if u == 0
    then return 0
    else deCasteljauRight buf deg u

  -- Do left Casteljau on our local buffer, [u,0] --> [u,v].
  -- As we have already re-scaled once, the new right edge is not v but ( v - u ) / ( 1 - u ).
  err2 <-
    if v == 1
    then return 0
    else do
      let v' = ( v - u ) / ( 1 - u )
      deCasteljauLeft buf deg v'

  -- Propagate the changes from the temporary buffer to the parent array
  forLoop 0 count \k -> do
    val <- Prim.readPrimArray buf k
    writeStrided coeffs k val

  return $ err1 + err2

-- | Restrict a Bernstein polynomial on [0, 1] to [0, t].
--
-- Returns accumulated floating-point error.
deCasteljauLeft
  :: Prim.MutablePrimArray s Double
  -> Int    -- ^ degree
  -> Double -- ^ split point @t@
  -> ST s 𝕀
deCasteljauLeft marr deg t = do
  let
    tInv = 1 - t

    loopJ j errAcc
      | j > deg   = return errAcc
      | otherwise = do

        -- NB: backwards iteration.
        let
          loopI i currentErr
            | i < j     = return currentErr
            | otherwise = do
              let
              val     <- Prim.readPrimArray marr i
              valPrev <- Prim.readPrimArray marr ( i - 1 )

              let
                !(# term1 , e1 #) = prodWithErr valPrev tInv
                !(# term2 , e2 #) = prodWithErr val t
                !(# newVal, e3 #) = addWithErr term1 term2

                stepErr = e1 + e2 + e3

              Prim.writePrimArray marr i newVal
              loopI (i - 1) (currentErr + stepErr)

        newErr <- loopI deg 0
        loopJ (j + 1) (errAcc + newErr)

  loopJ 1 0

-- | Restrict a Bernstein polynomial on [0, 1] to [t, 1].
--
-- Returns accumulated floating-point error.
deCasteljauRight
  :: Prim.MutablePrimArray s Double
  -> Int    -- ^ degree
  -> Double -- ^ split point @t@
  -> ST s 𝕀
deCasteljauRight marr deg t = do
  let
    tInv = 1.0 - t

    loopJ j errAcc
      | j > deg   = return errAcc
      | otherwise = do

        -- NB: forwards iteration.
        let
          limit = deg - j
          loopI i currentErr
            | i > limit = return currentErr
            | otherwise = do
              val     <- Prim.readPrimArray marr i
              valNext <- Prim.readPrimArray marr ( i + 1 )

              let
                !(# term1 , e1 #) = prodWithErr val tInv
                !(# term2 , e2 #) = prodWithErr valNext t
                !(# newVal, e3 #) = addWithErr term1 term2

                stepErr = e1 + e2 + e3

              Prim.writePrimArray marr i newVal
              loopI (i + 1) (currentErr + stepErr)

        newErr <- loopI 0 0
        loopJ (j + 1) (errAcc + newErr)

  loopJ 1 0

evaluateBernsteinModel
  :: forall { order } ( degs :: Shape order )
  .  ( KnownNat order
     , KnownShape degs
     , Representable Double ( ℝ order )
     )
  => BernsteinModel degs -> ℝ order -> 𝕀
evaluateBernsteinModel ( BernsteinModel coeffs startingErr ) x =
  runST \ @s -> do

    let
      !dims = knownShape @degs
      !dimsArray = fromList dims
      !maxDeg    = if null dims then 0 else maximum dims
      !totalLen  = Prim.sizeofPrimArray coeffs

    -- Copy input into our a local buffer.
    bufA <- Prim.newPrimArray @( ST s ) @Double totalLen
    Prim.copyPrimArray bufA 0 coeffs 0 totalLen

    -- A second buffer for ping/pong.
    bufB <- Prim.newPrimArray @( ST s ) @Double  totalLen

    -- Small scratch buffer for 1D evaluations.
    scratchBuf <- Prim.newPrimArray @( ST s ) @Double $ maxDeg + 1

    let
      -- We return the error accumulated in this step AND the buffer containing the result
      processDim
        :: Prim.MutablePrimArray s Double -- src
        -> Prim.MutablePrimArray s Double -- dst
        -> Int                            -- current (logical) size of src
        -> ( Int, Double )                -- @(deg, t)@
        -> ST s ( Int, 𝕀 )                -- @(new_size, err)@
      processDim src dst srcSize (deg, t) = do
        let
          !rowSize   = deg + 1
          !outputSize = srcSize `div` rowSize

        stepErr <- foldMLoop 0 outputSize 0 \ i accErr -> do
          let offset = i * rowSize
          -- eval1D reads from 'src' and uses 'scratchBuf' internally
          (val, evalErr) <- eval1D scratchBuf src offset deg t
          Prim.writePrimArray dst i val
          return $! accErr + evalErr
        return (outputSize, stepErr)

    -- Ping-pong
    ( finalBuf, _, _, totalRefinementErr)  <-
      forEachCoord @order ( bufA, bufB, totalLen, 0 ) \ ( Fin i ) ( curr, other, sz, err ) -> do
        let
          order = Prim.sizeofPrimArray dimsArray

          -- Process the last dimension first, because it has stride 1.
          --
          -- Apologies for the index math, caused by 'Fin' being 1-indexed
          -- while all the array indices are 0-based.
          !j     = order - fromIntegral i -- 0-based index from end
          !deg_j = dimsArray `Prim.indexPrimArray` j
          !x_j   = x `index` ( Fin $ fromIntegral j + 1 )

        ( newSize, stepErr ) <- processDim curr other sz ( deg_j, x_j )
        -- Swap buffers
        return ( other, curr, newSize, err + stepErr )

    finalVal <- Prim.readPrimArray finalBuf 0
    return $ 𝕀 finalVal finalVal + totalRefinementErr + startingErr


-- | Evaluates a 1D Bernstein polynomial at 't' using de Casteljau.
eval1D
  :: Prim.MutablePrimArray s Double -- ^ scratch buffer (size >= degree + 1)
  -> Prim.MutablePrimArray s Double -- ^ coefficients (only read, not written)
  -> Int
  -> Int
  -> Double -- @t@
  -> ST s ( Double, 𝕀 )
eval1D buf marr off deg t = do

  -- Copy the coefficients into the scratch buffer.
  let
    !count = deg + 1
  forLoop 0 count \ k -> do
    v <- Prim.readPrimArray marr ( off + k )
    Prim.writePrimArray buf k v

  let
    !tInv = 1 - t

  -- Triangle loop: Reduce degree from 'deg' down to 1
  finalErr <-
    foldMLoop 1 ( deg + 1 ) 0 \ r accErr -> do
      -- Inner loop: compute degree r coefficients
      -- We write to indices 0 .. (deg - r)
      let !limit = deg - r
      stepErr <- foldMLoop 0 ( limit + 1 ) 0 \ i currentErr -> do
        c0 <- Prim.readPrimArray buf i
        c1 <- Prim.readPrimArray buf $ i + 1

        let
          !(# p1 , e1 #) = prodWithErr c0 tInv
          !(# p2 , e2 #) = prodWithErr c1 t
          !(# val, e3 #) = addWithErr p1 p2

        Prim.writePrimArray buf i val
        return $ currentErr + e1 + e2 + e3

      return ( accErr + stepErr )

  finalVal <- Prim.readPrimArray buf 0
  return (finalVal, finalErr)


-- TODO: degree lowering, degree raising


{- TODO: Bézier clipping

Start with multivariate Taylor model in Bernstein basis

  1. Choose a dimension/variable.
  2. Perform column reduction to that dimension, taking max/min of all coefficients across other dimensions.
     This gives a 1D Taylor model in Bernstein basis.
  3. Use Andrew's monotone chain algorithm to compute a floor and roof,
     including the Taylor error term (negative part for floor, positive part for roof).
  4. Clip the floor and roof against y=0 (using interval arithmetic to account for rounding errors).
     For example, for the ceiling:
       - start with the first edge that starts above y=0 or crosses y=0,
       - end with the last edge that ends above y=0 or crosses y=0.
      This (potentially) refines the original interval in that dimension.

Then iterate across all dimensions. It may be beneficial to refine the Bézier patch
after handling each dimension, as long as we shrink enough in the dimension to
warrant doing this computation (e.g. if we shrink to 99% of the orginal size in the first
dimension, probably not worth doing de Casteljau).

-}

{- FUTURE ENHANCEMENTS

  - composition of Bernstein tensors/Bernstein models

-}

--------------------------------------------------------------------------------
-- Helpers

data Strided s = Strided
  { stridedArray  :: !( Prim.MutablePrimArray s Double )
  , stridedOffset :: !Int
  , stridedStride :: !Int
  }

readStrided :: Strided s -> Int -> ST s Double
readStrided ( Strided arr off str) i =
  Prim.readPrimArray arr ( off + i * str )

writeStrided :: Strided s -> Int -> Double -> ST s ()
writeStrided ( Strided arr off str ) i x =
  Prim.writePrimArray arr ( off + i * str ) x

{-# INLINEABLE forLoop #-}
{-# SPECIALISE INLINE forall s. forall. forLoop @( ST s ) #-}
-- | @C@-style for loop, from start to end (but not including the end).
forLoop :: Monad m => Int -> Int -> (Int -> m ()) -> m ()
forLoop start end body = go start
  where
    go !i | i >= end  = return ()
          | otherwise = body i >> go (i + 1)

{-# INLINE foldLoop #-}
foldLoop :: Int -> Int -> a -> ( Int -> a -> a ) -> a
foldLoop start end initVal body = go start initVal
  where
    go !i !acc
      | i >= end  = acc
      | otherwise = let acc' = body i acc in go (i + 1) acc'

{-# INLINEABLE foldMLoop #-}
{-# SPECIALISE INLINE forall s. forall. foldMLoop @( ST s ) #-}
foldMLoop :: Monad m => Int -> Int -> a -> ( Int -> a -> m a ) -> m a
foldMLoop start end initVal body = go start initVal
  where
    go !i !acc | i >= end  = return acc
               | otherwise = body i acc >>= go (i + 1)


zipWithCoeffs
  :: ( Double -> Double -> (# Double, 𝕀 #) )
  -> Prim.PrimArray Double
  -> Prim.PrimArray Double
  -> ( 𝕀, Prim.PrimArray Double )
zipWithCoeffs op = \ c1 c2 -> runST $ do
  let !len = Prim.sizeofPrimArray c1
  marr <- Prim.newPrimArray len

  -- Compute the hull of local floating-point errors
  totalFpErr <- foldMLoop 0 len ( 𝕀 (1/0) (-1/0) ) \ i accErr -> do
    let
      !v1 = Prim.indexPrimArray c1 i
      !v2 = Prim.indexPrimArray c2 i
      !(# res, err #) = op v1 v2
    Prim.writePrimArray marr i res
    return $! hull accErr err

  resArr <- Prim.unsafeFreezePrimArray marr
  return ( totalFpErr, resArr )
{-# INLINE zipWithCoeffs #-}
  -- inline to expose the 'op' function