content: {-# LANGUAGE FlexibleContexts #-}
  module Main (main) where
  import Language.Haskell.Exts
  import Data.Generics.Uniplate.Data
  import Control.Monad.State
  import Control.Applicative
  import Control.Arrow
  import qualified Control.Exception as Exc
  import qualified Language.Haskell.Exts.Pretty as PP
  import qualified Pretty as P
  import System.Process
  import System.Directory
  import System.Environment
  import System.FilePath
  import System.IO.Error
  import System.IO
  import Data.Maybe
  import Data.Char
  import Data.List.Split
  import qualified Data.IntMap as IMap
  import Text.RegexPR
  import Numeric
  import Data.List
  
  type Exp_ = Exp
  pprint tm
    = "<style> a.info{    position:relative;    z-index:24;    color:#000;    text-decoration:none}  a.info:hover{z-index:25;} a.info span{display: none} a.info:hover span{ /*the span will display just on :hover state*/    display:block;    position:absolute;    top:2em; left:2em; width:15em;    border:1px solid #0cf;    background-color:#cff; color:#000;    text-align: center} </style><pre>"
        ++ P.prettyPrint 120 tm ++ "</pre>"
   
  assignSCC :: Module -> Module
  assignSCC m = evalState (transformBiM f m) 0
    where f e
            = do st <- get
                 put (st + 1)
                 return $ addSCC st $ strip e
  
          strip :: Exp_ -> Exp_
          strip (SCCPragma _ e) = e
          strip e = e
  
          addSCC :: Int -> Exp_ -> Exp_
          addSCC _ e@(Var _) = e
          addSCC _ e@(Lit _) = e
          addSCC num e = Paren $ SCCPragma (show num) e
  
  addColour :: IMap.IntMap Float -> String -> String
  addColour m s
    = gsubRegexPRBy (pref ++ ".*?" ++ suf)
        (\ str -> pref ++ sccNumToColour str ++ suf)
        s
    where
          sccNumToColour :: String -> String
          sccNumToColour str
            = maybe transparentColour toColour (readSCCNumber str `IMap.lookup` m)
  
          readSCCNumber :: String -> Int
          readSCCNumber str = fst $ head $ reads $ drop (length pref) str
          pref = "color: #"
          suf = "\">"
  
  transparentColour = "00ffffff"
  
  toColour :: Float -> String
  toColour fl
    | 0 <= fl && fl < 1.0e-2 = transparentColour
    | 1.0e-2 <= fl && fl < yellow =
      (pad $ showHex (truncate $ fl * 255 * (1 / yellow)) "") ++ "ff00"
    | yellow <= fl && fl <= 1 =
      "ff" ++ (pad $ showHex (255 - (truncate $ (fl - yellow) * 255 / (1 - yellow))) "") ++ "00"
    where
          pad :: String -> String
          pad str = if length str == 1 then '0' : str else str
          yellow = 0.1
  
  parseModuleFromFile :: FilePath -> IO Module
  parseModuleFromFile path = fromParseResult <$> parseFile path
  
  ind l n = if length l > n then Just $ l !! n else Nothing
  
  outputHTML file profName tm
    = do putStrLn "parsing profiling results"
         profFile <- lines <$> readFile profName
         let prof = filter ((all isDigit) . head) $ map (wordsBy (==' ')) $ takeWhile (/="") $ drop 9 profFile
         let profMap = IMap.fromList $ zip (map (read . (!!0)) prof) (map ((/100) . read . (!!2)) prof)
         putStrLn "printing output html file"
         let html = addColour profMap $ pprint tm
         writeFile (file ++ ".html") html
  
  main :: IO ()
  main
    = do args <- getArgs
         let mode = args !! 0
         let file = args !! 1
         let run = args !! 2
         unless ("-h" `isPrefixOf` mode || "-p" `isPrefixOf` mode)
           (error "mode should be -h or -px")
         let programArgs = ind args 3
         let inpFile = ind args 4
         let bak = file ++ ".bak"
         let profName = takeBaseName run ++ ".prof"
         (do removeFile bak `Exc.catch`
               (\ e -> unless (isDoesNotExistError e) $ ioError e)
             renameFile file bak
             putStrLn "started parsing original file"
             m <- parseModuleFromFile bak
             let tm = assignSCC m
             putStrLn "writing modified file"
             writeFile file $ PP.prettyPrint tm
  
             let buildStr = "ghc -prof -rtsopts -O2 --make " ++ run ++ ".hs"
             putStrLn buildStr
             buildCommand <- runCommand buildStr
             waitForProcess buildCommand
  
             let runRTS = ["+RTS", mode, "-RTS"]
             let runArgs = runRTS ++ maybe [] words programArgs
             inp <- maybe (return Inherit)
                      (\file -> UseHandle <$> openFile file ReadMode) inpFile
             putStrLn $ run ++ " " ++ unwords runArgs
             (_,_,_,runCommand) <- createProcess (proc ("./" ++ run) runArgs){std_in = inp}
             waitForProcess runCommand
  
             (when ("-px" `isPrefixOf` mode) (outputHTML file profName tm))
             renameFile file (file ++ ".scc"))
           `Exc.finally`
           (do copyFile bak file
               removeFile bak)
  