text:
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}

module Main where

import Control.Monad.Except
  ( ExceptT,
    MonadError (..),
    MonadIO (liftIO),
    MonadTrans (lift),
    runExceptT,
  )
import Data.Array (Array, listArray)
import Data.Complex (Complex (..), imagPart, realPart)
import Data.Functor ((<&>))
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Maybe (isJust, isNothing)
import Data.Ratio (denominator, numerator, (%))
import Numeric (readFloat, readHex, readOct)
import System.Environment (getArgs)
import System.IO
  ( Handle,
    IOMode (ReadMode, WriteMode),
    hClose,
    hFlush,
    hGetLine,
    hPrint,
    hPutStr,
    openFile,
    stderr,
    stdin,
    stdout,
  )
import Text.ParserCombinators.Parsec
  ( ParseError,
    Parser,
    alphaNum,
    anyChar,
    between,
    char,
    digit,
    endBy,
    hexDigit,
    letter,
    many,
    many1,
    noneOf,
    notFollowedBy,
    octDigit,
    oneOf,
    option,
    parse,
    sepBy,
    sepEndBy,
    skipMany,
    skipMany1,
    space,
    string,
    try,
    (<|>),
  )

arrayList :: Array i e -> [e]
arrayList = foldr (:) []

until_ :: Monad m => (a -> Bool) -> m a -> (a -> m ()) -> m ()
until_ pred prompt action = do
  result <- prompt
  if pred result
    then return ()
    else action result >> until_ pred prompt action

main :: IO ()
main = do
  args <- getArgs
  if null args then runRepl else runOne args

runOne :: [String] -> IO ()
runOne args = do
  env <-
    primitiveBindings
      >>= flip bindVars [("args", List $ map String $ drop 1 args)]
  runIOThrows (show <$> eval env (List [Atom "load", String $ head args]))
    >>= hPutStr stderr

runRepl :: IO ()
runRepl =
  primitiveBindings
    >>= until_ (== "quit") (readPrompt "Lisp>>> ") . evalAndPrint

flushStr :: String -> IO ()
flushStr str = putStr str >> hFlush stdout

readPrompt :: String -> IO String
readPrompt prompt = flushStr prompt >> getLine

evalString :: Env -> String -> IO String
evalString env expr =
  runIOThrows $ show <$> (liftThrows (readExpr expr) >>= eval env)

evalAndPrint :: Env -> String -> IO ()
evalAndPrint env expr = evalString env expr >>= putStrLn

spaces :: Parser ()
spaces = skipMany1 space

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

type Env = IORef [(String, IORef LispVal)]

nullEnv :: IO Env
nullEnv = newIORef []

primitiveBindings :: IO Env
primitiveBindings =
  nullEnv
    >>= flip
      bindVars
      ( map (makeFunc PrimitiveFunc) primitives
          ++ map (makeFunc IOFunc) ioPrimitives
      )
  where
    makeFunc constructor (var, func) = (var, constructor func)

readOrThrow :: Parser a -> String -> ThrowsError a
readOrThrow parser input = case parse parser "lisp" input of
  Left err -> throwError $ Parser err
  Right val -> return val

readExpr :: String -> ThrowsError LispVal
readExpr = readOrThrow parseExpr

readExprList :: String -> ThrowsError [LispVal]
readExprList = readOrThrow (endBy parseExpr spaces)

data LispVal
  = Nil
  | Atom String
  | List [LispVal]
  | DottedList [LispVal] LispVal
  | Vector (Array Int LispVal)
  | Number Integer
  | Float Double
  | Ratio Rational
  | Complex (Complex Double)
  | String String
  | Char Char
  | Bool Bool
  | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
  | Func
      { params :: [String],
        varargs :: Maybe String,
        body :: [LispVal],
        closure :: Env
      }
  | IOFunc ([LispVal] -> IOThrowsError LispVal)
  | Port Handle

dotted2List :: [LispVal] -> LispVal -> LispVal
dotted2List xs y = List $ xs ++ [y]

showVal :: LispVal -> String
showVal Nil = "nil"
showVal (Atom a) = a
showVal (Number n) = show n
showVal (Float f) = show f
showVal (Ratio r) = show (numerator r) ++ "/" ++ show (denominator r)
showVal (Complex c) = show (realPart c) ++ "+" ++ show (imagPart c) ++ "i"
showVal (String s) = "\"" ++ s ++ "\""
showVal (Char c) =
  "#\\" ++ case c of
    ' ' -> "space"
    '\n' -> "newline"
    _ -> [c]
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (List xs) = "(" ++ unwordsList xs ++ ")"
showVal (DottedList heads tail) =
  "(" ++ unwordsList heads ++ " . " ++ showVal tail ++ ")"
showVal (Vector v) = "#(" ++ unwordsList (foldr (:) [] v) ++ ")"
showVal (PrimitiveFunc _) = "<primitive>"
showVal (Func params varargs body _) =
  "(lambda ("
    ++ unwords (map show params)
    ++ ( case varargs of
           Nothing -> ""
           Just arg -> " . " ++ arg
       )
    ++ ") ...)"
showVal (Port _) = "<IO port>"
showVal (IOFunc _) = "<IO primitive>"

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

instance Show LispVal where
  show = showVal

data LispError
  = NumArgs Integer [LispVal]
  | TypeMismatch String LispVal
  | Parser ParseError
  | BadSpecialForm String LispVal
  | NotFunction String String
  | UnboundVar String String
  | ExpectedAtom LispVal
  | Default String

showError :: LispError -> String
showError (NumArgs expected found) =
  "Expected " ++ show expected
    ++ " arg(s), found value(s): "
    ++ unwordsList found
showError (TypeMismatch expected found) =
  "Invalid type: expected " ++ expected
    ++ ", found "
    ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (UnboundVar message varname) = message ++ ": " ++ varname
showError (ExpectedAtom found) = "Expected atom, found: " ++ show found
showError (Default message) = message

instance Show LispError where
  show = showError

type ThrowsError = Either LispError

instance MonadFail ThrowsError where
  fail message = error $ "Unexpected error: " ++ message

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

extractValue :: ThrowsError a -> a
extractValue val' = let (Right val) = val' in val

type IOThrowsError = ExceptT LispError IO

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err) = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = runExceptT (trapError action) <&> extractValue

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

escapes :: [(Char, Char)]
escapes =
  [ ('\'', '\''),
    ('"', '"'),
    ('\\', '\\'),
    ('a', '\a'),
    ('b', '\b'),
    ('f', '\f'),
    ('n', '\n'),
    ('r', '\r'),
    ('t', '\t'),
    ('v', '\v')
  ]

parseEscape :: Parser Char
parseEscape = do
  char '\\'
  c <- oneOf $ map fst escapes
  let (Just escaped) = lookup c escapes
  return escaped

parseChar :: Parser LispVal
parseChar =
  let alphabetic = do x <- anyChar; notFollowedBy alphaNum; return x
   in do
        string "#\\"
        x <-
          try $
            (string "space" >> return ' ')
              <|> (string "newline" >> return '\n')
              <|> alphabetic
        return $ Char x

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

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

parseNumber :: Parser LispVal
parseNumber = parseBin <|> parseOct <|> parseDec <|> parseHex

parseBin :: Parser LispVal
parseBin = do
  try $ string "#b"
  x <- many1 $ oneOf "01"
  let d = readBin x where readBin x = readBin' 0 x
      readBin' v [] = v
      readBin' v (x : xs) =
        let v' = v * 2 + (if x == '1' then 1 else 0)
         in readBin' v' xs
  return $ Number d

parseOct :: Parser LispVal
parseOct = do
  try $ string "#o"
  x <- many1 octDigit
  let [(d, _)] = readOct x
  return $ Number d

parseDec :: Parser LispVal
parseDec = parseDec' <|> parseDec''
  where
    parseDec' = Number . read <$> many1 digit
    parseDec'' = do
      try $ string "#d"
      parseDec'

parseHex :: Parser LispVal
parseHex = do
  try $ string "#x"
  x <- many1 hexDigit
  let [(d, _)] = readHex x
  return $ Number d

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

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

parseComplex :: Parser LispVal
parseComplex =
  let getDouble :: Parser Double
      getDouble = do
        f <- option Nil (try parseFloat)
        case f of
          (Float f) ->
            return $ realToFrac f
          _ -> do
            (Number n) <- parseDec
            return $ fromIntegral n
   in do
        x <- getDouble
        char '+'
        y <- getDouble
        char 'i'
        return $ Complex (x :+ y)

parseList :: Parser LispVal
parseList = between beg end parseList'
  where
    beg = char '(' >> skipMany space
    end = skipMany space >> char ')'
    parseList' = do
      heads <- sepEndBy parseExpr spaces
      tail <- option Nil (char '.' >> spaces >> parseExpr)
      return $ case tail of
        Nil -> List heads
        _ -> DottedList heads tail

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

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

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

parseVector :: Parser LispVal
parseVector = do
  string "#("
  arr <- sepBy parseExpr spaces
  char ')'
  let vec = listArray (0, length arr - 1) arr
  return $ Vector vec

parseExpr :: Parser LispVal
parseExpr =
  parseAtom
    <|> parseString
    <|> try parseComplex
    <|> try parseFloat
    <|> try parseRatio
    <|> try parseNumber
    <|> try parseBool
    <|> try parseChar
    <|> parseQuoted
    <|> parseQuasiquote
    <|> parseUnquoted
    <|> parseVector
    <|> parseList

eval :: Env -> LispVal -> IOThrowsError LispVal
eval env (Atom var) = getVar env var
eval _ (List [Atom "quote", val]) = return val
eval _ (List [Atom "quasiquote", val]) =
  throwError $ Default "` is not support"
eval _ (List [Atom "unquoted", val]) =
  throwError $ Default ", is not support"
eval env (List [Atom "if", pred, conseq, alt]) = do
  result <- eval env pred
  case result of
    Bool True -> eval env conseq
    Bool False -> eval env alt
    _ -> throwError $ TypeMismatch "boolean" result
eval env form@(List (Atom "cond" : clauses)) =
  if null clauses
    then
      throwError $
        BadSpecialForm "No true clause in cond expression" form
    else case head clauses of
      List [Atom "else", expr] -> eval env expr
      List [test, expr] ->
        eval env $
          List [Atom "if", test, expr, List (Atom "cond" : tail clauses)]
      _ -> throwError $ BadSpecialForm "Ill-formed cond expression" form
eval env form@(List (Atom "case" : key : clauses)) =
  if null clauses
    then
      throwError $
        BadSpecialForm "No true clause in case expression: " form
    else case head clauses of
      List (Atom "else" : exprs) -> eval' env exprs
      List (List datums : exprs) -> do
        result <- eval env key
        let equality =
              any
                ( \x ->
                    let (Right (Bool res)) =
                          eqv [result, x]
                     in res
                )
                datums
        if equality
          then eval' env exprs
          else eval env $ List (Atom "case" : key : tail clauses)
      _ -> throwError $ BadSpecialForm "Ill-formed case expression" form
  where
    eval' env exprs = mapM (eval env) exprs <&> last
eval env (List [Atom "set!", Atom var, form]) =
  eval env form >>= setVar env var
eval env (List [Atom "define", Atom var, form]) =
  eval env form >>= defineVar env var
eval env (List (Atom "define" : List (Atom var : params) : body)) =
  makeNormalFunc env params body >>= defineVar env var
eval
  env
  ( List
      (Atom "define" : DottedList (Atom var : params) varargs : body)
    ) = makeVarArgsFunc varargs env params body >>= defineVar env var
eval env (List (Atom "lambda" : List params : body)) =
  makeNormalFunc env params body
eval env (List (Atom "lambda" : DottedList params varargs : body)) =
  makeVarArgsFunc varargs env params body
eval env (List (Atom "lambda" : varargs : body)) =
  makeVarArgsFunc varargs env [] body
eval env (List [Atom "load", String filename]) =
  load filename >>= fmap last . mapM (eval env)
eval env (List (func' : args')) = do
  func <- eval env func'
  args <- mapM (eval env) args'
  apply func args
eval _ val@(List _) =
  throwError $ BadSpecialForm "Unrecognized special form" val
eval _ val@(DottedList _ _) = throwError $ Default ". is not support"
eval _ val = return val

apply :: LispVal -> [LispVal] -> IOThrowsError LispVal
apply (Func params varargs body closure) args =
  if length params /= length args && isNothing varargs
    then throwError $ NumArgs (toInteger $ length params) args
    else
      liftIO (bindVars closure $ zip params args)
        >>= bindVarArgs varargs
        >>= evalBody
  where
    evalBody env = last <$> mapM (eval env) body
    bindVarArgs arg env = case arg of
      Just argName -> liftIO $ bindVars env [(argName, List remainingArgs)]
      Nothing -> return env
    remainingArgs = drop (length params) args
