text:
{-# LANGUAGE Safe #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Functor.Classes
-- Copyright   :  (c) Ross Paterson 2013
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  libraries@haskell.org
-- Stability   :  experimental
-- Portability :  portable
--
-- Liftings of the Prelude classes 'Eq', 'Ord', 'Read' and 'Show' to
-- unary and binary type constructors.
--
-- These classes are needed to express the constraints on arguments of
-- transformers in portable Haskell.  Thus for a new transformer @T@,
-- one might write instances like
--
-- > instance (Eq1 f) => Eq1 (T f) where ...
-- > instance (Ord1 f) => Ord1 (T f) where ...
-- > instance (Read1 f) => Read1 (T f) where ...
-- > instance (Show1 f) => Show1 (T f) where ...
--
-- If these instances can be defined, defining instances of the base
-- classes is mechanical:
--
-- > instance (Eq1 f, Eq a) => Eq (T f a) where (==) = eq1
-- > instance (Ord1 f, Ord a) => Ord (T f a) where compare = compare1
-- > instance (Read1 f, Read a) => Read (T f a) where
-- >   readPrec     = readPrec1
-- >   readListPrec = readListPrecDefault
-- > instance (Show1 f, Show a) => Show (T f a) where showsPrec = showsPrec1
--
-- @since 4.9.0.0
-----------------------------------------------------------------------------

module Data.Functor.Classes (
    -- * Liftings of Prelude classes
    -- ** For unary constructors
    Eq1(..), eq1,
    Ord1(..), compare1,
    Read1(..), readsPrec1, readPrec1,
    liftReadListDefault, liftReadListPrecDefault,
    Show1(..), showsPrec1,
    -- ** For binary constructors
    Eq2(..), eq2,
    Ord2(..), compare2,
    Read2(..), readsPrec2, readPrec2,
    liftReadList2Default, liftReadListPrec2Default,
    Show2(..), showsPrec2,
    -- * Helper functions
    -- $example
    readsData, readData,
    readsUnaryWith, readUnaryWith,
    readsBinaryWith, readBinaryWith,
    showsUnaryWith,
    showsBinaryWith,
    -- ** Obsolete helpers
    readsUnary,
    readsUnary1,
    readsBinary1,
    showsUnary,
    showsUnary1,
    showsBinary1,
  ) where

import Control.Applicative (Alternative((<|>)), Const(Const))

import Data.Functor.Identity (Identity(Identity))
import Data.Proxy (Proxy(Proxy))
import Data.List.NonEmpty (NonEmpty(..))
import Data.Ord (Down(Down))
import Data.Complex (Complex((:+)))

import GHC.Tuple (Solo (..))
import GHC.Read (expectP, list, paren)

import Text.ParserCombinators.ReadPrec (ReadPrec, readPrec_to_S, readS_to_Prec)
import Text.Read (Read(..), parens, prec, step)
import Text.Read.Lex (Lexeme(..))
import Text.Show (showListWith)

-- $setup
-- >>> import Prelude
-- >>> import Data.Complex (Complex (..))
-- >>> import Text.ParserCombinators.ReadPrec

-- | Lifting of the 'Eq' class to unary type constructors.
--
-- @since 4.9.0.0
class Eq1 f where
    -- | Lift an equality test through the type constructor.
    --
    -- The function will usually be applied to an equality function,
    -- but the more general type ensures that the implementation uses
    -- it to compare elements of the first container with elements of
    -- the second.
    --
    -- @since 4.9.0.0
    liftEq :: (a -> b -> Bool) -> f a -> f b -> Bool

-- | Lift the standard @('==')@ function through the type constructor.
--
-- @since 4.9.0.0
eq1 :: (Eq1 f, Eq a) => f a -> f a -> Bool
eq1 = liftEq (==)

-- | Lifting of the 'Ord' class to unary type constructors.
--
-- @since 4.9.0.0
class (Eq1 f) => Ord1 f where
    -- | Lift a 'compare' function through the type constructor.
    --
    -- The function will usually be applied to a comparison function,
    -- but the more general type ensures that the implementation uses
    -- it to compare elements of the first container with elements of
    -- the second.
    --
    -- @since 4.9.0.0
    liftCompare :: (a -> b -> Ordering) -> f a -> f b -> Ordering

-- | Lift the standard 'compare' function through the type constructor.
--
-- @since 4.9.0.0
compare1 :: (Ord1 f, Ord a) => f a -> f a -> Ordering
compare1 = liftCompare compare

-- | Lifting of the 'Read' class to unary type constructors.
--
-- Both 'liftReadsPrec' and 'liftReadPrec' exist to match the interface
-- provided in the 'Read' type class, but it is recommended to implement
-- 'Read1' instances using 'liftReadPrec' as opposed to 'liftReadsPrec', since
-- the former is more efficient than the latter. For example:
--
-- @
-- instance 'Read1' T where
--   'liftReadPrec'     = ...
--   'liftReadListPrec' = 'liftReadListPrecDefault'
-- @
--
-- For more information, refer to the documentation for the 'Read' class.
--
-- @since 4.9.0.0
class Read1 f where
    {-# MINIMAL liftReadsPrec | liftReadPrec #-}

    -- | 'readsPrec' function for an application of the type constructor
    -- based on 'readsPrec' and 'readList' functions for the argument type.
    --
    -- @since 4.9.0.0
    liftReadsPrec :: (Int -> ReadS a) -> ReadS [a] -> Int -> ReadS (f a)
    liftReadsPrec rp rl = readPrec_to_S $
        liftReadPrec (readS_to_Prec rp) (readS_to_Prec (const rl))

    -- | 'readList' function for an application of the type constructor
    -- based on 'readsPrec' and 'readList' functions for the argument type.
    -- The default implementation using standard list syntax is correct
    -- for most types.
    --
    -- @since 4.9.0.0
    liftReadList :: (Int -> ReadS a) -> ReadS [a] -> ReadS [f a]
    liftReadList rp rl = readPrec_to_S
        (list $ liftReadPrec (readS_to_Prec rp) (readS_to_Prec (const rl))) 0

    -- | 'readPrec' function for an application of the type constructor
    -- based on 'readPrec' and 'readListPrec' functions for the argument type.
    --
    -- @since 4.10.0.0
    liftReadPrec :: ReadPrec a -> ReadPrec [a] -> ReadPrec (f a)
    liftReadPrec rp rl = readS_to_Prec $
        liftReadsPrec (readPrec_to_S rp) (readPrec_to_S rl 0)

    -- | 'readListPrec' function for an application of the type constructor
    -- based on 'readPrec' and 'readListPrec' functions for the argument type.
    --
    -- The default definition uses 'liftReadList'. Instances that define
    -- 'liftReadPrec' should also define 'liftReadListPrec' as
    -- 'liftReadListPrecDefault'.
    --
    -- @since 4.10.0.0
    liftReadListPrec :: ReadPrec a -> ReadPrec [a] -> ReadPrec [f a]
    liftReadListPrec rp rl = readS_to_Prec $ \_ ->
        liftReadList (readPrec_to_S rp) (readPrec_to_S rl 0)

-- | Lift the standard 'readsPrec' and 'readList' functions through the
-- type constructor.
--
-- @since 4.9.0.0
readsPrec1 :: (Read1 f, Read a) => Int -> ReadS (f a)
readsPrec1 = liftReadsPrec readsPrec readList

-- | Lift the standard 'readPrec' and 'readListPrec' functions through the
-- type constructor.
--
-- @since 4.10.0.0
readPrec1 :: (Read1 f, Read a) => ReadPrec (f a)
readPrec1 = liftReadPrec readPrec readListPrec

-- | A possible replacement definition for the 'liftReadList' method.
-- This is only needed for 'Read1' instances where 'liftReadListPrec' isn't
-- defined as 'liftReadListPrecDefault'.
--
-- @since 4.10.0.0
liftReadListDefault :: Read1 f => (Int -> ReadS a) -> ReadS [a] -> ReadS [f a]
liftReadListDefault rp rl = readPrec_to_S
    (liftReadListPrec (readS_to_Prec rp) (readS_to_Prec (const rl))) 0

-- | A possible replacement definition for the 'liftReadListPrec' method,
-- defined using 'liftReadPrec'.
--
-- @since 4.10.0.0
liftReadListPrecDefault :: Read1 f => ReadPrec a -> ReadPrec [a]
                        -> ReadPrec [f a]
liftReadListPrecDefault rp rl = list (liftReadPrec rp rl)

-- | Lifting of the 'Show' class to unary type constructors.
--
-- @since 4.9.0.0
class Show1 f where
    -- | 'showsPrec' function for an application of the type constructor
    -- based on 'showsPrec' and 'showList' functions for the argument type.
    --
    -- @since 4.9.0.0
    liftShowsPrec :: (Int -> a -> ShowS) -> ([a] -> ShowS) ->
        Int -> f a -> ShowS

    -- | 'showList' function for an application of the type constructor
    -- based on 'showsPrec' and 'showList' functions for the argument type.
    -- The default implementation using standard list syntax is correct
    -- for most types.
    --
    -- @since 4.9.0.0
    liftShowList :: (Int -> a -> ShowS) -> ([a] -> ShowS) ->
        [f a] -> ShowS
    liftShowList sp sl = showListWith (liftShowsPrec sp sl 0)

-- | Lift the standard 'showsPrec' and 'showList' functions through the
-- type constructor.
--
-- @since 4.9.0.0
showsPrec1 :: (Show1 f, Show a) => Int -> f a -> ShowS
showsPrec1 = liftShowsPrec showsPrec showList

-- | Lifting of the 'Eq' class to binary type constructors.
--
-- @since 4.9.0.0
class Eq2 f where
    -- | Lift equality tests through the type constructor.
    --
    -- The function will usually be applied to equality functions,
    -- but the more general type ensures that the implementation uses
    -- them to compare elements of the first container with elements of
    -- the second.
    --
    -- @since 4.9.0.0
    liftEq2 :: (a -> b -> Bool) -> (c -> d -> Bool) -> f a c -> f b d -> Bool

-- | Lift the standard @('==')@ function through the type constructor.
--
-- @since 4.9.0.0
eq2 :: (Eq2 f, Eq a, Eq b) => f a b -> f a b -> Bool
eq2 = liftEq2 (==) (==)

-- | Lifting of the 'Ord' class to binary type constructors.
--
-- @since 4.9.0.0
class (Eq2 f) => Ord2 f where
    -- | Lift 'compare' functions through the type constructor.
    --
    -- The function will usually be applied to comparison functions,
    -- but the more general type ensures that the implementation uses
    -- them to compare elements of the first container with elements of
    -- the second.
    --
    -- @since 4.9.0.0
    liftCompare2 :: (a -> b -> Ordering) -> (c -> d -> Ordering) ->
        f a c -> f b d -> Ordering

-- | Lift the standard 'compare' function through the type constructor.
--
-- @since 4.9.0.0
compare2 :: (Ord2 f, Ord a, Ord b) => f a b -> f a b -> Ordering
compare2 = liftCompare2 compare compare

-- | Lifting of the 'Read' class to binary type constructors.
--
-- Both 'liftReadsPrec2' and 'liftReadPrec2' exist to match the interface
-- provided in the 'Read' type class, but it is recommended to implement
-- 'Read2' instances using 'liftReadPrec2' as opposed to 'liftReadsPrec2',
-- since the former is more efficient than the latter. For example:
--
-- @
-- instance 'Read2' T where
--   'liftReadPrec2'     = ...
--   'liftReadListPrec2' = 'liftReadListPrec2Default'
-- @
--
-- For more information, refer to the documentation for the 'Read' class.
-- @since 4.9.0.0
class Read2 f where
    {-# MINIMAL liftReadsPrec2 | liftReadPrec2 #-}

    -- | 'readsPrec' function for an application of the type constructor
    -- based on 'readsPrec' and 'readList' functions for the argument types.
    --
    -- @since 4.9.0.0
    liftReadsPrec2 :: (Int -> ReadS a) -> ReadS [a] ->
        (Int -> ReadS b) -> ReadS [b] -> Int -> ReadS (f a b)
    liftReadsPrec2 rp1 rl1 rp2 rl2 = readPrec_to_S $
        liftReadPrec2 (readS_to_Prec rp1) (readS_to_Prec (const rl1))
                      (readS_to_Prec rp2) (readS_to_Prec (const rl2))

    -- | 'readList' function for an application of the type constructor
    -- based on 'readsPrec' and 'readList' functions for the argument types.
    -- The default implementation using standard list syntax is correct
    -- for most types.
    --
    -- @since 4.9.0.0
    liftReadList2 :: (Int -> ReadS a) -> ReadS [a] ->
        (Int -> ReadS b) -> ReadS [b] -> ReadS [f a b]
    liftReadList2 rp1 rl1 rp2 rl2 = readPrec_to_S
       (list $ liftReadPrec2 (readS_to_Prec rp1) (readS_to_Prec (const rl1))
                             (readS_to_Prec rp2) (readS_to_Prec (const rl2))) 0

    -- | 'readPrec' function for an application of the type constructor
    -- based on 'readPrec' and 'readListPrec' functions for the argument types.
    --
    -- @since 4.10.0.0
    liftReadPrec2 :: ReadPrec a -> ReadPrec [a] ->
        ReadPrec b -> ReadPrec [b] -> ReadPrec (f a b)
    liftReadPrec2 rp1 rl1 rp2 rl2 = readS_to_Prec $
        liftReadsPrec2 (readPrec_to_S rp1) (readPrec_to_S rl1 0)
                       (readPrec_to_S rp2) (readPrec_to_S rl2 0)

    -- | 'readListPrec' function for an application of the type constructor
    -- based on 'readPrec' and 'readListPrec' functions for the argument types.
    --
    -- The default definition uses 'liftReadList2'. Instances that define
    -- 'liftReadPrec2' should also define 'liftReadListPrec2' as
    -- 'liftReadListPrec2Default'.
    --
    -- @since 4.10.0.0
    liftReadListPrec2 :: ReadPrec a -> ReadPrec [a] ->
        ReadPrec b -> ReadPrec [b] -> ReadPrec [f a b]
    liftReadListPrec2 rp1 rl1 rp2 rl2 = readS_to_Prec $ \_ ->
        liftReadList2 (readPrec_to_S rp1) (readPrec_to_S rl1 0)
                      (readPrec_to_S rp2) (readPrec_to_S rl2 0)

-- | Lift the standard 'readsPrec' function through the type constructor.
--
-- @since 4.9.0.0
readsPrec2 :: (Read2 f, Read a, Read b) => Int -> ReadS (f a b)
readsPrec2 = liftReadsPrec2 readsPrec readList readsPrec readList

-- | Lift the standard 'readPrec' function through the type constructor.
--
-- @since 4.10.0.0
readPrec2 :: (Read2 f, Read a, Read b) => ReadPrec (f a b)
readPrec2 = liftReadPrec2 readPrec readListPrec readPrec readListPrec

-- | A possible replacement definition for the 'liftReadList2' method.
-- This is only needed for 'Read2' instances where 'liftReadListPrec2' isn't
-- defined as 'liftReadListPrec2Default'.
--
-- @since 4.10.0.0
liftReadList2Default :: Read2 f => (Int -> ReadS a) -> ReadS [a] ->
    (Int -> ReadS b) -> ReadS [b] ->ReadS [f a b]
liftReadList2Default rp1 rl1 rp2 rl2 = readPrec_to_S
    (liftReadListPrec2 (readS_to_Prec rp1) (readS_to_Prec (const rl1))
                       (readS_to_Prec rp2) (readS_to_Prec (const rl2))) 0

-- | A possible replacement definition for the 'liftReadListPrec2' method,
-- defined using 'liftReadPrec2'.
--
-- @since 4.10.0.0
liftReadListPrec2Default :: Read2 f => ReadPrec a -> ReadPrec [a] ->
    ReadPrec b -> ReadPrec [b] -> ReadPrec [f a b]
liftReadListPrec2Default rp1 rl1 rp2 rl2 = list (liftReadPrec2 rp1 rl1 rp2 rl2)

-- | Lifting of the 'Show' class to binary type constructors.
--
-- @since 4.9.0.0
class Show2 f where
    -- | 'showsPrec' function for an application of the type constructor
    -- based on 'showsPrec' and 'showList' functions for the argument types.
    --
    -- @since 4.9.0.0
    liftShowsPrec2 :: (Int -> a -> ShowS) -> ([a] -> ShowS) ->
        (Int -> b -> ShowS) -> ([b] -> ShowS) -> Int -> f a b -> ShowS

    -- | 'showList' function for an application of the type constructor
    -- based on 'showsPrec' and 'showList' functions for the argument types.
    -- The default implementation using standard list syntax is correct
    -- for most types.
    --
    -- @since 4.9.0.0
    liftShowList2 :: (Int -> a -> ShowS) -> ([a] -> ShowS) ->
        (Int -> b -> ShowS) -> ([b] -> ShowS) -> [f a b] -> ShowS
    liftShowList2 sp1 sl1 sp2 sl2 =
        showListWith (liftShowsPrec2 sp1 sl1 sp2 sl2 0)

-- | Lift the standard 'showsPrec' function through the type constructor.
--
-- @since 4.9.0.0
showsPrec2 :: (Show2 f, Show a, Show b) => Int -> f a b -> ShowS
showsPrec2 = liftShowsPrec2 showsPrec showList showsPrec showList

-- Instances for Prelude type constructors

-- | @since 4.9.0.0
instance Eq1 Maybe where
    liftEq _ Nothing Nothing = True
    liftEq _ Nothing (Just _) = False
    liftEq _ (Just _) Nothing = False
    liftEq eq (Just x) (Just y) = eq x y

-- | @since 4.9.0.0
instance Ord1 Maybe where
    liftCompare _ Nothing Nothing = EQ
    liftCompare _ Nothing (Just _) = LT
    liftCompare _ (Just _) Nothing = GT
    liftCompare comp (Just x) (Just y) = comp x y

-- | @since 4.9.0.0
instance Read1 Maybe where
    liftReadPrec rp _ =
        parens (expectP (Ident "Nothing") *> pure Nothing)
        <|>
        readData (readUnaryWith rp "Just" Just)

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.9.0.0
instance Show1 Maybe where
    liftShowsPrec _ _ _ Nothing = showString "Nothing"
    liftShowsPrec sp _ d (Just x) = showsUnaryWith sp "Just" d x

-- | @since 4.9.0.0
instance Eq1 [] where
    liftEq _ [] [] = True
    liftEq _ [] (_:_) = False
    liftEq _ (_:_) [] = False
    liftEq eq (x:xs) (y:ys) = eq x y && liftEq eq xs ys

-- | @since 4.9.0.0
instance Ord1 [] where
    liftCompare _ [] [] = EQ
    liftCompare _ [] (_:_) = LT
    liftCompare _ (_:_) [] = GT
    liftCompare comp (x:xs) (y:ys) = comp x y `mappend` liftCompare comp xs ys

-- | @since 4.9.0.0
instance Read1 [] where
    liftReadPrec _ rl = rl
    liftReadListPrec  = liftReadListPrecDefault
    liftReadList      = liftReadListDefault

-- | @since 4.9.0.0
instance Show1 [] where
    liftShowsPrec _ sl _ = sl

-- | @since 4.10.0.0
instance Eq1 NonEmpty where
  liftEq eq (a :| as) (b :| bs) = eq a b && liftEq eq as bs

-- | @since 4.10.0.0
instance Ord1 NonEmpty where
  liftCompare cmp (a :| as) (b :| bs) = cmp a b `mappend` liftCompare cmp as bs

-- | @since 4.10.0.0
instance Read1 NonEmpty where
  liftReadsPrec rdP rdL p s = readParen (p > 5) (\s' -> do
    (a, s'') <- rdP 6 s'
    (":|", s''') <- lex s''
    (as, s'''') <- rdL s'''
    return (a :| as, s'''')) s

-- | @since 4.10.0.0
instance Show1 NonEmpty where
  liftShowsPrec shwP shwL p (a :| as) = showParen (p > 5) $
    shwP 6 a . showString " :| " . shwL as

-- | @since 4.9.0.0
instance Eq2 (,) where
    liftEq2 e1 e2 (x1, y1) (x2, y2) = e1 x1 x2 && e2 y1 y2

-- | @since 4.9.0.0
instance Ord2 (,) where
    liftCompare2 comp1 comp2 (x1, y1) (x2, y2) =
        comp1 x1 x2 `mappend` comp2 y1 y2

-- | @since 4.9.0.0
instance Read2 (,) where
    liftReadPrec2 rp1 _ rp2 _ = parens $ paren $ do
        x <- rp1
        expectP (Punc ",")
        y <- rp2
        return (x,y)

    liftReadListPrec2 = liftReadListPrec2Default
    liftReadList2     = liftReadList2Default

-- | @since 4.9.0.0
instance Show2 (,) where
    liftShowsPrec2 sp1 _ sp2 _ _ (x, y) =
        showChar '(' . sp1 0 x . showChar ',' . sp2 0 y . showChar ')'

-- | @since 4.15
instance Eq1 Solo where
  liftEq eq (Solo a) (Solo b) = a `eq` b

-- | @since 4.9.0.0
instance (Eq a) => Eq1 ((,) a) where
    liftEq = liftEq2 (==)

-- | @since 4.15
instance Ord1 Solo where
  liftCompare cmp (Solo a) (Solo b) = cmp a b

-- | @since 4.9.0.0
instance (Ord a) => Ord1 ((,) a) where
    liftCompare = liftCompare2 compare

-- | @since 4.15
instance Read1 Solo where
    liftReadPrec rp _ = readData (readUnaryWith rp "Solo" Solo)

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.9.0.0
instance (Read a) => Read1 ((,) a) where
    liftReadPrec = liftReadPrec2 readPrec readListPrec

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.15
instance Show1 Solo where
    liftShowsPrec sp _ d (Solo x) = showsUnaryWith sp "Solo" d x

-- | @since 4.9.0.0
instance (Show a) => Show1 ((,) a) where
    liftShowsPrec = liftShowsPrec2 showsPrec showList

-- | @since 4.16.0.0
--
-- >>> eq2 ('x', True, "str") ('x', True, "str")
-- True
--
instance Eq a => Eq2 ((,,) a) where
    liftEq2 e1 e2 (u1, x1, y1) (v1, x2, y2) =
        u1 == v1 &&
        e1 x1 x2 && e2 y1 y2

-- | @since 4.16.0.0
--
-- >>> compare2 ('x', True, "aaa") ('x', True, "zzz")
-- LT
instance Ord a => Ord2 ((,,) a) where
    liftCompare2 comp1 comp2 (u1, x1, y1) (v1, x2, y2) =
        compare u1 v1 `mappend`
        comp1 x1 x2 `mappend` comp2 y1 y2

-- | @since 4.16.0.0
--
-- >>> readPrec_to_S readPrec2 0 "('x', True, 2)" :: [((Char, Bool, Int), String)]
-- [(('x',True,2),"")]
--
instance Read a => Read2 ((,,) a) where
    liftReadPrec2 rp1 _ rp2 _ = parens $ paren $ do
        x1 <- readPrec
        expectP (Punc ",")
        y1 <- rp1
        expectP (Punc ",")
        y2 <- rp2
        return (x1,y1,y2)

    liftReadListPrec2 = liftReadListPrec2Default
    liftReadList2     = liftReadList2Default

-- | @since 4.16.0.0
--
-- >>> showsPrec2 0 ('x', True, 2 :: Int) ""
-- "('x',True,2)"
--
instance Show a => Show2 ((,,) a) where
    liftShowsPrec2 sp1 _ sp2 _ _ (x1,y1,y2)
        = showChar '(' . showsPrec 0 x1
        . showChar ',' . sp1 0 y1
        . showChar ',' . sp2 0 y2
        . showChar ')'

-- | @since 4.16.0.0
instance (Eq a, Eq b) => Eq1 ((,,) a b) where
    liftEq = liftEq2 (==)

-- | @since 4.16.0.0
instance (Ord a, Ord b) => Ord1 ((,,) a b) where
    liftCompare = liftCompare2 compare

-- | @since 4.16.0.0
instance (Read a, Read b) => Read1 ((,,) a b) where
    liftReadPrec = liftReadPrec2 readPrec readListPrec

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.16.0.0
instance (Show a, Show b) => Show1 ((,,) a b) where
    liftShowsPrec = liftShowsPrec2 showsPrec showList

-- | @since 4.16.0.0
--
-- >>> eq2 ('x', True, "str", 2) ('x', True, "str", 2 :: Int)
-- True
--
instance (Eq a, Eq b) => Eq2 ((,,,) a b) where
    liftEq2 e1 e2 (u1, u2, x1, y1) (v1, v2, x2, y2) =
        u1 == v1 &&
        u2 == v2 &&
        e1 x1 x2 && e2 y1 y2

-- | @since 4.16.0.0
--
-- >>> compare2 ('x', True, "str", 2) ('x', True, "str", 3 :: Int)
-- LT
--
instance (Ord a, Ord b) => Ord2 ((,,,) a b) where
    liftCompare2 comp1 comp2 (u1, u2, x1, y1) (v1, v2, x2, y2) =
        compare u1 v1 `mappend`
        compare u2 v2 `mappend`
        comp1 x1 x2 `mappend` comp2 y1 y2

-- | @since 4.16.0.0
--
-- >>> readPrec_to_S readPrec2 0 "('x', True, 2, 4.5)" :: [((Char, Bool, Int, Double), String)]
-- [(('x',True,2,4.5),"")]
--
instance (Read a, Read b) => Read2 ((,,,) a b) where
    liftReadPrec2 rp1 _ rp2 _ = parens $ paren $ do
        x1 <- readPrec
        expectP (Punc ",")
        x2 <- readPrec
        expectP (Punc ",")
        y1 <- rp1
        expectP (Punc ",")
        y2 <- rp2
        return (x1,x2,y1,y2)

    liftReadListPrec2 = liftReadListPrec2Default
    liftReadList2     = liftReadList2Default

-- | @since 4.16.0.0
--
-- >>> showsPrec2 0 ('x', True, 2 :: Int, 4.5 :: Double) ""
-- "('x',True,2,4.5)"
--
instance (Show a, Show b) => Show2 ((,,,) a b) where
    liftShowsPrec2 sp1 _ sp2 _ _ (x1,x2,y1,y2)
        = showChar '(' . showsPrec 0 x1
        . showChar ',' . showsPrec 0 x2
        . showChar ',' . sp1 0 y1
        . showChar ',' . sp2 0 y2
        . showChar ')'

-- | @since 4.16.0.0
instance (Eq a, Eq b, Eq c) => Eq1 ((,,,) a b c) where
    liftEq = liftEq2 (==)

-- | @since 4.16.0.0
instance (Ord a, Ord b, Ord c) => Ord1 ((,,,) a b c) where
    liftCompare = liftCompare2 compare

-- | @since 4.16.0.0
instance (Read a, Read b, Read c) => Read1 ((,,,) a b c) where
    liftReadPrec = liftReadPrec2 readPrec readListPrec

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.16.0.0
instance (Show a, Show b, Show c) => Show1 ((,,,) a b c) where
    liftShowsPrec = liftShowsPrec2 showsPrec showList

-- | @since 4.9.0.0
instance Eq2 Either where
    liftEq2 e1 _ (Left x) (Left y) = e1 x y
    liftEq2 _ _ (Left _) (Right _) = False
    liftEq2 _ _ (Right _) (Left _) = False
    liftEq2 _ e2 (Right x) (Right y) = e2 x y

-- | @since 4.9.0.0
instance Ord2 Either where
    liftCompare2 comp1 _ (Left x) (Left y) = comp1 x y
    liftCompare2 _ _ (Left _) (Right _) = LT
    liftCompare2 _ _ (Right _) (Left _) = GT
    liftCompare2 _ comp2 (Right x) (Right y) = comp2 x y

-- | @since 4.9.0.0
instance Read2 Either where
    liftReadPrec2 rp1 _ rp2 _ = readData $
         readUnaryWith rp1 "Left" Left <|>
         readUnaryWith rp2 "Right" Right

    liftReadListPrec2 = liftReadListPrec2Default
    liftReadList2     = liftReadList2Default

-- | @since 4.9.0.0
instance Show2 Either where
    liftShowsPrec2 sp1 _ _ _ d (Left x) = showsUnaryWith sp1 "Left" d x
    liftShowsPrec2 _ _ sp2 _ d (Right x) = showsUnaryWith sp2 "Right" d x

-- | @since 4.9.0.0
instance (Eq a) => Eq1 (Either a) where
    liftEq = liftEq2 (==)

-- | @since 4.9.0.0
instance (Ord a) => Ord1 (Either a) where
    liftCompare = liftCompare2 compare

-- | @since 4.9.0.0
instance (Read a) => Read1 (Either a) where
    liftReadPrec = liftReadPrec2 readPrec readListPrec

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.9.0.0
instance (Show a) => Show1 (Either a) where
    liftShowsPrec = liftShowsPrec2 showsPrec showList

-- Instances for other functors defined in the base package

-- | @since 4.9.0.0
instance Eq1 Identity where
    liftEq eq (Identity x) (Identity y) = eq x y

-- | @since 4.9.0.0
instance Ord1 Identity where
    liftCompare comp (Identity x) (Identity y) = comp x y

-- | @since 4.9.0.0
instance Read1 Identity where
    liftReadPrec rp _ = readData $
         readUnaryWith rp "Identity" Identity

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.9.0.0
instance Show1 Identity where
    liftShowsPrec sp _ d (Identity x) = showsUnaryWith sp "Identity" d x

-- | @since 4.9.0.0
instance Eq2 Const where
    liftEq2 eq _ (Const x) (Const y) = eq x y

-- | @since 4.9.0.0
instance Ord2 Const where
    liftCompare2 comp _ (Const x) (Const y) = comp x y

-- | @since 4.9.0.0
instance Read2 Const where
    liftReadPrec2 rp _ _ _ = readData $
         readUnaryWith rp "Const" Const

    liftReadListPrec2 = liftReadListPrec2Default
    liftReadList2     = liftReadList2Default

-- | @since 4.9.0.0
instance Show2 Const where
    liftShowsPrec2 sp _ _ _ d (Const x) = showsUnaryWith sp "Const" d x

-- | @since 4.9.0.0
instance (Eq a) => Eq1 (Const a) where
    liftEq = liftEq2 (==)
-- | @since 4.9.0.0
instance (Ord a) => Ord1 (Const a) where
    liftCompare = liftCompare2 compare
-- | @since 4.9.0.0
instance (Read a) => Read1 (Const a) where
    liftReadPrec = liftReadPrec2 readPrec readListPrec

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault
-- | @since 4.9.0.0
instance (Show a) => Show1 (Const a) where
    liftShowsPrec = liftShowsPrec2 showsPrec showList

-- Proxy unfortunately imports this module, hence these instances are placed
-- here,
-- | @since 4.9.0.0
instance Eq1 Proxy where
  liftEq _ _ _ = True

-- | @since 4.9.0.0
instance Ord1 Proxy where
  liftCompare _ _ _ = EQ

-- | @since 4.9.0.0
instance Show1 Proxy where
  liftShowsPrec _ _ _ _ = showString "Proxy"

-- | @since 4.9.0.0
instance Read1 Proxy where
  liftReadPrec _ _ = parens (expectP (Ident "Proxy") *> pure Proxy)

  liftReadListPrec = liftReadListPrecDefault
  liftReadList     = liftReadListDefault

-- | @since 4.12.0.0
instance Eq1 Down where
    liftEq eq (Down x) (Down y) = eq x y

-- | @since 4.12.0.0
instance Ord1 Down where
    liftCompare comp (Down x) (Down y) = comp x y

-- | @since 4.12.0.0
instance Read1 Down where
    liftReadsPrec rp _ = readsData $
         readsUnaryWith rp "Down" Down

-- | @since 4.12.0.0
instance Show1 Down where
    liftShowsPrec sp _ d (Down x) = showsUnaryWith sp "Down" d x

-- | @since 4.16.0.0
--
-- >>> eq1 (1 :+ 2) (1 :+ 2)
-- True
--
-- >>> eq1 (1 :+ 2) (1 :+ 3)
-- False
--
instance Eq1 Complex where
    liftEq eq (x :+ y) (u :+ v) = eq x u && eq y v

-- | @since 4.16.0.0
--
-- >>> readPrec_to_S readPrec1 0 "(2 % 3) :+ (3 % 4)" :: [(Complex Rational, String)]
-- [(2 % 3 :+ 3 % 4,"")]
--
instance Read1 Complex where
    liftReadPrec rp _  = parens $ prec complexPrec $ do
        x <- step rp
        expectP (Symbol ":+")
        y <- step rp
        return (x :+ y)
      where
        complexPrec = 6

    liftReadListPrec = liftReadListPrecDefault
    liftReadList     = liftReadListDefault

-- | @since 4.16.0.0
--
-- >>> showsPrec1 0 (2 :+ 3) ""
-- "2 :+ 3"
--
instance Show1 Complex where
    liftShowsPrec sp _ d (x :+ y) = showParen (d > complexPrec) $
        sp (complexPrec+1) x . showString " :+ " . sp (complexPrec+1) y
      where
        complexPrec = 6

-- Building blocks

-- | @'readsData' p d@ is a parser for datatypes where each alternative
-- begins with a data constructor.  It parses the constructor and
-- passes it to @p@.  Parsers for various constructors can be constructed
-- with 'readsUnary', 'readsUnary1' and 'readsBinary1', and combined with
-- @mappend@ from the @Monoid@ class.
--
-- @since 4.9.0.0
readsData :: (String -> ReadS a) -> Int -> ReadS a
readsData reader d =
    readParen (d > 10) $ \ r -> [res | (kw,s) <- lex r, res <- reader kw s]

-- | @'readData' p@ is a parser for datatypes where each alternative
-- begins with a data constructor.  It parses the constructor and
-- passes it to @p@.  Parsers for various constructors can be constructed
-- with 'readUnaryWith' and 'readBinaryWith', and combined with
-- '(<|>)' from the 'Alternative' class.
--
-- @since 4.10.0.0
readData :: ReadPrec a -> ReadPrec a
readData reader = parens $ prec 10 reader

-- | @'readsUnaryWith' rp n c n'@ matches the name of a unary data constructor
-- and then parses its argument using @rp@.
--
-- @since 4.9.0.0
readsUnaryWith :: (Int -> ReadS a) -> String -> (a -> t) -> String -> ReadS t
readsUnaryWith rp name cons kw s =
    [(cons x,t) | kw == name, (x,t) <- rp 11 s]

-- | @'readUnaryWith' rp n c'@ matches the name of a unary data constructor
-- and then parses its argument using @rp@.
--
-- @since 4.10.0.0
readUnaryWith :: ReadPrec a -> String -> (a -> t) -> ReadPrec t
readUnaryWith rp name cons = do
    expectP $ Ident name
    x <- step rp
    return $ cons x

-- | @'readsBinaryWith' rp1 rp2 n c n'@ matches the name of a binary
-- data constructor and then parses its arguments using @rp1@ and @rp2@
-- respectively.
--
-- @since 4.9.0.0
readsBinaryWith :: (Int -> ReadS a) -> (Int -> ReadS b) ->
    String -> (a -> b -> t) -> String -> ReadS t
readsBinaryWith rp1 rp2 name cons kw s =
    [(cons x y,u) | kw == name, (x,t) <- rp1 11 s, (y,u) <- rp2 11 t]

-- | @'readBinaryWith' rp1 rp2 n c'@ matches the name of a binary
-- data constructor and then parses its arguments using @rp1@ and @rp2@
-- respectively.
--
-- @since 4.10.0.0
readBinaryWith :: ReadPrec a -> ReadPrec b ->
    String -> (a -> b -> t) -> ReadPrec t
readBinaryWith rp1 rp2 name cons = do
    expectP $ Ident name
    x <- step rp1
    y <- step rp2
    return $ cons x y

-- | @'showsUnaryWith' sp n d x@ produces the string representation of a
-- unary data constructor with name @n@ and argument @x@, in precedence
-- context @d@.
--
-- @since 4.9.0.0
showsUnaryWith :: (Int -> a -> ShowS) -> String -> Int -> a -> ShowS
showsUnaryWith sp name d x = showParen (d > 10) $
    showString name . showChar ' ' . sp 11 x

-- | @'showsBinaryWith' sp1 sp2 n d x y@ produces the string
-- representation of a binary data constructor with name @n@ and arguments
-- @x@ and @y@, in precedence context @d@.
--
-- @since 4.9.0.0
showsBinaryWith :: (Int -> a -> ShowS) -> (Int -> b -> ShowS) ->
    String -> Int -> a -> b -> ShowS
showsBinaryWith sp1 sp2 name d x y = showParen (d > 10) $
    showString name . showChar ' ' . sp1 11 x . showChar ' ' . sp2 11 y

-- Obsolete building blocks

-- | @'readsUnary' n c n'@ matches the name of a unary data constructor
-- and then parses its argument using 'readsPrec'.
--
-- @since 4.9.0.0
{-# DEPRECATED readsUnary "Use 'readsUnaryWith' to define 'liftReadsPrec'" #-}
readsUnary :: (Read a) => String -> (a -> t) -> String -> ReadS t
readsUnary name cons kw s =
    [(cons x,t) | kw == name, (x,t) <- readsPrec 11 s]

-- | @'readsUnary1' n c n'@ matches the name of a unary data constructor
-- and then parses its argument using 'readsPrec1'.
--
-- @since 4.9.0.0
{-# DEPRECATED readsUnary1 "Use 'readsUnaryWith' to define 'liftReadsPrec'" #-}
readsUnary1 :: (Read1 f, Read a) => String -> (f a -> t) -> String -> ReadS t
readsUnary1 name cons kw s =
    [(cons x,t) | kw == name, (x,t) <- readsPrec1 11 s]

-- | @'readsBinary1' n c n'@ matches the name of a binary data constructor
-- and then parses its arguments using 'readsPrec1'.
--
-- @since 4.9.0.0
{-# DEPRECATED readsBinary1
      "Use 'readsBinaryWith' to define 'liftReadsPrec'" #-}
readsBinary1 :: (Read1 f, Read1 g, Read a) =>
    String -> (f a -> g a -> t) -> String -> ReadS t
readsBinary1 name cons kw s =
    [(cons x y,u) | kw == name,
        (x,t) <- readsPrec1 11 s, (y,u) <- readsPrec1 11 t]

-- | @'showsUnary' n d x@ produces the string representation of a unary data
-- constructor with name @n@ and argument @x@, in precedence context @d@.
--
-- @since 4.9.0.0
{-# DEPRECATED showsUnary "Use 'showsUnaryWith' to define 'liftShowsPrec'" #-}
showsUnary :: (Show a) => String -> Int -> a -> ShowS
showsUnary name d x = showParen (d > 10) $
    showString name . showChar ' ' . showsPrec 11 x

-- | @'showsUnary1' n d x@ produces the string representation of a unary data
-- constructor with name @n@ and argument @x@, in precedence context @d@.
--
-- @since 4.9.0.0
{-# DEPRECATED showsUnary1 "Use 'showsUnaryWith' to define 'liftShowsPrec'" #-}
showsUnary1 :: (Show1 f, Show a) => String -> Int -> f a -> ShowS
showsUnary1 name d x = showParen (d > 10) $
    showString name . showChar ' ' . showsPrec1 11 x

-- | @'showsBinary1' n d x y@ produces the string representation of a binary
-- data constructor with name @n@ and arguments @x@ and @y@, in precedence
-- context @d@.
--
-- @since 4.9.0.0
{-# DEPRECATED showsBinary1
      "Use 'showsBinaryWith' to define 'liftShowsPrec'" #-}
showsBinary1 :: (Show1 f, Show1 g, Show a) =>
    String -> Int -> f a -> g a -> ShowS
showsBinary1 name d x y = showParen (d > 10) $
    showString name . showChar ' ' . showsPrec1 11 x .
        showChar ' ' . showsPrec1 11 y

{- $example
These functions can be used to assemble 'Read' and 'Show' instances for
new algebraic types.  For example, given the definition

> data T f a = Zero a | One (f a) | Two a (f a)

a standard 'Read1' instance may be defined as

> instance (Read1 f) => Read1 (T f) where
>     liftReadPrec rp rl = readData $
>         readUnaryWith rp "Zero" Zero <|>
>         readUnaryWith (liftReadPrec rp rl) "One" One <|>
>         readBinaryWith rp (liftReadPrec rp rl) "Two" Two
>     liftReadListPrec = liftReadListPrecDefault

and the corresponding 'Show1' instance as

> instance (Show1 f) => Show1 (T f) where
>     liftShowsPrec sp _ d (Zero x) =
>         showsUnaryWith sp "Zero" d x
>     liftShowsPrec sp sl d (One x) =
>         showsUnaryWith (liftShowsPrec sp sl) "One" d x
>     liftShowsPrec sp sl d (Two x y) =
>         showsBinaryWith sp (liftShowsPrec sp sl) "Two" d x y

-}


text:
-- | The method of moments can be used to estimate a number of commonly used distributions.  This module is still under construction as I work out the best way to handle morphisms from the Moments3 type to types of other distributions.  For more information, see the wikipedia entry: <https://en.wikipedia.org/wiki/Method_of_moments_(statistics)>

module HLearn.Models.Distributions.Univariate.Geometric
    ( Geometric
    )
    where

import Control.DeepSeq
import GHC.TypeLits
import qualified Data.Vector.Unboxed as U
import Data.Vector.Unboxed.Deriving

import qualified Statistics.Distribution as S
import Statistics.Distribution.Geometric

import HLearn.Algebra
import HLearn.Models.Distributions.Common
import HLearn.Models.Distributions.Univariate.Internal.Moments
import HLearn.Models.Distributions.Visualization.Gnuplot

-------------------------------------------------------------------------------
-- Geometric

newtype Geometric sample prob = Geometric {  moments :: (Moments3 sample) }
    deriving (Read,Show,Eq,Ord,Monoid,Group)

instance (Num sample) => HomTrainer (Geometric sample prob) where
    type Datapoint (Geometric sample prob) = sample
    train1dp dp = Geometric $ train1dp dp

instance (Num sample) => Probabilistic (Geometric sample prob) where
    type Probability (Geometric sample prob) = prob

instance (Integral sample, Floating prob) => PDF (Geometric sample prob) where
    pdf dist dp = p*(1-p)^^dp
        where
            p = geo_p dist

instance
    ( PDF (Geometric sample prob)
    , Show prob
    , Show sample
    , Ord sample
    , Ord prob
    , Fractional prob
    , RealFrac prob
    , Integral sample
    ) => PlottableDistribution (Geometric sample prob)
        where

    plotType _ = Points

    samplePoints dist = [min..max]
        where
            min = 0
            max = maximum [20,round $ 3*(fromIntegral $ mean dist)]

geo_p :: (Fractional prob, Integral sample) => Geometric sample prob -> prob
geo_p (Geometric dist) = 1/((fromIntegral $ m1 dist)/(fromIntegral $ m0 dist) +1)

instance (Integral sample, RealFrac prob) => Mean (Geometric sample prob) where
    mean dist = round $ 1/(geo_p dist)

instance (Integral sample, Fractional prob) => Variance (Geometric sample prob) where
    variance dist = (1-p)/(p*p)
        where
            p = geo_p dist

ext: hs


text:
module Y2017.M02.D24.Exercise where

-- below imports from cabal package 'statistics'

import Statistics.Distribution
import Statistics.Distribution.Binomial

{--
AHA!

Whenever you think you have to invent something, stop, pause, and think, ...

SOMEBODY ELSE HAS ALREADY DONE IT! USE THE DARK SIDE OF THE FORCE AND TOGETHER
WE CAN RULE THE GALAXY AS FATHER AND SON!

See? All you need is love, a psychopathic father that murders young children
at the Jedi Academy, and 'midichlorians.' Lots of 'midichlorians.'

Where was I?

Ah, yes: 'already invented.'

So, we know from Rosalind.Traits that any descendant from an Aa Bb/Aa Bb
pairing will have a 25% probability of itself being Aa Bb (you can verify
this if you'd like), so the question arises, what is the probability of at
least n offspring at generation x (or Generation X, if you prefer) also being
Aa Bb.

So.

Today's Haskell problem: with the help of the above imports, create a binomial
probability distribution at 25%, run a scenario where you are in generation
x, where the xth generation has x^2 offspring, and you are looking for at least
n of the offspring to have the desired characteristic (in the abstract case,
the what the characteristic is doesn't matter, we need what the probably outcome
of having it.
--}

type Trials = Int

aAbB :: Trials -> BinomialDistribution
aAbB trials = undefined

genx :: Int -> BinomialDistribution -> Double
genx atleastN bin = undefined

-- what is the result for gen = 2, n = at least 1?

-- Now solve for gen = 7, n = at least 35.

-- hint: what is 'trials' for gen = x?
-- hint-hint: remember solving for atLeast n outcomes from your stats classes?


text:
{-# LANGUAGE ConstraintKinds #-}

module Quadrature
( getGaussPoints
, numGaussPoints
, integrate
, integrateFunc
) where

import qualified Numeric.LinearAlgebra.HMatrix as HMat

type FrElNuFi a = (Fractional a,HMat.Element a,HMat.Numeric a,HMat.Field a)
type FlElNuFi a = (Floating a,HMat.Element a,HMat.Numeric a,HMat.Field a)

-- Function for computing the required number of Gauss points for a polynomial
numGaussPoints :: Integral a => a -> a
numGaussPoints order = ceiling (0.5 * (fromIntegral order + 1.0))

-- Function for getting Gauss points for integration
-- Given the number of points it returns a list of the points and the weights
getGaussPoints :: Floating a => Int -> ([a], [a])
getGaussPoints 1 = ([0.0], [2.0])
getGaussPoints 2 = ([-1.0/sqrt 3.0, 1.0/sqrt 3.0], [1.0, 1.0])
getGaussPoints 3 = ([-sqrt (3.0/5.0), 0.0, sqrt (3.0/5.0)], [5.0/9.0, 8.0/9.0, 5.0/9.0])
getGaussPoints 4 = ([-0.861136311594053,-0.339981043584856,0.339981043584856,0.861136311594053],[0.347854845137454,0.652145154862546,0.652145154862546,0.347854845137454])
getGaussPoints _ = error "Gauss points only up to 4 point quadrature available."

-- Function for integrating some integrand over an element
-- Usage: integrate dim ngpts integrand
integrate :: FrElNuFi a => Int -> Int -> ([a] -> HMat.Matrix a) -> HMat.Matrix a
integrate 1 ngpts integrand = foldl1 HMat.add $ zipWith HMat.scale wts (map (\x -> integrand [x]) xi)
  where
    gpts = getGaussPoints ngpts
    xi = fst gpts
    wts = snd gpts
integrate _ _ _ = error "Integration only works in 1d for now."

-- Function for integrating a function from -1 to 1
integrateFunc :: Floating a => Int -> (a -> a) -> a
integrateFunc ngpts func = sum $ zipWith (*) wts (map func pts)
  where
    gpts = getGaussPoints ngpts
    pts = fst gpts
    wts = snd gpts


text:
{-# LANGUAGE RecordWildCards #-}

module Main where

import           Control.Monad                  ( when
                                                , unless
                                                )
import qualified Data.Text.IO                  as TIO
import qualified Data.ByteString.Lazy          as BL
                                                ( readFile )
import           Data.Csv                       ( decodeByName )
import           QuoteData
import           Statistics
import           StatReport
import           Charts
import           Params
import           Data.Bool
import           HtmlReport

main :: IO ()
main = cmdLineParser >>= work

work :: Params -> IO ()
work params = do
  csvData <- BL.readFile (fname params)
  case decodeByName csvData of
    Left  err         -> putStrLn err
    Right (_, quotes) -> generateReports params quotes

generateReports :: (Functor t, Foldable t) => Params -> t QuoteData -> IO ()
generateReports Params {..} quotes = do
  TIO.putStr $ statReport statInfo'
  when prices $ plotChart title quotes [Open, Close, High, Low] fname_prices
  when volumes $ plotChart title quotes [Volume] fname_volumes
  unless no_text $ TIO.putStr $ statReport statInfo'
  when html $ BL.writeFile fname_html $ htmlReport title quotes statInfo' images
 where
  statInfo' = statInfo quotes
  withCompany pref = if company /= "" then pref ++ company else ""
  img_suffix    = withCompany "_" ++ ".svg"
  fname_prices  = "prices" ++ img_suffix
  fname_volumes = "volumes" ++ img_suffix
  title         = "Historical Quotes" ++ withCompany " for "
  images        = concat
    $ zipWith (bool []) [[fname_prices], [fname_volumes]] [prices, volumes]
  fname_html = "report" ++ withCompany "_" ++ ".html"


text:
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Strict       #-}
module STC.PowerMethod where

import           Control.Monad                  as M
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX as A
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.FourierSeries
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           Graphics.Gnuplot.Simple
import           Image.IO
import           STC.CompletionField
import           STC.Convolution
import           STC.DFTArray
import           STC.Multiplication
import           STC.Plan
import           STC.Reversal
import           STC.Utils
import           System.FilePath                ((</>))
import           Text.Printf
import           Utils.Array
import           Utils.List
import           Utils.Parallel
import           Utils.Time
import FourierPinwheel as FP
import           Foreign.CUDA.Driver            as CUDA

powerMethod ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethod _ _ _ _ _ _ 0 !arr = return arr
powerMethod !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numStep !input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  convolvedArr <-
    STC.Convolution.convolve Source plan coefficients harmonicsArray input
  let !biasedConvolvedArr = parMapDFTArray (VS.zipWith (*) bias) convolvedArr
      !maxMag =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x)^2)) . getDFTArrayVector $
        biasedConvolvedArr
      !normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (maxMag :+ 0))) biasedConvolvedArr
  when writeFlag .
    plotDFTArrayPower
      (folderPath </> (printf "Source_%03d.png" numStep))
      cols
      rows $
    convolvedArr
  powerMethod
    plan
    folderPath
    writeFlag
    coefficients
    harmonicsArray
    bias
    (numStep - 1)
    normalizedBiasedConvolvedArr

{-# INLINE computeContour #-}
computeContour ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> String
  -> DFTArray
  -> IO DFTArray
computeContour !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numIteration !suffix !input@(DFTArray rows cols _ _ _) = do
  eigenVector <-
    powerMethod
      plan
      folderPath
      writeFlag
      coefficients
      harmonicsArray
      bias
      numIteration
      input
  source <- STC.Convolution.convolve Source plan coefficients harmonicsArray eigenVector
  sink <- STC.Convolution.convolve Sink plan coefficients harmonicsArray eigenVector
  completion <- completionField plan source sink
  plotDFTArrayPower (folderPath </> "SinkPower.png") rows cols sink
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols
    completion
  return completion

powerMethod' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int -> [[(Complex Double)]]
  -> DFTArray
  -> IO DFTArray
powerMethod' _ _ _ _ _ _ 0 _ !arr = return arr
powerMethod' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numStep thetaRHarmonics !input@(DFTArray rows cols thetaFreqs rFreqs _) = do
  printCurrentTime (show numStep)
  convolvedArr <- convolve' Source plan coefficients harmonicsArray input
  let -- sigma = 15
      -- gaussian =
      --   L.map (\x -> (exp (-(x ^ 2) / (2 * sigma ^ 2))) :+ 0) thetaFreqs
      convolvedArr' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $
        convolvedArr
      -- convolvedArr1 =
      --   DFTArray rows cols thetaFreqs rFreqs .
      --   L.zipWith (\g vec -> VS.map (* g) vec) gaussian . getDFTArrayVector $
      --   convolvedArr
  plotThetaDimension folderPath (printf "FourierSeries_%d_" numStep) (29, 12) .
    R.map magnitude $
    convolvedArr'
  let !biasedConvolvedArr = parMapDFTArray (VS.zipWith (*) bias) convolvedArr
      !maxMag =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x) ^ 2)) .
        getDFTArrayVector $
        biasedConvolvedArr
      !normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (maxMag :+ 0))) biasedConvolvedArr
  when writeFlag .
    plotDFTArrayPower
      (folderPath </> (printf "Source_%03d.png" numStep))
      cols
      rows $
    convolvedArr
  powerMethod'
    plan
    folderPath
    writeFlag
    coefficients
    harmonicsArray
    bias
    (numStep - 1)
    thetaRHarmonics
    normalizedBiasedConvolvedArr

{-# INLINE computeContour' #-}
computeContour' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> String -> [[Complex Double]]
  -> DFTArray
  -> IO DFTArray
computeContour' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !numIteration !suffix thetaRHarmonics !input@(DFTArray rows cols _ _ _) = do
  eigenVector <-
    powerMethod'
      plan
      folderPath
      writeFlag
      coefficients
      harmonicsArray
      bias
      numIteration
      thetaRHarmonics
      input
  source <- convolve' Source plan coefficients harmonicsArray eigenVector
  sink <- convolve' Sink plan coefficients harmonicsArray eigenVector
  completion <- completionField' plan source sink
  plotDFTArrayPower (folderPath </> "Sink.png") rows cols sink
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols
    completion
  let point = (32, -4)
      source' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics .
        getDFTArrayVector $
        source
      sink' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics .
        getDFTArrayVector $
        sink
      completion' =
        fromUnboxed (Z :. 72 :. cols :. rows) .
        VS.convert .
        VS.concat .
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $
        completion
  plotThetaDimension folderPath ("FourierSeries_Completion_") point .
    R.map magnitude $
    completion'
  plotThetaDimension
    folderPath
    ("FourierSeries_Source_")
    point .
    R.map magnitude $
    source'
  plotThetaDimension
    folderPath
    ("FourierSeries_Sink_")
    point .
    R.map magnitude $
    sink'
  return completion

powerMethodG' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> VS.Vector (Complex Double)
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethodG' _ _ _ _ _ _ _ 0 !arr = return arr
powerMethodG' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !gaussian !numStep !input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  convolvedArr <- convolveG' Source plan coefficients harmonicsArray gaussian input
  let !biasedConvolvedArr = parMapDFTArray (VS.zipWith (*) bias) convolvedArr
      !maxMag =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x)^2)) .
        getDFTArrayVector $
        biasedConvolvedArr
      !normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (maxMag :+ 0))) biasedConvolvedArr
  when writeFlag .
    plotDFTArrayPower
      (folderPath </> (printf "Source_%03d.png" numStep))
      cols
      rows $
    convolvedArr
  powerMethodG'
    plan
    folderPath
    writeFlag
    coefficients
    harmonicsArray
    bias
    gaussian
    (numStep - 1)
    normalizedBiasedConvolvedArr

{-# INLINE computeContourG' #-}
computeContourG' ::
     DFTPlan
  -> FilePath
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> VS.Vector (Complex Double)
  -> Int
  -> String
  -> DFTArray
  -> IO DFTArray
computeContourG' !plan !folderPath !writeFlag !coefficients !harmonicsArray !bias !gaussian !numIteration !suffix !input@(DFTArray rows cols _ _ _) = do
  eigenVector <-
    powerMethodG'
      plan
      folderPath
      writeFlag
      coefficients
      harmonicsArray
      bias
      gaussian
      numIteration
      input
  source <- convolveG' Source plan coefficients harmonicsArray gaussian eigenVector
  sink <- convolveG' Sink plan coefficients harmonicsArray gaussian eigenVector
  completion <- completionField' plan source sink
  plotDFTArrayPower (folderPath </> "Sink.png") rows cols sink
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols
    completion
  return completion

powerMethodSparse ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> [R.Array U DIM2 (Complex Double)]
  -> [R.Array U DIM2 (Complex Double)]
powerMethodSparse _ _ _ _ _ _ _ 0 !arr = arr
powerMethodSparse !field !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !xs !numStep !input =
  let !convolvedArr =
        convolveSparse
          field
          coefficients
          harmonicsArray
          thetaFreqs
          rFreqs
          cutoff
          xs
          input
      !maxMag =
        sqrt .
        L.sum .
        parMap rdeepseq (VU.sum . VU.map (\x -> (magnitude x) ^ 2) . toUnboxed) $
        convolvedArr
      !normalizedConvolvedArr =
        parMap rseq (computeS . R.map (/ (maxMag :+ 0))) convolvedArr
  in powerMethodSparse
       field
       coefficients
       harmonicsArray
       thetaFreqs
       rFreqs
       cutoff
       xs
       (numStep - 1)
       normalizedConvolvedArr

{-# INLINE computeContourSparse #-}
computeContourSparse ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM2 (Complex Double)]
  -> IO [R.Array U DIM2 (Complex Double)]
  -- -> IO DFTArray
computeContourSparse !plan !folderPath !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
      ys =
        powerMethodSparse
          Source
          coefficients
          harmonicsArray
          thetaFreqs
          rFreqs
          cutoff
          xs
          numIteration
          input
      !eigenSourceSparse =
        sparseArrayToDFTArray
          rows
          cols
          (R.toList thetaFreqs)
          (R.toList rFreqs)
          xs
          ys
  harmonicsArrayDFT <-
    fmap (listArray (bounds harmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    harmonicsArray
  printCurrentTime "Source"
  source <-
    STC.Convolution.convolve Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower_Sparse.png") rows cols source
  printCurrentTime "Sink"
  let sink =
        parZipWithDFTArray
          (\vec (rFreq, thetaFreq) -> VS.map (* (cis (thetaFreq * pi))) vec)
          source
          ((,) <$> (R.toList rFreqs) <*> (R.toList thetaFreqs))
  -- sink <- convolve Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SinkPower_Sparse.png") rows cols sink
  printCurrentTime "Completion"
  completion <- completionField plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> ("CompletionPower_Sparse.png"))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return ys

{-# INLINE computeContourSparseG #-}
computeContourSparseG ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> VS.Vector (Complex Double)
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM2 (Complex Double)]
  -> IO DFTArray
computeContourSparseG !plan !folderPath !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !gaussian !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  blurredVecF <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VU.convert . toUnboxed) . IA.elems $
    harmonicsArray
  blurredHarmonicsArray <-
    (listArray (bounds harmonicsArray) .
     L.map
       (\vec ->
          let arr = fromUnboxed (Z :. cols :. rows) . VS.convert $ vec
          in computeS . R.traverse arr id $ \f idx@(Z :. i :. j) ->
               let r =
                     sqrt . fromIntegral $
                     (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2
               in if r <= 0
                    then 0
                    else f idx)) <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) blurredVecF
  let !eigenSourceSparse =
        sparseArrayToDFTArray
          rows
          cols
          (R.toList thetaFreqs)
          (R.toList rFreqs)
          xs .
        powerMethodSparse
          Source
          coefficients
          blurredHarmonicsArray
          thetaFreqs
          rFreqs
          cutoff
          xs
          numIteration $
        input
  harmonicsArrayDFT <-
    fmap (listArray (bounds blurredHarmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    blurredHarmonicsArray
  printCurrentTime "Source"
  source <-
    STC.Convolution.convolve Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower.png") rows cols source
  printCurrentTime "Sink"
  let sink =
        parZipWithDFTArray
          (\vec (rFreq, thetaFreq) -> VS.map (* (cis (thetaFreq * pi))) vec)
          source
          ((,) <$> (R.toList rFreqs) <*> (R.toList thetaFreqs))
  -- sink <- convolve Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SinkPower.png") rows cols sink
  printCurrentTime "Completion"
  completion <- completionField plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return completion

powerMethodSparse' ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> [R.Array U DIM1 (Complex Double)]
  -> IO [R.Array U DIM1 (Complex Double)]
powerMethodSparse' _ _ _ _ _ _ _ 0 !arr = return arr
powerMethodSparse' !field !coefficients !harmonicsArray !phiFreqs !rhoFreqs  !cutoff !xs !numStep !input =
  let !convolvedArr =
        convolveSparse'
          field
          coefficients
          harmonicsArray
          phiFreqs
          rhoFreqs
          cutoff
          xs
          input
      !maxMag =
        sqrt . L.sum .
        parMap rdeepseq (VU.sum . VU.map (\x -> (magnitude x) ^ 2) . toUnboxed) $
        convolvedArr
      !normalizedConvolvedArr =
        parMap rseq (computeS . R.map (\x -> x / (maxMag :+ 0))) convolvedArr
  in do
        powerMethodSparse'
          field
          coefficients
          harmonicsArray
          phiFreqs
          rhoFreqs
          cutoff
          xs
          (numStep - 1)
          normalizedConvolvedArr

{-# INLINE computeContourSparse' #-}
computeContourSparse' ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> [[Complex Double]]
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM1 (Complex Double)]
  -> IO [R.Array U DIM1 (Complex Double)]
  -- -> IO DFTArray
computeContourSparse' !plan !folderPath !coefficients !harmonicsArray thetaRHarmonics !phiFreqs !rhoFreqs !cutoff !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  -- !eigenSourceSparse <-
  --   sparseArrayToDFTArray' rows cols (R.toList phiFreqs) (R.toList rhoFreqs) xs <$>
    -- powerMethodSparse'
    --   Source
    --   coefficients
    --   harmonicsArray
    --   phiFreqs
    --   rhoFreqs
    --   cutoff
    --   xs
    --   numIteration
    --   input
  ys <-
    powerMethodSparse'
      Source
      coefficients
      harmonicsArray
      phiFreqs
      rhoFreqs
      cutoff
      xs
      numIteration
      input
  let eigenSourceSparse =
        sparseArrayToDFTArray'
          rows
          cols
          (R.toList phiFreqs)
          (R.toList rhoFreqs)
          xs
          ys
  print (bounds harmonicsArray)
  print . L.length . IA.elems $ harmonicsArray
  harmonicsArrayDFT <-
    fmap (listArray (bounds harmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    harmonicsArray
  printCurrentTime "Source"
  source <-
    convolve' Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower_Sparse.png") rows cols source
  -- plotDFTArrayThetaR
  --   (folderPath </> "Source.png")
  --   rows
  --   cols
  --   thetaRHarmonics
  --   source
  printCurrentTime "Sink"
  -- sink <- convolve' Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  let sink =
        parZipWithDFTArray
          (\vec thetaFreq -> VS.map (* (cis (thetaFreq * pi))) vec)
          source
          (R.toList phiFreqs)
  plotDFTArrayPower (folderPath </> "SinkPower_Sparse.png") rows cols sink
  -- plotDFTArrayThetaR (folderPath </> "Sink.png") rows cols thetaRHarmonics sink
  printCurrentTime "Completion"
  completion <- completionField' plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower (folderPath </> ("CompletionPower_Sparse.png")) rows cols -- .
    -- parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  -- let point = (31, -4)
  --     source' =
  --       fromUnboxed (Z :. 72 :. cols :. rows) .
  --       VS.convert .
  --       VS.concat .
  --       computeFourierSeriesThetaR thetaRHarmonics .
  --       getDFTArrayVector $
  --       source
  --     sink' =
  --       fromUnboxed (Z :. 72 :. cols :. rows) .
  --       VS.convert .
  --       VS.concat .
  --       computeFourierSeriesThetaR thetaRHarmonics .
  --       getDFTArrayVector $
  --       sink
  --     completion' =
  --       fromUnboxed (Z :. 72 :. cols :. rows) .
  --       VS.convert .
  --       VS.concat .
  --       computeFourierSeriesThetaR thetaRHarmonics .
  --       getDFTArrayVector $
  --       completion
  -- plotThetaDimension
  --   folderPath
  --   ("FourierSeries_Completion_")
  --   point .
  --   R.map magnitude $
  --   completion'
  -- plotThetaDimension
  --   folderPath
  --   ("FourierSeries_Source_")
  --   point .
  --   R.map magnitude $
  --   source'
  -- plotThetaDimension
  --   folderPath
  --   ("FourierSeries_Sink_")
  --   point .
  --   R.map magnitude $
  --   sink'
  -- return completion
  return ys

{-# INLINE computeContourSparseG' #-}
computeContourSparseG' ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> [[Complex Double]]
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> VS.Vector (Complex Double)
  -> Double
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM1 (Complex Double)]
  -> IO DFTArray
computeContourSparseG' !plan !folderPath !coefficients !harmonicsArray !thetaRHarmonics !phiFreqs !rhoFreqs !cutoff !gaussian !std !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  blurredVecF <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VU.convert . toUnboxed) . IA.elems $
    harmonicsArray
  blurredHarmonicsArray <-
    (listArray (bounds harmonicsArray) .
     L.map
       (\vec ->
          let arr = fromUnboxed (Z :. cols :. rows) . VS.convert $ vec
          in computeS . R.traverse arr id $ \f idx@(Z :. i :. j) ->
               let r =
                     sqrt . fromIntegral $
                     (i - div cols 2) ^ 2 + (j - div rows 2) ^ 2
               in if r <= 0
                    then 0
                    else f idx)) <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) blurredVecF
  !eigenSourceSparse <-
    sparseArrayToDFTArray'
      rows
      cols
      (R.toList phiFreqs)
      (R.toList rhoFreqs)
      xs <$>
    powerMethodSparse'
      Source
      coefficients
      blurredHarmonicsArray
      phiFreqs
      rhoFreqs
      cutoff
      xs
      numIteration
      input
  harmonicsArrayDFT <-
    fmap (listArray (bounds blurredHarmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    blurredHarmonicsArray
  printCurrentTime "Source"
  source <-
    convolve' Source plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SourcePower.png") rows cols source
  plotDFTArrayThetaR
    (folderPath </> "Source.png")
    rows
    cols
    thetaRHarmonics
    source
  printCurrentTime "Sink"
  sink <- convolve' Sink plan coefficients harmonicsArrayDFT eigenSourceSparse
  plotDFTArrayPower (folderPath </> "SinkPower.png") rows cols sink
  plotDFTArrayThetaR (folderPath </> "Sink.png") rows cols thetaRHarmonics sink
  printCurrentTime "Completion"
  completion <- completionField' plan source sink
  let mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return completion

-- # INLINE computeContourSparse''' #
computeContourSparse''' ::
     DFTPlan
  -> FilePath
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> [[Complex Double]]
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double -> VS.Vector (Complex Double)
  -> [(Double, Double)]
  -> Int
  -> String
  -> [R.Array U DIM1 (Complex Double)]
  -> IO DFTArray
computeContourSparse''' !plan !folderPath !coefficients !harmonicsArray !thetaRHarmonics !phiFreqs !rhoFreqs !cutoff !gaussian !xs !numIteration !suffix !input = do
  printCurrentTime $ printf "Power Method %d iterations" numIteration
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
  blurredVecF <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VU.convert . toUnboxed) . IA.elems $
    harmonicsArray
  blurredHarmonicsArray <-
    (listArray (bounds harmonicsArray) .
     L.map
       (\vec ->
          let arr = fromUnboxed (Z :. cols :. rows) . VS.convert $ vec
          in computeS . R.traverse arr id $ \f idx@(Z :. i :. j) ->
               if (i == div cols 2) && (j == div rows 2)
                 then 0
                 else f idx)) <$>
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) blurredVecF
  !eigenSourceSparse <-
    sparseArrayToDFTArray'''
      rows
      cols
      (R.toList phiFreqs)
      (R.toList rhoFreqs)
      xs <$>
    powerMethodSparse'
      Source
      coefficients
      blurredHarmonicsArray
      phiFreqs
      rhoFreqs
      cutoff
      xs
      numIteration
      input
  harmonicsArrayDFT <-
    fmap (listArray (bounds blurredHarmonicsArray)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) .
    L.map (VS.convert . toUnboxed . computeUnboxedS . makeFilter2D) . IA.elems $
    blurredHarmonicsArray
  xs <-
    M.mapM
      (\(i, (p, q)) -> do
         source <- convolve' Source plan coefficients harmonicsArrayDFT q
         plotDFTArrayPower
           (folderPath </> (printf "SourcePower_%d.png" i))
           rows
           cols
           source
         sink <- convolve' Sink plan coefficients harmonicsArrayDFT p
         plotDFTArrayPower
           (folderPath </> (printf "SinkPower_%d.png" i))
           rows
           cols
           sink
         completion <- completionField' plan source sink
         plotDFTArrayPower
           (folderPath </> (printf "CompletionPower_%d.png" i))
           rows
           cols
           completion
         when
           (i == 0)
           (do let point = (31, -4)
                   source' =
                     fromUnboxed (Z :. 72 :. cols :. rows) .
                     VS.convert .
                     VS.concat .
                     computeFourierSeriesThetaR thetaRHarmonics .
                     getDFTArrayVector $
                     source
                   sink' =
                     fromUnboxed (Z :. 72 :. cols :. rows) .
                     VS.convert .
                     VS.concat .
                     computeFourierSeriesThetaR thetaRHarmonics .
                     getDFTArrayVector $
                     sink
                   completion' =
                     fromUnboxed (Z :. 72 :. cols :. rows) .
                     VS.convert .
                     VS.concat .
                     computeFourierSeriesThetaR thetaRHarmonics .
                     getDFTArrayVector $
                     completion
               plotThetaDimension
                 folderPath
                 ("FourierSeries_Completion_")
                 point .
                 R.map magnitude $
                 completion'
               plotThetaDimension
                 folderPath
                 ("FourierSeries_Source_")
                 point .
                 R.map magnitude $
                 source'
               plotThetaDimension
                 folderPath
                 ("FourierSeries_Sink_")
                 point .
                 R.map magnitude $
                 sink')
         return completion) .
    L.zip [0 :: Int ..] $
    eigenSourceSparse
  let (DFTArray rows cols thetaFreqs rFreqs _) = L.head xs
      completion =
        DFTArray rows cols thetaFreqs rFreqs .
        L.foldl1' (L.zipWith (VS.zipWith (+))) .
        L.map
          (\arr ->
             let ys = getDFTArrayVector arr
                 s = L.maximum . L.map (VS.maximum . VS.map magnitude) $ ys
             in L.map (VS.map (/ (s :+ 0))) ys) $
        xs
      mm =
        L.maximum . L.map (VS.maximum . VS.map magnitude) . getDFTArrayVector $
        completion
  print mm
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower_%s.png" suffix))
    rows
    cols .
    parMapDFTArray (VS.map (/ (mm :+ 0))) $
    completion
  return completion

powerMethodPinwheelBasis' ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethodPinwheelBasis' _ _ _ _ _ _ _ _ _ _ _ _ 0 arr = return arr
powerMethodPinwheelBasis' plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numPoints numR2Freqs delta periodR2 numBatch numStep input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  let convolvedArr = convoluvePinhweelBasis' coefficients harmonicsArray input
  biasedConvolvedArr <-
    multiplyPinwheelBasisBatch1 plan numBatch dftBias convolvedArr
  let s =
        L.maximum .
        parMap rdeepseq (VS.maximum . VS.map (\x -> (magnitude x) )) .
        getDFTArrayVector $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
      -- envelope =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
      --     let r =
      --           sqrt . fromIntegral $
      --           (i - div numR2Freqs 2) ^ 2 + (j - div numR2Freqs 2) ^ 2
      --     in if r == 0
      --          then 0
      --          else r ** (-0.75) :+ 0
      -- normalizedBiasedConvolvedArr =
      --   parMapDFTArray (VS.zipWith (*) envelope) normalizedBiasedConvolvedArr'
  when
    writeFlag
    (do let eigenMat =
              A.transpose . A.use . fromDFTArray . getDFTArrayVector $
              convolvedArr
        eigenR2 <-
          computeFourierSeriesR2StreamAcc
            ptxs
            numR2Freqs
            numPoints
            (L.length . getDFTArrayVector $ convolvedArr)
            periodR2
            delta
            numBatch
            eigenMat
        plotImageRepa (folderPath </> (printf "Source_%03d.png" numStep)) .
          ImageRepa 8 .
          fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
          -- VU.map sqrt .
          toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
          eigenR2
        -- let biasMat =
        --       A.transpose . A.use . fromDFTArray . getDFTArrayVector $
        --       normalizedBiasedConvolvedArr
        -- biasR2 <-
        --   computeFourierSeriesR2StreamAcc
        --     ptxs
        --     numR2Freqs
        --     numPoints
        --     (L.length . getDFTArrayVector $ convolvedArr)
        --     periodR2
        --     delta
        --     1
        --     biasMat
        -- plotImageRepa (folderPath </> (printf "Biased_%03d.png" numStep)) .
        --   ImageRepa 8 .
        --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
        --   VU.map sqrt .
        --   toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
        --   biasR2
    )
  powerMethodPinwheelBasis'
    plan
    folderPath
    ptxs
    writeFlag
    coefficients
    harmonicsArray
    dftBias
    numPoints
    numR2Freqs
    delta
    periodR2
    numBatch
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourPinwheelBasis' ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO (R.Array U DIM4 (Complex Double))
computeContourPinwheelBasis' plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numStep numBatch numPoints numR2Freqs delta periodR2 thetaFreqs rFreqs input = do
  eigenSource' <-
    powerMethodPinwheelBasis'
      plan
      folderPath
      ptxs
      writeFlag
      coefficients
      harmonicsArray
      dftBias
      numPoints
      numR2Freqs
      delta
      periodR2
      numBatch
      numStep
      input
  let eigenSource =
        convoluvePinhweelBasis' coefficients harmonicsArray eigenSource'
      numLogPolarFreqs = thetaFreqs * rFreqs
      eigenSourceMat =
        A.transpose . A.use . fromDFTArray . getDFTArrayVector $ eigenSource
  eigenSourceR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freqs
      numPoints
      numLogPolarFreqs
      periodR2
      delta
      numBatch
      eigenSourceMat
  let eigenSourceR2 =
        R.reshape (Z :. rFreqs :. thetaFreqs :. numPoints :. numPoints) $
        eigenSourceR2'
      eigenSinkR2 =
        timeReversalRepa
          (L.map fromIntegral . getListFromNumber $ thetaFreqs)
          eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  plotImageRepa (folderPath </> "Source.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    -- VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSourceR2
  plotImageRepa (folderPath </> "Sink.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    -- VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSinkR2
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    -- VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  return completionR2

powerMethodPinwheelBasis ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO DFTArray
powerMethodPinwheelBasis _ _ _ _ _ _ _ _ _ _ _ _ 0 arr = return arr
powerMethodPinwheelBasis plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numPoints numR2Freqs delta periodR2 numBatch numStep input@(DFTArray rows cols _ _ _) = do
  printCurrentTime (show numStep)
  let convolvedArr = convoluvePinhweelBasis coefficients harmonicsArray input
  biasedConvolvedArr <-
    multiplyPinwheelBasisBatch1 plan numBatch dftBias convolvedArr
  let s =
        L.maximum .
        parMap rdeepseq (VS.maximum . VS.map (\x -> (magnitude x) )) .
        getDFTArrayVector $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapDFTArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
      -- envelope =
      --   VU.convert .
      --   toUnboxed . computeS . fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
      --     let r =
      --           sqrt . fromIntegral $
      --           (i - div numR2Freqs 2) ^ 2 + (j - div numR2Freqs 2) ^ 2
      --     in if r == 0
      --          then 0
      --          else r ** (-0.75) :+ 0
      -- normalizedBiasedConvolvedArr =
      --   parMapDFTArray (VS.zipWith (*) envelope) normalizedBiasedConvolvedArr'
  when
    writeFlag
    (do let eigenMat =
              A.transpose . A.use . fromDFTArray . getDFTArrayVector $
              convolvedArr
        eigenR2 <-
          computeFourierSeriesR2StreamAcc
            ptxs
            numR2Freqs
            numPoints
            (L.length . getDFTArrayVector $ convolvedArr)
            periodR2
            delta
            numBatch
            eigenMat
        plotImageRepa (folderPath </> (printf "Source_%03d.png" numStep)) .
          ImageRepa 8 .
          fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
          VU.map sqrt .
          toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
          eigenR2
        let biasMat =
              A.transpose . A.use . fromDFTArray . getDFTArrayVector $
              normalizedBiasedConvolvedArr
        biasR2 <-
          computeFourierSeriesR2StreamAcc
            ptxs
            numR2Freqs
            numPoints
            (L.length . getDFTArrayVector $ convolvedArr)
            periodR2
            delta
            numBatch
            biasMat
        plotImageRepa (folderPath </> (printf "Biased_%03d.png" numStep)) .
          ImageRepa 8 .
          fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
          VU.map sqrt .
          toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
          biasR2
    )
  powerMethodPinwheelBasis
    plan
    folderPath
    ptxs
    writeFlag
    coefficients
    harmonicsArray
    dftBias
    numPoints
    numR2Freqs
    delta
    periodR2
    numBatch
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourPinwheelBasis ::
     DFTPlan
  -> FilePath
  -> [PTX]
  -> Bool
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> DFTArray
  -> IO (R.Array U DIM4 (Complex Double))
computeContourPinwheelBasis plan folderPath ptxs writeFlag coefficients harmonicsArray dftBias numStep numBatch numPoints numR2Freqs delta periodR2 thetaFreqs rFreqs input = do
  eigenSource' <-
    powerMethodPinwheelBasis
      plan
      folderPath
      ptxs
      writeFlag
      coefficients
      harmonicsArray
      dftBias
      numPoints
      numR2Freqs
      delta
      periodR2
      numBatch
      numStep
      input
  let eigenSource =
        convoluvePinhweelBasis coefficients harmonicsArray eigenSource'
      numLogPolarFreqs = thetaFreqs * rFreqs
      eigenSourceMat =
        A.transpose . A.use . fromDFTArray . getDFTArrayVector $ eigenSource
  eigenSourceR2' <-
    computeFourierSeriesR2StreamAcc
      ptxs
      numR2Freqs
      numPoints
      numLogPolarFreqs
      periodR2
      delta
      numBatch
      eigenSourceMat
  let eigenSourceR2 =
        R.reshape (Z :. rFreqs :. thetaFreqs :. numPoints :. numPoints) $
        eigenSourceR2'
      eigenSinkR2 =
        timeReversalRepa
          (L.map fromIntegral . getListFromNumber $ thetaFreqs)
          eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  plotImageRepa (folderPath </> "Source.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSourceR2
  plotImageRepa (folderPath </> "Sink.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSinkR2
  plotImageRepa (folderPath </> "Completion.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  return completionR2

powerMethodFourierPinwheel ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
powerMethodFourierPinwheel _ _ _ _ _ _ _ _ _ 0 arr = return arr
powerMethodFourierPinwheel plan folderPath writeFlag harmonicsArray dftBias numPoints delta periodR2 numBatch numStep input = do
  printCurrentTime (show numStep)
  convolvedArr <- FP.convolve harmonicsArray input
  biasedConvolvedArr <- multiplyBias4D plan dftBias convolvedArr
  let s =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x) ^ 2)) . getFPArray $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapFPArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
  when
    (writeFlag && mod numStep 4 == 0)
    (do _ <-
          plotFPArray
            plan
            (folderPath </> (printf "Source_%03d.png" numStep))
            convolvedArr
        -- _ <- plotFPArray
        --        plan
        --        (folderPath </> (printf "Bias_%03d.png" numStep))
        --        normalizedBiasedConvolvedArr
        -- plotFPArrayFreqency
        --   (folderPath </> (printf "SourceFreq_%03d.png" numStep))
        --   convolvedArr
        -- plotFPArrayFreqency
        --   (folderPath </> (printf "BiasFreq1_%03d.png" numStep))
        --   normalizedBiasedConvolvedArr'
        -- plotFPArrayFreqency
        --   (folderPath </> (printf "BiasFreq2_%03d.png" numStep))
        --   normalizedBiasedConvolvedArr
        return ())
  powerMethodFourierPinwheel
    plan
    folderPath
    writeFlag
    harmonicsArray
    dftBias
    numPoints
    delta
    periodR2
    numBatch
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourFourierPinwheel ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double -> Double
  -> FPArray (VS.Vector (Complex Double))
  -> String
  -> [Int]
  -> IO (R.Array U DIM4 (Complex Double))
computeContourFourierPinwheel plan folderPath writeFlag harmonicsArray dftBias numStep numBatch numPoints delta periodR2 periodEnv input suffix deviceIDs = do
  eigenSource' <-
    powerMethodFourierPinwheel
      plan
      folderPath
      writeFlag
      harmonicsArray
      dftBias
      numPoints
      delta
      periodR2
      numBatch
      numStep
      input
  eigenSource <- FP.convolve harmonicsArray eigenSource'
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  -- eigenSourceR2 <-
  --   plotFPArray plan (folderPath </> printf "Source_%s.png" suffix) eigenSource
  eigenSourceR2 <-
    plotFPArrayAcc
      ptxs
      (folderPath </> printf "Source_%s.png" suffix)
      numPoints
      delta
      periodR2
      numBatch
      eigenSource
  let numRFreq = getFPArrayNumRFreq eigenSource
      numThetaFreq = getFPArrayNumThetaFreq eigenSource
  eigenSinkR2 <-
    computeUnboxedP $
    timeReversalRepa
      (L.map fromIntegral . getListFromNumber $ numThetaFreq)
      eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  -- plotImageRepa (folderPath </> printf "Source_%s.png" suffix) .
  --   ImageRepa 8 .
  --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
  --   -- VU.map (\x -> x^2) .
  --   toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
  --   eigenSourceR2
  (sumP . R.map (\x -> magnitude x ** 2) . rotate4D2 $ eigenSinkR2) >>= sumP >>=
    plotImageRepa (folderPath </> printf "Sink_%s.png" suffix) .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All)
  completionR2' <-
    (sumP . R.map (\x -> magnitude x ** 2) . rotate4D2 $ completionR2) >>= sumP
  plotImageRepa (folderPath </> printf "Completion_%s.png" suffix) .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) $
    completionR2'
  plotImageRepa (folderPath </> printf "Completion1_%s.png" suffix) .
    ImageRepa 8 . computeS . extend (Z :. (1 :: Int) :. All :. All) . R.map sqrt $
    completionR2'
  return completionR2

powerMethodFourierPinwheelDiscrete ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Double
  -> Double
  -> Int
  -> (Int, Int)
  -> Int
  -> FPArray (VS.Vector (Complex Double))
  -> IO (FPArray (VS.Vector (Complex Double)))
powerMethodFourierPinwheelDiscrete  _ _ _ _ _ _ _ _ _ idx 0 arr = return arr
powerMethodFourierPinwheelDiscrete plan folderPath writeFlag harmonicsArray dftBias numPoints delta periodR2 numBatch idx numStep input' = do
  printCurrentTime (show numStep)
  -- input <- multiplyRFunction plan (periodR2^2 / 2) input'
  convolvedArr <- FP.convolve harmonicsArray input'
  biasedConvolvedArr <- multiplyBiasDiscrete plan dftBias convolvedArr
  let s =
        L.sum .
        parMap rdeepseq (VS.sum . VS.map (\x -> (magnitude x) ^ 2)) . getFPArray $
        biasedConvolvedArr
      normalizedBiasedConvolvedArr =
        parMapFPArray (VS.map (/ (s :+ 0))) biasedConvolvedArr
  when
    (writeFlag && mod numStep 2 == 0
    )
    (do arrR2 <-
          plotFPArray
            plan
            (folderPath </> (printf "Source_%03d.png" numStep))
            convolvedArr
        plotRThetaDist
          (folderPath </> (printf "ThetaDist_%03d.png" numStep))
          (folderPath </> (printf "RDist_%03d.png" numStep))
          (getFPArrayNumXFreq convolvedArr)
          72
          100
          periodR2
          idx
          arrR2
        _ <- plotFPArray
               plan
               (folderPath </> (printf "Bias_%03d.png" numStep))
               normalizedBiasedConvolvedArr
        return ()
        -- let biasMat = A.transpose . toMatrixAcc $ normalizedBiasedConvolvedArr
        -- biasR2 <-
        --   computeFourierSeriesR2StreamAcc
        --     ptxs
        --     (getFPArrayNumXFreq convolvedArr)
        --     numPoints
        --     (getFPArrayNumRFreq convolvedArr *
        --      getFPArrayNumThetaFreq convolvedArr)
        --     periodR2
        --     delta
        --     numBatch
        --     biasMat
        -- plotImageRepa (folderPath </> (printf "Biased_%03d.png" numStep)) .
        --   ImageRepa 8 .
        --   fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
        --   VU.map sqrt .
        --   toUnboxed . sumS . R.map (\x -> (magnitude x) ** 2) . rotate3D $
        --   biasR2
     )
  powerMethodFourierPinwheelDiscrete
    plan
    folderPath
    writeFlag
    harmonicsArray
    dftBias
    numPoints
    delta
    periodR2
    numBatch
    idx
    (numStep - 1)
    normalizedBiasedConvolvedArr

computeContourFourierPinwheelDiscrete ::
     DFTPlan
  -> FilePath
  -> Bool
  -> FPData (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> FPArray (VS.Vector (Complex Double))
  -> String
  -> (Int, Int)
  -> IO (R.Array U DIM4 (Complex Double))
computeContourFourierPinwheelDiscrete plan folderPath writeFlag harmonicsArray dftBias numStep numBatch numPoints delta periodR2 input suffix idx = do
  eigenSource' <-
    powerMethodFourierPinwheelDiscrete
      plan
      folderPath
      writeFlag
      harmonicsArray
      dftBias
      numPoints
      delta
      periodR2
      numBatch
      idx
      numStep
      input
  eigenSource <- FP.convolve harmonicsArray eigenSource'
  eigenSourceR2 <- plotFPArray plan (folderPath </> printf "Source_%s.png" suffix)  eigenSource
  let numThetaFreq = getFPArrayNumThetaFreq eigenSource
      eigenSinkR2 =
        timeReversalRepa
          (L.map fromIntegral . getListFromNumber $ numThetaFreq)
          eigenSourceR2
  completionR2 <- completionFieldRepa plan eigenSourceR2 eigenSinkR2
  plotImageRepa (folderPath </> printf "Sink_%s.png" suffix) .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    eigenSinkR2
  plotImageRepa (folderPath </> printf "Completion_%s.png" suffix) .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoints :. numPoints) .
    VU.map sqrt .
    toUnboxed . sumS . sumS . R.map (\x -> (magnitude x) ** 2) . rotate4D2 $
    completionR2
  return completionR2

max_stars_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_issues_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_forks_repo_name: XinhuaZhang/Stochastic-Completion-Field