module Main where

import Data.DecisionTree

import Data.List
import qualified Data.Map as Map
import qualified Data.Foldable as Foldable

-- ### A decision tree matching package descriptions to packages ###
-- The data was taken from the hackage package summary and cleaned up slightly

-- Build and draw the tree
hackageTree :: Int -> IO (DTree String [String])
hackageTree n = do
    assocs <- readFile "assocs.txt"
    let mp = Map.fromListWith (\a b -> nub (a++b)) $ (read assocs :: [(String,[String])])
    let testAtts = [Att (\a -> a `elem` (Foldable.concat $ Map.lookup k mp)) k | k <- Map.keys mp]

    packages <- readFile "packages.txt"
    let testRows = read packages :: [String]

    let splitter = minSplit $ \ tl fl -> abs (length tl - length fl)

    let testTree = fmap snd $
                   maxDecisions n $
                   prune (\(_,as) -> entropy as == 0) $
                   runSplitter splitter (testAtts,testRows)
    return testTree

-- ### An example from my homework ###

data Student = Student {
    firstLastYear :: Bool,
    male :: Bool,
    worksHard :: Bool,
    drinks :: Bool,
    firstThisYear :: Bool
    } deriving (Eq,Ord,Show)

students = [richard,alan,alison,jeff,gail,simon]

richard = Student True True False True True
alan = Student True True True False True
alison = Student False False True False True
jeff = Student False True False True False
gail = Student True False True True True
simon = Student False True True True False

matthew = Student False True False True True
mary = Student False False True True False

-- Sigh
atts = [
    Att firstLastYear "firstLastYear",
    Att male "male",
    Att worksHard "worksHard",
    Att drinks "drinks"
    ]

target = firstThisYear

sumEntropy a b = (entropy $ map target a) + (entropy $ map target b)

tree = fmap (map target . snd) $ prune (\(_,as) -> entropy as == 0) $ runSplitter (minSplit sumEntropy) (atts,students)