/* Modified from SUPOLAR, Copyright (c) Colorado School of Mines, 2011.*/
/* All rights reserved.                                                */

/* SUPOLAR_PS: $Revision: 1.9 $ ; $Date: Sat, 23 Jul 2016 00:10:26 -0700 */

#include <omp.h>
#include <stdio.h>
#include <gsl/gsl_statistics_float.h>
#include "su.h"
#include "segy.h"
#include "header.h"

/*********************** self documentation *****************************/
char *sdoc[] = {
"                                                                       ",
" SUPOLAR_PS - POLarization analysis of three-component data            ",
"            - to provide P and S arrivals                              ",
"            - Modification of SUPOLAR                                  ",
"                                                                       ",
" supolar_PS <stdin [optional parameters]                               ",
"                                                                       ",
" Required parameters:                                                  ",
"    -ntr (number of traces) needs to be specified in the SU file       ",
"                                                                       ",
" Optional parameters:                                                  ",
"    dt=(from header)  time sampling interval in seconds                ",
"    wl=0.1            correlation window length in seconds             ",
"    win=boxcar        correlation window shape, choose \"boxcar\",     ",
"                      \"hanning\", \"bartlett\", or \"welsh\"          ",
"    file=polar        base of output file name(s)                      ",
"    rl=1              1 = rectilinearity evaluating 2 eigenvalues,     ",
"                      2, 3 = rectilinearity evaluating 3 eigenvalues   ",
"    rlq=1.0           contrast parameter for rectilinearity            ",
"    dir=1             1 = 3 components of direction of polarization    ",
"                          (the only three-component output file)       ",
"    tau=0             1 = global polarization parameter                ",
"    ellip=0           1 = principal, subprincipal, and transverse      ",
"                          ellipticities e21, e31, and e32              ",
"    pln=0             1 = planarity measure                            ",
"    f1=0              1 = flatness or oblateness coefficient           ",
"    l1=0              1 = linearity coefficient                        ",
"    amp=0             1 = amplitude parameters: instantaneous,         ",
"                          quadratic, and eigenresultant ir, qr, and er ",
"    theta=0           1, 2, 3 = incidence angle of principal axis      ",
"    phi=0             1, 2, 3 = horizontal azimuth of principal axis   ",
"    angle=rad         unit of angles theta and phi, choose \"rad\",    ",
"                      \"deg\", or \"gon\"                              ",
"    all=0             1, 2, 3 = set all output flags to that value     ",
"    verbose=0         1 = echo additional information                  ",
"    kwl=5*(1/dt)      kurtosis window length, samples                  ",
"                                                                       ",
" Notes:                                                                ",
"    Three adjacent traces are considered as one three-component        ",
"    dataset.                                                           ",
"    Correct calculation of angles theta and phi requires the first of  ",
"    these traces to be the vertical component, followed by the two     ",
"    horizontal components (e.g. Z, N, E, or Z, inline, crossline).     ",
"    Significant signal energy on Z is necessary to resolve the 180 deg ",
"    ambiguity of phi (options phi=2,3 only).                           ",
"                                                                       ",
"    Each calculated polarization attribute is written into its own     ",
"    SU file. These files get the same base name (set with \"file=\")   ",
"    and the parameter flag as an extension (e.g. polar.rl).            ",
"                                                                       ",
"    In case of a tapered correlation window, the window length wl may  ",
"    have to be increased compared to the boxcar case, because of their ",
"    smaller effective widths (Bartlett, Hanning: 1/2, Welsh: 1/3).     ",
"                                                                       ",
" Range of values:                                                      ",
"    parameter     option  interval                                     ",
"    rl            1, 2    0.0 ... 1.0   (1.0: linear polarization)     ",
"    rl            3      -1.0 ... 1.0                                  ",
"    tau, l1       1       0.0 ... 1.0   (1.0: linear polarization)     ",
"    pln, f1       1       0.0 ... 1.0   (1.0: planar polarization)     ",
"    e21, e31, e32 1       0.0 ... 1.0   (0.0: linear polarization)     ",
"    theta         1      -pi/2... pi/2  rad                            ",
"    theta         2, 3    0.0 ... pi/2  rad                            ",
"    phi           1      -pi/2... pi/2  rad                            ",
"    phi           2      -pi  ... pi    rad   (see notes above)        ",
"    phi           3       0.0 ... 2 pi  rad   (see notes above)        ",
"                                                                       ",
"                                                                       ",
NULL};

/*
 * Author: Nils Maercklin,
 *         GeoForschungsZentrum (GFZ) Potsdam, Germany, 1998-2001.
 *         E-mail: nils@gfz-potsdam.de
 *
 *
 * References:
 *    Jurkevics, A., 1988: Polarization analysis of three-component
 *       array data. Bulletin of the Seismological Society of America,
 *       vol. 78, no. 5.
 *    Kanasewich, E. R., 1981: Time Sequence Analysis in Geophysics.
 *       The University of Alberta Press.
 *    Kanasewich, E. R., 1990: Seismic Noise Attenuation.
 *       Handbook of Geophysical Exploration, Pergamon Press, Oxford.
 *    Meyer, J. H. 1988: First Comparative Results of Integral and
 *       Instantaneous Polarization Attributes for Multicomponent Seismic
 *       Data. Institut Francais du Petrole.
 *    Press, W. H., Teukolsky, S. A., Vetterling, W. T., and Flannery, B. P.
 *       1996: Numerical Recipes in C - The Art of Scientific Computing.
 *       Cambridge University Press, Cambridge.
 *    Samson, J. C., 1973: Description of the Polarisation States of Vector
 *       Processes: Application to ULF Electromagnetic Fields.
 *       Geophysical Journal vol. 34, p. 403-419.
 *    Sheriff, R. E., 1991: Encyclopedic Dictionary of Exploration
 *       Geophysics. 3rd ed., Society of Exploration Geophysicists, Tulsa.
 *
 * Trace header fields accessed: ns, dt
 * Trace header fields modified: none
 */
/**************** end self doc *******************************************/

/* prototypes of functions used internally */
float calc_ellip(float *d, int d1, int d2);
float calc_er(float *d);
float calc_f1(float *d);
float calc_l1(float *d);
float calc_phi(float **v, int opt);
float calc_plan(float *d);
float calc_rl(float d[], float q, int opt);
float calc_tau(float *d);
float calc_theta(float **v, int opt);
float covar(float *data1, float *data2, int istart, int iwl, float *w);
float calc_pfilt(float rl, float theta);
float calc_sfilt(float rl, float theta);
float kurtosiswindow(float *data, float *data_kwl, int it, int kwl, int nt);
void ampparams(float **indata, float *data_ir, float *data_qr, int nt, int iwl);
void calc_dir(float **data3c_dir, float **v, int it);
void calc_window(float *w, int iwl, int iwin);
void fputdata(FILE *fileptr, FILE *headerptr, float outdata[], int nt);
void fputdata3c(FILE *fileptr, FILE *headerptr, float **outdata3c, int nt);

/* window shape identifiers */
#define WBOXCAR 0
#define WBARTLETT 1
#define WHANNING 2
#define WWELSH 3

segy tr;                /* SEG-Y record (trace and header) */

