module Optimiser(optimise) where

import Data.List
import qualified Data.Map.Strict as Map

-- import Debug.Trace

import AST


type Optimisation = [Instruction] -> [Instruction]


optimisations :: [Optimisation]
optimisations =
    [collectSimilar, nullOps, collectAdds, propagateKnowns, firstSets, propagateSlides, uselessEnd, specialLoops, normaliseOrder]

composedOpts :: Optimisation
composedOpts = foldl1 (.) (reverse optimisations)
-- composedOpts = foldl1 (.) (map (traceShowId .) $ reverse optimisations)
-- composedOpts = foldl1 (.) (map ((\r -> traceShow (take 7 r) r) .) $ reverse optimisations)

optimise :: Program -> Program
optimise (Program inss) =
    Program $ repeated composedOpts inss


collectSimilar :: Optimisation
collectSimilar [] = []
collectSimilar (IAdd x off1 : IAdd y off2 : rest) | off1 == off2 =
    collectSimilar $ IAdd (x + y) off1 : rest
collectSimilar (ISlide off1 : ISlide off2 : rest) =
    collectSimilar $ ISlide (off1 + off2) : rest
collectSimilar (ILoop inss off : rest) = ILoop (collectSimilar inss) off : collectSimilar rest
collectSimilar (ins : rest) = ins : collectSimilar rest

nullOps :: Optimisation
nullOps [] = []
nullOps (IAdd 0 _ : rest) = nullOps rest
nullOps (ICopy _ _ 0 : rest) = nullOps rest
nullOps (ISlide 0 : rest) = nullOps rest
nullOps (ILoop inss off : rest) = ILoop (nullOps inss) off : nullOps rest
nullOps (ins : rest) = ins : nullOps rest

collectAdds :: Optimisation
collectAdds [] = []
collectAdds (ICopy from1 to1 m1 : ICopy from2 to2 m2 : rest)
    | from1 == from2, m1 == m2 = ICopy from1 to1 m1 : collectAdds (ICopy to1 to2 1 : rest)
collectAdds inss =
    let adds = map (\(IAdd v o) -> (v, o)) $ takeWhile isIAdd inss
        dests = nub $ map snd adds
        collected = [(sum $ map fst $ filter ((== d) . snd) adds, d) | d <- dests]
        rest = drop (length adds) inss
        rest' = case rest of
            [] -> []
            (i:is) -> i : collectAdds is
    in map (uncurry IAdd) collected ++ rest'

propagateKnowns :: Optimisation
propagateKnowns [] = []
propagateKnowns allinss@(ins : _) | isISet ins || isIStart ins =
    let defaultZero = isIStart ins

        go :: [Instruction] -> Map.Map Offset Byte -> ([Instruction], Map.Map Offset Byte)
        go (IStart : rest) mp = go rest mp
        go (ISet val off : rest) mp = go rest $ Map.insert off val mp
        go inss@(IAdd val off : rest) mp = case Map.lookup off mp of
            Nothing -> (inss, mp)
            Just origval -> go rest $ Map.insert off (origval + val) mp
        go inss@(ICopy from to mult : rest) mp =
            if defaultZero
                then let fromval = maybe 0 id $ Map.lookup from mp
                         toval = maybe 0 id $ Map.lookup to mp
                     in go rest $ Map.insert to (toval + mult * fromval) mp
                else case (Map.lookup from mp, Map.lookup to mp) of
                         (Nothing, _) -> (inss, mp)
                         (_, Nothing) -> (inss, mp)
                         (Just fromval, Just toval) ->
                             go rest $ Map.insert to (toval + mult * fromval) mp
        go inss mp = (inss, mp)

        compareSetAdd :: Instruction -> Instruction -> Ordering
        compareSetAdd (ISet _ _) (IAdd _ _) = LT
        compareSetAdd (IAdd _ _) (ISet _ _) = GT
        compareSetAdd i1 i2 = compare (offsetOf i1) (offsetOf i2)

        (after, valuemap) = go allinss Map.empty
        sets = sortBy compareSetAdd
                   $ Map.foldlWithKey
                       (\l off val -> if not defaultZero || val /= 0 then ISet val off : l else l)
                       [] valuemap
    in if isIStart ins then IStart : sets ++ propagateKnowns after else sets ++ propagateKnowns after
