/* Copyright (c) 2011-2012, Jérémy Fix. All rights reserved. */

/* Redistribution and use in source and binary forms, with or without */
/* modification, are permitted provided that the following conditions are met: */

/* * Redistributions of source code must retain the above copyright notice, */
/* this list of conditions and the following disclaimer. */
/* * Redistributions in binary form must reproduce the above copyright notice, */
/* this list of conditions and the following disclaimer in the documentation */
/* and/or other materials provided with the distribution. */
/* * None of the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. */

/* THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS "AS IS" AND */
/* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED */
/* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE */
/* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE */
/* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL */
/* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR */
/* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER */
/* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, */
/* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE */
/* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. */

#define DEBUGMODE true

#ifndef UKF_SR_STATE_NDIM_H
#define UKF_SR_STATE_NDIM_H

#include <gsl/gsl_linalg.h> // For the Cholesky decomposition
#include <gsl/gsl_math.h>
#include <gsl/gsl_blas.h>

#include "ukf_types.h"
#include "ukf_math.h"

void print_mat(const gsl_matrix * A)
{
    for(unsigned int i = 0 ; i < A->size1 ; ++i)
    {
        for(unsigned int j = 0 ; j < A ->size2 ; ++j)
            printf("%e ", gsl_matrix_get(A,i,j));

        printf("\n");
    }
}

void print_vec(const gsl_vector * A)
{
    for(unsigned int i = 0 ; i < A->size ; ++i)
    {
        printf("%e ", gsl_vector_get(A,i));
    }
    printf("\n");
}

namespace ukf
{

    /**
      * @short MUST NOT BE USED !!!!!!! Square root UKF for state estimation, additive noise case
      * The notations follow "Sigma-Point Kalman Filters for Probabilistic Inference in Dynamic State-Space Models",p115, PhD, van Der Merwe
      */
    namespace srstate
    {
        /**
          * @short Allocation of the vectors/matrices and initialization
          *
          */
        void ukf_init(ukf_param &p, ukf_state &s)
        {
            // Parameters for the sigma points of the process equation
            p.nbSamples = 2 * p.n + 1;
            p.lambda = p.alpha * p.alpha * (p.n + p.kpa) - p.n;
            p.gamma = sqrt(p.n + p.lambda);

            // Parameters for the sigma points of the observation equation
            p.nbSamplesMeasure = 4 * p.n + 1;
            p.lambda_aug = p.alpha * p.alpha * (2*p.n + p.kpa) - 2*p.n;
            p.gamma_aug = sqrt(2*p.n + p.lambda_aug);

            // Init the matrices used to iterate
            s.xi = gsl_vector_alloc(p.n);
            gsl_vector_set_zero(s.xi);

            // Images of the sigma points
            s.xi_prediction = gsl_matrix_alloc(p.n, p.nbSamples);
            gsl_matrix_set_zero(s.xi_prediction);

            s.xi_mean = gsl_vector_alloc(p.n);
            gsl_vector_set_zero(s.xi_mean);

            // Cholesky factor of the variance/covariance matrix Pxxi
            s.Sxi = gsl_matrix_alloc(p.n,p.n);
            gsl_matrix_set_identity(s.Sxi);
            gsl_matrix_scale(s.Sxi, sqrt(p.prior_x));

            // Cholesky factor of the variance/covariance matrix of the process noise
            s.Svi = gsl_matrix_alloc(p.n,p.n);
            gsl_matrix_set_identity(s.Svi);
            gsl_matrix_scale(s.Svi, sqrt(p.process_noise));

            s.yi_prediction = gsl_matrix_alloc(p.no, p.nbSamplesMeasure);
            gsl_matrix_set_zero(s.yi_prediction);

            s.yi_mean = gsl_vector_alloc(p.no);
            gsl_vector_set_zero(s.yi_mean);

            s.ino_yi = gsl_vector_alloc(p.no);
            gsl_vector_set_zero(s.ino_yi);

            // Cholesky factor of the variance/covariance matrix Pyyi
            s.Syi = gsl_matrix_alloc(p.no, p.no);
            gsl_matrix_set_zero(s.Syi);

            s.Sni = gsl_matrix_alloc(p.no,p.no);
            gsl_matrix_set_identity(s.Sni);
            gsl_matrix_scale(s.Sni, sqrt(p.measurement_noise));

            s.Pxyi = gsl_matrix_alloc(p.n, p.no);
            gsl_matrix_set_zero(s.Pxyi);

            s.sigmaPoint = gsl_vector_alloc(p.n);
            gsl_vector_set_zero(s.sigmaPoint);

            s.sigmaPoints = gsl_matrix_alloc(p.n, p.nbSamples);
            gsl_matrix_set_zero(s.sigmaPoints);

            s.sigmaPointMeasure = gsl_vector_alloc(p.n);
            gsl_vector_set_zero(s.sigmaPoint);

            s.sigmaPointsMeasure = gsl_matrix_alloc(p.n, p.nbSamplesMeasure);
            gsl_matrix_set_zero(s.sigmaPointsMeasure);

            s.U = gsl_matrix_alloc(p.n, p.no);
            gsl_matrix_set_zero(s.U);

            // Weights used to update the statistics
            s.wm_j  = gsl_vector_alloc(p.nbSamples); // Weights used to compute the mean of the sigma points images
            s.wc_j = gsl_vector_alloc(p.nbSamples); // Weights used to update the covariance matrices

            // Set the weights
            gsl_vector_set(s.wm_j, 0, p.lambda / (p.n + p.lambda));
            gsl_vector_set(s.wc_j, 0, p.lambda / (p.n + p.lambda) + (1.0 - p.alpha*p.alpha + p.beta));
            for(int j = 1 ; j < p.nbSamples; j ++)
            {
                gsl_vector_set(s.wm_j, j, 1.0 / (2.0 * (p.n + p.lambda)));
                gsl_vector_set(s.wc_j, j, 1.0 / (2.0 * (p.n + p.lambda)));
            }

            // Set the weights
            s.wm_aug_j  = gsl_vector_alloc(p.nbSamplesMeasure); // Weights used to compute the mean of the sigma points images
            s.wc_aug_j = gsl_vector_alloc(p.nbSamplesMeasure); // Weights used to update the covariance matrices
            gsl_vector_set(s.wm_aug_j, 0, p.lambda_aug / (2*p.n + p.lambda_aug));
            gsl_vector_set(s.wc_aug_j, 0, p.lambda_aug / (2*p.n + p.lambda_aug) + (1.0 - p.alpha*p.alpha + p.beta));
            for(int j = 1 ; j < p.nbSamplesMeasure; j ++)
            {
                gsl_vector_set(s.wm_aug_j, j, 1.0 / (2.0 * (2*p.n + p.lambda_aug)));
                gsl_vector_set(s.wc_aug_j, j, 1.0 / (2.0 * (2*p.n + p.lambda_aug)));
            }

            s.Ki = gsl_matrix_alloc(p.n, p.no);
            s.Ki_T = gsl_matrix_alloc(p.no, p.n);

            // Allocate temporary matrices
            s.temp_n = gsl_vector_alloc(p.n);
            s.temp_no = gsl_vector_alloc(p.no);

            s.temp_n_1 = gsl_matrix_alloc(p.n,1);
            s.temp_1_n = gsl_matrix_alloc(1,p.n);
            s.temp_n_n = gsl_matrix_alloc(p.n, p.n);
            s.temp_n_no = gsl_matrix_alloc(p.n, p.no);
            s.temp_no_1 = gsl_matrix_alloc(p.no,1);
            s.temp_1_no = gsl_matrix_alloc(1,p.no);
            s.temp_no_no = gsl_matrix_alloc(p.no, p.no);

            s.temp_3n_n = gsl_matrix_alloc(3*p.n, p.n);
            s.temp_2nno_no = gsl_matrix_alloc(2*p.n + p.no, p.no);
        }

