{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Graphics.Rendering.Plot.Figure
-- Copyright   :  (c) A. V. H. McPhail 2010
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  portable
--
-- Creation and manipulation of 'Figure's
--
-- The same problem of leaked instances as at <http://hackage.haskell.org/packages/archive/graphviz/2999.10.0.1/doc/html/Data-GraphViz-Commands.html#t%3AGraphvizCanvas> occurs here.
--
--
-- /with/, /set/, /clear/, /new/, and /add/ are the operations that can
-- be performed on various elements of a figure.
--
-- /glib/\//data-accessor/ abstractions (verbs/modifiers) are planned for future implementations

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

module Graphics.Rendering.Plot.Figure (
                                       module Data.Colour.Names
                                       -- * Top level operation
                                      , Figure(), FigureState()
                                       -- * Default options
                                      , withTextDefaults
                                      , withLineDefaults
                                      , withPointDefaults
                                      , withBarDefaults
                                      -- * Figures
                                      , newFigure
                                      -- ** Formatting
                                      , setBackgroundColour
                                      , setFigurePadding
                                      , withTitle
                                      , withSubTitle
                                      , setPlots
                                      , withPlot, withPlots
                                      -- * Sub-plots
                                      ,  Plot()
                                      -- ** Colour
                                      , setPlotBackgroundColour
                                      -- ** Plot elements
                                      , Border
                                      , setBorder
                                      , setPlotPadding
                                      , withHeading
                                      -- ** Series data
                                      , Function(), VectorFunction(), Series(), MinMaxSeries(), ErrorSeries()
                                      , Surface()
                                      , SeriesLabel()
                                      , Abscissa(), Ordinate(), Dataset()
                                      , FormattedSeries(), SeriesType(..)
                                      , line, point, linepoint
                                      , impulse, step
                                      , area
                                      , bar
                                      , hist
                                      , candle, whisker
                                      , setDataset
                                      -- * Annotations
                                      , Location, Head, Fill
                                      , Annote()
                                      , arrow
                                      , oval
                                      , rect
                                      , glyph
                                      , text
                                      , cairo
                                      , withAnnotations
                                      -- ** Plot type
                                      , setSeriesType
                                      , setAllSeriesTypes
                                      -- ** Formatting
                                      , PlotFormats()
                                      , withSeriesFormat
                                      , withAllSeriesFormats
                                      -- * Range
                                      , Scale(..)
                                      , setRange
                                      , setRangeFromData
                                      -- * Axes
                                      , Axis
                                      , AxisType(..),AxisSide(..),AxisPosn(..)
                                      , clearAxes
                                      , clearAxis
                                      , addAxis
                                      , withAxis
                                      -- * BarSetting
                                      , BarSetting(..)
                                      , barSetting
                                      -- * Data Sampling
                                      , SampleData
                                      , sampleData
                                      -- * Legend
                                      , Legend
                                      , LegendBorder
                                      , LegendLocation(..), LegendOrientation(..)
                                      , clearLegend
                                      , setLegend
                                      , withLegendFormat
                                      -- ** Formatting
                                      , Tick(..), TickValues(..), GridLines
                                      , TickFormat(..)
                                      , setTicks
                                      , setGridlines
                                      , setTickLabelFormat
                                      , setTickLabels
                                      , withTickLabelsFormat
                                      , withAxisLabel
                                      , withAxisLine
                                      , withGridLine
                                       -- * Lines
                                      , Line(), LineFormat()
                                      , DashStyle,Dash(..),LineWidth
                                      , clearLineFormat
                                      , setDashStyle
                                      , setLineWidth
                                      , setLineColour
                                      -- * Points
                                      , Point(), PointFormat()
                                      , Glyph(..)
                                      , PointSize
                                      , setGlyph
                                      , setPointSize
                                      , setPointColour
                                      -- * Bars
                                      , Bar(), BarFormat()
                                      , clearBarFormat
                                      , setBarWidth
                                      , setBarColour
                                      , setBarBorderWidth
                                      , setBarBorderColour
                                      -- * Text labels
                                      , Text()
                                      , FontFamily,FontSize,Color
                                      -- | A text element must exist for formatting to work
                                      , clearText
                                      , clearTextFormat
                                      , setText
                                      , setFontFamily
                                      , setFontStyle
                                      , setFontVariant
                                      , setFontWeight
                                      , setFontStretch
                                      , setFontSize
                                      , setFontColour
                                      ) where

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

--import Data.Packed.Vector
--import Numeric.LinearAlgebra.Linear

--import Data.Word
--import Data.Colour.SRGB
import Data.Colour.Names

import qualified Data.Array.IArray as A

--import qualified Graphics.Rendering.Cairo as C
--import qualified Graphics.Rendering.Pango as P

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

import Prelude hiding(min,max)

import Graphics.Rendering.Plot.Figure.Text
import Graphics.Rendering.Plot.Figure.Line
import Graphics.Rendering.Plot.Figure.Point
import Graphics.Rendering.Plot.Figure.Bar
import Graphics.Rendering.Plot.Figure.Plot

import Graphics.Rendering.Plot.Types
import Graphics.Rendering.Plot.Defaults

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

-- | perform some actions on the text defaults, must be run before other text element modifications
withTextDefaults :: Text () -> Figure ()
withTextDefaults m = do
  o <- getDefaults
  let to' = _textoptions o
  let (FontText to _) = execText m to' (FontText to' "")
  modifyDefaults $ \s -> s { _textoptions = to }

