{-# OPTIONS -Wall #-}

module Main where

import Data.List (sort, nub)
import qualified Data.Map.Strict as M
import LPFPCore

-------------------------
-- Wave History method --
-------------------------

-- What is the meaning of timStep and posStep?
-- According to the latex definitions in the notes,
-- dt should be the time between nt = 0 and nt = 2.
data WaveHistory = WaveHistory { timInitial :: R
                               , timStep    :: R
                               , posInitial :: R
                               , posStep    :: R
                               , velWave    :: R
                               , waveValue  :: M.Map (Int,Int) R
                               } deriving (Show)

-- works for position too
tim :: R -> R -> Int -> R
tim t0 dt nt = t0 + dt / 2 * fromIntegral nt

-- Wave equation with a source term.
-- Assume boundaries are fixed at Ey = 0.
-- Ey is on even nx, even nt
-- Bz is on odd  nx, odd  nt
-- We keep track of Ey from 2 .. 198
historyUpdateEy :: ((R,R) -> R) -> WaveHistory -> WaveHistory
historyUpdateEy source wh
    = let dt = timStep wh
          dx = posStep wh
          t0 = timInitial wh
          x0 = posInitial wh
--          x nx' = x0 + dx / 2 * fromIntegral nx'
--          t nt' = t0 + dt / 2 * fromIntegral nt'
          stPoints = M.keys (waveValue wh)
          nt = maximum [nt' | (_,nt') <- stPoints]     -- this should be odd
          nxMax = maximum [nx' | (nx',_) <- stPoints]  -- this should be odd
          wv = waveValue wh
          f :: (Int,Int) -> R
          f (nx',nt') = case M.lookup (nx',nt') wv of
                          Nothing  -> error $ "bad lookup: " ++ show (nx',nt')
                          Just val -> val
          -- The most important functions is newEy:
          newEy :: Int -> R    -- newEy is for nt+1
          newEy nx = f (nx,nt-1) - f (nx+1,nt) + f (nx-1,nt) - dt / epsilon0 * source (tim x0 dx nx, tim t0 dt nt)  -- source = Jy
          wvNew = foldr (\nx -> M.insert (nx,nt+1) (newEy nx)) wv [2, 4 .. nxMax - 1]
      in wh { waveValue = wvNew }

-- Wave equation with a source term.
-- Assume boundaries are fixed at Ey = 0.
-- Ey is on even nx, even nt
-- Bz is on odd  nx, odd  nt
-- We keep track from 1 .. 199
-- We actually keep track of c Bz, which has the same units as Ey.
historyUpdateBz :: WaveHistory -> WaveHistory
historyUpdateBz wh
    = let stPoints = M.keys (waveValue wh)
          nt = maximum [nt' | (_,nt') <- stPoints]     -- this should be even
          nxMax = maximum [nx' | (nx',_) <- stPoints]  -- this should be odd
          wv = waveValue wh
          f :: (Int,Int) -> R
          f (nx',nt') = case M.lookup (nx',nt') wv of
                          Nothing -> if nx' == 0 || nx' == nxMax + 1
                                     then 0
                                     else error $ "bad lookup: " ++ show (nx',nt',nxMax,nt)
                          Just val -> val
          -- The most important functions is newBz:
          newBz :: Int -> R    -- newBz is for nt+1, which should be odd
          newBz nx = f (nx,nt-1) - f (nx+1,nt) + f (nx-1,nt)
          wvNew = foldr (\nx -> M.insert (nx,nt+1) (newBz nx)) wv [1, 3 .. nxMax]
      in wh { waveValue = wvNew }

historyUpdate :: ((R,R) -> R) -> WaveHistory -> WaveHistory
historyUpdate source = historyUpdateBz . historyUpdateEy source

lineToText :: (R,R) -> String
lineToText (x,y) = "L " ++ show x ++ "," ++ show (-y) ++ " "

framePathText :: [(R,R)] -> String
framePathText xys = 'M' : drop 1 (concatMap lineToText xys)

