{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DefaultSignatures     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE StandaloneDeriving    #-}

-- | Module : LambdaCube.GL.Input.Type
--
-- This module provides types for describing the matrix and vector ABI for
-- passing vectors and matrices to OpenGL shaders as "uniforms".
--
-- For a given shader input 'InputType', there is associated (via
-- 'witnessType') a type-level description of the GL calls necessary to upload
-- pointers to the GPU.
--
-- To make it so that a Haskell value can be uploaded, instances of 'GLData'
-- and 'Uniformable' should be provided.  These describe how to marshall a
-- haskell type into a suitable pointer for a given context.
--
-- For example, users of the hmatrix package, might find these instances useful:
--
-- > import GHC.TypeLits
-- > import LambdaCube.GL.Input.Type
-- > import Numeric.LinearAlgebra
-- > import Numeric.LinearAlgebra.Devel
-- >
-- > instance Uniformable (Matrix Float) where
-- >     uniformContexts _ = contexts floatMatrices
-- >
-- > instance Uniformable (Vector Float) where
-- >     uniformContexts _ = contexts $ do
-- >         supports TypeFloat
-- >         supports TypeV2F
-- >         supports TypeV3F
-- >         supports TypeV4F
-- >
-- > instance (KnownNat r, KnownNat c) => GLData (Matrix Float) (GLMatrix r c Float) where
-- >     marshalUniform abi mat = case matrixDimensions abi of
-- >         (r,c) | fromIntegral (natVal r) /= rows mat -> Nothing
-- >               | fromIntegral (natVal c) /= cols mat -> Nothing
-- >         _ -> let isRowOrder = case orderOf mat of
-- >                         RowMajor    -> 1
-- >                         ColumnMajor -> 0
-- >              in Just $ MarshalGLMatrix
-- >                      $ \f -> apply mat (\ptr -> f 1 isRowOrder ptr) (\r c sr sc ptr -> ptr)
-- >
-- > instance KnownNat n => GLData (Vector Float) (GLVector n Float) where
-- >     marshalUniform abi vec
-- >         | natVal (vectorLength abi) /= fromIntegral (size vec) = Nothing
-- >         | otherwise = Just $ MarshalGLVector
-- >                            $ \f -> apply vec (\ptr -> f 1 ptr) (\n ptr -> ptr)
module LambdaCube.GL.Input.Type where

import Control.Monad
import Control.Monad.State
import Data.Functor.Identity
import Data.Typeable
import Foreign
import GHC.TypeLits

import Data.Dependent.Map as DMap
import Data.Dependent.Sum
import Data.GADT.Compare

import Graphics.GL.Core33
import LambdaCube.IR (InputType(..))


-- | A 'Uniformable' type /a/ has a runtime dictionary of available instances
-- of 'GLData' /a/ /c/.  For example, if you have a matrix type that does
-- not include type-level dimension information,  you may want to set
--
-- > uniformContexts _ = contexts floatMatrices
--
-- Use 'contexts' and 'supports' to list individual instances.  Typically, for
-- dimension-specific types, 'uniformContexts' will be a singleton map.
--
-- Note that if a context is known at compile-time, then the compiler can
-- lookup an appropriate instance.  This class is provided so that the lookup
-- can happen at runtime instead.
class Uniformable a where
    uniformContexts :: proxy a -> DMap TypeTag (UniformContext a)

-- | A light wrapper over 'execState' to make it easier to specify the
-- 'uniformContexts' DMap.
contexts :: State (DMap TypeTag (UniformContext a)) () -> DMap TypeTag (UniformContext a)
contexts = flip execState DMap.empty

-- | Inserts an instance into the 'uniformContexts' DMap.  See 'floatMatrices'
-- for an example.
supports :: GLData a c => TypeTag c -> State (DMap TypeTag (UniformContext a)) ()
supports ty = modify' (DMap.insert ty UniformContext)


-- | A set of GLData instances for all Float matrix types.  It is defined:
--
-- > floatMatrices = do
-- >     supports TypeM22F
-- >     supports TypeM23F
-- >     supports TypeM24F
-- >     supports TypeM32F
-- >     supports TypeM33F
-- >     supports TypeM34F
-- >     supports TypeM42F
-- >     supports TypeM43F
-- >     supports TypeM44F
--
-- Use 'contexts' to turn this into a 'DMap' suitable to implemented 'uniformContexts'.
floatMatrices :: (GLData a (GLMatrix 2 2 Float),
                  GLData a (GLMatrix 2 3 Float),
                  GLData a (GLMatrix 2 4 Float),
                  GLData a (GLMatrix 3 2 Float),
                  GLData a (GLMatrix 3 3 Float),
                  GLData a (GLMatrix 3 4 Float),
                  GLData a (GLMatrix 4 2 Float),
                  GLData a (GLMatrix 4 3 Float),
                  GLData a (GLMatrix 4 4 Float)) =>
                 State (DMap TypeTag (UniformContext a)) ()
floatMatrices = do
    supports TypeM22F
    supports TypeM23F
    supports TypeM24F
    supports TypeM32F
    supports TypeM33F
    supports TypeM34F
    supports TypeM42F
    supports TypeM43F
    supports TypeM44F

-- | A runtime witness for an instance of 'GLData' /a/ /c/.  To obtain a
-- witness given a specific context 'TypeTag' /c/, perform a lookup into the
-- 'uniformContexts' DMap.
data UniformContext a c = GLData a c => UniformContext

-- | Type-check (at runtime) a specific GLData instance.
knownContext :: (GLData a have, Typeable have, Typeable want) => TypeTag have -> TypeTag want -> Maybe (UniformContext a want)
knownContext known ty = do
    Refl <- withTypes known ty <$> eqT
    return UniformContext

-- | Lookup a GLData instance, given a Uniformable value.
resolveContext :: Uniformable a => a -> TypeTag want -> Maybe (UniformContext a want)
resolveContext val ty = DMap.lookup ty . uniformContexts .  mkproxy $ val
 where
    mkproxy :: val -> Proxy val
    mkproxy _ = Proxy


-- | This type classifies a shader "uniform" input based on the OpenGL calls
-- necessary to upload it to the GPU.  The primitive types are described as if
-- they are 1-vectors.  Vector types are described by 'GLVector' and matrix
-- types are described by 'GLMatrix'.
--
-- Currently, Texture uniforms are not described and there is a, probably
-- useless, run-time distinction between Bool and Word which have the same
-- OpenGL ABI.
--
-- Use 'witnessType' to obtain this from a runtime-only value.
data TypeTag t where
    TypeBool  :: TypeTag (GLVector 1 Word32)
    TypeV2B   :: TypeTag (GLVector 2 Word32)
    TypeV3B   :: TypeTag (GLVector 3 Word32)
    TypeV4B   :: TypeTag (GLVector 4 Word32)
    TypeWord  :: TypeTag (GLVector 1 Word32)
    TypeV2U   :: TypeTag (GLVector 2 Word32)
    TypeV3U   :: TypeTag (GLVector 3 Word32)
    TypeV4U   :: TypeTag (GLVector 4 Word32)

    TypeInt   :: TypeTag (GLVector 1 Int32)
    TypeV2I   :: TypeTag (GLVector 2 Int32)
    TypeV3I   :: TypeTag (GLVector 3 Int32)
    TypeV4I   :: TypeTag (GLVector 4 Int32)

    TypeFloat :: TypeTag (GLVector 1 Float)
    TypeV2F   :: TypeTag (GLVector 2 Float)
    TypeV3F   :: TypeTag (GLVector 3 Float)
    TypeV4F   :: TypeTag (GLVector 4 Float)

    TypeM22F  :: TypeTag (GLMatrix 2 2 Float)
    TypeM23F  :: TypeTag (GLMatrix 3 2 Float)
    TypeM24F  :: TypeTag (GLMatrix 4 2 Float)
    TypeM32F  :: TypeTag (GLMatrix 2 3 Float)
    TypeM33F  :: TypeTag (GLMatrix 3 3 Float)
    TypeM34F  :: TypeTag (GLMatrix 4 3 Float)
    TypeM42F  :: TypeTag (GLMatrix 2 4 Float)
    TypeM43F  :: TypeTag (GLMatrix 3 4 Float)
    TypeM44F  :: TypeTag (GLMatrix 4 4 Float)

instance GEq TypeTag where
    geq a b = do
        Refl <- geq (glABI a) (glABI b)
        guard $ isBoolTag a == isBoolTag b
        return Refl

instance GCompare TypeTag where
    gcompare a b = case compare (isBoolTag a) (isBoolTag b) of
        LT -> GLT
        EQ -> gcompare (glABI a) (glABI b)
        GT -> GGT

-- | This function provides the only term-only information in the 'TypeTag'
-- type.  Booleans are treated similarly to unsigned ints by OpenGL, so the
-- distinction is not made in the type of the tag and the 'GEq' instance of
-- 'TypeTag' must use this helper to help distinguish based on on the runtime
-- information.
isBoolTag :: TypeTag t -> Bool
isBoolTag TypeBool = True
isBoolTag TypeV2B  = True
isBoolTag TypeV3B  = True
isBoolTag TypeV4B  = True
isBoolTag _        = False


-- | Obtain a type-level description of a uniform type context.
--
-- See 'unwitnessType' for the inverse operation.
witnessType :: InputType -> Maybe (Some TypeTag)
witnessType Bool  = Just $ This TypeBool
witnessType V2B   = Just $ This TypeV2B
witnessType V3B   = Just $ This TypeV3B
witnessType V4B   = Just $ This TypeV4B
witnessType Word  = Just $ This TypeWord
witnessType V2U   = Just $ This TypeV2U
witnessType V3U   = Just $ This TypeV3U
witnessType V4U   = Just $ This TypeV4U
witnessType Int   = Just $ This TypeInt
witnessType V2I   = Just $ This TypeV2I
witnessType V3I   = Just $ This TypeV3I
witnessType V4I   = Just $ This TypeV4I
witnessType Float = Just $ This TypeFloat
witnessType V2F   = Just $ This TypeV2F
witnessType V3F   = Just $ This TypeV3F
witnessType V4F   = Just $ This TypeV4F
witnessType M22F  = Just $ This TypeM22F
witnessType M23F  = Just $ This TypeM23F
witnessType M24F  = Just $ This TypeM24F
witnessType M32F  = Just $ This TypeM32F
witnessType M33F  = Just $ This TypeM33F
witnessType M34F  = Just $ This TypeM34F
witnessType M42F  = Just $ This TypeM42F
witnessType M43F  = Just $ This TypeM43F
witnessType M44F  = Just $ This TypeM44F
witnessType _     = Nothing

-- | Discard type-level input information.  Inverse of 'witnessType'.
unwitnessType :: TypeTag c -> InputType
unwitnessType TypeBool  = Bool
unwitnessType TypeV2B   = V2B
unwitnessType TypeV3B   = V3B
unwitnessType TypeV4B   = V4B
unwitnessType TypeWord  = Word
unwitnessType TypeV2U   = V2U
unwitnessType TypeV3U   = V3U
unwitnessType TypeV4U   = V4U
unwitnessType TypeInt   = Int
unwitnessType TypeV2I   = V2I
unwitnessType TypeV3I   = V3I
unwitnessType TypeV4I   = V4I
unwitnessType TypeFloat = Float
unwitnessType TypeV2F   = V2F
unwitnessType TypeV3F   = V3F
unwitnessType TypeV4F   = V4F
unwitnessType TypeM22F  = M22F
unwitnessType TypeM23F  = M23F
unwitnessType TypeM24F  = M24F
unwitnessType TypeM32F  = M32F
unwitnessType TypeM33F  = M33F
unwitnessType TypeM34F  = M34F
unwitnessType TypeM42F  = M42F
unwitnessType TypeM43F  = M43F
unwitnessType TypeM44F  = M44F


-- | A function used to upload a "uniform" primitive value or vector to the
-- GPU.  If /n/ is 1, then it uploads a primitive value of type /typ/.
-- Otherwise, it uploads a vector of length /n/.
--
-- The arguments to the function are:
--
--   * An integer naming the uniform input slot.
--
--   * An element count which, if greater than 1, will bulk-upload an array of
--   values.
--
--   * A pointer to the value or values to be uploaded.
--
-- Use 'glUniform' to obtain this for a given type context.
data GLVector (n :: Nat) typ
    = GLVector (Int32 -> GLsizei -> Ptr typ -> IO ())


-- | A function used to upload a "uniform" /r/ × /c/ matrix of /typ/ to the
-- GPU.  The arguments are:
--
--   * An integer naming the uniform input slot.
--
--   * An element count which, if greater than 1, will bulk-upload an array.
--
--   * A flag that is GL_TRUE if the values at the pointer are arranged in
--   row-major order.  Set this to GL_FALSE to indicate column-major ordering.
--   Row-major means that the first /c/ contiguous values at the pointer
--   constitute the first row of the matrix.
--
--   * A pointer to the matrix values.
--
-- IMPORTANT: This type flouts the usual graphics convention of width×height,
-- which is used in the naming of the constructors for 'InputType' and
-- 'TypeTag', and instead instead follows the opposite, matrix-math convention,
-- of rows×columns.  Thus, @'TypeM42F'@ is associated with @(GLMatrix 2 4
-- Float)@.
--
-- Use 'glUniform' to obtain this for a given type context.
data GLMatrix (r :: Nat) (c :: Nat) typ
    = GLMatrix (Int32 -> GLsizei -> GLboolean -> Ptr typ -> IO ())

-- | Obtain a suitable 'GLVector' or 'GLMatrix' OpenGL API function for a given
-- uniform type.
glUniform :: TypeTag a -> a
glUniform TypeBool  = GLVector glUniform1uiv
glUniform TypeWord  = GLVector glUniform1uiv
glUniform TypeInt   = GLVector glUniform1iv
glUniform TypeFloat = GLVector glUniform1fv
glUniform TypeV2B   = GLVector glUniform2uiv
glUniform TypeV2U   = GLVector glUniform2uiv
glUniform TypeV2I   = GLVector glUniform2iv
glUniform TypeV2F   = GLVector glUniform2fv
glUniform TypeV3B   = GLVector glUniform3uiv
glUniform TypeV3U   = GLVector glUniform3uiv
glUniform TypeV3I   = GLVector glUniform3iv
glUniform TypeV3F   = GLVector glUniform3fv
glUniform TypeV4B   = GLVector glUniform4uiv
glUniform TypeV4U   = GLVector glUniform4uiv
glUniform TypeV4I   = GLVector glUniform4iv
glUniform TypeV4F   = GLVector glUniform4fv
glUniform TypeM22F  = GLMatrix glUniformMatrix2fv
glUniform TypeM23F  = GLMatrix glUniformMatrix2x3fv
glUniform TypeM24F  = GLMatrix glUniformMatrix2x4fv
glUniform TypeM32F  = GLMatrix glUniformMatrix3x2fv
glUniform TypeM33F  = GLMatrix glUniformMatrix3fv
glUniform TypeM34F  = GLMatrix glUniformMatrix3x4fv
glUniform TypeM42F  = GLMatrix glUniformMatrix4x2fv
glUniform TypeM43F  = GLMatrix glUniformMatrix4x3fv
glUniform TypeM44F  = GLMatrix glUniformMatrix4fv

-- | 'GLData' /a/ /c/ is true when /a/ can be used as GL shader "uniform"
-- input for a 'GLVector' or 'GLMatrix' uniform context /c/.
--
-- For simple (non-matrix) types that implement 'Storable', a default
-- implementation, implemented by 'marshalUniformStorable', is provided, so in
-- this case, simply declare the instance.  If a 'Storable' column-major matrix
-- type, 'marshalColumnMajor' can be used.
class (Typeable c, Typeable a) => GLData a c where
    -- | Provide a pointer suitable for passing to the OpenGL api.
    marshalUniform :: GLABI c -- ^ Description of the variable type of the GL uniform.
                      -> a         -- ^ The value to upload to the GPU.
                      -> Maybe (MarshalGL c)

    default marshalUniform :: ( c ~ GLVector n typ
                              , Storable a
                              ) => GLABI c -> a -> Maybe (MarshalGL c)
    marshalUniform _ a = Just (marshalUniformStorable a)

-- | A suitable default for 'Storable' uniform non-matrix types.
marshalUniformStorable :: Storable a => a -> MarshalGL (GLVector n typ)
marshalUniformStorable a = MarshalGLVector $ \f -> with a (f 1 . castPtr)

-- | A suitable default for 'Storable' uniform column-major matrix types.
-- Column-major means that columns are stored as contiguous regions of memory
-- at the pointer.
marshalColumnMajor :: Storable a => a -> MarshalGL (GLMatrix r c typ)
marshalColumnMajor a = MarshalGLMatrix $ \f -> with a (f 1 GL_FALSE . castPtr)


-- | Run-time information about a "uniform" type context.  A proxy for the
-- pointer type is provided directly.  You can use 'vectorLength' or
-- 'matrixDimensions' to extract proxies for size information.
--
-- You can obtain this from a 'TypeTag' using 'glABI'.
data GLABI m where
    IsGLVector :: KnownNat n => GLPointerType typ -> GLABI (GLVector n typ)
    IsGLMatrix :: (KnownNat r, KnownNat c) => GLPointerType typ -> GLABI (GLMatrix r c typ)

instance GEq GLABI where
    geq a@(IsGLVector x) b@(IsGLVector y) = do
        Refl <- geq x y
        Refl <- withTypes (vectorLength a) (vectorLength b) <$> eqT
        return Refl
    geq a@(IsGLMatrix x) b@(IsGLMatrix y) = do
        Refl <- geq x y
        let (ar,ac) = matrixDimensions a
            (br,bc) = matrixDimensions b
        Refl <- withTypes (asTypeOf ar RowCount) (asTypeOf br RowCount) <$> eqT
        Refl <- withTypes (asTypeOf ac ColumnCount) (asTypeOf bc ColumnCount) <$> eqT
        return Refl
    geq _ _ = Nothing

instance GCompare GLABI where
    gcompare IsGLVector{} IsGLMatrix{} = GLT
    gcompare IsGLMatrix{} IsGLVector{} = GGT
    gcompare a@(IsGLVector x) b@(IsGLVector y) = case gcompare x y of
        GLT -> GLT
        GEQ -> case withTypes (vectorLength a) (vectorLength b) <$> eqT of
            Just Refl -> GEQ
            Nothing   -> case compare (natVal $ vectorLength a) (natVal $ vectorLength b) of
                    LT -> GLT
                    GT -> GGT
        GGT -> GGT
    gcompare a@(IsGLMatrix x) b@(IsGLMatrix y) = case gcompare x y of
        GLT -> GLT
        GEQ -> case matrixDimensions a of
                (ar,ac) -> case matrixDimensions b of
                  (br,bc) -> case withTypes ar br <$> eqT of
                    Just Refl -> case withTypes ac bc <$> eqT of
                        Just Refl -> GEQ
                        Nothing   -> if natVal ac < natVal bc then GLT else GGT
                    Nothing -> if natVal ar < natVal br then GLT else GGT
        GGT -> GGT


-- | Run-time encoding of the type of passed pointer to the OpenGL api for a
-- given uniform.
data GLPointerType typ where
    GLPrimUInt  :: GLPointerType Word32
    GLPrimInt   :: GLPointerType Int32
    GLPrimFloat :: GLPointerType Float

deriving instance Show (GLPointerType typ)

instance GEq GLPointerType where
    geq GLPrimUInt  GLPrimUInt  = Just Refl
    geq GLPrimInt   GLPrimInt   = Just Refl
    geq GLPrimFloat GLPrimFloat = Just Refl

instance GCompare GLPointerType where
    gcompare GLPrimUInt  GLPrimUInt  = GEQ
    gcompare GLPrimUInt  _           = GLT
    gcompare GLPrimFloat GLPrimFloat = GEQ
    gcompare GLPrimFloat _           = GGT
    gcompare GLPrimInt   GLPrimUInt  = GLT
    gcompare GLPrimInt   GLPrimInt   = GEQ
    gcompare GLPrimInt   GLPrimFloat = GGT

-- | Convenience proxy for the number of dimensions in a vector.
data VectorLength (n :: Nat) = VectorLength

-- | Extract the number of dimensions in a vector from a 'GLABI'.
vectorLength :: GLABI (GLVector n typ) -> VectorLength n
vectorLength IsGLVector{} = VectorLength

-- | Convenience proxy for the number of rows in a matrix.
data RowCount (n :: Nat) = RowCount

-- | Convenience proxy for the number of columns in a matrix.
data ColumnCount (n :: Nat) = ColumnCount

-- | Extract row and column counts from a 'GLABI'.
matrixDimensions :: GLABI (GLMatrix r c typ) -> (RowCount r, ColumnCount c)
matrixDimensions IsGLMatrix{} = (RowCount,ColumnCount)

-- | Convenience accessor for the pointer type proxy of a 'GLABI'.
ptrType :: GLABI (f t) -> GLPointerType t
ptrType (IsGLVector p) = p
ptrType (IsGLMatrix p) = p


-- | Provides a pointer and element count to a given continuation.
--
-- If the pointer refers to matrix data, then a 'GLboolean' is also passed.
-- This is GL_TRUE when the matrix data is in row-major format (i.e. matrix
-- rows are contiguous blocks of memory.)
data MarshalGL c where

    MarshalGLVector :: (forall b. (GLsizei -> Ptr typ -> IO b) -> IO b)
                        -> MarshalGL (GLVector n typ)

    MarshalGLMatrix :: (forall b. (GLsizei -> GLboolean -> Ptr typ -> IO b) -> IO b)
                        -> MarshalGL (GLMatrix r c typ)


-- | Obtain a type-level description for a given input context.  This is
-- similar to 'glUniform' except that it is used only as description and does
-- not provide the actual OpenGL API function.
glABI :: TypeTag a -> GLABI a
glABI TypeBool  = IsGLVector GLPrimUInt
glABI TypeWord  = IsGLVector GLPrimUInt
glABI TypeInt   = IsGLVector GLPrimInt
glABI TypeFloat = IsGLVector GLPrimFloat
glABI TypeV2B   = IsGLVector GLPrimUInt
glABI TypeV2U   = IsGLVector GLPrimUInt
glABI TypeV2I   = IsGLVector GLPrimInt
glABI TypeV2F   = IsGLVector GLPrimFloat
glABI TypeV3B   = IsGLVector GLPrimUInt
glABI TypeV3U   = IsGLVector GLPrimUInt
glABI TypeV3I   = IsGLVector GLPrimInt
glABI TypeV3F   = IsGLVector GLPrimFloat
glABI TypeV4B   = IsGLVector GLPrimUInt
glABI TypeV4U   = IsGLVector GLPrimUInt
glABI TypeV4I   = IsGLVector GLPrimInt
glABI TypeV4F   = IsGLVector GLPrimFloat
glABI TypeM22F  = IsGLMatrix GLPrimFloat
glABI TypeM23F  = IsGLMatrix GLPrimFloat
glABI TypeM24F  = IsGLMatrix GLPrimFloat
glABI TypeM32F  = IsGLMatrix GLPrimFloat
glABI TypeM33F  = IsGLMatrix GLPrimFloat
glABI TypeM34F  = IsGLMatrix GLPrimFloat
glABI TypeM42F  = IsGLMatrix GLPrimFloat
glABI TypeM43F  = IsGLMatrix GLPrimFloat
glABI TypeM44F  = IsGLMatrix GLPrimFloat

-- | This is a convenience utility to provide explicit context to
-- type-equality.  For example, to compare types denoted by proxies /a/ and
-- /b/, use
--
-- > case withTypes a b <$> eqT of
-- >   Just Refl -> typesAreEqual
-- >   Nothing   -> typesArenotEqual
--
-- This utility is not OpenGL related, but it is used by the GCompare instances
-- implemented here.
withTypes :: p (a::k) -> q (b::k) -> f a b -> f a b
withTypes _ _ x = x