// Copyright 2019 Xanadu Quantum Technologies Inc.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "rlhafnian.h"
#include <stdio.h>
void
evals(double z[], double complex vals[], int n)
{
    /*
     * Calculates the eigenvalues of the real n x n matrix z
     * returns the (complex) eigenvalues in the array vals
     * For the calculation of the eigenvalues it uses lapacke
     */
    lapack_int info;

    char jobvs = 'N';
    char sort  = 'N';

    lapack_int lda  = n;
    lapack_int sdim = 0;
    lapack_int ldvs = n;
    double     wr[n];
    double     wi[n];

    info = LAPACKE_dgees(LAPACK_ROW_MAJOR, jobvs, sort,
                                NULL, n, &(z[0]), lda,
                                &sdim, &(wr[0]), &(wi[0]), NULL, ldvs);
    assert(info == 0);
    for (int i = 0; i < n; i++)
    {
        vals[i] = wr[i] + I * wi[i];
    }
}


void
powtrace(double z[], int n, double traces[], int l)
{
    /*
     * given a complex matrix z of dimensions n x n
     * it calculates traces[k] = tr(z^j) for 1 <= j <= l
     * It does it by first finding the eigenvalues of the matrix z
     * using the orutine evals
     */
    double complex vals[n];
    double complex pvals[n];

    evals(z, vals, n);
    int            i, j;
    double complex sum;
    for (j = 0; j < n; j++)
    {
        pvals[j] = vals[j];
    }
    for (i = 0; i < l; i++)
    {
        sum = 0.0;
        for (j = 0; j < n; j++)
        {
            sum += pvals[j];
        }
        traces[i] = creal(sum);
        for (j = 0; j < n; j++)
        {
            pvals[j] = pvals[j] * vals[j];
        }
    }
}


void
dec2bin(char *dst, unsigned long long int x, sint len)
{
    /*
     * Convert decimal number in  x to character vector dst of length len
     * representing binary number
     */

    char i;      // this variable cannot be unsigned

    for (i = len - 1; i >= 0; --i)
    {
        *dst++ = x >> i & 1;
    }
}


sint
find2(char *dst, sint len, sint *pos)
{
    /* Given a string of length len
     * it finds in which positions it has a one
     * and stores its position i, as 2*i and 2*i+1 in consecutive slots
     * of the array pos.
     * It also returns (twice) the number of ones in array dst
     */
    sint sum = 0;
    sint j   = 0;

    for (sint i = 0; i < len; i++)
    {
        if (1 == dst[i])
        {
            sum++;
            pos[2 * j]     = 2 * i;
            pos[2 * j + 1] = 2 * i + 1;
            j++;
        }
    }
    return 2 * sum;
}


double
do_chunk(double mat[], int n, unsigned long long int X,
            unsigned long long int chunksize)
{
    /*
     * This function calculates adds parts X to X+chunksize of Cygan and Pilipczuk formula for the
     * Hafnian of matrix mat
     */

    double res = 0.0;

    for (unsigned long long int x = X; x < X + chunksize; x++)
    {
        double summand = 0.0;
        sint   m = n / 2;
        int    i, j, k;
        double factor;
        double powfactor;


        char dst[m];
        sint pos[n];
        char sum;
        dec2bin(dst, x, m);
        sum = find2(dst, m, pos);
        double B[sum * sum];
        for (i = 0; i < sum; i++)
        {
            for (j = 0; j < sum; j++)
            {
                B[i * sum + j] = mat[pos[i] * n + ((pos[j]) ^ 1)];
            }
        }
        double traces[m];
        powtrace(B, sum, traces, m);
        char cnt;
        sint cntindex;
        cnt      = 1;
        cntindex = 0;
        double comb[2][m + 1];
        for (i = 0; i < n / 2 + 1; i++)
        {
            comb[0][i] = 0.0;
        }
        comb[0][0] = 1.0;

        for (i = 1; i <= n / 2; i++)
        {
            factor    = traces[i - 1] / (2 * i);
            powfactor = 1.0;
            //the following line is significantly different than in Octave
            //It basically reassings  the first column of comb to the second and viceversa
            cnt      = -cnt;
            cntindex = (1 + cnt) / 2;
            for (j = 0; j < n / 2 + 1; j++)
            {
                comb[1 - cntindex][j] = comb[cntindex][j];
            }
            for (j = 1; j <= (n / (2 * i)); j++)
            {
                powfactor = powfactor * factor / j;
                for (k = i * j + 1; k <= n / 2 + 1; k++)
                {
                    comb[1 - cntindex][k - 1] =
                        comb[1 - cntindex][k - 1] + comb[cntindex][k - i * j -
                                                                                 1] * powfactor;
                }
            }
        }
        if (((sum / 2) % 2) == (n / 2 % 2))
        {
            summand = comb[1 - cntindex][n / 2];
        }
        else
        {
            summand = -comb[1 - cntindex][n / 2];
        }
        res += summand;
    }

    return res;
}


