{-# LANGUAGE        OverloadedLists #-}
{-# LANGUAGE        BlockArguments  #-}
{-# OPTIONS_HADDOCK not-home        #-}

{-|
Module      : QM
Description : The Quantum Monad
Stability   : experimental

Definition of the quantum monad (`QM`) and helper functions
-}
module Lib.QM (
    -- * The Quantum Monad
      QM
    , run
    , runDebug
    , eval
    , put
    , get
    , modify
    , io

    -- * Quantum State representation 
    , QState(..)

    -- * QBit
    , Ix
    , QBit(..)

    -- * Helpers
    , checkState
    , getState
    , stateSize
    ) where

import Numeric.LinearAlgebra
    ( size, toList, C, Vector, magnitude )
import Data.List ( intercalate )
import qualified Control.Monad.Random as Rand ( fromList, evalRandIO )

-- | Internal index type, to indicate it is a qubit index.
type Ix = Int

-- | Pointer to a qubit in QState.
-- Represents the linking function in QLambda
newtype QBit = Ptr { QBit -> Ix
link :: Ix }
    deriving Ix -> QBit -> ShowS
[QBit] -> ShowS
QBit -> String
(Ix -> QBit -> ShowS)
-> (QBit -> String) -> ([QBit] -> ShowS) -> Show QBit
forall a.
(Ix -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [QBit] -> ShowS
$cshowList :: [QBit] -> ShowS
show :: QBit -> String
$cshow :: QBit -> String
showsPrec :: Ix -> QBit -> ShowS
$cshowsPrec :: Ix -> QBit -> ShowS
Show

-- | The program state, a complex vector representation of
-- the qubits in the system
newtype QState = QState { QState -> Vector C
state :: Vector C }

instance Show QState where
    show :: QState -> String
show (QState q :: Vector C
q) = "== QState: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ " ==\n"
                   String -> ShowS
forall a. [a] -> [a] -> [a]
++ String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate "\n" ((C -> String) -> [C] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map C -> String
forall a. Show a => a -> String
show ([C] -> [String]) -> [C] -> [String]
forall a b. (a -> b) -> a -> b
$ Vector C -> [C]
forall a. Storable a => Vector a -> [a]
toList Vector C
q) String -> ShowS
forall a. [a] -> [a] -> [a]
++ "\n"
        where s :: String
s = Ix -> String
forall a. Show a => a -> String
show (Ix -> String) -> Ix -> String
forall a b. (a -> b) -> a -> b
$ Vector C -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Vector C
q

instance Eq QState where
    == :: QState -> QState -> Bool
(==) (QState q1 :: Vector C
q1) (QState q2 :: Vector C
q2) = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (C -> C -> Bool) -> [C] -> [C] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith C -> C -> Bool
(~=) (Vector C -> [C]
forall a. Storable a => Vector a -> [a]
toList Vector C
q1) (Vector C -> [C]
forall a. Storable a => Vector a -> [a]
toList Vector C
q2)

-- | Compareas two complex numbers for equality to the 6th decimal
(~=) :: C -> C -> Bool
~= :: C -> C -> Bool
(~=) a :: C
a b :: C
b = Double
bm Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
eqMargin Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
am Bool -> Bool -> Bool
&& Double
am Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
bm Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eqMargin
  where am :: Double
am = C -> Double
forall a. RealFloat a => Complex a -> a
magnitude C
a
        bm :: Double
bm = C -> Double
forall a. RealFloat a => Complex a -> a
magnitude C
b
        eqMargin :: Double
eqMargin = 0.000001

-- | The Quantum Monad
-- Keeps a state of the complex vector representation while allowing
-- pseudo-random number generation
newtype QM a = QM { QM a -> QState -> IO (a, QState)
runQM :: QState -> IO (a, QState) }

instance Show (QM a) where
    show :: QM a -> String
show _q :: QM a
_q = "Please use the function 'run' to perform the simulation"

instance Functor QM where
    fmap :: (a -> b) -> QM a -> QM b
fmap f :: a -> b
f m :: QM a
m = QM a
m QM a -> (a -> QM b) -> QM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= b -> QM b
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> QM b) -> (a -> b) -> a -> QM b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f

instance Applicative QM where
    pure :: a -> QM a
pure a :: a
a = (QState -> IO (a, QState)) -> QM a
forall a. (QState -> IO (a, QState)) -> QM a
QM \s :: QState
s -> (a, QState) -> IO (a, QState)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a,QState
s)
    {-# INLINE pure #-}
    QM af :: QState -> IO (a -> b, QState)
af <*> :: QM (a -> b) -> QM a -> QM b
<*> QM ax :: QState -> IO (a, QState)
ax = (QState -> IO (b, QState)) -> QM b
forall a. (QState -> IO (a, QState)) -> QM a
QM \s :: QState
s -> do
        (f :: a -> b
f, s' :: QState
s') <- QState -> IO (a -> b, QState)
af QState
s
        (x :: a
x, s'' :: QState
s'') <- QState -> IO (a, QState)
ax QState
s'
        (b, QState) -> IO (b, QState)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
x, QState
s'')

instance Monad QM where
    return :: a -> QM a
return = a -> QM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    m :: QM a
m >>= :: QM a -> (a -> QM b) -> QM b
>>= k :: a -> QM b
k = (QState -> IO (b, QState)) -> QM b
forall a. (QState -> IO (a, QState)) -> QM a
QM \s :: QState
s -> do
        (a :: a
a, s' :: QState
s') <- QM a -> QState -> IO (a, QState)
forall a. QM a -> QState -> IO (a, QState)
runQM QM a
m QState
s
        QM b -> QState -> IO (b, QState)
forall a. QM a -> QState -> IO (a, QState)
runQM (a -> QM b
k a
a) QState
s'

-- | Perform IO action inside the quantum monad
io :: IO a -> QM a
io :: IO a -> QM a
io m :: IO a
m = (QState -> IO (a, QState)) -> QM a
forall a. (QState -> IO (a, QState)) -> QM a
QM \s :: QState
s -> do
    a
a <- IO a
m
    (a, QState) -> IO (a, QState)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a,QState
s)
{-# INLINE io #-}

-- | Replace the quantum state
put :: QState -> QM ()
put :: QState -> QM ()
put s :: QState
s = (QState -> IO ((), QState)) -> QM ()
forall a. (QState -> IO (a, QState)) -> QM a
QM \_ -> ((), QState) -> IO ((), QState)
forall (m :: * -> *) a. Monad m => a -> m a
return ((), QState
s)
{-# INLINE put #-}

-- | Fetch the quantum state
get :: QM QState
get :: QM QState
get = (QState -> IO (QState, QState)) -> QM QState
forall a. (QState -> IO (a, QState)) -> QM a
QM \s :: QState
s -> (QState, QState) -> IO (QState, QState)
forall (m :: * -> *) a. Monad m => a -> m a
return (QState
s,QState
s)
{-# INLINE get #-}

-- | Apply a function to the quantum state
modify :: (QState -> QState) -> QM ()
modify :: (QState -> QState) -> QM ()
modify f :: QState -> QState
f = (QState -> IO ((), QState)) -> QM ()
forall a. (QState -> IO (a, QState)) -> QM a
QM \s :: QState
s -> ((), QState) -> IO ((), QState)
forall (m :: * -> *) a. Monad m => a -> m a
return ((), QState -> QState
f QState
s)
{-# INLINE modify #-}

-- | Run quantum program
eval :: QM a -> IO (a, QState)
eval :: QM a -> IO (a, QState)
eval qm :: QM a
qm = QM a -> QState -> IO (a, QState)
forall a. QM a -> QState -> IO (a, QState)
runQM QM a
qm (Vector C -> QState
QState [])
{-# INLINE eval #-}

-- | Run quantum program, discarding it's final state
run :: QM a -> IO a
run :: QM a -> IO a
run qm :: QM a
qm = (a, QState) -> a
forall a b. (a, b) -> a
fst ((a, QState) -> a) -> IO (a, QState) -> IO a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QM a -> IO (a, QState)
forall a. QM a -> IO (a, QState)
eval QM a
qm
{-# INLINE run #-}

-- | Run quantum program, prints the final quantum state and returns the result
runDebug :: QM a -> IO a
runDebug :: QM a -> IO a
runDebug qm :: QM a
qm = do
    (a :: a
a, s :: QState
s) <- QM a -> IO (a, QState)
forall a. QM a -> IO (a, QState)
eval QM a
qm
    QState -> IO ()
forall a. Show a => a -> IO ()
print QState
s
    a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

-- | Given a QState, returns how many qbits it consists of.
stateSize :: QState -> Ix
stateSize :: QState -> Ix
stateSize q :: QState
q = case Vector C -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size (QState -> Vector C
state QState
q) of
                0 -> 0
                x :: IndexOf Vector
x -> Ix -> Ix
forall a. Integral a => a -> Ix
log2 Ix
IndexOf Vector
x

-- | Normal integer log with base 2.
log2 :: Integral a => a -> Ix
log2 :: a -> Ix
log2 = Double -> Ix
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Ix) -> (a -> Double) -> a -> Ix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase 2 (Double -> Double) -> (a -> Double) -> a -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- | Returns the quantum state together with it's size
getState :: QM (QState, Int)
getState :: QM (QState, Ix)
getState = do
    QState
s <- QM QState
get
    let size :: Ix
size = QState -> Ix
stateSize QState
s
    (QState, Ix) -> QM (QState, Ix)
forall (m :: * -> *) a. Monad m => a -> m a
return (QState
s, Ix
size)

-- | Print the quantum state during operation
checkState :: QM ()
checkState :: QM ()
checkState = do
    QState
state <- QM QState
get
    IO () -> QM ()
forall a. IO a -> QM a
io (IO () -> QM ()) -> IO () -> QM ()
forall a b. (a -> b) -> a -> b
$ QState -> IO ()
forall a. Show a => a -> IO ()
print QState
state