        /**
          * @short Free of memory allocation
          *
          */
        void ukf_free(ukf_param &p, ukf_state &s)
        {
            gsl_vector_free(s.xi);
            gsl_matrix_free(s.xi_prediction);
            gsl_vector_free(s.xi_mean);
            gsl_matrix_free(s.Sxi);
            gsl_matrix_free(s.Svi);

            gsl_matrix_free(s.yi_prediction);
            gsl_vector_free(s.yi_mean);
            gsl_vector_free(s.ino_yi);
            gsl_matrix_free(s.Syi);
            gsl_matrix_free(s.Sni);

            gsl_matrix_free(s.Pxyi);

            gsl_vector_free(s.sigmaPoint);
            gsl_matrix_free(s.sigmaPoints);

            gsl_vector_free(s.sigmaPointMeasure);
            gsl_matrix_free(s.sigmaPointsMeasure);

            gsl_matrix_free(s.U);

            gsl_vector_free(s.wm_j);
            gsl_vector_free(s.wc_j);

            gsl_vector_free(s.wm_aug_j);
            gsl_vector_free(s.wc_aug_j);

            gsl_matrix_free(s.Ki);
            gsl_matrix_free(s.Ki_T);

            gsl_vector_free(s.temp_n);
            gsl_vector_free(s.temp_no);

            gsl_matrix_free(s.temp_n_1);
            gsl_matrix_free(s.temp_1_n);
            gsl_matrix_free(s.temp_n_n);

            gsl_matrix_free(s.temp_n_no);
            gsl_matrix_free(s.temp_no_1);
            gsl_matrix_free(s.temp_1_no);
            gsl_matrix_free(s.temp_no_no);

            gsl_matrix_free(s.temp_3n_n);
            gsl_matrix_free(s.temp_2nno_no);
        }

