text:
module ForSyDe.Shallow.Example.Synchronous.Equalizer.Test (
  -- * Test data
  -- ** Test signals
  zeros, ones, increm, modCos,
  bassUp, bassDn, trebleUp, trebleDn,
  overrides, distorsionFlags,
  -- ** Test filter coefficients
  lpCoeff, bpCoeff, hpCoeff,
  lpCoeffD, bpCoeffD, hpCoeffD,
  -- ** Test vectors
  testV1, testV2, testV3,
  -- * Test functions
  -- ** Signal processing blocks
  testDFT, testFFT, testDftAgainstFft, testFIR,
  audioFilterI, audioFilterD,
  -- ** Equalizer modules
  testAudioFilterD,
  testButtonControl,
  testDistortionControl,
  testAnalyzer,
  testEqualizer,
  -- * Test scripts
  testFilterScript,
  testAnalyzerScript,
  testFilterAndAnalyzerScript
  )where

import Data.Complex
import System.IO
import System.FilePath
import ForSyDe.Shallow
import ForSyDe.Shallow.Utility.FIR
import ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioAnalyzer
import ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter
import ForSyDe.Shallow.Example.Synchronous.Equalizer.ButtonControl
import ForSyDe.Shallow.Example.Synchronous.Equalizer.DistortionControl
import ForSyDe.Shallow.Example.Synchronous.Equalizer.Equalizer
import ForSyDe.Shallow.Example.Synchronous.Equalizer.EqualizerTypes

---------------
-- TEST DATA --
---------------

-- | test signal for 'testButtonControl'
bassUp = signal [
  Prst Active, Prst Active, Abst, Abst,
  Prst Active, Prst Active, Abst, Abst,
  Prst Active, Prst Active, Abst, Abst,
  Prst Active, Prst Active, Abst, Abst ]

-- | test signal for 'testButtonControl'
bassDn = signal [
  Abst, Abst, Prst Active, Abst,
  Abst, Abst, Prst Active, Abst,
  Abst, Abst, Prst Active, Abst,
  Abst, Abst, Prst Active, Abst ]

-- | test signal for 'testButtonControl'
trebleUp = signal [
  Abst, Abst, Abst, Abst,
  Abst, Abst, Abst, Prst Active,
  Abst, Abst, Abst, Abst,
  Abst, Abst, Abst, Prst Active ]

-- | test signal for 'testButtonControl'
trebleDn = signal [
  Abst, Abst, Abst, Prst Active,
  Abst, Abst, Abst, Prst Active,
  Abst, Abst, Abst,  Prst Active,
  Abst, Abst, Abst, Prst Active ]

-- | test signal for 'testButtonControl'
overrides = signal [
  Abst, Abst, Abst, Abst,
  Prst Lock, Abst, Abst, Abst,
  Abst, Prst CutBass, Abst, Abst,
  Prst Release,  Abst, Abst, Abst ]

-- | example signal for testing
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.DistorsionControl.distortionControl'
distorsionFlags = signal [
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Fail,
  Abst, Abst, Abst, Prst Pass,
  Abst, Abst, Abst, Prst Pass ]

-- | Low-pass filter coefficients
lpCoeff = vector [
  0.01392741661548,
  0.01396895728902,
  0.01399870011280,
  0.01401657422649,
  0.01402253700635,
  0.01401657422649,
  0.01399870011280,
  0.01396895728902,
  0.01392741661548 ]

-- | Low-pass filter coefficients
lpCoeffD = vector [
   0.03898579822345,
   0.09739504968381,
   0.15360490491115,
   0.19416166962179,
   0.20893350067585,
   0.19416166962179,
   0.15360490491115,
   0.09739504968381,
   0.03898579822345
   ]

-- | Band-pass filter coefficients
bpCoeff = vector [
  0.06318761339784,
  0.08131651217682,
  0.09562326700432,
  0.10478344432968,
  0.10793629404886,
  0.10478344432968,
  0.09562326700432,
  0.08131651217682,
  0.06318761339784 ]

-- | Band-pass filter coefficients
bpCoeffD = vector [
  -0.07845593083988,
   0.00000000000000,
  -0.30707118658796,
   0.00000000000000,
   0.58268794919522,
   0.00000000000000,
  -0.30707118658796,
   0.00000000000000,
  -0.07845593083988
   ]

-- | High-pass filter coefficients
hpCoeff = vector [
  -0.07883878579454,
  -0.09820015927379,
  -0.11354603917221,
  -0.12339860164118,
   0.87320570334018,
  -0.12339860164118,
  -0.11354603917221,
  -0.09820015927379,
  -0.07883878579454 ]

-- | High-pass filter coefficients
hpCoeffD = vector [
   0.03898579822345,
  -0.09739504968381,
   0.15360490491115,
  -0.19416166962179,
   0.20893350067585,
  -0.19416166962179,
   0.15360490491115,
  -0.09739504968381,
   0.03898579822345
  ]

-- | Infinite signal of zeroes
zeros = infiniteS id 0.0

-- | Finite signal of ones, having its length given as input
ones pts = takeS pts $ infiniteS id 1.0

-- | Incremental signal of 400 samples
increm = takeS 400 (infiniteS (+1/200) 0.0)

-- | Additive sinusoidal signal of 400 samples, with 3 sine components.
modCos = mapSY modCosF increm
  where
    modCosF x = cos10 x + cos50 x + cos90 x
    cos10 x = cos (2*pi*10*x)
    cos50 x = cos (2*pi*50*x)
    cos90 x = cos (2*pi*90*x)

-- | example input vector to test the ForSyDe implementation of FFT and DFT
testV1 = mapV toComplex (vector [1, 2])

-- | example input vector to test the ForSyDe implementation of FFT and DFT
testV2 = mapV toComplex (vector [1, 2,3,4])

-- | example input vector to test the ForSyDe implementation of FFT and DFT
testV3 = mapV toComplex (vector [1, 2, 3, 4, 5, 6, 7, 8])

--------------------
-- TEST FUNCTIONS --
--------------------

-- | tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.ButtonControl.buttonControl'
-- function.
testButtonControl = buttonControl overrides bassUp bassDn trebleUp trebleDn
--testButtonInterface = buttonInterface  bassUp bassDn trebleUp trebleDn

toComplex a = a:+0

-- | tests ForSyDe's 'ForSyDe.Shallow.dft' against 'testV3'
testDFT = dft 8 testV3

-- | tests ForSyDe's 'ForSyDe.Shallow.fft' against 'testV3'
testFFT = fft 8 testV3

-- | tests both 'ForSyDe.Shallow.dft' and 'ForSyDe.Shallow.fft'
testDftAgainstFft = zipWithV (-) testDFT testFFT

-- | tests ForSyDe's 'ForSyDe.Shallow.firSY' with 4 coefficients
-- against an impulse signal.
testFIR = firSY coeff s
  where
    coeff = vector [0.1, -0.2, 0.5, 0.2]
    s = signal [1.0, 0, 0, 0, 0]

pts = 4

-- | tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.DistorsionControl.distortionControl'
-- function against 'distorsionFlags'.
testDistortionControl = distortionControl distorsionFlags

-- | Tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioAnalyzer'
-- against a 'audioFilter'ed dummy signal.
testAnalyzer = audioAnalyzer 2 (audioFilter lpCoeff bpCoeff hpCoeff bass treble audioIn)
  where audioIn = ones (pts * 4)
        bass    = zeros
        treble  = zeros

-- | Instantiates
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter'.
audioFilterI = audioFilter lpCoeff bpCoeff hpCoeff zeros zeros

-- | Instantiates
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter'.
audioFilterD = audioFilter lpCoeffD bpCoeffD hpCoeffD zeros zeros

-- | Tests
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter'
-- against a test signal.
testAudioFilterD = audioFilter lpCoeffD bpCoeffD hpCoeffD zeros zeros s
  where s = signal [0.1,0.05,0.3,0.2,-0.5,0.2,0.1,0.3,0.1,-0.1,0.0,0.2]

-- | Tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.Equalizer.equalizer'
-- function against a dummy signal consisting of all ones.
testEqualizer = equalizer lpCoeff bpCoeff hpCoeff 2 bassUp bassDn trebleUp trebleDn audioIn
  where audioIn = ones (pts * 4)

--testButtonInterface = buttonInterface  bassUp bassDn trebleUp trebleDn

-- fs = 200

--lpCoeffF8 = mapV real2Fixed8 lpCoeffD
--bpCoeffF8 = mapV real2Fixed8 bpCoeffD
--hpCoeffF8 = mapV real2Fixed8 hpCoeffD
--sF8    = mapSY real2Fixed8 s

--lpCoeffF16 = mapV real2Fixed16 lpCoeffD
--bpCoeffF16 = mapV real2Fixed16 bpCoeffD
--hpCoeffF16 = mapV real2Fixed16 hpCoeffD
--sF16    = mapSY real2Fixed16 s

--outF16    = audioFilter lpCoeffF16 bpCoeffF16 hpCoeffF16 (mapSY real2Fixed16 zeros) (mapSY real2Fixed16 zeros) sF16
--outF8    = audioFilter lpCoeffF8 bpCoeffF8 hpCoeffF8 (mapSY real2Fixed8 zeros)(mapSY real2Fixed8 zeros) sF8

--writeAudioOut = writeFile "Test/AudioOut.for") . writeS

-------------
-- SCRIPTS --
-------------

-- | Tests the 'audioFilterD' function, against the contents of a
-- @.dat@ file. Dumps in another @.dat@ file. An example input file is
-- provided in the @files/test-equalizer@ folder in the root path of
-- the project.
testFilterScript inFile = do
  contents <- readFile inFile
  writeFile audioOut (writeS (audioFilterD (readS contents)))
  putStrLn $ "\nDone. Data dumped in " ++ audioOut
  where
    audioOut = takeDirectory inFile ++ "/AudioOutFSD.ext"

-- | Tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioAnalyzer'
-- function, against the contents of a @.dat@ file. Dumps in another
-- @.dat@ file. An example input file is provided in the
-- @files/test-equalizer@ folder in the root path of the project.
testAnalyzerScript inFile = do
  fftInfile  <- openFile inFile ReadMode
  fftOutfile <- openFile outFile WriteMode
  contents   <- hGetContents fftInfile
--    hPutStr fftOutfile (writeS (audioAnalyzer 6 ((readS contents2) :: Signal Double)))
  putStr (show (audioAnalyzer pts ((readS contents) :: Signal Double)))
  hPutStr fftOutfile (writeS (audioAnalyzer pts ((readS contents) :: Signal Double)))
--    hClose fftOutfile
  putStr $ "\nDone. Data dumped in " ++ outFile ++ "\n"
  where
    k = 8
    n = 2 ^ k
    sig = takeS (2*(2^k)) input
    input = 0.1 :- 0.2 :- input
    outFile = takeDirectory inFile ++ "/fftOut.dat"

-- | Chains and tests the
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioFilter' and
-- 'ForSyDe.Shallow.Example.Synchronous.Equalizer.AudioFilter.audioAnalyzer'
-- functions, against the contents of a @.dat@ file. Dumps in another
-- @.dat@ file. An example input file is provided in the
-- @files/test-equalizer@ folder in the root path of the project.
testFilterAndAnalyzerScript inFile = do
  -- Test AudioFilter
  putStr "\n-->Test AudioFilter \n"
  filterInfile     <- openFile inFile ReadMode
  filterContents   <- hGetContents filterInfile
  putStr (show (audioFilterD (readS filterContents)))
  writeFile audioOut (writeS (audioFilterD (readS filterContents)))
  putStrLn $ "\nDone. Data dumped in " ++ audioOut