-- | perform some actions on the line defaults, must be run before other line element modifications
withLineDefaults :: Line () -> Figure ()
withLineDefaults m = do
  o <- getDefaults
  let lo' = _lineoptions o
  let (TypeLine lo _) = execLine m lo' (TypeLine lo' black)
  modifyDefaults $ \s -> s { _lineoptions = lo }

-- | perform some actions on the point defaults, must be run before other point modifications
withPointDefaults :: Point () -> Figure ()
withPointDefaults m = do
  o <- getDefaults
  let po' = _pointoptions o
  let (FullPoint po _) = execPoint m po' (FullPoint po' defaultGlyph)
  modifyDefaults $ \s -> s { _pointoptions = po }

-- | perform some actions on the bar defaults, must be run before other point modifications
withBarDefaults :: Bar () -> Figure ()
withBarDefaults m = do
  o <- getDefaults
  let bo' = _baroptions o
  let (TypeBar bo _) = execBar m bo' (TypeBar bo' black)
  modifyDefaults $ \s -> s { _baroptions = bo }

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

-- | create a new blank 'Figure'
newFigure :: Figure ()
newFigure = putFigure $ Figure defaultFigureBackgroundColour
                          defaultFigurePadding NoText NoText
                          (A.listArray ((1,1),(1,1)) [Nothing])
{-
newLineFigure :: DataSeries                      -- ^ the y series
              -> FigureData
newLineFigure d@(DS_1toN _ _) = let ((xmin,xmax),(ymin,ymax)) = calculateRanges d
                                    plot = Plot False defaultPlotPadding NoText
                                           (defaultRanges xmin xmax ymin ymax)
                                           [defaultXAxis,defaultYAxis]
                                           Nothing Line d []
                                in Figure defaultFigurePadding NoText NoText
                                       (A.listArray ((1,1),(1,1)) [Just plot])
-}
{-
-- | create a new 'Figure'
newFigure :: PlotType -> DataSeries -> Figure ()
newFigure Line   d@(DS_1toN _ _) = putFigure $ newLineFigure d
--newFigure _      _               = error "Figure type not implemented"
-}

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

-- | set the background colour of the figure
setBackgroundColour :: Color -> Figure ()
setBackgroundColour c = modifyFigure $ \s -> s { _back_clr = c }

-- | set the padding of the figure
setFigurePadding :: Double -> Double -> Double -> Double -> Figure ()
setFigurePadding l r b t = modifyFigure $ \s ->
                             s { _fig_pads = Padding l r b t }

-- | operate on the title
withTitle :: Text () -> Figure ()
withTitle m = do
  o <- getDefaults
  modifyFigure $ \s ->
    s { _title = execText m (_textoptions o) (_title s) }

-- | operate on the sub-title
withSubTitle :: Text () -> Figure ()
withSubTitle m = do
  o <- getDefaults
  modifyFigure $ \s ->
    s { _subtitle = execText m (_textoptions o) (_title s) }

-- | set the shape of the plots, losing all current plots
setPlots :: Int      -- ^ rows
         -> Int      -- ^ columns
         -> Figure ()
setPlots r c = modifyFigure $ \s ->
                 s { _plots = A.listArray ((1,1),(r,c))
                      (replicate (r*c) Nothing) }

-- | perform some actions on the specified subplot
withPlot :: (Int,Int) -> Plot () -> Figure ()
withPlot i m = do
  o <- getDefaults
  s <- getSupplies
  modifyFigure $ \p ->
    p { _plots = let plots = _plots p
                     plot' = plots A.! i
                     plot = case plot' of
                              Nothing -> emptyPlot
                              Just p' -> p'
                                                    -- we revert supplies to the original here
                                                    -- since we might want the same colour
                                                    -- order for all plots
                                                    -- HOWEVER: need a better execPlot group
                 in plots A.// [(i,Just $ execPlot m s o plot)] }

-- | perform some actions all subplots
withPlots :: Plot () -> Figure ()
withPlots m = do
  o <- getDefaults
  s <- getSupplies
  modifyFigure $ \p ->
    p { _plots = let plots = _plots p
                     plot p' = case p' of
                                 Nothing  -> emptyPlot
                                 Just p'' -> p''
                 in plots A.// map (\(i,e) ->
                          (i,Just $ execPlot m s o (plot e))) (A.assocs plots) }

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