int main(int argc, char **argv)
{
    /* output file pointers */
    FILE *rlfp=NULL, *taufp=NULL, *e21fp=NULL;
    FILE *e31fp=NULL, *e32fp=NULL, *plnfp=NULL;
    FILE *f1fp=NULL, *l1fp=NULL;
    FILE *thetafp=NULL, *phifp=NULL, *dirfp=NULL;
    FILE *erfp=NULL, *irfp=NULL, *qrfp=NULL;
    FILE *headerfp=NULL, *pfiltfp=NULL, *sfiltfp=NULL;
    FILE *nfiltfp=NULL, *efiltfp=NULL;
 //   FILE *pkur=NULL, *skur=NULL;
        /* temporary file for trace headers */
			 /* (one 3C station only) */

    char *file=NULL;         /* base of output file name(s) */
    char *fname=NULL;        /* complete output file name */
    char *angle=NULL;        /* unit used for angles theta and phi */
    char *win=NULL;          /* shape of used time window */
    float fangle=0.0;   /* unit conversion factor applied to angles theta and phi */
    int iwin=0;           /* time window shape identifier */

    /* flags (see selfdoc) */
    int rl,theta,phi,tau,ellip,pln,f1,l1,dir,amp,verbose,all;

    int i,j,icomp;      /* indices for components (in loops) */
    int it;             /* index for time sample in main loop */
    int iwl;            /* correlation window length in samples */
    int nstat;          /* number of 3-component datasets */
    int nt;             /* number of time samples in one trace */
//    int kwl;            /* kurtosis window length in seconds */

    float **data3c;     /* three-component data ([1..3][0..nt-1]) */
    float **a;          /* covariance matrix (a[1..3][1..3]) */
    float **v;          /* eigenvectors of covariance matrix (v[1..3][1..3]) */
    float *d;           /* the corresponding eigenvalues (d[1..3]) */
    float *w;           /* time window weights for correlation window */
    float dt;           /* sampling interval in seconds */
    float rlq;          /* contrast factor of rectilinearity */
    float wl;           /* correlation window length in seconds */

    float *data_e21=NULL;    /* main ellipticity */
    float *data_e31=NULL;    /* second ellipticity */
    float *data_e32=NULL;    /* transverse ellipticity */
    float *data_er=NULL;     /* eigenresultant */
    float *data_f1=NULL;     /* flatness coefficient */
    float *data_ir=NULL;     /* instantaneous resultant */
    float *data_l1=NULL;     /* linearity coefficient */
    float *data_phi=NULL;    /* horizontal azimuth phi */
    float *data_pln=NULL;    /* planarity */
    float *data_qr=NULL;     /* quadratic resultant */
    float *data_rl=NULL;     /* rectilinearity factor */
    float *data_tau=NULL;    /* polarization parameter tau */
    float *data_theta=NULL;  /* inclination angle theta */
    float *data_pfilt=NULL;  /* P (vertical) polarization filter */
    float *data_sfilt=NULL;  /* S (horizontal) polarization filter */
    float *data_zfilt=NULL;  /* Z Filtered Trace */
    float *data_nfilt=NULL;  /* N Filtered Trace */
    float *data_efilt=NULL;  /* E Filtered Trace */
//    float *data_kwl=NULL;   /* Data for Kurtosis Window */
//    float *data_pkur=NULL;   /* Kurtosis detector for P */
//    float *data_skur=NULL;   /* Kurtosis detector for S */
    float **data3c_dir=NULL; /* 3 components of direction of polarization ([1..3][0..nt-1]) */

    /* initialize */
    initargs(argc, argv);
    requestdoc(1);

    /* get info from first trace */
    if(!gettr(&tr)) err("can't get first trace");
    nt = tr.ns;

    /* get parameters ... */
    if (!getparstring("file", &file)) file="polar";
    if (!getparstring("angle", &angle)) angle="rad";
    if (!getparstring("win", &win)) win="boxcar";
    if (!getparfloat("wl", &wl)) wl = 0.1;
    if (!getparfloat("dt", &dt)) dt = ((double) tr.dt)/1000000.0;
    if (!getparfloat("rlq", &rlq)) rlq = 1.0;
    if (!getparint("verbose", &verbose)) verbose = 0;
//    if (!getparint("kwl", &kwl)) kwl = 5 * ((int) 1/dt);

    /* ... and output flags */
    if (!getparint("all", &all))       all = 0;
    if (!getparint("rl", &rl))          rl = (all) ? all : 1;
    if (!getparint("dir", &dir))       dir = (all) ? 1 : 1;
    if (!getparint("theta", &theta)) theta = (all) ? all : 0;
    if (!getparint("phi", &phi))       phi = (all) ? all : 0;
    if (!getparint("tau", &tau))       tau = (all) ? 1 : 0;
    if (!getparint("ellip", &ellip)) ellip = (all) ? 1 : 0;
    if (!getparint("pln", &pln))       pln = (all) ? 1 : 0;
    if (!getparint("f1", &f1))          f1 = (all) ? 1 : 0;
    if (!getparint("l1", &l1))          l1 = (all) ? 1 : 0;
    if (!getparint("amp", &amp))       amp = (all) ? 1 : 0;

    checkpars();

    /* get time window shape */
    if      (STREQ(win, "boxcar"))   iwin=WBOXCAR;
    else if (STREQ(win, "bartlett")) iwin=WBARTLETT;
    else if (STREQ(win, "hanning"))  iwin=WHANNING;
    else if (STREQ(win, "welsh"))    iwin=WWELSH;
    else err("unknown win=%s", win);

    /* get unit conversion factor for angles */
    if      (STREQ(angle, "rad")) fangle=1.0;
    else if (STREQ(angle, "deg")) fangle=180.0/PI;
    else if (STREQ(angle, "gon")) fangle=200.0/PI;
    else err("unknown angle=%s", angle);

    /* convert seconds to samples */
    if (!dt) {
        dt = 0.004;
        warn("dt not set, assuming dt=0.004");
    }

    iwl = NINT(wl/dt);

    /* data validation */
    if (iwl<1) err("wl=%g must be positive", wl);
    if (iwl>nt) err("wl=%g too long for trace", wl);
    if (!strlen(file)) err("file= not set and default overridden");

    /* echo some information */
    if (verbose && (theta || phi)) warn("computing angles in %s", angle);
    if (verbose) warn("%s window length = %d samples\n", win, iwl);

    if (rl && theta) warn("computing filtered phase");

    /* open temporary file for trace headers */
    headerfp = etmpfile();

    /* set filenames and open files */
    fname = malloc( strlen(file)+7 );
    sprintf(fname, "%s.rl", file);    if (rl) rlfp = efopen(fname, "w");
    sprintf(fname, "%s.theta", file); if (theta) thetafp = efopen(fname, "w");
    sprintf(fname, "%s.phi", file);   if (phi) phifp = efopen(fname, "w");
    sprintf(fname, "%s.tau", file);   if (tau) taufp = efopen(fname, "w");
    sprintf(fname, "%s.e21", file);   if (ellip) e21fp = efopen(fname, "w");
    sprintf(fname, "%s.e31", file);   if (ellip) e31fp = efopen(fname, "w");
    sprintf(fname, "%s.e32", file);   if (ellip) e32fp = efopen(fname, "w");
    sprintf(fname, "%s.pln", file);   if (pln) plnfp = efopen(fname, "w");
    sprintf(fname, "%s.f1", file);    if (f1) f1fp = efopen(fname, "w");
    sprintf(fname, "%s.l1", file);    if (l1) l1fp = efopen(fname, "w");
    sprintf(fname, "%s.dir", file);   if (dir) dirfp = efopen(fname, "w");
    sprintf(fname, "%s.er", file);    if (amp) erfp = efopen(fname, "w");
    sprintf(fname, "%s.ir", file);    if (amp) irfp = efopen(fname, "w");
    sprintf(fname, "%s.qr", file);    if (amp) qrfp = efopen(fname, "w");
    sprintf(fname, "%s.pfilt", file); if (rl && theta) pfiltfp = efopen(fname, "w");
    sprintf(fname, "%s.sfilt", file); if (rl && theta) sfiltfp = efopen(fname, "w");
    sprintf(fname, "%s.nfilt", file); if (rl && theta) nfiltfp = efopen(fname, "w");
    sprintf(fname, "%s.efilt", file); if (rl && theta) efiltfp = efopen(fname, "w");
  //  sprintf(fname, "%s.pkur", file); if (rl && theta) pkur = efopen(fname, "w");
  //  sprintf(fname, "%s.skur", file); if (rl && theta) skur = efopen(fname, "w");
    free(fname);

    /* allocate space for input data and analysis matrices */
    /* index ranges used here: data3c[1..3][0..nt-1],      */
    /* a[1..3][1..3], v[1..3][1..3], d[1..3]               */
    data3c = ealloc2float(nt,3); data3c-=1;
    a = ealloc2float(3,3); a[0]-=1; a-=1;
    v = ealloc2float(3,3); v[0]-=1; v-=1;
    d = ealloc1float(3); d-=1;

    /* calculate time window weights */
    w = ealloc1float(iwl);
    memset((void *) w, 0, iwl*FSIZE);
    calc_window(w, iwl, iwin);

    /* allocate and zero out space for output data */
    if (rl) {
        data_rl = ealloc1float(nt);
        memset((void *) data_rl, 0, nt*FSIZE);
    }
    if (theta) {
        data_theta = ealloc1float(nt);
        memset((void *) data_theta, 0, nt*FSIZE);
    }
    if (phi) {
        data_phi = ealloc1float(nt);
        memset((void *) data_phi, 0, nt*FSIZE);
    }
    if (tau) {
        data_tau = ealloc1float(nt);
        memset((void *) data_tau, 0, nt*FSIZE);
    }
    if (ellip) {
        data_e21 = ealloc1float(nt);
        data_e31 = ealloc1float(nt);
        data_e32 = ealloc1float(nt);
        memset((void *) data_e21, 0, nt*FSIZE);
        memset((void *) data_e31, 0, nt*FSIZE);
        memset((void *) data_e32, 0, nt*FSIZE);
    }
    if (pln) {
        data_pln = ealloc1float(nt);
        memset((void *) data_pln, 0, nt*FSIZE);
    }
    if (f1) {
        data_f1 = ealloc1float(nt);
        memset((void *) data_f1, 0, nt*FSIZE);
    }
    if (l1) {
        data_l1 = ealloc1float(nt);
        memset((void *) data_l1, 0, nt*FSIZE);
    }
    if (amp) {
        data_er = ealloc1float(nt);
        data_ir = ealloc1float(nt);
        data_qr = ealloc1float(nt);
        memset((void *) data_er, 0, nt*FSIZE);
        memset((void *) data_ir, 0, nt*FSIZE);
        memset((void *) data_qr, 0, nt*FSIZE);
    }
    if (dir) {
        data3c_dir = ealloc2float(nt,3); data3c_dir-=1;
        for (i=1;i<=3;i++) memset((void *) data3c_dir[i], 0, nt*FSIZE);
    }
    if (rl && theta) {
        data_pfilt = ealloc1float(nt);
        memset((void *) data_pfilt, 0, nt*FSIZE);
        data_sfilt = ealloc1float(nt);
        memset((void *) data_pfilt, 0, nt*FSIZE);
/*        data_3Cfilt = ealloc2float(nt,3);
        for (i=1;i<=3;i++) memset((void *) data_3Cfilt[i], 0, nt*FSIZE); */
        data_zfilt = ealloc1float(nt);
        data_nfilt = ealloc1float(nt);
        data_efilt = ealloc1float(nt);
        memset((void *) data_zfilt, 0, nt*FSIZE);
        memset((void *) data_nfilt, 0, nt*FSIZE);
        memset((void *) data_efilt, 0, nt*FSIZE);

 //       data_pkur = ealloc1float(nt);
//        data_skur = ealloc1float(nt);
//        memset((void *) data_pkur, 0, nt*FSIZE);
//        memset((void *) data_skur, 0, nt*FSIZE);
        /* Allocate data for kurtosis window arrays */
//        data_kwl = ealloc1float(iwl);
//        memset((void *) data_kwl, 0, kwl*FSIZE);
    }

/* ************************ BEGIN CALCULATION ******************************* */

    /* loop over traces */
    icomp=0;
    nstat=0;
    // Need to convert this do while loop into a for loop so as to be easier
    // to parallelize
    warn("Trace Start Time: %d %d %d %d %d", tr.year, tr.day, tr.hour, tr.minute, tr.sec);
    do {
        /* store trace header in temporary file and read data */
        efwrite(&tr, HDRBYTES, 1, headerfp);
        icomp++;
        memcpy((void *)data3c[icomp], (const void *) tr.data, nt*FSIZE);

        /* process 3-component dataset */
        if (icomp==3) {
            erewind(headerfp);
            icomp = 0;
            nstat++;
            if (verbose)
                fprintf(stderr,"%s: analyzing station %d \r",argv[0], nstat);

            /* start loop over samples */

            for (it=iwl/2;it<nt-iwl/2;it++) {
                //warn("Sample %d", it);
                /* covariance matrix */

                for (i=1;i<=3;i++)
                {
                    for (j=i;j<=3;j++)
                    {
                        a[i][j]=a[j][i]=covar(data3c[i], data3c[j], it-iwl/2, iwl, w);
                    }
                }

                /* compute eigenvalues and vectors */
                eig_jacobi(a,d,v,3);
                sort_eigenvalues(d,v,3);

                /* polarization parameters */
                if (rl) data_rl[it]=calc_rl(d,rlq,rl);
                if (theta) data_theta[it]=calc_theta(v, theta) * fangle;
                if (phi) data_phi[it]=calc_phi(v, phi) * fangle;
                if (tau) data_tau[it]=calc_tau(d);
                if (ellip) {
                    data_e21[it]=calc_ellip(d,2,1);
                    data_e31[it]=calc_ellip(d,3,1);
                    data_e32[it]=calc_ellip(d,3,2);
                }
                if (pln) data_pln[it]=calc_plan(d);
                if (f1) data_f1[it]=calc_f1(d);
                if (l1) data_l1[it]=calc_l1(d);
                if (amp) data_er[it]=calc_er(d);
                if (dir) calc_dir(data3c_dir,v,it);
                if (rl && theta)
                {
                    data_zfilt[it] = data3c[1][it] * calc_pfilt(rl, theta);
                    data_nfilt[it] = data3c[2][it] * calc_sfilt(rl, theta);
                    data_efilt[it] = data3c[3][it] * calc_sfilt(rl, theta);
                    data_pfilt[it] = data_zfilt[it];
                    data_sfilt[it] = (data_nfilt[it] + data_efilt[it]) / 2;
//                    data_pkur[it] = kurtosiswindow(data_pfilt,data_kwl,it - kwl/2,kwl,nt);
//                    data_skur[it] = kurtosiswindow(data_sfilt,data_kwl,it - kwl/2,kwl,nt);
               }

            } /* end loop over samples */

            /* compute amplitude parameters */
            if (amp) ampparams(data3c, data_ir, data_qr, nt, iwl);

/* *************************** END CALCULATION ****************************** */
/* ***************************** BEGIN WRITE ******************************** */

            /* write polarization attributes to files */
            if (rl) fputdata(rlfp, headerfp, data_rl, nt);
            if (theta) fputdata(thetafp, headerfp, data_theta, nt);
            if (phi) fputdata(phifp, headerfp, data_phi, nt);
            if (tau) fputdata(taufp, headerfp, data_tau, nt);
            if (ellip) {
                fputdata(e21fp, headerfp, data_e21, nt);
                fputdata(e31fp, headerfp, data_e31, nt);
                fputdata(e32fp, headerfp, data_e32, nt);
            }
            if (pln) fputdata(plnfp, headerfp, data_pln, nt);
            if (f1) fputdata(f1fp, headerfp, data_f1, nt);
            if (l1) fputdata(l1fp, headerfp, data_l1, nt);
            if (amp) {
                fputdata(erfp, headerfp, data_er, nt);
                fputdata(irfp, headerfp, data_ir, nt);
                fputdata(qrfp, headerfp, data_qr, nt);
            }
            if (dir) fputdata3c(dirfp, headerfp, data3c_dir, nt);
            if (rl && theta)
            {
                fputdata(pfiltfp, headerfp, data_pfilt, nt);
                fputdata(sfiltfp, headerfp, data_sfilt, nt);
                fputdata(nfiltfp, headerfp, data_nfilt, nt);
                fputdata(efiltfp, headerfp, data_efilt, nt);
               // fputdata(pkur, headerfp, data_pkur, nt);
               // fputdata(skur, headerfp, data_skur, nt);
            }

/* ****************************** END WRITE ********************************* */

        } /* end of processing three-component dataset */

    } while (gettr(&tr)); /* end loop over traces */

    if (verbose) {
        fprintf(stderr,"\n");
        if (icomp) warn("last %d trace(s) skipped", icomp);
    }

    /* close files */
    efclose(headerfp);
    if (rl) efclose(rlfp);
    if (theta) efclose(thetafp);
    if (phi) efclose(phifp);
    if (tau) efclose(taufp);
    if (ellip) {
        efclose(e21fp);
        efclose(e31fp);
        efclose(e32fp);
    }
    if (pln) efclose(plnfp);
    if (f1) efclose(f1fp);
    if (l1) efclose(l1fp);
    if (amp) {
        efclose(erfp);
        efclose(irfp);
        efclose(qrfp);
    }
    if (dir) efclose(dirfp);
    if (rl && theta) {
        efclose(pfiltfp);
        efclose(sfiltfp);
//        efclose(pkur);
//        efclose(skur);
        }
    return(CWP_Exit());
}

/**********************************************************************/
/* Functions used internally                                          */
/**********************************************************************/

/* calculate time window weights for a smooth time window,        */
/* after Sheriff, 1991, p. 335; tapered windows are suggested by  */
/* Jurkevics, 1988                                                */

void calc_window(float *w, int iwl, int iwin)
{
    int i,j;      /* loop indices within window */
    float m;      /* half of time window (iwl/2) */

    m = (float) iwl/2;

    switch (iwin) {
        case WBOXCAR :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = 1.0;
            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
        case WBARTLETT :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = 1.0 - (fabs((float) i) / m);

            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
        case WHANNING :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = 0.5 + 0.5*cos(PI*fabs((float) i) / m);
            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
        case WWELSH :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = fabs( fabs((float) i) / m - 1.0);
                w[j] *= w[j];
            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
    }
}

/* covariance of two time sequences *data1 and *data2,   */
/* calculated in a window of length iwl,                 */
/* weighted with factors *w and starting at index istart */

float covar(float *data1, float *data2, int istart, int iwl, float *w)
{
    register int i, j;
    float cov=0.0;
    float mean1=0.0;
    float mean2=0.0;

    for (i=istart,j=0;i<(istart+iwl);i++,j++) {
        mean1 += data1[i] * w[j];
        mean2 += data2[i] * w[j];
    }
    mean1 = mean1 / (float) iwl;
    mean2 = mean2 / (float) iwl;

    for (i=istart,j=0;i<(istart+iwl);i++,j++) {
	/* modified by Mega G. Baitoff 6 May 2015 */
        /* cov += (data1[i]-mean1) * (data2[i]-mean2) * w[j]; */
        cov += (data1[i] * w[j] - mean1) * (data2[i] * w[j] - mean2);
    }
    cov = cov / (float) iwl;
    return cov;
}

