/**
 * \author Sylvain Marsat, University of Maryland - NASA GSFC
 *
 * \brief C code for structures EOBNRv2HM reduced order model (non-spinning version).
 * See CQG 31 195010, 2014, arXiv:1402.4146 for details on the reduced order method.
 * See arXiv:1106.1021 for the EOBNRv2HM model.
 *
 * Borrows from the SEOBNR ROM LAL code written by Michael Puerrer and John Veitch.
 *
 * Put the untared data in the directory designated by the environment variable ROM_DATA_PATH.
 *
 * Parameter range:
 *   q = 1-12 (almost)
 *   No spin
 *   Mtot >= 10Msun for fstart=8Hz
 *
 */

#define _XOPEN_SOURCE 500

#ifdef __GNUC__
#define UNUSED __attribute__ ((unused))
#else
#define UNUSED
#endif

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <complex.h>
#include <time.h>
#include <unistd.h>
#include <getopt.h>
#include <stdbool.h>
#include <string.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_bspline.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_min.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_complex.h>

#include "constants.h"
#include "EOBNRv2HMROMstruct.h"
#include "EOBNRv2HMROM.h"

/************************************************************************/
/********************* Functions for list structures ********************/

/***************** Functions for the SplineList structure ****************/

/* Prepend a node to a linked list of splines, or create a new head */
SplineList* SplineList_AddElementNoCopy(
	   SplineList* appended,  /* List structure to prepend to */
	   gsl_spline* spline,  /* spline to contain */
           gsl_interp_accel* accel,  /* accelerator to contain */
	   int i /* index in the list */)
{
    SplineList* splinelist;
    /* Check if the node with this index already exists */
    splinelist = appended;
    while( splinelist ){
      if( i == splinelist->i ){
	break;
      }
      splinelist = splinelist->next;
    }
    if( splinelist ){ /* We don't allow for the case where the index already exists*/
      printf("Error: Tried to add an already existing index to a SplineList");
      return(NULL);
    } else { /* In that case, we do NOT COPY the input spline, which therefore can't be
                 used anywhere else; this will be acceptable as these operations will only be done
                 when initializing the data */
      splinelist = malloc( sizeof(SplineList) );
    }
    splinelist->i = i;
    if( spline ){
      splinelist->spline = spline;
    } else {
      splinelist->spline = NULL;
    }
    if( accel ){
      splinelist->accel = accel;
    } else {
      splinelist->accel = NULL;
    }
    if( appended ){
      splinelist->next = appended;
    } else {
        splinelist->next = NULL;
    }
    return splinelist;
}
/* Get the element of a SplineList with a given index */
SplineList* SplineList_GetElement(
	      SplineList* const splinelist,  /* List structure to get element from */
              const int i ) /* Index looked for */
{
    if( !splinelist ) return NULL;

    SplineList* itr = splinelist;
    while( itr->i != i ){
        itr = itr->next;
        if( !itr ) return NULL;
    }
    return itr; /* The element returned is itself a pointer to a SplineList */
}
/* Delete list from given pointer to the end of the list */
void SplineList_Destroy( SplineList* splinelist ) /* Head of linked list to destroy */
{
  SplineList* pop;
  while( (pop = splinelist) ){
    if( pop->spline ){ /* Internal spline and accelerator are freed */
      gsl_spline_free( pop->spline );
    }
    if( pop->accel ){
      gsl_interp_accel_free( pop->accel );
    }
    /* Notice that the index i is not freed, like in SphHarmTimeSeries struct indices l and m */
    splinelist = pop->next;
    free( pop );
  }
}

/***************** Functions for the EOBNRHMROMdata structure ****************/
ListmodesEOBNRHMROMdata* ListmodesEOBNRHMROMdata_AddModeNoCopy(
	   ListmodesEOBNRHMROMdata* appended,  /* List structure to prepend to */
	   EOBNRHMROMdata* data,  /* data to contain */
	   const int l, /* major mode number */
	   const int m  /* minor mode number */)
{
    ListmodesEOBNRHMROMdata* list;
    /* Check if the node with this mode already exists */
    list = appended;
    while( list ){
      if( l == list->l && m == list->m ){
	break;
      }
      list = list->next;
    }
    if( list ){ /* We don't allow for the case where the mode already exists in the list*/
      printf("Error: Tried to add an already existing mode to a ListmodesEOBNRHMROMdata ");
      return(NULL);
    } else { /* In that case, we do NOT COPY the input interpolated data, which therefore can't be
		used anywhere else; this will be acceptable as these operations will only be done
		when interpolating the initialization data */
      list = malloc( sizeof(ListmodesEOBNRHMROMdata) );
    }
    list->l = l;
    list->m = m;
    if( data ){
      list->data = data;
    } else {
      list->data = NULL;
    }
    if( appended ){
      list->next = appended;
    } else {
        list->next = NULL;
    }
    return list;
}
/* Get the element of a ListmodesEOBNRHMROMdata with a given index */
ListmodesEOBNRHMROMdata* ListmodesEOBNRHMROMdata_GetMode(
	   ListmodesEOBNRHMROMdata* const list,  /* List structure to get a particular mode from */
	   int l, /*< major mode number */
	   int m  /*< minor mode number */ )
{
    if( !list ) return NULL;

    ListmodesEOBNRHMROMdata *itr = list;
    while( itr->l != l || itr->m != m ){
        itr = itr->next;
        if( !itr ) return NULL;
    }
    return itr; /* The element returned is itself a pointer to a ListmodesEOBNRHMROMdata */
}
void ListmodesEOBNRHMROMdata_Destroy(
	   ListmodesEOBNRHMROMdata* list  /* List structure to destroy; notice that the data is destroyed too */
)
{
  ListmodesEOBNRHMROMdata* pop;
  while( (pop = list) ){
    if( pop->data ){ /* Destroying the EOBNRHMROMdata data */
      EOBNRHMROMdata_Cleanup( pop->data );
    }
    /* Notice that the mode indices l and m are not freed, like in SphHarmTimeSeries struct indices l and m */
    list = pop->next;
    free( pop );
  }
}

/***************** Functions for the EOBNRHMROMdata_interp structure ****************/
ListmodesEOBNRHMROMdata_interp* ListmodesEOBNRHMROMdata_interp_AddModeNoCopy(
	   ListmodesEOBNRHMROMdata_interp* appended,  /* List structure to prepend to */
	   EOBNRHMROMdata_interp* data_interp,  /* data to contain */
	   int l, /* major mode number */
	   int m  /* minor mode number */)
{
    ListmodesEOBNRHMROMdata_interp* list;
    /* Check if the node with this mode already exists */
    list = appended;
    while( list ){
      if( l == list->l && m == list->m ){
	break;
      }
      list = list->next;
    }
    if( list ){ /* We don't allow for the case where the mode already exists in the list*/
      printf("Error: Tried to add an already existing mode to a ListmodesEOBNRHMROMdata_interp ");
      return(NULL);
    } else { /* In that case, we do NOT COPY the input interpolated data, which therefore can't be
		used anywhere else; this will be acceptable as these operations will only be done
		when interpolating the initialization data */
      list = malloc( sizeof(ListmodesEOBNRHMROMdata_interp) );
    }
    list->l = l;
    list->m = m;
    if( data_interp ){
      list->data_interp = data_interp;
    } else {
      list->data_interp = NULL;
    }
    if( appended ){
      list->next = appended;
    } else {
        list->next = NULL;
    }
    return list;
}
/* Get the element of a ListmodesEOBNRHMROMdata with a given index */
ListmodesEOBNRHMROMdata_interp* ListmodesEOBNRHMROMdata_interp_GetMode(
	   ListmodesEOBNRHMROMdata_interp* const list,  /* List structure to get a particular mode from */
	   int l, /*< major mode number */
	   int m  /*< minor mode number */ )
{
    if( !list ) return NULL;

    ListmodesEOBNRHMROMdata_interp *itr = list;
    while( itr->l != l || itr->m != m ){
        itr = itr->next;
        if( !itr ) return NULL;
    }
    return itr; /* The element returned is itself a pointer to a ListmodesEOBNRHMROMdata_interp */
}
void ListmodesEOBNRHMROMdata_interp_Destroy(
	   ListmodesEOBNRHMROMdata_interp* list  /* List structure to destroy; notice that the data is destroyed too */
)
{
  ListmodesEOBNRHMROMdata_interp* pop;
  while( (pop = list) ){
    if( pop->data_interp ){ /* Destroying the EOBNRHMROMdata_interp data */
      EOBNRHMROMdata_interp_Cleanup( pop->data_interp );
    }
    /* Notice that the mode indices l and m are not freed, like in SphHarmTimeSeries struct indices l and m */
    list = pop->next;
    free( pop );
  }
}

/***************** Functions for the EOBNRHMROMdata_coeff structure ****************/
ListmodesEOBNRHMROMdata_coeff* ListmodesEOBNRHMROMdata_coeff_AddModeNoCopy(
	   ListmodesEOBNRHMROMdata_coeff* appended,  /* List structure to prepend to */
	   EOBNRHMROMdata_coeff* data_coeff,  /* data to contain */
	   int l, /* major mode number */
	   int m  /* minor mode number */)
{
    ListmodesEOBNRHMROMdata_coeff* list;
    /* Check if the node with this mode already exists */
    list = appended;
    while( list ){
      if( l == list->l && m == list->m ){
	break;
      }
      list = list->next;
    }
    if( list ){ /* We don't allow for the case where the mode already exists in the list*/
      printf("Error: Tried to add an already existing mode to a ListmodesEOBNRHMROMdata_coeff ");
      return(NULL);
    } else { /* In that case, we do NOT COPY the input interpolated data, which therefore can't be
		used anywhere else; this will be acceptable as these operations will only be done
		when interpolating the initialization data */
      list = malloc( sizeof(ListmodesEOBNRHMROMdata_coeff) );
    }
    list->l = l;
    list->m = m;
    if( data_coeff ){
      list->data_coeff = data_coeff;
    } else {
      list->data_coeff = NULL;
    }
    if( appended ){
      list->next = appended;
    } else {
        list->next = NULL;
    }
    return list;
}
/* Get the element of a ListmodesEOBNRHMROMdata_coeff with a given index */
ListmodesEOBNRHMROMdata_coeff* ListmodesEOBNRHMROMdata_coeff_GetMode(
	   ListmodesEOBNRHMROMdata_coeff* const list,  /* List structure to get a particular mode from */
	   int l, /*< major mode number */
	   int m  /*< minor mode number */ )
{
    if( !list ) return NULL;

    ListmodesEOBNRHMROMdata_coeff *itr = list;
    while( itr->l != l || itr->m != m ){
        itr = itr->next;
        if( !itr ) return NULL;
    }
    return itr; /* The element returned is itself a pointer to a ListmodesEOBNRHMROMdata_coeff */
}
void ListmodesEOBNRHMROMdata_coeff_Destroy(
	   ListmodesEOBNRHMROMdata_coeff* list  /* List structure to destroy; notice that the data is destroyed too */
)
{
  ListmodesEOBNRHMROMdata_coeff* pop;
  while( (pop = list) ){
    if( pop->data_coeff ){ /* Destroying the EOBNRHMROMdata_coeff data */
      EOBNRHMROMdata_coeff_Cleanup( pop->data_coeff );
    }
    /* Notice that the mode indices l and m are not freed, like in SphHarmTimeSeries struct indices l and m */
    list = pop->next;
    free( pop );
  }
}


