{-# OPTIONS -Wall #-}

-- based on textbook "Numerical Electromagnetics"
-- by Inan and Marshall (Cambridge 2011)

-- Rather than a state-based update, this method uses
-- a big store that gets added to as needed.

module TM2D where

import Control.Monad (foldM)
import qualified Data.Map.Strict as M
import GIF

-- Inan/Marshall p. 84, Table 4.1

-- TM mode (Ez, Hx, Hy)

--    t    x    y
-- Ez odd  even even
-- Hx even even odd
-- Hy even odd  even

data XY = XY { xVal :: Int
             , yVal :: Int
             } deriving (Eq, Show)

instance Ord XY where
    compare (XY x1 y1) (XY x2 y2)
        = if y1 > y2
          then LT   -- because we want the high y values at the top
          else if y1 < y2
               then GT
               else if x1 < x2
                    then LT
                    else if x1 > x2
                         then GT
                         else EQ

type Store = M.Map (Int,Int,Int) Double

getNum :: Store -> (Int,Int,Int) -> (Store,Double)
getNum st (t,x,y)
    = case M.lookup (t,x,y) st of
        Just val -> (st,val)
        Nothing  -> if t <= 0
                    then (M.insert (t,x,y) 0 st, 0)
                    else if odd t && even x && even y
                         then let (st1,hxr) = getNum st  (t-1,x  ,y+1)
                                  (st2,hxl) = getNum st1 (t-1,x  ,y-1)
                                  (st3,hyr) = getNum st2 (t-1,x+1,y  )
                                  (st4,hyl) = getNum st3 (t-1,x-1,y  )
                                  (st5,ez ) = getNum st4 (t-2,x,  y  )
                                  -- j = if x == 200 && y == -10 || x == 200 && y == 10
                                  --     then 10 * sin (2*pi*c*(fromIntegral t * dt / 2)/lambda)
                                  --     else 0
                                  -- j = if x == 0 && y == -100 || x == 0 && y == 100
                                  --     then 10 * sin (2*pi*c*(fromIntegral t * dt / 2)/lambda)
                                  --     else 0
                                  j = if x == 200 && y == 0  -- 2D traveling wave, no interference
                                      then 10 * sin (2*pi*c*(fromIntegral t * dt / 2)/lambda)
                                      else 0
                                  result = ez + dt / myEps * ( (hyr - hyl) / dx
                                                             - (hxr - hxl) / dy - j )
                                  newst = M.delete (t-2,x,y) $ M.insert (t,x,y) result st5
                              in (newst,result)
                         else if even t && even x && odd y
                              then let (st1,ezr) = getNum st  (t-1,x,y+1)
                                       (st2,ezl) = getNum st1 (t-1,x,y-1)
                                       (st3,hx ) = getNum st2 (t-2,x,y  )
                                       result = hx - dt / myMu * (ezr - ezl) / dy
                                       newst = M.delete (t-2,x,y) $ M.insert (t,x,y) result st3
                                   in (newst,result)
                              else if even t && odd x && even y
                                   then let (st1,ezr) = getNum st  (t-1,x+1,y)
                                            (st2,ezl) = getNum st1 (t-1,x-1,y)
                                            (st3,hy ) = getNum st2 (t-2,x,y  )
                                            result = hy - dt / myMu * (- (ezr - ezl) / dx )
                                            newst = M.delete (t-2,x,y) $ M.insert (t,x,y) result st3
                                        in (newst,result)
                                   else error "bad trouble"


-- 10 * sin (2*pi*c*(fromIntegral t * dt / 2)/lambda)
-- = 10 * sin (2*pi*(30cm/ns)*(0.0625ns)/108cm)
-- = 10 * sin (2*pi*0.01736)

type Plane = M.Map XY Double

examplePlane :: Plane
examplePlane = M.fromList [(XY 1 1,15),(XY 1 2, -30),(XY 2 1, 4),(XY 2 2, -28)]

c :: Double
c = 30  -- cm/ns

myEps :: Double
myEps = 1 / c

myMu :: Double
myMu = 1 / c

dx :: Double
dx = 5.4   -- cm

dy :: Double
dy = 5.4   -- cm

dt :: Double
dt = 0.125   -- close to  dx / c / sqrt 2, which is 0.13 ns

-- 278 MHz
-- c = lambda / T = lambda f
-- lambda = c / f
lambda :: Double
lambda = 108   -- cm, 20 cells/wavelength

loX :: Int
--loX = -200  -- 200 x 200 cells
loX = 0  -- 200 x 200 cells

hiX :: Int
--hiX = 200
hiX = 400

loY :: Int
loY = -200

hiY :: Int
hiY = 200

-- 5 wavelengths (5.4 m) at 30 cm/ns
-- time is 540 cm / (30 cm/ns) = 18 ns
-- for a disturbance at the center to reach the edge
-- 18 ns / 0.13 ns = about 140 time steps
-- lets do 256
nTimeSteps :: Int
nTimeSteps = 256

twoDigit :: Int -> String
twoDigit n = reverse $ take 2 $ reverse ("0" ++ show n)

threeDigit :: Int -> String
threeDigit n = reverse $ take 3 $ reverse ("00" ++ show n)

getEzPlane :: Plane -> [Double]
getEzPlane = map snd . M.toList

-- x changes most rapidly; start at high y values since they are at the top
-- I don't think this order matters.
xysGIF :: [(Int,Int)]
xysGIF = [(x,y) | y <- [hiY,hiY-2..loY],  x <- [loX,loX+2..hiX]]

framePlane :: (Store,Int) -> (Store,Plane)
framePlane (st,t)
    = foldl (\(st',state) (x,y) -> let (st'',e') = getNum st' (t,x,y)
                                   in (st'',M.insert (XY x y) e' state)) (st,M.empty) xysGIF

oneFrame :: (Store,[[Double]])
         -> Int
         -> (Store,[[Double]])
oneFrame (store,zss) n
    = let (store1,state1) = framePlane (store,2*n-1)
          zs = getEzPlane state1
      in (store1, zss ++ [zs])

oneFrameLoud :: (Store,[[Double]])
             -> Int
             -> IO (Store,[[Double]])
oneFrameLoud (store,zss) n
    = let (store1,state1) = framePlane (store,2*n-1)
          zs = getEzPlane state1
      in do writeFile "debug.txt" $ show n ++ "\n" ++ show zs ++ "\n"
            print n
            return (store1, zss ++ [zs])

oneFrameSquared :: (Store,[[Double]])
         -> Int
         -> (Store,[[Double]])
oneFrameSquared (store,zss) n
    = let (store1,state1) = framePlane (store,2*n-1)
          zs = getEzPlane state1
      in (store1, zss ++ [map (**2) zs])

test2 :: Image
test2 = Image { imageColorTable = redBlueColorTable
              , imageWidth  = 201
              , imageHeight = 201
              , imageData = scale127list $ getEzPlane $ snd $ framePlane (M.empty, 51)
              }

main2 :: IO ()
main2 = makeImageGIF "test2.gif" test2

anim4 :: Animation
anim4 = Animation { animColorTable = redBlueColorTable
                  , animWidth      = 201
                  , animHeight     = 201
                  , animTimeStep   =  10
                  , animData = scale127 $ snd $ foldl oneFrame (M.empty,[]) [0..255]
                  }

main4 :: IO ()
main4 = makeAnimGIF "anim4.gif" anim4

anim5 :: Animation
anim5 = Animation { animColorTable = blackYellowColorTable
                  , animWidth      = 201
                  , animHeight     = 201
                  , animTimeStep   =  10
                  , animData = scaleNN $ snd $ foldl oneFrameSquared (M.empty,[]) [0..255]
                  }

-- Takes about 18 minutes to run.
main5 :: IO ()
main5 = makeAnimGIF "anim5.gif" anim5

anim6 :: Animation
anim6 = Animation { animColorTable = blackYellowColorTable
                  , animWidth      = 201
                  , animHeight     = 201
                  , animTimeStep   =  10
                  , animData = scaleNNpower 0.2 $ snd $ foldl oneFrameSquared (M.empty,[]) [0..255]
                  }

-- This is awesome.
main6 :: IO ()
main6 = makeAnimGIF "anim6.gif" anim6

-- Two-slit interference (changed the source in getNum)
main7 :: IO ()
main7 = makeAnimGIF "anim7.gif" anim6

-- Two-slit interference, sources at y=-100 and y=100.
main8 :: IO ()
main8 = makeAnimGIF "anim8.gif" anim6

-- uses scale127power for brigher colors
-- interference from y=-100 and y=100
anim9 :: Animation
anim9 = Animation { animColorTable = redBlueColorTable
                  , animWidth      = 201
                  , animHeight     = 201
                  , animTimeStep   =  10
                  , animData = scale127power 0.2 $ snd $ foldl oneFrame (M.empty,[]) [0..400]
                  }

main9 :: IO ()
main9 = makeAnimGIF "anim9.gif" anim9

-- prints time step as it goes
anim10 :: IO Animation
anim10 = do ad <- fmap (scale127power 1 . snd) $ foldM oneFrameLoud (M.empty,[]) [0..50]
            return Animation { animColorTable = redBlueColorTable
                             , animWidth      = 201
                             , animHeight     = 201
                             , animTimeStep   =  10
                             , animData       = ad
                             }

main10 :: IO ()
main10 = anim10 >>= makeAnimGIF "anim10.gif"

-- 1/3 power instead of 0.2, 512 time steps
anim11 :: Animation
anim11 = Animation { animColorTable = redBlueColorTable
                   , animWidth      = 201
                   , animHeight     = 201
                   , animTimeStep   =  10
                   , animData = scale127power (1/3) $ snd $ foldl oneFrame (M.empty,[]) [0..511]
                   }

main11 :: IO ()
main11 = makeAnimGIF "anim11.gif" anim11

anim12 :: Animation
anim12 = Animation { animColorTable = redBlueColorTable
                   , animWidth      = 201
                   , animHeight     = 201
                   , animTimeStep   =  10
                   , animData = scale127power (1/3) $ snd $ foldl oneFrame (M.empty,[]) [0..99]
                   }

main12 :: IO ()
main12 = makeAnimGIF "anim12.gif" anim12

-- anim4, but with 1/3 power scaling
anim13 :: Animation
anim13 = Animation { animColorTable = redBlueColorTable
                  , animWidth      = 201
                  , animHeight     = 201
                  , animTimeStep   =  10
                  , animData = scale127power (1/3) $ snd $ foldl oneFrame (M.empty,[]) [0..255]
                  }

main13 :: IO ()
main13 = makeAnimGIF "anim13.gif" anim13