#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <math.h>
#include <mpi.h>
#include "parmt_polarity.h"
#ifdef PARMT_USE_INTEL
#include <mkl_cblas.h>
#else
#include <cblas.h>
#endif
#include "ttimes.h"
#include "iscl/array/array.h"
#include "iscl/geodetic/geodetic.h"
#include "iscl/memory/memory.h"

#define LDG 8
#ifndef MAX
#define MAX(x,y) (((x) > (y)) ? (x) : (y))
#endif

static void fillBasis(const double i, const double phi,
                      double *__restrict__ gam,
                      double *__restrict__ phat,
                      double *__restrict__ phihat);
static double computeContraction3x3(const double *__restrict__ a,
                                    const double *__restrict__ M,
                                    const double *__restrict__ b);
static void setM3x3(const int k, double *__restrict__ M);
static int computePadding64f(const int n);
static int performPolaritySearch64f(const int nmt, const int ldm,
                                    const int nobs,
                                    const int blockSize, const int mblock,
                                    const int Mrows, const int Kcols,
                                    const double *__restrict__ Dmat,
                                    const double *__restrict__ G,
                                    const double *__restrict__ Sigma,
                                    const double *__restrict__ mts,
                                    double *__restrict__ phi);

/*!
 * @brief Driver routine for computing the polarity Green's functions from the
 *        ttimes ak135 global travel time table.
 *
 * @param[in] globalComm     Global MPI communicator.
 * @param[in] parms          Contains the polarity modeling parameters.
 * @param[in] data           The SAC data whose header information will define
 *                           the polarity and channel information.
 *
 * @param[out] polarityData  On exit contains the corresponding polarity Green's
 *                           functions that can be used by the grid-search to
 *                           estimate polarities from a given moment tensor.
 *
 * @result 0 indicates success.
 */
int parmt_polarity_computeTTimesGreens(
    const MPI_Comm globalComm,
    const struct parmtPolarityParms_struct parms,
    const struct parmtData_struct data,
    struct polarityData_struct *polarityData)
{
    char kt0[8], kcmpnm[8], stat[8];
    double G6[6], *cmpazs, *cmpincs, *deps, *evlas, *evlos,
           *GxxBuf, *GyyBuf, *GzzBuf, *GxyBuf, *GxzBuf, *GyzBuf,
           *stlas, *stlos, cmpinc, cmpincSAC, cmpaz, stla, stlo;
    int *icomps, *observation, *polarity, *waveType, icomp, ierr, ierrAll,
        iloc, iobs, ipol, it, iwav, jloc, k, kt, myid, nPolarity, nprocs;
    size_t lenos;
    const int master = 0;
    const int nTimeVars = 11;
    const enum sacHeader_enum timeVarNames[11]
       = {SAC_CHAR_KA,
          SAC_CHAR_KT0, SAC_CHAR_KT1, SAC_CHAR_KT2, SAC_CHAR_KT3,
          SAC_CHAR_KT4, SAC_CHAR_KT5, SAC_CHAR_KT6, SAC_CHAR_KT7,
          SAC_CHAR_KT8, SAC_CHAR_KT9};
    //------------------------------------------------------------------------//
    //
    // Initialize
    ierr = 0;
    cmpazs = NULL;
    cmpincs = NULL;
    deps = NULL;
    evlas = NULL;
    evlos = NULL;
    GxxBuf = NULL;
    GyyBuf = NULL;
    GzzBuf = NULL;
    GxyBuf = NULL;
    GxzBuf = NULL;
    GyzBuf = NULL;
    icomps = NULL;
    observation = NULL;
    polarity = NULL;
    stlas = NULL;
    stlos = NULL;
    waveType = NULL;
    memset(polarityData, 0, sizeof(struct polarityData_struct));
    // ttimes uses read-only file-io and i don't want to fix it
    MPI_Comm_rank(globalComm, &myid);
    MPI_Comm_size(globalComm, &nprocs);
    //if (myid != master){goto WAIT;}
    // Verify there is a chance for something to do
    if (data.nobs < 1 || data.nlocs < 1)
    {
        if (data.nobs < 1){fprintf(stderr, "%s: No observations\n", __func__);}
        if (data.nlocs < 1){fprintf(stderr, "%s: No locations\n", __func__);}
        return 0;
    }
    // Extract the depths in the grid-search from the first waveform
    iobs = 0;
    deps = memory_calloc64f(data.nlocs);
    evlas = memory_calloc64f(data.nlocs);
    evlos = memory_calloc64f(data.nlocs);
    for (iloc=0; iloc<data.nlocs; iloc++)
    {
        k = iobs*data.nlocs + iloc;
        ierr  = sacio_getFloatHeader(SAC_FLOAT_EVDP, data.sacGxx[k].header,
                                     &deps[iloc]);
        ierr += sacio_getFloatHeader(SAC_FLOAT_EVLA, data.sacGxx[k].header,
                                     &evlas[iloc]);
        ierr += sacio_getFloatHeader(SAC_FLOAT_EVLO, data.sacGxx[k].header,
                                     &evlos[iloc]);
        if (ierr != 0)
        {
            fprintf(stderr, "%s: Unable to get event coordinates on %d\n",
                    __func__, myid);
            goto ERROR;
        }
    }
    // Compute the number of usable polarities (i.e. observations with P picks
    // or S picks with first motions)
    nPolarity = 0;
    if (myid == master)
    {
        polarity = memory_calloc32i(data.nobs);
        waveType = memory_calloc32i(data.nobs);
        icomps = memory_calloc32i(data.nobs);
        stlas = memory_calloc64f(data.nobs);
        stlos = memory_calloc64f(data.nobs);
        cmpazs = memory_calloc64f(data.nobs);
        cmpincs = memory_calloc64f(data.nobs);
        observation = array_set32i(data.nobs, -1, &ierr);
        for (iobs=0; iobs<data.nobs; iobs++)
        {
            // Ensure the essential preliminary information is defined
            ierr  = sacio_getFloatHeader(SAC_FLOAT_CMPINC,
                                         data.data[iobs].header, &cmpincSAC);
            ierr += sacio_getFloatHeader(SAC_FLOAT_CMPAZ,
                                         data.data[iobs].header, &cmpaz);
            ierr += sacio_getFloatHeader(SAC_FLOAT_STLA,
                                         data.data[iobs].header, &stla);
            ierr += sacio_getFloatHeader(SAC_FLOAT_STLO,
                                         data.data[iobs].header, &stlo);
            ierr += sacio_getCharacterHeader(SAC_CHAR_KCMPNM,
                                             data.data[iobs].header, kcmpnm);
            if (ierr != 0)
            {
                fprintf(stderr, "%s: Failed to get header information\n",
                        __func__);
                continue;
            }
            // Figure out the component
            ierr = parmt_utils_getComponent(kcmpnm, cmpincSAC, &icomp);
            if (ierr != 0)
            {
                fprintf(stderr, "%s: Failed to classify component\n", __func__);
            }
            // SAC to SEED convention
            cmpinc = cmpincSAC - 90.0;
            // Figure out the component
/*
            lenos = MAX(1, strlen(kcmpnm));
            icomp = 1;
            if (kcmpnm[lenos-1] == 'Z' || kcmpnm[lenos-1] == '1')
            {
                icomp = 1;
            }
            else if (kcmpnm[lenos-1] == 'N' || kcmpnm[lenos-1] == '2')
            {
                icomp = 2;
            }
            else if (kcmpnm[lenos-1] == 'E' || kcmpnm[lenos-1] == '3')
            {
                icomp = 3;
            }
            else
            {
                fprintf(stderr, "%s: Cannot classify component %s\n", __func__, kcmpnm);
                continue;
            }
*/
            // Get the primary pick
            for (it=0; it<nTimeVars; it++)
            {
                ierr = sacio_getCharacterHeader(timeVarNames[it],
                                                data.data[iobs].header,
                                                kt0);
                if (ierr == 0){break;}
            }
            lenos = strlen(kt0);
            // Ensure the pick is defined
            if (ierr == 0 && lenos > 1)
            {
                iwav = 1;
                if (kt0[0] == 'P' || kt0[0] == 'p')
                {
                    iwav = 1;
                }
                else if (kt0[0] == 'S' || kt0[0] == 's')
                {
                    iwav = 2;
                }
                else
                {
                    // surface waves will commonly be processed and not
                    // have polarities
                    if (kt0[0] == 'R' || kt0[0] == 'r' ||
                        kt0[0] == 'L' || kt0[0] == 'l')
                    {
                        continue;
                    }
                    fprintf(stderr, "%s: t0 phase is not a P or S phase %s\n",
                            __func__, kt0);
                    continue;
                }
                if (kt0[lenos-1] == '+')
                {
                    ipol = 1;
                }
                else if (kt0[lenos-1] == '-')
                {
                    ipol =-1;
                }
                else
                {
                    fprintf(stderr, "%s: could not classify polarity %s\n",
                            __func__, kt0);
                    continue;
                }
                // let user know something happened
                sacio_getCharacterHeader(SAC_CHAR_KSTNM,
                                         data.data[iobs].header, stat);
                fprintf(stdout, "%s: Polarity for %s is %d\n",
                        __func__, stat, ipol);
//printf("%f %f %f %f %d %d %d %d\n", stla, stlo, cmpinc, cmpaz, icomp, iobs, iwav, ipol);
                // save information for modeling
                stlas[nPolarity] = stla;
                stlos[nPolarity] = stlo;
                cmpincs[nPolarity] = cmpinc;
                cmpazs[nPolarity] = cmpaz;
                icomps[nPolarity] = icomp;
                observation[nPolarity] = iobs;
                waveType[nPolarity] = iwav;
                polarity[nPolarity] = ipol;
                nPolarity = nPolarity + 1;
                //printf("%d %d %f %f\n", iwav, ipol, cmpinc, cmpaz);
           } // End basic header info check
        } // Loop on observations
    } // End check on myid == master
    // Tell all other processes about the forward modeling information
    MPI_Bcast(&nPolarity, 1, MPI_INT, master, globalComm);
    if (nPolarity == 0)
    {
        if (myid == master)
        {
            fprintf(stdout, "%s: There are no polarities\n", __func__);
        }
        ierr = 0;
        goto ERROR;
    }
    if (myid != master)
    {
        stlas = memory_calloc64f(nPolarity);
        stlos = memory_calloc64f(nPolarity);
        cmpincs = memory_calloc64f(nPolarity);
        cmpazs = memory_calloc64f(nPolarity);
        icomps = memory_calloc32i(nPolarity);
        observation = memory_calloc32i(nPolarity);
        waveType = memory_calloc32i(nPolarity);
        polarity = memory_calloc32i(nPolarity);
    }
    else
    {
        fprintf(stdout, "%s: Warning - i'm setting wts to unity for now\n",
                __func__);
    }
    MPI_Bcast(stlas,   nPolarity, MPI_DOUBLE, master, globalComm);
    MPI_Bcast(stlos,   nPolarity, MPI_DOUBLE, master, globalComm);
    MPI_Bcast(cmpincs, nPolarity, MPI_DOUBLE, master, globalComm);
    MPI_Bcast(cmpazs,  nPolarity, MPI_DOUBLE, master, globalComm);
    MPI_Bcast(icomps,      nPolarity, MPI_INT, master, globalComm);
    MPI_Bcast(observation, nPolarity, MPI_INT, master, globalComm);
    MPI_Bcast(waveType,    nPolarity, MPI_INT, master, globalComm);
    MPI_Bcast(polarity,    nPolarity, MPI_INT, master, globalComm);
    // Set space
    ierrAll = 0;
    GxxBuf = memory_calloc64f(data.nlocs*nPolarity);
    GyyBuf = memory_calloc64f(data.nlocs*nPolarity);
    GzzBuf = memory_calloc64f(data.nlocs*nPolarity);
    GxyBuf = memory_calloc64f(data.nlocs*nPolarity);
    GxzBuf = memory_calloc64f(data.nlocs*nPolarity);
    GyzBuf = memory_calloc64f(data.nlocs*nPolarity);
    polarityData->Gxx = memory_calloc64f(data.nlocs*nPolarity);
    polarityData->Gyy = memory_calloc64f(data.nlocs*nPolarity);
    polarityData->Gzz = memory_calloc64f(data.nlocs*nPolarity);
    polarityData->Gxy = memory_calloc64f(data.nlocs*nPolarity);
    polarityData->Gxz = memory_calloc64f(data.nlocs*nPolarity);
    polarityData->Gyz = memory_calloc64f(data.nlocs*nPolarity);
    // Set the data
    polarityData->polarity = memory_calloc64f(nPolarity);
    for (ipol=0; ipol<nPolarity; ipol++)
    {
        polarityData->polarity[ipol] = (double) polarity[ipol];
    }
    // Set the weights
    // TODO fix me
    polarityData->wts = array_set64f(nPolarity, 1.0, &ierr);
    // Compute the forward modeling matrix columns
    for (jloc=0; jloc<data.nlocs; jloc=jloc+nprocs)
    {
        iloc = jloc + myid;
        if (iloc >= data.nlocs){continue;}
        // Compute the forward modeling matrix for this observation group
        for (ipol=0; ipol<nPolarity; ipol++)
        {
            kt = iloc*nPolarity + ipol;
            ierr = parmt_polarity_computeGreensRowFromTtimes(
                                waveType[ipol], icomps[ipol],
                                evlas[iloc], evlos[iloc], deps[iloc],
                                stlas[ipol], stlos[ipol],
                                cmpincs[ipol], cmpazs[ipol],
                                parms.ttimesTablesDir,
                                parms.ttimesModel,
                                G6);
            if (ierr != 0)
            {
                fprintf(stderr,
                        "%s: Error computing polarities %d %d on PID %d\n",
                       __func__, iloc, ipol, myid);
                ierrAll = ierrAll + 1;
                continue;
            }
            // Save it
            GxxBuf[kt] = G6[0];
            GyyBuf[kt] = G6[1];
            GzzBuf[kt] = G6[2];
            GxyBuf[kt] = G6[3];
            GxzBuf[kt] = G6[4];
            GyzBuf[kt] = G6[5];
            //printf("%f %f %f %e %e %e %e %e %e\n", stlas[ipol], stlos[ipol], deps[iloc], G6[0], G6[1], G6[2], G6[3], G6[4], G6[5]);
        } // Loop on the polarities
        //printf("\n");
    } // Loop on the locations
    ierr = ierrAll;
    if (ierr != 0)
    {
        fprintf(stderr, "%s: Error computing polarities from ttimes\n",
                __func__);
        ierr = 1;
        goto ERROR;
    }
    polarityData->nPolarity = nPolarity;
    polarityData->nlocs = data.nlocs;
    polarityData->nobs = data.nobs;
    polarityData->obsMap = array_copy32i(nPolarity, observation, &ierr);
    MPI_Allreduce(GxxBuf, polarityData->Gxx, data.nlocs*nPolarity,
                  MPI_DOUBLE, MPI_SUM, globalComm);
    MPI_Allreduce(GyyBuf, polarityData->Gyy, data.nlocs*nPolarity,
                  MPI_DOUBLE, MPI_SUM, globalComm);
    MPI_Allreduce(GzzBuf, polarityData->Gzz, data.nlocs*nPolarity,
                  MPI_DOUBLE, MPI_SUM, globalComm);
    MPI_Allreduce(GxyBuf, polarityData->Gxy, data.nlocs*nPolarity,
                  MPI_DOUBLE, MPI_SUM, globalComm);
    MPI_Allreduce(GxzBuf, polarityData->Gxz, data.nlocs*nPolarity,
                  MPI_DOUBLE, MPI_SUM, globalComm);
    MPI_Allreduce(GyzBuf, polarityData->Gyz, data.nlocs*nPolarity,
                  MPI_DOUBLE, MPI_SUM, globalComm);
    // Finally assemble Green's functions into modeling matrix
//WAIT:;
ERROR:;
    MPI_Barrier(globalComm);
    // free memory
    memory_free64f(&GxxBuf);
    memory_free64f(&GyyBuf);
    memory_free64f(&GzzBuf);
    memory_free64f(&GxyBuf);
    memory_free64f(&GxzBuf);
    memory_free64f(&GyzBuf);
    memory_free64f(&deps);
    memory_free64f(&evlas);
    memory_free64f(&evlos);
    memory_free64f(&stlas);
    memory_free64f(&stlos);
    memory_free64f(&cmpincs);
    memory_free64f(&cmpazs);
    memory_free32i(&observation);
    memory_free32i(&polarity);
    memory_free32i(&icomps);
    memory_free32i(&waveType);
    return ierr;
}
//============================================================================//
/*!
 * @brief Computes the row of the Green's functions matrix for modeling
 *        a polarity from the ttimes model.
 *
 * @param[in] data      Contains the pick type and channel information.
 * @param[in] wavetype  P_WAVE (1) indicates a P-wave.
 * @param[in] wavetype  S_WAVE (2) indicates an S-wave.
 * @param[in] evdp      Depth of the event in kilometers.
 * @param[in] dirnm     Directory containing the iasp-tau binary files.
 * @param[in] model     Name of the model, e.g., ak135.
 *
 * @param[out] G        On exit contains the Green's functions so that for
 *                      the i'th row of G, \f$ G_i \cdot \textbf{m} \f$
 *                      computes an estimate of polarity.  This is packed
 *                      in order
 *                      \f$
 *                        \{m_{xx}, m_{yy}, m_{zz}, m_{xy}, m_{xz}, m_{yz} \}
 *                      \f$
 *                      and must have a dimension of at least [6].
 *
 * @result 0 indicates success.
 *
 * @copyright ISTI distributed under the Apache 2 license.
 *
 */
