module FST where
import DFA
import Data.List (tails)

type FST st = ([st],   -- states
               [Char], -- input alphabet
               [Char], -- output alphabet
               st->Char->(st,[Char]), -- transition function
               st)     -- start state

-- | deltaStarFST m: extended transition function of FST m.
deltaStarFST :: (st -> Char -> (st,[Char])) -> (st -> [Char] -> (st,[Char]))
deltaStarFST delta q [] = (q,[])
deltaStarFST delta q (x:xs) = (q'',u'++u'') where
  (q' ,u' ) = delta q x
  (q'',u'') = deltaStarFST delta q' xs


-- | evalFST m w: run FST m on input w, and compute the output string
evalFST :: FST st -> [Char] -> [Char]
evalFST (qs, sigma, gamma, delta, s) w = u 
  where (q,u) = deltaStarFST delta s w

-- | composeFST m2 m1: returns an FST m such that m(w) = m2(m1(w))
composeFST :: FST st2 -> FST st1 -> FST (st1,st2)
composeFST (qs2,sigma2,gamma2,delta2,s2) 
           (qs1,sigma1,gamma1,delta1,s1) | sigma2==gamma1 =
  ([(q1,q2) | q1 <- qs1, q2 <- qs2],sigma1,gamma2,delta,(s1,s2)) 
  where delta (q1,q2) a = 
          let (q1',u) = delta1 q1 a
              (q2',v) = deltaStarFST delta2 q2 u
          in ((q1',q2'),v)

-- | Run FST on input string, and feed the output into DFA
composeFSTDFA :: FST st1 -> DFA st2 -> DFA (st1,st2)
composeFSTDFA (qsT,sigma,gamma,deltaT,sT) 
              (qs1,gamma1,delta1,s1,inF1) | gamma==gamma1 = 
  (qs,sigma,delta,s,inF) 
  where qs = [(qT,q1) | qT <- qsT, q1 <- qs1]
        s = (sT,s1)
        inF (qT,q1) = inF1 q1
        delta (qT,q1) a = 
          let (qT',w) = deltaT qT a
              delta1Star q [] = q
              delta1Star q (a:w) = delta1Star (delta1 q a) w
          in (qT', delta1Star q1 w)

-- NonDeterministic Finite State Transducers

type NFST st = ([st],     -- states
                [Char],   -- input alphabet
                [Char],   -- output alphabet
                st->Maybe Char->[(st,Maybe Char)], -- transition function
                st,       -- start state
                st->Bool) -- accepting states

composeNFST :: NFST st1 -> NFST st2 -> NFST (st1,st2)
composeNFST (qs1, sigma1, gamma1, delta1, s1, inF1) 
            (qs2, sigma2, gamma2, delta2, s2, inF2) | sigma2==gamma1 =
  let qs = [(q1,q2) | q1 <- qs1, q2 <- qs2]
      s = (s1,s2)
      delta (q1,q2) a = do
        (r1,w) <- delta1 q1 a
        (r2,u) <- delta2 q2 w
        return ((r1,r2),u)
      inF (q1,q2) = inF1 q1 && inF2 q2
  in (qs, sigma1, gamma2, delta, s, inF)

fst2nfst :: FST st -> NFST (st, String)
fst2nfst (qs, sigma, gamma, delta, s) =
  let qs1 = [(fst(delta q a), w) | q<-qs, a<-sigma, w<-tails(snd(delta q a)) ]
      delta1 (q,"") (Just a) = [(delta q a, Nothing)]
      delta1 (q,x:xs) Nothing = [((q,xs), Just x)]
      delta1 _ _ = []
      s1 = (s, "")
      inF1 (q,w) = (w == "")
  in (qs1, sigma, gamma, delta1, s1, inF1)