#include <QPainter>
#include <QMouseEvent>

#include <QDebug>
#include <QAction>
#include <QFont>
#include <QGuiApplication>
#include <QInputDialog>
#include <QMessageBox>
#include <QFileDialog>
#include <QRegularExpression>
#include <cmath>

#include <string>
#include <regex>
#include <iostream>

#include "SegmentMonitor.h"

SegmentMonitor::SegmentMonitor(QWidget *parent)
  : QWidget(parent),
    m_step(40),
    m_segmentHeight(4),
    m_duration(1),
    m_idxColor(0),
    m_annotMode(false),
    m_prevPosition(-1)
{
  resize(75000, 140);

  // contextual menu containing speaker as manually annotated
  QFont boldFont = QGuiApplication::font();
  boldFont.setBold(true);

  m_newSpeakerAct = new QAction(tr("Add new speaker"), this);
  m_endAct = new QAction(tr("End"), this);
  
  m_endAct->setFont(boldFont);
  m_newSpeakerAct->setFont(boldFont);

  setContextMenuPolicy(Qt::CustomContextMenu);
  m_speakers = new QMenu;
  
  m_speakers->addAction(m_newSpeakerAct);
  m_speakers->addAction(m_endAct);
  m_speakers->addSeparator();

  m_colors << QColor(219, 213, 185) << QColor(224, 199, 30);
  m_colors << QColor(186, 22, 63) << QColor(255, 250, 129);
  m_colors << QColor(181, 225, 174) << QColor(154, 206, 223);
  m_colors << QColor(249, 149, 51) << QColor(252, 169, 133);
  m_colors << QColor(193, 179, 215) << QColor(116, 161, 97);
  m_colors << QColor(133, 202, 93) << QColor(72, 181, 163);
  m_colors << QColor(69, 114, 147) << QColor(249, 150, 182);
  m_colors << QColor(117, 137, 191) << QColor(84, 102, 84);
  reset();

  connect(this, SIGNAL(customContextMenuRequested(const QPoint&)),
	  this, SLOT(showContextMenu(const QPoint&)));
  connect(m_newSpeakerAct, SIGNAL(triggered()), this, SLOT(addNewSpeaker()));
  connect(m_endAct, SIGNAL(triggered()), this, SLOT(endOfAnnotation()));
}

void SegmentMonitor::reset()
{
  for (int i(0); i < m_segmentationList.size(); i++)
    m_segmentationList[i].clear();
  m_segmentationList.clear();
  m_currList.clear();
  m_lastStart = 0;
  m_currPosition = 0;
  m_idxColor = 0;
  m_currText = QString();
  m_currSpeaker = QString();
  m_segColors.insert(QString(), QColor(255, 255, 255));
  m_segColors.insert("S", QColor(255, 255, 255));
}

///////////////
// modifiers //
///////////////

void SegmentMonitor::setAnnotMode(bool annotMode)
{
  m_annotMode = annotMode;
  if (annotMode)
    m_segmentHeight = 8;
  else
    m_segmentHeight = 4;
}

///////////////
// accessors //
///////////////

qreal SegmentMonitor::getRatio() const
{
  return m_ratio;
}

///////////
// slots //
///////////

void SegmentMonitor::processShot(Segment *shot)
{
  m_currList.push_back(shot);

  QString camera = shot->getLabel();
  if (!camera.isEmpty() && !m_segColors.contains(camera))
    m_segColors.insert(camera, m_colors[m_idxColor++ % m_colors.size()]);
}