{-# LANGUAGE BangPatterns, LambdaCase #-}

module Input
   ( Options(..)
   , FractalType(..)
   , AA(..)
   , Normalization(..)
   , Color(..)
   , Animation(..)
   , processArgs
   , pPrintOptions
   , defaultOptions
   )
where

import           Data.Complex
import           System.IO
import           Data.Word
import           System.Environment
import           Data.Char
import qualified Graphics.Image                as I

data FractalType =
   Mandelbrot | Julia deriving (Show)
data AA
   = AAEnabled | AADisabled deriving (Show)
data Normalization =
   Linear | Sigmoid Double Double | Periodic Double | Sine Double deriving (Show)
data Color =
   Greyscale | Hue | Gradient Bool [I.Pixel I.RGB Double] deriving (Show)
data Animation =
   NoAnimation | Power Double | Zoom Double (Maybe Int) |
   Iterations Int | Theta Double (Maybe Double) | LinearC (Complex Double) |
   Grid Double Double Double Double Int Int deriving (Show)
data Options = Options
   { fractalType     :: FractalType
   , resolution      :: (Int, Int)
   , center          :: Complex Double
   , range           :: Double
   , iterations      :: Int
   , power           :: Double
   , aa              :: AA
   , normalization   :: Normalization
   , color           :: Color
   , animation       :: Animation
   , cvalue          :: Complex Double
   , setColor        :: I.Pixel I.RGB Double
   , frameCount      :: Int
   , startingFrame   :: Int
   } deriving (Show)

defaultOptions = Options Mandelbrot
                         (500, 500)
                         (0 :+ 0)
                         2
                         500
                         2
                         AAEnabled
                         Linear
                         Hue
                         NoAnimation
                         (0 :+ 0)
                         (I.PixelRGB 0 0 0)
                         100
                         1

-- Primary function that produces a set of options from a ist of arguments
processArgs :: [String] -> Options
processArgs = getOptionsFromArgs . separateArgs . normalizeArgs

-- Makes sure all arguments are lowercase
normalizeArgs :: [String] -> [String]
normalizeArgs = map $ map toLower

-- Group into arguments and their subarguments
separateArgs :: [String] -> [[String]]
separateArgs []       = []
separateArgs (x : xs) = case lookup (take 3 x) arguments of
   Nothing -> error $ "Invalid argument \"" ++ x ++ "\""
   otherwise ->
      let (subArgs, args) = span
             (\y -> case lookup (take 3 y) arguments of
                Nothing   -> True
                otherwise -> False
             )
             xs
      in  (x : subArgs) : separateArgs args

-- Lookup and process all arguments
getOptionsFromArgs :: [[String]] -> Options
getOptionsFromArgs = foldl
   (\opt (arg : subargs) ->
      let func = lookup (take 3 arg) arguments
      in  case fmap (\f -> f (subargs, opt)) func of
             Just x    -> x
             otherwise -> error ""
   )
   defaultOptions

-- Function to convert a string like "23,4,198" to a color
colorFromString :: String -> I.Pixel I.RGB Double
colorFromString s =
   case map ((/ 255) . read) . words $ map (\case ',' -> ' '; x -> x) s
     of (r : g : b : []) -> if r <= 1 && g <= 1 && b <= 1
           then I.PixelRGB r g b
           else error "All color values must be on the interval [0,255]"
        otherwise -> error "Invalid color format for gradient"

-- Array of all command line arguments and functions to interpret their subarguments
-- It is assumed that arguments have been converted to lowercase and shortened to 3 characters
arguments :: [(String, (([String], Options) -> Options))]
arguments =
   [ ( "-fr"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("man" : []) -> opt { fractalType = Mandelbrot }
        ("jul" : []) -> opt { fractalType = Julia }
        otherwise    -> error $ "Invalid arguments for -fractalType: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-re"
     , \case
        (c : r : [], opt) -> opt { resolution = (read r, read c) }
        x                 -> error $ "Invalid arguments for -resolution: \"" ++ (show x) ++ "\""
     )
   , ( "-ce"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("rec" : r : i : []) -> opt { center = read r :+ read i }
        ("pol" : m : p : []) -> opt { center = mkPolar (read m) ((/ 180) . (* pi) $ read p) }
        otherwise            -> error $ "Invalid arguments for -center: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-ra"
     , \case
        (n : [], opt) -> opt { range = read n }
        x             -> error $ "Invalid arguments for -range: \"" ++ (show x) ++ "\""
     )
   , ( "-it"
     , \case
        (n : [], opt) -> opt { iterations = read n }
        x             -> error $ "Invalid arguments for -iterations: \"" ++ (show x) ++ "\""
     )
   , ( "-po"
     , \case
        (n : [], opt) -> opt { power = read n }
        x             -> error $ "Invalid arguments for -power: \"" ++ (show x) ++ "\""
     )
   , ( "-aa"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("ena" : []) -> opt { aa = AAEnabled }
        ("dis" : []) -> opt { aa = AADisabled }
        otherwise    -> error $ "Invalid arguments for -aa: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-no"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("lin"         : []) -> opt { normalization = Linear }
        ("sig"     : c : []) -> opt { normalization = Sigmoid (read c) 2.5 }
        ("sig" : c : p : []) -> opt { normalization = Sigmoid (read c) (read p) }
        ("per"     : p : []) -> opt { normalization = Periodic (read p) }
        ("sin"     : p : []) -> opt { normalization = Sine (read p) }
        otherwise -> error $ "Invalid arguments for -normalization: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-co"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("gre"             : []  ) -> opt { color = Greyscale }
        ("hue"             : []  ) -> opt { color = Hue }
        ("gra" : ('l' : _) : args) -> if length args < 2
           then error "Insufficient colors for a gradient to be formed"
           else opt { color = (Gradient False $ map colorFromString args) }
        ("gra" : args) -> if length args < 2
           then error "Insufficient colors for a gradient to be formed"
           else opt { color = (Gradient True $ map colorFromString args) }
        otherwise -> error $ "Invalid arguments for -color: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-an"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("non"           : []) -> opt { animation = NoAnimation }
        ("pow" : fn : fr : []) -> opt { animation = Power (read fn), frameCount = (read fr) }
        ("zoo" : fn : fr : []) -> opt { animation = Zoom (read fn) Nothing, frameCount = (read fr) }
        ("zoo" : fn : fr : fi : []) ->
           opt { animation = Zoom (read fn) (Just $ read fi), frameCount = (read fr) }
        ("ite" : fn : fr : []) -> opt { animation = Iterations (read fn), frameCount = (read fr) }
        ("the" : fn : fr : []) -> opt { animation = Theta (read fn) Nothing, frameCount = (read fr) }
        ("the" : fn : fr : fm : []) ->
           opt { animation = Theta (read fn) (Just $ read fm), frameCount = (read fr) }
        ("lin" : ty : xs) -> case ((take 3 ty) : xs) of
           ("rec" : r : i : fr : []) ->
              opt { animation = LinearC (read r :+ read i), frameCount = (read fr) }
           ("pol" : m : p : fr : []) ->
              opt { animation = LinearC (mkPolar (read m) ((/ 180) . (* pi) $ read p))
                  , frameCount = (read fr) }
        ("gri" : ri : rl : ii : il : rc : ic : []) ->
            opt { animation = Grid (read ri) (read rl) (read ii) (read il) (read rc) (read ic)
                , frameCount = (read rc) * (read ic) }
        otherwise -> error $ "Invalid arguments for -animation: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-cv"
     , \(x : xs, opt) -> case (take 3 x) : xs of
        ("rec" : r : i : []) -> opt { cvalue = read r :+ read i }
        ("pol" : m : p : []) -> opt { cvalue = mkPolar (read m) ((/ 180) . (* pi) $ read p) }
        otherwise            -> error $ "Invalid arguments for -cvalue: \"" ++ (show $ x : xs) ++ "\""
     )
   , ( "-se"
     , \case
        ([x], opt) -> opt { setColor = colorFromString x }
        x          -> error $ "Invalid arguments for -setColor: \"" ++ (show x) ++ "\""
     )
   , ( "-st"
     , \case
        ([x], opt) -> opt { startingFrame = read x }
        x          -> error $ "Invalid arguments for -startingFrame: \"" ++ (show x) ++ "\""
     )
   ]

pPrintOptions :: Options -> String
pPrintOptions options =
   "Fractal Type:    " ++ (show $ fractalType   options) ++ '\n' :
   "Resolution:      " ++ (show $ resolution    options) ++ '\n' :
   "Center:          " ++ (show $ center        options) ++ '\n' :
   "Range:           " ++ (show $ range         options) ++ '\n' :
   "Iterations:      " ++ (show $ iterations    options) ++ '\n' :
   "Power:           " ++ (show $ power         options) ++ '\n' :
   "Antialiasing:    " ++ (show $ aa            options) ++ '\n' :
   "Normalization:   " ++ (show $ normalization options) ++ '\n' :
   "Color:           " ++ (show $ color         options) ++ '\n' :
   "Animation:       " ++ (show $ animation     options) ++ '\n' :
   "C-Value:         " ++ (show $ cvalue        options) ++ '\n' :
   "Set Color:       " ++ (show $ setColor      options) ++ '\n' :
   "Frames:          " ++ (show $ frameCount    options) ++ '\n' :
   "Starting Frame:  " ++ (show $ startingFrame options)


-- File created: 2009-07-11 20:29:49

module Haschoo.Parser (runParser, programValue, value, number) where

import Control.Applicative                  ((<$>))
import Control.Arrow                        (first)
import Control.Monad                        (liftM2)
import Data.Char                            (digitToInt, toLower)
import Data.Complex                         (Complex((:+)), mkPolar)
import Data.Maybe                           (fromJust, fromMaybe, isJust)
import Data.Ratio                           ((%))
import Numeric                              (readInt)
import Text.ParserCombinators.Parsec hiding (runParser)

import Haschoo.Types (ScmValue(..), toScmString, toScmVector)
import Haschoo.Utils (void, (.:))

runParser :: Parser a -> SourceName -> String -> Either String a
runParser = (either (Left . show) Right .:) . parse

programValue :: Maybe SourcePos -> Parser (Maybe (ScmValue, String, SourcePos))
programValue pos = do
   maybe (return ()) setPosition pos
   optional atmosphere
   v <- (Just <$> value) <|> (eof >> return Nothing)
   case v of
        Nothing -> return Nothing
        Just v' -> liftM2 (Just .: (,,) v') getInput getPosition

