plot.hs

{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE FlexibleContexts          #-}
import Diagrams.Prelude hiding (radius)
import Diagrams.Backend.PGF.CmdLine

import Data.List

--f x | x < -1    = x * sin x
--    | x <  0    = 5 * x
--    | x <  2    = x * (x + 2) * (x - 2)
--    | otherwise = 1
--specialf d = [ jumpLeft f (-1), (Corner,0), jumpLeft f 2 ] -- f

--g x | x < -1    = ((x + 2)^2+1)/2
--    | x <  2    = fromIntegral $ floor (-sin x) + 1
--    | x == 2    = 4
--    | otherwise = x*(x-2)*(x-3)
--specialg d = [ (Corner,-1), jumpRight g 0, (Removable (g (2-eps)),2) ] -- g

--h x | x < -2    = x * sin x + (2 + 2*sin(-2))
--    | x == -2   = 3
--    | x <  0    = -(x+2)
--    | otherwise = cos x
--specialh d = [ jumpRem h (-2), jumpLeft h 0 ]

-- s t = (64-16*(t-1)^2) / 6

f x | x <  -2   = 5 * cos x
    | x <  -1   = -3*x-4
    | x <=  1   = -abs (x+1/2) - (1/2)
    | otherwise = sin (-4*x)
specialf d = [jumpLeft f (-2), jumpRight f 1]


--integrate eps f l h =
--  where
--    w = h - l
--    n = w / eps
--    is = [i | i <- [0..n-1]]


eps = 0.0001
radius = 0.10 :: Double


data Special
    = JumpLeft Double
    | JumpRight Double
    | JumpRem Double Double
    | Removable Double
    | Corner
    | Stop
    | Hole (Double,Double)
    deriving (Show, Eq, Read)

jumpLeft :: (Double -> Double) -> Double -> (Special, Double)
jumpLeft f x = (JumpLeft (f (x - eps)), x)

jumpRight :: (Double -> Double) -> Double -> (Special, Double)
jumpRight f x = (JumpRight (f (x + eps)), x)

jumpRem f x = (JumpRem (f (x-eps)) (f (x+eps)), x)

-- Search for special points.
--special :: (Double -> Double) -> (Double, Double) -> [(Special, Double)]

getThree :: (Double,Double) -> Double -> Double -> (P2 Double,P2 Double,P2 Double)
getThree p@(x,y) y' y'' = (p2 p, p2 (x,y'), p2 (x,y''))

getTwo :: (Double,Double) -> Double -> (P2 Double,P2 Double)
getTwo p@(x,y) y' = (p2 p, p2 (x,y'))

getOne :: (Double,Double)           -> P2 Double
getOne = p2

hole, closed :: Diagram B -> Diagram B
closed = fc black
hole = lwG (radius/2)


plotSpecial :: Special -> (Double, Double) -> Diagram B
plotSpecial (JumpLeft  y) p = circle radius `place` p' # closed
                           <> circle radius `place` q' # hole
                        where (p',q') = getTwo p y
plotSpecial (JumpRight y) p = circle radius `place` p' # closed
                           <> circle radius `place` q' # hole
                        where (p',q') = getTwo p y
