{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Copyright   : Written by David Himmelstrup
-- License     : Unlicense
-- Maintainer  : lemmih@gmail.com
-- Stability   : experimental
-- Portability : POSIX
module Reanimate.LaTeX
  ( latexCfg,
    TexEngine (..),
    TexConfig (..),
    latex,
    latexWithHeaders,
    latexChunks,
    xelatex,
    xelatexWithHeaders,
    ctex,
    ctexWithHeaders,
    latexAlign,

    -- * Font configurations
    chalkduster,
    calligra,
    noto,
    helvet,
    libertine,
  )
where

import Control.Lens
import qualified Data.ByteString as B
import Data.Hashable
import Data.Monoid
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import GHC.Generics (Generic)
import Graphics.SvgTree
  ( Tree,
    clipPathRef,
    clipRule,
    mapTree,
    parseSvgFile,
    strokeColor,
    pattern ClipPathTree,
    pattern None,
  )
import Reanimate.Animation (SVG)
import Reanimate.Cache
import Reanimate.External
import Reanimate.Misc
import Reanimate.Parameters
import Reanimate.Svg
import System.FilePath
  ( replaceExtension,
    takeFileName,
    (</>),
  )
import System.IO.Unsafe (unsafePerformIO)

-- | TeX backends. They have different features and capabilities.
data TexEngine = LaTeX | XeLaTeX | LuaLaTeX
  deriving (Generic, Hashable, Eq, Ord, Read, Show)

-- | TeX configurations can load packages and set up environments for tex scripts.
data TexConfig = TexConfig
  { texConfigEngine :: TexEngine,
    texConfigHeaders :: [T.Text],
    texConfigPostScript :: [T.Text]
  }
  deriving (Generic, Hashable, Read, Show, Eq, Ord)

-- | Render TeX script using a given configuration.
latexCfg :: TexConfig -> T.Text -> SVG
latexCfg (TexConfig engine headers postscript) =
  gen headers postscript
  where
    gen =
      case engine of
        LaTeX -> someTexWithHeaders engine "latex" "dvi" []
        XeLaTeX -> someTexWithHeaders engine "xelatex" "xdv" ["-no-pdf"]
        LuaLaTeX -> someTexWithHeaders engine "lualatex" "pdf" []

-- | Invoke latex and import the result as an SVG object. SVG objects are
--   cached to improve performance.
--
--   Example:
--
--   > latex "$e^{i\\pi}+1=0$"
--
--   <<docs/gifs/doc_latex.gif>>
latex :: T.Text -> Tree
latex = latexWithHeaders []

-- | Invoke latex with extra script headers.
latexWithHeaders :: [T.Text] -> T.Text -> Tree
latexWithHeaders = someTexWithHeaders LaTeX "latex" "dvi" [] []

someTexWithHeaders ::
  TexEngine ->
  String ->
  String ->
  [String] ->
  [T.Text] ->
  [T.Text] ->
  T.Text ->
  Tree
someTexWithHeaders _engine _exec _dvi _args _headers _postscript tex
  | pNoExternals = mkText tex
someTexWithHeaders engine exec dvi args headers postscript tex =
  (unsafePerformIO . (cacheMem . cacheDiskSvg) (latexToSVG engine dvi exec args))
    script
  where
    script = mkTexScript exec args headers (T.unlines (postscript ++ [tex]))

-- | Invoke latex and separate results.
latexChunks :: [T.Text] -> [Tree]
latexChunks chunks | pNoExternals = map mkText chunks
latexChunks chunks = worker (svgGlyphs $ latex $ T.concat chunks) chunks
  where
    merge lst = mkGroup [fmt svg | (fmt, _, svg) <- lst]
    worker [] [] = []
    worker _ [] = error "latex chunk mismatch"
    worker everything (x : xs) =
      let width = length $ svgGlyphs (latex x)
       in merge (take width everything) : worker (drop width everything) xs

-- | Invoke xelatex and import the result as an SVG object. SVG objects are
--   cached to improve performance. Xelatex has support for non-western scripts.
xelatex :: Text -> Tree
xelatex = xelatexWithHeaders []

-- | Invoke xelatex with extra script headers.
xelatexWithHeaders :: [T.Text] -> T.Text -> Tree
xelatexWithHeaders = someTexWithHeaders XeLaTeX "xelatex" "xdv" [] ["-no-pdf"]

-- | Invoke xelatex with "\usepackage[UTF8]{ctex}" and import the result as an
--   SVG object. SVG objects are cached to improve performance. Xelatex has
--   support for non-western scripts.
--
--   Example:
--
--   > ctex "中文"
--
--   <<docs/gifs/doc_ctex.gif>>
ctex :: T.Text -> Tree
ctex = ctexWithHeaders []

-- | Invoke xelatex with extra script headers + ctex headers.
ctexWithHeaders :: [T.Text] -> T.Text -> Tree
ctexWithHeaders headers = xelatexWithHeaders ("\\usepackage[UTF8]{ctex}" : headers)

-- | Invoke latex and import the result as an SVG object. SVG objects are
--   cached to improve performance. This wraps the TeX code in an 'align*'
--   context.
--
--   Example:
--
--   > latexAlign "R = \\frac{{\\Delta x}}{{kA}}"
--
--   <<docs/gifs/doc_latexAlign.gif>>
latexAlign :: Text -> Tree
latexAlign tex = latex $ T.unlines ["\\begin{align*}", tex, "\\end{align*}"]

postprocess :: Tree -> Tree
postprocess =
  simplify
    . lowerTransformations
    . scaleXY 0.1 (-0.1)
    . removeClipPaths
    . lowerIds
    . mapTree clearDrawAttr
  where
    clearDrawAttr t = t & strokeColor .~ Last Nothing

enginePostprocess :: TexEngine -> Tree -> Tree
enginePostprocess LuaLaTeX svg = translate 0 (svgHeight svg) svg
enginePostprocess _ svg = svg

removeClipPaths :: SVG -> SVG
removeClipPaths = mapTree worker
  where
    worker ClipPathTree {} = None
    worker t = t & clipRule .~ Last Nothing & clipPathRef .~ Last Nothing

-- executable, arguments, header, tex
latexToSVG :: TexEngine -> String -> String -> [String] -> Text -> IO Tree
latexToSVG engine dviExt latexExec latexArgs tex = do
  latexBin <- requireExecutable latexExec
  withTempDir $ \tmp_dir -> withTempFile "tex" $ \tex_file ->
    withTempFile "svg" $ \svg_file -> do
      let dvi_file =
            tmp_dir </> replaceExtension (takeFileName tex_file) dviExt
      B.writeFile tex_file (T.encodeUtf8 tex)
      runCmd
        latexBin
        ( latexArgs
            ++ [ "-interaction=nonstopmode",
                 "-halt-on-error",
                 "-output-directory=" ++ tmp_dir,
                 tex_file
               ]
        )
      if dviExt == "pdf"
        then do
          pdf2svg <- requireExecutable "pdf2svg"
          runCmd
            pdf2svg
            [dvi_file, svg_file]
        else do
          dvisvgm <- requireExecutable "dvisvgm"
          runCmd
            dvisvgm
            [ dvi_file,
              "--precision=5",
              "--exact", -- better bboxes.
              "--no-fonts", -- use glyphs instead of fonts.
              "--verbosity=0",
              "-o",
              svg_file
            ]
      svg_data <- B.readFile svg_file
      case parseSvgFile svg_file svg_data of
        Nothing -> error "Malformed svg"
        Just svg ->
          return $
            enginePostprocess engine $
              postprocess $ unbox $ replaceUses svg

mkTexScript :: String -> [String] -> [Text] -> Text -> Text
mkTexScript latexExec latexArgs texHeaders tex =
  T.unlines $
    [ "% " <> T.pack (unwords (latexExec : latexArgs)),
      "\\documentclass[preview]{standalone}",
      "\\usepackage{amsmath}",
      "\\usepackage{gensymb}"
    ]
      ++ texHeaders
      ++ [ "\\usepackage[english]{babel}",
           "\\linespread{1}",
           "\\begin{document}",
           tex,
           "\\end{document}"
         ]

{- Packages used by manim.

\\\usepackage{amsmath}\n\
\\\usepackage{amssymb}\n\
\\\usepackage{dsfont}\n\
\\\usepackage{setspace}\n\
\\\usepackage{relsize}\n\
\\\usepackage{textcomp}\n\
\\\usepackage{mathrsfs}\n\
\\\usepackage{calligra}\n\
\\\usepackage{wasysym}\n\
\\\usepackage{ragged2e}\n\
\\\usepackage{physics}\n\
\\\usepackage{xcolor}\n\
\\\usepackage{textcomp}\n\
\\\usepackage{xfrac}\n\
\\\usepackage{microtype}\n\
-}

-- | Chalkduster configuration. Depends on lualatex.
--   Font files are automatically downloaded.
--
--   Example:
--
-- @
-- `latexCfg` `chalkduster` "chalkduster"
-- @
--
--   <<docs/gifs/doc_chalkduster.gif>>
chalkduster :: TexConfig
chalkduster =
  TexConfig
    { texConfigEngine = XeLaTeX,
      texConfigHeaders =
        [ "\\usepackage[no-math]{fontspec}",
          "\\setmainfont[Mapping=tex-text,Path={" <> chalkdusterFont <> "/},Extension=.ttf]{Chalkduster}",
          "\\usepackage[defaultmathsizes]{mathastext}"
        ],
      texConfigPostScript = []
    }
  where
    chalkdusterFont =
      T.pack $
        zipArchive
          "https://www.ffonts.net/Chalkduster.font.zip"
          "Wplv4RjuFiI0hDQnAM5MVHl2evrZqWstRLdVAfBomCM="

-- | Calligra configuration.
--
--   Example:
--
-- @
-- `latexCfg` `calligra` "calligra"
-- @
--
--   <<docs/gifs/doc_calligra.gif>>
calligra :: TexConfig
calligra =
  TexConfig
    { texConfigEngine = LaTeX,
      texConfigHeaders = ["\\usepackage{calligra}"],
      texConfigPostScript = ["\\calligra"]
    }

-- | Noto configuration.
--
--   Example:
--
-- @
-- `latexCfg` `noto` "noto"
-- @
--
--   <<docs/gifs/doc_noto.gif>>
noto :: TexConfig
noto =
  TexConfig
    { texConfigEngine = LaTeX,
      texConfigHeaders = ["\\usepackage{noto}"],
      texConfigPostScript = []
    }

-- | Helvetica configuration.
--
--   Example:
--
-- @
-- `latexCfg` `helvet` "helvet"
-- @
--
--   <<docs/gifs/doc_helvet.gif>>
helvet :: TexConfig
helvet =
  TexConfig
    { texConfigEngine = LaTeX,
      texConfigHeaders = ["\\usepackage{helvet}"],
      texConfigPostScript = []
    }

-- | Libertine configuration.
--
--   Example:
--
-- @
-- `latexCfg` `libertine` "libertine"
-- @
--
--   <<docs/gifs/doc_libertine.gif>>
libertine :: TexConfig
libertine =
  TexConfig
    { texConfigEngine = LaTeX,
      texConfigHeaders = ["\\usepackage{libertine}"],
      texConfigPostScript = []
    }