{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{- LANGUAGE ImpredicativeTypes #-}
module Petzval.Optimization where

import Linear
import Control.Monad
import Control.Monad.Loops
import Control.Monad.State
import Control.Monad.Writer (MonadWriter, tell)
import Control.Lens
import Control.Lens.Unsound (adjoin)
import Numeric.AD.Mode (auto)
import Numeric.AD.Mode.Reverse.Double
import Petzval.Merit
import Petzval.Optics
import Petzval.Types
import Control.Monad.Writer.Class
import Data.Default
import qualified Data.List as L

import Debug.Trace

import Numeric.LinearAlgebra hiding (Element, dot, (??))
import qualified Numeric.LinearAlgebra as L


-- | A set of modifiable parts of a lens system, expressed as a traversal.
-- The recommended way of generating a variable set is with the following construction:
-- @
-- vars = ix 1 . roc `adjoin` ix 2 . thickness
-- @
type VariableSet = forall mat a. Traversal' [Element mat a] a
type AdMode s = ReverseDouble s

extractVars :: VariableSet -> [Element mat a] -> [a]
extractVars vars system = system ^. partsOf vars

setVars :: VariableSet -> [Element mat a] -> [a] -> [Element mat a]
setVars vars system vals = system & partsOf vars .~ vals

gradientAt :: VariableSet -- ^ The set of independent variables
           -> (forall a. Calcuable a => [Element mat a] -> a) -- ^ merit function
           -> [Element mat Double] -- ^ The system
           -> (Double, [Double]) -- ^ The gradient
gradientAt vars merit system =  grad' (merit . setVars vars (system & each.liftFp %~ auto)) (extractVars vars system)
           
data DLSState = DLSState { _damping :: Double
                         , _lastSum :: Double
                         , _curPt :: Vector Double
                         , _dampScale :: Double
                         , _found :: Bool
                         , _curIter :: Integer
                         }

data DLSCfg = DLSCfg { eps1 :: Double -- ^ Cutoff for the derivative of the metric function
                     , eps2 :: Double -- ^ Cutoff for no longer making progress
                     , maxSteps :: Integer -- ^ Maximum number of steps to iterate
                     }

instance Default DLSCfg where
    def = DLSCfg 1e-8 1e-4 1000

makeLenses ''DLSState

           {-
evalMeritJ :: MeritFunction -> VariableSet -> [Element SellemeierMat Double] -> [(Double, [Double])]
evalMeritJ fn vars system = 
    let ts sys = TracedSystem {
               _system=sys
             , _fieldAngles=fieldAngles fn
             , _wavelengths = wavelengths fn
             }
        evalAt pt = fmap ($lts) $ components fn
            where lts = ts $ setVars vars (system & each.liftFp %~ auto) pt
    in unzip . jacobian' . evalAt $ extractVars vars system
-}

optimizeDLS :: (MonadWriter [[Double]] m -- ^ This yields a list of intermediate values of the merit function
               , MonadIO m)
            => DLSCfg
            -> VariableSet -- ^ The set of independent variables
            -> MeritFunction -- ^ merit function
            -> [Element SellemeierMat Double] -- ^ The system
            -> m [Element SellemeierMat Double]

optimizeDLS cfg vars merit system = fmap (setVars vars system . toList) $ evalStateT doOptimize initialState 
  where
    initialState = let pt = fromList $ extractVars vars system
                       (y,j) = traceShowId $ jacobianAt pt
                       a :: Matrix Double = tr j L.<> j
                       damping :: Double = 1e-3 * (maximum . toList . takeDiag) a
                       lastSum = sum . map (^2) . evalMerit merit $ system
                   in DLSState { _damping=damping
                               , _lastSum=lastSum
                               , _curPt=pt
                               , _dampScale = 2
                               , _found = False
                               , _curIter = maxSteps cfg
                               }
    isDone = orM [use found, uses curIter (<0)]
    doOptimize = (untilM_ (curIter -= 1 >> optimizeStep) isDone) >> use curPt
    -- optimizeStep :: m Double, where the return value is the change in merit
    optimizeStep = do
      mu <- use damping
      liftIO . print $ "Damping: " ++ show mu
      lastPt <- use curPt
      let (y, a) = jacobianAt lastPt
      let g = tr a #> y
      let dPt :: Vector Double = -(inv $ tr a L.<> a + (scalar mu * ident (cols a)) ) #> g
      let newPt = lastPt + dPt
      let oldMerit = sumSq y
      let newMerit = sumSq . fromList . evalMerit merit . setVars vars system . toList $ newPt
      let dL = (dPt `L.dot` (scalar mu * dPt - g)) / 2
      let gain = (oldMerit - newMerit) / dL
      if gain > 0
        then do curPt .= lastPt `add` dPt
                damping .= mu * max (1/3) (1 - (2*gain - 1) ^ 3)
                dampScale .= 2
                found ||= (norm_2 g <= eps1 cfg || norm_2 dPt <= eps2 cfg * (norm_2 lastPt + eps2 cfg))
                curPt .= newPt
        else do scale <- use dampScale
                d <- use damping
                dampScale *= 2
                liftIO $ print ("Rescaling", scale, d)
                damping *= scale
        
      tell [toList y]
      
      return $ newMerit - oldMerit
    jacobian' :: ([Double] -> [Double]) -> [Double] -> [(Double, [Double])]
    jacobian' fn vars = zip (fn vars) jacobian
        where deltas = map (/1000) vars
              jacobian = L.transpose $ varSets vars id
              varSets :: [Double] -> ([Double] -> [Double]) -> [[Double]]
              varSets [] _ = []
              varSets (var:vars) varPfx = head : rest
                  where fnWithD d = fn . varPfx $ var + d : vars
                        delta = var / 1000
                        head = map (/ (2 * delta)) $ zipWith (-) (fnWithD delta) (fnWithD (-delta))
                        rest = varSets vars (varPfx . (var:))
    jacobianAt :: Vector Double -> (Vector Double, Matrix Double)
    jacobianAt pt = let eval pt = undefined
                        (y,a) = unzip . jacobian' (\adpt -> evalMerit merit $ setVars vars (system & each.liftFp %~ auto) adpt) $ toList pt
                    in (fromList y, fromLists a)

sumSq :: Vector Double -> Double
sumSq x = L.dot x x
-- instances