from PyQt6.QtWidgets import (
    QMainWindow,
    QLabel,
    QPushButton,
    QFormLayout,
    QWidget,
    QHBoxLayout,
)
from PyQt6 import QtWidgets
import numpy as np
import sympy
from typing import *

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg

matplotlib.use("QtAgg")

import widgets
import task
import algorithm


class ParamsField(widgets.InputWidget):
    _layout: QFormLayout
    fields: Dict[str, widgets.InputWidget]

    callback_internal: Callable

    def __init__(self, parent, fields: List[str]):
        super().__init__(parent)

        self._initialize()
        self._create(fields)

        super().init()

    def _initialize(self):
        self.fields = {}
        self._layout = QFormLayout(self)
        self._layout.setContentsMargins(0, 0, 0, 0)

    def _create(self, fields: List[str]):
        fields = sorted(fields)

        for field in fields:
            val = "0.0"
            if not self.value.get(field) is None:
                val = str(self.value.get(field))

            self.fields[field] = widgets.InputField(
                self, label_text=f"{field:>4}: ", lineedit_text=val
            )
            self.fields[field].set_callback(
                widgets.make_callback_chain(
                    [widgets.float_check_callback, self._callback]
                )
            )
            self._layout.addRow(self.fields[field])

    def recreate(self, fields: List[str]):
        for field in self.fields:
            self.fields[field].deleteLater()

        self.fields = {}
        self._create(fields)
        self.callback_internal()

    def _get_value_internal(self) -> Dict[str, float]:
        value = {}
        for field in self.fields:
            value[field] = float(self.fields[field].get_value())

        return value

    def _set_value_internal(self, value: Dict[str, float]):
        for field in self.fields:
            self.fields[field].set_value(value[field])

    def _set_callback_internal(self, callback):
        self.callback_internal = callback

    def _callback(self, widget: widgets.InputWidget):
        self.callback_internal()
        return True


class PointsField(widgets.InputWidget):
    _layout: QFormLayout
    points: List[widgets.InputField]

    min: float
    max: float

    DIVISIONS = 1_000_000

    callback_internal: Callable[[None], None]

    def __init__(self, parent, start: int, points: List[float], min: float, max: float):
        super().__init__(parent)

        self._initialize()
        self._create(start, points, min, max)

        super().init()

    def _initialize(self):
        self.points = []
        self.min = 0.0
        self.max = 1.0
        self._layout = QFormLayout(self)
        self._layout.setContentsMargins(0, 0, 0, 0)

    def _create(self, start: int, points: List[float], min: float, max: float):
        self.min = min
        self.max = max

        for i in range(0, len(points)):
            value = int((points[i] - min) / (max - min) * self.DIVISIONS)
            self.points.append(
                widgets.InputSlider(
                    self,
                    label_text=f"{i + start:>4}: ",
                    value=value,
                    min_value=0,
                    max_value=self.DIVISIONS,
                )
            )
            self.points[i].set_callback(
                widgets.make_callback_chain(
                    [widgets.float_check_callback, self._callback]
                )
            )
            self._layout.addRow(self.points[i])

    def recreate(self, start: int, points: List[float], min: float, max: float):
        for point in self.points:
            point.deleteLater()

        self.points = []
        self._create(start, points, min, max)
        self.callback_internal()

    def _get_value_internal(self) -> List[float]:
        value = []
        for point in self.points:
            value.append(
                float(point.get_value()) / self.DIVISIONS * (self.max - self.min)
                + self.min
            )

        return value

    def _set_value_internal(self, value: List[float]):
        for i in range(0, len(value)):
            self.points[i].set_value(
                (value[i] - self.min) / (self.max - self.min) * self.DIVISIONS
            )

    def _set_callback_internal(self, callback):
        self.callback_internal = callback

    def _callback(self, widget: widgets.InputWidget):
        self.callback_internal()
        return True


