#include <time.h>

#include <lapack.h>

#include <lapacke.h>
#include <mpi.h>
#include <petscconf.h>
#include <petscerror.h>
#include <petscmat.h>
#include <petscoptions.h>
#include <petscsys.h>
#include <petscsystypes.h>
#include <petsctime.h>
#include <petscvec.h>
#include <petscviewer.h>
#include <petscviewerhdf5.h>

#include <petscdm.h>
#include <petscdmda.h>
#include <petscksp.h>
#include <petscviewertypes.h>
#include <stdio.h>
#include <stdlib.h>

static char help[] = "Example PETSc program\n\n";

void swap(double *a, double *b) {
    double t = *a;
    *a = *b;
    *b = t;
}

void print_all_hostnames() {
    int rank, size, name_len;
    char proc_name[MPI_MAX_PROCESSOR_NAME];

    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    MPI_Get_processor_name(proc_name, &name_len);

    char *all_names = NULL;
    if (rank == 0)
        all_names = malloc(size * MPI_MAX_PROCESSOR_NAME * sizeof(char));

    MPI_Gather(proc_name, MPI_MAX_PROCESSOR_NAME, MPI_CHAR, all_names, MPI_MAX_PROCESSOR_NAME, MPI_CHAR, 0,
               MPI_COMM_WORLD);

    if (rank == 0) {
        for (int i = 0; i < size; i++) {
            printf("Rank %d: %s\n", i, &all_names[i * MPI_MAX_PROCESSOR_NAME]);
        }
        free(all_names);
    }
}

PetscErrorCode ArnoldiIteration(Mat A, Vec b, PetscInt n, PetscInt m, Vec *Q, double *h);

