#include <math.h>
#include <unistd.h>
#include <stdio.h>
#include <ctype.h>
#include <stdlib.h>
#include <time.h>
#include <string.h>
//#include <pthread.h>
#include <omp.h>
#include <complex.h>
#include <fftw3.h>
#include <gsl/gsl_interp.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_roots.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_errno.h>
#include "21CMMC.h"

void ComputeInitialConditions(struct UserParams user_params, struct CosmoParams cosmo_params, struct InitialConditions boxes) {

    /*
     Generates the initial conditions: gaussian random density field (DIM^3) as well as the equal or lower resolution velocity fields, and smoothed density field (HII_DIM^3).
     See INIT_PARAMS.H and ANAL_PARAMS.H to set the appropriate parameters.
     Output is written to ../Boxes

     Author: Andrei Mesinger
     Date: 9/29/06
     */

    fftwf_plan plan;

    unsigned long long ct;
    int n_x, n_y, n_z, i, j, k, ii;
    float k_x, k_y, k_z, k_mag, p, a, b, k_sq;
    double pixel_deltax;

    float f_pixel_factor;

    gsl_rng * r;

    /************  INITIALIZATION **********************/

    // Removed all references to threads as 21CMMC is always a single core implementation

    printf("%d\n",cosmo_params.RANDOM_SEED);
    // seed the random number generators
//    r = gsl_rng_alloc(gsl_rng_mt19937);
//    gsl_rng_set(r, cosmo_params.RANDOM_SEED);

/*
    // allocate array for the k-space and real-space boxes
    HIRES_box = (fftwf_complex *) fftwf_malloc(sizeof(fftwf_complex)*user_params.KSPACE_NUM_PIXELS);
    HIRES_box_saved = (fftwf_complex *) fftwf_malloc(sizeof(fftwf_complex)*user_params.KSPACE_NUM_PIXELS);

    // now allocate memory for the lower-resolution box
    // use HII_DIM from ANAL_PARAMS
    LOWRES_density = (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);
    LOWRES_vx = (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);
    LOWRES_vy= (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);
    LOWRES_vz = (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);

    if(SECOND_ORDER_LPT_CORRECTIONS){
        LOWRES_vx_2LPT = (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);
        LOWRES_vy_2LPT = (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);
        LOWRES_vz_2LPT = (float *) malloc(sizeof(float)*HII_TOT_NUM_PIXELS);
    }

    // find factor of HII pixel size / deltax pixel size
    f_pixel_factor = DIM/(float)HII_DIM;
*/

    /************  END INITIALIZATION ******************/

    /************ CREATE K-SPACE GAUSSIAN RANDOM FIELD ***********/
/*
    for (n_x=0; n_x<DIM; n_x++){
        // convert index to numerical value for this component of the k-mode: k = (2*pi/L) * n
        if (n_x>MIDDLE)
            k_x =(n_x-DIM) * DELTA_K;  // wrap around for FFT convention
        else
            k_x = n_x * DELTA_K;

        for (n_y=0; n_y<DIM; n_y++){
            // convert index to numerical value for this component of the k-mode: k = (2*pi/L) * n
            if (n_y>MIDDLE)
                k_y =(n_y-DIM) * DELTA_K;
            else
                k_y = n_y * DELTA_K;

            // since physical space field is real, only half contains independent modes
            for (n_z=0; n_z<=MIDDLE; n_z++){
                // convert index to numerical value for this component of the k-mode: k = (2*pi/L) * n
                k_z = n_z * DELTA_K;

                // now get the power spectrum; remember, only the magnitude of k counts (due to issotropy)
                // this could be used to speed-up later maybe
                k_mag = sqrt(k_x*k_x + k_y*k_y + k_z*k_z);
                p = power_in_k(k_mag);

                // ok, now we can draw the values of the real and imaginary part
                // of our k entry from a Gaussian distribution
                a = gsl_ran_ugaussian(r);
                b = gsl_ran_ugaussian(r);
                HIRES_box[C_INDEX(n_x, n_y, n_z)] = sqrt(VOLUME*p/2.0) * (a + b*I);
            }
        }
    }
*/
    /*****  Adjust the complex conjugate relations for a real array  *****/

//    adj_complex_conj(HIRES_box);

    /*** Let's also create a lower-resolution version of the density field  ***/

/*    memcpy(HIRES_box_saved, HIRES_box, sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);

    if (DIM != HII_DIM)
        filter(HIRES_box, 0, L_FACTOR*BOX_LEN/(HII_DIM+0.0));
    // FFT back to real space
    plan = fftwf_plan_dft_c2r_3d(DIM, DIM, DIM, (fftwf_complex *)HIRES_box, (float *)HIRES_box, FFTW_ESTIMATE);
    fftwf_execute(plan);
    // now sample the filtered box
    for (i=0; i<HII_DIM; i++){
        for (j=0; j<HII_DIM; j++){
            for (k=0; k<HII_DIM; k++){
                LOWRES_density[HII_R_INDEX(i,j,k)] =
                *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                   (unsigned long long)(j*f_pixel_factor+0.5),
                                                   (unsigned long long)(k*f_pixel_factor+0.5)))/VOLUME;
            }
        }
    }
*/

    /******* PERFORM INVERSE FOURIER TRANSFORM *****************/
    // add the 1/VOLUME factor when converting from k space to real space
/*
    memcpy(HIRES_box, HIRES_box_saved, sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);

    for (ct=0; ct<KSPACE_NUM_PIXELS; ct++){
        HIRES_box[ct] /= VOLUME;
    }
    plan = fftwf_plan_dft_c2r_3d(DIM, DIM, DIM, (fftwf_complex *)HIRES_box, (float *)HIRES_box, FFTW_ESTIMATE);
    fftwf_execute(plan);
    fftwf_destroy_plan(plan);
    fftwf_cleanup();

    for (i=0; i<DIM; i++){
        for (j=0; j<DIM; j++){
            for (k=0; k<DIM; k++){
                *((float *)HIRES_density + R_FFT_INDEX(i,j,k)) = *((float *)HIRES_box + R_FFT_INDEX(i,j,k));
            }
        }
    }

    for(ii=0;ii<3;ii++) {

        memcpy(HIRES_box, HIRES_box_saved, sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);
        // Now let's set the velocity field/dD/dt (in comoving Mpc)

        for (n_x=0; n_x<DIM; n_x++){
            if (n_x>MIDDLE)
                k_x =(n_x-DIM) * DELTA_K;  // wrap around for FFT convention
            else
                k_x = n_x * DELTA_K;

            for (n_y=0; n_y<DIM; n_y++){
                if (n_y>MIDDLE)
                    k_y =(n_y-DIM) * DELTA_K;
                else
                    k_y = n_y * DELTA_K;

                for (n_z=0; n_z<=MIDDLE; n_z++){
                    k_z = n_z * DELTA_K;

                    k_sq = k_x*k_x + k_y*k_y + k_z*k_z;

                    // now set the velocities
                    if ((n_x==0) && (n_y==0) && (n_z==0)){ // DC mode
                        HIRES_box[0] = 0;
                    }
                    else{
                        if(ii==0) {
                            HIRES_box[C_INDEX(n_x,n_y,n_z)] *= k_x*I/k_sq/VOLUME;
                        }
                        if(ii==1) {
                            HIRES_box[C_INDEX(n_x,n_y,n_z)] *= k_y*I/k_sq/VOLUME;
                        }
                        if(ii==2) {
                            HIRES_box[C_INDEX(n_x,n_y,n_z)] *= k_z*I/k_sq/VOLUME;
                        }
                        // note the last factor of 1/VOLUME accounts for the scaling in real-space, following the FFT
                    }
                }
            }
        }

        if (DIM != HII_DIM)
            filter(HIRES_box, 0, L_FACTOR*BOX_LEN/(HII_DIM+0.0));

        plan = fftwf_plan_dft_c2r_3d(DIM, DIM, DIM, (fftwf_complex *)HIRES_box, (float *)HIRES_box, FFTW_ESTIMATE);
        fftwf_execute(plan);
        // now sample to lower res
        // now sample the filtered box
        for (i=0; i<HII_DIM; i++){
            for (j=0; j<HII_DIM; j++){
                for (k=0; k<HII_DIM; k++){
                    if(ii==0) {
                        LOWRES_vx[HII_R_INDEX(i,j,k)] =
                        *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                           (unsigned long long)(j*f_pixel_factor+0.5),
                                                           (unsigned long long)(k*f_pixel_factor+0.5)));
                    }
                    if(ii==1) {
                        LOWRES_vy[HII_R_INDEX(i,j,k)] =
                        *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                           (unsigned long long)(j*f_pixel_factor+0.5),
                                                           (unsigned long long)(k*f_pixel_factor+0.5)));
                    }
                    if(ii==2) {
                        LOWRES_vz[HII_R_INDEX(i,j,k)] =
                        *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                           (unsigned long long)(j*f_pixel_factor+0.5),
                                                           (unsigned long long)(k*f_pixel_factor+0.5)));
                    }
                }
            }
        }
    }
    // write out file
*/
    /* *************************************************** *
     *              BEGIN 2LPT PART                        *
     * *************************************************** */

    // Generation of the second order Lagrangian perturbation theory (2LPT) corrections to the ZA
    // reference: Scoccimarro R., 1998, MNRAS, 299, 1097-1118 Appendix D
/*
    // Parameter set in ANAL_PARAMS.H
    if(SECOND_ORDER_LPT_CORRECTIONS){
        // use six supplementary boxes to store the gradients of phi_1 (eq. D13b)
        // Allocating the boxes
#define PHI_INDEX(i, j) ((int) ((i) - (j)) + 3*((j)) - ((int)(j))/2  )
        // ij -> INDEX
        // 00 -> 0
        // 11 -> 3
        // 22 -> 5
        // 10 -> 1
        // 20 -> 2
        // 21 -> 4

        fftwf_complex *phi_1[6];

        for(i = 0; i < 3; ++i){
            for(j = 0; j <= i; ++j){
                phi_1[PHI_INDEX(i, j)] = (fftwf_complex *) fftwf_malloc(sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);
            }
        }

        for(i = 0; i < 3; ++i){
            for(j = 0; j <= i; ++j){

                // read in the box
                memcpy(HIRES_box, HIRES_box_saved, sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);

                // generate the phi_1 boxes in Fourier transform
                for (n_x=0; n_x<DIM; n_x++){
                    if (n_x>MIDDLE)
                        k_x =(n_x-DIM) * DELTA_K;  // wrap around for FFT convention
                    else
                        k_x = n_x * DELTA_K;

                    for (n_y=0; n_y<DIM; n_y++){
                        if (n_y>MIDDLE)
                            k_y =(n_y-DIM) * DELTA_K;
                        else
                            k_y = n_y * DELTA_K;

                        for (n_z=0; n_z<=MIDDLE; n_z++){
                            k_z = n_z * DELTA_K;

                            k_sq = k_x*k_x + k_y*k_y + k_z*k_z;

                            float k[] = {k_x, k_y, k_z};
                            // now set the velocities
                            if ((n_x==0) && (n_y==0) && (n_z==0)){ // DC mode
                                phi_1[PHI_INDEX(i, j)][0] = 0;
                            }
                            else{
                                phi_1[PHI_INDEX(i, j)][C_INDEX(n_x,n_y,n_z)] = -k[i]*k[j]*HIRES_box[C_INDEX(n_x, n_y, n_z)]/k_sq/VOLUME;
                                // note the last factor of 1/VOLUME accounts for the scaling in real-space, following the FFT
                            }
                        }
                    }
                }
                // Now we can generate the real phi_1[i,j]
                plan = fftwf_plan_dft_c2r_3d(DIM, DIM, DIM, (fftwf_complex *)phi_1[PHI_INDEX(i, j)], (float *)phi_1[PHI_INDEX(i, j)], FFTW_ESTIMATE);
                fftwf_execute(plan);
            }
        }

        // Then we will have the laplacian of phi_2 (eq. D13b)
        // After that we have to return in Fourier space and generate the Fourier transform of phi_2
        int m, l;
        for (i=0; i<DIM; i++){
            for (j=0; j<DIM; j++){
                for (k=0; k<DIM; k++){
                    *( (float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i), (unsigned long long)(j), (unsigned long long)(k) )) = 0.0;
                    for(m = 0; m < 3; ++m){
                        for(l = m+1; l < 3; ++l){
                            *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i),(unsigned long long)(j),(unsigned long long)(k)) ) += ( *((float *)(phi_1[PHI_INDEX(l, l)]) + R_FFT_INDEX((unsigned long long) (i),(unsigned long long) (j),(unsigned long long) (k)))  ) * (  *((float *)(phi_1[PHI_INDEX(m, m)]) + R_FFT_INDEX((unsigned long long)(i),(unsigned long long)(j),(unsigned long long)(k)))  );
                            *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i),(unsigned long long)(j),(unsigned long long)(k)) ) -= ( *((float *)(phi_1[PHI_INDEX(l, m)]) + R_FFT_INDEX((unsigned long long)(i),(unsigned long long) (j),(unsigned long long)(k) ) )  ) * (  *((float *)(phi_1[PHI_INDEX(l, m)]) + R_FFT_INDEX((unsigned long long)(i),(unsigned long long)(j),(unsigned long long)(k) ))  );
                            *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i),(unsigned long long)(j),(unsigned long long)(k)) ) /= TOT_NUM_PIXELS;
                        }
                    }
                }
            }
        }

        plan = fftwf_plan_dft_r2c_3d(DIM, DIM, DIM, (float *)HIRES_box, (fftwf_complex *)HIRES_box, FFTW_ESTIMATE);
        fftwf_execute(plan);

        memcpy(HIRES_box_saved, HIRES_box, sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);

        // Now we can store the content of box in a back-up file
        // Then we can generate the gradients of phi_2 (eq. D13b and D9)
*/
        /***** Write out back-up k-box RHS eq. D13b *****/
/*
        // For each component, we generate the velocity field (same as the ZA part)

        // Now let's set the velocity field/dD/dt (in comoving Mpc)

        // read in the box
        // TODO correct free of phi_1

        for(ii=0;ii<3;ii++) {

            if(ii>0) {
                memcpy(HIRES_box, HIRES_box_saved, sizeof(fftwf_complex)*KSPACE_NUM_PIXELS);
            }
            // set velocities/dD/dt
            for (n_x=0; n_x<DIM; n_x++){
                if (n_x>MIDDLE)
                    k_x =(n_x-DIM) * DELTA_K;  // wrap around for FFT convention
                else
                    k_x = n_x * DELTA_K;

                for (n_y=0; n_y<DIM; n_y++){
                    if (n_y>MIDDLE)
                        k_y =(n_y-DIM) * DELTA_K;
                    else
                        k_y = n_y * DELTA_K;

                    for (n_z=0; n_z<=MIDDLE; n_z++){
                        k_z = n_z * DELTA_K;

                        k_sq = k_x*k_x + k_y*k_y + k_z*k_z;

                        // now set the velocities
                        if ((n_x==0) && (n_y==0) && (n_z==0)){ // DC mode
                            HIRES_box[0] = 0;
                        }
                        else{
                            if(ii==0) {
                                HIRES_box[C_INDEX(n_x,n_y,n_z)] *= k_x*I/k_sq;
                            }
                            if(ii==1) {
                                HIRES_box[C_INDEX(n_x,n_y,n_z)] *= k_y*I/k_sq;
                            }
                            if(ii==2) {
                                HIRES_box[C_INDEX(n_x,n_y,n_z)] *= k_z*I/k_sq;
                            }
                            // note the last factor of 1/VOLUME accounts for the scaling in real-space, following the FFT
                        }
                    }
                }
            }

            if (DIM != HII_DIM)
                filter(HIRES_box, 0, L_FACTOR*BOX_LEN/(HII_DIM+0.0));

            plan = fftwf_plan_dft_c2r_3d(DIM, DIM, DIM, (fftwf_complex *)HIRES_box, (float *)HIRES_box, FFTW_ESTIMATE);
            fftwf_execute(plan);
            // now sample to lower res
            // now sample the filtered box
            for (i=0; i<HII_DIM; i++){
                for (j=0; j<HII_DIM; j++){
                    for (k=0; k<HII_DIM; k++){
                        if(ii==0) {
                            LOWRES_vx_2LPT[HII_R_INDEX(i,j,k)] =
                            *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                               (unsigned long long)(j*f_pixel_factor+0.5),
                                                               (unsigned long long)(k*f_pixel_factor+0.5)));
                        }
                        if(ii==1) {
                            LOWRES_vy_2LPT[HII_R_INDEX(i,j,k)] =
                            *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                               (unsigned long long)(j*f_pixel_factor+0.5),
                                                               (unsigned long long)(k*f_pixel_factor+0.5)));
                        }
                        if(ii==2) {
                            LOWRES_vz_2LPT[HII_R_INDEX(i,j,k)] =
                            *((float *)HIRES_box + R_FFT_INDEX((unsigned long long)(i*f_pixel_factor+0.5),
                                                               (unsigned long long)(j*f_pixel_factor+0.5),
                                                               (unsigned long long)(k*f_pixel_factor+0.5)));
                        }
                    }
                }
            }
        }
        // deallocate the supplementary boxes
        for(i = 0; i < 3; ++i){
            for(j = 0; j <= i; ++j){
                fftwf_free(phi_1[PHI_INDEX(i,j)]);
            }
        }
    }
*/
    /* *********************************************** *
     *               END 2LPT PART                     *
     * *********************************************** */

    // deallocate
//    fftwf_free(HIRES_box);
//    fftwf_free(HIRES_box_saved);
}

