-- | This module computes the circuit for a single Trotter step of the
-- parallel GSE circuit described in arXiv:2310.12256.

module TrotterStep where

import Quipper

import Data.List
import System.Random

import Scheduler

-- ----------------------------------------------------------------------
-- * Gates

-- | Apply a /T/(θ) gate. This is the gate
--
-- > \[1, 0]
-- > \[0, [exp -iθ]].
-- 
-- It is equivalent to [exp iθ/2Z] = 'expZt' (-iθ/2) up to a
-- phase. Note that /T/(-π/4) is the usual /T/-gate.
gse_T_at :: Double -> Qubit -> Circ ()
gse_T_at theta q = named_rotation_at "T(%)" theta q

-- | Apply a fermionic swap gate to qubits /a/ and /b/.
-- Specially, this implements the unitary
--
-- * |00⟩ ↦ |00⟩
-- * |01⟩ ↦ |10⟩
-- * |10⟩ ↦ |01⟩
-- * |11⟩ ↦ −|11⟩
--
-- Note that we do not explicitly reorder the qubits in the output.
-- So for example,
--
-- > (a',b') <- fermionic_swap a b
-- 
-- produces this circuit:
--
-- > a ───×──•──── a'
-- >      |  |
-- > b ───×──•──── b'
--
-- and not this one:
-- 
-- > a ───×──•──── b'
-- >      |  |
-- > b ───×──•──── a'

fermionic_swap :: Qubit -> Qubit -> Circ (Qubit, Qubit)
fermionic_swap q r = do
  swap_at q r
  gate_Z_at r `controlled` q
  return (q, r)