void
haf(double mat[], int n, double res[])
{
    /*
     * This is a wrapper for the function hafnian. Instead of returning the Hafnian of the matrix
     * mat by value it does by reference with two doubles for the real an imaginary respectively.
     */
    telem result = hafnian(mat, n);

    res[0] = creal(result);
    res[1] = cimag(result);
}


void
dhaf(double *mat, int n, double *res)
{
    haf(mat, n, res);
}


telem
hafnian(double *mat, int n)
{
    /*
     * Add all the terms necessary to calculate the Hafnian of the matrix mat of size n
     * It additionally uses open MP to parallelize the summation
     */
    assert(n % 2 == 0);
    sint  m   = n / 2;
    sint  mm  = m / 2;
    telem res = 0.0;
    unsigned long long int pow1    = ((unsigned long long int)pow(2.0, (double)m));  //-1;
    unsigned long long int workers =
        ((unsigned long long int)pow(2.0, (double)mm));
    workers = MIN(workers, pow1);
    unsigned long long int chunksize = pow1 / workers;

#pragma omp parallel
    {
        telem summand;
#pragma omp for
        for (unsigned long long int X = 1; X <= pow1; X += chunksize)
        {
            summand = do_chunk(mat, n, X, chunksize);
#pragma omp critical
            res += summand;
        }         //end for X
    }
    return res;
}


telem
dhafnian(double *mat, int n)
{
    return hafnian((double *)mat, n);
}


// Functions for computing Hafnians WITH considering loops (diagonal elements). These functions are named with a suffix "_loops" to differentiate them from previous ones.

