-- | Utilities for full-precision raytracing

-- {-# OPTIONS_HADDOCK ignore-exports #-}
module Petzval.Trace
  ( Ray(..)
  , _dir, _pos
  , createRay
  , HitRecord(..)
  , TraceError(..)
  , raytrace
  -- * Ray patterns
  , hexapolarPattern
  , spiralPattern
  ) where

import Linear
import Petzval.System
import Petzval.Optics
import Numeric.AD.Mode (Scalar, Mode, auto)
import Control.Lens

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

_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

orLeft :: Maybe a -> b -> Either b a
orLeft  = maybe Left  (const . Right)

-- | 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
          -> 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 (pz * tan h') (pz * cos h')
      
        source = dy ^+^ dz
        target = V3 (px * pr) (py * pr) pz



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{_roc, _outsideRadius} (Ray pos dir) =
  toMaybe (hit1 && hit2) (Ray npos dir, Just normal)
  where origin = dir & _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
        normal = normalize $ origin ^+^ dir ^* dist
        npos = pos ^+^ dir ^* dist
        hit2 = (quadrance $ npos ^. _xy) <= _outsideRadius^2

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 :: V3 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 (rebased relative to the beginning of the optical system)
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 = trace' 1 ray system
  where -- trace' :: Double -> Ray a -> [Element BakedIOR a] -> Either TraceError [HitRecord a]
        trace' n1 ray@(Ray pos dir) (element:elements) = do
          let stopP = isStop element
          (nray, mnorm) <- hitTest element ray `orLeft` (if stopP then ElementMissed else HitStop)
          let mat@(BakedIOR _ n2)  = maybe (BakedIOR n1 n1) id $ element ^? material
          nray' <- maybe (Right nray) (\normal -> refract mat normal nray `orLeft` TIR) mnorm
          
          let opl = distance pos (nray ^. _pos) * auto n1
          (fray, rest) <- trace' n2 (nray'&_pos._z -~ element ^. thickness) elements
        
          return (fray & _pos._z +~ element ^. thickness, HitRecord { pos=(nray' ^. _pos), opl} : rest)

-- | 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]