/*
 * @BEGIN LICENSE
 *
 * sointegrals by Psi4 Developer, a plugin to:
 *
 * Psi4: an open-source quantum chemistry software package
 *
 * Copyright (c) 2007-2019 The Psi4 Developers.
 *
 * The copyrights for code used from other parties are included in
 * the corresponding files.
 *
 * This file is part of Psi4.
 *
 * Psi4 is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * Psi4 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 Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License along
 * with Psi4; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * @END LICENSE
 */

#include "psi4/psi4-dec.h"
#include "psi4/libpsi4util/PsiOutStream.h"
#include "psi4/liboptions/liboptions.h"
#include "psi4/libmints/wavefunction.h"
#include "psi4/libmints/molecule.h"
#include "psi4/libmints/basisset.h"
#include "psi4/libmints/matrix.h"
#include "psi4/libmints/mintshelper.h"
#include "psi4/libmints/factory.h"
#include "psi4/libmints/integral.h"
#include "psi4/libpsi4util/process.h"

#include <math.h>
#include <stdlib.h>
#include <vector>
#include <stdio.h>
#include <cstddef>
#include <time.h>
#include <malloc.h>
#include <fstream>
#include <string>

#ifdef _OPENMP
#include "omp.h"
#endif

using namespace std;

namespace psi{ namespace mointegrals {

extern "C" PSI_API
int read_options(std::string name, Options& options)
{
    if (name == "MOINTEGRALS"|| options.read_globals()) {
        /*- The amount of information printed
            to the output file -*/
        options.add_int("PRINT", 1);
        /*- Whether to compute two-electron integrals -*/
        options.add_bool("DO_TEI", true);
    }

    return true;
}


extern "C" PSI_API
SharedWavefunction mointegrals(SharedWavefunction ref_wfn, Options& options)
{
    int print = options.get_int("PRINT");
    int doTei = options.get_bool("DO_TEI");

    std::shared_ptr<BasisSet> aoBasis = ref_wfn->basisset();
    int nbf = {aoBasis->nbf()};
    time_t start, end;


    ifstream num_ao_file_read("num_ao", ios::binary);
    unsigned long long int num_ao = 0;
    num_ao_file_read.read((char*) &num_ao, sizeof(long long int));
    num_ao_file_read.close();

    // Now build the MO two-electron integrals
    SharedMatrix C_psi4 = ref_wfn->Ca(); // MO coefficients obtatined from input wfn
    int active_size = C_psi4->colspi()[0]; // # of Active MO


    start = time(NULL);

    // Using 2D array to store the MO coefficients
    double C[nbf][active_size];

    for (int i = 0; i < active_size; i++){
        for (int j = 0; j < nbf; j++){
            C[j][i] = C_psi4->get(j, i);
        }
    }

    // Here we divide the AO to MO transform into multiple task
    unsigned int array_size = active_size * active_size * active_size * active_size;

    double *mo_eri;
    mo_eri = (double *) calloc(array_size, sizeof(double));

    ifstream ao_eri_value_file("ao_eri_value", ios::binary);
    ifstream ao_index1_file("ao_eri_index1", ios::binary);
    ifstream ao_index2_file("ao_eri_index2", ios::binary);
    ifstream ao_index3_file("ao_eri_index3", ios::binary);
    ifstream ao_index4_file("ao_eri_index4", ios::binary);

    if (!ao_eri_value_file){printf("\n\tReading ao_eri_value file failed!!\n\n");}
    if (!ao_index1_file){printf("\n\tReading ao index 1 failed!!\n\n");}
    if (!ao_index2_file){printf("\n\tReading ao index 2 failed!!\n\n");}
    if (!ao_index3_file){printf("\n\tReading ao index 3 failed!!\n\n");}
    if (!ao_index4_file){printf("\n\tReading ao index 4 failed!!\n\n");}

    int* ao_index_1;
    int* ao_index_2;
    int* ao_index_3;
    int* ao_index_4;
    double* ao_eri;

    unsigned int cycle = 1;

    while (num_ao  > 4e8 ){

        unsigned long long int num_ao_part = 3e8;
        num_ao -= num_ao_part;

        ao_index_1 = (int*) calloc(num_ao_part, sizeof(int));
        ao_index_2 = (int*) calloc(num_ao_part, sizeof(int));
        ao_index_3 = (int*) calloc(num_ao_part, sizeof(int));
        ao_index_4 = (int*) calloc(num_ao_part, sizeof(int));
        ao_eri = (double*) calloc(num_ao_part, sizeof(double));

        // Check the ptr is successfully built
        if (ao_eri == NULL){
            printf("\tThings go wrong at malloc ao_eri\n\n");
            fflush(stdout);
        }

        if (ao_index_1 == NULL) {
            printf("\tThings go wrong at malloc ao_index_1\n\n");
            fflush(stdout);
        }
        if (ao_index_2 == NULL) {
            printf("\tThings go wrong at malloc ao_index_2\n\n");
            fflush(stdout);
        }

        if (ao_index_3 == NULL) {
            printf("\tThings go wrong at malloc ao_index_3\n\n");
            fflush(stdout);
        }

        if (ao_index_4 == NULL) {
            printf("\tThings go wrong at malloc ao_index_4\n\n");
            fflush(stdout);
        }

        // Start read the information of ao_eri from files
        // AO eri value


        for (int i = 0; i < num_ao_part; i++){
            ao_eri_value_file.read((char *) &ao_eri[i], sizeof(double));
        }

        for (int i = 0; i < num_ao_part; i++){
            ao_index1_file.read((char *) &ao_index_1[i], sizeof(int));
        }

        // AO index 2
        for (int i = 0; i < num_ao_part; i++){
            ao_index2_file.read((char *) &ao_index_2[i], sizeof(int));
        }


        // AO index 3
        for (int i = 0; i < num_ao_part; i++){
            ao_index3_file.read((char *) &ao_index_3[i], sizeof(int));
        }

        //AO index 4
        for (int i = 0; i < num_ao_part; i++){
            ao_index4_file.read((char *) &ao_index_4[i], sizeof(int));
        }

        // Start compute the MO two electron integral
        #pragma omp parallel for collapse(4) schedule(static) reduction(+:mo_eri[:array_size])
        for (int i = 0; i < active_size; i++){
            for (int j = 0; j < active_size; j++){
                for (int k = 0; k < active_size; k++){
                    for (int l = 0; l < active_size; l++){
                        int ij = j * (j + 1) / 2 + i;
                        int kl = l * (l + 1) / 2 + k;

                        if ((i <= j) and (k <= l) and (ij <= kl)){
                            double coupling = 0.0;
                            for (int N = 0; N < num_ao_part; N++){
                                int p = ao_index_1[N];
                                int q = ao_index_2[N];
                                int r = ao_index_3[N];
                                int s = ao_index_4[N];
                                coupling += C[p][i] * C[q][j] * C[r][k] * C[s][l] * ao_eri[N];
                            }
                            mo_eri[(l * active_size * active_size * active_size) + (k * active_size * active_size) + (j * active_size) + i] += coupling;
                        }
                    }
                }
            }
        }

        free(ao_eri);
        free(ao_index_1);
        free(ao_index_2);
        free(ao_index_3);
        free(ao_index_4);

        cycle += 1;
        fflush(stdout);

    }

    // Last part
    ao_index_1 = (int*) calloc(num_ao, sizeof(int));
    ao_index_2 = (int*) calloc(num_ao, sizeof(int));
    ao_index_3 = (int*) calloc(num_ao, sizeof(int));
    ao_index_4 = (int*) calloc(num_ao, sizeof(int));
    ao_eri = (double*) calloc(num_ao, sizeof(double));

    // Check the ptr is successfully built
    if (ao_eri == NULL){
        printf("\tThings go wrong at malloc ao_eri\n\n");
        fflush(stdout);
    }

    if (ao_index_1 == NULL) {
        printf("\tThings go wrong at malloc ao_index_1\n\n");
        fflush(stdout);
    }

    if (ao_index_2 == NULL) {
        printf("\tThings go wrong at malloc ao_index_2\n\n");
        fflush(stdout);
    }

    if (ao_index_3 == NULL) {
        printf("\tThings go wrong at malloc ao_index_3\n\n");
        fflush(stdout);
    }

    if (ao_index_4 == NULL) {
        printf("\tThings go wrong at malloc ao_index_4\n\n");
        fflush(stdout);
    }

    // Start read the information of ao_eri from files
    // AO eri
    for (int i = 0; i < num_ao; i++){
        ao_eri_value_file.read((char *) &ao_eri[i], sizeof(double));
    }

    ao_eri_value_file.close();

    // AO index 1
    for (int i = 0; i < num_ao; i++){
        ao_index1_file.read((char *) &ao_index_1[i], sizeof(int));
    }

    ao_index1_file.close();

    // AO index 2
    for (int i = 0; i < num_ao; i++){
        ao_index2_file.read((char *) &ao_index_2[i], sizeof(int));
    }

    ao_index2_file.close();


    // AO index 3
    for (int i = 0; i < num_ao; i++){
        ao_index3_file.read((char *) &ao_index_3[i], sizeof(int));
    }

    ao_index3_file.close();

    //AO index 4
    for (int i = 0; i < num_ao; i++){
        ao_index4_file.read((char *) &ao_index_4[i], sizeof(int));
    }

    ao_index4_file.close();

    // Start compute the MO two electron integrals
    #pragma omp parallel for collapse(4) schedule(static) reduction(+:mo_eri[:array_size])
    for (int i = 0; i < active_size; i++){
        for (int j = 0; j < active_size; j++){
            for (int k = 0; k < active_size; k++){
                for (int l = 0; l < active_size; l++){
                     int ij = j * (j + 1) / 2 + i;
                     int kl = l * (l + 1) / 2 + k;

                     if (i <= j and k <= l and ij <= kl){
                         double coupling = 0.0;
                         for (int N = 0; N < num_ao; N++){
                             int p = ao_index_1[N];
                             int q = ao_index_2[N];
                             int r = ao_index_3[N];
                             int s = ao_index_4[N];
                             coupling += C[p][i] * C[q][j] * C[r][k] * C[s][l] * ao_eri[N];
                         }
                         mo_eri[(l * active_size * active_size * active_size) + (k * active_size * active_size) + (j * active_size) + i] += coupling;
                     }

                }
            }
        }
    }

    free(ao_eri);
    free(ao_index_1);
    free(ao_index_2);
    free(ao_index_3);
    free(ao_index_4);

    end = time(NULL);

    double diff_time_mo_sec = difftime(end, start);
    double diff_time_mo;

    printf("\tCalculating MO two-electron integrals is Done!!\n");
    printf("\tTotal Cycles: %d\n", cycle);

    if (diff_time_mo_sec > 3600){
        diff_time_mo = diff_time_mo_sec/3600;
        printf("\tTotal time taken for calculating MO two-electron integrals: %.1f hr\n", diff_time_mo);
    }
    else{
        diff_time_mo = diff_time_mo_sec;
        printf("\tTotal time taken for calculating MO two-electron integrals: %.1f sec\n", diff_time_mo);
    }

    fflush(stdout);

    string mo_eri_filename = "mo_eri";
    ofstream mo_eri_file(mo_eri_filename, ios::binary);

    #pragma single
    for (int i = 0; i < active_size; i++){
        for (int j = 0; j < active_size; j++){
            for (int k = 0; k < active_size; k++){
                for (int l = 0; l < active_size; l++){
                    // psi::outfile->Printf("\t(%2d %2d | %2d %2d) =:\t%20.10lf\n", i, j, k, l, mo_eri[(l * active_size * active_size * active_size) + (k * active_size * active_size) + (j * active_size) + i]);
                    mo_eri_file.write((char*) &mo_eri[(l * active_size * active_size * active_size) + (k * active_size * active_size) + (j * active_size) + i], sizeof(double));
                }
            }
        }
    }

    mo_eri_file.close();

    free(mo_eri);
    fflush(stdout);

    return ref_wfn;
}

}} // End namespaces