%D \module
%D   [       file=mp-grap.mpiv,
%D        version=2012.10.16, % 2008.09.08 and earlier,
%D          title=\CONTEXT\ \METAPOST\ graphics,
%D       subtitle=graph packagesupport,
%D         author=Hans Hagen \& Alan Braslau,
%D           date=\currentdate,
%D      copyright={PRAGMA ADE \& \CONTEXT\ Development Team}]
%C
%C This module is part of the \CONTEXT\ macro||package and is
%C therefore copyrighted by \PRAGMA. See licen-en.pdf for
%C details.

if known context_grap : endinput ; fi

boolean context_grap ; context_grap := true ;

% Instead we could include graph here and then clean it up as well as use private
% variables in the grap_ namespace. After all, graph is frozen.

input graph.mp ;

vardef roundd(expr x, d) =
    if abs d > 4 :
        if d > 0 :
            x
        else :
            0
        fi
    elseif d > 0 :
        save i ; i = floor x ;
        i + round(Ten_to[d]*(x-i))/Ten_to[d]
    else :
        round(x/Ten_to[-d])*Ten_to[-d]
    fi
enddef ;

Ten_to0 =     1 ;
Ten_to1 =    10 ;
Ten_to2 =   100 ;
Ten_to3 =  1000 ;
Ten_to4 = 10000 ;

def sFe_base = enddef ;

if unknown Fe_plus :
    picture Fe_plus ; Fe_plus := textext("+") ; % btex + etex ;
fi ;

vardef format    (expr f,x) = dofmt_.Feform_(f,x) enddef ;
vardef Mformat   (expr f,x) = dofmt_.Meform (f,x) enddef ;
vardef formatstr (expr f,x) = dofmt_.Feform_(f,x) enddef ;
vardef Mformatstr(expr f,x) = dofmt_.Meform(f,x) enddef ;

vardef escaped_format(expr s) =
    "" for n=1 upto length(s) : &
        if ASCII substring (n,n+1) of s = 37 :
            "@"
        else :
            substring (n,n+1) of s
        fi
    endfor
enddef ;

vardef dofmt_@#(expr f, x) =
    textext("\MPgraphformat{" & escaped_format(f) & "}{" & (if string x : x else: decimal x fi) & "}")
  % textext(mfun_format_number(escaped_format(f),x))
enddef ;

% note that suffix @# is ignored above...

vardef strfmt(expr f, x) =
    "\MPgraphformat{" & escaped_format(f) & "}{" & (if string x : x else: decimal x fi) & "}"
enddef ;

% We redefine autogrid from graph.mp adding the possibility of differing X and Y
% formats. Autoform is defined in graph.mp (by default "%g").

% graph.mp: string Autoform; Autoform = "%g";
% graph.mp:
% graph.mp: vardef autogrid(suffix tx, ty) text w =
% graph.mp:   Gneedgr_:=false;
% graph.mp:   if str tx<>"": for x=auto.x: tx(Autoform,x) w; endfor fi
% graph.mp:   if str ty<>"": for y=auto.y: ty(Autoform,y) w; endfor fi
% graph.mp: enddef;

% string Autoform_X ; Autoform_X := "@.0e" ;
% string Autoform_Y ; Autoform_Y := "@.0e" ;

vardef autogrid(suffix tx, ty) text w =
    Gneedgr_ := false ;
    if str tx <> "" :
        for x=auto.x :
            tx (
                if string Autoform_X :
                    if Autoform_X <> "" :
                        Autoform_X
                    else :
                        Autoform
                    fi
                else :
                    Autoform
                fi,
                x
            ) w ;
        endfor
    fi
    if str ty <> "" :
        for y=auto.y :
            ty (
                if string Autoform_Y :
                    if Autoform_Y <> "" :
                        Autoform_Y
                    else :
                        Autoform
                    fi
                else :
                    Autoform
                fi,
                y
            ) w ;
        endfor
    fi
enddef ;

% A couple of extensions:

% Define a function plotsymbol() returning a picture: 10 different shapes,
% unfilled outline, interior filled with different shades of the background.
% This allows overlapping points on a plot to be more distinguishable.

% grap_symsize := fontsize defaultfont ; % can be redefined
%
% dynamic version:

vardef grap_symsize =
  % fontsize defaultfont
  % .8ExHeight
    .35BodyFontSize
enddef ;

path grap_sym[] ; % (internal) symbol path

grap_sym[0] := (0,0) ;                      % point
grap_sym[1] := fullcircle ;                 % circle
grap_sym[2] := (up -- down) scaled .5 ;     % vertical bar

