#define _GNU_SOURCE
#include "assign_protein_type.h"
#include "error/check_error.h"
#include "error/check_h5_error.h"
#include "error/check_ncbi_error.h"
#include "model/gi_type_data.h"
#include "model/gi_type_data_init.h"
#include "model/sequence_data.h"
#include "model/sequence_data_init.h"
#include <blast.h>
#include <error.h>
#include <hdf5_hl.h>
#include <ncbi.h>
#include <readdb.h>
#include <salpacc.h>
#include <search.h>
#include <stdbool.h>

/*
 * BLAST database containing all of the influenza protein sequences.
 */
#define SEQDB "/home/don/exp004/influenzadb/influenzadb"

/*
 * BLAST reference database of prototypical protein types.
 */
#define REFDB "/home/don/exp004/influenzadb/proteinnames"

#define BUFFER_LEN 50

void
assign_protein_type (hid_t file_id)
{
  /*
   * Iterate through the records for which no protein type has been
   * assigned.  Create a BioSeq Pointer to the data and then use this
   * as input to the BLAST run against the reference database of
   * prototypical sequence records.
   */

  /*
   * Make BLAST messages reportable.
   */
  ErrSetMessageLevel (SEV_WARNING);

  /*
   * Open the BLAST sequence database.
   */
  ReadDBFILEPtr seqdb = readdb_new (SEQDB, true);

  /*
   * Get default BLAST options.
   */
  BLAST_OptionsBlkPtr options = BLASTOptionNew ("blastp", true);

  /*
   * Set to single hit mode.
   */
  options->window_size = 0;
  options->multiple_hits_only = false;

  ValNodePtr error_returns = NULL;

  /*
   * Read the data from HDF5 influenza.faa.
   */
  hsize_t faa_nfields;
  hsize_t faa_nrecords;
  herr_t status = H5TBget_table_info (file_id, "influenza.faa", &faa_nfields,
				      &faa_nrecords);
  if (status < 0)
    check_h5_error (status, __FILE__, __LINE__);

  sequence_data* faa_buf = malloc (sizeof(sequence_data) * faa_nrecords);

  size_t faa_size;
  size_t faa_offset[SEQUENCE_DATA_FIELD_NUM];
  size_t faa_sizes[SEQUENCE_DATA_FIELD_NUM];
  hid_t faa_field_type[SEQUENCE_DATA_FIELD_NUM];
  sequence_data_init (&faa_size, faa_offset, faa_sizes, faa_field_type);

  status = H5TBread_table (file_id, "influenza.faa", faa_size, faa_offset,
			   faa_sizes, faa_buf);
  if (status < 0)
    check_h5_error (status, __FILE__, __LINE__);

  /*
   * Allocate memory for the new table.
   */
  gi_type_data* new_buf = malloc (sizeof (gi_type_data) * faa_nrecords);
  if (new_buf == NULL)
    check_error (__FILE__, __LINE__);

  /*
   * Read the data from HDF5 gi_type_data.
   */
  hsize_t gi_nfields = 0;
  hsize_t gi_nrecords = 0;
  size_t gi_size = 0;
  size_t gi_offset[GI_TYPE_DATA_FIELD_NUM];
  size_t gi_sizes[GI_TYPE_DATA_FIELD_NUM];
  hid_t gi_field_type[GI_TYPE_DATA_FIELD_NUM];
  gi_type_data_init (&gi_size, gi_offset, gi_sizes, gi_field_type);

  gi_type_data* old_buf = NULL;

  /*
   * If the table is already present read the values into memory and
   * then clear the table.
   */
  if (H5LTfind_dataset (file_id, "gi_type_data") == 1)
    {

      printf ("Updating gi_type_data.\n");

      status = H5TBget_table_info (file_id, "gi_type_data", &gi_nfields,
				   &gi_nrecords);
      if (status < 0)
	check_h5_error (status, __FILE__, __LINE__);

      printf ("  Using gi_type_data cache of %i records.\n", (int)gi_nrecords);

      old_buf = malloc (sizeof(gi_type_data) * gi_nrecords);

      status = H5TBread_table (file_id, "gi_type_data", gi_size, gi_offset,
			       gi_sizes, old_buf);
      if (status < 0)
	check_h5_error (status, __FILE__, __LINE__);

      status = H5TBdelete_record (file_id, "gi_type_data", 0, gi_nrecords);
      if (status < 0)
	check_h5_error (status, __FILE__, __LINE__);

    }

  /*
   * If the table is not already present create it.
   */
  else
    {

      printf ("Creating gi_type_data.\n");

      const char* gi_type_data_field_names[GI_TYPE_DATA_FIELD_NUM] =
	GI_TYPE_DATA_FIELD_NAMES;

      hsize_t chunk_size = 10;
      int *fill_data = NULL;
      int compress = 0;

      status = H5TBmake_table ("gi_type_data", file_id,
			       "gi_type_data",
			       GI_TYPE_DATA_FIELD_NUM, 0,
			       gi_size, gi_type_data_field_names,
			       gi_offset, gi_field_type,
			       chunk_size, fill_data, compress,
			       NULL);
      if (status < 0)
	check_h5_error (status, __FILE__, __LINE__);

    }

  /*
   * Copy the contents of the old table into a hash.
   */
  struct hsearch_data htab;
  bzero (&htab, sizeof (htab));
  if (hcreate_r (gi_nrecords * 2, &htab) == 0)
    error_at_line (EXIT_FAILURE, 0, __FILE__, __LINE__,
		   "Allocation of cache failed.");
  ENTRY e, *ep;

  for (int i = 0; i < (int)gi_nrecords; i++)
    {
      char gi_chr[25];
      snprintf (gi_chr, 25, "%i", old_buf[i].gi);
      e.key = strdup (gi_chr);
      e.data = &old_buf[i];
      if (hsearch_r (e, ENTER, &ep, &htab) == 0)
	error_at_line (EXIT_FAILURE, 0, __FILE__, __LINE__,
		       "Allocation failed.");
    }

  /*
   * Assign protein types to records for which the field is empty.
   */
  printf ("Records to process: %i\n", (int)faa_nrecords);
  int written = 0;
  for (int i = 0; i < (int)faa_nrecords; i++)
    {
      new_buf[i].gi = faa_buf[i].gi;
      strncpy (new_buf[i].type, "", sizeof (new_buf[i].type));
      strncpy (new_buf[i].protein, "", sizeof (new_buf[i].protein));

      char gi_chr[25];
      snprintf (gi_chr, 25, "%i", faa_buf[i].gi);
      e.key = gi_chr;
      e.data = NULL;

      /*
       * A record was not found in the cache for this gi.
       */
      if (hsearch_r (e, FIND, &ep, &htab) == 0)
	{

	  /*
	   * Read the sequence from the database by GI.
	   */
	  Int4 sequence_number = readdb_gi2seq (seqdb, faa_buf[i].gi, NULL);
	  BioseqPtr bsp = readdb_get_bioseq (seqdb, sequence_number);
	  if (bsp == NULL)
	    {
	      error_at_line (EXIT_FAILURE, 0, __FILE__, __LINE__,
			     "Unable to find BLAST record for gi|%i.  Ensure "
			     "the BLAST database is up-to-date with the HDF5 "
			     "record set.  See the BLAST formatdb.log file "
			     "for details.\n",
			     faa_buf[i].gi);
	    }

	  SeqAlignPtr seqalign = BioseqBlastEngine (bsp,
						    "blastp",
						    REFDB,
						    options,
						    NULL,
						    &error_returns,
						    NULL);

	  /*
	   * BLAST reported an error.  Write it out and continue processing.
	   */
	  if (error_returns != NULL)
	    {
	      CharPtr msg = BlastErrorToString (error_returns);
	      printf ("Warning: An error has been reported by the NCBI Toolkit "
		      "API for sequence gi|%i: %s",
		      faa_buf[i].gi, msg);
	      free (msg);
	    }

	  /*
	   * A hit was found.  Record the first hit as the protein type.
	   * Skip the first 6 characters and eat the "lcl|x_".
	   */
	  else if (seqalign != NULL)
	    {
	      Char target_id_buf[BUFFER_LEN + 1];
	      SeqIdPtr target_id = SeqAlignId (seqalign, 1);
	      SeqIdWrite (target_id, target_id_buf, PRINTID_FASTA_SHORT,
			  BUFFER_LEN);

	      // Species Type
	      new_buf[i].type[0] = target_id_buf[4];
	      new_buf[i].type[1] = '\0';

	      // Protein Type
	      strncpy (new_buf[i].protein, &target_id_buf[6],
		       sizeof (new_buf[i].protein));
	    }

	  /*
	   * BLAST did not find any hits.
	   */
	  else
	    {
	      printf ("Warning: Unable to identify protein type for sequence "
		      "gi|%i\n", faa_buf[i].gi);
	    }

	  /*
	   * Clean up memory for the next ieration.
	   */
	  seqalign = SeqAlignSetFree (seqalign);
	  bsp = BioseqFree (bsp);

	} // End existing entry not found.

      /*
       * Hash table entry found.  Keep the old value.
       */
      else
	{
	  gi_type_data* old_value = (gi_type_data*)ep->data;
	  new_buf[i].gi = old_value->gi;
	  strncpy (new_buf[i].type, old_value->type, sizeof (new_buf[i].type));
	  strncpy (new_buf[i].protein, old_value->protein, sizeof (new_buf[i].protein));
	}

      /*
       * Write the data out to the file.
       */
      if ( (i % 1000 == 0) && (i > 0) )
	{
	  status = H5TBappend_records (file_id, "gi_type_data", 1000,
				      gi_size, gi_offset, gi_sizes,
				      &new_buf[i-1000]);
	  if (status < 0)
	    check_h5_error (status, __FILE__, __LINE__);

	  status = H5Fflush (file_id, H5F_SCOPE_GLOBAL);
	  if (status < 0)
	    check_h5_error (status, __FILE__, __LINE__);

	  written = i;

	  printf ("Processed %i of %i records.\n", i, (int)faa_nrecords);
	}

    }

  /*
   * Write out records from the last bin if it was less than 1000
   * records in size.
   */
  if ((int)faa_nrecords < 1000)
    {
      status = H5TBappend_records (file_id, "gi_type_data", faa_nrecords,
				  gi_size, gi_offset, gi_sizes,
				  new_buf);
    }

  else
    {
      status = H5TBappend_records (file_id, "gi_type_data", faa_nrecords - written,
				  gi_size, gi_offset, gi_sizes,
				   &new_buf[written]);
    }

  if (status < 0)
    check_h5_error (status, __FILE__, __LINE__);

  status = H5Fflush (file_id, H5F_SCOPE_GLOBAL);
  if (status < 0)
    check_h5_error (status, __FILE__, __LINE__);

  free (faa_buf);
  free (old_buf);
  free (new_buf);
  hdestroy_r (&htab);

  options = BLASTOptionDelete (options);
  readdb_destruct (seqdb);

  return;
}