{-# LANGUAGE MultiWayIf        #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Unsafe            #-}

{-|
Module      : Text.Pandoc.Filter.Pyplot
Description : Pandoc filter to create Matplotlib figures from code blocks
Copyright   : (c) Laurent P René de Cotret, 2019
License     : MIT
Maintainer  : laurent.decotret@outlook.com
Stability   : stable
Portability : portable

This module defines a Pandoc filter @makePlot@ that can be
used to walk over a Pandoc document and generate figures from
Python code blocks.

The syntax for code blocks is simple, Code blocks with the @.pyplot@
attribute will trigger the filter. The code block will be reworked into a Python
script and the output figure will be captured, along with a high-resolution version
of the figure and the source code used to generate the figure.

To trigger pandoc-pyplot, the following is __required__:

    * @.pyplot@: Trigger pandoc-pyplot but let it decide on a filename

Here are the possible attributes what pandoc-pyplot understands:

    * @target=...@: Filepath where the resulting figure should be saved.
    * @directory=...@ : Directory where to save the figure.
    * @caption="..."@: Specify a plot caption (or alternate text).
    * @dpi=...@: Specify a value for figure resolution, or dots-per-inch. Default is 80DPI.
    * @include=...@: Path to a Python script to include before the code block. Ideal to avoid repetition over many figures.

Here are some example blocks in Markdown:

@
This is a paragraph

```{.pyplot caption="This is a caption."}
import matplotlib.pyplot as plt

plt.figure()
plt.plot([0,1,2,3,4], [1,2,3,4,5])
plt.title('This is an example figure')
```
@

This filter was originally designed to be used with [Hakyll](https://jaspervdj.be/hakyll/). 
In case you want to use the filter with your own Hakyll setup, you can use a transform 
function that works on entire documents:

@
import Text.Pandoc.Filter.Pyplot (plotTransform)

import Hakyll

-- Unsafe compiler is required because of the interaction
-- in IO (i.e. running an external Python script).
makePlotPandocCompiler :: Compiler (Item String)
makePlotPandocCompiler =
  pandocCompilerWithTransformM
    defaultHakyllReaderOptions
    defaultHakyllWriterOptions
    (unsafeCompiler . plotTransform)
@

-}
module Text.Pandoc.Filter.Pyplot
    ( makePlot
    , plotTransform
    , PandocPyplotError(..)
      -- For testing purposes only
    , makePlot'
    , directoryKey
    , captionKey
    , dpiKey
    , includePathKey
    , saveFormatKey
    ) where

import           Control.Monad                 ((>=>))

import           Data.List                     (intersperse)

import qualified Data.Map.Strict               as Map
import           Data.Maybe                    (fromMaybe)
import           Data.Monoid                   ((<>))
import qualified Data.Text                     as T
import qualified Data.Text.IO                  as T
import           Data.Version                  (showVersion)

import           Paths_pandoc_pyplot           (version)

import           System.Directory              (createDirectoryIfMissing,
                                                doesFileExist)
import           System.FilePath               (isValid, makeValid,
                                                replaceExtension, takeDirectory)

import           Text.Pandoc.Definition
import           Text.Pandoc.Walk              (walkM)

import           Text.Pandoc.Filter.FigureSpec (FigureSpec (..),
                                                SaveFormat (..), addPlotCapture,
                                                figurePath, hiresFigurePath,
                                                saveFormatFromString)
import           Text.Pandoc.Filter.Scripting

-- | Possible errors returned by the filter
data PandocPyplotError
    = ScriptError Int -- ^ Running Python script has yielded an error
    | InvalidTargetError FilePath -- ^ Invalid figure path
    | BlockingCallError -- ^ Python script contains a block call to 'show()'
    deriving (Eq)

-- | Translate filter error to an error message
instance Show PandocPyplotError where
    show (ScriptError exitcode) =
        "Script error: plot could not be generated. Exit code " <> (show exitcode)
    show (InvalidTargetError fname) = "Target filename " <> fname <> " is not valid."
    show BlockingCallError = "Script contains a blocking call to show, like 'plt.show()'"

-- | Keys that pandoc-pyplot will look for in code blocks. These are only exported for testing purposes.
directoryKey, captionKey, dpiKey, includePathKey, saveFormatKey :: String
directoryKey = "directory"

captionKey = "caption"

dpiKey = "dpi"

includePathKey = "include"

saveFormatKey = "format"

-- | list of all keys related to pandoc-pyplot.
inclusionKeys :: [String]
inclusionKeys = [directoryKey, captionKey, dpiKey, includePathKey, saveFormatKey]

-- | Determine inclusion specifications from Block attributes.
-- Note that the target key is required, but all other parameters are optional
parseFigureSpec :: Block -> IO (Maybe FigureSpec)
parseFigureSpec (CodeBlock (id', cls, attrs) content)
    | "pyplot" `elem` cls = Just <$> figureSpec
    | otherwise = return Nothing
  where
    attrs' = Map.fromList attrs
    filteredAttrs = filter (\(k, _) -> k `notElem` inclusionKeys) attrs
    dir = makeValid $ Map.findWithDefault "generated" directoryKey attrs'
    format = fromMaybe (PNG) $ saveFormatFromString $ Map.findWithDefault "png" saveFormatKey attrs'
    includePath = Map.lookup includePathKey attrs'
    figureSpec :: IO FigureSpec
    figureSpec = do
        includeScript <- fromMaybe (return "") $ T.readFile <$> includePath
        let header = "# Generated by pandoc-pyplot " <> ((T.pack . showVersion) version)
            fullScript = mconcat $ intersperse "\n" [header, includeScript, T.pack content]
            caption' = Map.findWithDefault mempty captionKey attrs'
            dpi' = read $ Map.findWithDefault "80" dpiKey attrs'
            blockAttrs' = (id', filter (/= "pyplot") cls, filteredAttrs)
        return $ FigureSpec caption' fullScript format dir dpi' blockAttrs'
parseFigureSpec _ = return Nothing

-- | Check figure specifications for common mistakes
validateSpec :: FigureSpec -> Maybe PandocPyplotError
validateSpec spec
    | not (isValid path) = Just $ InvalidTargetError path
    | hasBlockingShowCall rendered = Just $ BlockingCallError
    | otherwise = Nothing
  where
    path = figurePath spec
    rendered = script spec

-- | Run the Python script. In case the file already exists, we can safely assume
-- there is no need to re-run it.
runScriptIfNecessary :: FigureSpec -> IO ScriptResult
runScriptIfNecessary spec = do
    createDirectoryIfMissing True . takeDirectory $ figurePath spec
    fileAlreadyExists <- doesFileExist $ figurePath spec
    if fileAlreadyExists
        then return ScriptSuccess
        else do
            result <- runTempPythonScript $ addPlotCapture spec
            case result of
                ScriptFailure code -> return $ ScriptFailure code
                ScriptSuccess
                    -- Save the original script into a separate file
                    -- so it can be inspected
                    -- Note : using a .txt file allows to view source directly
                    --        in the browser, in the case of HTML output
                 -> do
                    let sourcePath = replaceExtension (figurePath spec) ".txt"
                    T.writeFile sourcePath $ script spec
                    return ScriptSuccess

-- | Main routine to include Matplotlib plots.
-- Code blocks containing the attributes @.pyplot@ are considered
-- Python plotting scripts. All other possible blocks are ignored.
makePlot' :: Block -> IO (Either PandocPyplotError Block)
makePlot' block = do
    parsed <- parseFigureSpec block
    case parsed of
        Nothing -> return $ Right block
        Just spec ->
            case validateSpec spec of
                Just err -> return $ Left err
                Nothing -> do
                    result <- runScriptIfNecessary spec
                    case result of
                        ScriptFailure code -> return $ Left $ ScriptError code
                        ScriptSuccess -> do
                            let relevantAttrs = blockAttrs spec
                                sourcePath = replaceExtension (figurePath spec) ".txt"
                                hiresPath = hiresFigurePath spec
                                srcTarget = Link nullAttr [Str "Source code"] (sourcePath, "")
                                hiresTarget = Link nullAttr [Str "high res."] (hiresPath, "")
                                -- TODO: use pandoc-types Builder module
                                caption' =
                                    [ Str $ caption spec
                                    , Space
                                    , Str "("
                                    , srcTarget
                                    , Str ","
                                    , Space
                                    , hiresTarget
                                    , Str ")"
                                    ]
                            -- To render images as figures with captions, the target title
                            -- must be "fig:"
                            -- Janky? yes
                                image = Image relevantAttrs caption' (figurePath spec, "fig:")
                            return $ Right $ Para $ [image]

-- | Highest-level function that can be walked over a Pandoc tree.
-- All code blocks that have the '.pyplot' parameter will be considered
-- figures.
makePlot :: Block -> IO Block
makePlot = makePlot' >=> either (fail . show) return

-- | Walk over an entire Pandoc document, changing appropriate code blocks
-- into figures.
plotTransform :: Pandoc -> IO Pandoc
plotTransform = walkM makePlot