int main(int argc, char **argv) {
    PetscFunctionBeginUser;
    PetscInitialize(&argc, &argv, (char *)0, help);

    print_all_hostnames();

    PetscInt n, l;
    char matrix_name[PETSC_MAX_PATH_LEN];

    PetscBool flg;
    PetscOptionsGetInt(NULL, NULL, "-n", &n, &flg);
    if (!flg)
        n = -1;

    PetscOptionsGetInt(NULL, NULL, "-l", &l, &flg);
    if (!flg)
        l = 10;

    PetscOptionsGetString(NULL, NULL, "-matrix_path", matrix_name, sizeof(matrix_name), &flg);
    if (!flg)
        snprintf(matrix_name, sizeof(matrix_name), "./matrices/laplacian/laplacian-discretization-3d.mat");

    int rank, total;
    MPI_Comm_rank(PETSC_COMM_WORLD, &rank);
    MPI_Comm_size(PETSC_COMM_WORLD, &total);

    char hostname[64];
    PetscCall(PetscGetHostName(hostname, sizeof(hostname)));
    if (rank == 0)
        PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Running on %s, rank %d of %d\n", hostname, rank, total));

    PetscLogDouble load_start_time;
    PetscCall(PetscTime(&load_start_time));

    Mat A;
    MatCreate(PETSC_COMM_WORLD, &A);
    PetscViewer v;
    PetscCall(PetscViewerCreate(PETSC_COMM_WORLD, &v));
    PetscCall(PetscViewerSetType(v, PETSCVIEWERHDF5));
    PetscCall(PetscViewerPushFormat(v, PETSC_VIEWER_HDF5_MAT));
    PetscCall(PetscViewerSetFromOptions(v));
    PetscCall(PetscViewerFileSetMode(v, FILE_MODE_READ));
    PetscCall(PetscViewerFileSetName(v, matrix_name));

    PetscCall(MatSetOptionsPrefix(A, "a_"));
    PetscCall(PetscObjectSetName((PetscObject)A, "A"));

    PetscCall(MatLoad(A, v));

    // Print matrix info
    {
        PetscInt m, n;
        MatGetSize(A, &m, &n);
        PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Matrix loaded from %s\n", matrix_name));
        PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Matrix size: %d x %d\n", m, n));
    }

    MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);

    PetscLogDouble load_start_end;
    PetscCall(PetscTime(&load_start_end));

    PetscLogDouble load_time = load_start_end - load_start_time;
    if (rank == 0)
        PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Load time: %f seconds\n", load_time));

    if (n == -1) {
        MatGetSize(A, &n, NULL);
    }

    Vec b;
    VecCreate(PETSC_COMM_WORLD, &b);
    VecSetSizes(b, PETSC_DECIDE, n);
    VecSetType(b, VECMPI);

    VecSet(b, 1.0);

    // seed random number generator
    // srand((unsigned int)time(NULL));

    // fill b with random values
    // for (PetscInt i = 0; i < n; i++) {
    //     double val = (double)rand() / RAND_MAX;
    //     VecSetValue(b, i, val, INSERT_VALUES);
    // }

    // VecAssemblyBegin(b);
    // VecAssemblyEnd(b);

    // VecSetValue(b, 0, 1.0, INSERT_VALUES);

    // MatView(A, PETSC_VIEWER_STDOUT_WORLD);
    // VecView(b, PETSC_VIEWER_STDOUT_WORLD);

    // PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Allocating memory for Krylov subspace basis\n"));

    Vec *Q;
    PetscMalloc1(l, &Q);

    for (PetscInt i = 0; i < l; i++) {
        VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE, n, &Q[i]);
    }

    // PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Constructing Hessenberg matrix\n"));

    // Mat H;
    // PetscCall(MatCreate(PETSC_COMM_SELF, &H));
    // PetscCall(PetscObjectSetName((PetscObject)H, "hessenberg"));
    // PetscCall(MatSetSizes(H, PETSC_DECIDE, PETSC_DECIDE, l + 1, l));
    // PetscCall(MatSetType(H, MATDENSE));

    // PetscCall(MatCreate(PETSC_COMM_WORLD, &H));

    // MatSetType(H, MATMPIAIJ);

    double *h = (double *)malloc((l + 1) * l * sizeof(double));

    for (PetscInt ii = 0; ii < (l + 1) * l; ii++) {
        h[ii] = 0.0;
    }

    // PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Starting iteration\n"));

    PetscLogDouble arnoldi_start_time;
    PetscCall(PetscTime(&arnoldi_start_time));
    {
        PetscCall(ArnoldiIteration(A, b, l, n, Q, h));
    }
    PetscLogDouble arnoldi_end_time;
    PetscCall(PetscTime(&arnoldi_end_time));
    PetscLogDouble arnoldi_time = arnoldi_end_time - arnoldi_start_time;
    if (rank == 0)
        PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Arnoldi time: %f seconds\n", arnoldi_time));

    // PetscCall(MatSetValue(H, 2, 3, -1, INSERT_VALUES));
    // PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
    // PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));

    // PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Done\n"));

    // int rank;
    // PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));

    if (rank == 0) {
        // print Hessenberg matrix
        printf("H = \n");
        for (int i = 0; i < l + 1; i++) {
            for (int j = 0; j < l; j++) {
                printf("%.2f\t", h[i * (l + 1) + j]);
            }
            printf("\n");
        }

        double *wr = (double *)malloc(l * sizeof(double));
        double *wi = (double *)malloc(l * sizeof(double));
        double *z = (double *)malloc(l * l * sizeof(double));
        double *work = (double *)malloc(3 * l * sizeof(double));
        int info;

        PetscLogDouble lapack_start_time;
        PetscCall(PetscTime(&lapack_start_time));
        {
            // call LAPACK function "DHSEQR" to compute the eigenvalues of the Hessenberg matrix
            LAPACKE_dhseqr(LAPACK_ROW_MAJOR, 'E', 'I', l, 1, l, h, l + 1, wr, wi, z, l);
        }

        PetscLogDouble lapack_end_time;
        PetscCall(PetscTime(&lapack_end_time));
        PetscLogDouble lapack_time = lapack_end_time - lapack_start_time;
        PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] LAPACK time: %f seconds\n", lapack_time));

        // print Hessenberg matrix
        printf("H = \n");
        for (int i = 0; i < l + 1; i++) {
            for (int j = 0; j < l; j++) {
                printf("%.2f\t", h[i * (l + 1) + j]);
            }
            printf("\n");
        }

        // sort eigenvalues
        for (int i = 0; i < l; i++) {
            for (int j = i + 1; j < l; j++) {
                if (wr[i] > wr[j]) {
                    swap(&wr[i], &wr[j]);
                    swap(&wi[i], &wi[j]);
                    for (int k = 0; k < l; k++) {
                        swap(&z[i * l + k], &z[j * l + k]);
                    }
                }
            }
        }

        // print eigenvalues
        printf("Eigenvalues = \n");
        for (int i = 0; i < l; i++) {
            printf("%.3f + %.3f i\n", wr[i], wi[i]);
        }
    }

    // print eigenvectors
    // printf("Eigenvectors = \n");
    // for (int i = 0; i < l; i++) {
    //     for (int j = 0; j < l; j++) {
    //         printf("%f ", z[i * l + j]);
    //     }
    //     printf("\n");
    // }

    PetscCall(MatDestroy(&A));
    PetscCall(VecDestroy(&b));
    // PetscCall(MatDestroy(&H));
    PetscFinalize();
    return 0;
}

