#ifndef NOMINMAX
#define NOMINMAX
#endif
#include "DvhComparePage.h"
#include "DvhPlotWidget.h"
#include "dcmIO.h"

#include <QLineEdit>
#include <QListWidget>
#include <QTableWidget>
#include <QTextEdit>
#include <QProgressBar>
#include <QCheckBox>
#include <QPushButton>
#include <QLabel>
#include <QHBoxLayout>
#include <QVBoxLayout>
#include <QGridLayout>
#include <QGroupBox>
#include <QHeaderView>
#include <QAbstractItemView>
#include <QFileDialog>
#include <QMessageBox>
#include <QSettings>
#include <QFileInfo>
#include <QDir>
#include <QCoreApplication>
#include <QLibrary>
#include <QColor>
#include <QPair>
#include <QFont>

#include <cmath>
#include <cstring>

namespace {

struct Rgb {
  int r, g, b;
};

static QVector<Rgb> dvhPalette()
{
  return {
    {255, 0, 255}, {227, 23, 13}, {128, 42, 42}, {255, 128, 0},
    {255, 153, 18}, {127, 255, 0}, {34, 139, 34}, {64, 224, 205},
    {3, 168, 158}, {0, 0, 255}, {218, 112, 214}, {138, 43, 226}
  };
}

static Qt::PenStyle dosePenStyle(int gid)
{
  static const Qt::PenStyle styles[] = {
    Qt::SolidLine, Qt::DashLine, Qt::DotLine, Qt::DashDotLine, Qt::DashDotDotLine
  };
  return styles[gid % 5];
}

static bool isMetaPath(const QString& path)
{
  const QString l = path.toLower();
  return l.endsWith(QStringLiteral(".mhd")) || l.endsWith(QStringLiteral(".mha"))
      || l.endsWith(QStringLiteral(".nii")) || l.endsWith(QStringLiteral(".nii.gz"));
}

static bool isDcmPath(const QString& path)
{
  const QString l = path.toLower();
  return l.endsWith(QStringLiteral(".dcm")) || l.endsWith(QStringLiteral(".dicom"));
}

} // namespace

DvhComparePage::DvhComparePage(QWidget* parent)
  : QWidget(parent)
{
  buildUi();
}

DvhComparePage::~DvhComparePage()
{
  freeResults();
}

