/*
    Theseus - maximum likelihood superpositioning of macromolecular structures

    Copyright (C) 2004-2010 Douglas L. Theobald

    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.,
    59 Temple Place, Suite 330,
    Boston, MA  02111-1307  USA

    -/_|:|_|_\-
*/

#include "CovMat_local.h"
#include "CovMat.h"


void
PrintCovMatGnuPlot(const double **mat, const int dim, char *outfile)
{
    FILE           *fp = NULL;
    int             i, j;

    fp = myfopen(outfile, "w");
/*     for (i = 0; i < dim; ++i) */
/*     { */
/*         for (j = 0; j < dim; ++j) */
/*             fprintf(fp, "% 11.8f ", mat[i][j]); */
/*  */
/*         fprintf(fp, "\n"); */
/*     } */
/*     fprintf(fp, "\n"); */

    for (i = 0; i < dim; ++i)
    {
        for (j = 0; j < dim; ++j)
            fprintf(fp, "%4d %4d % 14.8f\n", i, j, mat[i][j]);

        fprintf(fp, "\n");
    }
    fprintf(fp, "\n");

    fclose(fp);
}


void
SetupCovWeighting(CdsArray *cdsA)
{
    int             i;
    const int       vlen = cdsA->vlen;

    /* set up matrices and initialize to identity for full covariance matrix weighting */
    if (cdsA->CovMat == NULL)
        cdsA->CovMat = MatAlloc(vlen, vlen);

    if (cdsA->WtMat == NULL)
        cdsA->WtMat = MatAlloc(vlen, vlen);

    for (i = 0; i < vlen; ++i)
        cdsA->CovMat[i][i] = cdsA->WtMat[i][i] = 1.0;

    if (cdsA->tmpvecK == NULL)
        cdsA->tmpvecK = malloc(vlen * sizeof(double));

    if (cdsA->tmpmatKK1 == NULL)
        cdsA->tmpmatKK1 = MatAlloc(vlen, vlen);

    if (cdsA->tmpmatKK2 == NULL)
        cdsA->tmpmatKK2 = MatAlloc(vlen, vlen);

    if (cdsA->tmpmatK3a == NULL)
        cdsA->tmpmatK3a = MatAlloc(vlen, 3);

    if (cdsA->tmpmatK3b == NULL)
        cdsA->tmpmatK3b = MatAlloc(vlen, 3);

    if (cdsA->tmpmat3K == NULL)
        cdsA->tmpmat3K = MatAlloc(3, vlen);
}


/* returns 1 if all variances are about zero (< DBL_EPSILON) */
int
CheckZeroVariances(CdsArray *cdsA)
{
    Algorithm      *algo = cdsA->algo;
    int             i, zeroflag = 1;

    if (algo->varweight != 0)
    {
        for (i = 0; i < cdsA->vlen; ++i)
            if (cdsA->var[i] > DBL_EPSILON)
                zeroflag = 0;
    }
    else if (algo->covweight != 0)
    {
        for (i = 0; i < cdsA->vlen; ++i)
            if (cdsA->CovMat[i][i] > DBL_EPSILON)
                zeroflag = 0;
    }

    return(zeroflag);

/*     if (zeroflag == 1) */
/*     { */
/*         double var = cdsA->stats->wRMSD_from_mean * cdsA->stats->wRMSD_from_mean; */
/*  */
/*      if (algo->varweight != 0) */
/*      { */
/*          memsetd(cdsA->var, var, cdsA->vlen); */
/*      } */
/*      else if (algo->covweight != 0) */
/*      { */
/*          for (i = 0; i < cdsA->vlen; ++i) */
/*              cdsA->CovMat[i][i] = var; */
/*      } */
/*     } */
}


void
CalcBfactC(CdsArray *cdsA)
{
    int         i, j;
    double      trBS, occsum;

    for (i = 0; i < cdsA->cnum; ++i)
    {
        trBS = occsum = 0.0;
        for (j = 0; j < cdsA->vlen; ++j)
        {
            if (cdsA->cds[i]->o[j] > 0)
            {
                occsum += 1.0;
                trBS += cdsA->cds[i]->prvar[j] / cdsA->var[j];
                /*printf("trBS[%d] = % f\n", j, cdsA->cds[i]->prvar[j] / cdsA->var[j]);*/
            }
        }

        cdsA->cds[i]->bfact_c = occsum / trBS;
        /*printf("bfact_c[%d] = % f\n", i, cdsA->cds[i]->bfact_c);*/
    }
}


