/*
    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 <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <math.h>
#include <float.h>
#include "CovMat.h"
#include "HierarchVars.h"
#include "pdbMalloc.h"
#include "pdbStats.h"
#include "pdbIO.h"
#include "pdbUtils.h"
#include "ProcGSLSVD.h"
#include "ProcGSLSVDNu.h"
#include "CovMat.h"
#include "MultiPose.h"
#include "MultiPose2MSA.h"
#include "DLTmath.h"
#include <gsl/gsl_rng.h>
#include "theseuslib.h"


int
test_charmm(CdsArray *cdsA)
{
    int         nensem, natomx, ls, cnum, vlen;
    double     *xbuf = NULL, *ybuf = NULL, *zbuf = NULL;
    int         i;

    nensem = cnum = cdsA->cnum;
    natomx = vlen = cdsA->vlen;

    xbuf = malloc(nensem*natomx*sizeof(double));
    ybuf = malloc(nensem*natomx*sizeof(double));
    zbuf = malloc(nensem*natomx*sizeof(double));

    for (i = 0; i < cnum; ++i)
    {
        memcpy(&xbuf[i*vlen], cdsA->cds[i]->x, vlen*sizeof(double));
        memcpy(&ybuf[i*vlen], cdsA->cds[i]->y, vlen*sizeof(double));
        memcpy(&zbuf[i*vlen], cdsA->cds[i]->z, vlen*sizeof(double));
    }

    theseus_(xbuf, ybuf, zbuf, &natomx, &nensem, &ls);

    for (i = 0; i < cnum; ++i)
    {
        memcpy(cdsA->cds[i]->x, &xbuf[i*vlen], vlen*sizeof(double));
        memcpy(cdsA->cds[i]->y, &ybuf[i*vlen], vlen*sizeof(double));
        memcpy(cdsA->cds[i]->z, &zbuf[i*vlen], vlen*sizeof(double));
    }

    free(xbuf);
    free(ybuf);
    free(zbuf);

    return(1);
}


/* Expects xbuf, etc. to be allocated vectors of length natomx*nensem */
int
theseus_(double *xbuf, double *ybuf, double *zbuf, int *natomx, int *nensem, int *ls)
{
    CdsArray    *cdsA = NULL;
    int         cnum = *nensem, vlen = *natomx;
    int         i, j;

    cdsA = CdsArrayInit();

    if (*ls)
    {
        algo->leastsquares = 1;
        algo->varweight = 0;
        algo->hierarch = 0;
    }

    CdsArrayAlloc(cdsA, cnum, vlen);

    /* DLT debug - shouldn't need to do this, I think prob is in SuperPose() */
    for (j = 0; j < cnum; ++j)
        for (i = 0; i < vlen; ++i)
            cdsA->cds[j]->nu[i] = 1;

    for (i = 0; i < cnum; ++i)
    {
        memcpy(cdsA->cds[i]->x, &xbuf[i*vlen], vlen*sizeof(double));
        memcpy(cdsA->cds[i]->y, &ybuf[i*vlen], vlen*sizeof(double));
        memcpy(cdsA->cds[i]->z, &zbuf[i*vlen], vlen*sizeof(double));
    }

    MultiPoseLib(cdsA);

    for (i = 0; i < cnum; ++i)
    {
        memcpy(&xbuf[i*vlen], cdsA->cds[i]->x, vlen*sizeof(double));
        memcpy(&ybuf[i*vlen], cdsA->cds[i]->y, vlen*sizeof(double));
        memcpy(&zbuf[i*vlen], cdsA->cds[i]->z, vlen*sizeof(double));
    }

    CdsArrayDestroy(&cdsA);

    return(1);
}


