{-# 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 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


{-# 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, [])


{-# 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 DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE Strict                #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.LeakyTanh
Description : Hyperbolic tangent nonlinear layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.LeakyTanh
  ( LeakyTanh(..)
  , SpecLeakyTanh (..)
  , specLeakyTanh1D
  , specLeakyTanh2D
  , specLeakyTanh3D
  , specLeakyTanh
  , leakyTanhLayer
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Constraint                (Dict (..))
import           Data.Proxy
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build
import           Grenade.Types
import           Grenade.Utils.Conversion
import           Grenade.Utils.Vector

-- | A LeakyTanh layer. A layer which can act between any shape of the same dimension, performing a tanh function. The maximum value is given in Promille, i.e. 995 is 0.995, and is positive and negative. I.e. LeakyTanh(x) = min 0.995 (max (-0.995) x).
data LeakyTanh (v :: Nat) = LeakyTanh
  deriving (Generic,NFData,Show)

instance UpdateLayer (LeakyTanh maxVal) where
  type Gradient (LeakyTanh maxVal) = ()
  runUpdate _ l@LeakyTanh{} _ = l

instance RandomLayer (LeakyTanh maxVal) where
  createRandomWith _ _ = return LeakyTanh

instance Serialize (LeakyTanh maxVal) where
  put _ = return ()
  get = return LeakyTanh

instance (a ~ b, SingI a, KnownNat maxVal) => Layer (LeakyTanh maxVal) a b where
  type Tape (LeakyTanh maxVal) a b = S a
  runForwards _ (S1DV v) = (S1DV v, S1DV $ mapVector (tanhMax maxVal) v)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S2DV v) = (S2DV v, S2DV $ mapVector (tanhMax maxVal) v)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S1D a) = (S1D a, S1D $ LAS.dvmap (tanhMax maxVal) a)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S2D a) = (S2D a, S2D $ LAS.dmmap (tanhMax maxVal) a)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S3D a) = (S3D a, S3D $ LAS.dmmap (tanhMax maxVal) a)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S1DV v) (S1DV gs) = ((), S1DV $ zipWithVector (\t g -> tanhMax' maxVal t * g) v gs)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S2DV v) (S2DV gs) = ((), S2DV $ zipWithVector (\t g -> tanhMax' maxVal t * g) v gs)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S1D a) (S1D g) = ((), S1D $ LAS.dvmap (tanhMax' maxVal) a * g)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S2D a) (S2D g) = ((), S2D $ LAS.dmmap (tanhMax' maxVal) (tanh' a) * g)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S3D a) (S3D g) = ((), S3D $ LAS.dmmap (tanhMax' maxVal) (tanh' a) * g)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards l x y = runBackwards l x (toLayerShape x y)

tanhMax :: (Ord a, Fractional a) => a -> a -> a
tanhMax m v
  | v > m = leaky * (v - m) + m
  | v < -m = leaky * (v - m) - m
  | otherwise = v
  where
    leaky = 0.05

tanhMax' :: (Ord a, Fractional a, Floating a) => a -> a -> a
tanhMax' m v
  | v > m = leaky
  | v < -m = -leaky
  | otherwise = max 0.005 (tanh' v)
  where
    leaky = 0.05

tanh' :: (Floating a) => a -> a
tanh' t = 1 - s ^ (2 :: Int)  where s = tanh t

-------------------- DynamicNetwork instance --------------------

instance (KnownNat maxVal) => FromDynamicLayer (LeakyTanh maxVal) where
  fromDynamicLayer inp _ LeakyTanh = SpecNetLayer $ SpecLeakyTanh maxVal (tripleFromSomeShape inp)
    where maxVal = (/1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)

instance ToDynamicLayer SpecLeakyTanh where
  toDynamicLayer _ _ (SpecLeakyTanh maxVal (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     reifyNat (round $ 1000 * maxVal) $ \(_ :: (KnownNat maxVal) => Proxy maxVal) ->
     case (rows, cols, depth) of
         (_, 1, 1) -> case (unsafeCoerce (Dict :: Dict()) :: Dict ()) of
           Dict -> return $ SpecLayer (LeakyTanh :: LeakyTanh maxVal) (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> case (unsafeCoerce (Dict :: Dict()) :: Dict ()) of
           Dict -> return $ SpecLayer (LeakyTanh :: LeakyTanh maxVal) (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _         -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer (LeakyTanh :: LeakyTanh maxVal) (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh1D :: RealNum -> Integer -> SpecNet
specLeakyTanh1D maxVal i = specLeakyTanh3D maxVal (i, 1, 1)

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh2D :: RealNum -> (Integer, Integer) -> SpecNet
specLeakyTanh2D maxVal (i, j) = specLeakyTanh3D maxVal (i, j, 1)

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh3D :: RealNum -> (Integer, Integer, Integer) -> SpecNet
specLeakyTanh3D maxVal = SpecNetLayer . SpecLeakyTanh maxVal

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh :: RealNum -> (Integer, Integer, Integer) -> SpecNet
specLeakyTanh maxVal = SpecNetLayer . SpecLeakyTanh maxVal

-- | Add a LeakyTanh layer to your build.
leakyTanhLayer :: RealNum -> BuildM ()
leakyTanhLayer maxVal
  | maxVal <= 0 || maxVal > 1 = error "The maxVal parameter of tanhMaxLayer has to be in (0,1]"
  | otherwise = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecLeakyTanh maxVal

-------------------- GNum instances --------------------

instance GNum (LeakyTanh maxVal) where
  _ |* _ = LeakyTanh
  _ |+ _ = LeakyTanh
  sumG _ = LeakyTanh


{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE UndecidableInstances        #-}

module BlueRipple.Model.MRP where

import qualified BlueRipple.Data.Keyed         as K
import qualified BlueRipple.Data.DataFrames    as BR
import           BlueRipple.Data.CountFolds
import qualified Control.Foldl                 as FL
import           Control.Monad                  ( join )
import qualified Control.Monad.State           as State
import qualified Data.Array                    as A
import           Data.Function                  ( on )
import qualified Data.List                     as L
import qualified Data.Set                      as S
import qualified Data.IntMap                   as IM
import qualified Data.Map                      as M
import           Data.Maybe                     ( isJust
                                                , fromJust
                                                )
import qualified Data.Text                     as T
import qualified Data.Profunctor               as P
import qualified Data.Serialize                as SE
import           Data.Serialize.Text ()
import qualified Data.Vector                   as V
import qualified Data.Vector.Storable          as VS

import qualified Frames                        as F
import qualified Frames.Melt                   as F
import qualified Frames.InCore                 as FI
import qualified Data.Vinyl                    as V
import qualified Data.Vinyl.TypeLevel          as V

import qualified Control.MapReduce             as MR
import qualified Frames.Transform              as FT
import qualified Frames.Folds                  as FF
import qualified Frames.MapReduce              as FMR
import qualified Frames.Enumerations           as FE
--import qualified Frames.Misc                  as FU
import qualified Frames.Serialize              as FS

import qualified Knit.Report                   as K hiding (elements)
import qualified Polysemy.Error                as P
                                                ( mapError )
import qualified Polysemy                      as P
                                                ( raise )
import           Text.Pandoc.Error             as PE
import qualified Text.Blaze.Colonnade          as BC
import qualified Text.Show

import qualified Numeric.LinearAlgebra         as LA

import qualified Data.IndexedSet               as IS
import qualified Numeric.GLM.ProblemTypes      as GLM
import qualified Numeric.GLM.ModelTypes        as GLM
import qualified Numeric.GLM.FunctionFamily    as GLM
import           Numeric.GLM.MixedModel        as GLM
import qualified Numeric.GLM.Bootstrap         as GLM
import qualified Numeric.GLM.Report            as GLM
import qualified Numeric.GLM.Predict           as GLM
import qualified Numeric.GLM.Confidence        as GLM
import qualified Numeric.SparseDenseConversions
                                               as SD

import qualified Relude.Extra as Relude

import qualified Statistics.Types              as ST
import           GHC.Generics                   ( Generic, Rep )

{- Moved to BlueRipple.Data.CountFolds
-- map reduce folds for counting
type Count = "Count" F.:-> Int
type Successes = "Successes" F.:-> Int
type MeanWeight = "MeanWeight" F.:-> Double
type VarWeight = "VarWeight" F.:-> Double
type WeightedSuccesses = "WeightedSuccesses" F.:-> Double
type UnweightedSuccesses =  "UnweightedSuccesses" F.:-> Int

binomialFold
  :: (F.Record r -> Bool) -> FL.Fold (F.Record r) (F.Record '[Count, Successes])
binomialFold testRow =
  let successesF = FL.premap (\r -> if testRow r then 1 else 0) FL.sum
  in  (\n s -> n F.&: s F.&: V.RNil) <$> FL.length <*> successesF

countFold
  :: forall k r d
   . ( Ord (F.Record k)
     , FI.RecVec (k V.++ '[Count, Successes])
     , k F.⊆ r
     , d F.⊆ r
     )
  => (F.Record d -> Bool)
  -> FL.Fold (F.Record r) [F.FrameRec (k V.++ '[Count, Successes])]
countFold testData = MR.mapReduceFold
  MR.noUnpack
  (FMR.assignKeysAndData @k)
  (FMR.foldAndAddKey $ binomialFold testData)

type CountCols = '[Count, UnweightedSuccesses, WeightedSuccesses, MeanWeight, VarWeight]

zeroCount :: F.Record CountCols
zeroCount = 0 F.&: 0 F.&: 0 F.&: 1 F.&: 0 F.&: V.RNil

weightedBinomialFold
  :: (F.Record r -> Bool)
  -> (F.Record r -> Double)
  -> FL.Fold (F.Record r) (F.Record CountCols)
weightedBinomialFold testRow weightRow =
  let wSuccessesF =
        FL.premap (\r -> if testRow r then weightRow r else 0) FL.sum
      successesF  = FL.premap (\r -> if testRow r then 1 else 0) FL.sum
      meanWeightF = FL.premap weightRow FL.mean
      varWeightF  = FL.premap weightRow FL.variance
      f s ws mw = if mw < 1e-6 then realToFrac s else ws / mw -- if meanweight is 0 but
  in  (\n s ws mw vw -> n F.&: s F.&: f s ws mw F.&: mw F.&: vw F.&: V.RNil)
      <$> FL.length
      <*> successesF
      <*> wSuccessesF
      <*> meanWeightF
      <*> varWeightF

weightedCountFold
  :: forall k r d
   . (Ord (F.Record k)
     , FI.RecVec (k V.++ CountCols)
     , k F.⊆ r
     , k F.⊆ (k V.++ r)
     , d F.⊆ (k V.++ r)
     )
  => (F.Record r -> Bool) -- ^ count this row?
  -> (F.Record d -> Bool) -- ^ success ?
  -> (F.Record d -> Double) -- ^ weight
  -> FL.Fold (F.Record r) (F.FrameRec (k V.++ CountCols))
weightedCountFold {- filterData testData weightData-} = weightedCountFoldGeneral (F.rcast @k)
{-  FMR.concatFold $ FMR.mapReduceFold
    (MR.filterUnpack filterData)
    (FMR.assignKeysAndData @k)
    (FMR.foldAndAddKey $ weightedBinomialFold testData weightData)
-}

weightedCountFoldGeneral
  :: forall k r d
   . (Ord (F.Record k)
     , FI.RecVec (k V.++ CountCols)
     , k F.⊆ (k V.++ r)
     , d F.⊆ (k V.++ r)
     )
  => (F.Record r -> F.Record k)
  -> (F.Record r -> Bool) -- ^ include this row?
  -> (F.Record d -> Bool) -- ^ success ?
  -> (F.Record d -> Double) -- ^ weight
  -> FL.Fold (F.Record r) (F.FrameRec (k V.++ CountCols))
weightedCountFoldGeneral getKey filterData testData weightData =
  FL.prefilter filterData $ FMR.concatFold $ FMR.mapReduceFold
    (MR.Unpack $ \r ->  [getKey r `V.rappend` r])
    (FMR.assignKeysAndData @k @d)
    (FMR.foldAndAddKey $ weightedBinomialFold testData weightData)

-}

getFraction r =
  let n = F.rgetField @Count r
  in  if n == 0
        then 0
        else realToFrac (F.rgetField @Successes r) / realToFrac n

getFractionWeighted r =
  let n = F.rgetField @Count r
  in  if n == 0 then 0 else F.rgetField @WeightedSuccesses r / realToFrac n

data RecordColsProxy k = RecordColsProxy deriving (Show, Enum, Bounded, A.Ix, Eq, Ord)
type instance GLM.GroupKey (RecordColsProxy k) = F.Record k

--type instance GLM.GroupKey (Proxy k) = F.Record k

recordToGroupKey
  :: forall k r . (k F.⊆ r) => F.Record r -> RecordColsProxy k -> F.Record k
recordToGroupKey r _ = F.rcast @k r

glmErrorToPandocError :: GLM.GLMError -> PE.PandocError
glmErrorToPandocError x = PE.PandocSomeError $ show x

type GroupCols gs cs = gs V.++ cs
type CountedCols gs cs = GroupCols gs cs V.++ CountCols

---
newtype  SimplePredictor ps = SimplePredictor { unSimplePredictor :: F.Record ps }

instance (Show (F.Record ps)) => Show (SimplePredictor ps) where
  show (SimplePredictor x) = "SimplePredictor " ++ show @String x

instance (Eq (F.Record ps)) => Eq (SimplePredictor ps) where
  (SimplePredictor x) == (SimplePredictor y) = x == y

instance (Ord (F.Record ps)) => Ord (SimplePredictor ps) where
  compare (SimplePredictor x) (SimplePredictor y) = compare x y

instance (Ord (F.Record ps), K.FiniteSet (F.Record ps)) => Enum (SimplePredictor ps) where
  toEnum n =
    let im = IM.fromList $ zip [0..] $ SimplePredictor <$> S.toAscList K.elements
    in fromMaybe (error "Bad n given to toEnum :: Int -> SimplePredictor") $ IM.lookup n im
  fromEnum a =
    let m = M.fromList $ zip (SimplePredictor <$> S.toAscList K.elements) [0..]
    in fromMaybe (error "Bad SimplePredictor given to fromEnum :: SimplePredictor -> Int.  This means something is very bad.") $ M.lookup a m

instance (K.FiniteSet (F.Record ps)) => Bounded (SimplePredictor ps) where
  minBound = fromMaybe (error "Empty Set of Records used as predictor") $ viaNonEmpty head $ SimplePredictor <$> S.toList K.elements
  maxBound = fromMaybe (error "Empty Set of Records used as predictor") $ viaNonEmpty last $ SimplePredictor <$> S.toList K.elements

allSimplePredictors :: K.FiniteSet (F.Record ps) => [SimplePredictor ps]
allSimplePredictors = SimplePredictor <$> S.toList K.elements

type SimpleEffect ps = GLM.WithIntercept (SimplePredictor ps)

simplePredictor :: forall ps rs. (ps F.⊆ rs
                                     , Eq (F.Record ps)
                                     )
                    => F.Record rs -> SimplePredictor ps -> Double
simplePredictor r p = if F.rcast @ps r == unSimplePredictor p then 1 else 0

predMap :: forall cs. (K.FiniteSet (F.Record cs), Ord (F.Record cs), cs F.⊆ cs)
  => F.Record cs -> M.Map (SimplePredictor cs) Double
predMap r =  M.fromList $ Relude.fmapToSnd (simplePredictor r) allSimplePredictors

catPredMaps :: forall cs.  (K.FiniteSet (F.Record cs), Ord (F.Record cs), cs F.⊆ cs)
  => M.Map (F.Record cs) (M.Map (SimplePredictor cs) Double)
catPredMaps = M.fromList $ fmap (\k -> (unSimplePredictor k,predMap (unSimplePredictor k))) allSimplePredictors

data  LocationHolder c f a =  LocationHolder { locName :: T.Text
                                             , locKey :: Maybe (F.Rec f LocationCols)
                                             , catData :: M.Map (F.Rec f c) a
                                             } deriving (Generic)

deriving instance (V.RMap c
                  , V.ReifyConstraint Show F.ElField c
                  , V.RecordToList c
                  , Show a) => Show (LocationHolder c F.ElField a)

instance (SE.Serialize a
         , Ord (F.Rec FS.SElField c)
         , SE.GSerializePut
           (Rep (F.Rec FS.SElField c))
         , SE.GSerializeGet (Rep (F.Rec FS.SElField c))
         , (Generic (F.Rec FS.SElField c))
         ) => SE.Serialize (LocationHolder c FS.SElField a)

lhToS :: (Ord (F.Rec FS.SElField c)
         , V.RMap c
         )
      => LocationHolder c F.ElField a -> LocationHolder c FS.SElField a
lhToS (LocationHolder n lkM cdm) = LocationHolder n (fmap FS.toS lkM) (M.mapKeys FS.toS cdm)

lhFromS :: (Ord (F.Rec F.ElField c)
           , V.RMap c
         ) => LocationHolder c FS.SElField a -> LocationHolder c F.ElField a
lhFromS (LocationHolder n lkM cdm) = LocationHolder n (fmap FS.fromS lkM) (M.mapKeys FS.fromS cdm)

type LocationCols = '[BR.StateAbbreviation]
locKeyPretty :: F.Record LocationCols -> T.Text
locKeyPretty r =
  let stateAbbr = F.rgetField @BR.StateAbbreviation r
  in stateAbbr

--type ASER = '[BR.SimpleAgeC, BR.SexC, BR.CollegeGradC, BR.SimpleRaceC]
predictionsByLocation ::
  forall ps r rs. ( V.RMap ps
                  , V.ReifyConstraint Show V.ElField ps
                  , V.RecordToList ps
                  , Ord (F.Record ps)
                  , Ord (SimplePredictor ps)
                  , FI.RecVec (ps V.++ CountCols)
                  , F.ElemOf (ps V.++ CountCols) Count
                  , F.ElemOf (ps V.++ CountCols) MeanWeight
                  , F.ElemOf (ps V.++ CountCols) UnweightedSuccesses
                  , F.ElemOf (ps V.++ CountCols) VarWeight
                  , F.ElemOf (ps V.++ CountCols) WeightedSuccesses
                  , K.FiniteSet (F.Record ps)
                  , Show (F.Record (LocationCols V.++ ps V.++ CountCols))
                  , Show (F.Record ps)
                  , Enum (SimplePredictor ps)
                  , Bounded (SimplePredictor ps)
                  , (ps V.++ CountCols) F.⊆ (LocationCols V.++ ps V.++ CountCols)
                  , ps F.⊆ (ps V.++ CountCols)
                  , ps F.⊆ (LocationCols V.++ ps V.++ CountCols)
                  , F.ElemOf rs BR.StateAbbreviation
                  , K.KnitEffects r
               )
  => GLM.MinimizeDevianceVerbosity
  -> F.FrameRec rs
  -> FL.Fold (F.Record rs) (F.FrameRec (LocationCols V.++ ps V.++ CountCols))
  -> [SimpleEffect ps]
  -> M.Map (F.Record ps) (M.Map (SimplePredictor ps) Double)
  -> K.Sem r [LocationHolder ps V.ElField Double]
predictionsByLocation verbosity ccesFrame countFold predictors catPredMap =
  P.mapError glmErrorToPandocError  $ K.wrapPrefix "predictionsByLocation" $ do
    K.logLE K.Diagnostic "Inferring"
    (mm, rc, ebg, bu, vb, bs) <- inferMR @LocationCols @ps @ps
                                 verbosity
                                 countFold
                                 predictors
                                 simplePredictor
                                 ccesFrame

    let states = FL.fold FL.set $ fmap (F.rgetField @BR.StateAbbreviation) ccesFrame
        allStateKeys = (F.&: V.RNil) <$> FL.fold FL.list states
        predictLoc l = LocationHolder (locKeyPretty l) (Just l) catPredMap
        toPredict = [LocationHolder "National" Nothing catPredMap] <> fmap predictLoc allStateKeys
        predict (LocationHolder n lkM cpms) = P.mapError glmErrorToPandocError $ do
          let predictFrom catKey predMap =
                let groupKeyM = fmap (`V.rappend` catKey) lkM --lkM >>= \lk -> return $ lk `V.rappend` catKey
                    emptyAsNationalGKM = case groupKeyM of
                                           Nothing -> Nothing
                                           Just k -> k <$ GLM.categoryNumberFromKey rc k (RecordColsProxy @(LocationCols V.++ ps))
                in (fmap snd . GLM.runLogOnGLMException Nothing)
                   $ GLM.predictFromBetaB mm (`M.lookup` predMap) (const emptyAsNationalGKM) rc ebg bu vb
          cpreds <- M.traverseWithKey predictFrom cpms
          return $ LocationHolder n lkM cpreds
    traverse predict toPredict

---

-- TODO: Add bootstraps back in, make optional
-- Maybe set up way to serialize all this at this level??

-- ls: are keys for locations (groups with no matching fixed-effect)
-- cs: are keys for fixedEffects, one for each constructor in b
-- b: Fixed Effect type
-- g: group type
inferMR
  :: forall ls cs ks b g f rs effs
   . ( Foldable f
     , Functor f
     , K.KnitEffects effs
     , F.RDeleteAll ls (ls V.++ ks V.++ CountCols) ~ (ks V.++ CountCols)
     , (ls V.++ ks V.++ CountCols) ~ (ls V.++ (ks V.++ CountCols))
     , Ord (F.Record ls)
     , ls F.⊆ (ls V.++ ks V.++ CountCols)
     , (ks V.++ CountCols) F.⊆ (ls V.++ ks V.++ CountCols)
     , ks F.⊆ (ks V.++ CountCols)
     , (ls V.++ cs) F.⊆ (ls V.++ (ks V.++ CountCols))
     , Show (F.Record (ls V.++ ks V.++ CountCols))
     , FI.RecVec (ls V.++ (ks V.++ CountCols))
     , F.ElemOf (ks V.++ CountCols) Count
     , F.ElemOf (ks V.++ CountCols) MeanWeight
     , F.ElemOf (ks V.++ CountCols) VarWeight
     , F.ElemOf (ks V.++ CountCols) WeightedSuccesses
     , F.ElemOf (ks V.++ CountCols) UnweightedSuccesses
     , F.ElemOf (ls V.++ ks V.++ CountCols) Count
     , F.ElemOf (ls V.++ ks V.++ CountCols) MeanWeight
     , F.ElemOf (ls V.++ ks V.++ CountCols) VarWeight
     , F.ElemOf (ls V.++ ks V.++ CountCols) WeightedSuccesses
     , F.ElemOf (ls V.++ ks V.++ CountCols) UnweightedSuccesses
     , Show (F.Record (ls V.++ cs))
     , K.FiniteSet (F.Record ks)
     , Ord b
     , Show b
     , Enum b
     , Bounded b
     , Ord (F.Record (GroupCols ls cs))
     , g ~ RecordColsProxy (GroupCols ls cs)
     )
  => GLM.MinimizeDevianceVerbosity
  -> FL.Fold (F.Record rs) (F.FrameRec (ls V.++ ks V.++ CountCols))
  -> [GLM.WithIntercept b] -- fixed effects to fit
  -> (F.Record (ls V.++ ks V.++ CountCols) -> b -> Double)  -- how to get a fixed effect value from a record
  -> f (F.Record rs)
  -> K.Sem
       effs
       ( GLM.MixedModel b g
       , GLM.RowClassifier g
       , GLM.EffectsByGroup g b
       , GLM.BetaVec
       , LA.Vector Double
       , [(GLM.BetaVec, LA.Vector Double)]
       )
inferMR verbosity cf fixedEffectList getFixedEffect rows =
  P.mapError glmErrorToPandocError
  $ (fmap snd . GLM.runLogOnGLMException Nothing)
  $ K.wrapPrefix "inferMR"
  $ do
    let
      addZeroCountsF = FMR.concatFold $ FMR.mapReduceFold
                       FMR.noUnpack
                       (FMR.splitOnKeys @ls)
                       ( FMR.makeRecsWithKey id
                         $ FMR.ReduceFold
                         $ const
                         $ K.addDefaultRec @ks zeroCount
                       )
      counted = FL.fold FL.list $ FL.fold addZeroCountsF $ FL.fold cf rows
--    K.logLE K.Diagnostic $ T.intercalate "\n" $ fmap (T.pack . show) counted
    let
      vCounts = VS.fromList $ fmap (F.rgetField @Count) counted
      designEffect mw vw = 1 + (vw / (mw * mw))
      vWeights = VS.fromList $ fmap
        (\r ->
            let mw = F.rgetField @MeanWeight r
                vw = F.rgetField @VarWeight r
            in  case (mw < 1e-12, vw < 1e-12) of
              (True, _) -> 0
              (False, True) -> 1
              _ -> 1/sqrt (designEffect mw vw)
        )
        counted -- VS.replicate (VS.length vCounts) 1.0

      fixedEffects = GLM.FixedEffects $ IS.fromList fixedEffectList
      groups       = IS.fromList [RecordColsProxy]
      (observations, fixedEffectsModelMatrix, rcM) = FL.fold
        (lmePrepFrame getFractionWeighted
         fixedEffects
         groups
         getFixedEffect
         (recordToGroupKey @(GroupCols ls cs))
        )
        counted
--        K.logLE K.Diagnostic $ "vCounts=" <> (T.pack $ show vCounts)
--        K.logLE K.Diagnostic $ "vWeights=" <> (T.pack $ show vWeights)
--        K.logLE K.Diagnostic $ "mX=" <> (T.pack $ show fixedEffectsModelMatrix)
--        K.logLE K.Diagnostic $ "vY=" <> (T.pack $ show observations)
    let regressionModelSpec = GLM.RegressionModelSpec
          fixedEffects
          fixedEffectsModelMatrix
          observations
    rowClassifier <- case rcM of
      Left  msg -> K.knitError msg
      Right x   -> return x
--        K.logLE K.Diagnostic $ "rc=" <> (T.pack $ show rowClassifier)
    let effectsByGroup =
          M.fromList [(RecordColsProxy, IS.fromList [GLM.Intercept])]
    fitSpecByGroup <- GLM.fitSpecByGroup @b @g fixedEffects
                      effectsByGroup
                      rowClassifier
    let lmmControls = GLM.LMMControls GLM.LMM_BOBYQA 1e-6 Nothing
        lmmSpec     = GLM.LinearMixedModelSpec
                      (GLM.MixedModelSpec regressionModelSpec fitSpecByGroup)
                      lmmControls
        cc = GLM.PIRLSConvergenceCriterion GLM.PCT_Deviance 1e-6 20
        glmmControls = GLM.GLMMControls GLM.UseCanonical 10 cc
        glmmSpec = GLM.GeneralizedLinearMixedModelSpec
                   lmmSpec
                   vWeights
                   (GLM.Binomial vCounts)
                   glmmControls
        mixedModel = GLM.GeneralizedLinearMixedModel glmmSpec
    randomEffectsModelMatrix <- GLM.makeZ fixedEffectsModelMatrix
                                fitSpecByGroup
                                rowClassifier
--        K.logLE K.Diagnostic $ "smZ=" <> (T.pack $ show randomEffectsModelMatrix)
    let randomEffectCalc = GLM.RandomEffectCalculated
          randomEffectsModelMatrix
          (GLM.makeLambda fitSpecByGroup)
        th0         = GLM.setCovarianceVector fitSpecByGroup 1 0
        -- mdVerbosity = MDVNone
    GLM.checkProblem mixedModel randomEffectCalc
    K.logLE K.Info "Fitting data..."
    ((th, pd, sigma2, beta, mzvu, mzvb, cs), vMuSol, cf) <- GLM.minimizeDeviance
                                                            verbosity
                                                            ML
                                                            mixedModel
                                                            randomEffectCalc
                                                            th0
    vb <- K.knitMaybe "b-vector (random effects coefficients) is zero in MRP.inferMR" $ GLM.useMaybeZeroVec Nothing Just mzvb
    GLM.report mixedModel
      randomEffectsModelMatrix
      beta
      (SD.toSparseVector vb)
    let fes = GLM.fixedEffectStatistics mixedModel sigma2 cs beta
    K.logLE K.Diagnostic $ "FixedEffectStatistics: " <> show fes
    epg <- GLM.effectParametersByGroup @g @b rowClassifier effectsByGroup vb
--        K.logLE K.Diagnostic
--          $  "EffectParametersByGroup: "
--          <> (T.pack $ show epg)
    gec <- GLM.effectCovariancesByGroup effectsByGroup mixedModel sigma2 th
    K.logLE K.Diagnostic
      $  "EffectCovariancesByGroup: "
      <> show gec
    rebl <- GLM.randomEffectsByLabel epg rowClassifier
    K.logLE K.Diagnostic
      $  "Random Effects:\n"
      <> GLM.printRandomEffectsByLabel rebl
    smCondVar <- GLM.conditionalCovariances mixedModel
                 cf
                 randomEffectCalc
                 th
                 beta
                 mzvu
    let bootstraps                           = []
    let GLM.FixedEffectStatistics _ mBetaCov = fes
    let f r = do
                let obs = getFractionWeighted r
                predictCVCI <- GLM.predictWithCI
                  mixedModel
                  (Just . getFixedEffect r)
                  (Just . recordToGroupKey @(GroupCols ls cs) r)
                  rowClassifier
                  effectsByGroup
                  beta
                  vb
                  (ST.mkCL 0.95)
                  (GLM.NaiveCondVarCI mBetaCov smCondVar)
                return (r, obs, predictCVCI)
    fitted <- traverse f (FL.fold FL.list counted)
    K.logLE K.Diagnostic
      $  "Fitted:\n"
      <> (T.intercalate "\n" $ fmap show fitted)
    fixedEffectTable <- GLM.printFixedEffects fes
    K.logLE K.Diagnostic $ "FixedEffects:\n" <> fixedEffectTable
    let GLM.FixedEffectStatistics fep _ = fes
    return
      (mixedModel, rowClassifier, effectsByGroup, beta, vb, bootstraps) -- fes, epg, rowClassifier, bootstraps)

lmePrepFrame
  :: forall p g rs
   . (GLM.PredictorC p, GLM.GroupC g)
  => (F.Record rs -> Double) -- ^ observations
  -> GLM.FixedEffects p
  -> IS.IndexedSet g
  -> (F.Record rs -> p -> Double) -- ^ predictors
  -> (F.Record rs -> g -> GLM.GroupKey g)  -- ^ classifiers
  -> FL.Fold
       (F.Record rs)
       ( LA.Vector Double
       , LA.Matrix Double
       , Either T.Text (GLM.RowClassifier g)
       ) -- ^ (X,y,(row-classifier, size of class))
lmePrepFrame observationF fe groupIndices getPredictorF classifierLabelF
  = let
      makeInfoVector
        :: M.Map g (M.Map (GLM.GroupKey g) Int)
        -> M.Map g (GLM.GroupKey g)
        -> Either T.Text (V.Vector (GLM.ItemInfo g))
      makeInfoVector indexMaps groupKeys =
        let
          g (grp, groupKey) =
            GLM.ItemInfo
              <$> (   maybe
                      (Left $ "Failed on " <> show (grp, groupKey))
                      Right
                  $   M.lookup grp indexMaps
                  >>= M.lookup groupKey
                  )
              <*> pure groupKey
        in  fmap V.fromList $ traverse g $ M.toList groupKeys
      makeRowClassifier
        :: Traversable f
        => M.Map g (M.Map (GLM.GroupKey g) Int)
        -> f (M.Map g (GLM.GroupKey g))
        -> Either T.Text (GLM.RowClassifier g)
      makeRowClassifier indexMaps labels = do
        let sizes = fmap M.size indexMaps
        indexed <- traverse (makeInfoVector indexMaps) labels
        return $ GLM.RowClassifier groupIndices
                                   sizes
                                   (V.fromList $ FL.fold FL.list indexed)
                                   indexMaps
      getPredictorF' _   GLM.Intercept     = 1
      getPredictorF' row (GLM.Predictor x) = getPredictorF row x
      predictorF row = LA.fromList $ case fe of
        GLM.FixedEffects indexedFixedEffects ->
          getPredictorF' row <$> IS.members indexedFixedEffects
        GLM.InterceptOnly -> [1]
      getClassifierLabels :: F.Record rs -> M.Map g (GLM.GroupKey g)
      getClassifierLabels r =
        M.fromList $ Relude.fmapToSnd (classifierLabelF r) $ IS.members groupIndices
      foldObs   = LA.fromList <$> FL.premap observationF FL.list
      foldPred  = LA.fromRows <$> FL.premap predictorF FL.list
      foldClass = FL.premap getClassifierLabels FL.list
      g (vY, mX, ls) =
        ( vY
        , mX
        , makeRowClassifier
          (snd $ State.execState (addAll ls) (M.empty, M.empty))
          ls
        )
    in
      g <$> ((,,) <$> foldObs <*> foldPred <*> foldClass)

addOne
  :: GLM.GroupC g
  => (g, GLM.GroupKey g)
  -> State.State (M.Map g Int, M.Map g (M.Map (GLM.GroupKey g) Int)) ()
addOne (grp, label) = do
  (nextIndexMap, groupIndexMaps) <- State.get
  let groupIndexMap = fromMaybe M.empty $ M.lookup grp groupIndexMaps
  case M.lookup label groupIndexMap of
    Nothing -> do
      let index         = fromMaybe 0 $ M.lookup grp nextIndexMap
          nextIndexMap' = M.insert grp (index + 1) nextIndexMap
          groupIndexMaps' =
            M.insert grp (M.insert label index groupIndexMap) groupIndexMaps
      State.put (nextIndexMap', groupIndexMaps')
      return ()
    _ -> return ()

addMany
  :: (GLM.GroupC g, Traversable h)
  => h (g, GLM.GroupKey g)
  -> State.State (M.Map g Int, M.Map g (M.Map (GLM.GroupKey g) Int)) ()
addMany = traverse_ addOne

addAll
  :: GLM.GroupC g
  => [M.Map g (GLM.GroupKey g)]
  -> State.State (M.Map g Int, M.Map g (M.Map (GLM.GroupKey g) Int)) ()
addAll = traverse_ (addMany . M.toList)

-- useful data folds
weightedSumF :: (Real w, Real v, Fractional y) => FL.Fold (w, v) y
weightedSumF =
  (/)
    <$> FL.premap (\(w, v) -> realToFrac w * realToFrac v) FL.sum
    <*> P.dimap fst realToFrac FL.sum

weightedSumRecF
  :: forall w v y rs
   . ( V.KnownField w
     , V.KnownField v
     , Fractional y
     , F.ElemOf rs w
     , F.ElemOf rs v
     , Real (V.Snd w)
     , Real (V.Snd v)
     )
  => FL.Fold (F.Record rs) y
weightedSumRecF =
  FL.premap (\r -> (F.rgetField @w r, F.rgetField @v r)) weightedSumF

sumProdIfRecF
  :: forall w v t y rs
   . ( V.KnownField w
     , V.KnownField v
     , V.KnownField t
     , Fractional y
     , F.ElemOf rs w
     , F.ElemOf rs v
     , F.ElemOf rs t
     , Real (V.Snd w)
     , Real (V.Snd v)
     )
  => (V.Snd t -> Bool)
  -> FL.Fold (F.Record rs) y
sumProdIfRecF test = FL.prefilter (test . F.rgetField @t) $ FL.premap
  (\r -> realToFrac (F.rgetField @w r) * realToFrac (F.rgetField @v r))
  FL.sum