// File: host/src/GraphWidget.cpp
#include "GraphWidget.h"
#include <QDialogButtonBox>
#include <QFormLayout>
#include <QInputDialog>
#include <QMenu>
#include <cmath>
#include <limits>

// Simple Linear Regression Helper
static void linearRegression(const QVector<double> &x, const QVector<double> &y,
                             double &m, double &c) {
  double sumX = 0, sumY = 0, sumXY = 0, sumX2 = 0;
  int n = x.size();
  for (int i = 0; i < n; i++) {
    sumX += x[i];
    sumY += y[i];
    sumXY += x[i] * y[i];
    sumX2 += x[i] * x[i];
  }
  m = (n * sumXY - sumX * sumY) / (n * sumX2 - sumX * sumX);
  c = (sumY - m * sumX) / n;
}

GraphWidget::GraphWidget(QWidget *parent) : QWidget(parent) {
  mainLayout = new QVBoxLayout(this);
  mainLayout->setContentsMargins(0, 0, 0, 0);
  mainLayout->setSpacing(0);

  // Toolbar
  toolbar = new QWidget(this);
  toolbar->setStyleSheet("background-color: #2D2D2D;");
  QHBoxLayout *toolLayout = new QHBoxLayout(toolbar);
  toolLayout->setContentsMargins(5, 2, 5, 2);

  btnScaleX = new QPushButton("Scale X", this);
  btnScaleY = new QPushButton("Scale Y", this);
  btnScaleBoth = new QPushButton("Scale Both", this);
  btnCenter = new QPushButton("Center", this);
  btnAnalyze = new QPushButton("Analyze", this);

  QString btnStyle = "QPushButton { background-color: #444; color: white; "
                     "border: 1px solid #555; padding: 3px 8px; border-radius: "
                     "3px; } QPushButton:hover { background-color: #555; }";
  btnScaleX->setStyleSheet(btnStyle);
  btnScaleY->setStyleSheet(btnStyle);
  btnScaleBoth->setStyleSheet(btnStyle);
  btnCenter->setStyleSheet(btnStyle);
  btnAnalyze->setStyleSheet(btnStyle);

  toolLayout->addWidget(btnScaleX);
  toolLayout->addWidget(btnScaleY);
  toolLayout->addWidget(btnScaleBoth);
  toolLayout->addWidget(btnCenter);
  toolLayout->addStretch();
  toolLayout->addWidget(btnAnalyze);

  mainLayout->addWidget(toolbar);

  plot = new QCustomPlot(this);
  mainLayout->addWidget(plot);

  plot->setBackground(QBrush(QColor(25, 25, 25)));

  auto styleAxis = [](QCPAxis *axis) {
    axis->setBasePen(QPen(Qt::white));
    axis->setTickPen(QPen(Qt::white));
    axis->setSubTickPen(QPen(Qt::white));
    axis->setTickLabelColor(Qt::white);
    axis->setLabelColor(Qt::white);
    axis->grid()->setPen(QPen(QColor(60, 60, 60), 0, Qt::DotLine));
    axis->grid()->setSubGridVisible(true);
    axis->grid()->setSubGridPen(QPen(QColor(40, 40, 40), 0, Qt::DotLine));
  };

  styleAxis(plot->xAxis);
  styleAxis(plot->yAxis);
  styleAxis(plot->yAxis2);

  // --- Setup Graphs ---
  graphReal = plot->addGraph();
  graphReal->setPen(QPen(QColor(0, 255, 255), 2));
  graphReal->setLineStyle(QCPGraph::lsLine);
  graphReal->setScatterStyle(
      QCPScatterStyle(QCPScatterStyle::ssCircle, QColor(0, 255, 255), 3));

  graphImag = plot->addGraph(plot->xAxis, plot->yAxis2);
  graphImag->setPen(QPen(QColor(255, 0, 255), 2));
  graphImag->setLineStyle(QCPGraph::lsLine);
  graphImag->setScatterStyle(
      QCPScatterStyle(QCPScatterStyle::ssTriangle, QColor(255, 0, 255), 3));

  graphHilbert = plot->addGraph(plot->xAxis, plot->yAxis2);
  QPen pen3(Qt::green);
  pen3.setWidth(2);
  pen3.setStyle(Qt::DashLine);
  graphHilbert->setPen(pen3);

  graphNyquistCorr = plot->addGraph(plot->xAxis, plot->yAxis);
  graphNyquistCorr->setPen(QPen(QColor(255, 165, 0), 2));
  graphNyquistCorr->setLineStyle(QCPGraph::lsLine);
  graphNyquistCorr->setScatterStyle(
      QCPScatterStyle(QCPScatterStyle::ssCross, 4));
  graphNyquistCorr->setName("De-embedded (True Cell)");

  graphAmp = plot->addGraph();
  graphAmp->setPen(QPen(QColor(50, 255, 50), 2));
  graphAmp->setLineStyle(QCPGraph::lsLine);
  graphAmp->setScatterStyle(QCPScatterStyle(QCPScatterStyle::ssDisc, 3));
  graphAmp->setName("Current");

  graphExtrapolated = plot->addGraph(plot->xAxis, plot->yAxis);
  graphExtrapolated->setLineStyle(QCPGraph::lsNone);
  graphExtrapolated->setScatterStyle(
      QCPScatterStyle(QCPScatterStyle::ssStar, QColor(255, 215, 0), 12));
  graphExtrapolated->setPen(QPen(QColor(255, 215, 0), 3));
  graphExtrapolated->setName("Rs (Extrapolated)");

  graphLSVBlank = plot->addGraph();
  QPen penBlank(QColor(150, 150, 150));
  penBlank.setWidth(2);
  penBlank.setStyle(Qt::DashLine);
  graphLSVBlank->setPen(penBlank);
  graphLSVBlank->setName("Blank (Tap Water)");

  graphLSVSample = plot->addGraph();
  graphLSVSample->setPen(QPen(Qt::yellow, 2));
  graphLSVSample->setName("Sample (Bleach)");

  graphLSVDiff = plot->addGraph();
  graphLSVDiff->setPen(QPen(Qt::cyan, 3));
  graphLSVDiff->setName("Diff (Chlorine)");

  graphFit = plot->addGraph();
  graphFit->setPen(QPen(Qt::red, 2));
  graphFit->setName("Fit");

  graphNyquistRaw = graphReal;

  plot->yAxis2->setVisible(true);
  plot->yAxis2->setTickLabels(true);

  connect(plot->yAxis, SIGNAL(rangeChanged(QCPRange)), plot->yAxis2,
          SLOT(setRange(QCPRange)));
  connect(plot->yAxis2, SIGNAL(rangeChanged(QCPRange)), plot->yAxis,
          SLOT(setRange(QCPRange)));

  plot->setInteractions(QCP::iRangeDrag | QCP::iRangeZoom |
                        QCP::iSelectPlottables);

  // Right click drag to zoom
  plot->axisRect()->setRangeDrag(Qt::Horizontal | Qt::Vertical);
  plot->axisRect()->setRangeZoom(Qt::Horizontal | Qt::Vertical);

  plot->legend->setVisible(true);
  QFont legendFont = font();
  legendFont.setPointSize(9);
  plot->legend->setFont(legendFont);
  plot->legend->setBrush(QBrush(QColor(40, 40, 40, 200)));
  plot->legend->setBorderPen(QPen(Qt::white));
  plot->legend->setTextColor(Qt::white);

  selectionRect = new QCPItemRect(plot);
  selectionRect->setPen(QPen(Qt::yellow, 1, Qt::DashLine));
  selectionRect->setBrush(QBrush(QColor(255, 255, 0, 50)));
  selectionRect->setVisible(false);

  connect(btnScaleX, &QPushButton::clicked, this, &GraphWidget::scaleX);
  connect(btnScaleY, &QPushButton::clicked, this, &GraphWidget::scaleY);
  connect(btnScaleBoth, &QPushButton::clicked, this, &GraphWidget::scaleBoth);
  connect(btnCenter, &QPushButton::clicked, this, &GraphWidget::centerView);
  connect(btnAnalyze, &QPushButton::clicked, this, &GraphWidget::startAnalyze);

  // Custom mouse handling for selection
  connect(plot, &QCustomPlot::mousePress, [this](QMouseEvent *event) {
    if (isSelecting && event->button() == Qt::LeftButton) {
      selStartX = plot->xAxis->pixelToCoord(event->pos().x());
      selStartY = plot->yAxis->pixelToCoord(event->pos().y());
      selectionRect->topLeft->setCoords(selStartX, selStartY);
      selectionRect->bottomRight->setCoords(selStartX, selStartY);
      selectionRect->setVisible(true);
      plot->replot();
    }
  });

  connect(plot, &QCustomPlot::mouseMove, [this](QMouseEvent *event) {
    if (isSelecting && (event->buttons() & Qt::LeftButton)) {
      double x = plot->xAxis->pixelToCoord(event->pos().x());
      double y = plot->yAxis->pixelToCoord(event->pos().y());
      selectionRect->bottomRight->setCoords(x, y);
      plot->replot();
    }
  });

  connect(plot, &QCustomPlot::mouseRelease, [this](QMouseEvent *event) {
    if (isSelecting && event->button() == Qt::LeftButton) {
      isSelecting = false;
      plot->setInteractions(QCP::iRangeDrag | QCP::iRangeZoom);

      double x2 = plot->xAxis->pixelToCoord(event->pos().x());
      double y2 = plot->yAxis->pixelToCoord(event->pos().y());

      // Collect data in rect
      QVector<double> xData, yData;
      double minX = std::min(selStartX, x2);
      double maxX = std::max(selStartX, x2);
      double minY = std::min(selStartY, y2);
      double maxY = std::max(selStartY, y2);

      // Iterate visible graphs to find data
      QList<QCPGraph *> graphs = {graphReal, graphAmp, graphLSVSample,
                                  graphLSVDiff};
      for (auto g : graphs) {
        if (!g->visible())
          continue;
        for (auto it = g->data()->begin(); it != g->data()->end(); ++it) {
          if (it->key >= minX && it->key <= maxX && it->value >= minY &&
              it->value <= maxY) {
            xData.append(it->key);
            yData.append(it->value);
          }
        }
        if (!xData.isEmpty())
          break; // Only fit one graph
      }

      if (xData.size() > 2) {
        // Show Dialog
        QDialog dlg(this);
        dlg.setWindowTitle("Fit Data");
        QFormLayout *layout = new QFormLayout(&dlg);
        QComboBox *type = new QComboBox();
        type->addItems({"Linear", "Exponential", "Logarithmic", "Polynomial"});
        QSpinBox *order = new QSpinBox();
        order->setRange(1, 10);
        order->setValue(2);
        QCheckBox *inverse = new QCheckBox("Inverse");

        layout->addRow("Type:", type);
        layout->addRow("Poly Order:", order);
        layout->addRow(inverse);

        QDialogButtonBox *btns = new QDialogButtonBox(QDialogButtonBox::Ok |
                                                      QDialogButtonBox::Cancel);
        layout->addRow(btns);
        connect(btns, &QDialogButtonBox::accepted, &dlg, &QDialog::accept);
        connect(btns, &QDialogButtonBox::rejected, &dlg, &QDialog::reject);

        if (dlg.exec() == QDialog::Accepted) {
          performFit(xData, yData, type->currentIndex(), order->value(),
                     inverse->isChecked());
        }
      }

      selectionRect->setVisible(false);
      plot->replot();
    }
  });

  configureRawPlot();
}

