Tracking.hs

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Online operations for structure recognizer.
--
-- See "Swarm.Game.Scenario.Topography.Structure.Recognition.Precompute" for
-- details of the structure recognition process.
module Swarm.Game.Scenario.Topography.Structure.Recognition.Tracking (
  RecognitionActiveStatus (..),
  entityModified,
  entityModifiedLoggable,
) where

import Control.Arrow (left, (&&&))
import Control.Lens ((%~), (&), (.~), (^.))
import Control.Monad (foldM, forM_, guard, unless)
import Control.Monad.Extra (findM)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Maybe (MaybeT (..), runMaybeT)
import Control.Monad.Trans.Writer.Strict
import Data.Either (partitionEithers)
import Data.Either.Extra (maybeToEither)
import Data.Function (on)
import Data.HashMap.Strict qualified as HM
import Data.HashSet qualified as HS
import Data.Hashable (Hashable)
import Data.IntSet qualified as IS
import Data.IntSet.NonEmpty (NEIntSet)
import Data.IntSet.NonEmpty qualified as NEIS
import Data.List (sortOn)
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Ord (Down (..))
import Data.Semigroup (Max (..), Min (..))
import Data.Tuple (swap)
import Linear (V2 (..))
import Swarm.Game.Location (Location)
import Swarm.Game.Scenario.Topography.Structure.Named (name)
import Swarm.Game.Scenario.Topography.Structure.Recognition
import Swarm.Game.Scenario.Topography.Structure.Recognition.Log
import Swarm.Game.Scenario.Topography.Structure.Recognition.Precompute (GenericEntLocator, ensureStructureIntact)
import Swarm.Game.Scenario.Topography.Structure.Recognition.Prep (binTuplesHM)
import Swarm.Game.Scenario.Topography.Structure.Recognition.Registry
import Swarm.Game.Scenario.Topography.Structure.Recognition.Type
import Swarm.Game.Scenario.Topography.Terraform
import Swarm.Game.Universe
import Text.AhoCorasick

data RecognitionActiveStatus
  = RecognizeNewStructures
  | -- | Do not add new recognitions to the registry.
    -- This is useful if one needs to construct a larger structure
    -- for which other smaller structures contained within it
    -- would otherwise be recognized first, precluding the larger
    -- structure from ever being recognized.
    -- Removing elements of a previously recognized structure
    -- will still cause it to be removed from the registry.
    DisableNewRecognition
  deriving (Show, Eq, Ord, Enum, Bounded)

-- | A hook called from the centralized entity update function,
-- 'Swarm.Game.Step.Util.updateEntityAt'.
entityModified ::
  (Monad s, Hashable a, Eq b) =>
  GenericEntLocator s a ->
  CellModification a ->
  Cosmic Location ->
  RecognizerAutomatons b a ->
  RecognitionState b a ->
  s (RecognitionState b a)
entityModified entLoader modification cLoc autoRecognizer oldRecognitionState = do
  (val, accumulatedLogs) <-
    runWriterT $
      entityModifiedLoggable RecognizeNewStructures entLoader modification cLoc autoRecognizer oldRecognitionState
  return $
    val
      & recognitionLog %~ (reverse accumulatedLogs <>)

-- | This handles structure detection upon addition of an entity,
-- and structure de-registration upon removal of an entity.
-- Also handles atomic entity swaps.
entityModifiedLoggable ::
  (Monoid (f (SearchLog a)), Monad m, Hashable a, Eq b, Applicative f) =>
  RecognitionActiveStatus ->
  (Cosmic Location -> m (AtomicKeySymbol a)) ->
  CellModification a ->
  Cosmic Location ->
  RecognizerAutomatons b a ->
  RecognitionState b a ->
  WriterT (f (SearchLog a)) m (RecognitionState b a)
