{-# OPTIONS -Wall #-}

module GradientVectorField where

-- Prerequisites: diagrams-lib and diagrams-svg
-- Do one of the following to install the prerequisites:
-- cabal install --lib diagrams-lib diagrams-svg
-- stack install diagrams-lib diagrams-svg

import LPFPCore
    ( R, Vec, Position, VectorField
    , (^/), xComp, yComp, cart, phiHat, rVF )
import Diagrams.Prelude
    ( Diagram, V2(..), PolyType(..), PolyOrientation(..), PolygonOpts(..)
    , (#), (@@), dims, p2, r2, arrowAt, position, fc, black, white
    , blend, none, lw, rotate, deg, rad, scale, polygon, sinA )
import Diagrams.Backend.SVG

vfSVG :: ((R,R) -> Position)
      -> (Vec -> (R,R))
      -> FilePath     -- file name
      -> R            -- scale factor in units per meter
      -> [(R,R)]      -- positions to use
      -> VectorField
      -> IO ()
vfSVG toPos fromVec fileName unitsPerMeter pts vf
    = let vf2d = r2 . fromVec . (^/ unitsPerMeter) . vf . toPos
          pic  = mconcat [arrowAt (p2 pt) (vf2d pt) | pt <- pts]
      in renderSVG fileName (dims (V2 1024 1024)) pic

vfSVGxy :: FilePath     -- file name
        -> R            -- scale factor
        -> [(R,R)]      -- positions to use
        -> VectorField
        -> IO ()
vfSVGxy = vfSVG (\(x,y) -> cart x y 0) (\v -> (xComp v, yComp v))

phiHatSVG :: IO ()
phiHatSVG
    = vfSVGxy "phiHatSVG.svg" 1
      [(r * cos ph, r * sin ph) | r  <- [1,2]
                                , ph <- [0,pi/4..2*pi]] phiHat

rVFsvg :: IO ()
rVFsvg
    = vfSVGxy "rVFsvg.svg" 2
      [(r * cos ph, r * sin ph) | r  <- [1,2]
                                , ph <- [0,pi/4..2*pi]] rVF

vfGrad :: (R -> R)
       -> ((R,R) -> Position)
       -> (Vec -> (R,R))
       -> FilePath
       -> Int    -- n for n x n
       -> VectorField
       -> IO ()
vfGrad curve toPos fromVec fileName n vf 
  = let step = 2 / fromIntegral n            
        xs = [-1+step/2, -1+3*step/2 .. 1-step/2]
        pts = [(x, y) | x <- xs, y <- xs]
        array = [(pt,magRad $ fromVec $ vf $ toPos pt) | pt <- pts] 
        maxMag = maximum (map (fst . snd) array) 
        scaledArrow m th = scale step $ arrowMagRad (curve (m/maxMag)) th
        pic = position [(p2 pt, scaledArrow m th) | (pt,(m,th)) <- array] 
     in renderSVG fileName (dims (V2 1024 1024)) pic

magRad :: (R,R) -> (R,R)
magRad (x,y) = (sqrt (x*x + y*y), atan2 y x)

-- magnitude from 0 to 1
arrowMagRad :: R  -- magnitude
            -> R  -- angle in radians, counterclockwise from x axis
            -> Diagram B
arrowMagRad mag th
    = let r      = sinA (15 @@ deg) / sinA (60 @@ deg)
          myType = PolyPolar [120 @@ deg, 0 @@ deg, 45 @@ deg, 30 @@ deg,
                              45 @@ deg, 0 @@ deg, 120 @@ deg]
                   [1,1,r,1,1,r,1,1]
          myOpts = PolygonOpts myType NoOrient (p2 (0,0))
      in scale 0.5 $ polygon myOpts # lw none # fc (blend mag black white) #
         rotate (th @@ rad)

rVFGrad :: IO ()
rVFGrad = vfGrad id
          (\(x,y) -> cart x y 0)
          (\v -> (xComp v,yComp v))
          "rVFGrad.svg" 20
          rVF