{-# LANGUAGE ForeignFunctionInterface #-}

module Numeric.LinearAlgebra.HArpack
       ( eigs
       , eigsN
       , eigsS
       , linearOperator
       , SparseMatrix (..)
       , Problem (..)
       , SWhich (..)
       , NWhich (..)
       , Result (..)
       , Dimension
       , HowMany
       , Tolerance
       , Iterations ) where

import           Control.Monad
import           Control.Monad.Loops
import           Data.Complex
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           Data.Vector.Storable.Mutable (IOVector)
import qualified Data.Vector.Storable.Mutable as MV
import           Foreign
import           Foreign.C.String
import           Foreign.C.Types
import           Foreign.ForeignPtr.Unsafe    as FPUnsafe
import           System.IO.Unsafe

-- | Functions of type 'LinearOperator' represent a matrix vector product.
-- The values in the second vector should be overwritten by operator applied to
-- the first.

data SparseMatrix =
  SparseMatrix { dim ::Int
               , indexes :: [((Int,Int),Double)] }
  deriving Show

data LinearOperator =
  LinearOperator
    { applyOperator :: IOVector CDouble -> IOVector CDouble -> IO () }

linearOperator :: [((Int,Int),Double)] -> LinearOperator
linearOperator l = LinearOperator $ \x y ->
  do MV.set y 0
     forM_ l $ \((i,j),v) -> do
       v' <- (* realToFrac v) <$> MV.read x j
       MV.modify y (+ v') i

--------------------------------------------------------------------------------
-- | Which eigenvalues to calculate:
-- - The first character refers to 'S'ymmetric or 'N'on-Symmetric;
-- - The second refers to 'L'argest or 'S'mallest;
-- - The third refers to 'M'agnitude, 'A'lgebraic, 'R'eal, 'I'mmaginary;
-- - The option 'BE' means 'BE'tween, for central eigenvalues.

data SWhich = SLM | SSM | SLA | SSA | SBE

instance Show SWhich where
  show SLM = "LM"
  show SSM = "SM"
  show SLA = "LA"
  show SSA = "SA"
  show SBE = "BE"

----------

data NWhich = NLM | NSM | NLR | NSR | NLI | NSI

instance Show NWhich where
  show NLM = "LM"
  show NSM = "SM"
  show NLR = "LR"
  show NSR = "SR"
  show NLI = "LI"
  show NSI = "SI"

----------

data Problem = Symmetric SWhich | NonSymmetric NWhich
     deriving (Show)

data Eigenpairs = Error String
                | EReal [(Double,Vector Double)]
                | EComplex [(Complex Double, Vector (Complex Double))]

--------------------------------------------------------------------------------

-- | Encapsulates the arrays and data needed to invoke _naupd routines in ARPACK

data ArpackSetup =
  ArpackSetup { aIDO    :: Ptr CInt
              , aBMAT   :: Ptr CChar
              , aN      :: Ptr CInt
              , aWHICH  :: Ptr CChar
              , aNEV    :: Ptr CInt
              , aTOL    :: Ptr CDouble
              , aRESID  :: Ptr CDouble
              , aNCV    :: Ptr CInt
              , aV      :: Ptr CDouble
              , aLDV    :: Ptr CInt
              , aIPARAM :: Ptr CInt
              , aIPNTR  :: Ptr CInt
              , aWORKD  :: Ptr CDouble
              , aWORKL  :: Ptr CDouble
              , aLWORKL :: Ptr CInt
              , aINFO   :: Ptr CInt }
  deriving (Show)

type Dimension = Int
type HowMany = Int
type Tolerance = Double
type Iterations = Int

-- | The 'setupArpack' function allocates the arrays and data
-- needed by ARPACK to solve an eigenvalue problem

setupArpack :: Problem
            -> Dimension
            -> HowMany
            -> Tolerance
            -> Iterations
            -> IO ArpackSetup

setupArpack (NonSymmetric which') n nev tol maxIter =
  do let ncv = min n (2 * nev + 1)
         lWorkl = 3 * ncv * ncv + 6 * ncv
         ldv = n

     idoPtr <- ptrInt 0
     nPtr   <- ptrInt n
     nevPtr <- ptrInt nev
     ncvPtr <- ptrInt ncv
     ldvPtr <- ptrInt ldv

     bmatPtr  <- ptrCharArray "I"
     whichPtr <- ptrCharArray (show which')

     tolPtr   <- ptrDouble tol

     residPtr <- mallocDoubleArray n
     vPtr     <- mallocDoubleArray (n * ncv)
     workdPtr <- mallocDoubleArray (3 * n)
     worklPtr <- mallocDoubleArray lWorkl

     iParamPtr <- mallocIntArray 11

     pokeElemOff iParamPtr 0 1
     pokeElemOff iParamPtr 2 (fromIntegral maxIter)
     pokeElemOff iParamPtr 3 1
     pokeElemOff iParamPtr 6 1

     iPntrPtr  <- mallocIntArray 14

     lworklPtr <- ptrInt lWorkl
     infoPtr   <- ptrInt 0

     setup <- return $
              ArpackSetup idoPtr bmatPtr nPtr whichPtr nevPtr tolPtr
                          residPtr ncvPtr vPtr ldvPtr iParamPtr iPntrPtr
                          workdPtr worklPtr lworklPtr infoPtr

     dnaupd setup
     return setup

setupArpack (Symmetric which) n nev tol maxIter =
  do let ncv = min n (2 * nev + 1)
         lWorkl = ncv * ncv + 8 * ncv
         ldv = n

     idoPtr <- ptrInt 0
     nPtr   <- ptrInt n
     nevPtr <- ptrInt nev
     ncvPtr <- ptrInt ncv
     ldvPtr <- ptrInt ldv

     bmatPtr  <- ptrCharArray "I"
     whichPtr <- ptrCharArray (show which)

     tolPtr   <- ptrDouble tol

     residPtr <- mallocDoubleArray n
     vPtr     <- mallocDoubleArray (n * ncv)
     workdPtr <- mallocDoubleArray (3 * n)
     worklPtr <- mallocDoubleArray lWorkl

     iParamPtr <- mallocIntArray 11

     pokeElemOff iParamPtr 0 1
     pokeElemOff iParamPtr 2 (fromIntegral maxIter)
     pokeElemOff iParamPtr 3 1
     pokeElemOff iParamPtr 6 1

     iPntrPtr  <- mallocIntArray 11

     lworklPtr <- ptrInt lWorkl
     infoPtr   <- ptrInt 0

     setup <- return $
              ArpackSetup idoPtr bmatPtr nPtr whichPtr nevPtr tolPtr
                          residPtr ncvPtr vPtr ldvPtr iParamPtr iPntrPtr
                          workdPtr worklPtr lworklPtr infoPtr

     dsaupd setup

     return setup

-- |Invoke ARPACK's iteration through the reverse communication interface.
iterateArpack :: Problem -> LinearOperator -> ArpackSetup -> IO ()
iterateArpack (NonSymmetric _) f setup = do
  workdXElem <- peekElemOff (aIPNTR setup) 0
  workdYElem <- peekElemOff (aIPNTR setup) 1

  n <- peek (aN setup)

  xPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdXElem - 1)
  yPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdYElem - 1)

  let x = MV.unsafeFromForeignPtr xPtr 0 (fromIntegral n)
      y = MV.unsafeFromForeignPtr yPtr 0 (fromIntegral n)

  applyOperator f x y

  dnaupd setup

iterateArpack (Symmetric _) f setup = do
  workdXElem <- peekElemOff (aIPNTR setup) 0
  workdYElem <- peekElemOff (aIPNTR setup) 1

  n <- peek (aN setup)

  xPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdXElem - 1)
  yPtr <- newForeignPtr_ $ advancePtr (aWORKD setup) (fromIntegral workdYElem - 1)

  let x = MV.unsafeFromForeignPtr xPtr 0 (fromIntegral n)
      y = MV.unsafeFromForeignPtr yPtr 0 (fromIntegral n)

  applyOperator f x y

  dsaupd setup

arpack :: LinearOperator -> Problem -> Dimension -> HowMany -> Tolerance ->
          Iterations -> IO ArpackSetup
arpack op problem n nev tol mxItr = do
  ar <- setupArpack problem n nev tol mxItr
  iterateArpack problem op ar
  whileM_ ((== 1) <$> peek (aIDO ar)) (iterateArpack problem op ar)
  return ar

--------------------------------------------------------------------------------

data D = DReal (IOVector CDouble)
       | DComplex (IOVector CDouble) (IOVector CDouble)

data Sigma = SReal (Ptr CDouble)
           | SComplex (Ptr CDouble) (Ptr CDouble)

data ArpackResults =
  ArpackResults { eRVEC    :: Ptr CInt
                , eHOWMNY  :: Ptr CChar
                , eSELECT  :: Ptr CInt
                , eD       :: D
                , eZ       :: IOVector CDouble
                , eLDZ     :: Ptr CInt
                , eSIGMA   :: Sigma
                , eWORKEV  :: Maybe (Ptr CDouble)
                , eSetup   :: ArpackSetup }

errCheck :: ArpackResults -> IO Bool
errCheck ar = do
  info <- (peek . aINFO . eSetup) ar
  print info
  return (info < 0)

parseArpackOutput :: Problem -> ArpackSetup -> IO ArpackResults
parseArpackOutput (NonSymmetric _) setup = do
  ncv <- peek (aNCV setup)
  nev <- peek (aNEV setup)
  n <- peek (aN setup)

  rvecPtr <- ptrInt 1
  howmnyPtr <- ptrCharArray "A"
  drVec <- newVectorDouble (nev + 1)
  diVec <- newVectorDouble (nev + 1)
  selectPtr <- mallocIntArray (fromIntegral ncv)
  zVec <- newVectorDouble (n*(nev+1))
  ldzPtr <- ptrInt (fromIntegral n)
  sigmarPtr <- mallocDouble
  sigmaiPtr <- mallocDouble
  workevPtr <- mallocDoubleArray (fromIntegral (3*ncv))

  let (drfPtr,_,_) = MV.unsafeToForeignPtr drVec
      drPtr = FPUnsafe.unsafeForeignPtrToPtr drfPtr
      (difPtr,_,_) = MV.unsafeToForeignPtr diVec
      diPtr = FPUnsafe.unsafeForeignPtrToPtr difPtr
      (zfPtr,_,_) = MV.unsafeToForeignPtr zVec
      zPtr = FPUnsafe.unsafeForeignPtrToPtr zfPtr

  c_dneupd rvecPtr howmnyPtr selectPtr drPtr diPtr zPtr ldzPtr sigmarPtr
    sigmaiPtr workevPtr (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
    (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
    (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup) (aLWORKL setup)
    (aINFO setup)

  return $ ArpackResults rvecPtr howmnyPtr selectPtr (DComplex drVec diVec) zVec
               ldzPtr (SComplex sigmarPtr sigmaiPtr) (Just workevPtr) setup

parseArpackOutput (Symmetric _) setup = do
  ncv <- peek (aNCV setup)
  nev <- peek (aNEV setup)
  n <- peek (aN setup)

  rvecPtr <- ptrInt 1
  howmnyPtr <- ptrCharArray "A"
  dVec <- newVectorDouble nev
  selectPtr <- mallocIntArray (fromIntegral ncv)
  zVec <- newVectorDouble (n*(nev+1))
  ldzPtr <- ptrInt (fromIntegral n)
  sigmaPtr <- mallocDouble

  let (drfPtr,_,_) = MV.unsafeToForeignPtr dVec
      dPtr = FPUnsafe.unsafeForeignPtrToPtr drfPtr
      (zfPtr,_,_) = MV.unsafeToForeignPtr zVec
      zPtr = FPUnsafe.unsafeForeignPtrToPtr zfPtr

  c_dseupd rvecPtr howmnyPtr selectPtr dPtr zPtr ldzPtr
    sigmaPtr (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
    (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
    (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup) (aLWORKL setup)
    (aINFO setup)

  return $ ArpackResults rvecPtr howmnyPtr selectPtr (DReal dVec) zVec
               ldzPtr (SReal sigmaPtr) Nothing setup

getComplexEigenvalue :: ArpackResults -> Int -> IO (Complex Double)
getComplexEigenvalue results index = do
  let DComplex reEigs imEigs = eD results
  re <- MV.read reEigs index
  im <- MV.read imEigs index
  return (realToFrac re :+ realToFrac im)

getRealEigenvector :: ArpackResults -> Int -> IO (Vector (Complex Double))
getRealEigenvector results index = do
  n <- (peek . aN . eSetup) results
  ldz <- peek $ eLDZ results
  let z = eZ results
  evec <- V.freeze $ MV.slice (index * fromIntegral ldz) (fromIntegral n) z
  return $ V.map ((:+0) . realToFrac) evec

getComplexEigenvector
  :: ArpackResults -> Int -> Bool -> IO (Vector (Complex Double))
getComplexEigenvector results index conjugate = do
  n <- (peek . aN . eSetup) results
  ldz <- peek $ eLDZ results
  let z = eZ results
  reEvec <- V.freeze $ MV.slice (index * fromIntegral ldz) (fromIntegral n) z
  imEvec <- V.freeze $ MV.slice ((index+1) * fromIntegral ldz) (fromIntegral n) z
  if conjugate
     then return $
          V.zipWith (:+) (V.map realToFrac reEvec)
                         (V.map (negate . realToFrac) imEvec)
     else return $
          V.zipWith (:+) (V.map realToFrac reEvec)
                         (V.map realToFrac imEvec)

getEigenpair
  :: ArpackResults -> Int -> IO (Complex Double, V.Vector (Complex Double))
getEigenpair results index = do
  eig <- getComplexEigenvalue results index
  let evec = case compare (imagPart eig) 0 of
        EQ -> getRealEigenvector results index
        LT -> getComplexEigenvector results (index - 1) True
        GT -> getComplexEigenvector results index False
  evec' <- evec
  return (eig,evec')

-- | The function 'eigs' computes a few eigenvalues of a real valued
-- unsymmetric matrix by invoking the ARPACK library. First element is true if
-- ARPACK converged. If converged the second element is a list of eigenpairs.

eigsN :: LinearOperator
      -> NWhich
      -> Dimension
      -> HowMany
      -> Tolerance
      -> Iterations
      -> Maybe [(Complex Double, Vector (Complex Double))]
eigsN f which' n' nev' tol' iters = unsafePerformIO $ do
  setup <- arpack f (NonSymmetric which') n' nev' tol' iters
  results <- parseArpackOutput (NonSymmetric which') setup
  err <- errCheck results
  if err
    then return Nothing
    else do
      pairs <- mapM (getEigenpair results) [0..(nev'-1)]
      return (Just pairs)

data Result = RError
            | RReal [(Double,Vector Double)]
            | RComplex [(Complex Double, Vector (Complex Double))]
              deriving Show

eigs :: SparseMatrix
     -> Problem
     -> HowMany
     -> Result
eigs m (NonSymmetric w) h =
  case eigsN (linearOperator (indexes m)) w (dim m) h (1e-10) 6000 of
  Nothing -> RError
  Just x -> RComplex x
eigs m (Symmetric w) h =
  case eigsS (linearOperator (indexes m)) w (dim m) h (1e-10) 6000 of
  Nothing -> RError
  Just x -> RReal x

--------------------------------------------------------------------------------

getSymmetricEigenvalue :: ArpackResults -> Int -> IO Double
getSymmetricEigenvalue results index = do
  let DReal eigV = eD results
  eig <- MV.read eigV index
  return (realToFrac eig)

getSymmetricEigenvector :: ArpackResults -> Int -> IO (Vector Double)
getSymmetricEigenvector results index = do
  n <- fromIntegral <$> (peek . aN . eSetup) results
  ldz <- peek $ eLDZ results
  let z = eZ results
  evec <- V.freeze $ MV.slice (index * fromIntegral ldz) n z
  return $ V.map realToFrac evec

getSymmetricEigenpair :: ArpackResults -> Int -> IO (Double, Vector Double)
getSymmetricEigenpair results index = do
  eig <- getSymmetricEigenvalue results index
  evec <- getSymmetricEigenvector results index
  return (eig,evec)

eigsS :: LinearOperator
      -> SWhich
      -> Dimension
      -> HowMany
      -> Tolerance
      -> Iterations
      -> Maybe [(Double,Vector Double)]
eigsS f which n nev tol' iters = unsafePerformIO $ do
  setup <- arpack f (Symmetric which) n nev tol' iters
  results <- parseArpackOutput (Symmetric which) setup
  err <- errCheck results
  if err
    then return Nothing
    else do
      pairs <- mapM (getSymmetricEigenpair results) [0..(nev-1)]
      return (Just pairs)

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

mallocInt :: IO (Ptr CInt)
mallocInt = malloc

ptrInt :: Int -> IO (Ptr CInt)
ptrInt k =
  do ptr <- malloc
     poke ptr (fromIntegral k)
     return ptr

mallocIntArray :: Int -> IO (Ptr CInt)
mallocIntArray = mallocArray

----------

mallocDouble :: IO (Ptr CDouble)
mallocDouble = malloc

ptrDouble :: Double -> IO (Ptr CDouble)
ptrDouble x =
  do ptr <- malloc
     poke ptr (realToFrac x)
     return ptr

mallocDoubleArray :: Int -> IO (Ptr CDouble)
mallocDoubleArray = mallocArray

newVectorDouble :: CInt -> IO (MV.IOVector CDouble)
newVectorDouble = MV.new . fromIntegral

----------

mallocChar :: IO (Ptr CChar)
mallocChar = malloc

mallocCharArray :: Int -> IO (Ptr CChar)
mallocCharArray = mallocArray

ptrCharArray :: String -> IO (Ptr CChar)
ptrCharArray s =
  do ptr <- mallocCharArray (length s + 1)
     pokeArray ptr (map castCharToCChar s)
     return ptr

--------------------------------------------------------------------------------
--------------------------------------------------------------------------------

foreign import ccall unsafe "arpack.h dnaupd_"
  c_dnaupd :: Ptr CInt    -- ^ reverse communication flag, to be started with 0
           -> Ptr CChar   -- ^ 'I' for standard problem, 'G' for generalized
           -> Ptr CInt    -- ^ dimension of the eigenproblem
           -> Ptr CChar   -- ^ which eigenpairs to compute
           -> Ptr CInt    -- ^ N : number of eigenpairs to compute
           -> Ptr CDouble -- ^ tolerance
           -> Ptr CDouble -- ^ residual vector
           -> Ptr CInt    -- ^ NCV : number of columns of V matrix
           -> Ptr CDouble -- ^ N x NCV output array
           -> Ptr CInt    -- ^ leading dimension of V
           -> Ptr CInt    -- ^ array for internal parameters, length 11
           -> Ptr CInt    -- ^ pointer of starting locations, length 14
           -> Ptr CDouble -- ^ workd : work array
           -> Ptr CDouble -- ^ workl : work array
           -> Ptr CInt    -- ^ length of workl
           -> Ptr CInt    -- ^ info
           -> IO ()

dnaupd :: ArpackSetup -> IO ()
dnaupd setup =
  c_dnaupd (aIDO setup) (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
           (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
           (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup)
           (aLWORKL setup) (aINFO setup)

foreign import ccall unsafe "arpack.h dneupd_"
  c_dneupd :: Ptr CInt    -- ^ whether to calculate eigenvectors
           -> Ptr CChar   -- ^ whether to calculate Ritz or Schur vectors
           -> Ptr CInt    -- ^ selects the vectors to be computed
           -> Ptr CDouble -- ^ real part of eigenvalues
           -> Ptr CDouble -- ^ imaginary part of eigenvalues
           -> Ptr CDouble -- ^ approximate eigenvectors
           -> Ptr CInt    -- ^ leading dimension of former array
           -> Ptr CDouble -- ^ real part of the shift
           -> Ptr CDouble -- ^ imaginary part of the shift
           -> Ptr CDouble -- ^ work array
           -> Ptr CChar -> Ptr CInt -> Ptr CChar -> Ptr CInt -> Ptr CDouble
           -> Ptr CDouble -> Ptr CInt -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> Ptr CInt -> Ptr CDouble -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> IO ()

--------------------------------------------------------------------------------

foreign import ccall unsafe "arpack.h dsaupd_"
  c_dsaupd :: Ptr CInt    -- ^ reverse communication flag, to be started with 0
           -> Ptr CChar   -- ^ 'I' for standard problem, 'G' for generalized
           -> Ptr CInt    -- ^ dimension of the eigenproblem
           -> Ptr CChar   -- ^ which eigenpairs to compute
           -> Ptr CInt    -- ^ N : number of eigenpairs to compute
           -> Ptr CDouble -- ^ tolerance
           -> Ptr CDouble -- ^ residual vector
           -> Ptr CInt    -- ^ NCV : number of columns of V matrix
           -> Ptr CDouble -- ^ N x NCV output array
           -> Ptr CInt    -- ^ leading dimension of V
           -> Ptr CInt    -- ^ array for internal parameters, length 11
           -> Ptr CInt    -- ^ pointer of starting locations, length 14
           -> Ptr CDouble -- ^ workd : work array
           -> Ptr CDouble -- ^ workl : work array
           -> Ptr CInt    -- ^ length of workl
           -> Ptr CInt    -- ^ info
           -> IO ()

dsaupd :: ArpackSetup -> IO ()
dsaupd setup =
  c_dsaupd (aIDO setup) (aBMAT setup) (aN setup) (aWHICH setup) (aNEV setup)
           (aTOL setup) (aRESID setup) (aNCV setup) (aV setup) (aLDV setup)
           (aIPARAM setup) (aIPNTR setup) (aWORKD setup) (aWORKL setup)
           (aLWORKL setup) (aINFO setup)

foreign import ccall unsafe "arpack.h dseupd_"
  c_dseupd :: Ptr CInt    -- ^ whether to calculate eigenvectors
           -> Ptr CChar   -- ^ whether to calculate Ritz or Schur vectors
           -> Ptr CInt    -- ^ selects the vectors to be computed
           -> Ptr CDouble -- ^ eigenvalues
           -> Ptr CDouble -- ^ approximate eigenvectors
           -> Ptr CInt    -- ^ leading dimension of former array
           -> Ptr CDouble -- ^ shift
           -> Ptr CChar -> Ptr CInt -> Ptr CChar -> Ptr CInt -> Ptr CDouble
           -> Ptr CDouble -> Ptr CInt -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> Ptr CInt -> Ptr CDouble -> Ptr CDouble -> Ptr CInt -> Ptr CInt
           -> IO ()


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

module Data.CDF
  ( Centile(..)
  , renderCentile
  , briefCentiles
  , stdCentiles
  , nEquicentiles
  , CDFError (..)
  , CDF(..)
  , CDFIx (..)
  , KnownCDF (..)
  , liftCDFVal
  , unliftCDFVal
  , centilesCDF
  , filterCDF
  , zeroCDF
  , projectCDF
  , projectCDF'
  , indexCDF
  , DirectCDF
  , cdf
  , mapToCDF
  , Divisible (..)
  , Combine (..)
  , stdCombine1
  , stdCombine2
  , CDF2
  , collapseCDFs
  , cdf2OfCDFs
  --
  , module Data.SOP.Strict
  ) where

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

import Data.Aeson (FromJSON(..), ToJSON(..))
import Data.SOP.Strict
import Data.Time.Clock (NominalDiffTime)
import Data.Vector qualified as Vec
import Statistics.Sample qualified as Stat
import Text.Printf (printf)

import Ouroboros.Consensus.Util.Time (secondsToNominalDiffTime)

-- | Centile specifier: a fractional in range of [0; 1].
newtype Centile =
  Centile { unCentile :: Double }
  deriving (Eq, Generic, FromJSON, ToJSON, Show)
  deriving anyclass NFData

renderCentile :: Int -> Centile -> String
renderCentile width = \case
  Centile x    -> printf ("%0."<>show (width-2)<>"f") x

briefCentiles :: [Centile]
briefCentiles =
  [ Centile 0.5, Centile 0.9, Centile 1.0 ]

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

-- | Given a N-large population, produce centiles for each element, except for min and max.
--   We don't need min and max, because CDF already has range.
nEquicentiles :: Int -> [Centile]
nEquicentiles n =
  if reindices == indices
  then Centile <$> centiles
  else error $ printf "centilesForN:  reindices for %d: %s, indices: %s" n (show reindices) (show indices)
 where
   reindices = centiles <&> runCentile n
   centiles  = [ step * (fromIntegral i + 0.5) | i <- indices ]
   indices   = if n > 2
               then [1 .. n - 2] -- ignore first and last indices, standing for min and max.
               else []
   step :: Double
   step = 1.0 / fromIntegral n

-- | Given a centile of N-large population, produce index of the population element referred by centile.
{-# INLINE runCentile #-}
runCentile :: Int -> Double -> Int
runCentile n centile = floor (fromIntegral n * centile)
                       & min (n - 1)

{-# INLINE vecCentile #-}
vecCentile :: Vec.Vector a -> Int -> Centile -> a
vecCentile vec n (Centile c) = vec Vec.! runCentile n c

--
-- * Parametric CDF (cumulative distribution function)
--
data CDF p a =
  CDF
  { cdfSize      :: Int
  , cdfAverage   :: Double
  , cdfStddev    :: Double
  , cdfRange     :: (a, a)
  , cdfSamples  :: [(Centile, p a)]
  }
  deriving (Eq, Functor, Generic, Show)
  deriving anyclass NFData

instance (FromJSON (p a), FromJSON a) => FromJSON (CDF p a)
instance (  ToJSON (p a),   ToJSON a) => ToJSON   (CDF p a)

-- * Singletons
--
data CDFIx p where
  CDFI :: CDFIx I
  CDF2 :: CDFIx (CDF I)

class KnownCDF a where
  cdfIx :: CDFIx a

instance KnownCDF      I  where cdfIx = CDFI
instance KnownCDF (CDF I) where cdfIx = CDF2

centilesCDF :: CDF p a -> [Centile]
centilesCDF = fmap fst . cdfSamples

zeroCDF :: (Num a) => CDF p a
zeroCDF =
  CDF
  { cdfSize    = 0
  , cdfAverage = 0
  , cdfStddev  = 0
  , cdfRange   = (0, 0)
  , cdfSamples = mempty
  }

filterCDF :: ((Centile, p a) -> Bool) -> CDF p a -> CDF p a
filterCDF f d =
  d { cdfSamples = cdfSamples d & filter f }

indexCDF :: Int -> CDF p a -> p a
indexCDF i d = snd $ cdfSamples d !! i

projectCDF :: Centile -> CDF p a -> Maybe (p a)
projectCDF p = fmap snd . find ((== p) . fst) . cdfSamples

projectCDF' :: String -> Centile -> CDF p a -> p a
projectCDF' desc p =
  maybe (error er) snd . find ((== p) . fst) . cdfSamples
 where
   er = printf "Missing centile %f in %s" (show $ unCentile p) desc

--
-- * Trivial instantiation: samples are value points
--
type DirectCDF a = CDF I a

liftCDFVal :: Real a => a -> CDFIx p -> p a
liftCDFVal x = \case
  CDFI -> I x
  CDF2 -> CDF { cdfSize    = 1
              , cdfAverage = fromRational (toRational x)
              , cdfStddev  = 0
              , cdfRange   = (x, x)
              , cdfSamples = []
              , .. }

unliftCDFVal :: Divisible a => CDFIx p -> p a -> a
unliftCDFVal CDFI (I x) =x
unliftCDFVal CDF2 CDF{cdfRange = (mi, ma)} = (mi + ma) `divide` 2

cdf :: (Real a, KnownCDF p) => [Centile] -> [a] -> CDF p a
cdf centiles (sort -> sorted) =
  CDF
  { cdfSize        = size
  , cdfAverage     = Stat.mean   doubleVec
  , cdfStddev      = Stat.stdDev doubleVec
  , cdfRange       = (mini, maxi)
  , cdfSamples =
    (    (Centile   0, liftCDFVal mini ix) :) .
    (<> [(Centile 1.0, liftCDFVal maxi ix) ]) $
    centiles <&>
      \spec ->
        let sample = if size == 0 then 0
                     else vecCentile vec size spec
        in (,) spec (liftCDFVal sample ix)
  }
 where ix          = cdfIx
       vec         = Vec.fromList sorted
       size        = length vec
       doubleVec   = fromRational . toRational <$> vec
       (,) mini maxi =
         if size == 0
         then (0,           0)
         else (vec Vec.! 0, Vec.last vec)

mapToCDF :: Real a => (b -> a) -> [Centile] -> [b] -> DirectCDF a
mapToCDF f pspecs xs = cdf pspecs (f <$> xs)

type CDF2 a = CDF (CDF I) a

data CDFError
  = CDFIncoherentSamplingLengths  [Int]
  | CDFIncoherentSamplingCentiles [[Centile]]
  | CDFEmptyDataset
  deriving Show

-- | Avoiding `Fractional`
class Real a => Divisible a where
  divide :: a -> Double -> a

instance Divisible Double where
  divide = (/)

instance Divisible Int where
  divide x by = round $ fromIntegral x / by

instance Divisible Integer where
  divide x by = round $ fromIntegral x / by

instance Divisible Word32 where
  divide x by = round $ fromIntegral x / by

instance Divisible Word64 where
  divide x by = round $ fromIntegral x / by

instance Divisible NominalDiffTime where
  divide x by = x / secondsToNominalDiffTime by

-- * Combining population stats
data Combine p a
  = Combine
    { cWeightedAverages :: !([(Int, Double)] -> Double)
    , cStddevs          :: !([Double] -> Double)
    , cRanges           :: !([(a, a)] -> (a, a))
    , cWeightedSamples  :: !([(Int, a)] -> a)
    , cCDF              :: !([p a] -> Either CDFError (CDF I a))
    }

stdCombine1 :: forall a. (Divisible a) => [Centile] -> Combine I a
stdCombine1 cs =
  Combine
  { cWeightedAverages = weightedAverage
  , cRanges           = outerRange
  , cStddevs          = maximum          -- it's an approximation
  , cWeightedSamples  = weightedAverage
  , cCDF              = Right . cdf cs . fmap unI
  }
  where
    weightedAverage :: forall b. (Divisible b) => [(Int, b)] -> b
    weightedAverage xs = (`divide` (fromIntegral . sum $ fst <$> xs)) . sum $
                          xs <&> \(size, avg) -> fromIntegral size * avg
    outerRange      xs = (,) (minimum $ fst <$> xs)
                             (maximum $ snd <$> xs)

stdCombine2 :: Divisible a => [Centile] -> Combine (CDF I) a
stdCombine2 cs =
  let c@Combine{..} = stdCombine1 cs in
  Combine
  { cCDF = collapseCDFs c
  , ..
  }

-- | Collapse:  Given a ([Value] -> CDF I) function, and a list of (CDF I), produce a (CDF I).
--
collapseCDFs :: forall a. Combine I a -> [CDF I a] -> Either CDFError (CDF I a)
collapseCDFs _ [] = Left CDFEmptyDataset
collapseCDFs Combine{..} xs = do
  unless (all (head lengths ==) lengths) $
    Left $ CDFIncoherentSamplingLengths lengths
  unless (null incoherent) $
    Left $ CDFIncoherentSamplingCentiles (fmap fst <$> incoherent)
  pure CDF
    { cdfSize    = sum sizes
    , cdfAverage = xs <&> cdfAverage & cWeightedAverages . zip sizes
    , cdfRange   = xs <&> cdfRange   & cRanges
    , cdfStddev  = xs <&> cdfStddev  & cStddevs
    , cdfSamples = coherent <&>
                   fmap (I . cWeightedSamples . zip sizes . fmap unI)
    }
 where
   sizes   = xs <&> cdfSize
   samples = xs <&> cdfSamples
   lengths = length <$> samples

   centileOrdered :: [[(Centile, I a)]] -- Each sublist must (checked) have the same Centile.
   centileOrdered = transpose samples

   coherent :: [(Centile, [I a])]
   (incoherent, coherent) = partitionEithers $ centileOrdered <&>
     \case
       [] -> error "cdfOfCDFs:  empty list of centiles, hands down."
       xxs@((c, _):(fmap fst -> cs)) -> if any (/= c) cs
                                        then Left xxs
                                        else Right (c, snd <$> xxs)

-- | Polymorphic, but practically speaking, intended for either:
--    1. given a ([I]     -> CDF I) function, and a list of (CDF I),       produce a CDF (CDF I), or
--    2. given a ([CDF I] -> CDF I) function, and a list of (CDF (CDF I)), produce a CDF (CDF I)
cdf2OfCDFs :: forall a p. Combine p a -> [CDF p a] -> Either CDFError (CDF (CDF I) a)
cdf2OfCDFs _ [] = Left CDFEmptyDataset
cdf2OfCDFs Combine{..} xs = do
  unless (all (head lengths ==) lengths) $
    Left $ CDFIncoherentSamplingLengths lengths
  unless (null incoherent) $
    Left $ CDFIncoherentSamplingCentiles (fmap fst <$> incoherent)

  cdfSamples <- mapM sequence -- ..to  Either CDFError [(Centile, CDF I a)]
                  (coherent <&> fmap cCDF :: [(Centile, Either CDFError (CDF I a))])
  pure CDF
    { cdfSize    = sum sizes
    , cdfAverage = xs <&> cdfAverage & cWeightedAverages . zip sizes
    , cdfRange   = xs <&> cdfRange   & cRanges
    , cdfStddev  = xs <&> cdfStddev  & cStddevs
    , cdfSamples = cdfSamples
    }
 where
   sizes   = xs <&> cdfSize
   samples = xs <&> cdfSamples
   lengths = length <$> samples

   centileOrdered :: [[(Centile, p a)]]
   centileOrdered = transpose samples

   (incoherent, coherent) = partitionEithers $ centileOrdered <&>
     \case
       [] -> error "cdfOfCDFs:  empty list of centiles, hands down."
       xxs@((c, _):(fmap fst -> cs)) -> if any (/= c) cs
                                        then Left xxs
                                        else Right (c, snd <$> xxs)


module Mnist
  ( MnistData(..)
  , MnistLabels(..)
  , MnistImages(..)
  , saveMnist
  , readMnist
  , zipMnist
  ) where

import qualified Codec.Compression.GZip     as GZip
import           Control.Arrow
import           Control.DeepSeq
import           Control.Monad.State
import           Data.Binary                as B
import           Data.Binary.Get.Internal   as BGI
import qualified Data.ByteString            as BSS
import qualified Data.ByteString.Lazy       as BS
import qualified Data.ByteString.Lazy.UTF8  as UTF8
import           Data.List.Split
import           Debug.Trace
import           Network.HTTP.Client
import           Numeric.LinearAlgebra.Data
import           System.Directory
import           System.FilePath.Posix
import           Util

newtype MnistLabels =
  MnistLabels [Word8]
  deriving (Show)

newtype MnistImages =
  MnistImages [Matrix Z]
  deriving (Show)

newtype MnistData =
  MnistData [(Word8, Matrix Z)]
  deriving (Show)

zipMnist :: MnistLabels -> MnistImages -> MnistData
zipMnist (MnistLabels labels) (MnistImages images) =
  list `deepseq` MnistData list
  where
    list = labels `zip` images

markMinistLabels = 2049 :: Int

markMnistImages = 2051 :: Int

putAsWord32 :: Integral a => a -> Put
putAsWord32 a = B.put (fromIntegral a :: Word32)

getAsIntegral :: Integral a => Get a
getAsIntegral = fromIntegral <$> (B.get :: Get Word32)

instance B.Binary MnistLabels where
  put (MnistLabels labels) = do
    B.put markMinistLabels
    putAsWord32 $ length labels
    B.put $ BS.pack labels
  get = do
    mark <- getAsIntegral
    guard $ mark == markMinistLabels
    size <- getAsIntegral
    let total = trace ("Reading labels: " ++ show size) size
    labels <- BGI.readN total BSS.unpack
    return $ MnistLabels $ labels

instance B.Binary MnistImages where
  put (MnistImages (images@(hm:_))) = do
    putAsWord32 markMnistImages
    putAsWord32 $ length images
    putAsWord32 $ rows hm
    putAsWord32 $ cols hm
    B.put $ (BS.pack . map fromIntegral . concat . concat . map toLists) images
  get = do
    mark <- getAsIntegral
    guard $ mark == markMnistImages
    size <- getAsIntegral
    nRows <- getAsIntegral
    nCols <- getAsIntegral
    let total =
          trace
            ("Reading images: " ++
             show size ++ " of " ++ show nRows ++ "x" ++ show nCols)
            size *
          nRows *
          nCols
    zs <- map fromIntegral <$> BGI.readN total BSS.unpack
    let images = map (nRows >< nCols) $ chunksOf (nRows * nCols) zs
    return $ MnistImages images

download :: Manager -> String -> IO BS.ByteString
download manager url = do
  request <- parseRequest url
  timestamp $ "Downloading " ++ url
  response <- httpLbs request manager
  return $ responseBody response

saveMnist :: FilePath -> String -> [String] -> Manager -> IO [FilePath]
saveMnist rootDir urlBase filenames manager = do
  isDir <- doesDirectoryExist rootDir
  _ <-
    if isDir
      then return ()
      else createDirectory rootDir
  mapM saveFile filenames
  where
    saveFile filename = do
      let url = urlBase ++ filename
      let file = rootDir </> filename
      timestamp $ "Checking file " ++ file
      e <- doesFileExist file
      if e
        then return ()
        else download manager url >>= (BS.writeFile file)
      return file

readMnist :: FilePath -> IO (Either MnistLabels MnistImages)
readMnist file = do
  timestamp $ "decoding " ++ file ++ " ..."
  bs <- GZip.decompress <$> BS.readFile file
  let a =
        left (const $ B.decode bs) $ right (\(_, _, a) -> a) $ B.decodeOrFail bs
  timestamp $ "decoded " ++ file
  return a


{-# LANGUAGE FlexibleInstances #-}

module Isumi.Math.AppxEq
  ( AppxEq(..)
  , (~/=)
  ) where

import           Data.Complex
import qualified Data.Vector.Unboxed as UV

class AppxEq a where
  (~=) :: a -> a -> Bool

infix 4 ~=

infix 4 ~/=
(~/=) :: AppxEq a => a -> a -> Bool
x ~/= y = not $ x ~= y

instance AppxEq Int where
  x ~= y = x == y

instance AppxEq Double where
  x ~= y = abs (x - y) < 0.000001

instance AppxEq a => AppxEq (Complex a) where
  (r :+ v) ~= (r2 :+ v2) = (r ~= r2) && (v ~= v2)

instance AppxEq a => AppxEq (Maybe a) where
  Just x ~= Just y = x ~= y
  Nothing ~= Nothing = True
  _ ~= _ = False

instance AppxEq a => AppxEq [a] where
  xs ~= ys = and $ zipWith (~=) xs ys

instance (UV.Unbox a, AppxEq a) => AppxEq (UV.Vector a) where
  xs ~= ys = UV.toList xs ~= UV.toList ys

instance (AppxEq a, AppxEq b) => AppxEq (a, b) where
  (x1, y1) ~= (x2, y2) = x1 ~= x2 && y1 ~= y2


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, DeriveFunctor, MultiParamTypeClasses, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : First attempt at approximate mean
-- | Creator: Xiao Ling
-- | Created: 11/29/2015
-- | see    : https://github.com/snoyberg/conduit
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module Play1129 where

import System.Random

import Data.List
import Data.Random
import Data.Conduit
import qualified Data.Conduit.List as Cl

import Data.Vector (Vector,fromList,cons,snoc,(!?))
import qualified Data.Vector       as V

import Control.Monad.Identity
import Control.Monad.State

import Test.QuickCheck ((==>), (.&&.))
import qualified Test.QuickCheck as T

import Core
import Statistics

{-----------------------------------------------------------------------------
   1. make some examples
   2. make
   3. see if cassava can be put into a conduit
   4. output Row type

   see : http://stackoverflow.com/questions/2376981/haskell-types-frustrating-a-simple-average-function
------------------------------------------------------------------------------}

-- * first figure out how to sample from a list with probability p
-- * it's a dice with 1/p faces

--type MaxSize      = Int
--type Store a      = (Vector a, MaxSize)

--toMedian :: Counter -> Sink a (StateT (Store a) RVar) (Vector a)
--toMedian c = do
--  mi <- await
--  case mi of
--    Nothing -> lift get >>= return . fst
--    Just i  -> do
--      (v,s) <- lift get
--      case V.length v < s of
--        True -> do
--          let v' = i `cons` v
--          put

{-----------------------------------------------------------------------------
    II. Approximate Median try ii
        too much unnnecessary stuff that adds no value
------------------------------------------------------------------------------}

-- * TODO: make this into an array since we'll be doing (!!) which is O(n) in list
-- * A reservoir is a list, its max allowed size, and its current size

-- * this deign is really bad, cannot gaurantee that currentSize is upToDate with length [a]
-- * dependent types would be good, or we can jsut get rid of current size
-- * or hide it in Core
data Reservoir a = R { run :: [a], size :: Int, currentSize :: Int }
  deriving (Show,Functor,Eq)

instance Num a => T.Arbitrary (Reservoir a) where

  -- * by construction, an arbitrary reservoir is always 1 item away from full
  arbitrary = do
    s <- T.arbitrary :: T.Gen (T.Positive Int)
    let s' = T.getPositive s
    let n = s' - 1   :: Int
    return $ R (replicate n 0) s' n

  shrink (R xs s n) = (\s' -> let n = max 0 (s' - 1) in (R (replicate n 0) s' n)) <$> T.shrink s

-- * constructor an empty resoivor of max size `s`
reservoir :: Int -> Reservoir a
reservoir s = R [] s 0

-- * check if reservoir is full
full :: Reservoir a -> Bool
full (R xs s n) = n >= s

-- * insert item, note the order
(>+) :: Reservoir a -> a -> Reservoir a
(>+) t@(R xs s n) x | full t     = t
                    | otherwise  = R (x:xs) s $ succ n

-- * remove the `i`th item from `r`
-- * if i > size then no item removed
(>-) :: Reservoir a -> Int -> Reservoir a
r@(R xs s n) >- i | i > n     = r
                  | otherwise = undefined      -- * really need to use array here

-- * sample an item `x` from the stream with uniform probability place in reservoir
-- * keep a conter `c` of elements seen
sampl' :: Counter -> Sink a (StateT (Reservoir a) RVar) [a]
sampl' i = do
  ma <- await
  case ma of
    Nothing -> lift get >>= \r -> return $ run r
    Just a  -> do
      lift $ store' i a
      sampl' $ succ i

store' :: Counter -> a -> StateT (Reservoir a) RVar ()
store' i a = do
  r <- get
  case full r of
    False -> put $ r >+ a
    _     -> return ()
      where p = (read . show $ size r :: Float)/i

-- * test m/2 - e*m < rank(y) < m/2 + e*m

-- * quickCheck Resoivor

-- * full resoivor is full
prop_fullRes :: Reservoir Int -> T.Property
prop_fullRes r = size r > 0 ==> full (r>+1) == True

---- * inserting into non-full resoivor resutls in additional element
prop_insertNotFull :: Reservoir Int -> T.Property
prop_insertNotFull r = size r > 0 ==> currentSize (r>+1) == currentSize r + 1

-- * inserting into full resoivor leaves it unchanged
prop_insertFull :: Reservoir Int -> T.Property
prop_insertFull r = size r > 0 ==> currentSize (r'>+1) == currentSize r'
  where r' = r >+ 1

props :: T.Property
props =  prop_fullRes .&&. prop_insertFull .&&. prop_insertFull

testReservoir :: IO ()
testReservoir = T.quickCheck props

{-----------------------------------------------------------------------------
    II. Approximate Median Naive
------------------------------------------------------------------------------}

-- * setting problem: can you make this streaming? assuming you know m?
-- * can you do running uniform distribution?

-- * Find approximate eps e-median of list `xs`
-- * Setting: list given before hand for some parameter t
-- * Use : runRVar (naive xs e d) StdRandom
medianNaive :: (Floating a, Ord a) => [a] -> Eps -> Delta -> RVar a
medianNaive xs e d | m <= t    = return $ median xs
             | otherwise = fmap median $ sampl'' xs m t
              where
                m = length xs
                t = round $ 7/(e^2) * log (2/d)

-- * use vector here due to (!!)
-- * sample `t` items from the list `[0..m-1]` *with replacement*
-- * Use : runRVar (sample m t) StdRandom
sampl'' :: [a] -> Int -> Int -> RVar [a]
sampl'' xs m t | m < 0 || t < 0   = return []
          | otherwise        = (\is -> [ xs !! i | i <- is]) <$> mis
            where mis = fmap sort . replicateM t $ uniform 0 (m-1)


module Main (main) where

import Autoencoder

import Control.Applicative
import Control.Monad
import Control.Monad.ST
import Data.Array.ST
import Data.Char
import Data.Classifier (Classifier)
import Data.Classifier.NaiveBayes (NaiveBayes)
import Data.Counter (Counter)
import Data.Maybe
import Data.Monoid
import Data.STRef
import Data.Text (Text)
import Data.Time.Clock
import Numeric.LinearAlgebra.HMatrix (Vector)
import Prelude
import System.Environment (getArgs)
import System.Exit (exitFailure)
import System.Random (randomR, getStdRandom, RandomGen)
import Text.Read (readMaybe)
import qualified AI.HNN.FF.Network as Neural
import qualified Data.Classifier as Classifier
import qualified Data.Classifier.NaiveBayes as NaiveBayes
import qualified Data.Counter as Counter
import qualified Data.Map as Map
import qualified Data.Text as Text
import qualified Numeric.LinearAlgebra.HMatrix as Vector
import qualified System.IO as IO

main :: IO ()
main = getArgs >>= \case
  ["bayes", filename] ->
    void $ evaluateBayes filename
  ["neural", filename, readMaybe -> Just trainTimes, readMaybe -> Just layers] ->
    neural filename trainTimes layers
  ["generate_autoencoder", filename] -> do
    IO.hSetBuffering IO.stdout IO.NoBuffering
    t <- getCurrentTime
    autoencoder filename >>= print
    getCurrentTime >>= print . (`diffUTCTime` t)
  ["apply_autoencoder", encoder, filename] -> do
    (v, (e, _)) <- read <$> readFile encoder :: IO (Counter Text, Autoencoder)
    applyAutoencoder v e filename 10 []
  ["everything", filename] -> do
    void $ evaluateBayes filename
    neural filename 10 []
    neural filename 100 []
    neural filename 1000 []
    neural filename 10 [10]
    neural filename 100 [10]
    neural filename 1000 [10]
    neural filename 10 [100]
    neural filename 100 [100]
    neural filename 1000 [100]
    neural filename 10000 [10]
    neural filename 10 [100, 50]
    neural filename 100 [100, 50]
    neural filename 1000 [100, 50]
  _ -> do
    putStrLn "Invalid arguments"
    exitFailure

applyAutoencoder :: Counter Text -> Encoder -> FilePath -> Int -> [Int] -> IO ()
applyAutoencoder vocab encoder path times layers = do
  shuffled <- readFile path >>= getStdRandom . shuffle . extractData
  let (train, test) = splitAt (length shuffled `div` 2) shuffled
  let trainVectors = map (\(x, y) -> (encode encoder x, y)) $ classifierToVector boolToVector vocab $ mconcat $ map rowToClassifier train
  let testVectors = map (\(x, y) -> (encode encoder x, y)) $ classifierToVector boolToVector vocab $ mconcat $ map rowToClassifier test
  case trainVectors of
    [] -> putStrLn "No data"
    (v, _) : _ -> do
      startTime <- getCurrentTime
      network <- Neural.createNetwork (Vector.size v) layers 1
      IO.hSetBuffering IO.stdout IO.NoBuffering
      void $ iterateM 9 (0 :: Int, network, 0) $ \(n, net, timeTakenEvaluating) -> do
        let newNet = trainNet net trainVectors
        print layers
        print (n + times)
        print newNet
        endTime <- getCurrentTime
        print $ (endTime `diffUTCTime` startTime) - timeTakenEvaluating
        evalStartTime <- getCurrentTime
        print $ Counter.fromList $ map (\(x, y) -> (y, Vector.cmap r $ Neural.output newNet tanh x)) testVectors
        evalEndTime <- getCurrentTime
        print $ evalEndTime `diffUTCTime` evalStartTime
        return (n + times, newNet, timeTakenEvaluating + (evalEndTime `diffUTCTime` evalStartTime))
  where r = (fromIntegral :: Int -> Double) . (round :: Double -> Int)
        trainNet = Neural.trainNTimes times 0.8 Neural.tanh Neural.tanh'

autoencoder :: FilePath -> IO (Counter Text, Autoencoder)
autoencoder path = do
  shuffled <- readFile path >>= getStdRandom . shuffle . extractData
  let (train, test) = splitAt (length shuffled `div` 2) shuffled
  let trainClassifier = mconcat $ map rowToClassifier train
  let vocab = vocabulary trainClassifier
  let trainVectors = classifierToVector boolToVector vocab trainClassifier
  let testVectors = classifierToVector boolToVector vocab $ mconcat $ map rowToClassifier test
  case trainVectors of
    [] -> error "no data"
    v -> do
      res <- generateAutoencoderIO (map fst v) 1000 1000
      return (vocab, res)

evaluateBayes :: FilePath -> IO (Counter (Bool, Maybe Bool))
evaluateBayes fp = do
  file <- readFile fp
  startTime <- getCurrentTime
  let res = extractData file
  let classifier = createClassifier res
  print classifier
  printTimeDiff startTime
  evalStartTime <- getCurrentTime
  let counted = Counter.fromList $ applyNaiveBayes classifier res
  print counted
  printTimeDiff evalStartTime
  return counted

printTimeDiff :: UTCTime -> IO ()
printTimeDiff s = do
  e <- getCurrentTime
  print $ e `diffUTCTime` s

bayes :: String -> Counter (Bool, Maybe Bool)
bayes file = Counter.fromList results
  where res = extractData file
        results = applyNaiveBayes (createClassifier res) res

neural :: FilePath -> Int -> [Int] -> IO ()
neural path times layers = do
  shuffled <- readFile path >>= getStdRandom . shuffle . extractData
  let (train, test) = splitAt (length shuffled `div` 2) shuffled
  let trainClassifier = mconcat $ map rowToClassifier train
  let vocab = vocabulary trainClassifier
  let trainVectors = classifierToVector boolToVector vocab trainClassifier
  let testVectors = classifierToVector boolToVector vocab $ mconcat $ map rowToClassifier test
  case trainVectors of
    [] -> putStrLn "No data"
    (v, _) : _ -> do
      startTime <- getCurrentTime
      network <- Neural.createNetwork (Vector.size v) layers 1
      IO.hSetBuffering IO.stdout IO.NoBuffering
      void $ iterateM 9 (0 :: Int, network, 0) $ \(n, net, timeTakenEvaluating) -> do
        let newNet = trainNet net trainVectors
        print layers
        print (n + times)
        print newNet
        endTime <- getCurrentTime
        print $ (endTime `diffUTCTime` startTime) - timeTakenEvaluating
        evalStartTime <- getCurrentTime
        print $ Counter.fromList $ map (\(x, y) -> (y, Vector.cmap r $ Neural.output newNet tanh x)) testVectors
        evalEndTime <- getCurrentTime
        print $ evalEndTime `diffUTCTime` evalStartTime
        return (n + times, newNet, timeTakenEvaluating + (evalEndTime `diffUTCTime` evalStartTime))
  where r = (fromIntegral :: Int -> Double) . (round :: Double -> Int)
        trainNet = Neural.trainNTimes times 0.8 Neural.tanh Neural.tanh'

iterateM :: Monad m => Int -> a -> (a -> m a) -> m [a]
iterateM 0 _ _ = return []
iterateM n a f = do
  v <- f a
  next <- iterateM (n - 1) v f
  return $ v : next

boolToVector :: Bool -> Vector Double
boolToVector True = 1
boolToVector False = -1

vocabulary :: Ord b => Classifier a b -> Counter b
vocabulary = Map.foldr (mappend . mconcat) mempty . Classifier.toMap

counterToVector :: Ord a => Counter a -> Counter a -> Vector Double
counterToVector (Counter.toMap -> vocab) (Counter.toMap -> m) =
  Vector.vector $ map snd $ Map.toAscList $ Map.mergeWithKey (\_ v _ -> Just $ fromIntegral v) (const mempty) (fmap (const 0)) m vocab

classifierToVector :: (Ord a, Ord b) => (a -> Vector Double) -> Counter b -> Classifier a b -> [(Vector Double, Vector Double)]
classifierToVector f vocab (Classifier.toMap -> m) =
  Map.foldrWithKey (\k v a -> fmap ((,) <$> counterToVector vocab <*> pure (f k)) v <> a) [] m

applyNaiveBayes :: NaiveBayes Bool Text -> [Row] -> [(Bool, Maybe Bool)]
applyNaiveBayes classifier rows =
  foldl (\ a t -> collect (NaiveBayes.remove (rowToClassifier t) classifier) a t) [] rows

collect :: NaiveBayes Bool Text -> [(Bool, Maybe Bool)] -> Row -> [(Bool, Maybe Bool)]
collect cls acc (b, _, _, _, c) =
  (b, tested) : acc
  where tested = NaiveBayes.test cls $ Counter.fromList $ process c

type Row = (Bool, Text, Text, Text, Text)

extractData :: String -> [Row]
extractData = mapMaybe readMaybe . lines

createClassifier :: [Row] -> NaiveBayes Bool Text
createClassifier = mconcat . map (NaiveBayes.fromClassifier . rowToClassifier)

rowToClassifier :: Row -> Classifier Bool Text
rowToClassifier (b, _, _, _, c) = Classifier.singleton b $ Counter.fromList $ process c

process :: Text -> [Text]
process = filter (not . Text.null) .
          map (Text.map toLower . Text.filter isAlpha) .
          concatMap (Text.splitOn ".") .
          Text.splitOn " " .
          Text.filter (not . (== '-'))

-- shuffle from the haskell wiki @ https://wiki.haskell.org/Random_shuffle
shuffle :: RandomGen g => [a] -> g -> ([a], g)
shuffle xs gen = runST $ do
    g <- newSTRef gen
    ar <- newListArray' xs
    xs' <- forM [1..n] $ \i -> do
            j <- randomRST (i,n) g
            vi <- readArray ar i
            vj <- readArray ar j
            writeArray ar j vi
            return vj
    gen' <- readSTRef g
    return (xs', gen')
  where
    n = length xs
    newListArray' :: [a] -> ST s (STArray s Int a)
    newListArray' = newListArray (1, n)
    randomRST lohi g = do
      (a, s') <- randomR lohi <$> readSTRef g
      writeSTRef g s'
      return a


{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveDataTypeable        #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE QuasiQuotes               #-}
{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE TupleSections             #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_GHC  -fplugin=Polysemy.Plugin  #-}

module MRP.WWC (post) where

import qualified Control.Foldl                 as FL
import qualified Data.List as L
import qualified Data.Set as S
import qualified Data.Map                      as M
import qualified Data.Text                     as T
import qualified Data.Serialize                as SE

import           Graphics.Vega.VegaLite.Configuration as FV
import qualified Graphics.Vega.VegaLite.Compat as FV
import qualified Frames as F
import qualified Data.Vinyl as V

import qualified Graphics.Vega.VegaLite        as GV
import qualified Knit.Report                   as K
import qualified Polysemy.Error                as P (mapError)
import           Text.Pandoc.Error             as PE
import qualified Text.Blaze.Colonnade          as BC

import           Data.String.Here               ( i )

import qualified Colonnade                     as C
import qualified Text.Blaze.Html5.Attributes   as BHA

import           BlueRipple.Configuration
import           BlueRipple.Utilities.KnitUtils
import           BlueRipple.Utilities.TableUtils

import qualified Numeric.GLM.ProblemTypes      as GLM
import qualified Numeric.GLM.ModelTypes      as GLM
import qualified Numeric.GLM.Bootstrap            as GLM
import qualified Numeric.GLM.Predict            as GLM
import qualified Numeric.GLM.Confidence            as GLM

import qualified Statistics.Types              as ST
import GHC.Generics (Generic)

import BlueRipple.Data.DataFrames
import qualified BlueRipple.Data.DemographicTypes as BR
import qualified BlueRipple.Data.ElectionTypes as ET
import qualified BlueRipple.Model.MRP as BR
import MRP.Common
import BlueRipple.Data.CCES as CCES

import qualified PreferenceModel.Common as PrefModel

brIntro :: T.Text
brIntro = [i|
In our research pieces so far, we've looked only at aggregate data, that
is data which comes from adding together a large number of people: census
counts or election results, for example.  In this piece we look at some
per-person data, namely the CCES survey, which surveys about 60,000 people
in every election cycle. For each survey response,
the CCES includes geographic and demographic information along with opinion about
various political questions, whether the person is registered to vote, and whether they
voted and who for in elections for Governor, House, Senate and President, whenever
each is applicable.

The geographic information makes this interesting.  This allows
us to start estimating a variety of things at the state level, something we couldn't do using
only aggregate data.  We do this using multi-level regression
(the "MR" of "MRP", which stands for Multi-level Regression with Post-stratification), a technique
we explain in more detail [here][MRP:Methods].

Following up on our earlier posts about
[changing voter preference][PrefModel:AcrossTime] and
the [white working class][PrefModel:WWCV] (WWC),
we look again at WWC voters in 2016 and
2018, this time estimating voter preference in each state and looking at
the changes, with a particular focus on the battleground states in the
mid-west: Indiana, Michigan, Ohio, Pennsylvania and Wisconsin.

[MRP:Methods]: <${brGithubUrl (postPath PostMethods)}>
[PrefModel:WWCV]: <${brGithubUrl (PrefModel.postPath PrefModel.PostWWCV)}>
[PrefModel:AcrossTime]: <${brGithubUrl (PrefModel.postPath PrefModel.PostAcrossTime)}>
|]

glmErrorToPandocError :: GLM.GLMError -> PE.PandocError
glmErrorToPandocError x = PE.PandocSomeError $ T.pack $ show x

--type BR.LocationCols = '[StateAbbreviation]
type CatCols = '[]
--type CCESGroup = Proxy (BR.GroupCols Location

post :: (K.KnitOne r, K.Member GLM.RandomFu r, K.Member GLM.Async r)
     => M.Map T.Text T.Text -- state names from state abbreviations
--     -> K.Cached es [F.Record CCES.CCES_MRP]
     -> K.Sem r ()
post stateNameByAbbreviation = P.mapError glmErrorToPandocError $ K.wrapPrefix "Intro" $ return () {- $ do

  K.logLE K.Info $ "Working on Intro post..."
  let isWWC r = (F.rgetField @BR.SimpleRaceC r == BR.White) && (F.rgetField @BR.CollegeGradC r == BR.NonGrad)
      countWWCDemPres2016VotesF = BR.weightedCountFold @ByCCESPredictors @CCES.CCES_MRP @'[CCES.Pres2016VoteParty,CCES.CCESWeightCumulative]
                                  (\r -> (F.rgetField @CCES.Turnout r == CCES.T_Voted)
                                         && (F.rgetField @Year r == 2016)
                                         && isWWC r
                                         && (F.rgetField @CCES.Pres2016VoteParty r `elem` [ET.Republican, ET.Democratic]))
                                  ((== ET.Democratic) . F.rgetField @CCES.Pres2016VoteParty)
                                  (F.rgetField @CCES.CCESWeightCumulative)
      countWWCDemHouseVotesF y = BR.weightedCountFold @ByCCESPredictors @CCES.CCES_MRP @'[CCES.HouseVoteParty,CCES.CCESWeightCumulative]
                                 (\r -> (F.rgetField @CCES.Turnout r == CCES.T_Voted)
                                        && (F.rgetField @Year r == y)
                                        && isWWC r
                                        && (F.rgetField @CCES.HouseVoteParty r `elem` [ET.Republican, ET.Democratic]))
                                 ((== ET.Democratic) . F.rgetField @CCES.HouseVoteParty)
                                 (F.rgetField @CCES.CCESWeightCumulative)
  let --makeTableRows :: K.Sem r [WWCTableRow]
      makeWWCTableRows = do
        ccesFrameAll <- ccesDataLoader -- F.toFrame <$> P.raise (K.useCached ccesRecordListAllCA)
        (mm2016p, rc2016p, ebg2016p, bu2016p, vb2016p, bs2016p) <- BR.inferMR @BR.LocationCols @CatCols @[BR.SimpleAgeC
                                                                                ,BR.SexC
                                                                                ,BR.CollegeGradC
                                                                                ,BR.SimpleRaceC
                                                                                ]
                                                                   countWWCDemPres2016VotesF
                                                                   [GLM.Intercept, GLM.Predictor CCES.P_WWC]
                                                                   ccesPredictor
                                                                   (fmap F.rcast ccesFrameAll)

        (mm2016, rc2016, ebg2016, bu2016, vb2016, bs2016) <- BR.inferMR @BR.LocationCols @CatCols @[BR.SimpleAgeC
                                                                                ,BR.SexC
                                                                                ,BR.CollegeGradC
                                                                                ,BR.SimpleRaceC
                                                                                ]
                                                             (countWWCDemHouseVotesF 2016)
                                                             [GLM.Intercept, GLM.Predictor P_WWC]
                                                             ccesPredictor
                                                             ccesFrameAll

        (mm2018, rc2018, ebg2018, bu2018, vb2018, bs2018) <- BR.inferMR @BR.LocationCols @CatCols @[BR.SimpleAgeC
                                                                                ,BR.SexC
                                                                                ,BR.CollegeGradC
                                                                                ,BR.SimpleRaceC
                                                                                ]
                                                             (countWWCDemHouseVotesF 2018)
                                                             [GLM.Intercept, GLM.Predictor P_WWC]
                                                             ccesPredictor
                                                             ccesFrameAll

        let states = FL.fold FL.set $ fmap (F.rgetField @StateAbbreviation) ccesFrameAll
            toPredict = [("National", M.empty)] <> (fmap (\s -> (s,M.singleton BR.RecordColsProxy (s F.&: V.RNil))) $ S.toList states)
            wwc = M.singleton P_WWC 1
            addPredictions (s, groupMap) = P.mapError glmErrorToPandocError $ do
              let cl = ST.mkCL 0.95
                  ct bs = GLM.BootstrapCI GLM.BCI_Accelerated bs
              p2016p <- GLM.predictWithCI mm2016p (flip M.lookup wwc) (flip M.lookup groupMap) rc2016p ebg2016p bu2016p vb2016p cl (ct bs2016p)
              p2016 <- GLM.predictWithCI mm2016 (flip M.lookup wwc) (flip M.lookup groupMap) rc2016 ebg2016 bu2016 vb2016 cl (ct bs2016)
              p2018 <- GLM.predictWithCI mm2018 (flip M.lookup wwc) (flip M.lookup groupMap) rc2018 ebg2018 bu2018 vb2018 cl (ct bs2018)
              return  $ WWCTableRow s p2016p p2016 p2018
        L.sortOn deltaHouse <$> traverse addPredictions toPredict
  forWWCTable <-  K.retrieveOrMake "mrp/intro/wwcPrefTable" makeWWCTableRows -- cache this so we don't need to redo
  forWWCChart <-  GLM.eitherToSem $ traverse (\tr -> do
                                            let sa = stateAbbr tr
                                            fullState <- maybe (Left $ "Couldn't find " <> sa) Right $ M.lookup sa stateNameByAbbreviation
                                            return (fullState, significantDeltaHouse tr)) $ filter (\tr -> (stateAbbr tr) /= "National") forWWCTable
  brAddMarkDown brIntro
  _ <- K.addHvega Nothing Nothing $ (vlPctStateChoropleth "Significant Change in WWC Dem Voter Preference 2016 to 2018" (FV.ViewConfig 800 400 10) forWWCChart)
  let mwBG = ["IA", "MI", "OH", "PA", "WI"]
  brAddRawHtmlTable
    "WWC Democratic Voter Preference"
    (BHA.class_ "brTable")
    (colWWC_Change $ emphasizeStates mwBG <> emphasizeNational) forWWCTable

  brAddMarkDown brReadMore
-}

data WWCTableRow = WWCTableRow { stateAbbr :: T.Text
                               , pref2016Pres :: (Double, (Double, Double))
                               , pref2016House :: (Double, (Double, Double))
                               , pref2018House :: (Double, (Double, Double))
                               } deriving (Generic)

instance SE.Serialize WWCTableRow

deltaHouse tr = fst (pref2018House tr) - fst (pref2016House tr)
inStates s tr = stateAbbr tr `elem` s

emphasizeStates s = CellStyle (\tr _ -> if inStates s tr then highlightCellBlue else "")
emphasizeNational = CellStyle (\tr _ -> if  stateAbbr tr == "National" then highlightCellPurple else "")

significantGivenCI :: (Double, Double) -> (Double, Double) -> Double
significantGivenCI (loA, hiA) (loB, hiB) =
  let maxLo = max loA loB
      minHi = min hiA hiB
      overlap = max 0 (minHi - maxLo)
      shortest = min (hiA - loA) (hiB - loB)
  in overlap / shortest

significantDeltaHouse tr = if signifWWCT tr < 0.25 then deltaHouse tr else 0

signifWWCT x = significantGivenCI (snd $ pref2016House x) (snd $ pref2018House x)

colWWC_Change :: CellStyle WWCTableRow T.Text -> C.Colonnade C.Headed WWCTableRow BC.Cell
colWWC_Change cas =
  C.headed "State" (toCell cas "State" "State" (textToStyledHtml . stateAbbr))
  <> C.headed "2016 Pres (%)" (toCell cas "2016 Pres" "2016 Pres" (numberToStyledHtml "%2.1f" . (*100) . fst . pref2016Pres))
  <> C.headed "2016 House (%)" (toCell cas "2016 House" "2016 House" (numberToStyledHtml "%2.1f" . (*100) . fst . pref2016House))
  <> C.headed "Low (95%)" (toCell cas "Low (95%)" "Low (95%)" (numberToStyledHtml "%2.1f" . (*100) . fst . snd . pref2016House))
  <> C.headed "High (95%)" (toCell cas "High (95%)" "High (95%)" (numberToStyledHtml "%2.1f" . (*100) . snd . snd . pref2016House))
  <> C.headed "2018 House (%)" (toCell cas "2018 D Pref" "2018 House" (numberToStyledHtml "%2.1f" . (*100) . fst . pref2018House))
  <> C.headed "2018-2016 House" (toCell cas "2018-2016 House" "2018-2016 House" (numberToStyledHtml "%2.1f" . (*100) . deltaHouse))
  <> C.headed "CI Overlap" (toCell cas "CI Overlap" "CI Overlap" (numberToStyledHtml "%2.1f" . (*100) . signifWWCT))

usStatesTopoJSONUrl = "https://cdn.jsdelivr.net/npm/us-atlas@3/states-10m.json"
usStatesAlbersTopoJSONUrl = "https://cdn.jsdelivr.net/npm/us-atlas@3/states-albers-10m.json"

-- Lessons learned
-- The data is the geography
-- transform with lookup (properties.name) to add data
-- to figure out the right lookup name you can just plot the geography (no transform, no encoding) and look at loaded data
-- in vega editor
vlPctStateChoropleth :: Foldable f => T.Text -> FV.ViewConfig -> f (T.Text, Double) -> GV.VegaLite
vlPctStateChoropleth title vc stateData =
  let datGeo = GV.dataFromUrl usStatesTopoJSONUrl [GV.TopojsonFeature "states"]
      datVal = GV.dataFromRows [] $ concat $ fmap (\(s,x) -> GV.dataRow [("state", GV.Str s),("value", GV.Number x)] []) $ FL.fold FL.list stateData
      dataSets = GV.datasets [("stateNums",datVal)]
      projection = GV.projection [GV.PrType GV.AlbersUsa]
      transform = GV.transform . GV.lookup "properties.name" datVal "state" (GV.LuFields ["state","value"])
      mark = GV.mark GV.Geoshape []
      colorEnc = GV.color [GV.MName "value", GV.MmType GV.Quantitative, GV.MScale [GV.SScheme "redyellowgreen" [], GV.SDomain (GV.DNumbers [-0.5,0.5])]]
      tooltip = GV.tooltips [[GV.TName "state", GV.TmType GV.Nominal],[GV.TName "value", GV.TmType GV.Quantitative, GV.TFormat ".0%"]]
      enc = GV.encoding .  colorEnc . tooltip
  in FV.configuredVegaLite vc [FV.title title, datGeo, mark, projection, transform [], enc []]

vlTest :: FV.ViewConfig -> GV.VegaLite
vlTest vc =
  let
    datGeo = GV.dataFromUrl usStatesTopoJSONUrl [GV.TopojsonFeature "states"]
    projection = GV.projection [GV.PrType GV.AlbersUsa]
    mark = GV.mark GV.Geoshape [GV.MFill "lightgrey" ]
  in FV.configuredVegaLite vc [datGeo, mark, projection]


{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

--
-- Copyright (c) 2009-2011, ERICSSON AB
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--
--     * Redistributions of source code must retain the above copyright notice,
--       this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of the ERICSSON AB nor the names of its contributors
--       may be used to endorse or promote products derived from this software
--       without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-- FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-- DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-- CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-- OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-- OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--

module Feldspar.Core.Constructs.Complex
where

import Language.Syntactic
import Language.Syntactic.Constructs.Binding

import Data.Complex

import Feldspar.Core.Types
import Feldspar.Core.Interpretation

data COMPLEX a
  where
    MkComplex :: (Type a, RealFloat a) => COMPLEX (a :-> a :-> Full (Complex a))
    RealPart  :: (Type a, RealFloat a) => COMPLEX (Complex a :-> Full a)
    ImagPart  :: (Type a, RealFloat a) => COMPLEX (Complex a :-> Full a)
    Conjugate :: (Type a, RealFloat a) => COMPLEX (Complex a :-> Full (Complex a))
    MkPolar   :: (Type a, RealFloat a) => COMPLEX (a :-> a :-> Full (Complex a))
    Magnitude :: (Type a, RealFloat a) => COMPLEX (Complex a :-> Full a)
    Phase     :: (Type a, RealFloat a) => COMPLEX (Complex a :-> Full a)
    Cis       :: (Type a, RealFloat a) => COMPLEX (a :-> Full (Complex a))

instance Semantic COMPLEX
  where
    semantics MkComplex = Sem "complex"   (:+)
    semantics RealPart  = Sem "creal"     realPart
    semantics ImagPart  = Sem "cimag"     imagPart
    semantics Conjugate = Sem "conjugate" conjugate
    semantics MkPolar   = Sem "mkPolar"   mkPolar
    semantics Magnitude = Sem "magnitude" magnitude
    semantics Phase     = Sem "phase"     phase
    semantics Cis       = Sem "cis"       cis

instance Equality COMPLEX where equal = equalDefault; exprHash = exprHashDefault
instance Render   COMPLEX where renderArgs = renderArgsDefault
instance ToTree   COMPLEX
instance Eval     COMPLEX where evaluate = evaluateDefault
instance EvalBind COMPLEX where evalBindSym = evalBindSymDefault
instance Sharable COMPLEX

instance SizeProp (COMPLEX :|| Type)
  where
    sizeProp (C' s) = sizePropDefault s

instance AlphaEq dom dom dom env => AlphaEq COMPLEX COMPLEX dom env
  where
    alphaEqSym = alphaEqSymDefault

instance ( (COMPLEX :|| Type) :<: dom
         , OptimizeSuper dom)
      => Optimize (COMPLEX :|| Type) dom
  where
    constructFeatOpt _ (C' MkComplex) ((rp :$ a) :* (ip :$ b) :* Nil)
        | Just (C' RealPart) <- prjF rp
        , Just (C' ImagPart) <- prjF ip
        , alphaEq a b
        = return a

    constructFeatOpt _ (C' RealPart) ((mkc :$ r :$ _) :* Nil)
        | Just (C' MkComplex) <- prjF mkc
        = return r

    constructFeatOpt _ (C' ImagPart) ((mkc :$ _ :$ i) :* Nil)
        | Just (C' MkComplex) <- prjF mkc
        = return i

    constructFeatOpt _ (C' MkPolar) ((mag :$ a) :* (ph :$ b) :* Nil)
        | Just (C' Magnitude) <- prjF mag
        , Just (C' Phase)     <- prjF ph
        , alphaEq a b
        = return a

    constructFeatOpt _ (C' Magnitude) ((mkp :$ m :$ _) :* Nil)
        | Just (C' MkPolar) <- prjF mkp
        = return m

    constructFeatOpt _ (C' Phase) ((mkp :$ _ :$ p) :* Nil)
        | Just (C' MkPolar) <- prjF mkp
        = return p

    constructFeatOpt _ (C' Conjugate) ((conj :$ a) :* Nil)
        | Just (C' Conjugate) <- prjF conj
        = return a

    constructFeatOpt opts sym args = constructFeatUnOpt opts sym args

    constructFeatUnOpt opts x@(C' _) = constructFeatUnOptDefault opts x


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-|
Module      : Grenade.Core.Pooling
Description : Max Pooling layer for 2D and 3D images
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Pooling (
    Pooling (..)
  ) where

import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits hiding (natVal)
import           GHC.TypeLits
import           Control.DeepSeq (NFData)
import           GHC.Generics    (Generic)

import           Grenade.Core
import           Grenade.Layers.Internal.Pooling

import           Numeric.LinearAlgebra.Static as LAS hiding ((|||), build, toRows)

-- | A pooling layer for a neural network.
--
--   Does a max pooling, looking over a kernel similarly to the convolution network, but returning
--   maxarg only. This layer is often used to provide minor amounts of translational invariance.
--
--   The kernel size dictates which input and output sizes will "fit". Fitting the equation:
--   `out = (in - kernel) / stride + 1` for both dimensions.
--
data Pooling :: Nat -> Nat -> Nat -> Nat -> * where
  Pooling :: Pooling kernelRows kernelColumns strideRows strideColumns
  deriving (NFData, Generic)

instance Show (Pooling k k' s s') where
  show Pooling = "Pooling"

instance UpdateLayer (Pooling kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Pooling kernelRows kernelColumns strideRows strideColumns) = ()
  runUpdate _ Pooling _ = Pooling

instance RandomLayer (Pooling k k' s s') where
  createRandomWith _ _ = return Pooling

instance Serialize (Pooling kernelRows kernelColumns strideRows strideColumns) where
  put _ = return ()
  get = return Pooling

-- | A two dimentional image can be pooled.
instance ( KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputColumns - 1) * strideColumns) ~ (inputColumns - kernelColumns)
         ) => Layer (Pooling kernelRows kernelColumns strideRows strideColumns) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) where
  type Tape (Pooling kernelRows kernelColumns strideRows strideColumns) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) = S ('D2 inputRows inputColumns)
  runForwards Pooling (S2D input) =
    let height = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        width  = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ex = extract input
        r  = poolForward 1 height width kx ky sx sy ex
        rs = fromJust . create $ r
    in  (S2D input, S2D rs)
  runBackwards Pooling (S2D input) (S2D dEdy) =
    let height = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        width  = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ex = extract input
        eo = extract dEdy
        vs = poolBackward 1 height width kx ky sx sy ex eo
    in  ((), S2D . fromJust . create $ vs)

-- | A three dimensional image can be pooled on each layer.
instance ( KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , KnownNat channels
         , KnownNat (outputRows * channels)
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputColumns - 1) * strideColumns) ~ (inputColumns - kernelColumns)
         ) => Layer (Pooling kernelRows kernelColumns strideRows strideColumns) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) where
  type Tape (Pooling kernelRows kernelColumns strideRows strideColumns) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) = S ('D3 inputRows inputColumns channels)
  runForwards Pooling (S3D input) =
    let ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ch = fromIntegral $ natVal (Proxy :: Proxy channels)
        ex = extract input
        r  = poolForward ch ix iy kx ky sx sy ex
        rs = fromJust . create $ r
    in  (S3D input, S3D rs)
  runBackwards Pooling (S3D input) (S3D dEdy) =
    let ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        ch = fromIntegral $ natVal (Proxy :: Proxy channels)
        ex = extract input
        eo = extract dEdy
        vs = poolBackward ch ix iy kx ky sx sy ex eo
    in  ((), S3D . fromJust . create $ vs)

-------------------- GNum instance --------------------

instance GNum (Pooling k k' s s') where
  _ |* _ = Pooling
  _ |+ _ = Pooling
  gFromRational _ = Pooling


module HVX.DcpTests.OkXAffine where

import Numeric.LinearAlgebra

import HVX

main :: IO ()
main = do
  let zero = EConst $ konst 0.0 (2, 1)
      x = EVar "x"
      e = hexp $ x
      _ = zero >=~ e
      _ = e <=~ zero
  return ()