{-# LANGUAGE OverloadedStrings #-}

module Test.BarAggregator
(
  unitTests,
  properties
) where

import           ATrade.BarAggregator
import           ATrade.RoboCom.Types
import           ATrade.Types
import           Data.List
import           Data.List.NonEmpty   (NonEmpty (..))
import qualified Data.List.NonEmpty   as NE
import qualified Data.Map.Strict      as M
import qualified Data.Text            as T
import           Data.Time.Calendar
import           Data.Time.Clock
import           Safe

import           Hedgehog             as HH
import qualified Hedgehog.Gen         as Gen
import qualified Hedgehog.Range       as Range

import           Test.Tasty
import           Test.Tasty.Hedgehog
import           Test.Tasty.HUnit


unitTests = testGroup "BarAggregator" [
    testUnknownBarSeries
  , testOneTick
  , testTwoTicksInSameBar
  , testTwoTicksInDifferentBars
  ]

properties = testGroup "BarAggregator" [
    prop_allTicksInOneBar
  , prop_threeBars
  ]

secParams = InstrumentParameters "TEST_TICKER" 1 0.01

genTick :: T.Text -> UTCTime -> Int -> Gen Tick
genTick tickerId baseTime timeframe = do
  ts <- generateTimestampInsideBar baseTime timeframe
  val <- fromIntegral <$> Gen.int (Range.linear 1 1000000)
  vol <- Gen.integral (Range.linear 1 1000000)
  return $ Tick tickerId LastTradePrice ts (fromDouble $ val / 1000) vol
  where
    generateTimestampInsideBar base timeframe =
      flip addUTCTime base .
      fromRational .
      toRational .
      picosecondsToDiffTime <$> Gen.integral (Range.linear 0 (truncate 1e12 * fromIntegral timeframe))

mkAggregator :: TickerId -> Int -> BarAggregator
mkAggregator tickerId tf = mkAggregatorFromBars (M.singleton tickerId (BarSeries tickerId (BarTimeframe tf) [] secParams)) [(0, 86400)]

assertBarCorrespondence :: (MonadTest m) => Bar -> NE.NonEmpty Tick -> m ()
assertBarCorrespondence bar ticks = do
  barHigh bar === maximum (value <$> sortedTicks)
  barLow bar === minimum (value <$> sortedTicks)
  barOpen bar === value (NE.head sortedTicks)
  barClose bar === value (NE.last sortedTicks)
  barVolume bar === sum (volume <$> sortedTicks)
  where
    sortedTicks = NE.fromList . sortOn timestamp . NE.toList $ ticks

testUnknownBarSeries :: TestTree
testUnknownBarSeries = testCase "Tick with unknown ticker id" $ do
  let agg = BarAggregator M.empty M.empty [(0, 86400)]
  let (mbar, newagg) = handleTick tick agg
  mbar @?= Nothing
  (bars newagg) @?= M.empty
  where
    testTimestamp = (UTCTime (fromGregorian 1970 1 1) 100)
    tick = Tick {
            security = "TEST_TICKER",
            datatype = LastTradePrice,
            timestamp = testTimestamp,
            value = fromDouble 12.00,
            volume = 1 }

testOneTick :: TestTree
testOneTick = testCase "One tick" $ do
  let series = BarSeries "TEST_TICKER" (BarTimeframe 60) [] secParams
  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
  let (mbar, newagg) = handleTick tick agg
  mbar @?= Nothing
  (bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg)) @?= Just [Bar "TEST_TICKER" testTimestamp 12.00 12.00 12.00 12.00 1]
  where
    testTimestamp = (UTCTime (fromGregorian 1970 1 1) 60)
    tick = Tick {
            security = "TEST_TICKER",
            datatype = LastTradePrice,
            timestamp = testTimestamp,
            value = fromDouble 12.00,
            volume = 1 }

