-- | General utilities for dealing with full lens systems (i.e., the composition of multiple refracting surfaces)
module Petzval.System
  ( Pupil(..)
  , entrancePupil
  , seidel, Seidel(..)
  ) where

import Petzval.Optics
import Petzval.Optics.RTM
import Linear
import Numeric.AD.Mode
import Control.Lens
import Data.List

import qualified Debug.Trace


splitSystem :: [ Element mat a ] -> ([Element mat a], Maybe (Element mat a), [Element mat a])
splitSystem (s@Stop{}:rest) = ([], Just s, rest)
splitSystem [] = ([], Nothing, [])
splitSystem (s:rest) = (s:pfx, stop, sfx)
  where (pfx, stop, sfx) = splitSystem rest

-- | A pupil (real or virtual, entrance or exit)
data Pupil a = Pupil { radius :: a, position :: a }
  deriving (Show)

pupil (stop@Stop{_outsideRadius}) subsystem =
  let rtm = systemRTM subsystem
      V2 my mu = rtm !* V2 1 0
      V2 cy cu = (inv22 rtm) !* V2 0 (-1)
      --msg = intercalate " " ["ptrace: ", show (rtm !* V2 (_outsideRadius / my ) 0)]
  in Pupil { radius = _outsideRadius / my
           , position = -cy / cu }

-- | Compute the entrance pupil of a system viewed from infinity
entrancePupil :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a) => [Element BakedIOR a] -> Pupil a
entrancePupil system = pupil stop prefix
  where (prefix, (Just stop), _) = splitSystem system

-- | The set of seidel coefficients of a lens system
data Seidel a = Seidel
                   { sphr, coma, asti, fcur, dist :: a
                   -- , c1, c2 :: a
                   }
  deriving (Show)

instance Num a => Semigroup (Seidel a) where
  (<>) Seidel{sphr, coma, asti, fcur, dist} Seidel{sphr=sphr', coma=coma', asti=asti', fcur=fcur', dist=dist'} =
    Seidel { sphr = sphr + sphr'
           , coma = coma + coma'
           , asti = asti + asti'
           , fcur = fcur + fcur'
           , dist = dist + dist'
           --, c1 = c1 + c1'
           --, c2 = c2 + c2'
           }
  
instance Num a => Monoid (Seidel a) where
  mempty = Seidel { sphr=0, coma=0, asti=0, fcur=0, dist=0
                  -- , c1 = mempty, c2 = mempty
                  }    


-- | Initial matrix is [ h_ h; u_ u ]
seidel' :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a) => M22 a -> Element BakedIOR a -> (M22 a, Seidel a)
seidel' rays (s@Stop{}) = (thicknessRTM s !*! rays, mempty)
seidel' rays (s@Surface{_material=BakedIOR n1 n2,_curvature=curvature}) = (rays'', Seidel {sphr,coma,asti,fcur,dist})
  where rays' = refractionRTM s !*! rays
        rays'' = thicknessRTM s !*! rays'
        marg = column _y
        chief = column _x
        _i = to (\ray -> (ray ^. _x) / curvature + (ray ^. _y))
        a = auto n1 * (rays ^. marg._i)
        abar = auto n1 * (rays ^. chief._i)
        h = rays ^. marg._x
        δun = (rays' ^. marg._y / auto n2) - (rays ^. marg._y / auto n1)
        δ1n = auto (1/n2-1/n1)
        δ1n2 = auto (1/n2^2-1/n1^2)
        lagrange = auto n1 * det22 rays

        sphr = a ^ 2 * h * δun
        coma = -a * abar * h * δun
        asti = -abar^2 * h * δun
        fcur = lagrange^2 * curvature * δ1n
        -- Normally this is defined as abar/a * (asti*fcur)
        -- However, a is 0 whenever the marginal ray is normal to the surface
        -- Thus, we use this formula from Kidger
        dist = -abar^3 * h * δ1n2 + (rays^. chief._x) * abar * (2 * h * abar - rays ^. chief._x * a) * δ1n * curvature
        -- TODO: evaluate at other IORs
        -- cbas = h  (n2 - n1) / n1
        -- c1 = 
        
-- | Compute the seidel coefficients of a system given the pupil and field angle
seidel :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a) => Pupil a -> Double -> [Element BakedIOR a] -> [Seidel a]
seidel pupil fieldAngle =
  snd . mapAccumL seidel' rays
  where ubar = auto (tan (pi / 180 * fieldAngle))
        rays = V2 (V2 (-position pupil * ubar) (radius pupil))
                  (V2 ubar 0)