-----------------------------------------------------------------------------
-- |
-- Module      :  Graphics.Rendering.Plot.Figure.Plot
-- Copyright   :  (c) A. V. H. McPhail 2010
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Creation and manipulation of 'Plot's
--
-----------------------------------------------------------------------------

module Graphics.Rendering.Plot.Figure.Plot (
                                            Plot
                                           -- * Plot elements
                                           , Border
                                           , setBorder
                                           , setPlotBackgroundColour
                                           , setPlotPadding
                                           , withHeading
                                           -- * Series data
                                           , D.Abscissa(), D.Ordinate(), D.Dataset()
                                           , SeriesLabel
                                           , D.FormattedSeries()
                                           , D.line, D.point, D.linepoint
                                           , D.impulse, D.step
                                           , D.area
                                           , D.bar
                                           , D.hist
                                           , D.candle, D.whisker
                                           , setDataset
                                           -- * Annotations
                                           , Location, Head, Fill
                                           , AN.arrow
                                           , AN.oval
                                           , AN.rect
                                           , AN.glyph
                                           , AN.text
                                           , AN.cairo
                                           , withAnnotations
                                           -- ** Plot type
                                           , setSeriesType
                                           , setAllSeriesTypes
                                           -- ** Formatting
                                           , D.PlotFormats(..)
                                           , withSeriesFormat
                                           , withAllSeriesFormats
                                           -- * Range
                                           , Scale(..)
                                           , setRange
                                           , setRangeFromData
                                           -- * Axes
                                           , AX.Axis
                                           , AxisType(..),AxisSide(..),AxisPosn(..)
                                           , clearAxes
                                           , clearAxis
                                           , addAxis
                                           , withAxis
                                           -- * BarSetting
                                           , barSetting
                                           -- * SampleData
                                           , sampleData
                                           -- * Legend
                                           , L.Legend
                                           , LegendBorder
                                           , L.LegendLocation(..), L.LegendOrientation(..)
                                           , clearLegend
                                           , setLegend
                                           , withLegendFormat
                                            -- ** Formatting
                                            , Tick(..), TickValues(..), GridLines
                                           , TickFormat(..)
                                           , AX.setTicks
                                           , AX.setGridlines
                                           , AX.setTickLabelFormat
                                           , AX.setTickLabels
                                           , AX.withTickLabelsFormat
                                           , AX.withAxisLabel
                                           , AX.withAxisLine
                                           , AX.withGridLine
                                           ) where

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

--import Data.Eq.Unicode
--import Data.Bool.Unicode
--import Data.Ord.Unicode

import Numeric.LinearAlgebra.Data hiding(Range)

import qualified Data.Array.IArray as A

import Control.Monad.State
import Control.Monad.Reader
--import Control.Monad.Supply

import Prelude hiding(min,max)
import qualified Prelude as Prelude

import Graphics.Rendering.Plot.Types
import Graphics.Rendering.Plot.Defaults
import qualified Graphics.Rendering.Plot.Figure.Text as T
import qualified Graphics.Rendering.Plot.Figure.Plot.Data as D
import qualified Graphics.Rendering.Plot.Figure.Plot.Axis as AX
import qualified Graphics.Rendering.Plot.Figure.Plot.Legend as L
import qualified Graphics.Rendering.Plot.Figure.Plot.Annotation as AN

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

-- | whether to draw a boundary around the plot area
setBorder :: Border -> Plot ()
setBorder b = modify $ \s -> s { _border = b }

-- | set the plot background colour
setPlotBackgroundColour :: Color -> Plot ()
setPlotBackgroundColour c = modify $ \s -> s { _back_colr = c }

-- | set the padding of the subplot
setPlotPadding :: Double -> Double -> Double -> Double -> Plot ()
setPlotPadding l r b t = modify $ \s -> s { _plot_pads = Padding l r b t }

-- | set the heading of the subplot
withHeading :: Text () -> Plot ()
withHeading m = do
                o <- asks _textoptions
                modify $ \s -> s { _heading = execText m o (_heading s) }

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

