psrfitsio.h

//-*-C++-*-
/***************************************************************************
 *
 *   Copyright (C) 2006 - 2025 by Willem van Straten
 *   Licensed under the Academic Free License version 2.1
 *
 ***************************************************************************/
 
// psrchive/Util/fitsutil/psrfitsio.h
 
#ifndef __psrfitsio_h
#define __psrfitsio_h
 
#include "FITSError.h"
#include "ReferenceTo.h"
#include "fitsutil.h"
 
#include <fitsio.h>
 
#include <string>
#include <vector>
#include <complex>
 
extern bool psrfits_verbose;

template<typename T> struct FITS_traits { };
  
template<> struct FITS_traits<double>
{
  static inline int datatype() { return TDOUBLE; }
  static inline double null () { return 0; }
};

template<> struct FITS_traits<float>
{
  static inline int datatype() { return TFLOAT; }
  static inline float null () { return fits_nullfloat; }
};

template<> struct FITS_traits<short>
{
  static inline int datatype() { return TSHORT; }
  static inline short null () { return -1; }
};

template<> struct FITS_traits<int>
{
  static inline int datatype() { return TINT; }
  static inline int null () { return -1; }
};

template<> struct FITS_traits<long> {
  static inline int datatype() { return TLONG; }
  static inline long null () { return -1; }
};

template<> struct FITS_traits<long long> {
  static inline int datatype() { return TLONGLONG; }
  static inline long null () { return -1; }
};

template<> struct FITS_traits<unsigned> {
  static inline int datatype() { return TUINT; }
  static inline long null () { return 0; }
};

template<> struct FITS_traits<std::string>
{
  static inline int datatype() { return TSTRING; }
  static inline std::string null () { return ""; }
};

template<> struct FITS_traits<std::complex<float>>
{
  static inline int datatype() { return TCOMPLEX; }
  static inline float null () { return 0; }
};

template<> struct FITS_traits<std::complex<double>>
{
  static inline int datatype() { return TDBLCOMPLEX; }
  static inline double null () { return 0; }
};
 
template<typename T>
void* FITS_void_ptr (const T& ptr)
{
  return const_cast<T*> (&ptr);
}
 
void* FITS_void_ptr (const std::string&);

template<typename T>
void psrfits_update_key (fitsfile* fptr, const char* name, T data, const char* comment = 0)
{
  // status
  int status = 0;
  
  fits_update_key (fptr, FITS_traits<T>::datatype(),
                   const_cast<char*>(name), &data,
                   const_cast<char*>(comment), &status);
  
  if (status)
    throw FITSError (status, "psrfits_update_key", name);
}

void psrfits_update_key (fitsfile* fptr, const char* name,
                         const char* data,
                         const char* comment = 0);

void psrfits_update_key (fitsfile* fptr, const char* name,
                         const std::string& data,
                         const char* comment = 0);

void psrfits_update_key (fitsfile* fptr,
                         const char* name,
                         int column,
                         const std::string& value,
                         const char* comment = 0);

void psrfits_read_key (fitsfile* fptr,
                       const char* name,
                       int column,
                       std::string& value,
                       std::string& comment);

template<typename T>
void psrfits_read_key_work (fitsfile* fptr, const char* name, T* data, int* status)
{
  // no comment
  char* comment = 0;
  
  fits_read_key (fptr, FITS_traits<T>::datatype(), 
                 const_cast<char*>(name), data,
                 comment, status);
}

void psrfits_read_key_work (fitsfile* fptr, const char* name, std::string*,
                            int* status);

template<typename T>
void psrfits_read_key (fitsfile* fptr, const char* name, T* data)
{
  // status
  int status = 0;
  psrfits_read_key_work (fptr, name, data, &status);
  if (status)
    throw FITSError (status, "psrfits_read_key", name);
}


template<typename Class, typename T>
void psrfits_read_key (fitsfile* fptr, const char* name, Class* instance, void(Class::*set)(T))
{
  T val;
  psrfits_read_key (fptr, name, &val);
  (instance->*set)(val);
}


template<typename Class, typename T>
void psrfits_read_key (fitsfile* fptr, const char* name, Reference::To<Class>& instance, void(Class::*set)(T))
{
  T val;
  psrfits_read_key (fptr, name, &val);
  (instance->*set)(val);
}

template<typename T>
void psrfits_read_key (fitsfile* fptr, const char* name, T* data, T dfault, bool verbose = false)
{
  // status
  int status = 0;
  psrfits_read_key_work (fptr, name, data, &status);
  if (status)
  {
    if (verbose)
    {
      FITSError error (status, "psrfits_read_key", name);
      std::cerr << error.get_message() << std::endl;
      std::cerr << "psrfits_read_key: using default=" << dfault << std::endl;
    }
    *data = dfault;
  }
  else if (verbose)
    std::cerr << "psrfits_read_key: " << name << "=" << *data << std::endl;
}

void psrfits_update_tdim (fitsfile* ffptr, int column, unsigned dim);

void psrfits_update_tdim (fitsfile* ffptr, int column,
                          unsigned dim1, unsigned dim2);