apply (IOFunc func) args = func args
apply primitive args =
  let (PrimitiveFunc func) = primitive in liftThrows $ func args

primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives =
  [ ("+", numericBinop (+)),
    ("-", numericBinop (-)),
    ("*", numericBinop (*)),
    ("/", numericBinop div),
    ("mod", numericBinop mod),
    ("quotient", numericBinop quot),
    ("remainder", numericBinop rem),
    ("symbol?", unaryOp $ return . isSymbol),
    ("string?", unaryOp $ return . isString),
    ("number?", unaryOp $ return . isNumber),
    ("bool?", unaryOp $ return . isBool),
    ("list?", unaryOp $ return . isList),
    ("symbol->string", unaryOp symbol2string),
    ("string->symbol", unaryOp string2symbol),
    ("string-length", unaryOp stringLength),
    ("string-ref", stringRef),
    ("=", numBoolBinop (==)),
    ("<", numBoolBinop (<)),
    (">", numBoolBinop (>)),
    ("/=", numBoolBinop (/=)),
    (">=", numBoolBinop (>=)),
    ("<=", numBoolBinop (<=)),
    ("&&", boolBoolBinop (&&)),
    ("||", boolBoolBinop (||)),
    ("string=?", strBoolBinop (==)),
    ("string<?", strBoolBinop (<)),
    ("string>?", strBoolBinop (>)),
    ("string<=?", strBoolBinop (<=)),
    ("string>=?", strBoolBinop (>=)),
    ("car", car),
    ("cdr", cdr),
    ("cons", cons),
    ("eq?", eqv),
    ("eqv?", eqv),
    ("equal?", equal)
  ]

unaryOp ::
  (LispVal -> ThrowsError LispVal) ->
  [LispVal] ->
  ThrowsError LispVal
unaryOp func [param] = func param
unaryOp _ params = throwError $ NumArgs 1 params

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

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

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

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

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

symbol2string :: LispVal -> ThrowsError LispVal
symbol2string (Atom i) = return $ String i
symbol2string val = throwError $ TypeMismatch "atom" val

string2symbol :: LispVal -> ThrowsError LispVal
string2symbol (String s) = return $ Atom s
string2symbol val = throwError $ TypeMismatch "string" val

stringLength :: LispVal -> ThrowsError LispVal
stringLength (String s) = return . Number $ toInteger . length $ s
stringLength val = throwError $ TypeMismatch "string" val

stringRef :: [LispVal] -> ThrowsError LispVal
stringRef [String s, Number i]
  | i `elem` [0 .. toInteger (length s) - 1] =
    return . Char $ s !! fromInteger i
  | otherwise = throwError $ Default "Index out of bound"
stringRef [String _, val] = throwError $ TypeMismatch "number" val
stringRef [val, _] = throwError $ TypeMismatch "string" val
stringRef val = throwError $ NumArgs 2 val

numericBinop ::
  (Integer -> Integer -> Integer) ->
  [LispVal] ->
  ThrowsError LispVal
numericBinop _ [] = throwError $ NumArgs 2 []
numericBinop _ param@[_] = throwError $ NumArgs 2 param
numericBinop op params = mapM unpackNum params <&> Number . foldl1 op

unpackNum :: LispVal -> ThrowsError Integer
unpackNum (Number n) = return n
unpackNum val@(String s) =
  let parsed = reads s
   in if null parsed
        then throwError $ TypeMismatch "number" val
        else return $ fst . head $ parsed
unpackNum (List [n]) = unpackNum n
unpackNum val = throwError $ TypeMismatch "number" val

boolBinop ::
  (LispVal -> ThrowsError a) ->
  (a -> a -> Bool) ->
  [LispVal] ->
  ThrowsError LispVal
boolBinop unpacker op args
  | length args /= 2 = throwError $ NumArgs 2 args
  | otherwise = do
    lhs <- unpacker $ head args
    rhs <- unpacker $ last args
    return $ Bool $ lhs `op` rhs

numBoolBinop ::
  (Integer -> Integer -> Bool) ->
  [LispVal] ->
  ThrowsError LispVal
numBoolBinop = boolBinop unpackNum

strBoolBinop :: (String -> String -> Bool) -> [LispVal] -> ThrowsError LispVal
strBoolBinop = boolBinop unpackString

unpackString :: LispVal -> ThrowsError String
unpackString (String s) = return s
unpackString (Number n) = return $ show n
unpackString val = throwError $ TypeMismatch "string" val

boolBoolBinop :: (Bool -> Bool -> Bool) -> [LispVal] -> ThrowsError LispVal
boolBoolBinop = boolBinop unpackBool

unpackBool :: LispVal -> ThrowsError Bool
unpackBool (Bool b) = return b
unpackBool val = throwError $ TypeMismatch "boolean" val

car :: [LispVal] -> ThrowsError LispVal
car [List (x : xs)] = return x
car [DottedList (x : xs) _] = return x
car [val] = throwError $ TypeMismatch "pair" val
car val = throwError $ NumArgs 1 val

cdr :: [LispVal] -> ThrowsError LispVal
cdr [List (_ : xs)] = return $ List xs
cdr [DottedList [_] x] = return x
cdr [DottedList (_ : xs) x] = return $ DottedList xs x
cdr [val] = throwError $ TypeMismatch "pair" val
cdr val = throwError $ NumArgs 1 val

cons :: [LispVal] -> ThrowsError LispVal
cons [x, List xs] = return $ List (x : xs)
cons [x, DottedList xs y] = return $ DottedList (x : xs) y
cons [x, y] = return $ DottedList [x] y
cons val = throwError $ NumArgs 1 val

eqv :: [LispVal] -> ThrowsError LispVal
eqv [Nil, Nil] = return $ Bool True
eqv [Atom lhs, Atom rhs] = return $ Bool (lhs == rhs)
eqv [Number lhs, Number rhs] = return $ Bool (lhs == rhs)
eqv [Float lhs, Float rhs] = return $ Bool (lhs == rhs)
eqv [Ratio lhs, Ratio rhs] = return $ Bool (lhs == rhs)
eqv [Complex lhs, Complex rhs] = return $ Bool (lhs == rhs)
eqv [String lhs, String rhs] = return $ Bool (lhs == rhs)
eqv [Char lhs, Char rhs] = return $ Bool (lhs == rhs)
eqv [Bool lhs, Bool rhs] = return $ Bool (lhs == rhs)
eqv [Vector lhs, Vector rhs] =
  eqv [List $ arrayList lhs, List $ arrayList rhs]
eqv [DottedList lhs lhs', DottedList rhs rhs'] =
  eqv [dotted2List lhs lhs', dotted2List rhs rhs']
eqv [List lhs, List rhs] = eqList eqv lhs rhs
eqv [_, _] = return $ Bool False
eqv val = throwError $ NumArgs 2 val

eqList ::
  ([LispVal] -> ThrowsError LispVal) ->
  [LispVal] ->
  [LispVal] ->
  ThrowsError LispVal
eqList eq lhs rhs =
  return . Bool $ (length lhs == length rhs) && all eq' (zip lhs rhs)
  where
    eq' (lhs, rhs) = let (Right (Bool res)) = eq [lhs, rhs] in res

data Unpacker = forall a. Eq a => AnyUnpacker (LispVal -> ThrowsError a)

unpackers :: [Unpacker]
unpackers =
  [ AnyUnpacker unpackNum,
    AnyUnpacker unpackString,
    AnyUnpacker unpackBool
  ]

equal :: [LispVal] -> ThrowsError LispVal
equal [Vector lhs, Vector rhs] =
  equal [List $ arrayList lhs, List $ arrayList rhs]
equal [DottedList lhs lhs', DottedList rhs rhs'] =
  equal [dotted2List lhs lhs', dotted2List rhs rhs']
equal [List lhs, List rhs] = eqList equal lhs rhs
equal args@[lhs, rhs] = do
  primitiveEquals <- or <$> mapM (unpackEquals lhs rhs) unpackers
  (Bool eqvEquals) <- eqv args
  return $ Bool (primitiveEquals || eqvEquals)
equal val = throwError $ NumArgs 2 val

unpackEquals :: LispVal -> LispVal -> Unpacker -> ThrowsError Bool
unpackEquals lhs' rhs' (AnyUnpacker unpacker) =
  do
    lhs <- unpacker lhs'
    rhs <- unpacker rhs'
    return $ lhs == rhs
    `catchError` const (return False)

isBound :: Env -> String -> IO Bool
isBound envRef var = readIORef envRef <&> isJust . lookup var

getVar :: Env -> String -> IOThrowsError LispVal
getVar envRef var = do
  env <- liftIO $ readIORef envRef
  maybe
    (throwError $ UnboundVar "Getting an unbound variable" var)
    (liftIO . readIORef)
    (lookup var env)

setVar :: Env -> String -> LispVal -> IOThrowsError LispVal
setVar envRef var value = do
  env <- liftIO $ readIORef envRef
  maybe
    (throwError $ UnboundVar "Setting an unbound variable" var)
    (liftIO . (`writeIORef` value))
    (lookup var env)
  return value

defineVar :: Env -> String -> LispVal -> IOThrowsError LispVal
defineVar envRef var value = do
  alreadyDefined <- liftIO $ isBound envRef var
  if alreadyDefined
    then setVar envRef var value
    else liftIO $ do
      valueRef <- newIORef value
      env <- readIORef envRef
      writeIORef envRef ((var, valueRef) : env)
      return value

bindVars :: Env -> [(String, LispVal)] -> IO Env
bindVars envRef bindings =
  readIORef envRef
    >>= extendEnv bindings
    >>= newIORef
  where
    extendEnv bindings env = (++ env) <$> mapM addBinding bindings
    addBinding (var, value) = do
      ref <- newIORef value
      return (var, ref)

makeFunc ::
  Maybe String ->
  Env ->
  [LispVal] ->
  [LispVal] ->
  IOThrowsError LispVal
makeFunc varargs env params' body = do
  params <- mapM getParamName params'
  return $ Func params varargs body env
  where
    getParamName :: LispVal -> IOThrowsError String
    getParamName param = do
      case param of
        Atom i -> return i
        val -> throwError $ ExpectedAtom val

makeNormalFunc :: Env -> [LispVal] -> [LispVal] -> IOThrowsError LispVal
makeNormalFunc = makeFunc Nothing

makeVarArgsFunc ::
  LispVal ->
  Env ->
  [LispVal] ->
  [LispVal] ->
  IOThrowsError LispVal
makeVarArgsFunc (Atom varargs) = makeFunc $ Just varargs
makeVarArgsFunc val = \_ _ _ -> throwError $ ExpectedAtom val

ioPrimitives :: [(String, [LispVal] -> IOThrowsError LispVal)]
ioPrimitives =
  [ ("apply", applyProc),
    ("open-input-file", makePort ReadMode),
    ("open-output-file", makePort WriteMode),
    ("close-input-port", closePort),
    ("close-output-port", closePort),
    ("read", readProc),
    ("write", writeProc),
    ("read-contents", readContents),
    ("read-all", readAll)
  ]

applyProc :: [LispVal] -> IOThrowsError LispVal
applyProc [func, List args] = apply func args
applyProc (func : args) = apply func args
applyProc val = throwError $ BadSpecialForm "Ill-formed apply form" (List val)

makePort :: IOMode -> [LispVal] -> IOThrowsError LispVal
makePort mode [String filename] = Port <$> liftIO (openFile filename mode)
makePort _ [val] = throwError $ TypeMismatch "string" val
makePort _ val = throwError $ NumArgs 1 val

closePort :: [LispVal] -> IOThrowsError LispVal
closePort [Port port] = liftIO $ hClose port >> return (Bool True)
closePort _ = return $ Bool False

readProc :: [LispVal] -> IOThrowsError LispVal
readProc [] = readProc [Port stdin]
readProc [Port port] = liftIO (hGetLine port) >>= liftThrows . readExpr
readProc [val] = throwError $ TypeMismatch "port" val
readProc val =
  throwError . Default $ "Expected 0 or 1 argument, found: " ++ show val

writeProc :: [LispVal] -> IOThrowsError LispVal
writeProc [obj] = writeProc [obj, Port stdout]
writeProc [obj, Port port] = liftIO $ hPrint port obj >> return (Bool True)
writeProc [_, val] = throwError $ TypeMismatch "port" val
writeProc val =
  throwError . Default $ "Expected 0 or 1 argument, found: " ++ show val

readContents :: [LispVal] -> IOThrowsError LispVal
readContents [String filename] = String <$> liftIO (readFile filename)
readContents [val] = throwError $ TypeMismatch "string" val
readContents val = throwError $ NumArgs 1 val

load :: String -> IOThrowsError [LispVal]
load filename = liftIO (readFile filename) >>= liftThrows . readExprList