float kurtosiswindow(float *data, float *data_kwl, int it, int kwl, int nt)
    {
        int i, j;
        for (i=it, j=0; i<(it+kwl); i++, j++)
        {
            data_kwl[j] = data[i];
        }
        return gsl_stats_float_kurtosis(data_kwl,1,kwl);
    }

/**********************************************************************/
/* Functions defining polarization attributes and angles              */
/**********************************************************************/

/* rectilinearity factor rl (3 different definitions) */

float calc_rl(float *d, float rlq, int opt)
{
    float rl;

    if (d[1]!=0.0) {
        switch (opt) {
            default:
                /* case 1: */
                /* rl definition after Kanasewich, 1981 */
                rl = 1 - pow(fabs(d[2]/d[1]), rlq);
                break;
            case 2:
                /* rl definition after Jurkevics, 1988 */
                /* case used by Z. Ross et. al., 2014 */
                rl = 1 - pow( 0.5*(fabs(d[2]/d[1]) + fabs(d[3]/d[1]) ), rlq);
                break;
            case 3:
                /* rl definition after Meyer, 1988 */
                rl = 1 - pow(fabs(d[2]/d[1]) + fabs(d[3]/d[1]), rlq);
                break;
        }
        return rl;
    }
    else {
        return 0.0;
    }
}

/* vertical (incidence) angle theta (different definitions) */
/* assumes a right handed coordinate system starting        */
/* with the vertical component Z (Z=1, N,R=2, E,T=3)        */

float calc_theta(float **v, int opt)
{
    float theta, horiz;

    switch (opt) {
        default:
            /* case 1: */
            /* definition after Kanasewich, 1981 */
            /* interval -pi/2 <= theta <= pi/2 */
            if (v[1][1]) {
                horiz = sqrt( v[2][1]*v[2][1] + v[3][1]*v[3][1] );
                theta = atan( horiz / v[1][1] );
            }
            else {
                theta = 0.0;
            }
            break;
        case 2:
        case 3:
            /* definition after Jurkevics, 1988 */
            /* interval 0.0 <= theta <= pi/2 */
            /* case used by Z. Ross et. al., 2014 */
            theta = acos( fabs(v[1][1]) );
            break;
    }
    return theta;
}

/* horizontal azimuth phi (different definitions)    */
/* assumes a right handed coordinate system starting */
/* with the vertical component Z (Z=1, N,R=2, E,T=3) */

/* Note: The SIGN function is introduced to resolve  */
/* the 180 deg ambiguity by taking the positive      */
/* vertical component of v[1][1] (Jurkevics, 1988).  */
#define VSIGN ( (v[1][1]<0.0) ? -1.0 : 1.0 )

float calc_phi(float **v, int opt)
{
    float phi;

    switch (opt) {
        default:
            /* case 1: */
            /* definition after Kanasewich, 1981 */
            /* interval -pi/2 <= phi <= pi/2 */
            if (v[1][2]) {
                phi = atan( v[3][1] / v[2][1] );
            }
            else {
                phi = (v[3][1]>0.0) ? 0.5*PI : -0.5*PI;
            }
            break;
        case 2:
        case 3:
            /* definitions after Jurkevics, 1988 */
            /* interval -pi <= phi <= pi */
            if (v[2][1]) {
                phi = atan2( v[3][1]*VSIGN, v[2][1]*VSIGN);
            }
            else {
                phi = (v[3][1]>0.0) ? 0.5*PI*VSIGN : -0.5*PI*VSIGN;
            }

            /* interval 0.0 <= phi <= 2*pi */
            if (phi<0.0 && opt==3) phi += 2.0*PI;
            break;
    }
    return phi;
}
#undef VSIGN

/* global polarization parameter tau (Samson, 1973) */

float calc_tau(float *d)
{
    float x1, x2, x3, x4, tau;

    if (d[1]) {
        x1 = pow(( 1 - fabs(d[2]/d[1])), 2.);
        x2 = pow(( 1 - fabs(d[3]/d[1])), 2.);
        x3 = pow(( fabs(d[2]/d[1]) - fabs(d[3]/d[1]) ), 2.);
        x4 = pow(( 1 + fabs(d[2]/d[1]) + fabs(d[3]/d[1])), 2.);
        tau = sqrt( (x1+x2+x3) / (2*x4) );
        return tau;
    }
    else
        return 0.0;
}

/* ellipticities e_ik */

float calc_ellip(float *d, int d1, int d2)
{
    float ellip;

    if (d[d2]) {
        ellip = sqrt( fabs(d[d1]/d[d2]) );
        return ellip;
    }
    else
        return 0.0;
}

/* planarity after Jurkevics, 1988 */

float calc_plan(float *d)
{
    float pln;

    if (d[1]+d[2]) {
        pln = 1.0 - 2.0*d[3] / (d[1] + d[2]);
        return pln;
    }
    else
        return 0.0;
}

/* flatness coefficient f1 after Benhama et. al, 1988 */

float calc_f1(float *d)
{
    float f1,x1,x2;

    x1 = 3.0 * calc_ellip(d,3,1);
    x2 = 1.0 + calc_ellip(d,2,1) + calc_ellip(d,3,1);
    f1 = 1.0 - x1 / x2;
    return f1;
}

/* linearity coefficient l1 */

float calc_l1(float *d)
{
    float l1,x1,x2;

    x1 = 3. * ( calc_ellip(d,2,1) + calc_ellip(d,3,1) );
    x2 = 2. * ( 1 + calc_ellip(d,2,1) + calc_ellip(d,3,1));
    l1 = 1. - x1 / x2;
    return l1;
}

/* direction cosines or directivity functions (3 components) */

void calc_dir(float **data3c_dir, float **v, int it)
{
    int i;
    for (i=1;i<=3;i++) {
        data3c_dir[i][it]=v[i][1];
    }
}

/* amplitude parameters */

/* eigenresultant */

float calc_er(float *d)
{
    float er;
    er = sqrt(fabs(d[1]));
    return er;
}

/* instantaneous and quadratic resultant (Meyer, 1988) */

void ampparams(float **indata, float *data_ir, float *data_qr, int nt, int iwl)
{
    int i,it;
    float sqrsum;

    for (it=0;it<nt;it++) {
        /* instantaneous resultant */
        data_ir[it] = sqrt(indata[1][it]*indata[1][it]+ \
            indata[2][it]*indata[2][it]+indata[3][it]*indata[3][it]);

        /* quadratic resultant */
        if ( (it >= iwl/2) && (it < nt-iwl/2) ) {
            sqrsum=0.0;
            for (i=it-iwl/2;i<(it+iwl);i++) {
                sqrsum += indata[1][i]*indata[1][i]+ \
                    indata[2][i]*indata[2][i]+indata[3][i]*indata[3][i];
            }
            data_qr[it] = sqrsum / (float) iwl;
        }
    }
}

/* P, S detection parameters */

/* P-Wave energy filter */
float calc_pfilt(float rl, float theta)
{
    float pfilt;
    pfilt = rl*cos(theta);
    return pfilt;
}

/* S-Wave energy filter */
float calc_sfilt(float rl, float theta)
{
    float sfilt;
    sfilt = rl * (1 - cos(theta));
    return sfilt;
}

/**********************************************************************/
/* Functions for data output                                          */
/**********************************************************************/

/* write one-component data into file */

void fputdata(FILE *fileptr, FILE *headerptr, float *outdata, int nt)
{
    efread(&tr, 1, HDRBYTES, headerptr);
    erewind(headerptr);

    memcpy((void *)tr.data, (const void *) outdata, nt*FSIZE);

    fputtr(fileptr, &tr);
}

/* write three-component data into file */

void fputdata3c(FILE *fileptr, FILE *headerptr, float **outdata3c, int nt)
{
    int i;

    for(i=1;i<=3;i++) {
        efread(&tr, 1, HDRBYTES, headerptr);

        memcpy((void *)tr.data, (const void *) outdata3c[i], nt*FSIZE);

        fputtr(fileptr, &tr);
    }
    erewind(headerptr);
}

/* END OF FILE */


#include <cblas.h>
#include <stdlib.h>
#include <string.h>
#include "loss.h"

void expit(const double *x, double *out, int x_len) {
    for (int i = 0; i < x_len; i++) {
        if (x[i] > 0) {
            out[i] = 1. / (1. + exp(-x[i]));
        } else {
            out[i] = 1. - 1. / (1. + exp(x[i]));
        }
    }
}

void log_logistic(const double *x, double *out, int x_len) {
    for (int i = 0; i < x_len; i++) {
        if (x[i] > 0.0) {
            out[i] = -log(1.0 + exp(-x[i]));
        } else {
            out[i] = x[i] - log(1.0 + exp(x[i]));
        }
    }
}

void logistic(const double *x, double *out, int x_len) {
    for (int i = 0; i < x_len; i++) {
        if (x[i] > 0) {
            out[i] = 1. / (1. + exp(-x[i]));
        } else {
            out[i] = 1. - 1. / (1. + exp(x[i]));
        }
    }
}

double log_sum_exp(const double *x, int x_len) {
    double max_x = x[0], out = 0.0;
    for (int i = 1; i < x_len; i++) {
        if (x[i] > max_x) {
            max_x = x[i];
        }
    }
    for (int i = 0; i < x_len; i++) {
        out += exp(x[i] - max_x);
    }
    return max_x + log(out);
}

void logistic_loss_grad(const double *w,
                        const double *x_tr,
                        const double *y_tr,
                        double *loss_grad,
                        double eta,
                        int n_samples,
                        int n_features) {
    for (int i = 0; i < n_features + 1; i++) {
        if (isnan(w[i])) {
            printf("%f\n", w[i]);
            printf("warning: loss grad error!\n");
            break;
        }
    }
    int i, n = n_samples, p = n_features;
    double intercept = w[p], sum_z0 = 0.0;
    loss_grad[0] = 0.0;
    double *yz = malloc(sizeof(double) * n);
    double *z0 = malloc(sizeof(double) * n);
    double *logistic = malloc(sizeof(double) * n);
    for (i = 0; i < n; i++) { /** calculate yz */
        yz[i] = intercept;
    }
    //x_tr^T*w+
    cblas_dgemv(CblasRowMajor, CblasNoTrans, n, p, 1., x_tr, p, w, 1, 1., yz, 1);
    for (i = 0; i < n; i++) {
        yz[i] *= y_tr[i];
    }
    expit(yz, z0, n); // calculate z0 and final intercept
    /** calculate logistic logistic[i] = 1/(1+exp(-y[i]*(xi^T*w+c)))*/
    log_logistic(yz, logistic, n);
    /** calculate loss of data fitting part*/
    for (i = 0; i < n; i++) {
        z0[i] = (z0[i] - 1.) * y_tr[i];
        sum_z0 += z0[i];
        loss_grad[0] -= logistic[i];
    }
    /**calculate loss of regularization part (it does not have intercept)*/
    loss_grad[0] += 0.5 * eta * cblas_ddot(p, w, 1, w, 1);
    /** calculate gradient of coefficients*/
    memcpy(loss_grad + 1, w, sizeof(double) * p);
    /** x^T*z0 + eta*w, where z0[i]=(logistic[i] - 1.)*yi*/
    cblas_dgemv(CblasRowMajor, CblasTrans,
                n, p, 1., x_tr, p, z0, 1, eta, loss_grad + 1, 1);
    /** calculate gradient of intercept part*/
    loss_grad[p + 1] = sum_z0; // intercept part
    free(logistic);
    free(z0);
    free(yz);
}

void logistic_loss_grad_sparse(const double *w,
                               const double *x_tr,
                               const double *y_tr,
                               double *loss_grad,
                               double eta,
                               int n_samples,
                               int n_features) {
    //sub_p is number of nonzeros in w.
    int i, n = n_samples, p = n_features, sub_p = 0;
    double intercept = w[p], sum_z0 = 0.0;
    loss_grad[0] = 0.0;
    double *yz = malloc(sizeof(double) * n);
    double *z0 = malloc(sizeof(double) * n);
    double *logistic = malloc(sizeof(double) * n);
    for (i = 0; i < n; i++) { /** calculate yz */
        yz[i] = intercept;
    }
    double nonzero_w[p];
    int nonzero_w_ind[p];
    for (i = 0; i < p; i++) {
        if (w[i] != 0.0) {
            nonzero_w_ind[sub_p] = i;
            nonzero_w[sub_p] = w[i];
            sub_p++;
        }
    }
    //x_tr^T*w+ to use the sparsity of w
    for (i = 0; i < n; i++) {
        double tmp_val = 0.0;
        for (int j = 0; j < sub_p; j++) {
            tmp_val += x_tr[i * p + nonzero_w_ind[j]] * nonzero_w[j];
        }
        yz[i] = yz[i] + tmp_val;
    }
    for (i = 0; i < n; i++) {
        yz[i] *= y_tr[i];
    }
    expit(yz, z0, n); // calculate z0 and final intercept
    /** calculate logistic logistic[i] = 1/(1+exp(-y[i]*(xi^T*w+c)))*/
    log_logistic(yz, logistic, n);
    /** calculate loss of data fitting part*/
    for (i = 0; i < n; i++) {
        z0[i] = (z0[i] - 1.) * y_tr[i];
        sum_z0 += z0[i];
        loss_grad[0] -= logistic[i];
    }
    /**calculate loss of regularization part (it does not have intercept)*/
    loss_grad[0] += 0.5 * eta * cblas_ddot(p, w, 1, w, 1);
    /** calculate gradient of coefficients*/
    memcpy(loss_grad + 1, w, sizeof(double) * p);
    /** x^T*z0 + eta*w, where z0[i]=(logistic[i] - 1.)*yi*/
    cblas_dgemv(CblasRowMajor, CblasTrans,
                n, p, 1., x_tr, p, z0, 1, eta, loss_grad + 1, 1);
    /** calculate gradient of intercept part*/
    loss_grad[p + 1] = sum_z0; // intercept part
    free(logistic);
    free(z0);
    free(yz);
}

