{-# LANGUAGE NamedFieldPuns #-}
{-|
Description: First-order optical computation tools based on ray transfer matrices

Provides tools for computing various first-order properties of a lens
based on the paraxial approximation. 
-}
module Petzval.Optics.RTM
  (
  -- * Computation of ray transfer matrices from elements
    elementRTM, refractionRTM, thicknessRTM
  -- * Computation of ray transfer matrices for a system
  , lensRTM, systemRTM
  -- * First order properties of a lens system
  , rearPrincipalPlane, rearFocalLength, rearWorkingDistance
  -- * RTM manipulation
  , reverseRTM
  ) where


import Numeric.AD.Mode
import Linear.V2
import Linear.Matrix
import Petzval.Optics
import Control.Lens

-- | Compute an RTM for the optically active parts of a lens
-- system. This is the composition of the RTM for each element, but
-- ignores the thickness of the final element.
lensRTM :: (Fractional a, Mode a, Scalar a ~ Double) => [Element BakedIOR a] -> M22 a
-- | Compute the RTM of a complete optical system, including the
-- thickness of the final element.
systemRTM :: (Fractional a, Mode a, Scalar a ~ Double) => [Element BakedIOR a] -> M22 a
lensRTM [] = V2 (V2 1 0) (V2 0 1)
lensRTM [el] = refractionRTM el
lensRTM (el:rest) = lensRTM rest !*! elementRTM el
systemRTM [] = V2 (V2 1 0) (V2 0 1)
systemRTM els = foldl1 (flip (!*!)) . map elementRTM $ els


-- | Compute the refractive part of the RTM for an element
refractionRTM :: (Fractional a, Mode a, Scalar a ~ Double) => Element BakedIOR a -> M22 a
refractionRTM  (Surface{_thickness, _outsideRadius, _roc, _material=BakedIOR n1 n2}) =
  let ior' = auto (n1 / n2)
  in V2 (V2 1 0)
     (V2 ((ior' - fromIntegral 1) / _roc) ior')
refractionRTM Stop{} = V2 (V2 1 0) (V2 0 1)

-- | Compute the part of the RTM that represents the thickness of the element
thicknessRTM :: (Fractional a) => Element mat a -> M22 a
thicknessRTM element = V2 (V2 1 (element ^.thickness)) (V2 0 1)

-- | Compute the complete RTM of an element, i.e., \(thickness * refractive\)
elementRTM :: (Fractional a, Mode a, Scalar a ~ Double) => Element BakedIOR a -> M22 a
elementRTM el@Surface{} = thicknessRTM el !*! refractionRTM el
elementRTM el@Stop{} = thicknessRTM el

-- | Compute the position of the rear principal plane for a lens system relative to the final element.
rearPrincipalPlane :: Fractional a => M22 a -> a
rearPrincipalPlane rtm = let V2 y m = rtm !* V2 1 0
                     in (1 - y) / m
-- | Compute the rear focal length (distance from the rear principal plane to the focal point)
rearFocalLength :: Fractional a => M22 a -> a
rearFocalLength rtm = let V2 y m = rtm !* V2 1 0
                      in -1 / m
-- | Compute the rear working distance of a system (distance from the final element to the focal point)
rearWorkingDistance :: Fractional a => M22 a -> a
rearWorkingDistance rtm = let V2 y m = rtm !* V2 1 0
                          in -y / m

-- | Compute the RTM for the reverse of the system represented by @rtm@. This is equivalent to
-- \[
-- \begin{bmatrix} 1 & 0 \\ 0 & -1 \end{bmatrix}
-- rtm^{-1}
-- \begin{bmatrix} 1 & 0 \\ 0 & -1 \end{bmatrix}
-- \]
reverseRTM :: (Fractional a) => M22 a -> M22 a
reverseRTM rtm = let V2 (V2 a b) (V2 c d) = inv22 rtm
                 in V2 (V2 a (-b)) (V2 (-c) d)