/* Weighting by dimensional, axial Xi covariance matrix, here diagonal. */
/* Holding the superposition constant, calculates the covariance
   matrices and corresponding weight matrices, looping till 
   convergence. */
void
CalcCovariances(CdsArray *cdsA)
{
    Algorithm      *algo = cdsA->algo;

	if (algo->varweight != 0 || algo->leastsquares != 0)
	{
		if (algo->alignment == 1)
			VarianceCdsOcc(cdsA);
		else
			VarianceCds(cdsA);
	}
	else if (algo->covweight != 0)
		CalcCovMat(cdsA);
}


void
MVCovMat(CdsArray *cdsA)
{
    int             i, j;
    const int       cnum = cdsA->cnum, vlen = cdsA->vlen;
    const double    idf = 1.0 / (double)(cnum * vlen);
    const Cds  **cds = (const Cds **) cdsA->cds;
    const Cds   *cdsi;
    double        **MVCovMat = NULL;

    if (cdsA->MVCovMat == NULL)
        MVCovMat = cdsA->MVCovMat = MatAlloc(3, 3);
    else
        MVCovMat = cdsA->MVCovMat;

    for (i = 0; i < 3; ++i)
        for (j = 0; j < 3; ++j)
            MVCovMat[i][j] = 0.0;

    for (i = 0; i < cnum; ++i)
    {
        for (j = 0; j < vlen; ++j)
        {
            cdsi = (Cds *) cds[i];
            MVCovMat[0][0] += mysquare(cdsi->residual_x[j]);
            MVCovMat[1][1] += mysquare(cdsi->residual_y[j]);
            MVCovMat[2][2] += mysquare(cdsi->residual_z[j]);
            MVCovMat[0][1] += cdsi->residual_x[j] * cdsi->residual_y[j];
            MVCovMat[0][2] += cdsi->residual_x[j] * cdsi->residual_z[j];
            MVCovMat[1][2] += cdsi->residual_y[j] * cdsi->residual_z[j];
        }
    }

    MVCovMat[1][0] = MVCovMat[0][1];
    MVCovMat[2][0] = MVCovMat[0][2];
    MVCovMat[2][1] = MVCovMat[1][2];

    for (i = 0; i < 3; ++i)
        for (j = 0; j < 3; ++j)
            MVCovMat[i][j] *= idf;
}


void
CalcCovMat(CdsArray *cdsA)
{
    double          newx1, newy1, newz1, newx2, newy2, newz2;
    double          covsum;
    double         *cdskx, *cdsky, *cdskz;
    int             i, j, k;
    const int       cnum = cdsA->cnum, vlen = cdsA->vlen;
    const double    normalize = 1.0 / (3.0 * cnum);
    const Cds  **cds = (const Cds **) cdsA->cds;
    const Cds   *cdsk;
    double        **CovMat = cdsA->CovMat;
    const double   *avex = (const double *) cdsA->avecds->x,
                   *avey = (const double *) cdsA->avecds->y,
                   *avez = (const double *) cdsA->avecds->z;

    if (cdsA->CovMat == NULL)
    {
        printf("\nERROR654\n");
        exit(EXIT_FAILURE);
    }

    /* calculate covariance matrix of atoms across structures,
       based upon current superposition, put in CovMat */
    for (i = 0; i < vlen; ++i)
    {
        for (j = 0; j <= i; ++j)
        {
            covsum = 0.0;
            for (k = 0; k < cnum; ++k)
            {
                cdsk = cds[k];
                cdskx = cdsk->x;
                cdsky = cdsk->y;
                cdskz = cdsk->z;

                newx1 = cdskx[i] - avex[i];
                newy1 = cdsky[i] - avey[i];
                newz1 = cdskz[i] - avez[i];

                newx2 = cdskx[j] - avex[j];
                newy2 = cdsky[j] - avey[j];
                newz2 = cdskz[j] - avez[j];

                #ifdef FP_FAST_FMA
                covsum += fma(newx1, newx2, fma(newy1, newy2, newz1 * newz2));
                #else
                covsum += (newx1 * newx2 + newy1 * newy2 + newz1 * newz2);
                #endif
            }

            CovMat[i][j] = CovMat[j][i] = covsum * normalize; /* sample variance, ML biased not n-1 definition */
        }
    }

    for (i = 0; i < vlen; ++i)
        cdsA->var[i] = CovMat[i][i];
}


