/// ImageView.cpp
/**
*/

#if defined(__WXGTK__) || defined(__WXMOTIF__)
  #include <wx/wx.h>
#endif

#include <vector>
#include <algorithm>
#include <limits>

#include <wx/wxprec.h>
#include <wx/image.h>
#include <wx/filename.h>

#include <wx/event.h>
#include <wx/window.h>
#include <wx/slider.h>
#include <wx/dialog.h>
#include <wx/sizer.h>
#include <wx/dcbuffer.h>
#include <wx/log.h>
#include <wx/progdlg.h>

#include "DicomFile.h"
#include "ImageView.h"
#include "StringConvert.h"

using namespace jcs;

wxImage CreateImage(DicomFile& file);

enum {
  BUTTON_PLAY,
  TIMER_ID,
  SAVE_IMAGE,
  SAVE_ALL,
  WINDOW,
  RESET,
  CLOSE
};

bool useDefaultWC = true;
int width = 0;
int center = 0;

BEGIN_EVENT_TABLE(BitmapWindow, wxWindow)

EVT_PAINT			(BitmapWindow::OnPaint)
EVT_ERASE_BACKGROUND(BitmapWindow::OnErase)
EVT_LEFT_DOWN       (BitmapWindow::OnMouseDown)
EVT_MOTION          (BitmapWindow::OnMouseMotion)
EVT_LEFT_DCLICK     (BitmapWindow::OnDoubleClick)

END_EVENT_TABLE()

///
/**
 */
BitmapWindow::BitmapWindow(ImageViewerBase* parent, const wxBitmap& bitmap)
: wxWindow(parent, -1, wxDefaultPosition, wxDefaultSize), mBitmap(bitmap), mPoint(wxDefaultPosition) 
{
  SetClientSize(mBitmap.GetWidth(), mBitmap.GetHeight());
  myParent = parent;
}


///
/**
 */
void
BitmapWindow::SetBitmap(const wxBitmap& bitmap)
{
  mBitmap = bitmap;
  Refresh();
}


///
/**
 */
void
BitmapWindow::OnPaint(wxPaintEvent& event)
{
  wxPaintDC dc(this);
  dc.DrawBitmap(mBitmap, 0, 0, FALSE );
}


///
/**
 */
void
BitmapWindow::OnErase(wxEraseEvent& event)
{
}


///
/**
 */
void
BitmapWindow::OnMouseDown(wxMouseEvent& event)
{
  mPoint = event.GetPosition();
  useDefaultWC = false;
}


///
/**
 */
void
BitmapWindow::OnMouseMotion(wxMouseEvent& event)
{
  if (event.Dragging() && event.LeftIsDown()) {
    width += event.GetPosition().x - mPoint.x;
    center += event.GetPosition().y - mPoint.y;
    if (width < 0) { width = 0; }
    if (center < 0) { center = 0; }
    mPoint.x = event.GetPosition().x;
    mPoint.y = event.GetPosition().y;
    myParent->UpdateImage();
  }
}


///
/**
 */
void
BitmapWindow::OnDoubleClick(wxMouseEvent& event)
{
  useDefaultWC = true;
  myParent->UpdateImage();
}


BEGIN_EVENT_TABLE(ImageViewerBase, wxFrame)  //Dialog)

EVT_SCROLL		    (             ImageViewerBase::OnSlide)
EVT_BUTTON		    (BUTTON_PLAY, ImageViewerBase::OnPlay)
EVT_TIMER			    (TIMER_ID,    ImageViewerBase::OnTimer)
EVT_MENU			    (SAVE_IMAGE,  ImageViewerBase::OnSave)
EVT_MENU			    (SAVE_ALL,	  ImageViewerBase::OnSaveAll)
EVT_MENU          (WINDOW,      ImageViewerBase::OnWindow)
EVT_MENU          (RESET,       ImageViewerBase::OnReset)
EVT_MENU          (CLOSE,       ImageViewerBase::OnClose)
EVT_CONTEXT_MENU	(             ImageViewerBase::OnContextMenu)

END_EVENT_TABLE()


///
/**
 */
