/*  prog1p1.c -- Restoration of Analog Audio Recordings
    Written October 28, 2011 by Eric Olson for Math 761

    This program makes use of the fast Fourier transform given
    in the file fft.c.  The command

        gcc -o prog1p1 prog1p1.c fft.c -lm

    will compile this program in most POSIX/GNU environments.
    Please use a version of the GNU C compiler that supports 
    native complex data types.

*/

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <math.h>
#include <complex.h>

#define N 262144
#define alpha 1.75
#define beta 0.5
#define sigma 0.0004
#define epsilon 0.000001

/* The filter kfunc specifies the loss of high frequencies and tape
   print through by means of the constants alpha, beta and sigma. */

double kfunc(double t){
    double tm1=t;
    double tm2=t+0.5;
    return alpha*exp(-tm1*tm1/(2*sigma*sigma))/(sigma*sqrt(2*M_PI))
        +beta*exp(-tm2*tm2/(4*sigma*sigma))/(sigma*sqrt(4*M_PI));
}

void readsignal(complex double X[],char *p){
    float buf[N];
    int fd,i;
    printf("Reading %s...\n",p);
    if((fd=open(p,O_RDONLY))==-1){
        fprintf(stderr,"Can't open %s for read!\n",p);
        exit(1);
    }
    if(read(fd,buf,sizeof(float)*N)!=sizeof(float)*N){
        fprintf(stderr,"Error reading data from %s!\n",p);
        exit(2);
    }
    close(fd);
    for(i=0;i<N;i++) X[i]=buf[i];
}

void writesignal(complex double X[],char *p){
    float buf[N];
    int fd,i;
    printf("Writing %s...\n",p);
    for(i=0;i<N;i++) buf[i]=X[i];
    if((fd=open(p,O_WRONLY|O_CREAT|O_TRUNC,0644))==-1){
        fprintf(stderr,"Can't open %s for write!\n",p);
        exit(3);
    }
    if(write(fd,buf,sizeof(float)*N)!=sizeof(float)*N){
        fprintf(stderr,"Error writing data to %s!\n",p);
        exit(4);
    }
    close(fd);
}

/*  Compute relative error how well X approximates x */

double relerr(complex double X[],complex double x[]){
    double r1=0.0,r2=0.0;
    int i;
    for(i=1;i<N;i++) {
        complex double t=x[i];
        r1+=t*conj(t);
        t-=X[i];
        r2+=t*conj(t);
    }
    return sqrt(r2/r1);
}

complex double K[N],W[N],X[N];
main(){
    int i;
    printf("prog1p1.c -- Restoration of Analog Audio Recordings\n"
        "Written October 28, 2011 by Eric Olson for Math 761\n\n");

    readsignal(W,"signalw.raw");

    /*  The function kfilt(t) takes an argument of time in seconds
        whereas the K[i] is sampled at 44.1 khz.

            i=1 corresponds to t=1/44100 of a second and
            i=N-1 corresponds to t=-1/44100 of a second. */

    for(i=0;i<N;i++){
        double t;
        if(i<N/2) t=(double)i/44100.0;
        else t=(double)(i-N)/44100.0;
        K[i]=kfunc(t)/44100.0;
    }

    printf("Approximating X...\n");
    fft(K,N,1);
    fft(W,N,1);
    for(i=0;i<N;i++) {

        /*  This is the line that needs to be modified to implement
            the optimal Wiener-Kolmogorov filter. */

        complex double Ki=K[i];
        if(cabs(Ki)>epsilon) W[i]/=Ki;
	else W[i]=0;
    }
    fft(W,N,-1);

    writesignal(W,"wtox.raw");

    /*  Compute the error in the recovered signal */
        
    readsignal(X,"signalx.raw");

    printf("\nThe error is E=%g\n",relerr(W,X));
    exit(0);
}