module Main where

import Control.Lens
import Control.Lens.Unsound (adjoin)
import Control.Monad.Writer
-- import Criterion.Main
import Data.Default
import Data.Either
import Petzval.Optics
import Petzval.Optics.RTM
import Petzval.System
import Petzval.Merit
import Petzval.Trace
import Petzval.Calculations
import Petzval.Optimization
import Petzval.Types
import System.Environment (getArgs)
import Linear
import Numeric.AD.Mode (Scalar, Mode)
import qualified Numeric.LinearAlgebra as L

bk7 = SellemeierMat [ (1.03961212,  6.00069867e-3 )
                    , (0.231792344, 2.00179144e-2 )
                    , (1.01046945,  103.560653 )]
n_sk16 = SellemeierMat [ (1.343177740, 0.007046873)
                       , (0.241144399, 0.0229005000)
                       , (0.994317969, 92.75085260)]

n_ssk8 = SellemeierMat [ (1.44857867, 1.17965926e-01)
                       , (1.06937528, 8.69310149E-03)
                       , (4.21566593E-02, 1.11300666E+02) ]

system1 =
  [ Stop{_thickness = 0, _outsideRadius=5}
  , Surface{_material = bk7, _thickness = 10, _roc = 100, _outsideRadius=10}
  , Surface{_material = air, _thickness = 95, _roc = -100, _outsideRadius=10}
  ]

system2 =
  [
    Surface{_material = s_sk16, _outsideRadius=11.5, _roc=42.98790, _thickness = 4}
  , Surface{_material = air, _outsideRadius = 11.5, _roc = -248.07740, _thickness = 10.51018}
  , Surface{_material = s_f4, _outsideRadius = 9.852, _roc = -38.21035, _thickness = 2.5}
  , Surface{_material = air, _outsideRadius = 8.885, _roc = 43.95894, _thickness = 0}
  , Stop{_outsideRadius = 8.6762522794, _thickness = 9.86946}
  , Surface{_material=s_sk16, _outsideRadius = 11, _roc=656.66349, _thickness = 4.5}
  , Surface{_material = air, _outsideRadius = 11, _roc = -33.50754, _thickness = 86.48643}
  ]
  where s_sk16 = constMat 1.620408 :: SellemeierMat
        s_f4 = constMat 1.616589

merit :: (Material mat, Calcuable a) => [Element mat a] -> a
-- merit :: (Material mat) => [Element mat Double] -> Double
merit system = result 
  where baked = bake 0.5875618 system
        ep = entrancePupil baked
        result = sum
          . map (^2)
          . map (\ angle -> rmsSize
                  . toListOf (each._pos._xy)
                  -- . (\x -> x :: [Ray a])
                  . rights
                  . map fst
                  . map (raytrace baked . createRay Nothing ep angle)
                  $ hexapolarPattern 6)
          $ [0, 7.07, 15]

type Id a = a -> a


doOptimize steps = runWriterT $ optimizeDLS cfg vars merit system1
    where merit = mkMerit [ DynMerit $ TW 50 100 $ BFL Nothing
                          , DynMerit $ TW 0 1 $ SpotSize Nothing
                          ]
          vars :: VariableSet
          vars = (ix 1 . roc) `adjoin` (ix 2 . roc)
          vars2 = undefined -- foldl1 adjoin
                  -- [ ix 1.roc , ix 2.roc]
          cfg = def{maxSteps = steps}

              

main :: IO ()
main =
  let baked = bake 0.587618 system2
      ep = (entrancePupil baked){position=20.4747094540}
      fa = 20
  in do
    [arg] <- getArgs
    let steps = read arg
    print $ L.norm_2  (L.fromList [3.0 , 4.0] :: L.Vector Double)
    doOptimize steps >>= print
    putStrLn $ show ep
    putStrLn $ show $ rearFocalLength $ lensRTM baked
    putStrLn $ show $ mconcat (seidel ep fa baked :: [Seidel Double])
    
 {-   defaultMain [
      bgroup "merit" [ bench "system1" $ nf merit (system1 :: [Element SellemeierMat Double])
                     , bench "system2" $ nf merit (system2 :: [Element ConstMat Double])
                     ]
      , bgroup "ad" [ bench "system1" $ nf (gradientAt (ix 1.thickness) merit) system1,
                      bench "system2" $ nf (gradientAt (ix 1.thickness) merit) system2

                    ]
      ]
-}