entityModifiedLoggable activeStatus entLoader modification cLoc autoRecognizer oldRecognitionState = do
  case modification of
    Add newEntity -> doAddition newEntity oldRecognitionState
    Remove _ -> doRemoval oldRecognitionState
    Swap _ newEntity -> doRemoval oldRecognitionState >>= doAddition newEntity
 where
  entLookup = autoRecognizer ^. automatonsByEntity

  doAddition newEntity = case activeStatus of
    RecognizeNewStructures -> maybe return logAndRegister $ HM.lookup newEntity entLookup
    DisableNewRecognition -> return
   where
    logAndRegister finder s = do
      tell . pure . FoundParticipatingEntity $
        ParticipatingEntity
          newEntity
          (finder ^. inspectionOffsets)
      newFoundStructures <- registerRowMatches entLoader cLoc finder $ s ^. foundStructures
      return $ s & foundStructures .~ newFoundStructures

  doRemoval sOld =
    -- Entity was removed; may need to remove registered structure.
    maybe return logAndRemove structureAtLoc sOld
   where
    structureAtLoc = M.lookup cLoc $ foundByLocation $ sOld ^. foundStructures
    logAndRemove fs s = do
      tell $ pure $ StructureRemoved structureName
      return $ s & foundStructures %~ removeStructure fs
     where
      structureName = name . originalItem . entityGrid $ structureWithGrid fs

-- | In case this cell would match a candidate structure,
-- ensures that the entity in this cell is not already
-- participating in a registered structure.
--
-- Furthermore, treating cells in registered structures
-- as 'Nothing' has the effect of "masking" them out,
-- so that they can overlap empty cells within the bounding
-- box of the candidate structure.
candidateEntityAt ::
  (Monad s, Hashable a) =>
  GenericEntLocator s a ->
  FoundRegistry b a ->
  Cosmic Location ->
  s (AtomicKeySymbol a)
candidateEntityAt entLoader registry cLoc = runMaybeT $ do
  guard $ M.notMember cLoc $ foundByLocation registry
  MaybeT $ entLoader cLoc

-- | Excludes entities that are already part of a
-- registered found structure.
getWorldRow ::
  (Monad s, Hashable a) =>
  GenericEntLocator s a ->
  FoundRegistry b a ->
  Cosmic Location ->
  InspectionOffsets ->
  s [AtomicKeySymbol a]
getWorldRow entLoader registry cLoc (InspectionOffsets (Min offsetLeft) (Max offsetRight)) = do
  mapM getCandidate horizontalOffsets
 where
  getCandidate = candidateEntityAt entLoader registry
  horizontalOffsets = map mkLoc [offsetLeft .. offsetRight]
  mkLoc x = cLoc `offsetBy` V2 x 0

-- | This runs once per non-overlapping subset of found chunks
checkChunksCombination ::
  (Monoid (f (SearchLog a)), Applicative f, Monad m, Hashable a, Eq b) =>
  Cosmic Location ->
  InspectionOffsets ->
  NE.NonEmpty (RowChunkMatchingReference b a) ->
  [Position (NE.NonEmpty a)] ->
  WriterT (f (SearchLog a)) m [FoundStructure b a]
checkChunksCombination
  cLoc
  horizontalOffsets
  rowChunkReferences
  candidatesChunked = do
    tell . pure . FoundPiecewiseChunks . map swap . HM.toList $
      fmap NEIS.elems foundRowChunksLookup

    tell . pure . ChunkFailures $ candidateFailures

    tell . pure . ChunksMatchingExpected $
      map (modifyChunkedRowMatch $ fmap renderSharedNames) candidateExpected

    return structurePositionsToCheck
   where
    structurePositionsToCheck = concatMap mkFoundStructures candidateExpected

    candidateExpected = concatMap NE.toList candidateExpectedLists

    foundRowChunksLookup =
      fmap NEIS.fromList $
        binTuplesHM $
          map (pVal &&& pIndex) candidatesChunked

    (candidateFailures, candidateExpectedLists) =
      partitionEithers $
        map (checkCandidateAgainstObservedChunks horizontalOffsets foundRowChunksLookup) $
          NE.toList rowChunkReferences

    mkFoundStructures x =
      NE.toList $ NE.map mkFoundStructure . referencingRows . chunkStructure $ foundChunkRow x
     where
      mkFoundStructure r =
        PositionedStructure
          (cLoc `offsetBy` theOffset)
          (wholeStructure r)
       where
        theOffset = V2 (horizontalStructPos $ foundChunkRow x) (rowIndex r)

    modifyChunkedRowMatch f (ChunkedRowMatch x y) = ChunkedRowMatch x (f y)

