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

#include <stdio.h>
#include <stdlib.h>
#include <cblas.h>
#include <lapacke.h>
#include <time.h>

flattening_info* flattening_info_init(const MPI_tensor* ten, int flattening, int iscol, int t_v_block)
{
    // Tensor info
    flattening_info* fi = (flattening_info*) malloc(sizeof(flattening_info));
    int t_d = ten->d;

    int* t_nps = (int*) malloc(t_d*sizeof(int));
    int t_Nblocks = 1;
    for (int ii = 0; ii < t_d; ++ii){
        t_nps[ii] = ten->nps[ii];
        t_Nblocks = t_Nblocks * t_nps[ii];
    }

    int* t_t_block = (int*) malloc(t_d*sizeof(int));
    to_tensor_ind(t_t_block, t_v_block, t_nps, t_d);

    long t_N = 1;
    int* t_t_sizes = (int*) malloc(t_d*sizeof(int));
    int* t_t_index = (int*) malloc(t_d*sizeof(int));
    for (int ii = 0; ii < t_d; ++ii){
        int* partition_ii = ten->partitions[ii];
        t_t_index[ii] = partition_ii[t_t_block[ii]];
        t_t_sizes[ii] = partition_ii[t_t_block[ii]+1] - t_t_index[ii];
        t_N = t_N * t_t_sizes[ii];
    }

    // Flattening info
    int f_d = (iscol) ? flattening : t_d-flattening;

    int offset = (iscol) ? 0 : flattening;
    long f_N = 1;
    int f_Nblocks = 1;
    int* f_nps     = (int*) malloc(f_d * sizeof(int));
    int* f_t_block = (int*) malloc(f_d * sizeof(int));
    int* f_t_index = (int*) malloc(f_d * sizeof(int));
    int* f_t_sizes = (int*) malloc(f_d * sizeof(int));
    for (int ii = 0; ii < f_d; ++ii){
        f_nps[ii]     = t_nps[ii+offset];
        f_t_block[ii] = t_t_block[ii+offset];
        f_t_index[ii] = t_t_index[ii+offset];
        f_t_sizes[ii] = t_t_sizes[ii+offset];
        f_N = f_N * f_t_sizes[ii];
        f_Nblocks = f_Nblocks * f_nps[ii];
    }
    int f_v_block = to_vec_ind(f_t_block, f_nps, f_d);

    // Sketching dimensions info
    int s_d = t_d - f_d;
    long s_N = 1;

    offset = (iscol) ? flattening : 0;
    int* s_nps     = (int*) malloc(s_d * sizeof(int));
    int* s_t_index = (int*) malloc(s_d * sizeof(int));
    int* s_t_sizes = (int*) malloc(s_d * sizeof(int));
    for (int ii = 0; ii < s_d; ++ii){
        s_nps[ii]     = t_nps[ii+offset];
        s_t_index[ii] = t_t_index[ii+offset];
        s_t_sizes[ii] = t_t_sizes[ii+offset];
        s_N = s_N * s_t_sizes[ii];
    }

    // Assigning
    fi->t_d       = t_d;
    fi->t_N       = t_N;
    fi->t_Nblocks = t_Nblocks;
    fi->t_nps     = t_nps;
    fi->t_v_block = t_v_block;
    fi->t_t_block = t_t_block;
    fi->t_t_index = t_t_index;
    fi->t_t_sizes = t_t_sizes;

    fi->flattening = flattening;
    fi->iscol      = iscol;
    fi->f_d        = f_d;
    fi->f_N        = f_N;
    fi->f_Nblocks  = f_Nblocks;
    fi->f_nps      = f_nps;
    fi->f_v_block  = f_v_block;
    fi->f_t_block  = f_t_block;
    fi->f_t_index  = f_t_index;
    fi->f_t_sizes  = f_t_sizes;

    fi->s_d       = s_d;
    fi->s_N       = s_N;
    fi->s_nps     = s_nps;
    fi->s_t_index = s_t_index;
    fi->s_t_sizes = s_t_sizes;

    return fi;
}

void flattening_info_free(flattening_info* fi)
{
    free(fi->t_nps);     fi->t_nps = NULL;
    free(fi->t_t_block); fi->t_t_block = NULL;
    free(fi->t_t_index); fi->t_t_index = NULL;
    free(fi->t_t_sizes); fi->t_t_sizes = NULL;
    free(fi->f_nps);     fi->f_nps = NULL;
    free(fi->f_t_block); fi->f_t_block = NULL;
    free(fi->f_t_index); fi->f_t_index = NULL;
    free(fi->f_t_sizes); fi->f_t_sizes = NULL;
    free(fi->s_nps);     fi->s_nps = NULL;
    free(fi->s_t_index); fi->s_t_index = NULL;
    free(fi->s_t_sizes); fi->s_t_sizes = NULL;

    free(fi);
}

void flattening_info_update(flattening_info* fi, const MPI_tensor* ten, int t_v_block)
{
    // Tensor info
    int t_d = ten->d;
    int* t_nps = fi->t_nps;
    int* t_t_block = fi->t_t_block;
    to_tensor_ind(t_t_block, t_v_block, t_nps, t_d);

    long t_N = 1;
    int* t_t_sizes = fi->t_t_sizes;
    int* t_t_index = fi->t_t_index;
    for (int ii = 0; ii < t_d; ++ii){
        int* partition_ii = ten->partitions[ii];
        t_t_index[ii] = partition_ii[t_t_block[ii]];
        t_t_sizes[ii] = partition_ii[t_t_block[ii]+1] - t_t_index[ii];
        t_N = t_N * t_t_sizes[ii];
    }

    // Flattening info
    int flattening = fi->flattening;
    int iscol = fi->iscol;
    int f_d = fi->f_d;

    int offset = (iscol) ? 0 : flattening;
    long f_N = 1;
    int* f_nps     = fi->f_nps;
    int* f_t_block = fi->f_t_block;
    int* f_t_index = fi->f_t_index;
    int* f_t_sizes = fi->f_t_sizes;
    for (int ii = 0; ii < f_d; ++ii){
        f_nps[ii]     = t_nps[ii+offset];
        f_t_block[ii] = t_t_block[ii+offset];
        f_t_index[ii] = t_t_index[ii+offset];
        f_t_sizes[ii] = t_t_sizes[ii+offset];
        f_N = f_N * f_t_sizes[ii];
    }
    int f_v_block = to_vec_ind(f_t_block, f_nps, f_d);

    // Sketching dimensions info
    int s_d = fi->s_d;
    long s_N = 1;

    offset = (iscol) ? flattening : 0;
    int* s_nps     = fi->s_nps;
    int* s_t_index = fi->s_t_index;
    int* s_t_sizes = fi->s_t_sizes;
    for (int ii = 0; ii < s_d; ++ii){
        s_nps[ii]     = t_nps[ii+offset];
        s_t_index[ii] = t_t_index[ii+offset];
        s_t_sizes[ii] = t_t_sizes[ii+offset];
        s_N = s_N * s_t_sizes[ii];
    }

    // Assigning
    fi->t_d = t_d;
    fi->t_N = t_N;
    fi->t_nps = t_nps;
    fi->t_v_block = t_v_block;
    fi->t_t_block = t_t_block;
    fi->t_t_index = t_t_index;
    fi->t_t_sizes = t_t_sizes;

    fi->flattening = flattening;
    fi->iscol = iscol;
    fi->f_d = f_d;
    fi->f_N       = f_N;
    fi->f_nps     = f_nps;
    fi->f_v_block = f_v_block;
    fi->f_t_block = f_t_block;
    fi->f_t_index = f_t_index;
    fi->f_t_sizes = f_t_sizes;

    fi->s_d = s_d;
    fi->s_N = s_N;
    fi->s_nps = s_nps;
    fi->s_t_index = s_t_index;
    fi->s_t_sizes = s_t_sizes;
}

void flattening_info_f_update(flattening_info* fi, const MPI_tensor* ten, int f_v_block)
{
    // tensor info
    int t_d = fi->t_d;

    long t_N = -1;
    int t_v_block = -1;
    int* t_t_block = fi->t_t_block;
    int* t_t_index = fi->t_t_index;
    int* t_t_sizes = fi->t_t_sizes;
    for (int ii = 0; ii < t_d; ++ii){
        t_t_block[ii] = -1;
        t_t_index[ii] = 0;
        t_t_sizes[ii] = 0;
    }
    fi->t_N = t_N;
    fi->t_v_block = t_v_block;

    // flattening info
    int flattening = fi->flattening;
    int iscol = fi->iscol;
    int f_d = fi->f_d;
    int f_Nblocks = fi->f_Nblocks;
    int* f_nps = fi->f_nps;

    int* f_t_block = fi->f_t_block;
    to_tensor_ind(f_t_block, (long) f_v_block, f_nps, f_d);

    long f_N = 1;
    int* f_t_index = fi->f_t_index;
    int* f_t_sizes = fi->f_t_sizes;
    int offset = (iscol) ? 0 : flattening;
    for (int ii = 0; ii < f_d; ++ii){
        int* partition_ii = ten->partitions[ii + offset];
        f_t_index[ii] = partition_ii[f_t_block[ii]];
        f_t_sizes[ii] = partition_ii[f_t_block[ii]+1] - f_t_index[ii];
        f_N = f_N * f_t_sizes[ii];
    }
    fi->f_N = f_N;
    fi->f_v_block = f_v_block;

    // Sketching dimensions info
    int s_d = fi->s_d;
    long s_N = 0;

    int* s_t_index = fi->s_t_index;
    int* s_t_sizes = fi->s_t_sizes;

    for (int ii = 0; ii < s_d; ++ii){
        s_t_index[ii] = 0;
        s_t_sizes[ii] = 0;
    }
    fi->s_N = s_N;
}