values :: Parser [ScmValue]
values = optional atmosphere >> (many $ value `discard` atmosphere)

value :: Parser ScmValue
value = do
   quotes <-
      concat <$> many
         (choice [ string "'"
                 , string "`"
                 , liftM2 (:) (char ',') (string "@" <|> return "") ]
             `discard` atmosphere)

   val <- choice [ ident
                 , list
                 , vector
                 , bool
                 , number 10
                 , character
                 , quotedString ]

   return $ quote quotes val
 where
   quote []            = id
   quote ('\''    :qs) = quoteWith "quote"            qs
   quote ('`'     :qs) = quoteWith "quasiquote"       qs
   quote (',' :'@':qs) = quoteWith "unquote-splicing" qs
   quote (','     :qs) = quoteWith "unquote"          qs
   quote _             = error "Parser.quote :: the impossible happened"

   quoteWith s qs = ScmList . (ScmIdentifier s :) . (:[]) . quote qs

ident :: Parser ScmValue
ident = do
   -- peculiar needs the try due to negative numbers
   x <- try $ choice [ peculiar `discard` try delimiter
                     , ordinary `discard` delimiter ]
   return (ScmIdentifier $ map toLower x)
 where
   peculiar = choice [return <$> oneOf "+-", string "..."]
   ordinary = do
      x  <- oneOf initial
      xs <- many (oneOf (initial ++ "+-.@" ++ ['0'..'9']))
      return (x:xs)
    where
      initial = ['a'..'z'] ++ ['A'..'Z'] ++ "!$%&*/:<=>?^_~"

bool :: Parser ScmValue
bool = ScmBool . (`elem` "tT") <$> try (char '#' >> oneOf "ftFT")

character :: Parser ScmValue
character = do
   try (string "#\\")
   c <- try named <|> anyChar
   delimiter
   return (ScmChar c)
 where
   named = choice [ string "space"   >> return ' '
                  , string "newline" >> return '\n' ]

quotedString :: Parser ScmValue
quotedString =
   fmap toScmString
      . between (try $ char '"') (char '"')
      . many $
          (char '\\' >> (oneOf "\\\"" <?> concat [show "\\"," or ",show "\""]))
      <|> satisfy (/= '"')