void DvhComparePage::buildUi()
{
  // Original dvhViewer layout:
  //   [ ROI LISTS | DVH RESULTS | PATHS / CONTROL / OUTPUT ]
  auto* root = new QHBoxLayout(this);
  root->setContentsMargins(8, 8, 8, 8);
  root->setSpacing(8);

  // ---- Left: ROI LISTS ----
  auto* roiBox = new QGroupBox(tr("ROI LISTS"), this);
  auto* roiLay = new QGridLayout(roiBox);
  auto* allBtn = new QPushButton(tr("ALL"), this);
  auto* noneBtn = new QPushButton(tr("CLC"), this);
  auto* refreshBtn = new QPushButton(tr("DSP"), this);
  connect(allBtn, &QPushButton::clicked, this, &DvhComparePage::onSelectAllRoi);
  connect(noneBtn, &QPushButton::clicked, this, &DvhComparePage::onClearAllRoi);
  connect(refreshBtn, &QPushButton::clicked, this, &DvhComparePage::onRoiVisibilityChanged);
  roiLay->addWidget(allBtn, 0, 0);
  roiLay->addWidget(noneBtn, 0, 1);
  roiLay->addWidget(refreshBtn, 0, 2);
  m_roiList = new QListWidget(this);
  roiLay->addWidget(m_roiList, 1, 0, 5, 3);
  root->addWidget(roiBox, 1);

  // ---- Center: DVH RESULTS ----
  auto* plotBox = new QGroupBox(tr("DVH RESULTs"), this);
  auto* plotLay = new QVBoxLayout(plotBox);
  plotLay->setContentsMargins(4, 8, 4, 4);
  m_plot = new DvhPlotWidget(this);
  plotLay->addWidget(m_plot, 1);
  root->addWidget(plotBox, 7);

  // ---- Right: PATHS + CONTROL + OUTPUT (~20% narrower than before) ----
  auto* right = new QVBoxLayout();
  root->addLayout(right, 3);

  auto* pathBox = new QGroupBox(tr("PATHS"), this);
  auto* pathLay = new QGridLayout(pathBox);

  auto* imgLab = new QLabel(tr("Image  :"), this);
  imgLab->setFixedWidth(70);
  m_imgEdit = new QLineEdit(this);
  m_imgEdit->setPlaceholderText(tr("CT/参考图像：文件夹、.mhd、.nii"));
  auto* imgBtn = new QPushButton(tr("..."), this);
  imgBtn->setFixedSize(24, 24);
  connect(imgBtn, &QPushButton::clicked, this, &DvhComparePage::onBrowseImage);
  pathLay->addWidget(imgLab, 0, 0);
  pathLay->addWidget(m_imgEdit, 0, 1, 1, 10);
  pathLay->addWidget(imgBtn, 0, 11);

  auto* rtsLab = new QLabel(tr("Struc  :"), this);
  rtsLab->setFixedWidth(70);
  m_rtsEdit = new QLineEdit(this);
  m_rtsEdit->setPlaceholderText(tr("RTSTRUCT .dcm"));
  auto* rtsBtn = new QPushButton(tr("..."), this);
  rtsBtn->setFixedSize(24, 24);
  connect(rtsBtn, &QPushButton::clicked, this, &DvhComparePage::onBrowseRts);
  pathLay->addWidget(rtsLab, 1, 0);
  pathLay->addWidget(m_rtsEdit, 1, 1, 1, 10);
  pathLay->addWidget(rtsBtn, 1, 11);

  auto* addDose = new QPushButton(tr("+"), this);
  auto* rmDose = new QPushButton(tr("-"), this);
  addDose->setFixedWidth(36);
  rmDose->setFixedWidth(36);
  connect(addDose, &QPushButton::clicked, this, &DvhComparePage::onAddDose);
  connect(rmDose, &QPushButton::clicked, this, &DvhComparePage::onRemoveDose);
  pathLay->addWidget(addDose, 2, 9);
  pathLay->addWidget(rmDose, 2, 10);

  m_doseTable = new QTableWidget(0, 2, this);
  m_doseTable->setHorizontalHeaderLabels({tr("Name"), tr("Dose Path")});
  m_doseTable->horizontalHeader()->setStretchLastSection(true);
  m_doseTable->verticalHeader()->setVisible(false);
  m_doseTable->setSelectionBehavior(QAbstractItemView::SelectRows);
  m_doseTable->setSelectionMode(QAbstractItemView::SingleSelection);
  m_doseTable->setShowGrid(false);
  pathLay->addWidget(m_doseTable, 3, 0, 8, 12);
  right->addWidget(pathBox, 5);

  auto* ctrlBox = new QGroupBox(tr("CONTROL"), this);
  auto* ctrlLay = new QHBoxLayout(ctrlBox);
  m_progress = new QProgressBar(this);
  m_progress->setRange(0, 10);
  m_progress->setValue(0);
  m_calcBtn = new QPushButton(tr("CALCULATE"), this);
  m_calcBtn->setMinimumWidth(100);
  connect(m_calcBtn, &QPushButton::clicked, this, &DvhComparePage::onCalculate);
  ctrlLay->addWidget(m_progress, 8);
  ctrlLay->addWidget(m_calcBtn, 1);
  right->addWidget(ctrlBox, 1);

  auto* logBox = new QGroupBox(tr("OUT PUT INFO."), this);
  auto* logLay = new QVBoxLayout(logBox);
  m_log = new QTextEdit(this);
  m_log->setReadOnly(true);
  m_log->setFont(QFont(QStringLiteral("Consolas"), 9));
  logLay->addWidget(m_log);
  right->addWidget(logBox, 4);
}