/*****  Adjust the complex conjugate relations for a real array  *****/
/*
void adj_complex_conj(){
    int i, j, k;

    // corners
    HIRES_box[C_INDEX(0,0,0)] = 0;
    HIRES_box[C_INDEX(0,0,MIDDLE)] = crealf(HIRES_box[C_INDEX(0,0,MIDDLE)]);
    HIRES_box[C_INDEX(0,MIDDLE,0)] = crealf(HIRES_box[C_INDEX(0,MIDDLE,0)]);
    HIRES_box[C_INDEX(0,MIDDLE,MIDDLE)] = crealf(HIRES_box[C_INDEX(0,MIDDLE,MIDDLE)]);
    HIRES_box[C_INDEX(MIDDLE,0,0)] = crealf(HIRES_box[C_INDEX(MIDDLE,0,0)]);
    HIRES_box[C_INDEX(MIDDLE,0,MIDDLE)] = crealf(HIRES_box[C_INDEX(MIDDLE,0,MIDDLE)]);
    HIRES_box[C_INDEX(MIDDLE,MIDDLE,0)] = crealf(HIRES_box[C_INDEX(MIDDLE,MIDDLE,0)]);
    HIRES_box[C_INDEX(MIDDLE,MIDDLE,MIDDLE)] = crealf(HIRES_box[C_INDEX(MIDDLE,MIDDLE,MIDDLE)]);

    // do entire i except corners
    for (i=1; i<MIDDLE; i++){
        // just j corners
        for (j=0; j<=MIDDLE; j+=MIDDLE){
            for (k=0; k<=MIDDLE; k+=MIDDLE){
                HIRES_box[C_INDEX(i,j,k)] = conjf(HIRES_box[C_INDEX(DIM-i,j,k)]);
            }
        }

        // all of j
        for (j=1; j<MIDDLE; j++){
            for (k=0; k<=MIDDLE; k+=MIDDLE){
                HIRES_box[C_INDEX(i,j,k)] = conjf(HIRES_box[C_INDEX(DIM-i,DIM-j,k)]);
                HIRES_box[C_INDEX(i,DIM-j,k)] = conjf(HIRES_box[C_INDEX(DIM-i,j,k)]);
            }
        }
    } // end loop over i

    // now the i corners
    for (i=0; i<=MIDDLE; i+=MIDDLE){
        for (j=1; j<MIDDLE; j++){
            for (k=0; k<=MIDDLE; k+=MIDDLE){
                HIRES_box[C_INDEX(i,j,k)] = conjf(HIRES_box[C_INDEX(i,DIM-j,k)]);
            }
        }
    } // end loop over remaining j

}
*/


