% Annotate LaTeX Equations
% 
% (c) 2022 by ST John, https://github.com/st--/
% Licensed under MIT License
% 
\NeedsTeXFormat{LaTeX2e}[1994/06/01]
\ProvidesPackage{annotate-equations}
  [2023/05/06 v0.2.2 easily annotate equations using TikZ]

%%% lualatex compatibility, from https://tex.stackexchange.com/a/351520/171664
\RequirePackage{ifluatex}
\ifluatex
\RequirePackage{luatex85}
\RequirePackage{pdftexcmds}
  \makeatletter
  \let\pdfstrcmp\pdf@strcmp
  \let\pdffilemoddate\pdf@filemoddate
  \makeatother
\fi
%%%

\RequirePackage{tikz}
\RequirePackage{xcolor}

\usetikzlibrary{backgrounds}
\usetikzlibrary{arrows,shapes}
\usetikzlibrary{tikzmark} % for \tikzmarknode
\usetikzlibrary{calc} % for computing the midpoint between two nodes, e.g. at ($(p1.north)!0.5!(p2.north)$) 


%%%%% SETTINGS %%%%%

\newcommand{\eqnhighlightheight}{}  % colorbox will shrink to content
\renewcommand{\eqnhighlightheight}{\mathstrut} % colorbox will always have full height

\newcommand{\eqnhighlightshade}{17}  % light
%\renewcommand{\eqnhighlightshade}{47}  % dark

\newcommand{\eqnannotationstrut}{\strut} % Package default
\newcommand{\eqnannotationfont}{\sffamily\footnotesize}


\providecommand\EAmarkanchor{north} % default set to "above"
\providecommand\EAwesteast{east} % default set to "right"
\providecommand\EAlabelanchor{south} % default set to "label above"
% for pgfkeys, see https://tex.stackexchange.com/a/34318/171664
% for no-value keys, see https://tex.stackexchange.com/a/401848/171664
\pgfkeys{
    /eqnannotate/.is family, /eqnannotate,
    above/.code = {\renewcommand\EAmarkanchor{north}},
    below/.code = {\renewcommand\EAmarkanchor{south}},
    left/.code = {\renewcommand\EAwesteast{west}},
    right/.code = {\renewcommand\EAwesteast{east}},
    label above/.code = {\renewcommand\EAlabelanchor{south}},
    label below/.code = {\renewcommand\EAlabelanchor{north}},
}

\tikzset{annotate equations/arrow/.style={}}
\tikzset{annotate equations/text/.style={font=\eqnannotationfont}}

%%%%% %%%%%%%% %%%%%


