import copy
import traceback
import sys
from typing import Optional, Tuple, Callable, List,  Any
from dataclasses import dataclass, field

from PyQt5.QtWidgets import (
    QWidget, 
    QVBoxLayout,
    QHBoxLayout, 
    QLabel,
    QGraphicsRectItem, 
    QSpacerItem,
    QSizePolicy
    )
from PyQt5.QtCore import Qt
from loguru import logger
import pyqtgraph as pg
import pandas as pd
import numpy as np

from base.base import BasePlotWidget, GraphicPassport, PlotItems, PointPassport, UsefulGraphData, BaseController
from utils.json_tools import read_json
from utils import qt_settings as qts


@dataclass
class ChannelTimings():
    shift: float = 0
    TWC_time: float = 0.0
    ideal_time: float = 0.0
    client_time: float = 0.0
    TWC_start: float = 0.0
    TWC_end: float = 0.0
    worst_performance: float = 2
    worst_timeframe: list = field(default_factory=lambda: [0, 0])


class PlotWidget(BasePlotWidget):
    
    def __init__(self, controller:BaseController):
        super().__init__(controller=controller)
        self._plt_structures = read_json("params/plot_structure_params.json")
        self._plt_channels = None

    def build(self, data: list[GraphicPassport]) -> None:
        """
        Создает набор виджетов по предоставленному списку данных.
        """
        try:
            self._datalen = len(data)
            widgets_datapack = [self._build_widget(data_sample) for self._datastep, data_sample in enumerate(data)]

        except (KeyError, ValueError, TypeError) as e:
            logger.error(f"Возникла проблема при сборке данных графика: {e}")
            logger.debug(traceback.format_exc())
            error_label = QLabel("Произошла ошибка при формировании графиков. Пожалуйста, проверьте корректность данных.")
            widgets_datapack = [error_label]
        except Exception as e:
            logger.error(f"Непредвиденная ошибка при формировании графиков: {e}")
            logger.debug(traceback.format_exc())
            error_label = QLabel("Непредвиденная ошибка при формировании графиков.")
            widgets_datapack = [error_label]
        finally:
            self._mediator.notify(self, widgets_datapack)

    def build_raw_trace(self, data:list[pd.DataFrame, dict]) -> None:
        """
        Создаёт один виджет с одним графиком, где представлены все данные
        """
        container_widget, container_layout, pyqt_container = self._generate_widget_container()
        plot = self._build_raw_plotitem(data, pyqt_container)     
        container_layout.addWidget(plot)
        container_widget.setProperty("pyqt_container", pyqt_container)
        
        self._mediator.notify(self, [container_widget])
    
    def set_mode(self, mode:int) -> None:
        match mode:
            case 1:
                self._plt_channels = self._plt_structures["Online Path Scanner"]
            case 2: 
                self._plt_channels = self._plt_structures["Report Editor"]
            case 3: 
                self._plt_channels = self._plt_structures["Client Trace Watcher"]      

    def _build_raw_plotitem(self, 
                            data:list[pd.DataFrame, dict], 
                            pyqt_container:PlotItems) -> pg.GraphicsLayoutWidget:
        
        plot_item, legend = PlotItemGenerator._init_plot_item("Customer data")
        dataframe, events = data
        channels = dataframe.columns.to_list()
        for i, channel in enumerate(channels):
            plot = plot_item.plot(dataframe["time"], dataframe[channel], pen = qts.colors[i], fast = True)
            legend.addItem(plot, channel)
            pyqt_container.curves["real"].setdefault(channel, {})
            pyqt_container.curves["real"][channel] = plot

        for i in range(len(events["Squeeze"][0])):
            point_events = {}
            for key, item in events.items():
                point_events[key] = [item[0][i], item[1][i]]
                PlotItemGenerator._add_stage_regions(self._stage_colors, plot_item, point_events, pyqt_container.regions,20)

        plot_layout = pg.GraphicsLayoutWidget()
        plot_layout.addItem(plot_item)

        return plot_layout

    def _build_performance_label(self,
                               timings: ChannelTimings, 
                               qt_items: dict) -> QWidget:
        """
        Добавляет QLabel с информацией о производительности.
        """
        tesla_TWC = round((1 - timings.TWC_time/timings.client_time)*100, 2) if timings.client_time else 0.0
        tesla_ideal = round((1 - timings.ideal_time/timings.client_time)*100, 2) if timings.client_time else 0.0
        TWC_ideal = round((timings.ideal_time/timings.TWC_time)*100, 2) if timings.TWC_time else 0.0
        label_widget = QWidget()
        label_layout = QHBoxLayout(label_widget)
        start_label = QLabel("Сокращение длительности: ")
        real_label = QLabel(f"фактическое = {tesla_TWC} % ")
        if not tesla_TWC or not timings.TWC_time: real_label.setVisible(False)
        ideal_label = QLabel(f"идеальное = {tesla_ideal} % ")
        if not tesla_ideal: ideal_label.setVisible(False)
        kdip_label = QLabel(f"КДИП = {TWC_ideal}% ")
        if not TWC_ideal: kdip_label.setVisible(False)
        label_layout.addWidget(start_label, alignment=Qt.AlignLeft)
        label_layout.addWidget(real_label, alignment=Qt.AlignLeft)
        label_layout.addWidget(ideal_label, alignment=Qt.AlignLeft)
        label_layout.addWidget(kdip_label, alignment=Qt.AlignLeft)
        spacer = QSpacerItem(1, 1, QSizePolicy.Expanding, QSizePolicy.Minimum)
        label_layout.addSpacerItem(spacer)
        self.set_style(label_widget)

        qt_items["performance label"] = label_widget
        qt_items["real performance"] = real_label
        qt_items["ideal performance"] = ideal_label
        qt_items["real to ideal performance"] = kdip_label
        return label_widget

    def _build_widget(self, graphic_passport: GraphicPassport) -> QWidget:
        """
        Собирает графический виджет для одного набора данных.
        """
        container_widget, container_layout, pyqt_container = self._generate_widget_container()
        plot_layout = CustomPlotLayout(graphic_passport, len(self._plt_channels), self._stage_colors, self)
        plot_layout.build(pyqt_container, self._plt_channels)

        if plot_layout.property("performance"):
            perf_widget = self._build_performance_label(
                plot_layout.property("performance"), 
                pyqt_container.qt_items
                )
            container_layout.addWidget(perf_widget)
            
        container_layout.addWidget(plot_layout)
        container_widget.setProperty("pyqt_container", pyqt_container)
        return container_widget
    
    @staticmethod
    def _generate_widget_container() -> Tuple[QWidget, QVBoxLayout, PlotItems]:
        container_widget = QWidget()
        container_layout = QVBoxLayout(container_widget)
        pyqt_container = PlotItems({"real":{}, "ideal":{}}, {"real":{}, "ideal":{}}, {})
        return (container_widget, container_layout, pyqt_container)

    def _update_status(self, widgsteps:int, pointsteps:int, cur_widg:int, cur_point:int):
        if self._datalen:
            sycle_start = self._datastep/self._datalen*100 + 1
            period1 = 99/self._datalen
        else:
            sycle_start = 1
            period1 = 100

        period2 = period1/widgsteps if widgsteps != 0 else period1
        period3 = period2/pointsteps if pointsteps != 0 else period2

        progress = int(sycle_start + period2*cur_widg + period3*cur_point)
        self._controller.update_progress(progress)


