{-# LANGUAGE FlexibleContexts, GADTs, TypeOperators #-}
import Data.Manifold.TreeCover
import Data.Manifold.Web
import Data.Manifold.Web.Internal
import qualified Data.Graph as Graph
import Data.Manifold.Types
import Data.VectorSpace
import Linear.V2
import Data.AffineSpace
import Math.LinearMap.Category
import Data.Random
import Data.Random.Manifold

import Control.Arrow
import Control.Applicative (empty)
import Data.Foldable (toList, forM_)
import Data.Traversable (forM)

import Graphics.Dynamic.Plot.R2

import Diagrams.Prelude ( Point(P), r2, circle, (&), (#), (^.), (^&), _1, moveTo
                        , fc, lc, opacity, red, blue, white
                        , fromVertices )
import Diagrams.CubicSpline

prettyTreePlot :: ShadeTree ℝ² -> [DynamicPlottable]
prettyTreePlot tr = [ plot [ shapePlot $ circle 0.03 & moveTo p & opacity 0.2 | p <- ps ]
                    , plot $ onlyNodes tr ]
                 ++ [unitAspect]
 where ps = map P $ onlyLeaves_ tr
 
prettyWebPlot :: PointsWeb ℝ² () -> [DynamicPlottable]
prettyWebPlot w = [ shapePlot $ cubicSpline False [P r₁, P m, P r₂]
                  | ((r₁@(V2 x₁ y₁),()),(r₂@(V2 x₂ y₂),())) <- edg
                  , let m = V2 ((x₁+2*x₂)/3 + (y₂-y₁)/19)
                               ((y₁+2*y₂)/3 + (x₁-x₂)/19) ]
 where edg = map (gnodes *** gnodes) $ Graph.edges graph
       (graph, gnodes) = toGraph w

plotTreeAndWeb :: ShadeTree ℝ² -> IO ()
plotTreeAndWeb tr = do
    plotWindow $ prettyTreePlot tr # map (tint white)
              ++ prettyWebPlot web # map (tweakPrerendered $ opacity 0.7)
              ++ [plot [ plot [ plot p & tweakPrerendered (opacity 0.3)
                              , shapePlot (circle 0.03 & moveTo (P pp)) ]
                          & tint red
                       | (p@(Cutplane pp _),_) <- webBoundary web ]]
    return ()
 where web = fromShadeTree euclideanMetric tr

cartTree :: ShadeTree ℝ²
cartTree = fromLeafPoints
              [ (x^&y) | x<-[0, 0.15 .. 4]
                       , y<-[0, 0.2 .. 4]
                       , (x-2)^2 + (y-2)^2 < 4 ]

plotTreeAndWeb cartTree

honeycombTree = fromLeafPoints $
           [V2 x y | x<-[-2,-1..6], y<-[-1, 0.8..6]]
        ++ [V2 x y | x<-[-1.5,-0.5..6], y<-[-0.7, 1.1 ..6]]
        ++ [V2 x y | x<-[-1.5,-0.5..6], y<-[-0.1, 1.7 ..6]]
        ++ [V2 x y | x<-[-2,-1..6], y<-[0.2, 2.0..6]]

plotTreeAndWeb honeycombTree

hexagonTree = fromLeafPoints $
           [V2 x y | x<-[-2,-1.7..2], y<-[-1,-0.6..2]]
        ++ [V2 x y | x<-[-1.85,-1.55..2], y<-[-0.8,-0.4..2]]

plotTreeAndWeb hexagonTree

scatterTree = let
    tps₀ = [(0,0), (0,1), (1,1), (1,2), (2,2)]
    tps₁ = [p .+^ v^/3 | p<-tps₀, v <- [(0,0), (-1,1), (1,2)]]
    tps₂ = [p .+^ v^/4 | p<-tps₁, v <- [(0,0), (-1,1), (1,2)]]
    tps₃ = [p .+^ v^/5 | p<-tps₂, v <- [(0,0), (-2,1), (1,2)]]
    tps₄ = [p .+^ v^/7 | p<-tps₃, v <- [(0,1), (-1,1), (1,2)]]
     in fromLeafPoints $ r2<$>tps₄
plotTreeAndWeb scatterTree
((_,exampleTwiglet), exampleTwigEnviron) = twigsWithEnvirons scatterTree !! 5
plotWindow $ prettyTreePlot scatterTree
                 # map (tweakPrerendered (opacity 0.3) . tint white)
          ++ prettyTreePlot exampleTwiglet
                 # map (tint blue)
          ++ concat (prettyTreePlot . snd <$> exampleTwigEnviron)
                 # map (tweakPrerendered (opacity 0.3) . tint red)

randomTr <- fmap fromLeafPoints . forM [0..1000] $ \_->
     runRVar (sample $ (1^&1):±[1^&0, 0^&1]) StdRandom :: IO ℝ²

plotTreeAndWeb randomTr

let web = fromShadeTree euclideanMetric randomTr :: PointsWeb ℝ² ()
 in plotWindow [
      plotLatest [ plotMultiple
           [ lineSegPlot [(x,y) | V2 x y<-path]
           | path <- pathsTowards i (localFmapWeb _thisNodeCoord web) ]
          & plotDelay 2
       | i <- [0..] ]
     , plot (prettyWebPlot web) & tweakPrerendered (opacity 0.2)
     ]