void GraphWidget::scaleX() {
  plot->rescaleAxes(true);
  plot->replot();
}
void GraphWidget::scaleY() {
  plot->yAxis->rescale(true);
  plot->yAxis2->rescale(true);
  plot->replot();
}
void GraphWidget::scaleBoth() {
  plot->rescaleAxes(true);
  plot->replot();
}
void GraphWidget::centerView() { scaleBoth(); }

void GraphWidget::startAnalyze() {
  isSelecting = true;
  // FIX: Use default constructor for empty flags instead of 0
  plot->setInteractions(QCP::Interactions());
}

void GraphWidget::performFit(const QVector<double> &x, const QVector<double> &y,
                             int type, int order, bool inverse) {
  graphFit->data()->clear();
  QVector<double> xFit, yFit;

  // Simple Linear Fit Implementation for demo
  if (type == 0) { // Linear
    double m, c;
    QVector<double> yProc = y;
    if (inverse) {
      for (int i = 0; i < y.size(); i++)
        yProc[i] = 1.0 / y[i];
    }
    linearRegression(x, yProc, m, c);

    double minX = *std::min_element(x.begin(), x.end());
    double maxX = *std::max_element(x.begin(), x.end());

    for (int i = 0; i <= 100; i++) {
      double xv = minX + (maxX - minX) * i / 100.0;
      double yv = m * xv + c;
      if (inverse)
        yv = 1.0 / yv;
      xFit.append(xv);
      yFit.append(yv);
    }
  }
  // Add other fits here (Exp, Log, Poly) as needed

  graphFit->setData(xFit, yFit);
  graphFit->setVisible(true);
  plot->replot();
}