void SegmentMonitor::processSpokenFrame(qint64 position, const QString &text, const QString &speaker)
{
  // regular expression to detect subtitles corresponding to noise
  QRegularExpression noiseSource("\\(.*\\)");
  QRegularExpressionMatch match = noiseSource.match(text);

  if (position != 0 && !match.hasMatch()) {

    // possible boundary
    if (m_currSpeaker != speaker || m_currText != text) {

      // end of speech segment
      if (!m_currSpeaker.isEmpty()) {
	SpeechSegment *segment = new SpeechSegment(m_lastStart, m_currPosition, m_currText, m_currSpeaker);
	m_currList.push_back(segment);

	if (!m_segColors.contains(m_currSpeaker))
	  m_segColors.insert(m_currSpeaker, m_colors[m_idxColor++ % m_colors.size()]);
      }
      
      // beginning of new segment
      if (!speaker.isEmpty())
	m_lastStart = position;
    }

    // updating data
    m_currPosition = position;
    m_currText = text;
    m_currSpeaker = speaker;
  }
}

void SegmentMonitor::segmentationRetrieved()
{
  // retrieving speakers as annotated
  if (m_segmentationList.size() == 1) {

    QString currSpeaker;

    for (int i(0); i < m_currList.size(); i++)
      if (!m_annotSpeakers.contains((currSpeaker = m_currList[i]->getLabel())) && currSpeaker != "S") {
	m_annotSpeakers.push_back(currSpeaker);
	m_speakers->addAction(new QAction(currSpeaker, this));
      }
  }

  m_segmentationList.push_back(m_currList);
  m_selectedLabels.push_back(QString());
  m_currList.clear();
  m_step = height() / (m_segmentationList.size() + 1);
}

void SegmentMonitor::positionChanged(qint64 position)
{
  m_currPosition = position;
  update();
}

void SegmentMonitor::updateDuration(qint64 duration)
{
  m_duration = duration;
}

void SegmentMonitor::setWidth(int newWidth)
{
  m_width = newWidth;
  resize(m_width, height());
  m_ratio = static_cast<qreal>(m_width) / m_duration;
  update();
}

void SegmentMonitor::paintEvent(QPaintEvent *event)
{
  Q_UNUSED(event);
  int x1, x2;
  int y1, y2;
  QString segLabel;
  QPainter painter(this);
  QColor color;

  QFont font;
  font.setPointSize(m_segmentHeight * 2);
  painter.setFont(font);

  // filling background
  painter.fillRect(QRect(0, 0, m_width, height()), QBrush(Qt::white));

  // draw current position
  painter.setPen(Qt::gray);
  painter.drawLine(m_currPosition * m_ratio, 0, m_currPosition * m_ratio, height());

  // drawing segments
  for (int i(0); i < m_segmentationList.size(); i++) {

    y1 = m_step * (i + 1) - m_segmentHeight / 2;
    y2 = y1 + m_segmentHeight;

    // surrounding each segmentation
    painter.setPen(Qt::lightGray);
    painter.drawRect(0, y1, m_width, m_segmentHeight);

    painter.setPen(Qt::black);
    for (int j(0); j < m_segmentationList[i].size(); j++) {

      segLabel = m_segmentationList.at(i).at(j)->getLabel();
      color = m_segColors[segLabel];
      
      if (m_selectedLabels[i].isEmpty() || m_selectedLabels[i] == segLabel) {

	x1 = qRound(m_segmentationList.at(i).at(j)->getPosition() * m_ratio);

	// case of shot segmentation
	if (i == 0) {

	  // last shot
	  if (j == m_segmentationList[i].size() - 1)
	    x2 = qRound(m_duration * m_ratio);
	  else
	    x2 = qRound(m_segmentationList.at(i).at(j + 1)->getPosition() * m_ratio);
	}
      
	// case of speech segmentation
	else
	  x2 = qRound(m_segmentationList.at(i).at(j)->getEnd() * m_ratio);

	if (m_selectedLabels[i] == segLabel) {

	  if (segLabel == "S")
	    segLabel = QString::number(j + 1);

	  QRectF rectF = painter.boundingRect(QRectF(x1, y2 + m_segmentHeight, x2 - x1, m_segmentHeight * 4), Qt::AlignHCenter, segLabel);
	  
	  if (rectF.width() <= x2 - x1)
	    painter.drawText(rectF, Qt::AlignHCenter, segLabel);
	}

	QPointF p1(x1, y1);
	QPointF p2(x2, y2);
	painter.fillRect(QRectF(p1, p2), color);
	painter.drawRect(QRectF(p1, p2));
      }
    }
  }
}