double do_chunk_loops(double mat[], double C[], double D[], int n, unsigned long long int X, unsigned long long int chunksize)
{
    /*
     * This function calculates adds parts X to X+chunksize of Cygan and Pilipczuk formula for the
     * Hafnian of matrix mat
     */

    double res = 0.0;

#pragma omp parallel
    {
#pragma omp for

        for (unsigned long long int x = X * chunksize; x < (X + 1) * chunksize; x++)
        {
            double summand = 0.0;
            sint   m = n / 2;
            int    i, j, k;
            double factor;
            double powfactor;


            char dst[m];
            sint pos[n];
            sint sum;
            dec2bin(dst, x, m);
            sum = find2(dst, m, pos);

            double B[sum * sum], B_powtrace[sum * sum];
            double C1[sum], D1[sum];


            for (i = 0; i < sum; i++)
            {
                for (j = 0; j < sum; j++)
                {
                    B[i * sum + j]          = (double)mat[pos[i] * n + ((pos[j]) ^ 1)];
                    B_powtrace[i * sum + j] = (double)mat[pos[i] * n + ((pos[j]) ^ 1)];
                }
                C1[i] = C[pos[i]];
                D1[i] = D[pos[i]];
            }


            double traces[m];
            powtrace(B_powtrace, sum, traces, m);
            char cnt;
            sint cntindex;
            cnt      = 1;
            cntindex = 0;
            double comb[2][m + 1];
            for (i = 0; i < n / 2 + 1; i++)
            {
                comb[0][i] = 0.0;
            }
            comb[0][0] = 1.0;

            for (i = 1; i <= n / 2; i++)
            {
                factor    = traces[i - 1] / (2 * i);
                powfactor = 1.0;
                //the following line is significantly different than in Octave
                //It basically reassings  the first column of comb to the second and viceversa
                cnt      = -cnt;
                cntindex = (1 + cnt) / 2;
                for (j = 0; j < n / 2 + 1; j++)
                {
                    comb[1 - cntindex][j] = comb[cntindex][j];
                }

                for (j = 1; j <= (n / (2 * i)); j++)
                {
                    powfactor = powfactor * factor / j;
                    for (k = i * j + 1; k <= n / 2 + 1; k++)
                    {
                        comb[1 - cntindex][k - 1] = comb[1 - cntindex][k - 1] + comb[cntindex][k - i * j - 1] * powfactor;
                    }
                }

                //  % Combinations from open walks with loops at both ends
                // This is where the new part goes!

                factor = 0.0;

                for (int i = 0; i < sum; i++)
                {
                    factor += C1[i] * D1[i] / 2;
                }

                double complex tmp_c1[sum];
                double complex tmp = 0.0;

                for (int i = 0; i < sum; i++)
                {
                    tmp = 0.0;

                    for (int j = 0; j < sum; j++)
                    {
                        tmp += C1[j] * B[j * sum + i];
                    }

                    tmp_c1[i] = tmp;
                }

                for (int i = 0; i < sum; i++)
                {
                    C1[i] = tmp_c1[i];
                }


                powfactor = 1.0;
                //the following line is significantly different than in Octave
                //It basically reassings  the first column of comb to the second and viceversa
                cnt      = -cnt;
                cntindex = (1 + cnt) / 2;
                for (j = 0; j < n / 2 + 1; j++)
                {
                    comb[1 - cntindex][j] = comb[cntindex][j];
                }

                for (j = 1; j <= (n / (2 * i)); j++)
                {
                    powfactor = powfactor * factor / j;
                    for (k = i * j + 1; k <= n / 2 + 1; k++)
                    {
                        comb[1 - cntindex][k - 1] = comb[1 - cntindex][k - 1] + comb[cntindex][k - i * j - 1] * powfactor;
                    }
                }
            }


            if (((sum / 2) % 2) == (n / 2 % 2))
            {
                summand = comb[1 - cntindex][n / 2];
            }
            else
            {
                summand = -comb[1 - cntindex][n / 2];
            }
#pragma omp critical
            res += summand;
        }
    }

    return res;
}


double hafnian_loops(double *mat, int n)
{
    /*
     * Add all the terms necessary to calculate the Hafnian of the matrix mat of size n
     * It additionally uses open MP to parallelize the summation
     */
    assert(n % 2 == 0);
    sint m  = n / 2;
    sint mm = m / 2;
    // double res = 0.0;
    unsigned long long int pow1    = ((unsigned long long int)pow(2.0, (double)m)); //-1;
    unsigned long long int workers = ((unsigned long long int)pow(2.0, (double)mm));
    workers = MIN(workers, pow1);


    // double D_mat[n][n] ;
    double D[n], C[n];


    unsigned long long int chunksize = pow1;
    unsigned long long int rank      = 0;

    for (int i = 0; i < n; i++)
    {
        // for(int j = 0; j < n; j++)
        // {
        //     D_mat[i][j] = 0.0;
        // }
        D[i] = 0.0;
        C[i] = 0.0;
    }
    for (int i = 0; i < n; i++)
    {
        // D_mat[i][i] = mat[i * n + i];
        D[i] = mat[i * n + i];
    }


    for (int i = 0; i < n; i += 2)
    {
        C[i]     = D[i + 1];
        C[i + 1] = D[i];
    }



    double summand;

    summand = do_chunk_loops(mat, C, D, n, rank, chunksize);

    return summand;
}


double dhafnian_loops(double *mat, int n)
{
    return hafnian_loops((double *)mat, n);
}


void haf_loops(double mat[], int n, double res[])
{
    /*
     * This is a wrapper for the function hafnian. Instead of returning the Hafnian of the matrix
     * mat by value it does by reference with two doubles for the real an imaginary respectively.
     */

    double result = hafnian_loops(mat, n);

    res[0] = creal(result);
    res[1] = cimag(result);
}


void dhaf_loops(double *mat, int n, double *res)
{
    haf_loops((double *)mat, n, res);
}
