/*
    Theseus - maximum likelihood superpositioning of macromolecular structures

    Copyright (C) 2004-2015 Douglas L. Theobald

    This file is part of THESEUS.

    THESEUS 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 3 of
    the License, or (at your option) any later version.

    THESEUS 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 THESEUS in the file 'COPYING'. If not, see
    <http://www.gnu.org/licenses/>.

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

#include "MultiPosePth_local.h"
#include "MultiPosePth.h"


static void
*CalcRot_pth(void *rotdata_ptr);

static double
CalcRotations_pth(CdsArray *cdsA, RotData **rotdata, pthread_t *callThd,
                  pthread_attr_t *attr, const int thrdnum);


static void
*CalcRot_pth(void *rotdata_ptr)
{
    int             i;
    double          deviation = 0.0;
    RotData        *rotdata = (RotData *) rotdata_ptr;
    Cds            *cds = NULL;

    for (i = rotdata->start; i < rotdata->end; ++i)
    {
        cds = rotdata->cds[i];
        /* note that the avecds are already multiplied by the weight matrices */
//        deviation = CalcRMSDRotationalMatrix(cds, rotdata->tcds, cds->vlen, &cds->matrix[0][0], NULL);

        /* rotate the scratch cds with new rotation matrix */
        RotateCdsIp(cds, (const double **) cds->matrix);

        /* find global rmsd and average cds (both held in structure) */
        cds->wRMSD_from_mean = sqrt(deviation / (3 * rotdata->vlen));
    }

    pthread_exit((void *) 0);
}


static double
CalcRotations_pth(CdsArray *cdsA, RotData **rotdata, pthread_t *callThd,
                  pthread_attr_t *attr, const int thrdnum)
{
    Cds           **cds = cdsA->cds;
    const Cds      *avecds = cdsA->avecds;
    const double   *wts = (const double *) cdsA->w;
    Cds            *tcds = cdsA->tcds;
    double          deviation_sum = 0.0;
    int             i, rc = 0, incr;
    const int       cnum = cdsA->cnum;

    if (algo->covweight)
    {
        MatMultCdsMultMatDiag(tcds,
                              (const double **) cdsA->WtMat,
                              avecds);
    }
    else if (algo->varweight || algo->leastsquares)
    {
        MatDiagMultCdsMultMatDiag(tcds, wts, avecds);
    }

    incr = cnum / thrdnum;

    for (i = 0; i < thrdnum - 1; ++i)
    {
        rotdata[i]->cds = cds;
        rotdata[i]->tcds = tcds;
        rotdata[i]->start = i * incr;
        rotdata[i]->end = i*incr + incr;
        rotdata[i]->vlen = cdsA->vlen;

        rc = pthread_create(&callThd[i], attr, CalcRot_pth, (void *) rotdata[i]);

        if (rc)
        {
            printf("ERROR811: return code from pthread_create() %d is %d\n", i, rc);
            exit(EXIT_FAILURE);
        }
    }

    rotdata[thrdnum - 1]->cds = cds;
    rotdata[thrdnum - 1]->tcds = tcds;
    rotdata[thrdnum - 1]->start = (thrdnum - 1) * incr;
    rotdata[thrdnum - 1]->end = cnum;
    rotdata[thrdnum - 1]->vlen = cdsA->vlen;

    rc = pthread_create(&callThd[thrdnum - 1], attr, CalcRot_pth, (void *) rotdata[thrdnum - 1]);

    if (rc)
    {
        printf("ERROR811: return code from pthread_create() %d is %d\n", i, rc);
        exit(EXIT_FAILURE);
    }

    for (i = 0; i < thrdnum; ++i)
    {
        rc = pthread_join(callThd[i], (void **) NULL);

        if (rc)
        {
            printf("ERROR812: return code from pthread_join() %d is %d\n", i, rc);
            exit(EXIT_FAILURE);
        }
    }

    for (i = 0; i < cnum; ++i)
        deviation_sum += 3 * cdsA->vlen * cds[i]->wRMSD_from_mean * cds[i]->wRMSD_from_mean;

    return(deviation_sum);
}


