{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE BangPatterns #-}

module Main where

import Options.Applicative
import Data.Semigroup ((<>))
import Chart.Heatmap
import Control.Lens
import Control.Monad (forM_)
import Data.List (maximumBy)
import Data.Ord (comparing)
import Data.Colour.SRGB (sRGB, RGB(..))
import Data.Colour.RGBSpace.HSL (hsl)
import Data.Complex (magnitude, realPart, imagPart, Complex(..))
import Graphics.Rendering.Chart.Easy
import Graphics.Rendering.Chart.Backend.Cairo (toFile, FileOptions(..), FileFormat(..))
import Data.Array.CArray.Base (unsafeForeignPtrToCArray, toForeignPtr, CArray(..), ixmapWithIndP)
import Data.Fixed (mod')
import Math.FFT (dftRC)
import Numeric.Signal (analytic_signal, auto_correlation, deriv, Filterable)
import Data.Maybe (catMaybes)
import Sound.MIDI.File
import qualified Data.Vector.Storable as VS
import qualified Data.Vector as V
import qualified Sound.File.Sndfile as SF
import qualified Sound.File.Sndfile.Buffer.Vector as SFV
import qualified Sound.MIDI.File.Save as MS
import qualified Sound.MIDI.File.Event as ME
import qualified Sound.MIDI.Message.Channel as MC
import qualified Sound.MIDI.Message.Channel.Voice as MCV
import qualified Sound.MIDI.Message.Channel.Mode as MCM
import qualified Data.EventList.Relative.TimeBody as E

type Frequency = Double
type Amplitude = Double
type Step = Double

maxFrequency :: Double
maxFrequency = 7000

data NoteLetter = A
        | Ais
        | B
        | C
        | Cis
        | D
        | Dis
        | E
        | F
        | Fis
        | G
        | Gis
        deriving (Show, Eq)

type Octave = Int

data Note = Note NoteLetter Octave deriving (Show, Eq)

getBaseFreq :: NoteLetter -> Frequency
getBaseFreq A   = 27.50
getBaseFreq Ais = 29.14
getBaseFreq B   = 30.87
getBaseFreq C   = 16.35
getBaseFreq Cis = 17.32
getBaseFreq D   = 18.35
getBaseFreq Dis = 19.45
getBaseFreq E   = 20.60
getBaseFreq F   = 21.83
getBaseFreq Fis = 23.12
getBaseFreq G   = 24.50
getBaseFreq Gis = 25.96

getFreq :: Note -> Frequency
getFreq (Note letter n) = getBaseFreq letter * (fromIntegral n + 1)

getNote :: Frequency -> Maybe Note
getNote freq
  | divisor A `mod'` 1 < 0.02 || divisor A `mod'` 1 > 0.98 = Just
    (Note A (findN freq A))
  | divisor Ais `mod'` 1 < 0.02 || divisor Ais `mod'` 1 > 0.98 = Just
    (Note Ais (findN freq Ais))
  | divisor B `mod'` 1 < 0.02 || divisor B `mod'` 1 > 0.98 = Just
    (Note B (findN freq B))
  | divisor C `mod'` 1 < 0.02 || divisor C `mod'` 1 > 0.98 = Just
    (Note C (findN freq C))
  | divisor Cis `mod'` 1 < 0.02 || divisor Cis `mod'` 1 > 0.98 = Just
    (Note Cis (findN freq Cis))
  | divisor D `mod'` 1 < 0.02 || divisor D `mod'` 1 > 0.98 = Just
    (Note D (findN freq D))
  | divisor Dis `mod'` 1 < 0.02 || divisor Dis `mod'` 1 > 0.98 = Just
    (Note Dis (findN freq Dis))
  | divisor E `mod'` 1 < 0.02 || divisor E `mod'` 1 > 0.98 = Just
    (Note E (findN freq E))
  | divisor F `mod'` 1 < 0.02 || divisor F `mod'` 1 > 0.98 = Just
    (Note F (findN freq F))
  | divisor Fis `mod'` 1 < 0.02 || divisor Fis `mod'` 1 > 0.98 = Just
    (Note Fis (findN freq Fis))
  | divisor G `mod'` 1 < 0.02 || divisor G `mod'` 1 > 0.98 = Just
    (Note G (findN freq G))
  | divisor Gis `mod'` 1 < 0.02 || divisor Gis `mod'` 1 > 0.98 = Just
    (Note Gis (findN freq Gis))
  | otherwise = Nothing
 where
  findN f n = round $ logBase 2 (f / getBaseFreq n) :: Int
  divisor note = logBase 2 (freq / getBaseFreq note)

-- samples = N = number of samples
-- n = sample number
hannWindow :: Integer -> Integer -> Double
hannWindow samples n =
  0.5 * (1 - cos (2 * pi * fromIntegral n / (fromIntegral samples - 1)))

sinFunc :: Double -> Double -> Double -> Double -> Double
sinFunc samples cycle amplitude p =
  -amplitude * sin (2 * pi * cycle * p / samples)

cosFunc :: Double -> Double -> Double -> Double -> Double
cosFunc samples cycle amplitude p =
  amplitude * cos (2 * pi * cycle * p / samples)

performFFT :: VS.Vector Double -> Int -> IO (VS.Vector (Complex Double))
performFFT vector intSamples = do
  let (vec, x1, x2) = VS.unsafeToForeignPtr $ VS.take intSamples vector
  carray <- unsafeForeignPtrToCArray vec (x1, x2)
  let fft           = dftRC carray
  let (len, fftptr) = toForeignPtr fft
  return (VS.unsafeFromForeignPtr fftptr 0 len)

-- |Filter an array of maxima by partioning fftabs into bins of 100 elements, then finding the
-- largest value in that bin.
-- peaks are the list of indexes of the maxima.
filterMaxima :: V.Vector Int -> V.Vector Double -> V.Vector Int
filterMaxima peaks fftabs
  | V.null peaks = V.empty
  | otherwise = V.filter (> -1) $ V.map
    (\i ->
      let
        section = V.map
          (\x -> (x, fftabs V.! x))
          (V.filter (\x -> x < V.length fftabs && x >= i && x < i + binSize)
                    peaks
          )
      in  if V.null section
            then -1
            else fst $ V.maximumBy (comparing snd)
            -- list of (index, magnitude)
                                                   section
    )
    (V.fromList [0, binSize .. V.last peaks])
  where binSize = 100

-- | Find the maxima from fftabs by finding all points where the left and right points
-- are less than the current point.
findMaxima :: V.Vector Double -> V.Vector Int
findMaxima fftabs = V.filter
  (\x -> (diffs' V.! x) > 0 && (diffs' V.! (x + 1)) < 0)
  (V.fromList [1 .. (V.length diffs' - 2)])
 where
  diff f x = f x - f (x - 1)
  diffs' = V.map (diff (fftabs V.!)) (V.fromList [1 .. (V.length fftabs - 1)])

-- | Find maxima by finding points where that point is the maximum in it's neighbourhood.
findMaxima2 :: Filterable a => Int -> V.Vector a -> V.Vector (Int, a)
findMaxima2 radius v =
  V.filter (\(i, x) -> x == V.maximum (neighbours i))
    $ V.imap (\i x -> (i, x)) v
  where neighbours x0 = V.ifilter (\i _ -> abs (x0 - i) <= radius) v

findSignificantNotesForSection
  :: Int -> Int -> V.Vector (Complex Double) -> [(Frequency, Double)]
findSignificantNotesForSection sampleRate samples fftSection = fmap
  (\cycle ->
    ( fromIntegral cycle * fromIntegral sampleRate / fromIntegral samples
    , fftabs V.! cycle
    )
  )
  maxima'
 where
  fftabs  = V.map magnitude fftSection :: V.Vector Double
  peaks   = findMaxima fftabs
  peaks'  = filterMaxima peaks fftabs
  maxima  = V.toList peaks' -- tricky: the cycles returned is the cycles+1
  maxima' = filter (\value -> fftabs V.! value > 20) maxima -- Find 'large' maxima

midiPitch :: Note -> Int
midiPitch (Note letter n) =
  round $ 12 * logBase 2 (((2 ^ n) * getBaseFreq letter) / 440) + 69

-- | Find notes that standout in each stft section.
findSignificantNotes
  :: Int -> Int -> [(Step, [Complex Double])] -> String -> IO ()
findSignificantNotes sampleRate samples stft filePrefix = do
  notes <-
    mapM
      (\(step, fft) -> do
        let freqs =
              findSignificantNotesForSection sampleRate samples (V.fromList fft)
        let notes = fmap (\(f, _) -> getNote f)
                         (filter (\(f, _) -> f < maxFrequency) freqs)
        return $ map (\note -> (step, note)) $ catMaybes notes
      )
      stft :: IO [[(Step, Note)]]

  print $ map (map (\(s, note) -> (s, note, midiPitch note))) notes

  let
    sections = map
      (foldl
        (\acc (step, note) -> E.append
          acc
          (E.singleton
            (toElapsedTime 0)
            (ME.MIDIEvent
              (MC.Cons
                (MC.toChannel 1)
                (MC.Voice
                  (MCV.NoteOn (MC.toPitch (midiPitch note)) (MC.toVelocity 20))
                )
              )
            )
          )
        )
        E.empty
      )
      notes

  let
    t =
      E.append
          (foldl (\acc section -> E.append acc (E.delay 100 section))
                 E.empty
                 sections
          )
        $ E.singleton
            (toElapsedTime 100)
            (ME.MIDIEvent (MC.Cons (MC.toChannel 1) (MC.Mode MCM.AllNotesOff)))

  let midi = Cons Serial (Ticks 120) [t] :: T

  MS.toFile (filePrefix ++ ".midi") midi

--    toFile def{_fo_format=SVG} "diff.svg" $ do
--        layout_title .= "First derivative"
--        plot (points "diff" $ zip fftx (VS.toList fftabs))

--    print $ V.map (\x -> (x, fftabs VS.! x)) peaks'
    --toFile def{_fo_format=SVG} "diff2.svg" $ do
    --    layout_title .= "Second derivative"
    --    plot (points "diff" $ zip fftx diffs'')

    --let sinAmplitudes = take 500 $ fmap (\(cycle, _) -> (fromIntegral cycle, fftimg   VS.! cycle / fromIntegral (samples `div` 2))) maxima :: [(Double, Double)]
    --let cosAmplitudes = take 500 $ fmap (\(cycle, _) -> (fromIntegral cycle, fftreals VS.! cycle / fromIntegral (samples `div` 2))) maxima :: [(Double, Double)]

    --print "sin amps"
    --print sinAmplitudes
    --print "cos amps"
    --print cosAmplitudes

    --toFile def{_fo_format=SVG} "sine2.svg" $ do
    --    layout_title .= "curve"
    --    plot (line "data" [zip x y])
    --    plot
    --        (line "fit"
    --        [zip
    --            x
    --            (fmap
    --                (\p -> foldr (\(cycle, amplitude) total -> total + sinFunc (fromIntegral samples) cycle amplitude p) 0 sinAmplitudes +
    --                    foldr (\(cycle, amplitude) total -> total + cosFunc (fromIntegral samples) cycle amplitude p) 0 cosAmplitudes) x)])

    --let signal = VS.take samples $ analytic_signal vector
    --let signalReal = VS.toList $ VS.map realPart signal
    --let signalImag = VS.toList $ VS.map imagPart signal
    ----print $ VS.length signal
    --toFile def{_fo_format=SVG} "signal.svg" $ do
    --    layout_title .= "Hibert transform derivative"
    --    plot (line "Hilbert transform" [zip x (VS.toList $ VS.map imagPart signal)])
    --    plot (line "Original signal" [zip x intList])
    --    plot (line "Analytic signal" [zip x $ VS.toList (VS.map magnitude signal)])

autoCorrelate :: Filterable a => VS.Vector a -> Int -> IO ()
autoCorrelate vector sampleRate = do
  let result = auto_correlation 1000 vector
  let y      = map realToFrac $ VS.toList result :: [Double]
  --let values = [zip (map (\x -> x / fromIntegral sampleRate) [1..]) y] :: [[(Double, Double)]]
  let values = [zip [1 ..] y] :: [[(Double, Double)]]
  toFile def { _fo_format = SVG } "autocorrelation.svg" $ do
    layout_title .= "auto correlation"
    plot (line "" values)

  let localMaxima = findMaxima2 10 (VS.convert result)
  print $ V.map (\(i, x) -> (i, realToFrac x)) localMaxima
  print $ V.imap
    (\j (i, _) -> if j == 0 then j else i - fst (localMaxima V.! (j - 1)))
    localMaxima
  print $ fromIntegral sampleRate / fromIntegral
    (V.maximum
      (V.imap
        (\j (i, _) -> if j == 0 then j else i - fst (localMaxima V.! (j - 1)))
        localMaxima
      )
    )

graphSignal :: Filterable a => VS.Vector a -> Int -> IO ()
graphSignal vector sampleRate = do
  let intList = map realToFrac $ VS.toList vector :: [Double]
  let y       = intList
  let x       = map (\x -> x / fromIntegral sampleRate) [1 .. 200] :: [Double]
  toFile def { _fo_format = SVG } "signal.svg" $ do
    layout_x_axis . laxis_title .= "Elapsed time (s)"
    layout_title .= "signal curve"
    plot (line "" [zip x y])

graphFFTSignal :: VS.Vector (Complex Double) -> Int -> Int -> IO ()
graphFFTSignal fftvec samples sampleRate = do
  let fftreals = VS.toList $ VS.map realPart fftvec :: [Double]
  let fftimg   = VS.toList $ VS.map imagPart fftvec :: [Double]
  let fftabs   = VS.toList $ VS.map magnitude fftvec :: [Double]
  let fftx =
        map (\x -> fromIntegral (x * sampleRate) / fromIntegral samples)
            [0 .. 1000] :: [Double]
  toFile def { _fo_format = SVG } "fft.svg" $ do
    layout_title .= "fft curve"
    layout_x_axis . laxis_title .= "Frequency (Hz)"
    plot (points "abs" $ zip fftx fftabs)
    plot (points "real" $ zip fftx fftreals)
    plot (points "imaginary" $ zip fftx fftimg)

data Options = Options {
    sourceFile :: String,
    stepWidth :: Int
}

args :: ParserInfo Options
args = info
  (Options <$> strOption (long "src" <> metavar "SOURCE") <*> option
    auto
    (long "step-width" <> metavar "STEP_WIDTH" <> value 10000 <> showDefault)
  )
  (fullDesc <> progDesc "Desc" <> header "test")

main :: IO ()
main = do
  opts         <- execParser args
  (i, content) <-
    SF.readFile (sourceFile opts) :: IO (SF.Info, Maybe (SFV.Buffer Double))
  let sampleRate = SF.samplerate i
  let samples    = SF.frames i * SF.channels i
  print samples
  print i
  let v = fmap SFV.fromBuffer content
  case v of
    Nothing     -> print "Nothing!"
    Just vector -> do
      graphSignal vector sampleRate

      -- Run FFT on whole dataset and plot the result
      fftvec <- performFFT vector samples
      graphFFTSignal fftvec samples sampleRate
      findSignificantNotes sampleRate samples [(0, VS.toList fftvec)] "fft"

      -- get auto correlation
      autoCorrelate vector sampleRate

      -- plot histogram
      let (vec, x1, x2) = VS.unsafeToForeignPtr vector
      carray <- unsafeForeignPtrToCArray vec (x1, x2)
      --let stepWidth = fromIntegral samples :: Integer
      let stepWidth = 5000 :: Integer
      print $ "Step width: " ++ show stepWidth
      let cutSamples =
            (samples `div` fromIntegral stepWidth) * fromIntegral stepWidth :: Int

      -- windows of size stepWidth
      let stft =
            fmap (stftSection carray cutSamples stepWidth)
                 [0, stepWidth `div` 2 .. fromIntegral cutSamples] :: [ ( Step
                , [Complex Double]
                )
              ]

      findSignificantNotes sampleRate cutSamples stft "stft"

      -- (window start, fft for window)
      let plots' =
            concatMap
              (spectrogramIntensity (fromIntegral cutSamples) sampleRate)
              stft :: [(Step, Frequency, Amplitude)]

      let stftData  = generateStftData (fromIntegral stepWidth) plots'
      let stftData' = filter (\((_, _), (_, _), z) -> z > 1) stftData
      toFile def { _fo_format = PNG } "stft.png" $ do
        layout_title .= "curve"
        plot
          (liftEC $ do
            plot_rect_title .= "title"
            plot_rect_style .= \amp ->
              let RGB a b c = hsl ((1 - amp) * 240) 1 (amp * 0.6)
              in  solidFillStyle (opaque $ sRGB a b c)
            plot_rect_values .= stftData'
          )

      return ()

-- convert to plottable set of values
generateStftData
  :: Step
  -> [(Step, Frequency, Amplitude)]
  -> [((Double, Double), (Double, Double), Amplitude)]
generateStftData windowSize = map (generateColumn windowSize)

-- column in spectrogram
generateColumn
  :: Step
  -> (Step, Frequency, Amplitude)
  -> ((Double, Double), (Double, Double), Amplitude)
generateColumn windowSize (step, freq, amp) =
  ((step, windowSize + step), (freq, freq + 1), amp)

spectrogramIntensity
  :: Double -> Int -> (Step, [Complex Double]) -> [(Step, Frequency, Amplitude)]
spectrogramIntensity totalSamples sampleRate (step, fft) =
  let rate = fromIntegral sampleRate / totalSamples
  in  filter
        (\(_, f, _) -> f < 4000)
        (map (\(cycle, x) -> (step, cycle * rate, magnitude (x + 0) ** 2))
             (zip [0 ..] fft)
        )

ftSignals :: Double -> [Complex Double] -> [Double] -> [Double]
ftSignals totalSamples fft = fmap
  (\x ->
    foldr
        (\(cycle, amplitude) total ->
          total + sinFunc totalSamples cycle amplitude x
        )
        0
        sinAmplitudes
      + foldr
          (\(cycle, amplitude) total ->
            total + cosFunc totalSamples cycle amplitude x
          )
          0
          cosAmplitudes
  )
 where
  significantSignals = filter (\x -> magnitude x > 5) fft
  cycles             = zip [0 ..] significantSignals
  sinAmplitudes =
    fmap
      (\(cycle, signal) ->
        (fromIntegral cycle, realPart signal / (totalSamples / 2.0))
      )
      cycles :: [(Double, Double)]
  cosAmplitudes =
    fmap
      (\(cycle, signal) ->
        (fromIntegral cycle, imagPart signal / (totalSamples / 2.0))
      )
      cycles :: [(Double, Double)]

-- | Apply window function to the set of samples
windowedSection
  :: CArray Int Double -> Int -> Integer -> Integer -> CArray Int Double
windowedSection data_ totalSamples n stepWidth = ixmapWithIndP
  (0, totalSamples)
  id
  (\_ x i' -> if i' < fromIntegral n || i' >= fromIntegral (n + stepWidth)
    then 0
    else x * hannWindow (fromIntegral stepWidth) (fromIntegral i' - n)
  )
  data_

-- | n = start of step sample number
stftSection
  :: CArray Int Double -> Int -> Integer -> Integer -> (Step, [Complex Double])
stftSection !completeSample totalSamples stepWidth n =
  (fromIntegral n, VS.toList $ VS.unsafeFromForeignPtr fftptr 0 len)
 where
  fft = dftRC (windowedSection completeSample totalSamples n stepWidth)
  (len, fftptr) = toForeignPtr fft


module SynapseSpec
  ( spec
  ) where

import qualified Numeric.LinearAlgebra as A
import           Synapse
import           Test.Hspec
import           Test.Hspec.QuickCheck (prop)
import           Test.QuickCheck

runTest :: IO ()
runTest = hspec spec

spec :: Spec
spec = do
  describe "Double" propsDouble
  describe "Sigmoid" propsSigmoid
  describe "ReLU" propsReLU
  describe "Softmax" propsSoftmax
  describe "CrossEntropy" propsCrossEntropy
  describe "SoftmaxWithCross" propsSoftmaxWithCross
  describe "Affine" propsAffine

propsDouble = do
  prop "multiple" $
    forAll (vectorOf 3 (genSignal `suchThat` small)) $ \[a, b, c] ->
      let x3 a = a ^ 3
       in x3 a * x3 b * x3 c `shouldNotBe` 0
  prop "roundy" $
    forAll (genSignal `suchThat` small) $ \a ->
      let b = abs a
          x = foldr (\i j -> b ^ i + j) 0 [0 .. 6]
          y = roundy 4 x
          e = 0.0001
       in (y < x && (x - e) < y) `shouldBe` True
  where
    small a = a /= 0 && -1 < a && a < 1

propsSigmoid = do
  prop "make with exp" $
    forAll genSignal $ \n ->
      let r = sigmoid n
          x = 1.0 / (1.0 + exp (-n))
       in r `shouldBe` x
  prop "apply all elements" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals 1 (m * n)) $ \[vs] ->
        let r = sigmoidm $ (m A.>< n) vs
            rs = A.toList `concatMap` A.toRows r
            xs = sigmoid `map` vs
         in rs `shouldBe` xs
  prop "calc diff" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals 2 (m * n)) $ \[ys, ds] ->
        let r = sigmoidBackward ((m A.>< n) ys) ((m A.>< n) ds)
            rs = A.toList `concatMap` A.toRows r
            xs = zipWith (\y d -> d * (1.0 - y) * y) ys ds
         in rs `shouldBe` xs

propsReLU = do
  prop "take max" $
    forAll genSignal $ \v ->
      let r = relu v
          x = 0 `max` v
       in r `shouldBe` x
  prop "apply all elements" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals 1 (m * n)) $ \[vs] ->
        let r = relum $ (m A.>< n) vs
            rs = A.toList `concatMap` A.toRows r
            xs = relu `map` vs
         in rs `shouldBe` xs
  prop "calc diff" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals 2 (m * n)) $ \[ys, ds] ->
        let r = relumBackward ((m A.>< n) ys) ((m A.>< n) ds)
            rs = A.toList `concatMap` A.toRows r
            reverse x =
              if 0 < x
                then 1
                else 0
            xs = zipWith (\x d -> d * fromIntegral (reverse x)) ys ds
         in rs `shouldBe` xs

propsSoftmax = do
  prop "take max" $
    forAll (listOf1 genSignal) $ \vs ->
      let r = softmax $ A.vector vs
          c = maximum vs
          ds = map (\v -> exp $ v - c) vs
          s = sum ds
          xs = map (/ s) ds
       in A.toList r `shouldBe` xs
  prop "apply all rows" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals m n) $ \vss ->
        let r = softmaxm $ (m A.>< n) $ concat vss
            xs = map (softmax . A.vector) vss
         in A.toRows r `shouldBe` xs

propsCrossEntropy = do
  prop "take sum" $
    forAll (listOf1 genSignal) $ \vs ->
      forAll (genTeacher $ length vs) $ \ts ->
        let ys = softmax $ A.vector vs
            d = 1e-10
            r = crossEntropy (A.vector ts) ys
            cs = zipWith (\t y -> t * log (y + d)) ts $ A.toList ys
            x = -sum cs
         in r `shouldBe` x
  prop "apply all rows" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals m n) $ \vss ->
        forAll (vectorOf m $ genTeacher n) $ \tss ->
          let ys = map (softmax . A.vector) vss
              ts = A.vector `map` tss
              r = crossEntropym (A.fromRows ts) (A.fromRows ys)
              xs = zipWith crossEntropy ts ys
              x = sum xs / fromIntegral m
           in r `shouldBe` x

