find loops in graph

{-# LANGUAGE DisambiguateRecordFields,ViewPatterns #-}
import Data.List
import Control.Monad
import Control.Monad.State
import Debug.Trace
import qualified Data.IntMap as IntMap
{-

  Input : List of int
          Predicate : a->b->Bool

  Output  : type G a = (a, [G a]) : result = [G a] such that
    for each m in result (snd m) == filter (fst m) result

  Additional question : How to "show" cyclic structures . Deriving Show will print indefinetly
  show structure should start marking already visited nodes . This internal info may possibly be not shown

-}

-- Need to build instance of show , since this is cyclic structure
data LoopDetect  = LD {  visited :: [Int] , isLoop :: Bool } deriving (Show)
data (Show a,Eq a) => G a = G {val :: (a, [G a]) , idx :: Int , foundLoops :: [Int] } deriving (Eq)
type Pred a = a -> a -> Bool

type StMon a = State LoopDetect (G a)
type St a = ((G a),LoopDetect)

-- Build list , by tying the knot
-- There is a problem of printing cyclic structure
-- List (first argument) should be unique

func :: (Show a,Eq a) => [a] -> Pred a -> [G a]
func lst pred = res_lst_with_loops where
  tmp_lst =  zip lst j
  ap2' f x y = f x y
  ap3' f x y z = f x y z

  -- Build lazy infinite list of LD
  initLDlst = zipWith3 ap2' (repeat LD)  (repeat []) (repeat False)
  res_lst = zipWith4 ap3' (repeat G) tmp_lst [0..] (repeat [])-- Convert tuple to G datatype

  -- k list of func (a->Bool)
  -- where each function has it's first argument injected
  j = helper [] res_lst lst




  helper acc _ [] = reverse acc -- Termination condition
  helper acc glist lst@(x:xs) = helper (subList:acc) glist xs where
    subList = filter ((pred x) . fst . val) glist
    -- Here update list of loops
    --stateUpdate = detectLoops x -- Will update visited


  --res_lst_with_loops = seq res_lst (findLoops res_lst)
  --res_lst_with_loops = findLoops res_lst
  res_lst_with_loops = trace ((++) "Initial res_lst = " $! seq res_lst show res_lst ) $! findLoops res_lst

-- End of func :: [a] -> Pred a -> [G a]






--currNodeUpd :: (Show a , Eq a) => (G a) -> (LoopDetect -> (G a,LoopDetect))
currNodeUpd :: (Show a , Eq a) => (G a) -> State LoopDetect (G a)
currNodeUpd entry = state resFunc where
  _idx = idx entry
  resFunc st' = (entry,newSt) where
    --newSt = LD (_idx : visited st') (loops st') -- First version --> Don't remove state of loops
    -- Adding one node can only close one loop . -- Looks reasonable , but need to prove
    -- But every node can be part of multiple loops . Thus foundLoops field is a list
    newSt = LD (_idx : visited st') (isLoop st') -- Loops should be empty . Possible refactoring --> Convert to boolean

nextNodeJoinVis :: (Show a , Eq a) => LoopDetect -> (G a) -> State LoopDetect (G a)
nextNodeJoinVis currSt entry = state resFunc where
  _idx = idx entry
  resFunc st' = (entry,newSt) where
    --newSt = LD (_idx : visited st') (loops st') -- First version --> Don't remove state of loops
    -- Adding one node can only close one loop . -- Looks reasonable , but need to prove
    -- But every node can be part of multiple loops . Thus foundLoops field is a list
    newSt = LD (union  (visited st') (visited currSt)) (isLoop st')
-- Update state of the next node --> Check for loop
--
--nextNodeUpd :: (Show a , Eq a) => (G a) -> (LoopDetect -> (G a,LoopDetect))
nextNodeUpd :: (Show a , Eq a) => (G a) -> State LoopDetect (G a)
nextNodeUpd entry = state resFunc where
  _idx = idx entry
  resFunc st' = (newEntry,newSt) where
    _vis = visited st'
    _loopFound = _idx `elem` _vis
    _newLoops = if (_loopFound) then (_idx:foundLoops entry) else foundLoops entry
    newEntry = G { val = val entry , idx = idx entry , foundLoops = _newLoops }
    newSt = LD (_vis) (_idx `elem` _vis)

-- Problem :
-- Given no defined root --> Not a tree . I may traverse nodes twice
-- Once when iterating through list of nodes
-- and second time , when walking through links
-- Given I am 'tying the knot' I should not return for second pass

-- Assumption :
--  Every node should be visited at least once
--  even in presence of loops . Thus we may pass all nodes and mark as visited





findLoops :: (Show a,Eq a) => [G a] -> [G a]
findLoops lst = newLst where
  -- Build auxillary map to remove already visited

  --nodeMap = IntMap.fromList $ zip (map idx lst) lst -- Build map of unvisited nodes
  nodeMap = trace "Print nodeMap in findLoops " $! IntMap.fromList $! seq lst zip (map idx lst) lst -- Build map of unvisited nodes
  initSTvals = trace "Init st vals = " $! IntMap.fromList $! zip (map idx $! lst) (map (\x -> (x,(LD [] False)) ) $! lst) -- Build default state map

  --newLst = map (fst . snd ) $! IntMap.toList $! snd $! seq nodeMap fl' nodeMap $! initSTvals -- Map of State monad to traverse . Key is node idx
  newLst = map (fst . snd ) $! IntMap.toList $! snd $! (seq nodeMap  fl' $! trace "Node map is " $! nodeMap) $! initSTvals -- Map of State monad to traverse . Key is node idx



  --fl' :: (Show a,Eq a) => IntMap.IntMap (G a) -> [G a] -> IntMap.IntMap (State LoopDetect (G a)) -> IntMap.IntMap (G a)
  -- fl' ::
  --fl' :: (Show a,Eq a) => IntMap.IntMap (G a) -> IntMap.IntMap (St a) -> (IntMap.IntMap (G a),IntMap.IntMap (St a))
  fl' :: (Show a,Eq a) => IntMap.IntMap (G a) -> IntMap.IntMap (St a) -> (IntMap.IntMap (G a),IntMap.IntMap (St a))

  -- Termination case : All nodes are already traversed
  fl' not_visited  stateMap | IntMap.null not_visited = result
                            | otherwise = (newStateMap,new_not_visited) where
      -- Run state
      initSt = LD [] False
      -- Apply runState for a map and get last value
      -- For each entry result is already visited
      result = (IntMap.empty,stateMap)
    -- General case : At least one node to traverse

      pair = head $ IntMap.toList not_visited
      key = fst pair
      --x = snd pair
      x = trace "Current x = " $ seq pair snd pair

      -- How to build state transformer
      _idx = idx x
      _chLst = (snd . val) x

      (_newStateMap,_new_not_visited) = fl'' x stateMap not_visited -- Process single node
      (newStateMap,new_not_visited) = fl' _new_not_visited _newStateMap -- Process remainder




      -- Process children . Perform depth next step
      -- Check state of children


   -- Update state of a sub graph , taking as a first node some node
  fl'' :: (Eq a, Show a) => (G a) -> IntMap.IntMap (St a) -> IntMap.IntMap (G a) -> (IntMap.IntMap (St a),IntMap.IntMap (G a))

   -- Update due to a single state
  fl'' x stateMap (IntMap.null -> True) = (stateMap,IntMap.empty) -- Process empty not visited map
  fl'' x stateMap nonVisited = (newStateMap,newNonVisited) where

      --_idx = idx x
      _idx = trace "_idx = " $ seq x idx x
      _children = map idx  ((snd . val) x)
      (currVal,currSt) =  (IntMap.!) stateMap _idx

      -- Set update state as visited
      (currVal_1,currSt_1) =  (flip runState currSt $ (return currVal) >>= currNodeUpd)

      -- Check every child to be a loop
      -- first I apply runState on state of child + and value of node . I apply it for

      --(currVal_2,currSt_2) = flip runState currSt_2 $ foldl' gg (return currVal_2) $ map (\x -> (IntMap.!) stateMap x) $ _children where

      (currVal_2,currSt_2) = flip runState currSt_1 $ foldl' gg (return currVal_1)  _children where
          gg :: (Eq a,Show a) => StMon a -> Int -> StMon a
          gg _jj _idx1 = result where
            (_,chSt) =  (IntMap.!) stateMap _idx1
            result = return $ fst $ flip runState chSt $ _jj >>= nextNodeUpd

      -- notLoopChildrenIdx =  filter (\j -> notElem j $ foundLoops currVal_2 )   _children
      notLoopChildrenIdx =  filter (\j -> notElem j $ foundLoops currVal_2 ) $ seq _children $ trace "\nPrint _children\n"  _children
      notLoopChildren =  map (fst . ((IntMap.!) _newStateMap_2)) $ trace "\nPrint not loop index " notLoopChildrenIdx


       -- Update StateMap with current entry
      _newStateMap_1 = IntMap.adjust (\_ -> (currVal_2,currSt_2)) _idx stateMap


      -- Update state of each non loop child with visited node
      _newStateMap_2 = foldl' hh _newStateMap_1 notLoopChildrenIdx where
        --hh :: IntMap.IntMap (St a) -> Int -> IntMap.IntMap (St a)
        hh _map _idx = _newMap where
          (prevVal,prevSt) = _map IntMap.! _idx  -- Fetch previous value
          (tmpVal,tmpSt) = flip runState prevSt $ return prevVal >>= (nextNodeJoinVis currSt_2)
          _newMap = IntMap.adjust (\_ -> (tmpVal,tmpSt)) _idx stateMap






       -- Remove current entry from not visited
      _newNonVisited = IntMap.delete _idx nonVisited

      -- Apply this only to children not creating loops
      (newStateMap,newNonVisited) = foldl' kk ((trace ("Print new state map" ++ show _newStateMap_2 ++ "\n\n") _newStateMap_2),_newNonVisited) notLoopChildren where
        kk y x  = fl'' x (fst y) (snd y)

       -- Recursive call of fl'' to a children

-- Print graph
printGraph :: (Show a,Eq a) => G a -> String
printGraph entry
  | (snd . val) entry == [] =
    prefix ++ suffix
  | otherwise = result where
    result = prefix ++ body ++ suffix
    index_ :: Int
    index_ = idx entry
    prefix = "G {" ++ "val = " ++ "(" ++  (show $ (fst . val) entry) ++ "," ++ "["
    suffix = "])," ++ "idx = " ++ (show $ index_ ) ++ ",foundLoops = " ++ (show $ foundLoops entry) ++  " }"


    -- For each step of traversal
    {-
     - Loops detection is created once
     - Thus , show should operate on ready foundLoops field
     - for now check for each element of sons' list whether it's
     - already visited . Then body should reflect G a and word LOOP = <loop entry index>
     -}

    -- Print all children
--    chLst = (snd . val) entry
--    body = foldl' kk "" chLst where
--      kk acc ch | ((idx ch) `elem` (foundLoops entry) ) == True  = acc ++ " Loop = " ++ show (idx ch) ++ " "
--                | otherwise = printGraph ch
    -- DEBUG_444
    body = " children "
    -- END DEBUG_444


my_lst = [1,2]
_pred = \x y -> x * 1 /= y
test1 = func my_lst _pred

my_lst2 = [1,2,3]
test2 = func my_lst2 _pred

instance (Show a,Eq a) => Show (G a) where
  show = printGraph