penrose-0.1.1.0: src/Utils.hs

-- | "Utils" contains frequently used utility function, some parameters to "Runtime", and debugging functions.
-- | Utils should not import other modules of Penrose.

-- This is for the "typeclass synonym"
{-# LANGUAGE ConstraintKinds #-}
{-# OPTIONS_HADDOCK prune #-}

module Utils where
import Debug.Trace
import Data.Typeable
import Control.Arrow

default (Int, Float)

-- | A more concise typeclass for polymorphism for autodiff
-- | NiceFloating :: * -> Constraint
-- https://stackoverflow.com/questions/48631939/a-concise-way-to-factor-out-multiple-typeclasses-in-haskell/48631986#48631986
type Autofloat a = (RealFloat a, Floating a, Real a, Show a, Ord a)
type Autofloat' a = (RealFloat a, Floating a, Real a, Show a, Ord a, Typeable a)

type Pt2 a = (a, a)
type Property = String

--------------------------------------------------------------------------------
-- Errors in the system

-- | Errors from one of the Substance, Style, or Element compilers.
data CompilerError
  = SubstanceParse String -- ^ an error in Substance parsing
  | StyleParse String -- ^ an error in Style parsing
  | ElementParse String -- ^ an error in Element parsing
  | SubstanceTypecheck String -- ^ an error in Substance typechecking
  | PluginParse String -- ^ an error in plugin parsing 
  | PluginRun String -- ^ an error when running a plugin 
  | StyleTypecheck String -- ^ an error in Style typechecking 
  | ElementTypecheck String -- ^ an error in Element typechecking
  deriving (Show)

-- | Errors from the optimizer
data RuntimeError =
  RuntimeError String
  deriving (Show)


--------------------------------------------------------------------------------
-- Parameters of the system

stepsPerSecond :: Int
stepsPerSecond = 100000

picWidth, picHeight :: Int
picWidth = 800
picHeight = 700

ptRadius :: Float
ptRadius = 4 -- The size of a point on canvas

defaultWeight :: Floating a => a
defaultWeight = 1

-- Debug flags
-- debug = True
debug = False
debugStyle = False

-- used when sampling the inital state, make sure sizes satisfy subset constraints
subsetSizeDiff :: Floating a => a
subsetSizeDiff = 10.0

epsd :: Floating a => a -- to prevent 1/0 (infinity). put it in the denominator
epsd = 10 ** (-10)

--------------------------------------------------------------------------------
-- General helper functions

fromRight :: (Show a, Show b) => Either a b -> b
fromRight (Left x) = error ("Failed with error: " ++ show x)
fromRight (Right y) = y

divLine :: IO ()
divLine = putStr "\n--------\n\n"

-- don't use r2f outside of zeroGrad or addGrad, since it doesn't interact well w/ autodiff
r2f :: (Fractional b, Real a) => a -> b
r2f = realToFrac

-- | Wrap a list around anything
toList :: a -> [a]
toList x = [x]

-- | similar to fst and snd, get the third element in a tuple
trd :: (a, b, c) -> c
trd (_, _, x) = x

-- | transform from a 2-element list to a 2-tuple
tuplify2 :: [a] -> (a,a)
tuplify2 [x,y] = (x,y)

-- | apply a function to each element of a 2-tuple
app2 :: (a -> b) -> (a, a) -> (b, b)
app2 f (x, y) = (f x, f y)

-- | apply a function to each element of a 3-tuple
app3 :: (a -> b) -> (a, a, a) -> (b, b, b)
app3 f (x, y, z) = (f x, f y, f z)

-- | generic cartesian product of elements in a list
cartesianProduct :: [[a]] -> [[a]]
cartesianProduct = foldr f [[]] where f l a = [ x:xs | x <- l, xs <- a ]

-- | given a side of a rectangle, compute the length of the half half diagonal
halfDiagonal :: (Floating a) => a -> a
halfDiagonal side = 0.5 * dist (0, 0) (side, side)

-- | `compose2` is used to compose with a function that takes in
-- two arguments. As if now, it is used to compose `penalty` with
-- constraint functions
compose2 :: (b -> c) -> (a -> a1 -> b) -> a -> a1 -> c
compose2 = (.) . (.)

concat4 :: [([a], [b], [c], [d])] -> ([a], [b], [c], [d])
concat4 x = (concatMap fst4 x, concatMap snd4 x, concatMap thd4 x, concatMap frth4 x)
fst4 (a, _, _, _) = a
snd4 (_, a, _, _) = a
thd4 (_, _, a, _) = a
frth4 (_, _, _, a) = a

-- | Define ternary expressions in Haskell
data Cond a = a :? a

infixl 0 ?
infixl 1 :?

(?) :: Bool -> Cond a -> a
True  ? (x :? _) = x
False ? (_ :? y) = y

--------------------------------------------------------------------------------
-- Internal naming conventions

nameSep, labelWord :: String
nameSep = " "
labelWord = "label"

labelName :: String -> String
labelName name = name ++ nameSep ++ labelWord

-- | Given a Substance ID and a Style ID for one of its associated graphical primitive, generate a globally unique identifier for this primitive
uniqueShapeName :: String -> String -> String
uniqueShapeName subObjName styShapeName = subObjName ++ nameSep ++ styShapeName
 -- e.g. "B yaxis" (the concatenation should be unique), TODO add the two names as separate obj fields





--------------------------------------------------------------------------------
-- Debug Functions

-- Some debugging functions.
debugF :: (Show a) => a -> a
debugF x = if debug then traceShowId x else x
debugXY x1 x2 y1 y2 = if debug then trace (show x1 ++ " " ++ show x2 ++ " " ++ show y1 ++ " " ++ show y2 ++ "\n") else id

-- To send output to a file, do ./EXECUTABLE 2> FILE.txt

-- For Runtime use only
tr :: Show a => String -> a -> a
tr s x = if debug then trace "---" $ trace s $ traceShowId x else x -- prints in left to right order

-- For Style use only
trs :: Show a => String -> a -> a
trs s x = if debugStyle then trace "---" $ trace s $ traceShowId x else x -- prints in left to right order

trRaw :: Show a => String -> a -> a
trRaw s x = trace "---" $ trace s $ trace (show x ++ "\n") x -- prints in left to right order

trStr :: String -> a -> a
trStr s x = if debug then trace "---" $ trace s x else x -- prints in left to right order

--------------------------------------------------------------------------------
-- Lists-as-vectors utility functions

-- define operator precedence: higher precedence = evaluated earlier
infixl 6 +., -.
infixl 7 *., /.

-- assumes lists are of the same length
dotL :: (RealFloat a, Floating a) => [a] -> [a] -> a
dotL u v = if not $ length u == length v
           then error $ "can't dot-prod different-len lists: " ++ (show $ length u) ++ " " ++ (show $ length v)
           else sum $ zipWith (*) u v

(+.) :: Floating a => [a] -> [a] -> [a] -- add two vectors
(+.) u v = if not $ length u == length v
           then error $ "can't add different-len lists: " ++ (show $ length u) ++ " " ++ (show $ length v)
           else zipWith (+) u v

(-.) :: Floating a => [a] -> [a] -> [a] -- subtract two vectors
(-.) u v = if not $ length u == length v
           then error $ "can't subtract different-len lists: " ++ (show $ length u) ++ " " ++ (show $ length v)
           else zipWith (-) u v

negL :: Floating a => [a] -> [a]
negL = map negate

(*.) :: Floating a => a -> [a] -> [a] -- multiply by a constant
(*.) c v = map ((*) c) v

(/.) :: Floating a => [a] -> a -> [a]
(/.) v c = map ((/) c) v

p2v (x, y) = [x, y]

v2p [x, y] = (x, y)

infixl 6 +:, -:
infixl 7 *:, /:

(+:) :: Floating a => (a, a) -> (a, a) -> (a, a)
(+:) (x, y) (c, d) = (x + c, y + d)

(-:) :: Floating a => (a, a) -> (a, a) -> (a, a)
(-:) (x, y) (c, d) = (x - c, y - d)

(*:) :: Floating a => a -> (a, a) -> (a, a)
(*:) c (x, y) = (c * x, c * y)

(/:) :: Floating a => (a, a) -> a -> (a, a)
(/:) (x, y) c = (c / x, c / y)

dotv :: Floating a => (a, a) -> (a, a) -> a
dotv (x, y) (a, b) = x * a + y * b

-- Again (see below), we add epsd to avoid NaNs. This is a general problem with using `sqrt`.
norm :: Floating a => [a] -> a
norm v = sqrt ((sum $ map (^ 2) v) + epsd)

normsq :: Floating a => [a] -> a
normsq = sum . map (^ 2)

normalize :: Floating a => [a] -> [a]
normalize v = (1 / norm v) *. v

-- Find the angle between x-axis and a line passing points, reporting in radians
findAngle :: Floating a => (a, a) -> (a, a) -> a
findAngle (x1, y1) (x2, y2) = atan $ (y2 - y1) / (x2 - x1)

midpoint :: Floating a => (a, a) -> (a, a) -> (a, a) -- mid point
midpoint (x1, y1) (x2, y2) = ((x1 + x2) / 2, (y1 + y2) / 2)

midpointV :: Floating a => [a] -> [a] -> [a]
midpointV x y = (x +. y) /. 2.0

-- We add epsd to avoid NaNs in the denominator of the gradient of dist.
-- Now, grad dist (0, 0) (0, 0) is 0 instead of NaN.
dist :: Floating a => (a, a) -> (a, a) -> a -- distance
dist (x1, y1) (x2, y2) = sqrt ((x1 - x2)^2 + (y1 - y2)^2 + epsd)

distsq :: Floating a => (a, a) -> (a, a) -> a -- distance
distsq (x1, y1) (x2, y2) = (x1 - x2)^2 + (y1 - y2)^2

-- Closest point on line segment to a point
-- https://stackoverflow.com/questions/849211/shortest-distance-between-a-point-and-a-line-segment
closestpt_pt_seg :: (Floating a, Ord a) => (a, a) -> ((a, a), (a, a)) -> (a, a)
closestpt_pt_seg p@(px, py) (v@(vx, vy), w@(wx, wy)) =
                 let eps = 10 ** (-3)
                     lensq = distsq v w in
                 if lensq < eps then v                  -- line seg looks like a point
                 else let dir = w -: v
                          t = ((p -: v) `dotv` dir) / lensq -- project vector onto line seg and normalize
                          t' = clamp (0, 1) t in
                      v +: (t' *: dir) -- walk along vector of line seg

clamp :: (Floating a, Ord a) => (a, a) -> a -> a
clamp (l, r) x = max l $ min r x

--------------------------------------
-- Reflection capabilities to typecheck Computation functions
-- Typeable doesn't works with polymorphism (e.g. `id`) but works with `Floating a` by replacing it with `Double`
-- Adapted https://stackoverflow.com/questions/5144799/reflection-on-inputs-to-a-function-in-haskell
-- TODO unit test for function
-- Unfortunately it's pretty slow...

-- BTW, to check if something has the same type at runtime, can check equality of their typereps (or [Typerep], or [String], or String)
fnTypes :: Typeable a => a -> [TypeRep]
fnTypes x = split (typeOf x)
       where split t = case first tyConName (splitTyConApp t) of
                       (_     ,  []) -> [t] -- not an arrow
                       ("(->)", [x]) -> [x] -- return value type
                       ("(->)",   x) -> let current = init x
                                            next    = last x
                                        in current ++ split next
                       (_     ,   _) -> [t]

fnTypesStr :: Typeable a => a -> [String]
fnTypesStr = map show . fnTypes

-- If not an arrow type, then first list is empty (as expected)
inputsOutput :: Typeable a => a -> ([TypeRep], TypeRep)
inputsOutput x = let res = fnTypes x in
                 (init res, last res)

inputsOutputStr :: Typeable a => a -> ([String], String)
inputsOutputStr x = let (args, val) = inputsOutput x in
                    (map show args, show val)

-- other helpers

rot90 :: Floating a => (a, a) -> (a, a)
rot90 (x,y) = (-y,x)

lerpP :: Floating a => (a, a) -> (a, a) -> a -> (a, a)
lerpP (a,b) (c,d) k = let 
    lerpNum a b = a*(1.0-k) + b*k
    in (lerpNum a c, lerpNum b d)

mag :: Floating a => (a, a) -> a
mag (a,b) = sqrt $ magsq (a,b)

magsq :: Floating a => (a, a) -> a
magsq (a,b) = a**2 + b**2

normalize' :: Floating a => (a, a) -> (a, a)
normalize' (a,b) = let l = mag (a,b) in (a/l, b/l)

scaleP :: Floating a => a -> (a, a) -> (a, a)
scaleP k (a,b) = (a*k, b*k)

translate2 :: Floating a => (a, a) -> [(a, a)] -> [(a, a)]
translate2 v pts = map (+: v) pts

map2 :: (a -> b) -> (a, a) -> (b, b)
map2 f (a, b) = (f a, f b)

rotateList :: [a] -> [a]
rotateList l = take (length l) $ drop 1 (cycle l)

-- | Scale a value x in [lower, upper] linearly to x' lying in range [lower', upper'].
-- (Allow reverse lerping, i.e. upper < lower)
scaleLinear :: Autofloat a => a -> Pt2 a -> Pt2 a -> a
scaleLinear x (lower, upper) (lower', upper') = 
            if x < lower || x > upper 
            then trace ("invalid value " ++ show x ++ " to range " ++ show (lower, upper)) x
                 -- Values may be wrong in the intermediate stage due to optimization, so maybe clamp to range anyway
            else let (range, range') = (upper - lower, upper' - lower')
                 in ((x - lower) / range) * range' + lower'

-- n = number of interpolation points (not counting endpoints)
-- so with n = 1, we would have [x1, (x1+x2/2), x2]
lerp :: Autofloat a => a -> a -> Int -> [a]
lerp x1 x2 n = let dx = (x2 - x1) / (fromIntegral n + 1) in
               take (n + 2) $ iterate (+ dx) x1

lerp2 :: Autofloat a => Pt2 a -> Pt2 a -> Int -> [Pt2 a]
lerp2 (x1, y1) (x2, y2) n = let xs = lerp x1 x2 n 
                                ys = lerp y1 y2 n
                            in zip xs ys -- Interp both simultaneously
                            -- Interp the first, then the second
                            -- in zip xs (repeat y1) ++ zip (repeat x2) ys

cross :: Autofloat a => [a] -> [a] -> [a]
cross [x1, y1, z1] [x2, y2, z2] = [ y1 * z2   -   y2 * z1,
                                    x2 * z1   -   x1 * z2,
                                    x1 * y2   -   x2 * y1 ]

angleBetweenRad :: Autofloat a => [a] -> [a] -> a -- Radians
angleBetweenRad p q = acos ((p `dotL` q) / (norm p * norm q + epsd))

-- n is the "original" plane normal for pq
-- https://stackoverflow.com/questions/5188561/signed-angle-between-two-3d-vectors-with-same-origin-within-the-same-plane
angleBetweenSigned :: Autofloat a => [a] -> [a] -> [a] -> a -- Radians
angleBetweenSigned n p q = let sign = -1 * signum ((p `cross` q) `dotL` n) in
                           sign * angleBetweenRad p q

-- Unit rotation in the plane of the basis vectors (e1, e2) to time t (arc length)
-- Used for spherical geometry.
-- NOTE: expects e1 and e2 to be unit vectors
circPtInPlane :: Autofloat a => [a] -> [a] -> a -> [a]
circPtInPlane e1 e2 t = cos t *. e1 +. sin t *. e2

-- Assuming unit circle and unit basis vectors. Arc starts at e1.
slerp :: Autofloat a => Int -> a -> a -> [a] -> [a] -> [[a]]
slerp n t0 t1 e1 e2 = let dt = (t1 - t0) / (fromIntegral n + 1)
                          ts = take (n + 2) $ iterate (+ dt) t0
                       in map (circPtInPlane e1 e2) ts -- Travel along the arc

-- Project q onto p, without normalization
proj :: Autofloat a => [a] -> [a] -> [a]
proj p q = let unit_p = normalize p
           in (q `dotL` unit_p) *. unit_p