propsSoftmaxWithCross = do
  prop "composite" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals m n) $ \vss ->
        forAll (vectorOf m $ genTeacher n) $ \tss ->
          let tx = A.fromRows $ map A.vector tss
              vx = A.fromRows $ map A.vector vss
              r = softmaxWithCross tx vx
              y = softmaxm vx
              c = crossEntropym tx y
           in r `shouldBe` (y, c)
  prop "calc diff" $
    forAll genMN $ \(m, n) ->
      forAll (genSignals m n) $ \vss ->
        forAll (vectorOf m $ genTeacher n) $ \tss ->
          let ts = concat tss
              tx = (m A.>< n) ts
              y = softmaxm $ A.fromLists vss
              ys = concat $ A.toLists y
              r = softmaxWithCrossBackward tx y
              d = (m A.>< n) $ zipWith (\y t -> (y - t) / fromIntegral m) ys ts
           in reduce r `shouldBe` reduce d

propsAffine = do
  prop "forward" $
    forAll (choose (1, 10)) $ \l ->
      forAll genMN $ \(m, n) ->
        forAll (genSignals m n) $ \wss ->
          forAll (vectorOf n genSignal) $ \bs ->
            forAll (genSignals l m) $ \xss ->
              let x = A.fromLists xss
                  w = A.fromLists wss
                  b = A.vector bs
                  y = affinem w b x
                  rows = zipWith row (A.toRows x) (replicate l $ A.toColumns w)
                  r = A.fromLists $ map (zipWith (+) bs) rows
               in reduce y `shouldBe` reduce r
  prop "backward" $
    forAll (choose (1, 10)) $ \l ->
      forAll genMN $ \(m, n) ->
        forAll (genSignals m n) $ \wss ->
          forAll (genSignals l m) $ \xss ->
            forAll (genSignals l n) $ \dss ->
              let x = A.fromLists xss
                  w = A.fromLists wss
                  d = A.fromLists dss
                  (nx, nw, nb) = affinemBackward w x d
                  dx =
                    A.fromLists $
                    zipWith row (A.toRows d) (replicate l $ A.toRows w)
                  dw =
                    A.fromLists $
                    zipWith row (A.toColumns x) (replicate m $ A.toColumns d)
                  db = A.vector $ foldr (zipWith (+)) (replicate n 0.0) dss
               in (reduce nx, reduce nw, reduce $ A.asRow nb) `shouldBe`
                  (reduce dx, reduce dw, reduce $ A.asRow db)

