content: {-# OPTIONS_GHC -fno-warn-orphans #-}
  -----------------------------------------------------------------------------
  -- |
  -- Module     : Algebra.Graph.Test.Arbitrary
  -- Copyright  : (c) Andrey Mokhov 2016-2025
  -- License    : MIT (see the file LICENSE)
  -- Maintainer : andrey.mokhov@gmail.com
  -- Stability  : experimental
  --
  -- Generators and orphan Arbitrary instances for various data types.
  -----------------------------------------------------------------------------
  module Algebra.Graph.Test.Arbitrary (
      -- * Generators of arbitrary graph instances
      arbitraryGraph, arbitraryRelation, arbitraryAdjacencyMap,
      ) where
  
  import Data.List.NonEmpty (NonEmpty (..), toList)
  import Data.Maybe (catMaybes)
  import Test.QuickCheck
  
  import Algebra.Graph
  import Algebra.Graph.Export
  import Algebra.Graph.Label
  
  import qualified Algebra.Graph.Undirected                       as UG
  import qualified Algebra.Graph.Acyclic.AdjacencyMap             as AAM
  import qualified Algebra.Graph.AdjacencyIntMap                  as AIM
  import qualified Algebra.Graph.AdjacencyMap                     as AM
  import qualified Algebra.Graph.Bipartite.AdjacencyMap           as BAM
  import qualified Algebra.Graph.Bipartite.AdjacencyMap.Algorithm as BAMA
  import qualified Algebra.Graph.NonEmpty.AdjacencyMap            as NAM
  import qualified Algebra.Graph.Class                            as C
  import qualified Algebra.Graph.Labelled                         as LG
  import qualified Algebra.Graph.Labelled.AdjacencyMap            as LAM
  import qualified Algebra.Graph.NonEmpty                         as NonEmpty
  import qualified Algebra.Graph.Relation                         as Relation
  import qualified Algebra.Graph.Relation.Preorder                as Preorder
  import qualified Algebra.Graph.Relation.Reflexive               as Reflexive
  import qualified Algebra.Graph.Relation.Symmetric               as Symmetric
  import qualified Algebra.Graph.Relation.Transitive              as Transitive
  
  -- | Generate an arbitrary 'C.Graph' value of a specified size.
  arbitraryGraph :: (C.Graph g, Arbitrary (C.Vertex g)) => Gen g
  arbitraryGraph = sized expr
    where
      expr 0 = return C.empty
      expr 1 = C.vertex <$> arbitrary
      expr n = do
          left <- choose (0, n)
          oneof [ C.overlay <$> expr left <*> expr (n - left)
                , C.connect <$> expr left <*> expr (n - left) ]
  
  instance Arbitrary a => Arbitrary (Graph a) where
      arbitrary = arbitraryGraph
  
      shrink Empty         = []
      shrink (Vertex    _) = [Empty]
      shrink (Overlay x y) = [Empty, x, y]
                          ++ [Overlay x' y' | (x', y') <- shrink (x, y) ]
      shrink (Connect x y) = [Empty, x, y, Overlay x y]
                          ++ [Connect x' y' | (x', y') <- shrink (x, y) ]
  
  -- An Arbitrary instance for Graph.Undirected
  instance Arbitrary a => Arbitrary (UG.Graph a) where
      arbitrary = arbitraryGraph
  
  -- An Arbitrary instance for Acyclic.AdjacencyMap
  instance (Ord a, Arbitrary a) => Arbitrary (AAM.AdjacencyMap a) where
      arbitrary = AAM.shrink <$> arbitrary
  
      shrink g = shrinkVertices ++ shrinkEdges
        where
          shrinkVertices =
            let vertices = AAM.vertexList g
            in [ AAM.removeVertex x g | x <- vertices ]
  
          shrinkEdges =
            let edges = AAM.edgeList g
            in [ AAM.removeEdge x y g | (x, y) <- edges ]
  
  -- | Generate an arbitrary 'NonEmpty.Graph' value of a specified size.
  arbitraryNonEmptyGraph :: Arbitrary a => Gen (NonEmpty.Graph a)
  arbitraryNonEmptyGraph = sized expr
    where
      expr 0 = NonEmpty.vertex <$> arbitrary -- can't generate non-empty graph of size 0
      expr 1 = NonEmpty.vertex <$> arbitrary
      expr n = do
          left <- choose (1, n)
          oneof [ NonEmpty.overlay <$> expr left <*> expr (n - left)
                , NonEmpty.connect <$> expr left <*> expr (n - left) ]
  
  instance Arbitrary a => Arbitrary (NonEmpty.Graph a) where
      arbitrary = arbitraryNonEmptyGraph
  
      shrink (NonEmpty.Vertex    _) = []
      shrink (NonEmpty.Overlay x y) = [x, y]
          ++ [NonEmpty.Overlay x' y' | (x', y') <- shrink (x, y) ]
      shrink (NonEmpty.Connect x y) = [x, y, NonEmpty.Overlay x y]
          ++ [NonEmpty.Connect x' y' | (x', y') <- shrink (x, y) ]
  
  -- | Generate an arbitrary 'Relation'.
  arbitraryRelation :: (Arbitrary a, Ord a) => Gen (Relation.Relation a)
  arbitraryRelation = Relation.stars <$> arbitrary
  
  -- TODO: Implement a custom shrink method.
  instance (Arbitrary a, Ord a) => Arbitrary (Relation.Relation a) where
      arbitrary = arbitraryRelation
  
      shrink g = shrinkVertices ++ shrinkEdges
        where
           shrinkVertices =
             let vertices = Relation.vertexList g
             in  [ Relation.removeVertex v g | v <- vertices ]
  
           shrinkEdges =
             let edges = Relation.edgeList g
             in  [ Relation.removeEdge v w g | (v, w) <- edges ]
  
  -- TODO: Simplify.
  instance (Arbitrary a, Ord a) => Arbitrary (Reflexive.ReflexiveRelation a) where
      arbitrary = Reflexive.fromRelation . Relation.reflexiveClosure
          <$> arbitraryRelation
  
  instance (Arbitrary a, Ord a) => Arbitrary (Symmetric.Relation a) where
      arbitrary = Symmetric.toSymmetric <$> arbitraryRelation
  
  instance (Arbitrary a, Ord a) => Arbitrary (Transitive.TransitiveRelation a) where
      arbitrary = Transitive.fromRelation . Relation.transitiveClosure
          <$> arbitraryRelation
  
  instance (Arbitrary a, Ord a) => Arbitrary (Preorder.PreorderRelation a) where
      arbitrary = Preorder.fromRelation . Relation.closure
          <$> arbitraryRelation
  
  -- | Generate an arbitrary 'AdjacencyMap'. It is guaranteed that the
  -- resulting adjacency map is 'consistent'.
  arbitraryAdjacencyMap :: (Arbitrary a, Ord a) => Gen (AM.AdjacencyMap a)
  arbitraryAdjacencyMap = AM.stars <$> arbitrary
  
  instance (Arbitrary a, Ord a) => Arbitrary (AM.AdjacencyMap a) where
      arbitrary = arbitraryAdjacencyMap
  
      shrink g = shrinkVertices ++ shrinkEdges
        where
           shrinkVertices = [ AM.removeVertex v g | v <- AM.vertexList g ]
           shrinkEdges    = [ AM.removeEdge v w g | (v, w) <- AM.edgeList g ]
  
  -- | Generate an arbitrary non-empty 'NAM.AdjacencyMap'. It is guaranteed that
  -- the resulting adjacency map is 'consistent'.
  arbitraryNonEmptyAdjacencyMap :: (Arbitrary a, Ord a) => Gen (NAM.AdjacencyMap a)
  arbitraryNonEmptyAdjacencyMap = NAM.stars1 <$> nonEmpty
    where
      nonEmpty = do
          xs <- arbitrary
          case xs of
              [] -> do
                  x <- arbitrary
                  return ((x, []) :| []) -- There must be at least one vertex
              (x:xs) -> return (x :| xs)
  
  instance (Arbitrary a, Ord a) => Arbitrary (NAM.AdjacencyMap a) where
      arbitrary = arbitraryNonEmptyAdjacencyMap
  
      shrink g = shrinkVertices ++ shrinkEdges
        where
           shrinkVertices =
             let vertices = toList $ NAM.vertexList1 g
             in catMaybes [ NAM.removeVertex1 v g | v <- vertices ]
  
           shrinkEdges =
             let edges = NAM.edgeList g
             in  [ NAM.removeEdge v w g | (v, w) <- edges ]
  
  instance Arbitrary AIM.AdjacencyIntMap where
      arbitrary = AIM.stars <$> arbitrary
  
      shrink g = shrinkVertices ++ shrinkEdges
        where
           shrinkVertices = [ AIM.removeVertex x g | x <- AIM.vertexList g ]
           shrinkEdges    = [ AIM.removeEdge x y g | (x, y) <- AIM.edgeList g ]
  
  -- | Generate an arbitrary labelled 'LAM.AdjacencyMap'. It is guaranteed
  -- that the resulting adjacency map is 'consistent'.
  arbitraryLabelledAdjacencyMap :: (Arbitrary a, Ord a, Eq e, Arbitrary e, Monoid e) => Gen (LAM.AdjacencyMap e a)
  arbitraryLabelledAdjacencyMap = LAM.fromAdjacencyMaps <$> arbitrary
  
  instance (Arbitrary a, Ord a, Eq e, Arbitrary e, Monoid e) => Arbitrary (LAM.AdjacencyMap e a) where
      arbitrary = arbitraryLabelledAdjacencyMap
  
      shrink g = shrinkVertices ++ shrinkEdges
        where
           shrinkVertices =
             let vertices = LAM.vertexList g
             in  [ LAM.removeVertex v g | v <- vertices ]
  
           shrinkEdges =
             let edges = LAM.edgeList g
             in  [ LAM.removeEdge v w g | (_, v, w) <- edges ]
  
  -- | Generate an arbitrary labelled 'LAM.Graph' value of a specified size.
  arbitraryLabelledGraph :: (Arbitrary a, Arbitrary e) => Gen (LG.Graph e a)
  arbitraryLabelledGraph = sized expr
    where
      expr 0 = return LG.empty
      expr 1 = LG.vertex <$> arbitrary
      expr n = do
          label <- arbitrary
          left  <- choose (0, n)
          LG.connect label <$> expr left <*> expr (n - left)
  
  instance (Arbitrary a, Arbitrary e, Monoid e) => Arbitrary (LG.Graph e a) where
      arbitrary = arbitraryLabelledGraph
  
      shrink LG.Empty           = []
      shrink (LG.Vertex      _) = [LG.Empty]
      shrink (LG.Connect e x y) = [LG.Empty, x, y, LG.Connect mempty x y]
                               ++ [LG.Connect e x' y' | (x', y') <- shrink (x, y) ]
  
  -- TODO: Implement a custom shrink method.
  instance Arbitrary s => Arbitrary (Doc s) where
      arbitrary = mconcat . map literal <$> arbitrary
  
  instance (Arbitrary a, Num a, Ord a) => Arbitrary (Distance a) where
      arbitrary = (\x -> if x < 0 then distance infinite else distance (unsafeFinite x)) <$> arbitrary
  
  instance (Arbitrary a, Num a, Ord a) => Arbitrary (Capacity a) where
      arbitrary = (\x -> if x < 0 then capacity infinite else capacity (unsafeFinite x)) <$> arbitrary
  
  instance (Arbitrary a, Num a, Ord a) => Arbitrary (Count a) where
      arbitrary = (\x -> if x < 0 then count infinite else count (unsafeFinite x)) <$> arbitrary
  
  instance Arbitrary a => Arbitrary (Minimum a) where
      arbitrary = frequency [(10, pure <$> arbitrary), (1, pure noMinimum)]
  
  instance (Arbitrary a, Ord a) => Arbitrary (PowerSet a) where
      arbitrary = PowerSet <$> arbitrary
  
  instance (Arbitrary o, Arbitrary a) => Arbitrary (Optimum o a) where
      arbitrary = Optimum <$> arbitrary <*> arbitrary
  
  instance (Arbitrary a, Arbitrary b, Ord a, Ord b) => Arbitrary (BAM.AdjacencyMap a b) where
      arbitrary = BAM.toBipartite <$> arbitrary
      shrink = map BAM.toBipartite . shrink . BAM.fromBipartite
  
  instance (Arbitrary a, Arbitrary b) => Arbitrary (BAM.List a b) where
      arbitrary = sized go
        where
          go 0 = return BAM.Nil
          go 1 = do h <- arbitrary
                    return $ BAM.Cons h BAM.Nil
          go n = do f <- arbitrary
                    s <- arbitrary
                    (BAM.Cons f . BAM.Cons s) <$> go (n - 2)
  
  instance (Arbitrary a, Arbitrary b, Ord a, Ord b) => Arbitrary (BAMA.Matching a b) where
      arbitrary = BAMA.matching <$> arbitrary
  