/* Expects xbuf, etc. to be allocated vectors of length natomx*nensem */
int
theseus2_(double *xbuf, double *ybuf, double *zbuf, int *natomx, int *nensem, int *ls)
{
    CdsArray    *cdsA = NULL;
    int         cnum = *nensem, vlen = *natomx;
    int         i, j;

    cdsA = CdsArrayInit();

    if (*ls)
    {
        algo->leastsquares = 1;
        algo->varweight = 0;
        algo->hierarch = 0;
    }

    CdsArrayAlloc(cdsA, cnum, vlen);

    /* DLT debug - shouldn't need to do this, I think prob is in SuperPose() */
    for (j = 0; j < cnum; ++j)
        for (i = 0; i < vlen; ++i)
            cdsA->cds[j]->nu[i] = 1;

    for (i = 0; i < cnum; ++i)
    {
        free(cdsA->cds[i]->x);
        free(cdsA->cds[i]->y);
        free(cdsA->cds[i]->z);

        cdsA->cds[i]->x = &xbuf[i*vlen];
        cdsA->cds[i]->y = &ybuf[i*vlen];
        cdsA->cds[i]->z = &zbuf[i*vlen];
    }

    MultiPoseLib(cdsA);

    for (i = 0; i < cnum; ++i)
    {
        cdsA->cds[i]->x = NULL;
        cdsA->cds[i]->y = NULL;
        cdsA->cds[i]->z = NULL;
    }

    CdsArrayDestroy(&cdsA);

    return(1);
}


void
CalcS2(CdsArray *cdsA, const int nsell, double *bxij, double *byij, double *bzij,
       double *rij, double *s2, const int whoiam)
{
    int         i, j, k, m;
    int         cnum = cdsA->cnum;
    double     *x = NULL, *y = NULL, *z = NULL;
    double      sx2, sy2, sz2, sxy, sxz, syz, xij, yij, zij, rijk;

    for (k = 0; k < nsell; ++k)
    {
        sx2 = sy2 = sz2 = sxy = sxz = syz = 0.0;
        rij[k] = 0.0;
        for (m = 0; m < cnum; ++m)
        {
            x = cdsA->cds[m]->x;
            y = cdsA->cds[m]->y;
            z = cdsA->cds[m]->z;

            i = k;
            j = nsell + k;
            xij = x[i] - x[j];
            yij = y[i] - y[j];
            zij = z[i] - z[j];

            if (m == whoiam)
            {
                bxij[k] = xij;
                byij[k] = yij;
                bzij[k] = zij;
            }

            rij[k] += sqrt(xij*xij + yij*yij + zij*zij);

            sx2 += xij*xij;
            sy2 += yij*yij;
            sz2 += zij*zij;
            sxy += xij*yij;
            sxz += xij*zij;
            syz += yij*zij;
        }

        sx2 /= cnum;
        sy2 /= cnum;
        sz2 /= cnum;
        sxy /= cnum;
        sxz /= cnum;
        syz /= cnum;

        rij[k] /= cnum;
        rijk = rij[k];

        //        order parameter:
        s2[k] = (1.5/(rijk*rijk*rijk*rijk)) * (sx2*sx2 + sy2*sy2 + sz2*sz2 +
              2.0*(sxy*sxy + sxz*sxz + syz*syz)) - 0.5;

        //        components of force:
        xij = bxij[k];
        yij = byij[k];
        zij = bzij[k];
        bxij[k] = sx2*xij + sxy*yij + sxz*zij;
        byij[k] = sy2*yij + sxy*xij + syz*zij;
        bzij[k] = sz2*zij + sxz*xij + syz*yij;

        if (isnan(s2[k]))
        {
            printf("ERRORTH1> %3d: %8.3e %8.3e %8.3e %8.3e %8.3e %8.3e %8.3e\n",
                   k, rijk, sx2, sy2, sz2, sxy, sxz, syz);
            printf("ERRORTH2> %3d: %8.3e %8.3e %8.3e \n", k, xij, yij, zij);
        }
    }
}


