{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}

module Vector where

import Tensor
import Data.Monoid
import Data.Maybe
import Data.Complex
import Helper
import Data.Either
import System.Random

data Vector a = Scalar a | Row [Vector a] | Column [Vector a]

-- This Eq implementation ignores structural differences
instance (Eq a) => Eq (Vector a) where
    x == y = wx == wy && hx == hy && go 0 0
        where
            wx = width x
            wy = width y
            hx = height x
            hy = height y
            go i j
                | i == hx = True
                | j == wx = go (i + 1) 0
                | otherwise = atM x (i, j) == atM y (i, j) && go i (j + 1)

instance Functor Vector where
    fmap f (Scalar s) = Scalar $ f s
    fmap f (Row t) = Row $ map (fmap f) t
    fmap f (Column t) = Column $ map (fmap f) t

instance Applicative Vector where
    pure = Scalar
    Scalar s <*> t = fmap s t
    t <*> (Scalar s) = fmap ($s) t
    Row r <*> t = Row $ map (<*> t) r
    Column c <*> t = Column $ map (<*> t) c

instance Monad Vector where
    Scalar s >>= f = f s
    Row r >>= f = Row $ map (>>= f) r
    Column c >>= f = Column $ map (>>= f) c

instance Num a => Num (Vector a) where
    Row x + (Row y) = Row $ zipWith (+) x y
    Column x + (Column y) = Column $ zipWith (+) x y
    Scalar x + y = fmap (+ x) y
    x + (Scalar y) = fmap (+ y) x
    _+_ = error "Vector Addition: Dimension mismatch"
    Scalar x * y = fmap (* x) y
    x * Scalar y = fmap (* y) x
    Row x * (Column y) = getSum $ foldMap Sum $ zipWith (*) x y
    Column x * (Column y) = Column $ map (* Column y) x
    x * (Row y) = Row $ map (x *) y
    abs = fmap abs
    signum = fmap signum
    fromInteger = Scalar . fromInteger
    negate = fmap negate

instance Foldable Vector where
  foldr f acc (Scalar s) = f s acc
  foldr f acc (Row r) = foldr (flip (foldr f)) acc r
  foldr f acc (Column c) = foldr (flip (foldr f)) acc c

instance Traversable Vector where
    traverse f (Scalar s) = Scalar <$> f s
    traverse f (Row r) = Row <$> traverse (traverse f) r
    traverse f (Column c) = Column <$> traverse (traverse f) c

instance (Num a) => Semigroup (Vector a) where
    x <> y = (*) <$> x <*> y

instance (Num a) => Monoid (Vector a) where
    mempty = Scalar 1

instance (Cmplx a) => Cmplx (Vector a) where
    conj = fmap conj . transposeStruct

instance (Num a, Cmplx a) => Tensor Vector a where
    inner = dot
    -- TODO: norm for matrices
    norm t = (/ (sqrt.getSum $ foldMap (Sum . (**2) . abs) t)) <$> t
    fromList xs dim = go (foldr (con . Scalar) (Column []) xs) dim
        where
            go (Column c) n
                | n == 1 = Column [go (Row c) dim]
                | l <= fromIntegral dim = Column c
                | otherwise = go (Row $ take (l `div` fromIntegral n) c) dim
                    `con` go (Column $ drop (l `div` fromIntegral n) c) (n - 1)
                where l = length c

            go (Row r) n
                | n == 1 = Row [go (Column r) dim]
                | l <= fromIntegral dim = Row r
                | otherwise = go (Column $ take (l `div` fromIntegral n) r) dim
                    `con` go (Row $ drop (l `div` fromIntegral n) r) (n - 1)
                where l = length r

            go s n = error "error when adding tensor rank"
    rank (Scalar s) = 0
    rank (Row []) = 1
    rank (Column []) = 1
    rank (Row (r:rs)) = rank r + 1
    rank (Column (c:cs)) = rank c + 1

    isSquare t
        | (Scalar _) <- t = True
        | (Column c) <- t = getAll $ go (Column c) (length c)
        | (Row r) <- t = getAll $ go (Row r) (length r)
        where
            go (Column c) l
                | length c == l = foldMap (`go` l) c
                | otherwise = All False
            go (Row r) l
                | length r == l = foldMap (`go` l) r
                | otherwise = All False
            go (Scalar _) _ = All True

    conjTr = conj

instance (Floating a, Eq a, Show a, Ord a, UniformRange a, Cmplx a) => QuantumRegister Vector a where

    densityOp x = transposeStruct x * x

    initQubit zero one
        | abs zero ** 2 + abs one ** 2 == 1 = Right $ fromList [zero, one] 2
        | otherwise = Left "Sum of probability of |0> and |1> state does not equal to 1"

    initQubit0 = fromList [1, 0] 2

    initNumQubit0 n
        | n < 1 = Left "Invalid qubit number"
        | otherwise = Right $ foldr (|*|) initQubit0 $ replicate (fromIntegral n - 1) initQubit0

    toQuantumRegister q
        | isLeft num = Left "Input list length is not a power of 2"
        | otherwise = Right . norm $ fromList q 2
        where num = integralLog 2 (length q)

    illegalPeek (Scalar s) = show s
    illegalPeek q = go 0 0
        where
            w = width q
            h = height q
            go i j
                | i == h = mempty
                | j == w = "\t|\n" ++ go (i + 1) 0
                | j == 0 = "|\t" <> current <> "\t" <> go i (j + 1)
                | otherwise = current <> "\t" <> go i (j + 1)
                where
                    current = maybe "null" show (atM q (i, j))

    subSystem lst q
        | getAny $ foldMap (Any . (\ x -> x >= rank q || x < 0)) lst = Left "Index out of bounds"
        | not (isSquare q) = Left "Invalid Quantum Register"
        | otherwise = Right . fmap sqrt $ go lst (fmap (\x -> conj x * x) q)
        where
            go [] (Scalar s) = Scalar s
            go is (Row r)
                | null is || Prelude.head is /= 0 = foldr (+) (Scalar 0) preSum
                | otherwise = Row $ map (go $ fmap (flip (-) 1) (Prelude.tail is)) r
                where
                    preSum = map (go $ fmap (flip (-) 1) is) r
            go is (Column c)
                | null is || Prelude.head is /= 0 = foldr (+) (Scalar 0) preSum
                | otherwise = Column $ map (go $ fmap (flip (-) 1) (Prelude.tail is)) c
                where
                    preSum = map (go $ fmap (flip (-) 1) is) c
            go _ _ = error "Exceptions not handled by subSystem function"

    isEntangled lst q = do
            reduced <- reducedDensity lst q
            return $ abs (trace (reduced * reduced) - 1) > 1e-13

    mask0 = Row [
                Column[1, 0],
                Column[0, 0]
            ]

    mask1 = Row [
                Column[0, 0],
                Column[0, 1]
            ]

    pauliId = Row [
                Column [1, 0],
                Column [0, 1]
            ]

instance (RealFloat a, Show a, UniformRange a) => Gates Vector (Complex a) where
    pauliX = Row [
            Column [Scalar (0 :+ 0), Scalar (1 :+ 0)],
            Column [Scalar (1 :+ 0), Scalar (0 :+ 0)]
        ]

    pauliY = Row [
                Column [Scalar (0 :+ 0), Scalar (0 :+ 1)],
                Column [Scalar (0 :+ (-1)), Scalar (0 :+ 0)]
            ]

    pauliZ = Row [
                Column [Scalar (1 :+ 0), Scalar (0 :+ 0)],
                Column [Scalar (0 :+ 0), Scalar ((-1) :+ 0)]
            ]

    hadamard = Row [
                Column [Scalar (1 / sqrt 2 :+ 0), Scalar (1 / sqrt 2 :+ 0)],
                Column [Scalar (1 / sqrt 2 :+ 0), Scalar ((-1) / sqrt 2 :+ 0)]
            ]

at :: Integral b => a -> Vector a -> (b, b) -> a
at x y z = fromMaybe x (atM y z)

atM :: Integral b => Vector a -> (b, b) -> Maybe a
atM (Scalar s) (0, 0) = Just s
atM (Column (c:cs)) (i, j)
    | height (Column (c:cs)) > i = let offset = height c in
                                    if offset > i
                                    then atM c (i, j)
                                    else atM (Column cs) (i - offset, j)
    | otherwise = Nothing
atM (Row (r:rs)) (i, j)
    | width (Row (r:rs)) > j = let offset = width r in
                                    if offset > j
                                    then atM r (i, j)
                                    else atM (Row rs) (i, j - offset)
    | otherwise = Nothing
atM _ _ = Nothing

con :: Vector a -> Vector a -> Vector a
con x (Row y) = Row $ x : y
con x (Column y) = Column $ x : y
con _ _ = error "Appending to Scalar is prohibited"

append :: Vector a -> Vector a -> Vector a
append (Row x) (Row y) = Row $ x ++ y
append (Column x) (Column y) = Column $ x ++ y
append _ _ = error "Vector direction mismatch"

width :: Integral b => Vector a -> b
width (Scalar s) = 1
width (Row r) = getSum $ foldMap (Sum . width) r
width (Column c) = width $ Prelude.head c

height :: Integral b => Vector a -> b
height (Scalar s) = 1
height (Column c) = getSum $ foldMap (Sum . height) c
height (Row r) = height $ Prelude.head r

-- Transform structure retaining the original rank and shape
transformStruct :: Vector a -> Vector a
transformStruct (Row (Column c:cs)) = Column . go $ Column c:cs
    where
        go (Row [] : _) = []
        go (Column [] : _) = []
        go cs = Row (map (transformStruct . Vector.head) cs)
            : go (map Vector.tail cs)

transformStruct (Column (Row r:rs)) = Row . go $ Row r:rs
    where
        go (Row [] : _) = []
        go (Column [] : _) = []
        go rs = Column (map (transformStruct . Vector.head) rs)
            : go (map Vector.tail rs)

transformStruct (Column cs) = Column $ map transformStruct cs
transformStruct (Row rs) = Row $ map transformStruct rs
transformStruct t = t

-- transpose via structure
transposeStruct :: Vector a -> Vector a
transposeStruct (Scalar s) = Scalar s
transposeStruct (Row r) = Column $ map transposeStruct r
transposeStruct (Column c) = Row $ map transposeStruct c

-- transpose retaining structure
transpose :: Vector a -> Vector a
transpose = transformStruct . transposeStruct

head :: Vector a -> Vector a
head (Scalar s) = Scalar s
head (Row r) = Prelude.head r
head (Column c) = Prelude.head c

last :: Vector a -> Vector a
last (Scalar s) = Scalar s
last (Row r) = Prelude.last r
last (Column c) = Prelude.last c

tail :: Vector a -> Vector a
tail (Scalar s) = Scalar s
tail (Row r) = Row $ Prelude.tail r
tail (Column c) = Column $ Prelude.tail c

init :: Vector a -> Vector a
init (Scalar s) = Scalar s
init (Row r) = Row $ Prelude.init r
init (Column c) = Column $ Prelude.init c

dot :: Num a => Vector a -> Vector a -> Vector a
dot x y = transposeStruct x * y

reducedDensity :: (Num a, Cmplx a, Floating a) => [Int] -> Vector a -> Either String (Vector a)
reducedDensity lst q
    | getAny $ foldMap (Any . (\ x -> x >= rank q || x < 0)) lst = Left "Index out of bounds"
    | not (isSquare q) = Left "Invalid Quantum Register"
    | otherwise = Right $ go lst q (transposeStruct q)
    where
        go [] (Scalar s) (Scalar t) = Scalar (s * t)
        go is (Column c) (Row r)
            | null is || Prelude.head is /= 0 = (getSum . foldMap Sum) $ zipWith (go (flip (-) 1 <$> is)) c r
            | otherwise = Row $ map (Column . ($ c) . (\ b a -> map (\ x -> go (flip (-) 1 <$> Prelude.tail is) x b) a)) r
        go _ _ _ = error "Input is not a valid statevector"

isColumn :: Vector a -> Bool
isColumn (Column _) = True
isColumn _ = False

isRow :: Vector a -> Bool
isRow (Row _) = True
isRow _ = False

isScalar :: Vector a -> Bool
isScalar (Scalar _) = True
isScalar _ = False

trace :: (Num a) => Vector a -> a
trace (Scalar s) = s
trace (Row r) = (getSum . foldMap Sum) $ zipWith go [0, 1 ..] r
    where
        go i (Row r) = trace (r !! i)
        go i (Column c) = trace (c !! i)
        go i (Scalar s) = s
{-
    Tensor type:
    All one-dimensional tensors are column vectors. This means matrices,
    which are rank 2 tensors, compose of multiple column tensors, which
    may not follow conventional representation in other traversable types.
-}

-- data Tensor a = Scalar a | Tensor [Tensor a]
--     deriving(Eq)

-- instance (Show a) => Show (Tensor a) where
--     show (Scalar s) = show s
--     show (Tensor ((Tensor t):ts)) = "[\n" ++ unlines (go (Tensor t:ts)) ++ "]"
--         where go = fmap ("  " ++).concatMap(lines . show)
--     show (Tensor s) = show s ++ "\n"

