{-# LANGUAGE UnicodeSyntax, FlexibleContexts, ScopedTypeVariables #-}
module Rules where

import Prelude.Unicode
import Graph
import GraphRewriting.Rule
import GraphRewriting.Pattern
import GraphRewriting.Pattern.InteractionNet
import GraphRewriting.Graph.Read
import GraphRewriting.Graph.Write
import Data.List (transpose, elemIndex, delete)
import Control.Applicative
import Data.Monoid
import Control.Monad

import Data.Maybe (fromJust)


compileShare ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
compileShare = do
	Multiplexer {out = o, ins = is} ← node
	case is of
		[ ] → replace $ byNode Eraser {inp = o}
		[i] → rewire [[o,i]]
		ins → let (ins1, ins2) = splitAt (length ins `div` 2) ins in replace $ do
			(o1,o2) ← (,) <$> byEdge <*> byEdge
			byNode $ Duplicator {level = 0, inp = o, out1 = o1, out2 = o2}
			byNode $ Multiplexer {out = o1, ins = ins1}
			byNode $ Multiplexer {out = o2, ins = ins2}

withoutIdx ∷ [a] → Int → [a]
withoutIdx xs i = let (ys,zs) = splitAt i xs in ys ⧺ tail zs

insertIdx ∷ Int → a → [a] → [a]
insertIdx i x xs = let (l,r) = splitAt i xs in l ⧺ [x] ⧺ r

split ∷ Int → Int → [a] → [[a]]
split i n [] = replicate n []
split i n xs = let (x,xs') = splitAt i xs in x : split i n xs'

transpose' n [] = replicate n []
transpose' n xs = transpose xs

commute ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
commute = do
	n1 :-: n2 ← activePair
	require (n1 ≢ n2) -- TODO: replace by linear
	let ports1 = inspect n1 ∷ [Port]
	let ports2 = inspect n2 ∷ [Port]
	let (pp1,pp1idx) = head [(p,i) | (p,i) ← ports1 `zip` [0..], p ≡ pp n1]
	let (pp2,pp2idx) = head [(p,i) | (p,i) ← ports2 `zip` [0..], p ≡ pp n2]
	let aux1 = pp1 `delete` inspect n1
	let aux2 = pp2 `delete` inspect n2
	let es1 = length aux1
	let es2 = length aux2
	replace $ do
		edges ← replicateM (es1 * es2) byEdge
		let edges1 = split es1 es2 edges
		let edges2 = transpose' es1 edges1
		mconcat [byNode $ updateLevel n2 $ update (insertIdx pp1idx pp1 auxs) n1 | (pp1,auxs) ← zip aux2 edges1]
		mconcat [byNode $ updateLevel n1 $ update (insertIdx pp2idx pp2 auxs) n2 | (pp2,auxs) ← zip aux1 edges2]
	where updateLevel you me = case me of
		Duplicator {} → maybeLevelUp
		Delimiter  {} → maybeLevelUp
		_ → me
		where maybeLevelUp = case you of
			Delimiter  {} → if level you ≤ level me then me {level = level me + 1} else me
			Abstractor {} → me {level = level me + 1}
			_ → me

annihilate ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
annihilate = do
	n1 :-: n2 ← activePair
	require (n1 ≡ n2) -- TODO: ???
	let aux1 = pp n1 `delete` inspect n1
	let aux2 = pp n2 `delete` inspect n2
	rewire $ [[a1,a2] | (a1,a2) ← aux1 `zip` aux2]

annihilateDelimiters ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
annihilateDelimiters = do
	rewrite ← annihilate
	Delimiter {} ← liftReader . inspectNode =<< previous
	return rewrite

-- This rule doesn't trigger for constants with arguments
eliminateDelimiterConstant ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
eliminateDelimiterConstant = do
	c@Constant {args = as, name = n} :-: Delimiter {inp = iD} ← activePair
	require (inp c ≢ iD && as == [])
	replace $ byNode $ Constant {inp = iD, args = [], name = n}

eliminateDelimiterEraser ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
eliminateDelimiterEraser = do
	c@Eraser {} :-: Delimiter {inp = iD} ← activePair
	require (inp c ≢ iD)
	replace $ byNode $ Eraser {inp = iD}

eliminateDuplicator ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
eliminateDuplicator = do
	Eraser {inp = iE} ← node
	Duplicator {inp = iD, out1 = o1, out2 = o2} ← neighbour =<< previous
	require (iE ≡ o1 ∨ iE ≡ o2)
	if iE ≡ o1
		then rewire [[iD,o2]]
		else rewire [[iD,o1]]

eraser ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
eraser = do
	rewrite ← commute
	Eraser {} ← liftReader . inspectNode =<< previous
	return rewrite

duplicate ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
duplicate = do
	rewrite ← commute
	Duplicator {} ← liftReader . inspectNode =<< previous
	return rewrite

beta ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
beta = do
	Applicator {inp = ai, func = f, arg = a} :-: Abstractor {body = b, var = v} ← activePair
	replace $ do
		byNode $ Delimiter {level = 0, inp = ai, out = b}
		byNode $ Delimiter {level = 0, inp = a, out = v}

commuteDelimiter ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
commuteDelimiter = do
	rewrite ← commute
	Delimiter {} ← liftReader . inspectNode =<< previous
	return rewrite

applyConstant ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
applyConstant = do
	Applicator {inp = i, arg = a} :-: Constant {name = n, args = as} ← activePair
	replace $ byNode $ Constant {inp = i, name = n, args = as ++ [a]}

applyOperator ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
applyOperator = do
	Applicator {inp = i, arg = a} :-: Operator {ops = os, arity = ar, lmop = l, function = fn, name = n} ← activePair
	require (ar > length os)
	replace $ byNode $ Operator {inp = i, ops = os ++ [a], arity = ar, lmop = l, function = fn, name = n}

-- TODO: Require that the lmoPort is not on one of the unreduced ports yet
-- Do we only reduce operator args, if the operator has all args already?
reduceOperand ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
reduceOperand = do
	o@(Operator {ops = os, lmop = lmo}) ← node
	opid ← previous
	let ports = inspect o ∷ [Port]
	let lmoport = ports !! lmo
	-- only change the lmo port if it is on top or if it is attached to a Constant
	require (lmo == 0) <|> do {Constant {} ← nodeWith lmoport; return ()}
	port ← branch os -- get a pattern that matches each port in os
	-- we require that at least one node attached to the operator is not a constant
	requireFailure $ do
		Constant {} ← nodeWith port
		return ()
	-- we need to add one, since the input port is not part of os, but is part of the port numbering
	let unreducedport = 1 + fromJust (elemIndex port os)
	return $ updateNode opid (o {lmop = unreducedport})

execOperator ∷ forall n. (View [Port] n, View NodeLS n) ⇒ Rule n
execOperator = do
	Operator {inp = i, ops = os, arity = ar, function = fn, name = n} ← node
	opid ← previous
	require (length os == ar)
	-- check that all args are constants
	argss ← forM os $ \o → do
			c@Constant {} ← adverse o opid
			return c
	case fn (map name argss) of
		Nothing → mempty
		Just n' → replace $ byNode $ Constant {inp = i, args = [], name = n'}

caseNode ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
caseNode = do
	-- the order of constant and case here is important, otherwise strategies don't work
	Case {inp = i, alts = alts, names = names} :-: Constant {name = n, args = as} ← activePair
	let matchingport = alts !! (fromJust $ elemIndex n names)
	let nn = length alts
	let m = length as
	-- We generate m new applicator nodes with m+1 new edges connecting them
	if m > 0 then
		replace $ do
			es ← replicateM (m+1) byEdge
			byWire matchingport (es !! 0) -- We merge the first new edge with the matching port from the Case node
			byWire i (es !! m) -- We merge the last new edge with the input edge of the Case node
			mconcat [byNode $ Applicator {inp = es !! (i+1), func = es !! i, arg = as !! i} | i ← [0..m-1]]
			mconcat [byNode $ Eraser {inp = alts !! i} | i ← filter (/= fromJust (elemIndex n names)) [0..nn-1]]
			 else do
		replace $ do
			byWire i matchingport -- Attach the alternative directly to the input of the case node
			mconcat [byNode $ Eraser {inp = alts !! i} | i ← filter (/= fromJust (elemIndex n names)) [0..nn-1]]

-- | Not the readback semantics as defined in the paper. Just a non-semantics-preserving erasure of all
-- delimiters to make the graph more readable
readback ∷ (View [Port] n, View NodeLS n) ⇒ Rule n
readback = do
	Delimiter {inp = i, out = o} ← node
	rewire [[i,o]]