class PlotWidget(QWidget):
    figure: Figure
    axes: Axes
    canvas: FigureCanvasQTAgg

    pre_style: Callable[[Axes], None]
    post_style: Callable[[Axes], None]
    plots: List[Callable[[Axes], None]]

    def __init__(self, parent):
        super().__init__(parent)

        self._layout = QHBoxLayout(self)

        self.figure = plt.figure()
        self.axes = self.figure.add_subplot(111)

        self.figure.tight_layout(pad=3)
        self.figure.subplots_adjust(left=0.15, right=0.975, top=0.95)
        self.figure.patch.set_facecolor("#f0f0f0")

        self.canvas = FigureCanvasQTAgg(self.figure)
        self._layout.addWidget(self.canvas)

        self.pre_style = lambda _: None
        self.plots = []

    def set_pre_style(self, style: Callable[[Axes], None]):
        self.pre_style = style

    def set_post_style(self, style: Callable[[Axes], None]):
        self.post_style = style

    def clear_plots(self):
        self.plots = []

    def add_plot(self, plot: Callable[[Axes], None]):
        self.plots.append(plot)

    def redraw(self):
        self.axes.clear()
        self.pre_style(self.axes)
        for plot in self.plots:
            plot(self.axes)
        self.post_style(self.axes)

        self.canvas.draw()


