-- John Garvin
-- CS 429 Final Project
-- April 24, 2001

import "hask/Picture"
import Model (model)
import Random

epsilon :: Float
epsilon = 0.2

grav_const :: Float
grav_const = 0.01

collides :: Position -> Position -> Bool
collides (Pos (x1,y1)) (Pos (x2,y2)) = abs (x2 - x1) < epsilon && abs (y2 - y1) < epsilon

newtype Position = Pos (Float, Float)
    deriving Show
newtype Velocity = Vel (Float, Float)
    deriving Show

-- Workaround for bugs in (**)
(***) :: (Floating a, Ord a) => a -> a -> a
a *** b = case a of
	  0           -> 0
	  a' | a' < 0 -> (-a') ** b
	  _           -> a ** b

data GravParticle = Particle {idnum :: Int,
			  pos :: Position,
			  vel :: Velocity,
			  mass :: Float}
		deriving Show

drawP :: GravParticle -> Region
drawP p = let Pos c = pos p 
	      in Translate c $ Shape $ Ellipse 0.01 0.01

grav :: GravParticle -> GravParticle -> Velocity
grav p1 p2 =
    let (Pos (x1,y1)) = pos p1
	(Pos (x2,y2)) = pos p2
	mag = grav_const * (mass p2)
	      * (abs $ ((x2 - x1)***2 + (y2 - y1)***2) *** -1.5)
	      / (mass p1)
    in Vel (mag * (x2 - x1),
	    mag * (y2 - y1))

vsum :: [Velocity] -> Velocity
vsum vs = Vel (foldr (\(x1,y1) (x2,y2) -> (x1+x2,y1+y2)) (0,0) (map (\(Vel v) -> v) vs))

vavg :: [Velocity] -> Velocity
vavg vs = Vel (x/len, y/len)
    where len = fromInt (length vs)
	  (Vel (x,y)) = vsum vs

vscale :: Velocity -> Float -> Velocity
vscale (Vel (x,y)) c = Vel (c*x,c*y)

move :: Velocity -> Position -> Position
move (Vel (vx,vy)) (Pos (px,py)) = Pos (px + vx, py + vy)

fsum :: [Float] -> Float
fsum = foldr (+) 0.0

pavg :: [Position] -> Position
pavg ps = 
    let getx (Pos (x,y)) = x
	gety (Pos (x,y)) = y
	in Pos ((fsum (map getx ps) / fromInt (length ps)),
		 (fsum (map gety ps) / fromInt (length ps)))

-- Upon collision of two or more particles, they "stick together" to
-- form one big particle.  I implement this as follows: let the one
-- with the lowest idnum be the "Borg" which assimilates the rest.
-- Average out the positions and the velocities, and add up the
-- masses.

oneStep :: GravParticle -> [GravParticle] -> [GravParticle]
oneStep (this_p @ Particle {idnum = my_i, pos = my_pos, vel = my_vel, mass = my_m}) ps
    = let any_borgs = any (\p -> pos p `collides` my_pos && idnum p < my_i) ps
	  assimilees = filter (\p -> pos p `collides` my_pos && idnum p > my_i) ps
	  new_v = vsum (my_vel : 
			map (\p -> grav this_p p) ps)
	  in if any_borgs
	  then []
	  else let sum_m = sum (my_m : map mass assimilees) 
		   in (case assimilees of
		       [] -> [Particle {idnum = my_i, 
					pos = move new_v my_pos,
					vel = new_v,
					mass = my_m}]
		       _  -> [Particle {idnum = my_i, 
					pos = pavg (my_pos : map pos assimilees),
					vel = (vavg  (my_vel `vscale` my_m :
						     map (\p -> vscale (vel p) (mass p))
						     assimilees)) 
					`vscale` (1 / sum_m),
					mass = sum_m}])

rotate_90r :: Velocity -> Velocity
rotate_90r (Vel (x,y)) = Vel (y,-x)

rough_rotate :: Velocity -> Velocity
rough_rotate (Vel (x,y)) | x >= 0 && y >= 0 = Vel ( 1,-1)
rough_rotate (Vel (x,y)) | x >= 0 && y <  0 = Vel (-1, 1)
rough_rotate (Vel (x,y)) | x <  0 && y >= 0 = Vel (-1,-1)
rough_rotate _                              = Vel ( 1, 1)


gen_data r num fun = 
    let vals = randoms (StdGen r r)
	adj (x,y) = (x * 5 - 2.5, y * 5 - 2.5)
	in take num (fun vals 0)

whirl r = let vals = randoms (StdGen r r)
	      adj (x,y) = (x * 5 - 2.5, y * 5 - 2.5)
	      make_part (n1:n2:n3:ns) i = Particle {idnum = i,
						    pos = Pos $ adj (n1, n2),
						    vel = (rough_rotate $ Vel $ adj (n1, n2))
						    `vscale` 0.1,
						    mass = 1}
					  : make_part ns (i+1)
	      in take 30 (make_part vals 0)

rand r = let vals = randoms (StdGen r r)
	     adj (x,y) = (x * 5 - 2.5, y * 5 - 2.5)
	     make_part (n1:n2:n3:n4:ns) i = Particle {idnum = i,
						      pos = Pos $ adj (n1, n2),
						      vel = (Vel $ adj (n3, n4)) `vscale` 0.1,
						      mass = 1}
					    : make_part ns (i+1)
	     in take 30 (make_part vals 0)

data1 :: [GravParticle]
data1 = [Particle {idnum = 0, 
		   pos = Pos (-2,1),
		   vel = Vel (0.05,0),
		   mass = 1.0},
	 Particle {idnum = 1,
		   pos = Pos (-2,-1),
		   vel = Vel (0.05,0),
		   mass = 1.0}]
{-
	Particle {idnum = 2,
		  pos = Pos (0,1),
		  vel = Vel (0,0.0),
		  mass = 4.0}]
-}

test1 = model drawP oneStep data1

test2 r = model drawP oneStep (whirl r)

test3 r = model drawP oneStep (rand r)