-- | A simple "main" function for gse-new.

module Main where

import Quipper

import Quipper.Libraries.Decompose

import Scheduler
import TrotterStep

import Control.Monad
import Data.List
import Quipper.Utils.CommandLine
import System.Console.GetOpt
import System.Environment
import System.Exit

-- | There are several different things that the program can output.
data WhatToShow =
  Circuit
  | Schedule
  | Singletons
  | Pairs
  | Triples
  | Quads
  deriving Show

-- | A data structure for the command line options.
data Options = Options {
  what :: WhatToShow,
  format :: Format,
  gatebase :: GateBase,
  precision_bits :: Integer,
  orbitals :: Integer,
  parallel :: Integer,
  timestep :: Timestep,
  revert :: Bool
  }

-- | The default options.
defaultOptions :: Options
defaultOptions = Options {
  what = Circuit,
  format = Preview,
  gatebase = Logical,
  precision_bits = 3,
  orbitals = 7,
  parallel = 1,
  timestep = 1.0,
  revert = False
  }

-- | A helper function for showing the default of some option.
showDefault :: (Show a) => (Options -> a) -> String
showDefault func = show (func defaultOptions)

-- | The list of command line options, in the format required by 'getOpt'.
options :: [OptDescr (Options -> IO Options)]
options = [ 
  Option ['h'] ["help"]      (NoArg  help)                       $ "print usage info and exit",
  Option ['C'] ["circuit"]   (NoArg (what Circuit))              $ "output the whole circuit (default)",
  Option ['S'] ["schedule"]  (NoArg (what Schedule))             $ "output the schedule",
  Option ['T'] ["template"]  (ReqArg read_template "<template>") $ "output a particular circuit template",
  Option ['f'] ["format"]    (ReqArg read_format "<format>")     $ "output format for circuits (default: " ++ showDefault format ++ ")",
  Option ['g'] ["gatebase"]  (ReqArg read_gatebase "<gatebase>") $ "gates to decompose into (default: " ++ showDefault gatebase ++ ")",
  Option ['m'] ["orbitals"]  (ReqArg read_m "<N>")               $ "number of orbitals (default: " ++ showDefault orbitals ++ ")",
  Option ['p'] ["parallel"]  (ReqArg read_p "<N>")               $ "with -T: number of parallel templates (default: " ++ showDefault parallel ++ ")",
  Option ['b'] ["precision"] (ReqArg read_b "<N>")               $ "number of precision qubits (default: " ++ showDefault precision_bits ++ ")",
  Option ['t'] ["timestep"]  (ReqArg read_t "<D>")               $ "timestep (default: " ++ showDefault timestep ++ ")",
  Option ['r'] ["revert"]    (NoArg set_revert)                  $ "reverse the schedule"
  ]
    where
      what :: WhatToShow -> Options -> IO Options
      what w o = return o { what = w }

      set_revert :: Options -> IO Options
      set_revert o = return o { revert = True }
      
      read_format :: String -> Options -> IO Options
      read_format str o = do
        case match_enum format_enum str of
          [(_, f)] -> return o { format = f }
          [] -> optfail ("Unknown format -- " ++ str ++ "\n")
          _  -> optfail ("Ambiguous format -- " ++ str ++ "\n")

      read_gatebase :: String -> Options -> IO Options
      read_gatebase str o = do
        case match_enum gatebase_enum str of
          [(_, f)] -> return o { gatebase = f }
          [] -> optfail ("Unknown gate base -- " ++ str ++ "\n")
          _  -> optfail ("Ambiguous gate base -- " ++ str ++ "\n")

      read_b :: String -> Options -> IO Options
      read_b string o =
        case parse_int string of 
          Just n | n > 0 -> return o { precision_bits = n }
          _ -> optfail ("Invalid b (precision) -- " ++ string ++ "\n")
          
      read_m :: String -> Options -> IO Options
      read_m string o =
        case parse_int string of 
          Just n | n > 0 -> return o { orbitals = n }
          _ -> optfail ("Invalid m (orbitals) -- " ++ string ++ "\n")

      read_p :: String -> Options -> IO Options
      read_p string o =
        case parse_int string of 
          Just n | n > 0 -> return o { parallel = n }
          _ -> optfail ("Invalid p (parallel) -- " ++ string ++ "\n")

      read_t :: String -> Options -> IO Options
      read_t string o =
        case parse_double string of 
          Just t -> return o { timestep = t }
          _ -> optfail ("Invalid t (timestep) -- " ++ string ++ "\n")

      read_template :: String -> Options -> IO Options
      read_template str o = do
        case match_enum template_enum str of
          [(_, f)] -> return o { what = f }
          [] -> optfail ("Unknown template -- " ++ str ++ "\n")
          _  -> optfail ("Ambiguous template -- " ++ str ++ "\n")
            
      help :: Options -> IO Options
      help o = do
        usage
        exitSuccess

