{-# OPTIONS_GHC -fglasgow-exts #-}

module QIO.QioSyn where

import Data.Monoid as Monoid
import Complex


-- complex numbers
type CC = Complex RR

amp :: CC -> RR
amp k = (magnitude k)*(magnitude k)


-- real numbers
type RR = Float


-- Qubits are references
newtype Qbit = Qbit Int deriving (Num, Enum, Eq, Ord)


-- QIO and U as traces
type Rotation = ((Bool,Bool) -> CC)

data U = UReturn | Rot Qbit Rotation U
       | Swap Qbit Qbit U | Cond Qbit (Bool -> U) U | Ulet Bool (Qbit -> U) U

data QIO a = QReturn a | MkQbit Bool (Qbit -> QIO a) | ApplyU U (QIO a) 
           | Meas Qbit (Bool -> QIO a)


-- U functions
instance Monoid U where
    mempty = UReturn
    mappend UReturn u = u
    mappend (Rot x a u) u' = Rot x a (mappend u u')
    mappend (Swap x y u) u' = Swap x y (mappend u u')
    mappend (Cond x br u') u'' = Cond x br (mappend u' u'')
    mappend (Ulet b f u) u' = Ulet b f (mappend u u')

rot :: Qbit -> Rotation -> U
rot x r = Rot x r UReturn

swap :: Qbit -> Qbit -> U
swap x y = Swap x y UReturn

cond :: Qbit -> (Bool -> U) -> U
cond x br = Cond x br UReturn

ulet :: Bool -> (Qbit -> U) -> U
ulet b ux = Ulet b ux UReturn

urev :: U -> U
urev UReturn = UReturn
urev (Rot x r u) = urev u `mappend` rot x (rrev r)
urev (Swap x y u) = urev u `mappend` swap x y
urev (Cond x br u) = urev u `mappend` cond x (urev.br)
urev (Ulet b xu u) = urev u `mappend` ulet b (urev.xu)

unot :: Qbit -> U
unot x = rot x rnot

uhad :: Qbit -> U
uhad x = rot x rhad

uphase :: Qbit -> RR -> U
uphase x r = rot x (rphase r) 


--- QIO functions
instance Monad QIO where
    return = QReturn
    (QReturn a) >>= f = f a
    (MkQbit b g) >>= f = MkQbit b (\ x -> g x >>= f)
    (ApplyU u q) >>= f = ApplyU u (q >>= f)
    (Meas x g) >>= f = Meas x (\ b -> g b >>= f)

mkQbit :: Bool -> QIO Qbit
mkQbit b = MkQbit b return

applyU :: U -> QIO ()
applyU u =  ApplyU u (return ())

measQbit :: Qbit -> QIO Bool
measQbit x = Meas x return


-- rotations
rid :: Rotation
rid (x,y) = if x==y then 1 else 0

rnot :: Rotation
rnot (x,y) = if x==y then 0 else 1

rhad :: Rotation
rhad (x,y) = if x && y then -h else h where h = (1/sqrt 2)

rphase :: RR -> Rotation
rphase _ (False,False)  = 1
rphase r (True,True)    = exp(0:+r)
rphase _ (_,_)          = 0

rrev :: Rotation -> Rotation
rrev r (False,True)   = conjugate (r (True,False)) 
rrev r (True,False)   = conjugate (r (False,True))
rrev r xy             = conjugate (r xy)

instance Eq Rotation where
    f == g =    (f (False,False)  == g (False,False)) 
             && (f (False,True)   == g (False,True)) 
             && (f (True,False)   == g (True,False)) 
             && (f (True,True)    == g (True,True))
    f /= g =    (f (False,False)  /= g (False,False)) 
             || (f (False,True)   /= g (False,True)) 
             || (f (True,False)   /= g (True,False)) 
             || (f (True,True)    /= g (True,True))


-- show functions (for Qbit, Rotation and U)
instance Show Qbit where
	show (Qbit q) = "(Qbit:" ++ show q ++ ")"

instance Show Rotation where
    show f = "(" ++ (show (f (False,False))) ++ "," ++ (show (f (False,True))) ++ "," ++ (show (f (True,False))) ++ "," ++ (show (f (True,True))) ++ ")"

instance Show U where
	show u = show' u 0 (-1)

show' :: U -> Int -> Int -> String

show' (UReturn) x fv = ""

show' (Rot q a u) x fv = spaces x ++ "Rotate " ++ show q ++ " by " ++ show a ++ ".\n" ++ show' u x fv

show' (Swap q1 q2 u) x fv = spaces x ++ "Swap " ++ show q1 ++ " and " ++ show q2 ++ ".\n" ++ show' u x fv

show' (Cond q f u) x fv = spaces x ++ "Cond (if " ++ show q ++ " then \n" ++ spaces (x+1) ++ "(\n" ++ show' (f True) (x+1) fv ++ spaces (x+1) ++ ")\n" ++ spaces x ++ "else \n" ++ spaces (x+1) ++ "(\n" ++ show' (f False) (x+1) fv ++ spaces (x+1) ++ ")\n" ++ show' u x fv

show' (Ulet b f u) x fv = spaces x ++ "Ulet " ++ show b ++ " (\\" ++ show (Qbit fv) ++ "->\n " ++ show' (f (Qbit fv)) x (fv-1) ++ ")\n" ++ show' u x fv

spaces :: Int -> String
spaces 0 = ""
spaces (n+1) = "    " ++ spaces n