{- |
Module      :  $Header$
Description :  Family trees with lenses
Copyright   :  (c) 2012 Nathan "Taneb" van Doorn
License     :  MIT License

Maintainer  :  nvd124@gmail.com
Stability   :  unstable
Portability :  portable

This module is for Family Trees.
It's got lenses, which are pretty cool.
I would use "Data.Binary" to do saving and loading.

Better description coming soon!
-}
module Data.FamilyTree where

import Control.Arrow
import Control.Monad

import Data.Binary
import Data.Function
import Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HM
import Data.Lens.Common
import Data.IntMap (IntMap)
import qualified Data.IntMap as IM
import Data.IntSet (IntSet)
import qualified Data.IntSet as IS
import Data.Text (Text)
import qualified Data.Text as T 
import Data.Text.Encoding
import Data.Time

data Person = Person
  {name :: Maybe Text
  ,birthdate :: Maybe Day
  ,birthplace :: Maybe Location
  ,deathdate :: Maybe Day
  ,deathplace :: Maybe Location
  ,attributes :: HashMap Text Text
  ,attendedEvents :: IntSet
  } deriving (Eq, Show)

data Family = Family
  {head1 :: Maybe Int
  ,head2 :: Maybe Int
  ,relationship :: Maybe Relationship
  ,relationFrom :: Maybe Day
  ,relationTo :: Maybe Day
  ,children :: IntSet
  } deriving (Eq, Show)
  
data Event = Event 
  {eventInfo :: Text
  ,eventDate :: Maybe Day
  ,eventAttendees :: IntSet
  } deriving (Eq, Show)
  
data Location = Coord Double Double | PlaceName Text deriving (Eq, Show)

data Relationship = Marriage | Other Text deriving (Eq, Show)

data FamilyTree = FamilyTree
  {treeName :: Text
  ,people :: IntMap Person
  ,families :: IntMap Family
  ,events :: IntMap Event
  } deriving (Eq, Show)
  
instance Binary Location where
  put (Coord x y) = do
    put (0 :: Word8)
    put x
    put y
  put (PlaceName t) = do
    put (1 :: Word8)
    put (encodeUtf8 t)
  get = do
    tag <- getWord8
    case tag of
      0 -> do
        x <- get
        y <- get
        return $ Coord x y
      1 -> fmap (PlaceName . decodeUtf8) get

instance Binary Relationship where
  put (Other t) = do
    put (0 :: Word8)
    put (encodeUtf8 t)
  put Marriage = put (1 :: Word8)
  get = do
    tag <- getWord8
    case tag of
      0 -> fmap (Other . decodeUtf8) get
      1 -> return Marriage

instance Binary Person where
  put person = do
    put (fmap encodeUtf8 $ name person)
    put (fmap toModifiedJulianDay $ birthdate person)
    put (birthplace person)
    put (fmap toModifiedJulianDay $ deathdate person)
    put (deathplace person)
    put (map (join (***) encodeUtf8) . HM.toList $ attributes person)
    put (attendedEvents person)
  get = do
    n <- get
    bd <- get
    bp <- get
    dd <- get
    dp <- get
    a <- get
    e <- get
    return Person
      {name = fmap decodeUtf8 n
      ,birthdate = fmap ModifiedJulianDay bd
      ,birthplace = bp
      ,deathdate = fmap ModifiedJulianDay dd
      ,deathplace = dp
      ,attributes = HM.fromList $ map (join (***) decodeUtf8) a
      ,attendedEvents = e
      }

instance Binary Family where
  put fam = do
    put $ head1 fam
    put $ head2 fam
    put $ relationship fam
    put $ fmap toModifiedJulianDay $ relationFrom fam
    put $ fmap toModifiedJulianDay $ relationTo fam
    put $ children fam
  get = do
    h1 <- get
    h2 <- get
    r <- get
    rf <- get
    rt <- get
    c <- get
    return Family
      {head1 = h1
      ,head2 = h2
      ,relationship = r
      ,relationFrom = fmap ModifiedJulianDay rf
      ,relationTo = fmap ModifiedJulianDay rt
      ,children = c
      }
      
instance Binary Event where
  put evnt = do
    put . encodeUtf8 $ eventInfo evnt
    put $ fmap toModifiedJulianDay $ eventDate evnt
    put $ eventAttendees evnt
  get = do
    n <- get
    d <- get
    a <- get
    return Event
      {eventInfo = decodeUtf8 n
      ,eventDate = fmap ModifiedJulianDay d
      ,eventAttendees = a
      }

instance Binary FamilyTree where
  put tree = do
    put $ encodeUtf8 $ treeName tree
    put $ people tree
    put $ families tree
    put $ events tree
  get = do
    n <- get
    p <- get
    f <- get
    e <- get
    return FamilyTree
      {treeName = decodeUtf8 n
      ,people = p
      ,families = f
      ,events = e
      }