\newcommand*{\eqnhighlightcolorbox}[2]{%
% \colorbox sets the second argument in text mode, so for use within equations we wrap it in $ $ again
    \mathchoice% to get right font size in each mode:
        {\colorbox{#1}{$\displaystyle #2$}}%
        {\colorbox{#1}{$\textstyle #2$}}%
        {\colorbox{#1}{$\scriptstyle #2$}}%
        {\colorbox{#1}{$\scriptscriptstyle #2$}}%
}

%%% the fbox with 0pt rule fixes the height of eqnmark vs eqnmarkbox issue
\newcommand*{\eqnhighlightfbox}[2]{%
% \fbox sets the second argument in text mode, so for use within equations we wrap it in $ $ again
    \mathchoice% to get right font size in each mode:
        {\begingroup\setlength{\fboxrule}{0pt}\fbox{$\displaystyle\color{#1}#2$}\endgroup}%
        {\begingroup\setlength{\fboxrule}{0pt}\fbox{$\textstyle\color{#1}#2$}\endgroup}%
        {\begingroup\setlength{\fboxrule}{0pt}\fbox{$\scriptstyle\color{#1}#2$}\endgroup}%
        {\begingroup\setlength{\fboxrule}{0pt}\fbox{$\scriptscriptstyle\color{#1}#2$}\endgroup}%
}

% . is the current color

\newcommand*{\eqnhighlight}[2]{\begingroup\colorlet{currentcolor}{.}\eqnhighlightcolorbox{#1!\eqnhighlightshade}{\eqnhighlightheight #2}\endgroup}
\newcommand*{\eqncolor}[2]{\begingroup\colorlet{currentcolor}{.}\eqnhighlightfbox{#1}{\eqnhighlightheight #2}\endgroup}

%%% Arguments to \eqnmark[box]: [highlight color]{node name}{term to highlight}
\newcommand*{\eqnmarkbox}[3][currentcolor]{\addvalue{#2}{#1}\tikzmarknode{#2}{\eqnhighlight{#1}{#3}}}
\newcommand*{\eqnmark}[3][currentcolor]{\addvalue{#2}{#1}\tikzmarknode{#2}{\eqncolor{#1}{#3}}}

% Store current color in a dictionary (lookup table),
% from https://tex.stackexchange.com/a/48931/171664 :
\def\addvalue#1#2{\expandafter\gdef\csname eqnannotate@data@#1\endcsname{#2}}
\def\usevalue#1{%
  \ifcsname eqnannotate@data@#1\endcsname
    \csname eqnannotate@data@#1\expandafter\endcsname
  \else
    currentcolor%
  \fi
}


%%%%% Helpers for swapping north/south / west/east / -/+ depending on above/below / left/right etc.:
\newcommand*{\swapNorthSouth}[1]{%
    \ifnum\pdfstrcmp{#1}{south}=0 north\else south\fi
}
\newcommand*{\swapWestEast}[1]{%
    \ifnum\pdfstrcmp{#1}{east}=0 west\else east\fi
}
\newcommand*{\EAxshift}[1]{%
    \ifnum\pdfstrcmp{#1}{east}=0 -0.3ex\else 0.3ex\fi
}
%%%%%


\newcounter{eqnannotatenode}
\newcommand*{\eqnannotateCurrentNode}{eqnannotatenode\theeqnannotatenode}


\newcommand{\annotatetwo}[5][]{%
    \begingroup% %%% so we don't leak the \def's below
    \stepcounter{eqnannotatenode}%
    %%% #1: (optional) extra args for \node e.g. yshift=...
    \pgfkeys{/eqnannotate, #2}% %%% all configuration options
    \def\myEAmarkOne{#3}%
    \def\myEAmarkTwo{#4}%
    \colorlet{currentcolor}{.}
    \def\myEAtext{#5}%
    \def\myEAcolor{\usevalue{\myEAmarkOne}}%
    \begin{tikzpicture}[overlay,remember picture,>=stealth,nodes={align=left,inner ysep=1pt},<-]
        % default anchor is at center
		\node[anchor=\swapNorthSouth{\EAmarkanchor},color=\myEAcolor!85,
			annotate equations/text,#1
			]  % color blended with white to 85%, any (optional) extra args #1
            (\eqnannotateCurrentNode)   % use counter-based "local node"
            at ($(\myEAmarkOne.\EAmarkanchor)!0.5!(\myEAmarkTwo.\EAmarkanchor)$)  % centered between the two nodes
            {\myEAtext\eqnannotationstrut};
        % double arrow to two uses within the equation:
        \draw [<->,color=\myEAcolor, annotate equations/arrow] (\myEAmarkOne.\EAmarkanchor) |- ([yshift=0.1ex] \eqnannotateCurrentNode.\EAlabelanchor) -| (\myEAmarkTwo.\EAmarkanchor);  % from node 1 via annotation to node 2, with anchor #6 each
    \end{tikzpicture}%
    \endgroup% %%% close group again
}


%%% \extractfirst from https://tex.stackexchange.com/a/115733/171664
\RequirePackage{expl3}
\RequirePackage{xparse}
\ExplSyntaxOn
\NewDocumentCommand{\extractfirst}{mm}
 {
  \tl_set:Nx #1 {\clist_item:Nn #2 { 1 } }
 }
\ExplSyntaxOff
%%%

\newcommand{\annotate}[4][]{%
    \begingroup% %%% so we don't leak the \def's below
    \stepcounter{eqnannotatenode}%
    %
    %
    %
    % #1: (optional) extra args for \node e.g. yshift=...
    \pgfkeys{/eqnannotate, #2}
    \def\myEAmarks{#3}%
    \extractfirst\myEAmark\myEAmarks% %%% get first node for color and annotation
    \def\myEAtext{#4}%
    %
    %
    \colorlet{currentcolor}{.}
    \def\myEAcolor{\usevalue{\myEAmark}}%
    %
    \def\EAspace{ }  % workaround: did not find any other way of getting a space into \myEAlabelanchor without upsetting LaTeX/PGF/... somehow
    \edef\myEAlabelanchor{\EAlabelanchor\EAspace\EAwesteast}%
    %
    %
    \def\myEAxshift{\EAxshift{\EAwesteast}}%
    \begin{tikzpicture}[overlay,remember picture,>=stealth,nodes={align=left,inner ysep=1pt},<-]
		\node[anchor=\swapNorthSouth{\EAmarkanchor} \swapWestEast{\EAwesteast},
			color=\myEAcolor!85,annotate equations/text,#1] % TODO for some reason, passing #1 through command doesn't work...
                % anchor=west: align left edge of text on top of tikzmark in equation
                % should be north west for below and south west for above ...
            (\eqnannotateCurrentNode) at (\myEAmark.\EAmarkanchor)  % \EAmarkanchor north: above the equation, south: below
            {\myEAtext\eqnannotationstrut};
        \foreach \EAmark in \myEAmarks
        \draw [color=\myEAcolor, annotate equations/arrow] (\EAmark.\EAmarkanchor)  % arrow from the equation
                % \EAmarkanchor north: above the equation, south: below
            |- ([xshift=\myEAxshift,yshift=0.1ex] \eqnannotateCurrentNode.\myEAlabelanchor);
                % - south east: we want line to end at bottom right of annotation text;
                % - negative xshift makes it a little bit shorter;
                % - yshift for aesthetics (\strut is ever so slightly too tall).
    \end{tikzpicture}%
    \endgroup% %%% close group again
}

\endinput
%%
%% End of file