int
MultiPose_pth(CdsArray *baseA)
{
    int             i, round, innerround;
    int             slxn; /* index of random coord to select as first */
    const int       cnum = baseA->cnum;
    const int       vlen = baseA->vlen;
    double         *evals = malloc(3 * sizeof(double));
    Algorithm      *algo = NULL;
    Statistics     *stats = NULL;
    Cds           **cds = NULL;
    Cds            *avecds = NULL;
    CdsArray       *scratchA = NULL;

    gsl_rng               *r2 = NULL;
    const gsl_rng_type    *T = NULL;

    T = gsl_rng_ranlxs2;
    r2 = gsl_rng_alloc(T);

    // THREAD STUFF /////////////////////////////////////////////////////
    const int       thrdnum = algo->threads;
    RotData       **rotdata = malloc(thrdnum * sizeof(RotData *));
    AveData       **avedata = malloc(thrdnum * sizeof(AveData *));
    pthread_t      *callThd = malloc(thrdnum * sizeof(pthread_t));
    pthread_attr_t  attr;

    for (i = 0; i < thrdnum; ++i)
    {
        rotdata[i] = malloc(sizeof(RotData));
        avedata[i] = malloc(sizeof(AveData));
    }

    pthread_attr_init(&attr);
/*     pthread_attr_getstacksize (&attr, &stacksize); */
/*     printf("\nDefault stack size = %d", (int) stacksize); */
    pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);
    pthread_attr_setscope(&attr, PTHREAD_SCOPE_SYSTEM);
    // THREAD STUFF /////////////////////////////////////////////////////

    /* setup scratchA */
    scratchA = CdsArrayInit();
    CdsArrayAlloc(scratchA, cnum, vlen);
    CdsArraySetup(scratchA);

    baseA->scratchA = scratchA;

    /* duplicate baseA -- copy to scratchA */
    CdsArrayCopy(scratchA, baseA);

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

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

    slxn = gsl_rng_uniform_int(r2, cnum);
    CdsCopyAll(avecds, baseA->cds[slxn]);

    if (algo->dotrans)
    {
        CenMass(avecds);
        ApplyCenterIp(avecds);
    }

    /* 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)
    {
        if (algo->nullrun)
            break;

        ++round;
        algo->rounds = round;

        /* Find weighted center and translate all cds */
        if (algo->dotrans)
        {
            CalcTranslationsOp(baseA, 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

           Loops till convergence, holding constant the weights, variances, and covariances
           (and thus the translations too) */
        innerround = 0;
        do
        {
            ++innerround;
            algo->innerrounds += innerround;

            /* save the old rotation matrices to test convergence at bottom of loop */
            for (i = 0; i < cnum; ++i)
                MatCpySqr(cds[i]->last_matrix, (const double **) cds[i]->matrix, 3);

            /* find the optimal rotation matrices */
            if (algo->dorot)
            {
                if (algo->alignment)
                {
                    CalcRotations(scratchA);
                }
                else
                {
                    // THREAD STUFF /////////////////////////////////////////////////////
                    CalcRotations_pth(scratchA, rotdata, callThd, &attr, thrdnum);
                    // THREAD STUFF /////////////////////////////////////////////////////
                }
            }

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

            /* find global rmsd and average cds (both held in structure) */
            if (algo->doave)
            {
                if (algo->alignment)
                {
                    AveCdsNu(scratchA);
                    EM_MissingCds(scratchA);
                }
                else
                {
                    // THREAD STUFF /////////////////////////////////////////////////////
                    AveCds_pth(scratchA, avedata, callThd, &attr, thrdnum);
                    // THREAD STUFF /////////////////////////////////////////////////////
                }
            }

            //stats->wRMSD_from_mean = sqrt(deviation_sum / (3 * vlen * cnum));

            if (algo->noinnerloop)
                break;
            else if (innerround > 160)
            {
                putchar(',');
                fflush(NULL);
                break;
            }
        }
        while(CheckConvergenceInner(scratchA, algo->precision) == 0);

        if (algo->docovars)
        {
            CalcCovariances(scratchA);
            if (algo->varweight || algo->covweight)
                HierarchVars(scratchA);
        }

        CalcWts(scratchA);
    }

    outsidetheloops:

    CdsArrayDestroy(&scratchA);
    MyFree(evals);

    // THREAD STUFF /////////////////////////////////////////////////////
    pthread_attr_destroy(&attr);
    for (i = 0; i < thrdnum; ++i)
        MyFree(rotdata[i]);
    for (i = 0; i < thrdnum; ++i)
        MyFree(avedata[i]);
    MyFree(rotdata);
    MyFree(avedata);
    MyFree(callThd);
    // THREAD STUFF /////////////////////////////////////////////////////

    gsl_rng_free(r2);
    r2 = NULL;

    return(round);
}