int parmt_polarity_computeGreensRowFromData(const struct sacData_struct data,
                                            const int wavetype,
                                            const double evdp,
                                            const char *dirnm,
                                            const char *model,
                                            double *__restrict__ G)
{
    char kcmpnm[16];
    double cmpaz, cmpinc, cmpincSAC, evla, evlo, stla, stlo;
    int icomp, ierr;
    size_t lenos;
    memset(kcmpnm, 16, 16*sizeof(char));
    ierr  = sacio_getFloatHeader(SAC_FLOAT_CMPINC, data.header, &cmpincSAC);
    ierr += sacio_getFloatHeader(SAC_FLOAT_CMPAZ,  data.header, &cmpaz);
    ierr += sacio_getFloatHeader(SAC_FLOAT_EVLA,   data.header, &evla);
    ierr += sacio_getFloatHeader(SAC_FLOAT_EVLO,   data.header, &evlo);
    ierr += sacio_getFloatHeader(SAC_FLOAT_STLA,   data.header, &stla);
    ierr += sacio_getFloatHeader(SAC_FLOAT_STLO,   data.header, &stlo);
    ierr += sacio_getCharacterHeader(SAC_CHAR_KCMPNM, data.header, kcmpnm);
    if (ierr != 0)
    {
        fprintf(stderr, "%s: Failed to get header information\n", __func__);
        return -1;
    }
    // Figure out the component
    ierr = parmt_utils_getComponent(kcmpnm, cmpincSAC, &icomp);
    if (ierr != 0)
    {
        fprintf(stderr, "%s: Failed to classify component\n", __func__);
        return -1;
    }
    // SAC to SEED convention
    cmpinc = cmpincSAC - 90.0;
    // Figure out the component
    lenos = MAX(1, strlen(kcmpnm));
    icomp = 1;
    if (kcmpnm[lenos-1] == 'Z' || kcmpnm[lenos-1] == '1')
    {
        icomp = 1;
    }
    else if (kcmpnm[lenos-1] == 'N' || kcmpnm[lenos-1] == '2')
    {
        icomp = 2;
    }
    else if (kcmpnm[lenos-1] == 'E' || kcmpnm[lenos-1] == '3')
    {
        icomp = 3;
    }
    else
    {
        fprintf(stderr, "%s: Can't classify component %s\n", __func__, kcmpnm);
        return -1;
    }
    ierr = parmt_polarity_computeGreensRowFromTtimes(wavetype, icomp,
                                                     evla, evlo, evdp,
                                                     stla, stlo,
                                                     cmpinc, cmpaz,
                                                     dirnm, model,
                                                     G);
    if (ierr != 0)
    {
        fprintf(stderr, "%s: Failed to compute G\n", __func__);
    }
    return ierr;
}
//============================================================================//
/*!
 * @brief Computes a row of the Green's functions matrix from ttimes
 *
 * @param[in] wavetype   If 1 then this is a P wave.
 * @param[in] wavetype   If 2 then this is an S wave.
 * @param[in] icomp      If 1 then this is the vertical channel.
 * @param[in] icomp      If 2 then this is the north (1 or 2) channel.
 * @param[in] icomp      If 3 then this is the east (2 or 3) channel.
 * @param[in] evla       Event latitude (degrees).
 * @param[in] evlo       Event longitude (degrees).
 * @param[in] evdp       Event depth (km).
 * @param[in] stla       Station latitude (degrees).
 * @param[in] stlo       Station longitude (degrees).
 * @param[in] cmpaz      Component azimuth (0 north, +90 east).
 * @param[in] cmpinc     Component inclination (-90 up, 0 east/north, +90 down).
 * @param[in] dirnm      Directory containing the ttimes precomputed binary
 *                       files.  If NULL then the default as dicated by the
 *                       ttimes configuration will be used.
 * @param[in] model      Model name (e.g., ak135 or iasp91)
 *
 * @param[out] G         Row of matrix s.t. G*m produces estimates the polarity
 *                       at the station.  Here m is packed
 *                       \f$ \{m_{xx}, m_{yy}, m_{zz},
 *                             m_{xy}, m_{xz}, m_{yz} \} \f$
 *                       This must have dimension of at least 6.
 *
 * @author Ben Baker
 *
 * @copyright ISTI distributed under Apache 2
 *
 */