-- | A list of possible arguments for the -T option.
template_enum :: [(String, WhatToShow)]
template_enum = [
  ("singletons", Singletons),
  ("pairs", Pairs),
  ("triples", Triples),
  ("quads", Quads)
  ]

-- | Process /argv/-style command line options into an 'Options' structure.
dopts :: [String] -> IO Options
dopts argv =
  case getOpt Permute options argv of
    (o, [], []) -> (foldM (flip id) defaultOptions o)
    (_, _, []) -> optfail "Too many non-option arguments\n"
    (_, _, errs) -> optfail (concat errs)

-- | Print usage message to standard output.
usage :: IO ()
usage = do
  putStr (usageInfo header options) 
  putStr (show_enum "format" format_enum)
  putStr (show_enum "gatebase" gatebase_enum)
  putStr (show_enum "template" template_enum)
    where header = "Usage: gse [OPTION...]"

-- ----------------------------------------------------------------------
-- * The main function

-- | Main function: read options, then execute the appropriate task.
main :: IO()
main = do
  -- Read options.
  argv <- getArgs
  options <- dopts argv
  
  case what options of
    Circuit -> main_circuit options
    Schedule -> main_schedule options
    Singletons -> main_singletons options
    Pairs -> main_pairs options
    Triples -> main_triples options
    Quads -> main_quads options
    
-- | Main function for outputting the GSE circuit.
main_circuit :: Options -> IO()
main_circuit options = do

  -- Read parameters.
  let b = precision_bits options
      m = orbitals options
      t = timestep options
      r = revert options
      
  -- Generate the main circuit.
  let circuit = trotter_step r t fake_hamiltonian

  let bs = genericReplicate b qubit
  let qs = genericReplicate m qubit

  -- Decompose it.
  let circuit' = decompose_generic (gatebase options) circuit
  
  -- Print it.
  print_generic (format options) circuit' bs qs

-- | Main function for outputting the schedule.
main_schedule :: Options -> IO ()
main_schedule options = do
  let m = orbitals options
  let schedule = scheduler m
  sequence_ [print x | x <- schedule]

-- | Main function for outputting a singleton template.
main_singletons :: Options -> IO ()
main_singletons options = do
  let m = parallel options
  let b = precision_bits options
  let t = timestep options
  let bs = genericReplicate b qubit
  let qs = zip [1.1..] (genericReplicate m qubit)
  let circuit = do_singletons t
  let circuit' = decompose_generic (gatebase options) circuit
  print_generic (format options) circuit' bs qs

-- | Main function for outputting a pair template.
main_pairs :: Options -> IO ()
main_pairs options = do
  let m = parallel options
  let b = precision_bits options
  let t = timestep options
  let bs = genericReplicate b qubit
  let qs = zip (zip [1.1..] [1.2..]) (genericReplicate m (qubit,qubit))
  let circuit = do_pairs t
  let circuit' = decompose_generic (gatebase options) circuit
  print_generic (format options) circuit' bs qs

-- | Main function for outputting a triple template.
main_triples :: Options -> IO ()
main_triples options = do
  let m = parallel options
  let triples = m `div` 3
  let rest = m `mod` 3
  let b = precision_bits options
  let t = timestep options
  let bs = genericReplicate b qubit
  let qs = zip (zip3 [1.1..] [1.2..] [1.3..]) (genericReplicate triples (qubit,qubit,qubit))
  let noops = genericReplicate rest qubit
  let circuit = do_triples_extended t
  let circuit' = decompose_generic (gatebase options) circuit
  print_generic (format options) circuit' bs qs noops

  where
    do_triples_extended :: Timestep -> [Qubit] -> [((Double, Double, Double), (Qubit, Qubit, Qubit))] -> [Qubit] -> Circ ()
    do_triples_extended t bs qs noops = do
      do_triples t bs qs

-- | Main function for outputting a quad template.
main_quads :: Options -> IO ()
main_quads options = do
  let m = parallel options
  let b = precision_bits options
  let t = timestep options
  let bs = genericReplicate b qubit
  let qs = zip (zip3 [1.1..] [1.2..] [1.3..]) (genericReplicate m (qubit,qubit,qubit,qubit))
  let circuit = do_quads t
  let circuit' = decompose_generic (gatebase options) circuit
  print_generic (format options) circuit' bs qs
