{-# LANGUAGE ScopedTypeVariables #-}
-- | This module implements the composition of a Majority Judgment
-- from a tree of Majority Judgments: for the same question,
-- the same choices, the same judges and the same grades.
-- In that tree, a parent judgment is formed by the aggregation of its children judgments,
-- where a child judgment contributes only for a percentage of the parent judgment.
module Majority.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 Majority.Merit

-- * 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)

-- | @'opinionsBySection' cs js ss@ computes the 'Opinions' of the given 'Judges' @js@ about the given 'choice's @cs@,
-- from the 'grade' (specified or omitted) attributed to 'choice's
-- and the 'Share's (specified or omitted) attributed to 'Tree.Node'
-- in given 'Tree' @ss@.
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