testTwoTicksInSameBar :: TestTree
testTwoTicksInSameBar = testCase "Two ticks - same bar" $ do
  let series = BarSeries "TEST_TICKER" (BarTimeframe 60) [] secParams
  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
  let (mbar, newagg) = handleTick (tick testTimestamp1 12.00) agg
  mbar @?= Nothing
  let (mbar', newagg') = handleTick (tick testTimestamp2 14.00) newagg
  mbar' @?= Nothing
  (bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg')) @?= Just [Bar "TEST_TICKER" testTimestamp2 12.00 14.00 12.00 14.00 2]
  where
    testTimestamp1 = (UTCTime (fromGregorian 1970 1 1) 58)
    testTimestamp2 = (UTCTime (fromGregorian 1970 1 1) 59)
    tick ts val = Tick {
            security = "TEST_TICKER",
            datatype = LastTradePrice,
            timestamp = ts,
            value = fromDouble val,
            volume = 1 }

testTwoTicksInDifferentBars :: TestTree
testTwoTicksInDifferentBars = testCase "Two ticks - different bar" $ do
  let series = BarSeries "TEST_TICKER" (BarTimeframe 60) [] secParams
  let agg = mkAggregatorFromBars (M.fromList [("TEST_TICKER", series)]) [(0, 86400)]
  let (mbar, newagg) = handleTick (tick testTimestamp1 12.00) agg
  mbar @?= Nothing
  let (mbar', newagg') = handleTick (tick testTimestamp2 14.00) newagg
  mbar' @?= Just (Bar "TEST_TICKER" barEndTime 12.00 12.00 12.00 12.00 1)

  (bsBars <$> (M.lookup "TEST_TICKER" $ bars newagg')) @?= Just [Bar "TEST_TICKER" testTimestamp2 14.00 14.00 14.00 14.00 1, Bar "TEST_TICKER" barEndTime 12.00 12.00 12.00 12.00 1]
  where
    testTimestamp1 = UTCTime (fromGregorian 1970 1 1) 58
    barEndTime = UTCTime (fromGregorian 1970 1 1) 60
    testTimestamp2 = UTCTime (fromGregorian 1970 1 1) 61
    tick ts val = Tick {
            security = "TEST_TICKER",
            datatype = LastTradePrice,
            timestamp = ts,
            value = fromDouble val,
            volume = 1 }

prop_allTicksInOneBar :: TestTree
prop_allTicksInOneBar = testProperty "All ticks in one bar" $ property $ do
  tf <- forAll $ Gen.integral (Range.constant 1 86400)
  ticks <- forAll $ Gen.list (Range.linear 1 100) (genTick "TEST_TICKER" baseTime tf)
  let ticks' = sortOn timestamp ticks
  let (newbars, agg) = handleTicks ticks' (mkAggregator "TEST_TICKER" tf)
  let (Just lastBar) = currentBar "TEST_TICKER" agg
  HH.assert $ null newbars
  assertBarCorrespondence lastBar $ NE.fromList ticks
  where
    currentBar tickerId agg = headMay =<< (bsBars <$> M.lookup tickerId (bars agg))
    baseTime = UTCTime (fromGregorian 1970 1 1) 0

prop_threeBars :: TestTree
prop_threeBars = testProperty "Three bars" $ property $ do
  tf <- forAll $ Gen.integral (Range.constant 1 86400)

  ticks1 <- forAll $ Gen.list (Range.linear 1 100) (genTick "TEST_TICKER" baseTime tf)

  let secondBarBaseTime = addUTCTime (fromIntegral tf) baseTime
  ticks2 <- forAll $ Gen.list (Range.linear 1 100) (genTick "TEST_TICKER" secondBarBaseTime tf)

  let thirdBarBaseTime = addUTCTime (fromIntegral $ 2 * tf) baseTime
  ticks3 <- forAll $ Gen.list (Range.linear 1 100) (genTick "TEST_TICKER" thirdBarBaseTime tf)

  let ticks' = sortOn timestamp $ ticks1 <> ticks2 <> ticks3
  let ([secondBar, firstBar], agg) = handleTicks ticks' (mkAggregator "TEST_TICKER" tf)

  assertBarCorrespondence firstBar (NE.fromList ticks1)
  assertBarCorrespondence secondBar (NE.fromList ticks2)

  barTimestamp firstBar === secondBarBaseTime
  barTimestamp secondBar === thirdBarBaseTime

  let (Just lastBar) = currentBar "TEST_TICKER" agg
  assertBarCorrespondence lastBar (NE.fromList ticks3)
  where
    currentBar tickerId agg = headMay =<< (bsBars <$> M.lookup tickerId (bars agg))
    baseTime = UTCTime (fromGregorian 1970 1 1) 0