{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE NamedFieldPuns #-}
{-|
Basical optical types
-}
module Petzval.Optics 
  (
    -- * IOR calculations
    Material(..), SellemeierMat(..)
  , BakedIOR(..)
  , ConstMat(..)
    -- * Optical elements
  , Element(..)
  , isStop, isSurface
  , thickness
  , outsideRadius
  , material
  , roc
  , liftFp
  , specialize, bake
  ) where

import Control.Lens
import Data.Maybe

-- | An optical material, suitable to calculate indices of refraction
class Material mat where
  iorAtWavelength :: Double -> mat -> Double
  -- | A material with a constant IOR of 1
  air :: mat
  
-- | Calculates IOR based on the Sellemeier equation
data SellemeierMat = SellemeierMat [(Double, Double)]
  deriving (Show)
instance Material SellemeierMat where
  iorAtWavelength λ (SellemeierMat coeff) = sqrt $ 1 + (sum . map contrib $ coeff)
    where contrib (b,c) = b * w2 / (w2 - c)
          w2 = λ ^ 2
  air = SellemeierMat []

-- | A material with a constant IOR
newtype ConstMat = ConstMat Double
  deriving (Eq, Ord, Show)

instance Material ConstMat where
  iorAtWavelength _ (ConstMat ior) = ior
  air = ConstMat 1

-- | A representation of IOR as (previous,new) IOR at a surface
data BakedIOR = BakedIOR Double Double
  deriving (Eq, Ord, Show)

-- | An optical element
data Element mat fp =
  -- | Refractive surface
  Surface { _thickness :: fp
          , _outsideRadius :: fp
          , _roc :: fp
          , _material :: mat
          }
  -- | Aperture stop
  | Stop {
      _thickness :: fp
      , _outsideRadius :: fp
      }
                    deriving (Show)
makeLenses ''Element

-- | Determine if an element is a stop
isStop :: Element mat a -> Bool
isStop Stop{} = True
isStop _ = False
-- | Determine if an element is a refractive surface
isSurface :: Element mat a -> Bool
isSurface Surface{} = True
isSurface _ = False
    {-
-- | The space after the current element
thickness :: Lens (Element mat a) (Element mat a) a a
thickness inj (s@Surface{_thickness}) = (\nt -> s{_thickness=nt}) <$> inj _thickness
thickness inj (s@Stop{_thickness}) = (\nt -> s{_thickness=nt}) <$> inj _thickness
-- | The outside radius of the element. Rays that intersect the element beyond this radius are considered to have missed.
outsideRadius :: Lens (Element mat a) (Element mat a) a a
outsideRadius inj (s@Surface{_outsideRadius}) = (\nt -> s{_outsideRadius=nt}) <$> inj _outsideRadius
outsideRadius inj (s@Stop{_outsideRadius}) = (\nt -> s{_outsideRadius=nt}) <$> inj _outsideRadius

infinity :: (RealFloat a) => a
infinity = encodeFloat (floatRadix 0 - 1) (snd $ floatRange 0)
-- | The radius of curvature of an element. This is 1 for stops
roc :: RealFloat a => Lens (Element mat a) (Element mat a) a a
roc inj (s@Surface{_roc}) = (\nt -> s{_roc=nt}) <$> inj _roc
roc inj (s@Stop{}) = (const s) <$> inj infinity

-- | The material that a surface transitions into
--
-- For stops, this returns None and ignores being set.
material :: (Material mat, Material mat') => Lens (Element mat a) (Element mat' a) (Maybe mat) (Maybe mat')
material inj (s@Surface{_thickness, _outsideRadius, _roc, _material}) = (\nt -> Surface{_material=fromMaybe air nt, _thickness, _outsideRadius, _roc}) <$> inj (Just _material)
material inj (s@Stop{_thickness, _outsideRadius}) = (const Stop{_thickness, _outsideRadius}) <$> inj Nothing
-}
-- | Translate a lens element from one FP type to another. E.g., can be used to convert from scalars to the types in an automatic differentiation tool.
liftFp :: (Applicative f) => (fp -> f fp') -> Element m fp -> f (Element m fp')
liftFp inj (s@Surface{_thickness, _outsideRadius, _roc, _material}) = (\t' or' roc' -> Surface{_thickness=t', _outsideRadius=or', _roc=roc', _material}) <$> inj _thickness <*> inj _outsideRadius <*> inj _roc
liftFp inj (s@Stop{_thickness, _outsideRadius}) = (\t' or' -> Stop{_thickness=t', _outsideRadius=or'}) <$> inj _thickness <*> inj _outsideRadius
{-# INLINE liftFp #-}

-- | Specialize a system for a specific wavelength
specialize :: (Material mat) => Double -> [Element mat a] -> [Element ConstMat a]
specialize wavelength = (each.material) %~ (ConstMat . iorAtWavelength wavelength)
-- | Annotate each material with the incoming index of refraction
bake :: (Material mat) => Double -> [Element mat a] -> [Element BakedIOR a]
bake wavelength mat =
  snd $
  mapAccumLOf (each.material) (\n1 (ConstMat n2) -> (n2, BakedIOR n1 n2)) 1 $
  specialize wavelength mat