content: {- TooManyCells.Paths.Distance
  Gregory W. Schwartz
  
  Collects functions pertaining to calculating distances within the tree (a type
  of "pseudotime").
  -}
  
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module TooManyCells.Paths.Distance
      ( linearNodeDistance
      , linearItemDistance
      , labelItemDistance
      ) where
  
  -- Remote
  import BirchBeer.Types
  import BirchBeer.Utility (getGraphLeavesCycles)
  import Control.Monad (join)
  import Data.Function (on)
  import Data.List (sortBy, genericLength)
  import Data.Maybe (fromMaybe)
  import Data.Tuple (swap)
  import Safe (headMay, atMay)
  import qualified Control.Foldl as Fold
  import qualified Control.Lens as L
  import qualified Data.Foldable as F
  import qualified Data.Graph.Inductive as G
  import qualified Data.Map.Strict as Map
  import qualified Data.Set as Set
  import qualified Data.Sparse.Common as S
  import qualified Data.Text as T
  
  -- Local
  import TooManyCells.MakeTree.Types
  import TooManyCells.Paths.Types
  
  -- | Convert path back to root.
  pathBackToRoot :: (G.DynGraph gr) => G.Node -> gr a Double -> gr a Double
  pathBackToRoot n gr = flip flipPath gr $ G.esp 0 n gr
  
  -- | Flip the path of a graph.
  flipPath :: (G.DynGraph gr) => [G.Node] -> gr a b -> gr a b
  flipPath [] !gr         = gr
  flipPath [n1] !gr       = gr
  flipPath (n1:n2:ns) !gr = flipPath ns $ flipEdge n1 n2 gr
  
  -- | Flip an edge in a graph, keeping the label.
  flipEdge :: (G.DynGraph gr) => G.Node -> G.Node -> gr a b -> gr a b
  flipEdge n1 n2 gr = G.insEdge (n2, n1, label) . G.delEdge (n1, n2) $ gr
    where
      label = maybe (error "No node found on path.") snd
            . F.find ((== n2) . fst)
            $ G.lsuc gr n1
  
  -- | Flip edges.
  flipEdges :: G.Node -> G.Context a Double -> G.Context a Double
  flipEdges n1 (a1, 0, l, a2) = ((0, n1) : a2, 0, l, removeN a1)
    where
      removeN = filter (\(_, x) -> x /= n1)
  flipEdges _ (a1, n, l, a2)  = (a2, n, l, a1)
  
  -- | Get the start of the diameter of a subtree. Can specify which node.
  treeDiameterStart
      :: G.Graph gr
      => PathDistance -> G.Node -> gr (G.Node, b) Double -> G.Node
  treeDiameterStart pd n gr =
      fst
          . fromMaybe (error "Empty tree when finding diameter.")
          . headMay
          . reverse
          . sortBy (compare `on` snd)
          . fmap (getDistance pd gr n . fst)
          . F.toList
          . getGraphLeavesCycles [] gr
          $ n
  
  -- | Get the shallowest leaf.
  shallowLeaf :: (G.Graph gr, Show a)
              => PathDistance -> gr (G.Node, a) Double -> G.Node
  shallowLeaf pd gr = fst
                    . fromMaybe (error "No nodes in tree.")
                    . headMay
                    . sortBy (compare `on` snd)
                    . fmap (getDistance pd gr 0 . fst)
                    . F.toList
                    . getGraphLeavesCycles [] gr
                    $ 0
  
  -- | Get the distance of each item from a starting point. Make graph undirected.
  linearItemDistance
      :: (TreeItem a, Show a)
      => ShallowFlag -> FlipFlag -> PathDistance -> ClusterGraph a -> [(a, Double)]
  linearItemDistance shallow direction pd (ClusterGraph gr) =
      concatMap (uncurry assignItems)
          . linearNodeDistance shallow direction pd
          . G.undir
          . G.emap (fromMaybe 0 . L.view edgeDistance)
          $ gr
    where
      assignItems n v =
          fmap (, v)
              . fromMaybe []
              . fmap (maybe [] F.toList . snd)
              . G.lab gr
              $ n
  
  -- | Flip the distances to reverse direction.
  flipDistance :: (Fractional b, Ord b) => [(a, b)] -> [(a, b)]
  flipDistance = Fold.fold subtractMaxAll
    where
      subtractMaxAll = (\m -> fmap (subtractMax m))
                   <$> Fold.premap snd Fold.maximum
                   <*> Fold.list
      subtractMax m = L.over L._2 (\x -> abs $ x - fromMaybe 0 m)
  
  -- | Get the distance of each leaf from a starting point.
  linearNodeDistance
      :: (G.Graph gr, Show a)
      => ShallowFlag
      -> FlipFlag
      -> PathDistance
      -> gr (G.Node, a) Double
      -> [(G.Node, Double)]
  linearNodeDistance (ShallowFlag shallow) (FlipFlag direction) pd gr =
      flipDirection direction
          . fmap (getDistance pd gr (start shallow) . fst)
          . F.toList
          . getGraphLeavesCycles [] gr
          $ 0
    where
      flipDirection True  = id
      flipDirection False = flipDistance
      start True          = shallowLeaf pd gr
      start False         = treeDiameterStart pd 0 gr
  
  -- | Get the distance of a base node n1 to another node n2.
  getDistance
      :: (G.Graph gr)
      => PathDistance -> gr a Double -> G.Node -> G.Node -> (G.Node, Double)
  getDistance PathStep gr n1 n2 = (n2, genericLength $ G.esp n1 n2 gr)
  getDistance PathModularity gr n1 n2 = (n2, fromMaybe 0 $ G.spLength n1 n2 gr)
  
  -- | Label path distances.
  labelItemDistance
      :: TreeItem a
      => LabelMap -> [(a, Double)] -> [(Label, Double)]
  labelItemDistance (LabelMap lm) =
      fmap (L.over L._1 (flip (Map.findWithDefault (Label "NA") . getId) lm))
  