oneCurve :: R -> R -> (Int,Int,Int) -> [(R,R)] -> String
oneCurve t dt rgb xys
    = let fpt = framePathText xys
          pre = "<g stroke-linecap=\"butt\" fill=\"rgb(0,0,0)\" stroke-miterlimit=\"10.0\" stroke-width=\"0.03\" fill-opacity=\"0.0\" stroke-opacity=\"1.0\" stroke-linejoin=\"miter\" stroke=\"rgb" ++ show rgb ++ "\">"
          middle = "<path d=\"" ++ fpt ++ "\" visibility=\"hidden\"> <set attributeName=\"visibility\" attributeType=\"CSS\" to=\"visible\" begin=\"" ++ show t ++ "s\" dur=\"" ++ show dt ++ "s\" /> </path>\n"
          post = "</g>"
      in pre ++ middle ++ post

makeSVGAnimation :: FilePath -> WaveHistory -> IO ()
makeSVGAnimation filename wh
    = let dt = timStep wh
          dx = posStep wh
          t0 = timInitial wh
          x0 = posInitial wh
          wv = waveValue wh
          nts = sort $ nub $ [nt | (_,nt) <- M.keys wv]
          keysTime :: Int -> [(Int,Int)]
          keysTime nt = sort $ filter (\(_,nt') -> nt' == nt) $ M.keys wv
          f :: (Int,Int) -> R
          f (nx',nt') = case M.lookup (nx',nt') wv of
                          Nothing  -> error $ "makeSVGAnimation: bad lookup: " ++ show (nx',nt')
                          Just val -> val
          xyPairs :: Int -> [(R,R)]
          xyPairs nt = [(tim x0 dx nx', f (nx',nt')) | (nx',nt') <- keysTime nt]
          matterEy = concat [oneCurve (tim t0 dt nt) dt (0,0,255) (xyPairs nt) | nt <- nts, even nt]
          matterBz = concat [oneCurve (tim t0 dt nt) dt (255,0,0) (xyPairs nt) | nt <- nts, odd nt]
          svgStart =    "<?xml version=\"1.0\" standalone=\"no\"?>\n"
                     ++ "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n"
                     ++ "  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n"
                     ++ "<svg width=\"32cm\" height=\"16cm\" viewBox=\"-5 -2 25 2\"\n"
                     ++ "     xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n"
          svgEnd   =    "</svg>\n"
      in writeFile filename $ svgStart ++ matterEy ++ matterBz ++ svgEnd

makeSVGAnimation60 :: FilePath -> WaveHistory -> IO ()
makeSVGAnimation60 filename wh
    = let dt = timStep wh
          dx = posStep wh
          t0 = timInitial wh
          x0 = posInitial wh
          wv = waveValue wh
          nts = sort $ nub $ [nt | (_,nt) <- M.keys wv]
          keysTime :: Int -> [(Int,Int)]
          keysTime nt = sort $ filter (\(_,nt') -> nt' == nt) $ M.keys wv
          f :: (Int,Int) -> R
          f (nx',nt') = case M.lookup (nx',nt') wv of
                          Nothing  -> error $ "makeSVGAnimation: bad lookup: " ++ show (nx',nt')
                          Just val -> val
          xyPairs :: Int -> [(R,R)]
          xyPairs nt = [(tim x0 dx nx', f (nx',nt')) | (nx',nt') <- keysTime nt]
          curvesEy = [oneCurve (tim t0 dt nt) dt (0,0,255) (xyPairs nt) | nt <- nts, even nt]
          curvesBz = [oneCurve (tim t0 dt nt) dt (255,0,0) (xyPairs nt) | nt <- nts, odd  nt]
          ntMaxEven = maximum [nt | nt <- nts, even nt]
          ntMaxOdd  = maximum [nt | nt <- nts, odd  nt]
          matterEy = concat curvesEy ++ oneCurve (tim t0 dt ntMaxEven) 60 (0,0,255) (xyPairs ntMaxEven)
          matterBz = concat curvesBz ++ oneCurve (tim t0 dt ntMaxOdd ) 60 (255,0,0) (xyPairs ntMaxOdd )
          svgStart =    "<?xml version=\"1.0\" standalone=\"no\"?>\n"
                     ++ "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n"
                     ++ "  \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n"
                     ++ "<svg width=\"32cm\" height=\"16cm\" viewBox=\"-5 -2 25 2\"\n"
                     ++ "     xmlns=\"http://www.w3.org/2000/svg\" version=\"1.1\">\n"
          svgEnd   =    "</svg>\n"
      in writeFile filename $ svgStart ++ matterEy ++ matterBz ++ svgEnd

-- x0 = 0, x1 = 200 * dt / 2 = 10
-- Every 0.1s of animation time, the wave proceeds one step in position (from nx = 100
-- to nx = 98, say.  It should take 5 seconds to get to the boundary.
-- That's about right.  I want it to take 10 seconds.

initialWaveZero :: R -> WaveHistory
initialWaveZero dt
    = let dx = dt
          v  = 1
      in WaveHistory { timInitial = 0
                     , timStep    = dt
                     , posInitial = 0
                     , posStep    = dx
                     , velWave    = v
                     , waveValue  = M.fromList $ [((nx,0),0) | nx <- [2, 4 .. 198]]
                                    ++ [((nx,1),0) | nx <- [1, 3 .. 199]]
                     }

-- Can depend on x and t.
sourceTerm :: (R,R) -> R
sourceTerm (x,t)
    = if x > 9.99 && x < 10.01
      then 1e-10 * sin t
      else 0

initialWave2 :: R -> Int -> R -> WaveHistory
initialWave2 dt nxmax boxLength
    = let dx = 2 * boxLength / fromIntegral nxmax
          v  = dx / dt
      in WaveHistory { timInitial = 0
                     , timStep    = dt
                     , posInitial = 0
                     , posStep    = dx
                     , velWave    = v
                     , waveValue  = M.fromList $ [((nx,0),0) | nx <- [2, 4 .. nxmax - 2]]
                                    ++ [((nx,1),0) | nx <- [1, 3 .. nxmax - 1]]
                     }

main08 :: IO ()
main08 = let wh = last $ take 150 $ iterate (historyUpdate sourceTerm) (initialWave2 0.1 600 20)
         in makeSVGAnimation60 "w08.svg" wh

-- w08.svg is awesome.  It freezes for a minute just as the wave hits the boundaries.

-- If we were working in SI units, our source frequency would be 1 rad/s.
-- f = 1/(2 pi) Hz, and T = 2 pi seconds.
-- If 0.1s were a physical time, and 20m was a physical length,
-- we'd have dx = 2 * 20 / 600 = 2/30 m
-- and v = dx / dt = 2/3 m/s.
-- This gives a wavelength of lambda = v T = 4 pi/3 meters.
-- So, we get about 20 / (4 pi / 3) = 15 / pi wavelengths per box.
-- A little less than 5 wavelengths per box.
-- That seems correct.

-- Standing wave
-- initial Ey:
-- Ey(x,t) = sin k x cos omega t 
-- Ey(x,t) = sin(k x - omega t) + sin(k x + omega t)
-- omega = v k

initialWaveStanding :: R -> Int -> R -> WaveHistory
initialWaveStanding dt nxmax boxLength
    = let dx = 2 * boxLength / fromIntegral nxmax
          x0 = 0
          v  = dx / dt
          k = 2 * pi * 1.5 / boxLength
          ey x = sin (k * (x - x0))
          omega = v * k
          bz (x,t) = 0.5 * sin(k * (x - x0) - omega * t) - 0.5 * sin(k * (x - x0) + omega * t)
      in WaveHistory { timInitial = 0
                     , timStep    = dt
                     , posInitial = x0
                     , posStep    = dx
                     , velWave    = v
                     , waveValue  = M.fromList $ [((nx,0),ey (tim x0 dx nx)) | nx <- [2, 4 .. nxmax - 2]]
                                    ++ [((nx,1),bz (tim x0 dx nx,dt/2)) | nx <- [1, 3 .. nxmax - 1]]
                     }

main11 :: IO ()
main11 = let wh = last $ take 600 $ iterate (historyUpdate (const 0)) (initialWaveStanding 0.1 200 20)
         in makeSVGAnimation "w11.svg" wh

main :: IO ()
main = main11

-- w11.svg is a nice example of an EM standing wave.  Worth studying and talking about.