module pf

import StdBool
import StdFile
from StdFunc import flip
import StdMisc
import StdString
import StdTuple

import GenEq

import Control.Applicative
import Control.Monad
import Data.Either
from Data.Func import $
import Data.Functor
import Data.List
import Data.Maybe
import System.CommandLine

import Yard

:: Expr
	= Lambda Ident Expr
	| Ident Ident
	| Literal Literal
	| (@) infixl Expr Expr

:: Literal
	= LitInt Int
	| LitBool Bool

:: Ident :== String

:: Token
	= TBackSlash
	| TParenOpen
	| TParenClose
	| TArrow
	| TBool Bool
	| TInt Int
	| TIdent Ident

derive gEq Token; instance == Token where == a b = a === b

tokenize :: ([Char] -> [Token])
tokenize = flip tok []
where
	tok :: [Char] [Token] -> [Token]
	tok []           tks = reverse tks
	tok ['\\':cs]    tks = tok cs [TBackSlash:tks]
	tok ['(':cs]     tks = tok cs [TParenOpen:tks]
	tok [')':cs]     tks = tok cs [TParenClose:tks]
	tok ['-':'>':cs] tks = tok cs [TArrow:tks]
	tok ['T':'r':'u':'e':cs]     tks = tok cs [TBool True:tks]
	tok ['F':'a':'l':'s':'e':cs] tks = tok cs [TBool False:tks]
	tok [c:cs]       tks
	| isSpace c          = tok cs tks
	| isDigit c          = tok rest [TInt (toInt (toString digs)):tks]
		with (digs,rest) = span isDigit [c:cs]
	| isIdentChar c      = tok rest [TIdent (toString ids):tks]
		with (ids,rest) = span isIdentChar [c:cs]

	isIdentChar :: Char -> Bool
	isIdentChar c = isAlpha c || c == '_'

parse :: String -> Maybe Expr
parse s = case runParser expr (tokenize $ fromString s) of
	(Right e, []) -> Just e
	_             -> Nothing
where
	simple :: Parser Token Expr
	simple = liftM2 lambda (item TBackSlash *> some ident <* item TArrow) expr
		<|> Literal <$> LitInt  <$> int
		<|> Literal <$> LitBool <$> bool
		<|> item TParenOpen *> expr <* item TParenClose
		<|> Ident <$> ident
	where
		lambda :: [Ident] Expr -> Expr
		lambda [x]    e = Lambda x e
		lambda [x:xs] e = Lambda x (lambda xs e)

	expr :: Parser Token Expr
	expr = liftM2 (\f xs -> app [Ident f:xs]) ident (some simple)
		<|> simple
	where
		app :: [Expr] -> Expr
		app [e] = e
		app es = app (init es) @ last es

	ident :: Parser Token Ident
	ident = (\(TIdent t) -> t) <$> satisfy (\t -> t=:(TIdent _))

	int :: Parser Token Int
	int = (\(TInt i) -> i) <$> satisfy (\t -> t=:(TInt _))

	bool :: Parser Token Bool
	bool = (\(TBool b) -> b) <$> satisfy (\t -> t=:(TBool _))

hasIdent :: Ident Expr -> Bool
hasIdent id (Ident n) = id == n
hasIdent id (Literal _) = False
hasIdent id (Lambda n e) = id <> n && hasIdent id e
hasIdent id (f @ x) = hasIdent id f || hasIdent id x

print :: Expr -> String
print e = foldl (+++) "" (pr False e [])
where
	pr :: Bool Expr [String] -> [String]
	pr p l=:(Lambda _ _) st =
		[if p "(\\" "\\":intersperse " " vars] ++
		[" -> ":pr False rhs [if p ")" "":st]]
	where
		(vars,rhs) = getLambdas l

		getLambdas :: Expr -> ([Ident], Expr)
		getLambdas (Lambda x rhs) = let (is,e) = getLambdas rhs in ([x:is],e)
		getLambdas e = ([], e)

	pr _ (Ident id)  st = [id:st]
	pr _ (Literal l) st = pLit l st
	pr p (f @ x)     st = [if p "(" "":pr False f [" ":pr True x [if p ")" "":st]]]

	pLit :: Literal [String] -> [String]
	pLit (LitInt i)  st = [toString i:st]
	pLit (LitBool b) st = [toString b:st]

optim :: Expr -> Expr
optim (Lambda id e)
| hasIdent id e = case optim e of
	Ident id` -> if (id == id`) (Ident "id") (Ident "const" @ e)
	e=:(f @ x) -> optim (optApp e)
	e -> Lambda id e
| otherwise = Ident "const" @ optim e
where
	optApp :: Expr -> Expr
	optApp (f @ x)
	| hasIdent id f && hasIdent id x = Ident "ap" @ Lambda id f @ Lambda id x
	| otherwise = case moveOutside id (f @ x) of
		a=:(f @ Ident id`) -> if (id == id`) f (Ident "const" @ a)
		a                  -> Ident "const" @ a
optim (Ident "flip" @ Ident "ap" @ Ident "id") = Ident "join"
optim (f @ x) = optim f @ optim x
optim e = e

moveOutside :: Ident Expr -> Expr
moveOutside _ (Ident id) = Ident id
moveOutside _ (Literal l) = Literal l
moveOutside id (f @ x)
| hasIdent id f && hasIdent id x = abort "moveOutside on @ with id in both sides"
| hasIdent id f = case moveOutside id f of
	Ident id -> Ident "flip" @ Ident "id" @ x @ Ident id
	g @ Ident id -> Ident "flip" @ g @ x @ Ident id
| hasIdent id x = case moveOutside id x of
	e=:(Ident id) -> f @ e
	g @ e=:(Ident id) -> Ident "(o)" @ f @ g @ e
moveOutside _ (Lambda x e) = Lambda x e // TODO

Start w
# ([prg:cmd],w) = getCommandLine w
| length cmd <> 1 = err ("Usage: " +++ prg +++ " EXPRESSION") w
# e = parse (hd cmd)
| isNothing e = err "Expression could not be parsed." w
# e = fromJust e
# (io,w) = stdio w
# io = io <<< "Request: " <<< print e <<< "\n"
# io = io <<< "Result:  " <<< print (optim e) <<< "\n"
# (_,w) = fclose io w
= w
where
	err :: a *World -> *World | <<< a
	err e w = snd (fclose (stderr <<< e <<< "\n") w)