//
// Created by pedram pakseresht on 2/18/21.
//

static char help[] = "Simple Petsc program";

#include <petsc.h>
#include "incompressibleFlow.h"
#include "mesh.h"
#include "particleInertial.h"
#include "particles.h"
#include "particleInitializer.h"
#include "petscviewer.h"

static PetscErrorCode uniform_u(PetscInt Dim, PetscReal time, const PetscReal *X, PetscInt Nf, PetscScalar *u, void *ctx) {
    //PetscInt d;
    //for (d = 0; d < Dim; ++d)
    //    u[d] = 1.0;
    //return 0;
    if (X[1]==1.0){
        u[0] = 1.0;
    }
    else {
        u[0] = 0.0;
    }
        u[1] =0.0;
    // u[0] = X[1] - X[1]*X[1] ;
    // u[1] = 0.0;
    return 0;
}

static PetscErrorCode uniform_u_t(PetscInt Dim, PetscReal time, const PetscReal *X, PetscInt Nf, PetscScalar *u, void *ctx) {
    u[0] = 0.0;
    u[1] = 0.0;
    return 0;
}

static PetscErrorCode uniform_p(PetscInt Dim, PetscReal time, const PetscReal *X, PetscInt Nf, PetscScalar *p, void *ctx) {
    p[0] = 0.0;
    return 0;
}

static PetscErrorCode uniform_T(PetscInt Dim, PetscReal time, const PetscReal *X, PetscInt Nf, PetscScalar *T, void *ctx) {
    T[0] = 0.0;
    return 0;
}
static PetscErrorCode uniform_T_t(PetscInt Dim, PetscReal time, const PetscReal *X, PetscInt Nf, PetscScalar *T, void *ctx) {
    T[0] = 0.0;
    return 0;
}

static PetscErrorCode SetInitialConditions(TS ts, Vec u) {

    PetscErrorCode (*initFuncs[3])(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar u[], void *ctx) = {uniform_u,uniform_p,uniform_T};
    // u here is the solution vector including velocity, temperature and pressure fields.
    DM dm;
    PetscReal t;
    PetscErrorCode ierr;
    PetscFunctionBegin;

    ierr = TSGetDM(ts, &dm);
    CHKERRQ(ierr);
    ierr = TSGetTime(ts, &t);

    CHKERRQ(ierr);

    ierr = DMProjectFunction(dm, 0.0, initFuncs, NULL, INSERT_ALL_VALUES, u);
    CHKERRQ(ierr);
    // get the flow to apply the completeFlowInitialization method
    ierr = IncompressibleFlow_CompleteFlowInitialization(dm, u);
    CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

static PetscErrorCode MonitorFlowAndParticleError(TS ts, PetscInt step, PetscReal crtime, Vec u, void *ctx) {
    // PetscErrorCode (*exactFuncs[3])(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *ctx);
    // void *ctxs[3];
    DM dm;
    PetscDS ds;
    Vec v;
    PetscReal ferrors[3];
    PetscInt f;
    PetscErrorCode ierr;
    PetscInt num;

    PetscFunctionBeginUser;
    ierr = TSGetDM(ts, &dm);
    CHKERRQ(ierr);
    ierr = DMGetDS(dm, &ds);
    CHKERRQ(ierr);

    // get the particle data from the context
    ParticleData particlesData = (ParticleData)ctx;
    PetscInt particleCount;
    ierr = DMSwarmGetSize(particlesData->dm, &particleCount);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // compute the average particle location
    const PetscReal *coords;
    PetscInt dims;
    PetscReal avg[3] = {0.0, 0.0, 0.0};
    ierr = DMSwarmGetField(particlesData->dm, DMSwarmPICField_coor, &dims, NULL, (void **)&coords);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    for (PetscInt p = 0; p < particleCount; p++) {
        for (PetscInt n = 0; n < dims; n++) {
            avg[n] += coords[p * dims + n] / particleCount;  // PetscReal
        }
    }
    ierr = DMSwarmRestoreField(particlesData->dm, DMSwarmPICField_coor, &dims, NULL, (void **)&coords);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ierr = PetscPrintf(PETSC_COMM_WORLD,
                       "Timestep: %04d time = %-8.4g \t L_2 Error: [%2.3g, %2.3g, %2.3g] ParticleCount: %d\n",
                       (int)step,
                       (double)crtime,
                       (double)ferrors[0],
                       (double)ferrors[1],
                       (double)ferrors[2],
                       particleCount,
                       (double)avg[0],
                       (double)avg[1]);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "Avg Particle Location: [%2.3g, %2.3g, %2.3g]\n", (double)avg[0], (double)avg[1], (double)avg[2]);

    CHKERRABORT(PETSC_COMM_WORLD, ierr);

     ierr = VecViewFromOptions(u, NULL, "-vec_view");
     CHKERRABORT(PETSC_COMM_WORLD, ierr);

    //  ierr = VecView(u, NULL);CHKERRABORT(PETSC_COMM_WORLD, ierr);
    // ierr = PetscPrintf(PETSC_COMM_WORLD, "current time=%g\n", crtime);CHKERRQ(ierr);

    ierr = DMSetOutputSequenceNumber(particlesData->dm, step, crtime);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

     /*
     if (step == 0 ){
        ierr = VecViewFromOptions(u, NULL, "-vec_init");
        CHKERRABORT(PETSC_COMM_WORLD, ierr);
     } else {
        ierr = VecViewFromOptions(u, NULL, "-vec_view");
         CHKERRABORT(PETSC_COMM_WORLD, ierr);
     }
     */

    if (step == 0) {
        ierr = ParticleViewFromOptions(particlesData, NULL, "-particle_init");
        CHKERRABORT(PETSC_COMM_WORLD, ierr);
    } else {
        ierr = ParticleViewFromOptions(particlesData, NULL, "-particle_view");
        CHKERRABORT(PETSC_COMM_WORLD, ierr);
    }

    PetscFunctionReturn(0);
}