void logistic_predict(const double *x_te,
                      const double *wt,
                      double *pred_prob,
                      double *pred_label,
                      double threshold,
                      int n,
                      int p) {
    openblas_set_num_threads(1);
    int i;
    cblas_dgemv(CblasRowMajor, CblasNoTrans,
                n, p, 1., x_te, p, wt, 1, 0., pred_prob, 1);
    for (i = 0; i < n; i++) {
        pred_prob[i] += wt[p]; // intercept
    }
    expit(pred_prob, pred_prob, n);
    for (i = 0; i < n; i++) {
        if (pred_prob[i] >= threshold) {
            pred_label[i] = 1.;
        } else {
            pred_label[i] = -1.;
        }
    }
}

void least_square_loss_grad(const double *w,
                            const double *x_tr,
                            const double *y_tr,
                            double *loss_grad,
                            double eta,
                            int n_samples,
                            int n_features) {
    for (int i = 0; i < n_features + 1; i++) {
        if (isnan(w[i])) {
            printf("%f\n", w[i]);
            printf("warning: loss grad error!\n");
            break;
        }
    }
    int i, n = n_samples, p = n_features;
    double *yz = malloc(sizeof(double) * n);
    double *w0 = malloc(sizeof(double) * n);
    cblas_dcopy(n, y_tr, 1, yz, 1); // y_tr --> yz
    cblas_dscal(p + 2, 0.0, loss_grad, 1);
    cblas_daxpy(p, eta, w, 1, loss_grad + 1, 1); // l2_lambda*w --> loss_grad
    //Order,TransA, M, N, alpha, A, lda, x, incX, beta, Y, incY
    // Y<-alpha*AX + beta*Y, where A=MxN, Y=N
    cblas_dgemv(CblasRowMajor, CblasNoTrans,
                n, p, 1., x_tr, p, w, 1, -1., yz, 1); //Xw - y
    for (i = 0; i < n; i++) { w0[i] = w[p]; }
    cblas_daxpy(n, 1., w0, 1, yz, 1);
    loss_grad[0] = cblas_ddot(n, yz, 1, yz, 1) / (2. * n);
    loss_grad[0] += (.5 * eta) * cblas_ddot(p, w, 1, w, 1); //loss
    cblas_dgemv(CblasRowMajor, CblasTrans,
                n, p, 1. / n, x_tr, p, yz, 1, 1., loss_grad + 1, 1);
    for (i = 0; i < n; i++) {
        w0[i] = w[p];
        loss_grad[p + 1] = yz[i];
    }
    loss_grad[p + 1] /= n;
    free(w0), free(yz);
}

void least_square_predict(const double *x_te,
                          const double *wt,
                          double *pred_prob,
                          double *pred_label,
                          double threshold,
                          int n_samples,
                          int p_features) {
    openblas_set_num_threads(1);
    int i, n = n_samples, p = p_features;
    cblas_dgemv(CblasRowMajor, CblasNoTrans,
                n, p, 1., x_te, p, wt, 1, 0., pred_prob, 1);
    for (i = 0; i < n; i++) {
        pred_prob[i] += wt[p]; // intercept
    }
    // this is not useful.
    for (i = 0; i < n; i++) {
        if (pred_prob[i] >= threshold) {
            pred_label[i] = 1.;
        } else {
            pred_label[i] = -1.;
        }
    }
}


#include <jni.h>
#include <assert.h>
#include <lapacke.h>

extern void dtrmv_(char *uplo, char *trans, char *diag, int *n, double *A,
                   int *lda, double *x, int *incx);

#define jniRowMajor 101
#define jniColMajor 102

#define jniNoTrans 111
#define jniTrans   112
#define jniConjTrans    113

#define jniUpper   121
#define jniLower   122

#define jniNonUnit 131
#define jniUnit    132

#define jniLeft    141
#define jniRight   142

JNIEXPORT jint Java_JAMAJni_EigenvalueDecomposition_dgeev
(JNIEnv *env, jclass klass, jint matrix_layout, jchar jobvl, jchar jobvr,
 jint n, jdoubleArray a, jint lda, jdoubleArray wr, jdoubleArray wi,
 jdoubleArray vl, jint ldvl, jdoubleArray vr, jint ldvr){

    double *aElems, *wrElems, *wiElems, *vlElems, *vrElems;
    int info;

    aElems = (*env)-> GetDoubleArrayElements (env, a, NULL);
    wrElems = (*env)-> GetDoubleArrayElements (env, wr, NULL);
    wiElems = (*env)-> GetDoubleArrayElements (env, wi, NULL);
    vlElems = (*env)-> GetDoubleArrayElements (env, vl, NULL);
    vrElems = (*env)-> GetDoubleArrayElements (env, vr, NULL);

    assert(aElems && wrElems && wiElems && vlElems && vrElems);

    info = LAPACKE_dgeev((int) matrix_layout, (char) jobvl, (char) jobvr, (lapack_int) n, aElems, (lapack_int) lda, wrElems, wiElems, vlElems, (lapack_int) ldvl, vrElems, (lapack_int) ldvr);

    (*env)-> ReleaseDoubleArrayElements (env, a, aElems, 0);
    (*env)-> ReleaseDoubleArrayElements (env, vl, vlElems, 0);
    (*env)-> ReleaseDoubleArrayElements (env, vr, vrElems, 0);
    (*env)-> ReleaseDoubleArrayElements (env, wr, wrElems, 0);
    (*env)-> ReleaseDoubleArrayElements (env, wi, wiElems, 0);

    return info;

}

JNIEXPORT void JNICALL Java_JAMAJni_EigenvalueDecomposition_dsyev
(JNIEnv *env, jclass obj, jint layout, jchar jobz, jchar uplo, jint n,
 jdoubleArray ja, jint lda, jdoubleArray jw)
{
    jdouble *a = (*env)-> GetDoubleArrayElements (env, ja, NULL);
    jdouble *w = (*env)-> GetDoubleArrayElements (env, jw, NULL);

    assert(a && w);

    LAPACKE_dsyev((int) layout, (char) jobz, (char) uplo, (lapack_int) n, a, (lapack_int) lda, w);

    (*env)-> ReleaseDoubleArrayElements (env, ja, a, 0);
    (*env)-> ReleaseDoubleArrayElements (env, jw, w, 0);
}

JNIEXPORT void JNICALL Java_JAMAJni_EigenvalueDecomposition_dtrsm
(JNIEnv *env, jclass obj, jint layout, jchar jobz, jchar uplo, jint n,
 jdoubleArray ja, jint lda, jdoubleArray jw)
{
    jdouble *a = (*env)-> GetDoubleArrayElements (env, ja, NULL);
    jdouble *w = (*env)-> GetDoubleArrayElements (env, jw, NULL);

    assert(a && w);

    LAPACKE_dsyev((int) layout, (char) jobz, (char) uplo, (lapack_int) n, a, (lapack_int) lda, w);

    (*env)-> ReleaseDoubleArrayElements (env, ja, a, 0);
    (*env)-> ReleaseDoubleArrayElements (env, jw, w, 0);
}


#include <math.h>
#include <stdio.h>
#include <gsl/gsl_rng.h>
#include "matrix_operations.c"
#include "ellipse/parametric_function_roots.c"
#include "polynomial.c"

void characteristic_ellipse_matrix(double *X, double *R, double phi, double exponent)
{

    // rotation matrix
    double Q[2][2] = {{ 0 }};
    double Qt[2][2] = {{ 0 }};
    Q[0][0] = cos(phi);
    Q[0][1] = sin(phi);
    Q[1][0] = -sin(phi);
    Q[1][1] = cos(phi);
    matrix_transpose(&Qt[0][0], &Q[0][0], 2, 2);

    // radii matrix
    double O[2][2] = {{ 0 }};
    double diag_vals[2];
    diag_vals[0] = pow(R[0], exponent * -2.0);
    diag_vals[1] = pow(R[1], exponent * -2.0);
    set_diagonal(&O[0][0], diag_vals, 2, 2);

    // characteristic ellipse matrix
    double X_temp[2][2] = {{ 0 }};
    matrix_multiply(&X_temp[0][0], &O[0][0], &Q[0][0], 2, 2, 2);
    matrix_multiply(X, &Qt[0][0], &X_temp[0][0], 2, 2, 2);

}

