{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
module Hdoc.DTC.Analyze.Collect where

import Control.Applicative (Applicative(..), liftA2)
import Control.Monad
import Data.Bool
import Data.Default.Class (Default(..))
import Data.Either (Either(..))
import Data.Foldable (Foldable(..), any)
import Data.Function (($), (.))
import Data.Functor ((<$>), (<$))
import Data.Map.Strict (Map)
import Data.Maybe (Maybe(..), fromMaybe)
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
import Data.Sequence (Seq)
import Data.TreeSeq.Strict (Tree(..))
import Text.Show (Show(..))
import qualified Control.Monad.Trans.Reader as R
import qualified Data.HashMap.Strict as HM
import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Text.Lazy as TL
import qualified Data.Tree as Tree
import qualified Data.TreeSeq.Strict as TS
import qualified Data.TreeMap.Strict as TM
import qualified Majority.Judgment as MJ

import qualified Hdoc.TCT.Cell as TCT
import Hdoc.DTC.Document as DTC
import qualified Hdoc.XML as XML

-- * Type 'Reader'
newtype Reader = Reader
 { reader_section :: Either Head Section
 }
instance Default Reader where
	def = Reader
	 { reader_section = Left def
	 }

-- * Type 'All'
data All = All
 { all_figure    :: HM.HashMap TL.Text (Map XML.Pos (Maybe Title))
 , all_grades    :: HM.HashMap Ident (Seq [Grade])
 , all_index     :: HM.HashMap XML.Pos Terms
 , all_irefs     :: TM.TreeMap Word (Seq (Either Head Section))
 , all_judges    :: HM.HashMap Ident (Seq Judges)
 , all_judgment  :: HM.HashMap Judgment [Tree.Tree (Maybe MJ.Share, [Choice])]
 , all_notes     :: TS.Trees (Seq [Para])
 , all_reference :: HM.HashMap Ident (Seq Reference)
 , all_rrefs     :: HM.HashMap Ident (Seq ((TCT.Location, XML.Pos), Either Head Section))
 , all_section   :: HM.HashMap Title (Seq (Either Head Section))
 , all_tag       :: HM.HashMap Title (Seq (TCT.Location, XML.Pos))
 } deriving (Show)
instance Default All where
	def = All
	 { all_figure    = def
	 , all_grades    = def
	 , all_index     = def
	 , all_irefs     = TM.empty
	 , all_judges    = def
	 , all_judgment  = def
	 , all_notes     = def
	 , all_reference = def
	 , all_rrefs     = def
	 , all_section   = def
	 , all_tag       = def
	 }
instance Semigroup All where
	x<>y = All
	 { all_figure    = hm_union all_figure
	 , all_grades    = hm_union all_grades
	 , all_index     = hm_union all_index
	 , all_irefs     = tm_union all_irefs
	 , all_judges    = hm_union all_judges
	 , all_judgment  = hm_union all_judgment
	 , all_notes     = ts_union (all_notes x) (all_notes y)
	 , all_reference = hm_union all_reference
	 , all_rrefs     = hm_union all_rrefs
	 , all_section   = hm_union all_section
	 , all_tag       = hm_union all_tag
	 } where
		hm_union get = HM.unionWith (<>) (get x) (get y)
		tm_union get = TM.union     (<>) (get x) (get y)
		ts_union :: TS.Trees (Seq [Para]) -> TS.Trees (Seq [Para]) -> TS.Trees (Seq [Para])
		ts_union sx sy = lx <> Seq.singleton union <> ry
			where
			filter   = null . TS.subTrees
			union    = TS.tree0 $ join $ TS.unTree <$> (rx <> ly)
			(rx, lx) = Seq.spanr filter sx
			(ly, ry) = Seq.spanl filter sy
instance Monoid All where
	mempty  = def
	mappend = (<>)
instance Semigroup (R.Reader Reader All) where
	(<>) = liftA2 (<>)
instance Monoid (R.Reader Reader All) where
	mempty  = pure def
	mappend = (<>)

-- * Class 'Collect'
class Collect a where
	collect :: a -> R.Reader Reader All
instance Collect Document where
	collect Document{..} =
		R.local (\ro -> ro{reader_section = Left head}) $
			return def
			 { all_section =
				HM.fromListWith (<>) $
				(\t -> (t, pure $ Left head))
				 <$> about_titles (head_about head)
			 } <>
			(<$> collect body) (\ro -> ro
			 { all_judgment =
				choicesBySectionByJudgment HM.empty $
				TS.Tree (choicesByJudgment $ head_judgments head) $
				choicesByJudgmentBySection body
			 })
instance Collect (TS.Trees BodyNode) where
	collect = foldMap $ \(TS.Tree b bs) ->
		case b of
		 BodyBlock blk ->
			collect blk
		 BodySection section@Section{..} ->
			R.local (\ro -> ro{reader_section = Right section}) $
				((\all -> all{all_notes = pure $ TS.Tree Seq.empty $ all_notes all}) <$>) $
				return def
				 { all_section =
					HM.fromListWith (<>) $
					(\(Alias t) -> (t, pure $ Right section))
					 <$> (Alias section_title : section_aliases)
				 } <>
				collect section_title <>
				collect bs
instance Collect Block where
	collect = \case
	 BlockPara p     -> collect p
	 BlockBreak{}    -> pure def
	 BlockToC{}      -> pure def
	 BlockToF{}      -> pure def
	 BlockAside{..}  -> foldMap collect blocks
	 BlockIndex{..}  -> return def{all_index = HM.singleton posXML terms}
	 BlockFigure{..} ->
		return def
		 { all_figure = HM.singleton type_ (Map.singleton posXML mayTitle)
		 } <> foldMap collect paras
	 BlockReferences{..} ->
		return def
		 { all_reference=
			HM.fromListWith (<>) $
			(<$> refs) $
			\ref@DTC.Reference{..} -> (reference_id, pure ref)
		 } <> foldMap collect refs
	 BlockGrades{attrs=CommonAttrs{id}, ..} ->
		return def{all_grades = HM.singleton (fromMaybe "" id) $ pure scale}
	 BlockJudges judges@Judges{judges_attrs=CommonAttrs{id}, ..} ->
		return def{all_judges = HM.singleton (fromMaybe "" id) $ pure judges}
instance Collect Para where
	collect = \case
	 ParaItem item -> collect item
	 ParaItems{..} -> foldMap collect items
instance Collect ParaItem where
	collect = \case
	 ParaPlain p    -> collect p
	 ParaArtwork{}  -> return def
	 ParaQuote{..}  -> foldMap collect paras
	 ParaComment{}  -> return def
	 ParaOL items   -> foldMap collect items
	 ParaUL items   -> foldMap (foldMap collect) items
	 ParaJudgment{} -> return def
instance Collect ListItem where
	collect ListItem{..} = foldMap collect paras
instance Collect Title where
	collect (Title t) = collect t
instance Collect Plain where
	collect = foldMap collect
instance Collect (Tree PlainNode) where
	collect (Tree n ts) =
		case n of
		 PlainBreak    -> return def
		 PlainText{}   -> return def
		 PlainGroup    -> collect ts
		 PlainB        -> collect ts
		 PlainCode     -> collect ts
		 PlainDel      -> collect ts
		 PlainI        -> collect ts
		 PlainSpan{}   -> collect ts
		 PlainSub      -> collect ts
		 PlainSup      -> collect ts
		 PlainSC       -> collect ts
		 PlainU        -> collect ts
		 PlainNote{..} ->
			return def
			 { all_notes = pure $ TS.tree0 $ pure note_paras
			 } <> foldMap collect note_paras
		 PlainQ        -> collect ts
		 PlainEref{}   -> collect ts
		 PlainIref{..} -> do
			Reader{..} <- R.ask
			case pathFromWords iref_term of
			 Nothing -> collect ts
			 Just path ->
				return def
				 { all_irefs = TM.singleton path $ pure reader_section
				 } <> collect ts
		 PlainTag{..} -> return def
			 { all_tag = HM.singleton (Title ts) $ pure (tag_locTCT, tag_posXML) }
		 PlainRref{..} -> do
			Reader{..} <- R.ask
			return def
			 { all_rrefs = HM.singleton rref_to $ pure ((rref_locTCT, rref_posXML), reader_section)
			 } <> collect ts
instance Collect Reference where
	collect Reference{..} =
		collect reference_about
instance Collect About where
	collect About{..} =
		foldMap collect about_titles


choicesByJudgment :: [Judgment] -> HM.HashMap Judgment (Maybe MJ.Share, [Choice])
choicesByJudgment js =
	HM.fromList $ (<$> js) $ \j@Judgment{..} ->
		(j,(judgment_importance, judgment_choices))

choicesByJudgmentBySection :: Body -> TS.Trees (HM.HashMap Judgment (Maybe MJ.Share, [Choice]))
choicesByJudgmentBySection bod = bod >>= \(Tree b bs) ->
	case b of
	 BodyBlock{} -> mempty
	 BodySection Section{..} ->
		pure $
			let choicesJ = choicesByJudgment section_judgments in
			Tree choicesJ $
				-- NOTE: if the 'BodySection' has a child which
				-- is not a 'BodySection' itself, then add "phantom" 'Judgment's
				-- which will inherit from this 'BodySection'.
				-- This enables judges to express something on material not in a sub 'BodySection'.
				let childrenBlocksJudgments =
					if (`any`bs) $ \case
						 Tree BodyBlock{} _ -> True
						 _ -> False
					then Seq.singleton $ Tree ((Nothing,[]) <$ choicesJ) Seq.empty
					else Seq.empty in
				childrenBlocksJudgments <>
				choicesByJudgmentBySection bs

choicesBySectionByJudgment :: -- TODO: see if this can be done using Reader and collect
 HM.HashMap Judgment [Tree.Tree (Maybe MJ.Share, [Choice])] ->
 TS.Tree (HM.HashMap Judgment (Maybe MJ.Share, [Choice])) ->
 HM.HashMap Judgment [Tree.Tree (Maybe MJ.Share, [Choice])]
choicesBySectionByJudgment inh (TS.Tree selfJ childrenJS) =
	HM.unionWith
	 (\selfS childrenS ->
		(<$> selfS) $ \(Tree.Node choices old) ->
			Tree.Node choices (old<>childrenS))
	 (selfSJ <> inh)
	 childrenSJ
	where
	selfSJ = (\cs -> [Tree.Node cs []]) <$> selfJ
	childrenSJ =
		foldl'
		 (\accJ childJ ->
			HM.unionWith (<>) accJ $
			choicesBySectionByJudgment
			 (([Tree.Node (Nothing,[]) []] <$ selfJ) <> inh)
			 childJ
		 )
		 HM.empty
		 childrenJS