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

import Control.Applicative (Applicative(..))
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
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 'Node'.
data Section grade
 =   Section
 {   sectionShare :: Maybe Share
     -- ^ A 'Share' within the parent '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 'Node'
     --   (defaulting to the 'grade' set on an ancestor 'Node' if any,
     --   or the |judge|'s default grade).
 } deriving (Eq, Show)

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

-- ** Type 'SectionByJudgeByChoice'
-- | Node value of a 'Tree' holding a 'Section', per 'judge', per choice.
type SectionByJudgeByChoice choice judge grade = HM.HashMap choice (SectionByJudge judge grade)

-- * 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 'Node'
-- in given 'Tree'.
opinionsBySection ::
 forall choice judge grade.
 Show choice =>
 Show judge =>
 Show grade =>
 Eq choice =>
 Hashable choice =>
 Eq judge =>
 Hashable judge =>
 Ord grade =>
 Choices choice ->
 Judges judge grade ->
 Tree (SectionByJudgeByChoice 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 (SectionByJudgeByChoice choice judge grade) ->
	      Either (ErrorSection choice judge grade)
	             (Tree (OpinionsByChoice choice judge grade))
	go defaultDistJC (Node currOpinJC childOpinJCS) =
		-- From current |Node|'s value.
			let currDistJC :: HM.HashMap choice (HM.HashMap judge (Distribution grade)) =
				-- Collect the 'Distribution' of current 'Node',
				-- and insert default 'Distribution'
				-- for each unspecified 'judge'
				-- of each (specified or unspecified) 'choice'.
				let specifiedDistJC =
					HM.mapWithKey (\choice ->
						HM.mapWithKey (\judge ->
							maybe (defaultDistJC HM.!choice HM.!judge) singleGrade .
							sectionGrade))
					 currOpinJC
				in
				HM.unionWith HM.union
				 specifiedDistJC
				 defaultDistJC
			in
		-- From children 'Node's.
			let defaultChildShareSJC = ([Nothing] <$ js) <$ defaultDistJC in
			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 (\childOpinJC ->
					let specifiedChildShareSJC = (pure . sectionShare <$>) <$> rootLabel childOpinJC in
					-- Fusion specified 'choice's into accum.
					HM.unionWith (HM.unionWith (List.++)) $
						-- Add default 'Share' for this '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 '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 'Node':
		-- current 'Distribution' is computed from current |Node|'s value ('currOpinJC')
		-- and inherited default 'Distribution' ('defaultDistJC').
		 [] -> Right $ 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 'Node's:
		-- current 'Opinions' is computed from the 'Opinions' of the children 'Node's.
		 _ -> do
			distJCS :: [Tree (HM.HashMap choice (HM.HashMap judge (Distribution grade)))] <-
				traverse (go $ currDistJC) childOpinJCS
				-- 'grade's set at current 'Node' ('currDistJC')
				-- become the new default 'grade's ('defaultDistJC')
				-- within its children '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 ->
					HM.mapWithKey (\judge ->
						Map.unionsWith (+) .
						List.zipWith (\share dist -> (share *) <$> dist)
						 (childShareSJC HM.!choice HM.!judge)))
				 distSJC
			Right $ Node distJC distJCS