double ellipse_overlap(double *rA, double *radiiA, double phiA, double *rB, double *radiiB, double phiB)
{

    // find XA^(-1) and XB^(1/2)
    double XA[2][2] = {{ 0 }};
    double XB[2][2] = {{ 0 }};
    characteristic_ellipse_matrix(&XA[0][0], &radiiA[0], phiA, -1.0);
    characteristic_ellipse_matrix(&XB[0][0], &radiiB[0], phiB, 0.5);

    // find A_AB
    double A_AB[2][2] = {{ 0 }};
    double A_temp[2][2] = {{ 0 }};
    matrix_multiply(&A_temp[0][0], &XA[0][0], &XB[0][0], 2, 2, 2);
    matrix_multiply(&A_AB[0][0], &XB[0][0], &A_temp[0][0], 2, 2, 2);

    // find r_AB
    double rAB[2];
    rAB[0] = rB[0] - rA[0];
    rAB[1] = rB[1] - rA[1];

    // find a_AB
    double a_AB[2];
    matrix_multiply(&a_AB[0], &XB[0][0], &rAB[0], 2, 2, 1);

    // extract elements of the matrix A_AB and a_AB
    double a11, a12, a21, a22;
    double b1, b2;
    a11 = A_AB[0][0];
    a12 = A_AB[0][1];
    a21 = A_AB[1][0];
    a22 = A_AB[1][1];
    b1 = a_AB[0];
    b2 = a_AB[1];

    // find coefficients for the parametric polynomial derivative used to find max
    double h[5];
    double z[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
    size_t n = 5;
    h[0] = h0(a11, a12, a21, a22, b1, b2);
    h[1] = h1(a11, a12, a21, a22, b1, b2);
    h[2] = h2(a11, a12, a21, a22, b1, b2);
    h[3] = h3(a11, a12, a21, a22, b1, b2);
    h[4] = h4(a11, a12, a21, a22, b1, b2);

    // find roots
    size_t m;
    m = find_roots(&z[0], &h[0], n);

    double F;
    int i;

    if (m > 1)
    {
        if (f(0, a11, a12, a21, a22, b1, b2) > f(1, a11, a12, a21, a22, b1, b2))
        {

           F = f(0, a11, a12, a21, a22, b1, b2);

        }
        else
        {

           F = f(1, a11, a12, a21, a22, b1, b2);

        }

        for(i=0; i<=m-2; i++)
        {

           if (( z[2 * i + 1] == 0 ) & (z[2 * i] > 0) & (z[2 * i] < 1))
           {

              F = f(z[2 * i], a11, a12, a21, a22, b1, b2);

           }

        }
    }
    else
    {
        F = 0.;
    }

    return F;

}

double container_square_overlap_potential(double *rA, double *radiiA, double phiA)
{

    double rB[2];
    double radiiB[2];
    double phiB;

    double top = 0;
    double bottom = 0;
    double left = 0;
    double right = 0;

    // top
    rB[0] = 0.5;
    rB[1] = 2.0;
    radiiB[0] = INFINITY;
    radiiB[1] = 1.0;
    phiB = 0.0;
    top = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    // bottom
    rB[0] = 0.5;
    rB[1] = -1.0;
    radiiB[0] = INFINITY;
    radiiB[1] = 1.0;
    phiB = 0.0;
    bottom = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    // left
    rB[0] = -1.0;
    rB[1] = 0.5;
    radiiB[0] = 1.0;
    radiiB[1] = INFINITY;
    phiB = 0.0;
    left = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    // right
    rB[0] = 2.0;
    rB[1] = 0.5;
    radiiB[0] = 1.0;
    radiiB[1] = INFINITY;
    phiB = 0.0;
    right = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    return fminf(top, fminf(bottom, fminf(left, right)));

}

size_t rsa_align_square(double *x, double *y,
    size_t npoints, double *radius, double phi, int step_limit,
    unsigned long randSeed)
{

    // Setup GSL random number generator
    const gsl_rng_type * T;
    gsl_rng * r;
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    // Set the seed
    // srand ( time(NULL) );
    // unsigned long randSeed = rand();
    gsl_rng_set(r, randSeed);

    // arrays for overlap_potential functions
    double rA[2];
    double rB[2];
    double radiiA[2];
    double radiiB[2];
    double F;
    double xn = 0.;
    double yn = 0.;

    radiiA[0] = radius[0];
    radiiA[1] = radius[1];

    // Set the initial position
    double C;
    C = fmaxf(radius[0], radius[1]);

    F = 0;
    while (F < 1.)
    {

        xn = gsl_rng_uniform (r);
        yn = gsl_rng_uniform (r);

        rA[0] = xn;
        rA[1] = yn;

        F = container_square_overlap_potential(&rA[0], &radiiA[0], phi);

    }

    x[0] = xn;
    y[0] = yn;

    size_t valid_pts;
    int k, flag, step;

    step = 0;
    valid_pts = 1;
    while ((valid_pts < npoints) & (step < step_limit))
    {

        // Generate new ellipse inside the container
        F = 0;
        while (F < 1.)
        {

            xn = gsl_rng_uniform (r);
            yn = gsl_rng_uniform (r);

            rA[0] = xn;
            rA[1] = yn;

            F = container_square_overlap_potential(&rA[0], &radiiA[0], phi);

        }

        // Determine if new ellipse overlaps with existing ellipses
        flag = 1;
        for (k = 0; k < valid_pts; k++)
        {

            rA[0] = x[k];
            rA[1] = y[k];
            rB[0] = xn;
            rB[1] = yn;
            radiiA[0] = radius[0];
            radiiA[1] = radius[1];
            radiiB[0] = radius[0];
            radiiB[1] = radius[1];

            F = ellipse_overlap(&rA[0], &radiiA[0], phi, &rB[0], &radiiB[0], phi);

            if (F < 1.0)
            {

                flag = 0;
                break;

            }
        }
        if (flag == 1)
        {

           x[valid_pts] = xn;
           y[valid_pts] = yn;
           valid_pts += 1;

        }

        step += 1;

    }

    gsl_rng_free (r);

    return valid_pts;

}

size_t rsa_square(double *x, double *y,
    size_t npoints, double *radius, double *phi, int step_limit,
    unsigned long randSeed)
{

    // Setup GSL random number generator
    const gsl_rng_type * T;
    gsl_rng * r;
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    // Set the seed
    // srand ( time(NULL) );
    // unsigned long randSeed = rand();
    gsl_rng_set(r, randSeed);

    // arrays for overlap_potential functions
    double rA[2];
    double rB[2];
    double radiiA[2];
    double radiiB[2];
    double phiA;
    double phiB;
    double F;
    double xn = 0;
    double yn = 0;
    double phin = 0;

    radiiA[0] = radius[0];
    radiiA[1] = radius[1];

    // Set the initial position
    double C;
    C = fmaxf(radius[0], radius[1]);

    F = 0;
    while (F < 1.)
    {

        xn = gsl_rng_uniform (r);
        yn = gsl_rng_uniform (r);

        rA[0] = xn;
        rA[1] = yn;
        phiA = 2 * M_PI * gsl_rng_uniform (r);

        F = container_square_overlap_potential(&rA[0], &radiiA[0], phiA);

    }

    x[0] = xn;
    y[0] = yn;
    phi[0] = phiA;

    size_t valid_pts;
    int k, flag, step;

    step = 0;
    valid_pts = 1;
    while ((valid_pts < npoints) & (step < step_limit))
    {

        // Generate new ellipse inside the container
        F = 0;
        while (F < 1.)
        {

            xn = gsl_rng_uniform (r);
            yn = gsl_rng_uniform (r);
            phin = 2 * M_PI * gsl_rng_uniform (r);

            rA[0] = xn;
            rA[1] = yn;
            phiA = phin;

            F = container_square_overlap_potential(&rA[0], &radiiA[0], phiA);

        }

        // Determine if new ellipse overlaps with existing ellipses
        flag = 1;
        for (k = 0; k < valid_pts; k++)
        {

            rA[0] = x[k];
            rA[1] = y[k];
            phiA = phi[k];

            rB[0] = xn;
            rB[1] = yn;
            phiB = phin;

            radiiA[0] = radius[0];
            radiiA[1] = radius[1];
            radiiB[0] = radius[0];
            radiiB[1] = radius[1];

            F = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

            if (F < 1.0)
            {

                flag = 0;
                break;

            }
        }
        if (flag == 1)
        {

           x[valid_pts] = xn;
           y[valid_pts] = yn;
           phi[valid_pts] = phin;
           valid_pts += 1;

        }

        step += 1;

    }

    gsl_rng_free (r);

    return valid_pts;

}


/* interpolation/akima.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/* Author:  G. Jungman
 */
#include <config.h>
#include <stdlib.h>
#include <math.h>
#include <gsl/gsl_errno.h>
#include "integ_eval.h"
#include <gsl/gsl_interp.h>

typedef struct
{
  double * b;
  double * c;
  double * d;
  double * _m;
} akima_state_t;

/* common creation */
static void *
akima_alloc (size_t size)
{
  akima_state_t *state = (akima_state_t *) malloc (sizeof (akima_state_t));

  if (state == NULL)
    {
      GSL_ERROR_NULL("failed to allocate space for state", GSL_ENOMEM);
    }

  state->b = (double *) malloc (size * sizeof (double));

  if (state->b == NULL)
    {
      free (state);
      GSL_ERROR_NULL("failed to allocate space for b", GSL_ENOMEM);
    }

  state->c = (double *) malloc (size * sizeof (double));

  if (state->c == NULL)
    {
      free (state->b);
      free (state);
      GSL_ERROR_NULL("failed to allocate space for c", GSL_ENOMEM);
    }

  state->d = (double *) malloc (size * sizeof (double));

  if (state->d == NULL)
    {
      free (state->c);
      free (state->b);
      free (state);
      GSL_ERROR_NULL("failed to allocate space for d", GSL_ENOMEM);
    }

  state->_m = (double *) malloc ((size + 4) * sizeof (double));

  if (state->_m == NULL)
    {
      free (state->d);
      free (state->c);
      free (state->b);
      free (state);
      GSL_ERROR_NULL("failed to allocate space for _m", GSL_ENOMEM);
    }

  return state;
}

/* common calculation */
static void
akima_calc (const double x_array[], double b[],  double c[],  double d[], size_t size, double m[])
{
  size_t i;

  for (i = 0; i < (size - 1); i++)
    {
      const double NE = fabs (m[i + 1] - m[i]) + fabs (m[i - 1] - m[i - 2]);
      if (NE == 0.0)
        {
          b[i] = m[i];
          c[i] = 0.0;
          d[i] = 0.0;
        }
      else
        {
          const double h_i = x_array[i + 1] - x_array[i];
          const double NE_next = fabs (m[i + 2] - m[i + 1]) + fabs (m[i] - m[i - 1]);
          const double alpha_i = fabs (m[i - 1] - m[i - 2]) / NE;
          double alpha_ip1;
          double tL_ip1;
          if (NE_next == 0.0)
            {
              tL_ip1 = m[i];
            }
          else
            {
              alpha_ip1 = fabs (m[i] - m[i - 1]) / NE_next;
              tL_ip1 = (1.0 - alpha_ip1) * m[i] + alpha_ip1 * m[i + 1];
            }
          b[i] = (1.0 - alpha_i) * m[i - 1] + alpha_i * m[i];
          c[i] = (3.0 * m[i] - 2.0 * b[i] - tL_ip1) / h_i;
          d[i] = (b[i] + tL_ip1 - 2.0 * m[i]) / (h_i * h_i);
        }
    }
}

static int
akima_init (void * vstate, const double x_array[], const double y_array[],
            size_t size)
{
  akima_state_t *state = (akima_state_t *) vstate;

  double * m = state->_m + 2; /* offset so we can address the -1,-2
                                 components */

  size_t i;
  for (i = 0; i <= size - 2; i++)
    {
      m[i] = (y_array[i + 1] - y_array[i]) / (x_array[i + 1] - x_array[i]);
    }

  /* non-periodic boundary conditions */
  m[-2] = 3.0 * m[0] - 2.0 * m[1];
  m[-1] = 2.0 * m[0] - m[1];
  m[size - 1] = 2.0 * m[size - 2] - m[size - 3];
  m[size] = 3.0 * m[size - 2] - 2.0 * m[size - 3];

  akima_calc (x_array, state->b, state->c, state->d, size, m);

  return GSL_SUCCESS;
}

static int
akima_init_periodic (void * vstate,
                     const double x_array[],
                     const double y_array[],
                     size_t size)
{
  akima_state_t *state = (akima_state_t *) vstate;

  double * m = state->_m + 2; /* offset so we can address the -1,-2
                                 components */

  size_t i;
  for (i = 0; i <= size - 2; i++)
    {
      m[i] = (y_array[i + 1] - y_array[i]) / (x_array[i + 1] - x_array[i]);
    }

  /* periodic boundary conditions */
  m[-2] = m[size - 1 - 2];
  m[-1] = m[size - 1 - 1];
  m[size - 1] = m[0];
  m[size] = m[1];

  akima_calc (x_array, state->b, state->c, state->d, size, m);

  return GSL_SUCCESS;
}

static void
akima_free (void * vstate)
{
  akima_state_t *state = (akima_state_t *) vstate;

  free (state->b);
  free (state->c);
  free (state->d);
  free (state->_m);
  free (state);
}

static
int
akima_eval (const void * vstate,
            const double x_array[], const double y_array[], size_t size,
            double x,
            gsl_interp_accel * a,
            double *y)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t index;

  if (a != 0)
    {
      index = gsl_interp_accel_find (a, x_array, size, x);
    }
  else
    {
      index = gsl_interp_bsearch (x_array, x, 0, size - 1);
    }

  /* evaluate */
  {
    const double x_lo = x_array[index];
    const double delx = x - x_lo;
    const double b = state->b[index];
    const double c = state->c[index];
    const double d = state->d[index];
    *y = y_array[index] + delx * (b + delx * (c + d * delx));
    return GSL_SUCCESS;
  }
}

static int
akima_eval_deriv (const void * vstate,
                  const double x_array[], const double y_array[], size_t size,
                  double x,
                  gsl_interp_accel * a,
                  double *dydx)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t index;

  DISCARD_POINTER(y_array); /* prevent warning about unused parameter */

  if (a != 0)
    {
      index = gsl_interp_accel_find (a, x_array, size, x);
    }
  else
    {
      index = gsl_interp_bsearch (x_array, x, 0, size - 1);
    }

  /* evaluate */
  {
    double x_lo = x_array[index];
    double delx = x - x_lo;
    double b = state->b[index];
    double c = state->c[index];
    double d = state->d[index];
    *dydx = b + delx * (2.0 * c + 3.0 * d * delx);
    return GSL_SUCCESS;
  }
}

static
int
akima_eval_deriv2 (const void * vstate,
                   const double x_array[], const double y_array[], size_t size,
                   double x,
                   gsl_interp_accel * a,
                   double *y_pp)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t index;

  DISCARD_POINTER(y_array); /* prevent warning about unused parameter */

  if (a != 0)
    {
      index = gsl_interp_accel_find (a, x_array, size, x);
    }
  else
    {
      index = gsl_interp_bsearch (x_array, x, 0, size - 1);
    }

  /* evaluate */
  {
    const double x_lo = x_array[index];
    const double delx = x - x_lo;
    const double c = state->c[index];
    const double d = state->d[index];
    *y_pp = 2.0 * c + 6.0 * d * delx;
    return GSL_SUCCESS;
  }
}

static
int
akima_eval_integ (const void * vstate,
                  const double x_array[], const double y_array[], size_t size,
                  gsl_interp_accel * acc,
                  double a, double b,
                  double * result)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t i, index_a, index_b;

  if (acc != 0)
    {
      index_a = gsl_interp_accel_find (acc, x_array, size, a);
      index_b = gsl_interp_accel_find (acc, x_array, size, b);
    }
  else
    {
      index_a = gsl_interp_bsearch (x_array, a, 0, size - 1);
      index_b = gsl_interp_bsearch (x_array, b, 0, size - 1);
    }

  *result = 0.0;

  /* interior intervals */

  for(i=index_a; i<=index_b; i++) {
    const double x_hi = x_array[i + 1];
    const double x_lo = x_array[i];
    const double y_lo = y_array[i];
    const double dx = x_hi - x_lo;
    if(dx != 0.0) {

      if (i == index_a || i == index_b)
        {
          double x1 = (i == index_a) ? a : x_lo;
          double x2 = (i == index_b) ? b : x_hi;
          *result += integ_eval (y_lo, state->b[i], state->c[i], state->d[i],
                                 x_lo, x1, x2);
        }
      else
        {
          *result += dx * (y_lo
                           + dx*(0.5*state->b[i]
                                 + dx*(state->c[i]/3.0
                                       + 0.25*state->d[i]*dx)));
        }
    }
    else {
      *result = 0.0;
      return GSL_EINVAL;
    }
  }

  return GSL_SUCCESS;
}

