/*

      Simple EM demo for fitting a mixture of two or more classes to data in two Boolean variables.

      Use: Use: emdemo seed #iterations #classes #vars #...00 #...01 #...10 ...
      where
        seed = integer used to initialize random number generator
        #iterations = how many EM iterations to run
        The #...00 etc. arguments are the data itself.  E.g. for two variables:
        #00 = how many data elements are <0, 0> 
        #01 = how many data elements are <0, 1>
        #10 = how many data elements are <1, 0>
        #11 = how many data elements are <1, 1>
        
      Assumes a naive Bayes model:

           Class 
           /  \      
        var0  var1

      Uses EM to estimate parameters of Bayesian model:
      - Probability of choice of class c0 or c1
                 P(C)
      - Probability tables for var0 and var1 conditioned on choice of class
                 P(V0 | C)
                 P(V1 | C)                
      After each iteration prints the estimated probability of class given a data element
                 P(C | V0 V1 ... Vn),
      the new parameters of the model, and the log likelihood of the data.

*/



#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#define MAX_NUMVARS 6
#define MAX_NUMCLASSES 10
#define MAX_NUMCOMBOS (1<<MAX_NUMVARS)

int data[MAX_NUMCOMBOS];                          /* data[combination] -- number of training instances of combination */
float data_class[MAX_NUMCOMBOS][MAX_NUMCLASSES];  /* data_class[combination][class] -- combination is tuple of Boolean vars */
float prior[MAX_NUMCLASSES];                      /* prior[class] - prob instance is class */
float cpt[MAX_NUMVARS][MAX_NUMCLASSES][2];        /* cpt[variable][parent class][false / true] */
int seed;                                         /* initialize random number generator */
int numit;                                        /* number of iterations */

float logmarginal[MAX_NUMCOMBOS];                 /* log P(combo) */
float loglike;                                    /* log likelihood of data */

int numvars;                                      /* number Boolean vars in data */
int numcombos;                                    /* 2^numvars - number of possible tuples */
int numclasses;                                   /* Size of Bayes net */

void normalize(int n, float d[])
{
    double f = 0.0;
    int i;
    for (i=0; i<n; i++) f += d[i];
    for (i=0; i<n; i++) d[i] /= f;
}

float frand()
{
    return (float) rand();
}

char * bool2str(int width, int x, char * s)
{
    char * p = s;
    while (width>0) {
        if ((x>>(width-1))&1) *(p++)='1';
        else *(p++)='0';
        width--;
    }
    *p = 0;
    return s;
}


void emdemo(void)
{
    int c,v,t,i,r,b,combo;
    float sum, prod;
    char buf[65];

    /* Randomly initialize prior and cpt */
    for (c=0;c<numclasses;c++) prior[c] = frand();
    normalize(numclasses, prior);
    for (v=0;v<numvars;v++){
        for (c=0;c<numclasses;c++){
            for (t=0;t<2;t++)
                cpt[v][c][t] = frand();
            normalize(2, cpt[v][c]);
        }
    }

    printf("Initialization\n");
    for (c=0;c<numclasses;c++) printf("    P(c%i)= %8.6f\n", c, prior[c]);
    for (v=0;v<numvars;v++)
        for (c=0;c<numclasses;c++)
            printf("    P(~v%i|c%d)= %8.6f  P(v%i|c%d)= %8.6f\n", v, c, cpt[v][c][0], v, c, cpt[v][c][1]);
    printf("\n");

    for (i=1; i<=numit; i++){

        /* Estimate missing data */
        /* P(class|data) = alpha P(class) P(data|class)                   */
        /*               = alpha P(class) P(var1|class) P(var2|class) ... */

        for (combo=0; combo<numcombos; combo++){
            for (c=0; c<numclasses; c++){
                data_class[combo][c] = prior[c];
                for (v=0; v<numvars; v++)
                    data_class[combo][c] *= cpt[v][c][(combo>>v) & 1];
            }
            normalize(numclasses, data_class[combo]);
        }


        printf("Iteration %d\n", i);
        printf("  E-step\n");
        for (combo=0;combo<numcombos;combo++){
            printf("  ");
            for (c=0;c<numclasses;c++)
                printf("  P(c%i|%s)=  %8.6f",c,  bool2str(numvars, combo, buf), data_class[combo][c]);
            printf("\n");
        }

        /* Recalculate prior and cpt */


        /* P(class) = alpha sum_{combo} P(class|combo) data[combo] */

        for (c=0;c<numclasses;c++){
            prior[c] = 0.0;
            for (combo=0; combo<numcombos; combo++)
                prior[c] += data[combo]*data_class[combo][c];
        }
        normalize(numclasses, prior);


        /* What is this exactly? */

        for (c=0;c<numclasses;c++){
            for (v=0; v<numvars; v++){
                cpt[v][c][0] = 0.0;
                cpt[v][c][1] = 0.0;
                for (combo=0; combo<numcombos; combo++)
                    cpt[v][c][(combo>>v)&1] += data[combo]*data_class[combo][c];
                normalize(2, cpt[v][c]);
            }
        }

        printf("  M-step\n");

        for (c=0;c<numclasses;c++) printf("    P(c%i)= %8.6f\n", c, prior[c]);
        for (v=0;v<numvars;v++)
            for (c=0;c<numclasses;c++)
                printf("    P(~v%i|c%d)= %8.6f  P(v%i|c%d)= %8.6f\n", v, c, cpt[v][c][0], v, c, cpt[v][c][1]);

        /* Compute the log likelihood of the data */

        /* P(combo) = sum_class P(combo | class) P(class)
                    = sum_class ( product_var P(var = x|class) ) P(class)
                              where x = truth value var gets in combo

           log likelihood data = sum_combo data[combo] * P(combo)
        */

        for (combo=0; combo<numcombos; combo++){
            sum = 0.0;
            for (c=0; c<numclasses; c++){
                prod = prior[c];
                for (v=0; v<numvars; v++)
                    prod *= cpt[v][c][(combo>>v)&1];
                sum += prod;
            }
            logmarginal[combo] = log(sum);
        }
        loglike = 0.0;
        for (combo=0; combo<4; combo++) loglike += data[combo] * logmarginal[combo];
        printf("  Log likelihood = %8.1f\n", loglike);

        printf("\n");
    }

}


int main(int argc, char** argv)
{
    int i;


    if (argc < 7){
        fprintf(stderr, "Use: emdemo seed #iterations #classes #vars #...00 #...01 #...10 ...\n");
        exit(1);
    }
    seed = atoi(argv[1]);
    srand(seed);
    numit = atoi(argv[2]);
    numclasses = atoi(argv[3]);
    numvars = atoi(argv[4]);
    numcombos = 1 << numvars;
    if (argc != numcombos + 5) {
        fprintf(stderr, "Bad arguments\n");
        exit(1);
    }
    for (i=0;i<numcombos;i++)
        data[i] = atoi(argv[i+5]);

    emdemo();
    return 0;
}