{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Matrix.Herm
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Hermitian views of matrices.
--

module Numeric.LinearAlgebra.Matrix.Herm (
    -- * Immutable interface

    -- ** Vector multiplication
    hermMulVector,
    hermMulVectorWithScale,
    addHermMulVectorWithScales,

    -- ** Matrix  multiplication
    hermMulMatrix,
    hermMulMatrixWithScale,
    addHermMulMatrixWithScales,

    -- ** Updates
    hermRank1Update,
    hermRank2Update,
    hermRankKUpdate,
    hermRank2KUpdate,

    -- * Mutable interface
    hermCreate,

    -- ** Vector multiplication
    hermMulVectorTo,
    hermMulVectorWithScaleTo,
    addHermMulVectorWithScalesM_,

    -- ** Matrix multiplication
    hermMulMatrixTo,
    hermMulMatrixWithScaleTo,
    addHermMulMatrixWithScalesM_,

    -- ** Updates
    hermRank1UpdateM_,
    hermRank2UpdateM_,
    hermRankKUpdateM_,
    hermRank2KUpdateM_,

    ) where

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

import Numeric.LinearAlgebra.Vector( Vector, RVector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V
import Numeric.LinearAlgebra.Matrix.Base( Matrix )
import Numeric.LinearAlgebra.Matrix.STBase( STMatrix, RMatrix )
import qualified Numeric.LinearAlgebra.Matrix.STBase as M
import Numeric.LinearAlgebra.Types
import qualified Foreign.BLAS as BLAS

-- | A safe way to create and work with a mutable Herm Matrix before returning
-- an immutable one for later perusal.
hermCreate :: (Storable e)
           => (forall s. ST s (Herm (STMatrix s) e))
           -> Herm Matrix e
hermCreate mh = runST $ do
    (Herm u ma) <- mh
    a <- M.unsafeFreeze ma
    return $ Herm u a

-- | @hermRank1Update alpha x a@ returns
-- @alpha * x * x^H + a@.
hermRank1Update :: (BLAS2 e)
                => Double -> Vector e -> Herm Matrix e -> Herm Matrix e
hermRank1Update alpha x (Herm uplo a) = runST $ do
    ma' <- M.newCopy a
    hermRank1UpdateM_ alpha x (Herm uplo ma')
    a' <- M.unsafeFreeze ma'
    return $ Herm uplo a'

-- | @hermRank2Update alpha x y a@ returns
-- @alpha * x * y^H + conj(alpha) * y * x^H + a@.
hermRank2Update :: (BLAS2 e)
                => e -> Vector e -> Vector e -> Herm Matrix e
                -> Herm Matrix e
hermRank2Update alpha x y (Herm uplo a) = runST $ do
    ma' <- M.newCopy a
    hermRank2UpdateM_ alpha x y (Herm uplo ma')
    a' <- M.unsafeFreeze ma'
    return $ Herm uplo a'

-- | @hermRankKUpdate alpha trans a beta c@ returns
-- @c := alpha * a * a^H + beta * c@ when @trans@ is @NoTrans@ and
-- @c := alpha * a^H * a + beta * c@ when @trans@ is @ConjTrans@.  The
-- function signals an error when @trans@ is @Trans@.
hermRankKUpdate :: (BLAS3 e)
                => e -> Trans -> Matrix e -> e -> Herm Matrix e
                -> Herm Matrix e
hermRankKUpdate alpha trans a beta (Herm uplo c) = runST $ do
    mc' <- M.newCopy c
    hermRankKUpdateM_ alpha trans a beta (Herm uplo mc')
    c' <- M.unsafeFreeze mc'
    return $ Herm uplo c'

-- | @hermRank2KUpdate alpha trans a b beta c@ returns
-- @c := alpha * a * b^H + conj(alpha) * b * a^H + beta * c@ when @trans@ is
-- @NoTrans@ and @c := alpha * b^H * a + conj(alpha) * a^H * b + beta * c@
-- when @trans@ is @ConjTrans@.  The function signals an error when @trans@
-- is @Trans@.
hermRank2KUpdate :: (BLAS3 e)
                 => e -> Trans -> Matrix e -> Matrix e -> e -> Herm Matrix e
                 -> Herm Matrix e
hermRank2KUpdate alpha trans a b beta (Herm uplo c) = runST $ do
    mc' <- M.newCopy c
    hermRank2KUpdateM_ alpha trans a b beta (Herm uplo mc')
    c' <- M.unsafeFreeze mc'
    return $ Herm uplo c'

-- | @hermRank1UpdateM_ alpha x a@ sets
-- @a := alpha * x * x^H + a@.
hermRank1UpdateM_ :: (RVector v, BLAS2 e)
                  => Double -> v e -> Herm (STMatrix s) e -> ST s ()
hermRank1UpdateM_ alpha x (Herm uplo a) = do
    nx <- V.getDim x
    (ma,na) <- M.getDim a
    let n = nx

    when ((not . and) [ nx == n, (ma,na) == (n,n) ]) $ error $
        printf ("hermRank1UpdateM_ _ <vector with dim %d>"
                 ++ " (Herm _ <matrix with dim (%d,%d)>):"
                 ++ " invalid dimensions") nx ma na

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        M.unsafeWith a $ \pa lda ->
            BLAS.her uplo n alpha px 1 pa lda

-- | @hermRank2UpdateM_ alpha x y a@ sets
-- @a := alpha * x * y^H + conj(alpha) * y * x^H + a@.
hermRank2UpdateM_ :: (RVector v1, RVector v2, BLAS2 e)
                  => e -> v1 e -> v2 e -> Herm (STMatrix s) e -> ST s ()
hermRank2UpdateM_ alpha x y (Herm uplo a) = do
    nx <- V.getDim x
    ny <- V.getDim y
    (ma,na) <- M.getDim a
    let n = nx

    when ((not . and) [ nx == n, ny == n, (ma,na) == (n,n) ]) $ error $
        printf ("hermRank2UpdateM_ _ <vector with dim %d>"
                 ++ " <vector with dim %d>"
                 ++ " (Herm _ <matrix with dim (%d,%d)>):"
                 ++ " invalid dimensions") nx ny ma na

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
        M.unsafeWith a $ \pa lda ->
            BLAS.her2 uplo n alpha px 1 py 1 pa lda

-- | @hermRankKUpdateM_ alpha trans a beta c@ sets
-- @c := alpha * a * a^H + beta * c@ when @trans@ is @NoTrans@ and
-- @c := alpha * a^H * a + beta * c@ when @trans@ is @ConjTrans@.  The
-- function signals an error when @trans@ is @Trans@.
hermRankKUpdateM_ :: (RMatrix m, BLAS3 e)
                  => e -> Trans -> m e -> e -> Herm (STMatrix s) e
                  -> ST s ()
hermRankKUpdateM_ alpha trans a beta (Herm uplo c) = do
    (ma,na) <- M.getDim a
    (mc,nc) <- M.getDim c
    let (n,k) = if trans == NoTrans then (ma,na) else (na,ma)

    when (trans == Trans) $ error $
        printf ("hermRankKUpdateM_ _ %s:"
                 ++ " trans argument must be NoTrans or ConjTrans")
               (show trans)

    when ((not . and) [ (mc,nc) == (n,n)
                      , case trans of NoTrans -> (ma,na) == (n,k)
                                      _       -> (ma,na) == (k,n)
                      ]) $ error $
            printf ("hermRankKUpdateM_ _ %s <matrix with dim (%d,%d)> _"
                    ++ " (Herm _ <matrix with dim (%d,%d)>):"
                    ++ " invalid dimensions") (show trans) ma na mc nc

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith c $ \pc ldc ->
            BLAS.herk uplo trans n k alpha pa lda beta pc ldc

-- | @hermRank2KUpdateM_ alpha trans a b beta c@ sets
-- @c := alpha * a * b^H + conj(alpha) * b * a^H + beta * c@ when @trans@ is
-- @NoTrans@ and @c := alpha * b^H * a + conj(alpha) * a^H * b + beta * c@
-- when @trans@ is @ConjTrans@.  The function signals an error when @trans@
-- is @Trans@.
hermRank2KUpdateM_ :: (RMatrix m1, RMatrix m2, BLAS3 e)
                   => e -> Trans -> m1 e -> m2 e -> e -> Herm (STMatrix s) e
                   -> ST s ()
hermRank2KUpdateM_ alpha trans a b beta (Herm uplo c) = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    (mc,nc) <- M.getDim c
    let (n,k) = if trans == NoTrans then (ma,na) else (na,ma)

    when (trans == Trans) $ error $
        printf ("hermRank2KUpdateM_ _ %s:"
                 ++ " trans argument must be NoTrans or ConjTrans")
               (show trans)

    when ((not . and) [ (mc,nc) == (n,n)
                      , (mb,nb) == (ma,na)
                      , case trans of NoTrans -> (ma,na) == (n,k)
                                      _       -> (ma,na) == (k,n)
                      ]) $ error $
            printf ("hermRank2KUpdateM_ _ %s <matrix with dim (%d,%d)>"
                    ++ " <matrix with dim (%d,%d)> _"
                    ++ " (Herm _ <matrix with dim (%d,%d)>):"
                    ++ " invalid dimensions") (show trans) ma na mb nb mc nc

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
        M.unsafeWith c $ \pc ldc ->
            BLAS.her2k uplo trans n k alpha pa lda pb ldb beta pc ldc

-- | @hermMulVector a x@ returns @a * x@.
hermMulVector :: (BLAS2 e)
              => Herm Matrix e
              -> Vector e
              -> Vector e
hermMulVector a x =
    V.create $ do
        n <- V.getDim x
        y <- V.new_ n
        hermMulVectorTo y a x
        return y

-- | @hermMulVectorWithScale alpha a x@ retunrs @alpha * a * x@.
hermMulVectorWithScale :: (BLAS2 e)
                       => e
                       -> Herm Matrix e
                       -> Vector e
                       -> Vector e
hermMulVectorWithScale alpha a x =
    V.create $ do
        n <- V.getDim x
        y <- V.new_ n
        hermMulVectorWithScaleTo y alpha a x
        return y

-- | @addHermMulVectorWithScales alpha a x y@
-- returns @alpha * a * x + beta * y@.
addHermMulVectorWithScales :: (BLAS2 e)
                           => e
                           -> Herm Matrix e
                           -> Vector e
                           -> e
                           -> Vector e
                           -> Vector e
addHermMulVectorWithScales alpha a x beta y =
    V.create $ do
        y' <- V.newCopy y
        addHermMulVectorWithScalesM_ alpha a x beta y'
        return y'

-- | @hermMulMatrix side a b@
-- returns @alpha * a * b@ when @side@ is @LeftSide@ and
-- @alpha * b * a@ when @side@ is @RightSide@.
hermMulMatrix :: (BLAS3 e)
              => Side -> Herm Matrix e
              -> Matrix e
              -> Matrix e
hermMulMatrix side a b =
    M.create $ do
        mn <- M.getDim b
        c <- M.new_ mn
        hermMulMatrixTo c side a b
        return c

-- | @hermMulMatrixWithScale alpha side a b@
-- returns @alpha * a * b@ when @side@ is @LeftSide@ and
-- @alpha * b * a@ when @side@ is @RightSide@.
hermMulMatrixWithScale :: (BLAS3 e)
                       => e
                       -> Side -> Herm Matrix e
                       -> Matrix e
                       -> Matrix e
hermMulMatrixWithScale alpha side a b =
    M.create $ do
        mn <- M.getDim b
        c <- M.new_ mn
        hermMulMatrixWithScaleTo c alpha side a b
        return c

-- | @addHermMulMatrixWithScales alpha side a b beta c@
-- returns @alpha * a * b + beta * c@ when @side@ is @LeftSide@ and
-- @alpha * b * a + beta * c@ when @side@ is @RightSide@.
addHermMulMatrixWithScales :: (BLAS3 e)
                           => e
                           -> Side -> Herm Matrix e
                           -> Matrix e
                           -> e
                           -> Matrix e
                           -> Matrix e
addHermMulMatrixWithScales alpha side a b beta c =
    M.create $ do
        c' <- M.newCopy c
        addHermMulMatrixWithScalesM_ alpha side a b beta c'
        return c'

-- | @hermMulVectorTo dst a x@ sets @dst := a * x@.
hermMulVectorTo :: (RMatrix m, RVector v, BLAS2 e)
                => STVector s e
                -> Herm m e
                -> v e
                -> ST s ()
hermMulVectorTo dst = hermMulVectorWithScaleTo dst 1

-- | @hermMulVectorWithScaleTo dst alpha a x@
-- sets @dst := alpha * a * x@.
hermMulVectorWithScaleTo :: (RMatrix m, RVector v, BLAS2 e)
                         => STVector s e
                         -> e
                         -> Herm m e
                         -> v e
                         -> ST s ()
hermMulVectorWithScaleTo dst alpha a x =
    addHermMulVectorWithScalesM_ alpha a x 0 dst

-- | @addHermMulVectorWithScalesM_ alpha a x beta y@
-- sets @y := alpha * a * x + beta * y@.
addHermMulVectorWithScalesM_ :: (RMatrix m, RVector v, BLAS2 e)
                             => e
                             -> Herm m e
                             -> v e
                             -> e
                             -> STVector s e
                             -> ST s ()
addHermMulVectorWithScalesM_ alpha (Herm uplo a) x beta y = do
    (ma,na) <- M.getDim a
    nx <- V.getDim x
    ny <- V.getDim y
    let n = ny

    when (ma /= na) $ error $
        printf ("addHermMulVectorWithScalesM_ _"
                ++ " (Herm %s <matrix with dim (%d,%d)>)"
                ++ " %s <vector with dim %d>"
                ++ " _"
                ++ " <vector with dim %d>: Herm matrix is not square")
               (show uplo) ma na
               nx ny

    when ((not . and) [ (ma,na) == (n,n)
                      , nx == n
                      , ny == n
                      ]) $ error $
        printf ("addHermMulVectorWithScalesM_ _"
                ++ " (Herm %s <matrix with dim (%d,%d)>)"
                ++ " %s <vector with dim %d>"
                ++ " _"
                ++ " <vector with dim %d>: dimension mismatch")
               (show uplo) ma na
               nx ny

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
            BLAS.hemv uplo n alpha pa lda px 1 beta py 1

-- | @hermMulMatrixTo dst side a b@
-- sets @dst := a * b@ when @side@ is @LeftSide@ and
-- @dst := b * a@ when @side@ is @RightSide@.
hermMulMatrixTo :: (RMatrix m1, RMatrix m2, BLAS3 e)
                => STMatrix s e
                -> Side -> Herm m1 e
                -> m2 e
                -> ST s ()
hermMulMatrixTo dst = hermMulMatrixWithScaleTo dst 1

-- | @hermMulMatrixWithScaleTo dst alpha side a b@
-- sets @dst := alpha * a * b@ when @side@ is @LeftSide@ and
-- @dst := alpha * b * a@ when @side@ is @RightSide@.
hermMulMatrixWithScaleTo :: (RMatrix m1, RMatrix m2, BLAS3 e)
                         => STMatrix s e
                         -> e
                         -> Side -> Herm m1 e
                         -> m2 e
                         -> ST s ()
hermMulMatrixWithScaleTo dst alpha side a b =
    addHermMulMatrixWithScalesM_ alpha side a b 0 dst

-- | @addHermMulMatrixWithScalesM_ alpha side a b beta c@
-- sets @c := alpha * a * b + beta * c@ when @side@ is @LeftSide@ and
-- @c := alpha * b * a + beta * c@ when @side@ is @RightSide@.
addHermMulMatrixWithScalesM_ :: (RMatrix m1, RMatrix m2, BLAS3 e)
                             => e
                             -> Side -> Herm m1 e
                             -> m2 e
                             -> e
                             -> STMatrix s e
                             -> ST s ()
addHermMulMatrixWithScalesM_ alpha side (Herm uplo a) b beta c = do
    (ma,na) <- M.getDim a
    (mb,nb) <- M.getDim b
    (mc,nc) <- M.getDim c
    let (m,n) = (mc,nc)

    when (ma /= na) $ error $
        printf ("addHermMulMatrixWithScalesM_ _"
                ++ " %s (Herm %s <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ " _"
                ++ " <matrix with dim (%d,%d)>: Herm matrix is not square")
               (show side) (show uplo) ma na
               mb nb
               mc nc
    when ((not . and) [ case side of LeftSide  -> (ma,na) == (m,m)
                                     RightSide -> (ma,na) == (n,n)
                      , (mb, nb ) == (m,n)
                      , (mc, nc ) == (m,n)
                      ]) $ error $
        printf ("addHermMulMatrixWithScalesM_ _"
                ++ " %s (Herm %s <matrix with dim (%d,%d)>)"
                ++ " <matrix with dim (%d,%d)>"
                ++ " _"
                ++ " <matrix with dim (%d,%d)>: dimension mismatch")
               (show side) (show uplo) ma na
               mb nb
               mc nc

    unsafeIOToST $
        M.unsafeWith a $ \pa lda ->
        M.unsafeWith b $ \pb ldb ->
        M.unsafeWith c $ \pc ldc ->
            BLAS.hemm side uplo m n alpha pa lda pb ldb beta pc ldc


#!/usr/bin/env stack
-- stack --resolver lts-15.04 runghc --package reanimate
module Main where

import           Codec.Picture.Types
import           Control.Lens
import           Control.Monad
import           Data.List
import           Data.Maybe
import qualified Data.Text                     as T
import           Data.Tuple
import qualified Data.Vector                   as V
import           Linear.V2
import           Linear.Vector
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix (Matrix, linearSolve, toLists,
                                                (><))
import           Reanimate
import           Reanimate.Math.Common         (barycentricCoords, isBetween,
                                                rayIntersect)

type Points = V.Vector (V2 Double)
type Edges = [(Int, Int, Int)]
data Mesh = Mesh { meshPoints :: Points, meshEdges :: Edges }
data MeshPair = MeshPair Points Points Edges
-- The points in a RelMesh are:
--   relMeshStatic ++ x where Ax = B
data RelMesh = RelMesh
  { relMeshStatic :: Points
  , relMeshEdges  :: Edges
  , relMeshA      :: Matrix Double
  , relMeshB      :: Matrix Double
  }
data RelMeshPair = RelMeshPair Points Edges (Matrix Double) (Matrix Double) (Matrix Double) (Matrix Double)
-- Linear interpolation on RelMesh gives smooth morph.
-- solveMesh :: RelMesh -> Mesh
-- mkRelative :: Mesh -> RelMesh
-- mkRelativePair :: MeshPair -> RelMeshPair
-- triangulate :: Polygon -> Polygon -> MeshPair ?
-- embed :: MeshPair -> MeshPair
-- compatible :: Mesh -> Mesh -> Maybe MeshPair
-- linearInterpolate :: MeshPair -> Double -> Mesh
-- convexInterpolate :: RelMeshPair -> Double -> RelMesh

main :: IO ()
main = reanimate morphAnimation

morphAnimation :: Animation
morphAnimation = playThenReverseA $ pauseAround 1 1 $ addStatic (mkBackground "black") $
  signalA (curveS 2) $ mkAnimation 5 $ \t -> lowerTransformations $ scale 3 $ pathify $ center $
  mkGroup
  [ translate (-1) 0 $ drawTrig (linearInterpolate meshPair t)
  , translate 1 0 $ drawTrig $ solveMesh $ convexInterpolate relPair t
  ]
  where
    relPair = mkRelativePair meshPair
    meshPair = fromJust $ compatible example1 example2

mkLineP :: P -> P -> SVG
mkLineP (V2 x1 y1) (V2 x2 y2) = mkLine (x1,y1) (x2,y2)

-- FIXME: Check that the triangles are all anticlockwise.
-- FIXME: Check that the edges connect all the points.
-- FIXME: Check that the edgse leave no gaps.
compatible :: Mesh -> Mesh -> Maybe MeshPair
compatible a b =
  if meshEdges a == meshEdges b && V.length (meshPoints a) == V.length (meshPoints b)
    then Just $ MeshPair (meshPoints a) (meshPoints b) (meshEdges a)
    else Nothing

linearInterpolate :: MeshPair -> Double -> Mesh
linearInterpolate (MeshPair aP bP edges) t = Mesh
    { meshPoints = V.zipWith (lerp (1-t)) aP bP
    , meshEdges = edges }

example1 :: Mesh
example1 = Mesh points edges
  where
    points = V.fromList
      [ V2 1 0
      , V2 (-1/2) (sqrt 3 / 2)
      , V2 (-1/2) (-sqrt 3 / 2)
      , 3 * points V.! 0 ^/ 4
      , 3 * points V.! 1 ^/ 4
      , 3 * points V.! 2 ^/ 4
      , points V.! 0 ^/ 2
      , points V.! 1 ^/ 2
      , points V.! 2 ^/ 2
      ]
    edges =
      [ (1,5,4), (1,2,5), (2,6,5), (2,3,6), (3,4,6), (3,1,4)
      , (4,8,7), (4,5,8), (5,9,8), (5,6,9), (6,7,9), (6,4,7), (7,8,9)]

example2 :: Mesh
example2 = Mesh points edges
  where
    points = V.fromList
      [ V2 1 0
      , V2 (-1/2) (sqrt 3 / 2)
      , V2 (-1/2) (-sqrt 3 / 2)
      , 3 * points V.! 2 ^/ 4
      , 3 * points V.! 0 ^/ 4
      , 3 * points V.! 1 ^/ 4
      , points V.! 1  ^/ 2
      , points V.! 2 ^/ 2
      , points V.! 0 ^/ 2
      ]
    edges = meshEdges example1

drawTrig (Mesh points gs) = withStrokeColor "grey" $ withFillColor "white" $
  withStrokeWidth (defaultStrokeWidth/2) $ mkGroup
  [ mkGroup
    [ mkGroup
      [ mkLine (ax, ay) (bx, by)
      , mkLine (bx, by) (cx, cy)
      , mkLine (cx, cy) (ax, ay)
      ]
    | (a, b, c) <- gs
    , let V2 ax ay = points V.! (a-1)
          V2 bx by = points V.! (b-1)
          V2 cx cy = points V.! (c-1)
    ]
  , mkGroup $ concat
    [ [ colored v $ translate ax ay $ mkCircle circleRadius
      , withStrokeWidth 0 $
        withStrokeColor "white" $ withFillColor "black" $ mkGroup
        [ translate ax ay $ ppNum v ]
      ]
    | v <- nub $ concat [ [a,b,c] | (a,b,c) <- gs]
    , let V2 ax ay = points V.! (v-1)
    ]]
  where
    colored n =
      let c = promotePixel $ turbo (fromIntegral n / fromIntegral (length gs-1))
      in withStrokeColorPixel c . withFillColorPixel c
    ppNum n = cachedNumbers !! n
     --scaleToHeight (circR*1.5) $ center $ latex $ T.pack $ "\\texttt{" ++ show n ++ "}"

cachedNumbers =
  [ scaleToHeight (circleRadius*1.5) $ center $ latex $ T.pack $ "\\texttt{" ++ show n ++ "}"
  | n <- [0 .. ] ]

circleRadius :: Double
circleRadius = 0.05

-- Anticlockwise. No duplicate vertices. length >= 3
type Polygon = [V2 Double]
type P = V2 Double

-- T = (U, G)
-- G = [Polygon]
-- U = nub $ concat G

findStarNeighbours :: Eq a => [(a,a,a)] -> a -> [(a, a)]
findStarNeighbours allTrig self =
  [ (b,c)
  | (a,b,c) <- allTrig
  , self == a
  ] ++
  [ (c,a)
  | (a,b,c) <- allTrig
  , self == b
  ] ++
  [ (a,b)
  | (a,b,c) <- allTrig
  , self == c
  ]

isInterior :: Eq a => [(a, a)] -> Bool
isInterior = isJust . getExteriorPoly

getExteriorPoly :: Eq a => [(a, a)] -> Maybe [a]
getExteriorPoly [] = Nothing
getExteriorPoly ((a,b):rest) = worker [a] a b rest
  where
    worker acc start this [] = do
      guard (start == this)
      return (reverse acc)
    worker acc start this xs =
      case lookup this xs of
        Just next -> worker (this:acc) start next (delete (this,next) xs)
        Nothing   ->
          case lookup this (map swap xs) of
            Just next -> worker (this:acc) start next (delete (next, this) xs)
            Nothing   -> Nothing

convexInterpolate :: RelMeshPair -> Double -> RelMesh
convexInterpolate (RelMeshPair static edges leftM leftB rightM rightB) t =
  RelMesh
  { relMeshStatic = static
  , relMeshEdges  = edges
  , relMeshA      = Matrix.scale (1-t) leftM +
                    Matrix.scale t rightM
  , relMeshB      = Matrix.scale (1-t) leftB +
                    Matrix.scale t rightB
  }

solveMesh :: RelMesh -> Mesh
solveMesh (RelMesh static edges m b) =
  case linearSolve m b of
    Nothing -> error "Failed to solve mesh"
    Just ret ->
      Mesh (static <> V.fromList (worker (toLists ret))) edges
  where
    worker []             = []
    worker ([x]:[y]:rest) = V2 x y : worker rest
    worker _              = error "invalid result"

mkRelative :: Mesh -> RelMesh
mkRelative (Mesh points edges) = RelMesh (V.fromList exteriorPoints) edges mM bM
  where
    mM = (s><s) (concat m)
    bM = (s><1) b
    (s,exterior, (m, b)) = toParameters points edges
    exteriorPoints =
      [ points V.! (i-1)
      | i <- exterior
      ]

mkRelativePair :: MeshPair -> RelMeshPair
mkRelativePair (MeshPair p1 p2 edges) =
  let RelMesh static _ leftM leftB = mkRelative (Mesh p1 edges)
      RelMesh _ _ rightM rightB = mkRelative (Mesh p2 edges)
  in RelMeshPair static edges leftM leftB rightM rightB

toParameters points groups = (length interior*2,exterior,unzip $ concat
  [ let lst = [(if i == j then -1 else t)
              | j <- interior
              , let t = fromMaybe 0 $ lookup (i,j) lam_ij_cache
              ]
        pos = negate $ sum
           [ pj ^* t
           | j <- exterior
           , let t = fromMaybe 0 $ lookup (i,j) lam_ij_cache
                 pj = points V.! (j-1)
           ]
    in [ (dupX lst, pos ^. _x)
       , (dupY lst, pos ^. _y)]
  | i <- interior ])
  where
    lam_ij_cache = lam_ij points groups
    dupX []     = []
    dupX (x:xs) = x:0:dupX xs
    dupY []     = []
    dupY (x:xs) = 0:x:dupY xs
    (interior, exterior) =
      partition (isInterior . findStarNeighbours groups) [1 .. length points]

lam_ij points groups =
  [ ((i, j), t)
  | i <- [1..V.length points]
  , (j, t) <- lam_j points groups i
  ]
lam_j points groups p =
  [ (nP, sum [ t | (j,k,t) <- mu, j == nP ] / fromIntegral (length nPoints))
  | let n = findStarNeighbours groups p
        nPoints = fromMaybe [] $ getExteriorPoly n
        mu = calcMu points groups p
  , nP <- nPoints ]
calcMu points groups p = concat
    [ [ (nP, nP, t1)
      , (a, nP, t2)
      , (b, nP, t3) ]
      -- (nP, a, b)
    | {-p <- [1..length points]-}
      let selfVert = points V.! (p-1)
          n = findStarNeighbours groups p
          nPoints :: [Int]
          nPoints = fromMaybe [] $ getExteriorPoly n
    , (i, nP) <- zip [1 .. ] nPoints
    , let vert = points V.! (nP-1)
    , let line = (vert, selfVert)
    , let (a,b,aP,bP) = head $
            [ (a,b,aP,bP)
            | (a,b) <- n
            , let aP = points V.! (a-1)
                  bP = points V.! (b-1)
                  segment = (points V.! (a-1), points V.! (b-1))
            , case rayIntersect line segment of
                Nothing -> False
                Just u  -> isBetween u segment
            , a /= nP
            , b /= nP ]
    -- , b == (nPoints ++ nPoints) !! i
    , let (t1,t2,t3) = barycentricCoords vert aP bP selfVert
    ]


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric #-}
-- |
-- Module    : Statistics.Distribution.Gamma
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The gamma distribution.  This is a continuous probability
-- distribution with two parameters, /k/ and ϑ. If /k/ is
-- integral, the distribution represents the sum of /k/ independent
-- exponentially distributed random variables, each of which has a
-- mean of ϑ.

module Statistics.Distribution.Gamma
    (
      GammaDistribution
    -- * Constructors
    , gammaDistr
    , gammaDistrE
    , improperGammaDistr
    , improperGammaDistrE
    -- * Accessors
    , gdShape
    , gdScale
    ) where

import Control.Applicative
import Data.Aeson           (FromJSON(..), ToJSON, Value(..), (.:))
import Data.Binary          (Binary(..))
import Data.Data            (Data, Typeable)
import GHC.Generics         (Generic)
import Numeric.MathFunctions.Constants (m_pos_inf, m_NaN, m_neg_inf)
import Numeric.SpecFunctions (incompleteGamma, invIncompleteGamma, logGamma, digamma)
import qualified System.Random.MWC.Distributions as MWC

import Statistics.Distribution.Poisson.Internal as Poisson
import qualified Statistics.Distribution as D
import Statistics.Internal

-- | The gamma distribution.
data GammaDistribution = GD {
      gdShape :: {-# UNPACK #-} !Double -- ^ Shape parameter, /k/.
    , gdScale :: {-# UNPACK #-} !Double -- ^ Scale parameter, ϑ.
    } deriving (Eq, Typeable, Data, Generic)

instance Show GammaDistribution where
  showsPrec i (GD k theta) = defaultShow2 "improperGammaDistr" k theta i
instance Read GammaDistribution where
  readPrec = defaultReadPrecM2 "improperGammaDistr" improperGammaDistrE

instance ToJSON GammaDistribution
instance FromJSON GammaDistribution where
  parseJSON (Object v) = do
    k     <- v .: "gdShape"
    theta <- v .: "gdScale"
    maybe (fail $ errMsgI k theta) return $ improperGammaDistrE k theta
  parseJSON _ = empty

instance Binary GammaDistribution where
  put (GD x y) = put x >> put y
  get = do
    k     <- get
    theta <- get
    maybe (fail $ errMsgI k theta) return $ improperGammaDistrE k theta

-- | Create gamma distribution. Both shape and scale parameters must
-- be positive.
gammaDistr :: Double            -- ^ Shape parameter. /k/
           -> Double            -- ^ Scale parameter, ϑ.
           -> GammaDistribution
gammaDistr k theta
  = maybe (error $ errMsg k theta) id $ gammaDistrE k theta

errMsg :: Double -> Double -> String
errMsg k theta
  =  "Statistics.Distribution.Gamma.gammaDistr: "
  ++ "k=" ++ show k
  ++ "theta=" ++ show theta
  ++ " but must be positive"

-- | Create gamma distribution. Both shape and scale parameters must
-- be positive.
gammaDistrE :: Double            -- ^ Shape parameter. /k/
            -> Double            -- ^ Scale parameter, ϑ.
            -> Maybe GammaDistribution
gammaDistrE k theta
  | k > 0 && theta > 0 = Just (GD k theta)
  | otherwise          = Nothing

-- | Create gamma distribution. Both shape and scale parameters must
-- be non-negative.
improperGammaDistr :: Double            -- ^ Shape parameter. /k/
                   -> Double            -- ^ Scale parameter, ϑ.
                   -> GammaDistribution
improperGammaDistr k theta
  = maybe (error $ errMsgI k theta) id $ improperGammaDistrE k theta

errMsgI :: Double -> Double -> String
errMsgI k theta
  =  "Statistics.Distribution.Gamma.gammaDistr: "
  ++ "k=" ++ show k
  ++ "theta=" ++ show theta
  ++ " but must be non-negative"

-- | Create gamma distribution. Both shape and scale parameters must
-- be non-negative.
improperGammaDistrE :: Double            -- ^ Shape parameter. /k/
                    -> Double            -- ^ Scale parameter, ϑ.
                    -> Maybe GammaDistribution
improperGammaDistrE k theta
  | k >= 0 && theta >= 0 = Just (GD k theta)
  | otherwise            = Nothing

instance D.Distribution GammaDistribution where
    cumulative = cumulative

instance D.ContDistr GammaDistribution where
    density    = density
    logDensity (GD k theta) x
      | x <= 0    = m_neg_inf
      | otherwise = log x * (k - 1) - (x / theta) - logGamma k - log theta * k
    quantile   = quantile

instance D.Variance GammaDistribution where
    variance (GD a l) = a * l * l

instance D.Mean GammaDistribution where
    mean (GD a l) = a * l

instance D.MaybeMean GammaDistribution where
    maybeMean = Just . D.mean

instance D.MaybeVariance GammaDistribution where
    maybeStdDev   = Just . D.stdDev
    maybeVariance = Just . D.variance

instance D.MaybeEntropy GammaDistribution where
  maybeEntropy (GD a l)
    | a > 0 && l > 0 =
      Just $
      a
      + log l
      + logGamma a
      + (1-a) * digamma a
    | otherwise = Nothing

instance D.ContGen GammaDistribution where
    genContVar (GD a l) = MWC.gamma a l

density :: GammaDistribution -> Double -> Double
density (GD a l) x
  | a < 0 || l <= 0   = m_NaN
  | x <= 0            = 0
  | a == 0            = if x == 0 then m_pos_inf else 0
  | x == 0            = if a < 1 then m_pos_inf else if a > 1 then 0 else 1/l
  | a < 1             = Poisson.probability (x/l) a * a / x
  | otherwise         = Poisson.probability (x/l) (a-1) / l

cumulative :: GammaDistribution -> Double -> Double
cumulative (GD k l) x
  | x <= 0    = 0
  | otherwise = incompleteGamma k (x/l)

quantile :: GammaDistribution -> Double -> Double
quantile (GD k l) p
  | p == 0         = 0
  | p == 1         = 1/0
  | p > 0 && p < 1 = l * invIncompleteGamma k p
  | otherwise      =
    error $ "Statistics.Distribution.Gamma.quantile: p must be in [0,1] range. Got: "++show p


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}
import           Control.Monad
import           Control.Monad.Random
import           Data.List                    (foldl')

import qualified Data.ByteString              as B
import           Data.Semigroup               ((<>))
import           Data.Serialize

import           GHC.TypeLits

import qualified Numeric.LinearAlgebra.Static as SA

import           Options.Applicative

import           Grenade

-- The definition for our simple feed forward network.
-- The type level lists represents the layers and the shapes passed through the layers.
-- One can see that for this demonstration we are using relu, tanh and logit non-linear
-- units, which can be easily substitute for each other in and out.
--
-- With around 100000 examples, this should show two clear circles which have been learned by the network.
type FFNet = Network '[ FullyConnected 2 40, Tanh, FullyConnected 40 30, Relu, FullyConnected 30 20, Relu, FullyConnected 20 10, Relu, FullyConnected 10 1, Logit ]
                     '[ 'D1 2, 'D1 40, 'D1 40, 'D1 30, 'D1 30, 'D1 20, 'D1 20, 'D1 10, 'D1 10, 'D1 1, 'D1 1]

randomNet :: IO FFNet
randomNet = randomNetworkInitWith HeEtAl  -- you might want to try `Xavier` or `UniformInit` instead of `HeEtAl`

netTrain :: FFNet -> LearningParameters -> Int -> IO FFNet
netTrain net0 rate n = do
    inps <- replicateM n $ do
      s  <- getRandom
      return $ S1D $ SA.randomVector s SA.Uniform * 2 - 1
    let outs = flip map inps $ \(S1D v) ->
                 if v `inCircle` (fromRational 0.50, 0.50)  || v `inCircle` (fromRational (-0.50), 0.50)
                   then S1D $ fromRational 1
                   else S1D $ fromRational 0

    let trained = foldl' trainEach net0 (zip inps outs)
    return trained

  where trainEach !network (i,o) = train rate network i o

netLoad :: FilePath -> IO FFNet
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $ runGet (get :: Get FFNet) modelData

renderClass :: IO ()
renderClass = do
  let testIns = [ [ (x,y)  | x <- [0..50] ]
                           | y <- [0..20] ]
  let outMat  = fmap (fmap (\(x,y) -> (render (x/25-1) (y/10-1)))) testIns
  putStrLn $ unlines outMat

  where
    render x y  | x == 0 && y == 0 = '+'
                | y == 0 = '-'
                | x == 0 = '|'
                | otherwise = let v = SA.vector [x,y] :: SA.R 2
                              in if v `inCircle` (fromRational 0.50, 0.50)  || v `inCircle` (fromRational (-0.50), 0.50)
                                 then '1'
                                 else ' '

netScore :: FFNet -> IO ()
netScore network = do
    let testIns = [ [ (x,y)  | x <- [0..50] ]
                             | y <- [0..20] ]
        outMat  = fmap (fmap (\(x,y) -> (render (x/25-1) (y/10-1) . normx) $ runNet network (S1D $ SA.vector [x / 25 - 1,y / 10 - 1]))) testIns
    putStrLn $ unlines outMat

  where
    render x y n'  | x == 0 && y == 0 = '+'
                   | y == 0 = '-'
                   | x == 0 = '|'
                   | n' <= 0.2  = ' '
                   | n' <= 0.4  = '.'
                   | n' <= 0.6  = '-'
                   | n' <= 0.8  = '='
                   | otherwise = '#'

normx :: S ('D1 1) -> Double
normx (S1D r) = SA.mean r

testValues :: FFNet -> IO ()
testValues network = do
  inps <- replicateM 1000 $ do
      s  <- getRandom
      return $ S1D $ SA.randomVector s SA.Uniform * 2 - 1
  let outs = flip map inps $ \(S1D v) ->
                 if v `inCircle` (fromRational 0.50, 0.50)  || v `inCircle` (fromRational (-0.50), 0.50)
                   then 1 :: Integer
                   else 0
  let ress = zip outs (map (round . normx . runNet network) inps)
      correct = length $ filter id $ map (uncurry (==)) ress
      incorrect = length $ filter id $ map (uncurry (/=)) ress
      falsePositives = length $ filter id $ map (uncurry (\shd nn -> shd == 0 && nn == 1)) ress
      falseNegatives = length $ filter id $ map (uncurry (\shd nn -> shd == 1 && nn == 0)) ress
  putStr $ show correct  ++ " | "
  putStr $ show incorrect ++ " | "
  putStr $ show falsePositives ++ " | "
  putStrLn $ show falseNegatives ++ " | "

inCircle :: KnownNat n => SA.R n -> (SA.R n, Double) -> Bool
v `inCircle` (o, r) = SA.norm_2 (v - o) <= r

data FeedForwardOpts = FeedForwardOpts Int LearningParameters

feedForward' :: Parser FeedForwardOpts
feedForward' =
  FeedForwardOpts <$> option auto (long "examples" <> short 'e' <> value 1000)
                  <*> (LearningParameters
                      <$> option auto (long "train_rate" <> short 'r' <> value 0.005)
                      <*> option auto (long "momentum" <> value 0.0)
                      <*> option auto (long "l2" <> value 0.0005)
                      )

main :: IO ()
main = do
  FeedForwardOpts examples rate <- execParser (info (feedForward' <**> helper) idm)

  putStrLn "| Nr | Correct | Incorrect | FalsePositives | FalseNegatives |"
  putStrLn "--------------------------------------------------------------"
  let nr = 100 :: Int
  mapM_ (\n -> do
    putStr $ "| " ++ show n  ++ " | "
    net0 <- randomNet
    net <- netTrain net0 rate examples
    -- netScore net
    testValues net) [1..nr]


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}
---------------------------------
-- |
-- Module      :  HMM
-- Copyright   :  (C) Kaushik Chakraborty, 2019
-- License     :  Apache v2 (see the file LICENSE)
-- Maintainer  :  Kaushik Chakraborty <git@kaushikc.org>
-- Stability   :  experimental
--
---------------------------------

module HMM
  (
    HMM(..),
    -- ** Smart Constructors
    mkHMM_, mkHMM, mkHMM1, sensorDiagonal,
    -- * Utility Functions
    normalise, inverse, reverseV, extract1, (!**!), unitColumn,
    -- ** Converting to/from HMatrix matrices
    toHM, fromHM,
    -- * Inference Algorithms
    -- ** Forward-Backward
    forward, backward, forwardBackward,
    -- ** Fixed Lag Smoothing
    fixedLagSmoothing,
    -- * Sample HMM Models
    umbrellaHMM,
    -- * Miscellaneous
    -- ** Type Synonyms
    R, TransitionModel, SensorDiagonal, SensorModel, Message, Distribution,
    -- ** Data Types
    Persistent(..),
    -- ** Test Functions
    runFLSAlgo, runFBAlgo,
    -- ** Lenses
    prior, tModel, sDist, sModel
) where

import Text.Printf (printf)

import Papa
import GHC.TypeNats
import Data.Proxy (Proxy(..))

import Control.Monad.State (State, runState)
import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)

import qualified Data.Vector as Vec

import Linear.V
import Linear

import qualified Numeric.LinearAlgebra.Static as HM
import qualified Numeric.LinearAlgebra as HMat
-----------------------------------------------------------------------------------

type R = Double
type TransitionModel (s :: Nat) a = V s (V s a)
type SensorDiagonal (s  :: Nat) a = V s (V s a)
type SensorModel (t :: Nat) (s :: Nat) a = V t (SensorDiagonal s a)
type Message (s :: Nat) a = V s (V 1 a)
type Distribution (s :: Nat) a = V s a

-- | Hidden Markov Models
data HMM (s :: Nat) (t :: Nat)  a b = HMM {
  -- | prior distribution which is used as initial forward message
  _prior :: Message s a,
  -- | transition model with @s@ states as @sxs@ matrix
  _tModel :: TransitionModel s a,
  -- | evidence value vector, each index maps to the corresponding evidence values
  _sDist :: V t b,
  -- | sensor model with @t@ evidence values having @sxs@ diagonal matrix capturing their ditributions for each state @s@
  _sModel :: SensorModel t s a
  } deriving (Show)

makeLenses ''HMM

mkHMM_ :: (KnownNat s, KnownNat t) => Maybe (Message s R) -- ^ Prior distribution on the initial state, @P(X0)@. If nothing then considered @(0.5,0.5)@
       -> TransitionModel s R -- ^ Transition Model as @sxs@ matrix, @s@ being the number of states
       -> V t b -- ^ evidence values vector map
       -> V t (V s R) -- ^ a @txs@ matrix for each evidence value @t@, a vector capturing conditional probabilities for each state @s@
       -> HMM s t R b
mkHMM_ mp xs evs = let p = fromMaybe (V $ Vec.replicate (dim xs) 0.5) mp
              in
              HMM p xs evs . (scaled <$>)

-- | A default HMM where both transition model states and evidence variables have boolean support
mkHMM :: KnownNat s => TransitionModel s R -> V 2 (V s R) -> HMM s 2 R Bool
mkHMM ts = mkHMM_ Nothing ts (V $ Vec.fromList [True, False])

-- | A default HMM with one state variable having boolean support
mkHMM1 :: V2 (V2 R) -> V2 (V2 R) -> HMM 2 2 R Bool
mkHMM1 ts ss = let f = (_V #) . over mapped toV
               in
                 uncurry mkHMM $ over both f (ts , ss)

-- | create a @sxs@ diagonal matrix with corresponding posterior probabilities from the sensor model of the @HMM@ for an input sensor value
sensorDiagonal :: (KnownNat s, KnownNat t, Eq b)
               => HMM s t a b
               -> b -- ^ sensor value
               -> Maybe (SensorDiagonal s a)
sensorDiagonal hmm e = findIndexOf (sDist.folded) (== e) hmm >>= (\i -> hmm ^? sModel . ix i)

-- | Multiply each number by a constant such that the sum is 1.0
normalise :: (Fractional a, Foldable f, Functor f) => f a -> f a
normalise xs | null xs = xs
             | otherwise =
                 let
                   s = sum xs
                 in
                   (/ s) <$> xs

-- | inverse a square matrix
inverse :: forall s. (KnownNat s) => V s (V s R) -> V s (V s R)
inverse  = fromHM . HM.inv . toHM

-- | reverse contents of a 'Linear.V' vector
reverseV :: forall s a. (KnownNat s) => V s a -> V s a
reverseV = V . Vec.fromList . Papa.reverse . toList

extract1 :: (KnownNat s) => V s (V 1 a) -> V s a
extract1 =  V . foldMap toVector

-- | Pointwise product of 2 @mxn@ 'V' vectors
infix 8 !**!
(!**!) :: (KnownNat m, KnownNat n) => V m (V n R) -> V m (V n R) -> V m (V n R)
as !**! bs = fromHM $ toHM as * toHM bs

-- | an unit vector having @s@ columns
unitColumn :: forall s. KnownNat s => Message s R
unitColumn = V $ Vec.replicate (fromIntegral $ natVal (Proxy :: Proxy s))  (toV $ V1 1.0)

-- | take a @sxt@ 'Linear.V.V' matrix and give corresponding 'Numeric.LinearAlgebra.Static.L' version
toHM :: (KnownNat s, KnownNat t) => V s (V t R) -> HM.L s t
toHM = HM.matrix . foldMap (Vec.toList . toVector)

-- | take a @sxt@ 'Numeric.LinearAlgebra.Static.L' matrix and give corresponding 'Linear.V.V' version
fromHM :: (KnownNat s, KnownNat t) => HM.L s t -> V s (V t R)
fromHM m = V $ V <$> Vec.fromList (fmap Vec.fromList $ HMat.toLists $ HM.extract m)

-- | Filtering message propagated forward
forward :: (KnownNat s) => HMM s t R b -> SensorDiagonal s R  -> Message s R -> Message s R
forward hmm ot f = normalise (ot !*! Linear.transpose (hmm ^. tModel) !*! f)

-- | Smoothing message propagated backward
backward :: (KnownNat s) => HMM s t R b -> SensorDiagonal s R -> Message s R -> Message s R
backward hmm ok1 bk2 = (hmm ^. tModel) !*! ok1 !*! bk2

-- | The forward–backward algorithm for smoothing: computing posterior prob- abilities of a sequence of states given a sequence of observations
forwardBackward :: forall s t u b . (KnownNat t, KnownNat s, KnownNat u, Eq b)
                => HMM s u R b -- ^ HMM model as a way to implement
                -> V t b -- ^ list of evidences for each time step
                -> V t (Distribution s R)
forwardBackward hmm evs = let
  -- reifying the number of evidences
  tNat = dim evs
  -- forward messages from time t .. 0
  fv :: V (t + 1) (Message s R)
  fv = V $ Vec.fromList $ foldl' (\m@(x:_) e -> forward hmm (sensorDiagonal hmm e ^?! _Just) x : m) [hmm ^. prior] evs
  -- getting rid of the prior message from the end of the list
  -- so now forward messages vector is from time t .. 1
  fv_0 :: V t (Message s R)
  fv_0 = V $ Vec.fromList $ fv ^.. taking tNat traversed

  --  backward messages from time 1 .. t
  bs :: V t (Message s R)
  bs = V $ Vec.fromList $ foldl' (\m@(x:_) e -> backward hmm (sensorDiagonal hmm e ^?! _Just) x : m) [unitColumn] $ reverseV evs
  -- reversing the backward messages
  -- so now the backward messages vector is from time t .. 1
  revBs :: V t (Message s R)
  revBs = reverseV bs
  in
    -- smoothing probabilities in reverse order of the list of evidences
    -- i.e. starting from time t .. 1
    liftA2 (\f b' -> extract1 $ normalise $ f !**! b') fv_0 revBs

-- | Persistent State
data Persistent (s :: Nat) a b d= Persistent {
  _t :: d, -- ^ current time
  _f_msg :: Message s a, -- ^ the forward message @P(Xt|e1:t)@
  _b :: V s (V s a), -- ^ the @d@-step backward transformation matrix
  _e_td_t :: Vec.Vector b -- ^ double-ended list of evidence from @t − d@ to @t@
  } deriving (Show)

makeLenses ''Persistent

-- | Initial persistent state where
--
--  * t = 1
--
--  * f_msg  = hmm.prior
--
--  * b = identity matrix
--
--  * e_td_t = empty vector
persistentInit :: forall s b d. (KnownNat s, Integral d) => Message s R -> Persistent s R b d
persistentInit p = Persistent { _t = 1, _f_msg = p, _b = identity, _e_td_t = Vec.empty}

-- | Online Algorithm for smoothing with a fixed time lag of @d@ steps
fixedLagSmoothing :: forall s u b d. (KnownNat s, KnownNat u, Eq b, Integral d)
                  => HMM s u R b -- ^ HMM model
                  -> d -- ^ length of lag
                  -> b -- ^ evidence at time @t@
                  -> MaybeT (State (Persistent s R b d)) (Distribution s R)
fixedLagSmoothing hmm d e = do
  e_td_t %= flip Vec.snoc e

  o_t <- uses e_td_t $ (^?! _Just) . sensorDiagonal hmm . Vec.last

  t' <- use t
  if t' > d then
    do
      e_td_t %= Vec.drop 1
      o_tmd <- uses e_td_t $ (^?! _Just) . sensorDiagonal hmm . Vec.head
      f_msg %= forward hmm o_tmd
      b %= \b' -> inverse o_tmd !*! inverse (hmm ^. tModel) !*! b' !*! (hmm ^. tModel) !*! o_t
      t += 1

      (f'', b'') <- liftA2 (,) (use f_msg) (use b)

      return $ extract1 $ normalise (f'' !**! (b'' !*! unitColumn))
  else
    do
      b %= \b' -> b' !*! (hmm ^. tModel) !*! o_t
      t += 1
      mzero

umbrellaHMM :: HMM 2 2 R Bool
umbrellaHMM = mkHMM1 (V2 (V2 0.7 0.3) (V2 0.3 0.7)) (V2 (V2 0.9 0.2) (V2 0.1 0.8))

runFLSAlgo  :: [Bool] -> Integer -> [Maybe (Distribution 2 String)]
runFLSAlgo bs d = let hmm = mkHMM1 (V2 (V2 0.7 0.3) (V2 0.3 0.7)) (V2 (V2 0.9 0.2) (V2 0.1 0.8))
                      initState = persistentInit (V $ Vec.fromList [V $ Vec.singleton 0.5, V $ Vec.singleton 0.5]) :: Persistent 2 R Bool Integer
                      algo = fixedLagSmoothing hmm d
                      a :: ([Maybe (Distribution 2 R)], Persistent 2 R Bool Integer)
                      a = foldl' (\(rs, s) x -> runState (runMaybeT $ algo x) s & _1 #%~ ((rs ++) . pure)) ([], initState) bs
                  in
                    (a ^. _1) & traverse.traverse.traverse %~ \(x :: R) -> printf "%.3f" x :: String

runFBAlgo :: [Bool] -> [Distribution 2 R]
runFBAlgo bs = reifyVectorNat (Vec.fromList bs) (toList . forwardBackward umbrellaHMM)


{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Presburger #-}

{- | Algorithms for zero-dimensional ideals.

   Since 0.4.0.0
-}
module Algebra.Algorithms.ZeroDim
  ( -- * Root finding for zero-dimensional ideal
    solveM,
    solve',
    solveViaCompanion,
    solveLinear,

    -- * Radical computation
    radical,
    isRadical,

    -- * Converting monomial ordering to Lex using FGLM algorithm
    fglm,
    fglmMap,

    -- ** Internal helper function
    solveWith,
    univPoly,
    reduction,
    matrixRep,
    subspMatrix,
    vectorRep,
  )
where

import Algebra.Algorithms.FGLM
import Algebra.Algorithms.Groebner
import Algebra.Instances ()
import Algebra.Internal hiding (OLt)
import qualified Algebra.Matrix as AM
import Algebra.Prelude.Core
import Algebra.Ring.Polynomial.Quotient
import Control.Lens
import Control.Monad.Loops (whileM_)
import Control.Monad.Random hiding (next)
import Control.Monad.Reader (runReaderT)
import Control.Monad.ST (runST)
import Data.Complex (Complex (..), magnitude)
import Data.Convertible (Convertible, convert)
import qualified Data.Matrix as M
import Data.Maybe (fromJust)
import Data.Reflection (Reifies)
import Data.STRef.Strict (newSTRef)
import qualified Data.Sized as SV
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Numeric.Algebra as NA
import qualified Numeric.LinearAlgebra as LA
import qualified Prelude as P

{- | Finds complex approximate  roots of given zero-dimensional ideal,
   using randomized altorithm.

   See also @'solve''@ and @'solveViaCompanion'@.
-}
solveM ::
  forall m r ord n.
  ( Normed r
  , Ord r
  , MonadRandom m
  , Field r
  , CoeffRing r
  , KnownNat n
  , IsMonomialOrder n ord
  , Convertible r Double
  , 0 < n
  ) =>
  Ideal (OrderedPolynomial r ord n) ->
  m [Sized n (Complex Double)]
solveM ideal =
  {-# SCC "solveM" #-}
  withKnownNat (sSucc (sNat :: SNat n)) $
    reifyQuotient (radical ideal) $ \pxy ->
      case standardMonomials' pxy of
        Just bs -> step 10 (length bs)
        Nothing -> error "Given ideal is not zero-dimensional"
  where
    step bd len =
      {-# SCC "solveM/step" #-}
      do
        coeffs <-
          {-# SCC "solveM/coeff-gen" #-}
          replicateM (sNatToInt (sSucc (sNat :: SNat n))) $ getRandomR (- bd, bd)
        let vars = one : SV.toList allVars
            f = sum $ zipWith (.*.) (map (NA.fromInteger :: Integer -> r) coeffs) vars
        case solveWith f ideal of
          Nothing -> step (bd * 2) len
          Just sols -> return sols

{- | @solveWith f is@ finds complex approximate roots of the given zero-dimensional @n@-variate polynomial system @is@,
   using the given relatively prime polynomial @f@.
-}
solveWith ::
  forall r n ord.
  ( DecidableZero r
  , Normed r
  , Ord r
  , Field r
  , CoeffRing r
  , 0 < n
  , IsMonomialOrder n ord
  , KnownNat n
  , Convertible r Double
  ) =>
  OrderedPolynomial r ord n ->
  Ideal (OrderedPolynomial r ord n) ->
  Maybe [Sized n (Complex Double)]
solveWith f0 i0 =
  {-# SCC "solveWith" #-}
  withKnownNat (sSucc (sNat :: SNat n)) $
    reifyQuotient (radical i0) $ \pxy ->
      let ideal = gBasis' pxy
          Just base = map (leadingMonomial . quotRepr) <$> standardMonomials' pxy
       in case {-# SCC "findOne" #-} elemIndex one base of
            Nothing -> Just []
            Just cind ->
              let f = modIdeal' pxy f0
                  vars =
                    sortBy (flip $ comparing snd) $
                      map (\on -> (on, leadingMonomial $ var on `asTypeOf` f0)) $
                        enumOrdinal $ sArity' f0
                  inds = flip map vars $
                    second $ \b ->
                      case findIndex (== b) base of
                        Just ind -> Right ind
                        Nothing ->
                          let Just g = find ((== b) . leadingMonomial) ideal
                              r = leadingCoeff g
                              answer = mapCoeff toComplex $ injectCoeff (recip r) * (toPolynomial (leadingTerm g) - g)
                           in Left answer
                  mf = AM.fromLists $ map (map toComplex) $ matrixRep f
                  (_, evecs) = LA.eig $ LA.tr mf
                  calc vec =
                    {-# SCC "calc" #-}
                    let c = vec LA.! cind
                        phi (idx, Right nth) acc = acc & ix idx .~ (vec LA.! nth) P./ c
                        phi (idx, Left g) acc = acc & ix idx .~ substWith (*) acc g
                     in if c == 0
                          then Nothing
                          else Just $ foldr ({-# SCC "rewrite-answer" #-} phi) (SV.replicate (sArity' f0) (error "indec!")) inds
               in sequence $ map calc $ LA.toColumns evecs

{- | @'solve'' err is@ finds numeric approximate root of the
   given zero-dimensional polynomial system @is@,
   with error <@err@.

   See also @'solveViaCompanion'@ and @'solveM'@.
-}
solve' ::
  forall r n ord.
  ( Field r
  , CoeffRing r
  , KnownNat n
  , 0 < n
  , IsMonomialOrder n ord
  , Convertible r Double
  ) =>
  Double ->
  Ideal (OrderedPolynomial r ord n) ->
  [Sized n (Complex Double)]
solve' err ideal =
  reifyQuotient ideal $ \ii ->
    if gBasis' ii == [one]
      then []
      else
        let vs =
              map (nub . LA.toList . LA.eigenvalues . AM.fromLists . map (map toComplex) . matrixRep . modIdeal' ii) $
                SV.toList allVars
            mul p q = toComplex p * q
         in [ xs
            | xs0 <- sequence vs
            , let xs = SV.unsafeFromList' xs0
            , all ((< err) . magnitude . substWith mul xs) $ generators ideal
            ]

{- | Given a zero-dimensional ideal \(I\),
   \('subspMatrix' i I\) computes a multiplication matrix
   \(m_{x_i} \mathbin{\upharpoonright} V_i \) by \(x_i\)
   restricted to the subspace \(V_i = \mathop{\mathrm{span}}(\left\{\ x_i^n \ \middle|\ 1 \leq n \right\})\) of \(k[\mathbf{X}]\).
-}
subspMatrix ::
  forall r n ord.
  (Show r, Ord r, Field r, CoeffRing r, KnownNat n, IsMonomialOrder n ord) =>
  Ordinal n ->
  Ideal (OrderedPolynomial r ord n) ->
  M.Matrix r
subspMatrix on ideal =
  let poly = univPoly on ideal
      v = var on :: OrderedPolynomial r ord n
      dim = totalDegree' poly
      cfs = [negate $ coeff (leadingMonomial $ pow v (j :: Natural)) poly | j <- [0 .. fromIntegral (dim - 1)]]
      leftTops = M.fromLists [replicate (dim - 1) zero]
      leftBots = fmap unwrapAlgebra (M.identity (dim - 1))
      lfts =
        leftTops
          M.<-> leftBots
      rights = M.colVector (V.fromList cfs)
   in lfts M.<|> rights

{- | @'solveViaCompanion' err is@ finds numeric approximate root of the
   given zero-dimensional polynomial system @is@,
   with error <@err@.

   See also @'solve''@ and @'solveM'@.
-}
solveViaCompanion ::
  forall r ord n.
  (Show r, Ord r, Field r, CoeffRing r, KnownNat n, IsMonomialOrder n ord, Convertible r Double) =>
  Double ->
  Ideal (OrderedPolynomial r ord n) ->
  [Sized n (Complex Double)]
solveViaCompanion err ideal =
  if calcGroebnerBasis ideal == [one]
    then []
    else
      let vs =
            map (nub . LA.toList . LA.eigenvalues . LA.fromLists . matToLists . fmap toComplex . flip subspMatrix ideal) $
              enumOrdinal (sNat :: SNat n)
          mul p q = toComplex p * q
       in [ xs
          | xs0 <- sequence vs
          , let xs = SV.unsafeFromList' xs0
          , all ((< err) . magnitude . substWith mul xs) $ generators ideal
          ]

matToLists :: M.Matrix a -> [[a]]
matToLists mat = [V.toList $ M.getRow i mat | i <- [1 .. M.nrows mat]]

matrixRep ::
  ( DecidableZero t
  , Eq t
  , Field t
  , KnownNat n
  , IsMonomialOrder n order
  , Reifies ideal (QIdeal (OrderedPolynomial t order n))
  ) =>
  Quotient (OrderedPolynomial t order n) ideal ->
  [[t]]
matrixRep f =
  {-# SCC "matrixRep" #-}
  case standardMonomials of
    Just [] -> []
    Just bases ->
      let anss = map (quotRepr . (f *)) bases
       in transpose $ map (\a -> map (flip coeff a . leadingMonomial . quotRepr) bases) anss
    Nothing -> error "Not finite dimension"

toComplex :: Convertible a Double => a -> Complex Double
toComplex a = convert a :+ 0

-- | Calculates n-th reduction of f: @f `div` < f, ∂_{x_n} f >@.
reduction ::
  (CoeffRing r, KnownNat n, IsMonomialOrder n ord, Field r) =>
  Ordinal n ->
  OrderedPolynomial r ord n ->
  OrderedPolynomial r ord n
reduction on f =
  {-# SCC "reduction" #-}
  let df = {-# SCC "differentiate" #-} diff on f
   in snd $ head $ f `divPolynomial` calcGroebnerBasis (toIdeal [f, df])

-- | Calculate the monic generator of \( k[X_0, ..., X_n] \cap k[X_i]\).
univPoly ::
  forall r ord n.
  (Ord r, Field r, CoeffRing r, KnownNat n, IsMonomialOrder n ord) =>
  Ordinal n ->
  Ideal (OrderedPolynomial r ord n) ->
  OrderedPolynomial r ord n
univPoly nth ideal =
  {-# SCC "univPoly" #-}
  reifyQuotient ideal $ \pxy ->
    if gBasis' pxy == [one]
      then one
      else
        let x = var nth
            monomDeg =
              fromIntegral $
                length $
                  fromJust $
                    standardMonomials' pxy
            p0 : pows =
              [ fmap WrapAlgebra $ vectorRep $ modIdeal' pxy (pow x i)
              | i <- [0 :: Natural .. monomDeg]
              ]
            step m ~(p : ps) =
              {-# SCC "univPoly/step" #-}
              case solveLinear m p of
                Nothing -> {-# SCC "recur" #-} step ({-# SCC "consCol" #-} m M.<|> M.colVector p) ps
                Just ans ->
                  let cur = fromIntegral $ V.length ans :: Natural
                   in {-# SCC "buildRelation" #-}
                      pow x cur
                        - sum
                          ( zipWith
                              (.*.)
                              (fmap unwrapAlgebra $ V.toList ans)
                              [pow x i | i <- [0 :: Natural .. cur P.- 1]]
                          )
         in step (M.colVector p0) pows

-- | Solves linear system. If the given matrix is degenerate, this returns @Nothing@.
solveLinear ::
  (Ord r, P.Fractional r) =>
  M.Matrix r ->
  V.Vector r ->
  Maybe (V.Vector r)
solveLinear mat vec =
  {-# SCC "solveLinear" #-}
  if ({-# SCC "uRank" #-} uRank u) < uRank u' || M.diagProd u == 0 || uRank u < M.ncols mat
    then Nothing
    else
      let ans = M.getCol 1 $ p P.* M.colVector vec
          lsol = {-# SCC "solveL" #-} solveL ans
          cfs = M.getCol 1 $ q P.* M.colVector ({-# SCC "solveU" #-} solveU lsol)
       in Just cfs
  where
    Just (u, l, p, q, _, _) = M.luDecomp' mat
    Just (u', _, _, _, _, _) = M.luDecomp' (mat M.<|> M.colVector vec)
    uRank = V.foldr (\a acc -> if a /= 0 then acc + 1 else acc) (0 :: Int) . M.getDiag
    solveL v = V.create $ do
      let stop = min (M.ncols l) (M.nrows l)
      mv <- MV.replicate (M.ncols l) 0
      forM_ [0 .. stop - 1] $ \i -> do
        MV.write mv i $ v V.! i
        forM_ [0, 1 .. min (i -1) (M.ncols l - 1)] $ \j -> do
          a <- MV.read mv i
          b <- MV.read mv j
          MV.write mv i $ a P.- (l M.! (i + 1, j + 1)) P.* b
      return mv
    solveU v = V.create $ do
      let stop = min (M.ncols u) (M.nrows u)
      mv <- MV.replicate (M.ncols u) 0
      forM_ [stop - 1, stop - 2 .. 0] $ \i -> do
        MV.write mv i $ v V.! i
        forM_ [i + 1, i + 2 .. M.ncols u -1] $ \j -> do
          a <- MV.read mv i
          b <- MV.read mv j
          MV.write mv i $ a P.- (u M.! (i + 1, j + 1)) P.* b
        a0 <- MV.read mv i
        MV.write mv i $ a0 P./ (u M.! (i + 1, i + 1))
      return mv

-- | Calculate the radical of the given zero-dimensional ideal.
radical ::
  forall r ord n.
  ( Ord r
  , CoeffRing r
  , KnownNat n
  , Field r
  , IsMonomialOrder n ord
  ) =>
  Ideal (OrderedPolynomial r ord n) ->
  Ideal (OrderedPolynomial r ord n)
radical ideal =
  {-# SCC "radical" #-}
  let gens = {-# SCC "calcGens" #-} map (\on -> reduction on $ univPoly on ideal) $ enumOrdinal (sNat :: SNat n)
   in toIdeal $ calcGroebnerBasis $ toIdeal $ generators ideal ++ gens

-- | Test if the given zero-dimensional ideal is radical or not.
isRadical ::
  forall r ord n.
  ( Ord r
  , CoeffRing r
  , KnownNat n
  , 0 < n
  , Field r
  , IsMonomialOrder n ord
  ) =>
  Ideal (OrderedPolynomial r ord n) ->
  Bool
isRadical ideal =
  let gens =
        map (\on -> reduction on $ univPoly on ideal) $
          enumOrdinal (sNat :: SNat n)
   in all (`isIdealMember` ideal) gens

-- * FGLM

{- | Calculate the Groebner basis w.r.t. lex ordering of the zero-dimensional ideal using FGLM algorithm.
   If the given ideal is not zero-dimensional this function may diverge.
-}
fglm ::
  ( Ord r
  , KnownNat n
  , Field r
  , IsMonomialOrder n ord
  , 0 < n
  ) =>
  Ideal (OrderedPolynomial r ord n) ->
  ([OrderedPolynomial r Lex n], [OrderedPolynomial r Lex n])
fglm ideal = reifyQuotient ideal $ \pxy ->
  fglmMap (vectorRep . modIdeal' pxy)

-- | Compute the kernel and image of the given linear map using generalized FGLM algorithm.
fglmMap ::
  forall k ord n.
  ( Ord k
  , Field k
  , 0 < n
  , IsMonomialOrder n ord
  , CoeffRing k
  , KnownNat n
  ) =>
  -- | Linear map from polynomial ring.
  (OrderedPolynomial k ord n -> V.Vector k) ->
  -- | The tuple of:
  --
  --     * lex-Groebner basis of the kernel of the given linear map.
  --
  --     * The vector basis of the image of the linear map.
  ( [OrderedPolynomial k Lex n]
  , [OrderedPolynomial k Lex n]
  )
fglmMap l = runST $ do
  env <- FGLMEnv l <$> newSTRef [] <*> newSTRef [] <*> newSTRef Nothing <*> newSTRef one
  flip runReaderT env $ do
    mainLoop
    whileM_ toContinue $ nextMonomial >> mainLoop
    (,) <$> look gLex <*> (map (changeOrder Lex) <$> look bLex)

mainLoop ::
  (DecidableZero r, Ord r, KnownNat n, Field r, IsMonomialOrder n o) =>
  Machine s r o n ()
mainLoop = do
  m <- look monomial
  let f = toPolynomial (one, changeMonomialOrderProxy Proxy m)
  lx <- image f
  bs <- mapM image =<< look bLex
  let mat = foldr (M.<|>) (M.fromList 0 0 []) $ map (M.colVector . fmap WrapAlgebra) bs
      cond
        | null bs =
          if V.all (== zero) lx
            then Just $ V.replicate (length bs) 0
            else Nothing
        | otherwise = solveLinear mat (fmap WrapAlgebra lx)
  case cond of
    Nothing -> do
      proced .== Nothing
      bLex @== (f :)
    Just cs -> do
      bps <- look bLex
      let g = changeOrder Lex $ f - sum (zipWith (.*.) (V.toList $ fmap unwrapAlgebra cs) bps)
      proced .== Just (changeOrder Lex f)
      gLex @== (g :)

toContinue ::
  forall s r o n.
  ( 0 < n
  , Ord r
  , KnownNat n
  , Field r
  ) =>
  Machine s r o n Bool
toContinue = do
  mans <- look proced
  case mans of
    Nothing -> return True
    Just g -> do
      let xLast = P.maximum allVars `asTypeOf` g
      return $ not $ leadingMonomial g `isPowerOf` leadingMonomial xLast

nextMonomial ::
  forall s r ord n.
  (CoeffRing r, KnownNat n) =>
  Machine s r ord n ()
nextMonomial = do
  m <- look monomial
  gs <- map leadingMonomial <$> look gLex
  let next =
        fst $
          maximumBy
            (comparing snd)
            [ (OrderedMonomial monom, fromEnum od)
            | od <- enumOrdinal (sNat :: SNat n)
            , let monom = beta (getMonomial m) od
            , all (not . (`divs` OrderedMonomial monom)) gs
            ]
  monomial .== next

beta :: KnownNat n => Monomial n -> Ordinal n -> Monomial n
beta xs o@(OLt k) =
  let n = sizedLength xs
   in (SV.take (sSucc k) $ xs & ix o +~ 1) SV.++ SV.replicate (n %- sSucc k) 0
beta _ _ = error "beta: Bug in ghc!"


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Core.Softmax
Description : Softmax loss layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Softmax (
    Softmax (..)
  , softmax
  , softmax'
  ) where

import           Data.Serialize

import           Control.DeepSeq              (NFData (..))
import           GHC.Generics                 (Generic)
import           GHC.TypeLits
import           Grenade.Core

import           Numeric.LinearAlgebra.Static as LAS

-- | A Softmax layer
--
--   This layer is like a logit layer, but normalises
--   a set of matricies to be probabilities.
--
--   One can use this layer as the last layer in a network
--   if they need normalised probabilities.
data Softmax = Softmax
  deriving (Show, Generic, NFData)

instance UpdateLayer Softmax where
  type Gradient Softmax = ()
  runUpdate _ _ _ = Softmax

instance RandomLayer Softmax where
  createRandomWith _ _ = return Softmax

instance ( KnownNat i ) => Layer Softmax ('D1 i) ('D1 i) where
  type Tape Softmax ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (softmax y))
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (softmax' y dEdy))

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

softmax :: KnownNat i => LAS.R i -> LAS.R i
softmax xs =
  let xs' = LAS.dvmap exp xs
      s   = LAS.dot xs' 1
  in  LAS.dvmap (/ s) xs'

softmax' :: KnownNat i => LAS.R i -> LAS.R i -> LAS.R i
softmax' x grad =
  let yTy = outer sm sm
      d   = diag sm
      g   = d - yTy
  in  g #> grad
    where
  sm = softmax x

-------------------- GNum instances --------------------

instance GNum Softmax where
  _ |* Softmax = Softmax
  _ |+ Softmax = Softmax
  gFromRational _ = Softmax


-- |
-- = Single-node reservoir
--
-- In this project, we exploit an established analogy between
-- spatially extended and delay systems (refs. 1-3). That makes possible
-- to employ DDEs as a reservoir substrate.
--
-- 1. Arecchi, F. T., et al. “Two-Dimensional Representation of
--    a Delayed Dynamical System.” Physical Review A, vol. 45, no. 7,
--    Jan. 1992, doi:10.1103/physreva.45.r4225.
-- 2. Appeltant, L., et al. “Information Processing Using a Single
--    Dynamical Node as Complex System.” Nature Communications, vol. 2,
--    2011, p. 468., doi:10.1038/ncomms1476.
-- 3. Virtual Chimera States for Delayed-Feedback Systems - Laurent Larger,
--    Bogdan Penkovsky, Yuri Maistrenko. Physical Review Letters - 08 / 2013.

{-# LANGUAGE BangPatterns #-}
module RC.NTC.Reservoir
  ( Reservoir (..)
  , genReservoir
  ) where

import           Numeric.LinearAlgebra
import qualified Data.Vector.Storable as V
import qualified Linear.V2 as V2
import qualified Numeric.DDE as DDE
import qualified Numeric.DDE.Model as DDEModel

import           RC.NTC.Types
import qualified RC.Helpers as H

-- | Substrate-specific low-level reservoir implementation
genReservoir :: DDEModel.RC -> Reservoir
genReservoir par@DDEModel.RC {
    DDEModel._filt = DDEModel.BandpassFiltering { DDEModel._tau = tau }
  } = Reservoir _r
  where
    _r sample = H.unflatten' nodes responseX
      where
        oversampling = 1 :: Int  -- No oversampling
        detuning = 1.0 :: Double  -- Delay detuning factor, 1 = no detuning
        nodes = rows sample
        delaySamples = round $ detuning * fromIntegral (oversampling * nodes)

        -- Matrix to timetrace
        trace1 = H.flatten' sample

        -- Duplicate the last element (DDE.integHeun2_2D consumes one extra input)
        trace = trace1 V.++ V.singleton (V.last trace1)

        -- Empirically chosen integration time step:
        -- twice faster than the system response time tau
        hStep = tau / 2

        (_, response) = DDE.integHeun2_2D [delaySamples] hStep (DDEModel.bandpassRhs par) (DDE.Input trace)
        responseX = V.map (\(V2.V2 x _) -> x) response


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}

-- HdpH implementaion of a Mandelbrot set creation
--
-- Based largely on code by Rob Stewart <R.Stewart@hw.ac.uk> who adapted it
-- from the monad-par implementation.
--
-- Author: Blair Archbald
-- Date: 2/3/2015
-----------------------------------------------------------------------------

module Main where

import Control.DeepSeq hiding (force)
import Control.Exception (evaluate)
import Control.Monad (when)
import Control.Parallel.HdpH
import Control.Parallel.HdpH.Strategies
import qualified Control.Monad.Par as MonadPar
import qualified Control.Monad.Par.Combinator as MonadPar.C
import qualified Control.Parallel.HdpH as HdpH (declareStatic)
import qualified Control.Parallel.HdpH.Strategies as Strategies (declareStatic)

import Data.Binary (Binary)
import Data.Complex
import Data.Functor ((<$>))
import Data.List (stripPrefix)
import Data.List.Split (splitOn)
import Data.Maybe (isJust,fromJust)
import Data.Monoid (mconcat)
import Data.Serialize
import Data.Typeable (Typeable)
import Data.Vector.Cereal () -- To get the generic deriving for cereal and vectors.
import qualified Data.Vector.Unboxed as V

import GHC.Generics (Generic)

import System.Clock
import System.Environment (getArgs)
import System.IO (stdout, stderr, hSetBuffering, BufferMode(..))
import System.Random (mkStdGen, setStdGen)

import Prelude

type MSTime = Double

data VecTree = Leaf (V.Vector Int)
             | MkNode VecTree VecTree
             deriving (Eq,Show,Generic,Typeable)
instance Serialize VecTree
instance NFData VecTree
instance ToClosure VecTree
  where locToClosure = $(here)

instance ToClosure (Int, Int)
  where locToClosure = $(here)

----------------------------------------------------------------------------
-- sequential mandel function

mandel :: Int -> Complex Double -> Int
mandel max_depth c = loop 0 0
  where
   loop i !z
    | i == max_depth       = i
    | magnitude z >= 2.0   = i
    | otherwise            = loop (i+1) (z*z + c)

checkSum :: VecTree -> Int
checkSum (Leaf vec) = V.foldl (+) 0 vec
checkSum (MkNode v1 v2) = checkSum v1 + checkSum v2

-----------------------------------------------
-- Mandel using monad-par
--  From: https://github.com/simonmar/monad-par/blob/master/examples/src/mandel.hs

monadparRunMandel :: Int -> Int -> Int -> Int -> MonadPar.Par VecTree
monadparRunMandel = monadparRunMandel' (-2) (-2) 2 2

--This is where the skeleton comes from.
monadparRunMandel' :: Double -> Double -- (minX, MinY)
                    -> Double -> Double -- (maxX, maxY)
                    -> Int -> Int -- (winX, winY)
                    -> Int -- Depth
                    -> Int -- Threshold
                    -> MonadPar.Par VecTree
monadparRunMandel' minX minY maxX maxY winX winY max_depth threshold =
  MonadPar.C.parMapReduceRangeThresh threshold (MonadPar.C.InclusiveRange 0 (winY-1))
     (\y ->
       do
          let vec = V.generate winX (\x -> mandelStep y x)
          seq (vec V.! 0) $ return (Leaf vec))
     (\ a b -> return$ MkNode a b)
     (Leaf V.empty)
  where
    mandelStep i j = mandel max_depth z
        where z = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
                  ((fromIntegral i * c_scale) / fromIntegral winX + minX)
    r_scale  =  maxY - minY  :: Double
    c_scale =   maxX - minX  :: Double

-----------------------------------------------
-- Mandel using the generic HpdH D&C skeleton

hdphDandCMandel :: Int -> Int -> Int -> Int -> Par VecTree
hdphDandCMandel = hdphDandCMandel' (-2) (-2) 2 2

hdphDandCMandel' :: Double -> Double -- (minX, MinY)
                 -> Double -> Double -- (maxX, maxY)
                 -> Int -> Int -- (winX, winY)
                 -> Int -- Depth
                 -> Int -- Threshold
                 -> Par VecTree
hdphDandCMandel' minX minY maxX maxY winX winY maxDepth threshold = do
  res <- parDivideAndConquer
          $(mkClosure [|dc_trivial threshold|])
          $(mkClosure [|dc_decompose|])
          $(mkClosure [|dc_combine|])
          $(mkClosure [|dc_algorithm (minX, minY, maxX, maxY, winX, winY, maxDepth)|])
          (toClosure (0,winY-1))
  return $ unClosure res

dc_trivial :: Int -> Thunk (Closure (Int, Int) -> Bool)
dc_trivial threshold = Thunk $ \bnds -> let (min,max) = unClosure bnds
                                        in   max - min <= threshold

dc_decompose :: Closure (Int, Int) -> [Closure (Int, Int)]
dc_decompose bnds = let (min, max) = unClosure bnds
                        mid = min + (max - min) `quot` 2
                     in [toClosure (min, mid), toClosure (mid+1, max)]

-- This could prove to be a bottleneck and reason to use Rob's Skeleton
-- (TODO: compare performance between approaches)
dc_combine :: Closure a -> [Closure VecTree] -> Closure VecTree
dc_combine _ ts = toClosure $ foldl1 MkNode $ map unClosure ts

dc_algorithm :: (Double, Double, -- (minX, minY)
                 Double, Double, -- (maxX, maxY)
                 Int, Int, -- (winX, winY)
                 Int) -- maxDepth
              -> Thunk (Closure (Int,Int) -> Par (Closure VecTree))
dc_algorithm (minX, minY, maxX, maxY, winX, winY, maxDepth) = Thunk $ \bnds ->
  let (min,max) = unClosure bnds
      v = foldl go V.empty [min..max]
  in  do vec <- force v
         return $ toClosure (Leaf vec)
  where
    go a y = a V.++ V.generate winX (\x -> mandelStep y x)
    mandelStep i j = mandel maxDepth (calcZ i j)
    calcZ i j = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
                ((fromIntegral i * c_scale) / fromIntegral winX + minX)
    r_scale =  maxY - minY  :: Double
    c_scale =  maxX - minX  :: Double

---------------------------------------------------------------------------
-- Timed Version
hdphDandCMandelTimed :: Int -> Int -> Int -> Int -> Par VecTree
hdphDandCMandelTimed = hdphDandCMandelTimed' (-2) (-2) 2 2

hdphDandCMandelTimed' :: Double -> Double -- (minX, MinY)
                 -> Double -> Double -- (maxX, maxY)
                 -> Int -> Int -- (winX, winY)
                 -> Int -- Depth
                 -> Int -- Threshold
                 -> Par VecTree
hdphDandCMandelTimed' minX minY maxX maxY winX winY maxDepth threshold = do
  res <- parDivideAndConquer
          $(mkClosure [|dc_trivial threshold|])
          $(mkClosure [|dc_decompose|])
          $(mkClosure [|dc_combine|])
          $(mkClosure [|dc_algorithm_timed (minX, minY, maxX, maxY, winX, winY, maxDepth)|])
          (toClosure (0,winY-1))
  return $ unClosure res

dc_algorithm_timed :: (Double, Double, -- (minX, minY)
                       Double, Double, -- (maxX, maxY)
                       Int, Int, -- (winX, winY)
                       Int) -- maxDepth
                   -> Thunk (Closure (Int,Int) -> Par (Closure VecTree))
dc_algorithm_timed (minX, minY, maxX, maxY, winX, winY, maxDepth) = Thunk $ \bnds ->
  let (min,max) = unClosure bnds
      v = foldl go V.empty [min..max]
  in  do s <- io $ getTime Monotonic
         vec <- force v
         e <- io $ getTime Monotonic
         io $ putStrLn $ show min ++ "," ++ show max ++ ":" ++ show (timeDiffMSecs s e)
         return $ toClosure (Leaf vec)
  where
    go a y = a V.++ V.generate winX (\x -> mandelStep y x)
    mandelStep i j = mandel maxDepth (calcZ i j)
    calcZ i j = ((fromIntegral j * r_scale) / fromIntegral winY + minY) :+
                ((fromIntegral i * c_scale) / fromIntegral winX + minX)
    r_scale =  maxY - minY  :: Double
    c_scale =  maxX - minX  :: Double
-----------------------------------------------------------------------------
-- initialisation, argument processing and 'main'

timeDiffMSecs :: TimeSpec -> TimeSpec -> MSTime
timeDiffMSecs (TimeSpec s1 n1) (TimeSpec s2 n2) = fromIntegral (t2 - t1)
                                                          /
                                                  fromIntegral (10 ^ 6)
  where t1 = (fromIntegral s1 * 10 ^ 9) + fromIntegral n1
        t2 = (fromIntegral s2 * 10 ^ 9) + fromIntegral n2

timeIOMs :: IO a -> IO (a, MSTime)
timeIOMs action = do s  <- getTime Monotonic
                     x  <- action
                     e  <- getTime Monotonic
                     return (x, timeDiffMSecs s e)

-----------------------------------------------------
-- Static Closures

$(return []) -- Bring the types into scope so that reify works.
declareStatic :: StaticDecl
declareStatic =
  mconcat
    [HdpH.declareStatic,         -- declare Static deserialisers
     Strategies.declareStatic,   -- from imported modules
     declare (staticToClosure :: StaticToClosure VecTree),
     declare (staticToClosure :: StaticToClosure (Int,Int)),

     declare $(static 'dc_trivial),
     declare $(static 'dc_decompose),
     declare $(static 'dc_combine),
     declare $(static 'dc_algorithm),

     declare $(static 'dc_algorithm_timed),

     declare $(static 'mandel)]
---------------------------------------------------------------------------

-- parse (optional) arguments in this order:
-- * version to run
-- * X value for Mandel
-- * Y value for Mandel
-- * Depth value for Mandel
-- * Threshold for Mandel
parseArgs :: [String] -> (Int, Int, Int, Int, Int, Int)
parseArgs []     = (defVers, defX, defY, defDepth,defThreshold,defExpected)
parseArgs (s:ss) =
  let go :: Int -> [String] -> (Int, Int, Int, Int,Int, Int)
      go v []              = (v, defX, defY, defDepth,defThreshold,defExpected)
      go v [s1]            = (v, defX, read s1,  defDepth,defThreshold,defExpected)
      go v [s1,s2]         = (v, read s1,  read s2,  defDepth,defThreshold,defExpected)
      go v [s1,s2,s3]      = (v, read s1,  read s2,  read s3, defThreshold,defExpected)
      go v [s1,s2,s3,s4] = (v, read s1,  read s2,  read s3, read s4,defExpected)
      go v (s1:s2:s3:s4:s5:_) = (v, read s1,  read s2,  read s3, read s4, read s5)
  in case stripPrefix "v" s of
       Just s' -> go (read s') ss
       Nothing -> go defVers (s:ss)

-- defaults from Simon Marlow from monad-par example
defVers, defX, defY, defDepth, defThreshold, defExpected :: Int
defVers      = 0
defX         = 1024
defY         = 1024
defDepth     = 256
defThreshold = 1
defExpected  = 0

main :: IO ()
main = do
  hSetBuffering stdout LineBuffering
  hSetBuffering stderr LineBuffering
  register Main.declareStatic

  hdphCfg    <- flip updateConf defaultRTSConf =<< getArgs
  (conf, as) <- case hdphCfg of
    Left err -> error $ "Could not initialise HdpH config: " ++ show err
    Right a  -> return a

  let (version, valX, valY, valDepth,valThreshold,expected) = parseArgs as
  case version of
      1 -> do (p, t) <- timeIOMs $ evaluate =<< return (MonadPar.runPar
                (monadparRunMandel valX valY valDepth valThreshold))
              printOutput (Just p,t)
      2 -> do res <- timeIOMs $ evaluate =<< runParIO conf
                                (hdphDandCMandel valX valY valDepth valThreshold)
              printOutput res
      3 -> do res <- timeIOMs $ evaluate =<< runParIO conf
                                (hdphDandCMandelTimed valX valY valDepth valThreshold)
              printOutput res
      _ -> return ()

printOutput :: (Maybe VecTree, MSTime) -> IO ()
printOutput (p,t) = case p of
                     Just pixels -> do
                       putStrLn $ "CHECKSUM: " ++ show (checkSum pixels)
                       putStrLn $ "RUNTIME: "  ++ show t ++ "ms"
                     Nothing -> return ()


{-# LANGUAGE BangPatterns , RankNTypes, GADTs, DataKinds #-}

{- | The 'Numerical.HBLAS.BLAS.Level3' module provides a fully general
yet type safe Level3 BLAS API.

When in doubt about the semantics of an operation,
consult your system's BLAS api documentation, or just read the documentation
for
<https://software.intel.com/sites/products/documentation/hpc/mkl/mklman/index.htm the Intel MKL BLAS distribution>

A few basic notes about how to invoke BLAS routines.

Many BLAS operations take one or more arguments of type 'Transpose'.
'Tranpose' has the following different constructors, which tell BLAS
routines what transformation to implicitly apply to an input matrix @mat@ with dimension @n x m@.

*  'NoTranspose' leaves the matrix @mat@ as is.

* 'Transpose' treats the @mat@ as being implicitly transposed, with dimension
    @m x n@. Entry @mat(i,j)@ being treated as actually being the entry
    @mat(j,i)@. For Real matrices this is also the matrix adjoint operation.
    ie @Tranpose(mat)(i,j)=mat(j,i)@

*  'ConjNoTranspose' will implicitly conjugate @mat@, which is a no op for Real ('Float' or 'Double') matrices, but for
'Complex Float' and 'Complex Double' matrices, a given matrix entry @mat(i,j)==x':+'y@
will be treated as actually being  @conjugate(mat)(i,j)=y':+'x@.

* 'ConjTranpose' will implicitly transpose and conjugate the input matrix.
ConjugateTranpose acts as matrix adjoint for both real and complex matrices.

The *gemm operations  work as follows (using 'sgemm' as an example):

* @'sgemm trLeft trRight alpha beta left right result'@, where @trLeft@ and @trRight@
are values of type 'Transpose' that respectively act on the matrices @left@ and @right@.

* the generalized matrix computation thusly formed can be viewed as being
@result = alpha * trLeft(left) * trRight(right) + beta * result@

the *gemv operations are akin to the *gemm operations, but with @right@ and @result@
being vectors rather than matrices.

the *trsv operations solve for @x@ in the equation @A x = y@ given @A@ and @y@.
The 'MatUpLo' argument determines if the matrix should be treated as upper or
lower triangular and 'MatDiag' determines if the triangular solver should treat
the diagonal of the matrix as being all 1's or not.  A general pattern of invocation
would be @'strsv' matuplo  tranposeMatA  matdiag  matrixA  xVector@.
A key detail to note is that the input vector is ALSO the result vector,
ie 'strsv' and friends updates the vector place.

-}

module Numerical.HBLAS.BLAS.Level3(
        dgemm
        ,sgemm
        ,cgemm
        ,zgemm

        ,chemm
        ,zhemm

        ,cherk
        ,zherk

        ,cher2k
        ,zher2k

        ,ssymm
        ,dsymm
        ,csymm
        ,zsymm

        ,ssyrk
        ,dsyrk
        ,csyrk
        ,zsyrk

        ,ssyr2k
        ,dsyr2k
        ,csyr2k
        ,zsyr2k

        ,strmm
        ,dtrmm
        ,ctrmm
        ,ztrmm

        ,strsm
        ,dtrsm
        ,ctrsm
        ,ztrsm
            ) where

import Numerical.HBLAS.UtilsFFI
import Numerical.HBLAS.BLAS.FFI.Level3
import Numerical.HBLAS.BLAS.Internal.Level3
import Control.Monad.Primitive
import Data.Complex

sgemm :: PrimMonad m=>  GemmFun Float  orient  (PrimState m) m
sgemm =  gemmAbstraction "sgemm"  cblas_sgemm_safe cblas_sgemm_unsafe (\x f -> f x )

dgemm :: PrimMonad m=>  GemmFun  Double orient  (PrimState m) m
dgemm = gemmAbstraction "dgemm"  cblas_dgemm_safe cblas_dgemm_unsafe  (\x f -> f x )

cgemm :: PrimMonad m=>  GemmFun (Complex Float) orient  (PrimState m) m
cgemm = gemmAbstraction "cgemm" cblas_cgemm_safe cblas_cgemm_unsafe  withRStorable_

zgemm :: PrimMonad m=>  GemmFun (Complex Double) orient  (PrimState m) m
zgemm = gemmAbstraction "zgemm"  cblas_zgemm_safe cblas_zgemm_unsafe withRStorable_

chemm :: PrimMonad m=>  HemmFun (Complex Float) orient (PrimState m) m
chemm = hemmAbstraction "chemm" cblas_chemm_safe cblas_chemm_unsafe withRStorable_

zhemm :: PrimMonad m=>  HemmFun (Complex Double) orient (PrimState m) m
zhemm = hemmAbstraction "zhemm" cblas_zhemm_safe cblas_zhemm_unsafe withRStorable_

cherk :: PrimMonad m=>  HerkFun Float (Complex Float) orient (PrimState m) m
cherk = herkAbstraction "cherk" cblas_cherk_safe cblas_cherk_unsafe (\x f -> f x)

zherk :: PrimMonad m=>  HerkFun Double (Complex Double) orient (PrimState m) m
zherk = herkAbstraction "zherk" cblas_zherk_safe cblas_zherk_unsafe (\x f -> f x)

cher2k :: PrimMonad m=>  Her2kFun Float (Complex Float) orient (PrimState m) m
cher2k = her2kAbstraction "cher2k" cblas_cher2k_safe cblas_cher2k_unsafe withRStorable_

zher2k :: PrimMonad m=>  Her2kFun Double (Complex Double) orient (PrimState m) m
zher2k = her2kAbstraction "zher2k" cblas_zher2k_safe cblas_zher2k_unsafe withRStorable_

ssymm :: PrimMonad m=>  SymmFun Float orient (PrimState m) m
ssymm = symmAbstraction "ssymm" cblas_ssymm_safe cblas_ssymm_unsafe (\x f -> f x)

dsymm :: PrimMonad m=>  SymmFun Double orient (PrimState m) m
dsymm = symmAbstraction "dsymm" cblas_dsymm_safe cblas_dsymm_unsafe (\x f -> f x)

csymm :: PrimMonad m=>  SymmFun (Complex Float) orient (PrimState m) m
csymm = symmAbstraction "csymm" cblas_csymm_safe cblas_csymm_unsafe withRStorable_

zsymm :: PrimMonad m=>  SymmFun (Complex Double) orient (PrimState m) m
zsymm = symmAbstraction "zsymm" cblas_zsymm_safe cblas_zsymm_unsafe withRStorable_

ssyrk :: PrimMonad m=>  SyrkFun Float orient (PrimState m) m
ssyrk = syrkAbstraction "ssyrk" cblas_ssyrk_safe cblas_ssyrk_unsafe (\x f -> f x)

dsyrk :: PrimMonad m=>  SyrkFun Double orient (PrimState m) m
dsyrk = syrkAbstraction "dsyrk" cblas_dsyrk_safe cblas_dsyrk_unsafe (\x f -> f x)

csyrk :: PrimMonad m=>  SyrkFun (Complex Float) orient (PrimState m) m
csyrk = syrkAbstraction "csyrk" cblas_csyrk_safe cblas_csyrk_unsafe withRStorable_

zsyrk :: PrimMonad m=>  SyrkFun (Complex Double) orient (PrimState m) m
zsyrk = syrkAbstraction "zsyrk" cblas_zsyrk_safe cblas_zsyrk_unsafe withRStorable_

ssyr2k :: PrimMonad m=>  Syr2kFun Float orient (PrimState m) m
ssyr2k = syr2kAbstraction "ssyr2k" cblas_ssyr2k_safe cblas_ssyr2k_unsafe (\x f -> f x)

dsyr2k :: PrimMonad m=>  Syr2kFun Double orient (PrimState m) m
dsyr2k = syr2kAbstraction "dsyr2k" cblas_dsyr2k_safe cblas_dsyr2k_unsafe (\x f -> f x)

csyr2k :: PrimMonad m=>  Syr2kFun (Complex Float) orient (PrimState m) m
csyr2k = syr2kAbstraction "csyr2k" cblas_csyr2k_safe cblas_csyr2k_unsafe withRStorable_

zsyr2k :: PrimMonad m=>  Syr2kFun (Complex Double) orient (PrimState m) m
zsyr2k = syr2kAbstraction "zsyr2k" cblas_zsyr2k_safe cblas_zsyr2k_unsafe withRStorable_

strmm :: PrimMonad m=> TrmmFun Float orient (PrimState m) m
strmm = trmmAbstraction "strmm" cblas_strmm_safe cblas_strmm_unsafe (\x f -> f x)

dtrmm :: PrimMonad m=> TrmmFun Double orient (PrimState m) m
dtrmm = trmmAbstraction "dtrmm" cblas_dtrmm_safe cblas_dtrmm_unsafe (\x f -> f x)

ctrmm :: PrimMonad m=> TrmmFun (Complex Float) orient (PrimState m) m
ctrmm = trmmAbstraction "ctrmm" cblas_ctrmm_safe cblas_ctrmm_unsafe withRStorable_

ztrmm :: PrimMonad m=> TrmmFun (Complex Double) orient (PrimState m) m
ztrmm = trmmAbstraction "ztrmm" cblas_ztrmm_safe cblas_ztrmm_unsafe withRStorable_

strsm :: PrimMonad m=> TrsmFun Float orient (PrimState m) m
strsm = trsmAbstraction "strsm" cblas_strsm_safe cblas_strsm_unsafe (\x f -> f x)

dtrsm :: PrimMonad m=> TrsmFun Double orient (PrimState m) m
dtrsm = trsmAbstraction "dtrsm" cblas_dtrsm_safe cblas_dtrsm_unsafe (\x f -> f x)

ctrsm :: PrimMonad m=> TrsmFun (Complex Float) orient (PrimState m) m
ctrsm = trsmAbstraction "ctrsm" cblas_ctrsm_safe cblas_ctrsm_unsafe withRStorable_

ztrsm :: PrimMonad m=> TrsmFun (Complex Double) orient (PrimState m) m
ztrsm = trsmAbstraction "ztrsm" cblas_ztrsm_safe cblas_ztrsm_unsafe withRStorable_


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module PLC1 (PLC1(..)) where

import Control.Applicative
import Control.Comonad
import Control.Comonad.Store
import Control.Exception (assert)
-- import Control.Monad
-- import Data.Distributive
import Data.Proxy
import Data.Unfoldable
import Data.Validity
import qualified Data.Vector as V
import Data.VectorSpace
import GHC.TypeLits
-- import Linear.Metric
-- import Linear.Vector
import qualified Numeric.LinearAlgebra as L
import qualified Test.QuickCheck as QC
import qualified Test.QuickCheck.Function as QC

import ComonadicVector
import Function
import Integration

-- https://hackage.haskell.org/package/polynomial-0.7.3

newtype PLC1 (n' :: Nat) c s = PLC1
  { getPLC1 :: ComonadicVector V.Vector s
  } deriving (Eq, Ord, Read, Show, Foldable, Functor, Traversable)

instance (KnownNat n', Validity s) => Validity (PLC1 n' c s) where
    validate (PLC1 xs) =
        mconcat
        [ n > 0 <?@> "length is positive"
        , length xs == n <?@> "vector has correct length "
        , xs <?!> "vector"
        ]
        where n = fromInteger (natVal (Proxy :: Proxy n'))
    isValid = isValidByValidating

instance KnownNat n' => Unfoldable (PLC1 n' c) where
    -- unfold xs = PLC1 <$> unfold xs
    -- unfold xs = (PLC1 . ComonadicVector 0 . V.fromListN n) <$> mklist n
    --     where n = fromInteger (natVal (Proxy :: Proxy n')) :: Int
    --           mklist i | i == 0 = pure []
    --                    | otherwise = (:) <$> xs <*> mklist (i-1)
    unfold xs = PLC1 . goodlength <$> unfold xs
        where n = fromInteger (natVal (Proxy :: Proxy n')) :: Int
              goodlength ys = assert (n > 0 && length ys == n) ys

instance (Applicative (ComonadicVector V.Vector), KnownNat n') =>
         Applicative (PLC1 n' c) where
  -- pure x = PLC1 (pure x)
  -- Note: 'pure' should use the neutral element of the monoid used
  -- for 'pos_' in '(<*>)'
  pure x = PLC1 (ComonadicVector (n - 1) (V.replicate n x))
      where n = fromInteger (natVal (Proxy :: Proxy n'))
  PLC1 fs <*> PLC1 xs = PLC1 (fs <*> xs)

-- instance (Alternative (ComonadicVector V.Vector), KnownNat n') =>
--          Alternative (PLC1 n' c) where
--   empty = PLC1 empty
--   PLC1 xs <|> PLC1 ys = PLC1 (xs <|> ys)

-- instance Monad (ComonadicVector V.Vector) => Monad (PLC1 n' c) where
--     -- (>>=) :: PLC1 n' c s -> (s -> PLC1 n' c t) -> PLC1 n' c t
--     PLC1 x >>= f = PLC1 (x >>= getPLC1 . f)
--
-- instance MonadPlus (ComonadicVector V.Vector) => MonadPlus (PLC1 n' c) where
--     mzero = PLC1 mzero
--     PLC1 xs `mplus` PLC1 ys = PLC1 (xs `mplus` ys)

-- instance Distributive (ComonadicVector V.Vector) =>
--          Distributive (PLC1 n' c) where
--   collect f arr = PLC1 (collect (getPLC1 . f) arr)

instance Comonad (ComonadicVector V.Vector) => Comonad (PLC1 n' c) where
  extract (PLC1 arr) = extract arr
  extend f (PLC1 arr) = PLC1 (extend (f . PLC1) arr)

instance ComonadStore Int (ComonadicVector V.Vector) =>
         ComonadStore Int (PLC1 n' c) where
  pos (PLC1 arr) = pos arr
  peek i (PLC1 arr) = peek i arr

instance (KnownNat n', QC.Arbitrary s) => QC.Arbitrary (PLC1 n' c s) where
  -- arbitrary = do
  --   arr <- QC.arbitrary
  --   return (PLC1 arr)
  -- shrink (PLC1 arr) = [PLC1 arr' | arr' <- QC.shrink arr]
  -- arbitrary = arbitraryDefault
  arbitrary = do
      xs <- V.generateM n (const QC.arbitrary)
      i <- QC.choose (0, V.length xs - 1)
      return (PLC1 (ComonadicVector i xs))
      where n = fromInteger (natVal (Proxy :: Proxy n'))
  shrink arr = [arr' | pos arr > 0] ++ traverse QC.shrink arr'
      where arr' = seek 0 arr

instance QC.CoArbitrary s => QC.CoArbitrary (PLC1 n' c s) where
    coarbitrary = QC.coarbitrary . getPLC1

instance QC.Function s => QC.Function (PLC1 n' c s) where
    function = QC.functionMap getPLC1 PLC1

-- instance Additive (ComonadicVector V.Vector) => Additive (PLC1 n' c) where
--     zero = PLC1 zero
--
-- instance (Additive (PLC1 n' c), Bounded c, RealFrac c) => Metric (PLC1 n' c) where
--     -- dot :: Num a => f a -> f a -> a
--     dot xs ys =
--         assert (length ys == n) $
--         case n of
--           0 -> 0
--           _ -> sum [w i * peek i xs * peek i ys | i <- [0 .. n-1]]
--         where n = length xs
--               xmin = minBound :: c
--               xmax = maxBound :: c
--               w i = realToFrac (weight n xmin xmax i)

instance ( KnownNat n'
         , AdditiveGroup (ComonadicVector V.Vector s)
         , AdditiveGroup s
         ) =>
         AdditiveGroup (PLC1 n' c s) where
  zeroV = pure zeroV
  negateV = fmap negateV
  (^+^) = liftA2 (^+^)

instance ( KnownNat n'
         , VectorSpace (ComonadicVector V.Vector s)
         , VectorSpace s
         ) =>
         VectorSpace (PLC1 n' c s) where
  type Scalar (PLC1 n' c s) = Scalar s
  (*^) x = fmap (x *^)

-- instance ( KnownNat n'
--          , Bounded c
--          , RealFrac c
--          , InnerSpace s
--          , VectorSpace (Scalar s)
--          , Fractional (Scalar (Scalar s))
--          ) =>
--          InnerSpace (PLC1 n' c s) where
--   -- xs <.> ys =
--   --   assert (length ys == length xs) $ sumV (weighted ((<.>) <$> xs <*> ys))

instance ( KnownNat n'
         , Integrable (PLC1 n')
         , IntegrableOk (PLC1 n') c s
         , Bounded c
         , RealFrac c
         , InnerSpace s
         , Fractional (Scalar s)
         ) => InnerSpace (PLC1 n' c s) where
  xs <.> ys = getPLC1 (extend weighted xs) <.> getPLC1 ys
    where
      weighted zs = weight n (xmin, xmax) (pos zs) *^ extract zs
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs

instance Function (PLC1 n') where
  type FunctionOk (PLC1 n') a b = ( KnownNat n'
                                  , Bounded a
                                  , RealFrac a
                                  , InnerSpace b
                                  , Floating (Scalar b)
                                  , Ord (Scalar b)
                                  , b ~ Double) -- TODO
  eval (PLC1 arr) x =
    case n of
      0 -> zeroV
      1 -> extract arr
      _ ->
        let y = ((x - cx) / rx + 1) / 2 * fromIntegral (n - 1)
            i = max 0 $ min (n - 2) $ floor y
            f0 = 1 - f1
            f1 = realToFrac (y - fromIntegral i)
        in f0 *^ peek i arr ^+^ f1 *^ peek (i + 1) arr
    where
      -- n = length arr
      n = fromInteger (natVal (Proxy :: Proxy n'))
      xmin = minBound `asTypeOf` x
      xmax = maxBound `asTypeOf` x
      cx = (xmax + xmin) / 2
      rx = (xmax - xmin) / 2

instance KnownNat n' => Discretization (PLC1 n') where
  discretized f = PLC1 (ComonadicVector 0 res)
    where
      n = fromInteger (natVal (Proxy :: Proxy n'))
      -- res = V.generate n go
      -- TODO: use (symmetric) triangular solver
      rhs = L.vector [go i | i <- [0..n-1]]
      mat = L.matrix n [ overlap n (xmin, xmax) i j
                       | i <- [0..n-1], j <- [0..n-1]]
      res = V.fromListN n $ L.toList $ mat L.<\> rhs
      (xmin, xmax) = (minBound, maxBound)
      _ = (f xmin, f xmax)      -- constrain types
      go i = integrate bf lo hi
        where
          bf x = b x *^ f x
          b x = basis n (xmin, xmax) i x
          (lo, hi) = support n (xmin, xmax) i

instance KnownNat n' => Integrable (PLC1 n') where
  type IntegrableOk (PLC1 n') a b = ()
  integral xs = sumV (extend weighted xs)
    where
      weighted ys = weight n (xmin, xmax) (pos ys) *^ extract ys
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs

instance Differentiable (PLC1 n') where
  type DifferentiableOk (PLC1 n') a b = ()
  type Dir (PLC1 n') = ()
  boundary () xs = extend bndry xs
    where
      bndry ys
        | n == 1     = zeroV *^ extract ys -- or zeroV
        | i == 0     = negateV (extract ys ^/ realToFrac dx)
        | i == n - 1 = extract ys ^/ realToFrac dx
        | otherwise  = zeroV *^ extract ys -- or zeroV
        where
          i = pos ys
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs
      rx = (xmax - xmin) / 2
      hdx = rx / fromIntegral (n - 1)
      dx = 2 * hdx
  derivative () xs = extend deriv xs
    where
      deriv ys
        | n == 1     = zeroV *^ extract ys -- or zeroV
        | i == 0     = (peek (i + 1) ys ^-^ extract ys) ^/ realToFrac dx
        | i == n - 1 = (extract ys ^-^ peek (i - 1) ys) ^/ realToFrac dx
        | otherwise  =
            (peek (i + 1) ys ^-^ peek (i - 1) ys) ^/ realToFrac (2 * dx)
        where
          i = pos ys
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs
      rx = (xmax - xmin) / 2
      hdx = rx / fromIntegral (n - 1)
      dx = 2 * hdx

triangle :: (Num a, Ord a) => a -> a
triangle x = max 0 (1 - abs x)

typed :: f a b -> a -> a
typed _ = id

bounds :: Bounded a => f a b -> (a, a)
bounds xs = (typed xs minBound, typed xs maxBound)

-- The basis functions are defined to have a maximum of '1'
basis :: (RealFrac a, Fractional b) => Int -> (a, a) -> Int -> a -> b
basis n (xmin, xmax) i x | n == 1    = 1
                         | otherwise = realToFrac (triangle y)
  where
    cx  = (xmin + xmax) / 2
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)
    xi  = cx + hdx * fromIntegral (2 * i + 1 - n)
    dx  = 2 * hdx
    y   = (x - xi) / dx

support :: (Fractional a, Ord a) => Int -> (a, a) -> Int -> (a, a)
support n (xmin, xmax) i | n == 1    = (xmin, xmax)
                         | otherwise = (max xmin (xi - dx), min xmax (xi + dx))
  where
    cx  = (xmin + xmax) / 2
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)
    xi  = cx + hdx * fromIntegral (2 * i + 1 - n)
    dx  = 2 * hdx

-- Volume
weight :: (RealFrac a, Fractional b) => Int -> (a, a) -> Int -> b
weight n (xmin, xmax) i | n == 1     = realToFrac (2 * rx)
                        | i == 0     = realToFrac hdx
                        | i == n - 1 = realToFrac hdx
                        | otherwise  = realToFrac (2 * hdx)
  where
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)

overlap :: (RealFrac a, Fractional b) => Int -> (a, a) -> Int -> Int -> b
overlap n (xmin, xmax) i j
    | n == 1                           = realToFrac (2 * rx)
    | i == j && (i == 0 || i == n - 1) = realToFrac (2 / 3 * hdx)
    | i == j                           = realToFrac (4 / 3 * hdx)
    | abs (i - j) == 1                 = realToFrac (1 / 3 * hdx)
    | otherwise                        = 0
  where
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)