QString DvhComparePage::defaultOpenDir() const
{
  QSettings settings(QStringLiteral("Manteia"), QStringLiteral("DoseCompare"));
  const QString last = settings.value(QStringLiteral("dvhLastOpenDir")).toString();
  if (!last.isEmpty() && QDir(last).exists())
    return last;
  return QCoreApplication::applicationDirPath();
}

void DvhComparePage::rememberOpenPath(const QString& path)
{
  const QString dir = QFileInfo(path).absolutePath();
  if (dir.isEmpty())
    return;
  QSettings settings(QStringLiteral("Manteia"), QStringLiteral("DoseCompare"));
  settings.setValue(QStringLiteral("dvhLastOpenDir"), dir);
}

QString DvhComparePage::resolveAlgoDllPath() const
{
  const QString base = QCoreApplication::applicationDirPath() + QStringLiteral("/algo/");
#ifdef _DEBUG
  const QString name = QStringLiteral("calcDVHLibraryD.dll");
#else
  const QString name = QStringLiteral("calcDVHLibrary.dll");
#endif
  return base + name;
}

void DvhComparePage::appendLog(const QString& msg)
{
  m_log->append(msg);
}

void DvhComparePage::onBrowseImage()
{
  const QString path = QFileDialog::getExistingDirectory(
    this, tr("选择图像目录（DICOM 序列）"), defaultOpenDir());
  if (!path.isEmpty()) {
    m_imgEdit->setText(QDir::toNativeSeparators(path));
    rememberOpenPath(path);
    return;
  }
  const QString file = QFileDialog::getOpenFileName(
    this, tr("选择图像文件"), defaultOpenDir(),
    tr("Image (*.mhd *.mha *.nii *.nii.gz);;All (*.*)"));
  if (file.isEmpty())
    return;
  m_imgEdit->setText(QDir::toNativeSeparators(file));
  rememberOpenPath(file);
}

void DvhComparePage::onBrowseRts()
{
  const QString file = QFileDialog::getOpenFileName(
    this, tr("选择 RT Structure"), defaultOpenDir(),
    tr("DICOM (*.dcm);;All (*.*)"));
  if (file.isEmpty())
    return;
  m_rtsEdit->setText(QDir::toNativeSeparators(file));
  rememberOpenPath(file);
}

void DvhComparePage::onAddDose()
{
  if (m_doseTable->rowCount() >= 5) {
    QMessageBox::information(this, tr("提示"), tr("最多支持 5 个剂量计划。"));
    return;
  }
  const QString file = QFileDialog::getOpenFileName(
    this, tr("选择剂量文件"), defaultOpenDir(),
    tr("Dose (*.dcm *.mhd *.mha *.nii *.nii.gz);;All (*.*)"));
  if (file.isEmpty())
    return;
  rememberOpenPath(file);
  const int row = m_doseTable->rowCount();
  m_doseTable->insertRow(row);
  m_doseTable->setItem(row, 0, new QTableWidgetItem(QStringLiteral("Plan%1").arg(row + 1)));
  m_doseTable->setItem(row, 1, new QTableWidgetItem(QDir::toNativeSeparators(file)));
}

void DvhComparePage::onRemoveDose()
{
  const auto rows = m_doseTable->selectionModel()->selectedRows();
  for (int i = rows.size() - 1; i >= 0; --i)
    m_doseTable->removeRow(rows[i].row());
}

void DvhComparePage::freeResults()
{
  if (!m_strucSet) {
    m_doseNum = 0;
    return;
  }
  for (int gid = 0; gid < m_doseNum; ++gid) {
    if (!m_strucSet[gid].struc)
      continue;
    for (int sid = 0; sid < m_strucSet[gid].strucNum; ++sid) {
      Struc* s = m_strucSet[gid].struc + sid;
      if (s->h_innDat) {
        free(s->h_innDat);
        s->h_innDat = nullptr;
      }
    }
    delete[] m_strucSet[gid].struc;
    m_strucSet[gid].struc = nullptr;
  }
  delete[] m_strucSet;
  m_strucSet = nullptr;
  m_doseNum = 0;
  m_keptDoses.clear();
  m_keptMasks.clear();
}