void psrfits_update_tdim (fitsfile* ffptr, int column,
                          unsigned dim1, unsigned dim2, unsigned dim3);

void psrfits_update_tdim (fitsfile* ffptr, int column,
                          unsigned dim1, unsigned dim2,
                          unsigned dim3, unsigned dim4);

void psrfits_update_tdim (fitsfile* ffptr, int column,
                          const std::vector<unsigned>& dims);
 
template<typename T>
void psrfits_write_col (fitsfile* fptr, const char* name, int row,
                        const std::vector<T>& data,
                        const std::vector<unsigned>& dims)
{
  //
  // Get the number of the named column
 
  int colnum = 0;
  int status = 0;
 
  if (psrfits_verbose)
    std::cerr << "psrfits_write_col calling fits_get_colnum"
      " name='" << name << "'" << std::endl;
  
  fits_get_colnum (fptr, CASEINSEN, const_cast<char*>(name), &colnum, &status);
 
  if (status)
    throw FITSError (status, "psrfits_write_col(vector<T>)",
                     "fits_get_colnum (name=%s)", name);
 
  psrfits_write_col (fptr, colnum, row, data, dims);
}
 
/* when it is not feasible to write a single block of data in one go */
template<typename Stream>
void psrfits_write_col (fitsfile* fptr, const char* name, int row,
                        const Stream* stream)
{
  // 
  // Get the number of the named column
 
  int colnum = 0;
  int status = 0;
 
  if (psrfits_verbose)
    std::cerr << "psrfits_write_col calling fits_get_colnum"
      " name='" << name << "'" << std::endl;
 
  fits_get_colnum (fptr, CASEINSEN, const_cast<char*>(name), &colnum, &status);
 
  if (status)
    throw FITSError (status, "psrfits_write_col(vector<T>)",
                     "fits_get_colnum (name=%s)", name);
 
  try
  {
    psrfits_write_col (fptr, colnum, row, stream);
  }
  catch (Error& error)
  {
    throw error += "psrfits_write_col(vector<T>," + std::string(name) + ")";
  }
}
 
template<typename T>
void psrfits_write_col (fitsfile* fptr, int colnum, int row,
                        const std::vector<T>& data,
                        const std::vector<unsigned>& dims)
{
  int status = 0;
 
  if (psrfits_verbose)
    std::cerr << "psrfits_write_col calling fits_modify_vector_len"
      " colnum=" << colnum << " size=" << data.size() << std::endl;
  
  fits_modify_vector_len (fptr, colnum, data.size(), &status);
 
  if (status)
    throw FITSError (status, "psrfits_write_col(vector<T>)",
                     "fits_modify_vector_len (col=%d size=%u)",
                     colnum, data.size());
 
  if (dims.size() > 1)
  {
    if (psrfits_verbose)
      std::cerr << "psrfits_write_col calling psrfits_update_tdim" << std::endl;
 
    psrfits_update_tdim (fptr, colnum, dims);
  }
 
  if (psrfits_verbose)
    std::cerr << "psrfits_write_col calling fits_write_col" << std::endl;
  
  fits_write_col (fptr, FITS_traits<T>::datatype(),
                  colnum, row,
                  1, data.size(),
                  const_cast<T*>(&(data[0])), &status);
 
  if (status)
    throw FITSError (status, "psrfits_write_col(vector<T>)",
                     "fits_write_col (type=%s col=%d row=%d size=%u)",
                     fits_datatype_str(FITS_traits<T>::datatype()),
                     colnum, row, data.size());
}
 
template<typename Stream>
void psrfits_write_col (fitsfile* fptr, int colnum, int row,
                        const Stream* stream)
{
  int status = 0;
 
  unsigned ndat = stream->get_ndat();
 
  if (psrfits_verbose)
    std::cerr << "psrfits_write_col (Stream) calling fits_modify_vector_len"
      " colnum=" << colnum << " size=" << ndat << std::endl;
 
  fits_modify_vector_len (fptr, colnum, ndat, &status);
 
  if (status)
    throw FITSError (status, "psrfits_write_col(vector<T>)",
                     "fits_modify_vector_len (col=%d size=%u)",
                     colnum, ndat);
 
  unsigned nwritten = 0;
 
  unsigned buffer_size = 1024 * 1024;
 
  std::vector< typename Stream::Type > tmp ( buffer_size );
  while (nwritten < ndat)
  {
    if (nwritten + buffer_size > ndat)
      buffer_size = ndat - nwritten;
 
    stream->get_data (nwritten, buffer_size, &tmp[0] );
 
    fits_write_col (fptr, FITS_traits< typename Stream::Type >::datatype(),
                    colnum, row,
                    nwritten + 1, buffer_size,
                    &(tmp[0]), &status);
 
    if (status)
      throw FITSError (status, "psrfits_write_col(Stream)",
                     "fits_write_col (type=%s col=%d row=%d size=%u)",
                     fits_datatype_str(FITS_traits< typename Stream::Type >::datatype()),
                     colnum, row, buffer_size);
 
    nwritten += buffer_size;
  }
}
 