        /**
          * @short UKF-additive (zero-mean) noise case, "Kalman Filtering and Neural Networks", p.233
          *
          */
        template <typename FunctProcess, typename FunctObservation>
                void ukf_iterate(ukf_param &p, ukf_state &s, FunctProcess f, FunctObservation h, gsl_vector* yi)
        {
            int i,j;
	    gsl_matrix_view mat_view;
	    gsl_vector_view vec_view;
            // ************************************************** //
            // ************ Compute the sigma points ************ //
            // ************************************************** //

            // 1- Compute the sigma points,
            // Equation (3.209)
            // sigmapoint_j = x_(i-1)
            // sigmapoint_j = x_(i-1) + gamma * Sxi_j for 1 <= j <= n
            // sigmapoint_j = x_(i-1) - gamma * Sxi_(j-(n+1)) for n+1 <= j <= 2n
            gsl_matrix_set_col(s.sigmaPoints, 0, s.xi);
            for(j = 1 ; j < p.n + 1 ; ++j)
                for(i = 0 ; i < p.n ; ++i)
                {
                gsl_matrix_set(s.sigmaPoints,i,j, s.xi->data[i] + p.gamma * gsl_matrix_get(s.Sxi, i, j-1));
                gsl_matrix_set(s.sigmaPoints,i,j+p.n, s.xi->data[i] - p.gamma * gsl_matrix_get(s.Sxi, i, j-1));
            }

            /**********************************/
            /***** Time update equations  *****/
            /**********************************/

            // Time update equations
            // 0 - Compute the image of the sigma points and the mean of these images
            gsl_vector_set_zero(s.xi_mean);
            for(j = 0 ; j < p.nbSamples ; ++j)
            {
                gsl_matrix_get_col(s.sigmaPoint, s.sigmaPoints, j);
                // Compute the image of the sigma points
                // Eq 3.210
		vec_view = gsl_matrix_column(s.xi_prediction,j);
                f(s.params, s.sigmaPoint, &vec_view.vector);

                // Update the mean, Eq (3.211)
                for(i = 0 ; i < p.n ; ++i)
                    s.xi_mean->data[i] += s.wm_j->data[j] * gsl_matrix_get(s.xi_prediction,i,j);
            }

            // 1- Compute the QR decomposition and keeps only the upper triangular part of R
            // Eq. 3.212
            // We first fill the matrix Sxk^-
            // with the transpose of what is in the brackets of qr { ... }
            for(j = 1 ; j < p.nbSamples; ++j)
            {
                for(i = 0 ; i < p.n ; ++i)
                {
                    gsl_matrix_set(s.temp_3n_n, j-1, i, sqrt(s.wc_j->data[1]) * (gsl_matrix_get(s.xi_prediction,i,j) - s.xi_mean->data[i]));
                }
            }
            // The last rows are filled with the cholesky factor of the process noise covariance
            for(j = 0 ; j < p.n ; ++j)
            {
	      mat_view = gsl_matrix_submatrix(s.temp_3n_n, p.nbSamples-1, 0, p.n, p.n);
                gsl_matrix_memcpy(&mat_view.matrix, s.Svi);
            }
            // We now perform the QR decomposition of s.temp_3n_n;
            gsl_linalg_QR_decomp(s.temp_3n_n, s.temp_n);
            // From this QR decomposition, we keep only the upper triangular part of R and copy it in s.Sxi
            for( j = 0 ; j < p.n ; ++j)
            {
                for(i = 0 ; i < j+1 ; ++i)
                    gsl_matrix_set(s.Sxi, i, j, gsl_matrix_get(s.temp_3n_n, i, j));

                for(i = j+1 ; i < p.n ; ++i)
                    gsl_matrix_set(s.Sxi, i, j, 0.0);
            }
            //gsl_matrix_transpose(s.Sxi);
            if(DEBUGMODE) print_mat(s.Sxi);
            if(DEBUGMODE) printf("\n");

            // 2- Perform the cholupdate, which is a rank one update
            // Eq. 3.213
            for(i = 0 ; i < p.n ; ++i)
                s.temp_n->data[i] = gsl_matrix_get(s.xi_prediction,i,0) - s.xi_mean->data[i];

            if(DEBUGMODE) printf("Before cholupdate 1 \n");
            gsl_matrix_transpose(s.Sxi);
            ukf::math::choleskyUpdate(s.Sxi, ukf::math::signof(s.wc_j->data[0]) * sqrt(fabs(s.wc_j->data[0])), s.temp_n);
            // Warning : Here, we keep Sxi as a triangular inferior matrix

            //gsl_matrix_transpose(s.Sxi);
            if(DEBUGMODE) printf("ok\n");
            if(DEBUGMODE) printf("Sxi after cholupdate : \n");
            if(DEBUGMODE) print_mat(s.Sxi);
            // ** Augment the sigma points ** //
            // Eq 3.124

            // 1 - Copy the images of the original sigma points
	    mat_view = gsl_matrix_submatrix(s.sigmaPointsMeasure,0,0, p.n, p.nbSamples);
            gsl_matrix_memcpy(&mat_view.matrix, s.xi_prediction);

            // 2 - Compute the additional sigma points
            for(j = 0 ; j < p.n ; ++j)
            {
                for(i = 0 ; i < p.n ; ++i)
                {
                    gsl_matrix_set(s.sigmaPointsMeasure, i ,j + p.nbSamples , gsl_matrix_get(s.xi_prediction, i, 0) + p.gamma_aug * gsl_matrix_get(s.Svi, i,j));
                    gsl_matrix_set(s.sigmaPointsMeasure, i ,j + p.nbSamples + p.n , gsl_matrix_get(s.xi_prediction, i, 0) - p.gamma_aug * gsl_matrix_get(s.Svi, i,j));
                }
            }

            // Compute the image of the sigma points, i.e. the associated predicted observations
            // and their mean
            gsl_vector_set_zero(s.yi_mean);
            for(j = 0 ; j < p.nbSamplesMeasure ; ++j)
            {
                gsl_matrix_get_col(s.sigmaPointMeasure, s.sigmaPointsMeasure, j);
		vec_view = gsl_matrix_column(s.yi_prediction,j);
                h(s.sigmaPointMeasure, &vec_view.vector);

                // Update the mean , eq (3.176)
                for(i = 0 ; i < p.no ; ++i)
                    s.yi_mean->data[i] += s.wm_aug_j->data[j] * gsl_matrix_get(s.yi_prediction,i,j);
            }

            /*****************************************/
            /***** Measurement update equations  *****/
            /*****************************************/

            // Compute the Cholesky factor of the variance-covariance matrix of the observations
            // Eq. 3.217
            // First put in s.temp_2nno_no the transpose of what is in the brackets of qr{...} ,
            for(j = 0 ; j < 2*p.n ; ++j)
                for(i = 0 ; i < p.no ; ++i)
                    gsl_matrix_set(s.temp_2nno_no, j, i, sqrt(s.wc_aug_j->data[1])*(gsl_matrix_get(s.yi_prediction, i, j+1) - s.yi_mean->data[i]) );

            // Copy the Cholesky factor of the covariance of the observation noise
	    mat_view = gsl_matrix_submatrix(s.temp_2nno_no,2*p.n, 0, p.no, p.no);
            gsl_matrix_memcpy(&mat_view.matrix, s.Sni);

            // We now perform the QR decomposition of s.temp_2nno_no;
            gsl_linalg_QR_decomp(s.temp_2nno_no, s.temp_no);
            // From this QR decomposition, we keep only the upper triangular part of R and copy it in s.Syi
            for( j = 0 ; j < p.no ; ++j)
            {
                for(i = 0 ; i < j+1 ; ++i)
                    gsl_matrix_set(s.Syi, i, j, gsl_matrix_get(s.temp_2nno_no, i, j));

                for(i = j+1 ; i < p.no ; ++i)
                    gsl_matrix_set(s.Syi, i, j, 0.0);
            }
            if(DEBUGMODE) printf("Syi : \n");
            if(DEBUGMODE) print_mat(s.Syi);
            if(DEBUGMODE) printf("\n");
            //gsl_matrix_transpose(s.Syi);

            // 2- Perform the cholupdate, which is a rank one update
            // Eq. 3.218
            // Syi is triangular superior
            for(i = 0 ; i < p.no ; ++i)
                s.temp_no->data[i] = gsl_matrix_get(s.yi_prediction,i,0) - s.yi_mean->data[i];
            if(DEBUGMODE) print_vec(s.temp_no);
            if(DEBUGMODE) printf("Before cholupdate 2 \n");
            gsl_matrix_transpose(s.Syi);
            ukf::math::choleskyUpdate(s.Syi, ukf::math::signof(s.wc_aug_j->data[0])*sqrt(fabs(s.wc_aug_j->data[0])), s.temp_no);
            if(DEBUGMODE) printf("ok \n");
            // Now, Syi is the lower triangular cholesky factor
            //gsl_matrix_transpose(s.Syi);
            if(DEBUGMODE) printf("Syi after cholupdate : \n");
            if(DEBUGMODE) print_mat(s.Syi);

            // Compute the state/observation covariance
            // Eq (3.219)
            gsl_matrix_set_zero(s.Pxyi);
            for(j = 0 ; j < p.nbSamplesMeasure ; ++j)
            {
                for(i = 0 ; i < p.n ; ++i)
                    s.temp_n_1->data[i] = gsl_matrix_get(s.sigmaPointsMeasure,i,j) - s.xi_mean->data[i];

                for(i = 0 ; i < p.no ; ++i)
                    s.temp_1_no->data[i] = gsl_matrix_get(s.yi_prediction,i,j) - s.yi_mean->data[i];

                gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, s.wc_aug_j->data[j] , s.temp_n_1, s.temp_1_no, 1.0, s.Pxyi);
                //gsl_matrix_add(s.Pxyi, s.temp_n_no);
            }

            /*gsl_matrix_set_zero(s.Pxyi);
            for(j = 0 ; j < p.nbSamples ; ++j)
            {
                for(i = 0 ; i < p.n ; ++i)
                    s.temp_n_1->data[i] = gsl_matrix_get(s.xi_prediction,i,j) - s.xi_mean->data[i];

                for(i = 0 ; i < p.no ; ++i)
                    s.temp_1_no->data[i] = gsl_matrix_get(s.yi_prediction,i,j) - s.yi_mean->data[i];

                gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, s.wc_j->data[j] , s.temp_n_1, s.temp_1_no, 1.0, s.Pxyi);
                //gsl_matrix_add(s.Pxyi, s.temp_n_no);
            }*/

            // Compute the Kalman gain,
            // Eq (3.220)