checkCandidateAgainstObservedChunks ::
  Hashable e =>
  InspectionOffsets ->
  HM.HashMap (NE.NonEmpty e) NEIntSet ->
  RowChunkMatchingReference b e ->
  Either (ChunkMatchFailureReason e) (NE.NonEmpty (ChunkedRowMatch (ConsolidatedRowReferences b e) e))
checkCandidateAgainstObservedChunks horizontalOffsets foundRowChunksLookup (RowChunkMatchingReference r chunkPositionMap) =
  left (ChunkMatchFailureReason $ renderSharedNames r) $ do
    unless isKeysSubset . Left $
      NoKeysSubset $
        (FoundChunkComparison `on` HS.toList) foundChunksKeys referenceChunksKeys

    nonEmptyPairs <-
      maybeToEither EmptyIntersection $
        NE.nonEmpty sortedByAlignmentChoices

    let maybeViables = do
          possibles <- seedPossibleOffsets $ snd $ NE.head nonEmptyPairs
          foldM findCoveringOffsets possibles $ NE.map (snd . snd) nonEmptyPairs

    viableRowOffsets <- maybeToEither EmptyIntersection maybeViables
    return $ NE.map mkRowMatch $ NEIS.toList viableRowOffsets
 where
  theIntersection =
    HM.intersectionWith
      FoundAndExpectedChunkPositions
      foundRowChunksLookup
      modifiedChunkPositionMap
  intersectionWithSizeDifferences = HM.map (sizeDifference &&& id) theIntersection
   where
    sizeDifference x = (subtract `on` NEIS.size) (expectedPositions x) (foundPositions x)

  -- Remove the pairings that have fewer occurrences than the required number.
  -- The 'fst' element of the tuple is the difference between the "observed" and "required" count.
  withSufficientCoverage = HM.filter ((>= 0) . fst) intersectionWithSizeDifferences
  sortedByAlignmentChoices = sortOn (fst . snd) $ HM.toList withSufficientCoverage

  isKeysSubset = referenceChunksKeys `HS.isSubsetOf` foundChunksKeys

  mkRowMatch rowOffset =
    ChunkedRowMatch
      (map swap $ HM.toList theIntersection)
      (FoundRowFromChunk rowOffset horizontalStructurePosition r)
   where
    horizontalStructurePosition = fromIntegral rowOffset + getMin (startOffset horizontalOffsets)

  modifiedChunkPositionMap = fmap NEIS.fromList chunkPositionMap
  foundChunksKeys = HM.keysSet foundRowChunksLookup
  referenceChunksKeys = HM.keysSet chunkPositionMap

-- | Search for any structure row that happens to
-- contain the placed entity.
registerRowMatches ::
  (Monoid (f (SearchLog a)), Applicative f, Monad s, Hashable a, Eq b) =>
  GenericEntLocator s a ->
  Cosmic Location ->
  AutomatonInfo b a ->
  FoundRegistry b a ->
  WriterT (f (SearchLog a)) s (FoundRegistry b a)
