-- | Utilities for full-precision raytracing
-- {-# OPTIONS_HADDOCK ignore-exports #-}
{-# LANGUAGE FlexibleContexts, BangPatterns, DeriveAnyClass, DeriveGeneric #-}
module Petzval.Trace
  ( Ray(..)
  , _dir, _pos
  , createRay
  , HitRecord(..)
  , TraceError(..)
  , raytrace
  , raytrace1
  -- * Ray patterns
  , hexapolarPattern
  , spiralPattern
  ) where

import Linear
import Petzval.System
import Petzval.Optics
import Numeric.AD.Mode (Scalar, Mode, auto)
import Control.DeepSeq
import Control.Lens
import Control.Monad.State
import Control.Monad.Except
import Control.Monad.Writer
import GHC.Generics
import qualified Debug.Trace

-- | A ray. The first argument is the direction, and the second 
data Ray a = Ray (V3 a) (V3 a)
  deriving (Show, Eq, Generic, NFData)

_dir, _pos :: Lens' (Ray a) (V3 a)
-- | The direction of a ray
_dir = lens (\(Ray _ dir) -> dir) (\(Ray pos _) -> Ray pos)
-- | The position of a ray
_pos = lens (\(Ray pos _) -> pos) (\(Ray _ pos) dir -> Ray dir pos)

toMaybe :: Bool -> a -> Maybe a
toMaybe False = const Nothing
toMaybe True = Just

orError :: (MonadError e m) => Maybe a -> e -> m a
orError  = maybe throwError  (const . return)

forceRay :: Ray a -> Ray a
forceRay ray@(Ray (V3 !px !py !pz) (V3 !dx !dy !dz)) = ray

-- | Create a ray for a given field angle and pupil position.
--
-- * The first argument is the image plane position. If `Nothing`, the object is at infinity.
--
-- * The second argument is the entrance pupil to aim at
--
-- * The third argument is the field angle
--
-- * The fourth argument is the normalized pupil coordinates (in the range of \([-1,1]\))
createRay :: (RealFloat a, Mode a, Epsilon a)
          => Maybe a -- ^ The image plane position. If `Nothing`, the object is at infinity
          -> Pupil a -- ^ The entrance pupil to aim at
          -> a       -- ^ Field angle, in degrees
          -> V2 a    -- ^ Normalized pupil coordinates (in the range \([-1,1]\))
          -> Ray a        
createRay (Just objectPlane) Pupil{position=pz,radius=pr} h (V2 px py) = 
  Ray source (normalize $ target ^-^ source)
  where dz = pz - objectPlane 
        source = V3 0 (dz * tan h) objectPlane
        target = V3 (px * pr) (py * pr) pz
createRay Nothing Pupil{position=pz,radius=pr} h (V2 px py) = Ray source (normalize $ target ^-^ source)
  where 
        h' = (pi * (-abs h) / 180) -- field angle in rad
        dy = (V3 0 (cos h') (-sin h')) `project` (V3 0 (py * pr) 0)
        dz =  V3 0 (tan h') 1
      
        source = (dy ^-^ dz * 10) & _x .~ (px * pr)
        target = V3 (px * pr) (py * pr) pz

trace1 :: Show a => String -> a -> a
trace1 msg = (msg++) . show >>= Debug.Trace.trace


hitTest :: (Floating a, Ord a, Mode a, Epsilon a) =>  Element mat a -> Ray a -> Maybe (Ray a, Maybe (V3 a))
hitTest Stop{_outsideRadius} (Ray pos dir) =
  toMaybe pass $ (Ray npos dir, Nothing)
  where dz = pos ^. _z / dir ^. _z
        npos = pos ^-^ (dir ^* dz)
        pass = quadrance (npos ^. _xy) <= _outsideRadius ^ 2
hitTest Surface{_curvature, _outsideRadius} ray@(Ray pos dir) =
  toMaybe (hit1 && hit2) (Ray npos dir, Just normal)
  where _roc = 1/_curvature
        origin = pos & _z -~ _roc
        !a = dir `dot` dir
        !b = (dir `dot` origin) * 2
        !c = (origin `dot` origin) - _roc ^ 2
        !det = b^2 - 4 * a * c
        !hit1 = det >= 0
        !p2 = sqrt det
        !sa = (p2 - b) / 2 / a
        !sb = (-p2 - b) / 2 / a
        !s1 = min sa sb
        !s2 = max sa sb
        !dist = if s1 >= -0.001 then s1 else s2
        !normal0 = normalize $ origin ^+^ dir ^* dist
        !normal = if (normal0 ^. _z < 0) then -normal0 else normal0
        !npos = pos ^+^ dir ^* dist
        !hit2 = (quadrance $ npos ^. _xy) <= _outsideRadius^2
hitTest ImagingPlane{_thickness} (Ray pos dir) =
    Just (Ray npos dir, Nothing)
        where dz = pos ^. _z / dir ^. _z - _thickness
              npos = pos ^-^ (dir ^* dz)

refract :: (Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a) => BakedIOR -> V3 a -> Ray a -> Maybe (Ray a)
refract (BakedIOR n1 n2) normal (Ray pos incident) =
  let ni = normal `dot` incident
      mu = auto $ n1 / n2
      det = 1 - mu^2 * (1 - ni^2)
  in toMaybe (det >= 0) $ Ray pos $ mu *^ incident + (sqrt det - mu * ni) *^ normal

-- | The interaction of a ray with a particular element
data HitRecord a = HitRecord { pos :: Ray a -- ^ Position of the hit
                             , opl :: a    -- ^ Optical path length from the last hit to here
                             }
  deriving (Show)

-- | How a ray failed to complete a trace
data TraceError = HitStop       -- ^ Ray passed outside the aperture stop
                | ElementMissed -- ^ The ray missed an element completely
                | TIR           -- ^ The ray hit an element so obliquely that it
                                --   suffered from total internal reflection
  deriving (Show, Eq)

-- | Trace a ray through the give system. Returns the ray after the last element, relative to the vertex of the last element.
--
-- This is equivalent to `foldM raytrace1 ray system`, given an appropriate monad stack.
raytrace :: (Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a)
         => [Element BakedIOR a] -- ^ The system to trace
         -> Ray a -- ^ The initial ray
         -> (Either TraceError (Ray a), [HitRecord a])
raytrace system ray =
  runIdentity
  . flip evalStateT (1 :: Double)
  . runWriterT
  . runExceptT
  $ foldM raytrace1 ray system
  
-- | Trace a ray through a single element. Given an appropriate monad, this is a far more powerful interface to tracing than `raytrace`
raytrace1 :: ( Floating a, Ord a, Mode a, Scalar a ~ Double, Epsilon a
             , MonadState Double m
             , MonadWriter [HitRecord a] m -- ^ Tracing yields a list of 
             , MonadError TraceError m)  -- ^ This can fail
          => Ray a -> Element BakedIOR a -> m (Ray a)
raytrace1 ray element = do
      n1 <- get
      let stopP = isStop element
      
      (nray, mnorm) <- hitTest element ray `orError` (if stopP then HitStop else ElementMissed)
      let !mat@(BakedIOR _ n2)  = maybe (BakedIOR n1 n1) id $ element ^? material
      !nray' <- maybe (return nray) (\normal -> refract mat normal nray `orError` TIR) mnorm
      let !opl = distance (ray ^. _pos) (nray ^. _pos) * auto n1
      put n2
      tell [HitRecord { pos=nray', opl}]
      return $ nray' &_pos._z -~ element ^. thickness

-- | Spiral pattern. This is somewhat more irregular than the hexapolar pattern. The argument is the number of points
spiralPattern :: Floating a => Int -> [V2 a]
spiralPattern n = map (point . fromIntegral) [0..n-1]
  where npoints = fromIntegral n - 1
        point n = let r = sqrt (n / npoints)
                      theta = 2.3999632297286531 * n
                  in r *^ V2 (sin theta) (cos theta)
-- | A hexapolar pattern. This is rather optimally distributed; the argument is the number of rings
hexapolarPattern :: Floating a => Integer -> [V2 a]
hexapolarPattern nRings = (V2 0 0) : ( mconcat . map ring $ [0..nRings-1])
  where ring n = let point t = r *^ V2 (sin t) (cos t)
                     r = (fromIntegral n) / (fromIntegral nRings)
                 in map (point . (\x -> x / (fromIntegral) n * 3 * pi) . fromIntegral) [0 .. n-1]