-- Author: Danila Fedorin
import Data.List

-- Definitions
type Var = String
type Macro = String

data PenMode = Up | Down

data Expr = Variable Var
          | Number Int
          | Sum Expr Expr

data Cmd = Pen PenMode
         | Move Expr Expr
         | Define Macro [Var] Prog
         | Call Macro [Expr]

type Prog = [Cmd]

-- HW Code
-- define line(x1, y1, x2, y2) {
--     pen up; move (x1, y1);
--     pen down; move(x2, y2);
-- }
lineMacro = Define "line" [ "x1", "y1", "x2", "y2" ]
    [ Pen Up
    , Move (Variable "x1") (Variable "y1")
    , Pen Down
    , Move (Variable "x2") (Variable "y2")
    ]

-- define nix(x, y, w, h) {
--     call line(x, y, x + w, y + h);
--     call line(x + w, y, x, y + h);
-- }
nixMacro = Define "nix" [ "x", "y", "w", "h" ]
    [ Call "line"
        [ Variable "x"
        , Variable "y"
        , Sum (Variable "x") (Variable "w")
        , Sum (Variable "y") (Variable "h")
        ]
    , Call "line"
        [ Sum (Variable "x") (Variable "w")
        , Variable "y"
        , Variable "x"
        , Sum (Variable "y") (Variable "h")
        ]
    ]

steps :: Int -> Prog
steps n = base ++ stepList
    where
        base = [ Pen Up, Move (Number 0) (Number 0), Pen Down ]
        stepList = [1..n] >>= (\i -> [ Move (Number (i -1)) (Number i), Move (Number i) (Number i) ])

prettyPrintExpr :: Expr -> String
prettyPrintExpr (Variable v) = v
prettyPrintExpr (Number i) = show i
prettyPrintExpr (Sum l r) = prettyPrintExpr l ++ "+" ++ prettyPrintExpr r

prettyPrintPenMode :: PenMode -> String
prettyPrintPenMode Up = "up"
prettyPrintPenMode Down = "down"

prettyPrintCmd :: Cmd -> [String]
prettyPrintCmd (Pen m) = [ "pen " ++ prettyPrintPenMode m ++ ";" ]
prettyPrintCmd (Move l r) = [ "move(" ++ prettyPrintExpr l ++ ", " ++ prettyPrintExpr r ++ ");" ]
prettyPrintCmd (Define m vs cmds) = ("define " ++ m ++ "(" ++ intercalate ", " vs ++ ") {") : map ("    " ++) (cmds >>= prettyPrintCmd) ++ [ "}" ]
prettyPrintCmd (Call n xs) = [ "call " ++ n ++ "(" ++ intercalate ", " (map prettyPrintExpr xs) ++ ");" ]

prettyPrint :: Prog -> String
prettyPrint prog = intercalate "\n" $ (prog >>= prettyPrintCmd) ++ [ "" ]

macros :: Prog -> [Macro]
macros xs = xs >>= macrosC
    where
        macrosC (Define n vs cs) = n:(cs >>= macrosC)
        macrosC _ = []

sortWith :: Ord b => (a -> b) -> [a] -> [a]
sortWith f = sortBy (\l r -> compare (f l) (f r))

optE :: Expr -> Expr
optE expr = mergeAll $ sortWith exprOrdinal $ listTerms expr
    where
        listTerms (Sum l r) = listTerms l ++ listTerms r
        listTerms e = [ e ]
        exprOrdinal (Number _) = 2
        exprOrdinal _ = 1
        mergeTwo (Number i1) (Number i2) = Number (i1 + i2)
        mergeTwo l r = Sum l r
        mergeAll [x] = x
        mergeAll xs = foldr1 mergeTwo xs

optP :: Prog -> Prog
optP = map optC
    where
        optC (Move e1 e2) = Move (optE e1) (optE e2)
        optC (Call n exs) = Call n (map optE exs)
        optC c = c