-- |
-- Module:      Math.NumberTheory.DirichletCharacters
-- Copyright:   (c) 2018 Bhavik Mehta
-- Licence:     MIT
-- Maintainer:  Bhavik Mehta <bhavikmehta8@gmail.com>
--
-- Implementation and enumeration of Dirichlet characters.
--

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE ViewPatterns               #-}

module Math.NumberTheory.DirichletCharacters
  (
  -- * Roots of unity
    RootOfUnity
  -- ** Conversions
  , toRootOfUnity
  , fromRootOfUnity
  , toComplex
  -- * An absorbing semigroup
  , OrZero, pattern Zero, pattern NonZero
  , orZeroToNum
  -- * Dirichlet characters
  , DirichletCharacter
  -- ** Construction
  , indexToChar
  , indicesToChars
  , characterNumber
  , allChars
  , fromTable
  -- ** Evaluation
  , eval
  , evalGeneral
  , evalAll
  -- ** Special Dirichlet characters
  , principalChar
  , isPrincipal
  , orderChar
  -- ** Real Dirichlet characters
  , RealCharacter
  , isRealCharacter
  , getRealChar
  , toRealFunction
  , jacobiCharacter
  -- ** Primitive characters
  , PrimitiveCharacter
  , isPrimitive
  , getPrimitiveChar
  , induced
  , makePrimitive
  , WithNat(..)
  -- * Debugging
  , validChar
  ) where