/* Same as CalcCovMat() but weights by the occupancies */
void
CalcCovMatOcc(CdsArray *cdsA)
{
    double          newx1, newy1, newz1, newx2, newy2, newz2;
    double          covsum;
    double         *cdskx, *cdsky, *cdskz;
    int             i, j, k;
    const int       cnum = cdsA->cnum, vlen = cdsA->vlen;
    const Cds  **cds = (const Cds **) cdsA->cds;
    const Cds   *cdsk;
    double        **CovMat = cdsA->CovMat;
    const double   *avex = (const double *) cdsA->avecds->x,
                   *avey = (const double *) cdsA->avecds->y,
                   *avez = (const double *) cdsA->avecds->z;
    double         *occ, osum;

    if (cdsA->CovMat == NULL)
    {
        printf("\nERROR653\n");
        exit(EXIT_FAILURE);
    }

    /* calculate covariance matrix of atoms across structures,
       based upon current superposition, put in CovMat */
    for (i = 0; i < vlen; ++i)
    {
        for (j = 0; j <= i; ++j)
        {
            covsum = osum = 0.0;
            for (k = 0; k < cnum; ++k)
            {
                cdsk = cds[k];
                occ = cdsk->o;
                cdskx = cdsk->x;
                cdsky = cdsk->y;
                cdskz = cdsk->z;

                newx1 = cdskx[i] - avex[i];
                newy1 = cdsky[i] - avey[i];
                newz1 = cdskz[i] - avez[i];

                newx2 = cdskx[j] - avex[j];
                newy2 = cdsky[j] - avey[j];
                newz2 = cdskz[j] - avez[j];

                covsum += occ[i] * occ[j] * 
                          (newx1 * newx2 + newy1 * newy2 + newz1 * newz2);

                osum += occ[i] * occ[j];
            }

            if (osum > 0.0)
                CovMat[i][j] = CovMat[j][i] = covsum / osum; /* sample variance, ML biased not n-1 definition */
            else
                CovMat[i][j] = CovMat[j][i] = 0.0;
        }
    }
}


void
CalcStructCovMat(CdsArray *cdsA)
{
    double          invdf, cov_sum;
    int             i, j, k;
    const int       cnum = cdsA->cnum, vlen = cdsA->vlen;
    const Cds  **cds = (const Cds **) cdsA->cds;
    const Cds   *cdsi = NULL, *cdsj = NULL;
    double        **SCovMat = NULL;

    invdf = 1.0 / (double) vlen; /* ML, biased */

    if (cdsA->SCovMat == NULL)
        cdsA->SCovMat = MatAlloc(cnum, cnum);
    else
        memset(&cdsA->SCovMat[0][0], 0, cnum * cnum * sizeof(double));

    SCovMat = cdsA->SCovMat;

    /* calculate covariance matrix of structures across atoms, put in SCovMat */
    for (i = 0; i < cnum; ++i)
    {
        for (j = 0; j < cnum; ++j)
        {
            cov_sum = 0.0;
            for (k = 0; k < vlen; ++k)
            {
                cdsi = cds[i];
                cdsj = cds[j];
                cov_sum += cdsi->x[k] * cdsj->x[k];
                cov_sum += cdsi->y[k] * cdsj->y[k];
                cov_sum += cdsi->z[k] * cdsj->z[k];
            }

            SCovMat[i][j] = cov_sum * invdf; /* sample variance, ML biased not n-1 definition */
        }
    }
}


void
CalcFullCovMat(CdsArray *cdsA)
{
    double          newx1, newy1, newz1, newx2, newy2, newz2;
    double          invdf;
    double        **FullCovMat = cdsA->FullCovMat;
    const double   *avex = (const double *) cdsA->avecds->x,
                   *avey = (const double *) cdsA->avecds->y,
                   *avez = (const double *) cdsA->avecds->z;
    const int       vlen = cdsA->vlen, cnum = cdsA->cnum;
    int             i, j, k, m, n, p, q0, q1, q2;
    const Cds  **cds = (const Cds **) cdsA->cds;
    Cds         *cdsk;

    invdf = 1.0 / (double) cnum; /* ML, biased */

    AveCds(cdsA);

    if (FullCovMat == NULL)
        FullCovMat = cdsA->FullCovMat = MatAlloc(3 * cdsA->vlen, 3 * cdsA->vlen);

    for (i = 0; i < 3 * cdsA->vlen; ++i)
        for (j = 0; j < 3 * cdsA->vlen; ++j)
            FullCovMat[i][j] = 0.0;

    /* calculate covariance matrix based upon current superposition, put in FullCovMat */
    for (m = i = 0; i < vlen; ++i, m += 3)
    {
        for (n = j = 0; j <= i; ++j, n += 3)
        {
            for (k = 0; k < cnum; ++k)
            {
                cdsk = (Cds *) cds[k];
                newx1 = cdsk->x[i] - avex[i];
                newy1 = cdsk->y[i] - avey[i];
                newz1 = cdsk->z[i] - avez[i];

                newx2 = cdsk->x[j] - avex[j];
                newy2 = cdsk->y[j] - avey[j];
                newz2 = cdsk->z[j] - avez[j];

                q0 = n+0;
                q1 = n+1;
                q2 = n+2;

                p = m+0;
                FullCovMat[p][q0] += newx1 * newx2;
                FullCovMat[p][q1] += newx1 * newy2;
                FullCovMat[p][q2] += newx1 * newz2;

                p = m+1;
                FullCovMat[p][q0] += newy1 * newx2;
                FullCovMat[p][q1] += newy1 * newy2;
                FullCovMat[p][q2] += newy1 * newz2;

                p = m+2;
                FullCovMat[p][q0] += newz1 * newx2;
                FullCovMat[p][q1] += newz1 * newy2;
                FullCovMat[p][q2] += newz1 * newz2;
            }
        }
    }

    for (i = 0; i < 3 * vlen; ++i)
        for (j = 0; j <= i; ++j)
            FullCovMat[i][j] *= invdf;

    for (i = 0; i < 3 * vlen; ++i)
        for (j = 0; j < i; ++j)
            FullCovMat[j][i] = FullCovMat[i][j];
}