ImageViewerBase::ImageViewerBase(wxWindow* parent, const wxString& title)
: wxFrame(parent, -1, title, wxDefaultPosition, wxDefaultSize), mTimer(this, TIMER_ID)
//: wxDialog(parent, -1, title, wxDefaultPosition, wxDefaultSize), mTimer(this, TIMER_ID)
{
  useDefaultWC = true;
  //	CreateStatusBar();

  /*	wxImage image = CreateInitialImage();

  int size_x = image.GetWidth();
  int size_y = image.GetHeight();
  int size_z = GetNumberOfImages();
*/	
  wxBoxSizer* panelSizer;
  wxBoxSizer* sliderSizer;
  mpSliceSlider = new wxSlider(this, -1, 0 , 0, 100);
  panelSizer = new wxBoxSizer(wxVERTICAL);
  sliderSizer = new wxBoxSizer(wxHORIZONTAL);
  sliderSizer->Add(mpSliceSlider, 1);

  mpPlayButton = new wxButton(this, BUTTON_PLAY, _T("play"));
  sliderSizer->Add(mpPlayButton, 0);

  mpBitmap = new BitmapWindow(this, wxBitmap(256, 256));

  panelSizer->Add(mpBitmap, 1, wxEXPAND|wxALIGN_CENTER);
  panelSizer->Add(sliderSizer, 0, wxEXPAND);
  mpText = new wxStaticText(this, -1, _T("\n"));
  panelSizer->Add(mpText, 0);

  SetSizer(panelSizer);

  panelSizer->Fit(this);

  mTimer.Start(100);	
}


///
/**
 */
ImageViewer::ImageViewer(wxWindow* parent, const std::vector<wxString>& names, const wxString& show_first) 
: ImageViewerBase(parent, show_first), showFirst(show_first), filenames(names)
{
  wxInitAllImageHandlers();
  wxSize oldBitmapSize = mpBitmap->GetSize();
  DicomFile file(showFirst.mb_str(wxConvLocal));

  std::vector<wxString>::iterator it = find(filenames.begin(), filenames.end(), showFirst);
  if (filenames.size() < 2) {
    mpSliceSlider->Disable();
    mpPlayButton->Disable();
  }
  else {
    mpSliceSlider->SetRange(0, filenames.size() - 1);
    mpSliceSlider->SetValue(static_cast<int>(it - filenames.begin()));
  }
  UpdateImage();
}


///
/**
 */
void 
ImageViewerBase::OnSlide(wxScrollEvent& event)
{
  UpdateImage();
}


///
/**
 */
void
ImageViewerBase::OnPlay(wxCommandEvent& event)
{
  if (mpPlayButton->GetLabel() == _T("play")) {
    if (mpSliceSlider->GetValue() == mpSliceSlider->GetMax()) {
      mpSliceSlider->SetValue(0);
    }
    mpPlayButton->SetLabel(_T("stop"));
  }
  else {
    mpPlayButton->SetLabel(_T("play"));
  }
}


///
/**
 */
void
ImageViewerBase::OnTimer(wxTimerEvent& event)
{
  if (mpPlayButton->GetLabel() != _T("stop")) { return; }
  int fileno = mpSliceSlider->GetValue() + 1;
  if (mpSliceSlider->GetMax() < fileno) {
    mpPlayButton->SetLabel(_T("play"));
    return;
  }
  mpSliceSlider->SetValue(fileno);
  UpdateImage();
}


///
/**
 */
void
ImageViewer::OnWindow(wxCommandEvent& event) 
{
  DicomFile file(filenames.at(mpSliceSlider->GetValue()).mb_str(wxConvLocal));
  int maxpix;
  file.Find("LargestImagePixelValue", maxpix);
  width = maxpix;
  center = maxpix/2;
  useDefaultWC = false;
  UpdateImage();
}


///
/**
 */
void
ImageViewerBase::OnReset(wxCommandEvent& event)
{
  useDefaultWC = true;
  UpdateImage();
}


///
/**
 */
void
ImageViewerBase::SetImage(wxImage image)
{
  if (image.GetWidth() < 256) {
    float scaleFactor = 256.0/image.GetWidth();
    image.Rescale(256, static_cast<int>(image.GetHeight() * scaleFactor));
  }
  mpBitmap->SetBitmap(image);

  wxSize imageSize = wxSize(image.GetWidth(), image.GetHeight());
  if (mpBitmap->GetSize() != imageSize) {
    wxSize newSize = GetSize() + imageSize - mpBitmap->GetSize();
    mpBitmap->SetSize(imageSize);
    SetSize(newSize);
  }
  mImage = image;
}


