module Main
where


import IO (getLine)
import List
import Text.ParserCombinators.Parsec

{-
Desc: 

This program is inteneded to help writing clear programs
in matlab (vectorized).
Because when one writes A = repmat(reshape(permute(c,[3,1]),[1,m,1,...
its not really obvious what the author meant. Its hard to read, write
and debug.

This program defines a function that gets a "semi-matlab" [ more accuratly, 
math notation] code
that's readable, and produces a pure matlab code without any loops,
but since you know its equiv to whatever was the input, then
you don't debug the output of this program, but debug the input
to this program, which is much easier.

Examples: (input then output)
[binary : ~ilya/bin/vectorizer 
just paste the code and copy the result back]
















A(i1,i2,i3,i4,i5) = sum(t=1000,v=1000| B(t,i2,v).*sum(i=100|A(i,i2)))



B(i1)             = sum(t=1000,v=1000| B(t,i1,v).*sum(i=100|A(i,i1))


dx(i,j) = sum(:=dim_x|  (x(i,:)-x(j,:)).^2())


q(i,j)=sum(c=m | fi(i,c).*r(i,c,j))


grad_y(u,v)=sum(j=n|dy2(u,v,j).*(p1(u,v,j)-s(u,v,j)+p1(j,v,u)-s(j,v,u)))


sum1(k,e) = sum(c=m, i=n| ndelta(i,k).*fi(i,c).*rdy2(i,c,k,e).*(U(i,c)-pdivq(i,k).*B(i,c,k)))


















A[n,n,m,n,a](i1,i2,i3,i4,i5) = sum(t=1000,v=1000| B(t,i2,v).*sum(i=100|A(i,i2)))

A=repmat(reshape(             sum(permute(sum(B.*repmat(reshape(sum(permute(A,[2,1]),2),[1,n,1]),[1000,1,1000]),3),[2,1]),2)                  ,[1,n,1,1,1]),[n,1,m,n,a]);













B[n](i1)        = sum(t=1000,v=1000| B(t,i1,v).*sum(i=100|A(i,i1)))

% compare to the above, both the input and the output [I added the spaces]

B=     sum(permute(sum(B.*repmat(reshape(sum(permute(A,[2,1]),2),[1,n,1]),[1000,1,1000]),3),[2,1]),2);











dx[n,n](i,j) = sum(:=dim_x|  (x(i,:)-x(j,:)).^2())

dx=sum((repmat(reshape(x,[n,1,dim_x]),[1,n,1])-repmat(reshape(x,[1,n,dim_x]),[n,1,1])).^2,3);








q[n,n](i,j)=sum(c=m | fi(i,c).*r(i,c,j))

q=sum(repmat(reshape(fi,[n,1,m]),[1,n,1]).*permute(r,[1,3,2]),3);






grad_y[n,m](u,v)=sum(j=n|dy2(u,v,j).*(p1(u,v,j)-s(u,v,j)+p1(j,v,u)-s(j,v,u)))

grad_y=sum(dy2.*(p1-s+permute(p1,[3,2,1])-permute(s,[3,2,1])),3);




sum1[n,m](k,e) = sum(c=m, i=n| ndelta(i,k).*fi(i,c).*rdy2(i,c,k,e).*(U(i,c)-pdivq(i,k).*B(i,c,k)))



sum1=sum(permute(sum(repmat(reshape(permute(ndelta,[2,1]),[n,1,1,n]),[1,m,m,1]).*repmat(reshape(permute(fi,[2,1]),[1,m,1,n]),[n,1,m,1]).*permute(rdy2,[3,2,4,1]).*(repmat(reshape(permute(U,[2,1]),[1,m,1,n]),[n,1,m,1])-repmat(reshape(permute(pdivq,[2,1]),[n,1,1,n]),[1,m,m,1]).*repmat(reshape(permute(B,[3,2,1]),[n,m,1,n]),[1,1,m,1])),4),[1,3,2]),3);


This should be clear

-}





to_index :: (Show a,Eq a)=>[a]->a-> Integer
to_index [] elem = error ("to_index: elem " ++ (show elem)++" not found")
to_index (x:xs) elem 
	| x == elem =  0
	| x /= elem = 1+ to_index xs elem
			

from_index :: (Show a,Eq a)=>[a]->Integer-> a
from_index (x:xs) 0 =  x
from_index (x:xs) n = from_index xs (n-1)
from_index []     _ = error "from_index: index too big"

-- length of dimentions list (its infinite with inf 1's in the end)
dlen (x:xs) 
	| x == "1" = 0
	| x /= "1" = 1 + dlen xs


sort_using ind ind1 = 
	let a = map (to_index ind) ind1
	    b = sort a
	    c = map (from_index ind) b
	in map (to_index ind1) c


fill_gaps_with w (x:xs) (f:fx) 
	| x == f = x:fill_gaps_with w xs fx
	| x /= f = w:fill_gaps_with  w (x:xs) fx

fill_gaps_with w [] fx =  [ w | _ <- fx]
fill_gaps_with a b c = error ("fill_gaps_with: " ++ show a ++ show b ++ show c)


add_to_dims (x:xs) el  
	| x == "1" = el:xs
	| x /= "1" = x:add_to_dims xs el
			
add_to_dims [] el = error "add_to_dims: not a dims list"


sorted (x:[]) = True
sorted (x:y:xs) = x < y && sorted (y:xs)
sorted [] = True


-- gets free variables from an expression.
free :: Expr -> [Var]
free (Sum exprs ops) = nub $ concat $  map free exprs
free (Mat m vars) = vars
free (VecFn f var range d) = filter (/= var) (free d) -- var is not free inside, so filter it out.
free (Parens f) = free f


all_vars :: Expr -> [Var]
all_vars (Sum exprs ops) = concat $ map all_vars exprs
all_vars (Mat m vars) = vars
all_vars (VecFn f var range d) = all_vars d
all_vars (Parens f) = all_vars f


ones = "1":ones



type Var = String
type Range = Var
type Op = Var
type Fn = Var
data Expr = Sum [Expr] [Op] | VecFn Fn Var Range Expr | Mat Var [Var] | Parens Expr
	deriving Show

prep :: Expr -> [Var] -> [Var] -> String

prep (Sum (m:mx) (o:ox)) ind dims = 
	prep m ind dims ++ o ++ 
	prep (Sum mx ox) ind dims 
prep (Sum (m:[]) []) ind dims = 
	prep m ind dims
prep (Sum _ _) _ _ = error "prep:Sum wrong lengths of lists" 
-- but the above should not usually happen.

-- take care of parens in code.
prep (Parens x) ind dims = "(" ++ prep x ind dims ++")"

 



prep (VecFn fn x range f) ind dims = 
    let free_vars =  filter (/=x) (free f)
	new_dims = map (from_index dims)
		       (map (to_index ind) free_vars)

	g = fn++"("++prep f (free_vars ++ [x])
	                    (new_dims  ++ [range])
	           ++"," ++ show (length new_dims + 1)
              ++")"
    in repmat_reshape_permute g
		   ind free_vars dims



prep (Mat c []) ind dims = c  -- if c is a scalar.
prep (Mat c ind1) ind dims = 
    repmat_reshape_permute c ind ind1 dims 


--convinience
repmat_reshape_permute g ind ind1 dims 
	 = repmat_reshape (permute g ind ind1) ind ind1 dims



permute :: String -> [Var] -> [Var] -> String
permute c ind ind1 = 
    let	res = map (+1) (sort_using ind ind1)
    in 	if sorted res   -- indeces withrespect to ind.
	  then c --if sorted - then no need to permute.
	else
	  "permute(" ++ c ++ "," ++ show res ++ ")"


-- converst a matrix with indeces ind1 to matrix with indeces ind (using sizes dims)

repmat_reshape :: String -> [Var] -> [Var] -> [Var] -> String
repmat_reshape d ind ind1 dims = 
    let	s = sort $ map (to_index ind) ind1
	-- in this code I need dims to be infinite, so:
	-- and will not use dims anymore in this function.
	dims0 = dims ++ ones 
	-- s has the sorted indecis 
	f = fill_gaps_with (dlen dims0 + 1) s [0..toInteger(length ind)-1]
	g = map (from_index dims0) f
	h = [u `text_div` v | (u,v) <- zip dims0 g]
	    
        -- if the ones are contained just in the end, reshape
	-- is useless, so first - remove all 'ending ones'
	-- and then see if reshape is needed.

	reshape = "reshape(" ++ d ++ "," ++ format g ++ ")"
     in 
	if length ind1 == length ind 
	then d else
	"repmat(" ++ reshape ++ "," ++ format h ++ ")"


--to reduce size of produced code.
u `text_div` "1" = u
u `text_div` v 
  | u == v = "1"
  | u /= v = u ++ "/" ++ v

format :: [String]-> String
format x = "[" ++ format1 x ++ "]"

format1 (x:y:[]) = x ++"," ++ y
format1 (x:[]) = x
format1 (x:xs) = x ++ "," ++ format1 xs
format1 [] = ""



prep_vec f var range d ind dims = 
    let f = filter (/=var) (free d)
	new_dims = map (from_index dims) (map (to_index ind) f)
	in "<<" ++  show f ++ " " ++ (show $ take 5 new_dims) ++ ">>"




--------------------------------------------------
--PARSING-----------------------------------------
--Will parse a matlab expression. 
--It will not parse matlab fully, but just enough to
--get this running. I will learn how to parse and do it.
--so vectorized code will be easy in matlab


-- Note that : is a variable name, to trick matlab.
var = many1 $ oneOf ( ['a'..'z'] ++ ['A'..'Z'] ++ ['0'..'9']++"_"++":" )
op = many1 $ oneOf ( ['*','.','^','+','-','/'] )

sp = many space

sum_expr = 
     do	x <- not_sum_expr;sp -- can't be a sum.
	o <- op		 ;sp
	z <- expr -- if its a sum then the result is accamulated
	return 
	 (case z of
	 	(Sum mx ox) -> Sum (x:mx) (o:ox)
		u ->	       Sum [x,u] [o]) --if not its also:)

q_expr = 
     do	char '(' 	;sp
	u <- expr 	;sp
	char ')'
	return$  Parens u
	

{-vec_fn_expr = 
     do	f <- var ;sp
	char '(' ;sp
	v <- var ;sp
	char '=' ;sp 
	r <- var ;sp
	char '|' ;sp
	d <- expr;sp
	char ')' ;sp
	return (VecFn f v r d)
-}

-- something like "sum(i=n, b=m, ...)


vec_fn_expr = 
    do f <- var   ;sp
       char '('   ;sp
       inp <- sepBy1 var_range (char ',')  ;sp-- get list of parameters. 
       char '|'   ;sp
       d <- expr  ;sp
       char ')'	      
       return (vec_fn_aux f inp d)
	      
vec_fn_aux :: Var-> [(Var, Var)] -> Expr-> Expr

vec_fn_aux f [] d = d  --produce the reuslts from the input.
vec_fn_aux f ((var, range):xs) d = 
    VecFn f var range (vec_fn_aux f xs d)



var_range = 
    do sp
       v <- var ;sp
       char '=' ;sp
       r <- var ;sp
       return (v,r);  -- pair of variable and range




mat_expr = 
    do	m <- var  			;sp	
	char '('			;sp
	g <- sepBy var (char ',')	;sp
	char ')'			;sp
	return (Mat m g)

not_sum_expr = try q_expr <|> try vec_fn_expr <|> try mat_expr

expr =  try sum_expr <|> not_sum_expr

---------------------------------------------------
--It parses now. So last step: to add a shell,
--so that it would be just a bit user friendly.





trans = parse shell1 "" 
shell1 =
     do	sp;
	assign_to <- var; sp;
        sp;char '[' 			;sp
	dims <- sepBy var (char ',')    ;sp
	char ']' 			;sp
	char '(' 			;sp
	ind <- sepBy var (char ',')     ;sp
	char ')' 			;sp
	char '=' 			;sp 
	p <- expr
 	return $ 
	       assign_to ++ "=" ++ prep p ind dims ++ ";" 



forever x = x >> forever x


main :: IO ()
main=
    putStr "Vectorizer 4.0. type your expr, get matlab expr\n" >>
    forever
    (do  x<-getLine
	 case trans x of
           Right s -> putStr s
	   Left e  -> putStr (show e)
	 putStr "\n"
	)