{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies  #-}

-- | Provides spherical harmonic models of scalar-valued functions.
module Math.SphericalHarmonics
(
  SphericalHarmonicModel
, sphericalHarmonicModel
, scaledSphericalHarmonicModel
, evaluateModel
, evaluateModelCartesian
, evaluateModelGradient
, evaluateModelGradientCartesian
, evaluateModelGradientInLocalTangentPlane
)
where

import Data.Complex
import Data.VectorSpace hiding (magnitude)
import Math.SphericalHarmonics.AssociatedLegendre
import Numeric.AD

-- | Represents a spherical harmonic model of a scalar-valued function.
data SphericalHarmonicModel a = SphericalHarmonicModel [[(a, a)]]
  deriving (Functor)

-- | Creates a spherical harmonic model.
-- Result in an error if the length of the list is not a triangular number.
sphericalHarmonicModel :: (Fractional a) => [[(a, a)]] -- ^ A list of g and h coefficients for the model
                       -> SphericalHarmonicModel a -- ^ The spherical harmonic model
sphericalHarmonicModel cs | valid = SphericalHarmonicModel cs
                          | otherwise = error "The number of coefficients is not a triangular number."
  where
    valid = and $ zipWith (==) (fmap length cs) [1..length cs]

-- | Creates a spherical harmonic model, scaling coefficients for the supplied reference radius.
-- Result in an error if the length of the list is not a triangular number.
scaledSphericalHarmonicModel :: (Fractional a) => a -- ^ The reference radius
                       -> [[(a, a)]] -- ^ A list of g and h coefficients for the model
                       -> SphericalHarmonicModel a -- ^ The spherical harmonic model
scaledSphericalHarmonicModel r cs = sphericalHarmonicModel cs'
  where
    cs' = normalizeReferenceRadius r cs

instance(Fractional a, Eq a) => AdditiveGroup (SphericalHarmonicModel a) where
  zeroV = SphericalHarmonicModel [[(0,0)]]
  negateV = fmap negate
  (SphericalHarmonicModel m1) ^+^ (SphericalHarmonicModel m2) = SphericalHarmonicModel (combineCoefficients m1 m2)
    where
      combineCoefficients [] cs = cs
      combineCoefficients cs [] = cs
      combineCoefficients (c1:cs1) (c2:cs2) = zipWith addPairs c1 c2 : combineCoefficients cs1 cs2
      addPairs (g1, h1) (g2, h2) = (g1 + g2, h1 + h2)

instance (Fractional a, Eq a) => VectorSpace (SphericalHarmonicModel a) where
  type Scalar (SphericalHarmonicModel a) = a
  x *^ m = fmap (* x) m

normalizeReferenceRadius :: (Fractional a) => a -> [[(a, a)]] -> [[(a, a)]]
normalizeReferenceRadius r = zipWith (fmap . mapWholePair . transform) [0 :: Int ..]
  where
    transform n = (* (r ^ (2 + n)))

-- | Computes the scalar value of the spherical harmonic model at a specified spherical position.
evaluateModel :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
              -> a -- ^ Spherical radius
              -> a -- ^ Spherical colatitude (radian)
              -> a -- ^ Spherical longitude (radian)
              -> a -- ^ Model value
evaluateModel m r colat lon = evaluateModel' m r (cos colat) (cis lon)

-- | Computes the scalar value of the spherical harmonic model at a specified Cartesian position.
evaluateModelCartesian :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                       -> a -- ^ X position
                       -> a -- ^ Y position
                       -> a -- ^ Z position
                       -> a -- ^ Model value
evaluateModelCartesian m x y z = evaluateModel' m r cosColat cisLon
  where
    r = sqrt $ (x*x) + (y*y) + (z*z)
    cosColat = z / r
    cisLon = normalize $ mkPolar x y

evaluateModel' :: (RealFloat a, Ord a) => SphericalHarmonicModel a
               -> a -- r
               -> a -- cosColat
               -> Complex a -- cisLon
               -> a
evaluateModel' (SphericalHarmonicModel cs) r cosColat cisLon = sum $ zipWith (*) (iterate (/ r) (recip r)) (zipWith evaluateDegree [0..] cs)
  where
    sines = 1 : iterate (* cisLon) cisLon
    evaluateDegree n cs' = sum $ zipWith3 evaluateOrder (fmap (schmidtSemiNormalizedAssociatedLegendreFunction n) [0..n]) cs' sines
    evaluateOrder p (g, h) cisMLon = ((g * realPart cisMLon) + (h * imagPart cisMLon)) * (p (cosColat))

-- | Computes the gradient of the scalar value of the spherical harmonic model, in spherical coordinates, at a specified location.
evaluateModelGradient :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                      -> a -- ^ Spherical radius
                      -> a -- ^ Spherical colatitude (radian)
                      -> a -- ^ Spherical longitude (radian)
                      -> (a, a, a) -- ^ Radial, colatitudinal, and longitudinal components of gradient
evaluateModelGradient model r colat lon = makeTuple . fmap negate $ modelGrad [r, colat, lon]
  where
    modelGrad = grad (\[r', c', l'] -> evaluateModel (fmap auto model) r' c' l')

-- | Computes the gradient of the scalar value of the spherical harmonic model at a specified location, in Cartesian coordinates.
-- The result is expressed in right-handed coordinates centered at the origin of the sphere, with the positive Z-axis piercing the
-- north pole and the positive x-axis piercing the reference meridian.
evaluateModelGradientCartesian :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                               -> a -- ^ X position
                               -> a -- ^ Y position
                               -> a -- ^ Z position
                               -> (a, a, a) -- X, Y, and Z components of gradient
evaluateModelGradientCartesian model x y z = makeTuple . fmap negate $ modelGrad [x, y, z]
  where
    modelGrad = grad (\[x', y', z'] -> evaluateModelCartesian (fmap auto model) x' y' z')

-- | Computes the gradient of the scalar value of the spherical harmonic model at a specified location, in Cartesian coordinates.
-- The result is expressed in a reference frame locally tangent to the sphere at the specified location.
evaluateModelGradientInLocalTangentPlane :: (RealFloat a, Ord a) => SphericalHarmonicModel a -- ^ Spherical harmonic model
                                         -> a -- ^ Spherical radius
                                         -> a -- ^ Spherical colatitude (radian)
                                         -> a -- ^ Spherical longitude (radian)
                                         -> (a, a, a) -- ^ East, North, and up components of gradient
evaluateModelGradientInLocalTangentPlane model r colat lon = (e, n, u)
  where
    (r', colat', lon') = evaluateModelGradient model r colat lon
    e = lon' / (r * sin colat)
    n = -colat' / r -- negated because the colatitude increase southward
    u = r'

normalize :: (RealFloat a) => Complex a -> Complex a
normalize r@(x :+ y) | isInfinite m' = 0
                     | otherwise = (x * m') :+ (y * m')
  where
    m' = recip . magnitude $ r

mapWholePair :: (a -> b) -> (a, a) -> (b, b)
mapWholePair f (a, b) = (f a, f b)

makeTuple :: [a] -> (a, a, a)
makeTuple [x, y, z] = (x, y, z)