{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DerivingVia           #-}
{-# Language FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# Language MultiParamTypeClasses #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE UndecidableInstances  #-}
module Geometry.Intersections
  ( Intersections (..)
  , intersections, isIntersecting
  ) where

import Data.Complex
import Data.List (nub)

import Geometry.Base
import Geometry.Line
import Geometry.Circle
import Geometry.Polygon
import Geometry.Plot
import Geometry.Decorations

--------------------------------------------------------------------------------
-- | Class provides `intersections'` function returning a list (possible empty)
-- of all potential intersection points (co-dimension 1).
class (Manifold a, Manifold b) => Intersections a b where
  intersections' :: a -> b -> [Cmp]

-- | Filters intersection points, keeping ones belonging to given manifolds.
intersections :: (Figure a, Figure b, Intersections a b)
  => a -> b -> [Cmp]
intersections a b
  | isTrivial a = filter (isContaining b) [cmp $ refPoint a]
  | isTrivial b = filter (isContaining a) [cmp $ refPoint b]
  | otherwise =
    nub $
    filter (isContaining a) $
    filter (isContaining b) $
    intersections' a b

-- | Returns `True` if tho curves have intersection points.
isIntersecting :: (Figure a, Figure b, Intersections a b) => a -> b -> Bool
isIntersecting a b = not . null $ intersections a b

instance (Figure a, Figure b, Intersections a b) =>
  Intersections a (Maybe b) where
  intersections' a = maybe [] (intersections' a)

instance (Figure a, Figure b, Intersections b a) =>
  Intersections (Maybe a) b where
  intersections' = flip intersections'

instance {-# OVERLAPPING #-}
  (Figure a, Figure b, Intersections b a) =>
  Intersections (Maybe a) (Maybe b) where
  intersections' a = maybe [] (intersections' a)

instance (Figure a, Figure b, Intersections a b) =>
         Intersections a (Decorated b) where
  intersections' a = intersections a . fromDecorated

instance (Figure a, Figure b, Intersections b a) =>
         Intersections (Decorated a) b where
  intersections' = flip intersections'

instance {-# OVERLAPPING #-}
  (Figure a, Figure b, Intersections a b) =>
  Intersections (Decorated a) (Decorated b) where
  intersections' = intersections' . fromDecorated

--------------------------------------------------------------------------------

intersectionCC c1 c2 | d == 0 || b < 0 = []
                     | b == 0 = [x]
                     | b > 0 = [x, conjugate x]
  where
    r1 = radius c1
    r2 = radius c2
    d = magnitude (center c1 - center c2)
    a = r1**2-r2**2+d**2
    b = -(r2-r1-d)*(r2-r1+d)*(r2+r1-d)*(r2+r1+d)
    x = scale (1/(2*d)) $ a :+ sqrt b

intersectionCL c l | b > r = []
                   | b == r = [0:+b]
                   | b < r = [a:+b, (-a):+b]
  where
    b = getY (refPoint l)
    r = radius c
    a = sqrt (r**2 - b**2)

--------------------------------------------------------------------------------

instance Intersections Line Line where
  intersections' l1 l2
    | l2 `isContaining` refPoint l1 = [refPoint l1]
    | l1 `isContaining` refPoint l2 = [refPoint l2]
    | otherwise = lineIntersection l1 l2

instance Intersections Line Circle  where
  intersections' l cir =
    invt <$> intersectionCL (t cir) (t l)
    where
      t :: Trans x => x -> x
      t = rotate (-a) . translate' (negate (center cir))
      invt = translate' (center cir) . rotate a
      a = angle l

instance Intersections Line Polyline where
  intersections' l = foldMap (intersections l) . segments . closePolyline

--------------------------------------------------------------------------------

instance Intersections Circle Line where
  intersections' = flip intersections'

instance Intersections Circle Circle where
  intersections' cir1 cir2 =
    invt <$> intersectionCC (t cir1) (t cir2)
    where
      t = rotate (-a) . translate' (negate c1)
      invt = translate' c1 . rotate a
      c1 = center cir1
      c2 = center cir2
      a = azimuth c1 c2

instance Intersections Circle Polyline where
  intersections' c = foldMap (intersections c) . segments . closePolyline

--------------------------------------------------------------------------------

instance Intersections Polyline Line where
  intersections' = flip intersections'

instance Intersections Polyline Circle where
  intersections' = flip intersections'

instance Intersections Polyline Polyline where
  intersections' p = foldMap (intersections p) . segments . closePolyline

--------------------------------------------------------------------------------

deriving via Line instance
  (Manifold a, Intersections a Line) => Intersections a Ray

instance Intersections Ray Line where
  intersections' = intersections' . asLine

instance Intersections Ray Polyline where
  intersections' = intersections' . asLine

instance Intersections Ray Circle where
  intersections' = intersections' . asLine

--------------------------------------------------------------------------------

deriving via Line instance
  (Manifold a, Intersections a Line) => Intersections a Segment

instance Intersections Segment Line where
  intersections' = intersections' . asLine

instance Intersections Segment Polyline where
  intersections' = intersections' . asLine

instance Intersections Segment Circle where
  intersections' = intersections' . asLine

--------------------------------------------------------------------------------

deriving via Polyline instance
  (Manifold a, Intersections a Polyline) => Intersections a Polygon

instance Intersections Polygon Line  where
  intersections' = intersections' . asPolyline

instance Intersections Polygon Polyline  where
  intersections' = intersections' . asPolyline

instance Intersections Polygon Circle  where
  intersections' = intersections' . asPolyline

--------------------------------------------------------------------------------

deriving via Polygon instance
  (Manifold a, Intersections a Polyline) => Intersections a Triangle

instance Intersections Triangle Line  where
  intersections' = intersections' . asPolyline

instance Intersections Triangle Polyline  where
  intersections' = intersections' . asPolyline

instance Intersections Triangle Circle  where
  intersections' = intersections' . asPolyline

--------------------------------------------------------------------------------

deriving via Polygon instance
  (Manifold a, Intersections a Polyline) => Intersections a Rectangle

instance Intersections Rectangle Line  where
  intersections' = intersections' . asPolyline

instance Intersections Rectangle Polyline  where
  intersections' = intersections' . asPolyline

instance Intersections Rectangle Circle  where
  intersections' = intersections' . asPolyline

--------------------------------------------------------------------------------

instance Pnt a => Intersections (Plot a) Line where
  intersections' p l = intersectionsP p l $ intersections (asPolyline p) l

instance Pnt a => Intersections (Plot a) Circle where
  intersections' p c = intersectionsP p c $ intersections (asPolyline p) c

instance Pnt a => Intersections (Plot a) Polyline where
  intersections' p pl = intersectionsP p pl $ intersections (asPolyline p) pl

instance (Pnt a, Manifold b, Intersections (Plot a) b) => Intersections b (Plot a) where
  intersections' = flip intersections'

instance {-# OVERLAPPING #-}
  (Pnt a, Pnt b) => Intersections (Plot a) (Plot b) where
  intersections' p1 p2 = intersectionsP p1 p2 $
                         intersections (asPolyline p1) (asPolyline p2)

--------------------------------------------------------------------------------

deriving via (Plot a) instance
  (Manifold b, Affine a, Trans a, Intersections b (Plot a)) => Intersections b (ClosedPlot a)

instance Pnt a => Intersections (ClosedPlot a) Line  where
  intersections' p l = intersectionsP p l $ intersections (asPolyline p) l

instance Pnt a => Intersections (ClosedPlot a) Polyline  where
  intersections' p pl = intersectionsP p pl $ intersections (asPolyline p) pl

instance Pnt a => Intersections (ClosedPlot a) Circle  where
  intersections' p c = intersectionsP p c $ intersections (asPolyline p) c


{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE PatternSynonyms     #-}
{-# LANGUAGE QuasiQuotes         #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module FramesTutorial where

import qualified Control.Foldl     as L
import qualified Data.Foldable     as F
import qualified Data.Vector       as V
import           Flow
import           Frames
import           Frames.CSV        (RowGen (..), readTableOpt, rowGen)
import           Lens.Micro
import           Lens.Micro.Extras
import           Pipes             hiding (Proxy)
import qualified Statistics        as S

tableTypes'
  (rowGen "data/housing_cleaned.csv")
  { rowTypeName = "House"
  , columnNames =
      [ "crime"
      , "zone"
      , "industrial"
      , "charles"
      , "nox"
      , "rooms"
      , "age"
      , "dist_to_ec"
      , "highway_accessibility"
      , "tax"
      , "pt_ratio"
      , "blacks"
      , "lower_status"
      , "median_value"
      ]
  , separator = ","
  }

houseStream :: MonadSafe m => Producer House m ()
houseStream = readTableOpt houseParser "data/housing_cleaned.csv"

loadHouse :: IO (Frame House)
loadHouse = inCoreAoS houseStream

go = do
    df <- loadHouse
    {-let out = take 5 $ F.toList df-}
    let out =
            view blacks <$> df
            |> F.toList
            |> V.fromList
            |> S.describe
    print out


{-# LANGUAGE CPP #-}

-- | Binary lists are lists whose number of elements is a power of two.
--   This data structure is efficient for some computations like:
--
-- * Splitting a list in half.
--
-- * Appending two lists of the same length.
--
-- * Extracting an element from the list.
--
--   All the functions exported are total except for 'fromListWithDefault'.
--   It is impossible for the user of this library to create a binary list
--   whose length is /not/ a power of two.
--
--   Since many names in this module clash with the names of some "Prelude"
--   functions, you probably want to import this module this way:
--
-- > import Data.BinaryList (BinList,Exponent)
-- > import qualified Data.BinaryList as BL
--
--   Remember that binary lists are an instance of the 'Foldable' and 'Traversable'
--   classes. If you are missing a function here, look for functions using those
--   instances.
--
--   Note that some functions like 'replicate', 'generate', or 'take', don't use
--   the length of the list as argument, but the exponent of its length expressed
--   as a power of two. Throughout this document, this is referred
--   as the /length exponent/. For example, if the list has length 16, its length exponent
--   is 4 since 2^4 = 16. Therefore @replicate 4 0@ will create a list with 16 zeroes.
--   Keep this in mind when using this library. Note as well that this implies that
--   there is no need to check that the length argument is or is not a power of two.
--
module Data.BinaryList (
    -- * Type
    BinList
  , Exponent
    -- * Construction
  , singleton
  , append
  , replicate
  , replicateA
  , replicateAR
  , generate
  , generateM
    -- * Queries
  , lengthExponent
#if !MIN_VERSION_base(4,8,0)
  , length
#endif
  , lookup
  , head
  , last
    -- * Deconstruction
  , split
  , take
  , takeEnd
    -- * Transformation
  , replace
  , reverse
    -- * Tuples
  , joinPairs
  , disjoinPairs
    -- * Zipping and Unzipping
  , zip , unzip
  , zipWith
    -- * Lists
    -- ** From list
  , fromList
  , fromListWithDefault
  , fromListSplit
    -- ** To list
  , toListFilter
  , toListSegment
    -- * Others
  , traverseSegment
    -- * Example: Radix-2 FFT
    -- $fft
  ) where

import Prelude hiding
  ( map, head, last
  , reverse, replicate
  , take, lookup
  , zip, unzip, zipWith
  , foldr1, length, foldr
    )
#if MIN_VERSION_base(4,8,0)
import Prelude (length, foldr, foldr1)
#endif
import Data.Foldable (fold,toList)
import Foreign.Storable (sizeOf)
import Data.List (find)
import Data.BinaryList.Internal
import Control.Applicative.Backwards
import Control.Arrow ((***))
import Control.Monad.Trans.State
  ( StateT (..)
  , evalStateT ,evalState
  , runState
  , get ,modify )
import Control.Monad.Trans.Class (lift)
import Data.Functor.Identity (Identity (..))
import Control.Applicative.PhantomState
import Data.Word (Word8)

-- GHC-7.8 compatibility
#if !MIN_VERSION_base(4,8,0)

import qualified Prelude
import Data.Foldable (Foldable (..))
import Data.Traversable (Traversable (..))
import Control.Applicative (Applicative (..), (<$>))
import Data.Monoid (mappend)

-- | /O(1)/. Number of elements in the list.
length :: BinList a -> Int
length = (2^) . lengthExponent

#endif

-- | An exponent.
type Exponent = Word8

-- | /O(1)/. Build a list with a single element.
singleton :: a -> BinList a
singleton = ListEnd

-- | /O(1)/. Given a binary list @l@ with length @2^k@:
--
-- > lengthExponent l = k
--
lengthExponent :: BinList a -> Exponent
lengthExponent (ListNode n _ _) = n
lengthExponent (ListEnd _) = 0

{-# RULES
       "Data.BinaryList: length equality"
         forall xs ys . length xs == length ys = lengthExponent xs == lengthExponent ys
  #-}

-- | /O(log n)/. Lookup an element in the list by its index (starting from 0).
--   If the index is out of range, 'Nothing' is returned.
lookup :: BinList a -> Int -> Maybe a
lookup (ListNode n l r) i =
   let m = 2^(n-1) -- Number of elements in a single branch
   in  if i < m
          then lookup l i       -- Lookup in the left branch
          else lookup r $ i - m -- Lookup in the right branch
lookup (ListEnd x) 0 = Just x
lookup _ _ = Nothing

-- | /O(log n)/. Replace a single element in the list. If the index is
--   out of range, returns the original list.
replace :: Int -- ^ Index to look for
        -> a   -- ^ Element to insert
        -> BinList a -> BinList a
replace i0 y = go i0
  where
    go i (ListNode n l r) =
      let m = 2^(n-1)
      in  if i < m
             then ListNode n (go i l)          r
             else ListNode n       l (go (i-m) r)
    go 0 (ListEnd _) = ListEnd y
    go _ e = e

-- | /O(1)/. Append two binary lists. This is only possible
--   if both lists have the same length. If this condition
--   is not hold, 'Nothing' is returned.
append :: BinList a -> BinList a -> Maybe (BinList a)
append xs ys =
  let i = lengthExponent xs
  in  if i == lengthExponent ys
         then Just $ ListNode (i+1) xs ys
         else Nothing

-- | /O(1)/. Split a binary list into two sublists of half the length,
--   unless the list only contains one element. In that case, it
--   just returns that element.
split :: BinList a -> Either a (BinList a,BinList a)
split (ListNode _ l r) = Right (l,r)
split (ListEnd x) = Left x

-- | /O(log n)/. Calling @take n xs@ returns the first @min (2^n) (length xs)@ elements of @xs@.
take :: Exponent -> BinList a -> BinList a
take k xs@(ListNode n l _) = if k >= n then xs else take k l
take _ xs = xs

-- | /O(log n)/. Calling @takeEnd n xs@ returns the last @min (2^n) (length xs)@ elements of @xs@.
takeEnd :: Exponent -> BinList a -> BinList a
takeEnd k xs@(ListNode n _ r) = if k >= n then xs else takeEnd k r
takeEnd _ xs = xs

-- | Calling @replicateA n f@ builds a binary list collecting the results of
--   executing @2^n@ times the applicative action @f@.
replicateA :: Applicative f => Exponent -> f a -> f (BinList a)
replicateA n f = go n
  where
    go 0 = ListEnd <$> f
    go i = let b = go (i-1)
           in  ListNode <$> pure i <*> b <*> b

-- | The same as 'replicateA', but the actions are executed in reversed order.
replicateAR :: Applicative f => Exponent -> f a -> f (BinList a)
replicateAR n = forwards . replicateA n . Backwards

{-# RULES
      "Data.BinaryList: map reverse/replicateA"
         forall i f . map reverse (replicateA  i f) = replicateAR i f
  #-}

{-# RULES
      "Data.BinaryList: map reverse/replicateAR"
         forall i f . map reverse (replicateAR i f) = replicateA  i f
  #-}

-- | /O(log n)/. Calling @replicate n x@ builds a binary list with
--   @2^n@ occurences of @x@.
replicate :: Exponent -> a -> BinList a
replicate n = runIdentity . replicateA n . Identity

{-# RULES
      "Data.BinaryList: map/replicate"
         forall f n x . map f (replicate n x) = replicate n (f x)
  #-}

-- | /O(n)/. Build a binary list with the given length exponent (see 'lengthExponent')
--   by applying a function to each index.
generate :: Exponent -> (Int -> a) -> BinList a
generate l f = evalState (replicateA l $ fmap f get <* modify (+1)) 0

-- | Like 'generate', but the generator function returns a value in a 'Monad'.
--   Therefore, the result is as well contained in a 'Monad'.
generateM :: (Applicative m, Monad m) => Exponent -> (Int -> m a) -> m (BinList a)
generateM l f = evalStateT (replicateA l $ (get >>= lift . f) <* modify (+1)) 0

-- | /O(log n)/. Get the first element of a binary list.
head :: BinList a -> a
head (ListNode _ l _) = head l
head (ListEnd x) = x

-- | /O(log n)/. Get the last element of a binary list.
last :: BinList a -> a
last (ListNode _ _ r) = last r
last (ListEnd x) = x

-- | /O(n)/. Reverse a binary list.
reverse :: BinList a -> BinList a
{-# INLINE[2] reverse #-}
reverse (ListNode n l r) = ListNode n (reverse r) (reverse l)
reverse xs = xs

{-# RULES
      "Data.BinaryList: reverse/reverse"
         forall xs. reverse (reverse xs) = xs
  #-}

------------------------------
-- Transformations with tuples

-- | /O(n)/. Transform a list of pairs into a flat list. The
--   resulting list will have twice more elements than the
--   original.
joinPairs :: BinList (a,a) -> BinList a
{-# INLINE[1] joinPairs #-}
joinPairs (ListEnd (x,y)) = ListNode 1 (ListEnd x) (ListEnd y)
joinPairs (ListNode n l r) = ListNode (n+1) (joinPairs l) (joinPairs r)

-- | /O(n)/. Opposite transformation of 'joinPairs'. It halves
--   the number of elements of the input. As a result, when
--   applied to a binary list with a single element, it returns
--   'Nothing'.
disjoinPairs :: BinList a -> Maybe (BinList (a,a))
{-# INLINE [1] disjoinPairs #-}
disjoinPairs (ListEnd _) = Nothing
disjoinPairs xs = Just $ disjoinPairsNodes xs

disjoinPairsNodes :: BinList a -> BinList (a,a)
disjoinPairsNodes (ListNode _ (ListEnd x) (ListEnd y)) = ListEnd (x,y)
disjoinPairsNodes (ListNode n l r) = ListNode (n-1) (disjoinPairsNodes l) (disjoinPairsNodes r)
disjoinPairsNodes _ = error "disjoinPairsNodes: bug. Please, report this with an example input."

{-# RULES
      "Data.BinaryList: disjoinPairs/joinPairs"
         forall xs . disjoinPairs (joinPairs xs) = Just xs
  #-}

{-# RULES
      "Data.BinaryList: disjoinPairs/map/joinPairs"
         forall f xs . disjoinPairs (map f (joinPairs xs)) = Just (map (f *** f) xs)
  #-}

-- | Expression @pairBuilder f xs@ is equivalent to @joinPairs (map f xs)@, but does
--   not build any intermediate structure. Used for rewriting rules.
pairBuilder :: (a -> (b,b)) -> BinList a -> BinList b
{-# INLINE[0] pairBuilder #-}
pairBuilder f = go
  where
    go (ListEnd x) = let (a,b) = f x in ListNode 1 (ListEnd a) (ListEnd b)
    go (ListNode n l r) = ListNode (n+1) (go l) (go r)

{-# RULES
      "Data.BinaryList: joinPairs/map"
         forall f xs . joinPairs (map f xs) = pairBuilder f xs
  #-}

-- | Expression @zipAndJoing f g xs ys@ is equivalent to @pairBuilder f (zipWith g xs ys)@,
--   but does not build any intermediate structure. Used for rewriting rules.
zipAndJoin :: (c -> (d,d)) -> (a -> b -> c) -> BinList a -> BinList b -> BinList d
zipAndJoin f g = go
  where
    -- Recursion
    go xs@(ListNode n l r) ys@(ListNode n' l' r')
         -- If both lists have the same length, recurse assuming it
         -- to avoid comparisons.
       | n == n'   = ListNode (n+1) (goEquals l l') (goEquals r r')
         -- If the first list is larger, the second fits entirely in
         -- the left branch of the first.
       | n >  n'   = go l ys
         -- If the second list is larger, the first fits entirely in
         -- the left branch of the second.
       | otherwise = go xs l'
    go xs ys       = let (x,y) = f $ g (head xs) (head ys)
                     in  ListNode 1 (ListEnd x) (ListEnd y)
    -- Recursion assuming both lists have the same length
    goEquals (ListNode n l r) (ListNode _ l' r') =
                     ListNode (n+1) (goEquals l l') (goEquals r r')
    goEquals (ListEnd x) (ListEnd y) =
      let (x',y') = f $ g x y
      in  ListNode 1 (ListEnd x') (ListEnd y')
    goEquals _ _ = undefined -- This can't happen!

{-# RULES
      "Data.BinaryList: pairBuilder/zipWith"
         forall f g xs ys . pairBuilder f (zipWith g xs ys) = zipAndJoin f g xs ys
  #-}

------------------------
-- Zipping and Unzipping

-- | /O(n)/. Zip two binary lists using an operator.
zipWith :: (a -> b -> c) -> BinList a -> BinList b -> BinList c
{-# INLINE[1] zipWith #-}
zipWith f = go
  where
    -- Recursion
    go xs@(ListNode n l r) ys@(ListNode n' l' r')
         -- If both lists have the same length, recurse assuming it
         -- to avoid comparisons.
       | n == n'   = ListNode n (goEquals l l') (goEquals r r')
         -- If the first list is larger, the second fits entirely in
         -- the left branch of the first.
       | n >  n'   = go l ys
         -- If the second list is larger, the first fits entirely in
         -- the left branch of the second.
       | otherwise = go xs l'
    go xs ys       = ListEnd $ f (head xs) (head ys)
    -- Recursion assuming both lists have the same length
    goEquals (ListNode n l r) (ListNode _ l' r') =
                     ListNode n (goEquals l l') (goEquals r r')
    goEquals (ListEnd x) (ListEnd y) = ListEnd (f x y)
    goEquals _ _ = undefined -- this can't happen

-- | /O(n)/. Zip two binary lists in pairs.
zip :: BinList a -> BinList b -> BinList (a,b)
{-# INLINE zip #-}
zip = zipWith (,)

-- | /O(n)/. Unzip a binary list of pairs.
unzip :: BinList (a,b) -> (BinList a, BinList b)
{-# INLINE[1] unzip #-}
unzip (ListEnd (x,y)) = (ListEnd x, ListEnd y)
unzip (ListNode n l r) =
  let (la,lb) = unzip l
      (ra,rb) = unzip r
  in  (ListNode n la ra, ListNode n lb rb)

-- | Expression @unzipMap f xs@ is equivalent to @unzip (map f xs)@, but
--   does not create any intermediate structure.
unzipMap :: ((a,b) -> (c,d)) -> BinList (a,b) -> (BinList c,BinList d)
unzipMap f = go
  where
    go (ListEnd p) = ListEnd *** ListEnd $ f p
    go (ListNode n l r) =
      let (lc,ld) = go l
          (rc,rd) = go r
      in  (ListNode n lc rc, ListNode n ld rd)

{-# RULES
      "Data.BinaryList: unzip/map"
         forall f xs . unzip (map f xs) = unzipMap f xs
  #-}

-----------------------------
-- Transforming from/to lists

-- | /O(log n)/. Calculate the exponent of a positive integer number expressed
--   as a power of two.
exponentInBasisTwo :: Int -> Maybe Exponent
exponentInBasisTwo 1 = Just 0
exponentInBasisTwo n =
  if even n
     then fmap (+1) $ exponentInBasisTwo $ div n 2
     else Nothing

-- | /O(n)/. Build a binary list from a linked list. If the input list
--   has length different from a power of two, it returns 'Nothing'.
fromList :: [a] -> Maybe (BinList a)
fromList xs = fmap builder . exponentInBasisTwo $ Prelude.length xs
  where
    builder l = evalState (replicateA l $ StateT $ \(h:t) -> pure (h,t)) xs

-- | /O(1)/. This is the last exponent that has power of two defined in the type 'Int'.
--
-- /Note: This value is system dependent, since the type 'Int' varies in size/
-- /from system to system./
--
lastExponentOfTwo :: Exponent
lastExponentOfTwo = fromIntegral $ 8 * sizeOf (undefined :: Int) - 2

-- | /O(1)/. Calculate the next power of two exponent, if there is any. It is possible
--   to not find a next one since the type 'Int' is finite. If the input is
--   already a power of two, its exponent is returned.
nextExponentOfTwo :: Int -> Maybe Exponent
nextExponentOfTwo n = find (\i -> n <= 2^i) [0 .. lastExponentOfTwo]

-- | /O(n)/. Build a binary list from a linked list. If the input list
--   has length different from a power of two, fill to the next
--   power of two with a default element.
--
-- /Warning: this function crashes if the input list length is larger than any/
-- /power of two in the type 'Int'. However, this is very unlikely./
fromListWithDefault :: a -> [a] -> BinList a
fromListWithDefault e xs =
  let l = Prelude.length xs
  in  case nextExponentOfTwo l of
        Just n ->
          evalState (replicateA n $ StateT $
             \ys -> pure $ case ys of
                      (h:t) -> (h,t )
                      _     -> (e,[])
               ) xs
        _ -> error "[binary-list] fromListWithDefault: input list is too big."

-- | /O(n)/. Build a binary list from a linked list. It returns a binary list
--   with length @2 ^ n@ (where @n@ is the supplied 'Int' argument), and
--   the list of elements of the original list that were not used. If the
--   input list is shorter than @2 ^ n@, a default element will be used to
--   complete the binary list. This method for building binary lists is faster
--   than both 'fromList' and 'fromListWithDefault'.
fromListSplit :: a   -- ^ Default element
              -> Exponent -- ^ Length exponent
              -> [a] -- ^ Input list
              -> (BinList a, [a])
fromListSplit e n =
  runState $ replicateA n $ StateT $
    \xs -> pure $ case xs of
                    (h:t) -> (h,t )
                    _     -> (e,[])

-- | /O(n)/. Create a list from the elements of a binary list matching a given
--   condition.
toListFilter :: (a -> Bool) -> BinList a -> [a]
{-# INLINE toListFilter #-}
toListFilter c = foldr (\x -> if c x then (x:) else id) []

-- | /O(n)/. Create a list extracting a sublist of elements from a binary list.
toListSegment :: Int -> Int -> BinList a -> [a]
{-# INLINE toListSegment #-}
toListSegment s e xs = runPhantomState (traverseSegment (changeState . (:)) s e xs) []

-- | Apply an applicative action to every element in a segment of a binary list, from left to right.
traverseSegment :: Applicative f => (a -> f ()) -> Int -> Int -> BinList a -> f ()
{-# INLINE traverseSegment #-}
traverseSegment f s e xs
  | s > e = pure ()
  | e < 0 = pure ()
  | s >= length xs = pure ()
  | otherwise = traverseSegmentFromTo f (max 0 s) e xs

traverseSegmentFromTo :: Applicative f => (a -> f ()) -> Int -> Int -> BinList a -> f ()
{-# INLINE traverseSegmentFromTo #-}
traverseSegmentFromTo f = go
  where
    go s e (ListNode n l r) =
      let k = 2^(n-1)
      in  if s >= k
             -- Sublist is contained in right portion
             then go (s - k) (e - k) r
             else if e < k
                     -- Sublist is contained in left portion
                     then go s e l
                     -- Sublist is divided in both portions
                     else traverseSegmentFrom f s l *> traverseSegmentTo f (e - k) r
    go _ _ (ListEnd x) = f x

traverseSegmentFrom :: Applicative f => (a -> f ()) -> Int -> BinList a -> f ()
{-# INLINE traverseSegmentFrom #-}
traverseSegmentFrom f = go
  where
    go s (ListNode n l r) =
      let k = 2^(n-1)
      in  if s >= k
             -- Sublist is contained in right portion
             then go (s - k) r
             -- Sublist is divided in both portions, but right
             -- portion is taken entirely
             else go s l *> traverseFull f r
    go _ (ListEnd x) = f x

traverseSegmentTo :: Applicative f => (a -> f ()) -> Int -> BinList a -> f ()
{-# INLINE traverseSegmentTo #-}
traverseSegmentTo f = go
  where
    go e (ListNode n l r) =
      let k = 2^(n-1)
      in  if e < k
             -- Sublist is contained in left portion
             then go e l
             -- Sublist is divided in both portions, but left
             -- portion is taken entirely
             else traverseFull f l *> go (e - k) r
    go _ (ListEnd x) = f x

traverseFull :: Applicative f => (a -> f ()) -> BinList a -> f ()
{-# INLINE traverseFull #-}
traverseFull f = go
  where
    go (ListEnd x) = f x
    go (ListNode _ l r) = go l *> go r

------------------------------------------------
------------------------------------------------

-----------------------------
-- Show and Functor instances

instance Show a => Show (BinList a) where
  show = show . toList

{- Internal map

Although we encourage the use of 'fmap', we define fmap as a custom 'map'
function and inline 'fmap' to make them equivalent, so writing 'fmap' is
actually writing 'map'. We do this to use 'map' in rewriting rules.

-}

map :: (a -> b) -> BinList a -> BinList b
{-# INLINE[1] map #-}
map f = go
  where
    go (ListEnd x) = ListEnd (f x)
    go (ListNode n l r) = ListNode n (go l) (go r)

instance Functor BinList where
  {-# INLINE fmap #-}
  fmap = map

instance Foldable BinList where
  -- Folding
  foldr1 f = go
    where
      go (ListEnd x) = x
      go (ListNode _ l r) = f (go l) (go r)
  --
  fold = foldr1 mappend
  foldl1 = foldr1
  foldMap f = fold . fmap f
#if MIN_VERSION_base(4,8,0)
  length = (2^) . lengthExponent
#endif

instance Traversable BinList where
  sequenceA (ListEnd f) = ListEnd <$> f
  sequenceA (ListNode n l r) = ListNode <$> pure n <*> sequenceA l <*> sequenceA r

-----------------------------
-- Example: Radix-2 FFT

{- $fft

This is an example demonstrating the use of binary lists to calculate the Discrete
Fourier Transform of complex vectors with the Radix-2 Fast Fourier Transform algorithm.

> import Data.BinaryList (BinList)
> import qualified Data.BinaryList as BL
>
> import Data.Complex
> import Data.Maybe (fromJust)
>
> i :: Complex Double
> i = 0 :+ 1
>
> fft :: BinList (Complex Double) -> BinList (Complex Double)
> fft xs =
>   case BL.disjoinPairs xs of
>     Nothing -> xs
>     Just ps ->
>       let (evens,odds) = BL.unzip ps
>           n = BL.lengthExponent xs - 1
>           q = negate $ pi * i / fromIntegral (2^n)
>           twiddles = BL.generate n $ \k -> exp $ q * fromIntegral k
>           oddsfft = BL.zipWith (*) twiddles $ fft odds
>           evensfft = fft evens
>       in  fromJust $
>             BL.append (BL.zipWith (+) evensfft oddsfft)
>                       (BL.zipWith (-) evensfft oddsfft)

-}


module Data.Random.Distribution.Hypergeometric.Impl where

import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Uniform

import Numeric.SpecFunctions (logGamma)

d1 = 1.7155277699214135 -- 2*sqrt(2/e)
d2 = 0.8989161620588988 -- 3 - 2*sqrt(3/e)

rhyper :: (Integral a) => (a, a, a) -> RVarT m a
rhyper (sample, popsize, good) = return . fix2 . fix1 =<< loop
  where fix1 z = if good > bad then m - z else z
        fix2 z = if m < sample then good - z else z
        mingoodbad = min good bad
        bad = popsize - good
        maxgoodbad = max good bad
        m = min sample (popsize - sample)
        gamma z = sum $ map (logGamma . fromIntegral)
          [ z + 1,     mingoodbad - z + 1
          , m - z + 1, maxgoodbad - m + z + 1 ]
        d4 = fromIntegral mingoodbad / fromIntegral popsize
        d5 = 1 - d4
        d6 = fromIntegral m * d4 + 0.5
        d7 = sqrt $ fromIntegral (popsize - m) * fromIntegral sample * d4 * d5 / fromIntegral (popsize - 1) + 0.5
        d8 = d1 * d7 + d2
        d9 = floor $ fromIntegral (m + 1) * fromIntegral (mingoodbad + 1) / fromIntegral (popsize + 2)
        d10 = gamma d9
        d11 = fromIntegral $ min (min m mingoodbad + 1) (floor (d6 + 16 * d7))
        -- 16 for 16-decimal-digit precision in d1 and d2
        loop = do
          x <- doubleUniform 0 1
          y <- doubleUniform 0 1
          case () of
            _  | w < 0 || w >= d11    -> loop -- fast rejection
               | x * (4 - x) - 3 <= t -> return z -- fast acceptance
               | x * (x - t) >= 1     -> loop -- fast rejection
               | 2 * (log x) <= t     -> return z -- acceptance
               | otherwise            -> loop -- rejection
               where w = d6 + d8 * (y - 0.5) / x
                     z = floor w
                     t = d10 - gamma z


{-# LANGUAGE RecordWildCards #-}
module Q.Options.Black76
  (
    module Q.Options
  , Black76(..)
  , atmf
  , euOption
  , eucall
  , euput
  )
  where

import           Q.Options
import           Q.Types
import           Statistics.Distribution        (cumulative, density)
import           Statistics.Distribution.Normal (standard)

data Black76 = Black76 {
    b76F   :: Forward
  , b76DF  :: DF
  , b76T   :: YearFrac
  , b76Vol :: Vol
}

-- | At the money forward strike.
atmf :: Black76 -> Strike
atmf Black76{..} = Strike f
  where (Forward f) = b76F

-- | European option valuation with black 76
euOption :: Black76 -> OptionType -> Strike -> Valuation
euOption b76@Black76{..} cp k = Valuation premium delta vega gamma where
  (Forward f) = b76F
  n           = cumulative standard
  (Vol sigmaSqt) = scale b76T b76Vol
  d1          = (dPlus  b76F b76Vol k b76T)
  d2          = (dMinus b76F b76Vol k b76T)
  nd1         = n d1
  nd2         = n d2
  callDelta   = b76DF `discount` nd1
  putDelta    = b76DF `discount` (- (n (-d1)))
  vega        = Vega  $ b76DF `discount` (density standard d1 ) * f * sigmaSqt
  gamma       = Gamma $ b76DF `discount` (density standard d1) / (f * sigmaSqt)
  premium  = Premium $ case cp of
    Call -> b76DF `discount` (f * nd1 - nd2 * k')
    Put  -> b76DF `discount` (n (-d2) * k' - n (-d1) * f)
    where (Strike k') = k
  delta | cp == Call = Delta $ callDelta
        | cp == Put  = Delta $ putDelta

-- | see 'euOption'
euput b76 =  euOption b76 Put

-- | see 'euOption'
eucall b76 = euOption b76 Call

dPlus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) + (0.5 * sigma * sigma) * t)
dMinus (Forward f) (Vol sigma) (Strike k) (YearFrac t) =
  recip (sigma * sqrt t) * (log (f/k) - (0.5 * sigma * sigma) * t)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE RankNTypes            #-}

module Grenade.Utils.OneHot (
    oneHot
  , hotMap
  , makeHot
  , unHot
  , sample
  ) where

import qualified Control.Monad.Random as MR

import           Data.List ( group, sort )

import           Data.Map ( Map )
import qualified Data.Map as M

import           Data.Proxy
import           Data.Singletons.TypeLits

import           Data.Vector ( Vector )
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS

import           Numeric.LinearAlgebra ( maxIndex )
import           Numeric.LinearAlgebra.Devel
import           Numeric.LinearAlgebra.Static

import           Grenade.Core.Shape

-- | From an int which is hot, create a 1D Shape
--   with one index hot (1) with the rest 0.
--   Rerurns Nothing if the hot number is larger
--   than the length of the vector.
oneHot :: forall n. (KnownNat n)
       => Int -> Maybe (S ('D1 n))
oneHot hot =
  let len = fromIntegral $ natVal (Proxy :: Proxy n)
  in if hot < len
      then
        fmap S1D . create $ runSTVector $ do
        vec    <- newVector 0 len
        writeVector vec hot 1
        return vec
      else Nothing

-- | Create a one hot map from any enumerable.
--   Returns a map, and the ordered list for the reverse transformation
hotMap :: (Ord a, KnownNat n) => Proxy n -> [a] -> Either String (Map a Int, Vector a)
hotMap n as =
  let len  = fromIntegral $ natVal n
      uniq = [ c | (c:_) <- group $ sort as]
      hotl = length uniq
  in if hotl == len
      then
        Right (M.fromList $ zip uniq [0..], V.fromList uniq)
      else
        Left ("Couldn't create hotMap of size " ++ show len ++ " from vector with " ++ show hotl ++ " unique characters")

-- | From a map and value, create a 1D Shape
--   with one index hot (1) with the rest 0.
--   Rerurns Nothing if the hot number is larger
--   than the length of the vector or the map
--   doesn't contain the value.
makeHot :: forall a n. (Ord a, KnownNat n)
        => Map a Int -> a -> Maybe (S ('D1 n))
makeHot m x = do
  hot    <- M.lookup x m
  let len = fromIntegral $ natVal (Proxy :: Proxy n)
  if hot < len
      then
        fmap S1D . create $ runSTVector $ do
        vec    <- newVector 0 len
        writeVector vec hot 1
        return vec
      else Nothing

unHot :: forall a n. KnownNat n
      => Vector a -> S ('D1 n) -> Maybe a
unHot v (S1D xs)
  = (V.!?) v
  $ maxIndex (extract xs)

sample :: forall a n m. (KnownNat n, MR.MonadRandom m)
       => Double -> Vector a -> S ('D1 n) -> m a
sample temperature v (S1D xs) = do
  ix <- MR.fromList . zip [0..] . fmap (toRational . exp . (/ temperature) . log) . VS.toList . extract $ xs
  return $ v V.! ix


module Object.Vertex (Vertex (..), createVertex) where

import Numeric.Vector
import Types

createVertex :: Vec3d -> Vec3d -> Vec2d -> Vec3d -> Vec3d -> Vertex
createVertex pos color texcoord normal tangent = Vertex { v_getPos = pos
                                                        , v_getColor = color
                                                        , v_getTexCoord = texcoord
                                                        , v_getNormal = normal
                                                        , v_getTangent = tangent
                                                        , v_getBitangent = tangent × normal
                                                        }


{-# LANGUAGE FlexibleContexts, TypeOperators, RankNTypes, GeneralizedNewtypeDeriving, DeriveFunctor, PolyKinds, GADTs #-}

module Qubit where

import Linear.V2
import Linear.V1
import Linear.Vector
import Vec
import Data.Complex
import Data.Functor.Compose
import Data.Functor.Product
import Data.Functor.Identity
import Cat
import Data.Coerce

type Kron = Compose
type Spin = V2
type QuBit = V2
type C = Complex Double

type DSum = Product

newtype FKron f g a = FKron [(f a, g a)]

type Fock a =  Q [a] -- a "fock space" over a. No assumption of anti symmettric or symmettric.
newtype Anti a = Anti (Fock a)
-- for antisymmettric space, keep list sorted (canonical form).
-- liftAnti1 :: (a -> Q a) -> ([a] -> Q [a])
-- liftAnti2 :: (a,a) -> Q (a,a) -> ([a] -> Q [a])
--

-- this is fock system as a sparse indictaor matrix. We can also do the form of occupation number  Q (V Bool)

vac :: Fock a
vac = pure []

type SHO = Q Int

shovac :: SHO
shovac = pure 0
adag :: Int -> SHO
adag n = W [(n + 1, sqrt (fromInteger (toInteger n) + 1))]
agad 0 = mempty
agad n = W [(n - 1, sqrt n)]

raise = LinOp adag
lower = LinOp agad

-- boson system can be seen as collection of oscillators. Q (V Int).  Positive integers really.

type Boson v = Q (v Int) -- accepts a functor describing all the oscillators. A Vec for example.
-- bvac = pure zero -- if v is a Vector type, it probably has a zero.
--

-- traversable ordering. The vector is put into canocnical form with adag appearing first in the traversal coming first.
type Fermion v = Q (v Bool) -- fermion needs to prefix count to figure out signs. This means that v has to be traversable? Foldable?
fvac :: Applicative v => Fermion v
fvac = pure (pure False)
{-
-- what about systems with different types of particles.
particle-hole  Q ([a], [a])

AntiOp a = AntiOp [a] -> Q [a]
-- indexed linear operators. a -> AntiOp a
psi :: Ord a => a -> (a -> Anti a)
psidag :: Ord a => a -> (a -> Anti a)

phi :: Ord a => a -> (a -> Sym a)
phidag :: Ord a => a -> (a -> Sym a)

-- hmm. Something kind of funny about these.
lift1 :: (i -> Double) -> (i -> Fock i)
lift2 :: (i -> i -> Double) -> (i -> Fock i)
symlift -- mention that there is implcit symmettrization
2 natural lifts. cnnected by logarithms. UxUxUxU vs HIII + IHII + IIHI + IIIH

sum1 :: (i -> i -> Q i) -> (i -> Q i) -- takes an indexed operator and turns it into an operator
sum1 :: (i -> AntiOp i) -> (AntiOp i)
sum2
sum :: Addable b => (i -> b) -> b

turn all of this into a typeclass functions. Then do this form, but also full initial form for manipulation.

dyson

-}
-- instance (Additive f, Additive g) => Additive (Compose f g) where
-- instance Num f (g a) => Num (Compose f g a)

-- (Kron (Kron V2 V2) V2) Complex Double  -- a 3 qubit vector space
swap :: (Traversable f, Applicative g) => Kron f g a -> Kron g f a
swap = Compose . sequenceA . getCompose

-- kron :: (Num a, Functor f, Functor g) => f a -> g a -> Kron f g a
-- kron f g = Compose $ fmap (\amp1 -> fmap (\amp2 -> amp1 * amp2) g) f
kron' :: (Num a, Functor f, Functor g) => f a -> g a -> Kron f g a
kron' x y = Compose (outer x y)

densify ::  (Additive (Kron f g), Num a, Functor f, Functor g) => FKron f g a -> Kron f g a
densify (FKron xs) = sumV $ map (uncurry kron') xs

kron''' :: (Applicative f', Applicative g', Traversable f, Traversable g') =>
    (f a -> f' a) -> (g a -> g' a) -> (Kron f g a -> Kron f' g' a)
kron''' lf lg (Compose fga) = Compose $ sequenceA $ (fmap lf) $ sequenceA  $ (fmap lg fga)

type LinOp' f g a = f a -> g a
-- SVD let's you re free
-- but also [indicatro vectro, g a]
-- and perhaps randomize sampling.
    {-
sigx :: Spin C -> Spin C
sigx (V2 up down) = V2 down up

sigz :: Spin C -> Spin C
sigz (V2 (a :+ b) (c :+ d)) = V2 ()

sigz :: Spin C -> Spin C
sigz (V2 up down) = V2 up (-down)
-}
{-
class Add a where
    (+++) :: a -> a -> a
    negate' :: a -> a

instance Num a => Add a where
    +++ = +
instance    => Additive f
-}

flip' :: Spin a -> Spin a
flip' (V2 up down) = V2 down up

sigx = flip'

sigz :: Num a => Spin a -> Spin a
sigz (V2 up down) = V2 up (negate down)

-- phase :: Complex a => Spin a -> Spin a
type f ~> g = forall a. (Functor f, Functor g) => f a -> g a
type f *** g = Compose f g
-- (Functor f, Functor g, Functor h, Functor k)
-- we only need g to be a functor
-- see below. If we remove the type signature, this implementation generalizes to another type
par :: Functor g => (forall a. f a -> g a) -> (forall b. h b -> k b) -> (forall c. (Compose f h c) -> (Compose g k c))
par f g = Compose . (fmap g) . f . getCompose -- insufficient though. We can't write linear operators this way

-- I should used V1, not Identity
-- These are all also coerce.
rightUnitor :: Functor f => Compose f Identity a -> f a
rightUnitor (Compose f) = fmap runIdentity f

rightUnitor' :: f ~> Compose f Identity
rightUnitor' = Compose . fmap Identity

leftUnitor' :: f ~> Compose Identity f
leftUnitor' = Compose . Identity

leftUnitor :: Identity *** f ~> f
leftUnitor = runIdentity . getCompose

-- should be coerce. Confusing role problems
assoc :: Functor f => (Kron (Kron f g) h) ~> (Kron f (Kron g h))
assoc = Compose . (fmap Compose) . getCompose . getCompose

assoc' :: Functor f =>  (Kron f (Kron g h)) ~> (Kron (Kron f g) h)
assoc' (Compose x)  = Compose $ Compose $ (fmap getCompose x)

-- if the functor in question happens to be representable
    -- then we can use our previous defitions of monoidal catgoeries.
    -- Reader Functors
-- Num c => prod :: Reader a c -> Reader b c -> Reader (a,b) c
-- Num
-- newtype ReadCat c a b = ReadCat (Reader a c -> Reader b c)
-- newtype LinOp c a b = LinOp ((a -> c) -> (b -> c))
-- par (LinOp f) (LinOp g) = LinOp \h -> h (\(a,b) ->   )
-- FMapN
-- fmapN @n
{-Compose f g

assoc :: Compose (Compose f g) h a -> Compose f (Compose g h) a

assoc = Compose . (fmap Compose) . getCompose . getCompose

Compose f Id a -> f a

rightunit = (fmap runIdentity) . getCompose

The morphisms of the functor category are natural transformations

par :: (f ~> g) -> (h~> k) -> ((Compose f h) ~> (Compose g k))

-}

newtype Kron' k l f a = Kron' ((k (l f)) a) deriving Functor

-- deriving via compose
{-
newtype V2' f a = V2' (Compose V2 f a)  deriving Functor
newtype V1' f a = V1' (Compose V1 f a) deriving Functor
-}
{-
newtype V2' f a = V2' (V2 (f a))  deriving Functor
newtype V1' f a = V1' (f a) deriving Functor
-}

type V2' f a = Kron V2 f a
type V1' f a = Kron V1 f a

lowerv1 (Compose (V1 f)) = f -- It's a coerce. V1 is assuredly a newtype
{-
exvsa :: (Additive f, Num a) => V2' f a
exvsa = zero
-}

-- constant vecotrs a morphisms from unit
-- Using f ~ one = (V1 1)
exconst :: (a ~ Double, Functor f) => V1' f a -> V2' f a
exconst x = let x' = lowerv1 x in Compose $ V2 (1.4 *^ x') (7.3 *^ x')

one = V1 1
attemp = exconst (Compose (V1 one))

-- Applicative composes, so Additive probably should too.
-- Wait is this in the library?
-- Why is applicative g necessary?
{-
Ok. There doesn't appear to be a compose
But V2 etc have a Num instance
and liftU2 has a default implementation that uses Applicative
I probably don't need Applcative g if I define a liftU2 using g's liftU2

-}
-- I should check to make sure the lift makes sense
-- Why isn't this in the library? There must be a good reason.
-- Is there an additive instance for Product?

instance (Additive g, Additive f) => Additive (Compose f g) where
    (Compose v) ^+^ (Compose w) = Compose (liftU2 (^+^) v w) -- I  we might be able to used
    zero = zero -- what?
    liftU2 f (Compose x) (Compose y) = Compose $ liftU2 (liftU2 f) x y
    liftI2 f (Compose x) (Compose y) = Compose $ liftI2 (liftI2 f) x y

type LinOp1 f g a = forall k. Additive k => Kron f k a -> Kron g k a

-- This feels less right to me. But why?
type LinOp1' f g a = forall k. Functor k => Kron k f a -> Kron k g a
{-
It feels less right perhaps because it will change how we do things so little.
We just have to fmap.
No, but saying we always have direct access to the scalar isn't right.
-}

-- reassoc :: Kron f (Kron g k) a -> Kron (Kron f g) k a
-- reassoc = coerce

--reassoc' ::  Kron (Kron f g) k a -> Kron f (Kron g k) a
--reassoc' = -- coerce

-- yikes. coerce will help a bit. Quite a bit.
-- Also Fixing the unnecessary applicative instances.
-- Additive k' is expected. Applicative is coming from Additve (Kron )  constraint.

{-kron'' :: (Applicative g, Applicative k', Additive k', Functor f', Functor f, Functor g') =>
  (forall k. Additive k => Kron f k a -> Kron f' k a) -> (forall k. Additive k => Kron g k a -> Kron g' k a) -> (Kron (Kron f g) k') a -> (Kron (Kron f' g') k') a
-}

type Qubit = V2

kron'' :: (Additive f, Additive g, Additive k, Additive f', Additive g') =>
  LinOp1 f f' a -> LinOp1 g g' a -> Kron (Kron f g) k a -> Kron (Kron f' g') k a
kron'' lf lg fgk = let v = (Qubit.assoc fgk) in Qubit.assoc' (Compose $ fmap lg $ getCompose (lf v))

sigx' :: LinOp1 Qubit Qubit C
sigx' (Compose (V2 up down)) = Compose $ V2 down up

sigz' :: LinOp1 Qubit Qubit C
sigz' (Compose (V2 up down)) = Compose  $ V2 up ((-1) *^ down)

sigy' :: LinOp1 Qubit Qubit C
sigy' (Compose (V2 up down)) = Compose $ V2 ((-i) *^ down) (i *^ up) where i = 0 :+ 1

swap' :: (Traversable f, Applicative g) => LinOp1 (Kron f g) (Kron g f) a
swap' (Compose (Compose fgk)) = Compose $ Compose $ sequenceA fgk

cnot :: LinOp1 (Kron Qubit Qubit) (Kron Qubit Qubit) a
cnot (Compose (Compose (V2 (V2 up1 down1) v))) = Compose $ Compose $ V2 (V2 down1 up1) v

phase :: Double -> LinOp1 Qubit Qubit C
phase phi (Compose (V2 up down)) = Compose $ V2 up ((cis phi) *^ down)

mtest :: LinOp1 (Kron Qubit Qubit) (Kron Qubit Qubit) C
mtest = (left sigx') . (left sigy') . (right sigz') . swap'

left :: (Additive f, Additive k, Additive g) => LinOp1 f g a -> LinOp1 (Kron f k) (Kron g k) a
left l = kron'' l id -- Qubit.assoc' . l . Qubit.assoc

right :: (Additive k, Additive f, Additive g) => LinOp1 f g a -> LinOp1 (Kron k f) (Kron k g) a
right l = kron'' id l -- (Compose (Compose fkk)) = Compose $ Compose $ (fmap (getCompose . l . Compose) fkk)

leftUnit :: Kron V1 (Kron f k) a -> Kron f k a
leftUnit (Compose (V1 x)) = x

{-
(f a -> g a) is a natural transfromation
(Fractional a => f a -> g a) is a numerical transfromation

Horizontal composition in Functor land. Sure.
Even in applicative
The issue is that Good linear ops are going to hold numbers. So ti feels wrong to quanitfy over them
One option is to add a type class constraint
(Num a) or (Fractional a)
This will work because linear lifts these

-- this is the same as par written above.

-- Functor f => (t -> f a) -> (a -> b) -> t -> f b
The oppotite ordering
--  Functor f' => (f' b -> t') -> (a -> b) -> f' a -> t'
The more general type as inferred is
f b ~ t'

(f' a -> f a) -> (a -> b) -> (f' a -> f b)
vs
(f' b -> f b) -> (a -> b) -> (f' a -> f b)
vs.
forall

kron l l' x = l' (fmap l x)

(forall f. Additive f => Kron V2 f a -> Kron V3 f a) -> () -> (forall p. Kron f (Kron k p) -> Kron f' (Kron k' p) )
We want to allow the field a to unify across them all.
The scalar is generic in a certain sense, but not in another.

This is the lifted Kron'
We don't want Additive1, something went screwy there.
Additive (Kron' )

The problem with Num based Kron: If we define LinOp
(LinOp f g = forall Num a => f a -> g a
Why Num, Why not Fractional. This is a problem.
And really we want to be threading the scalar a throughout the entire system.

kron

And what if we unified the two orderings

Hmm. These are pretty similar to van Laarhoven lenses. This observation is porbably poison
Swtich the arguments
(a -> b) -> Lens
Functor f => (t -> f a) -> (t -> f b)
type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t
Lens t a t b

kron :: (Fractional a => f a -> g a) -> (Fractional k) -> LinOp
kron l l' (Compose fk) =  Compose (fmap l' (l fk))

Constrained natural transformation
kron  forall const. (const a => f a -> g a) -> (forall b. const (k b) => k b -> l b)

(f Double -> g Double) ->
(forall f. V2' f Double -> V2' f Double)

kron :: LinOp f g -> LinOp k l -> LinOp (Kron f k) (Kron g l)

-}

{-

-- can't build. Well we can build vzero
exv2' :: (Additive f, Num a) => V2' f a
exv2' = V2' (Compose ( ?  ))

-}
-- constants are given as morphisms
-- exv2' :: (Additive f, Num a) => V1' f a -> V2' f a
-- exv2' (V1' f) = (V2' (V2   )

{-
There is no additive instance for Compose?
Huh.

class (forall f. Additive f => Additive k f) => Additive1 k where

I'm doing something insane.
Yea, The 1 pattern. I've been here before and decided it was insane.
Something in purescript. Surely it wasn't Additive1 was it?

-}
{-
class Additive1 k where
    vadd :: (Additive f, Num a) => k f a -> k f a -> k f a
    vzero :: (Additive f, Num a) => k f a
instance Additive1 V2' where -- (Functor (V2' f)) =>

    vadd (V2' (V2 v w)) (V2' (V2 x y)) = V2' (V2 (v ^+^ x)  (w ^+^ y)   )
    vzero = V2' (V2 zero zero)

    --vadd (V2' v) (V2' w) = V2' (v ^+^ w)
    --vzero = V2' zero

instance Additive1 V1' where -- (Functor (V2' f)) =>
    vadd (V1' v) (V1' w) = V1' (v ^+^ w)
    vzero = V1' zero

instance (Additive1 k, Additive1 l) => Additive1 (Kron' k l) where
   vadd (Kron' x) (Kron' y) = Kron' (vadd x y)
   vzero = Kron' vzero
   -}
{-
Deriv functor for V2'
-}


module Statistics.GLM
    ( glm
    , logisticReg
    ) where

import Numeric.LinearAlgebra hiding ((<>))

data RegFamily = RegFamily { g :: Matrix Double -> Vector Double -> Vector Double
                           , gprime :: Matrix Double -> Vector Double -> Vector Double
                           , variance :: Matrix Double -> Vector Double -> Vector Double
                           }

logisticReg :: RegFamily
logisticReg = RegFamily g gprime variance
  where
    g a x = 1 / (exp(-(a #> x)) + 1)
    gprime a x = exp(-(a #> x)) / (exp(-(a #> x)) + 1)**2
    variance a x = g a x * (1 - g a x)

doIteration :: Matrix Double -> Vector Double -> Vector Double -> RegFamily -> Vector R
doIteration a x b (RegFamily g gprime variance) =
  let gp = gprime a x
      w = gp ** 2 / variance a x
      z = a #> x + (b - g a x) / gp
      eqleft = tr a <> (matrix 1 (toList w) * a)
      eqright = tr a #> (w * z)
  in eqleft <\> eqright

glm :: RegFamily -> Matrix Double -> Vector Double -> Vector Double -> [Vector Double]
glm rf a x b = x : glm rf a (doIteration a x b rf) b


{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DataKinds, PolyKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE UndecidableInstances #-}

module TensorHMatrix where

import Misc.Nats
import Numeric.LinearAlgebra as H
import Data.Vector.Storable as V
import Data.Default
import Data.Singletons.Prelude
import Generics.SOP.Constraint as C
import Misc.Constraints
import GenericTensor

import Prelude hiding (zipWith)

-- the kitchen sink of constraints
type Usable a = (Element a, Num a, Numeric a, Num (Vector a), Container Vector a, Floating a, Floating (Vector a))

data HTensor a (dims :: [k]) where
  Scalar :: !a -> HTensor a '[]
  Vector :: Vector a -> HTensor a '[n]
  Matrix :: Matrix a -> HTensor a '[n, m]

deriving instance (Show a, Element a) => Show (HTensor a dims)

instance (IntegralL dims, Default a, Usable a) => Default (HTensor a dims) where
  def = fill' def

--type instance IndexOf (HTensor l) = Rec LT l
--type instance IndexOf (HTensor l) = Rec (Const Int) l
--type instance IndexOf (HTensor l) = Vec (Length l) Int

{- Fails backwards fundep :(
instance Indexable (HTensor (n ': l) a) (HTensor l a) where
  Matrix m ! i = Vector (m ! i)
  Vector v ! i = Scalar (v ! i)
-}

instance Usable a => Tensor (HTensor a :: [k] -> *) where
  type N (HTensor a) = a

  scalar (Scalar a) = a

  transpose (Matrix m) = Matrix (tr' m)

  scale (Scalar a) (Vector v) = Vector (H.scale a v)

  dot (Vector a) (Vector b) = Scalar (a <.> b)

  outer (Vector a) (Vector b) = Matrix (H.outer a b)

  mv (Matrix m) (Vector v) = Vector $ m #> v

  mm (Matrix m1) (Matrix m2) = Matrix $ m1 <> m2

  select i (Vector v) = Scalar (v V.! i)

  oneHot :: forall (n :: k). (SingI n, IntegralN n) => Int -> HTensor a '[n]
  oneHot i = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(i, 1)]

  --fill :: forall (dims :: [k]). (IntegralL dims) => a -> HTensor a dims
  fill sDims a = case sDims of
    SNil -> Scalar a
    SCons n SNil -> Vector $ konst a (natVal' n)
    SCons m (SCons n SNil) -> Matrix $ konst a (natVal' m, natVal' n)
    _ -> error "HTensors only go up to dimension 2."

-- pretty much all of these instances should be automatically derivable
-- the only issue is that fromInteger needs to use natVal'

instance (Usable a, IntegralL dims) => Num (HTensor a dims) where
  Scalar a + Scalar b = Scalar (a + b)
  Vector a + Vector b = Vector (a + b)
  Matrix a + Matrix b = Matrix (a + b)

  Scalar a - Scalar b = Scalar (a - b)
  Vector a - Vector b = Vector (a - b)
  Matrix a - Matrix b = Matrix (a - b)

  Scalar a * Scalar b = Scalar (a * b)
  Vector a * Vector b = Vector (a * b)
  Matrix a * Matrix b = Matrix (a * b)

  negate (Scalar a) = Scalar (negate a)
  negate (Vector a) = Vector (negate a)
  negate (Matrix a) = Matrix (negate a)

  abs (Scalar a) = Scalar (abs a)
  abs (Vector a) = Vector (abs a)
  abs (Matrix a) = Matrix (abs a)

  signum (Vector a) = Vector (signum a)
  signum (Scalar a) = Scalar (signum a)
  signum (Matrix a) = Matrix (signum a)

  fromInteger = fill' . fromInteger

instance (IntegralL dims, Usable a) => Fractional (HTensor a dims) where
  Scalar a / Scalar b = Scalar (a / b)
  Vector a / Vector b = Vector (a / b)
  Matrix a / Matrix b = Matrix (a / b)

  recip (Scalar a) = Scalar (recip a)
  recip (Vector a) = Vector (recip a)
  recip (Matrix a) = Matrix (recip a)

  fromRational = fill' . fromRational

instance (IntegralL dims, Usable a) => Floating (HTensor a dims) where
  pi = fill' pi

  exp (Scalar a) = Scalar (exp a)
  exp (Vector a) = Vector (exp a)
  exp (Matrix a) = Matrix (exp a)

  log (Scalar a) = Scalar (log a)
  log (Vector a) = Vector (log a)
  log (Matrix a) = Matrix (log a)

  sqrt (Scalar a) = Scalar (sqrt a)
  sqrt (Vector a) = Vector (sqrt a)
  sqrt (Matrix a) = Matrix (sqrt a)

  Scalar a ** Scalar b = Scalar (a ** b)
  Vector a ** Vector b = Vector (a ** b)
  Matrix a ** Matrix b = Matrix (a ** b)

  logBase (Scalar a) (Scalar b) = Scalar (logBase a b)
  logBase (Vector a) (Vector b) = Vector (logBase a b)
  logBase (Matrix a) (Matrix b) = Matrix (logBase a b)

  sin (Scalar a) = Scalar (sin a)
  sin (Vector a) = Vector (sin a)
  sin (Matrix a) = Matrix (sin a)

  cos (Scalar a) = Scalar (cos a)
  cos (Vector a) = Vector (cos a)
  cos (Matrix a) = Matrix (cos a)

  tan (Scalar a) = Scalar (tan a)
  tan (Vector a) = Vector (tan a)
  tan (Matrix a) = Matrix (tan a)

  asin (Scalar a) = Scalar (asin a)
  asin (Vector a) = Vector (asin a)
  asin (Matrix a) = Matrix (asin a)

  acos (Scalar a) = Scalar (acos a)
  acos (Vector a) = Vector (acos a)
  acos (Matrix a) = Matrix (acos a)

  atan (Scalar a) = Scalar (atan a)
  atan (Vector a) = Vector (atan a)
  atan (Matrix a) = Matrix (atan a)

  sinh (Scalar a) = Scalar (sinh a)
  sinh (Vector a) = Vector (sinh a)
  sinh (Matrix a) = Matrix (sinh a)

  cosh (Scalar a) = Scalar (cosh a)
  cosh (Vector a) = Vector (cosh a)
  cosh (Matrix a) = Matrix (cosh a)

  tanh (Scalar a) = Scalar (tanh a)
  tanh (Vector a) = Vector (tanh a)
  tanh (Matrix a) = Matrix (tanh a)

  asinh (Scalar a) = Scalar (asinh a)
  asinh (Vector a) = Vector (asinh a)
  asinh (Matrix a) = Matrix (asinh a)

  acosh (Scalar a) = Scalar (acosh a)
  acosh (Vector a) = Vector (acosh a)
  acosh (Matrix a) = Matrix (acosh a)

  atanh (Scalar a) = Scalar (atanh a)
  atanh (Vector a) = Vector (atanh a)
  atanh (Matrix a) = Matrix (atanh a)

instance Usable a => ImpliesC (IntegralL dims) (CompC Floating (HTensor a) dims) where
  impliesC = Sub Dict

instance Usable a => ForallC (ImpliesC1 IntegralL (CompC Floating (HTensor a))) where
  forallC = Forall Dict

tmap :: (Container Vector a, Num a, Element b) => (a -> b) -> HTensor a dims -> HTensor b dims
tmap f (Scalar a) = Scalar $ f a
tmap f (Vector v) = Vector $ cmap f v
tmap f (Matrix m) = Matrix $ cmap f m

tZipWith :: (Storable a, Storable b, Storable c) => (a -> b -> c) -> HTensor a '[n] -> HTensor b '[n] -> HTensor c '[n]
tZipWith f (Vector a) (Vector b) = Vector (zipWith f a b)

{-
transpose :: (Numeric a) => HTensor '[n, m] a -> HTensor '[m, n] a
transpose (Matrix m) = Matrix (tr' m)

asCol :: Storable a => HTensor '[n] a -> HTensor '[n, FromInteger 1] a
asCol (Vector v) = Matrix $ asColumn v

asRow' :: Storable a => HTensor '[n] a -> HTensor '[FromInteger 1, n] a
asRow' (Vector v) = Matrix $ asRow v

dot :: Numeric a => HTensor '[n] a -> HTensor '[n] a -> a
dot (Vector a) (Vector b) = a <.> b

mv :: Numeric a => HTensor '[n, m] a -> HTensor '[m] a -> HTensor '[n] a
mv (Matrix m) (Vector v) = Vector $ m #> v

mm :: Numeric a => HTensor '[n, m] a -> HTensor '[m, k] a -> HTensor '[n, k] a
mm (Matrix m1) (Matrix m2) = Matrix $ m1 <> m2

gradDot :: Numeric a => HTensor '[n] a -> HTensor '[n] a -> a -> (HTensor '[n] a, HTensor '[n] a)
gradDot (Vector a) (Vector b) g = (Vector $ scale g b, Vector $ scale g a)

gradMV :: Numeric a => HTensor '[n, m] a -> HTensor '[m] a -> HTensor '[n] a -> (HTensor '[n, m] a, HTensor '[m] a)
gradMV m v g = (mm (asCol g) (asRow' v), mv (transpose m) g)

gradMM :: Numeric a => HTensor '[n, m] a -> HTensor '[m, k] a -> HTensor '[n, k] a -> (HTensor '[n, m] a, HTensor '[m, k] a)
gradMM a b g = (mm g (transpose b), mm (transpose a) g)

gradSelect :: forall a n. (SingI n, IntegralN n, Usable a) => Int -> HTensor '[n] a -> a -> HTensor '[n] a
gradSelect i _ a = Vector $ (konst 0 (natVal' (sing::Sing n))) // [(i, a)]
-}


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE BangPatterns #-}

module Numeric.MCMC.RiemannianLangevin (
            MarkovChain(..)
          , Options, createOptions
          , runChain, localMean, perturb
          ) where

import Numeric.MCMC.Langevin hiding (Options, runChain)
import Numeric.LinearAlgebra
import Control.Arrow
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.Reader
import System.Random.MWC
import System.Random.MWC.Distributions

-- | Options for the chain.  The target (expected to be a log density), the
--   right matrix product of its curvature and gradient, the inverse Fisher metric tensor,
--   its (Cholesky) square root, and the step size.
data Options = Options {
    _target              :: [Double] -> Double        -- Target (log density)
  , _curvatureXgradient  :: [Double] -> [Double]      -- Curvature right-multiplied by gradient
  , _invFisherMetric     :: [Double] -> Matrix Double -- Inverse Fisher metric tensor
  , _sqrtInvFisherMetric :: [Double] -> Matrix Double -- Square root of the tensor
  , _eps                 :: {-# UNPACK #-} !Double    -- Step size
  }

-- | A result with this type has a view of the chain options.
type ViewsOptions = ReaderT Options

-- | Construct Options (data constructor not exported).
-- FIXME sqrtm is not converging; alternatives?
createOptions :: ([Double] -> Double)     -- Target (log density)
              -> ([Double] -> [Double])   -- Gradient
              -> ([Double] -> [[Double]]) -- Hessian
              -> Double                   -- Step size
              -> Options                  -- Options
createOptions t g h =
    Options t curvatureXgradient invFisherMetric sqrtInvFisherMetric
  where curvatureXgradient xs =
            let mat = invFisherMetric xs <> fromColumns [fromList (g xs)]
            in  concat . toLists . trans $ mat
        invFisherMetric       = pinv . fromRows . map (fromList . map (* (-1))) . h
        sqrtInvFisherMetric x = let z = schur (invFisherMetric x)
                                in  fst z <> sqrtm (snd z) <> trans (fst z)
{-# INLINE createOptions #-}

-- | Non-isotropic Gaussian density.
nonIsoGauss :: [Double] -> [Double] -> Matrix Double -> Double
nonIsoGauss xs mu sig = exp val
  where val = -0.5*p*log (2*pi) - 0.5*ldet - 0.5*
            (trans (xsM `sub` muM) <> invSig <> (xsM `sub` muM)) @@> (0, 0)
        (xsM, muM) = (\f (a, b) -> (f a, f b)) (\l -> fromColumns [fromList l]) (xs, mu)
        p                   = fromIntegral $ cols sig
        (invSig, (ldet, _)) = invlndet sig
{-# INLINE nonIsoGauss #-}

-- | Mean function for the discretized Riemannian Langevin diffusion.
localMean :: Monad m
          => [Double]                -- Current state
          -> ViewsOptions m [Double] -- Localized mean of proposal distribution
localMean t = do
    Options _ c _ _ e <- ask
    return $! zipWith (+) t (map (* (0.5 * e^(2 :: Int))) (c t))
{-# INLINE localMean #-}

-- | Perturb the state, creating a new proposal.
perturb :: PrimMonad m
        => [Double]                   -- Current state
        -> Gen (PrimState m)          -- MWC PRNG
        -> ViewsOptions m [Double]    -- Resulting perturbation.
perturb t g = do
    Options _ _ _ s _ <- ask
    zs <- replicateM (length t) (lift $ standard g)
    t0 <- localMean t
    let adjustedBrownianMotion = s t <> fromColumns [fromList zs]
        abmList = concat . toLists . trans $ adjustedBrownianMotion
        perturbedState = zipWith (+) t0 t1
        t1 = map (* eps) abmList
    return $! perturbedState
{-# INLINE perturb #-}

-- | Perform a Metropolis accept/reject step.
metropolisStep :: PrimMonad m
               => MarkovChain                -- Current state
               -> Gen (PrimState m)          -- MWC PRNG
               -> ViewsOptions m MarkovChain -- New state
metropolisStep state g = do
    Options target _ iF _ e <- ask
    let (t0, nacc) = (theta &&& accepts) state
    zc           <- lift $ uniformR (0, 1) g
    proposal     <- perturb t0 g
    t0Mean       <- localMean t0
    proposalMean <- localMean proposal
    let mc = if   zc < acceptProb
             then (proposal, 1)
             else (t0,       0)

        acceptProb = if isNaN val then 1 else val where val = arRatio

        arRatio = exp . min 0 $
            target proposal + log (nonIsoGauss proposal t0Mean ((e^(2::Int)) `scale` iF t0))
          - target t0       - log (nonIsoGauss t0 proposalMean ((e^(2::Int)) `scale` iF proposal))

    return $! MarkovChain (fst mc) (nacc + snd mc)
{-# INLINE metropolisStep #-}

-- | Diffuse through states.
runChain :: Options         -- Options of the Markov chain.
         -> Int             -- Number of epochs to iterate the chain.
         -> Int             -- Print every nth iteration
         -> MarkovChain     -- Initial state of the Markov chain.
         -> Gen RealWorld   -- MWC PRNG
         -> IO MarkovChain  -- End state of the Markov chain, wrapped in IO.
runChain = go
  where go o n t !c g | n == 0 = return c
                      | n `rem` t /= 0 = do
                            r <- runReaderT (metropolisStep c g) o
                            go o (n - 1) t r g
                      | otherwise = do
                            r <- runReaderT (metropolisStep c g) o
                            print r
                            go o (n - 1) t r g
{-# INLINE runChain #-}


{-# LANGUAGE FlexibleContexts #-}

import Data.Function
import Data.Fixed
import Numeric.LinearAlgebra

grid :: Integer -> Matrix Double
grid n = build (300, 300) (builder `on` (+1) . round)
  where builder x y = let id = x + 10
                       in fromInteger . subtract 5
                        $ (id * y + n) * id `mod` 1000 `div` 100

main = getContents >>= putStrLn
                     . (\(x, y) -> show (x+1) ++ "," ++ show (y+1))
                     . maxIndex
                     . corr2 (konst 1 (3, 3))
                     . grid . read


module HLearn.Metrics.Mahalanobis
    where

import Control.DeepSeq
import qualified Data.Vector.Generic as VG

import qualified Data.Strict as Strict
import Foreign.Storable
import Numeric.LinearAlgebra hiding ((<>),dim)
import qualified Numeric.LinearAlgebra as LA

import HLearn.Algebra

-------------------------------------------------------------------------------
-- data types

data Mahalanobis dp = Mahalanobis
    { rawdp :: !dp
    , moddp :: !dp
    }
    deriving (Read,Show,Eq,Ord)

instance NFData dp => NFData (Mahalanobis dp) where
    rnf m = deepseq (rawdp m)
          $ rnf (moddp m)

mkIdentity :: dp -> Mahalanobis dp
mkIdentity dp = Mahalanobis dp dp

---------------------------------------

dim :: VG.Vector dp r => Mahalanobis (dp r) -> Int
dim m = VG.length $ rawdp m

---------------------------------------

class
    ( Storable (Ring dp)
    , Element (Ring dp)
    , Container Vector (Ring dp)
    , LA.Product (Ring dp)
    , Field (Ring dp)
    , Ring dp ~ Double
    ) => MatrixField dp

instance
    ( Storable (Ring dp)
    , Element (Ring dp)
    , Container Vector (Ring dp)
    , LA.Product (Ring dp)
    , Field (Ring dp)
    , Ring dp ~ Double
    ) => MatrixField dp

class MkMahalanobis params where
    type MetricDatapoint params -- = Datapoint params
    mkMahalanobis :: params -> MetricDatapoint params -> Mahalanobis (MetricDatapoint params)

-------------------------------------------------------------------------------
-- algebra

instance HasRing dp => HasRing (Mahalanobis dp) where
    type Ring (Mahalanobis dp) = Ring dp

instance
    ( RealFrac r
    , Floating r
    , Ring (dp r) ~ r
    , HasRing (dp r)
    , VG.Vector dp r
    ) => MetricSpace (Mahalanobis (dp r)) where
    {-# INLINABLE distance #-}

    distance !m1 !m2 = {-# SCC distance #-} sqrt $ go 0 (dim m1-1)
        where
            go tot (-1) = tot
            go tot i = go (tot+(((rawdp m1) `VG.unsafeIndex` i)-((rawdp m2) `VG.unsafeIndex` i))
                              *(((moddp m1) `VG.unsafeIndex` i)-((moddp m2) `VG.unsafeIndex` i))) (i-1)

--     isFartherThanWithDistance !m1 !m2 !dist = {-# SCC isFartherThanWithDistance #-}
--         go 0 (dim m1-1)
--         where
--             dist2=dist*dist
--
--             go tot (-1) = Strict.Just $ sqrt tot
--             go tot i = if tot'>dist2
--                 then Strict.Nothing
--                 else go tot' (i-1)
--                 where
--                     tot' = tot+(((rawdp m1) `VG.unsafeIndex` i)-((rawdp m2) `VG.unsafeIndex` i))
--                               *(((rawdp m1) `VG.unsafeIndex` i)-((rawdp m2) `VG.unsafeIndex` i))


{-# LANGUAGE AllowAmbiguousTypes, RankNTypes, UnicodeSyntax, NoMonomorphismRestriction #-}
{-# LANGUAGE BangPatterns, FlexibleInstances #-}

module Optimizer where

import Utils
import Style
import GenOptProblem
import Numeric.AD
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Reverse
import Numeric.AD.Internal.Sparse
import qualified Numeric.LinearAlgebra as L
import Debug.Trace
import System.Random
import System.Console.ANSI
import Data.List (foldl')

default (Int, Float)

------ Opt types, util functions, and params

type ObjFn1 a = forall a . (Autofloat a) => [a] -> a

-- used for duf
type ObjFn2 a = forall a . (Autofloat a) => [a] -> [a] -> a

type GradFn a = forall a . (Autofloat a) => [a] -> [a]

----- Various consts

nanSub :: (Autofloat a) => a
nanSub = 0

infinity :: Floating a => a
infinity = 1/0 -- x/0 == Infinity for any x > 0 (x = 0 -> Nan, x < 0 -> -Infinity)
-- all numbers are smaller than infinity except infinity, to which it's equal

----- Hyperparameters

weightGrowthFactor :: (Autofloat a) => a -- for EP weight
weightGrowthFactor = 10

epsUnconstr :: Floating a => a
epsUnconstr = 10 ** (-2)

epStop :: Floating a => a -- for EP diff
epStop = 10 ** (-5)
-- epStop = 60 ** (-3)
-- epStop = 10 ** (-1)
-- epStop = 0.05

-- Parameters for Armijo-Wolfe line search
-- NOTE: must maintain 0 < c1 < c2 < 1
c1 :: Floating a => a
c1 = 0.001 -- for Armijo, corresponds to alpha in backtracking line search (see below for explanation)
-- smaller c1 = shallower slope = less of a decrease in fn value needed = easier to satisfy
-- turn Armijo off: c1 = 0
-- Nocedal p38: "In practice, c1 is chosen to be quite small, say 10−4."

c2 :: Floating a => a
c2 = 0.9 -- for Wolfe, is the factor decrease needed in derivative value
-- new directional derivative value / old DD value <= c2
-- smaller c2 = smaller new derivative value = harder to satisfy
-- turn Wolfe off: c1 = 1 (basically backatracking line search only)
-- Nocedal p39: "Typical values of c2 are 0.9 when the search direction pk is chosen by a Newton or quasi-Newton method"

-- true = force linesearch halt if interval gets too small; false = no forced halt
intervalMin = True

useLineSearch :: Bool
useLineSearch = True

useAutodiff :: Bool
useAutodiff = True

constT :: Floating a => a
constT = 0.001

debugOpt = True
debugLineSearch = False
debugBfgs = True

trb :: String -> a -> a
trb s x = if debugBfgs then trace "---" $ trace s x else x -- prints in left to right order

tro :: String -> a -> a
tro s x = if debugOpt then trace "---" $ trace s x else x -- prints in left to right order

trl :: Show a => String -> a -> a
trl s x = if debugLineSearch then trace "---" $ trace s $ traceShowId x else x -- prints in left to right order

----- Convergence criteria

-- convergence criterion for EP
-- if you want to use it for UO, needs a different epsilon
epStopCond :: (Autofloat a) => [a] -> [a] -> a -> a -> Bool
epStopCond x x' fx fx' =
           tro ("EP: \n||x' - x||: " ++ (show $ norm (x -. x'))
           ++ "\n|f(x') - f(x)|: " ++ (show $ abs (fx - fx'))) $
           (norm (x -. x') <= epStop) || (abs (fx - fx') <= epStop)

unconstrainedStopCond :: (Autofloat a) => [a] -> Bool
unconstrainedStopCond gradEval = norm gradEval < epsUnconstr

---------------------------------------

-- Policies

stepPolicy :: State -> (Params, PolicyParams)
stepPolicy s =
    -- Check overall convergence first
    let epStatus = optStatus $ (paramsr s) in
    let pparams = policyParams s in
    case epStatus of

    -- Generate new objective function and replace the optimization and policy params accordingly
    EPConverged ->
                -- TODO: clean up the step incrementing
        let pparams' = pparams { policySteps = 1 + policySteps pparams } in
        let (policyRes, psNew) = (policyFn s) (objFns s) (constrFns s) pparams' in -- See what the policy function wants
        case policyRes of
            Nothing     -> (paramsr s, pparams' { policyState = psNew }) -- steps incremented, policy done

            Just newFns -> -- Policy keeps going
                let objFnNew = genObjfn (transr s) (filter isObjFn newFns) (filter isConstr newFns) (varyingPaths s)
                    -- TODO: check that these inputs are right
                    -- Change obj function and restart optimization
                    pparamsNew = pparams' { policyState = psNew,
                                            currFns = newFns }
                    paramsNew = Params { weight = initWeight,
                                         optStatus = NewIter,
                                         overallObjFn = objFnNew,
                                         bfgsInfo = defaultBfgsParams }
                in tro ("Step policy, EP converged, new params:\n" ++ show (paramsNew, pparamsNew, newFns)) $ (paramsNew, pparamsNew)

    -- If not converged, optimize as usual, don't change policy mid-optimization
    _ -> tro ("Step policy, EP not converged, new params:\n" ++ show (paramsr s, pparams)) $ (paramsr s, pparams)

---------------------------------------

-- Main optimization functions

step :: State -> State
step s = let (state', params') = stepShapes (oConfig s) (paramsr s) (varyingState s) (rng s)
             s'                = s { varyingState = state',
                                     paramsr = params' }
             -- NOTE: we intentionally discard the random generator here because
             -- we want to have consistent computation output in a single
             -- optimization session
             -- For the same reason, all subsequent step* functions such as
             -- stepShapes do not return the new random generator
             (!shapes', _, _)     = evalTranslation s'

             -- Check the state and see if the overall objective function should be changed
             -- The policy may change EPConverged to a new iteration before the frontend sees it
             (paramsNew, pparamsNew) = stepPolicy s'

             -- For debugging
             oldParams = paramsr s

         in tro ("Params: \n" ++ show oldParams ++ "\n:") $
            s' { shapesr = shapes',
                 paramsr = paramsNew,
                 policyParams = pparamsNew }
            -- note: trans is not updated in state

-- Note use of realToFrac to generalize type variables (on the weight and on the varying state)

-- implements exterior point algo as described on page 6 here:
-- https://www.me.utexas.edu/~jensen/ORMM/supplements/units/nlp_methods/const_opt.pdf
stepShapes :: OptConfig -> Params -> [Float] -> StdGen -> ([Float], Params)
stepShapes config params vstate g = -- varying state
         -- if null vstate then error "empty state in stepshapes" else
         let (epWeight, epStatus) = (weight params, optStatus params) in
         case epStatus of

         -- start the outer EP optimization and the inner unconstrained optimization, recording initial EPstate
         NewIter -> let status' = UnconstrainedRunning (map realToFrac vstate) in
                    (vstate', params { weight = initWeight, optStatus = status', bfgsInfo = defaultBfgsParams } )
         -- check *weak* convergence of inner unconstrained opt.
         -- if UO converged, set opt state to converged and update UO state (NOT EP state)
         -- if not, keep running UO (inner state implicitly stored)
         -- note convergence checks are only on the varying part of the state
         UnconstrainedRunning lastEPstate ->  -- doesn't use last EP state
           -- let unconstrConverged = unconstrainedStopCond gradEval in
           let unconstrConverged = epStopCond vstate vstate' (objFnApplied vstate) (objFnApplied vstate') in
               -- Two stopping conditions
               -- unconstrainedStopCond gradEval in
           if unconstrConverged then
              let status' = UnconstrainedConverged lastEPstate in -- update UO state only!
              (vstate', params { optStatus = status', bfgsInfo = defaultBfgsParams }) -- note vstate' (UO converged), not vstate
           else (vstate', params { bfgsInfo = bfgs' }) -- update UO state but not EP state; UO still running

         -- check EP convergence. if converged then stop, else increase weight, update states, and run UO again
         -- TODO some trickiness about whether unconstrained-converged has updated the correct state
         -- and whether to check WRT the updated state or not
         UnconstrainedConverged lastEPstate ->
           let epConverged = epStopCond lastEPstate (map r2f vstate) -- lastEPstate is last state for converged UO
                                   (objFnApplied lastEPstate) (objFnApplied (map r2f vstate)) in
           if epConverged then
              let status' = EPConverged in -- no more EP state
              (vstate, params { optStatus = status', bfgsInfo = defaultBfgsParams }) -- do not update UO state
           -- update EP state: to be the converged state from the most recent UO
           else let status' = UnconstrainedRunning (map realToFrac vstate) in -- increase weight
                let epWeight' = weightGrowthFactor * epWeight in
                -- trace ("Unconstrained converged. New weight: " ++ show epWeight') $
                      (vstate, params { weight = epWeight', optStatus = status', bfgsInfo = defaultBfgsParams })

         -- done; don't update obj state or params; user can now manipulate
         EPConverged -> (vstate, params { bfgsInfo = defaultBfgsParams } )

         -- TODO: implement EPConvergedOverride (for when the magnitude of the gradient is still large)

         -- TODO factor out--only unconstrainedRunning needs to run stepObjective, but EPconverged needs objfn
        where (vstate', gradEval, bfgs') = stepWithObjective config g params vstate
              objFnApplied = (overallObjFn params) g (r2f $ weight params)

-- Given the time, state, and evaluated gradient (or other search direction) at the point,
-- return the new state

stepT :: Float -> Float -> Float -> Float
stepT dt x dfdx = x - dt * dfdx

-- Calculates the new state by calculating the directional derivatives (via autodiff)
-- and timestep (via line search), then using them to step the current state.
-- Also partially applies the objective function.
stepWithObjective :: OptConfig -> StdGen -> Params -> [Float] -> ([Float], [Float], BfgsParams)
stepWithObjective config g params state =
          -- get timestep via line search, and evaluated gradient at the state
          let (t', gradEval, gradToUse, bfgs') = timeAndGrad config params objFnApplied state
              -- step each parameter of the state with the time and gradient
              state' = map (\(v, dfdv) -> stepT t' v dfdv) (zip state $ gradToUse)
              (fx, fx') = (objFnApplied state, objFnApplied state')
          in -- if fx' > fx then error ("Error: new energy is greater than old energy: " ++ show (fx', fx)) else
             tro ("\n----------------------------------------\n"
                   ++ "\nopt params: \n" ++ (show params)
                   ++ "\n||x' - x||: " ++ (show $ norm (state -. state'))
                   ++ "\n|f(x') - f(x)|: " ++
                  (show $ abs (fx' - fx))
                   ++ "\nf(x'): \n" ++ (show fx')
                   ++ "\ngradEval: \n" ++ (show gradEval)
                   ++ "\n||gradEval||: \n" ++ (show $ norm gradEval)
                   ++ "\ngradToUse: \n" ++ (show gradToUse)
                   ++ "\n||gradToUse||: \n" ++ (show $ norm gradToUse)
                   -- ++ "\nhessian: \n" ++ (show $ h)
                   ++ "\nbfgs': \n" ++ (show bfgs') -- TODO: use trb
                   ++ "\n timestep: \n" ++ (show t')
                   ++ "\n original state: \n" ++ (show state)
                   ++ "\n new state: \n" ++ (show state')
                  )
             (state', gradEval, bfgs')

          where objFnApplied :: ObjFn1 b
                objFnApplied = (overallObjFn params) g cWeight

                cWeight = r2f $ weight params
                -- realToFrac generalizes the type variable `a` to the type variable `b`, which timeAndGrad expects

-- a version of grad with a clearer type signature
appGrad :: (Autofloat a) => (forall b . (Autofloat b) => [b] -> b) -> [a] -> [a]
appGrad f l = grad f l

instance Show (Numeric.AD.Internal.On.On
                             (Numeric.AD.Internal.Reverse.Reverse
                                s (Numeric.AD.Internal.Sparse.Sparse a))) where
         show a = "error: not sure how to derive show for hessian element"

appHess :: (Autofloat a) => (forall b . (Autofloat b) => [b] -> b) -> [a] -> [[a]]
appHess f l = hessian f l

-- Precondition the gradient
gradP :: OptConfig -> BfgsParams -> [Double] -> ObjFn1 a -> [Float] -> ([Double], BfgsParams)
gradP config bfgsParams gradEval f state =
      let x_k = L.vector $ map r2f state
          grad_fx_k = L.vector gradEval
      in case optMethod config of
      GradientDescent -> (gradEval, bfgsParams)

      Newton -> -- Precondition gradient with the pseudoinverse of the hessian
          let h = appHess f state
              h_list = (map r2f $ concat h) :: [Double]
              hinv = L.pinv $ L.matrix (length gradEval) $ h_list
              gradPreconditioned = hinv L.#> (L.vector gradEval) in
          (L.toList gradPreconditioned, bfgsParams)

      BFGS -> -- Approximate inverse of hessian with the change in gradient (see Nocedal S9.1 p224)
           let grad_val = lastGrad bfgsParams
               h_val = invH bfgsParams
               state_val = lastState bfgsParams
           in case (h_val, grad_val, state_val) of
              (Nothing, Nothing, Nothing) -> -- First step. Initialize the approximation to the identity (not clear how else to approximate it)
                    -- k=0 steps from x_0 to x_1: so on the first step, we take a normal gradient descent step
                    let h_0 = L.ident $ length gradEval
                    in (gradEval, bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k, invH = Just h_0 })

              (Just h_km1, Just grad_fx_km1, Just x_km1) -> -- For x_{k+1}, to compute H_k, we need the (k-1) info
                    -- Our convention is that we are always working "at" k to compute k+1
                    -- x_0 doesn't require any H; x_1 (the first step) with k = 0 requires H_0
                    -- x_2 (the NEXT step) with k=1 requires H_1. For example>
                    -- x_2 = x_1 - alpha_1 H_1 grad f(x_1)   [GD step]
                    -- H_1 = V_0 H_0 V_0 + rho_0 s_0 s_0^T   [This is confusing because the book adds an extra +1 to the H index]
                    -- V_0 = I - rho_0 y_0 s_0^T
                    -- rho_0 = 1 / y_0^T s_0
                    -- s_0 = x_1 - x_0
                    -- y_0 = grad f(x_1) - grad f(x_0)

                    let s_km1 = x_k - x_km1
                        y_km1 = grad_fx_k - grad_fx_km1
                        rho_km1 = 1 / (y_km1 `L.dot` s_km1) -- Scalar
                        v_km1 = L.ident (length gradEval) - (rho_km1 `L.scale` y_km1 `L.outer` s_km1) -- Scaling can happen before outer
                        h_k = (L.tr (v_km1) L.<> h_km1 L.<> v_km1) + (rho_km1 `L.scale` s_km1 `L.outer` s_km1)
                        gradPreconditioned = h_k L.#> grad_fx_k
                    in (L.toList gradPreconditioned,
                        bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k, invH = Just h_k })

              _ -> error "invalid BFGS state"

      LBFGS -> -- Approximate the inverse of the Hessian times the gradient
               -- Only using the last `m` gradient/state difference vectors, not building the full h_k matrix (Nocedal p226)
            let grad_prev = lastGrad bfgsParams
                x_prev = lastState bfgsParams
                ss_val = s_list bfgsParams
                ys_val = y_list bfgsParams
                km1 = numUnconstrSteps bfgsParams -- Our current step is k; the last step is km1 (k_minus_1)
                m = memSize bfgsParams
            in case (grad_prev, x_prev, ss_val, ys_val, km1) of

               -- Perform normal gradient descent on first step
               (Nothing, Nothing, [], [], 0) ->
                   -- Store x_k, grad f(x_k) so we can compute s_k, y_k on next step
                   let bfgsParams' = bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k,
                                                  s_list = [], y_list = [], numUnconstrSteps = km1 + 1 } in
                   (gradEval, bfgsParams')

               (Just grad_fx_km1, Just x_km1, ss, ys, km1) ->
                   -- Compute s_{k-1} = x_k - x_{k-1} and y_{k-1} = (analogous with grads)
                   -- Unlike Nocedal, compute the difference vectors first instead of last (same result, just a loop rewrite)
                   -- Use the updated {s_i} and {y_i}. (If k < m, this reduces to normal BFGS, i.e. we use all the vectors so far)
                   let (s_km1, y_km1) = (x_k - x_km1, grad_fx_k - grad_fx_km1) -- Newest vectors added to front
                       (ss', ys') = (take m $ s_km1 : ss, take m $ y_km1 : ys) -- The limited-memory part: drop stale vectors
                       gradPreconditioned = lbfgs grad_fx_k ss' ys'
                       descentDirCheck = -gradPreconditioned `L.dot` grad_fx_k

                   -- Reset L-BFGS if the result is not a descent direction, and use steepest descent direction
                   -- https://github.com/JuliaNLSolvers/Optim.jl/issues/143 https://github.com/JuliaNLSolvers/Optim.jl/pull/144

                   in if descentDirCheck < epsd
                      then (L.toList gradPreconditioned, bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k,
                                                                      s_list = ss', y_list = ys', numUnconstrSteps = km1 + 1 })
                      else tro "L-BFGS did not find a descent direction. Resetting correction vectors." $
                           (gradEval, bfgsParams { lastState = Just x_k, lastGrad = Just grad_fx_k,
                                                   s_list = [], y_list = [], numUnconstrSteps = km1 + 1 })

                   -- TODO: check the curvature condition y_k^T s_k > 0 (8.7) (Nocedal 201)
                   -- https://github.com/JuliaNLSolvers/Optim.jl/issues/26

               _ -> error "invalid L-BFGS state"

type LVector = L.Vector L.R
type LMatrix = L.Matrix L.R

-- See Nocedal p225
-- expects ys and ss to be in order from most recent to oldest (k-1 ... k-m)
-- expects length ys == length ss, length ys > 0, length ys > 0
lbfgs :: LVector -> [LVector] -> [LVector] -> LVector
lbfgs grad_fx_k ss ys =
    let rhos = map calculate_rho $ zip ss ys -- The length of any list should be the number of stored vectors
        q_k = grad_fx_k
        (q_k_minus_m, alphas) = foldl' pull_q_back (q_k, []) (zip3 rhos ss ys) -- backward: for i = k-1 ... k-m
        -- Note the order of alphas will be from k-m through k-1 for the push_r_forward loop
        h_0_k = estimate_hess (head ys) (head ss)
        r_k_minus_m = h_0_k L.#> q_k_minus_m
        r_k = foldl' push_r_forward r_k_minus_m (zip (zip (reverse rhos) alphas) (zip (reverse ss) (reverse ys)))
        -- forward: for i = k-m .. k-1 (TODO: optimize out the reverses)
    in r_k -- is H_k * grad f(x_k)

           where calculate_rho :: (LVector, LVector) -> L.R
                 calculate_rho (s, y) = 1 / ((y `L.dot` s) + epsd)

                 pull_q_back :: (LVector, [L.R]) -> (L.R, LVector, LVector) -> (LVector, [L.R]) -- from i+1 to i
                 pull_q_back (q_i_plus_1, alphas) (rho_i, s_i, y_i) =
                        let alpha_i = rho_i * (s_i `L.dot` q_i_plus_1) -- scalar
                            q_i = q_i_plus_1 - (alpha_i `L.scale` y_i) -- scalar * vector
                        in (q_i, alpha_i : alphas) -- Note the order of alphas

                 -- Scale I by an estimate of the size of the Hessian along the most recent search direction (Nocedal p226)
                 estimate_hess :: LVector -> LVector -> LMatrix
                 estimate_hess y_km1 s_km1 =
                        let gamma_k = (s_km1 `L.dot` y_km1) / ((y_km1 `L.dot` y_km1) + epsd)
                        in gamma_k `L.scale` (L.ident (L.size y_km1))

                 push_r_forward :: LVector -> ((L.R, L.R), (LVector, LVector)) -> LVector -- from i to i+1
                 push_r_forward r_i ((rho_i, alpha_i), (s_i, y_i)) =
                        let beta_i = rho_i * (y_i `L.dot` r_i) -- scalar
                            r_i_plus_1 = r_i + (alpha_i - beta_i) `L.scale` s_i -- scalar * vector
                        in r_i_plus_1

estimateGradient :: ObjFn1 a -> [Float] -> [Float]
estimateGradient f state =
      let len = length state
          h = 0.001 -- Choice of h really matters!!! This is not an accurate estimate.
          fx = f state
      -- time is O(|state|^2)
          dfx i = ((f $ replace i ((state !! i) + h) state) - fx) / h
          dfxs = map dfx [0..(len-1)]
      in dfxs
      where replace pos newVal list = take pos list ++ newVal : drop (pos+1) list

-- Given the objective function, gradient function, timestep, and current state,
-- return the timestep (found via line search) and evaluated gradient at the current state.
-- the autodiff library requires that objective functions be polymorphic with Floating a
timeAndGrad :: OptConfig -> Params -> ObjFn1 a -> [Float] -> (Float, [Float], [Float], BfgsParams)
timeAndGrad config params f state =
            let gradEval = if useAutodiff then gradF state else estimateGradient f state
                (gradToUse_d, bfgs') = gradP config (bfgsInfo params) (map r2f gradEval :: [Double]) f state
                gradToUse = map r2f gradToUse_d
                -- Use line search to find a good timestep. If we use Newton's method, the descent direction uses the preconditioned gradient.
                descentDir = negL $ gradToUse
                timestep = let resT = if useLineSearch && useAutodiff
                                      then awLineSearch config f (duf descentDir) descentDir state
                                      else constT in -- hardcoded timestep
                           if isNaN resT then error "returned timestep is NaN" else resT

            in tr "timeAndGrad: " (timestep, gradEval, gradToUse, bfgs')

            where gradF :: GradFn a
                  gradF = appGrad f

                  -- directional derivative of f at x in the direction of u (descent direction, which may not have unit norm)
                  duf :: (Autofloat a) => [a] -> [a] -> a
                  duf u x = u `dotL` gradF x

linesearch_max :: Int
linesearch_max = 100 -- TODO what's a reasonable limit (if any)?

-- Implements Armijo-Wolfe line search as specified in Keenan's notes, converges on nonconvex fns as well
-- based off Lewis & Overton, "Nonsmooth optimization via quasi-Newton methods
-- duf = D_u(f), the directional derivative of f at descent direction u
-- D_u(x) = <gradF(x), u>. If u = -gradF(x) (as it is here), then D_u(x) = -||gradF(x)||^2

awLineSearch :: OptConfig -> ObjFn1 a -> ObjFn1 a -> [Float] -> [Float] -> Float
awLineSearch config f duf descentDir x0 =

     let (a0, b0, t0) = (0, infinity, 1) -- Unit step length should be tried first for quasi-Newton methods. (Nocedal 201)

         -- drop while a&w are not satisfied OR the interval is large enough
         (numUpdates, (af, bf, tf)) = head $ dropWhile intervalOK_or_notArmijoAndWolfe
                                      $ zip [0..] $ iterate update (a0, b0, t0)

     in tro ("Line search # updates: " ++ show numUpdates) $ tf

          where update :: (Float, Float, Float) -> (Float, Float, Float)
                update (a, b, t) =
                       let (a', b', sat) | not $ armijo t    = trl "not armijo" (a, t, False)
                                         | not $ wolfe t     = trl "not wolfe" (t, b, False)
                                         | otherwise         = (a, b, True) in
                       if sat then (a, b, t) -- if armijo and wolfe, then we use (a, b, t) as-is
                       else if b' < infinity then trl "b' < infinity" (a', b', (a' + b') / 2)
                       else trl "b' = infinity" (a', b', 2 * a')

                intervalOK_or_notArmijoAndWolfe :: (Int, (Float, Float, Float)) -> Bool
                intervalOK_or_notArmijoAndWolfe (numUpdates, (a, b, t)) =
                      not $
                      if armijo t && wolfe t then
                           trl ("stop: both sat. |descentDir at x0| = " ++ show (norm descentDir)) True
                      else if abs (b - a) < minInterval then
                           trl ("stop: interval too small. |descentDir at x0| = " ++ show (norm descentDir)) True
                      else if numUpdates > linesearch_max then
                           trl ("stop: number of line search updates exceeded max") True
                      else False

                armijo :: Float -> Bool
                armijo t = (f (x0 +. t *. descentDir)) <= (fAtx0 + c1 * t * dufAtx0)

                weakWolfe :: Float -> Bool -- Better for nonsmooth functions
                weakWolfe t = (duf (x0 +. t *. descentDir)) >= (c2 * dufAtx0)

                strongWolfe :: Float -> Bool
                strongWolfe t = (abs (duf (x0 +. t *. descentDir))) <= (c2 * abs dufAtx0)

                -- Descent direction (at x0) must have a negative dot product with the gradient of f at x0.
                dufAtx0 :: Float -- TODO: this is redundant, cache it
                dufAtx0 = let res = duf x0 in
                        trl ("<grad f(x0), descent direction at x0>: " ++ show res) $
                        if res > 0 then tro "WARNING: descent direction doesn't satisfy condition" res else res

                fAtx0 :: Float
                fAtx0 = f x0

                minInterval :: Float -- stop if the interval gets too small; might not terminate
                minInterval = if intervalMin then 10 ** (-10) else 0
                -- TODO: the line search is very sensitive to this parameter. Blows up with 10**(-5). Why?

                wolfe :: Float -> Bool
                wolfe = weakWolfe


module HNum.Special where

import           System.Random
import qualified Numeric.SpecFunctions         as SF
import           HNum.Vector
import           HNum.F

class FuncTools f => SpecialFunc f where
  erf :: Real a => f a -> f Double
  invErf :: Real a => f a -> f Double

instance SpecialFunc Vector where
  erf v = SF.erf . realToFrac <$> v
  invErf v = SF.invErf . realToFrac <$> v