Zigzag.hs

-- | Exact alpha-complex persistence across non-nested activation depths.
--
-- Each observed depth is rebuilt independently from stable labelled points.
-- Consecutive complexes are compared through the canonical union cospan
-- @K_i -> K_i union K_(i+1) <- K_(i+1)@.  Stable simplex labels, never DCEL
-- handles, own cross-mesh identity.  Homology remains the sole owner of the
-- zigzag interval decomposition.
module Moonlight.Triangulation.Zigzag
  ( ActivationPoint (..),
    ActivationSlice (..),
    ActivationCell (..),
    ActivationStage (..),
    ActivationZigzagError (..),
    ActivationComplex,
    activationComplexStage,
    activationComplexCells,
    activationComplexBasisAt,
    activationComplexChainComplex,
    AdjacentUnionWitness,
    adjacentUnionStage,
    adjacentCommonCells,
    adjacentUnionComplex,
    adjacentLeftInclusion,
    adjacentRightInclusion,
    ActivationZigzag,
    activationZigzagStages,
    activationZigzagComplexes,
    activationZigzagCorrespondences,
    activationZigzagDiagram,
    ZigzagInterval (..),
    ActivationInterval,
    activationAlphaZigzag,
    activationZigzagIntervals,
    activationAlphaPersistence,
  )
where

import Data.Bifunctor (first)
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NonEmpty
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import Moonlight.Core (firstDuplicate)
import Moonlight.Homology.Boundary
  ( BoundaryIncidence,
    FiniteChainComplex,
    emptyBoundaryIncidence,
    emptyBoundaryIncidenceOf,
    materializeIncidenceBoundary,
    mkFiniteChainComplexChecked,
    overlapBoundaryIncidence,
  )
import Moonlight.Homology.Chain
  ( HomologicalDegree (..),
    HomologyFailure,
  )
import Moonlight.Homology.Persistence
  ( FiniteChainMap,
    FiniteChainZigzag,
    ZigzagArrow (..),
    ZigzagFailure,
    ZigzagInterval (..),
    mkFiniteChainMapChecked,
    mkFiniteChainZigzag,
    rationalZigzagIntervals,
  )
import Moonlight.Homology.Topology qualified as CellComplex
import Moonlight.Triangulation.Alpha
  ( AlphaFiltrationError,
    alphaComplexAtRadius,
    alphaFiltration,
  )
import Moonlight.Triangulation.BulkLoad
  ( DuplicatePayloadPolicy (KeepFirstPayload),
    delaunayFromCoordinates,
  )
import Moonlight.Triangulation.CellSet (CellSelectionError)
import Moonlight.Triangulation.CellComplex
  ( DCELComplex,
    fromExactCellSet,
  )
import Moonlight.Triangulation.Dcel qualified as Dcel
import Moonlight.Triangulation.Handles.HandleDefs
  ( FaceId,
    UndirectedEdgeId,
  )
import Moonlight.Triangulation.Types
  ( BuildError,
    BuildResult (..),
    BuildStats (..),
    DelaunayTriangulation,
    Point,
    RadiusSquared,
    unitElementDefaults,
  )

-- | One stable observation identity and its two-dimensional embedding.  The
-- coordinates may change at every depth; the label is the correspondence.
data ActivationPoint label = ActivationPoint
  { activationPointLabel :: !label,
    activationPointCoordinates :: !Point
  }
  deriving stock (Eq, Show)

-- | One independently sampled depth and the exact squared-radius alpha
-- threshold to observe there.
data ActivationSlice depth label = ActivationSlice
  { activationSliceDepth :: !depth,
    activationSliceRadiusSquared :: !RadiusSquared,
    activationSlicePoints :: !(Vector (ActivationPoint label))
  }
  deriving stock (Eq, Show)

-- | Canonical abstract simplex identity induced by stable vertex labels.
-- Edge and face labels are stored in ascending order.
data ActivationCell label
  = ActivationVertex !label
  | ActivationEdge !label !label
  | ActivationFace !label !label !label
  deriving stock (Eq, Ord, Show)