int parmt_polarity_computeGreensRowFromTtimes(
    const int wavetype, const int icomp,
    const double evla, const double evlo, const double evdp,
    const double stla, const double stlo,
    const double cmpinc, const double cmpaz,
    const char *dirnm, const char *model,
    double *__restrict__ G)
{
    double aoiRec, az, azSrc, baz, bazRec, dist, delta, toaSrc;
    struct ttimesTravelTime_struct ttime;
    int ierr;
    ierr = 0;
    if (G == NULL)
    {
        fprintf(stderr, "%s: Error G is NULL\n", __func__);
        return -1;
    }
    memset(G, 0, 6*sizeof(double));
    if (evdp < 0.0 || evdp > ttimes_getMaxDepth())
    {
        fprintf(stderr, "%s: Error depth must be between [0,%f]\n", __func__,
                ttimes_getMaxDepth());
        return -1;
    }
    memset(&ttime, 0, sizeof(struct ttimesTravelTime_struct));
    geodetic_gps2distanceAzimuth(evla, evlo, stla, stlo,
                                 &dist, &delta, &az, &baz);
    if (wavetype == 1)
    {
        ierr = ttimes_getFirstPPhase(delta, evdp, dirnm, model, &ttime);
    }
    else if (wavetype == 2)
    {
        ierr = ttimes_getFirstSPhase(delta, evdp, dirnm, model, &ttime);
    }
    else
    {
        fprintf(stderr, "%s: Invalid phase type - must be 1 (P) or 2 (S)\n",
                __func__);
        return -1;
    }
    if (ierr != 0)
    {
        fprintf(stderr, "%s: Error computing theoretical traveltime info\n",
                 __func__);
        return -1;
    }
    // Compute the column in the Green's function matrix
    toaSrc = ttime.toang;
    azSrc  = az;
    bazRec = baz;
    aoiRec = ttime.aoi;
    //printf("%f %f %f %f %f %f %f\n", delta, stla, stlo, toaSrc, azSrc, aoiRec, bazRec);
    ierr = parmt_polarity_computeGreensMatrixRow(wavetype, icomp,
                                                 azSrc, toaSrc, bazRec, aoiRec,
                                                 cmpinc, cmpaz, G);

    if (ierr != 0)
    {
        fprintf(stderr, "%s: Failed to compute polarity for row\n", __func__);
        memset(G, 0, 6*sizeof(double));
        return -1;
    }
    return 0;
}
//============================================================================//
/*!
 * @brief Computes a column for the Green's function s.t. G*m produces an
 *        estimate of polarity on the icomp'th component measured in the
 *        far-field.  The moment tensor which would be applied to this row
 *        is packed:
 *          \f$ \{m_{xx}, m_{yy}, m_{zz}, m_{xy}, m_{xz}, m_{yz} \} \f$
 *        with convention North, East, Down (e.g., Jost and Herrmann).
 *        For more see Quantitative Seismology - Aki and Richards 2002,
 *        Eqn 4.96 on pg 111 and Source Mechanisms of Earthquakes: Theory
 *        and Practice - Udias et al. pg 100..
 *
 * @param[in] wavetype   =1 -> P wave
 *                       =2 -> S wave
 * @param[in] icomp      receiver component of motion:
 *                         =1 -> Vertical channel
 *                         =2 -> 1 or North channel
 *                         =3 -> 2 or East channel
 * @param[in] azSrc      source to receiver azimuth is measured positive from
 *                       north (degrees)
 * @param[in] toaSrc     take-off angle (measured positive from x3 where x3
 *                       points down) (degrees)
 * @param[in] bazRec     receiver to source back azimuth measured positive
 *                       from north (degrees)
 * @param[in] aoiRec     angle of incidence at receiver
 * @param[in] cmpaz      component azimuth (0 north, +90 east)
 * @param[in] cmpinc     component inclinantion (-90 up, 0 east/north, +90 down)
 *
 * @param[out] G         row of matrix s.t. G*m produces estimates the polarity
 *                       at the station.  Here m is packed
 *                       \f$ \{m_{xx}, m_{yy}, m_{zz},
 *                             m_{xy}, m_{xz}, m_{yz} \} \f$
 *
 * @bugs i've really only looked at the vertical p-teleseismic case - no idea
 *        about other phases, wavetypes
 *
 * @author Ben Baker
 *
 * @copyright ISTI distributed under Apache 2
 *
 */
int parmt_polarity_computeGreensMatrixRow(const int wavetype,
                                          const int icomp,
                                          const double azSrc,
                                          const double toaSrc,
                                          const double bazRec,
                                          const double aoiRec,
                                          const double cmpinc,
                                          const double cmpaz,
                                          double *__restrict__ G)
{
    double M[9], G63[18], up[6], ush[6], usv[6],
           gam[3], lhat[3], phat[3], phihat[3],
           cosba, cos_cmpaz, cost_rec,
           r11, r12, r13, r21, r22, r23, r31, r32, r33,
           sinba, sin_cmpaz,
           sint_rec,
           t1, t2, t3, theta, xsign, u1, u2, ue, un, uz;
    int k;
    const double pi180 = M_PI/180.0;
    const bool lrot = true;//false;
    const int P_WAVE = 1;
    const int S_WAVE = 2;
    //------------------------------------------------------------------------//
    //
    // Error check
    if (icomp < 1 || icomp > 3)
    {
        fprintf(stderr, "%s: Invalid component\n", __func__);
        return -1;
    }
    if (wavetype < P_WAVE || wavetype > S_WAVE)
    {
        fprintf(stderr, "%s: Invalid wavetype\n", __func__);
        return -1;
    }
    // Fill the basis at the source (Aki and Richards Eqn 4.88)
    fillBasis(toaSrc, azSrc, gam, phat, phihat);
    // Compute the contractions for u_p, u_sv, and u_sh (A&R Eqn 4.96)
    // for all 6 individual moment tensor terms
    for (k=0; k<6; k++)
    {
        // Set the moment tensor with the k'th mt term 1 others zero
        setM3x3(k, M);
        // Compute contraction
        up[k]  = computeContraction3x3(gam,    M, gam);
        usv[k] = computeContraction3x3(phat,   M, gam);
        ush[k] = computeContraction3x3(phihat, M, gam);
    }
    // Fill the basis at the receiver - notice the basis uses the forward
    // azimuth from the receiver to the source
    fillBasis(aoiRec, (bazRec + 180.0), lhat, phat, phihat);
    // Compute geometric factors at receiver
//printf("%f %f %f %f %f\n", toaSrc, azSrc, bazRec, aoiRec, cmpaz);
    theta = aoiRec*pi180;
    cosba     = cos(bazRec*pi180);
    cost_rec  = cos(theta);
    sinba     = sin(bazRec*pi180);
    sint_rec  = sin(theta);
    cos_cmpaz = cos(cmpaz*pi180);
    sin_cmpaz = sin(cmpaz*pi180);
    // Set 3D rotation matrix
    r11 = cost_rec;
    r21 = sint_rec;
    r31 = 0.0;
    r12 =-sint_rec*sinba;
    r22 = cost_rec*sinba;
    r32 =         -cosba;
    r13 =-sint_rec*cosba;
    r23 = cost_rec*cosba;
    r33 =          sinba;
    // Flip sign for receivers that acquire positive down
    xsign = 1.0;
    if (fabs(cmpinc - 90.0) < 1.e-4){xsign =-1.0;}
    // Compute 6 x 3 subforward modeling matrix with row for up (1 channel)
    if (wavetype == P_WAVE)
    {
        // Loop on mts terms
        for (k=0; k<6; k++)
        {
            // Extract the (north, east, down) component
            t3 =-up[k]*lhat[0]; // z-down -> z-up
            //t3 = up[k]*lhat[0]; // z-up
            t2 = up[k]*lhat[1]; // east
            t1 = up[k]*lhat[2]; // north
            // Not sure if i have to rotate LQT -> ZNE
            if (lrot)
            {
                uz = t1*r11 + t2*r21 + t3*r31; // Z
                ue = t1*r12 + t2*r22 + t3*r32; // E
                un = t1*r13 + t2*r23 + t3*r33; // N
            }
            else
            {
                uz = t3;
                ue = t2;
                un = t1;
            }
            // Rotate into (1, 2)
            u1 = un*cos_cmpaz + ue*sin_cmpaz;
            u2 =-un*sin_cmpaz + ue*cos_cmpaz;
            // Finish
            G63[k*3+0] = xsign*uz;
            //G63[k*3+0] =-xsign*uz;
            G63[k*3+1] = u1;
            G63[k*3+2] = u2;
        }
    }
    // SH wave
    else
    {
        // Loop on mts terms
        for (k=0; k<6; k++)
        {
            // Extract the (north, east, down) component
            t3 =-usv[k]*phat[0] - ush[k]*phihat[0]; // z-down -> z-up
            t2 = usv[k]*phat[1] + ush[k]*phihat[1]; // east
            t1 = usv[k]*phat[2] + ush[k]*phihat[2]; // north
            // Not sure if i have to rotate LQT -> ZNE
            if (lrot)
            {
                uz = t1*r11 + t2*r21 + t3*r31; // Z
                ue = t1*r12 + t2*r22 + t3*r32; // E
                un = t1*r13 + t2*r23 + t3*r33; // N
            }
            else
            {
                uz = t3;
                ue = t2;
                un = t1;
            }
            // Rotate into (1, 2)
            u1 = un*cos_cmpaz + ue*sin_cmpaz;
            u2 =-un*sin_cmpaz + ue*cos_cmpaz;
            // Finish
            G63[k*3+0] = xsign*uz;
            G63[k*3+1] = u1;
            G63[k*3+2] = u2;
        }
    }
    // Copy the result - vertical
    if (icomp == 1)
    {
        for (k=0; k<6; k++)
        {
            G[k] = G63[k*3+0];
        }
    }
    // 1 component
    else if (icomp == 2)
    {
        for (k=0; k<6; k++)
        {
            G[k] = G63[k*3+1];
        }
    }
    // 2 component
    else if (icomp == 3)
    {
        for (k=0; k<6; k++)
        {
            G[k] = G63[k*3+2];
        }
    }
    return 0;
}
//============================================================================//
/*!
 * @brief Fills in the basis vectors - Aki and Richards pg 108 Eqn 4.88.
 *
 * @param[in] i       Take-off angle (degrees).
 * @param[in] phi     Azimuth angle (degrees).
 * @param[in] gam     P-wave direction.  This has dimension [3].
 * @param[in] phat    SV-wave direction.  This has dimension [3].
 * @param[in] phihat  SH-wave direction.  This has dimension [3].
 *
 * @copyright ISTI distributed under the Apache 2 license.
 *
 */
static void fillBasis(const double i, const double phi,
                      double *__restrict__ gam,
                      double *__restrict__ phat,
                      double *__restrict__ phihat)
{
    double cosi, sini, cosp, sinp;
    const double pi180 = M_PI/180.0;
    cosi = cos(i*pi180);
    sini = sin(i*pi180);
    cosp = cos(phi*pi180);
    sinp = sin(phi*pi180);
    gam[0] = sini*cosp;
    gam[1] = sini*sinp;
    gam[2] = cosi;

    phat[0] = cosi*cosp;
    phat[1] = cosi*sinp;
    phat[2] =-sini;

    phihat[0] =-sinp;
    phihat[1] = cosp;
    phihat[2] = 0.0;
}
//============================================================================//
/*!
 * @brief Sets the moment tensor matrix where the k'th moment tensor
 *        term is 1 and others are zero.  Here moment tensor terms
 *        are counted {0,1,2,3,4,5} = {xx,yy,zz,xy,xz,yz}
 *
 * @param[in] k    Moment tensor index.  This is in the range [0,5] and follows
 *                 the mapping {0,1,2,3,4,5} = {xx,yy,zz,xy,xz,yz}.
 *
 * @param[out] M   The 3x3 NED moment tensor. This is an array of dimension [9].
 *
 */