readAll :: [LispVal] -> IOThrowsError LispVal
readAll [String filename] = List <$> load filename
readAll [val] = throwError $ TypeMismatch "string" val
readAll val = throwError $ NumArgs 1 val


text:
{-# LANGUAGE NoImplicitPrelude #-}
module LinearRegression1 where

import Protolude
import Data.Vector as V
import Numeric.LinearAlgebra as H

ext: hs

max_stars_repo_name: DataHaskell/data-haskell-examples

max_issues_repo_name: DataHaskell/data-haskell-examples

max_forks_repo_name: DataHaskell/data-haskell-examples


text:
-------------------------------------------------------------------------------
-- |
-- Module    :  Data.CartPole
-- Copyright :  (c) Sentenai 2017
-- License   :  BSD3
-- Maintainer:  sam@sentenai.com
-- Stability :  experimental
-- Portability: non-portable
--
-- Shared datatypes between Gym environments and the haskell implementation of
-- CartPole.
-------------------------------------------------------------------------------
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DataKinds #-}
module Data.CartPole
  ( StateCP(..)
  , Action(..)
  , Event
  ) where

import qualified Reinforce.Spaces.State as Spaces

import Data.Aeson
import Data.Aeson.Types
import Data.Hashable
import GHC.Generics
import Control.Exception (AssertionFailed(..), throw)
import Reinforce.Spaces
import Reinforce.Spaces.Action (Size)
-- import Numeric.LinearAlgebra.Static
import qualified Data.Logger as Logger
import qualified Data.Vector as V

-- | Specific datatype for a CartPole event
type Event = Logger.Event Double StateCP Action

-- | Cartpole can only go left or right has an action space
-- of "discrete 2" containing {0..n-1}.
--
-- FIXME: Migrate this to either a more generic "directions" actions
-- (would need things like "up", "down" versions as well) or a "discrete
-- actions" version. I'm a fan of the former.
data Action
  = GoLeft
  | GoRight
  deriving (Show, Eq, Enum, Bounded, Ord, Generic)

instance Hashable Action

instance DiscreteActionSpace Action where
  type Size Action = 2

instance ToJSON Action where
  toJSON :: Action -> Value
  toJSON GoLeft  = toJSON (0 :: Int)
  toJSON GoRight = toJSON (1 :: Int)

-- | The state of a cart on a pole in a CartPole environment
data StateCP = StateCP
  { position  :: Float  -- ^ position of the cart on the track
  , angle     :: Float  -- ^ angle of the pole with the vertical
  , velocity  :: Float  -- ^ cart velocity
  , angleRate :: Float  -- ^ rate of change of the angle
  } deriving (Show, Eq, Generic, Ord)

instance Hashable StateCP

instance Semigroup StateCP where
  (StateCP a0 b0 c0 d0) <> (StateCP a1 b1 c1 d1)
    = StateCP (a0+a1) (b0+b1) (c0+c1) (d0+d1)

instance Monoid StateCP where
  mempty = StateCP 0 0 0 0
  mappend = (<>)

instance Spaces.StateSpace StateCP where
  toVector (StateCP p a v r) = Spaces.toVector [p, a, v, r]
  fromVector vec =
    case getVals of
      Nothing -> throw $ AssertionFailed "malformed vector found"
      Just s -> return s

    where
      getVals :: Maybe StateCP
      getVals = StateCP
        <$> findField 0
        <*> findField 1
        <*> findField 2
        <*> findField 3

      findField :: Int -> Maybe Float
      findField i = realToFrac <$> vec V.!? i

instance FromJSON StateCP where
  parseJSON :: Value -> Parser StateCP
  parseJSON arr@(Array _)= do
    (p, a, v, r) <- parseJSON arr :: Parser (Float, Float, Float, Float)
    return $ StateCP p a v r
  parseJSON invalid    = typeMismatch "StateCP" invalid

instance StateSpaceStatic StateCP where
  type Size StateCP = 4
  -- toR = vector . V.toList . Spaces.toVector
  -- fromR = Spaces.fromVector . V.fromList . LA.toList . unwrap

max_stars_repo_name: juliendehos/reinforce

max_issues_repo_name: sentenai/reinforce

max_forks_repo_name: sentenai/reinforce


text:
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Copyright: (c) 2021 Tom Westerhout
-- SPDX-License-Identifier: Apache-2.0
-- Maintainer: Tom Westerhout <14264576+twesterhout@users.noreply.github.com>
--
-- DLPack: Open In Memory Tensor Structure
module DLPack
  ( dlVersion,
    DLDeviceType (..),
    KnownDLDeviceType,
    -- natToInt,
    -- reifyDeviceType,
    DLDevice (..),
    DLDataTypeCode (..),
    DLDataType (..),
    DLTensor (..),
    -- DLManagedTensor (..),
    IsDLDataType (..),
    IsDLTensor (..),
    -- tensorToFlatList,
    -- fold1,
    -- loop1,
    -- foldN,
    viaContiguousBuffer,
    ContiguousVectorView (..),
    StridedVectorView (..),
    dlAsStridedVector,
    dlAsContiguousVector,
    dlExpectDType,
    dlExpectNDim,
  )
where

import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad)
import Control.Monad.ST
import Data.Complex
import Data.Functor.Identity
import Data.Int
import Data.Kind (Type)
import Data.Primitive (Prim)
import Data.Primitive.PrimArray
import qualified Data.Primitive.Ptr as P
import Data.Proxy
import Data.Text (Text)
import qualified Data.Text as T
-- import Foreign.Marshal.Utils (with)

import Data.Word
import Foreign.C.Types
import Foreign.Marshal.Array (withArrayLen)
import Foreign.Ptr (FunPtr, Ptr, castPtr, nullPtr, plusPtr)
import Foreign.Storable
import qualified GHC.Exts as GHC
import GHC.Stack (HasCallStack)
import GHC.TypeLits
import Prelude hiding (init)

