-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE CPP #-}

module ProfileTests where

import Haxl.Prelude

import Haxl.Core
import Haxl.Core.Monad
import Haxl.Core.Stats

import Test.HUnit

import Control.DeepSeq (force)
import Control.Exception (evaluate)
import Data.Aeson
import Data.IORef
import qualified Data.HashMap.Strict as HashMap
#if MIN_VERSION_aeson(2,0,0)
import qualified Data.Aeson.KeyMap as KeyMap
#else
import qualified Data.HashMap.Strict as KeyMap
#endif
import Data.Int

import TestTypes
import TestUtils
import WorkDataSource
import SleepDataSource

mkProfilingEnv :: IO HaxlEnv
mkProfilingEnv = do
  env <- makeTestEnv False
  return env { flags = (flags env) { report = profilingReportFlags } }

-- expects only one label to be shown
labelToDataMap :: Profile -> HashMap.HashMap ProfileLabel ProfileData
labelToDataMap Profile{..} = HashMap.fromList hashKeys
  where
    labelKeys = HashMap.fromList [
      (k, l) | ((l, _), k) <- HashMap.toList profileTree]
    hashKeys = [ (l, v)
      | (k, v) <- HashMap.toList profile
      , Just l <- [HashMap.lookup k labelKeys]]

collectsdata :: Assertion
collectsdata = do
  e <- mkProfilingEnv
  _x <- runHaxl e $
          withLabel "bar" $
            withLabel "foo" $ do
              u <- env userEnv
              slp <- sum <$> mapM (\x -> withLabel "baz" $ return x) [1..5]
              -- do some non-trivial work that can't be lifted out
              -- first sleep though in order to force a Blocked result
              sleep slp `andThen` case fromJSON <$> KeyMap.lookup "A" u of
                Just (Success n) | sum [n .. 1000::Integer] > 0 -> return 5
                _otherwise -> return (4::Int)
  profCopy <- readIORef (profRef e)
  let
    profData = profile profCopy
    labelKeys = HashMap.fromList [
      (l, k) | ((l, _), k) <- HashMap.toList (profileTree profCopy)]
    getData k = do
      k2 <- HashMap.lookup k labelKeys
      HashMap.lookup k2 profData
  assertEqual "has data" 4 $ HashMap.size profData
  assertBool "foo allocates" $
    case profileAllocs <$> getData "foo" of
      Just x -> x > 10000
      Nothing -> False
  assertEqual "foo is only called once" (Just 1) $
    profileLabelHits <$> getData "foo"
  assertEqual "baz is called 5 times" (Just 5) $
    profileLabelHits <$> getData "baz"
  assertBool "bar does not allocate (much)" $
    case profileAllocs <$> getData "bar" of
      Just n -> n < 5000  -- getAllocationCounter can be off by +/- 4K
      _otherwise -> False
  let fooParents = case HashMap.lookup "foo" labelKeys of
        Nothing -> []
        Just kfoo ->
          [ kparent
          | ((_, kparent), k) <- HashMap.toList (profileTree profCopy)
          , k == kfoo]
  assertEqual "foo's parent" 1 (length fooParents)
  assertEqual "foo's parent is bar" (Just (head fooParents)) $
    HashMap.lookup ("bar", 0) (profileTree profCopy)

collectsLazyData :: Assertion
collectsLazyData = do
  e <- mkProfilingEnv
  _x <- runHaxl e $ withLabel "bar" $ do
          u <- env userEnv
          withLabel "foo" $ do
             let start = if KeyMap.member "A" u
                         then 10
                         else 1
             return $ sum [start..10000::Integer]
  profCopy <- readIORef (profRef e)
  -- check the allocations are attributed to foo
  assertBool "foo has allocations" $
    case profileAllocs <$> HashMap.lookup "foo" (labelToDataMap profCopy) of
      Just x -> x > 10000
      Nothing -> False

exceptions :: Assertion
exceptions = do
  env <- mkProfilingEnv
  _x <- runHaxl env $
          withLabel "outer" $
            tryToHaxlException $ withLabel "inner" $
              unsafeLiftIO $ evaluate $ force (error "pure exception" :: Int)
  profData <- labelToDataMap <$> readIORef (profRef env)
  assertBool "inner label not added" $
    not $ HashMap.member "inner" profData

  env2 <- mkProfilingEnv
  _x <- runHaxl env2 $
          withLabel "outer" $
            tryToHaxlException $ withLabel "inner" $
              throw $ NotFound "haxl exception"
  profData <- labelToDataMap <$> readIORef (profRef env2)
  assertBool "inner label added" $
    HashMap.member "inner" profData

-- Test that we correctly attribute work done in child threads when
-- using BackgroundFetch to the caller of runHaxl. This is important
-- for correct accounting when relying on allocation limits.
threadAlloc :: Integer -> Assertion
threadAlloc batches = do
  env' <- initEnv (stateSet mkWorkState stateEmpty) () :: IO (Env () ())
  let env = env'  { flags = (flags env') {
    report = setReportFlag ReportFetchStats defaultReportFlags } }
  a0 <- getAllocationCounter
  let
    wsize = 100000
    w = forM [wsize..(wsize+batches-1)] work
  _x <- runHaxl env $ sum <$> w
  a1 <- getAllocationCounter
  let
    lower = fromIntegral $ 1000000 * batches
    upper = fromIntegral $ 25000000 * batches
  assertBool "threadAlloc lower bound" $ (a0 - a1) > lower
  assertBool "threadAlloc upper bound" $ (a0 - a1) < upper
    -- the result was 16MB on 64-bit, or around 25KB if we miss the allocs
    -- in the child thread. For batched it should be similarly scaled.
    -- When we do not reset the counter for each batch was
    -- scaled again by number of batches.

  stats <- readIORef (statsRef env)
  assertEqual
    "threadAlloc: batches"
    [fromIntegral batches]
    (aggregateFetchBatches length stats)
  -- if we actually do more than 1 batch then the above test is not useful

data MemoType = Global | Local

-- Test that we correctly attribute memo work
memos:: MemoType -> Assertion
memos memoType = do
  env <- mkProfilingEnv
  let
    memoAllocs = 10000000 :: Int64
    doWork = unsafeLiftIO $ do
      a0 <- getAllocationCounter
      setAllocationCounter $ a0 - memoAllocs
      return (5 :: Int)
    mkWork
      | Global <- memoType = return (memo (1 :: Int) doWork)
      | Local <- memoType = memoize doWork
  _ <- runHaxl env $ do
    work <- mkWork
    andThen
      (withLabel "do" work)
      (withLabel "cached" work)
  profData <- labelToDataMap <$> readIORef (profRef env)
  case HashMap.lookup "do" profData of
    Nothing -> assertFailure "do not in data"
    Just ProfileData{..} -> do
      assertEqual "has correct memo id" profileMemos [ProfileMemo 1 False]
      assertBool "allocs are included in 'do'" (profileAllocs >= memoAllocs)
  case HashMap.lookup "cached" profData of
    Nothing -> assertFailure "cached not in data"
    Just ProfileData{..} -> do
      assertEqual "has correct memo id" profileMemos [ProfileMemo 1 True]
      assertBool "allocs are *not* included in 'cached'" (profileAllocs < 50000)
  (Stats memoStats) <- readIORef (statsRef env)
  assertEqual "exactly 1 memo/fetch" 1 (length memoStats)
  let memoStat = head memoStats
  putStrLn $ "memoStat=" ++ show memoStat
  assertEqual "correct call id" 1 (memoStatId memoStat)
  assertBool "allocs are big enough" $ memoSpace memoStat >= memoAllocs
  assertBool "allocs are not too big" $ memoSpace memoStat < memoAllocs + 100000

tests = TestList
  [ TestLabel "collectsdata" $ TestCase collectsdata
  , TestLabel "collectsdata - lazy" $ TestCase collectsLazyData
  , TestLabel "exceptions" $ TestCase exceptions
  , TestLabel "threads" $ TestCase (threadAlloc 1)
  , TestLabel "threads with batch" $ TestCase (threadAlloc 50)
  , TestLabel "memos - Global" $ TestCase (memos Global)
  , TestLabel "memos - Local" $ TestCase (memos Local)
  ]


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}

module SleepDataSource (
    Sleep, sleep,
  ) where

import Haxl.Prelude
import Prelude ()

import Haxl.Core
import Haxl.DataSource.ConcurrentIO

import Control.Concurrent
import Data.Hashable
import Data.Typeable

sleep :: Int -> GenHaxl u w Int
sleep n = dataFetch (Sleep n)

data Sleep deriving Typeable
instance ConcurrentIO Sleep where
  data ConcurrentIOReq Sleep a where
    Sleep :: Int -> ConcurrentIOReq Sleep Int

  performIO (Sleep n) = threadDelay (n*1000) >> return n

deriving instance Eq (ConcurrentIOReq Sleep a)
deriving instance Show (ConcurrentIOReq Sleep a)

instance ShowP (ConcurrentIOReq Sleep) where showp = show

instance Hashable (ConcurrentIOReq Sleep a) where
  hashWithSalt s (Sleep n) = hashWithSalt s n


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}

module StatsTests (tests) where

import Test.HUnit
import Data.List
import Data.Maybe

import Haxl.Prelude
import Haxl.Core
import Prelude()

import ExampleDataSource
import SleepDataSource
import Haxl.DataSource.ConcurrentIO

import Control.Monad (void)
import Data.IORef
import qualified Data.HashMap.Strict as HashMap

aggregateBatches :: Test
aggregateBatches = TestCase $ do
  let
    statsNoBatches = [ FetchStats { fetchDataSource = "foo"
                                  , fetchBatchSize = 7
                                  , fetchStart = 0
                                  , fetchDuration = 10
                                  , fetchSpace = 1
                                  , fetchFailures = 2
                                  , fetchIgnoredFailures = 0
                                  , fetchBatchId = n
                                  , fetchIds = [1,2] }
                                  | n <- reverse [1..10] ++ [11..20] ]
                     ++ [ FetchCall "A" ["B"] 1, FetchCall "C" ["D"] 2 ]
    fetchBatch = [ FetchStats { fetchDataSource = "batch"
                              , fetchBatchSize = 1
                              , fetchStart = 100
                              , fetchDuration = 1000 * n
                              , fetchSpace = 3
                              , fetchFailures = if n <= 3 then 1 else 0
                              , fetchIgnoredFailures = 0
                              , fetchBatchId = 123
                              , fetchIds = [fromIntegral n] } | n <- [1..50] ]
    agg (sz,bids) FetchStats{..} = (sz + fetchBatchSize, fetchBatchId:bids)
    agg _ _ = error "unexpected"
    agg' = foldl' agg (0,[])
    aggNoBatch = aggregateFetchBatches agg' (Stats statsNoBatches)
    expectedNoBatch = [(7, [n]) | n <- reverse [1..20] :: [Int]]
    aggBatch = aggregateFetchBatches agg' (Stats fetchBatch)
    expectedResultBatch = (50, [123 | _ <- [1..50] :: [Int]])
    aggInterspersedBatch =
      aggregateFetchBatches agg'
      (Stats $ intersperse (head fetchBatch) statsNoBatches)
    expectedResultInterspersed =
      (21, [123 | _ <- [1..21] :: [Int]]) : expectedNoBatch
  assertEqual "No batch has no change" expectedNoBatch aggNoBatch
  assertEqual "Batch is combined" [expectedResultBatch] aggBatch
  assertEqual
    "Grouping works as expected" expectedResultInterspersed aggInterspersedBatch

testEnv :: IO (Env () ())
testEnv = do
  -- To use a data source, we need to initialize its state:
  exstate <- ExampleDataSource.initGlobalState
  sleepState <- mkConcurrentIOState

  -- And create a StateStore object containing the states we need:
  let st = stateSet exstate (stateSet sleepState stateEmpty)

  -- Create the Env:
  env <- initEnv st ()
  return env{ flags = (flags env){
    report = setReportFlag ReportFetchStack profilingReportFlags } }