static void setM3x3(const int k, double *__restrict__ M)
{
    memset(M, 0, 9*sizeof(double));
    // mxx (fill mtt)
    if (k == 0)
    {
        M[4] = 1.0; //M[0][0] = 1.0;
    }
    // myy (fill mpp)
    else if (k == 1)
    {
        M[8] = 1.0; //M[1][1] = 1.0;
    }
    // mzz (fill mrr)
    else if (k == 2)
    {
        M[0] = 1.0; //M[2][2] = 1.0;
    }
    // mxy and myz (fill mtp)
    else if (k == 3)
    {
        M[5] =-1.0; //M[0][1] = 1.0;
        M[7] =-1.0; //M[1][0] = 1.0;
    }
    // mxz and mzx (fill mrp)
    else if (k == 4)
    {
        M[1] = 1.0; //M[0][2] = 1.0;
        M[3] = 1.0; //M[2][0] = 1.0;
    }
    // myz and mzy (fill mrp)
    else
    {
       M[2] =-1.0; //M[1][2] = 1.0;
       M[6] =-1.0; //M[2][1] = 1.0;
    }
/*
    // mxx
    if (k == 0)
    {
        M[0] = 1.0; //M[0][0] = 1.0;
    }
    // myy
    else if (k == 1)
    {
        M[4] = 1.0; //M[1][1] = 1.0;
    }
    // mzz
    else if (k == 2)
    {
        M[8] = 1.0; //M[2][2] = 1.0;
    }
    // mxy and myz
    else if (k == 3)
    {
        M[1] = 1.0; //M[0][1] = 1.0;
        M[3] = 1.0; //M[1][0] = 1.0;
    }
    // mxz and mzx
    else if (k == 4)
    {
        M[2] = 1.0; //M[0][2] = 1.0;
        M[6] = 1.0; //M[2][0] = 1.0;
    }
    // myz and mzy
    else
    {
       M[5] = 1.0; //M[1][2] = 1.0;
       M[7] = 1.0; //M[2][1] = 1.0;
    }
*/
    return;
}
//============================================================================//
/*!
 * @brief Computes the contraction in Aki and Richards Eqn 4.96
 *        for the given basis vectors and moment tensor.
 */
static double computeContraction3x3(const double *__restrict__ a,
                                    const double *__restrict__ M,
                                    const double *__restrict__ b)
{
    double res;
    int i, ij, j;
    res = 0.0;
    for (i=0; i<3; i++)
    {
        for (j=0; j<3; j++)
        {
            ij = 3*i + j;
            res = res + a[i]*M[ij]*b[j];
        }
    }
    return res;
}
//============================================================================//
/*!
 * @brief Tabulates the objective function over the earthquake locations.
 *
 * @param[in] locComm       Location MPI communicator.
 * @param[in] blockSize     Block-size for performing matrix-matrix
 *                          multiplications.
 * @param[in] polarityData  Contains the polarity data and the Green's functions.
 * @param[in] mtloc         Contains the local moment tensors in this rank's
 *                          grid search.
 *
 * @param[out] phi          The objective function tabulated for all locations
 *                          and moment tensors.  This is only accessed by
 *                          master process in locComm; and for this case has
 *                          dimension [nlocs x mtloc.nmtAll] with leading
 *                          dimension mtloc.nmtAll. Otherwise, this can be NULL.
 *
 * @result 0 indicates success.
 *
 * @copyright ISTI distributed under the Apache 2 license.
 *
 */
int polarity_performLocationSearch64f(const MPI_Comm locComm,
                                      const int blockSize,
                                      struct polarityData_struct polarityData,
                                      struct localMT_struct mtloc,
                                      double *__restrict__ phi)
{
    double *Dmat, *G, *Sigma, *phiLoc, *phiWork;
    int icol, ierr, ierrAll, iloc, ipol, jloc, jndx,
        kt, mylocID, nlocProcs, mblock, pad;
    const int nPolarity = polarityData.nPolarity;
    const int Mrows = nPolarity;
    const int Kcols = 6;
    const int nlocs = polarityData.nlocs;
    const int ldm = mtloc.ldm;
    const int nmt = mtloc.nmt;
    const int master = 0;
    // Initialize
    ierr = 0;
    ierrAll = 0;
    MPI_Comm_size(locComm, &nlocProcs);
    MPI_Comm_rank(locComm, &mylocID);
    // Compute padding for 64 bit data alignment in observations and mt blocks
    pad = computePadding64f(blockSize);
    mblock = blockSize + pad;
    // Set space
    phiLoc = memory_calloc64f(nmt);
    G = memory_calloc64f(nPolarity*LDG);
    Dmat = memory_calloc64f(nPolarity*mblock);
    Sigma = array_set64f(nPolarity, 1.0, &ierr); // Default to identity
    phiWork = memory_calloc64f(nlocs*mtloc.nmtAll);
    // Set the row major data matrix where each row is an observation
    for (ipol=0; ipol<nPolarity; ipol++)
    {
        for (icol=0; icol<mblock; icol++)
        {
            Dmat[ipol*mblock+icol] = polarityData.polarity[ipol];
            if (icol == 1){printf("%f\n", polarityData.polarity[ipol]);}
        }
    }
    // Input weights are given - copy them to diagonal weight matrix
    if (polarityData.wts != NULL)
    {
        ierr = array_copy64f_work(nPolarity, polarityData.wts, Sigma);
    }
    // Loop on the source locations in the grid search
    for (jloc=0; jloc<nlocs; jloc=jloc+nlocProcs)
    {
        iloc = jloc + mylocID;
        kt = iloc*nPolarity;
        // Assemble the row major Green's functions matrix
        array_zeros64f_work(LDG*nPolarity, G);
        #pragma omp simd
        for (ipol=0; ipol<nPolarity; ipol++)
        {
            G[LDG*ipol+0] = polarityData.Gxx[kt+ipol];
            G[LDG*ipol+1] = polarityData.Gyy[kt+ipol];
            G[LDG*ipol+2] = polarityData.Gzz[kt+ipol];
            G[LDG*ipol+3] = polarityData.Gxy[kt+ipol];
            G[LDG*ipol+4] = polarityData.Gxz[kt+ipol];
            G[LDG*ipol+5] = polarityData.Gyz[kt+ipol];
        }
        // Perform the polarity search for all mt's at this location
        ierr = performPolaritySearch64f(nmt, ldm,
                                        nPolarity,
                                        blockSize, mblock,
                                        Mrows, Kcols,
                                        Dmat, G, Sigma, mtloc.mts,
                                        phiLoc);
        if (ierr != 0)
        {
            fprintf(stderr, "%s: Error in mt polarity search\n", __func__);
            ierrAll = ierrAll + 1;
        }
        // Gather the moment tensor search results onto the master
        jndx = 0;
        if (mtloc.myid == master){jndx = iloc*mtloc.nmtAll;}
        if (mtloc.commSize == 1)
        {
            array_copy64f_work(mtloc.nmt, phiLoc, &phiWork[jndx]);
        }
        else
        {
            MPI_Gatherv(phiLoc, mtloc.nmt, MPI_DOUBLE,
                        &phiWork[jndx], mtloc.nmtProc, mtloc.offset,
                        MPI_DOUBLE, master, mtloc.comm);
        }
        fprintf(stdout, "%s: mylocID: %d min: %f max: %f\n",
                __func__, mylocID, array_min64f(nmt, phiLoc, &ierr),
                array_max64f(nmt, phiLoc, &ierr));

    }
    // Reduce the search onto the master
    //if (linLocComm && mylocID == master)
    {
         MPI_Reduce(phiWork, phi, nlocs*mtloc.nmtAll, MPI_DOUBLE,
                    MPI_SUM, master, locComm);
    }
    // Free memory
    memory_free64f(&phiWork);
    memory_free64f(&phiLoc);
    memory_free64f(&G);
    memory_free64f(&Dmat);
    memory_free64f(&Sigma);
    return 0;
}
//============================================================================//
/*!
 * @brief Performs the polarity search for a host of moment tensors.
 *
 * @param[in] nmt        Number of moment tensors.
 * @param[in] ldm        Leading dimension of moment tensor matrix.  This must
 *                       be at least 6.
 * @param[in] nPolarity  The number of polarities (observations).
 * @param[in] blockSize  Controls the number of forward problems GM to compute
 *                       simultaneously.
 * @param[in] mblock     The largest block size.  This serves as a leading
 *                       dimension.
 * @param[in] Mrows      Number of rows in Green's functions matrix.
 * @param[in] Kcols      Number of columns in Green's functions matrix.  This
 *                       should be 6.
 * @param[in] Dmat       The polarity data replicated to the max block size.
 *                       This has dimension [nPolarity x mblock] with leading
 *                       dimension mblock.
 * @param[in] G          Forward modeling matrix.  This has dimension
 *                       [nPolarity x LDG] with leading dimension LDG.
 * @param[in] Sigma      These are the data weights and has dimension
 *                       [nPolarity].
 * @param[in] mts        Matrix of moment tensors and has dimension [nmt x ldm].
 *                       with leading dimension ldm.
 *
 * @param[out] phi       The objective function (variance reduction) tabulated
 *                       for all the moment tensors and observations.  This has
 *                       dimension [nmt].
 *
 * @result 0 indicates success.
 *
 * @copyright ISTI distributed under the Apache 2 license.
 *
 */
static int performPolaritySearch64f(const int nmt, const int ldm,
                                    const int nPolarity,
                                    const int blockSize, const int mblock,
                                    const int Mrows, const int Kcols,
                                    const double *__restrict__ Dmat,
                                    const double *__restrict__ G,
                                    const double *__restrict__ Sigma,
                                    const double *__restrict__ mts,
                                    double *__restrict__ phi)
{
    double *U, *Usign, *res2, *sqrtSigmaWt, res, traceSigma;
    int i, ic, ierr, ipol, jmt, kmt, nmtBlocks, Ncols;
    const double one = 1.0;
    const double zero = 0.0;
    ierr = 0;
    nmtBlocks = (int) ((double) (nmt)/(double) (blockSize) + 1.0);
    if (nmtBlocks*blockSize < nmt)
    {
        fprintf(stderr, "%s: Internal error - all mts wont be visited\n",
                __func__);
        return -1;
    }
#ifdef __INTEL_COMPILER
    __assume_aligned(G, 64);
    __assume_aligned(Dmat, 64);
#endif
    // The VR in Chiang 2016 is:
    //   (1 - \frac{ \sum_i w_i (Pol_{i,obs} - Pol_{i,est})^2 }
    //             { \sum_i w_i (Pol_{i,obs}^2) } )
    // Because Pol_{obs} and Pol_{est} take values of + or -1  we can
    // reduce the demoniator so that the VR is
    //   (1 - \frac{ \sum_i w_i (Pol_{i,obs} - Pol_{i,est})^2 }
    //             { \sum_i w_i })
    // Furthermore, we can incorporate this term into the weighting
    // function s.t.  Sigma/trace(Sigma) to obtain
    //   (1 - \frac{ \sum_i \hat{Sigma} (Pol_{i,obs} - Pol_{i,est})^2)
    // Finally, because we want likelihoods we want to rescale the
    // residual squared from [0,4] to [0,1].  To do this we multiply
    // the weight by 0.5 s.t. the residual [-2,2] goes to [-1,1] which,
    // when squared, fits in the desired bounds of [0,1]; i.e. multiply by
    // (1/2)^2 = 0.25
    traceSigma = array_sum64f(nPolarity, Sigma, &ierr);
    if (ierr != 0){traceSigma = 1.0;}
    sqrtSigmaWt = memory_calloc64f(nPolarity);
#ifdef __INTEL_COMPILER
    __assume_aligned(sqrtSigmaWt, ISCL_MEMORY_ALIGN);
#endif
    for (i=0; i<nPolarity; i++){sqrtSigmaWt[i] = sqrt(0.25*Sigma[i]/traceSigma);}
    #pragma omp parallel \
     private (i, ic, ipol, jmt, kmt, Ncols, res, res2, U, Usign) \
     shared (G, Dmat, mts, nmtBlocks, phi, sqrtSigmaWt) \
     default (none) reduction(+:ierr)
    {
    U = memory_calloc64f(mblock*nPolarity);
    res2 = memory_calloc64f(mblock);
    Usign = memory_calloc64f(mblock*nPolarity);
#ifdef __INTEL_COMPILER
    __assume_aligned(res2, ISCL_MEMORY_ALIGN);
    __assume_aligned(Usign, ISCL_MEMORY_ALIGN);
#endif
    #pragma omp for
    for (kmt=0; kmt<nmtBlocks; kmt++)
    {
        jmt = kmt*blockSize;
        Ncols = MIN(blockSize, nmt - jmt); // Number of columns of M
        // Compute U = GM
        cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans,
                    Mrows, Ncols, Kcols, one, G, LDG,
                    &mts[ldm*jmt], ldm, zero, U, mblock);
        // Make the theoretical polarity +1 or -1 to match the data
        array_copysign64f_work(mblock*nPolarity, U, Usign);
        memset(res2, 0, (size_t) mblock*sizeof(double));
        // Compute the weighted residual
        for (ipol=0; ipol<nPolarity; ipol++)
        {
            #pragma omp simd aligned(Dmat, Usign, res2: ISCL_MEMORY_ALIGN)
            for (ic=0; ic<Ncols; ic++)
            {
                // Change residual range from [0,2] to [0,1] and weight
                res = sqrtSigmaWt[ipol]*( Dmat[ipol*mblock+ic]
                                        - Usign[ipol*mblock+ic]);
                res2[ic] = res2[ic] + res*res;
            }
        }
        // Compute the variance reduction and put it into objective function
        for (ic=0; ic<Ncols; ic++)
        {
            phi[jmt+ic] = 1.0 - res2[ic];
        }
    } // Loop on moment tensor blocks
    memory_free64f(&Usign);
    memory_free64f(&U);
    memory_free64f(&res2);
    } // end parallel section
    memory_free64f(&sqrtSigmaWt);
    return 0;
}

