#include <stdio.h>
#include <stdlib.h>
#include <mpc.h>
#define MAXNPROCS 20
int [host]Input();
void [host]Output();
void Partition(), SeqMult();
void [net SimpleNet(p)w] ParMult(float *, float *, float *, repl *, repl);
void [*]MxM(float *[host]x, float *[host]y, float *[host]z, int [host]n);

nettype Star(m, n[m]) {
  coord I=m;
  node {
    I>=0: n[I];
  };
};

int [host]N;

void [*]main(int [host]argc, char **[host]argv)
{
  float *[host]x, *[host]y, *[host]z;
  double [host]start;
  
  if(Input(&x, &y, argv[1])) {
    ([host]printf)("Wrong input\n");
    MPC_Exit(-1);
  }
  
  z=([host]calloc)(N*N, [host](sizeof(float)));
  start=([host]MPC_Wtime)();

  MxM(x, y, z, N);
  ([host]printf)
    ("It takes %f seconds to multiplicate\ntwo %dx%d dense matrices\n",
     (([host]MPC_Wtime)()-start), N, N);
  Output(z);
}

void [*]MxM(float *[host]x, float *[host]y, float *[host]z, int [host]n)
{
  repl nprocs, nrows[MAXNPROCS], dn;
  repl double *powers;
  
  MPC_Processors_static_info(&nprocs, &powers);
  dn=n;
  Partition(nprocs, powers, nrows, dn);
  {
    int [host]i;
    for(i=0; i<[host]nprocs; i++)
      ([host]printf)("proc=%d   nrows=%d\n", i, ([host]nrows)[i]);
  }
  {
    net Star(nprocs, nrows) w;
    float *[w]dx, *[w]dy, *[w]dz;
    repl [w]n;
    int [w]myn;
    int [w]sof;

    sof=[w](sizeof(float));
    n=[w]dn;
    myn=([w]nrows)[I coordof dx];
    
    dx=([w]calloc)(n*myn, sof);
    ([host]free)([host]dx);
    [host]dx=(void *)x;
    dy=([w]calloc)(n*n, sof);
    ([host]free)([host]dy);
    [host]dy=(void *)y;
    dz=([w]calloc)(n*myn, sof);
    ([host]free)([host]dz);
    [host]dz=(void *)z;

    ([([w]nprocs)w])ParMult(dx, dy, dz, [w]nrows, n);
  }
}

void [net SimpleNet(p)w] ParMult(float *dx, float *dy, float *dz, repl *r, repl n)
{
  repl s=0;
  int myn, i;
  int *d, *l, c;
  
  myn=r[I coordof r];
  ([(0)w])MPC_Bcast(&s, dy, 1, n*n, dy, 1);
  d=calloc(p, sizeof(int));
  l=calloc(p, sizeof(int));
  for(i=0, d[0]=0; i<p; i++)
    {
      l[i]=r[i]*n;
      if(i+1<p)
        d[i+1]=l[i]+d[i];
    }
  c=l[I coordof c];
  ([(0)w])MPC_Scatter(&s, dx, d, l, c, dx);
  ([w]SeqMult)(dx, dy, dz, myn, n);
  ([(0)w])MPC_Gather(&s, dz, d, l, c, dz);
}

void SeqMult(float *dx, float *dy, float *dz, int myr, int n)
{
  int i, j, k, ixn;
  double s;
  
  for(i=0; i<myr; i++)
    {
      ixn=i*n;
      for(j=0; j<n; j++)
        {
          for(k=0, s=0.0; k<n; k++)
            s += dx[ixn+k]*(double)(dy[k*n+j]);
          dz[ixn+j]=s;
        }
    }
}

void Partition(int p, double *v, int *r, int n)
{
  int sr, i;
  double sv;
  
  for(i=0, sv=0.0; i<p; i++)
    sv+=v[i];
  for(i=0, sr=0; i<p; i++)
    {
      r[i]=(int)(v[i]/sv*n);
      sr+=r[i];
    }
  if(sr!=n)
    r[0]+=n-sr;
}

int [host]Input(float **x, float **y, char *file)
{
  int i, j, ixN;
  FILE *fp;

  printf("N=");
  if(scanf("%d",&N)<1 || N<0)
    {
      printf("Problems with N. Bye.\n");
      return -1;
    }
  if(file==NULL)
    {
      printf("No input file.\nGenerating input ...\n", N);
      *x=calloc(N*N, sizeof(float));
      *y=calloc(N*N, sizeof(float));
      for(i=0; i<N; i++)
        {
          ixN=i*N;
          for(j=0; j<N; j++)
            if(i==j)
              {
              (*x)[ixN+j]=2.0;
              (*y)[ixN+j]=3.0;
              }
            else
              {
                (*x)[ixN+j]=0.001;
                (*y)[ixN+j]=0.001;
              }
        }
      return 0;
    }
  if((fp=fopen(file, "r"))==NULL)
    {
      printf("Cannot open the input file '%s'.\n", file);
      return 1;
    }
  *x=calloc(N*N, sizeof(float));
  *y=calloc(N*N, sizeof(float));
  for(i=0; i<N; i++)
    {
      ixN=i*N;
      for(j=0; j<N; j++)
        fscanf(fp, "%f", *x+ixN+j);
    }
  for(i=0; i<N; i++)
    {
      ixN=i*N;
      for(j=0; j<N; j++)
        fscanf(fp, "%f", *y+ixN+j);
    }
  fclose(fp);
  return 0;
}

void [host]Output(float *z)
{
  int i, j, ixN;
  double sumdiag=0.0;

  for(i=0; i<N; i++)
    {
      ixN=i*N;
      for(j=0; j<N; j++)
        sumdiag+=z[ixN+j];
    }
  printf("sumdiag/N=%.2f\n", sumdiag/N);
}