/*  Program to calculate partial F's, for each founder within pedigree */

/*  Data are expected to be in files entered as function arguments.  */
/*  Data entry files:
                argv[1]   mice.out     output file name
                argv[2]   newmice.txt  me/ma/pa input file
*/
/*      Founders MUST come first in the data set. */

#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <dos.h>
#include <malloc.h>

void filerror(char *s) {printf("\nCannot open file %s\n",s); exit(1);}

#define MAXWORD 40

void shrink(int);
void ffcalc(int);
void booksize(void);
int setrr(void);

FILE *parents,*mprn;
short **rr,**ff,**fk;
int na0,na,nrec,idmax,nanim,nfndr;
short  *mom,*pop,*ndx,*id,*ivalid;
short  *mom0,*pop0,*ndx0,*id0;
int sizeptr,sizeshort,sizedbl;


int main(int argc, char **argv)
{
        int i,j,k;
        char fprn[MAXWORD],parentsf[MAXWORD];

        sizeshort = sizeof(short);
        sizeptr = sizeof(mom);
        sizedbl = sizeof(double);

        strcpy(parentsf,argv[2]);
        strcpy(fprn,argv[1]);

        if((parents = fopen(parentsf,"r")) == NULL) filerror(parentsf);

        booksize();
        setrr();

        id[0] = mom[0] = pop[0] = ndx[0] = 0;
        id0[0] = mom0[0] = pop0[0] = ndx0[0] = 0;

        printf("\nidmax = %d, na0 = %d, nrec = %d, nfndr = %d",idmax,na0,nrec,nfndr);
/*  nfndr is the number of founders (incl founder 0).
        founder 0 is the code for calculating total inbreeding.
        na0 is the number of mice in analysis, plus 1  */

        for(i = 1; fscanf(parents,"%hd,%hd,%hd\n",id0 + i,mom0 + i, pop0 + i) == 3; i++) {
                if(i != ndx0[id0[i]]) {
                        printf("\nScrew-up in ID sequence.  i = %d, id = %d, ndx = %d",i,id0[i],ndx0[id0[i]]);
                        exit(1);
                }
        }
        fclose(parents);

        for(j = 0; j < nfndr; j++) {
                printf("\nCalculating partial Fs for founder # %d: %d\n",j,id[j]);
                na = na0;
                for(i = nfndr; i < na0; i++) {
                        mom[i] = mom0[i];
                        pop[i] = pop0[i];
                        id[i] = id0[i];
                        ndx[id[i]] = i;
                }
                if(j > 0) {
                        rr[j][j] = 20000;
                        for(k = 1; k < nfndr; k++) {
                                if(k != j) {
                                        rr[k][k] = 0;
                                }
                        }
                }
                else { /* Total F calculations */
                        for(k = 1; k < nfndr; k++) {
                                rr[k][k] = 20000;
                                mom[k] = mom0[k];
                                pop[k] = pop0[k];
                                id[k] = id0[k];
                                ndx[id[k]] = k;
                        }
                }
                ffcalc(j);
        }

        if((mprn = fopen(fprn,"w+")) == NULL) filerror(fprn);
        fprintf(mprn,
"Inbreeding coefficients, F(all), partial inbreeding coefficients, F(x),\n  and founder contributions, P(x), for %s\n   ID   DAM   SIRE    F(all) ",strtok(fprn,"."));
        for(j = 1; j < nfndr; j++) {
                fprintf(mprn,"F(%3d), ",id0[j]);
        }
        for(j = 1; j < nfndr; j++) {
                fprintf(mprn,"P(%3d), ",id0[j]);
        }
        fputs("\n\n",mprn);
        for(i = 1; i < na0; i++) {
                fprintf(mprn,"%5d, %5d, %5d",id0[i],mom0[i],pop0[i]);
                for(j = 0; j < nfndr; j++) {
                        fprintf(mprn,", %6.4f", 0.00005 * (double) ff[i][j]);
                }
                for(j = 1; j < nfndr; j++) {
                        fprintf(mprn,", %6.4f",0.00005 * (double) fk[i][j]);
                }
                fprintf(mprn,"\n");
        }
        fclose(mprn);

        printf("\n\n Output will be in file %s",argv[1]);

        return(1);
}