void SegmentMonitor::mouseDoubleClickEvent(QMouseEvent *event)
{
  if (event->button() == Qt::LeftButton) {

    bool inBound(false);
    qreal inf, sup;
    int start(-1), end(-1);
  
    bool found(false);
    int min, max;
 
    qreal position(event->x() / m_ratio);

    // test if clicked in segmentation regions
    int i(0);
    int j(0);

    while (i < m_segmentationList.size() && !inBound) {
      inf = m_step * (i + 1);
      sup = inf + m_segmentHeight;
      if (event->y() >= inf && event->y() <= sup)
	inBound = true;
      i++;
    }

    // binary search to retrieve index of segment clicked segment
    if (inBound) {

      QList<Segment *> segList(m_segmentationList[--i]);
    
      min = 0;
      max = segList.size() - 1;
      
      while (min <= max && !found) {
	j = (min + max) / 2;
      
	start = segList[j]->getPosition();
      
	// shot segmentation activated
	if (i == 0) {
	  // last shot
	  if (j == segList.size() - 1)
	    end = m_duration;
	  else
	    end = segList[j+1]->getPosition();
	}

	// speech segmentation activated
	else
	  end = segList[j]->getEnd();

	if (position >= start && position <= end)
	  found = true;
	else if (position < start)
	  max = j - 1;
	else
	  min = j + 1;
      }

      if (found && (m_selectedLabels[i].isEmpty() || m_selectedLabels[i] == segList[j]->getLabel())) {
	QList<QPair<qint64, qint64>> utterances;
	utterances.push_back(QPair<qint64, qint64>(start, end));
	emit playSegments(utterances);
      }
    }
  }
}

void SegmentMonitor::mousePressEvent(QMouseEvent *event)
{
  bool inBound(false);
  qreal inf, sup;
  int start(-1), end(-1);
  
  bool found(false);
  int min, max;
 
  qreal position(event->x() / m_ratio);

  // test if clicked in segmentation regions
  int i(0);
  int j(0);

  while (i < m_segmentationList.size() && !inBound) {
    inf = m_step * (i + 1);
    sup = inf + m_segmentHeight;
    if (event->y() >= inf && event->y() <= sup)
      inBound = true;
    i++;
  }
    
  // binary search to retrieve index of segment clicked segment
  if (inBound) {

    QList<Segment *> segList(m_segmentationList[--i]);
    
    min = 0;
    min = 0;
    max = segList.size() - 1;
      
    while (min <= max && !found) {
      j = (min + max) / 2;
      
      start = segList[j]->getPosition();
      
      // shot segmentation activated
      if (i == 0) {
	// last shot
	if (j == segList.size() - 1)
	  end = m_duration;
	else
	  end = segList[j+1]->getPosition();
      }

      // speech segmentation activated
      else
	end = segList[j]->getEnd();

      if (position >= start && position <= end)
	found = true;
      else if (position < start)
	max = j - 1;
      else
	min = j + 1;
    }

    if (found && !m_annotMode && event->button() == Qt::RightButton) {
      if (m_selectedLabels[i].isEmpty())
	m_selectedLabels[i] = segList[j]->getLabel();
      else if (m_selectedLabels[i] == segList[j]->getLabel())
	m_selectedLabels[i] = "";
      update();
    }
  }
}

