{-# LANGUAGE RecordWildCards #-}

module HEP.Kinematics.Variable.M2
    (
      M2Solution (..)

    , mkInput

    , m2XXSQP
    , m2CXSQP
    , m2XCSQP
    , m2CCSQP

    , m2XXAugLag
    , m2CXAugLag
    , m2XCAugLag
    , m2CCAugLag
    ) where

import HEP.Kinematics
import HEP.Kinematics.Vector.LorentzVector (setXYZT)

import Numeric.LinearAlgebra               (Vector, fromList, toList)
import Numeric.NLOPT

-- import Numeric.GSL.Differentiation
-- import Debug.Trace

-- | All the components are rescaled by '_scale'.
data InputKinematics = InputKinematics
                       { _p1     :: FourMomentum  -- ^ p1 = a1 + b1
                       , _p2     :: FourMomentum  -- ^ p2 = a2 + b2
                       , _q1     :: FourMomentum  -- ^ q1 = b1
                       , _q2     :: FourMomentum  -- ^ q2 = b2
                       , _ptmiss :: TransverseMomentum
                         -- | squared mass of the invisible particle
                       , _mInvSq :: !Double
                       , _scale  :: !Double }

-- | creates the input for the process of
--
--   A1 + A2 --> a1 B1 + a2 B2 --> a1 b1 C1 + a2 b2 C2.
--
mkInput :: [FourMomentum]      -- ^ [a1, a2]
        -> [FourMomentum]      -- ^ [b1, b2]
        -> TransverseMomentum  -- ^ pT(miss)
        -> Double              -- ^ M_{invisible}
        -> Maybe InputKinematics
mkInput as bs ptmiss mInv
    | length as /= 2 || length bs /= 2 = Nothing
    | otherwise                        = do
          let ps = zipWith (+) as bs
              p1 = head ps
              p2 = ps !! 1

              m1 = mass p1
              m2 = mass p2
              mInvSq = mInv * mInv
              scaleSq = m1 * m1 + m2 * m2 + 2 * mInvSq
          if scaleSq <= 0  -- why?
              then Nothing
              else do let scale = sqrt scaleSq
                          p1'     = p1        ^/ scale
                          p2'     = p2        ^/ scale
                          q1'     = head bs   ^/ scale
                          q2'     = (bs !! 1) ^/ scale
                          ptmiss' = ptmiss    ^/ scale
                          mInv'   = mInv       / scale
                      return $ InputKinematics { _p1     = p1'
                                               , _p2     = p2'
                                               , _q1     = q1'
                                               , _q2     = q2'
                                               , _ptmiss = ptmiss'
                                               , _mInvSq = mInv' * mInv'
                                               , _scale  = scale }

data M2Solution = M2Solution { _M2    :: Double
                             , _k1sol :: FourMomentum
                             , _k2sol :: FourMomentum
                             } deriving Show

getM2Solution :: InputKinematics -> Either Result Solution -> Maybe M2Solution
getM2Solution inp@InputKinematics {..} sol =
    case sol of
        Left _                   -> Nothing
        -- sol0@(Right (Solution m2 ks _)) -> do
        Right (Solution m2 ks _) -> do
            let Invisibles k1 k2 = mkInvisibles inp (vecToVars ks)
            -- traceM ("sol = " ++ show sol0)
            return $ M2Solution { _M2    = m2  * _scale
                                , _k1sol = k1 ^* _scale
                                , _k2sol = k2 ^* _scale }

{-
initialGuess :: InputKinematics -> Vector Double
initialGuess InputKinematics { _ptmiss = ptmiss } =
    fromList [0.5 * px ptmiss, 0.5 * py ptmiss, 0, 0]
-}
initialGuess :: InputKinematics -> Vector Double
initialGuess inp@InputKinematics {..} =
    let objf ks = let Invisibles k1 k2 = mkInvisibles inp (vecToVars ks)
                  in invariantMass [_p1, _p2, k1, k2]
        ks0 = fromList [0.5 * px _ptmiss, 0.5 * py _ptmiss, 0, 0]
        eps' = eps * 0.1
        stop = MaximumEvaluations 1000
               :| [ ObjectiveAbsoluteTolerance eps'
                  , ObjectiveRelativeTolerance eps']
        algorithm = NELDERMEAD objf [] Nothing
        problem = LocalProblem 4 stop algorithm
        sol = minimizeLocal problem ks0
    in case sol of
        Left _                     -> ks0
        Right (Solution _ ksSol _) -> ksSol

eps :: Double
eps = 1e-3

stopObjCond :: [StoppingCondition]
stopObjCond = [ObjectiveAbsoluteTolerance eps', ObjectiveRelativeTolerance eps']
  where eps' = eps * 1e-3

type MultivarFunc = Vector Double -> (Double, Vector Double)

m2SQP :: [InputKinematics -> MultivarFunc]  -- ^ constraint functions
      -> Maybe InputKinematics
      -> Maybe M2Solution
m2SQP _   Nothing                         = Nothing
m2SQP cfs (Just inp@InputKinematics {..}) =
    let -- objective function with gradient
        objfD = m2ObjF inp

        -- constraint function with gradient
        constraints' = ($ inp) <$> cfs
        constraints  = (\cf -> EqualityConstraint (Scalar cf) eps)
                       <$> constraints'

        stop = MaximumEvaluations 1000 :| stopObjCond
        algorithm = SLSQP objfD [] [] constraints
        problem = LocalProblem 4 stop algorithm

        sol = minimizeLocal problem (initialGuess inp)
    in getM2Solution inp sol

m2XXSQP, m2CXSQP, m2XCSQP, m2CCSQP :: Maybe InputKinematics -> Maybe M2Solution
m2XXSQP = m2SQP []
m2CXSQP = m2SQP [constraintA]
m2XCSQP = m2SQP [constraintB]
m2CCSQP = m2SQP [constraintA, constraintB]

m2AugLag :: [InputKinematics -> MultivarFunc]  -- ^ constraint functions
         -> Maybe InputKinematics
         -> Maybe M2Solution
m2AugLag _   Nothing                         = Nothing
m2AugLag cfs (Just inp@InputKinematics {..}) =
    let -- objective function without gradient
        -- objf  = m2ObjF inp
        objf  = getResultOnly (m2ObjF inp)

        -- constraint function without gradient
        constraints' = ($ inp) <$> cfs
        -- constraints  = (\cf -> EqualityConstraint (Scalar cf) eps)
        --                <$> constraints'
        constraints  = (\cf -> EqualityConstraint (Scalar (getResultOnly cf))
                               eps) <$> constraints'

        stop = MaximumEvaluations 5000 :| stopObjCond
        -- algorithm = VAR1 objf Nothing
        -- algorithm = LBFGS objf Nothing
        -- algorithm = SLSQP objf [] [] constraints
        algorithm = NELDERMEAD objf [] Nothing

        subproblem = LocalProblem 4 stop algorithm
        -- problem = AugLagProblem [] constraints (AUGLAG_EQ_LOCAL subproblem)
        problem = AugLagProblem constraints [] (AUGLAG_EQ_LOCAL subproblem)

        sol = minimizeAugLag problem (initialGuess inp)
    in getM2Solution inp sol
  where
    getResultOnly :: MultivarFunc -> Vector Double -> Double
    getResultOnly fD ks = let (result, _) = fD ks in result

m2XXAugLag, m2CXAugLag, m2XCAugLag, m2CCAugLag
    :: Maybe InputKinematics -> Maybe M2Solution
m2XXAugLag = m2AugLag []
m2CXAugLag = m2AugLag [constraintA]
m2XCAugLag = m2AugLag [constraintB]
m2CCAugLag = m2AugLag [constraintA, constraintB]

-- | (k1x, k1y, k1z, k2z).
type Variables = (Double, Double, Double, Double)

-- | unsafe transformation, but it would be fine.
vecToVars :: Vector Double -> Variables
vecToVars ks = let [k1x, k1y, k1z, k2z] = toList ks in (k1x, k1y, k1z, k2z)

data Invisibles = Invisibles { _k1 :: FourMomentum , _k2 :: FourMomentum }

mkInvisibles :: InputKinematics -> Variables -> Invisibles
mkInvisibles InputKinematics {..} (k1x, k1y, k1z, k2z) = Invisibles k1 k2
  where
    eInv1 = sqrt (k1x * k1x + k1y * k1y + k1z * k1z + _mInvSq)
    k1 = eInv1 `seq` setXYZT k1x k1y k1z eInv1

    k2x = px _ptmiss - k1x
    k2y = py _ptmiss - k1y
    eInv2 = sqrt (k2x * k2x + k2y * k2y + k2z * k2z + _mInvSq)
    k2 = eInv2 `seq` setXYZT k2x k2y k2z eInv2

-- | the unknowns are (k1x, k2x, k1z, k2z).
m2ObjF :: InputKinematics -> MultivarFunc
m2ObjF inp@InputKinematics {..} ks =
    if m1 < m2 then (m2, grad2) else (m1, grad1)
  where
    invs = mkInvisibles inp (vecToVars ks)
    (grad1, grad2, m1, m2) = m2Grad inp invs _p1 _p2 ks

m2Grad :: InputKinematics
       -> Invisibles
       -> FourMomentum  -- ^ p1 (or q1)
       -> FourMomentum  -- ^ p2 (or q2)
       -> Vector Double
       -> (Vector Double, Vector Double, Double, Double)
m2Grad InputKinematics {..} (Invisibles k1 k2) p1 p2 ks = (d1, d2, m1, m2)
  where
    (k1x, k1y, k1z, k2z) = vecToVars ks

    (e1, p1x, p1y, p1z) = epxpypz p1
    m1  = invariantMass [p1, k1]
    m1' = safeDivisor m1
    r1 = e1 / safeDivisor (energy k1)
    d1 = fromList $ ( / m1') <$> [ r1 * k1x - p1x
                                 , r1 * k1y - p1y
                                 , r1 * k1z - p1z
                                 , 0 ]

    (e2, p2x, p2y, p2z) = epxpypz p2
    m2  = invariantMass [p2, k2]
    m2' = safeDivisor m2
    r2 = e2 / safeDivisor (energy k2)
    d2 = fromList $ ( / m2') <$> [ r2 * (k1x - px _ptmiss) + p2x
                                 , r2 * (k1y - py _ptmiss) + p2y
                                 , 0
                                 , r2 * k2z - p2z ]

constraintF :: FourMomentum -> FourMomentum -> InputKinematics -> MultivarFunc
constraintF p1 p2 inp ks = (m1 - m2, grad1 - grad2)
  where
    invs = mkInvisibles inp (vecToVars ks)
    (grad1, grad2, m1, m2) = m2Grad inp invs p1 p2 ks

constraintA, constraintB :: InputKinematics -> MultivarFunc
constraintA inp@InputKinematics {..} = constraintF _p1 _p2 inp
constraintB inp@InputKinematics {..} = constraintF _q1 _q2 inp

safeDivisor :: Double -> Double
safeDivisor x | x >= 0    = max eps0 x
              | otherwise = min (-eps0) x
  where eps0 = 1.0e-8

{-
m2ObjF :: InputKinematics -> MultivarFunc
m2ObjF inp ks = (m, d)
  where
    ks0 = vecToVars ks
    m = fst $ m2ObjF' inp ks0
    d = m2ObjFDNumeric inp ks0

m2ObjF' :: InputKinematics -> Variables -> (Double, Vector Double)
m2ObjF' inp@InputKinematics {..} ks =
    if m1 < m2 then (m2, grad2) else (m1, grad1)
  where
    invs@(Invisibles k1 k2) = mkInvisibles inp ks
    m1 = invariantMass [_p1, k1]
    m2 = invariantMass [_p2, k2]
    (k1x, k1y, k1z, k2z) = ks
    (grad1, grad2, _) = m2Grad inp invs _p1 _p2 (fromList [k1x, k1y, k1z, k2z])

m2ObjFDNumeric :: InputKinematics -> Variables -> Vector Double
m2ObjFDNumeric inp@InputKinematics {..} = fDNumeric f
  where f ks = fst (m2ObjF' inp ks)

m2GradNumeric :: InputKinematics
              -> Invisibles
              -> FourMomentum  -- ^ p1 (or q1)
              -> FourMomentum  -- ^ p2 (or q2)
              -> Vector Double
              -> (Vector Double, Vector Double, Double)
m2GradNumeric inp@(InputKinematics {..}) _ p1 p2 ks = (d1, d2, m1 - m2)
  where
    ks0 = vecToVars ks

    m1F ks' = let Invisibles k1 _ = mkInvisibles inp ks'
              in invariantMass [p1, k1]
    m1 = m1F ks0
    d1 = fDNumeric m1F ks0

    m2F ks' = let Invisibles _ k2 = mkInvisibles inp ks'
              in invariantMass [p2, k2]
    m2 = m2F ks0
    d2 = fDNumeric m2F ks0

fDNumeric :: (Variables -> Double) -> Variables -> Vector Double
fDNumeric f (k1x, k1y, k1z, k2z) = fromList [fk1xD, fk1yD, fk1zD, fk2zD]
  where
    fk1x a = f (a, k1y, k1z, k2z)
    fk1y a = f (k1x, a, k1z, k2z)
    fk1z a = f (k1x, k1y, a, k2z)
    fk2z a = f (k1x, k1y, k1z, a)

    epsN = 1e-8

    fk1xD = fst $ derivCentral epsN fk1x k1x
    fk1yD = fst $ derivCentral epsN fk1y k1y
    fk1zD = fst $ derivCentral epsN fk1z k1z
    fk2zD = fst $ derivCentral epsN fk2z k2z
-}


-- | https://en.wikibooks.org/wiki/Write_Yourself_a_Scheme_in_48_Hours/Parsing

module Main where

import           Control.Monad
import           Data.Char                     (digitToInt, isDigit)
import           Data.Complex
import           Data.Ratio                    (denominator, numerator, (%))
import           Numeric                       (readDec, readFloat, readHex,
                                                readInt, readOct)
import           System.Environment
import           Text.ParserCombinators.Parsec hiding (spaces)

symbol :: Parser Char
symbol = oneOf "!$%&|*+-/:<=>?@^_~" <?> "a symbol"

-- readExpr :: String -> Either ParseError LispVal
-- readExpr = parse parseExpr "lisp"

-- Whitespace
spaces :: Parser ()
spaces = skipMany1 space

-- Return Values
data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             | Character Char
             | String String
             | Bool Bool

parseString :: Parser LispVal
parseString = do
  char '"'
  str <- many (parseHelper <|> noneOf "\\\"")
  char '"'
  return $ String str
  <?> "a string for parseString"
  where parseHelper :: Parser Char
        parseHelper = do
          char '\\'
          x <- oneOf "\\\"nrt"
          return $ case x of
            '\\' -> x
            '"'  -> x
            'n'  -> '\n'
            'r'  -> '\r'
            't'  -> '\t'

parseAtom :: Parser LispVal
parseAtom = do
  first <- letter <|> symbol
  rest <- many $ letter <|> symbol <|> digit
  return $ Atom (first : rest)
  <?> "a atom for parseAtom"

parseBool :: Parser LispVal
parseBool = do
  x <- try (string "#t") <|> string "#f"
  return $ if x == "#t"
           then Bool True
           else Bool False

parseNumber :: Parser LispVal
parseNumber = parseDigit <|> try parseOct <|> try parseDec <|> try parseHex <|> parseBin
  where parseOct = do
          char '#' >> oneOf "oO"
          str <- many1 octDigit
          let [(n, _)] = readOct str
          return $ Number n
        parseDec = do
          char '#' >> oneOf "dD"
          str <- many1 digit
          let [(n, _)] = readDec str
          return $ Number n
        parseHex = do
          char '#' >> oneOf "xX"
          str <- many1 hexDigit
          let [(n, _)] = readHex str
          return $ Number n
        parseBin = do
          char '#' >> oneOf "bB"
          str <- many1 $ oneOf "01"
          let [(n, _)] = readInt 2 (`elem` "01") digitToInt str
          return $ Number n
        parseDigit = do
          str <- many1 digit
          return $ (Number . read) str

parseChar :: Parser LispVal
parseChar = do
  string "#\\"
  x <- many $ letter <|> digit
  return $ case length x of
    0 -> Character '\n'
    1 -> Character $ head x
    _ -> case x of
      "newline" -> Character '\n'
      "tab"     -> Character '\t'
      "space"   -> Character ' '

parseFloat :: Parser LispVal
parseFloat = do
  x <- many1 digit
  char '.'
  y <- many1 digit
  return $ let [(n, _)] = readFloat (x ++ '.' : y) in Float n

parseRatio :: Parser LispVal
parseRatio = do
  x <- many1 digit
  char '/'
  y <- many1 digit
  return $ Ratio (read x % read y)

parseComplex :: Parser LispVal
parseComplex = do
  x <- try parseFloat <|> parseNumber
  char '+'
  y <- try parseFloat <|> parseNumber
  char 'i'
  return $ Complex (toDouble x :+ toDouble y)
  where toDouble :: LispVal -> Double
        toDouble (Float f) = f
        toDouble (Number n) = fromIntegral n

-- Recursive Parsers: Adding lists, dotted lists, and quoted datums

parseList :: Parser LispVal
parseList = do
  char '('
  x <- sepBy parseExpr spaces
  char ')'
  return $ List x

parseDottedList :: Parser LispVal
parseDottedList = do
  char '('
  x <- endBy parseExpr spaces
  y <- char '.' >> spaces >> parseExpr
  char ')'
  return $ DottedList x y

parseQuoted :: Parser LispVal
parseQuoted = do
  char '\''
  x <- parseExpr
  return $ List [Atom "quote", x]

parseBackquote :: Parser LispVal
parseBackquote = do
  char '`'
  x <- parseExpr
  return $ List [Atom "quasiquote", x]

parseUnquote :: Parser LispVal
parseUnquote = do
  char ','
  y <- parseExpr
  return $ List [Atom "unquote", y]

parseUnquoteSplicing :: Parser LispVal
parseUnquoteSplicing = do
  char ',' >> char '@'
  y <- parseExpr
  return $ List [Atom "unquote-splicing", y]

parseExpr :: Parser LispVal
parseExpr = parseAtom
            <|> try parseBool
            <|> try parseRatio
            <|> try parseComplex
            <|> try parseFloat
            <|> try parseNumber
            <|> parseChar
            <|> parseString
            <|> try parseList
            <|> parseDottedList
            <|> parseQuoted
            <|> parseBackquote
            <|> try parseUnquoteSplicing
            <|> parseUnquote

-- Beginning the Evaluator
showVal :: LispVal -> String
showVal (Atom x) = x
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (Ratio x) = show (numerator x) ++ ('/' : show (denominator x))
showVal (Complex x) = show (realPart x) ++ ('+' : show (imagPart x) ++ "i")
showVal (Float x) = show x
showVal (Number x) = show x
showVal (Character x) = "#\\" ++ show x
showVal (String x) = "\"" ++ x ++ "\""
showVal (List x) = "(" ++ unwordsList x ++ ")"
showVal (DottedList first rest)= "(" ++ unwordsList first ++ " . " ++ showVal rest ++ ")"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where
  show = showVal

readExpr :: String -> LispVal
readExpr input = case parse parseExpr "lisp" input of
  Left err -> String $ "No match: " ++ show err
  Right val -> val

-- Beginnings of an evaluator: Primitives
main :: IO ()
main = getArgs >>= print . eval . readExpr . head

eval :: LispVal -> LispVal
eval val@(String _) = val
eval val@(Bool _) = val
eval val@(Ratio _) = val
eval val@(Complex _) = val
eval val@(Float _) = val
eval val@(Number _) = val
eval val@(Character _) = val
eval (List [Atom "quote", val]) = val

-- Adding basic primitives
eval (List (Atom func : args)) = apply func $ map eval args

apply :: String -> [LispVal] -> LispVal
apply func args = maybe (Bool False) ($ args) $ lookup func primitives

primitives :: [(String, [LispVal] -> LispVal)]
primitives = [("+", numericBinop (+)),
              ("-", numericBinop (-)),
              ("*", numericBinop (*)),
              ("/", numericBinop div),
              ("mod", numericBinop mod),
              ("quotient", numericBinop quot),
              ("remainder", numericBinop rem),
              ("symbol?", unaryOp symbolp),
              ("string?", unaryOp stringp),
              ("number?", unaryOp numberp),
              ("bool?", unaryOp boolp),
              ("list?", unaryOp listp),
              ("pair?", unaryOp pairp)]

numericBinop :: (Integer -> Integer -> Integer) -> [LispVal] -> LispVal
numericBinop op args = Number $ foldl1 op $ map unpackNum args
  where unpackNum :: LispVal -> Integer
        unpackNum (Number n) = n
        unpackNum (String n) = let parsed = readDec n in
          if null parsed
          then 0
          else fst $ head parsed
        unpackNum (List [n]) = unpackNum n
        unpackNum _ = 0

unaryOp :: (LispVal -> LispVal) -> [LispVal] -> LispVal
unaryOp op [x] = op x

symbolp :: LispVal -> LispVal
symbolp (Atom _) = Bool True
symbolp _ = Bool False

stringp :: LispVal -> LispVal
stringp (String _) = Bool True
stringp _ = Bool False

numberp :: LispVal -> LispVal
numberp (Number _) = Bool True
numberp _ = Bool False

boolp :: LispVal -> LispVal
boolp (Bool _) = Bool True
boolp _ = Bool False

listp :: LispVal -> LispVal
listp (List _) = Bool True
listp _ = Bool False

pairp :: LispVal -> LispVal
pairp (List _) = Bool True
pairp (DottedList _ _) = Bool True
pairp _ = Bool False


{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE QuasiQuotes               #-}
{-# OPTIONS_GHC  -fplugin=Polysemy.Plugin  #-}
--{-# LANGUAGE AllowAmbiguousTypes       #-}

import qualified Control.Foldl                 as FL

import qualified Data.List                     as L
import qualified Data.Map                      as M
import qualified Data.Vector                   as VB
import qualified Data.Text                     as T

import qualified Text.Printf                   as PF

import qualified Data.Vinyl                    as V
import qualified Frames                        as F
import qualified Frames.CSV                    as F

import qualified Pipes                         as P
import qualified Pipes.Prelude                 as P

import qualified Numeric.LinearAlgebra         as LA

import qualified Frames.Visualization.VegaLite.Data
                                               as FV
import qualified Frames.Visualization.VegaLite.StackedArea
                                               as FV
import qualified Frames.Visualization.VegaLite.LineVsTime
                                               as FV
import qualified Frames.Visualization.VegaLite.ParameterPlots
                                               as FV
import qualified Frames.Visualization.VegaLite.Correlation
                                               as FV

import qualified Frames.Transform              as FT
import qualified Frames.Folds                  as FF
import qualified Frames.MapReduce              as MR
import qualified Frames.Enumerations           as FE
import qualified Frames.Utils                  as FU

import qualified Knit.Report                   as K
import           Polysemy.Error                 ( Error )

import           Data.String.Here               ( here )

import           BlueRipple.Configuration
import           BlueRippleUtilitie.KnitUtils
import           BlueRipple.Data.DataFrames
import           BlueRipple.Data.MRP
import qualified BlueRipple.Model.TurnoutAdjustment
                                               as TA

templateVars = M.fromList
  [ ("lang"     , "English")
  , ("author"   , "Adam Conner-Sax & Frank David")
  , ("pagetitle", "Preference Model & Predictions")
--  , ("tufte","True")
  ]

{- TODO
1. Why are rows still being dropped?  Which col is missing?  In general, write a diagnostic for showing what is missing...
Some answers:
The boring:  CountyFIPS.  I don't want this anyway.  Dropped
The weird:  Missing voted_rep_party.  These are non-voters or voters who didn't vote in house race.  A bit more than 1/3 of
survey responses.  Maybe that's not weird??
-}

main :: IO ()
main = do
  let template = K.FromIncludedTemplateDir "mindoc-pandoc-KH.html"
--  let template = K.FullySpecifiedTemplatePath "pandoc-templates/minWithVega-pandoc.html"
  pandocWriterConfig <- K.mkPandocWriterConfig template
                                               templateVars
                                               K.mindocOptionsF
  eitherDocs <-
    K.knitHtmls (Just "MRP_Basics.Main") K.logAll pandocWriterConfig $ do
      K.logLE K.Info "Loading data..."
      let csvParserOptions =
            F.defaultParser { F.quotingMode = F.RFC4180Quoting ' ' }
          tsvParserOptions = csvParserOptions { F.columnSeparator = "\t" }
          preFilterYears   = FU.filterOnMaybeField @CCESYear (`L.elem` [2016])
      ccesMaybeRecs <- loadToMaybeRecs @CCES_MRP_Raw @(F.RecordColumns CCES)
        tsvParserOptions
        preFilterYears
        ccesTSV
      ccesFrame <-
        fmap transformCCESRow
          <$> maybeRecsToFrame fixCCESRow (const True) ccesMaybeRecs
      let firstFew = take 4 $ FL.fold FL.list ccesFrame
      K.logLE K.Diagnostic
        $  "ccesFrame (first 4 rows):\n"
        <> (T.pack $ show firstFew)
      K.logLE K.Info "Inferring..."
      K.logLE K.Info "Knitting docs..."
  case eitherDocs of
    Right namedDocs ->
      K.writeAllPandocResultsWithInfoAsHtml "reports/html/MRP_Basics" namedDocs
    Left err -> putStrLn $ "pandoc error: " ++ show err


module BondCalculus.StochPyExtraction
  (stochPyModel, generateStochPy, reaction, ReactionVect, reactions, ReactionSystem(..), extractReactionSystem, fromIVPToStochPyModel, simulateStochPy) where

import BondCalculus.Symbolic
import BondCalculus.Vector
import BondCalculus.AST
import BondCalculus.Base
import qualified BondCalculus.AST as AST
import BondCalculus.Transitions (MTS)
import Data.Bifunctor
import BondCalculus.Simulation (Trace)
-- import qualified Data.HashMap.Strict as H
-- import BondCalculus.Simulation (Trace)
-- import Data.String.Utils
import BondCalculus.ODEExtraction (sympyExpr, IVP(..), sympySimplify, runPython)
-- , runPython)

-- import qualified Control.Exception as X
import qualified Data.Map as M
import qualified Data.List as L
import Data.Maybe
import BondCalculus.Processes

--import qualified Numeric.LinearAlgebra as LA

-- import System.IO.Unsafe
-- import qualified System.Process as OS
-- import System.Process (readProcess)
-- import Data.Maybe

type ReactionVect k = Vect (Tensor (ProcessVect Conc) (ProcessVect Conc)) k

reaction :: (Vector k (DirectionVect k), Vector k (ReactionVect k), Nullable (DirectionVect k), Expression k) =>
            [DirectionVect k] -> ReactionVect k
reaction = multilinear react' <$> filter (/=vectZero)
  where react' xs = vect( sourceVect xs :* targetVect xs )
        source (spec :* _) = spec
        targetVect xs = embed $ concretify $ foldl (<|>) (mkAbsBase Nil) (map target xs)
        target (_ :* spec') = spec'
        sourceVect xs = foldl (+>) vectZero (map (embed.source) xs)

reactions :: (Vector k (ProcessVect k), Show k, DoubleExpression k, ExpressionOver Double k) => (Species -> MTS) -> ConcreteAffinityNetwork Double -> ProcessVect k -> ReactionVect k
reactions tr network p = simplify $ dPdt'' reaction tr network p

type ConcVect = Vect String Conc
newtype ReactionSystem = ReactionSystem ([String], [(SymbolicExpr Double, ConcVect, ConcVect)], [Double])
  deriving (Show, Eq, Ord)

extractReactionSystem :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> [Double] -> ReactionSystem
extractReactionSystem env network p inits = ReactionSystem (vars, transitions, inits)
  where tr = tracesGivenNetwork network env
        vars = map (pretty.simplify.snd) (toList p)
        reacts = reactions tr network p
        transitions = [ (simplify coeff, toConcVect u, toConcVect v)
                      | (coeff, u :* v) <- toList reacts]
        toConcVect  = fromList . map (second (pretty.simplify)) . toList

fromIVPToStochPyModel :: IVP Double -> Either String String
fromIVPToStochPyModel ivp =
  if any isNothing exprs
  then Left "An ODE equation has an unbound variable"
  else Right $
  "Output_In_Conc: True\n" ++
  "Species_In_Conc: True\n\n" ++
  L.intercalate "\n" rules ++ "\n" ++
  L.intercalate "\n" initlines
    where
      species = ["S" ++ show i | i <- [(0::Integer)..]]
      varmap = M.fromList $ zip vars species
      IVP (vars,rhss,inits) = ivp
      mkgrowrule spec y =
        "grow" ++ spec ++ ":\n" ++
        "\t$pool > " ++ spec ++ "\n" ++
        "\t" ++ mkgrowrate y ++ "\n"
      mkgrowrate y = "abs(" ++ y ++ " + abs(" ++ y ++ "))/2"
      mkdecayrule spec y =
        "decay" ++ spec ++ ":\n" ++
        "\t" ++ spec ++ " > $pool\n" ++
        "\t" ++ mkdecayrate y ++ "\n"
      mkdecayrate y = "abs(-" ++ y ++ " + abs(" ++ y ++ "))/2"
      exprs = map (sympyExpr varmap) rhss
      growrules = zipWith mkgrowrule species (catMaybes exprs)
      decayrules = zipWith mkdecayrule species (catMaybes exprs)
      rules = zipWith (\x y -> x ++ "\n" ++ y) growrules decayrules
      initlines = [y ++ " = " ++ show y0 | (y,y0) <- zip species inits]

stochPyModel :: ReactionSystem -> Double -> Either String String
stochPyModel (ReactionSystem (vars, transitions, inits)) h =
  if any isNothing transitions'
  then Left "An ODE equation has an unbound variable"
  else Right $
  "Output_In_Conc: True\n" ++
  "Species_In_Conc: True\n\n" ++
  L.intercalate "\n" rules ++ "\n" ++
  "h = " ++ show h ++ "\n" ++
  L.intercalate "\n" initlines
    where
      species = ["S" ++ show i | i <- [(0::Integer)..]]
      reacts = ["R" ++ show i | i <- [(0::Integer)..]]
      varmap = M.fromList $ zip vars species
      varmaph = fmap (\s -> "(" ++ s ++ "*h)") varmap
      rules = [
            i ++ ":\n" ++
            "\t" ++ formatVector u ++ " > " ++ formatVector v ++ "\n" ++
            "\t" ++ rate ++ "\n"
          | (i, (rate, u, v)) <- reacts `zip` catMaybes transitions'
        ]
      transitions' = fmap (\(rate,u,v) -> do
                      rate' <- sympyExpr varmaph rate
                      u' <- vectorVarLookup u
                      v' <- vectorVarLookup v
                      return (sympySimplify ("(" ++ rate' ++ ")/h"),u',v')) transitions
      vectorVarLookup u =
        let us = map (second (`M.lookup` varmap)) $ toList u
        in if any (isNothing.snd) us
          then Nothing
          else Just (fromList [(k,v) | (k, Just v) <- us])
      formatVector u = if null us
                       then "$pool"
                       else L.intercalate " + "
                              ["{" ++ show k ++ "} " ++ i | (k,i) <- us]
        where us = toList u
      initlines = [y ++ " = " ++ show (y0/h) | (y,y0) <- zip species inits]

generateStochPy :: String -> AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> Double -> [Double] -> IO String
generateStochPy filename env network p h inits = case stochPyModel (extractReactionSystem env network p inits) h of
  Right script -> do
    -- putStrLn $ "Python script:\n\n" ++ script
    writeFile filename script
    return script
  Left _ -> undefined

simulateStochPy :: String -> Int -> AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> Double -> [Double] -> IO Trace
simulateStochPy method steps env network p h inits = case stochPyModel (extractReactionSystem env network p inits) h of
  Right stochpy -> do
    -- putStrLn $ "Python script:\n\n" ++ script
    writeFile "/tmp/stochpy.psc" stochpy
    let script = "import os, sys\n" ++
                 "f = open(os.devnull, 'w')\n" ++
                 "stdout = sys.stdout\n" ++
                 "sys.stdout = f\n" ++
                 "import stochpy\n" ++
                 "smod = stochpy.SSA()\n" ++
                 "smod.Model('/tmp/stochpy.psc')\n" ++
                 "smod.DoStochSim(method='" ++ method ++ "', end=" ++
                   show steps ++ ", mode='time')\n" ++
                 "sys.stdout=stdout\n" ++
                 "xss = smod.data_stochsim.getSimData(" ++ L.intercalate "," ["'S" ++ show i ++ "'" | (i,_) <- [0..] `zip` kis] ++ ")\n" ++
                 "xss[:,1:] *= " ++ show h ++ "\n" ++
                 "for xs in xss:\n" ++
                 "    print(' '.join(map('{:.18f}'.format, list(xs))))\n"
    writeFile "stochpy_script.py" script
    res <- runPython script
    let yss = (map (map read.words) $ lines res) :: [[Double]]
    return [ (head ys, toP (tail ys)) | ys <- yss ]
  Left _ -> undefined
  where toP ks = fromList [(k,i) | (k, (_,i)) <- ks `zip` kis]
        kis = toList p


{-# LANGUAGE BangPatterns #-}
module Polestar.Eval
  (termShift
  ,termTypeSubst
  ,termSubstD
  ,termSubst
  ,ValueBinding(..)
  ,getValueFromContext
  ,eval1
  ,eval
  ) where
import Polestar.Type
import Polestar.TypeCheck
import Data.Complex
import GHC.Float (expm1, log1p)

termSubstD :: Int -> Term -> Int -> Term -> Term
termSubstD !depth s !i t = case t of
  TmAbs name ty body -> TmAbs name (typeShift (-1) i ty) (termSubstD depth s (i + 1) body)
  TmTyAbs name bound body -> TmTyAbs name (typeShift (-1) i <$> bound) (termSubstD depth s (i + 1) body)
  TmRef j | j == i -> termShift depth 0 s
          | j > i -> TmRef (j - 1)
          | otherwise -> t
  TmApp u v -> TmApp (termSubstD depth s i u) (termSubstD depth s i v)
  TmTyApp u t -> TmTyApp (termSubstD depth s i u) (typeShift (-1) i t)
  TmLet name def body -> TmLet name (termSubstD depth s i def) (termSubstD depth s (i + 1) body) -- ?
  TmIf cond then_ else_ -> TmIf (termSubstD depth s i cond) (termSubstD depth s i then_) (termSubstD depth s i else_)
  TmPrim _ -> t
  TmCoerce x ty  -> TmCoerce (termSubstD depth s i x) (typeShift (-1) i ty)
  TmAlt name tys body -> TmAlt name (map (typeShift (-1) i) tys) (termSubstD depth s (i + 1) body)
  TmTuple components -> TmTuple $ termSubstD depth s i <$> components
  TmProj tuple j -> TmProj (termSubstD depth s i tuple) j
  TmCoherentTuple components -> TmCoherentTuple $ termSubstD depth s i <$> components

-- replaces occurrences of TRef j (j > i) with TRef (j-1), and TRef i with the given term
termSubst = termSubstD 0

natFromValue :: Term -> Either String Integer
natFromValue (TmPrim (PVZero)) = return 0
natFromValue (TmPrim (PVInt x)) | x >= 0 = return x
natFromValue _ = Left "type error (expected a natural number)"

intFromValue :: Term -> Either String Integer
intFromValue (TmPrim (PVZero)) = return 0
intFromValue (TmPrim (PVInt x)) = return x
intFromValue _ = Left "type error (expected an integer)"

nnrealFromValue :: Term -> Either String Double
nnrealFromValue (TmPrim (PVZero)) = return 0
nnrealFromValue (TmPrim (PVInt x)) | x >= 0 = return (fromIntegral x)
nnrealFromValue (TmPrim (PVReal x)) | x >= 0 = return x
nnrealFromValue _ = Left "type error (expected a non-negative real number)"

realFromValue :: Term -> Either String Double
realFromValue (TmPrim (PVZero)) = return 0
realFromValue (TmPrim (PVInt x)) = return (fromIntegral x)
realFromValue (TmPrim (PVReal x)) = return x
realFromValue _ = Left "type error (expected a real number)"

imaginaryFromValue :: Term -> Either String Double
imaginaryFromValue (TmPrim (PVZero)) = return 0
imaginaryFromValue (TmPrim (PVImaginary x)) = return x
imaginaryFromValue _ = Left "type error (expected an imaginary number)"

complexFromValue :: Term -> Either String (Complex Double)
complexFromValue (TmPrim (PVZero)) = return 0
complexFromValue (TmPrim (PVInt x)) = return (fromIntegral x)
complexFromValue (TmPrim (PVReal x)) = return (x:+0)
complexFromValue (TmPrim (PVImaginary x)) = return (0:+x)
complexFromValue (TmPrim (PVComplex z)) = return z
complexFromValue _ = Left "type error (expected a complex number)"

boolFromValue :: Term -> Either String Bool
boolFromValue (TmPrim (PVBool x)) = return x
boolFromValue _ = Left "type error (expected a boolean)"

applyBuiltinUnary :: Builtin -> Term -> Either String Term
applyBuiltinUnary f v = case f of
  BNegate -> case v of
    TmPrim PVZero -> return v
    TmPrim (PVInt x) -> return (TmPrim $ PVInt $ negate x)
    TmPrim (PVReal x) -> return (TmPrim $ PVReal $ negate x)
    TmPrim (PVImaginary x) -> return (TmPrim $ PVImaginary $ negate x)
    TmPrim (PVComplex x) -> return (TmPrim $ PVComplex $ negate x)
    _ -> Left "type error (expected an integer or a real number)"
  BLogicalNot -> (TmPrim . PVBool . not) <$> boolFromValue v
  BNatToInt -> (TmPrim . PVInt) <$> natFromValue v
  BNatToNNReal -> (TmPrim . PVReal . fromIntegral) <$> natFromValue v
  BIntToNat -> (\x -> if x >= 0 then TmPrim (PVInt x) else TmPrim PVZero) <$> intFromValue v
  BIntToReal -> (TmPrim . PVReal . fromIntegral) <$> intFromValue v
  BNNRealToReal -> (TmPrim . PVReal) <$> realFromValue v
  BRealToComplex -> (TmPrim . PVComplex . (:+0)) <$> realFromValue v
  BMkImaginary -> (TmPrim . PVImaginary) <$> realFromValue v
  BImaginaryToComplex -> (TmPrim . PVComplex . (0:+)) <$> imaginaryFromValue v
  BRealPart -> case v of
    TmPrim PVZero -> return v
    TmPrim (PVInt _) -> return v
    TmPrim (PVReal _) -> return v
    TmPrim (PVImaginary _) -> return $ TmPrim PVZero
    TmPrim (PVComplex z) -> return $ TmPrim (PVReal (realPart z))
    _ -> Left "type error (expected a complex number)"
  BImagPart -> case v of
    TmPrim PVZero -> return $ TmPrim PVZero
    TmPrim (PVInt _) -> return $ TmPrim PVZero
    TmPrim (PVReal _) -> return $ TmPrim PVZero
    TmPrim (PVImaginary x) -> return $ TmPrim (PVReal x)
    TmPrim (PVComplex z) -> return $ TmPrim (PVReal (imagPart z))
    _ -> Left "type error (expected a complex number)"
  BAbs -> case v of
    TmPrim PVZero -> return $ TmPrim PVZero
    TmPrim (PVInt x) -> return $ TmPrim (PVInt (abs x))
    TmPrim (PVReal x) -> return $ TmPrim (PVReal (abs x))
    TmPrim (PVImaginary x) -> return $ TmPrim (PVReal (abs x))
    TmPrim (PVComplex z) -> return $ TmPrim (PVReal (magnitude z))
    _ -> Left "type error (expected a complex number)"
  BSqrt -> case v of
    TmPrim PVZero -> return $ TmPrim PVZero
    TmPrim (PVInt x) | x >= 0 -> return $ TmPrim (PVReal (sqrt $ fromIntegral x))
    TmPrim (PVReal x) | x >= 0 -> return $ TmPrim (PVReal (sqrt x))
    TmPrim (PVImaginary x) -> return $ TmPrim (PVComplex (sqrt (0:+x)))
    TmPrim (PVComplex z) -> return $ TmPrim (PVComplex (sqrt z))
    _ -> Left "type error (expected a complex number)"
  BExp -> makeOverloadedFn [(TmPrim . PVReal . exp) <$> realFromValue v
                           ,(TmPrim . PVComplex . exp) <$> complexFromValue v
                           ]
  BExpm1 -> makeOverloadedFn [(TmPrim . PVReal . expm1) <$> realFromValue v
                             ,(TmPrim . PVComplex . expm1) <$> complexFromValue v
                             ]
  BLog -> makeOverloadedFn [(TmPrim . PVReal . log) <$> nnrealFromValue v
                           ,(TmPrim . PVComplex . log) <$> complexFromValue v
                           ]
  BLog1p -> makeOverloadedFn [(TmPrim . PVReal . log1p) <$> nnrealFromValue v
                             ,(TmPrim . PVComplex . log1p) <$> complexFromValue v
                             ]
  BSin -> makeOverloadedFn [(TmPrim . PVReal . sin) <$> realFromValue v
                           ,(TmPrim . PVImaginary . sinh) <$> imaginaryFromValue v
                           ,(TmPrim . PVComplex . sin) <$> complexFromValue v
                           ]
  BCos -> makeOverloadedFn [(TmPrim . PVReal . cos) <$> realFromValue v
                           ,(TmPrim . PVReal . cosh) <$> imaginaryFromValue v
                           ,(TmPrim . PVComplex . cos) <$> complexFromValue v
                           ]
  BTan -> makeOverloadedFn [(TmPrim . PVReal . tan) <$> realFromValue v
                           ,(TmPrim . PVComplex . tan) <$> complexFromValue v
                           ]
  BSinh -> makeOverloadedFn [(TmPrim . PVReal . sinh) <$> realFromValue v
                           ,(TmPrim . PVImaginary . sin) <$> imaginaryFromValue v
                           ,(TmPrim . PVComplex . sinh) <$> complexFromValue v
                           ]
  BCosh -> makeOverloadedFn [(TmPrim . PVReal . cosh) <$> realFromValue v
                           ,(TmPrim . PVReal . cos) <$> imaginaryFromValue v
                           ,(TmPrim . PVComplex . cosh) <$> complexFromValue v
                           ]
  BTanh -> makeOverloadedFn [(TmPrim . PVReal . tan) <$> realFromValue v
                            ,(TmPrim . PVImaginary . tanh) <$> imaginaryFromValue v
                            ,(TmPrim . PVComplex . tan) <$> complexFromValue v
                            ]
  BAsin -> makeOverloadedFn [(TmPrim . PVImaginary . asinh) <$> imaginaryFromValue v
                            ,(TmPrim . PVComplex . asin) <$> complexFromValue v
                            ]
  BAcos -> makeOverloadedFn [(TmPrim . PVComplex . acos) <$> complexFromValue v
                            ]
  BAtan -> makeOverloadedFn [(TmPrim . PVReal . atan) <$> realFromValue v
                            ,(TmPrim . PVComplex . atan) <$> complexFromValue v
                            ]
  BAsinh -> makeOverloadedFn [(TmPrim . PVReal . asinh) <$> realFromValue v
                             ,(TmPrim . PVComplex . asinh) <$> complexFromValue v
                             ]
  BAcosh -> makeOverloadedFn [(TmPrim . PVReal . acosh) <$> realFromValue v
                             ]
  BAtanh -> makeOverloadedFn [(TmPrim . PVImaginary . atan) <$> imaginaryFromValue v
                             ,(TmPrim . PVComplex . atanh) <$> complexFromValue v
                             ]
  _ -> Left "applyBuiltinUnary: not unary"

makeOverloadedFn :: [Either String Term] -> Either String Term
makeOverloadedFn [] = Left "no overload"
makeOverloadedFn [x] = x
makeOverloadedFn (Left _ : xs) = makeOverloadedFn xs
makeOverloadedFn (Right x : xs) = Right x

isZero :: Term -> Bool
isZero (TmPrim PVZero) = True
isZero _ = False

isNonNegative :: Term -> Bool
isNonNegative (TmPrim PVZero) = True
isNonNegative (TmPrim (PVInt x)) = x >= 0
isNonNegative (TmPrim (PVReal x)) = x >= 0
isNonNegative _ = False

applyBuiltinBinary :: Builtin -> Term -> Term -> Either String Term
applyBuiltinBinary f u v = case f of
  BAdd -> makeOverloadedFn [if isZero u && isZero v then return (TmPrim PVZero) else Left "not zero"
                           ,(TmPrim . PVInt) <$> ((+) <$> intFromValue u <*> intFromValue v)
                           ,(TmPrim . PVReal) <$> ((+) <$> realFromValue u <*> realFromValue v)
                           ,(TmPrim . PVImaginary) <$> ((+) <$> imaginaryFromValue u <*> imaginaryFromValue v)
                           ,(TmPrim . PVComplex) <$> ((+) <$> complexFromValue u <*> complexFromValue v)
                           ]
  BSub -> makeOverloadedFn [if isZero u && isZero v then return (TmPrim PVZero) else Left "not zero"
                           ,(TmPrim . PVInt) <$> ((-) <$> intFromValue u <*> intFromValue v)
                           ,(TmPrim . PVReal) <$> ((-) <$> realFromValue u <*> realFromValue v)
                           ,(TmPrim . PVImaginary) <$> ((-) <$> imaginaryFromValue u <*> imaginaryFromValue v)
                           ,(TmPrim . PVComplex) <$> ((-) <$> complexFromValue u <*> complexFromValue v)
                           ]
  BMul -> makeOverloadedFn [if isZero u || isZero v then return (TmPrim PVZero) else Left "not zero"
                           ,(TmPrim . PVInt) <$> ((*) <$> intFromValue u <*> intFromValue v)
                           ,(TmPrim . PVReal) <$> ((*) <$> realFromValue u <*> realFromValue v)
                           ,(TmPrim . PVImaginary) <$> ((*) <$> realFromValue u <*> imaginaryFromValue v)
                           ,(TmPrim . PVImaginary) <$> ((*) <$> imaginaryFromValue u <*> realFromValue v)
                           ,(TmPrim . PVReal . negate) <$> ((*) <$> imaginaryFromValue u <*> imaginaryFromValue v)
                           ,(TmPrim . PVComplex) <$> ((*) <$> complexFromValue u <*> complexFromValue v)
                           ]
  BDiv -> makeOverloadedFn [(TmPrim . PVReal) <$> ((/) <$> realFromValue u <*> realFromValue v)
                           ,(TmPrim . PVImaginary) <$> ((/) <$> imaginaryFromValue u <*> realFromValue v)
                           ,(TmPrim . PVImaginary . negate) <$> ((/) <$> realFromValue u <*> imaginaryFromValue v)
                           ,(TmPrim . PVReal) <$> ((/) <$> imaginaryFromValue u <*> imaginaryFromValue v)
                           ,(TmPrim . PVComplex) <$> ((/) <$> complexFromValue u <*> complexFromValue v)
                           ]
  BPow -> makeOverloadedFn [if isZero v then return (TmPrim (PVInt 1)) else Left "not zero"
                           ,(TmPrim . PVInt) <$> ((^) <$> intFromValue u <*> natFromValue v)
                           ,(TmPrim . PVReal) <$> ((**) <$> nnrealFromValue u <*> realFromValue v)
                           ,(TmPrim . PVReal) <$> ((^^) <$> realFromValue u <*> intFromValue v)
                           ,(TmPrim . PVComplex) <$> ((**) <$> complexFromValue u <*> complexFromValue v)
                           ]
  BTSubNat -> do u' <- natFromValue u
                 v' <- natFromValue v
                 let w | u' >= v' = u' - v'
                       | otherwise = 0
                 return (TmPrim (PVInt w))
  BLt -> (TmPrim . PVBool) <$> ((<) <$> realFromValue u <*> realFromValue v)
  BLe -> (TmPrim . PVBool) <$> ((<=) <$> realFromValue u <*> realFromValue v)
  BEqual -> (TmPrim . PVBool) <$> ((==) <$> complexFromValue u <*> complexFromValue v)
  BMax -> makeOverloadedFn [(TmPrim . PVInt) <$> (max <$> intFromValue u <*> intFromValue v)
                           ,(TmPrim . PVReal) <$> (max <$> realFromValue u <*> realFromValue v)
                           ]
  BMin -> makeOverloadedFn [if (isZero u && isNonNegative v) || (isZero v && isNonNegative u) then return (TmPrim PVZero) else Left "not zero"
                           ,(TmPrim . PVInt) <$> (min <$> intFromValue u <*> intFromValue v)
                           ,(TmPrim . PVReal) <$> (min <$> realFromValue u <*> realFromValue v)
                           ]
  BIntDiv -> (TmPrim . PVInt) <$> (div <$> intFromValue u <*> intFromValue v)
  BIntMod -> (TmPrim . PVInt) <$> (mod <$> intFromValue u <*> intFromValue v)
  BGcd -> (TmPrim . PVInt) <$> (gcd <$> intFromValue u <*> intFromValue v)
  BLcm -> (TmPrim . PVInt) <$> (lcm <$> intFromValue u <*> intFromValue v)
  BLogicalAnd -> (TmPrim . PVBool) <$> ((&&) <$> boolFromValue u <*> boolFromValue v)
  BLogicalOr -> (TmPrim . PVBool) <$> ((||) <$> boolFromValue u <*> boolFromValue v)
  _ -> Left "not implemented yet; sorry"

data ValueBinding = ValueBind !Term
                  | TypeBind
                  deriving (Eq,Show)

getValueFromContext :: [ValueBinding] -> Int -> Term
getValueFromContext ctx i
  | i < length ctx = case ctx !! i of
                       ValueBind x -> x
                       b -> error ("TRef: expected a variable binding, found " ++ show b)
  | otherwise = error "TRef: index out of bounds"

eval1 :: [ValueBinding] -> Term -> Either String Term
eval1 ctx t = case t of
  TmPrim _ -> return t
  TmAbs _ _ _ -> return t
  TmTyAbs _ _ _ -> return t
  TmRef i -> return $ getValueFromContext ctx i
  TmApp u v
    | isValue u && isValue v -> case u of
        TmAbs _name _ty body -> return $ termSubst v 0 body -- no type checking here
        TmPrim (PVBuiltin f) | isUnary f -> applyBuiltinUnary f v
        TmPrim (PVBuiltin _) -> return t -- partial application
        TmApp (TmPrim (PVBuiltin f)) u' | isBinary f -> applyBuiltinBinary f u' v
        TmApp (TmApp (TmTyApp (TmPrim (PVBuiltin BIterate)) ty) n) z -> do
          n' <- natFromValue n
          if n' == 0
            then return $ z
            else return $ TmApp (TmApp (TmApp (TmTyApp (TmPrim (PVBuiltin BIterate)) ty) (TmPrim (PVInt (n' - 1)))) (TmApp v z)) v
        _ -> Left "invalid function application (expected function type)"
    | isValue u -> TmApp u <$> (eval1 ctx v)
    | otherwise -> TmApp <$> (eval1 ctx u) <*> pure v
  TmTyApp u ty
    | isValue u -> case u of
        TmTyAbs _name _bound body -> return $ termTypeSubst ty 0 body
        _ -> Left "invalid type application (expected forall type)"
    | otherwise -> TmTyApp <$> eval1 ctx u <*> pure ty
  TmLet name def body
    | isValue def -> return $ termSubst def 0 body
    | otherwise -> TmLet name <$> eval1 ctx def <*> pure body
  TmIf cond then_ else_
    | isValue cond -> case cond of
        TmPrim (PVBool True) -> return then_  -- no type checking here
        TmPrim (PVBool False) -> return else_ -- no type checking here
        _ -> Left "if-then-else: condition must be boolean"
    | otherwise -> TmIf <$> (eval1 ctx cond) <*> pure then_ <*> pure else_
  TmCoerce x ty
    | isValue x -> return x
    | otherwise -> TmCoerce <$> eval1 ctx x <*> pure ty
  TmAlt name tys x
    | isValue x -> return $ termTypeSubst TyUnit 0 x -- dummy type
    | otherwise -> TmAlt name tys <$> eval1 ctx x
  TmTuple components -> case span isValue components of
    (_,[]) -> return t
    (v,w:ws) -> eval1 ctx w >>= \w' -> return (TmTuple (v ++ (w':ws)))
  TmProj tuple j
    | isValue tuple -> case tuple of
        TmTuple components | length components > j -> return $ components !! j
                           | otherwise -> Left "tuple too short"
        _ -> Left "projection: not a tuple"
    | otherwise -> TmProj <$> eval1 ctx tuple <*> pure j
  TmCoherentTuple components -> Left "coherenet tuple: not supported yet"

eval :: [ValueBinding] -> Term -> Either String Term
eval ctx t | isValue t = return t
           | otherwise = eval1 ctx t >>= eval ctx


module HNNWrapper
  (

  ) where

import AI.HNN.FF.Network
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.HMatrix
import Foreign.Storable
import System.Random
import qualified Data.Vector as V

randomizeMatrixValue :: V.Vector (Matrix a) -> IO (V.Vector (Matrix a))
randomizeMatrixValue = undefined

setMatrixValue :: (Element a) => (Int, Int) -> a -> Matrix a -> Matrix a
setMatrixValue (i,j) a matrix = fromRows newRows
  where
    oldRows = toRows matrix
    oldRow  = oldRows !! i
    newRow  = fromList $ replaceElement j a (toList oldRow)
    newRows = replaceElement i newRow oldRows

(!!+) :: Maybe [a] -> Int -> Maybe a
(!!+) Nothing     _ = Nothing
(!!+) (Just list) i = if i >= length list || i < 0 then Nothing else Just (list !! i)

replaceElement :: Int -> a -> [a] -> [a]
replaceElement i x xs = let (ys,zs) = splitAt i xs
                        in if length zs > 0 then ys ++ [x] ++ (tail zs)
                           else ys ++ [x]


{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Data.Elem.LAPACK.C
-- Copyright  : Copyright (c) , Patrick Perry <patperry@stanford.edu>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@stanford.edu>
-- Stability  : experimental
--
-- Low-level interface to LAPACK.
--

module Data.Elem.LAPACK.C
    where

import Control.Exception( assert )
import Control.Monad
import Data.Complex( Complex )
import Data.Elem.BLAS.Level3
import Data.Matrix.Class( TransEnum(..), SideEnum(..) )
import Foreign
import LAPACK.CTypes

import Data.Elem.LAPACK.Double
import Data.Elem.LAPACK.Zomplex

-- | The LAPACK typeclass.
class (BLAS3 e) => LAPACK e where
    geqrf :: Int -> Int -> Ptr e -> Int -> Ptr e -> IO ()
    gelqf :: Int -> Int -> Ptr e -> Int -> Ptr e -> IO ()
    unmqr :: SideEnum -> TransEnum -> Int -> Int -> Int -> Ptr e -> Int -> Ptr e -> Ptr e -> Int -> IO ()
    unmlq :: SideEnum -> TransEnum -> Int -> Int -> Int -> Ptr e -> Int -> Ptr e -> Ptr e -> Int -> IO ()
    larfg :: Int -> Ptr e -> Ptr e -> Int -> IO e

callWithWork :: (Storable e) => (Ptr e -> Int -> IO a) -> IO a
callWithWork call =
    alloca $ \pQuery -> do
        call pQuery (-1)
        ldWork <- peek (castPtr pQuery) :: IO Double
        let lWork = max 1 $ ceiling ldWork
        allocaArray lWork $ \pWork -> do
            call pWork lWork

checkInfo :: Int -> IO ()
checkInfo info = assert (info == 0) $ return ()

instance LAPACK Double where
    geqrf m n pA ldA pTau =
        checkInfo =<< callWithWork (dgeqrf m n pA ldA pTau)
    gelqf m n pA ldA pTau =
        checkInfo =<< callWithWork (dgelqf m n pA ldA pTau)
    unmqr s t m n k pA ldA pTau pC ldC =
        checkInfo =<< callWithWork (dormqr (cblasSide s) (cblasTrans t) m n k pA ldA pTau pC ldC)
    unmlq s t m n k pA ldA pTau pC ldC =
        checkInfo =<< callWithWork (dormlq (cblasSide s) (cblasTrans t) m n k pA ldA pTau pC ldC)
    larfg n alpha x incx = with 0 $ \pTau ->
        dlarfg n alpha x incx pTau >> peek pTau

instance LAPACK (Complex Double) where
    geqrf m n pA ldA pTau =
        checkInfo =<< callWithWork (zgeqrf m n pA ldA pTau)
    gelqf m n pA ldA pTau =
        checkInfo =<< callWithWork (zgelqf m n pA ldA pTau)
    unmqr s t m n k pA ldA pTau pC ldC =
        checkInfo =<< callWithWork (zunmqr (cblasSide s) (cblasTrans t) m n k pA ldA pTau pC ldC)
    unmlq s t m n k pA ldA pTau pC ldC =
        checkInfo =<< callWithWork (zunmlq (cblasSide s) (cblasTrans t) m n k pA ldA pTau pC ldC)
    larfg n alpha x incx = with 0 $ \pTau ->
        zlarfg n alpha x incx pTau >> peek pTau


{-|
Module      : HABQTlib.UnsafeAPI

This module contains functions for performing and simulating HABQT in Haskell.

__Caution__: functions in this module perform no input validation and are partial. For a safe API refer to "HABQTlib".
-}
module HABQTlib.UnsafeAPI where

import Control.Applicative (liftA2)
import Control.Monad (replicateM)
import Control.Monad.State.Lazy
import Data.Maybe (fromJust)
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.Data.Particle
import HABQTlib.MeasurementProcessing
import HABQTlib.RandomStates
import qualified Numeric.LinearAlgebra as LA
import Streaming
import qualified Streaming.Prelude as S
import qualified System.Random.MWC as MWC
import Text.Printf (printf)

-- | Tomography keeps track of the particle hierarchy and list of previous
-- measurement results, IO is used for verbose output and assorted random state
-- generation.
type TomState = StateT (ParticleHierarchy, [PureStateVector]) IO

-- | Tomography function takes a measurement result and returns state-dependent
-- Bayesian mean estimate of state and the optimal next POVM to perform.
type TomFun = PureStateVector -> TomState (DensityMatrix, PurePOVM)

-- | Given parameters such as output verbosity level and number of quantum
-- bits, set up the tomography function.
tomographyFun' ::
     QBitNum -- ^ Number of quantum bits under tomography
  -> MHMCiter -- ^ Number of MHMC iterations to perform when resampling
  -> OptIter -- ^ Number of POVM optimisation steps to perform
  -> OutputVerb -- ^ Verbosity of stdout output
  -> MWC.GenIO -- ^ IO generator for variates from "System.Random.MWC"
  -> TomFun
tomographyFun' nq mi oi outv gen nextResult = do
  (ph, ms) <- get
  let nextPH = updateParticleHierarchy nextResult ph
      dim = LA.rows . getStateVector $ nextResult
      effectiveSizes =
        V.map (liftA2 (/) effectiveSize (fromIntegral . ptsNumber)) nextPH
      ra = ResampleArgs outv gen dim mi
      resampleC es pts =
        if es < 0.5
          then resample ra (nextResult : ms) pts
          else return pts
  nextPH' <- lift $ V.zipWithM resampleC effectiveSizes nextPH
  sv0s <- liftIO $ replicateM nq (genPureSV 2)
  let nextEstimate = getMixedEstimate nextPH'
      nextPOVM = optimiseSingleQbPOVM oi sv0s nextPH'
  put (nextPH', nextResult : ms)
  return (nextEstimate, nextPOVM)

-- | Given a true state's density matrix and parameters, set up a simulation of
-- quantum tomography that outputs infidelity between mean estimates and true
-- state.
simulatedTomography' ::
     DensityMatrix -- ^ True state's density matrix
  -> QBitNum -- ^ Number of quantum bits under tomography
  -> MHMCiter -- ^ Number of MHMC iterations to perform when resampling
  -> OptIter -- ^ Number of POVM optimisation steps to perform
  -> OutputVerb -- ^ Verbosity of stdout output
  -> MWC.GenIO -- ^ IO generator for variates from "System.Random.MWC"
  -> StateT PurePOVM TomState Double
simulatedTomography' trueDM nq mi oi outv gen = do
  povm <- get
  nextResult <- liftIO $ simulateMeasuremet trueDM povm gen
  (nextEstimate, nextPOVM) <- lift $ tomographyFun' nq mi oi outv gen nextResult
  (nextPH, _) <- lift get
  put nextPOVM
  let fid = fidelityDM trueDM nextEstimate
  when (outv > NoOutput) . liftIO $ do
    let dim = 2 ^ nq
        rankFids =
          V.map
            (fidelityDM trueDM . snd . getWDM . reduceParticlesToMean)
            nextPH
        weightsAndFids =
          V.zip3 (V.enumFromN (1 :: Rank) dim) (V.map ptsWeight nextPH) rankFids
    putStrLn ""
    V.mapM_
      (\(a, b, c) ->
         printf "(Rank: %4d, Weight: %10.9f, Fidelity: %10.9f)\n" a b c)
      weightsAndFids
  return $ 1 - fid

-- | Stream simulated tomography results.
streamResults' ::
     QBitNum -- ^ Number of quantum bits under tomography
  -> Rank -- ^ Rank of true state
  -> NumberOfParticles -- ^ Number of particles (per rank) to use for tomography
  -> MHMCiter -- ^ Number of MHMC iterations to perform when resampling
  -> OptIter -- ^ Number of POVM optimisation steps to perform
  -> OutputVerb -- ^ Verbosity of stdout output
  -> Stream (Of Double) IO ()
streamResults' nq rank pn mi oi outv = do
  let dim = 2 ^ nq
  trueDM <- liftIO $ genDM dim rank
  ph <- liftIO $ initialiseParticleHierarchy dim pn
  gen <- liftIO MWC.createSystemRandom
  rPOVM <-
    liftIO $
    productPOVM <$>
    replicateM nq (mkAntipodalPOVM . fromJust . svToAngles <$> genPureSV 2)
  let tomS = S.repeatM (simulatedTomography' trueDM nq mi oi outv gen)
      tomS' = evalStateT (distribute tomS) rPOVM
      initInfid = 1 - fidelityDM trueDM (getMixedEstimate ph)
  S.yield initInfid
  evalStateT (distribute tomS') (ph, [])


-- Tasty makes it easy to test your code. It is a test framework that can
-- combine many different types of tests into one suite. See its website for
-- help: <http://documentup.com/feuerbach/tasty>.
import Test.Tasty
-- Hspec is one of the providers for Tasty. It provides a nice syntax for
-- writing tests. Its website has more info: <https://hspec.github.io>.
-- import Test.Tasty.Hspec
import Test.Tasty.HUnit

import ExplicitSimplexStream
import Persistence

import Numeric.LinearAlgebra

-- EXAMPLE STREAMS --

stream1 :: Stream Int
stream1 = addSimplex (addVertex (addVertex (addVertex initializeStream 1) 2) 3) (Simplex [1,2])

stream2 :: Stream Int
stream2 = addVertex (addSimplex (addVertex (addVertex initializeStream 1) 2) (Simplex [2,1])) 3

stream3 :: Stream Int
stream3 = addVertex (addVertex (addVertex (addVertex initializeStream 1) 2) 3) 4

stream4 :: Stream Int
stream4 = initializeStream

stream5 :: Stream Int
stream5 = addSimplex stream4 (Simplex [1,2,3,4])

----

main :: IO ()
main = do
  defaultMain (testGroup "SimplexStream tests" [addSingleVertexTest, initializeStreamTest, streamEqualityTest, streamNumVerticesEmptyTest, streamNumVertices4CellTest, streamGetSizeEmptyTest, streamGetSize4CellTest, streamGetSize4VertexTest, streamToOrderedSimplexListFourVerticesTest, streamToOrderedSimplexListThreeVerticesOneEdgeTest, getBoundaryMapTest, getBoundaryMapTest2, getBoundaryMapTest3, getHomologyDimensionTest, persistenceTest, persistenceTest2, persistenceTest3, persistenceTest4])

addSingleVertexTest :: TestTree
addSingleVertexTest = testCase "Testing addition of single vertex"
  (assertEqual "Should return true for search of vertex 5" (True) (isVertexInStream (addVertex (initializeStream :: Stream Int) 5) 5))

initializeStreamTest :: TestTree
initializeStreamTest = testCase "Testing initialization of stream"
  (assertEqual "Should return Simplices []" (Simplices [Simplex []] :: Stream Int) (initializeStream :: Stream Int))

-- Testing stream equality.
-- Order of Simplex in Stream data type _should not matter_
streamEqualityTest :: TestTree
streamEqualityTest = testCase "Testing quality of streams"
  (assertEqual "Should return Simplices []" (True) (stream1 == stream2))

-- Test numVertices
streamNumVerticesEmptyTest :: TestTree
streamNumVerticesEmptyTest = testCase "Testing number of vertices in empty stream"
  (assertEqual "Should return 0" (0) (numVertices stream4))

streamNumVertices4CellTest :: TestTree
streamNumVertices4CellTest = testCase "Testing number of vertices in a 4-cell"
  (assertEqual "Should return 4" (4) (numVertices stream5))

-- Test getSize
streamGetSizeEmptyTest :: TestTree
streamGetSizeEmptyTest = testCase "Testing get size on empty stream."
  (assertEqual "Should return 1 (the null-cell)" (1) (getSize stream4))

streamGetSize4CellTest :: TestTree
streamGetSize4CellTest = testCase "Testing get size on 4-cell stream."
  (assertEqual "Should return 16" (16) (getSize stream5))

streamGetSize4VertexTest :: TestTree
streamGetSize4VertexTest = testCase "Testing get size on stream with 4 vertices."
  (assertEqual "Should return 5 (four vertices + 1 null cell)." (5) (getSize stream3))

-- Test streamToOrderedSimplexList
streamToOrderedSimplexListFourVerticesTest :: TestTree
streamToOrderedSimplexListFourVerticesTest = testCase "Testing streamToOrderedSimplexList on stream with 4 vertices."
  (assertEqual "Should return object with lengths 0 and 1 with four vertices inside the 1 key." (OrderedSimplexList [SimplexListByDegree 0 [(Simplex [])], SimplexListByDegree 1 [(Simplex [1]), (Simplex [2]), (Simplex [3]), (Simplex [4])]]) (streamToOrderedSimplexList stream3))

streamToOrderedSimplexListThreeVerticesOneEdgeTest :: TestTree
streamToOrderedSimplexListThreeVerticesOneEdgeTest = testCase "Testing streamToOrderedSimplexList on stream with 3 vertices and one edge."
  (assertEqual "Should return object with lengths 0, 1, and 2 with three vertices inside the 1 key and 1 edge inside the 2 key." (OrderedSimplexList [SimplexListByDegree 0 [(Simplex [])], SimplexListByDegree 1 [(Simplex [1]), (Simplex [2]), (Simplex [3])], SimplexListByDegree 2 [(Simplex [1,2])]]) (streamToOrderedSimplexList stream1))

-- Test getBoundaryMap
simplexList1 :: SimplexListByDegree Int
simplexList2 :: SimplexListByDegree Int
simplexList1 = SimplexListByDegree 3 [(Simplex [1,2,3]), (Simplex [2,3,4])]
simplexList2 = SimplexListByDegree 2 [(Simplex [1,2]), (Simplex [2,3]), (Simplex [1,3]), (Simplex [2,4]), (Simplex [3,4])]
getBoundaryMapTest :: TestTree
getBoundaryMapTest = testCase "Testing getBoundaryMap."
  (assertEqual "Should return ..." (fromLists [[1,0],[1,1],[-1,0],[0,-1],[0,1]]) (getBoundaryMap simplexList2 simplexList1))

simplexList3 :: SimplexListByDegree Int
simplexList4 :: SimplexListByDegree Int
simplexList3 = SimplexListByDegree 1 [(Simplex [1]), (Simplex [2]), (Simplex [3]), (Simplex [4])]
simplexList4 = SimplexListByDegree 2 [(Simplex [1,2]), (Simplex [1,3]), (Simplex [1,4]), (Simplex [2,3]), (Simplex [3,4])]
getBoundaryMapTest2 :: TestTree
getBoundaryMapTest2 = testCase "Testing getBoundaryMap."
  (assertEqual "Should return ..." (fromLists [[-1,-1,-1,0,0],[1,0,0,-1,0],[0,1,0,1,-1],[0,0,1,0,1]]) (getBoundaryMap simplexList3 simplexList4))

simplexList5 :: SimplexListByDegree Int
simplexList6 :: SimplexListByDegree Int
simplexList5 = SimplexListByDegree 0 [(Simplex [])]
simplexList6 = SimplexListByDegree 1 [(Simplex [1]), (Simplex [2])]
getBoundaryMapTest3 :: TestTree
getBoundaryMapTest3 = testCase "Testing trivial case for getBoundaryMap."
  (assertEqual "Should return ..." (fromLists [[1, 1]]) (getBoundaryMap simplexList5 simplexList6))

-- Test getHomologyDimension
simplexList7 :: SimplexListByDegree Int
simplexList8 :: SimplexListByDegree Int
simplexList9 :: SimplexListByDegree Int
simplexList7 = SimplexListByDegree 0 [(Simplex [])]
simplexList8 = SimplexListByDegree 1 [(Simplex [1]), (Simplex [2]), (Simplex [3])]
simplexList9 = SimplexListByDegree 2 [(Simplex [1,2]), (Simplex [2,3]), (Simplex [1,3])]
map1 :: Matrix Double
map2 :: Matrix Double
map1 = getBoundaryMap simplexList7 simplexList8
map2 = getBoundaryMap simplexList8 simplexList9
getHomologyDimensionTest :: TestTree
getHomologyDimensionTest = testCase "Testing getHomologyDimension function."
  (assertEqual "Should return ..." (1) (getHomologyDimension map2 map1 0))

-- Test persistence
stream6 :: Stream Int
stream6 = addSimplex (addSimplex (addSimplex initializeStream (Simplex [1,2])) (Simplex [2,3])) (Simplex [1,3])
persistenceTest :: TestTree
persistenceTest = testCase "Testing persistence function."
  (assertEqual "Should return ..." (BettiVector [1,1]) (persistence stream6 1))

persistenceTest2 :: TestTree
persistenceTest2 = testCase "Testing persistence function on filled tetrahedron."
  (assertEqual "Should return ..." (BettiVector [1,0,0,0]) (persistence stream5 1))

persistenceTest3 :: TestTree
persistenceTest3 = testCase "Testing persistence function on four vertices."
  (assertEqual "Should return ..." (BettiVector [4]) (persistence stream3 1))

stream7 :: Stream Int
stream7 = subtractSimplex (addSimplex initializeStream (Simplex [1,2,3,4])) (Simplex [1,2,3,4])
persistenceTest4 :: TestTree
persistenceTest4 = testCase "Testing persistence function."
  (assertEqual "Should return ..." (BettiVector [1,0,1]) (persistence stream7 1))


--Hasquil registers
--Copyright Laurence Emms 2018

module Register (Quantum(..),
                 Classical(..)) where

import Data.Complex

data Quantum = QubitRegister Int |
               MetaQubitRegister String

instance Show Quantum where
    show (QubitRegister i) = show i
    show (MetaQubitRegister s) = s

data Classical a = Register Int |
                   Range Int Int |
                   RealConstant a |
                   ComplexConstant (Complex a) |
                   MetaRegister String

instance (Floating a, Show a, Ord a) => Show (Classical a) where
    show (Register i) = "[" ++ (show i) ++ "]"
    show (Range i j) = "[" ++ (show i) ++ "-" ++ (show j) ++ "]"
    show (RealConstant r) = show r
    show (ComplexConstant (p :+ q))
        | q >= 0 = (show p) ++ "+" ++ (show q) ++ "i"
        | otherwise = (show p) ++ (show q) ++ "i"
    show (MetaRegister s) = s


{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module Main where

import Prelude hiding (map)
import qualified Prelude as P

import Control.Monad
import Criterion.Config
import Criterion.Main
import Criterion.Monad
import Criterion.Plot
import Data.Array.Vector
import Nikola
import Statistics.Function (minMax)
import Statistics.KernelDensity (epanechnikovPDF)
import Statistics.Types
import System.Environment (getArgs)

main :: IO ()
main = do
    (cfg, _) <- parseArgs defaultConfig defaultOptions =<< System.Environment.getArgs
    withNewContext $ \_ -> do
    reify f >>= compileTopFun "f" >>= print
    samples <- withCompiledFunction f $ \f -> do
               replicateM 10000 $ do
               timeKernel f $ \f -> do
               call f [1..1024]
    withConfig cfg $
        plotAll [("map", (toU . P.map (realToFrac . fst)) samples)]
  where
    f :: Exp [Int] -> Exp [Int]
    f = map (\x -> x + 1)

plotAll :: [(String, Sample)] -> Criterion ()
plotAll descTimes = forM_ descTimes $ \(desc,times) -> do
  plotWith Timing $ \o -> plotTiming o desc times
  plotWith KernelDensity $ \o -> uncurry (plotKDE o desc extremes)
                                     (epanechnikovPDF 100 times)
  where
    extremes = case descTimes of
                 (_:_:_) -> toJust . minMax . concatU . P.map snd $ descTimes
                 _       -> Nothing
    toJust r@(lo :*: hi)
        | lo == infinity || hi == -infinity = Nothing
        | otherwise                         = Just r
        where infinity                      = 1/0


{-# LANGUAGE FlexibleContexts #-}
module CW (parseOpts, run) where

import           Control.Monad
import           Control.Monad.Error.Class
import           Data.Char
import           Data.Complex
import           Data.List
import qualified Filter as F
import           IQ
import           Options.Applicative
import           Pipes
import qualified Pipes.Prelude as P
import           Render
import           System.IO
import           ZeroStuff

data CWOptions = CWOptions
  { sampleRate :: Int
  , wpm :: Int
  , text :: String
  , outputFile :: String }

parseOpts :: Parser CWOptions
parseOpts = CWOptions
         <$> option auto
             ( long "sample-rate"
            <> short 's'
            <> metavar "SAMPLE_RATE"
            <> help "Sample rate of the produced IQ file" )
         <*> option auto
             ( long "wpm"
            <> short 'w'
            <> value 25
            <> showDefault
            <> metavar "WPM_RATE"
            <> help "Rate of keying in Words per Minute" )
         <*> strOption
             ( long "in-text"
            <> short 't'
            <> metavar "INPUT_TEXT"
            <> help "Input text message" )
         <*> strOption
             ( long "out-file"
            <> short 'o'
            <> metavar "OUTPUT_FILE"
            <> help "Output IQ file" )

data CW = Dah
        | Dit

cwMap :: MonadError String m => Char -> m [CW]
cwMap c
  | c == 'a' = pure [ Dah, Dit ]
  | c == 'b' = pure [ Dah, Dit, Dit, Dit ]
  | c == 'c' = pure [ Dah, Dit, Dah, Dit ]
  | c == 'd' = pure [ Dah, Dit, Dit ]
  | c == 'e' = pure [ Dit ]
  | c == 'f' = pure [ Dit, Dit, Dah, Dit ]
  | c == 'g' = pure [ Dah, Dah, Dit ]
  | c == 'h' = pure [ Dit, Dit, Dit, Dit ]
  | c == 'i' = pure [ Dit, Dit ]
  | c == 'j' = pure [ Dit, Dah, Dah, Dah ]
  | c == 'k' = pure [ Dah, Dit, Dah ]
  | c == 'l' = pure [ Dit, Dah, Dit, Dit ]
  | c == 'm' = pure [ Dah, Dah ]
  | c == 'n' = pure [ Dah, Dit ]
  | c == 'o' = pure [ Dah, Dah, Dah ]
  | c == 'p' = pure [ Dit, Dah, Dah, Dit ]
  | c == 'q' = pure [ Dah, Dah, Dit, Dah ]
  | c == 'r' = pure [ Dit, Dah, Dit ]
  | c == 's' = pure [ Dit, Dit, Dit ]
  | c == 't' = pure [ Dah ]
  | c == 'u' = pure [ Dit, Dit, Dah ]
  | c == 'v' = pure [ Dit, Dit, Dit, Dah ]
  | c == 'w' = pure [ Dit, Dah, Dah ]
  | c == 'x' = pure [ Dah, Dit, Dit, Dah ]
  | c == 'y' = pure [ Dah, Dit, Dah, Dah ]
  | c == 'z' = pure [ Dah, Dah, Dit, Dit ]
  | c == '1' = pure [ Dit, Dah, Dah, Dah, Dah ]
  | c == '2' = pure [ Dit, Dit, Dah, Dah, Dah ]
  | c == '3' = pure [ Dit, Dit, Dit, Dah, Dah ]
  | c == '4' = pure [ Dit, Dit, Dit, Dit, Dah ]
  | c == '5' = pure [ Dit, Dit, Dit, Dit, Dit ]
  | c == '6' = pure [ Dah, Dit, Dit, Dit, Dit ]
  | c == '7' = pure [ Dah, Dah, Dit, Dit, Dit ]
  | c == '8' = pure [ Dah, Dah, Dah, Dit, Dit ]
  | c == '9' = pure [ Dah, Dah, Dah, Dah, Dit ]
  | c == '0' = pure [ Dah, Dah, Dah, Dah, Dah ]
  | otherwise = throwError $ "Unknown CW character: '" ++ [c] ++ "'"

data KeyState = On | Off
data Keying = Keying
  { units :: Int
  , keyState :: KeyState }

cwKeyingMap :: CW -> Keying
cwKeyingMap Dah = Keying 3 On
cwKeyingMap Dit = Keying 1 On

keyLetter :: MonadError String m => Char -> m [Keying]
keyLetter l = intersperse (Keying 1 Off) <$> (map cwKeyingMap <$> cwMap l)

keyWord :: MonadError String m => String -> m [Keying]
keyWord word = intercalate  [Keying 3 Off] <$> sequence (keyLetter <$> word)

keyText :: String -> Either String [Keying]
keyText text = intercalate [Keying 7 Off] <$> sequence (keyWord <$> words text)

synthesizeKeying :: MonadIO m => Int -> Int -> Pipe Keying IQ m ()
synthesizeKeying sr wpm = forever $ do
  keying <- await

  let iq = case keying of
        Keying _ On -> 1.0 :+ 0.0
        Keying _ Off  -> 0.0 :+ 0.0
      n = round $ fromIntegral sr * (60.0 / (50.0 * fromIntegral wpm)) * fromIntegral (units keying)
  replicateM_ n $ yield iq

run :: CWOptions -> IO ()
run opts = case keyText $ toLower <$> text opts of
  Left err -> putStr err
  Right keyings -> do
    putStr "Synthesizing samples... "
    hFlush stdout

    let kernel = F.lowPass (sampleRate opts) $ F.FilterOptions 100 1000

    withBinaryFile (outputFile opts) WriteMode $ \f ->
      runEffect $ each keyings
           >-> synthesizeKeying (sampleRate opts) (wpm opts)
           >-> F.convolve kernel
           >-> cfileSink f

    putStrLn "Done."


{-# LANGUAGE ForeignFunctionInterface #-}
{-|
Module      : Grenade.Layers.Internal.Transpose
Description : Functions to quickly transpose a matrix
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Internal.Transpose (
  transpose4d
) where

import qualified Data.Vector.Storable        as U (unsafeFromForeignPtr0,
                                                   unsafeToForeignPtr0)

import           Foreign                     (mallocForeignPtrArray,
                                              withForeignPtr)
import           Foreign.C.Types             (CInt)
import           Foreign.Ptr                 (Ptr)
import           Numeric.LinearAlgebra       (Matrix, Vector, cmap, flatten)
import qualified Numeric.LinearAlgebra.Data  as D
import qualified Numeric.LinearAlgebra.Devel as U
import           System.IO.Unsafe            (unsafePerformIO)

import           Grenade.Types

-- | Transpose a 4d matrix, similarly to the tranpose function in numpy
transpose4d :: [Int]           -- ^ original dimensions
            -> Vector RealNum  -- ^ permutation vector
            -> Matrix RealNum  -- ^ input matrix
            -> Matrix RealNum  -- ^ transposed matrix
transpose4d dims@[n, c, h, w] permsV m
  = let outMatSize = n * c * h * w
        outW       = dims !! (round $ permsV D.! 3)
        vec        = flatten m
        permsV'    = cmap (U.fi . round) permsV  :: Vector CInt
        dimsV      = D.fromList $ map U.fi dims  :: Vector CInt
    in unsafePerformIO $ do
      outPtr        <- mallocForeignPtrArray outMatSize
      let (inPtr, _) = U.unsafeToForeignPtr0 vec
          (pPtr, _)  = U.unsafeToForeignPtr0 permsV'
          (dPtr, _)  = U.unsafeToForeignPtr0 dimsV

      withForeignPtr inPtr $ \inPtr' ->
        withForeignPtr pPtr $ \pPtr' ->
          withForeignPtr dPtr $ \dPtr' ->
            withForeignPtr outPtr $ \outPtr' ->
              transpose_4d inPtr' dPtr' pPtr' outPtr'

      let matVec = U.unsafeFromForeignPtr0 outPtr outMatSize
      return (U.matrixFromVector U.RowMajor (div outMatSize outW) outW matVec)
{-# INLINE transpose4d #-}

foreign import ccall unsafe
    transpose_4d
      :: Ptr RealNum -> Ptr CInt -> Ptr CInt -> Ptr RealNum -> IO ()


{-# LANGUAGE TypeFamilies #-}

module FPNLA.Operations.Parameters(
    -- * Elements
    Elt(..),
    -- * Strategies and contexts
    StratCtx(),

    -- * Result type
    -- | In BLAS it's common that operations in higher levels use operations in the lower levels, so, an operation in level three that by its signature manipulates matrices only, internally uses level two operations that manipulates vectors. In order to avoid the /show . read/ problem, the type of the vector (or any other internal data type) must appear in the signature of an operation.
    -- To solve the problem we use phantom types to pass the internally used types to the Haskell type system.
    ResM(),
    ResV(),
    ResS(),
    blasResultM,
    blasResultV,
    blasResultS,
    getResultDataM,
    getResultDataV,
    getResultDataS,

    -- * Miscellaneous
    TransType(..),
    UnitType(..),
    TriangType(..),
    unTransT,
    unUnitT,
    unTriangT,
    elemTrans_m,
    dimTrans_m,
    elemSymm,
    dimTriang,
    elemUnit_m,
    dimUnit_m,
    elemTransUnit_m,
    dimTransUnit_m,
    transTrans_m

) where

import FPNLA.Matrix(Matrix(..))
import Data.Complex (Complex, conjugate)
import Data.Tuple (swap)

-- | This class represents the elements that can be used in the BLAS operations.
-- The elements in BLAS are real or complex numbers, so we provide default instances for the Haskell 'Double', 'Float' and 'Complex' types.
class (Eq e, Floating e) => Elt e where
    -- | Returns the conjugate of a number. For real numbers it's the identity function and for complex numbers it's the common 'Complex.conjugate' function.
    getConjugate :: e -> e
    getConjugate = id

instance Elt Double
instance Elt Float
instance (RealFloat e) => Elt (Complex e) where
    getConjugate = conjugate

-- | This type family is used to represent the /context/ of an operation.
-- A particular implementation is a combination of an algorithm and a parallelism technique, and we call it a /strategy/. A particular strategy may need particular information to execute. For example, an operation that computes the matrix-matrix multiplication by splitting the matrices in blocks must require the size of the blocks.
-- With this context we allows to pass any additional information that the operation needs to execute as parameters, but maintaining a common signature.
-- The /s/ type parameter is the strategy so, there must exist a Haskell data type to represent a particular strategy.
type family StratCtx s :: *

-- | The 'ResM' data type is used as result of level three BLAS operations and returns a matrix /m/ of elements /e/ and contains the strategy /s/ and vector /v/ as phantom types.
data ResM s (v :: * -> *) m e = ResM { unResM :: m e } deriving (Show)
-- | The 'ResV' data type is used as result of level two BLAS operations and returns a vector /v/ of elements /e/ and contains the strategy /s/ as phantom types.
data ResV s v e = ResV { unResV :: v e } deriving (Show)
-- | The 'ResS' data type is used as result of level one BLAS operations and returns an scalar /e/ and contains the strategy /s/ as phantom types.
data ResS s e = ResS { unResS :: e } deriving (Show)

-- | Wrap a matrix into a 'ResM'.
blasResultM :: m e -> ResM s v m e
blasResultM = ResM
-- | Unwrap a matrix from a 'ResM'.
getResultDataM :: ResM s v m e -> m e
getResultDataM = unResM

-- | Wrap a vector into a 'ResV'.
blasResultV :: v e -> ResV s v e
blasResultV = ResV
-- | Unwrap a vector from a 'ResV'.
getResultDataV :: ResV s v e -> v e
getResultDataV = unResV

-- | Wrap a scalar into a 'ResS'.
blasResultS :: e -> ResS s e
blasResultS = ResS
-- | Unwrap a scalar from a 'ResS'.
getResultDataS :: ResS s e -> e
getResultDataS = unResS

-- | Indicates if a matrix must be considered as normal, transposed or transposed conjugated.
-- This is part of the common flags in the BLAS operation signatures and it's useful to work with a transposed matrix without really computing the transposed matrix.
data TransType m = Trans m | NoTrans m | ConjTrans m deriving (Eq, Show)
-- | Indicates if a matrix must be considered as unitary or not. An unitary matrix is a matrix that contains ones in the diagonal.
-- This is part of the common flags in the BLAS operation signatures.
data UnitType m = Unit m | NoUnit m deriving (Eq, Show)
-- | Indicates that a matrix is symmetric and with which triangular part of the matrix the operation is going to work ('Upper' or 'Lower').
-- The operation only will see the indicated part of the matrix and should not try to access the other part.
-- This is part of the common flags in the BLAS operation signatures.
data TriangType m = Lower m | Upper m deriving (Eq, Show)

-- | Given a data type flagged by a TransType, returns a pair containing the TransType constructor and the data type.
unTransT :: TransType a -> (b -> TransType b, a)
unTransT (Trans a) = (Trans, a)
unTransT (NoTrans a) = (NoTrans, a)
unTransT (ConjTrans a) = (ConjTrans, a)

-- | Given a data type flagged by a UnitType, returns a pair containing the UnitType constructor and the data type.
unUnitT :: UnitType a -> (b -> UnitType b, a)
unUnitT (Unit a) = (Unit, a)
unUnitT (NoUnit a) = (NoUnit, a)

-- | Given a data type flagged by a TriangType, returns a pair containing the TriangType constructor and the data type.
unTriangT :: TriangType a -> (b -> TriangType b, a)
unTriangT (Lower a) = (Lower, a)
unTriangT (Upper a) = (Upper, a)

-- | Given an /i,j/ position and a TransType flagged matrix, returns the element in that position without computing the transpose.
elemTrans_m :: (Elt e, Matrix m e) => Int -> Int -> TransType (m e) -> e
elemTrans_m i j (NoTrans m) = elem_m i j m
elemTrans_m i j (Trans m) = elem_m j i m
elemTrans_m i j (ConjTrans m) = getConjugate $ elem_m j i m

-- | Given a TransType flagged matrix, returns the dimension of the matrix without computing the transpose.
dimTrans_m :: (Matrix m e) => TransType (m e) -> (Int, Int)
dimTrans_m (NoTrans m) = dim_m m
dimTrans_m (ConjTrans m) = swap $ dim_m m
dimTrans_m (Trans m) = swap $ dim_m m

-- | Given an /i,j/ position and a TransType flagged matrix, returns the element in that position only accessing the part indicated by the TransType.
elemSymm :: (Matrix m e) => Int -> Int -> TriangType (m e) -> e
elemSymm i j (Upper m)
    | i > j = elem_m j i m
    | otherwise = elem_m i j m
elemSymm i j (Lower m)
    | i > j = elem_m i j m
    | otherwise = elem_m j i m

-- | Given a TransType flagged matrix, returns the dimension of the matrix.
dimTriang :: (Matrix m e) => TriangType (m e) -> (Int, Int)
dimTriang = dim_m . snd . unTriangT

-- | Given an /i,j/ position and a UnitType flagged matrix, returns the element in that position. If the matrix is flagged as Unit and /i == j/ (the element is in the diagonal) returns one.
elemUnit_m :: (Elt e, Matrix m e) => Int -> Int -> UnitType (m e) -> e
elemUnit_m i j (Unit m)
    | i == j = 1
    | otherwise = elem_m i j m
elemUnit_m i j (NoUnit m) = elem_m i j m

-- | Given a UnitType flagged matrix, returns the dimension of the matrix.
dimUnit_m :: (Matrix m e) => UnitType (m e) -> (Int, Int)
dimUnit_m (Unit m) = dim_m m
dimUnit_m (NoUnit m) = dim_m m

-- | Given an /i,j/ position and a TransType-UnitType flagged matrix, returns the element in that position without computing the transpose.
elemTransUnit_m :: (Elt e, Matrix m e) => Int -> Int -> TransType (UnitType (m e)) -> e
elemTransUnit_m i j (NoTrans pmA) = elemUnit_m i j pmA
elemTransUnit_m i j (Trans pmA) = elemUnit_m j i pmA
elemTransUnit_m i j (ConjTrans pmA) = getConjugate $ elemUnit_m j i pmA

-- | Given a TransType-UnitType flagged matrix, returns the dimension of the matrix.
dimTransUnit_m :: Matrix m e => TransType (UnitType (m e)) -> (Int, Int)
dimTransUnit_m = dimUnit_m . snd . unTransT

-- | Given a TransType flagged matrix, computes and returns its transpose.
transTrans_m :: (Elt e, Matrix m e) => TransType (m e) -> m e
transTrans_m (NoTrans m) = m
transTrans_m (ConjTrans m) = map_m getConjugate $ transpose_m m
transTrans_m (Trans m) = transpose_m m


{-# LANGUAGE AllowAmbiguousTypes    #-}
{-# LANGUAGE BangPatterns           #-}
{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE DeriveGeneric          #-}
{-# LANGUAGE DerivingVia            #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE NoStarIsType           #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE TypeInType             #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}
{-# OPTIONS_GHC -fno-warn-orphans   #-}

-- |
-- Module      : Numeric.Opto.Ref
-- Copyright   : (c) Justin Le 2019
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Abstract over different types for mutable references of values.
module Numeric.Opto.Ref (
    MR(..), ML(..)
  ) where

import           Control.Monad.Primitive
import           Data.Complex
import           Data.Mutable
import           GHC.TypeNats
import qualified Data.Vector.Generic                 as VG
import qualified Data.Vector.Generic.Mutable.Sized   as SMVG
import qualified Data.Vector.Generic.Sized           as SVG
import qualified Data.Vector.Storable                as VS
import qualified Numeric.LinearAlgebra               as HU
import qualified Numeric.LinearAlgebra.Devel         as HU
import qualified Numeric.LinearAlgebra.Static        as H
import qualified Numeric.LinearAlgebra.Static.Vector as H

instance VG.Vector v a => Mutable s (SVG.Vector v n a) where
    type Ref s (SVG.Vector v n a) = SVG.MVector (VG.Mutable v) n s a
    thawRef         = SVG.thaw
    freezeRef       = SVG.freeze
    copyRef         = SVG.copy
    moveRef         = SMVG.move
    cloneRef        = SMVG.clone
    unsafeThawRef   = SVG.unsafeThaw
    unsafeFreezeRef = SVG.unsafeFreeze

instance HU.Element a => Mutable s (HU.Matrix a) where
    type Ref s (HU.Matrix a) = HU.STMatrix s a
    thawRef x   = stToPrim $ HU.thawMatrix x
    freezeRef v = stToPrim $ HU.freezeMatrix v
    copyRef v x = stToPrim $ HU.setMatrix v 0 0 x
    moveRef         = undefined
    cloneRef        = undefined
    unsafeThawRef   = undefined
    unsafeFreezeRef = undefined

-- | Mutable ref for hmatrix's statically sized vector types, 'H.R' and
-- 'H.C'.
newtype MR s n a = MR { getMR :: SVG.MVector VS.MVector n s a }

instance KnownNat n => Mutable s (H.R n) where
    type Ref s (H.R n) = MR s n Double

    thawRef = fmap MR . thawRef . H.rVec
    freezeRef = fmap H.vecR . freezeRef . getMR
    copyRef (MR v) x = copyRef v (H.rVec x)
    moveRef         = undefined
    cloneRef        = undefined
    unsafeThawRef   = undefined
    unsafeFreezeRef = undefined

instance KnownNat n => Mutable s (H.C n) where
    type Ref s (H.C n) = MR s n (Complex Double)

    thawRef = fmap MR . thawRef . H.cVec
    freezeRef = fmap H.vecC . freezeRef . getMR
    copyRef (MR v) x = copyRef v (H.cVec x)
    moveRef         = undefined
    cloneRef        = undefined
    unsafeThawRef   = undefined
    unsafeFreezeRef = undefined

-- | Mutable ref for hmatrix's statically sized matrix types, 'H.L' and
-- 'H.M'.
newtype ML s n k a = ML { getML :: SVG.MVector VS.MVector (n * k) s a }

instance (KnownNat n, KnownNat k) => Mutable s (H.L n k) where
    type Ref s (H.L n k) = ML s n k Double

    thawRef = fmap ML . thawRef . H.lVec
    freezeRef = fmap H.vecL . freezeRef . getML
    copyRef (ML v) x = copyRef v (H.lVec x)
    moveRef         = undefined
    cloneRef        = undefined
    unsafeThawRef   = undefined
    unsafeFreezeRef = undefined

instance (KnownNat n, KnownNat k) => Mutable s (H.M n k) where
    type Ref s (H.M n k) = ML s n k (Complex Double)

    thawRef = fmap ML . thawRef . H.mVec
    freezeRef = fmap H.vecM . freezeRef . getML
    copyRef (ML v) x = copyRef v (H.mVec x)
    moveRef         = undefined
    cloneRef        = undefined
    unsafeThawRef   = undefined
    unsafeFreezeRef = undefined


module Main where

import           Control.Arrow                         (first, (&&&))

import           Criterion
import           Criterion.IO                          (readRecords)
import           Criterion.Main                        (defaultConfig,
                                                        defaultMainWith)
import           Criterion.Types
import           Statistics.Resampling.Bootstrap       (estPoint)

import qualified Data.ByteString.Lazy                  as BL
import           Data.Csv                              as Csv
import           Data.List                             (isPrefixOf, sortBy, zip4)
import           Data.Ord                              (comparing)
import qualified Data.Vector.Unboxed                   as V

import           Test.QuickCheck

import           Graphics.Gnuplot.Advanced
import qualified Graphics.Gnuplot.Frame                as Frame
import qualified Graphics.Gnuplot.Frame.Option         as Opt
import qualified Graphics.Gnuplot.Frame.OptionSet      as Opts
import qualified Graphics.Gnuplot.Graph.TwoDimensional as Graph2D
import qualified Graphics.Gnuplot.LineSpecification    as LineSpec
import qualified Graphics.Gnuplot.Plot.TwoDimensional  as Plot2D
import qualified Graphics.Gnuplot.Terminal.PNG         as PNG

import           Quickselect

range :: [Int]
range = [0, 1000 .. 50000]

testData :: [(Int, IO (V.Vector Int))]
testData = map f range
  where f n = (n, generate $ (V.fromList <$> vectorOf n arbitrary))

rawOut = "../data/kth-element-haskell.raw"
csvOut = "../data/kth-element-haskell.csv"
pngOut = "../data/kth-element-haskell.png"

config :: Config
config = defaultConfig { resamples = 3, rawDataFile = Just rawOut }

main :: IO ()
main = do
  defaultMainWith config [
    bgroup "main" [
        bgroup "find" $ map (run quickselect) testData
        , bgroup "cheat" $ map (run sortselect) testData
        ]
    ]
  (Right records) <- readRecords rawOut
  let records' = map (\(Analysed r) -> r) records
  let find' = extractOutcomes (filterReports "main/find/" records')
      cheat' = extractOutcomes (filterReports "main/cheat/" records')
      results = zip3 find' cheat' range
  BL.writeFile csvOut (encode results)
  plot (PNG.cons pngOut) resultPlot
  return ()
  where run f (n, xs) = env xs $ \ ~(xs') -> bench (show n) $ nf (f n) xs'
        filterReports name = filter ((name `isPrefixOf`) . reportName)
        extractOutcomes = map (estPoint . anMean . reportAnalysis)

resultPlot :: Frame.T (Graph2D.T Int Double)
resultPlot =
  let lineSpec title =
        Graph2D.lineSpec $
        LineSpec.title title $
        LineSpec.deflt
      lineSpec1 = lineSpec "find"
      lineSpec2 = lineSpec "cheat"
      frameOpts =
        Opts.xLabel "input size" $
        Opts.yLabel "execution time (seconds)" $
        Opts.title "Performance of Kth Element" $
        Opts.add (Opt.custom "datafile separator" "") ["\",\""] $
        Opts.deflt
      path1 = fmap lineSpec1 $ Plot2D.pathFromFile Graph2D.lines csvOut 3 1
      path2 = fmap lineSpec2 $ Plot2D.pathFromFile Graph2D.lines csvOut 3 2
  in Frame.cons frameOpts $ mconcat [path1, path2]