for i = 3 upto 9 :                          % polygons
    grap_sym[i] :=
        for j = 0 upto i-1 :
            (up scaled .5) rotated (360j/i) --
        endfor cycle ;
endfor

grap_sym[12] := grap_sym[2] rotated +90 ;   % horizontal line
grap_sym[22] := grap_sym[2] rotated +45 ;   % backslash
grap_sym[32] := grap_sym[2] rotated -45 ;   % slash
grap_sym[13] := grap_sym[3] rotated 180 ;   % down  triangle
grap_sym[23] := grap_sym[3] rotated -90 ;   % right triangle
grap_sym[33] := grap_sym[3] rotated +90 ;   % left  triangle
grap_sym[14] := grap_sym[4] rotated +45 ;   % square
grap_sym[15] := grap_sym[5] rotated 180 ;   % down pentagon
grap_sym[16] := grap_sym[6] rotated +90 ;   % turned hexagon
grap_sym[17] := grap_sym[7] rotated 180 ;
grap_sym[18] := grap_sym[8] rotated +22.5 ;

numeric l ;

for j = 5 upto 9 :
    l := length(grap_sym[j]) ;
    pair p[] ;
    for i = 0 upto l :
        p[i] = whatever [point i             of grap_sym[j],
                         point (i+2   mod l) of grap_sym[j]] ;
        p[i] = whatever [point (i+1   mod l) of grap_sym[j],
                         point (i+l-1 mod l) of grap_sym[j]] ;
    endfor
    grap_sym[20+j] := for i = 0 upto l : point i of grap_sym[j]--p[i]--endfor cycle ;
endfor

path s    ; s := grap_sym[4] ;
path q    ; q := s scaled .25 ;
numeric l ; l := length(s) ;

pair p[] ;

grap_sym[24] := for i = 0 upto l-1 :
     hide(
         p[i]   = whatever [point i   of s, point (i+1   mod l) of s] ;
         p[i]   = whatever [point i   of q, point (i-1+l mod l) of q] ;
         p[i+l] = whatever [point i   of s, point (i+1   mod l) of s] ;
         p[i+l] = whatever [point i+1 of q, point (i+2   mod l) of q] ;
     )
     point i of q -- p[i] -- p[i+l] --
endfor cycle ;

grap_sym[34] := grap_sym[24] rotated 45 ;

% usage: gdraw p plot plotsymbol(1,red,1) ;       % a filled red circle
% usage: gdraw p plot plotsymbol(4,blue,0) ;      % a blue square
% usage: gdraw p plot plotsymbol(14,green,0.5) ;  % a 50% filled green diamond

def plotsymbol(expr n,c,f) = % (number,color,color|number)
    if known grap_sym[n] :
        image(
            path p ; p := grap_sym[n] scaled grap_symsize ;
            undraw p withpen currentpen scaled 2 ;
            if cycle p : fill p withcolor
                if color f and known f :
                    f
                elseif numeric f and known f and color c and known c :
                    f[background,c]
                elseif numeric f and known f :
                    f[background,black]
                else :
                    background
                fi ;
            fi
            draw p if color c and known c : withcolor c fi ;
        )
    else :
        nullpicture
    fi
enddef ;

% The following extensions are not specific to graph and could be moved to metafun...

% convert a polygon path to a smooth path (useful, e.g. as a guide to the eye)

def smoothpath (suffix $) =
    if path $ :
    (for i=0 upto length $ :
        if i>0 : .. fi
        (point i of $)
     endfor )
    fi
enddef ;

% return a path of a function func(x) with abcissa running from f to t over n intervals

def makefunctionpath (expr f, t, n) (text func) =
    (for x=f step ((t-f)/n) until t :
        if x<>f : .. fi
        (x, func)
    endfor )
enddef ;

% shift a path, point by point
%
% example:
%
%   p1 := addnoisetopath(p0,(.1normaldeviate,.1normaldeviate)) ;

vardef addnoisetopath (suffix p) (text t) =
    if path p :
    hide(pair p_i)
    (for i=0 upto length p :
        if i>0 : -- fi
        hide(p_i := point i of p; x := xpart p_i; y := ypart p_i)z shifted t
    endfor)
    fi
enddef ;

% return a new path of a function func(x) using the same abcissa as an existing path

vardef functionpath (suffix p) (text t) =
    (for i=0 upto length p :
        if i>0 : .. fi
        (hide(x := xpart(point i of p))x,t)
     endfor )
enddef ;

