#include <alloca.h>
#include <stdio.h>
#include <math.h>
#include <float.h>
#include <stdlib.h>
#include <immintrin.h>
#include <string.h>

typedef struct {
  double x;
  double y;
} point;

double sqr(double x)
{
  return x*x;
}

double dist(double x[], double y[], int i, int j) {
  return sqrt(sqr(x[i]-x[j])+
	      sqr(y[i]-y[j]));
}

double DistSqrd(point cities[], int i, int j) {
  return (sqr(cities[i].x-cities[j].x)+
	  sqr(cities[i].y-cities[j].y));
}

void swap(double *p, double *q)
{
  double tmp=*p;
  *p = *q;
  *q = tmp;
}

// tourx and toury aligned on 32-byte boundaries
void tsp(point cities[], double tourx[], double toury[], int ncities)
{
  int i,kk;
  tourx[0] = cities[ncities-1].x;
  toury[0] = cities[ncities-1].y;
  for (i=1; i<ncities; i++) {
    tourx[i]=cities[i-1].x;
    toury[i]=cities[i-1].y;
  }

  for (i=1; i<ncities; i++) {
    double ThisX = tourx[i-1];
    double ThisY = toury[i-1];
    double CloseDist = DBL_MAX;
    int min_k = i;
    int k     = i;
    if ((i & 3) != 0 || ncities - i < 8) {
      // scalar prolog
      int nlast = ncities - i < 8 ? ncities : (i + 4) & (-4);
      for (; k < nlast; ++k) {
        double ThisDist = sqr(tourx[k]-ThisX)+sqr(toury[k]-ThisY);
        if (ThisDist <= CloseDist) {
          CloseDist = ThisDist;
          min_k = k;
        }
      }
    }

    int ni = (ncities - k) / 4;
    if (ni > 0) {
      // main loop
      __m256d thisx_v   = _mm256_set1_pd(ThisX);
      __m256d thisy_v   = _mm256_set1_pd(ThisY);
      __m256d minDist_v = _mm256_set1_pd(CloseDist);
      __m256d minK_v    = _mm256_set1_pd((double)min_k);
      __m256d k_v       = _mm256_setr_pd((double)(k+0), (double)(k+1), (double)(k+2), (double)(k+3));
      __m256d kInc_v    = _mm256_set1_pd((double)4);
      const __m256d* tourx_v = (const __m256d*)&tourx[k];
      const __m256d* toury_v = (const __m256d*)&toury[k];
      k += ni*4;
      for (kk = 0; kk < ni; ++kk) {
        __m256d dx_v   = _mm256_sub_pd(tourx_v[kk], thisx_v);
        __m256d dy_v   = _mm256_sub_pd(toury_v[kk], thisy_v);
        __m256d dist_v = _mm256_add_pd(_mm256_mul_pd(dx_v, dx_v),_mm256_mul_pd(dy_v, dy_v));
        __m256d le_v   = _mm256_cmp_pd(dist_v, minDist_v, _CMP_LE_OQ);
        minDist_v      = _mm256_blendv_pd(minDist_v, dist_v, le_v);
        minK_v         = _mm256_blendv_pd(minK_v,    k_v,    le_v);
        k_v = _mm256_add_pd(k_v, kInc_v);
      }
      // fold
      __m256d minDist_vh = _mm256_permute2f128_pd(minDist_v, minDist_v, 1);
      __m256d minK_vh    = _mm256_permute2f128_pd(minK_v,    minK_v,    1);
      __m256d le_v       = _mm256_cmp_pd(minDist_v, minDist_vh, _CMP_LE_OQ);
      minDist_v          = _mm256_blendv_pd(minDist_vh, minDist_v, le_v);
      minK_v             = _mm256_blendv_pd(minK_vh,    minK_v,    le_v);
      minDist_vh         = _mm256_permute_pd(minDist_v, 1);
      minK_vh            = _mm256_permute_pd(minK_v,    1);
      le_v               = _mm256_cmp_pd(minDist_v, minDist_vh, _CMP_LE_OQ);
      minDist_v          = _mm256_blendv_pd(minDist_vh, minDist_v, le_v);
      minK_v             = _mm256_blendv_pd(minK_vh,    minK_v,    le_v);
      // CloseDist = (double)_mm256_castpd256_pd128(minDist_v);
      // min_k = (double)_mm256_castpd256_pd128(minK_v);
      memcpy(&CloseDist, &minDist_v, sizeof(double));
      double min_kd;
      memcpy(&min_kd, &minK_v, sizeof(double));
      min_k = (int)min_kd;
    }
    // scalar epilog
    for (; k < ncities; ++k) {
      double ThisDist = sqr(tourx[k]-ThisX)+sqr(toury[k]-ThisY);
      if (ThisDist <= CloseDist) {
        CloseDist = ThisDist;
        min_k = k;
      }
    }
    swap(&tourx[i],&tourx[min_k]);
    swap(&toury[i],&toury[min_k]);
  }
}

