import Data.List
import Data.Char

data Trm = Trm Head [Trm] deriving (Eq) -- applied λ-terms à la Wadsworth
data Head = Var String | Abs String Trm | Symb String Int deriving (Eq)

arity (Symb _ n) = n -- arity of head; length l should be >= arity h for any Trm h l 
arity _ = 0

instance Show Trm where -- pretty print λ-term
  show = oshow False
oshow b (Trm h@(Symb x i) l) = let (l1,l2) = splitAt i l in if i > length l1
  then error ("symbol "++show h++" has arity "++show i++" but only has "++show (length l)++" arguments")
  else (if b then parens else id) ((show h)++"["++intercalate "," (map (oshow False) l1)++"]"++intercalate " " (map (oshow True) l2))
oshow _ (Trm h []) = show h
oshow b (Trm h l) = (if b then parens else id) ((case h of
    (Var _) -> id
    _       -> parens) (show h)++" "++intercalate " " (map (oshow True) l))
parens s = "("++s++")"
instance Show Head where
  show (Var x) = showvar x
  show (Symb x _) = case x of
    (csym:s) -> csym:(map subscript s)
    _             -> x
  show (Abs x t) = "\x03BB"++showvar x++"."++show t -- λ
showvar x = let (v,s) = break (\c -> (c == '_')) x in v++(map subscript (drop 1 s))
subscript d = toEnum (fromEnum d - fromEnum '0' + fromEnum '\x2080')::Char -- renaming keeps names of vars, affixing serial numbers 

va x = Trm (Var x) [] -- skeuomorphisms from Church into Wadsworth; terms should be closed; no vars with underscores or starting with a hash
ab x t = Trm (Abs x t) [] -- λ-abstraction λx.t
sy x i = Trm (Symb x i) [] -- c/i symbol x with arity i
ap t s = hap t [s] -- application (t s) 
hap (Trm h l) m = Trm h (l++m) -- vector application (h l1 ... li m1 ... mj) 

-- β reductions; should be applied to (morally) closed λ-terms only; (λxy.x)y fails but (λxy.x)z is fine (substitution is naïve, no α-conversion)
                                                                                      
whnf t@(Trm (Abs x v) (u:l)) = let Trm h s = subst x u v in t:(whnf (Trm h (s++l))) -- lo-reduction sequence to whnf (not under λ)
whnf t = [t] 
subst x s (Trm h l) = let (Trm h' l') = (case h of (Var y)   | x == y -> s -- naïve substitution
                                                   (Abs y u) | x /= y -> Trm (Abs y (subst x s u)) []
                                                   _                  -> Trm h []) in (Trm h' (l'++(map (subst x s) l)))

nf = rnf (\x -> 1) where -- lo-reduction sequence to nf, including renaming steps λ-introducing underscored indices in vars, at the head
  rnf f t = let
    whnft       = whnf t
    s@(Trm h l) = last whnft in whnft++case h of 
      (Abs x _) -> let v = x++"_"++show (f x) in map (ab v) (rnf (update f x) (ap s (va v)))
      _         -> case l of
        [] -> []
        _ -> tail (map (Trm h) (exchange (map (rnf f) l)))

-- from λ-calculus to first-order term rewriting by generating supercombinators and their rules

symc = '\x03BA' -- first symbol of combinator names, κ kappa (arbitrary; just because it precedes λ and has the right sound for a kombinator)
lift t = rlift ([],t) -- lift λ-term to supercombinator terms, constructed on the fly (monadic, but direct here)
rlift (r,(Trm h l)) = let
  (rh,th) = (case h of
    (Abs x t) -> let (r',t')   = rlift (r,t)
                     (sc,args) = split x 1 t'
                     (n,r'')   = rlookup sc r' in (r'',Trm (Symb (symc:(show n)) (length args)) args)
    _         -> (r,Trm h []))
  (rl,tl) = llift (rh,l) in (rl,hap th tl)
llift p@(r,[]) = p
llift (r,t:l) = let
  (r',t')  = rlift (r,t)
  (r'',l') = llift (r',l) in (r'',t':l')