///
/**
 */
void
ImageViewer::UpdateImage()
{
  DicomFile file(filenames.at(mpSliceSlider->GetValue()).mb_str(wxConvLocal));
  wxImage image = CreateImage(file);
  SetImage(image);

  wxFileName name(filenames.at(mpSliceSlider->GetValue()));
  SetTitle(name.GetFullName());
  wxLogStatus(name.GetFullName());
  mpText->SetLabel(wxString::Format(wxT("Window width: %d. Window center: %d.\nClick and drag to adjust, double click to reset."), width, center));
}


///
/**
 */
void
ImageViewerBase::OnContextMenu(wxContextMenuEvent& event)
{
  wxMenu menu;
  menu.Append(SAVE_IMAGE, _T("Save this image as image file"));
  menu.Append(SAVE_ALL, _T("Save entire series as image files"));
  menu.Append(WINDOW, _T("Auto window level by pixel values"));
  menu.Append(RESET, _T("Reset window level to DICOM file"));
  menu.Append(CLOSE, _T("Close window"));
  PopupMenu(&menu, ScreenToClient(event.GetPosition()));
}


///
/**
 */
void
ImageViewerBase::OnSave(wxCommandEvent& event)
{
  wxFileDialog* dlg = new wxFileDialog(this,
  _T("Save image"), _T(""), _T(""),
  _T("PNG file(*.png)|*.png|Jpeg file(*.jpg)|*.jpg|Bitmap file(*.bmp)|*.bmp|TIFF file (*.tif)|*.tif"),
  wxFD_SAVE);

  if (dlg->ShowModal() == wxID_OK) {

    wxString savefilename = dlg->GetPath();

    if (!mImage.SaveFile(savefilename)) {
      wxMessageBox(_T("No handler for this file type."),
      _T("File was not saved"), wxOK | wxCENTRE, this);
    }
  }
  dlg->Destroy();
}


///
/**
 */
void
ImageViewerBase::OnSaveAll(wxCommandEvent& event)
{

  wxFileDialog* dlg = new wxFileDialog(this,
  _T("Choose image file base name"), _T(""), _T(""),
  _T("PNG file(*.png)|*.png|Jpeg file(*.jpg)|*.jpg|Bitmap file(*.bmp)|*.bmp|TIFF file (*.tif)|*.tif"),
  wxFD_SAVE);

  if (dlg->ShowModal() == wxID_OK) {
    int nfiles = mpSliceSlider->GetMax() + 1;
    wxProgressDialog* p_dlg = new wxProgressDialog(_T("Progress"), 
    _T("Saving files..."), nfiles, this, wxPD_CAN_ABORT | wxPD_APP_MODAL);

    int ndigits = itos(nfiles - 1).size();
    for (int i = 0; i < nfiles; ++i) {
      mpSliceSlider->SetValue(i);
      UpdateImage();
      wxFileName filename(dlg->GetPath());
      wxString base = filename.GetName();
      base.Append(wxString(itos(i, ndigits).c_str(), wxConvLocal));
      filename.SetName(base);
      if (!mImage.SaveFile(filename.GetFullPath())) {
        wxMessageBox(_T("No handler for this file type."),
        _T("File was not saved"),
        wxOK | wxCENTRE, this);
        break;
      }
      if (!p_dlg->Update(i)) { break; }
    }
    p_dlg->Destroy();
  }

  dlg->Destroy();
}


///
/**
 */
void
ImageViewerBase::OnClose(wxCommandEvent& event)
{
  Close();
}


///
/**
 */
