module HGraph.Undirected.Layout.SpringModel
        (setup, step, positions)
where

import HGraph.Undirected
import qualified Data.Array.Unboxed as U
import qualified Data.Map as M
import Data.Foldable
import Linear

data SpringModel g a = 
  SpringModel
  { SpringModel g a -> g Int
smGraph :: g Int
  , SpringModel g a -> Array Int (V2 Double)
smPoints :: U.Array Int (V2 Double)
  , SpringModel g a -> Array Int (V2 Double)
smVelocities :: U.Array Int (V2 Double)
  , SpringModel g a -> Double
smSpringStrength    :: Double
  , SpringModel g a -> Double
smRepulsionStrength :: Double
  , SpringModel g a -> Double
smFriction :: Double
  , SpringModel g a -> Double
smCenterAttractionStrength :: Double
  , SpringModel g a -> Map Int a
smIndexToLabel :: M.Map Int a
  }

setup :: Double -> Double -> Double -> Double -> g a -> SpringModel g a
setup Double
centerAttraction Double
springStrength Double
repulsiveStrength Double
friction g a
gr =
  let (g Int
gi, [(Int, a)]
itol) = g a -> (g Int, [(Int, a)])
forall (t :: * -> *) a.
UndirectedGraph t =>
t a -> (t Int, [(Int, a)])
linearizeVertices g a
gr
      n :: Int
n = g Int -> Int
forall (t :: * -> *) b a.
(UndirectedGraph t, Integral b) =>
t a -> b
numVertices g Int
gi Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  in SpringModel :: forall (g :: * -> *) a.
g Int
-> Array Int (V2 Double)
-> Array Int (V2 Double)
-> Double
-> Double
-> Double
-> Double
-> Map Int a
-> SpringModel g a
SpringModel
  { smGraph :: g Int
smGraph = g Int
gi
  , smPoints :: Array Int (V2 Double)
smPoints = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
U.array (Int
0, Int
n) ([(Int, V2 Double)] -> Array Int (V2 Double))
-> [(Int, V2 Double)] -> Array Int (V2 Double)
forall a b. (a -> b) -> a -> b
$ [Int] -> [V2 Double] -> [(Int, V2 Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..Int
n] (V2 Double -> [V2 Double]
forall a. a -> [a]
repeat V2 Double
0)
  , smVelocities :: Array Int (V2 Double)
smVelocities = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
U.array (Int
0, Int
n) ([(Int, V2 Double)] -> Array Int (V2 Double))
-> [(Int, V2 Double)] -> Array Int (V2 Double)
forall a b. (a -> b) -> a -> b
$ [Int] -> [V2 Double] -> [(Int, V2 Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..Int
n] (V2 Double -> [V2 Double]
forall a. a -> [a]
repeat V2 Double
0)
  , smSpringStrength :: Double
smSpringStrength = Double
springStrength
  , smRepulsionStrength :: Double
smRepulsionStrength = Double
repulsiveStrength
  , smFriction :: Double
smFriction = Double
friction
  , smCenterAttractionStrength :: Double
smCenterAttractionStrength = Double
centerAttraction
  , smIndexToLabel :: Map Int a
smIndexToLabel = [(Int, a)] -> Map Int a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(Int, a)]
itol
  }

step :: Double -> SpringModel g a -> SpringModel g a
step Double
dt SpringModel g a
layout = SpringModel g a
layout{smPoints :: Array Int (V2 Double)
smPoints = Array Int (V2 Double)
points', smVelocities :: Array Int (V2 Double)
smVelocities = Array Int (V2 Double)
velocities'}
  where
    points :: Array Int (V2 Double)
points = SpringModel g a -> Array Int (V2 Double)
forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double)
smPoints SpringModel g a
layout
    velocities :: Array Int (V2 Double)
velocities = SpringModel g a -> Array Int (V2 Double)
forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double)
smVelocities SpringModel g a
layout
    points' :: Array Int (V2 Double)
points' = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
U.array (Array Int (V2 Double) -> (Int, Int)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
U.bounds Array Int (V2 Double)
points)
      [ (Int
i, (Array Int (V2 Double)
points Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
i) V2 Double -> V2 Double -> V2 Double
forall a. Num a => a -> a -> a
+ (Array Int (V2 Double)
velocities Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
i))
      | Int
i <- Array Int (V2 Double) -> [Int]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i]
U.indices Array Int (V2 Double)
points 
      ]
    velocities' :: Array Int (V2 Double)
velocities' = (Int, Int) -> [(Int, V2 Double)] -> Array Int (V2 Double)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
U.array (Array Int (V2 Double) -> (Int, Int)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
U.bounds Array Int (V2 Double)
points)
                          [ (Int
i, ((SpringModel g a -> Double
forall (g :: * -> *) a. SpringModel g a -> Double
smFriction SpringModel g a
layout) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
dt) Double -> V2 Double -> V2 Double
forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ (Array Int (V2 Double)
velocities Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
i) V2 Double -> V2 Double -> V2 Double
forall a. Num a => a -> a -> a
+ (SpringModel g a -> Int -> V2 Double
forall (g :: * -> *) a.
Adjacency g =>
SpringModel g a -> Int -> V2 Double
accelerate SpringModel g a
layout Int
i))
                          | Int
i <- Array Int (V2 Double) -> [Int]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i]
U.indices (Array Int (V2 Double) -> [Int]) -> Array Int (V2 Double) -> [Int]
forall a b. (a -> b) -> a -> b
$ Array Int (V2 Double)
points]

