{-# LANGUAGE TemplateHaskell, GeneralizedNewtypeDeriving, Rank2Types, NoMonomorphismRestriction #-}
import Test.QuickCheck
import Test.QuickCheck.Gen.Unsafe
import Data.List
import Data.Int
import Data.Word
import Data.Version
import System.Exit
import Data.Complex
import Text.ParserCombinators.ReadP (readP_to_S)

newtype Path a = Path [a] deriving (Show, Functor)

instance Arbitrary a => Arbitrary (Path a) where
  arbitrary = do
    x <- arbitrary
    fmap Path (pathFrom 100 x)
    where
      pathFrom n x =
        fmap (x:) $
        case shrink x of
          [] -> return []
          _ | n == 0 -> return []
          ys -> oneof [pathFrom (n-1) y | y <- ys]

  shrink (Path xs) = map Path [ ys | ys <- inits xs, length ys > 0 && length ys < length xs ]

path :: (a -> Bool) -> Path a -> Bool
path p (Path xs) = all p xs

somePath :: (a -> Bool) -> Path a -> Property
somePath p = expectFailure . withMaxSuccess 1000 . path (not . p)

newtype Extremal a = Extremal { getExtremal :: a } deriving (Show, Eq, Ord, Num, Enum, Real, Integral)

instance (Arbitrary a, Bounded a) => Arbitrary (Extremal a) where
  arbitrary =
    fmap Extremal $
    frequency
      [(1, return minBound),
       (1, return maxBound),
       (8, arbitrary)]
  shrink (Extremal x) = map Extremal (shrink x)

smallProp :: Integral a => Path a -> Bool
smallProp = path (\x -> (x >= -100 || -100 `asTypeOf` x >= 0) && x <= 100)

largeProp :: Integral a => Path a -> Property
largeProp = somePath (\x -> x < -1000000 || x > 1000000)

prop_int :: Path Int -> Bool
prop_int = smallProp

prop_int32 :: Path Int32 -> Property
prop_int32 = largeProp

prop_word :: Path Word -> Property
prop_word = largeProp

prop_word32 :: Path Word32 -> Property
prop_word32 = largeProp

prop_integer :: Path Integer -> Bool
prop_integer = smallProp

prop_small :: Path (Small Int) -> Bool
prop_small = smallProp

prop_large :: Path (Large Int) -> Property
prop_large = largeProp

prop_smallWord :: Path (Small Word) -> Bool
prop_smallWord = smallProp

prop_largeWord :: Path (Large Word) -> Property
prop_largeWord = largeProp

data Choice a b = Choice a b deriving Show
instance (Arbitrary a, Arbitrary b) => Arbitrary (Choice a b) where
  arbitrary = do
    Capture eval <- capture
    return (Choice (eval arbitrary) (eval arbitrary))

idemProp :: (Eq a, Arbitrary a, Arbitrary b) => (b -> a) -> Choice a b -> Bool
idemProp f (Choice x y) = x == f y

prop_fixed_length :: Arbitrary a => Path (Fixed a) -> Bool
prop_fixed_length (Path xs) = length xs == 1

prop_fixed_idem = idemProp getFixed
prop_blind_idem = idemProp getBlind

prop_ordered_list = path (\(Ordered xs) -> sort xs == xs)
prop_nonempty_list = path (\(NonEmpty xs) -> not (null xs))

pathInt, somePathInt ::
  (Arbitrary (f (Extremal Int)), Show (f (Extremal Int)),
   Arbitrary (f Integer), Show (f Integer),
   Arbitrary (f (Extremal Int8)), Show (f (Extremal Int8)),
   Arbitrary (f (Extremal Int16)), Show (f (Extremal Int16)),
   Arbitrary (f (Extremal Int32)), Show (f (Extremal Int32)),
   Arbitrary (f (Extremal Int64)), Show (f (Extremal Int64)),
   Arbitrary (f (Extremal Word)), Show (f (Extremal Word)),
   Arbitrary (f (Extremal Word8)), Show (f (Extremal Word8)),
   Arbitrary (f (Extremal Word16)), Show (f (Extremal Word16)),
   Arbitrary (f (Extremal Word32)), Show (f (Extremal Word32)),
   Arbitrary (f (Extremal Word64)), Show (f (Extremal Word64))) =>
  Bool -> (forall a. f a -> a) -> (forall a. Integral a => a -> Bool) -> Property
pathInt word f p =
  conjoin
    [counterexample "Int" (path ((p :: Int -> Bool) . getExtremal . f)),
     counterexample "Integer" (path ((p :: Integer -> Bool) . f)),
     counterexample "Int8" (path ((p :: Int8 -> Bool) . getExtremal . f)),
     counterexample "Int16" (path ((p :: Int16 -> Bool) . getExtremal . f)),
     counterexample "Int32" (path ((p :: Int32 -> Bool) . getExtremal . f)),
     counterexample "Int64" (path ((p :: Int64 -> Bool) . getExtremal . f)),
     counterexample "Word" (not word .||. path ((p :: Word -> Bool) . getExtremal . f)),
     counterexample "Word8" (not word .||. path ((p :: Word8 -> Bool) . getExtremal . f)),
     counterexample "Word16" (not word .||. path ((p :: Word16 -> Bool) . getExtremal . f)),
     counterexample "Word32" (not word .||. path ((p :: Word32 -> Bool) . getExtremal . f)),
     counterexample "Word64" (not word .||. path ((p :: Word64 -> Bool) . getExtremal . f))]
somePathInt word f p = expectFailure (pathInt word f (not . p))

prop_positive = pathInt True getPositive (> 0)
prop_positive_bound = somePathInt True getPositive (== 1)

prop_nonzero = pathInt True getNonZero (/= 0)
prop_nonzero_bound_1 = somePathInt True getNonZero (== 1)
prop_nonzero_bound_2 = somePathInt True getNonZero (== -1)

prop_nonnegative = pathInt True getNonNegative (>= 0)
prop_nonnegative_bound = somePathInt True getNonNegative (== 0)

prop_negative = pathInt False getNegative (< 0)
prop_negative_bound = somePathInt False getNegative (== -1)

prop_nonpositive = pathInt True getNonPositive (<= 0)
prop_nonpositive_bound = somePathInt True getNonPositive (== 0)

reachesBound :: (Bounded a, Integral a, Arbitrary a) =>
  a -> Property
reachesBound x = withMaxSuccess 1000 (expectFailure (x < 3 * (maxBound `div` 4)))

prop_reachesBound_Int8 = reachesBound :: Int8 -> Property
prop_reachesBound_Int16 = reachesBound :: Int16 -> Property
prop_reachesBound_Int32 = reachesBound :: Int32 -> Property
prop_reachesBound_Int64 = reachesBound :: Int64 -> Property
prop_reachesBound_Word8 = reachesBound :: Word8 -> Property
prop_reachesBound_Word16 = reachesBound :: Word16 -> Property
prop_reachesBound_Word32 = reachesBound :: Word32 -> Property
prop_reachesBound_Word64 = reachesBound :: Word64 -> Property

-- Shrinking should not loop.
noShrinkingLoop :: (Eq a, Arbitrary a) => Path a -> Bool
noShrinkingLoop (Path (x:xs)) = x `notElem` xs

prop_no_shrinking_loop_Unit = noShrinkingLoop :: Path () -> Bool
prop_no_shrinking_loop_Bool = noShrinkingLoop :: Path Bool -> Bool
prop_no_shrinking_loop_Ordering = noShrinkingLoop :: Path Ordering -> Bool
prop_no_shrinking_loop_Maybe = noShrinkingLoop :: Path (Maybe Int) -> Bool
prop_no_shrinking_loop_Either = noShrinkingLoop :: Path (Either Int Int) -> Bool
prop_no_shrinking_loop_List = noShrinkingLoop :: Path [Int] -> Bool
prop_no_shrinking_loop_Ratio = noShrinkingLoop :: Path Rational -> Bool
prop_no_shrinking_loop_Complex = noShrinkingLoop :: Path (Complex Double) -> Bool
prop_no_shrinking_loop_Fixed = noShrinkingLoop :: Path (Fixed Int) -> Bool
prop_no_shrinking_loop_Pair = noShrinkingLoop :: Path (Int, Int) -> Bool
prop_no_shrinking_loop_Triple = noShrinkingLoop :: Path (Int, Int, Int) -> Bool
prop_no_shrinking_loop_Integer = noShrinkingLoop :: Path Integer -> Bool
prop_no_shrinking_loop_Int = noShrinkingLoop :: Path Int -> Bool
prop_no_shrinking_loop_Int8 = noShrinkingLoop :: Path Int8 -> Bool
prop_no_shrinking_loop_Int16 = noShrinkingLoop :: Path Int16 -> Bool
prop_no_shrinking_loop_Int32 = noShrinkingLoop :: Path Int32 -> Bool
prop_no_shrinking_loop_Int64 = noShrinkingLoop :: Path Int64 -> Bool
prop_no_shrinking_loop_Word = noShrinkingLoop :: Path Word -> Bool
prop_no_shrinking_loop_Word8 = noShrinkingLoop :: Path Word8 -> Bool
prop_no_shrinking_loop_Word16 = noShrinkingLoop :: Path Word16 -> Bool
prop_no_shrinking_loop_Word32 = noShrinkingLoop :: Path Word32 -> Bool
prop_no_shrinking_loop_Word64 = noShrinkingLoop :: Path Word64 -> Bool
prop_no_shrinking_loop_Char = noShrinkingLoop :: Path Char -> Bool
prop_no_shrinking_loop_Float = noShrinkingLoop :: Path Float -> Bool
prop_no_shrinking_loop_Double = noShrinkingLoop :: Path Double -> Bool
prop_no_shrinking_loop_Version = noShrinkingLoop :: Path Version -> Bool
prop_no_shrinking_loop_ExitCode = noShrinkingLoop :: Path ExitCode -> Bool

-- Check that shrinking a Double always produces a shrinking candidate.
prop_shrink_candidate_double :: Property
prop_shrink_candidate_double =
  forAllShrink gen shrink $ \x ->
    x > 0 ==>
    not (null (shrink x))
  where
    gen :: Gen Double
    gen = oneof [arbitrary, fmap fromInteger arbitrary]

-- Bad shrink: infinite list
--
-- remove unexpectedFailure in prop_B1, shrinking should not loop forever.
data B1 = B1 Int deriving (Eq, Show)

instance Arbitrary B1 where
    arbitrary = fmap B1 arbitrary
    shrink x = x : shrink x

prop_B1 :: B1 -> Property
prop_B1 (B1 n) = expectFailure $ n === n + 1

return []
main = do True <- $forAllProperties (quickCheckWithResult stdArgs { maxShrinks = 10000 }); return ()


module Parser where
import Text.ParserCombinators.Parsec
import System.Environment
import Control.Monad
import Utilities.Number
import Numeric
import Data.Ratio
import Data.Complex
import Data.Array
import Utilities.Types

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

spaces1 :: Parser ()
spaces1 = skipMany1 space

escapedChars :: Parser Char
escapedChars = do
    char '\\' -- a backslash
    x <- oneOf "\\\"nrt" -- return either backslash or doublequote
    return $ case x of
        '\\' -> x
        '"' -> x
        'n' -> '\n'
        'r' -> '\r'
        't' -> '\t'

parseCharacter :: Parser LispVal
parseCharacter = do
    try $ string "#\\"
    value <- try (string "newline" <|> string "space")
        <|> do { x <- anyChar; notFollowedBy alphaNum; return [x]}
    return $ Character $ case value of
        "space" -> ' '
        "newline" -> '\n'
        _ -> head value

parseBool :: Parser LispVal
parseBool = do
    char '#'
    (char 't' >> return (Bool True)) <|> (char 'f' >> return (Bool False))

parseString :: Parser LispVal
parseString = do
    char '"'
    x <- many $ escapedChars <|> noneOf "\"\\"
    char '"'
    return $ String x

parseAtom :: Parser LispVal
parseAtom = do
    first <- letter <|> symbol
    rest <- many (letter <|> digit <|> symbol)
    let atom = first:rest
    return $ case atom of
        "#t" -> Bool True
        "#f" -> Bool False
        v -> Atom v

parseDecimal1 :: Parser LispVal
parseDecimal1 = Number . read <$> many1 digit

parseDecimal2 :: Parser LispVal
parseDecimal2 = do
    try $ string "#d"
    x <- many1 digit
    (return . Number . read) x

parseHex :: Parser LispVal
parseHex = do
    try $ string "#x"
    x <- many1 hexDigit
    return $ Number (hex2dig x)

parseOct :: Parser LispVal
parseOct = do
    try $ string "#o"
    x <- many1 octDigit
    return $ Number (oct2dig x)

parseBin :: Parser LispVal
parseBin = do
    try $ string "#o"
    x <- many1 (oneOf "10")
    return $ Number (bin2dig x)

parseNumber :: Parser LispVal
parseNumber = parseDecimal1 <|> parseDecimal2 <|> parseHex <|> parseOct <|> parseBin

parseFloat :: Parser LispVal
parseFloat = do
    x <- many1 digit
    char '.'
    y <- many1 digit
    return $ Float (fst . head $ readFloat (x ++ "." ++ y))

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 <|> parseDecimal1 <|> parseDecimal2
    char '+'
    y <- try parseFloat <|> parseDecimal1 <|> parseDecimal2
    char 'i'
    return $ Complex (toDouble x :+ toDouble y)

-- parseList :: Parser LispVal
-- parseList = liftM List $ sepBy parseExpr spaces1

-- parseDottedList :: Parser LispVal
-- parseDottedList = do
--     head <- endBy parseExpr spaces1
--     tail <- char '.' >> spaces1 >> parseExpr
--     return $ DottedList head tail

parseList :: Parser LispVal
parseList = do
    char '(' >> spaces
    head <- parseExpr `sepEndBy` spaces1
    do
        char '.' >> spaces1
        tail <- parseExpr
        spaces >> char ')'
        return $ DottedList head tail
        <|> (spaces >> char ')' >> (return $ List head))

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

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

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

parseVector :: Parser LispVal
parseVector = do
    arrayValues <- sepBy parseExpr spaces1
    return $ Vector (listArray (0, length arrayValues - 1) arrayValues)

parseExpr :: Parser LispVal
parseExpr = parseAtom
    <|> parseString
    <|> try parseComplex
    <|> try parseFloat
    <|> try parseRatio
    <|> try parseNumber
    <|> parseQuoted
    <|> try (do
        string "#("
        x <- parseVector
        char ')'
        return x)
    <|> parseList
    <|> try parseBool
    <|> try parseCharacter
    <|> parseQuasiQuoted
    <|> parseUnQuote


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

module Main where

import           Control.Monad
import           Control.Monad.Error
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 -> ThrowsError LispVal
readExpr input = case parse parseExpr "lisp" input of
  Left err -> throwError $ Parser err
  Right val -> return val

-- Beginnings of an evaluator: Primitives
-- main :: IO ()
-- main = getArgs >>= print . eval . readExpr . head
main :: IO ()
main = do
  args <- getArgs
  evaled <- return $ liftM show $ readExpr (head args) >>= eval
  putStrLn $ extractValue $ trapError evaled

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

-- Adding basic primitives
eval (List (Atom func : args)) = mapM eval args >>= apply func

apply :: String -> [LispVal] -> ThrowsError LispVal
apply func args = maybe (throwError $ NotFunction "No such function: " func)
                  ($ args)
                  (lookup func primitives)

primitives :: [(String, [LispVal] -> ThrowsError 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),
              ("symbol->string", unaryOp symbol2String),
              ("string->symbol", unaryOp string2Symbol)]

numericBinop :: (Integer -> Integer -> Integer) -> [LispVal] -> ThrowsError LispVal
numericBinop _ [] = throwError $ NumArgs 2 []
numericBinop _ singleVal@[_] = throwError $ NumArgs 2 singleVal
numericBinop op args = mapM unpackNum args >>= return . Number . foldl1 op
  where unpackNum :: LispVal -> ThrowsError Integer
        unpackNum (Number n) = return n
        -- unpackNum (String n) = let parsed = readDec n in
          -- if null parsed
          -- then 0
          -- else fst $ head parsed
        -- unpackNum (List [n]) = unpackNum n
        unpackNum notNum = throwError $ TypeMismatch "number" notNum

unaryOp :: (LispVal -> ThrowsError LispVal) -> [LispVal] -> ThrowsError LispVal
unaryOp op [x] = op x
unaryOp _ wrongArg = throwError $ NumArgs 1 wrongArg

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

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

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

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

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

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

symbol2String :: LispVal -> ThrowsError LispVal
symbol2String (Atom x) = return $ String x
symbol2String mismatch = throwError $ TypeMismatch "atom" mismatch

string2Symbol :: LispVal -> ThrowsError LispVal
string2Symbol (String x) = return $ Atom x
string2Symbol mismatch = throwError $ TypeMismatch "string" mismatch

-- Error Checking and Exceptions
data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

showError :: LispError -> String
showError (UnboundVar message varname)  = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func)    = message ++ ": " ++ show func
showError (NumArgs expected found)      = "Expected " ++ show expected
                                          ++ " args; found values " ++ unwordsList found
