-- | Display function graphs
{-# LANGUAGE ScopedTypeVariables #-}
module Graphics.Diagrams.FunctionGraphs where

import Graphics.Diagrams

import Data.List (zip5)

-------------------------

-- | Make a function discrete
dotty :: (Num a, Real b) => (a -> b) -> (Double -> Double)
dotty f x
    | abs (fromInteger rx - x) < 0.02 = realToFrac $ f $ fromInteger rx
    | otherwise = sqrt (-1)
  where
    rx = round x

-- | Draw coordinate system
coords :: Point -> Point -> Diagram
coords (a1,a2) (b1,b2) = union 
    [ (a1,0) ~~ (b1,0)
    , (b1-1,0.5) ~~ (b1,0)
    , (b1-1,-0.5) ~~ (b1,0)
    , (0.5,b2-1) ~~ (0,b2)
    , (-0.5,b2-1) ~~ (0,b2)
    , (0,a2) ~~ (0,b2)
    , (1,-0.5) ~~ (1,0.5)
    ]

-- | Display a function
showFun :: (Fractional a, Real b) => Bool -> Point -> Point -> (a -> b) -> Diagram
showFun dotty a@(a1,a2) b@(b1,b2) f 
    = clip a b $ rectangle a b `fill` black <|> coords a b `stroke` yellow 
        <|> union rr `stroke` white `fill` white
 where
    l = map (\x->(x, realToFrac $ f $ realToFrac x)) [a1 - 0.1, a1 .. b1 + 0.1]
    ds = zipWith (-) (map snd l) (map snd $ drop 1 l)

    rr = if dotty then 
            [circle 0.1 `move` (x, y) | x <- map fromIntegral [round a1 .. round b1 :: Integer]
                , let y = realToFrac $ f $ realToFrac x, y <= b2 && y >= a2]
        else lines ++ dots

    lines = [e >-< f | (e,True,f) <- zip3 (drop 1 l) extralines (drop 2 l)]
    dots = [circle 0.1 `move` b | (x,b,y)
            <- zip3 extralines (drop 2 l) (drop 1 extralines)
            , not (x && y), snd b <= b2 && snd b >= a2]
    extralines = [abs y< (50* abs x) `max` 0.1 && abs y< (50* abs z) `max` 0.1 &&
                 (x * y >= 0 || y * z >= 0)
                 | (x,e,y,f,z)<-zip5 ds (drop 1 l) (drop 1 ds) (drop 2 l) (drop 2 ds)]


-- showFun a b f = showArc (fst a, fst b) a b (\t -> (t, f t))

-- | Display an arc
showArc :: (Fractional a, Real b, Real c) => 
    (Double,Double) -> Point -> Point -> (a -> (b, c)) -> Diagram
showArc (k1,k2) a@(a1,a2) b@(b1,b2) f
    = clip a b $ rectangle a b `fill` black <|> coords a b `stroke` yellow 
        <|> union (lines ++ dots) `stroke` white `fill` white
 where
    l = map (\(x,y)-> (realToFrac x, realToFrac y)) 
                $ map (f . realToFrac) [k1-0.1, k1.. k2+0.1]
    ds = zipWith minu l (drop 1 l)

    (a,b) `minu` (c,d) = (c-a, d-b)
    abs (a,b) = sqrt (a*a+b*b)
    (a,b) <<< (c,d) = a <= c && b <= d
    (a,b) .* (c,d) = a*c + b*d

    lines = [e >-< f | (e,True,f) <- zip3 (drop 1 l) extralines (drop 2 l)]
    dots = [circle 0.1 `move` z | (x,z,y)
            <- zip3 extralines (drop 2 l) (drop 1 extralines)
            , not (x && y), a <<< z && z <<< b]
    extralines = [abs y< (10* abs x) `max` 0.1 && abs y< (10* abs z) `max` 0.1 &&
                 (x .* y >= 0 || y .* z >= 0)
                 | (x,e,y,f,z)<-zip5 ds (drop 1 l) (drop 1 ds) (drop 2 l) (drop 2 ds)]