class CustomPlotLayout(pg.GraphicsLayoutWidget):
    def __init__(self, 
                 graphic_passport: GraphicPassport, 
                 widget_steps: int, 
                 colors: dict, 
                 parent: PlotWidget = None) -> None:
        super().__init__()
        self._plotter = PlotItemGenerator(graphic_passport, widget_steps, colors, parent)
        self.setProperty("performance", None)

    def build(self, pyqt_container:PlotItems, plt_channels:dict) -> None:

        for widget_num, (channel, description) in enumerate(plt_channels.items()):
            plot_item, plot_timings = self._plotter.generate_plot_item(widget_num, channel, description, pyqt_container)
            if widget_num == 0:
                main_plot = plot_item
            else:
                plot_item.setXLink(main_plot)

            if description["Settings"].get("performance", False):
                self.setProperty("performance", plot_timings)
            
            self.addItem(plot_item, widget_num, 0)

        timings = ChannelTimings()
        if self.property('performance'):timings:ChannelTimings = self.property('performance')
        navigator = NavigatorPlot(timings.worst_timeframe, main_plot)
        if navigator is not None:
            self.addItem(navigator, widget_num+1, 0)

class PlotItemGenerator:

    def __init__(self, 
                 graphic_passport: GraphicPassport, 
                 widget_steps: int, 
                 colors: dict, parent:PlotWidget = None) -> None:
        self._stage_colors = colors
        self._ideal_mode = graphic_passport.dataframe is None
        self._parent = parent
        
        if not self._ideal_mode:
            dataframe_headers = graphic_passport.dataframe.columns.tolist()
            point_steps = len(graphic_passport.points_pocket)
        else: 
            dataframe_headers = []
            point_steps = 1
            
        self._datapack = [
            graphic_passport.dataframe, 
            dataframe_headers, 
            graphic_passport.useful_data,
            graphic_passport.points_pocket, 
            widget_steps, 
            point_steps
        ]

    def generate_plot_item(self, 
                    widget_num: int,
                    channel: str, 
                    description: dict[str, Any], 
                    pyqt_container: PlotItems) -> Tuple[pg.PlotItem, ChannelTimings]:
    
        dataframe, dataframe_headers, useful_data, points_pocket, widget_steps, point_steps = self._datapack

        plot_item, legend = self._init_plot_item(title=channel)
        settings:dict = description["Settings"]
        timings = ChannelTimings()
        timings.client_time = useful_data.client_time

        ideal_df = pd.DataFrame({})

        # TODO: рассчитать корректный параметр range 
        if settings.get("mirror ME", False) and not self._ideal_mode:
            dataframe = self._shift_data(
                "ME", 
                description["Real_signals"], 
                dataframe, 
                lambda x: useful_data.range_ME-x
            )

        # Итерация по точкам
        for cur_point, data in enumerate(points_pocket):
            point_data: PointPassport = data
            ideal_data = copy.deepcopy(point_data.ideal_data)
            is_last = (cur_point == len(points_pocket) - 1)
            is_first = (cur_point == 0)
            
            if self._ideal_mode:
                timings, point_data.events, point_data.timeframe = self._generate_synthetic_events(timings, ideal_data)
            else:
                if settings.get("force compensation FE", False):
                    force = point_data.useful_data["force"]
                    k_hardness = useful_data.k_hardness
                    signals = description["Real_signals"]
                    dataframe = self._apply_force_compensation(force, k_hardness, dataframe, point_data.timeframe, signals)

                if settings.get("stages", False):
                    self._add_stage_regions(self._stage_colors, plot_item, point_data.events, pyqt_container.regions, 75)

                if settings.get("force accuracy", False):
                    force = point_data.useful_data["force"]
                    self._add_force_accuracy_region(point_data.events["Welding"], force, plot_item)

            if settings.get("ideals", False) and settings.get("mirror ME", False):
                for stage in point_data.events.keys():
                    ideal_data[stage] = self._shift_data("ME", description["Ideal_signals"], ideal_data[stage], lambda x: useful_data.range_ME-x)

            if settings.get("workpiece", False):
                self._add_workpiece(point_data, plot_item)

            if settings.get("ideals", False):
                self._add_ideal_stage_regions(self._stage_colors, plot_item, ideal_data, point_data.events, pyqt_container.regions, 100)
                ideal_df = self._modify_ideal_df(ideal_df, ideal_data, point_data.events)

            if settings.get("performance", False):
                timings = self._calc_performance(timings, point_data, ideal_data, is_first, is_last)
                            
            self._parent._update_status(widget_steps, point_steps, widget_num, cur_point)

        # Добавляем реальные сигналы
        if settings.get("ideals", False):
            self._add_signals(plot_item, ideal_df, description["Ideal_signals"], legend, pyqt_container.curves["ideal"])

        if not self._ideal_mode:
            self._add_signals(plot_item, dataframe, description["Real_signals"], legend, pyqt_container.curves["real"])
        return (plot_item, timings)
    
    @staticmethod
    def _shift_data(valid_str: str,
                    signals: list[dict],
                    dataframe: pd.DataFrame, 
                    func: Callable) -> pd.DataFrame:
        keys = dataframe.keys()
        for signal in signals:
            if valid_str in signal["name"] and signal["name"] in keys:
                dataframe[signal["name"]] = dataframe[signal["name"]].apply(func)
        return dataframe

    @staticmethod
    def _init_plot_item(title: str) -> tuple[pg.PlotItem, pg.LegendItem]:
        plot_item = pg.PlotItem(title=title)
        plot_item.sigXRangeChanged.connect(lambda: PlotItemGenerator._update_plots_downsample(plot_item))
        # Оптимизация отображения графиков
        plot_item.showGrid(x=True, y=True)
        plot_item.setClipToView(True)
        legend = plot_item.addLegend(offset=(70, 20))
        return plot_item, legend

    @staticmethod
    def _create_stage_region(colors:dict,
                             stage: str,
                             start_timestamp: float,
                             finish_timestamp: float,
                             transparency: int) -> Optional[pg.LinearRegionItem]:
        """
        Создает регион для определённого этапа, если заданы временные рамки.
        """
        if start_timestamp is not None and finish_timestamp is not None:
            region = pg.LinearRegionItem([start_timestamp, finish_timestamp], movable=False)
            color = colors.get(stage, [100, 100, 100, 100])
            region.setBrush(pg.mkBrush(color[:3] + [transparency]))
            return region
        return None

    @staticmethod
    def _add_stage_regions(colors:dict,
                           plot_item: pg.PlotItem,
                           point_events: dict,
                           reg_items: dict,
                           transparency: int = 75) -> None:
        """
        Добавляет регионы для реальных этапов
        """
        stages = point_events.keys()
        #if all(stage in dataframe_headers for stage in stages):
        for stage in stages:
            start_t, end_t = point_events[stage]
            region = PlotItemGenerator._create_stage_region(colors, stage, start_t, end_t, transparency)
            if region is not None:
                region.setZValue(-20)
                plot_item.addItem(region)
                reg_items["real"].setdefault(stage, [])
                reg_items["real"][stage].append(region)

    @staticmethod
    def _add_ideal_stage_regions(colors: dict, 
                                 plot_item: pg.PlotItem,
                                 ideal_data: dict[str, Any],
                                 point_events: dict[str, list[float]],
                                 reg_items: dict,
                                 transparency: int = 125) -> None:
        """
        Добавляет регионы для идеальных этапов.
        """
        ideal_timings = ideal_data["Ideal timings"]
        stages = list(point_events.keys())
        for i, stage in enumerate(stages):
            start_t = point_events[stage][0]
            end_t = start_t + ideal_timings[i]
            region = PlotItemGenerator._create_stage_region(colors, stage, start_t, end_t, transparency)
            if region:
                region.setZValue(-10)
                plot_item.addItem(region)
                reg_items["ideal"].setdefault(stage, [])
                reg_items["ideal"][stage].append(region) 

    @staticmethod
    def _modify_ideal_df(ideal_df: pd.DataFrame,
                        ideal_data: dict[str, Any],
                        point_events: dict[str, list[float]]) -> None:
        """
        Добавляет идеальные сигналы для каждого этапа.
        """
        for stage in point_events.keys():
            if not ideal_df.empty:
                last_time = ideal_df['time'].iloc[-1]
                separator_row = {col: np.nan if col != "time" else last_time + 0.01 for col in ideal_df.columns}
                separator_df = pd.DataFrame([separator_row])
                worker_df = ideal_data[stage].copy(deep=True)
                worker_df["time"] = worker_df["time"]+point_events[stage][0]
                ideal_df = pd.concat([ideal_df, separator_df, worker_df], ignore_index=True)
            else: 
                ideal_df = ideal_data[stage].copy()
        return ideal_df
    
    @staticmethod
    def _add_signals(plot_item: pg.PlotItem,
                    dataframe: pd.DataFrame,
                    real_signals: list[dict[str, Any]],
                    legend: pg.LegendItem, 
                    curve_items: dict) -> None:
        """
        Добавляет реальные сигналы из dataframe на виджет.
        """
        dataframe_headers = dataframe.columns.tolist()
        for signal in real_signals:
            if signal["name"] in dataframe_headers:
                plot = plot_item.plot(dataframe["time"], dataframe[signal["name"]], pen=signal["pen"], fast=True)
                plot.setZValue(0)
                legend.addItem(plot, signal["name"])
                curve_items.setdefault(signal["name"], {})
                curve_items[signal["name"]] = plot

    @staticmethod
    def _update_plots_downsample(plot_item:pg.PlotItem):
        visible_range = plot_item.getViewBox().viewRange()[0]
        diapason = visible_range[1] - visible_range[0]
        if diapason >= 30: plot_item.setDownsampling(ds=50, auto=True, mode='peak')
        elif diapason >= 10: plot_item.setDownsampling(ds=20, auto=True, mode='peak')
        elif diapason >=4: plot_item.setDownsampling(ds=10, auto=True, mode='peak')
        else: plot_item.setDownsampling(ds=1, auto=True, mode='peak')

    @staticmethod
    def _add_force_accuracy_region(event:list, 
                                   force: float, 
                                   plot_item:pg.PlotItem) -> None:
        modifier = 0.05
        x1 = event[0]
        dx = event[1] - x1
        y1 = force*(1-modifier)
        dy = force*(2*modifier)

        rect_item = QGraphicsRectItem(x1, y1, dx, dy)
        rect_item.setZValue(-5)
        rect_item.setBrush(pg.mkBrush((0,255,0, 50)))
        rect_item.setPen(pg.mkPen('black', width=0))
        plot_item.addItem(rect_item)

    @staticmethod
    def _add_workpiece(point_data:PointPassport, 
                       plot_item: pg.PlotItem) -> None:
        x1 = point_data.events["Closing"][0]
        dx = point_data.events["Relief"][1] - x1
        y1 = point_data.useful_data["part_pos"]*1000
        dy = point_data.useful_data["thickness"]*1000

        rect_item = QGraphicsRectItem(x1, y1, dx, dy)
        rect_item.setZValue(-5)
        rect_item.setBrush(pg.mkBrush('grey'))
        rect_item.setPen(pg.mkPen('black', width=3))
        plot_item.addItem(rect_item)
    
    def _calc_performance(self, 
                          timings:ChannelTimings, 
                          point_data: PointPassport, 
                          ideal_data:dict,
                          is_first:bool, 
                          is_last:bool) -> ChannelTimings:
        if is_first:
            if not self._ideal_mode: 
                timings.TWC_start = point_data.events["Closing"][0]
            ideal_delta = ideal_data["Ideal cycle"]
        elif is_last:
            if not self._ideal_mode: 
                timings.TWC_end = point_data.events["Relief"][1]
                timings.TWC_time = timings.TWC_end - timings.TWC_start
                timings.worst_timeframe = [timings.TWC_start, timings.TWC_end]
            ideal_delta = sum(ideal_data["Ideal timings"][0:3])
        else:
            ideal_delta = ideal_data["Ideal cycle"]

        timings.ideal_time += ideal_delta
        #curr_perf = ideal_delta/TWC_delta if TWC_delta != 0 else 1

        if False: #curr_perf < timings.worst_performance:
            timings.worst_performance = curr_perf
            timings.worst_timeframe = point_data.timeframe
        return timings
    
    @staticmethod
    def _generate_synthetic_events(timings:ChannelTimings, 
                                   ideal_data:dict) -> Tuple[ChannelTimings, dict, list[float]]:
        timings.worst_timeframe = point_timeframe = [timings.shift, timings.shift+ ideal_data["Ideal cycle"]]
        point_events = {}
        keys = list(ideal_data.keys())
        shift = 0
        for i, time in enumerate(ideal_data["Ideal timings"]):
            point_events[keys[i]] = [timings.shift+shift, timings.shift+time+shift]
            shift += time
        timings.shift +=ideal_data["Ideal cycle"]
        return (timings, point_events, point_timeframe)

    @staticmethod
    def _apply_force_compensation(force: float, 
                                  k_hardness: float,
                                  dataframe:pd.DataFrame, 
                                  point_timeframe:list, 
                                  real_signals:list[dict]) -> pd.DataFrame:
        F_comp = - force/k_hardness
        point_idxs = dataframe[(dataframe["time"] >= point_timeframe[0]) & (dataframe["time"] <= point_timeframe[1])].index
        dataframe.loc[point_idxs] = PlotItemGenerator._shift_data("FE", real_signals, dataframe.loc[point_idxs], lambda x: x + F_comp)
        return dataframe