split x n t@(Trm h l) = if (elem x (vars t)) -- split combinator term into minimal bound context / skeleton and maximal x-free subterms
  then if arity h >= length l
    then case h of
      (Var y) | (x==y)                 -> (Trm (Var "#0") [],[])
      (Symb y i) | arity h == length l -> let (ts,ls) = lsplit x n l in (Trm (Symb y i) ts,ls) 
      _                                -> error ("unexpected head in term "++show t)
    else let (l',t') = unsnoc l
             (tl,ll) = split x n (Trm h l')
             (tr,lr) = split x (n+length ll) t' in (ap tl tr,ll++lr) 
  else (Trm (Var ('#':(show n))) [],[t])
lsplit x n [] = ([],[])
lsplit x n (t:l) = let
  (sc,arg)   = split x n t
  (scs,args) = lsplit x (n+(length arg)) l in (sc:scs,arg++args) 
vars (Trm h l) = concat (map vars l)++case h of -- variables in a combinator term, so fail if there are λs
  (Var x)   -> [x]
  (Abs _ _) -> error ("cannot compute variables for λ-abstraction "++show h)
  _         -> []

-- supercombinator reduction for trs as list [r0,...,rn-1], with ri a combinator over vars "#0".."#n", the rhs of rule for symbol Symb i n

swhnf r t@(Trm h@(Symb (symc:x) i) l) = if (length l < i) -- reduction to whnf of combinator terms, h t1 ... tn with h var or sc with arity h <= n
  then error ("combinator"++show h++" requires "++show i++"but only has "++show (length l)++" arguments in "++show t)
  else (if (length l == i) then [t] else (let (args,l') = splitAt (i+1) l in case (r !? ((read x)::Int)) of
    Nothing    -> error ("no rewrite rule for combinator "++show h)
    (Just rhs) -> t:(swhnf r (hap (lsubst args rhs) l'))))
swhnf _ t@(Trm (Var _) _) = [t]
lsubst l t = lsubstn 1 l where -- substitute vector l1 ... ln l0 for vars "#1" ... "#n" "#0" 
  lsubstn _ [s] = subst "#0" s t
  lsubstn n (s:ss) = subst ('#':(show n)) s (lsubstn (n+1) ss)

swchnf r t = let -- reduction to Cagman--Hindley weak nf 
  swhnft      = swhnf r t
  s@(Trm h l) = last swhnft in swhnft++case l of
    [] -> []
    _  -> tail (map (Trm h)  (exchange (map (swchnf r) l)))

snf r = rsnf (\x -> 1) where -- reduction to nf, supplying fresh vars if head is stuck sc (not enough args); = nf for λ
  rsnf f t = let
    swhnft = swhnf r t
    s@(Trm h l) = last swhnft in swhnft++case h of
      (Symb x n) -> let v = x++"_"++show (f x) in map (ab v) (rsnf (update f x) (ap s (va v)))
      _          -> case l of
        [] -> []
        _  -> tail (map (Trm h) (exchange (map (rsnf f) l)))

expand r (Trm h l) = Trm (hexpand h) (map (expand r) l) where -- map each supercombinator to its λ-abstracted rhs; homomorphically extend 
  hexpand (Symb (symc:rl) i) = absvec i i (Abs "#0" (expand r (r!!((read rl)::Int))))
  hexpand h@(Abs _ _) = error ("head "++show h++" is an abstraction, not a supercombinator")
  hexpand h = h
  absvec 0 i h = h
  absvec n i h = Abs ('#':(show (1+i-n))) (Trm (absvec (n-1) i h) [])

-- generic auxiliary functions

(!?) l i = if (0 <= i && i < length l) then Just(l!!i) else Nothing -- partial !! (in Data.List?)

exchange [] = [] -- exchanging parallelism for length; left--outer
exchange [l] = [[x] | x <- l]
exchange (l:xl) = let exl = exchange xl in map (\x -> x:(head exl)) l ++ tail (map (\y -> last l:y) exl)

update f x y = f y + (if (x==y) then 1 else 0) -- update f by incrementing at x

unsnoc [x] = ([],x) -- unsnoc (in Data.List?)
unsnoc (x:l) = let (l',x') = unsnoc l in (x:l',x')
                                    
rlookup t = rlookupt where -- pair (index of t in l',l') where l' = if t elem l then l else l++[t]
  rlookupt [] = (0,[t])
  rlookupt l@(r:rs) = if (t==r) then (0,l) else let (i,rs') = rlookupt rs in (i+1,r:rs')  

-- for experimentation and / or debugging purposes 

delta = ab "x" (ap (va "x") (va "x")) -- three λ-terms in example 2 of submission δ,Ω,Μ
omega = ap delta delta
example2M = ab "y" (ap (ab "z" (va "y")) omega) -- λ-term M in example 2 of submission
example4trs = fst (lift example2M) -- supercombinator system of example 4 in submission by lifting M
example4trm = snd (lift example2M) -- supercombinator term of example 4 in submission by lifting M
one =  ab "y" (ab "z" (ap (va "y") (va "z"))) -- Church numeral 1
two = ab "y" (ab "z" (ap (va "y") (ap (va "y") (va "z")))) -- Church numeral 2
n256 = ap delta (ap delta two) -- application is exponentiation on Church-numerals so (2^2)^(2^2) evaluates to 256  
cli = ab "x" (va "x") -- the three combinators of CL, I,K,S
clk = ab "x" (ab "y" (va "x"))
cls = ab "x" (ab "y" (ab "z" (ap (ap (va "x") (va "z")) (ap (va "y") (va "z")))))
t 0 = delta
t n = let yn = "y"++(show n) in ab yn (ap (t (n-1)) (ap (va yn) (va yn)))