{-# LANGUAGE NoMonomorphismRestriction #-}
module PTS.AST
  ( Name
  , Term (..)
  , Stmt (..)
  -- , freevars2constants
  , freevars
  , freshvar
  , freshvarl
  , handlePos
  , allvars
  , subst
  , freshCommonVar) where

import Control.Applicative hiding (Const)
import Control.Monad.Reader

import Data.List (nub, (\\), intersperse)

import Tools.Errors.Class

import Parametric.Error
import Parametric.AST (Name)

-- Syntax

-- the string in NatOp is an identifier for the function. It is necessary
-- to check equivalence of terms (the functions cannot be directly compared)

data (Show c, Eq c) => Term c
  = Nat Int
  | NatOp Name (Int -> Int -> Int) (Term c) (Term c)
  | IfZero (Term c) (Term c) (Term c)
  | Var Name
  | Const c
  | App (Term c) (Term c)
  | Lam Name (Term c) (Term c)
  | Pi Name (Term c) (Term c)
  | Pos Position (Term c)


data Stmt c
  = Bind Name (Term c)
  | Term (Term c)


handlePos f p t = annotatePos p $ Pos p <$> f t

infixl 2 >>>
(>>>) = flip (.)

freevars :: (Eq c, Show c) => Term c -> [String]
freevars (Var x) = [x]
freevars (App t1 t2) = nub $ freevars t1 ++ freevars t2
freevars (NatOp _ _ t1 t2) = nub $ freevars t1 ++ freevars t2
freevars (IfZero t1 t2 t3) = nub $ freevars t1 ++ freevars t2  ++ freevars t3
freevars (Lam x t1 t2) = nub $ freevars t1 ++ ((freevars t2) \\ [x])
freevars (Pi x t1 t2) = nub $ freevars t1 ++ ((freevars t2) \\ [x])
freevars (Pos p t) = freevars t
freevars _ = []

freshvar :: (Eq c, Show c) => Term c -> String -> String
freshvar t x = freshvarl (freevars t) x

freshvarl :: [String] -> String -> String
freshvarl xs x = freshvarhelper 0
  where
  freshvarhelper i = if (x++((show i)) `elem` xs) then freshvarhelper (i+1) else (x++((show i)))


freshCommonVar 
  :: (Eq c, Show c) 
  => Name -> Name -> Term c -> Term c -> (Name, Term c, Term c)

freshCommonVar n1 n2 b1 b2 = (n', b1', b2') where
  n' | n1 == n2 = n1
     | n1 `notElem` fv2 = n1
     | n2 `notElem` fv1 = n2
     | otherwise = freshvarl (fv1 ++ fv2) n1
     
  b1' | n' == n1 = b1
      | otherwise = subst b1 n1 (Var n')
  
  b2' | n' == n2 = b2
      | otherwise = subst b2 n2 (Var n')
  
  fv1 = freevars b1
  fv2 = freevars b2

-- substitution (generates fresh variables if needed to prevent accidential capture)

avoidCapture :: (Show c, Eq c) => Term c -> Name -> Name -> Term c -> (Name, Term c)
avoidCapture s x y t = (x', s') where
  x' | x `elem` fvt = freshvarl fv x
     | otherwise = x

  s' | x == x' = s
     | otherwise = subst s x (Var x')
  
  fvs = freevars s
  fvt = freevars t
  fv = y : fvs ++ fvt
  
subst :: (Show c, Eq c) => Term c -> String -> Term c-> Term c
subst (Var y) x t' | y == x = t'
subst (App t1 t2) x t' = App (subst t1 x t') (subst t2 x t')
subst (NatOp i o t1 t2) x t' = NatOp i o (subst t1 x t') (subst t2 x t')
subst (IfZero t1 t2 t3) x t' = IfZero (subst t1 x t') (subst t2 x t')  (subst t3 x t')
subst q@(Lam y t1 t2) x t' | x /= y =
  let (newy, newt2) = avoidCapture t2 y x t' 
  in Lam newy (subst t1 x t') (subst newt2 x t')
subst (Lam y t1 t2) x t' | x == y = Lam y (subst t1 x t') t2
subst q@(Pi y t1 t2) x t' | x /= y =
  let (newy, newt2) = avoidCapture t2 y x t'
  in Pi newy (subst t1 x t') (subst newt2 x t')
subst (Pi y t1 t2) x t' | x == y = Pi y (subst t1 x t') t2
subst (Pos p t) x t' = Pos p (subst t x t') -- delete pos annotation here?
subst t _ _ = t
  
allvars :: (Eq c, Show c) => Term c -> [String]
allvars (Var x) = [x]
allvars (App t1 t2) = nub $ allvars t1 ++ allvars t2
allvars (NatOp _ _ t1 t2) = nub $ allvars t1 ++ allvars t2
allvars (IfZero t1 t2 t3) = nub $ allvars t1 ++ allvars t2  ++ allvars t3
allvars (Lam x t1 t2) = nub $ allvars t1 ++ allvars t2 ++ [x]
allvars (Pi x t1 t2) = nub $ allvars t1 ++ allvars t2 ++ [x]
allvars (Pos p t) = allvars t
allvars _ = []
