    module Hoogle.DataBase.TypeSearch.Graphs where

    import Hoogle.DataBase.TypeSearch.Graph
    import Hoogle.DataBase.TypeSearch.Binding
    import Hoogle.DataBase.TypeSearch.Result
    import Hoogle.DataBase.Instances
    import Hoogle.DataBase.Aliases
    import Hoogle.DataBase.TypeSearch.TypeScore
    import Hoogle.Item.All
    import Hoogle.TypeSig.All

    import Data.Binary.Defer
    import Data.Binary.Defer.Index
    import qualified Data.IntMap as IntMap
    import qualified Data.Heap as Heap
    import General.Code
    import Control.Monad.State

    -- for resGraph, the associated ArgPos is the arity of the function

    data Graphs = Graphs
        {entryInfo :: Index EntryInfo
        ,argGraph :: Graph -- the arguments
        ,resGraph :: Graph -- the results
        }

    instance Show Graphs where
        show (Graphs a b c) = "== Arguments ==\n\n" ++ show b ++
                              "\n== Results ==\n\n" ++ show c

    instance BinaryDefer Graphs where
        put (Graphs a b c) = put3 a b c
        get = do
            res@(Graphs a b c) <- get3 Graphs
            getDeferPut a
            return res

    ---------------------------------------------------------------------
    -- GRAPHS CONSTRUCTION

    newGraphs :: Aliases -> Instances -> [(Link Entry, TypeSig)] -> Graphs
    newGraphs as is xs = Graphs (newIndex $ map snd entries) argGraph resGraph
        where
            entries = [ (t2, e2{entryInfoEntries = sortOn linkKey $ map snd ys})
                      | ys@(((t2,e2),_):_) <- sortGroupFst $ map (\(e,t) -> (normType as is t, e)) xs]

            argGraph = newGraph (concat args)
            resGraph = newGraph res

            (args,res) = unzip
                [ initLast $ zipWith (\i t -> (lnk, i, t)) [0..] $ fromTFun t
                | (i, (t, e)) <- zip [0..] entries, let lnk = newLink i e]

    normType :: Aliases -> Instances -> TypeSig -> (Type, EntryInfo)
    normType as is t = (t3, EntryInfo [] (length (fromTFun t3) - 1) c2 a)
        where TypeSimp c2 t2 = normInstances is t
              (a,t3) = normAliases as t2

    ---------------------------------------------------------------------
    -- GRAPHS SEARCHING

    -- sorted by TypeScore
    graphsSearch :: Aliases -> Instances -> Graphs -> TypeSig -> [ResultReal]
    graphsSearch as is gs t = resultsCombine is query ans
        where
            ans = mergesBy (compare `on` resultArgBind . snd) $
                  f Nothing (resGraph gs) res :
                  zipWith (\i -> f (Just i) (argGraph gs)) [0..] args

            f a g = map ((,) a) . graphSearch g
            (args,res) = initLast $ fromTFun ts
            (ts,query) = normType as is t

    data S = S
        {infos :: IntMap.IntMap (Maybe ResultAll) -- Int = Link EntryInfo
        ,pending :: Heap.Heap Int Result
        ,todo :: [(Maybe ArgPos, ResultArg)]
        ,instances :: Instances
        ,query :: EntryInfo
        }

    resultsCombine :: Instances -> EntryInfo -> [(Maybe ArgPos, ResultArg)] -> [ResultReal]
    resultsCombine is query xs = flattenResults $ evalState delResult s0
        where s0 = S IntMap.empty Heap.empty xs is query

    -- Heap -> answer
    delResult :: State S [Result]
    delResult = do
        pending <- gets pending
        todo <- gets todo
        case todo of
            [] -> concatMapM f $ Heap.elems pending
            t:odo -> do
                let (res,hp) = Heap.popWhile (costBinding $ resultArgBind $ snd t) pending
                modify $ \s -> s{todo=odo, pending=hp}
                ans1 <- concatMapM f res
                uncurry addResult t
                ans2 <- delResult
                return $ ans1 ++ ans2
        where
            f r = do
                infos <- gets infos
                (Just res,infos) <- return $ IntMap.updateLookupWithKey
                    (\_ _ -> Just Nothing) (linkKey $ fst3 r) infos
                if isNothing res then return [] else do
                    modify $ \s -> s{infos=infos}
                    return [r]

    -- todo -> heap/info
    addResult :: Maybe ArgPos -> ResultArg -> State S ()
    addResult arg val = do
        let entId = linkKey $ resultArgEntry val
        infs <- gets infos
        is <- gets instances
        query <- gets query
        let def = newResultAll query (resultArgEntry val)
        case IntMap.lookup entId infs of
            Just Nothing -> return ()
            Nothing | isNothing def -> modify $ \s -> s{infos = IntMap.insert entId Nothing $ infos s}
            x -> do
                let inf = fromJust $ fromMaybe def x
                (inf,res) <- return $ addResultAll is query (arg,val) inf
                res <- return $ map (costTypeScore . thd3 &&& id) res
                modify $ \s -> s
                    {infos = IntMap.insert entId (Just inf) $ infos s
                    ,pending = Heap.insertList res (pending s)
                    }