static int computePadding64f(const int n)
{
    size_t mod, pad;
    int ipad;
    // Set space and make G matrix
    pad = 0;
    mod = ((size_t) n*sizeof(double))%64;
    if (mod != 0)
    {
        pad = (64 - mod)/sizeof(double);
    }
    ipad = (int) pad;
    return ipad;
}


#include <config.h>
#include <stddef.h>
#include <stdlib.h>
#include <math.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_complex.h>

#include <gsl/gsl_fft_real.h>
#include <gsl/gsl_fft_real_float.h>

#define BASE_DOUBLE
#include "templates_on.h"
#include "real_pass.h"
#include "real_init.c"
#include "real_main.c"
#include "real_pass_2.c"
#include "real_pass_3.c"
#include "real_pass_4.c"
#include "real_pass_5.c"
#include "real_pass_n.c"
#include "real_radix2.c"
#include "real_unpack.c"
#include "templates_off.h"
#undef  BASE_DOUBLE

#define BASE_FLOAT
#include "templates_on.h"
#include "real_pass.h"
#include "real_init.c"
#include "real_main.c"
#include "real_pass_2.c"
#include "real_pass_3.c"
#include "real_pass_4.c"
#include "real_pass_5.c"
#include "real_pass_n.c"
#include "real_radix2.c"
#include "real_unpack.c"
#include "templates_off.h"
#undef  BASE_FLOAT


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <gsl/gsl_rng.h>
#include <assert.h>

#include "core_allvars.h"
#include "core_proto.h"

void init_galaxy(int p, int halonr)
{
  int j, step;

	assert(halonr == Halo[halonr].FirstHaloInFOFgroup);

  Gal[p].Type = 0;

  Gal[p].GalaxyNr = GalaxyCounter;
  GalaxyCounter++;

  Gal[p].HaloNr = halonr;
  Gal[p].MostBoundID = Halo[halonr].MostBoundID;
  Gal[p].SnapNum = Halo[halonr].SnapNum - 1;

  Gal[p].mergeType = 0;
  Gal[p].mergeIntoID = -1;
  Gal[p].mergeIntoSnapNum = -1;
  Gal[p].dT = -1.0;

  for(j = 0; j < 3; j++)
  {
    Gal[p].Pos[j] = Halo[halonr].Pos[j];
    Gal[p].Vel[j] = Halo[halonr].Vel[j];
  }

  Gal[p].Len = Halo[halonr].Len;
  Gal[p].Vmax = Halo[halonr].Vmax;
  Gal[p].Vvir = get_virial_velocity(halonr);
  Gal[p].Mvir = get_virial_mass(halonr);
  Gal[p].Rvir = get_virial_radius(halonr);

  Gal[p].deltaMvir = 0.0;

  Gal[p].ColdGas = 0.0;
  Gal[p].StellarMass = 0.0;
  Gal[p].BulgeMass = 0.0;
  Gal[p].HotGas = 0.0;
  Gal[p].EjectedMass = 0.0;
  Gal[p].BlackHoleMass = 0.0;
  Gal[p].ICS = 0.0;

  Gal[p].MetalsColdGas = 0.0;
  Gal[p].MetalsStellarMass = 0.0;
  Gal[p].MetalsBulgeMass = 0.0;
  Gal[p].MetalsHotGas = 0.0;
  Gal[p].MetalsEjectedMass = 0.0;
  Gal[p].MetalsICS = 0.0;

  for(step = 0; step < STEPS; step++)
  {
    Gal[p].SfrDisk[step] = 0.0;
    Gal[p].SfrBulge[step] = 0.0;
    Gal[p].SfrDiskColdGas[step] = 0.0;
    Gal[p].SfrDiskColdGasMetals[step] = 0.0;
    Gal[p].SfrBulgeColdGas[step] = 0.0;
    Gal[p].SfrBulgeColdGasMetals[step] = 0.0;
  }

  Gal[p].DiskScaleRadius = get_disk_radius(halonr, p);
  Gal[p].MergTime = 999.9;
  Gal[p].Cooling = 0.0;
  Gal[p].Heating = 0.0;
  Gal[p].r_heat = 0.0;
  Gal[p].QuasarModeBHaccretionMass = 0.0;
  Gal[p].TimeSinceMajorMerger = -1.0;
  Gal[p].TimeSinceMinorMerger = -1.0;
  Gal[p].OutflowRate = 0.0;
	Gal[p].TotalSatelliteBaryons = 0.0;

	// infall properties
  Gal[p].infallMvir = -1.0;
  Gal[p].infallVvir = -1.0;
  Gal[p].infallVmax = -1.0;

  //Raouf New things add to code
  Gal[p].Qjet = 0.0;
  Gal[p].Rcocoon = 0.0;
  Gal[p].Rshocked = 0.0;
  Gal[p].t_AGN_returne = 0.0;
  Gal[p].t_AGN_on = 0.0;
  Gal[p].Tshocked = 0.0;
  Gal[p].Mshocked = 0.0;
  for(j = 0; j < 7; j++)
  {
      Gal[p].RadioLuminosity[j] = 0.0;
  }
  for(j = 0; j < 600; j++)
  {
      Gal[p].RadioLuminosity_lifetime[j] = -999.0;
  }
  for(j = 0; j < 600; j++)
  {
      Gal[p].Rshocked_lifetime[j] = -999.0;
  }

  Gal[p].RadioAGNaccretionRate = 0.0;
  Gal[p].rho_zero_Makino = 0.0;
  Gal[p].rho_zero_Capelo = 0.0;
  Gal[p].rho_zero_iso = 0.0;
  Gal[p].b_gas = 0.0;
  Gal[p].Rs = 0.0;
  Gal[p].concentration = 0.0;
  Gal[p].Temp_Gas = 0.0;
  Gal[p].Lx_bol = 0.0;
  Gal[p].R_index = 0.0;
  Gal[p].Q_index = 0.0;
  Gal[p].R_cool = 0.0;
  Gal[p].fcool = 0.0;
  Gal[p].t_static = 0.0;
  Gal[p].t_AGN_off = 0.0;
  Gal[p].time_to_next_on = 0.0;
  Gal[p].delta = 0.0;
  Gal[p].t_cool_Makino = 0.0;
}

double get_disk_radius(int halonr, int p)
{
  double SpinMagnitude, SpinParameter;

	if(Gal[p].Vvir > 0.0 && Gal[p].Rvir > 0.0)
	{
		// See Mo, Shude & White (1998) eq12, and using a Bullock style lambda.
		SpinMagnitude = sqrt(Halo[halonr].Spin[0] * Halo[halonr].Spin[0] +
			Halo[halonr].Spin[1] * Halo[halonr].Spin[1] + Halo[halonr].Spin[2] * Halo[halonr].Spin[2]);

		// trim the extreme tail of the spin distribution for more a realistic r_s
		if(SpinMagnitude > 1.5)
			SpinMagnitude = 1.5;

		SpinParameter = SpinMagnitude / (1.414 * Gal[p].Vvir * Gal[p].Rvir);
		return (SpinParameter / 1.414) * Gal[p].Rvir;
	}
	else
		return 0.1 * Gal[p].Rvir;

}

double get_metallicity(double gas, double metals)
{
  double metallicity;

  if(gas > 0.0 && metals > 0.0)
  {
    metallicity = metals / gas;
    if(metallicity < 1.0)
      return metallicity;
    else
      return 1.0;
  }
  else
    return 0.0;

}

double dmax(double x, double y)
{
  if(x > y)
    return x;
  else
    return y;
}

double get_virial_mass(int halonr)
{
  if(halonr == Halo[halonr].FirstHaloInFOFgroup && Halo[halonr].Mvir >= 0.0)
    return Halo[halonr].Mvir;   /* take spherical overdensity mass estimate */
  else
    return Halo[halonr].Len * PartMass;
}

double get_virial_velocity(int halonr)
{
	double Rvir;

	Rvir = get_virial_radius(halonr);

  if(Rvir > 0.0)
		return sqrt(G * get_virial_mass(halonr) / Rvir);
	else
		return 0.0;
}

double get_virial_radius(int halonr)
{
  // return Halo[halonr].Rvir;  // Used for Bolshoi

  double zplus1, hubble_of_z_sq, rhocrit, fac;

  zplus1 = 1 + ZZ[Halo[halonr].SnapNum];
  hubble_of_z_sq =
    Hubble * Hubble *(Omega * zplus1 * zplus1 * zplus1 + (1 - Omega - OmegaLambda) * zplus1 * zplus1 +
    OmegaLambda);

  rhocrit = 3 * hubble_of_z_sq / (8 * M_PI * G);
  fac = 1 / (200 * 4 * M_PI / 3.0 * rhocrit);

  return cbrt(get_virial_mass(halonr) * fac);
}


/*
 * Copyright (c) 2011-2018 The University of Tennessee and The University
 *                         of Tennessee Research Foundation.  All rights
 *                         reserved.
 *
 * @precisions normal z -> s d c
 *
 */
#ifndef _DPLASMA_CORES_H_
#define _DPLASMA_CORES_H_

#include <math.h>
#include <cblas.h>
#include <lapacke.h>
#include "dplasma_complex.h"
#include "parsec/data_dist/matrix/matrix.h"
#include <core_blas.h>