// ... (Rest of configuration methods same as before) ...
void GraphWidget::configureRawPlot() {
  plot->xAxis->setLabel("Frequency (Hz)");
  plot->xAxis->setScaleType(QCPAxis::stLogarithmic);
  QSharedPointer<QCPAxisTickerLog> logTicker(new QCPAxisTickerLog);
  plot->xAxis->setTicker(logTicker);
  plot->xAxis->setNumberFormat("eb");

  plot->yAxis->setLabel("Magnitude (Ohms)");
  plot->yAxis->setScaleType(QCPAxis::stLogarithmic);
  QSharedPointer<QCPAxisTickerLog> logTickerY(new QCPAxisTickerLog);
  plot->yAxis->setTicker(logTickerY);
  plot->yAxis->setNumberFormat("eb");

  plot->yAxis2->setLabel("Phase (Rad)");
  plot->yAxis2->setScaleType(QCPAxis::stLinear);
  QSharedPointer<QCPAxisTicker> linTicker(new QCPAxisTicker);
  plot->yAxis2->setTicker(linTicker);
  plot->yAxis2->setNumberFormat("f");
  plot->yAxis2->setVisible(true);

  graphReal->setName("Magnitude");
  graphImag->setName("Phase");
  graphHilbert->setName("Hilbert");

  graphReal->setVisible(true);
  graphImag->setVisible(true);
  graphHilbert->setVisible(true);
  graphNyquistCorr->setVisible(false);
  graphAmp->setVisible(false);
  graphExtrapolated->setVisible(false);
  graphLSVBlank->setVisible(false);
  graphLSVSample->setVisible(false);
  graphLSVDiff->setVisible(false);
  graphFit->setVisible(false);

  plot->replot();
}

