{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}

module Types where

import Syntax
import Data.List
import Control.Applicative
import Debug.Trace

---------------------------
-- TYPES FOR TransformeR --
---------------------------

class RefTy t 
  where tyOK :: t -> Bool

data RecTy = RecTy [(Label, Ty)] -- labeled fields
  deriving (Show, Eq)

data Ty = NumTy (Number, Number) -- lower and upper bound
        | CatTy [Category] -- set of catagories
  deriving (Show, Eq)

type Label = String

type Gamma = [(Name, Either RecTy Ty)] -- type environment

mapsTo g l t = (l, t) : g

---------------------
-- FORMATION RULES --
---------------------

instance RefTy Ty
  where tyOK (NumTy (n1, n2)) = n1 <= n2
        tyOK (CatTy cs) = isSet cs

instance RefTy RecTy
  where tyOK (RecTy lts) = all tyOK taus && isSet ls
          where (ls, taus) = unzip lts

---------------------
-- SUBTYPING RULES --
---------------------

instance Ord TransTy
  where TransTy rtauIn rtauOut <= TransTy rtauIn' rtauOut' 
          = rtauIn == rtauIn' && rtauOut <= rtauOut'

instance Ord RecTy
  where RecTy lts <= RecTy lts' 
          = (length lts == length lts') 
            && labs == labs'
            && all id (zipWith (<=) taus taus')
            where
              (labs, taus) = unzip lts
              (labs', taus') = unzip lts'

instance Ord Ty 
  where NumTy (lb, ub) <= NumTy (lb', ub') = lb' <= lb && ub <= ub'
        CatTy cs <= CatTy cs' = all (\c -> elem c cs') cs

-------------------------
-- TYPE SYSTEM/CHECKER --
-------------------------

data TransTy = TransTy RecTy RecTy 
  deriving (Show, Eq)

tyTrans :: Transformation -> TransTy
tyTrans (TRANS arg sigIn sigOut b) 
  | tyOK tauIn && tyOK tauOut && tau' <= tauOut 
              = TransTy tauIn tauOut
  | otherwise = undefined
  where (Left tau', _) = ty [(arg, Left tauIn)] b
        tauOut    = tySig sigOut
        tauIn     = tySig sigIn

readSig :: Sig -> [(Name, Ty)]
readSig (Sig sig) = map (\(l, s) -> (l, tyDes s)) sig
  where tyDes (INTERVAL lb ub) = NumTy (lb, ub)
        tyDes (SET ns) = CatTy ns

tySig :: Sig -> RecTy
tySig = RecTy . readSig

ty :: Gamma -> Expression -> (Either RecTy Ty, Gamma)
ty gamma exp =
  let ty' (PROJ e l) = (Right fTau, gamma')
        where (rTau, gamma') = ty gamma e
              Just fTau = lookup l rTauR  
              Left (RecTy rTauR) = rTau
      ty' (OP opr es) 
        | tyOK tauOut = (Right tauOut, gamma')
        where tauOut = tyJoin opr taus
              (gamma', taus) = foldl tyNext (gamma, []) es
              tyNext (g, tau) e = let (t, g) = ty gamma e in (g, t:tau)
      ty' (LIT (NUM n)) = (Right $ NumTy (n, n), gamma)
      ty' (LIT (CAT c)) = (Right $ CatTy [c], gamma)
      ty' (LIT (MAP fvs)) = (ty' . REC) $ map (\(f, v) -> (f, LIT v)) fvs
      ty' (REC fields) 
        | tyOK tauRec = (Left tauRec, gamma')
        | otherwise = undefined
            where tauRec = RecTy taus
                  (taus, gamma')          = foldl tyNext ([], gamma) fields
                  tyNext (taus, g) (l, e) = 
                      let (Right t, g') = ty g e in (taus ++ [(l,t)], g')
      ty' (VAR x) = let Just tau = lookup x gamma in (tau, gamma)
      ty' (MUTATE x l e) = (Left tau', mapsTo gamma' x (Left tau'))
        where Just (Left (RecTy lts))= lookup x gamma
              (Right sigma, gamma') = ty' e
              tau' = if elem l $ map fst lts
                     then RecTy $ map (\(l', t) -> if l == l' then (l, sigma) else (l, t)) lts
                     else RecTy $ lts ++ [(l, sigma)]
      ty' (SEQ e1 e2) = ty gamma' e2 
        where (_, gamma') = ty' e1
      ty' (ASSIGN x e) = (t, mapsTo gamma' x t)
        where (t, gamma') = ty' e
  in ty' exp

-----------------------------
-- RULES FOR JOINING TYPES --
-----------------------------

tyJoin SUM [Right (NumTy (lb1, ub1)), Right (NumTy (lb2, ub2))]
      | lb1 <= ub1 && lb2 <= ub2 = NumTy (lb1 + lb2, ub1 + ub2)
      | otherwise = undefined
tyJoin CONCAT [Right (CatTy cs1), Right (CatTy cs2)]
      | isSet cs1 && isSet cs2 = CatTy . nub $ cs1 ++ cs2
      | otherwise = undefined
tyJoin _ _ = undefined


----------------------------
-- Mother's little helper --
----------------------------

isSet xs = length xs == length (nub xs)