static PetscErrorCode ParticleSetInitialVelocity(DM particleDm) {

    PetscFunctionBeginUser;

    PetscErrorCode  ierr;

    /*
    PetscScalar *vel;
    PetscInt dim, Npb, p, n;

    ierr = PetscOptionsBegin(PETSC_COMM_WORLD, "", "Box Particle Initialization Options", NULL);CHKERRQ(ierr);
    ierr = PetscOptionsInt("-Npb", "The initial number of particles per box dimension", NULL, Npb, &Npb, NULL);CHKERRQ(ierr);
    ierr = PetscOptionsEnd();

    ierr = PetscPrintf(PETSC_COMM_WORLD, "***Npb=%D\n***", Npb);CHKERRQ(ierr);
    ierr = DMSwarmGetField(particleDm, ParticleVelocity, NULL, NULL, (void **) &vel);CHKERRQ(ierr);

    ierr = DMGetDimension(particleDm, &dim);CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "***dim=%D\n***", dim);CHKERRQ(ierr);

    Npb *= dim;
    for (p = 0; p < Npb; ++p) {
        for (n = 0; n < dim; n++) {
            vel[p+n] = 0.0;
        }
    }
    ierr = DMSwarmRestoreField(particleDm, ParticleVelocity, NULL, NULL, (void **) &vel);CHKERRQ(ierr);
    */

    Vec vel;
    ierr = DMSwarmCreateGlobalVectorFromField(particleDm,ParticleVelocity, &vel);CHKERRQ(ierr);
    ierr = VecSet(vel, 0.0);CHKERRQ(ierr);
    //ierr = VecView(vel, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
    ierr = DMSwarmDestroyGlobalVectorFromField(particleDm,ParticleVelocity, &vel);CHKERRQ(ierr);

    PetscFunctionReturn(0);
}

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

    DM dm;     // domain definition
    TS ts;     // time-stepper
    PetscErrorCode ierr;

    PetscBag parameterBag; // constant flow parameters
    Vec flowField;         // flow solution vector
    FlowData flowData;

    PetscReal t;
    PetscInt Dim=2;
    PetscReal dt=0.1;       // dt for time stepper
    PetscInt max_steps=2; // maximum time steps

    PetscReal Re=1.0;       // Reynolds number
    PetscReal St=1.0;       // Strouhal number
    PetscReal Pe=1.0;       // Peclet number
    PetscReal Mu=1.0;      // viscosity
    PetscReal K_input=1.0;  // thermal conductivity
    PetscReal Cp = 1.0;    // heat capacity
    // PetscInt Np=10;

    // initialize Petsc ...
    ierr = PetscInitialize(&argc, &argv, NULL, "");CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "*** Set up a simple flow *** \n");CHKERRQ(ierr);

    // setup the ts
    ierr = TSCreate(PETSC_COMM_WORLD, &ts);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = CreateMesh(PETSC_COMM_WORLD, &dm, PETSC_TRUE, Dim);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = TSSetDM(ts, dm);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    //output the mesh
    ierr = DMViewFromOptions(dm, NULL, "-dm_view");
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Setup the flow data
    ierr = FlowCreate(&flowData);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // setup problem
    ierr = IncompressibleFlow_SetupDiscretization(flowData, dm);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    IncompressibleFlowParameters flowParameters;

    // changing non-dimensional parameters manually here ...
    flowParameters.strouhal = St;
    flowParameters.reynolds = Re;
    flowParameters.peclet = Pe;
    flowParameters.mu = Mu;
    flowParameters.k = K_input;
    flowParameters.cp = Cp;

    // print out the parameters
    ierr = PetscPrintf(PETSC_COMM_WORLD, "*** non-dimensional parameters ***\n");CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "St=%g\n", flowParameters.strouhal);CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "Re=%g\n", flowParameters.reynolds);CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "Pe=%g\n", flowParameters.peclet);CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "Mu=%g\n", flowParameters.mu);CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "K=%g\n", flowParameters.k);CHKERRQ(ierr);
    ierr = PetscPrintf(PETSC_COMM_WORLD, "Cp=%g\n", flowParameters.cp);CHKERRQ(ierr);

    // Start the problem setup
    PetscScalar constants[TOTAL_INCOMPRESSIBLE_FLOW_PARAMETERS];
    ierr = IncompressibleFlow_PackParameters(&flowParameters, constants);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = IncompressibleFlow_StartProblemSetup(flowData, TOTAL_INCOMPRESSIBLE_FLOW_PARAMETERS, constants);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Override problem with source terms, boundary, and set the exact solution
    PetscDS prob;
    ierr = DMGetDS(dm, &prob);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Setup Boundary Conditions
    // Note: DM_BC_ESSENTIAL is a Dirichlet BC.
    PetscInt id;
    id = 3;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "top wall velocity", "marker", VEL, 0, NULL, (void (*)(void))uniform_u, (void (*)(void))uniform_u_t, 1, &id, NULL);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    id = 1;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "bottom wall velocity", "marker", VEL, 0, NULL, (void (*)(void))uniform_u, (void (*)(void))uniform_u_t, 1, &id, NULL);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    id = 2;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "right wall velocity", "marker", VEL, 0, NULL, (void (*)(void))uniform_u, (void (*)(void))uniform_u_t, 1, &id, NULL);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    id = 4;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "left wall velocity", "marker", VEL, 0, NULL, (void (*)(void))uniform_u, (void (*)(void))uniform_u_t, 1, &id, NULL);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    id = 3;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "top wall temp", "marker", TEMP, 0, NULL, (void (*)(void))uniform_T, (void (*)(void))uniform_T_t, 1, &id, NULL);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    id = 1;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "bottom wall temp", "marker", TEMP, 0, NULL, (void (*)(void))uniform_T, (void (*)(void))uniform_T_t, 1, &id, parameterBag);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    id = 2;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "right wall temp", "marker", TEMP, 0, NULL, (void (*)(void))uniform_T, (void (*)(void))uniform_T_t, 1, &id, parameterBag);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    id = 4;
    ierr = PetscDSAddBoundary(
        prob, DM_BC_ESSENTIAL, "left wall temp", "marker", TEMP, 0, NULL, (void (*)(void))uniform_T, (void (*)(void))uniform_T_t, 1, &id, parameterBag);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ierr = IncompressibleFlow_CompleteProblemSetup(flowData, ts);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Name the flow field
    ierr = PetscObjectSetName(((PetscObject)flowData->flowField), "Numerical Solution");
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ierr = SetInitialConditions(ts, flowData->flowField);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ierr = TSGetTime(ts, &t);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    //  *** added by Pedram for dt and number of time steps ***
    ierr = TSSetTimeStep(ts,dt);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ierr = TSSetMaxSteps(ts,max_steps);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    // ***********************************

    ierr = DMSetOutputSequenceNumber(dm, 0, t);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // checks for convergence ...
    ierr = DMTSCheckFromOptions(ts, flowData->flowField);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ParticleData particles;

    // Setup the particle domain
    ierr = ParticleInertialCreate(&particles, Dim);
    //ierr = ParticleInertialCreate(&particles, Dim);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // link the flow to the particles
    ierr = ParticleInitializeFlow(particles, flowData);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // name the particle domain
    ierr = PetscObjectSetOptionsPrefix((PetscObject)(particles->dm), "particles_");
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = PetscObjectSetName((PetscObject)particles->dm, "Particles");
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // initialize the particles position and velocity
    ierr = ParticleInitialize(dm, particles->dm);CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = ParticleSetInitialVelocity(particles->dm);CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // setup the flow monitor to also check particles
    ierr = TSMonitorSet(ts, MonitorFlowAndParticleError, particles, NULL);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = TSSetFromOptions(ts);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Setup particle position integrator
    TS particleTs;
    ierr = TSCreate(PETSC_COMM_WORLD, &particleTs);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = PetscObjectSetOptionsPrefix((PetscObject)particleTs, "particle_");
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    ierr = ParticleInertialSetupIntegrator(particles, particleTs, flowData);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Solve the one way coupled system
    ierr = TSSolve(ts, flowData->flowField);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);

    // Cleanup
    ierr = DMDestroy(&dm);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = TSDestroy(&ts);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = TSDestroy(&particleTs);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = FlowDestroy(&flowData);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = ParticleInertialDestroy(&particles);
    CHKERRABORT(PETSC_COMM_WORLD, ierr);
    ierr = PetscFinalize();
    exit(ierr);

    // _lower 0.25,0.25 _upper 0.75,0.75
    // Npb 10
    // _ts_dt 0.1
    // _layout_type box
    // _ts_convergence_estimate
    // _num_refine 1
    // -particle_ts_monitor_cancel
    // -ts_monitor_cancel

    // -particle_layout_type box
    // -particle_lower 0.1,0.1
    // -particle_upper 0.3,0.3
    // -Npb 5
    // -particle_ts_max_steps 10
    // -particle_ts_dt 0.05

    /*
    -ksp_type
    fgmres
    -ksp_gmres_restart
    10
    -ksp_rtol
    1.0e-9
    -ksp_error_if_not_converged
    -pc_type
    fieldsplit
    -pc_fieldsplit_0_fields
    0,2
      -pc_fieldsplit_1_fields
    1
    -pc_fieldsplit_type
    schur
    -pc_fieldsplit_schur_factorization_type
    full
    -fieldsplit_0_pc_type
    lu
    -fieldsplit_pressure_ksp_rtol
    1e-10
    -fieldsplit_pressure_pc_type
    jacobi
     */

    /*
     * -Npb 1 -particle_lower 0.05,0.1 -particle_upper 0.3,0.9 -particle_layout_type box -vec_init hdf5:/Users/pedram/scratch/sol.h5 -vec_view hdf5:/Users/pedram/scratch/sol.h5::append -dm_view hdf5:/Users/pedram/scratch/sol.h5 -dm_refine 2 -vel_petscspace_degree 2 -pres_petscspace_degree 1 -temp_petscspace_degree 1 -ksp_type fgmres -ksp_rtol 1.0e-9 -ksp_atol 1.0e-12 -ksp_error_if_not_converged -pc_type fieldsplit -pc_fieldsplit_0_fields 0,2 -pc_fieldsplit_1_fields 1 -pc_fieldsplit_type schur -pc_fieldsplit_schur_factorization_type full -fieldsplit_0_pc_type lu -fieldsplit_pressure_ksp_rtol 1e-10 -fieldsplit_pressure_pc_type jacobi -dm_plex_separate_marker -snes_monitor -ksp_monitor -snes_converged_reason -ksp_converged_reason -particle_view hdf5:/Users/pedram/scratch/solP.h5::append -particle_init hdf5:/Users/pedram/scratch/solP.h5
     */

}


#ifndef PROXIMAL_MATRIX_H_GUARD
#define PROXIMAL_MATRIX_H_GUARD

#include <math.h>
#include <float.h>

#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_eigen.h>

void project_sdc(gsl_matrix *X);
void project_spectral_norm_ball(gsl_matrix *X);
void prox_neg_log_det(gsl_matrix *X, const double rho);
void prox_quad(gsl_vector *x, const double rho, gsl_matrix *A, gsl_matrix *b);

#endif


#include <lapacke.h>
#include <mpi.h>
#include <stdio.h>

#include "../include/paralleltt.h"

#include <unistd.h>

// X is r1 x n x r2, transpose is r2 x n x r1
double* train_transpose(const double* X, int r1, int n, int r2){
    double* X_T = (double*) calloc(r1*n*r2, sizeof(double));

    for (int ii = 0; ii < r1; ++ii){
        for (int jj = 0; jj < n; ++jj){
            for (int kk = 0; kk < r2; ++kk){
                X_T[kk + jj*r2 + ii*n*r2] = X[ii + jj*r1 + kk*n*r1];
            }
        }
    }

    return X_T;
}

void two_sketches_to_train(sketch* s1, sketch* s2, double** train_ptr)
{
    double* train = *train_ptr;
    int head = 0;
    if ((s1 == NULL) || (s2 == NULL)){
        sketch* s = (s1 == NULL) ? s2 : s1;
        MPI_tensor* ten = s->ten;
        MPI_Comm comm = ten->comm;
        int iscol = s->iscol;
        int flattening = s->flattening;

        flattening_info* fi = flattening_info_init(ten, flattening, iscol, 0);

        int d = ten->d;
        int n = ten->n[(iscol) ? 0 : d-1];
        int r = s->r;

        matrix_tt* train_mat = (iscol) ? matrix_tt_wrap(n, r, train) : matrix_tt_wrap(r, n, train);
        matrix_tt_fill_zeros(train_mat);

        int* op = s->owner_partition;
        int rank = ten->rank;

        for (int ii = op[rank]; ii < op[rank+1]; ++ii){
            // X submatrix
            matrix_tt* X_submat = own_submatrix(s, ii, 0);

            // train submatrix
            flattening_info_f_update(fi, ten, ii);
            int i0 = fi->f_t_index[0];
            int i1 = i0 + fi->f_t_sizes[0];

            matrix_tt* train_submat = (iscol) ? submatrix(train_mat, i0, i1, 0, r) : submatrix(train_mat, 0, r, i0, i1);
            train_submat->transpose = (iscol) ? 0 : 1;
            matrix_tt_copy_data(train_submat, X_submat);

            free(train_submat);
            free(X_submat);
        }
        matrix_tt_reduce(comm, rank, train_mat, NULL, head);

        flattening_info_free(fi);
        free(train_mat);
        return;
    }

    int iscol = s1->iscol ? 1 : 0;
    int iscol2 = s2->iscol ? 1 : 0;

    if (iscol != iscol2){
        printf("two_sketches_to_train: Cannot convert a column and a row to a train\n");
        return;
    }

    int flattening1 = s1->flattening;
    int flattening2 = s2->flattening;

    // We want s2 to be the taller sketch, so that we only communicate s1 if necessary (which should be cheaper!)
    if (((iscol) && (flattening1 > flattening2)) || (!(iscol) && (flattening2 > flattening1))){
        sketch* tmp = s1;
        s1 = s2;
        s2 = tmp;

        flattening1 = s1->flattening;
        flattening2 = s2->flattening;
    }

    // We want the flattening dimensions to be one off of each other
    MPI_tensor* ten = s1->ten;
    flattening_info* fi1 = flattening_info_init(ten, flattening1, iscol, 0);
    flattening_info* fi2 = flattening_info_init(ten, flattening2, iscol, 0);
    if (fi2->f_d - fi1->f_d != 1){
        printf("two_sketches_to_train: The difference between the two flattening dimensions is not correct!\n");
        return;
    }

    MPI_Comm comm = ten->comm;
    int size = ten->comm_size;
    int rank = ten->rank;

    int r1 = s1->r;
    int r2 = s2->r;
    int n_index = (iscol) ? flattening1 : flattening2;
    int n = ten->n[n_index];
    matrix_tt* train_mat = matrix_tt_wrap(r1, n*r2, train);
    matrix_tt_fill_zeros(train_mat);

    // Loop over all blocks of fi2
    matrix_tt* s_mat1 = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    for (int block2 = 0; block2 < fi2->f_Nblocks; ++block2){
        flattening_info_f_update(fi2, ten, block2);
        int owner2 = s_get_owner(s2, block2);

        // Get the corresponding block1
        int* f_t_block1 = fi2->f_t_block + ((iscol) ? 0 : 1);
        int* f_nps1 = fi2->f_nps + ((iscol) ? 0 : 1);
        int block1 = 0;
        for (int ii = (fi1->f_d) - 1; ii >= 0; --ii){
            block1 = block1*f_nps1[ii] + f_t_block1[ii];
        }

        flattening_info_f_update(fi1, ten, block1);
        int owner1 = s_get_owner(s1, block1);
        sendrecv_sketch_block(s_mat1, s1, fi1, owner2, 0);

        if (rank == owner2){

            // Get submatrix of s2
            matrix_tt* s_mat2 = own_submatrix(s2, block2 , 0);

            // Multiply
            matrix_tt* train_submat = submatrix(train_mat, 0, 0, 0, 0);
            matrix_tt* s_submat2 = submatrix(s_mat2, 0, 0, 0, 0);
            s_submat2->transpose = (iscol) ? 0 : 1;
            for (int kk = 0; kk < r2; ++kk){
                int jj0 = (fi2->f_t_index[(iscol) ? (fi2->f_d) - 1 : 0]) + kk * n;
                int sz = fi2->f_t_sizes[(iscol) ? (fi2->f_d) - 1 : 0];
                submatrix_update(train_submat, 0, r1, jj0, jj0+sz);

                s_submat2->X = s_mat2->X + s_mat2->offset + kk*(s_mat2->lda);
                if (iscol){
                    s_submat2->m = s_mat1->m;
                    s_submat2->lda = s_mat1->m;
                    s_submat2->n = sz;
                }
                else{
                    s_submat2->m = sz;
                    s_submat2->lda = sz;
                    s_submat2->n = s_mat1->m;
                }

                s_mat1->transpose = 1;
                matrix_tt_dgemm(s_mat1, s_submat2, train_submat, 1.0, 1.0);
            }
            free(train_submat);
            free(s_mat2);
            free(s_submat2);
        }
    }
    free(s_mat1);

    matrix_tt_reduce(comm, rank, train_mat, NULL, head);

    // If it was a row sketch, we need to transpose
    if ((!iscol) && (rank == head)){
        double* train_cp = train_transpose(train, r1, n, r2);
        *train_ptr = train_cp;

        matrix_tt_free(train_mat); train_mat = NULL;
    }
    else{
        free(train_mat); train_mat = NULL;
    }

    flattening_info_free(fi1);
    flattening_info_free(fi2);
}