void GraphWidget::configureNyquistPlot() {
  plot->xAxis->setLabel("Real (Z')");
  plot->xAxis->setScaleType(QCPAxis::stLinear);
  QSharedPointer<QCPAxisTicker> linTicker(new QCPAxisTicker);
  plot->xAxis->setTicker(linTicker);
  plot->xAxis->setNumberFormat("f");

  plot->yAxis->setLabel("-Imaginary (-Z'')");
  plot->yAxis->setScaleType(QCPAxis::stLinear);
  plot->yAxis->setTicker(linTicker);
  plot->yAxis->setNumberFormat("f");

  plot->yAxis2->setVisible(false);

  graphReal->setName("Measured (Raw)");
  graphReal->setLineStyle(QCPGraph::lsLine);
  graphReal->setScatterStyle(QCPScatterStyle(QCPScatterStyle::ssCircle, 4));

  graphReal->setVisible(true);
  graphImag->setVisible(false);
  graphHilbert->setVisible(false);
  graphNyquistCorr->setVisible(true);
  graphAmp->setVisible(false);
  graphExtrapolated->setVisible(true);
  graphLSVBlank->setVisible(false);
  graphLSVSample->setVisible(false);
  graphLSVDiff->setVisible(false);
  graphFit->setVisible(false);

  plot->replot();
}

void GraphWidget::configureAmperometricPlot() {
  plot->xAxis->setLabel("Sample Index");
  plot->xAxis->setScaleType(QCPAxis::stLinear);
  QSharedPointer<QCPAxisTicker> linTicker(new QCPAxisTicker);
  plot->xAxis->setTicker(linTicker);
  plot->xAxis->setNumberFormat("f");

  plot->yAxis->setLabel("Current (uA)");
  plot->yAxis->setScaleType(QCPAxis::stLinear);
  plot->yAxis->setTicker(linTicker);
  plot->yAxis->setNumberFormat("f");

  plot->yAxis2->setVisible(false);

  graphAmp->setName("Current");
  graphAmp->setVisible(true);

  graphReal->setVisible(false);
  graphImag->setVisible(false);
  graphHilbert->setVisible(false);
  graphNyquistCorr->setVisible(false);
  graphExtrapolated->setVisible(false);
  graphLSVBlank->setVisible(false);
  graphLSVSample->setVisible(false);
  graphLSVDiff->setVisible(false);
  graphFit->setVisible(false);

  plot->replot();
}

