mathgenealogy-1.1.0: Main.hs

-----------------------------------------------------------------------------
-- |
-- Module      :  Main
-- Copyright   :  (C) Peter Robinson 2010-2012
-- License     :  GPL-2
--
-- Maintainer  :  Peter Robinson <thaldyron@gmail.com>
-- Stability   :  experimental
-- Portability :  portable 
--
-----------------------------------------------------------------------------
module Main
where
import Control.Monad
import Control.Applicative
import Control.Concurrent
import Control.Exception
import Data.GraphViz
import Data.GraphViz.Attributes.Complete
import qualified Data.ByteString.Lazy.Char8 as B
import qualified Data.List as L
import System.Cmd(system)
import System.IO(hFlush,stdout)
import System.Exit
import System.Directory
import Data.Maybe
import Data.Text.Lazy(Text)
import Data.Text.Lazy.Encoding 
import qualified Data.Text.Lazy as T
import System.Console.CmdArgs
import Prelude hiding(catch)

import Extract
import Entry
import Graph

data OutputFile = FilePDF | FilePNG | FileSVG
  deriving(Eq,Data,Typeable)

instance Show OutputFile where
  show FilePDF = ".pdf"
  show FilePNG = ".png"
  show FileSVG = ".svg"

outputFileToGraphviz :: OutputFile -> String
outputFileToGraphviz FilePDF = " -Tpdf "
outputFileToGraphviz FilePNG = " -Tpng "
outputFileToGraphviz FileSVG = " -Tsvg "


instance Default OutputFile where
  def = FilePDF

-- command line arguments:
data MathGenealogy = MathGenealogy 
  { filePrefix   :: String 
  , keepDotFile  :: Bool
  , graphvizArgs :: String
  , onlyDotFile  :: Bool
  , verbose      :: Bool
  , includeTheses:: Bool
  , startURL     :: String
  , outputFile   :: OutputFile
  }
  deriving(Eq,Data,Typeable,Show)

mathGenealogy = MathGenealogy
  { startURL     = def &= args &= typ "URL" 
  , filePrefix   = "output" &= typ "PREFIX" 
  , keepDotFile  = False 
--  , graphvizArgs = " -Tpdf -Gcharset=utf8 " &= typ "<graphviz parameters>" &= opt " -Tpdf -Gcharset=utf8 "
  , graphvizArgs = " -Gcharset=utf8 " &= typ "<graphviz parameters>" -- &= opt " -Gcharset=utf8 "
  , onlyDotFile  = False &= help "Only create .dot file"
  , verbose      = False &= help "Print data to terminal."
  , includeTheses= False &= help "Include PhD thesis in output"
  , outputFile   = enum [ FilePDF &= help "create PDF file"
                        , FilePNG &= help "create PNG file"
                        , FileSVG &= help "create SVG file"
                        ]
  } &= summary "Mathematics Genealogy Visualizer" 
    &= details[ "Run the program with a start-URL, for example:",
             "# mathgenealogy http://genealogy.math.ndsu.nodak.edu/id.php?id=18231"]

-- | GraphViz style attributes for edges and nodes.
-- See: http://hackage.haskell.org/packages/archive/graphviz/2999.12.0.4/doc/html/Data-GraphViz-Attributes.html
-- TODO: figure out if GraphViz supports CSS files or similar. 
edgeAtt = const [ Dir Forward
                , color LightBlue2
                ] 

nodeAtt (num,e) = 
  [ textLabel (T.pack (show e)) 
  , Shape BoxShape
  , FontSize (if num == 1 then 18 else 14)
  , styles [rounded,filled,bold]
  , color LightYellow1
  , FontName (T.pack "ZapfChancery-MediumItalic") -- "Helvetica")
  ]


main = do
  args <- cmdArgs mathGenealogy
  let targs = args
  when (null $ startURL args) $ 
    throw (userError "Missing start-URL. Run 'mathgenealogy --help' for help.")

  gvExecutable <- (do m <- findExecutable "dot" 
                      if m == Nothing 
                        then throw (userError "Error - Couldn't find 'dot' program. Did you install graphviz?")
                        else return (fromJust m))

  let traverseEntries :: [Text] -> IO [Entry]
      traverseEntries = traverseEntries' 0 [] []
        where
          traverseEntries' _ acc _        []         = return $ L.nub acc
          traverseEntries' c acc prevUrls (url:urls) = do
            e <- (if includeTheses args then id else removeThesis)
                 `liftM` downloadEntry url
            threadDelay 1000000
            if verbose args 
              then print e
              else do 
                putStr $ show c ++ " "
                hFlush stdout
            let newUrls = urlAdvisors e
            traverseEntries' (c+1) (e:acc) (url:prevUrls) 
                             ((newUrls L.\\ prevUrls) ++ urls)

  putStrLn "Downloading entries from http://genealogy.math.ndsu.nodak.edu..."
  putStrLn "(this might take a few minutes)"
  theGraph <- entryGraph <$> traverseEntries [T.pack $ startURL args]
  putStrLn "done. :)"
  let dotFileName = filePrefix args ++ ".dot"         
  putStr $ "Writing DOT-file " ++ dotFileName ++ "..."
  let output = printDotGraph $ 
                 graphToDot nonClusteredParams{ fmtNode = nodeAtt 
                                              , fmtEdge = edgeAtt
                                              } theGraph
  B.writeFile dotFileName (encodeUtf8 output)
    `catch` (\(e::IOException) -> do { print e ; throw e })
  putStrLn "done. :)"

  unless (onlyDotFile args) $ do
    putStr "Generating graphics file..."
    let command = gvExecutable ++ " " ++ outputFileToGraphviz (outputFile args) ++ graphvizArgs args
                               ++ " " ++ dotFileName ++ " > " 
                               ++ filePrefix args ++ show (outputFile args)
    print command
    result <- system command
    when (isExitFailure result) $ do
      print $ "Error running the graphviz (dot) program. I tried: " ++ command
      throw $ userError "Existing."
    putStrLn "done. :)"
  unless (keepDotFile args) $ removeFile dotFileName
  where
    isExitFailure (ExitFailure _) = True
    isExitFailure _               = False

    
downloadEntry :: Text -> IO Entry
downloadEntry = liftM parseEntry . getTags