-- | Finite operational probability witnesses with explicit numerical
-- validation. The module does not implement a quantum reconstruction theorem.
module Foundational.Quantum
  ( normalizeProbabilities
  , isProbabilityDistribution
  , probabilityTolerance
  , genProbabilityDistribution
  , Distribution (..)
  , mkDistribution
  , validDistribution
  , LabelEvaluation (..)
  , mkLabelEvaluation
  , validLabelEvaluation
  , sameLabelsDistinctEvaluations
  , StochasticChannel (..)
  , mkStochasticChannel
  , applyChannel
  , identityChannel
  , composeChannels
  , tensorDistribution
  , coarseGrain
  , operationalTolerance
  , genStochasticChannelTriple
  , Effect (..)
  , mkEffect
  , evaluateEffect
  , MeasurePrepareInstrument
  , mkMeasurePrepareInstrument
  , instrumentProbabilities
  , instrumentOutput
  , twoPathInterference
  ) where

import Data.Complex (Complex, imagPart, magnitude, realPart)
import Data.List (transpose)
import Foundational.Representation (OutcomeLabel)
import Test.QuickCheck (Gen, choose, chooseInt, vectorOf)

-- | Normalize nonnegative finite weights with a finite, positive aggregate.
-- Negative zero is accepted as a zero weight.
normalizeProbabilities :: [Double] -> Maybe [Double]
normalizeProbabilities weights
  | null weights = Nothing
  | not (all validWeight weights) = Nothing
  | total <= 0 || not (finite total) = Nothing
  | otherwise = Just (map (/ total) weights)
  where
    total = sum weights

-- | Absolute tolerance used when recognizing an already normalized sum.
probabilityTolerance :: Double
probabilityTolerance = 1.0e-12

-- | Shared tolerance for the finite operational layer.
operationalTolerance :: Double
operationalTolerance = probabilityTolerance

-- | Check finite nonnegative weights whose finite aggregate is within
-- 'probabilityTolerance' of one.
isProbabilityDistribution :: [Double] -> Bool
isProbabilityDistribution probabilities =
  not (null probabilities)
    && all validWeight probabilities
    && finite total
    && abs (total - 1) <= probabilityTolerance
  where
    total = sum probabilities

-- | Generate a normalized finite distribution for property-based examples.
genProbabilityDistribution :: Gen [Double]
genProbabilityDistribution = do
  size <- chooseInt (1, 6)
  weights <- vectorOf size (choose (1.0e-6, 10.0))
  pure (maybe [1] id (normalizeProbabilities weights))

-- | A finite normalized distribution. The public constructor permits
-- adversarial tests; functions validate inputs at every boundary.
newtype Distribution = Distribution
  { distributionWeights :: [Double]
  }
  deriving (Eq, Show)

mkDistribution :: [Double] -> Maybe Distribution
mkDistribution weights
  | isProbabilityDistribution weights = Just (Distribution weights)
  | otherwise = Nothing

validDistribution :: Distribution -> Bool
validDistribution = isProbabilityDistribution . distributionWeights

-- | A normalized numerical evaluation attached to a nonempty list of
-- probability-free labels of one phantom-typed system.
data LabelEvaluation system = LabelEvaluation
  { evaluatedLabels :: [OutcomeLabel system]
  , labelDistribution :: Distribution
  }
  deriving (Eq, Show)

-- | Pair typed labels with a normalized distribution of the same dimension.
mkLabelEvaluation ::
  [OutcomeLabel system] ->
  [Double] ->
  Maybe (LabelEvaluation system)
mkLabelEvaluation labels weights
  | null labels = Nothing
  | length labels /= length weights = Nothing
  | otherwise = LabelEvaluation labels <$> mkDistribution weights

-- | Revalidate a label evaluation constructed through the public record
-- constructor.
validLabelEvaluation :: LabelEvaluation system -> Bool
validLabelEvaluation evaluation =
  not (null labels)
    && length labels == length weights
    && validDistribution distribution
  where
    labels = evaluatedLabels evaluation
    distribution = labelDistribution evaluation
    weights = distributionWeights distribution

-- | Check that the same probability-free labels carry two distinct valid
-- normalized evaluations.
sameLabelsDistinctEvaluations ::
  LabelEvaluation system ->
  LabelEvaluation system ->
  Bool