template <class T> wxImage
GetImage(DicomFile& file)
{
  int columns;
  int rows;
  //file.Find(DT_IMAGECOLUMNS, columns);
  file.Find("Columns", columns);
  //file.Find(DT_IMAGEROWS, rows);
  file.Find("Rows", rows);

  std::vector<T> data;
  if (file.PixelData(data, columns*rows*sizeof(T)) == 0) {
    wxLogError(_T("Unable to retrieve pixel data."));
  }

  //	int center, width;
  if (useDefaultWC) {
    file.Find("WindowCenter", center);
    file.Find("WindowWidth", width);
  }

  int rescale_slope, rescale_intercept;
  if (file.Find("RescaleSlope", rescale_slope) &&
      file.Find("RescaleIntercept", rescale_intercept)) {
    if (rescale_slope != 0) {
      width /= rescale_slope;
      center = (center - rescale_intercept)/rescale_slope;
    }
  }

  // rescale to unsigned 8 bit
  T data_max = std::numeric_limits<T>::max();
  T data_min = std::numeric_limits<T>::min();
  if ((center + width/2) < data_max) {
    data_max = static_cast<T>(center + width/2);
  }

  if ((center - width/2) > data_min) {
    data_min = static_cast<T>(center - width/2);
  }

  wxUint8 char_max = std::numeric_limits<wxUint8>::max();
  int max_index = data.size();

  if ((data_max - data_min) != 0) {
    float scaling = static_cast<float>(char_max)/static_cast<float>(data_max - data_min);
    T* c_ptr = &data.front();
    for (int i = 0; i < max_index; ++i) {
      if (*c_ptr > data_max) { *c_ptr = data_max; }
      else if (*c_ptr < data_min) { *c_ptr = data_min; }
      *c_ptr = static_cast<T>((*c_ptr) * scaling - data_min * scaling);
      ++c_ptr;
    }
  }

  unsigned char* my_data = (unsigned char*) malloc(columns*rows*3);
  unsigned char* new_ptr = my_data;
  T* old_ptr = &data.front();

  for (int i = 0; i < max_index; ++i) {
    unsigned char pixel = static_cast<unsigned char>(*old_ptr++);
    *new_ptr++ = pixel;
    *new_ptr++ = pixel;
    *new_ptr++ = pixel;
  }

  return wxImage(columns, rows, my_data, false);
}


///
/**
 */
wxImage
GetRgbImage(DicomFile& file)
{
  int columns;
  int rows;
  //file.Find(DT_IMAGECOLUMNS, columns);
  file.Find("Columns", columns);
  //file.Find(DT_IMAGEROWS, rows);
  file.Find("Rows", rows);

  unsigned char* data = (unsigned char*) malloc(columns*rows*3);
  file.PixelData(data, columns*rows*sizeof(unsigned char)*3);
  return wxImage(columns, rows, data, false);

}


///
/**
 */
wxImage
CreateImage(DicomFile& file)
{
  wxImage image;
  std::string pi;
  
  if (!file.Find("PhotometricInterpretation", pi)) {
    wxLogError(_T("File does not contain image data."));
    image = wxImage(256, 256);
    return image;
  }
  if (pi == "RGB") {
    image = GetRgbImage(file);
  }
  else {

    int bits_allocated, pixel_rep = 0;

    file.Find("BitsAllocated", bits_allocated);
    file.Find("PixelRepresentation", pixel_rep);

    switch (bits_allocated + pixel_rep) {

    case 9 :
      image = GetImage<wxInt8>(file);
      break;

    case 8 :
      image = GetImage<wxUint8>(file);
      break;

    case 33 :
      image = GetImage<wxInt32>(file);
      break;

    case 32 :
      image = GetImage<wxUint32>(file);
      break;

    case 17 :
      image = GetImage<wxInt16>(file);
      break;

    case 16 :
      default :
      image = GetImage<wxUint16>(file);
    }
  }
  return image;
}


///
/**
 */
MultiFrameImageViewer::MultiFrameImageViewer(wxWindow* parent, const wxString& name) 
: ImageViewerBase(parent, name)
{
  DicomFile file(name.mb_str(wxConvLocal));
  //file.Find(DT_IMAGECOLUMNS, columns);
  file.Find("Columns", columns);
  //file.Find(DT_IMAGEROWS, rows);
  file.Find("Rows", rows);
  file.Find("NumberOfFrames", frames);

  goodVolume = CreateVolume(file);

  mpSliceSlider->SetRange(0, volume.size() - 1);
  mpSliceSlider->SetValue(0);

  UpdateImage();
}


///
/**
 */
void 
MultiFrameImageViewer::UpdateImage()
{
  int frame = mpSliceSlider->GetValue();
  if (goodVolume) {
    SetImage(GetImage(frame));
  }
  wxLogStatus(_T("%d"), frame);
}


///
/**
 */