class NavigatorPlot(pg.PlotItem):
    def __init__(self, 
                 time_region:tuple[float, float],
                 main_plot: pg.PlotItem):
        super().__init__()
        self._init_navigator(time_region, main_plot)
        self._init_syncranisation(main_plot)

    @staticmethod
    def _downsample_data(x:list, y:list, max_points=5000):
        """
        Понижает разрешение данных до заданного количества точек для улучшения производительности навигатора.
        """
        if len(x) > max_points:
            factor = len(x) // max_points
            x_downsampled = x[::factor]
            y_downsampled = y[::factor]
            return x_downsampled, y_downsampled
        return x, y
    
    @staticmethod
    def _sync_main_plot_with_navigator(main_plot: pg.PlotItem,
                                    region: pg.LinearRegionItem) -> None:
        """
        Синхронизирует область просмотра основного графика с регионом навигатора.
        """
        x_min, x_max = region.getRegion()
        if main_plot:
            main_plot.blockSignals(True)
            main_plot.setXRange(x_min, x_max, padding=0)
            main_plot.blockSignals(False)

    @staticmethod
    def _sync_navigator_with_main(main_plot: pg.PlotItem, region:pg.LinearRegionItem):
        """
        Синхронизирует регион навигатора с областью просмотра основного графика.
        """
        if region:
            x_min, x_max = main_plot
            region.blockSignals(True)  # Предотвращаем рекурсию
            region.setRegion([x_min, x_max])
            region.blockSignals(False)

    def _init_syncranisation(self, main_plot: pg.PlotItem) -> None:
        """
        Связывает изменения навигатора и других графиков друг с другом
        """
        self._sync_main_plot_with_navigator(main_plot, self.ROI_region)
        self.ROI_region.sigRegionChanged.connect(
            lambda: self._sync_main_plot_with_navigator(main_plot, self.ROI_region)
            )
        main_plot.sigXRangeChanged.connect(
            lambda _, plot=main_plot, region=self.ROI_region: self._sync_navigator_with_main(main_plot=plot, region=region)
            )

    def _init_navigator(self, 
                        time_region:tuple[float, float],
                        main_plot: pg.PlotItem) -> None:
        """
        Создаёт график-навигатор, отображающий все данные в уменьшенном масштабе.
        """
        self.setTitle("Navigator")
        self.setFixedHeight(100) 
        # Получение кривых из main_plot
        for curve in main_plot.listDataItems():
            # Извлекаем данные из кривой
            x, y = curve.getData()
            curve_name = curve.opts.get("name", None)
            signal_pen = curve.opts.get("pen", None)
            x_downsampled, y_downsampled = self._downsample_data(x, y, max_points=1000)
            self.plot(x_downsampled, y_downsampled, pen=signal_pen, name=curve_name)

        self.ROI_region = pg.LinearRegionItem(movable=True, brush=pg.mkBrush(0, 0, 255, 100), pen=pg.mkPen(width=4))
        maxBound = 50
        if len(x) !=0: maxBound = x[-1]
        self.ROI_region.setBounds([0, maxBound])
        self.ROI_region.setRegion(time_region)
        self.addItem(self.ROI_region)
        self.getViewBox().setLimits(xMin=0, xMax=maxBound)