#ifndef CTETRA_SUBMESH_H
#define CTETRA_SUBMESH_H

#include <math.h>
#include <stdio.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include "ecache.h"

void submesh_ijk_to_k(int na, int nb, int nc, int i, int j, int k, double k_opt[3]);

int submesh_ijk_index(int na, int nb, int nc, int i, int j, int k);

int **subcells_around_ijk(int na, int nb, int nc, int i, int j, int k, int *subcell_num);

void get_k_orig(double k_opt[3], int G_order[3], int G_neg[3], double k_orig[3]);

void MinMaxVals(EnergyCache *Ecache, double *emin, double *emax);

void OptimizeGs(gsl_matrix *R, int G_order[3], int G_neg[3]);

#endif // CTETRA_SUBMESH_H


/* linalg/balance.c
 *
 * Copyright (C) 2006 Patrick Alken
 *
 * 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.
 */

/* Balance a general matrix by scaling the rows and columns, so the
 * new row and column norms are the same order of magnitude.
 *
 * B =  D^-1 A D
 *
 * where D is a diagonal matrix
 *
 * This is necessary for the unsymmetric eigenvalue problem since the
 * calculation can become numerically unstable for unbalanced
 * matrices.
 *
 * See Golub & Van Loan, "Matrix Computations" (3rd ed), Section 7.5.7
 * and Wilkinson & Reinsch, "Handbook for Automatic Computation", II/11 p320.
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>

#include <gsl/gsl_linalg.h>

#define FLOAT_RADIX       2.0
#define FLOAT_RADIX_SQ    (FLOAT_RADIX * FLOAT_RADIX)

int
gsl_linalg_balance_matrix(gsl_matrix * A, gsl_vector * D)
{
  const size_t N = A->size1;

  if (N != D->size)
    {
      GSL_ERROR ("vector must match matrix size", GSL_EBADLEN);
    }
  else
    {
      double row_norm,
             col_norm;
      int not_converged;
      gsl_vector_view v;

      /* initialize D to the identity matrix */
      gsl_vector_set_all(D, 1.0);

      not_converged = 1;

      while (not_converged)
        {
          size_t i, j;
          double g, f, s;

          not_converged = 0;

          for (i = 0; i < N; ++i)
            {
              row_norm = 0.0;
              col_norm = 0.0;

              for (j = 0; j < N; ++j)
                {
                  if (j != i)
                    {
                      col_norm += fabs(gsl_matrix_get(A, j, i));
                      row_norm += fabs(gsl_matrix_get(A, i, j));
                    }
                }

              if ((col_norm == 0.0) || (row_norm == 0.0))
                {
                  continue;
                }

              g = row_norm / FLOAT_RADIX;
              f = 1.0;
              s = col_norm + row_norm;

              /*
               * find the integer power of the machine radix which
               * comes closest to balancing the matrix
               */
              while (col_norm < g)
                {
                  f *= FLOAT_RADIX;
                  col_norm *= FLOAT_RADIX_SQ;
                }

              g = row_norm * FLOAT_RADIX;

              while (col_norm > g)
                {
                  f /= FLOAT_RADIX;
                  col_norm /= FLOAT_RADIX_SQ;
                }

              if ((row_norm + col_norm) < 0.95 * s * f)
                {
                  not_converged = 1;

                  g = 1.0 / f;

                  /*
                   * apply similarity transformation D, where
                   * D_{ij} = f_i * delta_{ij}
                   */

                  /* multiply by D^{-1} on the left */
                  v = gsl_matrix_row(A, i);
                  gsl_blas_dscal(g, &v.vector);

                  /* multiply by D on the right */
                  v = gsl_matrix_column(A, i);
                  gsl_blas_dscal(f, &v.vector);

                  /* keep track of transformation */
                  gsl_vector_set(D, i, gsl_vector_get(D, i) * f);
                }
            }
        }

      return GSL_SUCCESS;
    }
} /* gsl_linalg_balance_matrix() */

/*
gsl_linalg_balance_accum()
  Accumulate a balancing transformation into a matrix.
This is used during the computation of Schur vectors since the
Schur vectors computed are the vectors for the balanced matrix.
We must at some point accumulate the balancing transformation into
the Schur vector matrix to get the vectors for the original matrix.

A -> D A

where D is the diagonal matrix

Inputs: A - matrix to transform
        D - vector containing diagonal elements of D
*/

int
gsl_linalg_balance_accum(gsl_matrix *A, gsl_vector *D)
{
  const size_t N = A->size1;

  if (N != D->size)
    {
      GSL_ERROR ("vector must match matrix size", GSL_EBADLEN);
    }
  else
    {
      size_t i;
      double s;
      gsl_vector_view r;

      for (i = 0; i < N; ++i)
        {
          s = gsl_vector_get(D, i);
          r = gsl_matrix_row(A, i);

          gsl_blas_dscal(s, &r.vector);
        }

      return GSL_SUCCESS;
    }
} /* gsl_linalg_balance_accum() */


/**
 *
 * @file example_sgelqf.c
 *
 *  PLASMA testing routines
 *  PLASMA is a software package provided by Univ. of Tennessee,
 *  Univ. of California Berkeley and Univ. of Colorado Denver
 *
 * @brief Example using LQ factorization
 *
 * @version 2.6.0
 * @author Bilel Hadri
 * @date 2010-11-15
 * @generated s Tue Jan  7 11:45:20 2014
 *
 **/
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

#include <plasma.h>
#include <cblas.h>
#include <lapacke.h>
#include <core_blas.h>

#ifndef max
#define max(a, b) ((a) > (b) ? (a) : (b))
#endif
#ifndef min
#define min(a, b) ((a) < (b) ? (a) : (b))
#endif

int check_orthogonality(int, int, int, float*);
int check_factorization(int, int, float*, float*, int, float*);

int IONE=1;
int ISEED[4] = {0,0,0,1};   /* initial seed for slarnv() */

int main ()
{

    int cores = 2;
    int M     = 10;
    int N     = 15;
    int LDA   = 10;
    int K = min(M, N);
    int info;
    int info_ortho, info_factorization;
    int i,j;
    int LDAxN = LDA*N;

    float *A1 = (float *)malloc(LDA*N*sizeof(float));
    float *A2 = (float *)malloc(LDA*N*sizeof(float));
    float *Q  = (float *)malloc(LDA*N*sizeof(float));
    PLASMA_desc *T;

    /* Check if unable to allocate memory */
    if ((!A1)||(!A2)||(!Q)){
        printf("Out of Memory \n ");
        return EXIT_SUCCESS;
    }

    /* Plasma Initialization */
    PLASMA_Init(cores);
    printf("-- PLASMA is initialized to run on %d cores. \n",cores);

    /* Allocate T */
    PLASMA_Alloc_Workspace_sgelqf(M, N, &T);

    /* Initialize A1 and A2 */
    LAPACKE_slarnv_work(IONE, ISEED, LDAxN, A1);
    for (i = 0; i < M; i++)
        for (j = 0; j < N; j++)
            A2[LDA*j+i] = A1[LDA*j+i] ;

    /* Factorization QR of the matrix A2 */
    info = PLASMA_sgelqf(M, N, A2, LDA, T);

    /* Building the economy-size Q */
    memset((void*)Q, 0, LDA*N*sizeof(float));
    for (i = 0; i < K; i++)
        Q[LDA*i+i] = 1.0;

    PLASMA_sorglq(M, N, K, A2, LDA, T, Q, LDA);

    /* Check the orthogonality, factorization and the solution */
    info_ortho = check_orthogonality(M, N, LDA, Q);
    info_factorization = check_factorization(M, N, A1, A2, LDA, Q);

    if ((info_ortho != 0)|(info_factorization != 0)|(info != 0))
       printf("-- Error in SGELQF example ! \n");
    else
       printf("-- Run of SGELQF example successful ! \n");

    free(A1); free(A2); free(Q); free(T);

    PLASMA_Finalize();

    return EXIT_SUCCESS;
}

/*-------------------------------------------------------------------
 * Check the orthogonality of Q
 */

int check_orthogonality(int M, int N, int LDQ, float *Q)
{
    float alpha, beta;
    float normQ;
    int info_ortho;
    int i;
    int minMN = min(M, N);
    float eps;
    float *work = (float *)malloc(minMN*sizeof(float));

    eps = LAPACKE_slamch_work('e');
    alpha = 1.0;
    beta  = -1.0;

    /* Build the idendity matrix USE DLASET?*/
    float *Id = (float *) malloc(minMN*minMN*sizeof(float));
    memset((void*)Id, 0, minMN*minMN*sizeof(float));
    for (i = 0; i < minMN; i++)
        Id[i*minMN+i] = (float)1.0;

    /* Perform Id - Q'Q */
    if (M >= N)
        cblas_ssyrk(CblasColMajor, CblasUpper, CblasTrans, N, M, alpha, Q, LDQ, beta, Id, N);
    else
        cblas_ssyrk(CblasColMajor, CblasUpper, CblasNoTrans, M, N, alpha, Q, LDQ, beta, Id, M);

    normQ = LAPACKE_slansy_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), 'u', minMN, Id, minMN, work);

    printf("============\n");
    printf("Checking the orthogonality of Q \n");
    printf("||Id-Q'*Q||_oo / (N*eps) = %e \n",normQ/(minMN*eps));

    if ( isnan(normQ / (minMN * eps)) || (normQ / (minMN * eps) > 10.0) ) {
        printf("-- Orthogonality is suspicious ! \n");
        info_ortho=1;
    }
    else {
        printf("-- Orthogonality is CORRECT ! \n");
        info_ortho=0;
    }

    free(work); free(Id);

    return info_ortho;
}

