module MyPhoto.Actions.Outliers
    ( rmOutliers
    ) where

import           Control.Concurrent ( getNumCapabilities )
import           Control.Concurrent.Async ( mapConcurrently )
import           Control.Concurrent.MSem as MS
import           Control.Monad
import           GHC.Conc ( numCapabilities )
import           System.Console.GetOpt
import           System.IO.Temp
import           System.Directory
import           System.Exit
import           System.FilePath
import           System.Process
import           Graphics.Netpbm
import qualified Data.ByteString as B
import           Data.Maybe (maybe)
import qualified Statistics.Sample as S
import qualified Data.Vector as V (fromList)

import MyPhoto.Model
import MyPhoto.Utils

help :: String
help = let
    header = "Usage: rmoutliers [OPTION...] files..."
  in usageInfo header options

data Options
  = Options
  { optMaxDistance :: Double
  , optSize        :: Int
  , optHelp        :: Bool
  } deriving Show

defaultOptions :: Options
defaultOptions
  = Options
  { optMaxDistance = 200
  , optSize        = 6
  , optHelp        = False
  }

options :: [OptDescr (Options -> Options)]
options =
  [ Option ['d'] ["maxDistance"]
      (OptArg ((\ d opts -> opts { optMaxDistance = d }) . maybe (optMaxDistance defaultOptions) read)
        "DISTANCE")
      "maximum distance to allow"
  , Option ['s'] ["size"]
      (OptArg ((\ s opts -> opts { optSize = s }) . maybe (optSize defaultOptions) read)
        "SIZE")
      "size of reduced image"
  , Option ['h'] ["help"]
      (NoArg (\ opts -> opts { optHelp = True }))
      "print help"
  ]

getMyOpts :: [String] -> IO (Options, [String])
getMyOpts argv = case getOpt Permute options argv of
                   (o,n,[]  ) -> return (foldl (flip id) defaultOptions o, n)
                   (_,_,errs) -> ioError (userError (concat errs ++ help))

computImgVec :: Int -> FilePath -> Img -> IO (Img, [Int])
computImgVec size tmpdir img = do
  let ppmFile = tmpdir </> (takeFileName img ++ ".ppm")

  (_, _, _, pHandle) <- createProcess (proc "convert" [img, "-resize", show size ++ "x" ++ show size ++ "!", ppmFile])
  exitCode <- waitForProcess pHandle
  unless (exitCode == ExitSuccess) $
    fail ("Resize failed with " ++ show exitCode)

  ppmResult <- parsePPM <$> B.readFile ppmFile
  case ppmResult of
    Right ([ppmImg],_) -> let
        imgVec = (pixelDataToIntList . ppmData) ppmImg
      in return (img, imgVec)
    Left err           -> fail ("PPM parsing failed with " ++ err)
    _                  -> fail "PPM parsing failed"

computImgsVecs :: Int -> FilePath -> [Img] -> IO [(Img, [Int])]
computImgsVecs size tmpdir imgs = do
  capabilities <- getNumCapabilities
  sem <- MS.new (capabilities - 2)
  mapConcurrently (MS.with sem . computImgVec size tmpdir) imgs

rmOutliersImpl :: [String] -> [Img] -> PActionBody
rmOutliersImpl _    []            = return (Right [])
rmOutliersImpl args imgs@(img1:_) = let
    dropByDistances :: Double -> [(Img, [Int])] -> [Int] -> IO [(Img, Double)]
    dropByDistances _ []                                 _       = return []
    dropByDistances maxDist ((img, curVec):imgsWithVecs) lastVec = let
        calculateDistance :: [Int] -> [Int] -> Double
        calculateDistance vec1 vec2 =
          sqrt .
          fromIntegral .
          sum .
          map ((\v -> v*v) . uncurry (-)) $ zip vec1 vec2
        dist = calculateDistance lastVec curVec
      in
        if dist < maxDist
        then do
          fmap ((img, dist) :) (dropByDistances maxDist imgsWithVecs curVec)
        else do
          putStrLn ("drop " ++ img ++ " with dist " ++ show dist)
          dropByDistances maxDist imgsWithVecs lastVec
    printStatsOnDists :: [(Img, Double)] -> IO ()
    printStatsOnDists imgsWithoutOutliers = let
        dists = V.fromList (map snd imgsWithoutOutliers)
      in do
        putStrLn $ "Variance = " ++ show (S.variance dists)
        putStrLn $ "Standard deviation = " ++ show (S.stdDev dists)
        putStrLn $ "Mean = " ++ show (S.mean dists)
  in do
    (opts, _) <- getMyOpts args
    print opts
    if optHelp opts
    then return (Left help)
    else withTempDirectory (takeDirectory img1) "_outliers.tmp"
           (\tmpdir -> do
               imgsWithVecs <- computImgsVecs (optSize opts) tmpdir imgs
               imgsWithoutOutliers <- dropByDistances (optMaxDistance opts) imgsWithVecs []
               printStatsOnDists imgsWithoutOutliers
               return (Right (map fst imgsWithoutOutliers))
             )

rmOutliers :: PrePAction
rmOutliers args     = logSeparator "rmoutliers" <> PAction (rmOutliersImpl args)


{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wwarn #-}

module Data.Distribution
  ( ToRealFrac(..)
  , Distribution(..)
  , computeDistribution
  , computeDistributionStats
  , mapToDistribution
  , zeroDistribution
  , dPercIx
  , dPercSpec
  , dPercSpec'
  , PercSpec(..)
  , renderPercSpec
  , Percentile(..)
  , pctFrac
  , stdPercentiles
  -- Aux
  , spans
  ) where

import Prelude (String, (!!), error, head, last, show)
import Cardano.Prelude hiding (head, show)

import Control.Arrow
import Data.Aeson (FromJSON(..), ToJSON(..))
import Data.Foldable qualified as F
import Data.List (span)
import Data.Vector (Vector)
import Data.Vector qualified as Vec
import Statistics.Sample qualified as Stat
import Text.Printf (PrintfArg, printf)

data Distribution a b =
  Distribution
  { dSize         :: Int
  , dAverage      :: a
  , dPercentiles  :: [Percentile a b]
  }
  deriving (Functor, Generic, Show)

instance (FromJSON a, FromJSON b) => FromJSON (Distribution a b)
instance (  ToJSON a,   ToJSON b) => ToJSON   (Distribution a b)

newtype PercSpec a = Perc { psFrac :: a } deriving (Eq, Generic, Show)

instance (FromJSON a) => FromJSON (PercSpec a)
instance (  ToJSON a) => ToJSON   (PercSpec a)

dPercIx :: Int -> Distribution a b -> b
dPercIx i d = pctSample $ dPercentiles d !! i

dPercSpec :: Eq (PercSpec a) => PercSpec a -> Distribution a b -> Maybe b
dPercSpec p = fmap pctSample . find ((== p) . pctSpec) . dPercentiles

dPercSpec' :: (Show a, Eq (PercSpec a)) => String -> PercSpec a -> Distribution a b -> b
dPercSpec' desc p =
  maybe (error er) pctSample . find ((== p) . pctSpec) . dPercentiles
 where
   er = printf "Missing centile %f in %s" (show $ psFrac p) desc

renderPercSpec :: PrintfArg a => Int -> PercSpec a -> String
renderPercSpec width = \case
  Perc x    -> printf ("%0."<>show (width-2)<>"f") x

data Percentile a b =
  Percentile
  { pctSpec        :: !(PercSpec a)
  , pctSample      :: !b
  }
  deriving (Functor, Generic, Show)

instance (FromJSON a, FromJSON b) => FromJSON (Percentile a b)
instance (  ToJSON a,   ToJSON b) => ToJSON   (Percentile a b)

pctFrac :: Percentile a b -> a
pctFrac = psFrac . pctSpec

stdPercentiles :: [PercSpec Float]
stdPercentiles =
  [ Perc 0.01, Perc 0.05
  , Perc 0.1, Perc 0.2, Perc 0.3, Perc 0.4
  , Perc 0.5, Perc 0.6
  , Perc 0.7, Perc 0.75
  , Perc 0.8, Perc 0.85, Perc 0.875
  , Perc 0.9, Perc 0.925, Perc 0.95, Perc 0.97, Perc 0.98, Perc 0.99
  , Perc 0.995, Perc 0.997, Perc 0.998, Perc 0.999
  , Perc 0.9995, Perc 0.9997, Perc 0.9998, Perc 0.9999
  ]

zeroDistribution :: Num a => Distribution a b
zeroDistribution =
  Distribution
  { dSize        = 0
  , dAverage     = 0
  , dPercentiles = mempty
  }

-- | For a list of distributions, compute a distribution of averages and rel stddev
-- (aka. coefficient of variance).
computeDistributionStats ::
    forall a v
  . ( v ~ Double -- 'v' is fixed by Stat.stdDev
    , Num a
    )
  => String -> [Distribution a v]
  -> Either String (Distribution a v, Distribution a v)
computeDistributionStats desc xs = do
  when (null xs) $
    Left $ "Empty list of distributions in " <> desc
  let distPcts    = dPercentiles <$> xs
      pctDistVals = transpose distPcts
  unless (all (pctLen ==) (length <$> distPcts)) $
    Left ("Distributions with different percentile counts: " <> show (length <$> distPcts) <> " in " <> desc)
  pure $ (join (***) (Distribution (length xs) 0)
          :: ([Percentile a v], [Percentile a v]) -> (Distribution a v, Distribution a v))
       $ unzip (pctsMeanCoV <$> pctDistVals)
 where
   pctLen = length . dPercentiles $ head xs

   pctsMeanCoV :: [Percentile a v] -> (Percentile a v, Percentile a v)
   pctsMeanCoV xs' = join (***) (Percentile . pctSpec $ head xs')
     (mean, Stat.stdDev vec / mean)
    where
      vec = Vec.fromList $ pctSample <$> xs'
      mean = Stat.mean vec

mapToDistribution :: (Real v, ToRealFrac v a) => (b -> v) -> [PercSpec a] -> [b] -> Distribution a v
mapToDistribution f pspecs xs = computeDistribution pspecs (f <$> xs)

computeDistribution :: (Real v, ToRealFrac v a) => [PercSpec a] -> [v] -> Distribution a v
computeDistribution percentiles (sort -> sorted) =
  Distribution
  { dSize        = size
  , dAverage     = toRealFrac (F.sum sorted) / fromIntegral (size `max` 1)
  , dPercentiles =
    (Percentile     (Perc 0)   mini:) .
    (<> [Percentile (Perc 1.0) maxi]) $
    percentiles <&>
      \spec ->
        let sample = if size == 0
                     then 0
                     else sorted !! indexAtFrac (psFrac spec)
        in Percentile spec sample
  }
 where size = length sorted
       indexAtFrac f = floor (fromIntegral (size - 1) * f)
       (,) mini maxi =
         if size == 0
         then (0,           0)
         else (head sorted, last sorted)

class RealFrac b => ToRealFrac a b where
  toRealFrac :: a -> b

instance RealFrac b => ToRealFrac Int b where
  toRealFrac = fromIntegral

