{-# OPTIONS -Wall #-}

module AnimateFunction where

-------------------------------------------------------------------------
-- Make an SVG animation of a function of x and t.                     --
-- The main functions to use are animateFunction and animateFunctions. --
-------------------------------------------------------------------------

animateFunction :: FilePath -> ((R,R) -> R) -> (R,R,R) -> (R,R,R) -> IO ()
animateFunction filename f (xMin,dx,xMax) (tMin,dt,tMax)
    = let xs = [xMin, xMin + dx .. xMax]
          ts = [tMin, tMin + dt .. tMax]
          xyPairs :: R -> [(R,R)]
          xyPairs t = [(x, f (x,t)) | x <- xs]
          ys = concat [map snd (xyPairs t) | t <- ts]
          yMin = minimum ys
          yMax = maximum ys
          scaledPairs t = scalePairs xMin xMax yMin yMax (xyPairs t)
          svgStart =    "<?xml version=\"1.0\" standalone=\"no\"?>\n"
                     ++ "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n"
                     ++ "  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n"
                     ++ "<svg width=\"100%\" height=\"100%\" viewBox=\"-1.1 -1.1 2.2 2.2\" preserveAspectRatio=\"none\"\n"
                     ++ "     xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n"
          svgEnd   =    "</svg>\n"
          matter = concat [oneCurve t dt (0,0,255) (scaledPairs t) | t <- ts]
      in writeFile filename $ svgStart ++ matter ++ svgEnd

animateFunctions :: FilePath -> [((R,R) -> R,(Int,Int,Int))] -> (R,R,R) -> (R,R,R) -> IO ()
animateFunctions filename prs (xMin,dx,xMax) (tMin,dt,tMax)
    = let (fs,cs) = unzip prs
          anims = makeAnimations fs (xMin,dx,xMax) (tMin,dt,tMax)
          (yMin,yMax) = animationsRange anims
          scaledAnims = scaleAnimations xMin xMax yMin yMax anims
          scaledACs = zip scaledAnims cs
          svgStart =    "<?xml version=\"1.0\" standalone=\"no\"?>\n"
                     ++ "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n"
                     ++ "  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n"
                     ++ "<svg width=\"100%\" height=\"100%\" viewBox=\"-1.1 -1.1 2.2 2.2\" preserveAspectRatio=\"none\"\n"
                     ++ "     xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n"
          svgEnd   =    "</svg>\n"
          matter = animationsMatter scaledACs
      in writeFile filename $ svgStart ++ matter ++ svgEnd

--------------
-- Examples --
--------------

-- traveling wave
anim01 :: IO ()
anim01 = let k = 1
             omega = 1
         in animateFunction "anim01.svg" (\(x,t) -> cos (k * x - omega * t)) (0,0.1,10) (0,0.1,10)

-- standing wave
anim02 :: IO ()
anim02 = let k = 3 * pi / 10
             omega = 1
         in animateFunction "anim02.svg" (\(x,t) -> sin (k * x) * cos (omega * t)) (0,0.1,10) (0,0.1,10)

-- triangle pulse
anim03 :: IO ()
anim03 = let k = 1
             omega = 1
             triangle x
                 | x < 0      = 0
                 | x < 0.5    = 2 * x
                 | x < 1      = 2 - 2 * x
                 | otherwise  = 0
         in animateFunction "anim03.svg" (\(x,t) -> triangle (k * x - omega * t)) (0,0.1,10) (0,0.1,10)

-- Gaussian pulse
anim04 :: IO ()
anim04 = let k = 1
             v = 1
         in animateFunction "anim04.svg" (\(x,t) -> 1000 * exp (-k * (x - v * t)**2) - 300) (-5,0.1,5) (0,0.1,10)

-- narrowing Gaussian
anim05 :: IO ()
anim05 = let a t = 6 - t
         in  animateFunction "anim05.svg" (\(x,t) -> exp (-(x/a t)**2) / (a t * sqrt pi)) (-2,0.1,2) (1,0.1,5)

-- standing wave of E and B
anim06 :: IO ()
anim06 = animateFunctions "anim06.svg" [(\(x,t) -> sin (x - t) + sin (x + t),(0,0,255))
                                       ,(\(x,t) -> sin (x - t) - sin (x + t),(255,0,0))] (0,0.1,3*pi) (0,0.1,10)

---------------------
-- Supporting code --
---------------------

type R = Double

lineToText :: (R,R) -> String
lineToText (x,y) = "L " ++ show x ++ "," ++ show (-y) ++ " "

framePathText :: [(R,R)] -> String
framePathText xys = 'M' : tail (concatMap lineToText xys)

oneCurve :: R -> R -> (Int,Int,Int) -> [(R,R)] -> String
oneCurve t dt rgb xys
    = let fpt = framePathText xys
          pre = "<g stroke-linecap=\"butt\" fill=\"rgb(0,0,0)\" stroke-miterlimit=\"10.0\" stroke-width=\"0.005\" fill-opacity=\"0.0\" stroke-opacity=\"1.0\" stroke-linejoin=\"miter\" stroke=\"rgb" ++ show rgb ++ "\">"
          middle = "<path d=\"" ++ fpt ++ "\" visibility=\"hidden\"> <set attributeName=\"visibility\" attributeType=\"CSS\" to=\"visible\" begin=\"" ++ show t ++ "s\" dur=\"" ++ show dt ++ "s\" /> </path>\n"
          post = "</g>"
      in pre ++ middle ++ post

-- scale to the range -1 to 1
scale :: R -> R -> R -> R
scale xmin xmax x = 2 * (x - xmin) / (xmax - xmin) - 1

scalePair :: R -> R -> R -> R -> (R,R) -> (R,R)
scalePair xmin xmax ymin ymax (x,y)
    = (scale xmin xmax x, scale ymin ymax y)

scalePairs :: R -> R -> R -> R -> [(R,R)] -> [(R,R)]
scalePairs xmin xmax ymin ymax
    = map $ scalePair xmin xmax ymin ymax

type Curve = [(R,R)]  -- list of xy pairs
type Animation = ([(Curve,R)],R)  -- ([(curve,t)],dt)

curveRange :: [(R,R)] -> (R,R)
curveRange xys = let ys = map snd xys
                 in (minimum ys, maximum ys)

animationRange :: Animation -> (R,R)
animationRange (cts,_) = curveRange $ concatMap fst cts

animationsRange :: [Animation] -> (R,R)
animationsRange anims
    = let ranges = map animationRange anims
          (yMins,yMaxs) = unzip ranges
      in (minimum yMins, maximum yMaxs)

makeCurve :: ((R,R) -> R) -> (R,R,R) -> R -> Curve
makeCurve f (xMin,dx,xMax) t
    = let xs = [xMin, xMin + dx .. xMax]
      in [(x, f (x,t)) | x <- xs]

makeAnimation :: ((R,R) -> R) -> (R,R,R) -> (R,R,R) -> Animation
makeAnimation f (xMin,dx,xMax) (tMin,dt,tMax)
    = let ts = [tMin, tMin + dt .. tMax]
      in ([(makeCurve f (xMin,dx,xMax) t,t) | t <- ts],dt)

makeAnimations :: [(R,R) -> R] -> (R,R,R) -> (R,R,R) -> [Animation]
makeAnimations fs (xMin,dx,xMax) (tMin,dt,tMax)
    = map (\f -> makeAnimation f (xMin,dx,xMax) (tMin,dt,tMax)) fs

scaleCurve :: R -> R -> R -> R -> Curve -> Curve
scaleCurve = scalePairs

scaleAnimation :: R -> R -> R -> R -> Animation -> Animation
scaleAnimation xmin xmax ymin ymax (prs,dt)
    = let newPrs = map (\(c,t) -> (scaleCurve xmin xmax ymin ymax c, t)) prs
      in (newPrs,dt)

scaleAnimations :: R -> R -> R -> R -> [Animation] -> [Animation]
scaleAnimations xmin xmax ymin ymax
    = map (scaleAnimation xmin xmax ymin ymax)

-- use on a scaled animation
animationMatter :: Animation -> (Int,Int,Int) -> String
animationMatter (prs,dt) rgb
    = concat [oneCurve t dt rgb c | (c,t) <- prs]

animationsMatter :: [(Animation,(Int,Int,Int))] -> String
animationsMatter
    = concatMap (\(a,c) -> animationMatter a c)