{-# LANGUAGE UnicodeSyntax, ScopedTypeVariables, TypeFamilies #-}
module Spanning where

import Prelude.Unicode
import Lambda (Λ (..), Symbol (..), Params (..), V, combinations)
import Language.HaLex.Dfa (Dfa (..))
import Data.Map (Map)
import qualified Data.Map as Map
import Data.List ((\\),sort)
import Control.Monad.State
import Control.Applicative


-- | Turns a λ-DFA into a λ-spanning-tree. The I-symbol is used for indirection nodes
spanning ∷ ∀ state. Ord state => Params → Dfa state Symbol -> Λ
spanning params (Dfa symbols states start _ trans) = evalState (spanningΛ start) funcVarsVisited where

	-- Stateful computation that remembers which shared states have already been dispatched
	spanningΛ ∷ state → State (Map state (V,Bool)) Λ
	spanningΛ s = let descend = spanningLoc in do
		sharedStates ← get
		case Map.lookup s sharedStates of -- we check the status of s
			Nothing → descend s           -- s is not a shared state
			Just (f,True) → pure $ V f        -- s is a shared state that has already been dispatched
			Just (f,False) → do               -- s is a shared state that we will dispatch right now
				modify $ Map.insert s (f,True) -- we mark s as dispatched
				I f <$> descend s          -- we create an indirection node and descend below s

	spanningLoc ∷ state → State (Map state (V,Bool)) Λ
	spanningLoc s = let descend = spanningΛ in case outEdges s of
	   [(S_Λ,e)] → Λ (lookupMap s varMap) <$> descend e
	   [(S_A0,f), (S_A1,x)] → A <$> descend f <*> descend x
	   [(S_V,x)] → pure $ V (lookupMap x varMap)
	   [(S_S0,e),(S_S1,abs)] → S (lookupMap abs varMap) <$> descend e
	   [(S_F v, _)] → pure $ V v
	   [] → pure $ I "bh" (V "bh") -- blackhole
	   _ → error "This seems not to be a λ-DFA"

	varMap ∷ Map state V
	varMap = Map.fromList $ filter isAbstraction states `zip` vs where
		vs ∷ [V]
		vs = combinations "xyzabcderstuvw" \\ freeVariables
		isAbstraction ∷ state → Bool
		isAbstraction node = trans node S_Λ ≢ dummy

	dummy ∷ state -- the state whose outgoing edges all point to itself
	dummy = case [s | s ← states, all (\sym → trans s sym ≡ s) symbols] of
		[x] → x
		___ → error $ "couldn't identify dummy state"

	freeVariables ∷ [V]
	freeVariables = [v | S_F v ← symbols]

	fs ∷ [V]
	fs = combinations ['F'..'U'] \\ freeVariables

	-- same as funcVars only with the additional information wether we have already visited a node
	funcVarsVisited ∷ Map state (V,Bool)
	funcVarsVisited = fmap (\v → (v,False)) funcVars

	funcVars ∷ Map state V -- maps shared states to a unique function variable
	funcVars = Map.fromList $ {-filter (not . isVarOcc)-} sharedStates `zip` fs where
		isVarOcc s = any (\l → trans s l ≢ dummy) (filter isVarSymbol symbols) where
			isVarSymbol (S_V) = True
			isVarSymbol (S_F _) = True
			isVarSymbol _ = False

	-- states with multiple incoming (non-backlink) edges (shared subgraphs)
	sharedStates ∷ [state]
	sharedStates = Map.keys $ Map.filter (≥2) $ Map.fromListWith (+) (succList `zip` repeat 1) where
		succList = start : [target | (sym,target) ← concatMap outEdges states, not $ isBacklink sym]
		isBacklink (S_V)   = True
		isBacklink (S_F _) = True
		isBacklink (S_S1)  = True
		isBacklink _______ = False

	outEdges ∷ state → [(Symbol, state)]
	outEdges source = sort [(sym, target) | sym ← symbols, let target = trans source sym, target ≢ dummy]

lookupMap ∷ Ord k ⇒ k → Map k a → a
lookupMap = Map.findWithDefault (error "lookupMap: key not found")