module Imports where

import           Control.Applicative
import           Control.Arrow
import           Control.Monad
import           Control.Monad.Cont
import           Control.Monad.Error
import           Control.Monad.Fix
import           Control.Monad.Identity
import           Control.Monad.Loops
import           Control.Monad.RWS
import           Control.Monad.Reader
import           Control.Monad.State
import           Control.Monad.Writer
import qualified Data.Aeson as A
import           Data.Array
import           Data.Bits
import           Data.Bool
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import           Data.Char
import           Data.Complex
import           Data.Dynamic
import           Data.Either
import           Data.Eq
import           Data.Fixed
import           Data.Function
import           Data.Graph
import           Data.Int
import           Data.Ix
import           Data.List
import           Data.List.Split
import qualified Data.Map as M
import           Data.Maybe
import           Data.Monoid
import           Data.Ord
import           Data.Ratio
import qualified Data.Set as S
import           Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TL
import           Data.Tree
import           Data.Tuple
import           Data.Typeable
import           Data.Word
import           Debug.SimpleReflect
import           Prelude hiding (IO,putStr,putStrLn,getLine,readLn,print,readIO,readFile,writeFile,appendFile)
import           Safe hiding (abort)
import           ShowFun
import           System.Random
import           Text.Printf

-- Convenience for golfing.
asS :: (String -> String) -> Text -> Text
asS f = T.pack . f . T.unpack

asS' = flip asS


module Statistics
    ( runJSONFold
    , receivedCreatedF
    , module Statistics.Block
    , module Statistics.Chart
    , module Statistics.CSV
    , module Statistics.Focus
    , module Statistics.Graph
    , module Statistics.MemPool
    , module Statistics.Report
    , module Statistics.Throughput
    , module Statistics.Tx
    ) where

import           Control.Foldl (Fold (..))
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as M

