{-# 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
  , curvature
  , 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
  air = constMat 1
  constMat :: Double -> 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 []
  constMat ior = SellemeierMat [(ior * ior - 1, 0)]

-- | 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
  constMat = ConstMat

-- | 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
          , _curvature :: fp
          , _material :: mat
          }
  -- | Aperture stop
  | Stop {
      _thickness :: fp
      , _outsideRadius :: fp
      }
  -- | Imaging plane. This is a hack to make sure that rays go all the way to the target plane
  | ImagingPlane { _thickness :: fp }
                    deriving (Show)
makeLenses ''Element

roc :: Fractional fp => Traversal' (Element mat fp) fp
roc = curvature . iso (1/) (1/)

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

-- | 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@Stop{_thickness, _outsideRadius}) = Stop <$> inj _thickness <*> inj _outsideRadius
liftFp inj (s@Surface{_thickness, _outsideRadius, _curvature, _material}) = Surface <$> inj _thickness <*> inj _outsideRadius <*> inj _curvature <*> pure _material
{-# 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 system =
  snd $
  mapAccumLOf (each.material) (\n1 (ConstMat n2) -> (n2, BakedIOR n1 n2)) 1 $
  specialize wavelength system