{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses      #-}

-- ----------------------------------------------------------------------------

{- |
  Module     : Holumbus.Index.CompactDocuments
  Copyright  : Copyright (C) 2007- 2014 Sebastian M. Schlatt, Timo B. Huebel, Uwe Schmidt
  License    : MIT

  Maintainer : Uwe Schmidt (uwe@fh-wedel.de)
  Stability  : experimental

  A more space efficient substitute for Holumbus.Index.Documents
-}

-- ----------------------------------------------------------------------------

module Holumbus.Index.CompactDocuments
    (
      -- * Documents type
      Documents (..)
    , CompressedDoc
    , DocMap
    , URIMap

      -- * Construction
    , emptyDocuments
    , singleton

      -- * Conversion
    , simplify
    , toDocument
    , fromDocument
    , fromDocMap
    , toDocMap
    , mkCDoc
    )
where

import           Control.DeepSeq

import           Data.Binary                    (Binary)
import qualified Data.Binary                    as B
import qualified Data.ByteString                as BS
import qualified Data.ByteString.Lazy           as BL
import qualified Data.ByteString.Short          as SS
import           Data.Maybe                     (fromJust)
import qualified Data.StringMap.Strict          as M
import           Data.Typeable

import qualified Holumbus.ByteStringCompression as BZ
import           Holumbus.Index.Common

import           Text.XML.HXT.Core

#if sizeable == 1
import           Data.Size
#endif
-- ----------------------------------------------------------------------------
--
-- Document descriptions are always shorter as before when compressed
-- so no compressBZipSmart is neccessary

compress   :: BL.ByteString -> BL.ByteString
compress   = BZ.compressBZip

decompress :: BL.ByteString -> BL.ByteString
decompress = BZ.decompressBZip

-- ----------------------------------------------------------------------------

-- | The table which is used to map a document to an artificial id and vice versa.

data Documents a        = Documents
                          { idToDoc   :: ! (DocMap a)   -- ^ A mapping from a document id to
                                                        --   the document itself.
                          , docToId   :: ! URIMap       -- ^ A space efficient mapping from
                                                        --   the URI of a document to its id.
                          , lastDocId :: ! DocId        -- ^ The last used document id.
                          }
                          deriving (Show, Typeable)

-- ----------------------------------------------------------------------------

instance (Binary a, HolIndex i) => HolDocIndex Documents a i where
    defragmentDocIndex dt ix  = (dt1, ix1)
        where
          dt1                   = editDocIds editId dt
          ix1                   = updateDocIds' editId ix
          editId i              = fromJust . lookupDocIdMap i $ idMap
          idMap                 = fromListDocIdMap . flip zip (map mkDocId [1..]) . keysDocIdMap . toMap $ dt

    unionDocIndex dt1 ix1 dt2 ix2
        | s1 == 0               = (dt2, ix2)
        | s2 == 0               = (dt1, ix1)
        | s1 < s2               = unionDocIndex dt2 ix2 dt1 ix1
        | otherwise             = (dt, ix)
        where
          dt                    = unionDocs     dt1  dt2s
          ix                    = mergeIndexes  ix1  ix2s

          dt2s                  = editDocIds    add1 dt2
          ix2s                  = updateDocIds' add1 ix2

          add1                  = addDocId disp
          max1                  = maxKeyDocIdMap . toMap $ dt1
          min2                  = minKeyDocIdMap . toMap $ dt2
          disp                  = incrDocId $ subDocId max1 min2

          s1                    = sizeDocs dt1
          s2                    = sizeDocs dt2

-- ----------------------------------------------------------------------------

instance Binary a => HolDocuments Documents a where
  sizeDocs d                    = sizeDocIdMap (idToDoc d)

  lookupById  d i               = maybe (fail "") return
                                  . fmap toDocument
                                  . lookupDocIdMap i
                                  $ idToDoc d
  lookupByURI d u               = maybe (fail "") return
                                  . M.lookup  u
                                  $ docToId d

  -- this is a sufficient test, but if the doc ids don't form an intervall
  -- it may be too restrictive

  disjointDocs dt1 dt2
      | nullDocs dt1
        ||
        nullDocs dt2            = True
      | otherwise               = disjoint ( minKeyDocIdMap . idToDoc $ dt1
                                           , maxKeyDocIdMap . idToDoc $ dt1
                                           )
                                           ( minKeyDocIdMap . idToDoc $ dt2
                                           , maxKeyDocIdMap . idToDoc $ dt2
                                           )
      where
        disjoint p1@(x1, y1) p2@(x2, _y2)
            | x1 <= x2          = y1 < x2
            | otherwise         = disjoint p2 p1

  unionDocs dt1 dt2
      | disjointDocs dt1 dt2    = Documents
                                  { idToDoc     = unionDocIdMap (idToDoc dt1) (idToDoc dt2)
                                  , docToId     = M.union  (docToId dt1) (docToId dt2)
                                  , lastDocId   = lastDocId dt1 `max` lastDocId dt2
                                  }
      | otherwise               = error $
                                  "unionDocs: doctables are not disjoint: " ++
                                  show (didIntervall dt1) ++ ", " ++ show (didIntervall dt2)
      where
       didIntervall dt          = ( minKeyDocIdMap . idToDoc $ dt
                                  , maxKeyDocIdMap . idToDoc $ dt
                                  )

  makeEmpty _                   = emptyDocuments

  insertDoc ds d                = maybe reallyInsert (\oldId -> (oldId, ds)) (lookupByURI ds (uri d))
    where
    d'                          = fromDocument d
    reallyInsert                = newDocs `seq` (newId, newDocs)
      where
        newIdToDoc              = insertDocIdMap newId d' (idToDoc ds)
        newDocToId              = M.insert (uri d) newId  (docToId ds)
        newId                   = incrDocId (lastDocId ds)
        newDocs                 = Documents newIdToDoc newDocToId newId

  updateDoc ds i d              = ds
                                  { idToDoc = insertDocIdMap i d' (idToDoc ds)
                                  , docToId = M.insert (uri d) i (docToId (removeById ds i))
                                  }
      where
        d'                      = fromDocument d

  removeById ds d               = maybe ds reallyRemove (lookupById ds d)
    where
    reallyRemove (Document _ u _)
                                = Documents
                                  (deleteDocIdMap d (idToDoc ds))
                                  (M.delete u (docToId ds))
                                  (lastDocId ds)

  updateDocuments f d           = Documents updated (idToDoc2docToId updated) (lastId updated)
    where
    updated                     = mapDocIdMap (mapDocument f) (idToDoc d)

  filterDocuments p d           = Documents filtered (idToDoc2docToId filtered) (lastId filtered)
    where
    filtered                    = filterDocIdMap (p . toDocument) (idToDoc d)

  fromMap itd'                  = Documents itd (idToDoc2docToId itd) (lastId itd)
    where
    itd                         = toDocMap itd'

  toMap                         = fromDocMap . idToDoc

  editDocIds f d                = Documents
                                  { idToDoc     = newIdToDoc
                                  , docToId     = M.map f $ docToId d
                                  , lastDocId   = lastId newIdToDoc
                                  }
    where
    newIdToDoc                  = foldWithKeyDocIdMap (insertDocIdMap . f) emptyDocIdMap
                                  $ idToDoc d

-- ------------------------------------------------------------

-- Ignoring last document id when testing for equality

instance Eq a =>                Eq (Documents a)
    where
    (==) (Documents i2da d2ia _) (Documents i2db d2ib _)
                                = (i2da == i2db)
                                  &&
                                  (d2ia == d2ib)

-- ----------------------------------------------------------------------------

instance NFData a =>            NFData (Documents a)
    where
    rnf (Documents i2d d2i lid) = rnf i2d `seq` rnf d2i `seq` rnf lid

-- ----------------------------------------------------------------------------

instance (Binary a, XmlPickler a) =>
                                XmlPickler (Documents a)
    where
    xpickle                     = xpElem "documents" $
                                  xpWrap convertDoctable $
                                  xpWrap (fromListDocIdMap, toListDocIdMap) $
                                  xpList xpDocumentWithId
        where
        convertDoctable         = ( \ itd -> Documents itd (idToDoc2docToId itd) (lastId itd)
                                  , \ (Documents itd _ _) -> itd
                                  )
        xpDocumentWithId        = xpElem "doc" $
                                  xpPair (xpAttr "id" xpDocId) xpickle

-- ----------------------------------------------------------------------------

instance Binary a =>            Binary (Documents a)
    where
    put (Documents i2d _ lid)   = B.put lid >> B.put i2d
    get                         = do lid <- B.get
                                     i2d <- B.get
                                     return (Documents i2d (idToDoc2docToId i2d) lid)

-- ------------------------------------------------------------

-- | Create an empty table.

emptyDocuments                  :: Documents a
emptyDocuments                  = Documents emptyDocIdMap M.empty nullDocId

-- | Create a document table containing a single document.

singleton                       :: (Binary a) => Document a -> Documents a
singleton d                     = Documents
                                  (singletonDocIdMap firstDocId d')
                                  (M.singleton (uri d) firstDocId)
                                  firstDocId
    where
      d'                        = fromDocument d

-- | Simplify a document table by transforming the custom information into a string.

simplify                        :: (Binary a, Show a) => Documents a -> Documents String
simplify dt                     = Documents (simple (idToDoc dt)) (docToId dt) (lastDocId dt)
  where
  simple i2d                    = mapDocIdMap
                                  ( fromDocument
                                    . (\d -> Document (title d) (uri d) (maybe Nothing (Just . show) (custom d)))
                                    . toDocument
                                  ) i2d

-- ------------------------------------------------------------
--
-- The CompressedDoc is a candidate for a BS.ShortByteString available with bytestring 0.10.4,
-- then 5 machine words can be saved per value

newtype CompressedDoc a         = CDoc { unSDoc :: SS.ShortByteString }
                                  deriving (Eq, Show, NFData, Typeable)

mkCDoc                          :: BS.ByteString -> CompressedDoc a
mkCDoc s                        = CDoc $!! SS.toShort s

unCDoc                          :: CompressedDoc a -> BS.ByteString
unCDoc                          = SS.fromShort . unSDoc

-- in mkCDoc the ByteString is physically copied
-- to avoid sharing data with other possibly very large strings
-- this can occur e.g in the Binary instance with get

instance (Binary a, XmlPickler a) => XmlPickler (CompressedDoc a)
    where
    xpickle                     = xpWrap (fromDocument , toDocument) $
                                  xpickle

instance Binary a =>            Binary (CompressedDoc a)
    where
    put                         = B.put . unCDoc
    get                         = B.get >>= return . mkCDoc . BS.copy

-- to avoid sharing the data with with the input, the ByteString is physically copied
-- before return. This should be the single place where sharing is introduced,
-- else the copy must be moved to mkCDoc

toDocument      :: (Binary a) => CompressedDoc a -> Document a
toDocument      = B.decode . decompress . BL.fromStrict . unCDoc

fromDocument    :: (Binary a) => Document a -> CompressedDoc a
fromDocument    = mkCDoc . BL.toStrict . compress . B.encode

mapDocument     :: (Binary a) => (Document a -> Document a) -> CompressedDoc a -> CompressedDoc a
mapDocument f   = fromDocument . f . toDocument

-- ------------------------------------------------------------

type URIMap     = M.StringMap DocId

type DocMap a   = DocIdMap (CompressedDoc a)

toDocMap        :: (Binary a) => DocIdMap (Document a) -> DocMap a
toDocMap        = mapDocIdMap fromDocument

fromDocMap      :: (Binary a) => DocMap a -> DocIdMap (Document a)
fromDocMap      = mapDocIdMap toDocument

-- | Construct the inverse map from the original map.

idToDoc2docToId :: Binary a => DocMap a -> URIMap
idToDoc2docToId = foldWithKeyDocIdMap
                  (\i d r -> M.insert (uri . toDocument $ d) i r)
                  M.empty

-- | Query the 'idToDoc' part of the document table for the last id.

lastId          :: DocMap a -> DocId
lastId          = maxKeyDocIdMap

-- ------------------------------------------------------------
#if sizeable == 1

instance (Sizeable a, Typeable a) => Sizeable (Documents a) where
    dataOf _x                   = 3 .*. dataOfPtr
    statsOf x@(Documents i d l) = mkStats x <> statsOf i <> statsOf d <> statsOf l

instance (Sizeable a, Typeable a) => Sizeable (CompressedDoc a) where
    dataOf                      = dataOf  . unSDoc
    bytesOf                     = bytesOf . unSDoc
    statsOf                     = statsOf . unSDoc

#endif
-- ------------------------------------------------------------