-- instance Functor Tensor where
--     fmap f (Scalar s) = Scalar $ f s
--     fmap f (Tensor t) = Tensor $ map (fmap f) t

-- instance Applicative Tensor where
--     pure = Scalar
--     (Scalar s) <*> t = fmap s t
--     t <*> (Scalar s) = fmap ($ s) t
--     (Tensor []) <*> _ = Tensor []
--     (Tensor (Scalar s : ss)) <*> t = Tensor $ fmap s t : tensorList
--         where Tensor tensorList = Tensor ss <*> t
--     (Tensor (Tensor t : ts)) <*> t2 = Tensor $ (Tensor t <*> t2) : tensorList
--         where Tensor tensorList = Tensor ts <*> t2

-- instance Monad Tensor where
--     (Scalar s) >>= f = f s
--     (Tensor t) >>= f = Tensor $ map (>>= f) t

-- instance (Num a) => Num (Tensor a) where
--     (Scalar s) + t = fmap (+ s) t
--     t + (Scalar s) = fmap (+ s) t
--     (Tensor xs) + (Tensor ys) = Tensor $ zipWith (+) xs ys
--     (Scalar s) * t = fmap (* s) t
--     t * (Scalar s) = fmap (* s) t
--     _ * Tensor[] = Tensor []
--     (Tensor x) * (Tensor y) = getSum $ foldMap Sum (zipWith go x y)
--         where
--             go (Tensor a) b = Tensor $ map (* b) a
--             go a b = a * b
--     abs = fmap abs
--     signum = fmap signum
--     negate = fmap negate
--     fromInteger i = Scalar $ fromInteger i

-- instance Foldable Tensor where
--     foldr f acc (Tensor t) = foldr (flip $ foldr f) acc t
--     foldr f acc (Scalar s) = f s acc

-- instance Traversable Tensor where
--     traverse f (Scalar s) = Scalar <$> f s
--     traverse f (Tensor t) = Tensor <$> traverse (traverse f) t

-- -- TODO: revise function for higher-dimension tensors
-- fromList :: [a] -> Int -> Tensor a
-- fromList xs dim = go (foldr (con . Scalar) (Tensor[]) xs) dim
--     where
--         go (Tensor t) n
--             | n == 1 = Tensor [go (Tensor t) dim]
--             | length t <= dim = Tensor t
--             | otherwise = go (Tensor $ take (length t `div` n) t) dim
--                 `con` go (Tensor $ drop (length t `div` n) t) (n - 1)
--         go s n = error "error when adding tensor rank"

-- product :: (Num a) => Tensor a -> Tensor a -> Tensor a
-- product (Scalar s) t = Scalar s * t
-- product t (Scalar s) = t * Scalar s
-- product (Tensor []) _ = Tensor []
-- product _ (Tensor []) = Tensor []
-- product (Tensor (x:xs)) y = Tensor $ current : rest
--     where
--         current = Tensor.product x y
--         Tensor rest = Tensor.product (Tensor xs) y

-- norm :: (Floating a) => Tensor a -> Tensor a
-- norm t= (/ (sqrt.getSum $ foldMap (Sum . (**2) . abs) t)) <$> t

-- append :: Tensor a -> Tensor a -> Tensor a
-- append (Tensor x) (Tensor y) = Tensor $ x ++ y
-- append _ _ = error "Concatenating scalars is prohibited"

-- con :: Tensor a -> Tensor a -> Tensor a
-- con x (Tensor y) = Tensor $ x : y
-- con _ _ = error "Appending to Scalar is prohibited"

-- dot :: Num a => Tensor a -> Tensor a -> a
-- dot (Tensor x) (Tensor y) = getSum $ foldMap Sum (zipWith dot x y)
-- dot (Scalar x) (Scalar y) = x * y
-- dot _ _ = error "Length mismatch"

-- flatten2d :: Tensor a -> Tensor a
-- flatten2d (Tensor ((Tensor ((Tensor (t:ts)):ss)):xs)) = Tensor . go $ Tensor (Tensor (t:ts):ss):xs
--     where
--         go :: [Tensor a] -> [Tensor a]
--         go (Tensor (Tensor (Scalar s:ss):us):ys) = foldr (\a acc ->
--             let (Tensor result) = Tensor.transpose a
--                 in result ++ acc
--             ) [] (Tensor (Tensor (Scalar s:ss):us):ys)

--         go y = go flattened
--             where flattened = foldr (\a acc ->
--                     let (Tensor next) = a
--                         in Tensor (go next) : acc
--                     ) [] y
-- flatten2d t = t

-- head :: Tensor a -> Tensor a
-- head (Tensor ts) = Prelude.head ts
-- head t = t

-- last :: Tensor a -> Tensor a
-- last (Tensor ts) = Prelude.last ts
-- last t = t

-- tail :: Tensor a -> Tensor a
-- tail (Tensor ts) = Tensor $ Prelude.tail ts
-- tail t = t

-- init :: Tensor a -> Tensor a
-- init (Tensor ts) = Tensor $ Prelude.init ts
-- init t = t

-- transpose :: Tensor a -> Tensor a
-- transpose (Tensor (Tensor t:ts)) = Tensor $ go $ Tensor t:ts
--     where
--         go :: [Tensor a] -> [Tensor a]
--         go (Tensor [] : _) = []
--         go xs = Tensor (map Tensor.head xs) : go (map Tensor.tail xs)
-- transpose (Tensor (Scalar s:ss)) = Tensor $ map (Tensor.(:[])) (Scalar s:ss)
-- transpose t = t

-- isSquare :: Tensor a -> Bool
-- isSquare (Scalar _) = True
-- isSquare (Tensor []) = True
-- isSquare (Tensor t) = getAll $ go (Tensor t) (length t)
--     where
--         go (Tensor x) l
--             | length x == l = foldMap (`go` l) x
--             | otherwise = All False
--         go (Scalar _) _ = All True

-- rank :: Integral b => Tensor a -> b
-- rank (Scalar s) = 0
-- rank (Tensor []) = 1
-- rank (Tensor (t:ts)) = rank t + 1


