/* save-load-sound.c
 * atsa: ATS analysis implementation
 * Oscar Pablo Di Liscia / Pete Moss / Juan Pampin 
 */

#include "atsa.h"

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

static int my_filelength(FILE *fpt);
static double byte_swap(double *in);
static void byte_swap_header(ATS_HEADER *head, int swapMagic);

/* ats_load
 * ========
 * loads an ATS_SOUND from disk (whatever little/big endian byte order).
 * sound: pointer to ATS_SOUND structure
 * infile: pointer to input ats file
 * Returns the total file size in bytes (or <0 value upon error)
 */
int ats_load(ATS_SOUND *sound, const char *infile, int* sound_type)
{
    FILE* pATSFile;
	int nFileSize;
    ATS_HEADER sHeader;
    int bNeedByteSwap = 0;
    int n1, n2;
    int nTheoricalDataSize;
    double dMaxAmp, dMaxFreq, dTmp;
    int nFrameIdx, nPartIdx;

    fprintf(stdout, "Loading ATS sound from '%s' ...\n", infile);

    // Open the .ats file.
    pATSFile = fopen(infile, "rb");
    if (!pATSFile) 
    {
	fprintf(stderr, "Error: Could not open '%s'\n", infile);
	return -1;
    }

    // Determine real file length.
    nFileSize = my_filelength(pATSFile);
  
    // Load header from file.
    fread(&sHeader, sizeof(ATS_HEADER), 1, pATSFile);

    // Check file format and little/big endian (swap header byte if necessary).
    if (sHeader.mag != ATS_MAGIC_NUMBER) 
    { 
	double dMagNum = sHeader.mag;		
	sHeader.mag = byte_swap(&dMagNum);

	if ((int)sHeader.mag != (int)ATS_MAGIC_NUMBER) 
	{ 
	    fprintf(stderr, "Error: Could not find magic number in header\n");
	    fclose(pATSFile);
	    return -1;
	} 
	else 
	{
	    bNeedByteSwap=1;               //now we know that we must swap the bytes 
	    byte_swap_header(&sHeader, 0);    //swap bytes of header data first
	}  
    }

    // Check consistency between duration, sampling rate, frame size and number of frames.
    if (sHeader.fra * sHeader.fs != sHeader.sr * sHeader.dur)
    {
	sHeader.dur = sHeader.fra * sHeader.fs / sHeader.sr;
	fprintf(stdout, "Warning: Unconsistent sound duration ; "
		"trusting in frame description => fixed to %f s\n", sHeader.dur);
    }
    
    // Dump sound properties.
    fprintf(stdout, "  File type      = %d\n", (int)sHeader.typ);
    fprintf(stdout, "  File length    = %d bytes\n", nFileSize);
    fprintf(stdout, "  Duration       = %f s\n", sHeader.dur);
    fprintf(stdout, "  Sampling rate  = %d Hz\n", (int)sHeader.sr);
    fprintf(stdout, "  Nb partials    = %d\n", (int)sHeader.par);
    fprintf(stdout, "  Max. frequency = %f\n", sHeader.mf);
    fprintf(stdout, "  Max. amplitude = %f\n", sHeader.ma);
    fprintf(stdout, "  Noise ?        = %s\n", FILE_HAS_NOISE(sHeader.typ) ? "true" : "false");
    fprintf(stdout, "  Phase ?        = %s\n", FILE_HAS_PHASE(sHeader.typ) ? "true" : "false");
    fprintf(stdout, "  Nb frames      = %d\n", (int)sHeader.fra);
    fprintf(stdout, "  Frame size     = %d samples\n", (int)sHeader.fs);
    fprintf(stdout, "  Windows size   = %d samples\n", (int)sHeader.ws);

    // Get sound type.
    *sound_type = (int)sHeader.typ;

    // Get length of data from header and compare it to the real file length.
    n1 = 2;
    if (FILE_HAS_PHASE(sHeader.typ))
	n1 += 1;

    n2 = 0;
    if (FILE_HAS_NOISE(sHeader.typ))
	n2 += 1;

    nTheoricalDataSize =
	sizeof(ATS_HEADER)
	+ ((int)sHeader.par * (int)sHeader.fra * sizeof(double) * n1)
	+ (sizeof(double) * (int)sHeader.fra)
	+ (ATSA_CRITICAL_BANDS * (int)sHeader.fra * sizeof(double) * n2);

    if (nTheoricalDataSize != nFileSize)
    {
	fprintf(stderr, "Error: File size does not match header data\n");
	fclose(pATSFile);
	return -1;
    }

    // Reinitialize sound structure from loaded header contents.
    init_sound(sound,
	       (int)sHeader.sr,
	       (int)sHeader.fs,
	       (int)sHeader.ws,
	       (int)sHeader.fra, // WARNING: mem_realloc() used fra+1 !!!!!??????
	       (double)sHeader.dur,
	       (int)sHeader.par,
	       FILE_HAS_NOISE(sHeader.typ));

    // All right till now : load partials and residual ATS data.
    fseek(pATSFile,sizeof(ATS_HEADER),SEEK_SET);

    // For each frame :
    dMaxAmp = 0.0;
    dMaxFreq = 0.0;
    for (nFrameIdx = 0; nFrameIdx < sound->frames; ++nFrameIdx)
    {
	// 1) Load partials.
	fread(&dTmp,1,sizeof(double),pATSFile);  //read just one time value per frame
   
	if (bNeedByteSwap) {
	    sound->time[0][nFrameIdx]=byte_swap(&dTmp); 
	}
	else {
	    sound->time[0][nFrameIdx]=dTmp; 
	}

	// For each partial :
	for (nPartIdx=0; nPartIdx<sound->partials; ++nPartIdx) 
	{
	    // Copy time value from 1st partial.
	    sound->time[nPartIdx][nFrameIdx] = sound->time[0][nFrameIdx];

	    // Read amp value for each partial 
	    fread(&dTmp,1,sizeof(double),pATSFile);
	    if(bNeedByteSwap) {
		sound->amp[nPartIdx][nFrameIdx]=byte_swap(&dTmp); 
	    } 
	    else {
		sound->amp[nPartIdx][nFrameIdx]= dTmp; 
	    }
	    
	    // Update max amplitude if lower (ATSH did not save it !)
	    if (sound->amp[nPartIdx][nFrameIdx] > dMaxAmp)
		dMaxAmp = sound->amp[nPartIdx][nFrameIdx];

	     // Read freq value for each partial 
	    fread(&dTmp,1,sizeof(double),pATSFile);
	    if(bNeedByteSwap) {
		sound->frq[nPartIdx][nFrameIdx]=byte_swap(&dTmp); 
	    } 
	    else {
		sound->frq[nPartIdx][nFrameIdx]= dTmp; 
	    }
   
	    // Update max frequency if lower (ATSH did not save it !)
	    if (sound->frq[nPartIdx][nFrameIdx] > dMaxFreq)
		dMaxFreq = sound->frq[nPartIdx][nFrameIdx];

	    // Read phase data if any...
	    if(FILE_HAS_PHASE(sHeader.typ)) {
		fread(&dTmp,1,sizeof(double),pATSFile); 
		if(bNeedByteSwap) {
		    sound->pha[nPartIdx][nFrameIdx]=byte_swap(&dTmp); 
		} 
		else {
		    sound->pha[nPartIdx][nFrameIdx]= dTmp; 
		}
	    }
	} 
   
	// 2) Load residual if any.
	if (FILE_HAS_NOISE(sHeader.typ)) {
	    int nBandIdx;
	    for (nBandIdx=0; nBandIdx<ATSA_CRITICAL_BANDS; nBandIdx++) {
		fread(&dTmp,1,sizeof(double),pATSFile);
		if(bNeedByteSwap) {
		    sound->band_energy[nBandIdx][nFrameIdx]=byte_swap(&dTmp);  
		}
		else {
		    sound->band_energy[nBandIdx][nFrameIdx]=dTmp;
		}
	    }
	}
    }

    // Close the .ats file.
    fclose(pATSFile);

    // Set sound maxFreq and maxAmpl (as ATSH did not save them in the .ats file).
    sound->ampmax = dMaxAmp;
    sound->frqmax = dMaxFreq;

    // Setup the av array (partial averages and order = increasing frequency,
    // as enforced by ats_save, even in ATSH)
    set_av(sound);

    // Set optimized flag, as not saved in .ats file
    // Note: Set to "true" because ats_save assumes the sound was optimized before. 
    sound->optimized = 1;

    fprintf(stdout, "... done (%d bytes read).\n", nFileSize);

    return nFileSize;
}