void booksize()
{
        int i,d1,d2,d3;

        puts("\nCounting animals ...");
        idmax = 0;
        while( fscanf(parents,"%d,%d,%d\n",&i,&d2,&d3) == 3 ) {
                if(i > idmax) idmax = i;
        }

        rewind(parents);
        if(idmax > 32000) {
                puts("\nID too large.");
                exit(1);
        }
        if(( ndx = (short *)malloc((idmax + 1) * sizeshort) ) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        if(( ndx0 = (short *)malloc((idmax + 1) * sizeshort) ) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        for(i = 1; i <= idmax; i++) ndx0[i] = -1;
        i = 1;
        nfndr = 1;  /* Founder 0 is all founders -- full F */

        while( fscanf(parents,"%d,%d,%d\n",&d1,&d2,&d3) == 3 ) {
                if(d2 == 0 && d3 == 0) {
                        if(nfndr != i) {
                                puts("\n Founders not at top of file.");
                                exit(1);
                        }
                        nfndr++;
                }
                if(ndx0[d1] < 0) ndx0[d1] = i++;
        }
        rewind(parents);
        na0 = i;
        if(nfndr > 100) {
                puts("\nToo many founders.");
                exit(1);
        }
        if( (ivalid = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        for(i = 1; i < nfndr; i++) ivalid[i] = 1;

        if( (mom = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        if( (pop = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        if( (id = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        if( (mom0 = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        if( (pop0 = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        if( (id0 = (short *)malloc(na0 * sizeshort)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }
        return;
}


int setrr()
{
        int i,j;

        nrec = na0 + 1;

        if( (ff = (short **)malloc(na0 * sizeptr)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }

        for(i = 0; i < na0; i++) {
                if((ff[i] = (short *)malloc(nfndr * sizeshort)) == NULL) {
                        puts("\nNot enough memory.");
                        exit(0);
                }
        }
        for(i = 1; i < nfndr; i++) {
                for(j = 0; j < nfndr; j++) ff[i][j] = 0;
        }

        if( (fk = (short **)malloc(na0 * sizeptr)) == NULL) {
                puts("\nError in memory allocation.");
                exit(1);
        }

        for(i = 0; i < na0; i++) {
                if((fk[i] = (short *)malloc(nfndr * sizeshort)) == NULL) {
                        puts("\nNot enough memory.");
                        exit(0);
                }
        }

        for(i = 1; i < nfndr; i++) {
                fk[i][i] = 20000;
                for(j = 1; j < nfndr; j++) if(i != j) fk[i][j] = 0;
        }

        if( (rr = (short **)malloc(nrec * sizeptr)) == NULL) {
                puts("\nError in memory allocation. Probably insufficient memory.");
                exit(1);
        }

        for(i = 0; i < nrec; i++) {
                if((rr[i] = (short *)malloc((i + 1) * sizeshort)) == NULL) {
                        nrec = i - 10;
                        for(j = nrec; j < i; j++) free(rr[j]);
                }
        }
        if(nrec < nfndr + 1) {
                puts("\nNot enough memory for analysis.");
                exit(1);
        }
        for(i = 0; i < nrec; i++) rr[i][0] = 0;
        for(i = 1; i < nfndr; i++) {
                for(j = 1; j < i; j++) rr[i][j] = 0;
        }
        return(1);
}


void ffcalc(int fi)
{
        int i,i0,j,m,p,rmj,rpj,rmp;

        for(i = nfndr; i < na; i++) {
                if(i >= nrec) {
                        nanim = i - 1;
                        shrink(fi);
                        i = nanim + 1;
                }
                m = ndx[mom[i]];
                p = ndx[pop[i]];
                if(m < 0 || p < 0) {
                        printf("\nParent %d or parent %d does not exist in data matrix.\n",
                                mom[i],pop[i]);
                        exit(1);
                }
                for(j = 1; j < i; j++) {
                        if(m < j) rmj = rr[j][m];
                        else rmj = rr[m][j];
                        if(p < j) rpj = rr[j][p];
                        else rpj = rr[p][j];
                        rr[i][j] = (rmj + rpj) / 2;
                }
                if(m < p) rmp = rr[p][m] / 2;
                else rmp = rr[m][p] / 2;
                i0 = ndx0[id[i]];
                ff[i0][fi] = rmp;
                if(fi > 0) {
                	fk[i0][fi] = rr[i0][fi];
                    rr[i][i] = rmp + fk[i0][fi];
                }
                else rr[i][i] = rmp + 20000;
        }
}


void shrink(int fi)
{
        int i,j,k,l;

        puts("\nShrinking ... ");
        printf("\n%d animals out of %d already in matrix.",nanim,na);

        for(i = nfndr; i < na; i++) ivalid[i] = 0;
        for(i = nanim + 1; i < na; i++) {
                ivalid[i] = 1;
                k = ndx[mom[i]];   /* Keep parents in if kids not yet in */
                l = ndx[pop[i]];
                if(k < 0 || l < 0) {
                        printf("\nShrink: Parent %d or %d missing.",mom[i],pop[i]);
                        exit(1);
                }
                if(ivalid[k] == 0) ivalid[k] = 1;
                if(ivalid[l] == 0) ivalid[l] = 1;
        }
        for(i = nfndr; i < na; i++) ndx[id[i]] = -1;
        for(i = nfndr, j = nfndr - 1; i < na; i++) {
                if(ivalid[i] == 0) continue;
                j++;
                ndx[id[i]] = j;
                if(i == j) continue;
                id[j] = id[i];
                mom[j] = mom[i];
                pop[j] = pop[i];
                if(i > nanim) continue;

                for(k = 1, l = 1; k <= i; k++) {
                        if(ivalid[k]) rr[j][l++] = rr[i][k];
                }
        }
        na -= nanim; /* Accounts for ones still to be added to matrix */
        for(i = 1, j = 0; i <= nanim; i++) if(ivalid[i]) j++;
        nanim = j;   /* nanim is the number of mice already in the data matrix */
        if(nanim > nrec - 2) {
                puts("\nToo many living animals");
                exit(1);
        }
        na += nanim;   /*  na is 1 more than the number of mice in the matrix or still to be added. */
}