```{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Trustworthy #-}

{- |
Module      :  Physics.Learn.Volume
Copyright   :  (c) Scott N. Walck 2012-2014
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains functions for working with 'Volume's
and volume integrals over 'Volume's.
-}

module Physics.Learn.Volume
(
-- * Volumes
Volume(..)
, unitBall
, unitBallCartesian
, centeredBall
, ball
, northernHalfBall
, centeredCylinder
, shiftVolume
-- * Volume Integral
, volumeIntegral
)
where

import Data.VectorSpace
( VectorSpace
, InnerSpace
, Scalar
)
import Physics.Learn.CarrotVec
( Vec
, vec
, sumV
, (^+^)
, (^-^)
, (^*)
, (^/)
, (<.>)
, (><)
, magnitude
)
import Physics.Learn.Position
( Position
, Displacement
, Field
, cartesian
, cylindrical
, spherical
, shiftPosition
, displacement
)

-- | Volume is a parametrized function from three parameters to space,
--   lower and upper limits on the first parameter,
--   lower and upper limits for the second parameter
--   (expressed as functions of the first parameter),
--   and lower and upper limits for the third parameter
--   (expressed as functions of the first and second parameters).
data Volume = Volume { volumeFunc :: (Double,Double,Double) -> Position  -- ^ function from 3 parameters to space
, loLimit    :: Double                      -- ^ s_a
, upLimit    :: Double                      -- ^ s_b
, loCurve    :: Double -> Double            -- ^ t_a(s)
, upCurve    :: Double -> Double            -- ^ t_b(s)
, loSurf     :: Double -> Double -> Double  -- ^ u_a(s,t)
, upSurf     :: Double -> Double -> Double  -- ^ u_b(s,t)
}

-- | A unit ball, centered at the origin.
unitBall :: Volume
unitBall = Volume spherical 0 1 (const 0) (const pi) (\_ _ -> 0) (\_ _ -> 2*pi)

-- | A unit ball, centered at the origin.  Specified in Cartesian coordinates.
unitBallCartesian :: Volume
unitBallCartesian = Volume cartesian (-1) 1 (\x -> -sqrtTol (1 - x*x)) (\x -> sqrtTol (1 - x*x))
(\x y -> -sqrtTol (1 - x*x - y*y)) (\x y -> sqrtTol (1 - x*x - y*y))

-- | A ball with given radius, centered at the origin.
centeredBall :: Double -> Volume
centeredBall a = Volume spherical 0 a (const 0) (const pi) (\_ _ -> 0) (\_ _ -> 2*pi)

-- | Ball with given radius and center.
ball :: Double    -- ^ radius
-> Position  -- ^ center
-> Volume    -- ^ ball with given radius and center
ball a c = Volume (\(r,th,phi) -> shiftPosition (vec (r * sin th * cos phi) (r * sin th * sin phi) (r * cos th)) c)
0 a (const 0) (const pi) (\_ _ -> 0) (\_ _ -> 2*pi)

-- | Upper half ball, unit radius, centered at origin.
northernHalfBall :: Volume
northernHalfBall = Volume spherical 0 1 (const 0) (const (pi/2)) (\_ _ -> 0) (\_ _ -> 2*pi)

-- | Cylinder with given radius and height.  Circular base of the cylinder
--   is centered at the origin.  Circular top of the cylinder lies in plane z = h.
-> Double  -- height
-> Volume  -- cylinder
centeredCylinder r h = Volume cylindrical 0 r (const 0) (const (2*pi)) (\_ _ -> 0) (\_ _ -> h)

-- | A volume integral
volumeIntegral :: (VectorSpace v, Scalar v ~ Double) =>
Int          -- ^ number of intervals for first parameter   (s)
-> Int          -- ^ number of intervals for second parameter  (t)
-> Int          -- ^ number of intervals for third parameter   (u)
-> Field v      -- ^ scalar or vector field
-> Volume       -- ^ the volume
-> v            -- ^ scalar or vector result
volumeIntegral n1 n2 n3 field (Volume f s_l s_u t_l t_u u_l u_u)
= sumV \$ map sumV \$ map (map sumV) (zipCubeWith (^*) aveVals volumes)
where
pts = [[[f (s,t,u) | u <- linSpaced n3 (u_l s t) (u_u s t) ] | t <- linSpaced n2 (t_l s) (t_u s)] | s <- linSpaced n1 s_l s_u]
volumes = zipWith3 (zipWith3 (zipWith3 (\du dv dw -> du <.> (dv >< dw)))) dus dvs dws
dus = uncurry zipSub3 (shift1 pts)
dvs = uncurry zipSub3 (shift2 pts)
dws = uncurry zipSub3 (shift3 pts)
vals = map (map (map field)) pts
aveVals = ((uncurry zipAve3 . shift1) . (uncurry zipAve3 . shift2) . (uncurry zipAve3 . shift3)) vals

-- zipSquareWith :: (a -> b -> c) -> [[a]] -> [[b]] -> [[c]]
-- zipSquareWith = zipWith . zipWith

zipCubeWith :: (a -> b -> c) -> [[[a]]] -> [[[b]]] -> [[[c]]]
zipCubeWith = zipWith . zipWith . zipWith

-- zipSub :: [Vec] -> [Vec] -> [Vec]
-- zipSub = zipWith (^-^)

-- zipSub2 :: [[Vec]] -> [[Vec]] -> [[Vec]]
-- zipSub2 = zipWith \$ zipWith (^-^)

zipSub3 :: [[[Position]]] -> [[[Position]]] -> [[[Vec]]]
zipSub3 = zipCubeWith displacement

zipAve3 :: (VectorSpace v, Scalar v ~ Double) => [[[v]]] -> [[[v]]] -> [[[v]]]
zipAve3 = zipCubeWith ave

shift1 :: [a] -> ([a],[a])
shift1 pts = (pts, tail pts)

shift2 :: [[a]] -> ([[a]],[[a]])
shift2 pts2d = (pts2d, map tail pts2d)

shift3 :: [[[a]]] -> ([[[a]]],[[[a]]])
shift3 pts3d = (pts3d, map (map tail) pts3d)

-- | n+1 points
linSpaced :: Int -> Double -> Double -> [Double]
linSpaced n a b
| a < b      = let dx = (b - a) / fromIntegral n
in [a,a+dx..b]
| a ~~ b     = [ave a b]
| otherwise  = error \$ "linSpaced:  lower limit greater than upper limit:  (n,a,b) = " ++ show (n,a,b)

(~~) :: (InnerSpace v, Scalar v ~ Double) => v -> v -> Bool
a ~~ b = magnitude (a ^-^ b) < tolerance

tolerance :: Double
tolerance = 1e-10

ave :: (VectorSpace v, Scalar v ~ Double) => v -> v -> v
ave v1 v2 = (v1 ^+^ v2) ^/ 2

sqrtTol :: Double -> Double
sqrtTol x
| x >= 0              = sqrt x
| abs x <= tolerance  = 0
| otherwise           = error ("sqrtTol:  I can't take the sqrt of " ++ show x)

-- | Shift a volume by a displacement.
shiftVolume :: Displacement -> Volume -> Volume
shiftVolume d (Volume f sl su tl tu ul uu)
= Volume (shiftPosition d . f) sl su tl tu ul uu
```