void PSTT2_final_train(tensor_train* tt, sketch** sketches, int mid)
{
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    int head = 0;

    sketch* s1 = sketches[mid-1];
    int r1 = s1->r;
    sketch* s2 = sketches[mid];
    int r2 = s2->r;
    int n = tt->n[mid];

    MPI_tensor* ten = s1->ten;
    MPI_Comm comm = ten->comm;
    int size = ten->comm_size;
    int* schedule_rank = ten->schedule[rank];

    flattening_info* fi1 = flattening_info_init(ten, s1->flattening, s1->iscol, 0);
    flattening_info* fi2 = flattening_info_init(ten, s2->flattening, s2->iscol, 0);

    matrix_tt* train_mat = matrix_tt_wrap(r1*n, r2, tt->trains[mid]);
    matrix_tt_fill_zeros(train_mat);

    int tmp_n = 1;
    int* mid_partition = ten->partitions[mid];
    for (int ii = 0; ii < ten->nps[mid]; ++ii){
        int candidate = mid_partition[ii+1] - mid_partition[ii];
        tmp_n = (candidate > tmp_n) ? candidate : tmp_n;
    }
    matrix_tt* tmp = matrix_tt_init(r1, tmp_n);

    matrix_tt* recv_mat = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    matrix_tt* s_mat1 = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    matrix_tt* s_mat2 = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    matrix_tt* s_submat2 = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    matrix_tt* subtensor_mat = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    // Loop over schedule and stream
    for (int ii = 0; ii < ten->n_schedule; ++ii){
        int block = schedule_rank[ii];
        stream(ten, block);

        // Get the correct sketch matrices for multiplication
        for (int jj = 0; jj < size; ++jj){
            int* schedule_jj = ten->schedule[jj];
            int block_jj = schedule_jj[ii];

            if (block_jj != -1){
                // NOTE: Valgrind says I have a memory leak with these allocations.
                //       I think it's just confused by MPI, but maybe it's worth checking out
                flattening_info_update(fi1, ten, block_jj);
                sendrecv_sketch_block(recv_mat, s1, fi1, jj, 0);
                if (recv_mat->X != NULL){
                    matrix_tt* swtch = recv_mat;
                    recv_mat = s_mat1;
                    s_mat1 = swtch;
                }

                flattening_info_update(fi2, ten, block_jj);
                sendrecv_sketch_block(recv_mat, s2, fi2, jj, 0);
                if (recv_mat->X != NULL){
                    matrix_tt* swtch = recv_mat;
                    recv_mat = s_mat2;
                    s_mat2 = swtch;
                }
            }
        }

        flattening_info_update(fi1, ten, block);
        flattening_info_update(fi2, ten, block);

        // Multiply
        if (block != -1){
            int jj0 = fi1->t_t_index[mid];
            int sz = fi1->t_t_sizes[mid];
            submatrix_update(tmp, 0, r1, 0, sz);
            matrix_tt_wrap_update(subtensor_mat, fi1->f_N, fi1->s_N, get_X(ten));
            s_mat1->transpose = 1;
            submatrix_update(train_mat, jj0*r1, (jj0 + sz)*r1, 0, r2);
            submatrix_update(subtensor_mat, 0, 0, 0, 0);

            s_submat2->transpose = s_mat2->transpose;
            s_submat2->offset = 0;
            s_submat2->lda = s_mat2->lda;
            s_submat2->X_size = s_mat2->X_size;
            s_submat2->X = s_mat2->X + s_mat2->offset;

            for (int jj = 0; jj < fi2->f_N; ++jj){
                matrix_tt_reshape(r1, sz, tmp);
                submatrix_update(subtensor_mat, 0, fi1->f_N, jj*sz, (jj+1)*sz);

                matrix_tt_dgemm(s_mat1, subtensor_mat, tmp, 1.0, 0.0);
                matrix_tt_reshape(r1*sz, 1, tmp);

                submatrix_update(s_submat2, jj, jj+1, 0, r2);
                matrix_tt_dgemm(tmp, s_submat2, train_mat, 1.0, 1.0);
            }
        }

    }
    free(recv_mat);
    free(s_mat1);
    free(s_mat2);
    free(s_submat2);
    free(subtensor_mat);

    flattening_info_free(fi1);
    flattening_info_free(fi2);

    submatrix_update(train_mat, 0, n*r1, 0, r2);
    matrix_tt_reduce(comm, rank, train_mat, NULL, head);
    matrix_tt_free(tmp);

    free(train_mat);
}

int PSTT2_get_mid(MPI_tensor* ten, int mid)
{
    int d = ten->d;
    if ((mid > 0) && (mid < d)){
        return mid;
    }

    int* n = ten->n;

    long n_left = 1;
    long n_right = 1;
    for (int ii = 0; ii < d; ++ii){
        n_right = (long) n_right * n[ii];
    }

    for (int ii = 0; ii < d; ++ii){
        n_left = (long) n_left * n[ii];
        n_right = (long) n_right / n[ii];
        if ( (mid == -1) && (n_right < n_left) ){
            mid = ii;
        }
    }

    // Just in case, but this should never actually happen
    if (mid == -1){
        printf("You entered a weird tensor (probably n_i = 1 identically), setting mid = d-1\n");
        mid = d-1;
    }

    return mid;
}

void PSTT2(tensor_train* tt, MPI_tensor* ten, int mid)
{
    mid = PSTT2_get_mid(ten, mid);
    int rank = ten->rank;
    int BUF = 2;
    int* r = tt->r;
    int d = ten->d;

    // Create sketches
    sketch** sketches = (sketch**) calloc(d-1, sizeof(sketch*));
    for (int ii = 0; ii < d-1; ++ii){
        int iscol = (ii < mid) ? 1 : 0; // Column sketch below mid, row sketch above it
        sketches[ii] = sketch_init(ten, ii+1, r[ii+1], BUF, iscol);
    }

    // Sketch the tensor
    multi_perform_sketch(sketches, d-1);

    for (int ii = 0; ii < d-1; ++ii){
        sketch_qr(sketches[ii]);
    }

    // multiply out the sketches
    for (int ii = 0; ii < mid; ++ii){
        sketch* s1 = (ii == 0) ? NULL : sketches[ii-1];
        sketch* s2 = sketches[ii];
        two_sketches_to_train(s1, s2, tt->trains + ii);
    }

    for (int ii = mid; ii < d-1; ++ii){
        sketch* s1 = (ii == d-2) ? NULL : sketches[ii+1];
        sketch* s2 = sketches[ii];
        two_sketches_to_train(s1, s2, tt->trains + ii + 1);
    }

    // Get the middle train
    PSTT2_final_train(tt, sketches, mid);

    // free
    for (int ii = 0; ii < d-1; ++ii){
        sketch_free(sketches[ii]); sketches[ii] = NULL;
    }
    free(sketches);
}

void PSTT2_onepass_final_train(tensor_train* tt, sketch** sketches, int mid)
{
    sketch* sketch_mid = sketches[mid];
    sketch* Q_right = sketches[mid+1];
    sketch** Q_left = sketches + mid - 1; //Note that we are using the pointer here to feed into sketch_to_tensor

    // Multiply the middle not-QRed sketch against the sketch to its right
    int p = sketch_mid->r; // this is r_left + buf
    int r_left = tt->r[mid];
    int n_mid = tt->n[mid];
    int r_right = tt->r[mid+1];
    double* b = (double*) malloc(p*n_mid*r_right*sizeof(double)); // Named for the Ax=b solve that we will do later
    two_sketches_to_train(sketch_mid, Q_right, &b);

    MPI_tensor* ten_left = sketch_to_tensor(Q_left);

    sketch* sketch_left = sketch_init_with_Omega(ten_left, ten_left->d - 1, p, 0, 0, sketch_mid->Omegas); // buf = 0, is_col = 0
    perform_sketch(sketch_left);
    double* A = (double*) calloc(p * r_left, sizeof(double));

    two_sketches_to_train(sketch_left, NULL, &A); // This reduces the matrix A

    // Solve least squares
    matrix_tt* A_mat = matrix_tt_wrap(p, r_left, A);
    matrix_tt* b_mat = matrix_tt_wrap(p, n_mid*r_right, b);
    matrix_tt* x_mat = matrix_tt_wrap(r_left, n_mid*r_right, tt->trains[mid]);

    int head = 0;
    if (ten_left->rank == head){
        matrix_tt_dgels(x_mat, A_mat, b_mat);
    }

    // Free
    sketch_free(sketch_left); sketch_left = NULL;
    MPI_tensor_free(ten_left); ten_left = NULL;
    matrix_tt_free(A_mat); A_mat = NULL;
    matrix_tt_free(b_mat); b_mat = NULL;
    free(x_mat);
}

