/*=========================================================================

  Program:   Insight Segmentation & Registration Toolkit
  Module:    itkMRCImageIO.cxx
  Language:  C++
  Date:      $Date$
  Version:   $Revision$

  Copyright (c) Insight Software Consortium. All rights reserved.
  See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

#include "itkMRCImageIO.h"
#include "itkByteSwapper.h"
#include "itkIntTypes.h"

#include "itkMetaDataObject.h"
#include "itkIOCommon.h"
#include "itkGenericUtilities.h"


#include <numeric>
#include <algorithm>
#include <fstream>

#include <itksys/SystemTools.hxx>

namespace itk
{


const char *MRCImageIO::m_MetaDataHeaderName = "MRCHeader";

MRCImageIO::MRCImageIO()
  : StreamingImageIOBase()
{
  this->SetNumberOfComponents(1);
  this->SetNumberOfDimensions(3);
  this->SetFileTypeToBinary();

  this->AddSupportedReadExtension(".mrc");
  this->AddSupportedReadExtension(".rec");

  this->AddSupportedWriteExtension(".mrc");
  this->AddSupportedWriteExtension(".rec");
}

void MRCImageIO::PrintSelf(std::ostream& os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
}

bool MRCImageIO::CanReadFile(const char* filename)
{
  std::string fname = filename;
  if ( fname != "" &&
       ( fname.find(".mrc") < fname.length() ||
         fname.find(".rec") < fname.length() ) )
    {
    return true;
    }

  std::ifstream file;

  try
    {
    // this may throw an expection, but we just return false
    this->OpenFileForReading(file, filename);
    }
  catch (...)
    {
    return false;
    }

  itkDebugMacro( << "Reading Magic numbers " << filename );
  char map[4];
  char stamp[4];

  // special offset to magic number
  file.seekg(208);

  if( !this->ReadBufferAsBinary( file, static_cast<void*>(&map), 4 ) ||
      !this->ReadBufferAsBinary( file, static_cast<void*>(&stamp), 4 ))
    {
    return false;
    }

  // check the magic number
  if ( strncmp(map, magicMAP, 4) != 0 )
    {
    return false;
    }

  return true;


}

MRCImageIO::SizeType MRCImageIO::GetHeaderSize( void ) const
{
  if ( m_MRCHeader.IsNull() )
    {
    itkExceptionMacro(<< "Must read info first");
    }

  return m_MRCHeader->GetExtendedHeaderSize() + m_MRCHeader->GetHeaderSize();
}

void MRCImageIO::ReadImageInformation( void ) {
  std::ifstream file;


  this->InternalReadImageInformation( file );

  if ( m_MRCHeader->IsOriginalHeaderBigEndian() )
    {
    this->SetByteOrderToBigEndian();
    }
  else
    {
    this->SetByteOrderToLittleEndian();
    }

  const MRCHeaderObject::Header & header = m_MRCHeader->GetHeader();

  // fixed types defined by header
  switch  ( header.mode )
    {
    case MRCHeaderObject::MRCHEADER_MODE_UINT8:
      {
      // todo: the format is unclear weather this is signed or
      // unsigned, it would be best to check the min and max in the
      // header to see what makes since
      this->SetComponentType( UCHAR );
      this->SetNumberOfComponents( 1 );
      this->SetPixelType( SCALAR );
      break;
      }
    case MRCHeaderObject::MRCHEADER_MODE_IN16:
      {
      this->SetComponentType( SHORT );
      this->SetNumberOfComponents( 1 );
      this->SetPixelType( SCALAR );
      break;
      }
    case MRCHeaderObject::MRCHEADER_MODE_FLOAT:
      {
      this->SetComponentType( FLOAT );
      this->SetNumberOfComponents( 1 );
      this->SetPixelType( SCALAR );
      break;
      }
    case MRCHeaderObject::MRCHEADER_MODE_COMPLEX_INT16:
      {
      // ITK does not support short complex well
      // but if the program has gotten this far we can just write it out
      this->SetComponentType( SHORT );
      this->SetNumberOfComponents( 2 );
      this->SetPixelType( COMPLEX );
      break;
      }
    case MRCHeaderObject::MRCHEADER_MODE_COMPLEX_FLOAT:
      {
      this->SetComponentType( FLOAT );
      this->SetNumberOfComponents( 2 );
      this->SetPixelType( COMPLEX );
      break;
      }
    case MRCHeaderObject::MRCHEADER_MODE_UINT16:
      {
      this->SetComponentType( USHORT );
      this->SetNumberOfComponents( 1 );
      this->SetPixelType( SCALAR );
      break;
      }
    case MRCHeaderObject::MRCHEADER_MODE_RGB_BYTE:
      {
      this->SetComponentType( UCHAR );
      this->SetNumberOfComponents( 3 );
      this->SetPixelType( RGB );
      break;
      }
    default:
      {
      itkExceptionMacro(<< "Unrecognized mode");
      }
    }


  if ( header.xlen == 0 &&
       header.ylen == 0 &&
       header.zlen == 0 )
    {
    // if the spacing was not set in the header then this is the
    // default
    m_Spacing[0] = 1.0;
    m_Spacing[1] = 1.0;
    m_Spacing[2] = 1.0;
    }
  else
    {
    m_Spacing[0] = header.xlen / float(header.mx);
    m_Spacing[1] = header.ylen / float(header.my);
    m_Spacing[2] = header.zlen / float(header.mz);
    }

  // copy the origin
  m_Origin[0] = header.xorg;
  m_Origin[1] = header.yorg;
  m_Origin[2] = header.zorg;

  // copy the size of the dimensions
  m_Dimensions[0] = header.nx;
  m_Dimensions[1] = header.ny;
  m_Dimensions[2] = header.nz;

  // add the MRCHeader object into the metadata dictionary to that it
  // can be accessed
  MetaDataDictionary &thisDic=this->GetMetaDataDictionary();


  std::string classname( this->GetNameOfClass() );
  EncapsulateMetaData<std::string>( thisDic, ITK_InputFilterName, classname );
  EncapsulateMetaData<MRCHeaderObject::ConstPointer>( thisDic, m_MetaDataHeaderName, MRCHeaderObject::ConstPointer(m_MRCHeader) );


  return;
}

// methods to load the data into the MRCHeader member variable
void MRCImageIO::InternalReadImageInformation(std::ifstream &file) {
  char *buffer = 0;

  try
    {
    m_MRCHeader = MRCHeaderObject::New();

    itkDebugMacro(<< "Reading Information ");

    this->OpenFileForReading(file, this->m_FileName.c_str());

    buffer = new char[m_MRCHeader->GetHeaderSize()];
    if( !this->ReadBufferAsBinary( file, static_cast<void*>(buffer), m_MRCHeader->GetHeaderSize()) )
      {
      itkExceptionMacro(<<"Header Read failed: Wanted "
                        << m_MRCHeader->GetHeaderSize()
                        << " bytes, but read "
                        << file.gcount() << " bytes.");
      }

    // convert the raw buffer into the header
    if ( !m_MRCHeader->SetHeader(reinterpret_cast<const MRCHeaderObject::Header*>(buffer) ) )
      {
      itkExceptionMacro(<< "Unrecognized header");
      }

    delete [] buffer;
    buffer = 0;

    buffer = new char[m_MRCHeader->GetExtendedHeaderSize()];
    if( !this->ReadBufferAsBinary( file, static_cast<void*>(buffer),  m_MRCHeader->GetExtendedHeaderSize()) )
      {
      itkExceptionMacro(<<"Extended Header Read failed.");
      }


    m_MRCHeader->SetExtendedHeader(buffer);
    }
  catch (...)
    {
    // clean up dynamic allocation
    if (buffer)
      delete [] buffer;
    buffer = 0;
    throw;
    }

  delete [] buffer;
}


void MRCImageIO
::Read(void *buffer)
{
  std::ifstream file;

  if( this->RequestedToStream( ) )
    {
    // open and stream read
    this->OpenFileForReading(file, this->m_FileName.c_str());

    this->StreamReadBufferAsBinary(file, buffer);

    }

  else
    {

    // open the file
    this->OpenFileForReading(file, this->m_FileName.c_str());

    // seek base the header
    std::streampos dataPos = static_cast<std::streampos>( this->GetHeaderSize() );
    file.seekg( dataPos, std::ios::beg );

    if ( file.fail() )
      {
      itkExceptionMacro(<<"Failed seeking to data position");
      }


    // read the image
    this->ReadBufferAsBinary( file, buffer, this->GetImageSizeInBytes() );
    }

  int size = this->GetComponentSize();
  switch( size )
    {
    case 1:
      break;
    case 2:
      this->GetByteOrder() == BigEndian ?
        ByteSwapper<uint16_t>::SwapRangeFromSystemToBigEndian((uint16_t *)buffer, this->GetImageSizeInComponents() ) :
        ByteSwapper<uint16_t>::SwapRangeFromSystemToLittleEndian((uint16_t *)buffer, this->GetImageSizeInComponents() );
      break;
    case 4:
      this->GetByteOrder() == BigEndian ?
        ByteSwapper<uint32_t>::SwapRangeFromSystemToBigEndian((uint32_t *)buffer, this->GetImageSizeInComponents() ) :
        ByteSwapper<uint32_t>::SwapRangeFromSystemToLittleEndian((uint32_t *)buffer, this->GetImageSizeInComponents() );
      break;
    default:
      itkExceptionMacro(<< "Unknown component size");
    }
}


bool MRCImageIO::CanWriteFile(const char* fname)
{
  std::string filename = fname;
  if ( filename.length() > 4 &&
       (filename.find(".mrc") == filename.length() - 4 ||
        filename.find(".rec") == filename.length() - 4 ) )
    {
    return true;
    }
  return false;
}

void MRCImageIO::UpdateHeaderFromImageIO( void )
{

  MRCHeaderObject::Header header;
  memset(&header, 0, sizeof(MRCHeaderObject::Header));

  itkAssertOrThrowMacro ( this->GetNumberOfDimensions() != 0, "Invalid Dimension for Writting" );
  if ( this->GetNumberOfDimensions() > 3 )
    {
    itkExceptionMacro(<<"MRC Writer can not write more than 3-dimensional images");
    }

  // magic number
  strncpy(header.cmap, magicMAP, 4);

  if ( ByteSwapper<void*>::SystemIsBigEndian() )
    header.stamp[0] = 17;
  else
    header.stamp[0] = 68;


  header.alpha = 90;
  header.beta = 90;
  header.gamma = 90;

  header.mapc = MRCHeaderObject::MRCHEADER_MAP_X;
  header.mapr = MRCHeaderObject::MRCHEADER_MAP_Y;
  header.maps = MRCHeaderObject::MRCHEADER_MAP_Z;

  header.mx = header.nx = m_Dimensions[0];
  header.my = header.ny = (  this->GetNumberOfDimensions() >= 2 ) ? m_Dimensions[1] : 1;
  header.mz = header.nz = (  this->GetNumberOfDimensions() >= 3 ) ? m_Dimensions[2] : 1;

  header.mode = -1;
  if (this->GetPixelType() == SCALAR)
    {
    if (this->GetComponentType() == UCHAR)
      {
      header.mode = MRCHeaderObject::MRCHEADER_MODE_UINT8;
      }
    else if (this->GetComponentType() == SHORT)
      {
      header.mode = MRCHeaderObject::MRCHEADER_MODE_IN16;
      }
    else if (this->GetComponentType() == FLOAT)
      {
      header.mode = MRCHeaderObject::MRCHEADER_MODE_FLOAT;
      }
    else if (this->GetComponentType() == USHORT)
      {
      header.mode = MRCHeaderObject::MRCHEADER_MODE_UINT16;
      }
    }
  else if (this->GetPixelType() == COMPLEX)
    {
    // ITK does not support short complex well
    // but if we have gotten this far, it's done
    if (this->GetComponentType() == SHORT)
      {
      header.mode = MRCHeaderObject::MRCHEADER_MODE_COMPLEX_INT16;
      }
    else
      if (this->GetComponentType() == FLOAT)
        {
        header.mode = MRCHeaderObject::MRCHEADER_MODE_COMPLEX_FLOAT;
        }
    }
  else if (this->GetPixelType() == RGB
           && this->GetComponentType() == UCHAR)
    {
    header.mode = MRCHeaderObject::MRCHEADER_MODE_RGB_BYTE;
    }

  if (header.mode == -1)
    {
    itkExceptionMacro(<< "Unsupported pixel type: " << this->GetPixelTypeAsString( this->GetPixelType())
                      << " " << this->GetComponentTypeAsString(this->GetComponentType())
                      << std::endl
                      << "Supported pixel types include unsigned byte, unsigned short, signed short, float, rgb unsigned char, float complex"
      );
    }

  header.nxstart = 0;
  header.nystart = 0;
  header.nzstart = 0;
  
  header.xlen = m_Spacing[0]*float(header.mx);
  header.ylen = ( this->GetNumberOfDimensions() >= 2 ) ? m_Spacing[1]*float(header.my) : 1;
  header.zlen = ( this->GetNumberOfDimensions() >= 3 ) ? m_Spacing[2]*float(header.mz) : 1;
  
  header.xorg = m_Origin[0];
  header.yorg = ( this->GetNumberOfDimensions() >= 2 ) ? m_Origin[1] : 0;
  header.zorg = ( this->GetNumberOfDimensions() >= 3 ) ? m_Origin[2] : 0;

  // the SetHeader method is used to set the all the internal variable
  // of the header object correctly, and the data is verified
  m_MRCHeader = MRCHeaderObject::New();
  if (!m_MRCHeader->SetHeader(&header))
    {
    itkExceptionMacro(<< "Unexpected error setting header");
    }
}

void MRCImageIO::WriteImageInformation( const void * buffer )
{

  this->UpdateHeaderFromImageIO();

  this->UpdateHeaderWithMinMaxMean( buffer );

  std::ofstream file;

  this->OpenFileForWriting(file, this->m_FileName.c_str(), true);

  // write the header
  const MRCHeaderObject::Header & header = m_MRCHeader->GetHeader();
  file.write(static_cast<const char*>((void*)&(header)), 1024);
}

void MRCImageIO
::UpdateHeaderWithMinMaxMean( const void * bufferBegin )
{

  // fixed types defined by header
  const MRCHeaderObject::Header & header = m_MRCHeader->GetHeader();

  switch ( header.mode )
    {
    case 0:
      {
      // scalar unsigned char
      this->UpdateHeaderWithMinMaxMean(static_cast<const unsigned char*>(bufferBegin));
      break;
      }
    case 1:
      {
      // scalar short
      this->UpdateHeaderWithMinMaxMean(static_cast<const short*>(bufferBegin));
      break;
      }
    case 2:
      {
      // scalar float
      this->UpdateHeaderWithMinMaxMean(static_cast<const float*>(bufferBegin));
      break;
      }
    case 3:
      {
      // complex short

      // What is the best way to map complex to float?
      // just set resonable values
      m_MRCHeader->m_Header.amin = -1.0f;
      m_MRCHeader->m_Header.amax = 1.0f;
      m_MRCHeader->m_Header.amean = 0.0f;
      break;
      }
    case 4:
      {
      // complex float

      // What is the best way to map complex to float?
      // just set resonable values
      m_MRCHeader->m_Header.amin = -1.0f;
      m_MRCHeader->m_Header.amax = 1.0f;
      m_MRCHeader->m_Header.amean = 0.0f;
      break;
      }
    case 6:
      {
      // scalar unsigned short
      this->UpdateHeaderWithMinMaxMean(static_cast<const unsigned short*>(bufferBegin));
      break;
      }
    case 16:
      {
      // RGB of unsigned char

      // just set resonable values
      m_MRCHeader->m_Header.amin = 0.0f;
      m_MRCHeader->m_Header.amax = 255.0f;
      m_MRCHeader->m_Header.amean = 127.5f;
      break;
      }
    default:
      {
      itkExceptionMacro(<< "Unrecognized mode");
      }
    }
}

void MRCImageIO
::Write(const void* buffer)
{

  if( this->RequestedToStream() )
    {


    // we assume that GetActualNumberOfSplitsForWriting is called before
    // this methods and it will remove the file if a new header needs to
    // be written
    if (!itksys::SystemTools::FileExists( m_FileName.c_str() ))
      {
      this->WriteImageInformation( buffer );


      std::ofstream file;
      // open and stream write
      this->OpenFileForWriting(file, this->m_FileName.c_str(), false);

      // write one byte at the end of the file to allocate (this is a
      // nifty trick which should not write the entire size of the file
      // just allocate it, if the system supports sparse files)
      std::streampos seekPos = this->GetImageSizeInBytes() + this->GetHeaderSize() - 1;
      file.seekp( seekPos, std::ios::cur );
      file.write("\0", 1);
      file.seekp( 0 );

      }
    else
      {

      if ( m_MRCHeader.IsNull() )
        {
        // need to determin the size of the header in the file by
        // reading the header into m_MRCHeader

        // we assume that GetActualNumberOfSplitsForWriting is called
        // to verify that the header and file is compatible with the
        // region to be written, but this call will still overwrite
        // the internal m_MRCHeader variable

        std::ifstream file;
        this->InternalReadImageInformation( file );

        }

      // TODO:
      // how are we suppose to update min, max mean? we could be
      // overwriting data, and we don't know what's already there if
      // we were pasting
      }

    std::ofstream file;
    // open and stream write
    this->OpenFileForWriting(file, this->m_FileName.c_str(), false);

    this->StreamWriteBufferAsBinary(file, buffer);

    }

  else
    {

    // this will truncate file and write header
    this->WriteImageInformation( buffer );

    std::ofstream file;
    // open the file
    this->OpenFileForWriting(file, this->m_FileName.c_str(), false);

    // seek pass the header
    std::streampos dataPos = static_cast<std::streampos>( this->GetHeaderSize() );
    file.seekp( dataPos, std::ios::beg );

    if ( file.fail() )
      {
      itkExceptionMacro(<<"Failed seeking to data position");
      }

    // read the image
    if (!this->WriteBufferAsBinary( file, buffer, this->GetImageSizeInBytes() ))
      {
      itkExceptionMacro(<< "Could not write file: " << m_FileName);
      }
    }

}

} // namespace itk