/* Expects xbuf, etc. to be allocated vectors of length natomx*nensem */
/* Expects bxij, byij, bzij, il, jl, rij, s2, etc. to be length nsell */
int
theseuss2_(const double *xbuf, const double *ybuf, const double *zbuf,
           const int *natomx, const int *nensem, const int *nsell,
           double *bxij, double *byij, double *bzij,
           const int *il, const int *jl, double *rij, double *s2,
           const int *ls, const int *whoiam)
{
    CdsArray    *cdsA = NULL;
    int         cnum = *nensem, len = *natomx, vlen = *nsell*2;
    int         i, j, ilj, jlj;

    cdsA = CdsArrayInit();

    if (*ls)
    {
        algo->leastsquares = 1;
        algo->varweight = 0;
        algo->hierarch = 0;
    }
//printf(" ENSS2ML>: cnum=%d len=%d vlen=%d nsell=%d\n", cnum, len, vlen, *nsell);
//fflush(NULL);
    CdsArrayAlloc(cdsA, cnum, vlen);

    /* DLT debug - shouldn't need to do this, I think prob is in SuperPose() */
    for (j = 0; j < cnum; ++j)
        for (i = 0; i < vlen; ++i)
            cdsA->cds[j]->nu[i] = 1;

    for (i = 0; i < cnum; ++i)
    {
        for (j = 0; j < *nsell; ++j)
        {
            ilj = il[j] - 1;
            cdsA->cds[i]->x[j] = xbuf[i*len + ilj];
            cdsA->cds[i]->y[j] = ybuf[i*len + ilj];
            cdsA->cds[i]->z[j] = zbuf[i*len + ilj];

            jlj = jl[j] - 1;
            cdsA->cds[i]->x[*nsell + j] = xbuf[i*len + jlj];
            cdsA->cds[i]->y[*nsell + j] = ybuf[i*len + jlj];
            cdsA->cds[i]->z[*nsell + j] = zbuf[i*len + jlj];
        }
    }

    MultiPoseLib(cdsA);

    CalcS2(cdsA, vlen, bxij, byij, bzij, rij, s2, *whoiam);
//VecPrint(bxij, *nsell);
//VecPrint(rij, *nsell);
//VecPrint(s2, *nsell);
//fflush(NULL);
//WriteCdsFile(cdsA->cds[0], "charmm.pdb");
    CdsArrayDestroy(&cdsA);

    return(1);
}


/* A version of MultiPose for a general library, very pared down */
int
MultiPoseLib(CdsArray *cdsA)
{
    int             i, round, innerround, slxn; /* index of random coord to select as first */
    const int       cnum = cdsA->cnum;
    Algorithm      *algo = NULL;
    Statistics     *stats = NULL;
    Cds        **cds = NULL;
    Cds         *avecds = NULL;
    gsl_rng               *r2 = NULL;
    const gsl_rng_type    *T = NULL;

    T = gsl_rng_ranlxs2;
    r2 = gsl_rng_alloc(T);

    /* setup cdsA */
    CdsArraySetup(cdsA);

    /* setup local aliases based on cdsA */
    cds = cdsA->cds;
    avecds = cdsA->avecds;

    if (algo->covweight)
        SetupCovWeighting(cdsA); /* DLT debug */

    stats->hierarch_p1 = 1.0;
    stats->hierarch_p2 = 1.0;

    /* randomly select a structure to use as the initial mean structure */
    //slxn = (int) (genrand_real2() * cnum);
    slxn = gsl_rng_uniform_int(r2, cnum);

    CdsCopyAll(avecds, cdsA->cds[slxn]);

    if (algo->alignment)
        CalcDf(cdsA);

//    WriteCdsFile(cdsA->cds[0], "charmm_inp0.pdb");
//    WriteCdsFile(cdsA->cds[1], "charmm_inp1.pdb");

    /* The EM algorithm */
    /* The outer loop:
       (1) First calculates the translations
       (2) Does inner loop -- calc rotations and average till convergence
       (3) Holding the superposition constant, calculates the covariance
           matrices and corresponding weight matrices, looping till
           convergence when using a dimensional/axial covariance matrix */
    round = 0;
    while(1)
    {
        ++round;
        algo->rounds = round;

        /* Estimate Translations: Find weighted center and translate all cds */
        //CalcTranslationsIp(cdsA, algo); // DLT OP

        for (i = 0; i < cnum; ++i)
            ApplyCenterIp(cds[i]);

        /* save the translation vector for each coord in the array */
        for (i = 0; i < cnum; ++i)
            memcpy(cds[i]->translation, cds[i]->center, 3 * sizeof(double));

        /* Inner loop:
           (1) Calc rotations given weights/weight matrices
           (2) Rotate cds with new rotations
           (3) Recalculate average

           Loop till convergence, holding constant the variances, covariances, and translations */
        innerround = 0;
        do
        {
            ++innerround;

            /* find the optimal rotation matrices */
            if (algo->alignment)
                CalcRotationsNu(cdsA);
            else
                CalcRotations(cdsA);

            if ((innerround == 1) && CheckConvergenceOuter(cdsA, round, algo->precision))
                goto outsidetheloops;

            /* rotate the scratch cds with new rotation matrix */
            for (i = 0; i < cnum; ++i)
                RotateCdsIp(cds[i], (const double **) cds[i]->matrix);

            /* find global rmsd and average cds (both held in structure) */
            if (algo->alignment)
            {
                AveCdsNu(cdsA);
                EM_MissingCds(cdsA);
            }
            else
            {
                AveCds(cdsA);
            }
        }
        while((CheckConvergenceInner(cdsA, algo->precision) == 0) && (innerround < 160));

        /* Holding the superposition constant, calculate the covariance
           matrix and corresponding weight matrix, looping till convergence. */
        CalcCovariances(cdsA);
        if (algo->varweight || algo->covweight)
            HierarchVars(cdsA);

//        VecPrint(cdsA->var, cdsA->vlen);

        /* calculate the weights/weight matrices */
        CalcWts(cdsA);
    }

    outsidetheloops:

//    WriteCdsFile(cdsA->cds[0], "charmm_out.pdb");
//    fflush(NULL);
//    CalcStats(cdsA);
    printf(" ENSS2ML>: THESEUS rounds: %d\n", round);

    gsl_rng_free(r2);
    r2 = NULL;

    return(round);
}


