import Data.List
import Control.Monad

-- import Data.Either.Utils

maybeToEither :: e -> Maybe a -> Either e a
maybeToEither e Nothing  = Left e
maybeToEither _ (Just x) = return x

distinct :: Eq a => [a] -> Bool
distinct xs = length (nub xs) == length xs


type Check a = Either String a
type Type = String
type Var  = String
type FSym = String
type FSym_Info = ([Type], Type)
type Vars = Var -> Check Type
type Sig = FSym -> Check FSym_Info
data Term = 
  Var Var 
  | Fun FSym [Term]
  deriving (Eq, Show)

failure :: String -> Either String a
failure = Left

assert :: Bool -> String -> Either String ()
assert p e = if p then return () else failure e

type_check :: Sig -> Vars -> Term -> Check Type
type_check sig vars (Var x) = vars x
type_check sig vars t@(Fun f ts) = do
  (tys_in,ty_out) <- sig f
  tys_ts <- mapM (type_check sig vars) ts
  assert (tys_ts == tys_in) (show t ++ " ill-typed") 
  return ty_out

infer_type :: Sig -> Type -> Term -> Check [(Var,Type)]
infer_type sig ty (Var x) = return [(x,ty)]
infer_type sig ty t@(Fun f ts) = do
  (tys_in,ty_out) <- sig f
  assert (length tys_in == length ts) "lengths don't match"
  assert (ty_out == ty) "problem with expected type"
  vars_l <- mapM (uncurry (infer_type sig)) (zip tys_in ts)
  let vars = nub (concat vars_l) 
  assert (distinct (map fst vars)) "conflicting types of variables"
  return vars
  
type_check_eqn sigma (Var x, r) = failure "var as lhs"
type_check_eqn sigma (l @ (Fun f _), r) = do
  (_,ty) <- sigma f
  vars <- infer_type sigma ty l
  ty_r <- type_check 
     sigma 
     (\ x -> maybeToEither 
         (x ++ " is unknown variable") 
         (lookup x vars))
     r
  assert (ty == ty_r) "types of lhs and rhs don't match"


type Subst = Var -> Term
type Subst_List = [(Var,Term)]
  
  
-- Exercise 1 - Matching Algorithm

data Match_Error = Fun_Var FSym Var | Fun_Clash FSym FSym | Var_Clash Var Term Term

type Match_Result a = Either Match_Error a

-- existing implementation
match :: Term -> Term -> Maybe Subst
match l t = do
  xt_list <- match_list [(l,t)]
  Just (\ x -> case lookup x xt_list of
     Just s -> s
     Nothing -> Var x)

match_list :: [(Term,Term)] -> Maybe Subst_List
match_list [] = return []
match_list ((Fun f ls, Fun g ts) : pairs) =
  if f == g then match_list (zip ls ts ++ pairs) -- decompose
  else Nothing                                   -- fun-clash
match_list ((Fun _ _, Var _) : _) = Nothing      -- fun-var
match_list ((Var x, t) : pairs) = do
  xt_list <- match_list pairs
  case lookup x xt_list of
    Nothing -> Just ((x,t) : xt_list)            -- keep
    Just s -> if s == t then Just xt_list        -- remove duplicate
    else Nothing                                 -- var-clash
  
-- write new implementation: using copy-and-paste from above  
match_A :: Term -> Term -> Check Subst
match_A = undefined

-- write new implementation: using copy-and-paste from above  
match_B :: Term -> Term -> Match_Result Subst
match_B = undefined



-- Exercise 3 - Processing Programs


data Data_Definition = Data Type [(FSym, FSym_Info)]
data Function_Definition = Function FSym FSym_Info [(Term,Term)]
type Functional_Prog = ([Data_Definition],[Function_Definition])

type Sig_List = [(FSym, FSym_Info)]
type Defs = Sig_List
type Cons = Sig_List
type Equations = [(Term,Term)]

data Prog_Info = Prog_Info [Type] Cons Defs Equations deriving Show

initial_prog_info = Prog_Info [] [] [] []

equations :: Prog_Info -> Equations
equations (Prog_Info _ _ _ eqs) = eqs

process_data_definition :: Prog_Info -> Data_Definition -> Check Prog_Info
process_data_definition pi@(Prog_Info tys cons defs eqs) (Data ty new_cs) = do
  -- check that type is fresh
  assert (not (elem ty tys)) "type is not new"
  let new_tys = ty : tys
  -- check distinctness of new constructor names
  assert (distinct (map fst new_cs)) "constructors not distinct"
  -- check fresh constructor names
  assert (all (\ (c,_) -> lookup c (cons ++ defs) == Nothing) new_cs) "constructors not new"
  -- check types of constructors
  assert (all (\ (_,(tys_in,ty_out)) -> ty_out == ty  
                 && all (\ ty -> elem ty new_tys) tys_in) new_cs) 
     "problems in types of constructors"
  -- check existence of non-recursive constructor
  assert (any (\ (_,(tys_in,_)) -> all (/= ty) tys_in) new_cs)
    "no non-recursive constructor"
  return (Prog_Info new_tys (new_cs ++ cons) defs eqs)

process_data_definitions :: Prog_Info -> [Data_Definition] -> Check Prog_Info
process_data_definitions = foldM process_data_definition

linear :: Term -> Bool
linear t = undefined

check_equation :: 
  Sig_List ->       -- defined symbols, including f
  Sig_List ->       -- constructors
  FSym ->           -- f
  FSym_Info ->      -- type of f
  (Term, Term) ->   -- equation (l,r)
  Check ()
check_equation defs cons f (tys_in, ty_out) (l,r) = do
  -- check that l is linear
  -- check that types of l and r are identical
  -- check that l is of form f(pat_1,..,pat_n)
  -- check var-condition 
  return ()
  
  
process_function_definition :: Prog_Info -> Function_Definition -> Check Prog_Info
process_function_definition = undefined
  
process_function_definitions :: Prog_Info -> [Function_Definition] -> Check Prog_Info
process_function_definitions = foldM process_function_definition

process_program :: Functional_Prog -> Check Prog_Info
process_program (data_defs, fun_defs) = do
  pi <- process_data_definitions initial_prog_info data_defs
  process_function_definitions pi fun_defs
  
example_prog :: Functional_Prog
example_prog = (
  [
   Data "Nat" [("Zero",([],"Nat")), ("Succ",(["Nat"],"Nat"))],
   Data "List" [("Nil",([],"List")), ("Cons",(["Nat","List"],"List"))]
  ],
  [
    Function "add" (["Nat","Nat"],"Nat") [
      (Fun "add" [Fun "Succ" [Var "x"], Var "y"], 
       Fun "Succ" [Fun "add" [Var "x", Var "y"]]),
      (Fun "add" [Fun "Zero" [], Var "y"], 
       Var "y")
      ],
    Function "append" (["List","List"],"List") [
      (Fun "append" [Fun "Cons" [Var "x",Var "xs"], Var "ys"], 
       Fun "Cons" [Var "x", Fun "append" [Var "xs", Var "ys"]]),
      (Fun "append" [Fun "Nil" [], Var "ys"], 
       Var "ys")
      ]
  ])


test = process_program example_prog