{-# LANGUAGE BangPatterns #-}

import Graphics.Gloss
import Graphics.Gloss.Data.Vector
import System.Random
import Control.Monad

data MassPoint = MassPoint
    { pointPosition, pointVelocity :: {-# UNPACK #-} !Vector
    , pointMass :: {-# UNPACK #-} !Float
    } deriving (Show, Eq)

data Box = Box
   { boxXMin, boxXMax, boxYMin, boxYMax:: {-# UNPACK #-} !Float }
    deriving (Show, Eq)

data Tree
    = Tree
    { treeSubs :: [Tree]
    , width, height :: {-# UNPACK #-} !Float
    , massCenter :: {-# UNPACK #-} !Vector
    , totalMass :: {-# UNPACK #-} !Float
    , treeBox :: {-# UNPACK #-} !Box
    }
    deriving (Show, Eq)

main :: IO ()
main = do
    world <- randomPoints
    simulate (InWindow "BH" (windowSize, windowSize) (0, 0))
        white 20 (emptyTree, State [] world True) draw
        (\_ _ (_, state) ->
            let tree = head $ buildTree (boxForPoints $ points state) (points state)  in
            (tree, step tree state))
    where emptyTree = Tree [] 0 0 (0, 0) 0 (Box 0 0 0 0)

windowSize :: Int
windowSize = 800

data State = State { lastAccel :: [Vector]
                   , points :: [MassPoint]
                   , isFirstStep :: Bool
                   } deriving (Show, Eq)

step :: Tree -> State -> State
step tree !state = State newAccel movedPoints False -- HACK
    where newAccel = map (flip accelOf tree) (points state)
          lastAccel' = if isFirstStep state then newAccel else lastAccel state
          movedPoints = zipWith move (points state) $ zip lastAccel' newAccel

-- | leapfrog integration
{-# INLINE move #-}
move :: MassPoint -> (Vector, Vector) -> MassPoint
move !mp (lastAccel, newAccel) = mp
    { pointPosition = pointPosition mp + mulSV dt (pointVelocity mp) + mulSV (0.5 * dt^2) lastAccel
    , pointVelocity = pointVelocity mp + mulSV (0.5 * dt) (lastAccel + newAccel)
    }

dt :: Float
dt = 1.0

accelOf :: MassPoint -> Tree -> Vector
accelOf !mp tree
    | null (treeSubs tree) =
        let !r = sqrt (dotV between between)
            !between = massCenter tree - pointPosition mp -- checked
        in if r <= radius then (0,0) else mulSV (graviConst * totalMass tree / (eps + r)^2) $ normalizeV between
    | isFarAway tree mp = accelOf mp tree { treeSubs = [] }
    | otherwise = foldr ((+) . accelOf mp) (0, 0) (treeSubs tree)

graviConst :: Float
graviConst = 1e-2 -- fake

mass :: Float
mass = 100

radius :: Float
radius = 0.1

eps :: Float
eps = 1.0

{-# INLINE isFarAway #-}
isFarAway :: Tree -> MassPoint -> Bool
isFarAway tree !mp = d/r < 1
    where
        !d = max (width tree) (height tree)
        !r = norm (pointPosition mp - massCenter tree)

buildTree :: Box -> [MassPoint] -> [Tree]
buildTree _ [] = []
buildTree box [p] = [Tree [] w h (pointPosition p) (pointMass p) box]
    where
        !h = boxYMax box - boxYMin box
        !w = boxXMin box - boxXMax box
buildTree box ps = [Tree subs w h mc tm box]
    where
        !h = boxYMax box - boxYMin box
        !w = boxXMin box - boxXMax box
        subs = concat $ zipWith buildTree subBoxes pss
        !pss = foldr (insertPoint subBoxes) (replicate 4 []) ps
        insertPoint :: [Box] -> MassPoint -> [[MassPoint]] -> [[MassPoint]]
        insertPoint (box':subBoxes') point (pointBox:pointBoxes)
            | inBox box' (pointPosition point) = (point:pointBox):pointBoxes
            | otherwise = pointBox : insertPoint subBoxes' point pointBoxes
        insertPoint [] _ _ = error "bug in box finding"
        !subBoxes = splitBox box
        !mc = mulSV (1 / tm) $ sum $ map (\pm -> pointMass pm `mulSV` pointPosition pm) ps
        !tm = sum $ map pointMass ps

{-# INLINE boxForPoints #-}
boxForPoints :: [MassPoint] -> Box
boxForPoints !ps = Box
    (minimum (map (fst . pointPosition) ps))
    (maximum (map (fst . pointPosition) ps))
    (minimum (map (snd . pointPosition) ps))
    (maximum (map (snd . pointPosition) ps))

{-# INLINE inBox #-}
inBox :: Box -> Vector -> Bool
inBox !box !(x, y) =
    x >= boxXMin box && x <= boxXMax box &&
    y >= boxYMin box && y <= boxYMax box

{-# INLINE splitBox #-}
splitBox :: Box -> [Box]
splitBox !box =
    [ Box (boxXMin box) centerX        centerY       (boxYMax box) -- left top
    , Box centerX       (boxXMax box)  centerY       (boxYMax box)  -- right top
    , Box (boxXMin box) centerX        (boxYMin box) centerY -- left botton
    , Box centerX       (boxXMax box)  (boxXMin box) centerY -- right botton
    ]
    where
        !w = boxXMax box - boxXMin box
        !h = boxYMax box - boxYMin box
        !w' = w / 2
        !h' = h / 2
        !centerX = boxXMin box + w'
        !centerY = boxYMin box + h'

randomPoints :: IO [MassPoint]
randomPoints = replicateM 300 $ do
    angle <- randomRIO (0, 2*pi)
    dist <- randomRIO (0, fromIntegral $ windowSize `div` 2)
    let pos = (sin angle * dist, cos angle * dist)
    let vel = (0, 0) -- normalizeV (snd pos, -fst pos)
    return $ MassPoint pos vel mass

draw :: (Tree, State) -> Picture
draw (tree, state) = Pictures $ drawTree tree : map drawPoint (points state)

drawPoint :: MassPoint -> Picture
drawPoint p = uncurry Translate (pointPosition p) $ Color red $ Circle (4 * fromIntegral windowSize / 800)

drawTree :: Tree -> Picture
drawTree tree =
    Pictures $ drawBox (treeBox tree) : map drawTree (treeSubs tree)

drawBox :: Box -> Picture
drawBox !box = Line
    [(boxXMin box, boxYMin box), (boxXMax box, boxYMin box),
     (boxXMax box, boxYMax box), (boxXMin box, boxYMax box),
     (boxXMin box, boxYMin box)]

norm :: Vector -> Float
norm v = sqrt (dotV v v)