            // We want to solve : Kxy . Sy Sy^T = Pxy
            // Which can be solved using : Sy Sy^T Kxy^T = Pxy^T
            // We have a number of no systems to solve of the form S x = Pxy^T[:,i] ; S^T Kxy[:,i] = x with S being triangular
            // To solve this, we make use of gsl_blas_dtrsv
            for(int i = 0 ; i < p.n ; ++i)
            {
	      vec_view = gsl_matrix_row(s.Pxyi, i);
                gsl_vector_memcpy(s.temp_no, &vec_view.vector);
                gsl_blas_dtrsv(CblasLower,CblasNoTrans,CblasNonUnit,s.Syi,s.temp_no);
                gsl_blas_dtrsv(CblasLower,CblasTrans,CblasNonUnit,s.Syi,s.temp_no);
                // And then copy the result in the ith row of Ki
		vec_view = gsl_matrix_row(s.Ki, i);
                gsl_vector_memcpy(&vec_view.vector, s.temp_no);
            }
            if(DEBUGMODE) printf("Ki : \n");
            if(DEBUGMODE) print_mat(s.Ki);

            // Correction : correct the estimation of the state
            // Eq. 3.221
            // Compute the innovations
            for(i = 0 ; i < p.no ; ++i)
                s.ino_yi->data[i] = yi->data[i] - s.yi_mean->data[i];
            // And correct the current state estimate
            gsl_blas_dgemv(CblasNoTrans, 1.0, s.Ki, s.ino_yi, 1.0, s.xi);

            // Compute the matrix U = Kk Syk
            gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0 , s.Ki, s.Syi, 0, s.U);
            if(DEBUGMODE) printf("Matrix U : \n");
            if(DEBUGMODE) print_mat(s.U);

            // Perform the cholesky updates of Sxi
            // We perform p.no updates rank-1 update with the columns of U
            if(DEBUGMODE) printf("Before cholupdate 3 \n");
            for(i = 0 ; i < p.no ; ++i)
            {
	      vec_view = gsl_matrix_column(s.U, i);
                gsl_vector_memcpy(s.temp_n, &vec_view.vector);
                ukf::math::choleskyUpdate(s.Sxi, -1.0, s.temp_n);
                if(DEBUGMODE) printf("After cholupdate %i\n", i);
                if(DEBUGMODE) print_mat(s.Sxi);
            }

            //printf("After cholupdate %i\n", i);
            //print_mat(s.Sxi);
            //if(DEBUGMODE) printf("\n");

            //
            //            // Update of the process noise
            //            switch(p.process_noise_type)
            //            {
            //            case ukf::UKF_PROCESS_FIXED:
            //                //nothing to do
            //                break;
            //            case ukf::UKF_PROCESS_RLS:
            //                gsl_matrix_memcpy(s.Pvvi, s.Pxxi);
            //                gsl_matrix_scale(s.Pvvi, 1.0/p.process_noise-1.0);
            //                gsl_matrix_memcpy(s.cholPvvi, s.Pvvi);
            //                gsl_linalg_cholesky_decomp(s.cholPvvi);
            //                for(j = 0 ; j < p.n ; j++)
            //                    for(k = j+1 ; k < p.n ; k++)
            //                        gsl_matrix_set(s.cholPvvi,j,k,0.0);
            //                break;
            //            default:
            //                printf("Warning : Unrecognized process noise type\n");
            //            }

        }

        /**
          * @short Evaluation of the output from the sigma points
          *
          */
        template <typename FunctProcess, typename FunctObservation>
                void ukf_evaluate(ukf_param &p, ukf_state &s, FunctProcess f, FunctObservation h, gsl_vector* yi)
        {

            //          int i,j,k;
            //            // ************************************************** //
            //            // ************ Compute the sigma points ************ //
            //            // ************************************************** //
            //
            //            // 0 - Compute the Cholesky decomposition of s.Pxxi
            //            gsl_matrix_memcpy(s.cholPxxi, s.Pxxi);
            //            gsl_linalg_cholesky_decomp(s.cholPxxi);
            //            // Set all the elements of cholPvvi strictly above the diagonal to zero
            //            for(j = 0 ; j < p.n ; j++)
            //                for(k = j+1 ; k < p.n ; k++)
            //                    gsl_matrix_set(s.cholPxxi,j,k,0.0);
            //
            //            // 1- Compute the sigma points,
            //            // Equation (3.170)
            //            // sigmapoint_j = x_(i-1)
            //            // sigmapoint_j = x_(i-1) + gamma * sqrt(P_i-1)_j for 1 <= j <= n
            //            // sigmapoint_j = x_(i-1) - gamma * sqrt(P_i-1)_(j-(n+1)) for n+1 <= j <= 2n
            //            gsl_matrix_set_col(s.sigmaPoints, 0, s.xi);
            //            for(j = 1 ; j < p.n+1 ; ++j)
            //                for(i = 0 ; i < p.n ; ++i)
            //                {
            //                    gsl_matrix_set(s.sigmaPoints,i,j, s.xi->data[i] + p.gamma * gsl_matrix_get(s.cholPxxi, i, j-1));
            //                    gsl_matrix_set(s.sigmaPoints,i,j+p.n, s.xi->data[i] - p.gamma * gsl_matrix_get(s.cholPxxi, i, j-1));
            //                }
            //
            //            /**********************************/
            //            /***** Time update equations  *****/
            //            /**********************************/
            //
            //            // Time update equations
            //            // 0 - Compute the image of the sigma points and the mean of these images
            //            gsl_vector_set_zero(s.xi_mean);
            //            for(j = 0 ; j < p.nbSamples ; ++j)
            //            {
            //                gsl_matrix_get_col(s.sigmaPoint, s.sigmaPoints, j);
            //                f(s.params, s.sigmaPoint, &gsl_matrix_column(s.xi_prediction,j).vector);
            //
            //                // Update the mean, Eq (3.172)
            //                for(i = 0 ; i < p.n ; ++i)
            //                    s.xi_mean->data[i] += s.wm_j->data[j] * gsl_matrix_get(s.xi_prediction,i,j);
            //            }
            //
            //            // 1 - Compute the covariance of the images and add the process noise,
            //            // Equation (3.173)
            //            // Warning, s.Pxxi will now hold P_xk^-
            //            gsl_matrix_set_zero(s.Pxxi);
            //            for(j = 0 ; j < p.nbSamples ; ++j)
            //            {
            //                for(i = 0 ; i < p.n ; ++i)
            //                    s.temp_n_1->data[i] = gsl_matrix_get(s.xi_prediction,i,j) - s.xi_mean->data[i];
            //
            //                gsl_blas_dgemm(CblasNoTrans, CblasTrans, s.wc_j->data[j] , s.temp_n_1, s.temp_n_1, 0, s.temp_n_n);
            //                gsl_matrix_add(s.Pxxi, s.temp_n_n);
            //            }
            //            // Add the covariance of the evolution noise
            //            gsl_matrix_add(s.Pxxi, s.Pvvi);
            //
            //            // Augment sigma points
            //            // Equation 3.174
            //            // First put the images of the initial sigma points
            //            gsl_matrix_memcpy(&gsl_matrix_submatrix(s.sigmaPointsMeasure, 0, 0, p.n, p.nbSamples).matrix, s.xi_prediction);
            //            // And add the additional sigma points eq. (7.56)
            //            for(j = 0 ; j < p.n ; ++j)
            //            {
            //                for(i = 0 ; i < p.n ; ++i)
            //                {
            //                    gsl_matrix_set(s.sigmaPointsMeasure, i, j+p.nbSamples, gsl_matrix_get(s.xi_prediction,i,0)+p.gamma_aug*gsl_matrix_get(s.cholPvvi,i,j));
            //                    gsl_matrix_set(s.sigmaPointsMeasure, i, j+p.nbSamples+p.n, gsl_matrix_get(s.xi_prediction,i,0)-p.gamma_aug*gsl_matrix_get(s.cholPvvi,i,j));
            //                }
            //            }
            //
            //            // Compute the image of the sigma points through the observation equation
            //            // eq (3.175)
            //            gsl_vector_set_zero(yi);
            //            for(j = 0 ; j < p.nbSamplesMeasure ; ++j)
            //            {
            //                gsl_matrix_get_col(s.sigmaPointMeasure, s.sigmaPointsMeasure, j);
            //                h(s.sigmaPointMeasure, &gsl_matrix_column(s.yi_prediction,j).vector);
            //
            //                // Update the mean , eq (3.176)
            //                for(i = 0 ; i < p.no ; ++i)
            //                    yi->data[i] += s.wm_aug_j->data[j] * gsl_matrix_get(s.yi_prediction,i,j);
            //            }
        }
    } // srstate
} // ukf