/*------------------------------------------------------------
 *  Check the factorization QR
 */

int check_factorization(int M, int N, float *A1, float *A2, int LDA, float *Q)
{
    float Anorm, Rnorm;
    float alpha, beta;
    int info_factorization;
    int i,j;
    float eps;

    eps = LAPACKE_slamch_work('e');

    float *Ql       = (float *)malloc(M*N*sizeof(float));
    float *Residual = (float *)malloc(M*N*sizeof(float));
    float *work              = (float *)malloc(max(M,N)*sizeof(float));

    alpha=1.0;
    beta=0.0;

    if (M >= N) {
        /* Extract the R */
        float *R = (float *)malloc(N*N*sizeof(float));
        memset((void*)R, 0, N*N*sizeof(float));
        LAPACKE_slacpy_work(LAPACK_COL_MAJOR,'u', M, N, A2, LDA, R, N);

        /* Perform Ql=Q*R */
        memset((void*)Ql, 0, M*N*sizeof(float));
        cblas_sgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, M, N, N, (alpha), Q, LDA, R, N, (beta), Ql, M);
        free(R);
    }
    else {
        /* Extract the L */
        float *L = (float *)malloc(M*M*sizeof(float));
        memset((void*)L, 0, M*M*sizeof(float));
        LAPACKE_slacpy_work(LAPACK_COL_MAJOR,'l', M, N, A2, LDA, L, M);

    /* Perform Ql=LQ */
        memset((void*)Ql, 0, M*N*sizeof(float));
        cblas_sgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, M, N, M, (alpha), L, M, Q, LDA, (beta), Ql, M);
        free(L);
    }

    /* Compute the Residual */
    for (i = 0; i < M; i++)
        for (j = 0 ; j < N; j++)
            Residual[j*M+i] = A1[j*LDA+i]-Ql[j*M+i];

    Rnorm = LAPACKE_slange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), M, N, Residual, M, work);
    Anorm = LAPACKE_slange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), M, N, A2, LDA, work);

    if (M >= N) {
        printf("============\n");
        printf("Checking the QR Factorization \n");
        printf("-- ||A-QR||_oo/(||A||_oo.N.eps) = %e \n",Rnorm/(Anorm*N*eps));
    }
    else {
        printf("============\n");
        printf("Checking the LQ Factorization \n");
        printf("-- ||A-LQ||_oo/(||A||_oo.N.eps) = %e \n",Rnorm/(Anorm*N*eps));
    }

    if (isnan(Rnorm / (Anorm * N *eps)) || (Rnorm / (Anorm * N * eps) > 10.0) ) {
        printf("-- Factorization is suspicious ! \n");
        info_factorization = 1;
    }
    else {
        printf("-- Factorization is CORRECT ! \n");
        info_factorization = 0;
    }

    free(work); free(Ql); free(Residual);

    return info_factorization;
}


/*
 * Kaczmarz down-up sweep on diagonally banded matrix. The essential loop is:
 *
 * for each row i
 *     x = x + w*(b(i) - A(i,:)*x)*A(i,:)'
 * end
 *
 * The matrix is given in band storage format, where each row (stored
 * contiguously in memory) of the array R stores a diagonal of the matrix with
 * offset idx(j), such that
 *
 * 	   A(i,i+idx(j)) = R(i,j)
 *
 * use (from MATLAB):
 * 		y = sweepR_mex(R,idx,x,b,w,dir)
 *
 * 		R			- matrix of diagonals of matrix A
 *		idx			- offsets of diagonals
 *		x			- initial guess
 * 		b 			- right hand side (source)
 *		w 			- relaxation parameter (0 <= w <= 2)
 *		ns 		    - number of RHS
 * 		n_threads	- OPTIONAL argument to control the number of execution
 * 		threads solving CARP blocks in parallel. The number of threads can also
 * 		be defined via an environment variable (OMP_NUM_THREADS), but this
 * 		optional argument takes precedence. The default number of threads is
 * 		one. Take care if using more than one MATLAB worker per node: each
 * 		MATLAB worker will use OMP_NUM_THREADS, so if there are four workers on
 * 		a node, there will be 4 x OMP_NUM_THREADS parallel CARP sweeps.
 *  compile :
 *	   mex -largeArrayDims sweep_MT_mex.c -DDEFINEUNIX -lmwblas CFLAGS="\$CFLAGS -fopenmp" LDFLAGS="\$LDFLAGS -fopenmp"
 *
 * Author: Mathias Louboutin from Art Petrenko sweepR_mex.c
 *         Seismic Laboratory for Imaging and Modeling
 *         Department of Earth, Ocean, and Atmosperic Sciences
 *         The University of British Columbia
 *
 * Date: March, 2015

 * You may use this code only under the conditions and terms of the
 * license contained in the file LICENSE provided with this source
 * code. If you do not agree to these terms you may not use this
 * software.
*/

#include <stdlib.h> /* for getenv */
#include <stddef.h> /* for size_t type */
#include <string.h> /* for memcpy */
#include <pthread.h> /* for threading */
#include <omp.h>

/* The following section allows this file to compile on Mac OS X 10.8.5. Pass
 * the flag -DARCH_MACI64 to the compiler to activate it. */
#ifdef ARCH_MACI64
#include <mach/error.h>
typedef wchar_t char16_t;
#else /* not ARCH_MACI64 */
#include <error.h>
#endif /* ARCH_MACI64 */

#include <math.h>
#include <mex.h>
#include <matrix.h>
#include <blas.h>

struct copy_init_guess_data_t
{
	double *copy_src_real, *copy_dst_real;
	double *copy_src_imag, *copy_dst_imag;
	long n_to_copy;
};

struct sweep_data_t
{
	long start_row, end_row, ncol, ny, nx, haloWidth, main_diagonal_offset;
	double *Rr, *Ri, *yr, *yi, *br, *bi;
	long *idx;
	double w;
	int dir, n_threads, ns;
};

struct average_data_t
{
	double *copy_src_real, *copy_dst_real;
	double *copy_src_imag, *copy_dst_imag;
	double *halo_1_real, *halo_2_real;
	double *halo_1_imag, *halo_2_imag;
	double *halo_dst_real, *halo_dst_imag;
	long n_to_copy, n_in_halo;
};

struct thread_data_t
{
	struct copy_init_guess_data_t copy_init_guess_data;
	struct sweep_data_t sweep_data;
	struct average_data_t average_data;
	pthread_barrier_t *barrier;
};

