// apply.c  "QUANTIFICATION" NESTED SAMPLING APPLICATION
// (GNU General Public License software, (C) Sivia and Skilling 2006)
// Problem:     5 quantities q, which can be ON (+ve) or OFF (zero),
//              as measured by linear data with Gaussian noise.
// Inputs:
//  Prior(q)    is prob(ON) = PrON  with  prob(q|ON) = 1/(1+q)^2
//  Likelihood  is exp(-chisquared/2)
//                 chisquared = SUM residual^2
//                 residual   = mock data - actual data
//                 mock data  = [Expt response].[quantities q]
// Outputs:
//  Evidence    is Z = INTEGRAL Likelihood(q) * Prior(q) dq
//  Posterior   is P(q) = Likelihood(q) * Prior(q) / Z
//  Information is H = INTEGRAL P(q) log(P(q)/Prior(q)) dq
/*__________________________________________________________________*/
#define n    100   // # Objects
#define MAX 1000   // # iterates
#define M      5   // # switches
#define Mplus  8   // power-of-2 >= M
#define N      3   // # data
#include "tree.c"  // binary tree procedures void PutRate, int GetRate
const double PrON[M] ={0.1, 0.2, 0.3, 0.2, 0.1}; // prior prob(ON)
const double Data[N] ={3,6,9};                   // actual data
const double Expt[N][M] ={{0,1,2,3,4},           // 1st data response
                          {0,0,3,2,1},           // 2nd data response
                          {3,2,1,2,3}};          // 3rd data response