static const gsl_interp_type akima_type =
{
  "akima",
  5,
  &akima_alloc,
  &akima_init,
  &akima_eval,
  &akima_eval_deriv,
  &akima_eval_deriv2,
  &akima_eval_integ,
  &akima_free
};

const gsl_interp_type * gsl_interp_akima = &akima_type;

static const gsl_interp_type akima_periodic_type =
{
  "akima-periodic",
  5,
  &akima_alloc,
  &akima_init_periodic,
  &akima_eval,
  &akima_eval_deriv,
  &akima_eval_deriv2,
  &akima_eval_integ,
  &akima_free
};

const gsl_interp_type * gsl_interp_akima_periodic = &akima_periodic_type;


#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <math.h>
#include <string.h>
#include <stdbool.h>

#ifdef __APPLE__

#include <cblas.h>
#include <libiomp/omp.h>

#define set_num_threads(x) openblas_set_num_threads(x)
#define get_num_threads() openblas_get_num_threads()
#else
#include <mkl.h>
#include <omp.h>
#define set_num_threads(x) mkl_set_num_threads(x)
#define get_num_threads() mkl_get_num_threads()
#endif

double norm_inf_mat_Y11(double *X, int n);

int main(int argc, char **argv) {

    char *files[] = {
            //"1200_m1_t2_hi_c363417000.7.mat.bin",
            //"1200_m1_t7_hi_c536234414.4.mat.bin",
            //"2340_m2_a1e15_to_c3291.0.mat.bin",
            //"5900_m1_t1_to_c3205771604.1.mat.bin",
            //"1200_m1_t3_hi_c70487569.2.mat.bin",
            //"2340_m1_t1_to_c1446918556.5.mat.bin",
            //"2340_m2_a1e5_to_c8.0.mat.bin",
            //"5900_m2_a1e15_to_c4198.3.mat.bin",
            "494_bus.mtx.bin",
            //"bcsstk13.mtx.bin",
            //"bcsstk27.mtx.bin",
            //"ex9.mtx.bin",
            //"msc01050.mtx.bin",
            //"bcsstk15.mtx.bin",
            //"cage9.mtx.bin",
            //"gyro.mtx.bin",
            "tomography.mtx.bin"
    };
    int num_of_files = sizeof(files) / sizeof(char *);

    char *file_dir = NULL;
    int option = 0, omp_threads = 1;
    while ((option = getopt(argc, argv, "f:t:")) != -1) {
        switch (option) {
            case 'f':
                file_dir = optarg;
                break;
            case 't':
                omp_threads = atoi(optarg);
                break;
            default:
                printf("Usage: mpi_mexp -f string -t num_threads \n");
                return 0;
        }
    }
    if (file_dir == NULL) {
        printf("Usage: mpi_mexp -f string -t num_threads \n");
        return 0;
    }

    set_num_threads(omp_threads);
    omp_set_num_threads(omp_threads);

    int numTasks, rank;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &numTasks);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    int qProcs = (int) sqrt(numTasks);
    if (qProcs - sqrt(numTasks) != 0) {
        if (rank == 0)
            printf("np must be a square number\n");
        MPI_Finalize();
        return 0;
    }

    int q_root = qProcs - 1;
    int world_root = numTasks - 1;
    int n, dims[2] = {qProcs, qProcs}, periods[2] = {0, 0}, reorder = 0, coords[2], rc, i;
    MPI_Comm cartComm;
    MPI_Cart_create(MPI_COMM_WORLD, 2, dims, periods, reorder, &cartComm);
    MPI_Comm_rank(cartComm, &rank);
    MPI_Cart_coords(cartComm, rank, 2, coords);

    MPI_Comm rowComm, colComm;
    MPI_Comm_split(MPI_COMM_WORLD, coords[0], rank, &rowComm);
    MPI_Comm_split(MPI_COMM_WORLD, coords[1], rank, &colComm);

    int row_rank, col_rank;
    MPI_Comm_rank(colComm, &col_rank);
    MPI_Comm_rank(rowComm, &row_rank);

    char filename[1024];
    while (num_of_files-- > 0) {

        memset(filename, 0, strlen(filename));
        strcpy(filename, file_dir);
        strcat(filename, files[num_of_files]);

        MPI_Barrier(MPI_COMM_WORLD);

        if (rank == 0)
            printf("\n%s\n", filename);
        MPI_File matFile;
        rc = MPI_File_open(cartComm, filename, MPI_MODE_RDONLY, MPI_INFO_NULL, &matFile);
        if (rc && (rank == 0)) {
            printf("Unable to open file %s\n", filename);
            fflush(stdout);
        }

        MPI_Status status;
        MPI_File_read(matFile, &n, 1, MPI_INT, &status);

        int blockSize = 0, lastBlock = 0;
        blockSize = (n + qProcs) / qProcs;  /* number of rows in _block_ */
        lastBlock = n + 1 - (qProcs - 1) * blockSize;

        if (rank == 0) {
            printf("size: %d, %d ,%d\n", n, blockSize, lastBlock);
        }

        double *LocalA = (double *) malloc(blockSize * blockSize * sizeof(double));

        MPI_Datatype readFileType;
        MPI_Type_vector(blockSize, blockSize, n, MPI_DOUBLE, &readFileType);
        MPI_Type_commit(&readFileType);
        MPI_Offset offset = (MPI_Offset) (1 * sizeof(int) +
                                          sizeof(double) * (blockSize * coords[1] + blockSize * n * coords[0]));
        MPI_File_set_view(matFile, offset, MPI_DOUBLE, readFileType,
                          "native", MPI_INFO_NULL);

        MPI_File_read_at(matFile, 0, LocalA, blockSize * blockSize, MPI_DOUBLE, &status);
        MPI_Type_free(&readFileType);
        if (coords[0] == q_root) {
            memset(LocalA + (lastBlock - 1) * blockSize, 0,
                   (blockSize + 1 - lastBlock) * blockSize * sizeof(double));
        }
        if (coords[1] == q_root) {
            int kk = 0;
            for (kk = 0; kk < blockSize; ++kk) {
                memset(LocalA + kk * blockSize + lastBlock - 1, 0, (blockSize + 1 - lastBlock) * sizeof(double));
            }
        }

        double loc_d_var;
        double *loc_vec_var = (double *) malloc(blockSize * sizeof(double));
        memset(loc_vec_var, 0, blockSize * sizeof(double));
        offset = (MPI_Offset) (1 * sizeof(int) + sizeof(double) * n * n);
        MPI_File_set_view(matFile, offset, MPI_DOUBLE, MPI_DOUBLE,
                          "native", MPI_INFO_NULL);
        offset = coords[0] * blockSize;
        MPI_File_read_at_all(matFile, offset, loc_vec_var, coords[0] == qProcs - 1 ? lastBlock : blockSize, MPI_DOUBLE,
                             &status);

        MPI_File_close(&matFile);

        if (coords[1] == qProcs - 1) {
            cblas_dcopy(blockSize, loc_vec_var, 1, LocalA + lastBlock - 1, blockSize);
        }

        double mpi_start = MPI_Wtime();
        //compute inf norm of matA
        for (i = 0; i < (col_rank == q_root ? lastBlock - 1 : blockSize); ++i) {
            loc_vec_var[i] = cblas_dasum((row_rank == q_root ? lastBlock - 1 : blockSize), LocalA + i * blockSize,
                                         1);
        }

        if (row_rank == q_root) {
            MPI_Reduce(MPI_IN_PLACE, loc_vec_var, blockSize, MPI_DOUBLE, MPI_SUM, q_root, rowComm);
        } else {
            MPI_Reduce(loc_vec_var, loc_vec_var, blockSize, MPI_DOUBLE, MPI_SUM, q_root, rowComm);
        }

        if (row_rank == q_root) {
            loc_d_var = 0;
            for (i = 0; i < (col_rank == q_root ? lastBlock - 1 : blockSize); ++i) {
                if (loc_vec_var[i] > loc_d_var)
                    loc_d_var = loc_vec_var[i];
            }
            if (col_rank == q_root) {
                MPI_Reduce(MPI_IN_PLACE, &loc_d_var, 1, MPI_DOUBLE, MPI_MAX, q_root, colComm);
            } else {
                MPI_Reduce(&loc_d_var, &loc_d_var, 1, MPI_DOUBLE, MPI_MAX, q_root, colComm);
            }
        }

        MPI_Bcast(&loc_d_var, 1, MPI_DOUBLE, world_root, MPI_COMM_WORLD);

        //init Y martix
        for (i = 0; i < (col_rank == q_root ? lastBlock - 1 : blockSize); ++i) {
            cblas_dscal((row_rank == q_root ? lastBlock - 1 : blockSize), -1 / loc_d_var, LocalA + i * blockSize,
                        1);
        }
        if (row_rank == q_root) {
            cblas_dscal((col_rank == q_root ? lastBlock - 1 : blockSize), 1 / loc_d_var, LocalA + lastBlock - 1,
                        blockSize);
        }
        if (row_rank == col_rank) {
            for (i = 0; i < (row_rank == q_root ? lastBlock - 1 : blockSize); ++i) {
                LocalA[i * blockSize + i] += 1;
            }
        }
        if (rank == world_root) {
            LocalA[(lastBlock - 1) * (blockSize + 1)] = 1;
        }

        //loop
        double tolerance = 1e-8;
        int loopBreak = 0, iteration = 0;
        double *col_mat = (double *) malloc(qProcs * blockSize * blockSize * sizeof(double));

        double *row_mat = (double *) malloc(qProcs * blockSize * blockSize * sizeof(double));

        while (true) {
            //norm mat
            //compute inf norm of matA
            for (i = 0; i < (col_rank == q_root ? lastBlock - 1 : blockSize); ++i) {
                loc_vec_var[i] = cblas_dasum((row_rank == q_root ? lastBlock - 1 : blockSize),
                                             LocalA + i * blockSize, 1);
            }

            if (row_rank == q_root) {
                MPI_Reduce(MPI_IN_PLACE, loc_vec_var, blockSize, MPI_DOUBLE, MPI_SUM, q_root, rowComm);
            } else {
                MPI_Reduce(loc_vec_var, loc_vec_var, blockSize, MPI_DOUBLE, MPI_SUM, q_root, rowComm);
            }

            if (row_rank == q_root) {
                loc_d_var = 0;
                for (i = 0; i < (col_rank == q_root ? lastBlock - 1 : blockSize); ++i) {
                    if (loc_vec_var[i] > loc_d_var)
                        loc_d_var = loc_vec_var[i];
                }
                if (col_rank == q_root) {
                    MPI_Reduce(MPI_IN_PLACE, &loc_d_var, 1, MPI_DOUBLE, MPI_MAX, q_root, colComm);
                } else {
                    MPI_Reduce(&loc_d_var, &loc_d_var, 1, MPI_DOUBLE, MPI_MAX, q_root, colComm);
                }

                size_t index = cblas_idamax((col_rank == q_root ? lastBlock - 1 : blockSize),
                                            LocalA + lastBlock - 1,
                                            blockSize);
                double norm2b = 0;
                MPI_Reduce(&LocalA[index * blockSize + lastBlock - 1], &norm2b, 1, MPI_DOUBLE, MPI_MAX, q_root,
                           colComm);

                if (col_rank == q_root) {
                    loopBreak = (loc_d_var / norm2b < tolerance) ? 1 : 0;
                }

            }
            MPI_Bcast(&loopBreak, 1, MPI_INT, world_root, MPI_COMM_WORLD);
            if (loopBreak > 0)
                break;

            iteration++;

            //gather loc mat
            MPI_Allgather(LocalA, blockSize * blockSize, MPI_DOUBLE, col_mat, blockSize * blockSize, MPI_DOUBLE,
                          colComm);
            MPI_Allgather(LocalA, blockSize * blockSize, MPI_DOUBLE, row_mat, blockSize * blockSize, MPI_DOUBLE,
                          rowComm);
            cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, blockSize, blockSize, blockSize,
                        1, row_mat, blockSize, col_mat, blockSize, 0.0, LocalA, blockSize);
            for (i = 1; i < qProcs; ++i) {
                cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, blockSize, blockSize, blockSize,
                            1, row_mat + i * (blockSize * blockSize), blockSize,
                            col_mat + i * (blockSize * blockSize), blockSize, 1, LocalA, blockSize);
            }
        }

        free(row_mat);
        free(col_mat);
        free(loc_vec_var);

        if (rank == world_root) {
            printf("NumProc: %d\n", numTasks);
            printf("mpi_mexp_iter: %d\t", iteration);
            printf("mpi_mexp_time: %f\n", MPI_Wtime() - mpi_start);

        }

        if (row_rank == q_root) {
            for (i = 0; i < (row_rank == qProcs - 1 ? lastBlock : blockSize); ++i) {
                if (fabs(LocalA[i * blockSize + lastBlock - 1] - 1) > 1e-3) {
                    printf("MPI_MEXP_ERR %f \n", LocalA[i * blockSize + lastBlock - 1]);
                    break;
                }
            }
        }
        free(LocalA);

        /*
        * MPI_CG
        */
        rc = MPI_File_open(cartComm, filename, MPI_MODE_RDONLY, MPI_INFO_NULL, &matFile);
        if (rc && (rank == 0)) {
            printf("Unable to open file %s\n", filename);
            fflush(stdout);
        }

        MPI_File_read(matFile, &n, 1, MPI_INT, &status);
        blockSize = (n + qProcs - 1) / qProcs;  /* number of rows in _block_ */
        lastBlock = n - (qProcs - 1) * blockSize;

        LocalA = (double *) malloc(blockSize * blockSize * sizeof(double));

        MPI_Type_vector(blockSize, blockSize, n, MPI_DOUBLE, &readFileType);
        MPI_Type_commit(&readFileType);
        offset = (MPI_Offset) (1 * sizeof(int) +
                               sizeof(double) * (blockSize * coords[1] + blockSize * n * coords[0]));
        MPI_File_set_view(matFile, offset, MPI_DOUBLE, readFileType,
                          "native", MPI_INFO_NULL);

        MPI_File_read_at(matFile, 0, LocalA, blockSize * blockSize, MPI_DOUBLE, &status);
        MPI_Type_free(&readFileType);
        if (coords[0] == qProcs - 1) {
            memset(LocalA + lastBlock * blockSize, 0, (blockSize - lastBlock) * blockSize * sizeof(double));
        }
        if (coords[1] == qProcs - 1) {
            int kk = 0;
            for (kk = 0; kk < blockSize; ++kk) {
                memset(LocalA + kk * blockSize + lastBlock, 0, (blockSize - lastBlock) * sizeof(double));
            }
        }

        double *LocalB = (double *) malloc(blockSize * sizeof(double));
        memset(LocalB, 0, blockSize * sizeof(double));
        offset = (MPI_Offset) (1 * sizeof(int) + sizeof(double) * n * n);
        MPI_File_set_view(matFile, offset, MPI_DOUBLE, MPI_DOUBLE,
                          "native", MPI_INFO_NULL);
        offset = coords[1] * blockSize;
        MPI_File_read_at_all(matFile, offset, LocalB, coords[1] == qProcs - 1 ? lastBlock : blockSize, MPI_DOUBLE,
                             &status);

        MPI_File_close(&matFile);
        mpi_start = MPI_Wtime();

        double alpha, beta, rho_new = 0.0, rho_old = 0.0;
        iteration = 0;
        double *loc_Vr = LocalB; //(double *) malloc(blockSize * sizeof(double));
        double *loc_Vp = (double *) malloc(blockSize * sizeof(double));
        double *loc_Vw = (double *) malloc(blockSize * sizeof(double));
        double *loc_Vx = (double *) malloc(blockSize * sizeof(double));

        //localr <- localB
        memset(loc_Vx, 0, blockSize * sizeof(double));
        //cblas_dcopy(blockSize, LocalB, 1, loc_Vr, 1);
        double b_norm2 = 0;
        if (col_rank == qProcs - 1) {
            double loc_b_dot = cblas_ddot((row_rank == qProcs - 1 ? lastBlock : blockSize), loc_Vr, 1, loc_Vr, 1);
            MPI_Reduce(&loc_b_dot, &rho_new, 1, MPI_DOUBLE, MPI_SUM, qProcs - 1, rowComm);
            b_norm2 = sqrt(rho_new);
        }

        loopBreak = 0;
        while (true) {
            if (rank == world_root) {
                loopBreak = (sqrt(rho_new) / b_norm2 < tolerance) ? 1 : 0;
            }
            MPI_Bcast(&loopBreak, 1, MPI_INT, world_root, MPI_COMM_WORLD);
            if (loopBreak > 0)
                break;

            iteration++;
            if (iteration == 1) {
                cblas_dcopy(blockSize, loc_Vr, 1, loc_Vp, 1);
            } else {
                if (rank == world_root) {
                    beta = rho_new / rho_old;
                }
                MPI_Bcast(&beta, 1, MPI_DOUBLE, world_root, MPI_COMM_WORLD);
                cblas_dscal((row_rank == qProcs - 1 ? lastBlock : blockSize), beta, loc_Vp, 1);
                cblas_daxpy((row_rank == qProcs - 1 ? lastBlock : blockSize), 1.0, loc_Vr, 1, loc_Vp, 1);
            }

            cblas_dgemv(CblasRowMajor, CblasNoTrans, (col_rank == qProcs - 1 ? lastBlock : blockSize),
                        (row_rank == qProcs - 1 ? lastBlock : blockSize),
                        1.0, LocalA, blockSize, loc_Vp, 1, 0.0, loc_Vw, 1);

            if (col_rank == row_rank) {
                MPI_Reduce(MPI_IN_PLACE, loc_Vw, blockSize, MPI_DOUBLE, MPI_SUM, col_rank, rowComm);
            } else {
                MPI_Reduce(loc_Vw, loc_Vw, blockSize, MPI_DOUBLE, MPI_SUM, col_rank, rowComm);
            }

            MPI_Bcast(loc_Vw, (row_rank == qProcs - 1 ? lastBlock : blockSize), MPI_DOUBLE, row_rank, colComm);

            if (col_rank == qProcs - 1) {
                double loc_pw_dot = cblas_ddot((row_rank == qProcs - 1 ? lastBlock : blockSize), loc_Vp, 1, loc_Vw,
                                               1), glo_pw_dot;
                MPI_Reduce(&loc_pw_dot, &glo_pw_dot, 1, MPI_DOUBLE, MPI_SUM, qProcs - 1, rowComm);
                if (rank == world_root) {
                    alpha = rho_new / glo_pw_dot;
                }
            }
            MPI_Bcast(&alpha, 1, MPI_DOUBLE, world_root, MPI_COMM_WORLD);

            cblas_daxpy((row_rank == qProcs - 1 ? lastBlock : blockSize), -alpha, loc_Vw, 1, loc_Vr, 1);

            if (coords[0] == qProcs - 1) {
                rho_old = rho_new;
                cblas_daxpy((row_rank == qProcs - 1 ? lastBlock : blockSize), alpha, loc_Vp, 1, loc_Vx, 1);
                double loc_r_dot = cblas_ddot(blockSize, loc_Vr, 1, loc_Vr, 1);

                MPI_Reduce(&loc_r_dot, &rho_new, 1, MPI_DOUBLE, MPI_SUM, qProcs - 1, rowComm);
            }
        }

        if (rank == world_root) {
            printf("mpi_cg_it: %d\t", iteration);
            printf("mpi_cg_time: %f\n", MPI_Wtime() - mpi_start);

        }
        free(LocalA);
        free(loc_Vp);
        free(loc_Vr);
        free(loc_Vw);

        if (col_rank == qProcs - 1) {
            for (rc = 0; rc < (row_rank == qProcs - 1 ? lastBlock : blockSize); ++rc) {
                if (fabs(loc_Vx[rc] - 1) > 1e-3) {
                    printf("MPI_CG_ERR %f \n", loc_Vx[rc]);
                    break;
                }
            }
        }
        free(loc_Vx);
    }

    MPI_Comm_free(&cartComm);
    MPI_Comm_free(&colComm);
    MPI_Comm_free(&rowComm);

    MPI_Finalize();
}