// void tsp(point cities[], double tourx[], double toury[], int ncities)
// {
//   int i;
//   double *ClosePt=tourx;
//   double CloseDist;
//   double *p;
//   
//   for (i=1; i<ncities; i++) {
//     point c=cities[i-1];
//     tourx[i]=c.x;
//     toury[i]=c.y;
//   }
//   tourx[0] = cities[ncities-1].x;
//   toury[0] = cities[ncities-1].y;
//   
//   for (i=1; i<ncities; i++) {
//     double ThisX = tourx[i-1];
//     double ThisY = toury[i-1];
//     CloseDist = DBL_MAX;
//     for (p=&tourx[ncities];;) {
//       asm("vbroadcastsd %[closeDist], %%ymm4\n"
//           "	vbroadcastsd %[thisx], %%ymm6\n"
//           "	vbroadcastsd %[thisy], %%ymm7\n"
//           "1:	vsubpd -32(%[p]), %%ymm6, %%ymm5\n"
//           "	vsubpd -32(%[p],%[d]), %%ymm7, %%ymm8\n"
//           "	subq $32, %[p]\n"
//           "	vmulpd %%ymm5, %%ymm5, %%ymm5\n"
//           "	vmulpd %%ymm8, %%ymm8, %%ymm8\n"
//           "	prefetchnta -320(%[p])\n"
//           "	prefetchnta -320(%[p],%[d])\n"
//           "	vaddpd %%ymm8, %%ymm5, %%ymm5\n"
//           "	vcmpltpd %%ymm4, %%ymm5, %%ymm5\n"
//           "     vmovmskpd %%ymm5, %%rcx\n"
//           "	test %%rcx, %%rcx\n"
//           "	je 1b\n"
//           "	bsr %%rcx, %%rcx\n"
//           "	leaq (%[p],%%rcx, 8), %[p]"
//           : [p] "+r"(p)
//           : [closeDist] "m"(CloseDist), [thisx] "m"(tourx[i-1]), [thisy] "m"(toury[i-1]), [d] "r"(((long)toury)-((long)tourx))
//           : "%ymm4", "%ymm5", "%ymm6", "%ymm7", "%ymm8", "%rcx");
//       // double ThisDist;
//       // do {
//       //   p--;
//       //   ThisDist = sqr(*p-ThisX);
//       // } while (ThisDist>CloseDist);
//       // ThisDist += sqr(toury[p-tourx]-ThisY);
//       if (p < tourx+i)
//         break;
//       CloseDist = sqr(*p-ThisX)+sqr(toury[p-tourx]-ThisY);
//       ClosePt = p;
//     }
//     swap(&tourx[i],ClosePt);
//     swap(&toury[i],&toury[ClosePt-tourx]);
//   }
// }

int main(int argc, char *argv[])
{
  int i, ncities;
  point *cities;
  double *tourx;
  double *toury;
  FILE *psfile;
  double sumdist = 0.0;

  if (argc!=2) {
    fprintf(stderr, "usage: %s <ncities>", argv[0]);
    exit(1);
  }
  ncities = atoi(argv[1]);
  cities = alloca(ncities*sizeof(point));
  tourx=alloca((ncities+4)*sizeof(double))+4*sizeof(double);
  toury=alloca(ncities*sizeof(double));
  for (i=0; i<ncities; i++) {
    cities[i].x = ((double)(random()))/(double)(1U<<31);
    cities[i].y = ((double)(random()))/(double)(1U<<31);
  }
  tsp(cities,tourx,toury,ncities);
  psfile = fopen("tsp.eps","w");
  fprintf(psfile, "%%!PS-Adobe-2.0 EPSF-1.2\n%%%%BoundingBox: 0 0 300 300\n");
  fprintf(psfile, "1 setlinejoin\n0 setlinewidth\n");
  fprintf(psfile, "%f %f moveto\n",
	  300.0*tourx[0], 300.0*toury[0]);
  for (i=1; i<ncities; i++) {
    fprintf(psfile, "%f %f lineto\n",
	    300.0*tourx[i], 300.0*toury[i]);
    sumdist += dist(tourx, toury, i-1, i);
  }
  fprintf(psfile,"stroke\n");
  printf("sumdist = %f\n", sumdist);
  exit(0);
}