showError (TypeMismatch expected found) = "Invalid type: expected " ++ expected
                                          ++ ", found " ++ show found
showError (Parser parseErr)             = "Parse error at " ++ show parseErr

instance Show LispError where show = showError

instance Error LispError where
  noMsg = Default "An error has occurred"
  strMsg = Default

type ThrowsError = Either LispError

trapError :: (MonadError e m, Show e) => m String -> m String
trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val


-- Modulo destinado a leitura e escrita de arquivos.

module InputOutput
(getTraining,
 getTest,
 readIn,
 writeIn,
 isEmptyFile,
 definitiveAnswer,
 printEpoch
 ) where

import Data.List
import Data.Char
import Data.Maybe
import System.Directory
import System.Random
import Types
import Text.Printf
import Numeric.LinearAlgebra.HMatrix

type Image = [Double]
type Sample = (Int, Image)

-- Transforma uma String em uma lista
strToArr::String->[Double]
strToArr line = [(read x::Double)| x<-(words line)]

-- Transforma uma String em uma matrix
strToMatrix::String->[[Double]]
strToMatrix mt = [(strToArr x)| x<-(lines mt)]

-- Transforma uma lista em uma String
arrToStr::[Double]->String
arrToStr [] = ""
arrToStr (h:t) = (show h::String) ++ " " ++ (arrToStr t)