list :: Parser ScmValue
list =
   between (try $ char '(') (optional atmosphere >> char ')') $ do
      vals <- values
      if null vals
         then return$ ScmList vals
         else do
            optional atmosphere
            dot <- optionMaybe (char '.')
            if isJust dot
               then do
                  atmosphere
                  end <- value
                  return$ case end of
                               ScmList l -> ScmList (vals ++ l)
                               _         -> mkDotted vals end
               else return$ ScmList vals
 where
   -- Flatten (1 . (2 . 3)) to (1 2 . 3)
   mkDotted xs (ScmDottedList ys z) = mkDotted (xs ++ ys) z
   mkDotted xs x                    = ScmDottedList xs x

vector :: Parser ScmValue
vector = do
   try (char '#' >> char '(')
   toScmVector <$> values `discard` (optional atmosphere >> char ')')

number :: Int -> Parser ScmValue
number defRadix = do
   (radix,exact) <- try prefix

   let p = if isJust radix || isJust exact then id else try
   n <- p $ complex (fromMaybe False exact) (fromMaybe defRadix radix)

   delimiter

   return$ if fromMaybe True exact
              then n
              else case n of
                        ScmInt x -> ScmReal (fromInteger  x)
                        ScmRat x -> ScmReal (fromRational x)
                        _        -> n
 where
   prefix = do
      r <- radix
      e <- exactness
      -- They can be in either order
      flip (,) e <$> if isJust r
                        then return r
                        else radix
    where
      radix     = f [('B',2), ('O',8), ('D',10), ('X',16)]
      exactness = f [('E',True), ('I',False)]
      f xs = optionMaybe . try $ do
         char '#'
         let xs' = map (first toLower) xs ++ xs
         fromJust . (`lookup` xs') <$> oneOf (map fst xs')

   complex exact radix =
      choice [ try $ do n <- sreal False radix
                        ncChar 'i'
                        return (mkImaginary n)

             , try $ do a  <- real exact radix
                        at <- optionMaybe (char '@')
                        if isJust at
                           then mkComplex mkPolar a <$> real False radix
                           else do
                              b <- optionMaybe $ try imaginaryUnit <|>
                                      do n <- sreal False radix
                                         ncChar 'i'
                                         return n

                              return $ case b of
                                            Nothing -> a
                                            Just n  -> mkComplex (:+) a n

             , mkImaginary <$> imaginaryUnit ]
    where
      imaginaryUnit = do
         neg <- sign
         ncChar 'i'
         return (ScmInt neg)

      mkImaginary     = mkComplex (:+) (ScmInt 0)
      mkComplex f a b = ScmComplex $ f (toDouble a) (toDouble b)

   toDouble (ScmInt  i) = fromInteger  i
   toDouble (ScmRat  r) = fromRational r
   toDouble (ScmReal r) = r
   toDouble _           = error "number.toDouble :: internal error"

   real exact radix = do
      neg <- optionMaybe sign
      applySign (fromMaybe 1 neg) <$> ureal exact radix

   sreal exact radix = do
      neg <- sign
      applySign neg <$> ureal exact radix

   applySign neg (ScmInt  n) = ScmInt  (fromIntegral neg * n)
   applySign neg (ScmRat  n) = ScmRat  (fromIntegral neg * n)
   applySign neg (ScmReal n) = ScmReal (fromIntegral neg * n)
   applySign _   _           = error "number.applySign :: іnternal error"

   ureal exact radix = choice [ string "nan.#" >> return (ScmReal $ 0/0)
                              , string "inf.#" >> return (ScmReal $ 1/0)
                              , decimal exact radix
                              , do a <- uint radix
                                   b <- optionMaybe $ char '/' >> uint radix
                                   case b of
                                        Nothing ->
                                           if radix == 10
                                              then tryExponent exact a
                                              else return a
                                        Just n  -> return (mkRatio a n) ]

   mkRatio (ScmInt a) (ScmInt b) = ScmRat (a % b)
   mkRatio _          _          = error "number.mkRatio :: internal error"

   decimal exact radix | radix /= 10 = fail "Decimal outside radix 10"
                       | otherwise   =
      tryExponent exact =<<
         (try . choice) [ do char '.'
                             n <- many1 (digitN 10)
                             skipMany (char '#')
                             return $ readDecimal exact "0" n

                        , do a <- many1 (digitN 10)
                             char '.'
                             b <- many (digitN 10)
                             skipMany (char '#')
                             -- read doesn't like "123." so add a 0 if
                             -- necessary
                             return $
                                readDecimal exact a (if null b then "0" else b)

                        , do n <- many1 (digitN 10)
                             hashes  <- map (const '0') <$> many1 (char '#')
                             char '.'
                             hashes2 <- many (char '#')
                             return . inexactHashes (hashes ++ hashes2) .
                                ScmInt $ readInteger 10 (n ++ hashes)
                        ]

   tryExponent exact n = do
      ex <- optionMaybe $ do
               oneOf "esfdlESFDL" -- Ignore the exponent: all Double
               neg <- optionMaybe sign
               xs  <- many1 (digitN 10)
               return$ fromMaybe 1 neg * readInteger 10 xs

      return$
         case ex of
              Nothing -> n
              Just e | not exact -> ScmReal (10^^e * toDouble n)
                     | otherwise ->
                        case n of
                             ScmInt x -> ScmRat (10^^e * fromInteger x)
                             ScmRat x -> ScmRat (10^^e * x)
                             _        ->
                                error$ "Parser.tryExponent :: "
                                    ++ "the impossible happened"

   uint radix = do
      n <- many1 (digitN radix)
      hashes <- map (const '0') <$> many (char '#')
      return . inexactHashes hashes . ScmInt $ readInteger radix (n ++ hashes)

   -- If any # were present, the value is inexact (R5RS 6.2.4)
   inexactHashes :: String -> ScmValue -> ScmValue
   inexactHashes [] = id
   inexactHashes _  = ScmReal . toDouble

   digitN :: Int -> Parser Char
   digitN 2  = oneOf "01"
   digitN 8  = octDigit
   digitN 10 = digit
   digitN 16 = hexDigit
   digitN _  = error "number.digitN :: internal error"

   sign = (\c -> if c == '+' then 1 else -1) <$> oneOf "+-"

   -- These read functions all assume correctly formed input

   readInteger :: Int -> String -> Integer
   readInteger radix =
      fst.head . readInt (fromIntegral radix) (const True) digitToInt

   readDecimal :: Bool -> String -> String -> ScmValue
   readDecimal False as bs = ScmReal . read . concat $ [as, ".", bs]
   readDecimal True  as bs =
      ScmRat $ fromInteger (readInteger 10 as)
             + fromInteger (readInteger 10 bs) / 10 ^ length bs