/*__________________________________________________________________*/
typedef struct
{
    double  q[M];          // quantities, 0=OFF else prob(q)=1/(1+q)^2
    double  Tree[2*Mplus]; // binary rate tree
    double  logL;          // logLikelihood = ln Prob(data | q)
    double  logWt;         // log(Weight), with SUM(Wt) = Evidence Z
} Object;
/*__________________________________________________________________*/
double logLhood(        // logLikelihood function
double* q)              // quantities
{
    double resid;       // residual = (mock - actual) data
    int    j, k;        // component and data counters
    double C = 0.0;     // chisquared
    for( k = 0; k < N; k++ )
    {
        resid = -Data[k];
        for( j = 0; j < M; j++ )
            resid += Expt[k][j] * q[j];
        C += resid * resid;
    }
    return  -0.5 * C + sqrt(DBL_EPSILON) * UNIFORM;
}                // jitter eliminates ties between likelihood values
/*__________________________________________________________________*/
void Prior(            // Set Object according to prior
Object* Obj)           // Object being set
{
    double u;
    int    i;
// Initialise empty Tree of transition rates
    Obj->Tree[0] = Mplus;
    for( i = 1; i < 2*Mplus; i++ )
        Obj->Tree[i] = 0.0;
// Initialise object
    for( i = 0; i < M; i++ )
    {
        u = (UNIFORM < PrON[i]) ? UNIFORM : 0.0;
        Obj->q[i] = u / (1.0 - u);
        PutRate(Obj->Tree, i, (u > 0.0) ? 1.-PrON[i] : PrON[i]);
    }
    Obj->logL = logLhood(Obj->q);
}
/*__________________________________________________________________*/
double TryQ(            // revised trial q[i]
double* q,              // quantities
int     i,              // id of quantity q[i] to be varied
double  logLstar)       // constraint
{
    double resid;       // residual = (mock - actual) data
    int    j, k;        // component and data counters
    double A, B, C, D;  // quadratic coeffs and discriminant
    double u;           // controlling variable for q[i]
    double min, max;    // range, of q[i] then of u
// "L > Lstar" is quadratic interval "A*x*x + 2*B*x + C <= 0.0"
    A = B = 0.0;    C = 2.0 * logLstar;  // minus chisquared on entry
    for( k = 0; k < N; k++ )
    {
        resid = -Data[k];
        for( j = 0; j < M; j++ )
            resid += Expt[k][j] * q[j];
        A += Expt[k][i] * Expt[k][i];
        B += Expt[k][i] * resid;
        C += resid * resid;
    }
// Find controlling interval
    if( A > 0.0 )               // q[i] does affect data
    {
    // Solve quadratic for (qmin,qmax) relative to q[i]
        D = B * B - A * C;      // discriminant
        if( D > 0.0 )           // distinct real roots
        {
            if( B > 0.0 )
            { min = -B - sqrt(D);  max = C / min;  min /= A; }
            else
            { max = -B + sqrt(D);  min = C / max;  max /= A; }
    // Reset (qmin,qmax) relative to origin q=0
            min += q[i];   max += q[i];
    // Restrict (qmin,qmax) to non-negative values
            if( max <= 0.0 )  max = min = 0.0;
            else if( min < 0.0 )    min = 0.0;
    // Transform to controlling interval (umin,umax)
            min /= 1.0 + min;   max /= 1.0 + max;
        }
        else
            min = max = 0.0;    // no real roots, so null interval
    }
    else
    {                           // q[i] unmeasured, so
        min = 0.0;  max = 1.0;  // all u are in range
    }
// Accept/Reject
    if( q[i] == 0.0 )                  // entry state OFF
        u = (UNIFORM < max - min)      // accept ON, sample u
          ? min + (max - min) * UNIFORM : 0.0;
    else                               // entry state ON
        u = (min > 0.0)                // reject OFF, re-sample u
          ? min + (max - min) * UNIFORM : 0.0;
    return  u / (1.0 - u);             // trial quantity
}
/*__________________________________________________________________*/
void Explore(          // Evolve object within likelihood constraint
Object* Obj,           // Object being evolved
double  logLstar)      // Likelihood constraint L > Lstar
{
    double Interval = 30.0; // pre-judged time to equilibrate
    double t = 0.0;         // evolution time, initialized 0
    double qold;            // entry value
    double logLtry;         // trial loglikelihood
    int    i;               // quantity being changed
    while( (t += -log(UNIFORM) / Obj->Tree[1]) < Interval )
    {
        i = GetRate(Obj->Tree);
        qold = Obj->q[i];   // enable recovery of entry state
        Obj->q[i] = TryQ(Obj->q, i, logLstar);  // trial state
        logLtry = logLhood(Obj->q);
        if( logLtry > logLstar )
        {
            Obj->logL = logLtry;                // accept
            PutRate(Obj->Tree, i,
                     (Obj->q[i]>0.) ? 1.-PrON[i] : PrON[i]);
        }
        else
            Obj->q[i] = qold;                   // reject
    }
}
/*__________________________________________________________________*/
void Results(     // Posterior properties, here statistics of s
Object* Samples,  // Objects defining posterior
int     nest,     // # Samples
double  logZ)     // Evidence (= total weight = SUM[Samples] Weight)
{
    double post[M] = {0,0,0,0,0}; // posterior prob(ON)
    double mean[M] = {0,0,0,0,0}; // quantity mean (when ON)
    double var[M]  = {0,0,0,0,0}; // variance (when ON)
    double w;                     // Proportional weight
    int    i, k;                  // Sequence and sample counters
    for( i = 0; i < nest; i++ )
    {
        w = exp(Samples[i].logWt - logZ);
        for( k = 0; k < M; k++ )
            if( Samples[i].q[k] > 0.0 )
            {
                post[k] += w;
                mean[k] += w * Samples[i].q[k];
                var[k]  += w * Samples[i].q[k]*Samples[i].q[k];
            }
    }
    for( k = 0; k < M; k++ )
    {
        mean[k] /= post[k];
        var[k] = var[k] / post[k] - mean[k] * mean[k];
        printf("%d: Posterior(ON) =%3.0f%%,", k, 100.*post[k]);
        printf(" q =%6.2f +-%6.2f\n", mean[k], sqrt(var[k]));
    }
}