#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include "gpk_std.h"
#include "gpklib/avg.h"
#include "gpklib/ieee754.h"

#define D if(0)

LOCAL int v_wt_compare(C void *a, C void *b)
{
 C struct v_wt *aa = (C struct v_wt *)a;
 C struct v_wt *bb = (C struct v_wt *)b;

 return bb->v>aa->v ? -1 : (bb->v<aa->v ? 1 : 0);
}


LOCAL double ssqrt(double x) { return x>0 ? sqrt(x) : 0;}
inline double xx(double x) { return x*x;}

struct a_w avg_linear(C struct v_wt *in, int n)
{
 int i;
 double wtsum = 0, sum = 0, ss = 0;
 for(i=0;i<n;i++)	{
	wtsum += in[i].wt;
	C double tmp = in[i].v - in[0].v;
	sum += in[i].wt * tmp;
	ss += in[i].wt * xx(tmp);
	}
 struct a_w rv;
 rv.avg = in[0].v + sum/wtsum;
 rv.width = ssqrt(ss/wtsum - xx(sum/wtsum));
 rv.wt = wtsum;
 return rv;
}


/* Chop specifies how much weight to trim off each end.
	It should be between 0 and 0.5. */
struct a_w avg_robust(C struct v_wt *in, int n, double chop)
{
 struct a_w rv;
 assert(chop >= 0);
 rv.avg = quiet_nan(0);
 rv.width = quiet_nan(0);
 rv.wt = 0.0;
 if(n == 0)	{ return rv;}
 if(n == 1)	{ rv.avg = in[0].v; rv.width = 0; rv.wt = in[0].wt; return rv;}

 if(chop == 0)
	return avg_linear(in, n);

 D fprintf(stderr, "new struct v_wt\n");
 struct v_wt *z = new struct v_wt[n];
 int i;
 double wtsum = 0;
 for(i=0; i<n; i++)	{
	D fprintf(stderr, "copying %d = %g %g\n", i, in[i].wt, in[i].v);
	if(!finite(in[i].wt) || isnan(in[i].v))	{
		delete [] z;
		return rv;
		}
	wtsum += in[i].wt;
	assert(in[i].wt >= 0.0);
	z[i] = in[i];
	}


 D fprintf(stderr, "qsort\n");
 qsort(z, n, sizeof(struct v_wt), v_wt_compare);
 assert(z[n-1].v >= z[0].v);


 assert(chop>=0 && chop<=0.5);
 C double fc1 = chop*wtsum;
 C double fc2 = wtsum-fc1;
 C double wcf = exp(.2402 + 3.474*chop);
 assert(wcf > 0.0);
 D fprintf(stderr, "sum1\n");
 double tmpws = 0.0;
 for(i=0; i<n && tmpws<fc1; tmpws += z[i].wt, i++)
	;
 D fprintf(stderr, "sum2\n");

 assert(fc2 > fc1);
 assert(fc1 >= 0);
 assert(wtsum >= fc2);
 assert(tmpws >= fc1);
 assert(i <= n);
 C double tmpws0 = tmpws;
 double s = i>0 ? z[i-1].v*(tmpws0-fc1) : 0;
 int j;
 for(j=i; j<n && tmpws<=fc2; tmpws+= z[j].wt, j++)
	s += z[j].v*z[j].wt;
 assert(tmpws >= fc2);
 assert(j > 0  && j<=n);
 s -= z[j-1].v*(tmpws - fc2);
 rv.avg =  s/(fc2 - fc1);

 D fprintf(stderr, "sum3\n");
 assert(finite(rv.avg));
 assert(i-1 < n);
 assert(j <= n);
 double sd = i>0 ? fabs(z[i-1].v-rv.avg)*(tmpws0-fc1) : 0;
 int k;
 OPENMP(parallel for reduction(+: sd) lastprivate(k) schedule(guided, 128))
 for(k=i; k<j; k++)
	sd += fabs(z[k].v-rv.avg)*z[k].wt;
 assert(k>0 and k-1<n);
 sd -= fabs(z[k-1].v-rv.avg)*(tmpws-fc2);
 D fprintf(stderr, "delete\n");
 delete[] z;
 rv.width = n>0 ? wcf * sd / ((fc2-fc1)*sqrt(1.0-1.0/n)) : 0;
 rv.wt = wtsum * (1 - 0.5*2*chop);
 assert(rv.width >= 0.0);
 D fprintf(stderr, "return\n");
 return rv;
}


struct a_w avg_robust(const float *z, int n)
{
 struct v_wt *zz = new struct v_wt[n];
 struct a_w rv;
 assert(zz);
 int i;
 OPENMP(parallel for schedule(guided, 512))
 for(i=0; i<n; i++)	{
	zz[i].v = z[i];
	zz[i].wt = 1.0;
	}
 rv = avg_robust(zz, n);
 D fprintf(stderr, "Avg done\n");
 delete [] zz;
 D fprintf(stderr, "Outer return\n");
 return rv;
}