instance {-# OVERLAPPABLE #-} (RealFrac b, Real a) => ToRealFrac a b where
  toRealFrac = realToFrac

spans :: forall a. (a -> Bool) -> [a] -> [Vector a]
spans f = go []
 where
   go :: [Vector a] -> [a] -> [Vector a]
   go acc [] = reverse acc
   go acc xs =
     case span f $ dropWhile (not . f) xs of
       ([], rest) -> go acc rest
       (ac, rest) ->
         go (Vec.fromList ac:acc) rest


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeApplications #-}

module ProblemSpec where

import Commons
import Control.Applicative (liftA2)
import Control.Concurrent
import Control.Monad (forM_, replicateM, replicateM_, unless, when)
import Data.Array
import Data.Complex (Complex (..))
import qualified Data.IntMap.Strict as IM
import Data.List (intercalate, sort)
import Data.List.Split (splitOn)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (fromList)
import Data.Maybe (fromJust)
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Debug.Trace (traceShowId)
import GHC.IO.Exception (ExitCode (..))
import HashedExpression.Derivative
import HashedExpression.Internal.CollectDifferential
import HashedExpression.Internal.Expression
import HashedExpression.Internal.Inner
import HashedExpression.Internal.Node
import HashedExpression.Internal.Normalize (normalize)
import HashedExpression.Internal.Utils
import HashedExpression.Interp
import HashedExpression.Operation
import HashedExpression.Prettify (showExp, showExpDebug)
import HashedExpression.Problem
import HashedExpression.Value
import System.Process (readProcess, readProcessWithExitCode)
import Test.HUnit
import Test.Hspec
import Test.QuickCheck
import Var
import Prelude hiding ((^))
import qualified Prelude

---- |
--isOneAfterAnother :: MemMap -> [Int] -> Bool
--isOneAfterAnother memMap nIds = all isOk xs
--  where
--    xs = zip nIds (tail nIds)
--    isOk (cur, nxt) =
--      let (offsetCur, _, shapeCur) =
--            fromJust $ IM.lookup cur (entryMap memMap)
--          (offsetNxt, _, _) = fromJust $ IM.lookup nxt (entryMap memMap)
--       in offsetCur + Prelude.product shapeCur == offsetNxt

-- |
prop_constructProblemNoConstraint :: SuiteScalarR -> Expectation
prop_constructProblemNoConstraint (Suite exp valMap) = do
  let names = Map.keys valMap
      pdMap = partialDerivativeMaps . collectDifferentials . exteriorDerivative (Set.fromList names) $ exp
      constructResult = constructProblem exp names (Constraint [])
  case constructResult of
    NoVariables -> return () -- it is possible that the random expression doesn't have any variables
    ProblemInvalid reason ->
      assertFailure $ "Can't construct problem: " ++ reason
    ProblemValid Problem {..} -> do
      let vars = map varName variables
      vars `shouldBe` Map.keys pdMap -- vars should be keys of partial differential map
      let ok variable
            | Just pId <- Map.lookup (varName variable) pdMap,
              pId == partialDerivativeId variable =
              True
            | otherwise = False
      assertBool "partial derivative ids aren't correct" $ all ok variables

--      assertBool "variables are not allocated consecutively" $
--        isOneAfterAnother memMap (map nodeId variables)

-- |
makeValidBoxConstraint :: (String, Shape) -> IO ConstraintStatement
makeValidBoxConstraint (name, shape) =
  case shape of
    [] -> do
      let x = variable name
      val1 <- VScalar <$> generate arbitrary
      val2 <- VScalar <$> generate arbitrary
      generate $
        elements [x .<= val1, x .>= val2, x `between` (val1, val2)]
    [size] -> do
      let x = variable1D @Default1D name
      val1 <- V1D . listArray (0, size - 1) <$> generate (vectorOf size arbitrary)
      val2 <- V1D . listArray (0, size - 1) <$> generate (vectorOf size arbitrary)
      generate $ elements [x .<= val1, x .>= val2, x `between` (val1, val2)]
    [size1, size2] -> do
      let x = variable2D @Default2D1 @Default2D2 name
      val1 <- V2D . listArray ((0, 0), (size1 - 1, size2 - 1)) <$> generate (vectorOf (size1 * size2) arbitrary)
      val2 <- V2D . listArray ((0, 0), (size1 - 1, size2 - 1)) <$> generate (vectorOf (size1 * size2) arbitrary)
      generate $ elements [x .<= val1, x .>= val2, x `between` (val1, val2)]

--        [size1, size2, size3] -> do TODO -- add 3D for tests
--            val1 <-
--                V3D . listArray ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1)) <$>
--                generate (vectorOf (size1 * size2 * size3) arbitrary)
--            val2 <-
--                V3D . listArray ((0, 0, 0), (size1 - 1, size2 - 1, size3 - 1)) <$>
--                generate (vectorOf (size1 * size2 * size3) arbitrary)
--            generate $
--                elements [x .<= val1, x .>= val2, x `between` (val1, val2)]

-- |
prop_constructProblemBoxConstraint :: SuiteScalarR -> Expectation
prop_constructProblemBoxConstraint (Suite exp valMap) = do
  let names = Map.keys valMap
  let df@(Expression dfN dfMp) =
        collectDifferentials . exteriorDerivative (Set.fromList names) $ exp
      pdMap = partialDerivativeMaps df
      vars = Map.keys pdMap
      varsWithShape = Map.toList $ Map.map (`retrieveShape` dfMp) pdMap
  bcs <- mapM makeValidBoxConstraint varsWithShape
  sampled <- generate $ sublistOf bcs
  let constraints = Constraint sampled
  let constructResult = constructProblem exp names constraints
  case constructResult of
    NoVariables -> return () -- it is possible that the random expression doesn't have any variables
    ProblemInvalid reason ->
      assertFailure $ "Can't construct problem: " ++ reason
    ProblemValid Problem {..} -> do
      let vars = map varName variables
      vars `shouldBe` Map.keys pdMap -- vars should be keys of partial differential map
      let ok variable
            | Just pId <- Map.lookup (varName variable) pdMap,
              pId == partialDerivativeId variable =
              True
            | otherwise = False
      assertBool "partial derivative ids aren't correct" $
        all ok variables
      --      assertBool "variables are not allocated consecutively" $
      --        isOneAfterAnother memMap (map nodeId variables)
      case (sampled, boxConstraints) of
        (_ : _, []) ->
          assertFailure
            "Valid box constraints but not appear in the problem"
        (_, bs) -> length sampled `shouldBe` length bs

-- |
makeValidScalarConstraint :: IO ConstraintStatement
makeValidScalarConstraint = do
  sc <-
    fst <$> generate (sized (genScalarR @Default1D @Default2D1 @Default2D2))
  val1 <- VScalar <$> generate arbitrary
  val2 <- VScalar <$> generate arbitrary
  generate $ elements [sc .<= val1, sc .>= val2, sc `between` (val1, val2)]

-- |
prop_constructProblemScalarConstraints :: SuiteScalarR -> Expectation
prop_constructProblemScalarConstraints (Suite exp valMap) = do
  let names = Map.keys valMap
  let df@(Expression dfN dfMp) =
        collectDifferentials . exteriorDerivative (Set.fromList names) $ exp
      pdMap = partialDerivativeMaps df
      vars = Map.keys pdMap
      varsWithShape = Map.toList $ Map.map (`retrieveShape` dfMp) pdMap
  -- box constraints
  bcs <- mapM makeValidBoxConstraint varsWithShape
  sampled <- generate $ sublistOf bcs
  -- scalar constraints
  numScalarConstraint <- generate $ elements [2 .. 4]
  scc <- replicateM numScalarConstraint makeValidScalarConstraint
  let constraints = Constraint $ sampled ++ scc
  let constructResult = constructProblem exp names constraints
  case constructResult of
    NoVariables -> return () -- it is possible that the random expression doesn't have any variables
    ProblemInvalid reason ->
      assertFailure $ "Can't construct problem: " ++ reason
    ProblemValid Problem {..} -> do
      let vars = map varName variables
      vars `shouldBe` Map.keys pdMap -- vars should be keys of partial differential map
      let ok variable
            | Just pId <- Map.lookup (varName variable) pdMap,
              pId == partialDerivativeId variable =
              True
            | otherwise = False
      assertBool "partial derivative ids aren't correct" $
        all ok variables
      --      assertBool "variables are not allocated consecutively" $
      --        isOneAfterAnother memMap (map nodeId variables)
      case (scc, scalarConstraints) of
        ([], _) -> return ()
        (_ : _, []) ->
          assertFailure
            "Having scalar constraints but not present in problem"
        (_, sConstraints) -> do
          let isOk sc =
                length (constraintPartialDerivatives sc)
                  `shouldBe` length vars
          assertBool "Empty constraint ?" $ not (null sConstraints)
          mapM_ isOk sConstraints

-- | List of hand-written problems and the expected result whether this problem is valid
problemsRepo :: [(ProblemResult, Bool)]
problemsRepo =
  [ ( let [x, y, z, t] = map (variable2D @128 @128) ["x", "y", "z", "t"]
          f = x <.> y + z <.> t
          constraints =
            Constraint
              [ x .>= VFile (TXT "x_lb.txt"),
                y .<= VFile (TXT "y_ub.txt"),
                x <.> z .>= VScalar 3
              ]
          vars = ["x", "y", "z", "t"]
       in constructProblem f vars constraints,
      True
    ),
    ( let [x, y, z, t] = map (variable2D @128 @128) ["x", "y", "z", "t"]
          f = x <.> y + z <.> t
          constraints =
            Constraint
              [ x .>= VScalar 5,
                y .<= VFile (TXT "y_ub.txt"),
                x <.> z .>= VScalar 3
              ]
          vars = ["x", "y", "z", "t"]
       in constructProblem f vars constraints,
      False
    ),
    ( let [x, y, z, t] = map (variable2D @128 @128) ["x", "y", "z", "t"]
          f = x <.> y + z <.> t
          constraints =
            Constraint [x .>= VNum 5, y .<= VNum 10, x <.> z .>= VNum 18]
          vars = ["x", "y", "z", "t"]
       in constructProblem f vars constraints,
      True
    )
  ]

spec :: Spec
spec =
  describe "Hash Solver spec " $ do
    specify "test hand-written problems" $ forM_ problemsRepo $ \(problemResult, expected) ->
      case (problemResult, expected) of
        (ProblemInvalid _, True) ->
          assertFailure "This problem is valid but fail to construct"
        (ProblemValid _, False) ->
          assertFailure "This problem is invalid but success to construct"
        _ -> return ()
    specify "valid problem should be constructed successfully" $
      property prop_constructProblemNoConstraint
    specify "valid box constrained problem should be constructed successfully" $
      property prop_constructProblemBoxConstraint
    specify "valid scalar constraints problem should be successfully successfully" $
      property prop_constructProblemScalarConstraints


