module Widget.Percentile
  ( Percentile()
  , mkPercentile
  ) where

import Import
import qualified Format.Percentile as P
import qualified Format.FastTime   as D

import Control.Monad (when)
import Statistics.Distribution.Empirical

data Percentile = Percentile (Zipper EmpiricalDistribution) (Zipper NominalDiffTime) (Zipper NominalDiffTime) (Zipper Bool) Window

instance Widget Percentile where
  update                        = updatePercentile
  handle                        = handlePercentile
  redraw                        = redrawPercentile
  window (Percentile _ _ _ _ w) = w

mkPercentile :: [(EmpiricalDistribution, NominalDiffTime)] -> Window -> Percentile
mkPercentile p = Percentile (listToZipper $ map (\(ed, _) -> ed) p)
                            (listToZipper $ map (\(_, c) -> c) p)
                            (listToZipper [toEnum 0])
                            (listToZipper [True])

updatePercentile :: NominalDiffTime -> Percentile -> Curses Percentile
updatePercentile deltaT (Percentile z bs ts v w) =
  let ts' = replace (curs ts + deltaT) ts
      p   = Percentile z bs ts' v w
  in redrawPercentile p >> return p

handlePercentile :: TimerAction -> Percentile -> Curses Percentile
handlePercentile Advance (Percentile z bs ts v w) =
  return $ Percentile (next z) (next bs) (push 0 ts) (push True v) w

handlePercentile Reverse p@(Percentile z bs ts v w) =
  return $
    if null $ left z
      then p
      else
        let c   = curs ts
            ts' = trunc $ prev ts
            v'  = if (not $ curs $ prev $ prev v) && (not $ null $ left $ prev v)
                    then trunc $ prev v
                   else replace True $ trunc $ prev v
        in Percentile (prev z) (prev bs) (replace (curs ts' + c) ts') v' w

handlePercentile Skip    (Percentile z bs ts v w) =
  return $ Percentile (next z) (next bs) (push (toEnum 0) ts) (push False $ replace False v) w

handlePercentile Reset   (Percentile z bs ts v w) =
  return $ mkPercentile (zipWith (\ed c -> (ed, c)) (zipperToList z) (zipperToList bs)) w

handlePercentile ValidInvalid (Percentile z bs ts v w) =
  return $ Percentile z bs ts (replace (not $ curs v) v) w

handlePercentile _ t = return t

redrawPercentile :: Percentile -> Curses ()
redrawPercentile p@(Percentile z bs t v w) = do
  cid <- newColorID ColorBlack ColorDefault 1

  updateWindow w $ do
    (rows, columns) <- windowSize

    erase

    moveCursor 0 0
    drawLineV Nothing rows

    let (lt, ct, rt) = sizeZipperToScreen rows t
        (lb, cb, rb) = sizeZipperToScreen rows bs
        (lv, cv, rv) = sizeZipperToScreen rows v
        (lz, cz, rz) = sizeZipperToScreen rows z

        t'  = (lt ++ ct:rt)
        bs' = (lb ++ cb:rb)
        v'  = (lv ++ cv:rv)
        distrs = (lz ++ cz:rz)

        percentileF distr best current =
          if current == 0
            then "----"
            else if current >= best && best > 0
                   then show $ P.percentile distr current
                   else show $ D.fastTime $ current - best
        percentileFs = zipWith3 percentileF distrs bs' t'

    mapM_ (go cid)
      $ zip3 [0..] v' percentileFs

  where go cid (a, valid, p) = do
          when valid $ do
            moveCursor a 1
            drawString p


{-# LANGUAGE GeneralizedNewtypeDeriving, FlexibleInstances #-}

module SMC where

import PPL ( Sample(sample), Infer(score, observe, logscore), Model, runModel, sampleWeighted, WeightedT )
import Randomness ( Random, Weighted, Weight, parallel )
import Control.Monad ( ap, replicateM, zipWithM )
import Data.Either ( partitionEithers )
import Particle ( ParticleT(..), reweight )
import Control.Monad.Trans.Class (lift)
import Data.Monoid (Product(..))
import Control.Monad.Trans.Writer (WriterT, tell, runWriterT)
import Numeric.Log (Log(..))
import Distribution ( Dist (samplefrom), logpdf, Uniform (Uniform) )
import Statistics.Resampling (Bootstrap(resamples))
import Control.Monad.ST (ST)
import Control.Parallel (pseq)
import Data.List ( partition, sortOn, groupBy )
import Debug.Trace ( trace )
import Control.Monad.RWS (gets)

type Address = Int

instance (Monoid w, Sample m) => Sample (ParticleT l w m) where
  sample = Compute . fmap return . sample

instance Sample m => Infer (ParticleT l (Product Weight) m) where
  score w = Done (Product $ Exp $ log w) ()
  logscore w = Done (Product w) ()
  observe d x = Done (Product $ logpdf d x) ()

-- SMC with seperate score and resample statements
type SMCModel = ParticleT () (Product Weight) Random
type AddressSMC = ParticleT Address (Product Weight) Random

avg :: [Weighted Double] -> Double
avg xs = let total = sum (map snd xs) in
         exp $ ln $ sum [Exp (log a) * w / total | (a, w) <- xs]

stats :: Functor f => f [Weighted (Double, Double) ] -> f Double
stats = fmap (avg . map (\((a,b),c) -> (a,c)))

weight :: ParticleT l (Product Weight) m a -> Weight
weight (Done w _) = getProduct w
weight (Resample l w _) = getProduct w
weight (Compute _) = error "Compute has no weight"

-- A stochastic monad is one that can perform some sort of stochastic compute step
class Monad m => Stochastic m  where
  compute :: m a -> Random (m a)

instance Stochastic (ParticleT l (Product Weight) Random) where
  -- run a computation until it reaches a resample or is done
  compute (Compute ma) = ma >>= compute
  compute p = return p

class Stochastic m => Resampleable m where
  resample :: m ()

instance Resampleable SMCModel where
  resample = Resample () mempty $ return ()

class Infer m => SMC m where
  resampleSMC :: Int -> [m a] -> Random [Weighted a]
  normalisingConstant :: Int -> [m a] -> Random Weight

class Infer m => SampleWeighted m where
  doResample :: [m a] -> Random [m a]

-- SMC where every resample is explicit
newtype ExplicitResample a = ExplicitResample { unexplicit :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, Infer, SMC)

-- SMC where every score triggers a resample
newtype AlwaysResample a = AlwaysResample { unalways :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, SMC)

-- SMC where weight zero particles are resampled
newtype AliveFilter a = AliveFilter { unalive :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, SMC)

-- SMC where subsets of particles are resampled based on addresses
newtype AlphaSMC a = AlphaSMC { unalpha :: AddressSMC a } deriving (Functor, Applicative, Monad, Sample, Infer)

-- SMC where subsets of particles are resampled based on addresses
newtype Adaptive a = Adaptive { unadapt :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable, Show)

-- SMC with Startified Resampling
newtype Stratified a = Stratified { unstrat :: SMCModel a } deriving (Functor, Applicative, Monad, Sample, Resampleable)

-- SMC with Subset Resampling
newtype Labelled a = Labelled { unlab :: AddressSMC a } deriving (Functor, Applicative, Monad, Sample, Infer)

-- these all have implicit resamples:
instance Infer AliveFilter where
  score w = AliveFilter (score w >> resample)
  observe d x = AliveFilter (observe d x >> resample)

instance Infer Adaptive where
  score w = Adaptive (score w >> resample)
  observe d x = Adaptive (observe d x >> resample)

instance Infer AlwaysResample where
  score w = AlwaysResample (score w >> resample)
  observe d x = AlwaysResample (observe d x >> resample)

instance Infer Stratified where
  score w = Stratified (score w >> resample)
  observe d x = Stratified (observe d x >> resample)

-- Generalized Newtype Deriving isn't quite enough to do these
instance Stochastic ExplicitResample where
  compute (ExplicitResample a) = ExplicitResample <$> compute a

instance Stochastic Labelled where
  compute (Labelled a) = Labelled <$> compute a

instance Stochastic AlwaysResample where
  compute (AlwaysResample a) = AlwaysResample <$> compute a

instance Stochastic Adaptive where
  compute (Adaptive a) = Adaptive <$> compute a

instance Stochastic Stratified where
  compute (Stratified a) = Stratified <$> compute a

resampleLab :: [Labelled a] -> Random [Labelled a]
resampleLab ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unlab) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return $ map Labelled ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- labResample $ map getweight ms'
        -- recursive call to continue inference
        resampleLab $ map Labelled resampled

labResample :: [Weighted (AddressSMC a)] -> Random [AddressSMC a]
labResample ms = let grouped = groupBy (\(x,_) (y,_) -> address x == address y) ms in
    fmap concat $ parallel $ map stratResample grouped

resampleStrat :: [Stratified a] -> Random [Stratified a]
resampleStrat ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unstrat) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return $ map Stratified ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- stratResample $ map getweight ms'
        -- recursive call to continue inference
        resampleStrat $ map Stratified resampled

getSamples :: [Weight] -> [Weighted a] -> [a]
getSamples [] _ = []
getSamples (r:rs) values = let (y:ys) = dropWhile (\(x,w) -> w < r) values
                           in fst y : getSamples rs (y:ys)

residual :: [Weighted a] -> ([a], [Weighted a])
residual [] = ([], [])
residual ((a, w) : rest) = let (xs, wxs) = residual rest
                           in let w' = floor w
                           in (replicate w' a ++ xs, (a, w - fromIntegral w') : wxs)

stratResample :: Functor m => [Weighted (ParticleT l (Product Weight) m a)] -> Random [ParticleT l (Product Weight) m a]
stratResample xs = do
    -- get sum of weights
    let weight = sum $ map snd xs
    let n = length xs
    let avgweight = weight / Exp (log $ fromIntegral n)
    -- residuals
    let (guaranteed, rest) = residual [(a, w / avgweight) | (a,w) <- xs]
    let n' = n - length guaranteed
    -- rest
    let weighted = let (as, ws) = unzip rest in zip as (scanl1 (+) ws)
    let avgrest = snd (last weighted) / Exp (log $ fromIntegral n')

    -- list of n uniform[i-1/n,i/n] random values
    let between a b = do
          r <- samplefrom (Uniform (fromIntegral a) (fromIntegral b))
          return $ avgrest * Exp (log r)

    rs <- zipWithM between [0..n'-1] [1..n']
    return $ map (reweight $ Product avgweight) (getSamples rs weighted ++ guaranteed)
{-
stratResample :: [Weighted (SMCModel a)] -> Random [SMCModel a]
stratResample xs = do
    -- [(a1, w1), (a2, w2), ...] becomes [(a1, w1), (a2, w1+w2), ...]
    let weighted = let (as, ws) = unzip xs in zip as (scanl1 (+) ws)
    -- get sum of weights
    let weight = snd $ last weighted
    let n = length weighted
    let avgweight = weight / Exp (log $ fromIntegral n)

    -- list of n uniform[i-1/n,i/n] random values
    let between a b = do
          r <- samplefrom (Uniform (fromIntegral a) (fromIntegral b))
          return $ avgweight * Exp (log r)
    rs <- zipWithM between [0..n-1] [1..n]

    -- list of k weighted samples
    let result = getSamples rs weighted
    return $ map (reweight $ Product avgweight) result
-}

class Infer m => Addressable m where
  resample' :: Address -> m ()

  score' :: Address -> Double -> m ()
  score' a w = score w >> resample' a

  observe' :: Dist d v => Address -> d -> v -> m ()
  observe' a d v = observe d v >> resample' a

-- subset SMC specific functions
instance Addressable Labelled where
  resample' a = Labelled $ Resample a mempty $ return ()

-- alpha-SMC specific functions
instance Addressable AlphaSMC where
  resample' a = AlphaSMC $ Resample a mempty $ return ()

-- get address from a Dome or Resample
address :: Bounded l => ParticleT l w m a -> l
address (Resample a _ _) = a
address (Done _ _) = maxBound
address (Compute _ ) = error "compute has no address"

-- helper function used by alpha-SMC
resampleMinSubset :: [AddressSMC a] -> Random [AddressSMC a]
resampleMinSubset ps = do
  let mini = minimum $ map address ps
  let (minima, rest) = partition ((== mini) . address) ps
  let w = sum $ map weight minima
  let n = length minima
  minima' <- sampleWeighted n $ map getweight minima
  let normalized = map (reweight $ Product (w / Exp (log $ fromIntegral n))) minima'
  return (normalized ++ rest)

resampleAlpha :: Int -> [AlphaSMC a] -> Random [AlphaSMC a]
resampleAlpha n ms = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unalpha) ms
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return $ map AlphaSMC ms'
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- resampleMinSubset ms'
        -- recursive call to continue inference
        resampleAlpha n $ map AlphaSMC resampled