-- | Vertices of the expanded zigzag alternate between observed depths and
-- their adjacent common ambient unions.
data ActivationStage depth
  = ObservedDepth !depth
  | AdjacentDepthUnion !depth !depth
  deriving stock (Eq, Ord, Show)

data ActivationZigzagError depth label
  = ActivationDuplicateDepth !depth
  | ActivationDuplicateLabel !depth !label
  | ActivationCoordinateCollapse !depth !Int !Int
  | ActivationBuildFailed !depth !BuildError
  | ActivationFiltrationFailed !depth !AlphaFiltrationError
  | ActivationSelectionFailed !depth !CellSelectionError
  | ActivationInnerFaceInvalid !(ActivationStage depth) !FaceId
  | ActivationChainComplexFailed !(ActivationStage depth) !HomologyFailure
  | ActivationZigzagFailed !ZigzagFailure
  | ActivationStageIndexMissing !Int
  deriving stock (Eq, Show)

data ActivationBasis label = ActivationBasis
  { activationVertexBasis :: ![ActivationCell label],
    activationEdgeBasis :: ![ActivationCell label],
    activationFaceBasis :: ![ActivationCell label]
  }

-- | A checked, canonically labelled simplicial interpretation of one observed
-- alpha complex or adjacent union.
data ActivationComplex depth label = ActivationComplex
  { activationComplexStage :: !(ActivationStage depth),
    activationComplexCells :: !(Set (ActivationCell label)),
    activationComplexBasis :: !(ActivationBasis label),
    activationComplexChainComplex :: !(FiniteChainComplex Int)
  }

-- | The canonical labelled basis corresponding to the numeric basis of
-- 'activationComplexChainComplex' at one degree.
activationComplexBasisAt ::
  ActivationComplex depth label ->
  HomologicalDegree ->
  [ActivationCell label]
activationComplexBasisAt complexValue degreeValue =
  activationBasisAt degreeValue (activationComplexBasis complexValue)

-- | The two checked inclusion legs proving how adjacent observations meet in
-- their union.  The common-cell set is a derived inspection surface.
data AdjacentUnionWitness depth label = AdjacentUnionWitness
  { adjacentCommonCells :: !(Set (ActivationCell label)),
    adjacentUnionComplex :: !(ActivationComplex depth label),
    adjacentLeftInclusion :: !(FiniteChainMap Int),
    adjacentRightInclusion :: !(FiniteChainMap Int)
  }

adjacentUnionStage :: AdjacentUnionWitness depth label -> ActivationStage depth
adjacentUnionStage = activationComplexStage . adjacentUnionComplex

-- | The expanded stages, admitted observations, correspondence witnesses, and
-- authoritative Homology diagram travel together.
data ActivationZigzag depth label = ActivationZigzag
  { activationZigzagComplexes :: !(NonEmpty (ActivationComplex depth label)),
    activationZigzagCorrespondences :: !(Vector (AdjacentUnionWitness depth label)),
    activationZigzagDiagram :: !(FiniteChainZigzag Int)
  }

-- | The labelled view of the Homology diagram.  It is derived from the
-- admitted observations and their union witnesses rather than stored as a
-- second sequence that could disagree with them.
activationZigzagStages :: ActivationZigzag depth label -> Vector (ActivationStage depth)
activationZigzagStages zigzag =
  let firstComplex :| remainingComplexes = activationZigzagComplexes zigzag
   in Vector.fromList
        ( activationComplexStage firstComplex
            : concat
              ( zipWith
                  (\witness rightComplex ->
                      [adjacentUnionStage witness, activationComplexStage rightComplex]
                  )
                  (Vector.toList (activationZigzagCorrespondences zigzag))
                  remainingComplexes
              )
        )

-- | The Homology interval carrier with endpoints resolved into the activation
-- stage vocabulary.  This is a specialization, not a parallel record.
type ActivationInterval depth = ZigzagInterval (ActivationStage depth)

