{-# LANGUAGE ScopedTypeVariables, TypeOperators #-}
import Prelude ()
import Control.Category.Constrained.Prelude
import Control.Arrow.Constrained

import Data.Manifold
import Data.Manifold.TreeCover
import Data.Manifold.Web
import Data.Manifold.Shade
import Data.Manifold.Types
import Data.Function.Affine
import Math.LinearMap.Category
import Math.LinearMap.Category.Derivatives
import Data.VectorSpace
import Linear (V2(..), _x, _y)
import Control.Lens
:opt no-lint

test x = x`seq`putStrLn "Ok"

import Graphics.Dynamic.Plot.R2
import Data.Colour.Names
import Data.Colour.Manifold
import Diagrams.Prelude ((^&))

colourscheme :: Shade' ℝ -> Shade (Colour ℝ)
colourscheme = f `seq` \(Shade' y e) -> f (Shade y $ dualNorm e)
 where Just (f :: Shade ℝ -> Shade (Colour ℝ))
          = rangeWithinVertices (1, maxc)
                              [ (0, black) ]
       Just maxc = toInterior (grey :: Colour ℝ)
test colourscheme

colourscheme_nabla :: Shade' (V2 ℝ+>ℝ) -> Shade (Colour ℝ)
colourscheme_nabla = f `seq` \(Shade' y e) -> f (Shade y $ dualNorm e)
 where Just (f :: Shade (V2 ℝ+>ℝ) -> Shade (Colour ℝ))
          = rangeWithinVertices (zeroV, neutral)
                              [ (V2 2 0-+|>1, red)
                              , (V2 0 2-+|>1, green) ]
       Just neutral = toInterior (grey :: Colour ℝ)
test colourscheme_nabla

colourscheme_laplace :: Shade' (V2 ℝ⊗〃+>ℝ) -> Shade (Colour ℝ)
colourscheme_laplace = f `seq` \(Shade' y e) -> f (Shade y $ dualNorm e)
 where Just (f :: Shade (V2 ℝ⊗〃+>ℝ) -> Shade (Colour ℝ))
          = rangeWithinVertices (zeroV, neutral)
                              [ ((squareV (V2 1 0)^+^squareV (V2 0 1)) -+|> 1, white)
                              , ((squareV (V2 10 10)^-^squareV (V2 10 (-10))) -+|> 1, red)
                              , ((squareV (V2 10 0)^-^squareV (V2 0 10)) -+|> 1, blue) ]
       Just neutral = toInterior (grey :: Colour ℝ)
test colourscheme_laplace

import Data.Manifold.Function.LocalModel

osc :: PointsWeb ℝ (Shade' ℝ)
osc = fromWebNodes euclideanMetric
           $ [(x, sin x|±|[0.001]) | x<-[0,0.2 .. pi]]
fst <$> (localModels_CGrid osc :: [(ℝ, QuadraticModel ℝ ℝ)])

plotWindow [ plot osc  & legendName "sin"
           , plot (linIsoTransformShade (lfun ($ 1)) <$> differentiateUncertainWebFunction osc)
                       & legendName "sin' (nodes)"
           , plot (linIsoTransformShade (lfun ($ squareV 1)) <$> differentiate²UncertainWebFunction osc)
                       & legendName "sin'' (nodes)"
           , plot (fromWebNodes euclideanMetric
               $ second (linIsoTransformShade (lfun ($ 1)) . fst . snd . quadraticModel_derivatives)
                              <$> localModels_CGrid osc )
                       & legendName "sin' (links)"
           , plot (fromWebNodes euclideanMetric
               $ second (linIsoTransformShade (lfun ($ squareV 1)) . snd . snd . quadraticModel_derivatives)
                              <$> localModels_CGrid osc )
                       & legendName "sin'' (links)"
           , continFnPlot (negate . sin)
                       & legendName "-sin"
           ]

derivatives1Plot, derivatives2Plot :: PointsWeb ℝ (Shade' ℝ²) -> [DynamicPlottable]
derivatives1Plot signal
         = [ plot [ projectShade(lensEmbedding _x) . fst <$>signalLM
                  , projectShade(lensEmbedding _y) . fst <$>signalLM :: PointsWeb ℝ (Shade' ℝ) ]
           , plot [ projectShade(lensEmbedding (1*∂_x/∂id)) . fst . snd <$>signalLM
                  , projectShade(lensEmbedding (1*∂_y/∂id)) . fst . snd <$>signalLM :: PointsWeb ℝ (Shade' ℝ) ]
                  & legendName "∂"
           ]
 where signalLM = fmap quadraticModel_derivatives . fromWebNodes euclideanMetric $ localModels_CGrid signal
derivatives2Plot signal
         = [ plot [ projectShade(lensEmbedding _x)<$>signal
                  , projectShade(lensEmbedding _y)<$>signal :: PointsWeb ℝ (Shade' ℝ) ]
           , plot [ projectShade(lensEmbedding (1*∂_x/∂id)) . fst . snd <$>signalLM
                  , projectShade(lensEmbedding (1*∂_y/∂id)) . fst . snd <$>signalLM :: PointsWeb ℝ (Shade' ℝ) ]
                  & legendName "∂"
           , plot [ projectShade(lensEmbedding (1*∂_x/∂id.∂id)) . snd . snd <$>signalLM
                  , projectShade(lensEmbedding (1*∂_y/∂id.∂id)) . snd . snd <$>signalLM :: PointsWeb ℝ (Shade' ℝ) ]
                  & legendName "∂²"
           ]
 where signalLM = fmap quadraticModel_derivatives . fromWebNodes euclideanMetric $ localModels_CGrid signal

twinSignal :: PointsWeb ℝ (Shade' ℝ²)
twinSignal = fromWebNodes euclideanMetric
               $ [(x, sin (x+sin x) ^& cos (x-cos x)
                      |±|[2e-3^&0, 0^&(5e-3*(x+0.1))]) | x<-[0,0.2 .. 2*pi]]

plotWindow $ derivatives2Plot twinSignal

jumpFunction :: PointsWeb ℝ (Shade' ℝ²)
jumpFunction = fromWebNodes euclideanMetric
               $ [(x, signum (x-1) ^& signum (x-2)
                      |±|[1e-2^&0, 0^&1e-2]) | x<-[0,0.3 .. 3]]
plotWindow $ derivatives1Plot jumpFunction

gaußianPeak :: [ℝ²] -> ℝ -> PointsWeb ℝ² (Shade' ℝ)
gaußianPeak ps δy = fromWebNodes euclideanMetric
       [ (xy, exp (-2*magnitudeSq xy) |±|[δy * if x > 0 then 10 else 1]) | xy@(V2 x _) <- ps ]

δ = 0.01
grid = [V2 x y | x<-[-2,-1.8..2], y<-[-1,-0.8..2]]
hexGrid = [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]]

forM_ [grid{-, hexGrid-}] $ \g -> do
 let gaußianPeak_loc = fmap quadraticModel_derivatives
                 . fromWebNodes euclideanMetric
                 . localModels_CGrid $ gaußianPeak g δ

 plotWindow [plot . fmap (colourscheme . fst) $ gaußianPeak_loc, dynamicAxes]
 plotWindow [plot . fmap (colourscheme_nabla . fst . snd) $ gaußianPeak_loc, dynamicAxes]
 plotWindow [plot . fmap (colourscheme_laplace . snd . snd) $ gaußianPeak_loc, dynamicAxes]

plotWindow [plot . fmap colourscheme $ gaußianPeak grid δ ]
plotWindow [plot . fmap colourscheme_nabla . differentiateUncertainWebFunction
              $ gaußianPeak grid δ ]
plotWindow [plot . fmap (colourscheme_nabla . projectShade (lensEmbedding (1*∂id/∂_x)))
   . differentiateUncertainWebFunction . differentiateUncertainWebFunction
       $ gaußianPeak grid δ ]
plotWindow [plot . fmap (colourscheme_nabla . projectShade (lensEmbedding ((1*∂id/∂_x)/∂id)))
   . differentiateUncertainWebFunction . differentiateUncertainWebFunction
       $ gaußianPeak grid δ ]