sameLabelsDistinctEvaluations first second =
  validLabelEvaluation first
    && validLabelEvaluation second
    && evaluatedLabels first == evaluatedLabels second
    && labelDistribution first /= labelDistribution second

-- | A row-stochastic channel. Each input row contains a distribution over
-- outputs.
newtype StochasticChannel = StochasticChannel
  { channelRows :: [[Double]]
  }
  deriving (Eq, Show)

mkStochasticChannel :: [[Double]] -> Maybe StochasticChannel
mkStochasticChannel rows =
  case rows of
    [] -> Nothing
    firstRow : remainingRows
      | null firstRow -> Nothing
      | not (all ((== length firstRow) . length) remainingRows) -> Nothing
      | not (all isProbabilityDistribution rows) -> Nothing
      | otherwise -> Just (StochasticChannel rows)

-- | Apply a row-stochastic channel to a finite distribution.
applyChannel :: StochasticChannel -> Distribution -> Maybe Distribution
applyChannel channel state
  | not (validChannel channel) = Nothing
  | not (validDistribution state) = Nothing
  | length inputs /= length rows = Nothing
  | otherwise =
      mkDistribution
        [ sum (zipWith (*) inputs outputColumn)
        | outputColumn <- transpose rows
        ]
  where
    rows = channelRows channel
    inputs = distributionWeights state

-- | The identity channel on a nonempty carrier.
identityChannel :: Int -> Maybe StochasticChannel
identityChannel size
  | size <= 0 = Nothing
  | otherwise =
      mkStochasticChannel
        [ [if input == output then 1 else 0 | output <- [0 .. size - 1]]
        | input <- [0 .. size - 1]
        ]

-- | Sequential composition. @composeChannels first second@ applies
-- @first@ and then @second@.
composeChannels ::
  StochasticChannel ->
  StochasticChannel ->
  Maybe StochasticChannel
composeChannels first second
  | not (validChannel first && validChannel second) = Nothing
  | outputSize first /= inputSize second = Nothing
  | otherwise =
      mkStochasticChannel
        [ [ sum
              [ firstWeight * secondWeight
              | (firstWeight, secondWeight) <-
                  zip firstRow (map (!! output) (channelRows second))
              ]
          | output <- [0 .. outputSize second - 1]
          ]
        | firstRow <- channelRows first
        ]

-- | Parallel composition of independent finite distributions.
tensorDistribution ::
  Distribution ->
  Distribution ->
  Maybe Distribution
tensorDistribution left right
  | not (validDistribution left && validDistribution right) = Nothing
  | otherwise =
      mkDistribution
        [ leftWeight * rightWeight
        | leftWeight <- distributionWeights left
        , rightWeight <- distributionWeights right
        ]

-- | Push a finite distribution along a total map represented by target
-- indices. The target size is explicit so empty fibres are retained.
coarseGrain ::
  Int ->
  [Int] ->
  Distribution ->
  Maybe Distribution
coarseGrain targetSize assignment state
  | targetSize <= 0 = Nothing
  | not (validDistribution state) = Nothing
  | length assignment /= length weights = Nothing
  | any (\target -> target < 0 || target >= targetSize) assignment = Nothing
  | otherwise =
      mkDistribution
        [ sum
            [ weight
            | (weight, assignedTarget) <- zip weights assignment
            , assignedTarget == target
            ]
        | target <- [0 .. targetSize - 1]
        ]
  where
    weights = distributionWeights state

-- | Generate three compatible binary channels and one binary distribution.
-- This powers an associativity check without hiding dimensions in the test.
genStochasticChannelTriple ::
  Gen
    ( StochasticChannel
    , StochasticChannel
    , StochasticChannel
    , Distribution
    )
genStochasticChannelTriple = do
  first <- binaryChannel
  second <- binaryChannel
  third <- binaryChannel
  stateWeights <- genProbabilityDistribution
  let binaryStateWeights =
        case stateWeights of
          firstWeight : secondWeight : _ ->
            normalizeOrFallback [firstWeight, secondWeight]
          firstWeight : _ ->
            normalizeOrFallback [firstWeight, 1]
          [] -> [0.5, 0.5]
  pure
    ( first
    , second
    , third
    , Distribution binaryStateWeights
    )
  where
    binaryChannel = do
      firstWeight <- choose (1.0e-6, 1 - 1.0e-6)
      secondWeight <- choose (1.0e-6, 1 - 1.0e-6)
      pure
        (StochasticChannel
          [ [firstWeight, 1 - firstWeight]
          , [secondWeight, 1 - secondWeight]
          ])