{-# LANGUAGE BangPatterns                             #-}
{-# LANGUAGE DeriveFunctor                            #-}
{-# LANGUAGE DeriveGeneric                            #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GADTs                                    #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE NoStarIsType                             #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeFamilyDependencies                   #-}
{-# LANGUAGE TypeInType                               #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

module Numeric.Wavelet.Continuous (
  -- * Continuous Wavelet Transform
    CWD(..), mapCWD
  , CWDLine(..), mapCWDLine
  , LNorm(..), CWDOpts(..), defaultCWDO
  , cwd
  , cwdReal
  -- * Wavelets
  , AWavelet(..), mapAWavelet
  , morlet, morletFunc
  , meyer, meyerFunc
  , fbsp, fbspFunc
  ) where

import           Data.Complex
import           Data.Finite
import           Data.Foldable
import           Data.Maybe
import           Data.Ord
import           Data.Proxy
import           Data.Type.Equality
import           Data.Vector.Generic.Sized    (Vector)
import           GHC.Generics                 (Generic)
import           GHC.TypeLits.Compare
import           GHC.TypeNats
import           Math.FFT.Base
import           Numeric.Wavelet.Internal.FFT
import qualified Data.Vector.Generic          as UVG
import qualified Data.Vector.Generic.Sized    as VG
import qualified Data.Vector.Sized            as V

newtype CWD v n m a b = CWD { cwdLines :: V.Vector m (CWDLine v n a b) }
  deriving (Show, Functor)

data CWDLine v n a b = CWDLine
    { cwdlData  :: Vector v n b
    , cwdlScale :: Finite (n `Div` 2 + 1) -- ^ Scale factor, in number of ticks.
    , cwdlFreq  :: a                      -- ^ The frequency associated with this scale, in inverse tick
    , cwdlCoI   :: Finite (n `Div` 2 + 1) -- ^ How many items are /outside/ of the Cone of Influence, on each side.
    }
  deriving (Show, Functor)

mapCWDLine
    :: (UVG.Vector v b, UVG.Vector v c)
    => (b -> c)
    -> CWDLine v n a b
    -> CWDLine v n a c
mapCWDLine f (CWDLine d s q c) = CWDLine (VG.map f d) s q c

mapCWD
    :: (UVG.Vector v b, UVG.Vector v c)
    => (b -> c)
    -> CWD v n m a b
    -> CWD v n m a c
mapCWD f (CWD l) = CWD (mapCWDLine f <$> l)

data LNorm = L1 | L2
  deriving (Show, Eq, Ord, Enum, Bounded, Generic)

data CWDOpts n = CWDO
    { cwdoNorm     :: LNorm                     -- ^ wavelet normalization
    , cwdoMinScale :: Finite (n `Div` 2 + 1)    -- ^ min scale (period)
    , cwdoMaxScale :: Finite (n `Div` 2 + 1)    -- ^ max scale (period)
    }
  deriving (Show, Eq, Ord, Generic)

defaultCWDO :: KnownNat n => CWDOpts n
defaultCWDO = CWDO
    { cwdoNorm     = L2
    , cwdoMinScale = minBound
    , cwdoMaxScale = maxBound
    }

-- | 'cwd' for complex-valued analytic wavelets.
cwd :: forall v n m a b.
     ( UVG.Vector v (Complex b)
     , KnownNat n
     , KnownNat m
     , FFTWReal b
     , 1 <= n
     , RealFloat a
     )
    => AWavelet v a (Complex b)
    -> CWDOpts n
    -> Vector v n (Complex b)
    -> CWD v n m a (Complex b)
cwd AW{..} CWDO{..} xs = CWD . VG.generate $ \i ->
      let s   = scaleOf i
          dt  = 1/s
      in  case awVector dt of
            VG.SomeSized (wv :: Vector v q (Complex b))
              | Just Refl <- isLE (Proxy @1) (Proxy @q)
              , Just Refl <- isLE (Proxy @((q-1)`Div`2)) (Proxy @(q-1))
              -> let ys :: Vector v (n + q - 1) (Complex b)
                     ys  = (* (realToFrac (normie dt) :+ 0)) `VG.map` convolve xs wv
                     coi = fromMaybe maxBound . packFinite . round @a @Integer $ sqrt 2 * s
                     s'  = fromMaybe maxBound . packFinite . round @a @Integer $ s
                     ys' :: Vector v n (Complex b)
                     ys' = VG.slice @_ @((q - 1)`Div`2) @n @((q-1)-((q-1)`Div`2)) Proxy ys
                 in  CWDLine ys' s' (awFreq / s) coi
            _ -> error "Bad scale: wavelet vector is empty?"
  where
    n = natVal (Proxy @n)
    m = natVal (Proxy @m)
    normie :: a -> a
    normie = case cwdoNorm of
      L1 -> sqrt
      L2 -> id
    minScale = fromIntegral cwdoMinScale `max` 1
    maxScale = (fromIntegral cwdoMaxScale `min` (fromIntegral n / (2 * sqrt 2)))
         `max` (minScale + 1)
    scaleStep = (log maxScale - log minScale) / (fromIntegral m - 1)
    scaleOf :: Finite m -> a
    scaleOf i = exp $ log minScale + fromIntegral i * scaleStep

-- | 'cwd' for real-valued analytic wavelets.
cwdReal
    :: forall v n m a b.
     ( UVG.Vector v b
     , UVG.Vector v (Complex b)
     , KnownNat n
     , KnownNat m
     , FFTWReal b
     , 1 <= n
     , RealFloat a
     )
    => AWavelet v a b
    -> CWDOpts n
    -> Vector v n b
    -> CWD v n m a b
cwdReal aw cwdo = mapCWD realPart
                . cwd (mapAWavelet (:+ 0) aw) cwdo
                . VG.map (:+ 0)

-- | Analytical Wavelet
data AWavelet v a b = AW
    { awVector :: a -> v b    -- ^ generate a vector within awRange with a given dt
    , awFreq   :: a           -- ^ Dominant frequency component
    , awRange  :: a           -- ^ range away from zero outside of which wavelet can be considered negligible
    }
  deriving Functor

mapAWavelet
    :: (UVG.Vector v b, UVG.Vector v c)
    => (b -> c)
    -> AWavelet v a b
    -> AWavelet v a c
mapAWavelet f (AW v q r) = AW (UVG.map f . v) q r

morlet
    :: (UVG.Vector v (Complex a), RealFloat a)
    => a
    -> AWavelet v a (Complex a)
morlet σ = AW{..}
  where
    awRange  = 4
    (!awFreq, mf) = morletFunc_ σ
    awVector = renderFunc awRange mf

morletFunc :: RealFloat a => a -> a -> Complex a
morletFunc = snd . morletFunc_

morletFunc_ :: RealFloat a => a -> (a, a -> Complex a)
morletFunc_ σ = (q, f)
  where
    f t = (c * exp(-t*t/2) :+ 0) * (exp (0 :+ (σ * t)) - (exp (-σ2/2) :+ 0))
    !c   = pi ** (-1/4) * (1 + exp (-σ2) - 2 * exp (-3/4*σ2)) ** (-1/2)
    !σ2  = σ * σ
    !q   = converge 20 iter σ / (2 * pi)
    iter w = σ / (1 - exp (-σ * w))

meyer
    :: (UVG.Vector v a, RealFloat a)
    => AWavelet v a a
meyer = AW{..}
  where
   awRange  = 6
   awVector = renderFunc awRange meyerFunc
   awFreq   = 4 * pi / 3

meyerFunc :: RealFloat a => a -> a
meyerFunc t
    | isNaN ψ || isInfinite ψ = 0
    | otherwise               = ψ
  where
    t' = t - 0.5
    t'3 = t'**3
    sinTerm = sin(4*pi/3*t') / pi
    ψ1 = (4/3/pi*t'*cos(2*pi/3*t') - sinTerm)
       / (t' - 16/9 * t'3)
    ψ2 = (8/3/pi*t'*cos(8*pi/3*t') + sinTerm)
       / (t' - 64/9 * t'3)
    ψ = ψ1 + ψ2

fbsp
    :: (UVG.Vector v (Complex a), FFTWReal a)
    => Int      -- ^ m, >= 1
    -> a        -- ^ f_b, bandwidth
    -> a        -- ^ f_c, wavelet center frequency
    -> AWavelet v a (Complex a)
fbsp m fb fc = AW{..}
  where
    awRange  = 4/fb
    awVector = renderFunc awRange (fbspFunc m fb fc)
    awFreq   = autoDeriveFreq awRange awVector

autoDeriveFreq
    :: (UVG.Vector v (Complex a), FFTWReal a)
    => a
    -> (a -> v (Complex a))
    -> a
autoDeriveFreq r fv = case fv 0.001 of
    VG.SomeSized v ->
      let vv    = zip [1..] . map magnitude . VG.toList $ fft v
          (i,_) = maximumBy (comparing snd) vv
      in  fromInteger i / (r * 2)
    _ -> error "bad vector"

fbspFunc :: RealFloat a => Int -> a -> a -> a -> Complex a
fbspFunc m fb fc t =
    ((sqrt fb * sinc (fb * t / fromIntegral m))^m :+ 0) * exp (0 :+ (2 * pi * fc * t))
  where
    sinc x = sin x / x

-- | Render the effective range of a wavelet (based on 'awRange'), centered
-- around zero.  Takes a timestep.
renderFunc
    :: (UVG.Vector v b, RealFrac a)
    => a              -- ^ range about zero
    -> (a -> b)       -- ^ func
    -> a              -- ^ dt
    -> v b
renderFunc r f dt = UVG.generate (round n) $ \i ->
    f (fromIntegral i * dt - r)
  where
    n  = r * 2 / dt

converge
    :: (Fractional a, Ord a)
    => Int      -- ^ maximum iterations
    -> (a -> a) -- ^ function to find the fixed point convergence
    -> a        -- ^ starting value
    -> a
converge n f = go 0
  where
    go !i !x
        | i >= n                = 0
        | abs (x - y) < 0.00001 = x
        | otherwise             = go (i + 1) y
      where
        !y = f x


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}

import qualified Control.Foldl                as F
import           Control.Lens
import qualified Data.Attoparsec.Text         as A
import qualified Data.ByteString.Lazy         as BS
import           Data.Maybe                   (fromJust)
import           Data.Semigroup               ((<>))
import           Data.Serialize
import           Data.Singletons              (SingI)
import           Data.Singletons.Prelude.List (Head, Last)
import           GHC.Exts
import           GHC.TypeLits
import           Grenade
import           Grenade.Recurrent
import qualified Numeric.LinearAlgebra.Static as SA
import           Options.Applicative
import           Pipes
import qualified Pipes.ByteString             as PBS
import           Pipes.Group
import qualified Pipes.Prelude                as P
import qualified Pipes.Prelude.Text           as PT
import           System.IO                    (BufferMode (..), IOMode (..),
                                               hSetBuffering, stdout, withFile)

-- The defininition for our simple recurrent network.
-- This file just trains a network to generate a repeating sequence
-- of 0 0 1.
--
-- The F and R types are Tagging types to ensure that the runner and
-- creation function know how to treat the layers.
type R = Recurrent
type F = FeedForward
type TRNNIn = 4
type TRNNOut = 50
type TFFNOut = 1
type TRNN = R (LSTM TRNNIn TRNNOut)
type TFFN = F (FullyConnected TRNNOut TFFNOut)
-- type NJetInputs = 4
-- type NJetOutputs = 14
-- type JetLayer = Concat ('D1 4) Trivial ('D1 4) TracksLayer

type TracksNetShape = '[ TRNN, TFFN ]

type TracksNet =
  RecurrentNetwork
    TracksNetShape
    '[ 'D1 TRNNIn, 'D1 TRNNOut, 'D1 TFFNOut ]

type TracksInput = RecurrentInputs TracksNetShape

data Opts = Opts Int (Maybe String) LearningParameters

feedForward' :: Parser Opts
feedForward' =
  Opts
    <$> option auto (long "examples" <> short 'e' <> value 40000)
    <*> optional (strOption (long "input" <> short 'i'))
    <*> ( LearningParameters
          <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
          <*> option auto (long "momentum" <> value 0.9)
          <*> option auto (long "l2" <> value 0.0005)
        )

parseJet :: A.Parser [(S ('D1 TRNNIn), Maybe (S ('D1 1)))]
parseJet = do
  nB <- A.double
  tracks <- A.many1 . A.count 4 $ A.char '\t' *> A.double

  return $
    appLast Nothing (Just . S1D $ SA.konst nB)
    $ S1D . SA.fromList <$> tracks

  where
    appLast _ _ []     = []
    appLast _ y [z]    = [(z, y)]
    appLast x y (z:zs) = (z, x) : appLast x y zs

-- can't export this....
trainRecurrentF
  :: (Num (RecurrentInputs layers), SingI (Last shapes), Monad m)
  => LearningParameters
  -> m (RecurrentNetwork layers shapes, RecurrentInputs layers)
  -> F.FoldM
     m
     [(S (Head shapes), Maybe (S (Last shapes)))]
     (RecurrentNetwork layers shapes, RecurrentInputs layers)
trainRecurrentF lps start = F.FoldM step start done
  where
    step (!rnet, !rins) samp =
      return $ trainRecurrent lps rnet rins samp
    done = return

runRecurrentP
  :: Monad m
  => RecurrentNetwork layers shapes
  -> RecurrentInputs layers
  -> Pipe [S (Head shapes)] [S (Last shapes)] m b
runRecurrentP net inps = do
  xs <- await
  let (_, _, ys) = runRecurrentNetwork net inps xs
  yield ys
  runRecurrentP net inps