void flattening_info_print(flattening_info* fi)
{
    printf("\n~~~~~~~~~~~~~~~ Flattening Info ~~~~~~~~~~~~~~~\n");
    int t_d = fi->t_d;
    printf("t_d = %d\n", t_d);

    printf("t_N = %ld\n", fi->t_N);

    printf("t_Nblocks = %d\n", fi->t_Nblocks);

    if (t_d>0){
        printf("t_nps = [%d", fi->t_nps[0]);
        for (int ii = 1; ii < t_d; ++ii){printf(", %d", fi->t_nps[ii]); }
        printf("]\n");
    }
    else{ printf("t_nps = []\n"); }

    printf("t_v_block = %d\n", fi->t_v_block);

    if (t_d>0){
        printf("t_t_block = [%d", fi->t_t_block[0]);
        for (int ii = 1; ii < t_d; ++ii){printf(", %d", fi->t_t_block[ii]);}
        printf("]\n");
    }
    else{ printf("t_t_block = []\n"); }

    if (t_d>0){
        printf("t_t_index = [%d", fi->t_t_index[0]);
        for (int ii = 1; ii < t_d; ++ii){printf(", %d", fi->t_t_index[ii]);}
        printf("]\n");
    }
    else{ printf("t_t_index = []\n"); }

    if (t_d>0){
        printf("t_t_sizes = [%d", fi->t_t_sizes[0]);
        for (int ii = 1; ii < t_d; ++ii){printf(", %d", fi->t_t_sizes[ii]);}
        printf("]\n");
    }
    else{ printf("t_t_sizes = []\n"); }

    printf("flattening = %d\n", fi->flattening);

    printf("iscol = %d\n", fi->iscol);

    int f_d = fi->f_d;
    printf("f_d = %d\n", f_d);

    printf("f_N = %ld\n", fi->f_N);

    printf("f_Nblocks = %d\n", fi->f_Nblocks);

    if (f_d>0){
        printf("f_nps = [%d", fi->f_nps[0]);
        for (int ii = 1; ii < f_d; ++ii){printf(", %d", fi->f_nps[ii]);}
        printf("]\n");
    }
    else{ printf("f_nps = []\n"); }

    printf("f_v_block = %d\n", fi->f_v_block);

    if (f_d>0){
        printf("f_t_block = [%d", fi->f_t_block[0]);
        for (int ii = 1; ii < f_d; ++ii){printf(", %d", fi->f_t_block[ii]);}
        printf("]\n");
    }
    else{ printf("f_t_block = []\n"); }

    if (f_d>0){
        printf("f_t_index = [%d", fi->f_t_index[0]);
        for (int ii = 1; ii < f_d; ++ii){printf(", %d", fi->f_t_index[ii]);}
        printf("]\n");
    }
    else{ printf("f_t_index = []\n"); }

    if (f_d>0){
        printf("f_t_sizes = [%d", fi->f_t_sizes[0]);
        for (int ii = 1; ii < f_d; ++ii){printf(", %d", fi->f_t_sizes[ii]);}
        printf("]\n");
    }
    else{ printf("f_t_sizes = []\n"); }

    int s_d = fi->s_d;
    printf("s_d = %d\n", s_d);

    printf("s_N = %ld\n", fi->s_N);

    if (s_d>0){
        printf("s_nps = [%d", fi->s_nps[0]);
        for (int ii = 1; ii < s_d; ++ii){printf(", %d", fi->s_nps[ii]);}
        printf("]\n");
    }
    else{ printf("s_nps = []\n"); }

    if (s_d>0){
        printf("s_t_index = [%d", fi->s_t_index[0]);
        for (int ii = 1; ii < s_d; ++ii){printf(", %d", fi->s_t_index[ii]);}
        printf("]\n");
    }
    else{ printf("s_t_index = []\n"); }

    if (s_d>0){
        printf("s_t_sizes = [%d", fi->s_t_sizes[0]);
        for (int ii = 1; ii < s_d; ++ii){printf(", %d", fi->s_t_sizes[ii]);}
        printf("]\n");
    }
    else{ printf("s_t_sizes = []\n"); }
    printf("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n");
}

int* get_sketch_owners(MPI_tensor* ten, int flattening, int iscol)
{
    int np_flattening = 1;
    int* nps = ten->nps;
    int ii0 = (iscol) ? 0 : flattening;
    int ii1 = (iscol) ? flattening : ten->d;

    for (int ii = ii0; ii < ii1; ++ii){
        np_flattening = np_flattening * nps[ii];
    }

    int* owners = get_partition(ten->comm_size, np_flattening);
    return owners;
}

void get_sketch_height(MPI_tensor* ten, int* owners, int flattening, int iscol, long* X_height_ptr, long* buf_height_ptr)
{
    *X_height_ptr = 0;
    *buf_height_ptr = 0;

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

    for (int rank = 0; rank < ten->comm_size; ++rank){
        long X_height_tmp = 0;
        for (int jj = owners[rank]; jj < owners[rank + 1]; ++jj){
            flattening_info_f_update(fi, ten, jj);
            X_height_tmp = X_height_tmp + fi->f_N;
            *buf_height_ptr = (fi->f_N > *buf_height_ptr) ? fi->f_N : *buf_height_ptr;
        }

        *X_height_ptr = (X_height_tmp > *X_height_ptr) ? X_height_tmp : *X_height_ptr;
    }

    flattening_info_free(fi);
}

matrix_tt** get_sketch_Omega(const MPI_tensor* ten, int flattening, int r, int buf, int iscol)
{
    // Seed random numbers
    double t = (double) time(NULL);
    MPI_Bcast(&t, 1, MPI_DOUBLE, 0, ten->comm);
    srand((unsigned long) t);

    int ii0 = iscol ? flattening : 0;
    int n_Omega = iscol ? (ten->d) - flattening : flattening;

    matrix_tt** Omegas = (matrix_tt**) malloc(n_Omega * sizeof(matrix_tt*));
    for (int ii = 0; ii < n_Omega; ++ii){
        Omegas[ii] = matrix_tt_init(ten->n[ii+ii0], r+buf);
        matrix_tt_dlarnv(Omegas[ii]);
    }

    return Omegas;
}

void get_KR_info(int *d_KR, int *stride_KR, matrix_tt** KRs, MPI_tensor* ten, int flattening, int r, int buf, int iscol)
{
    int max_m = 1000;

    int assigned = 0;
    *d_KR = 0;
    int N_KR = 1;

    int ii0 = iscol ? flattening : 0;
    int n_Omega = iscol ? (ten->d) - flattening : flattening;

    KRs[2] = matrix_tt_init(1, r+buf);

    for (int ii = 0; ii < n_Omega; ++ii){
        if (assigned == 0){
            int* partition_ii = ten->partitions[ii + ii0];
            int sz_ii = 0;
            for (int jj = 0; jj < ten->nps[ii+ii0]; ++jj){
                int candidate_n = partition_ii[jj+1] - partition_ii[jj];
                sz_ii = (sz_ii > candidate_n) ? sz_ii : candidate_n;
            }
            if (N_KR * sz_ii < max_m){
                *d_KR = ii+1;
                N_KR = N_KR * sz_ii;
            }
            else{
                *stride_KR = max_m / N_KR;
                KRs[0] = matrix_tt_init(N_KR, r+buf);
                KRs[1] = matrix_tt_init(*stride_KR, r+buf);
                N_KR = N_KR * (*stride_KR);
                KRs[3] = matrix_tt_init(N_KR, r+buf);
                assigned = 1;
            }
        }
    }
    if (assigned == 0){
        *stride_KR = 0;
        KRs[0] = matrix_tt_init(N_KR, r+buf);
        KRs[1] = matrix_tt_init(1, r+buf);
        KRs[3] = matrix_tt_init(N_KR, r+buf);
    }
}

sketch* sketch_init(MPI_tensor* ten, int flattening, int r, int buf, int iscol)
{
    sketch* s = (sketch*) malloc(sizeof(sketch));

    s->ten = ten;
    s->flattening = flattening;
    s->r = r;
    s->buf = buf;
    s->iscol = iscol;
    s->owner_partition = get_sketch_owners(ten, flattening, iscol);
    long recv_buf_height = 0;
    get_sketch_height(ten, s->owner_partition, flattening, iscol, &(s->lda), &recv_buf_height);
    s->Omegas = get_sketch_Omega(ten, flattening, r, buf, iscol);
    s->X_size = (s->lda)*(r+buf);
    s->X = (double*) calloc(s->X_size, sizeof(double));
    s->scratch = (double*) calloc(s->X_size, sizeof(double));
    s->recv_buf = (double*) calloc(recv_buf_height * (r+buf), sizeof(double));
    s->fi = flattening_info_init(ten, flattening, iscol, 0);
    s->KRs = (matrix_tt**) calloc(4, sizeof(matrix_tt*));
    get_KR_info(&(s->d_KR), &(s->stride_KR), s->KRs, ten, flattening, r, buf, iscol);

    return s;
}

matrix_tt** copy_sketch_Omega(matrix_tt** Omegas, MPI_tensor* ten, int flattening, int r, int buf, int iscol)
{
    int ii0 = iscol ? flattening : 0;
    int n_Omega = iscol ? (ten->d) - flattening : flattening;

    matrix_tt** Omegas_cp = (matrix_tt**) malloc(n_Omega * sizeof(matrix_tt*));
    for (int ii = 0; ii < n_Omega; ++ii){
        submatrix_update(Omegas[ii], 0, ten->n[ii+ii0], 0, r+buf);
        Omegas_cp[ii] = matrix_tt_copy(Omegas[ii]);
    }

    return Omegas_cp;
}

sketch* sketch_init_with_Omega(MPI_tensor* ten, int flattening, int r, int buf, int iscol, matrix_tt** Omegas)
{
    sketch* s = (sketch*) malloc(sizeof(sketch));
    s->ten = ten;
    s->flattening = flattening;
    s->r = r;
    s->buf = buf;
    s->iscol = iscol;
    s->owner_partition = get_sketch_owners(ten, flattening, iscol);
    long recv_buf_height;
    get_sketch_height(ten, s->owner_partition, flattening, iscol, &(s->lda), &recv_buf_height);
    s->Omegas = copy_sketch_Omega(Omegas, ten, flattening, r, buf, iscol);
    s->X_size = (s->lda)*(r+buf);
    s->X = (double*) calloc(s->X_size, sizeof(double));
    s->scratch = (double*) calloc(s->X_size, sizeof(double));
    s->recv_buf = (double*) calloc(recv_buf_height * (r+buf), sizeof(double));
    s->fi = flattening_info_init(ten, flattening, iscol, 0);
    s->KRs = (matrix_tt**) calloc(4, sizeof(matrix_tt*));
    get_KR_info(&(s->d_KR), &(s->stride_KR), s->KRs, ten, flattening, r, buf, iscol);

    return s;
}

// NOTE: does not free the MPI_tensor. We assume this is shared, and should not be freed this way.
void sketch_free(sketch* s)
{
    MPI_tensor* ten = s->ten; s->ten = NULL;
    int d = ten->d;
    int n_Omega = s->iscol ? d - s->flattening : s->flattening;
    for (int ii = 0; ii < n_Omega; ++ii){
        matrix_tt_free(s->Omegas[ii]); s->Omegas[ii] = NULL;
    }
    for (int ii = 0; ii < 4; ++ii){
        matrix_tt_free(s->KRs[ii]); s->KRs[ii] = NULL;
    }
    free(s->KRs); s->KRs = NULL;

    free(s->owner_partition); s->owner_partition = NULL;
    free(s->X);               s->X = NULL;
    free(s->scratch);         s->scratch = NULL;
    free(s->Omegas);          s->Omegas = NULL;
    free(s->recv_buf);        s->recv_buf = NULL;
    flattening_info_free(s->fi); s->fi = NULL;
    free(s);
}

void sketch_print(sketch* s)
{
    printf("Printing sketch at %p\n", s);

    printf("\nten address: %p\n", s->ten);
    printf("flattening = %d\n",s->flattening);
    printf("r = %d\n", s->r);
    printf("buf = %d\n", s->buf);
    printf("iscol = %d\n", s->iscol);

    int* op = s->owner_partition;
    MPI_tensor* ten = s->ten;
    printf("owner_partition = [%d", op[0]);
    for (int ii = 1; ii < (ten->comm_size) + 1; ++ii){
        printf(", %d", op[ii]);
    }
    printf("]\n");

    printf("lda = %ld\n", s->lda);
    printf("X_size = %ld\n", s->X_size);

    printf("\nX = \n");
    int r = s->r; int buf = s->buf; long X_size = s->X_size; long lda = s->lda;
    matrix_tt* X_mat = matrix_tt_wrap(lda, r+buf, s->X);
    matrix_tt_print(X_mat, 1);
    free(X_mat); X_mat = NULL;

    printf("\nscratch = \n");
    matrix_tt* scratch_mat = matrix_tt_wrap(lda, r+buf, s->scratch);
    matrix_tt_print(scratch_mat, 1);
    free(scratch_mat); scratch_mat = NULL;

    int n_Omega = s->iscol ? (s->ten)->d-s->flattening : s->flattening;
    for (int ii = 0; ii < n_Omega; ++ii){
        printf("\nOmegas[%d] = \n", ii);
        matrix_tt_print(s->Omegas[ii], 1);
    }
}

// Performs C[ii + n_ii * jj,kk] = A[ii, kk] * B[jj, kk]
void submatrix_khatri_rao_outer_product(matrix_tt* A, matrix_tt* B, matrix_tt* C)
{
    for (int kk = 0; kk < A->n; ++kk){
        int A_offset = A->offset + kk*(A->lda);
        int B_offset = B->offset + kk*(B->lda);
        int C_offset = C->offset + kk*(C->lda);

        cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans,
                A->m, B->m, 1,
                1.0,
                A->X + A_offset, A->m,
                B->X + B_offset, 1,
                0.0,
                C->X + C_offset, A->m);
    }
}

// General idea of the algorithm:
// Split the Omegas into 3 parts: ii < d_KR, ii == d_KR, ii > d_KR
// For ii <  d_KR, we will pre-multiply the sketch matrices Omega[ii] (Call this KR_1)
// For ii == d_KR, we will loop through blocks of Omega[d_KR]
// For ii > d_KR, we will loop and multiply to get a row vector       (Call this KR_3)

// Then, at each point in the loop, we will perform
//     KR_2 = Omegas[ii] * KR_3 and
//     KR_4 = KR_1 * KR_2,
// where * is the Khatri-Rao product. Finally, we multiply X_mat by KR_4 to update the sketch, and repeat.
void subtensor_khatri_rao(sketch* s, matrix_tt* C, flattening_info* fi, double beta, matrix_tt* X_mat)
{
    int r = s->r + s->buf;
    matrix_tt** Omegas = s->Omegas;

    MPI_tensor* ten = s->ten;
    int d_KR = s->d_KR;
    matrix_tt* KR_1;
    matrix_tt* KR_4;
    if ((d_KR == 0) || (d_KR%2 == 1)){
        KR_1 = s->KRs[0];
        KR_4 = s->KRs[3];
    }
    else{
        KR_1 = s->KRs[3];
        KR_4 = s->KRs[0];
    }
    int s_d = fi->s_d;

    int N_kk = 1;

    int h = 0;

    int Omega_offset_KR;

//    printf("Starting loop to get KR_1\n");
    for (int ii = 0; ii < s_d; ++ii){
        int ii0 = fi->s_t_index[ii];
        int ii1 = ii0 + fi->s_t_sizes[ii];

        submatrix_update(Omegas[ii], ii0, ii1, 0, r);

        if (ii > d_KR){
            N_kk = N_kk * fi->s_t_sizes[ii]; // Number of elements in outer loop
        }
        else if (ii == d_KR){
            Omega_offset_KR = ii0;
        }

        if ((ii == 0) && (d_KR > 0)){
            matrix_tt_reshape(ii1-ii0, r, KR_1);
            matrix_tt_copy_data(KR_1, Omegas[0]);
            h = ii1 - ii0;
        }
        else if (ii < d_KR){ // Multiply inner matrices
            matrix_tt* tmp = KR_1;
            KR_1 = KR_4;
            KR_4 = tmp;

            h = h*(ii1-ii0);
            matrix_tt_reshape(h, r, KR_1);
            submatrix_khatri_rao_outer_product(KR_4, Omegas[ii], KR_1);
        }
    }
//    printf("Finished loop to get KR_1\n");

    if (fi->iscol){
        matrix_tt_wrap_update(X_mat, fi->f_N, fi->s_N, get_X(ten));
    }
    else{
        matrix_tt_wrap_update(X_mat, fi->s_N, fi->f_N, get_X(ten));
    }
    X_mat->transpose = (fi->iscol ? 0 : 1);

    if (d_KR == s_d){ // If we have done everything, just multiply (the sketch is relatively small)
        matrix_tt_dgemm(X_mat, KR_1, C, 1.0, beta);
        return;
    }

    int stride_KR = s->stride_KR;
    matrix_tt* KR_2 = s->KRs[1];
    matrix_tt* KR_3 = s->KRs[2];

    int* t_kk = ten->t_kk;
    int size_KR = fi->s_t_sizes[d_KR];
    int N_ll = 1 + (size_KR - 1) / stride_KR;

    int tensor_offset = 0;

    for (int kk = 0; kk < N_kk; ++kk) {
        to_tensor_ind(t_kk, (long) kk, fi->s_t_sizes + d_KR + 1, s_d - 1 - d_KR);

        for (int ii = d_KR + 1; ii < s_d; ++ii){
            if (ii == d_KR + 1){
                for (int jj = 0; jj < r; ++jj){
                    KR_3->X[jj] = matrix_tt_element(Omegas[ii], t_kk[ii-d_KR-1], jj);
                }
            }
            else{
                for (int jj = 0; jj < r; ++jj){
                    KR_3->X[jj] *= matrix_tt_element(Omegas[ii], t_kk[ii-d_KR-1], jj);
                }
            }
        }

        for (int ll = 0; ll < N_ll; ++ll){
            double bb = ((kk==0) && (ll==0)) ? beta : 1.0;
            int ii0 = Omega_offset_KR + ll * stride_KR;
            int ii1 = ((ll+1)*stride_KR > size_KR) ? size_KR + Omega_offset_KR: (ll+1)*stride_KR + Omega_offset_KR;

            submatrix_update(Omegas[d_KR], ii0, ii1, 0, r);
            if (d_KR + 1 == s_d){ // If we only need KR_1 and Omegas[ii]
                matrix_tt_reshape((KR_1->m) * (ii1-ii0), r, KR_4);
                submatrix_khatri_rao_outer_product(KR_1, Omegas[d_KR], KR_4);
            }
            else if (d_KR == 0){ // If we only need Omegas[ii] and KR_3
                matrix_tt_reshape(Omegas[d_KR]->m, r, KR_4);
                submatrix_khatri_rao_outer_product(Omegas[d_KR], KR_3, KR_4);
            }
            else{ // We need everything...
                matrix_tt_reshape(ii1-ii0, r, KR_2);
                submatrix_khatri_rao_outer_product(Omegas[d_KR], KR_3, KR_2);
                matrix_tt_reshape((ii1-ii0) * (KR_1->m), r, KR_4);
                submatrix_khatri_rao_outer_product(KR_1, KR_2, KR_4);
            }

            if (fi->iscol) {
                submatrix_update(X_mat, 0, fi->f_N, tensor_offset, tensor_offset + KR_4->m);
            }
            else {
                submatrix_update(X_mat, tensor_offset, tensor_offset + KR_4->m, 0, fi->f_N);
            }

            matrix_tt_dgemm(X_mat, KR_4, C, 1.0, bb);
            tensor_offset = tensor_offset + KR_4->m;
        }
    }
}

int get_owner(int block, int* owner_partition, int comm_size)
{
    int color = -1;
    for (int ii = 0; ii < comm_size; ++ii){
        if ((block >= owner_partition[ii]) && (block < owner_partition[ii+1])){
            color = ii;
        }
    }

    return color;
}

int s_get_owner(sketch* s, int f_v_block){
    int* op = s->owner_partition;
    MPI_tensor* ten = s->ten;
    int size = ten->comm_size;

    // Binary search (basically wikipedia)
    int a = 0;
    int b = size-1;
    int c;

    while (a <= b){
        c = (a+b)/2;
        if (f_v_block < op[c]){
            b = c-1;
        }
        else if(f_v_block >= op[c+1]){
            a = c+1;
        }
        else{
            return c;
        }
    }

    return -1;
}

void own_submatrix_update(matrix_tt* mat, sketch* s, int f_v_block, int with_buf){
    if (f_v_block == -1){
        return;
    }

    int sketch_offset = 0;
    int sketch_lda = s->lda;
    MPI_tensor* ten = s->ten;
    int world_rank = ten->rank;
    int* owner_partition = s->owner_partition;

    flattening_info* fi = s->fi;
    if ((f_v_block < owner_partition[world_rank]) || (f_v_block >= owner_partition[world_rank+1])){
        printf("r%d own_submatrix: This core does not own f_v_block = %d\n", world_rank, f_v_block);
    }

    for (int block = s->owner_partition[world_rank]; block < f_v_block; ++block){
        flattening_info_f_update(fi, ten, block);
        sketch_offset = sketch_offset + fi->f_N;
    }
    flattening_info_f_update(fi, ten, f_v_block);

    matrix_tt_wrap_update(mat, sketch_lda, s->r + s->buf, s->X);
    int r = (with_buf) ? s->r + s->buf : s->r;
    submatrix_update(mat, sketch_offset, sketch_offset + fi->f_N, 0, r);
}

matrix_tt* own_submatrix(sketch* s, int f_v_block, int with_buf){
    matrix_tt* mat = (matrix_tt*) malloc(sizeof(matrix_tt));
    own_submatrix_update(mat, s, f_v_block, with_buf);
    return mat;
}

void subtensor_sketch_multiply(sketch* s, flattening_info* fi, matrix_tt** holders)
{
    MPI_tensor* ten = s->ten;
    int* owner_partition = s->owner_partition;
    int world_size = ten->comm_size;
    int world_rank = ten->rank;

    // Actually sketch the thing
    matrix_tt* sketch_mat = holders[0];
    double beta = 0.0;
    if (ten->current_part != -1){
        int current_color = get_owner(fi->f_v_block, owner_partition, world_size);
        if (current_color == world_rank){
            own_submatrix_update(sketch_mat, s, fi->f_v_block, 1);
            beta = 1.0;
        }
        else{
            matrix_tt_wrap_update(sketch_mat, fi->f_N, s->r + s->buf, s->scratch);
            beta = 0.0;
        }

        subtensor_khatri_rao(s, sketch_mat, fi, beta, holders[1]);
    }
}

void subtensor_sketch_communicate(sketch* s, int stream_step, flattening_info* fi, matrix_tt** holders)
{
    MPI_tensor* ten = s->ten;
    int* owner_partition = s->owner_partition;
    int world_size = ten->comm_size;
    int world_rank = ten->rank;
    int* group_ranks = ten->group_ranks;

    flattening_info* fi_tmp = s->fi;
    matrix_tt* sketch_mat = holders[0];

    // For each flattening block
    for (int ii = 0; ii < fi->f_Nblocks; ++ii){
        int kk = 0;
        int owner_ii = get_owner(ii, owner_partition, world_size);
        int group_owner;
        // For each rank in the world
        for (int jj = 0; jj < world_size; ++jj){
            // If the rank owns the block
            if (owner_ii == jj){
                // Add it to the group
                group_ranks[kk] = jj;
                kk = kk+1;
                group_owner = jj;
            }
            else{
                int* schedule_jj = ten->schedule[jj];
                int stream_jj = schedule_jj[stream_step];
                if (stream_jj != -1){
                    flattening_info_update(fi_tmp, ten, schedule_jj[stream_step]);
                    // Else, if the rank just calculated something that adds to the block
                    if (ii == fi_tmp->f_v_block){
                        // Add it to the group
                        group_ranks[kk] = jj;
                        kk = kk+1;
                    }
                }
            }
        }

        if (world_rank == owner_ii){
            matrix_tt* comm_matrix = holders[1];
            own_submatrix_update(comm_matrix, s, ii, 1);
            matrix_tt* buf_mat = holders[2];
            matrix_tt_wrap_update(buf_mat, comm_matrix->m, comm_matrix->n, s->recv_buf);

            matrix_tt_group_reduce(ten->comm, world_rank, comm_matrix, buf_mat, owner_ii, group_ranks, kk);
        }
        else if (sketch_mat->X){
            matrix_tt* buf_mat = holders[2];
            matrix_tt_wrap_update(buf_mat, sketch_mat->m, sketch_mat->n, s->recv_buf);
            matrix_tt_group_reduce(ten->comm, world_rank,sketch_mat, buf_mat, owner_ii, group_ranks, kk);
        }

    }
}

void multi_perform_sketch(sketch** sketches, int n_sketch)
{
    int n_holders = 3;
    matrix_tt** holders = (matrix_tt**) malloc(n_holders*n_sketch*sizeof(matrix_tt*));
    for (int ii = 0; ii < n_holders*n_sketch; ++ii){
        holders[ii] = (matrix_tt*) calloc(1, sizeof(matrix_tt));
    }

    MPI_tensor* ten = sketches[0]->ten; // Tensor
    int rank = ten->rank;
    int* schedule_rank = ten->schedule[rank];

    flattening_info** fis = (flattening_info**) calloc(n_sketch, sizeof(flattening_info*));
    for (int ii = 0; ii < n_sketch; ++ii){
        fis[ii] = flattening_info_init(ten, sketches[ii]->flattening, sketches[ii]->iscol, 0);
    }

    for (int ii = 0; ii < ten->n_schedule; ++ii){
        stream(ten, schedule_rank[ii]);
        for (int jj = 0; jj < n_sketch; ++jj){
            matrix_tt** holders_jj = holders + jj*n_holders;
            flattening_info_update(fis[jj], ten, ten->current_part);
            holders_jj[0]->X = NULL;
            subtensor_sketch_multiply(sketches[jj], fis[jj], holders_jj);
        }

        for (int jj = 0; jj < n_sketch; ++jj){
            matrix_tt** holders_jj = holders + jj*n_holders;
            subtensor_sketch_communicate(sketches[jj], ii, fis[jj], holders_jj);
        }
    }

    for (int ii = 0; ii < n_holders*n_sketch; ++ii){
        free(holders[ii]); holders[ii] = NULL;
    }

    for (int ii = 0; ii < n_sketch; ++ii){
        flattening_info_free(fis[ii]); fis[ii] = NULL;
    }
    free(fis); fis = NULL;
    free(holders); holders = NULL;
}

void sketch_qr(sketch* sketch)
{

    MPI_tensor* ten = sketch->ten;
    int* owner_partition = sketch->owner_partition;
    int flattening = sketch->flattening;
    int iscol = sketch->iscol;
    int r = sketch->r;
    int buf = sketch->buf;

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

    // A little pre-processing
    int head = 0;
    flattening_info* fi = flattening_info_init(ten, flattening, iscol, 0);
    int* Ns = (int*) calloc(size, sizeof(int));

    for (int ii = 0; ii < size; ++ii){
        for (int jj = owner_partition[ii]; jj < owner_partition[ii+1]; ++jj){
            flattening_info_f_update(fi, ten, jj);
            Ns[ii] = Ns[ii] + fi->f_N;
        }
    }

    int N_rank = Ns[rank];
    int lda = sketch->lda;
    matrix_tt* Q_big = matrix_tt_wrap(lda, r+buf, sketch->X);
    matrix_tt* Q = submatrix(Q_big, 0, N_rank, 0, r+buf);
    matrix_tt* Q_head = NULL;
    matrix_tt* R = NULL;

    if (rank == head){
        Q_head = matrix_tt_init(size * (r+buf), r+buf);
        R = submatrix(Q_head, 0, r+buf, 0, r+buf);
    }
    else{
        R = matrix_tt_init(r+buf, r+buf);
    }

    matrix_tt_truncated_qr(Q, R, r+buf);

    // Gather the Rs
    for (int jj = 0; jj < r+buf; ++jj){
//        printf("r%d jj%d\n", rank, jj);
        if (rank == head){
            double* col = Q_head->X + jj*(Q_head->lda);
            MPI_Gather(col, r+buf, MPI_DOUBLE, col, r+buf, MPI_DOUBLE, head, comm);
        }
        else{
            MPI_Gather(R->X + jj*(R->lda), r+buf, MPI_DOUBLE, NULL, r+buf, MPI_DOUBLE, head, comm);
        }
    }

    // Take the QR of the Rs
    if (rank == head){
        matrix_tt_truncated_qr(Q_head, NULL, r+buf);
    }

    // Scatter the Q of the preceding step
    for (int jj = 0; jj < r+buf; ++jj){
        if (rank == head){
            double* col = Q_head->X + jj*(Q_head->lda);
            MPI_Scatter(col, r+buf, MPI_DOUBLE, col, r+buf, MPI_DOUBLE, head, comm);
        }
        else{
            MPI_Scatter(NULL, r+buf, MPI_DOUBLE, R->X + jj*(R->lda), r+buf, MPI_DOUBLE, head, comm);
        }
    }

    // Multiply to get the final Q
    matrix_tt* X_big = matrix_tt_wrap(lda, r, sketch->scratch);
    matrix_tt* new_X = submatrix(X_big, 0, N_rank, 0, r);
    matrix_tt* Q_head_sub = submatrix(R, 0, r+buf, 0, r);
    matrix_tt_dgemm(Q, Q_head_sub, new_X, 1.0, 0.0);

    // Switch so the QR lives in X
    double* tmp = sketch->scratch;
    sketch->scratch = sketch->X;
    sketch->X = tmp;

    free(X_big);              X_big = NULL;
    free(new_X);              new_X = NULL;
    free(Q_head_sub);         Q_head_sub = NULL;
    flattening_info_free(fi); fi = NULL;
    free(Ns);                 Ns = NULL;
    free(Q_big);              Q_big = NULL;
    free(Q);                  Q = NULL;
    if (rank == head){
        matrix_tt_free(Q_head); Q_head = NULL;
        free(R);             R = NULL;
    }
    else{
        matrix_tt_free(R); R = NULL;
    }
}

void perform_sketch(sketch* s)
{
    multi_perform_sketch(&s, 1);
}

// Gets the sketch block from the correct owner (stored in fi->f_v_block). Returns a matrix with the correct dimensions
void sendrecv_sketch_block(matrix_tt* mat, sketch* s, flattening_info* fi, int recv_rank, int with_buf)
{
    int f_v_block = fi->f_v_block;

    MPI_tensor* ten = s->ten;
    int rank = ten->rank;
    MPI_Comm comm = ten->comm;

    int owner = s_get_owner(s, f_v_block);
    int r = (with_buf) ? s->r + s->buf : s->r;

    if (rank == recv_rank){
        if (rank == owner){
            own_submatrix_update(mat, s, f_v_block, with_buf);
        }
        else{
            matrix_tt_wrap_update(mat, fi->f_N, r, s->scratch);
            matrix_tt_recv(comm, mat, owner);
        }
    }
    else{
        if(rank == owner){
            own_submatrix_update(mat, s, f_v_block, with_buf);
            matrix_tt_send(comm, mat, recv_rank);
        }
        mat->X = NULL;
    }
}

// Eats s, so be careful!
MPI_tensor* sketch_to_tensor(sketch** s_ptr)
{
    sketch* s = *s_ptr;

    MPI_tensor* sten = (MPI_tensor*) malloc(sizeof(MPI_tensor));
    MPI_tensor* ten = s->ten;
    int rank = ten->rank;
    int size = ten->comm_size;

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

    // Get the schedule from the owner_partition
    int n_schedule = 0;
    int* owners = s->owner_partition;
    for (int ii = 0; ii < size; ++ii){
        int sz = owners[ii+1] - owners[ii];
        n_schedule = (n_schedule > sz) ? n_schedule : sz;
    }

    int** schedule = (int**) malloc(size * sizeof(int*));
    int* inverse_schedule = (int*) malloc(fi->f_Nblocks * sizeof(int));
    for (int ii = 0; ii < size; ++ii){
        schedule[ii] = (int*) malloc(n_schedule * sizeof(int));

        int* schedule_ii = schedule[ii];
        int jj0 = owners[ii];
        int sz = owners[ii+1] - jj0;
        for (int jj = 0; jj < n_schedule; ++jj){
            if (jj < sz){
                schedule_ii[jj] = jj+jj0;
                inverse_schedule[jj+jj0] = ii*n_schedule + jj;
            }
            else{
                schedule_ii[jj] = -1;
            }
        }
    }

    int d = fi->f_d + 1;
    int soffset = (s->iscol) ? 0 : 1;
    int offset = (s->iscol) ? 0 : s->flattening;
    int end_ind = (s->iscol) ? d-1 : 0;

    // Inherit the properties of ten
    int* n = (int*) malloc(d * sizeof(int));
    int* nps = (int*) malloc(d * sizeof(int));
    int** partitions = (int**) malloc(d * sizeof(int*));
    for (int ii = 0; ii < d-1; ++ii){
        n[ii + soffset] = ten->n[ii + offset];
        nps[ii + soffset] = ten->nps[ii + offset];
        partitions[ii + soffset] = (int*) malloc( (nps[ii+soffset] + 1) * sizeof(int));

        int* spartition_ii = partitions[ii+soffset];
        int* partition_ii = ten->partitions[ii+offset];
        for (int jj = 0; jj < nps[ii+soffset] + 1; ++jj){
            spartition_ii[jj] = partition_ii[jj];
        }
    }
    // Set the special sketched dimension properties
    n[end_ind] = s->r;
    nps[end_ind] = 1;
    partitions[end_ind] = (int*) malloc( (nps[end_ind] + 1) * sizeof(int));
    int* spartition_end = partitions[end_ind];
    spartition_end[0] = 0; spartition_end[1] = n[end_ind];

    // Get the parameters giving the subtensor locations
    double** subtensors = (double**) malloc(n_schedule * sizeof(double*));
    int* schedule_rank = schedule[rank];
    int with_buf = 0;
    long scratch_offset = 0;
    for (int ii = 0; ii < n_schedule; ++ii){
        if (schedule_rank[ii] != -1){
            matrix_tt* X_submat = own_submatrix(s, schedule_rank[ii], with_buf);
            subtensors[ii] = s->scratch + scratch_offset;
            matrix_tt* scratch_submat = matrix_tt_wrap(X_submat->m, X_submat->n, subtensors[ii]);
            matrix_tt_copy_data(scratch_submat, X_submat);
            scratch_offset = scratch_offset + X_submat->n * X_submat->m;
            free(X_submat);
            free(scratch_submat);
        }
    }
    void* parameters = p_static_init(subtensors);

    // Assigning fields
    sten->d = d;
    sten->n = n;

    sten->comm = ten->comm;
    sten->rank = ten->rank;
    sten->comm_size = ten->comm_size;

    sten->schedule = schedule;
    sten->n_schedule = n_schedule;
    sten->inverse_schedule = inverse_schedule;

    sten->partitions = partitions;
    sten->nps = nps;

    sten->current_part = -1;

    sten->f_ten = NULL;
    sten->parameters = parameters;

    sten->X_size = s->X_size;
    sten->X = s->scratch;

    sten->ind1 = (int*) malloc(d * sizeof(int));
    sten->ind2 = (int*) malloc(d * sizeof(int));
    sten->tensor_part = (int*) malloc(d * sizeof(int));
    sten->group_ranks = (int*) malloc(ten->comm_size * sizeof(int));
    sten->t_kk = (int*) malloc(d * sizeof(int));

    // Freeing sketch things
    int n_Omega = s->iscol ? ten->d - s->flattening : s->flattening;
    for (int ii = 0; ii < n_Omega; ++ii){
        matrix_tt_free(s->Omegas[ii]); s->Omegas[ii] = NULL;
    }
    for (int ii = 0; ii < 4; ++ii){
        matrix_tt_free(s->KRs[ii]); s->KRs[ii] = NULL;
    }
    free(s->KRs); s->KRs = NULL;
    s->ten = NULL;
    free(s->owner_partition); s->owner_partition = NULL;
    free(s->X);               s->X = NULL;
    free(s->Omegas);          s->Omegas = NULL;
    free(s->recv_buf);        s->recv_buf = NULL;
    flattening_info_free(s->fi); s->fi = NULL;
    free(s);                  *s_ptr = NULL;

    flattening_info_free(fi);

    return sten;
}


/*
* -----------------------------------------------------------------
*  isat_lib.c
*  In Situ Adaptive Tabulation Library
*  Version: 2.0
*  Last Update: Nov 3, 2019
*
*  Programmer: Americo Barbosa da Cunha Junior
*              americo.cunhajr@gmail.com
* -----------------------------------------------------------------
*  Copyright (c) 2010-2019, Americo Barbosa da Cunha Junior
*  All rights reserved.
* -----------------------------------------------------------------
*  This is the implementation file for ISAT_LIB module, a
*  computational library with In Situ Adaptive Tabulation (ISAT)
*  algorithm routines.
* -----------------------------------------------------------------
*/

#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_errno.h>

#include "../include/thrm_lib.h"
#include "../include/ell_lib.h"
#include "../include/ode_lib.h"
#include "../include/isat_lib.h"

/*
*------------------------------------------------------------
*   isat_alloc
*
*   This function alocates memory for an ISAT workspace
*   structure and initialize its elements.
*
*   Output:
*   isat_mem - pointer to ISAT workspace
*
*   last update: Oct 9, 2019
*------------------------------------------------------------
*/

isat_wrk *isat_alloc()
{
    /* create ISAT workspace */
    isat_wrk *isat_mem = NULL;

    /* memory allocation for ISAT workspace */
    isat_mem = (isat_wrk *) malloc(sizeof(isat_wrk));
    if ( isat_mem == NULL )
        return NULL;

    /* initialize ISAT workspace elements */
    isat_mem->root      = NULL;
    isat_mem->lf        = 0;
    isat_mem->nd        = 0;
    isat_mem->add       = 0;
    isat_mem->grw       = 0;
    isat_mem->rtv       = 0;
    isat_mem->dev       = 0;
    isat_mem->hgt       = 0;
    isat_mem->maxleaves = 0;
    isat_mem->time_add  = 0.0;
    isat_mem->time_grw  = 0.0;
    isat_mem->time_rtv  = 0.0;
    isat_mem->time_dev  = 0.0;

    return isat_mem;
}
/*------------------------------------------------------------*/

/*
*------------------------------------------------------------
*   isat_free
*
*   This function release the memory used by ISAT workspace.
*
*   Input:
*   isat_bl - pointer to ISAT workspace
*
*   last update: Oct 9, 2019
*------------------------------------------------------------
*/

void isat_free(void **isat_bl)
{
    /* create ISAT workspace */
    isat_wrk *isat_mem = NULL;

    /* check if ISAT workspace memory block is NULL */
    if ( *isat_bl == NULL )
        return;

    isat_mem = (isat_wrk *) (*isat_bl);

    /* release the memory allocated for ISAT workspace elements */
    if( isat_mem->root != NULL )
    {
        bst_node_free((void **)&(isat_mem->root));
        isat_mem->root = NULL;
    }

    /* release the memory allocated for ISAT workspace */
    free(*isat_bl);
    *isat_bl = NULL;

    return;
}
/*------------------------------------------------------------*/

/*
*------------------------------------------------------------
*   isat_bst_init
*
*   This function initiates ISAT binary search tree.
*
*   Input:
*   isat_mem - pointer to ISAT workspace
*
*   last update: Oct 9, 2019
*------------------------------------------------------------
*/

int isat_bst_init(isat_wrk *isat_mem)
{
    /* check if ISAT workspece is allocated */
    if ( isat_mem == NULL )
        return GSL_EINVAL;

    /* memory allocation for ISAT binary search tree root */
    if ( isat_mem->root == NULL )
    {
        isat_mem->root = bst_node_alloc();
        if ( isat_mem->root == NULL )
        {
            free(isat_mem);
            isat_mem = NULL;
            return GSL_EINVAL;
        }
    }

    return GSL_SUCCESS;
}
/*------------------------------------------------------------*/

/*
*------------------------------------------------------------
*   isat_statistics
*
*   This function prints on screen ISAT workspace
*   statistics of usage.
*
*   Input:
*   isat_mem - pointer to ISAT workspace
*
*   last update: Oct 1, 2019
*------------------------------------------------------------
*/

void isat_statistics(isat_wrk *isat_mem)
{
    double time_add;
    double time_grw;
    double time_rtv;
    double time_dev;

    time_add = (double) (isat_mem->time_add / CLOCKS_PER_SEC) / isat_mem->add;
    time_grw = (double) (isat_mem->time_grw / CLOCKS_PER_SEC) / isat_mem->grw;
    time_rtv = (double) (isat_mem->time_rtv / CLOCKS_PER_SEC) / isat_mem->rtv;
    time_dev = (double) (isat_mem->time_dev / CLOCKS_PER_SEC) / isat_mem->dev;

    printf("\n ISAT statistics:");
    printf("\n # of adds         = %d",   isat_mem->add);
    printf("\n # of grows        = %d",   isat_mem->grw);
    printf("\n # of retrieves    = %d",   isat_mem->rtv);
    printf("\n # of dir. eval.   = %d",   isat_mem->dev);
    printf("\n # of leaves       = %d",   isat_mem->lf);
    printf("\n # of nodes        = %d",   isat_mem->nd);
    printf("\n tree height       = %d\n", isat_mem->hgt);

    printf("\n average values for CPU time (s):");
    printf("\n add: %+.6e",time_add);
    printf("\n grw: %+.6e",time_grw);
    printf("\n rtv: %+.6e",time_rtv);
    printf("\n dev: %+.6e",time_dev);

    return;
}
/*------------------------------------------------------------*/

/*
* -----------------------------------------------------------------
*   isat_input
*
*   This function receives ISAT input parameters from the user.
*
*   Input:
*   maxleaves - maximum number of ISAT tree leaves
*   etol      - ISAT error tolerance
*
*   Output:
*   success or error
*
*   last update: Nov 3, 2019
* -----------------------------------------------------------------
*/

int isat_input(unsigned int *maxleaves,double *etol)
{
    printf("\n Input ISAT parameters:\n");

    printf("\n maximum of tree leaves:");
    scanf("%d", maxleaves);
    printf("\n %d\n", *maxleaves);
    if( *maxleaves <= 0 )
	GSL_ERROR(" maxleaves must be a positive integer",GSL_EINVAL);

    printf("\n ISAT error tolerance:");
    scanf("%lf", etol);
    printf("\n %+.1e\n", *etol);
    if( *etol < 0.0 )
	GSL_ERROR(" etol must be grather than zero",GSL_EINVAL);

    return GSL_SUCCESS;
}
/*----------------------------------------------------------------*/

/*
*------------------------------------------------------------
*   isat_eoa_matrix
*
*   This function computes the EOA matrix in Cholesky form.
*
*   Input:
*   A    - gradient matrix
*   etol - error tolerance
*
*   Output:
*   L    - EOA Cholesky matrix
*
*   last update: Oct 9, 2019
*------------------------------------------------------------
*/

void isat_eoa_mtrx(gsl_matrix *A,
                   double etol,
                   gsl_matrix *L)
{
    unsigned int i;
    double  eps_max   = 0.5;
    gsl_matrix *Aetol = NULL;
    gsl_matrix *V     = NULL;
    gsl_vector *sigma = NULL;

    /* memory allocation */
    Aetol = gsl_matrix_calloc(A->size2,A->size2);
    V     = gsl_matrix_calloc(A->size2,A->size2);
    sigma = gsl_vector_calloc(A->size2);

    /* Aetol := A */
    gsl_matrix_memcpy(Aetol,A);

    /* Aetol := (1/etol).A */
    gsl_matrix_scale (Aetol,1.0/etol);

    /* Aetol = U*sigma*V^T */
    ell_psd2eig(Aetol,V,sigma);

    /* eliminate small and large singular values */
    for ( i = 0; i < sigma->size; i++ )
        sigma->data[i] = GSL_MAX(sigma->data[i],eps_max);

    /* V*sigma^2*V^T = L*L^T */
    ell_eig2chol(V,sigma,L);

    /* release allocated memory */
    gsl_vector_free(sigma);
    gsl_matrix_free(V);
    gsl_matrix_free(Aetol);
    sigma = NULL;
    V     = NULL;
    Aetol = NULL;

    return;
}
/*------------------------------------------------------------*/

/*
*------------------------------------------------------------
*   isat_lerror
*
*   This function computes the local error defined as
*
*   eps = ||Rl(phi)- R(phi)||_2 where
*
*   Rl(phi) = R(phi0) + A*(phi-phi0).
*
*   Input:
*   Rphi  - reaction mapping of phi
*   Rphi0 - reaction mapping of phi0
*   A     - mapping gradient matrix
*   phi   - query composition
*   phi0  - initial composition
*
*   Output:
*   eps   - local error
*
*   last update: Feb 19, 2010
*------------------------------------------------------------
*/

double isat_lerror(gsl_vector *Rphi,
                   gsl_vector *Rphi0,
                   gsl_matrix *A,
                   gsl_vector *phi,
                   gsl_vector *phi0)
{
    double eps;
    gsl_vector *Rlphi = NULL;

    /* memory allocation for Rlphi */
    Rlphi = gsl_vector_calloc(phi->size);

    /* Rlphi := R(phi0) + A*(phi-phi0) */
    linear_approx(phi,phi0,Rphi0,A,Rlphi);

    /* Rlphi := Rl(phi) - R(phi) */
    gsl_vector_sub(Rlphi,Rphi);

    /* eps := 2-norm(Rl(phi) - R(phi)) */
    eps = gsl_blas_dnrm2(Rlphi);

    /* releasing allocated memory */
    gsl_vector_free(Rlphi);
    Rlphi = NULL;

    return eps;
}
/*------------------------------------------------------------*/

/*
*------------------------------------------------------------
*   isat4
*
*   This function executes the 4-th version of the
*   In Situ Adaptive Tabulation (ISAT) algorithm.
*
*
*   Input:
*   isat_mem  - ISAT workspace
*   cvode_mem - ODE solver workspace
*   etol      - error tolerance
*   t0        - initial time
*   delta_t   - time step
*   phi       - query composition
*   A         - mapping gradient matrix
*   L         - EOA Cholesky matrix
*
*   Output:
*   Rphi      - reaction mapping
*   success or error
*
*   last update: Nov 3, 2019
*------------------------------------------------------------
*/

int isat4(isat_wrk *isat_mem,
          void *cvode_mem,
          double etol,
          double t0,
          double delta_t,
          gsl_vector *phi,
          gsl_matrix *A,
          gsl_matrix *L,
          gsl_vector *Rphi)
{

    /* CPU clock start counter */
    clock_t cpu_start = clock();

    int flag;
    unsigned int bst_side;
    bst_leaf *end_leaf     = NULL;
    bst_node *end_node     = NULL;

    /****** First call for ISAT algorithm ******/

    /* check if there is no leaf in the binary search tree */
    if( isat_mem->lf == 0 )
    {
        int flag;
        bst_leaf *first_leaf = NULL;

	/* memory allocation for BST first leaf */
        first_leaf = bst_leaf_alloc();
	if ( first_leaf == NULL )
                return GSL_ENOMEM;

        /* direct integration */
        flag = odesolver_reinit(t0,phi,cvode_mem);
        if ( flag != GSL_SUCCESS )
            return flag;

        flag = odesolver(cvode_mem,
                         delta_t,
                         Rphi);
        if ( flag != GSL_SUCCESS )
            return flag;

        /* compute the mapping gradient matrix */
        flag = gradient(cvode_mem,t0,delta_t,phi,Rphi,A);
        if ( flag != GSL_SUCCESS )
            return flag;

        /* compute EOA Cholesky matrix */
        isat_eoa_mtrx(A,etol,L);

        /* define first leaf elements */
        bst_leaf_set(phi,Rphi,A,L,first_leaf);
        isat_mem->root->r_leaf = first_leaf;

        /* update ISAT workspace counters */
        isat_mem->lf++;
        isat_mem->hgt = bst_height(isat_mem->root);

        return GSL_SUCCESS;
    }

    /****** Further calls for ISAT algorithm ******/

    /* search for the near composition in binary search tree */
    if( isat_mem->lf > 1 )
	bst_side = bst_search(isat_mem->root,
                              phi,
                              &end_node,
                              &end_leaf);
    else
    {
        end_node = isat_mem->root;
        end_leaf = isat_mem->root->r_leaf;
        bst_side = BST_RIGHT;
    }

    /* check if the near composition is inside the ellipsoid */
    flag = ell_pt_in(phi,end_leaf->phi,end_leaf->L);
    if( flag == ELL_TRUE )
    {
        /* compute the linear approximation */
        linear_approx(phi,end_leaf->phi,end_leaf->Rphi,end_leaf->A,Rphi);

        /* update ISAT workspace counters */
        isat_mem->rtv++;
	    isat_mem->time_rtv += clock() - cpu_start;

        return GSL_SUCCESS;
    }
    else
    {
        double lerror = 0.0;

        /* performe direct integration */
        flag = odesolver_reinit(t0,phi,cvode_mem);
        if ( flag != GSL_SUCCESS )
            return flag;

        flag = odesolver(cvode_mem,delta_t,Rphi);
        if ( flag != GSL_SUCCESS )
            return flag;

        /* compute ISAT local error */
        lerror = isat_lerror(Rphi,end_leaf->Rphi,
                             end_leaf->A,phi,end_leaf->phi);

        /* check if the local error is greater than etol */
        if( lerror < etol )
        {
            /* grow the ellipsoid */
            ell_pt_modify(phi,end_leaf->phi,end_leaf->L);

            /* update ISAT workspace counters */
            isat_mem->grw++;
	        isat_mem->time_grw += clock() - cpu_start;

            return GSL_SUCCESS;
        }
        else
        {
	    /* check if the maximum number of leaves is excced */
	    if( isat_mem->lf > isat_mem->maxleaves )
	    {
	        /*  update ISAT workspace counters */
	        isat_mem->dev++;
		    isat_mem->time_dev +=  clock() - cpu_start;
		    return GSL_SUCCESS;
	    }

            bst_leaf *new_leaf = NULL;

	    /* memory allocation */
            new_leaf = bst_leaf_alloc();
            if ( new_leaf == NULL )
                return GSL_ENOMEM;

            /* compute the mapping gradient matrix */
            flag = gradient(cvode_mem,t0,delta_t,phi,Rphi,A);
            if ( flag != GSL_SUCCESS )
                return flag;

            /* compute ellipsoid matrix */
            isat_eoa_mtrx(A,etol,L);

            /* define the new leaf elements */
            bst_leaf_set(phi,Rphi,A,L,new_leaf);

            /* check if the binary search tree has more than one leaf */
            if( isat_mem->lf > 1 )
            {
                bst_node *new_node = NULL;

		        /* memory allocation for the new node */
                new_node = bst_node_alloc();
                if ( new_node == NULL )
                    return GSL_ENOMEM;

		        /* define the new node leaves */
                bst_node_set(end_leaf,new_leaf,new_node);

		        /* add the new node to the binary search tree */
                bst_node_add(bst_side,end_node,new_node);
            }
	        else
                bst_node_set(end_leaf,new_leaf,end_node);

            /* update ISAT workspace counters */
            isat_mem->add++;
            isat_mem->lf++;
            isat_mem->nd++;
            isat_mem->hgt = bst_height(isat_mem->root);
	        isat_mem->time_add += clock() - cpu_start;

            return GSL_SUCCESS;
        }
    }
}
/*------------------------------------------------------------*/


#pragma once

#include <algorithm>
#include <array>
#include <cstdint>
#include <initializer_list>
#include <iterator>
#include <sstream>
#include <string>
#include <tuple>

#include <gsl/gsl>

#include "chainerx/axes.h"
#include "chainerx/constant.h"
#include "chainerx/dtype.h"
#include "chainerx/error.h"
#include "chainerx/macro.h"
#include "chainerx/shape.h"

namespace chainerx {

class Strides : public Dims {
    using BaseVector = Dims;

public:
    using const_iterator = BaseVector::const_iterator;
    using const_reverse_iterator = BaseVector::const_reverse_iterator;

    Strides() = default;

    ~Strides() = default;

    // Creates strides for contiguous array.
    Strides(const Shape& shape, Dtype dtype) : Strides{shape, GetItemSize(dtype)} {}
    Strides(const Shape& shape, int64_t item_size);

    // by iterators
    template <typename InputIt>
    Strides(InputIt first, InputIt last) {
        if (std::distance(first, last) > kMaxNdim) {
            throw DimensionError{"too many dimensions: ", std::distance(first, last)};
        }
        insert(begin(), first, last);
    }

    // by gsl:span
    explicit Strides(gsl::span<const int64_t> dims) : Strides{dims.begin(), dims.end()} {}

    // by initializer list
    Strides(std::initializer_list<int64_t> dims) : Strides{dims.begin(), dims.end()} {}

    // copy
    Strides(const Strides&) = default;
    Strides& operator=(const Strides&) = default;

    // move
    Strides(Strides&&) = default;
    Strides& operator=(Strides&&) = default;

    std::string ToString() const;

    int8_t ndim() const noexcept { return gsl::narrow_cast<int8_t>(size()); }

    const int64_t& operator[](int8_t index) const {
        if (!(0 <= index && static_cast<size_t>(index) < size())) {
            throw DimensionError{"Stride index ", index, " out of bounds for strides with ", size(), " size."};
        }
        return this->StackVector::operator[](index);
    }

    int64_t& operator[](int8_t index) {
        if (!(0 <= index && static_cast<size_t>(index) < size())) {
            throw DimensionError{"Stride index ", index, " out of bounds for strides with ", size(), " size."};
        }
        return this->StackVector::operator[](index);
    }

    // span
    gsl::span<const int64_t> span() const { return {*this}; }

    // Rearranges strides in the order specified by the axes.
    //
    // The size of given axes may be fewer than the size of strides.
    // In that case, new strides will be composed by only given axes.
    //
    // It is the caller's responsibility to ensure validity of permutation.
    // If the permutation is invalid, the behavior is undefined.
    Strides Permute(const Axes& axes) const {
        CHAINERX_ASSERT(axes.size() <= size());
        Strides new_strides{};
        for (int8_t axe : axes) {
            new_strides.emplace_back(operator[](axe));
        }
        return new_strides;
    }
};

std::ostream& operator<<(std::ostream& os, const Strides& strides);

void CheckEqual(const Strides& lhs, const Strides& rhs);

// Returns a pair of lower and upper byte offsets to store the data.
// This forumula always holds: lower <= 0 < item_size <= upper
std::tuple<int64_t, int64_t> GetDataRange(const Shape& shape, const Strides& strides, size_t item_size);

}  // namespace chainerx


#ifndef CIMPLE_CIMPLE_AUXILIARY_FUNCTIONS_H
#define CIMPLE_CIMPLE_AUXILIARY_FUNCTIONS_H

/**
 * @brief Macro computes number of an array.
 *
 * Remark: It is defined as a Macro so that it can be used independently of the type of the array
 * e.g. int a[2]; =>N_ELEMS(a)=2
 *      polytope b[3]; =>N_ELEMS(a)=2 (this is an array of 3 polytopes that's why it is not "polytope*")
 */
#define N_ELEMS(x)  (sizeof(x) / sizeof((x)[0]))

#include <stdio.h>
#include <string.h>
#include <pthread.h>
#include <stddef.h>
#include <math.h>
#include <stdbool.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>

/**
 * @brief Generates random numbers with normal distribution
 * @param mu
 * @param sigma
 * @return
 */
double randn (double mu,
              double sigma);

/**
 * @brief Timer simulating one time step in discrete control
 * @param arg
 * @return
 */
void * timer(void * arg);

/**
 * @brief Breaks infinte loop in a pthread
 * @param mtx
 * @return
 */
int needQuit(pthread_mutex_t *mtx);
#endif //CIMPLE_CIMPLE_AUXILIARY_FUNCTIONS_H


/* matrix/gsl_matrix_long_double.h
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman, Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#ifndef __GSL_MATRIX_LONG_DOUBLE_H__
#define __GSL_MATRIX_LONG_DOUBLE_H__

#include <stdlib.h>
#include <gsl/gsl_types.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_check_range.h>
#include <gsl/gsl_vector_long_double.h>

#undef __BEGIN_DECLS
#undef __END_DECLS
#ifdef __cplusplus
# define __BEGIN_DECLS extern "C" {
# define __END_DECLS }
#else
# define __BEGIN_DECLS /* empty */
# define __END_DECLS /* empty */
#endif

__BEGIN_DECLS

typedef struct
{
  size_t size1;
  size_t size2;
  size_t tda;
  long double * data;
  gsl_block_long_double * block;
  int owner;
} gsl_matrix_long_double;

typedef struct
{
  gsl_matrix_long_double matrix;
} _gsl_matrix_long_double_view;

typedef _gsl_matrix_long_double_view gsl_matrix_long_double_view;

typedef struct
{
  gsl_matrix_long_double matrix;
} _gsl_matrix_long_double_const_view;

typedef const _gsl_matrix_long_double_const_view gsl_matrix_long_double_const_view;

/* Allocation */

GSL_EXPORT
gsl_matrix_long_double *
gsl_matrix_long_double_alloc (const size_t n1, const size_t n2);

GSL_EXPORT
gsl_matrix_long_double *
gsl_matrix_long_double_calloc (const size_t n1, const size_t n2);

GSL_EXPORT
gsl_matrix_long_double *
gsl_matrix_long_double_alloc_from_block (gsl_block_long_double * b,
                                   const size_t offset,
                                   const size_t n1,
                                   const size_t n2,
                                   const size_t d2);

GSL_EXPORT
gsl_matrix_long_double *
gsl_matrix_long_double_alloc_from_matrix (gsl_matrix_long_double * m,
                                    const size_t k1,
                                    const size_t k2,
                                    const size_t n1,
                                    const size_t n2);

GSL_EXPORT
gsl_vector_long_double *
gsl_vector_long_double_alloc_row_from_matrix (gsl_matrix_long_double * m,
                                        const size_t i);

GSL_EXPORT
gsl_vector_long_double *
gsl_vector_long_double_alloc_col_from_matrix (gsl_matrix_long_double * m,
                                        const size_t j);

GSL_EXPORT void gsl_matrix_long_double_free (gsl_matrix_long_double * m);

/* Views */

GSL_EXPORT
_gsl_matrix_long_double_view
gsl_matrix_long_double_submatrix (gsl_matrix_long_double * m,
                            const size_t i, const size_t j,
                            const size_t n1, const size_t n2);

GSL_EXPORT
_gsl_vector_long_double_view
gsl_matrix_long_double_row (gsl_matrix_long_double * m, const size_t i);

GSL_EXPORT
_gsl_vector_long_double_view
gsl_matrix_long_double_column (gsl_matrix_long_double * m, const size_t j);

GSL_EXPORT
_gsl_vector_long_double_view
gsl_matrix_long_double_diagonal (gsl_matrix_long_double * m);

GSL_EXPORT
_gsl_vector_long_double_view
gsl_matrix_long_double_subdiagonal (gsl_matrix_long_double * m, const size_t k);

GSL_EXPORT
_gsl_vector_long_double_view
gsl_matrix_long_double_superdiagonal (gsl_matrix_long_double * m, const size_t k);

GSL_EXPORT
_gsl_matrix_long_double_view
gsl_matrix_long_double_view_array (long double * base,
                             const size_t n1,
                             const size_t n2);

GSL_EXPORT
_gsl_matrix_long_double_view
gsl_matrix_long_double_view_array_with_tda (long double * base,
                                      const size_t n1,
                                      const size_t n2,
                                      const size_t tda);

GSL_EXPORT
_gsl_matrix_long_double_view
gsl_matrix_long_double_view_vector (gsl_vector_long_double * v,
                              const size_t n1,
                              const size_t n2);

GSL_EXPORT
_gsl_matrix_long_double_view
gsl_matrix_long_double_view_vector_with_tda (gsl_vector_long_double * v,
                                       const size_t n1,
                                       const size_t n2,
                                       const size_t tda);

GSL_EXPORT
_gsl_matrix_long_double_const_view
gsl_matrix_long_double_const_submatrix (const gsl_matrix_long_double * m,
                                  const size_t i, const size_t j,
                                  const size_t n1, const size_t n2);

GSL_EXPORT
_gsl_vector_long_double_const_view
gsl_matrix_long_double_const_row (const gsl_matrix_long_double * m,
                            const size_t i);

GSL_EXPORT
_gsl_vector_long_double_const_view
gsl_matrix_long_double_const_column (const gsl_matrix_long_double * m,
                               const size_t j);

GSL_EXPORT
_gsl_vector_long_double_const_view
gsl_matrix_long_double_const_diagonal (const gsl_matrix_long_double * m);

GSL_EXPORT
_gsl_vector_long_double_const_view
gsl_matrix_long_double_const_subdiagonal (const gsl_matrix_long_double * m,
                                    const size_t k);

GSL_EXPORT
_gsl_vector_long_double_const_view
gsl_matrix_long_double_const_superdiagonal (const gsl_matrix_long_double * m,
                                      const size_t k);

GSL_EXPORT
_gsl_matrix_long_double_const_view
gsl_matrix_long_double_const_view_array (const long double * base,
                                   const size_t n1,
                                   const size_t n2);

GSL_EXPORT
_gsl_matrix_long_double_const_view
gsl_matrix_long_double_const_view_array_with_tda (const long double * base,
                                            const size_t n1,
                                            const size_t n2,
                                            const size_t tda);

GSL_EXPORT
_gsl_matrix_long_double_const_view
gsl_matrix_long_double_const_view_vector (const gsl_vector_long_double * v,
                                    const size_t n1,
                                    const size_t n2);

GSL_EXPORT
_gsl_matrix_long_double_const_view
gsl_matrix_long_double_const_view_vector_with_tda (const gsl_vector_long_double * v,
                                             const size_t n1,
                                             const size_t n2,
                                             const size_t tda);

/* Operations */

GSL_EXPORT long double   gsl_matrix_long_double_get(const gsl_matrix_long_double * m, const size_t i, const size_t j);
GSL_EXPORT void    gsl_matrix_long_double_set(gsl_matrix_long_double * m, const size_t i, const size_t j, const long double x);

GSL_EXPORT long double * gsl_matrix_long_double_ptr(gsl_matrix_long_double * m, const size_t i, const size_t j);
GSL_EXPORT const long double * gsl_matrix_long_double_const_ptr(const gsl_matrix_long_double * m, const size_t i, const size_t j);

GSL_EXPORT void gsl_matrix_long_double_set_zero (gsl_matrix_long_double * m);
GSL_EXPORT void gsl_matrix_long_double_set_identity (gsl_matrix_long_double * m);
GSL_EXPORT void gsl_matrix_long_double_set_all (gsl_matrix_long_double * m, long double x);

GSL_EXPORT int gsl_matrix_long_double_fread (FILE * stream, gsl_matrix_long_double * m) ;
GSL_EXPORT int gsl_matrix_long_double_fwrite (FILE * stream, const gsl_matrix_long_double * m) ;
GSL_EXPORT int gsl_matrix_long_double_fscanf (FILE * stream, gsl_matrix_long_double * m);
GSL_EXPORT int gsl_matrix_long_double_fprintf (FILE * stream, const gsl_matrix_long_double * m, const char * format);

GSL_EXPORT int gsl_matrix_long_double_memcpy(gsl_matrix_long_double * dest, const gsl_matrix_long_double * src);
GSL_EXPORT int gsl_matrix_long_double_swap(gsl_matrix_long_double * m1, gsl_matrix_long_double * m2);

GSL_EXPORT int gsl_matrix_long_double_swap_rows(gsl_matrix_long_double * m, const size_t i, const size_t j);
GSL_EXPORT int gsl_matrix_long_double_swap_columns(gsl_matrix_long_double * m, const size_t i, const size_t j);
GSL_EXPORT int gsl_matrix_long_double_swap_rowcol(gsl_matrix_long_double * m, const size_t i, const size_t j);
GSL_EXPORT int gsl_matrix_long_double_transpose (gsl_matrix_long_double * m);
GSL_EXPORT int gsl_matrix_long_double_transpose_memcpy (gsl_matrix_long_double * dest, const gsl_matrix_long_double * src);

GSL_EXPORT long double gsl_matrix_long_double_max (const gsl_matrix_long_double * m);
GSL_EXPORT long double gsl_matrix_long_double_min (const gsl_matrix_long_double * m);
GSL_EXPORT void gsl_matrix_long_double_minmax (const gsl_matrix_long_double * m, long double * min_out, long double * max_out);

GSL_EXPORT void gsl_matrix_long_double_max_index (const gsl_matrix_long_double * m, size_t * imax, size_t *jmax);
GSL_EXPORT void gsl_matrix_long_double_min_index (const gsl_matrix_long_double * m, size_t * imin, size_t *jmin);
GSL_EXPORT void gsl_matrix_long_double_minmax_index (const gsl_matrix_long_double * m, size_t * imin, size_t * jmin, size_t * imax, size_t * jmax);

GSL_EXPORT int gsl_matrix_long_double_isnull (const gsl_matrix_long_double * m);

GSL_EXPORT int gsl_matrix_long_double_add (gsl_matrix_long_double * a, const gsl_matrix_long_double * b);
GSL_EXPORT int gsl_matrix_long_double_sub (gsl_matrix_long_double * a, const gsl_matrix_long_double * b);
GSL_EXPORT int gsl_matrix_long_double_mul_elements (gsl_matrix_long_double * a, const gsl_matrix_long_double * b);
GSL_EXPORT int gsl_matrix_long_double_div_elements (gsl_matrix_long_double * a, const gsl_matrix_long_double * b);
GSL_EXPORT int gsl_matrix_long_double_scale (gsl_matrix_long_double * a, const double x);
GSL_EXPORT int gsl_matrix_long_double_add_constant (gsl_matrix_long_double * a, const double x);
GSL_EXPORT int gsl_matrix_long_double_add_diagonal (gsl_matrix_long_double * a, const double x);

/***********************************************************************/
/* The functions below are obsolete                                    */
/***********************************************************************/
GSL_EXPORT int gsl_matrix_long_double_get_row(gsl_vector_long_double * v, const gsl_matrix_long_double * m, const size_t i);
GSL_EXPORT int gsl_matrix_long_double_get_col(gsl_vector_long_double * v, const gsl_matrix_long_double * m, const size_t j);
GSL_EXPORT int gsl_matrix_long_double_set_row(gsl_matrix_long_double * m, const size_t i, const gsl_vector_long_double * v);
GSL_EXPORT int gsl_matrix_long_double_set_col(gsl_matrix_long_double * m, const size_t j, const gsl_vector_long_double * v);

/* inline functions if you are using GCC */

#ifdef HAVE_INLINE
extern inline
long double
gsl_matrix_long_double_get(const gsl_matrix_long_double * m, const size_t i, const size_t j)
{
#if GSL_RANGE_CHECK
  if (i >= m->size1)
    {
      GSL_ERROR_VAL("first index out of range", GSL_EINVAL, 0) ;
    }
  else if (j >= m->size2)
    {
      GSL_ERROR_VAL("second index out of range", GSL_EINVAL, 0) ;
    }
#endif
  return m->data[i * m->tda + j] ;
}

extern inline
void
gsl_matrix_long_double_set(gsl_matrix_long_double * m, const size_t i, const size_t j, const long double x)
{
#if GSL_RANGE_CHECK
  if (i >= m->size1)
    {
      GSL_ERROR_VOID("first index out of range", GSL_EINVAL) ;
    }
  else if (j >= m->size2)
    {
      GSL_ERROR_VOID("second index out of range", GSL_EINVAL) ;
    }
#endif
  m->data[i * m->tda + j] = x ;
}

extern inline
long double *
gsl_matrix_long_double_ptr(gsl_matrix_long_double * m, const size_t i, const size_t j)
{
#if GSL_RANGE_CHECK
  if (i >= m->size1)
    {
      GSL_ERROR_NULL("first index out of range", GSL_EINVAL) ;
    }
  else if (j >= m->size2)
    {
      GSL_ERROR_NULL("second index out of range", GSL_EINVAL) ;
    }
#endif
  return (long double *) (m->data + (i * m->tda + j)) ;
}

extern inline
const long double *
gsl_matrix_long_double_const_ptr(const gsl_matrix_long_double * m, const size_t i, const size_t j)
{
#if GSL_RANGE_CHECK
  if (i >= m->size1)
    {
      GSL_ERROR_NULL("first index out of range", GSL_EINVAL) ;
    }
  else if (j >= m->size2)
    {
      GSL_ERROR_NULL("second index out of range", GSL_EINVAL) ;
    }
#endif
  return (const long double *) (m->data + (i * m->tda + j)) ;
}

#endif

__END_DECLS

#endif /* __GSL_MATRIX_LONG_DOUBLE_H__ */