% Copyright 2019 by Christophe Jorssen
%
% This file may be distributed and/or modified
%
% 1. under the LaTeX Project Public License and/or
% 2. under the GNU Public License.
%
% See the file doc/generic/pgf/licenses/LICENSE for more details.
\def\pgfmathode@stopadd{\pgfmathode@stopadd}

\def\pgfmathode@addindextofunc[#1]#2{%
  \begingroup
    \gdef\pgfmathode@func{}%
    \pgfmathode@addindextofunc@i[#1]#2\pgfmathode@stopadd}

\def\pgfmathode@addindextofunc@i[#1]#2#3{%
    \ifx#1#2
      \expandafter\gdef\expandafter\pgfmathode@func\expandafter{%
        \pgfmathode@func #2[\pgfmathode@timeindex]}%
    \else
      \expandafter\gdef\expandafter\pgfmathode@func\expandafter{%
        \pgfmathode@func #2}%
    \fi
    \ifx#3\pgfmathode@stopadd
      \expandafter\gdef\expandafter\pgfmathode@func\expandafter{%
        \expandafter{\pgfmathode@func}}%
      \let\next\endgroup
    \else
      \def\next{\pgfmathode@addindextofunc@i[#1]#3}%
    \fi
    \next}

%**********************************************************************
% Know limitation: fk cannot be a time function. Should be possible to
% use \Sol[0] as time in the set of equations.
%**********************************************************************

% Numerically solve a differential system with 1st order Runge-Kutta
% (Euler method)
% dq1/dt = f1(t,q1,q2,...,qN)
% dq2/dt = f2(t,q1,q2,...,qN)
% ...
% dqN/dt = fN(t,q1,q2,...,qN)
%
% #1: A pgf array where the solution will be stored.
% At the end:
% #1 -> {{t(0),q1(t(0)),q2(t(0)),...,qN(t(0))},
%        {t(1),q1(t(1)),q2(t(1)),...,qN(t(1))},
%        ...
%        {t(Nstep),q1(t(Nstep)),q2(t(Nstep)),...,qN(t(Nstep))}}
% #2: A pgf array with f1,f2,...,fN
% {f1}
% #3: A pgf array with initial conditions.
% {t(0),q1(t(0)),q2(t(0)),...,qN(t(0))}
% #4: t(Nstep)
% #5: Nstep (number of steps)
%
% * First order example: dq/dt = -q + 5 with IC=(t=0,q(0)=0)
% \pgfmathodeRKI{\Sol}{{-\Sol[1]+5}}{{0,0}}{3.5}{100}
% * Second order example: {dq1/dt=q2, dq2/dt=-q2/2-sin(q1)} (damped
%   pendulum)
% \pgfmathodeRKI{\Sol}{{\Sol[2],-\Sol[2]/2-sin(deg(\Sol[1]))}}%
%   {{0,-1,5}}{15}{50}

\def\pgfmathode@generic@init@RK{%
  \begingroup
    % Sol array (#1) first element is given by the set of initial
    % conditions (#3). We need to enclose it in brace so that is a now
    % a "matrix" array.
    \gdef#1{{#3}}%
    % The set of rhs {f1(t,q1,q2,...,qN),...,fN(t,q1,q2,...,qN)} is
    % given in terms of #1[0],#1[1],...,#1[N]. We need to transform to
    % #1[\pgfmathode@timeindex][0],...,#1[\pgfmathode@timeindex][N], that is
    % add [\pgfmathode@timeindex] after every occurrence of #1 in #2.
    % No need to brace #2 since it is already surrounded by braces. The
    % result is stored in \pgfmathode@func
    \pgfmathode@addindextofunc[#1]#2%
    % Set time step. Actually the number of time steps is #5-1, so we
    % divide by (#5-1)+1 = #5.
    \pgfmathsetmacro\pgfmathode@Sol@t{#1[0][0]}%
    \pgfmathsetmacro\pgfmathode@Deltat{(#4-\pgfmathode@Sol@t)/#5}%
    % Compute the order of the system: it is equal to the number of IC
    % minus one (time).
    \pgfmathsetmacro\pgfmathode@dimfunc{dim(#3)-1}%
    \pgfmathsetmacro\pgfmathode@numstep{#5-1}}

\def\pgfmathodeRKI#1#2#3#4#5{%
    \pgfmathode@generic@init@RK
    % Start the time loop for Runge-Kutta 1st order
    \foreach \pgfmathode@timeindex in {0,...,\pgfmathode@numstep} {%
      \message{Step \pgfmathode@timeindex...}%
      % Start the function loop
      \foreach \pgfmathode@funcindex in {1,...,\pgfmathode@dimfunc} {%
        % The RKI algorithm
        % qk[t(i)] = Deltat*fk(t(i-1),q1(t(i-1)),...qN(t(i-1))) +
        %            qk[t(i-1)]
        % Note: this part can be really slow due to array management in
        % pgfmath.
        % TODO: find a better way (at least inside this loop)
        \pgfmathparse{%
          \pgfmathode@Deltat*(\pgfmathode@func[\pgfmathode@funcindex-1])+
          #1[\pgfmathode@timeindex][\pgfmathode@funcindex]}%
        % Store the result in a helper macro
        \expandafter\xdef\csname pgfmathode@Sol@\pgfmathode@funcindex
          \endcsname{\pgfmathresult}%
      }
      % Store step i
      \pgfmathparse{\pgfmathode@Sol@t+\pgfmathode@Deltat}%
      \xdef\pgfmathode@Sol@t{\pgfmathresult}%
      \xdef\pgfmathode@temp{}%
      \foreach \pgfmathode@funcindex in {1,...,\pgfmathode@dimfunc} {%
        \ifx\pgfmathode@temp\pgfutil@empty
          \xdef\pgfmathode@temp{%
            \csname pgfmathode@Sol@\pgfmathode@funcindex\endcsname}%
        \else
          \xdef\pgfmathode@temp{%
            \pgfmathode@temp,
            \csname pgfmathode@Sol@\pgfmathode@funcindex\endcsname}%
        \fi}
      \xdef#1{%
        {\expandafter\pgfutil@firstofone#1,%
          {\pgfmathode@Sol@t,\pgfmathode@temp}}}%
    }
  \endgroup}

% RKIV: far from working yet!
\def\pgfmathodeRKIV#1#2#3#4#5{%
    \pgfmathode@generic@init@RK
    % Start the time loop for Runge-Kutta 4th order
    \foreach \pgfmathode@timeindex in {0,...,\pgfmathode@numstep} {%
      \message{Step \pgfmathode@timeindex...}%
      % Start the function loop
      \foreach \pgfmathode@funcindex in {1,...,\pgfmathode@dimfunc} {%
        % The RKIV algorithm
        % qk[t(i)] = qk[t(i-1)] +
        %            (1/6) * (j1k + 2j2k + 2j3k + j4k)
        % where j1k = fk(t(i-1),q1(t(i-1)),...qN(t(i-1))) * Deltat
        %       j2k = fk(t(i-1)+.5*Deltat,
        %                q1(t(i-1))+.5*j11,...qN(t(i-1))+.5*j1N) * Deltat
        %       j3k = fk(t(i-1)+.5*Deltat,
        %                q1(t(i-1))+.5*j21,...qN(t(i-1))+.5*j2N) * Deltat
        %       j4k = fk(t(i-1)+Deltat,
        %                q1(t(i-1))+j31,...qN(t(i-1))+j3N) * Deltat
        %
        % Compute j1k
        \pgfmathparse{%
          \pgfmathode@Deltat*(\pgfmathode@func[\pgfmathode@funcindex-1])+
          #1[\pgfmathode@timeindex][\pgfmathode@funcindex]}%
        \expandafter\xdef\csname pgfmathode@RKIV@1@\pgfmathode@funcindex
          \endcsname{\pgfmathresult}%
        % Compute j2k
        \pgfmathparse{%
          \pgfmathode@Deltat*(\pgfmathode@func[\pgfmathode@funcindex-1])+
          #1[\pgfmathode@timeindex][\pgfmathode@funcindex]}%
        \expandafter\xdef\csname pgfmathode@RKIV@1@\pgfmathode@funcindex
          \endcsname{\pgfmathresult}%
      }
      % Store step i
      \pgfmathparse{\pgfmathode@Sol@t+\pgfmathode@Deltat}%
      \xdef\pgfmathode@Sol@t{\pgfmathresult}%
      \xdef\pgfmathode@temp{}%
      \foreach \pgfmathode@funcindex in {1,...,\pgfmathode@dimfunc} {%
        \ifx\pgfmathode@temp\pgfutil@empty
          \xdef\pgfmathode@temp{%
            \csname pgfmathode@Sol@\pgfmathode@funcindex\endcsname}%
        \else
          \xdef\pgfmathode@temp{%
            \pgfmathode@temp,
            \csname pgfmathode@Sol@\pgfmathode@funcindex\endcsname}%
        \fi}
      \xdef#1{%
        {\expandafter\pgfutil@firstofone#1,%
          {\pgfmathode@Sol@t,\pgfmathode@temp}}}%
    }
  \endgroup}


\endinput