void *do_sweep(void *thread_args_void)
{
    struct sweep_data_t *thread_args;
	/* Variables contained in thread_args_void struct */
	long start_row, end_row, ncol, ny, nx;
	/* Rr and Ri are pointers to short fat ncol-by-N matrices */
	double *Rr, *Ri, *yr, *yi, *br, *bi;
	long *idx;
	double w;
	int n_threads,ns;
	/* Temporary storage variables */
	double cr = 0, ci = 0;
	long offset, main_diagonal_offset;

	/* Assign local pointers to data locations in shared memory */
	thread_args = (struct sweep_data_t *) thread_args_void;
	start_row 	= thread_args->start_row;
	end_row 	= thread_args->end_row;
	ncol 		= thread_args->ncol;
	ny 			= thread_args->ny;
	nx			= thread_args->nx;
	main_diagonal_offset = thread_args->main_diagonal_offset;
	Rr 			= thread_args->Rr;
	Ri 			= thread_args->Ri;
	idx 		= thread_args->idx;
	yr 			= thread_args->yr;
	yi 			= thread_args->yi;
	br 			= thread_args->br;
	bi 			= thread_args->bi;
	w 			= thread_args->w;
	n_threads 	= thread_args->n_threads;
	ns 			= thread_args->ns;
	offset = (start_row == 0 ? 0 : - main_diagonal_offset);

	long i;
	int s;
	long j;
	long k;
	long indj;
	long toto;

	#pragma omp parallel for schedule(static,1) private(k,indj,i,j,cr,ci) num_threads(n_threads)
	for(s=0 ; s<ns; s++)
	{
		/* Kaczmarz sweep on one row block */
		for(i = start_row ; i<end_row ;i++ )
		{

			if (0 <= i + main_diagonal_offset && i + main_diagonal_offset < nx){
				cr = br[i + main_diagonal_offset+s*nx];
				ci = bi[i + main_diagonal_offset+nx*s];
			}
			else{
				error(1,0,"Discovery of whether the iterate vector is haloed failed.");
			}
			/* First loop over non-zero row elements calculates inner product
			 * of matrix row and CARP iterate */
			long diff = i - start_row + offset;
			long icol=i*ncol;

			for(j=0 ; j<ncol;j++)
			{
				/* i + idx[j] is the column index for the full Helmholtz matrix.
				 * k is the index into the vector representing the CARP iterate
				 * of the given block. */
				k = diff + idx[j];
				indj=icol + j;
				if(0<=k && k<ny)
				{
					cr -= Rr[indj]*yr[k+s*nx] - Ri[indj]*yi[k+s*nx];
					ci -= Rr[indj]*yi[k+s*nx] + Ri[indj]*yr[k+s*nx];

				}
			}

			cr*=w;
			ci*=w;
			/* Second loop over non-zero row elements updates Karkzmarz iterate */

			for(j=0 ; j<ncol;j++)
			{
				k = diff + idx[j];
				indj=icol + j;
				if(0<=k && k<ny)
				{
					yr[k+s*nx] +=   cr*Rr[indj] + ci*Ri[indj];
					yi[k+s*nx] +=  -cr*Ri[indj] + ci*Rr[indj];

				}
			}

		}
		cr = 0;
		ci = 0;
		for(i = end_row-1 ; i>start_row-1 ;i-- )
		{

			if (0 <= i + main_diagonal_offset && i + main_diagonal_offset < nx){
				cr = br[i + main_diagonal_offset+s*nx];
				ci = bi[i + main_diagonal_offset+nx*s];
			}
			else{
				error(1,0,"Discovery of whether the iterate vector is haloed failed.");
			}
			/* First loop over non-zero row elements calculates inner product
			 * of matrix row and CARP iterate */
			long diff = i - start_row + offset;
			long icol=i*ncol;

			for(j=0 ; j<ncol;j++)
			{
				/* i + idx[j] is the column index for the full Helmholtz matrix.
				 * k is the index into the vector representing the CARP iterate
				 * of the given block. */
				k = diff + idx[j];
				indj=icol + j;
				if(0<=k && k<ny)
				{
					cr -= Rr[indj]*yr[k+s*nx] - Ri[indj]*yi[k+s*nx];
					ci -= Rr[indj]*yi[k+s*nx] + Ri[indj]*yr[k+s*nx];

				}
			}

			cr*=w;
			ci*=w;
			/* Second loop over non-zero row elements updates Karkzmarz iterate */

			for(j=0 ; j<ncol;j++)
			{
				k = diff + idx[j];
				indj=icol + j;
				if(0<=k && k<ny)
				{
					yr[k+s*nx] +=   cr*Rr[indj] + ci*Ri[indj];
					yi[k+s*nx] +=  -cr*Ri[indj] + ci*Rr[indj];

				}
			}

		}
	}

	return NULL;
}

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
	/* structs to hold all arguments to each thread in one variable */
	struct sweep_data_t **thread_args_sweep = NULL;
	struct copy_init_guess_data_t **copy_init_guess_data = NULL;

	mwSize ncol, nx;
	ptrdiff_t ncolBlas, idxIncBlas = 1, maxIdxLoc = 0;

	double *Rr,*Ri,*idxd = NULL,*xr,*xi,*br,*bi,*yr,*yi;
	long *idx = NULL;
    double w = 0;
	int ns;
	mwSize n_threads = 1;
	char *n_threads_str = NULL;
	mwSize N=1, numGridPointsPerBlock, main_diagonal_offset;

	/* Flags that are set if memory is allocated within the MEX file */
	int Ri_alloc=0, xi_alloc=0, bi_alloc=0;
	mwSize *seg_bounds_hi, *seg_bounds_mid, *seg_bounds_row, *seg_bounds_lo;

	/* Read input arguments; initialize complex part to zero if input is real. */
	ns = lrint(mxGetScalar(prhs[5]));
	N = mxGetN(prhs[0]);
    ncol = mxGetM(prhs[0]);
	ncolBlas = (ptrdiff_t)ncol;
	Rr  = mxGetPr(prhs[0]);
    if(mxIsComplex(prhs[0])){
        Ri  = mxGetPi(prhs[0]);
    }
    else{
        Ri = mxCalloc(N*ncol,sizeof(double));
		Ri_alloc = 1;
    }
	idxd = mxGetPr(prhs[1]);
    nx  = mxGetM(prhs[2]);
	xr  = mxGetPr(prhs[2]);
    if(mxIsComplex(prhs[2])){
        xi  = mxGetPi(prhs[2]);
    }
    else{
        xi = mxCalloc(nx*ns,sizeof(double));
		xi_alloc = 1;
    }
	br  = mxGetPr(prhs[3]);
    if(mxIsComplex(prhs[3])){
        bi  = mxGetPi(prhs[3]);
    }
    else{
        bi = mxCalloc(nx*ns,sizeof(double));
		bi_alloc = 1;
    }
    if (mxGetM(prhs[3]) != nx){
    mexErrMsgIdAndTxt("SLIM_release_apps:tools:algorithms:ThreeDFreqModeling:sweepR_mex:NumElements",
				"The number of elements in the iterate and right hand side vectors must be equal.");
	}

	/* Allocate the final output vector */
	plhs[0]  = mxCreateDoubleMatrix(nx,ns, mxCOMPLEX);
	yr       = mxGetPr(plhs[0]);
    yi       = mxGetPi(plhs[0]);

	/* Check to make sure memory was allocated correctly */
	if (Rr==NULL || Ri==NULL || idxd==NULL || xr==NULL || xi==NULL || br==NULL || bi==NULL ||
	    yr==NULL || yi==NULL){
		mexErrMsgIdAndTxt("SLIM_release_apps:tools:algorithms:ThreeDFreqModeling:sweepR_mex:OutOfMemory",
				"Could not allocate memory for main computational variables.");
	}

	if ((idx = (long *) mxCalloc(ncol, sizeof(long))) == NULL){
		mexErrMsgIdAndTxt("SLIM_release_apps:tools:algorithms:ThreeDFreqModeling:sweepR_mex:OutOfMemory",
				"Could not allocate memory for main computational variables.");
	}
	mwSize i;
	for (i=0; i < ncol; i++){
		idx[i] = lrint(idxd[i]);
	}

	/* The default value for the number of threads can be overridden by an
	 * environment variable. */
	n_threads_str = getenv("OMP_NUM_THREADS");
	if (n_threads_str == NULL){
		n_threads = 1;
	}
	else{
		n_threads = strtol(n_threads_str, NULL, 10);
		if(n_threads < 1){
			n_threads = 1;
		}
	}
	/* The environment variable can in turn be overridden by an optional
	 * argument to the mexFunction. */
	if (nrhs >= 7){
		if(1 <= lrint(mxGetScalar(prhs[6]))){
			n_threads = lrint(mxGetScalar(prhs[6]));
		}
	}

	/* printf("Using %d threads \n",n_threads);
	/* Partition the iterate vector into blocks. Note that the below
	 * partitioning scheme is slighlty different from that in pCARPCG.m in this
	 * directory. The partitioning scheme of pCARPCG corresponds to
	 * distributing a three dimensional array with dimensions given by n
	 * according to Matlab's codistributor1d.defaultPartition(n(3)), and then
	 * vectorizing it. The partition scheme of the present file instead uses
	 * Matlab's codistributor1d.defaultPartion(prod(n)). In other words,
	 * pCARPCG divides the iterate into blocks along the slow dimension,
	 * whereas sweepR_mex.c does not take dimensionality into account, only the
	 * total number of gridpoints. This is done to avoid needing an extra input
	 * parameter with the the system dimensions. The seg_bounds_hi, _lo and
	 * _mid arrays contain indices into non-haloed vectors, while the
	 * seg_bounds_row array contains indices to the rows of the system matrix.
	 *
		 * yr_seg[i_thread-1]        overlap      yr_seg[i_thread]
	 * ------------------------|-----|-----|-------------------------------
	 *         .----------------^     |     ^-------------------.
	 *    seg_bounds_lo[i_thread], seg_bounds_mid[i_thread], seg_bounds_hi[i_thread]
	 */
	numGridPointsPerBlock = N;
	seg_bounds_hi  = (mwSize *)mxCalloc(2,sizeof(mwSize));
	seg_bounds_mid = (mwSize *)mxCalloc(2,sizeof(mwSize));
	seg_bounds_lo  = (mwSize *)mxCalloc(2,sizeof(mwSize));
	seg_bounds_row = (mwSize *)mxCalloc(2,sizeof(mwSize));
	if (N == nx){
		main_diagonal_offset = 0;
	}
	else{
		/* The vector is haloed. We are only able to correctly process matrices
		 * with a non-zero main diagonal and symmetric off-main diagonal offsets. */
		if (ncol % 2 != 1){
			mexErrMsgIdAndTxt("SLIM_release_apps:tools:algorithms:ThreeDFreqModeling:sweepR_mex:EvenNumberOfDiags",
					"Input iterate vector appears to be haloed but there is an even number of non-zero diagonals in the system matrix.");
		}
		main_diagonal_offset = idx[ncol/2];
		mwSize i;
		for (i = 1; i <= ncol/2; i++){
			if (idx[ncol/2 + i] - main_diagonal_offset != -(idx[ncol/2 - i] - main_diagonal_offset)){
				mexErrMsgIdAndTxt("SLIM_release_apps:tools:algorithms:ThreeDFreqModeling:sweepR_mex:DiagsNotSymmetric",
						"Input iterate vector appears to be haloed but the pattern of non-zero diagonals in the system matrix is not symmetric.");
			}
		}
	}

	seg_bounds_hi[0]  = 0;
	seg_bounds_mid[0] = 0;
	seg_bounds_row[0] = 0;
	seg_bounds_lo[0]  = 0;
	seg_bounds_lo[1]  = nx;
	seg_bounds_hi[1]  = nx;
	seg_bounds_mid[1] = nx;
	seg_bounds_row[1] = N;

	thread_args_sweep = (struct sweep_data_t **)mxCalloc(1, sizeof(struct sweep_data_t *));

	/* Set thread arguments */
	thread_args_sweep[0] = (struct sweep_data_t *)mxCalloc(1,sizeof(struct sweep_data_t));
	thread_args_sweep[0]->start_row 	= seg_bounds_row[0];
	thread_args_sweep[0]->end_row   	= seg_bounds_row[1];
	thread_args_sweep[0]->ncol 	 	= ncol;
	thread_args_sweep[0]->ny 	 	= seg_bounds_hi[1]-seg_bounds_lo[0];
	thread_args_sweep[0]->nx			= nx;
	thread_args_sweep[0]->main_diagonal_offset = main_diagonal_offset;
	thread_args_sweep[0]->Rr  		= Rr;
	thread_args_sweep[0]->Ri  		= Ri;
	thread_args_sweep[0]->idx 		= idx;
	thread_args_sweep[0]->yr	 	= yr;
	thread_args_sweep[0]->yi	 	= yi;
	thread_args_sweep[0]->br 		= br;
	thread_args_sweep[0]->bi 		= bi;
	thread_args_sweep[0]->w 		 = w;
	thread_args_sweep[0]->n_threads = n_threads;
	thread_args_sweep[0]->ns = ns;
	/* Set the initial guess directly in the output array too */
	memcpy((void *)yr, (void *)xr, sizeof(double)*nx*ns);
	memcpy((void *)yi, (void *)xi, sizeof(double)*nx*ns);
	do_sweep((void *)thread_args_sweep[0]);

	/* Free memory if it was allocated within the MEX file. */
	if (Ri_alloc){
		mxFree(Ri);
	}
	if (xi_alloc){
		mxFree(xi);
	}
	if (bi_alloc){
		mxFree(bi);
	}

	mxFree(idx);
	mxFree(thread_args_sweep);

	/* Don't think I need pthread_exit() here, because pthread_join is called above */
	return;
}