/* ats_save
 * ========
 * saves an optimized ATS_SOUND to disk.
 * sound: pointer to ATS_SOUND structure
 * outfile: pointer to output ats file
 * SMR_thres: partials with average SMR under this are not written
 * type: file type
 * Returns the total file size (nb of written bytes), or -1 in case of any error
 */
int ats_save(ATS_SOUND *sound, const char *outfile, double SMR_thres, int type)
{
  int frm, i, par, dead=0;
  double daux;
  ATS_HEADER header;
  size_t written = 0;
  FILE* ofile;

  fprintf(stdout, "Saving ATS sound to '%s' ...\n", outfile);
    
  /* refuse to save if sound was not optimized */
  if( sound->optimized == NIL ){
    fprintf(stderr, "Can not save a sound that has not been optimized before !\n");
    return -1;
  }

  /* open output file */
  ofile = fopen(outfile, "wb");
  if (!ofile) {
    fprintf(stderr, "Could not open %s for writing\n", outfile);
    return -1;
  }

  /* count how many partials are dead */
  for (i = 0; i < sound->partials; i++){
    if (sound->av[i].frq <= 0.0 || sound->av[i].smr < SMR_thres){
      dead++;
    }
  }

  /* fill header up */
  header.mag = ATS_MAGIC_NUMBER;
  header.sr = (double)sound->srate;
  header.fs = (double)sound->frame_size;
  header.ws = (double)sound->window_size;
  header.par = (double)(sound->partials - dead);
  header.fra = (double)sound->frames;
  header.ma = sound->ampmax;
  header.mf = sound->frqmax;
  header.dur = sound->dur;
  header.typ = (double)type;

  fprintf(stdout, "  File type        = %d\n", type);
  fprintf(stdout, "  Duration         = %f s\n", sound->dur);
  fprintf(stdout, "  Sampling rate    = %d Hz\n", sound->srate);
  fprintf(stdout, "  Nb kept partials = %d\n", sound->partials - dead);
  fprintf(stdout, "  Nb dead partials = %d (< %f RMS threshold or freq <= 0)\n", dead, SMR_thres);
  fprintf(stdout, "  Optimized        = %s\n", sound->optimized != NIL ? "true" : "false");
  fprintf(stdout, "  Max. frequency   = %f\n", sound->frqmax);
  fprintf(stdout, "  Max. amplitude   = %f\n", sound->ampmax);
  fprintf(stdout, "  Noise ?          = %s\n", FILE_HAS_NOISE(header.typ) ? "true" : "false");
  fprintf(stdout, "  Phase ?          = %s\n", FILE_HAS_PHASE(header.typ) ? "true" : "false");
  fprintf(stdout, "  Nb frames        = %d\n", sound->frames);
  fprintf(stdout, "  Frame size       = %d samples\n", sound->frame_size);
  fprintf(stdout, "  Windows size     = %d samples\n", sound->window_size);

  /* write header */
  fseek(ofile, 0, SEEK_SET);
  written += sizeof(ATS_HEADER) * fwrite(&header, sizeof(ATS_HEADER), 1, ofile);

  /* write frame data */
  for(frm=0; frm<sound->frames; frm++) {
    daux = sound->time[0][frm];
    written += sizeof(double) * fwrite(&daux, sizeof(double), 1, ofile);
    for(i=0; i<sound->partials; i++) {
      /* we ouput data in increasing frequency order and we check for dead partials */
      if(sound->av[i].frq > 0.0 && sound->av[i].smr >= SMR_thres){
	/* get partial number from sound */
	par = sound->av[i].track;
	/* output data to file */
	daux = sound->amp[par][frm];
	written += sizeof(double) * fwrite(&daux, sizeof(double), 1, ofile);
	daux = sound->frq[par][frm];
	written += sizeof(double) * fwrite(&daux, sizeof(double), 1, ofile);
	if(FILE_HAS_PHASE(header.typ)){
	  daux = sound->pha[par][frm];
	  written += sizeof(double) * fwrite(&daux, sizeof(double), 1, ofile);
	}
      }
    }
    /* write noise data */
    if(FILE_HAS_NOISE(header.typ)){
      for(i=0 ; i<ATSA_CRITICAL_BANDS ; i++){
	daux = sound->band_energy[i][frm];
	written += sizeof(double) * fwrite(&daux, sizeof(double), 1, ofile);
      }
    }
  }

  fclose(ofile);

  fprintf(stdout, "... done (%d bytes written).\n", (int)written);

  return (int)written;
}

