{-# OPTIONS_GHC -fglasgow-exts -fallow-undecidable-instances #-}

module QIO.QioClass where

import Data.Maybe as Maybe
import Data.Monoid as Monoid
import Control.Monad.State
import QIO.QioSyn
import QIO.Heap
import Complex

newtype UnitaryC =  U {unU :: Int -> HeapMap -> HeapMap}

instance Monoid UnitaryC where
    mempty = U (\ fv bs -> bs)
    mappend (U f) (U g) = U (\ fv h -> g fv (f fv h))

uRotC :: Qbit -> Rotation -> UnitaryC
uRotC x f | f==rnot    = U (\ _ h -> update h x (not (fromJust (h ? x))))
          | f==rid     = mempty
          | otherwise  = error "not classical"

uSwapC :: Qbit -> Qbit -> UnitaryC 
uSwapC x y = U (\ _ h -> hswap h x y )

uCondC :: Qbit -> (Bool -> UnitaryC) -> UnitaryC
uCondC x br = U (\ fv h -> update (unU (br (fromJust (h ? x))) fv (forget h x)) x (fromJust (h ? x)))

uLetC :: Bool -> (Qbit -> UnitaryC) -> UnitaryC
uLetC b ux = U (\ fv h -> unU (ux (Qbit fv)) (fv+1) (update h (Qbit fv) b))

runUC :: U -> UnitaryC
runUC UReturn = mempty
runUC (Rot x r u) = uRotC x r `mappend` runUC u
runUC (Swap x y u) = uSwapC x y `mappend` runUC u
runUC (Cond x us u) = uCondC x (runUC.us) `mappend` runUC u
runUC (Ulet b xu u) = uLetC b (runUC.xu) `mappend` runUC u

data StateC = StateC {fv :: Int, heap :: HeapMap}

initialStateC :: StateC 
initialStateC = StateC 0 initial

runQStateC :: QIO a -> State StateC a
runQStateC (QReturn a) = return a
runQStateC (MkQbit b xq) = do (StateC fv h) <- get 
                              put (StateC (fv+1) (update h (Qbit fv) b)) 
                              runQStateC (xq (Qbit fv))
runQStateC (ApplyU u q) = do (StateC fv h) <- get
                             put (StateC fv (unU (runUC u) fv h))
                             runQStateC q 
runQStateC (Meas x qs) = do (StateC _ h) <- get
                            runQStateC (qs (fromJust (h ? x)))

runC :: QIO a -> a
runC q = evalState (runQStateC q) initialStateC