accelerate :: SpringModel g a -> Int -> V2 Double
accelerate SpringModel g a
layout Int
v = V2 Double
springV V2 Double -> V2 Double -> V2 Double
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^+^ V2 Double
centerV V2 Double -> V2 Double -> V2 Double
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^+^ V2 Double
pushV
  where
    positions :: Array Int (V2 Double)
positions = SpringModel g a -> Array Int (V2 Double)
forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double)
smPoints SpringModel g a
layout
    pv :: V2 Double
pv = Array Int (V2 Double)
positions Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
v
    springV :: V2 Double
springV = [V2 Double] -> V2 Double
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
      [ V2 Double
du V2 Double -> Double -> V2 Double
forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a
^* (SpringModel g a -> Double
forall (g :: * -> *) a. SpringModel g a -> Double
smSpringStrength SpringModel g a
layout)
      | Int
u <- g Int -> Int -> [Int]
forall (t :: * -> *) a. Adjacency t => t a -> a -> [a]
neighbors (SpringModel g a -> g Int
forall (g :: * -> *) a. SpringModel g a -> g Int
smGraph SpringModel g a
layout) Int
v
      , let pu :: V2 Double
pu = Array Int (V2 Double)
positions Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
u
      , let du :: V2 Double
du = V2 Double
pu V2 Double -> V2 Double -> V2 Double
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^-^ V2 Double
pv
      ]
    centerV :: V2 Double
centerV = (V2 Double -> V2 Double
forall (f :: * -> *) a. (Functor f, Num a) => f a -> f a
negated V2 Double
pv) V2 Double -> Double -> V2 Double
forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a
^* (SpringModel g a -> Double
forall (g :: * -> *) a. SpringModel g a -> Double
smCenterAttractionStrength SpringModel g a
layout)
    pushV :: V2 Double
pushV = [V2 Double] -> V2 Double
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum
      [ if Double
d2 Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Double
0.0001 then
           (Double
100 Double -> Double -> Double
forall a. Num a => a -> a -> a
* SpringModel g a -> Double
forall (g :: * -> *) a. SpringModel g a -> Double
smRepulsionStrength SpringModel g a
layout) Double -> V2 Double -> V2 Double
forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ (Double -> V2 Double
forall a. Floating a => a -> V2 a
angle (Double -> V2 Double) -> Double -> V2 Double
forall a b. (a -> b) -> a -> b
$ Double
2Double -> Double -> Double
forall a. Num a => a -> a -> a
*Double
forall a. Floating a => a
pi Double -> Double -> Double
forall a. Num a => a -> a -> a
* ((Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
v) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> Int
forall a. Ix a => (a, a) -> Int
U.rangeSize ((Int, Int) -> Int) -> (Int, Int) -> Int
forall a b. (a -> b) -> a -> b
$ Array Int (V2 Double) -> (Int, Int)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
U.bounds Array Int (V2 Double)
positions)))
        else
           ((V2 Double
du V2 Double -> Double -> V2 Double
forall (f :: * -> *) a.
(Functor f, Fractional a) =>
f a -> a -> f a
^/ (Double -> Double
forall a. Floating a => a -> a
sqrt Double
d2)) V2 Double -> Double -> V2 Double
forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a
^* (SpringModel g a -> Double
forall (g :: * -> *) a. SpringModel g a -> Double
smRepulsionStrength SpringModel g a
layout)) V2 Double -> Double -> V2 Double
forall (f :: * -> *) a.
(Functor f, Fractional a) =>
f a -> a -> f a
^/ Double
d2
      | Int
u <- g Int -> [Int]
forall (t :: * -> *) a. UndirectedGraph t => t a -> [a]
vertices (g Int -> [Int]) -> g Int -> [Int]
forall a b. (a -> b) -> a -> b
$ SpringModel g a -> g Int
forall (g :: * -> *) a. SpringModel g a -> g Int
smGraph SpringModel g a
layout
      , let pu :: V2 Double
pu = Array Int (V2 Double)
positions Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
u
      , let du :: V2 Double
du = V2 Double
pv V2 Double -> V2 Double -> V2 Double
forall (f :: * -> *) a. (Additive f, Num a) => f a -> f a -> f a
^-^ V2 Double
pu
      , let d2 :: Double
d2 = (Double -> Double -> Double) -> Double -> V2 Double -> Double
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Double
s Double
x -> Double
s Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
xDouble -> Double -> Double
forall a. Num a => a -> a -> a
*Double
x) Double
0 V2 Double
du
      , Int
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
v
      ]

positions :: SpringModel g a -> [(a, V2 Double)]
positions SpringModel g a
layout =
  [ (SpringModel g a -> Map Int a
forall (g :: * -> *) a. SpringModel g a -> Map Int a
smIndexToLabel SpringModel g a
layout Map Int a -> Int -> a
forall k a. Ord k => Map k a -> k -> a
M.! Int
i, (SpringModel g a -> Array Int (V2 Double)
forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double)
smPoints SpringModel g a
layout) Array Int (V2 Double) -> Int -> V2 Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
U.! Int
i)
  | Int
i <- Array Int (V2 Double) -> [Int]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [i]
U.indices (Array Int (V2 Double) -> [Int]) -> Array Int (V2 Double) -> [Int]
forall a b. (a -> b) -> a -> b
$ SpringModel g a -> Array Int (V2 Double)
forall (g :: * -> *) a. SpringModel g a -> Array Int (V2 Double)
smPoints SpringModel g a
layout
  ]