fetchIdsSync :: Test
fetchIdsSync = TestCase $ do
  env <- testEnv
  _ <- runHaxl  env $
       sequence_
       [ void $ countAardvarks "abcabc" + (length <$> listWombats 3)
       , void $ listWombats 100
       , void $ listWombats 99
       , void $ countAardvarks "BANG4" `catch` \NotFound{} -> return 123
       ]
  -- expect a single DS stat
  (Stats stats) <- readIORef (statsRef env)
  let
    fetchStats = [x | x@FetchStats{} <- stats]
  assertEqual "Only 1 batch" 1 (length fetchStats)
  let
    [stat] = fetchStats
  assertEqual "No real failures" 0 (fetchFailures stat)
  assertEqual "1 ignored failure" 1 (fetchIgnoredFailures stat)

fetchIdsBackground :: Test
fetchIdsBackground = TestCase $ do
  env <- testEnv
  _ <- runHaxl  env $
       sequence_
       [ withLabel "short" $ sleep 1
       , withLabel "long" $ sleep 500 ]

  -- make sure that with memo'ing we still preserve the stack
  _ <- runHaxl  env $ withLabel "base"
    (memo (1 :: Int) $ withLabel "child" $ sleep 102)

  _ <- runHaxl env $ withLabel "short_cached" $ sleep 1

  -- expect a single DS stat
  (Stats stats) <- readIORef (statsRef env)
  (Profile p pt _) <- readIORef (profRef env)
  let
    keyMap =
      HashMap.fromList [ (label, k) | ((label,_), k) <- HashMap.toList pt]
    revMap = HashMap.fromList [(v,k) | (k,v) <- HashMap.toList pt]
    parentMap =
      HashMap.fromList $
      catMaybes
      [ case HashMap.lookup kp revMap of
          Just (lp,_) -> Just (label, lp)
          Nothing -> Nothing
      | ((label,kp), _) <- HashMap.toList pt]
    fetchMap =  HashMap.fromList [ (fid, x) | x@FetchStats{} <- stats
                                   , fid <- fetchIds x]
    get l = [ (prof, wasCached, fetchStat)
            | Just key <- [HashMap.lookup l keyMap]
            , Just prof <- [HashMap.lookup key p]
            , ProfileFetch fid _ wasCached <- profileFetches prof
            , Just fetchStat <- [HashMap.lookup fid fetchMap]]
    [(short, shortWC, shortFetch)] = get "short"
    [(long, longWC, longFetch)] = get "long"
    [(shortCached, shortCachedWC, shortCachedFetch)] = get "short_cached"

  assertEqual "3 batches" 3 (HashMap.size fetchMap)
  assertEqual "6 labels (inc MAIN)" 6 (HashMap.size keyMap)

  assertEqual "child parent is base"
    (Just "base")
    (HashMap.lookup "child" parentMap)

  assertEqual "base parent is MAIN"
    (Just "MAIN")
    (HashMap.lookup "base" parentMap)

  assertEqual "long parent is MAIN"
    (Just "MAIN")
    (HashMap.lookup "long" parentMap)

  assertBool "original fetches not cached (short)" (not shortWC)
  assertBool "original fetches not cached (long)" (not longWC)
  assertBool "was cached short" shortCachedWC

  assertEqual "one fetch short" 1 (length $ profileFetches short)
  assertEqual "one fetch long" 1 (length $ profileFetches long)
  assertEqual "one fetch short_cached" 1 (length $ profileFetches shortCached)

  assertBool "short fetch mapped properly" (fetchDuration shortFetch < 100000)
  assertEqual
    "short cached fetch mapped properly"
    (fetchDuration shortFetch)
    (fetchDuration shortCachedFetch)
  assertBool "long fetch was mapped properly" (fetchDuration longFetch > 100000)

ppStatsTest :: Test
ppStatsTest = TestCase $ do
  let
    r = ppStats (Stats [])
    mc = ppStats (Stats [MemoCall 0 0])
    fc = ppStats (Stats [FetchCall "" [] 0])
    fw = ppStats (Stats [FetchWait HashMap.empty 0 1])
    fs = ppStats (Stats [FetchStats "" 0 0 0 0 0 0 0 []])
  assertEqual "empty stats -> empty string" r ""
  assertEqual "memo call stats -> empty string" mc ""
  assertEqual "fetch call stats -> empty string" fc ""
  assertBool "fetch wait stats -> some data" (not $ null fw)
  assertBool "fetch stats -> some data" (not $ null fs)

tests = TestList [ TestLabel "Aggregate Batches" aggregateBatches
                 , TestLabel "Fetch IDs Sync" fetchIdsSync
                 , TestLabel "Fetch IDs Background" fetchIdsBackground
                 , TestLabel "ppStats" ppStatsTest ]


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE OverloadedStrings #-}
module TestBadDataSource (tests) where

import Haxl.Prelude as Haxl
import Prelude()

import Haxl.Core

import Data.IORef
import Test.HUnit
import Control.Exception
import System.Mem

import ExampleDataSource
import BadDataSource

testEnv impl fn = do
  -- Use allocation limits, just to make sure haxl properly behaves and
  -- doesn't reset this internally somewhere.
  -- `go` will disable this
  setAllocationCounter 5000000 -- 5 meg should be enough, uses ~100k atm
  enableAllocationLimit
  exstate <- ExampleDataSource.initGlobalState
  badstate <- BadDataSource.initGlobalState impl
  let st = stateSet exstate $ stateSet (fn badstate) stateEmpty
  initEnv st ()

wombats :: GenHaxl () () Int
wombats = length <$> listWombats 3

wombatsMany :: GenHaxl () () Int
wombatsMany = length <$> listWombats 7