void PSTT2_onepass(tensor_train* tt, MPI_tensor* ten, int mid)
{
    // Number of times in PSTT2: 5
    // Number of times in multi_perform_sketch: 3
    mid = PSTT2_get_mid(ten, mid);
    int rank = ten->rank;
    int BUF = 2;
    int* r = tt->r;
    int d = ten->d;

    // Create sketches
    sketch** sketches = (sketch**) calloc(d, sizeof(sketch*));
    for (int ii = 0; ii < mid; ++ii){
        int iscol = 1; // Column sketch below mid, row sketch above it
        sketches[ii] = sketch_init(ten, ii+1, r[ii+1], BUF, iscol);
    }

    // putting the BUF into the rank here, so that two_sketches_to_train and perform_sketch don't get confused in the future
    sketches[mid] = sketch_init(ten, mid, r[mid] + BUF, 0, 0);

    for (int ii = mid + 1; ii < d; ++ii){
        int iscol = 0;
        sketches[ii] = sketch_init(ten, ii, r[ii], BUF, iscol);
    }

    // Sketch the tensor
    multi_perform_sketch(sketches, d);

    // Find orthogonal basis
    for (int ii = 0; ii < mid; ++ii){
        sketch_qr(sketches[ii]);
    }
    for (int ii = mid+1; ii < d; ++ii){
        sketch_qr(sketches[ii]);
    }

    // multiply out the sketches
    for (int ii = 0; ii < mid; ++ii){
        sketch* s1 = (ii == 0) ? NULL : sketches[ii-1];
        sketch* s2 = sketches[ii];
        two_sketches_to_train(s1, s2, tt->trains + ii);
    }

    for (int ii = mid+1; ii < d; ++ii){
        sketch* s1 = (ii == d-1) ? NULL : sketches[ii+1];
        sketch* s2 = sketches[ii];
        two_sketches_to_train(s1, s2, tt->trains + ii);
    }

    // Get the middle train
    PSTT2_onepass_final_train(tt, sketches, mid);

    // free
    for (int ii = 0; ii < d; ++ii){
        if(sketches[ii] != NULL){
            sketch_free(sketches[ii]); sketches[ii] = NULL;
        }
    }
    free(sketches);
}


#pragma once

#include "movie_writer.h"

#include <boost/math/tools/roots.hpp>
#include <cxxopts.hpp>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtx/rotate_vector.hpp>
#include <glm/gtx/string_cast.hpp>
#include <gsl/gsl_integration.h>
#include <omp.h>
#include <opencv2/opencv.hpp>

#include <QApplication>
#include <QCheckBox>
#include <QGraphicsView>
#include <QGuiApplication>
#include <QHBoxLayout>
#include <QLineEdit>
#include <QMainWindow>
#include <QMenuBar>
#include <QPushButton>
#include <QVBoxLayout>
#include <QFileDialog>
#include <QGraphicsPixmapItem>
#include <QPixmap>

#include <array>
#include <chrono>
#include <cmath>
#include <filesystem>
#include <iostream>
#include <random>
#include <thread>
#include <vector>


#ifndef TNet_Matrix_h
#define TNet_Matrix_h

#include <stddef.h>
#include <stdlib.h>
#include <stdexcept>
#include <iostream>

#ifdef HAVE_ATLAS
extern "C"{
  #include <cblas.h>
  #include <clapack.h>
}
#endif

#include "Common.h"
#include "MathAux.h"
#include "Types.h"
#include "Error.h"

//#define TRACE_MATRIX_OPERATIONS
#define CHECKSIZE

namespace TNet
{

  //  class matrix_error : public std::logic_error {};
  //  class matrix_sizes_error : public matrix_error {};

  // declare the class so the header knows about it
  template<typename _ElemT> class Vector;
  template<typename _ElemT> class SubVector;
  template<typename _ElemT> class Matrix;
  template<typename _ElemT> class SubMatrix;

  // we need to declare the friend << operator here
  template<typename _ElemT>
    std::ostream & operator << (std::ostream & rOut, const Matrix<_ElemT> & rM);

  // we need to declare the friend << operator here
  template<typename _ElemT>
    std::istream & operator >> (std::istream & rIn, Matrix<_ElemT> & rM);

  // we need to declare this friend function here
  template<typename _ElemT>
   _ElemT TraceOfProduct(const Matrix<_ElemT> &A, const Matrix<_ElemT> &B); // tr(A B)

  // we need to declare this friend function here
  template<typename _ElemT>
	 _ElemT TraceOfProductT(const Matrix<_ElemT> &A, const Matrix<_ElemT> &B); // tr(A B^T)==tr(A^T B)

  /** **************************************************************************
   ** **************************************************************************
   *  @brief Provides a matrix class
   *
   *  This class provides a way to work with matrices in TNet.
   *  It encapsulates basic operations and memory optimizations.
   *
   */
  template<typename _ElemT>
    class Matrix
    {
    public:
      /// defines a transpose type

      struct HtkHeader
      {
        INT_32    mNSamples;
        INT_32    mSamplePeriod;
        INT_16    mSampleSize;
        UINT_16   mSampleKind;
      };

      /**
       * @brief Extension of the HTK header
       */
      struct HtkHeaderExt
      {
        INT_32 mHeaderSize;
        INT_32 mVersion;
        INT_32 mSampSize;
      };

      /// defines a type of this
      typedef Matrix<_ElemT>    ThisType;

      // Constructors

      /// Empty constructor
      Matrix<_ElemT> ():
        mpData(NULL), mMCols(0), mMRows(0), mStride(0)
#ifdef STK_MEMALIGN_MANUAL
        , mpFreeData(NULL)
#endif
      {}

      /// Copy constructor
      Matrix<_ElemT> (const Matrix<_ElemT> & rM, MatrixTrasposeType trans=NO_TRANS):
        mpData(NULL)
      { if(trans==NO_TRANS){ Init(rM.mMRows, rM.mMCols); Copy(rM); } else { Init(rM.mMCols,rM.mMRows); Copy(rM,TRANS); } }

      /// Copy constructor from another type.
      template<typename _ElemU>
      explicit Matrix<_ElemT> (const Matrix<_ElemU> & rM, MatrixTrasposeType trans=NO_TRANS):
        mpData(NULL)
      { if(trans==NO_TRANS){ Init(rM.Rows(), rM.Cols()); Copy(rM); } else { Init(rM.Cols(),rM.Rows()); Copy(rM,TRANS); } }

      /// Basic constructor
      Matrix(const size_t r, const size_t c, bool clear=true)
      { mpData=NULL; Init(r, c, clear);  }

	  Matrix<_ElemT> &operator = (const Matrix <_ElemT> &other) { Init(other.Rows(), other.Cols()); Copy(other); return *this; } // Needed for inclusion in std::vector

      /// Destructor
      ~Matrix()
      { Destroy(); }

      /// Initializes matrix (if not done by constructor)
      ThisType &
      Init(const size_t r,
           const size_t c, bool clear=true);

      /**
       * @brief Dealocates the matrix from memory and resets the dimensions to (0, 0)
       */
      void
      Destroy();

      ThisType &
      Zero();

      ThisType &
      Unit(); // set to unit.

      /**
       * @brief Copies the contents of a matrix
       * @param rM Source data matrix
       * @return Returns reference to this
       */
      template<typename _ElemU> ThisType &
        Copy(const Matrix<_ElemU> & rM, MatrixTrasposeType Trans=NO_TRANS);

      /**
       * @brief Copies the elements of a vector row-by-row into a matrix
       * @param rV Source vector
       * @param nRows Number of rows of returned matrix
       * @param nCols Number of columns of returned matrix
       *
       * Note that rV.Dim() must equal nRows*nCols
       */
      ThisType &
      CopyVectorSplicedRows(const Vector<_ElemT> &rV, const size_t nRows, const size_t nCols);

      /**
       * @brief Returns @c true if matrix is initialized
       */
      bool
		IsInitialized() const
      { return mpData != NULL; }

      /// Returns number of rows in the matrix
      inline size_t
		Rows() const
      {
        return mMRows;
      }

      /// Returns number of columns in the matrix
      inline size_t
      Cols() const
      {
        return mMCols;
      }

      /// Returns number of columns in the matrix memory
      inline size_t
		Stride() const
      {
        return mStride;
      }

      /**
       *  @brief Gives access to a specified matrix row without range check
       *  @return Pointer to the const array
       */
      inline const _ElemT*  __attribute__((aligned(16)))
       pData () const
      {
        return mpData;
      }

      /**
       *  @brief Gives access to a specified matrix row without range check
       *  @return Pointer to the non-const data array
       */
      inline _ElemT* __attribute__((aligned(16)))
       pData ()
      {
        return mpData;
      }

      /**
       *  @brief pData_workaround is a workaround that allows SubMatrix to get a
       *  @return pointer to non-const data even though the Matrix is const...
       */
    protected:
      inline _ElemT*  __attribute__((aligned(16)))
       pData_workaround () const
      {
        return mpData;
      }
    public:

      /// Returns size of matrix in memory
      size_t
      MSize() const
      {
        return mMRows * mStride * sizeof(_ElemT);
      }

      /// Checks the content of the matrix for nan and inf values
      void
      CheckData(const std::string file = "") const
      {
        for(size_t row=0; row<Rows(); row++) {
          for(size_t col=0; col<Cols(); col++) {
            if(isnan((*this)(row,col)) || isinf((*this)(row,col))) {
              std::ostringstream os;
              os << "Invalid value: " << (*this)(row,col)
                 << " in matrix row: " << row
                 << " col: " << col
                 << " file: " << file;
              Error(os.str());
            }
          }
        }
      }

      /**
       *  **********************************************************************
       *  **********************************************************************
       *  @defgroup RESHAPE Matrix reshaping rutines
       *  **********************************************************************
       *  **********************************************************************
       * @{
       */

      /**
       *  @brief Removes one row from the matrix. The memory is not reallocated.
       */
      ThisType &
      RemoveRow(size_t i);

      /** @} */

      /**
       *  **********************************************************************
       *  **********************************************************************
       *  @defgroup ACCESS Access functions and operators
       *  **********************************************************************
       *  **********************************************************************
       * @{
       */

      /**
       *  @brief Gives access to a specified matrix row without range check
       *  @return Subvector object representing the row
       */
      inline const SubVector<_ElemT>
      operator []  (size_t i) const
      {
        assert(i < mMRows);
        return SubVector<_ElemT>(mpData + (i * mStride), Cols());
      }

      inline SubVector<_ElemT>
      operator []  (size_t i)
      {
        assert(i < mMRows);
        return SubVector<_ElemT>(mpData + (i * mStride), Cols());
      }

      /**
       *  @brief Gives access to a specified matrix row without range check
       *  @return pointer to the first field of the row
       */
      inline  _ElemT*
      pRowData(size_t i)
      {
        assert(i < mMRows);
        return mpData + i * mStride;
      }

