text:
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DerivingVia               #-}
{-# OPTIONS_GHC -fno-warn-deprecations #-}

-- |
-- Module      : Data.Functor.Invariant.Inplicative
-- Copyright   : (c) Justin Le 2021
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Contains the classes 'Inply' and 'Inplicative', the invariant
-- counterparts to 'Apply'/'Divise' and 'Applicative'/'Divisible'.
--
-- @since 0.4.0.0
module Data.Functor.Invariant.Inplicative (
  -- * Typeclass
    Inply(..)
  , Inplicative(..)
  -- * Deriving
  , WrappedApplicativeOnly(..)
  , WrappedDivisibleOnly(..)
  -- * Invariant 'Day'
  , runDay
  , dather
  , runDayApply
  , runDayDivise
  -- * Assembling Helpers
  , gatheredN
  , gatheredNMap
  , gatheredN1
  , gatheredN1Map
  , gatheredNRec
  , gatheredNMapRec
  , gatheredN1Rec
  , gatheredN1MapRec
  , gatherN
  , gatherN1
  ) where

import           Control.Applicative
import           Control.Applicative.Backwards               (Backwards(..))
import           Control.Applicative.Lift                    (Lift(Pure, Other))
import           Control.Arrow                               (Arrow)
import           Control.Monad.Trans.Cont                    (ContT)
import           Control.Monad.Trans.Error                   (ErrorT(..))
import           Control.Monad.Trans.Except                  (ExceptT(..))
import           Control.Monad.Trans.Identity                (IdentityT(..))
import           Control.Monad.Trans.List                    (ListT(..))
import           Control.Monad.Trans.Maybe                   (MaybeT(..))
import           Control.Monad.Trans.RWS                     (RWST(..))
import           Control.Monad.Trans.Reader                  (ReaderT(..))
import           Control.Monad.Trans.State                   (StateT)
import           Control.Monad.Trans.Writer                  (WriterT(..))
import           Control.Natural
import           Data.Complex                                (Complex)
import           Data.Deriving
import           Data.Functor.Apply
import           Data.Functor.Bind.Class                     (Bind)
import           Data.Functor.Constant                       (Constant)
import           Data.Functor.Contravariant
import           Data.Functor.Contravariant.Divise
import           Data.Functor.Contravariant.Divisible
import           Data.Functor.Identity
import           Data.Functor.Invariant
import           Data.Functor.Invariant.Day
import           Data.Functor.Product                        (Product(..))
import           Data.Functor.Reverse                        (Reverse(..))
import           Data.Hashable                               (Hashable)
import           Data.Kind
import           Data.List.NonEmpty                          (NonEmpty)
import           Data.SOP hiding                             (hmap)
import           Data.Sequence                               (Seq)
import           Data.StateVar                               (SettableStateVar)
import           Data.Tagged                                 (Tagged)
import           Data.Tree                                   (Tree)
import           GHC.Generics                                (Generic)
import qualified Control.Monad.Trans.RWS.Strict as Strict    (RWST(..))
import qualified Control.Monad.Trans.State.Strict as Strict  (StateT)
import qualified Control.Monad.Trans.Writer.Strict as Strict (WriterT(..))
import qualified Data.HashMap.Lazy                           as HM
import qualified Data.IntMap                                 as IM
import qualified Data.Map                                    as M
import qualified Data.Monoid                                 as Monoid
import qualified Data.Semigroup                              as Semigroup
import qualified Data.Sequence.NonEmpty                      as NESeq
import qualified Data.Vinyl                                  as V
import qualified Data.Vinyl.Curry                            as V
import qualified Data.Vinyl.Functor                          as V
import qualified GHC.Generics                                as Generics

-- | The invariant counterpart of 'Apply' and 'Divise'.
--
-- Conceptually you can think of 'Apply' as, given a way to "combine" @a@ and
-- @b@ to @c@, lets you merge @f a@ (producer of @a@) and @f b@ (producer
-- of @b@) into a @f c@ (producer of @c@).  'Divise' can be thought of as,
-- given a way to "split" a @c@ into an @a@ and a @b@, lets you merge @f
-- a@ (consumer of @a@) and @f b@ (consumder of @b@) into a @f c@ (consumer
-- of @c@).
--
-- 'Inply', for 'gather', requires both a combining function and
-- a splitting function in order to merge @f b@ (producer and consumer of
-- @b@) and @f c@ (producer and consumer of @c@) into a @f a@.  You can
-- think of it as, for the @f a@, it "splits" the a into @b@ and @c@ with
-- the @a -> (b, c)@, feeds it to the original @f b@ and @f c@, and then
-- re-combines the output back into a @a@ with the @b -> c -> a@.
--
-- @since 0.4.0.0
class Invariant f => Inply f where
    -- | Like '<.>', '<*>', 'divise', or 'divide', but requires both
    -- a splitting and a recombining function.  '<.>' and '<*>' require
    -- only a combining function, and 'divise' and 'divide' require only
    -- a splitting function.
    --
    -- It is used to merge @f b@ (producer and consumer of @b@) and @f c@
    -- (producer and consumer of @c@) into a @f a@.  You can think of it
    -- as, for the @f a@, it "splits" the a into @b@ and @c@ with the @a ->
    -- (b, c)@, feeds it to the original @f b@ and @f c@, and then
    -- re-combines the output back into a @a@ with the @b -> c -> a@.
    --
    -- An important property is that it will always use @both@ of the
    -- ccomponents given in order to fulfil its job.  If you gather an @f
    -- a@ and an @f b@ into an @f c@, in order to consume/produdce the @c@,
    -- it will always use both the @f a@ or the @f b@ -- exactly one of
    -- them.
    --
    -- @since 0.4.0.0
    gather
        :: (b -> c -> a)
        -> (a -> (b, c))
        -> f b
        -> f c
        -> f a
    gather f g x y = invmap (uncurry f) g (gathered x y)
    -- | A simplified version of 'gather' that combines into a tuple.  You
    -- can then use 'invmap' to reshape it into the proper shape.
    --
    -- @since 0.4.0.0
    gathered
        :: f a
        -> f b
        -> f (a, b)
    gathered = gather (,) id
    {-# MINIMAL gather | gathered #-}

-- | The invariant counterpart of 'Applicative' and 'Divisible'.
--
-- The main important action is described in 'Inply', but this adds 'knot',
-- which is the counterpart to 'pure' and 'conquer'.  It's the identity to
-- 'gather'; if combine two @f a@s with 'gather', and one of them is
-- 'knot', it will leave the structure unchanged.
--
-- Conceptually, if you think of 'gather' as "splitting and re-combining"
-- along multiple forks, then 'knot' introduces a fork that is never taken.
--
-- @since 0.4.0.0
class Inply f => Inplicative f where
    knot :: a -> f a

-- | Interpret out of a contravariant 'Day' into any instance of 'Inply' by
-- providing two interpreting functions.
--
-- This should go in "Data.Functor.Invariant.Day", but that module is in
-- a different package.
--
-- @since 0.4.0.0
runDay
    :: Inply h
    => (f ~> h)
    -> (g ~> h)
    -> Day f g ~> h
runDay f g (Day x y a b) = gather a b (f x) (g y)

-- | Squash the two items in a 'Day' using their natural 'Inply'
-- instances.
--
-- This should go in "Data.Functor.Invariant.Day", but that module is in
-- a different package.
--
-- @since 0.4.0.0
dather
    :: Inply f
    => Day f f ~> f
dather (Day x y a b) = gather a b x y

-- | Ignores the contravariant part of 'gather'
instance Apply f => Inply (WrappedFunctor f) where
    gather f _ (WrapFunctor x) (WrapFunctor y) = WrapFunctor (liftF2 f x y)
    gathered (WrapFunctor x) (WrapFunctor y) = WrapFunctor (liftF2 (,) x y)
-- | @'knot' = 'pure'@
instance (Applicative f, Apply f) => Inplicative (WrappedFunctor f) where
    knot = pure

-- | Ignores the covariant part of 'gather'
instance Divise f => Inply (WrappedContravariant f) where
    gather _ g (WrapContravariant x) (WrapContravariant y) = WrapContravariant (divise g x y)
    gathered (WrapContravariant x) (WrapContravariant y) = WrapContravariant (divised x y)
-- | @'knot' _ = 'conquer'@
instance (Divisible f, Divise f) => Inplicative (WrappedContravariant f) where
    knot _ = conquer

-- | Ignores the covariant part of 'gather'
instance Divise f => Inply (WrappedDivisible f) where
    gather _ g (WrapDivisible x) (WrapDivisible y) = WrapDivisible (divise g x y)
    gathered (WrapDivisible x) (WrapDivisible y) = WrapDivisible (divised x y)
-- | @'knot' _ = 'conquer'@
instance (Divisible f, Divise f) => Inplicative (WrappedDivisible f) where
    knot _ = conquer

-- | Wrap an 'Applicative' that is not necessarily an 'Apply'.
newtype WrappedApplicativeOnly f a =
    WrapApplicativeOnly { unwrapApplicativeOnly :: f a }
  deriving (Generic, Eq, Show, Ord, Read, Functor, Foldable, Traversable)
  deriving newtype (Applicative, Monad)

deriveShow1 ''WrappedApplicativeOnly
deriveRead1 ''WrappedApplicativeOnly
deriveEq1 ''WrappedApplicativeOnly
deriveOrd1 ''WrappedApplicativeOnly

instance Invariant f => Invariant (WrappedApplicativeOnly f) where
    invmap f g (WrapApplicativeOnly x) = WrapApplicativeOnly (invmap f g x)
instance (Applicative f, Invariant f) => Apply (WrappedApplicativeOnly f) where
    x <.> y = x <*> y
-- | Ignores the contravariant part of 'gather'
instance (Applicative f, Invariant f) => Inply (WrappedApplicativeOnly f) where
    gather f _ (WrapApplicativeOnly x) (WrapApplicativeOnly y) = WrapApplicativeOnly (liftA2 f x y)
    gathered (WrapApplicativeOnly x) (WrapApplicativeOnly y) = WrapApplicativeOnly (liftA2 (,) x y)
-- | @'knot' = 'pure'@
instance (Applicative f, Invariant f) => Inplicative (WrappedApplicativeOnly f) where
    knot = pure

-- | Wrap an 'Divisible' that is not necessarily a 'Divise'.
newtype WrappedDivisibleOnly f a =
    WrapDivisibleOnly { unwrapDivisibleOnly :: f a }
  deriving (Generic, Eq, Show, Ord, Read, Functor, Foldable, Traversable)
  deriving newtype (Divisible, Contravariant)

deriveShow1 ''WrappedDivisibleOnly
deriveRead1 ''WrappedDivisibleOnly
deriveEq1 ''WrappedDivisibleOnly
deriveOrd1 ''WrappedDivisibleOnly

instance Invariant f => Invariant (WrappedDivisibleOnly f) where
    invmap f g (WrapDivisibleOnly x) = WrapDivisibleOnly (invmap f g x)
instance (Divisible f, Invariant f) => Divise (WrappedDivisibleOnly f) where
    divise g (WrapDivisibleOnly x) (WrapDivisibleOnly y) = WrapDivisibleOnly (divide g x y)
-- | Ignores the covariant part of 'gather'
instance (Divisible f, Invariant f) => Inply (WrappedDivisibleOnly f) where
    gather _ g (WrapDivisibleOnly x) (WrapDivisibleOnly y) = WrapDivisibleOnly (divide g x y)
    gathered (WrapDivisibleOnly x) (WrapDivisibleOnly y) = WrapDivisibleOnly (divided x y)
-- | @'knot' _ = 'conquer'@
instance (Divisible f, Invariant f) => Inplicative (WrappedDivisibleOnly f) where
    knot _ = conquer

funzip :: Functor f => f (a, b) -> (f a, f b)
funzip x = (fmap fst x, fmap snd x)

-- | @since 0.4.1.0
instance Inply f => Inply (MaybeT f) where
    gather f g (MaybeT x) (MaybeT y) = MaybeT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (MaybeT f) where
    knot x = MaybeT (knot (Just x))

-- | @since 0.4.1.0
instance (Inply f, Semigroup w) => Inply (WriterT w f) where
    gather f g (WriterT x) (WriterT y) = WriterT $
      gather (\case (a, q) -> \case (b, r) -> (f a b, q <> r))
             (\case (a, s) -> case g a of (b, c) -> ((b, s), (c, s)))
             x y
-- | @since 0.4.1.0
instance (Inplicative f, Monoid w) => Inplicative (WriterT w f) where
    knot x = WriterT (knot (x, mempty))

-- | @since 0.4.1.0
instance (Inply f, Semigroup w) => Inply (Strict.WriterT w f) where
    gather f g (Strict.WriterT x) (Strict.WriterT y) = Strict.WriterT $
      gather (\(~(a, q)) (~(b, r)) -> (f a b, q <> r))
             (\(~(a, s)) -> let ~(b, c) = g a in ((b, s), (c, s)))
             x y
-- | @since 0.4.1.0
instance (Inplicative f, Monoid w) => Inplicative (Strict.WriterT w f) where
    knot x = Strict.WriterT (knot (x, mempty))

-- | @since 0.4.1.0
instance Inply f => Inply (ReaderT r f) where
    gather f g (ReaderT x) (ReaderT y) = ReaderT $ \r ->
      gather f g (x r) (y r)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ReaderT r f) where
    knot x = ReaderT (\_ -> knot x)

-- | @since 0.4.1.0
instance Inply f => Inply (ExceptT e f) where
    gather f g (ExceptT x) (ExceptT y) = ExceptT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ExceptT e f) where
    knot x = ExceptT (knot (Right x))