-- Test AudioAnalyzer
  putStr "\n--> Test AudioAnalyzer \n"
  analyzerInfile   <- openFile audioOut ReadMode
  analyzerOutfile  <- openFile analyzerOut WriteMode
  analyzerContents <- hGetContents analyzerInfile
  putStr (show (audioAnalyzer pts ((readS analyzerContents) :: Signal Double)))
  hPutStr analyzerOutfile (writeS (audioAnalyzer pts ((readS analyzerContents) :: Signal Double)))
  putStrLn $ "\nDone. Data dumped in " ++ analyzerOut
  --fftInfile <- openFile "audioOut.txt" ReadMode
  --fftOutfile <- openFile "fftOut.txt" WriteMode
  --contents <- hGetContents fftInfile
  --putStr (writeS (((readS contents) :: Signal Double)))
  --hPutStr fftOutfile (writeS (((readS contents) :: Signal Double)))
  where
    audioOut    = takeDirectory inFile ++ "/audioOut.dat"
    analyzerOut = takeDirectory inFile ++ "/analyzerOut.dat"

max_stars_repo_name: Rojods/forsyde-shallow-examples

max_issues_repo_name: Rojods/forsyde-shallow-examples

max_forks_repo_name: Rojods/forsyde-shallow-examples


text:
{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE DeriveGeneric #-}
module FokkerPlanck.BrownianMotion
  ( Particle(..)
  , generatePath
  , moveParticle
  , thetaPlus
  , scalePlus
  , generateRandomNumber
  , generatePath'
  ) where

import           Control.DeepSeq
import           Control.Monad
import           Data.DList                          as DL
import           Data.Ix
import           Data.List                           as L
import           GHC.Generics                        (Generic)
import           Statistics.Distribution
import           Statistics.Distribution.Exponential
import           Statistics.Distribution.Normal
import           Statistics.Distribution.Uniform
import           System.Random.MWC
import           Text.Printf
import           Utils.Distribution

data Particle =
  Particle {-# UNPACK #-}!Double -- \phi
           {-# UNPACK #-}!Double -- \rho
           {-# UNPACK #-}!Double -- \theta
           {-# UNPACK #-}!Double -- r
           {-# UNPACK #-}!Double -- weight
  deriving (Show,Generic)

instance NFData Particle

{-# INLINE getParticleRho #-}
getParticleRho :: Particle -> Double
getParticleRho (Particle _ rho _ _ _ ) = rho

thetaCheck :: Double -> Double
thetaCheck theta =
  if theta < -pi
    then thetaCheck (theta + 2 * pi)
    else if theta >= pi
           then thetaCheck (theta - 2 * pi)
           else theta

{-# INLINE scaleCheck #-}
scaleCheck :: Double -> Double -> Double
scaleCheck maxScale scale =
  if scale >= (1 / maxScale) && scale <= maxScale
    then scale
    else error
           (printf
              "scaleCheck: %.2f is out of boundary (0,%.2f)\n"
              scale
              maxScale)

{-# INLINE thetaPlus #-}
thetaPlus :: Double -> Double -> Double
thetaPlus !x !y =
  let z = x + y
   in thetaCheck z

{-# INLINE scalePlus #-}
scalePlus :: Double -> Double -> Double
scalePlus x delta = x * exp delta

{-# INLINE scalePlusPeriodic #-}
scalePlusPeriodic :: Double -> Double -> Double -> Double
scalePlusPeriodic !logMaxScale !delta !x
  | z >= m = exp (z - 2 * m)
  | z < -m = exp (2 * m + z)
  | otherwise = exp z
  where
    !m = logMaxScale
    !z = log x + delta

{-# INLINE rhoCutoff #-}
rhoCutoff :: Double -> Double
rhoCutoff !rho =
  if rho < 0
    then 0
    else rho

{-# INLINE generateRandomNumber #-}
generateRandomNumber ::
     (Distribution d, ContGen d) => GenIO -> Maybe d -> IO Double
generateRandomNumber !gen dist =
  case dist of
    Nothing -> return 0
    Just d  -> genContVar d gen

{-# INLINE moveParticle #-}
moveParticle :: Double -> Particle -> Particle
moveParticle deltaT (Particle phi rho theta r0 v) =
  let !r = r0 * deltaT
      !x = theta - phi
      !cosX = cos x
      !newPhi = phi `thetaPlus` atan2 (r * sin x) (rho + r * cosX)
      !newRho = sqrt (rho * rho + r * r + 2 * r * rho * cosX)
  in Particle newPhi newRho theta r0 v

{-# INLINE diffuseParticle #-}
diffuseParticle :: Double -> Double -> Double -> Particle -> Particle
diffuseParticle !deltaTheta !deltaScale maxScale (Particle phi rho theta r v) =
  Particle
    phi
    rho
    (theta `thetaPlus` deltaTheta)
    (r `scalePlus` deltaScale)
    v

brownianMotion ::
     (Distribution d1, ContGen d1, Distribution d2, ContGen d2, Distribution d3)
  => GenIO
  -> Maybe d1
  -> Maybe d2
  -> Maybe d3
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Particle
  -> DList Particle
  -> IO (DList Particle)
brownianMotion randomGen thetaDist scaleDist poissonDist deltaT maxRho maxR tao stdR2 particle xs = do
  deltaTheta <- generateRandomNumber randomGen thetaDist
  deltaScale <- generateRandomNumber randomGen scaleDist
  let newParticle@(Particle phi rho theta r v) = moveParticle deltaT particle
      ys =
        if r <= maxR && r > 1 / maxR && rho <= maxRho && rho >= 1 / maxRho
          then DL.cons
                 (Particle phi rho theta r (1 - gaussian2DPolar rho stdR2))
                 xs
          else xs
  t <- genContVar (uniformDistr 0 1) randomGen :: IO Double
  if t > exp (1 / (-tao))
    then return ys
    else brownianMotion
           randomGen
           thetaDist
           scaleDist
           poissonDist
           deltaT
           maxRho
           maxR
           tao
           stdR2
           (diffuseParticle deltaTheta deltaScale maxR newParticle)
           ys

{-# INLINE generatePath #-}
generatePath ::
     (Distribution d1, ContGen d1, Distribution d2, ContGen d2, Distribution d3)
  => Maybe d1
  -> Maybe d2
  -> Maybe d3
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> GenIO
  -> IO (DList Particle)
generatePath thetaDist scaleDist poissonDist maxRho maxR tao deltaT stdR2 randomGen =
  brownianMotion
    randomGen
    thetaDist
    scaleDist
    poissonDist
    deltaT
    maxRho
    maxR
    tao
    stdR2
    (Particle 0 1 0 1 1)
    DL.empty

{-# INLINE moveParticle' #-}
moveParticle' ::  Particle -> Particle
moveParticle' (Particle phi rho theta r v) =
  let !x = theta - phi
      !cosX = cos x
      !newPhi = phi `thetaPlus` atan2 (r * sin x) (rho + r * cosX)
      !newRho = sqrt (rho * rho + r * r + 2 * r * rho * cosX)
   in Particle newPhi newRho theta r v

{-# INLINE diffuseParticle' #-}
diffuseParticle' :: (Distribution d, ContGen d)
  => GenIO ->  Double -> Maybe d -> Double -> Particle -> IO Particle
diffuseParticle' randomGen thetaSigma scaleDist deltaT (Particle phi rho theta r v) = do
  deltaScale <- generateRandomNumber randomGen scaleDist
  let r1 = r `scalePlus` deltaScale
  deltaTheta <-
    genContVar (normalDistr 0 (thetaSigma * sqrt (r1 * deltaT))) randomGen
  return $ Particle phi rho (theta `thetaPlus` deltaTheta) r1 v

brownianMotion' ::
     (Distribution d, ContGen d)
  => GenIO
  -> Double
  -> Maybe d
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Particle
  -> DList Particle
  -> IO (DList Particle)
brownianMotion' randomGen thetaSigma scaleDist poissonLambda deltaT maxRho maxR tao stdR2 particle xs = do
  let newParticle@(Particle phi rho theta r v) = moveParticle deltaT particle
      ys =
        if r <= maxR && r > 1 / maxR && rho <= maxRho && rho >= 1 / maxRho
          then DL.cons (Particle phi rho theta r (1 - gaussian2DPolar rho stdR2)) xs
          else xs
  t <- genContVar (uniformDistr 0 1) randomGen :: IO Double
  if t > exp (1 / (-tao / (r * deltaT)))
    then return ys
    else do
      diffusedParticle <-
        diffuseParticle' randomGen thetaSigma scaleDist deltaT newParticle
      brownianMotion'
        randomGen
        thetaSigma
        scaleDist
        poissonLambda
        deltaT
        maxRho
        maxR
        tao
        stdR2
        diffusedParticle
        ys

{-# INLINE generatePath' #-}
generatePath' ::
     (Distribution d, ContGen d)
  => Double
  -> Maybe d
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> GenIO
  -> IO (DList Particle)
generatePath' thetaSigma scaleDist poissonLambda maxRho maxR tao deltaT stdR2 randomGen = do
  brownianMotion'
    randomGen
    thetaSigma
    scaleDist
    poissonLambda
    deltaT
    maxRho
    maxR
    tao
    stdR2
    (Particle 0 1 0 1 1)
    DL.empty

ext: hs

max_stars_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_issues_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_forks_repo_name: XinhuaZhang/Stochastic-Completion-Field


text:
module FokkerPlanck.Histogramming where

import           Array.UnboxedArray          as UA
import           Data.Complex
import           Data.DList                  as DL
import           Data.List                   as L
import           FokkerPlanck.BrownianMotion
import           FokkerPlanck.Histogram

-- {-# INLINE sampleScale #-}
-- sampleScale ::
--      Double
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> Particle
--   -> Complex Double
-- sampleScale halfLogPeriod deltaLogRho phiFreq rhoFreq rFreq (Particle phi rho theta r) =
--   let n = floor $ (log r + halfLogPeriod) / deltaLogRho
--   in (deltaLogRho :+ 0) *
--      (L.sum .
--       L.map
--         (\particle ->
--            let (Particle newPhi newRho _ newR) = moveParticle particle
--            in cis $
--               pi * (rhoFreq * log newRho + rFreq * log newR) / halfLogPeriod +
--               phiFreq * newPhi) $
--       [(Particle phi rho theta (fromIntegral i * deltaLogRho)) | i <- [1 .. n]])

-- -- {-# INLINE countR2S1 #-}
-- -- countR2S1 ::
-- --      (Int, Int)
-- --   -> (Int, Int)
-- --   -> Int
-- --   -> [DList ParticleIndex]
-- --   -> Histogram Int
-- -- countR2S1 (xMin, xMax) (yMin, yMax) numOrientations xs =
-- --   let deltaTheta = 2 * pi / (fromIntegral numOrientations)
-- --       ys =
-- --         L.map (\(t, x, y) -> ((floor $ t / deltaTheta, round x, round y), 1)) .
-- --         DL.toList .
-- --         DL.concat . L.map (DL.map (\(x, y, t, s, t0, s0) -> (t, x, y))) $
-- --         xs
-- --       numTrajectories = L.length ys
-- --       arr =
-- --         UA.accum (+) 0 ((0, xMin, yMin), (numOrientations - 1, xMax, yMax)) ys
-- --    in Histogram
-- --         [(yMax - yMin + 1), (xMax - xMin + 1), numOrientations]
-- --         numTrajectories
-- --         (toUnboxedVector arr)

-- -- {-# INLINE countR2Z2T0S0 #-}
-- -- countR2Z2T0S0 ::
-- --      (Int, Int)
-- --   -> (Int, Int)
-- --   -> [Double]
-- --   -> [Double]
-- --   -> [Double]
-- --   -> [Double]
-- --   -> Double
-- --   -> [DList ParticleIndex]
-- --   -> Histogram (Complex Double)
-- -- countR2Z2T0S0 (xMin, xMax) (yMin, yMax) t0Freqs tFreqs s0Freqs sFreqs maxR xs =
-- --   let logMaxR =
-- --         if maxR == 1
-- --           then 1
-- --           else log maxR
-- --       ys =
-- --         DL.toList .
-- --         DL.concat .
-- --         L.map
-- --           (\((t0f, i), (tf, j), (s0f, k), (sf, l)) ->
-- --              DL.map
-- --                (\(x, y, t, s, t0, s0) ->
-- --                   let !v =
-- --                         exp
-- --                           (0 :+
-- --                            (-t0f * t0 + tf * t +
-- --                             (sf * s - s0f * s0) * 2 * pi / logMaxR))
-- --                       !x' = round x
-- --                       !y' = round y
-- --                   in ((j, l, i, k, x', y'), v)) .
-- --              DL.concat $
-- --              xs) $
-- --         [ (t0f, tf, s0f, sf)
-- --         | t0f <- (L.zip t0Freqs [1 ..])
-- --         , tf <- (L.zip tFreqs [1 ..])
-- --         , s0f <- (L.zip s0Freqs [1 ..])
-- --         , sf <- (L.zip sFreqs [1 ..])
-- --         ]
-- --       numTrajectories =
-- --         div
-- --           (L.length ys)
-- --           ((L.length t0Freqs) * (L.length tFreqs) * (L.length s0Freqs) *
-- --            (L.length sFreqs))
-- --       arr =
-- --         UA.accum
-- --           (+)
-- --           0
-- --           ( (1, 1, 1, 1, xMin, yMin)
-- --           , ( L.length tFreqs
-- --             , L.length sFreqs
-- --             , L.length t0Freqs
-- --             , L.length s0Freqs
-- --             , xMax
-- --             , yMax))
-- --           ys
-- --   in Histogram
-- --        [ (yMax - yMin + 1)
-- --        , (xMax - xMin + 1)
-- --        , (L.length s0Freqs)
-- --        , (L.length t0Freqs)
-- --        , (L.length sFreqs)
-- --        , (L.length tFreqs)
-- --        ]
-- --        numTrajectories
-- --        (toUnboxedVector arr)


text:
module Data.Random.Distribution.Hypergeometric.Impl where

import Data.Random.RVar
import Data.Random.Distribution
import Data.Random.Distribution.Uniform

import Numeric.SpecFunctions (logGamma)

d1 = 1.7155277699214135 -- 2*sqrt(2/e)
d2 = 0.8989161620588988 -- 3 - 2*sqrt(3/e)

rhyper :: (Integral a) => (a, a, a) -> RVarT m a
rhyper (sample, popsize, good) = return . fix2 . fix1 =<< loop
  where fix1 z = if good > bad then m - z else z
        fix2 z = if m < sample then good - z else z
        mingoodbad = min good bad
        bad = popsize - good
        maxgoodbad = max good bad
        m = min sample (popsize - sample)
        gamma z = sum $ map (logGamma . fromIntegral)
          [ z + 1,     mingoodbad - z + 1
          , m - z + 1, maxgoodbad - m + z + 1 ]
        d4 = fromIntegral mingoodbad / fromIntegral popsize
        d5 = 1 - d4
        d6 = fromIntegral m * d4 + 0.5
        d7 = sqrt $ fromIntegral (popsize - m) * fromIntegral sample * d4 * d5 / fromIntegral (popsize - 1) + 0.5
        d8 = d1 * d7 + d2
        d9 = floor $ fromIntegral (m + 1) * fromIntegral (mingoodbad + 1) / fromIntegral (popsize + 2)
        d10 = gamma d9
        d11 = fromIntegral $ min (min m mingoodbad + 1) (floor (d6 + 16 * d7))
        -- 16 for 16-decimal-digit precision in d1 and d2
        loop = do
          x <- doubleUniform 0 1
          y <- doubleUniform 0 1
          case () of
            _  | w < 0 || w >= d11    -> loop -- fast rejection
               | x * (4 - x) - 3 <= t -> return z -- fast acceptance
               | x * (x - t) >= 1     -> loop -- fast rejection
               | 2 * (log x) <= t     -> return z -- acceptance
               | otherwise            -> loop -- rejection
               where w = d6 + d8 * (y - 0.5) / x
                     z = floor w
                     t = d10 - gamma z

max_stars_repo_name: srijs/random-hypergeometric

max_issues_repo_name: srijs/random-hypergeometric

max_forks_repo_name: srijs/random-hypergeometric


text:
{-# LANGUAGE OverloadedStrings #-}

module Numeric.PHCPack.Parse (
  parseSolutions
  )
  where

import Data.Attoparsec
import Data.Attoparsec.ByteString.Char8 hiding (satisfy)
import Data.Attoparsec.Combinator
import Control.Applicative((<|>))

import Data.ByteString(ByteString)

import Data.ByteString.Char8(pack)
import qualified Data.ByteString as BS

import Numeric.PHCPack.Types
import qualified Data.Map as Map
import Data.Complex

space' :: Parser ()
space' = satisfy isHorizontalSpace >> return ()

spaces :: Parser ()
spaces = skipMany space' >> return ()

spaces1 :: Parser ()
spaces1 = space' >> spaces

parseSolutions :: ByteString -> Maybe [Solution]
parseSolutions = maybeResult . ((flip feed) BS.empty) . (parse solutionsParser)

solutionsParser :: Parser [Solution]
solutionsParser = do
  manyTill anyChar (try (string "START SOLUTIONS : " >> endOfLine))
  skipMany (space' <|> endOfLine)
  solutionCount <- decimal
  spaces1
  dimension <- decimal
  endOfLine

  sepLine
  solutions <- solution `sepBy1` bannerLine
  bannerLine
  skipMany anyChar

  return solutions

sepLine :: Parser ()
sepLine = skipMany (char '=') >> endOfLine

bannerLine :: Parser ()
bannerLine = string "==" >> manyTill anyChar (try endOfLine) >> return ()

solution :: Parser Solution
solution = do
  string "solution "
  solutionNumber <- decimal
  manyTill anyChar (try (string "the solution" >> manyTill anyChar (try endOfLine))) -- skip until "the solution for t :"
  solutions <- variableSolution `sepBy` endOfLine
  endOfLine
  return (Solution $ Map.fromList solutions)

-- variableSolution :: (Num c) => Parser (Unknown, Complex c)
variableSolution = do
  spaces
  name <- manyTill anyChar space
  spaces
  char ':'
  spaces
  real <- double
  spaces
  imag <- double
  return (name, real :+ imag)

max_stars_repo_name: mruegenberg/PHCPack-hs

max_issues_repo_name: mruegenberg/PHCPack-hs

max_forks_repo_name: mruegenberg/PHCPack-hs


text:
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict #-}
module STC.CompletionField where

import           Control.Monad.Parallel as MP
import           Data.Array.Repa        as R
import           Data.Complex
import           Data.List              as L
import           Data.Vector.Storable   as VS
import           DFT.Plan
import           Filter.Utils
import           STC.Convolution
import           STC.DFTArray
import           STC.Plan
import           Utils.Array
import           Utils.Parallel

{-# INLINE completionField #-}
completionField :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
completionField plan source@(DFTArray rows cols thetaFreqs rFreqs _) sink = do
  let numThetaFreq = L.length thetaFreqs
      numRFreq = L.length rFreqs
      sourceFilter =
        VS.convert .
        toUnboxed .
        computeUnboxedS .
        R.backpermute
          (Z :. numRFreq :. numThetaFreq :. cols :. rows)
          (\(Z :. a :. b :. c :. d) ->
             (Z :. (makeFilterHelper numRFreq a) :.
              (makeFilterHelper numThetaFreq b) :.
              c :.
              d)) .
        pad [rows, cols, numThetaFreq, numRFreq] 0 .
        dftArrayToRepa $
        source
      dftID = DFTPlanID DFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]
  dftSource <- dftExecute plan dftID sourceFilter
  dftSink <- dftExecute plan dftID . VS.concat . getDFTArrayVector $ sink
  arr <-
    fmap
      (fromUnboxed (Z :. numRFreq :. numThetaFreq :. cols :. rows) . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]) .
    VS.zipWith (*) dftSource $
    dftSink
  return $ repaToDFTArray thetaFreqs rFreqs $ arr

{-# INLINE completionField' #-}
completionField' :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
completionField' plan source@(DFTArray rows cols thetaFreqs rFreqs _) sink = do
  let numThetaFreq = L.length thetaFreqs
      sourceFilter =
        VS.convert .
        toUnboxed .
        computeUnboxedS .
        R.backpermute
          (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows)
          (\(Z :. _ :. b :. c :. d) ->
             (Z :. (0 :: Int) :. (makeFilterHelper numThetaFreq b) :. c :. d)) .
        dftArrayToRepa $
        source
      dftID = DFTPlanID DFT1DG [1, numThetaFreq, cols, rows] [0, 1]
  dftSource <- dftExecute plan dftID sourceFilter
  dftSink <- dftExecute plan dftID . VS.concat . getDFTArrayVector $ sink
  arr <-
    fmap
      (fromUnboxed (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows) .
       VS.convert) .
    dftExecute plan (DFTPlanID IDFT1DG [1, numThetaFreq, cols, rows] [0, 1]) .
    VS.zipWith (*) dftSource $
    dftSink
  return $ repaToDFTArray thetaFreqs rFreqs $ arr

{-# INLINE completionFieldRepa #-}
completionFieldRepa ::
     (R.Source s1 (Complex Double), R.Source s2 (Complex Double))
  => DFTPlan
  -> R.Array s1 DIM4 (Complex Double)
  -> R.Array s2 DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
completionFieldRepa plan source sink = do
  print . extent $ source
  let (Z :. numRFreq :. numThetaFreq :. cols :. rows) = extent source
      dftID = DFTPlanID DFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]
  sourceFilter <-
    fmap (VS.convert . toUnboxed) .
    computeUnboxedP .
    R.backpermute
      (extent source)
      (\(Z :. a :. b :. c :. d) ->
         Z :. makeFilterHelper numRFreq a :. makeFilterHelper numThetaFreq b :.
          c :.
          d) $
    source
  dftSource <- dftExecute plan dftID sourceFilter
  sinkU <- computeUnboxedP . delay $ sink
  dftSink <- dftExecute plan dftID . VS.convert . toUnboxed $ sinkU
  fmap (fromUnboxed (extent source) . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]) .
    VS.zipWith (*) dftSource $
    dftSink

-- {-# INLINE completionField #-}
-- completionField :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
-- completionField plan source@(DFTArray rows cols thetaFreqs rFreqs _) sink = do
--   let numThetaFreq = L.length thetaFreqs
--       numRFreq = L.length rFreqs
--       sourceFilter =
--         VS.convert .
--         toUnboxed .
--         computeUnboxedS .
--         R.backpermute
--           (Z :. numRFreq :. numThetaFreq :. cols :. rows)
--           (\(Z :. a :. b :. c :. d) ->
--              (Z :. (makeFilterHelper numRFreq a) :.
--               (makeFilterHelper numThetaFreq b) :.
--               c :.
--               d)) .
--         dftArrayToRepa $
--         source
--       dftID = DFTPlanID DFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]
--   dftSource <- dftExecute plan dftID $ sourceFilter
--   dftSink <- dftExecute plan dftID . VS.concat . getDFTArrayVector $ sink
--   arr <-
--     fmap
--       (fromUnboxed (Z :. numRFreq :. numThetaFreq :. cols :. rows) . VS.convert) .
--     dftExecute
--       plan
--       (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]) .
--     VS.zipWith (*) dftSource $
--     dftSink
--   return $ repaToDFTArray thetaFreqs rFreqs $ arr


text:
module Solve where

import Data.List
import Control.Applicative ((<|>), many)
import Control.Monad (void)
import Data.Char (isLetter, isDigit)
import qualified Debug.Trace as DT
import Numeric.LinearAlgebra.Data
import Data.Int

import Text.ParserCombinators.Parsec
import qualified Text.ParserCombinators.Parsec.Token as P
import Text.ParserCombinators.Parsec.Language

-- TYPE DEFINITIONS
data Point = MkPoint {
  x :: Int,
  y :: Int
} deriving (Show, Eq)

data Dimensions = MkDimensions {
  width  :: Int,
  height :: Int
} deriving (Show, Eq)

data Rectangle = MkRectangle {
  upperLeft :: Point,
  dimensions :: Dimensions
} deriving (Show, Eq)

data Claim  = MkClaim {
    index :: Int,
    rectangle :: Rectangle
} deriving (Show, Eq)

-- PARSING
-- of course, we could simply use a regular expression, but I've always wanted to play
-- around with the Parsec library for generating parsers, so bear with me :-)
numParser :: Parser Int
numParser = do
    n <- many1 digit
    return (read n)

pointParser :: Parser Point
pointParser = do
    x <- numParser
    void $ char ','
    y <- numParser
    return (MkPoint x y)

dimensionsParser :: Parser Dimensions
dimensionsParser = do
  width <- numParser
  void $ char 'x'
  height <- numParser
  return (MkDimensions width height)

rectangleParser :: Parser Rectangle
rectangleParser = do
    p <- pointParser
    void $ char ':'
    void $ char ' '
    dim <- dimensionsParser
    return (MkRectangle p dim)

claimParser :: Parser Claim
claimParser = do
    void $ char '#'
    idx <- numParser
    void $ char ' '
    void $ char '@'
    void $ char ' '
    rect <- rectangleParser
    return (MkClaim idx rect)

-- helper function to apply a parser
regularParse :: Parser a -> String -> Either ParseError a
regularParse p = parse p ""

-- MAIN routines
--
-- create an association list for the given claim
-- the association list contains a list of (x,y) coordinates plus the value (which equals the claim's index)
toAssocList :: Claim -> [((Int, Int), Int64)]
toAssocList c = map (\(i,j) -> ((i,j), (fromIntegral idx))) indices
  where idx = index c
        r = rectangle c
        x' = x $ upperLeft r
        y' = y $ upperLeft r
        width' = width $ dimensions r
        height' = height $ dimensions r
        indices = [(i,j) | i <- [(x'+1) .. (x' + width')],
                           j <- [(y'+1) .. (y' + height')]]

-- apply the given claim to the given matrix
applyClaim :: Claim -> Matrix Int64 -> Matrix Int64
applyClaim c m = accum m (\new old -> if old == 0 then new else -1) (toAssocList c)

-- apply the given list of claims to the given matrix
applyClaims :: [Either ParseError Claim ] -> Matrix Int64 -> Matrix Int64
applyClaims [] m = m
applyClaims (x:xs) m =
  case x of
    Right x' -> applyClaims xs (applyClaim x' m)
    Left _   -> applyClaims xs m -- leave m unchanged for erroneous claim

-- count the elements in the given matrix that are equal to the given value
countElems :: Matrix Int64 -> Int64 -> Int
countElems m val = length $ filter (\x -> x == val) $ toList $ flatten m

-- count the clashes (squares with overlapping claims)
countClashes :: Matrix Int64 -> Int
countClashes m = countElems m (-1)

-- prepare solution: parse input, build a 1001x1001 matrix, and apply all claims to it
prepSolution:: String -> ([Claim], Matrix Int64)
prepSolution input = (validClaims, applyClaims claims m)
  where claims = map (\line -> regularParse claimParser line) $ lines input
        m = (1001><1001) (repeat 0) -- 1001x1001, all elements are 0
        m2 = applyClaims claims m
        validClaims = map (\(Right x) -> x) $ filter (\x -> case x of Right _ -> True) claims

-- find the single claim that is intact
-- our approach is kind of brute force - for n claims, we check all matrix elements n times
findIntactClaim :: ([Claim], Matrix Int64) -> Claim
findIntactClaim (claims, m) = head $ filter (\c -> area c == countElems m (fromIntegral (index c))) claims
  where area c' = width (dimensions (rectangle c')) * height (dimensions (rectangle c'))

solveI :: String -> Int
solveI = countClashes . snd . prepSolution

solveII :: String -> Int
solveII = index . findIntactClaim . prepSolution

#if defined(STANDALONE)
main = do
    input <- readFile "input.txt"
    putStrLn $ show $ solveI  input
    putStrLn $ show $ solveII input
#endif


text:
module LossFunction
    ( meanSquaredError
    , crossEntropyError
    ) where

import Numeric.LinearAlgebra

meanSquaredError :: (Container c b, Num (c b), Floating b) => c b -> c b -> b
meanSquaredError x y = (*0.5) . sumElements $ (x-y)^2

crossEntropyError :: (Container c b, Num (c b), Floating b) => c b -> c b -> b
crossEntropyError x y = (*(-1)) . sumElements $ y * cmap (\x -> log(x+1e-7)) x


text:
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE RoleAnnotations #-}

module Deep.Expr where

import           Data.Void
import           Data.Proxy
import           Data.Constraint (Constraint, Dict(..))

import           Data.Bifunctor

import           GHC.Generics
import           Data.Generics.Uniplate.Direct
-- import           Language.Haskell.TH (Name)

import           Data.Complex

import           Control.Monad.State

import           Data.Type.Equality
import           Data.Typeable

import           Data.Char

import           GHC.Prim
import           Data.Coerce

-- Idea: Represent a pattern match of a particular type as
--    KnownSymbol name => Proxy name -> MatchType name
-- with 'MatchType' being a type family that gives the type of the match
-- associated to matching on the constructor named 'name'. Maybe these
-- could be stored in a heterogenous list.

-- | Idea: Deep embedding for a pattern match (possibly using a non-regular
-- type?). Convert types to a "canonical form" to use this (like 'Either's
-- and '(,)'s)

-- TODO: Should we have a phantom type argument with the original type to preserve (possibly) more type safety?

-- TODO: Should this just be a function type? Maybe it could be a function
-- from s -> t, where is a nested pair type, and we use projection
-- combinators to extract the values inside the function.

-- data ProdMatch s t where
--   ProdMatch :: forall a b r.
--     (GPURep a, GPURep b)
--        => (GPUExp a -> ProdMatch b r) -> ProdMatch  (a, b) r

--   OneProdMatch :: forall a r. (GPURep a) => (GPUExp a -> GPUExp r) -> ProdMatch a r
--   NullaryMatch :: forall a r. GPURep r => GPUExp r -> ProdMatch a r

-- data SumMatch s t where
--   SumMatch :: forall a b r.
--     (GPURep a, GPURep b, GPURepTy b ~ b)
--         => ProdMatch a r -> SumMatch b r -> SumMatch (Either a b) r

--   EmptyMatch :: forall r. SumMatch Void r

--   OneSumMatch :: forall a b. (GPURep a, GPURep b, GPURepTy a ~ a) => ProdMatch a b -> SumMatch a b

-- | The function this contains just is used to provide a semantics for
-- interpreting a GPUExp
-- type role SumMatch nominal nominal
newtype SumMatch  a b = MkSumMatch  { runSumMatch  :: a -> b }
-- type role ProdMatch nominal nominal
newtype ProdMatch a b = MkProdMatch { runProdMatch :: a -> b }

-- Done (case ... of A -> x; B -> y)  ==>  (case ... of A -> Done x; B -> Done y)
data Iter a b
  = Step b
  | Done a
  deriving (Functor, Generic)

-- | NOTE: The unique identifiers start at 1
-- type role Name nominal
newtype Name a = Name { getNameIdent :: Int } deriving (Eq, Show)

getNameUniq :: Name a -> Int
getNameUniq = getNameIdent

data EnvMapping where
  (:=>) :: forall a. Typeable a => (Name a) -> GPUExp a -> EnvMapping

newtype Env = Env [EnvMapping]

emptyEnv :: Env
emptyEnv = Env []

extendEnv :: Env -> EnvMapping -> Env
extendEnv (Env maps) m = Env (m:maps)

envLookup :: Typeable a => Env -> Name a -> Maybe (GPUExp a)
envLookup (Env maps) n = go maps
  where
    go [] = Nothing
    go ((n' :=> e):rest) =
      case namesEq n n' of
        Just Refl -> Just e
        Nothing   -> go rest

namesEq :: forall a b. (Typeable a, Typeable b) => Name a -> Name b -> Maybe (a :~: b)
namesEq (Name n) (Name n') =
  case eqT :: Maybe (a :~: b) of
    Just Refl
      | n == n' -> Just Refl
      | otherwise -> Nothing
    Nothing   -> Nothing

data GPUExp t where
  CaseExp :: forall t x r. (GPURep t, GPURepTy x ~ x, GPURepTy t ~ x) => GPUExp t -> GPUExp (SumMatch x r) -> GPUExp r

  SumMatchExp :: forall a b r.
    (GPURep a, GPURep b, GPURepTy b ~ b)
      => GPUExp (ProdMatch a r) -> GPUExp (SumMatch b r) -> GPUExp (SumMatch (Either a b) r)

  ProdMatchExp :: forall a b r.
    (GPURep a, GPURep b)
      => GPUExp (a -> ProdMatch b r) -> GPUExp (ProdMatch (a, b) r)

  NullaryMatch :: forall a r.
    GPURep a
      => GPUExp r -> GPUExp (ProdMatch a r)

  OneSumMatch :: forall a b.
    (GPURep a, GPURep b, GPURepTy a ~ a)
      => GPUExp (ProdMatch a b) -> GPUExp (SumMatch a b)

  OneProdMatch :: forall a b.
    (GPURep a)
      => GPUExp (a -> b) -> GPUExp (ProdMatch a b)

  EmptyMatch :: forall b. GPUExp (SumMatch Void b)

  FalseExp :: GPUExp Bool
  TrueExp :: GPUExp Bool

  Repped :: (GPURep a, GPURep (GPURepTy a)) => GPURepTy a -> GPUExp a

  Lam :: forall a b. (GPURep a, Typeable a) => Name a -> GPUExp b -> GPUExp (a -> b)
  Var :: forall a. Typeable a => Name a -> GPUExp a -- Constructed internally
  App :: forall a b. (GPURep a) => GPUExp (a -> b) -> GPUExp a -> GPUExp b

  -- Apply :: GPUExp (a -> b) -> GPUExp a -> GPUExp b

  CharLit :: Char -> GPUExp Char
  Lit :: (Typeable a, Show a, Num a) => a -> GPUExp a

  Ord :: GPUExp Char -> GPUExp Int

  Add :: Num a => GPUExp a -> GPUExp a -> GPUExp a
  Sub :: Num a => GPUExp a -> GPUExp a -> GPUExp a
  Mul :: Num a => GPUExp a -> GPUExp a -> GPUExp a

  -- Twice :: GPUExp (a -> a) -> GPUExp a -> GPUExp a

  FromEnum :: Enum a => GPUExp a -> GPUExp Int
  FromIntegral :: forall a b. (Typeable a, Typeable b, Integral a, Num b) => GPUExp a -> GPUExp b

  Sqrt :: Floating a => GPUExp a -> GPUExp a

  Equal :: Eq a => GPUExp a -> GPUExp a -> GPUExp Bool
  Lte :: Ord a => GPUExp a -> GPUExp a -> GPUExp Bool
  Gt :: Ord a => GPUExp a -> GPUExp a -> GPUExp Bool

  Not :: GPUExp Bool -> GPUExp Bool

  LeftExp :: forall a b. GPUExp a -> GPUExp (Either a b)
  RightExp :: forall a b. GPUExp b -> GPUExp (Either a b)

  PairExp :: forall a b. (GPURep a, GPURep b) => GPUExp a -> GPUExp b -> GPUExp (a, b)
  FstExp :: forall a b. GPUExp (a, b) -> GPUExp a
  SndExp :: forall a b. GPUExp (a, b) -> GPUExp b

  StepExp :: forall a b. GPUExp b -> GPUExp (Iter a b)
  DoneExp :: forall a b. GPUExp a -> GPUExp (Iter a b)

  UnitExp :: GPUExp ()

  IfThenElse :: GPUExp Bool -> GPUExp a -> GPUExp a -> GPUExp a

  TailRec :: forall a b. (GPURep a, GPURep b) => GPUExp (b -> Iter a b) -> GPUExp (b -> a)

    -- Typeable allows us to inspect the types. This can be useful if we
    -- want to, for instance, treat 'Complex Double' values differently
    -- than '(Double, Double)' values in a backend
  ConstructRep :: forall a. (Typeable a, GPURep a, GPURep (GPURepTy a)) => GPUExp (GPURepTy a) -> GPUExp a

  Construct :: a -> GPUExp a
  ConstructAp :: forall a b. (GPURep a) => GPUExp (a -> b) -> GPUExp a -> GPUExp b

  CastExp :: forall (a :: *) (b :: *). a ~ b => GPUExp a -> GPUExp b -- Inspiration from a mailing list posting from Conal Elliott

  UnconstructRep :: GPUExp a -> GPUExp (GPURepTy (M1 S ('MetaSel 'Nothing 'NoSourceUnpackedness 'NoSourceStrictness 'DecidedLazy) (Rec0 a) Void))

-- (GPURepTy
--     (M1
--        S
--        ('MetaSel
--           'Nothing 'NoSourceUnpackedness 'NoSourceStrictness 'DecidedLazy)
--        (Rec0 Int)
--        Void))

tailRecApp :: forall a b. (GPURep a, GPURep b) => GPUExp (b -> Iter a b) -> GPUExp b -> GPUExp a
tailRecApp body = App (TailRec body)

transformE :: (forall r. GPUExp r -> GPUExp r) -> GPUExp a -> GPUExp a
transformE tr0 = tr
  where
    tr :: forall s. GPUExp s -> GPUExp s
    tr = go . tr0

    go :: forall s. GPUExp s -> GPUExp s
    go (CaseExp x f) = CaseExp (tr x) (tr f)
    go (SumMatchExp x y) = SumMatchExp (tr x) (tr y)
    go (ProdMatchExp f) = ProdMatchExp (tr f)
    go (NullaryMatch x) = NullaryMatch (tr x)
    go (OneSumMatch p) = OneSumMatch (tr p)
    go (OneProdMatch f) = OneProdMatch (tr f)
    go EmptyMatch = EmptyMatch
    go FalseExp = FalseExp
    go TrueExp = TrueExp
    go e@(Repped _) = e
    go (Lam n x) = Lam n (tr x)
    go e@(Var _) = e
    go e@(Lit _) = e
    go (Add x y) = Add (tr x) (tr y)
    go (Sub x y) = Sub (tr x) (tr y)
    go (Mul x y) = Mul (tr x) (tr y)
    -- go (Twice f x) = Twice
    go (FromEnum x) = FromEnum (tr x)
    go (FromIntegral x) = FromIntegral (tr x)
    go (Sqrt x) = Sqrt (tr x)
    go (Equal x y) = Equal (tr x) (tr y)
    go (Lte x y) = Lte (tr x) (tr y)
    go (Gt x y) = Gt (tr x) (tr y)
    go (Not x) = Not (tr x)
    go (LeftExp x) = LeftExp (tr x)
    go (RightExp y) = RightExp (tr y)
    go (PairExp x y) = PairExp (tr x) (tr y)
    go (FstExp x) = FstExp (tr x)
    go (SndExp x) = SndExp (tr x)
    go (StepExp x) = StepExp (tr x)
    go (DoneExp y) = DoneExp (tr y)
    go (IfThenElse x y z) = IfThenElse (tr x) (tr y) (tr z)
    go (TailRec f) = TailRec (tr f)
    go e@(Construct _) = e
    go (ConstructAp f x) = ConstructAp (tr f) (tr x)
    go (Ord x) = Ord (tr x)
    go e@(CharLit _) = e
    go e@UnitExp = e
    go e@(ConstructRep _) = e
    go (CastExp x) = CastExp (tr (coerce x))
    go (App f x) = App (tr f) (tr x)

transformEA :: forall a m. Monad m
  => (forall r. GPUExp r -> m (GPUExp r)) -> GPUExp a -> m (GPUExp a)
transformEA tr0 = tr
  where
    tr :: forall s. GPUExp s -> m (GPUExp s)
    tr = go <=< tr0

    go :: forall s. GPUExp s -> m (GPUExp s)
    go (CaseExp x f) = CaseExp <$> (tr x) <*> (tr f)
    go (SumMatchExp x y) = SumMatchExp <$> (tr x) <*> (tr y)
    go (ProdMatchExp f) = ProdMatchExp <$> (tr f)
    go (NullaryMatch x) = NullaryMatch <$> (tr x)
    go (OneSumMatch p) = OneSumMatch <$> (tr p)
    go (OneProdMatch f) = OneProdMatch <$> (tr f)
    go EmptyMatch = pure EmptyMatch
    go FalseExp = pure FalseExp
    go TrueExp = pure TrueExp
    go e@(Repped _) = pure e
    go (Lam n x) = Lam n <$> (tr x)
    go e@(Var _) = pure e
    go e@(Lit _) = pure e
    go (Add x y) = Add <$> (tr x) <*> (tr y)
    go (Sub x y) = Sub <$> (tr x) <*> (tr y)
    go (Mul x y) = Mul <$> (tr x) <*> (tr y)
    -- go (Twice f x) = Twice
    go (FromEnum x) = FromEnum <$> (tr x)
    go (FromIntegral x) = FromIntegral <$> (tr x)
    go (Sqrt x) = Sqrt <$> (tr x)
    go (Equal x y) = Equal <$> (tr x) <*> (tr y)
    go (Lte x y) = Lte <$> (tr x) <*> (tr y)
    go (Gt x y) = Gt <$> (tr x) <*> (tr y)
    go (Not x) = Not <$> (tr x)
    go (LeftExp x) = LeftExp <$> (tr x)
    go (RightExp y) = RightExp <$> (tr y)
    go (PairExp x y) = PairExp <$> (tr x) <*> (tr y)
    go (FstExp x) = FstExp <$> (tr x)
    go (SndExp x) = SndExp <$> (tr x)
    go (StepExp x) = StepExp <$> (tr x)
    go (DoneExp y) = DoneExp <$> (tr y)
    go (IfThenElse x y z) = IfThenElse <$> (tr x) <*> (tr y) <*> (tr z)
    go (TailRec f) = TailRec <$> (tr f)
    go e@(Construct _) = pure e
    go (ConstructAp f x) = ConstructAp <$> (tr f) <*> (tr x)
    go (Ord x) = Ord <$> (tr x)
    go e@(CharLit _) = pure e
    go e@UnitExp = pure e
    go e@(ConstructRep _) = pure e
    go (CastExp x) = CastExp <$> tr (coerce x)
    go (App f x) = App <$> tr f <*> tr x

-- convertE :: Coercible a b => GPUExp a -> GPUExp b
-- convertE (CaseExp x f) = CaseExp x (convertE f)
-- convertE (SumMatchExp x y) = SumMatchExp _ _

-- foldEM :: forall a b m. Monad m => (forall x. GPUExp x -> b -> m b) -> b -> GPUExp a -> m b
-- foldEM f z = go
--   where
--     go :: forall s. GPUExp s -> m b
--     go (CaseExp x y) = f x =<< go y
--     go (SumMatchExp x y) = f x =<< go y
--     go (ProdMatchExp x) = f x z
--     go (NullaryMatch x) = f x z
--     go (OneSumMatch x) = f x z
--     go (OneProdMatch x) = f x z
--     go e@EmptyMatch = return z
--     go e@FalseExp = return z
--     go e@TrueExp = return z
--     go e@(Repped _) = return z
--     go (Lam n x) = f x z
--     go (Add x y) = f x =<< go y

-- -- Used for intermediate transformations
-- unLam :: GPUExp (a -> b) -> GPUExp a -> GPUExp b
-- unLam (Lam name body) = go
--   where
--     -- | go x   ==>  body[name|->x]
--     go x = undefined
-- unLam _ = error "unLam: Expected a 'Lam'"

-- subst :: Name -> a -> GPUExp b -> GPUExp b
-- subst name x = go
--   where
--     go = error "subst"

-- For convenience in Core transformation
deepFromInteger :: (Typeable b, Num b) => Integer -> GPUExp b
deepFromInteger = FromIntegral . Lit

-- lam :: forall a b. (GPURep a, Typeable a) => Name a -> (GPUExp a -> GPUExp b) -> GPUExp (a -> b)
-- lam name f = Lam name (f (Var name))

runIter :: forall a b. (GPURep a, GPURep b) => (b -> Iter a b) -> b -> a
runIter f = go
  where
    go x =
      case f x of
        Done r -> r
        Step x' -> go x'

-- unconstruct' :: forall t. (GPURep t, GPURep (GPURepTy t), GPURepTy (GPURepTy t) ~ GPURepTy t, Generic t, GPURep (Rep t Void), GPURepTy (Rep t Void) ~ GPURepTy t) => GPUExp (GPURepTy t) -> GPUExp t
-- unconstruct' (PairExp x y) = _ ( (construct x, construct y))

-- Typeable constraint is for ConstructRep
class Typeable t => GPURep t where
    -- Should we wrap these types in a 'Tagged t' in order to preserve more
    -- type safety?
  type GPURepTy t
  type GPURepTy t = GPURepTy (Rep t Void)

  -- | This should be unapplied type (without type arguments)

  construct :: t -> GPUExp (GPURepTy t)

  -- constructRep :: GPUExp (GPURepTy t) -> GPUExp t
  -- constructRep = ConstructRep
  -- {-# INLINABLE constructRep #-}

  -- NOTE: The INLINABLE pragmas make the unfoldings available to the Core
  -- plugin

  default construct :: (Generic t, GPURep (Rep t Void), GPURep (GPURepTy t), GPURepTy (Rep t Void) ~ GPURepTy t) => t -> GPUExp (GPURepTy t)
  construct = construct . (from :: t -> Rep t Void)
  {-# INLINABLE construct #-}

  -- default constructRep :: (Generic t, GPURep (Rep t Void), GPURepTy t ~ GPURepTy (Rep t Void)) => GPUExp (GPURepTy t) -> GPUExp t
  -- constructRep = undefined

  rep :: t -> GPUExp t

  -- default rep :: (Generic t, GPURepTy t ~ GPURepTy (Rep t Void), GPURep (Rep t Void), GPURep (GPURepTy t)) => t -> GPUExp t
  -- rep x = -- constructRep . construct
  --   let x' = from x :: Rep t Void
  --       y = rep x'
  --   in _
  default rep :: (GPURep (GPURepTy t)) => t -> GPUExp t
  rep x = ConstructRep (construct x :: GPUExp (GPURepTy t))
  {-# INLINABLE rep #-}

  rep' :: t -> GPURepTy t

  default rep' :: (Generic t, GPURepTy t ~ GPURepTy (Rep t Void), GPURep (Rep t Void))
    => t -> GPURepTy t
  rep' = rep' . (from :: t -> Rep t Void)
  {-# INLINABLE rep' #-}

  unrep' :: GPURepTy t -> t

  default unrep' :: (Generic t, GPURepTy t ~ GPURepTy (Rep t Void), GPURep (Rep t Void))
    => GPURepTy t -> t
  unrep' = (to :: Rep t Void -> t) . unrep'
  {-# INLINABLE unrep' #-}

-- TODO: Does this work for what we need (dealing with constructors for
-- user defined types)?

-- construct :: forall a x. (GPURep a, GPURep x, GPURepTy a ~ x) => a -> GPUExp a
-- construct = _ . rep . rep'

-- construct :: a -> GPUExp a
-- construct = Repped . rep'

-- construct :: (GPURep a, GPURep (GPURepTy a)) => GPUExp a -> GPUExp (GPURepTy a)
-- construct (Repped e) = rep e
-- construct _ = error "construct: Given non-Repped argument"

-- constructAp' :: forall a b. (GPURep a) => GPUExp (a -> b) -> GPUExp a -> GPUExp b
-- constructAp' f x = _

instance GPURep Int where
  type GPURepTy Int = Int
  construct = rep --Lit
  rep = Lit
  rep' = id
  unrep' = id
  -- constructRep = id
instance GPURep Integer where
  type GPURepTy Integer = Integer
  construct = rep --Lit
  rep = Lit
  rep' = id
  unrep' = id
  -- constructRep = id
instance GPURep Float where
  type GPURepTy Float = Float
  construct = rep --Lit
  rep = Lit
  rep' = id
  unrep' = id
  -- constructRep = id
instance GPURep Double where
  type GPURepTy Double = Double
  construct = rep --Lit
  rep = Lit
  rep' = id
  unrep' = id
  -- constructRep = id

instance GPURep Bool where
  type GPURepTy Bool = Bool

  construct = rep

  -- construct True = TrueExp
  -- construct False = FalseExp

  rep False = FalseExp
  rep True  = TrueExp
  rep' = id
  unrep' = id
  -- constructRep = id

instance GPURep Char where
  type GPURepTy Char = Char

  construct = rep --CharLit

  rep c = CharLit c
  rep' = id
  unrep' = id
  -- constructRep = id

instance GPURep a => GPURep (Complex a) where

instance GPURep a => GPURep (Maybe a) where

instance (GPURep a, GPURep b) => GPURep (Either a b) where
  -- type GPURepTy (Either a b) = Either (GPURepTy a) (GPURepTy b)
  type GPURepTy (Either a b) = Either a b

  construct (Left x) = LeftExp (rep x)
  construct (Right y) = RightExp (rep y)

  rep (Left x) = LeftExp (rep x)
  rep (Right y) = RightExp (rep y)

  rep' = id
  unrep' = id

instance (GPURep a, GPURep b) => GPURep (a, b) where
  -- type GPURepTy (a, b) =  (GPURepTy a, GPURepTy b)
  type GPURepTy (a, b) =  (a, b)

  construct (x, y) = PairExp (rep x) (rep y)

  rep (x, y) = PairExp (rep x) (rep y) --uncurry PairExp
  rep' = id
  unrep' = id

instance (GPURep a, GPURep b, GPURep c) => GPURep (a, b, c) where
instance (GPURep a, GPURep b, GPURep c, GPURep d) => GPURep (a, b, c, d) where

-- XXX: Should this instance exist?
instance (GPURep a, GPURep b) => GPURep (Iter a b) where

-- XXX: Should this instance exist?
instance (GPURep a, GPURep b) => GPURep (a -> b) where
  type GPURepTy (a -> b) = GPURepTy a -> GPURepTy b

  construct = error "construct for (a->b) instance called"

  rep = Construct
  rep' f x = rep' (f (unrep' x))
  unrep' = error "unrep' for (a->b) instance called"
  -- constructRep = ConstructRep

-- Generics instances
instance (Typeable (M1 i c f p), GPURep (f p), GPURep (GPURepTy (f p))) => GPURep (M1 i c f p) where
  type GPURepTy (M1 i c f p) = GPURepTy (f p)

  construct :: M1 i c f p -> GPUExp (GPURepTy (M1 i c f p))
  construct = construct . unM1

  rep' :: M1 i c f p -> GPURepTy (M1 i c f p)
  rep' (M1 x) = rep' x

  unrep' :: GPURepTy (M1 i c f p) -> M1 i c f p
  unrep' = M1 . unrep'

  rep (M1 x) = ConstructRep (construct x)

instance (Typeable ((p :+: q) x), GPURep (p x), GPURep (q x), GPURep (GPURepTy (p x)), GPURep (GPURepTy (q x))) => GPURep ((p :+: q) x) where
  type GPURepTy ((p :+: q) x) =  Either (GPURepTy (p x)) (GPURepTy (q x))

  construct (L1 x) = LeftExp (construct x)
  construct (R1 x) = RightExp (construct x)

  rep' (L1 x) = Left (rep' x)
  rep' (R1 y) = Right (rep' y)

  unrep' (Left x) = L1 (unrep' x)
  unrep' (Right y) = R1 (unrep' y)

  rep (L1 x) = ConstructRep (LeftExp (construct x))
  rep (R1 x) = ConstructRep (RightExp (construct x))

  -- constructRep (LeftExp x) = _

instance (Typeable ((p :*: q) x), GPURep (p x), GPURep (q x), GPURep (GPURepTy (p x)), GPURep (GPURepTy (q x))) => GPURep ((p :*: q) x) where
  type GPURepTy ((p :*: q) x) =  (GPURepTy (p x), GPURepTy (q x))

  construct (x :*: y) = PairExp (construct x) (construct y)

  rep' (x :*: y) = (rep' x, rep' y)
  unrep' (x, y) = (unrep' x :*: unrep' y)

  rep (x :*: y) = ConstructRep (PairExp (construct x) (construct y))

instance (Typeable (K1 i c p), GPURep c) => GPURep (K1 i c p) where
  type GPURepTy (K1 i c p) = c --GPURepTy c

  -- construct = construct . unK1

  construct = rep . unK1

  -- rep = _

  rep' (K1 x) = x
  unrep' = K1

instance Typeable (U1 p) => GPURep (U1 p) where
  type GPURepTy (U1 p) = ()

  construct _ = UnitExp

  rep' U1 = ()
  unrep' () = U1

instance GPURep () where
  type GPURepTy () = ()

  rep () = UnitExp

  rep' = id
  unrep' = id

class GenericRep repr where
    type GenericRepTy repr

    genericRep :: repr -> GPUExp (GenericRepTy repr)

    genericRep' :: repr -> GenericRepTy repr
    genericUnrep' :: GenericRepTy repr -> repr

instance forall p q. (GPURep (p Void), GPURep (q Void), GenericRep (p Void), GenericRep (q Void), GPURepTy (p Void) ~ GenericRepTy (p Void), GPURepTy (q Void) ~ GenericRepTy (q Void)) =>
  GenericRep ((p :+: q) Void) where

    type GenericRepTy ((p :+: q) Void) = Either (GPURepTy (p Void)) (GPURepTy (q Void))

    genericRep (L1 x) = LeftExp (genericRep x)
    genericRep (R1 x) = RightExp (genericRep x)

    genericRep' = bimap rep' rep' . toCanonical
    genericUnrep' = fromCanonical . bimap unrep' unrep'

instance (GPURep (GenericRepTy (p Void)), GPURep (GenericRepTy (q Void)), GPURep (p Void), GPURep (q Void), GenericRep (p Void), GenericRep (q Void), GPURepTy (p Void) ~ GenericRepTy (p Void), GPURepTy (q Void) ~ GenericRepTy (q Void)) =>
  GenericRep ((p :*: q) Void) where

    type GenericRepTy ((p :*: q) Void) = (GPURepTy (p Void), GPURepTy (q Void))

    genericRep (x :*: y) = PairExp (genericRep x) (genericRep y)

    genericRep' = bimap rep' rep' . toCanonical
    genericUnrep' = fromCanonical . bimap unrep' unrep'

instance (GenericRep (f Void)) =>
  GenericRep (M1 i c f Void) where
    type GenericRepTy (M1 i c f Void) = GenericRepTy (f Void)

    genericRep = genericRep . unM1

    genericRep' = genericRep' . unM1
    genericUnrep' = M1 . genericUnrep'

instance GenericRep (U1 x) where
    type GenericRepTy (U1 x) = ()

    genericRep _ = UnitExp

    genericRep' _ = ()
    genericUnrep' () = U1

instance GenericRep (K1 i c f) where
    type GenericRepTy (K1 i c f) = c

    genericRep = error "genericRep called for K1 instance"

    genericRep' = unK1
    genericUnrep' = K1

the :: a -> a
the = id

the_repr :: GPUExp a -> GPUExp a
the_repr = id

-- coerceWith :: Dict (Coercible a b) -> a -> b
-- coerceWith Dict x = coerce x

-- Should this just produce an error?
sumMatchAbs :: forall s t. (GPURepTy (GPURepTy s) ~ GPURepTy s, GPURep s) => Env -> GPUExp (SumMatch (GPURepTy s) t) -> s -> t
-- sumMatchAbs env (CastExp d e) = sumMatchAbs env (coerceWith d e)
sumMatchAbs env (SumMatchExp p q) =
  \x0 ->
    let x = rep' x0
    in
    case x of
      Left  a -> prodMatchAbs env p a
      Right b -> sumMatchAbs env q b
sumMatchAbs env EmptyMatch = \_ -> error "Non-exhaustive match"
sumMatchAbs env (OneSumMatch f) =
    prodMatchAbs env f . rep'

prodMatchAbs :: GPURep s => Env -> GPUExp (ProdMatch s t) -> s -> t
-- prodMatchAbs env (CastExp e) = prodMatchAbs env (coerce e)
prodMatchAbs env (ProdMatchExp f) =
  \pair ->
    case pair of
      (x, y) -> runProdMatch (gpuAbsEnv env f x) y --prodMatchAbs (f (rep x)) y

prodMatchAbs env (OneProdMatch f) = \x -> gpuAbsEnv env f x --gpuAbs (f (rep x))
prodMatchAbs env (NullaryMatch x) = \_ -> gpuAbsEnv env x

gpuAbs :: forall t. GPUExp t -> t
gpuAbs = gpuAbsEnv emptyEnv

gpuAbsEnv :: forall t. Env -> GPUExp t -> t
gpuAbsEnv env (Var v) =
  case envLookup env v of
    Just x -> gpuAbsEnv env x
    Nothing -> error $ "No binding for name " ++ show v
gpuAbsEnv env (CaseExp x f) = sumMatchAbs env f (gpuAbsEnv env x)
gpuAbsEnv env FalseExp = False
gpuAbsEnv env TrueExp  = True
gpuAbsEnv env UnitExp  = ()
gpuAbsEnv env (Repped x) = unrep' x
-- gpuAbs (Lam f) = gpuAbs . f . rep
-- gpuAbs (Apply f x) = gpuAbs f (gpuAbs x)
gpuAbsEnv env (Lit x) = x
gpuAbsEnv env (Ord x) = ord (gpuAbsEnv env x)
gpuAbsEnv env (Add x y) = gpuAbsEnv env x + gpuAbsEnv env y
gpuAbsEnv env (Sub x y) = gpuAbsEnv env x - gpuAbsEnv env y
gpuAbsEnv env (Mul x y) = gpuAbsEnv env x * gpuAbsEnv env y
gpuAbsEnv env (FromEnum x) = fromEnum (gpuAbsEnv env x)
gpuAbsEnv env (FromIntegral x) = fromIntegral (gpuAbsEnv env x)
gpuAbsEnv env (Sqrt x) = sqrt (gpuAbsEnv env x)
gpuAbsEnv env (Equal x y) = gpuAbsEnv env x == gpuAbsEnv env y
gpuAbsEnv env (Lte x y) = gpuAbsEnv env x <= gpuAbsEnv env y
gpuAbsEnv env (Gt x y) = gpuAbsEnv env x > gpuAbsEnv env y
gpuAbsEnv env (Not x) = not (gpuAbsEnv env x)
gpuAbsEnv env (LeftExp x) = Left (gpuAbsEnv env x)
gpuAbsEnv env (RightExp y) = Right (gpuAbsEnv env y)
gpuAbsEnv env (PairExp x y) = (gpuAbsEnv env x, gpuAbsEnv env y)
gpuAbsEnv env (StepExp x) = Step $ gpuAbsEnv env x
gpuAbsEnv env (DoneExp y) = Done $ gpuAbsEnv env y
gpuAbsEnv env (TailRec f) = \x ->
  case gpuAbsEnv env f x of
    Step x' -> gpuAbsEnv env (TailRec f) x'
    Done r  -> r
gpuAbsEnv env (IfThenElse cond t f)
  | gpuAbsEnv env cond = gpuAbsEnv env t
  | otherwise = gpuAbsEnv env f
gpuAbsEnv env (Construct x) = x
gpuAbsEnv env (ConstructAp f x) = gpuAbsEnv env f (gpuAbsEnv env x)

gpuAbsEnv env e@(SumMatchExp {}) = MkSumMatch (sumMatchAbs env e)
gpuAbsEnv env e@(OneSumMatch {}) = MkSumMatch (sumMatchAbs env e)
gpuAbsEnv env EmptyMatch = MkSumMatch absurd

gpuAbsEnv env e@(ProdMatchExp {}) = MkProdMatch (prodMatchAbs env e)
gpuAbsEnv env e@(OneProdMatch {}) = MkProdMatch (prodMatchAbs env e)
gpuAbsEnv env e@(NullaryMatch {}) = MkProdMatch (prodMatchAbs env e)

gpuAbsEnv env (FstExp x) = fst (gpuAbsEnv env x)
gpuAbsEnv env (SndExp x) = snd (gpuAbsEnv env x)

gpuAbsEnv env (ConstructRep x) = unrep' $ gpuAbsEnv env x

gpuAbsEnv env (Lam (name :: Name a) (body :: GPUExp b)) = \(arg :: a) ->

    let go0 :: forall x. GPUExp x -> GPUExp x
        go0 expr@(Var name2 :: GPUExp t') =
          case namesEq name name2 of
            Just Refl -> rep arg
            _ ->
              case envLookup env name2 of
                Just v  -> v
                Nothing -> expr
        go0 expr = expr

        go = transformE go0
    in

    gpuAbsEnv (extendEnv env (name :=> rep arg)) $ go body

gpuAbsEnv env (App f x) = gpuAbsEnv env f (gpuAbsEnv env x)
gpuAbsEnv env (CharLit c) = c
gpuAbsEnv env (CastExp x) = gpuAbsEnv env (coerce x)

class Canonical t where
  type GenericOp t :: (* -> *) -> (* -> *) -> * -> *

  toCanonical :: GenericOp t p q x -> t (p x) (q x)
  fromCanonical :: t (p x) (q x) -> GenericOp t p q x

instance Canonical Either where
  type GenericOp Either = (:+:)

  toCanonical (L1 x) = Left x
  toCanonical (R1 y) = Right y

  fromCanonical (Left x) = L1 x
  fromCanonical (Right y) = R1 y

instance Canonical (,) where
  type GenericOp (,) = (:*:)

  toCanonical (x :*: y) = (x, y)
  fromCanonical (x, y) = x :*: y


text:
{-# LANGUAGE TemplateHaskell #-}

import Test.QuickCheck
import System.Exit
import Data.Complex as C
import FFT

-- Helpers

allClose :: (Ord a, Fractional a) => [a] -> [a] -> Bool
allClose xs ys = and $ fmap (\diff -> diff < 1e-4) diffs
  where
    diffs = [abs $ x - y | (x,y) <- zip xs ys]

allCloseComplex :: (RealFloat a) => [Complex a] -> [Complex a] -> Bool
allCloseComplex xs ys = and $ fmap (\diff -> diff < 1e-4) diffs
  where
    diffs = [C.magnitude $ x - y | (x,y) <- zip xs ys]

-- Properties

prop_dftInverse :: [Complex Double] -> Bool
prop_dftInverse xs = allCloseComplex xs $ idft (dft xs)

prop_rdftInverse :: [Double] -> Bool
prop_rdftInverse xs = allClose xs $ irdft (rdft xs)

prop_fftInverse :: [Complex Double] -> Bool
prop_fftInverse xs = allCloseComplex xs $ ifft (fft xs)

prop_rfftInverse :: [Double] -> Bool
prop_rfftInverse xs = allClose xs $ irfft (rfft xs)

prop_fft :: [Complex Double] -> Bool
prop_fft xs = allCloseComplex (fft xs) (dft xs)

prop_rfft :: [Double] -> Bool
prop_rfft xs = allCloseComplex (rfft xs) (rdft xs)

-- https://hackage.haskell.org/package/QuickCheck-2.13.2/docs/Test-QuickCheck-All.html
return []
runTests :: IO Bool
runTests = $quickCheckAll

main :: IO ()
main = do
  success <- and <$> sequence [runTests]
  if success then exitSuccess else exitFailure

max_stars_repo_name: viswanathgs/haskell-fft

max_issues_repo_name: viswanathgs/haskell-fft

max_forks_repo_name: viswanathgs/haskell-fft


text:
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
module Main where

import Lib
import Numeric.Vector.Sized
import Numeric.Layer.Affine
import Numeric.Network
import Numeric.Layer.Affine
import Numeric.Layer.Sigmoid
import Numeric.Loss.MeanSquaredError
import qualified Data.Array.Accelerate.Interpreter as AI
import qualified Data.Array.Accelerate.LLVM.Native as Native

v :: SizedVector 3
v = use (fromList [1, 2, 3]) .+. use (fromList [2, 3, 4])

xs :: SizedMatrix 2 3
xs = use (fromList [1 .. 6])
-- 1 2 3
-- 4 5 6

ys :: SizedMatrix 3 4
ys = use (fromList [1 .. 12])
-- 1 2 3 4
-- 5 6 7 8
-- 9 10 11 12
--
u = use (fromList [1, 2]) :: SizedVector 2
zs = use (fromList [1, 2, 3, 4]) :: SizedMatrix 2 2

aff = affine xs v u

net :: Network '[Affine 3 2, Sigmoid 2, Affine 2 2, Sigmoid 2]
net = Affine ws b :~> Sigmoid () :~> Affine ws' b' :~> Last (Sigmoid ())
  where
    ws  = fromList [1, 2, 3, 4, 5, 6]
    b   = fromList [1, 2]
    ws' = fromList [-2, 2, 0, 1]
    b'  = fromList [0, 1]

main :: IO ()
main = do
  print net'
  where
    net' = train
      net
      (MeanSquaredError () :: MeanSquaredError 2)
      (use $ fromList [-1, -2, -3])
      (use $ fromList [1, 2])

  -- print (runUsing Native.run v)
  -- print (runUsing Native.run (xs >< ys))
  -- print (runUsing Native.run (xs #> v))
  -- print (runUsing Native.run (u <# zs))
  -- print (runUsing Native.run aff)

max_stars_repo_name: mixed-signals/mixed-signals

max_issues_repo_name: mixed-signals/mixed-signals

max_forks_repo_name: mixed-signals/mixed-signals


text:
module Strategy.TestOrientation where

import Prelude hiding (Right, Left)
import Numeric.LinearAlgebra ((<>))
import qualified Control.Monad.Writer as Writer
import qualified Data.DList as DList

import Cube (Side(..), Color(..))
import qualified Cube as Cube
import qualified Rotation as Rotation
import qualified RotationPath as RotationPath
import qualified Strategy.Orientation as Orientation

testSides :: Bool
testSides =
  let expected = [
          (Top, White)
        , (Front, Green)
        , (Right, Red)
        , (Back, Blue)
        , (Left, Orange)
        , (Bottom, Yellow)
        ]
      actual = Orientation.sides Cube.solvedCube
  in expected == actual

testIsColorOnSide :: Bool
testIsColorOnSide =
  let t = Orientation.isColorOnSide White Top Cube.solvedCube
      f = Orientation.isColorOnSide Blue Bottom Cube.solvedCube
  in t == True && f == False

testForColorOnTop :: Bool
testForColorOnTop = and $ do
  color <- [Cube.White ..]
  let (prefix, suffix) = Orientation.forColorOnTop color Cube.solvedCube
      cube = Rotation.rotate prefix Cube.solvedCube
      isOnTop = Orientation.isColorOnSide color Cube.Top cube
      isIdentity = (suffix <> prefix) == Rotation.identity
  return $ isIdentity && isOnTop

testPutColorOnTop :: Bool
testPutColorOnTop =
  let expectedSides = [
          (Top, Yellow)
        , (Front, Blue)
        , (Right, Red)
        , (Back, Green)
        , (Left, Orange)
        , (Bottom, White)
        ]
      expectedRotations = [Rotation.topToFront <> Rotation.topToFront]
      (cube, rotationLog) = Writer.runWriter $
        Orientation.putColorOnTop Yellow Cube.solvedCube
      actualSides = Orientation.sides cube
  in actualSides == expectedSides && expectedRotations == DList.toList rotationLog

testWithColorOnTop :: Bool
testWithColorOnTop =
  let (cube, rotationLog) = Writer.runWriter $
        Orientation.withColorOnTop Green (RotationPath.rotate Rotation.topL) Cube.solvedCube
      [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expectedCube = [
          (top,0,0,w),(top,0,1,w),(top,0,2,w),(top,1,0,w),(top,1,1,w),(top,1,2,w),(top,2,0,r),(top,2,1,r),(top,2,2,r)
        , (frn,0,0,g),(frn,0,1,g),(frn,0,2,g),(frn,1,0,g),(frn,1,1,g),(frn,1,2,g),(frn,2,0,g),(frn,2,1,g),(frn,2,2,g)
        , (rgt,0,0,y),(rgt,0,1,r),(rgt,0,2,r),(rgt,1,0,y),(rgt,1,1,r),(rgt,1,2,r),(rgt,2,0,y),(rgt,2,1,r),(rgt,2,2,r)
        , (bck,0,0,b),(bck,0,1,b),(bck,0,2,b),(bck,1,0,b),(bck,1,1,b),(bck,1,2,b),(bck,2,0,b),(bck,2,1,b),(bck,2,2,b)
        , (lft,0,0,o),(lft,0,1,o),(lft,0,2,w),(lft,1,0,o),(lft,1,1,o),(lft,1,2,w),(lft,2,0,o),(lft,2,1,o),(lft,2,2,w)
        , (bot,0,0,o),(bot,0,1,o),(bot,0,2,o),(bot,1,0,y),(bot,1,1,y),(bot,1,2,y),(bot,2,0,y),(bot,2,1,y),(bot,2,2,y)]
      expectedRotations = [Rotation.frontL]
  in cube == expectedCube && DList.toList rotationLog == expectedRotations

tests :: [(String, Bool)]
tests = [
    ("Strategy.Orientation.sides", testSides)
  , ("Strategy.Orientation.isColorOnSide", testIsColorOnSide)
  , ("Strategy.Orientation.forColorOnTop", testForColorOnTop)
  , ("Strategy.Orientation.putColorOnTop", testPutColorOnTop)
  , ("Strategy.Orientation.withColorOnTop", testWithColorOnTop)
  ]

max_stars_repo_name: runjak/hRubiks

max_issues_repo_name: runjak/hRubiks

max_forks_repo_name: runjak/hRubiks


text:
module Fractal.Math.Mandelbrot(
  mandel,
  mandel'
)where
{-

Math packet kommt alles mathe bezogene rein

-}

import Data.Complex
import Fractal.Math.Kernel (evaluateFractal)

type Point = Complex Double
type Limit = Double
type Iteration = Complex Double
type Magnitude = Double
type MaxIter = Int

mandelEquation :: Point -> Iteration -> Iteration
mandelEquation c (0.0 :+ 0.0) = c
mandelEquation c (0.0 :+ 1) = 0
mandelEquation c (0.0 :+ (-1)) = 0
mandelEquation c z = z * z + c

mandel :: MaxIter -> Limit -> [Point] -> [Int]
mandel maxIter l field = map mandelEvaluator field
  where mandelEvaluator = evaluateFractal mandelEquation l maxIter

mandel' :: MaxIter -> Limit -> Point -> Int
mandel' maxIter l field = mandelEvaluator field
  where mandelEvaluator = evaluateFractal mandelEquation l maxIter

max_stars_repo_name: sigmaticsMUC/FractalHaskell

max_issues_repo_name: sigmaticsMUC/FractalHaskell

max_forks_repo_name: sigmaticsMUC/FractalHaskell


text:
{-# LANGUAGE TemplateHaskell #-}
module Turtle (main) where
import Ros.Node
import Ros.Topic (cons, repeatM)
import Ros.Topic.Util (tee, filterBy, everyNew, interruptible, gate, share)
import Ros.Turtlesim.Pose
import Ros.Geometry_msgs.Twist
import qualified Ros.Geometry_msgs.Vector3 as V
import Ros.Logging

import System.IO (hFlush, stdout)
import Data.Default.Generics (def)
import Control.Applicative ((<$>))
import Control.Arrow ((&&&))
import Lens.Family ((.~), (^.), (&))
import Data.Complex

-- A type synonym for a 2D point.
type Point = Complex Float

-- A Topic of user-supplied waypoint trajectories.
getTraj :: Topic IO [Point]
getTraj = repeatM (do putStr "Enter waypoints: " >> hFlush stdout
                      $(logInfo "Waiting for new traj")
                      (map (uncurry (:+)) . read) `fmap` getLine)

-- Create Twist message with linear and angular velocity in arguments.
mkTwist :: Real a => a -> a -> Twist
mkTwist l a = def & linear  . V.x .~ realToFrac l
                  & angular . V.z .~ realToFrac a

wrapAngle theta
  | theta < 0    = theta + 2 * pi
  | theta > 2*pi = theta - 2 * pi
  | otherwise    = theta

angleDiff x y = wrapAngle (x' + pi - y') - pi
  where x' = if x < 0 then x + 2*pi else x
        y' = if y < 0 then y + 2*pi else y

-- Produce a unit value every time a goal is reached.
arrivalTrigger :: Topic IO Point -> Topic IO Pose -> Topic IO ()
arrivalTrigger goals poses = const () <$>
                             filterBy (arrived <$> goals) (p2v <$> poses)
  where arrived goal pose = magnitude (goal - pose) < 1.5
        p2v (Pose x y _ _ _) = x :+ y

-- Navigate to a goal given a current pose estimate.
navigate :: (Point, Pose) -> Twist
navigate (goal, pos) = mkTwist (min 2 (magnitude v)) angVel
  where v        = goal - (pos ^. x :+ pos ^. y)
        thetaErr = (angleDiff (phase v) (pos ^. theta)) * (180 / pi)
        angVel   = signum thetaErr * (min 2 (abs thetaErr))

main = runNode "HaskellBTurtle" $
       do enableLogging (Just Warn)
          poses <- subscribe "/turtle1/pose"
          (t1,t2) <- liftIO . tee . interruptible $ getTraj
          let goals = gate t1 (cons () (arrivalTrigger t2 poses))
          advertise "/turtle1/cmd_vel" $
                    fmap navigate (everyNew goals poses)

max_stars_repo_name: hpacheco/roshask

max_issues_repo_name: hpacheco/roshask

max_forks_repo_name: hpacheco/roshask


text:
module Syntax
  ( LispVal (..)
  , unwordsList
  , ThrowsError
  , trapError
  , extractValue,
  LispError (..)
  ) where
import           Data.Complex (Complex (..))
import Text.ParserCombinators.Parsec
import Control.Monad.Except

instance Show LispVal where
  show (Atom s)          = s
  show (List xs)         = "(" ++ unwordsList xs ++ ")"
  show (DottedList xs x) = "(" ++ unwordsList xs ++ " . " ++ show x ++ ")"
  show (Number x)        = show x
  show (String s)        = "\"" ++ s ++ "\""
  show (Bool True)       = "#t"
  show (Bool False)      = "#f"
  show (Character c)     = [c]
  show (Float x)         = show x
  show (Complex x)       = show x
  show (Rational x)      = show x

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Int
             | Bool Bool
             | Character Char
             | String String
             | Float Double
             | Complex (Complex Double)
             | Rational Rational
             | PrimitiveFunc [LispVal -> ThrowsError LispVal]

unwordsList :: [LispVal] -> String
unwordsList = unwords . map show

-- Error
type ThrowsError = Either LispError

data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

instance Show LispError where
  show (NumArgs expected found) = "Expected " ++ show expected ++ " args; found values " ++ unwordsList found
  show (TypeMismatch expected found) = "Invalid type: expected " ++ expected ++ ", found " ++ show found
  show (Parser parseErr) = "Parse error at " ++ show parseErr
  show (BadSpecialForm message form) = message ++ ": " ++ show form
  show (NotFunction message func) = message ++ ": " ++ show func
  show (UnboundVar message varname) = message ++ ": " ++ varname
  show (Default message) = message

trapError :: (Show a, MonadError a m) => m String -> m String
trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val

max_stars_repo_name: JuanBono/my-scheme

max_issues_repo_name: JuanBono/my-scheme

max_forks_repo_name: JuanBono/my-scheme


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

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

module Feldspar.Compiler.Imperative.FromCore.Literal
  (literal)
  where

import Data.Complex

import Language.Syntactic

import Feldspar.Core.Types as Core
import Feldspar.Core.Interpretation
import Feldspar.Core.Constructs.Literal

import Feldspar.Compiler.Imperative.Frontend
import Feldspar.Compiler.Imperative.Representation (Expression(..),Constant(..))
import qualified Feldspar.Compiler.Imperative.Representation as Rep (Type(..),
                                                                     Signedness(..),
                                                                     Size(..))
import Feldspar.Compiler.Imperative.FromCore.Interpretation

instance Compile (Literal :|| Core.Type) dom
  where
    compileExprSym (C' (Literal a)) info Nil = literal (infoType info) (infoSize info) a

    compileProgSym (C' (Literal a)) info loc Nil = case loc of
        Just l  -> literalLoc l (infoType info) (infoSize info) a
        Nothing -> return ()

literal :: TypeRep a -> Size a -> a -> CodeWriter (Expression ())
literal t@UnitType        sz a = return (ConstExpr $ literalConst t sz a)
literal t@BoolType        sz a = return (ConstExpr $ literalConst t sz a)
literal t@IntType{}       sz a = return (ConstExpr $ literalConst t sz a)
literal t@FloatType       sz a = return (ConstExpr $ literalConst t sz a)
literal t@DoubleType      sz a = return (ConstExpr $ literalConst t sz a)
literal t@ComplexType{}   sz a = return (ConstExpr $ literalConst t sz a)
literal t@ArrayType{}     sz a = return (ConstExpr $ literalConst t sz a)
literal t s a = do loc <- freshVar "x" t s
                   literalLoc loc t s a
                   return loc

literalConst :: TypeRep a -> Size a -> a -> Constant ()
literalConst UnitType        _  ()     = IntConst 0 (Rep.NumType Rep.Unsigned Rep.S32)
literalConst BoolType        _  a      = BoolConst a
literalConst trep@IntType{}  sz a      = IntConst (toInteger a) (compileTypeRep trep sz)
literalConst FloatType       _  a      = FloatConst a
literalConst DoubleType      _  a      = DoubleConst a
literalConst (ArrayType t)   _  a      = ArrayConst $ map (literalConst t (defaultSize t)) a
literalConst (ComplexType t) _  (r:+i) = ComplexConst re ie
  where re = literalConst t (defaultSize t) r
        ie = literalConst t (defaultSize t) i

literalLoc :: Expression () -> TypeRep a -> Size a -> a -> CodeWriter ()
literalLoc loc arr@ArrayType{} sz e
    = tellProg [copyProg (Just loc) [ConstExpr $ literalConst arr sz e]]

literalLoc loc (Tup2Type ta tb) (sa,sb) (a,b) =
    do literalLoc (StructField loc "member1") ta sa a
       literalLoc (StructField loc "member2") tb sb b

literalLoc loc (Tup3Type ta tb tc) (sa,sb,sc) (a,b,c) =
    do literalLoc (StructField loc "member1") ta sa a
       literalLoc (StructField loc "member2") tb sb b
       literalLoc (StructField loc "member3") tc sc c

literalLoc loc (Tup4Type ta tb tc td) (sa,sb,sc,sd) (a,b,c,d) =
    do literalLoc (StructField loc "member1") ta sa a
       literalLoc (StructField loc "member2") tb sb b
       literalLoc (StructField loc "member3") tc sc c
       literalLoc (StructField loc "member4") td sd d

literalLoc loc (Tup5Type ta tb tc td te) (sa,sb,sc,sd,se) (a,b,c,d,e) =
    do literalLoc (StructField loc "member1") ta sa a
       literalLoc (StructField loc "member2") tb sb b
       literalLoc (StructField loc "member3") tc sc c
       literalLoc (StructField loc "member4") td sd d
       literalLoc (StructField loc "member5") te se e

literalLoc loc (Tup6Type ta tb tc td te tf) (sa,sb,sc,sd,se,sf) (a,b,c,d,e,f) =
    do literalLoc (StructField loc "member1") ta sa a
       literalLoc (StructField loc "member2") tb sb b
       literalLoc (StructField loc "member3") tc sc c
       literalLoc (StructField loc "member4") td sd d
       literalLoc (StructField loc "member5") te se e
       literalLoc (StructField loc "member6") tf sf f

literalLoc loc (Tup7Type ta tb tc td te tf tg) (sa,sb,sc,sd,se,sf,sg) (a,b,c,d,e,f,g) =
    do literalLoc (StructField loc "member1") ta sa a
       literalLoc (StructField loc "member2") tb sb b
       literalLoc (StructField loc "member3") tc sc c
       literalLoc (StructField loc "member4") td sd d
       literalLoc (StructField loc "member5") te se e
       literalLoc (StructField loc "member6") tf sf f
       literalLoc (StructField loc "member7") tg sg g

literalLoc loc t sz a =
    do rhs <- literal t sz a
       assign (Just loc) rhs

max_stars_repo_name: kffaxen/feldspar-compiler

max_issues_repo_name: kffaxen/feldspar-compiler

max_forks_repo_name: kffaxen/feldspar-compiler


text:
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeFamilies              #-}

import           Control.Applicative
import           Control.Concurrent
import           Control.Monad.Random

import           Options.Applicative

import           Graphics.Image               hiding (map, on)

import qualified Numeric.LinearAlgebra        as LA
import qualified Numeric.LinearAlgebra.Static as H

import           Grenade

data SuperResOptions = SuperResOptions FilePath         -- onnx file
                                       FilePath         -- input image

pSuperRes :: Parser SuperResOptions
pSuperRes = SuperResOptions <$> argument str (metavar "onnx") <*> argument str (metavar "image")

loadSuperResImage :: FilePath -> IO (Maybe (S ('D3 224 224 1), [[RealNum]], [[RealNum]]))
loadSuperResImage path = do
  img <- readImageRGB VU path
  -- displayImage img
  return $ do
    guard $ dims img == (224, 224)
    let imgYCbCr = toImageYCbCr img
        imgY0    = map (\(PixelYCbCr y _ _ )      -> doubleToRealNum y )  . concat . toLists $ imgYCbCr
        imgCb    = map (map (\(PixelYCbCr _ cb _) -> doubleToRealNum cb)) . toLists $ imgYCbCr
        imgCr    = map (map (\(PixelYCbCr _ _ cr) -> doubleToRealNum cr)) . toLists $ imgYCbCr

    return (S3D (H.fromList imgY0), imgCb, imgCr)

displayHighResImage :: S ('D3 672 672 1) -> [[RealNum]] -> [[RealNum]] -> IO ()
displayHighResImage (S3D m) cbs crs = do
  let realNumToPixel = PixelX . realNumToDouble
      m'  = LA.toLists $ H.extract m              :: [[RealNum]]
      m'' = map (map realNumToPixel) m'           :: [[Pixel X Double]]
      img = fromLists m''                         :: Image VU X Double

      imgBs  = fromLists $ map (map realNumToPixel) cbs :: Image VU X Double
      imgRs  = fromLists $ map (map realNumToPixel) crs :: Image VU X Double

      imgBs' = resize Bilinear Edge (672, 672) imgBs :: Image VU X Double
      imgRs' = resize Bilinear Edge (672, 672) imgRs :: Image VU X Double

      finalImg = fromImagesX [(LumaYCbCr, img), (CBlueYCbCr, imgBs'), (CRedYCbCr, imgRs')] :: Image VU YCbCr Double

  displayImage finalImg
  threadDelay 10000000

main :: IO ()
main = do
    SuperResOptions netPath imgPath <- execParser (info (pSuperRes <**> helper) idm)
    res <- loadSuperResolution netPath

    inputM <- loadSuperResImage imgPath

    case (res, inputM) of
      (Right net, Just (input, cbs, crs))  -> do
        let S3D y = runNet net input
        displayHighResImage (S3D y) cbs crs

      (Left err, _) -> print err
      _             -> putStrLn "Failed to load network and file"

max_stars_repo_name: th-char/grenade

max_issues_repo_name: th-char/grenade

max_forks_repo_name: th-char/grenade