void GraphWidget::configureLSVPlot() {
  plot->xAxis->setLabel("Voltage (mV)");
  plot->xAxis->setScaleType(QCPAxis::stLinear);
  QSharedPointer<QCPAxisTicker> linTicker(new QCPAxisTicker);
  plot->xAxis->setTicker(linTicker);
  plot->xAxis->setNumberFormat("f");

  plot->yAxis->setLabel("Current (uA)");
  plot->yAxis->setScaleType(QCPAxis::stLinear);
  plot->yAxis->setTicker(linTicker);
  plot->yAxis->setNumberFormat("f");

  plot->yAxis2->setVisible(false);

  graphLSVBlank->setVisible(true);
  graphLSVSample->setVisible(true);
  graphLSVDiff->setVisible(true);

  graphReal->setVisible(false);
  graphImag->setVisible(false);
  graphHilbert->setVisible(false);
  graphNyquistCorr->setVisible(false);
  graphExtrapolated->setVisible(false);
  graphAmp->setVisible(false);
  graphFit->setVisible(false);

  plot->replot();
}

void GraphWidget::addBodeData(double freq, double val1, double val2) {
  if (plot->xAxis->scaleType() == QCPAxis::stLogarithmic && freq <= 0)
    return;
  graphReal->addData(freq, val1);
  graphImag->addData(freq, val2);
  graphReal->rescaleAxes(false);
  graphImag->rescaleAxes(false);
  graphHilbert->rescaleAxes(false);
  plot->replot();
}

void GraphWidget::addNyquistData(double r_meas, double i_meas, double r_corr,
                                 double i_corr, bool showCorr) {
  graphNyquistRaw->addData(r_meas, -i_meas);
  if (showCorr) {
    graphNyquistCorr->addData(r_corr, -i_corr);
  }
  graphNyquistRaw->rescaleAxes(false);
  if (showCorr) {
    graphNyquistCorr->rescaleAxes(true);
  }
  plot->replot();
}

void GraphWidget::addAmperometricData(double index, double current) {
  graphAmp->addData(index, current);
  graphAmp->rescaleAxes(false);
  plot->replot();
}

void GraphWidget::addLSVData(double voltage, double current,
                             LSVTrace traceType) {
  QCPGraph *target = nullptr;
  switch (traceType) {
  case LSV_BLANK:
    target = graphLSVBlank;
    break;
  case LSV_SAMPLE:
    target = graphLSVSample;
    break;
  case LSV_DIFF:
    target = graphLSVDiff;
    break;
  }

  if (target) {
    target->addData(voltage, current);
    target->rescaleAxes(false);
    plot->replot();
  }
}

void GraphWidget::addHilbertData(const QVector<double> &freq,
                                 const QVector<double> &hilbertImag) {
  if (plot->xAxis->label() != "Frequency (Hz)")
    return;
  graphHilbert->setData(freq, hilbertImag);
  graphHilbert->rescaleAxes(false);
  plot->replot();
}

void GraphWidget::setExtrapolatedPoint(double real, double imag) {
  graphExtrapolated->data()->clear();
  graphExtrapolated->addData(real, -imag);
  plot->replot();
}

void GraphWidget::clear() {
  graphReal->data()->clear();
  graphImag->data()->clear();
  graphHilbert->data()->clear();
  graphNyquistCorr->data()->clear();
  graphAmp->data()->clear();
  graphExtrapolated->data()->clear();
  graphFit->data()->clear();
  plot->replot();
}

void GraphWidget::clearLSV(LSVTrace traceType) {
  if (traceType == LSV_BLANK)
    graphLSVBlank->data()->clear();
  if (traceType == LSV_SAMPLE)
    graphLSVSample->data()->clear();
  if (traceType == LSV_DIFF)
    graphLSVDiff->data()->clear();
  plot->replot();
}