{-# OPTIONS_GHC -O2 -funbox-strict-fields #-}

module Vector where


class (Floating a, Ord a) => Scalar a


data (Scalar a) => Vector a = Vector !a !a
  deriving (Eq, Ord, Read, Show)

data (Scalar a) => Polar a = Polar !a !a
  deriving (Eq, Ord, Read, Show)


vmap :: (Scalar a) => (a -> a) -> Vector a -> Vector a
vmap f (Vector x y) = Vector (f x) (f y)

vzip :: (Scalar a) => (a -> a -> a) -> Vector a -> Vector a -> Vector a
vzip f (Vector x1 y1) (Vector x2 y2) = Vector (f x1 x2) (f y1 y2)

vfold :: (Scalar a) => (a -> a -> a) -> Vector a -> a
vfold f (Vector x y) = f x y


(|+|) :: (Scalar a) => Vector a -> Vector a -> Vector a
(|+|) = vzip (+)

(|-|) :: (Scalar a) => Vector a -> Vector a -> Vector a
(|-|) = vzip (-)


(|*|) :: (Scalar a) => Vector a -> Vector a -> a
v |*| w = vfold (+) (vzip (*) v w)


(|+) :: (Scalar a) => Vector a -> a -> Vector a
v |+ a = withPolarRadius v $ \r -> r + a

(+|) :: (Scalar a) => a -> Vector a -> Vector a
a +| v = withPolarRadius v $ \r -> a + r

(|-) :: (Scalar a) => Vector a -> a -> Vector a
v |- a = withPolarRadius v $ \r -> r - a

(-|) :: (Scalar a) => a -> Vector a -> Vector a
a -| v = withPolarRadius v $ \r -> a - r


(|*) :: (Scalar a) => Vector a -> a -> Vector a
v |* a = vmap (* a) v

(*|) :: (Scalar a) => a -> Vector a -> Vector a
a *| v = vmap (a *) v

(|/) :: (Scalar a) => Vector a -> a -> Vector a
v |/ a = vmap (/ a) v

(/|) :: (Scalar a) => a -> Vector a -> Vector a
a /| v = vmap (a /) v


vmag :: (Scalar a) => Vector a -> a
vmag v = sqrt (v |*| v)

vang :: (Scalar a) => Vector a -> a
vang (Vector x y) | x > 0 && y >= 0  = atan (y / x)
                  | x > 0 && y < 0   = atan (y / x) + 2 * pi
                  | x < 0            = atan (y / x) + pi
                  | x == 0 && y > 0  = pi / 2
                  | x == 0 && y < 0  = 3 * pi / 2
                  | x == 0 && y == 0 = 0


toVector :: (Scalar a) => Polar a -> Vector a
toVector (Polar r t) = r *| Vector (cos t) (sin t)

toPolar :: (Scalar a) => Vector a -> Polar a
toPolar v = Polar (vmag v) (vang v)

withPolar :: (Scalar a) => Vector a -> (Polar a -> Polar a) -> Vector a
withPolar v f = toVector (f (toPolar v))

withPolarRadius :: (Scalar a) => Vector a -> (a -> a) -> Vector a
withPolarRadius v f = withPolar v $ \(Polar r t) -> Polar (f r) t