--{-# LANGUAGE OverloadedStrings #-}
module Main (main) where

import System.FilePath ( (</>) )
import Data.Monoid ( (<>) )
import Control.Monad ( forM )
import qualified Data.Map as M
import qualified Database.HDBC.Sqlite3 as SQL ( connectSqlite3 )
import qualified Options.Applicative as O
import qualified Text.PrettyPrint.Leijen as PP
import Text.PrettyPrint.Leijen ( pretty, (<+>), (<$$>) )

import Main.Report (reportDatabase)
import Database.Priors ( calculatePriors )
import Database.HumanAssessment ( summarise, summariseTasks, getHumanMachineFiles, tally, mismatchState, expectedAbsence, prettyMatches )
import Database.Data ( PilotID, TaskID, NodeID, allEvidence )
import Util.Pretty ( prettyMap )


-- Filter pilot IDs.
filterPilots :: [PilotID] -> [(PilotID, a, b)] -> [(PilotID, a, b)]
filterPilots [] = id
filterPilots ps = filter (\(p,_,_) -> p `elem` ps)


-- | Filter node IDs.
filterNodes :: [NodeID] -> M.Map (TaskID, NodeID) m -> M.Map (TaskID, NodeID) m
filterNodes [] = id
filterNodes ns = M.filterWithKey (\(_,n) _ -> n `elem` ns)


-- | Filter task IDs.
filterTasks :: [TaskID] -> M.Map (TaskID, NodeID) m -> M.Map (TaskID, NodeID) m
filterTasks [] = id
filterTasks ts = M.filterWithKey (\(t,_) _ -> t `elem` ts)


data Command 
   = PriorsCommand PriorsArguments
   | HumanVsMachineCommand HumanVsMachineArguments
   | OutFileCommand OutFileArguments

data OutFileArguments = OutFileArguments
   { inDB :: FilePath
   , outHTML :: FilePath
   }

data PriorsArguments = PriorsArguments
   { priorsDatabases :: [FilePath]
   }

data HumanVsMachineArguments = HumanVsMachineArguments
   { humanDir :: FilePath
   , machineDir :: FilePath
   , pilots :: [PilotID] 
   , tasks :: [TaskID] --htmlOutput :: Maybe FilePath
   , nodes :: [NodeID]
   , showTotals :: Bool
   , analysisFlag :: Bool
   }

cli :: IO Command --Arguments
cli = O.execParser $ O.info (O.helper <*> commands) (O.progDesc "Advise-Me reporter and performance analysis tool")

   where

   commands :: O.Parser Command
   commands = O.hsubparser $ 
      O.command "priors" priorsParser <> 
      O.command "humanvsmachine" humanVsMachineParser <>
      O.command "outfiles" outFileParser

   outFileParser :: O.ParserInfo Command
   outFileParser = OutFileCommand <$> O.info (O.helper <*> outFileArgs) (O.progDesc "make a HTML report (legacy)")

   outFileArgs :: O.Parser OutFileArguments
   outFileArgs = OutFileArguments <$>
         ( O.strOption $
            O.short 'i' <> O.long "input" <>
            O.help "Input database"
         )
         <*>
         ( O.strOption $
            O.short 'o' <> O.long "output" <>
            O.help "Output file" <> O.value "out.html"
         )


   priorsParser :: O.ParserInfo Command
   priorsParser = PriorsCommand <$> O.info (O.helper <*> priorsArgs) (O.progDesc "calculation of node priors")

   priorsArgs :: O.Parser PriorsArguments
   priorsArgs = PriorsArguments <$> databaseFiles

   databaseFiles =
         ( O.some . O.strOption $ 
            O.short 'f' <> O.long "file" <>
            O.help "database file(s) to consider"
         )

   humanVsMachineParser :: O.ParserInfo Command
   humanVsMachineParser = HumanVsMachineCommand <$> O.info (O.helper <*> humanVsMachineArgs) (O.progDesc "human-vs-machine comparison")

   humanVsMachineArgs :: O.Parser HumanVsMachineArguments
   humanVsMachineArgs = HumanVsMachineArguments
      <$>
         ( O.strOption $
            O.long "human-directory" <>
            O.help "Directory containing the human assessments" <>
            O.value ("pilots" </> "assessments")
         )
      <*>
         ( O.strOption $
            O.long "machine-directory" <>
            O.help "Directory containing the machine assessments" <>
            O.value ("pilots" </> "processed")
         )
      <*>
         ( O.many . O.strOption $
            O.short 'p' <> O.long "pilot" <>
            O.help "Consider only the given pilots"
         )
      <*>
         ( O.many . O.strOption $
            O.short 't' <> O.long "task" <>
            O.help "Consider only the given tasks"
         )
      <*>
         ( O.many . O.strOption $
            O.short 'n' <> O.long "node" <>
            O.help "Consider only the given nodes"
         )
      <*>
         ( O.switch $
            O.short 's' <> O.long "node-summary" <>
            O.help "Show summary value for nodes per task"
         )
      <*> ( O.switch $
            O.short 'a' <> O.long "analysis" <>
            O.help "Show overview of mismatches"
          )

main :: IO ()
main = do
   cmd <- cli
   case cmd of
      PriorsCommand args -> mainPriors args
      HumanVsMachineCommand args -> mainHumanVsMachine args
      OutFileCommand args -> mainOutFiles args


mainOutFiles :: OutFileArguments -> IO ()
mainOutFiles arg = reportDatabase (inDB arg) (outHTML arg)


mainPriors :: PriorsArguments -> IO ()
mainPriors args = do
   priors <- fmap (foldr (<>) mempty) . forM (priorsDatabases args) $ \path -> 
      calculatePriors <$> (SQL.connectSqlite3 path >>= allEvidence)
   PP.putDoc $
      pretty "Priors:" <$$>
      pretty priors <$$>
      mempty


-- | Looks for spreadsheets with human assessments of the pilot data (in the
-- @tests/human-scored@ directory), and shows how they compare against the
-- machine assessments.
mainHumanVsMachine :: HumanVsMachineArguments -> IO ()
mainHumanVsMachine arg = do
   files <- filterPilots (pilots arg) <$> getHumanMachineFiles (machineDir arg) (humanDir arg)
   tallies' <- mapM (\(_,s,d) -> tally s d) files

   let tallies = filterNodes (nodes arg)
               . filterTasks (tasks arg)
               . foldl (M.unionWith (<>)) mempty
               $ tallies'
   let stally = summarise tallies

   let prettyTallies' = if showTotals arg 
         then prettyMap $ summariseTasks tallies 
         else prettyMap $ tallies

   PP.putDoc $ 
      PP.text "Considering pilots:" <+> 
      PP.hsep (map (\(p,_,_) -> PP.text p) files) <$$> 
      mempty <$$>
      pretty "Tallies:" <$$>
      prettyTallies' <$$>
      pretty "Summary:" <$$> 
      pretty stally <$$>
      mempty

   putStrLn ""

   if not $ analysisFlag arg 
      then mempty
      else PP.putDoc $
         pretty "Contradictions:" <$$>
         mempty <$$>
         prettyMatches (mismatchState stally ++ expectedAbsence stally)

   putStrLn ""