{-|
Module      : Nauty.Sparse.Internal
Description : Internal functions.
Copyright   : (c) Marcelo Garlet Milani, 2026
License     : MIT
Maintainer  : mgmilani@pm.me
Stability   : unstable

This module contains internal functions used by the "Nauty.Sparse6" module.
Except for test cases, you should not import this module.
-}

{-# LANGUAGE OverloadedStrings #-}

module Nauty.Sparse6.Internal where

import Control.Monad.Trans.State
import Data.Bits
import Data.List
import Data.Word
import Data.Maybe
import Control.Monad.Trans.Class
import Nauty.Internal.Encoding
import Nauty.Internal.Parsing
import Nauty.Internal.Utils
import qualified Data.Array.Unboxed as A
import qualified Data.ByteString.Lazy as B
import qualified Data.Text.Lazy as T
import qualified Data.Text.Lazy.Encoding as T
import qualified Data.Set as S

-- | A graph represented as an adjacency list.
data AdjacencyList = AdjacencyList
  { numberOfVertices :: Word64
  , adjacency        :: A.UArray Word64 Word64 -- ^ List of edges. Edges have the form {A[2i], A[2i+1]}, for each i from 0 to half of the size of the array.
  } deriving (Eq, Show)

-- | Parse a file of graphs in @sparse6@ format.
parse :: T.Text -> [Either T.Text AdjacencyList]
parse t = 
  let t' = header ">>sparse6<<" t
  in map graph $ T.lines t'

-- | Decode a single graph in @sparse6@ format.
graph :: T.Text -> Either T.Text AdjacencyList
graph str = (flip evalStateT) (T.encodeUtf8 str) $ do
  h <- consume 1
  if B.unpack h /= [fromIntegral $ fromEnum ':'] then
    lift $ Left $ "Expected ':', but found " <> (T.pack $ map (toEnum . fromIntegral) $ B.unpack h)
  else do
    n <- parseNumber
    if n == 0 then
      return AdjacencyList{ numberOfVertices = 0, adjacency = A.array (0,0) [(0,0)]}
    else do
      let k = numBits $ n - 1
      parseEdgeList n k

-- | Parse a graph in incremental @sparse6@ format.
-- The edges of the first graph are the symmetric difference between the edges of the second graph and the resulting graph.
symmetricDifference :: AdjacencyList -- ^ Reference graph
                    -> T.Text -- ^ symmetric difference between reference graph and the resulting graph, in incremental @sparse6@ format
                    -> Either T.Text AdjacencyList
symmetricDifference g hTxt = 
  (flip evalStateT) (T.encodeUtf8 hTxt) $ do
    h <- consume 1
    let hu = B.unpack h
    if hu == [fromIntegral $ fromEnum ';'] then do
      let n = (numberOfVertices g)
      let k = numBits $ n - 1
      g1 <- parseEdgeList n k
      let es0 = S.fromList $ snd $ toEdgeList g
          es1 = S.fromList $ snd $ toEdgeList g1
      return $ fromEdgeList (numberOfVertices g) (S.toList $ (es0 `S.difference` es1) `S.union` (es1 `S.difference` es0))
    else if hu == [fromIntegral $ fromEnum ':'] then do
      n <- parseNumber
      if n == 0 then
        return AdjacencyList{ numberOfVertices = 0, adjacency = A.array (0,0) [(0,0)]}
      else do
        let k = numBits $ n - 1
        parseEdgeList n k
    else
      lift $ Left $ "Expected ':' or ';', but found " <> (T.pack $ map (toEnum . fromIntegral) $ B.unpack h)

-- |State for encoding a graph.
data EncodeState = EncodeState
  { currentV    :: Word64
  , currentByte :: Word8
  , usedBits    :: Int
  , n2HasEdge   :: Bool
  , n1HasEdge   :: Bool
  , encoding    :: [Word8]
  } deriving (Eq, Show)

-- | Encode a graph in @sparse6@ format.
encode :: AdjacencyList -> T.Text
encode g = 
  ":"
  `T.append`
  (encodeNumber $ numberOfVertices g)
  `T.append`
  (encodeEdgeList (numberOfVertices g) $ 
      groupByTwo $ A.elems $ adjacency g)

-- | Encode a graph using the incremental @sparse6@ format.
-- The list of edges given is the symmetric difference between the current graph and the previous one in the list.
encodeSymmetricDifference :: Word64 -> [(Word64, Word64)] -> T.Text
encodeSymmetricDifference n es =
  ";"
  `T.append`
  (encodeEdgeList n es)

-- |Encode a list of edges.
encodeEdgeList :: Word64 -> [(Word64, Word64)] -> T.Text
encodeEdgeList n es = 
  let k = numBits $ n - 1 :: Int
      encodeEdges :: [(Word64, Word64)] -> StateT EncodeState Maybe T.Text
      encodeEdges [] = do
        st <- get
        if usedBits st == 0 then
          return $ T.pack $ map (\x -> toEnum $ fromIntegral $ x + 63) $ reverse $ 
            encoding st
        else
          let padding = 6 - usedBits st in
          if (n,k) `elem` [(2,1), (4,2), (8,3), (16,4)] 
             && padding >= k + 1 
             && n2HasEdge st
             && not (n1HasEdge st)
          then
            return $ T.pack $ map (\x -> toEnum $ fromIntegral $ x + 63) $ reverse $ 
              ((foldr 
                   (\i b -> setBit b i)
                   (shiftL (currentByte st) padding)
                   [0..padding - 2]) 
                 : encoding st)
          else
            return $ T.pack $ map (\x -> toEnum $ fromIntegral $ x + 63) $ reverse $ 
              ((foldr 
                 (\i b -> setBit b i)
                 (shiftL (currentByte st) padding)
                 [0..padding - 1]) 
               : encoding st)
      encodeEdges ((v', u) : es') = do
        if v' < u then
          encodeEdges ((u, v') : es')
        else do
          st <- get
          if v'  > (currentV st) + 1 then do
            if v' == n - 2 then
              put st{currentV = v', n2HasEdge = True}
            else if v' == n - 1 then
              put st{currentV = v', n1HasEdge = True}
            else
              put st{currentV = v'}
            zeroBit
            encodeVertex v'
            zeroBit
            encodeVertex u
            encodeEdges es'
          else if v' == (currentV st) + 1 then do
            if v' == n - 2 then
              put st{currentV = v', n2HasEdge = True}
            else if v' == n - 1 then
              put st{currentV = v', n1HasEdge = True}
            else
              put st{currentV = v'}
            oneBit
            encodeVertex u
            encodeEdges es'
          else if v' == (currentV st) then do
            zeroBit
            encodeVertex u
            encodeEdges es'
          else
            fail ""
      encodeVertex :: Word64 -> StateT EncodeState (Maybe) ()
      encodeVertex u = do
        st <- get
        if k < 6 - usedBits st then
          put st{ currentByte = (shiftL (currentByte st ) k) .|. (fromIntegral u)
                , usedBits = usedBits st + k }
        else do
          let (blocks, b', i') = encodeVertex' u k [] (currentByte st) (usedBits st)
          put st { currentByte = b'
                 , usedBits = i'
                 , encoding = blocks ++ encoding st
                 }
  in 
  case evalStateT (encodeEdges es) 
                 encodeStartState
  of
    Just txt -> txt
    Nothing -> fromJust $ evalStateT 
      ( encodeEdges $
          sortBy (\(v0, u0) (v1, u1) -> compare (u0, v0) (u1, v1)) $ 
            map 
            (\(v,u) -> if v < u then (u,v) else (v,u))
            es
      ) encodeStartState


-- |Encode a single vertex.
encodeVertex' :: Word64
              -> Int
              -> [Word8]
              -> Word8
              -> Int
              -> ([Word8], Word8, Int)
encodeVertex' u k' bs b i
  | k' == 6 - i = 
    let b' = (shiftL b k') .|. (fromIntegral u)
    in (b' : bs, 0, 0 )
  | k' < 6 - i = 
    let b' = (shiftL b k') .|. (fromIntegral u)
    in (bs, b', i + k' )
  | otherwise = 
    let k'' = (k' - 6 + i)
        u'  = shiftR u k''
        u'' = shiftR (shiftL u (64 - k'')) (64 - k'')
        b'  = (shiftL b (6 - i)) .|. (fromIntegral u')
    in encodeVertex' u'' (k' - 6 + i) (b' : bs) 0 0

-- | Set next bit to one.
oneBit :: (Monad m) => StateT EncodeState m ()
oneBit  = do
  st <- get
  if usedBits st == 6 then
    put st{ currentByte = 1
          , usedBits = 1
          , encoding = (currentByte st) : encoding st}
  else
    put st{ currentByte = setBit (shiftL (currentByte st) 1) 0
          , usedBits = usedBits st + 1
          }
-- | Set next bit to zero.
zeroBit :: (Monad m) => StateT EncodeState m ()
zeroBit = do
  st <- get
  if usedBits st == 6 then
    put st{ currentByte = 0
          , usedBits = 1
          , encoding = (currentByte st) : encoding st}
  else
    put st{ currentByte = shiftL (currentByte st) 1
          , usedBits = usedBits st + 1
          }

-- |Starting state when encoding a graph.
encodeStartState :: EncodeState
encodeStartState = EncodeState
  { currentV = 0
  , currentByte = 0
  , usedBits = 0
  , n2HasEdge = False
  , n1HasEdge = False
  , encoding = []
  }

-- |State when parsing a list of edges.
data EdgeListState = EdgeListState
  { elV        :: Word64
  , elU        :: Word64
  , uMissing   :: Int
  , edges      :: [(Word64, Word64)]
  } 
  deriving (Eq, Show)

-- |Parse a list of edges.
parseEdgeList :: Word64
              -> Word64
              -> StateT B.ByteString (Either T.Text) AdjacencyList
parseEdgeList n k = do
  bs <- get
  put B.empty
  let els = B.foldl' 
              (edgeList n k)
              EdgeListState
              { elV = 0
              , elU = 0
              , uMissing = 0
              , edges = []
              }
              bs
      es = edges els
  return $ AdjacencyList
    { numberOfVertices = n
    , adjacency = A.array (0, (2 * (fromIntegral $ length es) - 1)) $ zip [0..] $ 
        ungroupByTwo $ reverse es
    }

-- |Read a single byte.
consumeByte :: Word64
            -> Word64
            -> EdgeListState
            -> Word8
            -> Int
            -> EdgeListState
consumeByte n k st' b l' = consumeByte' st' l'
  where
    consumeByte' st 0 = st
    consumeByte' st l 
      | uMissing st == 0 =
          if testBit b (l - 1) then
            if elV st == n - 1 then
              st
            else
              consumeByte' st{elV = elV st + 1, uMissing = fromIntegral k} (l - 1)
          else
            consumeByte' st{uMissing = fromIntegral k} (l - 1)
      | uMissing st <= l =
          let u' = (shiftL (elU st) (uMissing st)) .|. (fromIntegral $ shiftR (shiftL b (8 - l)) (8 - (uMissing st)))
          in 
          consumeByte'
            (if u' > elV st then
               st
               { elV = u'
               , elU = 0
               , uMissing = 0
               }
            else
               st
               { elU = 0
               , uMissing = 0
               , edges = (elV st, u') : edges st
               })
            (l - (uMissing st))
      | otherwise =
          let u' = (shiftL (elU st) l) .|. (fromIntegral (shiftR (shiftL b (8 - l)) (8 - l)))
          in 
          st
          { elU = u'
          , uMissing = (uMissing st) - fromIntegral l
          }
-- |Bytes of the list of edges.
edgeList :: Word64 -> Word64 -> EdgeListState -> Word8 -> EdgeListState
edgeList n k st b = consumeByte n k st (b - 63) 6
        
-- | Create an adjacency matrix from a list of edges.
-- Vertices need to be in the range from 0 to n-1.
fromEdgeList :: Word64 -> [(Word64, Word64)] -> AdjacencyList
fromEdgeList n es = AdjacencyList
  { numberOfVertices = n
  , adjacency = A.array (0, fromIntegral $ 2 * (length es) - 1) $ zip [0..] $
    ungroupByTwo es
  }

-- | The list of edges of a graph together with the number of vertices.
toEdgeList :: AdjacencyList -> (Word64, [(Word64, Word64)])
toEdgeList g = (numberOfVertices g, groupByTwo $ A.elems $ adjacency g)

-- | Whether two vertices are adjacent.
areAdjacent :: AdjacencyList -> Word64 -> Word64 -> Bool
areAdjacent g v u = 
  let e = if v < u then (v,u) else (u,v) in
  e `elem` (groupByTwo $ A.elems (adjacency g))