#endif // UKF_SR_STATE_NDIM_H


/*
 * 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:43:59 EST 1999 */

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

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

/*
 * This function contains 156 FP additions, 42 FP multiplications,
 * (or, 140 additions, 26 multiplications, 16 fused multiply/add),
 * 46 stack variables, and 64 memory accesses
 */
static const fftw_real K195090322 = FFTW_KONST(+0.195090322016128267848284868477022240927691618);
static const fftw_real K980785280 = FFTW_KONST(+0.980785280403230449126182236134239036973933731);
static const fftw_real K555570233 = FFTW_KONST(+0.555570233019602224742830813948532874374937191);
static const fftw_real K831469612 = FFTW_KONST(+0.831469612302545237078788377617905756738560812);
static const fftw_real K382683432 = FFTW_KONST(+0.382683432365089771728459984030398866761344562);
static const fftw_real K923879532 = FFTW_KONST(+0.923879532511286756128183189396788286822416626);
static const fftw_real K707106781 = FFTW_KONST(+0.707106781186547524400844362104849039284835938);

/*
 * 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_real2hc_32(const fftw_real *input, fftw_real *real_output, fftw_real *imag_output, int istride, int real_ostride, int imag_ostride)
{
     fftw_real tmp7;
     fftw_real tmp135;
     fftw_real tmp31;
     fftw_real tmp83;
     fftw_real tmp14;
     fftw_real tmp148;
     fftw_real tmp34;
     fftw_real tmp82;
     fftw_real tmp123;
     fftw_real tmp143;
     fftw_real tmp71;
     fftw_real tmp99;
     fftw_real tmp126;
     fftw_real tmp144;
     fftw_real tmp76;
     fftw_real tmp100;
     fftw_real tmp22;
     fftw_real tmp136;
     fftw_real tmp38;
     fftw_real tmp80;
     fftw_real tmp29;
     fftw_real tmp137;
     fftw_real tmp41;
     fftw_real tmp79;
     fftw_real tmp116;
     fftw_real tmp140;
     fftw_real tmp54;
     fftw_real tmp96;
     fftw_real tmp119;
     fftw_real tmp141;
     fftw_real tmp59;
     fftw_real tmp97;
     ASSERT_ALIGNED_DOUBLE;
     {
	  fftw_real tmp1;
	  fftw_real tmp2;
	  fftw_real tmp3;
	  fftw_real tmp4;
	  fftw_real tmp5;
	  fftw_real tmp6;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp1 = input[0];
	  tmp2 = input[16 * istride];
	  tmp3 = tmp1 + tmp2;
	  tmp4 = input[8 * istride];
	  tmp5 = input[24 * istride];
	  tmp6 = tmp4 + tmp5;
	  tmp7 = tmp3 + tmp6;
	  tmp135 = tmp3 - tmp6;
	  tmp31 = tmp1 - tmp2;
	  tmp83 = tmp4 - tmp5;
     }
     {
	  fftw_real tmp10;
	  fftw_real tmp32;
	  fftw_real tmp13;
	  fftw_real tmp33;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp8;
	       fftw_real tmp9;
	       fftw_real tmp11;
	       fftw_real tmp12;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp8 = input[4 * istride];
	       tmp9 = input[20 * istride];
	       tmp10 = tmp8 + tmp9;
	       tmp32 = tmp8 - tmp9;
	       tmp11 = input[28 * istride];
	       tmp12 = input[12 * istride];
	       tmp13 = tmp11 + tmp12;
	       tmp33 = tmp11 - tmp12;
	  }
	  tmp14 = tmp10 + tmp13;
	  tmp148 = tmp13 - tmp10;
	  tmp34 = K707106781 * (tmp32 + tmp33);
	  tmp82 = K707106781 * (tmp33 - tmp32);
     }
     {
	  fftw_real tmp63;
	  fftw_real tmp121;
	  fftw_real tmp75;
	  fftw_real tmp122;
	  fftw_real tmp66;
	  fftw_real tmp124;
	  fftw_real tmp69;
	  fftw_real tmp125;
	  fftw_real tmp70;
	  fftw_real tmp72;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp61;
	       fftw_real tmp62;
	       fftw_real tmp73;
	       fftw_real tmp74;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp61 = input[31 * istride];
	       tmp62 = input[15 * istride];
	       tmp63 = tmp61 - tmp62;
	       tmp121 = tmp61 + tmp62;
	       tmp73 = input[7 * istride];
	       tmp74 = input[23 * istride];
	       tmp75 = tmp73 - tmp74;
	       tmp122 = tmp73 + tmp74;
	  }
	  {
	       fftw_real tmp64;
	       fftw_real tmp65;
	       fftw_real tmp67;
	       fftw_real tmp68;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp64 = input[3 * istride];
	       tmp65 = input[19 * istride];
	       tmp66 = tmp64 - tmp65;
	       tmp124 = tmp64 + tmp65;
	       tmp67 = input[27 * istride];
	       tmp68 = input[11 * istride];
	       tmp69 = tmp67 - tmp68;
	       tmp125 = tmp67 + tmp68;
	  }
	  tmp123 = tmp121 + tmp122;
	  tmp143 = tmp121 - tmp122;
	  tmp70 = K707106781 * (tmp66 + tmp69);
	  tmp71 = tmp63 + tmp70;
	  tmp99 = tmp63 - tmp70;
	  tmp126 = tmp124 + tmp125;
	  tmp144 = tmp125 - tmp124;
	  tmp72 = K707106781 * (tmp69 - tmp66);
	  tmp76 = tmp72 - tmp75;
	  tmp100 = tmp75 + tmp72;
     }
     {
	  fftw_real tmp18;
	  fftw_real tmp36;
	  fftw_real tmp21;
	  fftw_real tmp37;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp16;
	       fftw_real tmp17;
	       fftw_real tmp19;
	       fftw_real tmp20;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp16 = input[2 * istride];
	       tmp17 = input[18 * istride];
	       tmp18 = tmp16 + tmp17;
	       tmp36 = tmp16 - tmp17;
	       tmp19 = input[10 * istride];
	       tmp20 = input[26 * istride];
	       tmp21 = tmp19 + tmp20;
	       tmp37 = tmp19 - tmp20;
	  }
	  tmp22 = tmp18 + tmp21;
	  tmp136 = tmp18 - tmp21;
	  tmp38 = (K923879532 * tmp36) - (K382683432 * tmp37);
	  tmp80 = (K382683432 * tmp36) + (K923879532 * tmp37);
     }
     {
	  fftw_real tmp25;
	  fftw_real tmp39;
	  fftw_real tmp28;
	  fftw_real tmp40;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp23;
	       fftw_real tmp24;
	       fftw_real tmp26;
	       fftw_real tmp27;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp23 = input[30 * istride];
	       tmp24 = input[14 * istride];
	       tmp25 = tmp23 + tmp24;
	       tmp39 = tmp23 - tmp24;
	       tmp26 = input[6 * istride];
	       tmp27 = input[22 * istride];
	       tmp28 = tmp26 + tmp27;
	       tmp40 = tmp26 - tmp27;
	  }
	  tmp29 = tmp25 + tmp28;
	  tmp137 = tmp25 - tmp28;
	  tmp41 = (K923879532 * tmp39) + (K382683432 * tmp40);
	  tmp79 = (K382683432 * tmp39) - (K923879532 * tmp40);
     }
     {
	  fftw_real tmp46;
	  fftw_real tmp114;
	  fftw_real tmp58;
	  fftw_real tmp115;
	  fftw_real tmp49;
	  fftw_real tmp117;
	  fftw_real tmp52;
	  fftw_real tmp118;
	  fftw_real tmp53;
	  fftw_real tmp55;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp44;
	       fftw_real tmp45;
	       fftw_real tmp56;
	       fftw_real tmp57;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp44 = input[istride];
	       tmp45 = input[17 * istride];
	       tmp46 = tmp44 - tmp45;
	       tmp114 = tmp44 + tmp45;
	       tmp56 = input[9 * istride];
	       tmp57 = input[25 * istride];
	       tmp58 = tmp56 - tmp57;
	       tmp115 = tmp56 + tmp57;
	  }
	  {
	       fftw_real tmp47;
	       fftw_real tmp48;
	       fftw_real tmp50;
	       fftw_real tmp51;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp47 = input[5 * istride];
	       tmp48 = input[21 * istride];
	       tmp49 = tmp47 - tmp48;
	       tmp117 = tmp47 + tmp48;
	       tmp50 = input[29 * istride];
	       tmp51 = input[13 * istride];
	       tmp52 = tmp50 - tmp51;
	       tmp118 = tmp50 + tmp51;
	  }
	  tmp116 = tmp114 + tmp115;
	  tmp140 = tmp114 - tmp115;
	  tmp53 = K707106781 * (tmp49 + tmp52);
	  tmp54 = tmp46 + tmp53;
	  tmp96 = tmp46 - tmp53;
	  tmp119 = tmp117 + tmp118;
	  tmp141 = tmp118 - tmp117;
	  tmp55 = K707106781 * (tmp52 - tmp49);
	  tmp59 = tmp55 - tmp58;
	  tmp97 = tmp58 + tmp55;
     }
     {
	  fftw_real tmp15;
	  fftw_real tmp30;
	  fftw_real tmp131;
	  fftw_real tmp132;
	  fftw_real tmp133;
	  fftw_real tmp134;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp15 = tmp7 + tmp14;
	  tmp30 = tmp22 + tmp29;
	  tmp131 = tmp15 + tmp30;
	  tmp132 = tmp116 + tmp119;
	  tmp133 = tmp123 + tmp126;
	  tmp134 = tmp132 + tmp133;
	  real_output[8 * real_ostride] = tmp15 - tmp30;
	  imag_output[8 * imag_ostride] = tmp133 - tmp132;
	  real_output[16 * real_ostride] = tmp131 - tmp134;
	  real_output[0] = tmp131 + tmp134;
     }
     {
	  fftw_real tmp113;
	  fftw_real tmp129;
	  fftw_real tmp128;
	  fftw_real tmp130;
	  fftw_real tmp120;
	  fftw_real tmp127;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp113 = tmp7 - tmp14;
	  tmp129 = tmp29 - tmp22;
	  tmp120 = tmp116 - tmp119;
	  tmp127 = tmp123 - tmp126;
	  tmp128 = K707106781 * (tmp120 + tmp127);
	  tmp130 = K707106781 * (tmp127 - tmp120);
	  real_output[12 * real_ostride] = tmp113 - tmp128;
	  real_output[4 * real_ostride] = tmp113 + tmp128;
	  imag_output[4 * imag_ostride] = tmp129 + tmp130;
	  imag_output[12 * imag_ostride] = tmp130 - tmp129;
     }
     {
	  fftw_real tmp139;
	  fftw_real tmp155;
	  fftw_real tmp149;
	  fftw_real tmp151;
	  fftw_real tmp146;
	  fftw_real tmp150;
	  fftw_real tmp154;
	  fftw_real tmp156;
	  fftw_real tmp138;
	  fftw_real tmp147;
	  ASSERT_ALIGNED_DOUBLE;
	  tmp138 = K707106781 * (tmp136 + tmp137);
	  tmp139 = tmp135 + tmp138;
	  tmp155 = tmp135 - tmp138;
	  tmp147 = K707106781 * (tmp137 - tmp136);
	  tmp149 = tmp147 - tmp148;
	  tmp151 = tmp148 + tmp147;
	  {
	       fftw_real tmp142;
	       fftw_real tmp145;
	       fftw_real tmp152;
	       fftw_real tmp153;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp142 = (K923879532 * tmp140) + (K382683432 * tmp141);
	       tmp145 = (K923879532 * tmp143) - (K382683432 * tmp144);
	       tmp146 = tmp142 + tmp145;
	       tmp150 = tmp145 - tmp142;
	       tmp152 = (K923879532 * tmp141) - (K382683432 * tmp140);
	       tmp153 = (K382683432 * tmp143) + (K923879532 * tmp144);
	       tmp154 = tmp152 + tmp153;
	       tmp156 = tmp153 - tmp152;
	  }
	  real_output[14 * real_ostride] = tmp139 - tmp146;
	  real_output[2 * real_ostride] = tmp139 + tmp146;
	  imag_output[6 * imag_ostride] = tmp149 + tmp150;
	  imag_output[10 * imag_ostride] = tmp150 - tmp149;
	  imag_output[2 * imag_ostride] = tmp151 + tmp154;
	  imag_output[14 * imag_ostride] = tmp154 - tmp151;
	  real_output[10 * real_ostride] = tmp155 - tmp156;
	  real_output[6 * real_ostride] = tmp155 + tmp156;
     }
     {
	  fftw_real tmp95;
	  fftw_real tmp111;
	  fftw_real tmp110;
	  fftw_real tmp112;
	  fftw_real tmp102;
	  fftw_real tmp106;
	  fftw_real tmp105;
	  fftw_real tmp107;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp93;
	       fftw_real tmp94;
	       fftw_real tmp108;
	       fftw_real tmp109;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp93 = tmp31 - tmp34;
	       tmp94 = tmp80 + tmp79;
	       tmp95 = tmp93 + tmp94;
	       tmp111 = tmp93 - tmp94;
	       tmp108 = (K831469612 * tmp97) - (K555570233 * tmp96);
	       tmp109 = (K555570233 * tmp99) + (K831469612 * tmp100);
	       tmp110 = tmp108 + tmp109;
	       tmp112 = tmp109 - tmp108;
	  }
	  {
	       fftw_real tmp98;
	       fftw_real tmp101;
	       fftw_real tmp103;
	       fftw_real tmp104;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp98 = (K831469612 * tmp96) + (K555570233 * tmp97);
	       tmp101 = (K831469612 * tmp99) - (K555570233 * tmp100);
	       tmp102 = tmp98 + tmp101;
	       tmp106 = tmp101 - tmp98;
	       tmp103 = tmp41 - tmp38;
	       tmp104 = tmp83 + tmp82;
	       tmp105 = tmp103 - tmp104;
	       tmp107 = tmp104 + tmp103;
	  }
	  real_output[13 * real_ostride] = tmp95 - tmp102;
	  real_output[3 * real_ostride] = tmp95 + tmp102;
	  imag_output[5 * imag_ostride] = tmp105 + tmp106;
	  imag_output[11 * imag_ostride] = tmp106 - tmp105;
	  imag_output[3 * imag_ostride] = tmp107 + tmp110;
	  imag_output[13 * imag_ostride] = tmp110 - tmp107;
	  real_output[11 * real_ostride] = tmp111 - tmp112;
	  real_output[5 * real_ostride] = tmp111 + tmp112;
     }
     {
	  fftw_real tmp43;
	  fftw_real tmp91;
	  fftw_real tmp90;
	  fftw_real tmp92;
	  fftw_real tmp78;
	  fftw_real tmp86;
	  fftw_real tmp85;
	  fftw_real tmp87;
	  ASSERT_ALIGNED_DOUBLE;
	  {
	       fftw_real tmp35;
	       fftw_real tmp42;
	       fftw_real tmp88;
	       fftw_real tmp89;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp35 = tmp31 + tmp34;
	       tmp42 = tmp38 + tmp41;
	       tmp43 = tmp35 + tmp42;
	       tmp91 = tmp35 - tmp42;
	       tmp88 = (K980785280 * tmp59) - (K195090322 * tmp54);
	       tmp89 = (K195090322 * tmp71) + (K980785280 * tmp76);
	       tmp90 = tmp88 + tmp89;
	       tmp92 = tmp89 - tmp88;
	  }
	  {
	       fftw_real tmp60;
	       fftw_real tmp77;
	       fftw_real tmp81;
	       fftw_real tmp84;
	       ASSERT_ALIGNED_DOUBLE;
	       tmp60 = (K980785280 * tmp54) + (K195090322 * tmp59);
	       tmp77 = (K980785280 * tmp71) - (K195090322 * tmp76);
	       tmp78 = tmp60 + tmp77;
	       tmp86 = tmp77 - tmp60;
	       tmp81 = tmp79 - tmp80;
	       tmp84 = tmp82 - tmp83;
	       tmp85 = tmp81 - tmp84;
	       tmp87 = tmp84 + tmp81;
	  }
	  real_output[15 * real_ostride] = tmp43 - tmp78;
	  real_output[real_ostride] = tmp43 + tmp78;
	  imag_output[7 * imag_ostride] = tmp85 + tmp86;
	  imag_output[9 * imag_ostride] = tmp86 - tmp85;
	  imag_output[imag_ostride] = tmp87 + tmp90;
	  imag_output[15 * imag_ostride] = tmp90 - tmp87;
	  real_output[9 * real_ostride] = tmp91 - tmp92;
	  real_output[7 * real_ostride] = tmp91 + tmp92;
     }
}

fftw_codelet_desc fftw_real2hc_32_desc =
{
     "fftw_real2hc_32",
     (void (*)()) fftw_real2hc_32,
     32,
     FFTW_FORWARD,
     FFTW_REAL2HC,
     706,
     0,
     (const int *) 0,
};


/* specfunc/log.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Author:  G. Jungman */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>
#include "gsl_sf_log.h"