PetscErrorCode ArnoldiIteration(Mat A, Vec b, PetscInt n, PetscInt m, Vec *Q, double *h) {
    PetscFunctionBeginUser;

    int rank;
    MPI_Comm_rank(PETSC_COMM_WORLD, &rank);

    PetscScalar eps = 1e-12;

    Vec q;

    if (rank == 0)
        printf("n = %d, m = %d\n", n, m);

    /*

    eps = 1e-12
    h = np.zeros((n + 1, n))
    Q = np.zeros((A.shape[0], n + 1))
    # Normalize the input vector
    Q[:, 0] = b / np.linalg.norm(b, 2)  # Use it as the first Krylov vector
    for k in range(1, n + 1):
        v = np.dot(A, Q[:, k - 1])  # Generate a new candidate vector
        for j in range(k):  # Subtract the projections on previous vectors
            h[j, k - 1] = np.dot(Q[:, j].conj(), v)
            v = v - h[j, k - 1] * Q[:, j]
        h[k, k - 1] = np.linalg.norm(v, 2)
        if h[k, k - 1] > eps:  # Add the produced vector to the list, unless
            Q[:, k] = v / h[k, k - 1]
        else:  # If that happens, stop iterating.
            return Q, h


    */

    PetscCall(VecDuplicate(b, &q));
    PetscCall(VecCopy(b, q));
    PetscCall(VecNormalize(q, NULL));

    Q[0] = q;

    for (PetscInt k = 1; k < n + 1; k++) {

        if (rank == 0)
            PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Iteration %d\n", k));

        Vec v;
        PetscCall(VecDuplicate(b, &v));

        // v = A q_k
        PetscCall(MatMult(A, Q[k - 1], v));

        // Reorthogonalization using modified Gram-Schmidt
        // fill all rows from 0 to k - 1 at column k - 1
        for (PetscInt j = 0; j < k; j++) { // anche solo 3

            PetscScalar h_ij;
            // h_(j, k - 1) = q_j . v
            PetscCall(VecDot(Q[j], v, &h_ij));

            if (rank == 0)
                printf("h[%d, %d] = %f\n", j, k - 1, h_ij);

            h[j * (n + 1) + k - 1] = h_ij;

            // v -= h_ij * q_j
            PetscCall(VecAXPY(v, -h_ij, Q[j]));
        }

        // Normalize:

        // v = v / ||v||
        // h_ij = ||v||

        PetscScalar h_ij;
        PetscCall(VecNorm(v, NORM_2, &h_ij));

        if (rank == 0)
            printf("h[%d, %d] = %f\n", k, k - 1, h_ij);

        h[k * (n + 1) + k - 1] = h_ij;

        // Check for convergence
        if (h_ij > eps) {
            PetscCall(VecNormalize(v, NULL));
            Q[k] = v;
        } else {
            PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[Arnoldi] Early breakdown"));
            break;
        }
    }

    PetscFunctionReturn(PETSC_SUCCESS);
}