void DvhComparePage::onCalculate()
{
  m_log->clear();
  m_plot->setCurves({});
  freeResults();
  m_progress->setValue(0);

  const QString imgPath = m_imgEdit->text().trimmed();
  const QString rtsPath = m_rtsEdit->text().trimmed();
  if (imgPath.isEmpty() || rtsPath.isEmpty()) {
    QMessageBox::warning(this, tr("警告"), tr("请先指定图像与 RT Structure 路径。"));
    return;
  }

  QVector<QPair<QString, QString>> doses;
  for (int r = 0; r < m_doseTable->rowCount(); ++r) {
    const QString tag = m_doseTable->item(r, 0) ? m_doseTable->item(r, 0)->text().trimmed() : QString();
    const QString path = m_doseTable->item(r, 1) ? m_doseTable->item(r, 1)->text().trimmed() : QString();
    if (path.isEmpty())
      continue;
    doses.push_back({tag.isEmpty() ? QStringLiteral("Plan%1").arg(r + 1) : tag, path});
  }
  if (doses.isEmpty()) {
    QMessageBox::warning(this, tr("警告"), tr("请至少添加一个剂量文件路径。"));
    return;
  }
  if (doses.size() > 5) {
    QMessageBox::warning(this, tr("警告"), tr("剂量计划数量不能超过 5。"));
    return;
  }

  // Let user pick dose file if path cell empty was skipped; also allow browse via double-click later.
  // If path looks incomplete, try open dialog for empty ones already skipped.

  const QString dllPath = resolveAlgoDllPath();
  QLibrary lib(dllPath);
  if (!lib.load()) {
    QMessageBox::warning(this, tr("警告"),
                         tr("无法加载 DVH 算法库：\n%1\n%2").arg(dllPath, lib.errorString()));
    m_progress->setValue(10);
    return;
  }

  typedef int (*DvhAlgoFunc)(const float* ofs, const float* spc, const int* scz,
                             const float* dose, StrucSet& sSet);
  auto* dvhAlgoFunc = reinterpret_cast<DvhAlgoFunc>(lib.resolve("calcStrucDVH"));
  if (!dvhAlgoFunc) {
    QMessageBox::warning(this, tr("警告"), tr("DLL 中找不到 calcStrucDVH。"));
    lib.unload();
    return;
  }

  appendLog(tr("已加载 DLL。"));
  m_progress->setValue(1);

  int ret = 0;
  itk::Image<float, 3>::Pointer img;
  const std::string imgStd = imgPath.toLocal8Bit().constData();
  if (isMetaPath(imgPath))
    ret = itkMetaDataReader(imgStd, img);
  else
    ret = itkDicomSeriesReader(imgStd, img);

  if (ret < 0 || !img) {
    appendLog(tr("加载图像失败: %1").arg(ret));
    QMessageBox::warning(this, tr("错误"), tr("加载图像失败。"));
    lib.unload();
    return;
  }
  appendLog(tr("图像加载成功。"));
  m_progress->setValue(3);

  std::vector<itk::Image<float, 3>::Pointer> doseGroup;
  for (const auto& d : doses) {
    itk::Image<float, 3>::Pointer dose;
    const std::string p = d.second.toLocal8Bit().constData();
    if (isMetaPath(d.second))
      ret = itkMetaDataReader(p, dose);
    else if (isDcmPath(d.second))
      ret = itkDcmDoseReader(p, dose);
    else
      ret = -1;
    if (ret < 0 || !dose) {
      appendLog(tr("加载剂量失败: %1").arg(d.second));
      QMessageBox::warning(this, tr("错误"), tr("加载剂量失败：%1").arg(d.second));
      lib.unload();
      return;
    }
    doseGroup.push_back(dose);
  }
  m_keptDoses = doseGroup;
  appendLog(tr("剂量加载成功（%1）。").arg(doseGroup.size()));
  m_progress->setValue(5);

  std::vector<std::string> nameVec;
  std::vector<itk::Image<unsigned char, 3>::Pointer> maskVec;
  ret = itkRtStructureReader(rtsPath.toLocal8Bit().constData(), img, nameVec, maskVec);
  if (ret < 0 || maskVec.empty()) {
    appendLog(tr("加载 RT Structure 失败: %1").arg(ret));
    QMessageBox::warning(this, tr("错误"), tr("加载 RT Structure 失败。"));
    lib.unload();
    return;
  }
  m_keptMasks = maskVec;
  appendLog(tr("结构加载成功（%1 个 ROI）。").arg(static_cast<int>(maskVec.size())));
  m_progress->setValue(7);
  appendLog(tr("正在计算 DVH…"));

  const int doseGroupSz = static_cast<int>(doseGroup.size());
  m_doseNum = doseGroupSz;
  m_strucSet = new StrucSet[doseGroupSz];
  std::memset(m_strucSet, 0, sizeof(StrucSet) * doseGroupSz);

  for (int gid = 0; gid < doseGroupSz; ++gid) {
    m_strucSet[gid].tagName = doses[gid].first.toStdString();
    m_strucSet[gid].strucNum = static_cast<int>(maskVec.size());
    m_strucSet[gid].struc = new Struc[m_strucSet[gid].strucNum];
    std::memset(m_strucSet[gid].struc, 0, sizeof(Struc) * m_strucSet[gid].strucNum);

    for (int sid = 0; sid < static_cast<int>(maskVec.size()); ++sid) {
      Struc* struc = m_strucSet[gid].struc + sid;
      struc->strucName = nameVec[sid];
      struc->h_dat = maskVec[sid]->GetBufferPointer();
      for (int i = 0; i < 3; ++i) {
        struc->origin[i] = static_cast<float>(maskVec[sid]->GetOrigin()[i]);
        struc->spacing[i] = static_cast<float>(maskVec[sid]->GetSpacing()[i]);
        struc->size[i] = static_cast<int>(maskVec[sid]->GetLargestPossibleRegion().GetSize()[i]);
      }
      struc->datSz = struc->size[0] * struc->size[1] * struc->size[2];
    }

    float dsOrigin[3], dsSpacing[3];
    int dsSize[3];
    for (int i = 0; i < 3; ++i) {
      dsOrigin[i] = static_cast<float>(doseGroup[gid]->GetOrigin()[i]);
      dsSpacing[i] = static_cast<float>(doseGroup[gid]->GetSpacing()[i]);
      dsSize[i] = static_cast<int>(doseGroup[gid]->GetLargestPossibleRegion().GetSize()[i]);
    }
    float* dosePtr = doseGroup[gid]->GetBufferPointer();
    dvhAlgoFunc(dsOrigin, dsSpacing, dsSize, dosePtr, m_strucSet[gid]);

    appendLog(QStringLiteral(".................... %1 ....................").arg(doses[gid].first));
    for (int i = 0; i < m_strucSet[gid].strucNum; ++i) {
      const int datSz = m_strucSet[gid].struc[i].innSz;
      if (datSz <= 0 || !m_strucSet[gid].struc[i].h_innDat)
        continue;
      double sum = 0;
      for (int k = 0; k < datSz; ++k)
        sum += m_strucSet[gid].struc[i].h_innDat[k];
      sum /= datSz;
      const QString name = QString::fromStdString(m_strucSet[gid].struc[i].strucName);
      appendLog(tr("%1 / %2  min=%3  max=%4  ave=%5")
                  .arg(doses[gid].first, name)
                  .arg(m_strucSet[gid].struc[i].h_innDat[0], 0, 'f', 3)
                  .arg(m_strucSet[gid].struc[i].h_innDat[datSz - 1], 0, 'f', 3)
                  .arg(sum, 0, 'f', 3));
    }
  }

  m_progress->setValue(10);
  appendLog(tr("DVH 计算完成。"));
  lib.unload();

  rebuildRoiList();
  refreshPlot();
}

