-- | Finite coarse-graining and collision calculations.
--
-- These functions provide executable checks for finite inputs. They do not
-- implement a Boltzmann-Grad limit, propagation of chaos, or an H theorem.
module Foundational.Kinetic
  ( Collision (..)
  , Vec3 (..)
  , JointLaw (..)
  , addVec
  , applyMarkovStep
  , approximatelyEqualKinetic
  , collisionConserved
  , conservationTolerance
  , entropyBits
  , genConservingCollision
  , genHardSphereCase
  , hardSphereCollision
  , kacPairStep
  , kacRotatePair
  , kineticEnergy
  , marginalizeRight
  , validJointLaw
  , vecApproximatelyEqual
  ) where

import Foundational.Quantum (isProbabilityDistribution)
import Data.List (transpose)
import Test.QuickCheck (Gen, choose, chooseInt, elements, vectorOf)

-- | A finite three-dimensional velocity.
data Vec3 = Vec3
  { vecX :: Double
  , vecY :: Double
  , vecZ :: Double
  }
  deriving (Eq, Show)

-- | A finite joint law represented by rows of nonnegative weights.
newtype JointLaw = JointLaw
  { jointWeights :: [[Double]]
  }
  deriving (Eq, Show)

-- | Incoming and outgoing particles represented by scalar momentum and energy.
data Collision = Collision
  { incomingParticles :: [(Double, Double)]
  , outgoingParticles :: [(Double, Double)]
  }
  deriving (Eq, Show)

-- | Relative tolerance used for finite collision and mass checks.
conservationTolerance :: Double
conservationTolerance = 1.0e-12

-- | Validate a nonempty rectangular normalized finite joint law.
validJointLaw :: JointLaw -> Bool
validJointLaw (JointLaw rows) =
  case rows of
    [] -> False
    firstRow : remainingRows ->
      not (null firstRow)
        && all ((== length firstRow) . length) remainingRows
        && all validWeight weights
        && finite total
        && approximatelyEqualKinetic total 1
      where
        weights = concat rows
        total = sum weights

-- | Sum out the right coordinate of a valid finite joint law.
marginalizeRight :: JointLaw -> Maybe [Double]
marginalizeRight law@(JointLaw rows)
  | validJointLaw law =
      let marginal = map sum rows
       in if isProbabilityDistribution marginal then Just marginal else Nothing
  | otherwise = Nothing

-- | Check finite particle data and aggregate scalar conservation laws.
collisionConserved :: Collision -> Bool
collisionConserved collision =
  not (null (incomingParticles collision))
    && not (null (outgoingParticles collision))
    && all validParticle (incomingParticles collision <> outgoingParticles collision)
    && all finite [incomingMomentum, outgoingMomentum, incomingEnergy, outgoingEnergy]
    && approximatelyEqualKinetic incomingMomentum outgoingMomentum
    && approximatelyEqualKinetic incomingEnergy outgoingEnergy
  where
    incomingMomentum = totalMomentum (incomingParticles collision)
    outgoingMomentum = totalMomentum (outgoingParticles collision)
    incomingEnergy = totalEnergy (incomingParticles collision)
    outgoingEnergy = totalEnergy (outgoingParticles collision)

-- | Apply a finite row-stochastic Markov step to a row probability vector.
applyMarkovStep :: [[Double]] -> [Double] -> Maybe [Double]
applyMarkovStep transition probabilities
  | validTransition transition
      && length transition == length probabilities
      && isProbabilityDistribution probabilities =
      let output = map (sum . zipWith (*) probabilities) (transpose transition)
       in if isProbabilityDistribution output then Just output else Nothing
  | otherwise = Nothing

-- | Add two finite vectors componentwise.
addVec :: Vec3 -> Vec3 -> Vec3
addVec left right =
  Vec3
    (vecX left + vecX right)
    (vecY left + vecY right)
    (vecZ left + vecZ right)

subtractVec :: Vec3 -> Vec3 -> Vec3
subtractVec left right =
  Vec3
    (vecX left - vecX right)
    (vecY left - vecY right)
    (vecZ left - vecZ right)

scaleVec :: Double -> Vec3 -> Vec3
scaleVec scalar vector =
  Vec3
    (scalar * vecX vector)
    (scalar * vecY vector)
    (scalar * vecZ vector)

dotVec :: Vec3 -> Vec3 -> Double
dotVec left right =
  vecX left * vecX right
    + vecY left * vecY right
    + vecZ left * vecZ right

-- | Squared speed. The factor one half is omitted because it cancels in
-- invariant comparisons.
kineticEnergy :: Vec3 -> Double
kineticEnergy vector = dotVec vector vector

-- | Apply the elastic hard-sphere collision formula for a finite unit normal.
hardSphereCollision :: Vec3 -> Vec3 -> Vec3 -> Maybe (Vec3, Vec3)
hardSphereCollision velocity partner normal
  | all finite (components velocity <> components partner <> components normal)
      && approximatelyEqualKinetic (kineticEnergy normal) 1 =
      let relativeVelocity = subtractVec velocity partner
          transfer = dotVec relativeVelocity normal
          outgoing = subtractVec velocity (scaleVec transfer normal)
          outgoingPartner = addVec partner (scaleVec transfer normal)
          outputs = components outgoing <> components outgoingPartner
       in if finite transfer && all finite outputs
            then Just (outgoing, outgoingPartner)
            else Nothing
  | otherwise = Nothing