#ifndef ALM_I_APG_H
#define ALM_I_APG_H

#include "Matrix.h"
#include "APPROX2.h"
#include <string>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <stdio.h>      /* printf */
#include <time.h>
#include <fstream>
#include <algorithm>
#include <iomanip>
#include <ctime>
#include <math.h>

//This class solves problem of the form f(x)+g(x) under the constraint Mx=c;
// where f(x)=\sum_{j=1}^m lambda_f[j] \phi_j(<A_j,x>)
//and g(x)=sum_{i=1}^n g_i(x_i). We all assume that each \phi_j is 1-smooth.

// Each subproblem solves problem of the form f(x)+<Mx-c, lambda_s>+1/2beta_s\|Mx-c\|^2+g(x) by APG.
// This header file implements ASGARD_DL.

template<typename L, typename D>
class ALM_I_APG: public APPROX2<L, D>
{
private:

  std::vector<D> Au;
  std::vector<D> Az;

  std::vector<D> Mu;
  std::vector<D> Mz;

  std::vector<D> Mx_s;
  std::vector<D> Ax_s;
  std::vector<D> lambda_f;

protected:
  Matrix<L,D> data_A;

  Matrix<L, D> data_M;

  std::vector<D> x_s;

  std::vector<D> lambda_s;

  std::vector<D> M_tlambda_s;

  D beta_s;

  D epsilon_s;

  D m_s;

  D mp_s;

  D m_0;

  L m_1;

  L m_2;

  D c_s;

  D eta;

  D rho;

  D max_Lf_s;

  D max_M_s;

  D Lf;

  D L_m;

  L val_mu_f;
  L val_mu_g;

  std::vector<D> dual_alpha1;   //dual_alpha1[j]=gradient_of_phi_j(<A_j,x>)

  std::vector<D> dual_alpha2;  //dual_alpha2[j]=(<M_j,x>-c_j)

  std::vector<D> baralpha1;  //baralpha1=-sum_{j=1}^m_1 lambda_f[j]*dual_alpha1[j]*A_j

  std::vector<D> baralpha2;   //baralpha2=-sum_{j=1}^m_2 1/beta_s*dual_alpha2[j]*M_j

  D residual1;    //residual1=\dist(-M^\top y, \nabla f(x)+\partial g(x))

  D residual2;    //residual2=\|Mx-c\|

  D function_value;

  L print_every_N_ALM_I_APG;

  D running_time_ALM_I_APG;

  L nb_outer_iters;

  std::vector<D> gradient_of_f;

  ofstream samp_ALM_I_APG;
public:

  virtual inline D gradient_of_phi_j(D, L){return D(NULL);}

  virtual inline D value_of_g_i(D, L){return D(NULL);}
  virtual inline D value_of_phi_j(D, L){return D(NULL);}

  virtual inline D prox_of_g_i(D,D,D, L){return D(NULL);}

  virtual inline D feasible_dual(std::vector<D> &,std::vector<D> &) {return D(NULL);}

  virtual inline D value_of_phistar_i(D,L) {return D(NULL);}
  virtual inline D value_of_g_tilde_star(D, vector<D> &,vector<D> &){return D(NULL);}

  virtual inline void set_matrix_M(){}

  virtual inline void set_matrix_A(){}

  virtual inline D distance_to_subgradient_of_g(){return D(NULL);}
  //virtual inline D distance_to_subgradient_of_g(std::vector<D> &, std::vector<D> & ){return D(NULL);}

  ALM_I_APG(const char* Matrix_file, D val_lambda_f)
  :data_A(), data_M()
  {

  }

  ALM_I_APG()
  :data_A(), data_M()
  {

  }

  L get_nb_features(){
    return data_A.get_d();
  }

   //This function computes \nabla_i f(gamma u +z)
    inline D partial_i_of_f(L i){
      D res=0;
      for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
      {
        L j=data_A.col_idx[k];
        D tmp=lambda_f[j]*gradient_of_phi_j(this->gamma*Au[j]+Az[j], j);
        res+=data_A.A_t[k]*tmp;
      }
      for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
      {
        L j=data_M.col_idx[k];
        D tmp=(this->gamma*Mu[j]+Mz[j]-data_M.b[j])/beta_s+lambda_s[j];
        res+=data_M.A_t[k]*tmp;
      }
      return res;
    }

   //This function computes argmin{x1 t+x2 t^2/2+g_i(x3+t)}
   //inline D compute_prox(D x1, D x2, D x3, L i){
     //return prox_operA_tor(x1,x2,x3,i);
   //}

    void compute_x(){
        for(L i=0;i<this->n;i++){
           this->x[i]=this->gamma*this->u[i]+this->z[i];
           x_s[i]=this->x[i];
          }
        for(L j=0;j<m_2;j++)
            Mx_s[j]=this->gamma*this->Mu[j]+this->Mz[j];
        for(L j=0;j<m_1;j++)
            Ax_s[j]=this->gamma*this->Au[j]+this->Az[j];
    }

   inline void compute_primal_value() {

     D res=0;
     for(L i=0;i<this->n;i++){
       this->x[i]=this->gamma*this->u[i]+this->z[i];
       res+=value_of_g_i(this->x[i],i);
     }

     for(L j=0;j<m_1;j++){
       res+=lambda_f[j]*value_of_phi_j(this->gamma*Au[j]+Az[j],j);
     }
     for(L j=0;j<m_2;j++){
       D tmp=(this->gamma*Mu[j]+Mz[j]-data_M.b[j]+beta_s*lambda_s[j]);
       res+=tmp*tmp/2/beta_s;
     }
     this->primal_value=res;
   }

   inline void compute_dual_value(){
     for(L j=0;j<m_1;j++){
       dual_alpha1[j]=gradient_of_phi_j(this->gamma*Au[j]+Az[j], j);
     }
     for(L j=0;j<m_2;j++){
       dual_alpha2[j]=this->gamma*Mu[j]+Mz[j]-data_M.b[j]+beta_s*lambda_s[j];
     }
     for(L i=0;i<this->n;i++){
       baralpha1[i]=0;
       for(L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++){
         L j=data_A.col_idx[k];
         baralpha1[i]-=dual_alpha1[j]*data_A.A_t[k]*lambda_f[j];
       }
       baralpha2[i]=0;
       for(L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++){
         L j=data_M.col_idx[k];
         baralpha2[i]-=dual_alpha2[j]*data_M.A_t[k]/beta_s;
       }
     }
     D res=0;
     D scal=feasible_dual(baralpha1, baralpha2);
     for(L j=0;j<m_1;j++)
     {
       res-=value_of_phistar_i(scal*dual_alpha1[j],j)*lambda_f[j];
     }
     for(L j=0;j<m_2;j++)
       res-=scal*dual_alpha2[j]*(0.5*scal*dual_alpha2[j]+data_M.b[j]-beta_s*lambda_s[j])/beta_s;

     res-=value_of_g_tilde_star(scal,baralpha1,baralpha2);
     this->dual_value=res;
   }

   inline void compute_function_value() {
     D res=0;
     for(L i=0;i<this->n;i++){
       res+=value_of_g_i(x_s[i],i);
     }
     for(L j=0;j<m_1;j++){
       res+=lambda_f[j]*value_of_phi_j(Ax_s[j],j);
     }
     function_value=res;
   }

   D compute_prox(D x1,D x2,D x3, L i){
       D res=prox_of_g_i(x1+beta_s*(x3-x_s[i]), beta_s+x2,x3,i);
       if(fabs(res)>1e10)
       {cout<<"larger than 1e10:"<<res<<endl;
       cout<<beta_s<<"; "<<x1<<"; "<<x3<<"; "<<x_s[i]<<"; "<<i<<endl;
       D res2=0;
      for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
      {
        L j=data_A.col_idx[k];
        D tmp=lambda_f[j]*gradient_of_phi_j(this->gamma*Au[j]+Az[j], j);
        res2+=data_A.A_t[k]*tmp;
        cout<<"Ax["<<j<<"]="<<this->gamma*Au[j]+Az[j]<<endl;
      }
      for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
      {
        L j=data_M.col_idx[k];
        D tmp=(this->gamma*Mu[j]+Mz[j]-data_M.b[j])/beta_s+lambda_s[j];
        cout<<"Mu["<<j<<"]="<<Mu[j]<<"; gamma="<<this->gamma<<endl;
        res2+=data_M.A_t[k]*tmp;
      }
       }
       return res;
   }

   inline void compute_gradient_norm(){
     std::vector<D> Tx(this->n,0);
     this->do_single_step_prox(Tx);
     D res=0;
     for(L i=0;i<this->n;i++)
      res+=(Tx[i]-this->x[i])*(Tx[i]-this->x[i]);
     this->gradient_norm=sqrt(res);
   }