import           JSONLog
import           Statistics.Block
import           Statistics.Chart
import           Statistics.CSV
import           Statistics.Focus
import           Statistics.Graph
import           Statistics.MemPool
import           Statistics.Report
-- import Statistics.Relay
import           Statistics.Throughput
import           Statistics.Tx
import           Types
import           Universum
import           Util.Pipes (fold')

runJSONFold :: FilePath -> Fold IndexedJLTimedEvent a -> IO a
runJSONFold logDir fd = runParseLogs logDir $ fold' fd

receivedCreatedF :: Fold IndexedJLTimedEvent (Map TxHash (Maybe Timestamp))
receivedCreatedF = f <$> txFirstReceivedF <*> inBlockChainF
  where
    f :: Map TxHash Timestamp -> Map TxHash Timestamp -> Map TxHash (Maybe Timestamp)
    f rm cm = M.mapWithKey g rm
      where
        g :: TxHash -> Timestamp -> Maybe Timestamp
        g tx ts = maybe Nothing (\ts' -> Just $ ts' - ts) $ M.lookup tx cm


{-# LANGUAGE TypeFamilies, GADTs, EmptyDataDecls, RankNTypes, EmptyCase #-}
{-# LANGUAGE DataKinds, ConstraintKinds, PolyKinds #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}
{-# LANGUAGE TypeOperators, UndecidableInstances, ScopedTypeVariables #-}
{-# LANGUAGE FunctionalDependencies, DeriveFunctor, StandaloneDeriving #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

module Language.Grappa.GrappaInternals where

import Data.List
import Data.Typeable
import Data.Functor.Const
import Data.Proxy (Proxy(..))
import GHC.Exts (Constraint, IsList(..))

import Data.Aeson
import Data.Aeson.Types (Parser)
import Data.Vector (Vector)
import qualified Data.Vector as V
import qualified Numeric.LinearAlgebra as M

import Language.Grappa.Distribution

--
-- * Representing Grappa Types
--

-- | A type is "atomic" iff it is not an 'ADT'
type family IsAtomic a :: Bool where
  IsAtomic (ADT g) = 'False
  IsAtomic a = 'True

-- | The Grappa type of distributions. How it is represented depends on the
-- interpretation being used, so we do not instantiate it to an actual Haskell
-- type.
data Dist a deriving Typeable

-- | Is a type considered to be a base type in Grappa?
type family IsBaseType a where
  IsBaseType (a -> b) = 'False
  IsBaseType (ADT a) = 'False
  IsBaseType (Dist a) = 'False
  IsBaseType a = 'True

-- | An object-level representation of lists of Grappa types
data GrappaTypeListRepr (as :: [*]) where
  GrappaTypeListNil :: GrappaTypeListRepr '[]
  GrappaTypeListCons :: GrappaTypeRepr a -> GrappaTypeListRepr as ->
                        GrappaTypeListRepr (a ': as)

-- | An object-level representation of an application of a type constructor to 0
-- or more Grappa types
data GrappaTypeAppRepr (a :: k) where
  GrappaTypeAppBase :: Typeable a => GrappaTypeAppRepr a
  GrappaTypeAppApply :: GrappaTypeAppRepr (f :: * -> k) ->
                        GrappaTypeRepr a ->
                        GrappaTypeAppRepr (f a)

-- | An object-level representation of Grappa types
data GrappaTypeRepr (a :: *) where
  GrappaBaseType :: (IsAtomic a ~ 'True, IsBaseType a ~ 'True) =>
                    GrappaTypeAppRepr a -> GrappaTypeRepr a
  -- ^ A "base type" imported directly from Haskell, applied to 0 or more Grappa
  -- types; it must not be equal to any other Grappa type construct
  GrappaADTType :: GrappaADT adt => GrappaTypeAppRepr adt ->
                   GrappaTypeRepr (ADT adt)
  -- ^ A Grappa ADT type
  GrappaTupleType ::
    GrappaADT (TupleF as) =>
    GrappaTypeListRepr as -> GrappaTypeRepr (ADT (TupleF as))
  -- ^ A Grappa tuple type; note that we require 'GrappaADT', since the
  -- instance for 'TupleF' requires a class instance of 'GrappaTypeList' for
  -- @as@, not just a 'GrappaTypeListRepr'
  GrappaDistType :: GrappaTypeRepr a -> GrappaTypeRepr (Dist a)
  -- ^ A Grappa distribution type
  GrappaArrowType :: GrappaTypeRepr a -> GrappaTypeRepr b ->
                     GrappaTypeRepr (a -> b)
  -- ^ A Grappa function type

instance Show (GrappaTypeRepr a) where
  show _ = "(FIXME HERE: write a Show instance for GrappaTypeRepr!)"

-- | A typeclass indicating that a type is a valid Grappa type
class GrappaType a where
  grappaTypeRepr :: GrappaTypeRepr a

-- | A typeclass indicating that a list of types are all valid Grappa types. We
-- represent it as a type family as well as a type constraint, below, as the
-- type family helps GHC pull out all the 'GrappaType' instances for each elem
type family GrappaTypeListFam (as :: [*]) :: Constraint where
  GrappaTypeListFam '[] = ()
  -- To help GHC not have to unroll as many type applications, we put some
  -- helper cases to skip forward multiple types in a type list
  GrappaTypeListFam (a ': '[]) = GrappaType a
  GrappaTypeListFam (a ': b ': '[]) = (GrappaType a, GrappaType b)
  GrappaTypeListFam (a ': b ': c ': '[]) =
    (GrappaType a, GrappaType b, GrappaType c)
  GrappaTypeListFam (a ': b ': c ': d ': '[]) =
    (GrappaType a, GrappaType b, GrappaType c, GrappaType d)
  GrappaTypeListFam (a ': b ': c ': d ': e ': rest) =
    (GrappaType a, GrappaType b, GrappaType c, GrappaType d,
     GrappaType e, GrappaTypeList rest)

class (IsTypeList as, GrappaTypeListFam as) => GrappaTypeList as where
  grappaTypeListRepr :: GrappaTypeListRepr as

-- Instances for lists of types
instance GrappaTypeList '[] where
  grappaTypeListRepr = GrappaTypeListNil
instance GrappaType a => GrappaTypeList '[a] where
  grappaTypeListRepr = GrappaTypeListCons grappaTypeRepr GrappaTypeListNil
instance (GrappaType a, GrappaType b) => GrappaTypeList '[a,b] where
  grappaTypeListRepr =
    GrappaTypeListCons grappaTypeRepr
    (GrappaTypeListCons grappaTypeRepr GrappaTypeListNil)
instance (GrappaType a, GrappaType b, GrappaType c) =>
         GrappaTypeList '[a,b,c] where
  grappaTypeListRepr =
    GrappaTypeListCons grappaTypeRepr
    (GrappaTypeListCons grappaTypeRepr
     (GrappaTypeListCons grappaTypeRepr GrappaTypeListNil))
instance (GrappaType a, GrappaType b, GrappaType c, GrappaType d) =>
         GrappaTypeList '[a,b,c,d] where
  grappaTypeListRepr =
    GrappaTypeListCons grappaTypeRepr
    (GrappaTypeListCons grappaTypeRepr
     (GrappaTypeListCons grappaTypeRepr
      (GrappaTypeListCons grappaTypeRepr GrappaTypeListNil)))
instance (GrappaType a, GrappaType b, GrappaType c,
          GrappaType d, GrappaType e, GrappaTypeList rest) =>
         GrappaTypeList (a ': b ': c ': d ': e ': rest) where
  grappaTypeListRepr =
    GrappaTypeListCons grappaTypeRepr
    (GrappaTypeListCons grappaTypeRepr
     (GrappaTypeListCons grappaTypeRepr
      (GrappaTypeListCons grappaTypeRepr
       (GrappaTypeListCons grappaTypeRepr grappaTypeListRepr))))

-- Instances for representing the core Grappa type constructs
instance GrappaTypeList as => GrappaType (ADT (TupleF as)) where
  grappaTypeRepr = GrappaTupleType grappaTypeListRepr
instance GrappaType a => GrappaType (Dist a) where
  grappaTypeRepr = GrappaDistType grappaTypeRepr
instance (GrappaType a, GrappaType b) => GrappaType (a -> b) where
  grappaTypeRepr = GrappaArrowType grappaTypeRepr grappaTypeRepr
instance GrappaType a => GrappaType (Vector a) where
  grappaTypeRepr =
    GrappaBaseType $ GrappaTypeAppApply GrappaTypeAppBase grappaTypeRepr

-- Instances for representing base types; remember to add more instances here
-- when we want to support more Haskell types in Grappa!
instance GrappaType Bool where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase
instance GrappaType Int where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase
instance GrappaType R where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase
instance GrappaType Prob where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase

-- Instances for real and probability vectors and matrices
instance GrappaType RVector where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase
instance GrappaType RMatrix where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase
instance GrappaType ProbVector where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase
instance GrappaType ProbMatrix where
  grappaTypeRepr = GrappaBaseType GrappaTypeAppBase

-- | Test if a 'GrappaTypeRepr' is equal to 'R'
matchGrappaRType :: GrappaTypeRepr a -> Maybe (a :~: R)
matchGrappaRType (GrappaBaseType GrappaTypeAppBase) = eqT
matchGrappaRType _ = Nothing

-- | Test if a 'GrappaTypeRepr' is equal to 'Int'
matchGrappaIntType :: GrappaTypeRepr a -> Maybe (a :~: Int)
matchGrappaIntType (GrappaBaseType GrappaTypeAppBase) = eqT
matchGrappaIntType _ = Nothing

--
-- * Grappa ADTs
--

-- | Helper type: the identity
newtype Id a = Id {unId :: a} deriving (Functor, Show, Num, Floating, Fractional)

instance Applicative Id where
  pure = Id
  f <*> x = Id $ unId f $ unId x

-- FIXME: This approach to modeling ADTs does not let the type arguments change
-- in recursive uses of the ADT

-- | Builds an ADT type from an ADT type name
newtype ADT (adt :: (* -> *) -> * -> *) =
  ADT { unADT :: adt Id (ADT adt) }

-- | Class for mapping functions over ADTs
class TraversableADT (adt :: (* -> *) -> * -> *) where
  traverseADT :: Applicative m => (forall a. f a -> m (g a)) ->
                 adt f r -> m (adt g r)

-- | Map a function over an 'ADT'
mapADT :: TraversableADT adt => (forall a. f a -> g a) -> adt f r -> adt g r
mapADT f = unId . traverseADT (Id . f)

-- | Fold a function over all the constructor arguments of an ADT. Note that
-- this is /not/ the same as, say, a list fold, since here the "rest" argument
-- in a list is treated as a single argument, i.e., we do not recurse.
foldrADT :: TraversableADT adt =>
            (forall a. f a -> r) -> (r -> r -> r) -> r -> adt f (ADT adt) -> r
foldrADT resF f x = foldr f x . ctorArgsADT resF

-- | A function that can test if a given ADT uses a specific constructor
newtype CtorMatcher adt =
  CtorMatcher { applyCtorMatcher :: forall f. adt f (ADT adt) -> Bool }

-- | Typeclass of ADTs where we can enumerate all the constructors, by building
-- a list of ADT elements with 'Proxy's as arguments
class ReifyCtorsADT adt where
  reifyCtorsADT :: [adt Proxy (ADT adt)]

-- | Convert all the arguments of the constructor of an ADT into a list
ctorArgsADT :: TraversableADT adt => (forall a. f a -> r) ->
               adt f (ADT adt) -> [r]
ctorArgsADT resF = getConst . traverseADT (Const . (: []) . resF)

-- | Class for "valid" Grappa ADT types
class (TraversableADT adt, ReifyCtorsADT adt) =>
      GrappaADT (adt :: (* -> *) -> * -> *) where
  -- | Stronger form of 'TraversableADT', where we know that the argument types
  -- are themselves valid Grappa types
  gtraverseADT :: Applicative m =>
                  (forall a. GrappaType a => f a -> m (g a)) ->
                  adt f (ADT adt) -> m (adt g (ADT adt))

-- | Map a function on well-typed objects over an 'ADT'
gmapADT :: GrappaADT adt => (forall a. GrappaType a => f a -> g a) ->
           adt f (ADT adt) -> adt g (ADT adt)
gmapADT f = unId . gtraverseADT (Id . f)

-- | Helper type class for 'Show'ing 'ADT's
class ShowADT (adt :: (* -> *) -> * -> *) where
  showADT :: ADT adt -> String

instance ShowADT adt => Show (ADT adt) where
  show = showADT

-- NOTE: we don't want the generalized Show instance for an ADT, but usually
-- want adt-specific printing...
--
-- deriving instance (Show (adt f (ADT adt f))) => Show (ADT adt f)

--
-- * Heterogeneous Lists
--

-- | An @'HList' f [t0, t1, ..., tn]@ is a heterogeneous list of elements of
-- type @f t0@, @f t1@, etc.
data HList (f :: * -> *) (ts  :: [*]) where
  HListNil  :: HList f '[]
  HListCons :: f a -> HList f as -> HList f (a ': as)

-- | Convert a Grappa tuple to an 'HList'
tupleToHList :: TupleF ts f r -> HList f ts
tupleToHList Tuple0 = HListNil
tupleToHList tup@(Tuple1 _)           = consTupleToHList tup
tupleToHList tup@(Tuple2 _ _)         = consTupleToHList tup
tupleToHList tup@(Tuple3 _ _ _)       = consTupleToHList tup
tupleToHList tup@(Tuple4 _ _ _ _)     = consTupleToHList tup
tupleToHList tup@(TupleN _ _ _ _ _ _) = consTupleToHList tup

-- | Helper for 'tupleToHList'
consTupleToHList :: TupleF (t ': ts) f r -> HList f (t ': ts)
consTupleToHList tup =
  HListCons (tupleHead tup) (tupleToHList (tupleTail tup))

-- | Convert an 'HList' to a Grappa tuple
hListToTuple :: HList f ts -> TupleF ts f r
hListToTuple HListNil = Tuple0
hListToTuple (HListCons x xs) = tupleCons x (hListToTuple xs)

-- | Convert an 'HList' to a monomorphic list
hListToList :: (forall a. f a -> b) -> HList f ts -> [b]
hListToList _ HListNil = []
hListToList g (HListCons x xs) = g x : hListToList g xs

-- | Map a function over an 'HList'
mapHList :: (forall a. f a -> g a) -> HList f as -> HList g as
mapHList _ HListNil = HListNil
mapHList f (HListCons a as) = HListCons (f a) (mapHList f as)

-- | Traverse an 'HList'
traverseHList :: Applicative m => (forall a. f a -> m (g a)) ->
                 HList f ts -> m (HList g ts)
traverseHList _ HListNil = pure HListNil
traverseHList f (HListCons x xs) =
  pure HListCons <*> f x <*> traverseHList f xs

-- | Traverse an 'HList' in a monomorphic way
traverseHList' :: Applicative m => (forall a. f a -> m b) -> HList f ts -> m [b]
traverseHList' _ HListNil         = pure []
traverseHList' f (HListCons x xs) =
  pure (:) <*> f x <*> traverseHList' f xs

-- | Fold an 'HList' in a monomorphic way
foldHList :: (forall a. f a -> b -> b) -> HList f ts -> b -> b
foldHList _ HListNil acc         = acc
foldHList g (HListCons x xs) acc = g x $ foldHList g xs acc

-- | Compute the length of an 'HList'
hListLength :: HList f ts -> Int
hListLength l = foldHList (const (+ 1)) l 0

-- | Build an 'HList' of 'TypeListElem' proofs that each type in @as@ is an
-- element of @as@
buildHListElems :: HList f as -> HList (TypeListElem as) as
buildHListElems HListNil = HListNil
buildHListElems (HListCons _ rest) =
  HListCons TypeListElem_Base (mapHList TypeListElem_Cons
                               (buildHListElems rest))

{-
type family Concat (l1 :: [*]) (l2 :: [*]) where
  Concat '[] l = l
  Concat (x:xs) l = x : Concat xs l

hListConcat :: HList f as -> HList f bs -> HList f (Concat as bs)
hListConcat HListNil l = l
hListConcat (HListCons x xs) l = HListCons x (hListConcat xs l)
-}

--
-- * Tuple ADTs
--

-- | A generic tuple type as an ADT
data TupleF (ts :: [*]) f (r :: *) where
  Tuple0 :: TupleF '[] f r
  -- | Note: one-tuples are needed for recursive 'TupleN' constructions
  Tuple1 :: f a -> TupleF '[a] f r
  Tuple2 :: f a -> f b -> TupleF '[a, b] f r
  Tuple3 :: f a -> f b -> f c -> TupleF '[a, b, c] f r
  Tuple4 :: f a -> f b -> f c -> f d -> TupleF '[a, b, c, d] f r
  TupleN :: f a -> f b -> f c -> f d -> f e -> TupleF rest f r ->
            TupleF (a ': b ': c ': d ': e ': rest) f r

type family TupleConstraintF (cf :: * -> Constraint) (f :: * -> *)
     (ts :: [*]) :: Constraint where
  TupleConstraintF cf f '[] = ()
  TupleConstraintF cf f '[a] = (cf (f a))
  TupleConstraintF cf f '[a,b] = (cf (f a), cf (f b))
  TupleConstraintF cf f '[a,b,c] = (cf (f a), cf (f b), cf (f c))
  TupleConstraintF cf f '[a,b,c,d] = (cf (f a), cf (f b), cf (f c), cf (f d))
  TupleConstraintF cf f (a ': b ': c ': d ': e ': ts) =
    (cf (f a), cf (f b), cf (f c), cf (f d), cf (f e),
     TupleConstraintF cf f ts)

deriving instance TupleConstraintF Eq f ts => Eq (TupleF ts f r)

instance (IsTypeList ts, TupleConstraintF FromJSON f ts) =>
         FromJSON (TupleF ts f r) where
  parseJSON = withArray "Grappa tuple" (helper typeListProxy) where
    helper :: TupleConstraintF FromJSON f ts' => TupleF ts' Proxy r ->
              Vector Value -> Parser (TupleF ts' f r)
    helper Tuple0 _ = return Tuple0
    helper (Tuple1 _) vals = Tuple1 <$> parseJSON (vals V.! 0)
    helper (Tuple2 _ _) vals =
      Tuple2 <$> parseJSON (vals V.! 0) <*> parseJSON (vals V.! 1)
    helper (Tuple3 _ _ _) vals =
      Tuple3 <$> parseJSON (vals V.! 0) <*> parseJSON (vals V.! 1)
      <*> parseJSON (vals V.! 2)
    helper (Tuple4 _ _ _ _) vals =
      Tuple4 <$> parseJSON (vals V.! 0) <*> parseJSON (vals V.! 1)
      <*> parseJSON (vals V.! 2) <*> parseJSON (vals V.! 3)
    helper (TupleN _ _ _ _ _ tup) vals =
      TupleN <$> parseJSON (vals V.! 0) <*> parseJSON (vals V.! 1)
      <*> parseJSON (vals V.! 2) <*> parseJSON (vals V.! 3)
      <*> parseJSON (vals V.! 4) <*> helper tup (V.drop 5 vals)

-- | This says that @ts@ is a well-formed list of types, that we can reflect on
class IsTypeList ts where
  typeListProxy :: TupleF ts Proxy r

instance IsTypeList '[] where
  typeListProxy = Tuple0

instance IsTypeList '[a] where
  typeListProxy = Tuple1 Proxy

instance IsTypeList '[a, b] where
  typeListProxy = Tuple2 Proxy Proxy

instance IsTypeList '[a, b, c] where
  typeListProxy = Tuple3 Proxy Proxy Proxy

instance IsTypeList '[a, b, c, d] where
  typeListProxy = Tuple4 Proxy Proxy Proxy Proxy

instance IsTypeList rest => IsTypeList (a ': b ': c ': d ': e ': rest) where
  typeListProxy = TupleN Proxy Proxy Proxy Proxy Proxy typeListProxy

-- | "Proofs" that @t@ is in the list @ts@ of types
data TypeListElem ts t where
  TypeListElem_Base :: TypeListElem (t ': ts) t
  TypeListElem_Cons :: TypeListElem ts t -> TypeListElem (u ': ts) t

-- | Add an extra element to a tuple
tupleCons :: f t -> TupleF ts f r1 -> TupleF (t ': ts) f r2
tupleCons t Tuple0 = Tuple1 t
tupleCons t (Tuple1 a) = Tuple2 t a
tupleCons t (Tuple2 a b) = Tuple3 t a b
tupleCons t (Tuple3 a b c) = Tuple4 t a b c
tupleCons t (Tuple4 a b c d) = TupleN t a b c d Tuple0
tupleCons t (TupleN a b c d e rest) = TupleN t a b c d $ tupleCons e rest

-- | Get the first element of a tuple
tupleHead :: TupleF (t ': ts) f r -> f t
tupleHead (Tuple1 a) = a
tupleHead (Tuple2 a _) = a
tupleHead (Tuple3 a _ _) = a
tupleHead (Tuple4 a _ _ _) = a
tupleHead (TupleN a _ _ _ _ _) = a

-- | Remove the first element of a tuple
tupleTail :: TupleF (t ': ts) f r1 -> TupleF ts f r2
tupleTail (Tuple1 _) = Tuple0
tupleTail (Tuple2 _ b) = Tuple1 b
tupleTail (Tuple3 _ b c) = Tuple2 b c
tupleTail (Tuple4 _ b c d) = Tuple3 b c d
tupleTail (TupleN _ b c d e rest) =
  tupleCons b $ tupleCons c $ tupleCons d $ tupleCons e rest

-- | Project an element of a tuple
projectTuple :: TypeListElem ts t -> TupleF ts f r -> f t
projectTuple TypeListElem_Base tup = tupleHead tup
projectTuple (TypeListElem_Cons elemPf) tup =
  projectTuple elemPf $ tupleTail tup

-- | Project an element of an 'HList'
projectHList :: HList f as -> TypeListElem as a -> f a
projectHList (HListCons t _) TypeListElem_Base = t
projectHList (HListCons _ ts) (TypeListElem_Cons pf) = projectHList ts pf
projectHList HListNil tp_elem = case tp_elem of { }

-- Need a TraversableADT instance for each ADT type
instance TraversableADT (TupleF as) where
  traverseADT _ Tuple0 = pure Tuple0
  traverseADT f (Tuple1 a) = pure Tuple1 <*> f a
  traverseADT f (Tuple2 a b) = pure Tuple2 <*> f a <*> f b
  traverseADT f (Tuple3 a b c) = pure Tuple3 <*> f a <*> f b <*> f c
  traverseADT f (Tuple4 a b c d) = pure Tuple4 <*> f a <*> f b <*> f c <*> f d
  traverseADT f (TupleN a b c d e rest) =
    pure TupleN <*> f a <*> f b <*> f c <*> f d <*> f e <*> traverseADT f rest

untraverseTuple :: (Functor g, IsTypeList bs) =>
                   (forall a. g (f a) -> h a) ->
                   g (TupleF bs f r) -> TupleF bs h r
untraverseTuple = helper typeListProxy
  where
    helper :: Functor g => TupleF bs proxy r ->
              (forall a. g (f a) -> h a) -> g (TupleF bs f r) -> TupleF bs h r
    helper Tuple0 = \_ _ -> Tuple0
    helper tup@(Tuple1 _) = helperStep tup
    helper tup@(Tuple2 _ _) = helperStep tup
    helper tup@(Tuple3 _ _ _) = helperStep tup
    helper tup@(Tuple4 _ _ _ _) = helperStep tup
    helper tup@(TupleN _ _ _ _ _ _) = helperStep tup

    helperStep :: Functor g => TupleF (b ': bs) proxy r ->
                  (forall a. g (f a) -> h a) ->
                  g (TupleF (b ': bs) f r) -> TupleF (b ': bs) h r
    helperStep ts f gtup =
      tupleCons (f $ fmap tupleHead gtup)
      (helper (tupleTail ts) f $ fmap tupleTail gtup)

-- | Build a tuple from a polymorphic function for each element
buildTuple :: IsTypeList bs => (forall a. f a) -> TupleF bs f r
buildTuple f = untraverseTuple (\_ -> f) (Id typeListProxy)

-- Need a ReifyCtorsADT instance for each ADT type
instance IsTypeList ts => ReifyCtorsADT (TupleF ts) where
  reifyCtorsADT = [buildTuple Proxy]

-- | Map a binary function over tuples
mapTuple2 :: (forall a. f a -> g a -> h a) ->
             TupleF ts f r -> TupleF ts g r -> TupleF ts h r
mapTuple2 _ Tuple0 Tuple0 = Tuple0
mapTuple2 f (Tuple1 a1) (Tuple1 a2) = Tuple1 (f a1 a2)
mapTuple2 f (Tuple2 a1 b1) (Tuple2 a2 b2) = Tuple2 (f a1 a2) (f b1 b2)
mapTuple2 f (Tuple3 a1 b1 c1) (Tuple3 a2 b2 c2) =
  Tuple3 (f a1 a2) (f b1 b2) (f c1 c2)
mapTuple2 f (Tuple4 a1 b1 c1 d1) (Tuple4 a2 b2 c2 d2) =
  Tuple4 (f a1 a2) (f b1 b2) (f c1 c2) (f d1 d2)
mapTuple2 f (TupleN a1 b1 c1 d1 e1 rest1) (TupleN a2 b2 c2 d2 e2 rest2) =
  TupleN (f a1 a2) (f b1 b2) (f c1 c2) (f d1 d2) (f e1 e2) $
  mapTuple2 f rest1 rest2

-- Need a GrappaADT instance for each ADT type
instance GrappaTypeList as => GrappaADT (TupleF as) where
  gtraverseADT = helper where
    helper :: (GrappaTypeList bs, Applicative m) =>
              (forall a. GrappaType a => f a -> m (g a)) ->
              TupleF bs f r -> m (TupleF bs g r)
    helper _ Tuple0 = pure Tuple0
    helper f (Tuple1 a) = pure Tuple1 <*> f a
    helper f (Tuple2 a b) = pure Tuple2 <*> f a <*> f b
    helper f (Tuple3 a b c) = pure Tuple3 <*> f a <*> f b <*> f c
    helper f (Tuple4 a b c d) = pure Tuple4 <*> f a <*> f b <*> f c <*> f d
    helper f (TupleN a b c d e rest) =
      pure TupleN <*> f a <*> f b <*> f c <*> f d <*> f e <*> helper f rest

-- | Type synonym for Grappa tuples
type GTuple ts = ADT (TupleF ts)

-- | Map a type function over a list of types
type family MapF (f :: * -> *) (xs :: [*]) where
  MapF f '[] = '[]
  MapF f (x ': xs) = f x ': MapF f xs

-- | Defined type class for mappping a constraint function over a list; note
-- that we special-case small-sized lists, to help GHC do less unfolding
type family MapC (f :: * -> Constraint) (ts :: [*]) :: Constraint where
  MapC f '[] = ()
  MapC f '[a] = f a
  MapC f '[a,b] = (f a, f b)
  MapC f '[a,b,c] = (f a, f b, f c)
  MapC f '[a,b,c,d] = (f a, f b, f c, f d)
  MapC f (a ': b ': c ': d ': e ': rest) =
    (f a, f b, f c, f d, f e, MapC f rest)

instance MapC Show ts => ShowADT (TupleF ts) where
  showADT (ADT tup_body) = helper tup_body where
    helper :: forall ts' r. MapC Show ts' => TupleF ts' Id r -> String
    helper Tuple0 = "()"
    helper (Tuple1 (Id a)) = "(" ++ show a ++ ")"
    helper (Tuple2 (Id a) (Id b)) =
      "(" ++ show a ++ "," ++ show b ++ ")"
    helper (Tuple3 (Id a) (Id b) (Id c)) =
      "(" ++ show a ++ "," ++ show b ++ "," ++ show c ++ ")"
    helper (Tuple4 (Id a) (Id b) (Id c) (Id d)) =
      "(" ++ show a ++ "," ++ show b ++ "," ++ show c ++ "," ++ show d ++ ")"
    helper (TupleN (Id a) (Id b) (Id c) (Id d) (Id e) rest) =
      "(" ++ show a ++ "," ++ show b ++ "," ++ show c ++ "," ++ show d
      ++ "," ++ show e ++ "," ++ helper rest ++ ")"

--
-- * Boolean ADT
--

-- | The list type functor
data BoolF f r = TrueF | FalseF
  deriving (Eq, Show, Typeable)

-- Need a GrappaTypeRepr instance
instance GrappaType (ADT BoolF) where
  grappaTypeRepr = GrappaADTType GrappaTypeAppBase

-- Need a TraversableADT instance for each ADT type
instance TraversableADT BoolF where
  traverseADT _ TrueF  = pure TrueF
  traverseADT _ FalseF = pure FalseF

-- Need a ReifyCtorsADT instance for each ADT type
instance ReifyCtorsADT BoolF where
  reifyCtorsADT = [TrueF, FalseF]

-- Need a GrappaADT instance for each ADT type
instance GrappaADT BoolF where
  gtraverseADT _ TrueF  = pure TrueF
  gtraverseADT _ FalseF = pure FalseF

-- Also need CtorMatchers for each constructor
ctorMatcher__TrueF :: CtorMatcher BoolF
ctorMatcher__TrueF =
  CtorMatcher (\ x ->
               case x of
                 TrueF -> True
                 _ -> False)

ctorMatcher__FalseF :: CtorMatcher BoolF
ctorMatcher__FalseF =
  CtorMatcher (\ x ->
               case x of
                 FalseF -> True
                 _ -> False)

-- | Type synonym for lists, as Grappa ADTs
type GBool = ADT BoolF

fromHaskellBool :: Bool -> BoolF f r
fromHaskellBool True  = TrueF
fromHaskellBool False = FalseF

ifF :: BoolF f a -> t -> t -> t
ifF TrueF  t _ = t
ifF FalseF _ e = e

--
-- * List ADT
--

-- | The list type functor
data ListF a f r = Nil | Cons (f a) (f r)

-- Need a GrappaTypeRepr instance
instance GrappaType a => GrappaType (ADT (ListF a)) where
  grappaTypeRepr =
    GrappaADTType (GrappaTypeAppApply GrappaTypeAppBase grappaTypeRepr)

--deriving instance (Show (f a), Show (f r)) => Show (ListF a f r)

deriving instance (Eq (f a), Eq (f r)) => Eq (ListF a f r)

instance (FromJSON (f a), FromJSON (f r)) => FromJSON (ListF a f r) where
  parseJSON = withArray "Grappa list" $ \vals ->
    if V.null vals then return Nil else
      Cons <$> parseJSON (V.head vals) <*> parseJSON (Array $ V.tail vals)

-- Need a TraversableADT instance for each ADT type
instance TraversableADT (ListF a) where
  traverseADT _ Nil = pure Nil
  traverseADT f (Cons x xs) = pure Cons <*> f x <*> f xs

-- Need a ReifyCtorsADT instance for each ADT type
instance ReifyCtorsADT (ListF a) where
  reifyCtorsADT = [Nil, Cons Proxy Proxy]

-- Need a GrappaADT instance for each ADT type
instance GrappaType a => GrappaADT (ListF a) where
  gtraverseADT _ Nil = pure Nil
  gtraverseADT f (Cons x xs) = pure Cons <*> f x <*> f xs

-- Also need CtorMatchers for each constructor
ctorMatcher__Nil :: CtorMatcher (ListF a)
ctorMatcher__Nil =
  CtorMatcher (\ x ->
               case x of
                 Nil -> True
                 _ -> False)

ctorMatcher__Cons :: CtorMatcher (ListF a)
ctorMatcher__Cons =
  CtorMatcher (\ x ->
               case x of
                 Cons _ _ -> True
                 _ -> False)

-- | Type synonym for lists, as Grappa ADTs
type GList a = ADT (ListF a)

instance Show a => ShowADT (ListF a) where
  showADT = show . toList

listRepeatF :: a -> ADT (ListF a)
listRepeatF f = x
  where x = ADT (Cons (pure f) (pure x))

listIterateF :: a -> (a -> a) -> ADT (ListF a)
listIterateF x f = ADT (Cons (pure x) (pure (listIterateF (f x) f)))

enumFromByF :: (Num a) => a -> a -> ADT (ListF a)
enumFromByF from by = listIterateF from (+by)

enumFromToByF :: (Ord a, Num a) => a -> a -> a -> ADT (ListF a)
enumFromToByF from to by = go from
  where go n
          | n > to    = ADT Nil
          | otherwise = ADT (Cons (pure n) (pure (go (n + by))))

fromHaskellList :: [a] -> GList a
fromHaskellList = fromList

toHaskellList :: GList a -> [a]
toHaskellList = toList

-- | Convert a Grappa list with a non-'Id' functor to a Haskell list
toHaskellListF :: (f (ADT (ListF a)) -> ListF a f (ADT (ListF a))) ->
                  ListF a f (ADT (ListF a)) -> [f a]
toHaskellListF _ Nil = []
toHaskellListF unwrap (Cons x xs) = x : toHaskellListF unwrap (unwrap xs)

-- | Convert a Haskell list with a non-'Id' functor to a Grappa list
fromHaskellListF :: (ListF a f (ADT (ListF a)) -> f (ADT (ListF a))) ->
                    [f a] -> ListF a f (ADT (ListF a))
fromHaskellListF _ [] = Nil
fromHaskellListF wrap (x : xs) = Cons x (wrap $ fromHaskellListF wrap xs)

instance IsList (ADT (ListF a)) where
  type Item (ADT (ListF a)) = a
  fromList [] = ADT Nil
  fromList (x:xs) = ADT (Cons (Id x) (Id (fromList xs)))
  toList (ADT Nil) = []
  toList (ADT (Cons (Id x) (Id xs))) = x : toList xs

--
-- * Grappa-Specific Printing Methods
--

class GrappaShow t where
  grappaShow :: t -> String

instance GrappaShow Double where
  grappaShow = show

instance GrappaShow Prob where
  grappaShow = show

instance GrappaShow Int where
  grappaShow = show

instance GrappaShow a => GrappaShow (Id a) where
  grappaShow (Id x) = grappaShow x

instance GrappaShow (adt Id (ADT adt)) => GrappaShow (ADT adt) where
  grappaShow (ADT x) = grappaShow x

instance GrappaShow a => GrappaShow (Vector a) where
  grappaShow xs =
    "[|" ++ concat (intersperse ", " (map grappaShow $ V.toList xs)) ++ "|]"

instance GrappaShow RVector where
  grappaShow (RVector xs) =
    "[|" ++ concat (intersperse ", " (map grappaShow $ M.toList xs)) ++ "|]"

instance GrappaShow ProbVector where
  grappaShow (ProbVector xs) =
    "[|" ++ concat (intersperse ", "
                    (map (grappaShow . logRToProb) $ M.toList xs)) ++ "|]"

instance GrappaShow RMatrix where
  grappaShow m =
    "[R|" ++ concat (intersperse "\n" $ map grappaShow (toRowsM m)) ++ "|]"

instance GrappaShow ProbMatrix where
  grappaShow m =
    "[P|" ++ concat (intersperse "\n" $ map grappaShow (toRowsPM m)) ++ "|]"

instance MapC GrappaShow (MapF f ts) => GrappaShow (TupleF ts f r) where
  grappaShow Tuple0 = "()"
  grappaShow (Tuple1 a) = "(" ++ grappaShow a ++ ")"
  grappaShow (Tuple2 a b) = "(" ++ grappaShow a ++ "," ++ grappaShow b ++ ")"
  grappaShow (Tuple3 a b c) =
    "(" ++ grappaShow a ++ "," ++ grappaShow b ++ "," ++ grappaShow c ++ ")"
  grappaShow (Tuple4 a b c d) =
    "(" ++ grappaShow a ++ "," ++ grappaShow b ++ "," ++ grappaShow c ++
    "," ++ grappaShow d ++ ")"
  grappaShow (TupleN a b c d e _) =
    "(" ++ grappaShow a ++ "," ++ grappaShow b ++ "," ++ grappaShow c ++
    "," ++ grappaShow d ++ "," ++ grappaShow e ++ ", ...)"

instance (GrappaShow (f t), GrappaShowListContents (f r)) => GrappaShow (ListF t f r) where
  grappaShow Nil = "[]"
  grappaShow lst = "[" ++ go (showListContents lst) ++ "]"
    where go [] = ""
          go [x] = x
          go (x:xs) = x ++ "," ++ go xs

class GrappaShowListContents t where
  showListContents :: t -> [String]

instance (GrappaShow (f t), GrappaShowListContents (f r)) => GrappaShowListContents (ListF t f r) where
  showListContents Nil = []
  showListContents (Cons x xs) =
    grappaShow x : showListContents xs

instance GrappaShowListContents t => GrappaShowListContents (Id t) where
  showListContents (Id x) = showListContents x


{-| This module defines the 'Consume' typeclass, used for incrementally
    destructing inputs to random non-strict functions.

    Calling 'consume' on some value lazily returns an abstract type of 'Input',
    which contains all the entropy present in the original value. Paired with
    'Test.StrictCheck.Produce', these @Input@ values can be used to generate
    random non-strict functions, whose strictness behavior is dependent on the
    values given to them.
-}
module Test.StrictCheck.Consume
  ( -- * Incrementally consuming input
    Input
  , Inputs
  , Consume(..)
  -- * Manually writing 'Consume' instances
  , constructor
  , normalize
  , consumeTrivial
  , consumePrimitive
  -- * Generically deriving 'Consume' instances
  , GConsume
  , gConsume
  ) where

import Test.QuickCheck
import Generics.SOP
import Generics.SOP.NS

import Test.StrictCheck.Internal.Inputs

import Data.Complex

import Data.Foldable as Fold
import Data.List.NonEmpty (NonEmpty(..))
import Data.Tree     as Tree
import Data.Set      as Set
import Data.Map      as Map
import Data.Sequence as Seq
import Data.IntMap   as IntMap
import Data.IntSet   as IntSet

-- | Lazily monomorphize some input value, by converting it into an @Input@.
-- This is an incremental version of QuickCheck's @CoArbitrary@ typeclass.
-- It can also be seen as a generalization of the @NFData@ class.
--
-- Instances of @Consume@ can be derived automatically for any type implementing
-- the @Generic@ class from "GHC.Generics". Using the @DeriveAnyClass@
-- extension, we can say:
--
-- > import GHC.Generics as GHC
-- > import Generics.SOP as SOP
-- >
-- > data D x y
-- >   = A
-- >   | B (x, y)
-- >   deriving (GHC.Generic, SOP.Generic, Consume)
--
-- This automatic derivation follows these rules, which you can follow too if
-- you're manually writing an instance for some type which is not @Generic@:
--
-- For each distinct constructor, make a single call to 'constructor' with
-- a distinct @Int@, and a list of @Input@s, each created by recursively calling
-- 'consume' on every field in that constructor. For abstract types (e.g. sets),
-- the same procedure can be used upon an extracted list representation of the
-- contents.
class Consume a where
  -- | Convert an @a@ into an @Input@ by recursively destructing it using calls
  -- to @consume@
  consume :: a -> Input
  default consume :: GConsume a => a -> Input
  consume = gConsume

-- | Reassemble pieces of input into a larger Input: this is to be called on the
-- result of @consume@-ing subparts of input
constructor :: Int -> [Input] -> Input
constructor n !is =
  Input (Variant (variant n)) is

-- | Use the CoArbitrary instance for a type to consume it
--
-- This should only be used for "flat" types, i.e. those which contain no
-- interesting consumable substructure, as it's fully strict (non-incremental)
consumePrimitive :: CoArbitrary a => a -> Input
consumePrimitive !a =
  Input (Variant (coarbitrary a)) []

-- | Consume a type which has no observable structure whatsoever
--
-- This should only be used for types for which there is only one inhabitant, or
-- for which inhabitants cannot be distinguished at all.
consumeTrivial :: a -> Input
consumeTrivial !_ =
  Input mempty []

-- | Fully normalize something which can be consumed
normalize :: Consume a => a -> ()
normalize (consume -> input) = go input
  where
    go (Input _ is) = Fold.foldr seq () (fmap go is)

--------------------------------------------
-- Deriving Consume instances generically --
--------------------------------------------

-- | The constraints necessary to generically @consume@ something
type GConsume a = (Generic a, All2 Consume (Code a))

-- | Generic 'consume'
gConsume :: GConsume a => a -> Input
gConsume !(from -> sop) =
  constructor (index_SOP sop)
  . hcollapse
  . hcliftA (Proxy @Consume) (K . consume . unI)
  $ sop

---------------
-- Instances --
---------------

instance Consume (a -> b)  where consume = consumeTrivial
instance Consume (Proxy p) where consume = consumeTrivial

instance Consume Char     where consume = consumePrimitive
instance Consume Word     where consume = consumePrimitive
instance Consume Int      where consume = consumePrimitive
instance Consume Double   where consume = consumePrimitive
instance Consume Float    where consume = consumePrimitive
instance Consume Rational where consume = consumePrimitive
instance Consume Integer  where consume = consumePrimitive
instance (CoArbitrary a, RealFloat a) => Consume (Complex a) where
  consume = consumePrimitive

instance Consume ()
instance Consume Bool
instance Consume Ordering
instance Consume a => Consume (Maybe a)
instance (Consume a, Consume b) => Consume (Either a b)
instance Consume a => Consume [a]

instance Consume a => Consume (NonEmpty a) where
  consume (a :| as) = constructor 0 [consume a, consume as]

instance Consume a => Consume (Tree a) where
  consume (Node a as) = constructor 0 [consume a, consume as]

instance Consume v => Consume (Map k v) where
  consume = constructor 0 . fmap (consume . snd) . Map.toList

consumeContainer :: (Consume a, Foldable t) => t a -> Input
consumeContainer = constructor 0 . fmap consume . Fold.toList

instance Consume v => Consume (Seq v)    where consume = consumeContainer
instance Consume v => Consume (Set v)    where consume = consumeContainer
instance Consume v => Consume (IntMap v) where consume = consumeContainer
instance Consume IntSet where
  consume = consumeContainer . IntSet.toList

-- TODO: instances for the rest of Containers

instance (Consume a, Consume b) => Consume (a, b)
instance (Consume a, Consume b, Consume c) => Consume (a, b, c)
instance (Consume a, Consume b, Consume c, Consume d) => Consume (a, b, c, d)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e
         ) => Consume
  (a, b, c, d, e)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         ) => Consume
  (a, b, c, d, e, f)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g
         ) => Consume
  (a, b, c, d, e, f, g)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h
         ) => Consume
  (a, b, c, d, e, f, g, h)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i
         ) => Consume
  (a, b, c, d, e, f, g, h, i)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t, Consume u
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t, Consume u, Consume v
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t, Consume u, Consume v, Consume w
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t, Consume u, Consume v, Consume w, Consume x
          ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t, Consume u, Consume v, Consume w, Consume x
         , Consume y
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x, y)
instance ( Consume a, Consume b, Consume c, Consume d, Consume e, Consume f
         , Consume g, Consume h, Consume i, Consume j, Consume k, Consume l
         , Consume m, Consume n, Consume o, Consume p, Consume q, Consume r
         , Consume s, Consume t, Consume u, Consume v, Consume w, Consume x
         , Consume y, Consume z
         ) => Consume
  (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w, x, y, z)


-- |
-- Module      : Main
-- Description : Entry point module
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : jonatanhsundqvist@gmail.com
-- Stability   : experimental
-- Portability : POSIX

-- TODO | - JSON key bindings and config
--        - Sharing scores and recordings
--        - Buying premade assets, scores, sound fonts, etc. (user content)
--        - Audio!
--        - 3D audio (move listener around)
--        - Serialisation
--        - IO Lenses (?)
--        - Better way of dealing with state, sharing IO resources
--        - Restrict access to AppState, 'fenced off' IORefs

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Main where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Control.Monad      (forM_, when, forever, liftM)
import Control.Concurrent (threadDelay, forkIO)
import Control.Concurrent.MVar
import Text.Printf
import Data.IORef
import Data.Complex
import Data.StateVar
import Data.Functor
import qualified Data.Set  as S
import qualified Data.Map  as M

import qualified Graphics.Rendering.Cairo as Cairo
import           Graphics.UI.Gtk          as Gtk
import           Graphics.UI.Gtk          (AttrOp(..), on)

import Sound.OpenAL
-- import qualified Sound.ALUT   as Alut

-- Internal module imports
import           BattleHack.Types
import           BattleHack.Lenses
import qualified BattleHack.Render  as Render
import qualified BattleHack.Events  as Events
import qualified BattleHack.Audio   as Audio
import qualified BattleHack.Window  as Window

--------------------------------------------------------------------------------------------------------------------------------------------
-- Entry point
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
main :: IO ()
main = do
  -- Create and configure window
  (window, canvas) <- Window.create winsize

  -- Audio
  -- TODO: Utility function for accessing range indeces
  -- TODO: Don't hard-code range
  -- TODO: Check for errors
  -- Just (context, device) <- Audio.setup -- TODO: Return context as well (probably a good idea) (✓)
  -- [source] <- genObjectNames 1
  -- claviature'            <- Audio.makebuffersFromIndeces (zipWith const [0..] $ replicate 24 False)

  -- App state
  stateref <- newIORef (initalstate 24 origin' keysize') -- claviature')

  --
  -- mnotes <- newMVar . (-->piano.keys) <$> readIORef stateref
  mnotes <- readIORef stateref >>= (newMVar . (-->piano.keys))
  -- mnotes <- return [True, True, False, True] >>= newMVar -- (newMVar . (-->piano.keys))

  -- forkIO $ Audio.stream (1.0/30.0) source mnotes
  let audiofps = 5

  forkIO $ Audio.stream (1.0/audiofps) mnotes

  forkIO $ forever $ do
    putStrLn "New write iteration"
    playing <- (-->piano.keys) <$> readIORef stateref
    putMVar mnotes playing
    threadDelay . ceiling $ (1.0/audiofps) * 10^6
    printf "Next batch coming up!\n"
    printf "These notes: %s.\n\n\n\n" (show playing)

  -- Register event handlers
  Window.bindevents window canvas stateref

  -- Animation
  timeoutAdd (Events.onanimate canvas stateref) (1000 `div` fps)

  -- Enter main loop
  mainGUI
  where
    fps = 30                                   --
    origin'@(ox:+oy)  = 20:+20                 --
    keysize'@(sx:+sy) = (4:+13) * 40           --
    winsize           = (sx*7:+sy) + 2*origin' --

-- | Initial application state
-- TODO: Piano range
initalstate :: Int -> Vector -> Vector -> AppState
initalstate nkeys origin' keysize' = AppState { _piano = PianoSettings { _origin  = origin',
                                                                         _keysize = keysize',
                                                                         _indent  = 0.26,
                                                                         _mid     = 0.62,
                                                                         _active  = Nothing,
                                                                         _keys    = replicate nkeys False },

                                                _animation = AnimationData { _frame = 0,
                                                                             _fps   = 30 },

                                                _inputstate = InputState { _mouse=0:+0, _keyboard=S.empty },
                                                _bindings = M.fromList [("Escape", Cairo.liftIO mainQuit)] }

-- | Just a little hello world snippet to make sure everything is set up properly.
goodbyeWorld :: IO ()
goodbyeWorld = do
  putStrLn "Hello world!"
  putStrLn "Counting down to launch..."
  forM_ [10,9..0] ((>> threadDelay (10^6)) . print)
  putStrLn "Launching. What have I done!?"


-- Copyright 2016 TensorFlow authors.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{- | Rendering of TensorFlow operations as Haskell functions.

The basic type signature generated for each op is:

> {constraints} => {mandatory attrs} -> {input tensors} -> {output tensors}

where:

* @{mandatory attrs}@ is of the form @A_1 -> ... -> A_N@, where each @A@ is an
 op attribute that doesn't have a default and can't be inferred from other
 inputs.

* @{constraints}@ restrict the type parameters of the input and output tensors
 (for example: 'TensorType' or 'OneOf').

* @{input tensors}@ is of the form @T_1 -> ... -> T_N@, where each @T@ is of
the form @Tensor Ref a@ or @Tensor v a@ (or a list of one of those types),
and @a@ is either a concrete type or a (constrained) type variable.

* @{output tensors}@ is of the form @(T_1,...,T_N)@ for "pure" ops, and
@Build (T_1,...,T_N)@ for "stateful" ops.  An op is considered "stateful" if
it takes a @Tensor Ref@ or @Tensor v ResourceHandle@ as input, or if it's
explicitly marked \"Stateful\" in its @REGISTER_OP@ definition.  (If there
are no outputs, it is either @ControlNode@ or @Build ControlNode@.)
-}

module TensorFlow.OpGen
  ( OpGenFlags(..)
  , docOpList
  , flagParser)
  where

import Data.Foldable (toList)
import Data.Maybe (fromMaybe)
import Data.ProtoLens (def, showMessage)
import Data.List (sortOn)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NE
import Lens.Family2 ((^.), (.~), (&), view)
import Options.Applicative (Parser, help, long, strOption, value)
import Proto.Tensorflow.Core.Framework.OpDef
  ( OpList
  , OpDef
  , attr
  , inputArg
  , name
  , op
  , outputArg
  )
import Proto.Tensorflow.Core.Framework.Types (DataType(..))
import System.FilePath (takeBaseName)
import TensorFlow.OpGen.ParsedOp
import Text.PrettyPrint.Mainland
  ( Doc
  , (<>)
  , (<+>)
  , (</>)
  , (<+/>)
  , brackets
  , comma
  , commasep
  , dquotes
  , empty
  , enclose
  , flatten
  , folddoc
  , hang
  , indent
  , parens
  , sep
  , stack
  , strictText
  , tuple
  )
import qualified Data.Set as Set
import qualified Data.Text as Text

data OpGenFlags = OpGenFlags
     { outputFile :: String
     , prefix :: String
     , excludeList :: String
     }

flagParser :: Parser OpGenFlags
flagParser = OpGenFlags
     <$> strOption (mconcat [ long "output"
                            , help "File to write."
                            ])
     <*> strOption (mconcat [ long "prefix"
                            , help "Haskell package prefix to use"
                            ])
     <*> strOption (mconcat [ long "exclude_list"
                            , value ""
                            , help "Comma separated Ops names to ignore"
                            ])

docOpList :: OpGenFlags -> OpList -> Doc
docOpList flags opList =
  stack [ "{-# LANGUAGE ConstraintKinds #-}"
        , "{-# LANGUAGE DataKinds #-}"
        , "{-# LANGUAGE FlexibleContexts #-}"
        , "{-# LANGUAGE FlexibleInstances #-}"
        , "{-# LANGUAGE OverloadedStrings #-}"
        , "{-# LANGUAGE ScopedTypeVariables #-}"
          -- Avoids reports about shadowing standard library names.
        , "{-# OPTIONS_GHC -fno-warn-name-shadowing #-}"
          -- eqLengthGuard never returns false and dies instead.
        , "{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}"
        , "module" <+> strictText moduleName <+> "where"
        , empty
        , imports
        , empty
        , folddoc (\x y -> x </> empty </> y)
                  (map renderOpAndExtras $
                   sortOn (view name) $
                   filter (not . flip elem exclusions . view name) $
                   toList $ opList ^. op)
        ]
  where moduleName =
            Text.pack (prefix flags) <> "." <> camelCase
             -- Discards the optional trailing _ops_op_lib
            (fromMaybe shortName (Text.stripSuffix "_ops_op_lib" shortName))
        shortName = Text.pack (takeBaseName $ outputFile flags)
        exclusions = Text.splitOn "," $ Text.pack $ excludeList flags
        renderOpAndExtras o = renderOp (parseOp o) </> extras o

imports :: Doc
imports = stack [
      "import Data.ByteString (ByteString)"
    , "import Data.Complex (Complex)"
    , "import Data.Int (Int8, Int16, Int32, Int64)"
    , "import Data.Proxy (Proxy(Proxy))"
    , "import Data.Word (Word8, Word16, Word32, Word64)"
    , "import Lens.Family2 ((.~), (&))"
    , "import TensorFlow.Build"
    , "import TensorFlow.BuildOp"
    , "import TensorFlow.Tensor"
    , "import TensorFlow.Types"
    ]

renderHaskellName, renderTFName, renderQuotedTFName :: Name -> Doc
renderHaskellName = strictText . unHaskellName . haskellName
renderTFName = strictText . unTFName . tfName
renderQuotedTFName = dquotes . renderTFName

-- | Generate the source code for a single op.
-- For example:
--
-- -- | {haddock comment}
-- foo :: {type sig}
-- foo attr1 attr2 input1 input2 | eqLengthGuard [...] = {function body}
renderOp :: ParsedOp -> Doc
renderOp pOp = stack $
    [ haddocks
    -- Prevent unreasonably long compilation times on ghc-7.10, due
    -- to stack calling "-dump-hi" which (unnecessarily) includes the
    -- inlining information, and is large for ops with many arguments.
#if __GLASGOW_HASKELL__ < 800
    , "{-# NOINLINE" <+> n <+> "#-}"
#endif
    , n <+> "::" <+> hang 0 (typeSig empty pOp)
    , n <+> "=" <+> n <> "' id"
    , n' <+> "::" <+> hang 0 (typeSig "OpParams ->" pOp)
    , n' <+> hang 0 args <+> "|" <+> funcGuard listSizeAttrs
                <+> "=" </>  -- args are indented
                    -- the body needs to be indented wrt the name
                    indent indentation (functionBody pOp)
    ] ++ whereClause listSizeAttrs
  where
    n = renderHaskellName $ parsedOpName pOp
    n' = n <> "'"
    listSizeAttrs = inferredListSizeAttrs pOp
    args = sep $ "op'options"
               : (map renderHaskellName
                    $ map attrName (explicitInputAttrs pOp)
                    ++ map parsedArgName (parsedInputs pOp))
    haddocks = "-- |" <+> multilineComment (parsedOpSummary pOp) (parsedOpDescription pOp)

-- | A check that all lists of the given size have the given length.
-- For example:
--   eqLengthGuard [("N", [("input1", length input1), ("input2", length input2)])]
funcGuard :: [Attr (NonEmpty Name)] -> Doc
funcGuard attrs = "eqLengthGuard" <+> brackets (commasep entries)
      where
        entries =
            [ parens $ nAttr <> comma <+>
              brackets (commasep $ toList $
                            map renderTensorName (toList $ attrInfo a))
            | a <- attrs
            , let nAttr = renderQuotedTFName (attrName a)
            ]
        renderTensorName x = parens $ renderQuotedTFName x <> comma <+>
                        "length" <+> renderHaskellName x

-- | Define the implicit list length attributes.
-- For example:
--   where
--     n1 = fromIntegral (length input1) :: Int64
--     n2 = fromIntegral (length input2) :: Int64
whereClause :: [Attr (NonEmpty Name)] -> [Doc]
whereClause [] = []
whereClause as = [indent 2 $ "where" </> indent 2 (stack $ map defineLengthAttr as)]
  where
    defineLengthAttr a = renderHaskellAttrName a <+> "="
                            <+> "fromIntegral (length"
                            <+> renderHaskellName (NE.head $ attrInfo a)
                            <> ") :: Int64"

renderHaskellAttrName :: Attr a -> Doc
renderHaskellAttrName = renderHaskellName . attrName

functionBody :: ParsedOp -> Doc
functionBody pOp
    | parsedOpIsMonadic pOp
        = "build $ do"
            </> indent indentation (bindOpInputsVar
                        </> "buildOp" <+> outputListsSizes <+> opDef)
    | otherwise
        = "pureOp" <+> outputListsSizes <+> "$ do"
            </> indent indentation (bindOpInputsVar </> "return" <+> opDef)
  where
    outputListsSizes = brackets $ commasep
        [ renderHaskellName a
        | ParsedArg { parsedArgCase = ListArg { argLength = a } }
            <- parsedOutputs pOp
        ]
    opInputsVar = "op'inputs"
    bindOpInputsVar = opInputsVar <+> "<- fmap Prelude.concat $ Prelude.sequence"
                            <+> brackets (commasep $ map (\a -> "buildInputs" <+> a) tensorArgs)
    opDef = parens $ hang 0 $ stack $
        "opDef" <+> renderQuotedTFName (parsedOpName pOp) :
        -- Renders type parameter arguments.
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> inferredTypeExpr a
        | a <- inferredTypeAttrs pOp, let n = attrName a
        ] ++
        -- Renders mandatory attributes as function parameters.
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
        | a <- explicitInputAttrs pOp, let n = attrName a
        ] ++
        -- Renders sizes of tensor list types having number_attr.
        [ "& opAttr" <+> renderQuotedTFName n <+> ".~" <+> renderHaskellName n
        | a <- inferredListSizeAttrs pOp, let n = attrName a
        ] ++
        ["& op'options & opInputs .~" <+> opInputsVar]
    tensorArgs = renderTensorArg <$> parsedInputs pOp
    renderTensorArg = renderHaskellName . parsedArgName
    inferredTypeExpr a
        | typeParamIsList $ attrInfo a
            = "fromTensorTypes (Proxy :: Proxy" <+> renderHaskellAttrName a
                    <> ")"
        | otherwise = "tensorType (undefined ::" <+> renderHaskellAttrName a
                            <> ")"

-- | Write a comment with the inputs/outputs/attributes in proto format, for
-- debugging.
extras :: OpDef -> Doc
extras d = enclose "{-\n" "\n-}" $
            strictText $ Text.pack $
            showMessage ((def :: OpDef)
                        & inputArg .~ (d ^. inputArg)
                        & outputArg .~ (d ^. outputArg)
                        & attr .~ (d ^. attr))

-- | The type signature for an op.
-- Of the form:
-- forall t1 t2 v1 v2 . (TensorType t1, TensorType t2)
--      => {pre} Float -> Tensor t1 v1 -> Tensor t2 v2
-- where "Float" is an explicit input attribute, "Tensor t1 v1" is an input, and
-- "Tensor t2 v2" is an output.
typeSig :: Doc -> ParsedOp -> Doc
typeSig pre pOp = constraints
            <+/> pre </> signatureFold (map attrInput (explicitInputAttrs pOp)
                                ++ map tensorArgAndComment (parsedInputs pOp)
                                ++ [outputs])
  where
    constraints
        | null classConstraints = empty
        | otherwise = "forall" <+> sep typeParams <+> "." <+> tuple classConstraints <+> "=>"
    typeParams = [strictText v | k <- parsedInputs pOp ++ parsedOutputs pOp,
                  ArgSomeTensor v <- [argKind $ parsedArgCase k]]
                ++ [renderHaskellAttrName n | n <- inferredTypeAttrs pOp]
                ++ if parsedOpIsMonadic pOp then ["m'"] else []
    -- Use m' as the type parameter to avoid clashing with an attribute name.
    monadConstraint
        | parsedOpIsMonadic pOp = ["MonadBuild m'"]
        | otherwise = []
    classConstraints = monadConstraint ++ map tensorArgConstraint
                                                    (inferredTypeAttrs pOp)
    signatureFold = folddoc (\x y -> x </> "->" <+> y)
    attrInput a = renderAttrType (attrInfo a) <+> hang 0 ("-- ^" <+> attrComment a)
    renderAttrType (AttrSingle a) = renderAttrBaseType a
    renderAttrType (AttrList a) = brackets $ renderAttrBaseType a
    renderAttrBaseType = \case
        AttrBytes -> "ByteString"
        AttrInt64 -> "Data.Int.Int64"
        AttrFloat -> "Float"
        AttrBool -> "Bool"
        AttrType -> "DataType"
        AttrShape -> "Shape"
        AttrTensor -> "TensorProto"

    tensorArgAndComment t = tensorArg t <+> hang 0 ("-- ^" <+> argComment t)
    outputs = case parsedOutputs pOp of
        [] -> wrapOutput "ControlNode"
        -- TODO(judahjacobson): To improve indentation: `tensorArgAndComment a`
        [a] -> wrapOutput (tensorArg a) <+> "-- ^" <+> argComment a
        as -> wrapOutput (tuple (map tensorArg as)) <+/> resultComment as
    wrapOutput o
        | parsedOpIsMonadic pOp = "m'" <+> parens o
        | otherwise = o

-- | Render an op input or output.
-- For example: "Tensor Ref Int64", "Tensor v t"
tensorArg :: ParsedArg -> Doc
tensorArg p = case parsedArgCase p of
    SimpleArg { argType = t, argKind = k } -> tensorType t k
    ListArg { argType = t, argKind = k } -> brackets $ tensorType t k
    MixedListArg {argTypeAttr = t, argKind = k}
        -> "TensorList" <+> parens (kind k) <+> renderHaskellName t
  where
    kind k = case k of
                ArgTensorRef -> "Ref"
                ArgTensorValue -> "Value"
                ArgTensorBuild -> "Build"
                ArgSomeTensor v -> strictText v
    tensorType t k = let
        a = case t of
                ArgTypeFixed dt -> strictText $ dtTypeToHaskell dt
                ArgTypeAttr n -> renderHaskellName n
        in "Tensor" <+> kind k <+> a

attrComment :: Attr a -> Doc
attrComment a = argComment' (attrName a) (attrDescription a)

argComment :: ParsedArg -> Doc
argComment a = argComment' (parsedArgName a) (parsedArgDescription a)

argComment' :: Name -> Text.Text -> Doc
argComment' argName argDesc =
    bold (renderTFName argName) <> splitMultilineText (":" <+>) argDesc

bold :: Doc -> Doc
bold n = "__" <> n <> "__"

-- | Comment for the outputs of an op.
-- For example:
--   -- ^ (__output1__, __output2__)
--   --
--   -- * __output1__: description1
--   --
--   -- * __output2__: description2
resultComment :: [ParsedArg] -> Doc
resultComment os = stack $ flatten commentSummary : map commentDetails os
  where
    commentSummary = "-- ^" <+> tuple [bold (renderTFName $ parsedArgName o) | o <- os]
    commentDetails o =
        stack [ "--"
              , "-- *" <+> argComment o
              ]

-- | Constraints for a given type parameter.
-- E.g.: "TensorType t" or "OneOf [Int64, Float] t"
-- or "TensorTypes ts" or "OneOfs [..] ts".
tensorArgConstraint :: Attr TypeParam -> Doc
tensorArgConstraint a = case attrInfo a of
    TypeParam False Nothing -> "TensorType" <+> n
    TypeParam False (Just as) -> "OneOf" <+> typeList as <+> n
    TypeParam True Nothing -> "TensorTypes" <+> n
    TypeParam True (Just as) -> "OneOfs" <+> typeList as <+> n
  where
    n = renderHaskellAttrName a
    -- Produces a type-level list, e.g.: '[Int32,Int64,Float]
    typeList = ("'" <>) . brackets . commasep . map strictText .
                    Set.toList . Set.fromList .
                    map dtTypeToHaskell . toList

-- NOTE: The cases of this function should be kept in sync with
-- TensorFlow.Types.AllTensorTypes.
dtTypeToHaskell :: DataType -> Text.Text
dtTypeToHaskell DT_BOOL = "Bool"
dtTypeToHaskell DT_BFLOAT16 = "Data.Word.Word16"
dtTypeToHaskell DT_COMPLEX128 = "(Data.Complex.Complex Double)"
dtTypeToHaskell DT_COMPLEX64 = "(Data.Complex.Complex Float)"
dtTypeToHaskell DT_DOUBLE = "Double"
dtTypeToHaskell DT_FLOAT = "Float"
dtTypeToHaskell DT_INT16 = "Data.Int.Int16"
dtTypeToHaskell DT_INT32 = "Data.Int.Int32"
dtTypeToHaskell DT_INT64 = "Data.Int.Int64"
dtTypeToHaskell DT_INT8 = "Data.Int.Int8"
dtTypeToHaskell DT_QINT32 = "Data.Int.Int32"  -- TODO(gnezdo): make unique
dtTypeToHaskell DT_QINT8 = "Data.Word.Word8"  -- TODO(gnezdo): make unique
dtTypeToHaskell DT_QINT16 = "Data.Int.Int16"  -- TODO(gnezdo): make unique
dtTypeToHaskell DT_QUINT16 = "Data.Word.Word16"  -- TODO(gnezdo): make unique
dtTypeToHaskell DT_QUINT8 = "Data.Word.Word8"  -- TODO(gnezdo): make unique
dtTypeToHaskell DT_STRING = "Data.ByteString.ByteString"
dtTypeToHaskell DT_UINT16 = "Data.Word.Word16"
dtTypeToHaskell DT_UINT32 = "Data.Word.Word32"
dtTypeToHaskell DT_UINT64 = "Data.Word.Word64"
dtTypeToHaskell DT_HALF = "Data.Word.Word16"  -- TODO(gnezdo): make unique
dtTypeToHaskell DT_UINT8 = "Data.Word.Word8"
dtTypeToHaskell DT_RESOURCE = "ResourceHandle"
dtTypeToHaskell DT_VARIANT = "Variant"
dtTypeToHaskell x =
    Text.pack $ "Unsupported type in dtTypeToHaskell: " ++ show x

-- | haddockComment escapes TensorFlow doc strings into haddock.
-- TODO(gnezdo): deal with the markup.
haddockComment :: Text.Text -> Doc
haddockComment = strictText

-- | Generate a multiline comment.  For example:
--   summary'
--   --
--   -- detail_line1
--   -- detail_line2
--   -- ...
multilineComment :: Text.Text -> Text.Text -> Doc
multilineComment summary' detail =
    haddockComment summary' </>
    splitMultilineText insertParagraphAndComment detail
  where insertParagraphAndComment x = "--" </> "--" <+> x

-- | Converts the given multi-line detail string into
-- a multi-line haddock. Applies the given lead to the
-- first line. Returns an empty document for empty detail.
splitMultilineText :: (Doc -> Doc) -> Text.Text -> Doc
splitMultilineText lead detail =
  case Text.lines detail of
    [] -> empty
    (l : ls) -> stack $ lead (haddockComment l)
                      : map (("--" <+>) . haddockComment) ls

indentation :: Int
indentation = 4


module MyPhoto.Actions.Montage
    ( montage
    ) where

import           Control.Concurrent ( getNumCapabilities )
import           Control.Concurrent.Async ( mapConcurrently )
import           Control.Concurrent.MSem as MS
import           Control.Monad
import           GHC.Conc ( numCapabilities )
import           System.Console.GetOpt
import           System.IO.Temp
import           System.Directory
import           System.Exit
import           System.FilePath
import           System.Process
import           Graphics.Netpbm
import qualified Data.ByteString as B
import           Data.Maybe (maybe)
import qualified Statistics.Sample as S
import qualified Data.Vector as V (fromList)

import MyPhoto.Model
import MyPhoto.Utils

help :: String
help = let
    header = "Usage: montage OUTPUT_FILE"
  in usageInfo header options

data Options
  = Options
  { optOutputFile :: Maybe FilePath
  , optHelp       :: Bool
  } deriving Show

defaultOptions :: Options
defaultOptions
  = Options
  { optOutputFile = Nothing
  , optHelp       = False
  }

options :: [OptDescr (Options -> Options)]
options =
  [ Option ['o'] ["--output-file"]
      (OptArg (\ fp opts -> opts { optOutputFile = fp })
        "OUTPUT_FILE")
      "file to wriite the montage to (expected to be e.g. a PNG)"
  , Option ['h'] ["help"]
      (NoArg (\ opts -> opts { optHelp = True }))
      "print help"
  ]

getMyOpts :: [String] -> IO (Options, [String])
getMyOpts argv = case getOpt Permute options argv of
                   (o, n, []  ) -> return (foldl (flip id) defaultOptions o, n)
                   (_, _, errs) -> ioError (userError (concat errs ++ help))

montageImpl:: [String] -> [Img] -> PActionBody
montageImpl _    []            = return (Right [])
montageImpl args imgs@(img1:_) = do
  (opts, _) <- getMyOpts args

  if optHelp opts
  then return (Left help)
  else let
    (bn,ext) = splitExtensions img1
    outputFile = case optOutputFile opts of
          Just outputFile' -> outputFile'
          Nothing -> bn ++ "_" ++ show (length imgs) ++ "_MONTAGE.png"
    in do
      (_, _, _, pHandle) <- createProcess (proc "montage" (["-geometry", "100x100+2+2"] ++ imgs ++ [outputFile]))
      exitCode <- waitForProcess pHandle
      unless (exitCode == ExitSuccess) $
        fail ("Resize failed with " ++ show exitCode)
      return (Right imgs)

montage :: PrePAction
montage args = logSeparator "montage" <> PAction (montageImpl args)


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Numeric.LinearAlgebra hiding ((!))
import Numeric.LinearAlgebra.Devel
import Numeric.LinearAlgebra.SVD.SVDLIBC
import           Control.Monad
import           Control.Monad.Primitive
import           Control.Monad.IO.Class     (MonadIO(liftIO))
import           Control.Monad.Loops          (whileJust_)
import           Control.Monad.Trans          (lift)
import           Control.Monad.Trans.Resource
import           Control.Monad.Trans.State.Strict
import           Data.Conduit
import           Data.Conduit.Binary        (sourceFile)
import           Data.Conduit.List          (isolate)
import           Data.Conduit.Text   as CT
import           Data.Foldable              (forM_)
import           Data.HashMap.Strict        (HashMap)
import qualified Data.HashMap.Strict as HM
import           Data.IORef
import           Data.List                  (foldl')
import           Data.Maybe                 (fromJust,fromMaybe)
import           Data.IntMap                (IntMap)
import qualified Data.IntMap         as IM
import qualified Data.Set            as Set
import           Data.Text                  (Text)
import qualified Data.Text           as T
import qualified Data.Text.IO        as TIO
import qualified Data.Text.Lazy.IO   as TLIO
import           Data.Vector.Storable  (Vector, MVector(..),(!))
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV
import qualified Data.Vector.Unboxed as VU
import           Data.Word
import           Foreign.C.Types
import           Foreign.Ptr

foreign import ccall "mymain" c_mymain :: CInt -> Ptr CULong -> Ptr CULong  -> Ptr Double -> IO CInt

foreign import ccall "&callhaskell" p_callhaskell :: FunPtr (IO ())

foreign import ccall "dynamic" mkFunFromC :: FunPtr (IO ()) -> IO ()

foreign import ccall "wrapper" mkFunPtrFromHaskell :: IO () -> IO (FunPtr (IO ()))

foreign import ccall "registerfun" c_register :: FunPtr (IO ()) -> IO ()

count :: (MonadIO m) => Int -> Sink a m ()
count !n =
  (await >>=) $ mapM_ $ \_ -> do
    when (n `mod` 10000 == 0) $ liftIO $ print n
    count (n+1)

enumerate :: (Monad m) => Int -> Conduit a m (Int,a)
enumerate !n = do
  mx <- await
  forM_ mx $ \x -> do
    yield (n,x)
    enumerate (n+1)

buildMap :: (MonadIO m) => MVector (PrimState IO) Int -> MVector (PrimState IO) Int -> Sink (Int,Text) (StateT WordCountState m) ()
buildMap mcount md =
  whileJust_ await $ \(docid,t) -> do
    let ws = T.words t

    is <- lift $ mapM (voca mcount) ws

    let is' = Set.toList (Set.fromList is)
    lift $ doc md docid is'

voca mcount w = do
  (BoundedWordMap n m mrev,s) <- get -- readIORef mref
  case HM.lookup w m of
    Nothing -> do
      liftIO $ MV.write mcount n 1
      put $! (BoundedWordMap (n+1) (HM.insert w n m) (IM.insert n w mrev),s)
      -- writeIORef mref (BoundedWordMap (n+1) (HM.insert w n m) (IM.insert n w mrev))
      return n
    Just i -> liftIO $ MV.modify mcount (+1) i >> return i

doc :: MonadIO m => MVector (PrimState IO) Int -> Int -> [Int] -> StateT WordCountState m ()
doc md docid js = do
  (s,r) <- get
  put $! (s,IM.insert docid js r)
  liftIO $ mapM_ (MV.modify md (+1)) js

tfidf :: Int -> IntMap [Int] -> Vector Int -> Int -> Int -> Double
tfidf sz dmap occ_dwt t d = fromMaybe 0 calc
  where calc = do ws <- IM.lookup d dmap
                  guard (t `elem` ws)
                  return $ log (fromIntegral sz/fromIntegral (occ_dwt!t))

data BoundedWordMap = BoundedWordMap { newid :: !Int
                                     , wordmap :: !(HashMap Text Int)
                                     , revmap :: !(IntMap Text)
                                     }

type WordCountState = (BoundedWordMap,IntMap [Int])

emptyBWM = BoundedWordMap 0 HM.empty IM.empty

testcallback = putStrLn "I am in testcallback"

main = do
  myfunptr <- mkFunPtrFromHaskell testcallback
  c_register myfunptr

  let sz = 1000
  -- mref <- newIORef emptyBWM
  -- dref <- newIORef IM.empty
  mv <- V.thaw (V.replicate 1000000 (0 :: Int))
  md <- V.thaw (V.replicate 1000000 (0 :: Int))

  (_,s) <- runResourceT $ flip runStateT (emptyBWM,IM.empty) $
    sourceFile "1M.training" =$= CT.decode CT.utf8 =$= CT.lines =$= isolate sz =$= enumerate 0 $$ getZipSink $
      (,) <$> ZipSink (count 0) <*> ZipSink (buildMap mv md)
  --
  let (BoundedWordMap nword wmap rmap,docs) = s
  v <- V.freeze mv :: IO (Vector Int)
  occ_dwt <- V.freeze md :: IO (Vector Int)
  -- docs <- readIORef dref
  --

  mvvrows <- V.thaw (V.replicate 30000000 0)
  mvvcols <- V.thaw (V.replicate 30000000 0)
  mvvvals <- V.thaw (V.replicate 30000000 0)

  cref <- newIORef (0 :: Int)
  forM_ [0..sz-1] $ \d -> do
    let ts = fromJust (IM.lookup d docs)
    c <- readIORef cref
    forM_ (zip [c,c+1..] ts) $ \(i,t) -> do
      MV.write mvvrows i (fromIntegral t)
      MV.write mvvcols i (fromIntegral d)
      MV.write mvvvals i (tfidf sz docs occ_dwt t d)
    writeIORef cref (c+length ts)
  ssz <- readIORef cref
  vvrows <- V.take ssz <$> V.freeze mvvrows
  vvcols <- V.take ssz <$> V.freeze mvvcols
  vvvals <- V.take ssz <$> V.freeze mvvvals

  V.unsafeWith vvrows $ \pr ->
    V.unsafeWith vvcols $ \pc ->
      V.unsafeWith vvvals $ \pv -> do
        c_mymain (fromIntegral ssz) pr pc pv

  -- let csr  = CSR vvvals vvcols vvrows nword sz
  -- let (_,dd,_) = sparseSvd 100 csr
  -- print csr
  -- print dd
  {-
  let v3 :: VU.Vector (Int,Int,Double)
        =  VU.take 100 $ (VU.zip3 (V.convert vvrows) (V.convert vvcols) (V.convert vvvals))
  print v3
  -}
  freeHaskellFunPtr myfunptr

  return ()


{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor      #-}
{-# LANGUAGE FlexibleInstances  #-}
{-# LANGUAGE NoImplicitPrelude  #-}
{-# LANGUAGE TemplateHaskell    #-}

module Data.Configured.Types where

import           Control.Applicative         ((<*>))
import           Control.Concurrent.STM.TVar (TVar)
import           Control.Exception           (Exception, SomeException)
import           Control.Lens.TH             (makeLenses)
import           Data.Bool                   (Bool, bool, otherwise)
import qualified Data.ByteString             as BS
import qualified Data.ByteString.Lazy        as LBS
import           Data.Char                   (Char)
import           Data.Complex                (Complex)
import           Data.Data                   (Data)
import           Data.Eq                     (Eq, (==))
import           Data.Fixed                  (Fixed, HasResolution)
import           Data.Function               (($), (.))
import           Data.Functor                (Functor, fmap, (<$>))
import           Data.Hashable               (Hashable, hashWithSalt)
import           Data.HashMap.Strict         (HashMap)
import           Data.Int                    (Int, Int16, Int32, Int64, Int8)
import           Data.List                   (isSuffixOf)
import           Data.Maybe                  (Maybe (Just, Nothing))
import           Data.Ratio                  (Ratio, Rational, denominator,
                                              numerator, (%))
import           Data.Semigroup              ((<>))
import           Data.String                 (IsString, fromString)
import qualified Data.Text                   as T
import qualified Data.Text.Encoding          as T
import qualified Data.Text.Lazy              as LT
import qualified Data.Text.Lazy.Encoding     as LT
import           Data.Traversable            (traverse)
import           Data.Tuple                  (fst)
import           Data.Typeable               (Typeable)
import           Data.Word                   (Word, Word16, Word32, Word64,
                                              Word8)
import           Foreign.C.Types             (CDouble, CFloat)
import           Prelude                     (Double, Float, Integer, Integral,
                                              Num, RealFloat, String,
                                              fromInteger, fromRational)
import           System.IO                   (FilePath, IO)
import           Text.Show                   (Show, show)

type Name = T.Text

newtype Import = Import T.Text

data AutoConfig = AutoConfig { interval :: Int -- Time in seconds to check for updates
                             , onError  :: SomeException -> IO () -- Action on reload or change error.
                             } deriving (Typeable)
instance Show AutoConfig where
  show c = "AutoConfig {interval = " <> show (interval c) <> "}"

data Worth a = Required { worth :: a }
             | Optional { worth :: a }
             deriving (Functor, Typeable, Eq, Show)

data Value = Bool Bool | String T.Text | Number Rational | List [Value]
  deriving (Eq, Data, Show, Typeable)

data Pattern = Exact Name | Prefix Name
  deriving (Eq, Data, Show, Typeable)

instance IsString Pattern where
  fromString s | isSuffixOf ".*" s = Prefix . T.init . T.pack $ s
               | otherwise         = Exact $ T.pack s

instance Hashable Pattern where
    hashWithSalt salt (Exact n)  = hashWithSalt salt n
    hashWithSalt salt (Prefix n) = hashWithSalt salt n

type ChangeHandler = Name -> Maybe Value -> IO ()

data Config = Config
  { _group    :: Name
  , _paths    :: TVar [(Name, Worth FilePath)]
  , _conf     :: TVar (HashMap Name Value)
  , _handlers :: TVar (HashMap Pattern [ChangeHandler])
  }
makeLenses ''Config

type Settings = TVar SettingsData
data SettingsData = SettingsData
  { paths'    :: [Worth FilePath]
  , conf'     :: HashMap Name Value
  , handlers' :: [(Pattern, ChangeHandler)]
  }

instance (Hashable a) => Hashable (Worth a) where
  hashWithSalt s a = hashWithSalt s (worth a)

instance IsString (Worth FilePath) where
  fromString = Required

-- Exceptions

data ConfigError = ParseError FilePath String
  deriving (Show, Typeable)

instance Exception ConfigError

data KeyError = KeyError Name
  deriving (Show, Typeable)

instance Exception KeyError

-- Typeclasses

convertRational :: (Num a) => Rational -> Maybe a
convertRational r = bool Nothing (Just . fromInteger $ numerator r) $ denominator r == 1

instance Configured Value where
  convert = Just

class Configured a where
  convert :: Value -> Maybe a

instance Configured Bool where
  convert (Bool b) = Just b
  convert _        = Nothing

instance Configured Char where
  convert (String t) = fst <$> T.uncons t
  convert _          = Nothing

instance Configured Double where
  convert (Number r) = Just $ fromRational r
  convert _          = Nothing

instance Configured Float where
  convert (Number r) = Just $ fromRational r
  convert _          = Nothing

instance Configured Int where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Int8 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Int16 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Int32 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Int64 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Integer where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Word where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Word8 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Word16 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Word32 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured Word64 where
  convert (Number r) = convertRational r
  convert _          = Nothing

instance Configured BS.ByteString where
  convert = fmap T.encodeUtf8 . convert

instance Configured LBS.ByteString where
  convert = fmap LT.encodeUtf8 . convert

instance Configured T.Text where
  convert (String s) = Just s
  convert _          = Nothing

instance Configured CFloat where
  convert (Number r) = Just $ fromRational r
  convert _          = Nothing

instance Configured CDouble where
  convert (Number r) = Just $ fromRational r
  convert _          = Nothing

instance Configured LT.Text where
  convert = fmap LT.fromStrict . convert

instance Configured a => Configured [a] where
  convert (List l) = traverse convert l
  convert _        = Nothing

instance Integral a => Configured (Ratio a) where
  convert (Number r) = let n = numerator r
                           d = denominator r
                       in  Just $ (fromInteger n) % (fromInteger d)
  convert _          = Nothing

instance HasResolution a => Configured (Fixed a) where
  convert (Number r) = Just $ fromRational r
  convert _          = Nothing

instance RealFloat a => Configured (Complex a) where
  convert (Number r) = Just $ fromRational r
  convert _          = Nothing

instance (Configured a, Configured b) => Configured (a, b) where
  convert (List (a:b:_)) = (,) <$> convert a <*> convert b
  convert _              = Nothing

instance (Configured a, Configured b, Configured c) => Configured (a, b, c) where
  convert (List (a:b:c:_)) = (,,) <$> convert a <*> convert b <*> convert c
  convert _                = Nothing

instance (Configured a, Configured b, Configured c, Configured d) => Configured (a, b, c, d) where
  convert (List (a:b:c:d:_)) = (,,,) <$> convert a <*> convert b <*> convert c <*> convert d
  convert _                  = Nothing