import task
from algorithm import *

import sys
from PyQt6.QtWidgets import (
    QMessageBox,
    QMainWindow,
    QLabel,
    QComboBox,
    QPushButton,
    QFormLayout,
    QWidget,
    QLineEdit,
    QHBoxLayout,
)
from PyQt6.QtCore import Qt
import numpy as np
import scipy

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg

matplotlib.use("QtAgg")


class InputField(QWidget):
    def __init__(self, parent, label_text, placeholder_text, lineedit_text):
        super().__init__(parent)

        self._layout = QHBoxLayout(self)
        self._layout.setContentsMargins(0, 0, 0, 0)
        self.label = QLabel(label_text)
        self.lineedit = QLineEdit(lineedit_text)
        self.lineedit.setPlaceholderText(placeholder_text)

        self._layout.addWidget(self.label)
        self._layout.addWidget(self.lineedit)

    def get_value(self):
        return self.lineedit.text()

    def set_value(self, value):
        self.lineedit.setText(value)

    def on_change(self, func):
        self.lineedit.textChanged.connect(func)


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        self.a = float(-3)
        self.b = float(2)
        self.steps = int(1000)

        self.setupUI()

    def setupUI(self):
        self.setWindowTitle("Лабораторная работа №2")
        self.setFixedSize(800, 800)

        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)

        layout = QFormLayout()

        self.label1 = QLabel("Вариант 6", self)
        layout.addWidget(self.label1)
        self.label5 = QLabel(f"Уравнение {task.f_str()}", self)
        layout.addWidget(self.label5)

        self.label3 = QLabel("Диапазон:", self)
        layout.addWidget(self.label3)

        self.aInput = InputField(self, " a: ", "от", str(self.a))
        self.aInput.on_change(self.aInput_changed)
        layout.addWidget(self.aInput)

        self.bInput = InputField(self, " b: ", "до", str(self.b))
        self.bInput.on_change(self.bInput_changed)
        layout.addWidget(self.bInput)

        self.stepsInput = InputField(self, " шаги: ", "шаги", str(self.steps))
        self.stepsInput.on_change(self.stepsInput_changed)
        layout.addWidget(self.stepsInput)

        self.plot = plt.figure()
        self.canvas = FigureCanvasQTAgg(self.plot)
        self.plot_axes = self.plot.add_subplot(111)
        self.drawPlot(self.a, self.b)
        layout.addWidget(self.canvas)

        self.findIntervalsButton = QPushButton("Найти интервалы", self)
        self.findIntervalsButton.clicked.connect(self.findIntervalsButton_clicked)
        layout.addWidget(self.findIntervalsButton)

        self.intervals = []
        self.comboBoxIntervals = QComboBox(self)
        layout.addWidget(self.comboBoxIntervals)

        self.label2 = QLabel("Способ решения:", self)
        layout.addWidget(self.label2)

        self.comboBoxSolutions = QComboBox(self)
        self.comboBoxSolutions.addItem("Половинного деления")
        self.comboBoxSolutions.addItem("Хорд")
        self.comboBoxSolutions.addItem("Ньютона")
        self.comboBoxSolutions.addItem("Ньютона (мод)")
        self.comboBoxSolutions.addItem("Комбинированный")
        self.comboBoxSolutions.addItem("Итераций")
        layout.addWidget(self.comboBoxSolutions)

        self.buttonFindSolution = QPushButton("Решить", self)
        self.buttonFindSolution.clicked.connect(self.findSolutionButton_clicked)
        layout.addWidget(self.buttonFindSolution)

        self.label4 = QLabel("Ответ:", self)
        layout.addWidget(self.label4)

        self.labelResult = QLabel(self)
        layout.addWidget(self.labelResult)

        self.central_widget.setLayout(layout)

    def aInput_changed(self):
        try:
            a_str = self.aInput.get_value()
            if not (a_str == "" or a_str == "-"):
                self.a = float(a_str)
                self.drawPlot(self.a, self.b)
        except:
            self.aInput.set_value(str(self.a))

    def bInput_changed(self):
        try:
            b_str = self.bInput.get_value()
            if not (b_str == "" or b_str == "-"):
                self.b = float(b_str)
                self.drawPlot(self.a, self.b)
        except:
            self.bInput.set_value(str(self.b))

    def stepsInput_changed(self):
        try:
            steps_str = self.stepsInput.get_value()
            if not (steps_str == ""):
                self.steps = int(steps_str)
                self.steps = max(2, min(10000, self.steps))
                self.stepsInput.set_value(str(self.steps))
                self.drawPlot(self.a, self.b)
        except:
            self.stepsInput.set_value(str(self.steps))

    def redraw(self):
        self.canvas.draw()

    def findIntervalsButton_clicked(self):
        self.intervals = self.find_root_intervals(
            self.a, self.b, abs(self.b - self.a) / self.steps, task.f
        )
        self.comboBoxIntervals.clear()
        for a, b in self.intervals:
            self.comboBoxIntervals.addItem(f"({a:.5f}, {b:.5f})")

    def findSolutionButton_clicked(self):
        try:
            (a, b) = self.intervals[self.comboBoxIntervals.currentIndex()]
            solution = self.find_root(a, b, self.comboBoxSolutions.currentIndex())

            if solution is None:
                self.labelResult.setText("Решение не найдено")
            else:
                self.labelResult.setText(f"{solution:.12f}")
        except:
            self.labelResult.setText("Интервал не выбран")

    def drawPlot(self, a, b):
        steps = self.steps
        step = abs(b - a) / steps

        intervals = self.find_root_intervals(a, b, step, task.f)
        roots = []
        for interval in intervals:
            root = scipy.optimize.fsolve(task.f, interval)
            roots.extend(root)

        x_values = np.linspace(a, b, steps)
        y_values = task.f(x_values)
        ax = self.plot_axes
        ax.clear()
        ax.plot(x_values, y_values)
        ax.axhline(0, color="black", linewidth=1)
        ax.set_xlabel("x")
        ax.axvline(0, color="black", linewidth=1)
        ax.set_ylabel("y")
        ax.grid(True)
        ax.plot(x_values, y_values, label="y = f(x)")

        ax.scatter(roots, np.zeros_like(roots), color="red")
        for root in roots:
            ax.text(root, 0, f"x={root:.5f}", ha="left", va="bottom")

        self.redraw()

    def find_root(self, a, b, method):
        solution = None
        match method:
            case 0:
                solution = half_div_method(task.f, a, b)
            case 1:
                solution = chord_method(task.f, task.f_dd, a, b)
            case 2:
                solution = newtons_method(task.f, task.f_d, a, b)
            case 3:
                solution = modified_newtons_method(task.f, task.f_d, a, b)
            case 4:
                solution = combined_method(task.f, task.f_d, a, b)
            case 5:
                solution = iterative_method(task.f, task.f_d, a, b)

        return solution

    def find_root_intervals(self, a, b, step, f):
        x_range = np.array(np.arange(a, b, step), dtype=np.float64)
        y_range = f(x_range)
        solution_range = []
        for i in range(len(x_range) - 1):
            if y_range[i] * y_range[i + 1] < 0:
                solution_range.append((x_range[i], x_range[i + 1]))
        return solution_range
