{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators    #-}
module FokkerPlanck.Pinwheel
  ( module Filter.Pinwheel
  , module FokkerPlanck.Pinwheel
  ) where

import           Data.Array.Repa               as R
import           Data.Complex
import           Data.List                     as L
import           Data.Vector.Unboxed           as VU
import           Filter.Pinwheel
import           FokkerPlanck.Interpolation
import           Numeric.LinearAlgebra.Data    as NL
import           Numeric.LinearAlgebra.HMatrix
import           Types
import           Utils.Array
import           Utils.Coordinates
import           Utils.Parallel

{-# INLINE computeR2Z1T0ArrayRadial #-}
computeR2Z1T0ArrayRadial ::
     (R.Source r Double)
  => R.Array r DIM3 Double
  -> Int
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array D DIM4 (Complex Double)
computeR2Z1T0ArrayRadial radialArr xLen yLen scaleFactor thetaFreqs theta0Freqs =
  let pinwheelArr =
        traverse2
          (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
          (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
          (\(Z :. numThetaFreq) (Z :. numTheta0Freq) ->
             (Z :. numThetaFreq :. numTheta0Freq :. xLen :. yLen)) $ \f f0 (Z :. k :. l :. i :. j) ->
          pinwheel
            (f (Z :. k) - f0 (Z :. l))
            0
            (exp 1)
            0
            (i - center xLen)
            (j - center yLen)
   in radialCubicInterpolation radialArr scaleFactor pinwheelArr

{-# INLINE computeR2Z2T0S0ArrayRadial #-}
computeR2Z2T0S0ArrayRadial ::
     (R.Source r Double)
  => PinwheelType
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> R.Array D DIM6 (Complex Double)
computeR2Z2T0S0ArrayRadial pinwheelType radialArr xLen yLen scaleFactor rMax thetaFreqs scaleFreqs theta0Freqs scale0Freqs =
  let pinwheelArr =
        traverse4
          (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
          (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
          (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
          (fromListUnboxed (Z :. L.length scale0Freqs) scale0Freqs)
          (\(Z :. numThetaFreq) (Z :. numScaleFreq) (Z :. numTheta0Freq) (Z :. numScale0Freq) ->
             (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :.
              numScale0Freq :.
              xLen :.
              yLen)) $ \ft fs ft0 fs0 (Z :. t :. s :. t0 :. s0 :. i :. j) ->
          pinwheelFunc
            pinwheelType
            (ft (Z :. t) - ft0 (Z :. t0))
            (fs (Z :. s) - fs0 (Z :. s0))
            rMax
            0
            (i - center xLen)
            (j - center yLen)
  in radialCubicInterpolation radialArr scaleFactor pinwheelArr

{-# INLINE computeR2Z2T0S0ArrayRadial' #-}
computeR2Z2T0S0ArrayRadial' ::
     Int
  -> Int
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> R.Array D DIM6 (Complex Double)
computeR2Z2T0S0ArrayRadial' xLen yLen scaleFactor rMax thetaFreqs scaleFreqs theta0Freqs scale0Freqs =
  let pinwheelArr =
        traverse4
          (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
          (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
          (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
          (fromListUnboxed (Z :. L.length scale0Freqs) scale0Freqs)
          (\(Z :. numThetaFreq) (Z :. numScaleFreq) (Z :. numTheta0Freq) (Z :. numScale0Freq) ->
             (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :.
              numScale0Freq :.
              xLen :.
              yLen)) $ \ft fs ft0 fs0 (Z :. t :. s :. t0 :. s0 :. i :. j) ->
          pinwheel
            (ft (Z :. t) - ft0 (Z :. t0))
            (fs (Z :. s) - fs0 (Z :. s0))
            rMax
            0
            (i - center xLen)
            (j - center yLen)
   in pinwheelArr

{-# INLINE cutoff #-}
cutoff :: (R.Source r Double) => Int -> R.Array r DIM5 Double -> R.Array D DIM5 Double
cutoff r arr =
  R.traverse arr id $ \f idx@(Z :. _ :. _ :. _ :. _ :. e) ->
    if e > r
      then 0
      else f idx

{-# INLINE computeLocalEigenVector #-}
computeLocalEigenVector ::
     (R.Source r Double)
  => ParallelParams
  -> (Double -> Double -> Double -> Double -> Int -> Int -> Complex Double)
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
computeLocalEigenVector parallelParams pinwheelFunc radialArr xLen yLen rMax thetaFreqs scaleFreqs =
  let (Z :. numThetaFreq :. numScaleFreq :. _ :. _ :. _) = extent radialArr
      len = numThetaFreq * numScaleFreq
   in computeS .
      rotate4D .
      rotate4D .
      fromUnboxed (Z :. xLen :. yLen :. numThetaFreq :. numScaleFreq) .
      VU.concat .
      parMapChunk
        parallelParams
        rdeepseq
        (\(x, y) ->
           let r =
                 sqrt $
                 (fromIntegral $ x - center xLen) ^ 2 +
                 (fromIntegral $ y - center yLen) ^ 2
               interpolatedArray = cubicInterpolation radialArr r
               interpolatedPinwheel =
                 R.traverse3
                   interpolatedArray
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
                   (\sh _ _ -> sh) $ \f ft fs idx@(Z :. t :. s :. t0 :. s0) ->
                   (f idx :+ 0) *
                   pinwheelFunc
                     (ft (Z :. t) - ft (Z :. t0))
                     (fs (Z :. s) + fs (Z :. s0))
                     rMax
                     0
                     (x - center xLen)
                     (y - center yLen)
               (eigVal, eigVec) =
                 eig . (len >< len) . R.toList $ interpolatedPinwheel
               (val, vec) =
                 L.head . L.reverse . L.sortOn (magnitude . fst) $
                 L.zip (NL.toList eigVal) (toColumns $ eigVec)
               dominantVec = VU.fromList . L.map (* val) . NL.toList $ vec
            in -- VU.zipWith
               --   (*)
               --   (VU.fromList
               --      [ exp
               --        (0 :+
               --         (tf) *
               --         (angleFunctionRad
               --            (fromIntegral $ x - center xLen)
               --            (fromIntegral $ y - center yLen)))
               --      | tf <- thetaFreqs
               --      , sf <- scaleFreqs
               --      ])
               --   dominantVec
               dominantVec
        ) $
      [(x, y) | x <- [0 .. xLen - 1], y <- [0 .. yLen - 1]]

{-# INLINE computeLocalEigenVectorSink #-}
computeLocalEigenVectorSink ::
     (R.Source r Double)
  => ParallelParams
  -> (Double -> Double -> Double -> Double -> Int -> Int -> Complex Double)
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
computeLocalEigenVectorSink parallelParams pinwheelFunc radialArr xLen yLen rMax thetaFreqs scaleFreqs =
  let (Z :. numThetaFreq :. numScaleFreq :. _ :. _ :. _) = extent radialArr
      len = numThetaFreq * numScaleFreq
   in computeS .
      rotate4D .
      rotate4D .
      fromUnboxed (Z :. xLen :. yLen :. numThetaFreq :. numScaleFreq) .
      VU.concat .
      parMapChunk
        parallelParams
        rdeepseq
        (\(x, y) ->
           let r =
                 sqrt $
                 (fromIntegral $ x - center xLen) ^ 2 +
                 (fromIntegral $ y - center yLen) ^ 2
               interpolatedArray = cubicInterpolation radialArr r
               interpolatedPinwheel =
                 R.traverse3
                   interpolatedArray
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
                   (\sh _ _ -> sh) $ \f ft fs idx@(Z :. t :. s :. t0 :. s0) ->
                   (f idx :+ 0) * (exp (0 :+ (ft (Z :. t) - ft (Z :. t0)) * pi)) *
                   pinwheelFunc
                     (ft (Z :. t) - ft (Z :. t0))
                     (fs (Z :. s) + fs (Z :. s0))
                     rMax
                     0
                     (x - center xLen)
                     (y - center yLen)
               (eigVal, eigVec) =
                 eig . (len >< len) . R.toList $ interpolatedPinwheel
               (val, vec) =
                 L.head . L.reverse . L.sortOn (realPart . fst) $
                 L.zip (NL.toList eigVal) (toColumns $ eigVec)
            in VU.fromList . L.map (* val) . NL.toList $ vec) $
      [(x, y) | x <- [0 .. xLen - 1], y <- [0 .. yLen - 1]]


{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeOperators       #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Layers.FullyConnected where

import           Data.Constraint               (Dict (..))
import           Data.Kind                     (Type)
import           Data.Proxy
import           Data.Singletons               ()
import qualified Data.Vector.Storable          as V
import           GHC.TypeLits
import           Hedgehog
import qualified Numeric.LinearAlgebra.Static  as LA
import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix
import           Unsafe.Coerce                 (unsafeCoerce)

import           Grenade.Core
import           Grenade.Layers.FullyConnected
import           Grenade.Layers.Internal.BLAS
import           Grenade.Utils.ListStore

data OpaqueFullyConnected :: Type where
     OpaqueFullyConnected :: (KnownNat i, KnownNat o, KnownNat (i * o)) => FullyConnected i o -> OpaqueFullyConnected

instance Show OpaqueFullyConnected where
    show (OpaqueFullyConnected n) = show n

genOpaqueFullyConnected :: Gen OpaqueFullyConnected
genOpaqueFullyConnected = do
  input :: Integer <- choose 2 100
  output :: Integer <- choose 1 100
  variant <- choose 1 2
  let Just input' = someNatVal input
  let Just output' = someNatVal output
  case (input', output') of
    (SomeNat (Proxy :: Proxy i'), SomeNat (Proxy :: Proxy o')) ->
      case (unsafeCoerce (Dict :: Dict ()) :: Dict (KnownNat (i' * o'))) of
        Dict -> do
          wB <- randomVector
          bM <- randomVector
          wN <- uniformSample
          kM <- uniformSample
          let wInTmp = V.replicate (fromIntegral input) 0
              wBTmp = V.replicate (fromIntegral output) 0
              wNTmp = V.replicate (fromIntegral $ input * output) 0
          return . OpaqueFullyConnected $ case variant of
            1 -> (FullyConnected (FullyConnectedHMatrix wB wN) (ListStore 0 [Just $ FullyConnectedHMatrix bM kM]) (TempVectors wInTmp wBTmp wNTmp) :: FullyConnected i' o')
            2 ->
              let wB' = LA.extract wB
                  bM' = LA.extract bM
                  ext m = V.concat $ map LA.extract (LA.toRows m)
                  wN' = ext wN
                  kM' = ext kM
              in (FullyConnected (FullyConnectedBLAS wB' wN') (ListStore 0 [Just $ FullyConnectedBLAS bM' kM']) (TempVectors wInTmp wBTmp wNTmp) :: FullyConnected i' o')
            _ -> error "unexpected variant in genOpaqueFullyConnected in Test/FullyConnected.hs"

prop_fully_connected_forwards :: Property
prop_fully_connected_forwards = property $ do
    OpaqueFullyConnected (fclayer :: FullyConnected i o) <- blindForAll genOpaqueFullyConnected
    input :: S ('D1 i) <- blindForAll (S1D <$> randomVector)
    let (tape, output :: S ('D1 o)) = runForwards fclayer input
        backed :: (Gradient (FullyConnected i o), S ('D1 i))
                                    = runBackwards fclayer tape output
    backed `seq` success

tests :: IO Bool
tests = checkParallel $$(discover)


{-
   Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
-}

{-
  Units of measure extension to Fortran: backend
-}

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

module Camfort.Specification.Units.InferenceBackendFlint where

import           Control.Monad
import           Data.List (partition)
-- import           Debug.Trace (trace, traceM, traceShowM)
import           Foreign
import           Foreign.C.Types
import           Numeric.LinearAlgebra (atIndex, rows, cols)
import qualified Numeric.LinearAlgebra as H
import           System.IO.Unsafe (unsafePerformIO)

foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_init" fmpz_mat_init :: Ptr FMPZMat -> CLong -> CLong -> IO ()
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_set" fmpz_mat_set :: Ptr FMPZMat -> Ptr FMPZMat -> IO ()
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_entry" fmpz_mat_entry :: Ptr FMPZMat -> CLong -> CLong -> IO (Ptr CLong)
foreign import ccall unsafe "flint/fmpz.h fmpz_set_si" fmpz_set_si :: Ptr CLong -> CLong -> IO ()
foreign import ccall unsafe "flint/fmpz.h fmpz_get_si" fmpz_get_si :: Ptr CLong -> IO CLong
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_clear" fmpz_mat_clear :: Ptr FMPZMat -> IO ()
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_print_pretty" fmpz_mat_print_pretty :: Ptr FMPZMat -> IO ()
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_mul" fmpz_mat_mul :: Ptr FMPZMat -> Ptr FMPZMat -> Ptr FMPZMat -> IO ()

-- fmpz_mat_window_init(fmpz_mat_t window, const fmpz_mat_t mat, slong r1, slong c1, slong r2, slong c2)
--
-- Initializes the matrix window to be an r2 - r1 by c2 - c1 submatrix
-- of mat whose (0,0) entry is the (r1, c1) entry of mat. The memory
-- for the elements of window is shared with mat.
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_window_init" fmpz_mat_window_init :: Ptr FMPZMat -> Ptr FMPZMat -> CLong -> CLong -> CLong -> CLong -> IO ()

-- Frees the window (leaving underlying matrix alone).
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_window_clear" fmpz_mat_window_clear :: Ptr FMPZMat -> IO ()

-- r <- fmp_mat_rref B den A
--
-- Uses fraction-free Gauss-Jordan elimination to set (B, den) to the
-- reduced row echelon form of A and returns the rank of A. Aliasing
-- of A and B is allowed. r is rank of A.
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_rref" fmpz_mat_rref :: Ptr FMPZMat -> Ptr CLong -> Ptr FMPZMat -> IO CLong

-- r <- fmp_mat_inv B den A
--
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_inv" fmpz_mat_inv :: Ptr FMPZMat -> Ptr CLong -> Ptr FMPZMat -> IO CLong

-- fmpz_mat_hnf H A
--
-- H is the Hermite Normal Form of A.
foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_hnf" fmpz_mat_hnf :: Ptr FMPZMat -> Ptr FMPZMat -> IO ()

foreign import ccall unsafe "flint/fmpz_mat.h fmpz_mat_rank" fmpz_mat_rank :: Ptr FMPZMat -> IO CLong

data FMPZMat

instance Storable FMPZMat where
  sizeOf _ = 4 * sizeOf (undefined :: CLong)
  alignment _ = alignment (undefined :: CLong)
  peek _ = undefined
  poke _ = undefined

-- testFlint :: IO ()
-- testFlint = do
--   traceM "***********************8 testFlint 8***********************"
--   alloca $ \ a -> do
--     alloca $ \ b -> do
--       alloca $ \ den -> do
--         let n = 10
--         fmpz_mat_init a n n
--         fmpz_mat_init b n n
--         forM_ [0..n-1] $ \ i -> do
--           forM_ [0..n-1] $ \ j -> do
--             e <- fmpz_mat_entry a i j
--             fmpz_set_si e (2 * i + j)

--         -- fmpz_mat_mul b a a
--         r <- fmpz_mat_rref b den a
--         fmpz_mat_print_pretty a
--         fmpz_mat_print_pretty b
--         d <- peek den
--         traceM $ "r = " ++ show r ++ " den = " ++ show d
--         fmpz_mat_hnf b a
--         fmpz_mat_print_pretty b
--         fmpz_mat_clear a
--         fmpz_mat_clear b

rref :: H.Matrix Double -> (H.Matrix Double, Int, Int)
rref m = unsafePerformIO $ do
  alloca $ \ outputM -> do
    alloca $ \ inputM -> do
      alloca $ \ den -> do
        let numRows = fromIntegral $ rows m
        let numCols = fromIntegral $ cols m
        fmpz_mat_init outputM numRows numCols
        fmpz_mat_init inputM numRows numCols
        forM_ [0..numRows-1] $ \ i -> do
          forM_ [0..numCols-1] $ \ j -> do
            e <- fmpz_mat_entry inputM i j
            fmpz_set_si e (floor (m `atIndex` (fromIntegral i, fromIntegral j)))
        r <- fmpz_mat_rref outputM den inputM
        d <- peek den

        -- DEBUG:
        -- fmpz_mat_print_pretty outputM
        -- traceM $ "r = " ++ show r ++ " den = " ++ show d
        --

        lists <- forM [0..numRows-1] $ \ i -> do
          forM [0..numCols-1] $ \ j -> do
            e <- fmpz_mat_entry outputM i j
            fromIntegral `fmap` fmpz_get_si e
        let m' = H.fromLists lists
        fmpz_mat_clear inputM
        fmpz_mat_clear outputM
        return (m', fromIntegral d, fromIntegral r)

hnf :: H.Matrix Double -> H.Matrix Double
hnf m = unsafePerformIO $ do
  alloca $ \ outputM -> do
    alloca $ \ inputM -> do
      let numRows = fromIntegral $ rows m
      let numCols = fromIntegral $ cols m
      fmpz_mat_init outputM numRows numCols
      fmpz_mat_init inputM numRows numCols
      forM_ [0..numRows-1] $ \ i -> do
        forM_ [0..numCols-1] $ \ j -> do
          e <- fmpz_mat_entry inputM i j
          fmpz_set_si e (floor (m `atIndex` (fromIntegral i, fromIntegral j)))
      fmpz_mat_hnf outputM inputM
      r <- fmpz_mat_rank outputM

      -- DEBUG:
      -- fmpz_mat_print_pretty outputM
      -- traceM $ "rank = " ++ show r
      --

      lists <- forM [0..fromIntegral r-1] $ \ i -> do
        forM [0..numCols-1] $ \ j -> do
          e <- fmpz_mat_entry outputM i j
          fromIntegral `fmap` fmpz_get_si e
      let m' = H.fromLists lists
      fmpz_mat_clear inputM
      fmpz_mat_clear outputM
      return m'

inv :: H.Matrix Double -> Maybe (H.Matrix Double, Int)
inv m = unsafePerformIO $ do
  alloca $ \ outputM -> do
    alloca $ \ inputM -> do
      alloca $ \ den -> do
        let numRows = fromIntegral $ rows m
        let numCols = fromIntegral $ cols m
        fmpz_mat_init outputM numRows numCols
        fmpz_mat_init inputM numRows numCols
        forM_ [0..numRows-1] $ \ i -> do
          forM_ [0..numCols-1] $ \ j -> do
            e <- fmpz_mat_entry inputM i j
            fmpz_set_si e (floor (m `atIndex` (fromIntegral i, fromIntegral j)))
        r <- fmpz_mat_inv outputM den inputM
        d <- peek den

        -- DEBUG:
        fmpz_mat_print_pretty outputM
        -- traceM $ "r = " ++ show r ++ " den = " ++ show d
        --

        lists <- forM [0..numRows-1] $ \ i -> do
          forM [0..numCols-1] $ \ j -> do
            e <- fmpz_mat_entry outputM i j
            fromIntegral `fmap` fmpz_get_si e
        let m' = H.fromLists lists
        fmpz_mat_clear inputM
        fmpz_mat_clear outputM
        if r == 1 then
          return $ Just (m', fromIntegral d)
        else
          return Nothing

withMatrix :: H.Matrix Double -> ((CLong, CLong, Ptr FMPZMat) -> IO b) -> IO b
withMatrix m f = do
  alloca $ \ outputM -> do
    let numRows = fromIntegral $ rows m
    let numCols = fromIntegral $ cols m
    fmpz_mat_init outputM numRows numCols
    forM_ [0 .. numRows - 1] $ \ i ->
      forM_ [0 .. numCols - 1] $ \ j -> do
        e <- fmpz_mat_entry outputM i j
        fmpz_set_si e (floor (m `atIndex` (fromIntegral i, fromIntegral j)) :: CLong)
    x <- f (numRows, numCols, outputM)
    fmpz_mat_clear outputM
    return x

withBlankMatrix :: CLong -> CLong -> (Ptr FMPZMat -> IO b) -> IO b
withBlankMatrix numRows numCols f = do
  alloca $ \ outputM -> do
    fmpz_mat_init outputM (fromIntegral numRows) (fromIntegral numCols)
    x <- f outputM
    fmpz_mat_clear outputM
    return x

withWindow :: Ptr FMPZMat -> CLong -> CLong -> CLong -> CLong -> (Ptr FMPZMat -> IO b) -> IO b
withWindow underM r1 c1 r2 c2 f = do
  alloca $ \ window -> do
    fmpz_mat_window_init window underM r1 c1 r2 c2
    x <- f window
    fmpz_mat_window_clear window
    return x

flintToHMatrix :: CLong -> CLong -> Ptr FMPZMat -> IO (H.Matrix Double)
flintToHMatrix numRows numCols flintM = do
  lists <- forM [0..numRows-1] $ \ i -> do
    forM [0..numCols-1] $ \ j -> do
      e <- fmpz_mat_entry flintM i j
      fromIntegral `fmap` fmpz_get_si e
  return $ H.fromLists lists

pokeM :: Ptr FMPZMat -> CLong -> CLong -> CLong -> IO ()
pokeM flintM i j v = do
  e <- fmpz_mat_entry flintM i j
  fmpz_set_si e v

peekM :: Ptr FMPZMat -> CLong -> CLong -> IO CLong
peekM flintM i j = do
  e <- fmpz_mat_entry flintM i j
  fmpz_get_si e

copyMatrix :: Ptr FMPZMat -> Ptr FMPZMat -> CLong -> CLong -> CLong -> CLong -> IO ()
copyMatrix m1 m2 r1 c1 r2 c2 =
  forM_ [r1 .. r2-1] $ \ i ->
    forM_ [c1 .. c2-1] $ \ j ->
      peekM m2 i j >>= pokeM m1 i j

-- copyMatrix' m1 m2 r1 c1 r2 c2 =
--   withWindow m2 r1 c1 r2 c2 $ \ w2 ->
--     withWindow m1 r1 c1 r2 c2 $ \ w1 -> do
--       fmpz_mat_set w1 w2

--------------------------------------------------
-- 'normalising' Hermite Normal Form, for lack of a better name
--
-- The problem with HNF is that it will happily return a matrix where
-- the leading-coefficient diagonal contains integers > 1.
--
-- In some cases the leading-coefficient does not divide some of the
-- remaining numbers in the row.
--
-- This corresponds to the case where one of the solutions would be
-- fractional if we were dealing with rational matrices.
--
-- But we don't want fractional solutions. We need to bump the matrix
-- so that this situation does not arise.
--
-- This tends to happen due to implicit polymorphism.
--
-- We do this by adding another column that copies the column with the
-- problematic leading-coefficient.
--
-- We then set up a new row expressing the constraint that the old
-- column is equal to the new column, 1:1.
--
-- This prevents a 'fractional' solution.
--
-- Along the way we look for leading-coefficients > 1 that do divide
-- their entire row. Then we simply apply elementary row scaling to
-- the row so that the leading-coefficients is 1.
--
-- The result is a matrix where the leading-coefficients of all the
-- rows that matter are 1. It's not quite a 'reduced' form though
-- because it can be bigger than the original matrix.
--
-- As a result we also return a list of columns that were cloned this
-- way.
--
-- The resulting matrix can have non-zeroes above the
-- leading-coefficients in the same column, so I apply some elementary
-- row operations to fix that and bring things into RREF.
--
-- Running time is computation of Hermite Normal Form twice, plus
-- construction of a slightly larger matrix (possibly), plus scanning
-- through the matrix for leading-coefficients, and then again to
-- divide them out sometimes.

normHNF :: H.Matrix Double -> (H.Matrix Double, [Int])
normHNF m
  | (rows m, cols m) == (1, 1) = (H.ident (if H.atIndex m (0, 0) /= 0 then 1 else 0), [])
  | otherwise = (m', indices)
  where
    numCols = cols m
    indexLookup j | j < numCols = j
                  | otherwise = indices !! (j `mod` numCols)
    indices = map indexLookup $ concatMap snd rs1
    (rs1, (m',[]):_) = break (null . snd) results
    results = tail $ iterate (normHNF' . fst) (m, [])

normHNF' :: H.Matrix Double -> (H.Matrix Double, [Int])
normHNF' m = fmap (map fromIntegral) . unsafePerformIO $ withMatrix m normhnf

normhnf :: (CLong, CLong, Ptr FMPZMat) -> IO (H.Matrix Double, [CLong])
normhnf (numRows, numCols, inputM) = do
  withBlankMatrix numRows numCols $ \ outputM -> do
    fmpz_mat_hnf outputM inputM
    rank <- fmpz_mat_rank outputM
    -- scale the rows so that it is a RREF, returning the indices
    -- where leading-coefficients had to be scaled to 1.
    indices <- elemrowscale outputM rank numCols
    -- HNF allows non-zeroes above the leading-coefficients that
    -- were greater than 1, so also fix that to bring into RREF.
    elemrowadds outputM rank numCols indices

    -- column indices of leading co-efficients > 1
    lcoefs <- filter ((> 1) . head . snd) <$> forM [0 .. rank-1] (\ i -> do
      cs <- zip [0..] `fmap` sequence [ peekM outputM i j | j <- [0 .. numCols-1] ]
      let (_, (j, lcoef):rest) = span ((== 0) . snd) cs
      return ((i, j), lcoef:map snd rest))

    -- split the identified rows into two categories: those that can
    -- be divided out by the leading co-efficient, and those that
    -- cannot.
    let (multCands, consCands) = partition (\ (_, lcoef:rest) -> all ((== 0) . (`rem` lcoef)) rest) lcoefs

    -- apply elementary row scaling
    forM_ multCands $ \ ((i, j), lcoef:_) ->
      forM_ [j..numCols - 1] $ \ j' -> do
        x <- peekM outputM i j'
        pokeM outputM i j' $ x `div` lcoef

    -- identify columns that need additional constraints generated and
    -- their associated value d, which is the value of the non-leading
    -- co-efficient divided by the GCD of the row.
    let genColCons ((_, j), lcoef:rest) = (j, minNLcoef `div` gcd lcoef minNLcoef)
          where
            restABS   = map abs rest
            minNLcoef = minimum (filter (/= 0) restABS)
        genColCons _ = error "normhnf: genColCons: impossible: missing leading co-efficient"

    let consCols = map genColCons consCands

    -- generate operations that poke 1.0 and -d into columns of the
    -- given row (matrix parameter supplied later); d is the value of
    -- the non-leading co-efficient that isn't divisible by the
    -- leading-coefficient, divided by the GCD of the row.
    let ops = [ \ flintM i -> do pokeM flintM i j 1
                                 pokeM flintM i (numCols + k) (-d)
              | ((j, d), k) <- zip consCols [0..] ]
    let numOps = fromIntegral (length ops)
    let numRows' = numOps + rank
    let numCols' = numOps + numCols

    withBlankMatrix numRows' numCols' $ \ outputM' -> do
      -- create larger matrix containing outputM as submatrix
      copyMatrix outputM' outputM 0 0 rank numCols

      -- apply the operations on the extra space to write the
      -- additional rows and columns providing the additional
      -- constraints
      forM_ (zip [rank..] ops) $ \ (i, op) -> op outputM' i

      -- re-run HNF
      withBlankMatrix numRows' numCols' $ \ outputM'' -> do
        fmpz_mat_hnf outputM'' outputM'
        rank' <- fmpz_mat_rank outputM''
        -- scale the rows so that it is a RREF, returning the indices
        -- where leading-coefficients had to be scaled to 1.
        indices' <- elemrowscale outputM'' rank' numCols'
        -- HNF allows non-zeroes above the leading-coefficients that
        -- were greater than 1, so also fix that to bring into RREF.
        elemrowadds outputM'' rank' numCols' indices'
        -- convert back to HMatrix form
        h <- flintToHMatrix rank' numCols' outputM''
        return (h, map fst consCols)

-- find leading-coefficients that are greater than 1 and scale those
-- rows accordingly to reach RREF
--
-- precondition: matrix outputM is in HNF
elemrowscale :: Ptr FMPZMat -> CLong -> CLong -> IO [(CLong, CLong)]
elemrowscale outputM rank numCols = do
  -- column indices of leading co-efficients > 1
  lcoefs <- filter ((> 1) . head . snd) <$> forM [0 .. rank-1] (\ i -> do
    cs <- zip [0..] `fmap` sequence [ peekM outputM i j | j <- [0 .. numCols-1] ]
    let (_, (j, lcoef):rest) = span ((== 0) . snd) cs
    return ((i, j), lcoef:map snd rest))

  let multCands = filter (\ (_, lcoef:rest) -> all ((== 0) . (`rem` lcoef)) rest) lcoefs

  -- apply elementary row scaling
  forM multCands $ \ ((i, j), lcoef:_) -> do
    forM_ [j..numCols - 1] $ \ j' -> do
      x <- peekM outputM i j'
      pokeM outputM i j' $ x `div` lcoef
    return (i, j)

-- use indices to guide elementary row additions to reach RREF
--
-- precondition: matrix outputM is in HNF save for work done by
-- elemrowscale, and indices is a list of coordinates where we have
-- just scaled the leading-coefficient to 1 and now we must look to
-- see if there are any non-zeroes in the column above the
-- leading-coefficient, because that is allowed by HNF.
elemrowadds :: Ptr FMPZMat -> CLong -> CLong -> [(CLong, CLong)] -> IO ()
elemrowadds outputM _ numCols indices = do
  -- look for non-zero members of the columns above the
  -- leading-coefficient and wipe them out.
  forM_ indices $ \ (lcI, lcJ) -> do
    let j = lcJ
    forM_ [0..lcI-1] $ \ i -> do
      -- (i, j) ranges over the elements of the column above leading
      -- co-efficient (lcI, lcJ).
      x <- peekM outputM i j
      if x == 0 then
        pure () -- nothing to do at row i
      else do
        -- (i, j) is non-zero and row i must be cancelled x times
        let sf = x -- scaling factor = non-zero magnitude
        forM_ [lcJ..numCols-1] $ \ j' -> do
          -- (i,   j') ranges over the row where we discovered a non-zero
          -- (lcI, j') ranges over the row with the leading co-efficient
          x1 <- peekM outputM i j'
          x2 <- peekM outputM lcI j'
          -- add lower row scaled by sf to upper row
          pokeM outputM i j' (x1 - x2 * sf)

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

-- m1 :: H.Matrix Double
-- m1 = (8><6)
--  [ 1.0, 0.0, 0.0, -1.0,  0.0,  0.0
--  , 0.0, 1.0, 0.0,  0.0, -1.0,  0.0
--  , 0.0, 0.0, 1.0,  0.0,  0.0, -1.0
--  , 1.0, 0.0, 0.0, -1.0,  0.0,  0.0
--  , 0.0, 1.0, 0.0,  0.0, -1.0,  0.0
--  , 0.0, 0.0, 1.0,  0.0,  0.0, -1.0
--  , 0.0, 0.0, 0.0,  1.0, -4.0,  0.0
--  , 0.0, 0.0, 0.0,  0.0,  4.0, -3.0 ]

-- m2 :: H.Matrix Double
-- m2 = (8><6)
--  [ 1.0, 0.0, 0.0, -1.0,  0.0,  0.0
--  , 0.0, 1.0, 0.0,  0.0, -1.0,  0.0
--  , 0.0, 0.0, 1.0,  0.0,  0.0, -1.0
--  , 1.0, 0.0, 0.0, -1.0,  0.0,  0.0
--  , 0.0, 1.0, 0.0,  0.0, -1.0,  0.0
--  , 0.0, 0.0, 1.0,  0.0,  0.0, -1.0
--  , 0.0, 0.0, 0.0,  1.0, -6.0,  0.0
--  , 0.0, 0.0, 0.0,  0.0,  6.0, -4.0 ]


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}
{- Module      : Text.Loquate.Instances
   Description : Default Loquate Instances
-}
module Text.Loquate.Instances
  (
  ) where

--import qualified Data.ByteString as BS
--import qualified Data.ByteString.Lazy as LBS
import qualified Control.Exception.Base as E
import           Data.Array.Unboxed (Array, IArray, Ix, UArray, elems)
import           Data.Complex (Complex(..))
import           Data.Fixed
import           Data.Foldable
import           Data.Int (Int16, Int32, Int64, Int8)
import           Data.List.NonEmpty (NonEmpty)
import           Data.Ratio (Ratio, denominator, numerator)
import qualified Data.Text as T
import qualified Data.Text.Lazy as LT
import           Data.Text.Prettyprint.Doc (Pretty(..))
import           Data.Typeable
import           Data.Version (Version, showVersion)
import           Data.Void (Void)
import           Data.Word (Word16, Word32, Word64, Word8)
import           GHC.Stack
import           Numeric.Natural
import           Text.Loquate.Class
import           Text.Loquate.Doc

-- misc types

instance Lang l => Loquate l () where loq _ = pretty
instance Lang l => Loquate l Bool where loq _ = pretty
instance Lang l => Loquate l Version where loq _ = pretty . showVersion
instance Lang l => Loquate l Void where loq _ = pretty

-- textual types

instance Lang l => Loquate l Char where loq _ = pretty
instance {-# OVERLAPS #-} Lang l => Loquate l String where loq _ = pretty
-- instance Loquate l BS.ByteString where loq _ = pretty
-- instance Loquate l LBS.ByteString where loq _ = pretty
instance Lang l => Loquate l T.Text where loq _ = pretty
instance Lang l => Loquate l LT.Text where loq _ = pretty

-- numeric types

instance Lang l => Loquate l Double where loq _ = pretty
instance Lang l => Loquate l Float where loq _ = pretty
instance (Lang l, Typeable t, HasResolution t) => Loquate l (Fixed t) where loq _ = viaShow
-- | Integral types should loquate via Integer
instance Lang l => Loquate l Integer where loq _ = pretty
instance (Lang l, Loquate l Integer) => Loquate l Natural where loq l = loq l . toInteger

instance (Lang l, Loquate l Integer) => Loquate l Int where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Int8 where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Int16 where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Int32 where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Int64 where loq l = loq l . toInteger

instance (Lang l, Loquate l Integer) => Loquate l Word where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Word8 where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Word16 where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Word32 where loq l = loq l . toInteger
instance (Lang l, Loquate l Integer) => Loquate l Word64 where loq l = loq l . toInteger

-- polymorphic numeric types

instance Loquate l t => Loquate l (Complex t) where
  loq l (a :+ b) = loq l a <+> "+" <+> loq l b <> "i"

instance Loquate l t => Loquate l (Ratio t) where
  loq l x = loq l (numerator x) <> "/" <> loq l (denominator x)

-- polymorphic containers in dependencies of loquacious

instance Loquate l t => Loquate l (Maybe t) where
  loq _ Nothing  = mempty
  loq l (Just x) = loq l x

instance (Loquate l t, Loquate l s) => Loquate l (t, s) where
  loq l (a, b) = align $ tupled [loq l a, loq l b]

instance (Loquate l t, Loquate l s, Loquate l r) => Loquate l (t, s, r) where
  loq l (a, b, c) = align $ tupled [loq l a, loq l b, loq l c]

instance (Loquate l t) => Loquate l [t] where
  loq l xs = align . list $ loq l <$> xs

instance (Loquate l t) => Loquate l (NonEmpty t) where
  loq l xs = align . list $ loq l <$> toList xs

instance (Typeable i, Loquate l t) => Loquate l (Array i t) where
  loq l xs = align . list $ loq l <$> toList xs

instance (Loquate l t, IArray UArray t, Ix i, Typeable i) => Loquate l (UArray i t) where
  loq l xs = align . list $ loq l <$> elems xs

instance (Loquate l t, Loquate l s) => Loquate l (Either t s) where
  loq l (Left x)  = loq l x
  loq l (Right x) = loq l x

-- showing types rather than values

instance Lang l => Loquate l TypeRep where loq _ = viaShow

instance (Lang l, Typeable a) => Loquate l (Proxy a) where
  loq l _ = loq l $ typeRep (Proxy :: Proxy (Proxy a))

instance (Lang l, Typeable a, Typeable b) => Loquate l (a -> b) where
  loq l f = loq l $ typeOf f

instance (Lang l, Typeable a, Typeable b) => Loquate l (a :~: b) where
  loq l _ = loq l (typeRep (Proxy :: Proxy a)) <+> ":~:" <+> loq l (typeRep (Proxy :: Proxy b))

instance (Lang l, Typeable a, Typeable b) => Loquate l (a :~~: b) where
  loq l _ = loq l (typeRep (Proxy :: Proxy a)) <+> ":~~:" <+> loq l (typeRep (Proxy :: Proxy b))

-- call stacks!

instance Lang l => Loquate l CallStack where
  loq l = withFrozenCallStack ( align . vsep . fmap loqStack . getCallStack )
    where loqStack (func, srcLoc) = fromString func <+> "←" <+> loq l srcLoc

instance Lang l => Loquate l SrcLoc where
  loq l SrcLoc{..} = locFile <+> parens locPackage
    where
      locFile = fromString srcLocFile <> colon <> loq l srcLocStartLine <> colon <> loq l srcLocStartCol
      locPackage = fromString srcLocPackage <> colon <> fromString srcLocModule

-- Exception types

instance Lang l => Loquate l E.AllocationLimitExceeded where loq _ e = viaShow e
instance Lang l => Loquate l E.ArithException where loq _ e = viaShow e
instance Lang l => Loquate l E.ArrayException where loq _ e = viaShow e
instance Lang l => Loquate l E.AssertionFailed where loq _ e = viaShow e
instance Lang l => Loquate l E.AsyncException where loq _ e = viaShow e
instance Lang l => Loquate l E.BlockedIndefinitelyOnMVar where loq _ e = viaShow e
instance Lang l => Loquate l E.BlockedIndefinitelyOnSTM where loq _ e = viaShow e
instance Lang l => Loquate l E.CompactionFailed where loq _ e = viaShow e
instance Lang l => Loquate l E.Deadlock where loq _ e = viaShow e
instance Lang l => Loquate l E.FixIOException where loq _ e = viaShow e
instance Lang l => Loquate l E.IOException where loq _ e = viaShow e
instance Lang l => Loquate l E.NestedAtomically where loq _ e = viaShow e
instance Lang l => Loquate l E.NoMethodError where loq _ e = viaShow e
instance Lang l => Loquate l E.NonTermination where loq _ e = viaShow e
instance Lang l => Loquate l E.PatternMatchFail where loq _ e = viaShow e
instance Lang l => Loquate l E.RecConError where loq _ e = viaShow e
instance Lang l => Loquate l E.RecSelError where loq _ e = viaShow e
instance Lang l => Loquate l E.RecUpdError where loq _ e = viaShow e
instance Lang l => Loquate l E.SomeAsyncException where loq _ e = viaShow e
instance Lang l => Loquate l E.SomeException where loq _ e = viaShow e
instance Lang l => Loquate l E.TypeError where loq _ e = viaShow e


-- |
-- Module      : BattleHack.Types
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental
-- Portability : POSIX

-- Created September 12 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module BattleHack.Types where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import qualified Data.Set as S
import qualified Data.Map as M
import Sound.OpenAL

--------------------------------------------------------------------------------------------------------------------------------------------
-- Types
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
type Number = Double
type Vector = Complex Number
type Sample = Double

-- |
type Command = IO ()
type KeyMap  = M.Map String Command

-- |
data KeyLayout = KeyLeft | KeyRight | KeyBoth | KeyAccidental deriving (Show)

-- |
type Claviature = [(Source, Buffer)]

-- |
-- data Note = C Pitch | D Pitch | E Pitch | F Pitch | G Pitch | A Pitch | B Pitch deriving (Show)
-- data Pitch = Sharp | Flat | Natural                                             deriving (Show)

-- |
-- frame number, fps, etc.
data AnimationData = AnimationData { _frame :: Int,
                                     _fps   :: Int }

-- |
-- TODO: Slightly less confusing names (piano vs claviature)
data AppState = AppState { _piano      :: PianoSettings,
                          --  _claviature :: Claviature,
                          --  _source     :: Source,
                           _animation  :: AnimationData,
                           _bindings   :: KeyMap,
                           _inputstate :: InputState } -- deriving (Show)

-- |
-- TODO: Use custom key type or text (?)
data InputState =  InputState { _mouse    :: Vector,
                                _keyboard :: S.Set String }

-- |
data PianoSettings = PianoSettings { _origin  :: Vector,
                                     _keysize :: Vector,
                                     _indent  :: Number,
                                     _mid     :: Number,
                                     _active  :: Maybe Int,
                                     _keys    :: [Bool] } deriving (Show)


{-# Language DeriveGeneric     #-}
{-# Language OverloadedStrings #-}
{-# Language DeriveFunctor     #-}
{-# Language DeriveFoldable    #-}
{-# Language RecordWildCards   #-}

module Lib where

import           Control.Arrow                  ( (&&&) )
import           Data.Csv
import qualified Data.ByteString.Lazy          as BL
import qualified Data.Vector                   as V
import           Data.Either                    ( fromRight )
import           GHC.Generics
import           Statistics.Sample
import           System.IO.Unsafe
import           Data.Foldable                  ( toList
                                                , foldl'
                                                )
import           Data.List                      ( transpose )

-- The definition of a spambase entry.
-- It is SpamEntry a as that allows me to be able to map over the data easily.
type SpamEntry = SpamE Double -- Alias for nicer type signatures
data SpamE a =
  S {
    _make        ::  !a
  , _address     ::  !a
  , _all         ::  !a
  , _3d          ::  !a
  , _our         ::  !a
  , _over        ::  !a
  , _remove      ::  !a
  , _internet    ::  !a
  , _order       ::  !a
  , _mail        ::  !a
  , _receive     ::  !a
  , _will        ::  !a
  , _people      ::  !a
  , _report      ::  !a
  , _addresses   ::  !a
  , _free        ::  !a
  , _business    ::  !a
  , _email       ::  !a
  , _you         ::  !a
  , _credit      ::  !a
  , _your        ::  !a
  , _font        ::  !a
  , _000         ::  !a
  , _money       ::  !a
  , _hp          ::  !a
  , _hpl         ::  !a
  , _george      ::  !a
  , _650         ::  !a
  , _lab         ::  !a
  , _labs        ::  !a
  , _telnet      ::  !a
  , _857         ::  !a
  , _data        ::  !a
  , _415         ::  !a
  , _85          ::  !a
  , _technology  ::  !a
  , _1999        ::  !a
  , _parts       ::  !a
  , _pm          ::  !a
  , _direct      ::  !a
  , _cs          ::  !a
  , _meeting     ::  !a
  , _original    ::  !a
  , _project     ::  !a
  , _re          ::  !a
  , _edu         ::  !a
  , _table       ::  !a
  , _conference  ::  !a
  , _semicolon   ::  !a
  , _lparen      ::  !a
  , _pipe        ::  !a
  , _bang        ::  !a
  , _dollarsign  ::  !a
  , _hash        ::  !a
  , _caps_avg    ::  !a
  , _caps_max    ::  !a
  , _caps_total  ::  !a
  , _spam        ::  !Int
} deriving (Generic, Show, Functor, Foldable)
  -- Gives me mapping, toList, folding, decoding, etc., for free.

-- CSV decoding instanes for free because of Generic
instance FromField a => FromRecord (SpamE a)
instance ToField   a => ToRecord   (SpamE a)

--
-- This is written in a semi unidiomatic haskell style where I don't really
-- have a main function that assembles everything together.  Instead, I use
-- unsafePerformIO to grab the data and treat it as a pure value so I can just
-- hardcode the assignment here. It allows for faster iteratino time.
--

-- Decodes the file into a vector of spambase entries
{-# NOINLINE spamEntries #-}
spamEntries :: V.Vector SpamEntry
spamEntries = fromRight undefined . decode NoHeader $ mydata
 where
  -- Load the data from a file. Bit of a hacky way to do it, but it makes the
  -- rest of the code easier to incrementally write.
  mydata :: BL.ByteString
  mydata = unsafePerformIO $ BL.readFile "spambase.data"
  {-# NOINLINE mydata #-} -- if you inline this function, bad shit happens

-- This works by indexing the vector of SpamEntry entries
-- Then the vector is split into two vectors, one with all the even indices,
-- one with all odd (this effectively splits the data sets equally)
--
-- Then the extra indices are stripped from both vectors
test, train :: V.Vector SpamEntry
(test, train) = splitData spamEntries
 where
  splitData
    :: V.Vector SpamEntry
    -> (V.Vector SpamEntry, V.Vector SpamEntry)
  splitData e = both (snd <$>) $ V.partition (even . fst) (V.indexed e)

--
-- Probabilities of the prior
--
probSpam :: V.Vector SpamEntry -> Double
probSpam s = V.length (fst $ splitSpam s) // V.length s

pSpamTrain, pSpamTest :: Double
(pSpamTrain, pSpamTest) = both probSpam (train, test)

-- Takes the spam entries and collects the (mean, standard deviation) of all 57
-- features
stats :: Foldable t => V.Vector (t Double) -> [(Double, Double)]
stats entries = (mean &&& stdDev) <$> collectedStats
 where
  --  This is just data shuffling to get all 57 features from the different
  --  spamEntries into a list of feature-vectors
  collectedStats :: [V.Vector Double]
  collectedStats = V.fromList . normalize <$$> transpose $ toList <$> V.toList entries
  -- Normalize is a safety precaution to prevent zeroes
  normalize       = fmap (\x -> if x < 0.0001 then 0.0001 else x)

-- Splits into a tuple of (spam, non spam)
splitSpam :: V.Vector SpamEntry -> (V.Vector SpamEntry, V.Vector SpamEntry)
splitSpam = V.partition ((== 1) . _spam)

-- Just some globals to use the functions just defined above
trainSpamStat, trainNonStat, testSpamStat, testNonStat :: [(Double, Double)]
(trainSpamStat, trainNonStat) = both stats $ splitSpam train
(testSpamStat, testNonStat) = both stats $ splitSpam test

--  This is the N(xi;μ,σ) function from the text
normal :: Floating a => (a, a) -> a -> a
normal (μ, σ) x = 1 / root * exp (-(x - μ) ^ 2 / (2 * σ ^ 2))
  where root = sqrt (2 * pi) * σ

--
-- The meat of the library. This classifies things as spam or nonSpam
--
data Class = Spam | NonSpam deriving (Show)

classify :: SpamEntry -> Class
classify e = if testSpamLogs > testNonLogs then Spam else NonSpam
 where
  testSpamLogs = log pSpamTrain + sum (logify testSpamStat)
  testNonLogs  = log (1 - pSpamTrain) + sum (logify testNonStat)

  logify :: [(Double, Double)] -> [Double]
  logify stat = log <$> zipWith ($) (normal <$> stat) (toList e)

classify' :: SpamEntry -> Accuracy
classify' e = case (guess, actual) of
  (Spam   , Spam   ) -> TruePos
  (NonSpam, NonSpam) -> TrueNeg
  (Spam   , NonSpam) -> FalsePos
  (NonSpam, Spam   ) -> FalseNeg
 where
  guess  = classify e
  actual = if _spam e == 1 then Spam else NonSpam

--
-- Accuracy and analysis section
--
data Accuracy = TruePos | TrueNeg | FalsePos | FalseNeg deriving (Show)

-- The higher-ordered data structure is, again, so I can map over the data
-- easily.  foldl' (+) accuracyData 0 will sum up all of the items in this data
-- structure, for example.
type AccuracyData = ACount Int
data ACount a =
  ACount {
    _truePos  :: !a
  , _trueNeg  :: !a
  , _falsePos :: !a
  , _falseNeg :: !a
} deriving (Show, Functor, Foldable)

-- Default starting point
accuracyData :: AccuracyData
accuracyData = ACount 0 0 0 0

-- Updater functions: succ a <=> a+1
truePos :: AccuracyData -> AccuracyData
truePos a@ACount {..} = a { _truePos = succ _truePos }

trueNeg :: AccuracyData -> AccuracyData
trueNeg a@ACount {..} = a { _trueNeg = succ _trueNeg }

falsePos :: AccuracyData -> AccuracyData
falsePos a@ACount {..} = a { _falsePos = succ _falsePos }

falseNeg :: AccuracyData -> AccuracyData
falseNeg a@ACount {..} = a { _falseNeg = succ _falseNeg }

-- Go over all of the classified data; categorize it into the confusion matrix,
-- then tally up all of the results.
tally :: AccuracyData
tally = foldl' accum accuracyData classes
 where
  classes = classify' <$> test
  accum a e = case e of
    TruePos  -> truePos a
    TrueNeg  -> trueNeg a
    FalsePos -> falsePos a
    FalseNeg -> falseNeg a

--
-- The statistics:
--
accuracy :: Double
accuracy = allPos // total
 where
  allPos = _truePos tally + _trueNeg tally
  total  = foldl' (+) 0 tally

recall :: Double
recall = tp // (tp + _falseNeg tally) where tp = _truePos tally

precision :: Double
precision = tp // allPos
 where
  tp     = _truePos tally
  allPos = tp + _falsePos tally

--
-- Main
--
main = do
  putStrLn $ "The accuracy is: "         ++ show accuracy
  putStrLn $ "The recall is: "           ++ show recall
  putStrLn $ "The precision is: "        ++ show precision
  putStrLn $ "The confusion matrix is: " ++ show tally

--
-- Helper functions
--

-- Apply a function to both sides of a tuple at once
both :: (t -> b) -> (t, t) -> (b, b)
both f (a, b) = (f a, f b)

-- Nested mapping
infixl 5 <$$>
(<$$>) :: (Functor f2, Functor f1) => (a -> b) -> f1 (f2 a) -> f1 (f2 b)
(<$$>) = fmap fmap fmap

-- Division with implicit number conversion
(//) :: (Integral a1, Integral a, Fractional a2) => a -> a1 -> a2
a // b = fromIntegral a / fromIntegral b


{-#Language TupleSections, MultiWayIf#-}
module Stratification where

-- Внешние импорты
import Data.Matrix (Matrix(..))

import Data
import Point
import Data.Array as A
import Data.Map as M
import Data.Complex
import SymbolicImage
import Jac (jac, e)

str :: F Point -> Imagination -> Int -> Stratification
str fp im i = iterate
    (stepStratification fp)
    (toStratification im) !! i

-- Один шаг с подразбиением.
stepStratification :: F Point-> F Stratification
stepStratification fp (Stratification d1 d2 ls) = Stratification
    d1 (d2*2) $ concatMap celling $ cyclics $
    (shiftCeil d1 d2 fp #=) <$> ls
  where

{-
  where
    -- считаем граф
    !gr     = (shiftCeil d1 d2 fp ##=) <$> ls
    -- индех вершин
    !index  = formArray gr
    -- обратный индех
    !indexR = formMap index

    unpack !gr = (index A.!) <$> gr
    pack !gr = do
        (a, b) <- gr
        let a' = fromJust $ a`M.lookup`indexR
        return (a', a', MB.mapMaybe (`M.lookup`indexR) b)

(##=) !f !c = (c, f c)

-}

-- вспомогательные структуры данных
-- массив с ячейками п индексам.
formArray :: [(Ceil3, [Ceil3])] -> Array Int Ceil3
formArray gr = array (1, length a) (zip [1..] a)
  where a = plat gr

-- находим индекс по ячейке.
formMap :: Array Int Ceil3 -> Map Ceil3 Int
formMap  arr = M.fromList $ (\(a,b) -> (b,a)) <$> A.assocs arr

plat :: Ord c => [(c, [c])] -> [c]
plat a = myNub.concat $ (\(a, b) -> a:b) <$> a

toStratification :: Imagination -> Stratification
toStratification (Imagination d c) = Stratification d 1
    $ concatMap (\c -> [Ceil3 c 1, Ceil3 c 0]) c

-- подразбиение на более мелкие уровни
celling :: Ceil3 -> [Ceil3]
celling (Ceil3 c i) = [Ceil3 c (i*2), Ceil3 c (i*2-1)]

-- Тут что-то не то.
shiftCeil :: Diameter -> Diameter -> F Point -> Ceil3 -> [Ceil3]
shiftCeil d1 d2 fp (Ceil3 c l) = myNub $ do
    let Point x y = fromCeil d1 c
    func d1 d2 fp
        <$> mySequence x d1
        <*> mySequence y d1
        <*> mySequence (rad d2 l) (pi/2/d2)

func :: Double -> Double -> F Point -> Double -> Double -> Double -> Ceil3
func d1 d2 fp x y t = Ceil3 (ceil d1 fp x y) (int d2 (arcOf fp x y t))

arcOf :: F Point -> Double -> Double -> Double -> Double
arcOf fp x y t = arc $ toPoint $ jacobian fp x y * e t

ceil :: Double -> F Point -> Double -> Double -> Ceil
ceil d fp x y = toCeil d $ fp (Point x y)

jacobian :: F Point -> Double -> Double -> Matrix Double
jacobian fp x y = jac fp $ Point x y

arc :: Point -> Double
arc (Point x y) = if
    | x >  0    ->       phase $          x :+ y
    | x == 0    -> abs $ phase $          x :+ y
    | otherwise ->       phase $ negate $ x :+ y

int :: Double -> Double -> Int
int d r = ceiling $ r/pi*2*d

rad :: Double -> Int ->Double
rad d r = toEnum r * pi/2/d


module Statistics.Quantile.Util where

import Pipes
import qualified Pipes.Prelude as P

import Statistics.Quantile.Types
import System.IO

tmpDir :: FilePath
tmpDir = "/var/tmp/"

selectFromHandle :: Quantile
                 -> Selector IO
                 -> Handle
                 -> IO Double
selectFromHandle q (Selector select) hin = do
  select q $ streamHandle hin

streamHandle :: Handle -> Stream IO
streamHandle hin = Stream (P.fromHandle hin >-> P.map read)

takeStream :: Int -> Stream IO -> Stream IO
takeStream n (Stream p) = Stream $ p >-> P.take n


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, ConstraintKinds, ExistentialQuantification, GADTs, RankNTypes, MultiParamTypeClasses   #-}
{-# LANGUAGE ImpredicativeTypes, RankNTypes, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Count-Min Sketch
-- | Creator: Xiao Ling
-- | Created: 12/17/2015
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module CountMinSketch where

import Control.Monad.Random.Class
import Control.Monad.Random
import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Identity

import Data.Matrix
import Data.Conduit
import qualified Data.Conduit.List as Cl
import qualified Data.Matrix as M

import Core
import Utils
import Statistics

import System.IO.Unsafe

{-----------------------------------------------------------------------------
  Types
------------------------------------------------------------------------------}

type W        = Int              -- * num cols
type D        = Int              -- * num rows
type P        = Int              -- * some prime   --> change this to Integer
type Event    = Int              -- * some input event from
                                 -- * universe: {1, .., i, .. n}

type Sketch   = Matrix Int       -- * D x W matrix  of counters
type ABs      = Matrix Int       -- * 2 x D matrix  of a_j and b_j for j = 1..d
type ColIdx   = [Int]            -- * indices of counters to increment
type Hash     = Event -> ColIdx  -- * Maps event to d x 1 list of column indices

-- * Algorithm State
type AState   = (Sketch,Hash,(D,W,P))

-- * the algorithm uses `some m` equipped with state and randomness
-- * also consider `some'` `a`
type Some m   = (MonadRandom m, MonadState AState m)

{-----------------------------------------------------------------------------
    Count Min Sketch
------------------------------------------------------------------------------}

sk :: Batch Event IO AState
sk = sketch' ed 50

sketch' :: EpsDelta -> Event -> Batch Event IO AState
sketch' t n = sketch t n `using` eval

sketch :: Some m => EpsDelta -> Event -> Streaming Event m ()
sketch t n = lift (initDwp t n) $= inits False $= update

-- * there should be a combinator for this
update :: Some m => Streaming Event m ()
update = do
    mi <- await
    case mi of
        Nothing -> return ()
        Just i  -> lift (update' i) >> update

-- * there should be a combinator that does this
inits :: Some m => Bool -> Conduit Event m Event
inits False = lift inits' >> inits True
inits _     = pass

---- * a concrete `eval`uation of Some `m`
-- * run with vacuous inital condition to be thrown away
eval :: StateT AState (Rand StdGen) () -> IO AState
eval m = evalRandIO $ execStateT m $ (ones 1 1, const [], (0,0,0))

{-----------------------------------------------------------------------------
   Effectful Subroutines
   tenative conclusion, it works fine in here
--------------------------------------------------------------------------------}

initDwp :: Some m => EpsDelta -> Event -> m ()
initDwp (ED e d) n = setParam $ (round $ log2(1/d), round $ 2/e, round $ 2^n -1)

-- * initalize algorithm with zero counters and hash function
inits' :: Some m => m ()
inits' = do
    t       <- askParam
    hash    <- newHash t
    let sketch = newSketch t
    setHash   hash
    setSketch sketch

-- * on event `i`, update the `sketch`
update' :: Some m => Event -> m ()
update' i = do
    pm      <- askParam
    sketch  <- getSketch
    hash    <- askHash
    let incr = toMask (hash i) pm
    setSketch $ incr .+. sketch

-- * sketch1 == sketch1
-- * IO (sketch1', hash', t')
m1 = evalRandIO $ execStateT (update' i1) (sketch0, hash, tup)
m2 = evalRandIO $ execStateT (update' i2) (sketch1, hash, tup)

-- * Estimate
--estimate :: Some m => Event -> m ()
--estimate = do
--    ()

-- * intialize random cofficients a_j, b_j, for d x 2 matrix:
-- * [ a_1 ... a_j ... a_d ]
-- * [ b_1 ... b_j ... b_d ]
newAbs :: MonadRandom m => (D,W,P) -> m ABs
newAbs (d,_,p) = (M.fromList 2 d . take (2*d)) <$> getRandomRs (1,p)

{-----------------------------------------------------------------------------
    Pure Subroutines
------------------------------------------------------------------------------}

newDwp :: EpsDelta -> Event -> (D,W,P)
newDwp (ED e d) n = (round $ log2(1/d), round $ 2/e, round $ 2^n -1)

-- * initalize the sketch
newSketch :: (D,W,P) -> Sketch
newSketch (d,w,_) = zeros d w

-- * h_j (i) = (a_j x i + b_j mod p) mod w
-- * sumRow-wise  [ a_1 ... a_j ... a_d ]  .*.  [ i_1 ...  i_d ]     mod p  mod w
-- *              [ b_1 ... b_j ... b_d ]       [ 1   ...  1   ]
-- * result:      [ idx_1, ..., idx_d   ]
-- *               where the indices start at 1
newHash :: MonadRandom m => (D,W,P) -> m Hash
newHash t = do
    ab  <- newAbs t
    return $ toHash' t ab

-- * Note `is` maps some `i` from event universe onto d x 2 matrix:
-- * [ i_1 ... i_2 ]
-- * [ 1   ...  1  ]
toHash' :: (D, W, P) -> ABs -> Hash
toHash' (d, w, p) ab i = go $ sumR (ab .*. is) ... (\v -> v `mod` p `mod` w)
    where go = fmap (+1) . toList
          is = (i .+ (zeros 1 d)) <-> ones 1 d

-- * given hash function, construct a binary mask of size `w` x `d`
-- * to increment sketch
toMask :: ColIdx -> (D,W,P) -> Sketch
toMask ks (d,w,p) = go w ks
    where
        go w ks = let vec c w = zeros 1 (c-1) <|> single 1 <|> zeros 1 (w - c) in
            case ks of
                c:[] -> vec c w
                c:cs -> vec c w <-> go w cs

{-----------------------------------------------------------------------------
    Ad hoc testing
------------------------------------------------------------------------------}

[i1,i2,i3] = [12,0,50]

tup@(d,w,p) = newDwp ed 50

hash      = unsafePerformIO . evalRandIO $ newHash tup
sketch0   = newSketch tup

-- * update it once
ks1       = hash i1
incr1     = toMask ks1 tup
sketch1   = incr1 .+. sketch0

-- * update again
ks2       = hash i2
incr2     = toMask ks2 tup
sketch2   = incr2 .+. sketch1

{-----------------------------------------------------------------------------
    Managing State
    TODO : redesign so this is not necessary
------------------------------------------------------------------------------}

askParam :: Some m => m (D,W,P)
askParam = get >>= \(_,_,p) -> return p

setParam :: Some m => (D,W,P) -> m ()
setParam t = get >>= \(s,h,_) -> put (s,h,t)

getSketch :: Some m => m Sketch
getSketch = get >>= \(s,_,_) -> return s

setSketch :: Some m => Sketch -> m ()
setSketch s = get >>= \(_,h,t) -> put (s,h,t)

askHash :: Some m => m Hash
askHash = get >>= \(_,h,_) -> return h

setHash :: Some m => Hash -> m ()
setHash h = get >>= \(s,_,t) -> put (s,h,t)

{-----------------------------------------------------------------------------
    Find a lib function for this:
------------------------------------------------------------------------------}

-- * identity conduit, why can't I find this in the library?
pass :: Monad m => Conduit i m i
pass  = do
    mx <- await
    case mx of
        Nothing -> return ()
        Just x  -> yield x >> pass