-- Pushes back anything relevant for other parsers
delimiter :: Parser ()
delimiter = choice [ whitespaceOrComment
                   , void . lookAhead . try $ oneOf "()\""
                   , try eof ]

atmosphere :: Parser ()
atmosphere = skipMany whitespaceOrComment

whitespaceOrComment :: Parser ()
whitespaceOrComment =
   choice [ void newline
          , char '\r' >> optional newline
          , void space
          , char ';' >> skipMany (noneOf "\n\r") ]

ncChar :: Char -> Parser Char
ncChar c = toLower <$> satisfy ((== c) . toLower)

discard :: Parser a -> Parser b -> Parser a
discard a b = a >>= \a' -> b >> return a'


-- Copyright 2016 TensorFlow authors.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

-- | This module contains definitions for some built-in TensorFlow operations.
--
-- Note that certain, "stateful" ops like 'variable' and 'assign' return a
-- 'Build' action (e.g., @Build (Tensor Ref a)@ instead of a pure value; the
-- returned 'Tensor's are always rendered in the current 'Build' context.  This
-- approach helps us avoid problems with inlining or common subexpression
-- elimination, by writing
--
-- > do
-- >     v <- variable []
-- >     w <- assign v 3
-- >     render $ w * w
--
-- instead of
--
-- > let
-- >    v = variable []
-- >    w = assign v 3
-- > in w * w
--
-- since the latter could be reasonably transformed by the compiler into (or
-- vice versa)
--
-- > let
-- >    v = variable []
-- >    w = assign v 3
-- >    w' = assign v 3
-- > in w * w'
--
-- Ops should return a 'Build' action if their original 'OpDef' marks them as
-- stateful, or if they take any Refs as input.  (This mirrors the rules that
-- TensorFlow uses to avoid common subexpression elimination.)
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module TensorFlow.Ops
    ( CoreOps.add
    , CoreOps.add'
    , CoreOps.abs
    , CoreOps.abs'
    , CoreOps.addN
    , CoreOps.addN'
    , CoreOps.argMax
    , CoreOps.argMax'
    , CoreOps.assign
    , CoreOps.assign'
    , CoreOps.broadcastGradientArgs
    , CoreOps.broadcastGradientArgs'
    , CoreOps.cast
    , CoreOps.cast'
    , CoreOps.concat
    , CoreOps.concat'
    , constant
    , constant'
    , CoreOps.equal
    , CoreOps.equal'
    , expandDims
    , expandDims'
    , initializedVariable
    , initializedVariable'
    , zeroInitializedVariable
    , zeroInitializedVariable'
    , CoreOps.fill
    , CoreOps.fill'
    , CoreOps.identity
    , CoreOps.identity'
    , CoreOps.matMul
    , CoreOps.matMul'
    , matTranspose
    , matTranspose'
    , CoreOps.mean
    , CoreOps.mean'
    , CoreOps.mul
    , CoreOps.mul'
    , CoreOps.neg
    , CoreOps.neg'
    , CoreOps.oneHot
    , CoreOps.oneHot'
    , CoreOps.pack
    , CoreOps.pack'
    , placeholder
    , placeholder'
    , CoreOps.range
    , CoreOps.range'
    , reducedShape
    , reduceMean
    , reduceMean'
    , CoreOps.relu
    , CoreOps.relu'
    , CoreOps.reluGrad
    , CoreOps.reluGrad'
    , CoreOps.reshape
    , CoreOps.reshape'
    , restore
    , restoreFromName
    , save
    , scalar
    , scalar'
    , shape
    , shape'
    , CoreOps.sign
    , CoreOps.sign'
    , CoreOps.size
    , CoreOps.size'
    , CoreOps.softmax
    , CoreOps.softmax'
    , CoreOps.softmaxCrossEntropyWithLogits
    , CoreOps.softmaxCrossEntropyWithLogits'
    , CoreOps.sparseToDense
    , CoreOps.sparseToDense'
    , CoreOps.sub
    , CoreOps.sub'
    , CoreOps.sum
    , CoreOps.sum'
    , reduceSum
    , reduceSum'
    , CoreOps.transpose
    , CoreOps.transpose'
    , truncatedNormal
    , truncatedNormal'
    , CoreOps.variable
    , CoreOps.variable'
    , vector
    , vector'
    , zeros
    , CoreOps.zerosLike
    , CoreOps.zerosLike'
    , scalarize
    ) where

import Data.ByteString (ByteString)
import Data.Complex (Complex)
import Data.Int (Int32, Int64)
import Data.Word (Word16)
import Prelude hiding (abs, sum, concat)
import Data.ProtoLens (def)
import Data.Text.Encoding (encodeUtf8)
import Lens.Family2 ((.~), (&))
import Text.Printf (printf)
import Proto.Tensorflow.Core.Framework.Tensor
    ( TensorProto
    , dtype
    , tensorShape
    )
import qualified Proto.Tensorflow.Core.Framework.TensorShape
  as TensorShape
import TensorFlow.Build
import TensorFlow.BuildOp
import TensorFlow.ControlFlow (group)
import TensorFlow.Tensor
import TensorFlow.Types

import qualified TensorFlow.GenOps.Core as CoreOps

import qualified Prelude (abs)

-- TODO: Look into hs-boot refactoring to allow mutually recursive imports.
-- | Must be defined as an orphan because of the dependency order between Ops
-- and Tensor.
--
-- The indirect constraint "v ~ Value" helps disambiguate types, for example in
-- "neg 1 :: Tensor Value Float", it helps find the type of the subexpression
-- "1".
instance ( TensorType a
         , Num a
         , v ~ Build
         , OneOf '[ Double, Float, Int32, Int64
                  , Complex Float, Complex Double] a) => Num (Tensor v a) where
    (+) = CoreOps.add
    (*) = CoreOps.mul
    (-) = CoreOps.sub
    abs = CoreOps.abs
    fromInteger = scalar . fromInteger
    signum = CoreOps.sign
    negate = CoreOps.neg