template<typename T>
void psrfits_write_col( fitsfile *fptr, const char *name, int row,
                        const T& data )
{
  int colnum = 0;
  int status = 0;
  
  fits_get_colnum (fptr, CASEINSEN, const_cast<char*>(name), &colnum, &status);
  
  if (status)
    throw FITSError (status, "psrfits_write_col(T)",
                     "fits_get_colnum (name=%s)", name);
 
  try {
    psrfits_write_col (fptr, colnum, row, data);
  }
  catch (Error& error)
  {
    throw error += "psrfits_write_col(T," + std::string(name) + ")";
  }
}
 
template<typename T>
void psrfits_write_col( fitsfile *fptr, int colnum, int row, const T& data )
{
  int status = 0;
 
  fits_write_col (fptr, FITS_traits<T>::datatype(),
                  colnum, row,
                  1, 1,
                  FITS_void_ptr(data), &status);
 
  if (status)
    throw FITSError (status, "psrfits_write_col(T)",
                     "col=%d row=%d", colnum, row);
}
 
template<typename T>
void psrfits_read_col (fitsfile* fptr, const char* name,
                       std::vector< std::vector<T> >& data,
                       int row = 1, T null = FITS_traits<T>::null())
{
  //
  // Get the number of the named column
 
  int colnum = 0;
  int status = 0;
 
  fits_get_colnum (fptr, CASEINSEN, const_cast<char*>(name), &colnum, &status);
 
  int anynul = 0;
  int counter = 1;
 
  for (unsigned i = 0; i < data.size(); i++)
  {
    fits_read_col (fptr, FITS_traits<T>::datatype(),
                   colnum, row, 
                   counter, data[i].size(),
                   &null, &(data[i][0]),
                   &anynul, &status);
 
    if (status != 0)
      throw FITSError( status, "psrfits_read_col",
                       "%s colnum=%d firstrow=%d firstelem=%d nelements=%d",
                       name, colnum, row, counter, data[i].size() );
      
    counter += data[i].size();
  }
 
}
 
template<typename T>
void psrfits_read_col (fitsfile* fptr, const char* name, std::vector<T>& data,
                       int row = 1, T null = FITS_traits<T>::null())
{
  //
  // Get the number of the named column
 
  int colnum = 0;
  int status = 0;
 
  fits_get_colnum (fptr, CASEINSEN, const_cast<char*>(name), &colnum, &status);
 
  //
  // If the vector length is unspecified, read it from the file
 
  if (data.size() == 0)
  {
    int typecode = 0;
    long repeat = 0;
    long width = 0;
  
    fits_get_coltype (fptr, colnum, &typecode, &repeat, &width, &status);  
    if (status)
      throw FITSError (status, "psrfits_read_col",
                       "fits_get_coltype (%s)", name);
 
    data.resize( repeat );
  }
 
  int anynul = 0;
  fits_read_col (fptr, FITS_traits<T>::datatype(),
                 colnum, row,
                 1, data.size(),
                 &null, &(data[0]),
                 &anynul, &status);
 
  if (status)
    throw FITSError (status, "psrfits_read_col", name);
}

template< typename T >
void psrfits_read_col_work( fitsfile *fptr, const char *name, T *data,
                            int row, T null, int* status)
{ 
  int colnum = 0;
  fits_get_colnum (fptr, CASEINSEN, const_cast<char*>(name), &colnum, status);
  
  int anynul = 0;
  fits_read_col( fptr, FITS_traits<T>::datatype(),
                 colnum, row,
                 1, 1, &null, data, 
                 &anynul, status );
}

void psrfits_read_col_work( fitsfile *fptr, const char *name,
                            std::string *data,
                            int row, std::string& null, int* status);
 
 
template< typename T >
void psrfits_read_col( fitsfile *fptr, const char *name, T *data,
                       int row = 1, T null = FITS_traits<T>::null() )
{
  int status = 0;
  psrfits_read_col_work (fptr, name, data, row, null, &status);
  if (status)
    throw FITSError (status, "psrfits_read_call", name);
}

template< typename T >
void psrfits_read_col( fitsfile *fptr, const char *name, T *data,
                       int row, T null, T dfault, bool verbose)
{
  // status
  int status = 0;
  psrfits_read_col_work (fptr, name, data, row, null, &status);
  if (status)
  {
    if (verbose)
    {
      FITSError error (status, "psrfits_read_col", name);
      std::cerr << error.get_message() << std::endl;
      std::cerr << "psrfits_read_col: using default=" << dfault << std::endl;
    }
    *data = dfault;
  }
  else if (verbose)
    std::cerr << "psrfits_read_col: " << name << "=" << *data << std::endl;
 
}

void psrfits_init_hdu (fitsfile *fptr, const char *name, unsigned nrows=1);

bool psrfits_move_hdu (fitsfile *fptr, const char *name,
                       bool optional = false,
                       int table_type = BINARY_TBL, int version = 0);

void psrfits_clean_rows (fitsfile*);

void psrfits_set_rows (fitsfile*, unsigned nrow);

void psrfits_insert_row (fitsfile* fptr);

void psrfits_delete_col (fitsfile* fptr, const char* name);
 
#endif