-- |
-- Copyright : (C) 2013 Parallel Scientific Labs, LLC.
-- License   : GPLv2

{-# LANGUAGE BangPatterns #-}

module Math.Bursts
  ( -- * Data Structures
    Range(..)
  , Window
  , initialWindow
  , defaultAggregate
    -- * Analysis
  , sbtBurstAnalysis
  , sbtBurstAnalysisMany
  ) where

import Control.Monad
import Control.DeepSeq
import Data.Bits
import Data.List (foldl')

-- |Representation of ranges.
data Range = Range
    { rangeStart, rangeEnd, rangeMagn :: !Int
    }
    deriving (Eq, Ord, Show)

instance NFData Range where
  rnf (Range a b c) = a `seq` b `seq` c `seq` ()

-- |Join two ranges that are close to each other.
-- The first range is one with smaller rangeStart. The second range is one with biggest rangeEnd.
joinCloseRanges :: Range -> Range -> Range
joinCloseRanges (Range newStart _ m1) (Range _ newEnd m2) = Range newStart newEnd (max m1 m2)

-- |The window structure.
data Window a = Window
    { windowSamples :: [a] -- ^ Samples
    , windowAverage :: a   -- ^ Average to compute threshold.
    , windowThreeSigma1000 :: a
    -- ^ Three sigma (n*sigma, actually),  multiplied by 1000. Again, it is needed to compute a threshold.
    , windowThresholdAggreg	::	a -> a -> (a, a) -- ^ Threshold aggregation function.
    }

instance NFData a => NFData (Window a) where
  rnf (Window a b c _) = a `seq` b `seq` c `seq` ()

-- |Initial empty window.
initialWindow :: Num a => Int -> a -> a -> (a -> a -> (a,a))-> Window a
initialWindow len avg ts1000 aggreg
    | (len .&. (len-1)) /= 0 = error $ "length "++show len++" is not power of two."
    | otherwise = Window (replicate len 0) avg ts1000 aggreg

defaultAggregate :: (Integral a, Num a) => a -> a -> (a,a)
defaultAggregate average threeSigma1000 =
    (average*2, div (threeSigma1000*1414+500) 1000)

-- |Shifted Binary Tree burst analysis.
-- Returns Maybe Range. Just range indicates a burst detected.
-- Also returns a new window structure.
sbtBurstAnalysis :: (Ord a, Num a, Integral a, Show a) => a -> Window a -> (Window a, Maybe Range)
sbtBurstAnalysis sample window@(Window samples initialAvg initialThreeSigma1000 aggreg) = (newWindow, foundRange)
  where
    !foundRange = analyze 1 (2*initialAvg) (nextThreeSigma initialThreeSigma1000) samplesWithRanges
    newSamples = sample : init samples
    !newWindow = window { windowSamples = newSamples }
    samplesRanges = map (\i -> Range i i 1) [0..]
    nextThreeSigma ts1000 = div (ts1000*1404+500) 1000
    computeThreshold avg ts1000 = avg + div (ts1000 + 500) 1000
    initialThreshold = computeThreshold initialAvg initialThreeSigma1000
    samplesWithRanges = zipWith (\s r -> (False, s >= initialThreshold, s, r)) newSamples samplesRanges
    extractRange mbRange [] = fmap snd mbRange
    extractRange Nothing ((any, burst, _, r) : brss)
      | burst = extractRange (Just (True, r)) brss
      | any = extractRange (Just (False, r)) brss
      | otherwise = extractRange Nothing brss
    extractRange (Just (prevB, prevR)) ((any, burst, _, r) : brss)
      | prevB && not burst = Just prevR
      | prevB && burst = extractRange (Just (True, joinCloseRanges prevR r)) brss
      | burst = extractRange (Just (True, r)) brss
      | otherwise = Just prevR
    analyze magn average threeSigma1000 samplesWithRanges
      | length samplesWithRanges >= 4 = -- trace ("averaged samples "++show averagedSamples) $ trace ("samplesWithRanges "++show samplesWithRanges) $
          analyze (2*magn) nextAvg nextThreeSigma1000 averaged
      | otherwise = -- trace ("final samplesWithRanges "++show samplesWithRanges) $
          extractRange Nothing samplesWithRanges
      where
        (nextAvg, nextThreeSigma1000) = aggreg average threeSigma1000
        thisThreshold = computeThreshold average threeSigma1000
        avg ((any1, b1,s1,r1):bsrs@((any2, b2,s2,r2):_)) =
            (any1 || any2 || b1 || b2, thisB, z, avgR)
            : avg bsrs
          where
            avgR
              | b1 && b2 && thisB = (\r -> r { rangeMagn = magn}) $ joinCloseRanges r1 r2
              | b1 = r1
              | b2 = r2
              | any1 = r1
              | any2 = r2
              | otherwise = joinCloseRanges r1 r2
            thisB = z >= thisThreshold
            z = s1 + s2
        avg _ = []
        combineRanges False ((any1,b1,s1,r1):bsrs@((any2,b2,s2,r2):_))
            = headSample : combineRanges True bsrs
          where
            headSample
              | b1 && b2 = (True, True, s1, joinCloseRanges r1 r2)
              | b1 = (True, True, s1, r1)
              | b2 = (True, True, s1, r2)
              | otherwise = (any1 || any2, False, s1, if any2 then r2 else r1)
        combineRanges True ((any1,b1,s1,r1):((any2,b2,s2,r2):bsrs))
                = combineRanges False (headSample : bsrs)
              where
                headSample
                  | b1 && b2 = (True, True, s2, joinCloseRanges r1 r2)
                  | b1 = (True, True, s2, r1)
                  | b2 = (True, True, s2, r2)
                  | otherwise = (any1 || any2, False, s2, if any2 then r2 else r1)
        combineRanges _ rs = rs
        averagedSamples = avg samplesWithRanges
        averaged = combineRanges False averagedSamples
{-# INLINABLE sbtBurstAnalysis #-}
{-# SPECIALIZE sbtBurstAnalysis :: Int -> Window Int -> (Window Int, Maybe Range) #-}

sbtBurstAnalysisMany :: (Ord a, Num a, Integral a, Show a) => [a] -> Window a -> (Window a, Maybe Range)
sbtBurstAnalysisMany samples window = foldl' analyzeSample (window, Nothing) samples
  where
    analyzeSample (!window, _) sample = sbtBurstAnalysis sample window
{-# INLINABLE sbtBurstAnalysisMany #-}
{-# SPECIALIZE sbtBurstAnalysisMany :: [Int] -> Window Int -> (Window Int, Maybe Range) #-}