{-# language LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ConstraintKinds #-}
{-# language DerivingVia #-}
{-# language StandaloneDeriving #-}
{-# language ApplicativeDo #-}

-- | This module defines representation of all figures as SVG elements. as well as default style settings.
module Geometry.SVG
  ( Groupable, Group (..)
  , (<+>), group, row, rowSep
  , (<||>), beside, above
  , SVGable (..), SVGContext(..)
  , showSVG
  ) where

import Prelude hiding (writeFile, unwords)
import Graphics.Svg.Core (Attribute (..))
import Graphics.Svg (doctype, svg11_, with, prettyText, (<<-))
import Graphics.Svg (Element, toElement)
import Graphics.Svg.Elements
import Graphics.Svg.Attributes
import Data.Complex
import Data.Monoid
import Data.Maybe
import Data.Text (Text, pack)
import qualified Data.Text.Lazy as LT
import Data.Double.Conversion.Text (toShortest)

import Geometry.Base
import Geometry.Decorations
import Geometry.Point
import Geometry.Circle
import Geometry.Polygon
import Geometry.Plot
import Geometry.Line
import Geometry.Angle
import Geometry.Figures

showt :: Show a => a -> Text
showt = pack . show

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

find :: WithOptions a => (Option -> Maybe b) -> a -> Maybe b
find p d = extractOption p $ defaultOptions d <> options d

extractOption p = getFirst . foldMap (First . p) . getOptions

optFill          = \case {Fill x -> Just x; _ -> Nothing }
optStroke        = \case {Stroke x -> Just x; _ -> Nothing }
optInvisible     = \case {Invisible x -> Just x; _ -> Nothing }
optThickness     = \case {Thickness x -> Just x; _ -> Nothing }
optMultiStroke   = \case {MultiStroke x -> Just x; _ -> Nothing }
optDashing       = \case {Dashing x -> Just x; _ -> Nothing }
optLabelText     = \case {LabelText x -> Just x; _ -> Nothing }
optLabelPosition = \case {LabelPosition x -> Just x; _ -> Nothing }
optLabelOffset   = \case {LabelOffset x -> Just x; _ -> Nothing }
optLabelCorner   = \case {LabelCorner x -> Just x; _ -> Nothing }
optLabelAngle    = \case {LabelAngle x -> Just x; _ -> Nothing }
optSegmentMark   = \case {SegmentMark x -> Just x; _ -> Nothing }

------------------------------------------------------------
{- | Class for objects that could be represented as SVG elements.
Differs from @Graphics.Svg.ToElement@ class in adding
parameter context to a main render function `toSVG`.
-}
class SVGable a where
  toSVG :: a -> SVGContext -> Element
  toSVG _ _ = mempty

  fmtSVG :: a -> Text
  fmtSVG = mempty

instance SVGable Double where
  fmtSVG n = if n ~= 0 then "0" else toShortest n

instance SVGable Cmp where
  fmtSVG = fmtSVG . xy

instance SVGable XY where
  fmtSVG (x, y) = fmtSVG x <> "," <> fmtSVG y <> " "

instance {-# OVERLAPPABLE #-}
  (Foldable t, SVGable a) => SVGable (t a) where
  toSVG = foldMap toSVG

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

instance SVGable a => SVGable (Decorated a) where
  toSVG d = case find optInvisible d of
    Just True -> mempty
    _ -> toSVG (fromDecorated d) . updateOptions (options d)

attributes :: (WithOptions f, Figure f) => f -> SVGContext -> [Attribute]
attributes f ctx = mconcat fmt . getOptions $ opts
  where
    opts = options f <> figureOptions ctx
    fmt = [ attr optStroke Stroke_
          , attr optFill Fill_
          , attr optThickness Stroke_width_
          , attr optDashing Stroke_dasharray_ ]

    attr opt a x = maybeToList $ (\s -> a <<- pack s) <$> findOp opt x
    findOp p = getFirst . foldMap (First . p)

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

instance SVGable Point where
  toSVG p = circle_ . attr <> labelElement p
    where
      attr =
        attributes p <>
        const [ Cx_ <<- fmtSVG (getX p)
              , Cy_ <<- fmtSVG (getY p)
              , R_ <<- "3" ]

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

instance SVGable Label where
  toSVG l = labelElement l . updateOptions (options l)

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

instance SVGable Circle where
  toSVG c = circle_ . attr <> labelElement c
    where
      (x :+ y) = center c
      attr = attributes c <>
             const [ Cx_ <<- fmtSVG x
                   , Cy_ <<- fmtSVG y
                   , R_ <<- fmtSVG (radius c) ]

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

instance SVGable Line where
  toSVG l = do
    bx <- figureBox
    let relabel s = s #: lpos ((s @-> 0.95) - cmp s)
    foldMap (toSVG . relabel) $ clipBy bx l

instance SVGable Ray where
  toSVG r = do
    bx <- figureBox
    let relabel s = s #: lpos ((s @-> 0.95) - cmp s)
    foldMap (toSVG . relabel) $ clipBy bx r

instance SVGable Segment where
  toSVG s = element . attr <> labelElement s
    where
      (a, b) = refPoints s
      element = if isTrivial s then mempty else line_
      attr = attributes s <>
             const [ X1_ <<- fmtSVG (getX a)
                   , Y1_ <<- fmtSVG (getY a)
                   , X2_ <<- fmtSVG (getX b)
                   , Y2_ <<- fmtSVG (getY b) ]

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

instance SVGable Polyline where
  toSVG p | isTrivial p = labelElement p
          | otherwise = polyline_ . attr <> labelElement p
    where
      attr = attributes p <>
             const [ Points_ <<- foldMap fmtSVG (vertices p) ]

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

instance SVGable Polygon where
  toSVG p | isTrivial p = labelElement p
          | otherwise = polygon_ . attr <> labelElement p
    where
      attr = attributes p <>
             const [ Points_ <<- foldMap fmtSVG (vertices p) ]

deriving via Polygon instance SVGable Triangle
deriving via Triangle instance SVGable RightTriangle
deriving via Polygon instance SVGable Rectangle

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

instance (AlmostEq a, Pnt a) => SVGable (Plot a) where
  toSVG = toSVG . asPolyline

instance (AlmostEq a, Pnt a) => SVGable (ClosedPlot a) where
  toSVG = toSVG . asPolyline

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

instance SVGable Angle where
  toSVG an' ctx = toSVG (rays <+> group arc) ctx' <>
                 labelElement an ctx'
    where
      an = reflex an'
      t = extractOption optLabelText (figureOptions ctx)
      labelText = case t of
        Just "#" -> show (angleValue an)
        Just v -> v
        Nothing -> ""
      ctx' = updateOptions (options an <> mkOptions [LabelText labelText]) ctx
      Just ns = extractOption optMultiStroke (figureOptions ctx')
      Angle p s e = an
      rays = mkPolyline [p + cmp s, p, p + cmp e] # scaleAt' p 20
      arc = [ plotManifold an # at' (p + cmp s) # scaleAt' p r
            | i <- [1..ns]
            , let r = 12 + fromIntegral i * 4 ]

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

labelElement :: (WithOptions f, Figure f) => f -> SVGContext -> Element
labelElement ff ctx = case lb of
                   Just s -> text $ toElement s
                   Nothing -> mempty
  where
    opts = figureOptions ctx
    f = Decorated (options ff <> opts, ff)

    lb = find optLabelText f
    loff = fromMaybe 0 $ find optLabelOffset f
    lpos = fromMaybe 0 $ find optLabelPosition f
    (cx, cy) = fromMaybe (0,0) $ find optLabelCorner f

    fontSize = 14
    textWidth = fromIntegral $ length lb
    text = text_ $ [ X_ <<- fmtSVG x
                   , Y_ <<- fmtSVG y
                   , Font_size_ <<- "16"
                   , Font_family_ <<- "CMU Serif"
                   , Font_style_ <<- "italic"
                   , Stroke_ <<- "none"
                   , Fill_ <<- "white"] <> offsetX <> offsetY

    x :+ y = lpos + d
    d = scale (fromIntegral fontSize) loff
    offsetX = case signum cx of
                -1 -> [ Text_anchor_ <<- "start" ]
                0 -> [ Text_anchor_ <<- "middle" ]
                1 -> [ Text_anchor_ <<- "end" ]
    offsetY = case signum cy of
                1 -> [ Dy_ <<- showt (-fontSize `div` 4 -1) ]
                0 -> [ Dy_ <<- showt (fontSize `div` 4 +1) ]
                -1 -> [ Dy_ <<- showt (fontSize - 2) ]

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

-- | Constrain for an object that could be included to a group.
type Groupable a = (SVGable a, Show a, Trans a, Figure a, Eq a)

-- | The group of inhomogeneous `Groupable` objects.
--
-- > let f = G . translate (1,0) . scale 0.7 . rotate 30 <>
-- >         G . translate (1,0) . scale 0.6 . rotate (-45)
-- > in G aSegment # iterate f # take 8 # rotate 90
-- << figs/compose.svg >>
--
data Group where
  EmptyFig :: Group
  G :: Groupable a => a -> Group
  Append :: Group -> Group -> Group

instance Eq Group where
  _ == _ = False

instance Semigroup Group where (<>) = Append
instance Monoid Group where mempty = EmptyFig

infixl 3 <+>
-- | The appending operator for groupable objects.
(<+>) :: (Groupable a, Groupable b) => a -> b -> Group
a <+> b = G a <> G b

instance Trans Group where
  transform _ EmptyFig = EmptyFig
  transform t (G f) = G $ transform t f
  transform t (Append x xs) = Append (transform t x) (transform t xs)

instance Show Group where
  show EmptyFig = mempty
  show (G a) = show a
  show (Append x xs) = show x <> show xs

instance SVGable Group where
  toSVG EmptyFig = mempty
  toSVG (G a) = toSVG a
  toSVG (Append a b) = toSVG a <> toSVG b

instance Figure Group where
  refPoint = left . lower . corner
  isTrivial _ = False
  box EmptyFig = mempty
  box (G f) = box f
  box (Append a b) = box a <> box b

-- | Returns a group of homogeneous list of objects.
--
-- > group $ take 5 $ iterate (rotate 5 . scale 1.1) aTriangle
-- <<figs/groups.svg>>
--
group :: Groupable a => [a] -> Group
group = foldMap G

-- | Places one figure beside the other.
--
-- > aTriangle `beside` aSquare
-- <<figs/beside.svg>>
--
beside :: (Groupable a, Groupable b) => a -> b -> Group
beside f1 f2 = f1 <+> f2 # superpose (refPoint f2) (right . lower . corner $ f1)

-- | The operator form of the `beside` function.
infixl 2 <||>
(<||>) :: (Groupable a, Groupable b) => a -> b -> Group
(<||>) = beside

-- | Places one figure above the other.
--
-- > aTriangle `above` aSquare
-- <<figs/above.svg>>
--
-- > let tr t = t `above` (t `beside` t)
-- > in G aCircle # iterate tr # take 5 # rotate 225 # scaleX 0.6
-- <<figs/serp.svg>>
--
above :: (Groupable a, Groupable b) => a -> b -> Group
above f1 f2 = f2 <+> f1 # superpose (refPoint f2) (left . upper . corner $ f2)

{- | Arranges a homogeneous list of figures in a row

> row [ regularPoly n | n <- [3..7] ]
<<figs/row.svg>>
-}
row :: Groupable f => [f] -> Group
row [] = EmptyFig
row fs = foldl1 beside (G <$> fs)

{- | Arranges a homogeneous list of figures in a row, separated by a given separator.

> rowSep (space 1) [ regularPoly n | n <- [3..7] ]
<<figs/rowSep.svg>>
-}
rowSep :: (Groupable s, Groupable f) => s -> [f] -> Group
rowSep s [] = EmptyFig
rowSep s fs = foldl1 (\g f -> g `beside` s `beside` f) (G <$> fs)
------------------------------------------------------------

-- | The record containing image parameters and options.
data SVGContext = SVGContext
  { imageSize :: Int
  , figureBox :: Rectangle
  , border :: Double
  , figureOptions :: Options }

updateOptions :: Options -> SVGContext -> SVGContext
updateOptions opts ctx = ctx {figureOptions = figureOptions ctx <> opts}

wrapSVG :: Element -> SVGContext -> Element
wrapSVG content ctx =
  doctype <>
  with (svg11_ content)
  [ Version_ <<- "1.1"
  , Width_ <<- showt (border ctx + figureWidth (figureBox ctx))
  , Height_ <<- showt (1.5*border ctx + figureHeight (figureBox ctx))
  , Style_ <<- "background : #444;" ]

-- | Creates a SVG contents for geometric objects.
showSVG :: (Figure a, SVGable a) => Int -> a -> LT.Text
showSVG size obj = prettyText (contents ctx)
  where
    brd = 30
    contents = toSVG obj' >>= wrapSVG
    ctx = SVGContext size fb brd mempty
    obj' = obj
           # superpose p0 (0 :: Cmp)
           # reflect 0
           # scale ((fromIntegral size - brd*3) / ((w `max` h) `min` paperSize))
           # translate' ((brd*1.5, brd*1.5) :: XY)
    fb = boxRectangle obj
         # superpose p0 (0 :: Cmp)
         # reflect 0
         # scale ((fromIntegral size - brd) / ((w `max` h) `min` paperSize))
         # translate' ((brd/2, brd/2) :: XY)
    paperSize = 50
    w = figureWidth obj
    h = figureHeight obj
    p0 = left . upper . corner $ obj

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


{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE GADTs #-}
module Lib where

import Numeric
import Prelude hiding (and,or,(>))
import Control.Monad.RWS
import System.Process
import System.Exit
import Data.List (intercalate)
import qualified Data.Map as M
import qualified Data.Vector.Unboxed as V
import Statistics.Sample.KernelDensity
import Text.Read (readMaybe)

--------------------------------------------
-- Basic code-gen infrastructure

type a ⊸ b = a -> b
data Expr t where
  Expr :: String -> Expr t -- constant expression
  Lam :: (Expr a -> Expr t) -> Expr (a -> t) -- lambda
  App :: Expr (a -> t) -> Expr a -> Expr t -- application
  BinOp :: (String -> String -> String) -> Expr a -> Expr b -> Expr c
  UnOp :: (String -> String) -> Expr a -> Expr b
  If :: Prop -> Expr a -> Expr a -> Expr a
  -- Let :: P (Expr a) -> (Expr a ⊸ Expr b) -> Expr b
  Thunk :: P (Expr t) -> Expr t

class Representable a where
  constant :: a -> Expr a

instance Representable Bool where
  constant True = Expr "true"
  constant False = Expr "false"

instance Representable Float where
  constant = Expr . show

instance Fractional (Expr Float) where
  fromRational x = constant (fromRat x)
  (/) = BinOp (infixOp "/")

instance Num (Expr Float) where
  abs = UnOp (\x -> "Math.abs(" ++ x ++ ")")
  (-) = BinOp (infixOp "-")
  (+) = BinOp (infixOp "+")
  fromInteger x = constant (fromInteger x)
  (*) = BinOp (infixOp "*")

instance Floating (Expr Float) where
  sqrt = UnOp (\x -> "Math.sqrt(" ++ x ++ ")")
  exp = UnOp (\x -> "Math.exp(" ++ x ++ ")")
  log = UnOp (\x -> "Math.log(" ++ x ++ ")")

render :: Expr a -> P String
render = \case
  UnOp op x -> do
    x' <- render x
    return (op x')
  BinOp op x y -> do
    x' <- render x
    y' <- render y
    return (op x' y')
  Expr x -> return x
  App (Lam f) x -> render (f x)
  Lam f -> do
    x <- newVar
    body <- render (f (Expr x))
    return ("function( " ++ x ++ ") { return " ++ body ++ "; }")
  Thunk f -> do
    (x,body) <- censor (\_ -> []) $ listen f
    x' <- render x
    return ("function() {\n" ++ unlines body ++ " return " ++ x' ++ ";\n }")
  App x y -> do
    x' <- render x
    y' <- render y
    return ("(" ++ x' ++ ")(" ++ y' ++ ")")
  If cond x y -> do
    c' <- render cond
    x' <- render x
    y' <- render y
    return ("(" ++ c'  ++ "?" ++ x' ++ ":" ++ y' ++ ")")
  -- Let e f -> do
  --   x <- newVar
  --   e' <- render =<< e -- works only if f uses its argument at most once.
  --   fx <- render (f (Expr x))
  --   return ("var " ++ x ++ " = "++ e' ++";\n" ++ fx)

infixl #
(#) :: Expr (a -> b) -> Expr a -> Expr b
(#) = App

parens :: [Char] -> [Char]
parens x = "("++ x ++")"
infixOp :: [Char] -> [Char] -> [Char] -> [Char]
infixOp op x y = parens x ++ op ++ parens y
binFunc :: [Char] -> [Char] -> [Char] -> [Char]
binFunc f x y = f ++ "("++ x ++ "," ++ y ++")"

unFunc ::  [Char] -> [Char] -> [Char]
unFunc f x = f ++ "("++ x ++ ")"

newtype P a = P (RWS () [String] State a) deriving (Functor,Applicative,Monad,MonadState State, MonadWriter [String])

data State = State {nextVar :: Int}

logtrace :: String -> Expr a -> P ()
logtrace msg x = do
  x' <- render x
  emit ("console.log(" ++ show "LOG: " ++ "+" ++ show msg ++ "+ \" \" +" ++ x' ++ ");")

compileModel :: String -> P (Expr a) -> String
compileModel mainFunc m = unlines (z ++ [x])
  where (x,_,z) = runRWS p () (State {nextVar = 0})
        (P p) = render (UnOp (unFunc mainFunc) (Thunk m))

and :: Prop -> Prop -> Prop
and = BinOp (infixOp "&&")

(∧) :: Prop -> Prop -> Prop
(∧) = and

iff :: Prop -> Prop -> Prop
iff = BinOp (infixOp "==")

or :: Prop -> Prop -> Prop
or = BinOp (infixOp "||")

not' :: Prop -> Prop
not' = UnOp (\x -> "(!(" ++ x ++ "))")

(-->) :: Prop -> Prop -> Prop
p --> q = not' p `or` q

-- | Allocate a new variable
newVar :: P String
newVar = do
  n <- gets nextVar
  modify $ \State{..} -> State {nextVar = n+1, ..}
  return ("v" ++ show n)

emit :: String -> P ()
emit x = tell [x]

-----------------------------------
-- Types

type Vec = [Expr Float]
data Distrib a

type Ind = Vec
type Mat = [Vec]
-- type Vec = Expr Vector
type Prop = Expr Bool
type Pred = Ind -> Prop
type Measure = Ind -> Expr Float
type Adj = Vec
type AP = Measure
type CN = Ind -> Prop
type VP = Ind -> Prop
type NP = VP -> Prop
type Quant = CN -> NP

----------------------------------------------------
-- Compositional semantics

observe :: Prop -> P ()
observe = hyp

squared :: Num a => a -> a
squared x = x*x

observeEqual :: Expr Float -> Expr Float -> P ()
observeEqual x y = do
  f <- render (negate (squared (x-y)))
  emit ("factor(" ++ f ++ ");")

hyp :: Prop -> P ()
hyp x = do
  x' <- render x
  emit ("hyp(" ++ x' ++ ");")

-- | Sample new individual
newInd :: P Ind
newInd = newIndSuch []

-- | Sample new individual which satisfies some additional predicates
newIndSuch :: [Pred] -> P Ind
newIndSuch hs = do
  x <- newVector
  forM_ hs $ \h -> hyp (h x)
  return x

numberOfDimensions :: Int
numberOfDimensions = 2

newVector :: P Vec
newVector = mapM (uncurry sampleGaussian) (replicate numberOfDimensions (0,1))

newNormedVector :: P Vec
newNormedVector = do
  xs <- mapM (uncurry sampleGaussian) (replicate numberOfDimensions (0,1))
  return ((/ norm xs) <$> xs)

cosineDistance :: Vec -> Vec -> Expr Float
cosineDistance x y = dotProd x y / (norm x * norm y)

norm2 :: Vec -> Expr Float
norm2 x = dotProd x x

norm :: Vec -> Expr Float
norm = sqrt . norm2

newMatrix = mapM (\_ -> newVector) (replicate numberOfDimensions ())
newMatrix :: P Mat

satisfyAll :: [Ind] -> [Pred] -> P ()
satisfyAll xs ps = forM_ xs $ \x -> forM ps $ \p -> hyp (p x)

dotProd :: Vec -> Vec -> Expr Float
dotProd x y = sum (zipWith (*) x y)

vecMatProd :: Vec -> Mat -> Vec
vecMatProd v = map (dotProd v)

type Scalar = (Vec,Expr Float)

newScalarA :: P Scalar
newScalarA = do
  bias <- sampleGaussian 0 1
  v <- newNormedVector
  return (v,bias)

sampleGaussian :: Expr Float -> Expr Float -> P (Expr Float)
sampleGaussian mu sigma = do
  v <- newVar
  let m = Expr v
  mu' <- render mu
  sigma' <- render sigma
  emit ("var " ++ v ++ " = gaussian("  ++ mu' ++ "," ++ sigma' ++ ");")
  return m

newClass :: P Mat
newClass = newMatrix

forClass :: Mat -> Adj -> AP
forClass cl adj x = dotProd (vecMatProd adj cl) x

-- -- Alternative for scalar adjectives. We take the measure to be greater than that of a "random" element of the class.
-- forClassScalar :: Mat -> Adj -> AP
-- forClassScalar cl a x = Let (newIndSuch [isClass cl]) (\y -> adjAP a x - adjAP a y)

isClass :: Mat -> Pred
isClass clas x = dotProd (head clas) x > 0

adjAP :: Adj -> AP
adjAP = dotProd

gaussian :: Expr Float -> Expr Float -> Expr Float
gaussian mean stdDev = BinOp (binFunc "gaussian") mean stdDev

vague :: Float -> Measure -> Measure
vague vagueness m x = m x + gaussian 0 (Expr (show vagueness))

newMeasure :: P Measure
newMeasure = do
  (v,bias) <- newScalarA
  return (\x -> bias + dotProd v x)

positive :: Expr Float -> Prop
positive x = greaterThan x 0

newPred :: P (Ind -> Prop)
newPred = do
  m <- newMeasure
  return (\x -> positive (m x))

sample :: Expr (Distrib Bool) -> Expr Bool
sample d = Expr "sample" # d

bernouilli :: Expr Float -> Expr (Distrib Bool)
bernouilli = UnOp (\x -> "Bernoulli({p:" ++ x ++ "})")

-- TODO: there is a choice here.

-- this is an "intuitionistic" probably (it does not exclude the strong implication)
probablyInt :: Expr Float -> Prop -> Prop
probablyInt v x = sample (bernouilli v) --> x

-- this is a "definite" probably (it excludes the strong implication)
probablyDef ::  Expr Float -> Prop -> Prop
probablyDef v x = If (sample (bernouilli v)) x (not' x)

--  "a --> b"
-- and "if a then b else (not b)"
-- are not the same!

expectedPortion :: P Prop -> Expr Float
expectedPortion = UnOp (unFunc "expectedPortion") . Thunk

genQProb :: Expr Float -> Quant
genQProb prob cn vp = expectedPortion p > prob
  where p = do x <- newInd
               observe (cn x)
               return (vp x)

many :: Quant
many = genQProb 0.6

most :: Quant
most = genQProb 0.7

few :: Quant
few cn vp = genQProb 0.8 cn (not' . vp)

some :: Quant
some = genQProb 0.1

every :: Quant
every = genQProb 0.99

is :: Measure -> Pred
is m x =  (m x) > 0

more :: Measure -> Ind -> Ind -> Prop
more m x y = (m x) > (m y)

greaterThan :: Expr Float -> Expr Float -> Prop
greaterThan = BinOp (infixOp ">")

(>) :: Expr Float -> Expr Float -> Prop
(>) = greaterThan

-- equal :: Expr Float -> Expr Float -> Prop
-- equal x y = 0.1 > (abs (x - y))

equal :: Expr Float -> Expr Float -> Prop
equal x y = 0.1 > (abs (x - y))

disjoint :: Pred -> Pred -> Pred
disjoint p q x = not' (p x) `or` not' (q x)

subsective :: Adj -> Mat -> Pred
subsective a cl x = isClass cl x `and` is (forClass cl a) x

-- An alternative semantics for subsective scalars.

anything :: Pred
anything _ = Expr "true"

plot :: String -> [Double] -> IO ()
plot prefix xs = do
  let (xs',ys') = kde 64 (V.fromList xs)
      fname = prefix ++ ".svg"
  let ls =
        ["set terminal svg size 350,262 fname 'Verdana' enhanced background rgb 'white'",
         "set output '" ++ fname ++ "'",
         "set key off", -- no legend
          "$data << EOD"] ++
        [show x ++ " " ++ show y | (x,y) <- zip (V.toList xs') (V.toList ys')] ++
        ["EOD", "plot '$data' with lines"] -- filledcurve
  putStrLn "Plotting results..."
  (code,output,errors) <- readProcessWithExitCode "gnuplot" ["-p"] (unlines ls)
  case code of
    ExitFailure _ -> do
      putStrLn "Gnuplot failed with input:"
      putStrLn (unlines ls)
      putStrLn "errors:"
      putStrLn output
      putStrLn errors
    ExitSuccess ->
      putStrLn ("Plot output to " ++ fname)
  return ()

class KnownTyp a where
  isContinuous :: Bool

instance KnownTyp Float where
  isContinuous = True

instance KnownTyp Bool where
  isContinuous = False

parseValues :: [String] -> IO [Double]
parseValues [] = return []
parseValues (v:vs) = case readMaybe v of
  Nothing -> putStrLn v >> parseValues vs
  Just x -> (x:) <$> parseValues vs

run :: forall a. KnownTyp a => P (Expr a) -> IO ()
run model = do
  putStrLn "Creating model..."
  rts <- readFile "../Frontend/RTS.wppl"
  let m = compileModel mainName model
      funname = "modelFunction"
      mainName = if isContinuous @a then "mainContinuous" else "mainDiscrete"
      fname = funname ++ ".wppl"
  writeFile fname (intercalate "\n\n" [rts,m])
  putStrLn "Running webppl..."
  (code,output,errors) <- readProcessWithExitCode "webppl" [fname] ""
  case code of
    ExitFailure _ -> do
      putStrLn "Webppl failed with errors:"
      putStrLn output
      putStrLn errors
    ExitSuccess -> do
      putStrLn "Success!"
      case isContinuous @a of
        True -> do
            values <- parseValues (lines output)
            plot funname values
        False -> putStrLn output


{-# LANGUAGE ConstraintKinds, DataKinds, FlexibleInstances,
             GeneralizedNewtypeDeriving, MultiParamTypeClasses,
             TypeFamilies #-}
-- | Several implementations of the monad in which our dual numbers live
-- and implementations of calculating gradient, derivative and value
-- of an objective function defined on dual numbers.
module HordeAd.Core.Engine
  ( DualMonadValue, primalValueGeneral, primalValue
  , DualMonadGradient, dReverseGeneral, dReverse, dForwardGeneral, dForward, prettyPrintDf
  , DualMonadForward, dFastForwardGeneral, dFastForward
  , generateDeltaVars, initializerFixed
  ) where

import Prelude

import           Control.Monad.Trans.State.Strict
import qualified Data.Array.DynamicS as OT
import           Data.Functor.Identity
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           Numeric.LinearAlgebra (Matrix, Numeric, Vector)
import qualified Numeric.LinearAlgebra as HM

import HordeAd.Core.DualClass
  (Dual, IsPrimal (..), IsPrimalWithScalar, bindInState, dVar)
import HordeAd.Core.DualNumber
import HordeAd.Core.PairOfVectors (DualNumberVariables, makeDualNumberVariables)
import HordeAd.Internal.Delta
  ( CodeOut (..)
  , DeltaState (..)
  , derivativeFromDelta
  , gradientFromDelta
  , ppBindings
  , toDeltaId
  )

-- * The dummy monad implementation that does not collect deltas.
-- It's intended for efficiently calculating the value of the function only.

-- An overcomplicated
-- @type DualMonadValue r = Identity@
-- to avoid @-Wno-orphans@ and @UndecidableInstances@.
newtype DualMonadValue r a = DualMonadValue
  { runDualMonadValue :: Identity a }
  deriving (Monad, Functor, Applicative)

instance IsScalar 'DModeGradient r
         => DualMonad 'DModeGradient r (DualMonadValue r) where
  returnLet (D u _u') = DualMonadValue $ Identity $ D u dZero

-- The general case, needed for old, hacky tests using only scalars.
primalValueGeneral
  :: forall r a. IsScalar 'DModeGradient r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadValue r a)
  -> Domains r
  -> a
-- Small enough that inline won't hurt.
{-# INLINE primalValueGeneral #-}
primalValueGeneral f (params0, params1, params2, paramsX) =
  let replicateZeros p = V.replicate (V.length p) dZero
      variables = makeDualNumberVariables
                    (params0, params1, params2, paramsX)
                    ( replicateZeros params0  -- dummy
                    , replicateZeros params1
                    , replicateZeros params2
                    , replicateZeros paramsX )
  in runIdentity $ runDualMonadValue $ f variables

primalValue
  :: forall r a. IsScalar 'DModeGradient r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadValue r (DualNumber 'DModeGradient a))
  -> Domains r
  -> a
-- Small enough that inline won't hurt.
{-# INLINE primalValue #-}
primalValue f parameters =
  let D value _ = primalValueGeneral f parameters
  in value

-- * The fully-fledged monad implementation for gradients
-- and the code that uses it to compute gradients and also derivatives.

newtype DualMonadGradient r a = DualMonadGradient
  { runDualMonadGradient :: State (DeltaState r) a }
  deriving (Monad, Functor, Applicative)

instance IsScalar 'DModeGradient r
         => DualMonad 'DModeGradient
                      r (DualMonadGradient r) where
  returnLet (D u u') = DualMonadGradient $ do
    st <- get
    let (!stNew, !dId) = bindInState u' st
    put stNew
    return $! D u (dVar dId)

initializeState :: forall r. IsScalar 'DModeGradient r
                => Domains r -> DeltaState r
initializeState (params0, params1, params2, paramsX) =
  DeltaState { deltaCounter0 = toDeltaId (V.length params0)
             , deltaCounter1 = toDeltaId (V.length params1)
             , deltaCounter2 = toDeltaId (V.length params2)
             , deltaCounterX = toDeltaId (V.length paramsX)
             , deltaBindings = []
             }

dReverseGeneral
  :: forall r. HasDelta r
  => r
  -> DualNumberVariables 'DModeGradient r
  -> (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> (Domains r, r)
-- The functions in which @dReverseGeneral@ inlines are not inlined themselves
-- in client code, so the bloat is limited.
{-# INLINE dReverseGeneral #-}
dReverseGeneral dt
                variables@(params0, _, params1, _, params2, _, paramsX, _)
                f =
  let dim0 = V.length params0
      dim1 = V.length params1
      dim2 = V.length params2
      dimX = V.length paramsX
      initialState = initializeState (params0, params1, params2, paramsX)
      (D value d, st) = runState (runDualMonadGradient (f variables))
                                 initialState
      inlineDerivative primCode primalArgs dualArgs =
        let D _ u' = outDualNumber primCode primalArgs dualArgs
        in u'
      gradient =
        gradientFromDelta inlineDerivative inlineDerivative inlineDerivative
                          inlineDerivative inlineDerivative
                          dim0 dim1 dim2 dimX st d dt
  in (gradient, value)

dReverse :: HasDelta r
   => r
   -> (DualNumberVariables 'DModeGradient r
       -> DualMonadGradient r (DualNumber 'DModeGradient r))
   -> Domains r
   -> (Domains r, r)
dReverse dt f parameters =
  let varDeltas = generateDeltaVars parameters
      variables = makeDualNumberVariables parameters varDeltas
  in dReverseGeneral dt variables f

-- This function uses @DualMonadGradient@ for an inefficient computation
-- of forward derivaties. See @dFastForwardGeneral@ for an efficient variant.
dForwardGeneral
  :: forall r. HasDelta r
  => DualNumberVariables 'DModeGradient r
  -> (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Domains r
  -> (r, r)
{-# INLINE dForwardGeneral #-}
dForwardGeneral variables@(params0, _, params1, _, params2, _, paramsX, _)
                f ds =
  let initialState = initializeState (params0, params1, params2, paramsX)
      (D value d, st) = runState (runDualMonadGradient (f variables))
                                 initialState
      inlineDerivative primCode primalArgs dualArgs =
        let D _ u' = outDualNumber primCode primalArgs dualArgs
        in u'
  in ( derivativeFromDelta inlineDerivative inlineDerivative inlineDerivative
                           inlineDerivative inlineDerivative
                           st d ds
     , value )

-- The direction vector ds is taken as an extra argument.
dForward
  :: HasDelta r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Domains r
  -> Domains r
  -> (r, r)
dForward f parameters ds =
  let varDeltas = generateDeltaVars parameters
      variables = makeDualNumberVariables parameters varDeltas
  in dForwardGeneral variables f ds

-- * A monad for efficiently computing forward derivatives.

newtype DualMonadForward r a = DualMonadForward
  { runDualMonadForward :: Identity a }
  deriving (Monad, Functor, Applicative)

instance IsScalar 'DModeDerivative r
         => DualMonad 'DModeDerivative
                      r (DualMonadForward r) where
  returnLet (D u u') = DualMonadForward $ Identity $ D u u'

-- This the efficient variant of forward derivative computation.
dFastForwardGeneral
  :: Dual 'DModeDerivative r ~ r
  => DualNumberVariables 'DModeDerivative r
  -> (DualNumberVariables 'DModeDerivative r
      -> DualMonadForward r (DualNumber 'DModeDerivative r))
  -> (r, r)
{-# INLINE dFastForwardGeneral #-}
dFastForwardGeneral variables f =
  let D value d = runIdentity $ runDualMonadForward $ f variables
  in (d, value)

-- The direction vector ds is taken as an extra argument.
dFastForward
  :: forall r. (Numeric r, Dual 'DModeDerivative r ~ r)
  => (DualNumberVariables 'DModeDerivative r
      -> DualMonadForward r (DualNumber 'DModeDerivative r))
  -> Domains r
  -> Domains r
  -> (r, r)
dFastForward f parameters (params0, params1, params2, paramsX) =
  let variables =
        makeDualNumberVariables
          parameters
          (V.convert params0, params1, params2, paramsX)  -- ds
      (derivative, value) = dFastForwardGeneral variables f
  in (derivative, value)

-- * Additional mechanisms

prettyPrintDf
  :: forall r. HasDelta r
  => Bool
  -> (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> Domains r
  -> String
prettyPrintDf reversed f parameters =
  let varDeltas = generateDeltaVars parameters
      variables = makeDualNumberVariables parameters varDeltas
      initialState = initializeState parameters
      (D _ deltaTopLevel, st) = runState (runDualMonadGradient (f variables))
                                         initialState
  in ppBindings reversed st deltaTopLevel

generateDeltaVars
  :: forall r. IsScalar 'DModeGradient r
  => Domains r
  -> ( Data.Vector.Vector (Dual 'DModeGradient r)
     , Data.Vector.Vector (Dual 'DModeGradient (Vector r))
     , Data.Vector.Vector (Dual 'DModeGradient (Matrix r))
     , Data.Vector.Vector (Dual 'DModeGradient (OT.Array r)) )
generateDeltaVars (params0, params1, params2, paramsX) =
  let vVar :: forall a v. (IsPrimalWithScalar 'DModeGradient a r, V.Vector v a)
           => v a -> Data.Vector.Vector (Dual 'DModeGradient a)
      vVar p = V.generate (V.length p) (dVar . toDeltaId)
      !v0 = vVar params0
      !v1 = vVar params1
      !v2 = vVar params2
      !vX = vVar paramsX
  in (v0, v1, v2, vX)

-- | Initialize parameters using a uniform distribution with a fixed range
-- taken from an argument.
--
-- Must be Double, because @randomVector@ only works on that.
--
-- This only works fine for nets with levels that have similar size and use
-- the same activation function. Otherwise, the range should vary per level.
-- A rule of thumb range for weights is @sqrt(6 / (F_in + F_out)@,
-- where @F_in + F_out@ is the sum of inputs and outputs of the largest level.
-- See https://github.com/pytorch/pytorch/issues/15314 and their newer code.
initializerFixed :: Int -> Double -> (Int, [Int], [(Int, Int)], [OT.ShapeL])
                 -> ((Int, Int, Int, Int), Int, Double, Domains Double)
initializerFixed seed range (nParams0, lParams1, lParams2, lParamsX) =
  let vParams1 = V.fromList lParams1
      vParams2 = V.fromList lParams2
      vParamsX = V.fromList lParamsX
      createRandomVector n seedV =
        HM.scale (2 * range)
        $ HM.randomVector seedV HM.Uniform n - HM.scalar 0.5
      params0Init = createRandomVector nParams0 seed
      params1Init =
        V.imap (\i nPV -> createRandomVector nPV (seed + nPV + i)) vParams1
      params2Init =
        V.imap (\i (rows, cols) ->
                 HM.reshape cols
                 $ createRandomVector (rows * cols) (seed + rows + i)) vParams2
      paramsXInit =
        V.imap (\i sh ->
                 let sz = product sh
                 in OT.fromVector sh
                    $ createRandomVector sz (seed + sz + i)) vParamsX
      totalParams = nParams0
                    + V.sum vParams1
                    + V.sum (V.map (uncurry (*)) vParams2)
                    + V.sum (V.map product vParamsX)
  in ( (nParams0, V.length vParams1, V.length vParams2, V.length vParamsX)
     , totalParams
     , range
     , (params0Init, params1Init, params2Init, paramsXInit) )

outDualNumber :: (IsPrimal d a, RealFloat a, Show a, Show (Dual d a))
              => CodeOut -> [a] -> [Dual d a]
              -> DualNumber d a
outDualNumber PlusOut [u, v] [u', v'] = D u u' + D v v'
outDualNumber MinusOut [u, v] [u', v'] = D u u' - D v v'
outDualNumber TimesOut [u, v] [u', v'] = D u u' * D v v'
outDualNumber NegateOut [u] [u'] = negate $ D u u'
outDualNumber AbsOut [u] [u'] = abs $ D u u'
outDualNumber SignumOut [u] [u'] = signum $ D u u'
outDualNumber DivideOut [u, v] [u', v'] = D u u' / D v v'
outDualNumber RecipOut [u] [u'] = recip $ D u u'
outDualNumber ExpOut [u] [u'] = exp $ D u u'
outDualNumber LogOut [u] [u'] = log $ D u u'
outDualNumber SqrtOut [u] [u'] = sqrt $ D u u'
outDualNumber PowerOut [u, v] [u', v'] = D u u' ** D v v'
outDualNumber LogBaseOut [u, v] [u', v'] = logBase (D u u') (D v v')
outDualNumber SinOut [u] [u'] = sin $ D u u'
outDualNumber CosOut [u] [u'] = cos $ D u u'
outDualNumber TanOut [u] [u'] = tan $ D u u'
outDualNumber AsinOut [u] [u'] = asin $ D u u'
outDualNumber AcosOut [u] [u'] = acos $ D u u'
outDualNumber AtanOut [u] [u'] = atan $ D u u'
outDualNumber SinhOut [u] [u'] = sinh $ D u u'
outDualNumber CoshOut [u] [u'] = cosh $ D u u'
outDualNumber TanhOut [u] [u'] = tanh $ D u u'
outDualNumber AsinhOut [u] [u'] = asinh $ D u u'
outDualNumber AcoshOut [u] [u'] = acosh $ D u u'
outDualNumber AtanhOut [u] [u'] = atanh $ D u u'
outDualNumber Atan2Out [u, v] [u', v'] = atan2 (D u u') (D v v')
outDualNumber primCode primalArgs dualArgs =
  error $ "outDualNumber: wrong number of arguments"
          ++ show (primCode, primalArgs, dualArgs)


{-# LANGUAGE FlexibleContexts, ScopedTypeVariables #-}

{-| Fast FFTs using FFTW -}
module SDR.FFT (
    -- * Windows
    hamming,
    hanning,
    blackman,

    -- * FFTs
    fftw',
    fftw,
    fftwReal',
    fftwReal,
    fftwParallel
    ) where

import           Control.Monad                as CM
import           Foreign.Storable
import           Foreign.Storable.Complex
import           Foreign.C.Types
import           Data.Complex
import           Foreign.ForeignPtr
import           Control.Concurrent           hiding (yield)
import qualified Data.Map                     as Map
import           Data.Coerce

import qualified Data.Vector.Generic          as VG
import qualified Data.Vector.Storable         as VS
import qualified Data.Vector.Storable.Mutable as VSM

import           Pipes
import           Numeric.FFTW

import           SDR.FilterDesign
import           SDR.VectorUtils

mallocForeignBufferAligned :: forall a. Storable a => Int -> IO (ForeignPtr a)
mallocForeignBufferAligned elems = do
    ptr <- fftwMalloc $ fromIntegral $ elems * sizeOf (undefined :: a)
    newForeignPtr fftwFreePtr ptr

-- | Creates a function that performs a complex to complex DFT.
fftw' :: (VG.Vector v (Complex Double))
     => Int -- ^ The size of the input and output buffers
     -> IO (v (Complex Double) -> IO (VS.Vector (Complex Double)))
fftw' samples = do
    ina <- mallocForeignBufferAligned samples
    out <- mallocForeignBufferAligned samples

    plan <- withForeignPtr ina $ \ip ->
        withForeignPtr out $ \op ->
            planDFT1d samples ip op Forward fftwEstimate

    return $ \inv' -> do
        out <- mallocForeignBufferAligned samples
        ina <- mallocForeignBufferAligned samples
        let inv = VSM.unsafeFromForeignPtr0 ina samples

        copyInto inv inv'

        let (fp, offset, length) = VSM.unsafeToForeignPtr inv

        withForeignPtr (coerce fp) $ \fpp ->
            withForeignPtr (coerce out) $ \op ->
                executeDFT plan fpp op

        return $ VS.unsafeFromForeignPtr0 out samples

-- | Creates a Pipe that performs a complex to complex DFT.
fftw :: (VG.Vector v (Complex Double))
     => Int -- ^ The size of the input and output buffers
     -> IO (Pipe (v (Complex Double)) (VS.Vector (Complex Double)) IO ())
fftw samples = do
    func <- fftw' samples
    return $ for cat $ \dat -> lift (func dat) >>= yield

-- | Creates a function that performs a real to complex DFT.
fftwReal' :: (VG.Vector v Double)
         => Int -- ^ The size of the input Vector
         -> IO (v Double -> IO (VS.Vector (Complex Double)))
fftwReal' samples = do
    --Allocate in and out buffers that wont be used because there doesnt seem to be a way to create a plan without them
    ina <- mallocForeignBufferAligned samples
    out <- mallocForeignBufferAligned samples

    plan <- withForeignPtr ina $ \ip ->
        withForeignPtr out $ \op ->
            planDFTR2C1d samples ip op fftwEstimate

    return $ \inv' -> do
        out <- mallocForeignBufferAligned ((samples `quot` 2) + 1)
        ina <- mallocForeignBufferAligned samples
        let inv = VSM.unsafeFromForeignPtr0 ina samples

        copyInto inv inv'
        let (fp, offset, length) = VSM.unsafeToForeignPtr inv

        withForeignPtr (coerce fp)  $ \fpp ->
            withForeignPtr (coerce out) $ \op ->
                executeDFTR2C plan fpp op

        return $ VS.unsafeFromForeignPtr0 out samples

-- | Creates a pipe that performs a real to complex DFT.
fftwReal :: (VG.Vector v Double)
         => Int -- ^ The size of the input Vector
         -> IO (Pipe (v Double) (VS.Vector (Complex Double)) IO ())
fftwReal samples = do
    func <- fftwReal' samples
    return $ for cat $ \dat -> lift (func dat) >>= yield

{-| Creates a pipe that uses multiple threads to perform complex to complex DFTs in
    a pipelined fashion. Each time a buffer is consumed, it is given to
    a pool of threads to perform the DFT. Then, if a thread has finished
    performing a previous DFT, the result is yielded.
-}
fftwParallel :: (VG.Vector v (Complex Double))
             => Int -- ^ The number of threads to use
             -> Int -- ^ The size of the input Vector
             -> IO (Pipe (v (Complex Double)) (VS.Vector (Complex Double)) IO ())
fftwParallel threads samples = do
    --plan the DFT
    ina <- mallocForeignBufferAligned samples
    out <- mallocForeignBufferAligned samples

    plan <- withForeignPtr ina $ \ip ->
        withForeignPtr out $ \op ->
            planDFT1d samples ip op Forward fftwEstimate

    --setup the channels and worker threads
    inChan <- newChan
    outMap <- newMVar Map.empty

    CM.replicateM threads $ forkIO $ forever $ do
        (idx, res) <- readChan inChan

        out <- mallocForeignBufferAligned samples
        ina <- mallocForeignBufferAligned samples
        let inv = VSM.unsafeFromForeignPtr0 ina samples

        copyInto inv res

        let (fp, offset, length) = VSM.unsafeToForeignPtr inv

        withForeignPtr (coerce fp) $ \fpp ->
            withForeignPtr (coerce out) $ \op ->
                executeDFT plan fpp op

        theMap <- takeMVar outMap
        putMVar outMap $ Map.insert idx (VS.unsafeFromForeignPtr0 out samples) theMap

    --build the pipe
    let pipe nextIn nextOut = do
            dat <- await
            lift $ writeChan inChan (nextIn, dat)

            theMap <- lift $ takeMVar outMap
            case Map.lookup nextOut theMap of
                Nothing  -> do
                    lift $ putMVar outMap theMap
                    pipe (nextIn + 1) nextOut
                Just dat -> do
                    lift $ putMVar outMap $ Map.delete nextOut theMap
                    yield dat
                    pipe (nextIn + 1) (nextOut + 1)

    return $ pipe 0 0


module Core.Language where

import Control.Monad.Except
import Data.Array
import Data.Complex
import Data.IORef
import Data.Ratio
import System.IO
import Text.Printf (printf)

import Core.Error

-- SCHEME DATA TYPES -----------------------------------------------------------

type SchemeEnv = IORef [(String, IORef SchemeVal)]

type SchemeFun = [SchemeVal] -> ThrowsError SchemeVal

data SchemeVal = Atom String
               | Boolean Bool
               | Char Char
               | String String
               | Integer Int
               | Rational Rational
               | Real Double
               | Complex (Complex Double)
               | List [SchemeVal]
               | ImproperList [SchemeVal] SchemeVal
               | Vector (Array Int SchemeVal)
               | Bytevector (Array Int Integer)
               | Primitive SchemeFun
               | Func { args :: [String],
                        vararg :: Maybe String,
                        body :: [SchemeVal],
                        closure :: SchemeEnv }
               | IOPrimitive ([SchemeVal] -> IOThrowsError SchemeVal)
               | IOPort Handle
               | Unspecified

-- STRING REPRESENTATIONS ------------------------------------------------------

instance Show SchemeVal where show = showVal

showVal (Atom name) = name
showVal (Boolean bool) = if bool then "#t" else "#f"
showVal (Char ch) = show ch
showVal (String str) = printf "\"%s\"" str
showVal (Integer n) = show n
showVal (Rational q) = printf "%d/%d" (numerator q) (denominator q)
showVal (Real x) = show x
showVal (Complex z) = printf "%f+%fi" (realPart z) (imagPart z)
showVal (List xs) = printf "(%s)" (unwordsList xs)
showVal (ImproperList xs x) = printf "(%s . %s)" (unwordsList xs) (showVal x)
showVal (Vector array) = printf "#(%s)" (unwordsArray array)
showVal (Bytevector array) = printf "#u8(%s)" (unwordsArray array)
showVal (Primitive _) = "<primitive>"
showVal (Func args vararg _ _) = showLambda args vararg
showVal (IOPrimitive _) = "<IO primitive>"
showVal (IOPort _) = "<IO port>"
showVal Unspecified = "<unspecified>"

-- UTILS -----------------------------------------------------------------------

-- throws a wrong number of arguments error
throwWrongNumArgs :: Int -> [SchemeVal] -> ThrowsError SchemeVal
throwWrongNumArgs = (throwError .) . (. unwordsList) . WrongNumArgs

-- throws a type mismatch error
throwTypeMismatch :: String -> SchemeVal -> ThrowsError SchemeVal
throwTypeMismatch = (throwError .) . (. show) . TypeMismatch

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

unwordsArray :: Show a => Array Int a -> String
unwordsArray = unwordsList . elems

showLambda :: [String] -> Maybe String -> String
showLambda = (. maybe [] (" . " ++)) . printf "(lambda (%s%s) ...)" . unwords


{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE TypeFamilies          #-}
-- |
-- Module     : Data.Matrix.Generic
-- Copyright  : Copyright (c) 2012 Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- License    : BSD3
-- Maintainer : Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- Stability  : experimental
--
-- Interface for generic immutable matrices. For matrix transposition
-- and conjugate transposition newtype wrappers are used.
module Data.Matrix.Generic (
    -- * Type class
    IsMatrix(..)
  , Mutable
    -- * Accessors
  , rows
  , cols
  , shape
  , (@!)
  , unsafeIndex
    -- * Converions to/from mutable
  , freeze
  , thaw
  , unsafeFreeze
  , unsafeThaw
    -- * Conversion to string
  , showMatrixWith
    -- * Newtype wrappers
  , Transposed(..)
  , Conjugated(..)
  ) where

import Control.Monad             (liftM)
import Control.Monad.Primitive
import Data.List                 (intercalate)
import Data.Complex              (Complex,conjugate)
import qualified Data.Matrix.Generic.Mutable as M
import           Data.Vector.Generic           (Mutable)

----------------------------------------------------------------
-- Type class
----------------------------------------------------------------

-- | Basic API for immutable matrices. Since there's many way to lay
--   matrix in memory there isn't many operation which work for all of
--   them.
--
--   Methods of this type class shouldn't be used directly.
class M.IsMMatrix (Mutable mat) a => IsMatrix mat a where
  -- | Number of rows
  basicRows :: mat a -> Int
  -- | Number of columns
  basicCols :: mat a -> Int
  -- | Read element from matrix
  basicUnsafeIndex :: mat a -> (Int,Int) -> a
  -- | Convert immutable matrix to mutable. Immutable matrix may not
  --   be used after operation.
  basicUnsafeThaw   :: PrimMonad m => mat a -> m (Mutable mat (PrimState m) a)
  -- | Convert mutable matrix to immutable. Mutable matrix may not be
  --   modified after operation.
  basicUnsafeFreeze :: PrimMonad m => Mutable mat (PrimState m) a -> m (mat a)

----------------------------------------------------------------
-- Accessors
----------------------------------------------------------------

-- | Number of rows
rows :: IsMatrix mat a => mat a -> Int
{-# INLINE rows #-}
rows = basicRows

-- | Number of columns
cols :: IsMatrix mat a => mat a -> Int
{-# INLINE cols #-}
cols = basicCols

-- | Shape of the matrix.
shape :: IsMatrix mat a => mat a -> (Int,Int)
{-# INLINE shape #-}
shape m = (rows m, cols m)

-- | Indexing operator without range checking.
unsafeIndex :: IsMatrix mat a
            => mat a            -- ^ Matrix
            -> (Int,Int)        -- ^ (row,column)
            -> a
{-# INLINE unsafeIndex #-}
unsafeIndex = basicUnsafeIndex

-- | Indexing operator with range checking
(@!) :: IsMatrix mat a
     => mat a                   -- ^ Matrix
     -> (Int,Int)               -- ^ (row,column)
     -> a
{-# INLINE (@!) #-}
m @! a@(i,j)
  | i < 0 || i >= rows m = error "Data.Matrix.Generic.@!: row index out of range"
  | j > 0 || j >= cols m = error "Data.Matrix.Generic.@!: column index out of range"
  | otherwise            = unsafeIndex m a

-- | Convert mutable matrix to immutable.
freeze :: (PrimMonad m, IsMatrix mat a) => Mutable mat (PrimState m) a -> m (mat a)
{-# INLINE freeze #-}
freeze m = do
  unsafeFreeze =<< M.clone m

-- | Convert immutable matrix to mutable.
thaw :: (PrimMonad m, IsMatrix mat a) => mat a -> m (Mutable mat (PrimState m) a)
{-# INLINE thaw #-}
thaw m = do
  M.clone =<< unsafeThaw m

-- | Convert mutable matrix to immutable. Mutable matrix may not be
--   modified after operation.
unsafeFreeze :: (PrimMonad m, IsMatrix mat a) => Mutable mat (PrimState m) a -> m (mat a)
{-# INLINE unsafeFreeze #-}
unsafeFreeze = basicUnsafeFreeze

-- | Convert immutable matrix to mutable. Immutable matrix may not
--   be used after operation.
unsafeThaw :: (PrimMonad m, IsMatrix mat a) => mat a -> m (Mutable mat (PrimState m) a)
{-# INLINE unsafeThaw #-}
unsafeThaw = basicUnsafeThaw

-- | Generic function for printing matrix. Mostly useful for debugging
--   purposes.
showMatrixWith :: IsMatrix m a => (a -> String) -> m a -> String
showMatrixWith f m
  = unlines
  $ (show (rows m) ++ " >< " ++ show (cols m))
  : [ intercalate "\t" [f $ unsafeIndex m (i,j) | j <- [0 .. cols m - 1]]
    | i <- [0 .. rows m - 1]
    ]

----------------------------------------------------------------
-- Newtype wrappers
----------------------------------------------------------------

-- | Transposed matrix or vector. Being newtype this wrapper type is
--   used to select different instances for multiplication.
newtype Transposed mat a = Transposed { unTranspose :: mat a }

type instance Mutable (Transposed mat) = M.TransposedM (Mutable mat)

instance IsMatrix mat a => IsMatrix (Transposed mat) a where
  basicRows (Transposed m) = cols m
  {-# INLINE basicRows #-}
  basicCols (Transposed m) = rows m
  {-# INLINE basicCols #-}
  basicUnsafeIndex (Transposed m) (i,j) = unsafeIndex m (j,i)
  {-# INLINE basicUnsafeIndex #-}
  basicUnsafeThaw (Transposed m) = M.TransposedM `liftM` unsafeThaw m
  {-# INLINE basicUnsafeThaw #-}
  basicUnsafeFreeze (M.TransposedM m) = Transposed `liftM` unsafeFreeze m
  {-# INLINE basicUnsafeFreeze #-}

-- | Conjugate-transposed matrix or vector. Being newtype this wrapper type is
--   used to select different instances for multiplication.
newtype Conjugated mat a = Conjugated { unConjugate :: mat a }

type instance Mutable (Conjugated mat) = M.ConjugatedM (Mutable mat)

instance (IsMatrix mat (Complex a), RealFloat a) => IsMatrix (Conjugated mat) (Complex a) where
  basicRows (Conjugated m) = cols m
  {-# INLINE basicRows #-}
  basicCols (Conjugated m) = rows m
  {-# INLINE basicCols #-}
  basicUnsafeIndex (Conjugated m) (i,j)
    | j >= i    = unsafeIndex m (j,i)
    | otherwise = conjugate $! unsafeIndex m (j,i)
  {-# INLINE basicUnsafeIndex #-}
  basicUnsafeThaw (Conjugated m) = M.ConjugatedM `liftM` unsafeThaw m
  {-# INLINE basicUnsafeThaw #-}
  basicUnsafeFreeze (M.ConjugatedM m) = Conjugated `liftM` unsafeFreeze m
  {-# INLINE basicUnsafeFreeze #-}


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

module Deep.Expr where

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

import           Data.Bifunctor

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

import           Data.Complex

import           Control.Monad.State

import           Data.Type.Equality
import           Data.Typeable

import           Data.Char

import           GHC.Prim
import           Data.Coerce

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

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

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

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

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

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

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

--   EmptyMatch :: forall r. SumMatch Void r

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

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

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

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

getNameUniq :: Name a -> Int
getNameUniq = getNameIdent

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

newtype Env = Env [EnvMapping]

emptyEnv :: Env
emptyEnv = Env []

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

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

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

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

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

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

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

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

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

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

  FalseExp :: GPUExp Bool
  TrueExp :: GPUExp Bool

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

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

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

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

  Ord :: GPUExp Char -> GPUExp Int

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

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

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

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

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

  Not :: GPUExp Bool -> GPUExp Bool

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

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

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

  UnitExp :: GPUExp ()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  construct :: t -> GPUExp (GPURepTy t)

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

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

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

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

  rep :: t -> GPUExp t

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

  rep' :: t -> GPURepTy t

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

  unrep' :: GPURepTy t -> t

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

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

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

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

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

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

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

instance GPURep Bool where
  type GPURepTy Bool = Bool

  construct = rep

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

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

instance GPURep Char where
  type GPURepTy Char = Char

  construct = rep --CharLit

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

instance GPURep a => GPURep (Complex a) where

instance GPURep a => GPURep (Maybe a) where

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

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

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

  rep' = id
  unrep' = id

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

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

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

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

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

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

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

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

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

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

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

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

  rep (M1 x) = ConstructRep (construct x)

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

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

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

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

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

  -- constructRep (LeftExp x) = _

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

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

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

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

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

  -- construct = construct . unK1

  construct = rep . unK1

  -- rep = _

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

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

  construct _ = UnitExp

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

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

  rep () = UnitExp

  rep' = id
  unrep' = id

class GenericRep repr where
    type GenericRepTy repr

    genericRep :: repr -> GPUExp (GenericRepTy repr)

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

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

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

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

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

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

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

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

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

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

    genericRep = genericRep . unM1

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

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

    genericRep _ = UnitExp

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

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

    genericRep = error "genericRep called for K1 instance"

    genericRep' = unK1
    genericUnrep' = K1

the :: a -> a
the = id

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

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

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

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

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

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

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

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

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

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

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

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

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

        go = transformE go0
    in

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

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

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

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

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

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

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

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

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


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

module Properties (tests) where

import Control.Applicative as A ((<$>))
import Criterion.Analysis
import Prelude ()
import Prelude.Compat
import Statistics.Types (Sample)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.QuickCheck (testProperty)
import Test.QuickCheck
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U

instance (Arbitrary a, U.Unbox a) => Arbitrary (U.Vector a) where
  arbitrary = U.fromList A.<$> arbitrary
  shrink    = map U.fromList . shrink . U.toList

outlier_bucketing :: Double -> Sample -> Bool
outlier_bucketing y ys =
  countOutliers (classifyOutliers xs) <= fromIntegral (G.length xs)
  where xs = U.cons y ys

outlier_bucketing_weighted :: Double -> Sample -> Bool
outlier_bucketing_weighted x xs =
  outlier_bucketing x (xs <> G.replicate (G.length xs * 10) 0)

tests :: TestTree
tests = testGroup "Properties" [
    testProperty "outlier_bucketing" outlier_bucketing
  , testProperty "outlier_bucketing_weighted" outlier_bucketing_weighted
  ]


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeInType          #-}
{-# LANGUAGE TypeOperators       #-}

#if __GLASGOW_HASKELL__ >= 805
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE ExplicitNamespaces #-}
#endif

module Eigen.Matrix.Mutable
  ( -- * Types
    MMatrix(..)
  , MMatrixXf
  , MMatrixXd
  , MMatrixXcf
  , MMatrixXcd
  , IOMatrix
  , STMatrix

    -- * Construction
  , new
  , replicate

    -- * Indexing
  , read
  , write

    -- * Modification
  , set
  , copy

    -- * Modification with Pointers
  , unsafeWith

    -- * Conversion to Vectors
  , vals
  , fromVector
  ) where

import Eigen.Internal
  ( Elem
  , C(..)
  , natToInt
  , Row(..)
  , Col(..)
  )

import Prelude hiding (replicate, read)
import Control.Monad.Primitive (PrimMonad(..))
import Data.Complex (Complex)
import Data.Kind (Type)
import Foreign.C.Types (CInt)
import Foreign.Ptr (Ptr)
import GHC.Exts (RealWorld)
import GHC.TypeLits (Nat, type (*), type (<=), KnownNat)
import qualified Data.Vector.Storable.Mutable as VSM

-- | A mutable matrix. See 'Eigen.Matrix.Matrix' for
--   details about matrix layout.
newtype MMatrix :: Nat -> Nat -> Type -> Type -> Type where
  MMatrix :: Vec (n * m) s a -> MMatrix n m s a

-- | Used internally to track the size and corresponding C type of the matrix.
newtype Vec :: Nat -> Type -> Type -> Type where
  Vec :: VSM.MVector s (C a) -> Vec n s a

-- | Alias for single precision mutable matrix
type MMatrixXf  n m s = MMatrix n m s Float
-- | Alias for double precision mutable matrix
type MMatrixXd  n m s = MMatrix n m s Double
-- | Alias for single precision mutable matrix of complex numbers
type MMatrixXcf n m s = MMatrix n m s (Complex Float)
-- | Alias for double precision mutable matrix of complex numbers
type MMatrixXcd n m s = MMatrix n m s (Complex Double)

-- | A mutable matrix where the state token is specialised to 'RealWorld'.
type IOMatrix n m a   = MMatrix n m RealWorld a
-- | This type does not differ from MSparseMatrix, but might be desirable for readability.
type STMatrix n m s a = MMatrix n m s a

-- | Create a mutable matrix of the given size and fill it with 0 as an initial value.
new :: (PrimMonad p, Elem a, KnownNat n, KnownNat m) => p (MMatrix n m (PrimState p) a)
{-# INLINE new #-}
new = replicate 0

-- | Create a mutable matrix of the given size and fill it with an initial value.
replicate :: forall n m p a. (PrimMonad p, Elem a, KnownNat n, KnownNat m) => a -> p (MMatrix n m (PrimState p) a)
{-# INLINE replicate #-}
replicate !val = do
  let !mm_rows = natToInt @n
      !mm_cols = natToInt @m
      !cval    = toC val
  _vals <- VSM.replicate (mm_rows * mm_cols) cval
  pure (MMatrix $! Vec $! _vals)

-- | Set all elements of the matrix to a given value.
set :: (PrimMonad p, Elem a) => MMatrix n m (PrimState p) a -> a -> p ()
{-# INLINE set #-}
set (MMatrix (Vec !vec)) !val =
  let !cval = toC val
  in VSM.set vec cval

-- | Copy a matrix.
copy :: (PrimMonad p, Elem a) => MMatrix n m (PrimState p) a -> MMatrix n m (PrimState p) a -> p ()
{-# INLINE copy #-}
copy (MMatrix (Vec m1)) (MMatrix (Vec m2)) = VSM.unsafeCopy m1 m2

-- | Yield the element at the given position.
read :: forall n m p a r c. (PrimMonad p, Elem a, KnownNat n, KnownNat r, KnownNat c, r <= n, c <= m)
  => Row r -> Col c -> MMatrix n m (PrimState p) a -> p a
{-# INLINE read #-}
read _ _ (MMatrix (Vec m)) =
  let !row = natToInt @r
      !col = natToInt @c
      !mm_rows = natToInt @n
  in VSM.unsafeRead m (col * mm_rows + row) >>= \ !val -> let !cval = fromC val in pure cval

-- | Replace the element at the given position.
write :: forall n m p a r c. (PrimMonad p, Elem a, KnownNat n, KnownNat r, KnownNat c, r <= n, c <= m)
  => Row r -> Col c -> MMatrix n m (PrimState p) a -> a -> p ()
{-# INLINE write #-}
write _ _ (MMatrix (Vec m)) !val =
  let !row = natToInt @r
      !col = natToInt @c
      !mm_rows = natToInt @n
      !cval = toC val
  in VSM.unsafeWrite m (col * mm_rows + row) cval

-- | Pass a pointer to the matrix's data to the IO action.
--   Modifying dat through the pointer is unsafe if the matrix
--   could have been frozen before the modification.
unsafeWith :: forall n m a b. (KnownNat n, KnownNat m, Elem a) => IOMatrix n m a -> (Ptr (C a) -> CInt -> CInt -> IO b) -> IO b
{-# INLINE unsafeWith #-}
unsafeWith (MMatrix (Vec m)) f =
  let !cmm_rows = toC $! natToInt @n
      !cmm_cols = toC $! natToInt @m
  in VSM.unsafeWith m $ \p -> f p cmm_rows cmm_cols

-- | Return a mutable storable 'VSM.MVector' of the corresponding C types to one's mutable matrix.
vals :: MMatrix n m s a -> VSM.MVector s (C a)
{-# INLINE vals #-}
vals (MMatrix (Vec x)) = x

-- | Create a mutable matrix from a mutable storable 'VSM.MVector'.
fromVector :: Elem a => VSM.MVector s (C a) -> MMatrix n m s a
{-# INLINE fromVector #-}
fromVector x = MMatrix (Vec x)


module Repl.Parser (parseExpr) where

import Data.Array
import Data.Char
import Data.Complex
import Data.Maybe
import Data.Ratio
import Numeric
import Text.ParserCombinators.Parsec

import Core.Language

parseExpr :: Parser SchemeVal
parseExpr = tryAll parsers

parsers = [parseAtom,
           parseAtomEnclosed,
           parseBoolean,
           parseChar,
           parseHexChar,
           parseSpecialChar,
           parseString,
           parseComplex,
           parseReal,
           parseRational,
           parseInteger,
           parseList,
           parseImproperList,
           parseVector,
           parseBytevector,
           parseQuote,
           parseUnquote,
           parseQuasiquote]

-- PARSERS ---------------------------------------------------------------------

parseAtom :: Parser SchemeVal
parseAtom = do
  str <- many1 $ letter <|> digit <|> symbol
  if isDigit (head str) then fail "digit not allowed" else return $ Atom str

parseAtomEnclosed :: Parser SchemeVal
parseAtomEnclosed = do
  char '|'
  str <- many1 $ letter <|> digit <|> symbol <|> hexChar <|> space
  char '|'
  return $ Atom str

parseBoolean :: Parser SchemeVal
parseBoolean = do
  char '#'
  c <- oneOf "tf"
  return $ Boolean $ c == 't'

parseChar :: Parser SchemeVal
parseChar = do
  string "#\\"
  x <- anyChar
  notFollowedBy alphaNum
  return $ Char x

parseHexChar :: Parser SchemeVal
parseHexChar = fmap (Char . chr . fromHex) (string "#\\x" >> many1 hexDigit)

parseSpecialChar :: Parser SchemeVal
parseSpecialChar = do
  string "#\\"
  str <- tryAll $ map (string . fst) specialChars
  return $ Char $ fromJust $ lookup str specialChars

parseString :: Parser SchemeVal
parseString = do
  char '"'
  str <- many $ tryAll [escChar, hexChar, noneOf "\"\\"]
  char '"'
  return $ String str

parseInteger :: Parser SchemeVal
parseInteger = Integer <$> tryAll [int, hex, dec, oct, bin]

parseRational :: Parser SchemeVal
parseRational = do
  a <- many1 digit
  char '/'
  b <- many1 digit
  return $ Rational $ read a % read b

parseReal :: Parser SchemeVal
parseReal = do
  a <- many1 digit
  char '.'
  b <- many1 digit
  return $ Real $ fst . head . readFloat $ a ++ "." ++ b

parseComplex :: Parser SchemeVal
parseComplex = do
  x <- tryAll [parseReal, parseInteger]
  char '+'
  y <- tryAll [parseReal, parseInteger]
  char 'i'
  return $ Complex $ fromNumber x :+ fromNumber y

parseList :: Parser SchemeVal
parseList = do
  char '('
  xs <- sepBy parseExpr spaces
  char ')'
  return $ List xs

parseImproperList :: Parser SchemeVal
parseImproperList = do
  char '('
  xs <- endBy parseExpr spaces
  char '.'
  x <- spaces >> parseExpr
  char ')'
  return $ ImproperList xs x

parseVector :: Parser SchemeVal
parseVector = do
  string "#("
  xs <- sepBy parseExpr spaces
  char ')'
  return $ Vector $ list2array xs

parseBytevector :: Parser SchemeVal
parseBytevector = do
  string "#u8("
  xs <- sepBy byte spaces
  char ')'
  return $ Bytevector $ list2array xs

parseQuote :: Parser SchemeVal
parseQuote = do
  char '\''
  x <- parseExpr
  return $ List [Atom "quote", x]

parseUnquote :: Parser SchemeVal
parseUnquote = do
  char ','
  x <- parseExpr
  return $ List [Atom "unquote", x]

parseQuasiquote :: Parser SchemeVal
parseQuasiquote = do
  char '`'
  x <- parseExpr
  return $ List [Atom "quasiquote", x]

-- UTILS -----------------------------------------------------------------------

tryAll = foldl1 (<|>) . map try

symbol = oneOf "!$%&|*+-/:<=>?@^_~"

hexChar = do
  string "\\x"
  x <- many1 hexDigit
  char ';'
  return $ chr $ fromHex x

escChar = do
  char '\\'
  x <- oneOf "\\\"|abnrt"
  return $ case x of
    '\\' -> x
    '"'  -> x
    '|'  -> '|'
    'a'  -> '\a'
    'b'  -> '\b'
    'n'  -> '\n'
    'r'  -> '\r'
    't'  -> '\t'

specialChars = [("alarm", '\a'), ("backspace", '\b'), ("delete", '\DEL'),
                ("escape", '\ESC'), ("newline", '\n'), ("null", '\NUL'),
                ("return", '\r'), ("space", ' '), ("tab", '\t')]

int = read <$> many1 digit

hex = fmap fromHex $ string "#x" >> many1 hexDigit
dec = fmap fromDec $ string "#d" >> many1 decDigit
oct = fmap fromOct $ string "#o" >> many1 octDigit
bin = fmap fromBin $ string "#b" >> many1 binDigit

fromHex = fst . head . readHex
fromDec = read
fromOct = fst . head . readOct
fromBin = foldl1 (\ a b -> 2 * a + b) . map digitToInt

decDigit = digit

binDigit = oneOf "01"

fromNumber :: SchemeVal -> Double
fromNumber (Real x) = x
fromNumber (Integer x) = fromIntegral x

list2array xs = listArray (0, length xs - 1) xs

byte = do
  x <- read <$> many1 digit
  if 0 <= x && x < 256 then return x else fail $ show x ++ " is not a byte"


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

--
-- 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.Constructs.Num
    ( NUM (..)
    ) where

import Language.Syntactic
import Language.Syntactic.Constructs.Binding

import Data.Complex (Complex(..))

import Feldspar.Range
import Feldspar.Core.Types
import Feldspar.Core.Interpretation
import Feldspar.Core.Constructs.Literal
import Feldspar.Core.Constructs.Integral
import Feldspar.Core.Constructs.Complex

data NUM a
  where
    Abs  :: (Type a, Num a, Num (Size a)) => NUM (a :-> Full a)
    Sign :: (Type a, Num a, Num (Size a)) => NUM (a :-> Full a)
    Add  :: (Type a, Num a, Num (Size a)) => NUM (a :-> a :-> Full a)
    Sub  :: (Type a, Num a, Num (Size a)) => NUM (a :-> a :-> Full a)
    Mul  :: (Type a, Num a, Num (Size a)) => NUM (a :-> a :-> Full a)

instance Semantic NUM
  where
    semantics Abs  = Sem "abs" abs
    semantics Sign = Sem "signum" signum
    semantics Add  = Sem "(+)" (+)
    semantics Sub  = Sem "(-)" (-)
    semantics Mul  = Sem "(*)" (*)

instance Equality NUM where equal = equalDefault; exprHash = exprHashDefault
instance Render   NUM where renderArgs = renderArgsDefault
instance ToTree   NUM
instance Eval     NUM where evaluate = evaluateDefault
instance EvalBind NUM where evalBindSym = evalBindSymDefault
instance Sharable NUM

instance AlphaEq dom dom dom env => AlphaEq NUM NUM dom env
  where
    alphaEqSym = alphaEqSymDefault

instance SizeProp (NUM :|| Type)
  where
    sizeProp (C' Abs)  (WrapFull a :* Nil)               = abs (infoSize a)
    sizeProp (C' Sign) (WrapFull a :* Nil)               = signum (infoSize a)
    sizeProp (C' Add)  (WrapFull a :* WrapFull b :* Nil) = infoSize a + infoSize b
    sizeProp (C' Sub)  (WrapFull a :* WrapFull b :* Nil) = infoSize a - infoSize b
    sizeProp (C' Mul)  (WrapFull a :* WrapFull b :* Nil) = infoSize a * infoSize b

instance ( (NUM     :|| Type) :<: dom
         , (Literal :|| Type) :<: dom
         , (INTEGRAL :|| Type) :<: dom
         , (COMPLEX :|| Type) :<: dom
         , OptimizeSuper dom
         )
      => Optimize (NUM :|| Type) dom
  where
    constructFeatOpt _ (C' Abs) (a :* Nil)
        | RangeSet r <- infoRange (getInfo a)
        , isNatural r
        = return a

    constructFeatOpt _ (C' Sign) (a :* Nil)
        | RangeSet ra <- infoRange (getInfo a)
        , 0 `rangeLess` ra
        = return (literalDecor 1)

    constructFeatOpt _ (C' Sign) (a :* Nil)
        | RangeSet ra <- infoRange (getInfo a)
        , ra `rangeLess` 0
        = return (literalDecor (-1))

    constructFeatOpt opts (C' Add) (a :* b :* Nil)
        | Just 0 <- viewLiteral b = return a
        | Just 0 <- viewLiteral a = return b
        | alphaEq a b = constructFeatOpt opts (c' Mul) (a :* literalDecor 2 :* Nil)

    constructFeatOpt opts s@(C' Add) (a :* (op :$ b :$ c) :* Nil)
        | Just al       <- viewLiteral a
        , Just (C' Add) <- prjF op
        , Just cl       <- viewLiteral c
        = constructFeat opts s (b :* literalDecor (al+cl) :* Nil)

    constructFeatOpt opts s@(C' Add) (a :* (op :$ b :$ c) :* Nil)
        | Just al       <- viewLiteral a
        , Just (C' Sub) <- prjF op
        , Just cl       <- viewLiteral c
        = constructFeat opts s (b :* literalDecor (al-cl) :* Nil)

    constructFeatOpt opts s@(C' Add) ((op :$ a :$ b) :* c :* Nil)
        | Just cl       <- viewLiteral c
        , Just (C' Add) <- prjF op
        , Just bl       <- viewLiteral b
        = constructFeat opts s (a :* literalDecor (bl+cl) :* Nil)

    constructFeatOpt opts s@(C' Add) ((op :$ a :$ b) :* c :* Nil)
        | Just cl       <- viewLiteral c
        , Just (C' Sub) <- prjF op
        , Just bl       <- viewLiteral b
        = constructFeat opts s (a :* literalDecor (cl-bl) :* Nil)

    constructFeatOpt opts (C' Add) ((op1 :$ a :$ b) :* (op2 :$ c :$ d) :* Nil)
        | Just (C' Add) <- prjF op1
        , Just (C' Add) <- prjF op2
        , Just bl       <- viewLiteral b
        , Just dl       <- viewLiteral d
        = do
            ac <- constructFeat opts (c' Add) (a :* c :* Nil)
            constructFeat opts (c' Add) (ac :* literalDecor (bl+dl) :* Nil)

    constructFeatOpt opts (C' Add) ((op1 :$ a :$ b) :* (op2 :$ c :$ d) :* Nil)
        | Just (C' Add) <- prjF op1
        , Just (C' Sub) <- prjF op2
        , alphaEq a c
        , alphaEq b d
        = constructFeat opts (c' Add) (a :* c :* Nil)

    -- x `mod` y + y * (x `div` y) ==> x
    -- Partial index calculations materialized from contractT . expandT 2
    -- in MultiDim.hs, which is a no-op.
    constructFeatOpt opts (C' Add) ((rem :$ a :$ b) :* (mul :$ c :$ (quot :$ d :$ e)) :* Nil)
        | Just (C' Rem)  <- prjF rem
        , Just (C' Mul)  <- prjF mul
        , Just (C' Quot) <- prjF quot
        , alphaEq a d
        , alphaEq c e
        , alphaEq b e
        = return a

    -- literal a - (b + literal c) ==> literal (a-c) - b
    -- constructFeatOpt opts s@(C' Sub) (a :* (op :$ b :$ c) :* Nil)
    --     | Just a'      <- viewLiteral a
    --     , Just (C' Add) <- prjF op
    --     , Just c'      <- viewLiteral c
    --     = constructFeat opts s (literalDecor (a'-c') :* b :* Nil)

    -- literal a - (b - literal c) ==> literal (a+c) - b
    -- constructFeatOpt opts s@(C' Sub) (a :* (op :$ b :$ c) :* Nil)
    --     | Just a'      <- viewLiteral a
    --     , Just (C' Sub) <- prjF op
    --     , Just c'      <- viewLiteral c
    --     = constructFeat opts s (literalDecor (a'+c') :* b :* Nil)

    -- (a + literal b) - literal c ==> a + literal (b - c)
    constructFeatOpt opts (C' Sub) ((op :$ a :$ b) :* c :* Nil)
        | Just cl         <- viewLiteral c
        , Just s@(C' Add) <- prjF op
        , Just bl         <- viewLiteral b
        = constructFeat opts s (a :* literalDecor (bl-cl) :* Nil)

    -- (a - literal b) - literal c ==> a - literal (b + c)
    constructFeatOpt opts s@(C' Sub) ((op :$ a :$ b) :* c :* Nil)
        | Just cl       <- viewLiteral c
        , Just (C' Sub) <- prjF op
        , Just bl       <- viewLiteral b
        = constructFeat opts s (a :* literalDecor (bl+cl) :* Nil)

    constructFeatOpt opts (C' Sub) ((op1 :$ a :$ b) :* (op2 :$ c :$ d) :* Nil)
        | Just (C' Add) <- prjF op1
        , Just (C' Sub) <- prjF op2
        , alphaEq a c
        , alphaEq b d
        = constructFeat opts (c' Add) (b :* d :* Nil)

    constructFeatOpt _ (C' Sub) (a :* b :* Nil)
        | Just 0 <- viewLiteral b = return a
        | alphaEq a b             = return $ literalDecor 0

    -- (x + yi) * i ==> -y + xi; (x + yi) * (-i) ==> y - xi
    constructFeatOpt opts (C' Mul) (a :* iunit :* Nil)
        | ComplexType FloatType <- infoType (getInfo iunit)
        , Just (0 :+ k)         <- viewLiteral iunit
        , abs k == 1
        = do
             ra <- constructFeat opts (c' RealPart) (a :* Nil)
             ia <- constructFeat opts (c' ImagPart) (a :* Nil)
             iainv <- constructFeatOpt opts (c' Mul) (literalDecor (-k) :* ia :* Nil)
             rainv <- constructFeatOpt opts (c' Mul) (literalDecor k :* ra :* Nil)
             constructFeatOpt opts (c' MkComplex) (iainv :* rainv :* Nil)

    constructFeatOpt _ (C' Mul) (a :* b :* Nil)
        | Just 0 <- viewLiteral a = return a
        | Just 1 <- viewLiteral a = return b
        | Just 0 <- viewLiteral b = return b
        | Just 1 <- viewLiteral b = return a

    constructFeatOpt opts s@(C' Mul) (a :* (op :$ b :$ c) :* Nil)
        | Just al       <- viewLiteral a
        , Just (C' Mul) <- prjF op
        , Just cl       <- viewLiteral c
        = constructFeat opts s (b :* literalDecor (al*cl) :* Nil)

    constructFeatOpt opts s@(C' Mul) ((op :$ a :$ b) :* c :* Nil)
        | Just cl       <- viewLiteral c
        , Just (C' Mul) <- prjF op
        , Just bl       <- viewLiteral b
        = constructFeat opts s (a :* literalDecor (bl*cl) :* Nil)

    constructFeatOpt opts (C' Mul) ((op1 :$ a :$ b) :* (op2 :$ c :$ d) :* Nil)
        | Just (C' Mul) <- prjF op1
        , Just (C' Mul) <- prjF op2
        , Just b'      <- viewLiteral b
        , Just d'      <- viewLiteral d
        = do
            ac <- constructFeat opts (c' Mul) (a :* c :* Nil)
            constructFeat opts (c' Mul) (ac :* literalDecor (b'*d') :* Nil)

    -- Cases to make sure literals end up to the right:
    constructFeatOpt opts (C' Add) (a :* b :* Nil)
        | Just _ <- viewLiteral a = constructFeatUnOpt opts (c' Add) (b :* a :* Nil)

    constructFeatOpt opts (C' Mul) (a :* b :* Nil)
        | Just _ <- viewLiteral a = constructFeatUnOpt opts (c' Mul) (b :* a :* Nil)

    constructFeatOpt opts a args = constructFeatUnOpt opts a args

    constructFeatUnOpt opts x@(C' _) = constructFeatUnOptDefault opts x