#define PLASMA_BLKADDR(A, type, m, n)  (type *)plasma_getaddr(A, m, n)
#define PLASMA_BLKLDD(A, k) ( ( (k) + (A).i/(A).mb) < (A).lm1 ? (A).mb : (A).lm%(A).mb )

#endif /* _DPLASMA_CORES_H_ */


#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <gsl/gsl_math.h>

#include "allvars.h"
#include "proto.h"

#ifdef MHM

extern struct kindata_in
{
  MyFloat Pos[3];
  MyFloat Hsml;
  MyFloat Density;
  MyFloat Energy;
  int Index;
  int Task;
}
 *KinDataIn, *KinDataGet;

void kinetic_feedback_mhm(void)
{
  long long ntot, ntotleft;
  int ndone;
  int *noffset, *nbuffer, *nsend, *nsend_local, *numlist, *ndonelist;
  int i, j, n;
  int maxfill;
  int level, ngrp, sendTask, recvTask;
  int nexport;
  MPI_Status status;

  noffset = mymalloc(sizeof(int) * NTask);	/* offsets of bunches in common list */
  nbuffer = mymalloc(sizeof(int) * NTask);
  nsend_local = mymalloc(sizeof(int) * NTask);
  nsend = mymalloc(sizeof(int) * NTask * NTask);
  ndonelist = mymalloc(sizeof(int) * NTask);

  for(n = 0, NumSphUpdate = 0; n < N_gas; n++)
    {
      if(P[n].Type == 0)
	{
	  if(P[n].Ti_endstep == All.Ti_Current)
	    NumSphUpdate++;
	}
    }

  numlist = mymalloc(NTask * sizeof(int) * NTask);
  MPI_Allgather(&NumSphUpdate, 1, MPI_INT, numlist, 1, MPI_INT, MPI_COMM_WORLD);
  for(i = 0, ntot = 0; i < NTask; i++)
    ntot += numlist[i];
  myfree(numlist);

  i = 0;			/* beginn with this index */
  ntotleft = ntot;		/* particles left for all tasks together */

  while(ntotleft > 0)
    {
      for(j = 0; j < NTask; j++)
	nsend_local[j] = 0;

      /* do local particles and prepare export list */

      for(nexport = 0, ndone = 0; i < N_gas && nexport < All.BunchSizeKinetic - NTask; i++)
	if(P[i].Type == 0)
	  if(P[i].Ti_endstep == All.Ti_Current)
	    {
	      ndone++;

	      for(j = 0; j < NTask; j++)
		Exportflag[j] = 0;

	      kinetic_evaluate(i, 0);

	      for(j = 0; j < NTask; j++)
		{
		  if(Exportflag[j])
		    {
		      KinDataIn[nexport].Pos[0] = P[i].Pos[0];
		      KinDataIn[nexport].Pos[1] = P[i].Pos[1];
		      KinDataIn[nexport].Pos[2] = P[i].Pos[2];

		      KinDataIn[nexport].Density = SphP[i].Density;
		      KinDataIn[nexport].Energy = SphP[i].FeedbackEnergy;

		      KinDataIn[nexport].Hsml = PPP[i].Hsml;
		      KinDataIn[nexport].Index = i;
		      KinDataIn[nexport].Task = j;
		      nexport++;
		      nsend_local[j]++;
		    }
		}
	    }

      qsort(KinDataIn, nexport, sizeof(struct kindata_in), kin_compare_key);

      for(j = 1, noffset[0] = 0; j < NTask; j++)
	noffset[j] = noffset[j - 1] + nsend_local[j - 1];

      MPI_Allgather(nsend_local, NTask, MPI_INT, nsend, NTask, MPI_INT, MPI_COMM_WORLD);

      /* now do the particles that need to be exported */

      for(level = 1; level < (1 << PTask); level++)
	{
	  for(j = 0; j < NTask; j++)
	    nbuffer[j] = 0;
	  for(ngrp = level; ngrp < (1 << PTask); ngrp++)
	    {
	      maxfill = 0;
	      for(j = 0; j < NTask; j++)
		{
		  if((j ^ ngrp) < NTask)
		    if(maxfill < nbuffer[j] + nsend[(j ^ ngrp) * NTask + j])
		      maxfill = nbuffer[j] + nsend[(j ^ ngrp) * NTask + j];
		}
	      if(maxfill >= All.BunchSizeDensity)
		break;

	      sendTask = ThisTask;
	      recvTask = ThisTask ^ ngrp;

	      if(recvTask < NTask)
		{
		  if(nsend[ThisTask * NTask + recvTask] > 0 || nsend[recvTask * NTask + ThisTask] > 0)
		    {
		      /* get the particles */
		      MPI_Sendrecv(&KinDataIn[noffset[recvTask]],
				   nsend_local[recvTask] * sizeof(struct kindata_in), MPI_BYTE,
				   recvTask, TAG_DENS_A,
				   &KinDataGet[nbuffer[ThisTask]],
				   nsend[recvTask * NTask + ThisTask] * sizeof(struct kindata_in),
				   MPI_BYTE, recvTask, TAG_DENS_A, MPI_COMM_WORLD, &status);
		    }
		}

	      for(j = 0; j < NTask; j++)
		if((j ^ ngrp) < NTask)
		  nbuffer[j] += nsend[(j ^ ngrp) * NTask + j];
	    }

	  for(j = 0; j < nbuffer[ThisTask]; j++)
	    kinetic_evaluate(j, 1);

	  level = ngrp - 1;
	}

      MPI_Allgather(&ndone, 1, MPI_INT, ndonelist, 1, MPI_INT, MPI_COMM_WORLD);

      for(j = 0; j < NTask; j++)
	ntotleft -= ndonelist[j];
    }

  myfree(ndonelist);
  myfree(nsend);
  myfree(nsend_local);
  myfree(nbuffer);
  myfree(noffset);

}

/*! This function represents the core of the SPH density computation. The
 *  target particle may either be local, or reside in the communication
 *  buffer.
 */
void kinetic_evaluate(int target, int mode)
{
  int j, n;
  int startnode, numngb, numngb_inbox;
  double h, h2, hinv, hinv3;
  double rho, weight, wk, energy, delta_v;
  double dx, dy, dz, r, r2, u;
  double dvx, dvy, dvz;
  MyFloat *pos;

  if(mode == 0)
    {
      pos = P[target].Pos;
      h = PPP[target].Hsml;
      rho = SphP[target].Density;
      energy = SphP[target].FeedbackEnergy;
    }
  else
    {
      pos = KinDataGet[target].Pos;
      h = KinDataGet[target].Hsml;
      rho = KinDataGet[target].Density;
      energy = KinDataGet[target].Energy;
    }

  h2 = h * h;
  hinv = 1.0 / h;
#ifndef  TWODIMS
  hinv3 = hinv * hinv * hinv;
#else
  hinv3 = hinv * hinv / boxSize_Z;
#endif

  startnode = All.MaxPart;
  numngb = 0;
  do
    {
      numngb_inbox = ngb_treefind_variable(&pos[0], h, &startnode);

      for(n = 0; n < numngb_inbox; n++)
	{
	  j = Ngblist[n];

	  dx = pos[0] - P[j].Pos[0];
	  dy = pos[1] - P[j].Pos[1];
	  dz = pos[2] - P[j].Pos[2];

#ifdef PERIODIC			/*  now find the closest image in the given box size  */
	  if(dx > boxHalf_X)
	    dx -= boxSize_X;
	  if(dx < -boxHalf_X)
	    dx += boxSize_X;
	  if(dy > boxHalf_Y)
	    dy -= boxSize_Y;
	  if(dy < -boxHalf_Y)
	    dy += boxSize_Y;
	  if(dz > boxHalf_Z)
	    dz -= boxSize_Z;
	  if(dz < -boxHalf_Z)
	    dz += boxSize_Z;
#endif
	  r2 = dx * dx + dy * dy + dz * dz;

	  if(r2 < h2)
	    {
	      r = sqrt(r2);

	      if(r > 0)
		{

		  u = r * hinv;

		  if(u < 0.5)
		    wk = hinv3 * (KERNEL_COEFF_1 + KERNEL_COEFF_2 * (u - 1) * u * u);
		  else
		    wk = hinv3 * KERNEL_COEFF_5 * (1.0 - u) * (1.0 - u) * (1.0 - u);

		  weight = wk * P[j].Mass / rho;

		  delta_v = sqrt(2 * energy * weight / P[j].Mass);

		  dvx = delta_v * (-dx) / r;
		  dvy = delta_v * (-dy) / r;
		  dvz = delta_v * (-dz) / r;

		  P[j].Vel[0] += dvx;
		  P[j].Vel[1] += dvy;
		  P[j].Vel[2] += dvz;

		  SphP[j].VelPred[0] += dvx;
		  SphP[j].VelPred[1] += dvy;
		  SphP[j].VelPred[2] += dvz;
		}
	    }
	}
    }
  while(startnode >= 0);
}

/*! This routine is a comparison kernel used in a sort routine to group
 *  particles that are exported to the same processor.
 */
int kin_compare_key(const void *a, const void *b)
{
  if(((struct kindata_in *) a)->Task < (((struct kindata_in *) b)->Task))
    return -1;

  if(((struct kindata_in *) a)->Task > (((struct kindata_in *) b)->Task))
    return +1;

  return 0;
}

#endif


#ifndef ALLVARS_H
#define ALLVARS_H

#include <stdio.h>
#include <stdbool.h>
#include <gsl/gsl_rng.h>
#include "core_simulation.h"
#include <stdint.h>

#define ABORT(sigterm)                                                  \
do {                                                                \
  printf("Error in file: %s\tfunc: %s\tline: %i\n", __FILE__, __FUNCTION__, __LINE__); \
  myexit(sigterm);                                                \
} while(0)

#define  STEPS 10       // Number of integration intervals between two snapshots
#define  MAXGALFAC 1
#define  ALLOCPARAMETER 10.0
#define  MAX_NODE_NAME_LEN 50
#define  ABSOLUTEMAXSNAPS 1000
#define STELLAR_TRACKING_TIME 100 // This is the time (Myr) we wish to track the stellar history for.

#define  MAXLEN      1024
#define  MAX_STRING_LEN      1024
#define  GRAVITY     6.672e-8
#define  SOLAR_MASS  1.989e33
#define  SOLAR_LUM   3.826e33
#define  RAD_CONST   7.565e-15
#define  AVOGADRO    6.0222e23
#define  BOLTZMANN   1.3806e-16
#define  GAS_CONST   8.31425e7
#define  C           2.9979e10
#define  PLANCK      6.6262e-27
#define  CM_PER_MPC  3.085678e24
#define  PROTONMASS  1.6726e-24
#define  HUBBLE      3.2407789e-18   /* in h/sec */

#define  SEC_PER_MEGAYEAR   3.155e13
#define  SEC_PER_YEAR       3.155e7

#define	 CUBE(x) (x*x*x)

