Subversion Repositories planix.SVN

#include <u.h>
#include <libc.h>
#include <bio.h>
#include <regexp.h>
#include "hash.h"

enum
{
        MAXTAB = 256,
        MAXBEST = 32,
};

typedef struct Table Table;
struct Table
{
        char *file;
        Hash *hash;
        int nmsg;
};

typedef struct Word Word;
struct Word
{
        Stringtab *s;   /* from hmsg */
        int count[MAXTAB];      /* counts from each table */
        double p[MAXTAB];       /* probabilities from each table */
        double mp;      /* max probability */
        int mi;         /* w.p[w.mi] = w.mp */
};

Table tab[MAXTAB];
int ntab;

Word best[MAXBEST];
int mbest;
int nbest;

int debug;

void
usage(void)
{
        fprint(2, "usage: bayes [-D] [-m maxword] boxhash ... ~ msghash ...\n");
        exits("usage");
}

void*
emalloc(int n)
{
        void *v;

        v = mallocz(n, 1);
        if(v == nil)
                sysfatal("out of memory");
        return v;
}

void
noteword(Word *w)
{
        int i;

        for(i=nbest-1; i>=0; i--)
                if(w->mp < best[i].mp)
                        break;
        i++;

        if(i >= mbest)
                return;
        if(nbest == mbest)
                nbest--;
        if(i < nbest)
                memmove(&best[i+1], &best[i], (nbest-i)*sizeof(best[0]));
        best[i] = *w;
        nbest++;
}

Hash*
hread(char *s)
{
        Hash *h;
        Biobuf *b;

        if((b = Bopenlock(s, OREAD)) == nil)
                sysfatal("open %s: %r", s);

        h = emalloc(sizeof(Hash));
        Breadhash(b, h, 1);
        Bterm(b);
        return h;
}

void
main(int argc, char **argv)
{
        int i, j, a, mi, oi, tot, keywords;
        double totp, p, xp[MAXTAB];
        Hash *hmsg;
        Word w;
        Stringtab *s, *t;
        Biobuf bout;

        mbest = 15;
        keywords = 0;
        ARGBEGIN{
        case 'D':
                debug = 1;
                break;
        case 'k':
                keywords = 1;
                break;
        case 'm':
                mbest = atoi(EARGF(usage()));
                if(mbest > MAXBEST)
                        sysfatal("cannot keep more than %d words", MAXBEST);
                break;
        default:
                usage();
        }ARGEND

        for(i=0; i<argc; i++)
                if(strcmp(argv[i], "~") == 0)
                        break;

        if(i > MAXTAB)
                sysfatal("cannot handle more than %d tables", MAXTAB);

        if(i+1 >= argc)
                usage();

        for(i=0; i<argc; i++){
                if(strcmp(argv[i], "~") == 0)
                        break;
                tab[ntab].file = argv[i];
                tab[ntab].hash = hread(argv[i]);
                s = findstab(tab[ntab].hash, "*nmsg*", 6, 1);
                if(s == nil || s->count == 0)
                        tab[ntab].nmsg = 1;
                else
                        tab[ntab].nmsg = s->count;
                ntab++;
        }

        Binit(&bout, 1, OWRITE);

        oi = ++i;
        for(a=i; a<argc; a++){
                hmsg = hread(argv[a]);
                nbest = 0;
                for(s=hmsg->all; s; s=s->link){
                        w.s = s;
                        tot = 0;
                        totp = 0.0;
                        for(i=0; i<ntab; i++){
                                t = findstab(tab[i].hash, s->str, s->n, 0);
                                if(t == nil)
                                        w.count[i] = 0;
                                else
                                        w.count[i] = t->count;
                                tot += w.count[i];
                                p = w.count[i]/(double)tab[i].nmsg;
                                if(p >= 1.0)
                                        p = 1.0;
                                w.p[i] = p;
                                totp += p;
                        }

                        if(tot < 5){            /* word does not appear enough; give to box 0 */
                                w.p[0] = 0.5;
                                for(i=1; i<ntab; i++)
                                        w.p[i] = 0.1;
                                w.mp = 0.5;
                                w.mi = 0;
                                noteword(&w);
                                continue;
                        }

                        w.mp = 0.0;
                        for(i=0; i<ntab; i++){
                                p = w.p[i];
                                p /= totp;
                                if(p < 0.01)
                                        p = 0.01;
                                else if(p > 0.99)
                                        p = 0.99;
                                if(p > w.mp){
                                        w.mp = p;
                                        w.mi = i;
                                }
                                w.p[i] = p;
                        }
                        noteword(&w);
                }

                totp = 0.0;
                for(i=0; i<ntab; i++){
                        p = 1.0;
                        for(j=0; j<nbest; j++)
                                p *= best[j].p[i];
                        xp[i] = p;
                        totp += p;
                }
                for(i=0; i<ntab; i++)
                        xp[i] /= totp;
                mi = 0;
                for(i=1; i<ntab; i++)
                        if(xp[i] > xp[mi])
                                mi = i;
                if(oi != argc-1)
                        Bprint(&bout, "%s: ", argv[a]);
                Bprint(&bout, "%s %f", tab[mi].file, xp[mi]);
                if(keywords){
                        for(i=0; i<nbest; i++){
                                Bprint(&bout, " ");
                                Bwrite(&bout, best[i].s->str, best[i].s->n);
                                Bprint(&bout, " %f", best[i].p[mi]);
                        }
                }
                freehash(hmsg);
                Bprint(&bout, "\n");
                if(debug){
                        for(i=0; i<nbest; i++){
                                Bwrite(&bout, best[i].s->str, best[i].s->n);
                                Bprint(&bout, " %f", best[i].p[mi]);
                                if(best[i].p[mi] < best[i].mp)
                                        Bprint(&bout, " (%f %s)", best[i].mp, tab[best[i].mi].file);
                                Bprint(&bout, "\n");
                        }
                }
        }
        Bterm(&bout);
}