#pragma once

#include <gsl/gsl_vector.h>
#include <gsl/gsl_statistics_double.h>
#include <cmath>
#include <memory>
#include <cassert>

namespace halo {

// A simple wrapper for gsl_vector to manage its memory, and also some user-friendly additions.
class GSLVector {
public:
  GSLVector(size_t Sz, bool Zeroed=true)
    : Vec(gsl_vector_alloc(Sz), gsl_vector_free) {
      if (Zeroed)
        gsl_vector_set_all(Vec.get(), 0);
    }

  // number of elements the vector can handle
  size_t capacity() const { return Vec->size; }

  // The stride is the step-size from one element to the next in physical memory,
  // measured in units of the appropriate datatype.
  size_t stride() const { return Vec->stride; }

  // clears the vector, setting everything to zero
  void clear() {
    gsl_vector_set_all(Vec.get(), 0);
  }

  // access underlying vector, without taking ownership.
  gsl_vector* vec() { return Vec.get(); }
  gsl_vector const* vec() const { return Vec.get(); }

  // access the underlying vector's data
  double* data() { return Vec->data; }
  double const* data() const {return Vec->data; }

  // read an element from the vector
  double get(size_t idx) const { return gsl_vector_get(Vec.get(), idx); }

  // write an element of the vector
  void set(size_t idx, double val) { gsl_vector_set(Vec.get(), idx, val); }

private:
  // a managed gsl_vector, which must be constructed with an explicit deleter.
  std::unique_ptr<gsl_vector, decltype(&gsl_vector_free)> Vec;
};

/// A value whose quantity depends on some random phenomenon.
/// Not all observations are used to compute properties of the quantity,
/// with priority given to the most recent N observations, for some N.
class RandomQuantity {
public:
  RandomQuantity(size_t N=10) : Obs(N) {}

  /// Records an observation of the random quantity.
  /// If the history becomes full, the oldest observation is dropped to make room
  /// for this one.
  void observe(double NewSample) {
    Count += 1;
    Obs.set(NextFree, NewSample);

    const size_t Len = Obs.capacity();
    NextFree = (NextFree + 1) % Len; // bump with wrap-around

    // If this is the first observation, fill the vector with this value
    // if (Count == 1)
    //   for (size_t i = NextFree; i < Len; i++)
    //     Obs.set(i, NewSample);
  }

  // adds observations from another random quantity into this one.
  void merge(RandomQuantity const& RQ) {
    for (size_t i = 0; i < RQ.size(); i++)
      observe(RQ.Obs.get(i));
  }

  // resets the random quantity, forgetting all previous observations.
  void clear() {
    Obs.clear();
    NextFree = 0;
    Count = 0;
  }

  // returns the most recent observation
  double last() const {
    assert(Count > 0);
    size_t LastIdx = NextFree == 0 ? Obs.capacity()-1 : NextFree-1;
    return Obs.get(LastIdx);
  }

  /// returns the maximum capacity of this random quantity's ability to
  /// remember observations.
  size_t capacity() const {
    return Obs.capacity();
  }

  /// returns the number of observations currently remembered
  /// by this random quantity.
  size_t size() const { return std::min(Count, Obs.capacity()); }

  /// returns the mean of the RQ's observations.
  double mean() const {
    return gsl_stats_mean(Obs.data(), Obs.stride(), size());
  }

  /// Given the mean of this RQ's observations, returns the variance.
  double variance(double mean) const {
    return gsl_stats_variance_m(Obs.data(), Obs.stride(), size(), mean);
  }

  private:
    GSLVector Obs;
    size_t NextFree{0};
    size_t Count{0};
};

} // end namespace halo


#ifndef INCLUDED_BLAS_TRAITS_H
#define INCLUDED_BLAS_TRAITS_H

#include <cblas.h>
#include <type_traits>

namespace blas{

    /*
    Specialize to define the blas traits of an object. Should
    be constructed from an instance of T
    */
    template< typename T >
    struct Traits {

    };

    template< typename T >
    constexpr auto blas_cast( T&& x ){
        return Traits<T>( std::forward<T>( x ) );
    }
    template< typename T >
    constexpr CBLAS_ORDER layout( T&& traits ){
        return (CBLAS_ORDER) traits.layout(); }
    template< typename T >
    constexpr CBLAS_TRANSPOSE trans( T&& traits ){
        return (CBLAS_TRANSPOSE) traits.trans(); }
    template< typename T >
    constexpr CBLAS_UPLO uplo( T&& traits ){
        return (CBLAS_UPLO) traits.uplo(); }
    template< typename T >
    constexpr auto nrows( T&& traits ){ return traits.nrows(); }
    template< typename T >
    constexpr auto ncols( T&& traits ){ return traits.ncols(); }
    template< typename T >
    constexpr auto array( T&& traits ){ return traits.array(); }
    template< typename T >
    constexpr auto ld( T&& traits ){ return traits.ld(); }

    template< typename T >
    struct MatrixTraits {
        /*
        using scalar_t = ...
        static CBLAS_ORDER layout( const T& );
        static CBLAS_TRANSPOSE trans( const T& );
        static CBLAS_UPLO uplo( const T& ); --if symmetric
        static std::size_t nrows( const T& );
        static std::size_t ncols( const T& );
        static scalar_t* array( T& );
        static const scalar_t* array( const T& );
        static std::size_t ld( const T& );
        */
    };

    template< typename T >
    struct VectorTraits{
        /*
        using scalar_t = ...
        static std::size_t length( const T& );
        static std::size_t inc( const T& );
        static scalar_t* array( T& );
        static const scalar_t* array( const T& );
        */
    };

    template< typename T >
    struct BlasCast : public std::false_type{
        /*
        When implemented - extend std::true_type (define ::value = true)
        static _ blas_cast( T& );
        static _ blas_cast( const T& );
        */
    };

} //namespace blas

#endif