      /**
       *  @brief Gives access to a specified matrix row without range check
       *  @return pointer to the first field of the row (const version)
       */
      inline const _ElemT*
      pRowData(size_t i) const
      {
        assert(i < mMRows);
        return mpData + i * mStride;
      }

      /**
       *  @brief Gives access to matrix elements (row, col)
       *  @return reference to the desired field
       */
      inline _ElemT&
		operator () (size_t r, size_t c)
      {
#ifdef PARANOID
        assert(r < mMRows && c < mMCols);
#endif
		return *(mpData + r * mStride + c);
	  }

      /**
       *  @brief Gives access to matrix elements (row, col)
       *  @return pointer to the desired field (const version)
       */
      inline const _ElemT
		operator () (size_t r, size_t c) const
      {
#ifdef PARANOID
        assert(r < mMRows && c < mMCols);
#endif
		return *(mpData + r * mStride + c);
	  }

      /**
       * @brief Returns a matrix sub-range
       * @param ro Row offset
       * @param r  Rows in range
       * @param co Column offset
       * @param c  Coluns in range
       * See @c SubMatrix class for details
       */
      SubMatrix<_ElemT>
      Range(const size_t    ro, const size_t    r,
            const size_t    co, const size_t    c)
      { return SubMatrix<_ElemT>(*this, ro, r, co, c); }

      const SubMatrix<_ElemT>
      Range(const size_t    ro, const size_t    r,
            const size_t    co, const size_t    c) const
      { return SubMatrix<_ElemT>(*this, ro, r, co, c); }
      /** @} */

      /**
       *  **********************************************************************
       *  **********************************************************************
       *  @defgroup MATH ROUTINES
       *  **********************************************************************
       *  **********************************************************************
       *  @{
       **/

      /**
       *  @brief Returns sum of all elements
       */
      _ElemT&
      Sum() const;

      ThisType &
      DotMul(const ThisType& a);

      ThisType &
      Scale(_ElemT alpha);

      ThisType &
      ScaleCols(const Vector<_ElemT> &scale); // Equivalent to (*this) = (*this) * diag(scale).

      ThisType &
      ScaleRows(const Vector<_ElemT> &scale); // Equivalent to (*this) = diag(scale) * (*this);

      /// Sum another matrix rMatrix with this matrix
      ThisType&
      Add(const Matrix<_ElemT>& rMatrix);

      /// Sum scaled matrix rMatrix with this matrix
      ThisType&
      AddScaled(_ElemT alpha, const Matrix<_ElemT>& rMatrix);

      /// Apply log to all items of the matrix
      ThisType&
      ApplyLog();

      /**
       * @brief Computes the determinant of this matrix
       * @return Returns the determinant of a matrix
       * @ingroup MATH
       *
       */
      _ElemT LogAbsDeterminant(_ElemT *DetSign=NULL);

      /**
       *  @brief Performs matrix inplace inversion
       */
      ThisType &
      Invert(_ElemT *LogDet=NULL, _ElemT *DetSign=NULL, bool inverse_needed=true);

      /**
       *  @brief Performs matrix inplace inversion in double precision, even if this object is not double precision.
       */
      ThisType &
      InvertDouble(_ElemT *LogDet=NULL, _ElemT *DetSign=NULL, bool inverse_needed=true){
        double LogDet_tmp, DetSign_tmp;
        Matrix<double> dmat(*this); dmat.Invert(&LogDet_tmp, &DetSign_tmp, inverse_needed); if(inverse_needed) (*this).Copy(dmat);
        if(LogDet) *LogDet = LogDet_tmp; if(DetSign) *DetSign = DetSign_tmp;
        return *this;
      }

      /**
       *  @brief Inplace matrix transposition. Applicable only to square matrices
       */
      ThisType &
      Transpose()
      {
        assert(Rows()==Cols());
        size_t M=Rows();
        for(size_t i=0;i<M;i++)
          for(size_t j=0;j<i;j++){
           _ElemT &a = (*this)(i,j), &b = (*this)(j,i);
		   std::swap(a,b);
        }
		return *this;
      }

      bool IsSymmetric(_ElemT cutoff = 1.0e-05) const;

      bool IsDiagonal(_ElemT cutoff = 1.0e-05) const;

      bool IsUnit(_ElemT cutoff = 1.0e-05) const;

      bool IsZero(_ElemT cutoff = 1.0e-05) const;

      _ElemT FrobeniusNorm() const; // sqrt of sum of square elements.

      _ElemT LargestAbsElem() const; // largest absolute value.

      friend _ElemT TNet::TraceOfProduct<_ElemT>(const Matrix<_ElemT> &A, const Matrix<_ElemT> &B); // tr(A B)
      friend _ElemT TNet::TraceOfProductT<_ElemT>(const Matrix<_ElemT> &A, const Matrix<_ElemT> &B); // tr(A B^T)==tr(A^T B)
      friend class SubMatrix<_ElemT>; // so it can get around const restrictions on the pointer to mpData.

      /** **********************************************************************
       *  **********************************************************************
       *  @defgroup BLAS_ROUTINES BLAS ROUTINES
       *  @ingroup MATH
       *  **********************************************************************
       *  **********************************************************************
       **/

      ThisType &
      BlasGer(const _ElemT alpha, const Vector<_ElemT>& rA, const Vector<_ElemT>& rB);

      ThisType &
	    Axpy(const _ElemT alpha, const Matrix<_ElemT> &rM, MatrixTrasposeType transA=NO_TRANS);

      ThisType &
      BlasGemm(const _ElemT alpha,
               const ThisType& rA, MatrixTrasposeType transA,
               const ThisType& rB, MatrixTrasposeType transB,
               const _ElemT beta = 0.0);

      /** @} */

      /** **********************************************************************
       *  **********************************************************************
       *  @defgroup IO Input/Output ROUTINES
       *  **********************************************************************
       *  **********************************************************************
       *  @{
       **/

      friend std::ostream &
      operator << <> (std::ostream & out, const ThisType & m);

      void PrintOut(char *file);
      void ReadIn(char *file);

      bool
      LoadHTK(const char* pFileName);

      /** @} */

    protected:
//      inline void swap4b(void *a);
//      inline void swap2b(void *a);

    protected:
      /// data memory area
      _ElemT*   mpData;

      /// these atributes store the real matrix size as it is stored in memory
      /// including memalignment
      size_t    mMCols;       ///< Number of columns
      size_t    mMRows;       ///< Number of rows
      size_t    mStride;      ///< true number of columns for the internal matrix.
                              ///< This number may differ from M_cols as memory
                              ///< alignment might be used

#ifdef STK_MEMALIGN_MANUAL
      /// data to be freed (in case of manual memalignment use, see Common.h)
      _ElemT*   mpFreeData;
#endif
    }; // class Matrix

    template<>  Matrix<float> &  Matrix<float>::Invert(float *LogDet, float *DetSign, bool inverse_needed); // state that we will implement separately for float and double.
    template<>  Matrix<double> &  Matrix<double>::Invert(double *LogDet, double *DetSign, bool inverse_needed);

  /** **************************************************************************
   ** **************************************************************************
   *  @brief Sub-matrix representation
   *
   *  This class provides a way to work with matrix cutouts in STK.
   *
   *
   */
  template<typename _ElemT>
    class SubMatrix : public Matrix<_ElemT>
    {
    typedef SubMatrix<_ElemT>    ThisType;

    public:
      /// Constructor
      SubMatrix(const Matrix<_ElemT>& rT, // Input matrix cannot be const because SubMatrix can change its contents.
                const size_t    ro,
                const size_t    r,
                const size_t    co,
                const size_t    c);

      /// The destructor
      ~SubMatrix<_ElemT>()
      {
#ifndef STK_MEMALIGN_MANUAL
        Matrix<_ElemT>::mpData = NULL;
#else
        Matrix<_ElemT>::mpFreeData = NULL;
#endif
      }

      /// Assign operator
      ThisType& operator=(const ThisType& rSrc)
      {
        //std::cout << "[PERFORMing operator= SubMatrix&^2]" << std::flush;
        this->mpData = rSrc.mpData;
        this->mMCols = rSrc.mMCols;
        this->mMRows = rSrc.mMRows;
        this->mStride = rSrc.mStride;
        this->mpFreeData = rSrc.mpFreeData;
        return *this;
      }

      /// Initializes matrix (if not done by constructor)
      ThisType &
      Init(const size_t r,
           const size_t c, bool clear=true)
      { Error("Submatrix cannot do Init"); return *this; }

      /**
       * @brief Dealocates the matrix from memory and resets the dimensions to (0, 0)
       */
      void
      Destroy()
      { Error("Submatrix cannot do Destroy"); }

    };

  //Create useful shortcuts
  typedef Matrix<BaseFloat> BfMatrix;
  typedef SubMatrix<BaseFloat> BfSubMatrix;

  /**
   * Function for summing matrices of different types
   */
  template<typename _ElemT, typename _ElemU>
  void Add(Matrix<_ElemT>& rDst,  const Matrix<_ElemU>& rSrc) {
    assert(rDst.Cols() == rSrc.Cols());
    assert(rDst.Rows() == rSrc.Rows());

    for(size_t i=0; i<rDst.Rows(); i++) {
      const _ElemU* p_src = rSrc.pRowData(i);
      _ElemT* p_dst = rDst.pRowData(i);
      for(size_t j=0; j<rDst.Cols(); j++) {
        *p_dst++ += (_ElemT)*p_src++;
      }
    }
  }

  /**
   * Function for summing matrices of different types
   */
  template<typename _ElemT, typename _ElemU>
  void AddScaled(Matrix<_ElemT>& rDst, const Matrix<_ElemU>& rSrc, _ElemT scale) {
    assert(rDst.Cols() == rSrc.Cols());
    assert(rDst.Rows() == rSrc.Rows());

    Vector<_ElemT> tmp(rDst[0]);

    for(size_t i=0; i<rDst.Rows(); i++) {
      tmp.Copy(rSrc[i]);
      rDst[i].BlasAxpy(scale, tmp);

      /*
      const _ElemU* p_src = rSrc.pRowData(i);
      _ElemT* p_dst = rDst.pRowData(i);
      for(size_t j=0; j<rDst.Cols(); j++) {
        *p_dst++ += (_ElemT)(*p_src++) * scale;
      }
      */
    }
  }

} // namespace STK

//*****************************************************************************
//*****************************************************************************
// we need to include the implementation
#include "Matrix.tcc"
//*****************************************************************************
//*****************************************************************************

/******************************************************************************
 ******************************************************************************
 * The following section contains specialized template definitions
 * whose implementation is in Matrix.cc
 */

//#ifndef TNet_Matrix_h
#endif


/**
 *
 * @file core_slarfb_gemm.c
 *
 *  PLASMA core_blas kernel
 *  PLASMA is a software package provided by Univ. of Tennessee,
 *  Univ. of California Berkeley and Univ. of Colorado Denver
 *
 * @version 2.6.0
 * @author Azzam Haidar
 * @date 2010-11-15
 * @generated s Tue Jan  7 11:44:49 2014
 *
 **/
