% Copyright (c) 2025 George Allison
%
% This work may be distributed and/or modified under the
% conditions of the LaTeX Project Public License, either version 1.3
% of this license or any later version.
% The latest version of this license is in
%   http://www.latex-project.org/lppl.txt
% and version 1.3 or later is part of all distributions of LaTeX
% version 2005/12/01 or later.
%
% This work has the LPPL maintenance status `maintained'.
%
% The Current Maintainer of this work is George Allison.
%
% This work consists of the files lua-regression.sty and lua-regression.lua
%
% This software is provided 'as is', without warranty of any kind,
% either expressed or implied, including, but not limited to, the
% implied warranties of merchantability and fitness for a
% particular purpose.

\ProvidesPackage{lua-regression}[2025/04/17 1.0.2  Lua Regression Plotting project]

\RequirePackage{ifthen}
\ifluatex
  \RequirePackage{luacode}
\else
  {\PackageError{lua-regression}
  {Not running under LuaLaTeX}
  {This package requires LuaLaTeX. Try compiling this document with\MessageBreak 'lualatex' instead of 'latex'. This is a fatal error; I'm aborting now.}%
  }\stop
\fi

% Required packages
\RequirePackage{pgfkeys}
\RequirePackage{pgfplots}
\usepgfplotslibrary{fillbetween}

% Define the key-value options
\pgfkeys{
    /luaregression/.is family, /luaregression,
    default/.style = {
        xcol=1,            % Default x-column index
        ycol=2,            % Default y-column index
        ci=false,  % Default: no error band
        z-threshold=null,  % Default Z-score threshold
        sig-figures=4,     % Default significant figures
        order=1,          % Default polynomial order
        plot=false,        % Default plotting behavior
        pgf-options={mark=none,smooth}, % Default PGF options
        eq=false, % Toggle for showing the equation
        r2=false,       % Toggle for showing R²
        debug=false,    % Debug toggle for csv
        bootstrap=1000, % Number of bootstrap samples for confidence intervals
        cicolor=blue,        % CI fill color
        cifillopacity=0.2,  % CI fill opacity
    },
    xcol/.estore in = \luaregressionxcol,
    ycol/.estore in = \luaregressionycol,
    ci/.estore in = \luaregressionci,
    z-threshold/.estore in = \luaregressionzthreshold,
    sig-figures/.estore in = \luaregressionsigfigures,
    order/.estore in = \luaregressionorder,
    plot/.estore in = \luaregressionplot,
    pgf-options/.estore in = \luaregressionpgfoptions,
    eq/.estore in = \luaregressionshowequation,
    r2/.estore in = \luaregressionshowrsquare,
    debug/.estore in = \luaregressiondebug,
    bootstrap/.estore in = \luaregressionbootstrapsamples,
    cicolor/.estore in = \luaregressioncicolor,
    cifillopacity/.estore in = \luaregressioncifillopacity,
}

% Define the macro
\newcommand{\luaregression}[2][]{%
    \pgfkeys{/luaregression, default, #1}% Parse the options
    \directlua{
        require("lua-regression")
        process_data_with_options(
            "#2",
            {
                ["xcol"] = tonumber("\luaregressionxcol"),
                ["ycol"] = tonumber("\luaregressionycol"),
                ["z_threshold"] = tonumber("\luaregressionzthreshold"),
                ["sig_figures"] = tonumber("\luaregressionsigfigures"),
                ["ci"] = ("\luaregressionci" == "true"),
                ["order"] = tonumber("\luaregressionorder"),
                ["debug"] = ("\luaregressiondebug" == "true"),
                ["bootstrap_samples"] = tonumber("\luaregressionbootstrapsamples"),
            }
        )
    }%
    \ifthenelse{\equal{\luaregressionplot}{true}}{%
        \ifx\addplot\undefined
            \PackageError{lua-regression}{'plot=true' requires a tikzpicture environment and pgfplots}%
            {Use '\\begin{tikzpicture} ... \\end{tikzpicture}' with '\\usepackage{pgfplots}'.}%
        \fi
        \expandafter\addplot\expandafter[\luaregressionpgfoptions] {\polyeq};%
        % Plot confidence band if ci=true
        \ifthenelse{\equal{\luaregressionci}{true}}{%
            \begingroup
            \edef\luareg@ci@addplot{%
            \noexpand\addplot[fill=\luaregressioncicolor,fill opacity=\luaregressioncifillopacity] fill between[of=quprpath and qlwrpath];%
            }
            \addplot[name path=qlwrpath,draw=none] coordinates {\qlwr};
            \addplot[name path=quprpath,draw=none] coordinates {\qupr};
            \luareg@ci@addplot
            \endgroup
        }{}%
        % Construct the legend entry dynamically
        \begingroup
        \def\legendentry{}%
        \ifthenelse{\equal{\luaregressionshowequation}{true}}{%
            \edef\legendentry{$y = \printeq$}%
        }{}%
        \ifthenelse{\equal{\luaregressionshowrsquare}{true}}{%
            \ifx\legendentry\empty
                \edef\legendentry{$R^2 = \polyR$}%
            \else
                \edef\legendentry{\legendentry\\$R^2 = \polyR$}%
            \fi
        }{}%
        \ifx\legendentry\empty
        \else
            \expandafter\addlegendentry\expandafter{\legendentry}%
        \fi
        \endgroup
    }{}%
}