personLens :: Int -> Lens FamilyTree Person
personLens n = lens ((IM.! n) . people) $
  \person familyTree ->
    let oldPerson = people familyTree IM.! n
        newattended = (IS.difference `on` attendedEvents)
          person oldPerson
        oldattended = (IS.difference `on` attendedEvents)
          oldPerson person 
    in familyTree
      {people = IM.insert n person (people familyTree)
      ,events = IS.foldr (\i -> IM.adjust
        (\event -> event 
           {eventAttendees = IS.delete i (eventAttendees event)}) i)
        (IS.foldr (\i -> IM.adjust
          (\event -> event 
             {eventAttendees = IS.insert i (eventAttendees event)}) i)
          (events familyTree) newattended) oldattended
      }
      
familyLens :: Int -> Lens FamilyTree Family
familyLens n = lens ((IM.! n) . families) $
  \family' familyTree ->
    familyTree
      {families = IM.insert n family' (families familyTree)
      }
      
eventLens :: Int -> Lens FamilyTree Event
eventLens n = lens ((IM.! n) . events) $
  \event familyTree ->
    familyTree
      {events = IM.insert n event (events familyTree)
      ,people =
        let oldEventPeople =
              eventAttendees (events familyTree IM.! n)
              IS.\\ eventAttendees event
            newEventPeople = 
              eventAttendees event
              IS.\\ eventAttendees (events familyTree IM.! n)
        in IS.foldr
          (IM.adjust (\ person -> person
             {attendedEvents = IS.insert n (attendedEvents person)}))
          (IS.foldr
            (IM.adjust (\person -> person 
              {attendedEvents = IS.delete n (attendedEvents person)}))
            (people familyTree)
            oldEventPeople)
          newEventPeople
      }
      
blankPerson :: Person
blankPerson = Person
  {name = Nothing
  ,birthdate = Nothing
  ,birthplace = Nothing
  ,deathdate = Nothing
  ,deathplace = Nothing
  ,attributes = HM.empty
  ,attendedEvents = IS.empty
  }
  
blankFamily :: Family
blankFamily = Family
  {head1 = Nothing
  ,head2 = Nothing
  ,relationship = Nothing
  ,relationFrom = Nothing
  ,relationTo = Nothing
  ,children = IS.empty
  }
  
blankEvent :: Event
blankEvent = Event
  {eventInfo = T.empty
  ,eventDate = Nothing
  ,eventAttendees = IS.empty
  }
  
addPerson :: FamilyTree -> (FamilyTree, Int)
addPerson familyTree =
  let ((n,_):_) = dropWhile (uncurry (==)) $
                    zip [1..] $ IM.keys $ people familyTree
  in (familyTree
       {people = IM.insert n blankPerson $ people familyTree}, n)
  
addFamily :: FamilyTree -> (FamilyTree, Int)
addFamily familyTree =
  let ((n,_):_) = dropWhile (uncurry (==)) $
                    zip [1..] $ IM.keys $ families familyTree
  in (familyTree 
       {families = IM.insert n blankFamily $ families familyTree}, n)

addEvent :: FamilyTree -> (FamilyTree, Int)
addEvent familyTree =
  let ((n,_):_) = dropWhile (uncurry (==)) $
                    zip [1..] $ IM.keys $ events familyTree
  in (familyTree {events = IM.insert n blankEvent $ events familyTree}, n)
  
deletePerson :: Int -> FamilyTree -> FamilyTree
deletePerson n familyTree =
  familyTree 
    {people = IM.delete n $ people familyTree
    ,families = IM.map
       (\fam -> fam
         {head1 = if head1 fam == Just n then Nothing else head1 fam
         ,head2 = if head2 fam == Just n then Nothing else head2 fam
         ,children = IS.delete n $ children fam
         }
         )
       (families familyTree)
    ,events = IM.map
      (\evnt -> evnt
        {eventAttendees = IS.delete n $ eventAttendees evnt}
        )
      (events familyTree)
    }
    
deleteFamily :: Int -> FamilyTree -> FamilyTree
deleteFamily n familyTree =
  familyTree
    {families = IM.delete n $ families familyTree}

deleteEvent :: Int -> FamilyTree -> FamilyTree
deleteEvent n familyTree =
  let relevantPeople = eventAttendees (events familyTree IM.! n)
  in familyTree
       {events = IM.delete n $ events familyTree
       ,people = IS.foldr (IM.adjust
         (\p -> p {attendedEvents = IS.delete n $ attendedEvents p}))
         (people familyTree) relevantPeople
       }