#include "error.h"

#include "chebyshev.h"
#include "cheb_eval.c"

/*-*-*-*-*-*-*-*-*-*-*-* Private Section *-*-*-*-*-*-*-*-*-*-*-*/

/* Chebyshev expansion for log(1 + x(t))/x(t)
 *
 * x(t) = (4t-1)/(2(4-t))
 * t(x) = (8x+1)/(2(x+2))
 * -1/2 < x < 1/2
 * -1 < t < 1
 */
static double lopx_data[21] = {
  2.16647910664395270521272590407,
 -0.28565398551049742084877469679,
  0.01517767255690553732382488171,
 -0.00200215904941415466274422081,
  0.00019211375164056698287947962,
 -0.00002553258886105542567601400,
  2.9004512660400621301999384544e-06,
 -3.8873813517057343800270917900e-07,
  4.7743678729400456026672697926e-08,
 -6.4501969776090319441714445454e-09,
  8.2751976628812389601561347296e-10,
 -1.1260499376492049411710290413e-10,
  1.4844576692270934446023686322e-11,
 -2.0328515972462118942821556033e-12,
  2.7291231220549214896095654769e-13,
 -3.7581977830387938294437434651e-14,
  5.1107345870861673561462339876e-15,
 -7.0722150011433276578323272272e-16,
  9.7089758328248469219003866867e-17,
 -1.3492637457521938883731579510e-17,
  1.8657327910677296608121390705e-18
};
static cheb_series lopx_cs = {
  lopx_data,
  20,
  -1, 1,
  10
};

