{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE UndecidableInstances   #-}

module Numeric.Optimization.TwoPhase.Tableau
    ( -- * Tableaus
      Tableau
    , IsTableau
    , mkPhaseI
    , mkPhaseII
    , maximizeEntering
    , maximizeOptimal
    , minimizeEntering
    , minimizeOptimal
    , phaseIEntering
    , phaseIOptimal
    , tableauInfeasible
    , tableauLeaving
    , tableauResult
    , tableauStep
    ) where

import           Data.Finite                                    (Finite)
import qualified Data.List as List
import           Data.Maybe                                     (fromJust)
import           Data.Ord                                       (comparing)
import qualified Data.Vector.Sized as Vec
import qualified Data.Vector.Storable.Sized as SVec
import           GHC.TypeLits
import           Numeric.LinearAlgebra                          (Extractor(..), (??))
import qualified Numeric.LinearAlgebra as LA
import           Numeric.LinearAlgebra.Static                   (L)
import qualified Numeric.LinearAlgebra.Static as LS
import qualified Numeric.LinearAlgebra.Static.Vector as LS
import           Text.Printf                                    (printf)
import           Unsafe.Coerce                                  (unsafeCoerce)

import           Numeric.Optimization.Problem
import           Numeric.Optimization.TwoPhase.Tableau.Build
import           Numeric.Optimization.TwoPhase.Tableau.Pivot
import           Numeric.Optimization.TwoPhase.Tableau.VarMap
import           Numeric.Optimization.TwoPhase.Types

type IsTableau p v s a c =
    ( KnownNat v
    , KnownNat s
    , KnownNat a
    , KnownNat c
    , KnownNat (Rows p c)
    , KnownNat (Cols p v s a)
    )

data Tableau (p :: Phase) (d :: Direction) v s a c =
    Tableau
        { varMap :: VarMap (Rows p c) (Cols p v s a)
        , table  :: L (Rows p c) (Cols p v s a)
        }

instance (IsTableau p v s a c) => Show (Tableau p d v s a c) where
    show (Tableau vs x) = unlines (cols : rows)
      where
        cols        = concatMap (printf "%8s" . show)
            . Vec.toList $ columnVars vs

        rows        = zipWith (printf "%s   %s") rowVals rowNames
        rowVals     = lines . LA.format "" (printf "%8.2f") $ LS.unwrap x
        rowNames    = Vec.toList . Vec.map show $ rowVars vs

mkPhaseI
    :: (IsTableau 'PhaseI v s a c)
    => Problem d v s a c
    -> Tableau 'PhaseI d v s a c
mkPhaseI x =
    Tableau (mkVarMap x) (mkTable x)

mkPhaseII
    :: forall d v s a c.
        ( IsTableau 'PhaseI  v s a c
        , IsTableau 'PhaseII v s a c
        )
    => Tableau 'PhaseI  d v s a c
    -> Tableau 'PhaseII d v s a c
mkPhaseII (Tableau vs x) =
    Tableau (VarMap newRows newColumns) newMatrix
  where
    newRows     = Vec.tail . unsafeCoerce $ rowVars vs

    newColumns  =
        let is = fromJust $ Vec.fromList indices
        in  Vec.backpermute (columnVars vs) is

    newMatrix   =
        let xs = LS.unwrap x ?? (Drop 1, Pos (LA.idxs indices))
        in  fromJust $ LS.create xs

    indices     = fromIntegral <$> findIndicesColumns isPhaseII vs
    isPhaseII n = n /= Objective PhaseI && not (isArtificial n)

-- Read the value associated with a variable in the tableau. If the variable is
-- in the basis (has a row in the tableau), it's value is the value in the RHS
-- of that row divided by the value in the variable's column in that row. If
-- the variable is not in the basis it has a value of 0 (non-basic).
--
readValue
    :: (IsTableau p v s a c)
    => VarName
    -> Tableau p d v s a c
    -> Double
readValue n (Tableau vs x) =
    case (valRowIx, valColIx, rhsColIx) of
        (Just i,  Just j, Just r) -> index (i, r) x / index (i, j) x
        (Nothing, Just _, Just _) -> 0
        _                         -> error "readValue: Variable not in tableau"
  where
    valRowIx    = elemIndexRows n vs
    valColIx    = elemIndexColumns n vs
    rhsColIx    = elemIndexColumns RHS vs

-- Read the value associated with the current objective and each decision
-- variable, and store them in the Vars type for this solver. This is what is
-- returned when optimization succeeds.
--
tableauResult
    :: (IsTableau p v s a c)
    => Tableau p d v s a c
    -> TwoPhaseResult d v
tableauResult x =
    TwoPhaseResult . Vec.zip names $ Vec.map toValue names
  where
    toValue i   = readValue i x
    names       = fromJust . Vec.fromList . findColumns isVar $ varMap x

    isVar n     =
        let objective = indexColumns 0 (varMap x)
        in  n == objective || isDecision n

-- A phase I tableau is infeasible if the value of the objective function is
-- greater than 0. This means that at least one artificial variable has to be
-- non-zero, so there can be no feasible BFS to start phase II from.
--
tableauInfeasible
    :: (IsTableau 'PhaseI v s a c)
    => Tableau 'PhaseI d v s a c
    -> Bool
tableauInfeasible x
    | phaseIOptimal x   = readValue (Objective PhaseI) x /= 0
    | otherwise         = False

-- A tableau is optimal if the given test passes for each cell corresponding to
-- a variable (everything apart from columns for objectives and the RHS column)
-- in the current objective.
--
tableauOptimal
    :: (IsTableau p v s a c)
    => (Double -> Bool)
    -> Tableau p d v s a c
    -> Bool
tableauOptimal f (Tableau vs x) =
    SVec.and . SVec.imap optimal . LS.rVec . head . Vec.toList $ LS.lRows x
  where
    coeffs                  = findIndicesColumns isCoeff vs
    isCoeff                 = not . isSpecial

    optimal i a
        | i `elem` coeffs   = f a
        | otherwise         = True

-- A phase I problem is always a minimization. This is a separate function
-- from 'minimizeOptimal' as the problem being solved in phase II may not be
-- a minimization problem.
--
phaseIOptimal
    :: (IsTableau 'PhaseI v s a c)
    => Tableau 'PhaseI d v s a c
    -> Bool
phaseIOptimal =
    tableauOptimal (<= 0)

-- A phase II problem which is a maximization is optimal when all coefficients
-- in the objective function are non-negative.
--
maximizeOptimal
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Max v s a c
    -> Bool
maximizeOptimal =
    tableauOptimal (>= 0)

-- A phase I problem which is a minimization is optimal when all coefficients
-- in the objective function are non-negative.
--
minimizeOptimal
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Min v s a c
    -> Bool
minimizeOptimal =
    tableauOptimal (<= 0)

-- The entering variable is the lowest indexed variable where the value in the
-- objective column passes a test, and the name of the variable passes another
-- test. The tests change depending on the phase, and whether the optimization
-- is a maximization or a minimization.
--
enteringVar
    :: (IsTableau p v s a c)
    => (Finite (Cols p v s a) -> Bool)
    -> (VarName -> Bool)
    -> Tableau p d v s a c
    -> Either TwoPhaseError (Finite (Cols p v s a))
enteringVar validCol validVar (Tableau vs _) =
    case filter canEnter columns of
        [] -> Left NoEntering
        xs -> Right . fst $ List.minimumBy (comparing snd) xs
  where
    canEnter (j, n) = validCol j && validVar n
    columns         = Vec.toList . Vec.indexed $ columnVars vs

-- In phase I, the entering variable must be a positive value, as the problem
-- is a minimization. Artificial variables are still in the table at this
-- point, so the valid variable test must make sure artificial variables cannot
-- enter the basis (as they should be non-basic).
--
phaseIEntering
    :: (IsTableau 'PhaseI v s a c)
    => Tableau 'PhaseI d v s a c
    -> Either TwoPhaseError (Finite (Cols 'PhaseI v s a))
phaseIEntering x =
    enteringVar validCol validVar x
  where
    validCol j = index (0, j) (table x) > 0
    validVar n = isDecision n || isSlack n

-- When maximizing, the variable to enter the basis is the lowest indexed
-- variable with a negative coefficient in the objective function row.
--
maximizeEntering
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Max v s a c
    -> Either TwoPhaseError (Finite (Cols 'PhaseII v s a))
maximizeEntering x =
    enteringVar validCol validVar x
  where
    validCol j = index (0, j) (table x) < 0
    validVar   = not . isSpecial

-- When minimizing, the variable to enter the basis is the lowest indexed
-- variable with a positive coefficient in the objective function row.
--
minimizeEntering
    :: (IsTableau 'PhaseII v s a c)
    => Tableau 'PhaseII 'Min v s a c
    -> Either TwoPhaseError (Finite (Cols 'PhaseII v s a))
minimizeEntering x =
    enteringVar validCol validVar x
  where
    validCol j = index (0, j) (table x) > 0
    validVar   = not . isSpecial

-- The variable to leave the basis is the lowest indexed variable which has
-- the minimum ratio between it's coefficient and the RHS value. A variable
-- with a coefficient of 0 is not allowed to leave (as the ratio is infinite).
-- This function is the same regardless of the direction of optimization.
--
tableauLeaving
    :: (IsTableau p v s a c)
    => Finite (Cols p v s a)
    -> Tableau p d v s a c
    -> Either TwoPhaseError (Finite (Rows p c))
tableauLeaving enter (Tableau vs x) =
    case toRatio <$> filter canLeave rows of
        [] -> Left Unbounded
        xs -> Right . fst $ List.minimumBy (comparing snd) xs
  where
    rows            = Vec.toList . Vec.indexed $ rowVars vs
    rhsColumn       = fromJust $ elemIndexColumns RHS vs

    toRatio (i, _)  =
        let val = index (i, enter) x
            rhs = index (i, rhsColumn) x
        in  (i, rhs / val)

    canLeave (i, n) =
        let val = index (i, enter) x
            rhs = index (i, rhsColumn) x
        in  val > 0 && rhs /= 0 && not (isSpecial n)

tableauStep
    :: (IsTableau p v s a c)
    => (Tableau p d v s a c
        -> Either TwoPhaseError (Finite (Cols p v s a)))
    -> Tableau p d v s a c
    -> Either TwoPhaseError (Tableau p d v s a c)
tableauStep tableauEntering t@(Tableau vs x) = do
    enter <- tableauEntering t
    leave <- tableauLeaving enter t

    let cell = (leave, enter)
    pure $ Tableau (updateRow cell vs) (pivot cell x)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      :  Spiral.Util.Pretty.LaTeX
-- Copyright   :  (c) 2017-2020 Drexel University
-- License     :  BSD-style
-- Maintainer  :  mainland@drexel.edu

module Spiral.Util.Pretty.LaTeX (
    Pretty(..)
  ) where

import Data.Complex (Complex(..))
import Data.Complex.Cyclotomic
import Data.List (intersperse)
import qualified Data.Map as Map
import qualified Data.Matrix as M
import Data.Modular (Mod, unMod)
import Data.Monoid ((<>))
import Data.Ratio (Ratio, numerator, denominator)
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Symbol (unintern)
import qualified Data.Text as T
import qualified Data.Vector as V
import GHC.TypeLits (KnownNat)
import Text.LaTeX
import Text.LaTeX.Base.Class (comm0,
                              comm1)
import Text.LaTeX.Packages.AMSFonts
import Text.LaTeX.Packages.AMSMath

import Spiral.Array (IArray,
                     Matrix,
                     (!))
import qualified Spiral.Array as A
import Spiral.Array.Shape
import Spiral.Exp
import Spiral.Permutation hiding (Inv)
import qualified Spiral.Permutation as P
import Spiral.SPL
import Spiral.Util.Name (Name(..))
import Spiral.Util.Pretty (Assoc(..),
                           Fixity(..),
                           HasFixity(..),
                           addPrec,
                           infixl_,
                           precOf)
import Spiral.Util.Uniq (Uniq(..))

class Pretty a where
  {-# MINIMAL pprPrec | ppr #-}
  ppr     :: a -> LaTeX
  pprPrec :: Int -> a -> LaTeX

  pprList     :: [a] -> LaTeX
  pprPrecList :: Int -> [a] -> LaTeX

  ppr       = pprPrec 0
  pprPrec _ = ppr

  pprPrecList _ = pprList
  pprList       = autoSquareBrackets . commasep . map ppr

parensIf :: Bool -> LaTeX -> LaTeX
parensIf False l = l
parensIf True l  = autoParens l

hsep :: LaTeX -> [LaTeX] -> LaTeX
hsep sep = mconcat . intersperse sep

commasep :: [LaTeX] -> LaTeX
commasep = hsep ","

mskip :: LaTeX -> LaTeX
mskip = comm1 "mskip"

nicefrac :: LaTeX -> LaTeX -> LaTeX
nicefrac e1 e2 = mempty ^: e1 <> (mskip "-2mu" <> "/" <> mskip "-1mu") !: e2
--nicefrac e1 e2 = e1 <> "/" <> e2

instance Pretty a => Pretty [a] where
  ppr     = pprList
  pprPrec = pprPrecList

instance Pretty a => Pretty (Set a) where
  ppr = autoBraces . commasep . map ppr . Set.toList

instance Pretty Bool where
  ppr True  = mathfrak "T"
  ppr False = mathfrak "F"

pprSignedIntegral :: (Integral a, Show a) => a -> LaTeX
pprSignedIntegral x
  | x < 0     = "-" <> pprSignedIntegral (-x)
  | otherwise = (fromString . show) x

pprRealFloat :: (RealFloat a, Show a) => a -> LaTeX
pprRealFloat x
    | isIntegral x = ppr (ceiling x :: Integer)
    | x < 0        = fromString $ showFloat (-x)
    | otherwise    = fromString $ showFloat x
  where
    isIntegral :: RealFrac a => a -> Bool
    isIntegral x = fromIntegral (ceiling x :: Integer) == x

instance Pretty Int where
  ppr = pprSignedIntegral

instance Pretty Integer where
  ppr = pprSignedIntegral

instance Pretty Float where
  ppr = pprRealFloat

instance Pretty Double where
  ppr = pprRealFloat

instance (Eq a, Num a, Pretty a) => Pretty (Ratio a) where
  ppr x
      | d == 1    = ppr n
      | otherwise = nicefrac (ppr n) (ppr d)
    where
      n = numerator x
      d = denominator x

-- | Pretty-print an imaginary number
pprIm :: (Eq a, Num a, Pretty a) => a -> LaTeX
pprIm 0    = mempty
pprIm 1    = "i"
pprIm (-1) = "-i"
pprIm i    = ppr i <> "i"

instance (Eq a, Num a, Pretty a) => Pretty (Complex a) where
  ppr (r :+ 0)    = ppr r
  ppr (0 :+ 1)    = "i"
  ppr (0 :+ (-1)) = "-i"
  ppr (0 :+ i)    = pprIm i
  ppr (r :+ i)    = case (T.unpack . render) (pprIm i) of
                      '-' : _ -> ppr r <> pprIm i
                      _ -> ppr r + pprIm i

instance Pretty Cyclotomic where
  pprPrec p (Cyclotomic n mp)
    | Map.null mp = "0"
    | null xs     = leadingTerm rat n ex
    | otherwise   = parensIf (p > addPrec) $
                    leadingTerm rat n ex <> mconcat (map (followingTerm n) xs)
    where
      ((ex,rat):xs) = Map.toList mp

      pprBaseExp :: Integer -> Integer -> LaTeX
      pprBaseExp n 1  = omega !: ppr n
      pprBaseExp n ex = omega !^ (ppr n, ppr ex)

      leadingTerm :: Rational -> Integer -> Integer -> LaTeX
      leadingTerm r _ 0 = ppr r
      leadingTerm r n ex
        | r == 1     = t
        | r == (-1)  = "-" <> t
        | r > 0      = ppr r <> t
        | r < 0      = "-" <> ppr (-r) <> t
        | otherwise  = mempty
        where
          t = pprBaseExp n ex

      followingTerm :: Integer -> (Integer, Rational) -> LaTeX
      followingTerm n (ex, r)
        | r == 1     = "+" <> t
        | r == (-1)  = "-" <> t
        | r > 0      = "+" <> ppr r <> t
        | r < 0      = "-" <> ppr (-r) <> t
        | otherwise  = mempty
        where
          t = pprBaseExp n ex

instance (Pretty i, KnownNat p) => Pretty (Mod i p) where
  ppr z = ppr (unMod z)

instance Pretty (Const a) where
  ppr (BoolC x)      = ppr x
  ppr (IntC x)       = ppr x
  ppr (IntegerC x)   = ppr x
  ppr (FloatC x)     = ppr x
  ppr (DoubleC x)    = ppr x
  ppr (RationalC x)  = ppr x
  ppr (ComplexC r i) = ppr (r :+ i)
  ppr (W _ 0 _)      = "1"
  ppr (W n 1 _)      = omega !: ppr n
  ppr (W n k _)      = omega !^ (ppr n, ppr k)
  ppr (CycC x)       = ppr x
  ppr (PiC k)        = ppr k <> pi_
  ppr (ModularC x)   = ppr x

instance Pretty Uniq where
  ppr (Uniq u) = ppr u

instance Pretty Name where
  ppr (Name sym Nothing)  = fromString (unintern sym)
  ppr (Name sym (Just u)) = fromString (unintern sym) !: ppr u

instance Pretty Var where
  ppr (Var n) = ppr n

instance (Pretty e, IArray r DIM2 e) => Pretty (Matrix r e) where
  ppr mat = bmatrix Nothing $ M.matrix m n f
    where
      Z :. m :. n = A.extent mat

      f :: (Int, Int) -> LaTeX
      f (i, j) = ppr (mat ! (i-1, j-1))

instance Pretty (Exp a) where
  pprPrec p (ConstE c) = pprPrec p c
  pprPrec p (VarE v)   = pprPrec p v

  pprPrec _ (UnopE Abs e) =
    autoBrackets "|" "|" (ppr e)

  pprPrec _ (UnopE Sqrt e) =
    tsqrt Nothing (ppr e)

  pprPrec p (UnopE op e) =
    unop p op e

  pprPrec p (BinopE op e1 e2) =
    infixop p op e1 e2

  pprPrec _ (IdxE ev eis) =
    ppr ev <> autoSquareBrackets (commasep (map ppr eis))

  pprPrec p (ComplexE er ei) =
    parensIf (p > addPrec) $
    ppr (er :+ ei)

  pprPrec _ (ReE e) =
    comm0 "Re" <> pprPrec 10 e

  pprPrec _ (ImE e) =
    comm0 "Im" <> pprPrec 10 e

  pprPrec p (BBinopE op e1 e2) =
    infixop p op e1 e2

  pprPrec p (IfE e1 e2 e3) =
    parensIf (p > 10) $
    operatorname (mathsf "if") <> ppr e1 <>
    operatorname (mathsf "then") <> ppr e2 <>
    operatorname (mathsf "else") <> ppr e3

instance Pretty Unop where
  ppr Neg    = "-"
  ppr Abs    = operatorname "abs"
  ppr Signum = operatorname "sgn"
  ppr Exp    = texp
  ppr Log    = tlog
  ppr Sqrt   = operatorname "sqrt"
  ppr Sin    = tsin
  ppr Cos    = tcos
  ppr Asin   = arcsin
  ppr Acos   = arccos
  ppr Atan   = arctan
  ppr Sinh   = tsinh
  ppr Cosh   = tcosh
  ppr Asinh  = operatorname "arsinh"
  ppr Acosh  = operatorname "arcosh"
  ppr Atanh  = operatorname "artanh"

instance Pretty Binop where
  ppr Add  = "+"
  ppr Sub  = "-"
  ppr Mul  = comm0 "cdot"
  ppr Quot = comm0 "quot"
  ppr Rem  = operatorname "rem"
  ppr Div  = operatorname "div"
  ppr Mod  = operatorname "mod"
  ppr FDiv = "/"

instance Pretty BBinop where
  ppr Eq = "="
  ppr Ne = comm0 "neq"
  ppr Lt = "<"
  ppr Le = comm0 "le"
  ppr Ge = comm0 "ge"
  ppr Gt = ">"

instance Pretty Permutation where
  ppr (L mn n)        =  "L" !^ (ppr n, ppr mn)
  ppr (J n)           = operatorname "J" !: ppr n
  ppr (CS n k)        = operatorname "S" !: ppr n <> autoParens (ppr k)
  ppr (CRT m n _ _)   = operatorname "CRT" !^ (ppr n, ppr m)
  ppr (CRT' m n _ _)  = operatorname "CRT'" !^ (ppr n, ppr m)
  ppr (Good m n _ _)  = operatorname "Good" !^ (ppr n, ppr m)
  ppr (Good' m n _ _) = operatorname "Good'" !^ (ppr n, ppr m)
  ppr (R p a)         = operatorname "R" !^ (ppr p, ppr a)
  ppr (P.Inv p _)     = ppr p ^: "-1"

data MatrixBinop = KOp
                 | DSOp
  deriving (Eq, Ord, Show)

instance HasFixity MatrixBinop where
  fixity KOp  = infixl_ 7
  fixity DSOp = infixl_ 6

instance Pretty MatrixBinop where
  ppr KOp  = comm0 "otimes"
  ppr DSOp = comm0 "oplus"

pprArgs :: Pretty a => [a] -> LaTeX
pprArgs = autoParens . commasep . map ppr

instance (Num e, Pretty e) => Pretty (SPL e) where
  pprPrec p m@E{}       = pprPrec p (toMatrix m)
  pprPrec _ (I n)       = operatorname "I" !: ppr n
  pprPrec _ (T e)       = pprPrec 10 e ^: comm0 "intercal"
  pprPrec _ (Inv e)     = pprPrec 10 e ^: "-1"
  pprPrec p (Pi pi)     = pprPrec p pi
  pprPrec _ (Rot alpha) = operatorname "R" !: ppr alpha
  pprPrec _ (Diag xs)   = operatorname "diag" <> pprArgs (V.toList xs)
  pprPrec p (KDiag _ e) = pprPrec p e

  pprPrec _ ((tl `Beside` tr) `Above` (bl `Beside` br)) =
    bmatrix Nothing $ M.matrix 2 2 f
      where
        f :: (Int, Int) -> LaTeX
        f (1, 1) = ppr tl
        f (1, 2) = ppr tr
        f (2, 1) = ppr bl
        f (2, 2) = ppr br
        f _      = error "can't happen"

  pprPrec _ (Above a b) =
    bmatrix Nothing $ M.matrix 2 1 f
      where
        f :: (Int, Int) -> LaTeX
        f (1, 1) = ppr a
        f (2, 1) = ppr b
        f _      = error "can't happen"

  pprPrec _ (Beside a b) =
    bmatrix Nothing $ M.matrix 1 2 f
      where
        f :: (Int, Int) -> LaTeX
        f (1, 1) = ppr a
        f (1, 2) = ppr b
        f _      = error "can't happen"

  pprPrec p (Kron a b)  = infixop p KOp a b
  pprPrec p (DSum a b)  = infixop p DSOp a b
  pprPrec _ (Prod a b)  = pprPrec 10 a <> pprPrec 10 b
  pprPrec _ (Circ xs)   = operatorname "circ" <> pprArgs (V.toList xs)
  pprPrec _ (Skew xs)   = operatorname "skew" <> pprArgs (V.toList xs)
  pprPrec _ (Toep xs)   = operatorname "toep" <> pprArgs (V.toList xs)
  pprPrec _ (Re a)      = comm0 "Re" <> autoParens (ppr a)
  pprPrec _ F2          = operatorname "F" !: "2"
  pprPrec _ (DFT n)     = operatorname "DFT" !: ppr n
  pprPrec _ (DFT' n)    = operatorname "DFT" !^ (ppr n, "-1")
  pprPrec _ (F n w)     = operatorname "DFT" !: ppr n <> autoParens (ppr w)
  pprPrec _ (F' n w)    = operatorname "DFT" !^ (ppr n, "-1") <> autoParens (ppr w)

unop :: (Pretty a, Pretty op, HasFixity op)
     => Int -- ^ precedence of context
     -> op  -- ^ operator
     -> a
     -> LaTeX
unop prec op x =
    parensIf (prec > precOf op) $
    ppr op <> pprPrec (precOf op) x

infixop :: (Pretty a, Pretty b, Pretty op, HasFixity op)
        => Int -- ^ precedence of context
        -> op  -- ^ operator
        -> a   -- ^ left argument
        -> b   -- ^ right argument
        -> LaTeX
infixop prec op l r =
    parensIf (prec > opPrec) $
    pprPrec leftPrec l <> ppr op <> pprPrec rightPrec r
  where
    leftPrec | opAssoc == RightAssoc = opPrec + 1
             | otherwise             = opPrec

    rightPrec | opAssoc == LeftAssoc = opPrec + 1
              | otherwise            = opPrec

    Fixity opAssoc opPrec = fixity op


{-# LANGUAGE BangPatterns, GeneralizedNewtypeDeriving, StandaloneDeriving #-}

import           Prelude hiding (mapM)

import           Options.Applicative
import           Data.Monoid ((<>))
import           Control.Monad.Trans.Class

import           Data.Vector (Vector)
import qualified Data.Vector.Generic as V
import           Statistics.Sample (mean)

import           Data.Traversable (mapM)
import qualified Data.Set as S
import           Data.Set (Set)
import qualified Data.Map.Strict as M

import           ReadData
import           SerializeText
import qualified RunSampler as Sampler
import           BayesStack.DirMulti
import           BayesStack.Models.Topic.LDA
import           BayesStack.UniqueKey

import qualified Data.Text as T
import qualified Data.Text.IO as TIO

import           System.FilePath.Posix ((</>))
import           Data.Binary
import qualified Data.ByteString as BS
import           Text.Printf

import           Data.Random
import           System.Random.MWC

data RunOpts = RunOpts { nodesFile       :: FilePath
                       , stopwords       :: Maybe FilePath
                       , nTopics         :: Int
                       , samplerOpts     :: Sampler.SamplerOpts
                       , hyperParams     :: HyperParams
                       }

runOpts :: Parser RunOpts
runOpts = RunOpts
    <$> strOption  ( long "nodes"
                  <> short 'n'
                  <> metavar "FILE"
                  <> help "File containing nodes and their associated items"
                   )
    <*> nullOption ( long "stopwords"
                  <> short 's'
                  <> metavar "FILE"
                  <> reader (pure . Just)
                  <> value Nothing
                  <> help "Stop word list"
                   )
    <*> option     ( long "topics"
                  <> short 't'
                  <> metavar "N"
                  <> value 20
                  <> help "Number of topics"
                   )
    <*> Sampler.samplerOpts
    <*> hyperOpts

hyperOpts = HyperParams
    <$> option     ( long "prior-theta"
                  <> value 1
                  <> help "Dirichlet parameter for prior on theta"
                   )
    <*> option     ( long "prior-phi"
                  <> value 0.1
                  <> help "Dirichlet parameter for prior on phi"
                   )

mapMKeys :: (Ord k, Ord k', Monad m, Applicative m)
         => (a -> m a') -> (k -> m k') -> M.Map k a -> m (M.Map k' a')
mapMKeys f g x = M.fromList <$> (mapM (\(k,v)->(,) <$> g k <*> f v) $ M.assocs x)

termsToItems :: M.Map NodeName [Term]
             -> (M.Map Node [Item], (M.Map Item Term, M.Map Node NodeName))
termsToItems nodes =
    let ((d', nodeMap), itemMap) =
            runUniqueKey' [Item i | i <- [0..]] $
            runUniqueKeyT' [Node i | i <- [0..]] $ do
                mapMKeys (mapM (lift . getUniqueKey)) getUniqueKey nodes
    in (d', (itemMap, nodeMap))

netData :: HyperParams -> M.Map Node [Item] -> Int -> NetData
netData hp nodeItems nTopics =
    NetData { dHypers           = hp
            , dItems            = S.unions $ map S.fromList $ M.elems nodeItems
            , dTopics           = S.fromList [Topic i | i <- [1..nTopics]]
            , dNodeItems        = M.fromList
                                  $ zip [NodeItem i | i <- [0..]]
                                  $ do (n,items) <- M.assocs nodeItems
                                       item <- items
                                       return (n, item)
            , dNodes            = M.keysSet nodeItems
            }

opts :: ParserInfo RunOpts
opts = info runOpts (  fullDesc
                    <> progDesc "Learn LDA model"
                    <> header "run-lda - learn LDA model"
                    )

instance Sampler.SamplerModel MState where
    estimateHypers = reestimate
    modelLikelihood = modelLikelihood
    summarizeHypers ms =
        "  phi  : "++show (dmAlpha $ snd $ M.findMin $ stPhis ms)++"\n"++
        "  theta: "++show (dmAlpha $ snd $ M.findMin $ stThetas ms)++"\n"

main :: IO ()
main = do
    args <- execParser opts
    stopWords <- case stopwords args of
                     Just f  -> S.fromList . T.words <$> TIO.readFile f
                     Nothing -> return S.empty
    printf "Read %d stopwords\n" (S.size stopWords)

    (nodeItems, (itemMap, nodeMap)) <- termsToItems
                            <$> readNodeItems stopWords (nodesFile args)

    Sampler.createSweeps $ samplerOpts args
    let sweepsDir = Sampler.sweepsDir $ samplerOpts args
    encodeFile (sweepsDir </> "item-map") itemMap
    encodeFile (sweepsDir </> "node-map") nodeMap

    let termCounts = V.fromListN (M.size nodeItems)
                     $ map length $ M.elems nodeItems :: Vector Int
    printf "Read %d nodes\n" (M.size nodeItems)
    printf "Mean items per node:  %1.2f\n" (mean $ V.map realToFrac termCounts)

    withSystemRandom $ \mwc->do
    let nd = netData (hyperParams args) nodeItems (nTopics args)
    encodeFile (sweepsDir </> "data") nd
    mInit <- runRVar (randomInitialize nd) mwc
    let m = model nd mInit
    Sampler.runSampler (samplerOpts args) m (updateUnits nd)
    return ()


{-# LANGUAGE RankNTypes,
             ScopedTypeVariables,
             ViewPatterns,
             TypeFamilies,
             TypeOperators,
             RankNTypes,
             MultiWayIf,
             DataKinds,
             LambdaCase,
             ConstraintKinds,
             FlexibleContexts #-}

module Networks where

import Optimierung
import Tuple
import Numeric.LinearAlgebra.Static.Backprop
import Numeric.Backprop
import GHC.TypeNats
import qualified Numeric.LinearAlgebra.Static as H
import qualified Control.Foldl as L
import qualified Data.Vector.Storable as V

type KN i o = (KnownNat i, KnownNat o)
type Fehler a b = forall z. Reifies z W => BVar z a -> BVar z a -> BVar z b
type Res p a b c = forall z. Reifies z W => BVar z (a :& b) -> BVar z p -> BVar z c

type Modell p a b =
       forall z. Reifies z W
    => BVar z p
    -> BVar z a
    -> BVar z b

withScalarActivation :: (Reifies z W, KnownNat i)
                     => (BVar z Double -> BVar z Double)
                     -> BVar z (R i :& Double)
                     -> BVar z (R i)
                     -> BVar z Double
withScalarActivation f (w :&& b) v = f $ w <.> v + b

withActivation :: (Reifies z W, KN i o)
               => (BVar z (R o) -> BVar z (R o))
               -> BVar z (L o i :& R o)
               -> BVar z (R i)
               -> BVar z (R o)
withActivation f (w :&& b) v = f $ w #> v + b

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

linear :: KnownNat i => Modell (R i :& Double) (R i) Double
linear = withScalarActivation id

linear2 :: KnownNat i => Modell (R i :& Double) (R i) Double
linear2 (w :&& b) x = w <.> x + b

sigmoid :: KN i o => Modell (L o i :& R o) (R i) (R o)
sigmoid = withActivation logistic

--todo: dvmap isn't actually made differentiable, manually make Op
leakyReLU :: KN i o => Modell (L o i :& R o) (R i) (R o)
leakyReLU = withActivation $ dvmap $ \x -> if (x < 0) then 0.01*x
                                                      else x

sigmoidNoBias :: KN i o => Modell (L o i) (R i) (R o)
sigmoidNoBias w v = logistic $ w #> v

sigmoid' :: KnownNat i => Modell (R i :& Double) (R i) Double
sigmoid' = withScalarActivation logistic

se :: Fehler Double Double
se x a = (x - a)^2

se' :: (KnownNat n, 1 <= n) => Fehler (R n) Double
se' a e = (d <.> d) / 2
    where d = e - a

crossentropy' :: (KnownNat n, 1 <= n) => Fehler (R n) Double
crossentropy' a y = vsum . nanToNum . negate $ entropy y a + entropy (1-y) (1-a)
    where entropy x y = x * log y

crossentropyUnsafe' :: (KnownNat n, 1 <= n) => Fehler (R n) Double
crossentropyUnsafe' a y = vsum . negate $ entropy y a + entropy (1-y) (1-a)
    where entropy x y = x * log y

crossentropy :: (KnownNat n, 1 <= n) => R n -> R n -> Double
crossentropy a y = (H.konst 1) H.<.> (entropy (-y) a - entropy (1-y) (1-a))
    where entropy x y = x * log y

logisticSample :: KN i o => (L o i :& R o) -> Int -> R i -> R o
logisticSample (w :& b) s v = H.zipWithVector (\x p -> if x < p then 1 else 0) r probs
    where probs = logistic $ w H.#> v + b
          r     = H.randomVector s H.Uniform

nonModelLogistic :: KN i o => (L o i :& R o) -> R i -> R o
nonModelLogistic (w :& b) v = logistic $ w H.#> v + b

vsum :: (Reifies z W, KnownNat n) => BVar z (R n) -> BVar z Double
vsum x = (konst 1) <.> x

rbmcd1 :: KN d h => Int -> (L d h :& R d :& R h) -> R d -> (L d h :& R d :& R h)
rbmcd1 s (w :& b :& c) v1 = (H.outer v1 h1 - H.outer v2 h2) :& (v1 - v2) :& (h1 - h2)
    where h1 = (logisticSample (H.tr w :& c) s v1)
          v2 = (nonModelLogistic (w :& b) h1)
          h2 = (nonModelLogistic (H.tr w :& c) v2)

gradientStep :: (Backprop p, Backprop b, Backprop c)
             => Modell p a b  -- ^ Modell zum Trainieren
             -> Fehler b c    -- ^ Fehlerfunktion
             -> p             -- ^ Parameter am Anfang des Schrittes
             -> (a, b)        -- ^ (Eingabedatum, erwartete Antwort)
             -> p             -- ^ Gradient des Parameters
gradientStep mod err param (vec, trg) = gradBP (\p -> err (mod p (auto vec)) (auto trg)) param

makeFoldAll :: Optimierung d p
         => d p
         -> (p -> a -> p)
         -> p
         -> L.Fold a p
makeFoldAll d f p0 = L.Fold (\x a -> regel d x (f (aus x) a)) (start p0) aus

uberwachtesFold :: (Optimierung d p, Backprop p, Backprop b, Backprop c)
                => d p
                -> Modell p a b
                -> Fehler b c
                -> p
                -> L.Fold (a, b) p
uberwachtesFold op mod err p0 = makeFoldAll op (gradientStep mod err) p0

isVecNaN :: KnownNat n => R n -> Bool
isVecNaN = V.foldl (\acc x -> acc || isNaN x) False . H.extract

nanToNum :: (KnownNat n, Reifies z W) => BVar z (R n) -> BVar z (R n)
nanToNum = vmap bvToNum

bvToNum :: Reifies z W => BVar z Double -> BVar z Double
bvToNum = liftOp1 . op1 $ \x -> ((if
                            | isNaN x      -> 0.0
                            | isInfinite x -> if x < 0 then -1.7976931348623157e+308
                                                       else  1.7976931348623157e+308
                            | otherwise    -> x), id)

supervised :: (Backprop a, Backprop p, Backprop b, Backprop c)
           => Fehler b c
           -> Modell p a b
           -> Res p a b c
supervised err mod (a :&& y) ps = err (mod ps a) y

class Optimizable m where
    type Params m
    type In     m
    gradient :: m -> Params m -> In m -> Params m

makeFoldOptimizable :: (Optimizable m, Optimierung d (Params m))
                    => m
                    -> d (Params m)
                    -> Params m
                    -> L.Fold (In m) (Params m)
makeFoldOptimizable m d p0 = L.Fold (\x a -> regel d x (step (aus x) a)) (start p0) aus
    where step = gradient m

(<~) :: (Backprop p, Backprop q)
     => Modell  p       b c
     -> Modell      q  a b
     -> Modell (p :& q) a c
(f <~ g) (p :&& q) = f p . g q

infixr 8 <~


{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2015
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-----------------------------------------------------------------------------

module Numeric.AD.Mode
  (
  -- * AD modes
    Mode(..)
  ) where

import Numeric.Natural
import Data.Complex
import Data.Int
import Data.Ratio
import Data.Word

infixr 7 *^
infixl 7 ^*
infixr 7 ^/

class (Num t, Num (Scalar t)) => Mode t where
  type Scalar t
  -- | allowed to return False for items with a zero derivative, but we'll give more NaNs than strictly necessary
  isKnownConstant :: t -> Bool
  isKnownConstant _ = False

  -- | allowed to return False for zero, but we give more NaN's than strictly necessary then
  isKnownZero :: t -> Bool
  isKnownZero _ = False

  -- | Embed a constant
  auto  :: Scalar t -> t

  -- | Scalar-vector multiplication
  (*^) :: Scalar t -> t -> t
  a *^ b = auto a * b

  -- | Vector-scalar multiplication
  (^*) :: t -> Scalar t -> t
  a ^* b = a * auto b

  -- | Scalar division
  (^/) :: Fractional (Scalar t) => t -> Scalar t -> t
  a ^/ b = a ^* recip b

  -- |
  -- @'zero' = 'lift' 0@
  zero :: t
  zero = auto 0

instance Mode Double where
  type Scalar Double = Double
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Float where
  type Scalar Float = Float
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Int where
  type Scalar Int = Int
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Integer where
  type Scalar Integer = Integer
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Int8 where
  type Scalar Int8 = Int8
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Int16 where
  type Scalar Int16 = Int16
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Int32 where
  type Scalar Int32 = Int32
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Int64 where
  type Scalar Int64 = Int64
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Natural where
  type Scalar Natural = Natural
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Word where
  type Scalar Word = Word
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Word8 where
  type Scalar Word8 = Word8
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Word16 where
  type Scalar Word16 = Word16
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Word32 where
  type Scalar Word32 = Word32
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Mode Word64 where
  type Scalar Word64 = Word64
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance RealFloat a => Mode (Complex a) where
  type Scalar (Complex a) = Complex a
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)

instance Integral a => Mode (Ratio a) where
  type Scalar (Ratio a) = Ratio a
  isKnownConstant _ = True
  isKnownZero x = 0 == x
  auto = id
  (^/) = (/)


{-# LANGUAGE FlexibleContexts #-}

module Main (main) where

import Control.Monad.IO.Class (liftIO)
import Data.Complex (Complex)
import Data.Foldable (toList)
import Text.PrettyPrint.Mainland
import Text.PrettyPrint.Mainland.Class

import Spiral
import Spiral.Backend.C
import Spiral.Config
import Spiral.Driver
import Spiral.Exp
import Spiral.OpCount
import Spiral.Program
import Spiral.RootOfUnity
import Spiral.SPL
import Spiral.SPL.Run
import Spiral.Util.Uniq

main :: IO ()
main = defaultMain $ \args -> do
    useComplexType <- asksConfig $ testDynFlag UseComplex
    n <- case args of
           [s] -> return (read s)
           _   -> return 4
    let f = formula n
    pprint f
    if useComplexType
      then toProgram "f" f >>= go
      else toProgram "f" (Re f) >>= go
  where
    go :: (Typed a, Num (Exp a)) => Program a -> Spiral ()
    go prog = do
      pprint prog
      ops <- countProgramOps prog
      resetUnique
      defs <- evalCg $ cgProgram prog
      outp <- asksConfig output
      case outp of
        Nothing -> return ()
        Just{}  -> writeOutput (toList defs)
      liftIO $ putDocLn $
          text "Multiplications:" <+> ppr (mulOps ops) </>
          text "      Additions:" <+> ppr (addOps ops) </>
          text "          Total:" <+> ppr (allOps ops)

-- The SPL formula for which we generate code and count operations.
formula :: Int -> SPL (Exp (Complex Double))
formula _ =
    Pi (L 8 4) ×
    (I 2 ⊗ (Pi (L 4 2) × (I 2 ⊗ F 2 w2) × diag [1.0, 1.0, 1.0, w4] × (F 2 w2 ⊗ I 2))) ×
    diag [1.0, 1.0, 1.0, 1.0, 1.0, w8, w4, w8^(3 :: Int)] ×
    (F 2 w2 ⊗ I 4)
  where
    w2, w4, w8 :: Exp (Complex Double)
    w2 = omega 2
    w4 = omega 4
    w8 = omega 8


{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE KindSignatures         #-}
{-# LANGUAGE RecordWildCards        #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeOperators          #-}

module Numeric.Optimization.TwoPhase.Tableau.Build
    ( IsBuilder
    , mkTable
    ) where

import           Prelude hiding                             ((++))

import qualified Data.List as List
import           Data.Maybe                                 (fromJust)
import qualified Data.Vector.Sized as Vec
import           Data.Vector.Storable.Sized                 (Vector, (++))
import qualified Data.Vector.Storable.Sized as SVec
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static               (L)
import qualified Numeric.LinearAlgebra.Static as LS
import qualified Numeric.LinearAlgebra.Static.Vector as LS

import           Numeric.Optimization.Problem
import           Numeric.Optimization.TwoPhase.Types

type IsBuilder v s a c =
    ( KnownNat v
    , KnownNat s
    , KnownNat a
    , KnownNat c
    )

data Builder (v :: Nat) (s :: Nat) (a :: Nat) (c :: Nat) = Builder
    { slackSupply       :: [Vector s Double]
    , artificialSupply  :: [Vector a Double]
    , rowsBuilt         :: [Vector (Cols 'PhaseI v s a) Double]
    , adjustmentsBuilt  :: [Vector (Cols 'PhaseI v s a) Double]
    }

mkBuilder :: (IsBuilder v s a c) => Builder v s a c
mkBuilder = Builder slackI artificialI [] []
  where
    slackI      = reverse . fmap LS.rVec . LS.toRows $ LS.eye
    artificialI = reverse . fmap LS.rVec . LS.toRows $ LS.eye

buildPhaseI
    :: forall v s a c.
        (IsBuilder v s a c)
    => Builder v s a c
    -> Builder v s a c
buildPhaseI acc = acc
    { rowsBuilt = row : rowsBuilt acc }
  where
    row = SVec.snoc (os ++ xs ++ ss ++ as) 0
    os  = SVec.fromTuple (1, 0)
    xs  = SVec.replicate @v 0
    ss  = SVec.replicate @s 0
    as  = SVec.replicate @a (-1)

buildPhaseII
    :: forall v s a c.
        (IsBuilder v s a c)
    => Builder v s a c
    -> Coeffs v
    -> Builder v s a c
buildPhaseII acc xs' = acc
    { rowsBuilt = row : rowsBuilt acc }
  where
    row = SVec.snoc (os ++ xs ++ ss ++ as) 0
    os  = SVec.fromTuple (0, 1)
    xs  = SVec.map negate xs'
    ss  = SVec.replicate @s 0
    as  = SVec.replicate @a 0

buildObjective :: (IsBuilder v s a c)
    => Builder v s a c -> Coeffs v -> Builder v s a c
buildObjective acc = buildPhaseI . buildPhaseII acc

buildLEQ
    :: forall v s a c.
        (IsBuilder v s a c)
    => Builder v s a c
    -> Coeffs v
    -> Double
    -> Builder v s a c
buildLEQ acc xs rhs = acc
    { slackSupply   = tail (slackSupply acc)
    , rowsBuilt     = row : rowsBuilt acc
    }
  where
    row = SVec.snoc (os ++ xs ++ ss ++ as) rhs
    os  = SVec.fromTuple (0, 0)
    ss  = head (slackSupply acc)
    as  = SVec.replicate @a 0

buildGEQ
    :: forall v s a c.
        (IsBuilder v s a c)
    => Builder v s a c
    -> Coeffs v
    -> Double
    -> Builder v s a c
buildGEQ acc xs' rhs' = acc
    { slackSupply       = tail (slackSupply acc)
    , artificialSupply  = tail (artificialSupply acc)
    , rowsBuilt         = row : rowsBuilt acc
    }
  where
    row = SVec.snoc (os ++ xs ++ ss ++ as) rhs
    os  = SVec.fromTuple (0, 0)
    xs  = SVec.map negate xs'
    ss  = head (slackSupply acc)
    as  = head (artificialSupply acc)
    rhs = negate rhs'

buildEQU
    :: forall v s a c.
        (IsBuilder v s a c)
    => Builder v s a c
    -> Coeffs v
    -> Double
    -> Builder v s a c
buildEQU acc xs rhs = acc
    { artificialSupply  = tail (artificialSupply acc)
    , rowsBuilt         = row : rowsBuilt acc
    , adjustmentsBuilt  = row : adjustmentsBuilt acc
    }
  where
    row = SVec.snoc (os ++ xs ++ ss ++ as) rhs
    os  = SVec.fromTuple (0, 0)
    ss  = SVec.replicate @s 0
    as  = head (artificialSupply acc)

buildConstraint :: (IsBuilder v s a c)
    => Builder v s a c -> Constraint v x y -> Builder v s a c
buildConstraint acc (xs :< y) = buildLEQ acc xs y
buildConstraint acc (xs :> y) = buildGEQ acc xs y
buildConstraint acc (xs := y) = buildEQU acc xs y

toMatrix :: forall v s a c. (IsBuilder v s a c)
    => Builder v s a c -> L (Rows 'PhaseI c) (Cols 'PhaseI v s a)
toMatrix acc =
    LS.rowsL . fromJust . Vec.fromList $ fmap LS.vecR adjusted
  where
    rows    = rowsBuilt acc
    adjusts = adjustmentsBuilt acc

    adjusted = case adjusts of
        [] -> rows
        xs ->
            let oldObjP1  = head rows
                newObjP1  = List.foldl1' (SVec.zipWith (+)) (oldObjP1 : xs)
            in  newObjP1 : tail rows

mkTable :: forall d v s a c. (IsBuilder v s a c)
    => Problem d v s a c -> L (Rows 'PhaseI c) (Cols 'PhaseI v s a)
mkTable = toMatrix . go mkBuilder
  where
    go :: Builder v s a c -> Problem d v x y z -> Builder v s a c
    go acc (Maximize xs)    = buildObjective acc xs
    go acc (Minimize xs)    = buildObjective acc xs
    go acc (SuchThat p c)   = go (buildConstraint acc c) p


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


{-# LANGUAGE RecordWildCards #-}

{-|
Module      :  HABQTlib.Data.Particle

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

-- Author: Ian Ross, for OpenBrain Ltd

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

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

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

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

class Nanable a where
  hasnan :: a -> Bool

instance Nanable Double where
  hasnan = isNaN

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

type Width = Int

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

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

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

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

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

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

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

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

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

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

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

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

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

type Dist a = M.Map a Double

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

logb :: Double -> Double
logb = logBase 2

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE OverloadedStrings     #-}
{-|
Module      : Grenade.Layers.Reshape
Description : Multipurpose reshaping layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Reshape
  ( Reshape(..)
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Serialize
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Data     as LA (flatten)
import           Numeric.LinearAlgebra.Static

import           Grenade.Core

import           Grenade.Onnx

-- | Reshape Layer
--
-- The Reshape layer can flatten any 2D or 3D image to 1D vector with the
-- same number of activations, as well as cast up from 1D to a 2D or 3D
-- shape.
--
-- Can also be used to turn a 3D image with only one channel into a 2D image
-- or vice versa.
data Reshape = Reshape
  deriving (Show,Generic,NFData)

instance UpdateLayer Reshape where
  type Gradient Reshape = ()
  runUpdate _ _ _ = Reshape
  reduceGradient _ = ()

instance RandomLayer Reshape where
  createRandomWith _ _ = return Reshape

instance (KnownNat a, KnownNat x, KnownNat y, a ~ (x * y)) => Layer Reshape ('D2 x y) ('D1 a) where
  type Tape Reshape ('D2 x y) ('D1 a) = ()
  runForwards _ (S2D y)   =  ((), fromJust' . fromStorable . flatten . extract $ y)
  runBackwards _ _ (S1D y) = ((), fromJust' . fromStorable . extract $ y)

instance (KnownNat a, KnownNat x, KnownNat y, KnownNat (x * z), KnownNat z, a ~ (x * y * z)) => Layer Reshape ('D3 x y z) ('D1 a) where
  type Tape Reshape ('D3 x y z) ('D1 a) = ()
  runForwards _ (S3D y)     = ((), fromJust' . fromStorable . flatten . extract $ y)
  runBackwards _ _ (S1D y) = ((), fromJust' . fromStorable . extract $ y)

instance (KnownNat y, KnownNat x, KnownNat z, z ~ 1) => Layer Reshape ('D3 x y z) ('D2 x y) where
  type Tape Reshape ('D3 x y z) ('D2 x y) = ()
  runForwards _ (S3D y)    = ((), S2D y)
  runBackwards _ _ (S2D y) = ((), S3D y)

instance (KnownNat y, KnownNat x, KnownNat z, z ~ 1) => Layer Reshape ('D2 x y) ('D3 x y z) where
  type Tape Reshape ('D2 x y) ('D3 x y z) = ()
  runForwards _ (S2D y)    = ((), S3D y)
  runBackwards _ _ (S3D y) = ((), S2D y)

instance (KnownNat a, KnownNat x, KnownNat y, a ~ (x * y)) => Layer Reshape ('D1 a) ('D2 x y) where
  type Tape Reshape ('D1 a) ('D2 x y) = ()
  runForwards _ (S1D y)   =  ((), fromJust' . fromStorable . extract $ y)
  runBackwards _ _ (S2D y) = ((), fromJust' . fromStorable . flatten . extract $ y)

instance (KnownNat a, KnownNat x, KnownNat y, KnownNat (x * z), KnownNat z, a ~ (x * y * z)) => Layer Reshape ('D1 a) ('D3 x y z) where
  type Tape Reshape ('D1 a) ('D3 x y z) = ()
  runForwards _ (S1D y)     = ((), fromJust' . fromStorable . extract $ y)
  runBackwards _ _ (S3D y)  = ((), fromJust' . fromStorable . flatten . extract $ y)

instance (KnownNat a, KnownNat b, KnownNat c, KnownNat w, KnownNat x, KnownNat y, KnownNat z, (a * b * c) ~ (w * x * y * z)) => Layer Reshape ('D3 a b c) ('D4 w x y z) where
  type Tape Reshape ('D3 a b c) ('D4 w x y z) = ()
  runForwards _ (S3D y)     = ((), fromJust' . fromStorable . flatten . extract $ y)
  runBackwards _ _ (S4D y)  = ((), fromJust' . fromStorable . flatten . extract $ y)

instance (KnownNat a, KnownNat b, KnownNat c, KnownNat d, KnownNat x, KnownNat y, KnownNat z, (a * b * c * d) ~ (x * y * z)) => Layer Reshape ('D4 a b c d) ('D3 x y z) where
  type Tape Reshape ('D4 a b c d) ('D3 x y z) = ()
  runForwards _ (S4D y)     = ((), fromJust' . fromStorable . flatten . extract $ y)
  runBackwards _ _ (S3D y)  = ((), fromJust' . fromStorable . flatten . extract $ y)

instance Serialize Reshape where
  put _ = return ()
  get = return Reshape

fromJust' :: Maybe x -> x
fromJust' (Just x) = x
fromJust' Nothing  = error "Reshape error: data shape couldn't be converted."

instance OnnxOperator Reshape where
  onnxOpTypeNames _ = ["Flatten", "Reshape"]

instance OnnxLoadableActivation Reshape where
  activationLayer = Reshape