activationAlphaZigzag ::
  (Ord depth, Ord label) =>
  NonEmpty (ActivationSlice depth label) ->
  Either (ActivationZigzagError depth label) (ActivationZigzag depth label)
activationAlphaZigzag slices = do
  case firstDuplicate (fmap activationSliceDepth (NonEmpty.toList slices)) of
    Just duplicateDepth -> Left (ActivationDuplicateDepth duplicateDepth)
    Nothing -> Right ()
  observedComplexes <- traverse activationComplexFromSlice slices
  buildActivationZigzag
    (NonEmpty.zip (fmap activationSliceDepth slices) observedComplexes)

activationZigzagIntervals ::
  ActivationZigzag depth label ->
  Either (ActivationZigzagError depth label) [ActivationInterval depth]
activationZigzagIntervals zigzag = do
  intervals <-
    first ActivationZigzagFailed
      (rationalZigzagIntervals (activationZigzagDiagram zigzag))
  traverse (traverse (requireStage (activationZigzagStages zigzag))) intervals

activationAlphaPersistence ::
  (Ord depth, Ord label) =>
  NonEmpty (ActivationSlice depth label) ->
  Either (ActivationZigzagError depth label) [ActivationInterval depth]
activationAlphaPersistence slices =
  activationAlphaZigzag slices >>= activationZigzagIntervals

activationComplexFromSlice ::
  Ord label =>
  ActivationSlice depth label ->
  Either (ActivationZigzagError depth label) (ActivationComplex depth label)
activationComplexFromSlice slice = do
  let depthValue = activationSliceDepth slice
      points = activationSlicePoints slice
      labels = fmap activationPointLabel points
  case firstDuplicate (Vector.toList labels) of
    Just duplicateLabel -> Left (ActivationDuplicateLabel depthValue duplicateLabel)
    Nothing -> Right ()
  buildResult <-
    first (ActivationBuildFailed depthValue)
      ( delaunayFromCoordinates
          unitElementDefaults
          (fmap activationPointCoordinates points)
          labels
          KeepFirstPayload
      )
  let inputCount = Vector.length points
      uniqueCount = statUniquePoints (buildStats buildResult)
  if uniqueCount /= inputCount
    then Left (ActivationCoordinateCollapse depthValue inputCount uniqueCount)
    else do
      let triangulation = buildTriangulation buildResult
      filtration <-
        first (ActivationFiltrationFailed depthValue) (alphaFiltration triangulation)
      selectedCells <-
        first (ActivationSelectionFailed depthValue)
          (alphaComplexAtRadius (activationSliceRadiusSquared slice) filtration)
      let stage = ObservedDepth depthValue
      cellSet <- activationCellsFromDcel stage triangulation (fromExactCellSet selectedCells)
      compileActivationComplex stage cellSet

activationCellsFromDcel ::
  Ord label =>
  ActivationStage depth ->
  DelaunayTriangulation label ->
  DCELComplex ->
  Either (ActivationZigzagError depth label) (Set (ActivationCell label))
activationCellsFromDcel stage triangulation complexValue = do
  let vertexCells =
        fmap
          (ActivationVertex . Dcel.vertexData triangulation)
          (CellComplex.vertices complexValue)
      edgeCells =
        fmap
          (activationEdgeFromDcel triangulation complexValue)
          (CellComplex.edges complexValue)
  faceCells <-
    traverse
      (activationFaceFromDcel stage triangulation)
      (CellComplex.faces complexValue)
  pure (Set.fromList (vertexCells <> edgeCells <> faceCells))

activationEdgeFromDcel ::
  Ord label =>
  DelaunayTriangulation label ->
  DCELComplex ->
  UndirectedEdgeId ->
  ActivationCell label
activationEdgeFromDcel triangulation complexValue edgeValue =
  let (firstVertex, secondVertex) = CellComplex.edgeBoundary complexValue edgeValue
   in orderedActivationEdge
        (Dcel.vertexData triangulation firstVertex)
        (Dcel.vertexData triangulation secondVertex)

activationFaceFromDcel ::
  Ord label =>
  ActivationStage depth ->
  DelaunayTriangulation label ->
  FaceId ->
  Either (ActivationZigzagError depth label) (ActivationCell label)
