{- LANGUAGE ImpredicativeTypes #-}
module Petzval.Merit
    (
      -- * Merit functions
      BFL(..)
    , SpotSize(..)
    , DynMerit(..)
    , Vign(..)
    , TW(..)
    , Merit(..)
    , MeritFunction
    , TraceConditions(..)
    , evalMerit
    , mkMerit) where

import Petzval.Calculations
import Petzval.Types
import Petzval.Trace
import Petzval.Optics (Element, BakedIOR, SellemeierMat, bake)
import Petzval.System
import Control.Lens
import Data.Either
import Data.Default
import Data.List (genericLength)
import Petzval.Optics.RTM
import Linear
import Data.Maybe
import Numeric.AD.Mode (auto)

import qualified Data.Map as Map

type MeritPart a = TracedSystem a -> a


data TraceConditions = TraceConditions { fieldAngle :: Double
                                       , wavelength :: Double
                                       }
                     deriving (Show)
    
data TracedSystem a =
  TracedSystem { _innerSystem :: [Element SellemeierMat a]
               , _fieldAngles :: [Double]
               , _wavelengths :: [Double]
               --, _tracePoints :: Map.Map TraceConditions [(Ray a, HitRecord a)]
               }
makeLenses ''TracedSystem

defWavelength :: (TracedSystem a) -> Double
defWavelength = fromMaybe 0.5875618 . listToMaybe  . _wavelengths
defFA :: TracedSystem a -> Double
defFA = fromMaybe 0 . listToMaybe . _fieldAngles
defTraceConditions :: TracedSystem a -> TraceConditions
defTraceConditions = TraceConditions <$> defFA <*> defWavelength

systemAtWavelength :: Double -> TracedSystem a -> [Element BakedIOR a]
systemAtWavelength λ =  bake λ . _innerSystem

class Merit m where
    calc :: Calcuable a => m -> TracedSystem a -> a


data SpotSize = SpotSize (Maybe TraceConditions)
                deriving (Show)

instance Merit SpotSize where
    calc (SpotSize Nothing) sys = calc (SpotSize (Just $ defTraceConditions sys)) sys
    calc (SpotSize (Just a)) sys = spotSize' a sys

data BFL = BFL (Maybe Double)
         deriving (Show)
instance Merit BFL where
    calc (BFL Nothing) = bfl
    calc (BFL (Just wl)) = bfl' wl
                         
spotSize' :: Calcuable a => TraceConditions -> MeritPart a
spotSize' conditions ts = rmsSize
                             . toListOf (each._pos._xy)
                             . rights
                             . map fst
                             . map (raytrace system)
                             . map (createRay Nothing ep fa)
                             $ hexapolarPattern 6
    where
      system = systemAtWavelength (wavelength conditions) ts
      fa = auto $ fieldAngle conditions
      ep = entrancePupil system

spotSize :: Calcuable a => MeritPart a
spotSize = defTraceConditions >>= spotSize'

-- | Vignetting at a field angle.
-- Computes the proportion of rays that make it through the system
data Vign = Vign TraceConditions
          deriving (Show)
instance Merit Vign where
    calc (Vign cond) ts = (/ genericLength rays)
                          . genericLength
                          . rights
                          . map fst
                          . map (raytrace system)
                          $ rays
        where rays = map (createRay Nothing ep fa) $ hexapolarPattern 6
              system = systemAtWavelength (wavelength cond) ts
              fa = auto $ fieldAngle cond
              ep = entrancePupil system

-- | Compute the back focal length
bfl :: Calcuable a => MeritPart a
-- | Compute the back focal length at a specific wavelength
bfl' :: (Calcuable a)
        => Double -- ^ Wavelength
     -> MeritPart a -- ^ Merit function

bfl' wavelength = rearFocalLength . systemRTM . systemAtWavelength wavelength
bfl = defWavelength >>= bfl'


data TW m = TW Double Double m
          deriving (Show)
instance Merit m => Merit (TW m) where
    calc (TW target weight merit) system = auto weight * (calc merit system - auto target)

component :: Calcuable a => Double -> Double -> MeritPart a -> MeritPart a
component target weight part system = let value = part system
                                      in (value - auto target) * auto weight


data MeritFunction = MeritFunction { components :: [DynMerit]
                                   , fnFieldAngles :: [Double]
                                   , fnWavelengths :: [Double]
                                   }
mkMerit :: [DynMerit] -> MeritFunction
mkMerit parts =  MeritFunction {
                   components = parts
                 , fnFieldAngles = []
                 , fnWavelengths = []
                 }

data DynMerit = forall m. (Show m, Merit m) => DynMerit m
instance Merit DynMerit where
    calc (DynMerit m) = calc m
                   
evalMerit :: Calcuable a => MeritFunction -> [Element SellemeierMat a] -> [a]
evalMerit fn system =
    let ts = TracedSystem {
               _innerSystem=system
             , _fieldAngles=fnFieldAngles fn
             , _wavelengths = fnWavelengths fn
             }
    in fmap (flip calc ts) $ components fn

      -- fmap ($ts) $ components fn