-- | Finite actions of cyclic groups on independently sized carriers.
module Foundational.Representation
  ( FiniteAction (..)
  , validCyclicAction
  , applyAction
  , genFiniteAction
  , representationTolerance
  , FiniteState (..)
  , validFiniteState
  , actOnFiniteState
  , averageFiniteState
  , invariantFiniteState
  , FiniteObservable (..)
  , invariantObservable
  , FiniteMap (..)
  , equivariantMap
  , OutcomeLabel (..)
  , composeOutcomeLabels
  , unitOutcomeLabel
  ) where

import Data.List (sort)
import Test.QuickCheck (Gen, chooseInt)

-- | A generator action of the cyclic group @C_n@. 'actionOrder' is the
-- domain-group order @n@, not necessarily the minimal order of its image;
-- nonfaithful actions are permitted. 'generatorImage' is a permutation of an
-- independently sized carrier.
data FiniteAction = FiniteAction
  { actionOrder :: Int
  , generatorImage :: [Int]
  }
  deriving (Eq, Show)

-- | Validate a full-carrier permutation satisfying @g^n = id@.
validCyclicAction :: FiniteAction -> Bool
validCyclicAction action =
  actionOrder action > 0
    && not (null (generatorImage action))
    && isPermutation action
    && all (hasDeclaredOrder action) [0 .. carrierSize action - 1]

-- | Apply an integer group power to a carrier point, or reject invalid input.
applyAction :: FiniteAction -> Int -> Int -> Maybe Int
applyAction action steps point
  | not (validCyclicAction action) = Nothing
  | point < 0 || point >= carrierSize action = Nothing
  | otherwise = Just (iterateGenerator (generatorImage action) (steps `mod` actionOrder action) point)

isPermutation :: FiniteAction -> Bool
isPermutation action =
  sort (generatorImage action) == [0 .. carrierSize action - 1]

carrierSize :: FiniteAction -> Int
carrierSize = length . generatorImage

hasDeclaredOrder :: FiniteAction -> Int -> Bool
hasDeclaredOrder action point =
  iterateGenerator (generatorImage action) (actionOrder action) point == point

iterateGenerator :: [Int] -> Int -> Int -> Int
iterateGenerator image steps point
  | steps <= 0 = point
  | otherwise = iterateGenerator image (steps - 1) (image !! point)

-- | Generate a valid action whose cyclic order and carrier size may differ.
genFiniteAction :: Gen FiniteAction
genFiniteAction = do
  size <- chooseInt (2, 6)
  cycleSize <- chooseInt (1, size)
  let image = [if point < cycleSize then (point + 1) `mod` cycleSize else point | point <- [0 .. size - 1]]
  pure (FiniteAction cycleSize image)

-- | Shared absolute tolerance for finite representation-state checks.
--
-- This tolerance controls normalization, invariance, and observable equality.
-- It is exported so tests and later finite models use the same boundary.
representationTolerance :: Double
representationTolerance = 1.0e-12

-- | A finite algebraic state represented by weights on the action carrier.
-- The type contains no measurement or probability-evaluation operation.
newtype FiniteState = FiniteState
  { stateWeights :: [Double]
  }
  deriving (Eq, Show)

-- | Check finiteness, nonnegativity, and normalization of finite weights.
validFiniteState :: FiniteState -> Bool
validFiniteState state =
  not (null weights)
    && all finiteNonnegative weights
    && finite total
    && abs (total - 1) <= representationTolerance
  where
    weights = stateWeights state
    total = sum weights

-- | Pull a finite state along a cyclic symmetry action.
actOnFiniteState :: FiniteAction -> Int -> FiniteState -> Maybe FiniteState
actOnFiniteState action steps state
  | not (validCyclicAction action) = Nothing
  | not (validFiniteState state) = Nothing
  | length (stateWeights state) /= carrierSize action = Nothing
  | otherwise = do
      pulledWeights <-
        traverse
          (\point -> do
            source <- applyAction action (-steps) point
            pure (stateWeights state !! source))
          [0 .. carrierSize action - 1]
      let pulledState = FiniteState pulledWeights
      if validFiniteState pulledState then Just pulledState else Nothing

