/*
    Theseus - maximum likelihood superpositioning of macromolecular structures

    Copyright (C) 2004-2014 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);
}


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

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

    return(zeroflag);

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


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

    for (i = 0; i < cdsA->cnum; ++i)
    {
        trBS = nusum = 0.0;
        for (j = 0; j < cdsA->vlen; ++j)
        {
            if (cdsA->cds[i]->nu[j])
            {
                nusum += 1;
                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 = nusum / trBS;
        /*printf("bfact_c[%d] = % f\n", i, cdsA->cds[i]->bfact_c);*/
    }
}


/* Holding the superposition constant, calculates the covariance
   matrices and corresponding weight matrices, looping till
   convergence. */
void
CalcCovariances(CdsArray *cdsA)
{
    if (algo->varweight || algo->leastsquares)
    {
        if (algo->alignment)
            VarianceCdsNu(cdsA);
        else
            VarianceCds(cdsA);
    }
    else if (algo->covweight)
        CalcCovMat(cdsA);
}


void
CalcCovMat(CdsArray *cdsA)
{
    double          newx1, newy1, newz1, newx2, newy2, newz2;
    double          avexi, aveyi, avezi, avexj, aveyj, avezj;
    double          covsum;
    double         *cdskx = NULL, *cdsky = NULL, *cdskz = NULL;
    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 = NULL;
    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)
    {
        avexi = avex[i];
        aveyi = avey[i];
        avezi = avez[i];

        for (j = 0; j <= i; ++j)
        {
            avexj = avex[j];
            aveyj = avey[j];
            avezj = avez[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] - avexi;
                newy1 = cdsky[i] - aveyi;
                newz1 = cdskz[i] - avezi;

                newx2 = cdskx[j] - avexj;
                newy2 = cdsky[j] - aveyj;
                newz2 = cdskz[j] - avezj;

                #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 nu missing flag */
void
CalcCovMatNu(CdsArray *cdsA)
{
    double          newx1, newy1, newz1, newx2, newy2, newz2;
    double          covsum;
    double         *cdskx = NULL, *cdsky = NULL, *cdskz = NULL;
    int             i, j, k;
    const int       cnum = cdsA->cnum, vlen = cdsA->vlen;
    const Cds     **cds = (const Cds **) cdsA->cds;
    const Cds      *cdsk = NULL;
    double        **CovMat = cdsA->CovMat;
    const double   *avex = (const double *) cdsA->avecds->x,
                   *avey = (const double *) cdsA->avecds->y,
                   *avez = (const double *) cdsA->avecds->z;
    int            *nu = NULL;
    double          nusum;

    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 = nusum = 0.0;
            for (k = 0; k < cnum; ++k)
            {
                cdsk = cds[k];
                nu = cdsk->nu;
                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 += nu[i] * nu[j] *
                          (newx1 * newx2 + newy1 * newy2 + newz1 * newz2);

                nusum += nu[i] * nu[j];
            }

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


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 = NULL;

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

    if (algo->doave)
        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];
}


/* 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];
}


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");
}