-- | set the axis range
setRange :: AxisType -> AxisSide -> Scale -> Double -> Double -> Plot ()
setRange XAxis sd sc min max = modify $ \s -> s { _ranges = setXRanges' (_ranges s) }
    where setXRanges' r
              | sc == Log && min <= 0 = error "non-positive logarithmic range"
              | otherwise          = setXRanges sd r
          setXRanges Lower (Ranges (Left _) yr)       = Ranges (Left (Range sc min max)) yr
          setXRanges Lower (Ranges (Right (_,xr)) yr) = Ranges (Right ((Range sc min max,xr))) yr
          setXRanges Upper (Ranges (Left xr) yr)      = Ranges (Right (xr,Range sc min max)) yr
          setXRanges Upper (Ranges (Right (_,xr)) yr) = Ranges (Right (Range sc min max,xr)) yr
setRange YAxis sd sc min max = modify $ \s -> s { _ranges = setYRanges' (_ranges s) }
    where setYRanges' r
              | sc == Log && min <= 0 = error "non-positive logarithmic range"
              | otherwise          = setYRanges sd r
          setYRanges Lower (Ranges xr (Left _))       = Ranges xr (Left (Range sc min max))
          setYRanges Lower (Ranges xr (Right (_,yr))) = Ranges xr (Right ((Range sc min max,yr)))
          setYRanges Upper (Ranges xr (Left yr))      = Ranges xr (Right (yr,Range sc min max))
          setYRanges Upper (Ranges xr (Right (_,yr))) = Ranges xr (Right ((Range sc min max,yr)))

-- | set the axis ranges to values based on dataset
setRangeFromData :: AxisType -> AxisSide -> Scale -> Plot ()
setRangeFromData ax sd sc = do
  ds <- gets _data
  let ((xmin,xmax),(ymin,ymax)) = calculateRanges ds
  case ax of
    XAxis -> setRange ax sd sc (if sc == Log then if xmin == 0 then 1 else xmin else xmin) xmax
    YAxis -> setRange ax sd sc (if sc == Log then if ymin == 0 then 1 else ymin else ymin) ymax

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

withAnnotations :: Annote () -> Plot ()
withAnnotations = annoteInPlot

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

-- | clear the axes of a subplot
clearAxes :: Plot ()
clearAxes = modify $ \s -> s { _axes = [] }

-- | clear an axis of a subplot
clearAxis :: AxisType -> AxisPosn -> Plot ()
clearAxis at axp = do
  ax <- gets _axes
  modify $ \s -> s { _axes =
    filter (\(Axis at' axp' _ _ _ _ _ _) -> not (at == at' && axp == axp')) ax }

-- | add an axis to the subplot
addAxis :: AxisType -> AxisPosn -> AX.Axis () -> Plot ()
addAxis at axp m  = do
  ax' <- gets _axes
  o <- ask
  let ax = execAxis m o (defaultAxis at axp)
  modify $ \s -> s { _axes = ax : ax' }

-- | operate on the given axis
withAxis :: AxisType -> AxisPosn -> AX.Axis () -> Plot ()
withAxis at axp m = do
  axes' <- gets _axes
  o <- ask
  modify $ \s -> s { _axes =
    map (\a@(Axis at' ap' _ _ _ _ _ _) -> if at == at' && axp == ap'
                                         then execAxis m o a
                                         else a) axes' }

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

barSetting :: BarSetting -> Plot ()
barSetting bc = modify $ \s -> s { _barconfig = bc }

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

sampleData :: SampleData -> Plot ()
sampleData sd = modify $ \s -> s { _sampledata = sd }

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

-- | clear the legend
clearLegend :: Plot ()
clearLegend = withLegend $ L.clearLegend

-- | set the legend location and orientation
setLegend :: L.LegendBorder -> L.LegendLocation -> L.LegendOrientation -> Plot ()
setLegend b l o = withLegend $ L.setLegend b l o

-- | format the legend text
withLegendFormat :: T.Text () -> Plot ()
withLegendFormat f = withLegend $ L.withLegendFormat f

-- | operate on the legend
withLegend :: L.Legend () -> Plot ()
withLegend = legendInPlot

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

-- | operate on the data
withData :: D.Data () -> Plot ()
withData = dataInPlot

{- | set the data series of the subplot

   The data series are either 'FormattedSeries' or plain data series.
   A plain data series must carry a 'SeriesType'.

  A dataset may or may not have an abscissa series, and if so, it is paired
  with either a list of ordinate series or a single ordinate series.

  The abscissa series (if present) is of type 'Vector Double'.

  An ordinate series be a function (@Double -> Double@) or a series of points,
  a 'Vector Double' with optional error series, y axis preference, and labels.

  To specify decoration options for an ordinate series, use the appropriate function, such
  as 'linespoints', with the ordinate series and decoration formatting ('LineFormat',
  'PointFormat', and 'BarFormat') as arguments.

> setDataset (ts,[linespoints (xs,(le,ue),Upper,"data") (([Dash,Dash],3,blue),(Diamond,green))])

  has abscissa @ts@ paired with a list of ordinate series, the single element of which is a
  'FormattedSeries', @linespoints@ where the ordinate is @xs@ with error series @le@ and @ue@,
  to be graphed against the upper y-range with label \"data\".  The line element is formatted
  to be dashed, of width 3, and blue and the point element is to be a green diamond.
-}
setDataset :: D.Dataset a => a -> Plot ()
setDataset d = withData $ D.setDataSeries d

-- | set the plot type of a given data series
setSeriesType :: Int -> SeriesType -> Plot ()
setSeriesType i t = withData $ D.setSeriesType t i

-- | change the plot type of all data series
setAllSeriesTypes :: SeriesType -> Plot ()
setAllSeriesTypes t = withData $ D.setAllSeriesTypes t

-- | format the plot elements of a given series
withSeriesFormat :: D.PlotFormats m => Int -> m () -> Plot ()
withSeriesFormat i f = withData $ D.withSeriesFormat i f

{- |
  format the plot elements of all series

      the operation to modify the formats is passed the series index.
      This allows, for example, colours to be selected from a list
      that gets indexed by the argument

> setColour = withAllSeriesFormats (\i -> do
>                                         setLineColour $ [black,blue,red,green,yellow] !! i
>                                         setLineWidth 1.0)
-}
withAllSeriesFormats :: D.PlotFormats m => (Int -> m ()) -> Plot ()
withAllSeriesFormats f = withData $ D.withAllSeriesFormats f

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

findMinMax :: Abscissae -> Ordinates -> (Double,Double)
findMinMax (AbsFunction _) (OrdFunction _ f _) =
    let v = f (linspace 100 (-1,1))
    in (minElement v,maxElement v)
findMinMax (AbsPoints _ x) (OrdFunction _ f _) =
    let v = f x
    in (minElement v,maxElement v)
-- what if errors go beyond plot?
findMinMax _ (OrdPoints _ (Plain o) _)    = (minElement o,maxElement o)
findMinMax _ (OrdPoints _ (Error o _) _)  = (minElement o,maxElement o)
findMinMax _ (OrdPoints _ (MinMax (o,p) _) _) =
  (Prelude.min (minElement o) (minElement p)
  ,Prelude.max (maxElement o) (maxElement p))

abscMinMax :: Abscissae -> (Double,Double)
abscMinMax (AbsFunction _) = defaultXAxisSideLowerRange
abscMinMax (AbsPoints _ x) = (minElement x,maxElement x)

ordDim :: Ordinates -> Int
ordDim (OrdFunction _ _ _)  = 1
ordDim (OrdPoints _ o _)    = size $ getOrdData o

calculateRanges :: DataSeries -> ((Double,Double),(Double,Double))
calculateRanges (DS_Y ys)      =
  let xmax = maximum $ map (\(DecSeries o _) ->
        fromIntegral $ ordDim o) $ A.elems ys
      ym = unzip $ map (\(DecSeries o _) ->
        findMinMax (AbsFunction id) o) $ A.elems ys
      ymm = (minimum $ fst ym,maximum $ snd ym)
  in ((0,xmax),ymm)
calculateRanges (DS_1toN x ys) =
  let ym = unzip $ map (\(DecSeries o _) -> findMinMax x o) $ A.elems ys
      ymm = (minimum $ fst ym,maximum $ snd ym)
      xmm = abscMinMax x
  in (xmm,ymm)
calculateRanges (DS_1to1 ys)   =
  let (xm',ym') = unzip $ A.elems ys
      ym = unzip $ map (\(x,(DecSeries o _)) -> findMinMax x o) (zip xm' ym')
      ymm = (minimum $ fst ym,maximum $ snd ym)
      xm = unzip $ map abscMinMax xm'
      xmm = (minimum $ fst xm,maximum $ snd xm)
  in (xmm,ymm)
calculateRanges (DS_Surf m)     =
    ((0,fromIntegral $ cols m),(fromIntegral $ rows m,0))

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


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
module ArrayFire.GraphicsSpec where

import Control.Exception
import Data.Complex
import Data.Word
import Foreign.C.Types
import GHC.Int
import Test.Hspec

import ArrayFire

spec :: Spec
spec =
  describe "Graphics tests" $ do
    it "Should create window" $ do
      (1 + 1) `shouldBe` 2


{-# LANGUAGE ScopedTypeVariables, UndecidableInstances #-}
{-|
  Array access, including cell arrays and structures.

  Functions here are primarily thin wrappers to the underlying Matlab functions, and the same memory-management semantics hold.
  In particular, created arrays must be freed, 'copyMXArray' and 'freeMXArray' are deep operations, other set operations do not make copies.
-}

module Foreign.Matlab.Array (
    -- * Array manipulation
    anyMXArray,
    MNullArray, castMNull,
    mxArrayClass,
    mxArrayIsComplex,
    mxArraySize,
    mxArraySetSize,
    mxArrayLength,
    freeMXArray,
    copyMXArray,
    mIndexOffset,

    -- * Array element access
    MXArrayComponent (mxArrayGetOffset, mxArraySetOffset
      , mxArrayGetOffsetList, mxArraySetOffsetList
      , mxScalarGet, isMXScalar
      , createMXArray, createMXScalar
      , createColVector, createRowVector),
    castMXArray,
    -- | array element access
    mxArrayGet, mxArraySet,
    -- | array list access
    mxArrayGetList, mxArraySetList,
    mxArrayGetAll, mxArraySetAll,
    mxArrayGetOffsetSafe, mxArrayGetFirst, mxArrayGetLast,
    -- mxArrayGetSafe, --TODO--
    fromListIO, cellFromListsIO,
    isMNull,

    -- * Struct access
    -- |Structs in Matlab are always arrays, and so can be accessed using most array accessors.
    -- |However, the modification functions such as 'mxArraySet' are not implemented because they could disrupt field data in the entire struct array, and so some specialized functions are necessary.
    MStructArray,
    createStruct,
    mStructFields,
    mStructGet, mStructSet,
    mStructSetFields,
    mStructAddField, mStructRemoveField,
    mxCellGetAllOfType, mxCellGetArraysOfType,

    -- ** Object access
    -- |Some structs are also validated (blessed) user objects.
    mObjectGetClass, mObjectSetClass
  ) where

import           Control.Monad
import           Data.Foldable (toList)
import           Foreign
import           Foreign.C.String
import           Foreign.C.Types
import           Data.Complex
import qualified Data.Map.Strict as DM
import           Data.Maybe (catMaybes)
import           Foreign.Matlab.Util
import           Foreign.Matlab.Internal
import           Foreign.Matlab.Types

#include <matrix.h>

-- |(Un)cast an array to a generic type.
anyMXArray :: MXArray a -> MAnyArray
anyMXArray a = unsafeCastMXArray a

-- |A NULL (empty) array
type MNullArray = MXArray MNull
-- |Safely cast a generic array to a NULL array, or return Nothing if the array is not NULL
castMNull :: MAnyArray -> MIO (Maybe MNullArray)
castMNull a
  | isMNull a = pure $ Just (unsafeCastMXArray a)
  | otherwise = pure Nothing

foreign import ccall unsafe mxGetClassID :: MXArrayPtr -> IO MXClassID
-- |Return the representation of the type of the elements of an array
mxArrayClass :: MXArray a -> IO MXClass
mxArrayClass a
  | isMNull a = pure $ MXClassNull
  | otherwise = withMXArray a mxGetClassID >.= mx2hs

ndims :: MWSize -> Ptr MWSize -> IO MSize
ndims n s = map ii =.< peekArray (ii n) s
--nsubs = ndims

withNSubs :: With MSubs (MWSize, Ptr MWSize) (IO a)
withNSubs l f = withArrayLen (map ii l) (\l a -> f (ii l, a))
withNDims :: With MSize (MWSize, Ptr MWSize) (IO a)
withNDims = withNSubs . realMSize

foreign import ccall unsafe mxGetNumberOfDimensions :: MXArrayPtr -> IO MWSize
foreign import ccall unsafe mxGetDimensions :: MXArrayPtr -> IO (Ptr MWSize)
-- |Get the size (dimensions) of an array
mxArraySize :: MXArray a -> MIO MSize
mxArraySize a = withMXArray a $ \a -> do
  n <- mxGetNumberOfDimensions a
  s <- mxGetDimensions a
  ndims n s

foreign import ccall unsafe mxSetDimensions :: MXArrayPtr -> Ptr MWSize -> MWSize -> IO CInt
-- |Set dimension array and number of dimensions
mxArraySetSize :: MXArray a -> MSize -> IO ()
mxArraySetSize a s = do
  r <- withMXArray a (\a -> withNDims s (\(nd,d) -> mxSetDimensions a d nd))
  when (r /= 0) $ fail "mxArraySetSize"

foreign import ccall unsafe mxGetNumberOfElements :: MXArrayPtr -> IO CSize
-- |Like `numel` in MATLAB.
mxArrayLength :: MXArray a -> MIO Int
mxArrayLength a = ii =.< withMXArray a mxGetNumberOfElements

foreign import ccall unsafe mxCalcSingleSubscript :: MXArrayPtr -> MWSize -> Ptr MWIndex -> IO MWIndex
-- |Convert an array subscript into an offset
mIndexOffset :: MXArray a -> MIndex -> MIO Int
mIndexOffset _ (MSubs []) = pure 0
mIndexOffset _ (MSubs [i]) = pure i
mIndexOffset a (MSubs i) = ii =.< withMXArray a (withNSubs i . uncurry . mxCalcSingleSubscript)

foreign import ccall unsafe mxDuplicateArray :: MXArrayPtr -> IO MXArrayPtr
-- |Make a deep copy of an array
copyMXArray :: MXArray a -> MIO (MXArray a)
copyMXArray a = withMXArray a mxDuplicateArray >>= mkMXArray

foreign import ccall unsafe mxDestroyArray :: MXArrayPtr -> IO ()
-- |Destroy an array and all of its contents.
freeMXArray :: MXArray a -> MIO ()
freeMXArray a = do
  withMXArray a mxDestroyArray
  mxArraySetSize a [0, 0]

-- | Create and populate an MXArray in one go. Named without 'mx' due to possible
-- | conformity to a typeclass function.
fromListIO :: (Foldable t, MXArrayComponent a) => t a -> MIO (MXArray a)
fromListIO xs = do
  arr <- createMXArray [length xs]
  mxArraySetAll arr xsList
  pure arr
  where
    xsList = toList xs

-- | Like fromListIO but wraps elements in a cell. Most useful for converting a list of strings
-- | to a MATLAB cell array of strings. Named in conjunction with `fromListIO`, which is used
-- | as part of the implementation.
cellFromListsIO :: (Traversable s, Foldable t, MXArrayComponent a) => s (t a) -> MIO (MXArray MCell)
cellFromListsIO xss = do
  listOfStructArrays <- sequence $ fromListIO <$> xss
  arr <- createMXArray [length xss]
  mxArraySetAll arr (toList $ (MCell . anyMXArray) <$> listOfStructArrays)
  pure arr

-- |The class of standardly typeable array elements
class MXArrayComponent a where
  -- |Determine whether the given array is of the correct type
  isMXArray :: MXArray a -> MIO Bool
  -- |Create an array and initialize all its data elements to some default value, usually 0 or []
  createMXArray :: MSize -> MIO (MXArray a)

  -- |Determine if an array is singleton. Equivalent to
  --
  -- >  liftM (all (1 ==)) . mxArraySize
  isMXScalar :: MXArray a -> MIO Bool

  mxArrayGetOffset :: MXArray a -> Int -> MIO a
  mxArraySetOffset :: MXArray a -> Int -> a -> MIO ()

  mxArrayGetOffsetList :: MXArray a -> Int -> Int -> MIO [a]
  mxArraySetOffsetList :: MXArray a -> Int -> [a] -> MIO ()

  -- |Get the value of the first data element in an array or, more specifically, the value that the array will be interpreted as in scalar context
  mxScalarGet :: MXArray a -> MIO a

  -- |Create a singleton (scalar) array having the specified value
  createMXScalar :: a -> MIO (MXArray a)
  -- |Create a column vector from the given list.
  createColVector :: [a] -> MIO (MXArray a)
  -- |Create a row vector from the given list.
  createRowVector :: [a] -> MIO (MXArray a)

  isMXArray _ = pure False
  isMXScalar a = liftM2 (&&) (isMXArray a) (all (1 ==) =.< mxArraySize a)

  mxArrayGetOffsetList a o n = mapM (mxArrayGetOffset a) [o..o+n-1]
  mxArraySetOffsetList a o = zipWithM_ (mxArraySetOffset a . (o+)) [0..]

  mxScalarGet a = mxArrayGetOffset a 0

  createMXScalar x = do
    a <- createMXArray [1]
    mxArraySetOffset a 0 x
    pure a
  createRowVector l = do
    a <- createMXArray [1,length l]
    mxArraySetOffsetList a 0 l
    pure a
  createColVector l = do
    a <- createMXArray [length l]
    mxArraySetOffsetList a 0 l
    pure a

-- |Get the value of the specified array element.  Does not check bounds.
mxArrayGet :: MXArrayComponent a => MXArray a -> MIndex -> MIO a
mxArrayGet a i = mIndexOffset a i >>= mxArrayGetOffset a
-- |Set an element in an array to the specified value.  Does not check bounds.
mxArraySet :: MXArrayComponent a => MXArray a -> MIndex -> a -> MIO ()
mxArraySet a i v = do
  o <- mIndexOffset a i
  mxArraySetOffset a o v

-- |@'mxArrayGetList' a i n@ gets the sequential list of @n@ items from array @a@ starting at index @i@.  Does not check bounds.
mxArrayGetList :: MXArrayComponent a => MXArray a -> MIndex -> Int -> MIO [a]
mxArrayGetList a i n = do
  o <- mIndexOffset a i
  n <- if n == -1 then subtract o =.< mxArrayLength a else pure n
  mxArrayGetOffsetList a o n
-- |@'mxArraySetList' a i l@ sets the sequential items in array @a@ starting at index @i@ to @l@.  Does not check bounds.
mxArraySetList :: MXArrayComponent a => MXArray a -> MIndex -> [a] -> MIO ()
mxArraySetList a i l = do
  o <- mIndexOffset a i
  mxArraySetOffsetList a o l

-- |Get a flat list of all elements in the array.
mxArrayGetAll :: MXArrayComponent a => MXArray a -> IO [a]
mxArrayGetAll a = mxArrayGetList a mStart (-1)

-- |Set a flat list of all elements in the array.
mxArraySetAll :: MXArrayComponent a => MXArray a -> [a] -> IO ()
mxArraySetAll a = mxArraySetList a mStart

mxArrayGetFirst :: MXArrayComponent a => MXArray a -> MIO (Either String a)
mxArrayGetFirst arr = mxArrayGetOffsetSafe arr 0

mxArrayGetLast :: MXArrayComponent a => MXArray a -> MIO (Either String a)
mxArrayGetLast arr = do
  arrLen <- mxArrayLength arr
  mxArrayGetOffsetSafe arr (arrLen - 1)

-- |Like mxArrayGetOffset but safe.
mxArrayGetOffsetSafe :: forall a. MXArrayComponent a => MXArray a -> Int -> MIO (Either String a)
mxArrayGetOffsetSafe arr ix
  | isMNull arr = pure $ Left "Couldn't get element of null array"
  | otherwise = do
    arrLen <- mxArrayLength arr
    safeGetElem arrLen ix
  where
    safeGetElem :: Int -> Int -> MIO (Either String a)
    safeGetElem aLen aIx
      | aIx < aLen = Right <$> mxArrayGetOffset arr aIx
      | otherwise = pure $ Left $ "Couldn't get element at index "
        <> (show aIx) <> " of " <> (show aLen) <> "-length array"

-- |Safely cast a generic array to a type, or return Nothing if the array does not have the proper type
castMXArray :: forall a. MXArrayComponent a => MAnyArray -> MIO (Maybe (MXArray a))
castMXArray a
  | isMNull a = pure Nothing
  | otherwise = do
      y <- isMXArray b
      pure $ if y then Just b else Nothing
      where
        b :: MXArray a
        b = unsafeCastMXArray a

-- | Extract all arrays of a given type from a Cell Array.
mxCellGetArraysOfType :: MXArrayComponent a => MXArray MCell -> MIO ([MXArray a])
mxCellGetArraysOfType ca = do
  cellVals <- (fmap . fmap) mCell (mxArrayGetAll ca)
  mxaMays :: [Maybe (MXArray a)] <- sequence $ castMXArray <$> cellVals
  pure $ catMaybes mxaMays

-- | A convenience function to extract all arrays of a given type from a Cell Array;
-- | may have larger dimensions than the original Cell Array due to flattening.
mxCellGetAllOfType :: MXArrayComponent a => MXArray MCell -> MIO [a]
mxCellGetAllOfType ca = do
  as <- mxCellGetArraysOfType ca
  join <$> (sequence $ mxArrayGetAll <$> as)

foreign import ccall unsafe mxGetData :: MXArrayPtr -> IO (Ptr a)

class (MXArrayComponent a, MType mx a, Storable mx) => MXArrayData mx a where
  withArrayData :: MXArray a -> (Ptr mx -> IO b) -> IO b
  withArrayDataOff :: MXArray a -> Int -> (Ptr mx -> IO b) -> IO b
  arrayDataGet :: MXArray a -> Int -> IO a
  arrayDataSet :: MXArray a -> Int -> a -> IO ()

  arrayDataGetList :: MXArray a -> Int -> Int -> IO [a]
  arrayDataSetList :: MXArray a -> Int -> [a] -> IO ()

  withArrayData a f = withMXArray a (mxGetData >=> f)
  withArrayDataOff a o f = withArrayData a (\p -> f (advancePtr p o))
  arrayDataGet a o   = withArrayDataOff a o (mx2hs .=< peek)
  arrayDataSet a o v = withArrayDataOff a o (\p -> poke p (hs2mx v))
  arrayDataGetList a o n = withArrayDataOff a o (map mx2hs .=< peekArray n)
  arrayDataSetList a o l = withArrayDataOff a o (\p -> pokeArray p (map hs2mx l))
#let arrayDataComponent = "\
mxArrayGetOffset = arrayDataGet ;\
mxArraySetOffset = arrayDataSet ;\
mxArrayGetOffsetList = arrayDataGetList ;\
mxArraySetOffsetList = arrayDataSetList\
"
--"

foreign import ccall unsafe mxIsLogical :: MXArrayPtr -> IO CBool
foreign import ccall unsafe mxCreateLogicalArray :: MWSize -> Ptr MWSize -> IO MXArrayPtr
foreign import ccall unsafe mxGetLogicals :: MXArrayPtr -> IO (Ptr MXLogical)
foreign import ccall unsafe mxCreateLogicalScalar :: CBool -> IO MXArrayPtr
foreign import ccall unsafe mxIsLogicalScalar :: MXArrayPtr -> IO CBool
foreign import ccall unsafe mxIsLogicalScalarTrue :: MXArrayPtr -> IO CBool
instance MXArrayComponent MLogical where
  isMXArray a = boolC =.< withMXArray a mxIsLogical
  createMXArray s = withNDims s (uncurry mxCreateLogicalArray) >>= mkMXArray
  createMXScalar = mxCreateLogicalScalar . cBool >=> mkMXArray
  isMXScalar a = boolC =.< withMXArray a mxIsLogicalScalar
  mxScalarGet a = boolC =.< withMXArray a mxIsLogicalScalarTrue
  #arrayDataComponent
instance MXArrayData MXLogical MLogical where
  withArrayData a f = withMXArray a (mxGetLogicals >=> f)

foreign import ccall unsafe mxIsChar :: MXArrayPtr -> IO CBool
foreign import ccall unsafe mxCreateCharArray :: MWSize -> Ptr MWSize -> IO MXArrayPtr
foreign import ccall unsafe mxGetChars :: MXArrayPtr -> IO (Ptr MXChar)
foreign import ccall unsafe mxCreateStringFromNChars :: CString -> MWSize -> IO MXArrayPtr
instance MXArrayComponent MChar where
  isMXArray a = boolC =.< withMXArray a mxIsChar
  createMXArray s = withNDims s (uncurry mxCreateCharArray) >>= mkMXArray
  createRowVector s =
    mkMXArray =<< withCStringLen s (\(s,n) -> mxCreateStringFromNChars s (ii n))
  #arrayDataComponent
instance MXArrayData MXChar MChar where
  withArrayData a f = withMXArray a (mxGetChars >=> f)

foreign import ccall unsafe mxCreateNumericArray :: MWSize -> Ptr MWSize -> MXClassID -> (#type mxComplexity) -> IO MXArrayPtr
createNumericArray :: MXClass -> Bool -> MWSize -> Ptr MWSize -> IO MXArrayPtr
createNumericArray t c n s = mxCreateNumericArray n s (hs2mx t) (if c then (#const mxCOMPLEX) else (#const mxREAL))

#let numarray t = "\
foreign import ccall unsafe mxIs%s :: MXArrayPtr -> IO CBool\n\
instance MXArrayComponent M%s where\n\
  isMXArray a = boolC =.< withMXArray a mxIs%s\n\
  createMXArray s = withNDims s (uncurry $ createNumericArray (mxClassOf (undefined :: M%s)) False) >>= mkMXArray\n\
  \
mxArrayGetOffset = arrayDataGet ;\
mxArraySetOffset = arrayDataSet ;\
mxArrayGetOffsetList = arrayDataGetList ;\
mxArraySetOffsetList = arrayDataSetList\
  \n\
instance MXArrayData MX%s M%s\
", #t, #t, #t, #t, #t, #t

foreign import ccall unsafe mxIsDouble :: MXArrayPtr -> IO CBool
foreign import ccall unsafe mxCreateDoubleScalar :: MXDouble -> IO MXArrayPtr
foreign import ccall unsafe mxGetScalar :: MXArrayPtr -> IO MXDouble
instance MXArrayComponent MDouble where
  isMXArray a = boolC =.< withMXArray a mxIsDouble
  createMXScalar = mxCreateDoubleScalar . hs2mx >=> mkMXArray
  mxScalarGet a = withMXArray a mxGetScalar
  createMXArray s = withNDims s (uncurry $ createNumericArray (mxClassOf (undefined :: Double)) False) >>= mkMXArray
  #arrayDataComponent
instance MXArrayData MXDouble MDouble

#numarray Single
#numarray Int8
#numarray Int16
#numarray Int32
#numarray Int64
#numarray Uint8
#numarray Uint16
#numarray Uint32
#numarray Uint64

foreign import ccall unsafe mxIsCell :: MXArrayPtr -> IO CBool
foreign import ccall unsafe mxCreateCellArray :: MWSize -> Ptr MWSize -> IO MXArrayPtr
foreign import ccall unsafe mxGetCell :: MXArrayPtr -> MWIndex -> IO MXArrayPtr
foreign import ccall unsafe mxSetCell :: MXArrayPtr -> MWIndex -> MXArrayPtr -> IO ()
instance MXArrayComponent MCell where
  isMXArray a = boolC =.< withMXArray a mxIsCell
  createMXArray s = withNDims s (uncurry mxCreateCellArray) >>= mkMXArray
  -- Get a the specified cell element (not a copy).
  mxArrayGetOffset a o = withMXArray a (\a -> mxGetCell a (ii o) >>= mkMXArray >.= MCell)
  -- Set an element in a cell array to the specified value. The cell takes ownership of the array: and no copy is made. Any existing value should be freed first.
  mxArraySetOffset a o (MCell v) = withMXArray a (\a -> withMXArray v (mxSetCell a (ii o)))

-- |A (array of) structs
type MStructArray = MXArray MStruct

foreign import ccall unsafe mxIsStruct :: MXArrayPtr -> IO CBool
foreign import ccall unsafe mxIsObject :: MXArrayPtr -> IO CBool

foreign import ccall unsafe mxCreateStructArray :: MWSize -> Ptr MWSize -> CInt -> Ptr CString -> IO MXArrayPtr
-- |Create an N-Dimensional structure array having the specified fields; initialize all values to 'MNullArray'
createStruct :: MSize -> [String] -> MIO MStructArray
createStruct s f =
  withNDims s (\(nd,d) ->
    mapWithArrayLen withCString f (\(f,nf) ->
      mxCreateStructArray nd d (ii nf) f))
  >>= mkMXArray

foreign import ccall unsafe mxGetNumberOfFields :: MXArrayPtr -> IO CInt
foreign import ccall unsafe mxGetFieldNameByNumber :: MXArrayPtr -> CInt -> IO CString
foreign import ccall unsafe mxGetFieldNumber :: MXArrayPtr -> CString -> IO CInt
-- |Get the names of the fields
mStructFields :: MStructArray -> MIO [String]
mStructFields a = withMXArray a $ \a -> do
  n <- mxGetNumberOfFields a
  forM [0..pred n] (mxGetFieldNameByNumber a >=> peekCString)

foreign import ccall unsafe mxGetField :: MXArrayPtr -> MWIndex -> CString -> IO MXArrayPtr
foreign import ccall unsafe mxSetField :: MXArrayPtr -> MWIndex -> CString -> MXArrayPtr -> IO ()
foreign import ccall unsafe mxGetFieldByNumber :: MXArrayPtr -> MWIndex -> CInt -> IO MXArrayPtr
foreign import ccall unsafe mxSetFieldByNumber :: MXArrayPtr -> MWIndex -> CInt -> MXArrayPtr -> IO ()

-- |Return the contents of the named field for the given element.
-- |Returns 'MNullArray' on no such field or if the field itself is NULL
mStructGet :: MStructArray -> MIndex -> String -> MIO MAnyArray
-- |Sets the contents of the named field for the given element. The input is stored in the array -- no copy is made.
mStructSet :: MStructArray -> MIndex -> String -> MXArray a -> MIO ()
mStructGet a i f = do
  o <- mIndexOffset a i
  withMXArray a (\a -> withCString f (mxGetField a (ii o) >=> mkMXArray))
mStructSet a i f v = do
  o <- mIndexOffset a i
  withMXArray a (\a -> withCString f (withMXArray v . mxSetField a (ii o)))

foreign import ccall unsafe mxAddField :: MXArrayPtr -> CString -> IO CInt
foreign import ccall unsafe mxRemoveField :: MXArrayPtr -> CInt -> IO ()
-- |Add a field to a structure array.
mStructAddField :: MStructArray -> String -> MIO ()
-- |Remove a field from a structure array. Does nothing if no such field exists.
mStructRemoveField :: MStructArray -> String -> MIO ()
mStructAddField a f = do
  i <- withMXArray a (withCString f . mxAddField)
  when (i < 0) $ fail "mxAddField"
mStructRemoveField a f = withMXArray a $ \a -> do
  i <- withCString f (mxGetFieldNumber a)
  if i < 0
    then fail "mxRemoveField"
    else mxRemoveField a i

structGetOffsetFields :: MStructArray -> [String] -> Int -> IO MStruct
structGetOffsetFields a f o =
  MStruct =.< withMXArray a (\a -> DM.fromList <$>
    (zipWithM (\f -> ((,) f) .=< (mxGetFieldByNumber a (ii o) >=> mkMXArray)) f [0..]))

-- |Set the fields of a struct index to the given value list.  The list corresponds to the field list and must match in size.
mStructSetFields :: MStructArray -> MIndex -> [MXArray a] -> MIO ()
mStructSetFields a i v = do
  o <- mIndexOffset a i
  withMXArray a (\a -> zipWithM_ (\v -> withMXArray v . mxSetFieldByNumber a (ii o)) v [0..])

instance MXArrayComponent MStruct where
  isMXArray a = liftM2 (||) (boolC =.< withMXArray a mxIsStruct) (boolC =.< withMXArray a mxIsObject)
  createMXArray s = createStruct s []
  mxArrayGetOffset a o = do
    f <- mStructFields a
    structGetOffsetFields a f o
  mxArraySetOffset = error "mxArraySet undefined for MStruct: use mStructSet"
  mxArrayGetOffsetList a o n = do
    f <- mStructFields a
    mapM (structGetOffsetFields a f) [o..o+n-1]
  createMXScalar (MStruct fv) = do
    a <- createStruct [1] f
    withMXArray a $ \a -> zipWithM_ (\i v -> withMXArray v (mxSetFieldByNumber a 0 i)) [0..] v
    pure a
    where
      (f,v) = unzip $ DM.toList fv

foreign import ccall unsafe mxGetClassName :: MXArrayPtr -> IO CString
-- |Determine if a struct array is a user defined object, and return its class name, if any.
mObjectGetClass :: MStructArray -> IO (Maybe String)
mObjectGetClass a = do
  b <- boolC =.< withMXArray a mxIsObject
  if b
    then Just =.< withMXArray a (mxGetClassName >=> peekCString)
    else pure Nothing

foreign import ccall unsafe mxSetClassName :: MXArrayPtr -> CString -> IO CInt
-- |Set classname of an unvalidated object array.  It is illegal to call this function on a previously validated object array.
mObjectSetClass :: MStructArray -> String -> IO ()
mObjectSetClass a c = do
  r <- withMXArray a (withCString c . mxSetClassName)
  when (r /= 0) $ fail "mObjectSetClass"

castReal :: MXArray (Complex a) -> MXArray a
castReal = unsafeCastMXArray

foreign import ccall unsafe mxGetImagData :: MXArrayPtr -> IO (Ptr a)
withRealDataOff :: MXArrayData mx a => MXArray (Complex a) -> Int -> (Ptr mx -> IO b) -> IO b
withRealDataOff = withArrayDataOff . castReal
withImagDataOff :: MXArrayData mx a => MXArray (Complex a) -> Int -> (Ptr mx -> IO b) -> IO b
withImagDataOff a o f = withMXArray a (mxGetImagData >=> \p -> f (advancePtr p o))

foreign import ccall unsafe mxIsComplex :: MXArrayPtr -> IO CBool
mxArrayIsComplex :: MXArray a -> IO Bool
mxArrayIsComplex a = boolC =.< withMXArray a mxIsComplex

-- |Complex array access.
instance (RealFloat a, MNumeric a, MXArrayData mx a) => MXArrayComponent (MComplex a) where
  isMXArray a = liftM2 (&&) (isMXArray (castReal a)) (mxArrayIsComplex a)
  createMXArray s = withNDims s (uncurry $ createNumericArray (mxClassOf (undefined :: a)) True) >>= mkMXArray

  mxArrayGetOffset a o = do
    r <- withRealDataOff a o (mx2hs .=< peek)
    c <- withImagDataOff a o (mx2hs .=< peek)
    pure $ r :+ c
  mxArraySetOffset a o (r :+ c) = do
    withRealDataOff a o (\p -> poke p (hs2mx r))
    withImagDataOff a o (\p -> poke p (hs2mx c))
  mxArrayGetOffsetList a o n = do
    r <- withRealDataOff a o (map mx2hs .=< peekArray n)
    c <- withImagDataOff a o (map mx2hs .=< peekArray n)
    pure $ zipWith (:+) r c
  mxArraySetOffsetList a o v = do
    withRealDataOff a o (\p -> pokeArray p (map hs2mx r))
    withImagDataOff a o (\p -> pokeArray p (map hs2mx c))
    where (r,c) = unzip $ map (\(r:+c) -> (r,c)) v


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

{-|
Module:      TextShow.Data.Ratio
Copyright:   (C) 2014-2017 Ryan Scott
License:     BSD-style (see the file LICENSE)
Maintainer:  Ryan Scott
Stability:   Provisional
Portability: GHC

'TextShow' instances for 'Ratio'.

/Since: 2/
-}
module TextShow.Data.Complex () where

import Data.Complex (Complex)

import TextShow.Classes (TextShow(..))
import TextShow.Data.Floating ()
import TextShow.TH.Internal (deriveTextShow1, makeShowbPrec)

-- | /Since: 2/
instance TextShow a => TextShow (Complex a) where
    {-# SPECIALIZE instance TextShow (Complex Float)  #-}
    {-# SPECIALIZE instance TextShow (Complex Double) #-}
    showbPrec = $(makeShowbPrec ''Complex)
    {-# INLINE showbPrec #-}

-- | /Since: 2/
$(deriveTextShow1 ''Complex)


module Numeric.TBLIS.Types
  ( TblisError (..),
    TblisLenType,
    TblisStrideType,
    TblisLabelType,
    TblisComm,
    TblisConfig,
    TblisType (..),
    SomeTblisType (..),
    IsTblisType (..),
    TblisScalar (..),
    TblisTensor (..),
    tblisToTypedTensor,
    TblisTypedTensor,
    -- mkTblisTensor,
  )
where

-- import Control.Monad.Primitive
-- import Control.Monad.ST

-- import Data.Primitive.PrimArray
-- import Data.Primitive.Types (Prim)

import Control.Monad (when)
import DLPack (IsDLDataType)
import Data.Complex (Complex)
import qualified Data.Complex
import Data.Kind (Type)
import Data.Proxy
import Data.Text (Text)
import Foreign.C.Types
import Foreign.Marshal.Array (peekArray)
import Foreign.Marshal.Utils (fromBool)
import Foreign.Ptr
import Foreign.Storable
import qualified GHC.Exts as GHC
import GHC.Stack
import GHC.TypeLits
import System.IO.Unsafe (unsafePerformIO)
import Prelude hiding (pred)

natToInt :: forall n. KnownNat n => Int
natToInt = fromIntegral $ GHC.TypeLits.natVal (Proxy @n)
{-# INLINE natToInt #-}

type TblisLenType = CPtrdiff

type TblisStrideType = CPtrdiff

type TblisLabelType = CChar

data TblisComm

data TblisConfig

newtype TblisError = TblisError Text
  deriving stock (Show, Eq)

-- instance Error TblisError where
--   strMsg = TblisError . toText
--   {-# INLINE strMsg #-}

data TblisType :: (Type -> Type) where
  TblisFloat :: TblisType Float
  TblisDouble :: TblisType Double
  TblisComplexFloat :: TblisType (Complex Float)
  TblisComplexDouble :: TblisType (Complex Double)

deriving stock instance Show (TblisType a)

deriving stock instance Eq (TblisType a)

data SomeTblisType where
  SomeTblisType :: IsTblisType a => TblisType a -> SomeTblisType

-- typedef enum
-- {
--     TYPE_SINGLE   = 0,
--     TYPE_FLOAT    = TYPE_SINGLE,
--     TYPE_DOUBLE   = 1,
--     TYPE_SCOMPLEX = 2,
--     TYPE_DCOMPLEX = 3
-- } type_t;
instance Enum (TblisType a) where
  fromEnum x = case x of
    TblisFloat -> 0
    TblisDouble -> 1
    TblisComplexFloat -> 2
    TblisComplexDouble -> 3
  toEnum = error "toEnum not implemented for TblisType"

class (Num a, Storable a, IsDLDataType a) => IsTblisType a where
  tblisTypeOf :: proxy a -> TblisType a

instance IsTblisType Float where tblisTypeOf _ = TblisFloat

instance IsTblisType Double where tblisTypeOf _ = TblisDouble

instance IsTblisType (Complex Float) where tblisTypeOf _ = TblisComplexFloat

instance IsTblisType (Complex Double) where tblisTypeOf _ = TblisComplexDouble

newtype TblisScalar a = TblisScalar a
  deriving stock (Read, Show)
  deriving newtype (Eq, Ord, Num)

instance IsTblisType a => Storable (TblisScalar a) where
  sizeOf _ = 16 + 8
  {-# INLINE sizeOf #-}
  alignment _ = 8
  {-# INLINE alignment #-}
  peek p = case tblisTypeOf (Proxy :: Proxy a) of
    TblisFloat -> TblisScalar <$> peekByteOff p 0
    TblisDouble -> TblisScalar <$> peekByteOff p 0
    TblisComplexFloat -> TblisScalar <$> peekByteOff p 0
    TblisComplexDouble -> TblisScalar <$> peekByteOff p 0
  {-# INLINE peek #-}
  poke p (TblisScalar x) = do
    pokeByteOff p 0 x
    pokeByteOff p 16
      . (fromIntegral :: Int -> CInt)
      . fromEnum
      . tblisTypeOf
      $ (Proxy :: Proxy a)
  {-# INLINE poke #-}

data TblisTensor a = TblisTensor
  { tblisTensorConj :: !Bool,
    tblisTensorScalar :: {-# UNPACK #-} !(TblisScalar a),
    tblisTensorData :: {-# UNPACK #-} !(Ptr a),
    tblisTensorNDim :: {-# UNPACK #-} !Int,
    tblisTensorLen :: {-# UNPACK #-} !(Ptr TblisLenType),
    tblisTensorStride :: {-# UNPACK #-} !(Ptr TblisStrideType)
  }

-- type_t type;
-- int conj;
-- tblis_scalar scalar;
-- void* data;
-- unsigned ndim;
-- len_type* len;
-- stride_type* stride;
instance IsTblisType a => Storable (TblisTensor a) where
  sizeOf _ = 4 + 4 + 24 + 8 + 8 + 8 + 8
  alignment _ = 8
  poke p x = do
    pokeByteOff p 0 . (fromIntegral :: Int -> CInt) . fromEnum $ tblisTypeOf (Proxy @a)
    pokeByteOff p 4 . (fromBool :: Bool -> CInt) $ tblisTensorConj x
    pokeByteOff p 8 $ tblisTensorScalar x
    pokeByteOff p 32 $ tblisTensorData x
    pokeByteOff p 40 . (fromIntegral :: Int -> CUInt) $ tblisTensorNDim x
    pokeByteOff p 48 $ tblisTensorLen x
    pokeByteOff p 56 $ tblisTensorStride x
  peek _ = error $ "peek is not implemented for TblisTensor"

newtype TblisTypedTensor (rank :: Nat) (a :: Type) = TblisTypedTensor (TblisTensor a)

tblisToTypedTensor :: forall r a. (HasCallStack, KnownNat r) => TblisTensor a -> TblisTypedTensor r a
tblisToTypedTensor t
  | tblisTensorNDim t == natToInt @r = TblisTypedTensor t
  | otherwise =
    error $
      "tensor has wrong rank: " <> show (tblisTensorNDim t)
        <> "; expected "
        <> show (natToInt @r)

conjugateScalar :: forall a. IsTblisType a => TblisScalar a -> TblisScalar a
conjugateScalar (TblisScalar z) = TblisScalar $
  case tblisTypeOf (Proxy @a) of
    TblisFloat -> z
    TblisDouble -> z
    TblisComplexFloat -> Data.Complex.conjugate z
    TblisComplexDouble -> Data.Complex.conjugate z
{-# INLINE conjugateScalar #-}

tensorShape :: TblisTensor a -> [Int]
tensorShape t = unsafePerformIO $ do
  let helper i = let !i' = fromIntegral i in i'
  fmap helper <$> peekArray (tblisTensorNDim t) (tblisTensorLen t)

tensorIndex :: IsTblisType a => TblisTensor a -> [Int] -> IO a
tensorIndex t index = do
  stride <- fmap fromIntegral <$> peekArray (tblisTensorNDim t) (tblisTensorStride t)
  let shape = tensorShape t
  -- <- fmap fromIntegral <$> peekArray (tblisTensorNDim t) (tblisTensorLen t)
  when (any id $ zipWith (\i n -> i < 0 || i >= n) index shape) $
    error $ "invalid index: " <> show index
  let !linearIndex = sum $ zipWith (*) stride index
  !s <- (tblisTensorScalar t *) . TblisScalar <$> peekElemOff (tblisTensorData t) linearIndex
  let (TblisScalar !z) = if (tblisTensorConj t) then conjugateScalar s else s
  pure z
{-# NOINLINE tensorIndex #-}

flattenTensorToList :: IsTblisType a => TblisTensor a -> [a]
flattenTensorToList t = unsafePerformIO $ do
  -- <- fmap fromIntegral <$> peekArray (tblisTensorNDim t) (tblisTensorLen t)
  let !shape = tensorShape t
      indices = sequence . fmap (\n -> [0 .. n - 1]) $ shape
      f i = do !x <- tensorIndex t i; pure x
  mapM f indices

-- tensorToList1D :: forall a. Prim a => TblisTypedTensor 1 a -> [a]
-- tensorToList1D t = runST $
--   withForeignPtr' (tensorData t) $ \p ->
--     go p [] (stride * (extent - 1))
--   where
--     !stride = indexPrimArray (tensorStrides t) 0
--     !extent = indexPrimArray (tensorShape t) 0
--     go :: PrimMonad m => Ptr a -> [a] -> Int -> m [a]
--     go !p acc !i
--       | i >= 0 = do
--         !x <- P.readOffPtr p i
--         go p (x : acc) (i - stride)
--       | otherwise = touch t >> pure acc
--
-- tensorFromList1D :: forall a. (HasCallStack, Prim a) => Int -> [a] -> Tensor 'CPU 1 a
-- tensorFromList1D n xs = unsafePerformIO $ do
--   t <- newTensor [n]
--   withForeignPtr (tensorData t) $ \p -> do
--     let go !i []
--           | i == n = pure ()
--           | otherwise = error $ "list is shorter than expected"
--         go !i (y : ys)
--           | i < n = P.writeOffPtr p i y >> go (i + 1) ys
--           | otherwise = error "list is longer than expected"
--     go 0 xs
--   pure t
--
-- instance IsTblisType a => GHC.IsList (TblisTypedTensor 1 a) where
--   type Item (TblisTypedTensor 1 a) = a
--   toList = tensorToList1D
--   fromList xs = tensorFromList1D (length xs) xs
--   fromListN n xs = tensorFromList1D n xs

splitAt' :: HasCallStack => Int -> [a] -> ([a], [a])
splitAt' = go
  where
    go 0 [] = ([], [])
    go 1 (x : xs) = ([x], xs)
    go m (x : xs) = (x : xs', xs'')
      where
        (xs', xs'') = go (m - 1) xs
    go _ [] = error "wrong list length"

chunksOf :: Int -> [a] -> [[a]]
chunksOf n = go
  where
    go [] = []
    go xs@(_ : _) = let (ys, zs) = splitAt' n xs in ys : go zs

-- listShape2D :: [[a]] -> [Int]
-- listShape2D [] = [0, 0]
-- listShape2D xs@(x : _) = [length xs, length x]
-- {-# INLINE listShape2D #-}

-- listShape3D :: [[[a]]] -> [Int]
-- listShape3D [] = [0, 0, 0]
-- listShape3D xs@(x : _) = length xs : listShape2D x
--
-- listShape4D :: [[[a]]] -> [Int]
-- listShape4D [] = [0, 0, 0, 0]
-- listShape4D xs@(x : _) = length xs : listShape3D x
--
-- listShape5D :: [[[[a]]]] -> [Int]
-- listShape5D [] = [0, 0, 0, 0, 0]
-- listShape5D xs@(x : _) = length xs : listShape4D x

instance IsTblisType a => GHC.IsList (TblisTypedTensor 1 a) where
  type Item (TblisTypedTensor 1 a) = a
  toList (TblisTypedTensor t) = flattenTensorToList t
  fromList _ = error "IsList instance of TblisTypedTensor does not implement fromList"

instance IsTblisType a => GHC.IsList (TblisTypedTensor 2 a) where
  type Item (TblisTypedTensor 2 a) = [a]
  toList (TblisTypedTensor t) = case tensorShape t of
    [_, d₁] -> chunksOf d₁ (flattenTensorToList t)
    _ -> error "impossible: tensor must be rank 2"
  fromList _ = error "IsList instance of TblisTypedTensor does not implement fromList"

instance IsTblisType a => GHC.IsList (TblisTypedTensor 3 a) where
  type Item (TblisTypedTensor 3 a) = [[a]]
  toList (TblisTypedTensor t) = case tensorShape t of
    [_, d₁, d₂] -> chunksOf d₁ . chunksOf d₂ $ flattenTensorToList t
    _ -> error "impossible: tensor must be rank 3"
  fromList _ = error "IsList instance of TblisTypedTensor does not implement fromList"

-- peek p =
--   TblisTensor <$> peekByteOff p 4

-- mkShape :: (Prim a, Integral a) => (Int -> Bool) -> Int -> [Int] -> PrimArray a
-- mkShape pred n xs = runST $ do
--   arr <- newAlignedPinnedPrimArray n
--   let go [] !i
--         | i == n = return ()
--         | otherwise = error $ "list length is smaller than the specified size"
--       go (y : ys) !i
--         | i < n && pred y = writePrimArray arr i (fromIntegral y) >> go ys (i + 1)
--         | i >= n = error $ "list length is greated than the specified size"
--         | otherwise = error $ "encountered an invalid element: " <> show y
--   go xs 0
--   unsafeFreezePrimArray arr
--
-- mkLenAndStride :: [Int] -> [Int] -> (PrimArray TblisLenType, PrimArray TblisStrideType)
-- mkLenAndStride shape strides = (mkShape (>= 0) n shape, mkShape (> 0) n strides)
--   where
--     !n = length shape
--
-- mkTblisTensor :: IsTblisType a => Ptr a -> [Int] -> [Int] -> TblisTensor a
-- mkTblisTensor p shape strides =
--   let (len, stride) = mkLenAndStride shape strides
--    in TblisTensor False (TblisScalar 1) p len stride


{-# LANGUAGE BinaryLiterals            #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE NegativeLiterals          #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ScopedTypeVariables       #-}

-- | Tests for the flat module
module Main where

import           Control.Monad
import           Data.Bits
import qualified Data.ByteString          as B
import qualified Data.ByteString.Lazy     as L
import qualified Data.ByteString.Short    as SBS
import           Data.Char
import           Data.Either
import           Data.Flat
import           Data.Flat.Bits
import           Data.Flat.Decoder
import qualified Data.Flat.Encoder        as E
import qualified Data.Flat.Encoder.Prim   as E
import qualified Data.Flat.Encoder.Strict as E
import           Data.Int
import           Data.Proxy
import qualified Data.Sequence            as Seq
import qualified Data.Text                as T
import           Data.Word
import           Numeric.Natural
import           System.Exit
import           Test.Data
import           Test.Data.Arbitrary()
import           Test.Data.Flat
import           Test.Data.Values         hiding (lbs, ns)
import           Test.E
import           Test.E.Arbitrary()
import           Test.E.Flat
import           Test.Tasty
import           Test.Tasty.HUnit
import           Test.Tasty.QuickCheck    as QC hiding (getSize)
import           Data.Flat.Endian
import           Data.FloatCast
import Data.Text.Arbitrary
-- import Test.QuickCheck.Arbitrary
import qualified Data.Complex as B
import qualified Data.Ratio as B

import qualified Data.Map                 as C
import qualified Data.Map.Strict                 as CS
import qualified Data.Map.Lazy                 as CL
import qualified Data.IntMap.Strict as CS
import qualified Data.IntMap.Lazy  as CL

-- import           Data.List
-- import           Data.Ord

#if MIN_VERSION_base(4,9,0)
import qualified Data.List.NonEmpty as BI
#endif

instance Arbitrary UTF8Text where
  arbitrary = UTF8Text <$> arbitrary
  shrink t = UTF8Text <$> shrink (unUTF8 t)

instance Arbitrary UTF16Text where
    arbitrary = UTF16Text <$> arbitrary
    shrink t = UTF16Text <$> shrink (unUTF16 t)

-- instance Flat [Int16]
-- instance Flat [Word8]
-- instance Flat [Bool]

main = do
-- #ifdef ghcjs_HOST_OS
--   print "GHCJS"
-- #endif

  -- printInfo

  -- print $ flat asciiStrT

  mainTest
  -- print $ flatRaw 18446744073709551615::Word64
  -- print $ B.unpack . flat $ (True,0::Word64,18446744073709551615::Word64)
  -- print (2^56::Word64,fromIntegral (1::Word8) `shiftL` 56 :: Word64,(18446744073709551615::Word64) `shiftR` 1)
  -- mainShow
  -- eWord64E id 0b

mainShow = do
  mapM_ (\_ -> generate (arbitrary :: Gen Int) >>= print) [1..10]
  exitFailure

mainTest = defaultMain tests

tests :: TestTree
tests = testGroup "Tests" [
  testPrimitives

  ,testEncDec

  ,testFlat
 ]

testPrimitives = testGroup "conversion/memory primitives" [
   testEndian
  ,testFloatingConvert
  --,testShifts -- ghcjs fails this
  ]

testEncDec = testGroup "encode/decode primitives" [
   testEncodingPrim
  ,testDecodingPrim
#ifdef TEST_DECBITS
  ,testDecBits
#endif
  ]

testFlat = testGroup "flat/unflat" [
  testSize
  ,testLargeEnum
  ,testContainers
  ,flatUnflatRT
  ,flatTests
  ]

-- Data.Flat.Endian tests (to run, need to modify imports and cabal file)
testEndian = testGroup "Endian" [
   conv toBE16 (2^10 + 3)  (2^9+2^8+4)
  ,conv toBE32 (2^18 + 3)  50332672
  ,conv toBE64 (2^34 + 3)  216172782180892672
  ,conv toBE16 0x1234 0x3412
  ,conv toBE32 0x11223344 0x44332211
  ,conv toBE64 0x0123456789ABCDEF 0xEFCDAB8967452301
  ]

testFloatingConvert = testGroup "Floating conversions" [
   conv floatToWord -0.15625 3189768192
  ,conv wordToFloat 3189768192 -0.15625
  ,conv doubleToWord -0.15625 13818169556679524352
  ,conv wordToDouble 13818169556679524352 -0.15625
  ,rt "floatToWord" (prop_float_conv :: RT Float)
  ,rt "doubleToWord" (prop_double_conv :: RT Double)
 ]

-- ghcjs bug on shiftR 0, see: https://github.com/ghcjs/ghcjs/issues/706
testShifts = testGroup "Shifts" $ map tst [0..33]
   where
 tst n = testCase ("shiftR " ++ show n) $
  let val = 4294967295::Word32
      s = val `shift` (-n)
      r = val `shiftR` n
  in r @?= s

-- shR = shiftR
-- shR = unsafeShiftR
shR val 0 = val
shR val n = shift val (-n)

testEncodingPrim = testGroup "Encoding Primitives" [
   encRawWith 1 E.eTrueF [0b10000001]
  ,encRawWith 3 (E.eTrueF >=> E.eFalseF >=> E.eTrueF) [0b10100001]

  -- Depends on endianess
  --,encRawWith 32 (E.eWord32E id $ 2^18 + 3) [3,0,4,0,1]
  -- ,encRawWith 64 (E.eWord64E id $ 0x1122334455667788) [0x88,0x77,0x66,0x55,0x44,0x33,0x22,0x11,1]
  --,encRawWith 65 (E.eTrueF >=> E.eWord64E id (2^34 + 3)) [1,0,0,0,2,0,0,128,129]
  --,encRawWith 65 (E.eFalseF >=> E.eWord64E id (2^34 + 3)) [1,0,0,0,2,0,0,0,129]

  -- Big Endian
  ,encRawWith 32 (E.eWord32BEF  $ 2^18 + 3) [0,4,0,3,1]
  ,encRawWith 64 (E.eWord64BEF  $ 2^34 + 3) [0,0,0,4,0,0,0,3,1]
  ,encRawWith 65 (E.eTrueF >=> E.eWord64BEF (2^34 + 3)) [128,0,0,2,0,0,0,1,129]
  ,encRawWith 65 (E.eFalseF >=> E.eWord64BEF (2^34 + 3))  [0,0,0,2,0,0,0,1,129]
  ]
  where
    encRawWith sz enc exp = testCase (unwords ["encode raw with size",show sz]) $ flatRawWith sz enc @?= exp

conv f v e = testCase (unwords ["conv",sshow v,showB . flat $ v,"to",sshow e]) $ f v @?= e

testDecodingPrim = testGroup "Decoding Primitives" [
   dec ((,,,) <$> dropBits 13 <*> dBool <*> dBool <*> dBool) [0b10111110,0b10011010] ((),False,True,False)
  ,dec ((,,,) <$> dropBits 1 <*> dBE16 <*> dBool <*> dropBits 6) [0b11000000
                                                                 ,0b00000001
                                                                 ,0b01000000] ((),2^15+2,True,())
  ,dec ((,,,) <$> dropBits 1 <*> dBE32 <*> dBool <*> dropBits 6) [0b11000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000001
                                                                 ,0b01000000] ((),2^31+2,True,())
  ,dec dBE64 [0b10000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000010
                                                                 ] (2^63+2)

  ,dec ((,,,) <$> dropBits 1 <*> dBE64 <*> dBool <*> dropBits 6) [0b11000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000000
                                                                 ,0b00000001
                                                                 ,0b01000000] ((),2^63+2,True,())

  ]
    where
      dec decOp v e = testCase (unwords ["decode",sshow v]) $ unflatRawWith decOp (B.pack v) @?= Right e

testDecBits = testGroup "Decode Bits" $ concat [
   decBitsN dBEBits8
  ,decBitsN dBEBits16
  ,decBitsN dBEBits32
  ,decBitsN dBEBits64
 ]
  where
          -- Test dBEBits8/16/32/64, extraction of up to 8/16/32/bits from various positions
          decBitsN :: forall a. (Num a,FiniteBits a,Show a,Flat a) => (Int -> Get a) -> [TestTree]
          decBitsN dec = let s = finiteBitSize (undefined::a)
                         in [decBits_ dec val numBitsToTake pre | numBitsToTake <- [0 .. s], val <- [0::a ,1+2^(s - 2)+2^(s - 5) ,fromIntegral $ (2^s::Integer) - 1],pre <- [0,1,7]]

          decBits_ :: forall a. (FiniteBits a,Show a,Flat a) => (Int -> Get a) -> a -> Int -> Int -> TestTree
          decBits_ deco val numBitsToTake pre =
            -- a sequence composed by pre zero bits followed by the val and zero bits till the next byte boundary
            let vs = B.pack . asBytes . fromBools $ replicate pre False ++ toBools (asBits val)
                len = B.length vs
                sz = finiteBitSize (undefined::a)
                dec :: Get a
                dec = do
                  dropBits pre
                  r <- deco numBitsToTake
                  dropBits (len*8-numBitsToTake-pre)
                  return r
                -- we expect the first numBitsToTake bits of the value
                expectedD@(Right expected) :: Decoded a = Right $ val `shR` (sz - numBitsToTake) -- ghcjs: shiftR fails, see: https://github.com/ghcjs/ghcjs/issues/706
                actualD@(Right actual) :: Decoded a = unflatRawWith dec vs
            in testCase (unwords ["take",show numBitsToTake,"bits from",show val,"of size",show sz,"with prefix",show pre,"sequence",showB vs,show expected,show actual,show $ val == actual,show $ expected == actual,show $ expected /= actual,show $ show expected == show actual,show $ flat expected == flat actual])
                $ actualD @?= expectedD

testSize = testGroup "Size" $ concat [
  sz () 0
  ,sz True 1
  ,sz One 2
  ,sz Two 2
  ,sz Three 2
  ,sz Four 3
  ,sz Five 3
  ,sz 'a' 8
  ,sz 'à' 16
  ,sz '经' 24
  ,sz (0::Word8) 8
  ,sz (1::Word8) 8
  ,concatMap (uncurry sz) ns
  ,concatMap (uncurry sz) nsI
  ,concatMap (uncurry sz) nsII
  ,sz (1.1::Float) 32
  ,sz (1.1::Double) 64
  ,sz "" 1
  ,sz "abc" (4+3*8)
  ,sz ((),(),Unit) 0
  ,sz (True,False,One,Five) 7
  ,sz map1 7
  ,sz bs (4+3*8)
  ,sz stBS bsSize
  ,sz lzBS bsSize
#ifndef ghcjs_HOST_OS
  ,sz shBS bsSize
#endif
  ,sz tx utf8Size
  ,sz (UTF8Text tx) utf8Size
  ,sz (UTF16Text tx) utf16Size
 ]
   where
    tx = T.pack "txt"
    utf8Size = 8+8+3*32+8
    utf16Size = 8+8+3*16+8
    bsSize = 8+8+3*8+8

sz v e = [testCase (unwords ["size of",sshow v]) $ getSize v @?= e]

-- E258_256 = 11111110 _257 = 111111110 _258 = 111111111
testLargeEnum = testGroup "test enum with more than 256 constructors" $ concat
  [
#ifdef ENUM_LARGE
      sz E258_256 8
    , sz E258_257 9
    , sz E258_258 9

    -- As encodes are inlined, this is going to take for ever if this is compiled with -O1 or -O2
    -- , encRaw (E258_256) [0b11111110]
    -- , encRaw (E258_257) [0b11111111,0b00000000]
    -- , encRaw (E258_258) [0b11111111,0b10000000]
    -- , encRaw (E258_256,E258_257,E258_258) [0b11111110,0b11111111,0b01111111,0b11000000]

    , map trip [E258_1,E258_256,E258_257,E258_258]
    , map trip [E256_1,E256_134,E256_256]
#endif
  ]

testContainers = testGroup "containers" [
    trip longSeq
    , trip dataMap
    , trip listMap
    -- , trip intMap
    ]

flatUnflatRT = testGroup "unflat (flat v) == v"
  [  rt "()" (prop_Flat_roundtrip:: RT ())
    ,rt "Bool" (prop_Flat_roundtrip::RT Bool)
    ,rt "Char" (prop_Flat_roundtrip:: RT Char)
    ,rt "Complex" (prop_Flat_roundtrip:: RT (B.Complex Float))
    ,rt "Either N Bool" (prop_Flat_roundtrip:: RT (Either N Bool))
    ,rt "Either Int Char" (prop_Flat_roundtrip:: RT (Either Int Char))
    ,rt "Int8" (prop_Flat_Large_roundtrip:: RTL Int8)
    ,rt "Int16" (prop_Flat_Large_roundtrip:: RTL Int16)
    ,rt "Int32" (prop_Flat_Large_roundtrip:: RTL Int32)
    ,rt "Int64" (prop_Flat_Large_roundtrip:: RTL Int64)
    ,rt "Int" (prop_Flat_Large_roundtrip:: RTL Int)
    ,rt "[Int16]" (prop_Flat_roundtrip:: RT [Int16])
    ,rt "String" (prop_Flat_roundtrip:: RT String)
#if MIN_VERSION_base(4,9,0)
    ,rt "NonEmpty" (prop_Flat_roundtrip:: RT (BI.NonEmpty Bool))
#endif
    ,rt "Maybe N" (prop_Flat_roundtrip:: RT (Maybe N))
    ,rt "Ratio" (prop_Flat_roundtrip:: RT (B.Ratio Int32))
    ,rt "Word8" (prop_Flat_Large_roundtrip:: RTL Word8)
    ,rt "Word16" (prop_Flat_Large_roundtrip:: RTL Word16)
    ,rt "Word32" (prop_Flat_Large_roundtrip:: RTL Word32)
    ,rt "Word64" (prop_Flat_Large_roundtrip:: RTL Word64)
    ,rt "Word" (prop_Flat_Large_roundtrip:: RTL Word)
    ,rt "Natural" (prop_Flat_roundtrip:: RT Natural)
    ,rt "Integer" (prop_Flat_roundtrip:: RT Integer)
    ,rt "Float" (prop_Flat_roundtrip:: RT Float)
    ,rt "Double" (prop_Flat_roundtrip:: RT Double)

    ,rt "Text" (prop_Flat_roundtrip:: RT T.Text)
    ,rt "UTF8 Text" (prop_Flat_roundtrip:: RT UTF8Text)
    ,rt "UTF16 Text" (prop_Flat_roundtrip:: RT UTF16Text)

    ,rt "ByteString" (prop_Flat_roundtrip:: RT B.ByteString)
    ,rt "Lazy ByteString" (prop_Flat_roundtrip:: RT L.ByteString)
#ifndef ghcjs_HOST_OS
    ,rt "Short ByteString" (prop_Flat_roundtrip:: RT SBS.ShortByteString)
#endif

    ,rt "Map.Strict" (prop_Flat_roundtrip:: RT (CS.Map Int Bool))
    ,rt "Map.Lazy" (prop_Flat_roundtrip:: RT (CL.Map Int Bool))
    ,rt "IntMap.Strict" (prop_Flat_roundtrip:: RT (CS.IntMap Bool))
    ,rt "IntMap.Lazy" (prop_Flat_roundtrip:: RT (CL.IntMap Bool))

    ,rt "Unit" (prop_Flat_roundtrip:: RT Unit)
    ,rt "Un" (prop_Flat_roundtrip:: RT Un )
    ,rt "N" (prop_Flat_roundtrip:: RT N )
    ,rt "E2" (prop_Flat_roundtrip:: RT E2 )
    ,rt "E3" (prop_Flat_roundtrip:: RT E3 )
    ,rt "E4" (prop_Flat_roundtrip:: RT E4 )
    ,rt "E8" (prop_Flat_roundtrip:: RT E8 )
    ,rt "E16" (prop_Flat_roundtrip:: RT E16 )
    ,rt "E17" (prop_Flat_roundtrip:: RT E17 )
    ,rt "E32" (prop_Flat_roundtrip:: RT E32 )
    ,rt "A" (prop_Flat_roundtrip:: RT A )
    ,rt "B" (prop_Flat_roundtrip:: RT B )
    -- ,rt "Tree Bool" (prop_Flat_roundtrip:: RT (Tree Bool))
    -- ,rt "Tree N" (prop_Flat_roundtrip:: RT (Tree N))
    ,rt "List N" (prop_Flat_roundtrip:: RT (List N))
    ]

rt n = QC.testProperty (unwords ["round trip",n])

flatTests = testGroup "flat/unflat Unit tests" $ concat [

  -- Expected errors
   errDec (Proxy::Proxy Bool) [] -- no data
  ,errDec (Proxy::Proxy Bool) [128] -- no filler
  ,errDec (Proxy::Proxy Bool) [128+1,1,2,4,8] -- additional bytes

  ,encRaw () []
  ,encRaw ((),(),Unit) []
  ,encRaw (Unit,'a',Unit,'a',Unit,'a',Unit) [97,97,97]
  ,a () [1]
  ,a True [128+1]
  ,a (True,True) [128+64+1]
  ,a (True,False,True) [128+32+1]
  ,a (True,False,True,True) [128+32+16+1]
  ,a (True,False,True,True,True) [128+32+16+8+1]
  ,a (True,False,True,True,True,True) [128+32+16+8+4+1]
  ,a (True,False,True,True,True,True,True) [128+32+16+8+4+2+1]
  ,a (True,False,True,True,(True,True,True,True)) [128+32+16+8+4+2+1,1]
  ,encRaw (True,False,True,True) [128+32+16]
  ,encRaw ((True,True,False,True,False),(False,False,True,False,True,True)) [128+64+16+1,64+32]
  ,encRaw ('\0','\1','\127') [0,1,127]
  ,encRaw (33::Word32,44::Word32) [33,44]
    --,s (Elem True) [64]
    --,s (NECons True (NECons False (Elem True))) [128+64+32+4]
  ,encRaw (0::Word8) [0]
  ,encRaw (1::Word8) [1]
  ,encRaw (255::Word8) [255]
  ,encRaw (0::Word16) [0]
  ,encRaw (1::Word16) [1]
  ,encRaw (255::Word16) [255,1]
  ,encRaw (256::Word16) [128,2]
  ,encRaw (65535::Word16) [255,255,3]
  ,encRaw (127::Word32) [127]
  ,encRaw (128::Word32) [128,1]
  ,encRaw (129::Word32) [129,1]
  ,encRaw (255::Word32) [255,1]
  ,encRaw (16383::Word32) [255,127]
  ,encRaw (16384::Word32) [128,128,1]
  ,encRaw (16385::Word32) [129,128,1]
  ,encRaw (32767::Word32) [255,255,1]
  ,encRaw (32768::Word32) [128,128,2]
  ,encRaw (32769::Word32) [129,128,2]
  ,encRaw (65535::Word32) [255,255,3]
  ,encRaw (2097151::Word32) [255,255,127]
  ,encRaw (2097152::Word32) [128,128,128,1]
  ,encRaw (2097153::Word32) [129,128,128,1]
  ,encRaw (4294967295::Word32) [255,255,255,255,15]
  ,encRaw (255::Word64) [255,1]
  ,encRaw (65535::Word64) [255,255,3]
  ,encRaw (4294967295::Word64) [255,255,255,255,15]
  ,encRaw (18446744073709551615::Word64)       [255,255,255,255,255,255,255,255,255,1]
  ,encRaw (False,18446744073709551615::Word64) [127,255,255,255,255,255,255,255,255,128,128]
  ,encRaw (255::Word) [255,1]
  ,encRaw (65535::Word) [255,255,3]
  ,encRaw (4294967295::Word) [255,255,255,255,15]
  ,tstI [0::Int8,2,-2]
  ,encRaw (127::Int8) [254]
  ,encRaw (-128::Int8) [255]
  ,tstI [0::Int16,2,-2,127,-128]
  ,tstI [0::Int32,2,-2,127,-128]
  ,tstI [0::Int64,2,-2,127,-128]
  ,encRaw (-1024::Int64) [255,15]
  ,encRaw (maxBound::Word8)       [255]
  ,encRaw (True,maxBound::Word8)  [255,128]
  ,encRaw (maxBound::Word16)      [255,255,3]
  ,encRaw (True,maxBound::Word16) [255,255,129,128]
  ,encRaw (maxBound::Word32)      [255,255,255,255,15]
  ,encRaw (True,maxBound::Word32) [255,255,255,255,135,128]
  ,encRaw (maxBound::Word64)      [255,255,255,255,255,255,255,255,255,1]
  ,encRaw (True,maxBound::Word64) [255,255,255,255,255,255,255,255,255,128,128]
  ,encRaw (minBound::Int64) [255,255,255,255,255,255,255,255,255,1]
  ,encRaw (maxBound::Int64) [254,255,255,255,255,255,255,255,255,1]
  ,tstI [0::Int,2,-2,127,-128]
  ,tstI [0::Integer,2,-2,127,-128,-256,-512]
  ,encRaw (-1024::Integer) [255,15]
  ,encRaw (0::Float)  [0,0,0,0]
  ,encRaw (-2::Float) [0b11000000,0,0,0]
  ,encRaw (0.085::Float) [0b00111101,0b10101110,0b00010100,0b01111011]
  ,encRaw (0::Double)  [0,0,0,0,0,0,0,0]
  ,encRaw (-2::Double) [0b11000000,0,0,0,0,0,0,0]
  ,encRaw (23::Double) [0b01000000,0b00110111,0,0,0,0,0,0]
  ,encRaw (-0.15625::Float)  [0b10111110,0b00100000,0,0]
  ,encRaw (-0.15625::Double) [0b10111111,0b11000100,0,0,0,0,0,0]
  ,encRaw (-123.2325E-23::Double) [0b10111011,0b10010111,0b01000111,0b00101000,0b01110101,0b01111011,0b01000111,0b10111010]
  ,encRaw (Left True :: Either Bool (Double, Double)) [0b01000000]
  ,encRaw (-2.1234E15 :: Double) [195,30,44,226,90,221,64,0]
  ,encRaw (1.1234E-22 :: Double) [59,96,249,241,120,219,249,174]
  ,encRaw ((False,-2.1234E15) :: (Bool,Double)) [97,143,22,113,45,110,160,0,0]
  ,encRaw ((True,-2.1234E15) :: (Bool,Double)) [225,143,22,113,45,110,160,0,0]
  ,encRaw ((-2.1234E15 , 1.1234E-22) :: (Double, Double)) $ [0b11000011,30,44,226,90,221,64,0] ++ [59,96,249,241,120,219,249,174]
  ,encRaw ((True,-2.1234E15 , 1.1234E-22) :: (Bool,Double, Double)) [0b11100001,143,22,113,45,110,160,0,29,176,124,248,188,109,252,215,0]
  ,encRaw (Right (-2.1234E15 , 1.1234E-22) :: Either Bool (Double, Double)) [0b11100001,143,22,113,45,110,160,0,29,176,124,248,188,109,252,215,0]
  ,encRaw (Left True:: Either Bool Direction) [0b01000000]
  ,encRaw (Right West :: Either Bool Direction) [0b11110000]

  ,map trip [minBound,maxBound::Word8]
  ,map trip [minBound,maxBound::Word16]
  ,map trip [minBound,maxBound::Word32]
  ,map trip [minBound,maxBound::Word64]
  ,map trip [minBound::Int8,maxBound::Int8]
  ,map trip [minBound::Int16,maxBound::Int16]
  ,map trip [minBound::Int32,maxBound::Int32]
  ,map trip [minBound::Int64,maxBound::Int64]
  ,map tripShow [0::Float,-0::Float,0/0::Float,1/0::Float]
  ,map tripShow [0::Double,-0::Double,0/0::Double,1/0::Double]
  ,encRaw '\0' [0]
  ,encRaw '\1' [1]
  ,encRaw '\127' [127]
  ,encRaw 'a' [97]
  ,encRaw 'à' [224,1]
  ,encRaw '经' [207,253,1]
  ,[trip [chr 0x10FFFF]]
  ,encRaw Unit []
  ,encRaw (Un False) [0]
  ,encRaw (One,Two,Three) [16+8]
  ,encRaw (Five,Five,Five) [255,128]
    --,s (NECons True (Elem True)) [128+64+16]
  ,encRaw "" [0]
#ifdef LIST_BIT
  ,encRaw "abc" [176,216,172,96]
  ,encRaw [False,True,False,True] [128
                               +32+16
                               +8
                               +2+1,0]
#elif defined(LIST_BYTE)
  ,s "abc" s3
  ,s (cs 600) s600
#endif
    -- Aligned structures
    --,s (T.pack "") [1,0]
    --,s (Just $ T.pack "abc") [128+1,3,97,98,99,0]
    --,s (T.pack "abc") (al s3)
    --,s (T.pack $ cs 600) (al s600)
  ,encRaw map1 [0b10111000]
  ,encRaw (B.pack $ csb 3) (bsl c3)
  ,encRaw (B.pack $ csb 600) (bsl s600)
  ,encRaw (L.pack $ csb 3) (bsl c3)
   -- Long LazyStrings can have internal sections shorter than 255
   --,s (L.pack $ csb 600) (bsl s600)
  ,[trip [1..100::Int16]]

-- See https://github.com/typelead/eta/issues/901
#ifndef ETA_VERSION
  ,[trip longAsciiStrT]
  ,[trip longBoolListT]
#endif

  ,[trip asciiTextT]
  ,[trip english]
  ,[trip "维护和平正"]
  ,[trip (T.pack "abc")]
  ,[trip unicodeText]
  ,[trip unicodeTextUTF8T]

  ,[trip longBS,trip longLBS]
#ifndef ghcjs_HOST_OS
  ,[trip longSBS]
  ,[trip unicodeTextUTF16T]
#endif
  ]
    where

      --al = (1:) -- prealign
      bsl = id -- noalign

      tstI = map ti

      ti v | v >= 0    = testCase (unwords ["Int",show v]) $ teq v (2 * fromIntegral v ::Word64)
           | otherwise = testCase (unwords ["Int",show v]) $ teq v (2 * fromIntegral (-v) - 1 ::Word64)

      teq a b = ser a @?= ser b

              --,testCase (unwords ["unflat raw",sshow v]) $ desRaw e @?= Right v]

      -- Aligned values unflat to the original value, modulo the added filler.
      a v e = [testCase (unwords ["flat",sshow v]) $ ser v @?= e
              ,testCase (unwords ["unflat",sshow v]) $ let Right v' = des e in v @?= v']
      -- a v e = [testCase (unwords ["flat postAligned",show v]) $ ser (postAligned v) @?= e
      --         ,testCase (unwords ["unflat postAligned",show v]) $ let Right (PostAligned v' _) = des e in v @?= v']

encRaw :: forall a. (Show a, Flat a) => a -> [Word8] -> [TestTree]
encRaw v e = [testCase (unwords ["flat raw",sshow v,show . B.unpack . flat $ v]) $ serRaw v @?= e]

trip :: forall a .(Show a,Flat a,Eq a) => a -> TestTree
trip v = testCase (unwords ["roundtrip",sshow v]) $
  -- direct comparison
  (unflat (flat v::B.ByteString)::Decoded a) @?= (Right v::Decoded a)

tripShow :: forall a .(Show a,Flat a,Eq a) => a -> TestTree
tripShow v = testCase (unwords ["roundtrip",sshow v]) $
  -- we use show to get Right NaN == Right NaN
  show (unflat (flat v::B.ByteString)::Decoded a) @?= show (Right v::Decoded a)

-- Test Data
lzBS = L.pack bs
stBS = B.pack bs
bs = [32,32,32::Word8]
s3 = [3,97,98,99,0]
c3a = [3,99,99,99,0] -- Array Word8
c3 = pre c3a
s600 = pre s600a
pre = (1:)
s600a = concat [[255],csb 255,[255],csb 255,[90],csb 90,[0]]
s600B = concat [[55],csb 55,[255],csb 255,[90],csb 90,[200],csb 200,[0]]
longSeq :: Seq.Seq Word8
longSeq = Seq.fromList lbs
longBS = B.pack lbs
longLBS = L.concat $ concat $ replicate 10 [L.pack lbs]
lbs = concat $ replicate 100 [234,123,255,0]
cs n = replicate n 'c' -- take n $ cycle ['a'..'z']
csb = map (fromIntegral . ord) . cs
map1 = C.fromList [(False,True),(True,False)]

ns :: [(Word64, Int)]
ns =  [( (-) (2 ^(i*7)) 1,fromIntegral (8*i)) | i <- [1 .. 10]]

nsI :: [(Int64, Int)]
nsI = nsI_
nsII :: [(Integer, Int)]
nsII = nsI_
nsI_ =  [( (-) (2 ^(((-) i 1)*7)) 1,fromIntegral (8*i)) | i <- [1 .. 10]]

#ifndef ghcjs_HOST_OS
shBS = SBS.toShort stBS
longSBS = SBS.toShort longBS
#endif

sshow = take 80 . show

showB = show . B.unpack

errDec :: forall a . (Flat a, Eq a, Show a) => Proxy a -> [Word8] -> [TestTree]
--errDec _ bs = [testCase "bad decode" $ let ev = (des bs::Decoded a) in ev @?= Left ""]
errDec _ bs = [testCase "bad decode" $ let ev = (des bs::Decoded a) in isRight ev @?= False]

ser :: Flat a => a -> [Word8]
ser = B.unpack . flat

des :: Flat a => [Word8] -> Decoded a
des = unflat

flatRawWith sz enc = B.unpack $ E.strictEncoder (sz+8) (E.Encoding $ enc >=> E.eFillerF)

serRaw :: Flat a => a -> [Word8]
-- serRaw = B.unpack . flatRaw
-- serRaw = L.unpack . flatRaw
serRaw = asBytes . bits

--desRaw :: Flat a => [Word8] -> Decoded a
--desRaw = unflatRaw . L.pack

type RT a = a -> Bool
type RTL a = Large a -> Bool

prop_Flat_roundtrip :: (Flat a, Eq a) => a -> Bool
prop_Flat_roundtrip = roundTripExt

prop_Flat_Large_roundtrip :: (Eq b, Flat b) => Large b -> Bool
prop_Flat_Large_roundtrip (Large x) = roundTripExt x

roundTrip x = unflat (flat x::B.ByteString) == Right x

-- Test roundtrip for both the value and the value embedded between bools
roundTripExt x = roundTrip x && roundTrip (True,x,False)

prop_double_conv d = wordToDouble (doubleToWord d) == d

prop_float_conv d = wordToFloat (floatToWord d) == d

{-
prop_common_unsigned :: (Num l,Num h,Flat l,Flat h) => l -> h -> Bool
prop_common_unsigned n _ = let n2 :: h = fromIntegral n
                           in flat n == flat n2
-}

-- e :: Stream Bool
-- e = unflatIncremental . flat $ stream1

-- el :: List Bool
-- el = unflatIncremental . flat $ infList

-- deflat = unflat

-- b1 :: BLOB UTF8
-- b1 = BLOB UTF8 (preAligned (List255 [97,98,99]))
-- -- b1 = BLOB (preAligned (UTF8 (List255 [97,98,99])))


{-| This module defines the 'Produce' typeclass, used for generating random
    values for testing in StrictCheck.

    'Produce' is a strict generalization of "Test.QuickCheck"'s 'Arbitrary'
    typeclass. Paired with 'Consume' (a generalization of 'CoArbitrary') it can
    be used to create random non-strict functions, whose strictness behavior is
    dependent on the values given to them.
-}

module Test.StrictCheck.Produce
  ( Produce(..)
  -- * Tools for writing 'Produce' instances
  , recur
  , build
  -- * Producing non-strict functions
  , returning
  , variadic
  -- * Integration with "Test.QuickCheck"'s @Arbitrary@
  , Lazy(..)
  , freely
  -- * Abstract types representing input to a function
  , Input
  , Inputs
  -- * The traversal distribution for processing @Input@s
  , draws
  ) where

import Test.QuickCheck hiding (variant)
import Test.QuickCheck.Gen.Unsafe

import Test.StrictCheck.Internal.Inputs
import Test.StrictCheck.Consume
import Test.StrictCheck.Curry

import Generics.SOP
import Data.Complex
import Data.Monoid ((<>))

-------------------------------------------------------
-- The user interface for creating Produce instances --
-------------------------------------------------------

-- TODO: parameterize over destruction pattern?

-- | Produce an arbitrary value of type @b@, such that destructing that value
-- incrementally evaluates some input to a function.
--
-- Writing instances of @Produce@ is very similar to writing instances of
-- QuickCheck's 'Arbitrary'. The distinction: when making a recursive call to
-- produce a subfield of a structure, __always__ use 'build' or 'recur', and
-- __never__ a direct call to 'produce' itself. This ensures that the input can
-- potentially be demanded at any step of evaluation of the produced value.
--
-- If, in the course of generating a value of type @b@, you need to generate a
-- random value of some other type, which is /not/ going to be a subpart of the
-- resultant @b@ (e.g. a length or depth), use a direct call to @arbitrary@ or
-- some other generator which does not consume input.
--
-- An example instance of @Produce@:
--
-- > data D a
-- >   = X a
-- >   | Y [Int]
-- >
-- > instance Produce a => Produce (D a) where
-- >   produce =
-- >     oneof [ fmap X recur
-- >           , fmap Y recur
-- >           ]
class Produce b where
  produce :: (?inputs::Inputs) => Gen b

theInputs :: (?inputs::Inputs) => [Input]
theInputs = destruct ?inputs

-- | Given an input-consuming producer, wrap it in an outer layer of input
-- consumption, so that this consumption can be interleaved when the producer is
-- called recursively to generate a subfield of a larger produced datatype.
build :: (?inputs::Inputs) => ((?inputs::Inputs) => Gen a) -> Gen a
build gen = do
  (v, is') <- draws theInputs
  vary v $ let ?inputs = Inputs is' in gen

-- | Destruct some inputs to generate an output. This function handles the
-- interleaving of input destruction with output construction. When producing a
-- data type, it should be called to produce each subfield -- *not* produce
-- itself.
recur :: (Produce a, ?inputs::Inputs) => Gen a
recur = build produce

---------------------------------------
-- How to make random lazy functions --
---------------------------------------

-- NOTE: This instance must be defined in this module, as it has to break the
-- abstraction of the Inputs type. No other instance needs to break this.
-- Incidentally, it also must break Gen's abstraction barrier, because it needs
-- to use promote to make a function.

instance (Consume a, Produce b) => Produce (a -> b) where
  produce = returning produce

-- | Create an input-consuming producer of input-consuming functions, given an
-- input-consuming producer for results of that function.
returning
  :: (Consume a, ?inputs::Inputs)
  => ((?inputs::Inputs) => Gen b)
  -> Gen (a -> b)
returning out =
  promote $ \a ->
    let ?inputs = Inputs (consume a : theInputs)
    in build out

-- | Create an input-consuming producer of input-consuming functions, of any
-- arity. This will usually be used in conjuntion with type application, to
-- specify the type(s) of the argument(s) to the function.
variadic ::
  forall args result.
  (All Consume args, Curry args, ?inputs::Inputs)
  => ((?inputs::Inputs) => Gen result)
  -> Gen (args ⋯-> result)
variadic out =
  fmap (curryAll @args @_ @(NP I)) . promote $ \args ->
    let ?inputs =
          Inputs . (++ theInputs) $
            hcollapse $ hcliftA (Proxy @Consume) (K . consume . unI) args
    in build out

-------------------------------------------------------------------------
-- Random destruction of the original input, as transformed into Input --
-------------------------------------------------------------------------

-- | Destruct a random subpart of the given 'Input's, returning the 'Variant'
-- corresponding to the combined information harvested during this process, and
-- the remaining "leaves" of the inputs yet to be destructed
--
-- To maximize the likelihood that different random consumption paths through
-- the same value will diverge (desirable when generating functions with
-- interesting strictness), @draws@ destructs the forest of @Input@s as a
-- depth-first random traversal with a budget sampled from a geometric
-- distribution with expectation 1.
draws :: [Input] -> Gen (Variant, [Input])
draws inputs = go [inputs]
  where
    -- Mutually recursive:
    go, inwardFrom :: [[Input]] -> Gen (Variant, [Input])

    go levels =
      oneof                               -- 50% choice between:
        [ return (mempty, concat levels)  -- stop consuming input, or
        , inwardFrom levels ]             -- keep consuming input

    inwardFrom levels =
      case levels of
        [            ] -> return mempty         -- if no more input: stop
        [  ] : outside -> inwardFrom outside    -- if nothing here: backtrack
        here : outside -> do                    -- if something here: go deeper
          (Input v inside, here') <- pick here
          vary v $ do
            (entropy, levels') <- go (inside : here' : outside)  -- back to 'go'
            return (v <> entropy, levels')

    -- Pick a random list element and the remaining list
    pick :: [a] -> Gen (a, [a])
    pick as = do
      index <- choose (0, length as - 1)
      let (before, picked : after) = splitAt index as
      return (picked, before ++ after)

---------------------------------------------
-- Integration with QuickCheck's Arbitrary --
---------------------------------------------

-- | We hook into QuickCheck's existing Arbitrary infrastructure by using
-- a newtype to differentiate our special way of generating things.
newtype Lazy a
  = Lazy { runLazy :: a }

instance Produce a => Arbitrary (Lazy a) where
  arbitrary = Lazy <$> freely produce

-- | Actually produce an output, given an input-consuming producer. If a
-- function is to be produced, it will be almost-certainly non-strict.
freely :: ((?inputs::Inputs) => Gen a) -> Gen a
freely p = let ?inputs = Inputs [] in p

---------------
-- Instances --
---------------

instance Produce ()       where produce = arbitrary
instance Produce Bool     where produce = arbitrary
instance Produce Ordering where produce = arbitrary

instance Produce Char     where produce = arbitrary
instance Produce Word     where produce = arbitrary
instance Produce Int      where produce = arbitrary
instance Produce Double   where produce = arbitrary
instance Produce Float    where produce = arbitrary
instance Produce Rational where produce = arbitrary
instance Produce Integer  where produce = arbitrary

instance (Arbitrary a, RealFloat a) => Produce (Complex a) where
  produce = arbitrary

instance Produce a => Produce (Maybe a) where
  produce =
    oneof [ return Nothing
          , Just <$> recur
          ]

instance (Produce a, Produce b) => Produce (Either a b) where
  produce =
    oneof [ Left <$> recur
          , Right <$> recur
          ]

instance (Produce a) => Produce [a] where
  produce =
    frequency [ (1, return [])
              , (1, (:) <$> recur
                        <*> recur)
              ]


import Test.Framework (defaultMain)

import qualified Statistics.Iteratee.Tests as T

main :: IO ()
main = defaultMain T.tests


{-# LANGUAGE CPP #-}
-- |
-- This module reexports most of the definitions from the \"base\" package,
-- which are meant to be imported unqualified.
--
-- For details check out the source.
module Lazy.Prelude
(
  module Exports,
  -- * Reimplementations of functions presented in versions of \"base\" newer than 4.6
  -- ** Data.Bool
  bool,
  -- ** Data.Function
  (&),
  -- ** Data.Functor
  ($>),
  -- ** Data.List
  isSubsequenceOf,
  sortOn,
  uncons,
  -- ** Debug.Trace
  traceShowId,
  traceM,
  traceShowM,
)
where

-- Reexports
-------------------------

import Control.Applicative as Exports
import Control.Arrow as Exports hiding (first, second)
import Control.Category as Exports
import Control.Concurrent as Exports
import Control.Comonad as Exports hiding (($>))
import Control.Exception as Exports
import Control.Monad as Exports hiding (mapM_, sequence_, forM_, msum, mapM, sequence, forM)
import Control.Monad.Fix as Exports hiding (fix)
import Control.Monad.ST as Exports
import Data.Bits as Exports
import Data.Bool as Exports hiding (bool)
import Data.Char as Exports
import Data.Complex as Exports
import Data.Data as Exports
import Data.Dynamic as Exports
import Data.Either as Exports
import Data.Fixed as Exports
import Data.Foldable as Exports
import Data.Functor as Exports hiding (($>))
import Data.Function as Exports hiding ((.), id, (&))
import Data.Int as Exports
import Data.IORef as Exports
import Data.Ix as Exports
import Data.List as Exports hiding (sortOn, isSubsequenceOf, uncons, concat, foldr, foldl1, maximum, minimum, product, sum, all, and, any, concatMap, elem, foldl, foldr1, notElem, or, find, maximumBy, minimumBy, mapAccumL, mapAccumR, foldl')
import Data.Maybe as Exports
import Data.Monoid as Exports
import Data.Ord as Exports
import Data.Ratio as Exports
import Data.STRef as Exports
import Data.String as Exports
import Data.Traversable as Exports
import Data.Tuple as Exports
import Data.Unique as Exports
import Data.Version as Exports
import Data.Word as Exports
import Debug.Trace as Exports hiding (traceShowId, traceM, traceShowM)
import Foreign.Storable as Exports
import Foreign.Ptr as Exports
import Foreign.ForeignPtr as Exports
import Foreign.StablePtr as Exports
import GHC.Conc as Exports hiding (withMVar, threadWaitWriteSTM, threadWaitWrite, threadWaitReadSTM, threadWaitRead)
import GHC.Exts as Exports (lazy, inline, sortWith, groupWith)
import GHC.Generics as Exports (Generic)
import GHC.IO.Exception as Exports
import Numeric as Exports
import Prelude as Exports hiding (concat, foldr, mapM_, sequence_, foldl1, maximum, minimum, product, sum, all, and, any, concatMap, elem, foldl, foldr1, notElem, or, mapM, sequence, id, (.))
import System.Environment as Exports
import System.Exit as Exports
import System.IO as Exports (Handle, hClose)
import System.IO.Error as Exports
import System.IO.Unsafe as Exports
import System.Mem as Exports
import System.Mem.StableName as Exports
import System.Timeout as Exports
import Text.ParserCombinators.ReadP as Exports (ReadP, ReadS, readP_to_S, readS_to_P)
import Text.ParserCombinators.ReadPrec as Exports (ReadPrec, readPrec_to_P, readP_to_Prec, readPrec_to_S, readS_to_Prec)
import Text.Printf as Exports (printf, hPrintf)
import Text.Read as Exports (Read(..), readMaybe, readEither)
import Unsafe.Coerce as Exports

-- Conditional imports for newer bases
#if MIN_VERSION_base(4,7,0)
import Data.Coerce as Exports
import Data.Proxy as Exports
#endif
#if MIN_VERSION_base(4,8,0)
import Data.Bifunctor as Exports
import Data.Functor.Identity as Exports
import Data.Void as Exports
import Numeric.Natural as Exports
#endif
#if MIN_VERSION_base(4,9,0)
import Control.Monad.Fail as Exports (MonadFail)
import Control.Monad.IO.Class as Exports
import Data.Functor.Classes as Exports
import Data.Functor.Compose as Exports
import Data.List.NonEmpty as Exports (NonEmpty(..))
import Data.Semigroup as Exports hiding ((<>), First(..), Last(..))
#endif

-- Conditional imports for reimplementations
#if MIN_VERSION_base(4,7,0)
import Data.Bool (bool)
import Debug.Trace (traceShowId, traceM, traceShowM)
import Data.Functor (($>))
#endif
#if MIN_VERSION_base(4,8,0)
import Data.Function ((&))
import Data.List (isSubsequenceOf, sortOn, uncons)
#endif

---------------------------------
-- Reimplementations for base-4.7
---------------------------------

#if !MIN_VERSION_base(4,7,0)

-- | Case analysis for the 'Bool' type.
-- @bool a b p@ evaluates to @a@ when @p@ is @False@, and evaluates to @b@
-- when @p@ is @True@.
bool :: a -> a -> Bool -> a
bool f t b = if b then t else f

{-|
Like 'traceShow' but returns the shown value instead of a third value.
-}
traceShowId :: (Show a) => a -> a
traceShowId a = trace (show a) a

{-|
Like 'trace' but returning unit in an arbitrary monad. Allows for convenient
use in do-notation. Note that the application of 'trace' is not an action in the
monad, as 'traceIO' is in the 'IO' monad.

> ... = do
>   x <- ...
>   traceM $ "x: " ++ show x
>   y <- ...
>   traceM $ "y: " ++ show y
-}
traceM :: (Monad m) => String -> m ()
traceM string = trace string $ return ()

{-|
Like 'traceM', but uses 'show' on the argument to convert it to a 'String'.

> ... = do
>   x <- ...
>   traceShowM $ x
>   y <- ...
>   traceShowM $ x + y
-}
traceShowM :: (Show a, Monad m) => a -> m ()
traceShowM = traceM . show

infixl 4 $>

-- | Flipped version of '<$'.
($>) :: Functor f => f a -> b -> f b
($>) = flip (<$)

#endif

---------------------------------
-- Reimplementations for base-4.8
---------------------------------

#if !MIN_VERSION_base(4,8,0)

infixl 1 &

-- | '&' is a reverse application operator.  This provides notational
-- convenience.  Its precedence is one higher than that of the forward
-- application operator '$', which allows '&' to be nested in '$'.
(&) :: a -> (a -> b) -> b
x & f = f x

-- | The 'isSubsequenceOf' function takes two lists and returns 'True' if the
-- first list is a subsequence of the second list.
--
-- @'isSubsequenceOf' x y@ is equivalent to @'elem' x ('subsequences' y)@.
--
-- ==== __Examples__
--
-- >>> isSubsequenceOf "GHC" "The Glorious Haskell Compiler"
-- True
-- >>> isSubsequenceOf ['a','d'..'z'] ['a'..'z']
-- True
-- >>> isSubsequenceOf [1..10] [10,9..0]
-- False
isSubsequenceOf :: (Eq a) => [a] -> [a] -> Bool
isSubsequenceOf []    _                    = True
isSubsequenceOf _     []                   = False
isSubsequenceOf a@(x:a') (y:b) | x == y    = isSubsequenceOf a' b
                               | otherwise = isSubsequenceOf a b

-- | Decompose a list into its head and tail. If the list is empty,
-- returns 'Nothing'. If the list is non-empty, returns @'Just' (x, xs)@,
-- where @x@ is the head of the list and @xs@ its tail.
uncons        :: [a] -> Maybe (a, [a])
uncons []     = Nothing
uncons (x:xs) = Just (x, xs)

-- | Sort a list by comparing the results of a key function applied to each
-- element.  @sortOn f@ is equivalent to @sortBy . comparing f@, but has the
-- performance advantage of only evaluating @f@ once for each element in the
-- input list.  This is called the decorate-sort-undecorate paradigm, or
-- Schwartzian transform.
sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f =
  map snd . sortBy (comparing fst) . map (\x -> let y = f x in y `seq` (y, x))

#endif