-- | Create a new ancilla and initialize it from a list of qubits with
-- positive (\"True\") and/or negative (̈\"False\") controls.
initialize_ancilla :: Integer -> [Signed Qubit] -> Circ (Integer, Qubit)
initialize_ancilla i ctrls = do
  q <- qinit False
  qnot_at q `controlled` ctrls
  return (i, q)

-- ----------------------------------------------------------------------
-- * Hamiltonian data

-- | An index for an orbital. These are usually denotes p,q,r,s.
type Orbital = Integer

-- $ One problem with Hamiltonian integral data is that it is not
-- always clear how the data must be interpreted. There seem to be two
-- competing ways of specifying the data. They are equivalent but not
-- the same.
--
-- Recall that the unitary transformation that must be implemented is
-- exp[-/itH/], where the Hamiltonian /H/ can be specified in two
-- ways.
--
-- 1. It can be specified as a sum
--
-- H = ∑_pq h_pq a†_p a_q + ∑_pqrs h_pqrs a†_p a†_q a_r a_s, (1)
--
-- where (p,q) in the first sum and (p,q,r,s) in the second sum run
-- over /all/ pairs, respectively quadruples, of orbital indices.
--
-- 2. It can be specified as a sum
--
-- H = ∑_pq(can) k_pq (a†_p a_q + a†_q a_p) + ∑_pqrs(can) k_pqrs (a†_p a†_q a_r a_s + a†_s a†_r a_q a_p), (2)
--
-- where (p,q) in the first sum and (p,q,r,s) in the second sum run
-- over /canonical representatives/ of pairs, respectively quadruples,
-- of orbital indices.
--
-- The two representations are equivalent but not identical. First,
-- note that by the commuting relations of Fermionic annihilation and
-- creation operators, we have
--
-- a†_p a†_q a_r a_s
-- = -a†_p a†_q a_s a_r
-- = -a†_q a†_p a_r a_s
-- = a†_q a†_p a_s a_r.
--
-- Thus, in the first sum, certain operators are enumerated more than
-- once. This induces equivalences (p,q,r,s) ∼ (p,q,s,r) ∼ (q,p,r,s) ∼
-- (q,p,s,r) on quadruples. In a canonical representative (p,q,r,s)
-- for a quadruple, we therefore require that p ≤ q and s ≤ r.
--
-- Moreover, if p = q or s = r, then a†_p a†_q a_r a_s = 0,
-- so in a canonical representative, we require p < q and s < r.
-- 
-- Second, because the Hamiltonian is hermitian, we can assume h_pq
-- = h_qp and h_pqrs = h_srqp. (Here we are assuming all
-- integral data to be real numbers, which is apparently permissible
-- assuming a suitable choice of basis and the absence of an external
-- electromagnetic field. Otherwise, h_pq and h_qp would be
-- conjugates.) Thus, we get additional equivalences (p,q) ∼ (q,p) and
-- (p,q,r,s) ∼ (s,r,q,p).
--
-- Therefore, in a canonical representative (p,q) for a pair, we may
-- assume that p ≤ q, and in a canonical representative for a
-- quadruple (p,q,r,s), we may assume that (p,q) ≤ₗₑₓ (s,r), where
-- ≤ₗₑₓ denotes the lexicographic order.
--
-- Then the relationship between the two representations is as follows.
-- Write (p,q) ∼ (p',q') when the two are equivalent, and similarly
-- (p,q,r,s) ∼ (p',q',r',s').
--
-- Write b_pq = a†_p a_q + a†_q a_p
-- and b_pqrs = a†_p a†_q a_r a_s + a†_s a†_r a_q a_p
--
-- For convenience, also write a_pq = a†_p a_q, and
-- a_pqrs = a†_p a†_q a_r a_s
--
-- and for (p,q) ∼ (p',q'), say the sign of the equivalence is positive if
-- b_pq = b_p'q', and negative if b_pq = -b_p'q', and similarly for pqrs.
-- 
-- Let
--
-- k_pq = ½ ∑_p'q'(equiv) ±h_p'q', where the sum ranges over all (p',q') such
-- that (p',q') ∼ (p,q) and the sign is that of the equivalence.
-- 
-- k_pqrs = ½ ∑_p'q'r's'(equiv) ±h_p'q'r's', where the sum ranges over all
-- (p',q',r',s') ∼ (p,q,r,s), similarly.
--
-- Then
--
-- ∑_pq(can) k_pq (a†_p a_q + a†_q a_p)
-- = ∑_pq(can) k_pq b_pq
-- = ∑_pq(can) (½ ∑_p'q'(equiv) ±h_p'q') b_pq
-- = ½ ∑_pq(can) ∑_p'q'(equiv) ±h_p'q' b_pq
-- = ½ ∑_pq(can) ∑_p'q'(equiv) h_p'q' b_p'q'
-- = ½ ∑_p'q' ∑_pq(can)(equiv) h_p'q' b_p'q'
-- = ½ ∑_p'q' h_p'q' b_p'q'     (because one canonical rep per class)
-- = ½ ∑_pq h_pq b_pq           (reindexing)
-- = ½ ∑_pq h_pq (a†_p a_q + a†_q a_p)
-- = ½ ∑_pq h_pq a†_p a_q + ½ ∑_pq h_pq a†_q a_p
-- = ½ ∑_pq h_pq a†_p a_q + ½ ∑_pq h_qp a†_q a_p  (since h hermitian)
-- = ½ ∑_pq h_pq a†_p a_q + ½ ∑_pq h_pq a†_p a_q  (reindexing)
-- = ∑_pq h_pq a†_p a_q
--
-- and similarly
--
-- ∑_pqrs(can) k_pqrs (a†_p a†_q a_r a_s + a†_s a†_r a_q a_p)
-- = ∑_pqrs(can) k_pqrs b_pqrs
-- = ∑_pqrs(can) (½ ∑_p'q'r's'(equiv) ±h_p'q'r's') b_pqrs
-- = ½ ∑_pqrs(can) ∑_p'q'r's'(equiv) ±h_p'q'r's' b_pqrs
-- = ½ ∑_pqrs(can) ∑_p'q'r's'(equiv) h_p'q'r's' b_p'q'r's'
-- = ½ ∑_p'q'r's' ∑_pqrs(can)(equiv) h_p'q'r's' b_p'q'r's'
-- = ½ ∑_p'q'r's' h_p'q'r's' b_p'q'r's'
-- = ½ ∑_pqrs h_pqrs b_pqrs
-- = ½ ∑_pqrs h_pqrs (a_pqrs + a_srqp)
-- = ½ ∑_pqrs h_pqrs a_pqrs + ½ ∑_pqrs h_pqrs a_srqp
-- = ½ ∑_pqrs h_pqrs a_pqrs + ½ ∑_pqrs h_srqp a_srqp
-- = ½ ∑_pqrs h_pqrs a_pqrs + ½ ∑_pqrs h_pqrs a_pqrs
-- = ∑_pqrs h_pqrs a_pqrs
-- = ∑_pqrs h_pqrs a†_p a†_q a_r a_s.
--
-- It follows that with this choice of k, (1) and (2) are equivalent,
-- i.e., define the same H.
--
-- Note that it has some idiosyncracies. For example, when p=q, we have
--
-- k_pp = ½ ∑_p'p'(equiv) ±h_p'p' = ½ h_pp,
--
-- because the equivalence class of (p,p) is just (p,p). But when p≠q, we have
--
-- k_pq = ½ ∑_p'q'(equiv) ±h_p'q' = ½ h_pq + ½ h_qp = h_pq.
--
-- These additional factors or 1 or ½ are a main source of confusion
-- in interpreting the Hamiltonian data. In the 4-index (2-electron)
-- cases, there are even more cases to consider.
-- 
-- For the converse direction, given k_pq and k_pqrs, we can define
-- suitable h_pq and h_pqrs. For these to be unique, we need an
-- additional assumption. We might as well normalize the data so that:
--
-- h_pq = ± h_p'q' when (p,q) ∼ (p',q') and ± is the sign of the equivalence.
-- And similarly for h_pqrs.
--
-- In other words, h_pq = ±h_p'q' whenever b_pq = ± b_p'q'.
--
-- In particular, since b_ppqr = 0 = -b_ppqr, we have h_ppqr = -h_ppqr,
-- therefore h_ppqr = 0. So the "impossible" terms are 0.
--
-- With this assumption, the h values are uniquely determined by the k
-- values and vice versa.
--
-- Indeed, from k_pq = ½ ∑_p'q'(equiv) ±h_p'q'
-- we get
--
-- k_pq = ½ ∑_p'q'(equiv) ±h_p'q'
-- = ½ ∑_p'q'(equiv) h_pq
-- = (½ ∑_p'q'(equiv)) h_pq,
--
-- where the term (½ ∑_p'q'(equiv)) only depends on the size of the
-- equivalence class. Write #(p,q) for the size of that equivalence class.
-- Then
--
-- k_pq = ½#(p,q) h_pq.
--
-- The inverse formula, in the special case where (p,q) is a canonical
-- representative of an equivalence class, is:
--
-- h_pq = 2/#(p,q) k_pq.
--
-- In general, we must use h_pq = ± 2/#(p,q) k_pq',
-- where pq' is the canonical representative of pq, and ± is the sign
-- of the equivalence.
--
-- This makes sense because the formula (1) "spreads out" the weights
-- over the members of each equivalence class, so needs proportionally
-- smaller coefficients, whereas the formula (2) "concentrates" the
-- weights one one representative for each class, so needs larger
-- coefficients.  Moreover, the additional factor of ½ or 2 comes from
-- whether the individual terms have been symmetrized as in (2), or
-- not, as in (1).

-- | A pair (p,q) of orbits.
newtype OrbitalPair = OrbitalPair (Orbital, Orbital)
  deriving (Eq, Show, Ord)

-- | A quad (p,q,r,s) of orbits.
newtype OrbitalQuad = OrbitalQuad (Orbital, Orbital, Orbital, Orbital)
  deriving (Eq, Show, Ord)

-- | Equivalence. Returns 1 or -1 if equivalent, 0 if vanishing, or 'Nothing' if not equivalent.
equiv2 :: OrbitalPair -> OrbitalPair -> Maybe Integer
equiv2 (OrbitalPair (p,q)) (OrbitalPair (p',q'))
  | (p,q) == (p',q') = Just 1
  | (p,q) == (q',p') = Just 1
  | otherwise = Nothing

-- | Equivalence. Returns 1 or -1 if equivalent, 0 if equivalent and vanishing, or 'Nothing' if not equivalent.
equiv4 :: OrbitalQuad -> OrbitalQuad -> Maybe Integer
equiv4 (OrbitalQuad (p,q,r,s)) (OrbitalQuad (p',q',r',s'))
  | positive && negative = Just 0
  | positive = Just 1
  | negative = Just (-1)
  | otherwise = Nothing
  where
    positive = (p',q',r',s') `elem` [(p,q,r,s),(q,p,s,r),(s,r,q,p),(r,s,p,q)]
    negative = (p',q',r',s') `elem` [(p,q,s,r),(q,p,r,s),(s,r,p,q),(r,s,q,p)]

-- | Check if a quad is vanishing.
is_vanishing :: OrbitalQuad -> Bool
is_vanishing (OrbitalQuad (p,q,r,s))
  | p == q || r == s = True
  | otherwise = False

-- | Compute the orbit of a pair. It is sorted lexicographically. For each element of the orbit, also return the sign.
orbit2 :: OrbitalPair -> [(OrbitalPair, Integer)]
orbit2 pq@(OrbitalPair (p,q)) = map fix (nub (sort orbit))
  where
    orbit = [(p,q),(q,p)]
    fix (p',q') = (pq', sign)
      where
        pq' = OrbitalPair (p',q')
        sign = case equiv2 pq pq' of
          Just sign -> sign
          Nothing -> error "orbit2"
    
-- | Compute the orbit of a quad. It is sorted lexicographically.
orbit4 :: OrbitalQuad -> [(OrbitalQuad, Integer)]
orbit4 pqrs@(OrbitalQuad (p,q,r,s)) = map fix (nub (sort orbit))
  where
    orbit = [(p,q,r,s),(q,p,s,r),(s,r,q,p),(r,s,p,q),(p,q,s,r),(q,p,r,s),(s,r,p,q),(r,s,q,p)]
    fix (p',q',s',r') = (pqrs', sign)
      where
        pqrs' = OrbitalQuad (p',q',r',s')
        sign = case equiv4 pqrs pqrs' of
          Just sign -> sign
          Nothing -> error "orbit4"

-- | Return the size of an equivalence class.
equiv2_size :: OrbitalPair -> Integer
equiv2_size pq = fromIntegral (length (orbit2 pq))

-- | Return the size of an equivalence class.
equiv4_size :: OrbitalQuad -> Integer
equiv4_size pqrs = fromIntegral (length (orbit4 pqrs))

-- | Return the canonical representative of a pair, along with the
-- sign of the equivalence. The sign is ±1, or 0 if vanishing.
canonical2 :: OrbitalPair -> (OrbitalPair, Integer)
canonical2 pq = head (orbit2 pq)

-- | Return the canonical representative of a quad, along with the
-- sign of the equivalence. The sign is ±1, or 0 if vanishing.
canonical4 :: OrbitalQuad -> (OrbitalQuad, Integer)
canonical4 pqrs = head (orbit4 pqrs)

-- | Hamiltonian data is an oracle that one can query to get the
-- integral data. The functions that have \"spread\" in their name use
-- the spread-out data model, i.e., they return the coefficients h_pq
-- and h_pqrs of formula (1) above. The functions that have \"can\" in
-- their name use the canonical data model, i.e., they return the
-- coefficients k_pq and k_pqrs of formula (2) above.
--
-- Note: in data model (1), we do not necessarily assume that the
-- equations h_pq = ±h_p'q' whenever b_pq = ± b_p'q' are satisfied. So
-- we always sum over equivalent pairs. And in data model (2), we also
-- define the function at non-canonical representatives, but we do not
-- assume that it is defined there.
data HamiltonianData = HamiltonianData {
  hamil1_spread :: OrbitalPair -> Double,
  hamil2_spread :: OrbitalQuad -> Double,
  hamil1_can :: OrbitalPair -> Double,
  hamil2_can :: OrbitalQuad -> Double
  }

-- | Conversion between the two data models.
hamil1_can_of_spread :: (OrbitalPair -> Double) -> (OrbitalPair -> Double)
hamil1_can_of_spread f = f'
  where
    f' pq = 1/2 * sum [fromIntegral sign * f pq' | (pq',sign) <- orbit2 pq]

-- | Conversion between the two data models.
hamil2_can_of_spread :: (OrbitalQuad -> Double) -> (OrbitalQuad -> Double)
hamil2_can_of_spread f = f'
  where
    f' pqrs = 1/2 * sum [fromIntegral sign * f pqrs' | (pqrs',sign) <- orbit4 pqrs]

-- | Conversion between the two data models.
hamil1_spread_of_can :: (OrbitalPair -> Double) -> (OrbitalPair -> Double)
hamil1_spread_of_can f = f'
  where
    f' pq = fromIntegral sign * 2 / fromIntegral (equiv2_size pq) * f pq'
      where
        (pq',sign) = canonical2 pq

-- | Conversion between the two data models.
hamil2_spread_of_can :: (OrbitalQuad -> Double) -> (OrbitalQuad -> Double)
hamil2_spread_of_can f = f'
  where
    f' pqrs = fromIntegral sign * 2 / fromIntegral (equiv4_size pqrs) * f pqrs'
      where
        (pqrs',sign) = canonical4 pqrs

-- To do: provide functions to read the Hamiltonian data from a file.

-- | Some fake data that we can use.
fake_hamiltonian :: HamiltonianData
fake_hamiltonian = HamiltonianData {
  hamil1_spread = h1,
  hamil2_spread = h2,
  hamil1_can = k1,
  hamil2_can = k2
  }
  where
    kk1 (OrbitalPair (p,q)) = random (p+123*q)
    kk2 (OrbitalQuad (p,q,r,s)) = random (p+13*q+17*r+11*s)
    random n = r
      where
        gen = mkStdGen (fromIntegral n)
        (r, _) = randomR (0,1) gen
    h1 = hamil1_spread_of_can kk1
    h2 = hamil2_spread_of_can kk2
    k1 = hamil1_can_of_spread h1
    k2 = hamil2_can_of_spread h2
        
-- | Convenience functions for extracting Hamiltonian values in the
-- canonical data model.
get_hamil1 :: HamiltonianData -> Orbital -> Orbital -> Double
get_hamil1 hamil p q = hamil1_can hamil (OrbitalPair (p,q))

get_hamil2 :: HamiltonianData -> Orbital -> Orbital -> Orbital -> Orbital -> Double
get_hamil2 hamil p q r s = hamil2_can hamil (OrbitalQuad (p,q,r,s))

-- ----------------------------------------------------------------------
-- * The circuit for a Trotter step.
  
-- | Perform a single controlled Trotter timestep of the Hamiltonian
-- simulation.
--
-- The parameters are:
--
-- • /revert/: whether to perform the stages in the opposite order.
-- 
-- • /t/: the size of the elementary timestep
--
-- • /h/: the Hamiltonian data
--
-- • /bs/: a list of /b/ qubits, encoding an integer /k/ in binary, in
--   big-headian bit order, i.e., the head of the list contains the
--   most significant digit.
--
-- • /qs/: a list of /m/ qubits [q₀, ..., qₘ₋₁], corresponding to
--   orbitals.
--
-- It generates a circuit that computes the operator $e^{-itkH}$ on
-- the qubits /qs/, controlled by the integer /k/ encoded by the /bs/.
trotter_step :: Bool -> Timestep -> HamiltonianData -> [Qubit] -> [Qubit] -> Circ [Qubit]
trotter_step revert t hamil bs qs = do
  (perm, qs') <- aux schedule [0..m-1] qs
  qs'' <- do_permutation perm [0..m-1] qs'
  return qs''
  where
    schedule = (if revert then reverse else id) (scheduler m)
    b = fromIntegral (length bs)
    m = fromIntegral (length qs)

    -- | Given the \"current permutation\" of the qubits, apply the
    -- given schedule and then return the new (possibly permuted) list
    -- of qubits as well as the new permutation.
    --
    -- Note: the input, say [1,4,3,2] [q1,q4,q3,q2] may as well be
    -- zipped, i.e., [(1,q1),(4,q4),(3,q3),(2,q2)]. Each logical qubit
    -- always stays with the same index, i.e., the list of qubits is
    -- permuted in the same way as the list of integers.
    aux [] perm qs = return (perm, qs)
    aux (stage:stages) perm qs = do
      let perm' = perm_of_stage stage
      qs' <- do_permutation perm perm' qs
      qs'' <- do_stage t hamil stage bs qs'
      aux stages perm' qs''

-- | Generates a fermionic permutation circuit.
do_permutation :: [Orbital] -> [Orbital] -> [Qubit] -> Circ [Qubit]
do_permutation perm_in perm_out qs = apply_swap_stages swap_stages qs
  where
    swap_stages = parallel_swaps perm_in perm_out

    apply_swap_stages [] qs = return qs
    apply_swap_stages (stage:stages) qs = do
      qs' <- apply_swap_stage stage qs
      qs'' <- apply_swap_stages stages qs'
      return qs''

    apply_swap_stage [] [] = return []
    apply_swap_stage (Id:rest) (q:qs) = do
      qs' <- apply_swap_stage rest qs
      return (q:qs')
    apply_swap_stage (Swap:rest) (q:r:qs) = do
      (q',r') <- fermionic_swap q r
      qs' <- apply_swap_stage rest qs
      return (q':r':qs')

-- | Generates a parallel rotation circuit correponding to a
-- stage, assuming the qubits are already in the indicated order.
do_stage :: Timestep -> HamiltonianData -> Stage Orbital -> [Qubit] -> [Qubit] -> Circ [Qubit]
do_stage t hamil stage bs qs = do
  comment ("Stage " ++ show stage)
  label qs [show i | i <- perm]
  do_singletons t bs singletons
  do_pairs t bs pairs
  do_triples t bs triples
  do_quads t bs quads
  return qs
  where
    perm = perm_of_stage stage
    qmap = zip perm qs
    qubit_of_orbital p = case lookup p qmap of
      Just q -> q
      Nothing -> error "qubit_of_orbital"
    singletons = do
      Singleton p <- stage
      let o0 = qubit_of_orbital p
      let h00 = 2 * get_hamil1 hamil p p  -- Multiply by 2 for symmetry reasons
      return (h00,o0)
    pairs = do
      Pair p q <- stage
      let [o0,o1] = map qubit_of_orbital [p,q]
      let h01 = get_hamil1 hamil p q
      let h0110 = 2 * get_hamil2 hamil p q q p -- Multiply by 2 for symmetry reasons
      return ((h01,h0110),(o1,o0))
    triples = do
      Triple p q r <- stage
      let [o0,o1,o2] = map qubit_of_orbital [p,q,r]
      let h0120 = get_hamil2 hamil p q r p
      let h0121 = get_hamil2 hamil p q r q
      let h0221 = get_hamil2 hamil p r r q
      return ((h0120,h0121,h0221),(o2,o1,o0))
    quads = do
      Quad p q r s <- stage
      let [o0,o1,o2,o3] = map qubit_of_orbital [p,q,r,s]
      let h0132 = get_hamil2 hamil p q r s
      let h0231 = get_hamil2 hamil p s r q
      let h0321 = get_hamil2 hamil p r s q
      return ((h0132,h0231,h0321),(o3,o2,o1,o0))

-- ----------------------------------------------------------------------
-- * Singletons

-- | Apply a singleton circuit. /t/ is the timestep, /h00/ is the
-- Hamiltonian data, /bs/ are the precision bits in big-headian bit
-- order, e.g., [b2,b1,b0], and /o0/ is the orbital qubit.
do_singleton :: Timestep -> Double -> [Qubit] -> Qubit -> Circ ()
do_singleton t h00 bs o0 = do
  with_computed wc $ \as -> do
    let angles = [(2^i * t * h00, x) | (i,x) <- as]
    sequence_ [gse_T_at θ x | (θ,x) <- angles]
    return ()
  where
    ibs = zip [0..] (reverse bs)
    wc = do
     as <- sequence [initialize_ancilla i [Signed b True, Signed o0 True] | (i, b) <- ibs] 
     return as

-- | Plural of 'do_singleton'.
do_singletons :: Timestep -> [Qubit] -> [(Double, Qubit)] -> Circ ()
do_singletons t bs qs = do
  with_computed (wcs qs) $ \ass -> do
    sequence_ [do_angles h00 o0 as | ((h00,o0),as) <- zip qs ass]
  where
    ibs = zip [0..] (reverse bs)
    wc o0 = do
     as <- sequence [initialize_ancilla i [Signed b True, Signed o0 True] | (i, b) <- ibs] 
     return as
    wcs qs = do
      sequence [wc o0 | (_,o0) <- qs]
    do_angles h00 o0 as = do
      let angles = [(2^i * t * h00, x) | (i,x) <- as]
      sequence_ [gse_T_at θ x | (θ,x) <- angles]
      return ()

-- ----------------------------------------------------------------------
-- * Pairs

-- | Apply a pair circuit. /t/ is the timestep, /h01/ and /h0110/
-- is the Hamiltonian data, /bs/ are the precision bits in big-headian
-- bit order, e.g., [b2,b1,b0], and (/o1/, /o0/) are the orbital
-- qubits in that order.
do_pair :: Timestep -> Double -> Double -> [Qubit] -> (Qubit, Qubit) -> Circ ()
do_pair t h01 h0110 bs (o1,o0) = do
  with_basis_change bc1 $ do
    with_computed wc2 $ \(a1s, a2s) -> do
      with_basis_change (bc3 a1s a2s) $ do
        with_basis_change (bc4 a1s a2s) $ do
          let angles1 = [(-2^i/2*t*h0110, x) | (i,x) <- a1s]
          let angles2 = [(2^i/2*t*h01, x) | (i,x) <- a2s]
          let anglesum = sum [θ | (θ,x) <- angles2]
          expZt_at anglesum o0
          sequence_ [expZt_at θ x | (θ,x) <- angles1 ++ angles2]
          return ()
  where
    -- If bs = [b2, b1, b0], then ibs = [(0,b0), (1,b1), (2,b2)]
    ibs = zip [0..] (reverse bs)
    bc1 = do
      qnot_at o1 `controlled` o0
    wc2 = do
      a1s <- sequence [initialize_ancilla i [Signed b True, Signed o1 False, Signed o0 True] | (i, b) <- ibs]
      a2s <- sequence [initialize_ancilla i [Signed b True, Signed o1 True] | (i, b) <- ibs]
      return (a1s, a2s)
    bc3 a1s a2s = do
      hadamard_at o0
    bc4 a1s a2s = do
      sequence_ [qnot_at a `controlled` (Signed o0 False) | (i, a) <- a2s]

-- | Plural of 'do_pair'.
do_pairs :: Timestep -> [Qubit] -> [((Double, Double), (Qubit, Qubit))] -> Circ ()
do_pairs t bs qs = do
  with_basis_change (bc1s qs) $ do
    with_computed (wc2s qs) $ \as -> do
      with_basis_change (bc3s qs as) $ do
        with_basis_change (bc4s qs as) $ do
          sequence_ [do_angles h01 h0110 o0 o1 a1s a2s | (((h01,h0110),(o1,o0)),(a1s,a2s)) <- zip qs as]
  where
    -- If bs = [b2, b1, b0], then ibs = [(0,b0), (1,b1), (2,b2)]
    ibs = zip [0..] (reverse bs)
    bc1 o0 o1 = do
      qnot_at o1 `controlled` o0
    bc1s qs = do
      sequence_ [bc1 o0 o1 | (_, (o1, o0)) <- qs]
    wc2 o0 o1 = do
      a1s <- sequence [initialize_ancilla i [Signed b True, Signed o1 False, Signed o0 True] | (i, b) <- ibs]
      a2s <- sequence [initialize_ancilla i [Signed b True, Signed o1 True] | (i, b) <- ibs]
      return (a1s, a2s)
    wc2s qs = do
      sequence [wc2 o0 o1 | (_, (o1, o0)) <- qs]
    bc3 o0 o1 a1s a2s = do
      hadamard_at o0
    bc3s qs as = do
      sequence_ [bc3 o0 o1 a1s a2s | ((_, (o1, o0)), (a1s, a2s)) <- zip qs as]
    bc4 o0 o1 a1s a2s = do
      sequence_ [qnot_at a `controlled` (Signed (o0 :: Qubit) False) | (i, a) <- a2s]
    bc4s qs as = do
      sequence_ [bc4 o0 o1 a1s a2s | ((_, (o1, o0)), (a1s, a2s)) <- zip qs as]

    do_angles h01 h0110 o0 o1 a1s a2s = do
      let angles1 = [(-2^i/2*t*h0110, x) | (i,x) <- a1s]
      let angles2 = [(2^i/2*t*h01, x) | (i,x) <- a2s]
      let anglesum = sum [θ | (θ,x) <- angles2]
      expZt_at anglesum o0
      sequence_ [expZt_at θ x | (θ,x) <- angles1 ++ angles2]
      return ()

-- ----------------------------------------------------------------------
-- * Triples

-- | Apply a triple circuit. /t/ is the time step, /h0120/ is the
-- relevant Hamiltonian data, /bs/ are the precision bits in
-- big-headian bit order, e.g., [b2,b1,b0], and /os/ are the orbitals
-- 2, 1, 0 in that order.
do_0120 :: Timestep -> Double -> [Qubit] -> (Qubit, Qubit, Qubit) -> Circ ()
do_0120 t h0120 bs (o2,o1,o0) = do
  with_basis_change bc1 $ do
    with_computed wc2 $ \a1s -> do
      with_basis_change (bc3 a1s) $ do
        let angles = [(2^i / 2 * t * h0120, x) | (i,x) <- a1s]
        let anglesum = sum [θ | (θ,x) <- angles]
        expZt_at anglesum o0
        sequence_ [expZt_at θ x | (θ,x) <- angles]
        return ()
  where
    -- If bs = [b2, b1, b0], then ibs = [(0,b0), (1,b1), (2,b2)]
    ibs = zip [0..] (reverse bs)
    bc1 = do
      qnot_at o2 `controlled` o1
      hadamard_at o1
    wc2 = do
      a1s <- sequence [initialize_ancilla i [Signed b True, Signed o2 True, Signed o0 True] | (i, b) <- ibs]
      return a1s
    bc3 a1s = do
      sequence_ [qnot_at a `controlled` (Signed o1 False) | (i, a) <- a1s]

-- | Like 'do_0120', but for triples of the form /h0121/.
do_0121 :: Timestep -> Double -> [Qubit] -> (Qubit, Qubit, Qubit) -> Circ ()
do_0121 t h0121 bs (o2,o1,o0) = do_0120 t h0121 bs (o2,o0,o1)

-- | Like 'do_0221', but for triples of the form /h0221/.
do_0221 :: Timestep -> Double -> [Qubit] -> (Qubit, Qubit, Qubit) -> Circ ()
do_0221 t h0221 bs (o2,o1,o0) = do_0120 t h0221 bs (o1,o0,o2)
      
-- | Plural of 'do_0120'.
do_0120s :: Timestep -> [Qubit] -> [(Double, (Qubit, Qubit, Qubit))] -> Circ ()
do_0120s t bs qs = do
  with_basis_change (bc1s qs) $ do
    with_computed (wc2s qs) $ \as -> do
       with_basis_change (bc3s qs as) $ do
         sequence_ [do_angles h0120 o0 o1 o2 a1s | ((h0120,(o2,o1,o0)),a1s) <- zip qs as]
  where
    -- If bs = [b2, b1, b0], then ibs = [(0,b0), (1,b1), (2,b2)]
    ibs = zip [0..] (reverse bs)
    bc1 o0 o1 o2 = do
      qnot_at o2 `controlled` o1
      hadamard_at o1
    bc1s qs = do
      sequence_ [bc1 o0 o1 o2 | (_, (o2, o1, o0)) <- qs]
    wc2 o0 o1 o2 = do
      a1s <- sequence [initialize_ancilla i [Signed b True, Signed o2 True, Signed o0 True] | (i, b) <- ibs]
      return a1s
    wc2s qs = do
      sequence [wc2 o0 o1 o2 | (_, (o2, o1, o0)) <- qs]
    bc3 o0 o1 o2 a1s = do
      sequence_ [qnot_at a `controlled` (Signed (o1 :: Qubit) False) | (i, a) <- a1s]
    bc3s qs as = do
      sequence_ [bc3 o0 o1 o2 a1s | ((_, (o2, o1, o0)), a1s) <- zip qs as]
    do_angles h0120 o0 o1 o2 a1s = do
      let angles = [(2^i / 2 * t * h0120, x) | (i,x) <- a1s]
      let anglesum = sum [θ | (θ,x) <- angles]
      expZt_at anglesum o0
      sequence_ [expZt_at θ x | (θ,x) <- angles]
      return ()
        
-- | Plural of 'do_0121'.
do_0121s :: Timestep -> [Qubit] -> [(Double, (Qubit, Qubit, Qubit))] -> Circ ()
do_0121s t bs qs = do_0120s t bs qs'
  where
    qs' = [(h, (o2,o0,o1)) | (h, (o2,o1,o0)) <- qs]

-- | Plural of 'do_0221'.
do_0221s :: Timestep -> [Qubit] -> [(Double, (Qubit, Qubit, Qubit))] -> Circ ()
do_0221s t bs qs = do_0120s t bs qs'
  where
    qs' = [(h, (o1,o0,o2)) | (h, (o2,o1,o0)) <- qs]

-- | Apply a triple circuit to the given triples. 
do_triples :: Timestep -> [Qubit] -> [((Double, Double, Double), (Qubit, Qubit, Qubit))] -> Circ ()
do_triples t bs qs = do
  do_0120s t bs qs0120
  do_0121s t bs qs0121
  do_0221s t bs qs0221
    where
      qs0120 = [(h0120, (o2,o1,o0)) | ((h0120, h0121, h0221), (o2, o1, o0)) <- qs]
      qs0121 = [(h0121, (o2,o1,o0)) | ((h0120, h0121, h0221), (o2, o1, o0)) <- qs]
      qs0221 = [(h0221, (o2,o1,o0)) | ((h0120, h0121, h0221), (o2, o1, o0)) <- qs]

-- ----------------------------------------------------------------------
-- * Quads

-- | Apply a quad circuit. /t/ is the time step, /hpqrs/ is the
-- relevant Hamiltonian data, /bs/ are the precision bits in
-- big-headian bit order, e.g., [b2,b1,b0], and /os/ are the orbitals
-- 3, 2, 1, 0 in that order.
do_quad :: Timestep -> Double -> Double -> Double -> [Qubit] -> (Qubit, Qubit, Qubit, Qubit) -> Circ ()
do_quad t h0132 h0231 h0321 bs (o3,o2,o1,o0) = do
  with_basis_change bc1 $ do
    with_computed wc2 $ \(a1s, a2s, a3s) -> do
      with_basis_change (bc3 a1s a2s a3s) $ do
        let angles = [(2^i / 2 * t * h0132, x) | (i,x) <- a1s]
              ++ [(2^i / 2 * t * h0231, x) | (i,x) <- a2s]
              ++ [(2^i / 2 * t * h0321, x) | (i,x) <- a3s]
        let anglesum = sum [θ | (θ,x) <- angles]
        expZt_at anglesum o0
        sequence_ [expZt_at θ x | (θ,x) <- angles]
        return ()
  where
    -- If bs = [b2, b1, b0], then ibs = [(0,b0), (1,b1), (2,b2)]
    ibs = zip [0..] (reverse bs)
    bc1 = do
      qnot_at o1 `controlled` o0
      qnot_at o2 `controlled` o0
      qnot_at o3 `controlled` o0
      hadamard_at o0
    wc2 = do
      a1s <- sequence [initialize_ancilla i [Signed b True, Signed o3 True, Signed o2 True, Signed o1 False] | (i, b) <- ibs]
      a2s <- sequence [initialize_ancilla i [Signed b True, Signed o3 True, Signed o2 False, Signed o1 True] | (i, b) <- ibs]
      a3s <- sequence [initialize_ancilla i [Signed b True, Signed o3 False, Signed o2 True, Signed o1 True] | (i, b) <- ibs]
      return (a1s, a2s, a3s)
    bc3 a1s a2s a3s = do
      sequence_ [qnot_at a `controlled` (Signed o0 False) | (i, a) <- a1s ++ a2s ++ a3s]
      
-- | The plural of 'quad'. In principle, we could just do this by
-- calling 'quad' multiple times, but we want all the rotations to be
-- overtly parallel, and so we interleave various basis changes,
-- rotations, and uncomputations. This also makes for a prettier
-- circuit.
do_quads :: Timestep -> [Qubit] -> [((Double, Double, Double), (Qubit, Qubit, Qubit, Qubit))] -> Circ ()
do_quads t bs qs = do
  with_basis_change (bc1s qs) $ do
    with_computed (wc2s qs) $ \as -> do
       with_basis_change (bc3s qs as) $ do
         sequence_ [do_angles h0132 h0231 h0321 o0 o1 o2 o3 a1s a2s a3s | (((h0132,h0231,h0321),(o3,o2,o1,o0)),(a1s,a2s,a3s)) <- zip qs as]
  where
    -- If bs = [b2, b1, b0], then ibs = [(0,b0), (1,b1), (2,b2)]
    ibs = zip [0..] (reverse bs)
    bc1 o0 o1 o2 o3 = do
      qnot_at o1 `controlled` o0
      qnot_at o2 `controlled` o0
      qnot_at o3 `controlled` o0
      hadamard_at o0
    bc1s qs = do
      sequence_ [bc1 o0 o1 o2 o3 | (_, (o3, o2, o1, o0)) <- qs]
    wc2 o0 o1 o2 o3 = do
      a1s <- sequence [initialize_ancilla i [Signed b True, Signed o3 True, Signed o2 True, Signed o1 False] | (i, b) <- ibs]
      a2s <- sequence [initialize_ancilla i [Signed b True, Signed o3 True, Signed o2 False, Signed o1 True] | (i, b) <- ibs]
      a3s <- sequence [initialize_ancilla i [Signed b True, Signed o3 False, Signed o2 True, Signed o1 True] | (i, b) <- ibs]
      return (a1s, a2s, a3s)
    wc2s qs = do
      sequence [wc2 o0 o1 o2 o3 | (_, (o3, o2, o1, o0)) <- qs]
    bc3 :: Qubit -> Qubit -> Qubit -> Qubit -> [(Integer, Qubit)] -> [(Integer, Qubit)] -> [(Integer, Qubit)] -> Circ ()
    bc3 o0 o1 o2 o3 a1s a2s a3s = do
      sequence_ [qnot_at a `controlled` (Signed o0 False) | (i, a) <- a1s ++ a2s ++ a3s]
    bc3s qs as = do
      sequence_ [bc3 o0 o1 o2 o3 a1s a2s a3s | ((_, (o3, o2, o1, o0)), (a1s, a2s, a3s)) <- zip qs as]
    do_angles h0132 h0231 h0321 o0 o1 o2 o3 a1s a2s a3s = do
      let angles = [(2^i / 2 * t * h0132, x) | (i,x) <- a1s]
            ++ [(2^i / 2 * t * h0231, x) | (i,x) <- a2s]
            ++ [(2^i / 2 * t * h0321, x) | (i,x) <- a3s]
      let anglesum = sum [θ | (θ,x) <- angles]
      expZt_at anglesum o0
      sequence_ [expZt_at θ x | (θ,x) <- angles]
      return ()
        
-- ----------------------------------------------------------------------
-- * Testing

-- | For testing.
preview4 :: IO ()
preview4 = do
  print_generic Preview (do_quad 1 2 3 4) [qubit,qubit,qubit] (qubit,qubit,qubit,qubit)

-- | For testing.
previews4 :: IO ()
previews4 = do
  print_generic Preview (do_quads 1) [qubit,qubit,qubit] [((2,3,4),(qubit,qubit,qubit,qubit)) | n <- [0..1]]

-- | For testing.
preview1 :: IO ()
preview1 = do
  print_generic Preview (do_singleton 1 2) [qubit,qubit,qubit] qubit

-- | For testing.
previews1 :: IO ()
previews1 = do
  print_generic Preview (do_singletons 1) [qubit,qubit,qubit] [(1,qubit),(2,qubit)]

-- | For testing.
preview2 :: IO ()
preview2 = do
  print_generic Preview (do_pair 1 2.111 3.123) [qubit,qubit,qubit] (qubit,qubit)

-- | For testing.
previews2 :: IO ()
previews2 = do
  print_generic Preview (do_pairs 1) [qubit,qubit,qubit] [((2.111, 3.123),(qubit,qubit)) | n <- [0..1]]

-- | For testing.
preview_0120 :: IO ()
preview_0120 = do
  print_generic Preview (do_0120 1 2.111) [qubit,qubit,qubit] (qubit,qubit,qubit)

-- | For testing.
preview_0121 :: IO ()
preview_0121 = do
  print_generic Preview (do_0121 1 2.111) [qubit,qubit,qubit] (qubit,qubit,qubit)

-- | For testing.
preview_0221 :: IO ()
preview_0221 = do
  print_generic Preview (do_0221 1 2.111) [qubit,qubit,qubit] (qubit,qubit,qubit)

-- | For testing.
preview_0120s :: IO ()
preview_0120s = do
  print_generic Preview (do_0120s 1) [qubit,qubit,qubit] [(2.111,(qubit,qubit,qubit)) | n <- [0..1]]

-- | For testing.
preview_0121s :: IO ()
preview_0121s = do
  print_generic Preview (do_0121s 1) [qubit,qubit,qubit] [(2.111,(qubit,qubit,qubit)) | n <- [0..1]]

-- | For testing.
preview_0221s :: IO ()
preview_0221s = do
  print_generic Preview (do_0221s 1) [qubit,qubit,qubit] [(2.111,(qubit,qubit,qubit)) | n <- [0..1]]

-- | For testing.
previews3 :: IO ()
previews3 = do
  print_generic Preview (do_triples 1) [qubit,qubit,qubit] [((2.111,3.123,10),(qubit,qubit,qubit)) | n <- [0..1]]