matTranspose :: TensorType a => Tensor e a -> Tensor Build a
matTranspose = matTranspose' id

matTranspose' :: TensorType a => OpParams -> Tensor v a -> Tensor Build a
matTranspose' params = flip (CoreOps.transpose' params) (vector [1, 0 :: Int32])

placeholder :: (MonadBuild m, TensorType a) => Shape -> m (Tensor Value a)
placeholder = placeholder' id

placeholder' :: forall m a . (MonadBuild m, TensorType a)
             => OpParams -> Shape -> m (Tensor Value a)
placeholder' params pShape
    -- Note: we don't use CoreOps.placeholder' since that op isn't stateful,
    -- and thus would be CSE'd.
    = build $ buildOp [] $ opDef "Placeholder"
                & opAttr "dtype" .~ tensorType (undefined :: a)
                & opAttr "shape" .~ pShape
                & params

-- | Creates a variable initialized to the given value.
-- Initialization happens next time session runs.
initializedVariable :: (MonadBuild m, TensorType a)
                    => Tensor v a -> m (Tensor Ref a)
initializedVariable = initializedVariable' id

initializedVariable' :: (MonadBuild m, TensorType a)
                    => OpParams -> Tensor v a -> m (Tensor Ref a)
initializedVariable' params initializer = do
    v <- CoreOps.variable' params []  -- The shape is not known initially.
    i <- CoreOps.assign' (opAttr "validate_shape" .~ False) v
                            initializer
    addInitializer =<< group i
    return v

-- | Creates a zero-initialized variable with the given shape.
zeroInitializedVariable
  :: (MonadBuild m, TensorType a, Num a) =>
     TensorFlow.Types.Shape -> m (Tensor TensorFlow.Tensor.Ref a)
zeroInitializedVariable = zeroInitializedVariable' id

zeroInitializedVariable'
  :: (MonadBuild m, TensorType a, Num a) =>
     OpParams -> TensorFlow.Types.Shape -> m (Tensor TensorFlow.Tensor.Ref a)
zeroInitializedVariable' params = initializedVariable' params . zeros

-- TODO: Support heterogeneous list of tensors.
save :: forall a m v . (Rendered (Tensor v), MonadBuild m, TensorType a)
        => ByteString    -- ^ File path.
        -> [Tensor v a]  -- ^ Tensors to save.
        -> m ControlNode
save path xs = build $ do
    let toByteStringTensor = scalar . encodeUtf8 . encodeOutput . renderedOutput
    let names = fmap toByteStringTensor xs
    let types = replicate (length xs) (tensorType (undefined :: a))
    names' <- buildInputs $ CoreOps.pack names
    xs' <- buildInputs xs
    path' <- buildInputs $ scalar path
    buildOp [] $ opDef "Save"
                    & opAttr "T" .~ types
                    & opInputs .~ (path' ++ names' ++ xs')

-- | Restore a tensor's value from a checkpoint file.
--
-- This version allows restoring from a checkpoint file that uses a different
-- tensor name than the variable.
restoreFromName :: forall a m . (MonadBuild m, TensorType a)
                => ByteString    -- ^ File path.
                -> ByteString    -- ^ Tensor name override.
                -> Tensor Ref a  -- ^ Tensor to restore.
                -> m ControlNode
restoreFromName path name x = build $ do
    path' <- buildInputs $ scalar path
    name' <- buildInputs $ scalar name
    restoreOp <- buildOp [] $ opDef "Restore"
                               & opAttr "dt" .~ tensorType (undefined :: a)
                               & opInputs .~ (path' ++ name')
    group =<< CoreOps.assign x (restoreOp :: Tensor Value a)

-- | Restore a tensor's value from a checkpoint file.
restore :: forall a m . (MonadBuild m, TensorType a)
        => ByteString    -- ^ File path.
        -> Tensor Ref a  -- ^ Tensor to restore.
        -> m ControlNode
restore path x = restoreFromName path name x
  where
    name = encodeUtf8 $ encodeOutput $ renderedOutput x

-- | Create a constant tensor.
--
-- The values should be in row major order, e.g.,
--
--   element 0:   index (0, ..., 0)
--   element 1:   index (0, ..., 1)
--   ...
constant :: TensorType a => Shape -> [a] -> Tensor Build a
constant = constant' id

constant' :: forall a . TensorType a => OpParams -> Shape -> [a] -> Tensor Build a
constant' params (Shape cShape) values
    | invalidLength = error invalidLengthMsg
    | otherwise = CoreOps.const' (params . (opAttr "value" .~ typedNode))
  where
    invalidLength = product cShape /= fromIntegral (length values)
    invalidLengthMsg = printf "invalid tensor length: expected %d got %d"
                              (product cShape)
                              (length values)
    typedNode :: TensorProto
    typedNode = def
                & dtype .~ tensorType (undefined :: a)
                & tensorShape.TensorShape.dim .~
                      [def & TensorShape.size .~ x | x <- cShape]
                & tensorVal .~ values

-- | Reshape a N-D tensor down to a scalar.
--
-- See `TensorFlow.GenOps.Core.reshape`.
scalarize :: TensorType a => Tensor v a -> Tensor Build a
scalarize t = CoreOps.reshape t (vector scalarShape)
    where
        scalarShape = [] :: [Int32]

-- | Sum a tensor down to a scalar
-- Seee `TensorFlow.GenOps.Core.sum`
reduceSum :: (OneOf '[ Double, Float, Int32, Int64
                     , Complex Float, Complex Double] a) =>
             Tensor v a -> Tensor Build a