% least-squares "fit" to a polynomial
%
% example:
%
%   path p[] ;
%   numeric a[] ; a0 := 1 ; a1 := .1 ; a2 := .01 ; a3 := .001 ; a4 := 0.0001 ;
%   p0 := makefunctionpath(0,5,10,polynomial_function(a,4,x)) ;
%   p1 := addnoisetopath(p0,(0,.001normaldeviate)) ;
%   gdraw p0 ;
%   gdraw p1 plot plotsymbol(1,black,.5) ;
%
%   numeric b[] ;
%   polynomial_fit(p1, b, 4, 1) ;
%   gdraw functionpath(p1,polynomial_function(b,4,x)) ;
%
%   numeric c[] ;
%   linear_fit(p1, c, 1) ;
%   gdraw functionpath(p1,linear_function(c,x)) dashed evenly ;

% a polynomial function:
%
% y = a0 + a1 * x + a2 * x^2 + ... + a[n] * x^n

vardef polynomial_function (suffix $) (expr n, x) =
    (for j=0 upto n : + $[j]*(x**j) endfor) % no ;
enddef ;

% find the determinant of a (n+1)*(n+1) matrix; indices run from 0 to n

vardef det (suffix $) (expr n) =
    hide(
        numeric determinant ; determinant := 1 ;
        save jj ; numeric jj ;
        for k=0 upto n :
            if $[k][k]=0 :
                jj := -1 ;
                for j=0 upto n :
                    if $[k][j]<>0 :
                        jj := j ;
                        exitif true ;
                    fi
                endfor
                if jj<0 :
                    determinant := 0 ;
                    exitif true ;
                fi
                for j=k upto n : % interchange the columns
                    temp := $[j][jj] ;
                    $[j][jj] := $[j][k] ;
                    $[j][k] := temp ;
                endfor
                determinant = -determinant ;
            fi
            exitif determinant=0 ;
            determinant := determinant * $[k][k] ;
            if k<n : % subtract row k from lower rows to get a diagonal matrix
                for j=k+1 upto n:
                    for i=k+1 upto n:
                        $[j][i] := $[j][i]-$[j][k]*$[k][i]/$[k][k] ;
                    endfor
                endfor
            fi
        endfor ;
    )
    determinant % no ;
enddef ;

numeric fit_chi_squared ;

% least-squares fit of a polynomial $ of order n to a path p (unweighted)
%
% reference: P. R. Bevington, "Data Reduction and Error Analysis for the Physical
% Sciences", McGraw-Hill, New York 1969.

vardef polynomial_fit (suffix p, $) (expr n) (text t) =
    if not path p :
        Gerr(p, "Cannot fit--not a path") ;
    elseif length p < n :
        Gerr(p, "Cannot fit--not enough points") ;
    else :
        fit_chi_squared := 0 ;
        % calculate sums of the data
        save sumx, sumy ; numeric sumx[], sumy[] ;
        save w ; numeric w ;
        for i=0 upto 2n :
            sumx[i] := 0 ;
        endfor
        for i=0 upto n :
            sumy[i] := 0 ;
        endfor
        for i=0 upto length p :
            clearxy; z = point i of p ;
            w := if length(t) > 0 : t else : 1 fi ; % weight
            x1 := w ;
            for j=0 upto 2n :
                sumx[j] := sumx[j] + x1 ;
                x1 := x1 * x ;
            endfor
            y1 := y * w ;
            for j=0 upto n :
               sumy[j] := sumy[j] + y1 ;
               y1 := y1 * x ;
            endfor
            fit_chi_squared := fit_chi_squared + y*y*w ;
        endfor
        % construct matrices and calculate the polynomial coefficients
        save m ; numeric m[][] ;
        for j=0 upto n :
            for k=0 upto n :
                m[j][k] := sumx[j+k] ;
            endfor
        endfor
        save delta ; numeric delta ;
        delta := det(m,n) ; % this destroys the matrix m[][], which is OK
        if delta = 0 :
            fit_chi_squared := 0 ;
            for j=0 upto n :
                $[j] := 0 ;
            endfor
        else :
            for i=0 upto n :
                for j=0 upto n :
                    for k=0 upto n :
                        m[j][k] := sumx[j+k] ;
                    endfor
                    m[j][i] := sumy[j] ;
                endfor
                $[i] := det(m,n) / delta ; % matrix m[][] gets destroyed...
            endfor
            for j=0 upto n :
                fit_chi_squared := fit_chi_squared - 2sumy[j]*$[j] ;
                for k=0 upto n :
                    fit_chi_squared := fit_chi_squared + $[j]*$[k]*sumx[j+k] ;
                endfor
            endfor
            % normalize by the number of degrees of freedom
            fit_chi_squared := fit_chi_squared / (length(p) - n) ;
        fi
    fi
