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 System.IO (hPrint, hPutStrLn, stderr)
import Linear
import Numeric.AD.Mode (Scalar, Mode)
import qualified Numeric.LinearAlgebra as L
import qualified Data.List as DL

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, _curvature = 1/100, _outsideRadius=10}
  , Surface{_material = air, _thickness = 95, _curvature = -1/100, _outsideRadius=10}
  ]

system2 =
  [
    Surface{_material = s_sk16, _outsideRadius=11.5, _curvature=1/42.98790, _thickness = 4}
  , Surface{_material = air, _outsideRadius = 11.5, _curvature = -1/248.07740, _thickness = 10.51018}
  , Surface{_material = s_f4, _outsideRadius = 9.852, _curvature = -1/38.21035, _thickness = 2.5}
  , Surface{_material = air, _outsideRadius = 8.885, _curvature = 1/43.95894, _thickness = 0}
  , Stop{_outsideRadius = 8.6762522794, _thickness = 9.86946}
  , Surface{_material=s_sk16, _outsideRadius = 11, _curvature=1/656.66349, _thickness = 4.5}
  , Surface{_material = air, _outsideRadius = 11, _curvature = -1/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 100 100 $ BFL Nothing
                          , DynMerit $ TW 0 1 $ SpotSize Nothing
                          , DynMerit $ TW 0 1 $ SpotSize (Just $ TraceConditions 7 0.587618)
                          , DynMerit $ TW 0 1 $ SpotSize (Just $ TraceConditions 15 0.587618)
                          , DynMerit $ TW 1 100 $ Vign (TraceConditions 15 wl)
                          ]
          wl = 0.5875618
          vars :: VariableSet
          vars = (ix 1 . curvature) `adjoin` (ix 2 . curvature) `adjoin` (ix 2.thickness)
          cfg = def{maxSteps = steps,eps1=0,eps2=0}

              

main :: IO ()
main =
  let baked = bake 0.587618 system2 :: [Element BakedIOR Double]
      ep = (entrancePupil baked){position=20.4747094540}
      fa = 20
  in do
    [arg] <- getArgs
    let steps = read arg
    hPrint stderr system1
    print $ L.norm_2  (L.fromList [3.0 , 4.0] :: L.Vector Double)
    (result, tr) <- doOptimize steps
    hPutStrLn stderr $ "Result BFL: " ++ (show . rearFocalLength . lensRTM . bake 0.5875618) result
    hPrint stderr result
    putStrLn "---"
    mapM_ (putStrLn . DL.intercalate "," . map show) tr
    -- 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

                    ]
      ]
-}