registerRowMatches entLoader cLoc (AutomatonInfo horizontalOffsets pwMatcher) registry = do
  tell $ pure $ StartSearchAt cLoc horizontalOffsets

  tell . pure . ExpectedChunks $
    NE.map (HM.keys . confirmationMap) rowChunkReferences

  entitiesRow <-
    lift $
      getWorldRow
        entLoader
        registry
        cLoc
        horizontalOffsets

  tell . pure . WorldRowContent $ entitiesRow

  let candidatesChunked = findAll pwSM entitiesRow
  unrankedCandidateStructures <- checkCombo candidatesChunked

  -- [STRUCTURE RECOGNIZER CONFLICT RESOLUTION]
  -- We only allow an entity to participate in one structure at a time,
  -- so multiple matches require a tie-breaker.
  -- The largest structure (by area) shall win.
  -- Sort by decreasing order of preference
  -- (see the Ord instance of 'FoundStructure').
  let rankedCandidates = sortOn Down unrankedCandidateStructures
  tell . pure . FoundCompleteStructureCandidates $
    map getStructInfo rankedCandidates

  -- We should not check all of the structures, which can be expensive.
  -- Instead, we ranked the candidates by preference a-priori
  -- and now choose the first one that is verified.
  maybeIntactStructure <- findM validateIntactness2d rankedCandidates

  forM_ maybeIntactStructure $
    tell . pure . RecognizedSingleStructure . getStructInfo

  return $ maybe id addFound maybeIntactStructure registry
 where
  PiecewiseRecognition pwSM rowChunkReferences = pwMatcher

  getStructInfo (PositionedStructure loc swg) = (distillLabel swg, loc)

  validateIntactness2d fs = do
    maybeIntactnessFailure <- lift $ ensureStructureIntact registry entLoader fs
    tell . pure . ChunkIntactnessVerification
      $ IntactPlacementLog
        maybeIntactnessFailure
      $ PositionedStructure (upperLeftCorner fs) (distillLabel . structureWithGrid $ fs)

    return $ null maybeIntactnessFailure

  checkCombo = checkChunksCombination cLoc horizontalOffsets rowChunkReferences

-- |
-- For a given "chunk", there could be multiple recurrences.
-- However, the position of each recurrence is unique
-- (i.e. the chunk cannot exist twice at the same location).
--
-- Either:
-- A) An observed chunk is "superfluous" w.r.t. matching the candidate, or
-- B) It is necessary for the match.
--
-- The lowest-numbered "reference position" (i.e. in the structure definition)
-- of a given chunk must align with exactly one "observed position".
--
-- The difference between the "observed" position of the chunk that aligns with the
-- lowest-numbered "reference position" shall be the global "row offset" applied to our observations.
-- This row offset value applies to all "chunks" (both identical and distinct) that comprise the row.
--
-- If a given chunk occurrence is necessary for the match, then we may attempt to use it to compute
-- the "row offset" by taking its position minus the lowest-numbered "reference position".
--
-- We can iterate over each occurrence position in ascending order.
-- In the ideal case, the first such candidate ends up being the the actual, valid, offset.
-- Otherwise, we know that all invalid offset candidates encountered before the first valid
-- offset constitute "superfluous" chunks.
--
-- Note that there may exist multiple valid "row offsets".
-- At most, there will be
--   {number of observed occurrences} minus {number of required occurrences}
-- such offsets.
--
-- = Performance notes
--
-- We only have to do this computation once, and only for the "smallest" size discrepancy
-- between occurrences and references of a chunk. This generates the "seed" pool of possible offsets.
-- All subsequent chunks will merely filter on this initial set.
seedPossibleOffsets :: (Int, FoundAndExpectedChunkPositions) -> Maybe NEIntSet
seedPossibleOffsets (sizeDifference, FoundAndExpectedChunkPositions found expected) =
  NEIS.nonEmptySet $ IS.fromList possibleOffsets
 where
  possibleOffsets =
    NE.take (sizeDifference + 1) $
      NE.map (subtract (NEIS.findMin expected)) $
        NEIS.toAscList found

-- | Return all of the offsets that are viable for repetitions of this chunk.
--
-- Note that if there are an equal number of observed occurrences
-- and expected occurrences, then there is only one possible offset.
-- If there are N expected and (N + 1) observed, then there are 2 possible offsets.
findCoveringOffsets :: NEIntSet -> FoundAndExpectedChunkPositions -> Maybe NEIntSet
findCoveringOffsets possibleOffsets x =
  NEIS.nonEmptySet $ NEIS.filter (isCoveredWithOffset x) possibleOffsets

isCoveredWithOffset :: FoundAndExpectedChunkPositions -> Int -> Bool
isCoveredWithOffset (FoundAndExpectedChunkPositions found expected) offset =
  NEIS.map (+ offset) expected `NEIS.isSubsetOf` found