/* Usefull functions taken from atsh/atsh-files.c =========================*/
static int my_filelength(FILE *fpt)
{
    int size = 0;

    fseek(fpt, 0, SEEK_END);
    size=(int)ftell(fpt);
    fseek(fpt, 0, SEEK_SET);

    return size;
}

static double byte_swap(double *in)
{
    int size;
    double k;
    unsigned char *sw_in, *sw_out;
    double *out;
    
    size = sizeof(double);
    sw_out = (unsigned char*)malloc(size);
    
    k     = (double)in[0];
    sw_in = (unsigned char *)&k;
    
    sw_out[0] = sw_in[7];
    sw_out[1] = sw_in[6];
    sw_out[2] = sw_in[5];
    sw_out[3] = sw_in[4];
    sw_out[4] = sw_in[3];
    sw_out[5] = sw_in[2];
    sw_out[6] = sw_in[1];
    sw_out[7] = sw_in[0];
    
    out = (double*)sw_out;
    
    return *out;
}

static void byte_swap_header(ATS_HEADER *head, int swapMagic)
{
    double tmp;

    if (swapMagic) 
    {  //may be already swapped
	tmp=head->mag;
	head->mag=byte_swap(&tmp);
    }
			 
    tmp=head->sr;
    head->sr = byte_swap(&tmp);
    tmp=head->fs;  
    head->fs= byte_swap(&tmp);  
    tmp=head->ws;
    head->ws= byte_swap(&tmp);  
    tmp=head->par;
    head->par= byte_swap(&tmp); 
    tmp=head->fra;
    head->fra= byte_swap(&tmp); 
    tmp=head->ma;
    head->ma= byte_swap(&tmp);  
    tmp=head->mf;
    head->mf= byte_swap(&tmp);  
    tmp=head->dur;
    head->dur= byte_swap(&tmp); 
    tmp=head->typ;
    head->typ= byte_swap(&tmp); 
}