void SegmentMonitor::mouseMoveEvent(QMouseEvent *event)
{
  bool inBound(false);
  bool firstHalf(false);
  qreal inf, sup;
  int start(-1), end(-1);
  
  bool found(false);
  int min, max;
 
  qreal position(event->x() / m_ratio);

  // test if clicked in segmentation regions
  int i(0);
  int j(0);

  while (i < m_segmentationList.size() && !inBound) {
    inf = m_step * (i + 1);
    sup = inf + m_segmentHeight;
    if (event->y() >= inf && event->y() <= sup)
      inBound = true;
    i++;
  }
    
  // binary search to retrieve index of clicked segment
  if (i == 2 && inBound) {

    QList<Segment *> segList(m_segmentationList[--i]);
    
    min = 0;
    min = 0;
    max = segList.size() - 1;
      
    while (min <= max && !found) {
      j = (min + max) / 2;
      
      start = segList[j]->getPosition();
      
      // shot segmentation activated
      if (i == 0) {
	// last shot
	if (j == segList.size() - 1)
	  end = m_duration;
	else
	  end = segList[j+1]->getPosition();
      }

      // speech segmentation activated
      else
	end = segList[j]->getEnd();

      if (position >= start && position <= end)
	found = true;
      else if (position < start)
	max = j - 1;
      else
	min = j + 1;
    }

    // reference speech segment clicked for annotation purpose
    if (found && m_annotMode) {

      if (m_prevPosition != -1) {

	qint64 start, end, mid, position, newPosition, diff;

	start = segList[j]->getPosition();
	end = segList[j]->getEnd();
	mid = (start + end) / 2;

	// normalized mouse position in milliseconds from the beginning
	position = event->x() / m_ratio;

	if (m_prevPosition != -1)
	  diff = position - m_prevPosition;
	else
	  diff = 0;

	// modifying beginning of segment
	if (position >= start && position < mid) {

	  firstHalf = true;

	  if (j > 0 && start + diff <= segList[j-1]->getEnd())
	    newPosition = start;
	  else
	    newPosition = start + diff;
	  
	  newPosition = static_cast<qint64>(ceil(newPosition / 40.0)) * 40;

	  if (end - newPosition >= 240) {
	    segList[j]->setPosition(newPosition);
	    emit resetSpeaker(start, end, newPosition, end, true, VideoFrame::Ref);
	  }
	}
	 
	// modifying end of segment
	else if (position >= mid && position <= end) {

	  firstHalf = false;

	  if (j < segList.size() - 1 && end + diff >= segList[j+1]->getPosition())
	    newPosition = end;
	  else
	    newPosition = end + diff;
	  
	  newPosition = static_cast<qint64>(floor(newPosition / 40.0)) * 40;
	  if (newPosition - start >= 240) {
	    segList[j]->setEnd(newPosition);
	    emit resetSpeaker(start, end, start, newPosition, true, VideoFrame::Ref);
	  }
	}
	
	update();
      }
    }

    // updating normalized mouse position
    if (firstHalf)
      m_prevPosition = static_cast<qint64>(ceil(position / 40.0)) * 40;
    else
      m_prevPosition = static_cast<qint64>(floor(position / 40.0)) * 40;
  }
}

void SegmentMonitor::mouseReleaseEvent(QMouseEvent *event)
{
  Q_UNUSED(event);
  m_prevPosition = -1;
}