#include <cblas.h>
#include "common.h"

/***************************************************************************//**
 *
 * @ingroup CORE_float
 *
 *  CORE_slarfb_gemm applies a complex block reflector H or its transpose H'
 *  to a complex M-by-N matrix C, from either the left or the right.
 *  this kernel is similar to the lapack slarfb but it do a full gemm on the
 *  triangular Vs assuming that the upper part of Vs is zero and ones are on
 *  the diagonal. It is also based on the fact that a gemm on a small block of
 *  k reflectors is faster than a trmm on the triangular (k,k) + gemm below.
 *
 *  NOTE THAT:
 *  Only Columnwise/Forward cases are treated here.
 *
 *******************************************************************************
 *
 * @param[in] side
 *         @arg PlasmaLeft  : apply Q or Q**T from the Left;
 *         @arg PlasmaRight : apply Q or Q**T from the Right.
 *
 * @param[in] trans
 *         @arg PlasmaNoTrans   : No transpose, apply Q;
 *         @arg PlasmaTrans : ConjTranspose, apply Q**T.
 *
 * @param[in] direct
 *         Indicates how H is formed from a product of elementary
 *         reflectors
 *         @arg PlasmaForward  : H = H(1) H(2) . . . H(k) (Forward)
 *         @arg PlasmaBackward : H = H(k) . . . H(2) H(1) (Backward)
 *
 * @param[in] storev
 *         Indicates how the vectors which define the elementary
 *         reflectors are stored:
 *         @arg PlasmaColumnwise
 *         @arg PlasmaRowwise
 *
 * @param[in] M
 *         The number of rows of the matrix C.
 *
 * @param[in] N
 *         The number of columns of the matrix C.
 *
 * @param[in] K
 *          The order of the matrix T (= the number of elementary
 *          reflectors whose product defines the block reflector).
 *
 * @param[in] V
 *          COMPLEX*16 array, dimension
 *              (LDV,K) if storev = 'C'
 *              (LDV,M) if storev = 'R' and side = 'L'
 *              (LDV,N) if storev = 'R' and side = 'R'
 *          The matrix V. See further details.
 *
 * @param[in] LDV
 *         The leading dimension of the array V.
 *         If storev = 'C' and side = 'L', LDV >= max(1,M);
 *         if storev = 'C' and side = 'R', LDV >= max(1,N);
 *         if storev = 'R', LDV >= K.
 *
 * @param[in] T
 *         The triangular K-by-K matrix T in the representation of the
 *         block reflector.
 *         T is upper triangular by block (economic storage);
 *         The rest of the array is not referenced.
 *
 * @param[in] LDT
 *         The leading dimension of the array T. LDT >= K.
 *
 * @param[in,out] C
 *         COMPLEX*16 array, dimension (LDC,N)
 *         On entry, the M-by-N matrix C.
 *         On exit, C is overwritten by H*C or H'*C or C*H or C*H'.
 *
 * @param[in] LDC
 *         The leading dimension of the array C. LDC >= max(1,M).
 *
 * @param[in,out] WORK
 *         (workspace) COMPLEX*16 array, dimension (LDWORK,K).
 *
 * @param[in] LDWORK
 *         The dimension of the array WORK.
 *         If side = PlasmaLeft,  LDWORK >= max(1,N);
 *         if side = PlasmaRight, LDWORK >= max(1,M).
 *
 *******************************************************************************
 *
 * @return
 *          \retval PLASMA_SUCCESS successful exit
 *          \retval <0 if -i, the i-th argument had an illegal value
 *
 ******************************************************************************/
#if defined(PLASMA_HAVE_WEAK)
#pragma weak CORE_slarfb_gemm = PCORE_slarfb_gemm
#define CORE_slarfb_gemm PCORE_slarfb_gemm
#endif
int CORE_slarfb_gemm(PLASMA_enum side, PLASMA_enum trans, int direct, int storev,
                     int M, int N, int K,
                     const float *V, int LDV,
                     const float *T, int LDT,
                     float *C, int LDC,
                     float *WORK, int LDWORK)
{
    static float zzero =  0.0;
    static float zone  =  1.0;
    static float mzone = -1.0;
    /* Check input arguments */
    if ((side != PlasmaLeft) && (side != PlasmaRight)) {
        coreblas_error(1, "Illegal value of side");
        return -1;
    }
    if ((trans != PlasmaNoTrans) && (trans != PlasmaTrans)) {
        coreblas_error(2, "Illegal value of trans");
        return -2;
    }
    if ((direct != PlasmaForward) && (direct != PlasmaBackward)) {
        coreblas_error(3, "Illegal value of direct");
        return -3;
    }
    if ((storev != PlasmaColumnwise) && (storev != PlasmaRowwise)) {
        coreblas_error(4, "Illegal value of direct");
        return -4;
    }
    if (M < 0) {
        coreblas_error(5, "Illegal value of M");
        return -5;
    }
    if (N < 0) {
        coreblas_error(6, "Illegal value of N");
        return -6;
    }
    if (K < 0) {
        coreblas_error(7, "Illegal value of K");
        return -7;
    }

    /* Quick return */
    if ((M == 0) || (N == 0) || (K == 0) )
        return PLASMA_SUCCESS;

    /* For Left case, switch the trans. noswitch for right case */
    if( side == PlasmaLeft){
        if ( trans == PlasmaNoTrans) {
            trans = PlasmaTrans;
        }
        else {
            trans = PlasmaNoTrans;
        }
    }

    /* main code */
    if (storev == PlasmaColumnwise )
    {
        if ( direct == PlasmaForward )
        {
            /*
             * Let  V =  ( V1 )    (first K rows are unit Lower triangular)
             *           ( V2 )
             */
            if ( side == PlasmaLeft )
            {
                /*
                 * Columnwise / Forward / Left
                 */
                /*
                 * Form  H * C  or  H' * C  where  C = ( C1 )
                 *                                     ( C2 )
                 *
                 * W := C' * V    (stored in WORK)
                 */
                cblas_sgemm(
                    CblasColMajor, CblasTrans, CblasNoTrans,
                    N, K, M,
                    (zone), C, LDC,
                    V, LDV,
                    (zzero), WORK, LDWORK );
                /*
                 * W := W * T'  or  W * T
                 */
                cblas_strmm(
                    CblasColMajor, CblasRight, CblasUpper,
                    (CBLAS_TRANSPOSE)trans, CblasNonUnit, N, K,
                    (zone), T, LDT, WORK, LDWORK );
                /*
                 * C := C - V * W'
                 */
                cblas_sgemm(
                    CblasColMajor, CblasNoTrans, CblasTrans,
                    M, N, K,
                    (mzone), V, LDV,
                    WORK, LDWORK,
                    (zone), C, LDC);
            }
            else {
                /*
                 * Columnwise / Forward / Right
                 */
                /*
                 * Form  C * H  or  C * H'  where  C = ( C1  C2 )
                 * W := C * V
                 */
                cblas_sgemm(
                    CblasColMajor, CblasNoTrans, CblasNoTrans,
                    M, K, N,
                    (zone), C, LDC,
                    V, LDV,
                    (zzero), WORK, LDWORK);
                /*
                 * W := W * T  or  W * T'
                 */
                cblas_strmm(
                    CblasColMajor, CblasRight, CblasUpper,
                    (CBLAS_TRANSPOSE)trans, CblasNonUnit, M, K,
                    (zone), T, LDT, WORK, LDWORK );
                /*
                 * C := C - W * V'
                 */
                cblas_sgemm(
                    CblasColMajor, CblasNoTrans, CblasTrans,
                    M, N, K,
                    (mzone), WORK, LDWORK,
                    V, LDV,
                    (zone), C, LDC);
            }
        }
        else {
            /*
             * Columnwise / Backward / Left or Right
             */
            coreblas_error(3, "Not implemented (ColMajor / Backward / Left or Right)");
            return PLASMA_ERR_NOT_SUPPORTED;
        }
    }
    else {
        if (direct == PlasmaForward) {
            /*
             * Rowwise / Forward / Left or Right
             */
            coreblas_error(3, "Not implemented (RowMajor / Backward / Left or Right)");
            return PLASMA_ERR_NOT_SUPPORTED;
        }
        else {
            /*
             * Rowwise / Backward / Left or Right
             */
            coreblas_error(3, "Not implemented (RowMajor / Backward / Left or Right)");
            return PLASMA_ERR_NOT_SUPPORTED;
        }
    }
    return PLASMA_SUCCESS;
}


/*
 * Copyright (c) 1997-1999 Massachusetts Institute of Technology
 *
 * 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 2 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

/* This file was automatically generated --- DO NOT EDIT */
/* Generated on Sun Nov  7 20:44:25 EST 1999 */

#include <fftw-int.h>
#include <fftw.h>

/* Generated by: ./genfft -magic-alignment-check -magic-twiddle-load-all -magic-variables 4 -magic-loopi -hc2real 14 */

/*
 * This function contains 62 FP additions, 38 FP multiplications,
 * (or, 60 additions, 36 multiplications, 2 fused multiply/add),
 * 24 stack variables, and 28 memory accesses
 */
static const fftw_real K445041867 = FFTW_KONST(+0.445041867912628808577805128993589518932711138);
static const fftw_real K1_801937735 = FFTW_KONST(+1.801937735804838252472204639014890102331838324);
static const fftw_real K1_246979603 = FFTW_KONST(+1.246979603717467061050009768008479621264549462);
static const fftw_real K867767478 = FFTW_KONST(+0.867767478235116240951536665696717509219981456);
static const fftw_real K1_949855824 = FFTW_KONST(+1.949855824363647214036263365987862434465571601);
static const fftw_real K1_563662964 = FFTW_KONST(+1.563662964936059617416889053348115500464669037);
static const fftw_real K2_000000000 = FFTW_KONST(+2.000000000000000000000000000000000000000000000);

/*
 * Generator Id's :
 * $Id: exprdag.ml,v 1.41 1999/05/26 15:44:14 fftw Exp $
 * $Id: fft.ml,v 1.43 1999/05/17 19:44:18 fftw Exp $
 * $Id: to_c.ml,v 1.25 1999/10/26 21:41:32 stevenj Exp $
 */