/* Chebyshev expansion for (log(1 + x(t)) - x(t))/x(t)^2
 *
 * x(t) = (4t-1)/(2(4-t))
 * t(x) = (8x+1)/(2(x+2))
 * -1/2 < x < 1/2
 * -1 < t < 1
 */
static double lopxmx_data[20] = {
 -1.12100231323744103373737274541,
  0.19553462773379386241549597019,
 -0.01467470453808083971825344956,
  0.00166678250474365477643629067,
 -0.00018543356147700369785746902,
  0.00002280154021771635036301071,
 -2.8031253116633521699214134172e-06,
  3.5936568872522162983669541401e-07,
 -4.6241857041062060284381167925e-08,
  6.0822637459403991012451054971e-09,
 -8.0339824424815790302621320732e-10,
  1.0751718277499375044851551587e-10,
 -1.4445310914224613448759230882e-11,
  1.9573912180610336168921438426e-12,
 -2.6614436796793061741564104510e-13,
  3.6402634315269586532158344584e-14,
 -4.9937495922755006545809120531e-15,
  6.8802890218846809524646902703e-16,
 -9.5034129794804273611403251480e-17,
  1.3170135013050997157326965813e-17
};
static cheb_series lopxmx_cs = {
  lopxmx_data,
  19,
  -1, 1,
  9
};