void SegmentMonitor::showContextMenu(const QPoint &point)
{
  // retrieving selected speech segment
  bool inBound(false);
  qreal inf, sup;
  int start(-1), end(-1);
  
  bool found(false);
  int min, max;
 
  qreal position(point.x() / m_ratio);

  // test if clicked in segmentation regions
  int i(0);
  int j(0);

  while (i < m_segmentationList.size() && !inBound) {
    inf = m_step * (i + 1);
    sup = inf + m_segmentHeight;
    if (point.y() >= inf && point.y() <= sup)
      inBound = true;
    i++;
  }
    
  // binary search to retrieve index of segment clicked segment
  if (i == 2 && inBound) {

    QList<Segment *> segList(m_segmentationList[--i]);
    
    min = 0;
    max = segList.size() - 1;
      
    while (min <= max && !found) {
      j = (min + max) / 2;
      
      start = segList[j]->getPosition();
      
      // shot segmentation activated
      if (i == 0) {
	// last shot
	if (j == segList.size() - 1)
	  end = m_duration;
	else
	  end = segList[j+1]->getPosition();
      }
      
      // speech segmentation activated
      else
	end = segList[j]->getEnd();
      
      if (position >= start && position <= end)
	found = true;
      else if (position < start)
	max = j - 1;
      else
	min = j + 1;
    }

    if (found && m_annotMode) {
      QPoint globalPos = mapToGlobal(point);
      QAction *selected = m_speakers->exec(globalPos);

      if (selected && selected != m_newSpeakerAct && selected != m_endAct) {
	emit setSpeaker(segList[j]->getPosition(), segList[j]->getEnd(), selected->text(), VideoFrame::Ref);
	segList[j]->setLabel(selected->text());
	update();
      }
    }
  }
}

void SegmentMonitor::addNewSpeaker()
{
  bool ok(false);
  QString speaker = QInputDialog::getText(this, "New speaker", "New speaker name:", QLineEdit::Normal, QString(), &ok);

  if (ok && !speaker.isEmpty()) {
    m_speakers->addAction(new QAction(speaker, this));
    m_segColors.insert(speaker, m_colors[m_idxColor++ % m_colors.size()]);
  }
}

///////////////////////
// auxiliary methods //
///////////////////////

