-- | This module contains code to generate the LaTeX for Figure 2 of
-- arXiv:2310.12256. 

module ShowSchedule where

import Data.List

import Scheduler

-- | A column of the circuit output.
data Column =
  GateColumn (Stage Integer)  -- ^ A column of gates.
  | SwapColumn [SwapOp]       -- ^ A column of fermionic swaps.
  | LabelColumn [Integer]     -- ^ A column of labels.
  deriving (Show, Eq)

-- | Given /m/ and a list of stages, calculate the "logical" list of
-- columns of the circuit.
expand_schedule :: Integer -> [Stage Integer] -> [Column]
expand_schedule m stages = all_columns
  where
    aux :: [Integer] -> Stage Integer -> ([Integer], [Column])
    aux perm stage = (perm', cols)
      where
        perm' = perm_of_stage stage
        swaps = parallel_swaps perm perm'
        cols = [SwapColumn x | x <- swaps] ++ [LabelColumn perm'] ++ [GateColumn stage]

    auxs :: [Integer] -> [Stage Integer] -> ([Integer], [Column])
    auxs perm [] = (perm, [])
    auxs perm (h:t) = (perm'', h' ++ t')
      where
        (perm', h') = aux perm h
        (perm'', t') = auxs perm' t

    perm0 = [0..m-1]
    (perm, columns) = auxs perm0 stages
    last_swaps = parallel_swaps perm perm0
    all_columns = columns ++ [SwapColumn x | x <- last_swaps] ++ [LabelColumn perm0]

-- | Given a list of columns, fix the /x/-coordinate of each column
-- and also return the total width.
find_x_coordinates :: [Column] -> ([(Double, Double, Double, Column)], Double)
find_x_coordinates cols = (out, x2)
  where
    aux :: Double -> Column -> ((Double, Double, Double, Column), Double)
    aux x0 g@(GateColumn _) = ((x0, x, x1, g), x1)
      where
        x1 = x0 + 1.5
        x = x0 + 0.75
    aux x0 g@(SwapColumn _) = ((x0, x, x1, g), x1)
      where
        x1 = x0 + 1
        x = x0 + 0.5
    aux x0 g@(LabelColumn _) = ((x0, x, x1, g), x1)
      where
        x1 = x0 + 1.5
        x = x0 + 0.75

    auxs :: Double -> [Column] -> ([(Double, Double, Double, Column)], Double)
    auxs x0 [] = ([], x0)
    auxs x0 (h:t) = (h':t' , x2)
      where
        (h', x1) = aux x0 h
        (t', x2) = auxs x1 t

    x0 = 1
    (out, x1) = auxs x0 cols
    x2 = x1 + 1

-- | Given a list and a predicate, find the leftmost element
-- satisfying the predicate. Split the list into 3 parts.
leftmost :: (a -> Bool) -> [a] -> Maybe ([a], a, [a])
leftmost p [] = Nothing
leftmost p (h:t)
  | p h = Just ([], h, t)
  | otherwise = case leftmost p t of
      Nothing -> Nothing
      Just (left, middle, right) -> Just (h:left, middle, right)    

-- | Given a list and a predicate, find the rightmost element
-- satisfying the predicate. Split the list into 3 parts.
rightmost :: (a -> Bool) -> [a] -> Maybe ([a], a, [a])
rightmost p as = case leftmost p (reverse as) of
  Nothing -> Nothing
  Just (left, middle, right) -> Just (reverse right, middle, reverse left)

-- | Given a width, and a sequence of xcols, output several sequences
-- of xcols, each of at most the given width. Be sure to only insert a
-- line break before a label stage, and duplicate the label stage. No
-- adjustment of x-coordinates yet.
divide_into_lines :: Double -> [(Double, Double, Double, Column)] -> [[(Double, Double, Double, Column)]]
divide_into_lines width [] = []
divide_into_lines width (xcol:xcols) = case rightmost p xcols of
  Nothing -> [xcols]
  Just (left, middle, right) -> first_line : other_lines
    where
      first_line = [xcol] ++ left ++ [middle]
      other_lines = divide_into_lines width ([middle] ++ right)
  where
    x0 = left_boundary xcols
    p (_,_,x1,g) = (x1 <= x0 + width) && (is_labelcol g)

    is_labelcol (LabelColumn _) = True
    is_labelcol _ = False

left_boundary :: [(Double, Double, Double, Column)] -> Double
left_boundary [] = error "left_boundary"
left_boundary ((x0,_,_,_):_) = x0

right_boundary :: [(Double, Double, Double, Column)] -> Double
right_boundary [] = error "right_boundary"
right_boundary xs = x1
  where
    (_,_,x1,_) = last xs

shift_xcol :: Double -> (Double, Double, Double, Column) -> (Double, Double, Double, Column)
shift_xcol dx (x0, x, x1, g) = (x0+dx, x+dx, x1+dx, g)

shift_xcols :: Double -> [(Double, Double, Double, Column)] -> [(Double, Double, Double, Column)]
shift_xcols dx xcols = [shift_xcol dx xcol | xcol <- xcols]

expand_xcols :: Double -> [(Double, Double, Double, Column)] -> [(Double, Double, Double, Column)]
expand_xcols width [] = []
expand_xcols width xcols = xcols'
  where
    x0 = left_boundary xcols
    x1 = right_boundary xcols
    x0' = 0
    x1' = width
    a = (x1' - x0') / (x1 - x0)
    b =  x0' - a * x0
    xcols' = map f xcols
    f (x0, x, x1, g) = (a*x0+b, a*x+b, a*x1+b, g)    

-- | Generate LaTeX from a seat.
latex_of_gate :: Double -> Integer -> Seat Integer -> ([String], Integer)
latex_of_gate x y (Noop p) = ([], y-1)
latex_of_gate x y (Singleton p) = (["\\gate{$S$}{" ++ show x ++ "," ++ show y ++ "};"], y-1)
latex_of_gate x y (Pair p q) = (["\\biggate{$P$}{" ++ show x ++ "," ++ show (y-1) ++ "}{" ++ show x ++ "," ++ show y ++ "};"], y-2)
latex_of_gate x y (Triple p q r) = (["\\biggate{$T$}{" ++ show x ++ "," ++ show (y-2) ++ "}{" ++ show x ++ "," ++ show y ++ "};"], y-3)
latex_of_gate x y (Quad p q r s) = (["\\biggate{$Q$}{" ++ show x ++ "," ++ show (y-3) ++ "}{" ++ show x ++ "," ++ show y ++ "};"], y-4)

-- | The plural of 'latex_of_gate'.
latex_of_gates :: Double -> Integer -> [Seat Integer] -> ([String], Integer)
latex_of_gates x y [] = ([], y)
latex_of_gates x y (h:t) = (h' ++ t', y'')
  where
    (h', y') = latex_of_gate x y h
    (t', y'') = latex_of_gates x y' t

-- | Generate LaTeX from a 'SwapOp'.
latex_of_swapop :: Double -> Integer -> SwapOp -> ([String], Integer)
latex_of_swapop x y Id = ([], y-1)
latex_of_swapop x y Swap = (["\\dotswapgate{" ++ show x ++ "}{" ++ show y ++ "}{" ++ show (y-1) ++ "};"], y-2)

latex_of_swapops :: Double -> Integer -> [SwapOp] -> ([String], Integer)
latex_of_swapops x y [] = ([], y)
latex_of_swapops x y (h:t) = (h' ++ t', y'')
  where
    (h', y') = latex_of_swapop x y h
    (t', y'') = latex_of_swapops x y' t

-- | Generate LaTeX from a single column.
latex_of_column :: Integer -> Double -> Column -> [String]
latex_of_column m x (GateColumn stage) = cols
  where
    (cols, _) = latex_of_gates x (m-1) stage
latex_of_column m x (SwapColumn swaps) = cols
  where
    (cols, _) = latex_of_swapops x (m-1) swaps
latex_of_column m x (LabelColumn labels) =
  ["\\wirelabel{\\footnotesize " ++ show i ++ "}{" ++ show x ++ "," ++ show y ++ "};" | (y,i) <- zip [m-1,m-2..] labels]
  
-- | Plural of 'latex_of_column'.
latex_of_columns :: Integer -> [(Double, Double, Double, Column)] -> [String]
latex_of_columns m xcols = concat [latex_of_column m x col | (_,x,_,col) <- xcols]

lines_of_xcols :: Integer -> [(Double, Double, Double, Column)] -> [String]
lines_of_xcols m [] = []
lines_of_xcols m xcols = lines
  where
    x1 = right_boundary xcols
    gatecode = latex_of_columns m xcols
    lines = ["\\["]
      ++ (indents 2 $
          ["\\scalebox{0.45}{"]
          ++ ["\\begin{qcircuit}[scale=0.4]"]
          ++ indents 2 ["\\grid{" ++ show x1 ++ "}{" ++ intercalate "," [show i | i <- [0..m-1]] ++ "}"]
          ++ indents 2 gatecode
          ++ ["\\end{qcircuit}"])
          ++ ["}"]
      ++ ["\\]"]

-- | Generate the LaTeX code for an entire schedule.
show_schedule :: Integer -> String
show_schedule m = unlines lines
  where
    stages = scheduler m
    cols = expand_schedule m stages
    (xcols, _) = find_x_coordinates cols
    lines = lines_of_xcols m xcols

-- | Generate the LaTeX code for an entire schedule, breaking the
-- diagram into parts of at most the given width.
show_schedule_lines :: Double -> Integer -> String
show_schedule_lines width m = unlines lines
  where
    stages = scheduler m
    cols = expand_schedule m stages
    (xcols, _) = find_x_coordinates cols
    xcolss = divide_into_lines width xcols
    lines = concat [process xcols | xcols <- xcolss]

    process xcols = lines_of_xcols m xcols'
      where
        xcols' = expand_xcols width xcols

-- | Indent a string by /n/ spaces.
indent :: Integer -> String -> String
indent n s = genericReplicate n ' ' ++ s

-- | Indent a paragraph by /n/ spaces.
indents :: Integer -> [String] -> [String]
indents n ss = map (indent n) ss

-- | Main
main_schedule :: Integer -> IO ()
main_schedule m = do
  putStr (show_schedule m)

-- | With line breaking.
main_schedule_lines :: Double -> Integer -> IO ()
main_schedule_lines width m = do
  putStr (show_schedule_lines width m)

main :: IO ()
main = main_schedule_lines 65 8
