NgramsTree.hs 2.39 KB
Newer Older
1
{-|
2
Module      : Gargantext.API.Ngrams.NgramsTree
3 4 5 6 7 8 9 10 11 12 13
Description : Tree of Ngrams
Copyright   : (c) CNRS, 2017-Present
License     : AGPL + CECILL v3
Maintainer  : team@gargantext.org
Stability   : experimental
Portability : POSIX

-}

{-# LANGUAGE TemplateHaskell   #-}

14
module Gargantext.API.Ngrams.NgramsTree
15 16 17
  where

import Data.Aeson.TH (deriveJSON)
18 19
import Data.HashMap.Strict (HashMap)
import Data.Maybe (catMaybes)
20
import Data.Set (Set)
21
import Data.Swagger
22 23
import Data.Text (Text)
import Data.Tree
24
import GHC.Generics (Generic)
25
import Gargantext.API.Ngrams.Types
26 27
import Gargantext.Core.Types (ListType(..), NodeId)
import Gargantext.Core.Utils.Prefix (unPrefix, unPrefixSwagger)
28 29 30 31 32
import Gargantext.Prelude
import Test.QuickCheck
import qualified Data.HashMap.Strict as HashMap
import qualified Data.List as List
import qualified Data.Set as Set
33

34 35 36
type Children = Text
type Root = Text

37 38 39 40 41
data NgramsTree = NgramsTree { mt_label :: Text
                             , mt_value :: Double
                             , mt_children :: [NgramsTree]
                             }
    deriving (Generic, Show)
42

43 44
toNgramsTree :: Tree (NgramsTerm,Double) -> NgramsTree
toNgramsTree (Node (NgramsTerm l,v) xs) = NgramsTree l v (map toNgramsTree xs)
45

46
deriveJSON (unPrefix "mt_") ''NgramsTree
47

48
instance ToSchema NgramsTree where
49
  declareNamedSchema = genericDeclareNamedSchema (unPrefixSwagger "mt_")
50
instance Arbitrary NgramsTree
51
  where
52
    arbitrary = NgramsTree <$> arbitrary <*> arbitrary <*> arbitrary
53

54 55 56 57
toTree :: ListType
       -> HashMap NgramsTerm (Set NodeId)
       -> HashMap NgramsTerm NgramsRepoElement
       -> [NgramsTree]
58
toTree lt vs m = map toNgramsTree $ unfoldForest buildNode roots
59
  where
60
    buildNode r = maybe ((r, value r),[])
61 62
                        (\x -> ((r, value r), mSetToList $ _nre_children x))
                        (HashMap.lookup r m)
63

64
    value l = maybe 0 (fromIntegral . Set.size) $ HashMap.lookup l vs
65

66
    rootsCandidates :: [NgramsTerm]
67 68
    rootsCandidates = catMaybes
                    $ List.nub
69
                    $ map (\(c, c') -> case _nre_root c' of
70
                                       Nothing -> Just c
71 72
                                       _ -> _nre_root c'
                          ) (HashMap.toList m)
73

74 75 76
    roots = map fst
          $ filter (\(_,l) -> l == lt)
          $ catMaybes
77 78
          $ map (\c -> (,) <$> Just c <*> (_nre_list <$> HashMap.lookup c m))
          $ rootsCandidates