import matplotlib
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel

matplotlib.use('QT5Agg')
import matplotlib.pyplot as plt

from models.data_key import DataKey
from utils import ui_utils


class AgeResultsWidget(QWidget):
    def __init__(self, results_dialog):
        QWidget.__init__(self)

        self.results_dialog = results_dialog

        layout = QHBoxLayout()
        layout.addLayout(self._create_widget())
        self.setLayout(layout)

        results_dialog.sample_tree.tree.currentItemChanged.connect(lambda i, j: self.replot_graph())
        results_dialog.configuration_changed.connect(self.replot_graph)

    def _create_widget(self):
        layout = QVBoxLayout()
        layout.addWidget(QLabel("Sample and spot name"))
        layout.addWidget(self._create_age_graph_and_point_selection())

        return layout

    def _create_age_graph_and_point_selection(self):

        graph_and_points = QWidget()
        layout = QVBoxLayout()

        fig = plt.figure()
        self.axes = plt.axes()

        graph_widget, self.canvas = ui_utils.create_figure_widget(fig, self)

        layout.addWidget(graph_widget)

        graph_and_points.setLayout(layout)

        return graph_and_points

    ###############
    ### Actions ###
    ###############

    def replot_graph(self):
        current_spot = self.results_dialog.sample_tree.current_spot()
        config = self.results_dialog.configuration_widget.current_config
        if config and current_spot:
            self.plot_cps_graph(current_spot, config)

    def plot_cps_graph(self, spot, config):
        axis = self.axes
        axis.clear()
        if spot is None:
            return
        axis.spines['top'].set_visible(False)
        axis.spines['right'].set_visible(False)
        xs = []
        ys = []
        errors = []
        if DataKey.AGES not in spot.data[config]:
            # TODO plot words on graph
            return
        ages = spot.data[config][DataKey.AGES]
        if len(ages) != 0:
            for i, age in enumerate(ages):
                if isinstance(age, str):
                    continue
                x = i + 1
                y, dy = age
                xs.append(x)
                if y is None:
                    ys.append(0)
                    errors.append(0)
                else:
                    ys.append(y)
                    errors.append(dy)
        else:
            # TODO plot some text
            return
        weighted_age, age_st_dev = spot.data[config][DataKey.WEIGHTED_AGE]
        if isinstance(weighted_age, str):
            string = "No weighted age"
        else:
            string = f"Weighted age: {weighted_age:.0f}, 1σ: {age_st_dev:.0f}"

        axis.errorbar(xs, ys, yerr=errors, linestyle="none", marker='o')
        axis.text(0.5, 1, string, transform=axis.transAxes, horizontalalignment="center")
        axis.set_xlabel("Scan number")
        axis.set_ylabel("Age (ka)")
        self.canvas.draw()

[((5, 0, 5, 24), 'matplotlib.use', 'matplotlib.use', ({(5, 15, 5, 23): '"""QT5Agg"""'}, {}), "('QT5Agg')", False, 'import matplotlib\n'), ((14, 8, 14, 30), 'PyQt5.QtWidgets.QWidget.__init__', 'QWidget.__init__', ({(14, 25, 14, 29): 'self'}, {}), '(self)', False, 'from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel\n'), ((18, 17, 18, 30), 'PyQt5.QtWidgets.QHBoxLayout', 'QHBoxLayout', ({}, {}), '()', False, 'from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel\n'), ((26, 17, 26, 30), 'PyQt5.QtWidgets.QVBoxLayout', 'QVBoxLayout', ({}, {}), '()', False, 'from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel\n'), ((34, 27, 34, 36), 'PyQt5.QtWidgets.QWidget', 'QWidget', ({}, {}), '()', False, 'from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel\n'), ((35, 17, 35, 30), 'PyQt5.QtWidgets.QVBoxLayout', 'QVBoxLayout', ({}, {}), '()', False, 'from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel\n'), ((37, 14, 37, 26), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((38, 20, 38, 30), 'matplotlib.pyplot.axes', 'plt.axes', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((40, 36, 40, 76), 'utils.ui_utils.create_figure_widget', 'ui_utils.create_figure_widget', ({(40, 66, 40, 69): 'fig', (40, 71, 40, 75): 'self'}, {}), '(fig, self)', False, 'from utils import ui_utils\n'), ((27, 25, 27, 55), 'PyQt5.QtWidgets.QLabel', 'QLabel', ({(27, 32, 27, 54): '"""Sample and spot name"""'}, {}), "('Sample and spot name')", False, 'from PyQt5.QtWidgets import QHBoxLayout, QDialog, QPushButton, QWidget, QVBoxLayout, QLabel\n')]