activationFaceFromDcel stage triangulation faceValue =
  case Dcel.innerFaceVertices triangulation faceValue of
    Nothing -> Left (ActivationInnerFaceInvalid stage faceValue)
    Just (firstVertex, secondVertex, thirdVertex) ->
      Right
        ( orderedActivationFace
            (Dcel.vertexData triangulation firstVertex)
            (Dcel.vertexData triangulation secondVertex)
            (Dcel.vertexData triangulation thirdVertex)
        )

orderedActivationEdge :: Ord label => label -> label -> ActivationCell label
orderedActivationEdge firstLabel secondLabel =
  uncurry ActivationEdge (ascendingPair firstLabel secondLabel)

orderedActivationFace :: Ord label => label -> label -> label -> ActivationCell label
orderedActivationFace firstLabel secondLabel thirdLabel =
  let (firstLow, firstHigh) = ascendingPair firstLabel secondLabel
      (secondLow, finalHigh) = ascendingPair firstHigh thirdLabel
      (finalLow, finalMiddle) = ascendingPair firstLow secondLow
   in ActivationFace finalLow finalMiddle finalHigh

ascendingPair :: Ord value => value -> value -> (value, value)
ascendingPair firstValue secondValue =
  if firstValue <= secondValue
    then (firstValue, secondValue)
    else (secondValue, firstValue)

compileActivationComplex ::
  Ord label =>
  ActivationStage depth ->
  Set (ActivationCell label) ->
  Either (ActivationZigzagError depth label) (ActivationComplex depth label)
compileActivationComplex stage cells = do
  let basis = activationBasisFromCells cells
  degreeOneBoundary <-
    first (ActivationChainComplexFailed stage)
      ( materializeIncidenceBoundary
          activationCellBoundary
          (activationEdgeBasis basis)
          (activationVertexBasis basis)
      )
  degreeTwoBoundary <-
    first (ActivationChainComplexFailed stage)
      ( materializeIncidenceBoundary
          activationCellBoundary
          (activationFaceBasis basis)
          (activationEdgeBasis basis)
      )
  let degreeZeroBoundary :: BoundaryIncidence Int
      degreeZeroBoundary =
        emptyBoundaryIncidenceOf
          (fromIntegral (length (activationVertexBasis basis)))
          0
      boundaryAt :: HomologicalDegree -> BoundaryIncidence Int
      boundaryAt (HomologicalDegree degreeIndex) =
        case degreeIndex of
          0 -> degreeZeroBoundary
          1 -> degreeOneBoundary
          2 -> degreeTwoBoundary
          _ -> emptyBoundaryIncidence
  chainComplex <-
    first (ActivationChainComplexFailed stage)
      (mkFiniteChainComplexChecked (HomologicalDegree 2) boundaryAt)
  pure
    ActivationComplex
      { activationComplexStage = stage,
        activationComplexCells = cells,
        activationComplexBasis = basis,
        activationComplexChainComplex = chainComplex
      }

activationBasisFromCells :: Set (ActivationCell label) -> ActivationBasis label
activationBasisFromCells cells =
  foldr insertCell (ActivationBasis [] [] []) (Set.toAscList cells)
 where
  insertCell :: ActivationCell label -> ActivationBasis label -> ActivationBasis label
  insertCell cellValue basis =
    case cellValue of
      ActivationVertex _ -> basis{activationVertexBasis = cellValue : activationVertexBasis basis}
      ActivationEdge _ _ -> basis{activationEdgeBasis = cellValue : activationEdgeBasis basis}
      ActivationFace _ _ _ -> basis{activationFaceBasis = cellValue : activationFaceBasis basis}

activationCellBoundary :: Num coefficient => ActivationCell label -> [(coefficient, ActivationCell label)]
activationCellBoundary cellValue =
  case cellValue of
    ActivationVertex _ -> []
    ActivationEdge firstLabel secondLabel ->
      [ (-1, ActivationVertex firstLabel),
        (1, ActivationVertex secondLabel)
      ]
    ActivationFace firstLabel secondLabel thirdLabel ->
      [ (1, ActivationEdge secondLabel thirdLabel),
        (-1, ActivationEdge firstLabel thirdLabel),
        (1, ActivationEdge firstLabel secondLabel)
      ]

