content: {-# LANGUAGE FlexibleContexts  #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TemplateHaskell   #-}
  {-# LANGUAGE TupleSections     #-}
  {-# OPTIONS_GHC -Wall  #-}
  
  module Utils where
  
  import           Control.Lens ((<&>), makeLenses)
  import           Control.Monad.State
  import           Control.Monad.Trans.Except
  import           Data.Bifunctor (first, second)
  import           Data.List (nub)
  import qualified Data.Map as M
  import           Data.Monoid ((<>))
  import qualified Data.Set as S
  import           Debug.Trace  (trace)
  import           Types
  
  
  showTrace :: Show b => b -> b
  showTrace = trace =<< show
  
  
  data TIState = TIState
    { _tiVNames :: Int
    , _tiTNames :: Int
    , _tiSubst  :: Subst
    }
  
  makeLenses ''TIState
  
  type TI = ExceptT String (State TIState)
  
  
  unravel :: Exp a -> Maybe (VName, [Exp a])
  unravel = go []
    where
      go acc (LCon a) = pure (a, acc)
      go acc (a :@ b) = go (b : acc) a
      go _ _ = Nothing
  
  
  unravelNative :: Exp a -> Maybe (String, [Exp a])
  unravelNative = go []
    where
      go acc (Lit (LitNative a _)) = pure (a, acc)
      go acc (a :@ b)              = go (b : acc) a
      go _ _                       = Nothing
  
  
  letters :: [String]
  letters = do
    b <- "":letters
    a <- ['a'..'z']
    pure $ a : b
  
  
  runTI :: TI a -> Either String a
  runTI = flip evalState (TIState 0 0 mempty) . runExceptT
  
  
  kind :: Type -> TI Kind
  kind (TVar x)  = pure $ tKind x
  kind (TCon x)  = pure $ tKind x
  kind (a :@@ b) = do
    ka <- kind a
    kb <- kind b
    let kerr kk = throwE $ mconcat
          [ "kind mismatch: '"
          , show b
          , " :: "
          , show kb
          , "' vs '"
          , show kk
          , "'\nwhen trying to apply '"
          , show a
          , " :: "
          , show ka
          , "'\n"
          ]
    case ka of
      kal :>> kar -> do
        when (kal /= kb) $ kerr kal
        pure kar
      KStar       -> kerr KStar
      KConstraint -> kerr KConstraint
  
  
  
  buildDataCon
      :: VName
      -> [Type]
      -> Maybe Type
      -> GenDataCon
  buildDataCon n@(VName s) ts t' =
    let ks = fmap (either error id . runTI . kind) ts
        k  = foldr (:>>) KStar ks
        tr = maybe (foldl (:@@) (TCon $ TName s k)
           . fmap TVar
           $ S.toList $ free ts) id t'
        t  = foldr (:->) tr ts
        ls = fmap fst $ zip (fmap VName letters) ts
     in GenDataCon n ([] :=> t) ([] :=> tr)
      . foldr lam
              (foldl (:@) (LCon n) $ fmap V ls)
      $ ls
  
  
  buildRecord
      :: VName
      -> [(VName, Type)]
      -> Maybe Type
      -> (GenDataCon, [(VName, (Qual Type, Exp VName))])
  buildRecord n fs t =
    let gen@(GenDataCon _ _ t' _) = buildDataCon n (fmap snd fs) t
     in (gen, )
      $ zip [0..] fs <&> \(fn, (f, ft)) ->
          let p = take (length fs) $ putBack $ splitAt fn $ repeat PWildcard
              putBack (as, bs) = as <> [PVar "p"] <> bs
           in (f,)
            $ ([] :=> unqualType t' :-> ft,)
            $ lam "z"
            $ case_ "z"
            $ [(PCon n p, "p")]
  
  
  getDictName :: TName -> Type
  getDictName n = TCon . TName (getDictName2 n) $ tKind n
  
  getDictName2 :: TName -> String
  getDictName2 n = "@" <> unTName n
  
  
  getDictTypeForPred :: Pred -> Type
  getDictTypeForPred (IsInst c t) = getDictName c :@@ t
  
  getDict :: Pred -> String
  getDict (IsInst c t) = "@" <> show c <> "@" <> show (normalizeType2 t)
  
  
  
  
  buildDictType
      :: Class
      -> (GenDataCon, [(VName, (Qual Type, Exp VName))])
  buildDictType c@(Class v n ms) =
    second (fmap (second $ first $ dictToConstraint c))
    $
    buildRecord
      (VName name)
      -- TODO(sandy): there is a bug here if there is a constraint on the method
      (fmap (second unqualType) $ M.assocs ms)
      $ Just ((TVar $ TName name $ tKind n) :@@ TVar v)
      -- (Just $ TCon (TName name KStar))
    where
      name = getDictName2 n
  
  
  buildDict :: GenDataCon -> InstRep Pred -> (VName, (Qual Type, Exp VName))
  buildDict gdc (InstRep (_ :=> i@(IsInst c t)) impls) =
      (VName dict,)
        -- TODO(sandy): buggy; doesn't do nested dicts
        -- TODO(sandy): also buggy. we should just run the type checker on this
        $ (sub (Subst $ M.fromList [("a", t)] ) $ gdcFinalType gdc,)
        $ foldl (:@) (LCon (VName dname)) $ M.elems impls
    where
      dict = getDict i
      dname = getDictName2 c
  
  
  dictToConstraint :: Class -> Qual Type -> Qual Type
  dictToConstraint (Class v n _) (qs :=> (_ :-> t)) =
    (IsInst n $ TVar v) : qs :=> t
  
  
  
  normalizeType :: Qual Type -> Qual Type
  normalizeType = schemeType . normalize . Scheme mempty
  
  
  normalizeType2 :: Type -> Type
  normalizeType2 = unqualType . normalizeType . ([] :=>)
  
  
  normalize :: Scheme -> Scheme
  normalize (Scheme _ body) =
      Scheme (fmap snd ord) $ normqual body
    where
      ord = zip (nub . S.toList $ free body) letters <&> \(old, l) ->
        (old, TName l $ tKind old)
      normqual (xs :=> zs) =
        fmap (\(IsInst c t) -> IsInst c $ normtype t) xs :=> normtype zs
  
      normtype (TCon a)    = TCon a
      normtype (a :@@ b)   = normtype a :@@ normtype b
      normtype (TVar a)    =
        case lookup a ord of
          Just x  -> TVar $ TName (unTName x) (tKind x)
          Nothing -> error "type variable not in signature"
  
  