/* Write a numpy array file, single density only, as floats or doubles */

/* MJR 11/2020, 6/2023 */

#include<stdio.h>
#include<stdlib.h>

#include "c2xsf.h"

#define MAX_HDR_LEN 256

void npy_write(FILE* outfile, struct grid *g){
  char magic[6]={0x93,'N','U','M','P','Y'};
  char *hdr,c;
  void *data;
  int i,j,k;
  long ndata;

  if ((!g)||(!g->data)) error_exit("No grid data to write");

  ndata=g->size[0]*g->size[1]*g->size[2];
  if (g->comps>1){ /* Transpose data */
    if (flags&HIPREC)
      data=malloc(sizeof(double)*ndata*g->comps);
    else
      data=malloc(sizeof(float)*ndata*g->comps);
    if (!data) error_exit("Malloc error in npy_write");
    k=0;
    if (flags&HIPREC)
      for(i=0;i<ndata;i++)
	for(j=0;j<g->comps;j++)
	  ((double*)data)[k++]=g->data[i+ndata*j];
    else
      for(i=0;i<ndata;i++)
	for(j=0;j<g->comps;j++)
	  ((float*)data)[k++]=g->data[i+ndata*j];
    ndata*=g->comps;
  }
  else if (flags&HIPREC){
    data=g->data;
  }
  else{
    data=malloc(sizeof(float)*ndata);
    if (!data) error_exit("Malloc error in npy_write");
    for(i=0;i<ndata;i++)
      ((float*)data)[i]=g->data[i];
  }
    
  if (fwrite(magic,1,6,outfile)!=6)
    error_exit("Error writing magic number");

  c=1;
  fwrite(&c,1,1,outfile);
  c=0;
  fwrite(&c,1,1,outfile);

  hdr=malloc(MAX_HDR_LEN);
  for(i=0;i<MAX_HDR_LEN;i++) hdr[i]=' ';

  if (g->comps==1)
    i=snprintf(hdr,MAX_HDR_LEN,"{'descr': '%sf%d', 'fortran_order': False,"
	       " 'shape': (%d, %d, %d) }",(self_little_endian())?"<":">",
	       (flags&HIPREC)?8:4,
	       g->size[0],g->size[1],g->size[2]);
  else
    i=snprintf(hdr,MAX_HDR_LEN,"{'descr': '%sf%d', 'fortran_order': False,"
	       " 'shape': (%d, %d, %d, %d) }",(self_little_endian())?"<":">",
	       (flags&HIPREC)?8:4,
	       g->size[0],g->size[1],g->size[2],g->comps);

  /* Remove terminating null */
  hdr[i]=' ';

  /* Some docs say round to multiple of 16, some to multiple of 64.
   * Here 64 is used, so round i+10 up to multiple of 64 */
  i=i+10;
  i=i+(64-(i&63));
  i=i-10;
  hdr[i-1]='\n';

  c=i&0xff;
  fwrite(&c,1,1,outfile);
  c=(i&0xff00)>>8;
  fwrite(&c,1,1,outfile);
  
  fwrite(hdr,1,i,outfile);
  free(hdr);
  
  if (flags&HIPREC)
    ndata*=sizeof(double);
  else
    ndata*=sizeof(float);
  
  i=fwrite(data,1,ndata,outfile);

  if (i!=ndata)
    error_exit("Error writing array");

  if ((g->comps!=1)||(!(flags&HIPREC))) free(data);
  
}