#if !MIN_VERSION_base(4,12,0)
import Control.Applicative                                 (liftA2)
#endif
import Data.Bits                                           (Bits(..))
import Data.Complex                                        (Complex(..), cis)
import Data.Foldable                                       (for_)
import Data.Functor.Identity                               (Identity(..))
import Data.List                                           (mapAccumL, foldl', sort, find, unfoldr)
import Data.Maybe                                          (mapMaybe, fromJust, fromMaybe)
#if MIN_VERSION_base(4,12,0)
import Data.Monoid                                         (Ap(..))
#endif
import Data.Proxy                                          (Proxy(..))
import Data.Ratio                                          ((%), numerator, denominator)
import Data.Semigroup                                      (Semigroup(..), Product(..))
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Vector                                         (Vector, (!))
import GHC.TypeNats.Compat                                 (Nat, SomeNat(..), natVal, someNatVal)
import Numeric.Natural                                     (Natural)

import Math.NumberTheory.ArithmeticFunctions               (totient)
import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.Class                      (KnownNat, Mod, getVal)
import Math.NumberTheory.Moduli.Internal                   (isPrimitiveRoot', discreteLogarithmPP)
import Math.NumberTheory.Moduli.Multiplicative             (MultMod(..), isMultElement)
import Math.NumberTheory.Moduli.Singleton                  (Some(..), cyclicGroupFromFactors)
import Math.NumberTheory.Powers.Modular                    (powMod)
import Math.NumberTheory.Primes                            (Prime(..), UniqueFactorisation, factorise, nextPrime)
import Math.NumberTheory.Utils.FromIntegral                (wordToInt)
import Math.NumberTheory.Utils

-- | A Dirichlet character mod \(n\) is a group homomorphism from \((\mathbb{Z}/n\mathbb{Z})^*\)
-- to \(\mathbb{C}^*\), represented abstractly by `DirichletCharacter`. In particular, they take
-- values at roots of unity and can be evaluated using `eval`.
-- A Dirichlet character can be extended to a completely multiplicative function on \(\mathbb{Z}\)
-- by assigning the value 0 for \(a\) sharing a common factor with \(n\), using `evalGeneral`.
--
-- There are finitely many possible Dirichlet characters for a given modulus, in particular there
-- are \(\phi(n)\) characters modulo \(n\), where \(\phi\) refers to Euler's `totient` function.
-- This gives rise to `Enum` and `Bounded` instances.
newtype DirichletCharacter (n :: Nat) = Generated [DirichletFactor]

-- | The group (Z/nZ)^* decomposes to a product (Z/2^k0 Z)^* x (Z/p1^k1 Z)^* x ... x (Z/pi^ki Z)^*
-- where n = 2^k0 p1^k1 ... pi^ki, and the pj are odd primes, k0 possibly 0. Thus, a group
-- homomorphism from (Z/nZ)^* is characterised by group homomorphisms from each of these factor
-- groups. Furthermore, for odd p, we have (Z/p^k Z)^* isomorphic to Z / p^(k-1)*(p-1) Z, an
-- additive group, where an isomorphism is specified by a choice of primitive root.
-- Similarly, for k >= 2, (Z/2^k Z)^* is isomorphic to Z/2Z * (Z / 2^(k-2) Z) (and for k < 2
-- it is trivial).  (See `lambda` for this isomorphism).
-- Thus, to specify a Dirichlet character, it suffices to specify the value of generators
-- of each of these cyclic groups, when primitive roots are given. This data is given by a
-- DirichletFactor.
-- We have the invariant that the factors must be given in strictly increasing order, and the
-- generator is as given by `generator`, and are each non-trivial. These conditions are verified
-- using `validChar`.
data DirichletFactor = OddPrime { _getPrime :: Prime Natural
                                , _getPower :: Word
                                , _getGenerator :: Natural
                                , _getValue :: RootOfUnity
                                }
                     | TwoPower { _getPower2 :: Int -- this ought to be Word, but many applications
                                                    -- needed to use wordToInt, so Int is cleaner
                                                    -- Required to be >= 2
                                , _getFirstValue :: RootOfUnity
                                , _getSecondValue :: RootOfUnity
                                }
                     | Two

instance Eq (DirichletCharacter n) where
  Generated a == Generated b = a == b

instance Eq DirichletFactor where
  TwoPower _ x1 x2 == TwoPower _ y1 y2 = x1 == y1 && x2 == y2
  OddPrime _ _ _ x == OddPrime _ _ _ y = x == y
  Two              == Two              = True
  _ == _ = False

-- | A representation of <https://en.wikipedia.org/wiki/Root_of_unity roots of unity>: complex
-- numbers \(z\) for which there is \(n\) such that \(z^n=1\).
newtype RootOfUnity =
  RootOfUnity { -- | Every root of unity can be expressed as \(e^{2 \pi i q}\) for some
                -- rational \(q\) satisfying \(0 \leq q < 1\), this function extracts it.
                fromRootOfUnity :: Rational }
  deriving (Eq)

instance Show RootOfUnity where
  show (RootOfUnity q)
    | n == 0    = "1"
    | d == 1    = "-1"
    | n == 1    = "e^(πi/" ++ show d ++ ")"
    | otherwise = "e^(" ++ show n ++ "πi/" ++ show d ++ ")"
    where n = numerator (2*q)
          d = denominator (2*q)

-- | Given a rational \(q\), produce the root of unity \(e^{2 \pi i q}\).
toRootOfUnity :: Rational -> RootOfUnity
toRootOfUnity q = RootOfUnity ((n `rem` d) % d)
  where n = numerator q
        d = denominator q
        -- effectively q `mod` 1
  -- This smart constructor ensures that the rational is always in the range 0 <= q < 1.

-- | This Semigroup is in fact a group, so @'stimes'@ can be called with a negative first argument.
instance Semigroup RootOfUnity where
  RootOfUnity q1 <> RootOfUnity q2 = toRootOfUnity (q1 + q2)
  stimes k (RootOfUnity q) = toRootOfUnity (q * fromIntegral k)

instance Monoid RootOfUnity where
  mappend = (<>)
  mempty = RootOfUnity 0

-- | Convert a root of unity into an inexact complex number. Due to floating point inaccuracies,
-- it is recommended to avoid use of this until the end of a calculation. Alternatively, with
-- the [cyclotomic](http://hackage.haskell.org/package/cyclotomic-0.5.1) package, one can use
-- @[polarRat](https://hackage.haskell.org/package/cyclotomic-0.5.1/docs/Data-Complex-Cyclotomic.html#v:polarRat)
-- 1 . @'fromRootOfUnity' to convert to a cyclotomic number.
toComplex :: Floating a => RootOfUnity -> Complex a
toComplex (RootOfUnity t)
  | t == 1/2 = (-1) :+ 0
  | t == 1/4 = 0 :+ 1
  | t == 3/4 = 0 :+ (-1)
  | otherwise = cis . (2*pi*) . fromRational $ t

-- | For primes, define the canonical primitive root as the smallest such. For prime powers \(p^k\),
-- either the smallest primitive root \(g\) mod \(p\) works, or \(g+p\) works.
generator :: (Integral a, UniqueFactorisation a) => Prime a -> Word -> a
generator p k
  | k == 1 = modP
  | otherwise = if powMod modP (p'-1) (p'*p') == 1 then modP + p' else modP
  where p' = unPrime p
        modP = case cyclicGroupFromFactors [(p,k)] of
                 Just (Some cg) -> head $ filter (isPrimitiveRoot' cg) [2..p'-1]
                 _ -> error "illegal"

-- | Implement the function \(\lambda\) from page 5 of
-- https://www2.eecs.berkeley.edu/Pubs/TechRpts/1984/CSD-84-186.pdf
lambda :: Integer -> Int -> Integer
lambda x e = ((powMod x (2*modulus) largeMod - 1) `shiftR` (e+1)) .&. (modulus - 1)
  where modulus = bit (e-2)
        largeMod = bit (2*e - 1)

-- | For elements of the multiplicative group \((\mathbb{Z}/n\mathbb{Z})^*\), a Dirichlet
-- character evaluates to a root of unity.
eval :: DirichletCharacter n -> MultMod n -> RootOfUnity
eval (Generated ds) m = foldMap (evalFactor m') ds
  where m' = getVal $ multElement m

-- | Evaluate each factor of the Dirichlet character.
evalFactor :: Integer -> DirichletFactor -> RootOfUnity
evalFactor m =
  \case
    OddPrime (toInteger . unPrime -> p) k (toInteger -> a) b ->
      discreteLogarithmPP p k a (m `rem` p^k) `stimes` b
    TwoPower k s b -> (if testBit m 1 then s else mempty)
                   <> lambda (thingy k m) k `stimes` b
    Two -> mempty

thingy :: (Bits p, Num p) => Int -> p -> p
thingy k m = if testBit m 1
                then bit k - m'
                else m'
  where m' = m .&. (bit k - 1)

-- | A character can evaluate to a root of unity or zero: represented by @Nothing@.
evalGeneral :: KnownNat n => DirichletCharacter n -> Mod n -> OrZero RootOfUnity
evalGeneral chi t = case isMultElement t of
                      Nothing -> Zero
                      Just x -> NonZero $ eval chi x

-- | Give the principal character for this modulus: a principal character mod \(n\) is 1 for
-- \(a\) coprime to \(n\), and 0 otherwise.
principalChar :: KnownNat n => DirichletCharacter n
principalChar = minBound

mulChars :: DirichletCharacter n -> DirichletCharacter n -> DirichletCharacter n
mulChars (Generated x) (Generated y) = Generated (zipWith combine x y)
  where combine :: DirichletFactor -> DirichletFactor -> DirichletFactor
        combine Two Two = Two
        combine (OddPrime p k g n) (OddPrime _ _ _ m) =
          OddPrime p k g (n <> m)
        combine (TwoPower k a n) (TwoPower _ b m) =
          TwoPower k (a <> b) (n <> m)
        combine _ _ = error "internal error: malformed DirichletCharacter"

-- | This Semigroup is in fact a group, so @stimes@ can be called with a negative first argument.
instance Semigroup (DirichletCharacter n) where
  (<>) = mulChars
  stimes = stimesChar

instance KnownNat n => Monoid (DirichletCharacter n) where
  mempty = principalChar
  mappend = (<>)

stimesChar :: Integral a => a -> DirichletCharacter n -> DirichletCharacter n
stimesChar s (Generated xs) = Generated (map mult xs)
  where mult :: DirichletFactor -> DirichletFactor
        mult (OddPrime p k g n) = OddPrime p k g (s `stimes` n)
        mult (TwoPower k a b) = TwoPower k (s `stimes` a) (s `stimes` b)
        mult Two = Two

-- | We define `succ` and `pred` with more efficient implementations than
-- @`toEnum` . (+1) . `fromEnum`@.
instance KnownNat n => Enum (DirichletCharacter n) where
  toEnum = indexToChar . fromIntegral
  fromEnum = fromIntegral . characterNumber
  succ x = makeChar x (characterNumber x + 1)
  pred x = makeChar x (characterNumber x - 1)

  enumFromTo x y       = bulkMakeChars x [fromEnum x..fromEnum y]
  enumFrom x           = bulkMakeChars x [fromEnum x..]
  enumFromThenTo x y z = bulkMakeChars x [fromEnum x, fromEnum y..fromEnum z]
  enumFromThen x y     = bulkMakeChars x [fromEnum x, fromEnum y..]

instance KnownNat n => Bounded (DirichletCharacter n) where
  minBound = indexToChar 0
  maxBound = indexToChar (totient n - 1)
    where n = natVal (Proxy :: Proxy n)

-- | We have a (non-canonical) enumeration of dirichlet characters.
characterNumber :: DirichletCharacter n -> Integer
characterNumber (Generated y) = foldl' go 0 y
  where go x (OddPrime p k _ a) = x * m + numerator (fromRootOfUnity a * fromIntegral m)
          where p' = fromIntegral (unPrime p)
                m = p'^(k-1)*(p'-1)
        go x (TwoPower k a b)   = x' * 2 + numerator (fromRootOfUnity a * 2)
          where m = bit (k-2) :: Integer
                x' = x `shiftL` (k-2) + numerator (fromRootOfUnity b * fromIntegral m)
        go x Two = x

-- | Give the dirichlet character from its number.
-- Inverse of `characterNumber`.
indexToChar :: forall n. KnownNat n => Natural -> DirichletCharacter n
indexToChar = runIdentity . indicesToChars . Identity

-- | Give a collection of dirichlet characters from their numbers. This may be more efficient than
-- `indexToChar` for multiple characters, as it prevents some internal recalculations.
indicesToChars :: forall n f. (KnownNat n, Functor f) => f Natural -> f (DirichletCharacter n)
indicesToChars = fmap (Generated . unroll t . (`mod` m))
  where n = natVal (Proxy :: Proxy n)
        (Product m, t) = mkTemplate n

-- | List all characters for the modulus. This is preferred to using @[minBound..maxBound]@.
allChars :: forall n. KnownNat n => [DirichletCharacter n]
allChars = indicesToChars [0..m-1]
  where m = totient $ natVal (Proxy :: Proxy n)

-- | The same as `indexToChar`, but if we're given a character we can create others more efficiently.
makeChar :: Integral a => DirichletCharacter n -> a -> DirichletCharacter n
makeChar x = runIdentity . bulkMakeChars x . Identity

-- | Use one character to make many more: better than indicesToChars since it avoids recalculating
-- some primitive roots
bulkMakeChars :: (Integral a, Functor f) => DirichletCharacter n -> f a -> f (DirichletCharacter n)
bulkMakeChars x = fmap (Generated . unroll t . (`mod` m) . fromIntegral)
  where (Product m, t) = templateFromCharacter x

-- We assign each natural a unique Template, which can be decorated (eg in `unroll`) to
-- form a DirichletCharacter. A Template effectively holds the information carried around
-- in a DirichletFactor which depends only on the modulus of the character.
data Template = OddTemplate { _getPrime'     :: Prime Natural
                            , _getPower'     :: Word
                            , _getGenerator' :: !Natural
                            , _getModulus'   :: !Natural
                            }
              | TwoPTemplate { _getPower2'    :: Int
                             , _getModulus'   :: !Natural
                             } -- the modulus is derivable from the other values, but calculation
                               -- may be expensive, so we pre-calculate it
                               -- morally getModulus should be a prefactored but seems to be
                               -- pointless here
              | TwoTemplate

templateFromCharacter :: DirichletCharacter n -> (Product Natural, [Template])
templateFromCharacter (Generated t) = traverse go t
  where go (OddPrime p k g _) = (Product m, OddTemplate p k g m)
          where p' = unPrime p
                m = p'^(k-1)*(p'-1)
        go (TwoPower k _ _) = (Product (2*m), TwoPTemplate k m)
          where m = bit (k-2)
        go Two = (Product 1, TwoTemplate)

mkTemplate :: Natural -> (Product Natural, [Template])
mkTemplate = go . sort . factorise
  where go :: [(Prime Natural, Word)] -> (Product Natural, [Template])
        go ((unPrime -> 2, 1): xs) = (Product 1, [TwoTemplate]) <> traverse odds xs
        go ((unPrime -> 2, wordToInt -> k): xs) = (Product (2*m), [TwoPTemplate k m]) <> traverse odds xs
          where m = bit (k-2)
        go xs = traverse odds xs
        odds :: (Prime Natural, Word) -> (Product Natural, Template)
        odds (p, k) = (Product m, OddTemplate p k (generator p k) m)
          where p' = unPrime p
                m = p'^(k-1)*(p'-1)

-- the validity of the producted dirichletfactor list here requires the template to be valid
unroll :: [Template] -> Natural -> [DirichletFactor]
unroll t m = snd (mapAccumL func m t)
  where func :: Natural -> Template -> (Natural, DirichletFactor)
        func a (OddTemplate p k g n) = (a1, OddPrime p k g (toRootOfUnity $ (toInteger a2) % (toInteger n)))
          where (a1,a2) = quotRem a n
        func a (TwoPTemplate k n) = (b1, TwoPower k (toRootOfUnity $ (toInteger a2) % 2) (toRootOfUnity $ (toInteger b2) % (toInteger n)))
          where (a1,a2) = quotRem a 2
                (b1,b2) = quotRem a1 n
        func a TwoTemplate = (a, Two)

-- | Test if a given Dirichlet character is prinicpal for its modulus: a principal character mod
-- \(n\) is 1 for \(a\) coprime to \(n\), and 0 otherwise.
isPrincipal :: DirichletCharacter n -> Bool
isPrincipal chi = characterNumber chi == 0

-- | Induce a Dirichlet character to a higher modulus. If \(d \mid n\), then \(a \bmod{n}\) can be
-- reduced to \(a \bmod{d}\). Thus, the multiplicative function on \(\mathbb{Z}/d\mathbb{Z}\)
-- induces a multiplicative function on \(\mathbb{Z}/n\mathbb{Z}\).
--
-- >>> :set -XTypeApplications
-- >>> chi = indexToChar 5 :: DirichletCharacter 45
-- >>> chi2 = induced @135 chi
-- >>> :t chi2
-- Maybe (DirichletCharacter 135)
induced :: forall n d. (KnownNat d, KnownNat n) => DirichletCharacter d -> Maybe (DirichletCharacter n)
induced (Generated start) = if n `rem` d == 0
                            then Just (Generated (combine (snd $ mkTemplate n) start))
                            else Nothing
  where n = natVal (Proxy :: Proxy n)
        d = natVal (Proxy :: Proxy d)
        combine :: [Template] -> [DirichletFactor] -> [DirichletFactor]
        combine [] _ = []
        combine ts [] = map newFactor ts
        combine (t:xs) (y:ys) = case (t,y) of
                                  (TwoTemplate, Two) -> Two: combine xs ys
                                  (TwoTemplate, _) -> Two: combine xs (y:ys)
                                  (TwoPTemplate k _, Two) -> TwoPower k mempty mempty: combine xs ys
                                  (TwoPTemplate k _, TwoPower _ a b) -> TwoPower k a b: combine xs ys
                                  (TwoPTemplate k _, _) -> TwoPower k mempty mempty: combine xs (y:ys)
                                  (OddTemplate p k _ _, OddPrime q _ g a) | p == q -> OddPrime p k g a: combine xs ys
                                  (OddTemplate p k g _, OddPrime q _ _ _) | p < q -> OddPrime p k g mempty: combine xs (y:ys)
                                  _ -> error "internal error in induced: please report this as a bug"
        newFactor :: Template -> DirichletFactor
        newFactor TwoTemplate = Two
        newFactor (TwoPTemplate k _) = TwoPower k mempty mempty
        newFactor (OddTemplate p k g _) = OddPrime p k g mempty
        -- rest (p,k) = OddPrime p k (generator p k) mempty

-- | The <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol> gives a real Dirichlet
-- character for odd moduli.
jacobiCharacter :: forall n. KnownNat n => Maybe (RealCharacter n)
jacobiCharacter = if odd n
                     then Just $ RealChar $ Generated $ map go $ snd $ mkTemplate n
                     else Nothing
  where n = natVal (Proxy :: Proxy n)
        go :: Template -> DirichletFactor
        go (OddTemplate p k g _) = OddPrime p k g $ toRootOfUnity ((toInteger k) % 2)
          -- jacobi symbol of a primitive root mod p over p is always -1
        go _ = error "internal error in jacobiCharacter: please report this as a bug"
          -- every factor of n should be odd

-- | A Dirichlet character is real if it is real-valued.
newtype RealCharacter n = RealChar { -- | Extract the character itself from a `RealCharacter`.
                                     getRealChar :: DirichletCharacter n
                                   }
                                   deriving Eq

-- | Test if a given `DirichletCharacter` is real, and if so give a `RealCharacter`.
isRealCharacter :: DirichletCharacter n -> Maybe (RealCharacter n)
isRealCharacter t@(Generated xs) = if all real xs then Just (RealChar t) else Nothing
  where real :: DirichletFactor -> Bool
        real (OddPrime _ _ _ a) = a <> a == mempty
        real (TwoPower _ _ b) = b <> b == mempty
        real Two = True

-- TODO: it should be possible to calculate this without eval/evalGeneral
-- and thus avoid using discrete log calculations: consider the order of m
-- inside each of the factor groups?
-- | Evaluate a real Dirichlet character, which can only take values \(-1,0,1\).
toRealFunction :: KnownNat n => RealCharacter n -> Mod n -> Int
toRealFunction (RealChar chi) m = case evalGeneral chi m of
                                    Zero -> 0
                                    NonZero t | t == mempty -> 1
                                    NonZero t | t == RootOfUnity (1 % 2) -> -1
                                    _ -> error "internal error in toRealFunction: please report this as a bug"
                                      -- A real character should not be able to evaluate to
                                      -- anything other than {-1,0,1}, so should not reach this branch

-- | Test if the internal DirichletCharacter structure is valid.
validChar :: forall n. KnownNat n => DirichletCharacter n -> Bool
validChar (Generated xs) = correctDecomposition && all correctPrimitiveRoot xs && all validValued xs
  where correctDecomposition = sort (factorise n) == map getPP xs
        getPP (TwoPower k _ _) = (two, fromIntegral k)
        getPP (OddPrime p k _ _) = (p, k)
        getPP Two = (two,1)
        correctPrimitiveRoot (OddPrime p k g _) = g == generator p k
        correctPrimitiveRoot _ = True
        validValued (TwoPower k a b) = a <> a == mempty && (bit (k-2) :: Integer) `stimes` b == mempty
        validValued (OddPrime (unPrime -> p) k _ a) = (p^(k-1)*(p-1)) `stimes` a == mempty
        validValued Two = True
        n = natVal (Proxy :: Proxy n)
        two = nextPrime 2

-- | Get the order of the Dirichlet Character.
orderChar :: DirichletCharacter n -> Integer
orderChar (Generated xs) = foldl' lcm 1 $ map orderFactor xs
  where orderFactor (TwoPower _ (RootOfUnity a) (RootOfUnity b)) = denominator a `lcm` denominator b
        orderFactor (OddPrime _ _ _ (RootOfUnity a)) = denominator a
        orderFactor Two = 1

-- | Test if a Dirichlet character is <https://en.wikipedia.org/wiki/Dirichlet_character#Primitive_characters_and_conductor primitive>.
isPrimitive :: DirichletCharacter n -> Maybe (PrimitiveCharacter n)
isPrimitive t@(Generated xs) = if all primitive xs then Just (PrimitiveCharacter t) else Nothing
  where primitive :: DirichletFactor -> Bool
        primitive Two = False
        -- for odd p, we're testing if phi(p^(k-1)) `stimes` a is 1, since this means the
        -- character can come from some the smaller modulus p^(k-1)
        primitive (OddPrime _ 1 _ a) = a /= mempty
        primitive (OddPrime (unPrime -> p) k _ a) = (p^(k-2)*(p-1)) `stimes` a /= mempty
        primitive (TwoPower 2 a _) = a /= mempty
        primitive (TwoPower k _ b) = (bit (k-3) :: Integer) `stimes` b /= mempty

-- | A Dirichlet character is primitive if cannot be 'induced' from any character with
-- strictly smaller modulus.
newtype PrimitiveCharacter n = PrimitiveCharacter { -- | Extract the character itself from a `PrimitiveCharacter`.
                                                    getPrimitiveChar :: DirichletCharacter n
                                                    }
                                                    deriving Eq

data WithNat (a :: Nat -> *) where
  WithNat :: KnownNat m => a m -> WithNat a

-- | This function also provides access to the new modulus on type level, with a KnownNat instance
makePrimitive :: DirichletCharacter n -> WithNat PrimitiveCharacter
makePrimitive (Generated xs) =
  case someNatVal (product mods) of
    SomeNat (Proxy :: Proxy m) -> WithNat (PrimitiveCharacter (Generated ys) :: PrimitiveCharacter m)
  where (mods,ys) = unzip (mapMaybe prim xs)
        prim :: DirichletFactor -> Maybe (Natural, DirichletFactor)
        prim Two = Nothing
        prim (OddPrime p' k g a) = case find works options of
                                     Nothing -> error "invalid character"
                                     Just (0,_) -> Nothing
                                     Just (i,_) -> Just (p^i, OddPrime p' i g a)
          where options = (0,1): [(i,p^(i-1)*(p-1)) | i <- [1..k]]
                works (_,phi) = phi `stimes` a == mempty
                p = unPrime p'
        prim (TwoPower k a b) = case find worksb options of
                                  Nothing -> error "invalid character"
                                  Just (2,_) | a == mempty -> Nothing
                                  Just (i,_) -> Just (bit i :: Natural, TwoPower i a b)
          where options = [(i, bit (i-2) :: Natural) | i <- [2..k]]
                worksb (_,phi) = phi `stimes` b == mempty

#if !MIN_VERSION_base(4,12,0)
newtype Ap f a = Ap { getAp :: f a }
  deriving (Eq, Functor, Applicative, Monad)

instance (Applicative f, Semigroup a) => Semigroup (Ap f a) where
  (<>) = liftA2 (<>)

instance (Applicative f, Semigroup a, Monoid a) => Monoid (Ap f a) where
  mempty = pure mempty
  mappend = (<>)
#endif

-- | Similar to Maybe, but with different Semigroup and Monoid instances.
type OrZero a = Ap Maybe a
pattern Zero :: OrZero a
pattern Zero = Ap Nothing

pattern NonZero :: a -> OrZero a
pattern NonZero x = Ap (Just x)
{-# COMPLETE Zero, NonZero #-}

-- | Interpret an `OrZero` as a number, taking the `Zero` case to be 0.
orZeroToNum :: Num a => (b -> a) -> OrZero b -> a
orZeroToNum _ Zero = 0
orZeroToNum f (NonZero x) = f x

-- | In general, evaluating a DirichletCharacter at a point involves solving the discrete logarithm
-- problem, which can be hard: the implementations here are around O(sqrt n).
-- However, evaluating a dirichlet character at every point amounts to solving the discrete
-- logarithm problem at every point also, which can be done together in O(n) time, better than
-- using a complex algorithm at each point separately. Thus, if a large number of evaluations
-- of a dirichlet character are required, `evalAll` will be better than `evalGeneral`, since
-- computations can be shared.
evalAll :: forall n. KnownNat n => DirichletCharacter n -> Vector (OrZero RootOfUnity)
evalAll (Generated xs) = V.generate (fromIntegral n) func
  where n = natVal (Proxy :: Proxy n)
        vectors = map mkVector xs
        func :: Int -> OrZero RootOfUnity
        func m = foldMap go vectors
          where go :: (Int, Vector (OrZero RootOfUnity)) -> OrZero RootOfUnity
                go (modulus,v) = v ! (m `mod` modulus)
        mkVector :: DirichletFactor -> (Int, Vector (OrZero RootOfUnity))
        mkVector Two = (2, V.fromList [Zero, mempty])
        mkVector (OddPrime p k (fromIntegral -> g) a) = (modulus, w)
          where
            p' = unPrime p
            modulus = fromIntegral (p'^k) :: Int
            w = V.create $ do
              v <- MV.replicate modulus Zero
              -- TODO: we're in the ST monad here anyway, could be better to use STRefs to manage
              -- this loop, the current implementation probably doesn't fuse well
              let powers = iterateMaybe go (1,mempty)
                  go (m,x) = if m' > 1
                                then Just (m', x<>a)
                                else Nothing
                    where m' = m*g `mod` modulus
              for_ powers $ \(m,x) -> MV.unsafeWrite v m (NonZero x)
              -- don't bother with bounds check since m was reduced mod p^k
              return v
        -- for powers of two we use lambda directly instead, since the generators of the cyclic
        -- groups aren't obvious; it's possible to get them though:
        -- 5^(lambda(5)^{-1} mod 2^(p-2)) mod 2^p
        mkVector (TwoPower k a b) = (modulus, w)
          where
            modulus = bit k
            w = V.generate modulus f
            f m
              | even m = Zero
              | otherwise = NonZero ((if testBit m 1 then a else mempty) <> lambda (toInteger m'') k `stimes` b)
              where m'' = thingy k m

-- somewhere between unfoldr and iterate
iterateMaybe :: (a -> Maybe a) -> a -> [a]
iterateMaybe f x = unfoldr (fmap (\t -> (t, f t))) (Just x)

-- | Attempt to construct a character from its table of values.
-- An inverse to `evalAll`, defined only on its image.
fromTable :: forall n. KnownNat n => Vector (OrZero RootOfUnity) -> Maybe (DirichletCharacter n)
fromTable v = if length v == fromIntegral n
                 then Generated <$> traverse makeFactor tmpl >>= check
                 else Nothing
  where n = natVal (Proxy :: Proxy n)
        n' = fromIntegral n :: Integer
        tmpl = snd (mkTemplate n)
        check :: DirichletCharacter n -> Maybe (DirichletCharacter n)
        check chi = if evalAll chi == v then Just chi else Nothing
        makeFactor :: Template -> Maybe DirichletFactor
        makeFactor TwoTemplate = Just Two
        makeFactor (TwoPTemplate k _) = TwoPower k <$> getValue (-1,bit k) <*> getValue (exp4 k, bit k)
        makeFactor (OddTemplate p k g _) = OddPrime p k g <$> getValue (toInteger g, toInteger (unPrime p)^k)
        getValue :: (Integer,Integer) -> Maybe RootOfUnity
        getValue (g,m) = getAp (v ! fromInteger (fromJust (chineseCoprime (g,m) (1,n' `quot` m)) `mod` n'))

exp4terms :: [Rational]
exp4terms = [4^k % product [1..k] | k <- [0..]]

-- For reasons that aren't clear to me, `exp4` gives the inverse of 1 under lambda, so it gives the generator
-- This is the same as https://oeis.org/A320814
-- In particular, lambda (exp4 n) n == 1 (for n >= 3)
-- I've verified this for 3 <= n <= 2000, so the reasoning in fromTable should be accurate for moduli below 2^2000
exp4 :: Int -> Integer
exp4 n = (`mod` bit n) $ sum $ map (`mod` bit n) $ map (\q -> numerator q * fromMaybe (error "error in exp4") (recipMod (denominator q) (bit n))) $ take n $ exp4terms


{-# LANGUAGE OverloadedStrings, RecordWildCards, NamedFieldPuns #-}
module Numeric.LP where

import           Control.Monad
import           Data.Array.IArray           (Array, (//))
import qualified Data.Array.IArray           as A
import           Data.Maybe
import qualified Data.Text.Lazy              as T
import qualified Data.Text.Lazy.Builder      as TB
import qualified Formatting                  as F
import           Numeric.LinearAlgebra       (I, Matrix, Vector, atIndex, cols,
                                              flatten, ident, minIndex, rows,
                                              subMatrix, subVector, (!), (><),
                                              (?), (|||), (¿))
import qualified Numeric.LinearAlgebra       as LA
import           Numeric.LinearAlgebra.Devel (atM', modifyMatrix, runSTMatrix,
                                              thawMatrix)
import           Text.UnicodeTable

data LP = LP {
    a        :: Matrix Double,
    lhs, rhs :: Vector I}

lp2Table :: Int -> Int -> LP -> Table
lp2Table nPrec nColW LP{..} = mkTable ta cws hs vs
    where rN = rows a
          cN = cols a
          taElem r c | r>0 && c>0 = F.bprint (F.fixed nPrec) $ a `atIndex` (r-1, c-1)
                     | r==0 && c>1 = TB.fromString $ "x"++show (rhs!(c-1))
                     | r==1 && c==0 =  TB.fromString "z"
                     | r>1 && c==0 = TB.fromString $ "x"++show (lhs!(r-1))
                     | otherwise = TB.fromString ""
          ta = A.array ((0,0), (rN, cN)) [((r,c), taElem r c) | r <- [0..rN], c <- [0..cN]]
          cws = A.listArray (0,cN) (replicate (cN+1) nColW)
          hs = konstArray0 (rN+2) NoLine // [(0, SingleLine), (1, DoubleLine), (rN+1, SingleLine)]
          vs = konstArray0 (cN+2) NoLine // [(0, SingleLine), (1, DoubleLine), (cN+1, SingleLine)]

konstArray0 :: Int -> e -> Array Int e
konstArray0 n x = A.listArray (0,n-1) (replicate n x)

instance Show LP where
    show = T.unpack . TB.toLazyText . fromTable . lp2Table 4 15

transfromToEquationalForm :: Matrix Double -> Matrix Double
transfromToEquationalForm a = a ||| (ident (rows a))

choosePivot :: LP -> Maybe (Int, Int)
choosePivot LP {..} =
    let c = subVector 1 (cols a-1) (flatten $ a ? [0])
        b = subVector 1 (cols a-1) (flatten $ a ¿ [0])
        a' = subMatrix (1,1) (rows a-1, cols a-1) a
    in case LA.find (>0) c of
            s:_ ->
                let as = flatten $ a' ¿ [s] -- column s as a vector
                in Just (1 + minIndex (b / as), 1+s)
            _ -> Nothing

-- see <https://en.wikipedia.org/wiki/Simplex_algorithm>
-- Let's solve the following problem:
--      maximize $c^t x$ subject to
--      Ax ≤ b
--      ∀i, xᵢ ≥ 0
--      x,c ∈ ℝⁿ, b ∈ ℝ^p, A ∈ ℝ^pxn

ex1 :: LP
ex1 = LP {a = (3><3)[0,2,-4, 2,-6,1, 8,3,-4],
          lhs = LA.fromList [0,1,4],
          rhs = LA.fromList [0,2,3]
        }

simplexStep :: LP -> Int -> Int -> LP
simplexStep (LP {..}) sr sc =
    let theta = atM' a sr sc -- save the pivot element
        pc = flatten $ a ¿ [sc] -- save the whole pivot column
        a' = runSTMatrix
            (do
                a'm <- thawMatrix a
                -- replace each row -- except the pivot row -- by that linear combination
                -- of itself and the pivot row which makes its pivot column entry zero.
                forM_ [r | r <- [0..(rows a - 1)], r /= sr] $ \r -> do
                    let m = (-1) * (a `atIndex` (r,sc)) / theta
                    forM_ [0..(cols a - 1)] $ \i -> do
                        modifyMatrix a'm r i (\x -> x + m * (a `atIndex` (sr,i)))
                -- divide the pivot row by the negative of the pivot
                forM_ [0..(cols a - 1)] $ \i -> do
                        modifyMatrix a'm sr i (\x -> x/(-theta))
                -- replace the pivot by the negative of its saved reciprocal
                modifyMatrix a'm sr sc (\_ -> 1/theta)
                -- replace the other elements of the pivot column by its saved values divided
                -- by the saved pivot element
                forM_ [r | r <- [0..(rows a - 1)], r /= sr] $ \j -> do
                        modifyMatrix a'm j sc (\_ -> pc!j/theta)
                return a'm)
    in LP a' lhs rhs

main :: IO ()
main = do
    putStrLn (show ex1)
    let (LP {..}) = ex1
    let (sr, sc) = fromJust $ choosePivot ex1
    putStrLn ("Pivot row/column: "++(show sr)++", "++(show sc))
    let theta = atM' a sr sc
    putStrLn ("Pivot element: "++(show theta))


{- |
Description :  DNA Markov models (substitution models or mutation models)
Copyright   :  (c) Dominik Schrempf 2017
License     :  GPLv3

Maintainer  :  dominik.schrempf@gmail.com
Stability   :  unstable
Portability :  non-portable (not tested)

This module provides rate matrices for the most common DNA substitution models
(they will be added according to need). They are also used as mutation models by
the boundary mutation model.

* Changelog

-}

module DNAModel
  ( DNAModelSpec(..)
  , DNAModel(..)
  , Nuc
  , getDNAModelInfoStr
  , rateMatrix
  , dnaModelSpecGetStateFreqVec
  , dnaModelScale
  ) where

import           Data.List             (intercalate)
import           Numeric.LinearAlgebra
import           RateMatrix

-- | The nucleotides.
data Nuc = A | C | G | T deriving (Eq, Show, Read, Ord, Bounded, Enum)

nNucs :: Int
nNucs = 1 + fromEnum (maxBound :: Nuc)

-- | The DNA model specifications. E.g., for the HKY model, we only have kappa.
data DNAModelSpec = JC |
                    -- | HKY model with transition to transversion ratio kappa.
                    HKY Double StationaryDist |
                    -- | GTR model with five params and state frequency vector.
                    GTR Double Double Double Double Double StationaryDist

instance Show DNAModelSpec where
  show JC        = "JC"
  show (HKY k f) = "HKY[" ++ show k ++ "]" ++ show f
  show (GTR a b c d e f) = "GTR[" ++ paramsStr ++ "]" ++ show f
    where paramsStr = intercalate "," (map show [a,b,c,d,e])

-- | Extract stationary frequency vector of a DNA model specification 'DNAModelSpec'.
dnaModelSpecGetStateFreqVec :: DNAModelSpec -> StationaryDist
dnaModelSpecGetStateFreqVec JC                = vector $ replicate nNucs 0.25
dnaModelSpecGetStateFreqVec (HKY _ f)         = f
dnaModelSpecGetStateFreqVec (GTR _ _ _ _ _ f) = f

-- | A rate matrix of a DNA models is called DNAModel.
data DNAModel = DNAModel { dnaRateMatrix :: RateMatrix
                         , dnaModelSpec  :: DNAModelSpec }

-- | The matrix of exchangeabilities; diagonal entries are not set.
exchangeabilityMatrix :: DNAModelSpec -> Matrix R
exchangeabilityMatrix JC = (4><4)
  [ 0.0, 1.0, 1.0, 1.0
  , 1.0, 0.0, 1.0, 1.0
  , 1.0, 1.0, 0.0, 1.0
  , 1.0, 1.0, 1.0, 0.0 ]
exchangeabilityMatrix (HKY k _) = (4><4)
  [ 0.0, 1.0,   k, 1.0
  , 1.0, 0.0, 1.0,   k
  ,   k, 1.0, 0.0, 1.0
  , 1.0,   k, 1.0, 0.0 ]
exchangeabilityMatrix (GTR a b c d e _) = (4><4)
  [ 0.0,   a,   b,   c
  ,   a, 0.0,   d,   e
  ,   b,   d, 0.0, 1.0
  ,   c,   e, 1.0, 0.0 ]
-- exchangeabilityMatrix _     = error "Model not yet supported."

-- | DNA model rate matrix normalized so that one mutation happens per unit time.
rateMatrix :: DNAModelSpec -> DNAModel
rateMatrix s = DNAModel rm s
  where
    f    = dnaModelSpecGetStateFreqVec s
    exch = exchangeabilityMatrix s
    rm   = normalizeRates f $ setDiagonal $ exch <> diag f
-- rateMatrix _ _       = error "Model not yet supported."

-- | Report a specific model.
getDNAModelInfoStr :: DNAModel -> String
getDNAModelInfoStr (DNAModel m s) =
  case s of
    JC        -> unlines
      [ "JC model with rate matrix"
      , show m ]
    (HKY k f) -> unlines
      [ "HKY model"
      , reportRateMatrix
      , reportExchangeabilityMatrix
      , reportStateFreqVec f
      , "And a kappa value of: " ++ show k ]
    (GTR a b c d e f) -> unlines
      [ "GTR model"
      , reportRateMatrix
      , reportExchangeabilityMatrix
      , reportStateFreqVec f
      , "And parameters: " ++ intercalate ", " (map show [a,b,c,d,e]) ]
      -- _ -> error "Model not yet supported."
  where reportRateMatrix   = "Rate matrix " ++ show m
        reportStateFreqVec f = "This corresponds to state frequencies (A, C, G, T): " ++ show f
        reportExchangeabilityMatrix = "Exchangeability matrix (diagonal set to 0.0) " ++
                                      show (exchangeabilityMatrix s)

-- | Scale a DNA model.
dnaModelScale :: Double -> DNAModel -> DNAModel
dnaModelScale s (DNAModel rm spec) = DNAModel (scale s rm) scaledSpec
  where scaledSpec = case spec of
                       JC              -> JC
                       HKY k f         -> HKY (s*k) f
                       GTR a b c d e f -> GTR (s*a) (s*b) (s*c) (s*d) (s*e) f


module Optimization where

import Numeric.LinearAlgebra

class Optimize a where
  --optimizeInit :: a -> Double -> Double -> (Int, Int)  -> a
  paramUpdate :: a -> a

data SGD = SGD {
  sgdLearningRate :: Double,
  sgdDecay :: Double,
  sgdMomentum :: Double,
  sgdParams :: [Matrix R],
  sgdGparams :: [Matrix R],
  sgdDeltaPre :: [Matrix R],
  sgdBatchSize :: Double
                 }

instance Optimize SGD where
  paramUpdate sgd =
    let learningRate = sgdLearningRate sgd in
    let decay = sgdDecay sgd in
    let momentum = sgdMomentum sgd in
    let params = sgdParams sgd in
    let gparams = sgdGparams sgd in
    let deltaPre = sgdDeltaPre sgd in
    let batchSize = sgdBatchSize sgd in
    let _sgd (px:pxs) (gx: gxs) (dx:dxs) paUpdate deUpdate gaZeros  =
          let r = rows gx in
          let c = cols gx in
          let lrMat = (r >< c) [learningRate..]:: (Matrix R) in
          let batchSizeMat = (r >< c) [batchSize..] :: (Matrix R) in
          let momMat = (r >< c) [momentum..] :: (Matrix R) in
          let gUp = gx * lrMat / batchSizeMat in
          let ugd = momMat * dx - gx in
          let pu = px + ugd in _sgd pxs gxs dxs (paUpdate++[pu]) (deUpdate++[ugd]) (gaZeros++[(r >< c) [0..]])
        _sgd [] [] [] paUpdate deUpdate gaZeros = SGD {
          sgdLearningRate=learningRate,
          sgdDecay=decay,
          sgdMomentum=momentum,
          sgdParams=paUpdate,
          sgdGparams=gaZeros,
          sgdDeltaPre=deUpdate,
          sgdBatchSize=0
          } in
    _sgd params gparams deltaPre [] [] []

data Adadelta = Adadelta {
  adaDecay :: Double,
  adaEpsilon :: Double,
  adaBatchSize :: Double,
  adaParams :: [Matrix R],
  adaGparams :: [Matrix R],
  adaAccGrad :: [Matrix R],
  adaAccDelta :: [Matrix R]
                         }

instance Optimize Adadelta where
  paramUpdate ada =
    let decay = adaDecay ada in
    let epsilon = adaEpsilon ada in
    let batchSize = adaBatchSize ada in
    let params = adaParams ada in
    let gparams = adaGparams ada in
    let accGrad = adaAccGrad ada in
    let accDelta = adaAccDelta ada in
    let _ada (px:pxs) (gx:gxs) (accgx:accgxs) (accdx:accdxs) paUpdate gradUpdate deltaUpdate gaZeros =
          let r = rows px in
          let c = cols px in
          let decayMat = (r >< c) [decay ..] :: (Matrix R) in
          let epsilonMat = (r >< c) [epsilon ..] :: (Matrix R) in
          let batchSizeMat = (r >< c) [batchSize ..] :: (Matrix R) in
          let oneMat = (r >< c) [1.0 ..]:: (Matrix R) in
          let gUp = accgx / batchSizeMat in
          let gradUp = decayMat * accgx + (oneMat - decayMat) * gUp * gUp in
          let ugd = (cmap (\x -> -(sqrt x)) (accdx + epsilonMat)) /  (cmap (\x -> sqrt x) (accgx + epsilonMat)) * gUp in
          let deltaUp = decayMat * accdx + (oneMat - decayMat) * ugd * ugd in
          let pu = px + ugd in
          _ada pxs gxs accgxs accdxs (pu:paUpdate) (gradUp:gradUpdate) (deltaUp:deltaUpdate) (((r >< c)[0.0 ..]):gaZeros)
        _ada [] [] [] [] paUpdate gradUpdate deltaUpdate gaZeros = Adadelta {
          adaDecay=decay,
          adaEpsilon=epsilon,
          adaBatchSize=batchSize,
          adaParams=(reverse paUpdate),
          adaGparams=(reverse gaZeros),
          adaAccGrad=(reverse gradUpdate),
          adaAccDelta=(reverse deltaUpdate)
                                                                        }
    in _ada params gparams accGrad accDelta [] [] [] []


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Normal
-- Copyright : (c) 2009 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The normal distribution.  This is a continuous probability
-- distribution that describes data that cluster around a mean.

module Statistics.Distribution.Normal
    (
      NormalDistribution
    -- * Constructors
    -- , normalDistr
    --, normalDistrE
    , standard
    ) where

import Data.Data             (Data, Typeable)
import GHC.Generics          (Generic)
import Numeric.MathFunctions.Constants (m_sqrt_2, m_sqrt_2_pi)
import Numeric.SpecFunctions (erfc, invErfc)

import qualified Statistics.Distribution as D
import Statistics.Internal

-- | The normal distribution.
data NormalDistribution = ND {
      mean       :: {-# UNPACK #-} !Double
    , stdDev     :: {-# UNPACK #-} !Double
    , ndPdfDenom :: {-# UNPACK #-} !Double
    , ndCdfDenom :: {-# UNPACK #-} !Double
    } deriving (Eq, Typeable, Data, Generic)

instance Show NormalDistribution where
  showsPrec i (ND m s _ _) = defaultShow2 "normalDistr" m s i
instance Read NormalDistribution where
  readPrec = defaultReadPrecM2 "normalDistr" normalDistrE

instance D.Distribution NormalDistribution where
    cumulative      = cumulative
    complCumulative = complCumulative

instance D.ContDistr NormalDistribution where
    logDensity    = logDensity
    quantile      = quantile
    complQuantile = complQuantile

-- | Standard normal distribution with mean equal to 0 and variance equal to 1
standard :: NormalDistribution
standard = ND { mean       = 0.0
              , stdDev     = 1.0
              , ndPdfDenom = log m_sqrt_2_pi
              , ndCdfDenom = m_sqrt_2
              }

-- | Create normal distribution from parameters.
--
-- IMPORTANT: prior to 0.10 release second parameter was variance not
-- standard deviation.
normalDistrE :: Double            -- ^ Mean of distribution
             -> Double            -- ^ Standard deviation of distribution
             -> Maybe NormalDistribution
normalDistrE m sd
  | sd > 0    = Just ND { mean       = m
                        , stdDev     = sd
                        , ndPdfDenom = log $ m_sqrt_2_pi * sd
                        , ndCdfDenom = m_sqrt_2 * sd
                        }
  | otherwise = Nothing

logDensity :: NormalDistribution -> Double -> Double
logDensity d x = (-xm * xm / (2 * sd * sd)) - ndPdfDenom d
    where xm = x - mean d
          sd = stdDev d

cumulative :: NormalDistribution -> Double -> Double
cumulative d x = erfc ((mean d - x) / ndCdfDenom d) / 2

complCumulative :: NormalDistribution -> Double -> Double
complCumulative d x = erfc ((x - mean d) / ndCdfDenom d) / 2

quantile :: NormalDistribution -> Double -> Double
quantile d p
  | p == 0         = -inf
  | p == 1         = inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.quantile: p must be in [0,1] range. Got: "++show p
  where x          = - invErfc (2 * p)
        inf        = 1/0

complQuantile :: NormalDistribution -> Double -> Double
complQuantile d p
  | p == 0         = inf
  | p == 1         = -inf
  | p == 0.5       = mean d
  | p > 0 && p < 1 = x * ndCdfDenom d + mean d
  | otherwise      =
    error $ "Statistics.Distribution.Normal.complQuantile: p must be in [0,1] range. Got: "++show p
  where x          = invErfc (2 * p)
        inf        = 1/0


module Gradient
    (numericalDiff
    ) where

import Numeric.LinearAlgebra

numericalDiff :: Fractional a => (a -> a) -> a -> a
numericalDiff f x = (f (x+h) - f (x-h)) / (2*h)
    where h = 1e-4


import Data.Complex

main = do
  let a = 1.0 :+ 2.0    -- complex number 1+2i
  let b = 4             -- complex number 4+0i
  -- 'b' is inferred to be complex because it's used in
  -- arithmetic with 'a' below.
  putStrLn $ "Add:      " ++ show (a + b)
  putStrLn $ "Subtract: " ++ show (a - b)
  putStrLn $ "Multiply: " ++ show (a * b)
  putStrLn $ "Divide:   " ++ show (a / b)
  putStrLn $ "Negate:   " ++ show (-a)
  putStrLn $ "Inverse:  " ++ show (recip a)
  putStrLn $ "Conjugate:" ++ show (conjugate a)


{-# OPTIONS_GHC -Wall #-}

-- | Ring classes. A distinguishment is made between Rings and Commutative Rings.
module Plankton.Ring
  ( Semiring
  , Ring
  , CRing
  ) where

import Data.Complex (Complex(..))
import Plankton.Additive
import Plankton.Distribution
import Plankton.Multiplicative
import Protolude (Bool(..), Double, Float, Int, Integer)

-- | Semiring
class (MultiplicativeAssociative a, MultiplicativeUnital a, Distribution a) =>
      Semiring a

instance Semiring Double

instance Semiring Float

instance Semiring Int

instance Semiring Integer

instance Semiring Bool

instance (AdditiveGroup a, Semiring a) => Semiring (Complex a)

-- | Ring
-- a summary of the laws inherited from the ring super-classes
--
-- > zero + a == a
-- > a + zero == a
-- > (a + b) + c == a + (b + c)
-- > a + b == b + a
-- > a - a = zero
-- > negate a = zero - a
-- > negate a + a = zero
-- > a + negate a = zero
-- > one `times` a == a
-- > a `times` one == a
-- > (a `times` b) `times` c == a `times` (b `times` c)
-- > a `times` (b + c) == a `times` b + a `times` c
-- > (a + b) `times` c == a `times` c + b `times` c
-- > a `times` zero == zero
-- > zero `times` a == zero
class ( AdditiveGroup a
      , MultiplicativeAssociative a
      , MultiplicativeUnital a
      , Distribution a
      ) =>
      Ring a

instance Ring Double

instance Ring Float

instance Ring Int

instance Ring Integer

instance (Ring a) => Ring (Complex a)

-- | CRing is a Ring with Multiplicative Commutation.  It arises often due to '*' being defined as a multiplicative commutative operation.
class (Multiplicative a, Ring a) =>
      CRing a

instance CRing Double

instance CRing Float

instance CRing Int

instance CRing Integer

instance (CRing a) => CRing (Complex a)


{-# LANGUAGE RecordWildCards #-}
{-
The shape of a triangle can be represented by two explicit points (ie vectors)
and a third implicit point at (0,0). A triangle defined by three points can be
converted by subtracting one of the points from the others:
  shape (p1,p2,p3) = (p2-p1, p3-p1)
This keeps the shape intact but drops the position of the triangle.

If P and Q are two such triangles, there exists a matrix A that transforms
P into Q: AP = Q
The 'A' matrix can be calculated by multiplying Q and the inverse of P:
  A = QP_inv

So, if we have a set of triangles, we can drop their positional information
and compute their shape transformation matrices. But can we do it in reverse?
Can we start with a set of transformation matrices (which don't give us the
positions of the triangles) and compute a set of points that represent triangles
with the right transformation?
In other words, if we have a set of triangles (with absolute position) and
a transformation matrix for each triangle, can we compute new absolute positions
such that AP=Q?

Solving this is not straightforward as moving a single point will change the
transformation of multiple triangles. Fortunately, though, this problem can be
expression as linear equations and solved quickly with LAPACK.

To express the problem as a linear system, we need 4 equations for each
triangle and 2 unknowns for each point.

The system has this shape: Mx=B
We're trying to find 'x' which are the positions of the new points.
'M' is a matrix that transforms points into transformation matrices.
'B' is the target transformation matrices.

With a bit of rewriting, we can reduce 'A = QP_inv' to 4 linear equations
with 6 unknowns. The 6 unknowns are the x and y positions for the three corners
in a triangle.
This is just another way of calculating the transformation matrix 'A' and we
can verify it by looking at the output of:
  * computeA src dst
  * applyCoeff (coeffOfB src) dst
The output should be the same because they're both computing 'A'.

'applyCoeff' simply multiplies the coefficients by a 6x1 matrix containing
the x and y positions for 'dst'.

So, this means we have roughly the right shape: Mx=B. M is the coefficients
we found for 'src', 'x' is dst, and 'B' is the transformation matrix that
turns 'src' into 'dst'.
We can get the coefficients for each triangle and place them together in a
single large matrix. Then we can put the desired transformation matrices in 'B'
and ask a linear solver to find 'x' which satisfies the equations.

Phew, step one is done. We now have the machinery for turning arbitrary
transformation matrices into proper, connected triangles. Next step is figuring
out which rotational matrices to use. At the two extremes, we have the identity
matrices (representing our starting shape) and the matrices computed from the
target triangles (represetning out target shape). How can these two matrices be
interpolated for a smooth morph?

There are several ways of interpolating matrices. Linear interpolation is
possible. But, it looks prettier if we separate rotation from shear. Then we
can use spherical linear interpolation on the rotation and linear interpolation
on the shear.

-}
module Reanimate.Morph.Rigid where

import           Data.Foldable (toList)
import           Data.Vector                   (Vector)
import qualified Data.Vector                   as V
import           Linear.Quaternion
import           Linear.Vector
import           Linear.V2
import           Linear.V3
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix (GMatrix, Matrix,
                                                toLists, (!), (><))

type P = V2 Double
type Trig = (P,P,P)
type RelTrig = (Int,Int,Int)

data Mesh = Mesh
  { meshPointsA :: Vector P
  , meshPointsB :: Vector P
  , meshOutline :: Vector Int
  -- , meshSteiner :: Vector Int
  , meshTriangles :: Vector RelTrig }

-- applyA (computeA a b) a = b + some_constant_translation
-- A = Q P_inv
computeA :: Trig -> Trig -> Matrix Double
computeA p q = matQ <> Matrix.inv matP
  where
    matP = trigToMatrix p
    matQ = trigToMatrix q

-- A = UDV = U(VV_t)DV = (UV)(V_tDV) = RS
-- R = UV
-- S = V_tDV
computeRS :: Matrix Double -> (Matrix Double, Matrix Double)
computeRS a = (r, s)
  where
    (u,d,vt) = Matrix.svd a
    v = Matrix.tr vt
    r = u <> v
    s = vt <> Matrix.diag d <> v

matrixToQuaternion :: Matrix Double -> Quaternion Double
matrixToQuaternion r =  q
  where
    w = 0.5 * sqrt (1 + r!0!0 + r!1!1 + 1)
    z = 1/(4*w) * (r!0!1 - r!1!0)
    q = Quaternion w (V3 0 0 z)

quaternionToMatrix :: Quaternion Double -> Matrix Double
quaternionToMatrix q = (2><2)
    [ 1 - 2*(qj*qj + qk*qk), 2*(qi*qj+qk*qr)
    , 2*(qi*qj - qk*qr), 1 - 2*(qi*qi + qk*qk) ]
  where
    [qr,qi,qj,qk] = toList q

-- A = R((1-t)I + tS)
computeA_RSt :: Matrix Double -> Matrix Double -> Double -> Matrix Double
computeA_RSt r s t = r_t <> (realToFrac (1-t) * Matrix.ident 2 + realToFrac t * s)
  where
    i = Quaternion 1 0
    q = slerp i (matrixToQuaternion r) t
    r_t = quaternionToMatrix q

applyA :: Matrix Double -> Trig -> Trig
applyA a p =
    case toLists (a <> matP) of
      [ [x1, x2], [y1, y2] ] -> (V2 0 0, V2 x1 y1, V2 x2 y2)
      _ -> error "invalid matrix"
  where
    matP = trigToMatrix p

coeffOfB :: Trig -> Matrix Double
coeffOfB p = (4><6)
    [ -a0-a2,      0, a0,   0, a2,  0
    , -a1-a3,      0, a1,   0, a3,  0
    ,      0, -a0-a2,  0,  a0,  0, a2
    ,      0, -a1-a3,  0,  a1,  0, a3 ]
  where
    [[a0,a1],[a2,a3]] = toLists (Matrix.inv (trigToMatrix p))

applyCoeff :: Matrix Double -> Trig -> Matrix Double
applyCoeff b (V2 x1 y1, V2 x2 y2, V2 x3 y3) = b <> matQ
  where
    matQ = (6><1) [ x1, y1, x2, y2, x3, y3]

trigToMatrix :: Trig -> Matrix Double
trigToMatrix (p1,p2,p3) = matP
  where
    V2 p12x p12y = p2-p1
    V2 p13x p13y = p3-p1
    matP = (2><2)
      [p12x, p13x
      ,p12y, p13y]

data Prep = Prep
  { prepPivot    :: (P, P)
  , prepPointsA  :: Vector P
  , prepPointsB  :: Vector P
  , prepRS       :: Vector (Matrix Double, Matrix Double)
  , prepRSRev    :: Vector (Matrix Double, Matrix Double)
  , prepUToB     :: GMatrix
  }

symmetric :: Bool
symmetric = True

prepare :: Mesh -> Prep
prepare Mesh{..} = Prep
    { prepPivot = (aOrigin, bOrigin)
    , prepPointsA =
        V.map (subtract aOrigin) $
        V.take pivotIdx meshPointsA <> V.drop (pivotIdx+1) meshPointsA
    , prepPointsB =
        V.map (subtract bOrigin) $
        V.take pivotIdx meshPointsB <> V.drop (pivotIdx+1) meshPointsB
    , prepRS = rsList
    , prepRSRev = rsRevList
    , prepUToB = Matrix.mkSparse uToB }
  where
    aOrigin = meshPointsA V.! pivotIdx
    bOrigin = meshPointsB V.! pivotIdx
    pivotIdx = case V.head meshTriangles of
        (a,_,_) -> a
    mkAbs p (a,b,c) = (p V.! a,p V.! b,p V.! c)
    absATrigs = V.map (mkAbs meshPointsA) meshTriangles
    absBTrigs = V.map (mkAbs meshPointsB) meshTriangles
    aList = V.zipWith computeA absATrigs absBTrigs
    rsList = V.map computeRS aList
    revList = V.zipWith computeA absBTrigs absATrigs
    rsRevList = V.map computeRS revList
    n = length meshTriangles
    -- nT = if symmetric then n*2 else n
    -- pivotToB = ((nT*4)><2) $ concat
    --   [ [ fromMaybe 0 (lookup (x,pivotIdx*2) bigM)
    --     , fromMaybe 0 (lookup (x,pivotIdx*2+1) bigM)]
    --   | x <- [0..(nT*4)-1] ]
    uToB =
      [ ((x,y-2),key) | ((x,y),key) <- bigM, y /= pivotIdx*2 && y /= pivotIdx*2+1 ]
    bigM =
      concat (zipWith worker [0..] (V.toList meshTriangles)) ++
      if symmetric
        then concat $ zipWith workerRev [n..] (V.toList meshTriangles)
        else []
    worker i src@(a,b,c) = concat $
      let effs = coeffOfB (mkAbs meshPointsA src) in
      [ [((i*4+h, e*2), effs!h!(j*2))
        ,((i*4+h, e*2+1), effs!h!(j*2+1))]
      | h <- [0..3]
      , (e,j) <- zip [a,b,c] [0..]
      ]
    workerRev i dst@(a,b,c) = concat $
      let effs = coeffOfB (mkAbs meshPointsB dst) in
      [ [((i*4+h, e*2), effs!h!(j*2))
        ,((i*4+h, e*2+1), effs!h!(j*2+1))]
      | h <- [0..3]
      , (e,j) <- zip [a,b,c] [0..]
      ]

interpolate :: Prep -> Double -> Vector P
interpolate Prep{..} t = V.fromList $
    pivot : worker (Matrix.toList solution)
  where
    -- solution = Matrix.cgSolve False prepUToB b
    solution = Matrix.cgx $ last solutions
    solutions = Matrix.cgSolve'
      False
      1e-9
      1e-9
      1000
      prepUToB
      b
      (Matrix.fromList $ concat [ [x,y] | V2 x y <- V.toList target ])
      -- 0
    target = if t < 0.5
      then prepPointsA
      else prepPointsB
    worker (x:y:xs) = V2 x y ^+^ pivot : worker xs
    worker _ = []
    pivot = case prepPivot of
      (src, dst) -> lerp t dst src
    n = V.length prepRS
    b = Matrix.vector $
      [ concat (toLists a)!!j
      | i <- [0..n-1]
      , let (r,s) = prepRS V.! i
      , let a = computeA_RSt r s t
      , j <- [0..3]
      ] ++
      [ concat (toLists a)!!j
      | symmetric
      , i <- [0..n-1]
      , let (r,s) = prepRSRev V.! i
      , let a = computeA_RSt r s (1-t)
      , j <- [0..3]
      ]


{-|
Module      : Numeric.MCMC.Diagnostics
Description : Diagnostics for Markov Chain Monte Carlo Sampling
Copyright   : (c) Adam Conner-Sax 2019
License     : BSD
Maintainer  : adam_conner_sax@yahoo.com
Stability   : experimental

Numeric.MCMC.Diagnostics contains functions to analyze the convergence of MCMC sampling and to estimate the distribution of expectations computed
using a set of MCMC chains.
-}
module Numeric.MCMC.Diagnostics
  ( summarize
  , mpsrf
  , ExpectationSummary(..)
  , mannWhitneyUTest
  )
where

import qualified Control.Foldl                 as FL
import qualified Statistics.Types              as S
import qualified Data.List                     as L
--import qualified Data.Vector.Generic           as V
import qualified Data.Vector.Unboxed           as VU
import qualified Statistics.Test.MannWhitneyU  as S
import qualified Numeric.LinearAlgebra         as LA

data ExpectationSummary a = ExpectationSummary { mean :: a, confidence :: (a, a), rHat :: a } deriving (Show)
-- all sub-chains should have same length here
-- do we need to throw away first half of all the chains??
{- |
Compute the exepctation, confidence interval and
PSRF statistic (from <http://www2.stat.duke.edu/~scs/Courses/Stat376/Papers/ConvergeDiagnostics/BrooksGelman.pdf>)
for a given confidence %, expectation function, and set of MCMC chains. In the version in the paper, each chain is
halved before the statistic is computed but it's not clear why.  Here we assume that is "burn-in" and the user may
choose to do that or not before calling this function.

All sub-chains should have the same length here, though that is not checked.
-}
summarize
  :: (RealFrac b, Ord b)
  => S.CL b -- ^ confidence interval to test
  -> (a -> b) -- ^ function to map a sample to a function under consideration
  -> [[a]] -- ^ list of chains
  -> Maybe (ExpectationSummary b) -- ^ @Maybe@ here in case chains aren't long enough
summarize confidence' expectationF chains = do
  let applied    = fmap (fmap expectationF) chains
      allApplied = concat applied
      chainInterval :: Ord b => [b] -> Maybe (b, b)
      chainInterval c =
        let n  = L.length c
            n' = round (realToFrac n * S.confidenceLevel confidence')
            d  = (n - n') `div` 2
            c' = L.take n' $ L.drop d $ L.sort c
        in if length c' >= 2 then (head &&& last) <$> nonEmpty c' else Nothing
--      intervalLength (a, b) = b - a
  intervals <- traverse chainInterval applied
  (lo, hi)  <- chainInterval allApplied
  let (mLo, mHi) = FL.fold
        ((,) <$> FL.premap fst FL.mean <*> FL.premap snd FL.mean)
        intervals
  let rHat' = (hi - lo) / (mHi - mLo)
      mean' = FL.fold FL.mean allApplied
  return $ ExpectationSummary mean' (lo, hi) rHat'

-- if the chain list is empty or chains have different sizes, etc. we return Nothing
mpsrf :: [a -> Double] -> [[a]] -> Maybe Double
mpsrf expectations chains = do
  let --mkVec :: a -> LA.Vector Double
      mkVec x = LA.fromList $ fmap ($ x) expectations
      m = length chains
  n <- length . head <$> nonEmpty chains --if m > 0 then Just (length $ head chains) else Nothing
  traverse_ (\c -> if length c == n then Just () else Nothing) chains -- check that all chains are length n
  let
    sumVecsF = FL.Fold (+) (LA.vector $ L.replicate (length expectations) 0) id
    sumMatsF = FL.Fold (+) (LA.scale 0 $ LA.ident (length expectations)) id
    meanVecsF =
      (\m' l -> LA.scale (1 / realToFrac l) m') <$> sumVecsF <*> FL.length
    avgVec     = FL.fold meanVecsF
    phiChains  = fmap (fmap mkVec) chains
    avgPhiEach = fmap avgVec phiChains
    diffProd :: LA.Vector Double -> LA.Vector Double -> LA.Matrix Double
    diffProd x y = (x - y) `LA.outer` (x - y)
    w1 = (\(ap, ps) -> (`diffProd` ap) <$> ps) <$> zip avgPhiEach phiChains
    w =
      LA.scale (1 / (realToFrac m * realToFrac (n - 1)))
        $ FL.fold sumMatsF
        $ fmap (FL.fold sumMatsF) w1
    avgPhiAll = avgVec avgPhiEach
    b1        = fmap (`diffProd` avgPhiAll) avgPhiEach
    b = LA.scale (realToFrac n / realToFrac (m - 1)) $ FL.fold sumMatsF b1
    mat       = LA.scale (1 / realToFrac n) $ LA.inv w LA.<> b
    eigen1    = LA.eigenvaluesSH (LA.trustSym mat) `LA.atIndex` 0
    rhat =
      (realToFrac (n - 1) / realToFrac n)
        + (realToFrac (m + 1) / realToFrac m)
        * eigen1
  return rhat

mannWhitneyUTest
  :: (Ord b, VU.Unbox b)
  => S.PValue Double
  -> (a -> b)
  -> [a]
  -> [a]
  -> Maybe S.TestResult
mannWhitneyUTest p f c1 c2 =
  let vOfF = VU.fromList . fmap f
  in  S.mannWhitneyUtest S.SamplesDiffer p (vOfF c1) (vOfF c2)


{-# LANGUAGE ParallelListComp #-}
module Main where

import Numeric.LinearAlgebra.Data (fromList, toList, cmap)
import Numeric.GSL.Fitting.Linear (linear_w)
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo
import Text.Printf

x = fromList [10, 100, 1000, 10000, 100000, 1000000]
y = fromList [4772, 21148, 134194, 1343814, 13389821, 148509325]
sigma = fromList [2112, 8194, 61719, 804762, 4853113, 42071732]

(intercept, slope, _, _, _, chi_sq) = linear_w x (cmap (\s -> 1 / s^2) sigma) y

errbars = liftEC $ do
    let values = [symErrPoint x' y' 0 dy
                  | x' <- toList x
                  | y' <- toList y
                  | dy <- toList sigma
                 ]
    plot_errbars_values .= values

main = toFile def "benchmark-fit.png" $ do
    setColors (map opaque [blue,red,green])
    layout_title .=
        printf "t = %.2f (ns/element) n + %.2f (ns) [chi^2 = %.2f]" slope intercept chi_sq
    layout_x_axis . laxis_generate .= autoScaledLogAxis def
    layout_y_axis . laxis_generate .= autoScaledLogAxis def
    plot errbars
    plot $ points "Benchmark data"
        (zip (toList x) (toList y))
    plot $ line "Best fit" [[
            (t, slope*t + intercept)
            | t <- map (10**) [0.5, 0.6 .. 6.5 ]
        ]]


#!/usr/bin/env stack
-- stack --install-ghc runghc --resolver lts-10.4 --package hmatrix --package lens -- -Wall

{-# LANGUAGE DataKinds       #-}
{-# LANGUAGE DeriveGeneric   #-}
{-# LANGUAGE TemplateHaskell #-}

-- Imports for network logic

import           Control.Lens hiding          ((<.>))
import           Numeric.LinearAlgebra.Static

-- Imports for utility functions

import           GHC.Generics (Generic)
import           Numeric.OneLiner

data Net = N { _weights1 :: L 250 784
             , _bias1    :: R 250
             , _weights2 :: L 10 250
             , _bias2    :: R 10
             }
  deriving (Generic)
makeLenses ''Net

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

softMax
    :: R 10
    -> R 10
softMax x = expx / konst (norm_1 expx)
  where
    expx = exp x

crossEntropy
    :: R 10
    -> R 10
    -> Double
crossEntropy targ res = -(log res <.> targ)

runNet
    :: Net
    -> R 784
    -> R 10
runNet n x = z
  where
    y = logistic $ (n ^. weights1) #> x + (n ^. bias1)
    z = softMax  $ (n ^. weights2) #> y + (n ^. bias2)

netErr
    :: R 784
    -> R 10
    -> Net
    -> Double
netErr x targ n = crossEntropy targ (runNet n x)

main :: IO ()
main = return ()

instance Num Net where
    (+)         = gPlus
    (-)         = gMinus
    (*)         = gTimes
    negate      = gNegate
    abs         = gAbs
    signum      = gSignum
    fromInteger = gFromInteger

instance Fractional Net where
    (/)          = gDivide
    recip        = gRecip
    fromRational = gFromRational


{-# LANGUAGE OverloadedStrings, DeriveGeneric #-}

module Data.Wavelets.Scaling where
import Prelude hiding (length,replicate,map,maximum,minimum)
-- import Data.Wavelets
import Numeric.Statistics
import Statistics.Function
import Statistics.Sample
import GHC.Generics
import Data.Vector.Storable -- hiding (map)
import Data.Packed.Matrix
-- import Linear

{-| Scale factors are designed to make it possible to use a scaled wavelet to reconstruct an approximation
    of the original data series. ordinary least squares is used so that translation is available for approximation,
    in addition to scaling.

|-}

data SeriesFactors = SeriesFactors {
      seriesMin :: Double ,
      seriesMax :: Double ,
      seriesMean :: Double ,
      seriesCount  :: Int } deriving (Show,Generic)

computeSeriesFactors :: Vector Double -> SeriesFactors
computeSeriesFactors v = SeriesFactors mn mx avg tot
    where (mn, mx,avg, tot) = minMaxAndAverage v

newtype OldSeriesFactors = OSF SeriesFactors  deriving (Show)
newtype NewSeriesFactors = NSF SeriesFactors  deriving (Show)
newtype NewSeriesMatrix = NSM (Matrix Double) deriving (Show)
newtype OldSeriesMatrix = OSM (Matrix Double) deriving (Show)

-- | x scaling matrix is B in Y = XB + e matrix formulation of ordinary least squares

newtype ScalingMatrix = SM (Matrix Double) deriving (Show)

type SeriesMatrix = Matrix Double

-- ScaleMatrix is the scaling matrix for the least squares problem
computeSeriesMatrix :: SeriesFactors -> SeriesMatrix
computeSeriesMatrix (SeriesFactors mn mx amean _ ) = (fromColumns vlst)
    where onesList = replicate 3 1
          nsf' = fromList [mn,mx,amean]
          vlst = [onesList,nsf']

computeNewSeriesMatrix :: NewSeriesFactors -> NewSeriesMatrix
computeNewSeriesMatrix (NSF sf) = NSM $ computeSeriesMatrix sf

computeSeriesResult :: SeriesFactors -> SeriesMatrix
computeSeriesResult (SeriesFactors mn mx avg _ ) = (fromColumns vlst)
    where nsf' = fromList [mn,mx,avg]
          vlst = [nsf']

computeOldSeriesMatrix :: OldSeriesFactors -> OldSeriesMatrix
computeOldSeriesMatrix (OSF sf) = OSM $ computeSeriesResult sf

-- | ordinary least squares estimation for the multivariate model
--   Y = X B + e        rows are observations, columns are elements
--   mean e = 0, cov e = kronecker s I

computeScalingMatrix :: NewSeriesFactors -> OldSeriesFactors -> ScalingMatrix
computeScalingMatrix nsf osf = let (OSM _Y) = computeOldSeriesMatrix osf
                                   (NSM _X) = computeNewSeriesMatrix nsf
                                   (_B,_ ,_ ) = ols _X _Y
                               in SM _B

-- | compute the average of the absolute value of a foldable type by setting it up like
-- foldl (\(a,t) b -> absAvgItrFcn a b (t+1) ) (0,0) foldableThing
absAvgItrFcn ::(Num a, Fractional a)=> a -> a -> a -> (a,a)
absAvgItrFcn a b t = avgtuple
   where
     diff = (a - abs(b))
     dv  = abs(diff/t)
     tot  = a + dv
     avgtuple = (tot,t)

-- | compute the maximum absolute value of a foldable type by setting it up like
-- foldl (\mx b -> absMaxItrFcn mx b ) 0 foldableThing
absMaxItrFcn :: (Num a, Ord a) => a -> a -> a
absMaxItrFcn a b = max a (abs b)

-- | compute the max min and average of a Vector in one pass
-- (avg, max, min, tot)
minMaxAndAverage :: Vector Double -> (Double, Double ,Double , Int )
minMaxAndAverage v = (mn,mx,av,l)
    where
      av = mean v
      (mn, mx) = minMax v
      l = length v

applyScalingMatrix :: ScalingMatrix -> Vector Double -> Vector Double
applyScalingMatrix (SM sm) v = map scaleFcn v
    where tCoef = sm @@> (0,0)
          sCoef = sm @@> (1,0)
          scaleFcn x = sCoef * x + tCoef


module SparseLinAlg (solve_tree) where

import Data.Bifunctor (bimap)
import Data.List (sort)
import Data.Maybe (fromJust)
import Data.Sparse.SpMatrix
import Data.Sparse.SpVector
import Numeric.LinearAlgebra.Sparse

import LinEq (Coeff, Var, Equation(..), remove_term, combine_terms, ltree_of_tree, equations_of_ltree)
import Sexp
import Tree
import Util (debug)

-- Here our equations have a slightly different form, where the LHS is
-- a linear combination of the variables and the RHS is a single
-- rational value. E.g., 1x - 1/2y = 1/2.
type MatTerm = (Coeff, Var)
type MatEq = ([MatTerm], Double)

-- mateq_of_equation :: Equation -> MatEq
-- mateq_of_equation (Equation (x, tms)) =
--   case remove_term Nothing $ combine_terms tms of
--     Just (c, tms') -> (second fromJust <$> tms', c)
--     Nothing -> (second fromJust <$> tms, 0)

mateq_of_equation :: Equation -> MatEq
mateq_of_equation (Equation (x, tms)) =
  case remove_term Nothing $ combine_terms tms of
    Just (c, tms') -> ((1, x) : (bimap negate fromJust <$> tms'), c)
    Nothing -> ((1, x) : (bimap negate fromJust <$> tms), 0)

constraint_matrix :: [MatEq] -> SpMatrix Double
constraint_matrix eqs =
  let l = concat $ f <$> zip [0..] eqs in
    debug ("l: " ++ show ((\(x, y, z) -> (x, y, z)) <$> l)) $
    fromListSM (n, n) l
  where
    n = length eqs

    f :: (Int, MatEq) -> [(Int, Int, Double)]
    f (x, (tms, _)) = g x <$> tms

    g :: Int -> MatTerm -> (Int, Int, Double)
    g x (c, y) = (x, y, c)
    -- g x (c, y) = (y, x, c)

-- tree_constraint_matrix :: Tree Bool -> SpMatrix Double
-- tree_constraint_matrix =
--   fmap fromRational . constraint_matrix . map mateq_of_equation .
--   equations_of_ltree . ltree_of_tree

-- tree_constraint_matrix :: Tree Bool -> SpMatrix Double
-- tree_constraint_matrix =
--   fmap fromRational .
--   constraint_matrix .
--   map (\eq -> let mateq = mateq_of_equation eq in
--                 debug ("eq: " ++ show eq) $
--                 debug ("mateq" ++ show mateq) $
--                 mateq) .
--   equations_of_ltree .
--   ltree_of_tree

matrix_of_mateqs :: [MatEq] -> SpMatrix Double
matrix_of_mateqs = constraint_matrix

rhs_of_mateqs :: [MatEq] -> SpVector Double
rhs_of_mateqs eqs = fromListDenseSV n $ snd <$> eqs
  where n = length eqs

solve_system_gmres :: SpMatrix Double -> SpVector Double -> IO (SpVector Double)
solve_system_gmres mat rhs =
  debug ("mat: " ++ show mat) $
  debug ("rhs: " ++ show rhs) $
  debug ("dense:") $
  -- let _ = unsafePerformIO $ prd mat in
  -- mat <\> (fromListDenseSV n )
  mat <\> rhs
  -- where
  --   n = nrows mat -- should also be the length of the rhs vector

solve_tree :: Tree Double -> IO (SpVector Double)
solve_tree (Leaf x) = return $ fromListDenseSV 1 [x]
solve_tree t =
  let lt = ltree_of_tree t
      eqs = sort $ equations_of_ltree lt
      mateqs = mateq_of_equation <$> eqs
      (mat, rhs) = (matrix_of_mateqs mateqs, rhs_of_mateqs mateqs) in
    debug ("eqs: " ++ show eqs) $
    debug ("mateqs: " ++ show mateqs) $
    debug ("ltree: " ++ toSexp lt) $ do
    prd mat
    prd rhs
    sol <- solve_system_gmres mat rhs
    prd sol
    return sol

-- infer_gmres :: Tree Bool -> Maybe Rational
-- infer_gmres t =
--   toRational <$>
--   (lookupSV 0 $ unsafePerformIO $ solve_matrix $ tree_constraint_matrix t)

-- [(0,[(0,1.0),(1,-0.5),(4,-0.5)]),
--  (1,[(0,-0.5),(1,1.0),(3,-0.5)]),
--  (2,[(1,-0.5),(3,1.0)]),
--  (3,[(0,-0.5),(2,-0.5),(4,1.0)]),
--  (4,[(2,1.0),(5,-0.5)]),
--  (5,[(0,-0.5),(2,-0.5),(5,1.0)])]


{-# OPTIONS_HADDOCK hide #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Types
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Basic type classes and enums.
--

module Numeric.LinearAlgebra.Types (

    -- * Matrix views
    Herm(..),
    Tri(..),

    -- * Matrix factorization views
    Chol(..),

    -- * Vector math types
    VNum,
    VFractional,
    VFloating,

    -- * BLAS element types
    BLAS1,
    BLAS2,
    BLAS3,

    -- * LAPACK element types
    LAPACK,

    -- * Enums
    Trans(..),
    Uplo(..),
    Side(..),
    Diag(..),

    -- * Algorithm-specific parameters
    CovMethod(..),

    -- * Re-export of Complex from Data.Complex
    module Data.Complex,

    -- * Re-export of Storable from Foreign.Storable
    module Foreign.Storable,

    ) where

import Foreign.VMath( VNum, VFractional, VFloating )
import Foreign.BLAS( Trans(..), Uplo(..), Side(..), Diag(..), BLAS1, BLAS2, BLAS3 )
import Foreign.LAPACK( LAPACK )
import Data.Complex( Complex(..) )
import Foreign.Storable( Storable() )

-- | A hermitian view of an underlying matrix.  The view can either be
-- of the upper or lower triangular part of the matrix.  The type arguments
-- are as follows:
--
--     * @m@: the underlyting matrix type.
--
--     * @e@: the element type of the matrix.
--
data Herm m e = Herm Uplo (m e) deriving (Show)

-- | A triangular view of an underlying matrix.  The view can either be
-- upper or lower triangular, with a unit or non-unit diagonal.  The type
-- arguments are as follows:
--
--     * @m@: the underlyting matrix type.
--
--     * @e@: the element type of the matrix.
--
data Tri m e = Tri Uplo Diag (m e) deriving (Show)

-- | A Cholesky decomposition view of a matrix.
data Chol m e = Chol Uplo (m e) deriving (Show)

-- | The method of scaling the sample covariance matrix.
data CovMethod =
      UnbiasedCov -- ^ This is the default behavior. Corresponds to a
                  -- scaling of @n/(n-1)@ in the unweighed case, and
                  -- @1/(1 - \\sum w_i^2)@ in the weighted case, where @w_i@
                  -- is the normalized weight. Note the unweighted and
                  -- weighted cases agree when @w_i = 1/n@.

    | MLCov       -- ^ Returns the centered second moment matrix without
                  -- scaling the result.
    deriving (Eq, Show)


{-# LANGUAGE BangPatterns #-}
-- # LANGUAGE UndecidableInstances #
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Data.Network.FullyConnected
  (FeedForwardFC, FullyConnected(FC), Layers,
   weights, biases, activation,
   unGrad, unGradBatch, Grad(GradBatch, Grad))
where

import Numeric.LinearAlgebra as LA hiding (fromList, toList)
import Numeric.LinearAlgebra.Devel
import qualified Data.Vector.Storable as V hiding (foldr, modify)
import System.Random
import Control.DeepSeq
-- import Data.Foldable as F
import Data.Bifunctor
import Data.Network as Network
import Data.Network.Activation
import Algorithm.Training
import Algorithm.GradientDescent
import Data.Network.Utils (squeeze, (.*))

-- Fully connected feedforward layer
data FullyConnected i a = FC {weights :: {-# UNPACK #-} !(Matrix a),
                              biases ::  {-# UNPACK #-} !(Vector a),
                              activation :: ActivationFunction i
                       }

instance (V.Storable a, NFData a) => NFData (FullyConnected i a) where
  rnf (FC w b a) = rnf w `deepseq` rnf b

instance Layer (Vector R) (FullyConnected (Vector R) R) where
  compute (FC weights biases _) input = biases + (input <# weights)
  run     l@(FC weights biases af) input = unAct af (compute l input)
  makeRandomLayer g rf i o af = let rnds = (randomRs (0,1) g) :: [Double]
                                    rnds' = (rf rnds)
                                    rw   = (i><o) rnds'
                                    rb   = o |> (repeat 0)
                                in (FC rw rb af)

instance Layer (Matrix R) (FullyConnected (Matrix R) R) where
  compute (FC weights biases _) input = (fromRows $ replicate (rows input) biases) + (input <> weights)
  run     l@(FC _ _ af) input = unAct af (compute l input)
  makeRandomLayer g rf i o af = let rnds = (randomRs (0,1) g) :: [Double]
                                    rnds' = (rf rnds)
                                    rw   = (i><o) rnds'
                                    rb   = o |> (repeat 0)
                                in (FC rw rb af)

instance (Container Matrix a, Container Vector a, Num a) => Show (FullyConnected i a) where
  show l = let (a,b) = size $ weights l
           in "Layer " ++ (show a) ++ "x" ++ (show b) ++ " " ++ (name $ activation l)

instance Floating a => Num (FullyConnected a Double) where
  (FC w b a) + (FC w' b' _) = FC (w + w') (b + b') a
  (FC w b a) - (FC w' b' _) = FC (w - w') (b - b') a
  (FC w b a) * (FC w' b' _) = FC (w * w') (b * b') a
  abs (FC w b a) = FC (abs w) (abs b) a
  signum (FC w b a) = FC (signum w) (signum b) a
  fromInteger x = FC (fromInteger x) (fromInteger x) (act id "identity")

instance Floating a => Fractional (FullyConnected a Double) where
  (FC w b a) / (FC w' b' _) = FC (w / w') (b / b') a
  recip (FC w b a) = FC (recip w) (recip b) a
  fromRational a = FC (fromRational a) (fromRational a) (act id "identity")

instance (Container Vector a, Container Matrix a) => Linear a (FullyConnected i) where
  scale x (FC w b a) = FC (scale x w) (scale x b) a

-- Container Layer instance?

instance (Layer i (FullyConnected i Double)) => Network i (Layers (FullyConnected i Double)) where
  runNetwork (Output l) i = run l i
  runNetwork (l :=> n)  i = runNetwork n (fst $ run l i)
  runLayers  (Output l) i = [run l i]
  runLayers  (l :=> n)  i = let y = (run l i)
                            in  y:(runLayers n (fst y))

  makeRandomNetwork g t nt i [] (o, af) = let l = makeRandomLayer g t i o af
                                          in  (Output l)
  makeRandomNetwork g t nt i ((h,haf):hs) o = let (g1,g2) = split g
                                                  l = makeRandomLayer g1 t i h haf
                                              in nt $ l :=> (makeRandomNetwork g2 t nt h hs o)

type FeedForwardFC i = Layers (FullyConnected i Double)

instance (Monad m) => Backpropagation Data (NetworkEvaluator Data) (FeedForwardFC Data) m where
  backward lf n (out,tar) das = do
    let δout = (derivate lf (tar, out)) -- dE/dy
        deltas = scanr (\(l, a') δ -> let w = weights l in a' * (w #> δ)) δout (zip (tail $ toList n) das)
    return (deltas) -- deltas for computing bias

instance (Monad m) => Backpropagation Batch (NetworkEvaluator Batch) (FeedForwardFC Batch) m where
  backward lf n (out,tar) das = do
    let δout = tr (derivate lf (tar, out)) -- dE/dy
        deltas = scanr (\(l, a') δ -> let w = weights l in (tr a') * (w <> δ)) δout (zip (tail $ toList n) das)
    return (deltas) -- deltas for computing bias

-- type instance Params lf = forall i. ((FeedForwardFC i), (i,i))

instance (Monad m) => GradientDescent (NetworkEvaluator Batch) ((FeedForwardFC Batch), (Batch,Batch)) m where
  -- parameters of loss function: Network(weights, biases, activation functions for each layer)
  -- and sample((input, output) pair)

  data Grad (NetworkEvaluator Batch) = GradBatch { unGradBatch :: ([Matrix Double], [Vector Double]) } deriving (Show)

  -- Gradients over parameters(weights, biases)
  grad lf (n, (i,t)) = do
    let (as, as') = unzip $ runLayers n i -- forward propagation: compute layers outputs and activation derivatives
        (out) = last as
    (ds) <- backward lf n (out, t) (init as') -- compute deltas with backpropagation
    let r  = fromIntegral $ rows i -- size of minibatch
    let gs = zipWith (\δ a -> tr (δ <> a)) ds (i:init as)
    return $ GradBatch ((recip r .*) <$> gs, ((recip r .*) . squeeze) <$> ds)

  move lr (n, (i,t)) (GradBatch (gs, ds)) = do
    -- update function
    -- divide by batch size
    let update = (\(FC w b af) g δ -> FC (w + (lr).*g) (b + (lr).*δ) af)
        n' = Network.fromList $ zipWith3 update (Network.toList n) gs ds -- new model
    return (n', (i,t))

instance (Monad m) => GradientDescent (NetworkEvaluator Data) ((FeedForwardFC Data), (Data,Data)) m where
  -- parameters of loss function: Network(weights, biases, activation functions for each layer)
  -- and sample((input, output) pair)

  data Grad (NetworkEvaluator Data) = Grad { unGrad :: ([Matrix Double], [Vector Double]) } deriving (Show)

  -- Gradients over parameters(weights, biases)
  grad lf (n, (i,t)) = do
    let aos = runLayers n i -- forward propagation: compute layers outputs and activation derivatives
        (as, as') = (unzip) aos -- juggle data to extract layers inputs and derivatives in right order
        out = last as
    (ds) <- backward lf n (out,t) (init as') -- compute gradients with backpropagation
    let gs = zipWith (\δ i -> tr (δ `outer` i)) ds (i:as)
    return $ Grad (gs, ds)

  move lr (n, (i,t)) (Grad (gs, ds)) = do
    -- update function
    let update = (\(FC w b af) g δ -> FC (w + lr.*g) (b + lr.* δ) af)
        n' = Network.fromList $ zipWith3 update (Network.toList n) gs ds -- new model
    return (n', (i,t))


{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
--------------------------------------------------------------------------------
-- |
-- Robust Principal Component Analysis Outlier Detection a la Netflix/Surus
---------------------------------------------------------------------------------
module RobustPCA
  ( -- * Datatypes
    LR(..)
  , LRPenalty(..)
  , MaxIters(..)
  , FitIters(..)
  , Periodicity(..)
  , RobustPCA(..)
  , Sparse(..)
  , SparsePenalty(..)
  , Tolerance(..)
    -- * Fitting
  , runRobustPCA
  , optimize
  ) where

import           Control.Foldl            (fold)
import           Control.Foldl.Statistics (LMVSK (..), fastLMVSKu, fastStdDev)
import           Control.Lens             (declareWrapped, makeLenses, use,
                                           (+=), (.=), (^.))
import           Control.Monad            (return)
import           Control.Monad.Identity   (runIdentity)
import           Control.Monad.Loops      (whileM_)
import           Control.Monad.State      (State, execStateT)
import           Data.Functor             (void, (<$>))
import           Data.List                (replicate, take)
import           Numeric.LinearAlgebra    (Matrix, R, cmap, cols, diag, flatten,
                                           rows, sumElements, thinSVD, toList,
                                           tr, (<>), (><))
import           Prelude                  (Double, Int, Show, abs, div,
                                           fromIntegral, max, mod, ($), (&&),
                                           (*), (**), (+), (-), (.), (/), (<),
                                           (>))
import           Utils                    (absL1Norm, frobeniusNorm,
                                           softThreshold)

---------------------------------------------------------------------------------
-- Datatypes --------------------------------------------------------------------
---------------------------------------------------------------------------------

declareWrapped [d|
  newtype LRPenalty     = LRPenalty Double     deriving Show
  newtype SparsePenalty = SparsePenalty Double deriving Show
  newtype Periodicity   = Periodicity Int      deriving Show
  newtype MaxIters      = MaxIters Int         deriving Show
  newtype FitIters      = FitIters Int         deriving Show
  newtype LR            = LR [R]               deriving Show
  newtype Sparse        = Sparse [R]           deriving Show
  newtype Tolerance     = Tolerance [R]        deriving Show
  |]

-- |
-- Optimization state
data RobustPCA
  = RobustPCA
  { _maxIters       :: MaxIters
    -- ^ maximum optimization iterations
  , _lrPenalty      :: LRPenalty
    -- ^ low rank fit penalty
  , _sparsePenalty  :: SparsePenalty
    -- ^ sparse fit penalty
  , _normalizedData :: Matrix R
    -- ^ normalized input data matrix
  , _iters          :: Int
    -- ^ iterations count
  , _obj            :: R
    -- ^ objective
  , _tol            :: R
    -- ^ tolerance
  , _diff           :: R
    -- ^ difference
  , _mean           :: R
    -- ^ input data mean
  , _stdDev         :: R
    -- ^ input data standard deviation
  , _mu             :: R
    -- ^ incoherence factor
  , _lr             :: Matrix R
    -- ^ recovered low rank matrix (underlying data 'approximation')
  , _sparse         :: Matrix R
    -- ^ sparse outliers matrix
  , _tolerance      :: Matrix R
    -- ^ tolerated deviation of real data from the 'fit'
  } deriving (Show)
makeLenses ''RobustPCA

-- | Fitting state monad
type RobustPCAM a = State RobustPCA a

---------------------------------------------------------------------------------
-- Fitting ----------------------------------------------------------------------
---------------------------------------------------------------------------------

-- |
-- Run robust PCA optimization
runRobustPCA
  :: [R]
  -- ^ input 'Real' data
  -> LRPenalty
  -- ^ low rank fit penalty
  -> SparsePenalty
  -- ^ sparse fit penalty
  -> Periodicity
  -- ^ input data periodicity
  -> MaxIters
  -- ^ maximum iteration count allowed
  -> (LR, Sparse, Tolerance, FitIters)
runRobustPCA originalData lpen spen p mi =
    (LR lr', Sparse sparse', Tolerance tol', FitIters itrs)
  where rs      = runIdentity $
                   execStateT optimize $
                    initState originalData lpen spen p mi
        mean'   = rs^.mean
        stdDev' = rs^.stdDev
        lr'     = toList $ flatten $ tr $ cmap ((mean' +) . (stdDev' *)) (rs^.lr)
        sparse' = toList $ flatten $ tr $ cmap (stdDev' *) (rs^.sparse)
        tol'    = toList $ flatten $ tr $ cmap (stdDev' *) (rs^.tolerance)
        itrs    = rs^.iters

---------------------------------------------------------------------------------
-- Internal Functions -----------------------------------------------------------
---------------------------------------------------------------------------------

initState
  :: [R]
  -> LRPenalty
  -> SparsePenalty
  -> Periodicity
  -> MaxIters
  -> RobustPCA
initState originalData lp sp (Periodicity p) mi =
    RobustPCA mi lp sp normalizedData' 0 obj' tol' diff' m sdev mu' z z z
  where (LMVSK l m v _ _) = fold fastLMVSKu originalData
        sdev                   = v ** 0.5
        zeroMeanUnitVariance x = (x - m) / sdev
        paddedLength           = l - (l `mod` p)
        normalizedData'        =
          tr $ ((paddedLength `div` p) >< p)
               (zeroMeanUnitVariance <$> take paddedLength originalData)
        r                      = rows normalizedData'
        c                      = cols normalizedData'
        z                      = (r >< c) (replicate (r * c) 0.0)
        mu'                    =
          fromIntegral c * fromIntegral r / (4 * absL1Norm normalizedData')
        obj'                   = 0.5 * ((frobeniusNorm normalizedData') ** 2.0)
        tol'                   = 1e-8 * obj'
        diff'                  = 2.0 * tol'

optimize :: RobustPCAM ()
optimize = whileM_ checkTolerance go
  where go = do
          nuclearNorm <- computeSparse
          l1Norm      <- computeLR
          l2Norm      <- computeTol
          objPrev     <- use obj
          let obj'     = 0.5 * l2Norm + nuclearNorm + l1Norm
          diff        .= abs (objPrev - obj')
          obj         .= obj'
          void updateMu
          iters += 1
        checkTolerance = do
          iters'       <- use iters
          MaxIters mi' <- use maxIters
          tol'         <- use tol
          diff'        <- use diff
          return $ diff' > tol' && iters' < mi'

computeSparse :: RobustPCAM R
computeSparse = do
  normalizedData'          <- use normalizedData
  SparsePenalty spen       <- use sparsePenalty
  mu'                      <- use mu
  lr'                      <- use lr
  let sparsePenalty'        = spen * mu'
  let normalizedDataMinusLr = normalizedData' - lr'
  let penalizedSparse       = softThreshold normalizedDataMinusLr sparsePenalty'
  sparse                   .= penalizedSparse
  return $ (absL1Norm penalizedSparse) * sparsePenalty'

computeLR :: RobustPCAM R
computeLR = do
  normalizedData'              <- use normalizedData
  LRPenalty lpen               <- use lrPenalty
  mu'                          <- use mu
  sparse'                      <- use sparse
  let lrPenalty'                = lpen * mu'
  let normalizedDataMinusSparse = normalizedData' - sparse'
  let (u, svs, v)               = thinSVD normalizedDataMinusSparse
  let penalizedD                = softThreshold svs lrPenalty'
  let d'                        = diag penalizedD
  lr                           .= u <> d' <> (tr v)
  return $ (sumElements penalizedD) * lrPenalty'

computeTol :: RobustPCAM R
computeTol = do
  normalizedData' <- use normalizedData
  lr'             <- use lr
  sparse'         <- use sparse
  let tolerance'   = normalizedData' - lr' - sparse'
  tolerance       .= tolerance'
  return $ (frobeniusNorm tolerance') ** 2

updateMu :: RobustPCAM ()
updateMu = do
  tolerance'  <- use tolerance
  let tdStdDev = fold fastStdDev (toList $ flatten tolerance')
  let m        = rows tolerance'
  let n        = cols tolerance'
  mu .= max 0.01 (tdStdDev * ((2.0 * (fromIntegral $ max m n)) ** 0.5))


{-# LANGUAGE GeneralisedNewtypeDeriving #-}
-- |
-- Module Orbital
--
-- A probability distribution for driver laptimes, initialised from
-- Elo-like ratings.
--
-- In the PDF, t = 0 is taken to be the ideal laptime in a track.
--
-- The name "orbital" alludes to this PDF originally having been, modulo
-- constant factors, the radial probability density of the 1s hydrogen atom
-- orbital. The corresponding distribution can be obtained by setting the
-- gamma shape parameter to 3.
module Orbital
    ( deltaWP
    , eloWP
    , initialRating
    , kFromRating
    , ratingFromK
    , OrbitalDistribution
    , orbitalDistr
    , orbitalK
    ) where

import Numeric.SpecFunctions (choose)
import Statistics.Distribution
import Statistics.Distribution.Gamma
import Data.MemoTrie
import Data.List (foldl')

-- | The base implementation of victory probabilities for performances
-- modelled by gamma distributions with the same (integer) shape.
ratioWP :: Int -> Double -> Double
ratioWP gsh = \w -> w^gsh * evalPoly (memoWpCoefficient gsh) (gsh-1) w

-- | A strict pair type.
data SP a b = SP !a !b

-- | A polynomial evaluator. For our immediate purposes, it performs slightly
-- better than the one available in 'Numeric.Polynomial'.
evalPoly :: (Int -> Double) -> Int -> Double -> Double
evalPoly coef n x = y
    where
    SP xn y = foldl' alg (SP 1 0) [0..n]
    alg (SP xx s) m = SP (x * xx) (coef m * xx + s)

-- | Coefficients of the polynomial used in the gamma winning probability
-- calculation.
wpCoefficient :: Int -> Int -> Double
wpCoefficient gsh m =
    (-1)^m * choose gsh m * choose (2*gsh-1) (gsh-1)
        * fromIntegral (gsh-m) / fromIntegral (gsh+m)

-- | Memoised version of 'wpCoefficient'.
memoWpCoefficient :: Int -> Int -> Double
memoWpCoefficient = memo2 wpCoefficient

-- | Initial rating for new racers. Defined as a constant here for
-- expediteness, taking into account that making it configurable isn't as
-- essential as it is making other engine parameters configurable.
initialRating :: Double
initialRating = 1500

-- | The base Elo exponent. In essence, a logistic growth factor for expected
-- scores/winning probabilities. The denominator used here, 400, is standard
-- for chess rating computations. Using it, a 0.6 expected score corresponds
-- to a ~70.4 rating gap, and a 0.9 expected score, to a ~382 gap.
eloAlpha :: Double
eloAlpha = log 10 / 400

-- | Adjusts the k values obtained from ratings depending on the gamma shape so
-- that the same rating gaps correspond to approximately the same winning
-- probabilities.
--
-- The formula used here is made of two factors:
--
-- - A base correction, which on its own would adjust the derivative of
--   @deltaWP gsh@ to match that of @deltaWP 1@ at zero. This factor, which is
--   the best adjustment for rating gaps close to zero, slightly increases
--   the absolute excess winning probabilities over the entire range of gaps.
--
-- - An additional correction, which reduces the aforementioned bias. Its value
--   is chosen to (approximately) minimise the sum of the squares of the
--   differences between @ratioWP gsh@ and @ratioWP 1@ over the entire range
--   of ratios.
eloGammaCorrection :: Int -> Double
eloGammaCorrection gsh = extraAdjustment * baseCorrection
    where
    gsh' = fromIntegral gsh
    baseCorrection = 4^(gsh-1) / (gsh' * choose (2*gsh-1) (gsh-1))
    fittedConstant = (2*sqrt 29 - 3)*pi
    extraAdjustment = 1 - (gsh'-1) / (gsh' * fittedConstant)

-- | Elo conversion factor. Amounts to eloAlpha for shape 1, in which case the
-- gamma model coincides with the conventional Elo system.
eloFactor :: Int -> Double
eloFactor gsh = eloGammaCorrection gsh * eloAlpha

-- | Memoised Elo conversion factor.
memoEloFactor :: Int -> Double
memoEloFactor = memo eloFactor

-- | Winning probabilities for (scaled) rating differences. Ratings are
-- essentially logarithms of the gamma distribution rate parameters.
deltaWP :: Int -> Double -> Double
deltaWP gsh = \d -> ratioWP gsh (1 / (1 + exp(- memoEloFactor gsh * d)))
-- If ru = log u, rv = log v, and d = ru - rv, then
-- u / (u + v) = 1 / (1 + exp (-d))
-- Another way of expressing the ratio is given by
-- 1 / (1 + exp(-d)) = (1 + tanh (d/2)) / 2

-- | Elo-based winning probability. This version takes the ratings of the two
-- racers. Only here for reference, at least for now.
eloWP :: Double -> Double -> Double
eloWP ru rv = deltaWP 1 (ru-rv)

-- | Rating that corresponds to a gamma rate parameter of 1.
--
-- As far as the results go, the effect of 'referenceRating' is superficial.
-- Since it adjusts all ratings in the same manner, changing it does not affect
-- the overall results.
--
-- There used to be a counterpart @referenceK@ definition, which was removed
-- for being redundant. It is worth noting that scaling the conversions to make
-- the gamma rate ("k") values larger slows down the calculations, probably by
-- making it harder for the gamma distribution algorithms.
referenceRating :: Double
referenceRating = initialRating

-- | k-to-rating conversion. Only for reference, at least for now.
ratingFromK :: Int -> Double -> Double
ratingFromK gsh u = referenceRating + log u / eloFactor gsh

-- | rating-to-k conversion.
kFromRating :: Int -> Double -> Double
kFromRating gsh r = exp ((r - referenceRating) * eloFactor gsh)
-- Using memoEloFactor instead of eloFactor here doesn't seem to improve
-- performance.

-- | The model distribution is a special case of the gamma distribution, with
-- theta = 1/k.
newtype OrbitalDistribution = OrbitalDistribution
    { orbitalGamma :: GammaDistribution }
    deriving (Eq, Show, Distribution, ContDistr, ContGen, MaybeEntropy
        , Variance, MaybeVariance, Mean, MaybeMean)

-- | Sets up the model distribution given a rating.
orbitalDistr
    :: Int    -- ^ Shape parameter to be used.
    -> Double -- ^ Elo-like rating.
    -> OrbitalDistribution
orbitalDistr gsh r = OrbitalDistribution
    (gammaDistr (fromIntegral gsh) (1 / kFromRating gsh r))

-- | Recovers the k factor from the distribution.
orbitalK :: OrbitalDistribution -> Double
orbitalK distr = 1 / gdScale (orbitalGamma distr)

-- | Recovers the shape parameter from the distribution.
orbitalShape :: OrbitalDistribution -> Double
orbitalShape distr = gdShape (orbitalGamma distr)