class MainWindow(QMainWindow):
    central_widget: QWidget

    function_field: widgets.InputField
    function_param_fields: ParamsField

    x_0_field: widgets.InputField
    x_n_field: widgets.InputField

    n_display: widgets.DisplayField
    n_slider: widgets.InputSlider

    points_field: PointsField

    margin_display: widgets.DisplayField
    margin_slider: widgets.InputSlider

    function_plot: PlotWidget
    error_plot: PlotWidget

    function: Callable[[float], float]
    interpolated_function: Callable[[float], float]
    function_params: Dict[str, float]

    x_0: float
    x_n: float
    n: int
    points: List[float]

    margin: float

    parsing_function: bool
    updating: bool

    def __init__(self):
        super().__init__()

        self.n = 2
        self.margin = 5
        self.x_0 = 0.0
        self.x_n = 1.0
        self.points = [0.5]
        self.function = None
        self.function_params = {}
        self.parsing_function = False

        self.setup_ui()
        self.reset()

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

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

        layout = QHBoxLayout(self.central_widget)

        left_row = QWidget(self.central_widget)
        left_row_layout = QFormLayout(left_row)
        right_row = QWidget(self.central_widget)
        right_row_layout = QFormLayout(right_row)

        self.function_field = widgets.InputField(self, label_text="Функция: ")
        self.function_field.set_callback(self._function_field_changed)
        left_row_layout.addWidget(self.function_field)

        self.function_param_fields = ParamsField(self, [])
        self.function_param_fields.set_callback(self._function_param_field_changed)
        left_row_layout.addWidget(self.function_param_fields)

        self.x_0_field = widgets.InputField(
            self, label_text="x_0: ", lineedit_text="0.0"
        )
        self.x_0_field.set_callback(
            widgets.make_callback_chain(
                [widgets.float_check_callback, self._x_0_field_changed]
            )
        )
        left_row_layout.addWidget(self.x_0_field)

        self.x_n_field = widgets.InputField(
            self, label_text="x_n: ", lineedit_text="1.0"
        )
        self.x_n_field.set_callback(
            widgets.make_callback_chain(
                [widgets.float_check_callback, self._x_n_field_changed]
            )
        )
        left_row_layout.addWidget(self.x_n_field)

        self.n_display = widgets.DisplayField(self, "n: ", "")
        left_row_layout.addWidget(self.n_display)

        self.n_slider = widgets.InputSlider(self, "", 1, 1, 20, 1, True)
        self.n_slider.set_callback(self._n_slider_changed)
        left_row_layout.addWidget(self.n_slider)

        self.points_field = PointsField(self, 1, [0.5], 0.0, 1.0 - 0.001)
        self.points_field.set_callback(self._points_field_changed)
        left_row_layout.addWidget(self.points_field)

        self.margin_display = widgets.DisplayField(self, "отступ (%): ", "")
        left_row_layout.addWidget(self.margin_display)

        self.margin_slider = widgets.InputSlider(self, "", 5, 1, 500, 1)
        self.margin_slider.set_callback(self._margin_slider_changed)
        left_row_layout.addWidget(self.margin_slider)

        reset_button = QPushButton("Сброс", self)
        reset_button.clicked.connect(self.reset)
        left_row_layout.addWidget(reset_button)

        label1 = QLabel("График функции", self)
        right_row_layout.addWidget(label1)
        self.function_plot = PlotWidget(self)
        self.function_plot.set_pre_style(self.function_plot_pre_style)
        self.function_plot.set_post_style(self.function_plot_post_style)
        right_row_layout.addWidget(self.function_plot)

        label2 = QLabel("График ошибки", self)
        right_row_layout.addWidget(label2)
        self.error_plot = PlotWidget(self)
        self.error_plot.set_pre_style(self.function_plot_pre_style)
        self.error_plot.set_post_style(self.function_plot_post_style)
        right_row_layout.addWidget(self.error_plot)

        left_row.setLayout(left_row_layout)
        right_row.setLayout(right_row_layout)

        layout.addWidget(left_row)
        layout.addWidget(right_row)

        self.central_widget.setLayout(layout)

    def function_plot_pre_style(self, axes: Axes):
        axes.set_xlabel("x")
        axes.set_ylabel("y")
        axes.axhline(0, color="black", linewidth=1)
        axes.axvline(0, color="black", linewidth=1)
        axes.grid(True)

    def function_plot_post_style(self, axes: Axes):
        axes.legend(loc="upper left")

    def reset(self):
        self.function_field.set_value(task.f_str())
        self.x_0_field.set_value(0.0)
        self.x_n_field.set_value(1.0)
        self.n_slider.set_value(5)
        self.margin_slider.set_value(300)

        QtWidgets.QApplication.processEvents()

        self.function_params = task.f_args()
        self.function_param_fields.set_value(self.function_params)

        QtWidgets.QApplication.processEvents()

        self.redraw_plots()

        QtWidgets.QApplication.processEvents()

    def redraw_plots(self):
        length = abs(self.x_n - self.x_0)
        x_margin = length / 2 * self.margin
        x_min = self.x_0 - x_margin
        x_max = self.x_n + x_margin

        visible_length = abs(x_max - x_min)
        plot_x_margin = visible_length / 2 * 0.025

        steps = max(
            min(int(500 * 1 / visible_length**1.1), 5000),
            min(int(750 * visible_length**1.5), 10000),
        )
        x_values = np.linspace(x_min, x_max, steps)
        y_values = self.function(x_values)
        y_interpolated = self.interpolated_function(x_values)

        self.function_plot.clear_plots()

        self.function_plot.add_plot(
            lambda axes: axes.set_xlim(x_min - plot_x_margin, x_max + plot_x_margin)
        )

        self.function_plot.add_plot(
            lambda axes: axes.plot(x_values, y_values, label="y = f(x)", color="red")
        )
        self.function_plot.add_plot(
            lambda axes: axes.plot(
                x_values,
                y_interpolated,
                label="y = L(x)",
                color="blue",
                linestyle="dashed",
            )
        )

        points = [self.x_0, self.x_n]
        points[1:1] = self.points
        points = np.array(points)

        self.function_plot.add_plot(
            lambda axes: axes.scatter(
                points, self.function(points), color="green", zorder=10
            )
        )

        def label_func(axes):
            for i in range(len(points)):
                axes.text(
                    points[i],
                    self.function(points[i]) + 0.05,
                    str(i),
                    ha="center",
                    va="bottom",
                    color="green",
                )

        self.function_plot.add_plot(label_func)

        self.error_plot.clear_plots()

        self.error_plot.add_plot(
            lambda axes: axes.set_xlim(x_min - plot_x_margin, x_max + plot_x_margin)
        )

        self.error_plot.add_plot(
            lambda axes: axes.plot(
                x_values,
                np.abs(y_values - y_interpolated),
                color="blue",
                label="|f(x) - L(x)|",
            )
        )

        self.function_plot.redraw()
        self.error_plot.redraw()

    def parse_function(self, func_str: str) -> Callable[[float], float]:
        func = sympy.sympify(func_str)
        params = []
        for name in func.atoms(sympy.Symbol):
            params.append(str(name))

        if len(params) == 0 or params.count("x") < 1:
            raise ValueError("No x in function")

        func_lambda = sympy.lambdify(params, func, cse=True)

        if func_lambda.__code__.co_argcount != len(params):
            raise ValueError("Bad arguments count")

        params = [p for p in params if p != "x"]
        old_params = self.function_params
        if set(old_params.keys()) != set(params):
            self.recreate_params_field(params)

        def f(x: float) -> float:
            return func_lambda(x=x, **self.function_params)

        return f

    def update_function(self):
        if self.parsing_function:
            return

        try:
            self.recalculate_interpolated_function()
            self.redraw_plots()
        except Exception as e:
            print(e)

    def recalculate_interpolated_function(self):
        points = [self.x_0, self.x_n]
        points[1:1] = self.points
        points = np.array(points)

        self.interpolated_function = algorithm.make_forward_interpolate_function(
            self.function, points
        )

        # self.interpolated_function = algorithm.make_backward_interpolate_function(
        #     self.function, points
        # )

    def recreate_params_field(self, params: List[str]):
        self.function_param_fields.recreate(params)

    def recreate_points_field(self):
        points = np.linspace(self.x_0, self.x_n, self.n + 1)
        points = [float(p) for p in points[1:-1]]
        self.points_field.recreate(
            1,
            points,
            self.x_0 + 1 / self.points_field.DIVISIONS,
            self.x_n - 1 / self.points_field.DIVISIONS,
        )

        self.points = points

    def _function_field_changed(self, w: widgets.InputWidget) -> bool:
        try:
            self.parsing_function = True
            self.function = self.parse_function(w.get_value())
            self.parsing_function = False
        except Exception as e:
            print(f"{e}")
            self.parsing_function = False
            return False

        self.update_function()

        return True

    def _function_param_field_changed(self, w: widgets.InputWidget) -> bool:
        self.function_params = self.function_param_fields.get_value()

        self.update_function()
        return True

    def _x_0_field_changed(self, w: widgets.InputWidget) -> bool:
        value = float(w.get_value())
        if value == self.x_n:
            w.value = w.previous_value
            return False

        self.x_0 = value

        self.recreate_points_field()

        self.update_function()
        return True

    def _x_n_field_changed(self, w: widgets.InputWidget) -> bool:
        value = float(w.get_value())
        if value == self.x_0:
            w.value = w.previous_value
            return False

        self.x_n = value

        self.recreate_points_field()

        self.update_function()
        return True

    def _n_slider_changed(self, w: widgets.InputWidget) -> bool:
        self.n = int(w.get_value())
        self.n_display.set_text(str(self.n))

        self.recreate_points_field()

        # self.update_function()
        return True

    def _points_field_changed(self, w: widgets.InputWidget) -> bool:
        self.points = w.get_value()

        self.update_function()
        return True

    def _margin_slider_changed(self, w: widgets.InputWidget) -> bool:
        self.margin = (float(w.get_value()) - 250) / 250.0
        self.margin_display.set_text(f"{self.margin * 100:.2f}%")

        self.redraw_plots()
        return True