-- | A bounded finite effect.
newtype Effect = Effect
  { effectWeights :: [Double]
  }
  deriving (Eq, Show)

mkEffect :: [Double] -> Maybe Effect
mkEffect weights
  | null weights = Nothing
  | all validEffectWeight weights = Just (Effect weights)
  | otherwise = Nothing

evaluateEffect :: Effect -> Distribution -> Maybe Double
evaluateEffect effect state
  | not (validDistribution state) = Nothing
  | not (validEffect effect) = Nothing
  | length effects /= length weights = Nothing
  | finite result
      && result >= negate operationalTolerance
      && result <= 1 + operationalTolerance = Just result
  | otherwise = Nothing
  where
    effects = effectWeights effect
    weights = distributionWeights state
    result = sum (zipWith (*) effects weights)

-- | A finite classical measure-and-prepare instrument. Its effect rows fix
-- outcome probabilities; its output states independently fix updates.
data MeasurePrepareInstrument = MeasurePrepareInstrument
  { instrumentEffects :: [Effect]
  , instrumentOutputs :: [Distribution]
  }
  deriving (Eq, Show)

mkMeasurePrepareInstrument ::
  [[Double]] ->
  [Distribution] ->
  Maybe MeasurePrepareInstrument
mkMeasurePrepareInstrument rawEffects outputs = do
  effects <- traverse mkEffect rawEffects
  if validInstrument effects outputs
    then Just (MeasurePrepareInstrument effects outputs)
    else Nothing

instrumentProbabilities ::
  MeasurePrepareInstrument ->
  Distribution ->
  Maybe Distribution
instrumentProbabilities instrument state = do
  values <- traverse (`evaluateEffect` state) (instrumentEffects instrument)
  mkDistribution values

instrumentOutput ::
  Int ->
  MeasurePrepareInstrument ->
  Maybe Distribution
instrumentOutput outcome instrument
  | outcome < 0 || outcome >= length (instrumentOutputs instrument) = Nothing
  | otherwise = Just (instrumentOutputs instrument !! outcome)

-- | Interference for two explicitly supplied complex path amplitudes:
-- |a+b|^2 - |a|^2 - |b|^2.
twoPathInterference ::
  Complex Double ->
  Complex Double ->
  Maybe Double
twoPathInterference first second
  | not (finiteComplex first && finiteComplex second) = Nothing
  | not (finite combined && finite separate && finite result) = Nothing
  | otherwise = Just result
  where
    squaredMagnitude value = magnitude value * magnitude value
    combined = squaredMagnitude (first + second)
    separate = squaredMagnitude first + squaredMagnitude second
    result = combined - separate

validChannel :: StochasticChannel -> Bool
validChannel channel =
  case channelRows channel of
    [] -> False
    firstRow : remainingRows ->
      not (null firstRow)
        && all ((== length firstRow) . length) remainingRows
        && all isProbabilityDistribution (firstRow : remainingRows)

inputSize :: StochasticChannel -> Int
inputSize = length . channelRows

outputSize :: StochasticChannel -> Int
outputSize channel =
  case channelRows channel of
    [] -> 0
    firstRow : _ -> length firstRow

validEffect :: Effect -> Bool
validEffect = maybe False (const True) . mkEffect . effectWeights

validEffectWeight :: Double -> Bool
validEffectWeight value = finite value && value >= 0 && value <= 1

validInstrument :: [Effect] -> [Distribution] -> Bool
validInstrument effects outputs =
  not (null effects)
    && length effects == length outputs
    && all validDistribution outputs
    && all ((== inputDimension) . length . effectWeights) effects
    && inputDimension > 0
    && all
      (\column ->
        finite (sum column)
          && abs (sum column - 1) <= operationalTolerance)
      (transpose (map effectWeights effects))
  where
    inputDimension =
      case effects of
        [] -> 0
        firstEffect : _ -> length (effectWeights firstEffect)

finiteComplex :: Complex Double -> Bool
finiteComplex value = finite (realPart value) && finite (imagPart value)

normalizeOrFallback :: [Double] -> [Double]
normalizeOrFallback weights = maybe [0.5, 0.5] id (normalizeProbabilities weights)

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

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