void SegmentMonitor::endOfAnnotation()
{
  QList<Segment *> shotSegList = m_segmentationList[0];
  QList<Segment *> spkSegList = m_segmentationList[1];
  int ans;

  /*********************************/
  /* I: merging adjacent segments? */
  /*********************************/

  ans = QMessageBox::question(this, tr("End of annotation"), tr("End of annotation. Merge adjacent segments?"), QMessageBox::Yes | QMessageBox::No);

  if (ans == QMessageBox::Yes) {

    QString prevSpeaker, currSpeaker;
    int prevStart, prevEnd, start, end;

    // merging close speech segments coming from same speaker
    for (int i(1); i < spkSegList.size(); i++) {
      
      prevSpeaker = spkSegList[i-1]->getLabel();
      prevStart = spkSegList[i-1]->getPosition();
      prevEnd = spkSegList[i-1]->getEnd();

      currSpeaker = spkSegList[i]->getLabel();
      start = spkSegList[i]->getPosition();
      end = spkSegList[i]->getEnd();
      
      if (currSpeaker == prevSpeaker && start - prevEnd <= 500)
	emit resetSpeaker(start, end, prevStart, end, false, VideoFrame::Ref);
    }
  }

  /*********************************************/
  /* II: exporting file containing utterances? */
  /*********************************************/

  ans = QMessageBox::question(this, tr("End of annotation"), tr("End of annotation. Export file containing labeled speech segments?"), QMessageBox::Yes | QMessageBox::No);

  if (ans == QMessageBox::Yes) {
    QString fName = QFileDialog::getSaveFileName(this, tr("Save File"),
						 tr("tools/spkDiarization/lblAcousticSegmentation"),
						 tr("Label (*.lbl)"));
    if (!fName.isEmpty()) {
      QFile datFile(fName + ".dat");

      if (!datFile.open(QIODevice::WriteOnly | QIODevice::Text))
	return;

      QTextStream lblOut(&datFile);

      qreal start(-1.0), end(-1.0);
      QString label;

      // test for estimation of speech segments based on subtitles
      bool sub(true);
      int i(0);

      while (i < spkSegList.size() && sub)
	if (spkSegList[i++]->getLabel() != "S")
	  sub = false;
      
      if (sub) {
	for (int i(0); i < spkSegList.size(); i++) {
	  start = spkSegList[i]->getPosition() / 1000.0;
	  end = spkSegList[i]->getEnd() / 1000.0;

	  lblOut << start << " " << end << " " << "speech" << "\n";
	}
      }

      else {
	qreal beginning(-1.0);

	for (int i(0); i < spkSegList.size(); i++) {
	  
	  label = spkSegList[i]->getLabel();
	  label.replace(QRegularExpression(" "), "_");

	  if (label != "S") {
	    start = spkSegList[i]->getPosition() / 1000.0;
	    end = spkSegList[i]->getEnd() / 1000.0;
	    lblOut << start << " " << end << " " << label << "\n";
	    
	    if (beginning == -1)
	      beginning = start;
	  }
	}

	QFile uemFile(fName + ".uem");

	QRegularExpression re(".*/(.+_.+)");
	QRegularExpressionMatch match = re.match(fName);
	QString baseName;
      
	if (match.hasMatch())
	  baseName = match.captured(1);

	if (!uemFile.open(QIODevice::WriteOnly | QIODevice::Text))
	  return;

	QTextStream uemOut(&uemFile);

	uemOut << baseName << " " << 1 << " " << beginning << " " << end << "\n";
      }
    }
  }

  /****************************************************************************/
  /* III: exporting file containing utterances included in alternating shots? */
  /****************************************************************************/

  ans = QMessageBox::question(this, tr("End of annotation"), tr("End of annotation. Export file containing utterances included in alternating shots?"), QMessageBox::Yes | QMessageBox::No);

  if (ans == QMessageBox::Yes) {
    QString fName = QFileDialog::getSaveFileName(this, tr("Save File"),
						 tr("tools/spkDiarization/shotPatterns"),
						 tr("CSV (*.csv)"));
    if (!fName.isEmpty()) {
      QFile csvFile(fName + ".csv");
      QList<QString> lblWindow;             // last four shot labels: used to detect shot pattern

      if (!csvFile.open(QIODevice::WriteOnly | QIODevice::Text))
	return;
      
      QTextStream csvOut(&csvFile);
      QString label;
      bool inPattern(false);                // indicates if currently visiting a pattern
      int j(0);

      // looping over the shot positions
      for (int i(0); i < shotSegList.size() - 1; i++) {
	  
	label = shotSegList[i]->getLabel(); // current shot label

	// updating shot label window
	lblWindow.push_back(label);
	if (lblWindow.size() > 4)
	  lblWindow.pop_front();
	  
	// new pattern detected
	if (!inPattern && testShotPattern(lblWindow)) {
	    
	  // setting spoken segment position
	  while (j > 0 && spkSegList[j]->getPosition() > shotSegList[i-3]->getPosition())
	    j--;
	  while (j < spkSegList.size() && spkSegList[j]->getEnd() < shotSegList[i-3]->getPosition())
	    j++;
	  
	  // writing out labels of alternating labels
	  csvOut << lblWindow[0] << "_" << lblWindow[1] << " ";
	  inPattern = true;
	}
	
	// end of pattern
	if (inPattern && !testShotPattern(lblWindow)) {

	  // writing out new line
	  csvOut << endl;
	  inPattern = false;
	}

	// moving into current pattern
	if (inPattern && testShotPattern(lblWindow))

	  // writing out positions of spoken segments covered by the pattern
	  while (j < spkSegList.size() && spkSegList[j]->getPosition() < shotSegList[i+1]->getPosition()) {
	    csvOut << (j + 1) << " ";
	    j++;
	  }

	// updating flag
	inPattern = testShotPattern(lblWindow);
      }
    }
  }
}

bool SegmentMonitor::testShotPattern(const QList<QString> &lblWindow)
{
  if (!lblWindow.contains("") && lblWindow.size() == 4)
    return
      lblWindow[0] == lblWindow[2] &&
      lblWindow[1] == lblWindow[3] &&
      lblWindow[0] != lblWindow[1];
  
  return false;
}