{-# LANGUAGE BangPatterns, DeriveDataTypeable, DeriveGeneric,
    FlexibleContexts #-}
-- |
-- Module    : Statistics.Sample.Powers
-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Very fast statistics over simple powers of a sample.  These can all
-- be computed efficiently in just a single pass over a sample, with
-- that pass subject to stream fusion.
--
-- The tradeoff is that some of these functions are less numerically
-- robust than their counterparts in the 'Statistics.Sample' module.
-- Where this is the case, the alternatives are noted.

module Statistics.Sample.Powers
    (
    -- * Types
      Powers

    -- * Constructor
    , powers

    -- * Descriptive functions
    , order
    , count
    , sum

    -- * Statistics of location
    , mean

    -- * Statistics of dispersion
    , variance
    , stdDev
    , varianceUnbiased

    -- * Functions over central moments
    , centralMoment
    , skewness
    , kurtosis

    -- * References
    -- $references
    ) where

import Control.Monad.ST
import Data.Aeson            (FromJSON, ToJSON)
import Data.Binary           (Binary(..))
import Data.Data             (Data, Typeable)
import Data.Vector.Binary    ()
import Data.Vector.Unboxed   ((!))
import GHC.Generics          (Generic)
import Numeric.SpecFunctions (choose)
import Prelude hiding (sum)
import Statistics.Function   (indexed)
import qualified Data.Vector          as V
import qualified Data.Vector.Generic  as G
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector.Unboxed.Mutable as MU
import qualified Statistics.Sample.Internal  as S

newtype Powers = Powers (U.Vector Double)
    deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON Powers
instance ToJSON Powers

instance Binary Powers where
    put (Powers v) = put v
    get = fmap Powers get

-- | O(/n/) Collect the /n/ simple powers of a sample.
--
-- Functions computed over a sample's simple powers require at least a
-- certain number (or /order/) of powers to be collected.
--
-- * To compute the /k/th 'centralMoment', at least /k/ simple powers
--   must be collected.
--
-- * For the 'variance', at least 2 simple powers are needed.
--
-- * For 'skewness', we need at least 3 simple powers.
--
-- * For 'kurtosis', at least 4 simple powers are required.
--
-- This function is subject to stream fusion.
powers :: G.Vector v Double =>
          Int                   -- ^ /n/, the number of powers, where /n/ >= 2.
       -> v Double
       -> Powers
powers k sample
  | k < 2     = error "Statistics.Sample.powers: too few powers"
  | otherwise = runST $ do
      acc <- MU.replicate l 0
      G.forM_ sample $ \x ->
        let loop !i !xk
              | i == l    = return ()
              | otherwise = do MU.write acc i . (+ xk) =<< MU.read acc i
                               loop (i+1) (xk * x)
        in loop 0 1
      fmap Powers $ U.unsafeFreeze acc
  where
    l = k + 1
{-# SPECIALIZE powers :: Int -> U.Vector  Double -> Powers #-}
{-# SPECIALIZE powers :: Int -> V.Vector  Double -> Powers #-}
{-# SPECIALIZE powers :: Int -> SV.Vector Double -> Powers #-}

-- | The order (number) of simple powers collected from a 'sample'.
order :: Powers -> Int
order (Powers pa) = U.length pa - 1

-- | Compute the /k/th central moment of a sample.  The central
-- moment is also known as the moment about the mean.
centralMoment :: Int -> Powers -> Double
centralMoment k p@(Powers pa)
    | k < 0 || k > order p =
                  error ("Statistics.Sample.Powers.centralMoment: "
                         ++ "invalid argument")
    | k == 0    = 1
    | otherwise = (/n) . S.sum . U.map go . indexed . U.take (k+1) $ pa
  where
    go (i , e) = (k `choose` i) * ((-m) ^ (k-i)) * e
    n = U.head pa
    m = mean p

-- | Maximum likelihood estimate of a sample's variance.  Also known
-- as the population variance, where the denominator is /n/.  This is
-- the second central moment of the sample.
--
-- This is less numerically robust than the variance function in the
-- 'Statistics.Sample' module, but the number is essentially free to
-- compute if you have already collected a sample's simple powers.
--
-- Requires 'Powers' with 'order' at least 2.
variance :: Powers -> Double
variance = centralMoment 2

-- | Standard deviation.  This is simply the square root of the
-- maximum likelihood estimate of the variance.
stdDev :: Powers -> Double
stdDev = sqrt . variance

-- | Unbiased estimate of a sample's variance.  Also known as the
-- sample variance, where the denominator is /n/-1.
--
-- Requires 'Powers' with 'order' at least 2.
varianceUnbiased :: Powers -> Double
varianceUnbiased p@(Powers pa)
    | n > 1     = variance p * n / (n-1)
    | otherwise = 0
  where n = U.head pa

-- | Compute the skewness of a sample. This is a measure of the
-- asymmetry of its distribution.
--
-- A sample with negative skew is said to be /left-skewed/.  Most of
-- its mass is on the right of the distribution, with the tail on the
-- left.
--
-- > skewness . powers 3 $ U.to [1,100,101,102,103]
-- > ==> -1.497681449918257
--
-- A sample with positive skew is said to be /right-skewed/.
--
-- > skewness . powers 3 $ U.to [1,2,3,4,100]
-- > ==> 1.4975367033335198
--
-- A sample's skewness is not defined if its 'variance' is zero.
--
-- Requires 'Powers' with 'order' at least 3.
skewness :: Powers -> Double
skewness p = centralMoment 3 p * variance p ** (-1.5)

-- | Compute the excess kurtosis of a sample.  This is a measure of
-- the \"peakedness\" of its distribution.  A high kurtosis indicates
-- that the sample's variance is due more to infrequent severe
-- deviations than to frequent modest deviations.
--
-- A sample's excess kurtosis is not defined if its 'variance' is
-- zero.
--
-- Requires 'Powers' with 'order' at least 4.
kurtosis :: Powers -> Double
kurtosis p = centralMoment 4 p / (v * v) - 3
    where v = variance p

-- | The number of elements in the original 'Sample'.  This is the
-- sample's zeroth simple power.
count :: Powers -> Int
count (Powers pa) = floor $ U.head pa

-- | The sum of elements in the original 'Sample'.  This is the
-- sample's first simple power.
sum :: Powers -> Double
sum (Powers pa) = pa ! 1

-- | The arithmetic mean of elements in the original 'Sample'.
--
-- This is less numerically robust than the mean function in the
-- 'Statistics.Sample' module, but the number is essentially free to
-- compute if you have already collected a sample's simple powers.
mean :: Powers -> Double
mean p@(Powers pa)
    | n == 0    = 0
    | otherwise = sum p / n
    where n     = U.head pa

-- $references
--
-- * Besset, D.H. (2000) Elements of statistics.
--   /Object-oriented implementation of numerical methods/
--   ch. 9, pp. 311–331.
--   <http://www.elsevier.com/wps/product/cws_home/677916>
--
-- * Anderson, G. (2009) Compute /k/th central moments in one
--   pass. /quantblog/. <http://quantblog.wordpress.com/2009/02/07/compute-kth-central-moments-in-one-pass/>


module GradientDescent where
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util

--  Performs gradient descent to learn theta
--  theta = GRADIENTDESCENT(x, y, theta, alpha, num_iters) updates theta by
--  taking num_iters gradient steps with learning rate alpha
gradientDescent :: (Matrix Double -> Matrix Double) -> Matrix Double -> Double -> Int -> Matrix Double
gradientDescent gradFunc initialTheta alpha num_iters = iterate initialTheta num_iters
            where
                iterate theta 0 = theta
                iterate theta iters = iterate (updateTheta theta) (iters-1)
                updateTheta t = t - (scalar alpha) * (gradFunc t)


module Deanie.Density (
    logDensityBernoulli
  , logDensityBeta
  , logDensityGamma
  , logDensityGaussian
  ) where

import Numeric.SpecFunctions (logGamma)

logDensityBernoulli :: Double -> Bool -> Double
logDensityBernoulli p x
    | p < 0 || p > 1 = log 0
    | x              = log p
    | otherwise      = log (1 - p)
{-# INLINE logDensityBernoulli #-}

logDensityBeta :: Double -> Double -> Double -> Double
logDensityBeta a b x
  | x <= 0 || x >= 1 = log 0
  | a < 0 || b < 0   = log 0
  | otherwise        = (a - 1) * log x + (b - 1) * log (1 - x)
{-# INLINE logDensityBeta #-}

-- | The log-density for the gamma distribution with rate parameter b.
--
--     f(x) = b ^ a * x ^ (a - 1) * x ^ (-b * x) / Gamma(a)
logDensityGamma :: Double -> Double -> Double -> Double
logDensityGamma a b x
  | a < 0 || b < 0 = log 0
  | otherwise      = a * log b + (a - 1) * log x - b * x - logGamma a
{-# INLINE logDensityGamma #-}

-- | The density for the normal distribution with specified mean and standard
--   deviation.
logDensityGaussian :: Double -> Double -> Double -> Double
logDensityGaussian m sd x
  | sd <= 0   = log 0
  | otherwise = negate (log sd) - (x - m) * (x - m) / (2 * sd * sd)
{-# INLINE logDensityGaussian #-}


{-|
Module: MachineLearning.Optimization
Description: Optimization
Copyright: (c) Alexander Ignatyev, 2016
License: BSD-3
Stability: experimental
Portability: POSIX

Optimization module.
-}

module MachineLearning.Optimization
(
    MinimizeMethod(..)
  , minimize
  , checkGradient
)

where

import MachineLearning.Types (R, Vector, Matrix)
import MachineLearning.Model as Model
import MachineLearning.Regularization (Regularization)
import qualified MachineLearning.Optimization.GradientDescent as GD
import qualified MachineLearning.Optimization.MinibatchGradientDescent as MGD
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

import qualified Numeric.GSL.Minimization as Min

data MinimizeMethod = GradientDescent R       -- ^ Gradient descent, takes alpha. Requires feature normalization.
                    | MinibatchGradientDescent Int Int R  -- ^ Minibacth Gradietn Descent, takes seed, batch size and alpha
                    | ConjugateGradientFR R R -- ^ Fletcher-Reeves conjugate gradient algorithm,
                                              -- takes size of first trial step (0.1 is fine) and tol (0.1 is fine).
                    | ConjugateGradientPR R R -- ^ Polak-Ribiere conjugate gradient algorithm.
                                              -- takes size of first trial step (0.1 is fine) and tol (0.1 is fine).
                    | BFGS2 R R               -- ^ Broyden-Fletcher-Goldfarb-Shanno (BFGS) algorithm,
                                              -- takes size of first trial step (0.1 is fine) and tol (0.1 is fine).

-- | Returns solution vector (theta) and optimization path.
-- Optimization path's row format:
-- [iter number, cost function value, theta values...]
minimize :: Model.Model a =>
            MinimizeMethod
         -> a       -- ^ model (Least Squares, Logistic Regression etc)
         -> R   -- ^ epsilon, desired precision of the solution
         -> Int     -- ^ maximum number of iterations allowed
         -> Regularization   -- ^ regularization parameter
         -> Matrix  -- ^ X
         -> Vector  -- ^ y
         -> Vector  -- ^ initial solution, theta
         -> (Vector, Matrix) -- ^ solution vector and optimization path

minimize (BFGS2 firstStepSize tol) = minimizeVD Min.VectorBFGS2 firstStepSize tol
minimize (ConjugateGradientFR firstStepSize tol) = minimizeVD Min.ConjugateFR firstStepSize tol
minimize (ConjugateGradientPR firstStepSize tol) = minimizeVD Min.ConjugatePR firstStepSize tol
minimize (GradientDescent alpha) = GD.gradientDescent alpha
minimize (MinibatchGradientDescent seed batchSize alpha) = MGD.minibatchGradientDescent seed batchSize alpha

minimizeVD method firstStepSize tol model epsilon niters reg x y initialTheta
  = Min.minimizeVD method epsilon niters firstStepSize tol costf gradientf initialTheta
  where costf = Model.cost model reg x y
        gradientf = Model.gradient model reg x y

-- | Gradient checking function.
-- Approximates the derivates of the Model's cost function
-- and calculates derivatives using the Model's gradient functions.
-- Returns norn_2 between 2 derivatives.
-- Takes model, regularization, X, y, theta and epsilon (used to approximate derivatives, 1e-4 is a good value).
checkGradient :: Model a => a -> Regularization -> Matrix -> Vector -> Vector -> R -> R
checkGradient model reg x y theta eps = LA.norm_2 $ grad1 - grad2
  where theta_m = LA.asColumn theta
        eps_v = V.replicate (V.length theta) eps
        eps_m = LA.diag eps_v
        -- +/- eps_m in case of zero theta
        theta_m1 = theta_m * (1 + eps_m) + eps_m
        theta_m2 = theta_m * (1 - eps_m) - eps_m
        cost1 = LA.vector $ map (cost model reg x y) $ LA.toColumns theta_m1
        cost2 = LA.vector $ map (cost model reg x y) $ LA.toColumns theta_m2
        eps_v1 = LA.takeDiag $ theta_m1 - theta_m
        eps_v2 = LA.takeDiag $ theta_m - theta_m2
        grad1 = (cost1 - cost2) / (eps_v1 + eps_v2)
        grad2 = gradient model reg x y theta


-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Eigen
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Eigenvalue decompositions.
--
module Numeric.LinearAlgebra.Matrix.Eigen (
    hermEigen,
    hermEigenvalues,
    hermEigenM_,
    hermEigenvaluesM_,
    ) where

import Control.Monad( when )
import Control.Monad.ST( ST, runST, unsafeIOToST )
import Foreign( allocaArray, nullPtr )
import Text.Printf( printf )

import Foreign.LAPACK( LAPACK, EigJob(..), EigRange(..) )
import qualified Foreign.LAPACK as LAPACK

import Numeric.LinearAlgebra.Types( Herm(..) )
import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import Numeric.LinearAlgebra.Vector( Vector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V

-- | Compute the eigenvalues and eigenvectors of a Hermitian matrix.
-- Return the eigenvalues are in ascending order in the result vector;
-- store the corresponding eigenvectors are in the columns of the result
-- matrix.
hermEigen :: (LAPACK e) => Herm Matrix e -> (Vector Double, Matrix e)
hermEigen (Herm uplo a) = runST $ do
    (_,n) <- M.getDim a
    ma <- M.newCopy a
    mw <- V.new_ n
    mz <- M.new_ (n,n)
    hermEigenM_ (Herm uplo ma) mw mz
    w <- V.unsafeFreeze mw
    z <- M.unsafeFreeze mz
    return (w,z)

-- | Return the eigenvalues of a Hermitian matrix in ascending order.
hermEigenvalues :: (LAPACK e) => Herm Matrix e -> Vector Double
hermEigenvalues (Herm uplo a) = V.create $ do
    (_,n) <- M.getDim a
    ma <- M.newCopy a
    w <- V.new_ n
    hermEigenvaluesM_ (Herm uplo ma) w
    return w

-- | Compute and copy the eigenvalues and eigenvectors of a mutable
-- Hermitian matrix.  This destroys the original Hermitian matrix.
hermEigenM_ :: (LAPACK e)
                  => Herm (STMatrix s) e
                  -> STVector s Double
                  -> STMatrix s e
                  -> ST s ()
hermEigenM_ (Herm uplo a) w z = do
    (ma,na) <- M.getDim a
    (mz,nz) <- M.getDim z
    nw <- V.getDim w
    let n = na

    when (ma /= na) $ error $
        printf ("hermEigenM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>): nonsquare matrix"
               ) ma na
    when ((not . and) [ (ma,na) == (n,n)
                      , nw == n
                      , (mz,nz) == (n,n)
                      ]) $ error $
        printf ("hermEigenM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>:"
                ++ " <matrix with dim (%d,%d)>"
                ++ " dimension mismatch")
               ma na nw mz nz

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith w $ \pw ->
        M.unsafeWith z $ \pz ldz ->
        allocaArray (2*n) $ \psuppz -> do
            _m <- LAPACK.heevr jobz range uplo n pa lda abstol pw pz ldz
                               psuppz
            return ()
  where
    jobz = EigVectors
    range = AllEigs
    abstol = 0

-- | Compute and copy the eigenvalues of a mutable Hermitian matrix.  This
-- destroys the original Hermitian matrix.
hermEigenvaluesM_ :: (LAPACK e)
                        => Herm (STMatrix s) e
                        -> STVector s Double
                        -> ST s ()
hermEigenvaluesM_ (Herm uplo a) w = do
    (ma,na) <- M.getDim a
    nw <- V.getDim w
    let n = na

    when (ma /= na) $  error $
        printf ("hermEigenvaluesM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>): nonsquare matrix"
               ) ma na

    when ((not . and) [ (ma,na) == (n,n)
                      , nw == n
                      ]) $ error $
        printf ("hermEigenvaluesM_"
                ++ " (Herm _ <matrix with dim (%d,%d)>)"
                ++ " <vector with dim %d>: dimension mismatch")
               ma na nw

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith w $ \pw -> do
            _m <- LAPACK.heevr jobz range uplo n pa lda abstol pw pz ldz
                               psuppz
            return ()

  where
    jobz = NoEigVectors
    range = AllEigs
    abstol = 0
    pz = nullPtr
    ldz = 1
    psuppz = nullPtr


module Main where

import Control.Applicative
import Graphics.Rendering.Chart.Backend.Diagrams (toFile)
import qualified Graphics.Rendering.Chart.Easy as Plt
import Kalman
import Numeric.LinearAlgebra
import Sim

fA1 :: (Floating a) => [a] -> [a]
fA1 [x, y, xh, yh, ax, ay] =
  [ x + xh * dt + c * ax * dt * dt,
    y + yh * dt + c * ay * dt * dt,
    xh + ax * dt,
    yh + ay * dt,
    ax,
    ay
  ]
  where
    dt = 0.1
    c = 0.5

fH1 :: (Floating a) => [a] -> [a]
fH1 [x, y, xh, yh, ax, ay] = [x, y]

cannonBall = toFile Plt.def "CannonBall.svg" $ do
  -- KF Initial conditions
  let x0 = [0, 0, 10, 10, 0, -9.8]
  let q = diagl [0.1, 0.1, 1e-3, 1e-3, 1e-3, 1e-3]
  let r = diagl (replicate 2 0.5)
  let s0 = Estimate (vector x0) (diagl (replicate 6 5))

  -- Fire the cannon ball
  let plant = Plant fA1 id fH1 id
  let sim = take 50 $ simulate plant (vector x0) 0
  let measurements = map snd sim
  let trajectory = map fst sim

  -- Run KF/EKF/UKF
  let sys = System fA1 fH1 q r
  let ekfs = scanl (ekf sys) s0 measurements
  let ukfs = scanl (ukf sys) s0 measurements

  -- Plot the results
  let vector2Chart = map ((\[x, y] -> (x, y)) . toList . subVector 0 2)
  let state2Chart xs = vector2Chart $ map sX xs

  Plt.plot (Plt.line "Plant" [vector2Chart trajectory])
  Plt.plot (Plt.line "UKF" [state2Chart ukfs])
  Plt.plot (Plt.line "EKF" [state2Chart ekfs])
  Plt.plot (Plt.points "measurements" (vector2Chart measurements))

main = cannonBall


{-# LANGUAGE FlexibleContexts #-}

module Main where

import Graphics.Plot

import Data.Packed.Vector
import Data.Packed.Matrix
import Numeric.Container
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.HMatrix (tr)

import System.Environment (getArgs)

-- TODO (a->b->c) -> Vector a -> Vector b -> Matrix c
outerOp :: (Element a, Element b) => (Matrix a -> Matrix b -> Matrix c) -> Vector a -> Vector b -> Matrix c
outerOp k x x' = k xp xq
    where xp = repmat (asColumn x) 1 (dim x')
          xq= repmat (asRow x') (dim x) 1

k_SE sigma l x x' = (sigma^2*) $ exp $ (-1/(2*l^2))*dx^2
    where dx = outerOp (-) x x'

k = k_SE 3 0.3

inferenceTrivial k x y ySigma2 domain = (mean, variance)
    where mean = k domain x <> invkxx <> y
          variance = k domain domain - (kdomainx <> invkxx <> k x domain)
          invkxx = inv $ k x x + diag ySigma2
          kdomainx = k domain x

inferenceCholesky k x y ySigma2 domain = (mean, variance)
    where mean = (tr kxdomain) <> alpha
          variance = k domain domain - (tr v <> v)
          alpha = tr l <\> (l <\> y)
          l = chol $ k x x + diag ySigma2
          v = l <\> kxdomain
          kxdomain = k x domain

plotProcessDistribution :: Vector Double -> Vector Double -> Vector Double -> IO ()
plotProcessDistribution domain m v = mplot [domain, m, m + v, m - v, m + 2 * v, m - 2 * v]

main :: IO ()
main = do
    datfile:[] <- getArgs
    dat <- fmap readMatrix $ readFile datfile
    let [x, y] = toColumns dat -- [x,y] or [x,f]
    let domain = linspace 128 (0.5, 10.5)

    let n = 1
    let (x',y') = (vjoin $ replicate n x, vjoin $ take n $ repeat y)

    let (m, v') = inferenceTrivial k x' y' (constant 1 $ dim y') domain
    let v = takeDiag v'

    plotProcessDistribution domain m v


module Main where

import Criterion.Main (defaultMain, bgroup, bench, nf, nfIO, whnf)
import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import qualified Numeric.Morpheus as M

main :: IO ()
main = do
  !a <- LA.rand 1000 1000
  !b <- LA.rand 1000 1000
  !x <- LA.flatten <$> LA.rand 1000 1000
  !y <- LA.flatten <$> LA.rand 1000 1000
  let !ata = LA.mTm a

  defaultMain [
    bgroup "Linear Algebra" [
        bench "random matrix" $ nfIO (LA.rand 1000 1000)
        , bench "inv" $ nf LA.inv a
        , bench "Cholesky decomposition" $ nf LA.chol ata
        , bench "vector dot product" $ whnf (LA.dot x) y
        , bench "matrix transpose" $ nf LA.tr a
        , bench "matrix multiplication" $ nf (a LA.<>) b
        , bench "sigmoid" $ nf sigmoid a
        , bench "ReLU" $ nf relu a
        , bench "sum by rows" $ nf (reduceByRows V.sum) a
        , bench "sum by columns" $ nf (reduceByColumns V.sum) a
        , bench "max index in rows" $ nf (reduceByRowsV (fromIntegral . LA.maxIndex) ) a
        , bench "max index in columns" $ nf (reduceByColumnsV (fromIntegral . LA.maxIndex) ) a
        , bench "sigmoid (morpheus)" $ nf M.sigmoid a
        , bench "sigmoid gradient (morpheus)" $ nf M.sigmoidGradient a
        , bench "ReLU (morpheus)" $ nf M.relu a
        , bench "sum by rows (morpheus)" $ nf M.rowSum a
        , bench "sum by columns (morpheus)" $ nf M.columnSum a
        , bench "max index in rows (morpheus)" $ nf M.rowMaxIndex a
        , bench "max index in columns (morpheus)" $ nf M.columnMaxIndex a

        ]
    ]

type Matrix = LA.Matrix LA.R
type Vector = LA.Vector LA.R
type R = LA.R

sigmoid :: Floating a => a -> a
sigmoid z = 1 / (1+exp(-z))

relu :: Matrix -> Matrix
relu x = x * (LA.step x)

reduceByRowsV :: (Vector -> R) -> Matrix -> Vector
reduceByRowsV f = LA.vector . map f . LA.toRows

reduceByColumnsV :: (Vector -> R) -> Matrix -> Vector
reduceByColumnsV f = LA.vector . map f . LA.toColumns

reduceByRows :: (Vector -> R) -> Matrix -> Matrix
reduceByRows f = LA.asColumn . reduceByRowsV f

reduceByColumns :: (Vector -> R) -> Matrix -> Matrix
reduceByColumns f = LA.asRow . reduceByColumnsV f


import Data.Map.Strict ((!))
import Numeric.LinearAlgebra hiding ((!))
import LeastSquares

leastSquaresEstimate :: Vector R
leastSquaresEstimate = result ! "x"
  where result = minimize (SumSquares (a*x-b)) [c*x:==:d]
        a = Mat $ (3><2) [1.0, 1.0, 1.0, 4.0, 2.0, -1.0]
        b = Vec $ vector [13.0, 27.0, 1.0]
        c = Mat $ (1><2) [3.0, 7.0]
        d = Vec $ vector [-1.0]
        x = Var "x" 2


module Elliptic
  where
import Carlson
import Data.Complex
import Internal

ellipticF' :: Double -> Cplx -> Cplx -> Cplx
ellipticF' err phi m
  | phi == 0 =
    toCplx 0
  | m == 1 && abs(realPart phi) == pi/2 =
    toCplx (0/0)
  | m == 1 && abs(realPart phi) < pi/2 =
    atanh(sin phi)
  | abs(realPart phi) <= pi/2 =
    if m == 0
      then
        phi
      else
        let sine = sin phi in
        let sine2 = sine*sine in
        let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
        sine * carlsonRF' err cosine2 oneminusmsine2 1
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticF' err (pi/2) m + ellipticF' err phi' m

ellipticF :: Cplx -> Cplx -> Cplx
ellipticF = ellipticF' 1e-15

ellipticE' :: Double -> Cplx -> Cplx -> Cplx
ellipticE' err phi m
  | phi == 0 =
    toCplx 0
  | abs(realPart phi) <= pi/2 =
    case m of
      0 -> phi
      1 -> sin phi
      _ ->
        let sine = sin phi in
        let sine2 = sine*sine in
        let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
        sine * (carlsonRF' err cosine2 oneminusmsine2 1 -
          m * sine2 / 3 * carlsonRD' err cosine2 oneminusmsine2 1)
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticE' err (pi/2) m + ellipticE' err phi' m

ellipticE :: Cplx -> Cplx -> Cplx
ellipticE = ellipticE' 1e-15

ellipticPI' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
ellipticPI' err phi n m
  | phi == 0 =
    toCplx 0
  | phi == pi/2 && n == 1 =
    0/0
  | phi == pi/2 && m == 0 =
    pi/2/sqrt(1-n)
  | phi == pi/2 && m == n =
    ellipticE' err (pi/2) m / (1-m)
  | phi == pi/2 && n == 0 =
    ellipticF' err (pi/2) m
  | abs(realPart phi) <= pi/2 =
    let sine = sin phi in
    let sine2 = sine*sine in
    let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
    sine * (carlsonRF' err cosine2 oneminusmsine2 1 +
      n * sine2 / 3 * carlsonRJ' err cosine2 oneminusmsine2 1 (1-n*sine2))
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticPI' err (pi/2) n m + ellipticPI' err phi' n m

ellipticPI :: Cplx -> Cplx -> Cplx -> Cplx
ellipticPI = ellipticPI' 1e-15

jacobiZeta' :: Double -> Cplx -> Cplx -> Cplx
jacobiZeta' err phi m =
  if m == 1
    then
      if abs(realPart phi) <= pi/2
        then sin phi
        else let (phi',_) = getPhiK phi in sin phi'
    else
      ellipticE' err phi m -
        ellipticE' err (pi/2) m / ellipticF' err (pi/2) m *
        ellipticF' err phi m

jacobiZeta :: Cplx -> Cplx -> Cplx
jacobiZeta = jacobiZeta' 1e-15


{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Math.NumberTheory.Primes.Testing
import Math.NumberTheory.Moduli
import Options.Applicative
import System.Random
import System.IO
import System.Exit
import Data.Bits
import Data.Function
import Data.List
import Data.Maybe
import Test.QuickCheck
import Debug.Trace
import Control.Monad
import Statistics.Sample
import Statistics.Quantile
import Control.Exception (evaluate)
import Data.Bifunctor
import GHC.Stack
import Text.Read
import Text.Printf
import Data.List.Split
import qualified Data.Vector as V
import qualified Data.Sequence as S
import qualified Data.IntMap as IM
import qualified Data.IntSet as IS
import Data.Monoid

-- Some convenient names for some parameters
type BitLength = Int
type Width = Int

-- http://stackoverflow.com/a/13301611/946226
rndPrime :: BitLength -> Gen Integer
rndPrime 1 = error "Cannot create distinct p and q for 1 bit"
rndPrime bits = do
    x <- fmap (.|. 1) $ choose (2^(bits - 1), 2^bits - 1)
    if isPrime x then return x else rndPrime bits

data RSAParams = RSAParams { bits :: Int, p, q, n, e, dp, dq, kp, kq :: Integer } deriving Show

instance Arbitrary RSAParams where
    arbitrary = sized (\s -> genRSAParams (min 1024 (max s 2)))

genRSAParams :: BitLength -> Gen RSAParams
genRSAParams bits = do
    p <- rndPrime bits
    q <- rndPrime bits
    let n = p*q
        e = 65537
        mdp = invertMod e (p-1)
        mdq = invertMod e (q-1)
    case (mdp, mdq) of
        (Just dp, Just dq) -> do
            let kp = (e*dp-1) `div` (p-1)
                kq = (e*dq-1) `div` (q-1)
            return (RSAParams {..})
        _ -> genRSAParams bits -- retry

data Step = S | M deriving (Eq, Ord, Show) -- square or multiply
type Leak = V.Vector Step

leftToRight :: BitLength -> Width -> Integer -> Leak
leftToRight bits w s | w <= 0 = error "width must be positive"
leftToRight bits w s = V.fromList $ go (bits-1)
   where
    go i | i < 0         = []
    go i | s `testBit` i = replicate (slideLength-1) S ++ M : go (i - slideLength)
      where
        Just slideLength = find (\l -> s `testBit` (i-l+1)) $ [w,w-1..1]
    go i | otherwise     = S : go (i-1)

-- Bits learned
leftToRightGMPBits :: BitLength -> Width -> Integer -> Integer
leftToRightGMPBits bits w s | w <= 0 = error "width must be positive"
leftToRightGMPBits bits w s = go (bits -1) 0
   where
    go i !acc | i < 0         = acc
    go i !acc | s `testBit` i = go (i - slideLength) (acc + learned)
      where
        Just slideLength = find (\l -> s `testBit` (i-l+1)) $ [w,w-1..1]
        learned | slideLength == 1 = 1
                | otherwise        = 2
    go i ! acc | otherwise    =  go (i-1) (acc + 1)

rightToLeft :: BitLength -> Width -> Integer -> Leak
rightToLeft bits w s | w <= 0 = error "width must be positive"
rightToLeft bits w s = V.fromList $ go 0 []
   where
    go i acc | i == bits     = acc
    go i acc | s `testBit` i = go (i+slideLength) (replicate (slideLength-1) S ++ M : acc)
      where slideLength = min w (bits - i)
    go i acc | otherwise     = go (i+1) (S : acc)

-- Bits learned
rightToLeftBits :: BitLength -> Width -> Integer -> Integer
rightToLeftBits bits w s | w <= 0 = error "width must be positive"
rightToLeftBits bits w s = go 0 0
   where
    go i !acc | i == bits     = acc
    go i !acc | s `testBit` i = go (i+slideLength) (1+acc)
      where slideLength = min w (bits - i)
    go i !acc | otherwise     = go (i+1) (1+acc)

-- A datatype to track known bits
data BitMask = BM { unknownBits :: !Integer, knownBits :: !Integer }
    deriving Eq

showBitMask bits bm = map go [(bits-1),(bits-2)..0]
  where go i | not (bm `isKnownBit` i) = '_'
             | bm `getKnownBit` i      = '1'
             | otherwise               = '0'

instance Show BitMask where
    show bm = showBitMask bits bm
      where bits = head [i | i <- [0..], 2^i > knownBits bm && 2^i > unknownBits bm]

isKnownBit :: BitMask -> Int -> Bool
bm `isKnownBit` i = not (unknownBits bm `testBit` i)

getKnownBit :: BitMask -> Int -> Bool
bm `getKnownBit` i = knownBits bm `testBit` i

possibleBits :: BitMask -> Int -> [Bool]
possibleBits bm i
    | bm `isKnownBit` i = [bm `getKnownBit` i]
    | otherwise         = [False, True]

matches :: Integer -> BitMask -> Bool
matches p bm = (p .|. unknownBits bm) == (knownBits bm .|. unknownBits bm)

isKnown :: BitMask -> Maybe Integer
isKnown bm | unknownBits bm == 0 = Just (knownBits bm)
           | otherwise           = Nothing

writeBit :: Int -> Bool -> Integer -> Integer
writeBit i True  p = p `setBit` i
writeBit i False p = p `clearBit` i

writeKnownBit :: Int -> Bool -> BitMask -> BitMask
writeKnownBit i b bm = BM (unknownBits bm `clearBit` i) (writeBit i b (knownBits bm))

nothingKnown :: BitLength -> BitMask
nothingKnown bits = BM (2^bits-1) 0

anyOddNumber :: BitLength -> BitMask
anyOddNumber bits = writeKnownBit 0 True (nothingKnown bits)

fromKnownNumber :: Integer -> BitMask
fromKnownNumber i = BM 0 i

-- can be optimized, I guess
allMatching :: BitMask -> [Integer]
allMatching bm = map sum $ mapM go [(bits-1),(bits-2)..0]
  where go i | not (bm `isKnownBit` i) = [2^i,0]
             | bm `getKnownBit` i      = [2^i]
             | otherwise               = [0]
        bits = head [i | i <- [0..], 2^i > knownBits bm && 2^i > unknownBits bm]

maskFromMatching :: [Integer] -> BitMask
maskFromMatching (m:ms) = foldl' go (fromKnownNumber m) ms
  where
    go :: BitMask -> Integer -> BitMask
    go bm i = BM (unknownBits bm .|. disagree) (knownBits bm `clear` disagree)
      where disagree = (unknownBits bm .|. knownBits bm) `xor` (unknownBits bm .|. i)

x `clear` y = complement (complement x .|. y)

countKnownBits :: BitLength -> BitMask -> Int
countKnownBits bits bm = bits - popCount (unknownBits bm)

-- A breadth first search limiting the count by n
bfSearch :: Int -> a -> (a -> [a]) -> (a -> Bool) -> Maybe Int
bfSearch limit s0 next done = go 0 (S.singleton s0)
  where
    go n _ | n >= limit = Nothing
    go n todo = case S.viewl todo of
        S.EmptyL -> Nothing
        s S.:< todo' | done s    -> Just n
                     | otherwise -> go (n+1) (todo' S.>< S.fromList (next s))

dfSearch :: Int -> a -> (a -> [a]) -> (a -> Bool) -> Maybe Int
dfSearch limit s0 next done = go 0 [s0]
  where
    go n _ | n >= limit = Nothing
    go n [] = Nothing
    go n (s:todo') | done s    = Just n
                   | otherwise = go (n+1) (next s ++ todo')

-- Currently unused
linearWorkList :: forall a. Eq a =>
    IS.IntSet -> IM.IntMap a -> (Int -> IM.IntMap a -> a) -> IM.IntMap a
linearWorkList todo s0 recalc = go todo s0
  where
    go :: IS.IntSet -> IM.IntMap a -> IM.IntMap a
    go todo s | IS.null todo         = s
    go todo s | oldValue == newValue = go todo' s
              | otherwise            = go newTodo newS
      where
        (k, todo') = IS.deleteFindMin todo
        oldValue = s IM.! k
        newValue = recalc k s
        newS = IM.insert k newValue s
        newTodo = (if (k-1) `IM.member` s then IS.insert (k-1) else id) $
                  (if (k+1) `IM.member` s then IS.insert (k+1) else id) $ todo'

type WindowLengthRanges = [(Int, (Int,Int))]

possibleWindowLengths :: BitLength -> Width -> Leak -> WindowLengthRanges
possibleWindowLengths bits w leak = IM.elems result
  where
    clamp width = max 1 (min w width)

    -- positions of the window end (from the back, i.e. bits)
    windowEnds = IM.fromList $ zip [0..] $ map ((bits - 1) - ) $
                 V.toList $ V.findIndices (== M) leak

    windows = IM.keysSet windowEnds
    maxKey = IS.findMax windows

    zipIM = IM.intersectionWith (,)

    windowMax = IM.mapWithKey (\k _ -> w_max k) windowEnds
    windowMin = IM.mapWithKey (\k _ -> w_min k) windowEnds
    result = windowEnds `zipIM` (windowMin `zipIM` windowMax)

    msbIsOne = False

    -- The first window spans precisely from the MSB to the first M
    firstWinLength =  clamp (bits - pos)
      where pos = windowEnds IM.! 0

    -- minimum length. computed from right to left
    w_min k | k == 0, msbIsOne = firstWinLength
            | k == maxKey      = 1
            | otherwise        = clamp (j + w_min_j + w -i)
        where
            i       = windowEnds IM.! k
            j       = windowEnds IM.! (k+1)
            w_min_j = windowMin IM.! (k+1)

    -- maximum length. computed from left to right
    w_max k | k == 0    = if msbIsOne then firstWinLength else w
            | otherwise = clamp (i + w_max_i - w -j)
        where
            i       = windowEnds IM.! (k-1)
            j       = windowEnds IM.! k
            w_max_i = windowMax IM.! (k-1)

countPossibilities :: BitLength -> Width -> WindowLengthRanges -> Integer
countPossibilities bits w lengths = go (bits-1) 0
  where
    vec = V.fromList lengths

    pos k = fst (vec V.! k)
    windowLengthRange k = [minWin..maxWin]
      where
        (_, (minWin, maxWin)) =  vec V.! k

    go soonestOne k | k == V.length vec = 2^0
    go soonestOne k = sum
        [ countMemo k winLength
        | winLength <- windowLengthRange k
        , let firstOneOfWindow = pos k + winLength -1
        , firstOneOfWindow <= soonestOne
        ]

    -- simple memoization
    countTable = V.fromList
        [ count k winLength | k <- [0..V.length vec -1] , winLength <- [1..w] ]
    countMemo k winLength = countTable V.!  (k * w + (winLength - 1))

    count k winLength = 2^unknown * go nextPossibleOne (k+1)
      where
        firstOneOfWindow = pos k + winLength -1
        unknown = max 0 (winLength - 2)
        nextPossibleOne = firstOneOfWindow - w

testCountPossibilities :: BitLength -> Width -> Integer -> Property
testCountPossibilities bits w s = fromIntegral (length candidates) === poss
  where
    leak = leftToRight bits w s
    windowLengths = possibleWindowLengths bits w leak
    poss = countPossibilities bits w windowLengths

    mask = getHardBits bits w windowLengths
    candidates = realCandidates bits w leak mask

getWindowStarts :: WindowLengthRanges -> IS.IntSet
getWindowStarts list = IS.fromList
    [ pos + maxWin - 1
    | (pos, (minWin, maxWin)) <- list
    ]

getHardBits :: BitLength -> Width -> WindowLengthRanges -> BitMask
getHardBits bits w final_list = go 0 (nothingKnown bits)
      where
        go i !bm | i == bits = bm
                 | i `IS.member` ones    = go (i+1) (writeKnownBit i True bm)
                 | i `IS.member` zeros   = go (i+1) (writeKnownBit i False bm)
                 | i `IS.member` inRange = go (i+1) bm
                 | otherwise             = go (i+1) (writeKnownBit i False bm)

        onesOn =     IS.fromList [ pos
                                 | (pos,(minWin, maxWin)) <- final_list ]
        onesStart =  IS.fromList [ pos + minWin - 1
                                 | (pos,(minWin, maxWin)) <- final_list
                                 , minWin == maxWin ]
        zerosAfter = IS.fromList [ j
                                 | (pos,(minWin, maxWin)) <- final_list
                                 , j <- [pos + maxWin - w .. pos - 1] ]
        inRange =    IS.fromList [ j
                                 | (pos,(minWin, maxWin)) <- final_list
                                 , j <- [pos .. pos + maxWin - 1] ]

        ones  = IS.unions [ onesOn, onesStart ]
        zeros = IS.unions [ zerosAfter ]

countHardBits :: BitLength -> Width -> WindowLengthRanges -> Int
countHardBits bits w final_list = bits - sum (map uncertain final_list)
  where uncertain (_, (minWin, maxWin)) | minWin == maxWin = max 0 (maxWin - 2)
                                        | otherwise        = max 0 (maxWin - 1)

testCountHardBits bits w final_list =
    countKnownBits bits (getHardBits bits w final_list) == countHardBits bits w final_list

knownBitsFromLeak :: BitLength -> Width -> Leak -> BitMask
knownBitsFromLeak bits w leak = getHardBits bits w $ possibleWindowLengths bits w leak

testKnownBitsFromLeak :: BitLength -> Positive Integer -> Bool
testKnownBitsFromLeak bits (Positive p) = p `matches` bm
  where
    leak = leftToRight bits 4 p
    bm = knownBitsFromLeak bits 4 leak

testKnownBitsFromLeak2 :: RSAParams -> Bool
testKnownBitsFromLeak2 params
    = dp params `matches` bitmask_dp && dq params `matches` bitmask_dq
  where
    leak_dp    = leftToRight       (bits params) 4 (dp params)
    bitmask_dp = knownBitsFromLeak (bits params) 4 leak_dp
    leak_dq    = leftToRight       (bits params) 4 (dq params)
    bitmask_dq = knownBitsFromLeak (bits params) 4 leak_dq

type Candidate = (Int, Integer, Integer, Integer, Integer)

-- | An oracle is a function that takes a possible suffix (given with length
-- and value) and indicates if this suffix is possible.
type Oracle = Int -> Integer -> Bool

selectPortion :: Double -> Int -> [Int]
selectPortion h to = go 0 0
  where
    go chosen i | i > to    = []
                | f < h     = i : go (chosen + 1) (i+1)
                | otherwise =     go  chosen      (i+1)
      where f = fromIntegral chosen / fromIntegral i

selectRandomPortion :: StdGen -> Double -> Int -> [Int]
selectRandomPortion rng h to = go rng 0
  where
    go rng i | i > to    = []
             | x < h     = i : go rng' (i+1)
             | otherwise =     go rng' (i+1)
      where (x, rng') = random rng

mkHardBitOracle :: BitLength -> Double -> Integer -> Oracle
mkHardBitOracle bits h input = oracle
  where
    rng = mkStdGen (fromIntegral (input+3))
    knownBits = IS.fromList $ selectRandomPortion rng h bits
    oracle i guess
        | i `IS.member` knownBits
        = guess `testBit` i == input `testBit` i
        | otherwise
        = True

mkBlockSumOracle :: Int -> Integer -> Oracle
mkBlockSumOracle bs input = oracle
  where
    oracle i guess | j `mod` bs == 0 = guess ==? input
                   | otherwise       = True
      where
        j = i + 1
        select k = (k `div` 2^(j-bs)) `mod` 2^bs
        (==?) = ((==) `on` popCount) `on` select

testOracle :: (Integer -> Oracle) -> Positive Integer -> Positive Int ->  Bool
testOracle mkO (Positive input) (Positive i) = o i (input `mod` 2^(i+1))
  where o = mkO input

mkLeftToRightOracle bits w input = oracle
  where
   s = leftToRight bits w input

   windowLengths = possibleWindowLengths bits w s
   bitmask = getHardBits bits w windowLengths
   windowStarts = getWindowStarts windowLengths

   shouldCheck i = (i-1) `IS.member` windowStarts

   oracle i guess
        | (guess `testBit` i) `notElem` possibleBits bitmask i
        = False
        | shouldCheck i
        , let leak' = leftToRight i w guess
        , V.drop (bits - i) s /= leak'
        = False
        | otherwise
        = True

-- Recovers the bits from the leaks for p and q
lsbRecovery :: Int -> RSAParams -> (Oracle, Oracle) -> Maybe Int
lsbRecovery limit params (op, oq) = dfSearch limit s0 next done
  where
    s0 :: Candidate
    s0 = ( 0 -- last bit guessed/set; this and all later bits are known
         , 1 -- known suffix of dp
         , 1 -- known suffix of dq
         , 1 -- known suffix of p
         , 1 -- known suffix of q
         )

    next :: Candidate -> [Candidate]
    next (i, dp, dq, p, q)
        | j == (bits params) = []
        | otherwise =
            filter prune
            [ (j, dp', dq', p', q')
            | (p',dp') <- filter (pruneP j) $ (,) <$> nextP <*> nextDP
            , (q',dq') <- filter (pruneQ j) $ (,) <$> nextQ <*> nextDQ
            ]
      where
        j = i+1
        nextDP = [dp, dp `setBit` j]
        nextDQ = [dq, dq `setBit` j]
        nextP  = [p, p `setBit` j]
        nextQ  = [q, q `setBit` j]

    pruneP i (p, dp) = and
        [ (e params*dp - 1) `mod` modulus == (kp params * (p-1)) `mod` modulus
        , op i dp
        ]
      where modulus = 2^(i+1)

    pruneQ i (q, dq) = and
        [ (e params*dq - 1) `mod` modulus == (kq params * (q-1)) `mod` modulus
        , oq i dq
        ]
      where modulus = 2^(i+1)

    prune :: Candidate -> Bool
    prune (i, dp, dq, p, q)
        = and [ (p * q) `mod` modulus == n params `mod` modulus
              ]
      where modulus = 2^(i+1)

    done :: Candidate -> Bool
    done (i, dp', dq', p', q')
        = i + 1 == bits params
        && p' * q' == n params
        {-
        && Just dp' == invertMod (e params) (p'-1)
        && Just dq' == invertMod (e params) (q'-1)
        && (dp' == dp params || traceShow ('P',p',q', dp', dq', params) True)
        && (dq' == dq params || traceShow ('Q',p',q', dp', dq', params) True)
        && (dp' `matches` bitmask_dp || traceShow (bits params, dp params) True)
        && (dq' `matches` bitmask_dq || traceShow (bits params, dq params) True)
        -}

runLSBRecovery :: Int -> Width -> RSAParams -> Maybe Int
runLSBRecovery limit w params =
    let op = mkLeftToRightOracle (bits params) w (dp params)
        oq = mkLeftToRightOracle (bits params) w (dq params)
    in lsbRecovery limit params (op, oq)

main :: IO ()
main = join . customExecParser (prefs showHelpOnEmpty) $
  info (helper <*> parser)
  (  fullDesc
  <> header "Sliding right experiment"
  <> progDesc "Slides stuff around"
  <> footer "Some parameter understand range syntax in the form 1,2,3,4 or 4,8,..,2048"
  )
  where
    parser :: Parser (IO ())
    parser =
      hsubparser (mconcat
        [ command "run_recover"        recoverCommand
        , command "left_to_right"      printLTRCommand
        , command "test_hard_rules"    testHardRulesCommand
        , command "measure_leakage"    measureLeakageCommand
        , command "collision_entropy"  measureCollisionEntropyCommand
        , command "hard_soft_scatter"  hardSoftScatterCommand
        , command "leak_preinput_size" leakPreinputCommand
        , command "entropy_scatter"    entropyScatterCommand
        ])

    bitsArg = (`div` 2) <$> -- internally, we want the key length of p and q
      option auto
      (  long "bits"
         <> metavar "BITS"
         <> help "Bit length of p and q combined (usually 1024 or 2048)"
         <> value 1024
         <> showDefault
      )

    bitsRangeArg = (map (`div` 2)) <$> -- internally, we want the key length of p and q
      option (eitherReader parseRange)
      (  long "bits"
         <> metavar "BITS,.."
         <> help "Bit length of p and q combined. Range syntax allowed"
         <> value [1024]
         <> showDefaultWith showRange
      )

    knownBitsRangeArg =
      option (eitherReader parseRange)
      (  long "known-bits"
         <> metavar "PROB,.."
         <> help "Proportion of known bits in range 0...1000. Range syntax allowed"
         <> value [0,100..1000]
         <> showDefaultWith showRange
      )

    blockSumRangeArg =
      option (eitherReader parseRange)
      (  long "sum-length"
         <> metavar "INT,.."
         <> help "Lenght of blocks of which sum is known"
         <> value [1,2,3,4]
         <> showDefaultWith showRange
      )

    widthArg = option auto
         (  long "width"
         <> metavar "WIDTH"
         <> help "Width of the window (e.g. 4 or 5)"
         <> value 4
         <> showDefault
         )

    widthRangeArg = option (eitherReader parseRange)
         (  long "width"
         <> metavar "WIDTH,.."
         <> help "Width of the window (e.g. 4 or 5). Range syntax allowed."
         <> value [4]
         <> showDefaultWith showRange
         )

    iterationsArg = option auto
         (  long "iterations"
         <> metavar "N"
         <> help "Number of iterations"
         <> value 100
         <> showDefault
         )

    limitArg = option auto
         (  long "limit"
         <> metavar "N"
         <> help "Search limit"
         <> value 100000
         <> showDefault
         )

    tsvArg = switch
         (  long "tsv"
         <> help "Tab separated output (no variance)"
         )

    printLTRCommand = info
        (pure printLTR <*> bitsArg <*> widthArg<*> bitSequenceArg)
        (progDesc "Run left-to-right and print sequence")
      where bitSequenceArg = argument auto (metavar "input_sequence")

    testHardRulesCommand = info
        (pure testHardRules <*> bitsArg <*> widthArg<*> verboseArg)
        (progDesc "Verify hard bit calculation (no output = all good)")
      where verboseArg = switch (long "verbose" <> help "Show all tested data")

    recoverCommand = info
        (pure runRecover <*> bitsArg <*> widthArg<*> iterationsArg <*> limitArg )
        (progDesc "Run key recovery, and calculate search tree size")

    measureLeakageCommand = info
        (pure measureLeakage <*> bitsRangeArg <*> widthRangeArg<*> tsvArg <*> iterationsArg)
        (progDesc "Calculate the Shannon entropy rate and hard bits")

    measureCollisionEntropyCommand = info
        (pure measureCollisionEntropy <*> bitsRangeArg <*> widthRangeArg<*> tsvArg <*> iterationsArg)
        (progDesc "Calculate the collision entropy rate")

    hardSoftScatterCommand = info
        (pure hardSoftScatter <*> bitsRangeArg <*> widthRangeArg<*> iterationsArg)
        (progDesc "Draw a hard/soft bit scatter plot")

    leakPreinputCommand = info
        (pure leakPreinput <*> widthArg)
        (progDesc "Number of input sequences for given square-and-mulitply sequence")

    entropyScatterCommand = info
        (pure runScatter <*> bitsArg <*> knownBitsRangeArg <*> widthRangeArg <*> blockSumRangeArg <*>iterationsArg <*> limitArg )
        (progDesc "Simulates recovery with various oracles of given entropy, measuing tree size")

runRecover :: BitLength -> Width -> Int -> Int -> IO ()
runRecover  bits w iter limit = do
    samples <- replicateM iter $ do
       params <- generate (genRSAParams bits)
       let res = runLSBRecovery limit w params
       putStrLn $ maybe "-" show res
       return res
    let (mean, variance) = meanVariance (V.fromList (map fromIntegral (catMaybes samples)))
    putStrLn $ "-- Mean: " ++ show mean ++ " Variance: " ++ show variance

popCountCollEntropyRate :: Int -> Double
popCountCollEntropyRate w = (/fromIntegral w) $ negate $ logBase 2 $ sum [
    (fromIntegral (binom w i) / 2^w)^2
    | i <- [0..w]
    ]

binom :: Int -> Int -> Int
binom n m = fac n `div` (fac m * fac (n-m))

fac :: Int -> Int
fac n = product [1..n]

runScatter :: BitLength -> [Int] -> [Width] -> [Int] -> Int -> Int -> IO ()
runScatter  bits knownBitsRatios widths blockSums iter limit = do
    printf "leak\th\tmean\tsolved\tcalc_mean\n"
    let todo = [ (p, o, "hard_bits")
               | ratio <- sort knownBitsRatios
               , let p = fromIntegral ratio / 1000 :: Double
               , let o = mkHardBitOracle bits p
               ] ++
               [ (h, o, "ltr")
               | w <- sortBy (flip compare) widths
               , let h = [undefined, 1, 0.786, 0.652, 0.545, 0.461, 0.395, 0.345] !! w
               , let o = mkLeftToRightOracle bits w
               ] ++
               [ (h, o, "popcnt")
               | bs <- sortBy (flip compare) blockSums
               , let h = popCountCollEntropyRate bs
               , let o = mkBlockSumOracle bs
               ]
    forM_ todo $ \(h, mkO, desc) -> do
        samples <- replicateM iter $ do
           params <- generate (genRSAParams bits)
           let op = mkO (dp params)
               oq = mkO (dq params)
               res = lsbRecovery limit params (op, oq)
           -- putStrLn $ maybe "-" show res
           return res
        let (mean, variance) = meanVariance (V.fromList (map fromIntegral (catMaybes samples)))
        let solved = fromIntegral (length (filter isJust samples)) / fromIntegral iter :: Double
        printf "%s\t%.3f\t%.1f\t%.3f\t%.1g\n"
            desc h mean solved (treeSize bits h)

treeSize :: BitLength -> Double -> Double
treeSize k h = k' * (1 + (1 - 2**(k'*(1-2*h)))/(1-2**(1-2*h)))
  where k' = 2*fromIntegral k

printLargePrime :: BitLength -> Width -> IO ()
printLargePrime bits w = generate (rndPrime bits) >>= print

printRSA :: BitLength -> Width -> IO ()
printRSA bits w = generate (genRSAParams bits) >>= print

showLeak :: Leak -> String
showLeak = map go . V.toList
  where go S = '_'
        go M = 'M'

printLTR :: BitLength -> Width -> Integer -> IO ()
printLTR bits w s = putStrLn $ showLeak $ leftToRight bits w s

realCandidates :: BitLength -> Width -> Leak -> BitMask -> [Integer]
realCandidates bits w leak mask = ms'
   where ms    = allMatching mask
         ms'   = filter (\m -> leftToRight bits w m == leak) ms

reallyKnownBits :: BitLength -> Width -> Leak -> BitMask -> BitMask
reallyKnownBits bits w leak mask = mask'
   where ms'   = realCandidates bits w leak mask
         mask' = maskFromMatching ms'

testKnownBits bits w n = do
    putStrLn (showBitMask bits (fromKnownNumber n))
    let leak = leftToRight bits w n
    putStrLn (showLeak leak)
    let mask = knownBitsFromLeak bits w leak
    putStrLn (showBitMask bits mask)
    let mask' = reallyKnownBits bits w leak mask
    putStrLn (showBitMask bits mask')

testHardRules :: BitLength -> Width -> Bool -> IO ()
testHardRules bits w verbose = forM_ [2^(bits-1)..(2^bits)-1] $ \n -> do
    let leak  = leftToRight bits w n
    let mask  = knownBitsFromLeak bits w leak
    let ms    = allMatching mask
    let ms'   = filter (\m -> leftToRight bits w m == leak) ms
    let mask' = maskFromMatching ms'

    when (verbose || mask /= mask' || n `notElem` ms) $ do
        putStr $ showBitMask bits (fromKnownNumber n) ++ ": "
        putStr $ showLeak leak ++ " "
        putStr $ showBitMask bits mask ++ " "
        putStr $ showBitMask bits mask' ++ " "
        putStr $ if mask == mask' then "✓" else "☹"
        unless (n `elem` ms) $ putStr " HORROR"
        putStr "\n"

measureCollisionEntropy :: [BitLength] -> [Width] -> Bool -> Int -> IO ()
measureCollisionEntropy bitsRange wRange tsv iter= do
    when tsv $ putStrLn $ intercalate "\t" $
        words "bits width entropy"
    sequence_ [ measureCollisionEntropyOne bits w tsv iter
              | bits <- bitsRange, w <- wRange ]

measureCollisionEntropyOne :: BitLength -> Width -> Bool -> Int -> IO ()
measureCollisionEntropyOne bits w tsv iter = do
    probs <- replicateM iter doOne
    -- There are serious problems here with numerical precision
    -- and the results seem to vary a log depending on how the logarithm
    -- is calculated
    let entropy = - (logBase 2.0 (sum probs) - fromIntegral (2*bits) - logBase 2.0 (fromIntegral iter)) / fromIntegral (2*bits)
    if tsv then
        putStrLn $ intercalate "\t" $ map show [2*bits, w] ++ map show [entropy]
    else putStrLn $ "Collision entropy rate:" ++ show entropy
  where
    stats :: V.Vector Double -> [Double]
    stats samples =
        [ minimum samples ] ++
        [ weightedAvg n 4 samples | n <- [1..3]] ++
        [ maximum samples ]
    doOne = do
        dp <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)
        dq <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)

        let sp = leftToRight bits w dp
            wlp = possibleWindowLengths bits w sp
            hardBitsp = countHardBits bits w wlp
            possp = countPossibilities bits w wlp

            sq = leftToRight bits w dq
            wlq = possibleWindowLengths bits w sq
            hardBitsq = countHardBits bits w wlq
            possq = countPossibilities bits w wlq

            totalBits = fromIntegral (2*bits)
            prob = fromIntegral possp * fromIntegral possq
        return prob

measureLeakage :: [BitLength] -> [Width] -> Bool -> Int -> IO ()
measureLeakage bitsRange wRange tsv iter= do
    when tsv $ putStrLn $ intercalate "\t" $
        words "bits width" ++
        concatMap (\n -> [n ++ "_q" ++ show i | i <- [0..4]])
            (words "rightToLeft hardBitsLearned informationGained GMP")
    sequence_ [ measureLeakageOne bits w tsv iter
              | bits <- bitsRange, w <- wRange ]

measureLeakageOne :: BitLength -> Width -> Bool -> Int -> IO ()
measureLeakageOne bits w tsv iter= do
    (samplesRTL, samplesHard, samplesReal, samplesGMP) <- V.unzip4 . V.fromList <$> replicateM iter doOne
    if tsv then do
        putStrLn $ intercalate "\t" $
            map show [2*bits, w] ++
            map show (concatMap stats [ samplesRTL, samplesHard, samplesReal, samplesGMP])
    else do
        let (mean, variance) = meanVariance samplesHard
        let good = V.sum (V.map (\n -> if n > 0.5 then 1 else 0) samplesHard) / fromIntegral iter
        putStrLn $ "Hard bits mean:  " ++ show mean ++
                   " variance: " ++ show variance ++
                   " >50%: " ++ show good
        let (mean, variance) = meanVariance samplesReal
        let good = V.sum (V.map (\n -> if n > 0.5 then 1 else 0) samplesReal) / fromIntegral iter
        putStrLn $ "Total data mean: " ++ show mean ++
                   " variance: " ++ show variance ++
                   " >50%: " ++ show good
  where
    stats :: V.Vector Double -> [Double]
    stats samples =
        [ minimum samples ] ++
        [ weightedAvg n 4 samples | n <- [1..3]] ++
        [ maximum samples ]
    doOne = do
        dp <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)
        dq <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)

        let sp = leftToRight bits w dp
            wlp = possibleWindowLengths bits w sp
            hardBitsp = countHardBits bits w wlp
            possp = countPossibilities bits w wlp

            sq = leftToRight bits w dq
            wlq = possibleWindowLengths bits w sq
            hardBitsq = countHardBits bits w wlq
            possq = countPossibilities bits w wlq

            totalBits = fromIntegral (2*bits)
            hardBitsFrac = fromIntegral (hardBitsp + hardBitsq) / totalBits

            realBitsFrac = ( totalBits - logBase 2.0 (fromIntegral possp) - logBase 2.0 (fromIntegral possq) ) / totalBits

            rtlBitsFrac = ( fromIntegral $ rightToLeftBits bits w dp + rightToLeftBits bits w dq
                          ) / totalBits

            gmpBitsFrac = ( fromIntegral $ leftToRightGMPBits bits w dp + leftToRightGMPBits bits w dq
                          ) / totalBits

        return (rtlBitsFrac, hardBitsFrac, realBitsFrac, gmpBitsFrac)

hardSoftScatter :: [BitLength] -> [Width] -> Int -> IO ()
hardSoftScatter bitsRange wRange iter = do
    putStrLn $ intercalate "\t" $ words "bits width hard soft"
    forM_ bitsRange $ \bits -> forM_ wRange $ \w -> do
        pairs <- replicateM iter (doOne bits w)
        forM_ pairs $ \(h,s) ->
            putStrLn $ intercalate "\t" (map show [bits, w] ++ map show [h,s])
  where
    doOne bits w = do
        dp <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)
        dq <- fmap (.|. 1) $ randomRIO (2^(bits - 1), 2^bits - 1)

        let sp = leftToRight bits w dp
            wlp = possibleWindowLengths bits w sp
            hardBitsp = countHardBits bits w wlp
            possp = countPossibilities bits w wlp

            sq = leftToRight bits w dq
            wlq = possibleWindowLengths bits w sq
            hardBitsq = countHardBits bits w wlq
            possq = countPossibilities bits w wlq

            totalBits = fromIntegral (2*bits) :: Double
            hardBitsFrac = fromIntegral (hardBitsp + hardBitsq) / totalBits

            realBitsFrac = ( totalBits - logBase 2.0 (fromIntegral possp) - logBase 2.0 (fromIntegral possq) ) / totalBits

            rtlBitsFrac = ( fromIntegral $ rightToLeftBits bits w dp + rightToLeftBits bits w dq
                          ) / totalBits

            gmpBitsFrac = ( fromIntegral $ leftToRightGMPBits bits w dp + leftToRightGMPBits bits w dq
                          ) / totalBits

        return (hardBitsFrac, realBitsFrac)

leakPreinput ::  Width -> IO ()
leakPreinput w = do
    line <- getLine
    let line' = fixFirst line
    unless (all (`elem` "SM_") line') $ do
        hPutStrLn stderr "Error: Input is not a sequence of 'S', '_' or 'M'"
        exitFailure
    let bits = length line'
        sp = V.fromList (map toStep line')
        wlp = possibleWindowLengths bits w sp
            --hardBitsp = countHardBits bits w wlp
        possp = countPossibilities bits w wlp
    print possp
  where
    toStep 'M' = M
    toStep 'S' = S
    toStep '_' = S

    -- The input may explicitly say that the first character is a 1.
    -- This code knows that, but still expects there to be a S or M.
    -- Recover whether this is a S or M from looking at the following characters.
    fixFirst ('1':xs) = if 'M' `elem` take (w - 1) xs then 'S':xs else 'M':xs
    fixFirst xs = xs

showRange :: [Int] -> String
showRange = intercalate "," . map show

parseRange :: String -> Either String [Int]
parseRange s = case splitOn "," s of
    [n1,n2,"..",n3] -> enumFromThenTo <$> readEither n1 <*> readEither n2 <*> readEither n3
    [n1,"..",n3] -> enumFromTo <$> readEither n1 <*> readEither n3
    ns -> mapM readEither ns


{-# LANGUAGE FunctionalDependencies, FlexibleInstances, FlexibleContexts, BangPatterns #-}

module AI.Learning.LinearRegression where

import Data.List (foldl')
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Util (ones)

import AI.Util.Matrix

{-

A linear model should contain the following information:

* Regression coefficients
* Covariance matrix of regression coefficients
* Residuals
* Confidence intervals for regression coefficients
* Confidence intervals for residuals
* T-Statistic for regression coefficients
* P-Values for regression coefficients
* F-Statistic for regression
* P-Value for regression
* MSE
* SST
* SSE
* PRESS Statistic
* R squared
* adjusted r squared
*

-}

-------------
-- Options --
-------------

-- |Type recording what kind of linear model we want to build. Choose from
--  ordinary least squares (OLS), ridge regression with parameter lambda
--  and LASSO with parameter lambda.
data LMType = OLS
            | Ridge Double
            | LASSO Double
            deriving (Show)

---- |Data type for linear regression options.
data LMOpts = LMOpts { fitIntercept :: Bool
                     , standardizeRegressors :: Bool }
                     deriving (Show)

---- |Standard options for linear regression (used by 'regress').
stdLMOpts = LMOpts { fitIntercept = True
                   , standardizeRegressors = False }

---------------------
---- Linear Models --
---------------------

-- |Data type for a linear model. It consists of the coefficient vector,
--  together with options that specify how to transform any data that is used
--  to make new predictions.
data LinearModel = LM { coefs   :: Vector Double
                      , lmIntercept :: Bool
                      , lmStandardize :: Bool
                      , lmMean :: Vector Double
                      , lmStd  :: Vector Double }
                      deriving (Show)

-- |Statistics structure for a linear regression.
data LMStats = LMStats { covBeta :: Maybe (Matrix Double)
                       , betaCI :: Maybe (Matrix Double)
                       , sst :: Double
                       , sse :: Double
                       , mse :: Double
                       , rSquare :: Double
                       , tBeta :: Maybe (Vector Double)
                       , pBeta :: Maybe (Vector Double)
                       , fRegression :: Maybe Double
                       , pRegression :: Maybe Double }
                       deriving (Show)

type Predictor = Matrix Double
type Response = Vector Double

-- |Fit an ordinary least squares linear model to data.
lm :: Matrix Double -> Vector Double -> LinearModel
lm = lmWith OLS stdLMOpts

-- |Fit a ridge regression (least squares with quadratic penalty) to data.
lmRidge :: Matrix Double -> Vector Double -> Double -> LinearModel
lmRidge x y lambda = lmWith (Ridge lambda) stdLMOpts x y

-- |Fit a linear model to data. The exact model fitted is specified by the
--  'LMType' argument, and the options are specified in 'LMOpts'.
lmWith :: LMType -> LMOpts -> Matrix Double -> Vector Double -> LinearModel
lmWith kind opts x y = LM { coefs = beta
                          , lmIntercept = fitIntercept opts
                          , lmStandardize = standardizeRegressors opts
                          , lmMean = mu
                          , lmStd = sigma }
    where
        (xx,mu,sigma) = lmPrepare opts x
        beta = case kind of
            OLS     -> regress xx y
            Ridge a -> ridgeRegress xx y (fitIntercept opts) a
            LASSO a -> error "LASSO not implemented."

-- |Prepare data according to the specified options structure. This may involve
--  centering and standardizing the data, or adding a column of constants.
lmPrepare :: LMOpts
          -> Matrix Double
          -> (Matrix Double, Vector Double, Vector Double)
lmPrepare opts x = (x3,mu,sigma)
    where
        (x1,mu,sigma) = standardize x
        x2            = if standardizeRegressors opts then x1 else x
        x3            = if fitIntercept opts then addOnes x2 else x2

---------------------------------
---- Predict from linear model --
---------------------------------

-- |Make predictions from a linear model.
lmPredict :: LinearModel -> Matrix Double -> Vector Double
lmPredict model x = x2 <> beta
    where
        beta    = coefs model
        xbar    = lmMean model
        sigma   = lmStd model
        x1      = if lmStandardize model
                    then eachRow (\x -> (x - xbar) / sigma) x
                    else x
        x2      = if lmIntercept model
                    then addOnes x1
                    else x1

-- |Calculate statistics for a linear regression.
lmStats :: LinearModel -> Matrix Double -> Vector Double -> LMStats
lmStats model x y =
    let ybar      = constant (mean y) (dim y)
        yhat      = lmPredict model x
        residuals = y - yhat
        covBeta   = Nothing
        betaCI    = Nothing
        sst       = sumVector $ (y - ybar) ^ 2
        sse       = sumVector $ residuals ^ 2
        mse       = sse / fromIntegral (rows x)
        rSq       = 1 - sse / sst
        tBeta     = Nothing
        pBeta     = Nothing
        fReg      = Nothing
        pReg      = Nothing
    in LMStats covBeta betaCI sst sse mse rSq tBeta pBeta fReg pReg

mseEvalFun :: Response -> Response -> Double
mseEvalFun actual predicted = mean $ (actual - predicted) ^ 2

lmPredFun :: Predictor -> Response -> Predictor -> Response -> Double
lmPredFun xtrain ytrain xtest ytest = mseEvalFun ytest ypred
    where
        ypred = lmPredict model xtest
        model = lm xtrain ytrain

--------------------------
---- Perform Regression --
--------------------------

---- |Regress a vector y against a matrix of predictors x, with the specified
----  options.
regress :: Matrix Double -- X
        -> Vector Double -- y
        -> Vector Double -- beta
regress x y
    | rows x /= dim y = error "Inconsistent dimensions -- REGRESS"
    | otherwise = let (_,n) = size x
                      (_,r) = qr x
                      rr = takeRows n r
                   in (trans rr <> rr) <\> trans x <> y

-- |Ridge regression. This adds a penalty term to OLS regression, which
--  discourages large coefficient values to prevent overfitting.
ridgeRegress :: Matrix Double  -- X
             -> Vector Double  -- y
             -> Bool           -- useConst?
             -> Double         -- lambda
             -> Vector Double  -- beta
ridgeRegress x y useConst lambda
    | rows x /= dim y = error "Inconsistent dimensions -- RIDGEREGRESS"
    | otherwise = let (_,n) = size x
                      (_,r) = qr x
                      rr = takeRows n r
                      ww = if useConst
                            then diag $ join [0, constant 1 (n-1)]
                            else ident n
                   in (trans rr <> rr + lambda `scale` ww) <\> trans x <> y

-----------------
---- Utilities --
-----------------

-- |De-mean a sample.
demean :: Matrix Double -> Matrix Double
demean x = eachRow (subtract $ mean x) x

-- |Standardize a sample to have zero mean and unit variance. Returns a triple
--  consisting of the standardized sample, the sample mean and the sample
--  standard deviation.
standardize :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
standardize m = (eachRow (\x -> (x - mu) / sigma) m, mu, sigma)
    where mu    = mean m
          sigma = std m

-- |Standardize a sample to have zero mean and unit variance. This variant of
--  the function discards the mean and standard deviation vectors, only
--  returning the standardized sample.
standardize_ :: Matrix Double -> Matrix Double
standardize_ x = a where (a,_,_) = standardize x

-------------------------
-- Mean, Variance etc. --
-------------------------

class Mean a b | a -> b where
    mean :: a -> b

class Floating b => Variance a b | a -> b where
    var :: a -> b
    var x = y * y where y = std x

    std :: a -> b
    std x = sqrt $ var x

instance Mean (Vector Double) Double where
    mean v = sumVector v / fromIntegral (dim v)

instance Variance (Vector Double) Double where
    var v  = mean $ (v - constant vbar (dim v)) ^ 2
        where vbar = mean v

instance Mean (Matrix Double) (Vector Double) where
    mean m = fromList $ mapCols mean m

instance Variance (Matrix Double) (Vector Double) where
    var m  = fromList $ mapCols var m


import Data.Complex

direction :: Char -> Complex Float
direction 'U' = 0 :+ 1
direction 'D' = 0 :+ (-1)
direction 'L' = (-1) :+ 0
direction 'R' = 1 :+ 0

clamp :: Float -> Float
clamp = max (-1) . min 1

advance :: Complex Float -> Complex Float -> Complex Float
advance pos dir =
  let (x :+ y) = pos + dir
  in clamp x :+ clamp y

complexToButton :: Complex Float -> Char
complexToButton ((-1) :+   1)  = '1'
complexToButton (  0  :+   1)  = '2'
complexToButton (  1  :+   1)  = '3'
complexToButton ((-1) :+   0)  = '4'
complexToButton (  0  :+   0)  = '5'
complexToButton (  1  :+   0)  = '6'
complexToButton ((-1) :+ (-1)) = '7'
complexToButton (  0  :+ (-1)) = '8'
complexToButton (  1  :+ (-1)) = '9'

part1 :: String -> String
part1 input = do
  line <- lines input
  let button = complexToButton $ foldl advance 0 $ map direction line
  return button

complexToButton2 :: Complex Float -> Char
complexToButton2 (  0  :+   2)  = '1'
complexToButton2 ((-1) :+   1)  = '2'
complexToButton2 (  0  :+   1)  = '3'
complexToButton2 (  1  :+   1)  = '4'
complexToButton2 ((-2) :+   0)  = '5'
complexToButton2 ((-1) :+   0)  = '6'
complexToButton2 (  0  :+   0)  = '7'
complexToButton2 (  1  :+   0)  = '8'
complexToButton2 (  2  :+   0)  = '9'
complexToButton2 ((-1) :+ (-1)) = 'A'
complexToButton2 (  0  :+ (-1)) = 'B'
complexToButton2 (  1  :+ (-1)) = 'C'
complexToButton2 (  0  :+ (-2)) = 'D'
complexToButton2 _              = '0' -- default

advance2 :: Complex Float -> Complex Float -> Complex Float
advance2 pos dir =
  let pos' = pos + dir
  in if complexToButton2 pos' == '0'
  then pos
  else pos'

part2 :: String -> String
part2 input = do
  line <- lines input
  let button = complexToButton2 $ foldl advance2 0 $ map direction line
  return button


module Corfunc (corfunc,fitdata) where

import Data.Complex
import Numeric.FFT

porod :: Double -> Double -> Double -> Double
porod k sigma q = (k*q**(-4))*exp (-q**2*sigma)

guinier :: Double -> Double -> Double -> Double
guinier a b q = a * exp (b*q**2)

beta :: [Double] -> [Double] -> Double
beta x y = rval x y * std y / std x

alpha :: [Double] -> [Double] -> Double
alpha x y = mean y - beta x y * mean x

mean :: [Double] -> Double
mean xs = sum xs / fromIntegral (length xs)

std :: [Double] -> Double
std xs = sqrt $ mean (map (**2) xs) - mean xs**2

rval :: [Double] -> [Double] -> Double
rval xs ys = sum (zipWith (*) xdiff ydiff) / sqrt (sum (map (**2) xdiff) * sum (map (**2) ydiff))
    where
      diff qs = map (\z -> z - mean qs) qs
      xdiff = diff xs
      ydiff = diff ys

fitguinier :: [Double] -> [Double] -> Double -> Double
fitguinier qs is = guinier a b
    where
      xs = map (**2) qs
      ys = map log is
      a = exp $ alpha xs ys
      b = beta xs ys

fitporod :: [Double] -> [Double] -> Double -> Double
fitporod qs is = porod k sigma
    where
      is' = filter (>0) is
      qs' = filterKey (>0) is qs
      xs = map (**2) qs'
      ys = map log . zipWith (*) is' . map (**4) $ qs'
      k = exp $ alpha xs ys
      sigma = 0 * beta xs ys

filterKey :: (a->Bool) -> [a] -> [b] -> [b]
filterKey f ks = map snd . filter (f.fst) . zip ks

takeKey :: (a->Bool) -> [a] -> [b] -> [b]
takeKey f ks = map snd . takeWhile (f . fst) . zip ks

dropKey :: (a->Bool) -> [a] -> [b] -> [b]
dropKey f ks = map snd . dropWhile (f . fst) . zip ks

fitdata :: [Double] -> [Double] -> Double -> Double
fitdata qs is = let d = interp qs is
                    maxq = 0.04
                    highq = dropWhile (< maxq) qs
                    highi = dropKey (< maxq) qs is
                    fitp = fitporod highq highi
                    minq = 0.0065*3
                    lowq = takeWhile (< minq) qs
                    lowi = takeKey (< minq) qs is
                    fitg = fitguinier lowq lowi
                    s1 = smooth d fitp (maxq,maximum qs)
                    s2 = smooth fitg s1 (minimum qs,minq)
                in
                  s2

-- | Smoothly transition between two functions over a range of functions
smooth :: (Double -> Double) -> (Double -> Double) -> (Double, Double) -> Double -> Double
smooth f g (start, stop) x
    | x <= start = f x
    | x >= stop = g x
    | otherwise = h*g x +(1-h)*f x
    where
      h = 1/(1+(x-stop)**2/(start-x)**2)

interpHead :: ((Double, Double), (Double, Double)) -> Double -> Double
interpHead ((x1,x2),(y1,y2)) q = (x2-q)/(x2-x1)*y1 + (q-x1)/(x2-x1)*y2

interp :: [Double] -> [Double] -> Double -> Double
interp xs ys x =
    let ranges = zip (zip xs (tail xs)) (zip ys (tail ys))
    in
      case dropWhile (\q -> snd (fst q) < x) ranges of
        [] -> last ys
        p:_ -> interpHead p x

dct :: [Double] -> [Double]
dct = map realPart . fft . map (:+ 0)

corfunc :: [Double] -> [Double] -> [(Double,Double)]
corfunc qs is = let fullData = fitdata qs is
                    factor = 16
                    deltaq = (qs!!1-head qs)/factor
                    qgrid = map (*deltaq) [0..factor*2^(floor . log . (/ deltaq) . maximum $ qs)]
                    igrid = zipWith (*) (map (**2) qgrid) $ map fullData qgrid
                    ys = dct igrid
                    xs = map (\q -> 2 * pi * q / deltaq /fromIntegral (length qgrid)) [0..fromIntegral $ length qgrid]
                in
                  zip xs ys


module FET where

import qualified Data.ByteString.Lazy.Char8             as BS
import           Data.Eq
import           Data.Function
import           Data.Functor
import           Data.Int
import           Data.List
import           Data.Ord
import           Prelude                                (Double, fromIntegral,
                                                         (*), (+), (-), (/),
                                                         (^), otherwise, (&&))
import           Statistics.Distribution                (complCumulative,
                                                         cumulative,
                                                         probability)
import           Statistics.Distribution.ChiSquared
import           Statistics.Distribution.Hypergeometric
import           Text.Show

--the Fisher's exact test can be computed from the hypergeometric distribution
--eg. test whether the presence / absence of a gene is statistically significant
--between two bacterial groups
--           Group 1 | Group 2
--          -------------------
-- has gene |   2    |    7
-- no gene  |   5    |    3
--          -------------------
-- the hypergeometric distribution is calculated in Statistics.Distribution.Hypergeometric
-- hypergeometric = m l k
-- From the package description: The parameters of the distribution describe
-- k elements chosen from a population of l, with m elements of one type,
-- and l-m of the other (all are positive integers).
-- In our case:
-- m is the sum of "has factor" instances, so 9
-- l is the total population size, so 17
-- k is the sum of a particular group, so 7 for Group 1
-- let d = distribution 9 17 7
-- This gives us the distribution, but we need to test the probability
-- that the presence of this gene is statistically different, using the two-tail
-- approach.
-- The one-tail probability in our example of Group 1 having 2 members with the
-- gene can be tested using the cumulative distribution function from
-- Statistics.Distribution, our previously calculated distribution "d", and
-- the number of "successes" (x), which is 2.
-- let oneTail = cumulative d 2 = 0.11703002879473468
-- However, we have no previous expectation for this group to have or not have
-- this gene, so what we really want is the two-tailed test.
-- Please see http://www.pmean.com/11/Fishers.html for an excellent,
-- detailed description.
-- In short, we need to sum all of the probabilities that are smaller than or
-- equal to the probability of the observed table (eg. x = 2 in our example).
-- The twoTailValue is calculated by generating all of the probabilities for
-- our sample, filtering the ones that have a probability too large, and
-- summing the remaining values.
--
-- For a total number of observations greater than 1000, we use the Chi Square
-- test, as the FET returns NaN.
-- ChiSqr statistic  for
--           Group 1 | Group 2
--          -------------------
-- has gene |   2    |    7
-- no gene  |   5    |    3
--          -------------------
-- is [(2*3)-(5*7)]^2 * (2+7+5+3) / (2+7)(5+3)(7+3)(2+5)
-- With 1 degree of freedom, we can calculate the Chi Square distribution
-- with chiSquared 1 from Statistics.Distribution
-- To determine the P value of a Chi Square statistic, we use the
-- complCumulative function and the
-- Chi Square distribution with 1 degree of freedom.
-- For Chi Square, the one-tail is simply half the two-tail
-- Statistics.Test.ChiSquared only return "Significant" or "NotSignificant"
-- and we wanted the actual P-value, hence the function.
-- Likewise there is no Statistics.Test.FET

newtype GroupOneA = GroupOneA { unGroupOneA :: Int } deriving (Eq, Show)
newtype GroupOneB = GroupOneB { unGroupOneB :: Int } deriving (Eq, Show)
newtype GroupTwoA = GroupTwoA { unGroupTwoA :: Int } deriving (Eq, Show)
newtype GroupTwoB = GroupTwoB { unGroupTwoB :: Int } deriving (Eq, Show)
newtype PValue = PValue {unPValue :: Double} deriving (Eq, Show)
newtype FETName = FETName {unFETName :: BS.ByteString} deriving(Eq, Show)
data FETMode = OneTail | TwoTail deriving (Eq,Show)

data FETResult = FETResult{groupOneA :: Int
                          ,groupOneB :: Int
                          ,groupTwoA :: Int
                          ,groupTwoB :: Int
                          ,pvalue    :: Double
                          ,fetName   :: BS.ByteString} deriving (Eq, Show)

fet :: FETName
    -> GroupOneA
    -> GroupOneB
    -> GroupTwoA
    -> GroupTwoB
    -> FETMode
    -> FETResult
fet (FETName fName) (GroupOneA goa) (GroupOneB gob) (GroupTwoA gta) (GroupTwoB gtb) mode
  | l > 0 && goa >= 0 && gob >= 0 && gta >= 0 && gtb >= 0 && k > 0 && gt > 0 =
    case mode of
        OneTail -> theResult{pvalue = cumD}
        TwoTail -> theResult{pvalue = twoTailValue}
  | otherwise = theResult
  where
    cumD = if l >= 1000
            then chiP / 2
            else cumulative d $ fromIntegral goa
    probX = probability d goa
    twoTailValue = if l >= 1000
                      then chiP
                      else sum . filter (<= probX) . fmap (probability d) $ [0..k]
    d = hypergeometric m l k
    chiP = complCumulative d' $ chiSquare goa gob gta gtb
    d' = chiSquared 1
    m = goa + gta
    l = m + gob + gtb
    k = goa + gob
    gt = gta + gtb
    theResult = FETResult{groupOneA = goa
                         ,groupOneB = gob
                         ,groupTwoA = gta
                         ,groupTwoB = gtb
                         ,pvalue = 1
                         ,fetName = fName
                         }

chiSquare :: Int
          -> Int
          -> Int
          -> Int
          -> Double
chiSquare a' b' c' d' = (x' * x' * (a + b + c + d)) /
                        ((a+b) * (c+d) * (b+d) * (a+c))
  where
    x' = (a * d) - (b * c)
    a = fromIntegral a'
    b = fromIntegral b'
    c = fromIntegral c'
    d = fromIntegral d'


module Algoritmo where

import Prelude as P
import Data.Complex
import Control.Monad.State
import Numeric.Statistics (percentile)
import Numeric.LinearAlgebra
import Data.Vector.Storable as VS hiding (modify, null)

import FFT
import ArrayPoking
import Bases

type L2  = Vector (Complex Double)
type Dim = Int
type Niveles = Int

secAnalitica :: Niveles -> L2 -> State [L2] ()
secAnalitica n zF =
  do let dim = VS.length zF
         uF  = conj . fft $ haarU dim
         vF  = conj . fft $ haarV dim
         x1  = downSampF (zF * vF) dim
         y1  = downSampF (zF * uF) dim

     -- [y3, x3, x2, x1]
     -- modify (x1 :)
     if n <= 1
       -- Haz nada, este es el ultimo nivel
       then modify (zF :)
       --
       else modify (x1 :) >> secAnalitica (n - 1) y1

secSintetica :: Niveles -> State [L2] L2
secSintetica n =
  do y2 : x2 : rest <- get
     let dim = VS.length y2
         dim_n = dim * 2
         x2' = upSampF x2 dim
         y2' = upSampF y2 dim
         y1  = x2' * (fft (haarV dim_n)) +
               y2' * (fft (haarU dim_n))
     put (y1 : rest)

     if null rest || n <= 1
       then return y1
       else secSintetica (n - 1)

{-
  Idea: Mantener los @k terminos mayores y lo demas setearlos a 0. Es
  molesto contar, asi que veremos un porcentaje.
  rate = 90.0, es decir 90% de calidad, debo de matar 10% de los
  candidatos.
-}
compresion :: Double -> State [L2] ()
compresion rate =
  do vec <- get
     let cut  = percentile (100 - rate) . sortIP
                . cmap magnitude $ VS.concat vec

         go c = if magnitude c < cut then 0 else c
     put $ P.map (VS.map go) vec

waveletAlgo :: Double -> Niveles -> Niveles -> L2 -> L2
waveletAlgo calidad div recon color = evalState go []
  where
    go :: State [L2] L2
    go = return (fft color)
         >>= secAnalitica div
         -- >> modify (P.map ifft)
         >> compresion calidad
         -- >> modify (P.map fft)
         >> secSintetica recon
         >>= return . ifft