main :: IO ()
main = do
    Opts nex maybeIn rate <- execParser (info (feedForward' <**> helper) idm)
    hSetBuffering stdout LineBuffering
    start <-
      case maybeIn of
        Nothing -> randomRecurrent
        Just s ->
          withFile s ReadMode $ \h -> do
            res <- decodeLazy <$> BS.hGetContents h
            case res of
              Left err -> error err
              Right x  -> return x

    putStrLn "Training network..."

    let p =
          PT.stdinLn
          >-> P.map (A.parseOnly parseJet)
          >-> P.concat

    (net :: TracksNet, inps :: TracksInput) <-
      F.impurely P.foldM (trainRecurrentF rate $ return start)
      . over (chunksOf 1000) (maps $ \x -> liftIO (putStrLn "trained 1000 more") >> x)
      $ p >-> P.take nex

    withFile "net.out" WriteMode $ \h ->
      runEffect $ PBS.fromLazy (encodeLazy (net, inps)) >-> PBS.toHandle h

    putStrLn ""

    runEffect $
      p
      >-> P.mapM (
        \x -> do
          putStr . show . fromS1D . fromJust . snd $ last x
          putStr ", "
          return x
        )
      >-> P.map (fmap fst)
      >-> runRecurrentP net inps
      >-> P.mapM_ (print . fromS1D . last)

  where
    fromS1D :: KnownNat n => S ('D1 n) -> Double
    fromS1D (S1D r) = head . toList $ SA.unwrap r


{-# LANGUAGE RecordWildCards #-}

import Control.Monad
import Data.Complex
import Control.Monad.Trans.Either
import Foreign.C.Types
import Data.Word

import qualified Data.Vector.Storable as VS hiding ((++))
import qualified Data.Vector.Generic  as VG hiding ((++))
import           Pipes
import qualified Pipes.Prelude        as P
import           Options.Applicative

import SDR.Filter
import SDR.RTLSDRStream
import SDR.Util
import SDR.FFT
import SDR.Plot
import SDR.Demod
import SDR.Pulse
import SDR.PipeUtils
import SDR.ArgUtils
import SDR.CPUID

import Graphics.DynamicGraph.Waterfall
import Graphics.DynamicGraph.Util

--The filter coefficients are stored in another module
import Coeffs

data Options = Options {
    frequency :: Word32
}

optParser :: Parser Options
optParser = Options
          <$> option (fmap fromIntegral parseSize) (
                 long "frequency"
              <> short 'f'
              <> metavar "FREQUENCY"
              <> help "Frequency to tune to"
              )

opt :: ParserInfo Options
opt = info (helper <*> optParser) (fullDesc <> progDesc "SDR library demo. Receive FM broadcast band radio." <> header "SDR demo")

samples      = 8192
audioSamples = 1024

{-
    sampling frequency of the input is 1280 khz
    this is decimated by a factor of 8 and then demodulated
    sampling frequency of demodulated signal is 160 khz
    need audio output at 48 khz
    resampling factor is 48/160 == 3/10
    FM pilot tone at 19khz (0.3958 * 48)
    start audio filter cutoff at 15khz (0.3125 * 48)
-}

doIt Options{..} = do

    info <- lift getCPUInfo

    --Initialize GLFW
    res            <- lift setupGLFW
    unless res (left "Unable to initilize GLFW")

    --Initialize the RTLSDR device
    str            <- sdrStream frequency 1280000 1 (fromIntegral samples * 2)

    --Initialize the graphs and FFTs
    rfFFT          <- lift $ fftw samples
    rfSpectrum     <- plotWaterfall 1024 480 samples 1000 jet_mod
    audioFFT       <- lift $ fftwReal audioSamples
    audioSpectrum  <- plotFillAxes 1024 480 ((audioSamples `quot` 2) + 1) jet (zeroAxes 1024 480 48 10)

    --Initialize the sound sink
    pulseSink      <- lift pulseAudioSink

    --Initialize the filters
    deci <- lift $ fastDecimatorC info 8 coeffsRFDecim
    resp <- lift $ fastResamplerR info 3 10 coeffsAudioResampler
    filt <- lift $ fastFilterSymR info coeffsAudioFilter

    let window0 = hanning samples :: VS.Vector Float
        window  = hanning samples

    --Build the pipeline
    lift $ runEffect $   str
                     >-> P.map convertCAVX
                     >-> foldl1 combine [

                            P.map (VG.zipWith (flip mult) window0 . VG.zipWith mult (fftFixup samples))
                        >-> P.map (VG.map (cplxMap (realToFrac :: Float -> CDouble)))
                        >-> rfFFT
                        >-> P.map (VG.map ((* (32 / fromIntegral samples)) . realToFrac . magnitude))
                        >-> rfSpectrum,

                            firDecimator deci samples
                        >-> fmDemod
                        >-> firResampler resp samples
                        >-> firFilter filt audioSamples
                        >-> foldl1 combine [

                                P.map (VG.zipWith (*) window)
                            >-> P.map (VG.map (realToFrac :: Float -> CDouble))
                            >-> audioFFT
                            >-> P.map (VG.map ((/ 10) . realToFrac . magnitude))
                            >-> audioSpectrum,

                                P.map (VG.map ((* 0.2) . realToFrac))
                            >-> pulseSink
                        ]
                    ]

main = execParser opt >>= eitherT putStrLn return . doIt


{-# LANGUAGE OverloadedStrings #-}

import           Chapter04.LinearAlgebra.ListBasedVectorProperties (linearAlgebraTests)
import           StatisticsTests                                   (statsSpecs)
import           Test.Tasty

main :: IO ()
main = do
  statsSpecs' <- statsSpecs
  let tests = testGroup "all tests" [ statsSpecs', linearAlgebraTests ]
  defaultMain tests

currentTests :: IO ()
currentTests = statsSpecs >>= \x ->
  defaultMain $ testGroup "current development modules" [ x ]


{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE CPP #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
#endif
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2012 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Orphans
-----------------------------------------------------------------------------
module Linear.Instances () where

import Control.Applicative
import Control.Monad.Fix
import Control.Monad.Zip
import Data.Complex
import Data.Foldable
import Data.Functor.Bind
import Data.HashMap.Lazy as HashMap
import Data.Hashable
import Data.Semigroup
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
import Data.Traversable

instance (Hashable k, Eq k) => Apply (HashMap k) where
  (<.>) = HashMap.intersectionWith id

instance (Hashable k, Eq k) => Bind (HashMap k) where
  -- this is needlessly painful
  m >>- f = HashMap.fromList $ do
    (k, a) <- HashMap.toList m
    case HashMap.lookup k (f a) of
      Just b -> [(k,b)]
      Nothing -> []

instance Functor Complex where
  fmap f (a :+ b) = f a :+ f b
  {-# INLINE fmap #-}

instance Apply Complex where
  (a :+ b) <.> (c :+ d) = a c :+ b d

instance Applicative Complex where
  pure a = a :+ a
  (a :+ b) <*> (c :+ d) = a c :+ b d

instance Bind Complex where
  (a :+ b) >>- f = a' :+ b' where
    a' :+ _  = f a
    _  :+ b' = f b
  {-# INLINE (>>-) #-}

instance Monad Complex where
  return a = a :+ a
  {-# INLINE return #-}

  (a :+ b) >>= f = a' :+ b' where
    a' :+ _  = f a
    _  :+ b' = f b
  {-# INLINE (>>=) #-}

instance MonadZip Complex where
  mzipWith = liftA2

instance MonadFix Complex where
  mfix f = (let a :+ _ = f a in a) :+ (let _ :+ a = f a in a)

instance Foldable Complex where
  foldMap f (a :+ b) = f a `mappend` f b
  {-# INLINE foldMap #-}

instance Traversable Complex where
  traverse f (a :+ b) = (:+) <$> f a <*> f b
  {-# INLINE traverse #-}

instance Foldable1 Complex where
  foldMap1 f (a :+ b) = f a <> f b
  {-# INLINE foldMap1 #-}

instance Traversable1 Complex where
  traverse1 f (a :+ b) = (:+) <$> f a <.> f b
  {-# INLINE traverse1 #-}


{-# LANGUAGE RecordWildCards #-}

{-|
Module      :  HABQTlib.Data.Particle

Data structures and functions that deal with storing and processing particle
hierarchies.

/Warning/: functions in this module assume that the 'ptsParticles' is non-empty
and 'NumberOfParticles', 'Dim', and 'Rank' are positive, no validation is
performed. If you use them directly, instead of employing API from
"HABQTlib", you must ensure those assumptions hold.
-}
module HABQTlib.Data.Particle
  ( Particles(..)
  , genParticles
  , updateParticles
  , ParticleHierarchy
  , initialiseParticleHierarchy
  , updateParticleHierarchy
  , getMixedEstimate
  , foldOverPts
  , reduceParticlesToMean
  , effectiveSize
  , ResampleArgs(..)
  , resampleMultinom
  , resample
  , ecdf
  , icdf
  , nudgeParticle
  ) where

import Control.Monad (when)
import Control.Newtype.Generics (over)
import Data.Bool.HT (if', select)
import qualified Data.Vector as V
import HABQTlib.Data
import HABQTlib.RandomStates
import Numeric.LinearAlgebra (Complex(..))
import qualified Numeric.LinearAlgebra as LA
import qualified System.Random.MWC as MWC
import Text.Printf (printf)

-- | A vector of weighed density matrices is stored along with their rank and
-- number. 'ptsWeight' corresponds to the collective weight of particles of
-- rank 'ptsRank' in the hierarchical model, it is not the sum of individual
-- weights of particles (that is normalised to unity after every update).
data Particles = Particles
  { ptsRank :: Rank
  , ptsWeight :: Weight
  , ptsNumber :: NumberOfParticles
  , ptsParticles :: V.Vector WeighedDensityMatrix
  } deriving (Show)

-- | Particle hierarchy is described by a vector of 'Particles'.
type ParticleHierarchy = V.Vector Particles

-- | Generates random particles (according to induced measure).
genParticles :: Dim -> Rank -> NumberOfParticles -> IO Particles
genParticles d r n =
  let w = 1 / fromIntegral n
   in Particles r 1 n . fmap (mkWDM w) <$> V.replicateM n (genDM d r)

-- | Summarise particles to a mean estimate (weighed by the corresponding
-- hierarchical weight of the rank).
reduceParticlesToMean :: Particles -> WeighedDensityMatrix
reduceParticlesToMean Particles {..} =
  let wdm = V.foldl1' (<+>) ptsParticles
      wdmw = over WeighedDensityMatrix (\(w, dm) -> (w * ptsWeight, dm)) wdm
   in truncateRank ptsRank wdmw

-- | Map density matrices, combine them with their weights, and then perform a
-- (strict left) fold of results.
foldOverPts ::
     (DensityMatrix -> a) -- ^ function to map over density matrices
  -> (Weight -> a -> b) -- ^ function to combine weights with results of mapping
  -> (c -> b -> c) -- ^ fold funciton
  -> c -- ^ seed value for folding
  -> Particles
  -> c
foldOverPts f wf fld z Particles {..} =
  let wm (WeighedDensityMatrix (w, dm)) = wf w (f dm)
   in V.foldl' (\l r -> fld l (wm r)) z ptsParticles

fullDataLogLikelihood :: [PureStateVector] -> DensityMatrix -> Double
fullDataLogLikelihood vs dm =
  let lps = map (log . (`pureStateLikelihood` dm)) vs
   in sum lps

-- | Given a measurement result, perform a Bayesian update over the particles.
updateParticles :: PureStateVector -> Particles -> Particles
updateParticles sv pts@Particles {..} =
  let updateF :: WeighedDensityMatrix -> WeighedDensityMatrix
      updateF (WeighedDensityMatrix (w, dm)) = WeighedDensityMatrix (wnew, dm)
        where
          wnew = w * pureStateLikelihood sv dm
      upts = V.map updateF ptsParticles
      uw = V.foldl' (\acc (WeighedDensityMatrix (w, _)) -> acc + w) 0 upts
      npts = V.map (over WeighedDensityMatrix (\(w, dm) -> (w / uw, dm))) upts
   in pts {ptsWeight = ptsWeight * uw, ptsParticles = npts}

-- | Helper function that generates random particles of each applicable rank.
initialiseParticleHierarchy :: Dim -> NumberOfParticles -> IO ParticleHierarchy
initialiseParticleHierarchy d n = V.generateM d (\r -> genParticles d (r + 1) n)

-- | Given a measurement result, update all particles, then normalise resulting
-- hierarchical weights to sum to unity.
updateParticleHierarchy ::
     PureStateVector -> ParticleHierarchy -> ParticleHierarchy
updateParticleHierarchy sv ph =
  let uph = V.map (updateParticles sv) ph
      wgts = V.map ptsWeight uph
      nwgts = V.map (/ V.sum wgts) wgts
   in V.zipWith (\x w -> x {ptsWeight = w}) uph nwgts

-- | Summarise the whole particle hierarchy into one mean Bayesian estimate.
getMixedEstimate :: ParticleHierarchy -> DensityMatrix
getMixedEstimate ph =
  let rankEstimates = V.map reduceParticlesToMean ph
      WeighedDensityMatrix (_, result) = V.foldl1' (<+>) rankEstimates
   in result

-- | Calculate the effective sample size of particles (weights don’t
-- necessarily have to be normalised).
effectiveSize :: Particles -> Double
effectiveSize Particles {..} =
  let ss = V.sum . V.map ((^ (2 :: Int)) . fst . getWDM) $ ptsParticles
      wa = V.foldl' (flip ((+) . fst . getWDM)) 0 ptsParticles
   in wa ^ (2 :: Int) / ss

-- | Nudges a particle by mixing the state together with some randomly
-- generated pure state. Relative weight of the random component determines how
-- “close” a nudged particle is to the original one.
nudgeParticle ::
     Dim
  -> Weight -- ^ Relative weight (from 0 to 1) of random component
  -> WeighedDensityMatrix
  -> IO WeighedDensityMatrix
nudgeParticle dim weightFraction (WeighedDensityMatrix (w, dm)) = do
  DensityMatrix nudgeDM <- svToDM <$> genPureSV dim
  let dmw = LA.scale (1 - (weightFraction :+ 0)) (getDensityMatrix dm)
      dmwn = LA.scale (weightFraction :+ 0) nudgeDM
  return $ WeighedDensityMatrix (w, DensityMatrix $ dmw + dmwn)

-- | Calculates values of empirical distribution function at data points.
ecdf :: V.Vector WeighedDensityMatrix -> V.Vector Double
ecdf = V.postscanl' (+) 0 . V.map (fst . getWDM)

-- | /O(log n)/ Given a non-empty sorted vector (typically an empirical cdf
-- evaluated at data points returned by ecdf) and a real number return the
-- (0-based) index of the least element of vector which is greater or equal to
-- the given real number (or the index of the last element, in case there is no
-- element smaller than the argument).
icdf :: V.Vector Double -> Double -> Int
icdf cdf x =
  let tIdx = V.length cdf - 1
      go (lIdx, hIdx) =
        let mIdx =
              truncate $ ((fromIntegral lIdx :: Double) + fromIntegral hIdx) / 2
         in select
              (go (lIdx, mIdx))
              [ (lIdx == hIdx, lIdx)
              , (lIdx + 1 == hIdx, if' (cdf V.! lIdx > x) lIdx hIdx)
              , (cdf V.! mIdx < x, go (mIdx, hIdx))
              ]
   in select (go (0, tIdx)) [(x <= V.head cdf, 0), (x > V.last cdf, tIdx)]

-- | Multinomial resampling of particle vector, which equalises weights of
-- particles.
resampleMultinom :: MWC.GenIO -> Particles -> IO Particles
resampleMultinom gen pts@Particles {..} = do
  us <- MWC.uniformVector gen ptsNumber
  let cdf = ecdf ptsParticles
      idxs = V.map (icdf cdf) us
      w = 1 / fromIntegral ptsNumber
      pointR = over WeighedDensityMatrix (\(_, dm) -> (w, dm))
      resampled = V.map (ptsParticles V.!) idxs
      normed = V.map pointR resampled
  return pts {ptsParticles = normed}

mhmcStep ::
     MWC.GenIO
  -> Dim
  -> Double
  -> [PureStateVector]
  -> Particles
  -> IO (Double, Particles)
mhmcStep gen dim rw ms pts@Particles {..} = do
  let cr wdm wdm' =
        exp
          (fullDataLogLikelihood ms (snd . getWDM $ wdm') -
           fullDataLogLikelihood ms (snd . getWDM $ wdm))
  newParticles <-
    V.mapM (fmap (truncateRank ptsRank) . nudgeParticle dim rw) ptsParticles
  us <- V.replicateM ptsNumber (MWC.uniform gen :: IO Double)
  let crs = V.zipWith cr ptsParticles newParticles
      change = V.zipWith (<=) us crs
      accRate =
        (fromIntegral . V.length . V.filter id) change / fromIntegral ptsNumber
      rwdms = V.zipWith3 if' change newParticles ptsParticles
      final = pts {ptsParticles = rwdms}
  return (accRate, final)

resampleMHMC ::
     ResampleArgs
  -> DensityMatrix
  -> Double
  -> Int
  -> [PureStateVector]
  -> Particles
  -> IO Particles
resampleMHMC ra@ResampleArgs {..} estimate wr iter mts pts = do
  (accRate, resampled) <- mhmcStep argGen argDim wr mts pts
  when (argOut == FullOutput) $
    printf
      "(Weight of new particle: %8.3g, MHMC acceptance rate: %8.3g)\n"
      wr
      accRate
  let (iter', wr') =
        select
          (iter + 1, wr)
          [ (accRate < 1e-2, (0, wr * 0.25))
          , (accRate < 1e-1, (0, wr * 0.5))
          , (iter < argMinIter, (iter + 1, wr))
          , (accRate < 0.33, (0, wr * 0.5))
          ]
  if iter > argMinIter
    then return resampled
    else resampleMHMC ra estimate wr' iter' mts resampled

-- | Arguments for the resampling function.
data ResampleArgs = ResampleArgs
  { argOut :: OutputVerb
  , argGen :: MWC.GenIO
  , argDim :: Dim
  , argMinIter :: MHMCiter
  }

-- | Resample particles. First does one multinomial step that equalises the
-- weights, then performs MHMC iterations adaptively refining the proposal
-- distribution based on acceptance rate. 'argMinIter' iterations are performed
-- for proposal distributions with significant acceptance rates.
resample :: ResampleArgs -> [PureStateVector] -> Particles -> IO Particles
resample ra@ResampleArgs {..} mts pts@Particles {..} = do
  let estimate = getMixedEstimate . V.singleton $ pts
      nudgeW = 0.95
  when (argOut == FullOutput) $ do
    putStrLn ""
    putStrLn $ "resampling rank " ++ show ptsRank
    putStrLn ""
  rm <- resampleMultinom argGen pts
  resampleMHMC ra estimate nudgeW 0 mts rm


{-# LANGUAGE OverloadedStrings #-}

module Main (main) where

import Types
import Parsing

-- import qualified Data.Text as T
import Data.Array
import Data.Ratio
import Data.Complex

import Test.Hspec
-- import System.IO.Unsafe

main :: IO ()
main = do
  hspec $ describe "src/Parser.hs" $ do

    it "Atom" $
      readExpr "bb-8?" `shouldBe` Right (Atom "bb-8?")

    it "Num Negative" $
      readExpr "-2187" `shouldBe` Right (Number (-2187))

    it "Num Positive" $
      readExpr "112233" `shouldBe` Right (Number 112233)

    it "Num Positive with Sign" $
      readExpr "+12345" `shouldBe` Right (Number 12345)

    it "Float Negative" $
      readExpr "-2187.1" `shouldBe` Right (Float (-2187.1))

    it "Float Positive" $
      readExpr "112233.1" `shouldBe` Right (Float 112233.1)

    it "Float Positive with Sign" $
      readExpr "+12345.3" `shouldBe` Right (Float 12345.3)

    it "Ratio Negative" $
      readExpr "-2187/3" `shouldBe` Right (Ratio (-2187 % 3))

    it "Ratio Positive" $
      readExpr "187/3" `shouldBe` Right (Ratio (187 % 3))

    it "Ratio Positive with Sign" $
      readExpr "+2187/3" `shouldBe` Right (Ratio (2187 % 3))

    it "Float Positive" $
      readExpr "112233.1" `shouldBe` Right (Float 112233.1)

    it "Float Positive with Sign" $
      readExpr "+12345.3" `shouldBe` Right (Float 12345.3)

    it "Complex Num with Sign" $
      readExpr "12345+3i" `shouldBe` Right (Complex (12345 :+ 3))

    it "Complex Num Positive with Sign" $
      readExpr "+12345+3i" `shouldBe` Right (Complex (12345 :+ 3))

    it "Complex Num Negative" $
      readExpr "-12345+3i" `shouldBe` Right (Complex ((-12345) :+ 3))

    it "Complex Num Negative" $
      readExpr "-12345+(-3)i" `shouldBe` Right (Complex ((-12345) :+ (-3)))

    it "String" $
      readExpr "\"Gen L Organa\"" `shouldBe` Right (String "Gen L Organa")

    it "Bool True" $
      readExpr "#t" `shouldBe` Right (Bool True)

    it "Bool False" $
      readExpr "#f" `shouldBe` Right (Bool False)

    it "S-Expr: homogenous list" $
      readExpr "(2 1 87)" `shouldBe`
      Right (List [Number 2, Number 1,Number 87])

    it "S-Expr: homogenous list quoted" $
      readExpr "'(2 1 87)" `shouldBe`
      Right (List [Atom "quote",List [Number 2, Number 1,Number 87]])

    it "S-Expr: heterogenous list" $
      readExpr "(stromTrooper \"Fn\" 2 1 87)" `shouldBe`
      Right (List [Atom "stromTrooper", String "Fn", Number 2, Number 1,Number 87])

    it "S-Expr: heterogenous list quoted" $
      readExpr "'(stromTrooper \"Fn\" 2 1 87)" `shouldBe`
      Right (List [Atom "quote", List [Atom "stromTrooper", String "Fn", Number 2, Number 1,Number 87]])

    it "S-Expr: vector" $
      readExpr "#(stromTrooper \"Fn\" 2 1 87)" `shouldBe`
      Right (Vector (listArray (0, 4) [Atom "stromTrooper", String "Fn", Number 2, Number 1,Number 87]))

    it "S-Expr: char space" $
      readExpr "(#\\space)" `shouldBe` Right (List [Char ' '])

    it "S-Expr: char newline" $
      readExpr "(#\\newline)" `shouldBe` Right (List [Char '\n'])

    it "S-Expr: char letter" $
      readExpr "(#\\y)" `shouldBe` Right (List [Char 'y'])

    it "S-Expr: single negative" $
      readExpr "(-42)" `shouldBe` Right (List [Number (-42)])

    it "S-Expr: (- num)" $
      readExpr "(- 42)" `shouldBe` Right (List [Atom "-", Number 42])

    it "S-Expr: prim call: numbers" $
      readExpr "(+ 1 2)" `shouldBe`
      Right (List [Atom "+", Number 1, Number 2])

    it "S-Expr: prim call: neg nums" $
      readExpr "(- -42 -42)" `shouldBe`
      Right (List [Atom "-", Number (-42), Number (-42)])

    it "S-Expr: prim call: atoms" $
      readExpr "(- rogue squadron)" `shouldBe`
      Right (List [Atom "-", Atom "rogue", Atom "squadron"])

    it "S-Expr: nested list" $
      readExpr "(lambda (x x) (+ x x))" `shouldBe`
      Right (List [Atom "lambda", List [Atom "x", Atom "x"], List [Atom "+", Atom "x", Atom "x"]])

    it "Comment: end-of/single line" $
      readExpr ";skip\nartoodetoo ;extra will throw\n;skip" `shouldBe` Right (Atom "artoodetoo")

    it "Comment: multi-line comment in single line" $
      readExpr "(c3 #|Hello how are you|# po)" `shouldBe` Right (List [Atom "c3",Atom "po"])

    it "Comment: multi-line line" $
      readExpr "#|Han\nShot\nFirst\n|#(c3 #| these are not the droids you're looking for|# po)\n #|Jar Jar Binks =?= Sith Lord |#" `shouldBe` Right (List [Atom "c3",Atom "po"])

    it "S-Expr: braces/brackets" $
      readExpr "(lambda {x x} [+ x x])" `shouldBe`
      Right (List [Atom "lambda", List [Atom "x", Atom "x"], List [Atom "+", Atom "x", Atom "x"]])

    it "S-Expr: bars" $
      readExpr "(lambda |x x| [+ x x])" `shouldBe`
      Right (List [Atom "lambda", List [Atom "x", Atom "x"], List [Atom "+", Atom "x", Atom "x"]])

  -- hspec $ describe "src/Eval.hs" $ do
  --   wStd "test/add.scm"              $ Number 3
  --   wStd "test/if_alt.scm"           $ Number 2
  --   wStd "test/let.scm"              $ Number 321
  --   wStd "test/eval_boolean.scm"     $ Bool True
  --   wStd "test/eval_boolean_ops.scm" $ Bool True
  --   wStd "test/eval_lambda.scm"      $ Number 5
  --   wStd "test/test_quote.scm"       $ List [Atom "xNotFound", Atom "yNotFound"]
  --   wStd "test/test_car.scm"         $ Number 1
  --   wStd "test/test_cdr.scm"         $ List [Number 2]
  --   wStd "test/test_cadadr.scm"      $ Number 42
  --   wStd "test/test_gt.scm"          $ List [ Bool True, Bool False]
  --   wStd "test/test_scope1.scm"      $ Number 413281
  --   wStd "test/test_scope2.scm"      $ Number 11
  --   wStd "test/test_args.scm"        $ Number 105065
  --   wStd "test/test_fold.scm"        $ Number 42
  --   runExpr Nothing "test/define.scm"        $ Number 4
  --   runExpr Nothing "test/define_order.scm"  $ Number 42
  --   runExpr Nothing "test/define_lambda.scm" $ String "smalltalk"
  --   runExpr Nothing "test/test_evalargs.scm" $ Number 1558
    --runExpr Nothing "test/test_fix.scm"      $ Number 3628800
    --runExpr Nothing "test/test_fix2.scm"      $ Number 5040

-- helper functions
-- run file w/ stdlib
-- wStd :: T.Text -> LispVal -> SpecWith ()
-- wStd = runExpr (Just "test/stdlib_mod.scm")

-- -- run expr w/o stdLib
-- tExpr :: T.Text -> T.Text -> LispVal -> SpecWith ()
-- tExpr note expr val =
--     it (T.unpack note) $ evalVal `shouldBe` val
--     where evalVal = (unsafePerformIO $ runASTinEnv basicEnv $ fileToEvalForm "" expr)

-- runExpr :: Maybe T.Text -> T.Text -> LispVal -> SpecWith ()
-- runExpr  std file val =
--     it (T.unpack file) $ evalVal  `shouldBe` val
--     where evalVal = unsafePerformIO $ evalTextTest std file

-- evalTextTest :: Maybe T.Text -> T.Text -> IO LispVal --REPL
-- evalTextTest (Just stdlib) file= do
--   stdlib <- getFileContents $ T.unpack  stdlib
--   f      <- getFileContents $ T.unpack file
--   runASTinEnv basicEnv $ textToEvalForm stdlib  f

-- evalTextTest Nothing file = do
--   f <- getFileContents $ T.unpack file
--   runASTinEnv basicEnv $ fileToEvalForm (T.unpack file) f

-- -- run text expr w/ file
-- tExprStd :: T.Text -> T.Text -> LispVal -> SpecWith ()
-- tExprStd note  expr  val =
--     it (T.unpack note ) $ evalVal `shouldBe` val
--     where evalVal = unsafePerformIO $ evalExprTest expr

-- evalExprTest ::  T.Text -> IO LispVal --REPL
-- evalExprTest expr = do
--   stdlib <- getFileContents $ T.unpack  "test/stdlib_mod.scm"
--   runASTinEnv basicEnv $ textToEvalForm stdlib  expr


module Numeric.Sundials.ODEOpts where

import           Data.Word (Word32)
import qualified Data.Vector.Storable as VS

import           Numeric.LinearAlgebra.HMatrix (Vector, Matrix)

type Jacobian = Double -> Vector Double -> Matrix Double

data ODEOpts = ODEOpts {
    maxNumSteps :: Word32
  , minStep     :: Double
  , relTol      :: Double
  , absTols     :: VS.Vector Double
  , initStep    :: Maybe Double
  , maxFail     :: Word32
  } deriving (Read, Show, Eq, Ord)

data SundialsDiagnostics = SundialsDiagnostics {
    aRKodeGetNumSteps               :: Int
  , aRKodeGetNumStepAttempts        :: Int
  , aRKodeGetNumRhsEvals_fe         :: Int
  , aRKodeGetNumRhsEvals_fi         :: Int
  , aRKodeGetNumLinSolvSetups       :: Int
  , aRKodeGetNumErrTestFails        :: Int
  , aRKodeGetNumNonlinSolvIters     :: Int
  , aRKodeGetNumNonlinSolvConvFails :: Int
  , aRKDlsGetNumJacEvals            :: Int
  , aRKDlsGetNumRhsEvals            :: Int
  } deriving Show


{-# LANGUAGE FlexibleContexts, FlexibleInstances, UndecidableInstances #-}

-- Pure Haskell implementation of standard BFGS algorithm with cubic
-- line search, with the BFGS step exposed to provide access to
-- intermediate Hessian estimates.

-- Author: Ian Ross, for OpenBrain Ltd

module Fuml.Optimisation.BFGS
       ( BFGSOpts (..)
       , LineSearchOpts (..)
       , bfgs, bfgsWith
       , bfgsInit
       , bfgsStep, bfgsStepWith )
where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import Data.Default
--import Debug.Trace

-- Type synonyms mostly just for documentation purposes...
type Point = Vector Double
type Gradient = Vector Double
type Direction = Vector Double
type Fn = Point -> Double
type GradFn = Point -> Vector Double
type Hessian = Matrix Double

--------------------------------------------------------------------------------
--
--  NaN CHECKING
--

class Nanable a where
  hasnan :: a -> Bool

instance Nanable Double where
  hasnan = isNaN

instance (Nanable a, Container c a) => Nanable (c a) where
  hasnan = not . null . (find hasnan)

--------------------------------------------------------------------------------
--
--  BFGS
--

-- Options for BFGS solver: point tolerance, gradient tolerance,
-- maximum iterations.
--
data BFGSOpts = BFGSOpts { ptol :: Double
                         , gtol :: Double
                         , maxiters :: Int } deriving Show

instance Default BFGSOpts where
  def = BFGSOpts 1.0E-7 1.0E-7 200

-- BFGS solver data to carry between steps: current point, function
-- value at point, gradient at point, current direction, current
-- Hessian estimate, maximum line search step.
--
data BFGS = BFGS { p :: Point
                 , fp :: Double
                 , g :: Gradient
                 , xi :: Direction
                 , h :: Matrix Double
                 , stpmax :: Double } deriving Show

-- Main solver interface with default options.
--
bfgs :: Fn -> GradFn -> Point -> Either String (Point, Hessian)
bfgs f df p0 = bfgsWith def f df p0 (ident $ size p0)

{- Collection solver state into an infinite list: useful as "take n $
-- bfgsCollect f df b0"
--
bfgsCollect :: Fn -> GradFn -> BFGS -> [BFGS]
bfgsCollect f df b0 = map snd $ iterate step (False,b0)
  where step (cvg,b) = if cvg then b else bfgsStep f df b

-}
-- Utility function to set up initial BFGS state for bfgsCollect.
--
bfgsInit :: Fn -> GradFn -> Point -> Either String BFGS
bfgsInit f df p0 = case (hasnan f0, hasnan g0) of
  (False, False) -> Right $ BFGS p0 f0 g0 (-g0) (ident n) (maxStep p0)
  _ -> Left $ nanMsg p0 (Just f0) (Just g0)
  where n = size p0 ; f0 = f p0 ; g0 = df p0

-- Main iteration routine: sets up initial BFGS state, then steps
-- until converged or maximum iterations exceeded.
--
bfgsWith :: BFGSOpts -> Fn -> GradFn -> Point -> Hessian -> Either String (Point, Hessian)
bfgsWith opt@(BFGSOpts _ _ maxiters') f df p0 h0 =
  case (hasnan f0, hasnan g0) of
    (False, False) -> go 0 b0
    errs -> Left $ nanMsg p0 (Just f0) (Just g0)
  where go iters b =
          if iters > maxiters'
          then Left "maximum iterations exceeded in bfgs"
          else case bfgsStepWith opt f df b of
            Left err -> Left err
            Right (True, b') -> Right (p b', h b')
            Right (False, b') -> go (iters+1) $ {- trace (show $ (iters, fp b')) -} b'
        f0 = f p0 ; g0 = df p0
        b0 = BFGS p0 f0 g0 (-g0) h0 (maxStep p0)

-- Do a BFGS step with default parameters.
--
bfgsStep :: Fn -> GradFn -> BFGS -> Either String (Bool, BFGS)
bfgsStep = bfgsStepWith def

-- Main BFGS step routine.  This is a more or less verbatim
-- translation of the description in Section 10.7 of Numerical Recipes
-- in C, 2nd ed.
--
bfgsStepWith :: BFGSOpts -> Fn -> GradFn -> BFGS -> Either String (Bool, BFGS)
bfgsStepWith (BFGSOpts ptol' gtol' _) f df (BFGS p fp g xi h stpmax) =
  case lineSearch f p fp g xi stpmax of
    Left err -> Left err
    Right (pn, fpn) ->
      if hasnan gn
      then Left $ nanMsg pn Nothing (Just gn)
      else if cvg
           then Right (True, BFGS pn fpn gn xi h stpmax)
           else Right (False, BFGS pn fpn gn xin hn stpmax)
      where gn = df pn ; dp = pn - p ; dg = gn - g ; hdg = h #> dg
            dpdg = dp `dot` dg ; dghdg = dg `dot` hdg
            hn = h + ((dpdg + dghdg) / dpdg^2) `scale` (dp `outer` dp) -
                 (1/dpdg) `scale` (h <> (dg `outer` dp) + (dp `outer` dg) <> h)
            xin = -hn #> gn
            cvg = maxabsratio dp p < ptol' || maxabsratio' (fpn `max` 1) gn p < gtol'

-- Generate error messages for NaN production in function and gradient
-- calculations.
--
nanMsg :: Point -> Maybe Double -> Maybe Gradient -> String
nanMsg p' fval grad = "NaNs produced: p = " ++ show p' ++
                     maybe "" (("  fval = " ++) . show) fval ++
                     maybe "" (("  grad = " ++) . show) grad

--------------------------------------------------------------------------------
--
--  LINE SEARCH
--

-- Options for line search routine: point tolerance, "acceptable
-- decrease" parameter.
--
data LineSearchOpts = LineSearchOpts { xtol :: Double
                                     , alpha :: Double } deriving Show

instance Default LineSearchOpts where
  def = LineSearchOpts 1.0E-7 1.0E-4

-- Maximum line search step length.
--
maxStep :: Point -> Double
maxStep p0 = (100 * (norm_2 p0 `max` n))
  where n = fromIntegral (size p0)

-- Line search with default parameters.
--
lineSearch :: Fn -> Point -> Double -> Gradient -> Direction ->
              Double -> Either String (Point, Double)
lineSearch = lineSearchWith def

-- Main line search routine.  This is kind of nasty to translate into
-- functional form because of the switching about between the
-- quadratic and cubic approximations.  It works, but it could be
-- prettier.
--
lineSearchWith :: LineSearchOpts -> Fn -> Point -> Double -> Gradient ->
                  Direction -> Double -> Either String (Point, Double)
lineSearchWith (LineSearchOpts xtol alpha) func xold fold g pin stpmax =
  go 1.0 Nothing
  where p = if pinnorm > stpmax then (stpmax/pinnorm) `scale` pin else pin
        pinnorm = norm_2 pin
        slope = g `dot` p
        lammin = xtol / (maxabsratio p xold)

        go :: Double -> Maybe (Double,Double) -> Either String (Point,Double)
        go lam pass = if hasnan fnew
                      then Left $ nanMsg xnew (Just fnew) Nothing
                      else case check xnew fnew of
                        Just xandf -> Right xandf
                        Nothing ->
                          case pass of
                            -- First time.
                            Nothing -> go (lambound $ quadlam fnew) $ Just (lam,fnew)
                            -- Subsequent times.
                            Just val2 -> case cubiclam fnew val2 of
                              Right newlam -> go (lambound newlam) $ Just (lam,fnew)
                              Left err -> Left err
          where xnew = xold + lam `scale` p
                fnew = func xnew

                -- Check for convergence or a "sufficiently large" step.
                check :: Vector Double -> Double -> Maybe (Vector Double,Double)
                check x f =
                  if lam < lammin then Just (xold,fold)
                  else if f <= fold + alpha * lam * slope
                       then Just (x,f)
                       else Nothing

                -- Keep step length within bounds.
                lambound lam' = max (0.1 * lam) (min lam' (0.5 * lam))

                -- Quadratic and cubic approximations to better step
                -- value.
                quadlam fnew = -slope / (2 * (fnew - fold - slope))
                cubiclam fnew (lam2,f2) =
                  if a == 0
                  then Right (-slope / (2 * b))
                  else if disc < 0
                       then Left "Roundoff problem in lineSearch"
                       else Right $ (-b + sqrt disc) / (3 * a)
                    where rhs1 = fnew - fold - lam * slope
                          rhs2 = f2 - fold - lam2 * slope
                          a = (rhs1 / lam^2 - rhs2 / lam2^2) / (lam - lam2)
                          b = (-lam2 * rhs1 / lam^2 + lam * rhs2 / lam2^2) /
                              (lam - lam2)
                          disc = b^2 - 3 * a * slope

-- Utility functions for ratio testing.
--
absratio :: Double -> Double -> Double
absratio n d = abs n / (abs d `max` 1)

absratio' :: Double -> Double -> Double -> Double
absratio' scale n d = abs n / (abs d `max` 1) / scale

maxabsratio :: Vector Double -> Vector Double -> Double
maxabsratio n d = maxElement $ zipVectorWith absratio n d

maxabsratio' :: Double -> Vector Double -> Vector Double -> Double
maxabsratio' scale n d = maxElement $ zipVectorWith (absratio' scale) n d

--------------------------------------------------------------------------------
--
--  TEST FUNCTIONS
--

-- Simple test function.  Minimum at (3,4):
--
-- *BFGS> bfgs tstf1 tstgrad1 (fromList [-10,-10])
-- Right (fromList [3.0,4.0])
--

{-
tstf1 :: Fn
tstf1 p = (x-3)^2 + (y-4)^2
  where [x,y] = toList p

tstgrad1 :: GradFn
tstgrad1 p = fromList [2*(x-3),2*(y-4)]
  where [x,y] = toList p

-- Rosenbrock's function.  Minimum at (1,1):
--
-- *BFGS> bfgs tstf2 tstgrad2 (fromList [-10,-10])
-- Right (fromList [0.9999999992103538,0.9999999985219549])

tstf2 :: Fn
tstf2 p = (1-x)^2 + 100*(y-x^2)^2
  where [x,y] = toList p

tstgrad2 :: GradFn
tstgrad2 p = fromList [2*(x-1) - 400*x*(y-x^2), 200*(y-x^2)]
  where [x,y] = toList p

-- Test function from Numerical Recipes.  Minimum at (-2.0,0.89442719):
--
-- *BFGS> bfgs tstfnr tstgradnr nrp0
-- Right (fromList [-1.9999999564447526,0.8944271925873616])

tstfnr :: Fn
tstfnr p = 10*(y^2*(3-x)-x^2*(3+x))^2+(2+x)^2/(1+(2+x)^2)
  where [x,y] = toList p

tstgradnr :: GradFn
tstgradnr p = fromList [20*(y^2*x3m-x^2*x3p)*(-y^2-6*x-3*x^2)+
                        2*x2p/(1+x2p^2)-2*x2p^3/(1+x2p^2)^2,
                        40*(y^2*x3m-x^2*x3p)*y*x3m]
  where [x,y] = toList p
        x3m = 3 - x
        x3p = 3 + x
        x2p = 2 + x

nrp0 :: Point
nrp0 = fromList [0.1,4.2]

-- Test function to check NaN handling.
--
-- *BFGS> bfgs nantstf nantstgrad (fromList [-10,-10])
-- Left "function application returned NaN"

nantstf :: Fn
nantstf p = log x + (x-3)^2 + (y-4)^2
  where [x,y] = toList p

nantstgrad :: GradFn
nantstgrad p = fromList [1/x+2*(x-3),2*(y-4)]
  where [x,y] = toList p
-}


{-# LANGUAGE BangPatterns           #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE LambdaCase             #-}
{-# LANGUAGE PolyKinds              #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TupleSections          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilyDependencies #-}

module Backprop.Learn.Model (
    module M, Backprop(..)
  -- * Running and Grad
  , runModel, runModelStoch, runModelStateless, runModelStochStateless
  , gradModel, gradModelStoch
  -- * Work with parameters
  , initParam, initParamNormal
  , encodeParam, decodeParam, decodeParamOrFail, saveParam, loadParam, loadParamOrFail
  -- * Iterated runners
  , iterateModel, iterateModelM, iterateModelStoch
  , scanModel, scanModelStoch
  -- * No final state
  , iterateModel_, iterateModelM_, iterateModelStoch_
  , scanModel_, scanModelStoch_
  -- * "Prime" runners
  , primeModel, primeModelStoch, selfPrime, selfPrimeM
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Model.Combinator  as M
import           Backprop.Learn.Model.Function    as M
import           Backprop.Learn.Model.Neural      as M
import           Backprop.Learn.Model.Neural.LSTM as M
import           Backprop.Learn.Model.Parameter   as M
import           Backprop.Learn.Model.Regression  as M
import           Backprop.Learn.Model.State       as M
import           Backprop.Learn.Model.Stochastic  as M
import           Backprop.Learn.Model.Types       as M
import           Control.Monad.Primitive
import           Control.Monad.ST
import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Foldable
import           Data.Functor.Identity
import           Data.Type.Functor.Product
import           Data.Type.Tuple
import           Data.Word
import           Numeric.Backprop
import           Statistics.Distribution
import qualified Data.Binary                      as Bi
import qualified Data.ByteString.Lazy             as BSL
import qualified Data.Vector.Unboxed              as VU
import qualified System.Random.MWC                as MWC

-- TODO: this can be more efficient by breaking out into separate functions
runModel
    :: forall p s a b. (AllConstrainedProd Backprop s, Backprop b)
    => Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> (b, TMaybe s)
runModel f mp x ms = evalBP0 go
  where
    go :: forall z. Reifies z W => BVar z (b, TMaybe s)
    go = case ms' of
        PNothing    -> T2 y $ auto PNothing
        PJust s' -> T2 y $ isoVar (PJust . TF) (getTF . fromPJust) s'
      where
        y :: BVar z b
        ms' :: PMaybe (BVar z) s
        (y, ms') = runLearn f (mapProd (auto . getTF) mp)
                              (auto x)
                              (mapProd (auto . getTF) ms)

-- TODO: this can be more efficient by breaking out into separate functions
runModelStoch
    :: forall p s a b m. (AllConstrainedProd Backprop s, Backprop b, PrimMonad m)
    => Model p s a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> a
    -> TMaybe s
    -> m (b, TMaybe s)
runModelStoch f g mp x ms = do
    -- TODO: is this the best way to handle this?
    seed <- MWC.uniformVector @_ @Word32 @VU.Vector g 2
    pure $ evalBP0 (runST (go seed))
  where
    go  :: forall q z. Reifies z W
        => VU.Vector Word32
        -> ST q (BVar z (b, TMaybe s))
    go seed = do
      g' <- MWC.initialize seed
      (y :: BVar z b, ms') <- runLearnStoch f g'
        (mapProd (auto . getTF) mp)
        (auto x)
        (mapProd (auto . getTF) ms)
      pure $ case ms' of
        PNothing    -> T2 y $ auto PNothing
        PJust s' -> T2 y $ isoVar (PJust . TF) (getTF . fromPJust) s'

runModelStateless
    :: Model p 'Nothing a b
    -> TMaybe p
    -> a
    -> b
runModelStateless f = \case
    PNothing       -> evalBP  (runLearnStateless f PNothing  )
    PJust (TF p) -> evalBP2 (runLearnStateless f . PJust) p

runModelStochStateless
    :: PrimMonad m
    => Model p 'Nothing a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> a
    -> m b
runModelStochStateless f g mp x = do
    seed <- MWC.uniformVector @_ @Word32 @VU.Vector g 2
    pure $ case mp of
      PNothing       -> evalBP  (\x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless f g' PNothing x'
        ) x
      PJust (TF p) -> evalBP2 (\p' x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless f g' (PJust p') x'
        ) p x

gradModel
    :: (Backprop a, Backprop b, AllConstrainedProd Backprop p)
    => Model p 'Nothing a b
    -> TMaybe p
    -> a
    -> (TMaybe p, a)
gradModel f = \case
    PNothing       ->       (PNothing,)    . gradBP  (runLearnStateless f   PNothing)
    PJust (TF p) -> first (PJust . TF) . gradBP2 (runLearnStateless f . PJust) p

gradModelStoch
    :: (Backprop a, Backprop b, AllConstrainedProd Backprop p, PrimMonad m)
    => Model p 'Nothing a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> a
    -> m (TMaybe p, a)
gradModelStoch f g mp x = do
    seed <- MWC.uniformVector @_ @Word32 @VU.Vector g 2
    pure $ case mp of
      PNothing       ->       (PNothing,)    $ gradBP  (\x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless f g' PNothing x'
        ) x
      PJust (TF p) -> first (PJust . TF) $ gradBP2 (\p' x' -> runST $ do
          g' <- MWC.initialize seed
          runLearnStochStateless f g' (PJust p') x'
        ) p x

iterateModel
    :: (Backprop b, AllConstrainedProd Backprop s)
    => (b -> a)         -- ^ loop
    -> Int              -- ^ num times
    -> Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> ([b], TMaybe s)
iterateModel l n f p x = runIdentity . iterateModelM (Identity . l) n f p x

iterateModel_
    :: (Backprop b, AllConstrainedProd Backprop s)
    => (b -> a)         -- ^ loop
    -> Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> [b]
iterateModel_ l f p = go
  where
    go !x !s = y : go (l y) s'
      where
        (y, s') = runModel f p x s

selfPrime
    :: (Backprop b, AllConstrainedProd Backprop s)
    => (b -> a)         -- ^ loop
    -> Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> [TMaybe s]
selfPrime l f p = go
  where
    go !x !s = s' : go (l y) s'
      where
        (y, s') = runModel f p x s

iterateModelM
    :: (Backprop b, AllConstrainedProd Backprop s, Monad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> m ([b], TMaybe s)
iterateModelM l n f p = go 0
  where
    go !i !x !s
      | i <= n    = do
          let (y, s') = runModel f p x s
          (ys, s'') <- flip (go (i + 1)) s' =<< l y
          pure (y : ys, s'')
      | otherwise = pure ([], s)

iterateModelM_
    :: (Backprop b, AllConstrainedProd Backprop s, Monad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> m [b]
iterateModelM_ l n f p x = fmap fst . iterateModelM l n f p x

selfPrimeM
    :: (Backprop b, AllConstrainedProd Backprop s, Monad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> Model p s a b
    -> TMaybe p
    -> a
    -> TMaybe s
    -> m (TMaybe s)
selfPrimeM l n f p x = fmap snd . iterateModelM l n f p x

iterateModelStoch
    :: (Backprop b, AllConstrainedProd Backprop s, PrimMonad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> Model p s a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> a
    -> TMaybe s
    -> m ([b], TMaybe s)
iterateModelStoch l n f g p = go 0
  where
    go !i !x !s
      | i <= n    = do
          (y , s' ) <- runModelStoch f g p x s
          (ys, s'') <- flip (go (i + 1)) s' =<< l y
          pure (y : ys, s'')
      | otherwise = pure ([], s)

iterateModelStoch_
    :: (Backprop b, AllConstrainedProd Backprop s, PrimMonad m)
    => (b -> m a)           -- ^ loop
    -> Int                  -- ^ num times
    -> Model p s a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> a
    -> TMaybe s
    -> m [b]
iterateModelStoch_ l n f g p x = fmap fst . iterateModelStoch l n f g p x

scanModel
    :: (Traversable t, Backprop b, AllConstrainedProd Backprop s)
    => Model p s a b
    -> TMaybe p
    -> t a
    -> TMaybe s
    -> (t b, TMaybe s)
scanModel f p = runState . traverse (state . runModel f p)

scanModel_
    :: (Traversable t, Backprop b, AllConstrainedProd Backprop s)
    => Model p s a b
    -> TMaybe p
    -> t a
    -> TMaybe s
    -> t b
scanModel_ f p xs = fst . scanModel f p xs

primeModel
    :: (Foldable t, Backprop b, AllConstrainedProd Backprop s)
    => Model p s a b
    -> TMaybe p
    -> t a
    -> TMaybe s
    -> TMaybe s
primeModel f p = execState . traverse_ (state . runModel f p)

scanModelStoch
    :: (Traversable t, Backprop b, AllConstrainedProd Backprop s, PrimMonad m)
    => Model p s a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> t a
    -> TMaybe s
    -> m (t b, TMaybe s)
scanModelStoch f g p = runStateT . traverse (StateT . runModelStoch f g p)

scanModelStoch_
    :: (Traversable t, Backprop b, AllConstrainedProd Backprop s, PrimMonad m)
    => Model p s a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> t a
    -> TMaybe s
    -> m (t b)
scanModelStoch_ f g p xs = fmap fst . scanModelStoch f g p xs

primeModelStoch
    :: (Foldable t, Backprop b, AllConstrainedProd Backprop s, PrimMonad m)
    => Model p s a b
    -> MWC.Gen (PrimState m)
    -> TMaybe p
    -> t a
    -> TMaybe s
    -> m (TMaybe s)
primeModelStoch f g p = execStateT . traverse_ (StateT . runModelStoch f g p)

initParam
    :: (Initialize p, ContGen d, PrimMonad m)
    => model ('Just p) s a b                    -- ^ ignored
    -> d
    -> MWC.Gen (PrimState m)
    -> m p
initParam _ = initialize

initParamNormal
    :: (Initialize p, PrimMonad m)
    => model ('Just p) s a b                    -- ^ ignored
    -> Double
    -> MWC.Gen (PrimState m)
    -> m p
initParamNormal _ = initializeNormal

encodeParam
    :: Bi.Binary p
    => model ('Just p) s a b                    -- ^ ignored
    -> p
    -> BSL.ByteString
encodeParam _ = Bi.encode

decodeParam
    :: Bi.Binary p
    => model ('Just p) s a b                    -- ^ ignored
    -> BSL.ByteString
    -> p
decodeParam _ = Bi.decode

decodeParamOrFail
    :: Bi.Binary p
    => model ('Just p) s a b                    -- ^ ignored
    -> BSL.ByteString
    -> Either String p
decodeParamOrFail _ = bimap thrd thrd . Bi.decodeOrFail

saveParam
    :: Bi.Binary p
    => model ('Just p) s a b                    -- ^ ignored
    -> FilePath
    -> p
    -> IO ()
saveParam p fp = BSL.writeFile fp . encodeParam p

loadParam
    :: Bi.Binary p
    => model ('Just p) s a b                    -- ^ ignored
    -> FilePath
    -> IO p
loadParam p fp = decodeParam p <$> BSL.readFile fp

loadParamOrFail
    :: Bi.Binary p
    => model ('Just p) s a b                    -- ^ ignored
    -> FilePath
    -> IO (Either String p)
loadParamOrFail p fp = decodeParamOrFail p <$> BSL.readFile fp

thrd :: (a,b,c) -> c
thrd (_,_,z) = z


{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Data.List
import Data.Bits
import qualified Data.Set as S
import qualified Data.Map.Strict as M
import qualified Data.Map.Lazy as LM
import Text.Printf
import Numeric.LinearAlgebra hiding (find, (<>))
import Data.Functor
import Options.Applicative
import Control.Monad
import Data.Monoid
import Data.Function
import Debug.Trace

-- no warning about Debug.Trace please
_trace = trace

type Width = Int

type Mealy a b = [(a, ((a, [b]), (a, [b])))]

data L = S | M deriving (Eq, Ord, Enum, Bounded)

instance Show L where
    showsPrec _ S = ('S':)
    showsPrec _ M = ('M':)
    showList []   = ('·':)
    showList xs   = foldr ((.) . shows) id xs

buildMealyRTL :: Width -> Mealy Int L
buildMealyRTL 1 = [ (0, ((0, [S]), (0, [M]))) ]
buildMealyRTL w =
    [ (0, ((0, [S]), (1, [M]))) ] ++
    [ (i, ((j, [S]), (j, [S]))) | i <- [1..w-1] , let j = (i+1) `mod` w]

buildMealy :: Width -> Mealy Int L
buildMealy 1 = [ (0, ((0, [S]), (0,[M]))) ]
buildMealy w = [ (n, edges n) | n <- [0..2^(w-1)-1] ]
  where
    -- The root
    edges 0 = ( (0, [S]), (1, []) )
    -- The intermediate nodes (scanning until the end of the window)
    edges n | n < 2^(w-2) = ( (2*n, []), (2*n+1, []) )
    -- The last bit of the window, we can output the leak
    edges n = ( (0, output (2*n)), (0, output (2*n+1)) )

    output n = replicate (w - lastbit - 1) S ++ [M] ++ replicate lastbit S
      where Just lastbit = find (\i -> n `testBit` i) [0..]

-- | This function reduces the number of nodes by commoning up
-- equivalent nodes. Needs to be iterated.
reduceMealyOnce :: (Ord a, Ord b) => Mealy a b -> Mealy a b
reduceMealyOnce m
    = [ (n, ((f n0,s0), (f n1,s1)))
      | (n, ((n0,s0), (n1,s1))) <- m
      , f n == n
      ]
  where
    f = (M.!) $
        M.fromList
        [ (x, head dup)
        | dupSet <- M.elems $ M.fromListWith S.union [ (e,S.singleton n) | (n,e) <- m ]
        , let dup = S.toList dupSet
        , x <- dup
        ]

reduceMealy :: (Ord a, Ord b) => Mealy a b -> Mealy a b
reduceMealy m | length m == length m' = m
              | otherwise             = reduceMealy m'
  where m' = reduceMealyOnce m

mealy2dot :: (Show a, Show b) => Mealy a b -> String
mealy2dot nodes = unlines $
    [ "digraph {" ] ++
    [ printf "%s -> %s [label=\"0/%s\"];\n" (show n) (show n0) (show s0) ++
      printf "%s -> %s [label=\"1/%s\"];" (show n) (show n1) (show s1)
    | (n, ((n0, s0), (n1,s1))) <- nodes ] ++
    [ "}" ]

type Markov a = [(a, [a])]

buildMarkov :: (Ord a, Ord b) => (a,[b]) -> Mealy a b -> Markov (a,[b])
buildMarkov start mealy = go (S.singleton start) S.empty []
  where
    go todo done
        | S.null todo = id
        | (node,text) `S.member` done = go todo' done
        | otherwise = (((node,text), [n0', n1']):) .
            go todo'' (S.insert (node,text) done)

      where
      ((node,text), todo') = S.deleteFindMin todo
      Just ((n0, s0), (n1, s1)) = lookup node mealy
      n0' = (n0, tail text ++ s0)
      n1' = (n1, tail text ++ s1)
      todo'' = S.insert n0' (S.insert n1' todo')

markov2dot :: (a -> String) -> (a -> String) -> Markov a -> String
markov2dot col s nodes = unlines $
    [ "digraph {"
    -- , "layout = neato;"
    , "margin = 0;"
    , "epsilon = 0.0001;"
    , "splines=true;"
    -- , "edge [len = 0.8];"
    , "edge [len = 1.3];"
    , "node [width=0.2];"
    , "node [height=0.2];"
    , "node [style=filled];"
    ] ++
    [ printf "%s [ color=%s]" (s n) (col n)
    | (n, _) <- nodes ] ++
    [ printf "%s -> %s;\n" (s n) (s n')
    | (n, succs) <- nodes , n' <- succs] ++
    [ "}" ]

type Dist a = M.Map a Double

stationary :: Ord a => Markov a -> Dist a
stationary = stationaryCond (const True)

findMaxIndex f xs = snd (maximum (zip (map f xs) [0..]))

stationaryCond :: Ord a => (a -> Bool) -> Markov a -> Dist a
stationaryCond nonDead markov
    -- | not ok1 = error "First eigenvalue not 1"
    | not ok2 = error "Eigenvector does not consist of probabilities"
    -- | not ok3 = error "Eigenvector does not have eigenvalue 1"
    | otherwise = M.fromList $ zip (map fst markov) (toList evec)
  where
    n = length markov

    transMatrix :: Matrix Double
    transMatrix = (n >< n) $
        [ if nonDead n then sum [ p | n' <- e, n' == n0 ] else 0
        | (n,e) <- markov
        , (n0,_) <- markov
        , let p = 1/fromIntegral (length e)
        ]

    (evals, evecs) = eig (tr transMatrix)
    -- Just i = findIndex isOne (toList evals)
    i = findMaxIndex realPart (toList evals)
    eval = evals `atIndex` i
    evecC = toColumns evecs !! i
    evecNorm = scale (1/sum (toList evecC)) evecC
    evec = cmap realPart evecNorm

    isOne c = magnitude (c - 1) < 0.0000001
    isProb c = abs (imagPart c) < 0.0000001 && realPart c >= 0 && realPart c <= 1

    ok1 = magnitude (eval - 1) < 0.001
    ok2 = all isProb (toList evecNorm)
    ok3 = norm_2 (evec - (evec <# transMatrix)) < 0.001

prodMarkov :: Markov a -> Markov (a,a)
prodMarkov markov =
    [ ((n1,n2), (,) <$> e1 <*> e2)
    | (n1, e1) <- markov
    , (n2, e2) <- markov
    ]

filterMarkov :: (a -> Bool) -> Markov a -> Markov (Maybe a)
filterMarkov p markov =
    (Nothing, [Nothing]) : [ (Just n, map (justIf p) e) | (n, e) <- markov, p n ]

justIf p n = if p n then Just n else Nothing

-- Removes all edges from nodes where p does not hold, and all nodes
-- not reachable
deadEnd :: Ord a => (a -> Bool) -> a -> Markov a -> Markov a
deadEnd p start markov = go S.empty [start]
  where
    edges = (M.fromList markov M.!)
    go done [] = []
    go done (x:todo) | x `elem` done = go done todo
                     | otherwise = (x, e) : go (S.insert x done) todo'
      where e = edges x
            todo' | p x      = e ++ todo
                  | otherwise = todo

liveNodes :: Ord a => a -> Markov a -> S.Set a
liveNodes start markov = go (S.singleton start)
  where
    edges = (M.fromList markov M.!)
    go live | S.size live == S.size live' = live
            | otherwise                   = go live'
      where live' = live `S.union` S.fromList [ n | (n,e) <- markov, any (`S.member` live) e ]

-- | This function reduces the number of nodes by commoning up
-- equivalent nodes. Needs to be iterated.
reduceMarkovOnce :: (Ord a, Ord b) => (a -> b) -> Markov a -> Markov a
reduceMarkovOnce l m
    = [ (n, map f e) | (n, e) <- m, f n == n ]
  where
    f = (M.!) $
        M.fromList
        [ (x, head dup)
        | dupSet <- M.elems $ M.fromListWith S.union [ ((l n,sort e),S.singleton n) | (n,e) <- m ]
        , let dup = S.toList dupSet
        , x <- dup
        ]

reduceMarkov :: (Ord a, Ord b) => (a -> b) -> Markov a -> Markov a
reduceMarkov leak m | length m == length m' = m
                    | otherwise             = reduceMarkov leak m'
  where m' = reduceMarkovOnce leak m

-- If two independent copies of the markov chain, starting in the given
-- distribution are run n steps, what is the probability that the output is the
-- same?
sampleCollisions :: forall a b. (Ord a, Eq b) => Int -> (a -> b) -> Markov a -> Dist a -> Double
sampleCollisions n f markov dist =
    sum [ p * fromIntegral (countCollisions n n1 n2)
        | ((n1, n2),p) <- M.toList $ prodDist dist
        ] / 2^(2*n)
 where
    edges = (M.fromList markov M.!)

    -- From these two states on for further n steps, how many collisions?
    countCollisions :: Int -> a -> a -> Integer
    countCollisions 0 _ _ = 1
    countCollisions i n1 n2
        | f n1 == f n2 = sum [ countMemo (i-1) (n1', n2')
                             | n1' <- edges n1
                             , n2' <- edges n2
                             ]
        | otherwise = 0

    -- Yay, memoization!
    countMemo = curry (LM.fromList [ ((i,(n1,n2)), countCollisions i n1 n2) | i <- [0..n], n1 <- map fst markov, n2 <- map fst markov] M.!)

-- A typical output function
leak :: (a,[b]) -> b
leak (_,s) = head s

colorNode :: (a,[L]) -> String
colorNode s = if leak s == S then "lightblue" else "green1"

showNode :: (Show a, Show b) => (a,[b]) -> String
showNode (node,s) = show s ++ show node

showNode2 :: (Show a, Show b) => ((a,[b]),(a,[b])) -> String
showNode2 (n1,n2) = showNode n1 ++ "_" ++ showNode n2

calcPnl :: forall a b. (Ord a, Ord b) => Int -> (a -> b) -> Markov a -> Dist a -> Dist (a, [b])
calcPnl n f markov dist = jointDist dist $ \x1 -> normDist $ leakDist n x1
 where
    edges = (M.fromList markov M.!)

    -- Calculate the distribution of leaks
    -- (of the given length, starting in the given node)
    -- This is the performance hot spot in this program! (until I introduce memoization)
    leakDist :: Int -> a -> Dist [b]
    leakDist 0 _    = M.singleton [] 1
    leakDist i node = M.mapKeysMonotonic (f node :) $ addDists $
                      [ leakDistMemo (i-1) n' | n' <- edges node ]

    -- Yay, memoization!
    leakDistMemo = curry (LM.fromList [ ((i,node), leakDist i node) | i <- [0..n], node <- map fst markov ] M.!)

calcPl :: forall a b. (Ord a, Ord b) => Int -> (a -> b) -> Markov a -> Dist a -> Dist [b]
calcPl n f markov dist =
    addDists [ fmap (p*) (normDist $ leakDist n x1) | (x1, p) <- M.toList dist ]
 where
    edges = (M.fromList markov M.!)

    -- Calculate the distribution of leaks
    -- (of the given length, starting in the given node)
    -- This is the performance hot spot in this program! (until I introduce memoization)
    leakDist :: Int -> a -> Dist [b]
    leakDist 0 _    = M.singleton [] 1
    leakDist i node = M.mapKeysMonotonic (f node :) $ addDists $
                      [ leakDistMemo (i-1) n' | n' <- edges node ]

    -- Yay, memoization!
    leakDistMemo = curry (LM.fromList [ ((i,node), leakDist i node) | i <- [0..n], node <- map fst markov ] M.!)

condEntropyLB :: (Ord a, Ord b) => Dist (a, [b]) -> Dist (a, [b]) -> Double
condEntropyLB pnl pnlbar = sum
    [ pxy * logb (px/pxy)
    | ((xi,ys),pxy) <- M.toList pnl
    , let px = pnlbar M.! (xi, init ys)
    ]

condEntropyUB :: Ord b => Dist [b] -> Dist [b] -> Double
condEntropyUB pl plbar = sum
    [ pxy * logb (px/pxy)
    | (ys,pxy) <- M.toList pl
    , let px = plbar M.! init ys
    ]

markovStepProb :: Ord a => (a -> Bool) -> Markov a -> Dist a -> Double
markovStepProb p markov dist = sum
    [ pn * sum [ 1 | n' <- e, p n'] / fromIntegral (length e)
    | (n,pn) <- M.toList dist
    , let e = edges n
    ]
  where
    edges = (M.fromList markov M.!)

logb :: Double -> Double
logb = logBase 2

samplesToDist :: Ord a => [a] -> Dist a
samplesToDist xs = normDist $ M.fromListWith (+) [ (x,1) | x <- xs ]

-- | Take the union of two distributions (not normalised)
addDist :: Ord a => Dist a -> Dist a -> Dist a
addDist = M.unionWith (+)

addDists :: Ord a => [Dist a]-> Dist a
addDists = M.unionsWith (+)

mapDist :: Ord b => (a -> b) -> Dist a -> Dist b
mapDist = M.mapKeysWith (+)

normDist :: M.Map a Double -> M.Map a Double
normDist m = fmap (/sum m) m

jointDist :: (Ord a, Ord b) => Dist a -> (a -> Dist b) -> Dist (a,b)
jointDist d1 f = M.unionsWith (error "Not disjoint") $
        [ M.mapKeysMonotonic (a,) $ fmap (p_a*) (f a)
        | (a, p_a) <- M.toList d1
        ]

prodDist :: Ord a => Dist a -> Dist (a,a)
prodDist d = M.fromListWith (error "input wrong") $
    [ ((x,y), px * py) | (x,px) <- M.toList d, (y,py) <- M.toList d ]

samplePred :: (a -> Bool) -> Dist a -> Double
samplePred ok d = sum [ p | (x,p) <- M.toList d, ok x ]

conditionalDist :: Ord a => Dist (Maybe a) -> Dist a
conditionalDist d = normDist $ M.fromListWith (error "input wrong") $
    [ (n, p) | (Just n, p) <- M.toList d ]

data WhichGraph = RTL | LTR deriving Eq
data WhatToDo = PrintMealy
              | PrintMarkov
              | PrintCollMarkov
              | PrintEntropy
              | PrintTableRow
              | PrintColl
    deriving Eq

run :: Int -> Int -> WhichGraph -> WhatToDo -> IO ()
run w n which todo = do
    let my = case which of
            RTL -> reduceMealy $ buildMealyRTL w
            LTR -> reduceMealy $ buildMealy w
    let n0 = (0, replicate w S)
    let mv = buildMarkov n0 my
    let sd = stationary mv
    let singletonDist = M.singleton n0 1
    let ud = M.fromList mv $> ((1::Double) / fromIntegral (length mv))
    let pnl    = calcPnl n leak mv sd
    let pnlbar = calcPnl (n-1) leak mv sd
    let eLB = condEntropyLB pnl pnlbar
    let pl     = calcPl n leak mv sd
    let plbar  = calcPl (n-1) leak mv sd
    let eUB = condEntropyUB pl plbar

    let prodMv   = prodMarkov mv
    let thinProdMV = reduceMarkov (\(n1,n2) -> leak n1 == leak n2) prodMv
        -- ^ This is just to feed less data to the linear algrebra system
    let collSD   = stationaryCond (\(n1,n2) -> leak n1 == leak n2) thinProdMV
    let collP    = samplePred (\(n1,n2) -> leak n1 == leak n2) collSD
    let collRate = - logb collP

    let simCollP = sampleCollisions n leak mv sd -- singletonDist
    let simCollisions = simCollP^(2::Int) * 2^n
    let simCollRate = - (logb simCollP) / fromIntegral n

    case todo of
        PrintMealy      -> putStrLn (mealy2dot my)
        PrintMarkov     -> putStrLn (markov2dot colorNode showNode mv)
        PrintCollMarkov -> error "unimplemented" -- putStrLn (markov2dot showNode2 collidingMv)
        PrintTableRow   -> printf "%d\t%d\t%f\t%f\t%f\n" w n eLB eUB collRate
        PrintColl       -> error "unimplemented"  -- print collisions
        PrintEntropy    -> printf "w: %2d. n: %3d. %.4f ≤ H(Y) ≤ %.4f. coll rate: %.4f. sim coll rate: %.4f \n"
                                w n eLB eUB collRate simCollRate

main :: IO ()
main = join . execParser $
  info (helper <*> parser)
  (  fullDesc
  <> header "Entropy calculation"
  )
  where
    parser :: Parser (IO ())
    parser = run
      <$> option auto
         (  long "width"
         <> short 'w'
         <> metavar "WIDTH"
         <> help "Width of the window (typical 4 or 5)"
         <> value 4
         <> showDefault
         )
      <*> option auto
         (  long "n"
         <> short 'n'
         <> metavar "N"
         <> help "approximation"
         <> value 10
         <> showDefault
         )
      <*> choice
        [ flag' LTR (long "ltr" <> help "Analyse left-to-right")
        , flag' RTL (long "rtl" <> help "Analyse right-to-left")
        , pure LTR
        ]
      <*> choice
        [ flag' PrintMealy    (long "mealy"  <> help "Print Mealy graph in .dot format")
        , flag' PrintMarkov   (long "markov" <> help "Print Markov graph in .dot format")
        , flag' PrintCollMarkov   (long "colliding-markov" <> help "Print colliding Markov graph in .dot format")
        , flag' PrintTableRow (long "table"  <> help "Print a tsv table row")
        , flag' PrintColl     (long "coll"   <> help "Print collision probability tsv row")
        , pure PrintEntropy
        ]

choice :: Alternative f => [f a] -> f a
choice = foldr1 (<|>)


{-# LANGUAGE FlexibleContexts #-}

{-| FM demodulation pipes -}
module SDR.Demod (
    fmDemodStr,
    fmDemodVec,
    fmDemod
    ) where

import Data.Complex
import qualified Data.Vector.Generic               as VG
import qualified Data.Vector.Fusion.Stream.Monadic as VFSM
import qualified Data.Vector.Fusion.Bundle.Monadic as VFBM
import qualified Data.Vector.Fusion.Bundle.Size    as VFBS

import SDR.VectorUtils
import Pipes

-- | FM demodulate a stream of complex samples
{-# INLINE fmDemodStr #-}
fmDemodStr :: (RealFloat a, Monad m)
           => Complex a          -- ^ The starting sample - i.e. the last sample in the last buffer
           -> VFSM.Stream m (Complex a) -- ^ The input stream
           -> VFSM.Stream m a           -- ^ The output stream
fmDemodStr = mapAccumMV func
    where
    {-# INLINE func #-}
    func last sample = return (sample, phase (sample * conjugate last))

-- | FM demodulate a vector of complex samples
{-# INLINE fmDemodVec #-}
fmDemodVec :: (RealFloat a, VG.Vector v (Complex a), VG.Vector v a)
           => Complex a     -- ^ The starting sample - i.e. the last sample in the last buffer
           -> v (Complex a) -- ^ The input Vector
           -> v a           -- ^ The output Vector
fmDemodVec init inp = VG.unstream $ flip VFBM.fromStream (VFBS.Exact $ VG.length inp) $ fmDemodStr init $ VFBM.elements $ VG.stream inp

-- | Pipe that performs FM demodulation
{-# INLINE fmDemod #-}
fmDemod :: (RealFloat a, VG.Vector v (Complex a), VG.Vector v a) => Pipe (v (Complex a)) (v a) IO ()
fmDemod = func 0
    where
    func lastSample = do
        dat <- await
        yield $ fmDemodVec lastSample dat
        func $ VG.unsafeIndex dat (VG.length dat - 1)