-- Transforma uma matriz em uma String
matrixToStr::[[Double]]->String
matrixToStr [] = ""
matrixToStr (h:t) = (arrToStr h) ++ "\n" ++ (matrixToStr t)

-- Dado o seu tamanho, gera uma lista
-- com valores randomicos
listRandom::Int->IO [Double]
listRandom 0 = return []
listRandom sz = do
    head <- randomIO::IO Double
    tail <- listRandom (sz-1)
    return (head:tail)

-- Dado o seu tamanho, gera uma matrix
-- com valores randomicos
matrixRandom::Int->Int->IO [[Double]]
matrixRandom 0 y = return []
matrixRandom x y = do
    head <- listRandom y
    tail <- matrixRandom (x-1) y
    return (head:tail)

-- Dado o seu caminho, le o arquivo
-- e o transforma em uma lista
getArchiveL::String->Int->IO [Double]
getArchiveL path sz = do
    emp <- isEmptyFile path
    if not emp then
        do
            archive <- (readFile path)
            return $ strToArr archive
    else
        do
            ret <- listRandom sz
            return ret

-- Dado o seu caminho, le o arquivo
-- e o transforma em uma matriz
getArchiveM::String->Int->Int->IO [[Double]]
getArchiveM path x y = do
    emp <- isEmptyFile path
    if not emp then
        do
            archive <- readFile path
            return $ strToMatrix archive
    else
        do
            ret <- matrixRandom x y
            return ret

-- Le os vetores e matrizes presentes nos arquivos para montar e retornar um Data
readIn::IO Data
readIn = do
    wH <- getArchiveM "Data/Weight_hidden.txt" 30 784
    bH <- getArchiveL "Data/Biases_hidden.txt" 30
    aH <- getArchiveL "Data/Activation_hidden.txt" 30
    zH <- getArchiveL "Data/Zeta_hidden.txt" 30

    wO <- getArchiveM "Data/Weight_Output.txt" 10 30
    bO <- getArchiveL "Data/Biases_Output.txt" 10
    aO <- getArchiveL "Data/Activation_Output.txt" 10
    zO <- getArchiveL "Data/Zeta_Output.txt" 10

    return $ (Data (fromLists wH) (fromList bH) (fromList aH) (fromList zH) (fromLists wO) (fromList bO) (fromList aO) (fromList zO))

-- Escreve o Data nos arquivos
writeIn::Data->IO()
writeIn elem = do
    writeFile "NedeReural_Haskell/Data/Weight_hidden.txt" (matrixToStr (toLists $ wHidden elem))
    writeFile "NedeReural_Haskell/Data/Biases_hidden.txt" (arrToStr (toList $ bHidden elem))
    writeFile "NedeReural_Haskell/Data/Activation_hidden.txt" (arrToStr (toList $ aHidden elem))
    writeFile "NedeReural_Haskell/Data/Zeta_hidden.txt" (arrToStr (toList $ zetaHidden elem))

    writeFile "NedeReural_Haskell/Data/Weight_Output.txt" (matrixToStr (toLists $ wOutput elem))
    writeFile "NedeReural_Haskell/Data/Biases_Output.txt" (arrToStr (toList $ bOutput elem))
    writeFile "NedeReural_Haskell/Data/Activation_Output.txt" (arrToStr (toList $ aOutput elem))
    writeFile "NedeReural_Haskell/Data/Zeta_Output.txt" (arrToStr (toList $ zetaOutput elem))

-- Verifica se o arquivo esta vazio
isEmptyFile::String->IO Bool
isEmptyFile path = do
    elem <- readFile path

    if (length elem) == 0 then return True
    else return False

-- Pega o numero que define a resposta do teste/treino
pickNumber::FilePath->Int
pickNumber "" = -1
pickNumber (h:t)
                | h == '-' = digitToInt $ t!!0
                | otherwise = (pickNumber t)

-- Constroi uma imagem
makeImage::FilePath->IO Image
makeImage path = do
    elem <- readFile path
    let ret = (strToArr elem)
    return ret

-- Constroi um Sample
makeSample::FilePath->String->IO Sample
makeSample archive path = do
    let number = (pickNumber archive)
    mt <- (makeImage $ path++archive)
    return (number, mt)

-- Cria uma lista de Sample dado os arquivos e o local
listSample::[FilePath]->String->IO [Sample]
listSample [] path = return []
listSample (h:t) path = do
                head <- makeSample h path
                tail <- listSample t path
                return $ head:tail

-- Retorna o TestSet da rede
getTest::IO [Sample]
getTest = do
    let caminho = "NedeReural_Haskell/Tests" -- Diretorio de testes
    elems <- (getDirectoryContents caminho) -- Colocar diretorio
    ret <- listSample (filter (/= ".") $ filter (/= "..") elems) caminho
    return ret

-- Retorna o TraninigSet da rede
getTraining::IO [Sample]
getTraining = do
    let caminho = "NedeReural_Haskell/Training" -- Diretorio de treino
    elems <- (getDirectoryContents caminho) -- Colocar diretorio
    ret <- listSample (filter (/= ".") $ filter (/= "..") elems) caminho
    return ret

-- Imprime na tela a taxa de acerto de uma epoca de teste
printEpoch :: Int -> Int -> Int -> IO()
printEpoch epochIndex correctCnt total = putStrLn $ "EPOCH #" ++ (show epochIndex) ++ " - " ++ (show correctCnt) ++ " / " ++ (show total)

-- Recebe a lista de ativacao como parametro
-- Retorna uma tupla: o melhor sigmoid e seu respectivo indice
getBestSigmoid :: [Double] -> (Double, Int)
getBestSigmoid activationValues = let first = maximum (activationValues)
                                      second = fromJust $ (elemIndex (maximum (activationValues)) (activationValues))
                                  in (first, second)

-- Recebe a tupla (sigmoid, indice)
-- Retorna a formatacao em string que representa
-- o numero e a probabilidade de ser esse numero
toPercentage :: (Double, Int) -> [Double] -> String
toPercentage sigmoid activationValues =
    "[" ++ (printf "%d" (snd sigmoid + 1)) ++ " - " ++ (printf "%.2f" ((fst sigmoid) * 100 / sum (activationValues))) ++ "%]"

-- Computa a resposta definitiva a partir do
-- array de valores de ativacao
definitiveAnswer :: Data -> IO String
definitiveAnswer activationValues = return ("Resposta definitiva: " ++ toPercentage (getBestSigmoid (toList $ aOutput activationValues)) (toList $ aOutput activationValues))


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses, RankNTypes #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.Types
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Types and some basic functions on them
--
-----------------------------------------------------------------------------

module AI.Network.RNN.Types where

import Control.DeepSeq
import Control.Monad.Random
import Data.List
import Data.Tuple
import Numeric.LinearAlgebra.HMatrix as M

import AI.Network.RNN.Util

import Debug.Trace

-- | Class for RNN evaluation
class (NFData a) => RNNEval a where
    type Size a
    -- | Eval one step
    evalStep   :: a -> Vector Double -> (a,Vector Double)
    -- | Transform to vector
    toVector   :: a -> Vector Double
    -- | Build from vector representation
    fromVector :: Size a -> Vector Double -> a
    -- | Transform to vector
    toDList   :: a -> [Double]
    toDList = toList . toVector
    -- | Build from vector representation
    fromDList :: Size a -> [Double] -> a
    fromDList sz = fromVector sz . M.fromList
    -- | Get the size
    rnnsize    :: a -> Size a
    -- | full size
    fullSize   :: a -> Int

instance (RNNEval a) => RNNEval [a] where
    type Size [a]   = [Size a]
    evalStep xs v = swap $ mapAccumL e v xs
        where e v1 x1 =
                let (x2,v2) = evalStep x1 v1
                in (v2,x2)
    toVector = vjoin . map toVector
    fromVector sz vs = snd $ mapAccumL e vs sz
        where e v1 s1 =
                let x1 = fromVector s1 v1
                    fs1 = fullSize x1
                in (subVector fs1 (size v1 - fs1) v1,x1)
    rnnsize = map rnnsize
    fullSize = sum . map fullSize

-- | Full size of a network: the total number of weights required
type FullSize sz = Size sz -> Int

-- | Eval several steps
evalSteps :: (Traversable t,RNNEval a) => a -> t (Vector Double) -> (a,t (Vector Double))
evalSteps = mapAccumL evalStep