-- | Average a supplied normalized state over all powers of a valid finite
-- cyclic action.
averageFiniteState :: FiniteAction -> FiniteState -> Maybe FiniteState
averageFiniteState action state
  | not (validCyclicAction action) = Nothing
  | not (validFiniteState state) = Nothing
  | length (stateWeights state) /= carrierSize action = Nothing
  | otherwise = do
      orbit <- traverse (\steps -> actOnFiniteState action steps state) [0 .. actionOrder action - 1]
      let divisor = fromIntegral (actionOrder action)
      let averaged =
            FiniteState
              [ sum [stateWeights transformed !! point | transformed <- orbit] / divisor
              | point <- [0 .. carrierSize action - 1]
              ]
      if validFiniteState averaged then Just averaged else Nothing

-- | Check invariance under the generator. For a valid cyclic action, this is
-- equivalent to invariance under every represented group element.
invariantFiniteState :: FiniteAction -> FiniteState -> Bool
invariantFiniteState action state =
  case actOnFiniteState action 1 state of
    Nothing -> False
    Just transformed -> approximateList (stateWeights state) (stateWeights transformed)

-- | Scalar-valued finite observables used only to test fixed-point behavior.
newtype FiniteObservable = FiniteObservable
  { observableValues :: [Double]
  }
  deriving (Eq, Show)

-- | Check that an observable is fixed by the cyclic generator.
invariantObservable :: FiniteAction -> FiniteObservable -> Bool
invariantObservable action observable
  | not (validCyclicAction action) = False
  | length values /= carrierSize action = False
  | not (all finite values) = False
  | otherwise =
      and
        [ case applyAction action 1 point of
            Nothing -> False
            Just transformed ->
              abs (values !! transformed - values !! point)
                <= representationTolerance
        | point <- [0 .. carrierSize action - 1]
        ]
  where
    values = observableValues observable

-- | A total map between finite action carriers.
newtype FiniteMap = FiniteMap
  { finiteMapValues :: [Int]
  }
  deriving (Eq, Show)

-- | Check the generator equivariance equation for a finite carrier map.
equivariantMap :: FiniteAction -> FiniteAction -> FiniteMap -> Bool
equivariantMap source target finiteMap
  | not (validCyclicAction source) = False
  | not (validCyclicAction target) = False
  | length values /= carrierSize source = False
  | any (\point -> point < 0 || point >= carrierSize target) values = False
  | otherwise =
      and
        [ case (applyAction source 1 point, applyAction target 1 (values !! point)) of
            (Just movedSource, Just movedTarget) ->
              values !! movedSource == movedTarget
            _ -> False
        | point <- [0 .. carrierSize source - 1]
        ]
  where
    values = finiteMapValues finiteMap

-- | A phantom-typed, probability-free word of outcome labels.
newtype OutcomeLabel system = OutcomeLabel
  { outcomeAtoms :: [String]
  }
  deriving (Eq, Show)

-- | Compose typed labels by concatenation. The result records the product
-- system type and performs no evaluation into weights or probabilities.
composeOutcomeLabels ::
  OutcomeLabel left ->
  OutcomeLabel right ->
  OutcomeLabel (left, right)
composeOutcomeLabels left right =
  OutcomeLabel (outcomeAtoms left <> outcomeAtoms right)

-- | The empty outcome word for the monoidal unit.
unitOutcomeLabel :: OutcomeLabel ()
unitOutcomeLabel = OutcomeLabel []

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

finiteNonnegative :: Double -> Bool
finiteNonnegative value = finite value && value >= 0

approximateList :: [Double] -> [Double] -> Bool
approximateList left right =
  length left == length right
    && and
      (zipWith
        (\leftValue rightValue ->
          finite leftValue
            && finite rightValue
            && abs (leftValue - rightValue) <= representationTolerance)
        left
        right)