reduceSum x = CoreOps.sum x allAxes
  where allAxes = CoreOps.range 0 (CoreOps.rank x :: Tensor Build Int32) 1

reduceSum' :: (OneOf '[ Double, Float, Int32, Int64
                      , Complex Float, Complex Double] a) =>
              OpParams -> Tensor v a -> Tensor Build a
reduceSum' params x = CoreOps.sum' params x allAxes
  where allAxes = CoreOps.range 0 (CoreOps.rank x :: Tensor Build Int32) 1

-- | Computes the mean of elements across dimensions of a tensor.
-- See `TensorFlow.GenOps.Core.mean`
reduceMean
  :: ( TensorType a
     , OneOf '[ Double, Float, Complex Float, Complex Double] a
     )
  => Tensor v a -> Tensor Build a
reduceMean = reduceMean' id

reduceMean'
  :: ( TensorType a
     , OneOf '[ Double, Float, Complex Float, Complex Double] a
     )
  => OpParams -> Tensor v a -> Tensor Build a
reduceMean' params x = CoreOps.mean' params x allAxes
  where allAxes = CoreOps.range 0 (CoreOps.rank x :: Tensor Build Int32) 1

-- | Create a constant vector.
vector :: TensorType a => [a] -> Tensor Build a
vector = vector' id

vector' :: TensorType a => OpParams -> [a] -> Tensor Build a
vector' params xs = constant' params [fromIntegral $ length xs] xs

-- | Create a constant scalar.
scalar :: TensorType a => a -> Tensor Build a
scalar = scalar' id

scalar' :: TensorType a => OpParams -> a -> Tensor Build a
scalar' params x = constant' params [] [x]

-- | Random tensor from the unit normal distribution with bounded values.
--
-- This is a type-restricted version of 'TensorFlow.GenOps.Core.truncatedNormal'.
truncatedNormal :: (MonadBuild m, OneOf '[Word16, Double, Float] a)
                => Tensor v Int64  -- ^ Shape.
                -> m (Tensor Value a)
truncatedNormal = CoreOps.truncatedNormal

truncatedNormal' :: (MonadBuild m, OneOf '[Word16, Double, Float] a)
                => OpParams -> Tensor v Int64  -- ^ Shape.
                -> m (Tensor Value a)
truncatedNormal' = CoreOps.truncatedNormal'

zeros :: forall a . (Num a, TensorType a) => Shape -> Tensor Build a
zeros (Shape s) = CoreOps.fill (vector s) (scalar 0)

shape :: TensorType t => Tensor v t -> Tensor Build Int32
shape = CoreOps.shape

shape' :: TensorType t => OpParams -> Tensor v t -> Tensor Build Int32
shape' = CoreOps.shape'

expandDims :: TensorType t => Tensor v1 t -> Tensor v2 Int32 -> Tensor Build t
expandDims = CoreOps.expandDims

expandDims' :: TensorType t => OpParams -> Tensor v1 t -> Tensor v2 Int32 -> Tensor Build t
expandDims' = CoreOps.expandDims'

-- | Helper function for reduction ops (translation of math_ops.reduced_shape).
reducedShape :: (OneOf '[ Int32, Int64 ] t1, OneOf '[ Int32, Int64 ] t2) =>
                Tensor v1 t1 -> Tensor v2 t2 -> Tensor Build Int32
reducedShape inputShape axes =
    let inputShape32 = toInt32 inputShape         -- [2, 3, 5, 7]
        axes32 = toInt32 axes                     -- [1, 2]
        toInt32 x = CoreOps.cast x :: Tensor Build Int32
        inputRank = CoreOps.size inputShape32     -- 4
        axesMod = (axes32 + inputRank) `CoreOps.mod` inputRank
        axesShape = shape axesMod                 -- [2]
    in CoreOps.dynamicStitch                      -- [2, 1, 1, 7]
         [CoreOps.range 0 inputRank 1,            -- [0, 1, 2, 3]
           axesMod]                               -- [1, 2]
         [inputShape32,                           -- [2, 3, 5, 7]
           CoreOps.fill axesShape 1]              -- [1, 1]