module Petzval.System
  ( Pupil(..)
  , entrancePupil
  , createRay
  , seidel, Seidel(..)
  , Ray
  ) where

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

import qualified Debug.Trace

data Ray a = Ray (V3 a) (V3 a)

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

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 }
entrancePupil :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a, Show a) => [Element BakedIOR a] -> Pupil a
entrancePupil system = pupil stop prefix
  where (prefix, (Just stop), _) = splitSystem system

createRay :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a) => Maybe a -> Pupil a -> a -> (a,a) -> Ray a
createRay (Just objectPlane) Pupil{position=pz,radius=pr} h (px, py) = 
  Ray source (normalize $ target ^-^ source)
  where dz = pz - objectPlane 
        source = V3 0 (dz * tan h) objectPlane
        target = V3 (px * pr) (py * pr) pz
createRay Nothing Pupil{position=pz,radius=pr} h (px, py) = Ray source (normalize $ target ^-^ source)
  where h' = (pi * (-abs h) / 180) -- field angle in rad
        dy = (V3 0 (cos h') (-sin h')) `project` (V3 0 (py * pr) 0)
        dz =  V3 0 (pz * tan h') (pz * cos h')
      
        source = dy ^+^ dz
        target = V3 (px * pr) (py * pr) pz

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
                  }
  mappend 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'
           }
    


-- | Initial matrix is [ h_ h; u_ u ]
seidel' :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a, Show 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,_roc=roc}) = Debug.Trace.trace msg (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) / roc + (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 / roc * δ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 / roc
        msg = intercalate "\t" [show δun, show a, show abar]
        -- TODO: evaluate at other IORs
        -- cbas = h  (n2 - n1) / n1
        -- c1 = 
        
        
  

seidel :: (RealFloat a, Mode a, Scalar a ~ Double, Epsilon a, Show 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)