resampleESS :: Int -> ([Adaptive a], Weight) -> Random ([Adaptive a], Weight)
resampleESS n (ms, w) = do
    -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map (compute . unadapt) ms
    let w1 = sum (map weight ms')
    let w2 = sum (map ((^2) . weight) ms')
    let ess = w1 * w1 / w2
    let w' = w * (w1 / Exp (log $ fromIntegral n))

    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return (map Adaptive ms', w')
      -- otherwise we do resampling
      else do
        if ess < Exp (log (1 * fromIntegral n))
        then do
          resampled <- sampleWeighted n (map getweight ms')
          resampleESS n (map Adaptive resampled, w')
        else
          resampleESS n (map (Adaptive . ignoreResample) ms', w)

instance SMC AlphaSMC where
  resampleSMC n ms = do
    ms' <- resampleAlpha n ms
    return (map (\(AlphaSMC (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleAlpha n ms
    let w = sum (map (\(AlphaSMC (Done w a)) -> getProduct w) ms')
    return $ w / Exp (log $ fromIntegral n)

instance SMC Labelled where
  resampleSMC n ms = do
    ms' <- resampleLab ms
    return (map (\(Labelled (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleLab ms
    return $ getProduct (sum $ map (\(Labelled (Done w a)) -> w) ms') / Exp (log $ fromIntegral $ length ms')

instance SMC Adaptive where
  resampleSMC n ms = do
    (ms', w) <- resampleESS n (ms, Exp 0)
    return (map (\(Adaptive (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = snd <$> resampleESS n (ms, Exp 0)

instance SMC Stratified where
  resampleSMC n ms = do
    ms' <- resampleStrat ms
    return (map (\(Stratified (Done w a)) -> (a, getProduct w)) ms')
  normalisingConstant n ms = do
    ms' <- resampleStrat ms
    return $ getProduct (sum $ map (\(Stratified (Done w a)) -> w) ms') / Exp (log $ fromIntegral $ length ms')

ignoreResample :: ParticleT l (Product Weight) Random a -> ParticleT l (Product Weight) Random a
ignoreResample (Resample _ w a) = reweight w a
ignoreResample p@(Done _ _) = p
ignoreResample (Compute _) = error "unexpected compute"

instance Stochastic AliveFilter where
  -- keep trying until we get a nonzero weight
  compute p@(AliveFilter (Compute ma)) = do
    a <- compute (Compute ma)
    if weight a == Exp (log 0) -- note that log 0 = -Infinity
    then compute p
    else return $ AliveFilter a
  compute p = return p

instance Stochastic AlphaSMC where
  -- keep trying until we get a nonzero weight
  compute p@(AlphaSMC (Compute ma)) = do
    a <- compute (Compute ma)
    if weight a == Exp (log 0) -- note that log 0 = -Infinity
    then compute p
    else return $ AlphaSMC a
  compute p = return p

-- get weights, clear internal weights, remove resample
getweight :: ParticleT l (Product Weight) m a -> Weighted (ParticleT l (Product Weight) m a)
getweight (Done w a) = (Done mempty a, getProduct w)
getweight (Resample l w a) = (a, getProduct w)
getweight (Compute a) = error "unexpected compuation"

-- check if particle is done
isDone :: ParticleT l w m a -> Bool
isDone (Done _ _) = True
isDone _ = False

smcmodel :: Int -> ([SMCModel a], Weight) -> Random ([SMCModel a], Weight)
smcmodel n (ms, w) = do
  -- do all the stochastic compuations so everything
    -- is either done or at a resample statement
    ms' <- parallel $ map compute ms
    let w' = w * sum (map weight ms') / Exp (log $ fromIntegral n)
    -- check if all particles are done...
    if all isDone ms'
      -- ... in which case return
      then return (ms', w')
      -- otherwise we do resampling
      else do
        -- we need to prune the resample we dealt with and reset weights
        resampled <- sampleWeighted n (map getweight ms')
        -- recursive call to continue inference
        smcmodel n (resampled, w')

instance SMC SMCModel where
  resampleSMC n ms = do
    (ms', w) <- smcmodel n (ms, Exp 0)
    return (map (\(Done w a) -> (a, getProduct w)) ms')
  normalisingConstant n ms = snd <$> smcmodel n (ms, Exp 0)

inferSMC :: SMC m => Int -> m a -> Random [Weighted a]
inferSMC n m = resampleSMC n (replicate n m)

inferConstant :: SMC m => Int -> m a -> Random Weight
inferConstant n m = normalisingConstant n (replicate n m)


{-# LANGUAGE ApplicativeDo                            #-}
{-# LANGUAGE DataKinds                                #-}
{-# LANGUAGE DeriveFunctor                            #-}
{-# LANGUAGE ExistentialQuantification                #-}
{-# LANGUAGE FlexibleContexts                         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving               #-}
{-# LANGUAGE KindSignatures                           #-}
{-# LANGUAGE LambdaCase                               #-}
{-# LANGUAGE RankNTypes                               #-}
{-# LANGUAGE RecordWildCards                          #-}
{-# LANGUAGE ScopedTypeVariables                      #-}
{-# LANGUAGE StandaloneDeriving                       #-}
{-# LANGUAGE TypeApplications                         #-}
{-# LANGUAGE TypeOperators                            #-}
{-# LANGUAGE UndecidableInstances                     #-}
{-# LANGUAGE ViewPatterns                             #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise       #-}

-- |
-- Module      : Data.Bin
-- Copyright   : (c) Justin Le 2018
-- License     : BSD3
--
-- Maintainer  : justin@jle.im
-- Stability   : experimental
-- Portability : non-portable
--
-- Tools for aggregating numeric values into a set of discrete bins
-- according to some binning specification.
--
-- See 'withBinner' for main usage information, and 'Bin' for the main
-- binned data type, and 'binFreq' for a common usage example.
--

module Data.Bin (
  -- * Specifying the binning
    BinSpec(..), linBS, logBS, gaussBS
  -- ** 'BinView'
  , BinView, binView, linView, logView, gaussView
  -- ** Inspecting 'BinSpec'
  , binSpecIntervals, binSpecCenters
  -- * Creating and manipulating bins
  , Bin, Binner, withBinner, fromFin, bins
  -- ** Inspecting bins
  , binFin, binRange, binMin, binMax, binMid
  , binFinExt, binFinComp
  -- ** Showing bins
  , displayBin, displayBinDouble
  -- *** In-depth inspection
  , Pointed(..), pointed, pElem, binIx, fromIx, expandFin, unexpandFin
  -- * Untyped
  , SomeBin(..), sameBinSpec
  -- * Handy use patterns
  , binFreq
  ) where

import           Control.Monad
import           Data.Finite
import           Data.Foldable
import           Data.Profunctor
import           Data.Proxy
import           Data.Reflection
import           Data.Tagged
import           Data.Type.Equality
import           GHC.TypeNats
import           Numeric.SpecFunctions
import           Text.Printf
import           Unsafe.Coerce
import qualified Data.Map              as M
import qualified Data.Vector.Sized     as SV

-- | A bidirectional "view" to transform the data type before binning.
--
-- See 'linView' for a linear binning, and 'logView' for a logarithmic
-- binning.  You can construct your own custom transformer using 'binView'.
--
-- This type is essentially 'Iso' from the /lens/ library, and any 'Iso''
-- from lens can be used here.  However, it is important that all of these
-- represent /monotonic/ isomorphisms.
type BinView a b = forall p. Profunctor p => p b b -> p a a

-- | Construct a 'BinView' based on "to" and "from" functions
--
-- It is important that the "to" and "from" functions be /inverses/ of each
-- other.  Furthermore, both "to" and "from" should be __monotonic__.
binView
    :: (a -> b)       -- ^ "to"
    -> (b -> a)       -- ^ "from"
    -> BinView a b
binView = dimap

-- | Linear binning
linView :: BinView a a
linView = binView id id

-- | Logarithmic binning (smaller bins at lower levels, larger bins at
-- higher levels).
logView :: Floating a => BinView a a
logView = binView log exp

-- | Binning based on a Gaussian Distribution.  Bins "by standard
-- deviation"; there are more bins the closer to the mean you get, and less
-- bins the farther away.
gaussView
    :: RealFrac a
    => a           -- ^ center / mean
    -> a           -- ^ standard deviation
    -> BinView a Double
gaussView μ σ = binView to from
  where
    to   = erf . realToFrac . (/ σ) . subtract μ
    from = (+ μ) . (* σ) . realToFrac . invErf

view :: BinView a b -> a -> b
view v = runForget (v (Forget id))

review :: BinView a b -> b -> a
review v = unTagged . v . Tagged

-- | Specification of binning.
--
-- A @'BinSpec' n a b@ will bin values of type @a@ into @n@ bins, according
-- to a scaling in type @b@.
--
-- Constructor is meant to be used with type application syntax to indicate
-- @n@, like @'BinSpec' @5 0 10 'linView'@
data BinSpec (n :: Nat) a b =
        BS { bsMin  :: a             -- ^ lower bound of values
           , bsMax  :: a             -- ^ upper bound of values
           , bsView :: BinView a b   -- ^ binning view
           }

-- | Convenient constructor for a 'BinSpec' for a linear scaling.
--
-- Meant to be used with type application syntax:
--
-- @
-- 'linBS' @5 0 10
-- @
linBS
    :: forall n a. ()
    => a                            -- ^ Lower bound
    -> a                            -- ^ Upper bound
    -> BinSpec n a a
linBS mn mx = BS mn mx linView

-- | Convenient constructor for a 'BinSpec' for a logarithmic scaling.
--
-- Meant to be used with type application syntax:
--
-- @
-- 'logBS' @5 0 10
-- @
logBS
    :: forall n a. Floating a
    => a                            -- ^ Lower bound
    -> a                            -- ^ Upper bound
    -> BinSpec n a a
logBS mn mx = BS mn mx logView

-- | Convenient constructor for a 'BinSpec' for a gaussian scaling.  Uses
-- the midpoint as the inferred mean.
--
-- Meant to be used with type application syntax:
--
-- @
-- 'gaussBS' @5 3 0 10
-- @
--
-- indicates that you want 5 bins.
gaussBS
    :: forall n a. RealFrac a
    => a                            -- ^ Standard Deviation
    -> a                            -- ^ Lower bound
    -> a                            -- ^ Upper bound
    -> BinSpec n a Double
gaussBS σ mn mx = BS mn mx (gaussView ((mn + mx)/2) σ)

-- | Data type extending a value with an extra "minimum" and "maximum"
-- value.
data Pointed a = Bot
               | PElem !a
               | Top
  deriving (Show, Eq, Ord, Functor)

-- | @since 0.1.3.0
instance (Enum a, Bounded a) => Enum (Pointed a) where
    toEnum i
      | i == 0 = Bot
      | i == (fromEnum (maxBound :: a) + 2) = Top
      | otherwise = PElem (toEnum (i - 1))
    fromEnum = \case
      Bot     -> 0
      PElem x -> fromEnum x + 1
      Top     -> fromEnum (maxBound :: a) + 2

-- | @since 0.1.3.0
instance Bounded (Pointed a) where
    minBound = Bot
    maxBound = Top

-- | Church-style deconstructor for 'Pointed', analogous to 'maybe',
-- 'either', and 'bool'.
--
-- @since 0.1.1.0
pointed
    :: b                -- ^ return if 'Bot'
    -> (a -> b)         -- ^ apply if 'PElem'
    -> b                -- ^ return if 'Top'
    -> Pointed a
    -> b
pointed bot pelem top = \case
    Bot     -> bot
    PElem x -> pelem x
    Top     -> top

-- | Extract the item from a 'Pointed' if it is neither the extra minimum
-- or maximum.
pElem :: Pointed a -> Maybe a
pElem = pointed Nothing Just Nothing

-- | A @'Bin' s n@ is a single bin index out of @n@ partitions of the
-- original data set, according to a 'BinSpec' represented by @s@.
--
-- All 'Bin's with the same @s@ follow the same 'BinSpec', so you can
-- safely use 'binRange' 'withBinner'.
--
-- It has useful 'Eq' and 'Ord' instances.
--
-- Actually has @n + 2@ partitions, since it also distinguishes values
-- that are outside the 'BinSpec' range.
newtype Bin s n = Bin { _binIx :: Pointed (Finite n) }
  deriving (Eq, Ord, Enum, Bounded)

-- | A more specific version of 'binFin' that indicates whether or not the
-- value was too high or too low for the 'BinSpec' range.
binIx :: Bin s n -> Pointed (Finite n)
binIx = _binIx

-- | Extract, potentially, the 'Bin' index.  Will return 'Nothing' if the
-- original value was outside the 'BinSpec' range.
--
-- See 'binIx' for a more specific version, which indicates if the original
-- value was too high or too low.  Also see 'binFinExt', which extends the
-- range of the 'Finite' to embed lower or higher values.
binFin :: Bin s n -> Maybe (Finite n)
binFin = pElem . binIx

-- | Like 'binFin', but return the true "n + 2" slot number of a 'Bin',
-- where 'minBound' is "below minimum" and 'maxBound' is "above maximum"
--
-- @since 0.1.1.0
binFinExt
    :: KnownNat n
    => Bin s n
    -> Finite (1 + n + 1)
binFinExt = expandFin . binIx

-- | Like 'binFin', but squishes or compresses "below minimum" to "above
-- maximum" bins into the 'Finite', counting them in the same bin as the
-- minimum and maximum bin, respectively.
--
-- @since 0.1.1.0
binFinComp
    :: KnownNat n
    => Bin s n
    -> Finite n
binFinComp = pointed minBound id maxBound . binIx

tick
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> b
tick BS{..} = totRange / fromIntegral (natVal (Proxy @n))
  where
    totRange = view bsView bsMax - view bsView bsMin

packPointed
    :: KnownNat n
    => Integer
    -> Pointed (Finite n)
packPointed n
    | n < 0     = Bot
    | otherwise = maybe Top PElem . packFinite $ n

mkBin_
    :: forall n a b. (KnownNat n, RealFrac b)
    => BinSpec n a b
    -> a
    -> Pointed (Finite n)
mkBin_ bs = packPointed
          . floor
          . (/ tick bs)
          . subtract (scaleIn (bsMin bs))
          . scaleIn
  where
    scaleIn = view (bsView bs)

mkBin
    :: forall n a b s. (KnownNat n, RealFrac b, Reifies s (BinSpec n a b))
    => a
    -> Bin s n
mkBin = Bin . mkBin_ (reflect (Proxy @s))

-- | The type of a "binning function", given by 'withBinner'.  See
-- 'withBinner' for information on how to use.
type Binner s n a = a -> Bin s n

-- | With a 'BinSpec', give a "binning function" that you can use to create
-- bins within a continuation.
--
-- @
-- 'withBinner' myBinSpec $ \toBin ->
--     show (toBin 2.8523)
-- @
--
-- Uses a Rank-N continution to ensure that you can only compare 'Bin's
-- constructed from the same 'BinSpec'/binning function.
withBinner
    :: (KnownNat n, RealFrac b)
    => BinSpec n a b
    -> (forall s. Reifies s (BinSpec n a b) => Binner s n a -> r)
    -> r
withBinner bs f = reify bs $ \(_ :: Proxy s) -> f @s mkBin

-- | Generate a vector of the boundaries deriving the bins from
-- a 'BinSpec'.  Can be useful for debugging.
binSpecIntervals
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> SV.Vector (n + 1) a
binSpecIntervals bs = SV.generate $ \i ->
    case strengthen i of
      Just (fromIntegral->i') -> scaleOut $ i' * t + scaleIn (bsMin bs)
      Nothing                 -> bsMax bs
  where
    t        = tick bs
    scaleIn  = view (bsView bs)
    scaleOut = review (bsView bs)

-- | Generate a vector of the scaled centers of each bin derived from
-- a 'BinSpec'.  Can be useful for debugging.
--
-- Note that for non-linear scales, this centering respects the scale.  For
-- example, for logarithmic scales, this is a "log-aware" center: it gives
-- the midpoint in log-space, not linear-space.
--
-- @since 0.1.3.0
binSpecCenters
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> SV.Vector n a
binSpecCenters bs = SV.generate $ \i ->
    scaleOut $ (fromIntegral i + 0.5) * t + scaleIn (bsMin bs)
  where
    t        = tick bs
    scaleIn  = view (bsView bs)
    scaleOut = review (bsView bs)

binRange_
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> Pointed (Finite n)
    -> (Maybe a, Maybe a)
binRange_ bs = \case
    Bot     -> ( Nothing         , Just (SV.head v))
    PElem i -> ( Just (v `SV.index` weaken i)
               , Just (v `SV.index` shift i )
               )
    Top     -> ( Just (SV.last v), Nothing         )
  where
    v        = binSpecIntervals @n bs

-- | Extract the minimum and maximum of the range indicabed by a given
-- 'Bin'.
--
-- A 'Nothing' value indicates that we are outside of the normal range of
-- the 'BinSpec', so is "unbounded" in that direction.
binRange
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> (Maybe a, Maybe a)
binRange = binRange_ (reflect (Proxy @s)) . binIx

-- | Extract the minimum of the range indicabed by a given 'Bin'.
--
-- A 'Nothing' value means that the original value was below the minimum
-- limit of the 'BinSpec', so is "unbounded" in the lower direction.
binMin
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> Maybe a
binMin = fst . binRange

-- | Extract the maximum of the range indicabed by a given 'Bin'.
--
-- A 'Nothing' value means that the original value was above the maximum
-- limit of the 'BinSpec', so is "unbounded" in the upper direction.
binMax
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> Maybe a
binMax = snd . binRange

-- | Extract the center of the range indicabed by a given 'Bin'.
--
-- A 'Nothing' value means that the original value was either below or above the maximum
-- limit of the 'BinSpec'.
--
-- Note that for non-linear scales, this centering respects the scale.  For
-- example, for logarithmic scales, this is a "log-aware" center: it gives
-- the midpoint in log-space, not linear-space.
--
-- @since 0.1.3.0
binMid
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => Bin s n
    -> Maybe a
binMid = binMid_ (reflect (Proxy @s)) . binIx

binMid_
    :: forall n a b. (KnownNat n, Fractional b)
    => BinSpec n a b
    -> Pointed (Finite n)
    -> Maybe a
binMid_ bs = fmap (v `SV.index`) . pElem
  where
    v = binSpecCenters @n bs

-- | Display the interval maintained by a 'Bin'.
displayBin
    :: forall n a b s. (KnownNat n, Fractional b, Reifies s (BinSpec n a b))
    => (a -> String)        -- ^ how to display a value
    -> Bin s n
    -> String
displayBin f b = printf "%s .. %s" mn' mx'
  where
    (mn, mx) = binRange b
    mn' = case mn of
            Nothing -> "(-inf"
            Just m  -> "[" ++ f m
    mx' = case mx of
            Nothing -> "+inf)"
            Just m  -> f m ++ ")"

-- | Display the interval maintained by a 'Bin', if the 'Bin' contains
-- a 'Double'.
displayBinDouble
    :: forall n b s. (KnownNat n, Fractional b, Reifies s (BinSpec n Double b))
    => Int                      -- ^ number of decimal places to round
    -> Bin s n
    -> String
displayBinDouble d = displayBin (printf ("%." ++ show d ++ "f"))

instance (KnownNat n, Show a, Fractional b, Reifies s (BinSpec n a b)) => Show (Bin s n) where
    showsPrec d b = showParen (d > 10) $
      showString "Bin " . showString (displayBin @n show b)

-- | Generate a histogram: given a container of @a@s, generate a frequency
-- map of how often values in a given discrete bin occurred.
--
-- @
-- xs :: [Double]
-- xs = [1..100]
--
-- main :: IO ()
-- main = withBinner (logBS @10 5 50) $ \toBin ->
--     mapM_ (\(b, n) -> putStrLn (displayBinDouble 4 b ++ "\t" ++ show n))
--   . M.toList
--   $ binFreq toBin xs
-- @
--
-- @
-- (-inf .. 5.0000)        4
-- [5.0000 .. 6.2946)      2
-- [6.2946 .. 7.9245)      1
-- [7.9245 .. 9.9763)      2
-- [9.9763 .. 12.5594)     3
-- [12.5594 .. 15.8114)    3
-- [15.8114 .. 19.9054)    4
-- [19.9054 .. 25.0594)    6
-- [25.0594 .. 31.5479)    6
-- [31.5479 .. 39.7164)    8
-- [39.7164 .. 50.0000)    10
-- [50.0000 .. +inf)       51
-- @
binFreq
    :: forall n t a s. Foldable t
    => Binner s n a
    -> t a
    -> M.Map (Bin s n) Int
binFreq toBin = M.unionsWith (+) . map go . toList
  where
    go :: a -> M.Map (Bin s n) Int
    go x = M.singleton (toBin x) 1

-- | Construct a 'Bin' if you know the bin number you want to specify, or
-- if the bin is over or under the maximum.
fromIx :: Pointed (Finite n) -> Bin s n
fromIx = Bin

-- | Construct a 'Bin' if you know the bin number you want to specify.  See
-- 'fromIx' if you want to specify bins that are over or under the maximum,
-- as well.
fromFin :: Finite n -> Bin s n
fromFin = fromIx . PElem

-- | An (ordered) list of all 'Bin's.
--
-- @since 0.1.3.0
bins :: KnownNat n => [Bin s n]
bins = Bin <$> binser []
  where
    binser = (Bot :)
           . foldr ((.) . (:) . PElem) id finites
           . (Top :)

-- | "Expand" a 'Pointed' containing a 'Finite' to a wider-ranged 'Finite'.
-- Used for 'binFinExt'
--
-- @since 0.1.2.0
expandFin :: KnownNat n => Pointed (Finite n) -> Finite (1 + n + 1)
expandFin = pointed minBound (weaken . shift) maxBound

-- | The inverse of 'expandFin': "re-pack" a 'Finite' back into
-- a 'Pointed' containing a narrower-ranged 'Finite'.
--
-- @since 0.1.2.0
unexpandFin :: KnownNat n => Finite (1 + n + 1) -> Pointed (Finite n)
unexpandFin x = case unshift x of
  Nothing -> Bot
  Just y  -> case strengthen y of
    Nothing -> Top
    Just z  -> PElem z

-- | A @'SomeBin' a n@ is @'Bin' s n@, except with the 'BinSpec' s hidden.
-- It's useful for returning out of 'withBinner'.
--
-- It has useful 'Eq' and 'Ord' instances.
--
-- To be able to "unify" two 'Bin's inside a 'SomeBin', use 'sameBinSpec'
-- to verify that the two 'SomeBin's were created with the same 'BinSpec'.
data SomeBin a n = forall s b. (Fractional b, Reifies s (BinSpec n a b))
    => SomeBin { getSomeBin :: Bin s n }

deriving instance (KnownNat n, Show a) => Show (SomeBin a n)

-- | Compares if the ranges match.  Note that this is less performant than
-- comparing the original 'Bin's, or extracting and using 'sameBinSpec'.
instance (KnownNat n, Eq a) => Eq (SomeBin a n) where
    SomeBin x == SomeBin y = binRange x == binRange y

-- | Lexicographical ordering -- compares the lower bound, then the upper
-- bounds.  Note that this is less performant than comparing the original
-- 'Bin's, or extracting and using 'sameBinSpec'
instance (KnownNat n, Ord a) => Ord (SomeBin a n) where
    compare (SomeBin x) (SomeBin y) = compare (binRange x) (binRange y)

-- | Verify that the two reified 'BinSpec' types refer to the same one,
-- allowing you to use functions like '==' and 'compare' on 'Bin's that you
-- get out of a 'SomeBin'.
sameBinSpec
    :: forall s t n a b p. (Reifies s (BinSpec n a b), Reifies t (BinSpec n a b), KnownNat n, Eq a, Fractional b)
    => p s
    -> p t
    -> Maybe (s :~: t)
sameBinSpec _ _ = do
    guard $ binSpecIntervals bs1 == binSpecIntervals bs2
    guard $ binSpecCenters bs1 == binSpecCenters bs2
    pure (unsafeCoerce Refl)
  where
    bs1 = reflect (Proxy @s)
    bs2 = reflect (Proxy @t)


{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Configurator.Instances () where

import Data.Configurator.Types.Internal
import Data.Complex (Complex)
import Data.Fixed (Fixed, HasResolution)
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Text.Encoding (encodeUtf8)
import Data.Ratio (Ratio, denominator, numerator)
import Data.Word (Word8, Word16, Word32, Word64)
import Foreign.C.Types (CDouble, CFloat)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as LB
import qualified Data.Text as T
import qualified Data.Text.Lazy as L

instance Configured Value where
    convert = Just

instance Configured Bool where
    convert (Bool v) = Just v
    convert _        = Nothing

convertNumberToNum :: (Num a) => Value -> Maybe a
convertNumberToNum (Number r)
    | denominator r == 1 = Just $ fromInteger $ numerator r
convertNumberToNum _ = Nothing

instance Configured Int where
    convert = convertNumberToNum

instance Configured Integer where
    convert = convertNumberToNum

instance Configured Int8 where
    convert = convertNumberToNum

instance Configured Int16 where
    convert = convertNumberToNum

instance Configured Int32 where
    convert = convertNumberToNum

instance Configured Int64 where
    convert = convertNumberToNum

instance Configured Word where
    convert = convertNumberToNum

instance Configured Word8 where
    convert = convertNumberToNum

instance Configured Word16 where
    convert = convertNumberToNum

instance Configured Word32 where
    convert = convertNumberToNum

instance Configured Word64 where
    convert = convertNumberToNum

convertNumberToFractional :: (Fractional a) => Value -> Maybe a
convertNumberToFractional (Number r) = Just $ fromRational r
convertNumberToFractional _ = Nothing

instance Configured Double where
    convert = convertNumberToFractional

instance Configured Float where
    convert = convertNumberToFractional

instance Configured CDouble where
    convert = convertNumberToFractional

instance Configured CFloat where
    convert = convertNumberToFractional

instance Integral a => Configured (Ratio a) where
    convert = convertNumberToFractional

instance RealFloat a => Configured (Complex a) where
    convert = convertNumberToFractional

instance HasResolution a => Configured (Fixed a) where
    convert = convertNumberToFractional

instance Configured T.Text where
    convert (String v) = Just v
    convert _          = Nothing

instance Configured Char where
    convert (String txt) | T.length txt == 1 = Just $ T.head txt
    convert _ = Nothing

    convertList = fmap T.unpack . convert

instance Configured L.Text where
    convert = fmap L.fromStrict . convert

instance Configured B.ByteString where
    convert = fmap encodeUtf8 . convert

instance Configured LB.ByteString where
    convert = fmap (LB.fromChunks . (:[])) . convert

instance (Configured a, Configured b) => Configured (a,b) where
    convert (List [a,b]) = (,) <$> convert a <*> convert b
    convert _            = Nothing

instance (Configured a, Configured b, Configured c) => Configured (a,b,c) where
    convert (List [a,b,c]) = (,,) <$> convert a <*> convert b <*> convert c
    convert _              = Nothing

instance (Configured a, Configured b, Configured c, Configured d)
    => Configured (a,b,c,d) where
    convert (List [a,b,c,d]) = (,,,) <$> convert a <*> convert b <*> convert c
                                     <*> convert d
    convert _                = Nothing


module SimpleR.Language.Syntax where

import qualified Data.Complex as C

type Complex = C.Complex Double

data IdentAnnot = IdentAnnot deriving (Ord, Eq, Show, Read)

data Ident = Ident
  { idName :: String
  , idPkg :: Maybe String
  , idAnnot :: Maybe IdentAnnot
  } deriving (Ord, Eq, Show, Read)

data TickAnnot = TickAnnot  deriving (Eq, Show, Read)

data Tick = Tick
  { tickAnnot :: TickAnnot
  }  deriving (Eq, Show, Read)

data Const =
    IntConst Int
  | DoubleConst Double
  | ComplexConst Complex
  | BoolConst Bool
  | StringConst String
  | NilConst
  | NaConst
   deriving (Eq, Show, Read)

data Param =
    Param Ident
  | Default Ident Expr
  | VarParam
   deriving (Eq, Show, Read)

data Arg =
    Arg Expr
  | Named Ident Expr
  | VarArg
   deriving (Eq, Show, Read)

data Expr =
    Var Ident
  | Const Const
  | Seq [Expr]
  | LamAbs [Param] Expr
  | LamApp Expr [Arg]
  | NativeLamApp Ident [Ident]
  | Return Expr
  | Assign Expr Expr

  | VecInd Expr [Arg]
  | VecIndAssign Expr [Arg] Expr
  | VecSub Expr [Arg]
  | VecSubAssign Expr [Arg] Expr
  | ListName Expr Expr
  | ListNameAssign Expr Expr Expr
  | ObjAttr Expr Expr
  | ObjAttrAssign Expr Expr Expr

  | If Expr Expr Expr
  | While Expr Expr
  | Break
  | Next

  | Error
  | Tix Tick Expr
   deriving (Eq, Show, Read)

data Program = Program [Expr]  deriving (Eq, Show, Read)


{-# LANGUAGE DefaultSignatures   #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE InstanceSigs        #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Data.Roles
  ( Representational(rep)
  , Phantom(phantom)

  -- * Utility functions to define Representable easily
  , Proxy2(..)
  , wrapped
  , slash
  , coercionOf
  , eta
  ) where

import           Control.Applicative
import           Control.Category      ((>>>))
import           Control.Monad.ST      (ST)

import           Data.Complex
import           Data.Monoid
import           Data.Semigroup

import           Data.Proxy

import           Data.Functor.Identity
import           Data.Functor.Compose
import qualified Data.Functor.Product  as F
import qualified Data.Functor.Sum      as F

import           Data.IntMap
import           Data.Map

import Control.Monad.Reader      (ReaderT(..))
import Control.Monad.Except      (ExceptT(..), runExceptT)
import Control.Monad.Cont        (ContT(..))
import Control.Monad.Trans.Maybe (MaybeT(..))
import qualified Control.Monad.Writer.Lazy   as L (WriterT(..), runWriterT)
import qualified Control.Monad.Writer.Strict as S (WriterT(..), runWriterT)
import qualified Control.Monad.State.Lazy    as L (StateT(..), runStateT)
import qualified Control.Monad.State.Strict  as S (StateT(..), runStateT)
import qualified Control.Monad.RWS.Lazy      as L (RWST(..), runRWST)
import qualified Control.Monad.RWS.Strict    as S (RWST(..), runRWST)

import           Data.Coerce
import           Data.Type.Coercion
import           Unsafe.Coerce

class Representational (t :: k1 -> k2) where
  -- | An argument is representational if you can lift a coercion of the argument into one of the whole
  rep :: Coercion a b -> Coercion (t a) (t b)
  default rep :: Phantom t => Coercion a b -> Coercion (t a) (t b)
  rep _ = phantom

class Representational t => Phantom (t :: k1 -> k2) where
  -- | An argument is phantom if you can 'coerce' the whole ignoring the argument
  phantom :: Coercion (t a) (t b)
  default phantom :: Coercible (t a) (t b) => Coercion (t a) (t b)
  phantom = Coercion

wrapped :: forall x y proxy (f :: x -> y) (g :: x -> y) (a :: x) (b :: x).
           (Coercible f g) => proxy f g -> Coercion (f a) (f b) -> Coercion (g a) (g b)
wrapped _ c = (Coercion :: Coercion (g a) (f a)) >>> c >>> (Coercion :: Coercion (f b) (g b))

data Proxy2 a b = Proxy2

slash ::  proxy (f a) (g a) -> Proxy2 f g
slash _ = Proxy2

coercionOf :: Coercible a b => proxy a b -> Coercion a b
coercionOf _ = Coercion

eta :: forall f g a. Coercion f g -> Coercion (f a) (g a)
eta Coercion = Coercion

-- * Data.Proxy

instance Representational Proxy
instance Phantom Proxy

-- * Data.Functor.{Const, Identity, Sum, Product, Compose}

instance Representational Const where rep Coercion = Coercion
instance Representational (Const a)
instance Phantom (Const a)

instance Representational Identity where rep Coercion = Coercion

instance Representational F.Sum     where rep Coercion = Coercion
instance Representational (F.Sum f) where rep Coercion = Coercion

instance Representational F.Product where rep Coercion = Coercion
instance Representational (F.Product f) where rep Coercion = Coercion

instance Representational Compose where rep Coercion = Coercion
instance (Representational f,
          Representational g)
      => Representational (Compose f g) where
  rep c = coercionOf getCompose >>> rep (rep c) >>> coercionOf Compose

-- * Data.Type.Coercion

instance Representational Coercion     where rep = unsafeCoerce
instance Representational (Coercion a) where rep Coercion = Coercion

-- * Prelude

instance Representational (->)       where rep Coercion = Coercion
instance Representational ((->) a)   where rep Coercion = Coercion

instance Representational Either     where rep Coercion = Coercion
instance Representational (Either a) where rep Coercion = Coercion

instance Representational (,)     where rep Coercion = Coercion
instance Representational ((,) a) where rep Coercion = Coercion

instance Representational (,,)       where rep Coercion = Coercion
instance Representational ((,,) a)   where rep Coercion = Coercion
instance Representational ((,,) a b) where rep Coercion = Coercion

instance Representational (,,,)         where rep Coercion = Coercion
instance Representational ((,,,) a)     where rep Coercion = Coercion
instance Representational ((,,,) a b)   where rep Coercion = Coercion
instance Representational ((,,,) a b c) where rep Coercion = Coercion

instance Representational (,,,,)           where rep Coercion = Coercion
instance Representational ((,,,,) a)       where rep Coercion = Coercion
instance Representational ((,,,,) a b)     where rep Coercion = Coercion
instance Representational ((,,,,) a b c)   where rep Coercion = Coercion
instance Representational ((,,,,) a b c d) where rep Coercion = Coercion

instance Representational (,,,,,)             where rep Coercion = Coercion
instance Representational ((,,,,,) a)         where rep Coercion = Coercion
instance Representational ((,,,,,) a b)       where rep Coercion = Coercion
instance Representational ((,,,,,) a b c)     where rep Coercion = Coercion
instance Representational ((,,,,,) a b c d)   where rep Coercion = Coercion
instance Representational ((,,,,,) a b c d e) where rep Coercion = Coercion

instance Representational []      where rep Coercion = Coercion
instance Representational Maybe   where rep Coercion = Coercion
instance Representational IO      where rep Coercion = Coercion
instance Representational (ST s)  where rep Coercion = Coercion

-- * Control.Applicative

instance Representational ZipList where rep Coercion = Coercion

instance Representational WrappedMonad where rep Coercion = Coercion
instance Representational m => Representational (WrappedMonad m) where
  rep c = wrapped (slash WrapMonad) (rep c)

instance Representational WrappedArrow where rep Coercion = Coercion
instance (Representational p) => Representational (WrappedArrow p) where
  rep c = wrapped (slash (slash WrapArrow)) (rep c)
instance (Representational (p a)) => Representational (WrappedArrow p a) where
  rep c = wrapped (slash WrapArrow) (rep c)

-- * Data.Complex

instance Representational Complex where rep Coercion = Coercion

-- * Data.Semigroup

instance Representational Arg     where rep Coercion = Coercion
instance Representational (Arg a) where rep Coercion = Coercion
instance Representational Product where rep Coercion = Coercion
instance Representational Sum     where rep Coercion = Coercion
instance Representational Dual    where rep Coercion = Coercion
instance Representational Endo    where rep Coercion = Coercion
instance Representational Option  where rep Coercion = Coercion

instance Representational Data.Semigroup.First where rep Coercion = Coercion
instance Representational Data.Semigroup.Last  where rep Coercion = Coercion
instance Representational WrappedMonoid        where rep Coercion = Coercion

-- * Data.Monoid

instance Representational Ap                  where rep Coercion = Coercion
instance Representational f => Representational (Ap f) where
  rep c = wrapped (slash Ap) (rep c)

instance Representational Data.Monoid.First   where rep Coercion = Coercion
instance Representational Data.Monoid.Last    where rep Coercion = Coercion

-- * containers
instance Representational (Map k) where rep Coercion = Coercion
instance Representational IntMap  where rep Coercion = Coercion

-- * mtl, transformers
instance Representational ReaderT     where rep Coercion = Coercion
instance Representational (ReaderT r) where rep Coercion = Coercion
instance (Representational m)
      => Representational (ReaderT r m) where
  rep c = coercionOf runReaderT >>> rep (rep c) >>> coercionOf ReaderT

instance (Representational m)
      => Representational (ExceptT e m) where
  rep c = coercionOf runExceptT >>> rep (rep c) >>> coercionOf ExceptT

instance Representational (ContT r)   where rep Coercion = Coercion
instance Representational (ContT r m) where rep Coercion = Coercion

instance Representational MaybeT where rep Coercion = Coercion
instance (Representational m)
      => Representational (MaybeT m) where
  rep c = coercionOf runMaybeT >>> rep (rep c) >>> coercionOf MaybeT

firstR :: Coercion a b -> Coercion (a,c) (b,c)
firstR = eta . rep

instance (Representational m)
      => Representational (L.WriterT w m) where
  rep c = coercionOf L.runWriterT >>> rep (firstR c) >>> coercionOf L.WriterT

instance (Representational m)
      => Representational (S.WriterT w m) where
  rep c = coercionOf S.runWriterT >>> rep (firstR c) >>> coercionOf S.WriterT

instance (Representational m)
      => Representational (L.StateT s m) where
  rep c = coercionOf L.runStateT >>> rep (rep (firstR c)) >>> coercionOf L.StateT

instance (Representational m)
      => Representational (S.StateT s m) where
  rep c = coercionOf S.runStateT >>> rep (rep (firstR c)) >>> coercionOf S.StateT

instance Representational L.RWST where rep Coercion = Coercion
instance (Representational m)
      => Representational (L.RWST e w s m) where
  --                                     c :: a ~R b
  --                               rep $ c :: (,,) a ~R (,,) b
  --                   eta . eta . rep $ c :: (a,w,s) ~R (b,w,s)
  --             rep . eta . eta . rep $ c :: m (a,w,s) ~R m (b,w,s)
  -- rep . rep . rep . eta . eta . rep $ c :: e -> s -> m (a,w,s) ~R e -> s -> m (b,w,s)
  rep c = coercionOf L.runRWST >>> (rep . rep . rep . eta . eta . rep) c >>> coercionOf L.RWST

instance Representational S.RWST where rep Coercion = Coercion
instance (Representational m)
      => Representational (S.RWST e w s m) where
  rep c = coercionOf S.runRWST >>> (rep . rep . rep . eta . eta . rep) c >>> coercionOf S.RWST


{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
module Myelin.Model where

import Control.Lens hiding ((.=), (*~))
import Data.Aeson
import Numeric.LinearAlgebra

import Myelin.Neuron

{-- TODO: Implement units
import Numeric.Units.Dimensional
import Numeric.Units.Dimensional.UnitNames
import Numeric.Units.Dimensional.Quantities
import Numeric.Units.Dimensional.SIUnits

-- TODO(Christian): Move Units somewhere else
type Siemens = Unit Metric DElectricConductance Float
type Volt = Unit Metric DElectricPotential Float
type Farrad = Unit Metric DCapacitance Float
type Seconds = Unit Metric DTime Float
type Ampere = Unit Metric DElectricCurrent Float
-}

type Label = String

-- | A Node is a vertex in the connectivity graph.
data Node = Population {
        _numNeurons :: Integer, -- ^ number of neurons in the population
        _neuronType :: NeuronType,
        _label :: Label, -- ^ human readable label for the population
        _id :: Int -- ^ internal identifier to keep track of the population
    }
    | SpikeSourceArray {
        _spikeTimes :: [Float],
        _id :: Int
    }
    | SpikeSourcePoisson {
        _rate :: Float,
        _start :: Integer,
        _id :: Int
    }
    deriving (Eq, Show)

makeLenses ''Node
makePrisms ''Node

type Delay = Float -- TODO: This is platform specific
type Probability = Float -- TODO: Very unsophisticated representation

data ParameterGenerator
  = Constant Float
  | GaussianRandom
    { mean :: Float
    , scale :: Float
    }
  deriving (Eq, Show)

-- | Biases that biases the output of a single projection
data Biases
  = BiasGenerator ParameterGenerator
  | Biases [Float]
  deriving (Eq, Show)

-- | Weight types that describes weights for a projection
data Weights
  = WeightGenerator ParameterGenerator
  | Weights (Matrix Float)
  deriving (Eq, Show)

makeLenses ''Weights
makePrisms ''Weights

-- | Types of connections in a projection from one neuron population to another
data ConnectionType
  = AllToAll
    { _biases :: Biases
    , _weights :: Weights
    }
  | OneToOne
    { _biases :: Biases
    , _weights :: Weights
    }
  deriving (Eq, Show)

makeLenses ''ConnectionType
makePrisms ''ConnectionType

-- | Determines the effect of a synapse to be excitatory of inhibitory
data SynapseEffect
  = Inhibitory
  | Excitatory
  deriving (Eq, Show)

makeLenses '' SynapseEffect
makePrisms '' SynapseEffect

-- | Effect of projections that determines the method with which 'Node's connect
data ProjectionEffect
  = Static SynapseEffect ConnectionType
-- | Dynamic .. TODO (Jens)
  deriving (Eq, Show)

makeLenses ''ProjectionEffect
makePrisms ''ProjectionEffect

-- | An edge between two units in the network with a given 'ProjectionEffect'
data Edge
  = DenseProjection
    { _effect :: ProjectionEffect
    , _input :: Node
    , _output :: Node
    }
  | MergeProjection
    { _effect :: ProjectionEffect
    , _inputs :: (Node, Node)
    , _output :: Node
    }
  | ReplicateProjection
    { _effect :: ProjectionEffect
    , _input :: Node
    , _outputs :: (Node, Node)
    }
  deriving (Eq, Show)


--
-- 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.
--

module Feldspar.Core.Frontend.Fractional
where

import Data.Complex

import Feldspar.Core.Constructs
import Feldspar.Core.Constructs.Fractional

import Feldspar.Core.Frontend.Literal
import Feldspar.Core.Frontend.Num

-- | Fractional types. The relation to the standard 'Fractional' class is
--
-- @instance `Fraction` a => `Fractional` (`Data` a)@
class (Fractional a, Numeric a) => Fraction a
  where
    fromRationalFrac :: Rational -> Data a
    fromRationalFrac = value . fromRational

    divFrac :: Data a -> Data a -> Data a
    divFrac = sugarSymF DivFrac

instance Fraction Float
instance Fraction Double

instance (Fraction a, RealFloat a) => Fraction (Complex a)

instance (Fraction a) => Fractional (Data a)
  where
    fromRational = fromRationalFrac
    (/)          = divFrac


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

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

module CountMinSketch where

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

import Data.Matrix
import Data.Foldable (toList)
import Data.Hashable
import Data.Conduit
import Data.Sequence (Seq,(|>),update,empty,index)
import qualified Data.Sequence as S
import qualified Data.Conduit.List as Cl

import Core
import Utils
import Statistics

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

-- * The table has width `W` and height `D`
-- * items in input stream comes from universe [n] = {1...n}
-- * and let p = 2^n - 1
-- * hash function i |-> a * i + b mod p mod w
type W        = Int
type D        = Int
type P        = Int
type A        = Int
type B        = Int

-- * Our `Table` is a table of `Counters` of size W x D
type Counters = Seq Int
type Table    = (Counters, W, D)

-- * hash function `HashF` = a * i + b mod p mod w
-- * which maps some `i :: a` to `a`
-- * hash functions `HashFs` is a sequence of hash functions
type HashF    = Int -> Int
type HashFs   = Seq HashF

-- * Our algorithm modifies the `Table`, `HashF`unctions, a `Counter`,
-- * and is equipped with randomness
--type Some m = (MonadState (Table, HashFs) m, MonadState Counter m, MonadRandom m)

type Some m   = (MonadState (Table, HashFs) m, MonadRandom m)

{-----------------------------------------------------------------------------
  Play with the correct data type before jumping into it
------------------------------------------------------------------------------}

-- * declare random a for a_1 ... a_d
as :: MonadRandom m => P -> m [A]
as p = take p <$> getRandomRs (1,p)

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

-- * A vacuous inital condition to be thrown away
em :: (Table,HashFs)
em = (,) (S.replicate 0 0, 0, 0) empty

-- * initalize a table given `eps`ilon and `delt`a
table :: Some m => EpsDelta -> m Table
table (ED eps del) = return $ (S.replicate (w * d) 0, w, d)
    where (w,d) = (ceiling $ 2/eps, ceiling $ log2 (1/del))

-- * pick `a :: A` and `b ::B` of hash function given `p ::P`
-- * that is we pick and b uniformly from [1..p]
ab :: Some m => P -> m (A,B)
ab p = (\(a:b:_) -> (a,b)) <$> take 2 <$> getRandomRs (1,p)

-- * construct a `hashF`unction given `p` and `w`
hashF :: Some m => P -> W -> m HashF
hashF p w = do
    (a,b) <- ab p
    return $ \i -> ((a * i + b) `mod` p) `mod` w

hs :: Some m => m [(A,B)]
hs = replicateM 10 $ ab 100

-- * construct `D` pairwise independent `hashF`unctions
-- * TODO : make sure they're actually different
hashFs :: Some m => D -> P -> W -> m HashFs
hashFs d p = S.replicateM d . hashF p

-- * initalize count min algorithm
inits :: Some m => EpsDelta -> P -> m ()
inits eps p = do
    t@(_,w,d) <- table eps
    hs        <- hashFs d p w
    put (t,hs)

update :: Some m => m ()
update = do
    ((cs,w,d), hs) <- get
    return ()

{-----------------------------------------------------------------------------
  CountMinSketch
------------------------------------------------------------------------------}


{-# LANGUAGE MultiParamTypeClasses, FlexibleContexts, FlexibleInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.NeuralNetwork.Backend.HMatrix (
  module Data.NeuralNetwork.Backend.HMatrix.Layers,
  ByHmatrix(..),
  ErrCode(..)
) where

import Data.NeuralNetwork
import Data.NeuralNetwork.Common
import Data.NeuralNetwork.Backend.HMatrix.Layers
import Numeric.LinearAlgebra (Vector, Matrix)
import Control.Monad.Except
import Data.Functor.Identity
import Data.Constraint (Dict(..))

type Err     = ExceptT ErrCode IO

-- the backend type
data ByHmatrix = ByHmatrix

instance (HeadSize z, TranslateBody Err s,
          Component (SpecToCom s),
          RunInEnv (Run (SpecToCom s)) Err)
       => Backend ByHmatrix (z :++ s) where
  type Env ByHmatrix = Err
  type ConvertFromSpec ByHmatrix (z :++ s) = SpecToCom s
  compile _ (a :++ l)= trans (hsize a) l
  witness _ _ = Dict

instance RunInEnv Identity Err where
  run = return . runIdentity

instance TranslateBody Err SpecFullConnect where
  type SpecToCom SpecFullConnect = RunLayer (S F (T (SinglC :. Vector)))
  trans (D1 s) (FullConnect n) = do u <- lift $ newFLayer s n
                                    return $ Stack u (Activation (relu, relu'))
  trans _ _ = throwError ErrMismatch

instance TranslateBody Err SpecConvolution where
  type SpecToCom SpecConvolution = RunLayer (S C (T (MultiC :. Matrix)))
  trans (D2 k s t) (Convolution n f p) = do u <- lift $ newCLayer k n f p
                                            return $ Stack u (Activation (relu, relu'))
  trans _ _ = throwError ErrMismatch

instance TranslateBody Err SpecReshape2DAs1D where
  type SpecToCom SpecReshape2DAs1D = RunLayer A
  trans (D2 _ _ _) _ = return As1D
  trans (D1 _)     _ = throwError ErrMismatch

instance TranslateBody Err SpecMaxPooling where
  type SpecToCom SpecMaxPooling = RunLayer M
  trans (D2 _ _ _) (MaxPooling n) = return (MaxP n)
  trans (D1 _)     _              = throwError ErrMismatch


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Recurrent.Layers.BasicRecurrent (
    BasicRecurrent (..)
  , randomBasicRecurrent
  ) where

import           Control.Monad.Random ( MonadRandom, getRandom )

import           Data.Singletons.TypeLits

#if MIN_VERSION_base(4,9,0)
import           Data.Kind (Type)
#endif

import           Numeric.LinearAlgebra.Static

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Recurrent.Core

data BasicRecurrent :: Nat -- Input layer size
                    -> Nat -- Output layer size
                    -> Type where
  BasicRecurrent :: ( KnownNat input
                    , KnownNat output
                    , KnownNat matrixCols
                    , matrixCols ~ (input + output))
                 => !(R output)   -- Bias neuron weights
                 -> !(R output)   -- Bias neuron momentum
                 -> !(L output matrixCols) -- Activation
                 -> !(L output matrixCols) -- Momentum
                 -> BasicRecurrent input output

data BasicRecurrent' :: Nat -- Input layer size
                     -> Nat -- Output layer size
                     -> Type where
  BasicRecurrent' :: ( KnownNat input
                     , KnownNat output
                     , KnownNat matrixCols
                     , matrixCols ~ (input + output))
                  => !(R output)   -- Bias neuron gradients
                  -> !(L output matrixCols)
                  -> BasicRecurrent' input output

instance Show (BasicRecurrent i o) where
  show BasicRecurrent {} = "BasicRecurrent"

instance (KnownNat i, KnownNat o, KnownNat (i + o)) => UpdateLayer (BasicRecurrent i o) where
  type Gradient (BasicRecurrent i o) = (BasicRecurrent' i o)

  runUpdate lp (BasicRecurrent oldBias oldBiasMomentum oldActivations oldMomentum) (BasicRecurrent' biasGradient activationGradient) =
    let newBiasMomentum = konst (learningMomentum lp) * oldBiasMomentum - konst (learningRate lp) * biasGradient
        newBias         = oldBias + newBiasMomentum
        newMomentum     = konst (learningMomentum lp) * oldMomentum - konst (learningRate lp) * activationGradient
        regulariser     = konst (learningRegulariser lp * learningRate lp) * oldActivations
        newActivations  = oldActivations + newMomentum - regulariser
    in BasicRecurrent newBias newBiasMomentum newActivations newMomentum

  createRandom = randomBasicRecurrent

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentUpdateLayer (BasicRecurrent i o) where
  type RecurrentShape (BasicRecurrent i o) = S ('D1 o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentLayer (BasicRecurrent i o) ('D1 i) ('D1 o) where

  type RecTape (BasicRecurrent i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
  -- Do a matrix vector multiplication and return the result.
  runRecurrentForwards (BasicRecurrent wB _ wN _) (S1D lastOutput) (S1D thisInput) =
    let thisOutput = S1D $ wB + wN #> (thisInput # lastOutput)
    in ((S1D lastOutput, S1D thisInput), thisOutput, thisOutput)

  -- Run a backpropogation step for a full connected layer.
  runRecurrentBackwards (BasicRecurrent _ _ wN _) (S1D lastOutput, S1D thisInput) (S1D dRec) (S1D dEdy) =
    let biasGradient        = (dRec + dEdy)
        layerGrad           = (dRec + dEdy) `outer` (thisInput # lastOutput)
        -- calcluate derivatives for next step
        (backGrad, recGrad) = split $ tr wN #> (dRec + dEdy)
    in  (BasicRecurrent' biasGradient layerGrad, S1D recGrad, S1D backGrad)

randomBasicRecurrent :: (MonadRandom m, KnownNat i, KnownNat o, KnownNat x, x ~ (i + o))
                     => m (BasicRecurrent i o)
randomBasicRecurrent = do
    seed1 <- getRandom
    seed2 <- getRandom
    let wB = randomVector seed1 Uniform * 2 - 1
        wN = uniformSample seed2 (-1) 1
        bm = konst 0
        mm = konst 0
    return $ BasicRecurrent wB bm wN mm


{-# LANGUAGE CPP               #-}
{-# LANGUAGE OverloadedStrings #-}

module Main ( main ) where

import           Control.Applicative ((<$>))
import qualified Data.Array.Repa     as R
import qualified Data.ByteString     as BS
import           Data.Complex        (Complex (..))
import           Foreign.C.Types     (CDouble, CLLong)
import           Language.J
import           Test.Tasty
import           Test.Tasty.HUnit

main :: IO ()
main = do
#ifdef linux_HOST_OS
    jenv <- jinit libLinux
#else
#ifdef darwin_HOST_OS
    jenv <- jinit (libMac [8,0,7])
#else
#ifdef mingw32_HOST_OS
    jenv <- jinit (libWindows [9,0,1])
#endif
#endif
#endif
    defaultMain $
        testGroup "J dl"
            [ testCase "Performs calculation and has sensible output" (jComp jenv)
            , testCase "Reads back type in the environment" (jType jenv)
            , testCase "Reads a string" (jStr jenv)
            , testCase "Reads a complex array" (jGetC jenv)
            , testCase "Sends an array to J" (jSetA jenv)
            , testCase "Sends a complex array to J" (jSetC jenv)
            , testCase "Uses J to perform a complex calculation" (regress jenv)
            , testCase "Writes strings to J values" (stringRoundtrip jenv)
            -- , testCase "Uses J for something Haskell would have a hard time with" (fill jenv)
            , testCase "Loads a library" (loadNoFail jenv)
            , testCase "Sends an integer array" (printRes jenv)
            ]

loadNoFail :: JEnv -> Assertion
loadNoFail jenv = do
#ifdef linux_HOST_OS
    jLoad jenv (linuxProfile "9.01")
#else
#ifdef darwin_HOST_OS
    jLoad jenv (macProfile [8,0,7])
#else
#ifdef mingw32_HOST_OS
    jLoad jenv (windowsProfile [9,0,1])
#endif
#endif
#endif
    bsDispatch jenv "load'tables/csv'"
    res <- bsOut jenv
    assertBool "Doesn't fail" $
        not ("error: " `BS.isInfixOf` res)

printRes :: JEnv -> Assertion
printRes jenv = do
    setJData jenv "plot_data" $ JIntArr $ R.copyS $ R.map (fromIntegral :: Int -> CLLong) $ R.fromListUnboxed (R.ix1 3) [1,10,2]
    bsDispatch jenv "plot_data"
    res <- bsOut jenv
    res @?= "1 10 2\n"

fill :: JEnv -> Assertion
fill jenv = do
    bsDispatch jenv "random_res =: ? 50 50 $ 1e10"
    res <- getJData jenv "random_res"
    extrExtent res @?= [50, 50]
    where extrExtent :: JData R.DIM2 -> [Int]
          extrExtent (JIntArr res) = R.listOfShape $ R.extent res

regress :: JEnv -> Assertion
regress jenv = do
    let hsArr0 = R.fromListUnboxed (R.ix1 3) [1.0,2.0,3.0]
        hsArr1 = R.fromListUnboxed (R.ix1 3) [2.0,4.0,6.0]
    setJData jenv "xs" (JDoubleArr $ R.copyS $ R.map (realToFrac :: Double -> CDouble) hsArr0)
    setJData jenv "ys" (JDoubleArr $ R.copyS $ R.map (realToFrac :: Double -> CDouble) hsArr1)
    bsDispatch jenv "reg_result =: ys %. xs ^/ i.2"
    res <- getJData jenv "reg_result"
    doubleVect res @?= [5.995204332975845e-15,1.9999999999999971]

stringRoundtrip :: JEnv -> Assertion
stringRoundtrip jenv = do
    setJData jenv "stringy_string" (JString "hello" :: JData R.Z)
    res <- unwrapStr <$> getJData jenv "stringy_string"
    res @?= "hello"

jSetA :: JEnv -> Assertion
jSetA jenv = do
    let hsArr = R.fromListUnboxed (R.ix1 3) [1,3,6]
    setJData jenv "b" (JIntArr $ R.copyS $ R.map (fromIntegral :: Int -> CLLong) hsArr)
    res <- getJData jenv "b"
    intList res @?= [1,3,6]

jSetC :: JEnv -> Assertion
jSetC jenv = do
    let hsArr = R.fromListUnboxed (R.ix1 2) [0.0:+0.0,1.0:+1.0 :: Complex Double]
    let conv = fmap realToFrac :: Complex Double -> Complex CDouble
    setJData jenv "c" (JComplexArr $ R.copyS $ R.map conv hsArr)
    res <- getJData jenv "c"
    complexVect res @?= [0.0:+0.0,1.0:+1.0]

jStr :: JEnv -> Assertion
jStr jenv = do
    bsDispatch jenv "str =: 'hello'"
    res <- getJData jenv "str"
    unwrapStr res @?= "hello"

jGetC :: JEnv -> Assertion
jGetC jenv = do
    bsDispatch jenv "c =: j./~ i.2"
    res <- getJData jenv "c"
    complexVect2 res @?= [0.0:+0.0, 0.0:+1.0, 1.0:+0.0, 1.0:+1.0]

jComp :: JEnv -> Assertion
jComp jenv = do
    bsDispatch jenv "harmonic =: (+/ % #) &.: %"
    bsDispatch jenv "c =: harmonic 1 3 6"
    res <- getJData jenv "c"
    doubleScalar res @?= [2.0]

unwrapStr :: JData R.Z -> BS.ByteString
unwrapStr (JString bs) = bs

doubleVect :: JData R.DIM1 -> [CDouble]
doubleVect (JDoubleArr arr) = R.toList arr

doubleScalar :: JData R.Z -> [CDouble]
doubleScalar (JDoubleArr arr) = R.toList arr

complexVect2 :: JData R.DIM2 -> [Complex CDouble]
complexVect2 (JComplexArr arr) = R.toList arr

complexVect :: JData R.DIM1 -> [Complex CDouble]
complexVect (JComplexArr arr) = R.toList arr

intList :: JData R.DIM1 -> [CLLong]
intList (JIntArr arr) = R.toList arr

jType :: JEnv -> Assertion
jType jenv = do
    bsDispatch jenv "a =: 6?6"
    res <- getJData jenv "a"
    length (intList res) @?= 6


import Debug.Trace
-- import Control.Monad.State.Lazy

import System.Environment (getArgs)
import System.IO (readFile)
import Data.Map.Strict (Map, (!), insert, elems, fromList, toList, findWithDefault, size, empty, member, findMin, findMax, singleton, filter)
import qualified Data.Map.Strict as M
import Data.Set (Set)
import qualified Data.Set as S
--import qualified Data.Array as A
import Data.List (permutations)
import Data.List.Split (splitOn)
import Data.Complex (Complex((:+)), realPart, imagPart)

import UI.NCurses

type Instructions = Map Integer Integer

data ComputerState = Ready | Running | Blocked | Done
  deriving (Show, Eq)

data Computer = Computer {
  state :: ComputerState,
  memory :: Map Integer Integer,
  iptr :: Integer, -- instruction pointer
  base :: Integer, -- base offset
  input :: [Integer],
  output :: [Integer]}

computer0 = Computer { state = Ready, memory = empty, iptr = 0, base = 0, input = [], output = [] }

run :: Computer -> Computer -- (Integer, Integer) -> Instructions -> [Integer] -> [Integer]
run c
  | state c == Ready = run $ c {state = Running}
  | state c == Done = c
  | state c == Blocked = if input c == [] then c else run $ c {state = Running}
  | otherwise = -- running
    --traceShow (iptr c, base c, memory c) $
    case instr `mod` 100 of
      1 -> -- add
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 + arg 2) instructions}
      2 -> -- multiply
        run $ c {iptr = i+4, memory = insert (addr 3) (arg 1 * arg 2) instructions}
      3 -> -- read input
        if (null $ input c)
        then c {state = Blocked}
        else run $ c {iptr = i+2, memory = insert (addr 1) (head $ input c) instructions, input = tail $ input c}
      4 -> -- output
        --traceShow (arg 1) $
        run $ c {iptr = i+2, output = output c ++ [arg 1]}
      5 -> -- jump-if-true
        run $ c {iptr = if arg 1 == 0 then i+3 else arg 2}
      6 -> -- jump-if-false
        run $ c {iptr = if arg 1 == 0 then arg 2 else i+3}
      7 -> -- less than
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 < arg 2 then 1 else 0) instructions}
      8 -> -- equals
        run $ c {iptr = i+4, memory = insert (addr 3) (if arg 1 == arg 2 then 1 else 0) instructions}
      9 -> -- set relative base
        run $ c {iptr = i+2, base = base c + arg 1}
      99 -> -- halt
        run $ c {state = Done}
      _ -> error "unknown opcode"
  where instructions = memory c
        i = iptr c
        instr = instructions!i
        ii x = findWithDefault 0 x instructions
        arg :: Integer -> Integer
        arg n = case (instr `mod` (100*10^n)) `div` (10*10^n) of
          0 -> ii $ ii $ i+n
          1 -> ii $ i+n
          2 -> ii ((ii $ i+n) + base c)
          _ -> error "bad argument mode"
        addr n = case (instr `mod` (100*10^n)) `div` (10*10^n) of
          0 -> ii $ i+n
          1 -> error "address in mode 1"
          2 -> --trace "address in mode 2" $
            ii (i+n) + base c
          _ -> error $ "address in unknown mode"

instance Ord a => Ord (Complex a) where
  compare a b | ar < br || ar == br && ai < bi = LT
              | ar > br || ar == br && ai > bi = GT
              | otherwise = EQ
    where (ar, ai) = toParts a
          (br, bi) = toParts b

toParts :: Complex a -> (a, a)
toParts c = (realPart c, imagPart c)

toIntParts :: RealFrac a => Complex a -> (Int, Int)
toIntParts c = (round $ realPart c, round $ imagPart c)

fromParts :: (a, a) -> Complex a
fromParts (a,b) = a :+ b

data Game = Game
  { m :: Map (Integer, Integer) Integer -- map
  , x :: Integer
  , y :: Integer
  , minx :: Integer
  , miny :: Integer
  , maxx :: Integer
  , maxy :: Integer
  , explore :: [Int] -- stack of directions to explore
  } deriving (Show)

game0 = Game { m = empty, x = 0, y = 0, minx = 0, maxx = 0, miny = 0, maxy = 0, explore = [] }

updateGame :: Game -> Int -> Integer -> Game
updateGame g move status = updateMaxes $ case status of
  0 -> g { m = insert (x',y') 0 $ m g }
  1 -> g { m = insert (x',y') 1 $ m g, x = x', y = y'}
  2 -> g { m = insert (x',y') 2 $ m g, x = x', y = y'}
  where x' | move == 3 = x g - 1
           | move == 4 = x g + 1
           | otherwise = x g
        y' | move == 1 = y g + 1
           | move == 2 = y g - 1
           | otherwise = y g
        updateMaxes g = g
          { maxx = max x' $ maxx g
          , maxy = max y' $ maxy g
          , minx = min x' $ minx g
          , miny = min y' $ miny g
          }

topMargin = 5

drawGame :: Window -> Game -> Curses()
drawGame w g = do
  updateWindow w $ do
    clear
    (width, height) <- windowSize
    let (width', height') = (maxx g - minx g, maxy g - miny g)
    resizeWindow (max width width') (max height height')
    sequence_ [ draw i j val
              | i <- [minx g .. maxx g]
              , j <- [miny g .. maxy g]
              , let val = findWithDefault 5 (i,j) $ m g
              ]
  render
    where draw i j val = do
            let (y',x') =(maxy g - j + topMargin, i - minx g)
            if x' < 0 then error $ "x < 0: " ++ show x'
              else if y' < 0 then error $ "y < 0: " ++ show y'
                   else moveCursor y' x'
            if (i,j) == (x g, y g)
              then drawString "D"
              else drawString $ case val of
                                  5 -> " "
                                  0 -> "#"
                                  1 -> "."
                                  2 -> "*"
                                  _ -> error $  "unknown map value: " ++ show val

suggestMove :: Game -> Int
suggestMove g = search candidates seen
  where seen :: Set (Integer, Integer)
        seen = S.fromList [(x g, y g)]
        candidates :: Set (Int, Int, Integer, Integer) -- steps, direction, x, y
        candidates = S.fromList [ (1, 1, x g, y g + 1)
                                , (1, 2, x g, y g - 1)
                                , (1, 3, x g - 1, y g)
                                , (1, 4, x g + 1, y g)
                                ]
        search :: Set (Int, Int, Integer, Integer) -> Set (Integer, Integer) -> Int
        search candidates seen
          | S.null candidates = 5 -- error "explored all"
          | otherwise = let ((steps, direction, x', y'), candidates') = S.deleteFindMin candidates in
              if (x', y') `S.member` seen then search candidates' seen
              else
                let seen' :: Set (Integer, Integer)
                    seen' = S.insert (x',y') seen
                    candidates'' =
                      foldr S.insert candidates' [(steps+1, direction, x'+i, y'+j)| (i,j) <- [(-1,0),(0,-1),(1,0),(0,1)]]
                in case M.lookup (x',y') (m g) of
                     Nothing -> direction -- not explored at all
                     Just 0 -> search candidates' seen' -- wall
                     _ -> search candidates'' seen'

play :: Game -> Computer -> IO (Game)
play game c = runCurses $ do
  setEcho False
  w <- defaultWindow
  loop w game $ run c
    where loop :: Window -> Game -> Computer -> Curses(Game)
          loop w g c
            | state c == Done = do
                drawGame w g
                updateWindow w $ do
                  moveCursor 0 10
                  drawString "Press q to quit"
                render
                let qloop = do
                      ev <- getEvent w Nothing
                      case ev of
                        Nothing -> qloop
                        Just (EventCharacter 'q') -> return g
                        Just (EventCharacter 'Q') -> return g
                        _ -> qloop
                qloop
            | state c == Blocked = do -- waiting for input
                drawGame w g
                updateWindow w $ do
                  moveCursor 0 10
                  drawString "Move"
                render
                let moveLoop delay = do
                      moveEvent <- getEvent w  $ delay
                      case moveEvent of
                        Just (EventSpecialKey KeyUpArrow) -> return 1
                        Just (EventSpecialKey KeyDownArrow) -> return 2
                        Just (EventSpecialKey KeyLeftArrow) -> return 3
                        Just (EventSpecialKey KeyRightArrow) -> return 4
                        Nothing -> return $ suggestMove g
                        _ -> moveLoop delay
                move <- moveLoop $ Just 0
                if move == 5
                  then loop w g $ c{state = Done}
                  else let c' = run $ c {input = [fromIntegral move], output = []} in
                         loop w (updateGame g move $ head $ output c') $ c'{ input = [], output = [] }

            | otherwise = loop w game $ run c

type Point = (Integer,Integer)
type SearchState = (Int, Point)

main = do
  [instructionFile] <- getArgs
  instructionStrings <- readFile instructionFile
  let instructions = fromList . zip [0 ..] $ map read $ splitOn "," instructionStrings
  putStrLn "Part 1"
  game <- play game0 (computer0{memory = instructions})
  let search :: Set SearchState -> Set Point ->  SearchState
      search stateHeap seen
        | (m game)!(i,j) == 2 = (steps, (i,j)) -- found the goal, return the number of steps needed to reach it.
        | (m game)!(i,j) == 0 = search stateHeap' seen' -- testing a wall, don't add to the frontier
        | otherwise = search stateHeap'' seen' -- testing a clear spot, expand the frontier and keep searching
        where ((steps, (i,j)), stateHeap') = S.deleteFindMin stateHeap
              stateHeap'' = foldr S.insert stateHeap' $ frontier
              frontier = [(steps+1, pt)
                         | pt <- [(i+1,j), (i-1,j), (i,j+1), (i,j-1)]
                         , not $ pt `S.member` seen']
              seen' = S.insert (i,j) seen
  let (steps, oxygenOrigin) = search (S.singleton (0, (0,0))) S.empty
  putStrLn $ show $ (steps, oxygenOrigin)
  putStrLn "Part 2"
  let spread :: Set SearchState -> Set Point -> SearchState
      spread stateHeap seen
        | null stateHeap = (0, (0,0)) -- base case
        | otherwise = max (steps, (i,j)) $ spread stateHeap'' seen' -- testing a clear spot, expand the frontier and keep searching
        where ((steps, (i,j)), stateHeap') = -- traceShow stateHeap $
                S.deleteFindMin stateHeap
              stateHeap'' = foldr S.insert stateHeap' $ frontier
              frontier = [(steps+1, pt)
                         | pt <- [(i+1,j), (i-1,j), (i,j+1), (i,j-1)]
                         , not $ pt `S.member` seen'
                         , not $ (m game)!pt == 0 ]
              seen' = S.insert (i,j) seen
  let (steps, lastSpot) = spread (S.singleton (0, oxygenOrigin)) S.empty
  putStrLn $ show $ (steps, lastSpot)


{-# LANGUAGE FlexibleContexts #-}
module Image.Transform where

import           Data.Array.Repa              as R
import           Data.Array.Repa.Stencil      as R
import           Data.Array.Repa.Stencil.Dim2 as R
import           Data.DList                   as DL
import           Data.List                    as L
import           Data.Vector.Unboxed          as VU
import           Numeric.LinearAlgebra        as NL

-- factor = 2^n, n = 0,1,..
-- the first factor in the list corresponds to the inner-most (right-most) dimension.
{-# INLINE downsample #-}
downsample ::
     (Source s e, Shape sh) => [Int] -> R.Array s sh e -> R.Array D sh e
downsample factorList arr
  | L.all (== 1) factorList = delay arr
  | L.any (< 1) newDList = error "Downsample factors are too large."
  | otherwise =
    R.backpermute
      (shapeOfList newDList)
      (shapeOfList . L.zipWith (*) factorList . listOfShape)
      arr
  where
    dList = listOfShape . extent $ arr
    newDList = L.zipWith div dList factorList

{-# INLINE downsampleUnsafe #-}
downsampleUnsafe ::
     (Source s e, Shape sh) => [Int] -> R.Array s sh e -> R.Array D sh e
downsampleUnsafe factorList arr =
  R.backpermute newSh (shapeOfList . L.zipWith (*) factorList . listOfShape) arr
  where
    dList = listOfShape $ extent arr
    newSh = shapeOfList $ L.zipWith div dList factorList

{-# INLINE upsample #-}
upsample ::
     (Source s e, Shape sh, Num e) => [Int] -> R.Array s sh e -> R.Array D sh e
upsample factorList arr
  | L.all (== 1) factorList = delay arr
  | L.any (< 1) factorList =
    error $ "upsample: factor must be >= 1.\n" L.++ show factorList
  | otherwise =
    R.backpermuteDft
      (fromFunction
         (shapeOfList $ L.zipWith (*) (listOfShape . extent $ arr) factorList)
         (const 0))
      (\shape ->
         if L.all (== 0) . L.zipWith (flip mod) factorList . listOfShape $ shape
           then Just .
                shapeOfList . L.zipWith (flip div) factorList . listOfShape $
                shape
           else Nothing)
      arr

{-# INLINE crop #-}
crop ::
     (Source s e, Shape sh)
  => [Int]
  -> [Int]
  -> R.Array s sh e
  -> R.Array D sh e
crop start len arr
  | L.any (< 0) start ||
      L.or (L.zipWith3 (\x y z -> x > (z - y)) start len dList) =
    error $
    "Crop out of boundary!\n" L.++ show start L.++ "\n" L.++ show len L.++ "\n" L.++
    show dList
  | L.length start /= L.length len || L.length start /= L.length dList =
    error $
    "crop: dimension error. \n start: " L.++ show (L.length start) L.++ " len:" L.++
    show (L.length len) L.++
    " arr:" L.++
    show (L.length dList)
  | otherwise =
    R.backpermute
      (shapeOfList len)
      (shapeOfList . L.zipWith (+) start . listOfShape)
      arr
  where
    dList = listOfShape $ extent arr

{-# INLINE cropUnsafe #-}
cropUnsafe ::
     (Source s e, Shape sh)
  => [Int]
  -> [Int]
  -> R.Array s sh e
  -> R.Array D sh e
cropUnsafe start len =
  R.backpermute
    (shapeOfList len)
    (shapeOfList . L.zipWith (+) start . listOfShape)

{-# INLINE pad #-}
pad :: (Source s e, Shape sh) => [Int] -> e -> R.Array s sh e -> R.Array D sh e
pad newDims padVal arr
  | L.all (== 0) diff = delay arr
  | otherwise =
    backpermuteDft
      (fromFunction (shapeOfList dimList) (const padVal))
      (\sh' ->
         let idx = L.zipWith (-) (listOfShape sh') diff
          in if L.or (L.zipWith (\i j -> i < 0 || (i >= j)) idx oldDimList)
               then Nothing
               else Just $ shapeOfList idx)
      arr
  where
    oldDimList = listOfShape . extent $ arr
    dimList = L.zipWith max newDims oldDimList
    diff =
      L.zipWith
        (\a b ->
           if a - b <= 0
             then 0
             else if a - b == 1
                    then 1
                    else div (a - b) 2)
        newDims
        oldDimList

{-# INLINE normalizeValueRange #-}
normalizeValueRange ::
     (Source r e, Shape sh, Num e, Fractional e, Ord e, Unbox e)
  => (e, e)
  -> R.Array r sh e
  -> R.Array D sh e
normalizeValueRange (minVal, maxVal) arr
  | maxV == minV = delay arr
  | otherwise =
    R.map (\x -> (x - minV) / (maxV - minV) * (maxVal - minVal) + minVal) $ arr
  where
    vec = computeUnboxedS . delay $ arr
    minV = VU.minimum . toUnboxed $ vec
    maxV = VU.maximum . toUnboxed $ vec

{-# INLINE computeDerivativeS #-}
computeDerivativeS ::
     (R.Source r e, Num e, Fractional e) => R.Array r DIM2 e -> R.Array D DIM3 e
computeDerivativeS arr =
  L.foldl' R.append (R.extend (Z :. R.All :. R.All :. (1 :: Int)) arr) ds
  where
    xStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. 0 :. (-1) -> Just (-1)
          Z :. 0 :. 1 -> Just 1
          _ -> Nothing
    yStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. -1 :. 0 -> Just (-1)
          Z :. 1 :. 0 -> Just 1
          _ -> Nothing
    xyStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. -1 :. -1 -> Just 1
          Z :. -1 :. 1 -> Just (-1)
          Z :. 1 :. -1 -> Just (-1)
          Z :. 1 :. 1 -> Just 1
          _ -> Nothing
    ds =
      L.map
        (\s ->
           extend (Z :. R.All :. R.All :. (1 :: Int)) . R.map (/ 2) $
           mapStencil2 BoundClamp s arr)
        [xStencil, yStencil, xyStencil]

{-# INLINE bicubicInterpolation #-}
bicubicInterpolation ::
     (R.Source r Double)
  => ((Int, Int) -> (Double, Double))
  -> (Int, Int)
  -> R.Array r DIM2 Double
  -> R.Array U DIM2 Double
bicubicInterpolation idxFunc (newCols, newRows) arr =
  let (Z :. cols :. rows) = extent arr
      derivative = computeDerivativeS arr
      derivative16 =
        R.traverse
          derivative
          (const (Z :. newCols :. newRows :. 4 :. (4 :: Int)))
          (\f (Z :. a :. b :. derivativeIdx :. pairIdx) ->
             let (x, y) = idxFunc (a, b)
                 (x', y') =
                   case pairIdx of
                     0 -> (floor x, floor y)
                     1 -> (floor x, ceiling y)
                     2 -> (ceiling x, floor y)
                     3 -> (ceiling x, ceiling y)
              in if (x' < 0) ||
                    (x' > (fromIntegral cols - 1)) ||
                    (y' < 0) || (y' > (fromIntegral rows - 1))
                   then 0
                   else f (Z :. x' :. y' :. derivativeIdx))
      mat = ((newRows * newCols) >< 16) . R.toList $ derivative16
      multiplicationResultArray =
        fromListUnboxed (Z :. newCols :. newRows :. 4 :. 4) .
        NL.toList . flatten $
        mat NL.<> maxtrixA
   in R.sumS . R.sumS $
      R.traverse
        multiplicationResultArray
        id
        (\f idx@(Z :. a :. b :. i :. j) ->
           let (x, y) = idxFunc (a, b)
               x' = x - fromIntegral (floor x :: Int)
               y' = y - fromIntegral (floor y :: Int)
            in f idx * (x' ^ i) * (y' ^ j))
  where
    maxtrixA =
      NL.fromColumns . L.map NL.fromList $
      [ [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [-3, 3, 0, 0, -2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [2, -2, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, -3, 3, 0, 0, -2, -1, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 2, -2, 0, 0, 1, 1, 0, 0]
      , [-3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, -3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0]
      , [9, -9, -9, 9, 6, 3, -6, -3, 6, -6, 3, -3, 4, 2, 2, 1]
      , [-6, 6, 6, -6, -3, -3, 3, 3, -4, 4, -2, 2, -2, -2, -1, -1]
      , [2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0]
      , [-6, 6, 6, -6, -4, -2, 4, 2, -3, 3, -3, 3, -2, -1, -2, -1]
      , [4, -4, -4, 4, 2, 2, -2, -2, 2, -2, 2, -2, 1, 1, 1, 1]
      ]

{-# INLINE bilinearInterpolation #-}
bilinearInterpolation ::
     (R.Source r Double)
  => ((Int, Int) -> (Double, Double))
  -> (Int, Int)
  -> R.Array r DIM2 Double
  -> R.Array U DIM2 Double
bilinearInterpolation idxFunc (newCols, newRows) arr =
  let (Z :. cols :. rows) = extent arr
  in computeS . R.traverse arr id $ \f (Z :. a :. b) ->
       let (x, y) = idxFunc (a, b)
           x1 = floor x
           x2 = ceiling x
           y1 = floor y
           y2 = ceiling y
           q11 = f (Z :. x1 :. y1)
           q12 = f (Z :. x1 :. y2)
           q21 = f (Z :. x2 :. y1)
           q22 = f (Z :. x2 :. y2)
           r1 =
             ((fromIntegral x2 - x) / fromIntegral (x2 - x1)) * q11 +
             ((x - fromIntegral x1) / fromIntegral (x2 - x1)) * q21
           r2 =
             ((fromIntegral x2 - x) / fromIntegral (x2 - x1)) * q12 +
             ((x - fromIntegral x1) / fromIntegral (x2 - x1)) * q22
           p =
             ((fromIntegral y2 - y) / fromIntegral (y2 - y1)) * r1 +
             ((y - fromIntegral y1) / fromIntegral (y2 - y1)) * r2
       in if (x < 0) ||
             (x > (fromIntegral cols - 1)) ||
             (y < 0) || (y > (fromIntegral rows - 1))
            then 0
            else if x1 == x2
                   then if y1 == y2
                          then q11
                          else ((fromIntegral y2 - y) / fromIntegral (y2 - y1)) *
                               q11 +
                               ((y - fromIntegral y1) / fromIntegral (y2 - y1)) *
                               q12
                   else if y1 == y2
                          then r1
                          else p

{-# INLINE resize2D #-}
resize2D ::
     (Source s Double)
  => (Int, Int)
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
resize2D newSize@(newNx, newNy) bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (ratioX * fromIntegral i, ratioY * fromIntegral j))
    newSize
    arr
  where
    (Z :. nx' :. ny') = extent arr
    ratioX = fromIntegral nx' / fromIntegral newNx
    ratioY = fromIntegral ny' / fromIntegral newNy

{-# INLINE resize25D #-}
resize25D ::
     (Source s Double)
  => (Int, Int)
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
resize25D newSize@(newNx, newNy) bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (ratioX * fromIntegral i, ratioY * fromIntegral j))
         newSize .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    ratioX = fromIntegral nx' / fromIntegral newNx
    ratioY = fromIntegral ny' / fromIntegral newNy

{-# INLINE resize2DFixedRatio #-}
resize2DFixedRatio ::
     (Source s Double)
  => Int
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
resize2DFixedRatio n bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (ratio * fromIntegral i, ratio * fromIntegral j))
    newSize
    arr
  where
    (Z :. nx' :. ny') = extent arr
    newSize =
      if nx' > ny'
        then (n, round $ fromIntegral ny' * ratio)
        else (round $ fromIntegral nx' * ratio, n)
    ratio = fromIntegral (max nx' ny') / fromIntegral n

{-# INLINE resize25DFixedRatio #-}
resize25DFixedRatio ::
     (Source s Double)
  => Int
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
resize25DFixedRatio n bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (ratio * fromIntegral i, ratio * fromIntegral j))
         newSize .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    newSize@(newNx, newNy) =
      if nx' > ny'
        then (n, round $ fromIntegral ny' * ratio)
        else (round $ fromIntegral nx' * ratio, n)
    ratio = fromIntegral (max nx' ny') / fromIntegral n

{-# INLINE rescale2D #-}
rescale2D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
rescale2D scale bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (fromIntegral i / scale, fromIntegral j / scale))
    (round $ scale * fromIntegral nx', round $ scale * fromIntegral ny')
    arr
  where
    (Z :. nx' :. ny') = extent arr

{-# INLINE rescale25D #-}
rescale25D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
rescale25D scale bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (fromIntegral i / scale, fromIntegral j / scale))
         (newNx, newNy) .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    newNx = round $ scale * fromIntegral nx'
    newNy = round $ scale * fromIntegral ny'

{-# INLINE rotate2D #-}
rotate2D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array U DIM2 Double
rotate2D theta (centerX, centerY) arr =
  bicubicInterpolation
    (\(i, j) ->
       let i' = fromIntegral i - centerX
           j' = fromIntegral j - centerY
        in ( i' * cos theta - j' * sin theta + centerX
           , i' * sin theta + j' * cos theta + centerY))
    (nx', ny')
    arr
  where
    (Z :. nx' :. ny') = extent arr

{-# INLINE rotate25D #-}
rotate25D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array U DIM3 Double
rotate25D theta (centerX, centerY) arr =
  fromUnboxed (Z :. nf' :. nx' :. ny') .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       -- bicubicInterpolation
       bilinearInterpolation
         (\(i, j) ->
            let i' = fromIntegral i - centerX
                j' = fromIntegral j - centerY
            in ( i' * cos theta - j' * sin theta + centerX
               , i' * sin theta + j' * cos theta + centerY))
         (nx', ny') .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr