#include <stdio.h>
#include <stdlib.h>

extern void dcopy_(const int *n, double *A, const int *inca, double *B, const int *incb);
extern void dgemm_(const char *transa, const char *transb, const int *m, const int *n, const int *k, double *alpha, double *A, const int *lda, double *B, const int *ldb, double *beta, double *C, const int *ldc);
extern void dlacpy_(const char *uplo, const int *m, const int *n, double *A, const int *lda, double *B, const int *ldb);
extern void dlaset_(const char *uplo, const int *m, const int *n, double *alpha, double *beta, double *A, const int *lda);
extern double dlamch_(const char *a);
extern double dlange_(const char *norm, const int *m, const int *n, double *A, const int *lda, double *work);
extern void dgtsv_(const int *n, const int *nrhs, double *DL, double *D, double *DU, double *BF, const int *ldb, int *info);
extern void dtrcon_(const char *norm, const char *uplo, const char *diag, const int *n, double *A, const int *lda, double *rcond, double *work, int *iwork, int *info);
extern void dtrsm_(const char *side, const char *uplo, const char *transa, const char *diag, const int *m, const int *n, double *alpha, double *A, const int *lda, double *B, const int *ldb);
extern void dlarnv_(const int *idist, const int *iseed, const int *n, double *X);

int ISEED[4] = {9, 10, 11, 12};

void test_gtsv_fact(int n, int nrhs, double *DL, double *D, double *DU) {
    int info = 0;
    double *B = (double*)malloc(n * nrhs * sizeof(double));

    const int idist = 2;
    const int n1 = n - 1;

    for (int j = 0; j < nrhs; j++)
         dlarnv_(&idist, ISEED, &n, &B[j * n]);

    double *DF = (double*)malloc(n * sizeof(double));
    double *DLF = (double*)malloc((n - 1) * sizeof(double));
    double *DUF = (double*)malloc((n - 1) * sizeof(double));
    double *BF = (double*)malloc(n * nrhs * sizeof(double));

    const int ione = 1;
    dcopy_(&n, D, &ione, DF, &ione); 
    dcopy_(&n1, DL, &ione, DLF, &ione); 
    dcopy_(&n1, DU, &ione, DUF, &ione); 
    dlacpy_("Full", &n, &nrhs, B, &n, BF, &n);

    dgtsv_(&n, &nrhs, DLF, DF, DUF, BF, &n, &info);    
    if (info != 0)
        printf("gtsv info error for n = %d, nrhs = %d, info = %d\n", n, nrhs, info);

    double *U = (double*)malloc(n * n * sizeof(double));
    double zero = 0.0;
    dlaset_("Full", &n, &n, &zero, &zero, U, &n);
    
    for (int i = 0; i < n; i++) {
        U[i + n * i] = DF[i];
        if (i < n - 1) U[i + n * (i + 1)] = DUF[i];
        if (i < n - 2) U[i + n * (i + 2)] = DLF[i];
    }

    double *work = (double*)malloc(3 * n * sizeof(double));
    int *iwork = (int*)malloc(n * sizeof(int));

    double rcond;
    // Compute condition number
    dtrcon_("1", "Upper", "Non-unit", &n, U, &n, &rcond, work, iwork, &info);
    
    double *A = (double*)malloc(n * n * sizeof(double));
    dlaset_("Full", &n, &n, &zero, &zero, A, &n);

    // Copy arrays D, DL, DU into A
    for (int i = 0; i < n; i++) {
        A[i + n * i] = D[i];
        if (i < n - 1) A[i + n * (i + 1)] = DU[i];
        if (i < n - 1) A[i + 1 + n * i] = DL[i];
    }


    double *PL = (double*)malloc(n * n * sizeof(double));
    dlaset_("Full", &n, &n, &zero, &zero, PL, &n);
    dlacpy_("Full", &n, &n, A, &n, PL, &n);

    double fone = 1.0;

    // Compute PL = A * U^(-1)
    dtrsm_("R", "U", "N", "N", &n, &n, &fone, U, &n, PL, &n);

    double mone = -1.0;
    double *A_DIFF = (double*)malloc(n * n * sizeof(double));
    dlaset_("Full", &n, &n, &zero, &zero, A_DIFF, &n);

    dlacpy_("Full", &n, &n, A, &n, A_DIFF, &n);

    // Multiply back PL * U and compare with A
    dgemm_("N", "N", &n, &n, &n, &fone, PL, &n, U, &n, &mone, A_DIFF, &n);
    double normPLU_A = dlange_("F", &n, &n, A_DIFF, &n, work);

    double eps = dlamch_("Eps");
    printf("n = %03d nrhs = %03d rcondU = %e |PLU - A| = %e |PLU - A| / eps = %lf\n", n, nrhs, rcond, normPLU_A, normPLU_A / eps);

    free(B);

    free(DF);
    free(DLF);
    free(DUF);
    free(BF);

    free(A);
    free(U);
    free(PL);
    free(A_DIFF);

    free(work);
    free(iwork);
}

int main()
{
    for (int n = 1; n <= 100; n++) {
        double *D = (double*)malloc(n * sizeof(double));
        double *DL = (double*)malloc((n - 1) * sizeof(double));
        double *DU = (double*)malloc((n - 1) * sizeof(double));
        const int idist = 2;
        const int n1 = n - 1;
        dlarnv_(&idist, ISEED, &n, D);
        dlarnv_(&idist, ISEED, &n1, DL);
        dlarnv_(&idist, ISEED, &n1, DU);
  
        for (int nrhs = 1; nrhs <= 100; nrhs++) {
            test_gtsv_fact(n, nrhs, DL, D, DU);
        }
        printf("------------------------------------\n");
        free(D);
        free(DL);
        free(DU);
    }
    return 0;
}