-- | Cost function: lower is better
cost :: (RNNEval a) => a -> TrainData b -> Double
cost rnn td =
    let (_,res) = evalSteps rnn $ tdInputs td
    in (sum $ zipWith err (tdOutputs td) res)
    where
      err ::  Vector Double -> Vector Double -> Double
      --err tgt res =let
      --  norm = cmap (\x->if x>=0.5 then 1 else 0) res
      --  in euclidian  tgt norm
      err = euclidian
       --sum ( zipWith (\c d -> (c- d)**2 ) (toList a) (toList b)) / fromIntegral (size a)
      errsf a b = sum ( zipWith (\c d -> if c==1 then (c- d) else 0) (toList a) (softmax $ toList b))
      sf :: Vector Double -> Vector Double -> Double
      --sf a b = sum ( zipWith (\c d -> if c==1 then log d else 0 ) (toList a) (softmax $ toList b))
      sf a b = log $ (softmaxV b) ! (maxIndex a)
       -- sum ( zipWith (\c d -> if c==1 then log d else 0 ) (toList a) (softmax $ toList b))

-- | Build a random network of a given size, with weights uniformly between 0 and 1
randomNetwork
    :: (RNNEval a,Monad m,RandomGen g)
    => Size a
    -> FullSize a
    -> RandT g m a
randomNetwork sz fsz = do
    s <- getRandom
    return $ fromVector sz $ randomVector s Uniform $ fsz sz

startNetwork :: (RNNEval a)
    => Size a
    -> FullSize a
    -> a
startNetwork sz fsz = fromVector sz $ konst 0.1 (fsz sz)

-- | Training data
data TrainData a = TrainData
    { tdInputs  :: ![Vector Double]
    , tdOutputs :: ![Vector Double]
    , tdRecSize :: !Int
    , tdData    :: a}

class Contains a where
    type ContainsSize a
    buildContains :: ContainsSize a -> a
    toCList :: a -> [Int]

--data Training a c = Training (ContainsSize a) c
--
--randomContains :: (Contains a) => ContainsSize a -> a
--randomContains = buildContains
--
--fromTraining :: (Contains a)=>Training a c -> a
--fromTraining (Training sz d)= buildContains sz
--
--fromTrainingToList :: (Contains a)=>Training a c -> [Int]
--fromTrainingToList (Training sz d)=toCList $ buildContains sz