plotSpecial (JumpRem y y') p = circle radius `place` p'  # closed
                            <> circle radius `place` q'  # hole
                            <> circle radius `place` q'' # hole
                        where (p',q',q'') = getThree p y y'
plotSpecial (Removable y) p = circle radius `place` p' # closed
                           <> circle radius `place` q' # hole
                        where (p',q') = getTwo p y
plotSpecial Corner   p = mempty --  circle radius `place` (getOne p) # closed
plotSpecial Stop     p = circle radius `place` (getOne p) # closed
plotSpecial (Hole p) _ = circle radius `place` (getOne p) # hole

plotCont :: (Double -> Double) -> (Double, Double) -> (Double, Double) -> Diagram B
plotCont f (xl,xh) (yl,yh)
    | w < 0     = mempty
    | otherwise = cubicSpline False ps # sty
  where
    sty = lwG (radius / 2)
    w = xh - xl -- - radius/2
    f' = dfdx f

    ps = [ p2 (x, f x)
         | let al = atan (f' (xl + eps))
               ah = atan (f' (xh - eps))
               xl' = xl + radius * cos al
               xh' = xh - radius * cos ah
          , x <- step xl' (w/128) xh'
          ]

    step l s h = go l
      where
        go x | x - h > 0     = [h]
             | h - x < (s/2) = [h]
             | otherwise     = x : go (x + s)

glueAll :: (Floating n, Ord n, Metric v) => [Located (Trail' Line v n)] -> Located (Trail' Loop v n)
glueAll [] = error "Need at least one line"
glueAll ts@(t:_) = (glueLine . lineFromSegments . concatMap (lineSegments . unLoc) $ ts) `at` loc t

shadeBetween :: (Double -> Double) -> (Double -> Double) -> (Double, Double) -> (Double, Double) -> Diagram B
shadeBetween f g (xl,xh) (yl,yh)
    | w < 0     = mempty
    | otherwise = sty . stroke . toPath $ glueAll
                    [ cubicSpline False (ps f)
                    , (xh ^& f xh) ~~ (xh ^& g xh)
                    , reverseLocLine $ cubicSpline False (ps g)
                    , (xl ^& g xl) ~~ (xl ^& f xl)
                    ]
  where
    sty = lw none . fc silver -- lwG (radius / 2)
    w = xh - xl

    ps f = [ p2 (x, f x)
           | x <- step xl (w/128) xh
           ]

    step l s h = go l
      where
        go x | x - h > 0     = [h]
             | h - x < (s/2) = [h]
             | otherwise     = x : go (x + s)

shade :: (Double -> Double) -> (Double, Double) -> (Double, Double) -> Diagram B
shade f (xl,xh) (yl,yh)
    | w < 0     = mempty
    | otherwise = sty . stroke . toPath $ glueAll
                    [ cubicSpline False ps
                    , (xh ^& f xh) ~~ (xh ^& 0   )
                    , (xh ^& 0   ) ~~ (xl ^& 0   )
                    , (xl ^& 0   ) ~~ (xl ^& f xl)
                    ]
  where
    sty = lw none . fc silver -- lwG (radius / 2)
    w = xh - xl

    ps = [ p2 (x, f x)
         | x <- step xl (w/128) xh
         ]

    step l s h = go l
      where
        go x | x - h > 0     = [h]
             | h - x < (s/2) = [h]
             | otherwise     = x : go (x + s)

-- plot :: (Double -> Double) -> (Double, Double) -> (Double, Double) -> Diagram B
plot arrows f special d@(xl,xh) r@(yl,lh) = as <> mconcat specials <> mconcat ps
  where
    ss = special d
    cs = zip<*>tail $ nub (xl : map snd ss ++ [xh])
    ps = [plotCont f c r | c <- cs]

    specials = [plotSpecial s p | (s,x) <- ss, let p = (x, f x)]

    as | arrows    = arrowl f xl <> arrowr f xh
       | otherwise = mempty

arrowl, arrowr :: (Double -> Double) -> Double -> Diagram B
arrowl = arrowhead (triangleHead # rotate ( 90 @@ deg))
arrowr = arrowhead (triangleHead # rotate (-90 @@ deg))

triangleHead = triangle (radius*2) # fc black # translateY (-radius) # scaleX 0.5

arrowhead :: Diagram B -> (Double -> Double) -> Double -> Diagram B
arrowhead s f x = s # fc black # rotate (a @@ rad) # moveTo (p2 (x, f x))
  where
    a = atan (f' x)
    f' = dfdx f

dfdx f x = (f (x + 0.0000001) - f x) / 0.0000001


axes :: (Double, Double) -> (Double, Double) -> Diagram B
axes (xl,xh) (yl,yh) = v1 ~~ v2 # sty
                    <> h1 ~~ h2 # sty
                    <> mconcat hts
                    <> mconcat vts
                    <> labels # fontSize 8
                    <> as
  where
    v1 = p2 (0,yl)
    v2 = p2 (0,yh)
    h1 = p2 (xl,0)
    h2 = p2 (xh,0)
    sty = lwG (radius / 4)

    r = radius

    as = (triangleHead # rotateBy (1/4)) `place` p2 (xl,0)
      <> (triangleHead # rotateBy (1/4) # reflectX) `place` p2 (xh,0)
      <> (triangleHead # reflectY) `place` p2 (0,yl)
      <> triangleHead `place` p2 (0,yh)

    hts = [p2 (-r,y) ~~ p2 (r,y) # sty | y <- [yl+1..yh-1] \\ [0]]
    vts = [p2 (x,-r) ~~ p2 (x,r) # sty | x <- [xl+1..xh-1] \\ [0]]

    labels = mempty
          --    text  "12" `place` p2 (12, -1.1)
          -- <> text "$x=40-6y^2-y^3$"  `place` p2 ( 8, 3)
          -- <> text "$x=y^3-26y+10$"   `place` p2 (-7,-4)

    labelsS = text  "1" `place` p2 (1, -0.5)
           <> text  "3" `place` p2 (3, -0.5)
           <> text "64" `place` p2 (-0.5, 64/6)
           <> text "32" `place` p2 (-0.5, 32/6)

-- main = mainWith ((plot True f specialf (-2.5,3) (-4,4) <> axes (-4,4) (-4,4)) # centerXY # pad 1.1)
main = mainWith ((plot True f specialf (-5.75,5.75) (-6,6) <> axes (-6,6) (-6,6)) # centerXY # pad 1.1)
-- main = mainWith ((plot g specialg (-6,4)   (-10,10) <> axes (-10,10) (-10,10)) # centerXY # pad 1.1)
-- main = mainWith ((plot h specialh (-10,10) (-10,10) <> axes (-10,10) (-10,10)) # centerXY # pad 1.1)
-- main = mainWith ((plot s (const []) (0,3) (-1,64/6) <> axes (-1,4) (-1,72/6)) # centerXY # pad 1.1)