go :: FetchImpl -> Test
go impl = TestCase $ flip finally disableAllocationLimit $ do
  -- test that a failed acquire doesn't fail the other requests
  ref <- newIORef False
  env <- testEnv impl $ \st ->
    st { failAcquire = throwIO (DataSourceError "acquire")
       , failRelease = writeIORef ref True }

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \DataSourceError{} -> wombats

  assertEqual "badDataSourceTest1" 3 x

  -- We should *not* have called release
  assertEqual "badDataSourceTest2" False =<< readIORef ref

  -- test that a failed dispatch doesn't fail the other requests
  ref <- newIORef False
  env <- testEnv impl $ \st ->
    st { failDispatch = throwIO (DataSourceError "dispatch")
       , failRelease = writeIORef ref True }

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \DataSourceError{} -> wombats

  assertEqual "badDataSourceTest3" x 3

  -- We *should* have called release
  assertEqual "badDataSourceTest4" True =<< readIORef ref

  -- test that a failed wait is a DataSourceError
  env <- testEnv impl $ \st ->
    st { failWait = throwIO (DataSourceError "wait") }

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \DataSourceError{} -> wombats

  assertEqual "badDataSourceTest5" x 3

  -- We *should* have called release
  assertEqual "badDataSourceTest6" True =<< readIORef ref

  -- test that a failed release is still a DataSourceError, even
  -- though the request will have completed successfully
  env <- testEnv impl $ \st ->
    st { failRelease = throwIO (DataSourceError "release") }

  let
    -- In background fetches the scheduler might happen to process the data
    -- source result (FetchError in this case) before it processes the exception
    -- from release. So we have to allow both cases.
    isBg = case impl of
      Background -> True
      BackgroundMVar -> True
      _ -> False
    releaseCatcher e
      | Just DataSourceError{} <- fromException e = wombats
      | Just FetchError{} <- fromException e =
          if isBg then wombats else Haxl.throw e
      | otherwise = Haxl.throw e

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` releaseCatcher

  assertEqual "badDataSourceTest7" x 3

  -- test that if we don't throw anything we get the result
  -- (which is a fetch error for this source)
  env <- testEnv impl id
  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \FetchError{} -> wombats

  assertEqual "badDataSourceTest8" x 3

tests = TestList
  [ TestLabel "badDataSourceTest async" (go Async)
  , TestLabel "badDataSourceTest background" (go Background)
  , TestLabel "badDataSourceTest backgroundMVar" (go BackgroundMVar)
  , TestLabel "badDataSourceTest backgroundFetchSeq" (go BackgroundSeq)
  , TestLabel "badDataSourceTest backgroundFetchPar" (go BackgroundPar)
  ]


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE CPP, OverloadedStrings, RebindableSyntax, MultiWayIf #-}
{-# LANGUAGE RecordWildCards #-}
module TestExampleDataSource (tests) where

import Haxl.Prelude as Haxl
import Prelude()

import Haxl.Core.Monad (unsafeLiftIO)
import Haxl.Core

import Test.HUnit
import Data.IORef
import Data.Maybe
import Control.Exception
import System.Environment
import System.FilePath

import ExampleDataSource
import LoadCache

testEnv :: IO (Env () ())
testEnv = do
  -- To use a data source, we need to initialize its state:
  exstate <- ExampleDataSource.initGlobalState

  -- And create a StateStore object containing the states we need:
  let st = stateSet exstate stateEmpty

  -- Create the Env:
  env <- initEnv st ()
  return env{ flags = (flags env){
    report = setReportFlag ReportFetchStats defaultReportFlags } }

tests = TestList [
  TestLabel "exampleTest" exampleTest,
  TestLabel "orderTest" orderTest,
  TestLabel "preCacheTest" preCacheTest,
  TestLabel "cachedComputationTest" cachedComputationTest,
  TestLabel "cacheResultTest" cacheResultTest,
  TestLabel "memoTest" memoTest,
  TestLabel "dataSourceExceptionTest" dataSourceExceptionTest,
  TestLabel "dumpCacheAsHaskell" dumpCacheTest,
  TestLabel "fetchError" fetchError
  ]

-- Let's test ExampleDataSource.

exampleTest :: Test
exampleTest = TestCase $ do
  env <- testEnv

  -- Run an example expression with two fetches:
  x <- runHaxl env $
     countAardvarks "abcabc" + (length <$> listWombats 3)

  assertEqual "runTests" x (2 + 3)

  -- Should be just one fetching round:
  Stats stats <- readIORef (statsRef env)
  putStrLn (ppStats (Stats stats))
  assertEqual "rounds" 1 (length stats)

  -- With two fetches:
  assertBool "reqs" $
    case stats of
       [FetchStats{..}] ->
         fetchDataSource == "ExampleDataSource" && fetchBatchSize == 2
       _otherwise -> False

-- Test side-effect ordering

orderTest = TestCase $ do
  env <- testEnv

  ref <- newIORef ([] :: [Int])

  let tick n = unsafeLiftIO (modifyIORef ref (n:))

  let left = do tick 1
                r <- countAardvarks "abcabc"
                tick 2
                return r

  let right = do tick 3
                 r <- length <$> listWombats 3
                 tick 4
                 return r

  x <- runHaxl env $ left + right
  assertEqual "TestExampleDataSource2" x (2 + 3)

  -- The order of the side effects is 1,3,2,4.  First we see 1, then
  -- left gets blocked, then we explore right, we see 3, then right
  -- gets blocked.  The data fetches are performed, then we see 2 and
  -- then 4.

  ys <- readIORef ref
  assertEqual "TestExampleDataSource: ordering" (reverse ys) [1,3,2,4]

preCacheTest = TestCase $ do
  env <- testEnv

  x <- runHaxl env $ do
    cacheRequest (CountAardvarks "xxx") (Right 3)
    cacheRequest (ListWombats 1000000) (Right [1,2,3])
    countAardvarks "xxx" + (length <$> listWombats 1000000)
  assertEqual "preCacheTest1" x (3 + 3)

  y <- Control.Exception.try $ runHaxl env $ do
    cacheRequest (CountAardvarks "yyy") $ except (NotFound "yyy")
    countAardvarks "yyy"
  assertBool "preCacheTest2" $
     case y of
       Left (NotFound "yyy") -> True
       _other -> False

-- Pretend CountAardvarks is a request computed by some Haxl code
cachedComputationTest = TestCase $ do
  env <- testEnv
  let env' = env { flags = (flags env){trace = 3} }

  let x = cachedComputation (CountAardvarks "ababa") $ do
        a <- length <$> listWombats 10
        b <- length <$> listWombats 20
        return (a + b)

  r <- runHaxl env' $ x + x + countAardvarks "baba"

  assertEqual "cachedComputationTest1" 62 r

  stats <- readIORef (statsRef env)
  assertEqual "fetches" 3 (numFetches stats)

cacheResultTest = TestCase $ do
  env <- testEnv
  ref <- newIORef 0
  let request = cacheResult (CountAardvarks "ababa") $ do
         modifyIORef ref (+1)
         readIORef ref
  r <- runHaxl env $ (+) <$> request <*> request
  assertEqual "cacheResult" 2 r

-- Pretend CountAardvarks is a request computed by some Haxl code
memoTest = TestCase $ do
  env <- testEnv
  let env' = env { flags = (flags env){trace = 3} }

  let x = memo (CountAardvarks "ababa") $ do
        a <- length <$> listWombats 10
        b <- length <$> listWombats 20
        return (a + b)

  r <- runHaxl env' $ x + x + countAardvarks "baba"

  assertEqual "memoTest1" 62 r

  stats <- readIORef (statsRef env)
  assertEqual "fetches" 3 (numFetches stats)

-- Test that the FetchError gets returned properly, and that we have
-- a failure logged in the stats.
fetchError = TestCase $ do
  env <- testEnv
  r <- runHaxl env $ Haxl.try $
    (++) <$> listWombats 1000000 <*> listWombats 1000001
  assertBool "fetchError1" $ case r of
    Left FetchError{} -> True
    Right _ -> False
  Stats stats <- readIORef (statsRef env)
  assertEqual "fetchError2" 2 (sum [ fetchFailures | FetchStats{..} <- stats ])

dataSourceExceptionTest = TestCase $ do
  env <- testEnv

  r <- runHaxl env $ Haxl.try $ countAardvarks "BANG"
  assertBool "exception1" $
    case r of
      Left (ErrorCall "BANG") -> True
      _ -> False
  r <- runHaxl env $ Haxl.try $ countAardvarks "BANG2"
  assertBool "exception2" $
    case r of
      Left (ErrorCall "BANG2") -> True
      _ -> False

  -- In this test, BANG3 is an asynchronous exception (ThreadKilled),
  -- so we should see that instead of the exception on the left.
  -- Furthermore, it doesn't get caught by Haxl.try, and we have to
  -- catch it outside of runHaxl.
  env <- testEnv
  r <- Control.Exception.try $ runHaxl env $ Haxl.try $
          (length <$> listWombats 100) + countAardvarks "BANG3"
  print r
  assertBool "exception3" $
    case (r :: Either AsyncException (Either SomeException Int)) of
       Left ThreadKilled -> True
       _ -> False

-- Test that we can load the cache from a dumped copy of it, and then dump it
-- again to get the same result.
dumpCacheTest = TestCase $ do
  env <- testEnv
  runHaxl env loadCache
  str <- runHaxl env dumpCacheAsHaskell
  lcPath <-loadCachePath
  loadcache <- readFile lcPath
  -- The order of 'cacheRequest ...' calls is nondeterministic and
  -- differs among GHC versions, so we sort the lines for comparison.
  assertEqual "dumpCacheAsHaskell" (sort $ lines loadcache) (sort $ lines str)
  where
    loadCachePath = do
      lcEnv <- lookupEnv "LOADCACHE"
      return $ fromMaybe (dropFileName __FILE__ </> "LoadCache.txt") lcEnv


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE CPP, OverloadedStrings #-}
module Main where

import Test.Framework (defaultMain)
import Test.Framework.Providers.HUnit (hUnitTestToTests)
import AllTests

main :: IO ()
main = defaultMain $ hUnitTestToTests allTests


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE CPP #-}

module TestTypes
   ( UserEnv
   , Haxl
   , HaxlEnv
   , lookupInput
   , Id(..)
   ) where

import Data.Aeson
import Data.Binary (Binary)
import qualified Data.Text as Text
#if MIN_VERSION_aeson(2,0,0)
import qualified Data.Aeson.KeyMap as KeyMap
import Data.Aeson.Key (toText)
#else
import qualified Data.HashMap.Strict as KeyMap
#endif
import Data.Hashable
import Data.Typeable

import Haxl.Core

#if !MIN_VERSION_aeson(2,0,0)
type Key = Text.Text

toText :: Key -> Text.Text
toText = id
#endif

type UserEnv = Object
type Haxl a = GenHaxl UserEnv () a
type HaxlEnv = Env UserEnv ()

lookupInput :: FromJSON a => Key -> Haxl a
lookupInput field = do
  mb_val <- env (KeyMap.lookup field . userEnv)
  case mb_val of
    Nothing ->
      throw (NotFound (Text.concat ["field ", toText field, " was not found."]))
    Just val ->
      case fromJSON val of
        Error str ->
          throw (UnexpectedType (Text.concat
            ["field ", toText field, ": ", Text.pack str]))
        Success a -> return a

newtype Id = Id Int
  deriving (Eq, Ord, Binary, Enum, Num, Integral, Real, Hashable, Typeable,
            ToJSON, FromJSON)

instance Show Id where
  show (Id i) = show i


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE CPP #-}
module TestUtils
  ( makeTestEnv
  , expectResultWithEnv
  , expectResult
  , expectFetches
  , testinput
  , id1, id2, id3, id4
  ) where

import TestTypes
import MockTAO
import Haxl.DataSource.ConcurrentIO

import Data.IORef
import Data.Aeson
import Test.HUnit
#if MIN_VERSION_aeson(2,0,0)
import qualified Data.Aeson.KeyMap as KeyMap
#else
import qualified Data.HashMap.Strict as KeyMap
#endif

import Haxl.Core

import Prelude()
import Haxl.Prelude

testinput :: Object
testinput = KeyMap.fromList [
  "A" .= (1 :: Int),
  "B" .= (2 :: Int),
  "C" .= (3 :: Int),
  "D" .= (4 :: Int) ]

id1 :: Haxl Id
id1 = lookupInput "A"

id2 :: Haxl Id
id2 = lookupInput "B"

id3 :: Haxl Id
id3 = lookupInput "C"

id4 :: Haxl Id
id4 = lookupInput "D"

makeTestEnv :: Bool -> IO HaxlEnv
makeTestEnv future = do
  tao <- MockTAO.initGlobalState future
  stio <- mkConcurrentIOState
  let st = stateSet stio $ stateSet tao stateEmpty
  env <- initEnv st testinput
  return env { flags = (flags env) {
    report = setReportFlag ReportFetchStats defaultReportFlags } }

expectResultWithEnv
  :: (Eq a, Show a) => a -> Haxl a -> HaxlEnv -> Assertion
expectResultWithEnv result haxl env = do
  a <- runHaxl env haxl
  assertEqual "result" result a

expectResult :: (Eq a, Show a) => a -> Haxl a -> Bool -> Assertion
expectResult result haxl future = do
  env <- makeTestEnv future
  expectResultWithEnv result haxl env

expectFetches :: (Eq a, Show a) => Int -> Haxl a -> Bool -> Assertion
expectFetches n haxl future = do
  env <- makeTestEnv future
  _ <- runHaxl env haxl
  stats <- readIORef (statsRef env)
  assertEqual "fetches" n (numFetches stats)


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}

module WorkDataSource (
    mkWorkState,
    work,
  ) where

import Haxl.Prelude
import Prelude ()

import Haxl.Core

import Control.Exception
import Data.Hashable
import Data.Typeable
import Control.Monad (void)
import Control.Concurrent.MVar

data Work a where
  Work :: Integer -> Work Integer
  deriving Typeable

deriving instance Eq (Work a)
deriving instance Show (Work a)
instance ShowP Work where showp = show

instance Hashable (Work a) where
   hashWithSalt s (Work a) = hashWithSalt s (0::Int,a)

instance DataSourceName Work where
  dataSourceName _ = "Work"

instance StateKey Work where
  data State Work = WorkState

newtype Service = Service (MVar [IO ()])

run :: Work a -> IO a
run (Work n) = evaluate (sum [1..n]) >> return n

mkService :: IO Service
mkService = Service <$> newMVar []

process :: Service -> IO ()
process (Service q) = do
  r <- swapMVar q []
  sequence_ r

enqueue :: Service -> Work a -> IO (IO (Either SomeException a))
enqueue (Service q) w = do
  res <- newEmptyMVar
  let r = do
        v <- Control.Exception.try $ run w
        putMVar res v
  modifyMVar_ q (return . (:) r)
  return (takeMVar res)

instance DataSource u Work where
  fetch = backgroundFetchAcquireReleaseMVar
    mkService
    (\_ -> return ())
    -- pretend we are ready so that process does the work
    (\_ _ m -> void $ tryPutMVar m ())
    process
    enqueue

mkWorkState :: State Work
mkWorkState = WorkState

work :: Integer -> GenHaxl u w Integer
work n = dataFetch (Work n)


-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
module WriteTests (tests) where

import Test.HUnit

import Control.Arrow
import Control.Concurrent
import Data.Either
import Data.Foldable
import Data.Hashable
import Data.IORef
import qualified Data.Text as Text

import Haxl.Core.Monad (mapWrites, mapWriteTree, flattenWT, WriteTree)
import Haxl.Core
import Haxl.Prelude as Haxl

-- A fake data source
data SimpleDataSource a where
  GetNumber :: SimpleDataSource Int

deriving instance Eq (SimpleDataSource a)
deriving instance Show (SimpleDataSource a)
instance ShowP SimpleDataSource where showp = show

instance Hashable (SimpleDataSource a) where
  hashWithSalt s GetNumber = hashWithSalt s (0 :: Int)

instance StateKey SimpleDataSource where
  data State SimpleDataSource = DSState

instance DataSourceName SimpleDataSource where
  dataSourceName _ = "SimpleDataSource"

instance DataSource u SimpleDataSource where
  fetch _st _flags _usr  = SyncFetch $ Haxl.mapM_ fetch1
    where
    fetch1 :: BlockedFetch SimpleDataSource -> IO ()
    fetch1 (BlockedFetch GetNumber m) =
      threadDelay 1000 >> putSuccess m 37

newtype SimpleWrite = SimpleWrite Text
  deriving (Eq, Show, Ord)

assertEqualIgnoreOrder ::
  (Eq a, Show a, Ord a) => String -> [a] -> [a] -> Assertion
assertEqualIgnoreOrder msg lhs rhs =
  assertEqual msg (sort lhs) (sort rhs)

doInnerWrite :: GenHaxl u (WriteTree SimpleWrite) Int
doInnerWrite = do
  tellWrite $ SimpleWrite "inner"
  return 0

doOuterWrite :: GenHaxl u (WriteTree SimpleWrite) Int
doOuterWrite = do
  tellWrite $ SimpleWrite "outer1"

  doWriteMemo <- newMemoWith doInnerWrite
  let doMemoizedWrite = runMemo doWriteMemo
  _ <- doMemoizedWrite
  _ <- doMemoizedWrite

  tellWrite $ SimpleWrite "outer2"

  return 1

doNonMemoWrites :: GenHaxl u (WriteTree SimpleWrite) Int
doNonMemoWrites = do
  tellWrite $ SimpleWrite "inner"
  tellWriteNoMemo $ SimpleWrite "inner not memo"
  return 0

runHaxlWithWriteList :: Env u (WriteTree w) -> GenHaxl u (WriteTree w) a -> IO (a, [w])
runHaxlWithWriteList env haxl = second flattenWT <$> runHaxlWithWrites env haxl

writeSoundness :: Test
writeSoundness = TestCase $ do
  let numReps = 4

  -- do writes without memoization
  env1 <- emptyEnv ()
  (allRes, allWrites) <- runHaxlWithWriteList env1 $
    Haxl.sequence (replicate numReps doInnerWrite)

  assertBool "Write Soundness 1" $
    allWrites == replicate numReps (SimpleWrite "inner")
  assertBool "Write Soundness 2" $ allRes == replicate numReps 0

  -- do writes with memoization
  env2 <- emptyEnv ()

  (memoRes, memoWrites) <- runHaxlWithWriteList env2 $ do
    doWriteMemo <- newMemoWith doInnerWrite
    let memoizedWrite = runMemo doWriteMemo

    Haxl.sequence (replicate numReps memoizedWrite)

  assertBool "Write Soundness 3" $
    memoWrites == replicate numReps (SimpleWrite "inner")
  assertBool "Write Soundness 4" $ memoRes == replicate numReps 0

  -- do writes with interleaved memo
  env3 <- emptyEnv ()

  (ilRes, ilWrites) <- runHaxlWithWriteList env3 $ do
    doWriteMemo <- newMemoWith doInnerWrite
    let memoizedWrite = runMemo doWriteMemo

    Haxl.sequence $ replicate numReps (doInnerWrite *> memoizedWrite)

  assertBool "Write Soundness 5" $
    ilWrites == replicate (2*numReps) (SimpleWrite "inner")
  assertBool "Write Soundness 6" $ ilRes == replicate numReps 0

  -- do writes with nested memo
  env4 <- emptyEnv ()

  (nestRes, nestWrites) <- runHaxlWithWriteList env4 $ do
    doWriteMemo' <- newMemoWith doOuterWrite
    let memoizedWrite' = runMemo doWriteMemo'

    Haxl.sequence (replicate numReps memoizedWrite')

  let expWrites =
        [ SimpleWrite "outer1"
        , SimpleWrite "inner"
        , SimpleWrite "inner"
        , SimpleWrite "outer2"
        ]
  assertBool "Write Soundness 7" $
    nestWrites == fold (replicate numReps expWrites)
  assertBool "Write Soundness 8" $ nestRes == replicate numReps 1

  -- do both kinds of writes without memoization
  env5 <- emptyEnv ()
  (allRes, allWrites) <- runHaxlWithWriteList env5 $
    Haxl.sequence (replicate numReps doNonMemoWrites)

  assertBool "Write Soundness 9" $
    allWrites == replicate numReps (SimpleWrite "inner") ++
      replicate numReps (SimpleWrite "inner not memo")
  assertBool "Write Soundness 10" $ allRes == replicate numReps 0

  -- do both kinds of writes with memoization
  env6 <- emptyEnv ()

  (memoRes, memoWrites) <- runHaxlWithWriteList env6 $ do
    doWriteMemo <- newMemoWith doNonMemoWrites
    let memoizedWrite = runMemo doWriteMemo

    Haxl.sequence (replicate numReps memoizedWrite)

  -- "inner not memo" only appears once in this test
  assertBool "Write Soundness 11" $
    memoWrites == replicate numReps (SimpleWrite "inner") ++
      [SimpleWrite "inner not memo"]
  assertBool "Write Soundness 12" $ memoRes == replicate numReps 0

writeLogsCorrectnessTest :: Test
writeLogsCorrectnessTest = TestLabel "writeLogs_correctness" $ TestCase $ do
  e <- emptyEnv ()
  (_ , wrts) <- runHaxlWithWriteList e doNonMemoWrites
  assertEqualIgnoreOrder "Expected writes" [SimpleWrite "inner",
    SimpleWrite "inner not memo"] wrts
  wrtsNoMemo <- readIORef $ writeLogsRefNoMemo e
  wrtsMemo <- readIORef $ writeLogsRef e
  assertEqualIgnoreOrder "WriteTree not empty" [] $ flattenWT wrtsNoMemo
  assertEqualIgnoreOrder "WriteTree not empty" [] $ flattenWT wrtsMemo

mapWritesTest :: Test
mapWritesTest = TestLabel "mapWrites" $ TestCase $ do
  let funcSingle (SimpleWrite s) = SimpleWrite $ Text.toUpper s
      func = mapWriteTree funcSingle
  env0 <- emptyEnv ()
  (res0, wrts0) <- runHaxlWithWriteList env0 $ mapWrites func doNonMemoWrites
  assertEqual "Expected computation result" 0 res0
  assertEqualIgnoreOrder "Writes correctly transformed" [SimpleWrite "INNER",
    SimpleWrite "INNER NOT MEMO"] wrts0

  -- Writes should behave the same inside and outside mapWrites
  env1 <- emptyEnv ()
  (res1, wrts1) <- runHaxlWithWriteList env1 $ do
    outer <- doOuterWrite
    outerMapped <- mapWrites func doOuterWrite
    return $ outer == outerMapped
  assertBool "Results are identical" res1
  assertEqualIgnoreOrder
    "Writes correctly transformed, non-transformed writes preserved"
    [ SimpleWrite "outer1", SimpleWrite "inner"
    , SimpleWrite "inner", SimpleWrite "outer2"
    , SimpleWrite "OUTER1", SimpleWrite "INNER"
    , SimpleWrite "INNER", SimpleWrite "OUTER2"
    ]
    wrts1

  -- Memoization behaviour should be unaffected
  env2 <- emptyEnv ()
  (_res2, wrts2) <- runHaxlWithWriteList env2 $ do
    writeMemo <- newMemoWith doNonMemoWrites
    let doWriteMemo = runMemo writeMemo
    _ <- mapWrites func doWriteMemo
    _ <- doWriteMemo
    return ()
  -- "inner not memo" should appear only once
  assertEqualIgnoreOrder
    "Write correctly transformed under memoization"
    [ SimpleWrite "INNER"
    , SimpleWrite "inner"
    , SimpleWrite "INNER NOT MEMO"
    ]
    wrts2

  -- Same as previous, but the non-mapped computation is run first
  env3 <- emptyEnv ()
  (_res3, wrts3) <- runHaxlWithWriteList env3 $ do
    writeMemo <- newMemoWith doNonMemoWrites
    let doWriteMemo = runMemo writeMemo
    _ <- doWriteMemo
    _ <- mapWrites func doWriteMemo
    return ()
  -- "inner not memo" should appear only once
  assertEqualIgnoreOrder
    "Flipped: Write correctly transformed under memoization"
    [ SimpleWrite "inner"
    , SimpleWrite "INNER"
    , SimpleWrite "inner not memo"
    ]
    wrts3

  -- inner computation performs no writes
  env4 <- emptyEnv ()
  (res4, wrts4) <- runHaxlWithWriteList env4 $
    mapWrites func (return (0 :: Int))
  assertEqual "No Writes: Expected computation result" 0 res4
  assertEqualIgnoreOrder "No writes" [] wrts4

  -- inner computation throws an exception
  env5 <- emptyEnv ()
  (res5, wrts5) <- runHaxlWithWriteList env5 $ mapWrites func $ try $ do
    _ <- doNonMemoWrites
    _ <- throw (NotFound "exception")
    return 0
  assertBool "Throw: Expected Computation Result" $ isLeft
    (res5 :: Either HaxlException Int)
  assertEqualIgnoreOrder
    "Datasource writes correctly transformed"
    [ SimpleWrite "INNER"
    , SimpleWrite "INNER NOT MEMO"
    ]
    wrts5

  -- inner computation calls a datasource
  env6 <- initEnv (stateSet DSState stateEmpty) ()
  (res6, wrts6) <- runHaxlWithWriteList env6 $ mapWrites func $ do
    _ <- doNonMemoWrites
    dataFetch GetNumber

  assertEqual "Datasource: Expected Computation Result" 37 res6
  assertEqualIgnoreOrder
    "Datasource writes correctly transformed"
    [ SimpleWrite "INNER"
    , SimpleWrite "INNER NOT MEMO"
    ]
    wrts6

  -- inner computation calls a datasource, flipped calls
  env7 <- initEnv (stateSet DSState stateEmpty) ()
  (res7, wrts7) <- runHaxlWithWriteList env7 $ mapWrites func $ do
    df <- dataFetch GetNumber
    _ <- doNonMemoWrites
    return df

  assertEqual "Flipped Datasource: Expected Computation Result" 37 res7
  assertEqualIgnoreOrder
    "Flipped: Datasource writes correctly transformed"
    [ SimpleWrite "INNER"
    , SimpleWrite "INNER NOT MEMO"
    ]
    wrts7

tests = TestList
  [ TestLabel "Write Soundness" writeSoundness,
    TestLabel "writeLogs_correctness" writeLogsCorrectnessTest,
    TestLabel "mapWrites" mapWritesTest
  ]


import Distribution.Simple
main = defaultMain


-------------------------------------------------------------------------------
-- |
-- Module    :  Main
-- License   :  MIT
-- Stability :  experimental
-- Portability: non-portable
--
-- Example of how to build an agent using OpenAI.Gym.Client
-------------------------------------------------------------------------------
{-# LANGUAGE NamedFieldPuns #-}
module Main where

import Prelude
import Control.Monad (replicateM_, when)
import Control.Monad.Catch
import Control.Exception.Base

import OpenAI.Gym
import Servant.Client
import Network.HTTP.Client

main :: IO ()
main = do
  manager <- newManager defaultManagerSettings
  out <- runClientM example (ClientEnv manager url)
  case out of
    Left err -> print err
    Right ok -> print ok

  where
    url :: BaseUrl
    url = BaseUrl Http "localhost" 5000 ""

example :: ClientM ()
example = do
  inst <- envCreate CartPoleV0
  Monitor{directory} <- withMonitor inst $
    replicateM_ episodeCount (agent inst)

  -- Upload to the scoreboard.
  -- TODO: Implement environment variable support.
  upload (Config directory "algo" "")

  where
    episodeCount :: Int
    episodeCount = 100

agent :: InstID -> ClientM ()
agent inst = do
  envReset inst
  go 0 False
  where
    maxSteps :: Int
    maxSteps = 200

    reward :: Int
    reward = 0

    go :: Int -> Bool -> ClientM ()
    go x done = do
      Action a <- envActionSpaceSample inst
      Outcome ob reward done _ <- envStep inst (Step a True)
      when (not done && x < 200) $ go (x + 1) done

withMonitor :: InstID -> ClientM () -> ClientM Monitor
withMonitor inst agent = do
  envMonitorStart inst configs
  agent
  envMonitorClose inst
  return configs
  where
    configs :: Monitor
    configs = Monitor "/tmp/random-agent-results" True False False


-- * reexports
module OpenAI.Gym
  ( module X
  ) where

import OpenAI.Gym.API as X
import OpenAI.Gym.Data as X


-------------------------------------------------------------------------------
-- |
-- Module    :  OpenAI.Gym.API
-- License   :  MIT
-- Stability :  experimental
-- Portability: non-portable
-------------------------------------------------------------------------------
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
module OpenAI.Gym.API where

import OpenAI.Gym.Prelude
import OpenAI.Gym.Data

type GymAPI
  = "v1" :> ( "envs" :> ( ReqBody '[JSON] GymEnv :> Post '[JSON] InstID
                     :<|> Get '[JSON] Environment
                     :<|> Capture "instance_id" InstID :> "reset" :> Post '[JSON] Observation
                     :<|> Capture "instance_id" InstID :> "step"  :> ReqBody '[JSON] Step :> Post '[JSON] Outcome
                     :<|> Capture "instance_id" InstID :> "action_space" :> Get '[JSON] Info
                     :<|> Capture "instance_id" InstID :> "action_space" :> "sample"   :> Get '[JSON] Action
                     :<|> Capture "instance_id" InstID :> "action_space" :> "contains" :> Capture "x" Int :> Get '[JSON] Object
                     :<|> Capture "instance_id" InstID :> "observation_space"  :> Get '[JSON] Info
                     :<|> Capture "instance_id" InstID :> "monitor" :> "start" :> ReqBody '[JSON] Monitor :> Post '[HTML] ()
                     :<|> Capture "instance_id" InstID :> "monitor" :> "close" :> Post '[HTML] ()
                     :<|> Capture "instance_id" InstID :> "close"   :> Post '[HTML] ())
         :<|> "upload" :> ReqBody '[JSON] Config :> Post '[HTML] ()
         :<|> "shutdown" :> Post '[HTML] ())

gymAPI :: Proxy GymAPI
gymAPI = Proxy

envCreate               :: GymEnv  -> ClientM InstID
envListAll              :: ClientM Environment
envReset                :: InstID  -> ClientM Observation
envStep                 :: InstID  -> Step    -> ClientM Outcome
envActionSpaceInfo      :: InstID  -> ClientM Info
envActionSpaceSample    :: InstID  -> ClientM Action
envActionSpaceContains  :: InstID  -> Int     -> ClientM Object
envObservationSpaceInfo :: InstID  -> ClientM Info
envMonitorStart         :: InstID  -> Monitor -> ClientM ()
envMonitorClose         :: InstID  -> ClientM ()
envClose                :: InstID  -> ClientM ()
upload                  :: Config  -> ClientM ()
shutdownServer          :: ClientM ()

(envCreate
  :<|> envListAll
  :<|> envReset
  :<|> envStep
  :<|> envActionSpaceInfo
  :<|> envActionSpaceSample
  :<|> envActionSpaceContains
  :<|> envObservationSpaceInfo
  :<|> envMonitorStart
  :<|> envMonitorClose
  :<|> envClose)
  :<|> upload
  :<|> shutdownServer
  = client gymAPI


-------------------------------------------------------------------------------
-- |
-- Module    :  OpenAI.Gym.Data
-- License   :  MIT
-- Stability :  experimental
-- Portability: non-portable
-------------------------------------------------------------------------------
{-# LANGUAGE DeriveGeneric #-}
module OpenAI.Gym.Data
  ( GymEnv (..)
  , InstID (..)
  , Environment (..)
  , Observation (..)
  , Step (..)
  , Outcome (..)
  , Info (..)
  , Action (..)
  , Monitor (..)
  , Config (..)
  ) where

import OpenAI.Gym.Prelude
import qualified Data.Text  as T
import qualified Data.Aeson as A

data GymEnv
  -- | Classic Control Environments
  = CartPoleV0               -- ^ Balance a pole on a cart (for a short time).
  | CartPoleV1               -- ^ Balance a pole on a cart.
  | AcrobotV1                -- ^ Swing up a two-link robot.
  | MountainCarV0            -- ^ Drive up a big hill.
  | MountainCarContinuousV0  -- ^ Drive up a big hill with continuous control.
  | PendulumV0               -- ^ Swing up a pendulum.

  -- Toy text games
  | FrozenLakeV0             -- ^ Swing up a pendulum.

  -- | Atari Games
  | PongRamV0                -- ^ Maximize score in the game Pong, with RAM as input
  | PongV0                   -- ^ Maximize score in the game Pong
  deriving (Eq, Enum, Ord)

instance Show GymEnv where
  show CartPoleV0              = "CartPole-v0"
  show CartPoleV1              = "CartPole-v1"
  show AcrobotV1               = "Acrobot-v1"
  show MountainCarV0           = "MountainCar-v0"
  show MountainCarContinuousV0 = "MountainCarContinuous-v0"
  show PendulumV0              = "Pendulum-v0"
  show FrozenLakeV0            = "FrozenLake-v0"
  show PongRamV0               = "Pong-ram-v0"
  show PongV0                  = "Pong-v0"

instance ToJSON GymEnv where
  toJSON env = object [ "env_id" .= show env ]

data InstID = InstID !Text
  deriving (Eq, Show, Generic)

instance ToHttpApiData InstID where
  toUrlPiece (InstID i) = i

instance ToJSON InstID where
  toJSON (InstID i) = toSingleton "instance_id" i

instance FromJSON InstID where
  parseJSON = parseSingleton InstID "instance_id"

newtype Environment = Environment { all_envs :: HashMap Text Text }
  deriving (Eq, Show, Generic)

instance ToJSON Environment
instance FromJSON Environment

data Observation = Observation !Value
  deriving (Eq, Show, Generic)

instance ToJSON Observation where
  toJSON (Observation v) = toSingleton "observation" v

instance FromJSON Observation where
  parseJSON = parseSingleton Observation "observation"

data Step = Step
  { action :: !Value
  , render :: !Bool
  } deriving (Eq, Generic, Show)

instance ToJSON Step

data Outcome = Outcome
  { observation :: !Value
  , reward      :: !Double
  , done        :: !Bool
  , info        :: !Object
  } deriving (Eq, Show, Generic)

instance ToJSON Outcome
instance FromJSON Outcome

data Info = Info !Object
  deriving (Eq, Show, Generic)

instance ToJSON Info where
  toJSON (Info v) = toSingleton "info" v

instance FromJSON Info where
  parseJSON = parseSingleton Info "info"

data Action = Action !Value
  deriving (Eq, Show, Generic)

instance ToJSON Action where
  toJSON (Action v) = toSingleton "action" v

instance FromJSON Action where
  parseJSON = parseSingleton Action "action"

data Monitor = Monitor
  { directory      :: !Text
  , force          :: !Bool
  , resume         :: !Bool
  , video_callable :: !Bool
  } deriving (Generic, Eq, Show)

instance ToJSON Monitor

data Config = Config
  { training_dir :: !Text
  , algorithm_id :: !Text
  , api_key      :: !Text
  } deriving (Generic, Eq, Show)

instance ToJSON Config


-------------------------------------------------------------------------------
-- |
-- Module    :  OpenAI.Gym.Client
-- License   :  MIT
-- Stability :  experimental
-- Portability: non-portable
-------------------------------------------------------------------------------
{-# LANGUAGE MultiParamTypeClasses #-}
{-# OPTIONS_GHC -Wno-orphans #-}
module OpenAI.Gym.Prelude
  ( module P
  , parseSingleton
  , toSingleton
  ) where

import Control.Monad              as P
import Control.Monad.Loops        as P
import Control.Monad.Trans.Except as P (runExceptT)
import Control.Monad.Trans.Reader as P
import Control.Monad.Trans.Class  as P
import Data.Aeson                 as P
import Data.HashMap.Strict        as P (HashMap)
import Data.Proxy                 as P
import Data.Text                  as P (Text)
import GHC.Generics               as P
import Network.HTTP.Client        as P hiding (Proxy, responseBody, responseStatus)
import Servant.API                as P
import Servant.Client             as P
import Servant.HTML.Lucid         as P (HTML)
import Control.Monad.IO.Class     as P
import Prelude                    as P
import Data.Aeson.Types (Parser)

instance MimeUnrender HTML () where
    mimeUnrender _ _ = return ()

parseSingleton :: FromJSON a => (a -> b) -> Text -> Value -> Parser b
parseSingleton fn f (Object v) = fn <$> v .: f
parseSingleton fn f _          = mempty

toSingleton :: ToJSON a => Text -> a -> Value
toSingleton f a = object [ f .= toJSON a ]


main :: IO ()
main = putStrLn "Test suite not yet implemented"


{-# OPTIONS_GHC -F -pgmF hspec-discover #-}

module Spec where


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.BubbleSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.BubbleSort

spec :: Spec
spec = do
    describe "bubbleSort" $ do
        it "returns empty list when sorting empty list" $ property $
            bubbleSort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> bubbleSort x == (bubbleSort . bubbleSort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = bubbleSort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = bubbleSort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in bubbleSort unsortedList == sortedList


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.HeapSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.HeapSort

spec :: Spec
spec = do
    describe "heapSort" $ do
        it "returns empty list when sorting empty list" $ property $
            heapSort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> heapSort x == (heapSort . heapSort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = heapSort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = heapSort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in heapSort unsortedList == sortedList


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.InsertionSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.InsertionSort

spec :: Spec
spec = do
    describe "insertionSort" $ do
        it "returns empty list when sorting empty list" $ property $
            insertionSort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> insertionSort x == (insertionSort . insertionSort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = insertionSort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = insertionSort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in insertionSort unsortedList == sortedList


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.MergeSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.MergeSort

spec :: Spec
spec = do
    describe "megeSort" $ do
        it "returns empty list when sorting empty list" $ property $
            mergeSort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> mergeSort x == (mergeSort . mergeSort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = mergeSort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = mergeSort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in mergeSort unsortedList == sortedList


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.QuickSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.QuickSort

spec :: Spec
spec = do
    describe "quicksort" $ do
        it "returns empty list when sorting empty list" $ property $
            quicksort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> quicksort x == (quicksort . quicksort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = quicksort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = quicksort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in quicksort unsortedList == sortedList


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.SelectionSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.SelectionSort

spec :: Spec
spec = do
    describe "selectionSort" $ do
        it "returns empty list when sorting empty list" $ property $
            selectionSort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> selectionSort x == (selectionSort . selectionSort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = selectionSort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = selectionSort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in selectionSort unsortedList == sortedList


{-# LANGUAGE ScopedTypeVariables #-}
module SortSpecs.ShellSortSpec where

import Test.Hspec
import Test.QuickCheck
import Sorts.ShellSort

spec :: Spec
spec = do
    describe "shellSort" $ do
        it "returns empty list when sorting empty list" $ property $
            shellSort [] == ([] :: [Int])

        it "returns same list if input was already sorted" $ property $
            \(x :: [Int]) -> shellSort x == (shellSort . shellSort $ x)

        it "returns list with smallest element at 0" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = shellSort x
                        in head sortedList == minimum sortedList

        it "returns list with largest element at the end" $ property $
            forAll (listOf1 arbitrary) $
                \(x :: [Int]) -> let sortedList = shellSort x
                        in last sortedList == maximum sortedList

        it "handle simple sorting of static value" $
            let (unsortedList :: [Int]) = [4, 2, 1, 7, 3]
                (sortedList :: [Int]) = [1, 2, 3, 4, 7]
            in shellSort unsortedList == sortedList


module HaskellAlgorithms
(
    module Sorts.BubbleSort
) where

import Sorts.BubbleSort


module BinaryTree.BinarySearchTree where

data BTree a = Empty | Node a (BTree a) (BTree a) deriving (Show)
data Side = LeftSide | RightSide deriving (Eq, Show)

-- Function to get the data associated with the node.
nodeKey :: BTree a -> Maybe a
nodeKey Empty = Nothing
nodeKey (Node x _ _) = Just x

-- Perform inorder walk of the binary search tree.
-- Cormen, Thomas H., et al. Introduction to algorithms.  pg. 288, MIT press, 2009.
inorderWalk :: (Eq a, Ord a) => BTree a -> [a]
inorderWalk Empty = []
inorderWalk (Node x l r) = (inorderWalk l) ++ [x] ++ (inorderWalk r)

-- Function to insert a value into the tree. Returns the new tree.
-- Cormen, Thomas H., et al. Introduction to algorithms.  pg. 294, MIT press, 2009.
bstInsert :: (Eq a, Ord a) => BTree a -> a -> BTree a
bstInsert Empty z = Node z Empty Empty
bstInsert (Node x l r) z
    | z < x = Node x (bstInsert l z) r
    | otherwise = Node x l (bstInsert r z)

-- Function to find the maximum value in the BST.
bstMax :: (Eq a, Ord a) => BTree a -> Maybe a
bstMax Empty = Nothing
bstMax (Node x Empty Empty) = Just x
bstMax (Node x l Empty) = Just x
bstMax (Node x l r) = bstMax r

-- Function to find the minimum value in the BST.
bstMin :: (Eq a, Ord a) => BTree a -> Maybe a
bstMin Empty = Nothing
bstMin (Node x Empty Empty) = Just x
bstMin (Node x Empty r) = Just x
bstMin (Node x l r) = bstMin l

-- Function to build BST from a list of values using a fold.
bstFromList :: (Eq a, Ord a) => [a] -> BTree a
bstFromList [] = Empty
bstFromList lst = foldl (\tree elem -> bstInsert tree elem) Empty lst

sampleTree = bstFromList [10, 7, 3, 11, 12, 1, 3, 2]

-- Function to check if a given tree is a Binary Search Tree.
-- Property:
--     x is a node in the BST. If y is a node in the left subtree of x then
--     y.key <= x.key. If y is a node in the right subtree of x then
--     y.key >= x.key.
--     Cormen, Thomas H., et al. Introduction to algorithms. MIT press, 2009.
isBST :: (Ord a, Eq a) => BTree a -> Bool
isBST Empty = True
isBST (Node x Empty Empty) = True
isBST (Node x Empty r) = (x < (nkey r)) && (isBST r) where nkey = (\(Node n ll rr) -> n)
isBST (Node x l Empty) = (x >= (nkey l)) && (isBST l) where nkey = (\(Node n ll rr) -> n)
isBST (Node x l r) = (x >= (nkey l)) && (x < (nkey r)) && (isBST l) && (isBST r) where nkey = (\(Node n ll rr) -> n)


module BinaryTree.BinaryTree where

import qualified Data.List as L

data BTree a = Empty | Node a (BTree a) (BTree a) deriving (Show)
data Side = LeftSide | RightSide deriving (Eq, Show)

-- Get subtree on specified side
getSubTree :: Side -> BTree a -> BTree a
getSubTree _ Empty = Empty
getSubTree s (Node _ l r) = if s == LeftSide then l else r

-- Get Left Subtree
getLeftTree :: BTree a -> BTree a
getLeftTree Empty = Empty
getLeftTree (Node _ l _) = l

-- Get Right Subtree
getRightTree :: BTree a -> BTree a
getRightTree Empty = Empty
getRightTree (Node _ _ r) = r

-- Get string representation of node Data
nodeShow :: (Show a) => BTree a -> String
nodeShow Empty = ""
nodeShow (Node val _ _) = show val

-- Depth first traversal
dfsList :: BTree a -> [a]
dfsList Empty = []
dfsList (Node n l r) = [n] ++ (dfsList l) ++ (dfsList r)

-- Breadth first traversal.
bfsList :: BTree a -> [a]
bfsList Empty = []
bfsList t = concat $ takeWhile (\l -> (length l) > 0) [getLevel i 0 t | i <- [0..]]

-- Get all nodes from a single level in the tree.
getLevel :: (Num b, Enum b, Eq b) => b -> b -> BTree a -> [a]
getLevel _ _ Empty = []
getLevel 0 _ (Node n l r) = [n]
getLevel level i (Node n l r)
    | i == level = [n]
    | otherwise = (getLevel level (i+1) l) ++ (getLevel level (i+1) r)

-- Get a list of lists of nodes in each level
getLevels :: BTree a -> [[a]]
getLevels t = takeWhile (\l -> (length l) > 0) [getLevel i 0 t | i <- [0..]]

-- Get the depth of the tree
getDepth :: BTree a -> Int
getDepth t = length $ getLevels t

-- Generate a Binary Tree from a list of values.
-- Assume list is in breadth first order.
fromList :: [a] -> BTree a
fromList lst = fromListInt 0 lst
-- Internal function to convert list to tree.
fromListInt :: Int -> [a] -> BTree a
fromListInt _ [] = Empty
fromListInt i lst@(x:xs) = Node x (fromListInt (2*i + 1) (drop (i+1) lst))
                                  (fromListInt (2*i + 2) (drop (i+2) lst))

-- Count number of nodes in the tree.
numNodes :: BTree a -> Int
numNodes t = length $ bfsList t

-- Pretty Print a Binary Tree
simplePrint :: (Show a) => BTree a -> String
simplePrint Empty = ""
simplePrint t = (nodeShow t) ++ " " ++ (simplePrint $ getLeftTree t) ++ (simplePrint $ getRightTree t)


module DataStructures.MaxHeap where

-- Construct a max heap (represented as a list) from a given input list
maxHeap :: Ord a => [a] -> [a]
maxHeap [] = []
maxHeap [x] = [x]
maxHeap xs = maxHeap' xs (div (length xs) 2)

-- Insert an element into a heap and then fix the heap
insert :: Ord a => [a] -> a -> [a]
insert [] e = [e]
insert xs e = upHeapify (xs ++ [e]) (length xs)

-- Recursive helper function for maxHeap
maxHeap' :: Ord a => [a] -> Int -> [a]
maxHeap' xs 0 = maxHeapify xs 0
maxHeap' xs i = maxHeap' (maxHeapify xs i) (i-1)

-- Given a heap and an index (i), move the element at i up the heap until it is in the right position
upHeapify :: Ord a => [a] -> Int -> [a]
upHeapify xs 0 = xs
upHeapify xs i | xs !! i <= xs !! ip = xs
               | otherwise = upHeapify (swap xs i ip) ip
                where
                    ip = parent i

-- Given a heap and an index (i), move the element at i down the heap until it is in the right position
maxHeapify :: (Ord a) => [a] -> Int -> [a]
maxHeapify xs i | not (hasAnyChild xs i) = xs
                | hasLeftChild xs i && hasRightChild xs i =
                    let largestChild = if xs !! (leftChild i) > xs !! (rightChild i) then leftChild i else rightChild i
                    in if xs !! largestChild > xs !! i then maxHeapify (swap xs largestChild i) largestChild else xs
                | hasLeftChild xs i =
                    if xs !! (leftChild i) > xs !! i then maxHeapify (swap xs (leftChild i) i) (leftChild i) else xs
                | otherwise =
                    if xs !! (rightChild i) > xs !! i then maxHeapify (swap xs (rightChild i) i) (rightChild i) else xs

-- Remove the largest element from the heap, and then fix the heap
extract :: Ord a => [a] -> [a]
extract [] = []
extract [_] = []
extract xs = maxHeapify ((last xs):take (length xs - 2) (tail xs)) 0

-- Delete an element from a heap and then fix the heap
-- If the element is not present, return the heap as is
delete :: Ord a => [a] -> a -> [a]
delete [] _ = []
delete xs e = let i' = (search xs e)
                  f = if xs !! (length xs - 1) > e then upHeapify else maxHeapify
                  in case i' of
                      Nothing -> xs
                      Just i  -> f (take ((length xs) - 1) (swap xs i ((length xs) - 1))) i

-- Find the index of an element in a heap
-- Returns Nothing if the element is not present
search :: Eq a => [a] -> a -> Maybe Int
search xs e = search' xs e 0

-- Recursive helper function for search
search' :: Eq a => [a] -> a -> Int -> Maybe Int
search' [] _ _ = Nothing
search' (x:xs) e i | e == x = Just i
                   | otherwise = search' xs e (i+1)

-- Swap the values stored in two positions in a list
swap :: [a] -> Int -> Int -> [a]
swap xs i1 i2 = map snd . foldr (\x a ->
        if fst x == i1 then ys !! i2 : a
        else if fst x == i2 then ys !! i1 : a
        else x : a) [] $ ys
    where ys = zip [0..] xs

-- Index where the left child of the node at the given index should be located
leftChild :: Int -> Int
leftChild i = 2 * i + 1

-- Index where the right child of the node at the given index should be located
rightChild :: Int -> Int
rightChild i = 2 * i + 2

-- Index where the parent of the node at the given index should be located
parent ::  Int -> Int
parent i = floor (fromIntegral (div (i-1) 2))

-- Helper functions to determine which, if any, child nodes are present in the heap for the node at a given index
hasLeftChild, hasRightChild, hasAnyChild :: [a] -> Int -> Bool
hasLeftChild xs i = leftChild i < length xs

hasRightChild xs i = rightChild i < length xs

hasAnyChild xs i = hasLeftChild xs i || hasRightChild xs i


module Graph.Dfs where

import Data.List

type Node = Int
type Edge = (Node,Node)
type Graph = ([Node],[Edge])

depthfirst :: Graph -> Node -> [Node]
depthfirst (v,e) n
    | [x|x<-v,x==n] == [] = []
    | otherwise = dfrecursive (v,e) [n]

dfrecursive :: Graph -> [Node] -> [Node]
dfrecursive ([],_) _ = []
dfrecursive (_,_) [] = []
dfrecursive (v,e) (top:stack)
    | [x|x<-v,x==top] == [] = dfrecursive (newv, e) stack
    | otherwise = top : dfrecursive (newv, e) (adjacent ++ stack)
    where
        adjacent = [x | (x,y)<-e,y==top] ++ [x | (y,x)<-e,y==top]
        newv = [x|x<-v,x/=top]

connectedcomponents :: Graph -> [[Node]]
connectedcomponents ([],_) = []
connectedcomponents (top:v,e)
    | remaining == [] = [connected]
    | otherwise = connected : connectedcomponents (remaining, e)
    where
        connected = depthfirst (top:v,e) top
        remaining = (top:v) \\ connected

dfsbipartite :: Graph -> [(Node, Int)] -> [Node] -> [Node] -> Bool
dfsbipartite ([],_) _ _ _ = True
dfsbipartite (_,_) [] _ _ = True
dfsbipartite (v,e) ((nv, 0):stack) odd even
    | [x|x<-v,x==nv] == [] = dfsbipartite (v, e) stack odd even
    | [] == intersect adjacent even = dfsbipartite (newv, e) ([(x,1)|x<-adjacent] ++ stack) odd (nv : even)
    | otherwise = False
    where
        adjacent = [x | (x,y)<-e,y==nv] ++ [x | (y,x)<-e,y==nv]
        newv = [x|x<-v,x/=nv]
dfsbipartite (v,e) ((nv, 1):stack) odd even
    | [x|x<-v,x==nv] == [] = dfsbipartite (v, e) stack odd even
    | [] == intersect adjacent odd = dfsbipartite (newv, e) ([(x,0)|x<-adjacent] ++ stack) (nv : odd) even
    | otherwise = False
    where
        adjacent = [x | (x,y)<-e,y==nv] ++ [x | (y,x)<-e,y==nv]
        newv = [x|x<-v,x/=nv]

bipartite :: Graph -> Bool
bipartite ([],_) = True
bipartite (top:v,e) = dfsbipartite (top:v, e) [(top,0)] [] []


{-# LANGUAGE TupleSections #-}
module Graph.DirectedGraph where

import qualified Data.Set as Set
import qualified Data.Map as Map

type Graph a = Map.Map a (Set.Set a)

empty :: Ord a => Graph a
empty = Map.empty

insertNode :: Ord a => a -> Graph a -> Graph a
insertNode u g
    | Map.member u g = g
    | otherwise      = Map.insert u Set.empty g

insertEdge :: Ord a => (a, a) -> Graph a -> Graph a
insertEdge (u, v) = Map.insertWith Set.union u (Set.singleton v) . insertNode u . insertNode v

deleteNode :: Ord a => a -> Graph a -> Graph a
deleteNode u = Map.delete u . Map.map (Set.delete u)

adjacentNodes :: Ord a => a -> Graph a -> Maybe (Set.Set a)
adjacentNodes = Map.lookup

nodes :: Ord a => Graph a -> Set.Set a
nodes = Map.keysSet

edges :: Ord a => Graph a -> [(a, a)]
edges = concat . Map.elems . Map.mapWithKey (\u -> map (u,) . Set.toList)

-- Creates graph from list of edges
fromList :: Ord a => [(a, a)] -> Graph a
fromList = foldr insertEdge Map.empty

-- Creates new graph by flipping every edge
transpose :: Ord a => Graph a -> Graph a
transpose = fromList . map (\(a, b)->(b, a)) . edges


module Maths.Factorial where

fac :: Integer -> Integer
fac 0 = 1
fac n = n * fac (n - 1)

main :: IO ()
main = do
    print (fac 4)


module Maths.Fibonacci where

fib :: Integer -> Integer
fib 0 = 0
fib 1 = 1
fib n = fib (n-1) + fib (n-2)

main :: IO ()
main = do
    print (fib 10)


module Maths.GraphDist where

-- calculates the distance from the origin
calcDistOrigin :: Floating a => a -> a -> a -> a
calcDistOrigin x y z = sqrt ((x ** 2) + (y ** 2) + (z ** 2))

-- calculates the distance between two points
calcDist :: Floating a => a -> a -> a -> a -> a -> a -> a
calcDist x1 x2 y1 y2 z1 z2 = sqrt ((x2 - x1) ** 2 + (y2 - y1) ** 2 + (z2 - z1) ** 2)

main :: IO ()
main = do
  print (calcDistOrigin 4 5 6)
  print (calcDist 4 6 5 9 6 10)


module Maths.KadaneAlgorithm where

-- Sample Input List
inputList :: [Integer]
inputList = [1, -2, 3, 4, -5, 6, -7, 8, 9, 10]

-- Helper function that updates values of maxSoFar and maxEndingHere
-- and call itself recursively over the values of the list
-- and when the list is empty, return the maxSoFar value
-- with the start and end indices.
maxSubArrayHelper :: Integer -> Integer -> Integer -> Integer -> Integer -> [Integer] -> (Integer, Integer, Integer)
maxSubArrayHelper maxSoFar _ _ start end [] = (maxSoFar, start, end)
maxSubArrayHelper maxSoFar maxEndingHere i start end (x:xs) =
    let i' = i + 1
        maxEndingHere' = maxEndingHere + x
    in
        if maxSoFar < maxEndingHere' then
            maxSubArrayHelper maxEndingHere' maxEndingHere' i' start i xs
        else if maxEndingHere' < 0 then
            maxSubArrayHelper maxSoFar 0 i' i' i' xs
        else
            maxSubArrayHelper maxSoFar maxEndingHere' i' start end xs

-- Initially maxSoFar (maximum sum till the previous iteration),
-- maxEndingHere (maximum sum till end index of the current iteration),
-- start (start index) and end (end index) are sent as 0
maxSubArray :: [Integer] -> (Integer, Integer, Integer)
maxSubArray = maxSubArrayHelper 0 0 0 0 0

-- Outputs (sum, start, end)
-- sum - sum from start to end indices of the input array
-- start:end - the subarray with max sum
main :: IO ()
main = do
    print (maxSubArray inputList)


module Maths.Palindrome where

palinedrome :: Eq a => [a] -> Bool
palinedrome xs = (xs == reverse xs)

main :: IO ()
main = do
  print (palinedrome [1, 3, 1])


module Misc.BinarySearch where

bsWithIndex :: (Ord a) => [a] -> a -> Int -> Maybe Int
bsWithIndex list n i
    | n == head list        = Just i
    | len == 1              = Nothing   -- only candidate in list is not the right elem
    | n < head ys           = bsWithIndex xs n i
    | otherwise             = bsWithIndex ys n (i + half)
    where
        len         = length list
        half        = len `div` 2
        (xs, ys)    = splitAt half list

bs :: (Ord a) => [a] -> a -> Int
bs list n = case bsWithIndex list n 0 of
    Just x  -> x
    Nothing -> -1

main :: IO ()
main = do
    let intList = [1,4,7,10,25,30]
    print $ bs intList 29  -- 29 -> -1 as not in list
    print $ bs intList 7   --  7 ->  2 as in list


module Misc.NQueens where
{-
The n-Queens search is a backtracking algorithm. The n-Queens problem fits n
number of Queens onto a Chess board. As only one Queen can fit per row, a
one-dimensional integer array is used to represent the Queen's offset on each
row.
-}

import Data.List (permutations)

main = nqueens 8
nqueens size = mapM_ (printBoard size) $ take 1 $ filter (evaluateBoard size) $ board_permutations size

--N sized Chess boards are represented as a one-dimension array.
board_permutations size = permutations [0..size - 1]

--Count the number of valid boards for a specified Chess board size.
count_boards size = length $ filter (evaluateBoard size) $ board_permutations size

--Show every valid board
nqueens_list size = mapM_ (printBoard size) $ filter (evaluateBoard size) $ board_permutations size

--Board printing function
printBoard size board = do
    printBoard2 size board
    putStrLn "" where
        printBoard2 _ [] = return ()
        printBoard2 size board = do
            let row = head board
            printRow size row
            printBoard2 size $ tail board

printRow size row = do
    let lstring = (replicate row ". ")
    let rstring = replicate (size - row - 1) ". "
    putStrLn $ concat (lstring ++ ["Q "] ++ rstring)
    return ()

--Recursively check that prior rows are valid.
evaluateBoard _ [] = True
evaluateBoard size rows = (evaluateBoard size $ cut_last rows) && validate size (cut_last rows) (last_row - 1) (last_row + 1) last_row where
    last_row = last rows

--Validate that a Queen on a row doesn't have conflicts with earlier rows.
validate _ [] _ _ _ = True
validate size rows left right position = if check_row == left || check_row == right || check_row == position then False else validate size (cut_last rows) (left - 1) (right + 1) position where
    check_row = last rows

cut_last x = reverse $ drop 1 $ reverse x


module Misc.Powerset where

powerset :: [a] -> [[a]]
powerset [] = [[]]
powerset (x:xs) = (powerset xs) ++ (map (\ys -> x:ys) (powerset xs))

xs = [1,2,3,4]

main :: IO ()
main = do
    putStrLn (show (powerset xs))


module Misc.TowersOfHanoi where

hanoi :: (Eq a, Num a) => a -> [Char] -> [Char] -> [Char] -> IO ()
hanoi 0 _ _ _ = return ()
hanoi n startPole intermediatePole endPole = do
    hanoi (n - 1) startPole endPole intermediatePole
    putStrLn ("Move from " ++ startPole ++ " to " ++ endPole)
    hanoi (n - 1) intermediatePole startPole endPole

main :: IO ()
main = do
    hanoi 3 "startPole" "intermediatePole" "endPole"


module ProjectEuler.Problem1.Problem1 where

solList = filter (\n -> (rem n 5 == 0) || (rem n 3 == 0)) [1..999]

main = do
    print $ sum solList


module ProjectEuler.Problem2.Problem2 where

fib :: Integer -> [Integer]
fib n
    | n < 0 = []
    | n == 1 = [0]
    | n == 2 = [0, 1]
    | otherwise = reverse $ foldl (\acc n -> (sum (take 2 acc)):acc) [1, 0] [3..n]

main = do
    print $ sum $ filter even $ takeWhile (<=4000000) (fib 100)


module ProjectEuler.Problem3.Problem3 where

largestPrimeFactor :: Integer -> Integer -> Integer
largestPrimeFactor divi val
  | val `mod` divi == 0 = if divi == val then val else largestPrimeFactor 2 $ val `div` divi
  | val `mod` divi /= 0 = largestPrimeFactor (divi + 1) val

main = do
  print $ largestPrimeFactor 2 600851475143


module ProjectEuler.Problem4.Problem4 where

isPalindrome :: Integer -> Bool
isPalindrome i =
  let strInt = show i
      front = take (length strInt `div` 2) strInt
      back = take (length strInt `div` 2)  $ reverse strInt
  in front == back

main = do
  print $ maximum [i * j | i <- [100..999], j <- [100..999], isPalindrome (i * j)]


module ProjectEuler.Problem5.Problem5 where

import Data.List

primeFactors' :: Integer -> Integer -> [Integer]
primeFactors' _ 1 = [1]
primeFactors' d n
  | d * d > n       = [n]
  | n `rem` d == 0  = d : primeFactors' d (n `div` d)
  | otherwise       = primeFactors' (d + 1) n

primeFactors :: Integer -> [Integer]
primeFactors x = primeFactors' 2 x

combinedPrimeFactorisationTo :: Integer -> [Integer]
combinedPrimeFactorisationTo n
  | n == 1    = [1]
  | otherwise = furtherPrimes ++ (nPrimes \\ furtherPrimes)
      where
        furtherPrimes = combinedPrimeFactorisationTo (n - 1)
        nPrimes       = primeFactors n

main :: IO ()
main = do
  print $ product $ combinedPrimeFactorisationTo 20


module ProjectEuler.Problem6.Problem6 where

sumSquare :: [Integer] -> Integer
sumSquare = sum . map (^2)

squareSum :: [Integer] -> Integer
squareSum = (^2) . sum

main :: IO ()
main = do
    let l = [1..100]
    putStrLn $ show $ (squareSum l) - (sumSquare l)


module ProjectEuler.Problem7.Problem7 where

primes :: [Integer]
primes = sieve [2..]
    where sieve (p:xs) = p : sieve [x | x <- xs, rem x p > 0]

main :: IO ()
main = do
    putStrLn $ show $ primes !! 10000 -- zero indexing


module Robotics.ComplementaryFilter.CompFilt where

import Robotics.ComplementaryFilter.TestData

-- Utility functions to extract X, Y, Z components from 3D vector.
getX :: (a, a, a) -> a
getX (x,_,_) = x

getY :: (a, a, a) -> a
getY (_,y,_) = y

getZ :: (a, a, a) -> a
getZ (_,_,z) = z

-- Extract accel data from list of floats
getAccel :: (RealFloat a) => [a] -> (a, a, a)
getAccel [] = (0, 0, 0)
getAccel s = if length s >= 6
                then (s!!0, s!!1, s!!2)
                else (0, 0, 0)

-- Extract gyro data from a lsit of floats
getGyro :: (RealFloat a) => [a] -> (a, a, a)
getGyro s = if length s >= 6
               then (s!!3, s!!4, s!!5)
               else (0, 0, 0)

-- Function to calculate tilt angle from accelerometer reading.
-- By default the tilt measurement is made around the Z axis.
accelTiltAngle :: (RealFloat a) => (a, a, a) -> a
accelTiltAngle (_, y, z) = (atan2 z y)*180.0/pi

-- Complementary filter, uses the scanl pattern.
compFilt :: (RealFloat a) => [a] -> [a] -> a -> a -> [a]
compFilt Ï‰s Î¸_accs Î± Î´t = scanl (\Î¸ (Ï‰, Î¸_acc) -> Î±*(Î¸ + Ï‰*Î´t) + (1-Î±)*Î¸_acc)
                                (head Î¸_accs)
                                (zip Ï‰s Î¸_accs)

-- Calculate tilts
calcTilt :: (RealFloat a) => [(a, a, a)] -> [(a, a, a)] -> a -> a -> [a]
calcTilt accel gyro Î± Î´t = compFilt (map getX gyro) (map accelTiltAngle accel) Î± Î´t

main = do
    let accels = map getAccel testData
    let gyros  = map getGyro testData
    let tilts = calcTilt accels gyros 0.95 0.01
    print tilts


module Sorts.BubbleSort where

listToSort :: [Int]
listToSort = [13, 2, 3, 14, 17, 4, 1, 5, 16, 12, 9, 10, 15, 8, 7, 11, 18, 19, 6, 20]

-- The bubble sort function
bubbleSort :: (Ord a) => [a] -> [a]
bubbleSort lst = if bpassed == lst then lst
                 else bubbleSort bpassed
                 where bpassed = bubblePass lst

-- A single pass of bubble sort
bubblePass :: (Ord a) => [a] -> [a]
bubblePass [] = [] -- Empty list is empty.
bubblePass [x] = [x] -- Singleton list is always trivially sorted.
bubblePass (x1:x2:xs) = if x1 > x2
                        then [x2] ++ (bubblePass ([x1] ++ xs))
                        else [x1] ++ (bubblePass ([x2] ++ xs))

main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (bubbleSort listToSort)


module Sorts.HeapSort where

listToSort :: [Integer]
listToSort = [13, 2, 3, 14, 17, 4, 1, 5, 16, 12, 9, 10, 15, 8, 7, 11, 18, 19, 6, 20]

-- Sort the input list with the heap sort algorithm
heapSort :: (Ord a) => [a] -> [a]
heapSort [] = []
heapSort [x] = [x]
heapSort xs = heapSort' (maxHeap xs) []

-- Recursive helper function for heapSort
heapSort' :: Ord a => [a] -> [a] -> [a]
heapSort' [] out = out
heapSort' h out = heapSort' (extract h) out ++ [head h]

-- Construct a max heap (represented as a list) from a given input list
maxHeap :: Ord a => [a] -> [a]
maxHeap [] = []
maxHeap [x] = [x]
maxHeap xs = maxHeap' xs (div (length xs) 2)

-- Recursive helper function for maxHeap
maxHeap' :: Ord a => [a] -> Int -> [a]
maxHeap' xs 0 = maxHeapify xs 0
maxHeap' xs i = maxHeap' (maxHeapify xs i) (i-1)

-- Given a heap and an index (i), move the element at i down the heap until it is in the right position
maxHeapify :: (Ord a) => [a] -> Int -> [a]
maxHeapify xs i | not (hasAnyChild xs i) = xs
                | hasLeftChild xs i && hasRightChild xs i =
                    let largestChild = if xs !! (leftChild i) > xs !! (rightChild i) then leftChild i else rightChild i
                    in if xs !! largestChild > xs !! i then maxHeapify (swap xs largestChild i) largestChild else xs
                | hasLeftChild xs i =
                    if xs !! (leftChild i) > xs !! i then maxHeapify (swap xs (leftChild i) i) (leftChild i) else xs
                | otherwise =
                    if xs !! (rightChild i) > xs !! i then maxHeapify (swap xs (rightChild i) i) (rightChild i) else xs

-- Remove the largest element from the heap, and then fix the heap
extract :: Ord a => [a] -> [a]
extract [] = []
extract [_] = []
extract xs = maxHeapify ((last xs):take (length xs - 2) (tail xs)) 0

-- Swap the values stored in two positions in a list
swap :: [a] -> Int -> Int -> [a]
swap xs i1 i2 = map snd . foldr (\x a ->
        if fst x == i1 then ys !! i2 : a
        else if fst x == i2 then ys !! i1 : a
        else x : a) [] $ ys
    where ys = zip [0..] xs

-- Index where the left child of the node at the given index should be located
leftChild :: Int -> Int
leftChild i = 2 * i + 1

-- Index where the right child of the node at the given index should be located
rightChild :: Int -> Int
rightChild i = 2 * i + 2

-- Helper functions to determine which, if any, child nodes are present in the heap for the node at a given index
hasLeftChild, hasRightChild, hasAnyChild :: [a] -> Int -> Bool
hasLeftChild xs i = leftChild i < length xs

hasRightChild xs i = rightChild i < length xs

hasAnyChild xs i = hasLeftChild xs i || hasRightChild xs i

main :: IO ()
main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (heapSort listToSort)


module Sorts.InsertionSort where

listToSort = [13, 2, 3, 14, 17, 4, 1, 5, 16, 12, 9, 10, 15, 8, 7, 11, 18, 19, 6, 20]

insertionSort:: (Ord a) => [a] -> [a]
insertionSort [] = [] -- Empty list is empty
insertionSort [x] = [x] -- Singleton lists are trivially sorted.
insertionSort (x:xs) = insert x (insertionSort xs)

-- Assumes that the second argument is an alread-sorted list,
-- and inserts the first argument in the appropriate position
insert :: (Ord a) => a -> [a] -> [a]
insert x [] = [x]
insert x lst@(y:ys) = if x <= y then x:lst else y:(insert x ys)

main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (insertionSort listToSort)


module Sorts.MergeSort where

listToSort = [13, 2, 3, 14, 17, 4, 1, 5, 16, 12, 9, 10, 15, 8, 7, 11, 18, 19, 6, 20]

mergeSort :: (Ord a) => [a] -> [a]
mergeSort [] = [] -- Empty list is empty
mergeSort [x] = [x] -- Singleton lists are trivially sorted.
mergeSort [x, y] = [(min x y), (max x y)]
mergeSort lst = merge (mergeSort leftL) (mergeSort rightL)
                where leftL = take splitPoint lst
                      rightL = drop splitPoint lst
                      splitPoint = (length lst) `div` 2

-- Function to execute a merge of two sorted lists
merge :: (Ord a) => [a] -> [a] -> [a]
merge l1 [] = l1
merge [] l2 = l2
merge lst1@(x:xs) lst2@(y:ys) = if x < y
                                then x:(merge xs lst2)
                                else y:(merge lst1 ys)

main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (mergeSort listToSort)


module Sorts.QuickSort where

listToSort :: [Int]
listToSort = [13, 2, 3, 14, 17, 4, 1, 5, 16, 12, 9, 10, 15, 8, 7, 11, 18, 19, 6, 20]

quicksort :: (Ord a) => [a] -> [a]
quicksort [] = [] -- Empty list is empty.
quicksort [x] = [x] -- Singleton list is always trivially sorted.
quicksort [x, y] = [(min x y), (max x y)]
quicksort (x:xs) =
  quicksort [a | a <- xs, a <= x] ++ [x] ++ quicksort [a | a <- xs, a > x]
  -- x is the pivot, left quicksort returns smaller sorted and right quicksort bigger sorted

main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (quicksort listToSort)


module Sorts.SelectionSort where

listToSort :: [Int]
listToSort = [13, 2, 3, 14, 17, 4, 1, 5, 16, 12, 9, 10, 15, 8, 7, 11, 18, 19, 6, 20]

-- The selection sort function
selectionSort :: (Ord a) => [a] -> [a]
selectionSort [] = []
selectionSort (x:xs) =
    let (y, ys) = leastUnsorted (x:xs)
    in y : selectionSort ys

-- select least element from unsorted list, return it and the rest of unsorted list
leastUnsorted :: (Ord a) => [a] -> (a, [a])
leastUnsorted [x] = (x, [])
leastUnsorted (x:xs) =
    let (y, ys) = leastUnsorted xs
    in if x <= y then (x, xs) else (y, x:ys)

main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (selectionSort listToSort)


module Sorts.ShellSort where
-- Shell sort uses insertion sort for sorting individual sublists
import Sorts.InsertionSort

-- Produces the sequence of gaps to use for sorting a list of size n
gaps :: Int -> [Int]
gaps n = reverse (takeWhile (< n) tokuda)

-- Produces Tokuda's sequence
tokuda :: [Int]
tokuda = [gap n | n <- [1..]]

-- Computes term k in Tokuda's sequence
gap :: (Integral a) => a -> a
gap k = ceiling (gap' k)

-- Recursize helper function for gap
gap' :: (Integral a) => a -> Double
gap' 1 = 1
gap' k = 2.25 * gap' (k - 1) + 1

-- Sort an input sequence using the shell sort algorithm
shellSort :: (Ord a) => [a] -> [a]
shellSort xs = shellSort' xs (gaps (length xs))

-- Recursive helper function for shellSort
shellSort' :: (Ord a) => [a] -> [Int] -> [a]
shellSort' [] _ = []
shellSort' [x] _ = [x]
shellSort' l [] = l
shellSort' l (g:gs) = shellSort' (combine [insertionSort (getSub l g i) | i <- [0..g-1]]) gs

-- Get the sublist of l formed by taking elements with gap g between them, starting at index i
getSub :: (Ord a) => [a] -> Int -> Int -> [a]
getSub [] _ _ = []
getSub l g i = [l !! e | e <- [i, i+g..length l - 1]]

-- Combine a list of sublists into a single list in the correct order
-- Required after individual sublists have been sorted, to rebuild the main list
combine :: [[a]] -> [a]
combine [] = []
combine l@(xs:_)
        | length xs == 0 = []
        | otherwise = [x | (x:_) <- l] ++ combine (map (drop 1) l)

main :: IO ()
main = do
    putStrLn $ "Unsorted: " ++ show listToSort
    putStrLn $ "Sorted: " ++ show (shellSort listToSort)


module SpecializedStructure.MergeFindSet where

import qualified Data.Map as Map
import qualified Data.Set as Set

type Mfset a = (Map.Map a a, Map.Map a Int)

fromSet :: Ord a => Set.Set a -> Mfset a
fromSet ns = (fathers, ranks)
    where
        fathers = Map.fromList [(i, i) | i <- Set.toList ns]
        ranks   = Map.fromList [(i, 0) | i <- Set.toList ns]

find :: Ord a => a -> Mfset a -> a
find x mfset = let father = fst mfset Map.! x
               in if father == x then x else find father mfset

merge :: Ord a => a -> a -> Mfset a -> Mfset a
merge x y mfset
    | x' == y'  = mfset
    | rx == ry  = (x' `connectTo` y', Map.adjust (+1) y' ranks)
    | rx < ry   = (x' `connectTo` y', ranks)
    | otherwise = (y' `connectTo` x', ranks)
    where
        (fathers, ranks) = mfset
        connectTo a b = Map.adjust (const b) a fathers
        findWithRank a = (\f -> (f, ranks Map.! f)) $ find a mfset
        (x', rx) = findWithRank x
        (y', ry) = findWithRank y


module Statistics.Center where

import qualified Data.Sort as S
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Algorithms.Intro as VGAI

-- Measures of central tendency.

arithmeticMean :: (Foldable t, Fractional a) => t a -> a
arithmeticMean vals = (sum vals) / (fromIntegral $ length vals)

geometricMean :: (Foldable t, Floating a) => t a -> a
geometricMean vals = (product vals) ** (1/(fromIntegral $ length vals))

harmonicMean :: (Foldable t, Functor t, Fractional a) => t a -> a
harmonicMean vals = (sum $ fmap (1/) vals) / (fromIntegral $ length vals)

-- For median, since the containers are sorted differently, we need to use
-- different methods

medianList :: (Fractional a, Ord a) => [a] -> a
medianList = medianListSorted . S.sort

medianVector
  :: (VG.Vector vector a, Foldable vector, Fractional a, Ord a)
  => vector a -> a
medianVector = medianVectorSorted . VG.modify VGAI.sort

-- When sorted, the two median algorithms are quite similar. We can reduce
-- duplication by adding an export list to the module and using more
-- higher-order functions but let's leave this for a different PR.

medianListSorted :: Fractional a => [a] -> a
medianListSorted vals
  | odd n = vals !! mid
  | otherwise = arithmeticMean $ take 2 $ drop (mid - 1) vals
  where
    n = length vals
    mid = n `div` 2

medianVectorSorted
  :: (VG.Vector vector a, Foldable vector, Fractional a)
  => vector a -> a
medianVectorSorted vals
  | odd n = vals VG.! mid
  | otherwise = arithmeticMean $ VG.take 2 $ VG.drop (mid - 1) vals
  where
    n = length vals
    mid = n `div` 2


module Statistics.Dispersion where

import Statistics.Center

variance :: (Foldable t, Functor t, Fractional a) => t a -> a
variance vals = (sum $ fmap (\x -> x * x) deviations) / n
    where n = (fromIntegral $ length vals)
          mu = arithmeticMean vals
          deviations = fmap (\x -> x-mu) vals

stdev :: (Foldable t, Functor t, Floating a) => t a -> a
stdev vals = sqrt $ variance vals