/* calculate covariance matrix weighted cds
   \Sigma^-1 * \CdsMat */
void
CalcCovCds(Cds *cds, const double **covmat)
{
    int             i, k;
    double         *covx = cds->covx,
                   *covy = cds->covy,
                   *covz = cds->covz;
    const double   *x = (const double *) cds->x,
                   *y = (const double *) cds->y,
                   *z = (const double *) cds->z;
    double          covmatik;

    for (i = 0; i < cds->vlen; ++i)
    {
        covx[i] = covy[i] = covz[i] = 0.0;

        for (k = 0; k < cds->vlen; ++k)
        {
            covmatik = covmat[i][k];

            covx[i] += (covmatik * x[k]);
            covy[i] += (covmatik * y[k]);
            covz[i] += (covmatik * z[k]);
        }
    }
}


/* Normalize the covariance matrix to form the correlation matrix 
   by dividing each element by the square root of the product of the
   corresponding diagonal elements.
   This makes a pearson correlation matrix.
   The diagonal elements are always equal to 1, while
   the off-diagonals range from -1 to 1.
*/
void
CovMat2CorMat(double **CovMat, const int size)
{
    int             i, j;

    for (i = 0; i < size; ++i)
        for (j = 0; j < i; ++j)
            CovMat[i][j] = CovMat[j][i] = CovMat[i][j] / sqrt(CovMat[i][i] * CovMat[j][j]);

    for (i = 0; i < size; ++i)
        CovMat[i][i] = 1.0;
}


void
CorMat2CovMat(double **CovMat, const double *vars, const int size)
{
    int             i, j;

    for (i = 0; i < size; ++i)
        for (j = 0; j < i; ++j)
            CovMat[i][j] = CovMat[j][i] = CovMat[i][j] * sqrt(vars[i] * vars[j]);

    for (i = 0; i < size; ++i)
        CovMat[i][i] = vars[i];
}


/* Normalizes a covariance matrix by dividing every cell by the 
   average variance */
double
NormalizeCovMat(double **mat, const int size)
{
    int             i, j;
    double          normalize;

    normalize = 0.0;
    for (i = 0; i < size; ++i)
        normalize += mat[i][i];

    normalize = size / normalize;

/*     normalize = 0.0; */
/*     for (i = 0; i < size; ++i) */
/*         for (j = 0; j < size; ++j) */
/*             normalize += mat[i][j]; */
/*     normalize = (double) size / normalize; */

    for (i = 0; i < size; ++i)
        for (j = 0; j < size; ++j)
            mat[i][j] *= normalize;

    /* fprintf(stderr, "\n     Mat[%3d][%3d] = %12.5f", size/2, size/2, mat[size/2][size/2]); */
/*     fprintf(stderr, "\n     norm = %12.5f", normalize); */
/*     fflush(NULL); */
    return(normalize);
}


void
PrintCovMat(CdsArray *cdsA)
{
    int             i, j;
    const double  **CovMat = (const double **) cdsA->CovMat;

    for (i = 0; i < cdsA->vlen; ++i)
    {
        printf("\n");
        for (j = 0; j < cdsA->vlen; ++j)
            printf("%8.3f ", CovMat[i][j]);
    }
    printf("\n");
}