   inline void set_v()
   {
     this->v.resize(this->n,0);
     D maxv=0;
     D minv=std::numeric_limits<double>::max();
     D sumv=0;
     D sumvi1=0;
     L sumw=0;
     L maxw=0;
     L minw=this->n;
     this->sumofLi=0;
     for(L j=0;j<m_1;j++)
     {
       sumw+=data_A.w_t[j];
       maxw=max(maxw,data_A.w_t[j]);
       minw=min(minw,data_A.w_t[j]);
     }
     cout<<"sumw="<<sumw<<";  maxw="<<maxw<<"; minw="<<minw<<endl;
     for(L i=0;i<this->n;i++)
     {
       D vi=0;
       D vi1=0;
       for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
       {
         L j=data_A.col_idx[k];
         vi+=(1.+(data_A.w_t[j]-1.)*(this->tau-1.)/max(this->n-1.,1.))*data_A.A_t[k]*data_A.A_t[k]*lambda_f[j];
         vi1+=data_A.A_t[k]*data_A.A_t[k]*lambda_f[j];
       }
       for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
       {
         L j=data_M.col_idx[k];
         vi+=(1.+(data_M.w_t[j]-1.)*(this->tau-1.)/max(this->n-1.,1.))*data_M.A_t[k]*data_M.A_t[k]/beta_s;
         vi1+=data_M.A_t[k]*data_M.A_t[k]/beta_s;
       }
       this->v[i]=vi;
       sumv+=vi;
       sumvi1+=vi1;
       if(maxv<vi) maxv=vi;
       if(minv>vi) minv=vi;
     }
     if(this->tau==this->n){
          for(L i=0;i<this->n;i++)
              this->v[i]=sumvi1;
     }
     this->sumofLi=sumvi1;
     cout<<"  max of v: "<<maxv<<" ;  min of v: "<<minv<<" ;  sumofv: "<<sumv<<" sumofLi="<<this->sumofLi<<endl;
   }

   inline void compute_m0(D beta0, D val_eta, D val_rho, L val_tau){
       cout<<"beta0="<<beta0<<"val_eta="<<val_eta<<"; val_rho="<<val_rho<<"; val_tau="<<val_tau<<endl;
     m_s=this->n/val_tau*20;

     //D tmp7=1-val_tau/this->n*sqrt(beta0/(max_Lf_s+max_M_s/beta0+beta0));
     //cout<<"tmp7: "<<tmp7<<"; "<<val_tau/(this->n+0.)*sqrt(beta0/(max_Lf_s+max_M_s/beta0+beta0))<<endl;
     //m_0=(2*log(val_rho)+log(val_eta)+log(2))/log(tmp7);

     cout<<" m_s="<<m_s<<endl;
   }

  inline void rescale_Matrix(){
    Lf=0;
    max_Lf_s=0;
    max_M_s=0;
    L_m=0;
    cout<<"this->n="<<this->n<<endl;
    for(L i=0;i<this->n;i++)
    {
      D vi=0;
      for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
      {
        L j=data_A.col_idx[k];
        Lf+=data_A.A_t[k]*data_A.A_t[k]*lambda_f[j];
        vi+=(1.+(data_A.w_t[j]-1.)*(this->tau-1.)/max(this->n-1.,1.))*data_A.A_t[k]*data_A.A_t[k]*lambda_f[j];
      }
      if(vi>max_Lf_s) max_Lf_s=vi;
      D vi2=0;
      for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
      {
        L j=data_M.col_idx[k];
        L_m+=data_M.A_t[k]*data_M.A_t[k];
        vi2+=(1.+(data_M.w_t[j]-1.)*(this->tau-1.)/max(this->n-1.,1.))*data_M.A_t[k]*data_M.A_t[k];
      }
      if(vi2>max_M_s) {max_M_s=vi2;cout<<vi2<<" ";}
    }
    cout<<"max_M_s"<<max_M_s<<endl;
    D scal=sqrt((Lf+1)/L_m);
    cout<<"scal="<<scal<<endl;
    if(scal>1)
    {
      for(L i=0;i<this->n;i++){
        for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
        {
          data_M.A_t[k]*=scal;
        }
      }
      for(L j=0;j<m_2;j++){
        for (L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++)
        {
          data_M.A[k]*=scal;
        }
      }
      for(L j=0;j<m_2;j++)
        data_M.b[j]*=scal;
      max_M_s*=scal*scal;
      L_m*=scal*scal;
    }
    cout<<"max_M_s="<<max_M_s<<"; L_m="<<L_m<<"; max_Lf_s="<<max_Lf_s<<"; L_f="<<Lf<<endl;

  }

   inline void set_p(){
     this->proba_vector.resize(this->n,(0.0+this->tau)/this->n);
     this->max_p=(0.0+this->tau)/this->n;
   }

   void compute_KKT_residual(){
       D res;
       for(L i=0; i<this->n;i++){
           res=0;
           for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
           {
                L j=data_A.col_idx[k];
                D tmp=lambda_f[j]*gradient_of_phi_j(Ax_s[j], j);
                res+=data_A.A_t[k]*tmp;
           }
           gradient_of_f[i]=res;
       }
       residual1=distance_to_subgradient_of_g();
       res=0;
       for(L j=0;j<m_2;j++)
           res+=(Mx_s[j]-data_M.b[j])*(Mx_s[j]-data_M.b[j]);
       residual2=sqrt(res);
   }

   void testing(){
       std::vector<D> nablaf(this->n);
       D res1=0;
       compute_Mty();
       for(L i=0;i<this->n;i++)
       {
          nablaf[i]=0;
          for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
           {
             L j=data_M.col_idx[k];
             D tmp=lambda_s[j];
             nablaf[i]+=data_M.A_t[k]*tmp;
            }
          res1+=(nablaf[i]-M_tlambda_s[i])*(nablaf[i]-M_tlambda_s[i]);
       }
       //samp_ALM2<<" res1 ="<<sqrt(res1)<<endl;
   }

   inline void update_z_coordinate( L i, D dz){
     this->z[i]+=dz;
     for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
     {
       L j=data_A.col_idx[k];
       Az[j]+=dz*data_A.A_t[k];
     }
     for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
     {
       L j=data_M.col_idx[k];
       Mz[j]+=dz*data_M.A_t[k];
       //if(j==2869) cout<<"gamma="<<this->gamma<<"; "<<dz<<"; "<<Mz[j]<<endl;
     }
   }

    inline void update_x_coordinate( L i, D dx){
     this->x[i]+=dx;
     x_s[i]=this->x[i];
     L j;
     for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
     {
       j=data_M.col_idx[k];
       Mx_s[j]+=dx*data_M.A_t[k];
     }
      for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
     {
       j=data_A.col_idx[k];
       Ax_s[j]+=dx*data_A.A_t[k];
     }
   }

   inline void update_u_coordinate( L i, D du){
     this->u[i]+=du;
     for (L k = data_A.ptr_t[i]; k < data_A.ptr_t[i + 1];k++)
     {
       L j=data_A.col_idx[k];
       Au[j]+=du*data_A.A_t[k];
     }
     for (L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++)
     {
       L j=data_M.col_idx[k];
       Mu[j]+=du*data_M.A_t[k];
     }
   }

   void compute_Au(){
     for(L j=0;j<m_1;j++)
       for(L k = data_A.ptr[j]; k < data_A.ptr[j + 1];k++){
         L kj=data_A.row_idx[k];
         Au[j]+=this->u[kj]*data_A.A[k];
       }
   }

   void compute_Mu(){
     for(L j=0;j<m_2;j++)
       for(L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++){
         L kj=data_M.row_idx[k];
         Mu[j]+=this->u[kj]*data_M.A[k];
       }
   }

   void compute_Az(){
     for(L j=0;j<m_1;j++)
       for(L k = data_A.ptr[j]; k < data_A.ptr[j + 1];k++){
         L kj=data_A.row_idx[k];
         Az[j]+=this->z[kj]*data_A.A[k];
       }
   }
   void compute_Mz(){
     for(L j=0;j<m_2;j++)
       for(L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++){
         L kj=data_M.row_idx[k];
         Mz[j]+=this->z[kj]*data_M.A[k];
       }
   }

   void compute_Au(vector<D> & x0){
     for(L j=0;j<m_1;j++)
       for(L k = data_A.ptr[j]; k < data_A.ptr[j + 1];k++){
         L kj=data_A.row_idx[k];
         Au[j]+=x0[kj]*data_A.A[k];
       }

   }

   void compute_Az(vector<D> & x0){
     for(L j=0;j<m_1;j++)
       for(L k = data_A.ptr[j]; k < data_A.ptr[j + 1];k++){
         L kj=data_A.row_idx[k];
         Az[j]+=x0[kj]*data_A.A[k];
       }
   }
   void compute_Mu(vector<D> & x0){
     for(L j=0;j<m_2;j++)
       for(L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++){
         L kj=data_M.row_idx[k];
         Mu[j]+=x0[kj]*data_M.A[k];
       }

   }

   void compute_Mz(vector<D> & x0){
     for(L j=0;j<m_2;j++)
       for(L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++){
         L kj=data_M.row_idx[k];
         Mz[j]+=x0[kj]*data_M.A[k];
       }
   }

     void compute_Mx(vector<D> & x0){
     for(L j=0;j<m_2;j++)
       for(L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++){
         L kj=data_M.row_idx[k];
         Mx_s[j]+=x0[kj]*data_M.A[k];
       }
   }

   void compute_Ax(vector<D> & x0){
     for(L j=0;j<m_1;j++)
       for(L k = data_A.ptr[j]; k < data_A.ptr[j + 1];k++){
         L kj=data_A.row_idx[k];
         Ax_s[j]+=x0[kj]*data_A.A[k];
       }
   }

   void compute_cs(){
       c_s=0;
       for(L j=0;j<m_2;j++)
           c_s+=data_M.b[j]*lambda_s[j];
   }

    void compute_Mty(){
        M_tlambda_s.clear();
        M_tlambda_s.resize(this->n,0);
        for(L i=0;i<this->n;i++)
        {
          for(L k = data_M.ptr_t[i]; k < data_M.ptr_t[i + 1];k++){
            L j=data_M.col_idx[k];
            M_tlambda_s[i]+=lambda_s[j]*data_M.A_t[k];
           }
        }

    }

    void update_y(){
      D res=0;
      D dl=0;
        for(L j=0;j<m_2;j++)
           {
             dl=(Mx_s[j]-data_M.b[j])/beta_s;
             lambda_s[j]+=dl;
             res+=(Mx_s[j]-data_M.b[j])*(Mx_s[j]-data_M.b[j]);
             L i;
             for (L k = data_M.ptr[j]; k < data_M.ptr[j + 1];k++)
             {
               i=data_M.row_idx[k];
               M_tlambda_s[i]+=dl*data_M.A[k];
             }
            }
       //samp_ALM2<<"lambda_{k+1}-lambda_k"<<sqrt(res)<<endl;
    }