{-# LANGUAGE DeriveGeneric,
             MultiParamTypeClasses,
             FlexibleInstances,
             FlexibleContexts,
             GeneralizedNewtypeDeriving,
             FunctionalDependencies,
             UndecidableInstances,
             TypeSynonymInstances,
             TypeFamilies
#-}
{-# OPTIONS_GHC -Wno-unused-top-binds #-}

module Main where

import Control.Exception.Safe (assert, MonadThrow)
import Control.Monad
import Control.Monad.Loops
import Control.Monad.IO.Class
import Control.Monad.Primitive
import Control.Monad.RWS.Class
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader (ReaderT(..))
import Control.Monad.Trans.RWS (RWST, evalRWST)
import Data.List (maximumBy, sort)
import Data.Ord (comparing)
import Data.Vector ((!), Vector)
import Data.IORef

import Numeric.Ketamine.Agent
import Numeric.Ketamine.Environment
import Numeric.Ketamine.Episode
import qualified Numeric.Ketamine.Outcome as O

import qualified Control.Monad.Trans.Reader as TR
import qualified Data.DList as D
import qualified Data.Map.Strict as Map
import qualified Data.Vector as V
import qualified Statistics.Distribution as Dist
import qualified Statistics.Distribution.Normal as N
import qualified System.Random.MWC as R

-- | The slot machine index whose arm will be pulled
type Action = Int
type Reward = Double
type Outcome = O.Outcome Action Reward

-- TODO: separate transformer stacks and log in Episode
type SharedState = RWST (IORef BanditState) (D.DList Outcome) EnvState IO
type Environment = EnT Action Outcome SharedState
type Agent = AgT Action Outcome SharedState

-------------------------------------------------------------------------------
-- | Episode

episodeLength :: Int
episodeLength = 100

narms :: Int
narms = 10

main :: IO ()
main = do
    casino <- envState narms
    bandit <- defaultBanditState narms
    let ep :: SharedState ()
        ep = runEpisode $ stepEnv >~> stepBandit 0.1 $ 9
    res <- snd <$> evalRWST ep bandit casino
    mapM_ print $ D.toList res

-------------------------------------------------------------------------------
-- | Environment
--
data EnvState =
    EnvState { arms :: Vector N.NormalDistribution , gen :: R.GenIO, pulls :: Int }

instance Show EnvState where
    show c = "EnvState" ++
      "{ means = " ++
        show (fmap Dist.mean . V.toList $ arms c) ++
          ", pulls = " ++ show (pulls c) ++ " }"

mkEnvState :: Int -> [Double] -> IO EnvState
mkEnvState n vars = do
    gen <- R.createSystemRandom
    means <- replicateM n $ R.uniform gen
    let arms = V.fromList $ zipWith N.normalDistr (reverse $ sort means) vars
    return $ EnvState arms gen 0

-- | Default config of a n-armed bandit
envState :: Int -> IO EnvState
envState n = mkEnvState n $ replicate n 1.0

genContVar :: Dist.ContGen d => d -> Environment Double
genContVar d = do
    g <- gets gen
    liftIO $ Dist.genContVar d g

stepEnv :: Action -> Environment ()
stepEnv = loop where
    loop act = do
        rwd <- genContVar =<< (! act) . arms <$> get
        modify $ \(EnvState a g p) -> EnvState a g (p+1)
        r <- ask
        liftIO $ modifyIORef r $ addStats act (Stats act rwd)
        n <- gets pulls
        if n >= episodeLength
            then return ()
            else outcome (O.Outcome act rwd) >>= loop

------------------------------------------------------------------------------
-- | Bandit

-- Statistics observed for a particular candidate.
data Stats = Stats
    { armCount :: !Int  -- ^ Number of times this candidate was observed
    , armTotal :: !Double -- ^ Total reward over all observations
    } deriving (Show, Eq)

instance Semigroup Stats where
    s <> s' = Stats { armCount = armCount s + armCount s'
                    , armTotal = armTotal s + armTotal s' }

instance Monoid Stats where
    mempty = Stats 0 0

-- | A record of statistics for all possibilities
type BanditState = Map.Map Action Stats

-- | Average reward over all observations for that arm.
mean :: Stats -> Reward
mean (Stats 0 _) = 0
mean ss = armTotal ss / toEnum (armCount ss)

addStats :: Action -> Stats -> BanditState -> BanditState
addStats = Map.insertWith (<>)

epsilonGreedy :: (R.Variate e, Ord e, Ord r) => [(a, r)] -> e -> SharedState a
epsilonGreedy acts =
    epsilonChoice (fst $ maximumBy (comparing snd) acts) acts

epsilonChoice :: (R.Variate e, Ord e) => a -> [(a, r)] -> e -> SharedState a
epsilonChoice a acts eps = do
    g <- gets gen
    compare eps <$> R.uniform g >>= \case
        LT -> pure a
        _  -> do
            i <- R.uniformR (0, length acts) g
            pure . fst . head $ drop (i-1) acts

banditState :: Int -> Stats -> IO (IORef BanditState)
banditState n s = do
    r <- newIORef mempty
    let init = Map.fromList $ take n $ zip [0..] (repeat s)
    writeIORef r init
    return r

defaultBanditState :: Int -> IO (IORef BanditState)
defaultBanditState n = banditState n mempty

stepBandit :: Float -> Outcome -> Agent ()
stepBandit eps = loop where
    loop o@(O.Outcome act rwd) = do
        tell . pure $ o
        r <- ask
        liftIO $ modifyIORef r $ addStats act (Stats act rwd)
        s <- liftIO $ readIORef r
        let rwds = Map.map mean s
        a <- lift $ epsilonGreedy (Map.toList rwds) eps
        action a >>= loop


{-# LANGUAGE DeriveGeneric    #-}
{-# LANGUAGE TypeApplications #-}
module Noether.Algebra.Actions.Acts where

import           Data.Complex

import           Noether.Lemmata.Prelude
import           Noether.Lemmata.TypeFu

import           Noether.Algebra.Single
import           Noether.Algebra.Tags

import           GHC.Generics

data ActsE
  = Acts_Magma MagmaE
  | ActsNamed Symbol ActsE
  | ActsTagged Type ActsE
  deriving (Generic)

class ActsK (lr :: Side) (op :: k) a b (s :: ActsE) where
  actK :: Proxy op -> Proxy s -> Proxy lr -> a -> b -> b

type family ActsS (lr :: Side) (op :: k) (a :: Type) (b :: Type) = (r :: ActsE)

instance MagmaK op a zm => ActsK lr op a a (Acts_Magma zm) where
  actK opP _ _ = binaryOpK opP (Proxy @zm)


{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
module Statistics.BBVI.Gradient
  ( gradientScore
  , gradientReparam
  , DistInvariant(..)
  , Samples
  )
where

import qualified Data.Vector                   as V
import           Statistics.BBVI.Propagator     ( DistCell(..)
                                                , Gradient
                                                , Memory
                                                , dist
                                                )
import           Statistics.BBVI.Class

type Samples = V.Vector

-- | invariant data for a variational distribution cell
data DistInvariant a = DistInvariant { weight :: !Double
                                     -- ^ total occurences of z in log
                                     -- likelihood
                                     , prior :: !a
                                     -- ^ prior on latent variable
                                     , rhoF :: !(Gradient ->
                                                 DistCell a ->
                                                 (Memory, V.Vector Double))
                                     -- ^ calculates change in memory,
                                     -- and stepsize
                                     }

-- | internal helper function for building gradient propagators
gradient
  :: (DistUtil a)
  => (DistCell a -> Double -> c -> Double -> V.Vector Double)
  -> DistInvariant a
  -> DistCell a
  -> (Double, V.Vector Double, Samples c)
  -> DistCell a3
gradient f (DistInvariant {..}) no (nFactors, like, samples) = U
  memory'
  (V.zipWith (*) rho' gr)
 where
  summed =
    V.foldl' (V.zipWith (+)) (V.replicate (nParams (dist no)) 0.0)
      $ V.zipWith (f no nFactors) samples like
  gr              = fmap (/ (fromIntegral $ V.length samples)) summed
  (memory', rho') = rhoF gr no

-- | helper function for transforming log joint to score gradient
-- propagator
gradientScore
  :: Dist a c
  => DistInvariant a
  -> DistCell a -- ^ current variational distribution
  -> (Double, V.Vector Double, Samples c)   -- ^ n factors in log
  -- joint calculation, log joint calculation on each sample of \(z\),
  -- samples
  -> DistCell a -- ^ update to variational distribution
gradientScore gp@(DistInvariant {..}) = gradient f gp
 where
  f (Node _time _memory d) nFactors s l = fmap
    (* (l + nFactors / weight * (logProb prior s - logProb d s)))
    (paramGradOfLogQ d s)
  f (U _ _) _ _ _ = error "called gradient update on update cell"

-- | helper function for transforming log joint to reparameterization
-- gradient propagator
gradientReparam
  :: (Differentiable a Double)
  => DistInvariant a
  -> DistCell a -- ^ current variational distribution
  -> (Double, V.Vector Double, Samples Double)
  -- ^ n factors in log joint calculation, gradient of joint wrt \(z\)
  -- for each sample, un-transformed samples
  -> DistCell a -- ^ update to variational distribution
gradientReparam gp@(DistInvariant {..}) = gradient f gp
 where
  f (Node _time _memory d) nFactors s g = fmap
    (* ( g
       + nFactors
       / weight
       * ( sampleGradOfLogQ prior (transform d s)
         - sampleGradOfLogQ d     (transform d s)
         )
       )
    )
    (gradTransform d s)
  f (U _ _) _ _ _ = err

err :: a
err = error "called gradient update on update cell"


{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE Trustworthy #-}
{-# OPTIONS_GHC -fenable-rewrite-rules #-}
----------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2011-2014
-- License     :  BSD3
--
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
--
-- Representable endofunctors over the category of Haskell types are
-- isomorphic to the reader monad and so inherit a very large number
-- of properties for free.
----------------------------------------------------------------------

module Data.Functor.Rep
  (
  -- * Representable Functors
    Representable(..)
  , tabulateAlg
  , tabulated
  -- * Logarithms
  , Logarithm(..)
  , contramapLogarithm
  , logarithmRep
  -- * Wrapped representable functors
  , Co(..)
  -- * Default definitions
  -- ** Functor
  , fmapRep
  -- ** Distributive
  , distributeRep
  , collectRep
  , cotraverseRep
  -- ** Apply/Applicative
  , apRep
  , pureRep
  , liftR2
  , liftR3
  -- ** Bind/Monad
  , bindRep
  -- ** MonadFix
  , mfixRep
  -- ** MonadZip
  , mzipRep
  , mzipWithRep
  -- ** MonadReader
  , askRep
  , localRep
  -- ** Extend
  , duplicatedRep
  , extendedRep
  -- ** Comonad
  , duplicateRep
  , extendRep
  , extractRep
  -- ** Comonad, with user-specified monoid
  , duplicateRepBy
  , extendRepBy
  , extractRepBy
  -- ** WithIndex
  , imapRep
  , ifoldMapRep
  , itraverseRep
  -- ** Representable
  , tabulateCotraverse1
  , indexLogarithm
  , cotraverse1Iso
  -- ** Generics
  , gcotraverse1
  , GRep
  , gindex
  , gtabulate
  , WrappedRep(..)
  ) where

import Control.Applicative
import Control.Applicative.Backwards
#if __GLASGOW_HASKELL__ >= 708
import Data.Coerce
#endif
import Control.Comonad
import Control.Comonad.Trans.Class
import Control.Comonad.Trans.Traced
import Control.Comonad.Cofree
import Control.Monad.Trans.Identity
import Control.Monad.Reader
#if MIN_VERSION_base(4,4,0)
import Data.Complex
#endif
import Data.Distributive
import Data.Foldable (Foldable(fold))
import Data.Functor.Bind
import Data.Functor.Identity
import Data.Functor.Compose
import Data.Functor.Extend
import Data.Functor.Product
import Data.Functor.Reverse
import Data.Functor1
import Data.Functor1.Applied
import qualified Data.Monoid as Monoid
import Data.Profunctor.Unsafe
import Data.Proxy
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import Data.Semigroup hiding (Product)
import Data.Tagged
import Data.Traversable (Traversable(sequenceA))
import Data.Void
import GHC.Generics hiding (Rep)
import Prelude

-- | A 'Functor' @f@ is 'Representable' if 'tabulate' and 'index' witness an
-- isomorphism to @(->) r@, for some @r@.
--
-- Alternatively, an instance can be derived from 'cotraverse1', by defining:
--
-- @
-- 'Rep' f = 'Logarithm' f
-- 'tabulate' = 'tabulateCotraverse1'
-- 'index' = 'indexLogarithm'
-- @
--
-- Instances without random access should implement 'cotraverse1', as it can
-- allow asymptotically faster zipping.
--
-- Every 'Distributive' 'Functor' is actually 'Representable'.
--
-- Every 'Representable' 'Functor' from Hask to Hask is a right adjoint.
--
-- @
-- 'tabulate' . 'index'  ≡ id
-- 'index' . 'tabulate'  ≡ id
-- 'tabulate' . 'return' ≡ 'return'
--
-- 'distribute1' . 'Applied' ≡ 'fmap' ('Applied' . 'Identity')
-- 'distribute1' ('Const' x) ≡ 'Const' x '<$' xs
--
-- 'cotraverse1' f ≡ 'fmap' f . 'distribute1'
-- 'collect1' f ≡ 'distribute1' . 'map1' f
-- 'cotraverseMap1' f g ≡ 'cotraverse1' f . 'map1' g
-- @

class Distributive f => Representable f where
  -- | If no definition is provided, this will default to 'GRep'.
  type Rep f :: *
  type Rep f = GRep f

  -- |
  -- @
  -- 'fmap' f . 'tabulate' ≡ 'tabulate' . 'fmap' f
  -- @
  --
  -- If no definition is provided, this will default to 'gtabulate'.
  tabulate :: (Rep f -> a) -> f a
  default tabulate :: (Generic1 f, GRep f ~ Rep f, GTabulate (Rep1 f))
                   => (Rep f -> a) -> f a
  tabulate = gtabulate

  -- | If no definition is provided, this will default to 'gindex'.
  index    :: f a -> Rep f -> a
  default index :: (Generic1 f, GRep f ~ Rep f, GIndex (Rep1 f))
                => f a -> Rep f -> a
  index = gindex

  -- | A more powerful version of 'cotraverse'
  --
  -- @'cotraverse1' f = 'fmap' f . 'distribute1'@
  cotraverse1 :: Functor1 w => (w Identity -> a) -> w f -> f a
  cotraverse1 f w = tabulateAlg (\g -> f $ map1Identity g w)

  -- | A more powerful version of 'distribute'
  --
  -- @
  -- 'distribute1' . 'Applied' ≡ 'fmap' ('Applied' . 'Identity')
  -- 'distribute1' ('Const' x) ≡ 'Const' x '<$' xs
  -- @
  distribute1 :: Functor1 w => w f -> f (w Identity)
  distribute1 = cotraverse1 id

  -- | A more powerful version of 'collect'
  --
  -- @'collect1' f ≡ 'distribute1' . 'map1' f@
  collect1 :: Functor1 w => (forall x. g x -> f x) -> w g -> f (w Identity)
  collect1 f = distribute1 . map1 f

  -- | @'cotraverseMap1' f g ≡ 'cotraverse1' f . 'map1' g@
  cotraverseMap1 ::
       Functor1 w => (w Identity -> a) -> (forall x. g x -> f x) -> w g -> f a
  cotraverseMap1 f g = cotraverse1 f . map1 g

tabulateAlg :: Representable f => ((forall x. f x -> x) -> a) -> f a
tabulateAlg f = tabulate $ \i -> f (`index` i)

-- | A default implementation of 'Rep' for a datatype that is an instance of
-- 'Generic1'. This is usually composed of 'Either', tuples, unit tuples, and
-- underlying 'Rep' values. For instance, if you have:
--
-- @
-- data Foo a = MkFoo a (Bar a) (Baz (Quux a)) deriving ('Functor', 'Generic1')
-- instance 'Representable' Foo
-- @
--
-- Then you'll get:
--
-- @
-- 'GRep' Foo = Either () (Either ('WrappedRep' Bar) ('WrappedRep' Baz, 'WrappedRep' Quux))
-- @
--
-- (See the Haddocks for 'WrappedRep' for an explanation of its purpose.)
type GRep f = GRep' (Rep1 f)

-- | A default implementation of 'tabulate' in terms of 'GRep'.
gtabulate :: (Generic1 f, GRep f ~ Rep f, GTabulate (Rep1 f))
          => (Rep f -> a) -> f a
gtabulate = to1 . gtabulate'

-- | A default implementation of 'index' in terms of 'GRep'.
gindex :: (Generic1 f, GRep f ~ Rep f, GIndex (Rep1 f))
       => f a -> Rep f -> a
gindex = gindex' . from1

type family GRep' (f :: * -> *) :: *
class GTabulate f where
  gtabulate' :: (GRep' f -> a) -> f a
class GIndex f where
  gindex' :: f a -> GRep' f -> a

type instance GRep' (f :*: g) = Either (GRep' f) (GRep' g)
instance (GTabulate f, GTabulate g) => GTabulate (f :*: g) where
  gtabulate' f = gtabulate' (f . Left) :*: gtabulate' (f . Right)
instance (GIndex f, GIndex g) => GIndex (f :*: g) where
  gindex' (a :*: _) (Left  i) = gindex' a i
  gindex' (_ :*: b) (Right j) = gindex' b j

type instance GRep' (f :.: g) = (WrappedRep f, GRep' g)
instance (Representable f, GTabulate g) => GTabulate (f :.: g) where
  gtabulate' f = Comp1 $ tabulate $ gtabulate' <$> fmap (curry f) WrapRep
instance (Representable f, GIndex g) => GIndex (f :.: g) where
  gindex' (Comp1 fg) (i, j) = gindex' (index fg (unwrapRep i)) j

type instance GRep' Par1 = ()
instance GTabulate Par1 where
  gtabulate' f = Par1 (f ())
instance GIndex Par1 where
  gindex' (Par1 a) () = a

type instance GRep' (Rec1 f) = WrappedRep f
#if __GLASGOW_HASKELL__ >= 708
-- Using coerce explicitly here seems a bit more readable, and
-- likely a drop easier on the simplifier.
instance Representable f => GTabulate (Rec1 f) where
  gtabulate' = coerce (tabulate :: (Rep f -> a) -> f a)
                 :: forall a . (WrappedRep f -> a) -> Rec1 f a
instance Representable f => GIndex (Rec1 f) where
  gindex' = coerce (index :: f a -> Rep f -> a)
                 :: forall a . Rec1 f a -> WrappedRep f -> a
#else
instance Representable f => GTabulate (Rec1 f) where
  gtabulate' = Rec1 #. tabulate .# (. WrapRep)
instance Representable f => GIndex (Rec1 f) where
  gindex' = (. unwrapRep) #. index .# unRec1
#endif

type instance GRep' (M1 i c f) = GRep' f
instance GTabulate f => GTabulate (M1 i c f) where
  gtabulate' = M1 #. gtabulate'
instance GIndex f => GIndex (M1 i c f) where
  gindex' = gindex' .# unM1

-- | On the surface, 'WrappedRec' is a simple wrapper around 'Rep'. But it plays
-- a very important role: it prevents generic 'Representable' instances for
-- recursive types from sending the typechecker into an infinite loop. Consider
-- the following datatype:
--
-- @
-- data Stream a = a :< Stream a deriving ('Functor', 'Generic1')
-- instance 'Representable' Stream
-- @
--
-- With 'WrappedRep', we have its 'Rep' being:
--
-- @
-- 'Rep' Stream = 'Either' () ('WrappedRep' Stream)
-- @
--
-- If 'WrappedRep' didn't exist, it would be:
--
-- @
-- 'Rep' Stream = Either () (Either () (Either () ...))
-- @
--
-- An infinite type! 'WrappedRep' breaks the potentially infinite loop.
newtype WrappedRep f = WrapRep { unwrapRep :: Rep f }

{-# RULES
"tabulate/index" forall t. tabulate (index t) = t #-}

-- | 'tabulate' and 'index' form two halves of an isomorphism.
--
-- This can be used with the combinators from the @lens@ package.
--
-- @'tabulated' :: 'Representable' f => 'Iso'' ('Rep' f -> a) (f a)@
tabulated :: (Representable f, Representable g, Profunctor p, Functor h)
          => p (f a) (h (g b)) -> p (Rep f -> a) (h (Rep g -> b))
tabulated = dimap tabulate (fmap index)
{-# INLINE tabulated #-}

-- * Logarithms

-- | Can be used as a value for 'Rep'
newtype Logarithm f = Logarithm { runLogarithm :: forall x. f x -> x }

contramapLogarithm :: (forall x. f x -> g x) -> Logarithm g -> Logarithm f
contramapLogarithm f (Logarithm g) = Logarithm (g . f)

-- | An index is equivalent to a function which gets the element at that index.
--
-- This can be used with the combinators from the @lens@ package.
--
-- @'logarithmRep' :: 'Representable' f => 'Iso' ('Logarithm' f) ('Rep' f)@
logarithmRep ::
  (Representable f, Representable g, Profunctor p, Functor h) =>
  p (Rep f) (h (Rep g)) -> p (Logarithm f) (h (Logarithm g))
logarithmRep =
  dimap (\(Logarithm f) -> f askRep) (fmap (\x -> Logarithm (`index` x)))

-- * Default definitions

fmapRep :: Representable f => (a -> b) -> f a -> f b
fmapRep f = cotraverse1 (f . runIdentity . runApplied) . Applied

pureRep :: Representable f => a -> f a
pureRep = tabulate . const

bindRep :: Representable f => f a -> (a -> f b) -> f b
bindRep m f = distribute f `apRep` m

mfixRep :: Representable f => (a -> f a) -> f a
mfixRep = cotraverseRep fix

data PairOf a b f = PairOf (f a) (f b)
instance Functor1 (PairOf a b) where
  map1 f (PairOf x y) = PairOf (f x) (f y)

mzipWithRep :: Representable f => (a -> b -> c) -> f a -> f b -> f c
mzipWithRep f as bs =
  cotraverse1 (\(PairOf (Identity a) (Identity b)) -> f a b) (PairOf as bs)

mzipRep :: Representable f => f a -> f b -> f (a, b)
mzipRep = mzipWithRep (,)

askRep :: Representable f => f (Rep f)
askRep = tabulate id

localRep :: Representable f => (Rep f -> Rep f) -> f a -> f a
localRep f m = tabulate (index m . f)

apRep :: Representable f => f (a -> b) -> f a -> f b
apRep = mzipWithRep ($)

distributeRep :: (Representable f, Functor w) => w (f a) -> f (w a)
distributeRep = cotraverseRep id

collectRep :: (Representable f, Functor w) => (a -> f b) -> w a -> f (w b)
collectRep f = distributeRep . fmap f

newtype Composed g a f = Composed { runComposed :: g (f a) }
instance Functor g => Functor1 (Composed g a) where
  map1 f = Composed . fmap f . runComposed

cotraverseRep :: (Representable f, Functor w) => (w a -> b) -> w (f a) -> f b
cotraverseRep f = cotraverse1 (f . fmap runIdentity . runComposed) . Composed

duplicateRepBy :: Representable f => (Rep f -> Rep f -> Rep f) -> f a -> f (f a)
duplicateRepBy plus w = tabulate (\m -> tabulate (index w . plus m))

extendRepBy :: Representable f => (Rep f -> Rep f -> Rep f) -> (f a -> b) -> f a -> f b
extendRepBy plus f w = tabulate (\m -> f (tabulate (index w . plus m)))

extractRepBy :: Representable f => Rep f -> f a -> a
extractRepBy = flip index

duplicatedRep :: (Representable f, Semigroup (Rep f)) => f a -> f (f a)
duplicatedRep = duplicateRepBy (<>)

extendedRep :: (Representable f, Semigroup (Rep f)) => (f a -> b) -> f a -> f b
extendedRep = extendRepBy (<>)

duplicateRep :: (Representable f, Monoid (Rep f)) => f a -> f (f a)
duplicateRep = duplicateRepBy mappend

extendRep :: (Representable f, Monoid (Rep f)) => (f a -> b) -> f a -> f b
extendRep = extendRepBy mappend

extractRep :: (Representable f, Monoid (Rep f)) => f a -> a
extractRep = extractRepBy mempty

imapRep :: Representable r => (Rep r -> a -> b) -> r a -> r b
imapRep f = mzipWithRep f askRep

ifoldMapRep :: forall r m a. (Representable r, Foldable r, Monoid m)
            => (Rep r -> a -> m) -> (r a -> m)
ifoldMapRep ix = fold . imapRep ix

itraverseRep :: forall r f a a'. (Representable r, Traversable r, Applicative f)
             => (Rep r -> a -> f a') -> (r a -> f (r a'))
itraverseRep ix = sequenceA . imapRep ix

newtype TabulateArg a f = TabulateArg (Logarithm f -> a)
instance Functor1 (TabulateArg a) where
  map1 f (TabulateArg g) = TabulateArg (g . contramapLogarithm f)

-- | Derive 'tabulate' given @'Rep' f ~ 'Logarithm' f@ and an
-- implementation of 'cotraverse1'
tabulateCotraverse1 :: Representable f => (Logarithm f -> a) -> f a
tabulateCotraverse1 =
  cotraverse1 (\(TabulateArg g) -> g (Logarithm runIdentity)) . TabulateArg

-- | Derive 'index', given @'Rep' f ~ 'Logarithm' f@
indexLogarithm :: f a -> Logarithm f -> a
indexLogarithm = flip runLogarithm

-- | Derive 'cotraverse1' via an isomorphism
cotraverse1Iso ::
     (Representable g, Functor1 w)
  => (forall x. f x -> g x)
  -> (forall x. g x -> f x)
  -> (w Identity -> a)
  -> w f
  -> f a
cotraverse1Iso t frm f = frm . cotraverseMap1 f t

gcotraverse1 ::
     (Representable (Rep1 f), Functor1 w, Generic1 f)
  => (w Identity -> a)
  -> w f
  -> f a
gcotraverse1 = cotraverse1Iso from1 to1

-- * Instances

instance Representable Proxy where
  type Rep Proxy = Void
  index Proxy = absurd
  tabulate _ = Proxy

instance Representable Identity where
  type Rep Identity = ()
  index (Identity a) () = a
  tabulate f = Identity (f ())

instance Representable (Tagged t) where
  type Rep (Tagged t) = ()
  index (Tagged a) () = a
  tabulate f = Tagged (f ())

instance Representable m => Representable (IdentityT m) where
  type Rep (IdentityT m) = Rep m
  index = index .# runIdentityT
  tabulate = IdentityT #. tabulate
  cotraverse1 = cotraverse1Iso runIdentityT IdentityT

instance Representable ((->) e) where
  type Rep ((->) e) = e
  index = id
  tabulate = id

instance Representable m => Representable (ReaderT e m) where
  type Rep (ReaderT e m) = (e, Rep m)
  index (ReaderT f) (e,k) = index (f e) k
  tabulate = ReaderT . fmap tabulate . curry
  cotraverse1 = cotraverse1Iso (Comp1 . runReaderT) (ReaderT . unComp1)

instance (Representable f, Representable g) => Representable (Compose f g) where
  type Rep (Compose f g) = (Rep f, Rep g)
  index (Compose fg) (i,j) = index (index fg i) j
  tabulate = Compose . tabulate . fmap tabulate . curry
  cotraverse1 = cotraverse1Iso (Comp1 . getCompose) (Compose . unComp1)

instance Representable w => Representable (TracedT s w) where
  type Rep (TracedT s w) = (s, Rep w)
  index (TracedT w) (e,k) = index w k e
  tabulate = TracedT . unCo . collect (Co #. tabulate) . curry
  cotraverse1 = cotraverse1Iso (Comp1 . runTracedT) (TracedT . unComp1)

instance (Representable f, Representable g) => Representable (Product f g) where
  type Rep (Product f g) = Either (Rep f) (Rep g)
  index (Pair a _) (Left i)  = index a i
  index (Pair _ b) (Right j) = index b j
  tabulate f = Pair (tabulate (f . Left)) (tabulate (f . Right))
  cotraverse1 = cotraverse1Iso (\(Pair x y) -> x :*: y) (\(x :*: y) -> Pair x y)

instance Representable f => Representable (Cofree f) where
  type Rep (Cofree f) = Seq (Rep f)
  index (a :< as) key = case Seq.viewl key of
      Seq.EmptyL -> a
      k Seq.:< ks -> index (index as k) ks
  tabulate f = f Seq.empty :< tabulate (\k -> tabulate (f . (k Seq.<|)))

  -- this could be derived via isomorphism to
  -- Identity :*: (f :.: Cofree f), but then the instance would be
  -- recursive which would prevent specialization
  cotraverse1 f = go
    where
      go w =
        f (map1Identity extract w) :<
        cotraverse1
          (go . map1 (runIdentity . unComp1) . runAppCompose)
          (AppCompose $ map1 (Comp1 . unwrap) w)

instance Representable f => Representable (Backwards f) where
  type Rep (Backwards f) = Rep f
  index = index .# forwards
  tabulate = Backwards #. tabulate
  cotraverse1 = cotraverse1Iso forwards Backwards

instance Representable f => Representable (Reverse f) where
  type Rep (Reverse f) = Rep f
  index = index .# getReverse
  tabulate = Reverse #. tabulate
  cotraverse1 = cotraverse1Iso getReverse Reverse

instance Representable Monoid.Dual where
  type Rep Monoid.Dual = ()
  index (Monoid.Dual d) () = d
  tabulate f = Monoid.Dual (f ())

instance Representable Monoid.Product where
  type Rep Monoid.Product = ()
  index (Monoid.Product p) () = p
  tabulate f = Monoid.Product (f ())

instance Representable Monoid.Sum where
  type Rep Monoid.Sum = ()
  index (Monoid.Sum s) () = s
  tabulate f = Monoid.Sum (f ())

#if MIN_VERSION_base(4,4,0)
instance Representable Complex where
  type Rep Complex = Bool
  index (r :+ i) key = if key then i else r
  tabulate f = f False :+ f True
#endif

instance Representable U1 where
  type Rep U1 = Void
  index U1 = absurd
  tabulate _ = U1

instance (Representable f, Representable g) => Representable (f :*: g) where
  type Rep (f :*: g) = Either (Rep f) (Rep g)
  index (a :*: _) (Left  i) = index a i
  index (_ :*: b) (Right j) = index b j
  tabulate f = tabulate (f . Left) :*: tabulate (f . Right)
  cotraverse1 f w =
    cotraverseMap1 f (\(a :*: _) -> a) w :*: cotraverseMap1 f (\(_ :*: b) -> b) w

newtype AppCompose w g f = AppCompose { runAppCompose :: w (f :.: g) }
instance Functor1 w => Functor1 (AppCompose w g) where
  map1 f = AppCompose . map1 (Comp1 . f . unComp1) . runAppCompose

instance (Representable f, Representable g) => Representable (f :.: g) where
  type Rep (f :.: g) = (Rep f, Rep g)
  index (Comp1 fg) (i, j) = index (index fg i) j
  tabulate = Comp1 . tabulate . fmap tabulate . curry
  cotraverse1 f w =
    Comp1 $
    cotraverse1 (cotraverseMap1 f (runIdentity . unComp1) . runAppCompose) $
    AppCompose w

instance Representable Par1 where
  type Rep Par1 = ()
  index (Par1 a) () = a
  tabulate f = Par1 (f ())

instance Representable f => Representable (Rec1 f) where
  type Rep (Rec1 f) = Rep f
  index = index .# unRec1
  tabulate = Rec1 #. tabulate
  cotraverse1 = cotraverse1Iso unRec1 Rec1

instance Representable f => Representable (M1 i c f) where
  type Rep (M1 i c f) = Rep f
  index = index .# unM1
  tabulate = M1 #. tabulate
  cotraverse1 = cotraverse1Iso unM1 M1

newtype Co f a = Co { unCo :: f a } deriving Functor

instance Representable f => Representable (Co f) where
  type Rep (Co f) = Rep f
  tabulate = Co #. tabulate
  index = index .# unCo
  cotraverse1 = cotraverse1Iso unCo Co

instance Representable f => Apply (Co f) where
  (<.>) = apRep

instance Representable f => Applicative (Co f) where
  pure = pureRep
  (<*>) = apRep

instance Representable f => Distributive (Co f) where
  distribute = distributeRep
  collect = collectRep

instance Representable f => Bind (Co f) where
  (>>-) = bindRep

instance Representable f => Monad (Co f) where
  return = pure
  (>>=) = bindRep

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 704
instance (Representable f, Rep f ~ a) => MonadReader a (Co f) where
  ask = askRep
  local = localRep
#endif

instance (Representable f, Semigroup (Rep f)) => Extend (Co f) where
  extended = extendedRep

instance (Representable f, Monoid (Rep f)) => Comonad (Co f) where
  extend = extendRep
  extract = extractRep

instance ComonadTrans Co where
  lower (Co f) = f

liftR2 :: Representable f => (a -> b -> c) -> f a -> f b -> f c
liftR2 = mzipWithRep

data TripleOf a b c f = TripleOf (f a) (f b) (f c)
instance Functor1 (TripleOf a b c) where
  map1 f (TripleOf a b c) = TripleOf (f a) (f b) (f c)

liftR3 :: Representable f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftR3 f fa fb fc =
  cotraverse1
    (\(TripleOf (Identity a) (Identity b) (Identity c)) -> f a b c)
    (TripleOf fa fb fc)


{-# LANGUAGE DataKinds, TypeFamilies #-}
module TestConditionalSynth (testTrees) where

import Prelude

import           Data.List (nub, sort, unfoldr)
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Storable
import           Foreign.Storable (Storable)
import           Foreign.Storable.Tuple ()
import           GHC.Exts (inline)
import           Numeric.LinearAlgebra (Vector)
import qualified Numeric.LinearAlgebra as HM
import           System.Random
import           Test.Tasty
import           Test.Tasty.HUnit hiding (assert)

import HordeAd
import HordeAd.Tool.MnistTools

testTrees :: [TestTree]
testTrees = [ conditionalSynthTests
            ]

-- * A neural net that learns to represent a list of input and result pairs
-- as a function that is a sum of conditionals with linear conditions
-- and linear or zero results.
--
-- The samples are easy: inputs and results are integers and inputs
-- are sorted and not repeating (with input repetition it would not be
-- a function or the output would need to be repeated as well).

bloat :: Int
bloat = 1

lenSynthV :: Int -> Int -> Data.Vector.Vector Int
lenSynthV width nSamples =
  V.fromList $ replicate width (nSamples * 2) ++ [width]
               ++ replicate (bloat * nSamples * 3) width
               ++ replicate 3 (bloat * nSamples)

-- If ai ranges over ps1, bi over ps2 and ci over ps4,
-- the value of the function on input x is
-- the sum of if x * ai + bi > 0 then (x * ai + bi) * ci else 0, which is
-- the sum of if x * ai + bi > 0 then x * ai * ci + bi * ci else 0
-- so each condition depends in a linear way on x and each result,
-- if not 0, depends in a linear way on input x.
--
-- To approximate the samples (a list of input and result pairs on which
-- parameters are trained or tested) using this code, divide the input
-- and multiply result appropriately, see @synthLossSquared@.
synthValue :: forall d r m. DualMonad d r m
           => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
           -> r
           -> DualNumber d (Vector r)
           -> DualNumber d (Vector r)
           -> DualNumber d (Vector r)
           -> m (DualNumber d r)
synthValue factivation x ps1@(D u _) ps2 ps3 = do
  activated <- factivation $ scale (HM.konst x (V.length u)) ps1 + ps2
  returnLet $ activated <.>! ps3

synthLossSquared :: DualMonad d r m
                 => (DualNumber d (Vector r)
                     -> m (DualNumber d (Vector r)))
                 -> r
                 -> DualNumber d (Vector r)
                 -> DualNumber d (Vector r)
                 -> DualNumber d (Vector r)
                 -> r
                 -> m (DualNumber d r)
synthLossSquared factivation x ps1 ps2 ps3 targ = do
  y <- inline synthValue factivation (x / 1000) ps1 ps2 ps3
  lossSquared (targ / 10000) y  -- smaller target to overcome @tanh@ clamping

-- Inlined to avoid the tiny overhead of calling an unknown function.
sumResultsDual :: forall d r m a. (DualMonad d r m, Storable a)
               => (a -> m (DualNumber d r))
               -> Vector a
               -> m (DualNumber d r)
{-# INLINE sumResultsDual #-}
sumResultsDual f as = do
  let g :: DualNumber d r -> a -> m (DualNumber d r)
      g !acc a = do
        u <- f a
        return $! acc + u
  sumUs <- V.foldM g 0 as
  returnLet sumUs

synthLossAll
  :: forall d r m. DualMonad d r m
  => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
  -> Data.Vector.Storable.Vector (r, r)
  -> DualNumber d (Vector r)
  -> DualNumber d (Vector r)
  -> DualNumber d (Vector r)
  -> m (DualNumber d r)
synthLossAll factivation samples ps1 ps2 ps3 = do
  let f :: (r, r) -> m (DualNumber d r)
      f (x, y) = inline synthLossSquared factivation x ps1 ps2 ps3 y
  sumResultsDual f samples

sumTrainableInputsS :: forall d r. IsScalar d r
                    => DualNumber d (Vector r)
                    -> Int
                    -> DualNumberVariables d r
                    -> Int
                    -> Data.Vector.Vector (DualNumber d r)
sumTrainableInputsS x offset variables width =
  let f :: Int -> DualNumber d r
      f i = sumTrainableInputsV x (offset + i) variables
  in V.generate width f

splitLayerV :: forall d r m. DualMonad d r m
            => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
            -> DualNumber d (Vector r)
            -> Int
            -> DualNumberVariables d r
            -> Int
            -> m ( DualNumber d (Vector r)
                 , DualNumber d (Vector r)
                 , DualNumber d (Vector r) )
splitLayerV factivation hiddenVec offset variables width = do
  let multiplied = sumTrainableInputsS hiddenVec offset variables width
      chunkWidth = width `div` 3
      activate :: Int -> m (DualNumber d (Vector r))
      activate n = do
        let v = V.slice (n * chunkWidth) chunkWidth multiplied
        factivation $ seq1 v + var1 variables (offset + width + n)
  a0 <- activate 0
  a1 <- activate 1
  a2 <- activate 2
  return (a0, a1, a2)

synthLossBareTotal
  :: forall d r m. DualMonad d r m
  => (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
  -> (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
  -> (DualNumber d (Vector r) -> m (DualNumber d (Vector r)))
  -> Int
  -> Data.Vector.Storable.Vector (r, r)
  -> DualNumberVariables d r
  -> m (DualNumber d r)
synthLossBareTotal factivation factivationHidden factivationMiddle
                   width samples variables = do
  let (inputs, outputs) = V.unzip samples
      nSamples = V.length samples
      sampleData = inputs <> outputs
      hiddenLayer1 = sumConstantDataL sampleData 0 variables width
                     + var1 variables width  -- bias
  nonlinearLayer1 <- factivationHidden hiddenLayer1
  let offsetMiddle = width + 1
  (ps1, ps2, ps3) <-
    inline splitLayerV factivationMiddle nonlinearLayer1
                       offsetMiddle variables (bloat * nSamples * 3)
  inline synthLossAll factivation samples ps1 ps2 ps3

-- * Tests and generation of random data

-- Pair samples sorted and made unique wrt first element of the pair.
integerPairSamples :: (Storable r, Num r)
                   => (Int, Int) -> Int -> Int
                   -> Data.Vector.Storable.Vector (r, r)
integerPairSamples range seed k =
  let rolls :: RandomGen g => g -> [Int]
      rolls = unfoldr (Just . uniformR range)
      (g1, g2) = split $ mkStdGen seed
      nubInputs :: [Int] -> [Int] -> [Int]
      nubInputs candidates rest =
        let len = length candidates
        in if len == k
           then candidates
           else let (candidatesExtra, restExtra) = splitAt (k - len) rest
                    candidatesUniq = nub $ sort $ candidates ++ candidatesExtra
                in nubInputs candidatesUniq restExtra
      inputs = nubInputs [] (rolls g1)
  in V.zip (V.fromListN k $ map fromIntegral inputs)
           (V.fromListN k $ map fromIntegral $ rolls g2)

gradSmartTestCase
  :: forall r. (HasDelta r, r ~ Double)
  => String
  -> (Int
      -> Data.Vector.Storable.Vector (r, r)
      -> DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Int -> Int -> Int -> Int -> r
  -> TestTree
gradSmartTestCase prefix lossFunction seedSamples
                  nSamples width nIterations expected =
  let makeSamples s =
        integerPairSamples (-1000, 1000) (seedSamples + s) nSamples
      samples = map makeSamples [42, 49 .. 7 * nIterations]
      testSamples = map makeSamples [-1, -2 .. - 100]
      nParams1 = lenSynthV width nSamples
      -- Values from -0.5 to 0.5. TODO: start biases at 1
      params1Init =
        V.imap (\i nPV -> HM.randomVector (33 + nPV + i) HM.Uniform nPV
                          - HM.scalar 0.5)
               nParams1
      parametersInit = (V.empty, params1Init, V.empty, V.empty)
      name = prefix ++ " "
             ++ unwords [ show seedSamples, show nSamples, show width
                        , show (V.length nParams1), show (V.sum nParams1) ]
      f = lossFunction width
  in testCase name $ do
       let (parametersResult, _) =
             sgdAdam f samples parametersInit
                     (initialStateAdam parametersInit)
           (_, values) =
             unzip $ map (\t -> dReverse 1 (f t) parametersResult) testSamples
       (sum values / 100) @?= expected

conditionalSynthTests:: TestTree
conditionalSynthTests = do
 let f = inline synthLossBareTotal reluAct tanhAct tanhAct
 testGroup "synthesizing a sum of linear conditionals matching samples"
  [ gradSmartTestCase "reluAct"
      f 42 10 10  100
      4.740275311294229
  , gradSmartTestCase "reluAct"
      f 42 10 10  10000
      3.83451707827233e-2
  , gradSmartTestCase "reluAct"
      f 42 10 10  100000
      3.135485708489271e-2
  , gradSmartTestCase "reluAct"
      f 42 10 100 100000
      3.2872191198993095e-2
  ]