{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE DeriveFoldable        #-}
{-# LANGUAGE DeriveTraversable     #-}
{-# LANGUAGE EmptyCase             #-}
{-# LANGUAGE EmptyDataDeriving     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# OPTIONS_GHC -fno-warn-orphans  #-}

-- |
-- Module      : Data.Mutable.Instances
-- Copyright   : (c) Justin Le 2020
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Exports 'Ref' data types for various common data types, and also the
-- tools for automatic derivation of instances.  See "Data.Mutable" for
-- more information.
module Data.Mutable.Instances (
    RecRef(..)
  , HListRef(..)
  , UnitRef(..)
  , VoidRef
  -- * Generic
  , GRef(..)
  , gThawRef, gFreezeRef
  , gCopyRef, gMoveRef, gCloneRef
  , gUnsafeThawRef, gUnsafeFreezeRef
  , GMutable (GRef_)
  -- * Higher-Kinded Data Pattern
  , thawHKD, freezeHKD
  , copyHKD, moveHKD, cloneHKD
  , unsafeThawHKD, unsafeFreezeHKD
  -- * Coercible
  , CoerceRef(..)
  , thawCoerce, freezeCoerce
  , copyCoerce, moveCoerce, cloneCoerce
  , unsafeThawCoerce, unsafeFreezeCoerce
  -- * Traversable
  , TraverseRef(..)
  , thawTraverse, freezeTraverse
  , copyTraverse, moveTraverse, cloneTraverse
  , unsafeThawTraverse, unsafeFreezeTraverse
  -- * Immutable
  , ImmutableRef(..), thawImmutable, freezeImmutable, copyImmutable
  -- * Instances for Generics combinators themselves
  , GMutableRef(..)
  , MutSumF(..)
  -- * Utility
  , MapRef
  ) where

import           Control.Applicative
import           Data.Complex
import           Data.Functor.Compose
import           Data.Functor.Identity
import           Data.Functor.Product
import           Data.Functor.Sum
import           Data.Generics.Product.Internal.HList      (HList(..))
import           Data.Kind
import           Data.Mutable.Internal
import           Data.Mutable.Internal.TH
import           Data.Ord
import           Data.Primitive.Array
import           Data.Primitive.ByteArray
import           Data.Primitive.PrimArray
import           Data.Primitive.SmallArray
import           Data.Primitive.Types
import           Data.Ratio
import           Data.Vinyl                    as V hiding (HList)
import           Data.Void
import           Data.Word
import           Foreign.C.Types
import           Foreign.Storable
import           Numeric.Natural
import qualified Data.Monoid                               as M
import qualified Data.Vector                               as V
import qualified Data.Vector.Generic                       as VG
import qualified Data.Vector.Generic.Mutable               as MVG
import qualified Data.Vector.Mutable                       as MV
import qualified Data.Vector.Primitive                     as VP
import qualified Data.Vector.Primitive.Mutable             as MVP
import qualified Data.Vector.Storable                      as VS
import qualified Data.Vector.Storable.Mutable              as MVS
import qualified Data.Vector.Unboxed                       as VU
import qualified Data.Vector.Unboxed.Mutable               as MVU
import qualified Data.Vinyl.ARec                           as V
import qualified Data.Vinyl.Functor                        as V
import qualified Data.Vinyl.TypeLevel                      as V

instance Mutable s Int
instance Mutable s Integer
instance Mutable s Natural
instance Mutable s (Ratio a)
instance Mutable s Float
instance Mutable s Double
instance Mutable s (Complex a)
instance Mutable s Bool
instance Mutable s Char

instance Mutable s Word
instance Mutable s Word8
instance Mutable s Word16
instance Mutable s Word64

instance Mutable s CChar
instance Mutable s CSChar
instance Mutable s CUChar
instance Mutable s CShort
instance Mutable s CUShort
instance Mutable s CInt
instance Mutable s CUInt
instance Mutable s CLong
instance Mutable s CULong
instance Mutable s CPtrdiff
instance Mutable s CSize
instance Mutable s CWchar
instance Mutable s CSigAtomic
instance Mutable s CLLong
instance Mutable s CULLong
instance Mutable s CBool
instance Mutable s CIntPtr
instance Mutable s CUIntPtr
instance Mutable s CIntMax
instance Mutable s CUIntMax
instance Mutable s CClock
instance Mutable s CTime
instance Mutable s CUSeconds
instance Mutable s CSUSeconds
instance Mutable s CFloat
instance Mutable s CDouble

instance Mutable s a => Mutable s (Identity a) where
    type Ref s (Identity a) = CoerceRef s (Identity a) a

instance Mutable s a => Mutable s (Const a b) where
    type Ref s (Const a b) = CoerceRef s (Const a b) a

instance Mutable s a => Mutable s (V.Const a b) where
    type Ref s (V.Const a b) = CoerceRef s (V.Const a b) a

instance Mutable s a => Mutable s (M.Product a) where
    type Ref s (M.Product a) = CoerceRef s (M.Product a) a

instance Mutable s a => Mutable s (M.Sum a) where
    type Ref s (M.Sum a) = CoerceRef s (M.Sum a) a

instance Mutable s a => Mutable s (Down a) where
    type Ref s (Down a) = CoerceRef s (Down a) a

instance Mutable s a => Mutable s (M.Dual a) where
    type Ref s (M.Dual a) = CoerceRef s (M.Dual a) a

instance Mutable s a => Mutable s (Maybe a) where
    type Ref s (Maybe a) = GRef s (Maybe a)

instance (Mutable s a, Mutable s b) => Mutable s (Either a b) where
    type Ref s (Either a b) = GRef s (Either a b)

instance (Mutable s (f a), Mutable s (g a)) => Mutable s (Product f g a) where
    type Ref s (Product f g a) = GRef s (Product f g a)

instance (Mutable s (f a), Mutable s (g a)) => Mutable s (Sum f g a) where
    type Ref s (Sum f g a) = GRef s (Sum f g a)

instance (Mutable s (f (g a))) => Mutable s (Compose f g a) where
    type Ref s (Compose f g a) = CoerceRef s (Compose f g a) (f (g a))

-- | Mutable linked list with mutable references in each cell.  See
-- 'Data.Mutable.MutBranch' documentation for an example of using this as
-- a mutable linked list.l
instance Mutable s a => Mutable s [a] where
    type Ref s [a] = GRef s [a]

-- | Meant for usage with higher-kinded data pattern (See 'X.HKD')
instance Mutable s a => Mutable s (V.Identity a) where
    type Ref s (V.Identity a) = RefFor s a
    thawRef (V.Identity x) = RefFor <$> thawRef x
    freezeRef (RefFor r) = V.Identity <$> freezeRef r
    copyRef (RefFor r) (V.Identity x) = copyRef r x
    moveRef (RefFor r) (RefFor v) = moveRef r v
    cloneRef = fmap RefFor . cloneRef . getRefFor
    unsafeThawRef (V.Identity x) = RefFor <$> unsafeThawRef x
    unsafeFreezeRef (RefFor r) = V.Identity <$> unsafeFreezeRef r

-- | Mutable reference is 'MV.MVector'.
instance Mutable s (V.Vector a) where
    type Ref s (V.Vector a) = MV.MVector s a
    thawRef         = VG.thaw
    freezeRef       = VG.freeze
    copyRef         = VG.copy
    moveRef         = MVG.move
    cloneRef        = MVG.clone
    unsafeThawRef   = VG.unsafeThaw
    unsafeFreezeRef = VG.unsafeFreeze

-- | Mutable reference is 'MVS.MVector'.
instance Storable a => Mutable s (VS.Vector a) where
    type Ref s (VS.Vector a) = MVS.MVector s a
    thawRef         = VG.thaw
    freezeRef       = VG.freeze
    copyRef         = VG.copy
    moveRef         = MVG.move
    cloneRef        = MVG.clone
    unsafeThawRef   = VG.unsafeThaw
    unsafeFreezeRef = VG.unsafeFreeze

-- | Mutable reference is 'MVU.MVector'.
instance VU.Unbox a => Mutable s (VU.Vector a) where
    type Ref s (VU.Vector a) = MVU.MVector s a
    thawRef         = VG.thaw
    freezeRef       = VG.freeze
    copyRef         = VG.copy
    moveRef         = MVG.move
    cloneRef        = MVG.clone
    unsafeThawRef   = VG.unsafeThaw
    unsafeFreezeRef = VG.unsafeFreeze

-- | Mutable reference is 'MVP.MVector'.
instance Prim a => Mutable s (VP.Vector a) where
    type Ref s (VP.Vector a) = MVP.MVector s a
    thawRef         = VG.thaw
    freezeRef       = VG.freeze
    copyRef         = VG.copy
    moveRef         = MVG.move
    cloneRef        = MVG.clone
    unsafeThawRef   = VG.unsafeThaw
    unsafeFreezeRef = VG.unsafeFreeze

instance Mutable s (Array a) where
    type Ref s (Array a) = MutableArray s a

    thawRef xs = thawArray xs 0 (sizeofArray xs)
    freezeRef rs = freezeArray rs 0 (sizeofMutableArray rs)
    copyRef rs xs = copyArray rs 0 xs 0 l
      where
        l = sizeofArray xs `min` sizeofMutableArray rs
    moveRef rs vs = copyMutableArray rs 0 vs 0 l
      where
        l = sizeofMutableArray vs `min` sizeofMutableArray rs
    cloneRef rs = cloneMutableArray rs 0 (sizeofMutableArray rs)
    unsafeThawRef   = unsafeThawArray
    unsafeFreezeRef = unsafeFreezeArray

instance Mutable s (SmallArray a) where
    type Ref s (SmallArray a) = SmallMutableArray s a

    thawRef xs = thawSmallArray xs 0 (sizeofSmallArray xs)
    freezeRef rs = freezeSmallArray rs 0 (sizeofSmallMutableArray rs)
    copyRef rs xs = copySmallArray rs 0 xs 0 l
      where
        l = sizeofSmallArray xs `min` sizeofSmallMutableArray rs
    moveRef rs vs = copySmallMutableArray rs 0 vs 0 l
      where
        l = sizeofSmallMutableArray vs `min` sizeofSmallMutableArray rs
    cloneRef rs = cloneSmallMutableArray rs 0 (sizeofSmallMutableArray rs)
    unsafeThawRef   = unsafeThawSmallArray
    unsafeFreezeRef = unsafeFreezeSmallArray

instance Mutable s ByteArray where
    type Ref s ByteArray = MutableByteArray s

    thawRef xs = do
        rs <- newByteArray (sizeofByteArray xs)
        copyByteArray rs 0 xs 0 (sizeofByteArray xs)
        pure rs
    freezeRef rs = do
        xs <- newByteArray (sizeofMutableByteArray rs)
        copyMutableByteArray xs 0 rs 0 (sizeofMutableByteArray rs)
        unsafeFreezeByteArray xs
    copyRef rs xs = copyByteArray rs 0 xs 0 l
      where
        l = sizeofByteArray xs `min` sizeofMutableByteArray rs
    moveRef rs vs = copyMutableByteArray rs 0 vs 0 l
      where
        l = sizeofMutableByteArray vs `min` sizeofMutableByteArray rs
    cloneRef rs = do
        vs <- newByteArray (sizeofMutableByteArray rs)
        copyMutableByteArray vs 0 rs 0 (sizeofMutableByteArray rs)
        pure vs
    unsafeThawRef   = unsafeThawByteArray
    unsafeFreezeRef = unsafeFreezeByteArray

instance Prim a => Mutable s (PrimArray a) where
    type Ref s (PrimArray a) = MutablePrimArray s a

    thawRef xs = do
        rs <- newPrimArray (sizeofPrimArray xs)
        copyPrimArray rs 0 xs 0 (sizeofPrimArray xs)
        pure rs
    freezeRef rs = do
        xs <- newPrimArray (sizeofMutablePrimArray rs)
        copyMutablePrimArray xs 0 rs 0 (sizeofMutablePrimArray rs)
        unsafeFreezePrimArray xs
    copyRef rs xs = copyPrimArray rs 0 xs 0 l
      where
        l = sizeofPrimArray xs `min` sizeofMutablePrimArray rs
    moveRef rs vs = copyMutablePrimArray rs 0 vs 0 l
      where
        l = sizeofMutablePrimArray vs `min` sizeofMutablePrimArray rs
    cloneRef rs = do
        vs <- newPrimArray (sizeofMutablePrimArray rs)
        copyMutablePrimArray vs 0 rs 0 (sizeofMutablePrimArray rs)
        pure vs
    unsafeThawRef   = unsafeThawPrimArray
    unsafeFreezeRef = unsafeFreezePrimArray

data VoidRef s
  deriving (Show, Read, Eq, Ord, Functor, Traversable, Foldable)

instance Mutable s Void where
    type Ref s Void = VoidRef s
    thawRef         = \case {}
    freezeRef       = \case {}
    copyRef         = \case {}
    moveRef         = \case {}
    cloneRef        = \case {}
    unsafeThawRef   = \case {}
    unsafeFreezeRef = \case {}

data UnitRef s = UnitRef
  deriving (Show, Read, Eq, Ord, Functor, Traversable, Foldable)

instance Applicative UnitRef where
    pure _  = UnitRef
    _ <*> _ = UnitRef

instance Monad UnitRef where
    return   = pure
    _ >>= _ = UnitRef

instance Mutable s () where
    type Ref s () = UnitRef s
    thawRef   _       = pure UnitRef
    freezeRef _       = pure ()
    copyRef _ _       = pure ()
    moveRef _ _       = pure ()
    cloneRef _        = pure UnitRef
    unsafeThawRef _   = pure UnitRef
    unsafeFreezeRef _ = pure ()

-- | A 'Ref' of a tuple is a tuple of 'Ref's, for easy accessing.
--
-- @
-- Ref s (Int, 'V.Vector' Double) = ('Data.Primitive.MutVar.MutVar' s Int, 'MV.MVector' s Double)
-- @
instance (Mutable s a, Mutable s b) => Mutable s (a, b) where
    type Ref s (a, b) = (Ref s a, Ref s b)
    thawRef   (!x, !y) = (,) <$> thawRef x   <*> thawRef y
    freezeRef (u , v ) = (,) <$> freezeRef u <*> freezeRef v
    copyRef   (u , v ) (!x, !y) = copyRef u x *> copyRef v y
    moveRef   (u , v ) ( x,  y) = moveRef u x *> moveRef v y
    cloneRef  (u , v ) = (,) <$> cloneRef u   <*> cloneRef v
    unsafeThawRef   (!x, !y) = (,) <$> unsafeThawRef x   <*> unsafeThawRef y
    unsafeFreezeRef (u , v ) = (,) <$> unsafeFreezeRef u <*> unsafeFreezeRef v

mutableTuples [3..12]

-- | 'Ref' for components in a vinyl 'Rec'.
newtype RecRef s f a = RecRef { getRecRef :: Ref s (f a) }

deriving instance Eq (Ref s (f a)) => Eq (RecRef s f a)
deriving instance Ord (Ref s (f a)) => Ord (RecRef s f a)

instance Mutable s (Rec f '[]) where
    type Ref s (Rec f '[]) = Rec (RecRef s f) '[]
    thawRef   _       = pure RNil
    freezeRef _       = pure RNil
    copyRef _ _       = pure ()
    moveRef _ _       = pure ()
    cloneRef _        = pure RNil
    unsafeThawRef _   = pure RNil
    unsafeFreezeRef _ = pure RNil

instance ( Mutable s (f a)
         , Mutable s (Rec f as)
         , Ref s (Rec f as) ~ Rec (RecRef s f) as
         ) => Mutable s (Rec f (a ': as)) where
    type Ref s (Rec f (a ': as)) = Rec (RecRef s f) (a ': as)
    thawRef   = \case
      x :& xs -> (:&) <$> (RecRef <$> thawRef x) <*> thawRef xs
    freezeRef = \case
      RecRef v :& vs -> (:&) <$> freezeRef v <*> freezeRef vs
    copyRef = \case
      RecRef v :& vs -> \case
        x :& xs -> copyRef v x >> copyRef vs xs
    moveRef = \case
      RecRef v :& vs -> \case
        RecRef r :& rs ->
          moveRef v r >> moveRef vs rs
    cloneRef = \case
      RecRef v :& rs -> (:&) <$> (RecRef <$> cloneRef v) <*> cloneRef rs
    unsafeThawRef   = \case
      x :& xs -> (:&) <$> (RecRef <$> unsafeThawRef x) <*> unsafeThawRef xs
    unsafeFreezeRef = \case
      RecRef v :& vs -> (:&) <$> unsafeFreezeRef v <*> unsafeFreezeRef vs

instance ( RecApplicative as
         , V.NatToInt (V.RLength as)
         , RPureConstrained (V.IndexableField as) as
         , Mutable s (Rec f as)
         , Ref s (Rec f as) ~ Rec (RecRef s f) as
         ) => Mutable s (ARec f as) where
    type Ref s (ARec f as) = ARec (RecRef s f) as

    thawRef         = fmap toARec . thawRef   . fromARec
    freezeRef       = fmap toARec . freezeRef . fromARec
    copyRef r x     = copyRef (fromARec r) (fromARec x)
    moveRef r v     = moveRef (fromARec r) (fromARec v)
    cloneRef        = fmap toARec . cloneRef . fromARec
    unsafeThawRef   = fmap toARec . unsafeThawRef   . fromARec
    unsafeFreezeRef = fmap toARec . unsafeFreezeRef . fromARec

-- | The mutable reference of the 'HList' type from generic-lens.
data HListRef :: Type -> [Type] -> Type where
    NilRef :: HListRef s '[]
    (:!>)  :: Ref s a -> HListRef s as -> HListRef s (a ': as)
infixr 5 :!>

instance Mutable s (HList '[]) where
    type Ref s (HList '[]) = HListRef s '[]
    thawRef   _       = pure NilRef
    freezeRef _       = pure Nil
    copyRef _ _       = pure ()
    moveRef _ _       = pure ()
    cloneRef _        = pure NilRef
    unsafeThawRef _   = pure NilRef
    unsafeFreezeRef _ = pure Nil

instance (Mutable s a, Mutable s (HList as), Ref s (HList as) ~ HListRef s as) => Mutable s (HList (a ': as)) where
    type Ref s (HList (a ': as)) = HListRef s (a ': as)
    thawRef   = \case
      x :> xs -> (:!>) <$> thawRef x <*> thawRef xs
    freezeRef = \case
      v :!> vs -> (:>) <$> freezeRef v <*> freezeRef vs
    copyRef = \case
      v :!> vs -> \case
        x :> xs -> copyRef v x >> copyRef vs xs
    moveRef = \case
      v :!> vs -> \case
        r :!> rs ->
          moveRef v r >> moveRef vs rs
    cloneRef = \case
      v :!> rs -> (:!>) <$> cloneRef v <*> cloneRef rs
    unsafeThawRef   = \case
      x :> xs -> (:!>) <$> unsafeThawRef x <*> unsafeThawRef xs
    unsafeFreezeRef = \case
      v :!> vs -> (:>) <$> unsafeFreezeRef v <*> unsafeFreezeRef vs

-- ListRefTuple instances

-- this one instance is the reason why we have to define a new typeclass
-- instead of using 'ListRef' -- because of @'ListRef' s () '[]@.
instance ListRefTuple s (UnitRef s) '[] where
    tupleToListRef _ = Nil
    listRefToTuple _ = UnitRef

instance (Ref s a ~ ra) => ListRefTuple s ra '[a] where
    tupleToListRef x = x :> Nil
    listRefToTuple (x :> _) = x

instance (Ref s a ~ ra, Ref s b ~ rb) => ListRefTuple s (ra, rb) '[a, b] where
    tupleToListRef (x, y) = x :> y :> Nil
    listRefToTuple (x :> y :> _) = (x, y)

listRefTuples [3..12]


{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Copyright: (c) 2021 Tom Westerhout
-- SPDX-License-Identifier: Apache-2.0
-- Maintainer: Tom Westerhout <14264576+twesterhout@users.noreply.github.com>
--
-- DLPack: Open In Memory Tensor Structure
module DLPack
  ( dlVersion,
    DLDeviceType (..),
    KnownDLDeviceType,
    -- natToInt,
    -- reifyDeviceType,
    DLDevice (..),
    DLDataTypeCode (..),
    DLDataType (..),
    DLTensor (..),
    -- DLManagedTensor (..),
    IsDLDataType (..),
    IsDLTensor (..),
    -- tensorToFlatList,
    -- fold1,
    -- loop1,
    -- foldN,
    viaContiguousBuffer,
    ContiguousVectorView (..),
    StridedVectorView (..),
    dlAsStridedVector,
    dlAsContiguousVector,
    dlExpectDType,
    dlExpectNDim,
  )
where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad)
import Control.Monad.ST
import Data.Complex
import Data.Functor.Identity
import Data.Int
import Data.Kind (Type)
import Data.Primitive (Prim)
import Data.Primitive.PrimArray
import qualified Data.Primitive.Ptr as P
import Data.Proxy
import Data.Text (Text)
import qualified Data.Text as T
-- import Foreign.Marshal.Utils (with)

import Data.Word
import Foreign.C.Types
import Foreign.Marshal.Array (withArrayLen)
import Foreign.Ptr (FunPtr, Ptr, castPtr, nullPtr, plusPtr)
import Foreign.Storable
import qualified GHC.Exts as GHC
import GHC.Stack (HasCallStack)
import GHC.TypeLits
import Prelude hiding (init)

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

-- | Underlying DLPack version
dlVersion :: (Int, Int)
dlVersion = (0, 5)
{-# INLINE dlVersion #-}

-- | The device type in 'DLDevice'.
data DLDeviceType
  = -- | CPU device
    DLCPU
  | -- | CUDA GPU device
    DLCUDA
  | -- | Pinned CUDA CPU memory allocated by @cudaMallocHost@
    DLCUDAHost
  | -- | OpenCL device
    DLOpenCL
  | -- | Vulkan buffer for next generation graphics
    DLVulkan
  | -- | Metal for Apple GPU
    DLMetal
  | -- | Verilog simulator buffer
    DLVPI
  | -- | ROCm GPU device
    DLROCM
  | -- | Pinned ROCm CPU memory allocated by @hipMallocHost@
    DLROCMHost
  | -- | Reserved extension device type
    DLExtDev
  | -- | CUDA managed/unified memory allocated by @cudaMallocManaged@
    DLCUDAManaged
  deriving stock (Read, Show, Eq)

class KnownDLDeviceType (device :: DLDeviceType) where
  reifyDeviceType :: DLDeviceType

instance KnownDLDeviceType 'DLCPU where
  reifyDeviceType = DLCPU
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLCUDA where
  reifyDeviceType = DLCUDA
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLCUDAHost where
  reifyDeviceType = DLCUDAHost
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLOpenCL where
  reifyDeviceType = DLOpenCL
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLVulkan where
  reifyDeviceType = DLVulkan
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLMetal where
  reifyDeviceType = DLMetal
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLVPI where
  reifyDeviceType = DLVPI
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLROCM where
  reifyDeviceType = DLROCM
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLROCMHost where
  reifyDeviceType = DLROCMHost
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLExtDev where
  reifyDeviceType = DLExtDev
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLCUDAManaged where
  reifyDeviceType = DLCUDAManaged
  {-# INLINE reifyDeviceType #-}

instance Enum DLDeviceType where
  toEnum x = case x of
    1 -> DLCPU
    2 -> DLCUDA
    3 -> DLCUDAHost
    4 -> DLOpenCL
    7 -> DLVulkan
    8 -> DLMetal
    9 -> DLVPI
    10 -> DLROCM
    11 -> DLROCMHost
    12 -> DLExtDev
    13 -> DLCUDAManaged
    _ -> error $ "invalid DLDeviceType: " <> show x
  {-# INLINE toEnum #-}
  fromEnum x = case x of
    DLCPU -> 1
    DLCUDA -> 2
    DLCUDAHost -> 3
    DLOpenCL -> 4
    DLVulkan -> 7
    DLMetal -> 8
    DLVPI -> 9
    DLROCM -> 10
    DLROCMHost -> 11
    DLExtDev -> 12
    DLCUDAManaged -> 13
  {-# INLINE fromEnum #-}

data DLDevice = DLDevice
  { dlDeviceType :: !DLDeviceType,
    dlDeviceId :: {-# UNPACK #-} !Int
  }
  deriving stock (Read, Show, Eq)

instance Storable DLDevice where
  sizeOf _ = 8
  {-# INLINE sizeOf #-}
  alignment _ = 4
  {-# INLINE alignment #-}
  peek p =
    DLDevice
      <$> (toEnum . (fromIntegral :: CInt -> Int) <$> peekByteOff p 0)
      <*> ((fromIntegral :: CInt -> Int) <$> peekByteOff p 4)
  {-# INLINE peek #-}
  poke p x = do
    pokeByteOff p 0 . (fromIntegral :: Int -> CInt) . fromEnum . dlDeviceType $ x
    pokeByteOff p 4 (dlDeviceId x)
  {-# INLINE poke #-}

data DLDataTypeCode
  = DLInt
  | DLUInt
  | DLFloat
  | DLOpaqueHandle
  | DLBfloat
  | DLComplex
  deriving stock (Read, Show, Eq)

instance Enum DLDataTypeCode where
  toEnum x = case x of
    0 -> DLInt
    1 -> DLUInt
    2 -> DLFloat
    3 -> DLOpaqueHandle
    4 -> DLBfloat
    5 -> DLComplex
    _ -> error $ "invalid DLDataTypeCode: " <> show x
  {-# INLINE toEnum #-}
  fromEnum x = case x of
    DLInt -> 0
    DLUInt -> 1
    DLFloat -> 2
    DLOpaqueHandle -> 3
    DLBfloat -> 4
    DLComplex -> 5
  {-# INLINE fromEnum #-}

data DLDataType = DLDataType
  { dlDataTypeCode :: !DLDataTypeCode,
    dlDataTypeBits :: {-# UNPACK #-} !Int,
    dlDataTypeLanes :: {-# UNPACK #-} !Int
  }
  deriving stock (Read, Show, Eq)

instance Storable DLDataType where
  sizeOf _ = 4
  {-# INLINE sizeOf #-}
  alignment _ = 4
  {-# INLINE alignment #-}
  peek p =
    DLDataType
      <$> (toEnum . (fromIntegral :: Word8 -> Int) <$> peekByteOff p 0)
      <*> ((fromIntegral :: Word8 -> Int) <$> peekByteOff p 1)
      <*> ((fromIntegral :: Word16 -> Int) <$> peekByteOff p 2)
  {-# INLINE peek #-}
  poke p x = do
    pokeByteOff p 0 . (fromIntegral :: Int -> Word8) . fromEnum . dlDataTypeCode $ x
    pokeByteOff p 1 . (fromIntegral :: Int -> Word8) . dlDataTypeBits $ x
    pokeByteOff p 2 . (fromIntegral :: Int -> Word16) . dlDataTypeLanes $ x
  {-# INLINE poke #-}

data DLTensor = DLTensor
  { dlTensorData :: {-# UNPACK #-} !(Ptr ()),
    dlTensorDevice :: {-# UNPACK #-} !DLDevice,
    dlTensorNDim :: {-# UNPACK #-} !Int,
    dlTensorDType :: {-# UNPACK #-} !DLDataType,
    dlTensorShape :: {-# UNPACK #-} !(Ptr Int64),
    dlTensorStrides :: {-# UNPACK #-} !(Ptr Int64),
    dlTensorByteOffset :: {-# UNPACK #-} !Word64
  }
  deriving stock (Show, Eq)

instance Storable DLTensor where
  sizeOf _ = 8 + 8 + 4 + 4 + 8 + 8 + 8
  {-# INLINE sizeOf #-}
  alignment _ = 8
  {-# INLINE alignment #-}
  peek p =
    DLTensor
      <$> peekByteOff p 0
      <*> peekByteOff p 8
      <*> ((fromIntegral :: CInt -> Int) <$> peekByteOff p 16)
      <*> peekByteOff p 20
      <*> peekByteOff p 24
      <*> peekByteOff p 32
      <*> peekByteOff p 40
  {-# INLINE peek #-}
  poke p x = do
    pokeByteOff p 0 (dlTensorData x)
    pokeByteOff p 8 (dlTensorDevice x)
    pokeByteOff p 16 . (fromIntegral :: Int -> CInt) . dlTensorNDim $ x
    pokeByteOff p 20 (dlTensorDType x)
    pokeByteOff p 24 (dlTensorShape x)
    pokeByteOff p 32 (dlTensorStrides x)
    pokeByteOff p 40 (dlTensorByteOffset x)
  {-# INLINE poke #-}

class IsDLDataType a where
  dlDataTypeOf :: proxy a -> DLDataType

fromTag :: forall a proxy. Storable a => DLDataTypeCode -> proxy a -> DLDataType
fromTag tag _ =
  DLDataType {dlDataTypeCode = tag, dlDataTypeBits = (sizeOf element * 8), dlDataTypeLanes = 1}
  where
    element :: a
    element = undefined
{-# INLINE fromTag #-}

-- {{{ Integral

instance IsDLDataType CChar where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CSChar where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUChar where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CShort where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUShort where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CInt where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUInt where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CLong where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CULong where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CPtrdiff where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CSize where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CLLong where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CULLong where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CBool where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CIntPtr where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUIntPtr where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word8 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word16 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word32 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word64 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int8 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int16 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int32 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int64 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

-- }}}

-- {{{ Floating point

instance IsDLDataType CFloat where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Float where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CDouble where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Double where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

-- }}}

-- {{{ Complex

instance IsDLDataType (Complex CFloat) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType (Complex Float) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType (Complex CDouble) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType (Complex Double) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

-- }}}

class (Monad m) => IsDLTensor m a where
  withDLTensor :: a -> (DLTensor -> m b) -> m b

data ContiguousVectorView a
  = ContiguousVectorView {-# UNPACK #-} !Int {-# UNPACK #-} !(Ptr a)

data StridedVectorView a
  = StridedVectorView {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !(Ptr a)

dlExpectDType :: forall a. IsDLDataType a => DLTensor -> Either Text ()
dlExpectDType t
  | expected == got = Right ()
  | otherwise =
    Left . T.pack $
      "expected DLTensor to be of type " <> show expected
        <> ", but it is of type "
        <> show got
  where
    !expected = dlDataTypeOf (Proxy @a)
    !got = dlTensorDType t
{-# INLINE dlExpectDType #-}

dlExpectDevice :: DLDevice -> DLTensor -> Either Text ()
dlExpectDevice expected t
  | expected == got = Right ()
  | otherwise =
    Left . T.pack $
      "expected the tensor to be on " <> show expected <> ", but it is on "
        <> show got
  where
    !got = dlTensorDevice t
{-# INLINE dlExpectDevice #-}

dlExpectNDim :: Int -> DLTensor -> Either Text ()
dlExpectNDim expected t
  | expected == got = Right ()
  | otherwise =
    Left . T.pack $
      "expected a " <> show expected <> "-dimensional tensor, but got a "
        <> show got
        <> "-dimensional one"
  where
    !got = dlTensorNDim t
{-# INLINE dlExpectNDim #-}

dlTensorLength :: DLTensor -> Int
dlTensorLength t
  | rank > 0 = fromIntegral $ go (P.indexOffPtr (dlTensorShape t) 0) 1
  | otherwise = 0
  where
    !rank = dlTensorNDim t
    go !acc !i
      | i < rank = go (acc * P.indexOffPtr (dlTensorShape t) i) (i + 1)
      | otherwise = acc
{-# INLINE dlTensorLength #-}

dlExpectShape :: PrimArray Int -> DLTensor -> Either Text ()
dlExpectShape expected t = do
  () <- dlExpectNDim rank t
  let go !i
        | i < rank =
          indexPrimArray expected i == fromIntegral (P.indexOffPtr (dlTensorShape t) i)
            && go (i + 1)
        | otherwise = True
  if go 0
    then Right ()
    else
      Left . T.pack $
        "expected a tensor of shape " <> show expected
          <> ", but got a tensor of shape "
          <> show shape
  where
    !rank = sizeofPrimArray expected
    shape = generatePrimArray (dlTensorNDim t) (\i -> P.indexOffPtr (dlTensorShape t) i)
{-# INLINE dlExpectShape #-}

dlAsStridedVector :: forall a. IsDLDataType a => DLTensor -> Either Text (StridedVectorView a)
dlAsStridedVector t = do
  () <- dlExpectDType @a t
  () <- dlExpectNDim 1 t
  let !size = fromIntegral $ P.indexOffPtr (dlTensorShape t) 0
      !stride = fromIntegral $ P.indexOffPtr (dlTensorStrides t) 0
  pure $ StridedVectorView size stride (castPtr (dlTensorData t))
{-# INLINE dlAsStridedVector #-}

dlAsContiguousVector :: forall a. IsDLDataType a => DLTensor -> Either Text (ContiguousVectorView a)
dlAsContiguousVector t = do
  (StridedVectorView size stride ptr) <- dlAsStridedVector t
  if stride == 1
    then Right $ ContiguousVectorView size ptr
    else Left . T.pack $ "expected a contiguous DLTensor, but got one with stride " <> show stride
{-# INLINE dlAsContiguousVector #-}

-- dlRowMajorStride :: DLTensor -> Maybe Int
-- dlRowMajorStride t = let !f = go 0 in f
--   where
--     go !i
--       | i < rank - 1 =
--         let !n₂ = P.indexOffPtr shape (i + 1)
--             !s₁ = P.indexOffPtr stride i
--             !s₂ = P.indexOffPtr stride (i + 1)
--          in if s₁ `div` s₂ == n₁ then go (i + 1) else False
--       | otherwise = True

-- data ShapeInfo = ShapeInfo {shapeSizes :: {-# UNPACK #-} !(PrimArray Int), shapeStrides :: {-# UNPACK #-} !(PrimArray Int)}
--   deriving stock (Show, Eq)

dlRowMajorStride :: DLTensor -> Maybe Int
dlRowMajorStride t
  | rank == 0 = Just 1
  | otherwise = fromIntegral <$> go (P.indexOffPtr stride 0) 1
  where
    !rank = dlTensorNDim t
    !shape = dlTensorShape t
    !stride = dlTensorStrides t
    go !s !i
      | i < rank =
        let !n = P.indexOffPtr shape i
            !s' = P.indexOffPtr stride i
         in if s `div` s' == n then go s' (i + 1) else Nothing
      | otherwise = Just s

--
-- isColumnMajorLike :: Int -> Ptr Int64 -> Ptr Int64 -> Bool
-- isColumnMajorLike !rank !shape !stride = let !f = go 0 in f
--   where
--     go !i
--       | i < rank - 1 =
--         let !n₁ = P.indexOffPtr shape i
--             !s₁ = P.indexOffPtr stride i
--             !s₂ = P.indexOffPtr stride (i + 1)
--          in if s₂ `div` s₁ == n₁ then go (i + 1) else False
--       | otherwise = True

-- toList1D :: forall a. Prim a => DLTensor 'DLCPU 1 a -> [a]
-- toList1D t = runST $ go [] (stride * (extent - 1))
--   where
--     !extent = dlExtent t 0
--     !stride = dlStride t 0
--     !p = dlTensorData t `plusPtr` fromIntegral (dlTensorByteOffset t)
--     go :: PrimMonad m => [a] -> Int -> m [a]
--     go acc !i
--       | i >= 0 = do !x <- P.readOffPtr p i; go (x : acc) (i - stride)
--       | otherwise = pure acc
--
-- chunksOf :: HasCallStack => Int -> [a] -> [[a]]
-- chunksOf k xs = undefined

-- tensorToFlatList :: forall a device rank. Prim a => DLTensor device rank a -> [a]
-- tensorToFlatList t = reverse $ runIdentity $ foldN 0 shape strides combine []
--   where
--     buildList p = (\i -> fromIntegral $ P.indexOffPtr p i) <$> [0 .. dlTensorNDim t - 1]
--     strides = buildList (dlTensorStrides t)
--     shape = buildList (dlTensorShape t)
--     combine xs !i =
--       let !p = dlTensorData t `plusPtr` fromIntegral (dlTensorByteOffset t)
--           !x = P.indexOffPtr (castPtr p) i
--        in return (x : xs)

withListN :: (Prim a, PrimMonad m) => Int -> [a] -> (Ptr a -> m b) -> m b
withListN n xs action = do
  array <- newAlignedPinnedPrimArray n
  let go i []
        | i == n = return ()
        | otherwise =
          error $
            "list is shorter than indicated: " <> show i <> " < " <> show n
      go !i (y : ys)
        | i < n = writePrimArray array i y >> go (i + 1) ys
        | otherwise =
          error $
            "list is shorter than indicated: " <> show (length xs) <> " > " <> show n
  go 0 xs
  action (mutablePrimArrayContents array)

viaContiguousBuffer ::
  (IsDLDataType a, PrimMonad m) =>
  Ptr a ->
  [Int] ->
  [Int] ->
  (DLTensor -> m b) ->
  m b
viaContiguousBuffer dataPtr shape strides action =
  withListN rank (fromIntegral <$> shape) $ \shapePtr ->
    withListN rank (fromIntegral <$> strides) $ \stridesPtr ->
      action
        DLTensor
          { dlTensorData = castPtr dataPtr,
            dlTensorDevice = DLDevice DLCPU 0,
            dlTensorNDim = rank,
            dlTensorDType = dlDataTypeOf dataPtr,
            dlTensorShape = shapePtr,
            dlTensorStrides = stridesPtr,
            dlTensorByteOffset = 0
          }
  where
    !rank = length shape
{-# INLINE viaContiguousBuffer #-}

class
  (Monad m, KnownNat d, IsDLDataType (ListElement a)) =>
  ListIsDLTensor (m :: Type -> Type) (d :: Nat) (a :: Type)
  where
  listWithDLTensor :: proxy d -> a -> (DLTensor -> m b) -> m b

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

listShape3 :: [[[a]]] -> (Int, Int, Int)
listShape3 [] = (0, 0, 0)
listShape3 xs@(x : _) = let (d₁, d₂) = listShape2 x in (length xs, d₁, d₂)
{-# INLINE listShape3 #-}

listShape4 :: [[[[a]]]] -> (Int, Int, Int, Int)
listShape4 [] = (0, 0, 0, 0)
listShape4 xs@(x : _) = let (d₁, d₂, d₃) = listShape3 x in (length xs, d₁, d₂, d₃)
{-# INLINE listShape4 #-}

listShape5 :: [[[[[a]]]]] -> (Int, Int, Int, Int, Int)
listShape5 [] = (0, 0, 0, 0, 0)
listShape5 xs@(x : _) = let (d₁, d₂, d₃, d₄) = listShape4 x in (length xs, d₁, d₂, d₃, d₄)
{-# INLINE listShape5 #-}

instance
  (PrimMonad m, ListElement [a] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 1 [a]
  where
  listWithDLTensor _ xs action = withListN d₀ xs $ \ptr ->
    viaContiguousBuffer ptr [d₀] [1] action
    where
      d₀ = length xs

instance
  (PrimMonad m, ListElement [[a]] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 2 [[a]]
  where
  listWithDLTensor _ xs action =
    withListN (d₀ * d₁) (mconcat xs) $ \ptr ->
      viaContiguousBuffer ptr [d₀, d₁] [d₁, 1] action
    where
      (d₀, d₁) = listShape2 xs

instance
  (PrimMonad m, ListElement [[[a]]] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 3 [[[a]]]
  where
  listWithDLTensor _ xs action =
    withListN (d₀ * d₁ * d₂) (mconcat . mconcat $ xs) $ \ptr ->
      viaContiguousBuffer ptr [d₀, d₁, d₂] [d₁ * d₂, d₂, 1] action
    where
      (d₀, d₁, d₂) = listShape3 xs

instance
  (PrimMonad m, ListElement [[[[a]]]] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 4 [[[[a]]]]
  where
  listWithDLTensor _ xs action =
    withListN (d₀ * d₁ * d₂ * d₃) (mconcat . mconcat . mconcat $ xs) $ \ptr ->
      viaContiguousBuffer ptr [d₀, d₁, d₂, d₃] [d₁ * d₂ * d₃, d₂ * d₃, d₃, 1] action
    where
      (d₀, d₁, d₂, d₃) = listShape4 xs

type family IsFinal a where
  IsFinal [a] = 'False
  IsFinal a = 'True

type family ListElement' a b where
  ListElement' a 'True = a
  ListElement' [a] 'False = ListElement' a (IsFinal a)

type ListElement (a :: Type) = ListElement' a (IsFinal a)

type family ListDimension' (n :: Nat) (a :: Type) where
  ListDimension' n [a] = ListDimension' (n + 1) a
  ListDimension' n a = n

type ListDimension (a :: Type) = ListDimension' 0 a

instance (Monad m, ListIsDLTensor m (ListDimension [t]) [t]) => IsDLTensor m [t] where
  withDLTensor = listWithDLTensor (Proxy :: Proxy (ListDimension [t]))

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

-- instance Prim a => GHC.IsList (DLTensor 'DLCPU 1 a) where
--   type Item (DLTensor 'DLCPU 1 a) = a
--   toList = toList1D
--   fromList _ = error "IsList instance of DLTensor does not implement toList"

-- instance Prim a => GHC.IsList (DLTensor 'DLCPU 2 a) where
--   type Item (DLTensor 'DLCPU 2 a) = [a]
--   toList t = case dlFlatten t of
--     Just flat -> chunksOf k . toList1D $ flat
--     Nothing -> error "IsList instance of DLTensor does not support weirdly strided tensors"
--     where
--       !k = dlExtent t 1
--   fromList _ = error "IsList instance of DLTensor does not implement toList"

-- instance Prim a => GHC.IsList (DLTensor 'DLCPU 3 a) where
--   type Item (DLTensor 'DLCPU 3 a) = [[a]]
--   toList t = case dlFlatten t of
--     Just flat -> chunksOf d₁ . chunksOf d₂ . toList1D $ flat
--     Nothing -> error "IsList instance of DLTensor does not support weirdly strided tensors"
--     where
--       !d₁ = dlExtent t 1
--       !d₂ = dlExtent t 2
--   fromList _ = error "IsList instance of DLTensor does not implement toList"

-- instance (IsDLDataType a, Storable a) => IsDLTensor IO 'DLCPU 1 a [a] where
--   withDLTensor xs action = withArrayLen xs $ \size dataPtr ->
--     viaContiguousBuffer dataPtr [size] [1] action
--   {-# INLINE withDLTensor #-}
--
-- instance (IsDLDataType a, Storable a) => IsDLTensor IO 'DLCPU 2 a [[a]] where
--   withDLTensor l action = do
--     let (!n, !m) = listShape2D l
--     withArrayLen (mconcat l) $ \size dataPtr -> do
--       when (n * m /= size) $ error "list has wrong shape"
--       viaContiguousBuffer dataPtr [n, m] [m, 1] action
--   {-# INLINE withDLTensor #-}
--
-- instance (IsDLDataType a, Storable a) => IsDLTensor IO 'DLCPU 3 a [[[a]]] where
--   withDLTensor l action = do
--     let (!d₁, !d₂, !d₃) = listShape3D l
--     withArrayLen (mconcat (mconcat l)) $ \size dataPtr -> do
--       when (d₁ * d₂ * d₃ /= size) $ error "list has wrong shape"
--       viaContiguousBuffer dataPtr [d₁, d₂, d₃] [d₂ * d₃, d₃, 1] action
--   {-# INLINE withDLTensor #-}


-- | Common infrastructure for CVode/ARKode
{-# OPTIONS_GHC -Wno-name-shadowing #-}
module Numeric.Sundials.Common where

import Foreign.C.Types
import Foreign.Ptr
import Foreign.Storable (peek, poke)
import Foreign.C.String
import Numeric.Sundials.Types
import qualified Numeric.Sundials.Foreign as T
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import Data.Maybe
import Numeric.LinearAlgebra.HMatrix hiding (Vector)
import GHC.Prim
import Control.Monad.IO.Class
import Control.Monad.Cont
import Control.Exception
import Katip
import Data.Aeson
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.ByteString as BS
import Control.Monad.Reader
import GHC.Generics (Generic)
import Foreign.ForeignPtr

-- | A collection of variables that we allocate on the Haskell side and
-- pass into the C code to be filled.
data CVars vec = CVars
  { c_diagnostics :: vec SunIndexType
    -- ^ Mutable vector to which we write diagnostic data while
    -- solving. Its size corresponds to the number of fields in
    -- 'SundialsDiagnostics'.
  , c_root_info :: vec CInt
    -- ^ Just a temporary vector (of the size equal to the number of event
    -- specs) that we use to get root info. Isn't used for output.
  , c_event_index :: vec CInt
    -- ^ For each event occurrence, this indicates which of the events
    -- occurred. Size: max_num_events.
  , c_event_time :: vec CDouble
    -- ^ For each event occurrence, this indicates the time of the
    -- occurrence. Size: max_num_events.
  , c_n_events :: vec CInt
    -- ^ Vector of size 1 that gives the total number of events occurred.
  , c_n_rows :: vec CInt
    -- ^ The total number of rows in the output matrix.
  , c_output_mat :: vec CDouble
    -- ^ The output matrix stored in the row-major order.
    -- Dimensions: (1 + dim) * (2 * max_events + nTs).
  , c_actual_event_direction :: vec CInt
    -- ^ Vector of size max_num_events that gives the direction of the
    -- occurred event.
  , c_local_error :: vec CDouble
    -- ^ Vector containing local error estimates. Size: the dimensionality
    -- of the system.
  , c_var_weight :: vec CDouble
    -- ^ Vector containing variable weights (derived from the tolerances).
    -- Size: the dimensionality of the system.
  , c_local_error_set :: vec CInt
    -- The flag (size 1) indicating whether c_local_error is filled with meaningful
    -- values. *Should be initialized with 0.*
  }

allocateCVars :: OdeProblem -> IO (CVars (VS.MVector RealWorld))
allocateCVars OdeProblem{..} = do
  let dim = VS.length odeInitCond
  c_diagnostics <- VSM.new 11
  c_root_info <- VSM.new $ V.length odeEvents
  c_event_index <- VSM.new odeMaxEvents
  c_event_time <- VSM.new odeMaxEvents
  c_actual_event_direction <- VSM.new odeMaxEvents
  c_n_events <- VSM.new 1
  c_n_rows <- VSM.new 1
  c_local_error <- VSM.new dim
  c_var_weight <- VSM.new dim
  c_local_error_set <- VSM.new 1
  VSM.write c_local_error_set 0 0
  c_output_mat <- VSM.new $
    (1 + dim) * (2 * odeMaxEvents + VS.length odeSolTimes)
  return CVars {..}

-- NB: the mutable CVars must not be used after this
freezeCVars :: CVars (VS.MVector RealWorld) -> IO (CVars VS.Vector)
freezeCVars CVars{..} = do
  c_diagnostics <- VS.unsafeFreeze c_diagnostics
  c_root_info <- VS.unsafeFreeze c_root_info
  c_event_index <- VS.unsafeFreeze c_event_index
  c_event_time <- VS.unsafeFreeze c_event_time
  c_actual_event_direction <- VS.unsafeFreeze c_actual_event_direction
  c_n_events <- VS.unsafeFreeze c_n_events
  c_n_rows <- VS.unsafeFreeze c_n_rows
  c_output_mat <- VS.unsafeFreeze c_output_mat
  c_local_error <- VS.unsafeFreeze c_local_error
  c_var_weight <- VS.unsafeFreeze c_var_weight
  c_local_error_set <- VS.unsafeFreeze c_local_error_set
  return CVars {..}

-- | Similar to 'CVars', except these are immutable values that are
-- accessed (read-only) by the C code and specify the system to be solved.
data CConsts = CConsts
  { c_dim :: SunIndexType -- ^ the dimensionality (number of variables/equations)
  , c_method :: CInt -- ^ the ODE method (specific to the solver)
  , c_n_sol_times :: CInt
  , c_sol_time :: VS.Vector CDouble
  , c_init_cond :: VS.Vector CDouble
  , c_rhs :: FunPtr OdeRhsCType
  , c_rhs_userdata :: Ptr UserData
  , c_rtol :: CDouble
  , c_atol :: VS.Vector CDouble
  , c_n_event_specs :: CInt
  , c_event_fn :: CDouble -> Ptr T.SunVector -> Ptr CDouble -> Ptr () -> IO CInt
  , c_apply_event
      :: CInt -- number of triggered events
      -> Ptr CInt -- event indices
      -> CDouble -- time
      -> Ptr T.SunVector -- y
      -> Ptr T.SunVector -- new y
      -> Ptr CInt -- (out) stop the solver?
      -> Ptr CInt -- (out) record the event?
      -> IO CInt
  , c_jac_set :: CInt
  , c_jac :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunMatrix
          -> Ptr () -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr T.SunVector
          -> IO CInt
  , c_sparse_jac :: CInt
      -- ^ If 0, use a dense matrix.
      -- If non-0, use a sparse matrix with that number of non-zero
      -- elements.
  , c_requested_event_direction :: VS.Vector CInt
  , c_next_time_event :: IO CDouble
  , c_max_events :: CInt
  , c_minstep :: CDouble
  , c_fixedstep :: CDouble
  , c_max_n_steps :: SunIndexType
  , c_max_err_test_fails :: CInt
  , c_init_step_size_set :: CInt
  , c_init_step_size :: CDouble
  }

data Solver = CVode | ARKode
  deriving Show

class Method method where
  methodToInt :: method -> CInt
  methodSolver :: Solver

withCConsts
  :: Method method
  => ODEOpts method
  -> OdeProblem
  -> (CConsts -> IO r)
  -> IO r
withCConsts ODEOpts{..} OdeProblem{..} = runContT $ do
  let
    dim = VS.length c_init_cond
    c_init_cond = coerce odeInitCond
    c_dim = fromIntegral dim
    c_n_sol_times = fromIntegral . VS.length $ odeSolTimes
    c_sol_time = coerce odeSolTimes
    c_rtol = relTolerance odeTolerances
    c_atol = either (VS.replicate dim) id $ absTolerances odeTolerances
    c_minstep = coerce minStep
    c_fixedstep = coerce fixedStep
    c_max_n_steps = fromIntegral maxNumSteps
    c_max_err_test_fails = fromIntegral maxFail
    c_init_step_size_set = fromIntegral . fromEnum $ isJust initStep
    c_init_step_size = coerce . fromMaybe 0 $ initStep
    c_n_event_specs = fromIntegral $ V.length odeEvents
    c_requested_event_direction = V.convert $ V.map (directionToInt . eventDirection) odeEvents
    c_event_fn t y_ptr out_ptr _ptr = do
      y <- sunVecVals <$> peek y_ptr
      let vals = V.convert $ V.map (\ev -> coerce (eventCondition ev) t y) odeEvents
      -- FIXME: We should be able to use poke somehow
      T.vectorToC vals (fromIntegral c_n_event_specs) out_ptr
      return 0
    c_apply_event n_events event_indices_ptr t y_ptr y'_ptr stop_solver_ptr record_event_ptr = do
      event_indices <- vecFromPtr event_indices_ptr (fromIntegral n_events)
      -- Apparently there's no safe version of
      y_vec <- peek y_ptr
      EventHandlerResult{..} <-
        odeEventHandler
          (coerce t :: Double)
          (coerce $ sunVecVals y_vec :: VS.Vector Double)
          (VS.map fromIntegral event_indices :: VS.Vector Int)
      poke y'_ptr $ SunVector
        { sunVecN = sunVecN y_vec
        , sunVecVals = coerce eventNewState
        }
      poke stop_solver_ptr . fromIntegral $ fromEnum eventStopSolver
      poke record_event_ptr . fromIntegral $ fromEnum eventRecord
      return 0
    c_max_events = fromIntegral odeMaxEvents
    c_next_time_event = coerce odeTimeBasedEvents
    c_jac_set = fromIntegral . fromEnum $ isJust odeJacobian
    c_jac t y _fy jacS _ptr _tmp1 _tmp2 _tmp3 = do
      case odeJacobian of
        Nothing   -> undefined
        Just jacI -> do j <- matrixToSunMatrix . jacI (coerce t) <$> (coerce $ sunVecVals <$> peek y)
                        case jacobianRepr of
                          DenseJacobian -> poke jacS j
                          SparseJacobian spat -> poke (castPtr jacS) (T.SparseMatrix spat j)
                        return 0
    c_sparse_jac = case jacobianRepr of
      SparseJacobian (T.SparsePattern spat) -> VS.sum (VS.map fromIntegral spat)
      DenseJacobian -> 0
    c_method = methodToInt odeMethod

  (c_rhs, c_rhs_userdata) <-
    case odeRhs of
      OdeRhsC ptr u -> return (ptr, u)
      OdeRhsHaskell fun -> do
        let
          funIO :: CDouble -> Ptr T.SunVector -> Ptr T.SunVector -> Ptr UserData -> IO CInt
          funIO t y f _ptr = do
            sv <- peek y
            r <- fun t (sunVecVals sv)
            poke f $ SunVector { sunVecN = sunVecN sv
                               , sunVecVals = r
                               }
            return 0
        funptr <- ContT $ bracket (mkOdeRhsC funIO) freeHaskellFunPtr
        return (funptr, nullPtr)
  return CConsts{..}

matrixToSunMatrix :: Matrix Double -> T.SunMatrix
matrixToSunMatrix m = T.SunMatrix { T.rows = nr, T.cols = nc, T.vals = vs }
  where
    nr = fromIntegral $ rows m
    nc = fromIntegral $ cols m
    vs = coerce . VS.concat $ toColumns m

-- Contrary to the documentation, it appears that CVodeGetRootInfo
-- may use both 1 and -1 to indicate a root, depending on the
-- direction of the sign change. See near the end of cvRootfind.
intToDirection :: Integral d => d -> Maybe CrossingDirection
intToDirection d =
  case d of
    1  -> Just Upwards
    -1 -> Just Downwards
    _  -> Nothing

-- | Almost inverse of 'intToDirection'. Map 'Upwards' to 1, 'Downwards' to
-- -1, and 'AnyDirection' to 0.
directionToInt :: Integral d => CrossingDirection -> d
directionToInt d =
  case d of
    Upwards -> 1
    Downwards -> -1
    AnyDirection -> 0

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

assembleSolverResult
  :: OdeProblem
  -> CInt
  -> CVars VS.Vector
  -> IO (Either ErrorDiagnostics SundialsSolution)
assembleSolverResult OdeProblem{..} ret CVars{..} = do
  let
    dim = VS.length odeInitCond
    n_rows = fromIntegral . VS.head $ c_n_rows
    output_mat = coerce . reshape (dim + 1) . subVector 0 ((dim + 1) * n_rows) $ c_output_mat
    (local_errors, var_weights) =
      if c_local_error_set VS.! 0 == 0
        then (mempty, mempty)
        else coerce (c_local_error, c_var_weight)
    diagnostics = SundialsDiagnostics
      (fromIntegral $ c_diagnostics VS.!0)
      (fromIntegral $ c_diagnostics VS.!1)
      (fromIntegral $ c_diagnostics VS.!2)
      (fromIntegral $ c_diagnostics VS.!3)
      (fromIntegral $ c_diagnostics VS.!4)
      (fromIntegral $ c_diagnostics VS.!5)
      (fromIntegral $ c_diagnostics VS.!6)
      (fromIntegral $ c_diagnostics VS.!7)
      (fromIntegral $ c_diagnostics VS.!8)
      (fromIntegral $ c_diagnostics VS.!9)
      (toEnum . fromIntegral $ c_diagnostics VS.! 10)
  return $
    if ret == T.cV_SUCCESS
      then
        Right $ SundialsSolution
          { actualTimeGrid = extractTimeGrid output_mat
          , solutionMatrix = dropTimeGrid output_mat
          , diagnostics
          }
      else
        Left ErrorDiagnostics
          { partialResults = output_mat
          , errorCode = fromIntegral ret
          , errorEstimates = local_errors
          , varWeights = var_weights
          }
  where
    -- The time grid is the first column of the result matrix
    extractTimeGrid :: Matrix Double -> VS.Vector Double
    extractTimeGrid = head . toColumns
    dropTimeGrid :: Matrix Double -> Matrix Double
    dropTimeGrid = fromColumns . tail . toColumns

-- | The common solving logic between ARKode and CVode
solveCommon
  :: (Method method, Katip m)
  => (CConsts -> CVars (VS.MVector RealWorld) -> ReportErrorFn -> IO CInt)
      -- ^ the CVode/ARKode solving function; mostly inline-C code
  -> ODEOpts method
  -> OdeProblem
  -> m (Either ErrorDiagnostics SundialsSolution)
solveCommon solve_c opts problem@(OdeProblem{..})

  | VS.null odeInitCond = -- 0-dimensional (empty) system

    return . Right $ SundialsSolution
      { actualTimeGrid = odeSolTimes
      , solutionMatrix = (VS.length odeSolTimes >< 0) []
      , diagnostics = emptyDiagnostics
      }

  | otherwise = do

    report_error <- logWithKatip
    liftIO $ do -- the rest is in the IO monad
    vars <- allocateCVars problem
    ret <- withCConsts opts problem $ \consts ->
      solve_c consts vars report_error
    frozenVars <- freezeCVars vars
    assembleSolverResult problem ret frozenVars

-- | An auxiliary function to construct a storable vector from a C pointer
-- and length.
--
-- There doesn't seem to be a safe version of 'VS.unsafeFromForeignPtr0',
-- nor a way to clone an immutable vector, so we emulate it via an
-- intermediate mutable vector.
vecFromPtr
  :: VS.Storable a
  => Ptr a
  -> Int
  -> IO (VS.Vector a)
vecFromPtr ptr n = do
  fptr <- newForeignPtr_ ptr
  let mv = VSM.unsafeFromForeignPtr0 fptr n
  VS.freeze mv -- this does the copying and makes the whole thing safe

----------------------------------------------------------------------
--                           Logging
----------------------------------------------------------------------

-- | The Katip payload for logging Sundials errors
data SundialsErrorContext = SundialsErrorContext
  { sundialsErrorCode :: !Int
  , sundialsErrorModule :: !T.Text
  , sundialsErrorFunction :: !T.Text
  } deriving Generic
instance ToJSON SundialsErrorContext
instance ToObject SundialsErrorContext
instance LogItem SundialsErrorContext where
  payloadKeys _ _ = AllKeys

type ReportErrorFn =
  (  CInt    -- error code
  -> CString -- module name
  -> CString -- function name
  -> CString -- the message
  -> Ptr ()  -- user data (ignored)
  -> IO ()
  )

logWithKatip
  :: Katip m
  => m ReportErrorFn
logWithKatip = do
  log_env <- getLogEnv
  return $
    \err_code c_mod_name c_func_name c_msg _userdata ->
    -- See Note [CV_TOO_CLOSE]
    if err_code == T.cV_TOO_CLOSE then pure () else do
    let
      toText :: CString -> IO T.Text
      toText = fmap T.decodeUtf8 . BS.packCString
    mod_name <- toText c_mod_name
    func_name <- toText c_func_name
    msg <- toText c_msg
    let
      severity :: Severity
      severity =
        if err_code <= 0
          then ErrorS
          else WarningS
      errCtx :: SundialsErrorContext
      errCtx = SundialsErrorContext
        { sundialsErrorCode = fromIntegral err_code
        , sundialsErrorModule = mod_name
        , sundialsErrorFunction = func_name
        }
    flip runReaderT log_env . unKatipT $ do
      logF errCtx "sundials" severity (logStr msg)


module Levity.Class.Functor
  ( FunctorL(..)
  ) where

import Control.Applicative (ZipList)
import Control.Exception (Handler)
import Control.Monad.ST (ST)
import Data.Complex (Complex)
import Data.Functor.Const (Const)
import Data.Functor.Identity (Identity)
import Data.List.NonEmpty (NonEmpty)
import Data.Monoid qualified as Monoid
import Data.Proxy (Proxy)
import Data.Semigroup qualified as Semigroup
import GHC.Base (Type)
import GHC.Conc (STM)
import GHC.Exts
import GHC.Generics qualified as Generic
import Prelude
import System.Console.GetOpt (ArgDescr, OptDescr, ArgOrder)
import Text.ParserCombinators.ReadP (ReadP)
import Text.ParserCombinators.ReadPrec (ReadPrec)

import Levity.Types.Base

type FunctorL :: (Type -> TYPE r) -> Constraint
class FunctorL f where
  {-# minimal fmapL #-}
  fmapL :: (a -> b) -> (f a -> f b)

instance Prelude.Functor f => FunctorL (Base1 f) where
  fmapL = Prelude.fmap

-- Manually-defined instances are because of: https://gitlab.haskell.org/ghc/ghc/-/issues/20680
deriving via (Base1 []) instance FunctorL []
deriving via (Base1 Prelude.Maybe) instance FunctorL Prelude.Maybe
deriving via (Base1 Prelude.IO) instance FunctorL Prelude.IO
deriving via (Base1 Generic.Par1) instance FunctorL Generic.Par1
deriving via (Base1 NonEmpty) instance FunctorL NonEmpty
deriving via (Base1 ReadP) instance FunctorL ReadP
deriving via (Base1 ReadPrec) instance FunctorL ReadPrec
deriving via (Base1 Down) instance FunctorL Down
deriving via (Base1 (Semigroup.Arg a)) instance FunctorL (Semigroup.Arg a)
deriving via (Base1 (ST s)) instance FunctorL (ST s)
deriving via (Base1 ((,) a)) instance FunctorL ((,) a)
deriving via (Base1 (Prelude.Either a)) instance FunctorL (Prelude.Either a)
deriving via (Base1 Complex) instance FunctorL Complex
deriving via (Base1 Semigroup.Min) instance FunctorL Semigroup.Min
deriving via (Base1 Semigroup.Max) instance FunctorL Semigroup.Max
deriving via (Base1 Semigroup.First) instance FunctorL Semigroup.First
deriving via (Base1 Semigroup.Last) instance FunctorL Semigroup.Last
deriving via (Base1 ArgOrder) instance FunctorL ArgOrder
deriving via (Base1 OptDescr) instance FunctorL OptDescr
deriving via (Base1 ArgDescr) instance FunctorL ArgDescr
deriving via (Base1 Identity) instance FunctorL Identity
deriving via (Base1 Handler) instance FunctorL Handler
deriving via (Base1 STM) instance FunctorL STM
deriving via (Base1 Monoid.First) instance FunctorL Monoid.First
deriving via (Base1 Monoid.Last) instance FunctorL Monoid.Last
deriving via (Base1 Monoid.Dual) instance FunctorL Monoid.Dual
deriving via (Base1 Monoid.Sum) instance FunctorL Monoid.Sum
deriving via (Base1 Monoid.Product) instance FunctorL Monoid.Product
deriving via (Base1 ZipList) instance FunctorL ZipList

deriving via (Base1 Generic.U1) instance FunctorL Generic.U1
deriving via (Base1 Generic.V1) instance FunctorL Generic.V1
instance (Prelude.Functor f) => FunctorL (Generic.Rec1 f) where { fmapL = Prelude.fmap }
deriving via (Base1 (Generic.URec (Ptr ()))) instance FunctorL (Generic.URec (Ptr ()))
deriving via (Base1 (Generic.URec Char)) instance FunctorL (Generic.URec Char)
deriving via (Base1 (Generic.URec Double)) instance FunctorL (Generic.URec Double)
deriving via (Base1 (Generic.URec Float)) instance FunctorL (Generic.URec Float)
deriving via (Base1 (Generic.URec Int)) instance FunctorL (Generic.URec Int)
deriving via (Base1 (Generic.K1 i c)) instance FunctorL (Generic.K1 i c)
deriving via (Base1 Proxy) instance FunctorL Proxy
deriving via (Base1 (Const m)) instance FunctorL (Const m)
instance (Functor f) => FunctorL (Monoid.Alt f) where { fmapL = fmap }
instance (Functor f) => FunctorL (Monoid.Ap f) where { fmapL = fmap }
deriving via (Base1 ((->) r)) instance FunctorL ((->) r)
instance (Functor f, Functor g) => FunctorL (f Generic.:+: g) where { fmapL = fmap }
instance (Functor f, Functor g) => FunctorL (f Generic.:*: g) where { fmapL = fmap }
instance FunctorL ((# , #) a) where { fmapL f (# a, b #) = (# a, f b #) }
instance FunctorL ((# ,, #) a b) where { fmapL f (# a, b, c #) = (# a, b, f c #) }


{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}

module Main where

import Control.Monad.Except
  ( ExceptT,
    MonadError (..),
    MonadIO (liftIO),
    MonadTrans (lift),
    runExceptT,
  )
import Data.Array (Array, listArray)
import Data.Complex (Complex (..), imagPart, realPart)
import Data.Functor ((<&>))
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Maybe (isJust, isNothing)
import Data.Ratio (denominator, numerator, (%))
import Numeric (readFloat, readHex, readOct)
import System.Environment (getArgs)
import System.IO
  ( Handle,
    IOMode (ReadMode, WriteMode),
    hClose,
    hFlush,
    hGetLine,
    hPrint,
    hPutStr,
    openFile,
    stderr,
    stdin,
    stdout,
  )
import Text.ParserCombinators.Parsec
  ( ParseError,
    Parser,
    alphaNum,
    anyChar,
    between,
    char,
    digit,
    endBy,
    hexDigit,
    letter,
    many,
    many1,
    noneOf,
    notFollowedBy,
    octDigit,
    oneOf,
    option,
    parse,
    sepBy,
    sepEndBy,
    skipMany,
    skipMany1,
    space,
    string,
    try,
    (<|>),
  )

arrayList :: Array i e -> [e]
arrayList = foldr (:) []

until_ :: Monad m => (a -> Bool) -> m a -> (a -> m ()) -> m ()
until_ pred prompt action = do
  result <- prompt
  if pred result
    then return ()
    else action result >> until_ pred prompt action

main :: IO ()
main = do
  args <- getArgs
  if null args then runRepl else runOne args

runOne :: [String] -> IO ()
runOne args = do
  env <-
    primitiveBindings
      >>= flip bindVars [("args", List $ map String $ drop 1 args)]
  runIOThrows (show <$> eval env (List [Atom "load", String $ head args]))
    >>= hPutStr stderr

runRepl :: IO ()
runRepl =
  primitiveBindings
    >>= until_ (== "quit") (readPrompt "Lisp>>> ") . evalAndPrint

flushStr :: String -> IO ()
flushStr str = putStr str >> hFlush stdout

readPrompt :: String -> IO String
readPrompt prompt = flushStr prompt >> getLine

evalString :: Env -> String -> IO String
evalString env expr =
  runIOThrows $ show <$> (liftThrows (readExpr expr) >>= eval env)

evalAndPrint :: Env -> String -> IO ()
evalAndPrint env expr = evalString env expr >>= putStrLn

spaces :: Parser ()
spaces = skipMany1 space

symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~"

type Env = IORef [(String, IORef LispVal)]

nullEnv :: IO Env
nullEnv = newIORef []

primitiveBindings :: IO Env
primitiveBindings =
  nullEnv
    >>= flip
      bindVars
      ( map (makeFunc PrimitiveFunc) primitives
          ++ map (makeFunc IOFunc) ioPrimitives
      )
  where
    makeFunc constructor (var, func) = (var, constructor func)

readOrThrow :: Parser a -> String -> ThrowsError a
readOrThrow parser input = case parse parser "lisp" input of
  Left err -> throwError $ Parser err
  Right val -> return val

readExpr :: String -> ThrowsError LispVal
readExpr = readOrThrow parseExpr

readExprList :: String -> ThrowsError [LispVal]
readExprList = readOrThrow (endBy parseExpr spaces)

data LispVal
  = Nil
  | Atom String
  | List [LispVal]
  | DottedList [LispVal] LispVal
  | Vector (Array Int LispVal)
  | Number Integer
  | Float Double
  | Ratio Rational
  | Complex (Complex Double)
  | String String
  | Char Char
  | Bool Bool
  | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
  | Func
      { params :: [String],
        varargs :: Maybe String,
        body :: [LispVal],
        closure :: Env
      }
  | IOFunc ([LispVal] -> IOThrowsError LispVal)
  | Port Handle

dotted2List :: [LispVal] -> LispVal -> LispVal
dotted2List xs y = List $ xs ++ [y]

showVal :: LispVal -> String
showVal Nil = "nil"
showVal (Atom a) = a
showVal (Number n) = show n
showVal (Float f) = show f
showVal (Ratio r) = show (numerator r) ++ "/" ++ show (denominator r)
showVal (Complex c) = show (realPart c) ++ "+" ++ show (imagPart c) ++ "i"
showVal (String s) = "\"" ++ s ++ "\""
showVal (Char c) =
  "#\\" ++ case c of
    ' ' -> "space"
    '\n' -> "newline"
    _ -> [c]
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (List xs) = "(" ++ unwordsList xs ++ ")"
showVal (DottedList heads tail) =
  "(" ++ unwordsList heads ++ " . " ++ showVal tail ++ ")"
showVal (Vector v) = "#(" ++ unwordsList (foldr (:) [] v) ++ ")"
showVal (PrimitiveFunc _) = "<primitive>"
showVal (Func params varargs body _) =
  "(lambda ("
    ++ unwords (map show params)
    ++ ( case varargs of
           Nothing -> ""
           Just arg -> " . " ++ arg
       )
    ++ ") ...)"
showVal (Port _) = "<IO port>"
showVal (IOFunc _) = "<IO primitive>"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where
  show = showVal

data LispError
  = NumArgs Integer [LispVal]
  | TypeMismatch String LispVal
  | Parser ParseError
  | BadSpecialForm String LispVal
  | NotFunction String String
  | UnboundVar String String
  | ExpectedAtom LispVal
  | Default String

showError :: LispError -> String
showError (NumArgs expected found) =
  "Expected " ++ show expected
    ++ " arg(s), found value(s): "
    ++ unwordsList found
showError (TypeMismatch expected found) =
  "Invalid type: expected " ++ expected
    ++ ", found "
    ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (UnboundVar message varname) = message ++ ": " ++ varname
showError (ExpectedAtom found) = "Expected atom, found: " ++ show found
showError (Default message) = message

instance Show LispError where
  show = showError

type ThrowsError = Either LispError

instance MonadFail ThrowsError where
  fail message = error $ "Unexpected error: " ++ message

trapError :: (MonadError e m, Show e) => m String -> m String
trapError action = catchError action $ return . show

extractValue :: ThrowsError a -> a
extractValue val' = let (Right val) = val' in val

type IOThrowsError = ExceptT LispError IO

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err) = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = runExceptT (trapError action) <&> extractValue

parseString :: Parser LispVal
parseString = do
  char '"'
  x <- many $ noneOf ['\\', '"'] <|> parseEscape
  char '"'
  return $ String x

escapes :: [(Char, Char)]
escapes =
  [ ('\'', '\''),
    ('"', '"'),
    ('\\', '\\'),
    ('a', '\a'),
    ('b', '\b'),
    ('f', '\f'),
    ('n', '\n'),
    ('r', '\r'),
    ('t', '\t'),
    ('v', '\v')
  ]

parseEscape :: Parser Char
parseEscape = do
  char '\\'
  c <- oneOf $ map fst escapes
  let (Just escaped) = lookup c escapes
  return escaped

parseChar :: Parser LispVal
parseChar =
  let alphabetic = do x <- anyChar; notFollowedBy alphaNum; return x
   in do
        string "#\\"
        x <-
          try $
            (string "space" >> return ' ')
              <|> (string "newline" >> return '\n')
              <|> alphabetic
        return $ Char x

parseAtom :: Parser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many $ letter <|> digit <|> symbol
  return $ Atom (first : rest)

parseBool :: Parser LispVal
parseBool = do
  char '#'
  (char 't' >> return (Bool True)) <|> (char 'f' >> return (Bool False))

parseNumber :: Parser LispVal
parseNumber = parseBin <|> parseOct <|> parseDec <|> parseHex

parseBin :: Parser LispVal
parseBin = do
  try $ string "#b"
  x <- many1 $ oneOf "01"
  let d = readBin x where readBin x = readBin' 0 x
      readBin' v [] = v
      readBin' v (x : xs) =
        let v' = v * 2 + (if x == '1' then 1 else 0)
         in readBin' v' xs
  return $ Number d

parseOct :: Parser LispVal
parseOct = do
  try $ string "#o"
  x <- many1 octDigit
  let [(d, _)] = readOct x
  return $ Number d

parseDec :: Parser LispVal
parseDec = parseDec' <|> parseDec''
  where
    parseDec' = Number . read <$> many1 digit
    parseDec'' = do
      try $ string "#d"
      parseDec'

parseHex :: Parser LispVal
parseHex = do
  try $ string "#x"
  x <- many1 hexDigit
  let [(d, _)] = readHex x
  return $ Number d

parseFloat :: Parser LispVal
parseFloat = do
  x <- many1 digit
  char '.'
  y <- many1 digit
  let [(d, _)] = readFloat (x ++ "." ++ y)
  return $ Float d

parseRatio :: Parser LispVal
parseRatio = do
  x <- many1 digit
  char '/'
  y <- many1 digit
  let d = read x % read y
  return $ Ratio d

parseComplex :: Parser LispVal
parseComplex =
  let getDouble :: Parser Double
      getDouble = do
        f <- option Nil (try parseFloat)
        case f of
          (Float f) ->
            return $ realToFrac f
          _ -> do
            (Number n) <- parseDec
            return $ fromIntegral n
   in do
        x <- getDouble
        char '+'
        y <- getDouble
        char 'i'
        return $ Complex (x :+ y)

parseList :: Parser LispVal
parseList = between beg end parseList'
  where
    beg = char '(' >> skipMany space
    end = skipMany space >> char ')'
    parseList' = do
      heads <- sepEndBy parseExpr spaces
      tail <- option Nil (char '.' >> spaces >> parseExpr)
      return $ case tail of
        Nil -> List heads
        _ -> DottedList heads tail

parseQuoted :: Parser LispVal
parseQuoted = do
  char '\''
  x <- parseExpr
  return $ List [Atom "quote", x]

parseQuasiquote :: Parser LispVal
parseQuasiquote = do
  char '`'
  x <- parseExpr
  return $ List [Atom "quasiquote", x]

parseUnquoted :: Parser LispVal
parseUnquoted = do
  char ','
  x <- parseExpr
  return $ List [Atom "unquoted", x]

parseVector :: Parser LispVal
parseVector = do
  string "#("
  arr <- sepBy parseExpr spaces
  char ')'
  let vec = listArray (0, length arr - 1) arr
  return $ Vector vec

parseExpr :: Parser LispVal
parseExpr =
  parseAtom
    <|> parseString
    <|> try parseComplex
    <|> try parseFloat
    <|> try parseRatio
    <|> try parseNumber
    <|> try parseBool
    <|> try parseChar
    <|> parseQuoted
    <|> parseQuasiquote
    <|> parseUnquoted
    <|> parseVector
    <|> parseList

eval :: Env -> LispVal -> IOThrowsError LispVal
eval env (Atom var) = getVar env var
eval _ (List [Atom "quote", val]) = return val
eval _ (List [Atom "quasiquote", val]) =
  throwError $ Default "` is not support"
eval _ (List [Atom "unquoted", val]) =
  throwError $ Default ", is not support"
eval env (List [Atom "if", pred, conseq, alt]) = do
  result <- eval env pred
  case result of
    Bool True -> eval env conseq
    Bool False -> eval env alt
    _ -> throwError $ TypeMismatch "boolean" result
eval env form@(List (Atom "cond" : clauses)) =
  if null clauses
    then
      throwError $
        BadSpecialForm "No true clause in cond expression" form
    else case head clauses of
      List [Atom "else", expr] -> eval env expr
      List [test, expr] ->
        eval env $
          List [Atom "if", test, expr, List (Atom "cond" : tail clauses)]
      _ -> throwError $ BadSpecialForm "Ill-formed cond expression" form
eval env form@(List (Atom "case" : key : clauses)) =
  if null clauses
    then
      throwError $
        BadSpecialForm "No true clause in case expression: " form
    else case head clauses of
      List (Atom "else" : exprs) -> eval' env exprs
      List (List datums : exprs) -> do
        result <- eval env key
        let equality =
              any
                ( \x ->
                    let (Right (Bool res)) =
                          eqv [result, x]
                     in res
                )
                datums
        if equality
          then eval' env exprs
          else eval env $ List (Atom "case" : key : tail clauses)
      _ -> throwError $ BadSpecialForm "Ill-formed case expression" form
  where
    eval' env exprs = mapM (eval env) exprs <&> last
eval env (List [Atom "set!", Atom var, form]) =
  eval env form >>= setVar env var
eval env (List [Atom "define", Atom var, form]) =
  eval env form >>= defineVar env var
eval env (List (Atom "define" : List (Atom var : params) : body)) =
  makeNormalFunc env params body >>= defineVar env var
eval
  env
  ( List
      (Atom "define" : DottedList (Atom var : params) varargs : body)
    ) = makeVarArgsFunc varargs env params body >>= defineVar env var
eval env (List (Atom "lambda" : List params : body)) =
  makeNormalFunc env params body
eval env (List (Atom "lambda" : DottedList params varargs : body)) =
  makeVarArgsFunc varargs env params body
eval env (List (Atom "lambda" : varargs : body)) =
  makeVarArgsFunc varargs env [] body
eval env (List [Atom "load", String filename]) =
  load filename >>= fmap last . mapM (eval env)
eval env (List (func' : args')) = do
  func <- eval env func'
  args <- mapM (eval env) args'
  apply func args
eval _ val@(List _) =
  throwError $ BadSpecialForm "Unrecognized special form" val
eval _ val@(DottedList _ _) = throwError $ Default ". is not support"
eval _ val = return val

apply :: LispVal -> [LispVal] -> IOThrowsError LispVal
apply (Func params varargs body closure) args =
  if length params /= length args && isNothing varargs
    then throwError $ NumArgs (toInteger $ length params) args
    else
      liftIO (bindVars closure $ zip params args)
        >>= bindVarArgs varargs
        >>= evalBody
  where
    evalBody env = last <$> mapM (eval env) body
    bindVarArgs arg env = case arg of
      Just argName -> liftIO $ bindVars env [(argName, List remainingArgs)]
      Nothing -> return env
    remainingArgs = drop (length params) args
apply (IOFunc func) args = func args
apply primitive args =
  let (PrimitiveFunc func) = primitive in liftThrows $ func args

primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives =
  [ ("+", numericBinop (+)),
    ("-", numericBinop (-)),
    ("*", numericBinop (*)),
    ("/", numericBinop div),
    ("mod", numericBinop mod),
    ("quotient", numericBinop quot),
    ("remainder", numericBinop rem),
    ("symbol?", unaryOp $ return . isSymbol),
    ("string?", unaryOp $ return . isString),
    ("number?", unaryOp $ return . isNumber),
    ("bool?", unaryOp $ return . isBool),
    ("list?", unaryOp $ return . isList),
    ("symbol->string", unaryOp symbol2string),
    ("string->symbol", unaryOp string2symbol),
    ("string-length", unaryOp stringLength),
    ("string-ref", stringRef),
    ("=", numBoolBinop (==)),
    ("<", numBoolBinop (<)),
    (">", numBoolBinop (>)),
    ("/=", numBoolBinop (/=)),
    (">=", numBoolBinop (>=)),
    ("<=", numBoolBinop (<=)),
    ("&&", boolBoolBinop (&&)),
    ("||", boolBoolBinop (||)),
    ("string=?", strBoolBinop (==)),
    ("string<?", strBoolBinop (<)),
    ("string>?", strBoolBinop (>)),
    ("string<=?", strBoolBinop (<=)),
    ("string>=?", strBoolBinop (>=)),
    ("car", car),
    ("cdr", cdr),
    ("cons", cons),
    ("eq?", eqv),
    ("eqv?", eqv),
    ("equal?", equal)
  ]

unaryOp ::
  (LispVal -> ThrowsError LispVal) ->
  [LispVal] ->
  ThrowsError LispVal
unaryOp func [param] = func param
unaryOp _ params = throwError $ NumArgs 1 params

isSymbol :: LispVal -> LispVal
isSymbol (Atom _) = Bool True
isSymbol _ = Bool False

isString :: LispVal -> LispVal
isString (String _) = Bool True
isString _ = Bool False

isNumber :: LispVal -> LispVal
isNumber (Number _) = Bool True
isNumber _ = Bool False

isBool :: LispVal -> LispVal
isBool (Bool _) = Bool True
isBool _ = Bool False

isList :: LispVal -> LispVal
isList (List _) = Bool True
isList (DottedList _ _) = Bool True
isList _ = Bool False

symbol2string :: LispVal -> ThrowsError LispVal
symbol2string (Atom i) = return $ String i
symbol2string val = throwError $ TypeMismatch "atom" val

string2symbol :: LispVal -> ThrowsError LispVal
string2symbol (String s) = return $ Atom s
string2symbol val = throwError $ TypeMismatch "string" val

stringLength :: LispVal -> ThrowsError LispVal
stringLength (String s) = return . Number $ toInteger . length $ s
stringLength val = throwError $ TypeMismatch "string" val

stringRef :: [LispVal] -> ThrowsError LispVal
stringRef [String s, Number i]
  | i `elem` [0 .. toInteger (length s) - 1] =
    return . Char $ s !! fromInteger i
  | otherwise = throwError $ Default "Index out of bound"
stringRef [String _, val] = throwError $ TypeMismatch "number" val
stringRef [val, _] = throwError $ TypeMismatch "string" val
stringRef val = throwError $ NumArgs 2 val

numericBinop ::
  (Integer -> Integer -> Integer) ->
  [LispVal] ->
  ThrowsError LispVal
numericBinop _ [] = throwError $ NumArgs 2 []
numericBinop _ param@[_] = throwError $ NumArgs 2 param
numericBinop op params = mapM unpackNum params <&> Number . foldl1 op

unpackNum :: LispVal -> ThrowsError Integer
unpackNum (Number n) = return n
unpackNum val@(String s) =
  let parsed = reads s
   in if null parsed
        then throwError $ TypeMismatch "number" val
        else return $ fst . head $ parsed
unpackNum (List [n]) = unpackNum n
unpackNum val = throwError $ TypeMismatch "number" val

boolBinop ::
  (LispVal -> ThrowsError a) ->
  (a -> a -> Bool) ->
  [LispVal] ->
  ThrowsError LispVal
boolBinop unpacker op args
  | length args /= 2 = throwError $ NumArgs 2 args
  | otherwise = do
    lhs <- unpacker $ head args
    rhs <- unpacker $ last args
    return $ Bool $ lhs `op` rhs

numBoolBinop ::
  (Integer -> Integer -> Bool) ->
  [LispVal] ->
  ThrowsError LispVal
numBoolBinop = boolBinop unpackNum

strBoolBinop :: (String -> String -> Bool) -> [LispVal] -> ThrowsError LispVal
strBoolBinop = boolBinop unpackString

unpackString :: LispVal -> ThrowsError String
unpackString (String s) = return s
unpackString (Number n) = return $ show n
unpackString val = throwError $ TypeMismatch "string" val

boolBoolBinop :: (Bool -> Bool -> Bool) -> [LispVal] -> ThrowsError LispVal
boolBoolBinop = boolBinop unpackBool

unpackBool :: LispVal -> ThrowsError Bool
unpackBool (Bool b) = return b
unpackBool val = throwError $ TypeMismatch "boolean" val

car :: [LispVal] -> ThrowsError LispVal
car [List (x : xs)] = return x
car [DottedList (x : xs) _] = return x
car [val] = throwError $ TypeMismatch "pair" val
car val = throwError $ NumArgs 1 val

cdr :: [LispVal] -> ThrowsError LispVal
cdr [List (_ : xs)] = return $ List xs
cdr [DottedList [_] x] = return x
cdr [DottedList (_ : xs) x] = return $ DottedList xs x
cdr [val] = throwError $ TypeMismatch "pair" val
cdr val = throwError $ NumArgs 1 val

cons :: [LispVal] -> ThrowsError LispVal
cons [x, List xs] = return $ List (x : xs)
cons [x, DottedList xs y] = return $ DottedList (x : xs) y
cons [x, y] = return $ DottedList [x] y
cons val = throwError $ NumArgs 1 val

eqv :: [LispVal] -> ThrowsError LispVal
eqv [Nil, Nil] = return $ Bool True
eqv [Atom lhs, Atom rhs] = return $ Bool (lhs == rhs)
eqv [Number lhs, Number rhs] = return $ Bool (lhs == rhs)
eqv [Float lhs, Float rhs] = return $ Bool (lhs == rhs)
eqv [Ratio lhs, Ratio rhs] = return $ Bool (lhs == rhs)
eqv [Complex lhs, Complex rhs] = return $ Bool (lhs == rhs)
eqv [String lhs, String rhs] = return $ Bool (lhs == rhs)
eqv [Char lhs, Char rhs] = return $ Bool (lhs == rhs)
eqv [Bool lhs, Bool rhs] = return $ Bool (lhs == rhs)
eqv [Vector lhs, Vector rhs] =
  eqv [List $ arrayList lhs, List $ arrayList rhs]
eqv [DottedList lhs lhs', DottedList rhs rhs'] =
  eqv [dotted2List lhs lhs', dotted2List rhs rhs']
eqv [List lhs, List rhs] = eqList eqv lhs rhs
eqv [_, _] = return $ Bool False
eqv val = throwError $ NumArgs 2 val

eqList ::
  ([LispVal] -> ThrowsError LispVal) ->
  [LispVal] ->
  [LispVal] ->
  ThrowsError LispVal
eqList eq lhs rhs =
  return . Bool $ (length lhs == length rhs) && all eq' (zip lhs rhs)
  where
    eq' (lhs, rhs) = let (Right (Bool res)) = eq [lhs, rhs] in res

data Unpacker = forall a. Eq a => AnyUnpacker (LispVal -> ThrowsError a)

unpackers :: [Unpacker]
unpackers =
  [ AnyUnpacker unpackNum,
    AnyUnpacker unpackString,
    AnyUnpacker unpackBool
  ]

equal :: [LispVal] -> ThrowsError LispVal
equal [Vector lhs, Vector rhs] =
  equal [List $ arrayList lhs, List $ arrayList rhs]
equal [DottedList lhs lhs', DottedList rhs rhs'] =
  equal [dotted2List lhs lhs', dotted2List rhs rhs']
equal [List lhs, List rhs] = eqList equal lhs rhs
equal args@[lhs, rhs] = do
  primitiveEquals <- or <$> mapM (unpackEquals lhs rhs) unpackers
  (Bool eqvEquals) <- eqv args
  return $ Bool (primitiveEquals || eqvEquals)
equal val = throwError $ NumArgs 2 val

unpackEquals :: LispVal -> LispVal -> Unpacker -> ThrowsError Bool
unpackEquals lhs' rhs' (AnyUnpacker unpacker) =
  do
    lhs <- unpacker lhs'
    rhs <- unpacker rhs'
    return $ lhs == rhs
    `catchError` const (return False)

isBound :: Env -> String -> IO Bool
isBound envRef var = readIORef envRef <&> isJust . lookup var

getVar :: Env -> String -> IOThrowsError LispVal
getVar envRef var = do
  env <- liftIO $ readIORef envRef
  maybe
    (throwError $ UnboundVar "Getting an unbound variable" var)
    (liftIO . readIORef)
    (lookup var env)

setVar :: Env -> String -> LispVal -> IOThrowsError LispVal
setVar envRef var value = do
  env <- liftIO $ readIORef envRef
  maybe
    (throwError $ UnboundVar "Setting an unbound variable" var)
    (liftIO . (`writeIORef` value))
    (lookup var env)
  return value

defineVar :: Env -> String -> LispVal -> IOThrowsError LispVal
defineVar envRef var value = do
  alreadyDefined <- liftIO $ isBound envRef var
  if alreadyDefined
    then setVar envRef var value
    else liftIO $ do
      valueRef <- newIORef value
      env <- readIORef envRef
      writeIORef envRef ((var, valueRef) : env)
      return value

bindVars :: Env -> [(String, LispVal)] -> IO Env
bindVars envRef bindings =
  readIORef envRef
    >>= extendEnv bindings
    >>= newIORef
  where
    extendEnv bindings env = (++ env) <$> mapM addBinding bindings
    addBinding (var, value) = do
      ref <- newIORef value
      return (var, ref)

makeFunc ::
  Maybe String ->
  Env ->
  [LispVal] ->
  [LispVal] ->
  IOThrowsError LispVal
makeFunc varargs env params' body = do
  params <- mapM getParamName params'
  return $ Func params varargs body env
  where
    getParamName :: LispVal -> IOThrowsError String
    getParamName param = do
      case param of
        Atom i -> return i
        val -> throwError $ ExpectedAtom val

makeNormalFunc :: Env -> [LispVal] -> [LispVal] -> IOThrowsError LispVal
makeNormalFunc = makeFunc Nothing

makeVarArgsFunc ::
  LispVal ->
  Env ->
  [LispVal] ->
  [LispVal] ->
  IOThrowsError LispVal
makeVarArgsFunc (Atom varargs) = makeFunc $ Just varargs
makeVarArgsFunc val = \_ _ _ -> throwError $ ExpectedAtom val

ioPrimitives :: [(String, [LispVal] -> IOThrowsError LispVal)]
ioPrimitives =
  [ ("apply", applyProc),
    ("open-input-file", makePort ReadMode),
    ("open-output-file", makePort WriteMode),
    ("close-input-port", closePort),
    ("close-output-port", closePort),
    ("read", readProc),
    ("write", writeProc),
    ("read-contents", readContents),
    ("read-all", readAll)
  ]

applyProc :: [LispVal] -> IOThrowsError LispVal
applyProc [func, List args] = apply func args
applyProc (func : args) = apply func args
applyProc val = throwError $ BadSpecialForm "Ill-formed apply form" (List val)

makePort :: IOMode -> [LispVal] -> IOThrowsError LispVal
makePort mode [String filename] = Port <$> liftIO (openFile filename mode)
makePort _ [val] = throwError $ TypeMismatch "string" val
makePort _ val = throwError $ NumArgs 1 val

closePort :: [LispVal] -> IOThrowsError LispVal
closePort [Port port] = liftIO $ hClose port >> return (Bool True)
closePort _ = return $ Bool False

readProc :: [LispVal] -> IOThrowsError LispVal
readProc [] = readProc [Port stdin]
readProc [Port port] = liftIO (hGetLine port) >>= liftThrows . readExpr
readProc [val] = throwError $ TypeMismatch "port" val
readProc val =
  throwError . Default $ "Expected 0 or 1 argument, found: " ++ show val

writeProc :: [LispVal] -> IOThrowsError LispVal
writeProc [obj] = writeProc [obj, Port stdout]
writeProc [obj, Port port] = liftIO $ hPrint port obj >> return (Bool True)
writeProc [_, val] = throwError $ TypeMismatch "port" val
writeProc val =
  throwError . Default $ "Expected 0 or 1 argument, found: " ++ show val

readContents :: [LispVal] -> IOThrowsError LispVal
readContents [String filename] = String <$> liftIO (readFile filename)
readContents [val] = throwError $ TypeMismatch "string" val
readContents val = throwError $ NumArgs 1 val

load :: String -> IOThrowsError [LispVal]
load filename = liftIO (readFile filename) >>= liftThrows . readExprList

readAll :: [LispVal] -> IOThrowsError LispVal
readAll [String filename] = List <$> load filename
readAll [val] = throwError $ TypeMismatch "string" val
readAll val = throwError $ NumArgs 1 val


{-# OPTIONS_GHC -fno-warn-dodgy-exports #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NoImplicitPrelude #-}

module Sound
(
    -- * Reexports
    -- $reexports
    module Control.Monad
    , module Data.Complex
    , module Data.Int
    -- * Input, output, and conversion
    , module Sound.Io.Snd
    , module Sound.Io
    , module Sound.IoPtr
    , module Sound.Sox
    , module Sound.IoSox
    , module Sound.Portaudio
    -- * Abstraction ladder
    , module Sound.Class
    , module Sound.Pair
    , module Sound.Compile
    , module Sound.InfList
    , module Sound.List
    -- * Buffer
    , module Sound.Buffer
    , module Sound.Ptr
    -- * Time, tempo
    , module Sound.Time
    , module Sound.Tempo
    -- * Sampling
    , module Sound.Sample
    -- * Decibels, contours, envelopes, limiting, clamping
    , module Sound.Amplitude
    , module Sound.Ramp
    -- * Wavetable
    , module Sound.Table
    -- * White noise
    , module Sound.Random
    -- * Frequency modulation
    , module Sound.Fm
    -- * Integration
    , module Sound.Int
    -- * Instrument
    , module Sound.Perform
    -- * MIDI input
    , module Sound.Midi
    -- * Error handling, type hints
    , module Sound.Hint
    -- * Filters: zeros and poles
    , module Sound.Filter
    -- * Unclear
    , module Sound.Function
    -- * Failed experiments
    , module Sound.Abstract
    , module Sound.GenBody
    , module Sound.IoFail
    , module Sound.StreamVector
)
where

import Data.Complex
import Data.Int

import Control.Applicative
import Control.Monad

import Sound.Abstract
import Sound.Amplitude
import Sound.Buffer
import Sound.Class
import Sound.Compile
import Sound.Filter
import Sound.Fm
import Sound.Function
import Sound.GenBody
import Sound.Hint
import Sound.InfList
import Sound.Int
import Sound.Io
import Sound.Io.Snd ()
import Sound.IoFail
import Sound.IoPtr
import Sound.IoSox
import Sound.List ()
import Sound.Midi
import Sound.Pair
import Sound.Perform
import Sound.Portaudio ()
import Sound.Ptr
import Sound.Ramp
import Sound.Random
import Sound.Sample
import Sound.Sox hiding (Rate)
import Sound.StreamVector
import Sound.Table
import Sound.Tempo
import Sound.Time

import qualified Prelude as P

{- $reexports
If you use this module unqualified, you have to hide Prelude imports to avoid name clashes.
The portable way to do that is by adding this import to your code:

@
import Prelude ()
@

This module replaces 'P..' and 'P.id' from "Prelude"
with '.' and 'id' from "Control.Category".
This module hides 'P.seq' from "Prelude".

This module reexports everything else from all modules listed here.
-}