   void Initialize(D beta_0, D epsilon_0,  D val_eta, D val_rho,L val_tau, vector<D> & x0,vector<D> & y0, D val_lambda_f){
        cout<<"start initializing"<<endl;
        set_matrix_M();
        set_matrix_A();

        this->tau=val_tau;
        m_1=data_A.get_n();
        m_2=data_M.get_n();
        cout<<"m_1="<<m_1<<endl;
        cout<<"m_2="<<m_2<<endl;
        lambda_f.resize(m_1,val_lambda_f);
        dual_alpha1.resize(m_1,0);
        dual_alpha2.resize(m_2,0);
        this->n=data_A.nfeatures;
        baralpha1.resize(this->n,0);
        baralpha2.resize(this->n,0);
        gradient_of_f.resize(this->n,0);

       rescale_Matrix();

       beta_s=beta_0;
       //if(max_Lf_s>0)
       //  beta_s=min(beta_0,max_M_s/max_Lf_s);

       epsilon_s=epsilon_0;
       compute_m0(beta_s,val_eta,val_rho,val_tau);

       eta=val_eta;
       rho=val_rho;

       x_s.resize(this->n,0);
       lambda_s.resize(m_2,0);
       for(L i=0;i<this->n;i++)
           x_s[i]=x0[i];

       for(L j=0;j<m_2;j++)
           lambda_s[j]=y0[j];

       M_tlambda_s.resize(this->n,0);

       compute_cs();
       compute_Mty();

       Au.clear();
       Au.resize(m_1,0);
       Az.clear();
       Az.resize(m_1,0);
       Ax_s.clear();
       Ax_s.resize(m_1,0);
       Mu.clear();
       Mu.resize(m_2,0);
       Mz.clear();
       Mz.resize(m_2,0);
       Mx_s.clear();
       Mx_s.resize(m_2,0);

       compute_Az(x0);
       compute_Mz(x0);
       compute_Ax(x0);
       compute_Mx(x0);
   }

   void reset_everything(){
       epsilon_s*=rho*rho*eta;
       m_s/=eta;
       beta_s*=eta;

       compute_cs();
       compute_Mty();

       Au.clear();
       Au.resize(m_1,0);
       Az.clear();
       Az.resize(m_1,0);
       Mu.clear();
       Mu.resize(m_2,0);
       Mz.clear();
       Mz.resize(m_2,0);
       compute_Az(x_s);
       compute_Mz(x_s);
   }

    void APPROX_restart(L K, vector<D> & x0, L val_tau, L eval, L p_N,  L max_nb, D eps, string filename){
        std::vector<D> tmp_x(this->n);
        for(L i=0;i<this->n;i++){
          tmp_x[i]=x0[i];
       }
       L k=0;
       while(k<ceil(max_nb/(K+0.0))){
          k++;
          Au.clear();
          Au.resize(m_1,0);
          Az.clear();
          Az.resize(m_1,0);
          Mu.clear();
          Mu.resize(m_2,0);
          Mz.clear();
          Mz.resize(m_2,0);
          compute_Az(tmp_x);
          compute_Mz(tmp_x);
          cout<<"p_N="<<tmp_x[0]<<" "<<tmp_x[1]<<" "<<tmp_x[2]<<" "<<p_N<<endl;
          this->APPROX_MU(tmp_x, val_tau, 0, 0, eval, p_N, K, eps, filename,1);

           for(L i=0;i<this->n;i++){
             tmp_x[i]=this->x[i];
          }
          cout<<"restarting..."<<k<<"; "<<(max_nb/(K+0.0))<<endl;
          if(this->delta<eps) k=ceil(max_nb/(K+0.0))+1;
        }

  }

   inline void compute_and_record_res(){
        if(nb_outer_iters%print_every_N_ALM_I_APG==0){
          compute_KKT_residual();
          compute_function_value();
          cout<<setprecision(9)<<"Iteration: "<<nb_outer_iters<<"; time="<<running_time_ALM_I_APG<<"; residual1="<<residual1<<" residual2="<<residual2<<"; function value="<<function_value<<endl;
          samp_ALM_I_APG<<setprecision(9)<<nb_outer_iters<<" "<<running_time_ALM_I_APG<<" "<<residual1<<" "<<residual2<<" "<<function_value<<" "<<endl;
        }
   }

   void ALM_I_APG_solve_with_APPROX(D beta_0, D epsilon_0,  D eta, D rho,vector<D> & x0,vector<D> & y0, L val_tau, L max_nb_outer, L p_N_1, L p_N_2, D val_lambda_f,string filename1, string filename2, D time){
      Initialize(beta_0, epsilon_0, eta, rho,val_tau, x0, y0, val_lambda_f);

      nb_outer_iters=0;
      string sampname2="results/APPROXMU_"+filename2;
      this->samp.open(sampname2.c_str());
      filename1="results/ALM2_"+filename1;
      samp_ALM_I_APG.open(filename1.c_str());
      running_time_ALM_I_APG=0;
      print_every_N_ALM_I_APG=p_N_1;
      compute_and_record_res();
      D start;
      start = std::clock();
      std::vector<D> record_x_s(this->n);
      std::vector<D> record_x_sp1(this->n);
      cout<<"m0="<<ceil(m_s/this->n*val_tau)<<endl;
      //D K=ceil(2*sqrt((max_Lf_s+max_M_s)/(beta_s+val_mu_g)+1));
      //cout<<"restart period="<<K<<endl;
      cout<<"ms="<<ceil(m_s/this->n*val_tau)<<"; beta_s="<<beta_s<<"; epsilon_s="<<epsilon_s<<endl;
      this->APPROX_MU(x_s, val_tau, val_mu_f, val_mu_g, 3, p_N_2, ceil(m_s/this->n*val_tau), epsilon_s, filename2,1);
      //APPROX_restart(K, x_s, val_tau, 3, p_N_2, ceil((m_s+mp_s)/this->n*val_tau), epsilon_s, filename2);
      compute_x();
      update_y();
      nb_outer_iters++;
      running_time_ALM_I_APG+=( std::clock() - start ) / (double) CLOCKS_PER_SEC;
      compute_and_record_res();
      start = std::clock();
      reset_everything();
      running_time_ALM_I_APG+=( std::clock() - start ) / (double) CLOCKS_PER_SEC;
      while(nb_outer_iters<max_nb_outer){
         start = std::clock();
         cout<<"ms="<<ceil(m_s/this->n*val_tau)<<"; beta_s="<<beta_s<<"; epsilon_s"<<epsilon_s<<endl;
         //D K=ceil(2*sqrt((max_Lf_s+max_M_s)/(beta_s+val_mu_g)+1));
         //cout<<"restart period="<<K<<endl;
        this->APPROX_MU(x_s, val_tau, val_mu_f, val_mu_g, 3, p_N_2, ceil(m_s/this->n*val_tau), epsilon_s, filename2,1);
        for(L i=0;i<this->n;i++){
          record_x_s[i]=x_s[i];
        }
         //APPROX_restart(K, x_s, val_tau, 3, p_N_2, ceil((m_s+mp_s)/this->n*val_tau), epsilon_s, filename2);
         compute_x();
         update_y();
         nb_outer_iters++;
         running_time_ALM_I_APG+=( std::clock() - start ) / (double) CLOCKS_PER_SEC;
         testing();
         compute_and_record_res();
         start = std::clock();
         reset_everything();
         running_time_ALM_I_APG+=( std::clock() - start ) / (double) CLOCKS_PER_SEC;
         if (running_time_ALM_I_APG> time){
         	break;
		 }
      }

   }

};

#endif


#ifndef GBNET_GRAPHORNOR
#define GBNET_GRAPHORNOR

#include <iostream>
#include <string>
#include <vector>
#include <utility>
#include <map>
#include <set>
#include <stdio.h>

#include <gsl/gsl_rng.h>

#include "RVNode.h"
#include "XNode.h"
#include "HNodeORNOR.h"
#include "YDataNode.h"
#include "SNode.h"
#include "TNode.h"
#include "ZNode.h"
#include "NodeDictionary.h"
#include "GraphBase.h"
#include "YNoiseNode.h"

namespace gbn
{
    const double SPRIOR[3 * 3] = {0.40, 0.40, 0.20,
                                  0.30, 0.40, 0.30,
                                  0.20, 0.40, 0.40};

    class GraphORNOR: public GraphBase
    {
        private:
        protected:
        public:
            const double XPROB[2] = {0.99, 0.01};
            const double XPROB_ACTIVE[2] = {0.10, 0.90};
            double SPROB[3][3] = {{0.900, 0.090, 0.010},
                                        {0.050, 0.900, 0.050},
                                        {0.010, 0.090, 0.900}};
            double YPROB[3] = {0.05, 0.90, 0.05};
            double YNOISE[3][3] = {{0.945, 0.050, 0.005},
                                   {0.050, 0.900, 0.050},
                                   {0.005, 0.050, 0.945}};
            const double YNOISE_ALPHA[3][3] = {{2000,  100,   10},
                                               { 100, 2000,  100},
                                               {  10,  100, 2000}};

            GraphORNOR (unsigned int);
            virtual ~GraphORNOR ();

            void build_structure (network_t, evidence_dict_t, prior_active_tf_set_t, bool = true,
                                  const double [3 * 3] = SPRIOR, double z_alpha = 25., double z_beta = 25., double z0_alpha = 25., double z0_beta = 25., double t_alpha = 25., double t_beta = 25., bool comp_yprob = true, bool const_params = false,
                                  double t_focus = 2., double t_lmargin = 2., double t_hmargin = 2., double zn_focus = 2., double zn_lmargin = 8., double zn_hmargin = 2.);
    };
}

#endif