-- | @since 0.4.1.0
instance Inply f => Inply (ErrorT e f) where
    gather f g (ErrorT x) (ErrorT y) = ErrorT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ErrorT e f) where
    knot x = ErrorT (knot (Right x))

-- | @since 0.4.1.0
instance Inply f => Inply (ListT f) where
    gather f g (ListT x) (ListT y) = ListT $
      gather (liftA2 f) (funzip . fmap g) x y
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (ListT f) where
    knot x = ListT (knot [x])

-- | @since 0.4.1.0
deriving via WrappedFunctor (RWST r w s m) instance (Bind m, Invariant m, Semigroup w) => Inply (RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (RWST r w s m) instance (Monad m, Bind m, Invariant m, Monoid w) => Inplicative (RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.RWST r w s m) instance (Bind m, Invariant m, Semigroup w) => Inply (Strict.RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.RWST r w s m) instance (Monad m, Bind m, Invariant m, Monoid w) => Inplicative (Strict.RWST r w s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (StateT s m) instance (Bind m, Invariant m) => Inply (StateT s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (StateT s m) instance (Monad m, Bind m, Invariant m) => Inplicative (StateT s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.StateT s m) instance (Bind m, Invariant m) => Inply (Strict.StateT s m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Strict.StateT s m) instance (Monad m, Bind m, Invariant m) => Inplicative (Strict.StateT s m)

-- | @since 0.4.1.0
instance Inply f => Inply (Generics.M1 i t f :: Type -> Type) where
    gather f g (Generics.M1 x) (Generics.M1 y) = Generics.M1 (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Generics.M1 i t f :: Type -> Type) where
    knot = Generics.M1 . knot
-- | @since 0.4.1.0
instance (Inply f, Inply g) => Inply (f Generics.:*: g) where
    gather f g (x1 Generics.:*: y1) (x2 Generics.:*: y2) =
        gather f g x1 x2 Generics.:*: gather f g y1 y2
-- | @since 0.4.1.0
instance (Inplicative f, Inplicative g) => Inplicative (f Generics.:*: g) where
    knot x = knot x Generics.:*: knot x
-- | @since 0.4.1.0
instance (Inply f, Inply g) => Inply (Product f g) where
    gather f g (Pair x1 y1) (Pair x2 y2) =
      gather f g x1 x2 `Pair` gather f g y1 y2
-- | @since 0.4.1.0
instance (Inplicative f, Inplicative g) => Inplicative (Product f g) where
    knot x = knot x `Pair` knot x
-- | @since 0.4.1.0
instance Inply f => Inply (Generics.Rec1 f :: Type -> Type) where
    gather f g (Generics.Rec1 x) (Generics.Rec1 y) = Generics.Rec1 (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Generics.Rec1 f :: Type -> Type) where
    knot = Generics.Rec1 . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Monoid.Alt f) where
    gather f g (Monoid.Alt x) (Monoid.Alt y) = Monoid.Alt (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Monoid.Alt f) where
    knot = Monoid.Alt . knot
-- | @since 0.4.1.0
instance Inply f => Inply (IdentityT f :: Type -> Type) where
    gather f g (IdentityT x) (IdentityT y) = IdentityT (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (IdentityT f :: Type -> Type) where
    knot = IdentityT . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Reverse f :: Type -> Type) where
    gather f g (Reverse x) (Reverse y) = Reverse (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Reverse f :: Type -> Type) where
    knot = Reverse . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Backwards f :: Type -> Type) where
    gather f g (Backwards x) (Backwards y) = Backwards (gather f g x y)
-- | @since 0.4.1.0
instance Inplicative f => Inplicative (Backwards f :: Type -> Type) where
    knot = Backwards . knot
-- | @since 0.4.1.0
instance Inply f => Inply (Lift f) where
    gather f g = \case
      Pure  x -> \case
        Pure  y -> Pure (f x y)
        Other y -> Other (invmap (f x) (snd . g) y)
      Other x -> \case
        Pure  y -> Other (invmap (`f` y) (fst . g) x)
        Other y -> Other (gather f g x y)
-- | @since 0.4.1.0
instance Inply f => Inplicative (Lift f) where
    knot = Pure

-- | @since 0.4.1.0
deriving via WrappedApplicativeOnly (Tagged a) instance Inply (Tagged a)
-- | @since 0.4.1.0
deriving via WrappedApplicativeOnly (Tagged a) instance Inplicative (Tagged a)

-- | @since 0.4.1.0
deriving via WrappedFunctor Identity instance Inply Identity
-- | @since 0.4.1.0
deriving via WrappedFunctor Identity instance Inplicative Identity
-- | @since 0.4.1.0
deriving via WrappedFunctor (Proxy :: Type -> Type) instance Inply Proxy
-- | @since 0.4.1.0
deriving via WrappedFunctor (Proxy :: Type -> Type) instance Inplicative Proxy
-- | @since 0.4.1.0
deriving via WrappedFunctor [] instance Inply []
-- | @since 0.4.1.0
deriving via WrappedFunctor [] instance Inplicative []
-- | @since 0.4.1.0
deriving via WrappedFunctor ((->) r) instance Inply ((->) r)
-- | @since 0.4.1.0
deriving via WrappedFunctor ((->) r) instance Inplicative ((->) r)
-- | @since 0.4.1.0
deriving via WrappedFunctor Maybe instance Inply Maybe
-- | @since 0.4.1.0
deriving via WrappedFunctor Maybe instance Inplicative Maybe
-- | @since 0.4.1.0
deriving via WrappedFunctor (Either e) instance Inply (Either e)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Either e) instance Inplicative (Either e)
-- | @since 0.4.1.0
deriving via WrappedFunctor IO instance Inply IO
-- | @since 0.4.1.0
deriving via WrappedFunctor IO instance Inplicative IO
-- | @since 0.4.1.0
deriving via WrappedFunctor Generics.Par1 instance Inply Generics.Par1
-- | @since 0.4.1.0
deriving via WrappedFunctor Generics.Par1 instance Inplicative Generics.Par1
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.U1 :: Type -> Type) instance Inply Generics.U1
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.U1 :: Type -> Type) instance Inplicative Generics.U1
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.K1 i c :: Type -> Type) instance Semigroup c => Inply (Generics.K1 i c)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.K1 i c :: Type -> Type) instance Monoid c => Inplicative (Generics.K1 i c)
-- | @since 0.4.1.0
deriving via WrappedFunctor Complex instance Inply Complex
-- | @since 0.4.1.0
deriving via WrappedFunctor Complex instance Inplicative Complex
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Min instance Inply Semigroup.Min
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Min instance Inplicative Semigroup.Min
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Max instance Inply Semigroup.Max
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Max instance Inplicative Semigroup.Max
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.First instance Inply Semigroup.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.First instance Inplicative Semigroup.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Last instance Inply Semigroup.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Last instance Inplicative Semigroup.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Option instance Inply Semigroup.Option
-- | @since 0.4.1.0
deriving via WrappedFunctor Semigroup.Option instance Inplicative Semigroup.Option
-- | @since 0.4.1.0
deriving via WrappedFunctor ZipList instance Inply ZipList
-- | @since 0.4.1.0
deriving via WrappedFunctor ZipList instance Inplicative ZipList
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.First instance Inply Monoid.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.First instance Inplicative Monoid.First
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Last instance Inply Monoid.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Last instance Inplicative Monoid.Last
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Dual instance Inply Monoid.Dual
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Dual instance Inplicative Monoid.Dual
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Sum instance Inply Monoid.Sum
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Sum instance Inplicative Monoid.Sum
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Product instance Inply Monoid.Product
-- | @since 0.4.1.0
deriving via WrappedFunctor Monoid.Product instance Inplicative Monoid.Product
-- | @since 0.4.1.0
deriving via WrappedFunctor NonEmpty instance Inply NonEmpty
-- | @since 0.4.1.0
deriving via WrappedFunctor NonEmpty instance Inplicative NonEmpty
-- | @since 0.4.1.0
deriving via WrappedFunctor Tree instance Inply Tree
-- | @since 0.4.1.0
deriving via WrappedFunctor Tree instance Inplicative Tree
-- | @since 0.4.1.0
deriving via WrappedFunctor Seq instance Inply Seq
-- | @since 0.4.1.0
deriving via WrappedFunctor Seq instance Inplicative Seq
-- | @since 0.4.1.0
deriving via WrappedFunctor NESeq.NESeq instance Inply NESeq.NESeq
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedArrow a b) instance Arrow a => Inply (WrappedArrow a b)
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedArrow a b) instance Arrow a => Inplicative (WrappedArrow a b)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Generics.V1 :: Type -> Type) instance Inply Generics.V1
-- | @since 0.4.1.0
deriving via WrappedFunctor IM.IntMap instance Inply IM.IntMap
-- | @since 0.4.1.0
deriving via WrappedFunctor (M.Map k) instance Ord k => Inply (M.Map k)
-- | @since 0.4.1.0
deriving via WrappedFunctor (HM.HashMap k) instance (Hashable k, Eq k) => Inply (HM.HashMap k)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Const w :: Type -> Type) instance Semigroup w => Inply (Const w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Const w :: Type -> Type) instance Monoid w => Inplicative (Const w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Constant w :: Type -> Type) instance Semigroup w => Inply (Constant w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (Constant w :: Type -> Type) instance Monoid w => Inplicative (Constant w)
-- | @since 0.4.1.0
deriving via WrappedFunctor (ContT r (m :: Type -> Type)) instance Inply (ContT r m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (ContT r (m :: Type -> Type)) instance Inplicative (ContT r m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedMonad m) instance Monad m => Inply (WrappedMonad m)
-- | @since 0.4.1.0
deriving via WrappedFunctor (WrappedMonad m) instance Monad m => Inplicative (WrappedMonad m)
-- | @since 0.4.1.0
deriving via WrappedFunctor ((,) w :: Type -> Type) instance Semigroup w => Inply ((,) w)
-- | @since 0.4.1.0
deriving via WrappedFunctor ((,) w :: Type -> Type) instance Monoid w => Inplicative ((,) w)

-- | @since 0.4.1.0
deriving via WrappedDivisible SettableStateVar instance Inply SettableStateVar
-- | @since 0.4.1.0
deriving via WrappedDivisible SettableStateVar instance Inplicative SettableStateVar
-- | @since 0.4.1.0
deriving via WrappedDivisible Predicate instance Inply Predicate
-- | @since 0.4.1.0
deriving via WrappedDivisible Predicate instance Inplicative Predicate
-- | @since 0.4.1.0
deriving via WrappedDivisible Comparison instance Inply Comparison
-- | @since 0.4.1.0
deriving via WrappedDivisible Comparison instance Inplicative Comparison
-- | @since 0.4.1.0
deriving via WrappedDivisible Equivalence instance Inply Equivalence
-- | @since 0.4.1.0
deriving via WrappedDivisible Equivalence instance Inplicative Equivalence
-- | @since 0.4.1.0
deriving via WrappedDivisible (Op r) instance Semigroup r => Inply (Op r)
-- | @since 0.4.1.0
deriving via WrappedDivisible (Op r) instance Monoid r => Inplicative (Op r)

-- | Convenient wrapper to build up an 'Inplicative' instance by providing
-- each component of it.  This makes it much easier to build up longer
-- chains because you would only need to write the splitting/joining
-- functions in one place.
--
-- For example, if you had a data type
--
-- @
-- data MyType = MT Int Bool String
-- @
--
-- and an invariant functor and 'Inplicative' instance @Prim@
-- (representing, say, a bidirectional parser, where @Prim Int@ is
-- a bidirectional parser for an 'Int'@), then you could assemble
-- a bidirectional parser for a @MyType@ using:
--
-- @
-- invmap (\(MyType x y z) -> I x :* I y :* I z :* Nil)
--        (\(I x :* I y :* I z :* Nil) -> MyType x y z) $
--   gatheredN $ intPrim
--                    :* boolPrim
--                    :* stringPrim
--                    :* Nil
-- @
--
-- Some notes on usefulness depending on how many components you have:
--
-- *    If you have 0 components, use 'knot' directly.
-- *    If you have 1 component, you don't need anything.
-- *    If you have 2 components, use 'gather' directly.
-- *    If you have 3 or more components, these combinators may be useful;
--      otherwise you'd need to manually peel off tuples one-by-one.
--
-- @since 0.4.1.0
gatheredN
    :: Inplicative f
    => NP f as
    -> f (NP I as)
gatheredN = \case
    Nil     -> knot Nil
    x :* xs -> gather
      (\y ys -> I y :* ys)
      (\case I y :* ys -> (y, ys))
      x
      (gatheredN xs)

-- | Given a function to "break out" a data type into a 'NP' (tuple) and one to
-- put it back together from the tuple, 'gather' all of the components
-- together.
--
-- For example, if you had a data type
--
-- @
-- data MyType = MT Int Bool String
-- @
--
-- and an invariant functor and 'Inplicative' instance @Prim@
-- (representing, say, a bidirectional parser, where @Prim Int@ is
-- a bidirectional parser for an 'Int'@), then you could assemble
-- a bidirectional parser for a @MyType@ using:
--
-- @
--   concaMapInplicative
--      (\(MyType x y z) -> I x :* I y :* I z :* Nil)
--      (\(I x :* I y :* I z :* Nil) -> MyType x y z)
--      $ intPrim
--     :* boolPrim
--     :* stringPrim
--     :* Nil
-- @
--
-- See notes on 'gatheredNMap' for more details and caveats.
--
-- @since 0.4.1.0
gatheredNMap
    :: Inplicative f
    => (NP I as -> b)
    -> (b -> NP I as)
    -> NP f as
    -> f b
gatheredNMap f g = invmap f g . gatheredN

-- | A version of 'gatheredN' for non-empty 'NP', but only
-- requiring an 'Inply' instance.
--
-- @since 0.4.1.0
gatheredN1
    :: Inply f
    => NP f (a ': as)
    -> f (NP I (a ': as))
gatheredN1 (x :* xs) = case xs of
    Nil    -> invmap ((:* Nil) . I) (\case I y :* _ -> y) x
    _ :* _ -> gather
      (\y ys -> I y :* ys)
      (\case I y :* ys -> (y, ys))
      x
      (gatheredN1 xs)

-- | A version of 'gatheredNMap' for non-empty 'NP', but only
-- requiring an 'Inply' instance.
--
-- @since 0.4.1.0
gatheredN1Map
    :: Inplicative f
    => (NP I (a ': as) -> b)
    -> (b -> NP I (a ': as))
    -> NP f (a ': as)
    -> f b
gatheredN1Map f g = invmap f g . gatheredN1

-- | A version of 'gatheredN' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of tuple components.
--
-- @since 0.4.1.0
gatheredNRec
    :: Inplicative f
    => V.Rec f as
    -> f (V.XRec V.Identity as)
gatheredNRec = \case
    V.RNil    -> knot V.RNil
    x V.:& xs -> gather
      (V.::&)
      (\case y V.::& ys -> (y, ys))
      x
      (gatheredNRec xs)

-- | A version of 'gatheredNMap' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of tuple components.
--
-- @since 0.4.1.0
gatheredNMapRec
    :: Inplicative f
    => (V.XRec V.Identity as -> b)
    -> (b -> V.XRec V.Identity as)
    -> V.Rec f as
    -> f b
gatheredNMapRec f g = invmap f g . gatheredNRec

-- | Convenient wrapper to 'gather' over multiple arguments using tine
-- vinyl library's multi-arity uncurrying facilities.  Makes it a lot more
-- convenient than using 'gather' multiple times and needing to accumulate
-- intermediate types.
--
-- For example, if you had a data type
--
-- @
-- data MyType = MT Int Bool String
-- @
--
-- and an invariant functor and 'Inplicative' instance @Prim@
-- (representing, say, a bidirectional parser, where @Prim Int@ is
-- a bidirectional parser for an 'Int'@), then you could assemble
-- a bidirectional parser for a @MyType@ using:
--
-- @
-- 'gatherN'
--   MT                                         -- ^ curried assembling function
--   (\(MT x y z) -> x ::& y ::& z ::& XRNil)   -- ^ disassembling function
--   (intPrim :: Prim Int)
--   (boolPrim :: Prim Bool)
--   (stringPrim :: Prim String)
-- @
--
-- Really only useful with 3 or more arguments, since with two arguments
-- this is just 'gather' (and with zero arguments, you can just use
-- 'knot').
--
-- The generic type is a bit tricky to understand, but it's easier to
-- understand what's going on if you instantiate with concrete types:
--
-- @
-- ghci> :t gatherN @MyInplicative @'[Int, Bool, String]
--      (Int -> Bool -> String -> b)
--   -> (b -> XRec Identity '[Int, Bool, String])
--   -> MyInplicative Int
--   -> MyInplicative Bool
--   -> MyInplicative String
--   -> MyInplicative b
-- @
--
-- @since 0.4.1.0
gatherN
    :: forall f as b. (Inplicative f, V.IsoXRec V.Identity as, V.RecordCurry as)
    => V.Curried as b
    -> (b -> V.XRec V.Identity as)
    -> V.CurriedF f as (f b)
gatherN f g = V.rcurry @as @f $
    invmap (V.runcurry' f . V.fromXRec) g
  . gatheredNRec

-- | A version of 'gatheredN1' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of components.
--
-- @since 0.4.1.0
gatheredN1Rec
    :: Inply f
    => V.Rec f (a ': as)
    -> f (V.XRec V.Identity (a ': as))
gatheredN1Rec (x V.:& xs) = case xs of
    V.RNil   -> invmap (V.::& V.RNil) (\case z V.::& _ -> z) x
    _ V.:& _ -> gather
      (V.::&)
      (\case y V.::& ys -> (y, ys))
      x
      (gatheredN1Rec xs)

-- | A version of 'gatheredNMap' using 'V.XRec' from /vinyl/ instead of
-- 'NP' from /sop-core/.  This can be more convenient because it doesn't
-- require manual unwrapping/wrapping of tuple components.
--
-- @since 0.4.1.0
gatheredN1MapRec
    :: Inplicative f
    => (V.XRec V.Identity (a ': as) -> b)
    -> (b -> V.XRec V.Identity (a ': as))
    -> V.Rec f (a ': as)
    -> f b
gatheredN1MapRec f g = invmap f g . gatheredN1Rec

-- | 'gatherN' but with at least one argument, so can be used with any
-- 'Inply'.
--
-- @since 0.4.1.0
gatherN1
    :: forall f a as b. (Inply f, V.IsoXRec V.Identity as, V.RecordCurry as)
    => V.Curried (a ': as) b
    -> (b -> V.XRec V.Identity (a ': as))
    -> V.CurriedF f (a ': as) (f b)
gatherN1 f g = V.rcurry @(a ': as) @f $
    invmap (V.runcurry' f . V.fromXRec) g
  . gatheredN1Rec

-- | Interpret out of a contravariant 'Day' into any instance of 'Apply' by
-- providing two interpreting functions.
--
-- In theory, this should not need to exist, since you should always be
-- able to use 'runDay' because every instance of 'Apply' is also an
-- instance of 'Inply'.  However, this can be handy if you are using an
-- instance of 'Apply' that has no 'Inply' instance.  Consider also
-- 'unsafeInplyCo' if you are using a specific, concrete type for @h@.
runDayApply
    :: forall f g h. Apply h
    => f ~> h
    -> g ~> h
    -> Day f g ~> h
runDayApply f g (Day x y j _) = liftF2 j (f x) (g y)

-- | Interpret out of a contravariant 'Day' into any instance of 'Divise'
-- by providing two interpreting functions.
--
-- In theory, this should not need to exist, since you should always be
-- able to use 'runDay' because every instance of 'Divise' is also an
-- instance of 'Inply'.  However, this can be handy if you are using an
-- instance of 'Divise' that has no 'Inply' instance.  Consider also
-- 'unsafeInplyContra' if you are using a specific, concrete type for @h@.
runDayDivise
    :: forall f g h. Divise h
    => f ~> h
    -> g ~> h
    -> Day f g ~> h
runDayDivise f g (Day x y _ h) = divise h (f x) (g y)

max_stars_repo_name: jonascarpay/functor-combinators

max_issues_repo_name: jonascarpay/functor-combinators

max_forks_repo_name: jonascarpay/functor-combinators


text:
{-# LANGUAGE RankNTypes, BangPatterns, ScopedTypeVariables, FlexibleInstances #-}

module Math.Probably.MALA where

import Math.Probably.MCMC
import qualified Math.Probably.PDF as PDF
import Math.Probably.RandIO
import Math.Probably.FoldingStats
import Math.Probably.Sampler
import qualified Math.Probably.PDF as PDF
import Control.Applicative

import Numeric.LinearAlgebra
import Numeric.AD
import Text.Printf
import System.IO
import Data.Maybe

import Statistics.Test.KolmogorovSmirnov
import Statistics.Test.MannWhitneyU
import qualified Data.Vector.Unboxed as U
import qualified Control.Monad.State.Strict as S
import qualified Data.Vector.Storable as VS

import Debug.Trace
import Data.IORef
import Control.Spoon

class IsNaN a where
  isItNaN :: a -> Bool

instance IsNaN Double where
  isItNaN = nanOrInf

instance IsNaN (Vector Double) where
  isItNaN = any nanOrInf . toList

instance IsNaN (Matrix Double) where
  isItNaN = or . (map (any nanOrInf . toList)) . toRows

checkNaN :: IsNaN a => String -> a -> String
checkNaN s x | isItNaN x = " "++s++" "
             | otherwise = ""

class Covariance a where
  mvnSampler :: Vector Double ->  Double -> a -> Sampler (Vector Double)
  mvnPDF :: Vector Double -> Double -> a -> PDF.PDF (Vector Double)
  covMul :: a -> Vector Double -> Vector Double
  restrict :: Vector Int -> a -> a

instance Covariance (Matrix Double,Matrix Double,Matrix Double) where
  mvnSampler mu sigma (cov, covInv, covChol)
     = multiNormalByChol mu (scale (sqrt sigma) covChol)
  mvnPDF mu sigma (cov, covInv, covChol)
     = PDF.multiNormalByInvFixCov (scale (recip sigma) covInv) mu
  covMul (cov, covInv, covChol) v = cov <> v

instance Covariance (Vector Double) where
  mvnSampler mu sigma varv
    = multiNormalIndep (scale sigma varv) mu
  mvnPDF mu sigma varv
    = PDF.multiNormalIndep (scale sigma varv) mu
  covMul varv v = VS.zipWith (*) varv  v
  restrict vixs v = VS.backpermute v vixs

data MalaPar = MalaPar { mpXi :: !(Vector Double),
                         mpPi :: !Double,
                         mpGradLast :: !(Vector Double),
                         mpSigma :: !Double,
                         mpCount :: !Double,
                         mpAccept :: !Double,
                         mpFreezeSigma :: !Bool,
                         mpThin :: !Int,
                         mpMaxGradLen :: !Double } --,
--                         mpLastRatio :: !Double }
  deriving Show

data Pair a b = Pair !a !b

data List a = Cons !a !(List a)
            | Nil

unPair (Pair x y) = (x,y)
unList (Cons x xs) = x: unList xs
unList Nil =[]

trunc t v = scale (t/(max t (norm2 v))) v

mala1 :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Bool
         -> MalaPar
         -> Sampler MalaPar
mala1 cov postgrad useCache
      (MalaPar xi piCached gradientiCached sigma tr tracc freeze thinn maxGrad) = do
  let (pi,gradienti) = if useCache
                          then (piCached, gradientiCached)
                          else let (p, gr_untr) = postgrad xi
                               in (p, trunc maxGrad gr_untr)
  let xstarMean = xi + scale (sigma/2) (cov `covMul` gradienti)
  xstar <- mvnSampler xstarMean sigma cov
  u <- unitSample
  let (!pstar, gradientStarUntrunc) = postgrad xstar
      gradientStar = trunc maxGrad gradientStarUntrunc
  let !revJumpMean = xstar + scale (sigma/2) (cov `covMul` gradientStar)
      ptop = mvnPDF revJumpMean sigma cov xi
      pbot = mvnPDF xstarMean sigma cov xstar
      ratio = exp $   pstar -pi + ptop - pbot
      tr' = max 1 tr
      sigmaNext = case () of
         _ | freeze -> sigma
      accept = tracc / tr
      freezeNext = freeze {-| freeze = True
                 | not freeze = tr > 100 && accept > 0.5 && accept < 0.6  -}
  if trace (show $ (tr, pstar ,pi , ratio, sigma)) $ u < ratio
     then return $ MalaPar xstar pstar (gradientStar)
                           (if freezeNext then sigma else (min 1.4 $ 1+kmala/tr')*sigma)
                           (tr+1) (tracc+1) freezeNext  thinn maxGrad
--                           sigma (tr+1) (tracc+1)
     else return $ MalaPar xi pi ( gradienti)
                           (if freezeNext then sigma else (max 0.7143 $ 1-kmala/tr')**1.3*sigma)
                           (tr+1) tracc freezeNext thinn maxGrad
--                           sigma (tr+1) tracc

blockMala1 :: Covariance a =>
  (Vector Double -> (Double,Vector Double)) ->
  Vector Double ->
  [(a, MalaPar, Vector Int)] ->
  Sampler (Vector Double, [(a, MalaPar, Vector Int)])
blockMala1 postgrad xi blocks = do
  let go totalV [] accblocks = return (totalV, reverse accblocks)
      go totalV ((cov, mp0, vixs) : blocks) accblocks = do
          let myPostGrad vpars
               = let (p, grad) = postgrad $ VS.update_ totalV vixs vpars
                 in (p, VS.backpermute grad vixs)
          mp1 <- mala1 cov myPostGrad False mp0
          let newTotalV = VS.update_ totalV vixs (mpXi mp1)
          go newTotalV blocks $ (cov, mp1, vixs) : accblocks
  go xi blocks []

runMalaBlocks :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         ->  Int -> Double -> Int -> Vector Double -> [Vector Int]
         -> RIO [Vector Double]
runMalaBlocks cov  postgrad nsam truncN thinN init  vixs = go nsam initBlocks init []  where
  initBlocks = flip map vixs $ \vix->
     (restrict vix cov,
      MalaPar (VS.backpermute init vix) 0 (fromList [0]) 0.001 0 0 False thinN truncN,
      vix)
  go 0 _ _ vs = return vs
  go n blocks0 v0 vs = do
        (!v1, blocks1) <- sample $ blockMala1 postgrad v0 blocks0
        let (_,mpLast,_) = last blocks1
        io $ print $ (n, mpPi mpLast, map (\(_,mp,_) -> mpSigma mp) blocks1)
        let newChainRes = if thinN == 0 || n `mod` thinN ==0
                             then v1 : vs
                             else vs
        go (n-1) blocks1 v1 newChainRes

{-runMala :: Matrix Double -> (Vector Double -> (Double,Vector Double))
         ->  Int -> Vector Double -> RIO [(Double,Vector Double)]
runMala cov postgrad nsam init = go nsam mp1 [] where
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return xs
  go n y xs = do y1 <- sample $ mala1 cov postgrad y
                 go (n-1) y1 $ (mpPi y1, mpXi y1):xs
  (pi, gradi) = postgrad init
  mp1 =  MalaPar init pi (gradi)  1 0 0 False -}

runMalaMP :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> [(Double,Vector Double)] -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMP cov  pdf nsam init xs0 = go nsam init xs0 where
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov pdf True y
--                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
--                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then let !xi = mpXi y1
                                               !pi = mpPi y1
                                               !more = (pi,xi)
                                           in more:xs
                                      else xs
                 go (n-1) y1 newChainRes

runMalaUntilBetter :: Covariance a => a ->(Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaUntilBetter cov  pdf nsam init = go nsam init [] where
  pTarget = mpPi init
  go 0 mpar xs = do io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
                    io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
                    return (mpar, xs)
  go n y xs = do y1 <- sample $ mala1 cov  pdf True y
                 io $ do putStrLn $ show (mpCount y1, mpPi y1, mpSigma y1)
                         hFlush stdout
                 let newChainRes = if mpThin y1 == 0 || round (mpCount y1) `mod` mpThin y1 ==0
                                      then (mpPi y1, mpXi y1):xs
                                      else xs
                 if mpPi y1 > pTarget
                    then return (y1, newChainRes)
                    else go (n-1) y1 $ newChainRes

{-runMalaMPaccept' :: Matrix Double -> (Vector Double -> (Double,Vector Double))
         ->  Int -> MalaPar -> RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept'  cov pdf nsam init  = do
  stseed <- S.get
  res <- io $ do
   seedref <- newIORef stseed
   resRef <- newIORef Nil
   let go !mpar
        | (round $ mpAccept mpar) >= nsam
          = return mpar
        | otherwise
          = do seed <- readIORef seedref
               putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
               let (!mpar1, !seed1) = unSam (mala1 cov pdf mpar) seed
               writeIORef seedref seed1
               modifyIORef resRef (Cons (Pair (mpPi mpar1) ( mpXi mpar1)))
               go mpar1
   mpar <- go init
   seed <- readIORef seedref
   reslist <- readIORef resRef
   return (seed, (mpar,reslist))
  S.put $ fst res
  let mp = fst $ snd res
  return $ (mp, map unPair $ unList $ snd $ snd res)

--  pi = pdf $ toList $ mpXi init -}

runMalaMPaccept :: Covariance a => a -> (Vector Double -> (Double,Vector Double))
         -> Int -> MalaPar ->  RIO (MalaPar, [(Double,Vector Double)])
runMalaMPaccept cov pdf nsam init = go init [] where
  go !mpar !xs
    | (round $ mpAccept mpar) >= nsam = do
         io $ putStrLn $ "MALA accept = "++show (mpAccept mpar/mpCount mpar)
         io $ putStrLn $ "MALA sigma = "++show (mpSigma mpar)
         return (mpar, xs)
    | mpAccept mpar < 0.5 && mpCount mpar > 20 =
         return (mpar, [])
    | otherwise = do
         !mpar1 <- sample $ mala1 cov  pdf True mpar
         io $ putStr ((show (round $ mpAccept mpar)) ++".") >> hFlush stdout
         go mpar1 $ (mpPi mpar1, mpXi mpar1):xs

kmala = 5

{-runMalaRioESS ::  Matrix Double -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Vector Double -> RIO [Vector Double]
runMalaRioESS cov pdf want_ess xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 grad0 sigma0 0 0 False
        nsam0 = want_ess*20
    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMP cov pdf nsam0 mp0 []
    let have_ess = min (mpAccept mp1) $ calcESSprim $ map snd xs1
    if have_ess > realToFrac want_ess
       then return $ map snd xs1
       else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                   samples_per_es = realToFrac nsam0/have_ess
                   to_do = round $ samples_per_es * need_ess
               (mp2, xs2) <- runMalaMP cov pdf to_do mp1 xs1
               return  $ map snd  xs2 -}

runMalaRioCodaESS ::  Covariance a => a-> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Vector Double -> RIO [Vector Double]
runMalaRioCodaESS cov pdf want_ess truncN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc truncN grad0) sigma0 0 0 False 0 truncN
        nsam0 = want_ess*1
    let converged mp  xs = do
         let have_ess = min (mpAccept mp) $ calcESSprim $ map snd xs
         io $ putStrLn $ "ESS=" ++show have_ess
         if have_ess > realToFrac want_ess
            then return $ map snd xs
            else do let need_ess =  max 1 $ realToFrac want_ess - have_ess
                        samples_per_es = realToFrac nsam0/have_ess
                        to_do = round $ samples_per_es * need_ess
                    io $ putStrLn $ "running converged for "++show to_do
                    (mp2, xs2) <- runMalaMP cov pdf to_do (mp {mpFreezeSigma = True}) xs
                    io $ putStrLn $ "All done"
                    return $ map snd xs2
    let go mp  n xs = do
            (mp2, xs2) <- runMalaMP cov  pdf n mp []
            let testres = mannWhitneyUtest TwoTailed 0.05 (U.fromList $ map fst xs2)
                                            (U.fromList $ map fst xs)

            if testres/= Just NotSignificant
               then do io$ putStrLn $ "not converged: "++show testres++" at "++show (mpPi mp2)
                       go mp2  (round $ realToFrac n*2) xs2
               else do io$ putStrLn "converged!"
                       converged mp2  (xs2++xs)

    io $ putStrLn $ "initial sigma = "++show (mpSigma mp0)
    (mp1, xs1) <- runMalaMPaccept cov  pdf nsam0 mp0
    case xs1 of
            [] -> return []
            _ -> go mp1  (round $ mpCount mp1) xs1

{-    case (spoon (invlndet cov), mbCholSH cov) of
       (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
       _ -> let cov' = PDF.posdefify cov in
            case (spoon (invlndet cov'), mbCholSH cov') of
              (Just (covInv, (lndt,_)), Just covChol) -> go_rest (covInv, covChol)
              _ ->  do io $ putStrLn "non-invertible covariance matrix"
                       return [] -}

runMalaRioSimple ::  Covariance a => a -> (Vector Double -> (Double,Vector Double))
                  ->  Int -> Double -> Int -> Vector Double -> RIO [Vector Double]
runMalaRioSimple cov pdf samples maxGrad thinN xi = do
    let (p0,grad0) = pdf xi
    let sigma0 = 2.7e-4 --1.5 / (realToFrac $ dim xi) -- determined empirically
    let mp0 = MalaPar xi p0 (trunc maxGrad grad0) sigma0 0 0 False thinN maxGrad

    (mp2, xs2) <- runMalaMP cov  pdf samples mp0 []
    io $ putStrLn $ "All done"
    return $ map snd xs2

calcCovariance :: Vector Double ->
                  Vector Double ->
                  (Vector Double -> (Double,Vector Double)) ->
                  (Vector Double -> Double) ->
                  Either (Vector Double) (Matrix Double, Matrix Double, Matrix Double)
calcCovariance vinit vnear postgrad posterior = finalcov where
   ndim = dim vinit
   finalcov
     | ndim > 0 -- > 20000 --FIXME
        = Left $ calcFDindepVars vinit vnear posterior --Left $ iCov vinit --
     | otherwise
        = hessToCov (calcFDhess vinit vnear postgrad) Nothing

iCov v = VS.replicate (VS.length v) 1 -- in (m,m,m)

calcFDindepVars v v' post = trace ("FDVars = "++show (VS.take 10 vars))  $ vars where
   hv =  mapVector (*1e-4) v --mapVector (max 1e-9 .  abs) $ v - v'
   n = dim v
   postv = post v
   postPlus i = post $ v VS.// [(i,v @>i + hv @> i)]
   postMinus i = post $ v VS.// [(i,v @>i - hv@> i)]
   vars = buildVector n fvar
   fvar i = negate $ recip $ (postPlus i - 2*postv + postMinus i)/((hv @> i)*(hv @> i))

calcFDhess v v' postgrad = hess2 where
   grad =  snd . postgrad
   gradi i = (@>i) . grad
   hv = mapVector (max 1e-9 . abs) $ v - v'
   n = dim v
   gradv = grad v
   grads =  fromRows $ flip map [0..(n-1)] $ \i ->
               grad (v VS.// [(i,v @>i + hv @> i)])
   fhess (i,j) | i<j = 0
               | otherwise = (grads @@>(j,i) - (gradv @> i))
                                       /(2*(hv @> j)) +
                             (grads @@>(i,j)- (gradv @> j))
                                       /(2*(hv @> i))

--   hess3 = scale (recip $ realToFrac n) $ sum $ map outerSelf $ grads

   hess1 = buildMatrix n n fhess
   hess2 = buildMatrix n n $ \(i,j) -> if i>=j then hess1 @@> (i,j)
                                               else hess1 @@> (j,i)

outerSelf v= v `outer` v

hessToCov hess mOriginal =
   let mTryPosdefify = case mOriginal of
            Nothing -> hessToCov (PDF.posdefify hess) (Just hess)
            Just horiginal -> let ds =  mapVector (negate . recip)
                                         $ takeDiag horiginal
                              in trace ("USING DIAGS"++show (VS.take 10 ds)) $ Left $ ds
   in
   case spoon $ inv $ negate $ hess of
     Just cov -> case mbCholSH cov of
                   Just cholm ->  trace ("invert success:"++show (VS.take 10 $ takeDiag cov) ++ "det="++show (det cov)) $ Right (cov, negate hess, cholm)
                   Nothing -> trace ("chol fail") mTryPosdefify
     Nothing -> trace ("inv fail") mTryPosdefify
                  {- Just cov -> case mbCholSH cov of
                                Just cholm ->  Right (cov, negate hessToCov,cholm)
                                Nothing -> case mbCholSH $ PDF.posdefify cov of
                                            Just cholm -> Right (cov, negate hess,cholm)
                                            Nothing -> vars
                  Nothing -> vars -}

acceptSM ampar  | mpCount ampar == 0 = "0/0"
               | otherwise = printf "%.3g" (rate::Double) ++ " ("++show yes++"/"++show total++")" where
   rate = realToFrac (yes) / realToFrac (total)
   yes = mpAccept ampar
   total = mpCount ampar

ext: hs

max_stars_repo_name: glutamate/probably

max_issues_repo_name: glutamate/probably

max_forks_repo_name: glutamate/probably


text:
module Q.Options.ImpliedVol.Normal where
import           Data.Default.Class (Default (..))
import           Numeric.IEEE (epsilon)
import           Numeric.Natural (Natural)
import           Numeric.RootFinding as R (RiddersParam (RiddersParam),
                                           Root (NotBracketed, Root, SearchFailed), Tolerance,
                                           ridders)
import           Q.Options.Bachelier (Bachelier (Bachelier), euOption)

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

-- | Method to use to calculate the normal implied vol.
data Method =
    Jackel        -- ^ Jackel analytical formula approximation.
  | ChoKimKwak    -- ^ J. Choi, K kim, and M. Kwak (2009)
  -- | Numerical root finding. Currently Ridders is used.
  | RootFinding
        Natural                  -- ^ Maximum number of iterations.
        R.Tolerance              -- ^ Tolerance (relative or absolute)
        (Double, Double, Double) -- ^ Triple of @(low bound, initial
                                 --   guess, upper bound)@. If initial
                                 --   guess if out of bracket middle
                                 --   of bracket is taken as.

instance Default Method where
  def = Jackel

-- | Default method implementation of 'euImpliedVolWith' using 'Jackel'.
euImpliedVol :: OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVol = euImpliedVolWith def

-- | Calcualte the bachelier option implied vol of a european option.
--
-- If the options premium does not have time value @'hasTimeValue'@ return 0.
euImpliedVolWith :: Method -> OptionType -> Forward -> Strike -> YearFrac -> Rate -> Premium -> Vol
euImpliedVolWith m cp f k t r p
  | hasTimeValue cp f k p df = euImpliedVolWith' m cp f k t r p
  | otherwise                = Vol $ 0
  where df = discountFactor t r

euImpliedVolWith' :: Method-> OptionType-> Forward-> Strike-> YearFrac-> Rate-> Premium-> Vol
euImpliedVolWith' Jackel cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p)
  -- Case where interest rate is not 0 we need undiscount. The paper is written
  -- for the undiscounted Bachelier option prices.
  | r /= 0
    = euImpliedVol cp (Forward f) (Strike k) (YearFrac t) (Rate 0) (Premium (p/df))
  -- Case of ATM. Solve directly.
  | abs (k - f) <= epsilon = Vol $ p * sqrt2Pi / (sqrt t)
  -- Case of ITM option. Calcualte vol of the out of the money option with Put-Call-Parity.
  | phiStarTilde >= 0
    = euImpliedVol (succ cp) (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p')
  -- General case for an out of the money option.
  | otherwise  = let
      ẋ      = if phiStarTilde < c then
                 let g = 1 / (phiStarTilde - 0.5)
                     ξ = (0.032114372355 - (g**2)*(0.016969777977 - (g**2)*(2.6207332461E-3-(9.6066952861E-5)*g**2)))
                         /
                         (1-(g**2)*(0.6635646938 - (g**2)*(0.14528712196 - 0.010472855461*g**2)))
                 in g * (1 / sqrt2Pi + ξ*g**2)
               else
                 let h = sqrt $ (-log (-phiStarTilde))
                 in (9.4883409779-h*(9.6320903635-h*(0.58556997323 + 2.1464093351*h)))
                    /
                    (1-h*(0.65174820867 + h*(1.5120247828 + 6.6437847132E-5*h)))
      c       = (-0.001882039271)
      x       = ẋ + (3*q * ẋ * ẋ * (2 - q * ẋ * (2 + ẋ*ẋ)))
                    /
                    (6 + q*ẋ * ((-12) + ẋ *(6*q + ẋ * ((-6)*q*ẋ*(3+ẋ*ẋ)))))
      phiXBarTilde = (cumulative standard ẋ) + (density standard ẋ)/ẋ
      q       =  (phiXBarTilde-phiStarTilde)/ (density standard ẋ)
    in Vol $ (abs (k - f)) / (abs (x * sqrt t))
  where phiStarTilde = negate $ (abs (p - (max (theta * (f - k)) 0))) / (abs (k - f))
        theta        = if cp == Call then 1 else -1
        phiTilde     = (-theta) * p / (k - f)
        p'           = cpi * df * (f - k) + p
        cpi          = fromIntegral $ fromEnum cp --call put indicartor.
        df           = exp $ (-r) * t
        sqrt2Pi      = 2.506628274631000

euImpliedVolWith' ChoKimKwak cp (Forward f) (Strike k) (YearFrac t) (Rate r) (Premium p) =
  let df              = exp $ (-r) * t
      forwardPremium  = p / df
      straddlePremium = case cp of Call -> 2 * forwardPremium - (f - k)
                                   Put  -> 2 * forwardPremium + (f - k)
      nu'             = (f - k) / straddlePremium
      nu              = max (-1 + epsilon) (min nu' (1 - epsilon))
      eta             | abs nu < sqrtEpsilon = 1
                      | otherwise            = nu / (atanh nu)
      heta            = h eta
  in Vol $ sqrt (pi / (2 * t)) * straddlePremium * heta

euImpliedVolWith' (RootFinding maxIter tol bracket)  cp (Forward forward) k t r (Premium p) =
  let f vol        = p' - p  where
        (Premium p') = vPremium $ euOption b cp k
        b            = Bachelier t (Forward forward) r (Vol vol)
      (lb, _, ub) = bracket
      root = ridders (RiddersParam (fromEnum maxIter) tol) (lb, ub) f
  in case root of (Root vol)   -> Vol vol
                  NotBracketed -> error "not bracketed"
                  SearchFailed -> error "search failed"

sqrtEpsilon :: Double
sqrtEpsilon = sqrt epsilon

h :: Floating p => p -> p
h eta = sqrt(eta) * (num / den) where
  a0          = 3.994961687345134e-1
  a1          = 2.100960795068497e+1
  a2          = 4.980340217855084e+1
  a3          = 5.988761102690991e+2
  a4          = 1.848489695437094e+3
  a5          = 6.106322407867059e+3
  a6          = 2.493415285349361e+4
  a7          = 1.266458051348246e+4

  b0          = 1.000000000000000e+0
  b1          = 4.990534153589422e+1
  b2          = 3.093573936743112e+1
  b3          = 1.495105008310999e+3
  b4          = 1.323614537899738e+3
  b5          = 1.598919697679745e+4
  b6          = 2.392008891720782e+4
  b7          = 3.608817108375034e+3
  b8          = -2.067719486400926e+2
  b9          = 1.174240599306013e+1

  num = a0 + eta * (a1 + eta * (a2 + eta * (a3 + eta * (a4 + eta * (a5 + eta * (a6 + eta * a7))))))
  den = b0 + eta * (b1 + eta * (b2 + eta * (b3 + eta * (b4 + eta * (b5 + eta * (b6 + eta * (b7 + eta * (b8 + eta * b9))))))))

max_stars_repo_name: ghais/HQu

max_issues_repo_name: ghais/HQu

max_forks_repo_name: ghais/HQu


text:
{-# language BangPatterns #-}
{-|
Module      : Fitness
Description : Fitness Function for Genetic Algorithm for Poly Regression
Copyright   : (c) Fabricio Olivetti de Franca, 2021
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX
|-}

module Fitness (evalFitness, decode) where

import GA

import Data.List (foldl')
import Data.List.Split hiding (split)
import Numeric.LinearAlgebra ((<\>),(#>),(|||), sumElements, size, Matrix, Vector, fromRows, fromList)

-- | Creates a fitness function
evalFitness :: Int            -- ^ number of terms
            -> [[Double]]     -- ^ matrix X
            -> Vector Double  -- ^ vector y
            -> Solution (Poly n)       -- ^ solution to evaluate
            -> Solution (Poly n)
evalFitness nTerms xss ys sol =
  case _fitness sol of
       Nothing -> sol{ _coeffs = Just betas, _fitness = Just f }
       _       -> sol
  where
    zss    = decode xss (map _getPoly $ _chromo sol)
    betas  = zss <\> ys
    ysHat  = zss #> betas
    f      = mse ys ysHat

-- | Decodes the polynomial into a transformed matrix
decode :: [[Double]] -> [(Bool, Int)] -> Matrix Double
decode xss chromo = fromRows (map (evalPoly chromo) xss)

-- | Eval a single row from the dataset
evalPoly :: [(Bool, Int)] -> [Double] -> Vector Double
evalPoly chromo xs = fromList $ foldl' joinTerms [1.0] terms
  where
    nVars = length xs
    ix    = cycle [0..nVars-1]
    xss   = cycle xs
    terms = zip3 chromo xss ix

    joinTerms ts   ((False,_), _, 0)   = 1.0 : ts
    joinTerms ts   ((False,_), _, _)   = ts
    joinTerms ts   ((True, k), x, 0)   = x^k : ts
    joinTerms (t:ts) ((True,k), x, ix) = t * x^k : ts

-- | Calculates the mean squared error
mse :: Vector Double -> Vector Double -> Double
mse ys ysHat = s / fromIntegral (size ys)
  where
    s = sumElements $ (ys - ysHat)^2

max_stars_repo_name: folivetti/gapoly

max_issues_repo_name: folivetti/gapoly

max_forks_repo_name: folivetti/gapoly


text:
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Fields (C, Field (..), Fp, Prime (..), Q, R, Val (..), modInv) where

import Data.Complex (Complex (..))
import Data.Ratio (denominator, numerator)

------------------------------------------------------------------------------
-- Fields

-- Value information
data Val = None | NumDen Integer Integer | ReIm Double Double

-- Field type class
class (Eq f, Fractional f, Show f) => Field f where
  {-# MINIMAL char, val #-}
  char :: f -> Integer
  val :: f -> Val

------------------------------------------------------------------------------
-- Characteristic zero

-- Rational numbers Q
newtype Q = Q Rational deriving (Eq, Fractional, Num, Ord)
instance Show Q where
  show q = show n ++ if d == 1 then "" else "/" ++ show d
    where
      NumDen n d = val q
instance Field Q where
  char = const 0
  val (Q q) = NumDen (numerator q) (denominator q)

-- Real numbers R
newtype R = R Double deriving (Eq, Fractional, Num, Ord)
instance Show R where
  show (R r) = if r == fromIntegral r' then show r' else show r
    where
      r' = floor r
instance Field R where
  char = const 0
  val = const None

-- Complex numbers C
newtype C = C (Complex Double) deriving (Eq, Fractional, Num)
instance Show C where
  show (C (r :+ i))
    | signum i == -1 = "(" ++ show (R r) ++ "-" ++ show (R (abs i)) ++ "i)"
    | signum i == 1 = "(" ++ show (R r) ++ "+" ++ show (R i) ++ "i)"
    | otherwise = show (R r)
instance Field C where
  char = const 0
  val (C (r :+ i)) = ReIm r i

------------------------------------------------------------------------------
-- Characteristic prime

-- Prime numbers
class (Enum p, Show p) => Prime p where
  prime :: p -> Integer
  prime _ = read . tail . show $ (toEnum 0 :: p)

-- Prime subfields Fp
newtype Fp p = Fp Integer

-- Fp field instance
instance Prime p => Field (Fp p) where
  char = const $ prime (undefined :: p)
  val = const None

-- Fp standard instances
instance Prime p => Bounded (Fp p) where
  minBound = 0
  maxBound = -1
instance Prime p => Enum (Fp p) where
  toEnum = fromIntegral
  fromEnum (Fp n) = fromInteger n
instance Eq (Fp p) where
  Fp n == Fp n' = fromInteger n == fromInteger n'
instance Prime p => Fractional (Fp p) where
  fromRational n = fromInteger (numerator n) / fromInteger (denominator n)
  recip (Fp n) = case modInv n $ prime (undefined :: p) of
    Right m -> fromIntegral m
    Left m -> error $ show m
instance Prime p => Integral (Fp p) where
  quotRem (Fp n) (Fp n') = (fromInteger q, fromInteger r)
    where
      (q, r) = quotRem n n'
  toInteger (Fp n) = fromInteger n
instance Prime p => Num (Fp p) where
  Fp n + Fp n' = fromInteger $ n + n'
  Fp n * Fp n' = fromInteger $ n * n'
  abs n = n
  signum n = if n == 0 then 0 else 1
  fromInteger n = Fp . mod n $ prime (undefined :: p)
  negate (Fp n) = fromInteger $ (-n)
instance Ord (Fp p) where
  Fp n <= Fp n' = n <= n'
instance Prime p => Real (Fp p) where
  toRational (Fp n) = fromInteger n
instance Show (Fp p) where
  show (Fp n) = show n

-- Extended Euclidean algorithm
extGCD :: Integral a => a -> a -> ((a, a), a)
extGCD 0 y = ((0, 1), y)
extGCD x y = ((t - s * q, s), g)
  where
    (q, r) = quotRem y x
    ((s, t), g) = extGCD r x

-- Modular inverse
modInv :: Integral a => a -> a -> Either a a
modInv x p = if g == 1 then return (mod y p) else Left g
  where
    ((y, _), g) = extGCD x p

max_stars_repo_name: Multramate/EllipticCurves

max_issues_repo_name: Multramate/EllipticCurves

max_forks_repo_name: Multramate/EllipticCurves


text:
module ML.Unsupervised.Clustering.KMeans (
  fit,
  predict,
  KMeans(..)
) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data

import Data.Ratio ((%))
import Data.List (minimumBy)
import qualified Data.Vector.Storable as V

import Debug.Trace

data KMeans a l = KMeans {
  centroids :: [(l, V.Vector a)]
} deriving (Show, Eq)

fit :: (Fractional a, Ord a, Eq l, Show l, Show a, V.Storable a) => [l] -> [V.Vector a] -> KMeans a l
fit labels points = fit' (KMeans initialCentroids)
  where
    fit' kmeans = let
                    ass = assignment kmeans
                    updatedKmeans = update kmeans ass in
                    if (trace ("updated assignment: " ++ show (assignment updatedKmeans)) assignment updatedKmeans) == ass
                      then updatedKmeans
                      else fit' updatedKmeans
    update kmeans ass = KMeans $ zip labels $ fmap (\l -> centerOfMass $ snd <$> filter (\x -> fst x == l) ass) labels
    assignment kmeans = fmap (\x -> (assign kmeans x, x)) points
    assign kmeans point = fst $ minimumBy (\x1 x2 -> snd x1 `compare` snd x2) $ (fmap (\(l, cent) -> (l, distance cent point))) $ centroids kmeans
    distance v1 v2 = V.sum $ V.zipWith (\x1 x2 -> (x2 - x1) * (x2 - x1)) v1 v2
    initialCentroids = zip labels points

    centerOfMass :: (Fractional a, V.Storable a) => [V.Vector a] -> V.Vector a
    centerOfMass (p:pts) = V.map (/ (fromRational $ (toInteger $ length pts + 1) % 1)) $ foldr sumVectors p pts

    dim = V.length $ head points
    sumVectors :: (Num a, V.Storable a) => V.Vector a -> V.Vector a -> V.Vector a
    sumVectors = V.zipWith (+)

predict :: (Ord a, Num a, V.Storable a) => KMeans a l -> V.Vector a -> l
predict kmeans point = fst $ minimumBy (\x1 x2 -> snd x1 `compare` snd x2) $ (fmap (\(l, cent) -> (l, distance cent point))) $ centroids kmeans
  where
    distance v1 v2 = V.sum $ V.zipWith (\x1 x2 -> (x2 - x1) * (x2 - x1)) v1 v2

fitMatrix :: (Element a) => [l] -> Matrix a -> KMeans a l
fitMatrix = undefined

max_stars_repo_name: asakul/small-ml

max_issues_repo_name: asakul/small-ml

max_forks_repo_name: asakul/small-ml


text:
{-# LANGUAGE OverloadedStrings #-}
-- |
-- Module      :  Data.Random.Distribution.MultivariateNormal
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
module Q.Util.File (write)
  where

import Numeric.LinearAlgebra
import Control.Monad
import qualified Data.ByteString.Char8 as C
import Data.Csv
import Data.Char (ord)
import qualified Data.ByteString.Lazy as B
import Data.Random

rowToRecord :: (Show t) => [t] -> Record
rowToRecord x = record $ map (C.pack . show) x

write :: (Show t) => [[t]] -> [String] -> FilePath -> IO ()
write m header path = do
  let out = (encodeWith opt s) where
        opt = defaultEncodeOptions { encDelimiter = fromIntegral (ord ','), encQuoting = QuoteNone }
        rows :: [Record]
        rows = map rowToRecord $ m
        header_ = record $ map C.pack  header
        s = if null header_ then rows else header_:rows
  B.writeFile path out

max_stars_repo_name: ghais/lowq

max_issues_repo_name: ghais/lowq

max_forks_repo_name: ghais/lowq


text:
{-# LANGUAGE RecordWildCards #-}

module View where

import Constants
import Model

import Control.Arrow ((&&&))
import Data.Array
import Data.Bifunctor
import Data.Complex
import Graphics.Gloss
import Graphics.Gloss.Interface.IO.Interact

(<>) = mappend

draw :: World -> IO Picture
draw (World {..}) = pure $ (uncurry scale scaling . uncurry translate coords $ tiling board)

shearTile :: (Float, Float) -> (Float, Float)
shearTile (x, y) = (a*x + b*y, c*y)
  where
    a = rad * sqrt 3 + 1
    b = a / 2
    c = 15.866

-- | draw a board at each point in the lattice
tiling :: Board -> Picture
tiling board = pictures
             . map (\(x,y) -> translate x y (boardWithTiles rad spc board))
             . map shearBoard
             $ lattice

boardWithTiles :: Float -> Float -> Board -> Picture
boardWithTiles r s board = borderBoard r s <> tiles
  where
      tiles = pictures
            . map (\(i,e) -> drawTile (shearTile (bimap fromIntegral fromIntegral i)) e)
            . assocs
            $ board

drawTile :: (Float, Float) -> Tile -> Picture
drawTile _ Empty = blank
drawTile coords Black = color black (uncurry translate coords (circle (rad / 2)))
drawTile coords White = color white (uncurry translate coords (circle (rad / 2)))

shearBoard :: (Float, Float) -> (Float, Float)
shearBoard (x,y) = (a*x + b*y, c*y)
  where
      a = 11 * (2 * rad * cos (pi / 6) + spc)
      c = 11 * 15.866 -- TODO
      --c = 11 * (1.5 * rad + spc)
      b = 11 * (rad * cos (pi / 6) + spc / 2)

-- | square unit lattice corresponding to the boards we need to show
lattice :: [(Float, Float)]
--lattice = (,) <$> [-8..8] <*> [-4..3]
lattice = (,) <$> [-6..5] <*> [-4..3]

-- | board with the border
borderBoard :: Float -> Float -> Picture
borderBoard r s =
    color black
      (translate (-r * cos (pi / 6)) (-r * sin (pi / 6))
        (background r s)) <>
    color orange
      (boardWithSpace r s)

-- | a rhombus that goes behind the board to serve as the background
background :: Float -> Float -> Picture
background r s = polygon
               [ (0,0)
               , (10 * w + d, 0)
               , (10 * 1.5 * w  + d + r / sqrt 3, 10 * 15.866 + r) -- TODO: figure out where i got this number
               , (10 * 0.5 * w + r / sqrt 3, 10 * 15.866 + r)
               ]
  where w = 2 * r * cos (pi / 6) + s
        d = 5 / sqrt 3 + 5 * sqrt 3

-- | board with space in between
boardWithSpace :: Float -> Float -> Picture
boardWithSpace r s = pictures
                   . map (\(p,n) -> translate (n * a * cos (pi / 3)) (n * a * sin (pi / 3)) p)
                   . zip (repeat (rowWithSpace r s))
                   $ [0..10]
  where a = 2 * r * cos (pi / 6) + s

-- | several hexagons of given radius in a row with space in between
rowWithSpace :: Float -> Float -> Picture
rowWithSpace r s = pictures
                 . map (\(p,n) -> translate (n * (2 * r * cos (pi / 6) + s)) 0 p)
                 . zip (repeat (hexagon (r :+ 0)))
                 $ [0..10]

-- | draw a regular hexagon with the given length from the center to a vertex
hexagon :: Complex Float -> Picture
hexagon r = rotate 90
          . polygon
          . take 6
          . map ((realPart &&& imagPart) . (* r) . (\k -> exp (0 :+ 2 * pi * k / 6)))
          $ [0..]

{-
row :: Picture
row = pictures $ map (\(p,n) -> translate (n*18.3205) 0 p) $ zip (repeat (hexagon 10)) [0..10]

board :: Picture
board = pictures $ map (\(p,n) -> translate (n * 9.1603) (n * 15.866) p) $ zip (repeat row) [0..10]
-}

max_stars_repo_name: buonuomo/torus-hex

max_issues_repo_name: buonuomo/torus-hex

max_forks_repo_name: buonuomo/torus-hex


text:
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Graphics.Rendering.Plot.Render
-- Copyright   :  (c) A. V. H. McPhail 2010
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Rendering 'Figure's
--
-----------------------------------------------------------------------------

module Graphics.Rendering.Plot.Render (
                                       -- * Rendering
                                       render
                                       -- ** Access to 'FigureState'
                                       , newFigureState
                                       , updateFigureState
                                       , renderFigureState
                                       -- ** Outputting to file
                                       , OutputType(..)
                                       , writeFigure
                                       , writeFigureState
                                       -- * Notes
                                       -- $notes
                                       ) where

-----------------------------------------------------------------------------
{- TODO

    store 'next colour' list in state
-}
-----------------------------------------------------------------------------

--import Data.Either

--import Data.Packed.Vector
--import Numeric.LinearAlgebra.Linear

--import Data.Word

--import Data.Maybe

--import Data.Colour.SRGB
--import Data.Colour.Names

--import qualified Data.Array.IArray as A

import qualified Graphics.Rendering.Cairo as C
import qualified Graphics.Rendering.Pango as P

--import Control.Monad.Reader
--import Control.Monad.State
--import Control.Monad.Trans

import Graphics.Rendering.Plot.Types
import Graphics.Rendering.Plot.Defaults

--import Graphics.Rendering.Plot.Figure.Text

import Graphics.Rendering.Plot.Render.Types
import Graphics.Rendering.Plot.Render.Text
import Graphics.Rendering.Plot.Render.Plot

--import qualified Text.Printf as Printf

--import Prelude hiding(min,max)
--import qualified Prelude(max)

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

-- | render a 'Figure'
render :: Figure ()    -- ^ the figure to be rendered
       -> (Int,Int)    -- ^ (width,height)
       -> C.Render ()  -- ^ a Cairo operation
render g = (\(w,h) -> do
  pc <- pango $ P.cairoCreateContext Nothing
  to <- pango $ getDefaultTextOptions pc
  let options' = Options defaultLineOptions
                         defaultPointOptions defaultBarOptions to
  let (FigureState options _ figure) =
        execFigure g (FigureState options' defaultSupply emptyFigure)
  evalRender (renderFigure figure) (RenderEnv pc options)
             (BoundingBox 0 0 (fromIntegral w) (fromIntegral h)))

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

-- | create 'FigureState' from a series of 'Figure' actions
newFigureState :: Figure () -> IO FigureState
newFigureState f = do
  pc <- P.cairoCreateContext Nothing
  to <- getDefaultTextOptions pc
  let options' = Options defaultLineOptions
                         defaultPointOptions defaultBarOptions to
  return $ execFigure f (FigureState options' defaultSupply emptyFigure)

-- | modify a 'FigureState' with some new actions
updateFigureState :: FigureState -> Figure () -> FigureState
updateFigureState s f = execFigure f s

-- | render a 'FigureState'
renderFigureState :: FigureState  -- ^ the figure state
                -> (Int,Int)    -- ^ (width,height)
                -> C.Render ()  -- ^ a Cairo operation
renderFigureState (FigureState options _ figure) = (\(w,h) -> do
  pc <- pango $ P.cairoCreateContext Nothing
  evalRender (renderFigure figure) (RenderEnv pc options)
             (BoundingBox 0 0 (fromIntegral w) (fromIntegral h)))

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

-- | output the 'Figure'
writeFigure :: OutputType    -- ^ output file type
            -> FilePath      -- ^ file path
            -> (Int,Int)     -- ^ (width,height)
            -> Figure ()     -- ^ the 'Figure' rendering operation
            -> IO ()
writeFigure PNG fn wh f = withImageSurface wh (writeSurfaceToPNG fn (render f wh))
writeFigure PS  fn wh f = writeSurface C.withPSSurface fn wh f
writeFigure PDF fn wh f = writeSurface C.withPDFSurface fn wh f
writeFigure SVG fn wh f = writeSurface C.withSVGSurface fn wh f

withImageSurface :: (Int,Int) -> (C.Surface -> IO ()) -> IO ()
withImageSurface (w,h) = C.withImageSurface C.FormatARGB32 w h

writeSurfaceToPNG :: FilePath -> C.Render () -> C.Surface -> IO ()
writeSurfaceToPNG fn r s = do
  C.renderWith s r
  C.surfaceWriteToPNG s fn

writeSurface :: (FilePath -> Double -> Double -> (C.Surface -> IO ()) -> IO ())
            -> FilePath -> (Int,Int) -> Figure () -> IO ()
writeSurface rw fn (w,h) f = rw fn (fromIntegral w) (fromIntegral h)
                               (flip C.renderWith (render f (w,h)))

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

-- | output the 'FigureState'
writeFigureState :: OutputType    -- ^ output file type
            -> FilePath      -- ^ file path
            -> (Int,Int)     -- ^ (width,height)
            -> FigureState   -- ^ a FigureState
            -> IO ()
writeFigureState PNG fn wh f = withImageSurface wh
  (writeSurfaceToPNG fn (renderFigureState f wh))
writeFigureState PS  fn wh f = writeSurfaceFS C.withPSSurface fn wh f
writeFigureState PDF fn wh f = writeSurfaceFS C.withPDFSurface fn wh f
writeFigureState SVG fn wh f = writeSurfaceFS C.withSVGSurface fn wh f

writeSurfaceFS :: (FilePath -> Double -> Double -> (C.Surface -> IO ()) -> IO ())
            -> FilePath -> (Int,Int) -> FigureState -> IO ()
writeSurfaceFS rw fn (w,h) f = rw fn (fromIntegral w) (fromIntegral h)
                                (flip C.renderWith (renderFigureState f (w,h)))

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

renderFigure :: FigureData -> Render ()
renderFigure (Figure b p t s d) = do
  cairo $ do
    C.save
    setColour b
    C.paint
    C.restore

  applyPads p

  tx <- bbCentreWidth
  ty <- bbTopHeight
  (_,th) <- renderText t Centre TTop tx ty
  bbLowerTop (th+textPad)

  sx <- bbCentreWidth
  sy <- bbTopHeight
  (_,sh) <- renderText s Centre TTop sx sy
  bbLowerTop (sh+textPad)

  renderPlots d

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

max_stars_repo_name: JackTheEngineer/plot


text:
module Numeric.TBLIS.Types
  ( TblisError (..),
    TblisLenType,
    TblisStrideType,
    TblisLabelType,
    TblisComm,
    TblisConfig,
    TblisType (..),
    SomeTblisType (..),
    IsTblisType (..),
    TblisScalar (..),
    TblisTensor (..),
    tblisToTypedTensor,
    TblisTypedTensor,
    -- mkTblisTensor,
  )
where

-- import Control.Monad.Primitive
-- import Control.Monad.ST

-- import Data.Primitive.PrimArray
-- import Data.Primitive.Types (Prim)

import Control.Monad (when)
import DLPack (IsDLDataType)
import Data.Complex (Complex)
import qualified Data.Complex
import Data.Kind (Type)
import Data.Proxy
import Data.Text (Text)
import Foreign.C.Types
import Foreign.Marshal.Array (peekArray)
import Foreign.Marshal.Utils (fromBool)
import Foreign.Ptr
import Foreign.Storable
import qualified GHC.Exts as GHC
import GHC.Stack
import GHC.TypeLits
import System.IO.Unsafe (unsafePerformIO)
import Prelude hiding (pred)

natToInt :: forall n. KnownNat n => Int
natToInt = fromIntegral $ GHC.TypeLits.natVal (Proxy @n)
{-# INLINE natToInt #-}

type TblisLenType = CPtrdiff

type TblisStrideType = CPtrdiff

type TblisLabelType = CChar

data TblisComm

data TblisConfig

newtype TblisError = TblisError Text
  deriving stock (Show, Eq)

-- instance Error TblisError where
--   strMsg = TblisError . toText
--   {-# INLINE strMsg #-}

data TblisType :: (Type -> Type) where
  TblisFloat :: TblisType Float
  TblisDouble :: TblisType Double
  TblisComplexFloat :: TblisType (Complex Float)
  TblisComplexDouble :: TblisType (Complex Double)

deriving stock instance Show (TblisType a)

deriving stock instance Eq (TblisType a)

data SomeTblisType where
  SomeTblisType :: IsTblisType a => TblisType a -> SomeTblisType

-- typedef enum
-- {
--     TYPE_SINGLE   = 0,
--     TYPE_FLOAT    = TYPE_SINGLE,
--     TYPE_DOUBLE   = 1,
--     TYPE_SCOMPLEX = 2,
--     TYPE_DCOMPLEX = 3
-- } type_t;
instance Enum (TblisType a) where
  fromEnum x = case x of
    TblisFloat -> 0
    TblisDouble -> 1
    TblisComplexFloat -> 2
    TblisComplexDouble -> 3
  toEnum = error "toEnum not implemented for TblisType"

class (Num a, Storable a, IsDLDataType a) => IsTblisType a where
  tblisTypeOf :: proxy a -> TblisType a

instance IsTblisType Float where tblisTypeOf _ = TblisFloat

instance IsTblisType Double where tblisTypeOf _ = TblisDouble

instance IsTblisType (Complex Float) where tblisTypeOf _ = TblisComplexFloat

instance IsTblisType (Complex Double) where tblisTypeOf _ = TblisComplexDouble

newtype TblisScalar a = TblisScalar a
  deriving stock (Read, Show)
  deriving newtype (Eq, Ord, Num)

instance IsTblisType a => Storable (TblisScalar a) where
  sizeOf _ = 16 + 8
  {-# INLINE sizeOf #-}
  alignment _ = 8
  {-# INLINE alignment #-}
  peek p = case tblisTypeOf (Proxy :: Proxy a) of
    TblisFloat -> TblisScalar <$> peekByteOff p 0
    TblisDouble -> TblisScalar <$> peekByteOff p 0
    TblisComplexFloat -> TblisScalar <$> peekByteOff p 0
    TblisComplexDouble -> TblisScalar <$> peekByteOff p 0
  {-# INLINE peek #-}
  poke p (TblisScalar x) = do
    pokeByteOff p 0 x
    pokeByteOff p 16
      . (fromIntegral :: Int -> CInt)
      . fromEnum
      . tblisTypeOf
      $ (Proxy :: Proxy a)
  {-# INLINE poke #-}

data TblisTensor a = TblisTensor
  { tblisTensorConj :: !Bool,
    tblisTensorScalar :: {-# UNPACK #-} !(TblisScalar a),
    tblisTensorData :: {-# UNPACK #-} !(Ptr a),
    tblisTensorNDim :: {-# UNPACK #-} !Int,
    tblisTensorLen :: {-# UNPACK #-} !(Ptr TblisLenType),
    tblisTensorStride :: {-# UNPACK #-} !(Ptr TblisStrideType)
  }

-- type_t type;
-- int conj;
-- tblis_scalar scalar;
-- void* data;
-- unsigned ndim;
-- len_type* len;
-- stride_type* stride;
instance IsTblisType a => Storable (TblisTensor a) where
  sizeOf _ = 4 + 4 + 24 + 8 + 8 + 8 + 8
  alignment _ = 8
  poke p x = do
    pokeByteOff p 0 . (fromIntegral :: Int -> CInt) . fromEnum $ tblisTypeOf (Proxy @a)
    pokeByteOff p 4 . (fromBool :: Bool -> CInt) $ tblisTensorConj x
    pokeByteOff p 8 $ tblisTensorScalar x
    pokeByteOff p 32 $ tblisTensorData x
    pokeByteOff p 40 . (fromIntegral :: Int -> CUInt) $ tblisTensorNDim x
    pokeByteOff p 48 $ tblisTensorLen x
    pokeByteOff p 56 $ tblisTensorStride x
  peek _ = error $ "peek is not implemented for TblisTensor"

newtype TblisTypedTensor (rank :: Nat) (a :: Type) = TblisTypedTensor (TblisTensor a)

tblisToTypedTensor :: forall r a. (HasCallStack, KnownNat r) => TblisTensor a -> TblisTypedTensor r a
tblisToTypedTensor t
  | tblisTensorNDim t == natToInt @r = TblisTypedTensor t
  | otherwise =
    error $
      "tensor has wrong rank: " <> show (tblisTensorNDim t)
        <> "; expected "
        <> show (natToInt @r)

conjugateScalar :: forall a. IsTblisType a => TblisScalar a -> TblisScalar a
conjugateScalar (TblisScalar z) = TblisScalar $
  case tblisTypeOf (Proxy @a) of
    TblisFloat -> z
    TblisDouble -> z
    TblisComplexFloat -> Data.Complex.conjugate z
    TblisComplexDouble -> Data.Complex.conjugate z
{-# INLINE conjugateScalar #-}

tensorShape :: TblisTensor a -> [Int]
tensorShape t = unsafePerformIO $ do
  let helper i = let !i' = fromIntegral i in i'
  fmap helper <$> peekArray (tblisTensorNDim t) (tblisTensorLen t)

tensorIndex :: IsTblisType a => TblisTensor a -> [Int] -> IO a
tensorIndex t index = do
  stride <- fmap fromIntegral <$> peekArray (tblisTensorNDim t) (tblisTensorStride t)
  let shape = tensorShape t
  -- <- fmap fromIntegral <$> peekArray (tblisTensorNDim t) (tblisTensorLen t)
  when (any id $ zipWith (\i n -> i < 0 || i >= n) index shape) $
    error $ "invalid index: " <> show index
  let !linearIndex = sum $ zipWith (*) stride index
  !s <- (tblisTensorScalar t *) . TblisScalar <$> peekElemOff (tblisTensorData t) linearIndex
  let (TblisScalar !z) = if (tblisTensorConj t) then conjugateScalar s else s
  pure z
{-# NOINLINE tensorIndex #-}

flattenTensorToList :: IsTblisType a => TblisTensor a -> [a]
flattenTensorToList t = unsafePerformIO $ do
  -- <- fmap fromIntegral <$> peekArray (tblisTensorNDim t) (tblisTensorLen t)
  let !shape = tensorShape t
      indices = sequence . fmap (\n -> [0 .. n - 1]) $ shape
      f i = do !x <- tensorIndex t i; pure x
  mapM f indices

-- tensorToList1D :: forall a. Prim a => TblisTypedTensor 1 a -> [a]
-- tensorToList1D t = runST $
--   withForeignPtr' (tensorData t) $ \p ->
--     go p [] (stride * (extent - 1))
--   where
--     !stride = indexPrimArray (tensorStrides t) 0
--     !extent = indexPrimArray (tensorShape t) 0
--     go :: PrimMonad m => Ptr a -> [a] -> Int -> m [a]
--     go !p acc !i
--       | i >= 0 = do
--         !x <- P.readOffPtr p i
--         go p (x : acc) (i - stride)
--       | otherwise = touch t >> pure acc
--
-- tensorFromList1D :: forall a. (HasCallStack, Prim a) => Int -> [a] -> Tensor 'CPU 1 a
-- tensorFromList1D n xs = unsafePerformIO $ do
--   t <- newTensor [n]
--   withForeignPtr (tensorData t) $ \p -> do
--     let go !i []
--           | i == n = pure ()
--           | otherwise = error $ "list is shorter than expected"
--         go !i (y : ys)
--           | i < n = P.writeOffPtr p i y >> go (i + 1) ys
--           | otherwise = error "list is longer than expected"
--     go 0 xs
--   pure t
--
-- instance IsTblisType a => GHC.IsList (TblisTypedTensor 1 a) where
--   type Item (TblisTypedTensor 1 a) = a
--   toList = tensorToList1D
--   fromList xs = tensorFromList1D (length xs) xs
--   fromListN n xs = tensorFromList1D n xs

splitAt' :: HasCallStack => Int -> [a] -> ([a], [a])
splitAt' = go
  where
    go 0 [] = ([], [])
    go 1 (x : xs) = ([x], xs)
    go m (x : xs) = (x : xs', xs'')
      where
        (xs', xs'') = go (m - 1) xs
    go _ [] = error "wrong list length"

chunksOf :: Int -> [a] -> [[a]]
chunksOf n = go
  where
    go [] = []
    go xs@(_ : _) = let (ys, zs) = splitAt' n xs in ys : go zs

-- listShape2D :: [[a]] -> [Int]
-- listShape2D [] = [0, 0]
-- listShape2D xs@(x : _) = [length xs, length x]
-- {-# INLINE listShape2D #-}

-- listShape3D :: [[[a]]] -> [Int]
-- listShape3D [] = [0, 0, 0]
-- listShape3D xs@(x : _) = length xs : listShape2D x
--
-- listShape4D :: [[[a]]] -> [Int]
-- listShape4D [] = [0, 0, 0, 0]
-- listShape4D xs@(x : _) = length xs : listShape3D x
--
-- listShape5D :: [[[[a]]]] -> [Int]
-- listShape5D [] = [0, 0, 0, 0, 0]
-- listShape5D xs@(x : _) = length xs : listShape4D x

instance IsTblisType a => GHC.IsList (TblisTypedTensor 1 a) where
  type Item (TblisTypedTensor 1 a) = a
  toList (TblisTypedTensor t) = flattenTensorToList t
  fromList _ = error "IsList instance of TblisTypedTensor does not implement fromList"

instance IsTblisType a => GHC.IsList (TblisTypedTensor 2 a) where
  type Item (TblisTypedTensor 2 a) = [a]
  toList (TblisTypedTensor t) = case tensorShape t of
    [_, d₁] -> chunksOf d₁ (flattenTensorToList t)
    _ -> error "impossible: tensor must be rank 2"
  fromList _ = error "IsList instance of TblisTypedTensor does not implement fromList"

instance IsTblisType a => GHC.IsList (TblisTypedTensor 3 a) where
  type Item (TblisTypedTensor 3 a) = [[a]]
  toList (TblisTypedTensor t) = case tensorShape t of
    [_, d₁, d₂] -> chunksOf d₁ . chunksOf d₂ $ flattenTensorToList t
    _ -> error "impossible: tensor must be rank 3"
  fromList _ = error "IsList instance of TblisTypedTensor does not implement fromList"

-- peek p =
--   TblisTensor <$> peekByteOff p 4

-- mkShape :: (Prim a, Integral a) => (Int -> Bool) -> Int -> [Int] -> PrimArray a
-- mkShape pred n xs = runST $ do
--   arr <- newAlignedPinnedPrimArray n
--   let go [] !i
--         | i == n = return ()
--         | otherwise = error $ "list length is smaller than the specified size"
--       go (y : ys) !i
--         | i < n && pred y = writePrimArray arr i (fromIntegral y) >> go ys (i + 1)
--         | i >= n = error $ "list length is greated than the specified size"
--         | otherwise = error $ "encountered an invalid element: " <> show y
--   go xs 0
--   unsafeFreezePrimArray arr
--
-- mkLenAndStride :: [Int] -> [Int] -> (PrimArray TblisLenType, PrimArray TblisStrideType)
-- mkLenAndStride shape strides = (mkShape (>= 0) n shape, mkShape (> 0) n strides)
--   where
--     !n = length shape
--
-- mkTblisTensor :: IsTblisType a => Ptr a -> [Int] -> [Int] -> TblisTensor a
-- mkTblisTensor p shape strides =
--   let (len, stride) = mkLenAndStride shape strides
--    in TblisTensor False (TblisScalar 1) p len stride

max_stars_repo_name: twesterhout/tblis-hs

max_issues_repo_name: twesterhout/tblis-hs

max_forks_repo_name: twesterhout/tblis-hs