enddef ;

% y = a0 + a1 * x
%
% of course a line is just a polynomial of order 1

vardef linear_function (suffix $) (expr x) = polynomial_function($,1,x)   enddef ;
vardef linear_fit   (suffix p, $) (text t) = polynomial_fit(p, $, 1, t) ; enddef ;

% and a constant is polynomial of order 0

vardef constant_function (suffix $) (expr x) = polynomial_function($,0,x)   enddef ;
vardef constant_fit   (suffix p, $) (text t) = polynomial_fit(p, $, 0, t) ; enddef ;

% y = a1 * exp(a0*x)
%
% exp and ln defined in metafun

vardef exponential_function (suffix $) (expr x) = $1*exp($0*x) enddef ;

% since we take a log, this only works for positive ordinates

vardef exponential_fit (suffix p, $) (text t) =
    save a ; numeric a[] ;
    save q ; path q ; % fit to the log of the ordinate
    for i=0 upto length p :
        if ypart(point i of p)>0 :
            augment.q(xpart(point i of p),ln(ypart(point i of p))) ;
        fi
    endfor
    linear_fit(q,a,t) ;
    $0 := a1 ;
    $1 := exp(a0) ;
enddef ;

% y = a1 * x**a0

vardef power_law_function (suffix $) (expr x) = $1*(x**$0) enddef ;

% since we take logs, this only works for positive abcissae and ordinates

vardef power_law_fit (suffix p, $) (text t) =
    save a ; numeric a[] ;
    save q ; path q ; % fit to the logs of the abcissae and ordinates
    for i=0 upto length p :
        if (xpart(point i of p)>0) and (ypart(point i of p)>0) :
            augment.q(ln(xpart(point i of p)),ln(ypart(point i of p))) ;
        fi
    endfor
    linear_fit(q,a,t) ;
    $0 := a1 ;
    $1 := exp(a0) ;
enddef ;

% gaussian: y = a2 * exp(-ln(2)*((x-a0)/a1)^2)
%
% a1 is the hwhm; sigma := a1/sqrt(2ln(2)) or a1/1.17741

numeric lntwo ; lntwo := ln(2) ; % brrr, why not inline it

vardef gaussian_function (suffix $) (expr x) =
    if $1 = 0 :
        if x = $0 : $2 else : 0 fi
    else :
        $2 * exp(-lntwo*(((x-$0)/$1)**2))
    fi
    if known $3 :
        + $3
    fi
enddef ;

% since we take a log, this only works for positive ordinates

vardef gaussian_fit (suffix p, $) (text t) =
    save a ; numeric a[] ;
    save q ; path q ; % fit to the log of the ordinate
    for i=0 upto length p :
        if ypart(point i of p)>0 :
            augment.q(xpart(point i of p), ln(ypart(point i of p))) ;
        fi
    endfor
    polynomial_fit(q,a,2,if t > 0 : ln(t) else : 0 fi) ;
    $1 := sqrt(-lntwo/a2) ;
    $0 := -.5a1/a2 ;
    $2 := exp(a0-.25*a1*a1/a2) ;
    $3 := 0 ; % polynomial_fit will NOT work with a non-zero background!
enddef ;

% lorentzian: y = a2 / (1 + ((x - a0)/a1)^2)

vardef lorentzian_function (suffix $) (expr x) =
    if $1 = 0 :
        if x = $0 : $2 else : 0 fi
    else :
        $2 / (1 + ((x - $0)/$1)**2)
    fi
    if known $3 :
        + $3
    fi
enddef ;

vardef lorentzian_fit (suffix p, $) (text t) =
    save a ; numeric a[] ;
    save q ; path q ; % fit to the inverse of the ordinate
    for i=0 upto length p :
        if ypart(point i of p)<>0 :
            augment.q(xpart(point i of p), 1/ypart(point i of p)) ;
        fi
    endfor
    polynomial_fit(q,a,2,if t <> 0 : 1/(t) else : 0 fi) ;
    $0 := -.5a1/a2 ;
    $2 := 1/(a0-.25a1*a1/a2) ;
    $1 := sqrt((a0-.25a1*a1/a2)/a2) ;
    $3 := 0 ; % polynomial_fit will NOT work with a non-zero background!
enddef ;