void fftw_hc2real_14(const fftw_real *real_input, const fftw_real *imag_input, fftw_real *output, int real_istride, int imag_istride, int ostride)
{
     fftw_real tmp3;
     fftw_real tmp13;
     fftw_real tmp6;
     fftw_real tmp14;
     fftw_real tmp26;
     fftw_real tmp35;
     fftw_real tmp23;
     fftw_real tmp34;
     fftw_real tmp12;
     fftw_real tmp16;
     fftw_real tmp20;
     fftw_real tmp33;
     fftw_real tmp9;
     fftw_real tmp15;
     fftw_real tmp1;
     fftw_real tmp2;
     ASSERT_ALIGNED_DOUBLE;
     tmp1 = real_input[0];
     tmp2 = real_input[7 * real_istride];
     tmp3 = tmp1 - tmp2;
     tmp13 = tmp1 + tmp2;
     {
	  fftw_real tmp4;
	  fftw_real tmp5;
	  fftw_real tmp24;
	  fftw_real tmp25;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp4 = real_input[2 * real_istride];
	  tmp5 = real_input[5 * real_istride];
	  tmp6 = tmp4 - tmp5;
	  tmp14 = tmp4 + tmp5;
	  tmp24 = imag_input[2 * imag_istride];
	  tmp25 = imag_input[5 * imag_istride];
	  tmp26 = tmp24 - tmp25;
	  tmp35 = tmp24 + tmp25;
     }
     {
	  fftw_real tmp21;
	  fftw_real tmp22;
	  fftw_real tmp10;
	  fftw_real tmp11;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp21 = imag_input[6 * imag_istride];
	  tmp22 = imag_input[imag_istride];
	  tmp23 = tmp21 - tmp22;
	  tmp34 = tmp21 + tmp22;
	  tmp10 = real_input[6 * real_istride];
	  tmp11 = real_input[real_istride];
	  tmp12 = tmp10 - tmp11;
	  tmp16 = tmp10 + tmp11;
     }
     {
	  fftw_real tmp18;
	  fftw_real tmp19;
	  fftw_real tmp7;
	  fftw_real tmp8;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp18 = imag_input[4 * imag_istride];
	  tmp19 = imag_input[3 * imag_istride];
	  tmp20 = tmp18 - tmp19;
	  tmp33 = tmp18 + tmp19;
	  tmp7 = real_input[4 * real_istride];
	  tmp8 = real_input[3 * real_istride];
	  tmp9 = tmp7 - tmp8;
	  tmp15 = tmp7 + tmp8;
     }
     {
	  fftw_real tmp38;
	  fftw_real tmp37;
	  fftw_real tmp29;
	  fftw_real tmp28;
	  ASSERT_ALIGNED_DOUBLE;
	  output[7 * ostride] = tmp3 + (K2_000000000 * (tmp6 + tmp9 + tmp12));
	  tmp38 = (K1_563662964 * tmp35) + (K1_949855824 * tmp33) + (K867767478 * tmp34);
	  tmp37 = tmp3 + (K1_246979603 * tmp6) - (K1_801937735 * tmp12) - (K445041867 * tmp9);
	  output[ostride] = tmp37 - tmp38;
	  output[13 * ostride] = tmp37 + tmp38;
	  {
	       fftw_real tmp40;
	       fftw_real tmp39;
	       fftw_real tmp36;
	       fftw_real tmp32;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp40 = (K867767478 * tmp33) + (K1_563662964 * tmp34) - (K1_949855824 * tmp35);
	       tmp39 = tmp3 + (K1_246979603 * tmp12) - (K1_801937735 * tmp9) - (K445041867 * tmp6);
	       output[5 * ostride] = tmp39 - tmp40;
	       output[9 * ostride] = tmp39 + tmp40;
	       tmp36 = (K1_563662964 * tmp33) - (K1_949855824 * tmp34) - (K867767478 * tmp35);
	       tmp32 = tmp3 + (K1_246979603 * tmp9) - (K445041867 * tmp12) - (K1_801937735 * tmp6);
	       output[11 * ostride] = tmp32 - tmp36;
	       output[3 * ostride] = tmp32 + tmp36;
	  }
	  output[0] = tmp13 + (K2_000000000 * (tmp14 + tmp15 + tmp16));
	  tmp29 = (K867767478 * tmp20) + (K1_563662964 * tmp23) - (K1_949855824 * tmp26);
	  tmp28 = tmp13 + (K1_246979603 * tmp16) - (K1_801937735 * tmp15) - (K445041867 * tmp14);
	  output[12 * ostride] = tmp28 - tmp29;
	  output[2 * ostride] = tmp28 + tmp29;
	  {
	       fftw_real tmp31;
	       fftw_real tmp30;
	       fftw_real tmp27;
	       fftw_real tmp17;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp31 = (K1_563662964 * tmp26) + (K1_949855824 * tmp20) + (K867767478 * tmp23);
	       tmp30 = tmp13 + (K1_246979603 * tmp14) - (K1_801937735 * tmp16) - (K445041867 * tmp15);
	       output[8 * ostride] = tmp30 - tmp31;
	       output[6 * ostride] = tmp30 + tmp31;
	       tmp27 = (K1_563662964 * tmp20) - (K1_949855824 * tmp23) - (K867767478 * tmp26);
	       tmp17 = tmp13 + (K1_246979603 * tmp15) - (K445041867 * tmp16) - (K1_801937735 * tmp14);
	       output[4 * ostride] = tmp17 - tmp27;
	       output[10 * ostride] = tmp17 + tmp27;
	  }
     }
}

fftw_codelet_desc fftw_hc2real_14_desc =
{
     "fftw_hc2real_14",
     (void (*)()) fftw_hc2real_14,
     14,
     FFTW_BACKWARD,
     FFTW_HC2REAL,
     323,
     0,
     (const int *) 0,
};


#include <math.h>
// Constants not defined in MSVC's math.h
#ifndef M_SQRT1_2
#define M_SQRT1_2 0.70710678118654746172
#endif
#ifndef M_2_SQRTPI
#define M_2_SQRTPI 1.12837916709551255856
#endif
#include <gsl/gsl_spline.h>
#include <galpy_potentials.h>
// ChandrasekharDynamicalFrictionForce: 8 arguments: amp,ms,rhm,gamma^2,
// lnLambda, minr^2, ro, rf
double ChandrasekharDynamicalFrictionForceAmplitude(double R,double z,
						    double phi,double t,
						    double r2,
						    struct potentialArg * potentialArgs,
						    double vR,double vT,
						    double vz){
  double sr,X,Xfactor,d_ind,forceAmplitude;
  double * args= potentialArgs->args;
  //Get args
  double amp= *args;
  double ms= *(args+9);
  double rhm= *(args+10);
  double gamma2= *(args+11);
  double lnLambda= *(args+12);
  double ro= *(args+14);
  double rf= *(args+15);
  double GMvs;
  double r= sqrt( r2 );
  double v2=  vR * vR + vT * vT + vz * vz;
  double v= sqrt( v2 );
  // Constant or variable Lambda
  if ( lnLambda < 0 ) {
    GMvs= ms/v/v;
    if ( GMvs < rhm )
      lnLambda= 0.5 * log ( 1. + r2 / gamma2 / rhm / rhm );
    else
      lnLambda= 0.5 * log ( 1. + r2 / gamma2 / GMvs / GMvs );
  }
  d_ind= (r-ro)/(rf-ro);
  d_ind= d_ind <  0 ? 0. : ( d_ind > 1 ? 1. : d_ind);
  sr= gsl_spline_eval(*potentialArgs->spline1d,d_ind,*potentialArgs->acc1d);
  X= M_SQRT1_2 * v / sr;
  Xfactor= erf ( X ) - M_2_SQRTPI * X * exp ( - X * X );
  forceAmplitude= - amp * Xfactor * lnLambda / v2 / v \
    * calcDensity(R,z,phi,t,potentialArgs->nwrapped,
		  potentialArgs->wrappedPotentialArg);
  // Caching
  *(args + 1)= R;
  *(args + 2)= z;
  *(args + 3)= phi;
  *(args + 4)= t;
  *(args + 5)= vR;
  *(args + 6)= vT;
  *(args + 7)= vz;
  *(args + 8)= forceAmplitude;
  return forceAmplitude;
}
double ChandrasekharDynamicalFrictionForceRforce(double R,double z, double phi,
						 double t,
						 struct potentialArg * potentialArgs,
						 double vR,double vT,
						 double vz){
  double forceAmplitude;
  double * args= potentialArgs->args;
  double r2=  R * R + z * z;
  if ( r2 < *(args+13) )  // r < minr, don't bother caching
    return 0.;
  //Get args
  double cached_R= *(args + 1);
  double cached_z= *(args + 2);
  double cached_phi= *(args + 3);
  double cached_t= *(args + 4);
  double cached_vR= *(args + 5);
  double cached_vT= *(args + 6);
  double cached_vz= *(args + 7);
  if ( R != cached_R || phi != cached_phi || z != cached_z || t != cached_t \
       || vR != cached_vR || vT != cached_vT || vz != cached_vz )
    forceAmplitude= ChandrasekharDynamicalFrictionForceAmplitude(R,z,phi,t,r2,
								 potentialArgs,
								 vR,vT,vz);
  else
    forceAmplitude= *(args + 8);
  return forceAmplitude * vR;
}
double ChandrasekharDynamicalFrictionForcezforce(double R,double z, double phi,
						 double t,
						 struct potentialArg * potentialArgs,
						 double vR,double vT,
						 double vz){
  double forceAmplitude;
  double * args= potentialArgs->args;
  double r2=  R * R + z * z;
  if ( r2 < *(args+13) )  // r < minr, don't bother caching
    return 0.;
  //Get args
  double cached_R= *(args + 1);
  double cached_z= *(args + 2);
  double cached_phi= *(args + 3);
  double cached_t= *(args + 4);
  double cached_vR= *(args + 5);
  double cached_vT= *(args + 6);
  double cached_vz= *(args + 7);
  if ( R != cached_R || phi != cached_phi || z != cached_z || t != cached_t \
       || vR != cached_vR || vT != cached_vT || vz != cached_vz )
    //LCOV_EXCL_START
    forceAmplitude= ChandrasekharDynamicalFrictionForceAmplitude(R,z,phi,t,r2,
								 potentialArgs,
								 vR,vT,vz);
    //LCOV_EXCL_STOP
  else
    forceAmplitude= *(args + 8);
  return forceAmplitude * vz;
}
double ChandrasekharDynamicalFrictionForcephiforce(double R,double z,
						   double phi,double t,
						   struct potentialArg * potentialArgs,
						   double vR,double vT,
						   double vz){
  double forceAmplitude;
  double * args= potentialArgs->args;
  double r2=  R * R + z * z;
  if ( r2 < *(args+13) )  // r < minr, don't bother caching
    return 0.;
  //Get args
  double cached_R= *(args + 1);
  double cached_z= *(args + 2);
  double cached_phi= *(args + 3);
  double cached_t= *(args + 4);
  double cached_vR= *(args + 5);
  double cached_vT= *(args + 6);
  double cached_vz= *(args + 7);
  if ( R != cached_R || phi != cached_phi || z != cached_z || t != cached_t \
       || vR != cached_vR || vT != cached_vT || vz != cached_vz )
    //LCOV_EXCL_START
    forceAmplitude= ChandrasekharDynamicalFrictionForceAmplitude(R,z,phi,t,r2,
								 potentialArgs,
								 vR,vT,vz);
    //LCOV_EXCL_STOP
  else
    forceAmplitude= *(args + 8);
  return forceAmplitude * vT * R;
}