{-# LANGUAGE CPP #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}

module Protolude (
  -- * Base functions
  module Base,
  identity,
  pass,
#if !MIN_VERSION_base(4,8,0)
  (&),
  scanl',
#endif
  -- * Function functions
  module Function,
  applyN,
  -- * List functions
  module List,
  map,
  uncons,
  unsnoc,
  -- * Data Structures
  module DataStructures,
  -- * Show functions
  module Show,
  show,
  print,
  -- * Bool functions
  module Bool,
  -- * Monad functions
  module Monad,
  liftIO1,
  liftIO2,
  -- * Functor functions
  module Functor,
  -- * Either functions
  module Either,
  -- * Applicative functions
  module Applicative,
  guarded,
  guardedA,
  -- * String conversion
  module ConvertText,
  -- * Debug functions
  module Debug,

  -- * Panic functions
  module Panic,
  -- * Exception functions
  module Exception,
  Protolude.throwIO,
  Protolude.throwTo,
  -- * Semiring functions
  module Semiring,

  -- * String functions
  module String,
  -- * Safe functions
  module Safe,
  -- * Eq functions
  module Eq,
  -- * Ord functions
  module Ord,
  -- * Traversable functions
  module Traversable,
  -- * Foldable functions
  module Foldable,
  -- * Semigroup functions
#if MIN_VERSION_base(4,9,0)
  module Semigroup,
#endif
  -- * Monoid functions
  module Monoid,
  -- * Bifunctor functions
  module Bifunctor,
  -- * Bifunctor functions
  module Hashable,

  -- * Deepseq functions
  module DeepSeq,

  -- * Tuple functions
  module Tuple,

  module Typeable,

#if MIN_VERSION_base(4,7,0)
  -- * Typelevel programming
  module Typelevel,
#endif

  -- * Monads
  module Fail,
  module State,
  module Reader,
  module Except,
  module Trans,
  module ST,
  module STM,

  -- * Integers
  module Int,
  module Bits,

  -- * Complex functions
  module Complex,

  -- * Char functions
  module Char,

  -- * Maybe functions
  module Maybe,

  -- * Generics functions
  module Generics,

  -- * ByteString functions
  module ByteString,
  LByteString,

  -- * Text functions
  module Text,
  LText,

  -- * Read functions
  module Read,
  readMaybe,
  readEither,

  -- * System functions
  module System,
  die,

  -- * Concurrency functions
  module Concurrency,

  -- * Foreign functions
  module Foreign,
) where

-- Protolude module exports.
import Protolude.Debug as Debug
import Protolude.List as List
import Protolude.Show as Show
import Protolude.Bool as Bool
import Protolude.Monad as Monad
import Protolude.Functor as Functor
import Protolude.Either as Either
import Protolude.Applicative as Applicative
import Protolude.ConvertText as ConvertText
import Protolude.Panic as Panic
import Protolude.Exceptions as Exception
import Protolude.Semiring as Semiring
import qualified Protolude.Conv as Conv

import Protolude.Base as Base hiding (
    putStr           -- Overriden by Show.putStr
  , putStrLn         -- Overriden by Show.putStrLn
  , print            -- Overriden by Protolude.print
  , show             -- Overriden by Protolude.show
  , showFloat        -- Custom Show instances deprecated.
  , showList         -- Custom Show instances deprecated.
  , showSigned       -- Custom Show instances deprecated.
  , showSignedFloat  -- Custom Show instances deprecated.
  , showsPrec        -- Custom Show instances deprecated.
  )
import qualified Protolude.Base as PBase

-- Used for 'show', not exported.
import Data.String (String)
import Data.String as String (IsString)

-- Maybe'ized version of partial functions
import Protolude.Safe as Safe (
    headMay
  , headDef
  , initMay
  , initDef
  , initSafe
  , tailMay
  , tailDef
  , tailSafe
  , lastDef
  , lastMay
  , foldr1May
  , foldl1May
  , foldl1May'
  , maximumMay
  , minimumMay
  , maximumDef
  , minimumDef
  , atMay
  , atDef
  )

-- Applicatives
import Control.Applicative as Applicative (
    Applicative(..)
  , Alternative(..)
  , Const(Const,getConst)
  , ZipList(ZipList,getZipList)
  , (<**>)
  , liftA
  , liftA2
  , liftA3
  , optional
  )

-- Base typeclasses
import Data.Eq as Eq (
    Eq(..)
  )
import Data.Ord as Ord (
    Ord(..)
  , Ordering(LT,EQ,GT)
  , Down(Down)
  , comparing
  )
import Data.Traversable as Traversable
import Data.Foldable as Foldable (
    Foldable,
    fold,
    foldMap,
    foldr,
    foldr',
    foldl,
    foldl',
    toList,
#if MIN_VERSION_base(4,8,0)
    null,
    length,
#endif
    elem,
    maximum,
    minimum,
    foldrM,
    foldlM,
    traverse_,
    for_,
    mapM_,
    forM_,
    sequence_,
    sequenceA_,
    asum,
    msum,
    concat,
    concatMap,
    and,
    or,
    any,
    all,
    maximumBy,
    minimumBy,
    notElem,
    find,
  )
import Data.Functor.Identity as Functor (
    Identity(Identity, runIdentity)
  )

#if MIN_VERSION_base(4,9,0)
import Data.List.NonEmpty as List (
    NonEmpty((:|))
  , nonEmpty
  )
import Data.Semigroup as Semigroup (
    Semigroup(sconcat, stimes)
  , WrappedMonoid
  , diff
  , cycle1
  , stimesMonoid
  , stimesIdempotent
  , stimesIdempotentMonoid
  , mtimesDefault
  )
#endif

#if !MIN_VERSION_base(4,15,0)
import Data.Semigroup as Semigroup (
  , Option(..)
  , option
  )
#endif

import Data.Monoid as Monoid

#if !MIN_VERSION_base(4,8,0)
import Protolude.Bifunctor as Bifunctor (Bifunctor(bimap, first, second))
#else
import Data.Bifunctor as Bifunctor (Bifunctor(bimap, first, second))
#endif

-- Deepseq
import Control.DeepSeq as DeepSeq (
    NFData(..)
  , ($!!)
  , deepseq
  , force
  )

-- Data structures
import Data.Tuple as Tuple (
    fst
  , snd
  , curry
  , uncurry
  , swap
  )

import Data.List as List (
    splitAt
  , break
  , intercalate
  , isPrefixOf
  , drop
  , filter
  , reverse
  , replicate
  , take
  , sortBy
  , sort
  , intersperse
  , transpose
  , subsequences
  , permutations
  , scanl
#if MIN_VERSION_base(4,8,0)
  , scanl'
#endif
  , scanr
  , iterate
  , repeat
  , cycle
  , unfoldr
  , takeWhile
  , dropWhile
  , group
  , inits
  , tails
  , zipWith
  , zip
  , unzip
  , genericLength
  , genericTake
  , genericDrop
  , genericSplitAt
  , genericReplicate
  )

#if !MIN_VERSION_base(4,8,0)
-- These imports are required for the scanl' rewrite rules
import GHC.Exts (build)
import Data.List (tail)
#endif

-- Hashing
import Data.Hashable as Hashable (
    Hashable
  , hash
  , hashWithSalt
  , hashUsing
  )

import Data.Map as DataStructures (Map)
import Data.Set as DataStructures (Set)
import Data.Sequence as DataStructures (Seq)
import Data.IntMap as DataStructures (IntMap)
import Data.IntSet as DataStructures (IntSet)

import Data.Typeable as Typeable (
    TypeRep
  , Typeable
  , typeOf
  , cast
  , gcast
#if MIN_VERSION_base(4,7,0)
  , typeRep
  , eqT
#endif
  )

#if MIN_VERSION_base(4,7,0)
import Data.Proxy as Typelevel (
    Proxy(..)
  )

import Data.Type.Coercion as Typelevel (
    Coercion(..)
  , coerceWith
  , repr
  )

import Data.Type.Equality as Typelevel (
    (:~:)(..)
  , type (==)
  , sym
  , trans
  , castWith
  , gcastWith
  )

#endif

#if MIN_VERSION_base(4,8,0)
import Data.Void as Typelevel (
    Void
  , absurd
  , vacuous
  )
#endif

import Control.Monad.Fail as Fail (
    MonadFail
  )

-- Monad transformers
import Control.Monad.State as State (
    MonadState
  , State
  , StateT(StateT)
  , put
  , get
  , gets
  , modify
  , state
  , withState

  , runState
  , execState
  , evalState

  , runStateT
  , execStateT
  , evalStateT
  )

import Control.Monad.Reader as Reader (
    MonadReader
  , Reader
  , ReaderT(ReaderT)
  , ask
  , asks
  , local
  , reader
  , runReader
  , runReaderT
  )

import Control.Monad.Trans.Except as Except (
    throwE
  , catchE
  )

import Control.Monad.Except as Except (
    MonadError
  , Except
  , ExceptT(ExceptT)
  , throwError
  , catchError
  , runExcept
  , runExceptT
  , mapExcept
  , mapExceptT
  , withExcept
  , withExceptT
  )

import Control.Monad.Trans as Trans (
    MonadIO
  , lift
  , liftIO
  )

-- Base types
import Data.Int as Int (
    Int
  , Int8
  , Int16
  , Int32
  , Int64
  )
import Data.Bits as Bits (
  Bits,
  (.&.),
  (.|.),
  xor,
  complement,
  shift,
  rotate,
#if MIN_VERSION_base(4,7,0)
  zeroBits,
#endif
  bit,
  setBit,
  clearBit,
  complementBit,
  testBit,
#if MIN_VERSION_base(4,7,0)
  bitSizeMaybe,
#endif
  bitSize,
  isSigned,
  shiftL,
  shiftR,
  rotateL,
  rotateR,
  popCount,
#if MIN_VERSION_base(4,7,0)
  FiniteBits,
  finiteBitSize,
  bitDefault,
  testBitDefault,
  popCountDefault,
#endif
#if MIN_VERSION_base(4,8,0)
  toIntegralSized,
  countLeadingZeros,
  countTrailingZeros,
#endif
  )
import Data.Word as Bits (
    Word
  , Word8
  , Word16
  , Word32
  , Word64
#if MIN_VERSION_base(4,7,0)
  , byteSwap16
  , byteSwap32
  , byteSwap64
#endif
  )

import Data.Either as Either (
    Either(Left,Right)
  , either
  , lefts
  , rights
  , partitionEithers
#if MIN_VERSION_base(4,7,0)
  , isLeft
  , isRight
#endif
  )

import Data.Complex as Complex (
    Complex((:+))
  , realPart
  , imagPart
  , mkPolar
  , cis
  , polar
  , magnitude
  , phase
  , conjugate
  )
import Data.Char as Char (
    Char
  , ord
  , chr
  , digitToInt
  , intToDigit
  , toUpper
  , toLower
  , toTitle
  , isAscii
  , isLetter
  , isDigit
  , isHexDigit
  , isPrint
  , isAlpha
  , isAlphaNum
  , isUpper
  , isLower
  , isSpace
  , isControl
  )
import Data.Bool as Bool (
  Bool(True, False),
  (&&),
  (||),
  not,
  otherwise
  )
import Data.Maybe as Maybe (
    Maybe(Nothing, Just)
  , maybe
  , isJust
  , isNothing
  , fromMaybe
  , listToMaybe
  , maybeToList
  , catMaybes
  , mapMaybe
  )

import Data.Function as Function (
    const
  , (.)
  , ($)
  , flip
  , fix
  , on
#if MIN_VERSION_base(4,8,0)
  , (&)
#endif
  )

-- Genericss
import GHC.Generics as Generics (
    Generic(..)
  , Generic1
  , Rep
  , K1(..)
  , M1(..)
  , U1(..)
  , V1
  , D1
  , C1
  , S1
  , (:+:)(..)
  , (:*:)(..)
  , (:.:)(..)
  , Rec0
  , Constructor(..)
  , Datatype(..)
  , Selector(..)
  , Fixity(..)
  , Associativity(..)
#if ( __GLASGOW_HASKELL__ >= 800 )
  , Meta(..)
  , FixityI(..)
  , URec
#endif
  )

-- ByteString
import qualified Data.ByteString.Lazy
import Data.ByteString as ByteString (ByteString)

-- Text
import Data.Text as Text (
    Text
  , lines
  , words
  , unlines
  , unwords
  )
import qualified Data.Text.Lazy

import Data.Text.IO as Text (
    getLine
  , getContents
  , interact
  , readFile
  , writeFile
  , appendFile
  )

import Data.Text.Lazy as Text (
    toStrict
  , fromStrict
  )

import Data.Text.Encoding as Text (
    encodeUtf8
  , decodeUtf8
  , decodeUtf8'
  , decodeUtf8With
  )

import Data.Text.Encoding.Error as Text (
    OnDecodeError
  , OnError
  , UnicodeException
  , lenientDecode
  , strictDecode
  , ignore
  , replace
  )

-- IO
import System.Environment as System (getArgs)
import qualified System.Exit
import System.Exit as System (
    ExitCode(..)
  , exitWith
  , exitFailure
  , exitSuccess
  )
import System.IO as System (
    Handle
  , FilePath
  , IOMode(..)
  , stdin
  , stdout
  , stderr
  , withFile
  , openFile
  )

-- ST
import Control.Monad.ST as ST (
    ST
  , runST
  , fixST
  )

-- Concurrency and Parallelism
import Control.Exception as Exception (
    Exception,
    toException,
    fromException,
#if MIN_VERSION_base(4,8,0)
    displayException,
#endif
    SomeException(SomeException)
  , IOException
  , ArithException(
    Overflow,
    Underflow,
    LossOfPrecision,
    DivideByZero,
    Denormal,
    RatioZeroDenominator
    )
  , ArrayException(IndexOutOfBounds, UndefinedElement)
  , AssertionFailed(AssertionFailed)
#if MIN_VERSION_base(4,7,0)
  , SomeAsyncException(SomeAsyncException)
  , asyncExceptionToException
  , asyncExceptionFromException
#endif
  , AsyncException(StackOverflow, HeapOverflow, ThreadKilled, UserInterrupt)
  , NonTermination(NonTermination)
  , NestedAtomically(NestedAtomically)
  , BlockedIndefinitelyOnMVar(BlockedIndefinitelyOnMVar)
  , BlockedIndefinitelyOnSTM(BlockedIndefinitelyOnSTM)
#if MIN_VERSION_base(4,8,0)
  , AllocationLimitExceeded(AllocationLimitExceeded)
#endif
#if MIN_VERSION_base(4,10,0)
  , CompactionFailed(CompactionFailed)
#endif
  , Deadlock(Deadlock)
  , NoMethodError(NoMethodError)
  , PatternMatchFail(PatternMatchFail)
  , RecConError(RecConError)
  , RecSelError(RecSelError)
  , RecUpdError(RecUpdError)
#if MIN_VERSION_base(4,9,0)
  , ErrorCall(ErrorCall, ErrorCallWithLocation)
#else
  , ErrorCall(ErrorCall)
#endif
#if MIN_VERSION_base(4,9,0)
  , TypeError(TypeError)
#endif
  , ioError
  , catch
  , catches
  , Handler(Handler)
  , catchJust
  , handle
  , handleJust
  , try
  , tryJust
  , evaluate
  , mapException
  , mask
  , mask_
  , uninterruptibleMask
  , uninterruptibleMask_
  , MaskingState(..)
  , getMaskingState
#if MIN_VERSION_base(4,9,0)
  , interruptible
#endif
  , allowInterrupt
  , bracket
  , bracket_
  , bracketOnError
  , finally
  , onException
  )
import qualified Control.Exception as PException

import Control.Monad.STM as STM (
    STM
  , atomically
#if !(MIN_VERSION_stm(2,5,0))
  , always
  , alwaysSucceeds
#endif
  , retry
  , orElse
  , check
  , throwSTM
  , catchSTM
  )

import Control.Concurrent.MVar as Concurrency (
    MVar
  , newEmptyMVar
  , newMVar
  , takeMVar
  , putMVar
  , readMVar
  , swapMVar
  , tryTakeMVar
  , tryPutMVar
  , isEmptyMVar
  , withMVar
#if MIN_VERSION_base(4,7,0)
  , withMVarMasked
#endif
  , modifyMVar_
  , modifyMVar
  , modifyMVarMasked_
  , modifyMVarMasked
#if MIN_VERSION_base(4,7,0)
  , tryReadMVar
  , mkWeakMVar
#endif
  , addMVarFinalizer
  )
import Control.Concurrent.Chan as Concurrency (
    Chan
  , newChan
  , writeChan
  , readChan
  , dupChan
  , getChanContents
  , writeList2Chan
  )
import Control.Concurrent.QSem as Concurrency (
    QSem
  , newQSem
  , waitQSem
  , signalQSem
  )
import Control.Concurrent.QSemN as Concurrency (
    QSemN
  , newQSemN
  , waitQSemN
  , signalQSemN
  )
import Control.Concurrent as Concurrency (
    ThreadId
  , forkIO
  , forkFinally
  , forkIOWithUnmask
  , killThread
  , forkOn
  , forkOnWithUnmask
  , getNumCapabilities
  , setNumCapabilities
  , threadCapability
  , yield
  , threadDelay
  , threadWaitRead
  , threadWaitWrite
#if MIN_VERSION_base(4,7,0)
  , threadWaitReadSTM
  , threadWaitWriteSTM
#endif
  , rtsSupportsBoundThreads
  , forkOS
#if MIN_VERSION_base(4,9,0)
  , forkOSWithUnmask
#endif
  , isCurrentThreadBound
  , runInBoundThread
  , runInUnboundThread
  , mkWeakThreadId
  , myThreadId
  )
import Control.Concurrent.Async as Concurrency (
    Async(..)
  , Concurrently(..)
  , async
  , asyncBound
  , asyncOn
  , withAsync
  , withAsyncBound
  , withAsyncOn
  , wait
  , poll
  , waitCatch
  , cancel
  , cancelWith
  , asyncThreadId
  , waitAny
  , waitAnyCatch
  , waitAnyCancel
  , waitAnyCatchCancel
  , waitEither
  , waitEitherCatch
  , waitEitherCancel
  , waitEitherCatchCancel
  , waitEither_
  , waitBoth
  , link
  , link2
  , race
  , race_
  , concurrently
  )

import Foreign.Ptr as Foreign (IntPtr, WordPtr)
import Foreign.Storable as Foreign (Storable)
import Foreign.StablePtr as Foreign (StablePtr)

-- Read instances hiding unsafe builtins (read)
import qualified Text.Read as Read
import Text.Read as Read (
    Read
  , reads
  )

-- Type synonymss for lazy texts
type LText = Data.Text.Lazy.Text
type LByteString = Data.ByteString.Lazy.ByteString

#if !MIN_VERSION_base(4,8,0)
infixl 1 &

(&) :: a -> (a -> b) -> b
x & f = f x
#endif

-- | The identity function, returns the give value unchanged.
identity :: a -> a
identity x = x

map :: Functor.Functor f => (a -> b) -> f a -> f b
map = Functor.fmap

uncons :: [a] -> Maybe (a, [a])
uncons [] = Nothing
uncons (x:xs) = Just (x, xs)

unsnoc :: [x] -> Maybe ([x],x)
unsnoc = Foldable.foldr go Nothing
  where
    go x mxs = Just (case mxs of
       Nothing -> ([], x)
       Just (xs, e) -> (x:xs, e))

-- | Apply a function n times to a given value
applyN :: Int -> (a -> a) -> a -> a
applyN n f = Foldable.foldr (.) identity (List.replicate n f)

-- | Parse a string using the 'Read' instance.
-- Succeeds if there is exactly one valid result.
--
-- >>> readMaybe ("123" :: Text) :: Maybe Int
-- Just 123
--
-- >>> readMaybe ("hello" :: Text) :: Maybe Int
-- Nothing
readMaybe :: (Read b, Conv.StringConv a String) => a -> Maybe b
readMaybe = Read.readMaybe . Conv.toS

-- | Parse a string using the 'Read' instance.
-- Succeeds if there is exactly one valid result.
-- A 'Left' value indicates a parse error.
--
-- >>> readEither "123" :: Either Text Int
-- Right 123
--
-- >>> readEither "hello" :: Either Text Int
-- Left "Prelude.read: no parse"
readEither :: (Read a, Conv.StringConv String e, Conv.StringConv e String) => e -> Either e a
readEither = first Conv.toS . Read.readEither . Conv.toS

-- | The print function outputs a value of any printable type to the standard
-- output device. Printable types are those that are instances of class Show;
-- print converts values to strings for output using the show operation and adds
-- a newline.
print :: (Trans.MonadIO m, PBase.Show a) => a -> m ()
print = liftIO . PBase.print

-- | Lifted throwIO
throwIO :: (Trans.MonadIO m, Exception e) => e -> m a
throwIO = liftIO . PException.throwIO

-- | Lifted throwTo
throwTo :: (Trans.MonadIO m, Exception e) => ThreadId -> e -> m ()
throwTo tid e = liftIO (PException.throwTo tid e)

-- | Do nothing returning unit inside applicative.
pass :: Applicative f => f ()
pass = pure ()

guarded :: (Alternative f) => (a -> Bool) -> a -> f a
guarded p x = Bool.bool empty (pure x) (p x)

guardedA :: (Functor.Functor f, Alternative t) => (a -> f Bool) -> a -> f (t a)
guardedA p x = Bool.bool empty (pure x) `Functor.fmap` p x

-- | Lift an 'IO' operation with 1 argument into another monad
liftIO1 :: MonadIO m => (a -> IO b) -> a -> m b
liftIO1 = (.) liftIO

-- | Lift an 'IO' operation with 2 arguments into another monad
liftIO2 :: MonadIO m => (a -> b -> IO c) -> a -> b -> m c
liftIO2 = ((.).(.)) liftIO

show :: (Show a, Conv.StringConv String b) => a -> b
show x = Conv.toS (PBase.show x)
{-# SPECIALIZE show :: Show  a => a -> Text  #-}
{-# SPECIALIZE show :: Show  a => a -> LText  #-}
{-# SPECIALIZE show :: Show  a => a -> String  #-}

#if MIN_VERSION_base(4,8,0)
-- | Terminate main process with failure
die :: Text -> IO a
die err = System.Exit.die (ConvertText.toS err)
#else
-- | Terminate main process with failure
die :: Text -> IO a
die err = hPutStrLn stderr err >> exitFailure
#endif

#if !MIN_VERSION_base(4,8,0)
-- This is a literal copy of the implementation in GHC.List in base-4.10.1.0.

-- | A strictly accumulating version of 'scanl'
{-# NOINLINE [1] scanl' #-}
scanl'           :: (b -> a -> b) -> b -> [a] -> [b]
scanl' = scanlGo'
  where
    scanlGo'           :: (b -> a -> b) -> b -> [a] -> [b]
    scanlGo' f !q ls    = q : (case ls of
                            []   -> []
                            x:xs -> scanlGo' f (f q x) xs)

{-# RULES
"scanl'"  [~1] forall f a bs . scanl' f a bs =
  build (\c n -> a `c` foldr (scanlFB' f c) (flipSeqScanl' n) bs a)
"scanlList'" [1] forall f a bs .
    foldr (scanlFB' f (:)) (flipSeqScanl' []) bs a = tail (scanl' f a bs)
 #-}

{-# INLINE [0] scanlFB' #-}
scanlFB' :: (b -> a -> b) -> (b -> c -> c) -> a -> (b -> c) -> b -> c
scanlFB' f c = \b g -> \x -> let !b' = f x b in b' `c` g b'

{-# INLINE [0] flipSeqScanl' #-}
flipSeqScanl' :: a -> b -> a
flipSeqScanl' a !_b = a
#endif


{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PackageImports #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DefaultSignatures #-}
module Data.Roles
  ( Representational(rep)
  , Phantom(phantom)
  ) where

import Control.Applicative
import Control.Monad.ST
import Data.Complex
import Data.Monoid
import Data.Proxy
import Data.Type.Coercion
import Data.Coerce
import Unsafe.Coerce
import Data.Map
import Data.IntMap

class Representational (t :: k1 -> k2) where
  -- | An argument is representational if you can lift a coercion of the argument into one of the whole
  rep :: Coercion a b -> Coercion (t a) (t b)
  default rep :: Phantom t => Coercion a b -> Coercion (t a) (t b)
  rep _ = phantom

class Representational t => Phantom (t :: k1 -> k2) where
  -- | An argument is phantom if you can 'coerce' the whole ignoring the argument
  phantom :: Coercion (t a) (t b)
  default phantom :: Coercible (t a) (t b) => Coercion (t a) (t b)
  phantom = Coercion

-- * Data.Proxy

instance Representational Proxy
instance Phantom Proxy

-- * Const

instance Representational Const where rep Coercion = Coercion
instance Representational (Const a)
instance Phantom (Const a)

-- * Data.Type.Coercion

instance Representational Coercion     where rep = unsafeCoerce
instance Representational (Coercion a) where rep Coercion = Coercion

-- * Prelude

instance Representational (->)       where rep Coercion = Coercion
instance Representational ((->) a)   where rep Coercion = Coercion

instance Representational Either     where rep Coercion = Coercion
instance Representational (Either a) where rep Coercion = Coercion

instance Representational (,)     where rep Coercion = Coercion
instance Representational ((,) a) where rep Coercion = Coercion

instance Representational (,,)       where rep Coercion = Coercion
instance Representational ((,,) a)   where rep Coercion = Coercion
instance Representational ((,,) a b) where rep Coercion = Coercion

instance Representational (,,,)         where rep Coercion = Coercion
instance Representational ((,,,) a)     where rep Coercion = Coercion
instance Representational ((,,,) a b)   where rep Coercion = Coercion
instance Representational ((,,,) a b c) where rep Coercion = Coercion

instance Representational (,,,,)           where rep Coercion = Coercion
instance Representational ((,,,,) a)       where rep Coercion = Coercion
instance Representational ((,,,,) a b)     where rep Coercion = Coercion
instance Representational ((,,,,) a b c)   where rep Coercion = Coercion
instance Representational ((,,,,) a b c d) where rep Coercion = Coercion

instance Representational (,,,,,)             where rep Coercion = Coercion
instance Representational ((,,,,,) a)         where rep Coercion = Coercion
instance Representational ((,,,,,) a b)       where rep Coercion = Coercion
instance Representational ((,,,,,) a b c)     where rep Coercion = Coercion
instance Representational ((,,,,,) a b c d)   where rep Coercion = Coercion
instance Representational ((,,,,,) a b c d e) where rep Coercion = Coercion

instance Representational []      where rep Coercion = Coercion
instance Representational Maybe   where rep Coercion = Coercion
instance Representational IO      where rep Coercion = Coercion
instance Representational (ST s)  where rep Coercion = Coercion

-- * containers
instance Representational (Map k) where rep Coercion = Coercion
instance Representational IntMap  where rep Coercion = Coercion

-- * Data.Complex

instance Representational Complex where rep Coercion = Coercion

-- * Data.Monoid

instance Representational Sum     where rep Coercion = Coercion
instance Representational Product where rep Coercion = Coercion
instance Representational Dual    where rep Coercion = Coercion
instance Representational Endo    where rep Coercion = Coercion
instance Representational First   where rep Coercion = Coercion
instance Representational Last    where rep Coercion = Coercion


{-# language InstanceSigs, DerivingStrategies #-}
{-# language PartialTypeSignatures #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}
{-| This module defines the 'Shaped' typeclass, which is used to generically
    manipulate values as fixed-points of higher-order functors in order to
    analyze their structure, e.g. while observing evaluation.

    If you just care about testing the strictness of functions over datatypes
    which are already instances of @Shaped@, you don't need to use this module.

    __Important note:__ To define new instances of 'Shaped' for types which
    implement 'GHC.Generic', __an empty instance will suffice__, as all the
    methods of 'Shaped' can be filled in by generic implementations. For
    example:

    > import GHC.Generics as GHC
    > import Generics.SOP as SOP
    >
    > data D = C deriving (GHC.Generic)
    >
    > instance SOP.Generic D
    > instance SOP.HasDatatypeInfo D
    >
    > instance Shaped D

    Using the @DeriveAnyClass@ extension, this can be shortened to one line:

    > data D = C deriving (GHC.Generic, SOP.Generic, SOP.HasDatatypeInfo, Shaped)

    Manual instances of 'Shaped' are necessary for types which do not or cannot
    implement GHC's @Generic@ typeclass, such as existential types, abstract
    types, and GADTs.

    This module is heavily based upon the approach in "Data.Functor.Foldable",
    which in turn is modeled after the paper "Functional Programming with
    Bananas, Lenses, Envelopes and Barbed Wire" (1991) by Erik Meijer, Maarten
    Fokkinga and Ross Paterson. If you don't yet understand recursion schemes
    and want to understand this module, it's probably a good idea to familiarize
    yourself with "Data.Functor.Foldable" before diving into this higher-order
    generalization.
-}
module Test.StrictCheck.Shaped
  ( Shaped(..)
  , module Test.StrictCheck.Shaped.Flattened
  -- * Fixed-points of 'Shape's
  , type (%)(..)
  -- * Folds and unfolds over fixed-points of @Shape@s
  , unwrap
  , interleave
  , (%)
  , fuse
  , translate
  , fold
  , unfold
  , unzipWith
  -- , reshape
  -- * Rendering 'Shaped' things as structured text
  , QName
  , Rendered(..)
  , RenderLevel(..)
  , renderfold
  -- * Tools for manually writing instances of 'Shaped'
  -- ** Implementing 'Shaped' for primitive types
  , Prim(..), unPrim
  , projectPrim
  , embedPrim
  , matchPrim
  , flatPrim
  , renderPrim
  , renderConstant
  -- ** Implementing 'Shaped' for container types
  , Containing(..)
  , projectContainer
  , embedContainer
  -- * Generic implementation of the methods of 'Shaped'
  , GShaped
  , GShape(..)
  , gProject
  , gEmbed
  , gMatch
  , gRender
  ) where

import Type.Reflection
import Data.Functor.Product
import Data.Bifunctor
import Data.Bifunctor.Flip
import Data.Coerce

import Generics.SOP hiding ( Shape )

import Data.Complex
-- import Data.List.NonEmpty (NonEmpty(..))

import Test.StrictCheck.Shaped.Flattened

-- TODO: provide instances for all of Base

-- | When a type @a@ is @Shaped@, we know how to convert it into a
-- representation parameterized by an arbitrary functor @f@, so that @Shape a f@
-- (the "shape of @a@ parameterized by @f@") is structurally identical to the
-- topmost structure of @a@, but with @f@ wrapped around any subfields of @a@.
--
-- Note that this is /not/ a recursive representation! The functor @f@ in
-- question wraps the original type of the field and /not/ a @Shape@ of that
-- field.
--
-- For instance, the @Shape@ of @Either a b@ might be:
--
-- > data EitherShape a b f
-- >   = LeftShape  (f a)
-- >   | RightShape (f b)
-- >
-- > instance Shaped (Either a b) where
-- >   type Shape (Either a b) = EitherShape a b
-- >   ...
--
-- The shape of a primitive type should be isomorphic to the primitive type,
-- with the functor parameter left unused.
class Typeable a => Shaped (a :: *) where
  -- | The @Shape@ of an @a@ is a type isomorphic to the outermost level of
  -- structure in an @a@, parameterized by the functor @f@, which is wrapped
  -- around any fields (of any type) in the original @a@.
  type Shape a :: (* -> *) -> *
  type Shape a = GShape a

  -- | Given a function to expand any @Shaped@ @x@ into an @f x@, expand an @a@
  -- into a @Shape a f@
  --
  -- That is: convert the top-most level of structure in the given @a@ into a
  -- @Shape@, calling the provided function on each field in the @a@ to produce
  -- the @f x@ necessary to fill that hole in the produced @Shape a f@.
  --
  -- Inverse to 'embed'.
  project :: (forall x. Shaped x => x -> f x) -> a -> Shape a f

  default project
    :: GShaped a
    => (forall x. Shaped x => x -> f x)
    -> a
    -> Shape a f
  project = gProject

  -- | Given a function to collapse any @f x@ into a @Shaped@ @x@, collapse a
  -- @Shape a f@ into merely an @a@
  --
  -- That is: eliminate the top-most @Shape@ by calling the provided function on
  -- each field in that @Shape a f@, and using the results to fill in the pieces
  -- necessary to build an @a@.
  --
  -- Inverse to 'project'.
  embed :: (forall x. Shaped x => f x -> x) -> Shape a f -> a

  default embed
    :: GShaped a
    => (forall x. Shaped x => f x -> x)
    -> Shape a f
    -> a
  embed = gEmbed

  -- | Given two @Shape@s of the same type @a@ but parameterized by potentially
  -- different functors @f@ and @g@, pattern-match on them to expose a uniform
  -- view on their fields (a 'Flattened' @(Shape a)@) to a continuation which
  -- may operate on those fields to produce some result
  --
  -- If the two supplied @Shape@s do not structurally match, only the fields of
  -- the first are given to the continuation. If they do match, the fields of
  -- the second are also given, along with type-level proof that the types of
  -- the two sets of fields align.
  --
  -- This very general operation subsumes equality testing, mapping, zipping,
  -- shrinking, and many other structural operations over @Shaped@ things.
  --
  -- It is somewhat difficult to manually write instances for this method, but
  -- consulting its generic implementation 'gMatch' may prove helpful.
  --
  -- See "Test.StrictCheck.Shaped.Flattened" for more information.
  match :: Shape a f -> Shape a g
        -> (forall xs. All Shaped xs
              => Flattened (Shape a) f xs
              -> Maybe (Flattened (Shape a) g xs)
              -> result)
        -> result

  default match :: GShaped a
        => Shape a f -> Shape a g
        -> (forall xs. All Shaped xs
              => Flattened (Shape a) f xs
              -> Maybe (Flattened (Shape a) g xs)
              -> result)
        -> result
  match = gMatch

  -- | Convert a @Shape a@ whose fields are some unknown constant type into a
  -- 'RenderLevel' filled with that type
  --
  -- This is a specialized pretty-printing mechanism which allows for displaying
  -- counterexamples in a structured format. See the documentation for
  -- 'RenderLevel'.
  render :: Shape a (K x) -> RenderLevel x

  default render :: (GShaped a, HasDatatypeInfo a)
          => Shape a (K x) -> RenderLevel x
  render = gRender

-- * Fixed-points of 'Shape's

-- | A value of type @f % a@ has the same structure as an @a@, but with the
-- structure of the functor @f@ interleaved at every field (including ones of
-- types other than @a@). Read this type aloud as "a interleaved with f's".
newtype (f :: * -> *) % (a :: *) :: * where
  Wrap :: f (Shape a ((%) f)) -> f % a

-- | Look inside a single level of an interleaved @f % a@. Inverse to the 'Wrap'
-- constructor.
unwrap :: f % a -> f (Shape a ((%) f))
unwrap (Wrap fs) = fs

-- * Folds and unfolds over fixed-points of @Shape@s

-- | Map a function across all the fields in a 'Shape'
--
-- This function may change the functor over which the @Shape@ is parameterized.
-- It can assume recursively that all the fields in the @Shape@ are themselves
-- instances of @Shaped@ (which they should be!). This means that you can nest
-- calls to @translate@ recursively.
translate :: forall a f g. Shaped a
          => (forall x. Shaped x => f x -> g x)
          -> Shape a f -> Shape a g
translate t d = match @a d d $ \flat _ ->
  unflatten $ mapFlattened @Shaped t flat

-- | The equivalent of a fold (catamorphism) over recursively 'Shaped' values
--
-- Given a function which folds an @f@ containing some @Shape x g@ into a @g x@,
-- recursively fold any interleaved @f % a@ into a @g a@.
fold :: forall a f g. (Functor f, Shaped a)
     => (forall x. Shaped x => f (Shape x g) -> g x)
     -> f % a -> g a
fold alg = alg . fmap (translate @a (fold alg)) . unwrap

-- | The equivalent of an unfold (anamorphism) over recursively 'Shaped' values
--
-- Given a function which unfolds an @f x@ into a @g@ containing some @Shape x
-- f@, corecursively unfold any @f a@ into an interleaved @g % a@.
unfold :: forall a f g. (Functor g, Shaped a)
       => (forall x. Shaped x => f x -> g (Shape x f))
       -> f a -> g % a
unfold coalg = Wrap . fmap (translate @a (unfold coalg)) . coalg

-- TODO: mapM, foldM, unfoldM, ...

-- | Fuse the interleaved @f@-structure out of a recursively interleaved @f %
-- a@, given some way of fusing a single level @f x -> x@.
--
-- This is a special case of 'fold'.
fuse
  :: (Functor f, Shaped a)
  => (forall x. f x -> x)
  -> (f % a -> a)
fuse e = e . fold (fmap (embed e))

-- | Interleave an @f@-structure at every recursive level of some @a@, given
-- some way of generating a single level of structure @x -> f x@.
--
-- This is a special case of 'unfold'.
interleave
  :: (Functor f, Shaped a)
  => (forall x. x -> f x)
  -> (a -> f % a)
interleave p = unfold (fmap (project p)) . p

-- | An infix synonym for 'interleave'
(%) :: forall a f. (Functor f, Shaped a)
    => (forall x. x -> f x)
    -> a -> f % a
(%) = interleave

-- | A higher-kinded @unzipWith@, operating over interleaved structures
--
-- Given a function splitting some @f x@ into a functor-product @Product g h x@,
-- recursively split an interleaved @f % a@ into two interleaved structures:
-- one built of @g@-shapes and one of @h@-shapes.
--
-- Note that @Product ((%) g) ((%) h) a@ is isomorphic to @(g % a, h % a)@; to
-- get the latter, pattern-match on the 'Pair' constructor of 'Product'.
unzipWith
  :: (All Functor [f, g, h], Shaped a)
  => (forall x. f x -> (g x, h x))
  -> (f % a -> (g % a, h % a))
unzipWith split =
  unPair . fold (crunch . pair . split)
  where
    crunch
      :: forall x g h.
      (Shaped x, Functor g, Functor h)
      => Product g h (Shape x (Product ((%) g) ((%) h)))
      -> Product ((%) g) ((%) h) x
    crunch =
      pair
      . bimap (Wrap . fmap (translate @x (fst . unPair)))
              (Wrap . fmap (translate @x (snd . unPair)))
      . unPair

    pair :: (l x, r x) -> Product l r x
    pair = uncurry Pair

    unPair :: Product l r x -> (l x, r x)
    unPair (Pair lx rx) = (lx, rx)

-- | TODO: document this strange function
{-
reshape :: forall b a f g. (Shaped a, Shaped b, Functor f)
        => (f (Shape b ((%) g)) -> g (Shape b ((%) g)))
        -> (forall x. Shaped x => f % x -> g % x)
        -> f % a -> g % a
reshape homo hetero d =
  case eqTypeRep (typeRep @a) (typeRep @b) of
    Nothing    -> hetero d
    Just HRefl ->
      Wrap
      $ homo . fmap (translate @a (reshape @b homo hetero))
      $ unwrap d
-}

----------------------------------
-- Rendering shapes for display --
----------------------------------

-- | Convert an @f % a@ into a structured pretty-printing representation,
-- suitable for further display/processing
renderfold
  :: forall a f. (Shaped a, Functor f)
  => f % a -> Rendered f
renderfold = unK . fold oneLevel
  where
    oneLevel :: forall x. Shaped x
             => f (Shape x (K (Rendered f)))
             -> K (Rendered f) x
    oneLevel = K . RWrap . fmap (render @x)

-- | A @QName@ is a qualified name
--
-- Note:
-- > type ModuleName   = String
-- > type DatatypeName = String
type QName = (ModuleName, DatatypeName, String)

-- | @RenderLevel@ is a functor whose outer shape contains all the information
-- about how to pretty-format the outermost @Shape@ of some value. We use
-- parametricity to make it difficult to construct incorrect 'render' methods,
-- by asking the user merely to produce a single @RenderLevel@ and stitching
-- nested @RenderLevel@s into complete 'Rendered' trees.
data RenderLevel x
  = ConstructorD QName [x]
  -- ^ A prefix constructor, and a list of its fields
  | InfixD QName Associativity Fixity x x
  -- ^ An infix constructor, its associativity and fixity, and its two fields
  | RecordD QName [(QName, x)]
  -- ^ A record constructor, and a list of its field names paired with fields
  | CustomD Fixity
    [Either (Either String (ModuleName, String)) (Fixity, x)]
  -- ^ A custom pretty-printing representation (i.e. for abstract types), which
  -- records a fixity and a list of tokens of three varieties: 1) raw strings,
  -- 2) qualified strings (from some module), or 3) actual fields, annotated
  -- with their fixity
  deriving (Eq, Ord, Show, Functor)

-- | @Rendered f@ is the fixed-point of @f@ composed with 'RenderLevel': it
-- alternates between @f@ shapes and @RenderLevel@s. Usually, @f@ will be the
-- identity functor 'I', but not always.
data Rendered f
  = RWrap (f (RenderLevel (Rendered f)))

----------------------------------------------------
-- Tools for manually writing instances of Shaped --
----------------------------------------------------

-- | The @Shape@ of a spine-strict container (i.e. a @Map@ or @Set@) is the same
-- as a container of demands on its elements. However, this does not have the
-- right /kind/ to be used as a @Shape@.
--
-- The @Containing@ newtype solves this problem. By defining the @Shape@ of some
-- container @(C a)@ to be @(C `Containing` a)@, you can use the methods
-- @projectContainer@ and @embedContainer@ to implement @project@ and @embed@
-- for your container type (although you will still need to manually define
-- @match@ and @render@).
newtype Containing h a f
  = Container (h (f a))
  deriving (Eq, Ord, Show)

-- | Generic implementation of @project@ for any container type whose @Shape@
-- is represented as a @Containing@ newtype
projectContainer :: (Functor c, Shaped a)
  => (forall x. Shaped x => x -> f x)
  -> c a -> Containing c a f
projectContainer p x = Container (fmap p x)

-- | Generic implementation of @embed@ for any container type whose @Shape@
-- is represented as a @Containing@ newtype
embedContainer :: (Functor c, Shaped a)
  => (forall x. Shaped x => f x -> x)
  -> Containing c a f -> c a
embedContainer e (Container x) = fmap e x

-- TODO: helper functions for matching and prettying containers

-- | The @Shape@ of a primitive type should be equivalent to the type itself.
-- However, this does not have the right /kind/ to be used as a @Shape@.
--
-- The @Prim@ newtype solves this problem. By defining the @Shape@ of some
-- primitive type @p@ to be @Prim p@, you can use the methods @projectPrim@,
-- @embedPrim@, @matchPrim@, and @prettyPrim@ to completely fill in the
-- definition of the @Shaped@ class for a primitive type.
--
-- __Note:__ It is only appropriate to use this @Shape@ representation when a
-- type really is primitive, in that it contains no interesting substructure.
-- If you use the @Prim@ representation inappropriately, StrictCheck will not be
-- able to inspect the richer structure of the type in question.
newtype Prim (x :: *) (f :: * -> *)
  = Prim x
  deriving (Eq, Ord, Show)
  deriving newtype (Num)

-- | Get the wrapped @x@ out of a @Prim x f@ (inverse to the @Prim@ constructor)
unPrim :: Prim x f -> x
unPrim (Prim x) = x

-- | Generic implementation of @project@ for any primitive type whose @Shape@ is
-- is represented as a @Prim@ newtype
projectPrim :: (forall x. Shaped x => x -> f x) -> a -> Prim a f
projectPrim _ = Prim

-- | Generic implementation of @embed@ for any primitive type whose @Shape@ is
-- is represented as a @Prim@ newtype
embedPrim :: (forall x. Shaped x => f x -> x) -> Prim a f -> a
embedPrim _ = unPrim

-- | Generic implementation of @match@ for any primitive type whose @Shape@ is
-- is represented as a @Prim@ newtype with an underlying @Eq@ instance
matchPrim :: Eq a => Prim a f -> Prim a g
           -> (forall xs. All Shaped xs
                => Flattened (Prim a) f xs
                -> Maybe (Flattened (Prim a) g xs)
                -> result)
           -> result
matchPrim (Prim a) (Prim b) k =
  k (flatPrim a)
     (if a == b then (Just (flatPrim b)) else Nothing)

-- | Helper for writing @match@ instances for primitive types which don't have
-- @Eq@ instance
--
-- This generates a @Flattened@ appropriate for using in the implementation of
-- @match@. For more documentation on how to use this, see the documentation of
-- 'match'.
flatPrim :: a -> Flattened (Prim a) g '[]
flatPrim x = Flattened (const (Prim x)) Nil

-- | Generic implementation of @render@ for any primitive type whose @Shape@ is
-- is represented as a @Prim@ newtype
renderPrim :: Show a => Prim a (K x) -> RenderLevel x
renderPrim (Prim a) = renderConstant (show a)

-- | Given some @string@, generate a custom pretty-printing representation which
-- just shows the string
renderConstant :: String -> RenderLevel x
renderConstant s = CustomD 11 [Left (Left s)]

-- TODO: What about demands for abstract types with > 1 type of unbounded-count field?

{-
withFieldsContainer ::
  forall c a f result.
     (forall r h.
        c (h a) ->
        (forall x. Shaped x
           => [h x]
           -> (forall g. [g x] -> c (g a))
           -> r)
        -> r)
  -> Containing c a f
  -> (forall xs. All Shaped xs
        => NP f xs
        -> (forall g. NP g xs -> Containing c a g)
        -> result)
  -> result
withFieldsContainer viaContaining (Container c) cont =
  viaContaining c $
    \list un ->
       withNP @Shaped list (Container . un) cont

-- TODO: Make this work for any number of lists of fields, by carefully using
-- unsafeCoerce to deal with unknown list lengths

withFieldsViaList ::
  forall demand f result.
     (forall r h.
        demand h ->
        (forall x. Shaped x
           => [h x]
           -> (forall g. [g x] -> demand g)
           -> r)
        -> r)
  -> demand f
  -> (forall xs. All Shaped xs
        => NP f xs
        -> (forall g. NP g xs -> demand g)
        -> result)
  -> result
withFieldsViaList viaList demand cont =
  viaList demand $
    \list un ->
       withNP @Shaped list un cont

withNP :: forall c demand result f x. c x
       => [f x]
       -> (forall g. [g x] -> demand g)
       -> (forall xs. All c xs
             => NP f xs -> (forall g. NP g xs -> demand g) -> result)
       -> result
withNP list unList cont =
  withUnhomogenized @c list $ \np ->
    cont np (unList . homogenize)

withConcatenated :: NP (NP f) xss -> (forall xs. NP f xs -> r) -> r
withConcatenated pop cont =
  case pop of
    Nil         -> cont Nil
    (xs :* xss) -> withConcatenated xss (withPrepended xs cont)
  where
    withPrepended ::
      NP f ys -> (forall zs. NP f zs -> r)
              -> (forall zs. NP f zs -> r)
    withPrepended pre k rest =
      case pre of
        Nil        -> k rest
        (x :* xs)  -> withPrepended xs (k . (x :*)) rest

homogenize :: All ((~) a) as => NP f as -> [f a]
homogenize      Nil  = []
homogenize (a :* as) = a : homogenize as

withUnhomogenized :: forall c a f r.
  c a => [f a] -> (forall as. (All c as, All ((~) a) as) => NP f as -> r) -> r
withUnhomogenized      []  k = k Nil
withUnhomogenized (a : as) k =
  withUnhomogenized @c as $ \np -> k (a :* np)
-}

---------------------------------------------------
-- Generic implementation of the Shaped methods --
---------------------------------------------------

-- | The 'Shape' used for generic implementations of 'Shaped'
--
-- This wraps a sum-of-products representation from "Generics.SOP".
newtype GShape a f
  = GS (NS (NP f) (Code a))

-- | The collection of constraints necessary for a type to be given a generic
-- implementation of the 'Shaped' methods
type GShaped a =
  ( Generic a
  , Shape a ~ GShape a
  , All2 Shaped (Code a)
  , SListI (Code a)
  , All SListI (Code a) )

-- | Generic 'project'
gProject :: GShaped a
         => (forall x. Shaped x => x -> f x)
         -> a -> Shape a f
gProject p !(from -> sop) =
  GS (unSOP (hcliftA (Proxy @Shaped) (p . unI) sop))

-- | Generic 'embed'
gEmbed :: GShaped a
       => (forall x. Shaped x => f x -> x)
       -> Shape a f -> a
gEmbed e !(GS d) =
  to (hcliftA (Proxy @Shaped) (I . e) (SOP d))

-- | Generic 'match'
gMatch :: forall a f g result. GShaped a
       => Shape a f -> Shape a g
       -> (forall xs. All Shaped xs
             => Flattened (Shape a) f xs
             -> Maybe (Flattened (Shape a) g xs)
             -> result)
       -> result
gMatch !(GS df) !(GS dg) cont =
  go @(Code a) df (Just dg) $ \flatF mflatG ->
    cont (flatGD flatF) (flatGD <$> mflatG)
  where
    go :: forall xss r.
      (All SListI xss, All2 Shaped xss)
       => NS (NP f) xss
       -> Maybe (NS (NP g) xss)
       -> (forall xs. All Shaped xs
             => Flattened (Flip SOP xss) f xs
             -> Maybe (Flattened (Flip SOP xss) g xs)
             -> r)
       -> r
    go (Z (fieldsF :: _ xs)) (Just (Z fieldsG)) k =
      k @xs (flatZ fieldsF)  (Just (flatZ fieldsG))
    go (Z (fieldsF :: _ xs)) _ k =   -- Nothing | Just (S _)
      k @xs (flatZ fieldsF)  Nothing
    go (S moreF) Nothing k =
      go moreF Nothing $ \(flatF :: _ xs) _ ->
        k @xs (flatS flatF) Nothing
    go (S moreF) (Just (Z _)) k =
      go moreF Nothing $ \(flatF :: _ xs) _ ->
        k @xs (flatS flatF) Nothing
    go (S moreF) (Just (S moreG)) k =
      go moreF (Just moreG) $ \(flatF :: _ xs) mflatG ->
        k @xs (flatS flatF) (flatS <$> mflatG)

    flatZ
      :: forall h xs xss. NP h xs -> Flattened (Flip SOP (xs : xss)) h xs
    flatZ = Flattened (Flip . SOP . Z)

    flatS
      :: forall h xs xs' xss.
      Flattened (Flip SOP xss) h xs
      -> Flattened (Flip SOP (xs' : xss)) h xs
    flatS (Flattened un fields) =
      Flattened (Flip . SOP . S . coerce . un) fields

    flatGD :: forall t h xs.
      Flattened (Flip SOP (Code t)) h xs -> Flattened (GShape t) h xs
    flatGD (Flattened un fields) =
      Flattened (GS . coerce . un) fields

-- | Generic 'render'
gRender :: forall a x. (HasDatatypeInfo a, GShaped a)
         => Shape a (K x) -> RenderLevel x
gRender (GS demand) =
  case info of
    ADT m d cs ->
      renderC m d demand cs
    Newtype m d c ->
      renderC m d demand (c :* Nil)
  where
    info = datatypeInfo (Proxy @a)

    renderC :: forall as. ModuleName -> DatatypeName
            -> NS (NP (K x)) as
            -> NP ConstructorInfo as
            -> RenderLevel x
    renderC m d subShape constructors =
      case (subShape, constructors) of
        (Z demandFields, c :* _) ->
          case c of
            Constructor name ->
              ConstructorD (m, d, name) $
                hcollapse demandFields
            Infix name associativity fixity ->
              case demandFields of
                (K a :* K b :* Nil) ->
                  InfixD (m, d, name) associativity fixity a b
            Record name fieldsInfo ->
              RecordD (m, d, name) $
                zip ( hcollapse
                    . hliftA (\(FieldInfo f) -> K (m, d, f))
                    $ fieldsInfo )
                    (hcollapse demandFields)
        (S another, _ :* different) ->
          renderC m d another different

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

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

instance (Typeable a, Typeable b) => Shaped (a -> b) where
  type Shape (a -> b) = Prim (a -> b)
  project = projectPrim
  embed = embedPrim
  match (Prim f) (Prim g) k = k (flatPrim f) (Just $ flatPrim g)
  render _ = renderConstant ("<function> :: " ++ show (typeRep @(a -> b)))

instance Shaped Char where
  type Shape Char = Prim Char
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance Shaped Word where
  type Shape Word = Prim Word
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance Shaped Int where
  type Shape Int = Prim Int
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance Shaped Double where
  type Shape Double = Prim Double
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance Shaped Float where
  type Shape Float = Prim Float
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance Shaped Rational where
  type Shape Rational = Prim Rational
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance Shaped Integer where
  type Shape Integer = Prim Integer
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

instance (Typeable a, Eq a, Show a) => Shaped (Complex a) where
  type Shape (Complex a) = Prim (Complex a)
  project = projectPrim
  embed   = embedPrim
  match   = matchPrim
  render  = renderPrim

-- instance Generic (NonEmpty a)
-- instance HasDatatypeInfo (NonEmpty a)
-- instance Shaped a => Shaped (NonEmpty a) where

-- Tree
-- Map k
-- Seq
-- Set
-- IntMap
-- IntSet

instance (Shaped a, Shaped b) => Shaped (a, b)
instance (Shaped a, Shaped b, Shaped c) => Shaped (a, b, c)
instance (Shaped a, Shaped b, Shaped c, Shaped d) => Shaped (a, b, c, d)
instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e
         ) => Shaped
  (a, b, c, d, e)
instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
         ) => Shaped
  (a, b, c, d, e, f)
instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
         , Shaped g
         ) => Shaped
  (a, b, c, d, e, f, g)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h
--          ) => Shaped
--   (a, b, c, d, e, f, g, h)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t, Shaped u
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t, Shaped u, Shaped v
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t, Shaped u, Shaped v, Shaped w
--          ) => Shaped
--   (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r, s, t, u, v, w)
-- instance ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t, Shaped u, Shaped v, Shaped w, Shaped x
--           ) => Shaped
--   (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 ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t, Shaped u, Shaped v, Shaped w, Shaped x
--          , Shaped y
--          ) => Shaped
--   (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 ( Shaped a, Shaped b, Shaped c, Shaped d, Shaped e, Shaped f
--          , Shaped g, Shaped h, Shaped i, Shaped j, Shaped k, Shaped l
--          , Shaped m, Shaped n, Shaped o, Shaped p, Shaped q, Shaped r
--          , Shaped s, Shaped t, Shaped u, Shaped v, Shaped w, Shaped x
--          , Shaped y, Shaped z
--          ) => Shaped
--   (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)


{-# LANGUAGE DeriveFunctor             #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings         #-}
module ML.Gym.Value where

import           ML.Gym.DType
import           ML.Gym.Util

import           Control.Applicative
import qualified Control.Exception        as E
import           Control.Monad            (join, void, (>=>))
import qualified CPython                  as Py
import qualified CPython.Constants        as Py
import qualified CPython.Protocols.Object as Py
import qualified CPython.Types            as Py
import qualified CPython.Types.Exception  as Py
import qualified CPython.Types.Module     as Py
import qualified Data.ByteString          as BS
import           Data.Complex             (Complex)
import           Data.Foldable            (toList)
import           Data.Maybe               (fromMaybe)
import qualified Data.Text                as T
import           System.Environment       (getArgs, getProgName)
import           System.Exit
import           System.IO
import           System.Random            (randomRIO)

import           Debug.Trace

data GymValue a
  = GymScalar a
  | Gym1D [a]
  | Gym2D [[a]]
  | Gym3D [[[a]]]
  | GymTuple [GymValue a]
  deriving (Show, Eq, Functor)

gymValueTo1D :: GymValue a -> [a]
gymValueTo1D (GymScalar x) = [x]
gymValueTo1D (Gym1D xs)    = xs
gymValueTo1D (Gym2D xs)    = concat xs
gymValueTo1D (Gym3D xs)    = concat $ concat xs
gymValueTo1D (GymTuple xs) = concatMap gymValueTo1D xs

gymValueTo2D :: GymValue a -> [[a]]
gymValueTo2D (GymScalar x) = [[x]]
gymValueTo2D (Gym1D xs)    = [xs]
gymValueTo2D (Gym2D xs)    = xs
gymValueTo2D (Gym3D xs)    = concat xs
gymValueTo2D (GymTuple xs) = map gymValueTo1D xs

gymValueTo3D :: GymValue a -> [[[a]]]
gymValueTo3D (GymScalar x) = [[[x]]]
gymValueTo3D (Gym1D xs)    = [[xs]]
gymValueTo3D (Gym2D xs)    = [xs]
gymValueTo3D (Gym3D xs)    = xs
gymValueTo3D (GymTuple xs) = map gymValueTo2D xs

fromBool :: Num a => Bool -> a
fromBool True  = 1
fromBool False = 0


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

{-# OPTIONS_GHC -fno-warn-orphans #-}

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

{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE EmptyDataDecls            #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE ForeignFunctionInterface  #-}
{-# LANGUAGE FunctionalDependencies    #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE KindSignatures            #-}
{-# LANGUAGE LambdaCase                #-}
{-# LANGUAGE MagicHash                 #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE Rank2Types                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilyDependencies    #-}
{-# LANGUAGE TypeInType                #-}
{-# LANGUAGE UnboxedTuples             #-}
{-# LANGUAGE UndecidableInstances      #-}

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

-- | Internal module to Eigen.
--   Here we define all foreign function calls,
--   and some typeclasses integral to the public and private interfaces
--   of the library.
module Eigen.Internal where --   FIXME: Explicit export list

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

--import           Control.Monad            (when)
import           Data.Binary              (Binary(put,get))
import           Data.Binary.Get          (getByteString, getWord32be)
import           Data.Binary.Put          (putByteString, putWord32be)
import           Data.Bits                (xor)
import           Data.Complex             (Complex((:+)))
import           Data.Kind                (Type)
import           Data.Proxy               (Proxy(Proxy))
--import           Foreign.C.String         (CString, peekCString)
import           Foreign.C.Types          (CInt(CInt), CFloat(CFloat), CDouble(CDouble), CChar)
import           Foreign.ForeignPtr       (ForeignPtr, castForeignPtr, withForeignPtr)
import           Foreign.Ptr              (Ptr, castPtr, plusPtr)
import           Foreign.Storable         (Storable(sizeOf, alignment, poke, peek, peekByteOff, peekElemOff, pokeByteOff, pokeElemOff))
import           GHC.TypeLits             (natVal, KnownNat, Nat)
import           System.IO.Unsafe         (unsafeDupablePerformIO)
import qualified Data.Vector.Storable     as VS
import qualified Data.ByteString          as BS
import qualified Data.ByteString.Internal as BSI

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

-- | Like 'Proxy', but specialised to 'Nat'.
data Row (r :: Nat) = Row
-- | Like 'Proxy', but specialised to 'Nat'.
data Col (c :: Nat) = Col

-- | Used internally. Given a 'KnownNat' constraint, turn the type-level 'Nat' into an 'Int'.
natToInt :: forall n. KnownNat n => Int
{-# INLINE natToInt #-}
natToInt = fromIntegral (natVal @n Proxy)

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

-- | Cast to and from a C-FFI type
--   'Cast' is a closed typeclass with an associated injective type family.
--   It is closed in the sense that we provide only four types
--   with instances for it; and intend for eigen to only be used
--   with those four types. The injectivity of the type family is
--   then useful for avoiding MPTCs. 'Cast' has two functions; 'toC'
--   and 'fromC', where 'toC' goes from a Haskell type to its associated
--   C type for internal use, with the C FFI, and 'fromC' goes from the
--   associated C type to the Haskell type.
class Cast (a :: Type) where
  type family C a = (result :: Type) | result -> a
  toC   :: a -> C a
  fromC :: C a -> a

instance Cast Int where
  type C Int = CInt
  toC = CInt . fromIntegral
  {-# INLINE toC #-}
  fromC (CInt x) = fromIntegral x
  {-# INLINE fromC #-}

instance Cast Float where
  type C Float = CFloat
  toC = CFloat
  {-# INLINE toC #-}
  fromC (CFloat x) = x
  {-# INLINE fromC #-}

instance Cast Double where
  type C Double = CDouble
  toC = CDouble
  {-# INLINE toC #-}
  fromC (CDouble x) = x
  {-# INLINE fromC #-}

instance Cast a => Cast (Complex a) where
  type C (Complex a) = CComplex (C a)
  toC (a :+ b) = CComplex (toC a) (toC b)
  {-# INLINE toC #-}
  fromC (CComplex a b) = (fromC a) :+ (fromC b)
  {-# INLINE fromC #-}

-- | WARNING! 'toC' is lossy for any Int greater than (maxBound :: Int32)!
instance Cast a => Cast (Int, Int, a) where
  type C (Int, Int, a) = CTriplet a
  {-# INLINE toC #-}
  toC (x, y, z) = CTriplet (toC x) (toC y) (toC z)
  {-# INLINE fromC #-}
  fromC (CTriplet x y z) = (fromC x, fromC y, fromC z)

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

-- | Complex number for FFI with the same memory layout as std::complex\<T\>
data CComplex a = CComplex !a !a deriving (Show)

instance Storable a => Storable (CComplex a) where
    sizeOf _ = sizeOf (undefined :: a) * 2
    alignment _ = alignment (undefined :: a)
    poke p (CComplex x y) = do
        pokeElemOff (castPtr p) 0 x
        pokeElemOff (castPtr p) 1 y
    peek p = CComplex
        <$> peekElemOff (castPtr p) 0
        <*> peekElemOff (castPtr p) 1

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

-- | FIXME: Doc
data CTriplet a where
  CTriplet :: Cast a => !CInt -> !CInt -> !(C a) -> CTriplet a

deriving instance (Show a, Show (C a)) => Show (CTriplet a)

instance (Storable a, Elem a) => Storable (CTriplet a) where
    sizeOf _ = sizeOf (undefined :: a) + sizeOf (undefined :: CInt) * 2
    alignment _ = alignment (undefined :: CInt)
    poke p (CTriplet row col val) = do
        pokeElemOff (castPtr p) 0 row
        pokeElemOff (castPtr p) 1 col
        pokeByteOff p (sizeOf (undefined :: CInt) * 2) val
    peek p = CTriplet
        <$> peekElemOff (castPtr p) 0
        <*> peekElemOff (castPtr p) 1
        <*> peekByteOff p (sizeOf (undefined :: CInt) * 2)

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

-- | `Elem` is a closed typeclass that encompasses the properties
--   eigen expects its values to possess, and simplifies the external
--   API quite a bit.
class (Num a, Cast a, Storable a, Storable (C a), Code (C a)) => Elem a

instance Elem Float
instance Elem Double
instance Elem (Complex Float)
instance Elem (Complex Double)

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

-- | Encode a C Type as a CInt
--
--   Hack used in FFI wrapper functions when constructing FFI calls
class Code a where; code :: a -> CInt
instance Code CFloat             where; code _ = 0
instance Code CDouble            where; code _ = 1
instance Code (CComplex CFloat)  where; code _ = 2
instance Code (CComplex CDouble) where; code _ = 3

-- | Hack used in constructing FFI calls.
newtype MagicCode = MagicCode CInt deriving Eq

instance Binary MagicCode where
    put (MagicCode _code) = putWord32be $ fromIntegral _code
    get = MagicCode . fromIntegral <$> getWord32be

-- | Hack used in constructing FFI calls.
magicCode :: Code a => a -> MagicCode
magicCode x = MagicCode (code x `xor` 0x45696730)

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

-- | Machine size of a 'CInt'.
intSize :: Int
intSize = sizeOf (undefined :: CInt)

-- | FIXME: Doc
encodeInt :: CInt -> BS.ByteString
encodeInt x = BSI.unsafeCreate (sizeOf x) $ (`poke` x) . castPtr

-- | FIXME: Doc
decodeInt :: BS.ByteString -> CInt
decodeInt (BSI.PS fp fo fs)
    | fs == sizeOf x = x
    | otherwise = error "decodeInt: wrong buffer size"
    where x = performIO $ withForeignPtr fp $ peek . (`plusPtr` fo)

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

-- | 'Binary' instance for 'Data.Vector.Storable.Mutable.Vector'
instance Storable a => Binary (VS.Vector a) where
    put vs = put (BS.length bs) >> putByteString bs where
        (fp,fs) = VS.unsafeToForeignPtr0 vs
        es = sizeOf (VS.head vs)
        bs = BSI.fromForeignPtr (castForeignPtr fp) 0 (fs * es)

    get = get >>= getByteString >>= \bs -> let
        (fp,fo,fs) = BSI.toForeignPtr bs
        es = sizeOf (VS.head vs)
        -- `plusForeignPtr` is used qualified here to just remind a reader
        -- that it is defined internally within eigen
        vs = VS.unsafeFromForeignPtr0 (Eigen.Internal.plusForeignPtr fp fo) (fs `div` es)
        in return vs

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

-- | FIXME: Doc
data CSparseMatrix a
-- | FIXME: Doc
type CSparseMatrixPtr a = Ptr (CSparseMatrix a)

-- | FIXME: Doc
data CSolver a
-- | FIXME: Doc
type CSolverPtr a = Ptr (CSolver a)

-- {-# INLINE unholyPerformIO #-}
-- unholyPerformIO :: IO a -> a
-- unholyPerformIO (IO m) = case m realWorld# of (# _, r #) -> r

-- | FIXME: replace with unholyPerformIO (?)
performIO :: IO a -> a
performIO = unsafeDupablePerformIO

-- | FIXME: Doc
plusForeignPtr :: ForeignPtr a -> Int -> ForeignPtr b
plusForeignPtr fp fo = castForeignPtr fp1 where
    vs :: VS.Vector CChar
    vs = VS.unsafeFromForeignPtr (castForeignPtr fp) fo 0
    (fp1, _) = VS.unsafeToForeignPtr0 vs

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


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Microbenchmark where

import Test.HUnit hiding (State)
import Test.Framework
import Test.Framework.Providers.HUnit

import           System.Random
import           Control.Monad.State
import           Control.Monad
import           Control.Lens

import GHC.Conc.Sync (setNumCapabilities)

import qualified Data.HashMap.Strict       as HM
import qualified Data.HashSet              as Set
import           Data.List
import qualified Data.Vector               as V
import           Data.IORef
import           Data.Time.Clock.POSIX
import           Data.Aeson                as AE
import           Data.ByteString.Lazy      as BS (writeFile)
import           Kvservice_Types
import           Kvstore.KVSTypes

import           Requests
import           TestSetup
import           ServiceConfig
import           Versions

import Options.Applicative
import Data.Semigroup ((<>))

import           Statistics.Sample (mean)

import           Debug.Trace

valueTemplate v = "value-" ++ show v
fieldTemplate f = "field-" ++ show f
keyTemplate k = "key-" ++ show k
tableTemplate t = "table-" ++ show t

data RangeGen = RangeGen Int Int StdGen
instance RandomGen RangeGen where
  next (RangeGen lo hi g) = let (i,g') = randomR (lo,hi) g in (i, RangeGen lo hi g')
  split (RangeGen lo hi g) = let (g1,g2) = split g in (RangeGen lo hi g1, RangeGen lo hi g2)

newtype LinearGen = LinearGen Int
instance RandomGen LinearGen where
  next (LinearGen i) = (i, LinearGen $ i+1)
  split (LinearGen i) = (LinearGen i, LinearGen i)

data BenchmarkState g = BenchmarkState {
                        _fieldCount :: Int,
                        _fieldSelection :: RangeGen,
                        _valueSizeGen :: RangeGen,
                        _tableCount :: Int,
                        _tableSelection :: RangeGen,
                        _keySelection :: g,
                        _operationSelection :: RangeGen,
                        _fieldCountSelection :: RangeGen,
                        _scanCountSelection :: RangeGen
                      }

makeLenses ''BenchmarkState
-- lens does not work with RankNTypes :(

createValue :: forall g. RandomGen g => StateT (BenchmarkState g) IO String
createValue = do
  bmState <- get
  let valSizeGenerator = view valueSizeGen bmState
      (size,valSizeGenerator') = next valSizeGenerator
  put $ over valueSizeGen (const valSizeGenerator') bmState
  (return . concatMap valueTemplate . take size) [1,2..]

createRequest :: forall g. RandomGen g => StateT (BenchmarkState g) IO KVRequest
createRequest = do
  bmState <- get
  let opSelector = view operationSelection bmState
      (op,opSelector') = next opSelector
      bmState' = over operationSelection (const opSelector') bmState

      tableSelector = view tableSelection bmState
      (tableId,tableSelector') = next tableSelector
      table = tableTemplate tableId
      bmState'' = over tableSelection (const tableSelector') bmState'

      keySelector = view keySelection bmState
      (keyId,keySelector') = next keySelector
      key = keyTemplate keyId
      bmState''' = over keySelection (const keySelector') bmState''
  put bmState'''
  case op of
        0 -> prepareINSERT table key <$> createINSERTEntry
        1 -> prepareUPDATE table key <$> createUPDATEEntry
        2 -> prepareREAD table key . Set.map fieldTemplate <$> getFields
        3 -> prepareSCAN table key <$> getScanCount <*> (Set.map fieldTemplate <$> getFields)
        4 -> return $ prepareDELETE table key
        _ -> error $ "No such operation: " ++ show op
    where
      getFieldsAndValues :: forall g. RandomGen g => [Int] -> StateT (BenchmarkState g) IO (HM.HashMap String String)
      getFieldsAndValues = fmap HM.fromList . mapM (\ i -> (,) <$> pure (fieldTemplate i) <*> createValue)
      createINSERTEntry :: forall g. RandomGen g => StateT (BenchmarkState g) IO (HM.HashMap String String)
      createINSERTEntry = getFieldsAndValues . flip take  [1,2..] . view fieldCount =<< get
      getFields :: forall g. RandomGen g => StateT (BenchmarkState g) IO (Set.HashSet Int)
      getFields = do
        s <- get
        let fieldCountGen = view fieldCountSelection s
            (fieldCount,fieldCountGen') = next fieldCountGen
            s' = over fieldCountSelection (const fieldCountGen') s
            fieldSel = view fieldSelection s'
            (fieldSel', fields) = foldl (\ (sel,l) _ -> let (f,sel') = next sel in (sel',l ++ [f]) )
                                       (fieldSel,[])
                                       $ take fieldCount [1,2..]
            s'' = over fieldSelection (const fieldSel') s'
        put s''
        return $ Set.fromList fields
      createUPDATEEntry :: forall g. RandomGen g => StateT (BenchmarkState g) IO (HM.HashMap String String)
      createUPDATEEntry = getFieldsAndValues . Set.toList =<< getFields
      getScanCount = do
        s <- get
        let scanCountSel = view scanCountSelection s
            (scanCount,scanCountSel') = next scanCountSel
        put $ over scanCountSelection (const scanCountSel') s
        return scanCount

workload :: forall g. RandomGen g => Int -> StateT (BenchmarkState g) IO (V.Vector KVRequest)
workload operationCount = V.fromList <$> mapM (const createRequest) [1..operationCount]

showState :: KVSState MockDB -> IO String
showState (KVSState cache dbRef _ _) = do
  db <- readIORef dbRef
  return $ "Cache:\n" ++ show cache ++ "\nDB:\n" ++ show db

loadDB keyCount = do
  s <- initState
  -- fill the db first
  (requests,_) <- runStateT (workload keyCount) $ BenchmarkState
                                            10 -- _fieldCount
                                            (RangeGen 0 10 $ mkStdGen 0) -- _fieldSelection
                                            (RangeGen 5 10 $ mkStdGen 0) -- _valueSizeGen
                                            1 -- _tableCount
                                            (RangeGen 1 1 $ mkStdGen 0) -- _tableSelection
                                            (LinearGen 1) -- _keySelection
                                            (RangeGen 0 0 $ mkStdGen 0) -- _operationSelection (INSERT only)
                                            (RangeGen 3 10 $ mkStdGen 0) -- _fieldCountSelection
                                            (RangeGen 5 10 $ mkStdGen 0) -- _scanCountSelection
  -- traceM "requests (INSERT):"
  -- mapM (\i -> traceM $ show i ++ "\n" ) requests
  (responses, s') <- flip runStateT s $ ?execRequests requests
  responses `seq` assertEqual "wrong number of responses" keyCount $ length responses
  return s'

  -- traceM "state after init:"
  -- traceM =<< showState s'
  -- traceM "done with insert."

reqBenchmarkState keyCount = BenchmarkState
                                10 -- _fieldCount
                                (RangeGen 0 10 $ mkStdGen 0) -- _fieldSelection
                                (RangeGen 5 10 $ mkStdGen 0) -- _valueSizeGen
                                1 -- _tableCount
                                (RangeGen 1 1 $ mkStdGen 0) -- _tableSelection
                                (RangeGen 1 keyCount $ mkStdGen 0) -- _keySelection
                                (RangeGen 1 3 $ mkStdGen 0) -- _operationSelection (no INSERT, no DELETE)
                                (RangeGen 3 10 $ mkStdGen 0) -- _fieldCountSelection
                                (RangeGen 5 10 $ mkStdGen 0) -- _scanCountSelection

currentTimeMillis = round . (* 1000) <$> getPOSIXTime

-- then run some requests
-- runRequests :: (?execRequests :: ExecReqFn)
--             => Int -> Int -> BenchmarkState RangeGen -> KVSState MockDB -> IO (KVSState MockDB, Integer)
runRequests operationCount keyCount bmState s = do
  (requests,_) <- runStateT (workload operationCount) bmState
  -- traceM $ "requests:"
  -- mapM (\i -> traceM $ show i ++ "\n" ) requests
  start <- currentTimeMillis
  (responses, s') <- requests `seq` flip runStateT s $ ?execRequests requests
  stop <- currentTimeMillis
  let execTime = stop - start
  -- traceM "???????????????????????????????????????"
  -- traceM $ "responses:"
  -- mapM (\i -> traceM $ show i ++ "\n" ) responses
  responses `seq` assertEqual "wrong number of responses" operationCount $ length responses
  return (s',execTime)

runSingleBatch :: (?execRequests :: ExecReqFn) => Int -> Int -> IO ()
runSingleBatch keyCount reqCount = do
  (_, execTime) <- runRequests reqCount keyCount (reqBenchmarkState keyCount) =<< loadDB keyCount
  traceM $ "exec time: " ++ show execTime
  return ()

runMultipleBatches :: (?execRequests :: ExecReqFn) => Int -> Int -> Int -> IO Double
runMultipleBatches keyCount reqCount batchCount = do
   s <- loadDB keyCount
   let bmState = reqBenchmarkState keyCount
   (_, execTimes) <- foldM (\ (st,execTimes) _ -> do
                                (st',execTime) <- runRequests reqCount keyCount bmState st
                                return $ st' `seq` (st',execTimes ++ [execTime]))
                           (s,[])
                           $ take batchCount [1,2..]
   let meanExecTime = mean $ V.fromList $ map fromIntegral execTimes
   -- traceM $ "mean execution time: " ++ show meanExecTime ++ " ms"
   return meanExecTime

runMultipleBatches_ :: (?execRequests :: ExecReqFn) => Int -> Int -> Int -> IO ()
runMultipleBatches_ keyCount reqCount batchCount = (\e -> traceM $ "mean execution time: " ++ show e ++ " ms") =<< runMultipleBatches keyCount reqCount batchCount

data ScalabilityResult = ScalabilityResult { version :: String
                                           , results :: [Double] }

scalability name numThreads = do
    results <- foldM (\res n -> do
                      _ <- setNumCapabilities n
                      r <- runMultipleBatches 2000 20 20
                      return $ res ++ [r])
                    []
                    $ take numThreads [1,2..]
    return (name,results)

buildSuite (BenchmarkConfig maxThreadCount) =
  [testCase "Done!" runAsSingleTest]
  where
    runAsSingleTest = do
      results <- forM versions $ \(version, name) -> let ?execRequests = version in scalability name maxThreadCount
      BS.writeFile "scalability.json" $ AE.encode $ HM.fromList results

data BenchmarkConfig = BenchmarkConfig { maxThreadCount :: Int }

benchmarkOptionsParser = BenchmarkConfig <$> (read <$> strOption
                                                         ( short 't'
                                                        <> long "num-threads"
                                                        <> metavar "NUM_THREADS"
                                                        <> help "Maximum number of threads."
                                                        <> showDefault
                                                        <> value "8" ))


-- |
-- Copyright: (C) 2013 Amgen, Inc.
--

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}

module Data.Vector.SEXP.Base where

import Control.Memory.Region

import Foreign.R.Type
import Foreign.R (SEXP, SomeSEXP)

import Data.Singletons (SingI)

import Data.Complex (Complex)
import Data.Word (Word8)
import Data.Int (Int32)
import Foreign.Storable (Storable)

-- | Function from R types to the types of the representations of each element
-- in the vector.
type family ElemRep s (a :: SEXPTYPE) where
  ElemRep s 'Char    = Word8
  ElemRep s 'Logical = Logical
  ElemRep s 'Int     = Int32
  ElemRep s 'Real    = Double
  ElemRep s 'Complex = Complex Double
  ElemRep s 'String  = SEXP s 'Char
  ElemRep s 'Vector  = SomeSEXP s
  ElemRep s 'Expr    = SomeSEXP s
  ElemRep s 'Raw     = Word8

-- | 'ElemRep' in the form of a relation, for convenience.
type E s a b = ElemRep s a ~ b

-- | Constraint synonym for all operations on vectors.
type VECTOR s ty a = (Storable a, IsVector ty, SingI ty)

-- | Constraint synonym for all operations on vectors.
type SVECTOR ty a = (Storable a, IsVector ty, SingI ty, ElemRep V ty ~ a)


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
module Bench where
import           Graphics.Rendering.Chart.Easy
import           Graphics.Rendering.Chart.Backend.Diagrams
import           System.FilePath
import           System.Directory
import qualified Data.Vector                   as Vec
import qualified Statistics.Sample             as Stat
import qualified Data.Csv                      as Csv
import qualified Data.ByteString.Lazy          as B
import qualified Data.Map                      as Map
import           Data.Type.Natural
import           Data.List                      ( intercalate )
import qualified Data.Store                    as Store
                                         hiding ( size )
import qualified Data.ByteString               as Bs
import           Options.Applicative
import           Data.Semigroup                 ( (<>) )

import           CodeGen.Type
import           ParPattern
import           CodeGen
import           Example.Mergesort
import           Example.DotProd
import           Example.Wordcount
import           CodeGen.Data

class RandomGenData a where
    genData :: Int -> Int -> IO a

class GenBenchData a where
    mkBenchData :: Int -> BenchRawData a

instance (GenBenchData a, GenBenchData b) => GenBenchData (a, b) where
    mkBenchData len = RPair (mkBenchData len) (mkBenchData len)

instance GenBenchData [a] where
    mkBenchData len = RList len

instance RandomGenData a => RandomGenData (a, a) where
    genData x y = do
        first  <- genData x y
        second <- genData x y
        return (first, second)

instance RandomGenData [Int] where
    genData x _y = return [x]

type Benchable a = (Serialise a, GenBenchData a)

benchmarkCompile
    :: Benchable a
    => FilePath
    -> [Int]
    -> [Int]
    -> [Int]
    -> (Int -> ArrowPipe a b)
    -> IO ()
benchmarkCompile path seeds unrolls sizes aexpr = mapM_ benchmarkUnRolls sizes
  where
    benchmarkUnRolls size = mapM_ (benchmarkSeeds size) unrolls
    benchmarkSeeds size unroll = mapM_ (benchmarkList size unroll) seeds
    benchmarkList size unroll seed = do
        let x = mkBenchData (2 ^ size)
        codeGenCompileWithHeader
            x
            (runPipe1 zero (aexpr unroll))
            (path </> (intercalate "_" [show size, show unroll, show seed]))

benchmarkRun :: FilePath -> Int -> [Int] -> [Int] -> [Int] -> IO ()
benchmarkRun path time seeds unrolls sizes =
    rawData >>= (Bs.writeFile (path </> rawDataName) . Store.encode)
  where
    rawData = fmap (zip sizes) $ mapM benchmarkUnRolls sizes
    benchmarkUnRolls size =
        fmap (zip unrolls) $ mapM (benchmarkSeeds size) unrolls
    benchmarkSeeds size unroll =
        toBenchData <$> (concat <$> mapM (benchmarkList size unroll) seeds)
    benchmarkList size unroll seed = mapM
        (const $ codeGenBenchRun
            (path </> (intercalate "_" [show size, show unroll, show seed]))
        )
        [1 .. time]

benchmarkRun2 :: FilePath -> Int -> [Int] -> [Int] -> [Int] -> IO ()
benchmarkRun2 path time seeds unrolls sizes =
    (toRawData <$> rawData)
        >>= (B.writeFile (path </> rawDataName2) . Csv.encode)
  where
    toRawData xs =
        [ RawData size' 0 unroll' (fst val)
        | (size'  , rest) <- xs
        , (unroll', val ) <- rest
        ]

    rawData = fmap (zip sizes) $ mapM benchmarkUnRolls sizes
    benchmarkUnRolls size =
        fmap (zip unrolls) $ mapM (benchmarkSeeds size) unrolls
    benchmarkSeeds size unroll =
        toBenchData <$> (concat <$> mapM (benchmarkList size unroll) seeds)
    benchmarkList size unroll seed = mapM
        (const $ codeGenBenchRun
            (path </> (intercalate "_" [show size, show unroll, show seed]))
        )
        [1 .. time]

rawDataName :: String
rawDataName = "rawData"

rawDataName2 :: String
rawDataName2 = "rawData.csv"

benchmarkCollectData :: FilePath -> IO ()
benchmarkCollectData path = do
    rawDataRead <- Bs.readFile (path </> rawDataName)
    let rawData = case Store.decode rawDataRead of
            Left  x -> error $ show x
            Right y -> y
    removeFile (path </> rawDataName)
    writeBenchmarkCSV path rawData
    writeBenchmarkPlot path rawData
    writeBenchmarkSpeedUpPlot path rawData

benchmarkCollectData2 :: FilePath -> IO ()
benchmarkCollectData2 path = do
    rawData <- rawDataDecode path
    -- removeFile (path </> rawDataName2)
    -- writeBenchmarkCSV path rawData
    writeBenchmarkPlot path rawData
    writeBenchmarkSpeedUpPlot path rawData

benchmarking
    :: Benchable a
    => FilePath
    -> [Int]
    -> [Int]
    -> [Int]
    -> (Int -> ArrowPipe a b)
    -> IO ()
benchmarking path seeds unrolls sizes expr = do
    benchmarkCompile path seeds unrolls sizes expr
    benchmarkRun path 1 seeds unrolls sizes
    benchmarkCollectData path

mean :: [Double] -> Double
mean = (Stat.mean . Vec.fromList)

stdDev :: [Double] -> Double
stdDev = (Stat.stdDev . Vec.fromList)

toBenchData :: [Double] -> (Double, Double)
toBenchData xs = (mean xs, stdDev xs)

writeBenchmarkPlot :: FilePath -> [(Int, [(Int, (Double, Double))])] -> IO ()
writeBenchmarkPlot path sourceData = toFile def imagePath $ do
    layout_title .= "benchmark"
    (plotLines . toPlottableData) sourceData
    where imagePath = path </> "bench.svg"

-- the first int is size
writeBenchmarkSpeedUpPlot
    :: FilePath -> [(Int, [(Int, (Double, Double))])] -> IO ()
writeBenchmarkSpeedUpPlot path sourceData = toFile def imagePath $ do
    layout_title .= "Speedup"
    (plotLines . toPlottableData . toSpeedUpData) sourceData
    where imagePath = path </> "benchSpeedup.svg"

-- the first int is k
plotLines :: [(Int, [(Int, Double)])] -> EC (Layout Int Double) ()
plotLines = mapM_ (uncurry plotLine)
  where
    plotLine :: Int -> [(Int, Double)] -> EC (Layout Int Double) ()
    plotLine title xs = plot (line ("k" ++ show title) [xs])

toSpeedUpData
    :: [(Int, [(Int, (Double, Double))])] -> [(Int, [(Int, (Double, Double))])]
toSpeedUpData = fmap (\(x, y) -> (x, helper helper2 y))
  where
    helper2
        :: (Int, (Double, Double))
        -> (Int, (Double, Double))
        -> (Int, (Double, Double))
    helper2 (size, parTime) (_, serTime) =
        (size, (fst serTime / fst parTime, snd parTime))

    helper f xs = fmap (\x -> f x (head xs)) (tail xs)

toPlottableData
    :: [(Int, [(Int, (Double, Double))])] -> [(Int, [(Int, Double)])]
toPlottableData = helper . flattenMap
  where
    helper :: (Ord a, Ord b) => [(a, b, Double)] -> [(b, [(a, Double)])]
    helper source = Map.toList $ foldl
        (\map (key1, key2, val) -> Map.insertWith (++) key2 [(key1, val)] map)
        (Map.empty)
        source

    flattenMap :: [(Int, [(Int, (Double, Double))])] -> [(Int, Int, Double)]
    flattenMap source =
        [ (key1, key2, d1) | (key1, val) <- source, (key2, (d1, _d2)) <- val ]

writeBenchmarkCSV :: FilePath -> [(Int, [(Int, (Double, Double))])] -> IO ()
writeBenchmarkCSV path sourceData = B.writeFile csvPath csvRaw
  where
    records :: [BenchData]
    records = concatMap (\(x, y) -> toRecordHelper x y) sourceData

    toRecordHelper :: Int -> [(Int, (Double, Double))] -> [BenchData]
    toRecordHelper x xs = fmap
        (\(unroll, (meanVal, stdVal)) -> BenchData (x, unroll, meanVal, stdVal))
        xs

    csvRaw  = Csv.encodeByName csvTitle records

    csvPath = path </> "bench.csv"

newtype BenchData = BenchData (Int, Int, Double, Double)

instance Csv.ToNamedRecord BenchData where
    toNamedRecord (BenchData (size, unroll, meanVal, stdVal)) = Csv.namedRecord
        [ "size" Csv..= size
        , "unroll" Csv..= unroll
        , "mean" Csv..= meanVal
        , "std" Csv..= stdVal
        ]

data RawData = RawData { size :: Int, seed :: Int, unroll :: Int, time :: Double }

instance Csv.FromRecord RawData where
    parseRecord v =
        RawData <$> v Csv..! 0 <*> v Csv..! 1 <*> v Csv..! 2 <*> v Csv..! 3

instance Csv.ToRecord RawData where
    toRecord RawData {..} = Csv.record
        [ Csv.toField size
        , Csv.toField seed
        , Csv.toField unroll
        , Csv.toField time
        ]

rawDataDecode :: FilePath -> IO [(Int, [(Int, (Double, Double))])]
rawDataDecode path = do
    rawData <- B.readFile filePath
    let decoder =
            (Csv.decode Csv.NoHeader rawData) :: Either
                    String
                    (Vec.Vector RawData)
    return $ rawDatadecode' decoder
  where
    filePath = path </> rawDataName2
    rawDatadecode'
        :: Either String (Vec.Vector RawData)
        -> [(Int, [(Int, (Double, Double))])]
    rawDatadecode' x = case x of
        Left  msg        -> error $ "failed " ++ msg
        Right sourceData -> (toList . toMap'' . toMap' . toMap) sourceData

    toMap = Vec.foldl
        (\map record ->
            Map.insertWith (++) (size record, unroll record) [time record] map
        )
        Map.empty

    toMap'' = Map.foldrWithKey
        (\(mySize, myUnroll) val acc ->
            Map.insertWith (Map.union) mySize (Map.singleton myUnroll val) acc
        )
        Map.empty

    toMap' = Map.map toBenchData

    toList = (Map.toList . Map.map Map.toList)

csvTitle :: Vec.Vector Csv.Name
csvTitle = Vec.fromList ["size", "unroll", "mean", "std"]

runParser :: Parser (Bool, Bool)
runParser =
    (,)
        <$> switch (long "collect" <> short 'c' <> help "whether to collect")
        <*> switch (long "run" <> short 'r' <> help "whether to run or codegen")

genBenchmarkFile
    :: Benchable a
    => FilePath
    -> [Int]
    -> [Int]
    -> (Int -> ArrowPipe a b)
    -> IO ()
genBenchmarkFile path unrolls sizes expr =
    benchmarkEntry path 1 [1] unrolls sizes expr

benchmarkEntry
    :: Benchable a
    => FilePath
    -> Int
    -> [Int]
    -> [Int]
    -> [Int]
    -> (Int -> ArrowPipe a b)
    -> IO ()
benchmarkEntry path time seeds unrolls sizes expr = execParser opts >>= \x ->
    case x of
        (True , _    ) -> benchmarkCollectData2 path
        (False, True ) -> benchmarkRun2 path time seeds unrolls sizes
        (False, False) -> benchmarkCompile path seeds unrolls sizes expr
    where opts = info (runParser <**> helper) mempty

benchmarkNewEntry :: FilePath -> Int -> [Int] -> [Int] -> [Int] -> IO ()
benchmarkNewEntry mainPath time seeds unrolls sizes =
    execParser opts >>= \x -> case x of
        (True, _) ->
            mapM_ (\path -> benchmarkCollectData (mainPath </> path)) names
        (False, True) -> mapM_
            (\path -> benchmarkRun (mainPath </> path) time seeds unrolls sizes)
            names
        (False, False) ->
            mapM_ (\name -> myHelper name (mainPath </> name)) names
  where
    names = ["intcount", "mergesort", "dotprod"]
    opts  = info (runParser <**> helper) mempty

    myHelper "intcount" path =
        benchmarkCompile path seeds unrolls sizes wordCount
    myHelper "mergesort" path =
        benchmarkCompile path seeds unrolls sizes mergeSort
    myHelper "dotprod" path = benchmarkCompile path seeds unrolls sizes dotProd
    myHelper _         _    = error "benchmark example not implemented"


{-# 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


module Main where
  import Data.Picture
  import System.Environment
  import Data.Either
  import Numeric.LinearAlgebra (rows, cols)

  data Options = Options { file :: FilePath
                         , output :: FilePath
                         , argInvert :: Bool
                         , argGrayscale :: Bool
                         , argRotate :: Double
                         , argFade :: Double
                         , argContrast :: Double
                         , argGamma :: Int
                         , argBrightness :: Double
                         , argCompress :: Int
                         , argScale :: Double
                         }

  opts = Options { file = ""
                 , output = "output.png"
                 , argInvert = False
                 , argGrayscale = False
                 , argRotate = 0
                 , argFade = 100
                 , argContrast = 0
                 , argGamma = 1
                 , argBrightness = 0
                 , argCompress = 0
                 , argScale = 1
                 }

  main :: IO ()
  main = do
    args <- getArgs

    if null args
      then do
        putStrLn "Usage: picedit <input> [OPTIONS]"
        putStrLn "Options:"
        putStrLn "  --contrast <n> - a number between -255 and 255"
        putStrLn "  --brightness <n> - a number between -255 and 255"
        putStrLn "  --gamma <n>"
        putStrLn "  --fade <n> - a number between 0 and 100"
        putStrLn "  --rotate <n> - rotate image by n degrees"
        putStrLn "  --grayscale - turn the image grayscale"
        putStrLn "  --invert - invert (negative) the image"
        putStrLn "  --scale <n> - scale the image using nearest-neighbor interpolation"
        putStrLn "  --compress <n> - approximate the (width - n)-th rank of image using SVD, note: this is not size compression, a number between 0 (no compression) and image width (full compression)"
        putStrLn "  --output <filename> - output name, defaults to output.png"
      else do
        let options = parseArgs args opts

        pic <- readPicture (file options)
        Right other <- readPicture ("output.png")

        case pic of
          Left err -> print err
          Right p -> do
            let edited = rotate (argRotate options) Nothing
                       . fade (argFade options / 100)
                       . contrast (argContrast options)
                       . gamma (argGamma options)
                       . brightness (argBrightness options)
                       . conditionalFn grayscale (argGrayscale options)
                       . conditionalFn invert (argInvert options)
                       . scale (argScale options)
                       . compress (argCompress options) $ p
            writePicturePng (output options) edited

        return ()

    where
      conditionalFn f True = f
      conditionalFn f False = id

  parseArgs :: [String] -> Options -> Options
  parseArgs [] opts = opts
  parseArgs ("--invert":rest) opts = parseArgs rest (opts { argInvert = True })
  parseArgs ("--grayscale":rest) opts = parseArgs rest (opts { argGrayscale = True })
  parseArgs ("--rotate":n:rest) opts = parseArgs rest (opts { argRotate = read n })
  parseArgs ("--fade":n:rest) opts = parseArgs rest (opts { argFade = read n })
  parseArgs ("--contrast":n:rest) opts = parseArgs rest (opts { argContrast = read n })
  parseArgs ("--brightness":n:rest) opts = parseArgs rest (opts { argBrightness = read n })
  parseArgs ("--gamma":n:rest) opts = parseArgs rest (opts { argGamma = read n })
  parseArgs ("--compress":n:rest) opts = parseArgs rest (opts { argCompress = read n })
  parseArgs ("--scale":n:rest) opts = parseArgs rest (opts { argScale = read n })
  parseArgs ("--output":n:rest) opts = parseArgs rest (opts { output = n })
  parseArgs (name:rest) opts = parseArgs rest (opts { file = name })