--  This file is part of seismic_repa.
--
--  seismic_repa is free software: you can redistribute it and/or modify
--  it under the terms of the GNU General Public License as published by
--  the Free Software Foundation, either version 3 of the License, or
--  (at your option) any later version.
--
--  seismic_repa is distributed in the hope that it will be useful,
--  but WITHOUT ANY WARRANTY; without even the implied warranty of
--  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
--  GNU General Public License for more details.
--
--  You should have received a copy of the GNU General Public License
--  along with seismic_repa.  If not, see <http://www.gnu.org/licenses/>.


-- Implementation of the forward wave propagation
{-# LANGUAGE PackageImports, BangPatterns, QuasiQuotes, PatternGuards, TemplateHaskell #-}
{-# OPTIONS -Wall -fno-warn-missing-signatures -fno-warn-incomplete-patterns #-}

module Seismic
where

import Data.Array.Repa as R
import Data.Array.Repa.Stencil
import qualified Data.Vector.Unboxed as V

-- record type to hold the matrices
data SeismicData = SeismicData {ppf :: !(Array DIM2 Float),
                                apf :: !(Array DIM2 Float),
                                vel :: !(Array DIM2 Float)
                                }

-- some constants
fc = 125.0
v = 2000.0
dt = h/(5.0*2000.0)
h = 2000.0/(9.0*fc)
fat = (dt*dt)/(h*h*12.0)
tf = (2.0 * sqrt pi) / fc
a = sqrt 0.5 * pi * fc

-- produces correct values
genMatrices :: Int -> Int -> SeismicData
genMatrices _ 0 = undefined
genMatrices 0 _ = undefined
genMatrices x y = SeismicData {ppf = ppff, apf = apff, vel = vell}
  where
    ppff = R.fromList (Z :. (y::Int) :. (x::Int)) [0 |  _ <- [0..x*y-1]]
    apff = R.fromList (Z :. (y::Int) :. (x::Int)) [0 |  _ <- [0..x*y-1]]
    vell = R.fromList (Z :. (y::Int) :. (x::Int)) [v*v*fat | _ <- [0..x*y-1]] -- values are correct

-- produces correct values
genSeismicPulseVector :: Int -> Array DIM1 Float
genSeismicPulseVector 0 = undefined
genSeismicPulseVector t = pulseVector
	where
    pulseVector = R.fromList (Z :. (t::Int)) [ pulse $ (tt*dt) -tf |  tt <- [0.0..fromIntegral(t-1)]]
    pulse ts = (1.0 - 2.0 * ts * a * ts * a) * exp (-ts * a * ts * a)

-- loop n times
doNTimes :: SeismicData -> Array DIM1 Float -> Int -> Int -> Array DIM2 Float
doNTimes !sd !pv !t !ctr | ctr == t = apf sd
                         | otherwise = doNTimes (seismic (ppf sd) apfw (vel sd) ) pv t $ ctr+1
  where
    {-# INLINE apfw #-}
    apfw = addPulse (apf sd) $ pv `R.unsafeIndex` (Z:. ctr)
    {-# INLINE addPulse #-}
    addPulse apff pulse = force2 $ fromVector (Z:. xSize :. ySize) ( (toVector apff) `V.unsafeUpd`
                              [(posV, val (toVector apff) + pulse)])
    ex = extent $ apf sd
    xSize = (head $ listOfShape ex)
    ySize = (last $ listOfShape ex)
    posV = (xSize * (xSize `div` 2)) + (ySize `div` 2) -- index is correct
    val axs = axs `V.unsafeIndex` posV -- extract the value of apf at posV

-- the real calculation
seismic :: Array DIM2 Float ->  Array DIM2 Float -> Array DIM2 Float -> SeismicData
seismic
  !cppf@(Array _ [Region RangeAll (GenManifest _)])
  !capf@(Array _ [Region RangeAll (GenManifest _)])
  !cvel@(Array _ [Region RangeAll (GenManifest _)])
   = SeismicData { ppf = capf, apf = npf, vel = cvel}
  where
    {-# INLINE stencil #-}
    stencil = mapStencil2 (BoundConst 0.0) [stencil2|   0  0  -1  0  0
                                                        0  0  16  0  0
                                                        -1 16 -60 16 -1
                                                        0  0  16  0  0
                                                        0  0  -1  0  0 |] capf
    {-# INLINE npf #-}
    npf = capf `deepSeqArray` force2 $ (-^) (R.map (*2) capf) $ (-^) cppf $ (*^) cvel stencil
