{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict #-}
module STC.Convolution where

import           Control.Arrow
import           Control.Monad              as M
import           Control.Monad.Parallel     as MP (bindM2, mapM)
import           Data.Array.IArray          as IA
import           Data.Array.Repa            as R
import           Data.Complex
import           Data.Ix
import           Data.List                  as L
import           Data.Vector.Storable       as VS
import           Data.Vector.Unboxed        as VU
import           Debug.Trace
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.FourierSeries
import           Image.IO
import           STC.DFTArray
import           Types
import           Utils.Parallel
import FokkerPlanck.GPUKernel
import qualified Data.Array.Accelerate              as A
import qualified Data.Array.Accelerate.Data.Complex as A
import           Data.Array.Accelerate.LLVM.PTX
import Debug.Trace
import Utils.List
import Utils.BLAS

data Field
  = Source
  | Sink
  deriving (Read, Show)

{-# INLINE dftHarmonicsArray #-}
dftHarmonicsArray ::
     DFTPlan
  -> Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> IO (IA.Array (Int, Int) (VS.Vector (Complex Double)))
dftHarmonicsArray !plan !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff = do
  let !arr =
        computeHarmonicsArray
          numRows
          deltaRow
          numCols
          deltaCol
          phiFreqs
          rhoFreqs
          thetaFreqs
          rFreqs
          halfLogPeriod
          cutoff
  fmap (listArray (bounds arr)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
    L.map
      (VS.convert .
       toUnboxed .
       computeUnboxedS . makeFilter2D . fromUnboxed (Z :. numCols :. numRows)) .
    IA.elems $
    arr

{-# INLINE dftHarmonicsArrayG #-}
dftHarmonicsArrayG ::
     DFTPlan
  -> Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double -> VS.Vector (Complex Double)
  -> IO (IA.Array (Int, Int) (VS.Vector (Complex Double)))
dftHarmonicsArrayG !plan !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff !gaussian = do
  let !arr =
        computeHarmonicsArray
          numRows
          deltaRow
          numCols
          deltaCol
          phiFreqs
          rhoFreqs
          thetaFreqs
          rFreqs
          halfLogPeriod
          cutoff
  xs <-
    fmap (L.map (VS.zipWith (*) gaussian)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
    IA.elems $
    arr
  ys <- dftExecuteBatchP plan (DFTPlanID IDFT1DG [numCols, numRows] [0, 1]) xs
  fmap (listArray (bounds arr)) .
    dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
    L.map
      (\vec ->
         let arr = fromUnboxed (Z :. numCols :. numRows) . VS.convert $ vec
         in VS.convert .
            toUnboxed . computeUnboxedS . makeFilter2D . R.traverse arr id $ \f idx@(Z :. i :. j) ->
              if (i == div numCols 2 && j == div numRows 2)
                 then 0
                 else f idx) $
    ys

{-# INLINE convolve #-}
convolve ::
     Field
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> IO DFTArray
convolve !field !plan !coefficients !harmonicsArray !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
  dftVecs <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
  let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
      idx = (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)
  fmap (DFTArray rows cols thetaFreqs rFreqs) .
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
    parMap
      rdeepseq
      (\((!r, !rFreq), (!theta, !thetaFreq)) ->
         L.foldl'
           (\(!vec) (((!rho, !rhoFreq), (!phi, !phiFreq)), inputVec) ->
              VS.zipWith
                (+)
                vec
                (VS.map (* (coefficients R.! (Z :. r :. theta :. rho :. phi))) .
                 VS.zipWith
                   (*)
                   (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq) $
                 inputVec))
           initVec .
         L.zip idx $
         dftVecs) $
    idx

{-# INLINE convolve' #-}
convolve' ::
     Field
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> IO DFTArray
convolve' !field !plan !coefficients !harmonicsArray !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
  dftVecs <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
  let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
      idxTheta = L.zip [0 ..] thetaFreqs
  fmap (DFTArray rows cols thetaFreqs rFreqs) .
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
    parMap
      rdeepseq
      (\(!theta, !thetaFreq) ->
         VS.map
           (\x ->
              case field of
                Source -> x
                Sink -> x * cis (pi * thetaFreq)) .
         L.foldl'
           (\vec1 (rho, rhoFreq) ->
              L.foldl'
                (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                   VS.zipWith
                     (+)
                     vec2
                     (VS.map
                        (* (coefficients R.!
                            (Z :. (0 :: Int) :. theta :. rho :. phi))) .
                      VS.zipWith
                        (*)
                        (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq 0) $
                      inputVec))
                vec1 .
              L.zip idxTheta $
              dftVecs)
           initVec .
         L.zip [0 ..] $
         rFreqs) $
    idxTheta

{-# INLINE convolveG' #-}
convolveG' ::
     Field
  -> DFTPlan
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> VS.Vector (Complex Double)
  -> DFTArray
  -> IO DFTArray
convolveG' !field !plan !coefficients !harmonicsArray !gaussianFilter !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
  dftVecs <-
    L.map (VS.zipWith (*) gaussianFilter) <$>
    dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
  let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
      idxTheta = L.zip [0 ..] thetaFreqs
  fmap (DFTArray rows cols thetaFreqs rFreqs) .
    dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
    parMap
      rdeepseq
      (\(!theta, !thetaFreq) ->
         L.foldl'
           (\vec1 (rho, rhoFreq) ->
              L.foldl'
                (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                   VS.zipWith
                     (+)
                     vec2
                     (VS.map
                        (* (case field of
                              Source ->
                                coefficients R.!
                                (Z :. (0 :: Int) :. theta :. rho :. phi)
                              Sink ->
                                (coefficients R.!
                                 (Z :. (0 :: Int) :. theta :. rho :. phi)) *
                                (cis $ (-(phiFreq + thetaFreq)) * pi))) .
                      VS.zipWith
                        (*)
                        (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq 0) $
                      inputVec))
                vec1 .
              L.zip idxTheta $
              dftVecs)
           initVec .
         L.zip [0 ..] $
         rFreqs) $
    idxTheta

-- {-# INLINE convolve'' #-}
-- convolve'' ::
--      Field
--   -> DFTPlan
--   -> R.Array U DIM4 (Complex Double)
--   -> IA.Array (Int, Int) (VS.Vector (Complex Double))
--   -> DFTArray
--   -> IO DFTArray
-- convolve'' !field !plan !coefficients !harmonicsArray !arr@(DFTArray rows cols phiFreqs rhoFreqs vecs) = do
--   let !initVec = VS.replicate (VS.length . L.head $ vecs) 0
--       idx = L.zip [0 ..] phiFreqs
--   return .
--     (DFTArray rows cols phiFreqs rhoFreqs) .
--     parMap
--       rdeepseq
--       (\(theta, thetaFreq) ->
--          L.foldl'
--            (\vec1 (rho, rhoFreq) ->
--               L.foldl'
--                 (\vec2 (phi, phiFreq) ->
--                    VS.zipWith (+) vec2 .
--                    VS.map
--                      (* (case field of
--                            Source ->
--                              coefficients R.!
--                              (Z :. (0 :: Int) :. theta :. rho :. phi)
--                            Sink ->
--                              (coefficients R.!
--                               (Z :. (0 :: Int) :. theta :. rho :. phi)) *
--                              (cis $ (-(thetaFreq + phiFreq)) * pi))) $
--                    (getHarmonics harmonicsArray phiFreq rhoFreq 0 0))
--                 vec1
--                 idx)
--            initVec .
--          L.zip [0 ..] $
--          rhoFreqs) $
--     idx

-- {-# INLINE dftHarmonicsArrayGPU #-}
-- dftHarmonicsArrayGPU ::
--      DFTPlan
--   -> Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IO (Acc (A.Array A.DIM4 (A.Complex Double)))
-- dftHarmonicsArrayGPU !plan !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff = do
--   let !arr =
--         computeHarmonicsArrayGPU
--           numRows
--           deltaRow
--           numCols
--           deltaCol
--           phiFreqs
--           rhoFreqs
--           thetaFreqs
--           rFreqs
--           halfLogPeriod
--           cutoff
--       rangeFunc2 xs ys =
--         (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
--       (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
--       (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
--   fmap
--     (A.use .
--      A.fromList
--        (A.Z A.:. (rfUB - rfLB + 1) A.:. (tfUB - tfLB + 1) A.:. numCols A.:.
--         numRows) .
--      VS.toList . VS.concat) .
--     dftExecuteBatchP plan (DFTPlanID DFT1DG [numCols, numRows] [0, 1]) .
--     L.map
--       (VS.convert .
--        toUnboxed .
--        computeUnboxedS . makeFilter2D . fromUnboxed (Z :. numCols :. numRows)) .
--     IA.elems $
    -- arr

-- {-# INLINE convolveGPU #-}
-- convolveGPU ::
--      PTX
--   -> Field
--   -> DFTPlan
--   -> Acc (A.Array A.DIM4 (A.Complex Double))
--   -> Acc (A.Array A.DIM4 (A.Complex Double))
--   -> Acc (A.Array A.DIM4 (A.Complex Double))
--   -> DFTArray
--   -> IO DFTArray
-- convolveGPU !ptx !field !plan !coefficients !coefficientsSink !harmonics !arr@(DFTArray rows cols thetaFreqs rFreqs vecs) = do
--   dftVecs <- dftExecuteBatchP plan (DFTPlanID DFT1DG [cols, rows] [0, 1]) vecs
--   let !numRFreq = L.length rFreqs
--       !numThetaFreq = L.length thetaFreqs
--       idx =
--         [ (r, theta)
--         | r <- [0 .. (numRFreq - 1)]
--         , theta <- [0 .. (numThetaFreq - 1)]
--         ]
--   case field of
--     Source ->
--       fmap (DFTArray rows cols thetaFreqs rFreqs) .
--       dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
--       L.map
--         (\(r, theta) ->
--            VS.fromList . A.toList $
--            (runNWith ptx (convolveKernel coefficients harmonics))
--              (A.fromList A.Z [theta])
--              (A.fromList A.Z [r])
--              (A.fromList
--                 (A.Z A.:. numRFreq A.:. numThetaFreq A.:. cols A.:. rows) .
--               VS.toList . VS.concat $
--               dftVecs)) $
--       idx
--     -- Sink ->
--     --   fmap (DFTArray rows cols thetaFreqs rFreqs) .
--     --   dftExecuteBatchP plan (DFTPlanID IDFT1DG [cols, rows] [0, 1]) .
--     --   L.map
--     --     (\(r, theta) ->
--     --        VS.fromList .
--     --        A.toList .
--     --        (runNWith ptx (convolveKernel coefficientsSink harmonics))
--     --          (A.fromList (Z :. (1 :: Int)) theta)
--     --          (A.fromList (Z :. (1 :: Int)) r)
--     --          (A.fromList (Z :. numRFreq :. numThetaFreq :. cols :. rows) .
--     --           VS.toList . VS.concat $
--     --           dftVecs)) $
--     --   idx

{-# INLINE convolveSingle #-}
convolveSingle ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> Double
  -> R.Array U DIM2 (Complex Double)
  -> VU.Vector (Complex Double)
convolveSingle !field !coefficients !harmonicsArray !thetaFreqs !rFreqs !x !y !input =
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
      rowCenter = div rows 2
      colCenter = div cols 2
      (Z :. _ :. _ :. numRFreq :. numThetaFreq) = extent coefficients
      -- harmonics =
      --   R.traverse2 thetaFreqs rFreqs (\_ _ -> extent coefficients) $ \fThetaFreq fRFreq (Z :. r :. theta :. rho :. phi) ->
      --     (harmonicsArray IA.!
      --      ( round (fRFreq (Z :. rho) - fRFreq (Z :. r))
      --      , round $ (fThetaFreq (Z :. phi) - (fThetaFreq (Z :. theta))))) R.!
      --     (Z :. (x + colCenter) :. (y + rowCenter))
      harmonics =
        R.traverse2 thetaFreqs rFreqs (\_ _ -> extent coefficients) $ \fPhiFreq fRhoFreq (Z :. r :. theta :. rho :. phi) ->
          let !tf = fPhiFreq (Z :. phi) - fPhiFreq (Z :. theta)
              !rf = fRhoFreq (Z :. rho) - fRhoFreq (Z :. r)
          in (x :+ y) ** (tf :+ 0) *
             ((x ^ 2 + y ^ 2) :+ 0) **
             (((-tf - 1) :+ 2 * pi * rf / (log $ (sqrt 2) * 512)) / 2)
      product =
        R.traverse2 (R.zipWith (*) coefficients harmonics) input const $ \fCoefficient fInput idx@(Z :. _ :. _ :. rho :. phi) ->
          fCoefficient idx * fInput (Z :. rho :. phi)
      arr =
        case field of
          Source -> product
          Sink ->
            error
              "convolveSingle: You should not use it to compute a sink field."
  in if x == 0 && y == 0
       then VU.replicate (numRFreq * numThetaFreq) 0
       else toUnboxed . sumS . sumS $ arr

-- {-# INLINE convolveSparse #-}
convolveSparse ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> [R.Array U DIM2 (Complex Double)]
  -> [R.Array U DIM2 (Complex Double)]
convolveSparse !field !coefficients !harmonicsArray !thetaFreqs !rFreqs !cutoff !xs !inputs =
  L.map (fromUnboxed (extent . L.head $ inputs)) .
  parMap
    rdeepseq
    (\(x, y) ->
       L.foldl1' (VU.zipWith (+)) .
       L.zipWith
         (\(x', y') input ->
            let !x'' = x - x'
                !y'' = y - y'
                (!a, !b) =
                  if (x'' ^ 2 + y'' ^ 2) <= cutoff ^ 2
                    then (x'', y'')
                    else (0, 0)
            in convolveSingle
                 field
                 coefficients
                 harmonicsArray
                 thetaFreqs
                 rFreqs
                 a
                 b
                 input)
         xs $
       inputs) $
  xs

{-# INLINE convolveSingle' #-}
convolveSingle' ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> Double
  -> R.Array U DIM1 (Complex Double)
  -> VU.Vector (Complex Double)
convolveSingle' !field !coefficients !harmonicsArray !phiFreqs !rhoFreqs !x !y !input =
  let (Z :. cols :. rows) = extent (harmonicsArray IA.! (0, 0))
      rowCenter = div rows 2
      colCenter = div cols 2
      (Z :. _ :. _ :. numRhoFreq :. numPhiFreq) = extent coefficients
      -- harmonics =
      --   R.traverse2 phiFreqs rhoFreqs (\_ _ -> extent coefficients) $ \fPhiFreq fRhoFreq (Z :. _ :. theta :. rho :. phi) ->
      --     (harmonicsArray IA.!
      --      ( round (fRhoFreq (Z :. rho))
      --      , round (fPhiFreq (Z :. phi) - (fPhiFreq (Z :. theta))))) R.!
      --     (Z :. (x + colCenter) :. (y + rowCenter))
      harmonics =
        R.traverse2 phiFreqs rhoFreqs (\_ _ -> extent coefficients) $ \fPhiFreq fRhoFreq (Z :. _ :. theta :. rho :. phi) ->
          let !tf = fPhiFreq (Z :. phi) - fPhiFreq (Z :. theta)
              !rf = fRhoFreq (Z :. rho)
          in (x :+ y) ** (tf :+ 0) *
             ((x ^ 2 + y ^ 2) :+ 0) ** (((-tf - 1) :+ rf) / 2)
      product =
        R.traverse2 (R.zipWith (*) coefficients harmonics) input const $ \fCoefficient fInput idx@(Z :. _ :. _ :. _ :. phi) ->
          fCoefficient idx * fInput (Z :. phi)
      arr =
        case field of
          Source -> product
          Sink ->
            R.traverse2 product phiFreqs const $ \fProd fPhiFreq idx@(Z :. _ :. theta :. _ :. phi) ->
              fProd idx *
              (cis $ (-(fPhiFreq (Z :. phi) - fPhiFreq (Z :. theta))) * pi)
  in if x == 0 && y == 0
       then VU.replicate (numRhoFreq * numPhiFreq) 0
       else toUnboxed . sumS . sumS $ arr

convolveSparse' ::
     Field
  -> R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> [(Double, Double)]
  -> [R.Array U DIM1 (Complex Double)]
  -> [R.Array U DIM1 (Complex Double)]
convolveSparse' !field !coefficients !harmonicsArray !phiFreqs !rhoFreqs !cutoff !xs !inputs =
  L.map (fromUnboxed (extent . L.head $ inputs)) .
  parMap
    rdeepseq
    (\(x, y) ->
       L.foldl1' (VU.zipWith (+)) .
       L.zipWith
         (\(x', y') input ->
            let !x'' = x - x'
                !y'' = y - y'
                (!a, !b) =
                  if (x'' ^ 2 + y'' ^ 2) <= cutoff ^ 2
                    then (x'', y'')
                    else (0, 0)
            in convolveSingle'
                 field
                 coefficients
                 harmonicsArray
                 phiFreqs
                 rhoFreqs
                 a
                 b
                 input)
         xs $
       inputs) $
  xs

{-# INLINE gaussianFilter2D #-}
gaussianFilter2D :: DFTPlan -> Int -> Double -> IO (VS.Vector (Complex Double))
gaussianFilter2D plan numPoint stdG =
  dftExecute plan (DFTPlanID DFT1DG [numPoint, numPoint] [0, 1]) .
  VU.convert .
  toUnboxed . computeS . makeFilter2D . fromFunction (Z :. numPoint :. numPoint) $ \(Z :. i :. j) ->
    let x = fromIntegral $ i - div numPoint 2
        y = fromIntegral $ j - div numPoint 2
    in (exp (-(x ^ 2 + y ^ 2) / (2 * stdG ^ 2))) :+ 0

{-# INLINE convoluvePinhweelBasis #-}
convoluvePinhweelBasis ::
     R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
convoluvePinhweelBasis coefficients harmonicsArray arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) =
 let initVec = VS.replicate (VS.length . L.head $ vecs) 0
     idxs = (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)
  in DFTArray xFreq yFreq thetaFreqs rFreqs .
     parMap
      rdeepseq
       (\((r, rFreq), (theta, thetaFreq)) ->
          L.foldl'
            (\vec (((rho, rhoFreq), (phi, phiFreq)), inputVec) ->
               VS.zipWith (+) vec .
               VS.map (* (coefficients R.! (Z :. r :. theta :. rho :. phi))) .
               VS.zipWith
                 (*)
                 (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq) $
               inputVec)
            initVec .
          L.zip idxs $
          vecs) $
     idxs

{-# INLINE convoluvePinhweelBasis' #-}
convoluvePinhweelBasis' ::
     R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
convoluvePinhweelBasis' coefficients harmonicsArray arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) =
  let idxTheta = L.zip [0 ..] thetaFreqs
      initVec = VS.replicate (VS.length . L.head $ vecs) 0
  in DFTArray xFreq yFreq thetaFreqs rFreqs .
     parMap
       rdeepseq
       (\(!theta, !thetaFreq) ->
          L.foldl'
            (\vec1 (rho, rhoFreq) ->
               L.foldl'
                 (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                    VS.zipWith
                      (+)
                      vec2
                      (VS.map
                         (* (coefficients R.!
                             (Z :. (0 :: Int) :. theta :. rho :. phi))) .
                       VS.zipWith
                         (*)
                         (getHarmonics
                            harmonicsArray
                            phiFreq
                            rhoFreq
                            thetaFreq
                            0) $
                       inputVec))
                 vec1 .
               L.zip idxTheta $
               vecs)
            initVec .
          L.zip [0 ..] $
          rFreqs) $
     idxTheta

{-# INLINE convolveFull'' #-}
convolveFull'' ::
     R.Array U DIM4 (Complex Double)
  -> IA.Array (Int, Int) (VS.Vector (Complex Double))
  -> DFTArray
  -> DFTArray
convolveFull'' !coefficients !harmonicsArray !arr@(DFTArray r2Freq _ thetaFreqs rFreqs vecs) =
  let idxTheta = L.zip [0 ..] thetaFreqs
      !initVec = VS.replicate (VS.length . L.head $ vecs) 0
  in DFTArray r2Freq r2Freq thetaFreqs rFreqs .
     parMap
       rdeepseq
       (\(!theta, !thetaFreq) ->
          L.foldl'
            (\vec1 (rho, rhoFreq) ->
               L.foldl'
                 (\(!vec2) ((!phi, !phiFreq), inputVec) ->
                    VS.zipWith
                      (+)
                      vec2
                      (VS.map
                         (* (coefficients R.!
                             (Z :. (0 :: Int) :. theta :. rho :. phi)))
                         (getHarmonics
                            harmonicsArray
                            phiFreq
                            rhoFreq
                            thetaFreq
                            0)))
                 vec1 .
               L.zip idxTheta $
               vecs)
            initVec .
          L.zip [0 ..] $
          rFreqs) $
     idxTheta

{-# INLINE convoluvePinhweelBasisTest #-}
convoluvePinhweelBasisTest ::
     VS.Vector (Complex Double)
  -> [VS.Vector (Complex Double)]
  -> [VS.Vector (Complex Double)]
  -> DFTArray
  -> IO DFTArray
convoluvePinhweelBasisTest coefficients harmonicsArray1 harmonicsArray2 arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) = do
  let vecs1 = parZipWith rdeepseq (VS.zipWith (*)) vecs harmonicsArray1
      m = L.length thetaFreqs * L.length rFreqs
      n = xFreq * yFreq
  vec2 <- gemmBLAS m n m coefficients . VS.concat $ vecs1
  let arr =
        fromUnboxed
          (Z :. (L.length rFreqs) :. (L.length thetaFreqs) :. xFreq :. yFreq) .
        VS.convert $
        vec2
      vecs2 =
        parMap
          rdeepseq
          (\(i, j) ->
             VS.convert . toUnboxed . computeS . R.slice arr $
             (Z :. i :. j :. All :. All))
          [ (i, j)
          | i <- [0 .. L.length rFreqs - 1]
          , j <- [0 .. L.length thetaFreqs - 1]
          ] :: [VS.Vector (Complex Double)]
      vecs3 = parZipWith rdeepseq (VS.zipWith (*)) vecs2 harmonicsArray2
  return . DFTArray xFreq yFreq thetaFreqs rFreqs $ vecs3

{-# INLINE convoluvePinhweelBasisTest' #-}
convoluvePinhweelBasisTest' ::
     R.Array U DIM4 (Complex Double)
  -> [VS.Vector (Complex Double)]
  -> [VS.Vector (Complex Double)]
  -> DFTArray
  -> IO DFTArray
convoluvePinhweelBasisTest' coefficients harmonicsArray1 harmonicsArray2 arr@(DFTArray xFreq yFreq thetaFreqs rFreqs vecs) = do
  let (Z :. numRFreq :. numThetaFreq :. numRhoFreq :. numPhiFreq) =
        extent coefficients
      coefVec = VU.convert . toUnboxed $ coefficients
      vecs1 =
        parZipWith rdeepseq (VS.zipWith (*)) harmonicsArray1 .
        L.concat . L.replicate numRhoFreq $
        vecs
      m = numRFreq * numThetaFreq
      n = xFreq * yFreq
      k = numRhoFreq * numPhiFreq
  print (m, n, k)
  vec2 <- gemmBLAS m n k coefVec . VS.concat $ vecs1
  let arr =
        fromUnboxed (Z :. numRFreq :. numThetaFreq :. xFreq :. yFreq) .
        VS.convert $
        vec2
      vecs2 =
        parMap
          rdeepseq
          (\(i, j) ->
             VS.convert . toUnboxed . computeS . R.slice arr $
             (Z :. i :. j :. All :. All))
          [(i, j) | i <- [0 .. numRFreq - 1], j <- [0 .. numThetaFreq - 1]] :: [VS.Vector (Complex Double)]
      vecs3 = parZipWith rdeepseq (VS.zipWith (*)) vecs2 harmonicsArray2
  return . DFTArray xFreq yFreq thetaFreqs rFreqs $ vecs3


{-# LANGUAGE AllowAmbiguousTypes, DataKinds, RankNTypes, TypeFamilies,
             TypeOperators #-}
{-# OPTIONS_GHC -Wno-missing-export-lists #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | Shaped tensor-based implementation of Recurrent Neural Network
-- for classification of MNIST digits. Sports 2 hidden layers.
module HordeAd.Tool.MnistRnnShaped where

import Prelude

import           Control.Monad (foldM)
import qualified Data.Array.DynamicS as OT
import           Data.Array.Internal (valueOf)
import qualified Data.Array.Shape
import qualified Data.Array.Shaped as OSB
import qualified Data.Array.ShapedS as OS
import           Data.Proxy (Proxy)
import qualified Data.Vector.Generic as V
import           GHC.TypeLits (KnownNat)
import           Numeric.LinearAlgebra (Vector)
import qualified Numeric.LinearAlgebra as HM

-- until stylish-haskell accepts NoStarIsType
import qualified GHC.TypeLits

import HordeAd.Core.DualNumber
import HordeAd.Core.Engine
import HordeAd.Core.PairOfVectors (DualNumberVariables, varS)
import HordeAd.Tool.MnistData

zeroStateS
  :: (IsScalar d r, OS.Shape sh)
  => (DualNumber d (OS.Array sh r)  -- state
      -> a)
  -> a
zeroStateS f = f (konstS 0)

unrollLastS :: forall state c w m d r k rest.
               (DualMonad d r m, KnownNat k, OS.Shape rest)
            => (state -> OS.Array rest r -> w -> m (c, state))
            -> (state -> OS.Array (k : rest) r -> w -> m (c, state))
unrollLastS f s0 xs w = do
  let g :: (c, state) -> OS.Array rest r -> m (c, state)
      g (_, s) x = f s x w
  foldM g (undefined, s0) $ OSB.toList $ OS.unravel xs

type LayerWeigthsRNN in_width out_width d r =
  ( DualNumber d (OS.Array '[out_width, in_width] r)   -- input weight
  , DualNumber d (OS.Array '[out_width, out_width] r)  -- state weight
  , DualNumber d (OS.Array '[out_width] r) )           -- bias

rnnMnistLayerS
  :: forall in_width out_width batch_size d r m.
     ( DualMonad d r m
     , KnownNat in_width, KnownNat out_width, KnownNat batch_size )
  => DualNumber d (OS.Array '[out_width, batch_size] r)  -- in state
  -> DualNumber d (OS.Array '[in_width, batch_size] r)  -- in
  -> LayerWeigthsRNN in_width out_width d r
  -> m ( DualNumber d (OS.Array '[out_width, batch_size] r)  -- out
       , DualNumber d (OS.Array '[out_width, batch_size] r) )  -- out state
rnnMnistLayerS s x (wX, wS, b) = do
  let y = wX <>$ x + wS <>$ s + asColumnS b
  yTanh <- returnLet $ tanh y
  return (yTanh, yTanh)

rnnMnistTwoS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => DualNumber d (OS.Array '[2 GHC.TypeLits.* out_width, batch_size] r)
       -- initial state
  -> OS.Array '[SizeMnistWidth, batch_size] r
  -> ( LayerWeigthsRNN SizeMnistWidth out_width d r
     , LayerWeigthsRNN out_width out_width d r )
  -> m ( DualNumber d (OS.Array '[out_width, batch_size] r)
       , DualNumber d (OS.Array '[2 GHC.TypeLits.* out_width, batch_size] r) )
           -- final state
rnnMnistTwoS s x ((wX, wS, b), (wX2, wS2, b2)) = do
  let s1 = sliceS @0 @out_width s
      s2 = sliceS @out_width @out_width s
  (vec1, s1') <- rnnMnistLayerS s1 (constant x) (wX, wS, b)
  (vec2, s2') <- rnnMnistLayerS s2 vec1 (wX2, wS2, b2)
  s3 <- returnLet $ appendS s1' s2'
  return (vec2, s3)

rnnMnistZeroS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => OS.Array '[SizeMnistHeight, SizeMnistWidth, batch_size] r
  -- All below is the type of all paramters of this nn. The same is reflected
  -- in the length function below and read from variables further down.
  -> ( LayerWeigthsRNN SizeMnistWidth out_width d r
     , LayerWeigthsRNN out_width out_width d r )
  -> DualNumber d (OS.Array '[SizeMnistLabel, out_width] r)
  -> DualNumber d (OS.Array '[SizeMnistLabel] r)
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
rnnMnistZeroS xs ((wX, wS, b), (wX2, wS2, b2)) w3 b3 = do
  (out, _s) <- zeroStateS (unrollLastS rnnMnistTwoS) xs
                          ((wX, wS, b), (wX2, wS2, b2))
  returnLet $ w3 <>$ out + asColumnS b3

rnnMnistLenS
  :: forall out_width. KnownNat out_width
  => Proxy out_width -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
rnnMnistLenS _ =
  ( 0
  , []
  , []
  , [ Data.Array.Shape.shapeT @'[out_width, SizeMnistWidth]
    , Data.Array.Shape.shapeT @'[out_width, out_width]
    , Data.Array.Shape.shapeT @'[out_width]
    , Data.Array.Shape.shapeT @'[out_width, out_width]
    , Data.Array.Shape.shapeT @'[out_width, out_width]
    , Data.Array.Shape.shapeT @'[out_width]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel, out_width]
    , Data.Array.Shape.shapeT @'[SizeMnistLabel]
    ]
  )

rnnMnistS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => OS.Array '[SizeMnistHeight, SizeMnistWidth, batch_size] r
  -> DualNumberVariables d r
  -> m (DualNumber d (OS.Array '[SizeMnistLabel, batch_size] r))
rnnMnistS xs variables = do
  let wX = varS variables 0
      wS = varS variables 1
      b = varS variables 2
      wX2 = varS variables 3
      wS2 = varS variables 4
      b2 = varS variables 5
      w3 = varS variables 6
      b3 = varS variables 7
  rnnMnistZeroS @out_width xs ((wX, wS, b), (wX2, wS2, b2)) w3 b3

rnnMnistLossFusedS
  :: forall out_width batch_size d r m.
     (DualMonad d r m, KnownNat out_width, KnownNat batch_size)
  => Proxy out_width
  -> MnistDataBatchS batch_size r
  -> DualNumberVariables d r
  -> m (DualNumber d r)
rnnMnistLossFusedS _ (glyphS, labelS) variables = do
  let xs = OS.transpose @'[2, 1, 0] glyphS
  result <- rnnMnistS @out_width xs variables
  let targets2 = HM.tr $ HM.reshape (valueOf @SizeMnistLabel)
                       $ OS.toVector labelS
  vec <- lossSoftMaxCrossEntropyL targets2 (fromS2 result)
  returnLet $ scale (recip $ fromIntegral (valueOf @batch_size :: Int))
            $ sumElements0 vec

rnnMnistTestS
  :: forall out_width batch_size r.
     (IsScalar 'DModeGradient r, KnownNat out_width, KnownNat batch_size)
  => Proxy r -> Proxy out_width
  -> MnistDataBatchS batch_size r
  -> Domains r
  -> r
rnnMnistTestS _ _ (glyphS, labelS) parameters =
  let xs = OS.transpose @'[2, 1, 0] glyphS
      outputS = primalValue (rnnMnistS @out_width xs) parameters
      outputs = map OS.toVector $ OSB.toList $ OS.unravel
                $ OS.transpose @'[1, 0] $ outputS
      labels = map OS.toVector $ OSB.toList $ OS.unravel labelS
      matchesLabels :: Vector r -> Vector r -> Int
      matchesLabels output label | V.maxIndex output == V.maxIndex label = 1
                                 | otherwise = 0
  in fromIntegral (sum (zipWith matchesLabels outputs labels))
     / fromIntegral (valueOf @batch_size :: Int)


module RayLib where

import Numeric.LinearAlgebra
import qualified Graphics.Image as G
import qualified Debug.Trace as Debug

data Ray = Ray {rayOrigin :: Vector Double,
                rayDirection :: Vector Double}
    deriving (Show, Eq)

eps :: Double
eps = 0.000001

data Isect = Isect {iT :: Double,
                    iNormal :: Vector Double,
                    iObject :: Geometry,
                    iMaterial :: Material}

instance Eq Isect where
    (Isect iT1 iNormal1 _ _) == (Isect iT2 iNormal2 _ _) = (iT1 == iT2) && (iNormal1 == iNormal2)

instance Ord Isect where
    compare (Isect iT1 _ _ _) (Isect iT2 _ _ _) = compare iT1 iT2

instance Show Isect where
    show (Isect t n _ _) = "time: " ++ show t ++ ", normal: " ++ show n

class At a where
    at :: a -> Ray -> Vector Double

instance At Double where
    at t (Ray p d) = p + (fromList [t] :: Vector Double) * d

instance At Isect where
    at (Isect t _ _ _) (Ray p d) = p + (fromList [t] :: Vector Double) * d

type Color = Vector Double

data MatParam = MatParam Color
    deriving (Show, Read, Eq)

data Material = Material {ka :: MatParam,
                          ks :: MatParam,
                          kd :: MatParam,
                          ke :: MatParam,
                          kr :: MatParam,
                          kt :: MatParam,
                          shininess :: Double}
    deriving (Show, Read, Eq)

data Light =    PointLight       {lColor :: Color,
                                  lPosition :: Vector Double,
                                  constantTerm :: Double,
                                  linearTerm :: Double,
                                  quadraticTerm :: Double}
            |   DirectionalLight {lColor :: Color,
                                  lOrientation :: Vector Double}
    deriving (Show, Read, Eq)

type Lights = [Light]

lightColor :: Light -> Color
lightColor (PointLight c _ _ _ _) = c
lightColor (DirectionalLight c _) = c

getDirection :: Light -> Vector Double -> Vector Double
getDirection (PointLight _ pos _ _ _) p = normalize (pos - p)
getDirection (DirectionalLight _ o) _ = -o

distanceAttenuation :: Light -> Vector Double -> Double
distanceAttenuation (PointLight _ pos c b a) p = min 1.0 atten
    where atten = (c * dist * dist + b * dist + a) ** (-1.0)
          dist = norm_2 (p - pos)
distanceAttenuation (DirectionalLight _ _) p = 1.0

shadowAttenuation :: Scene -> Light -> Vector Double -> Double
shadowAttenuation scene (PointLight _ pos _ _ _) p =
    let shadowDir = normalize (pos - p) in
    let ray = Ray p shadowDir in
    let isect = intersect scene ray in
    case isect of
        Nothing -> 1
        Just (Isect t _ _ _) ->
            if (t < 0) || (t > norm_2 (pos - p)) then
                1
            else
                0
shadowAttenuation scene dl@(DirectionalLight _ _) p =
    let shadowDir = normalize (getDirection dl p) in
    let ray = Ray p shadowDir in
    let isect = intersect scene ray in
    case isect of
        Nothing -> 1
        Just (Isect t _ _ _) ->
            if t < 0 then
                1
            else
                0

data Geometry = Sphere {sphereCenter :: Vector Double,
                        sphereRadius :: Double,
                        sphereMaterial :: Material}
    deriving (Show, Read, Eq)

type Objects = [Geometry]

data Scene = Scene {scLights :: Lights,
                    scObjects :: Objects,
                    scAmbient :: Color}
    deriving (Show, Read, Eq)

data Camera = Camera {eye :: Vector Double,
                      viewDir :: Vector Double,
                      upDir :: Vector Double,
                      aspect :: Double,
                      fov :: Double}
    deriving (Show, Read, Eq)

rayThrough :: Camera -> Double -> Double -> Ray
rayThrough (Camera eye viewDir upDir aspect fov) xx yy =
    let z = -viewDir in
    let y = upDir in
    let x = cross y z in
    let m = fromRows [x,y,z] in
    let fov' = fov * pi / 180.0 in
    let normalizedHeight = 2 * tan (fov' / 2.0) in
    let u = (flatten (m <> (fromColumns [(fromList [1,0,0])]))) * (fromList [normalizedHeight * aspect]) in
    let v = (flatten (m <> (fromColumns [(fromList [0,1,0])]))) * (fromList [normalizedHeight]) in
    let look = (flatten (m <> (fromColumns [(fromList [0,0,-1])]))) in
    let x' = fromList [xx - 0.5] in
    let y' = fromList [yy - 0.5] in
    let dir = normalize (look + x' * u + y' * v) in
    Ray eye dir

intersectObj :: Ray -> Geometry -> Maybe Isect
intersectObj r@(Ray rayPos rayDir) s@(Sphere center radius material) =
    let v = center - rayPos in
    let b = v <.> rayDir in
    let discriminant = b * b - v <.> v + radius * radius in
    if discriminant < 0.0 then
        Nothing
    else
        let det = sqrt discriminant in
        let (t1,t2) = (b - det, b + det) in
        if t2 < 0.0 then
            Nothing
        else if t1 > 0.0 then
            pure (Isect t1 (normalize (at t1 r)) s material)
        else
            pure (Isect t2 (normalize (at t2 r)) s material)

intersect :: Scene -> Ray -> Maybe Isect
intersect scene@(Scene _ objects _) r =
    let isects = filter (/= Nothing) . map (intersectObj r) $ objects in
    if length isects == 0 then
        Nothing
    else
        minimum isects

perLightDiffuse :: Light -> Vector Double -> Vector Double -> Vector Double -> Double
perLightDiffuse light isectPoint rayDir normal =
    let lightDir = getDirection light isectPoint in
    let ddot = lightDir <.> (if (rayDir <.> normal > 0) then -normal else normal) in
    let dmax = max ddot 0.0 in
    dmax

perLightSpecular :: Light -> Vector Double -> Vector Double -> Vector Double -> Double
perLightSpecular light isectPoint rayDir normal =
    let lightDir = getDirection light isectPoint in
    let reflect = lightDir - (fromList [2 * (lightDir <.> normal)]) * normal in
    if norm_2 reflect == 0 then
        0
    else
        let reflectNorm = normalize reflect in
        let sdot = rayDir <.> reflect in
        let smax = max sdot 0.0 in
        smax

combineLightTerms :: Scene -> Vector Double -> Vector Double -> Vector Double -> Double -> Light -> Color
combineLightTerms scene isectPoint rayDir normal alpha light =
    let diffuse = perLightDiffuse light isectPoint rayDir normal in
    let specular = perLightSpecular light isectPoint rayDir normal in
    let color = lightColor light in
    let atten = (distanceAttenuation light isectPoint) * (shadowAttenuation scene light isectPoint) in
    color * (fromList [(diffuse + specular) * atten])

shade :: Scene -> Ray -> Isect -> Color
shade scene@(Scene lights objects ambient) r@(Ray p d) i@(Isect t normal _ material) =
    let (Material kaM ksM kdM keM krM ktM shiny) = material in
    let (MatParam ke) = keM in
    let emissive = ke in
    let (MatParam ka) = kaM in
    let amb = ka * ambient in
    let isectPoint = at t r in
    let alpha = shiny in
    emissive + amb + (sum . map (combineLightTerms scene isectPoint d normal alpha) $ lights)

data Screen = Screen {screenWidth :: Int,
                      screenHeight :: Int}
    deriving (Show, Read, Eq)

clamp :: Double -> Double -> Vector Double -> Vector Double
clamp a b = fromList . map (\x -> max a (min b x)) . toList

trace :: Scene -> Camera -> Double -> Double -> Color
trace scene camera x y =
    let ray = rayThrough camera x y in
    let color = traceRay scene ray 0 0 in
    clamp 0 1 color

tracePixel :: Scene -> Screen -> Camera -> Int -> Int -> Color
tracePixel scene (Screen width height) camera i j =
    let x = (fromIntegral i :: Double) / (fromIntegral width :: Double) in
    let y = (fromIntegral j :: Double) / (fromIntegral height :: Double) in
    let color = trace scene camera x y in
    (fromList [255.0]) * color

reflection :: Scene -> Ray -> Isect -> Int -> Color
reflection scene ray@(Ray p d) isect@(Isect t n _ mat) depth =
    let reflectR = d - (fromList [2 * d <.> n]) * n in
    if reflectR == 0 then
        fromList ([0,0,0])
    else
        let isectPointApprox = at (t - eps) ray in
        let reflRay = Ray isectPointApprox (normalize reflectR) in
        let (Material _ _ _ _ (MatParam kr)  _ _) = mat in
        kr * (traceRay scene reflRay (depth - 1) 0)

traceRay :: Scene -> Ray -> Int -> Double -> Color
traceRay scene ray@(Ray p d) depth t =
    let mIsect = intersect scene ray in
    case mIsect of
        Nothing -> fromList[0,0,0]
        (Just isect) -> let colorC = shade scene ray isect in
                        if depth == 0 then
                            colorC
                        else
                            let colorRefl = reflection scene ray isect depth in
                            colorC + colorRefl

type Image = ([[Color]], Int, Int)

traceImage :: Scene -> Screen -> Camera -> Image
traceImage scene screen@(Screen width height) camera =
    let (w, h) = (width - 1, height - 1) in
    let img = [[tracePixel scene screen camera i j | j <- [0..h]] | i <- [0..w]] in
    (img, width, height)

rgbify :: [Double] -> G.Pixel G.RGB Double
rgbify vec = G.PixelRGB (r / 255.0) (g / 255.0) (b / 255.0)
    where r = vec !! 0
          g = vec !! 1
          b = vec !! 2

render :: Image -> G.Image G.VU G.RGB Double
render (image, width, height) = G.fromLists . format $ image
    where format = map (map rgbify) . map (map toList)


{-# LANGUAGE FlexibleContexts #-}

------------------------------------------------
-- |
-- Module    :  Numeric.Matrix.Integral.HNFLLL
-- Copyright :  (c) Jun Yoshida 2019
-- License   :  BSD3
--
-- Compute Hermite normal forms in the LLL-based method.
-- The original "pseudo-code" is found in the paper
--   George Havas, Bohdan S. Majewski & Keith R. Matthews (1998) Extended GCD and Hermite Normal Form Algorithms via Lattice Basis Reduction, Experimental Mathematics, 7:2, 125-136, DOI: 10.1080/10586458.1998.10504362
--
------------------------------------------------

module Numeric.Matrix.Integral.HNFLLL (
  hnfLLLST,
  hnfLLL,
  ) where

import Control.Monad (when, forM_)
import Control.Monad.ST
import Control.Monad.Loops (whileM_)

import Data.STRef

import qualified Data.Maybe as M

import Data.Int

import Numeric.LinearAlgebra ((!))
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Devel

-- | Type class which guarantees the type can be embedded in the real number field
class (Num t,Eq t,Ord t) => Realizable t where
  toReal :: t -> Double

instance Realizable Double where
  toReal = id

instance Realizable Int64 where
  toReal = fromIntegral

-- | Type class to qualify a type to be used in LLL-based algorithm
class (Realizable t, LA.Indexable (LA.Vector t) t, LA.Container LA.Vector t) => LLLQualified t

instance LLLQualified Double

instance LLLQualified Int64

-- | Data Type to carry resources in the computation of LLL-based algorithm
data HNFLLLResource s t = HNFLLL {
  matA :: STMatrix s t,
  matU :: STMatrix s t,
  lambda :: STMatrix s Double,
  focusC1 :: STRef s (Maybe Int),
  focusC2 :: STRef s (Maybe Int) }

-- | The function which is called Reduce2 in the original "pseudo-code."
reduceLLL :: (LLLQualified t) => Int -> Int -> HNFLLLResource s t -> ST s ()
reduceLLL k i resRef = do
  ma <- freezeMatrix (matA resRef)
  case M.listToMaybe (LA.find (/=0) (ma!i)) of
    (Just j) -> do
      writeSTRef (focusC1 resRef) (Just j)
      when ( (ma!i!j) < 0 ) $ do
        rowOper (SCAL (-1) (Row i) AllCols) (matA resRef)
        rowOper (SCAL (-1) (Row i) AllCols) (matU resRef)
        rowOper (SCAL (-1) (Row (j+1)) (ColRange 0 j)) (lambda resRef)
        rowOper (SCAL (-1) (FromRow (j+2)) (Col (j+1))) (lambda resRef)
    Nothing -> writeSTRef (focusC1 resRef) Nothing
  case M.listToMaybe (LA.find (/=0) (ma!k)) of
    (Just j) -> writeSTRef (focusC2 resRef) (Just j)
    Nothing -> writeSTRef (focusC2 resRef) Nothing
  c1 <- readSTRef (focusC1 resRef)
  q <- fromIntegral.floor <$> do
    case c1 of
      (Just j) ->
        (/) <$> (toReal <$> readMatrix (matA resRef) k j) <*> (toReal <$> readMatrix (matA resRef) i j)
      Nothing -> do
        lki <- readMatrix (lambda resRef) (k+1) (i+1)
        di <- readMatrix (lambda resRef) (i+1) (i+1)
        if 2 * abs lki > di
          then return $ fromIntegral (round (lki / di))
          else return 0
  when (q /= 0) $ do
    rowOper (AXPY (negate q) i k AllCols) (matA resRef)
    rowOper (AXPY (negate q) i k AllCols) (matU resRef)
    rowOper (AXPY (toReal $ negate q) (i+1) (k+1) AllCols) (lambda resRef)

-- | The function which is called Swap1 in the original "pseudo-code."
swapLLL :: (Num t, LA.Element t) => Int -> HNFLLLResource s t -> ST s ()
swapLLL k resRef = do
  rowOper (SWAP k (k-1) AllCols) (matA resRef)
  rowOper (SWAP k (k-1) AllCols) (matU resRef)
  rowOper (SWAP k (k+1) (ColRange 0 (k-1))) (lambda resRef)
  ml <- freezeMatrix (lambda resRef)
  forM_ [(k+2)..(LA.cols ml-1)] $ \i -> do
    writeMatrix (lambda resRef) i k $
      ((ml!i!k)*(ml!(k+1)!k) + (ml!i!(k+1))*(ml!(k-1)!(k-1)))/(ml!k!k)
    writeMatrix (lambda resRef) i (k+1) $
      ((ml!i!k)*(ml!(k+1)!(k+1)) + (ml!i!(k+1))*(ml!(k+1)!k))/(ml!k!k)
  writeMatrix (lambda resRef) k k $
    ((ml!(k+1)!k)*(ml!(k+1)!k) + (ml!(k+1)!(k+1))*(ml!(k-1)!(k-1)))/(ml!k!k)

-- | Recognize whether a swap should occur or not.
shouldSwapST :: (Num t, LA.Element t) => Int -> HNFLLLResource s t -> ST s Bool
shouldSwapST k resRef = do
  ml <- freezeMatrix $ lambda resRef
  mc1 <- readSTRef $ focusC1 resRef
  mc2 <- readSTRef $ focusC2 resRef
  case (mc1,mc2) of
    (Just c1,Just c2) -> return $ c1 <= c2
    (Just _,Nothing) -> return True
    (Nothing,Just _) -> return False
    (Nothing,Nothing) ->
      return $ (ml!(k+1)!k)*(ml!(k+1)!k) + (ml!(k+1)!(k+1))*(ml!(k-1)!(k-1)) < (ml!k!k)*(ml!k!k)

-- | ST to compute the Hermite normal form for the matrix in stMatA.
-- The first argument is also subject to the operation that is executed on stMatA, so one gets the transformation unimodular matrix by passing the identity.
-- See the source of hnfLLL function for an example of usage.
hnfLLLST :: (LLLQualified t) => STMatrix s t -> STMatrix s t -> ST s ()
hnfLLLST stMatU stMatA = do
  m <- LA.rows <$> freezeMatrix stMatA
  newLambda <- thawMatrix $ LA.ident (m+1)
  newC1 <- newSTRef Nothing
  newC2 <- newSTRef Nothing
  let resRef = HNFLLL stMatA stMatU newLambda newC1 newC2
  kRef <- newSTRef (1 :: Int)
  whileM_ ((<m) <$> readSTRef kRef) $ do
    k <- readSTRef kRef
    reduceLLL k (k-1) resRef
    shouldSwap <- shouldSwapST k resRef
    if shouldSwap
      then swapLLL k resRef >> when (k>1) (writeSTRef kRef (k-1))
      else (forM_ [0..k-2] $ \i -> reduceLLL k i resRef) >> writeSTRef kRef (k+1)
  setMatrix stMatA 0 0 =<< (LA.flipud <$> freezeMatrix stMatA)
  setMatrix stMatU 0 0 =<< (LA.flipud <$> freezeMatrix stMatU)

-- | Execute the LLL-based algorithm on a given matrix.
hnfLLL :: LLLQualified t => LA.Matrix t -> (LA.Matrix t, LA.Matrix t)
hnfLLL mtx = runST $ do
  stMatA <- thawMatrix mtx
  stMatU <- thawMatrix $ LA.ident (LA.rows mtx)
  hnfLLLST stMatU stMatA
  h <- freezeMatrix stMatA
  u <- freezeMatrix stMatU
  return (u,h)


-- |
-- Module    : Numeric.SpecFunctions.Extra
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Less common mathematical functions.
module Numeric.SpecFunctions.Extra (
    bd0
  , chooseExact
  , logChooseFast
  , logGammaAS245
  , logGammaCorrection
  ) where

import Numeric.MathFunctions.Constants (m_NaN,m_pos_inf)
import Numeric.SpecFunctions.Internal  (chooseExact,logChooseFast,logGammaCorrection)

-- | Evaluate the deviance term @x log(x/np) + np - x@.
bd0 :: Double                   -- ^ @x@
    -> Double                   -- ^ @np@
    -> Double
bd0 x np
  | isInfinite x || isInfinite np || np == 0 = m_NaN
  | abs x_np >= 0.1*(x+np)                   = x * log (x/np) - x_np
  | otherwise                                = loop 1 (ej0*vv) s0
  where
    x_np = x - np
    v    = x_np / (x+np)
    s0   = x_np * v
    ej0  = 2*x*v
    vv   = v*v
    loop j ej s = case s + ej/(2*j+1) of
                    s' | s' == s   -> s'  -- FIXME: Comparing Doubles for equality!
                       | otherwise -> loop (j+1) (ej*vv) s'

-- | Compute the logarithm of the gamma function Γ(/x/).  Uses
-- Algorithm AS 245 by Macleod.
--
-- Gives an accuracy of 10-12 significant decimal digits, except
-- for small regions around /x/ = 1 and /x/ = 2, where the function
-- goes to zero.  For greater accuracy, use 'logGammaL'.
--
-- Returns ∞ if the input is outside of the range (0 < /x/ ≤ 1e305).
logGammaAS245 :: Double -> Double
-- Adapted from http://people.sc.fsu.edu/~burkardt/f_src/asa245/asa245.html
logGammaAS245 x
    | x <= 0    = m_pos_inf
    -- Handle positive infinity. logGamma overflows before 1e308 so
    -- it's safe
    | x > 1e308 = m_pos_inf
    -- Normal cases
    | x < 1.5   = a + c *
                  ((((r1_4 * b + r1_3) * b + r1_2) * b + r1_1) * b + r1_0) /
                  ((((b + r1_8) * b + r1_7) * b + r1_6) * b + r1_5)
    | x < 4     = (x - 2) *
                  ((((r2_4 * x + r2_3) * x + r2_2) * x + r2_1) * x + r2_0) /
                  ((((x + r2_8) * x + r2_7) * x + r2_6) * x + r2_5)
    | x < 12    = ((((r3_4 * x + r3_3) * x + r3_2) * x + r3_1) * x + r3_0) /
                  ((((x + r3_8) * x + r3_7) * x + r3_6) * x + r3_5)
    | x > 3e6   = k
    | otherwise = k + x1 *
                  ((r4_2 * x2 + r4_1) * x2 + r4_0) /
                  ((x2 + r4_4) * x2 + r4_3)
  where
    (a , b , c)
        | x < 0.5   = (-y , x + 1 , x)
        | otherwise = (0  , x     , x - 1)

    y      = log x
    k      = x * (y-1) - 0.5 * y + alr2pi
    alr2pi = 0.918938533204673

    x1 = 1 / x
    x2 = x1 * x1

    r1_0 =  -2.66685511495;   r1_1 =  -24.4387534237;    r1_2 = -21.9698958928
    r1_3 =  11.1667541262;    r1_4 =    3.13060547623;   r1_5 =   0.607771387771
    r1_6 =  11.9400905721;    r1_7 =   31.4690115749;    r1_8 =  15.2346874070

    r2_0 = -78.3359299449;    r2_1 = -142.046296688;     r2_2 = 137.519416416
    r2_3 =  78.6994924154;    r2_4 =    4.16438922228;   r2_5 =  47.0668766060
    r2_6 = 313.399215894;     r2_7 =  263.505074721;     r2_8 =  43.3400022514

    r3_0 =  -2.12159572323e5; r3_1 =    2.30661510616e5; r3_2 =   2.74647644705e4
    r3_3 =  -4.02621119975e4; r3_4 =   -2.29660729780e3; r3_5 =  -1.16328495004e5
    r3_6 =  -1.46025937511e5; r3_7 =   -2.42357409629e4; r3_8 =  -5.70691009324e2

    r4_0 = 0.279195317918525;  r4_1 = 0.4917317610505968;
    r4_2 = 0.0692910599291889; r4_3 = 3.350343815022304
    r4_4 = 6.012459259764103


{-# LANGUAGE ConstraintKinds, FlexibleContexts, FlexibleInstances, GADTs   #-}
{-# LANGUAGE KindSignatures, MultiParamTypeClasses                         #-}
{-# LANGUAGE NoMonomorphismRestriction, TypeFamilies, TypeSynonymInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Algebra.Matrix (Matrix(..), delta, companion,
                       gaussReduction, maxNorm, rankWith, det,
                       inverse, inverseWith) where
import           Algebra.Internal
import qualified Algebra.LinkedMatrix as LM
import           Algebra.Prelude.Core hiding (maxNorm, zero)

import           Control.Lens                (both, view, (%~), (&), _3)
import           Control.Monad               (when)
import qualified Data.Matrix                 as DM
import qualified Data.Vector                 as V
import           GHC.Exts                    (Constraint)
import           Numeric.Algebra             (Additive, DecidableZero)
import           Numeric.Algebra             (Monoidal, Multiplicative)
import           Numeric.Algebra             (Unital)
import qualified Numeric.Algebra             as NA
import qualified Numeric.Decidable.Zero      as NA
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as LA
import qualified Prelude                     as P

{-# ANN module "HLint: ignore Redundant bang pattern" #-}

class Matrix mat where
  type Elem mat a :: Constraint
  cmap  :: (Elem mat a, Elem mat b) => (a -> b) -> mat a -> mat b
  empty :: Elem mat b => mat b
  fromLists :: Elem mat a => [[a]] -> mat a
  fromCols  :: Elem mat a => [V.Vector a] -> mat a
  fromCols [] = zero 0 0
  fromCols xs = foldr1 (<||>) $ map colVector xs
  fromRows :: Elem mat a => [V.Vector a] -> mat a
  fromRows [] = zero 0 0
  fromRows xs = foldr1 (<-->) $ map rowVector xs
  toCols :: Elem mat a => mat a -> [V.Vector a]
  toCols m = map (`getCol` m) [1..ncols m]
  toRows :: Elem mat a => mat a -> [V.Vector a]
  toRows m = map (`getRow` m) [1..nrows m]
  ncols :: mat a -> Int
  nrows :: mat a -> Int
  identity :: Elem mat a => Int -> mat a
  diag     :: Elem mat a => V.Vector a -> mat a
  getDiag :: Elem mat a => mat a -> V.Vector a
  trace :: Elem mat a => mat a -> a
  diagProd :: Elem mat a => mat a -> a
  zero :: Elem mat a => Int -> Int -> mat a
  colVector :: Elem mat a => V.Vector a -> mat a
  rowVector :: Elem mat a => V.Vector a -> mat a
  getCol :: Elem mat a => Int -> mat a -> V.Vector a
  getRow :: Elem mat a => Int -> mat a -> V.Vector a
  switchRows :: Elem mat a => Int -> Int -> mat a -> mat a
  scaleRow :: Elem mat a => a -> Int -> mat a -> mat a
  combineRows :: Elem mat a => Int -> a -> Int -> mat a -> mat a
  trans :: Elem mat a => mat a -> mat a
  buildMatrix :: Elem mat a => Int -> Int -> ((Int, Int) -> a) -> mat a
  index :: Elem mat a => Int -> Int -> mat a -> Maybe a
  index i j m = if 1 <= i && i <= nrows m && 1 <= j && j <= ncols m
                then Just $ m ! (i, j)
                else Nothing
  (!) :: Elem mat a => mat a -> (Int, Int) -> a
  (<||>) :: Elem mat a => mat a -> mat a -> mat a
  (<-->) :: Elem mat a => mat a -> mat a -> mat a
  nonZeroRows :: (DecidableZero a, Elem mat a) => mat a -> [Int]
  nonZeroRows = map fst . filter (V.any (not . NA.isZero) . snd) . zip [1..] . toRows
  nonZeroCols :: (DecidableZero a, Elem mat a) => mat a -> [Int]
  nonZeroCols = map fst . filter (V.any (not . NA.isZero) . snd) . zip [1..] . toCols

instance Matrix DM.Matrix where
  type Elem DM.Matrix a = P.Num a
  empty = DM.zero 0 0
  cmap  = fmap
  fromLists = DM.fromLists
  ncols = DM.ncols
  nrows = DM.nrows
  trans = DM.transpose
  identity = DM.identity
  diag v = DM.matrix (V.length v) (V.length v) $ \(i, j) ->
    if i == j then v V.! (i-1) else 0
  getDiag = DM.getDiag
  trace = DM.trace
  diagProd = DM.diagProd
  zero = DM.zero
  colVector = DM.colVector
  rowVector = DM.rowVector
  getCol = DM.getCol
  getRow = DM.getRow
  switchRows = DM.switchRows
  combineRows = DM.combineRows
  scaleRow = DM.scaleRow
  buildMatrix = DM.matrix
  (!) = (DM.!)
  (<||>) = (DM.<|>)
  (<-->) = (DM.<->)

swapIJ :: Eq a => a -> a -> a -> a
swapIJ i j k
  | k == i = j
  | k == j = i
  | otherwise = k

instance Matrix LA.Matrix where
  type Elem LA.Matrix a = (P.Num a, LA.Numeric  a, LA.Element a, LA.Container LA.Vector a)
  empty = LA.fromLists [[]]
  fromLists = LA.fromLists
  cmap = LA.cmap
  ncols = LA.cols
  nrows = LA.rows
  trans = LA.tr
  identity = LA.ident
  fromCols = LA.fromColumns . map (LA.fromList . V.toList)
  diag = LA.diag . LA.fromList . V.toList
  getDiag = V.fromList . LA.toList . LA.takeDiag
  trace = LA.sumElements . LA.takeDiag
  diagProd = LA.prodElements . LA.takeDiag
  zero i j = LA.konst 0 (i, j)
  colVector = LA.asColumn . LA.fromList . V.toList
  rowVector = LA.asRow . LA.fromList . V.toList
  toCols = map (V.fromList . LA.toList) . LA.toColumns
  toRows = map (V.fromList . LA.toList) . LA.toRows
  getCol i = V.fromList . LA.toList . (!! (i - 1)) . LA.toColumns
  getRow i = V.fromList . LA.toList . (!! (i - 1)) . LA.toRows
  switchRows i j m = m LA.? map (swapIJ (i-1) (j-1)) [0.. nrows m - 1]
  combineRows j s i m = LA.mapMatrixWithIndex (\(k,l) v -> if k == j - 1 then s P.* (m ! (i,l+1)) P.+ v else v) m
  buildMatrix w h f = LA.build (w, h) (\i j -> f (toIntLA i+1, toIntLA j+1))
  scaleRow a i = (fst .) $ LA.mutable $ \(k, l) m -> do
    v <- LA.readMatrix m k l
    when (k == i - 1) $
      LA.writeMatrix m k l (a P.* v)
  m ! (i, j) = m `LA.atIndex` (i - 1, j - 1)
  m <||> n = LA.fromColumns $ LA.toColumns m ++ LA.toColumns n
  m <--> n = LA.fromRows $ LA.toRows m ++ LA.toRows n

toIntLA :: LA.Container LA.Matrix e => e -> Int
toIntLA e = fromIntegral $ LA.toZ ((1 LA.>< 1) [e]) `LA.atIndex` (0, 0)

instance Matrix LM.Matrix where
  type Elem LM.Matrix a = (CoeffRing a, Unital a, Monoidal a, Multiplicative a, Additive a)
  cmap = LM.cmap
  (!) m pos = m LM.! (pos & both %~ pred)
  index i j = LM.index (i-1) (j-1)
  empty = LM.empty
  buildMatrix h w f = LM.fromList [((i, j), f (i,j)) | j <- [1..w], i <- [1..h]]
  trans = LM.transpose
  combineRows j s i = LM.combineRows s (i-1) (j-1)
  switchRows i j = LM.switchRows (i-1) (j-1)
  scaleRow a i = LM.scaleRow a (pred i)
  zero = LM.zeroMat
  identity = LM.identity
  diag = LM.diag
  getDiag = LM.getDiag
  diagProd = LM.diagProd
  trace = LM.trace
  getRow = LM.getRow . pred
  getCol = LM.getCol . pred
  nrows = LM.nrows
  ncols = LM.ncols
  fromLists = LM.fromLists
  (<||>) = (LM.<||>)
  (<-->) = (LM.<-->)
  rowVector = LM.rowVector
  colVector = LM.colVector
  nonZeroRows = LM.nonZeroRows
  nonZeroCols = LM.nonZeroCols

delta :: (NA.Monoidal r, NA.Unital r) => Int -> Int -> r
delta i j | i == j = NA.one
          | otherwise = NA.zero

companion :: (KnownNat n, CoeffRing r, Matrix mat,
              Elem mat r, IsMonomialOrder n ord)
          => Ordinal n -> OrderedPolynomial r ord n -> mat r
companion odn poly =
  let deg = totalDegree' poly
      vx  = leadingMonomial (var odn `asTypeOf` poly)
  in buildMatrix deg deg $ \(j, k) ->
  if 1 <= k && k <= deg - 1
  then delta j (k+1)
  else NA.negate $ coeff (NA.pow vx (fromIntegral $ j-1 :: NA.Natural)) poly

-- | @gaussReduction a = (a', p)@ where @a'@ is row echelon form and @p@ is pivoting matrix.
gaussReduction :: (Matrix mat, Elem mat a, Normed a, Eq a, NA.Field a)
               => mat a -> (mat a, mat a)
gaussReduction mat =
  let (a, b, _) = gaussReduction' mat in (a, b)

-- | @gaussReduction a = (a', p)@ where @a'@ is row echelon form and @p@ is pivoting matrix.
gaussReduction' :: (Matrix mat, Elem mat a, Normed a, Eq a, NA.Field a)
               => mat a -> (mat a, mat a, a)
gaussReduction' mat = {-# SCC "gaussRed" #-} go 1 1 mat (identity $ nrows mat) NA.one
  where
    go i j a p acc
      | i > nrows mat || j > ncols mat = (a, p, acc)
      | otherwise =
        let (k, new) = maximumBy (comparing $ norm . snd) [(l, a ! (l, j)) | l <- [i..nrows mat]]
        in if new == NA.zero
           then go i (j + 1) a p NA.zero
           else let prc l a0 p0
                      | l == i = prc (l+1) a0 p0
                      | l > nrows mat = (a0, p0)
                      | otherwise     =
                        let coe = NA.negate (a0 ! (l, j))
                            a'' = combineRows l coe i a0
                            p'' = combineRows l coe i p0
                        in prc (l+1) a'' p''
                    (a', p') = prc 1 (scaleRow (NA.recip new) i $ switchRows i k a)
                                     (scaleRow (NA.recip new) i $ switchRows i k p)
                    offset = if i == k then id else NA.negate
                in go (i+1) (j+1) a' p' (offset $ acc NA.* new)

det :: (Elem mat a, Eq a, NA.Field a, Normed a, Matrix mat)
    => mat a -> a
det = view _3 . gaussReduction'

maxNorm :: (Elem mat a, Normed a, Matrix mat) => mat a -> Norm a
maxNorm = maximum . concat . map (map norm . V.toList) . toRows

rankWith :: (Elem mat r, CoeffRing r, Matrix mat)
         => (mat r -> mat r) -> mat r -> Int
rankWith gauss = length . nonZeroRows . gauss

inverse :: (Elem mat a, Eq a, NA.Field a, Normed a, Matrix mat)
        => mat a -> mat a
inverse = snd . gaussReduction

inverseWith :: (mat a -> (mat a, mat a)) -> mat a -> mat a
inverseWith = (snd .)


{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
module VAD(
      vad
    ) where

--  let framed = frameWithWinAndOverlap 256 hann 10 s
--      sp = mapS (fft . U.map (:+ 0)) framed
--

import Prelude hiding(splitAt,(:),foldr1,tail,(++))
import Internal
import Signal
import qualified Data.Vector.Unboxed as U
import Data.Vector.Unboxed((!),Unbox(..))
import Data.List.Stream
import Common
import Windows
import Data.Complex
import Transform
import qualified Trace as F
import Fixed
import GHC.TypeLits
import SpecialInt
import Data.Int

ltseF :: Sample a
      => Int
      -> Signal (U.Vector a)
      -> Signal (U.Vector a)
ltseF n (Signal s) =
    let (before,remaining) = splitAt n s
        _lste b (h:t) =
            let (future,tl) = splitAt n t
                valueMax = U.zipWith max (foldr1 (U.zipWith max) (h:future))  (foldl1' (U.zipWith max) b)
            in
            valueMax:(_lste (tail before ++ [h]) t)
    in
    Signal (_lste before remaining)

lstdD :: Int
      -> U.Vector Double
      -> U.Vector Double
      -> Double
lstdD winSize noiseEnergy lste =
        let d a b | b == 0 = fromDouble 1e12
                  | otherwise = a / b
        in
        10 * log(U.sum (U.zipWith (d) lste noiseEnergy) / fromIntegral winSize) / log 10

theTotal :: (SingI r, SingI s, SingI n)
         => Fixed Int32 n s r
         -> U.Vector (Fixed Int32 n s r)
         -> U.Vector (Fixed Int32 n s r)
         -> Fixed Int40 n s r
theTotal m lste noiseEnergy =
           let  d a b | b == 0 = convert m
                      | otherwise = convert a / convert b
           in
           U.sum (U.zipWith (d) lste noiseEnergy)
{-# INLINE [0] theTotal #-}

lstdF :: (SingI n, SingI s, SingI r)
      => Int
      -> U.Vector (Fixed Int32 n s r)
      -> U.Vector (Fixed Int32 n s r)
      -> (Fixed Int32 n s r)
lstdF winSize noiseEnergy lste = 10 * log(convert $ theTotal maxBound lste noiseEnergy/ fromIntegral winSize) / log 10

getDecisionF :: (SingI n, SingI s, SingI r, SingI (n + n))
             => Int
             -> U.Vector (Fixed Int32 (n + n) s r)
             -> [(U.Vector (Fixed Int32 (n + n) s r), U.Vector (Fixed Int32 (n + n) s r) )]
             -> [Fixed Int16 n s r]
getDecisionF winSize energy ((c,currentE):r) =
    let tv = fromDouble 31.0
        tn = fromDouble 27.0
        l = lstdF winSize energy c
        result | l >= tv = 1 : getDecisionF winSize energy r
               | l <= tn = 0 : getDecisionF winSize currentE r
               | otherwise = 0 : getDecisionF winSize energy r
    in
    result
getDecisionF winSize energy [] = 0:getDecisionF winSize energy []

getDecisionD :: Int
             -> U.Vector Double
             -> [(U.Vector Double, U.Vector Double)]
             -> [Double]
getDecisionD winSize energy ((c,currentE):r) =
    let tv = fromDouble 31.0
        tn = fromDouble 27.0
        l = lstdD winSize energy c
        result | l >= tv = 1 : getDecisionD winSize energy r
               | l <= tn = 0 : getDecisionD winSize currentE r
               | otherwise = 0 : getDecisionD winSize energy r
    in
    result
getDecisionD winSize energy [] = 0:getDecisionD winSize energy []

bandEnergy :: Complex Double -> Double
bandEnergy (x :+ y) = x*x + y*y

bandEnergyF :: (SingI n, SingI s, SingI r, SingI (n + n)) => Complex (Fixed Int16 n s r) -> Fixed Int32 (n + n) s r
bandEnergyF (x :+ y) = amul x x + amul y y

class VAD a where
    vad :: (Sample a, FFT a)
        => Sampled Time a
        -> Sampled Time a

instance (SingI n, SingI s, SingI r, SingI (n + n)) => VAD (Fixed Int16 n s r) where
    vad s =
        let winSize = 256
            overlap = 20
            n = 2
            framed = frameWithWinAndOverlap winSize overlap hann s
            energy = mapS (U.map bandEnergyF . fft . U.map (:+ 0)) (getSignal framed)
            noiseEnergy0 = U.generate winSize (const (fromDouble 0.00001))
            lt = ltseF n energy
        in
        --mapS (lstd winSize noiseEnergy0) lt
        Sampled (period framed) (onSamples (getDecisionF winSize noiseEnergy0) (zipS lt (dropS n energy)))

instance VAD Double where
    vad s =
        let winSize = 256
            overlap = 20
            n = 2
            framed = frameWithWinAndOverlap winSize overlap hann s
            energy = mapS (U.map bandEnergy . fft . U.map (:+ 0)) (getSignal framed)
            noiseEnergy0 = U.generate winSize (const (fromDouble 0.00001))
            lt = ltseF n energy
        in
        --mapS (lstd winSize noiseEnergy0) lt
        Sampled (period framed) (onSamples (getDecisionD winSize noiseEnergy0) (zipS lt (dropS n energy)))


{-|
Module : GradientDescent
|-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}

module GradientDescent
  ( sigmoid
  , sigmoidPrime
  , matrixSigmoid
  , matrixSigmoidPrime
  , getOutcomes
  , gradientDescent
  ) where

import Numeric.LinearAlgebra
import qualified Data.List.Split as S
import System.Random
import Data.Array.IO
import Control.Monad
import GHC.Float
import Network

-- | Shuffle list
shuffle :: [a] -> IO [a]
shuffle xs = do
        ar <- newArray n xs
        forM [1..n] $ \i -> do
            j <- randomRIO (i,n)
            vi <- readArray ar i
            vj <- readArray ar j
            writeArray ar j vi
            return vj
  where
    n = length xs
    newArray :: Int -> [a] -> IO (IOArray Int a)
    newArray n = newListArray (1, n)

-- | Sigmoid Function
sigmoid :: Double -> Double
sigmoid x = 1.0 / (1.0 + exp (-x))

-- | Apply Sigmoid Function for every element of matrix
matrixSigmoid :: Matrix Double -> Matrix Double
matrixSigmoid = cmap sigmoid

-- | Sigmoid Prime Function
sigmoidPrime :: Double -> Double
sigmoidPrime x = sigmoid x * (1 - sigmoid x)

-- | Apply Sigmoid Prime Function for every element of matrix
matrixSigmoidPrime :: Matrix Double -> Matrix Double
matrixSigmoidPrime = cmap sigmoidPrime

-- | Cost function
costPrime :: Matrix Double -> Matrix Double  -> Matrix Double
costPrime outputActivations y = outputActivations - y

-- | Feed Forward algorithm
feedForward :: Net -- ^ Neural Network
            -> Matrix Double -- ^ Image Matrix
            -> Matrix Double
feedForward net x = foldl (\z (w,b)-> matrixSigmoid ((w <> z) + b)) x (zip (weights net) (biases net))

-- | Evaluate Neural Network predictions
evaluate :: Net -- ^ Neural Network for evaluation
         -> [(Matrix Double, Int)] -- ^ List of Test Images with their corresponding label
         -> Double -- ^ Returns Count of positive predictions
evaluate net testData = sum $ fmap (\(x,y) -> if x==y then 1.0 else 0.0) outcomes
  where
    testImages = map fst testData
    testLabels = map snd testData
    outcomes = zip (fmap (head . fmap maxIndex . toColumns . feedForward net) testImages) testLabels

-- | Get Neural Network predictions
getOutcomes :: Net -- ^ Neural Network
            -> [(Matrix Double, Int)] -- ^ List of Test Images with their corresponding labels
            -> [([Float], (Int, Int))] -- ^ Returns List of Test Images with their corresponding label and predictions
getOutcomes net testData = zip listImages outcomes
  where
    testImages = map fst testData
    testLabels = map snd testData
    listImages = head . toLists . tr . cmap double2Float <$> testImages
    outcomes = zip (fmap (head . fmap maxIndex . toColumns . feedForward net) testImages) testLabels

-- | Get List of neuron activations
neuronActivation :: [Matrix Double] -- ^ Weights of Neural Network
                 -> [Matrix Double] -- ^ Biases of Neural Network
                 -> Matrix Double   -- ^ Image Matrix
                 -> [Matrix Double]
neuronActivation [] [] x = []
neuronActivation (w:ws) (b:bs) x = z : neuronActivation ws bs (matrixSigmoid z)
  where
    z = (w <> x) + b

-- | Calculate deltas
calDelta :: Net -- ^ Neural Network
          -> [Matrix Double] -- ^ Weights of Neural Network
          -> [Matrix Double] -- ^ List of Neuron Activations
          -> Matrix Double   -- ^ delta
          -> Int -- ^ current layer
          -> [Matrix Double]
calDelta net w z delta l
  | l == numLayers net = [delta]
  | otherwise = calDelta net (tail w) (tail z) d (l + 1) ++ [delta]
    where
      sp = matrixSigmoidPrime (head z)
      d = tr (head w) <> delta * sp

-- | Back Propagation algorithm
backProp :: Net -- ^ Neural Network
         -> Matrix Double -- ^ Image Matrix
         -> Matrix Double -- ^ Image vectorized label
         -> ([Matrix Double],[Matrix Double])
backProp net x y = (nablaW, nablaB)
  where
    zs = neuronActivation (weights net) (biases net) x
    activations = x : (matrixSigmoid <$> zs)
    delta = (costPrime (last activations) y * matrixSigmoidPrime (last zs)) :: Matrix Double
    lDelta = calDelta net (reverse $ weights net) (tail $ reverse zs) delta 2
    nablaB = lDelta
    nablaW = (\ (x, y) -> x <> tr y) <$> zip lDelta activations

-- | Calculate Nabla for Weights and Biases
calNabla :: Net -- ^ Neural Network
         -> [(Matrix Double, Matrix Double)] -- ^ Batch of Images with their corresponding labels
         -> [Matrix Double] -- ^ List of Nabla Matrix's of Weights
         -> [Matrix Double] -- ^ List of Nabla Matrix's of Biases
         -> ([Matrix Double], [Matrix Double])
calNabla net [] nW nB = (nW, nB)
calNabla net (x:xs) nW nB = calNabla net xs nablaW nablaB
  where
    (deltaNablaW, deltaNablaB) = uncurry (backProp net) x
    nablaB = uncurry (+) <$> zip nB deltaNablaB
    nablaW = uncurry (+) <$> zip nW deltaNablaW

-- | Update network's weights and biases for each batch
updateMiniBatch :: Net -- ^ Neural Network
                -> [(Matrix Double, Matrix Double)] -- ^ Batch of Images with their corresponding labels
                -> Double -- ^ Alpha
                -> Net
updateMiniBatch net miniBatch alpha = newNet
  where
    nB = fmap (\x -> konst (0 :: Double) ((fst $ size x) :: Int, (snd $ size x) :: Int) :: Matrix Double) (biases net)
    nW = fmap (\x -> konst (0 :: Double) ((fst $ size x) :: Int, (snd $ size x) :: Int) :: Matrix Double) (weights net)
    (nablaW, nablaB) = calNabla net miniBatch nW nB
    newWeights = (\(x,y) -> x - cmap (\ p -> (alpha / fromIntegral (length miniBatch)) * p) y ) <$> zip (weights net) nablaW
    newBiases = (\(x,y) -> x - cmap (\ p -> (alpha / fromIntegral (length miniBatch)) * p) y ) <$> zip (biases net) nablaB
    newNet = Net (numLayers net) (sizes net) newBiases newWeights

-- | Train Neural Network
gradientDescent :: Net -- ^ Neural Network to train
                -> [(Matrix Double, Matrix Double)] -- ^ List of Training Images with column vectorized labels
                -> Int -- ^ Number of training iterations
                -> Int -- ^ Size of Image batches
                -> Double -- ^ Alpha
                -> [(Matrix Double, Int)] -- ^ List of Test Images with Labels
                -> IO Net
gradientDescent net trainingData epochs miniBatchSize alpha testData = loop net 0
  where
    nTest = fromIntegral (length testData)
    loop currentNet epoch = do
      let evaluation = 100.0 * evaluate currentNet testData / nTest
      let evaluationMsg = "Epoch: " ++ show epoch ++ " = " ++ show evaluation ++ "% Out Of " ++ show nTest ++ " Test Images"
      print evaluationMsg
      if epoch == epochs then return currentNet
      else do
        shuffledTrainingData <- shuffle trainingData
        let miniBatches = S.chunksOf miniBatchSize shuffledTrainingData
        let newNet = foldl (\z x-> updateMiniBatch z x alpha) currentNet miniBatches
        loop newNet (epoch + 1)


{-# LANGUAGE LambdaCase, BangPatterns #-}
module Stock where

import Data.Random.Normal
import System.Random
import Data.Random
import Diff
import Numeric.SpecFunctions
import Numeric.LinearAlgebra
import Util

ndigits :: Int
ndigits = 4

roundToN :: Int -> Double -> Double
roundToN n number = (fromInteger (round (number * (10^n)))) / (10.0^n)

-- constants
r :: Double
r = 0.06

variance :: Double
variance = 0.2

days :: Int
days = 360

yperiod :: Double
yperiod = (fromIntegral days) / 360.0

tdiff :: Int
tdiff = 40

tstep :: Double
tstep = (fromIntegral tdiff) / 360.0

s0 :: Double
s0 = 36.0000

k :: Double
k = 40.000

d1 :: Double
d1 = (log (s0 / k) + (r + variance / 2) * yperiod) / (sqrt (variance * yperiod))

d2 :: Double
d2 = d1 - sqrt (variance * yperiod)

-- c = s_0 * N(D_1) - exp (-r * T) * K * N(D_2)
n :: Double -> Double
n = cdf StdNormal

c :: Double
c = s0 * (n d1) - (exp (-r * yperiod)) * k * (n d2)

p :: Double
p = c + k * exp (-r * yperiod) - s0

--- basis functions
l :: Int -> String -> DiffLang
l n x
  | n >= 0    = (Exp ((Const (-0.5)) :*: (Var x))) :*: (Exp (Var x)) :*: (Const (1/(factorial n))) :*: (derivN n (laguerre n x) x)
  | otherwise = error "undefined basis function"

rzs :: Integer -> [Double]
rzs rand = normals' (0.0,1.0) (mkStdGen (fromIntegral rand))

sRand :: Integer -> Int -> Double
sRand j n =
  let zs = rzs j
  in s zs n

s :: [Double] -> Int -> R
s zs 0 = s0
s zs n
  | n > 0 = roundToN ndigits $ (s zs (n - 1)) * (exp $ (r - variance/2) * tstep + ((sqrt (variance * tstep)) * (zs!!n)))

cf :: [R] -> [R]
cf = map (\stdj -> max (k - stdj) 0)

cv :: R -> R
cv = (* (exp (-r * tstep)))

eval :: DiffLang -> R -> R
eval (Var _) x = x
eval (Const n) _ = n
eval (e1 :+: e2) x = eval e1 x + eval e2 x
eval (e1 :*: e2) x = eval e1 x * eval e2 x
eval (e1 :^: e2) x = eval e1 x ** eval e2 x
eval (Exp e) x     = exp (eval e x)

-- compute coefficients {beta_i} using (n-1)'th day's cash as predicator values,
-- n'th days cv as ys
regColN :: ([Vector R], [Vector R]) -> Int -> Matrix R
regColN (paths, cfs) n
  | length paths > n && n > 1 = regCol n laguerreBasis s t
  | otherwise = error "too big a N"
    where st :: ([R], [R], [R])
          st = unzip3 $ filter (\(a,b,c) -> a /= 0) $ zip3 (toList (cfs!!(n-1))) (toList (paths!!(n-1))) (toList (cfs!!n))
          s :: Vector R
          s = fromList $ sndOf3 st
          t :: Vector R
          t = fromList $ thdOf3 st

laguerreJ :: Int -> R -> R
laguerreJ j = eval (laguerre j "X")

laguerreBasis :: [R -> R]
laguerreBasis = map laguerreJ [0..(days `div` tdiff)]

regCol :: Int -> [R -> R] -> Vector R -> Vector R -> Matrix R -- n x 1 result
regCol n basis s t = linearSolveLS a b
  where a :: Matrix R
        -- a = fromColumns $ map (\j -> (fromList . map (basis!!j) . toList) s) [0..n-1]
        a = fromRows $ map (\x -> (fromList . map (\j -> (basis!!j) x)) [0..n-1]) (toList s)
        b :: Matrix R
        b = (asColumn . fromList . map cv) (toList t)

evalPoly :: R -> [R] -> R
evalPoly x betas = foldl (\acc beta -> acc + beta * x) 0 betas

evalPolyL :: [R] -> [R] -> [R]
evalPolyL xs betas = map (\x -> evalPoly x betas) xs

epsilon :: Double
epsilon = 0.00000001

cfP :: ([Vector R], [Vector R]) -> Int -> [Vector R]
cfP (paths, cfs) n
  | length paths > n && n > 1 = [newCFNP, finalCFN]
  | otherwise = error "undefined at the date!"
  where cfn = map (\p -> if fst p == 0 then 0 else snd p) $ zip oldCFP predCF
        finalCFN = fromList $ map (\(s,t) -> if s <= epsilon then 0.0 else t) (zip newCFN oldCFN)
        newCFN = map (\(s,t) -> if t >= s then t else 0.0) (zip oldCFP cfn)
        newCFNP = fromList $ map (\(s,t) -> if t >= s then 0.0 else s) (zip oldCFP cfn)
        -- newCFNP = fromList $ map (\(s,t) -> if t <= epsilon then 0.0 else s) (zip oldCFP cfn)
        oldCFN = toList $ cfs!!n
        oldCFP = toList $ cfs!!(n-1)
        predCF = map (\x -> (foldl (\acc -> \case { (n,beta) -> beta * (laguerreBasis!!n) x + acc }) 0 betas)) (toList (paths!!(n-1)))
        betas :: [(Int,R)]
        betas = zip [0..n-1] ((toList . head . toColumns . regColN (paths, cfs)) n)

cfGenRStep :: ([Vector R], [Vector R]) -> Int -> [Vector R]
cfGenRStep (pathsC, cfsC) n = take (n - 1) cfsC ++ cfP (pathsC, cfsC) n ++ drop (n + 1) cfsC

cfGenD :: (Matrix R, Matrix R) -> [Matrix R]
cfGenD (paths, cfs) = map (fromColumns) $ foldl (\cfsCC n ->
                               let cfs' = head cfsCC
                               in (take (n - 1) cfs' ++ cfP (pathsC, cfs') n ++ drop (n + 1) cfs') : cfsCC)
                      [toColumns cfs] (reverse [2..(days `div` tdiff)])
  where pathsC = toColumns paths

cfGenR :: (Matrix R, Matrix R) -> [Vector R]
cfGenR (paths, cfs) = foldl (\cfs' n -> take (n - 1) cfs' ++ cfP (pathsC, cfs') n ++ drop (n + 1) cfs') cfsC (reverse [2..(days `div` tdiff)])
  where pathsC = toColumns paths
        cfsC = toColumns cfs

tableaux :: Integer -> Integer -> [[R]]
tableaux rand size =
  map (\r -> map (sRand r) [0..(days `div` tdiff)]) ss
  where ss :: [Integer]
        ss = seeds rand size

seeds :: Integer -> Integer -> [Integer]
seeds rand length = take (fromIntegral length) $ map (\x -> (ceiling (abs x * 100000000000))) (rzs rand)

tableauxM :: Integer -> Integer -> (Matrix R, Matrix R)
tableauxM rand size = (matrix (days `div` tdiff + 1) paths, matrix (days `div` tdiff + 1) cfs)
  where paths = concat (tableaux rand size)
        cfs   = cf paths

cfGen :: (Matrix R, Matrix R) -> Matrix R
cfGen = fromColumns . tail . cfGenR

numNonZ :: Matrix R -> Int
numNonZ = sum . map (length . (filter (/= 0)) . toList) . toColumns

aOptionValue :: Matrix R -> [R]
aOptionValue m = map (\lst -> if null lst then 0.0 else (exp (-r * (fromIntegral (fst (last lst))) * tstep)) * (snd (last lst)))
                 $ map (filter (\case {(n,x) -> x /= 0}) . zip [1..(days `div` tdiff)] . toList) (toRows m)

eOptionValue :: Matrix R -> [R]
eOptionValue = map (* (exp (-r * (fromIntegral (days `div` tdiff + 1)) * tstep))) . cf . toList . last . toColumns

cStarMin :: Integer -> (Matrix R, Matrix R) -> R
cStarMin p (cf,m) = (/dist) . sum $ map (\i -> (vs!!(i-1) - vbar) * (qs!!(i-1) - qbar))  [1..(fromIntegral p)]
  where vs = aOptionValue cf
        qs = eOptionValue m
        vbar = sum vs / (fromIntegral (length vs))
        qbar = sum qs / (fromIntegral (length qs))
        dist = sum $ map (\i -> (qs!!(i-1) - qbar)**2) [1..(fromIntegral p)]

rbar :: Integer -> (Matrix R, Matrix R) -> R -> R
rbar pass (cf,m) c = (/(fromIntegral pass)) . sum $ map (\i -> (vs!!(i-1)) - c * ((qs!!(i-1)) - p)) [1..(fromIntegral pass)]
  where vs = aOptionValue cf
        qs = eOptionValue m
        qbar = sum qs / (fromIntegral (length qs))

output :: Integer -> Integer -> R
output rand size = rbar size (cfs, fst m) (cStarMin size (cfs, fst m))
  where cfs = cfGen m
        m = tableauxM rand size

genxs :: Integer -> [Double]
genxs n = map (\i -> let x = aOptionValue (cfGen (tableauxM i 100))
                     in sum x / (fromIntegral (length x)) ) [100..(100+n)]


module Lib where

import Codec.Picture
import System.Directory
import System.IO
import Data.Array.Repa as R
import Data.Array.Repa.Repr.Vector
import Data.Array.Repa.FFTW
import Data.Complex (Complex(..), realPart, magnitude)
import Prelude as P
import Control.Monad

normalize :: Monad m => Array U DIM2 Double -> m (Array D DIM2 Double)
normalize arr = do
    minn <- foldAllP min 1 arr
    let arrmin = R.map (+ (-minn)) arr
    maxx <- foldAllP max 0 arrmin
    let arrnorm = R.map (/maxx) arrmin
    return arrnorm

myfilter tresh v
    | v < tresh = fromIntegral v * 0.6
    | otherwise = 0

rowFilter :: Monad m => Array D DIM1 Double -> m (Array D DIM1 Double)
rowFilter row = do
    rowComplex <- computeP $ R.map (\x -> x :+ 0 ) row
    let fftF = fft $ rowComplex
        (Z :. w) = extent fftF
        barier = round (fromIntegral w/2)
        filtered = fromListUnboxed (Z :. w) $ (P.map (myfilter barier) [1..w])
        fftFilt = R.zipWith (*) fftF filtered
    ifftF <- fmap ifft $ computeP fftFilt
    return $ R.map realPart ifftF

getRow :: Array D DIM2 Double -> Int -> Array D DIM1 Double
getRow array n = slice array (Any :. n :. All)

mapRows :: Monad m => (Array D DIM1 Double -> m (Array D DIM1 Double)) -> Array D DIM2 Double -> m (Array U DIM2 Double)
mapRows func array = do
    let (Z :. w :. h) = extent array
    rows <- mapM (\num -> do
        let row = getRow array num
        func row) [0,1..(w-1)]
    let hugeRow = toList $ foldr1 append rows
    return $ fromListUnboxed (Z :. w :. h) hugeRow

reconstruct :: Array D DIM2 Double -> Double -> Int -> IO ()
reconstruct img p orgW = do
    let (Z :. w :. h) = extent img
        angleStep = pi / fromIntegral h
        anglesList = takeWhile (<pi) [a * angleStep | a <- [0..]]
        orgWNum = fromIntegral orgW

        listOfIndicies (x, y) = let
            list = P.map (\a -> round $ (x * sin a + y * cos a)/p + orgWNum/2) anglesList
            list_zip = zip list [0..(h-1)]
            in [(a,b) | (a,b) <- list_zip, a >= 0, a < w]

        render (x, y) = let
            list = listOfIndicies (x, y)
            pixelList = P.map (\(p, h) -> img ! (Z :. p :. h)) list
            pixelSum = sum pixelList
            avg = pixelSum / (fromIntegral $ length pixelList)
            in if avg > 0 then avg else 0

    let imageIndicies = [orgWNum/(-2)..orgWNum/2-1]
        img' = fromListUnboxed (Z :. orgW :. orgW) (P.map render [(a,b) | a <- imageIndicies , b <- imageIndicies])
    img <- normalize img'
    writePng "res/reconstruct.png" $ generateImage (\x y -> dToPx (img ! (Z :. x :. y))) orgW orgW

getY a p r x = (x, round ((p - (fromIntegral (x - r) * (cos a))) / (sin a)) + r)
getX a p r y = (round ((p - (fromIntegral (y - r) * (sin a))) / (cos a)) + r, y)

filterCoords r (x, y) = ((fromIntegral x - r)**2 + (fromIntegral y - r)**2) < r**2

getLineAvg :: Double -> Double -> Int -> (Double -> Double -> Int -> Int -> (Int, Int)) -> Array U DIM2 Double -> Double
getLineAvg a p w getCoord img =
    let pixelList = P.map (getCoord a p (div w 2)) [0..w-1]
        r = (fromIntegral w) / 2
        pixelList' = filter (filterCoords r) pixelList
        len = length pixelList'
        ret = (sum $ P.map (\(x, y) -> img ! (Z :. x :. y)) pixelList') / fromIntegral len
    in if isNaN ret then 0 else ret

getDetectorValue :: Double -> Double -> Int -> Array U DIM2 Double -> Double
getDetectorValue a p w img
    | a < pi/4 || a > (3 * pi)/4 = getLineAvg a p w getX img
    | otherwise = getLineAvg a p w getY img

processImage :: String -> Int -> Int -> Int -> IO ()
processImage fname nsteps nrays opening' = do
    (w, h, img) <- imageToArray fname
    putStrLn "Image loaded"
    let opening = (fromIntegral opening'/ 180) * pi
        r = fromIntegral w/2
        tresh = sin (opening/2) * r
        p = (sin (opening/2) * r) * 2 / (fromIntegral nrays)
        step = pi / fromIntegral nsteps
        angle = takeWhile (<pi) [x * step | x <- [0..]]
    putStrLn "Calculating projections"
    let listOfP = P.map (\v -> ((fromIntegral v) * p) - tresh) [0..nrays-1]
        rest = P.map (\a -> P.map (\p -> getDetectorValue a p w img) listOfP) angle
        trans = fromListUnboxed (Z :. length rest :. length (head rest)) $ concat $ rest
        trans' = transpose trans
    x <- computeUnboxedP trans'
    nofilter <- normalize x
    arraySaveToImage nofilter "nofilter"
    result' <- mapRows rowFilter nofilter
    result <- normalize result'
    arraySaveToImage result "filter"
    reconstruct result p w

dToPx x
    | x < 0 = PixelYA8 0 255
    | x > 1 = PixelYA8 255 255
    | otherwise = PixelYA8 (round (x * 255)) 255

arraySaveToImage :: Array D DIM2 Double -> String -> IO ()
arraySaveToImage arr fname = do
    let (Z :. w :. h) = extent arr
    writePng ("res/" P.++ fname P.++ ".png") $ generateImage (\x y -> dToPx $ arr ! (Z :. x :. y)) w h

rgbToGrey :: PixelRGB8 -> Double
rgbToGrey (PixelRGB8 r g b) =
  let rgb = P.zipWith (*) [0.2126, 0.7152, 0.0722] $ P.map (fromIntegral) [r, g, b]
  in fromInteger . round . sum $ rgb :: Double

imageToArray :: String -> IO (Int, Int, Array U DIM2 Double)
imageToArray fname = do
    Right imgRGB <- readImage fname
    let (imgGray@(Image w h _)) = convertRGB8 imgRGB
        arr = fromFunction (Z :. w :. h) (\(Z :. x :. y) -> rgbToGrey $ pixelAt imgGray x y)
    img <- computeUnboxedP arr
    return (w, h, img)


{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE BangPatterns #-}
{- # OPTIONS_GHC -ddump-simpl # -}
module Img where

import Data.Complex
import Data.Colour.SRGB
import qualified Data.ByteString.Lazy as BS

import Codec.Picture

import RNA

-- Functional evaluator

-- more efficient than Data.Fixed.fmod (I believe)
foreign import ccall unsafe "math.h fmod" c_fmod :: Double -> Double -> Double

fmod :: Double -> Double -> Double
x `fmod` y =
  let z = c_fmod x y in
  if z < 0 then z + y else z

data C = C {-# UNPACK #-} !Double {-# UNPACK #-} !Double {-# UNPACK #-} !Double

type Img = Complex Double -> Maybe C
type Pattern = Complex Double -> C

toImg :: RNA -> Img
toImg r pos =
    if magnitude pos > 1.0
    then Nothing
    else Just (i pos)
  where
    i = go r

tween :: Double -> C -> C -> C
tween x (C a1 b1 c1) (C a2 b2 c2) = C
  ((1 - x) * a1 + x * a2)
  ((1 - x) * b1 + x * b2)
  ((1 - x) * c1 + x * c2)

go :: RNA -> Pattern
go _ (!_pos) | False = undefined

go (Solid color) _pos = C r g b
  where
    RGB r g b = color

go Blend{} _pos = error "blend"

go (Checker x r1 r2) pos = do
    let !(tmpx :+ tmpy) = realToFrac x * ((1.0 :+ 0.0) + pos)
    if abs ((tmpx + 1) `fmod` 2 - 1) + abs (tmpy `fmod` 2 - 1) < 1
      then i2 pos
      else i1 pos
  where
    i1 = go r1
    i2 = go r2

go (Rotate x r1) pos = do
    let p' = phase pos + x
    let pos' = mkPolar (magnitude pos) p'
    i1 pos'
  where
    i1 = go r1

go (Invert r1) pos = do
    let mag = magnitude pos
    let mag' = 1 - mag
    let pos' = mkPolar mag' (phase pos)
    i1 pos'
  where
    i1 = go r1

go (Swirl x r1) pos = do
    let mag = magnitude pos
    let phase' = phase pos + (1.0 - mag) * x
    let pos' = mkPolar mag phase'
    i1 pos'
  where
    i1 = go r1

go (Dilated r r1) pos = do
    let p = phase pos
    let mag = magnitude pos
    let phase' = p + 1/fromIntegral r * sin (p * fromIntegral r)
    let pos' = mkPolar mag phase'
    i1 pos'
  where
    i1 = go r1

go (Rays r r1 r2) pos =
    if (phase pos / pi * fromIntegral r + 0.5) `fmod` 2 < 1
    then i2 pos
    else i1 pos
  where
    i1 = go r1
    i2 = go r2

go (Gradient r1 r2) pos =
    tween (magnitude pos) (i1 pos) (i2 pos)
  where
    i1 = go r1
    i2 = go r2

go (Ontop x r1 r2) pos = do
    let pos1 = realToFrac (1/x) * pos
    if magnitude pos < x
    then i1 pos1 else i2 pos
  where
    i1 = go r1
    i2 = go r2

img2Png :: Img -> BS.ByteString
img2Png i = encodePng $ generateImage colorAt w h
  where
    w = 1000
    h = 1000
    colorAt :: Int -> Int -> PixelRGBA8
    colorAt x y = do
      let x' = 2 * fromIntegral x / fromIntegral w - 1
      let y' = 2 * fromIntegral y / fromIntegral h - 1
      case i (y' :+ x') of
        Nothing -> PixelRGBA8 0 0 0 0
        Just (C r g b) -> PixelRGBA8 (clamp r) (clamp g) (clamp b) 255

    clamp :: Double -> Pixel8
    clamp = fromIntegral . max (0::Int) . min 255 . round . (* 256)


{-# LANGUAGE Rank2Types, FlexibleContexts #-}
module Numeric.AD.Lagrangian.Internal where
import Numeric.Optimization.Algorithms.HagerZhang05
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as S
import Numeric.AD
import GHC.IO                   (unsafePerformIO)
import Numeric.AD.Types
import Numeric.AD.Internal.Classes
import Numeric.LinearAlgebra.Algorithms
import qualified Data.Packed.Vector as V
import qualified Data.Packed.Matrix as M
import Numeric.AD.Internal.Tower

infixr 1 <=>
-- | Build a 'Constraint' from a function and a constant
(<=>) :: (forall s r. (Mode s, Mode r) => [AD2 s r a] -> AD2 s r a) -> a -> Constraint a
g <=> c = (FU g,c)

-- | A constraint of the form @g(x, y, ...) = c@. See '<=>' for constructing a 'Constraint'.
type Constraint a = (FU a, a)

type AD2 s r a = AD s (AD r a)

-- | A newtype wrapper for working with the rank 2 types constraint functions.
newtype FU a = FU {unFU :: forall s r. (Mode s, Mode r) => [AD2 s r a] -> AD2 s r a}

-- | This is the lagrangian multiplier solver. It is assumed that the
--   objective function and all of the constraints take in the
--   same amount of arguments.
minimize :: Double
      -> (forall s r. (Mode s, Mode r) => [AD2 s r Double] -> AD2 s r Double)
        -- ^ The function to minimize
      -> [Constraint Double]
      -- ^ The constraints as pairs @g \<=\> c@ which represent equations
      --   of the form @g(x, y, ...) = c@
      -> Int
      -- ^ The arity of the objective function which should equal the arity of
      --   the constraints.
      -> Either (Result, Statistics) (S.Vector Double, S.Vector Double)
      -- ^ Either an explanation of why the gradient descent failed or a pair
      --   containing the arguments at the minimum and the lagrange multipliers
minimize tolerance toMin constraints argCount = result where
    -- The function to minimize for the langrangian is the squared gradient
    obj argsAndLams =
        squaredGrad (lagrangian toMin constraints argCount) argsAndLams

    constraintCount = length constraints

    -- perhaps this should be exposed
    guess = U.replicate (argCount + constraintCount) (1.0 :: Double)

    result = case unsafePerformIO $
                    optimize
                        (defaultParameters { printFinal = False })
                        tolerance
                        guess
                        (VFunction (lowerFU obj . U.toList))
                        (VGradient (U.fromList . grad obj . U.toList))
                        Nothing of
       (vs, ToleranceStatisfied, _) -> Right (S.take argCount vs,
                                              S.drop argCount vs)
       (_, x, y) -> Left (x, y)

-- | Finding the maximum is the same as the minimum with the objective function inverted
maximize :: Double
      -> (forall s r. (Mode s, Mode r) => [AD2 s r Double] -> AD2 s r Double)
        -- ^ The function to maximize
      -> [Constraint Double]
      -- ^ The constraints as pairs @g \<=\> c@ which represent equations
      --   of the form @g(x, y, ...) = c@
      -> Int
      -- ^ The arity of the objective function which should equal the arity of
      --   the constraints.
      -> Either (Result, Statistics) (S.Vector Double, S.Vector Double)
      -- ^ Either an explanation of why the gradient descent failed or a pair
      --   containing the arguments at the minimum and the lagrange multipliers
maximize tolerance toMax constraints argCount =
    minimize tolerance (negate1 . toMax) constraints argCount

lagrangian :: (Num a, Mode s, Mode r)
           => (forall s r. (Mode s, Mode r) => [AD2 s r a] -> AD2 s r a)
           -> [Constraint a]
           -> Int
           -> [AD2 s r a]
           -> AD2 s r a
lagrangian f constraints argCount argsAndLams = result where
    args = take argCount argsAndLams
    lams = drop argCount argsAndLams

    -- g(x, y, ...) = c <=> g(x, y, ...) - c = 0
    appliedConstraints = fmap (\(FU f, c) -> f args - (auto . auto) c) constraints

    -- L(x, y, ..., lam0, ...) = f(x, y, ...) + lam0 * (g0 - c0) ...
    result = f args + (sum . zipWith (*) lams $ appliedConstraints)

squaredGrad :: Num a
            => (forall s. Mode s => [AD s a] -> AD s a)
            -> [a] -> a
squaredGrad f vs = sum . fmap (\x -> x*x) . grad f $ vs

-- | WARNING. Experimental.
--   This is not a true feasibility test for the function. I am not sure
--   exactly how to implement that. This just checks the feasiblility at a point.
--   If this ever returns false, 'solve' can fail.
feasible :: (forall s r. (Mode s, Mode r) => [AD2 s r Double] -> AD2 s r Double)
         -> [Constraint Double]
         -> [Double]
         -> Bool
feasible toMin constraints points = result where
    obj argsAndLams =
        squaredGrad (lagrangian toMin constraints $ length points) argsAndLams

    hessianMatrix = M.fromLists . hessian obj $ points

    -- make sure all of the eigenvalues are positive
    result = all (>0) . V.toList . eigenvaluesSH $ hessianMatrix