-- | Apply one finite Kac rotation to a scalar velocity pair.
kacRotatePair :: Double -> (Double, Double) -> Maybe (Double, Double)
kacRotatePair angle (first, second)
  | all finite [angle, first, second] =
      let cosine = cos angle
          sine = sin angle
          outgoing = first * cosine + second * sine
          outgoingSecond = -first * sine + second * cosine
       in if all finite [cosine, sine, outgoing, outgoingSecond]
            then Just (outgoing, outgoingSecond)
            else Nothing
  | otherwise = Nothing

-- | Update a selected distinct pair in a finite scalar Kac state.
kacPairStep :: Int -> Int -> Double -> [Double] -> Maybe [Double]
kacPairStep firstIndex secondIndex angle velocities
  | firstIndex < 0
      || secondIndex < 0
      || firstIndex == secondIndex
      || firstIndex >= length velocities
      || secondIndex >= length velocities
      || not (all finite velocities) = Nothing
  | otherwise = do
      (first, second) <-
        kacRotatePair angle
          (velocities !! firstIndex, velocities !! secondIndex)
      pure (replaceAt secondIndex second (replaceAt firstIndex first velocities))

-- | Compute Shannon entropy in bits for an already normalized finite
-- distribution.
entropyBits :: [Double] -> Maybe Double
entropyBits probabilities
  | isProbabilityDistribution probabilities =
      Just (negate (sum (map contribution probabilities)))
  | otherwise = Nothing
  where
    contribution probability
      | probability == 0 = 0
      | otherwise = probability * logBase 2 probability

-- | Generate distinct finite scalar collision states with conserved totals.
genConservingCollision :: Gen Collision
genConservingCollision = do
  size <- chooseInt (2, 5)
  particles <- vectorOf size ((,) <$> choose (-10, 10) <*> choose (0, 20))
  momentumTransfer <- choose (0.25, 1.0)
  energyTransfer <- choose (0, snd (particles !! 1))
  let (firstMomentum, firstEnergy) = particles !! 0
  let (secondMomentum, secondEnergy) = particles !! 1
  let outgoing =
        (firstMomentum + momentumTransfer, firstEnergy + energyTransfer)
          : (secondMomentum - momentumTransfer, secondEnergy - energyTransfer)
          : drop 2 particles
  pure (Collision particles outgoing)

-- | Generate bounded finite velocity pairs and exact coordinate-axis normals.
genHardSphereCase :: Gen (Vec3, Vec3, Vec3)
genHardSphereCase = do
  first <- genVec3
  second <- genVec3
  normal <-
    elements
      [ Vec3 1 0 0
      , Vec3 (-1) 0 0
      , Vec3 0 1 0
      , Vec3 0 (-1) 0
      , Vec3 0 0 1
      , Vec3 0 0 (-1)
      ]
  pure (first, second, normal)

-- | Compare finite scalars with the exported relative tolerance.
approximatelyEqualKinetic :: Double -> Double -> Bool
approximatelyEqualKinetic left right =
  finite left
    && finite right
    && abs (left - right)
      <= conservationTolerance * max 1 (max (abs left) (abs right))

-- | Compare finite vectors componentwise with the exported tolerance.
vecApproximatelyEqual :: Vec3 -> Vec3 -> Bool
vecApproximatelyEqual left right =
  and
    [ approximatelyEqualKinetic (vecX left) (vecX right)
    , approximatelyEqualKinetic (vecY left) (vecY right)
    , approximatelyEqualKinetic (vecZ left) (vecZ right)
    ]

validTransition :: [[Double]] -> Bool
validTransition rows =
  case rows of
    [] -> False
    firstRow : remainingRows ->
      not (null firstRow)
        && length rows == length firstRow
        && all ((== length firstRow) . length) remainingRows
        && all isProbabilityDistribution rows

validWeight :: Double -> Bool
validWeight weight = finite weight && weight >= 0

validParticle :: (Double, Double) -> Bool
validParticle (momentum, energy) =
  finite momentum && finite energy && energy >= 0

totalMomentum :: [(Double, Double)] -> Double
totalMomentum = sum . map fst

totalEnergy :: [(Double, Double)] -> Double
totalEnergy = sum . map snd

components :: Vec3 -> [Double]
components vector = [vecX vector, vecY vector, vecZ vector]

replaceAt :: Int -> value -> [value] -> [value]
replaceAt index replacement values =
  take index values <> [replacement] <> drop (index + 1) values

genVec3 :: Gen Vec3
genVec3 =
  Vec3
    <$> choose (-100, 100)
    <*> choose (-100, 100)
    <*> choose (-100, 100)

finite :: Double -> Bool
finite value = not (isInfinite value || isNaN value)