{-# LANGUAGE NoMonomorphismRestriction #-}

import System.Random
import Control.Monad.State
import Data.Maybe
import Data.List
import Data.Tuple

import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo

import Data.Colour.Names
import Diagrams.Prelude
import Diagrams.Backend.SVG.CmdLine


-----------------------------------------
-- Network construction and properties --
-----------------------------------------

type Node = Int
type Edge = (Node, Node)
type Adjacency = [Node]
type Network = [Adjacency]

hasNode :: Node -> Network -> Bool
hasNode v n = v < (length n)

addNode :: Network -> Network
addNode n = n <> [[]]

addHalfedge :: Node -> Node -> Network -> Network
addHalfedge v1 v2 n = (take v1 n) <> [(n !! v1) ++ [v2]] ++ (drop (v1 + 1) n)

addEdge :: Edge -> Network -> Network
addEdge (v1, v2) n = addHalfedge v1 v2 (addHalfedge v2 v1 n)

getNodes :: Network -> [Node]
getNodes n = [0..(length n - 1)]

getEdges :: Network -> [Edge]
getEdges n = foldr (<>) [] $ map (\i -> map (\j -> (fst i, j)) (snd i)) (zip (getNodes n) n)

getNeighbors :: Network -> Node -> Adjacency
getNeighbors n v = n !! v

emptyNetwork :: Int -> Network
emptyNetwork 0 = []
emptyNetwork l = addNode (emptyNetwork (l - 1))

-- | Take a network and add z short edges to each bond.
addShortEdges :: Int -> Network -> Network
addShortEdges 0 n = n
addShortEdges 1 n = n
addShortEdges z n = let 
                      l = length n
                      addZEdge v = addEdge (v, mod (v + quot z 2) l)
                    in addShortEdges (z - 2) (foldr addZEdge n [0..(l - 1)])

-- | Lists all possible long edges.
possibleLongEdges l z = let
                          d = quot z 2 + 1
                          e v = (\u -> (v, mod (v + u) l)) <$> [d..(l - d)]
                        in foldr (<>) [] $ e <$> [0..(l - 1)]

-- | Take a random element from a list, removing it from the list.
randomChoice :: RandomGen g => State ([a], g) (Maybe a)
randomChoice = get >>= (\(xs, g) -> if (length xs == 0) then (return Nothing) else (let (r, g2) = randomR (0, length xs - 1) g in (put ((take r xs) ++ (drop (r+1) xs), g2)) >> return (Just (xs !! r))))

-- | Returns a list of randomly chosen values with no repeats.
randomChoices :: RandomGen g => Int -> State ([a], g) [a]
randomChoices n = state (\si -> 
    let (maybeVals, sf) = runState (sequence $ replicate n randomChoice) si
    in  (catMaybes maybeVals, sf)
  )

-- | Adds the appropriate number of long edges to a network.
addLongEdges :: (RealFrac a, RandomGen g) => Int -> Int -> a -> Network -> State g Network
addLongEdges l z p n = state (\g ->
    let
      nle = floor (p * fromIntegral (l * quot z 2))
      (es, (_, g2)) = runState (randomChoices nle) (possibleLongEdges l z, g)
    in (foldr addEdge n $ es, g2)
  )

constructNetwork :: (RealFrac a, RandomGen g) => Int -> Int -> a -> State g Network
constructNetwork l z p = addLongEdges l z p $ addShortEdges z $ emptyNetwork l


----------------------------
-- Finding shortest paths --
----------------------------

pathLengths :: Network -> (Int, [Node], [Maybe Int]) -> (Int, [Node], [Maybe Int])
pathLengths n (level, curNodes, knownDists) =
  let
    setVertexFound node dists =
      if isNothing (head after) 
      then (before ++ [Just level] ++ (drop 1 after))
      else dists
        where (before, after) = splitAt node dists
    newDists = foldr setVertexFound knownDists curNodes
  in (level + 1, nub $ filter (\x -> isNothing (newDists !! x)) $ foldr (<>) [] $ getNeighbors n <$> curNodes, newDists)

findPathLengthsFromNode :: Network -> Node -> [Int]
findPathLengthsFromNode n v = catMaybes $ (\(_,_,d) -> d) $ head $ dropWhile (\(_,nodes,_) -> length nodes > 0) $ iterate (pathLengths n) (0, [v], replicate (length n) Nothing)

findAllPathLengths :: Network -> [Int]
findAllPathLengths n = foldr (<>) [] $ (findPathLengthsFromNode n) <$> getNodes n

findAveragePathLength :: RealFrac a => Network -> a
findAveragePathLength n = let xs = findAllPathLengths n in fromIntegral (sum xs) / fromIntegral (length xs)

shortestPath f n i j 0 | i == j                   = 0
                     | elem (i, j) $ getEdges n = 1
                     | elem (j, i) $ getEdges n = 1
                     | otherwise                = length n

shortestPath f n i j k = min (f n i j (k - 1)) (f n i k (k - 1) + f n k j (k - 1))

data NaturalTree a = NNode a (NaturalTree a) (NaturalTree a)

NNode a tl tr !!! 0 = a 
NNode a tl tr !!! n =
   if odd n
     then tl !!! top
     else tr !!! (top-1)
        where top = n `div` 2

instance Functor NaturalTree where
   fmap f (NNode a tl tr) = NNode (f a) (fmap f tl) (fmap f tr)

naturals r n =
   NNode n
     ((naturals $! r2) $! (n+r))
     ((naturals $! r2) $! (n+r2))
        where r2 = 2*r

shortestPaths n =
  let
    nodes = getNodes n
    memo = fmap (\y -> fmap (\z -> fmap z (naturals 1 0)) y) (fmap (\x -> fmap x (naturals 1 0)) (fmap (shortestPath shortestPath' n) (naturals 1 0)))
    shortestPath' n a b c = memo !!! a !!! b !!! c
  in 
    (\i j -> shortestPath' n i j (length n - 1)) <$> nodes <*> nodes
  

----------------------------
-- Betweenness centrality --
----------------------------

----------------------
-- Drawing routines --
----------------------

drawNode :: Node -> Diagram B
drawNode n = named n $ fc black $ circle 0.1

drawEdge :: Edge -> (Diagram B -> Diagram B)
drawEdge (n1, n2) = withName n1 $ \b1 -> withName n2 $ \b2 -> Diagrams.Prelude.atop ((location b1 ~~ location b2))

drawNetwork :: Network -> Diagram B
drawNetwork n = atPoints (trailVertices $ regPoly (length n) 1) (map drawNode $ getNodes n) Diagrams.Prelude.# applyAll (map drawEdge $ getEdges n)


-----------------------
-- Plotting routines --
-----------------------

lengthHistogram :: Colour Double -> [Int] -> EC (Layout Double Int) ()
lengthHistogram color lengths = plot $ fmap histToPlot $ liftEC $ do
    plot_hist_fill_style .= def {_fill_color = (withOpacity color 0.5)}
    plot_hist_line_style .= def {_line_color = (withOpacity color 1.0)}
    plot_hist_bins .= 20
    plot_hist_values .= ((fromIntegral <$> lengths) :: [Double])
    plot_hist_norm_func .= const id


main = do
  g1 <- newStdGen
  g2 <- newStdGen
  void $ renderableToFile (def { _fo_format = SVG}) "01.07.hist.svg" $ fillBackground def $ toRenderable $ do
    lengthHistogram blue  $ shortestPaths $ evalState (constructNetwork 200 2 0.02) g1
    lengthHistogram red   $ shortestPaths $ evalState (constructNetwork 200 2 0.20) g2

--  mainWith (drawNetwork n)