natToInt :: forall n. KnownNat n => Int
natToInt = fromIntegral $ GHC.TypeLits.natVal (Proxy @n)
{-# INLINE natToInt #-}

-- | Underlying DLPack version
dlVersion :: (Int, Int)
dlVersion = (0, 5)
{-# INLINE dlVersion #-}

-- | The device type in 'DLDevice'.
data DLDeviceType
  = -- | CPU device
    DLCPU
  | -- | CUDA GPU device
    DLCUDA
  | -- | Pinned CUDA CPU memory allocated by @cudaMallocHost@
    DLCUDAHost
  | -- | OpenCL device
    DLOpenCL
  | -- | Vulkan buffer for next generation graphics
    DLVulkan
  | -- | Metal for Apple GPU
    DLMetal
  | -- | Verilog simulator buffer
    DLVPI
  | -- | ROCm GPU device
    DLROCM
  | -- | Pinned ROCm CPU memory allocated by @hipMallocHost@
    DLROCMHost
  | -- | Reserved extension device type
    DLExtDev
  | -- | CUDA managed/unified memory allocated by @cudaMallocManaged@
    DLCUDAManaged
  deriving stock (Read, Show, Eq)

class KnownDLDeviceType (device :: DLDeviceType) where
  reifyDeviceType :: DLDeviceType

instance KnownDLDeviceType 'DLCPU where
  reifyDeviceType = DLCPU
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLCUDA where
  reifyDeviceType = DLCUDA
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLCUDAHost where
  reifyDeviceType = DLCUDAHost
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLOpenCL where
  reifyDeviceType = DLOpenCL
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLVulkan where
  reifyDeviceType = DLVulkan
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLMetal where
  reifyDeviceType = DLMetal
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLVPI where
  reifyDeviceType = DLVPI
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLROCM where
  reifyDeviceType = DLROCM
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLROCMHost where
  reifyDeviceType = DLROCMHost
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLExtDev where
  reifyDeviceType = DLExtDev
  {-# INLINE reifyDeviceType #-}

instance KnownDLDeviceType 'DLCUDAManaged where
  reifyDeviceType = DLCUDAManaged
  {-# INLINE reifyDeviceType #-}

instance Enum DLDeviceType where
  toEnum x = case x of
    1 -> DLCPU
    2 -> DLCUDA
    3 -> DLCUDAHost
    4 -> DLOpenCL
    7 -> DLVulkan
    8 -> DLMetal
    9 -> DLVPI
    10 -> DLROCM
    11 -> DLROCMHost
    12 -> DLExtDev
    13 -> DLCUDAManaged
    _ -> error $ "invalid DLDeviceType: " <> show x
  {-# INLINE toEnum #-}
  fromEnum x = case x of
    DLCPU -> 1
    DLCUDA -> 2
    DLCUDAHost -> 3
    DLOpenCL -> 4
    DLVulkan -> 7
    DLMetal -> 8
    DLVPI -> 9
    DLROCM -> 10
    DLROCMHost -> 11
    DLExtDev -> 12
    DLCUDAManaged -> 13
  {-# INLINE fromEnum #-}

data DLDevice = DLDevice
  { dlDeviceType :: !DLDeviceType,
    dlDeviceId :: {-# UNPACK #-} !Int
  }
  deriving stock (Read, Show, Eq)

instance Storable DLDevice where
  sizeOf _ = 8
  {-# INLINE sizeOf #-}
  alignment _ = 4
  {-# INLINE alignment #-}
  peek p =
    DLDevice
      <$> (toEnum . (fromIntegral :: CInt -> Int) <$> peekByteOff p 0)
      <*> ((fromIntegral :: CInt -> Int) <$> peekByteOff p 4)
  {-# INLINE peek #-}
  poke p x = do
    pokeByteOff p 0 . (fromIntegral :: Int -> CInt) . fromEnum . dlDeviceType $ x
    pokeByteOff p 4 (dlDeviceId x)
  {-# INLINE poke #-}

data DLDataTypeCode
  = DLInt
  | DLUInt
  | DLFloat
  | DLOpaqueHandle
  | DLBfloat
  | DLComplex
  deriving stock (Read, Show, Eq)

instance Enum DLDataTypeCode where
  toEnum x = case x of
    0 -> DLInt
    1 -> DLUInt
    2 -> DLFloat
    3 -> DLOpaqueHandle
    4 -> DLBfloat
    5 -> DLComplex
    _ -> error $ "invalid DLDataTypeCode: " <> show x
  {-# INLINE toEnum #-}
  fromEnum x = case x of
    DLInt -> 0
    DLUInt -> 1
    DLFloat -> 2
    DLOpaqueHandle -> 3
    DLBfloat -> 4
    DLComplex -> 5
  {-# INLINE fromEnum #-}

data DLDataType = DLDataType
  { dlDataTypeCode :: !DLDataTypeCode,
    dlDataTypeBits :: {-# UNPACK #-} !Int,
    dlDataTypeLanes :: {-# UNPACK #-} !Int
  }
  deriving stock (Read, Show, Eq)

instance Storable DLDataType where
  sizeOf _ = 4
  {-# INLINE sizeOf #-}
  alignment _ = 4
  {-# INLINE alignment #-}
  peek p =
    DLDataType
      <$> (toEnum . (fromIntegral :: Word8 -> Int) <$> peekByteOff p 0)
      <*> ((fromIntegral :: Word8 -> Int) <$> peekByteOff p 1)
      <*> ((fromIntegral :: Word16 -> Int) <$> peekByteOff p 2)
  {-# INLINE peek #-}
  poke p x = do
    pokeByteOff p 0 . (fromIntegral :: Int -> Word8) . fromEnum . dlDataTypeCode $ x
    pokeByteOff p 1 . (fromIntegral :: Int -> Word8) . dlDataTypeBits $ x
    pokeByteOff p 2 . (fromIntegral :: Int -> Word16) . dlDataTypeLanes $ x
  {-# INLINE poke #-}

data DLTensor = DLTensor
  { dlTensorData :: {-# UNPACK #-} !(Ptr ()),
    dlTensorDevice :: {-# UNPACK #-} !DLDevice,
    dlTensorNDim :: {-# UNPACK #-} !Int,
    dlTensorDType :: {-# UNPACK #-} !DLDataType,
    dlTensorShape :: {-# UNPACK #-} !(Ptr Int64),
    dlTensorStrides :: {-# UNPACK #-} !(Ptr Int64),
    dlTensorByteOffset :: {-# UNPACK #-} !Word64
  }
  deriving stock (Show, Eq)

instance Storable DLTensor where
  sizeOf _ = 8 + 8 + 4 + 4 + 8 + 8 + 8
  {-# INLINE sizeOf #-}
  alignment _ = 8
  {-# INLINE alignment #-}
  peek p =
    DLTensor
      <$> peekByteOff p 0
      <*> peekByteOff p 8
      <*> ((fromIntegral :: CInt -> Int) <$> peekByteOff p 16)
      <*> peekByteOff p 20
      <*> peekByteOff p 24
      <*> peekByteOff p 32
      <*> peekByteOff p 40
  {-# INLINE peek #-}
  poke p x = do
    pokeByteOff p 0 (dlTensorData x)
    pokeByteOff p 8 (dlTensorDevice x)
    pokeByteOff p 16 . (fromIntegral :: Int -> CInt) . dlTensorNDim $ x
    pokeByteOff p 20 (dlTensorDType x)
    pokeByteOff p 24 (dlTensorShape x)
    pokeByteOff p 32 (dlTensorStrides x)
    pokeByteOff p 40 (dlTensorByteOffset x)
  {-# INLINE poke #-}

class IsDLDataType a where
  dlDataTypeOf :: proxy a -> DLDataType

fromTag :: forall a proxy. Storable a => DLDataTypeCode -> proxy a -> DLDataType
fromTag tag _ =
  DLDataType {dlDataTypeCode = tag, dlDataTypeBits = (sizeOf element * 8), dlDataTypeLanes = 1}
  where
    element :: a
    element = undefined
{-# INLINE fromTag #-}

-- {{{ Integral

instance IsDLDataType CChar where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CSChar where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUChar where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CShort where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUShort where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CInt where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUInt where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CLong where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CULong where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CPtrdiff where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CSize where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CLLong where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CULLong where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CBool where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CIntPtr where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CUIntPtr where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word8 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word16 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word32 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word64 where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Word where
  dlDataTypeOf = fromTag DLUInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int8 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int16 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int32 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int64 where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Int where
  dlDataTypeOf = fromTag DLInt
  {-# INLINE dlDataTypeOf #-}

-- }}}

-- {{{ Floating point

instance IsDLDataType CFloat where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Float where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType CDouble where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType Double where
  dlDataTypeOf = fromTag DLFloat
  {-# INLINE dlDataTypeOf #-}

-- }}}

-- {{{ Complex

instance IsDLDataType (Complex CFloat) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType (Complex Float) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType (Complex CDouble) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

instance IsDLDataType (Complex Double) where
  dlDataTypeOf = fromTag DLComplex
  {-# INLINE dlDataTypeOf #-}

-- }}}

class (Monad m) => IsDLTensor m a where
  withDLTensor :: a -> (DLTensor -> m b) -> m b

data ContiguousVectorView a
  = ContiguousVectorView {-# UNPACK #-} !Int {-# UNPACK #-} !(Ptr a)

data StridedVectorView a
  = StridedVectorView {-# UNPACK #-} !Int {-# UNPACK #-} !Int {-# UNPACK #-} !(Ptr a)

dlExpectDType :: forall a. IsDLDataType a => DLTensor -> Either Text ()
dlExpectDType t
  | expected == got = Right ()
  | otherwise =
    Left . T.pack $
      "expected DLTensor to be of type " <> show expected
        <> ", but it is of type "
        <> show got
  where
    !expected = dlDataTypeOf (Proxy @a)
    !got = dlTensorDType t
{-# INLINE dlExpectDType #-}

dlExpectDevice :: DLDevice -> DLTensor -> Either Text ()
dlExpectDevice expected t
  | expected == got = Right ()
  | otherwise =
    Left . T.pack $
      "expected the tensor to be on " <> show expected <> ", but it is on "
        <> show got
  where
    !got = dlTensorDevice t
{-# INLINE dlExpectDevice #-}

dlExpectNDim :: Int -> DLTensor -> Either Text ()
dlExpectNDim expected t
  | expected == got = Right ()
  | otherwise =
    Left . T.pack $
      "expected a " <> show expected <> "-dimensional tensor, but got a "
        <> show got
        <> "-dimensional one"
  where
    !got = dlTensorNDim t
{-# INLINE dlExpectNDim #-}

dlTensorLength :: DLTensor -> Int
dlTensorLength t
  | rank > 0 = fromIntegral $ go (P.indexOffPtr (dlTensorShape t) 0) 1
  | otherwise = 0
  where
    !rank = dlTensorNDim t
    go !acc !i
      | i < rank = go (acc * P.indexOffPtr (dlTensorShape t) i) (i + 1)
      | otherwise = acc
{-# INLINE dlTensorLength #-}

dlExpectShape :: PrimArray Int -> DLTensor -> Either Text ()
dlExpectShape expected t = do
  () <- dlExpectNDim rank t
  let go !i
        | i < rank =
          indexPrimArray expected i == fromIntegral (P.indexOffPtr (dlTensorShape t) i)
            && go (i + 1)
        | otherwise = True
  if go 0
    then Right ()
    else
      Left . T.pack $
        "expected a tensor of shape " <> show expected
          <> ", but got a tensor of shape "
          <> show shape
  where
    !rank = sizeofPrimArray expected
    shape = generatePrimArray (dlTensorNDim t) (\i -> P.indexOffPtr (dlTensorShape t) i)
{-# INLINE dlExpectShape #-}

dlAsStridedVector :: forall a. IsDLDataType a => DLTensor -> Either Text (StridedVectorView a)
dlAsStridedVector t = do
  () <- dlExpectDType @a t
  () <- dlExpectNDim 1 t
  let !size = fromIntegral $ P.indexOffPtr (dlTensorShape t) 0
      !stride = fromIntegral $ P.indexOffPtr (dlTensorStrides t) 0
  pure $ StridedVectorView size stride (castPtr (dlTensorData t))
{-# INLINE dlAsStridedVector #-}

dlAsContiguousVector :: forall a. IsDLDataType a => DLTensor -> Either Text (ContiguousVectorView a)
dlAsContiguousVector t = do
  (StridedVectorView size stride ptr) <- dlAsStridedVector t
  if stride == 1
    then Right $ ContiguousVectorView size ptr
    else Left . T.pack $ "expected a contiguous DLTensor, but got one with stride " <> show stride
{-# INLINE dlAsContiguousVector #-}

-- dlRowMajorStride :: DLTensor -> Maybe Int
-- dlRowMajorStride t = let !f = go 0 in f
--   where
--     go !i
--       | i < rank - 1 =
--         let !n₂ = P.indexOffPtr shape (i + 1)
--             !s₁ = P.indexOffPtr stride i
--             !s₂ = P.indexOffPtr stride (i + 1)
--          in if s₁ `div` s₂ == n₁ then go (i + 1) else False
--       | otherwise = True

-- data ShapeInfo = ShapeInfo {shapeSizes :: {-# UNPACK #-} !(PrimArray Int), shapeStrides :: {-# UNPACK #-} !(PrimArray Int)}
--   deriving stock (Show, Eq)

dlRowMajorStride :: DLTensor -> Maybe Int
dlRowMajorStride t
  | rank == 0 = Just 1
  | otherwise = fromIntegral <$> go (P.indexOffPtr stride 0) 1
  where
    !rank = dlTensorNDim t
    !shape = dlTensorShape t
    !stride = dlTensorStrides t
    go !s !i
      | i < rank =
        let !n = P.indexOffPtr shape i
            !s' = P.indexOffPtr stride i
         in if s `div` s' == n then go s' (i + 1) else Nothing
      | otherwise = Just s

--
-- isColumnMajorLike :: Int -> Ptr Int64 -> Ptr Int64 -> Bool
-- isColumnMajorLike !rank !shape !stride = let !f = go 0 in f
--   where
--     go !i
--       | i < rank - 1 =
--         let !n₁ = P.indexOffPtr shape i
--             !s₁ = P.indexOffPtr stride i
--             !s₂ = P.indexOffPtr stride (i + 1)
--          in if s₂ `div` s₁ == n₁ then go (i + 1) else False
--       | otherwise = True

-- toList1D :: forall a. Prim a => DLTensor 'DLCPU 1 a -> [a]
-- toList1D t = runST $ go [] (stride * (extent - 1))
--   where
--     !extent = dlExtent t 0
--     !stride = dlStride t 0
--     !p = dlTensorData t `plusPtr` fromIntegral (dlTensorByteOffset t)
--     go :: PrimMonad m => [a] -> Int -> m [a]
--     go acc !i
--       | i >= 0 = do !x <- P.readOffPtr p i; go (x : acc) (i - stride)
--       | otherwise = pure acc
--
-- chunksOf :: HasCallStack => Int -> [a] -> [[a]]
-- chunksOf k xs = undefined

-- tensorToFlatList :: forall a device rank. Prim a => DLTensor device rank a -> [a]
-- tensorToFlatList t = reverse $ runIdentity $ foldN 0 shape strides combine []
--   where
--     buildList p = (\i -> fromIntegral $ P.indexOffPtr p i) <$> [0 .. dlTensorNDim t - 1]
--     strides = buildList (dlTensorStrides t)
--     shape = buildList (dlTensorShape t)
--     combine xs !i =
--       let !p = dlTensorData t `plusPtr` fromIntegral (dlTensorByteOffset t)
--           !x = P.indexOffPtr (castPtr p) i
--        in return (x : xs)

withListN :: (Prim a, PrimMonad m) => Int -> [a] -> (Ptr a -> m b) -> m b
withListN n xs action = do
  array <- newAlignedPinnedPrimArray n
  let go i []
        | i == n = return ()
        | otherwise =
          error $
            "list is shorter than indicated: " <> show i <> " < " <> show n
      go !i (y : ys)
        | i < n = writePrimArray array i y >> go (i + 1) ys
        | otherwise =
          error $
            "list is shorter than indicated: " <> show (length xs) <> " > " <> show n
  go 0 xs
  action (mutablePrimArrayContents array)

viaContiguousBuffer ::
  (IsDLDataType a, PrimMonad m) =>
  Ptr a ->
  [Int] ->
  [Int] ->
  (DLTensor -> m b) ->
  m b
viaContiguousBuffer dataPtr shape strides action =
  withListN rank (fromIntegral <$> shape) $ \shapePtr ->
    withListN rank (fromIntegral <$> strides) $ \stridesPtr ->
      action
        DLTensor
          { dlTensorData = castPtr dataPtr,
            dlTensorDevice = DLDevice DLCPU 0,
            dlTensorNDim = rank,
            dlTensorDType = dlDataTypeOf dataPtr,
            dlTensorShape = shapePtr,
            dlTensorStrides = stridesPtr,
            dlTensorByteOffset = 0
          }
  where
    !rank = length shape
{-# INLINE viaContiguousBuffer #-}

class
  (Monad m, KnownNat d, IsDLDataType (ListElement a)) =>
  ListIsDLTensor (m :: Type -> Type) (d :: Nat) (a :: Type)
  where
  listWithDLTensor :: proxy d -> a -> (DLTensor -> m b) -> m b

listShape2 :: [[a]] -> (Int, Int)
listShape2 [] = (0, 0)
listShape2 xs@(x : _) = (length xs, length x)
{-# INLINE listShape2 #-}

listShape3 :: [[[a]]] -> (Int, Int, Int)
listShape3 [] = (0, 0, 0)
listShape3 xs@(x : _) = let (d₁, d₂) = listShape2 x in (length xs, d₁, d₂)
{-# INLINE listShape3 #-}

listShape4 :: [[[[a]]]] -> (Int, Int, Int, Int)
listShape4 [] = (0, 0, 0, 0)
listShape4 xs@(x : _) = let (d₁, d₂, d₃) = listShape3 x in (length xs, d₁, d₂, d₃)
{-# INLINE listShape4 #-}

listShape5 :: [[[[[a]]]]] -> (Int, Int, Int, Int, Int)
listShape5 [] = (0, 0, 0, 0, 0)
listShape5 xs@(x : _) = let (d₁, d₂, d₃, d₄) = listShape4 x in (length xs, d₁, d₂, d₃, d₄)
{-# INLINE listShape5 #-}

instance
  (PrimMonad m, ListElement [a] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 1 [a]
  where
  listWithDLTensor _ xs action = withListN d₀ xs $ \ptr ->
    viaContiguousBuffer ptr [d₀] [1] action
    where
      d₀ = length xs

instance
  (PrimMonad m, ListElement [[a]] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 2 [[a]]
  where
  listWithDLTensor _ xs action =
    withListN (d₀ * d₁) (mconcat xs) $ \ptr ->
      viaContiguousBuffer ptr [d₀, d₁] [d₁, 1] action
    where
      (d₀, d₁) = listShape2 xs

instance
  (PrimMonad m, ListElement [[[a]]] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 3 [[[a]]]
  where
  listWithDLTensor _ xs action =
    withListN (d₀ * d₁ * d₂) (mconcat . mconcat $ xs) $ \ptr ->
      viaContiguousBuffer ptr [d₀, d₁, d₂] [d₁ * d₂, d₂, 1] action
    where
      (d₀, d₁, d₂) = listShape3 xs

instance
  (PrimMonad m, ListElement [[[[a]]]] ~ a, IsDLDataType a, Prim a) =>
  ListIsDLTensor m 4 [[[[a]]]]
  where
  listWithDLTensor _ xs action =
    withListN (d₀ * d₁ * d₂ * d₃) (mconcat . mconcat . mconcat $ xs) $ \ptr ->
      viaContiguousBuffer ptr [d₀, d₁, d₂, d₃] [d₁ * d₂ * d₃, d₂ * d₃, d₃, 1] action
    where
      (d₀, d₁, d₂, d₃) = listShape4 xs

type family IsFinal a where
  IsFinal [a] = 'False
  IsFinal a = 'True

type family ListElement' a b where
  ListElement' a 'True = a
  ListElement' [a] 'False = ListElement' a (IsFinal a)

type ListElement (a :: Type) = ListElement' a (IsFinal a)

type family ListDimension' (n :: Nat) (a :: Type) where
  ListDimension' n [a] = ListDimension' (n + 1) a
  ListDimension' n a = n

type ListDimension (a :: Type) = ListDimension' 0 a

instance (Monad m, ListIsDLTensor m (ListDimension [t]) [t]) => IsDLTensor m [t] where
  withDLTensor = listWithDLTensor (Proxy :: Proxy (ListDimension [t]))

-- listShape2D :: [[a]] -> [Int]
-- listShape2D [] = [0, 0]
-- listShape2D xs@(x : _) = [length xs, length x]
--
-- listShape3D :: [[[a]]] -> [Int]
-- listShape3D [] = [0, 0, 0]
-- listShape3D xs@(x : _) = length xs : listShape2D x
--
-- listShape4D :: [[[a]]] -> [Int]
-- listShape4D [] = [0, 0, 0, 0]
-- listShape4D xs@(x : _) = length xs : listShape3D x
--
-- listShape5D :: [[[[a]]]] -> [Int]
-- listShape5D [] = [0, 0, 0, 0, 0]
-- listShape5D xs@(x : _) = length xs : listShape4D x

-- instance Prim a => GHC.IsList (DLTensor 'DLCPU 1 a) where
--   type Item (DLTensor 'DLCPU 1 a) = a
--   toList = toList1D
--   fromList _ = error "IsList instance of DLTensor does not implement toList"

-- instance Prim a => GHC.IsList (DLTensor 'DLCPU 2 a) where
--   type Item (DLTensor 'DLCPU 2 a) = [a]
--   toList t = case dlFlatten t of
--     Just flat -> chunksOf k . toList1D $ flat
--     Nothing -> error "IsList instance of DLTensor does not support weirdly strided tensors"
--     where
--       !k = dlExtent t 1
--   fromList _ = error "IsList instance of DLTensor does not implement toList"

-- instance Prim a => GHC.IsList (DLTensor 'DLCPU 3 a) where
--   type Item (DLTensor 'DLCPU 3 a) = [[a]]
--   toList t = case dlFlatten t of
--     Just flat -> chunksOf d₁ . chunksOf d₂ . toList1D $ flat
--     Nothing -> error "IsList instance of DLTensor does not support weirdly strided tensors"
--     where
--       !d₁ = dlExtent t 1
--       !d₂ = dlExtent t 2
--   fromList _ = error "IsList instance of DLTensor does not implement toList"

-- instance (IsDLDataType a, Storable a) => IsDLTensor IO 'DLCPU 1 a [a] where
--   withDLTensor xs action = withArrayLen xs $ \size dataPtr ->
--     viaContiguousBuffer dataPtr [size] [1] action
--   {-# INLINE withDLTensor #-}
--
-- instance (IsDLDataType a, Storable a) => IsDLTensor IO 'DLCPU 2 a [[a]] where
--   withDLTensor l action = do
--     let (!n, !m) = listShape2D l
--     withArrayLen (mconcat l) $ \size dataPtr -> do
--       when (n * m /= size) $ error "list has wrong shape"
--       viaContiguousBuffer dataPtr [n, m] [m, 1] action
--   {-# INLINE withDLTensor #-}
--
-- instance (IsDLDataType a, Storable a) => IsDLTensor IO 'DLCPU 3 a [[[a]]] where
--   withDLTensor l action = do
--     let (!d₁, !d₂, !d₃) = listShape3D l
--     withArrayLen (mconcat (mconcat l)) $ \size dataPtr -> do
--       when (d₁ * d₂ * d₃ /= size) $ error "list has wrong shape"
--       viaContiguousBuffer dataPtr [d₁, d₂, d₃] [d₂ * d₃, d₃, 1] action
--   {-# INLINE withDLTensor #-}

max_stars_repo_name: twesterhout/dlpack-hs

max_issues_repo_name: twesterhout/dlpack-hs

max_forks_repo_name: twesterhout/dlpack-hs


text:
{-#LANGUAGE BangPatterns #-}

module Tests.Correlation
    ( tests ) where

import Control.Arrow (Arrow(..))
import qualified Data.Vector as V
import Statistics.Correlation
import Statistics.Correlation.Kendall
import Test.QuickCheck ((==>),Property,counterexample)
import Test.Tasty
import Test.Tasty.QuickCheck
import Test.Tasty.HUnit

import Tests.ApproxEq

----------------------------------------------------------------
-- Tests list
----------------------------------------------------------------

tests :: TestTree
tests = testGroup "Correlation"
    [ testProperty "Pearson correlation"           testPearson
    , testProperty "Spearman correlation is scale invariant" testSpearmanScale
    , testProperty "Spearman correlation, nonlinear"         testSpearmanNonlinear
    , testProperty "Kendall test -- general"       testKendall
    , testCase     "Kendall test -- special cases" testKendallSpecial
    ]

----------------------------------------------------------------
-- Pearson's correlation
----------------------------------------------------------------

testPearson :: [(Double,Double)] -> Property
testPearson sample
  = (length sample > 1) ==> (exact ~= fast)
  where
    (~=) = eql 1e-12
    exact = exactPearson $ map (realToFrac *** realToFrac) sample
    fast  = pearson $ V.fromList sample

exactPearson :: [(Rational,Rational)] -> Double
exactPearson sample
  = realToFrac cov / sqrt (realToFrac (varX * varY))
  where
    (xs,ys) = unzip sample
    n       = fromIntegral $ length sample
    -- Mean
    muX  = sum xs / n
    muY  = sum ys / n
    -- Mean of squares
    muX2 = sum (map (\x->x*x) xs) / n
    muY2 = sum (map (\x->x*x) ys) / n
    -- Covariance
    cov  = sum (zipWith (*) [x - muX | x<-xs] [y - muY | y<-ys]) / n
    varX = muX2 - muX*muX
    varY = muY2 - muY*muY

----------------------------------------------------------------
-- Spearman's correlation
----------------------------------------------------------------

-- Test that Spearman correlation is scale invariant
testSpearmanScale :: [(Double,Double)] -> Double -> Property
testSpearmanScale xs a
  = and [ length xs > 1       -- Enough to calculate underflow
        , a /= 0
        , not (isNaN c1)
        , not (isNaN c2)
        , not (isNaN c3)
        , not (isNaN c4)
        ]
  ==> ( counterexample (show xs2)
      $ counterexample (show xs3)
      $ counterexample (show xs4)
      $ counterexample (show (c1,c2,c3,c4))
      $ and [ c1 == c4
           , c1 == signum a * c2
           , c1 == signum a * c3
           ]
      )
  where
    xs2 = map ((*a) *** id  ) xs
    xs3 = map (id   *** (*a)) xs
    xs4 = map ((*a) *** (*a)) xs
    c1 = spearman $ V.fromList xs
    c2 = spearman $ V.fromList xs2
    c3 = spearman $ V.fromList xs3
    c4 = spearman $ V.fromList xs4

-- Test that Spearman correlation allows to transform sample with
testSpearmanNonlinear :: [(Double,Double)] -> Property
testSpearmanNonlinear sample0
  = and [ length sample0 > 1
        , not (isNaN c1)
        , not (isNaN c2)
        , not (isNaN c3)
        , not (isNaN c4)
        ]
  ==> ( counterexample (show sample0)
      $ counterexample (show sample1)
      $ counterexample (show sample2)
      $ counterexample (show sample3)
      $ counterexample (show sample4)
      $ counterexample (show (c1,c2,c3,c4))
      $ and [ c1 == c2
            , c1 == c3
            , c1 == c4
            ]
      )
  where
    -- We need to stretch sample into [-10 .. 10] range to avoid
    -- problems with under/overflows etc.
    stretch xs
      | a == b = xs
      | otherwise = [ (x - a - 10) * 20 / (a - b) | x <- xs ]
      where
        a = minimum xs
        b = maximum xs
    sample1 = uncurry zip $ (stretch *** stretch) $ unzip sample0
    sample2 = map (exp *** id ) sample1
    sample3 = map (id  *** exp) sample1
    sample4 = map (exp *** exp) sample1
    c1 = spearman $ V.fromList sample1
    c2 = spearman $ V.fromList sample2
    c3 = spearman $ V.fromList sample3
    c4 = spearman $ V.fromList sample4

----------------------------------------------------------------
-- Kendall's correlation
----------------------------------------------------------------

testKendall :: [(Double, Double)] -> Bool
testKendall xy | isNaN r1 = isNaN r2
               | otherwise = r1 == r2
  where
    r1 = kendallBruteForce xy
    r2 = kendall $ V.fromList xy

testKendallSpecial :: Assertion
testKendallSpecial = vs @=? map (\(xs, ys) -> kendall $ V.fromList $ zip xs ys) d
  where
    (d, vs) = unzip testData
    testData :: [(([Double], [Double]), Double)]
    testData = [ (([1, 2, 3, 1, 2], [1, 2, 1, 5, 2]), -0.375)
               , (([1, 1, 1, 3, 3], [3, 3, 3, 2, 5]), 0)
               ]

kendallBruteForce :: [(Double, Double)] -> Double
kendallBruteForce xy = (n_c - n_d) / sqrt ((n_0 - n_1) * (n_0 - n_2))
  where
    allPairs = f xy
    (n_c, n_d, n_1, n_2) = foldl g (0,0,0,0) allPairs
    n_0 = fromIntegral.length $ allPairs
    g (!nc, !nd, !n1, !n2) ((x1, y1), (x2, y2))
      | (x2 - x1) * (y2 - y1) > 0 = (nc+1, nd, n1, n2)
      | (x2 - x1) * (y2 - y1) < 0 = (nc, nd+1, n1, n2)
      | otherwise = if x1 == x2
                       then if y1 == y2
                               then (nc, nd, n1+1, n2+1)
                               else (nc, nd, n1+1, n2)
                       else (nc, nd, n1, n2+1)
    f (x:xs) = zip (repeat x) xs ++ f xs
    f _ = []


text:
{-# LANGUAGE RankNTypes #-}
import Numeric.ADLens.Lens
import Numeric.ADLens.Basic
import Numeric.ADLens.List
import Numeric.ADLens.HMatrix
import Numeric.ADLens.Accelerate as Acc

import Numeric.LinearAlgebra

import Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.Native  as CPU

{-
dostuff = CPU.run $ dotp (use xs) (use ys)

fromList (Z:.10) [0..]
-}

type L1 = Matrix Double
type L2 = Matrix Double
type L3 = Matrix Double

type Input = Vector Double
type Output = Vector Double
type Weights = (L1,(L2,(L3,())))

class TupleSum a where
	tupsum :: a -> a -> a
instance TupleSum () where
	tupsum _ _ = ()
instance (Num a, TupleSum b) => TupleSum (a,b) where
	tupsum (a,x) (b,y) = (a + b, tupsum x y)

-- A dense relu neural network example
swaplayer :: Lens' ((Matrix t, b), Vector t) (b, (Matrix t, Vector t))
swaplayer = first' swap' . assoc'

mmultlayer :: Numeric t => Lens' (b, (Matrix t, Vector t)) (b, Vector t)
mmultlayer = second' mdot'

relulayer :: Lens' (b, Vector Double) (b, Vector Double)
relulayer = second' $ cmap' relu'

uselayer :: Lens' ((Matrix Double, b), Vector Double) (b, Vector Double)
uselayer = swaplayer . mmultlayer . relulayer

runNetwork :: Lens' (Weights, Input) ((), Output)
runNetwork =  uselayer . uselayer . uselayer

main :: IO ()
main = do
   putStrLn "Starting Tests"
   print $ grad (pow' 2) 1
   print $ grad (pow' 4) 1
   print $ grad (map' (pow' 2) . sum') $ [1 .. 5]
   print $ grad (map' (pow' 4) . sum') $ [1 .. 5]
   print $ map (\x -> 4 * x ^ 3 )  [1 .. 5]
   l1 <- randn 3 4
   l2 <- randn 2 3
   l3 <- randn 1 2
   let weights = (l1,(l2,(l3,())))
   print $ view runNetwork (weights, vector [1,2,3,4])
   putStrLn "The neural network gradients"
   print $ set runNetwork (weights, vector [1,2,3,4]) ((), vector [1])


text:
import Mnist
import SampleWeight
import Numeric.LinearAlgebra

main = do
    sw <- loadSW
    print $ sw !! 3
    putStrLn "Done"


text:
{-# LINE 8 "SpectrumAnalysis.lhs" #-}
--  This code was automatically generated by lhs2tex --code, from the file
--  HSoM/SpectrumAnalysis.lhs.  (See HSoM/MakeCode.bat.)
{-# LINE 18 "SpectrumAnalysis.lhs" #-}
{-#  LANGUAGE Arrows  #-}

module Euterpea.Music.Signal.SpectrumAnalysis where

import Euterpea
import Euterpea.Experimental (fftA)

import Data.Complex (Complex ((:+)), polar)
import Data.Maybe (listToMaybe, catMaybes)
{-# LINE 560 "SpectrumAnalysis.lhs" #-}
dft :: RealFloat a => [Complex a] -> [Complex a]
dft xs =
  let  lenI = length xs
       lenR = fromIntegral lenI
       lenC = lenR :+ 0
  in [  let i = -2 * pi * fromIntegral k / lenR
        in (1/lenC) * sum [  (xs!!n) * exp (0 :+ i * fromIntegral n)
                             | n <- [0,1..lenI-1] ]
        | k <- [0,1..lenI-1] ]
{-# LINE 583 "SpectrumAnalysis.lhs" #-}
mkTerm :: Int -> Double -> [Complex Double]
mkTerm num n = let f = 2 * pi / fromIntegral num
               in [  sin (n * f * fromIntegral i) / n :+ 0
                     | i <- [0,1..num-1] ]

mkxa, mkxb, mkxc :: Int-> [Complex Double]
mkxa num = mkTerm num 1
mkxb num = zipWith (+) (mkxa num) (mkTerm num 3)
mkxc num = zipWith (+) (mkxb num) (mkTerm num 5)
{-# LINE 610 "SpectrumAnalysis.lhs" #-}
printComplexL :: [Complex Double] -> IO ()
printComplexL xs  =
  let  f (i,rl:+im) =
            do  putStr (spaces (3 - length (show i))  )
                putStr (show i       ++ ":  ("        )
                putStr (niceNum rl  ++ ", "           )
                putStr (niceNum im  ++ ")\n"          )
  in mapM_ f (zip [0..length xs - 1] xs)

niceNum :: Double -> String
niceNum d =
  let  d' = fromIntegral (round (1e10 * d)) / 1e10
       (dec, fra)  = break (== '.') (show d')
       (fra',exp)  = break (== 'e') fra
  in  spaces (3  - length dec) ++ dec ++ take 11 fra'
      ++ exp ++ spaces (12 - length fra' - length exp)

spaces :: Int -> String
spaces  n = take n (repeat ' ')
{-# LINE 679 "SpectrumAnalysis.lhs" #-}
mkPulse :: Int -> [Complex Double]
mkPulse n = 100 : take (n-1) (repeat 0)
{-# LINE 721 "SpectrumAnalysis.lhs" #-}
x1 num = let f = pi * 2 * pi / fromIntegral num
         in map (:+ 0) [  sin (f * fromIntegral i)
                          | i <- [0,1..num-1] ]
{-# LINE 757 "SpectrumAnalysis.lhs" #-}
mkPolars :: [Complex Double] -> [Complex Double]
mkPolars = map ((\(m,p)-> m:+p) . polar)

max_stars_repo_name: bsdr/Euterpea


text:
-- perceptron simulation code
-- created by Tomohiro KIDA 2018/05/08/Tue

module Main
    ( main
    , linearFunc
    , linearFunc'
    , lossFunc
    , gradientDescent
    , plotDataPng
    , simPerceptron
    ) where

import Graphics.Gnuplot.Simple
import System.Random
import Numeric.LinearAlgebra

xxs :: [[Double]]
xxs = [[1,4,6],[1,1,4],[1,3,2],[1,7,3]]
ts :: [Int]
ts = [-1,-1,1,1]
a :: (Fractional frac) => frac
a  = 0.1

-- f (x) = w (tr x)
linearFunc :: Matrix R -> Matrix R -> Double
linearFunc w x
    | fst (size w) == 1 && fst (size x) == 1 = head . head . toLists $ w <> (tr x)
    | otherwise                              = -1 -- error

-- linear function for plot
-- the left hand side is x2
linearFunc' :: Matrix R -> Double -> Double
linearFunc' ws x =
    let w = (toLists ws)!!0
    in (-x)*(w!!1)/(w!!2) - ((w!!0)/(w!!2))

lossFunc :: Double -> Double -> Double
lossFunc t f = max 0.0 (-t*f)

-- return one itaration of gradientDescent
gradientDescent :: Int -> Matrix R -> Matrix R
gradientDescent rn ws =
    let ti      = realToFrac $ ts !! rn
        xi      = fromLists [xxs!!rn] :: Matrix R
        li      = lossFunc ti $ linearFunc ws xi
        wsn     = if li > 0 then ws - ( a*(realToFrac (-ti))*xi ) else ws
    in wsn

-- plot data and png
plotDataPng :: Matrix R -> IO ()
plotDataPng ws = do
    plotPathsStyle atribute plotstyle
    where
        lines     = map (\x->(x,(linearFunc' ws x))) [0..10]
        points_t  = map ((\[x,y]->(x,y)).tail) $ (\xs->[(xs!!0),(xs!!1)]) xxs
        points_f  = map ((\[x,y]->(x,y)).tail) $ (\xs->[(xs!!2),(xs!!3)]) xxs
        label     = [(XLabel "x"), (YLabel "y")]
        range     = [(XRange (0,8)), (YRange (0,8))]
        size      = [Aspect (Ratio 0.7)]
        title     = [(Title "simPerceptron")]
        png       = [(PNG   "simPerceptron.png")]
        atribute  = (label++range++title++png)
        plotstyle = [(defaultStyle {plotType = Points }, points_t),
                     (defaultStyle {plotType = Points }, points_f),
                     (defaultStyle {plotType = Lines  }, lines )]

-- simulation perceptron by grandient descent
simPerceptron :: StdGen -> Matrix R -> Int -> IO ()
simPerceptron gen ws count = do
    let (rn,ng) = randomR (0,3) gen :: (Int, StdGen) -- get one rand data of 4
    let wsn = gradientDescent rn ws -- update ws
    putStrLn $ show count ++ show wsn
    if (count > 0)
        then simPerceptron ng wsn (count-1) -- return recursion
        else plotDataPng wsn -- plot

main :: IO ()
main = do
    gen <- getStdGen
    let ws = (1><3) [-1,2,3] :: Matrix R
    simPerceptron gen ws 100 -- gradientDescent simulation

max_stars_repo_name: TomohiroKida/perceptron

max_issues_repo_name: TomohiroKida/perceptron

max_forks_repo_name: TomohiroKida/perceptron


text:
import Numeric.Container hiding (linspace)
import Numeric.LinearAlgebra.Util
import Numeric.LinearAlgebra.Data hiding (linspace)
import Numeric.LinearAlgebra.Algorithms
import Numeric.LinearAlgebra.Devel
import Data.List hiding (find)
import DSP.Basic

--For Charts
import Graphics.Rendering.Chart hiding (Matrix, Vector)
import Data.Colour
import Data.Colour.Names
import Data.Default.Class
import Graphics.Rendering.Chart.Backend.Cairo
import Control.Lens

{- 1.1.1 Method of Finite Differences
# Eigenvalue problem in one dimension. -}

--Create the tridiagonal differential matrix, with corresponding sparse matrix, if needed.
d :: Matrix Double
d = buildMatrix 100 100 (\(i,j) -> if (i,j) == (i,i) then 2 else if (i,j) == (i, i-1) then -1 else if (i,j) == (i,i+1) then -1 else 0)
d_Sparse = mkSparse [((i,j), (d @@> (i,j))) | (i,j) <- (find (/= 0) d)]

--Create the potential matrix with corresponding sparse matrix, if needed.
v :: Matrix Double
v = buildMatrix 100 100 (\(i,j) -> ((fromIntegral i)+50)^2)
v_Sparse = mkSparse [((i,j), (d @@> (i,j))) | (i,j) <- (find (/= 0) v)]

--The energy eigenvalues in a sorted List:
e = sort(toList(mapVector realPart (fst(eig(d_Sparse+v_Sparse)))))

--The corresponding eigenfunctions
psi = toColumns (mapMatrix realPart (snd(eig(d_Sparse+v_Sparse))))

--Plot
chart = toRenderable layout
  where

	--Grid for Visualization
	x = linspace 0.0 1.0 100
	--List of Tuples [(Grid, psi) ] for Plotting
	points1 = zip ( x) (map (\x -> abs(x)^2) (toList (psi !! 1)))
	points2 = zip ( x) (map (\x -> abs(x)^2) (toList (psi !! 2)))
	points3 = zip ( x) (map (\x -> abs(x)^2) (toList (psi !! 3)))
	points4 = zip ( x) (map (\x -> abs(x)^2) (toList (psi !! 4)))
	psi1 = plot_lines_values .~ [points1]
              $ plot_lines_style  . line_color .~ opaque blue
              $ plot_lines_title .~ "psi1"
              $ def
	psi2 = plot_lines_values .~ [points2]
              $ plot_lines_style  . line_color .~ opaque red
              $ plot_lines_title .~ "psi2"
              $ def

	psi3 = plot_lines_values .~ [points3]
              $ plot_lines_style  . line_color .~ opaque green
              $ plot_lines_title .~ "psi3"
              $ def

	psi4 = plot_lines_values .~ [points4]
              $ plot_lines_style  . line_color .~ opaque black
              $ plot_lines_title .~ "psi4"
              $ def

	layout = layout_title .~ "Eigenfunctions"
           $ layout_plots .~ [toPlot psi1,
                              toPlot psi2,
							  toPlot psi3,
							  toPlot psi4]
           $ def

main = renderableToFile def{_fo_format=PDF} "eigenfunctions_quadratic_v.pdf" chart

max_stars_repo_name: davidschlegel/hartree-fock

max_issues_repo_name: davidschlegel/hartree-fock

max_forks_repo_name: davidschlegel/hartree-fock


text:
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.Gelu
Description : Gaussian Error Linear Unit (GELU)
Copyright   : (c) Manuel Schneckenreither, 2020
License     : BSD2
Stability   : experimental

This module implements the Gaussian Error Linear Unit (GELU) activiation function. See

Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).

Available at: https://arxiv.org/pdf/1606.08415.pdf

As in the paper we simply let μ = 0 and σ = 1. Futher, we use the simplified and thus fast representation: x * sigmoid (1.702 * x)

-}
module Grenade.Layers.Gelu (
    Gelu (..)
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Serialize
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS

import           Grenade.Core

-- | A Gaussion Error Linear Unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
--
--   Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).
data Gelu = Gelu
  deriving (Generic, NFData, Show)

instance UpdateLayer Gelu where
  type Gradient Gelu = ()
  runUpdate _ _ _ = Gelu
  reduceGradient _ = ()

instance RandomLayer Gelu where
  createRandomWith _ _ = return Gelu

instance Serialize Gelu where
  put _ = return ()
  get = return Gelu

geluForwardFast :: Floating x => x -> x
geluForwardFast x = x / (e ** (-1.702 * x) + 1) -- = x * sigmoid (1.702 * x)
  where
    e = 2.71828

geluBackwardFast :: Floating x => x -> x
geluBackwardFast x = (e ** (1.702 * x) * (1 + e ** (1.702 * x) + 1.702 * x)) / (1 + e ** (1.702 * x)) ** 2
  where
    e = 2.71828

instance (KnownNat i) => Layer Gelu ('D1 i) ('D1 i) where
  type Tape Gelu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (gelu y))
    where
      gelu = LAS.dvmap geluForwardFast
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (gelu' y * dEdy))
    where
      gelu' = LAS.dvmap geluBackwardFast

instance (KnownNat i, KnownNat j) => Layer Gelu ('D2 i j) ('D2 i j) where
  type Tape Gelu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Gelu ('D3 i j k) ('D3 i j k) where

  type Tape Gelu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (gelu y))
    where
      gelu = LAS.dmmap geluForwardFast
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (gelu' y * dEdy))
    where
      gelu' = LAS.dmmap geluBackwardFast

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade


text:
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module TicTacToe.AiInterface where

import Data.Maybe
import Control.Monad.State hiding (state)
import Control.Monad.Reader
import Control.Exception.Base (assert)
import System.Random (StdGen, split, randomR, Random)
import System.Random.Shuffle (shuffle')

import Numeric.LinearAlgebra

import TicTacToe.Core (move, result)
import TicTacToe.Domain (Result(..), Board, Move, CellPos)

class Player a where
  predictedAction :: BoardMatrix -> a -> Action

data GameState = GameState {
  reward' :: Int,
  done' :: Bool,
  board :: BoardMatrix
}

type BoardMatrix = Matrix Z
type Action = (Int, Int)

--step :: (Player a) => BoardMatrix -> Action -> a -> GameState
--step m action opponent =
--  case aiResult m action of
--    Error   -> err
--    Win     -> win
--    Draw    -> draw
--    _       -> opponentMoveToState (fromJust (aiMove' m action)) opponent
--
--opponentMoveToState :: (Player a) => BoardMatrix -> a -> GameState
--opponentMoveToState m opponent =
--  let pAction = predictedAction m opponent
--  in
--    case aiResult m pAction of
--      Win         -> loss
--      Draw        -> draw
--      Unfinished  -> unfinished m pAction
--
--aiMove' :: BoardMatrix -> Action -> Maybe BoardMatrix
--aiMove' matrix action =
--  let board           = matrixToBoard matrix
--      (pl, cellPos)   = actionToPlayerCellPos action
--  in
--    case move board pl cellPos of
--      Just b  -> Just $ boardToMatrix b
--      _       -> Nothing
--
--aiResult :: BoardMatrix -> Action -> Result
--aiResult matrix action =
--  result board pl cellPos
--  where
--    board           = matrixToBoard matrix
--    (pl, cellPos)   = actionToPlayerCellPos action
--
--actionToPlayerCellPos :: Action -> (Move, CellPos)
--actionToPlayerCellPos = assert False undefined
--
--matrixToBoard :: BoardMatrix -> Board
--matrixToBoard = assert False undefined
--
--boardToMatrix :: Board -> BoardMatrix
--boardToMatrix = assert False undefined
--
--matrixToState :: BoardMatrix -> GameState
--matrixToState = assert False undefined
--
--emptyGameState :: GameState
--emptyGameState = assert False undefined
--
--err :: GameState
--err = assert False undefined
--
--loss :: GameState
--loss = assert False undefined
--
--win :: GameState
--win = assert False undefined
--
--draw :: GameState
--draw = assert False undefined
--
--unfinished :: BoardMatrix -> Action -> GameState
--unfinished m a = assert False undefined

-- NN stuff

type Input  = Vector R
type W1     = Matrix R
type W2     = Matrix R
type W3     = Matrix R
type Output = Vector R
type B      = Matrix R

type Hidden1  = Matrix R
type Hidden2  = Matrix R

type Reward = R

type ActionIndex = Int

data Model = Model {
  w1        :: W1,
  w2        :: W2,
  w3        :: W3,
  b1        :: B,
  b2        :: B,
  b3        :: B
}

newtype Target = Target {
  unTarget :: Model
}

data NeuralNetwork = NN {
  epsilon       :: Double,
  learningRate  :: Double,
  model         :: Model,
  targetModel   :: Target,
  input         :: Input,
  buffer        :: [BufferData],
  adamIter      :: Int,
  adamW1        :: AdamState,
  adamW2        :: AdamState,
  adamW3        :: AdamState,
  adamB1        :: AdamState,
  adamB2        :: AdamState,
  adamB3        :: AdamState,
  randomGen     :: StdGen
}

data AdamState = AdamState {
  mass      :: Matrix R,
  velocity  :: Matrix R
}

data BufferData = BufferD {
  state     :: Input,
  action    :: ActionIndex,
  reward    :: Reward,
  nextState :: Input,
  done      :: Bool
}

data HyperParameters = Hyp {
  player            :: Move,
  epsilonDecay      :: Double,
  epsilonMin        :: Double,
  gamma             :: Double,
  learningRateDecay :: Double,
  episodes          :: Int,
  batchSize         :: Int,
  bufferSize        :: Int,
  stateSize         :: Int,
  adamP             :: AdamParameters,
  winReward         :: Reward,
  drawReward        :: Reward,
  lossReward        :: Reward,
  unfinishedReward  :: Reward
}

data AdamParameters = AdamParameters {
  eps       :: Double,
  beta1     :: Double,
  beta2     :: Double
}

type Network' a = ReaderT HyperParameters (State NeuralNetwork) a

newtype Network a = Network {
  unNetwork :: Network' a
} deriving (Monad, Applicative, Functor, MonadReader HyperParameters,
  MonadState NeuralNetwork)

learn1 :: Network ()
learn1 = do
  state <- gets input
  actionIndex <- act state
  (nextState, reward, done) <- aiMove actionIndex
  return ()

aiMove :: ActionIndex -> Network (Input, Reward, Bool)
aiMove action = do
  currentState <- gets input
  let stateAfterAction = update currentState action
  if (win stateAfterAction) then do
      reward <- asks winReward
      return (stateAfterAction, reward, True)
    else do
      opponentAction <- act $ reverseState $ stateAfterAction
      let stateAfterOpponent = update stateAfterAction opponentAction
      tmp stateAfterOpponent
        where
          tmp :: Input -> Network (Input, Reward, Bool)
          tmp state
            | loss state  = do
              reward <- asks lossReward
              return (state, reward, True)
            | draw state  = do
              reward <- asks drawReward
              return (state, reward, True)
            | otherwise   = do
              reward <- asks unfinishedReward
              return (state, reward, False)

update :: Input -> ActionIndex -> Input
update = assert False undefined

reverseState :: Input -> Input
reverseState = assert False undefined

win :: Input -> Bool
win = assert False undefined

loss :: Input -> Bool
loss = assert False undefined

draw :: Input -> Bool
draw = assert False undefined

replay :: Network ()
replay = do
  epsMin  <- asks epsilonMin
  epsD    <- asks epsilonDecay
  bSize   <- asks batchSize
  buffer  <- gets buffer
  sample  <- randomSample bSize buffer
  forM_ sample (\bufferD -> do
    target <- calcTarget bufferD
    model <- gets model
    model' <- fit model (state bufferD) target
    s <- get
    put s{ model=model' }
    s@NN{epsilon = eps} <- get
    when (eps > epsMin) $ put s{epsilon = eps * epsD})

calcTarget :: BufferData -> Network Output
calcTarget (BufferD state action reward nextState done) = do
  model <- gets model
  targetModel <- gets targetModel
  gamma <- asks gamma
  predCurRew <- predict model state
  if done then
      return $ fromList $ replace (toList predCurRew) action reward
    else do
      predFutRew <- predict model nextState
      targetPredFutRew <- predict (unTarget targetModel) nextState
      return $ fromList $ replace (toList predCurRew) action $ reward + gamma * targetPredFutRew ! maxIndex predFutRew
    where
      replace xs index value =
        let (before, _:after) = splitAt (index - 1) xs
        in before ++ (value:after)

fit :: Model -> Input -> Output -> Network Model
fit m@(Model w1 w2 w3 b1 b2 b3) state target = do
  let state' = asColumn state -- [9x1]
      target' = asColumn target -- [9x1]
      hidden1 = relu $ b1 + (w1 <> state') -- [24x9] x [9x1] = [24x1]
      hidden2 = relu $ add b2 $ w2 <> hidden1 -- [24x24] x [24x1] = [24x1]
      prediction = add b3 $ w3 <> hidden2 -- [9x24] x [24x1] = [9x1]
      loss = meanSquaredError target' prediction -- [9x1]
      g1 = w1 <> state' + b1 -- [24x9] x [9x1] = [24x1]
      g2 = w2 <> hidden1 + b2 -- [24x24] x [24x1] = [24x1]
      dg3 = - (target' - prediction) -- [9x1]
      dg2 = drelu g1 * (tr' dg3 <> w3) -- [1x9] x [9x24] = [1x24]
      dg1 = drelu g2 * dg2 <> w2 -- [1x24] x [24x24] = [1x24]
      dw3 = dg3 <> tr' hidden2 -- [9x1] x [1x24] = [9x24]
      db3 = dg3 -- [9x1]
      dw2 = tr' dg2 <> tr' hidden1 -- [24x1] x [1x24] = [24x24]
      db2 = tr' dg2 -- [24x1]
      dw1 = tr' dg1 <> tr' state' -- [24x1] x [1x9] = [24x9]
      db1 = tr' dg1 -- [24x1]
      dm  = Model dw1 dw2 dw3 b1 b2 b3
      drelu g = fromLists $ map (map cutOff) $ toLists g
      cutOff el = if el <= 0 then 0 else 1
  adam m dm

adam :: Model -> Model -> Network Model
adam (Model w1 w2 w3 b1 b2 b3) (Model dw1 dw2 dw3 db1 db2 db3) = do
  adamP   <- asks adamP
  t       <- gets adamIter
  learningRate  <- gets learningRate
  w1A <- gets adamW1
  w2A <- gets adamW2
  w3A <- gets adamW3
  b1A <- gets adamB1
  b2A <- gets adamB2
  b3A <- gets adamB3
  let (w1', w1A') = adam' w1 dw1 w1A adamP t learningRate
      (w2', w2A') = adam' w2 dw2 w2A adamP t learningRate
      (w3', w3A') = adam' w3 dw3 w3A adamP t learningRate
      (b1', b1A') = adam' b1 db1 b1A adamP t learningRate
      (b2', b2A') = adam' b2 db2 b2A adamP t learningRate
      (b3', b3A') = adam' b3 db3 b3A adamP t learningRate
      t'          = t + 1
      model'      = Model w1 w2 w3 b1 b2 b3
  state   <- get
  put state{
    adamW1  = w1A',
    adamW2  = w2A',
    adamW3  = w3A',
    adamB1  = b1A',
    adamB2  = b2A',
    adamB3  = b3A',
    adamIter  = t'
    }
  return model'

adam' :: Matrix R -> Matrix R -> AdamState -> AdamParameters -> Int -> Double
  -> (Matrix R, AdamState)
adam' x dx (AdamState m v) (AdamParameters eps beta1 beta2) t learningRate =
  let m'  = scale beta1 m + scale (1 - beta1) dx
      mt  = m' / scalar (1 - beta1^t)
      v'  = scale beta2 v + scale (1 - beta2) (dx^2)
      vt  = v' / scalar (1 - beta2^t)
      x'  = x + scale (- learningRate) mt / (sqrtm vt + scalar eps)
  in  (x', AdamState m' v')

meanSquaredError :: Floating a => a -> a -> a
meanSquaredError target prediction =
  0.5 * ((target - prediction)**2)

act :: Input -> Network ActionIndex
act state = do
  eps <- gets epsilon
  (rNumber :: Double) <- randomNumber (0,1)
  if rNumber <= eps then
      randomElement $ possibleActions state
    else do
      model <- gets model
      maxIndex <$> predict model state

possibleActions :: Input -> [ActionIndex]
possibleActions =
  find (\el -> el > (- 0.5) && el < 0.5)

predict :: Model -> Input -> Network Output
predict (Model w1 w2 w3 b1 b2 b3) input = do
  let input' = asColumn input -- [9x1]
      hidden1 = relu $ add b1 $ w1 <> input' -- [24x9] x [9x1] = [24x1]
      hidden2 = relu $ add b2 $ w2 <> hidden1 -- [24x24] x [24x1] = [24x1]
  return (flatten $ add b3 $ w3 <> hidden2) -- [9x24] x [24x1] = [9x1]

remember :: BufferData -> Network ()
remember bData = do
  s@NN{ buffer=buffer }   <- get
  bSize                   <- asks bufferSize
  if length buffer < bSize
    then put s{ buffer = bData : buffer }
    else put s{ buffer = bData : init buffer }

relu :: Matrix R -> Matrix R
relu =
  fromLists . map (map (max 0)) . toLists

initHyp :: HyperParameters
initHyp = assert False undefined

initNN :: Reader HyperParameters NeuralNetwork
initNN = assert False undefined

randomSample :: Int -> [a] -> Network [a]
randomSample n xs =
  take n <$> randomShuffle xs

randomShuffle :: [a] -> Network [a]
randomShuffle xs = do
  s@NN{ randomGen=gen }  <- get
  let shuffled  = shuffle' xs (length xs) gen
      (gen', _) = split gen
  put s{ randomGen=gen' }
  return shuffled

randomElement :: [a] -> Network a
randomElement xs = do
  randomIndex :: Int <- randomNumber (0, length xs)
  return $ xs !! randomIndex

randomNumber :: (Random a) => (a, a) -> Network a
randomNumber (beg, end) = do
  s@NN{ randomGen=gen }  <- get
  let (rNumber, gen') = randomR (beg, end) gen
  put s{randomGen=gen'}
  return rNumber

max_stars_repo_name: horttanainen/tictactoe

max_issues_repo_name: horttanainen/tictactoe

max_forks_repo_name: horttanainen/tictactoe


text:
{-# LANGUAGE RankNTypes, FlexibleContexts #-}
module Numeric.ADLens.HMatrix
    ( mdot'
    , add'
    , dot'
    , sumElements'
    , cmap'
    ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel (zipVectorWith)
import Numeric.ADLens.Lens
-- import Data.Vector as V

dot' :: (Container Vector t, Numeric t) => Lens' (Vector t, Vector t) t
dot' = lens'' $ \(v1,v2) -> (v1 <.> v2, \ds -> (scale ds v2, scale ds v1))

mdot' :: (Product t, Numeric t) => Lens' (Matrix t, Vector t) (Vector t)
mdot' = lens'' $ \(a,v) -> (a #> v, \dv -> (outer dv v, dv <# a))

add' :: Additive c => Lens' (c, c) c
add' = lens'' $ \(v1,v2) -> (add v1 v2, \dv -> (dv, dv))

-- I need konst I think?
sumElements' :: (Container Vector t, Numeric t) => Lens' (Vector t) t
sumElements' = lens'' $ \v -> (sumElements v, \ds -> scalar ds)

reshape' :: Container Vector t => Int -> Lens' (Vector t) (Matrix t)
reshape' n = lens'' $ \v -> (reshape n v,  \dm -> flatten dm)

-- conjugate transpose not trace
tr'' ::  (Transposable m mt, Transposable mt m) => Lens' m mt
tr'' = lens'' $ \x -> (tr x, \dt -> tr dt)

flatten' :: (Num t, Container Vector t) => Lens' (Matrix t) (Vector t)
flatten' = lens'' $ \m -> let s = fst $ size m in
                          (flatten m,  \dm -> reshape s dm)

norm_2' :: (Container c R, Normed (c R), Linear R c) => Lens' (c R) R
norm_2' = lens'' $ \v -> let nv = norm_2 v in (nv, \dnv -> scale (2 * dnv / nv) v )

cmap' :: (Element b, Container Vector e) => (Lens' e b) -> Lens' (Vector e) (Vector b)
cmap' l = lens'' $ \c -> (cmap f c, \dc -> zipVectorWith f' c dc) where
        f = view l
        f' = set l

{-
maxElement' :: Container c e => Lens' (c e) e
maxElement' = lens'' $ \v -> let i = maxIndex v in (v ! i, dv -> scalar 0)
-}

det' :: Field t => Lens' (Matrix t) t
det' = lens'' $ \m -> let (minv, (lndet, phase)) = invlndet m in
                    let detm = phase * exp detm in
                    (detm, \ds -> (scale (ds * detm) minv))

diag' :: (Num a, Element a) => Lens' (Vector a) (Matrix a)
diag' = lens'' $ \v -> (diag v, takeDiag)

takeDiag' :: (Num a, Element a) => Lens' (Matrix a) (Vector a)
takeDiag' = lens'' $ \m -> (takeDiag m, diag)


text:
{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}

import           Control.Monad
import           Control.Monad.Random
import           Data.Constraint              (Dict (..))
import           Data.List                    (foldl')
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.Prelude.List
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static as SA
import           Options.Applicative
import           Unsafe.Coerce                (unsafeCoerce)

import           Grenade

-- | The definition for a feed forward network using the dynamic module. Note the nested networks. This network clearly is over-engeneered for this example!
netSpec :: SpecNet
netSpec = specFullyConnected 2 40 |=> specTanh1D 40 |=> specDropout 40 0.95 Nothing |=> netSpecInner |=> specFullyConnected 20 30 |=> specRelu1D 30 |=> specFullyConnected 30 20 |=> specRelu1D 20 |=> specFullyConnected 20 10 |=> specRelu1D 10 |=> specFullyConnected 10 1 |=> specLogit1D 1 |=> specNil1D 1
  where netSpecInner = specFullyConnected 40 30 |=> specRelu1D 30 |=> specFullyConnected 30 20 |=> specReshape1D2D 20 (2, 10) |=> specReshape2D1D (2, 10) 20 |=> specNil1D 20

-- | Specifications can be built using the following interface also. Here one does not have to carefully pay attention to match the dimension inputs and outputs as when using specification directly.
buildNetViaInterface :: IO SpecConcreteNetwork
buildNetViaInterface =
  buildModelWith UniformInit (BuildSetup { printResultingSpecification = False })  $
  inputLayer1D 2 >>                            -- 1. Every model has to start with an input dimension
  fullyConnected 10 >> dropout 0.89 >> relu >> -- 2. Layers are simply added as desired. The input of each layer is determined automatically,
  fullyConnected 20 >> relu >>                 --    whereas the output is specified as in `fullyConnected 20`. Thus, the layer empits 20 signals
  fullyConnected 4 >> tanhLayer >>             --    which arethe input of the next layer.
  networkLayer (
    inputLayer1D 4 >> fullyConnected 10 >> relu >> fullyConnected 4 >> sinusoid
    ) >>
  reshape (4, 1, 1) >>                         -- 3. Reshape is ignored if the input and output is the same
  fullyConnected 1                             -- 4. The output is simply the number of signals the alst layer emits.

netTrain ::
     (SingI (Last shapes), MonadRandom m, KnownNat len1, KnownNat len2, RunnableNetwork layers shapes, Head shapes ~ 'D1 len1, Last shapes ~ 'D1 len2)
  => Network layers shapes
  -> Optimizer o
  -> Int
  -> m (Network layers shapes)
netTrain net0 op n = do
    inps <- replicateM n $ do
      s  <- getRandom
      return $ S1D $ SA.randomVector s SA.Uniform * 2 - 1
    let outs = flip map inps $ \(S1D v) ->
                 if v `inCircle` (fromRational 0.50, 0.50)  || v `inCircle` (fromRational (-0.50), 0.50)
                   then S1D $ fromRational 1
                   else S1D $ fromRational 0

    let (trained, _) = foldl' trainEach (net0, 0) (zip inps outs)
    return trained

  where trainEach (!network, _) (i,o) = train op network i o quadratic'

netScore :: (RunnableNetwork layers shapes, KnownNat len, Head shapes ~ 'D1 len, Last shapes ~ 'D1 1) => Network layers shapes -> IO ()
netScore network = do
    let testIns = [ [ (x,y)  | x <- [0..50] ]
                             | y <- [0..20] ]
        outMat  = fmap (fmap (\(x,y) -> (render (x/25-1) (y/10-1) . normx) $ runNet network (S1D $ SA.vector [x / 25 - 1,y / 10 - 1]))) testIns
    putStrLn $ unlines outMat

  where
    render x y n'  | x == 0 && y == 0 = '+'
                   | y == 0 = '-'
                   | x == 0 = '|'
                   | n' <= 0.2  = ' '
                   | n' <= 0.4  = '.'
                   | n' <= 0.6  = '-'
                   | n' <= 0.8  = '='
                   | otherwise = '#'

normx :: S ('D1 1) -> RealNum
normx (S1D r) = SA.mean r

testValues :: (RunnableNetwork layers shapes, KnownNat len, Head shapes ~ 'D1 len, Last shapes ~ 'D1 1) => Network layers shapes -> IO ()
testValues network = do
  inps <- replicateM 1000 $ do
      s  <- getRandom
      return $ S1D $ SA.randomVector s SA.Uniform * 2 - 1
  let outs = flip map inps $ \(S1D v) ->
                 if v `inCircle` (fromRational 0.50, 0.50)  || v `inCircle` (fromRational (-0.50), 0.50)
                   then 1 :: Integer
                   else 0
  let ress = zip outs (map (round . normx . runNet network) inps)
      correct = length $ filter id $ map (uncurry (==)) ress
      incorrect = length $ filter id $ map (uncurry (/=)) ress
      falsePositives = length $ filter id $ map (\ (shd, nn) -> shd == 0 && nn == 1) ress
      falseNegatives = length $ filter id $ map (\ (shd, nn) -> shd == 1 && nn == 0) ress
  putStr $ show correct  ++ " | "
  putStr $ show incorrect ++ " | "
  putStr $ show falsePositives ++ " | "
  putStrLn $ show falseNegatives ++ " | "

inCircle :: KnownNat n => SA.R n -> (SA.R n, RealNum) -> Bool
v `inCircle` (o, r) = SA.norm_2 (v - o) <= r

data FeedForwardOpts = FeedForwardOpts Int (Optimizer 'Adam)

feedForward' :: Parser FeedForwardOpts
feedForward' =
  FeedForwardOpts <$> option auto (long "examples" <> short 'e' <> value 1000)
                  <*> (OptAdam
                       <$> option auto (long "alpha" <> short 'r' <> value 0.001)
                       <*> option auto (long "beta1" <> value 0.9)
                       <*> option auto (long "beta2" <> value 0.999)
                       <*> option auto (long "epsilon" <> value 1e-4)
                       <*> option auto (long "lambda" <> value 1e-3))

main :: IO ()
main = do

  FeedForwardOpts examples rate <- execParser (info (feedForward' <**> helper) idm)
  putStrLn "| Nr | Correct | Incorrect | FalsePositives | FalseNegatives |"
  putStrLn "--------------------------------------------------------------"
  let nr = 10 :: Int
  mapM_
    (\n -> do
       putStr $ "| " ++ show n ++ " | "
       SpecConcreteNetwork1D1D (net0 :: Network layers shapes) <- networkFromSpecificationWith HeEtAl netSpec
      -- We need to specify the actual number of output nodes, as our functions requiere that!
       case (unsafeCoerce (Dict :: Dict ()) :: Dict (('D1 1) ~ Last shapes)) of
         Dict -> do
           net <- netTrain net0 rate examples
           unsafeCoerce $ -- only needed as GADTs are enabled, which disallowes the type to escape and thus prevents the type inference to work. The result is not needed anyways.
             testValues net)
    [1 .. nr]

  -- Features of dynamic networks:
  SpecConcreteNetwork1D1D (net' :: Network layers shapes) <- networkFromSpecificationWith HeEtAl netSpec
  net2 <- netTrain net' rate examples
  let spec' = networkToSpecification net2
  putStrLn "String represenation of the network specification: "
  print spec'
  let serializedSpec = encode spec'   -- only the specification (not the weights) are serialized here! The weights can be serialized using the networks serialize instance!
  let _ = encode net2                 -- E.g. like this.
  case decode serializedSpec of
    Left err -> print err
    Right spec2 -> do
      print spec2
      SpecConcreteNetwork1D1D (net3 :: Network layers3 shapes3) <- networkFromSpecificationWith HeEtAl spec2
      net4 <- foldM (\n _ -> netTrain n rate examples) net3 [(1 :: Int)..30]
      case (unsafeCoerce (Dict :: Dict ()) :: Dict (('D1 1) ~ Last shapes3)) of
        Dict -> netScore net4

  -- There is a also nice interface available also
  SpecConcreteNetwork1D1D newNet <- buildNetViaInterface
  print (networkToSpecification newNet)


text:
module LinearRegression where
import Util
import GradientDescent
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util
import Debug.Trace

-- The cost function J
-- Compute regularized cost for linear regression with multiple variables
-- computes the cost of using theta as the parameter for linear
-- regression to fit the data points in X and y, returns double of cost and gradient
computeCost :: Matrix Double -> Matrix Double -> Matrix Double -> Double -> (Double, Matrix Double)
computeCost x y theta lambda = (cost, gradient)
    where
      m = fromIntegral $ rows y
      tempTheta = accum theta (*) [((0,0),0)]
      cost = ( 1/(2*m) ) * sumElements (((x <> theta) - y) ^2 ) + (lambda/(2*m)) * sumElements (tempTheta ^ 2)
      gradient = ctrans $ (scalar (1/m)) * ( sumColumns  ((x <> theta - y) * x) ) + (scalar (lambda/m)) * (ctrans tempTheta)

linearRegressionGD x y initial_theta alpha num_iters lambda = gradientDescent gradFunc initial_theta alpha num_iters
    where
      gradFunc t = snd (computeCost x y t lambda)

-- The prediction function
predict xi theta = xi <> theta

-- Generate learning curves
learningCurvesLinearRegression ::
                   Matrix Double ->
                   Matrix Double ->
                   Matrix Double ->
                   Matrix Double ->
                   Double ->
                   Double ->
                   Int ->
                   [(Double,Double,Int)]
learningCurvesLinearRegression x y xval yval lambda alpha iter = learningCurves trainFunc costFunc x y xval yval 1
  where
    initial_theta = ((cols x)><1) (repeat 1)
    trainFunc tx ty = linearRegressionGD tx ty initial_theta alpha iter lambda
    costFunc tx ty tt = fst $ computeCost tx ty tt 0

max_stars_repo_name: thade/haskellml

max_issues_repo_name: thade/haskellml

max_forks_repo_name: thade/haskellml