-- |
-- Copyright  : (c) 2012 Daniël de Kok
-- License    : Apache 2
--
-- Maintainer : Daniël de Kok <me@danieldk.eu>
-- Stability  : experimental
--
-- Approximate randomization test (Noreen, 1989)

{-# LANGUAGE DoAndIfThenElse #-}

module Main where

import           Control.Monad (liftM, when)
import           Control.Monad.Mersenne.Random (evalRandom)
import           Data.List (intersperse)
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Unboxed as V
import           Data.Word (Word64)
import           Statistics.Test.ApproxRand
import           Statistics.Test.Types (TestType(..))
import           Statistics.Types (Sample)
import           System.Console.GetOpt
import           System.Environment (getArgs)
import           System.Exit (exitFailure)
import           System.Random.Mersenne.Pure64 (PureMT, newPureMT, pureMT)
import           Text.Printf (printf)

import           ChartHistogram
import           SampleIO
import           TextHistogram

main :: IO ()
main = do
  -- Read command-line options and arguments.
  (opts, args) <- getOptions

  -- Read score files
  let col = pred $ optColumn opts
  v1 <- liftM V.fromList $ readFileCol (head args) col
  v2 <- liftM V.fromList $ readFileCol (args !! 1) col

  let stat = optTestStatistic opts

  prng <- case optPRNGSeed opts of
    Just seed -> return $ pureMT seed
    Nothing   -> newPureMT

  if optPrintStats opts then
    printStats opts stat prng v1 v2
  else
    applyTest opts stat prng v1 v2

applyTest :: Options -> TestStatistic -> PureMT -> Sample ->
  Sample -> IO ()
applyTest opts stat prng v1 v2 = do
  putStrLn $ printf "Iterations: %d" $ optIterations opts
  putStrLn $ printf "Sample sizes: %d %d" (V.length v1) (V.length v2)

  let testOptions = TestOptions (optTestType opts)
                                stat
                                (optIterations opts)
                                (optSigP opts)

  let testType = optTestType opts

  let pTest = optSigP opts
  let pTail = case testType of
                OneTailed -> pTest
                TwoTailed -> pTest / 2

  -- Test information
  putStrLn $ "Test type: " ++ show testType
  putStrLn $ printf "Test significance: %f" pTest
  putStrLn $ printf "Tail significance: %f" pTail

  -- Approximate randomization testing.
  let test = approxRandTest testOptions v1 v2
  let result = evalRandom test prng

  -- Print test statistic for the samples.
  putStrLn $ printf "Test statistic: %f" $ trStat result

  case trSignificance result of
    Significant    p -> putStrLn $ printf "Significant: %f" p
    NotSignificant p -> putStrLn $ printf "Not significant: %f" p

  when (optPrintHistogram opts) $ do
    putStrLn ""
    printHistogram 21 result
  case (optWriteHistogram opts) of
    Just fn ->
      writeHistogram testOptions 31 result fn
    Nothing ->
      return ()

printStats :: Options -> TestStatistic -> PureMT -> Sample ->
  Sample -> IO ()
printStats opts stat prng v1 v2 =
  VG.mapM_ (putStrLn . printf "%f") $
    evalRandom (approxRandStats stat (optIterations opts) v1 v2) prng

data Options = Options {
  optColumn         :: Int,
  optPrintHistogram :: Bool,
  optIterations     :: Int,
  optPRNGSeed       :: Maybe Word64,
  optPrintStats     :: Bool,
  optSigP           :: Double,
  optTestStatistic  :: TestStatistic,
  optTestType       :: TestType,
  optWriteHistogram :: Maybe String
}

defaultOptions :: Options
defaultOptions = Options {
  optColumn         = 1,
  optIterations     = 10000,
  optPRNGSeed       = Nothing,
  optPrintHistogram = False,
  optPrintStats     = False,
  optSigP           = 0.01,
  optTestStatistic  = meanDifference,
  optTestType       = TwoTailed,
  optWriteHistogram = Nothing
}

options :: [OptDescr (Options -> Options)]
options =
  chartHistogramOption backendFormats : mandatoryOptions

mandatoryOptions :: [OptDescr (Options -> Options)]
mandatoryOptions =
  [ Option ['c'] ["column"]
      (ReqArg (\arg opt -> opt { optColumn = read arg }) "NUMBER")
      "column number (starting at 1)",
    Option ['h']    ["print-histogram"]
      (NoArg (\opt -> opt { optPrintHistogram = True }))
      "print a histogram of randomized sample statistics",
    Option ['i'] ["iterations"]
      (ReqArg (\arg opt -> opt { optIterations = read arg }) "NUMBER")
      "number of iterations",
    Option ['o'] ["one-tailed"]
      (NoArg (\opt -> opt { optTestType = OneTailed }))
      "perform a one-tailed test",
    Option ['p'] []
      (ReqArg (\arg opt -> opt {optSigP = read arg }) "NUMBER")
      "significant p-value",
    Option []    ["print-scores"]
      (NoArg (\opt -> opt { optPrintStats = True }))
      "output statistics of permuted vectors",
    Option ['s'] ["seed"]
      (ReqArg (\arg opt -> opt { optPRNGSeed = Just $ read arg}) "NUMBER")
      "pseudorandom number generator seed",
    Option ['t'] ["test-statistic"]
      (ReqArg (\arg opt -> opt { optTestStatistic = parseStatistic arg}) "NAME")
      "test statistic (mean_diff, var_ratio)"
  ]

chartHistogramOption :: [String] -> OptDescr (Options -> Options)
chartHistogramOption formats =
  Option ['w'] ["write-histogram"]
    (ReqArg (\arg opt -> opt { optWriteHistogram = Just arg}) "FILENAME")
    $ printf "write a histogram (supported file extensions: %s)" formatStr
  where
    formatStr = concat $ intersperse ", " formats

getOptions :: IO (Options, [String])
getOptions = do
  args <- getArgs
  case getOpt Permute options args of
    (actions, nonOpts, [])   -> do
      when (length nonOpts /= 2) usageExit
      let opts = foldl (flip ($)) defaultOptions actions
      return (opts, nonOpts)
    (_,    _,          _)    ->
      usageExit
  where
  usageExit = do
    putStrLn $ usageInfo header options
    exitFailure
    where
      header = "Usage: approx-rand-test [OPTION...] scores scores2"

parseStatistic :: String -> TestStatistic
parseStatistic "mean_diff" = meanDifference
parseStatistic "var_ratio" = varianceRatio
parseStatistic _           = error "Unknown test statistic"


{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE BangPatterns         #-}

module TestChapter2Bandit (
       testChapter2
      )where

import           Control.Monad
import           Control.Monad.Trans.State (runStateT)
import           Control.Concurrent

import           Data.Configurator
import           Data.Configurator.Types

import           Data.Random
import           Data.Random.Distribution
import           Data.Random.Distribution.Uniform
import           Data.Random.Distribution.Bernoulli
import           Data.Random.Source.IO

import           Data.Text (Text)
import           Data.Time (getCurrentTime)
import           Graphics.Matplotlib

import           Numeric.LinearAlgebra (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA

import           System.Console.AsciiProgress(Options(..), Stats(..),
                                              displayConsoleRegions, complete,
                                              getProgressStats, def, newProgressBar, tickN)
-- project
import Utils
import Chapter2Bandit

------------------------------------------------------------------------------------------
testChapter2 :: FilePath -> IO ()
testChapter2 configPath = do
  print "Bandit Experiment Starting, will take several minutes "

  (config, _) <- autoReload autoConfig [Required configPath]
  (karm :: Int) <- require config "kArms"
  --
  (bGreedy :: Bool) <- require config "enable.bGreedyExperiment"
  (bUCB :: Bool) <- require config "enable.bUCBExperiment"
  (bGradient :: Bool) <- require config "enable.bGradientExperiment"
  (bCompareTogether :: Bool) <- require config "enable.bCompareTogether"
  (bFigure2_6 :: Bool) <- require config "enable.bFigure2_6"

  -- do the experiments in order
  -- EpsilonGreedy Experiment
  (greedyRewardsCurves, greedyActionCurves) <- if bGreedy then doEpsilonGreedyTests config
                                                          else pure (mp, mp)
  -- UCB Experiment
  (ucbRewardsCurves, ucbActionCurves) <- if bUCB then doUCBTests config else pure (mp, mp)

  -- Gradient Experiment
  (gradientRewardsCurves, gradientActionCurves) <- if bGradient then doGradientTests config
                                                      else pure (mp, mp)

  when bCompareTogether (drawExperimentsCurves config
                              (greedyRewardsCurves % ucbRewardsCurves % gradientRewardsCurves,
                               greedyActionCurves % ucbActionCurves % gradientActionCurves))

  when bFigure2_6 (drawFigure2_6 config)
  pure ()

------------------------------------------------------------------------------------------
-- | doEpsilonGreedyTest Method: compare different epsilons, different initial values,
--   also non-stationary (const step size, weight recent reward more) and stationary environment.
doEpsilonGreedyTests :: Config -> IO (Matplotlib, Matplotlib)
doEpsilonGreedyTests config = do
  startTime <- getCurrentTime
  print (show startTime ++ ": Start Epsilon Greedy Experiment: ")
  (karm :: Int) <- require config "kArms"
  (totalStep :: Int) <- require config "totalStep"
  (totalBandits :: Int) <- require config "totalBandits"
  (ges :: [Double]) <- require config "greedy.epsilons"
  (initValues :: [Double]) <- require config "greedy.initialOptimalValues"
  (stepValues :: [Double]) <- require config "greedy.stepSizes"
  -- run experiments
  let arrParams = zip3 ges initValues stepValues
  !results <- mapM (\ (ge, initValue, stepValue) -> do
                      let !greedyRVar = bernoulli ge
                      !ret <- replicateM totalBandits
                                (oneBanditGo karm totalStep initValue stepValue (EGreedy greedyRVar))
                      pure $ (sum ret) / (fromIntegral totalBandits)
                   )
                  arrParams
  -- draw it
  (ars, oas) <- drawEpsilonGreedy config results
  finishTime <- getCurrentTime
  print $ (show finishTime) ++ ": Finished Epsilon Greedy Experiments"
  pure (ars, oas)

------------------------------------------------------------------------------------------
-- | doUCBTest Method
doUCBTests :: Config -> IO (Matplotlib, Matplotlib)
doUCBTests config = do
  startTime <- getCurrentTime
  print (show startTime ++ ": Start UCB Experiment: ")
  -- read all sample average related params
  (karm :: Int) <- require config "kArms"
  (totalStep :: Int) <- require config "totalStep"
  (totalBandits :: Int) <- require config "totalBandits"
  (sss :: [Double]) <- require config "ucb.stepSizes"
  (eds :: [Double]) <- require config "ucb.exploreDegree"
  -- run experiments
  let arrParams = zip sss eds
  !results <- mapM (\ (stepSize, ed) -> do
                     let policy = UCB ed
                     !ret <- replicateM totalBandits (oneBanditGo karm totalStep 0.0 stepSize policy)
                     pure $ (sum ret) / (fromIntegral totalBandits)
                  ) arrParams
  -- draw it
  (ars, oas) <- drawUCB config results
  finishTime <- getCurrentTime
  print $ (show finishTime) ++ ": Finished UCB Experiments"
  pure (ars, oas)

------------------------------------------------------------------------------------------
-- | doUCBTest Method
doGradientTests :: Config -> IO (Matplotlib, Matplotlib)
doGradientTests config = do
  startTime <- getCurrentTime
  print (show startTime ++ ": Start Gradient Experiment: ")
  -- read all sample average related params
  (karm :: Int) <- require config "kArms"
  (totalStep :: Int) <- require config "totalStep"
  (totalBandits :: Int) <- require config "totalBandits"
  (sss :: [Double]) <- require config "gradient.stepSizes"
  (bBaselines :: [Bool]) <- require config "gradient.baselines"
  -- run experiments
  let arrParams = zip sss bBaselines
  !results <- mapM (\ (stepSize, bBase) -> do
                      let policy = Gradient bBase
                      !ret <- replicateM totalBandits (oneBanditGo karm totalStep 0.0 stepSize policy)
                      pure $ (sum ret) / (fromIntegral totalBandits)
                   ) arrParams
  (ars, oas) <- drawGradient config results
  finishTime <- getCurrentTime
  print $ (show finishTime) ++ ": Finished Gradient Experiments"
  pure (ars, oas)

------------------------------------------------------------------------------------------
oneBanditGo :: Int -> Int -> Double -> Double -> Policy -> IO (Vector Double)
oneBanditGo karm totalStep initValue stepSize policy = do
  trueValues <- generateRandomList karm stdNormal
  let initBandit = mkBandit karm totalStep initValue stepSize trueValues policy
  (averageRewards, bandit) <- runStateT (loopSteps totalStep) initBandit
  -- join two result lists together for convinent
  pure (LA.fromList $ averageRewards ++ (_bestTakes bandit))

------------------------------------------------------------------------------------------
-- drawings
drawFigure2_1 :: Int -> IO ()
drawFigure2_1 karm = do
  kArmRVars::[RVar Double] <- map (flip normal 1.0) <$> generateRandomList karm stdNormal
  kArmDatas <- mapM (generateRandomList 200) kArmRVars
  let figure2_1 = mp % (violinplot kArmDatas @@ [o2 "showmeans" True, o2 "showmedians" True])
                     % xlabel "Action"
                     % ylabel "Reward"
                     % title "Figure 2-1"
                     % grid True
  code figure2_1 >> onscreen figure2_1
  pure ()

drawEpsilonGreedy :: Config -> [Vector Double] -> IO (Matplotlib, Matplotlib)
drawEpsilonGreedy config results = do
  (totalStep :: Int) <- require config "totalStep"
  (ges :: [Double]) <- require config "greedy.epsilons"
  (initValues :: [Double]) <- require config "greedy.initialOptimalValues"

  let rewardCurves = foldl (goPlot 0 totalStep) mp (zip3 ges initValues results)
      bestActions = foldl (goPlot totalStep totalStep) mp (zip3 ges initValues results)
      --subplots @@ [o2 "nrows" 2, o2 "sharey" True]
      -- % setSubplot 0 -- % setSubplot 1
      !figureReward = rewardCurves
                       % title "Epsilon-Greedy Average Reward Comparison"
                       % xlabel "Step"
                       % ylabel "Average Reward"
                       % yticks [0.0, 0.2 .. 1.6]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
      !figureAction = bestActions
                       % title "Epsilon-Greedy Optimal Action Comparison"
                       % xlabel "Step"
                       % ylabel "Optiomal Actions"
                       % yticks [0.0::Double, 0.2 .. 1.0]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
                       % tightLayout
  --
  code figureReward >> code figureAction -- avoid Matplotlib's bug
  onscreen figureReward >> onscreen figureAction
  threadDelay 500000
  pure (rewardCurves, bestActions)
  where
  goPlot :: Int -> Int -> Matplotlib -> (Double, Double, Vector Double) -> Matplotlib
  goPlot startIdx totalLen acc (greedy, initValue, theData) =
    acc % plot [1..totalLen] (LA.subVector startIdx totalLen theData)
               @@ [o2 "label" ("epsilon="++(show greedy)++", Q="++(show initValue))]

drawUCB :: Config -> [Vector Double] -> IO (Matplotlib, Matplotlib)
drawUCB config results = do
  (totalStep :: Int) <- require config "totalStep"
  (eds :: [Double]) <- require config "ucb.exploreDegree"
  let rewardCurves = foldl (goPlot 0 totalStep) mp (zip eds results)
      bestActions = foldl (goPlot totalStep totalStep) mp (zip eds results)
      !figureReward = rewardCurves
                       % title "UCB Average Reward Comparison"
                       % xlabel "Step"
                       % ylabel "Average Reward"
                       % yticks [0.0, 0.2 .. 1.6]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
      !figureAction = bestActions
                       % title "UCB Optimal Action Comparison"
                       % xlabel "Step"
                       % ylabel "Optiomal Actions"
                       % yticks [0.0::Double, 0.2 .. 1.0]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
                       % tightLayout
  --
  code figureReward >> code figureAction -- avoid Matplotlib's bug
  onscreen figureReward >> onscreen figureAction
  threadDelay 500000
  pure (rewardCurves, bestActions)
  where
  goPlot :: Int -> Int -> Matplotlib -> (Double, Vector Double) -> Matplotlib
  goPlot startIdx totalLen acc (exploreDegree, theData) =
    acc % plot [1..totalLen] (LA.subVector startIdx totalLen theData)
               @@ [o2 "label" ("ucb, c=" ++ (show exploreDegree))]

drawGradient :: Config -> [Vector Double] -> IO (Matplotlib, Matplotlib)
drawGradient config results = do
  (totalStep :: Int) <- require config "totalStep"
  (bBaselines :: [Bool]) <- require config "gradient.baselines"
  (sss :: [Double]) <- require config "gradient.stepSizes"
  let rewardCurves = foldl (goPlot 0 totalStep) mp (zip3 sss bBaselines results)
      bestActions = foldl (goPlot totalStep totalStep) mp (zip3 sss bBaselines results)
      !figureReward = rewardCurves
                       % title "Gradient Average Reward Comparison"
                       % xlabel "Step"
                       % ylabel "Average Reward"
                       % yticks [0.0, 0.2 .. 1.6]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
      !figureAction = bestActions
                       % title "Gradient Optimal Action Comparison"
                       % xlabel "Step"
                       % ylabel "Optiomal Actions"
                       % yticks [0.0::Double, 0.2 .. 1.0]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
                       % tightLayout
  --
  code figureReward >> code figureAction -- avoid Matplotlib's bug
  onscreen figureReward >> onscreen figureAction
  threadDelay 500000
  pure (rewardCurves, bestActions)
  where
  goPlot :: Int -> Int -> Matplotlib -> (Double, Bool, Vector Double) -> Matplotlib
  goPlot startIdx totalLen acc (alpha, bBaseline, theData) =
    acc % plot [1..totalLen] (LA.subVector startIdx totalLen theData)
               @@ [o2 "label" ("alpha=" ++ (show alpha) ++ (bBaseline?(",with", ",without")))]

drawExperimentsCurves :: Config -> (Matplotlib, Matplotlib) -> IO ()
drawExperimentsCurves config (rewards, actions) = do
  let !figureReward = rewards
                       % title "Average Reward Comparison"
                       % xlabel "Step"
                       % ylabel "Average Reward"
                       % yticks [0.0, 0.2 .. 1.6]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
      !figureAction = actions
                       % title "Optimal Action Comparison"
                       % xlabel "Step"
                       % ylabel "Optiomal Actions"
                       % yticks [0.0::Double, 0.2 .. 1.0]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
                       % tightLayout
  --
  code figureReward >> code figureAction -- avoid Matplotlib's bug
  onscreen figureReward >> onscreen figureAction
  threadDelay 1000000

------------------------------------------------------------------------------------------
------------------------------------------------------------------------------------------
greedyEpsilons = [(-7) .. (-1)]
-- greedy with different initial values and a constant stepSize 0.1
greedyInitValues = [(-2) .. 3]
gradientAlpha  = [(-5) .. 2]
ucbExplore = [(-4) .. 3]

drawFigure2_6 :: Config -> IO ()
drawFigure2_6 config = do
  startTime <- getCurrentTime
  print (show startTime ++ ": Start Figure2_6 Experiment: ")
  -- read all sample average related params
  (karm :: Int) <- require config "kArms"
  (totalStep :: Int) <- require config "totalStep"
  (totalBandits :: Int) <- require config "totalBandits"

  -- 1. different epsilons
  let eGreedyParams = zip3 (map (2.0 ^^) greedyEpsilons) (repeat 0.0) (repeat (-1.0))
  !eGreedyResults <-
    mapM (\ (ge, initValue, stepValue) -> do
            let !greedyRVar = bernoulli ge
            !ret <- replicateM totalBandits
                      (oneBanditGo karm totalStep initValue stepValue (EGreedy greedyRVar))
            pure $ (sum ret) / (fromIntegral totalBandits)
         )  eGreedyParams
  let !eGreedyAverage = calcAverage 0 totalStep eGreedyResults
      !eGreedyAction = calcAverage totalStep totalStep eGreedyResults
      !eGreedyAverageCurve = mp % plot greedyEpsilons eGreedyAverage
                                 @@ [o2 "label" "epsilon-greedy"]
      !eGreedyActionCurve = mp % plot greedyEpsilons eGreedyAction
                                @@ [o2 "label" "epsilon-greedy"]

  -- 2. different e-greedy init values with epsilon = 0.0 and stepSize = 0.1
  let eGreedyInitValParams = zip3 (repeat 0.0) (map (2.0 ^^) greedyInitValues) (repeat 0.1)
  !eGreedyInitValResults <-
    mapM (\ (ge, initValue, stepValue) -> do
            let !greedyRVar = bernoulli ge
            !ret <- replicateM totalBandits
                      (oneBanditGo karm totalStep initValue stepValue (EGreedy greedyRVar))
            pure $ (sum ret) / (fromIntegral totalBandits)
         )  eGreedyInitValParams
  let !eGreedyInitValAverage = calcAverage 0 totalStep eGreedyInitValResults
      !eGreedyInitValActions = calcAverage totalStep totalStep eGreedyInitValResults
      eGreedyInitValAverageCurve = mp % plot greedyInitValues eGreedyInitValAverage
                                             @@ [o2 "label" "greedy-OpticalValue"]
      eGreedyInitValActionCurve = mp % plot greedyInitValues eGreedyInitValActions
                                            @@ [o2 "label" "greedy-OpticalValue"]

  -- 3. ucb with stepSize 0.1
  let ucbParams = zip (repeat 0.1) (map (2.0 ^^) ucbExplore)
  !ucbResults <- mapM (\ (stepSize, ed) -> do
                         let policy = UCB ed
                         !ret <- replicateM totalBandits
                                   (oneBanditGo karm totalStep 0.0 stepSize policy)
                         pure $ (sum ret) / (fromIntegral totalBandits)
                      ) ucbParams
  let !ucbAverage = calcAverage 0 totalStep ucbResults
      !ucbAction = calcAverage totalStep totalStep ucbResults
      ucbAverageCurve = mp % plot ucbExplore ucbAverage @@ [o2 "label" "UCB"]
      ucbActionCurve = mp % plot ucbExplore ucbAction @@ [o2 "label" "UCB"]

  -- 4. gradient method with different alphas and all with baseline
  let gradientParams = zip (map (2.0 ^^) gradientAlpha) (repeat True)
  !gradientResults <- mapM (\ (stepSize, bBase) -> do
                              let policy = Gradient bBase
                              !ret <- replicateM totalBandits
                                        (oneBanditGo karm totalStep 0.0 stepSize policy)
                              pure $ (sum ret) / (fromIntegral totalBandits)
                           ) gradientParams
  let !gradientAverage = calcAverage 0 totalStep gradientResults
      !gradientAction = calcAverage totalStep totalStep gradientResults
      gradientAverageCurve = mp % plot gradientAlpha gradientAverage @@ [o2 "label" "Gradient"]
      gradientActionCurve = mp % plot gradientAlpha gradientAction @@ [o2 "label" "Gradient"]
  -- 5. now draw it
  let !figureReward = eGreedyAverageCurve
                       % gradientAverageCurve
                       % ucbAverageCurve
                       % eGreedyInitValAverageCurve
                       % title "A Parameter study of the various bandit algorithms"
                       % xlabel "Paramter (2^x): e / Alpah / c / Q"
                       % ylabel "Average Reward"
                       % yticks [0.9, 0.1 .. 1.6]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
      !figureAction = eGreedyActionCurve
                       % gradientActionCurve
                       % ucbActionCurve
                       % eGreedyInitValActionCurve
                       % title "A Parameter study of the various bandit algorithms"
                       % xlabel "Paramter (2^x): e / Alpah / c / Q"
                       % ylabel "Optiomal Actions"
                       % yticks [0.2::Double, 0.1 .. 1.0]
                       % legend @@ [o2 "fancybox" True,
                                    o2 "shadow" True, o2 "loc" "lower right"]
                       % grid True
                       % tightLayout
  code figureReward >> code figureAction
  onscreen figureReward >> onscreen figureAction
  threadDelay 500000
  pure ()
  where
  calcAverage startIdx len =
    map ((/ (fromIntegral len)) . sum . LA.toList . LA.subVector startIdx len)


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

module Main where

import           Control.Monad.Compat
import           Control.Monad.Trans.Except

import           Criterion.Analysis
import           Criterion.IO.Printf
import           Criterion.Main
import           Criterion.Measurement           (initializeTime, measured)
import           Criterion.Monad
import           Criterion.Report
import           Criterion.Types

import           Data.Aeson
import qualified Data.ByteString.Lazy            as BS
import qualified Data.Vector                     as V

import           Prelude.Compat

import           Statistics.Resampling.Bootstrap
import           Statistics.Types

import           System.Directory
import           System.Environment

data SingleTest = SingleTest String [(Double,Int)]
        deriving Show

instance FromJSON SingleTest where
 parseJSON (Object v) =
    SingleTest
           <$> v .: "name"
           <*> fmap f (v .: "results")
      where f res = [ (t,i) | (Just t,i) <- res `zip` [1..] ]
 parseJSON _ = mzero

main :: IO ()
main = do
   initializeTime -- you need this
   [inpFile,outFile] <- getArgs
   inp <- BS.readFile inpFile
   let Just rawTestResults = decode inp
   let config = defaultConfig { verbosity = Verbose, reportFile = return "tmp.html", csvFile = return outFile }
   fileExists <- doesFileExist  outFile
   when fileExists $ removeFile outFile
   withConfig config $ do
           writeCsv ["Name","Mean","MeanLB","MeanUB","Stddev","StddevLB", "StddevUB"::String]
           rs <- sequence
                [ analyseOne n nm $ V.fromList $ [ measured { measTime = t, measIters = fromIntegral i } | (t,i) <- res ]
                | (n,SingleTest nm res) <- [0..] `zip` rawTestResults
                ]
           report rs
           return ()

analyseOne :: Int -> String -> V.Vector Measured -> Criterion Report
analyseOne i desc meas = do
  erp <- runExceptT $ analyseSample i desc meas
  case erp of
    Left err -> printError "*** Error: %s\n" err
    Right rpt@Report{..} -> do
      let SampleAnalysis{..} = reportAnalysis
          OutlierVariance{..} = anOutlierVar
      writeCsv (desc,
              estPoint anMean, estPoint anMean - confIntLDX (estError anMean) , estPoint anMean + confIntUDX (estError anMean),
              estPoint anStdDev, estPoint anStdDev - confIntLDX (estError anStdDev),
              estPoint anStdDev + confIntUDX (estError anStdDev))
      return rpt


{-# LANGUAGE TemplateHaskell #-}
module Turtle3 (main) where
import Prelude hiding (dropWhile)
import Control.Applicative
import Control.Arrow
import System.IO (hFlush, stdout)
import Data.Complex
import Ros.Logging
import Ros.Node
import Ros.Topic (repeatM, force, dropWhile, metamorphM, yieldM)
import Ros.Topic.Util (everyNew, interruptible, forkTopic, topicOn, subsample)
import Ros.Util.PID (pidTimedIO)
import AngleNum
import Ros.Turtlesim.Pose
import Ros.Geometry_msgs.Twist
import qualified Ros.Geometry_msgs.Vector3 as V
import Data.Default.Generics (def)
import Lens.Family ((.~), (^.), (&))

-- A type synonym for a 2D point.
type Point = Complex Float

-- A Topic of user-supplied waypoint trajectories.
getTraj :: Topic IO [Point]
getTraj = repeatM (do putStr "Enter waypoints: " >> hFlush stdout
                      $(logInfo "Waiting for new traj")
                      map (uncurry (:+)) . read <$> getLine)

-- Create Twist message with linear and angular velocity in arguments.
mkTwist :: Real a => a -> a -> Twist
mkTwist l a = def & linear  . V.x .~ realToFrac l
                  & angular . V.z .~ realToFrac a

-- Produce a new goal 'Point' every time a 'Pose' topic reaches the
-- vicinity of the previous goal.
destinations :: (Functor m, Monad m) =>
                Topic m Point -> Topic m Pose -> Topic m Point
destinations goals poses = metamorphM (start (p2v <$> poses)) goals
  where start t g = yieldM g (go g t)
        go g t g' = let keepGoing x = magnitude (g - x) > 0.5
                    in force (dropWhile keepGoing t) >>= yieldM g' . go g'
        p2v (Pose x y _ _ _) = x :+ y

-- Compute linear and angular error to goal.
toGoal :: (Pose,Point) -> (Float, Float)
toGoal (pos,goal) = (magnitude v, thetaErr)
  where v = goal - (pos ^. x :+ pos ^. y)
        thetaErr = toDegrees $ angle (phase v) - angle (pos ^. theta)

-- Run a PID loop on angular velocity to converge to a bearing for the
-- goal.
steering :: Topic IO (Float,Float) -> Topic IO Twist
steering = topicOn snd (mkTwist . fst) (($ 0) <$> pidTimedIO 0.01 0.5 0)

-- Compute position and bearing error to goal, clamp linear and
-- angular velocities, and generate a steering command.
navigate :: Topic IO (Pose,Point) -> Topic IO Twist
navigate = steering . fmap ((clamp *** id) . toGoal)
  where clamp x = signum x * min 2 (abs x)

main = runNode "HaskellBTurtle" $
       do enableLogging (Just Warn)
          poses <- subscribe "/turtle1/pose"
          let goals = destinations (interruptible getTraj) poses
              commands = subsample 15 $ navigate $ everyNew poses goals
          advertise "/turtle1/cmd_vel" commands


module View
  ( view
  ) where

import           Data.Maybe
import           Graphics.Rendering.OpenGL    (($=))
import qualified Graphics.Rendering.OpenGL    as GL
import qualified Graphics.UI.GLFW             as GLFW
import           Lens.Micro.Platform          ((^.))

--import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Static

import           Drawable
import           Matrix
import           MyState
import           Rotation
import           Shaders
import           Uniform

getColor Red   = GL.Color3 1.0 0.0 (0.0 :: GL.GLfloat)
getColor Green = GL.Color3 0.0 1.0 (0.0 :: GL.GLfloat)
getColor Blue  = GL.Color3 0.0 0.0 (1.0 :: GL.GLfloat)

--View
setColor program color = setUniform program "color" $ getColor color

viewMatrix :: Camera -> Matrix4
viewMatrix camera = rot <> translation
  where
    translation = translate . (* (-1)) $ camera ^. cameraPosition
    rot = rotationMatrix $ mkXYRotationCamera (-x) (-y)
    (x, y) = camera ^. cameraRotation

model :: Vector3 -> Matrix4
model = translate

lightPosition = GL.Vertex3 (-0.0) 100.0 0.0 :: GL.Vertex3 GL.GLfloat

reset =
  GL.clearColor $= GL.Color4 0 0 0 1 >>
  GL.clear [GL.ColorBuffer, GL.DepthBuffer]

drawCube mU pU state pos = do
  let m = model pos
  seq m $ return ()
  setUniform4fv m mU
  draw . head $ state ^. cubes

grid space =
  map
    (* realToFrac space)
    (vec3 <$> [-5 .. 5] <*> [-5 .. 5] <*> [-5 .. 5])

drawCubes shaders state = do
  program <- activateProgram shaders Normal
  setColor program $ state ^. color
  setUniform program "useVertexColor" (1 :: GL.GLuint)
  setUniform program "LightPosition_worldspace" lightPosition
  let m = model $ state ^. cubePos
      v = viewMatrix $ state ^. camera
      p = projection $ state ^. aspectRatio
      vp = p <> v <> m
  mU <- getUniform program "M"
  vU <- getUniform program "V"
  pU <- getUniform program "P"
  setUniform4fv v vU
  setUniform4fv p pU
  mapM_ (drawCube mU pU state) (grid $ state ^. spacing)

drawTerrain state program = do
  setColor program $ state ^. color
  setUniform program "useVertexColor" (1 :: GL.GLuint)
  setUniform program "LightPosition_worldspace" lightPosition
  let m = model $ state ^. cubePos
      v = viewMatrix $ state ^. camera
      p = projection $ state ^. aspectRatio
      vp = p <> v <> m
  getUniform program "V" >>= setUniform4fv v
  getUniform program "P" >>= setUniform4fv p
  getUniform program "M" >>= setUniform4fv m
  draw . head $ state ^. terrain

view :: Shaders -> MyState -> IO ()
view shaders state = do
  startTime <- GLFW.getTime
  reset
  drawCubes shaders state
  activateProgram shaders Normal >>= drawTerrain state
  activateProgram shaders Wireframe >>= drawTerrain state
  endTime <- GLFW.getTime
  let fps = ((1 /) . fromJust) ((-) <$> endTime <*> startTime)
  putStrLn $ "FPS: " ++ show fps
  putStrLn $
    "Triangles per second: " ++ show (fromIntegral (length (grid 3)) * 12 * fps)


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}
module Workflow
    ( WorkflowStage, prev, quantiles, waves, task, meanT, varT
    , Task(..)
    , Indices, Nodes
    , WFType(Sequential, Parallel, Complex)
    , wrap, calculateCDF, calculateCDF', distributionBounds
    , stageOrder,applyDE,applyRS,reverseEdges
    ) where
import Data.Array.Unboxed (listArray)
import Data.Complex (Complex)
import qualified Data.Map as Map
import qualified Data.List as List
import Data.Aeson
import Data.Aeson.Types
import qualified Data.Vector as Vector
import Data.Maybe (fromMaybe)

import Workflow.Toolset

-- | List of indices of the stages in Workflow Nodes Map
type Indices = [Int]

-- | Map structure, stores all stages of the workflow
type Nodes = Map.Map Int WorkflowStage

-- | A stage of the workflow - can be either an atomic task or a workflow
data WorkflowStage = Stage
    { prev      :: Indices          -- ^ previous stages
    , quantiles :: [Double]         -- ^ values of CDF at certain points
    , waves     :: [Complex Double] -- ^ values of characteristic functions at certain points
    , task      :: Task             -- ^ task (or sub-workflow) for execution at this stage
    , meanT     :: Double           -- ^ mean execution time estimation
    , varT      :: Double           -- ^ estimation of execution time variance
    } deriving (Eq)

-- | Representation of the task - it can be a single task or a package
data Task = AtomicTask
    { sample :: Array1D Double -- ^ adjusted data (T' = mu' + sigma' (T(X) - mu(X)) / sigma(X) )
    } | Workflow
    { nodes  :: Nodes   -- ^ list of stages and their predecessors
    , leafs  :: Indices -- ^ last tasks in the workflow
                        --   (workflow is fully executed if these tasks are finished)
    , wftype :: WFType  -- ^ whether this is primitive workflow or complex one
    } deriving (Eq)

-- | Type of the workflow.
--   Depending on the type of the workflow, distribution is calculated in different ways
data WFType = Sequential -- ^ can apply RS reduction
            | Parallel   -- ^ can apply RF reduction
            | Complex    -- ^ cannot reduce workflow
            deriving (Eq,Show)

-- | Wrap the task into workflow stage
--   Arguments: current task, and an array of previous (required) stages
wrap :: Task -> Indices -> WorkflowStage
wrap t@(AtomicTask {sample=arr}) preds = Stage
    { prev      = preds
    , task      = t
    , meanT     = mean arr
    , varT      = var arr
    , quantiles = undefined
    , waves     = undefined
    }
wrap t@(Workflow {nodes=ns,leafs=xs}) preds = Stage
    { prev      = preds
    , task      = t
    , meanT     = foldr1 max (xs >>= f meanT)
    , varT      = foldr1 max (xs >>= f varT)
    , quantiles = undefined
    , waves     = undefined
    }
  where
    f r i = g r ((Map.!) ns i)
    g r x@(Stage {prev=[]}) = [r x]
    g r x = map (+ r x) . concat . map (f r) . prev $ x

-- |
-- Calculate probability distribution of the workflow execution time
-- Arguments: workflow, bounds size (in sigmas), grid density k : n = 2^k
-- Return: ((a,b),wf); (a,b) - integration bounds, wf - resulting workflow
calculateCDF :: WorkflowStage -> Double -> Int -> ((Double,Double),WorkflowStage)
calculateCDF s q k = (bnds, ccdf (map (\x' -> x'-m*c) x, t) s) where
    grid@(x,t) = linSpace k bnds
    bnds@(a,b) = distributionBounds s q
    c = (b-a) / (fromIntegral . length . fst $ grid)
    m = fromIntegral . (round :: Double -> Int) $ meanT s / c

-- |
-- Calculate probability distribution of the workflow execution time
-- Arguments: workflow, grid
-- Return: ((a,b),wf); (a,b) - integration bounds, wf - resulting workflow
calculateCDF' :: WorkflowStage -> ([Double],[Double]) -> ((Double,Double),WorkflowStage)
calculateCDF' s grd = (bnds, ccdf (map (\x' -> x'-m*c) x, t) s) where
    grid@(x,t) = grd
    bnds@(a,b) = (head x, last x)
    c = (b-a) / (fromIntegral . length . fst $ grid)
    m = fromIntegral . (round :: Double -> Int) $ meanT s / c

-- | calculate cdf with grid as input
ccdf :: ([Double],[Double]) -> WorkflowStage -> WorkflowStage
ccdf g@(xGrid,tGrid) s = s{
    quantiles = qu,
    waves = wa,
    task  = newt
    } where
        n = length xGrid
        n2 = length tGrid
        c = (last xGrid - head xGrid) / fromIntegral (n-1)
        m = fromIntegral . (round :: Double -> Int) $ meanT s / c
        xg = map (+m*c) xGrid
        qu = df newt
        wa = cf newt
        newt = parseInner . task $ s
        parseInner t@(AtomicTask {}) = t
        parseInner t@(Workflow {nodes=ns}) = t{nodes = Map.map (ccdf g) ns}
        df (AtomicTask {sample=arr}) = distributionF arr xg
        df (Workflow {nodes=_ns,wftype=Sequential})
            = charToDist (map (+m*c) xGrid, tGrid) wa
        df (Workflow {nodes=ns,wftype=Parallel})
            = Map.fold (zipWith (*) . quantiles) (replicate n 1) ns
        df (Workflow {nodes=ns,leafs=lfs,wftype=Complex})
            = foldr (\(phi1,eX1) (f2) ->
                let m1 = min m $ fromIntegral . (round :: Double -> Int) $ eX1 / c
                    f1 = charToDist (map (+m1*c) xGrid, tGrid) phi1
                    md = round $ m - m1
                    f1' = drop md f1 ++ replicate md 1
                in zipWith (*) f2 f1'
                ) (replicate n 1) (lfs >>= f) where
                f i = gx ((Map.!) ns i)
                gx x@(Stage {prev=[]}) = [(waves x, meanT x)]
                gx x = map ((\(a,b) (e,d) -> (zipWith (*) a e,b+d)) (waves x, meanT x))
                    . concat . map f . prev $ x
        cf (AtomicTask {sample=arr}) = characteristicF arr tGrid
        cf (Workflow {nodes=ns,wftype=Sequential})
            = Map.fold (zipWith (*) . waves) (replicate n2 1) ns
        cf (Workflow {})
            = distToChar (xg, tGrid) qu

-- | Approximate bounds, such that total distribution almost fully lies in these bounds.
--   "k" stands for k-sigma rule (i.e. interval width is 2k*sigma).
--   (a,b) : 0 <= a <= b
distributionBounds :: WorkflowStage -> Double -> (Double,Double)
distributionBounds x q = (max 0 $ m - s, m + s) where
    m = meanT x
    s = q * (sqrt . varT $ x)

-- | Apply DE-reduction (delete redundant edges)
applyDE :: Task -> Task
applyDE t@(AtomicTask {}) = t
applyDE wf@(Workflow {nodes = ns}) = wf{nodes = Map.mapWithKey (\i s -> modify ((Map.!) orders i) s ) ns} where
    orders = stageOrder wf
    modify _ s@(Stage {prev=[]}) = s
    modify _ s@(Stage {prev=[_]}) = s
    modify [_] s = s
    modify ((o,_):ords) s = s{prev = fst $ foldr f ([],[]) ords} where
        f (_,is) (os,iss) = ( mergeOrderedLists os fis
                            ,foldr mergeOrderedLists iss . map (prevTillOrder (o-1)) $ fis
                            ) where fis = filter (not . flip isInOrderedList iss) is
    modify [] s = s
    prevTillOrder o i = if so <= o then [] else foldr mergeOrderedLists xs ys where
        s = (Map.!) ns i
        xs = prev s
        ys = map (prevTillOrder o) xs
        so = fst . last $ (Map.!) orders i

-- | Apply RS-reduction (Reduce Sequential stages to single stage)
applyRS :: Task -> Task
applyRS t@(AtomicTask {}) = t
applyRS wf@(Workflow {nodes = ns}) = wf{nodes = Map.union changed unchanged} where
    re = reverseEdges wf
    candidates =
        Map.filterWithKey (\i _ -> ((1/=) . length . (Map.!) re $ i) ||
        ((1>) . length . prev . (Map.!) ns . head . (Map.!) re $ i))
        . Map.filter ((1==) . length . (Map.!) re . head . prev )
        . Map.filter ((1==) . length . prev) $ ns
    appendPath Stage {prev=[]} = []
    appendPath Stage {prev=[i]} = if (1==) . length . (Map.!) re $ i then i:(appendPath $ (Map.!) ns i) else []
    appendPath Stage {} = []
    allSeqs = Map.mapWithKey (\i s -> i:(appendPath s)) candidates
    unchanged = Map.filterWithKey (\i _ -> isInOrderedList i unchKeys) ns where
        unchKeys = Map.keys ns List.\\ (concat . Map.elems $ allSeqs)
    changed = Map.map createSubWF allSeqs
    createSubWF xs = wrap Workflow {
        nodes = nnns,
        leafs = [head xs],
        wftype = Sequential
        } preds where
            nns = Map.filterWithKey (\i _ -> any (i==) xs) ns
            nnns = Map.adjust (\s -> s{prev=[]}) (last xs) nns
            preds = prev . (Map.!) ns . last $ xs

-- | helper function: return a map with tasks' execution orders
-- required to proceed with Delete Edge (DE) reduction
-- Arguments: Workflow
-- return map (Index of node -> value), where value: (order of this stage o, indices of previous stages with order o-1)
stageOrder :: Task -> Map.Map Int [(Int,Indices)]
stageOrder AtomicTask {} = Map.empty
stageOrder Workflow {leafs = lfs, nodes = ns} = List.foldl' ind' Map.empty lfs where
    ind' m i = m `seq` ind m ((Map.!) ns i) i
    ind m (Stage {prev=[]}) i = if Map.member i m then m else Map.insert i [(1, [])] m
    ind m (Stage {prev=xs}) i = if Map.member i m then m else
                Map.insert i (List.foldr insertNode [] . map createNode $ xs) cm where
                    cm = List.foldl' ind' m xs
                    createNode j = ((1+) . fst . last . (Map.!) cm $ j,[j])
    insertNode x [] = [x]
    insertNode x@(o,i) xxs@((o1,i1):xs)
        | o == o1   = (o1, mergeOrderedLists i i1):xs
        | o <  o1   = x:xxs
        | otherwise = (o1,i1):(insertNode x xs)

-- test it: giveWorkflow >>= mapM_ print . Map.toList . stageOrder . task

-- | helper function: return a map with "next" edges created from "prev" field of (next) stages
reverseEdges :: Task -> Map.Map Int Indices
reverseEdges AtomicTask {} = Map.empty
reverseEdges Workflow{nodes = ns} =
    Map.foldWithKey f (Map.fromAscList . zip (Map.keys ns) . repeat $ []) ns where
        f i Stage{prev = xs} m = foldr (Map.adjust (i:)) m xs

-- | helper function: merge two ordered lists into one ordered list - O(n+m)
mergeOrderedLists :: (Ord a) => [a] -> [a] -> [a]
mergeOrderedLists [] xs = xs
mergeOrderedLists xs [] = xs
mergeOrderedLists (x1:xs1) (x2:xs2) | x1 == x2  = x1 :(mergeOrderedLists xs1 xs2)
                                    | x1 > x2   = x2 :(mergeOrderedLists (x1:xs1) xs2)
                                    | otherwise = x1 :(mergeOrderedLists xs1 (x2:xs2))

-- | helper function: if x is a member of an ordered list - O(n)
isInOrderedList :: (Ord a) => a -> [a] -> Bool
isInOrderedList _ [] = False
isInOrderedList x [y] = x == y
isInOrderedList x (y:ys) | x == y    = True
                         | x >  y    = isInOrderedList x ys
                         | otherwise = False

instance Show Task where
    show AtomicTask {} = "AtomicTask"
    show Workflow {wftype = tp, nodes = ns, leafs = lfs} =
        show tp ++ " Workflow: leafs = " ++ show lfs ++
        "; nodes = " ++ show (Map.toList ns) ++ "."

instance Show WorkflowStage where
    show Stage {prev = [], task = (AtomicTask {})} = "AtomicTask"
    show Stage {prev = pr, task = (AtomicTask {})} = "AtomicTask (prev = " ++ show pr ++ ")"
    show Stage {prev = [], task = wf@Workflow {}} = show wf
    show Stage {prev = pr, task = (Workflow {wftype = tp, nodes = ns})} =
        show tp ++ " Workflow (size = " ++ show (Map.size ns) ++ ", prev = " ++ show pr ++ ")"

instance FromJSON Task where
    parseJSON (Object v) = do
        mleafs <- v .:? "leafs"
        mtasks <- v .:? "tasks"
        case (,) <$> mleafs <*> mtasks of
          Nothing -> AtomicTask . (\a -> listArray (1,length a) a) <$> v .: "sample"
          Just (lfs, tasks) -> return $ if Map.size tasks == 1
                               then task . head $ Map.elems tasks
                               else Workflow tasks lfs (inferWFType tasks lfs)
      where
       -- Safely escape (hopefully) impossible workflows
       inferWFType stages _ | Map.null stages = Complex
       -- check if fully sequential workflow
       inferWFType stages [oneleaf] = goChain stages oneleaf 1
       -- check if fully parallel workflow
       inferWFType stages lfs = if (length lfs == length stages)
                                    && (all (\s -> prev s == prev st) stgs)
                                then Parallel else Complex
                              where
                                st:stgs = Map.elems stages
       -- go through sequence of nodes to check if it is a sequential wf
       goChain stages lf n |   n == Map.size stages
                            && all (flip Map.notMember stages) (prev (stages Map.! lf)) = Sequential
                           |  p <- (prev (stages Map.! lf))
                              ,  n < Map.size stages
                              && length p == 1
                              && Map.member (head p) stages = goChain stages (head p) (n+1)
                           | otherwise = Complex
    parseJSON invalid = typeMismatch "Task" invalid

data IdTask = IdTask Int WorkflowStage

instance FromJSON IdTask where
    parseJSON ov@(Object v) = IdTask
        <$> v .: "id"
        <*> parseJSON ov
    parseJSON invalid = typeMismatch "(Int, Task)" invalid

instance FromJSON Nodes where
    parseJSON (Array v) = Map.fromList . map (\(IdTask i t) -> (i, t)) . Vector.toList <$> mapM parseJSON v
    parseJSON invalid = typeMismatch "AtomicTask" invalid

instance FromJSON WorkflowStage where
    parseJSON ov@(Object v) = do
      t  <- parseJSON ov
      mdeps <- v .:? "deps"
      return $ wrap t (fromMaybe [] mdeps)
    parseJSON invalid = typeMismatch "WorkflowStage" invalid


--
-- Copyright (c) 2009-2011, ERICSSON AB
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--
--     * Redistributions of source code must retain the above copyright notice,
--       this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of the ERICSSON AB nor the names of its contributors
--       may be used to endorse or promote products derived from this software
--       without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-- FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-- DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-- CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-- OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-- OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--

module Feldspar.Core.Frontend.Floating where

import qualified Prelude
import Prelude (Float,Double)
import Data.Complex

import Feldspar.Core.Constructs
import Feldspar.Core.Constructs.Floating
import Feldspar.Core.Frontend.Fractional

-- Make new class, with "Data" in all the types

infixr 8 **

class (Fraction a, Prelude.Floating a) => Floating a where
  pi        :: Data a
  pi        =  sugarSymF Pi
  exp       :: Data a -> Data a
  exp       =  sugarSymF Exp
  sqrt      :: Data a -> Data a
  sqrt      =  sugarSymF Sqrt
  log       :: Data a -> Data a
  log       =  sugarSymF Log
  (**)      :: Data a -> Data a -> Data a
  (**)      =  sugarSymF Pow
  logBase   :: Data a -> Data a -> Data a
  logBase   =  sugarSymF LogBase
  sin       :: Data a -> Data a
  sin       =  sugarSymF Sin
  tan       :: Data a -> Data a
  tan       =  sugarSymF Tan
  cos       :: Data a -> Data a
  cos       =  sugarSymF Cos
  asin      :: Data a -> Data a
  asin      =  sugarSymF Asin
  atan      :: Data a -> Data a
  atan      =  sugarSymF Atan
  acos      :: Data a -> Data a
  acos      =  sugarSymF Acos
  sinh      :: Data a -> Data a
  sinh      =  sugarSymF Sinh
  tanh      :: Data a -> Data a
  tanh      =  sugarSymF Tanh
  cosh      :: Data a -> Data a
  cosh      =  sugarSymF Cosh
  asinh     :: Data a -> Data a
  asinh     =  sugarSymF Asinh
  atanh     :: Data a -> Data a
  atanh     =  sugarSymF Atanh
  acosh     :: Data a -> Data a
  acosh     =  sugarSymF Acosh

instance Floating Float
instance Floating Double

instance (Fraction a, Prelude.RealFloat a) => Floating (Complex a)


{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleContexts, FlexibleInstances, UndecidableInstances, OverloadedLists #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeApplications #-}

{-# LANGUAGE DeriveGeneric #-}

{-# LANGUAGE GADTs #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedLists #-}

module Test where

import Control.Applicative
import Control.Monad
import Control.Monad.Fix
import Control.Lens hiding (index, inside)
import Data.Data
import Data.Distributive
import Data.Foldable as Foldable
--import Data.Profunctor.Rep (Representable(..))
import Data.Proxy
import Data.Reflection as R
import Data.Tagged
import Data.Traversable (sequenceA)
import Data.Vector as V
import GHC.TypeLits
import GHC.Generics (Generic)
import GHC.Generics (Generic1)

import Unsafe.Coerce
{-
import Prelude hiding ((*), (+), (-), negate, subtract,zipWith)
import qualified Prelude
import qualified Data.Semiring as R
-}

import qualified Control.Lens.Internal.Context as I

import Linear.Vector
-- import Linear.Covector
-- import Linear.Metric
-- import Linear.Conjugate

import Linear.V

import Data.Foldable as Foldable

import Data.Coerce (coerce)
--import Data.Data.Lens
import Data.Functor.Rep --(Representable(..))

-- TensorLike t = (Traversable t, Additive t, Dim t, Dims t, Representable t, Foldable t, Finite t)
-- also Elt a, Basis t

{-
-- a recursive typeclass for index flattening?
https://github.com/chessai/eigen/blob/89bb642f7fa0dbde0579b72c796b8706a3cf0c5b/src/Eigen/Internal.hsc#L81
class Cast (a :: Type) where
  type family C a = (result :: Type) | result -> a
  toC   :: a -> C a
  fromC :: C a -> a

instance Cast t => Cast (Dim n t) where fromC $! VS.unsafeIndex vals $! n * (fromC t + row

-}

-- TODO: add more base types. maybe replace TF OneOf with either of these approaches:
-- https://github.com/GU-CLASP/TypedFlow/blob/master/TypedFlow/Types.hs#L330-L443
-- http://hackage.haskell.org/package/hmatrix-0.19.0.0/docs/Numeric-LinearAlgebra.html#t:RealOf
-- | Class representing base elements of a tensor.
class (a ~ Float) => Elt a where

instance Elt Float

-- inductive typeclass for V a (V b (V c .... ))
instance Elt a => Dims a where reflectDims _ = []
-- this would make no sense:
--instance Elt a => Dim a where reflectDim _ = 0

class Dims s where reflectDims :: proxy s -> [Int]

-- TODO: remove overlapping: https://mzabani.github.io/posts/2018-08-13.html
instance {-# Overlapping #-} (Dim (V n v), Dims v) => Dims (V n v) where
  reflectDims _ = reflectDim (Proxy :: Proxy (V n v)) : reflectDims (Proxy :: Proxy v)

instance {-# Overlapping #-} (Dim (T n v), Dims v) => Dims (T n v) where
  reflectDims _ = reflectDim (Proxy :: Proxy (T n v)) : reflectDims (Proxy :: Proxy v)

instance {-# Overlapping #-} (Dim (F n v), Dims v) => Dims (F n v) where
  reflectDims _ = reflectDim (Proxy :: Proxy (F n v)) : reflectDims (Proxy :: Proxy v)

dims :: forall a. (Dims a) => a -> [Int]
dims _ = reflectDims (Proxy :: Proxy a)

-- TODO: figure out biplate or use recursion schemes?
--instance (KnownNat n, Data a) => Plated (V n a) where plate = biplate

-- TODO: update linear version or fork.
vLens :: Int -> Lens' (V n a) a
vLens i = \f (V v) -> (\a -> V (v V.// [(i, a)])) <$> f (v V.! i)

instance ( 1 <= n) => Field1  (V n a) (V n a) a a where _1  = vLens  0
instance ( 2 <= n) => Field2  (V n a) (V n a) a a where _2  = vLens  1
instance ( 3 <= n) => Field3  (V n a) (V n a) a a where _3  = vLens  2
instance ( 4 <= n) => Field4  (V n a) (V n a) a a where _4  = vLens  3
instance ( 5 <= n) => Field5  (V n a) (V n a) a a where _5  = vLens  4
instance ( 6 <= n) => Field6  (V n a) (V n a) a a where _6  = vLens  5
instance ( 7 <= n) => Field7  (V n a) (V n a) a a where _7  = vLens  6
instance ( 8 <= n) => Field8  (V n a) (V n a) a a where _8  = vLens  7
instance ( 9 <= n) => Field9  (V n a) (V n a) a a where _9  = vLens  8

-- | nested functions:
-- TODO: change to other rep?
{-
HS.build ::
  (KnownNat n, KnownNat m) =>
  (HS.ℝ -> HS.ℝ -> HS.ℝ) -> HS.L m n

withDims
  :: forall n m a. (Elem a, KnownNat n, KnownNat m)
  => (Int -> Int -> VS.Vector (C a)) -> Matrix n m a
withDims f =
  let !r = natToInt @n
      !c = natToInt @m
  in Matrix $ Vec $ f r c

-- TODO: is this rep better for constructing tensors?
--newtype F' n t = F' { unF' :: Tagged n Int -> t } deriving (Functor,Applicative,Generic,Generic1)
type F' n t = Tagged n Int -> Vector e
type F2' a b t = (KnownNat a, KnownNat b) => F' a (F' b t)
type F3' a b c t = (KnownNat a, KnownNat b, KnownNat c) => F' a (F' b (F' c t))

instance Num t => Num (Tagged n Int -> t) where
  (+) = liftA2 (+)
  (-) = liftA2 (-)
  (*) = liftA2 (*)

-}

newtype F (n :: Nat) t = F { unF :: Int -> t } deriving (Functor,Applicative,Generic,Generic1)

instance Additive (F n)
instance Dim n => Dim (F n t) where reflectDim _ = reflectDim (Proxy :: Proxy n)

instance Num a => Num (V.Vector a) where
  as + bs = V.zipWith (+) as bs
  {-# INLINE (+) #-}
  as - bs = V.zipWith (-) as bs
  {-# INLINE (-) #-}
  as * bs = V.zipWith (*) as bs
  {-# INLINE (*) #-}
  negate = fmap negate
  {-# INLINE negate #-}
  abs = fmap abs
  {-# INLINE abs #-}
  signum = fmap signum
  {-# INLINE signum #-}
  fromInteger = pure . fromInteger
  {-# INLINE fromInteger #-}

instance Num a => Num (F n a) where
  F as + F bs = F $ liftA2 (+) as bs
  {-# INLINE (+) #-}
  F as - F bs = F $ liftA2 (-) as bs
  {-# INLINE (-) #-}
  F as * F bs = F $ liftA2 (*) as bs
  {-# INLINE (*) #-}
  negate = fmap negate
  {-# INLINE negate #-}
  abs = fmap abs
  {-# INLINE abs #-}
  signum = fmap signum
  {-# INLINE signum #-}
  fromInteger = pure . fromInteger
  {-# INLINE fromInteger #-}

type F2 a b t = F a (F b t)
type F3 a b c t = F a (F b (F c t))

f26 :: F2 2 6 Float
f26 = pure 1

f265 :: F3 2 6 5 Float
f265 = pure 1

-- | nested 'Tagged's:
newtype T (n :: Nat) t = T { unT :: Tagged n t } deriving (Eq,Ord,Show,Functor,Applicative,Foldable,Traversable,Generic,Generic1)

instance Dim n => Additive (T n)
instance Dim n => Dim (T n t) where reflectDim _ = reflectDim (Proxy :: Proxy n)

instance Dim n => Additive (Tagged n)

instance Num a => Num (T n a) where
  (+) = liftA2 (+)
  (-) = liftA2 (-)
  (*) = liftA2 (*)
  negate = fmap negate
  {-# INLINE negate #-}
  abs = fmap abs
  {-# INLINE abs #-}
  signum = fmap signum
  {-# INLINE signum #-}
  fromInteger = pure . fromInteger
  {-# INLINE fromInteger #-}

type T2 a b t = T a (T b t)
type T3 a b c t = T a (T b (T c t))

-- TODO: sequence works but distribute doesn't for some reason
-- λ: t26 * (sequenceA $ t26 + t26)
-- Tagged (Tagged 2.0)

t26 :: T2 2 6 Float
t26 = pure 1

-- test: dims t265 == [2,6,5]
t265 :: T3 2 6 5 Float
t265 = pure 1

-- | nested 'Vector's:

type V2 a b t = (KnownNat a, KnownNat b) => V a (V b t)
type V3 a b c t = (KnownNat a, KnownNat b, KnownNat c) => V a (V b (V c t))

{-
 -
 -
Representable f =>

instance {-# Overlapping #-} (Representable v, KnownNat n) => Representable (V n v) where
  reflectDims _ = reflectDim (Proxy :: Proxy (V n v)) : reflectDims (Proxy :: Proxy v)

type Rep ?  = Idxs ds -> a

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

-}
-- TODO why isn't the applicative working?
b265 :: V3 2 6 5 Bool
b265 = pure . pure . pure $ False

v5 :: V 5 Float
v5 = return 1

v52 :: V2 5 2 Float
v52 = pure 1

v25 :: V2 2 5 Float
v25 = pure 1

v265 :: V3 2 6 5 Float
v265 = pure 1

v615 :: V3 6 1 5 Float
v615 = pure 1

v22 :: V 2 (V 2 Int)
v22 = pure 1

v222 :: V 2 (V 2 (V 2 Int))
v222 = pure 5

v312 :: V 3 (V 1 (V 2 Int))
v312 = pure 2

v4312 :: V 4 (V 3 (V 1 (V 2 Int)))
v4312 = pure 2

-- test: dims v4321 == [4,3,2,1]
v4321 :: V 4 (V 3 (V 2 (V 1 Int)))
v4321 = pure 2

v3512 :: V 3 (V 5 (V 1 (V 2 Int)))
v3512 = pure 2

{-
v232 :: V 2 (V 3 (V 2 Int))
v232 = unsafeCoerce
  [[[13, 14],
    [15, 16],
    [17, 18]],
   [[19, 20],
    [21, 22],
    [23, 24]]]

v223 :: V 2 (V 2 (V 3 Int))
v223 = unsafeCoerce
  [[[ 1,  2,  3],
    [ 4,  5,  6]],
   [[ 7,  8,  9],
    [10, 11, 12]]]
-}

----Tensile.Operators API
--

class Finite v where
  type Size v :: Nat -- this should allow kind k, for Reifies k Int
  toV :: v a -> V (Size v) a
  default toV :: Foldable v => v a -> V (Size v) a
  toV = V . V.fromList . Foldable.toList
  fromV :: V (Size v) a -> v a

instance Finite (V (n :: Nat)) where
  type Size (V n) = n
  toV = id
  fromV = id

-- TODO: remove
add :: (Additive f, Num a) => f a -> f a -> f a
add = (^+^)

sub :: (Additive f, Num a) => f a -> f a -> f a
sub = (^-^)

-- test: dim (scalar 1.7 :: V 1 Float) == 1
scalar :: (Additive f, Finite f, Elt a) => a -> f a
scalar a = fromV (V $ pure a)

-- (Tensor f, Ring a)
(.+.) :: (Additive f, Num a) => f a -> f a -> f a
(.+.) = add

(.-.) :: (Additive f, Num a) => f a -> f a -> f a
(.-.) = sub

-- | Right scalar product
(.*) :: (Functor f, Num a) => f a -> a -> f a
f .* a = fmap (* a) f

-- | Left scalar product
(*.) :: (Functor f, Num a) => a -> f a -> f a
a *. f = fmap (* a) f

-- | Right vector product
(.*!) :: (Functor m, Foldable r, Additive r, Num a) => m (r a) -> r a -> m a
m .*! v = fmap (\r -> Foldable.sum $ liftI2 (*) r v) m

-- (!*) :: (Metric r, Additive f, Num a) => r a -> r (f a) -> f a
-- f !* g = dot f <$> distribute g
--
-- | Left vector product
(!*.) :: (Foldable t, Additive f, Additive t, Num a) => t a -> t (f a) -> f a
f !*. g = sumV $ liftI2 (*.) f g

-- | Matrix product.
(.*.) :: (Functor m, Foldable t, Additive t, Additive n, Num a) => m (t a) -> t (n a) -> m (n a)
f .*. g = fmap (\f' -> Foldable.foldl' add zero $ liftI2 (*.) f' g) f

{-
infixl 6 :+:, :-:
infixl 7 :+:,.*:,:+.,*:,:+

sumV :: (Foldable f, Additive v, Num a) => f (v a) -> v a
sumV = Foldable.foldl' add zero

dot :: (Foldable f, Num a) => f a -> f a -> a
dot x y = Foldable.sum $ liftI2 R.(*) x y

-- | Product of a matrix of weights with a batched vector .
weight
  :: (Functor m, Foldable t, Additive t, Applicative t, Additive n, Num a, Traversable n)
  => n (t a) -> m (t a) -> m (n a)
weight m v = v .*. (sequenceA m) -- use distribute here instead?

--transpose :: (Distributive g, Functor f) => f (g a) -> g (f a)
--transpose = distribute

-}

-- | A generalization of 'Control.Lens.inside' that works over any corepresentable 'Functor'.
-- Useful for performing operations along different axes. See https://medium.com/@aerinykim/tensorflow-101-what-does-it-mean-to-reduce-axis-9f39e5c6dea2
-- for a tutorial on tensor axes.

-- λ: :t v4312 ^. (_1)
-- v4312 ^. (_1) :: V 3 (V 1 (V 2 Int))
-- λ: :t v4312 ^. (_1 . _1)
-- v4312 ^. (_1 . _1) :: V 1 (V 2 Int)
-- λ: :t v4312 ^. (_1 . _1 . _1)
-- v4312 ^. (_1 . _1 . _1) :: V 2 Int
-- λ: :t v4312 ^. (_1 . _1 . _1 ._1)
-- v4312 ^. (_1 . _1 . _1 ._1) :: Int
--
-- λ: :t v4312 ^. inside (_1 )
-- v4312 ^. inside (_1 ) :: V 4 (V 1 (V 2 Int))
-- λ: :t v4312 ^. inside (_1 . _1)
-- v4312 ^. inside (_1 . _1) :: V 4 (V 2 Int)
-- λ: :t v4312 ^. inside (_1 . _1 . _1)
-- v4312 ^. inside (_1 . _1 . _1) :: V 4 Int
--
-- λ: :t v4312 ^. (inside _1)
-- v4312 ^. (inside _1) :: V 4 (V 1 (V 2 Int))
-- λ: :t v4312 ^. (inside . inside $ _1)
-- v4312 ^. (inside . inside $ _1) :: V 4 (V 3 (V 2 Int))
-- λ: :t v4312 ^. (inside . inside . inside $ _1)
-- v4312 ^. (inside . inside . inside $ _1) :: V 4 (V 3 (V 1 Int))

inside
  :: Functor f
  => Representable g
  => LensLike (Context a b) s t a b
  -> (g a -> f (g b)) -> g s -> f (g t)
inside l = \afb s -> sbt s <$> afb (sa s)
  where
    sa = liftR $ \s -> I.ipos (accessors l s)

    sbt = liftR2 $ \s b -> I.ipeek b (accessors l s)

    -- Isomorphic to s -> (a, (b -> t))
    accessors l = l (Context id)

    liftR = fmapRep

-- TODO: push (V 1) constructor down n levels to add an extra dimension at axis (@n@) of size 1.
--
-- expandDim :: forall n s e. (KnownLen s, KnownPeano n) => Tensor s t -> Tensor (Take n s ++ (1 ': Drop n s)) t
-- expandDim (T x) = (T (funcall "tf.expand_dims" [x, named "axis" (integer (listLen @s - peanoInt @n))]))

--reduceSum0 :: ∀ s' n e. KnownLen s' => Tensor (n ': s') t -> Tensor s' t
reduceSum :: (Foldable f, Additive v, Num a) => f (v a) -> v a
reduceSum = sumV

-- dot prod along arbitrary axis
--
-- λ: :t v265 `dot` v265
-- v265 `dot` v265 :: V 6 (V 5 Float)
--
-- λ: :t liftI2 dot v265 v265
-- liftI2 dot v265 v265 :: V 2 (V 5 Float)
--
-- | For two tensors of the same dimension, the Hadamard product is a tensor of the same dimension as the operands, with elements given by pointwise multiplication.

--hadamard :: (Num a, Additive f) => f a -> f a -> f a
--hadamard = liftI2 (Prelude.(*))

{-
-- https://en.wikipedia.org/wiki/Kronecker_product
If A is an m × n matrix and B is a p × q matrix, then the Kronecker product A ⊗ B is the mp × nq block matrix formed by multiplying every scalar value in m by the matrix m.

Transposition and conjugate transposition are distributive over the Kronecker product:

transpose (m ⊗ n) == (transpose m) ⊗ (transpose n)
adjoint (m ⊗ n) == (adjoint m) ⊗ (adjoint n)

TODO: replace w/ version w/ fewer restrictions on a.
need to recurse to leaves and then build up matrices. zylomorph?
-}
kronecker :: (Functor f, Functor g, Num a) => f a -> g a -> f (g a)
kronecker = outer

-- $setup
-- >>> import Data.Complex
-- >>> import Data.IntMap
-- >>> import Debug.SimpleReflect.Vars
-- >>> import Linear.V

-- need Linear.Conjugate for this
-- | Hermitian conjugate or conjugate transpose
--
-- >>> adjoint (V2 (V2 (1 :+ 2) (3 :+ 4)) (V2 (5 :+ 6) (7 :+ 8)))
-- V2 (V2 (1.0 :+ (-2.0)) (5.0 :+ (-6.0))) (V2 (3.0 :+ (-4.0)) (7.0 :+ (-8.0)))
-- adjoint :: (Functor m, Distributive n, Conjugate a) => m (n a) -> n (m a)
-- adjoint = collect (fmap conjugate)
-- {-# INLINE adjoint #-}


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RebindableSyntax      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeSynonymInstances  #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Array.Accelerate.Data.Complex
-- Copyright   : [2015..2019] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- Complex numbers, stored in the usual C-style array-of-struct representation,
-- for easy interoperability.
--
module Data.Array.Accelerate.Data.Complex (

  -- * Rectangular from
  Complex(..), pattern (::+),
  real,
  imag,

  -- * Polar form
  mkPolar,
  cis,
  polar,
  magnitude, magnitude',
  phase,

  -- * Conjugate
  conjugate,

) where

import Data.Array.Accelerate.Array.Sugar
import Data.Array.Accelerate.Classes
import Data.Array.Accelerate.Data.Functor
import Data.Array.Accelerate.Pattern
import Data.Array.Accelerate.Prelude
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Type

import Data.Complex                                                 ( Complex(..) )
import Prelude                                                      (($))
import qualified Data.Complex                                       as C
import qualified Prelude                                            as P

infix 6 ::+
pattern (::+) :: Elt a => Exp a -> Exp a -> Exp (Complex a)
pattern r ::+ i <- (deconstructComplex -> (r, i))
  where (::+) = constructComplex
{-# COMPLETE (::+) #-}

-- Use an array-of-structs representation for complex numbers if possible.
-- This matches the standard C-style layout, but we can use this representation only at
-- specific types (not for any type 'a') as we can only have vectors of primitive type.
-- For other types, we use a structure-of-arrays representation. This is handled by the
-- ComplexRepr. We use the GADT ComplexR and function complexR to reconstruct
-- information on how the elements are represented.
--
instance Elt a => Elt (Complex a) where
  type EltRepr (Complex a) = ComplexRepr (EltRepr a)
  {-# INLINE eltType     #-}
  {-# INLINE [1] toElt   #-}
  {-# INLINE [1] fromElt #-}
  eltType = case complexR tp of
      ComplexRvec s -> TupRsingle $ VectorScalarType $ VectorType 2 s
      ComplexRtup   -> TupRunit `TupRpair` tp `TupRpair` tp
    where
      tp = eltType @a
  toElt = case complexR $ eltType @a of
    ComplexRvec _ -> \(Vec2 r i)   -> toElt r :+ toElt i
    ComplexRtup   -> \(((), r), i) -> toElt r :+ toElt i
  fromElt (r :+ i) = case complexR $ eltType @a of
    ComplexRvec _ -> Vec2 (fromElt r) (fromElt i)
    ComplexRtup   -> (((), fromElt r), fromElt i)

type family ComplexRepr a where
  ComplexRepr Half   = Vec2 Half
  ComplexRepr Float  = Vec2 Float
  ComplexRepr Double = Vec2 Double
  ComplexRepr Int    = Vec2 Int
  ComplexRepr Int8   = Vec2 Int8
  ComplexRepr Int16  = Vec2 Int16
  ComplexRepr Int32  = Vec2 Int32
  ComplexRepr Int64  = Vec2 Int64
  ComplexRepr Word   = Vec2 Word
  ComplexRepr Word8  = Vec2 Word8
  ComplexRepr Word16 = Vec2 Word16
  ComplexRepr Word32 = Vec2 Word32
  ComplexRepr Word64 = Vec2 Word64
  ComplexRepr a      = Tup2 a a

data ComplexR a c where
  ComplexRvec :: VecElt a => SingleType a -> ComplexR a (Vec2 a)
  ComplexRtup :: ComplexR a (Tup2 a a)

complexR :: TupleType a -> ComplexR a (ComplexRepr a)
complexR (TupRsingle (SingleScalarType (NumSingleType (FloatingNumType TypeHalf  )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (FloatingNumType TypeFloat )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (FloatingNumType TypeDouble)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt   )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt8  )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt16 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt32 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeInt64 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord  )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord8 )))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord16)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord32)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NumSingleType (IntegralNumType TypeWord64)))) = ComplexRvec singleType
complexR (TupRsingle (SingleScalarType (NonNumSingleType TypeChar)))                  = ComplexRtup
complexR (TupRsingle (SingleScalarType (NonNumSingleType TypeBool)))                  = ComplexRtup
complexR (TupRsingle (VectorScalarType (_)))                                          = ComplexRtup
complexR TupRunit                                                                     = ComplexRtup
complexR TupRpair{}                                                                   = ComplexRtup

constructComplex :: forall a. Elt a => Exp a -> Exp a -> Exp (Complex a)
constructComplex r i = case complexR $ eltType @a of
  ComplexRvec _ ->
    let
      r', i' :: Exp (EltRepr a)
      r' = coerce @a @(EltRepr a) r
      i' = coerce i
      v :: Exp (Vec2 (EltRepr a))
      v = V2 r' i'
    in
      coerce @(Vec2 (EltRepr a)) @(Complex a) $ v
  ComplexRtup   -> coerce $ T2  r i

deconstructComplex :: forall a. Elt a => Exp (Complex a) -> (Exp a, Exp a)
deconstructComplex c = case complexR $ eltType @a of
  ComplexRvec _ -> let V2 r i = coerce @(Complex a) @(Vec2 (EltRepr a)) c in (coerce r, coerce i)
  ComplexRtup   -> let T2 r i = coerce c in (r, i)

coerce :: EltRepr a ~ EltRepr b => Exp a -> Exp b
coerce (Exp e) = Exp e

instance (Lift Exp a, Elt (Plain a)) => Lift Exp (Complex a) where
  type Plain (Complex a) = Complex (Plain a)
  lift (r :+ i) = lift r ::+ lift i

instance Elt a => Unlift Exp (Complex (Exp a)) where
  unlift (r ::+ i) = r :+ i

instance Eq a => Eq (Complex a) where
  r1 ::+ c1 == r2 ::+ c2 = r1 == r2 && c1 == c2
  r1 ::+ c1 /= r2 ::+ c2 = r1 /= r2 || c1 /= c2

instance RealFloat a => P.Num (Exp (Complex a)) where
  (+)    = lift2 ((+) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (-)    = lift2 ((-) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  (*)    = lift2 ((*) :: Complex (Exp a) -> Complex (Exp a) -> Complex (Exp a))
  negate = lift1 (negate :: Complex (Exp a) -> Complex (Exp a))
  signum z@(x ::+ y) =
    if z == 0
       then z
       else let r = magnitude z
             in x/r ::+ y/r
  abs z         = magnitude z ::+ 0
  fromInteger n = fromInteger n ::+ 0

instance RealFloat a => P.Fractional (Exp (Complex a)) where
  fromRational x  = fromRational x ::+ 0
  z / z'          = (x*x''+y*y'') / d ::+ (y*x''-x*y'') / d
    where
      x  :+ y   = unlift z
      x' :+ y'  = unlift z'
      --
      x'' = scaleFloat k x'
      y'' = scaleFloat k y'
      k   = - max (exponent x') (exponent y')
      d   = x'*x'' + y'*y''

instance RealFloat a => P.Floating (Exp (Complex a)) where
  pi                = pi ::+ 0
  exp (x ::+ y)     = let expx = exp x
                       in expx * cos y ::+ expx * sin y
  log z             = log (magnitude z) ::+ phase z
  sqrt z@(x ::+ y)  =
    if z == 0
      then 0
      else u ::+ (y < 0 ? (-v, v))
    where
      T2 u v = x < 0 ? (T2 v' u', T2 u' v')
      v'     = abs y / (u'*2)
      u'     = sqrt ((magnitude z + abs x) / 2)

  x ** y =
    if y == 0 then 1 else
    if x == 0 then if exp_r > 0 then 0 else
                   if exp_r < 0 then inf ::+ 0
                                else nan ::+ nan
              else if isInfinite r || isInfinite i
                     then if exp_r > 0 then inf ::+ 0 else
                          if exp_r < 0 then 0
                                       else nan ::+ nan
                     else exp (log x * y)
    where
      r     ::+ i  = x
      exp_r ::+ _  = y
      --
      inf = 1 / 0
      nan = 0 / 0

  sin (x ::+ y)  = sin x * cosh y ::+ cos x * sinh y
  cos (x ::+ y)  = cos x * cosh y ::+ (- sin x * sinh y)
  tan (x ::+ y)  = (sinx*coshy ::+ cosx*sinhy) / (cosx*coshy ::+ (-sinx*sinhy))
    where
      sinx  = sin x
      cosx  = cos x
      sinhy = sinh y
      coshy = cosh y

  sinh (x ::+ y) = cos y * sinh x ::+ sin  y * cosh x
  cosh (x ::+ y) = cos y * cosh x ::+ sin y * sinh x
  tanh (x ::+ y) = (cosy*sinhx ::+ siny*coshx) / (cosy*coshx ::+ siny*sinhx)
    where
      siny  = sin y
      cosy  = cos y
      sinhx = sinh x
      coshx = cosh x

  asin z@(x ::+ y) = y' ::+ (-x')
    where
      x' ::+ y' = log (((-y) ::+ x) + sqrt (1 - z*z))

  acos z                    = y'' ::+ (-x'')
    where
      x'' ::+ y''  = log (z + ((-y') ::+ x'))
      x'  ::+ y'   = sqrt (1 - z*z)

  atan z@(x ::+ y) = y' ::+ (-x')
    where
      x' ::+ y' = log (((1-y) ::+ x) / sqrt (1+z*z))

  asinh z =  log (z + sqrt (1+z*z))
  acosh z =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
  atanh z =  0.5 * log ((1.0+z) / (1.0-z))

instance (FromIntegral a b, Num b, Elt (Complex b)) => FromIntegral a (Complex b) where
  fromIntegral x = fromIntegral x ::+ 0

-- | @since 1.2.0.0
--
instance Functor Complex where
  fmap f (r ::+ i) = f r ::+ f i

-- | The non-negative magnitude of a complex number
--
magnitude :: RealFloat a => Exp (Complex a) -> Exp a
magnitude (r ::+ i) = scaleFloat k (sqrt (sqr (scaleFloat mk r) + sqr (scaleFloat mk i)))
  where
    k     = max (exponent r) (exponent i)
    mk    = -k
    sqr z = z * z

-- | As 'magnitude', but ignore floating point rounding and use the traditional
-- (simpler to evaluate) definition.
--
-- @since 1.3.0.0
--
magnitude' :: RealFloat a => Exp (Complex a) -> Exp a
magnitude' (r ::+ i) = sqrt (r*r + i*i)

-- | The phase of a complex number, in the range @(-'pi', 'pi']@. If the
-- magnitude is zero, then so is the phase.
--
phase :: RealFloat a => Exp (Complex a) -> Exp a
phase z@(r ::+ i) =
  if z == 0
    then 0
    else atan2 i r

-- | The function 'polar' takes a complex number and returns a (magnitude,
-- phase) pair in canonical form: the magnitude is non-negative, and the phase
-- in the range @(-'pi', 'pi']@; if the magnitude is zero, then so is the phase.
--
polar :: RealFloat a => Exp (Complex a) -> Exp (a,a)
polar z =  T2 (magnitude z) (phase z)

-- | Form a complex number from polar components of magnitude and phase.
--
#if __GLASGOW_HASKELL__ <= 708
mkPolar :: forall a. RealFloat a => Exp a -> Exp a -> Exp (Complex a)
#else
mkPolar :: forall a. Floating a => Exp a -> Exp a -> Exp (Complex a)
#endif
mkPolar = lift2 (C.mkPolar :: Exp a -> Exp a -> Complex (Exp a))

-- | @'cis' t@ is a complex value with magnitude @1@ and phase @t@ (modulo
-- @2*'pi'@).
--
#if __GLASGOW_HASKELL__ <= 708
cis :: forall a. RealFloat a => Exp a -> Exp (Complex a)
#else
cis :: forall a. Floating a => Exp a -> Exp (Complex a)
#endif
cis = lift1 (C.cis :: Exp a -> Complex (Exp a))

-- | Return the real part of a complex number
--
real :: Elt a => Exp (Complex a) -> Exp a
real (r ::+ _) = r

-- | Return the imaginary part of a complex number
--
imag :: Elt a => Exp (Complex a) -> Exp a
imag (_ ::+ i) = i

-- | Return the complex conjugate of a complex number, defined as
--
-- > conjugate(Z) = X - iY
--
conjugate :: Num a => Exp (Complex a) -> Exp (Complex a)
conjugate z = real z ::+ (- imag z)


{-# LANGUAGE CPP #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module HVX.Internal.Primitives
  ( Var
  , Fun(..)
  , Expr(..)
  , getFun
  , getJacobian
  , getProperties
  ) where

import Numeric.LinearAlgebra

import HVX.Internal.Matrix
import HVX.Internal.DCP

type Var = String

data Expr vex mon where
  EConst :: Mat -> Expr 'Affine 'Const
  EVar   :: Var -> Expr 'Affine 'Nondec
  EFun   ::
#ifndef DISABLE_EXPR_CXT
      ValidVex vex =>
#endif
      Fun v1 m1 -> Expr v2 m2 -> Expr vex mon
  EAdd   ::
#ifndef DISABLE_EXPR_CXT
      ValidVex vex =>
#endif
    Expr v1 m1 -> Expr v2 m2 -> Expr vex mon

instance Show (Expr vex mon) where
  show (EConst mat) = show mat
  show (EVar   var) = var
  show (EFun   f a) = show f ++ "(" ++ show a ++ ")"
  show (EAdd   a b) = show a ++ " + " ++ show b

getProperties :: (GetVex vex, GetMon mon) => Expr vex mon -> String
getProperties expr = getVex expr ++ " " ++ getMon expr

data Fun (vex :: Vex) (mon :: Mon) where
  Mul                :: Mat -> Fun 'Affine 'Nonmon
  Abs                :: Fun 'Convex 'Nonmon
  Neg                :: Fun 'Affine 'Noninc
  Log                :: Fun 'Concave 'Nondec
  Exp                :: Fun 'Convex 'Nondec
  LogSumExp          :: Fun 'Convex 'Nondec
  Max                :: Fun 'Convex 'Nondec
  Min                :: Fun 'Concave 'Nondec
  Norm               :: Double -> Fun 'Convex 'Nonmon
  Berhu              :: Double -> Fun 'Convex 'Nonmon
  Huber              :: Double -> Fun 'Convex 'Nonmon
  Quadform           :: Mat -> Fun 'Convex 'Nonmon
  PowBaseP0          :: Fun 'Affine 'Const
  PowBaseP01         :: Double -> Fun 'Concave 'Nondec
  PowBaseP1          :: Fun 'Affine 'Nondec
  PowBaseP1InfEven   :: Integer -> Fun 'Convex 'Nonmon
  PowBaseP1InfNotInt :: Double -> Fun 'Convex 'Nondec

instance Show (Fun vex mon) where
  show (Mul _) = "mul"
  show Abs = "abs"
  show Neg = "-"
  show Log = "log"
  show Exp = "exp"
  show LogSumExp = "log_sum_exp"
  show Max = "max"
  show Min = "min"
  show (Norm p) = "norm" ++ show p
  show (Berhu m) = "berhu" ++ show m
  show (Huber m) = "huber" ++ show m
  show (Quadform m) = "quadform" ++ show m
  show PowBaseP0 = "pow0"
  show (PowBaseP01 p) = "pow" ++ show p
  show PowBaseP1 = "pow1"
  show (PowBaseP1InfEven p) = "pow" ++ show p
  show (PowBaseP1InfNotInt p) = "pow" ++ show p

getFun :: Fun vex mon -> Mat -> Mat
getFun (Mul a) x = a <> x
getFun Abs x = abs x
getFun Neg x = negate x
getFun Exp x = exp x
getFun Log x
  | anyMat (<= 0) x = error "Cannot take log of number <= 0."
  | otherwise = log x
getFun LogSumExp x = reduceMat (log . sum) $ exp x
getFun Max x = reduceMat maximum x
getFun Min x = reduceMat minimum x
getFun (Norm p) x = scalarMat $ lpnorm p x
getFun (Berhu m) x =
  cmap (\y -> if abs y <= m then abs y else (abs y ** 2 + m ** 2) / (2 * m)) x
getFun (Huber m) x =
  cmap (\y -> if abs y <= m then abs y ** 2 else 2 * m * abs y - m ** 2) x
getFun (Quadform m) x = tr x <> m <> x
getFun PowBaseP0 x = konst 1.0 ((rows x), (cols x))
getFun (PowBaseP01 p) x
  | allMat (0 <=) x = matrixPow p x
  | otherwise = error "Cannot raise negative number to power p with 0 < p < 1."
getFun PowBaseP1 x = x
getFun (PowBaseP1InfEven p) x = matrixPow (fromIntegral p) x
getFun (PowBaseP1InfNotInt p) x
  | allMat (0 <=) x = matrixPow p x
  | otherwise = error "Cannot raise negative number to a nonintegral power of p for p > 1."

getJacobian :: Fun vex mon -> Mat -> Mat
getJacobian (Mul a) _ = a
getJacobian Abs x = diagMat $ signum x
getJacobian Neg x = (-1) * ident (rows x)
getJacobian Log x = diagMat $ 1 / x
getJacobian Exp x = diagMat $ exp x
getJacobian LogSumExp x = tr $ scale (1 / sumElements (exp x)) $ exp x
getJacobian Max x = tr $ ei (rows x) i
  where (i, _) = maxIndex x
getJacobian Min x = tr $ ei (rows x) i
  where (i, _) = minIndex x
getJacobian (Norm p) x = tr $ if x == zeroVec n
  then
    zeroVec n
  else
    scale (1 / denominator) numerator
  where denominator = lpnorm p x ** (p - 1)
        numerator = x * abs x ** (pMat - 2)
        pMat = (1><1) [p]
        n = rows x
getJacobian (Berhu m) x = diagMat $
  cmap (\y -> if abs y <= m then signum y else y / m) x
getJacobian (Huber m) x = diagMat $
  cmap (\y -> if abs y <= m then 2 * y else 2 * m * signum y) x
getJacobian (Quadform m) x = scale 2 $ tr x <> m
getJacobian PowBaseP0 x = konst 0.0 (n, n)
  where n = rows x
getJacobian (PowBaseP01 p) x = scale p $ diagMat $ matrixPow (p - 1) x
getJacobian PowBaseP1 x = ident n
  where n = rows x
getJacobian (PowBaseP1InfEven intP) x = scale p $ diagMat $ matrixPow (p - 1) x
  where p = fromIntegral intP
getJacobian (PowBaseP1InfNotInt p) x = scale p $ diagMat $ matrixPow (p - 1) x


module Sound.Hommage.Filter
 (
 -- * Filters
   lowpass
 , highpass
 , bandpass
 , stretchpass
 , morphpass
-- , vocoder

-- , lowpassrect

 -- * FilterSpec
-- , specfilter
 , FilterSpec
 , shiftFilterSpec

 -- * FilterSpec and Lists
 , filterSpecToList
 , filterSpecsToLists
 , interpolFilterLists

 -- * Other
 , ffreq
 , ffreqi
 , average512

 -- * Filterfunctions
 , lowpassFS
 , highpassFS
 , bandpassFS
 , lowp0FS
 , lowp1FS
 , bandp0FS
 , bandp1FS
 , highp0FS
 , highp1FS

 )
 where

--import Sound.Hommage.DFTFilter
import Data.Complex

--import Sound.Hommage.FFT
---------------------------------------------------------------------------------------------------
lowpass :: [Double] -> [Double] -> [Double]
lowpass s c = concat $ map filterSpecToList $
 zipWith lowpassFS (average512 s) (average512 c)

highpass :: [Double] -> [Double] -> [Double]
highpass s c = concat $ map filterSpecToList $
 zipWith highpassFS (average512 s) (average512 c)

bandpass :: [Double] -> [Double] -> [Double] -> [Double]
bandpass r s c = concat $ map filterSpecToList $
 zipWith3 bandpassFS (average512 r) (average512 s) (average512 c)

stretchpass:: FilterSpec -> [Double] -> [Double]
stretchpass fs c = concat $ map filterSpecToList $
 map (flip shiftFilterSpec fs) $ map (2.0**) $ average512 c

morphpass:: FilterSpec -> FilterSpec -> [Double] -> [Double]
morphpass f1 f2 c =
 let l1 = filterSpecToList f1
     l2 = filterSpecToList f2
 in concat $ map (interpolFilterLists l1 l2) $ average512 c
---------------------------------------------------------------------------------------------------
{-

-- | first argument is width of filter, second is cutoff
lowpassfilter :: [Double] -> [Double] -> [Double] -> [Double]
lowpassfilter s c =

lowpassrect :: [Double] -> [Double] -> [Double] -> [Double]
lowpassrect s c =
 let i = ([]:) $ map filterSpecToList $ zipWith lowpassFS (average512 s) (average512 c)
 in rect_filter' 9 i

{-
specfilterFFT :: [FilterSpec] -> [Double] -> [Double]
specfilterFFT fs = fftt 8 (concat . zipWith fun fs )

fun :: FilterSpec -> [Complex Double] -> [Complex Double]
fun fs ((c0 :+ _) : cs) = zipWith f ((c0 :+ 0) : cs) (filterSpecToList fs)
 where
  f (x:+y) v = (x * v) :+ (y * v)
-}

-- | first argument is width of filter, second is cutoff
highpassfilter :: [Double] -> [Double] -> [Double] -> [Double]
highpassfilter s c =
 let i = concat $ map filterSpecToList $ zipWith highpassFS (average512 s) (average512 c)
 in dftfilter i

-- | first argument is width of filter, second is slope, third is cutoff
bandpassfilter :: [Double] -> [Double] -> [Double] -> [Double] -> [Double]
bandpassfilter r s c =
 let i = concat $ map filterSpecToList $ zipWith3 bandpassFS (average512 r) (average512 s) (average512 c)
 in dftfilter i

stretchfilter :: FilterSpec -> [Double] -> [Double] -> [Double]
stretchfilter fs c =
 let i = concat $ map filterSpecToList $ map (flip shiftFilterSpec fs) $ map (2.0**) $ average512 c
 in dftfilter i

morphfilter :: FilterSpec -> FilterSpec -> [Double] -> [Double] -> [Double]
morphfilter f1 f2 c =
 let l1 = filterSpecToList f1
     l2 = filterSpecToList f2
     i = concat $ map (interpolFilterLists l1 l2) $ average512 c
 in dftfilter i

---------------------------------------------------------------------------------------------------
specfilter :: [FilterSpec] -> [Double] -> [Double]
specfilter fs =
 let i = concat $ map filterSpecToList fs
 in dftfilter i
-}
---------------------------------------------------------------------------------------------------
-- | for filter coefficient number: (0.0 .. 1.0) -> (1.0 .. 512.0)
ffreq :: Double -> Double
ffreq k = r
 where
  k' = abs k
  r | k' > 1.0  = 512.0
    | otherwise = 2 ** (9.0 * k')

-- | floor of 'ffreq'
ffreqi :: Double -> Int
ffreqi = floor . ffreq

-- | the resulting list is 1 \/ 512 as long as the input list.
--   512 elements are read and their average is the next output value.
average512 :: [Double] -> [Double]
average512 = loop 512 0.0
 where
  loop 0 s xs     = (s / 512.0) : loop 512 0.0 xs
  loop k s (x:xs) = loop (k-1) (s+x) xs
  loop _ _ _      = []
--mkFilterFun :: (Double -> FilterSpec) -> (Double, Double) -> FilterFun Double
--mkFilterFun filt (cutoff, fmod) = withFilterEnv . mkFilterEnv . filt . (+cutoff) . (*fmod)
---------------------------------------------------------------------------------------------------
-- | Filtering starts at coeff nr 0 and value 0.0 (constant coeff is always zero).
--   The elements in FilterSpec define the next value and how many
--   coeffs it takes to reach this value.
type FilterSpec   = [(Int, Double)]

-- | 0..1: lower, 1..: higher
shiftFilterSpec :: Double -> FilterSpec -> FilterSpec
shiftFilterSpec d fs = map f fs
 where
  f (n,v) = (round (d * fromIntegral n), v)

-- | Converting a FilterSpec to a list of 512 Doubles
filterSpecToList :: FilterSpec -> [Double]
filterSpecToList xs = take 512 $ loop 0.0 xs
 where
  loop c []          = repeat c
  loop c ((n,d): ls) = let f = (d-c) / fromIntegral n in take n (iterate (+f) c) ++ loop d ls

-- | Converting and interpolating a sequence of FilterSpecs to a list of lists. The first elemet of the
-- tuple, an Int, describes how many lists it takes to reach the given FilterSpec.
filterSpecsToLists :: [(Int, FilterSpec)] -> [[Double]]
filterSpecsToLists = loop (filterSpecToList  [])
 where
  loop l []         = []
  loop l ((n,f):xs) = let f' = filterSpecToList f
                      in fun n l f' ++ loop f' xs
  fun n f1 f2       = let d = 1.0 / fromIntegral n
                      in take n $ map (interpolFilterLists f1 f2) $ iterate (+d) 0.0

interpolFilterLists :: [Double] -> [Double] -> Double -> [Double]
interpolFilterLists f1 f2 n | n >= 1.0  = f2
                            | n <= 0.0  = f1
                            | otherwise = let n' = 1.0 - n in
                                          zipWith (\a b -> a * n' + b * n) f1 f2

---------------------------------------------------------------------------------------------------
lowpassFS :: Double -> Double -> FilterSpec
lowpassFS d c = [(1,1.0), (round fd, 1.0), (round (fd * abs d), 0.0)]
 where
  fd = ffreq c

highpassFS :: Double -> Double -> FilterSpec
highpassFS d c = [(round fd, 0.0), (round (fd * abs d), 1.0)]
 where
  fd = ffreq c

-- | range, curve, cutoff
bandpassFS :: Double -- ^ width of freq-window
           -> Double --
           -> Double -- ^ cutoff
           -> FilterSpec
bandpassFS r s f = [(l1, 0.0), (l2, 1.0), (l3, 1.0), (l4, 0.0)]
 where

  m1 = abs r + 1
  m2 = m1 + abs s
  a = ffreqi (f / m2)
  b = ffreqi (f / m1)
  c = ffreqi (f * m1)
  d = ffreqi (f * m2)

  l1 = a
  l2 = b - a
  l3 = c - b
  l4 = d - c

bandpassFS' :: Double -- ^ width of freq-window
           -> Double --
           -> Double -- ^ cutoff
           -> FilterSpec
bandpassFS' r s f = [(1 + round a, 0.0), (1 + round b, 1.0), (1 + round c, 1.0), (1 + round d, 0.0)]
 where
  r' = 1.0 + abs r
  s' = 1.0 + abs s
  fd = ffreq f
  a  = fd
  b  = fd * s' - a
  c  = b * r'  - b
  d  = c * s'  - c
---------------------------------------------------------------------------------------------------
lowp0FS :: Double -> FilterSpec
lowp0FS f = [(0, 1.0), (fr, 1.0), (1, 0.0)]
 where
  fr = ffreqi f

lowp1FS :: Double -> FilterSpec
lowp1FS f = [(0, 1.0), (fr, 1.0), (fr, 0.0)]
 where
  fr = ffreqi f

bandp0FS :: Double -> FilterSpec
bandp0FS f = [(f0-1, 0.0), (f0, 1.0), (fr, 1.0), (1, 0.0)]
 where
  fr = ffreqi f
  f0 = fr - (div fr 2) + 1

bandp1FS :: Double -> FilterSpec
bandp1FS f = [(fh, 0.0), (fh, 1.0), (fh, 0.0)]
 where
  fr = ffreqi f
  fh = div fr 2

highp0FS :: Double -> FilterSpec
highp0FS f = [(fr, 0.0), (1, 1.0)]
 where
  fr = ffreqi f

highp1FS :: Double -> FilterSpec
highp1FS f = [(fh, 0.0), (fh, 1.0)]
 where
  fr = ffreqi f
  fh = div fr 2
---------------------------------------------------------------------------------------------------
{-
-- | first argument is control signal
vocoder :: [Double] -> [Double] -> [Double]
vocoder xs ys = dftsynthese $ zipWith f (dftanalyse xs) (dftanalyse ys)
 where
  f x (a :+ b) = let z = magnitude x in (a*z :+ b*z)
-}
---------------------------------------------------------------------------------------------------


{-# LANGUAGE Strict #-}
module FourierPinwheel.Multiplication where

import           Control.Concurrent.Async
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Storable     as VS
import           DFT.Plan
import           FourierPinwheel.Array
import           Utils.Parallel
import Filter.Utils
import Utils.Distribution

{-# INLINE multiplyBias #-}
multiplyBias ::
     DFTPlan
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyBias plan biasF (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs2) = do
  let planIDForward =
        DFTPlanID DFT1DG [numThetaFreq, numXFreq, numYFreq] [0, 1, 2]
      planIDBackward =
        DFTPlanID IDFT1DG [numThetaFreq, numXFreq, numYFreq] [0, 1, 2]
  vecs2F <- dftExecuteBatchP plan planIDForward vecs2
  let vecs3F = parMap rdeepseq (VS.zipWith (*) biasF) vecs2F
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq <$>
    dftExecuteBatchP plan planIDBackward vecs3F

{-# INLINE multiplyBias4D #-}
multiplyBias4D ::
     DFTPlan
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyBias4D plan biasF (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs2) = do
  let planIDForward =
        DFTPlanID
          DFT1DG
          [numRFreq, numThetaFreq, numXFreq, numYFreq]
          [0, 1, 2, 3]
      planIDBackward =
        DFTPlanID
          IDFT1DG
          [numRFreq, numThetaFreq, numXFreq, numYFreq]
          [0, 1, 2, 3]
  vecs2F <- dftExecute plan planIDForward . VS.concat $ vecs2
  let vecs3F = VS.zipWith (*) biasF vecs2F
  arr <-
    fromUnboxed (Z :. numRFreq :. numThetaFreq :. numXFreq :. numYFreq) .
    VS.convert <$>
    dftExecute plan planIDBackward vecs3F
  return .
    FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
    parMap
      rdeepseq
      (\rFreq ->
         VS.convert . toUnboxed . computeS . R.slice arr $ (Z :. rFreq :. All :. All :. All)) $
    [0 .. numRFreq - 1]

{-# INLINE multiplyBiasDiscrete #-}
multiplyBiasDiscrete ::
     DFTPlan
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyBiasDiscrete plan bias (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs2F) = do
  let planIDForward = DFTPlanID DFT1DG [numThetaFreq, numXFreq, numYFreq] [1, 2]
      planIDBackward =
        DFTPlanID IDFT1DG [numThetaFreq, numXFreq, numYFreq] [1, 2]
  vecs2 <-
    parMap rdeepseq (VS.zipWith (*) bias) <$>
    dftExecuteBatchP plan planIDBackward vecs2F
  FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq <$>
    dftExecuteBatchP plan planIDForward vecs2

{-# INLINE multiplyRFunction #-}
multiplyRFunction ::
     DFTPlan
  -> Double
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
multiplyRFunction plan periodEnv (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs) = do
  let freqConst = 2 * pi / log periodEnv
      filterVec =
        computeUnboxedS . makeFilter1D . fromFunction (Z :. numRFreq) $ \(Z :. i') ->
          let i = fromIntegral (i' - div numRFreq 2)
           in 1 / (1 :+ freqConst * fromIntegral i) -- gaussian1D i 10 :+ 0 --1 / (1 :+ (freqConst * fromIntegral i)^2)
      vec = VS.convert . VS.concat $ vecs
  filterF <-
    dftExecute plan (DFTPlanID DFT1DG [numRFreq] [0]) . VS.convert . toUnboxed $
    filterVec
  vecF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [numRFreq, numThetaFreq, numXFreq, numYFreq] [0])
      vec
  let arrF =
        fromUnboxed (Z :. numRFreq :. numThetaFreq :. numXFreq :. numYFreq) .
        VS.convert $
        vecF
      convolvedArrFR =
        R.traverse2
          arrF
          (fromUnboxed (Z :. numRFreq) . VS.convert $ filterF)
          const $ \fArr fVec idx@(Z :. r :. _ :. _ :. _) ->
          fArr idx * fVec (Z :. r)
  convolvedArrF <- computeUnboxedP convolvedArrFR
  convolvedVec <-
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, numXFreq, numYFreq] [0]) .
    VS.convert . toUnboxed $
    convolvedArrF
  let convolvedArr =
        fromUnboxed (Z :. numRFreq :. numThetaFreq :. numXFreq :. numYFreq) .
        VS.convert $
        convolvedVec
      convolvedVecs =
        parMap
          rdeepseq
          (\i ->
             VS.convert . toUnboxed . computeS . R.slice convolvedArr $
             (Z :. i :. All :. All :. All))
          [0 .. numRFreq - 1]
  return
    (FPArray
       numXFreq
       numYFreq
       numRFreq
       numThetaFreq
       numRhoFreq
       numPhiFreq
       convolvedVecs)

-- {-# INLINE multiplyRFunction #-}
-- multiplyRFunction :: FPArray (VS.Vector (Complex Double)) -> FPArray (VS.Vector (Complex Double))
-- multiplyRFunction (FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq vecs') =
--   let vecs = L.reverse vecs'
--       zeroVec = VS.replicate (numThetaFreq * numXFreq * numYFreq) 0
--       convolvedSGNPositive = L.init . L.scanl' (VS.zipWith (+)) zeroVec $ vecs
--       convolvedSGNNegative = L.tail . L.scanr (VS.zipWith (+)) zeroVec $ vecs
--       convolvedSGN =
--         parZipWith
--           rdeepseq
--           (VS.zipWith (\x y -> (x - y) * (0 :+ (-pi))))
--           convolvedSGNPositive
--           convolvedSGNNegative
--    in FPArray numXFreq numYFreq numRFreq numThetaFreq numRhoFreq numPhiFreq .
--       L.reverse $
--       convolvedSGN