/*-*-*-*-*-*-*-*-*-*-*-* Functions with Error Codes *-*-*-*-*-*-*-*-*-*-*-*/

#ifndef HIDE_INLINE_STATIC
int
gsl_sf_log_e(const double x, gsl_sf_result * result)
{
  /* CHECK_POINTER(result) */

  if(x <= 0.0) {
    DOMAIN_ERROR(result);
  }
  else {
    result->val = log(x);
    result->err = 2.0 * GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
}

int
gsl_sf_log_abs_e(const double x, gsl_sf_result * result)
{
  /* CHECK_POINTER(result) */

  if(x == 0.0) {
    DOMAIN_ERROR(result);
  }
  else {
    result->val = log(fabs(x));
    result->err = 2.0 * GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
}
#endif

int
gsl_sf_complex_log_e(const double zr, const double zi, gsl_sf_result * lnr, gsl_sf_result * theta)
{
  /* CHECK_POINTER(lnr) */
  /* CHECK_POINTER(theta) */

  if(zr != 0.0 || zi != 0.0) {
    const double ax = fabs(zr);
    const double ay = fabs(zi);
    const double min = GSL_MIN(ax, ay);
    const double max = GSL_MAX(ax, ay);
    lnr->val = log(max) + 0.5 * log(1.0 + (min/max)*(min/max));
    lnr->err = 2.0 * GSL_DBL_EPSILON * fabs(lnr->val);
    theta->val = atan2(zi, zr);
    theta->err = GSL_DBL_EPSILON * fabs(lnr->val);
    return GSL_SUCCESS;
  }
  else {
    DOMAIN_ERROR_2(lnr, theta);
  }
}

int
gsl_sf_log_1plusx_e(const double x, gsl_sf_result * result)
{
  /* CHECK_POINTER(result) */

  if(x <= -1.0) {
    DOMAIN_ERROR(result);
  }
  else if(fabs(x) < GSL_ROOT6_DBL_EPSILON) {
    const double c1 = -0.5;
    const double c2 =  1.0/3.0;
    const double c3 = -1.0/4.0;
    const double c4 =  1.0/5.0;
    const double c5 = -1.0/6.0;
    const double c6 =  1.0/7.0;
    const double c7 = -1.0/8.0;
    const double c8 =  1.0/9.0;
    const double c9 = -1.0/10.0;
    const double t  =  c5 + x*(c6 + x*(c7 + x*(c8 + x*c9)));
    result->val = x * (1.0 + x*(c1 + x*(c2 + x*(c3 + x*(c4 + x*t)))));
    result->err = GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
  else if(fabs(x) < 0.5) {
    double t = 0.5*(8.0*x + 1.0)/(x+2.0);
    gsl_sf_result c;
    cheb_eval_e(&lopx_cs, t, &c);
    result->val = x * c.val;
    result->err = fabs(x * c.err);
    return GSL_SUCCESS;
  }
  else {
    result->val = log(1.0 + x);
    result->err = GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
}

int
gsl_sf_log_1plusx_mx_e(const double x, gsl_sf_result * result)
{
  /* CHECK_POINTER(result) */

  if(x <= -1.0) {
    DOMAIN_ERROR(result);
  }
  else if(fabs(x) < GSL_ROOT5_DBL_EPSILON) {
    const double c1 = -0.5;
    const double c2 =  1.0/3.0;
    const double c3 = -1.0/4.0;
    const double c4 =  1.0/5.0;
    const double c5 = -1.0/6.0;
    const double c6 =  1.0/7.0;
    const double c7 = -1.0/8.0;
    const double c8 =  1.0/9.0;
    const double c9 = -1.0/10.0;
    const double t  =  c5 + x*(c6 + x*(c7 + x*(c8 + x*c9)));
    result->val = x*x * (c1 + x*(c2 + x*(c3 + x*(c4 + x*t))));
    result->err = GSL_DBL_EPSILON * fabs(result->val);
    return GSL_SUCCESS;
  }
  else if(fabs(x) < 0.5) {
    double t = 0.5*(8.0*x + 1.0)/(x+2.0);
    gsl_sf_result c;
    cheb_eval_e(&lopxmx_cs, t, &c);
    result->val = x*x * c.val;
    result->err = x*x * c.err;
    return GSL_SUCCESS;
  }
  else {
    const double lterm = log(1.0 + x);
    result->val = lterm - x;
    result->err = GSL_DBL_EPSILON * (fabs(lterm) + fabs(x));
    return GSL_SUCCESS;
  }
}

/*-*-*-*-*-*-*-*-*-* Functions w/ Natural Prototypes *-*-*-*-*-*-*-*-*-*-*/

#include "eval.h"

double gsl_sf_log(const double x)
{
  EVAL_RESULT(gsl_sf_log_e(x, &result));
}

double gsl_sf_log_abs(const double x)
{
  EVAL_RESULT(gsl_sf_log_abs_e(x, &result));
}

double gsl_sf_log_1plusx(const double x)
{
  EVAL_RESULT(gsl_sf_log_1plusx_e(x, &result));
}

double gsl_sf_log_1plusx_mx(const double x)
{
  EVAL_RESULT(gsl_sf_log_1plusx_mx_e(x, &result));
}


/**
 * \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
     */

}


#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);
}