{-@ LIQUID "--reflection"     @-}

module SGD.SGD where 

import           Prelude  hiding ( head, tail, sum)
import           Monad.PrM 
import           Data.Dist 
import           Data.Derivative

{-@ type StepSize = {v:Double | 0.0 <= v } @-}
type StepSize = Double
{-@ data StepSizes = SSEmp | SS StepSize StepSizes @-}
data StepSizes = SSEmp | SS StepSize StepSizes
type DataPoint = (Double, Double)
type Weight = Double
type LossFunction = DataPoint -> Weight -> Double

type Set a = [a]
{-@ type DataSet = {v:Set DataPoint| 1 < lend v && 1 < len v } @-}
type DataSet = Set DataPoint
type DataPrM = PrM DataPoint


{-@ reflect sgd @-}
{-@ sgd :: zs:{DataSet | 1 < len zs && 1 < lend zs } -> Weight -> ss:StepSizes -> LossFunction 
       -> PrM Weight / [ sslen ss, 0 ] @-}
sgd :: DataSet -> Weight -> StepSizes -> LossFunction -> PrM Weight
sgd _  w0 SSEmp    _ = ppure w0
sgd zs w0 (SS α a) f = 
  choice (one / lend zs)
         (bind uhead (sgdRecUpd zs w0 α a f))
         (bind utail (sgdRecUpd zs w0 α a f)) 
 where
  uhead = ppure (head zs)
  utail = unif (tail zs)


{-@ reflect sgdRecUpd @-}
{-@ sgdRecUpd :: zs:{DataSet | 1 < len zs && 1 < lend zs } -> Weight -> StepSize -> ss:StepSizes -> LossFunction 
       -> DataPoint -> PrM Weight / [ sslen ss, 1 ] @-}
sgdRecUpd :: DataSet -> Weight -> StepSize -> StepSizes -> LossFunction -> DataPoint -> PrM Weight
sgdRecUpd zs w0 α a f z = bind (sgd zs w0 a f) (pureUpdate z α f)

{-@ reflect pureUpdate @-}
{-@ pureUpdate :: DataPoint -> StepSize -> LossFunction -> Weight -> PrM Weight @-}
pureUpdate :: DataPoint -> StepSize -> LossFunction -> Weight -> PrM Weight 
pureUpdate zs a f = ppure . update zs a f


{-@ measure SGD.SGD.update :: DataPoint -> StepSize -> LossFunction -> Weight -> Weight @-}
{-@ update :: x1:DataPoint -> x2:StepSize -> x3:LossFunction -> x4:Weight 
           -> {v:Weight | v = SGD.SGD.update x1 x2 x3 x4 } @-}
update :: DataPoint -> StepSize -> LossFunction -> Weight -> Weight
update z α f w = w - α * (grad (f z) w) 


-------------------------------------------------------------------------------
-- | Helper Definitions -------------------------------------------------------
-------------------------------------------------------------------------------


{-@ measure lend @-}
{-@ lend :: xs:[a] -> {v:Double| 0.0 <= v } @-}
lend :: [a] -> Double
lend []       = 0
lend (_ : xs) = 1 + lend xs


{-@ reflect one @-}
{-@ one :: {v:Double| v = 1.0 } @-}
one :: Double
one = 1



{-@ reflect head @-}
{-@ head :: {xs:[a] | len xs > 0 } -> a @-}
head :: [a] -> a
head (z : _) = z

{-@ reflect tail @-}
{-@ tail :: {xs:[a] | len xs > 0 } -> {v:[a] | len v == len xs - 1 && lend v == lend xs - 1 } @-}
tail :: [a] -> [a]
tail (_ : zs) = zs

{-@ measure sslen @-}
sslen :: StepSizes -> Int 
{-@ sslen :: StepSizes -> Nat @-}
sslen SSEmp = 0 
sslen (SS _ ss) = 1 + sslen ss