{-# LANGUAGE ScopedTypeVariables #-}
module Hjugement.Section where

import Control.Applicative (Applicative(..), Alternative(..))
import Data.Bool
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Foldable (Foldable(..), any)
import Data.Function (($), (.))
import Data.Functor ((<$>), (<$))
import Data.Hashable (Hashable(..))
import Data.Maybe (Maybe(..), isNothing, maybe, fromMaybe)
import Data.Ord (Ord(..))
import Data.Traversable (Traversable(..))
import Data.Tree as Tree
import Prelude (Num(..), Fractional(..), toRational)
import Text.Show (Show(..))
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS
import qualified Data.List as List
import qualified Data.Map.Strict as Map

import Hjugement.MJ

-- * Type 'Section'
-- | An opinion of a 'judge' about a 'choice' at a specific section 'Tree.Node'.
data Section grade
 =   Section
 {   sectionShare :: Maybe Share
     -- ^ A 'Share' within the parent 'Tree.Node'
     --   (defaulting to a 'Share' computed as the remaining 'Share' to reach 1
     --   divided by the number of defaulted 'Share's).
 ,   sectionGrade :: Maybe grade
     -- ^ A 'grade' attributed to the current 'Tree.Node'
     --   (defaulting to the 'grade' set on an ancestor 'Tree.Node' if any,
     --   or the |judge|'s default grade).
 } deriving (Eq,Show)

-- ** Type 'SectionByJudge'
type SectionByJudge judge grade = HM.HashMap judge (Section grade)

-- ** Type 'SectionNode'
-- | Node value of a 'Tree' holding a 'Section', per 'judge', per 'choice'.
data SectionNode choice judge grade
 =   SectionNode
 {   sectionNodeShare       :: Maybe Share
     -- ^ A default 'sectionShare' for judges not specifying their own.
 ,   sectionByJudgeByChoice :: HM.HashMap choice (SectionByJudge judge grade)
 } deriving (Eq,Show)

-- * Type 'ErrorSection'
data ErrorSection choice judge grade
 =   ErrorSection_unknown_choices (HS.HashSet choice)
     -- ^ When some 'choice's are not known.
 |   ErrorSection_unknown_judges (HM.HashMap choice (HS.HashSet judge))
     -- ^ When some 'judge's are not known.
 |   ErrorSection_invalid_shares (HM.HashMap choice (HM.HashMap judge [Share]))
     -- ^ When at least one of the 'Share's is not positive, or when their sum is not 1.
 deriving (Eq,Show)

-- | Compute the 'Opinions' of the given |Judges| about the given 'Choices',
-- from the 'grade' (specified or omitted) attributed to 'Choice's
-- and the 'Share's (specified or omitted) attributed to 'Tree.Node'
-- in given 'Tree'.
opinionsBySection ::
 forall choice judge grade.
 Eq choice =>
 Eq judge =>
 Hashable choice =>
 Hashable judge =>
 Ord grade =>
 Choices choice ->
 Judges judge grade ->
 Tree (SectionNode choice judge grade) ->
 Either (ErrorSection choice judge grade)
        (Tree (OpinionsByChoice choice judge grade))
opinionsBySection cs js = go ((singleGrade <$> js) <$ HS.toMap cs)
	where
	go :: HM.HashMap choice (HM.HashMap judge (Distribution grade)) ->
	      Tree (SectionNode choice judge grade) ->
	      Either (ErrorSection choice judge grade)
	             (Tree (OpinionsByChoice choice judge grade))
	go defaultDistJC (Tree.Node (SectionNode _sectionNodeShare currOpinJC) childOpinJCS) =
		-- From current |Tree.Node|'s value.
			let currDistJC :: HM.HashMap choice (HM.HashMap judge (Distribution grade)) =
				-- Collect the 'Distribution' of current 'Tree.Node',
				-- and insert default 'Distribution'
				-- for each unspecified 'judge'
				-- of each (specified or unspecified) 'choice'.
				let specifiedDistJC =
					HM.mapWithKey (\choice ->
						let defaultDistJ = defaultDistJC HM.!choice in
						HM.mapWithKey (\judge ->
							maybe (defaultDistJ HM.!judge) singleGrade .
							sectionGrade))
					 currOpinJC
				in
				HM.unionWith HM.union
				 specifiedDistJC
				 defaultDistJC
			in
		-- From children 'Tree.Node's.
			let maybeChildShareSJC :: HM.HashMap choice (HM.HashMap judge [Maybe Share]) =
				-- Collect the (specified or explicitely (with 'Nothing') unspecified) 'Share's by section,
				-- and insert all unspecified 'Share's when a 'choice' or a 'judge' is unspecified.
				foldr (\(Tree.Node SectionNode{sectionNodeShare, sectionByJudgeByChoice} _) ->
					let defaultChildShareSJC = ([sectionNodeShare] <$ js) <$ defaultDistJC in
					let specifiedChildShareSJC =
						(<$> sectionByJudgeByChoice) $
						(pure . (<|> sectionNodeShare) . sectionShare <$>) in
					-- Fusion specified 'choice's into accum.
					HM.unionWith (HM.unionWith (List.++)) $
						-- Add default 'Share' for this 'Tree.Node',
						-- for each unspecified 'judge' of specified and unspecified 'choice'.
						HM.unionWith HM.union
						 specifiedChildShareSJC
						 defaultChildShareSJC)
				 HM.empty
				 childOpinJCS
			in
			let childShareSJC :: HM.HashMap choice (HM.HashMap judge [Share]) =
				-- Replace unspecified shares of each child 'Tree.Node'
				-- by an even default: the total remaining 'Share'
				-- divided by the number of unspecified 'Share's.
				(<$> maybeChildShareSJC) $ \maybeShareSJ ->
					(<$> maybeShareSJ) $ \maybeShareS ->
						let specifiedShare    = sum $ fromMaybe 0 <$> maybeShareS in
						let unspecifiedShares = toRational $ List.length $ List.filter isNothing maybeShareS in
						let defaultShare      = (1 - specifiedShare) / unspecifiedShares in
						fromMaybe defaultShare <$> maybeShareS
			in
		case childOpinJCS of
		-- Test for unknown choices.
		 _ | unknownChoices <- currOpinJC`HM.difference`defaultDistJC
		   , not $ null unknownChoices ->
			Left $ ErrorSection_unknown_choices $
				HS.fromMap $ (() <$) $ unknownChoices
		-- Test for unknown judges.
		 _ | unknownJudgesC <- HM.filter (not . null) $
		                       HM.intersectionWith HM.difference
		                        currOpinJC
		                        defaultDistJC
		   , not $ null unknownJudgesC ->
			Left $ ErrorSection_unknown_judges $
				HS.fromMap . (() <$) <$> unknownJudgesC
		-- Handle no child 'Tree.Node':
		-- current 'Distribution' is computed from current |Tree.Node|'s value ('currOpinJC')
		-- and inherited default 'Distribution' ('defaultDistJC').
		 [] -> Right $ Tree.Node currDistJC []
		-- Test for invalid shares.
		 _ | invalidSharesJC <-
		       HM.filter (not . null) $
		       HM.filter (\ss -> any (< 0) ss || sum ss /= 1)
		       <$> childShareSJC
		   , not $ null invalidSharesJC ->
			Left $ ErrorSection_invalid_shares invalidSharesJC
		-- Handle children 'Tree.Node's:
		-- current 'Opinions' is computed from the 'Opinions' of the children 'Tree.Node's.
		 _ -> do
			distJCS :: [Tree (HM.HashMap choice (HM.HashMap judge (Distribution grade)))] <-
				traverse (go $ currDistJC) childOpinJCS
				-- 'grade's set at current 'Tree.Node' ('currDistJC')
				-- become the new default 'grade's ('defaultDistJC')
				-- within its children 'Tree.Node's.
			let distSJC :: HM.HashMap choice (HM.HashMap judge [Distribution grade]) =
				-- Collect the 'Distribution's by section.
				foldr (\distJC ->
					let newDistSJC = (pure <$>) <$> rootLabel distJC in
					HM.unionWith (HM.unionWith (List.++)) newDistSJC)
				 HM.empty
				 distJCS
			let distJC :: HM.HashMap choice (HM.HashMap judge (Distribution grade)) =
				-- Compute the current 'Distribution' by scaling (share *) and merging (+)
				-- the children 'Distribution's.
				HM.mapWithKey (\choice ->
					let childShareSJ = childShareSJC HM.!choice in
					HM.mapWithKey (\judge ->
						let childShareS = childShareSJ HM.!judge in
						Map.unionsWith (+) .
						List.zipWith
						 (\share dist -> (share *) <$> dist)
						 childShareS))
				 distSJC
			Right $ Tree.Node distJC distJCS