buildActivationZigzag ::
  Ord label =>
  NonEmpty (depth, ActivationComplex depth label) ->
  Either (ActivationZigzagError depth label) (ActivationZigzag depth label)
buildActivationZigzag observedEntries@((_, firstComplex) :| remainingEntries) = do
  correspondences <-
    traverse
      ( \((leftDepth, leftComplex), (rightDepth, rightComplex)) ->
          adjacentUnionWitness leftDepth rightDepth leftComplex rightComplex
      )
      (zip (NonEmpty.toList observedEntries) remainingEntries)
  let arrows =
        correspondences
          >>= ( \witness ->
                  [ ForwardArrow (adjacentLeftInclusion witness),
                    BackwardArrow (adjacentRightInclusion witness)
                  ]
              )
  diagram <-
    first ActivationZigzagFailed
      ( mkFiniteChainZigzag
          (activationComplexChainComplex firstComplex)
          arrows
      )
  pure
    ActivationZigzag
      { activationZigzagComplexes = fmap snd observedEntries,
        activationZigzagCorrespondences = Vector.fromList correspondences,
        activationZigzagDiagram = diagram
      }

adjacentUnionWitness ::
  Ord label =>
  depth ->
  depth ->
  ActivationComplex depth label ->
  ActivationComplex depth label ->
  Either (ActivationZigzagError depth label) (AdjacentUnionWitness depth label)
adjacentUnionWitness leftDepth rightDepth leftComplex rightComplex = do
  let unionStage =
        AdjacentDepthUnion leftDepth rightDepth
      unionCells =
        Set.union
          (activationComplexCells leftComplex)
          (activationComplexCells rightComplex)
  unionComplex <- compileActivationComplex unionStage unionCells
  leftInclusion <- activationInclusion leftComplex unionComplex
  rightInclusion <- activationInclusion rightComplex unionComplex
  pure
    AdjacentUnionWitness
      { adjacentCommonCells =
          Set.intersection
            (activationComplexCells leftComplex)
            (activationComplexCells rightComplex),
        adjacentUnionComplex = unionComplex,
        adjacentLeftInclusion = leftInclusion,
        adjacentRightInclusion = rightInclusion
      }

activationInclusion ::
  Ord label =>
  ActivationComplex depth label ->
  ActivationComplex depth label ->
  Either (ActivationZigzagError depth label) (FiniteChainMap Int)
activationInclusion sourceComplex targetComplex =
  first ActivationZigzagFailed
    ( mkFiniteChainMapChecked
        (activationComplexChainComplex sourceComplex)
        (activationComplexChainComplex targetComplex)
        (inclusionAtDegree sourceComplex targetComplex)
    )

inclusionAtDegree ::
  Ord label =>
  ActivationComplex depth label ->
  ActivationComplex depth label ->
  HomologicalDegree ->
  BoundaryIncidence Int
inclusionAtDegree sourceComplex targetComplex degreeValue =
  overlapBoundaryIncidence
    (1 :: Int)
    (activationComplexBasisAt sourceComplex degreeValue)
    (activationComplexBasisAt targetComplex degreeValue)

activationBasisAt :: HomologicalDegree -> ActivationBasis label -> [ActivationCell label]
activationBasisAt (HomologicalDegree degreeIndex) basis =
  case degreeIndex of
    0 -> activationVertexBasis basis
    1 -> activationEdgeBasis basis
    2 -> activationFaceBasis basis
    _ -> []

requireStage ::
  Vector (ActivationStage depth) ->
  Int ->
  Either (ActivationZigzagError depth label) (ActivationStage depth)
requireStage stages stageIndex =
  maybe
    (Left (ActivationStageIndexMissing stageIndex))
    Right
    (stages Vector.!? stageIndex)