template <class T> void
GetVolume(DicomFile& file, std::vector<std::vector<wxUint8> > &volume)
{
  wxBusyCursor busy;
  long columns, rows, frames;
  //file.Find(DT_IMAGECOLUMNS, columns);
  file.Find("Columns", columns);
  //file.Find(DT_IMAGEROWS, rows);
  file.Find("Rows", rows);
  file.Find("NumberOfFrames", frames);

  volume.resize(frames);
  for (long i = 0; i < frames; ++i) {
    volume[i].resize(rows*columns*3);
  }

  std::vector<T> data;
  file.PixelData(data, columns * rows * frames * sizeof(T));

  // rescale to unsigned 8 bit
  typename std::vector<T>::iterator frame_start = data.begin();
  typename std::vector<T>::iterator frame_end;
  wxUint8 char_max = std::numeric_limits<wxUint8>::max();
  
  while (frame_start < data.end()) {
    frame_end = frame_start + columns*rows;
    T data_max = *std::max_element(frame_start, frame_end);
    T data_min = *std::min_element(frame_start, frame_end);
    int max_index = columns*rows;
    if ((data_max - data_min) != 0) {
      float scaling = static_cast<float>(char_max)/static_cast<float>(data_max - data_min);
      T* c_ptr = &*frame_start;
      for (long i = 0; i < columns*rows; ++i) {  //, ++c_ptr) {
        *c_ptr = static_cast<T>((*c_ptr) * scaling - data_min * scaling);
        ++c_ptr;
      }
    }
    frame_start = frame_end;
  }

  T* old_ptr = &data.front();

  for (long i = 0; i < frames; ++i) {
    unsigned char* new_ptr = &volume[i].front();
    for (long j = 0; j < (columns*rows); ++j){
      unsigned char pixel = static_cast<unsigned char>(*old_ptr++);
      *new_ptr++ = pixel;
      *new_ptr++ = pixel;
      *new_ptr++ = pixel;
    }
  }
}


///
/**
 */
wxImage
MultiFrameImageViewer::GetImage(int frame)
{
  wxUint8* ptr = &volume[frame].front();
  return wxImage(columns, rows, ptr, true);
}


///
/**
 */
// nb - need to find a multiframe rbg image to test this!
void
GetRgbVolume(DicomFile& file, std::vector<std::vector<wxUint8> > &volume)
{
  wxBusyCursor busy;
  long columns, rows, frames;
  //file.Find(DT_IMAGECOLUMNS, columns);
  file.Find("Columns", columns);
  //file.Find(DT_IMAGEROWS, rows);
  file.Find("Rows", rows);
  file.Find("NumberOfFrames", frames);

  std::vector<unsigned char> data;
  file.PixelData(data, columns * rows * frames * sizeof(unsigned char));
  unsigned char* old_ptr = &data.front();

  volume.resize(frames);
  for (long i = 0; i < frames; ++i) {
    volume[i].resize(rows*columns*3);
  }

  for (long i = 0; i < frames; ++i) {
    unsigned char* new_ptr = &volume[i].front();
    memcpy(new_ptr, old_ptr, rows*columns*sizeof(unsigned char));
    old_ptr += rows*columns*sizeof(unsigned char);
  }
}


///
/**
 */
bool
MultiFrameImageViewer::CreateVolume(DicomFile& file)
{
  std::string pi;

  if (!file.Find("PhotometricInterpretation", pi)) {
    wxLogError(_T("File does not contain image data."));
    return 0;
  }

  if (pi == "RGB") {
    GetRgbVolume(file, volume);
    return 1;
  }
  else {

    int bits_allocated, pixel_rep = 0;

    file.Find("BitsAllocated", bits_allocated);
    file.Find("PixelRepresentation", pixel_rep);

    switch (bits_allocated + pixel_rep) {

    case 8 :
      GetVolume<wxUint8>(file, volume);
      break;
    case 9 :
      GetVolume<wxInt8>(file, volume);
      break;
    case 16 :
      GetVolume<wxUint16>(file, volume);
      break;
    case 17 :
      GetVolume<wxInt16>(file, volume);
      break;
    case 32 :
      GetVolume<wxUint32>(file, volume);
      break;
    case 33 :
      GetVolume<wxInt32>(file, volume);
      break;
    default :
      wxLogMessage(_("BitsAllocated and PixelRepresentation values (%d, %d) not supported."), bits_allocated, pixel_rep);
    }
  }
  return 1;
}