int
MultiPoseLibSimp(CdsArray *cdsA)
{
    int             i, round, innerround, slxn; /* index of random coord to select as first */
    const int       cnum = cdsA->cnum;
    Algorithm      *algo = NULL;
    Statistics     *stats = NULL;
    Cds        **cds = NULL;
    Cds         *avecds = NULL;
    gsl_rng               *r2 = NULL;
    const gsl_rng_type    *T = NULL;

    T = gsl_rng_ranlxs2;
    r2 = gsl_rng_alloc(T);

    /* setup cdsA */
    CdsArraySetup(cdsA);

    /* setup local aliases based on cdsA */
    cds = cdsA->cds;
    avecds = cdsA->avecds;

    stats->hierarch_p1 = 1.0;
    stats->hierarch_p2 = 1.0;

    /* randomly select a structure to use as the initial mean structure */
    slxn = gsl_rng_uniform_int(r2, cnum);
    CdsCopyAll(avecds, cdsA->cds[slxn]);

    /* The CEM algorithm */
    /* The outer loop:
       (1) First calculates the translations
       (2) Does inner loop -- calc rotations and average till convergence
       (3) Holding the superposition constant, calculates the covariance
           matrices and corresponding weight matrices */
    round = 0;
    while(1)
    {
        ++round;
        algo->rounds = round;

        /* Estimate Translations: Find weighted center and translate all cds */
        //CalcTranslationsIp(cdsA, algo); // DLT OP

        for (i = 0; i < cnum; ++i)
            ApplyCenterIp(cds[i]);

        /* save the translation vector for each coord in the array */
        for (i = 0; i < cnum; ++i)
            memcpy(cds[i]->translation, cds[i]->center, 3 * sizeof(double));

        /* Inner loop:
           (1) Calc rotations given weights/weight matrices
           (2) Rotate cds with new rotations
           (3) Recalculate average

           Loop till convergence, holding constant the variances, covariances, and translations */
        innerround = 0;
        do
        {
            ++innerround;

            /* find the optimal rotation matrices */
            CalcRotations(cdsA);

            if ((innerround == 1) && (CheckConvergenceOuter(cdsA, round, algo->precision)))
                return(round);

            /* rotate the scratch cds with new rotation matrix */
            for (i = 0; i < cnum; ++i)
                RotateCdsIp(cds[i], (const double **) cds[i]->matrix);

            /* find global rmsd and average cds (both held in structure) */
            AveCds(cdsA);
        }
        while((CheckConvergenceInner(cdsA, algo->precision) == 0) && (innerround < 160));

        /* Holding the superposition constant, calculate the covariance
           matrix and corresponding weight matrix, looping till convergence. */
        CalcCovariances(cdsA);
        if (algo->varweight || algo->covweight)
            HierarchVars(cdsA);

        /* calculate the weights/weight matrices */
        CalcWts(cdsA);
    }

    gsl_rng_free(r2);
    r2 = NULL;

    return(round);
}