#ifdef NDEBUG
#define XASSERT(EXP, ...)                                do{} while(0)
#else
#define XASSERT(EXP, ...)                                              \
    do { if (!(EXP)) {                                                  \
            printf("Error in file: %s\tfunc: %s\tline: %d with expression `"#EXP"'\n", __FILE__, __FUNCTION__, __LINE__); \
            printf(__VA_ARGS__);                                        \
            fflush(stdout);                                             \
            ABORT(EXIT_FAILURE);                                         \
        } \
    } while (0)
#endif

#ifdef NDEBUG
#define XPRINT(EXP, ...)                                do{} while(0)
#else
#define XPRINT(EXP, ...)                                              \
    do { if (!(EXP)) {                                                  \
            printf("Warning in file: %s\tfunc: %s\tline: %d with expression `"#EXP"'\n", __FILE__, __FUNCTION__, __LINE__); \
            printf(__VA_ARGS__);                                        \
            fflush(stdout);                                             \
        } \
    } while (0)
#endif

// This structure contains the properties used within the code
struct GALAXY
{

  int32_t FileNr;

  int   SnapNum;
  int   Type;

  int   GalaxyNr;
  int   CentralGal;
  int   HaloNr;
  int   TreeNr;
  long long  MostBoundID;

  int   mergeType;  //0=none; 1=minor merger; 2=major merger; 3=disk instability; 4=disrupt to ICS
  int   mergeIntoID;
  int   mergeIntoSnapNum;
  float   dT;
  float Age;
  // (sub)halo properties
  float Pos[3];
  float Vel[3];
  int   Len;
  float Mvir;
  float deltaMvir;
  float CentralMvir;
  float Rvir;
  float Vvir;
  float Vmax;

  // baryonic reservoirs
  float ColdGas;
  float StellarMass;
  float BulgeMass;
  float HotGas;
  float EjectedMass;
  float EjectedMassSN;
  float EjectedMassQSO;
  float BlackHoleMass;
  float ICS;

  // metals
  float MetalsColdGas;
  float MetalsStellarMass;
  float MetalsBulgeMass;
  float MetalsHotGas;
  float MetalsEjectedMass;
  float MetalsICS;

  // to calculate magnitudes
  float SfrDisk[STEPS];
  float SfrBulge[STEPS];
  float SfrDiskColdGas[STEPS];
  float SfrDiskColdGasMetals[STEPS];
  float SfrBulgeColdGas[STEPS];
  float SfrBulgeColdGasMetals[STEPS];

  // misc
  float DiskScaleRadius;
  float MergTime;
  double Cooling;
  double Heating;
  float r_heat;
  float QuasarModeBHaccretionMass;
  float TimeOfLastMajorMerger;
  float TimeOfLastMinorMerger;
  float OutflowRate;
	float TotalSatelliteBaryons;

  // infall properties
  float infallMvir;
  float infallVvir;
  float infallVmax;

  int32_t *GridType;
  int32_t *GridFoFHaloNr;
  int32_t *GridHistory;
  float   *GridColdGas;
  float   *GridHotGas;
  float   *GridEjectedMass;
  float   *GridDustColdGas;
  float   *GridDustHotGas;
  float   *GridDustEjectedMass;
  float   *GridStellarMass;
  float   *GridBHMass;
  float   *GridSFR;
  float   *GridZ;
  float   *GridFoFMass;
  float   *GridHaloMass;
  float   *EjectedFraction;
  float   *EjectedFractionSN;
  float   *EjectedFractionQSO;
  int32_t *LenHistory;
  float   *GridOutflowRate;
  float   *GridInfallRate;
  int32_t *QuasarActivity;
  int32_t *QuasarSubstep;
  float   *DynamicalTime;
  int32_t *LenMergerGal;
  float   *GridReionMod;
  float   *GridNgamma_HI;
  float   *Gridfesc;
  float   *ColdCrit;

  int IsMerged;
  double GrandSum;
  int IsMalloced;

  // Delayed SN Properties
  float reheated_mass;
  float ejected_mass;
  float mass_stars_recycled;
  float mass_metals_new;
  float NSN; // Number of supernova within a time step.  Can be fractional.
  double *SN_Stars;

  double StellarAge_Numerator;
  double StellarAge_Denominator;
  double Total_SN_SF_Time;
  double Total_SN_Stars;

  // Dust Properties
  float DustColdGas;
  float DustHotGas;
  float DustEjectedMass;

  // Quasar Activity Tracking

  int32_t QuasarActivityToggle;
  float TargetQuasarTime;
  float QuasarBoostActiveTime;
  float QuasarFractionalPhotons;

  // Stellar Age Tracking To Determine Nion.
  float *Stellar_Stars; // Keep this separate from the Stars of the delayed SN.
                        // Its possible tat we want to do tracking of stellar ages to determine Nion but not do delayed SN (or vice versa).
  double Total_Stellar_SF_Time;
  double Total_Stellar_Stars;

  // UV Magnitudes
  float *MUV;

}
*Gal, *HaloGal, *MergedGal;

// auxiliary halo data
struct halo_aux_data
{
  int DoneFlag;
  int HaloFlag;
  int NGalaxies;
  int FirstGalaxy;
}
*HaloAux;

// For the fully self consistent model we read in a list of reionization modifiers for each halo.
// These lists only contain those halos with a reionization modifier less than 0.9999.
// We will read if a list for each of the previous snapshots below the one we are currently on.
// E.g., if we are applying reionization feedback for snapshot 40, we will read in lists for snapshot 0 to 39.
// There is one list for each tree file.
struct REIONMOD_STRUCT
{

  int32_t NumLists; // Will be reionization snapshot number.
  struct REIONMOD_LIST
  {
    int32_t NHalos_Found; // Number of halos that have been matched.
    int32_t NHalos_Ionized; // Number of halos in the list.
    int64_t *HaloID; // This is a unique ID for each tree file.  The most significant 32 bits (left-most) is the tree number, and the least significant (right-most) bits is the halonr within the tree.
    float *ReionMod; // Reionization Modifer for each halo.
  }*ReionMod_List;
}*ReionList;

// Grid for doing photoionization feedback.
// This will contain multiple grids, one for each redshift.
struct GRID_STRUCT
{

  int32_t GridSize;
  uint64_t NumCellsTotal;
  int32_t NumGrids;

  double *ReionRedshift; // This is the redshift the grid cell was ionized at.
  struct PHOTO_GRID
  {
    double *PhotoRate; // Photoionization Rate (s^-1).
    int32_t valid_grid; // This will control whether data has been read in for this snapshot. Useful for the early snapshots where we don't have any ionization.
  }*PhotoGrid;
}*Grid;

struct SELFCON_GRID_STRUCT
{

  int32_t GridSize;
  int32_t NumCellsTotal;
  float Nion_HI_Total;

  float *Nion_HI;
  int32_t *GalCount;

}*SelfConGrid;

extern int    FirstFile;    // first and last file for processing
extern int    LastFile;

extern int count_onehalo;

extern int    Ntrees;      // number of trees in current file
extern int    NumGals;     // Total number of galaxies stored for current tree
extern int    MaxGals;     // Maximum number of galaxies allowed for current tree
extern int    FoF_MaxGals;

extern int    GalaxyCounter;     // unique galaxy ID for main progenitor line in tree

extern int    LastSnapShotNr;

extern char   OutputDir[512];
extern char   GalaxyOutputDir[512];
extern char   GridOutputDir[512];
extern char   RunPrefix[512];
extern char   TreeName[512];
extern char   TreeExtension[512];
extern char   SimulationDir[512];
extern char   IonizationDir[512];
extern char   FileWithSnapList[512];
extern char   PhotoionDir[512];
extern char   PhotoionName[512];
extern char   ReionRedshiftName[512];
extern int    ReionSnap;

extern int    TotHalos;
extern int    TotGalaxies;
extern int    *TreeNgals;

extern int    *FirstHaloInSnap;

extern int    *TreeNHalos;
extern int    *TreeFirstHalo;

#ifdef MPI
extern int ThisTask, NTask, nodeNameLen;
extern char *ThisNode;
#else
extern int ThisTask, NTask;
#endif

extern int IMF;

extern double Omega;
extern double OmegaLambda;
extern double PartMass;
extern double BoxSize;
extern int GridSize;
extern int self_consistent;
extern double Hubble_h;
extern double EnergySNcode, EnergySN;

// recipe flags
extern int    QuasarRecipeOn;
extern int    ReionizationOn;
extern int    SupernovaRecipeOn;
extern int    DiskInstabilityOn;
extern int    AGNrecipeOn;
extern int    SFprescription;

// recipe parameters
extern double RecycleFraction;
extern double Yield;
extern double ReIncorporationFactor;
extern double ThreshMajorMerger;
extern double BaryonFrac;
extern double SfrEfficiency;
extern double FeedbackReheatingEpsilon;
extern double FeedbackEjectionEfficiency;
extern double RadioModeEfficiency;
extern double QuasarModeEfficiency;
extern double BlackHoleGrowthRate;
extern double Reionization_z0;
extern double Reionization_zr;
extern double ThresholdSatDisruption;

// Parameters for the gridding with self_consistent
extern int fescPrescription;

extern double alpha;
extern double beta;
extern double delta;

extern double fescMH_alpha;
extern double fescMH_beta;

extern double quasar_baseline;
extern double quasar_boosted;
extern double N_dyntime;

extern double MH_low;
extern double fesc_low;
extern double MH_high;
extern double fesc_high;

extern double UnitLength_in_cm,
  UnitTime_in_s,
  UnitVelocity_in_cm_per_s,
  UnitMass_in_g,
  RhoCrit,
  UnitPressure_in_cgs,
  UnitDensity_in_cgs,
  UnitCoolingRate_in_cgs,
  UnitEnergy_in_cgs,
  UnitTime_in_Megayears,
  sage_G,
  sage_Hubble,
  a0, ar;

extern int    ListOutputSnaps[ABSOLUTEMAXSNAPS];

extern double ZZ[ABSOLUTEMAXSNAPS];
extern double AA[ABSOLUTEMAXSNAPS];
extern double Age[ABSOLUTEMAXSNAPS];

extern int    MAXSNAPS;
extern int    Snaplistlen;

extern gsl_rng *random_generator;

extern int TreeID;
extern int FileNum;
extern int neg_cell;
extern int pos_cell;

extern double ejectedmass_total;
extern double metalsejectedmass_total;

extern int MergedNr;
extern int TotMerged;
extern int *TreeNMergedgals;
extern int MaxMergedGals;     // Maximum number of galaxies allowed for current tree

extern int zeromass_count;
extern int suppression_count;
extern int previous_tree;
extern int lowmass_halo;

extern double smallest_mass;

extern double IMF_norm;
extern double IMF_slope;
extern double Eta_SNII;
extern double m_SNII;
extern int IRA;
extern int TimeResolutionSN;
extern int SN_Array_Len;
extern int Time_SFH;

extern double alpha_energy;
extern double V_energy;
extern double beta_energy;
extern double alpha_mass;
extern double V_mass;
extern double beta_mass;
extern double epsilon_mass_max;

extern int RescaleSN;
extern int mergedgal_mallocs;
extern int gal_mallocs;
extern int ismerged_mallocs;

extern int mergedgal_frees;
extern int gal_frees;
extern int ismerged_frees;

extern int *gal_to_free;

extern int outside_box;
extern int inside_box;

extern int count_Mvir;
extern int count_Len;

extern long count_gal;

extern int32_t LowSnap;
extern int32_t HighSnap;

extern float Ngamma_HI_Total;

// Delayed Supernovae Lookup Tables

extern float *IMF_massgrid_eta;
extern float *IMF_massgrid_m;
extern float m_IMFbins_low;
extern float m_IMFbins_high;
extern float m_IMFbins_delta;
extern float N_massbins;

extern float *coreburning_times;
extern float coreburning_tbins_low;
extern float coreburning_tbins_high;
extern float coreburning_tbins_delta;
extern float N_tbins;

// Ionizing Photon Lookup Tables

extern int32_t PhotonPrescription;
extern float *stars_tbins;
extern float *stars_Ngamma;
extern double TimeResolutionStellar;
extern int32_t StellarTracking_Len;
extern int32_t HaloPartCut;

// UV Magnitude

extern int32_t calcUVmag;
extern float *stars_LUV;

#endif  // #ifndef ALLVARS_H


/* Copyright (C) 2021 Atsushi Togo */
/* All rights reserved. */

/* This file is part of phonopy. */

/* Redistribution and use in source and binary forms, with or without */
/* modification, are permitted provided that the following conditions */
/* are met: */

/* * Redistributions of source code must retain the above copyright */
/*   notice, this list of conditions and the following disclaimer. */

/* * Redistributions in binary form must reproduce the above copyright */
/*   notice, this list of conditions and the following disclaimer in */
/*   the documentation and/or other materials provided with the */
/*   distribution. */

/* * Neither the name of the phonopy project nor the names of its */
/*   contributors may be used to endorse or promote products derived */
/*   from this software without specific prior written permission. */

/* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS */
/* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT */
/* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS */
/* FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE */
/* COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, */
/* INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, */
/* BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; */
/* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER */
/* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT */
/* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN */
/* ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE */
/* POSSIBILITY OF SUCH DAMAGE. */

#ifndef __phono3py_H__
#define __phono3py_H__

#if defined(MKL_LAPACKE) || defined(SCIPY_MKL_H)
#include <mkl.h>
#else
#include <lapacke.h>
#endif
#include "phonoc_array.h"

long ph3py_get_interaction(
    Darray *fc3_normal_squared, const char *g_zero, const Darray *frequencies,
    const lapack_complex_double *eigenvectors, const long (*triplets)[3],
    const long num_triplets, const long (*bz_grid_addresses)[3],
    const long D_diag[3], const long Q[3][3], const double *fc3,
    const long is_compact_fc3, const double (*svecs)[3],
    const long multi_dims[2], const long (*multi)[2], const double *masses,
    const long *p2s_map, const long *s2p_map, const long *band_indices,
    const long symmetrize_fc3_q, const double cutoff_frequency);
long ph3py_get_pp_collision(
    double *imag_self_energy,
    const long relative_grid_address[24][4][3], /* thm */
    const double *frequencies, const lapack_complex_double *eigenvectors,
    const long (*triplets)[3], const long num_triplets,
    const long *triplet_weights, const long (*bz_grid_addresses)[3], /* thm */
    const long *bz_map,                                              /* thm */
    const long bz_grid_type, const long D_diag[3], const long Q[3][3],
    const double *fc3, const long is_compact_fc3, const double (*svecs)[3],
    const long multi_dims[2], const long (*multi)[2], const double *masses,
    const long *p2s_map, const long *s2p_map, const Larray *band_indices,
    const Darray *temperatures, const long is_NU, const long symmetrize_fc3_q,
    const double cutoff_frequency);
long ph3py_get_pp_collision_with_sigma(
    double *imag_self_energy, const double sigma, const double sigma_cutoff,
    const double *frequencies, const lapack_complex_double *eigenvectors,
    const long (*triplets)[3], const long num_triplets,
    const long *triplet_weights, const long (*bz_grid_addresses)[3],
    const long D_diag[3], const long Q[3][3], const double *fc3,
    const long is_compact_fc3, const double (*svecs)[3],
    const long multi_dims[2], const long (*multi)[2], const double *masses,
    const long *p2s_map, const long *s2p_map, const Larray *band_indices,
    const Darray *temperatures, const long is_NU, const long symmetrize_fc3_q,
    const double cutoff_frequency);
void ph3py_get_imag_self_energy_at_bands_with_g(
    double *imag_self_energy, const Darray *fc3_normal_squared,
    const double *frequencies, const long (*triplets)[3],
    const long *triplet_weights, const double *g, const char *g_zero,
    const double temperature, const double cutoff_frequency,
    const long num_frequency_points, const long frequency_point_index);
void ph3py_get_detailed_imag_self_energy_at_bands_with_g(
    double *detailed_imag_self_energy, double *imag_self_energy_N,
    double *imag_self_energy_U, const Darray *fc3_normal_squared,
    const double *frequencies, const long (*triplets)[3],
    const long *triplet_weights, const long (*bz_grid_addresses)[3],
    const double *g, const char *g_zero, const double temperature,
    const double cutoff_frequency);
void ph3py_get_real_self_energy_at_bands(
    double *real_self_energy, const Darray *fc3_normal_squared,
    const long *band_indices, const double *frequencies,
    const long (*triplets)[3], const long *triplet_weights,
    const double epsilon, const double temperature,
    const double unit_conversion_factor, const double cutoff_frequency);
void ph3py_get_real_self_energy_at_frequency_point(
    double *real_self_energy, const double frequency_point,
    const Darray *fc3_normal_squared, const long *band_indices,
    const double *frequencies, const long (*triplets)[3],
    const long *triplet_weights, const double epsilon, const double temperature,
    const double unit_conversion_factor, const double cutoff_frequency);
void ph3py_get_collision_matrix(
    double *collision_matrix, const Darray *fc3_normal_squared,
    const double *frequencies, const long (*triplets)[3],
    const long *triplets_map, const long *map_q,
    const long *rotated_grid_points, const double *rotations_cartesian,
    const double *g, const long num_ir_gp, const long num_gp,
    const long num_rot, const double temperature,
    const double unit_conversion_factor, const double cutoff_frequency);
void ph3py_get_reducible_collision_matrix(
    double *collision_matrix, const Darray *fc3_normal_squared,
    const double *frequencies, const long (*triplets)[3],
    const long *triplets_map, const long *map_q, const double *g,
    const long num_gp, const double temperature,
    const double unit_conversion_factor, const double cutoff_frequency);
void ph3py_get_isotope_scattering_strength(
    double *gamma, const long grid_point, const double *mass_variances,
    const double *frequencies, const lapack_complex_double *eigenvectors,
    const long num_grid_points, const long *band_indices, const long num_band,
    const long num_band0, const double sigma, const double cutoff_frequency);
void ph3py_get_thm_isotope_scattering_strength(
    double *gamma, const long grid_point, const long *ir_grid_points,
    const long *weights, const double *mass_variances,
    const double *frequencies, const lapack_complex_double *eigenvectors,
    const long num_ir_grid_points, const long *band_indices,
    const long num_band, const long num_band0,
    const double *integration_weights, const double cutoff_frequency);
void ph3py_distribute_fc3(double *fc3, const long target, const long source,
                          const long *atom_mapping, const long num_atom,
                          const double *rot_cart);
void ph3py_rotate_delta_fc2(double (*fc3)[3][3][3],
                            const double (*delta_fc2s)[3][3],
                            const double *inv_U,
                            const double (*site_sym_cart)[3][3],
                            const long *rot_map_syms, const long num_atom,
                            const long num_site_sym, const long num_disp);
void ph3py_get_permutation_symmetry_fc3(double *fc3, const long num_atom);
void ph3py_get_permutation_symmetry_compact_fc3(
    double *fc3, const long p2s[], const long s2pp[], const long nsym_list[],
    const long perms[], const long n_satom, const long n_patom);
void ph3py_transpose_compact_fc3(double *fc3, const long p2s[],
                                 const long s2pp[], const long nsym_list[],
                                 const long perms[], const long n_satom,
                                 const long n_patom, const long t_type);
long ph3py_get_triplets_reciprocal_mesh_at_q(
    long *map_triplets, long *map_q, const long grid_point, const long mesh[3],
    const long is_time_reversal, const long num_rot,
    const long (*rec_rotations)[3][3], const long swappable);
long ph3py_get_BZ_triplets_at_q(long (*triplets)[3], const long grid_point,
                                const long (*bz_grid_addresses)[3],
                                const long *bz_map, const long *map_triplets,
                                const long num_map_triplets,
                                const long D_diag[3], const long Q[3][3],
                                const long bz_grid_type);
long ph3py_get_integration_weight(
    double *iw, char *iw_zero, const double *frequency_points,
    const long num_band0, const long relative_grid_address[24][4][3],
    const long mesh[3], const long (*triplets)[3], const long num_triplets,
    const long (*bz_grid_addresses)[3], const long *bz_map,
    const long bz_grid_type, const double *frequencies1, const long num_band1,
    const double *frequencies2, const long num_band2, const long tp_type,
    const long openmp_per_triplets, const long openmp_per_bands);
void ph3py_get_integration_weight_with_sigma(
    double *iw, char *iw_zero, const double sigma, const double sigma_cutoff,
    const double *frequency_points, const long num_band0,
    const long (*triplets)[3], const long num_triplets,
    const double *frequencies, const long num_band, const long tp_type);
long ph3py_get_grid_index_from_address(const long address[3],
                                       const long mesh[3]);
void ph3py_get_gr_grid_addresses(long gr_grid_addresses[][3],
                                 const long D_diag[3]);
long ph3py_get_reciprocal_rotations(long rec_rotations[48][3][3],
                                    const long (*rotations)[3][3],
                                    const long num_rot,
                                    const long is_time_reversal);
long ph3py_transform_rotations(long (*transformed_rots)[3][3],
                               const long (*rotations)[3][3],
                               const long num_rot, const long D_diag[3],
                               const long Q[3][3]);
long ph3py_get_snf3x3(long D_diag[3], long P[3][3], long Q[3][3],
                      const long A[3][3]);
long ph3py_transform_rotations(long (*transformed_rots)[3][3],
                               const long (*rotations)[3][3],
                               const long num_rot, const long D_diag[3],
                               const long Q[3][3]);
long ph3py_get_ir_grid_map(long *ir_grid_map, const long D_diag[3],
                           const long PS[3], const long (*grg_rotations)[3][3],
                           const long num_rot);
long ph3py_get_bz_grid_addresses(long (*bz_grid_addresses)[3], long *bz_map,
                                 long *bzg2grg, const long D_diag[3],
                                 const long Q[3][3], const long PS[3],
                                 const double rec_lattice[3][3],
                                 const long type);
long ph3py_rotate_bz_grid_index(const long bz_grid_index,
                                const long rotation[3][3],
                                const long (*bz_grid_addresses)[3],
                                const long *bz_map, const long D_diag[3],
                                const long PS[3], const long bz_grid_type);
void ph3py_symmetrize_collision_matrix(double *collision_matrix,
                                       const long num_column,
                                       const long num_temp,
                                       const long num_sigma);
void ph3py_expand_collision_matrix(double *collision_matrix,
                                   const long *rot_grid_points,
                                   const long *ir_grid_points,
                                   const long num_ir_gp,
                                   const long num_grid_points,
                                   const long num_rot, const long num_sigma,
                                   const long num_temp, const long num_band);
long ph3py_get_neighboring_gird_points(
    long *relative_grid_points, const long *grid_points,
    const long (*relative_grid_address)[3], const long mesh[3],
    const long (*bz_grid_addresses)[3], const long *bz_map,
    const long bz_grid_type, const long num_grid_points,
    const long num_relative_grid_address);
long ph3py_get_thm_integration_weights_at_grid_points(
    double *iw, const double *frequency_points, const long num_band0,
    const long num_band, const long num_gp,
    const long (*relative_grid_address)[4][3], const long D_diag[3],
    const long *grid_points, const long (*bz_grid_addresses)[3],
    const long *bz_map, const long bz_grid_type, const double *frequencies,
    const long *gp2irgp_map, const char function);

#endif