genSignal :: Gen Double
genSignal = (arbitrary :: Gen Double) `suchThat` (\a -> -10 <= a && a <= 10)

genSignals :: Int -> Int -> Gen [[Double]]
genSignals m n = vectorOf m (vectorOf n genSignal)

genMN :: Gen (Int, Int)
genMN = do
  m <- choose (1, 10)
  n <- choose (1, 10)
  return (m, n)

genTeacher :: Int -> Gen [Double]
genTeacher n = do
  i <- choose (0, n - 1)
  return $ replicate i 0.0 ++ 1 : replicate (n - i - 1) 0.0

roundy n v = fromIntegral (floor $ v * 10 ^ n) / 10.0 ^^ n

reduce a = map (roundy 8) $ concat $ A.toLists a

cell arow bcol = sum $ zipWith (*) (A.toList arow) (A.toList bcol)

row arow = map (cell arow)


{-# LANGUAGE RecordWildCards #-}

import Control.Monad.Trans.Except
import Data.Word
import Data.Complex
import Data.Maybe
import Data.Monoid

import Control.Error.Util
import Pipes                as P
import Pipes.Prelude        as P
import Options.Applicative
import Data.Vector.Storable as VS hiding ((++))
import Data.Vector.Generic  as VG hiding ((++))

import SDR.Util             as U
import SDR.RTLSDRStream
import SDR.Filter
import SDR.ArgUtils
import SDR.FilterDesign
import SDR.Pulse
import SDR.PipeUtils
import SDR.CPUID

data Options = Options {
    frequency    :: Word32,
    bandwidth    :: Maybe Int
}

optParser :: Parser Options
optParser = Options
          <$> option (fmap fromIntegral parseSize) (
                 long "frequency"
              <> short 'f'
              <> metavar "FREQUENCY"
              <> help "Frequency to tune to"
              )
          <*> optional (option (fmap fromIntegral parseSize) (
                 long "bandwidth"
              <> short 'b'
              <> metavar "BANDWIDTH"
              <> help "Filter bandwidth. From 0 to 32K."
              ))

opt :: ParserInfo Options
opt = info (helper <*> optParser) (fullDesc <> progDesc "Receive AM Radio" <> header "RTLSDR AM")

size = 4096

doIt Options{..} = do

    info            <- lift getCPUInfo

    str             <- sdrStream (defaultRTLSDRParams (frequency + 256000) 1024000) 1 (fromIntegral $ size * 2)

    let coeffsDecim :: [Float]
        coeffsDecim =  VS.toList $ windowedSinc 71 0.4 blackman
    deci            <- lift $ fastDecimatorC info 2 coeffsDecim

    let coeffsFilt  :: [Float]
        coeffsFilt  =  VS.toList $ windowedSinc 71 (fromIntegral (fromMaybe 16000 bandwidth) / 32000) blackman
    filt            <- lift $ fastFilterC info coeffsFilt

    let coeffsResp  :: [Float]
        coeffsResp  =  VS.toList $ windowedSinc 71 0.25 blackman
    resp            <- lift $ fastResamplerR info 3 2 coeffsResp

    pulseSink       <- lift pulseAudioSink

    lift $ runEffect $   str
                     >-> P.map (interleavedIQUnsignedByteToFloatFast info)
                     >-> P.map (VG.zipWith (*) (quarterBandUp size))
                     >-> firDecimator deci size
                     >-> firDecimator deci size
                     >-> firDecimator deci size
                     >-> firDecimator deci size
                     >-> firDecimator deci size
                     >-> firFilter  filt size
                     >-> P.map (VG.map magnitude)
                     >-> dcBlockingFilter
                     >-> firResampler resp size
                     >-> P.map (VG.map (* 6))
                     >-> pulseSink

main = execParser opt >>= exceptT putStrLn return . doIt


{- -}

module Neural.Digits where

import Control.DeepSeq
import Control.Monad
import Numeric.LinearAlgebra
import System.Random

import Neural.FF
import Neural.LayerSpawners
import Neural.MNIST

-- 1st layer fn is a dummy. I might replace with a dedicated dummy layer fn
spawnDigitClassifierNet :: IO (FFNetwork)
spawnDigitClassifierNet = spawnFFNetwork [(fullyConnectedLayer, (0,0), (28, 28), layerOpts),
                                          (convLayer, (28, 28), (5,5), layerOpts {
                                              layerDepth = 5,
                                              pooling = maxPooling,
                                              poolingBP = maxPoolingBP,
                                              poolSize = (2,2)
                                              }),
                                          (fullyConnectedLayer, (12, 12), (10, 1), layerOpts)]

trainDigitClassifierNet :: FFNetwork -> [((Matrix Double), Int)] -> [((Matrix Double), Int)] -> Int -> IO (FFNetwork, [(Matrix Double, Int)])
trainDigitClassifierNet net trainingSet testSet nEpochs = go net trainingSet 1
  where go net ts i | i <= nEpochs = do
                        (newNet, newTS) <- getStdRandom (\gen -> sgd net gen ts 9 16 0.1 5.0)
                        let nCorrect = (evaluate newNet testSet 5.0)

                        putStrLn $ "Epoch " ++ (show i) ++ " of " ++ (show nEpochs) ++ ": "
                          ++ (show nCorrect) ++ "/" ++ (show $ length testSet)

                        {-
                        -- print out features
                        mapM_ (\layer ->
                                if (length $ weights layer) > 1
                                then (printMNISTImage $
                                      (\(w_1:w_) ->
                                        foldl (\w w_ -> w ||| (konst 0.0 (5,3) :: Matrix Double) ||| w_) w_1 w_)
                                      $ weights layer) >> putStrLn ""
                                else putStrLn "") $ layers net -}

                        go newNet newTS $ i + 1

                    | otherwise = return (net, trainingSet)

testDigitClassifierNet :: FFNetwork -> [((Matrix Double), Int)] -> IO ()
testDigitClassifierNet net testSet = do
  let tests = take 3 testSet
  putStrLn "Example tests: "
  printMNISTImage (foldl (\mat (input,expected) -> mat ||| input) ((\(x,y) -> x) (head tests)) (tail tests))

  putStrLn $ "Guesses: " ++
    (show $ map (\(input,expected) -> convertOutput $ ffOut net [input]) tests)

  putStrLn $ "Full evaluation: " ++ (show $ evaluate net testSet 1.0) ++ "/" ++ (show $ length testSet)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
-- | Fourier amplitude sensitivity test, a quasi-monte carlo method
--
-- does not yet play well with "ErrorProp.MonteCarlo", though a function
-- like 'ErrorProp.MonteCarlo.sample' should be straightforward to
-- implement. However, 'fast' provides the core of such a function
module ErrorProp.QuasiMC where

import qualified Data.Map as M
import Math.Polynomial
import Math.Polynomial.Bernoulli
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Generic as VG
import Data.List

import Data.Numbers.Primes
import Data.VectorSpace

import ErrorProp.QuasiMC.Cukier (cukierOmegas)

import Statistics.Sample
import Data.Maybe

import ErrorProp.MonteCarlo
import ErrorProp.Common

f1 alpha2 = constPoly 1 ^+^ ((-1)^(1 + alpha2 `mod` 2) * (2*pi)^(2*alpha2)) *^ bx
  where bx = bernoulliPoly !! (2*alpha2)

-- the explicit equations for F2,F4,F6 (equation 4.14,15,16) are correct
prop_f11 x = f11 x ~= f1 1 `evalPoly` x
  where f11 x = 1 + 2 * pi^2 * (x^2 - x + 1/6)

prop_f12 x = f12 x ~= f1 2 `evalPoly` x
  where f12 x = 1 + pi^4/45 * (1 - 30* x^2 *  (1-x)^2)

prop_f13 x = f13 x ~= f1 3 `evalPoly` x
  where f13 x = 1 + (2*pi^6/945) * (1 - 21*x^2 + 105*x^4 -126*x^5 + 42*x^6)

x ~= y = (2 * abs(x-y) / (1 + abs x + abs y)) < (1e-8 :: Double)
infix 4 ~=

-- eqn 4.10. Is the function integrated to generate the "best"
-- lattice points
f2 alpha2 xs = V.product $ V.map (f1 alpha2 `evalPoly`) xs

-- | Korbov's construction for good lattice points. Is an LCG
--
-- 4.32
korbov l n s = V.iterateN s (\p -> (p * l) `mod` n) 1

-- | try all values of L up to n/2
korbovLExhaustive alpha2 n s =
    head $
    sort [ (abs (r1-1) :: Double, l)
            | l <- [1 .. n `div` 2 ],
              let r1 = rank1 (korbov l n s) n (V.replicate s 0) (f2 alpha2) ]

-- | assume L has k prime factors. Provides candidate L values,
-- and the error in the integral used to decide that L is the best.
-- The error is absolute and relative, since the target function
-- integrates to 1.
korbovLCompositeK
  :: Int -- ^ alpha/2, roughness of the function being integrated
  -> Int -- ^ k
  -> Int -- ^ n, number of points in the quadrature
  -> Int -- ^ s, dimension
  -> [(Double, Int)] -- ^ (error, best l)
korbovLCompositeK alpha2 k n s =
    sort [ (abs (rankL l - 1) :: Double, l)
            | l <- go k ps 1 ]
  where
    rankL l = rank1 (korbov l n s) n (V.replicate s 0) (f2 alpha2)
    go 0 _ n = [n]
    go k ps1 n = do
      q : ps1' @ (_:_) <- tails ps1
      go (k-1) ps1' (n*q)

    ps = takeWhile ( < round ( (fromIntegral n / 2) ** (1/fromIntegral k)))
                primes

{- | rank1 shifted lattice rule

Given

> f :: [0,1)^s -> R

calculate

> Q(Z,N,C)f = (1/N) sum_j=0^N f { j z/N + c }

is an approximation to the integral of @f@ (@If@) over that region.

equation 4.42 Sloan&Joe 1994 pg 89

If C is multivariant uniform [0,1)^s, we can do a CI for the actual
integral

P( |Q - I| < sigma v) >= 1 - (1/v^2)

with sigma^2 estimated by sum_k=1^q (Q - Qbar)^2 / ( q (q-1))

is unbiased. Quantiles are biased (Lemieux 2008). Bootstrap/jacknife resample to estimate bias.

trapezoidal / rectangle rule is good for the periodic integral here,
but possibly try "GA Evans, JR Webster 1999"

-}
rank1 :: (Integral i, RealFrac a, VG.Vector v a, VG.Vector vi i,
            VG.Vector vi a)
  => vi i -- ^ z
  -> Int -- ^ N
  -> v a -- ^ c
  -> (v a -> a) -- ^ f
  -> a
rank1 z n c f = go 0 0
  where
    zn = VG.convert $ VG.map (\zi -> fromIntegral zi / fromIntegral n) z

    go !j !accum
      | j == VG.length zn = accum
      | otherwise = go (j+1) $ f $ VG.zipWith (+) c
              $ VG.map (\zni -> fracPart $ fromIntegral j*zni) zn

{- |

>>> map fracPart [12.3, -2.5]
[0.3, 0.5]

-}
fracPart x = x - fromInteger (floor x)

-- | fourier amplitude sensitivity test
fast ::
  (RealFrac b, VG.Vector V.Vector b) =>
    V.Vector Int -- ^ N incommensurate frequencies (see 'korbov', 'cukierOmegas', or others)
  -> (V.Vector b -> Double) -- ^ a function of N variables which will be sampled on @[0, 1)^N@
  -> V.Vector Double -- ^ fraction of the total variance attributed to the i'th variable
fast omegas f =
    let omegaMax = VG.maximum omegas
        q = omegaMax

        xj k = VG.map (\w -> fracPart $ fromIntegral (w*k) / fromIntegral (2*q+1)) omegas
        sks = V.enumFromTo (-q) q
        fks = VG.map (f . xj) sks

        withK k = VG.map (\w -> mean
                              $ VG.zipWith (\s f -> f * k (w*s)) sks fks)
                omegas

        nq = fromIntegral (2*q+1)

        a i = withK (cos . (*(2*i*pi/nq)) . fromIntegral)
        b i = withK (sin . (*(2*i*pi/nq)) . fromIntegral)

        vjs = foldl1 (V.zipWith (+))
              $ map (V.map (^2))
              [  if odd j then b fij else a fij
                  | j <- [1 .. 2],
                    let fij = fromIntegral j ]

        -- normalizeFac = 2 / variance fks -- XXX the
        -- correct normalization should be some analytical
        -- expression... like above, not the value below
        normalizeFac = recip $ V.sum vjs

    in V.map (*normalizeFac) vjs

fastKi s np k =
  let n = primes !! np
      (_, l) = korbovLExhaustive 1 n s
      omegas = korbov l n s

  in fast omegas (\v -> v V.! k)

fastKis np s = [ (k,fastKi s np k)  | k <- [0 .. s-1]]
{- properties.

> ki k = fast omegas (\x -> x V.! k)

> ki k V.! j == 1 if k==j else 0

-}

-- | `pickNP err s` uses 'fastKis' to pick the parameter np, such that the
-- value that should be 1 is within err of 1, for a problem of dimension s
pickNP err s = head [ np
    | np <- [ 2 .. ],
     all ok (fastKis np s) ]
  where ok (i, v) = abs (1 - v V.! i) < err

sampleQMC :: Double -- ^ error tolerance (see 'pickNP')
  -> MV Double
  -> [(String, Double)] -- variable name, fraction of the total variance
sampleQMC err mv =
  let n = wheelSieve 5 !! pickNP err s
      (_, l) = korbovLExhaustive 1 n s
      omegas = korbov l n s
      s = M.size (mv_vars mv)

      f v = mv_sample mv (updateMV v (mv_vars mv))
  in M.keys (mv_vars mv) `zip` VG.toList (fast omegas f)

updateMV :: (Floating b, V.Unbox b) => V.Vector b -> M.Map String (Maybe b) -> M.Map String b
updateMV v m = snd $ M.mapAccum (\ n e -> (n+1, fromMaybe 1 e * toUnitVar (v V.! n))) 0 m
  where
        toUnitVar x = sqrt 12 * (x - 0.5)


{-# LANGUAGE PartialTypeSignatures, ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-orphans -funbox-strict-fields #-}
-- | Provides @'Matrix'@ and @'MMatrix'@ instances for
--  @'LA.Matrix'@ type of @hmatrix@ package.
module Algebra.Matrix.HMatrix (STMatrix) where
import           Algebra.Matrix.Generic
import           Algebra.Matrix.Generic.Mutable (MMatrix (..))
import           Algebra.Prelude.Core
import           Control.Monad                  (forM_)
import           Control.Monad.Primitive
import qualified Data.Coerce                    as DC
import qualified Data.Vector.Storable           as SV
import qualified Numeric.LinearAlgebra          as LA
import qualified Numeric.LinearAlgebra.Devel    as LA

type instance Mutable LA.Matrix = STMatrix
type instance Row LA.Matrix = LA.Vector
type instance Column LA.Matrix = LA.Vector
type instance Row STMatrix = LA.Vector
type instance Column STMatrix = LA.Vector

-- | Wrapper type for @hmatrix@'s @'LA.STMatrix'@.
data STMatrix s a = STMat { stmRows :: !Int
                          , stmCols :: !Int
                          , unwrapSTM :: LA.STMatrix s a
                          }

instance (Num a, LA.Container LA.Matrix a) => Matrix LA.Matrix (WrapNum a) where
  basicRowCount = LA.rows
  basicColumnCount = LA.cols
  unsafeFromRows rs = DC.coerce $ LA.fromRows (DC.coerce rs :: [LA.Vector a])
  unsafeFromColumns cs = DC.coerce $ LA.fromColumns (DC.coerce cs :: [LA.Vector a])
  basicUnsafeIndexM m i j =
    return $ WrapNum $ (DC.coerce m :: LA.Matrix a) `LA.atIndex` (i, j)
  unsafeFreeze = stToPrim . LA.unsafeFreezeMatrix . unwrapSTM
  unsafeThaw m = stToPrim $ STMat (LA.rows m) (LA.cols m) <$> LA.unsafeThawMatrix m
  basicUnsafeGetRowM m i = return $ DC.coerce $ getRowLA  i (DC.coerce m :: LA.Matrix a)
  basicUnsafeGetColumnM m i =
    return $ DC.coerce $ head $ LA.toColumns (DC.coerce m LA.¿ [i] :: LA.Matrix a)
  unsafeGenerate w h f =
    DC.coerce $
    LA.fromRows [ SV.generate w $ \j -> unwrapNum $ f i j
                | i <- [0..h - 1]
                ]

instance (Integral a, LA.Container LA.Matrix a) => Matrix LA.Matrix (WrapIntegral a) where
  basicRowCount = LA.rows
  basicColumnCount = LA.cols
  unsafeFromRows rs = DC.coerce $ LA.fromRows (DC.coerce rs :: [LA.Vector a])
  unsafeFromColumns cs = DC.coerce $ LA.fromColumns (DC.coerce cs :: [LA.Vector a])
  basicUnsafeIndexM m i j =
    return $ WrapIntegral $ (DC.coerce m :: LA.Matrix a) `LA.atIndex` (i, j)
  unsafeFreeze = stToPrim . LA.unsafeFreezeMatrix . unwrapSTM
  unsafeThaw m = stToPrim $ STMat (LA.rows m) (LA.cols m) <$> LA.unsafeThawMatrix m
  basicUnsafeGetRowM m i = return $ DC.coerce $ getRowLA  i (DC.coerce m :: LA.Matrix a)
  basicUnsafeGetColumnM m i =
    return $ DC.coerce $ head $ LA.toColumns (DC.coerce m LA.¿ [i] :: LA.Matrix a)
  unsafeGenerate w h f =
    DC.coerce $
    LA.fromRows [ SV.generate w $ \j -> unwrapIntegral $ f i j
                | i <- [0..h - 1]
                ]

instance (Fractional a, LA.Container LA.Matrix a) => Matrix LA.Matrix (WrapFractional a) where
  basicRowCount = LA.rows
  basicColumnCount = LA.cols
  unsafeFromRows rs = DC.coerce $ LA.fromRows (DC.coerce rs :: [LA.Vector a])
  unsafeFromColumns cs = DC.coerce $ LA.fromColumns (DC.coerce cs :: [LA.Vector a])
  basicUnsafeIndexM m i j =
    return $ WrapFractional $ (DC.coerce m :: LA.Matrix a) `LA.atIndex` (i, j)
  unsafeFreeze = stToPrim . LA.unsafeFreezeMatrix . unwrapSTM
  unsafeThaw m = stToPrim $ STMat (LA.rows m) (LA.cols m) <$> LA.unsafeThawMatrix m
  basicUnsafeGetRowM m i = return $ DC.coerce $ getRowLA  i (DC.coerce m :: LA.Matrix a)
  basicUnsafeGetColumnM m i =
    return $ DC.coerce $ head $ LA.toColumns (DC.coerce m LA.¿ [i] :: LA.Matrix a)
  unsafeGenerate w h f =
    DC.coerce $
    LA.fromRows [ SV.generate w $ \j -> unwrapFractional $ f i j
                | i <- [0..h - 1]
                ]

getRowLA :: LA.Element t => Int -> LA.Matrix t -> LA.Vector t
getRowLA i m = head $ LA.toRows $ m LA.? [i]
getColLA :: LA.Element t => Int -> LA.Matrix t -> LA.Vector t
getColLA i m = head $ LA.toColumns $ m LA.¿ [i]

unwrapSTMNum :: STMatrix s (WrapNum a) -> LA.STMatrix s a
unwrapSTMNum = DC.coerce . unwrapSTM
{-# INLINE unwrapSTMNum #-}

castSTMatrixNum :: LA.STMatrix s (WrapNum a) -> LA.STMatrix s a
castSTMatrixNum = DC.coerce
{-# INLINE castSTMatrixNum #-}

instance (Num a, LA.Container LA.Matrix a) => MMatrix STMatrix (WrapNum a) where
  basicUnsafeNew n m = stToPrim $ STMat n m <$> LA.newMatrix 0 n m
  basicInitialise _ = return ()
  basicRowCount = stmRows
  basicColumnCount = stmCols
  unsafeGetRow i m =
    stToPrim $ DC.coerce . getRowLA i <$> LA.unsafeFreezeMatrix (unwrapSTMNum m)
  unsafeGetColumn i m = stToPrim $ DC.coerce . getColLA i <$> LA.unsafeFreezeMatrix (unwrapSTMNum m)
  unsafeFill n m a = stToPrim $ STMat n m <$> LA.newMatrix a n m
  unsafeFromRow r = stToPrim $ STMat 1 (SV.length r) <$> LA.unsafeThawMatrix (LA.asRow r)
  unsafeFromRows rs = stToPrim $ do
    let m = LA.fromRows (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeFromColumn c = stToPrim $ STMat (SV.length c) 1 <$> LA.unsafeThawMatrix (LA.asColumn c)
  unsafeFromColumns rs = stToPrim $ do
    let m = LA.fromColumns (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeCopy (STMat r c m)  (STMat _ _ m') = stToPrim $ do
    m0 <- LA.unsafeFreezeMatrix m'
    LA.setMatrix (DC.coerce m :: LA.STMatrix w a) r c (DC.coerce m0 :: LA.Matrix a)
  unsafeRead (STMat _ _ m) i j = stToPrim $ LA.readMatrix m i j
  unsafeWrite (STMat _ _ m) i j x = stToPrim $ LA.writeMatrix m i j x
  basicSet (STMat r c m) x = stToPrim $ LA.modifyMatrix m r c $ const x
  basicUnsafeSwapRows (STMat _ _ m) i j =
    stToPrim $ LA.rowOper (LA.SWAP i j LA.AllCols) $ castSTMatrixNum m
  unsafeScaleRow (STMat _ _ m) i c =
    stToPrim $ LA.rowOper (LA.SCAL (unwrapNum c) (LA.Row i) LA.AllCols) $ castSTMatrixNum m
  basicUnsafeIMapRowM (STMat _ c m) i f = forM_ [0..c-1] $ \j ->
    stToPrim . LA.writeMatrix m i j =<< f j =<< stToPrim (LA.readMatrix m i j)
  toRows (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toRows <$> LA.unsafeFreezeMatrix (castSTMatrixNum m)
  toColumns (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toColumns <$> LA.unsafeFreezeMatrix (castSTMatrixNum m)

unwrapSTMFractional :: STMatrix s (WrapFractional a) -> LA.STMatrix s a
unwrapSTMFractional = DC.coerce . unwrapSTM
{-# INLINE unwrapSTMFractional #-}

castSTMatrixFractional :: LA.STMatrix s (WrapFractional a) -> LA.STMatrix s a
castSTMatrixFractional = DC.coerce
{-# INLINE castSTMatrixFractional #-}

instance (Fractional a, LA.Container LA.Matrix a) => MMatrix STMatrix (WrapFractional a) where
  basicUnsafeNew n m = stToPrim $ STMat n m <$> LA.newMatrix 0 n m
  basicInitialise _ = return ()
  basicRowCount = stmRows
  basicColumnCount = stmCols
  unsafeGetRow i m =
    stToPrim $ DC.coerce . getRowLA i <$> LA.unsafeFreezeMatrix (unwrapSTMFractional m)
  unsafeGetColumn i m = stToPrim $ DC.coerce . getColLA i <$> LA.unsafeFreezeMatrix (unwrapSTMFractional m)
  unsafeFill n m a = stToPrim $ STMat n m <$> LA.newMatrix a n m
  unsafeFromRow r = stToPrim $ STMat 1 (SV.length r) <$> LA.unsafeThawMatrix (LA.asRow r)
  unsafeFromRows rs = stToPrim $ do
    let m = LA.fromRows (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeFromColumn c = stToPrim $ STMat (SV.length c) 1 <$> LA.unsafeThawMatrix (LA.asColumn c)
  unsafeFromColumns rs = stToPrim $ do
    let m = LA.fromColumns (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeCopy (STMat r c m)  (STMat _ _ m') = stToPrim $ do
    m0 <- LA.unsafeFreezeMatrix m'
    LA.setMatrix (DC.coerce m :: LA.STMatrix w a) r c (DC.coerce m0 :: LA.Matrix a)
  unsafeRead (STMat _ _ m) i j = stToPrim $ LA.readMatrix m i j
  unsafeWrite (STMat _ _ m) i j x = stToPrim $ LA.writeMatrix m i j x
  basicSet (STMat r c m) x = stToPrim $ LA.modifyMatrix m r c $ const x
  basicUnsafeSwapRows (STMat _ _ m) i j =
    stToPrim $ LA.rowOper (LA.SWAP i j LA.AllCols) $ castSTMatrixFractional m
  unsafeScaleRow (STMat _ _ m) i c =
    stToPrim $ LA.rowOper (LA.SCAL (unwrapFractional c) (LA.Row i) LA.AllCols) $ castSTMatrixFractional m
  basicUnsafeIMapRowM (STMat _ c m) i f = forM_ [0..c-1] $ \j ->
    stToPrim . LA.writeMatrix m i j =<< f j =<< stToPrim (LA.readMatrix m i j)
  toRows (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toRows <$> LA.unsafeFreezeMatrix (castSTMatrixFractional m)
  toColumns (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toColumns <$> LA.unsafeFreezeMatrix (castSTMatrixFractional m)

unwrapSTMIntegral :: STMatrix s (WrapIntegral a) -> LA.STMatrix s a
unwrapSTMIntegral = DC.coerce . unwrapSTM
{-# INLINE unwrapSTMIntegral #-}

castSTMatrixIntegral :: LA.STMatrix s (WrapIntegral a) -> LA.STMatrix s a
castSTMatrixIntegral = DC.coerce
{-# INLINE castSTMatrixIntegral #-}

instance (Integral a, LA.Container LA.Matrix a) => MMatrix STMatrix (WrapIntegral a) where
  basicUnsafeNew n m = stToPrim $ STMat n m <$> LA.newMatrix 0 n m
  basicInitialise _ = return ()
  basicRowCount = stmRows
  basicColumnCount = stmCols
  unsafeGetRow i m =
    stToPrim $ DC.coerce . getRowLA i <$> LA.unsafeFreezeMatrix (unwrapSTMIntegral m)
  unsafeGetColumn i m = stToPrim $ DC.coerce . getColLA i <$> LA.unsafeFreezeMatrix (unwrapSTMIntegral m)
  unsafeFill n m a = stToPrim $ STMat n m <$> LA.newMatrix a n m
  unsafeFromRow r = stToPrim $ STMat 1 (SV.length r) <$> LA.unsafeThawMatrix (LA.asRow r)
  unsafeFromRows rs = stToPrim $ do
    let m = LA.fromRows (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeFromColumn c = stToPrim $ STMat (SV.length c) 1 <$> LA.unsafeThawMatrix (LA.asColumn c)
  unsafeFromColumns rs = stToPrim $ do
    let m = LA.fromColumns (DC.coerce rs :: [LA.Vector a])
    STMat (LA.rows m) (LA.cols m) . DC.coerce <$> LA.unsafeThawMatrix m
  unsafeCopy (STMat r c m)  (STMat _ _ m') = stToPrim $ do
    m0 <- LA.unsafeFreezeMatrix m'
    LA.setMatrix (DC.coerce m :: LA.STMatrix w a) r c (DC.coerce m0 :: LA.Matrix a)
  unsafeRead (STMat _ _ m) i j = stToPrim $ LA.readMatrix m i j
  unsafeWrite (STMat _ _ m) i j x = stToPrim $ LA.writeMatrix m i j x
  basicSet (STMat r c m) x = stToPrim $ LA.modifyMatrix m r c $ const x
  basicUnsafeSwapRows (STMat _ _ m) i j =
    stToPrim $ LA.rowOper (LA.SWAP i j LA.AllCols) $ castSTMatrixIntegral m
  unsafeScaleRow (STMat _ _ m) i c =
    stToPrim $ LA.rowOper (LA.SCAL (unwrapIntegral c) (LA.Row i) LA.AllCols) $ castSTMatrixIntegral m
  basicUnsafeIMapRowM (STMat _ c m) i f = forM_ [0..c-1] $ \j ->
    stToPrim . LA.writeMatrix m i j =<< f j =<< stToPrim (LA.readMatrix m i j)
  toRows (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toRows <$> LA.unsafeFreezeMatrix (castSTMatrixIntegral m)
  toColumns (STMat _ _ m) =
    stToPrim $ DC.coerce . LA.toColumns <$> LA.unsafeFreezeMatrix (castSTMatrixIntegral m)

instance {-# OVERLAPPING #-} (LA.Container LA.Matrix a, Show a) => Show (LA.Matrix (WrapNum a)) where
  showsPrec d = showsPrec d . (DC.coerce :: LA.Matrix (WrapNum a) -> LA.Matrix a)

instance {-# OVERLAPPING #-} (LA.Container LA.Matrix a, Show a) => Show (LA.Matrix (WrapIntegral a)) where
  showsPrec d = showsPrec d . (DC.coerce :: LA.Matrix (WrapIntegral a) -> LA.Matrix a)

instance {-# OVERLAPPING #-} (LA.Container LA.Matrix a, Show a) => Show (LA.Matrix (WrapFractional a)) where
  showsPrec d = showsPrec d . (DC.coerce :: LA.Matrix (WrapFractional a) -> LA.Matrix a)


module STCR2Z1T0ImageSparseCode where

import           Control.Monad                as M
import           Control.Monad.Parallel       as MP
import           Data.Array.Repa              as R
import           Data.Binary                  (decodeFile)
import           Data.Complex
import           Data.Conduit
import           Data.Conduit.List            as CL
import           Data.List                    as L
import           Data.Maybe
import           Data.Vector.Generic          as VG
import           Data.Vector.Storable         as VS
import           DFT.Plan
import           Filter.Gaussian
import           Filter.Pinwheel
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           Image.Transform
import           PVPFile.IO
import           STC
import           STC.OrientationScaleAnalysis
import           System.Directory
import           System.Environment
import           System.FilePath
import           System.Random
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:sigmaStr:taoStr:lenStr:initialScaleStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:histFilePath:alphaStr:pinwheelFlagStr:imagePath:reconPath:numIterationStr:writeSourceFlagStr:thresholdStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      sigma = read sigmaStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      initialScale = read initialScaleStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      writeSourceFlag = read writeSourceFlagStr :: Bool
      threshold = read thresholdStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCR2Z1T0ImageSparseCode"
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  imgHeader <- readPVPHeader imagePath
  reconHeader <- readPVPHeader reconPath
  let m = min (nBands imgHeader) (nBands reconHeader)
  print m
  xs <-
    runConduitRes $
    pvpFileSource imagePath .| (CL.drop (m - 1) >> CL.map id) .| CL.head
  ys <-
    runConduitRes $
    pvpFileSource reconPath .| (CL.drop (m - 1) >> CL.map id) .| CL.head
  let img' = pvpOutputData2Array . fromJust $ xs
      recon = pvpOutputData2Array . fromJust $ ys
      (Z :. cols :. rows :. _) = extent img'
  print . extent $ img'
  plotImageRepa
    (folderPath </> "recon.png")
    (ImageRepa 8 . computeS . rotate3DR $ recon)
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z1T0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          sigma
          tao
          1
          theta0Freqs
          thetaFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length theta0Freqs
             , L.length thetaFreqs
             ]
             0)
  arrR2Z1T0 <-
    computeUnboxedP $
    computeR2Z1T0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      thetaFreqs
      theta0Freqs
  -- let arr3d =
  --       rotate3D . R.slice arrR2Z1T0 $
  --       (Z :. All :. (L.length theta0Freqs - 1) :. All :. All)
  -- MP.mapM_
  --   (\i ->
  --      plotImageRepaComplex
  --        (folderPath </> "GreensR2Z1T0_" L.++ show (i + 1) L.++ ".png") .
  --      ImageRepa 8 .
  --      computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr3d $
  --      (Z :. All :. All :. i))
  --   [0 .. (L.length thetaFreqs) - 1]
  plan <- makeR2Z1T0Plan emptyPlan arrR2Z1T0
  img'' <-
    fmap (fromUnboxed (Z :. cols :. rows)) .
    dropPixel (4 / 5) . toUnboxed . R.sumS $
    img'
  let s = R.sumAllS img'
      img = R.map (/ s) . R.sumS $ img'
      initialDistSource =
        R.traverse
          img
          (\(Z :. nx :. ny) -> (Z :. (L.length theta0Freqs) :. nx :. ny))
          (\f (Z :. _ :. i :. j) -> f (Z :. i :. j) :+ 0)
  undefined
  -- powerMethod1
  --   plan
  --   folderPath
  --   cols
  --   rows
  --   numOrientation
  --   thetaFreqs
  --   theta0Freqs
  --   arrR2Z1T0
  --   numIteration
  --   writeSourceFlag
  --   ""
  --   threshold
  --   initialDistSource


{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE QuasiQuotes     #-}
{-# LANGUAGE TemplateHaskell #-}
module Main where

import           Codec.Compression.GZip       ( decompress )
import           Control.Monad
import qualified Data.Array.Repa              as R
import qualified Data.Array.Repa.Stencil      as Stencil
import qualified Data.Array.Repa.Stencil.Dim2 as Stencil
import qualified Data.Massiv.Array            as A
import qualified Data.ByteString              as BS
import qualified Data.ByteString.Lazy         as BL
import           Data.List                    as L
import qualified Numeric.LinearAlgebra        as LA
import           Prelude                      as P
import           System.Directory             ( getCurrentDirectory )
import           System.Environment
import           System.IO
import           System.Random                as Rand
import           Control.DeepSeq              as DS

import           Foreign.Ptr                  ( Ptr )
import           System.IO.Unsafe             ( unsafePerformIO )
import qualified Data.Vector.Storable as U    ( unsafeToForeignPtr0, unsafeFromForeignPtr0 )
import qualified Numeric.LinearAlgebra.Devel as U
import           Foreign                      ( mallocForeignPtrArray, withForeignPtr )

main :: IO ()
main = do
    args <- getArgs

    imgs <- loadData
    kernels <- initKernels
    case args of
      "--massiv":_  -> do
          data_ <- return . BL.toStrict =<< getData "train-images-idx3-ubyte.gz"
          let imgs_ :: A.Array A.U A.Ix3 Double
              imgs_ = A.fromLists' A.Par $
                  (getImage_ data_) <$> [0..49999]
              {-# INLINE imgs_ #-}
              kernels_ :: A.Array A.U A.Ix3 Double
              kernels_ = A.computeAs A.U $
                  A.uniformRangeArray (mkStdGen 42) (0, 1) (A.ParN 3) (A.Sz3 20 5 5)
                  -- TODO: convert to standard normal
                  -- distribution using Box-Mueller transform
          runMassiv imgs_ kernels_
      "--hmatrix":_ -> runHmatrix imgs kernels
      "--cublas":_  -> runCublas
      "--repa":_    -> runRepa imgs kernels
      "--cuda":_    -> runCuda
      _             -> putStrLn "Error"

runMassiv :: A.Array A.U A.Ix3 Double -> A.Array A.U A.Ix3 Double -> IO ()
runMassiv imgs kernels = do
--    let stencil = A.makeConvolutionStencilFromKernel . (ks !!)
--        {-# INLINE stencil #-}
--
--        base = A.computeAs A.U $ A.applyStencil A.noPadding (stenc 0) x
--        conv_ = ((\i -> (\s -> A.computeAs A.U $ A.applyStencil A.noPadding s i) <$> stencils) <$> is)
--        conv = (foldl1' (\p ch -> A.computeAs A.U $ A.append' 3 p ch) conv) :: A.Array A.U A.Ix3 Double
    imgs `deepseq` putStrLn "Done"

-- loadDataMassiv :: IO (A.Array A.U A.Ix3 Double)
-- loadDataMassiv = do
--     trainImgs <- getData "train-images-idx3-ubyte.gz"
--     let imgs = BL.toStrict trainImgs
--         imgArr :: A.Array A.U A.Ix3 Double
--         imgArr = foldl1 (\acc i -> A.computeAs A.U $ A.append' 3 acc i) [getImage_ n imgs | n <- [0..49999]]
--     return imgArr

getImage_ :: BS.ByteString -> Int -> [[Double]]
getImage_ imgs n = chunkList 28 [normalize $ BS.index imgs (16 + n*784 + s) | s <- [0..783]]

-- loadDataMassiv :: IO (A.Array A.U A.Ix3 Double)
-- loadDataMassiv = do
--     trainImgs <- getData "train-images-idx3-ubyte.gz"
--     let imgs = BL.toStrict trainImgs
--         imgArr :: A.Array A.U A.Ix3 Double
--         imgArr = foldl1 (\acc i -> A.computeAs A.U $ A.append' 3 acc i) [getImage_ n imgs | n <- [0..49999]]
--     return imgArr

-- getImage_ :: Int -> BS.ByteString -> A.Array A.U A.Ix3 Double
-- getImage_ n imgs = A.resize' (A.Sz (1 A.:> 28 A.:. 28)) $
--     A.fromList A.Par [normalize $ BS.index imgs (16 + n*784 + s) | s <- [0..783]]

-- time:   1.54s user
-- space: ~2GB?
runHmatrix :: [LA.Matrix Double] -> [LA.Matrix Double] -> IO ()
runHmatrix imgs kernels =
    let kcols = LA.fromRows $ P.map LA.flatten kernels
        icols = (LA.tr . im2col 5 5 1 1) <$> imgs
        res = P.map (kcols LA.<>) icols
     in res `deepseq` putStrLn $ "Done, size: " P.++ (show . length $ res)

    --let dCircKernel = doublyBlockedCirculant . zeroPad (28, 28) . flipKernel $ k
    --sequence_ $
    --    (LA.dispShort 10 10 3) <$> [dCircKernel LA.<> img | img <- imgs]

-- time:   47.93s user
-- space: ~5GB
runRepa :: [LA.Matrix Double] -> [LA.Matrix Double] -> IO ()
runRepa imgs kernels = print "Done"
--    let ks = hmatrix2repa <$> kernels
--        is = hmatrix2repa <$> imgs
--
--        imgMat = head is
--        kernel = [stencil2| 10 10 10 10 10
--                            10 10 10 10 10
--                            10 10 10 10 10
--                            10 10 10 10 10
--                            10 10 10 10 10 |] :: Stencil.Stencil DIM2 Double
--    res <- forM is $ \i ->
--        forM (replicate 20 kernel) $ \k ->
--            computeP ((mapStencil2 (BoundConst 0) k i) :: Array PC5 DIM2 Double)
--    print (res :: [[Array U DIM2 Double]])

runCublas = print ""

runCuda = print ""

--convolveP :: Array U DIM2 Double  -- ^ 28x28    Input image       (Unboxed Vector)
--          -> Array C DIM3 Double  -- ^ 20x5x5   Kernel tensor     (Cursored Array)
--          -> Array U DIM3 Double  -- ^ 20x24x24 Convolution layer (Unboxed Vector)

{- Kernel stuff -}
initKernels :: IO [LA.Matrix Double]
initKernels = replicateM 20 $ LA.randn 5 5

flipKernel :: LA.Matrix Double -> LA.Matrix Double
flipKernel = LA.flipud

doublyBlockedCirculant :: LA.Matrix Double  -- ^ input kernel
                       -> LA.Matrix Double
doublyBlockedCirculant m = LA.fromBlocks . circ $ circOfRows m

circOfRows :: LA.Matrix Double -> [LA.Matrix Double]
circOfRows m = [ LA.fromLists $ circ row | row <- rows ]
    where rows = reverse $ LA.toLists m

circ :: [a] -> [[a]]
circ xs = take (length xs) $ iterate (rotate 1) xs -- `(length xs)` is wrong here..

rotate :: Int -> [a] -> [a]
rotate n xs = drop diff xs <> take diff xs
    where diff = length xs - n

zeroPad :: (Int, Int) -> LA.Matrix Double -> LA.Matrix Double
zeroPad (r,c) m = LA.fromBlocks $
    [[LA.konst 0 (r-r', 1), LA.row $ replicate (c-c') 0],
     [m, LA.row $ replicate (c-c') 0]]
         where (r', c') = LA.size m

zeroPadTo :: Int -> Int -> [LA.Matrix Double] -> [LA.Matrix Double]
zeroPadTo r c kernels = [ LA.fromBlocks [[LA.konst 0 (r', 1), LA.row $ replicate r' 0], [k, LA.row $ replicate r' 0]] | k <- kernels]
    where (r', c') = (r-5, c-5)

getSizeZP :: Int -> Int -> (Int, Int)
getSizeZP m n = ((n-m+1)^2, n^2)

{- Load data -}
loadData :: IO [LA.Matrix Double]
loadData = do
    trainImgs <- getData "train-images-idx3-ubyte.gz"
    let imgs = BL.toStrict trainImgs
        trainData = chunkList 10 [getImage n imgs | n <- [0..49999]]
    return $ LA.fromColumns <$> trainData

getImage :: Int -> BS.ByteString -> LA.Vector Double
getImage n imgs = LA.fromList [normalize $ BS.index imgs (16 + n*784 + s) | s <- [0..783]]

getData :: FilePath -> IO BL.ByteString
getData path = do
    currentDir <- getCurrentDirectory
    fileData <- decompress <$> BL.readFile (currentDir P.<> "/mnist_dataset/" P.<> path)
    return fileData

chunkList :: Int -> [a] -> [[a]]
chunkList n xs = takeWhile (not.null) $ unfoldr (Just . splitAt n) xs

normalize :: (Integral a, Floating b) => a -> b
normalize x = fromIntegral x/255

hmatrix2repa :: LA.Matrix Double -> R.Array R.U R.DIM2 Double
hmatrix2repa mat = R.fromListUnboxed (R.Z R.:. (r::Int) R.:. (c::Int)) . LA.toList . LA.flatten $ mat
    where (r, c) = LA.size mat

foreign import ccall unsafe im2col_cpu :: Ptr Double -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr Double -> IO ()

im2col_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> LA.Matrix Double -> LA.Matrix Double
im2col_c channels height width kernelRows kernelColumns strideRows strideColumns dataImg =
    let vec = LA.flatten dataImg
        rowOut = (height - kernelRows) `div` strideRows + 1
        colOut = (width - kernelColumns) `div` strideColumns + 1
        kernelSize = kernelRows * kernelColumns
        numberOfPatches = rowOut * colOut
    in unsafePerformIO $ do
        outPtr <- mallocForeignPtrArray (numberOfPatches * kernelSize * channels)
        let (inPtr, _) = U.unsafeToForeignPtr0 vec

        withForeignPtr inPtr $ \inPtr' ->
            withForeignPtr outPtr $ \outPtr' ->
                im2col_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns outPtr'

        let matVec = U.unsafeFromForeignPtr0 outPtr (numberOfPatches * kernelSize * channels)
        return $ U.matrixFromVector U.RowMajor numberOfPatches (kernelSize * channels) matVec

im2col :: Int -> Int -> Int -> Int -> LA.Matrix Double -> LA.Matrix Double
im2col kernelRows kernelColumns strideRows strideColumns dataIm =
  let channels = 1
      height = LA.rows dataIm
      width  = LA.cols dataIm
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns dataIm


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE BangPatterns #-}

module Numeric.MCMC.Langevin (
          MarkovChain(..), Options(..)
        , runChain
        ) where

import Control.Monad
import Control.Monad.Trans
import Control.Monad.Reader
import Control.Monad.Primitive
import Control.Arrow
import System.Random.MWC
import System.Random.MWC.Distributions
import Data.List
import Statistics.Distribution
import Statistics.Distribution.Normal hiding (standard)
import GHC.Float

-- | State of the Markov chain.  Current parameter values are held in 'theta',
--   while accepts counts the number of proposals accepted.
data MarkovChain = MarkovChain { theta   :: [Double]
                               , accepts :: {-# UNPACK #-} !Int }

-- | Options for the chain.  The target (expected to be a log density), its
--   gradient, and a step size tuning parameter.
data Options = Options { _target  :: [Double] -> Double
                       , _gTarget :: [Double] -> [Double]
                       , _eps     :: {-# UNPACK #-} !Double }

-- | A result with this type has a view of the chain options.
type ViewsOptions = ReaderT Options

-- | Display the current state.
instance Show MarkovChain where
    show config = filter (`notElem` "[]") $ show (map double2Float (theta config))

-- | Density function for an isotropic Gaussian.  The (identity) covariance
--   matrix is multiplied by the scalar 'sig'.
isoGauss :: [Double] -> [Double] -> Double -> Double
isoGauss xs mu sig = foldl1' (*) (zipWith density nds xs)
    where nds = map (`normalDistr` sig) mu
{-# INLINE isoGauss #-}

-- | Mean function for the discretized Langevin diffusion.
localMean :: Monad m
          => [Double]                -- Current state
          -> ViewsOptions m [Double] -- Localized mean
localMean t = do
    Options _ gTarget e <- ask
    return $! zipWith (+) t (map (* (0.5 * e^(2 :: Int))) (gTarget t))
{-# INLINE localMean #-}

-- | Perturb the state, creating a new proposal.
perturb :: PrimMonad m
        => [Double]                   -- Current state
        -> Gen (PrimState m)          -- MWC PRNG
        -> ViewsOptions m [Double]    -- Resulting perturbation.
perturb t g = do
    Options _ _ e <- ask
    zs    <- replicateM (length t) (lift $ standard g)
    t0    <- localMean t
    let perturbedState = zipWith (+) t0 t1
        t1 = map (* e) zs
    return $! perturbedState
{-# INLINE perturb #-}

-- | Perform a Metropolis accept/reject step.
metropolisStep :: PrimMonad m
               => MarkovChain                -- Current state
               -> Gen (PrimState m)          -- MWC PRNG
               -> ViewsOptions m MarkovChain -- New state
metropolisStep state g = do
    Options target _ e <- ask
    let (t0, nacc) = (theta &&& accepts) state
    zc           <- lift $ uniformR (0, 1) g
    proposal     <- perturb t0 g
    t0Mean       <- localMean t0
    proposalMean <- localMean proposal
    let mc = if   zc < acceptProb
             then (proposal, 1)
             else (t0,       0)

        acceptProb = if isNaN val then 0 else val where val = arRatio

        arRatio = exp . min 0 $
            target proposal + log (isoGauss proposal t0Mean (e^(2 :: Int)))
          - target t0       - log (isoGauss t0 proposalMean (e^(2 :: Int)))

    return $! MarkovChain (fst mc) (nacc + snd mc)
{-# INLINE metropolisStep #-}

-- | Diffuse through states.
runChain :: Options         -- Options of the Markov chain.
         -> Int             -- Number of epochs to iterate the chain.
         -> Int             -- Print every nth iteration
         -> MarkovChain     -- Initial state of the Markov chain.
         -> Gen RealWorld   -- MWC PRNG
         -> IO MarkovChain  -- End state of the Markov chain, wrapped in IO.
runChain = go
  where go o n t !c g | n == 0 = return c
                      | n `rem` t /= 0 = do
                            r <- runReaderT (metropolisStep c g) o
                            go o (n - 1) t r g
                      | otherwise = do
                            r <- runReaderT (metropolisStep c g) o
                            print r
                            go o (n - 1) t r g
{-# INLINE runChain #-}


{-# LANGUAGE PatternGuards #-}
-- |
-- Module    : Statistics.Matrix
-- Copyright : 2011 Aleksey Khudyakov, 2014 Bryan O'Sullivan
-- License   : BSD3
--
-- Basic matrix operations.
--
-- There isn't a widely used matrix package for Haskell yet, so
-- we implement the necessary minimum here.

module Statistics.Matrix
    ( -- * Data types
      Matrix(..)
    , Vector
      -- * Conversion from/to lists/vectors
    , fromVector
    , dimension
    -- , center
    , multiplyV
    , transpose
    , norm
    , column
    -- , row
    , for
    , unsafeIndex
    ) where

import Prelude hiding (exponent, map, sum)
import qualified Data.Vector.Unboxed as U

import Statistics.Function (for, square)
import Statistics.Matrix.Types
import Statistics.Sample.Internal (sum)

----------------------------------------------------------------
-- Conversion to/from vectors/lists
----------------------------------------------------------------

-- | Convert from a row-major vector.
fromVector :: Int               -- ^ Number of rows.
           -> Int               -- ^ Number of columns.
           -> U.Vector Double   -- ^ Flat list of values, in row-major order.
           -> Matrix
fromVector r c v
  | r*c /= len = error "input size mismatch"
  | otherwise  = Matrix r c 0 v
  where len    = U.length v

----------------------------------------------------------------
-- Other
----------------------------------------------------------------

-- | Return the dimensions of this matrix, as a (row,column) pair.
dimension :: Matrix -> (Int, Int)
dimension (Matrix r c _ _) = (r, c)

-- | Matrix-vector multiplication.
multiplyV :: Matrix -> Vector -> Vector
multiplyV m v
  | cols m == c = U.generate (rows m) (sum . U.zipWith (*) v . row m)
  | otherwise   = error $ "matrix/vector unconformable " ++ show (cols m,c)
  where c = U.length v

-- | Calculate the Euclidean norm of a vector.
norm :: Vector -> Double
norm = sqrt . sum . U.map square

-- | Return the given column.
column :: Matrix -> Int -> Vector
column (Matrix r c _ v) i = U.backpermute v $ U.enumFromStepN i c r
{-# INLINE column #-}

-- | Return the given row.
row :: Matrix -> Int -> Vector
row (Matrix _ c _ v) i = U.slice (c*i) c v

unsafeIndex :: Matrix
            -> Int              -- ^ Row.
            -> Int              -- ^ Column.
            -> Double
unsafeIndex = unsafeBounds U.unsafeIndex

-- | Given row and column numbers, calculate the offset into the flat
-- row-major vector, without checking.
unsafeBounds :: (Vector -> Int -> r) -> Matrix -> Int -> Int -> r
unsafeBounds k (Matrix _ cs _ v) r c = k v $! r * cs + c
{-# INLINE unsafeBounds #-}

transpose :: Matrix -> Matrix
transpose m@(Matrix r0 c0 e _) = Matrix c0 r0 e . U.generate (r0*c0) $ \i ->
  let (r,c) = i `quotRem` r0
  in unsafeIndex m c r


{-# LANGUAGE ScopedTypeVariables #-}

module Main (main) where

import Data.Complex
import Data.Monoid ((<>))
import Text.PrettyPrint.Mainland
import Control.Monad.IO.Class (liftIO)

import Spiral
import Spiral.Exp
import Spiral.FFT.CooleyTukey
import Spiral.SPL
import Spiral.Util.Pretty.Maple

main :: IO ()
main = defaultMain $ \args -> do
    n <- case args of
           [s] -> return (read s)
           _   -> return 4
    let dft_n :: SPL (Exp (Complex Double))
        dft_n = dit n
    --liftIO $ check dft_n
    liftIO $ putDocLn $ text "with(LinearAlgebra);"
    liftIO $ putDocLn $ text "A :=" <+> ppr (DFT n :: SPL (Exp (Complex Double))) <> semi
    liftIO $ putDocLn $ text "B :=" <+> ppr dft_n <> semi


{-# OPTIONS_GHC -Wall            #-}

{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE OverloadedLists     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Heat2D (main, a1, a2) where

import Data.Number.Symbolic
import Data.Proxy

import qualified Naperian as N
import qualified GHC.TypeLits as M
import           Data.Functor
import           System.IO

import           Numeric.Sundials
import           Numeric.LinearAlgebra

kx, ky :: Floating a => a
kx = 0.5
ky = 0.75

-- x mesh spacing
-- y mesh spacing
dx :: Floating a => a
dx = 1 / (fromIntegral nx - 1)

dy :: Floating a => a
dy = 1 / (fromIntegral ny - 1)

c1, c2 :: Floating a => a
c1 = kx/dx/dx
c2 = ky/dy/dy

preA2 :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
preA2 = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f p q i j k l | i == 0               = 0
                        | j == 0               = 0
                        | i == p - 1           = 0
                        | j == q - 1           = 0
                        | k == i - 1 && l == j = 1
                        | k == i     && l == j = -2
                        | k == i + 1 && l == j = 1
                        | otherwise            = 0

a2Num :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
a2Num = N.binary (*) (N.Scalar c2) preA2

a2 :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
          N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
a2 = N.binary (*) (N.Scalar $ var "a") preA2

preA1 :: forall b m n . (M.KnownNat m, M.KnownNat n, Num b) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] b
preA1 = N.Prism $ N.Prism $ N.Prism $ N.Prism $ N.Scalar $
      N.viota @m <&> (\(N.Fin x) ->
      N.viota @n <&> (\(N.Fin w) ->
      N.viota @m <&> (\(N.Fin v) ->
      N.viota @n <&> (\(N.Fin u) ->
      (f m n x w v u)))))
        where
          m = fromIntegral $ M.natVal (undefined :: Proxy m)
          n = fromIntegral $ M.natVal (undefined :: Proxy n)
          f :: Int -> Int -> Int -> Int -> Int -> Int -> b
          f p q i j k l | i == 0                   = 0
                        | j == 0                   = 0
                        | i == p - 1               = 0
                        | j == q - 1               = 0
                        | k == i     && l == j - 1 = 1
                        | k == i     && l == j     = -2
                        | k == i     && l == j + 1 = 1
                        | otherwise                = 0

a1Num :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a) =>
        N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] a
a1Num = N.binary (*) (N.Scalar c1) preA1

a1 :: forall a m n . (M.KnownNat m, M.KnownNat n, Floating a, Eq a) =>
           N.Hyper '[N.Vector n, N.Vector m, N.Vector n, N.Vector m] (Sym a)
a1 = N.binary (*) (N.Scalar $ var "b") preA1

h :: forall m n a . (M.KnownNat m, M.KnownNat n, Floating a) =>
                     N.Hyper '[N.Vector m, N.Vector n] a
h = N.Prism $ N.Prism $ N.Scalar $
     N.viota @n <&> (\(N.Fin x) ->
     N.viota @m <&> (\(N.Fin w) ->
     sin (pi * (fromIntegral w) * dx)
             * sin (2 * pi * (fromIntegral x) * dy)))

-- spatial mesh size
nx, ny :: Int
nx = 30
ny = 60

bigA :: Matrix Double
bigA = fromLists $
       fmap (N.elements . N.Prism . N.Prism . N.Scalar) $
       N.elements $ N.crystal $ N.crystal $
       N.binary (+) (a2Num @Double @60 @30) (a1Num @Double @60 @30)

b :: Vector Double
b = fromList $ N.elements (h @30 @60 @Double)

t0, tf :: Double
t0 = 0.0
tf = 0.3

bigNt :: Int
bigNt = 20

dTout :: Double
dTout = (tf - t0) / (fromIntegral bigNt)

ts :: [Double]
ts = map (dTout *) $ map fromIntegral [1..bigNt]

sol :: Matrix Double
-- sol = odeSolveV SDIRK_5_3_4' Nothing 1.0e-5 1.0e-10 (const bigU') (assoc (nx * ny) 0.0 [] :: Vector Double) (fromList $ ts)
--   where
--     bigU' bigU = bigA #> bigU + b
sol = undefined

main :: IO ()
main = do
  h1 <- openFile "heat2e.000.txt" WriteMode
  mapM_ (hPutStrLn h1) $ map (concatMap (' ':)) $ map (map show) $ toLists sol
  hClose h1
  mapM_ (\i -> putStrLn $ show $ sqrt $ (sol!i) <.> (sol!i) / (fromIntegral nx) / (fromIntegral ny)) ([0 .. length ts - 1] :: [Int])


{-# LANGUAGE ForeignFunctionInterface #-}
module Grenade.Layers.Internal.Convolution (
    im2col
  , col2im
  , col2vid
  , vid2col
  ) where

import qualified Data.Vector.Storable        as U (unsafeFromForeignPtr0,
                                                   unsafeToForeignPtr0)

import           Foreign                     (mallocForeignPtrArray, withForeignPtr)
import           Foreign.Ptr                 (Ptr)
import           Numeric.LinearAlgebra       (Matrix, cols, flatten, rows)
import qualified Numeric.LinearAlgebra.Devel as U
import           System.IO.Unsafe            (unsafePerformIO)

import           Grenade.Types

col2vid :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2vid kernelRows kernelColumns strideRows strideColumns height width dataCol =
  let channels = cols dataCol `div` (kernelRows * kernelColumns)
  in  col2im_c channels height width kernelRows kernelColumns strideRows strideColumns dataCol

col2im :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im kernelRows kernelColumns strideRows strideColumns height width dataCol =
  let channels = 1
  in  col2im_c channels height width kernelRows kernelColumns strideRows strideColumns dataCol

col2im_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im_c channels height width kernelRows kernelColumns strideRows strideColumns dataCol =
  let vec = flatten dataCol
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (height * width * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        col2im_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (height * width * channels)
    return $ U.matrixFromVector U.RowMajor (height * channels) width matVec

foreign import ccall unsafe
    col2im_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()

vid2col :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
vid2col kernelRows kernelColumns strideRows strideColumns height width dataVid =
  let channels = rows dataVid `div` height
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns dataVid

im2col :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2col kernelRows kernelColumns strideRows strideColumns dataIm =
  let channels = 1
      height = rows dataIm
      width  = cols dataIm
  in  im2col_c channels height width kernelRows kernelColumns strideRows strideColumns dataIm

im2col_c :: Int -> Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2col_c channels height width kernelRows kernelColumns strideRows strideColumns dataIm =
  let vec    = flatten dataIm
      rowOut = (height - kernelRows) `div` strideRows + 1
      colOut = (width - kernelColumns) `div` strideColumns + 1
      kernelSize      = kernelRows * kernelColumns
      numberOfPatches = rowOut * colOut
  in unsafePerformIO $ do
    outPtr <- mallocForeignPtrArray (numberOfPatches * kernelSize * channels)
    let (inPtr, _) = U.unsafeToForeignPtr0 vec

    withForeignPtr inPtr $ \inPtr' ->
      withForeignPtr outPtr $ \outPtr' ->
        im2col_cpu inPtr' channels height width kernelRows kernelColumns strideRows strideColumns outPtr'

    let matVec = U.unsafeFromForeignPtr0 outPtr (numberOfPatches * kernelSize * channels)
    return $ U.matrixFromVector U.RowMajor numberOfPatches (kernelSize * channels) matVec

foreign import ccall unsafe
    im2col_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> Ptr RealNum -> IO ()