{-# LANGUAGE TypeFamilies, ScopedTypeVariables, DeriveDataTypeable, NamedFieldPuns, BangPatterns #-}

module Main (
  main
) where

import Control.Applicative
import Control.DeepSeq
import Control.Exception
import Control.Monad
import Control.Monad.ST
import Data.Array hiding (elems)
import Data.Array.ST
import Data.Data
import Data.Graph.Etage
import Data.Graph.Inductive hiding (edges, defaultGraphSize)
import qualified Data.Map as M
import Data.List
import Data.Maybe
import Data.Ratio
import Data.Time.Clock.POSIX
import GHC.Arr
import System.Console.GetOpt
import System.Environment
import System.Exit
import System.IO
import System.Random
import System.Timeout
import Text.Printf

import Control.Etage

defaultGraphSize :: Int
defaultGraphSize = 6

minCollectTimeout :: Int
minCollectTimeout = 500000 -- microseconds

initialCollectTimeout :: Int
initialCollectTimeout = 5000000 -- microseconds

data Option = InputGraph String | OutputGraph String | OutputDot String | GraphSize Int | Help deriving (Eq, Show)

options :: [OptDescr Option]
options = [
    Option "g" ["graph"]  (ReqArg InputGraph "filename")            "read graph grom a file, default is to generate one randomly",
    Option "o" ["output"] (ReqArg OutputGraph "filename")           "save graph to a file",
    Option "d" ["dot"]    (ReqArg OutputDot "filename")             "save graph to a file in a GraphViz format",
    Option "s" ["size"]   (ReqArg (GraphSize . readParam) "number") ("size of the randomly generated graph, default is " ++ show defaultGraphSize),
    Option "h" ["help"]   (NoArg Help)                              "show this help"
  ]
    where readParam param = case reads param of
                              [(p,[])] -> if p < 1
                                            then error $ "invalid parameter `" ++ param ++ "'"
                                            else p
                              _        -> error $ "invalid parameter `" ++ param ++ "'"

isInputGraph :: Option -> Bool
isInputGraph (InputGraph _) = True
isInputGraph _              = False

isGraphSize :: Option -> Bool
isGraphSize (GraphSize _) = True
isGraphSize _             = False

isOutputGraph :: Option -> Bool
isOutputGraph (OutputGraph _) = True
isOutputGraph _               = False

isOutputDot :: Option -> Bool
isOutputDot (OutputDot _) = True
isOutputDot _             = False

generateGraph :: Int -> IO (Gr String Double)
generateGraph graphSize = do
  when (graphSize < 1) $ throwIO $ AssertionFailed $ "Graph size out of bounds " ++ show graphSize
  let ns = map (\n -> (n, show n)) [1..graphSize]
  edges <- fmap concat $ forM [1..graphSize] $ \node -> do
    nedges <- randomRIO (0, graphSize)
    others <- fmap (filter (node /=) . nub) $ forM [1..nedges] $ \_ -> randomRIO (1, graphSize)
    gen <- getStdGen
    let weights = randomRs (1, 10) gen
    return $ zip3 (repeat node) others weights
  return $ mkGraph ns edges

data TestNeuron a b = TestNeuron Int (Array (Node, Node) (b, [Node])) deriving (Typeable)

instance (Show a, Data a, Show b, Data b, Real b, Bounded b, NFData b) => Neuron (TestNeuron a b) where
  type NeuronFromImpulse (TestNeuron a b) = NoImpulse
  type NeuronForImpulse (TestNeuron a b) = GraphImpulse a b
  data NeuronOptions (TestNeuron a b) = NodeOptions {
      graphSize :: Int,
      knownPaths :: Array (Node, Node) (b, [Node])
    }

  mkDefaultOptions = return NodeOptions {
      graphSize = undefined,
      knownPaths = undefined
    }

  grow NodeOptions { graphSize, knownPaths } = return $ TestNeuron graphSize knownPaths
  
  live nerve (TestNeuron graphSize knownPaths) = do
    before <- getPOSIXTime
    pathsLazy <- stToIO $ newArray ((1, 1), (graphSize, graphSize)) (maxBound, [])
    collectTimeout <- collectPaths initialCollectTimeout pathsLazy
    pathsLazy' <- stToIO $ unsafeFreezeSTArray pathsLazy
    let !paths = pathsLazy' `deepseq` pathsLazy'
    after <- getPOSIXTime
    putStrLn $ "Etage search time for shortest paths: " ++ show (after - before - fromRational (fromIntegral collectTimeout % 1000000)) ++ " (" ++ printf "%fs" ((fromIntegral collectTimeout :: Double) / 1000000) ++ " timeout)" -- we correct for the last timeout
    let paths'      = M.fromList $ assocs paths
        knownPaths' = M.fromList $ assocs knownPaths
        shortest    = (fromIntegral . sum . map fromEnum . M.elems $ M.intersectionWith (\(l, p) (l', p') -> l == l' && p == p') knownPaths' paths') / fromIntegral (M.size knownPaths') :: Float
    putStrLn $ "Found " ++ printf "%.2f %%" (shortest * 100) ++ " shortest paths."
    dissolving ()
      where collectPaths :: Int -> STArray RealWorld (Node, Node) (b, [Node]) -> IO Int
            collectPaths collectTimeout arr = do
              before <- getPOSIXTime
              impulse <- timeout collectTimeout $ getForNeuron nerve
              case impulse of
                Nothing -> return collectTimeout
                Just i  -> do
                  let timestamp       = impulseTimestamp i
                      -- TODO: Improve timeout handling. Timeout only after the first TopologyChange?
                      collectTimeout' = max ((collectTimeout + round ((timestamp - before) * 2 * 1000000)) `div` 2) minCollectTimeout
                  case i of
                    TopologyChange {}                                                 -> collectPaths collectTimeout' arr
                    AddOutEdges {}                                                    -> collectPaths collectTimeout' arr
                    TopologyUpdate { destination = (node, _), path = (LP path, len) } -> do
                      let sourceNode = fst . head $ path
                      stToIO $ writeArray arr (sourceNode, node) (len, map fst path)
                      collectPaths collectTimeout' arr

main :: IO ()
main = do
  prepareEnvironment
  
  args <- getArgs
  opts <- case getOpt Permute options args of
            (o, [], [])  -> return o
            (_, p:_, []) -> throwIO $ ErrorCall $ "unrecognized option `" ++ p ++ "'"
            (_, _, errs) -> throwIO $ ErrorCall $ head $ lines $ head errs

  when (Help `elem` opts) $ do
    programName <- getProgName
    putStrLn $ "Usage:\n" ++ usageInfo (programName ++ " [option ...]" ++ "\n\nOptions:") options
    exitWith ExitSuccess
  
  (graph, graphSize) <- case find isInputGraph opts of
                          Just (InputGraph inputGraph) -> do
                            when (any isGraphSize opts) $ throwIO $ ErrorCall $ "conflicting options `" ++ "--graph" ++ "' and `" ++ "--size" ++ "'"
                            putStrLn $ "Reading graph from \"" ++ inputGraph ++ "\"."
                            [line1, line2] <- lines <$> readFile inputGraph
                            let g = mkGraph (read line1) (read line2)
                            forceStrictGraph g
                            return (g, noNodes g)
                          _                            -> do
                            let GraphSize s = fromMaybe (GraphSize defaultGraphSize) $ find isGraphSize opts
                            putStrLn $ "Generating a random graph of size " ++ show s ++ "."
                            g <- generateGraph s
                            forceStrictGraph g
                            return (g, s)

  case find isOutputGraph opts of
    Just (OutputGraph outputGraph) -> do
      putStrLn $ "Writing graph to \"" ++ outputGraph ++ "\"."
      writeFile outputGraph $ (show . labNodes $ graph) ++ "\n" ++ (show . labEdges $ graph) ++ "\n"
    _                              -> return ()

  case find isOutputDot opts of
    Just (OutputDot outputDot) -> do
      putStrLn $ "Writing graph in dot format to \"" ++ outputDot ++ "\"."
      writeFile outputDot $ graphviz graph "Etage" (8.27, 11.69) (1, 1) Landscape
    _                          -> return ()
  
  putStrLn $ "Graph contains " ++ show graphSize ++ " nodes."
  
  before <- getPOSIXTime
  let lazyPaths = dijkstraShortestPaths graph graphSize
      !paths    = lazyPaths `deepseq` lazyPaths
  after <- getPOSIXTime
  putStrLn $ "Dijkstra search time for shortest paths: " ++ show (after - before)

  incubate $ do
    nerveTest <- (growNeuron :: NerveOnlyFor (TestNeuron String Double)) (\o -> o { graphSize, knownPaths = paths })
    pathsNerves <- shortestPaths graph
    
    mapM_ (`attachTo` [TranslatableFor nerveTest]) $ M.elems pathsNerves
    
    sendTopologyChange pathsNerves

forceStrictGraph :: (NFData a, NFData b, Graph gr) => gr a b -> IO ()
forceStrictGraph g = labNodes g `deepseq` labEdges g `deepseq` return ()

dijkstraShortestPaths :: forall gr a b. (Graph gr, Bounded b, Real b) => gr a b -> Int -> Array (Node, Node) (b, [Node])
dijkstraShortestPaths graph graphSize = runSTArray buildPaths
  where buildPaths :: ST s (STArray s (Node, Node) (b, [Node]))
        buildPaths = do
          arr <- newArray ((1, 1), (graphSize, graphSize)) (maxBound, [])
          forM_ (nodes graph) $ \sourceNode ->
            forM_ (spTree sourceNode graph) $ \(LP (n@(node, len):ns)) ->
              writeArray arr (sourceNode, node) (len, reverse . map fst $ n:ns)
          return arr