void DvhComparePage::rebuildRoiList()
{
  m_roiList->clear();
  m_roiChecks.clear();
  if (!m_strucSet || m_doseNum <= 0)
    return;

  for (int i = 0; i < m_strucSet[0].strucNum; ++i) {
    auto* item = new QListWidgetItem(m_roiList);
    auto* check = new QCheckBox(QString::fromStdString(m_strucSet[0].struc[i].strucName), m_roiList);
    check->setChecked(true);
    m_roiList->addItem(item);
    m_roiList->setItemWidget(item, check);
    m_roiChecks.push_back(check);
    connect(check, &QCheckBox::toggled, this, &DvhComparePage::onRoiVisibilityChanged);
  }
}

void DvhComparePage::onSelectAllRoi()
{
  for (QCheckBox* c : m_roiChecks)
    c->setChecked(true);
}

void DvhComparePage::onClearAllRoi()
{
  for (QCheckBox* c : m_roiChecks)
    c->setChecked(false);
}

void DvhComparePage::onRoiVisibilityChanged()
{
  refreshPlot();
}

void DvhComparePage::refreshPlot()
{
  QVector<DvhCurve> curves;
  if (!m_strucSet || m_doseNum <= 0) {
    m_plot->setCurves(curves);
    return;
  }

  const auto palette = dvhPalette();
  float maxDose = 0.f;
  QVector<float> sampleBuf(DVH_MAX_SPN);

  for (int gid = 0; gid < m_doseNum; ++gid) {
    for (int i = 0; i < m_strucSet[gid].strucNum; ++i) {
      if (i < m_roiChecks.size() && !m_roiChecks[i]->isChecked())
        continue;
      Struc* struc = m_strucSet[gid].struc + i;
      if (!struc->h_innDat || struc->innSz <= 0)
        continue;

      int dvhSz = struc->innSz;
      float* dvhDat = struc->h_innDat;
      if (dvhSz > DVH_MAX_SPN) {
        sampleBuf[0] = dvhDat[0];
        sampleBuf[DVH_MAX_SPN - 1] = dvhDat[dvhSz - 1];
        const float span = (struc->innSz - 1) * 1.f / (DVH_MAX_SPN - 1);
        for (int k = 1; k < DVH_MAX_SPN - 1; ++k) {
          const float p = k * span;
          const int id0 = static_cast<int>(std::floor(p));
          const int id1 = id0 + 1;
          const float w0 = id1 - p;
          const float w1 = p - id0;
          sampleBuf[k] = w0 * dvhDat[id0] + w1 * dvhDat[id1];
        }
        dvhDat = sampleBuf.data();
        dvhSz = DVH_MAX_SPN;
      }

      DvhCurve curve;
      curve.name = QStringLiteral("%1 (%2)")
                     .arg(QString::fromStdString(struc->strucName))
                     .arg(QString::fromStdString(m_strucSet[gid].tagName));
      const Rgb rgb = palette[i % palette.size()];
      curve.color = QColor(rgb.r, rgb.g, rgb.b);
      curve.style = dosePenStyle(gid);
      curve.points.reserve(dvhSz + 1);
      curve.points.push_back(QPointF(0.0, 1.0));
      for (int j = 0; j < dvhSz; ++j) {
        curve.points.push_back(QPointF(dvhDat[j], 1.0 - (j + 1) / (1.0 * dvhSz)));
        if (dvhDat[j] > maxDose)
          maxDose = dvhDat[j];
      }
      curves.push_back(curve);
    }
  }

  m_plot->setXMax(maxDose > 0 ? maxDose * 1.2 : 1.0);
  m_plot->setCurves(curves);
}