propagateKnowns (ILoop inss off : rest) = ILoop (propagateKnowns inss) off : propagateKnowns rest
propagateKnowns (ins : rest) = ins : propagateKnowns rest

firstSets :: Optimisation
firstSets [] = []
firstSets (IStart : rest) =
    let pre = takeWhile (\ins -> isIAdd ins || isISet ins) rest
        post = drop (length pre) rest
        dests = nub $ map offsetOf pre
        collected = [ISet (accumSetAdd 0 (filter ((== d) . offsetOf) pre)) d | d <- dests]
    in IStart : collected ++ post
firstSets inss = inss

propagateSlides :: Optimisation
propagateSlides [] = []
propagateSlides (ISlide off : rest) = propagateSlides (incOffsets off rest) ++ [ISlide off]
propagateSlides (ILoop inss off : rest) = ILoop (propagateSlides inss) off : propagateSlides rest
propagateSlides (ins : rest) = ins : propagateSlides rest

uselessEnd :: Optimisation
uselessEnd [] = []
uselessEnd inss = reverse $ dropWhile isUseless $ reverse inss
  where
    isUseless :: Instruction -> Bool
    isUseless (IInput _) = False
    isUseless (IOutput _) = False
    isUseless (ILoop lp _) = all isUseless lp
    isUseless _ = True

specialLoops :: Optimisation
specialLoops [] = []
specialLoops (ILoop [IAdd v off2] off1 : rest)
    | off1 /= off2 = ILoop [] off1 : specialLoops rest
    | odd v = ISet 0 off1 : specialLoops rest
specialLoops (ILoop inss off : rest)
    | all isIAdd inss,
      sum (map (\(IAdd v _) -> v) $ filter ((== off) . offsetOf) inss) == -1 =
        let others = map (\(IAdd v o) -> (v, o)) $ filter ((/= off) . offsetOf) inss
            dests = nub $ map snd others
            copies = [ICopy off d (sum $ map fst $ filter ((== d) . snd) others) | d <- dests]
        in copies ++ ISet 0 off : specialLoops rest
specialLoops (ILoop inss off : rest) = ILoop (specialLoops inss) off : specialLoops rest
specialLoops (ins : rest) = ins : specialLoops rest

normaliseOrder :: Optimisation
normaliseOrder [] = []
normaliseOrder (ICopy from1 to1 m1 : ICopy from2 to2 m2 : rest)
    | m2 < m1, to1 /= from2, to2 /= from1 = ICopy from2 to2 m2 : normaliseOrder (ICopy from1 to1 m1 : rest)
normaliseOrder (ILoop inss off : rest) = ILoop (normaliseOrder inss) off : normaliseOrder rest
normaliseOrder (ins : rest) = ins : normaliseOrder rest


repeated :: Optimisation -> Optimisation
repeated opt = \inss -> let inss' = opt inss
                        in if inss == inss' then inss else repeated opt inss'

incOffsets :: Offset -> [Instruction] -> [Instruction]
incOffsets _ [] = []
incOffsets inc (IAdd v off : rest) = IAdd v (off + inc) : incOffsets inc rest
incOffsets inc (ISet v off : rest) = ISet v (off + inc) : incOffsets inc rest
incOffsets inc (ICopy from to mult : rest) = ICopy (from + inc) (to + inc) mult : incOffsets inc rest
incOffsets inc (ISlide off : rest) = ISlide off : incOffsets inc rest
incOffsets inc (ILoop inss off : rest) = ILoop (incOffsets inc inss) (off + inc) : incOffsets inc rest
incOffsets inc (IInput off : rest) = IInput (off + inc) : incOffsets inc rest
incOffsets inc (IOutput off : rest) = IOutput (off + inc) : incOffsets inc rest
incOffsets inc (IStart : rest) = IStart : incOffsets inc rest

accumSetAdd :: Byte -> [Instruction] -> Byte
accumSetAdd acc [] = acc
accumSetAdd _ (ISet v' _ : rest) = accumSetAdd v' rest
accumSetAdd acc (IAdd v' _ : rest) = accumSetAdd (acc + v') rest
accumSetAdd _ _ = unreachable

unreachable :: a
unreachable = error "Unreachable"