/* integration/qmomof.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000, 2007 Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_errno.h>

static void
compute_moments (double par, double * cheb);

static int
dgtsl (size_t n, double *c, double *d, double *e, double *b);

gsl_integration_qawo_table *
gsl_integration_qawo_table_alloc (double omega, double L,
                                  enum gsl_integration_qawo_enum sine,
                                  size_t n)
{
  gsl_integration_qawo_table *t;
  double * chebmo;

  if (n == 0)
    {
      GSL_ERROR_VAL ("table length n must be positive integer",
                        GSL_EDOM, 0);
    }

  t = (gsl_integration_qawo_table *)
    malloc (sizeof (gsl_integration_qawo_table));

  if (t == 0)
    {
      GSL_ERROR_VAL ("failed to allocate space for qawo_table struct",
                        GSL_ENOMEM, 0);
    }

  chebmo = (double *)  malloc (25 * n * sizeof (double));

  if (chebmo == 0)
    {
      free (t);
      GSL_ERROR_VAL ("failed to allocate space for chebmo block",
                        GSL_ENOMEM, 0);
    }

  t->n = n;
  t->sine = sine;
  t->omega = omega;
  t->L = L;
  t->par = 0.5 * omega * L;
  t->chebmo = chebmo;

  /* precompute the moments */

  {
    size_t i;
    double scale = 1.0;

    for (i = 0 ; i < t->n; i++)
      {
        compute_moments (t->par * scale, t->chebmo + 25*i);
        scale *= 0.5;
      }
  }

  return t;
}

int
gsl_integration_qawo_table_set (gsl_integration_qawo_table * t,
                                    double omega, double L,
                                    enum gsl_integration_qawo_enum sine)
{
  t->omega = omega;
  t->sine = sine;
  t->L = L;
  t->par = 0.5 * omega * L;

  /* recompute the moments */

  {
    size_t i;
    double scale = 1.0;

    for (i = 0 ; i < t->n; i++)
      {
        compute_moments (t->par * scale, t->chebmo + 25*i);
        scale *= 0.5;
      }
  }

  return GSL_SUCCESS;
}

int
gsl_integration_qawo_table_set_length (gsl_integration_qawo_table * t,
                                           double L)
{
  /* return immediately if the length is the same as the old length */

  if (L == t->L)
    return GSL_SUCCESS;

  /* otherwise reset the table and compute the new parameters */

  t->L = L;
  t->par = 0.5 * t->omega * L;

  /* recompute the moments */

  {
    size_t i;
    double scale = 1.0;

    for (i = 0 ; i < t->n; i++)
      {
        compute_moments (t->par * scale, t->chebmo + 25*i);
        scale *= 0.5;
      }
  }

  return GSL_SUCCESS;
}

void
gsl_integration_qawo_table_free (gsl_integration_qawo_table * t)
{
  RETURN_IF_NULL (t);
  free (t->chebmo);
  free (t);
}

static void
compute_moments (double par, double *chebmo)
{
  double v[28], d[25], d1[25], d2[25];

  const size_t noeq = 25;

  const double par2 = par * par;
  const double par4 = par2 * par2;
  const double par22 = par2 + 2.0;

  const double sinpar = sin (par);
  const double cospar = cos (par);

  size_t i;

  /* compute the chebyschev moments with respect to cosine */

  double ac = 8 * cospar;
  double as = 24 * par * sinpar;

  v[0] = 2 * sinpar / par;
  v[1] = (8 * cospar + (2 * par2 - 8) * sinpar / par) / par2;
  v[2] = (32 * (par2 - 12) * cospar
          + (2 * ((par2 - 80) * par2 + 192) * sinpar) / par) / par4;

  if (fabs (par) <= 24)
    {
      /* compute the moments as the solution of a boundary value
         problem using the asyptotic expansion as an endpoint */

      double an2, ass, asap;
      double an = 6;
      size_t k;

      for (k = 0; k < noeq - 1; k++)
        {
          an2 = an * an;
          d[k] = -2 * (an2 - 4) * (par22 - 2 * an2);
          d2[k] = (an - 1) * (an - 2) * par2;
          d1[k + 1] = (an + 3) * (an + 4) * par2;
          v[k + 3] = as - (an2 - 4) * ac;
          an = an + 2.0;
        }

      an2 = an * an;

      d[noeq - 1] = -2 * (an2 - 4) * (par22 - 2 * an2);
      v[noeq + 2] = as - (an2 - 4) * ac;
      v[3] = v[3] - 56 * par2 * v[2];

      ass = par * sinpar;
      asap = (((((210 * par2 - 1) * cospar - (105 * par2 - 63) * ass) / an2
                - (1 - 15 * par2) * cospar + 15 * ass) / an2
               - cospar + 3 * ass) / an2
              - cospar) / an2;
      v[noeq + 2] = v[noeq + 2] - 2 * asap * par2 * (an - 1) * (an - 2);

      dgtsl (noeq, d1, d, d2, v + 3);

    }
  else
    {
      /* compute the moments by forward recursion */
      size_t k;
      double an = 4;

      for (k = 3; k < 13; k++)
        {
          double an2 = an * an;
          v[k] = ((an2 - 4) * (2 * (par22 - 2 * an2) * v[k - 1] - ac)
                  + as - par2 * (an + 1) * (an + 2) * v[k - 2])
            / (par2 * (an - 1) * (an - 2));
          an = an + 2.0;
        }
    }

  for (i = 0; i < 13; i++)
    {
      chebmo[2 * i] = v[i];
    }

  /* compute the chebyschev moments with respect to sine */

  v[0] = 2 * (sinpar - par * cospar) / par2;
  v[1] = (18 - 48 / par2) * sinpar / par2 + (-2 + 48 / par2) * cospar / par;

  ac = -24 * par * cospar;
  as = -8 * sinpar;

  if (fabs (par) <= 24)
    {
      /* compute the moments as the solution of a boundary value
         problem using the asyptotic expansion as an endpoint */

      size_t k;
      double an2, ass, asap;
      double an = 5;

      for (k = 0; k < noeq - 1; k++)
        {
          an2 = an * an;
          d[k] = -2 * (an2 - 4) * (par22 - 2 * an2);
          d2[k] = (an - 1) * (an - 2) * par2;
          d1[k + 1] = (an + 3) * (an + 4) * par2;
          v[k + 2] = ac + (an2 - 4) * as;
          an = an + 2.0;
        }

      an2 = an * an;

      d[noeq - 1] = -2 * (an2 - 4) * (par22 - 2 * an2);
      v[noeq + 1] = ac + (an2 - 4) * as;
      v[2] = v[2] - 42 * par2 * v[1];

      ass = par * cospar;
      asap = (((((105 * par2 - 63) * ass - (210 * par2 - 1) * sinpar) / an2
                + (15 * par2 - 1) * sinpar
                - 15 * ass) / an2 - sinpar - 3 * ass) / an2 - sinpar) / an2;
      v[noeq + 1] = v[noeq + 1] - 2 * asap * par2 * (an - 1) * (an - 2);

      dgtsl (noeq, d1, d, d2, v + 2);

    }
  else
    {
      /* compute the moments by forward recursion */
      size_t k;
      double an = 3;
      for (k = 2; k < 12; k++)
        {
          double an2 = an * an;
          v[k] = ((an2 - 4) * (2 * (par22 - 2 * an2) * v[k - 1] + as)
                  + ac - par2 * (an + 1) * (an + 2) * v[k - 2])
            / (par2 * (an - 1) * (an - 2));
          an = an + 2.0;
        }
    }

  for (i = 0; i < 12; i++)
    {
      chebmo[2 * i + 1] = v[i];
    }

}

static int
dgtsl (size_t n, double *c, double *d, double *e, double *b)
{
  /* solves a tridiagonal matrix A x = b

     c[1 .. n - 1]   subdiagonal of the matrix A
     d[0 .. n - 1]   diagonal of the matrix A
     e[0 .. n - 2]   superdiagonal of the matrix A

     b[0 .. n - 1]   right hand side, replaced by the solution vector x */

  size_t k;

  c[0] = d[0];

  if (n == 0)
    {
      return GSL_SUCCESS;
    }

  if (n == 1)
    {
      b[0] = b[0] / d[0] ;
      return GSL_SUCCESS;
    }

  d[0] = e[0];
  e[0] = 0;
  e[n - 1] = 0;

  for (k = 0; k < n - 1; k++)
    {
      size_t k1 = k + 1;

      if (fabs (c[k1]) >= fabs (c[k]))
        {
          {
            double t = c[k1];
            c[k1] = c[k];
            c[k] = t;
          };
          {
            double t = d[k1];
            d[k1] = d[k];
            d[k] = t;
          };
          {
            double t = e[k1];
            e[k1] = e[k];
            e[k] = t;
          };
          {
            double t = b[k1];
            b[k1] = b[k];
            b[k] = t;
          };
        }

      if (c[k] == 0)
        {
          return GSL_FAILURE ;
        }

      {
        double t = -c[k1] / c[k];

        c[k1] = d[k1] + t * d[k];
        d[k1] = e[k1] + t * e[k];
        e[k1] = 0;
        b[k1] = b[k1] + t * b[k];
      }

    }

  if (c[n - 1] == 0)
    {
      return GSL_FAILURE;
    }

  b[n - 1] = b[n - 1] / c[n - 1];

  b[n - 2] = (b[n - 2] - d[n - 2] * b[n - 1]) / c[n - 2];

  for (k = n ; k > 2; k--)
    {
      size_t kb = k - 3;
      b[kb] = (b[kb] - d[kb] * b[kb + 1] - e[kb] * b[kb + 2]) / c[kb];
    }

  return GSL_SUCCESS;
}


// Copyright (C) 2021 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#ifndef _HEBench_ClearText_MatMult_H_7e5fa8c2415240ea93eff148ed73539b
#define _HEBench_ClearText_MatMult_H_7e5fa8c2415240ea93eff148ed73539b

#include <gsl/gsl>

#include "hebench/api_bridge/cpp/hebench.hpp"

#include "clear_benchmark.h"

template <class T>
class MatMult_Benchmark : public ClearTextBenchmark
{
private:
    HEBERROR_DECLARE_CLASS_NAME(MatMult_Benchmark)

public:
    MatMult_Benchmark(hebench::cpp::BaseEngine &engine,
                      const hebench::APIBridge::BenchmarkDescriptor &bench_desc,
                      const hebench::APIBridge::WorkloadParams &bench_params);
    ~MatMult_Benchmark() override;

    hebench::APIBridge::Handle encode(const hebench::APIBridge::PackedData *p_parameters) override;
    void decode(hebench::APIBridge::Handle encoded_data, hebench::APIBridge::PackedData *p_native) override;

    hebench::APIBridge::Handle load(const hebench::APIBridge::Handle *p_local_data, std::uint64_t count) override;
    void store(hebench::APIBridge::Handle remote_data,
               hebench::APIBridge::Handle *p_local_data, std::uint64_t count) override;

    hebench::APIBridge::Handle operate(hebench::APIBridge::Handle h_remote_packed,
                                       const hebench::APIBridge::ParameterIndexer *p_param_indexers) override;

    //std::int64_t classTag() const override { return BaseBenchmark::classTag() | ClearTextBenchmark::tag; }

protected:
    std::array<std::uint64_t, 2> m_rows;
    std::array<std::uint64_t, 2> m_cols;

private:
    static void matMult(gsl::span<T> &result,
                        const gsl::span<const T> &M0, const gsl::span<const T> &M1,
                        std::size_t columns);
};

#include "inl/bench_matmult.inl"

#endif // defined _HEBench_ClearText_MatMult_H_7e5fa8c2415240ea93eff148ed73539b


// Copyright (c) 2021 Stig Rune Sellevag
//
// This file is distributed under the MIT License. See the accompanying file
// LICENSE.txt or http://www.opensource.org/licenses/mit-license.php for terms
// and conditions.

#ifndef SCILIB_LINALG_BLAS1_VECTOR_NORM2_H
#define SCILIB_LINALG_BLAS1_VECTOR_NORM2_H

#ifdef USE_MKL
#include <mkl.h>
#else
#include <cblas.h>
#endif

#include <experimental/mdspan>

namespace Sci {
namespace Linalg {

namespace stdex = std::experimental;

template <stdex::extents<>::size_type ext_x, class Layout_x, class Accessor_x>
inline double
norm2(stdex::mdspan<double, stdex::extents<ext_x>, Layout_x, Accessor_x> x)
{
    const BLAS_INT n = static_cast<BLAS_INT>(x.size());
    const BLAS_INT incx = static_cast<BLAS_INT>(x.stride(0));

    return cblas_dnrm2(n, x.data(), incx);
}

template <stdex::extents<>::size_type ext_x, class Layout_x, class Accessor_x>
inline double norm2(
    stdex::mdspan<const double, stdex::extents<ext_x>, Layout_x, Accessor_x> x)
{
    const BLAS_INT n = static_cast<BLAS_INT>(x.size());
    const BLAS_INT incx = static_cast<BLAS_INT>(x.stride(0));

    return cblas_dnrm2(n, x.data(), incx);
}

template <class Layout, class Allocator>
inline double norm2(const Sci::Vector<double, Layout, Allocator>& x)
{
    return norm2(x.view());
}

} // namespace Linalg
} // namespace Sci

#endif // SCILIB_LINALG_BLAS1_VECTOR_NORM2_H