content: module LambdaCube.Convert.ToDeBruijnRec {-(convertGPOutput)-} (convert) where
  
  import GHC.TypeLits
  
  import Debug.Trace
  
  import LambdaCube.Language.APIType (FlatTuple(..),Frequency(..))
  import LambdaCube.Language.DSLType (GPU,Tuple(..),TupleIdx(..))
  import qualified LambdaCube.Language.APIType as T
  import qualified LambdaCube.Language.DSLType as T hiding (Shadow)
  import qualified LambdaCube.Language.PrimFun as T
  import qualified LambdaCube.Language.HOAS as H
  import qualified LambdaCube.Core.DeBruijn as U
  import LambdaCube.Core.DeBruijnRec
  import LambdaCube.Core.APIType
  import LambdaCube.Convert.PrimFun
  import LambdaCube.Core.Type as G
  
  import Data.Reify
  import Data.Typeable
  import Data.Dynamic
  import Control.Applicative hiding (Const)
  import Data.Map (Map,(!))
  import qualified Data.Map as Map
  import qualified Data.IntMap as IM
  import Control.Monad.State
  import Data.List (elemIndex)
  import LambdaCube.Core.Util.BiMap
  import qualified Data.Vector as V
  
  data ConvertState
      = ConvertState
      { lambdas   :: [Unique]
      , expMap    :: Map Unique (TypedExp Unique, Maybe Unique)
      , cvtMap    :: Map Unique (Ty,U.Exp)
      }
  
  type C = State ConvertState
  
  cvt :: Unique -> C Unique
  cvt idx = do
      s <- get
      case expMap s ! idx of
          (_,Just eId)    -> return eId
          (e,Nothing)     -> convertExp idx e
  
  cvtList :: Unique -> C [Unique]
  cvtList idx = do
      s <- get
      case expMap s ! idx of
          ((E ty e),_)  -> case e of
              Nil         -> return []
              Cons x xs   -> (:) <$> cvt x <*> cvtList xs
              _           -> error "internal error (cvtList)"
  
  cvtTuple :: Unique -> C [Unique]
  cvtTuple idx = do
      s <- get
      case expMap s ! idx of
          ((E ty e),_)  -> case e of
              Nil         -> return []
              Cons x xs   -> (:) <$> cvt x <*> cvtTuple xs
              _           -> error "internal error (cvtTuple)"
  
  emit :: Unique -> Ty -> U.Exp -> C Unique
  emit idx ty e = do
      ConvertState{..} <- get
      let nextIdx = Map.size cvtMap
          cvtMap' = Map.insert nextIdx (ty,e) cvtMap
          expMap' = Map.adjust (\(e,Nothing) -> (e,Just nextIdx)) idx expMap
      put $ ConvertState lambdas expMap' cvtMap'
      return nextIdx
  
  convertExp :: Unique -> TypedExp Unique -> C Unique
  convertExp eIdx (E ty exp) = case exp of
      Lam a b -> do
          let add e = do
                  cm <- gets cvtMap
                  let idx = Map.size cm
                  modify (\s -> s {cvtMap = Map.insert idx (Unknown "Body",e) cm})
                  return idx
          ls <- gets lambdas
          modify (\s -> s {lambdas = a : ls})
          v <- cvt a
          f <- cvt b
          modify (\s -> s {lambdas = ls})
          emit eIdx ty =<< U.Lam <$> add (U.Body f)
  
      Var                     -> do
          ls <- gets lambdas
          let Just i = elemIndex eIdx ls
          emit eIdx ty $ U.Var i ""
  
      -- Exp
      Const a                 -> emit eIdx ty $ U.Const a
      PrimVar a               -> emit eIdx ty $ U.PrimVar a
      Uni a                   -> emit eIdx ty $ U.Uni a
  
      Tup a                   -> emit eIdx ty =<< (\t -> U.Tup (reverse t)) <$> cvtTuple a
      Prj a b                 -> emit eIdx ty =<< U.Prj a <$> cvt b
      Cond a b c              -> emit eIdx ty =<< U.Cond <$> cvt a <*> cvt b <*> cvt c
      PrimApp a b             -> emit eIdx ty =<< U.PrimApp a <$> cvt b
      Sampler a b c           -> emit eIdx ty =<< U.Sampler a b <$> cvt c
      Loop a b c d            -> emit eIdx ty =<< U.Loop <$> cvt a <*> cvt b <*> cvt c <*> cvt d
  
      -- special tuple expressions
      VertexOut a b c d       -> emit eIdx ty =<< U.VertexOut <$> cvt a <*> cvt b <*> cvtList c <*> cvtList d
      GeometryOut a b c d e   -> emit eIdx ty =<< U.GeometryOut <$> cvt a <*> cvt b <*> cvt c <*> cvtList d <*> cvtList e
      FragmentOut a           -> emit eIdx ty =<< U.FragmentOut <$> cvtList a
      FragmentOutDepth a b    -> emit eIdx ty =<< U.FragmentOutDepth <$> cvt a <*> cvtList b
      FragmentOutRastDepth a  -> emit eIdx ty =<< U.FragmentOutRastDepth <$> cvtList a
  
      -- GP
      Fetch a b c             -> emit eIdx ty $ U.Fetch a b c
      Transform a b           -> emit eIdx ty =<< U.Transform <$> cvt a <*> cvt b
      Reassemble a b          -> emit eIdx ty =<< U.Reassemble <$> cvt a <*> cvt b
      Rasterize a b           -> emit eIdx ty =<< U.Rasterize a <$> cvt b
      FrameBuffer a           -> emit eIdx ty $ U.FrameBuffer a
      Accumulate a b c d e    -> emit eIdx ty =<< U.Accumulate a <$> cvt b <*> cvt c <*> cvt d <*> cvt e
      PrjFrameBuffer a b c    -> emit eIdx ty =<< U.PrjFrameBuffer a b <$> cvt c
      PrjImage a b c          -> emit eIdx ty =<< U.PrjImage a b <$> cvt c
  
      -- Texture
      TextureSlot a b         -> emit eIdx ty $ U.TextureSlot a b
      Texture a b c d         -> emit eIdx ty =<< U.Texture a b c <$> cvtList d
  
      -- Interpolated
      Flat a                  -> emit eIdx ty =<< U.Flat <$> cvt a
      Smooth a                -> emit eIdx ty =<< U.Smooth <$> cvt a
      NoPerspective a         -> emit eIdx ty =<< U.NoPerspective <$> cvt a
  
      GeometryShader a b c d e f  -> emit eIdx ty =<< U.GeometryShader a b c <$> cvt d <*> cvt e <*> cvt f
  
      -- FragmentFilter
      PassAll                 -> emit eIdx ty $ U.PassAll
      Filter a                -> emit eIdx ty =<< U.Filter <$> cvt a
  
      -- GPOutput
      ImageOut a b c          -> emit eIdx ty =<< U.ImageOut a b <$> cvt c
      ScreenOut a             -> emit eIdx ty =<< U.ScreenOut <$> cvt a
      MultiOut a              -> emit eIdx ty =<< U.MultiOut <$> cvtList a
      _ -> error $ "internal error (convertExp) " ++ show exp
  
  convert :: H.GPOutput a -> IO (U.Exp,U.DAG)
  convert e = do
      g@(Graph l r) <- reifyGraph e
      let s = ConvertState [] (Map.fromList [(i,(e,Nothing)) | (i,e) <- l]) Map.empty
          (r',s') = runState (cvt r) s
          tel = Map.toList $ cvtMap s'
          (tl,el) = unzip [((i,t),(i,e)) | (i,(t,e)) <- tel]
          emap  = IM.fromList el
          bimap = BiMap (Map.fromList [(e,i) | (i,e) <- el]) emap
          tmap  = IM.fromList tl
          cmap  = IM.fromList [(i,2) | (i,_) <- tl] -- HINT: hack
      print r >> mapM print l
      print "----"
      print r' >> mapM print tel
      return (emap IM.! r',U.DAG bimap tmap cmap V.empty V.empty)
  
  saveCore :: String -> H.GPOutput a -> IO ()
  saveCore name e = convert e >>= writeFile (name ++ ".lcore") . show
  
  deriving instance Typeable H.VertexOut
  deriving instance Typeable H.GeometryOut
  deriving instance Typeable H.FragmentOut
  deriving instance Typeable H.Exp
  deriving instance Typeable H.GPOutput
  
  instance GPU b => NewVar (H.Exp a b) where
    mkVar = H.Var
  
  class NewVar a where
    mkVar :: Dynamic -> a
  
  instance Eq (H.Exp a b) where
      _ == _ = False
  
  data TypedExp a = E Ty (Exp a) deriving (Eq,Ord,Show,Typeable)
  
  addType (e :: H.Exp a t) f = E (genTy (undefined :: t)) <$> f
  noType s f = E (Unknown s) <$> f
  
  instance (GPU a,MuRef (H.Exp f a),Typeable (H.Exp f a), NewVar (H.Exp f a), Typeable b, MuRef b, DeRef (H.Exp f a) ~ DeRef ((H.Exp f a) -> b),DeRef b ~ DeRef ((H.Exp f a) -> b)) => MuRef ((H.Exp f a) -> b) where
    type DeRef ((H.Exp f a) -> b) = TypedExp
  
    mapDeRef f (fn :: (H.Exp f a) -> b) = let v = mkVar $ toDyn fn 
                    in addType (undefined :: H.Exp f a) $ Lam <$> f v <*> f (fn v)
  
  
  instance MuRef (H.Exp a b) where
    type DeRef (H.Exp a b) = TypedExp
  
    mapDeRef f x@(H.Var t)        = addType x $ pure Var
    mapDeRef f x@(H.Const t)      = addType x $ pure (Const $ T.toValue t)
    mapDeRef f x@(H.PrimVar i)    = addType x $ pure (PrimVar $ fst $ T.toInput i)
    mapDeRef f x@(H.Uni i)        = addType x $ pure (Uni $ fst $ T.toInput i)
    mapDeRef f x@(H.Cond c t e)   = addType x $ Cond <$> f c <*> f t <*> f e
    mapDeRef f x@(H.PrimApp p a)  = addType x $ PrimApp (convertPrimFun p) <$> f a
    mapDeRef f x@(H.Tup t)        = addType x $ Tup <$> f t
    mapDeRef f x@(H.Prj idx (e :: H.Exp stage e') :: H.Exp stage' t') = addType x $ Prj (genTupLen (prjToInt idx) (undefined :: e')) <$> f e
    mapDeRef f x@(H.Sampler a b c)   = addType x $ Sampler a b <$> f c
    mapDeRef f x@(H.TextureSlot n t) = noType "TextureSlot" $ pure $ TextureSlot n $ convertTextureType t
    mapDeRef f x@(H.Texture t s m d) = noType "Texture" $ Texture (convertTextureType t) (T.toValue s) (convertMipMap m) <$> f d
    mapDeRef f x@(H.Loop a b c d)   = addType x $ Loop <$> f a <*> f b <*> f c <*> f d
    mapDeRef f x@(H.Fetch a b c)    = noType "Fetch" $ pure $ Fetch a (convertFetchPrimitive b) (T.toInputList c)
    mapDeRef f x@(H.Transform a b)  = noType "Transform" $ Transform <$> f a <*> f b
    mapDeRef f x@(H.Reassemble a b) = noType "Reassemble" $ Reassemble <$> f a <*> f b
    mapDeRef f x@(H.Rasterize a b)  = noType "Rasterize" $ Rasterize (convertRasterContext a) <$> f b
    mapDeRef f x@(H.FrameBuffer a)  = noType "FrameBuffer" $ pure $ FrameBuffer $ convertFrameBuffer a
    mapDeRef f x@(H.Accumulate a b c d e) = noType "Accumulate" $ Accumulate (convertAccumulationContext a) <$> f b <*> f c <*> f d <*> f e
    mapDeRef f x@(H.PrjFrameBuffer a b c) = noType "PrjFrameBuffer" $ PrjFrameBuffer a (prjToInt b) <$> f c
    mapDeRef f x@(H.PrjImage a b c) = noType "PrjImage" $ PrjImage a (toInt b) <$> f c
  
  instance (Typeable a, MuRef a,DeRef [a] ~ DeRef a) => MuRef [a] where
    type DeRef [a] = TypedExp
  
    mapDeRef f (x:xs) = noType "ConsL" $ liftA2 Cons (f x) (f xs)
    mapDeRef f []     = noType "NilL" $ pure Nil
  
  instance MuRef (Tuple (H.Exp stage) t) where
    type DeRef (Tuple (H.Exp stage) t) = TypedExp
  
    mapDeRef f NilTup             = noType "NilT" $ pure Nil
    mapDeRef f (es `SnocTup` e)   = noType "ConsT" $ liftA2 Cons (f e) (f es) -- inverted list
  
  instance MuRef (T.Interpolated (H.Exp f) a) where
    type DeRef (T.Interpolated (H.Exp f) a) = TypedExp
  
    mapDeRef f (T.Flat e)          = noType "Flat" $ Flat <$> f e
    mapDeRef f (T.Smooth e)        = noType "Smooth" $ Smooth <$> f e
    mapDeRef f (T.NoPerspective e) = noType "NoPerspective" $ NoPerspective <$> f e
  
  instance MuRef (H.FlatExp f a) where
    type DeRef (H.FlatExp f a) = TypedExp
  
    mapDeRef f (e:.xs) = noType "ConsFE" $ liftA2 Cons (f e) (f xs)
    mapDeRef f ZT = noType "NilFE" $ pure Nil
  
  instance MuRef (H.InterpolatedFlatExp f a) where
    type DeRef (H.InterpolatedFlatExp f a) = TypedExp
  
    mapDeRef f (e:.xs) = noType "ConsIFE" $ liftA2 Cons (f e) (f xs)
    mapDeRef f ZT = noType "NilIFE" $ pure Nil
  
  instance GPU a => MuRef (H.FragmentFilter a) where
    type DeRef (H.FragmentFilter a) = TypedExp
  
    mapDeRef f H.PassAll    = noType "PassAll" $ pure PassAll
    mapDeRef f (H.Filter a) = noType "Filter" $ Filter <$> f a
  
  instance MuRef (H.FragmentOut a) where
    type DeRef (H.FragmentOut a) = TypedExp
  
    mapDeRef f (H.FragmentOut a) = noType "FragmentOut" $ FragmentOut <$> f a
    mapDeRef f (H.FragmentOutDepth a b) = noType "FragmentOutDepth" $ FragmentOutDepth <$> f a <*> f b
    mapDeRef f (H.FragmentOutRastDepth a) = noType "FragmentOutRastDepth" $ FragmentOutRastDepth <$> f a
  
  instance MuRef (H.VertexOut a b) where
    type DeRef (H.VertexOut a b) = TypedExp
  
    mapDeRef f (H.VertexOut a b c d) = noType "VertexOut" $ VertexOut <$> f a <*> f b <*> f c <*> f d
  
  instance MuRef (H.GeometryOut a b c) where
    type DeRef (H.GeometryOut a b c) = TypedExp
  
    mapDeRef f (H.GeometryOut a b c d e) = noType "GeometryOut" $ GeometryOut <$> f a <*> f b <*> f c <*> f d <*> f e
  
  instance MuRef (H.GeometryShader a b c d e f g) where
    type DeRef (H.GeometryShader a b c d e f g) = TypedExp
  
    mapDeRef f (H.GeometryShader a b c d e g) = noType "GeometryShader" $ GeometryShader (toInt a) (convertOutputPrimitive b) c <$> f d <*> f e <*> f g
  
  instance MuRef (H.GPOutput a) where
    type DeRef (H.GPOutput a) = TypedExp
  
    mapDeRef f (H.ImageOut a b c) = noType "ImageOut" $ ImageOut a b <$> f c
    mapDeRef f (H.ScreenOut a) = noType "ScreenOut" $ ScreenOut <$> f a
    mapDeRef f (H.MultiOut a) = noType "MultiOut" $ MultiOut <$> f a
  
  main = do
          let g1 :: H.Exp V Float
              g1 = H.Const 13
          reifyGraph g1 >>= print
  
  toInt :: KnownNat n => T.NatNum n -> Int
  toInt (a :: T.NatNum n) = fromInteger $ natVal a
  
  prjIdx i lyt = i--length lyt - i - 1
  
  
  prjToInt :: TupleIdx t e -> Int
  prjToInt ZeroTupIdx     = 0
  prjToInt (SuccTupIdx i) = 1 + prjToInt i
  
  genTupLen :: GPU a => Int -> a -> Int
  genTupLen i a = sum $ map tySize $ take i rt
    where
      rt = reverse t
      Tuple t = genTy a
  
  type Layout = [[Ty]]
  
  genTy :: GPU a => a -> Ty
  genTy = T.tupleType
  
  convertTextureDataType :: T.TextureDataType t ar -> TextureDataType
  convertTextureDataType (T.Float a)  = FloatT (T.toColorArity a)
  convertTextureDataType (T.Int a)    = IntT (T.toColorArity a)
  convertTextureDataType (T.Word a)   = WordT (T.toColorArity a)
  convertTextureDataType T.Shadow     = ShadowT
  
  convertTextureType :: T.TextureType dim mip arr layerCount t ar -> TextureType
  convertTextureType (T.Texture1D a b)    = Texture1D (convertTextureDataType a) (toInt b)
  convertTextureType (T.Texture2D a b)    = Texture2D (convertTextureDataType a) (toInt b)
  convertTextureType (T.Texture3D a)      = Texture3D (convertTextureDataType a)
  convertTextureType (T.TextureCube a)    = TextureCube (convertTextureDataType a)
  convertTextureType (T.TextureRect a)    = TextureRect (convertTextureDataType a)
  convertTextureType (T.Texture2DMS a b)  = Texture2DMS (convertTextureDataType a) (toInt b)
  convertTextureType (T.TextureBuffer a)  = TextureBuffer (convertTextureDataType a)
  
  convertMipMap :: T.MipMap t -> MipMap
  convertMipMap (T.NoMip)         = NoMip
  convertMipMap (T.Mip a b)       = Mip a b
  convertMipMap (T.AutoMip a b)   = AutoMip a b
  
  convertRasterContext :: T.RasterContext p -> RasterContext
  convertRasterContext (T.PointCtx a b c)         = PointCtx a b c
  convertRasterContext (T.LineCtx a b)            = LineCtx a b
  convertRasterContext (T.TriangleCtx a b c d)    = TriangleCtx a b c d
  
  convertBlending :: T.Blending c -> Blending
  convertBlending T.NoBlending        = NoBlending
  convertBlending (T.BlendLogicOp a)  = BlendLogicOp a
  convertBlending (T.Blend a b c)     = Blend a b c
  
  convertFetchPrimitive :: T.FetchPrimitive a -> FetchPrimitive
  convertFetchPrimitive v = case v of
      T.Points                    -> Points
      T.Lines                     -> Lines
      T.Triangles                 -> Triangles
      T.LinesAdjacency            -> LinesAdjacency
      T.TrianglesAdjacency        -> TrianglesAdjacency
  
  convertOutputPrimitive :: T.OutputPrimitive a -> OutputPrimitive
  convertOutputPrimitive v = case v of
      T.TrianglesOutput   -> TrianglesOutput
      T.LinesOutput       -> LinesOutput
      T.PointsOutput      -> PointsOutput
  
  convertAccumulationContext :: T.AccumulationContext b -> AccumulationContext
  convertAccumulationContext (T.AccumulationContext n ops) = AccumulationContext n $ cvt ops
    where
      cvt :: FlatTuple T.NoConstraint T.FragmentOperation b -> [FragmentOperation]
      cvt ZT                          = []
      cvt (T.DepthOp a b:.xs)         = DepthOp a b : cvt xs
      cvt (T.StencilOp a b c :. xs)   = StencilOp a b c : cvt xs
      cvt (T.ColorOp a b :. xs)       = ColorOp (convertBlending a) (T.toValue b) : cvt xs
  
  convertFrameBuffer :: T.FrameBuffer layerCount t -> [Image]
  convertFrameBuffer = cvt
    where
      cvt :: T.FrameBuffer layerCount t -> [Image]
      cvt ZT                          = []
      cvt (T.DepthImage a b:.xs)      = DepthImage (toInt a) b : cvt xs
      cvt (T.StencilImage a b:.xs)    = StencilImage (toInt a) b : cvt xs
      cvt (T.ColorImage a b:.xs)      = ColorImage (toInt a) (T.toValue b) : cvt xs
      cvt (T.UnclearedImage a:.xs)    = UnclearedImage (toInt a) : cvt xs
  