from __future__ import annotations
import os
from typing import Tuple, Union, Optional, List
from dataclasses import dataclass, field

import numpy as np
import pandas as pd
from loguru import logger

from base.base import (
    BaseKukaDataParser, BaseKukaTextParser, 
    BaseTraceStageDetector, BaseTextStageDetector,
    BaseRawTraceProcessor, KukaDataHead, 
    KukaTXT, Settings
)


class KukaDataParser(BaseKukaDataParser):
    """
    Класс для парсинга данных Кука.
    
    Читает заголовочный файл (.dat) и соответствующий файл данных (.r64),
    объединяет их в pandas.DataFrame.
    """
    
    def parse(self, head_filepath: str) -> pd.DataFrame:
        """
        Основной метод парсинга. Читает заголовочный файл и файл данных.
        
        :param head_filepath: Путь к заголовочному файлу.
        :return: DataFrame с объединёнными данными.
        """
        header = self._parse_header_file(head_filepath)
        body_filepath = os.path.join(os.path.dirname(head_filepath), header.filename)
        dataframe = self._parse_body_file(body_filepath, header)
        return dataframe

    def _parse_header_file(self, filepath: str) -> KukaDataHead:
        """
        Парсит заголовочный файл (.dat) и извлекает информацию о каналах.
        
        :param filepath: Путь к заголовочному файлу.
        :return: Объект KukaDataHead с информацией о данных.
        """
        with open(filepath, 'r', encoding='cp1252') as file:
            data_head = KukaDataHead(0, "", {})
            is_inside_channel = False
            self._current_channel = None  # Текущее название канала
            for line in file:
                line = line.strip()
                if line in ('#BEGINCHANNELHEADER', "#BEGINGLOBALHEADER"):
                    is_inside_channel = True
                elif line in ('#ENDCHANNELHEADER', "#ENDGLOBALHEADER"):
                    is_inside_channel = False
                else:
                    if is_inside_channel:
                        self._parse_header_line(line, data_head)
        return data_head

    def _parse_header_line(self, line: str, data_head: KukaDataHead) -> None:
        """
        Обрабатывает отдельную строку заголовочного файла и обновляет объект data_head.
        
        :param line: Строка из файла.
        :param data_head: Объект KukaDataHead для обновления.
        """
        tag, value = line.split(',')
        match tag:
            case '102':
                data_head.rob_ID = value
            case '200':
                self._current_channel = str(value)
                data_head.channels[self._current_channel] = {}
            case '202':
                data_head.channels[self._current_channel]['unit'] = str(value)
            case '211':
                data_head.filename = str(value)
            case '220':
                data_head.channels[self._current_channel]['len'] = int(value)
            case '221':
                data_head.channels[self._current_channel]['num'] = int(value)
            case '241':
                # Переименовано в "multiplier" для повышения читаемости
                data_head.channels[self._current_channel]['multiplier'] = float(value)

    def _parse_body_file(self, filepath: str, data_head: KukaDataHead) -> pd.DataFrame:
        """
        Парсит файл данных (.r64) и объединяет его с временной осью.
        
        :param filepath: Путь к файлу данных.
        :param data_head: Заголовочная информация.
        :return: DataFrame с данными.
        """
        time_axis = self._build_time_axis(data_head)
        multipliers, channel_names = self._extract_channels_data(data_head)
        raw_data = self._read_r64_file(filepath)

        if raw_data.size % len(channel_names) != 0:
            raise ValueError(f"Количество записей в {filepath} не кратно количеству найденных каналов ({len(channel_names)})")
        try:
            data_reshaped = raw_data.reshape(-1, len(channel_names))
        except ValueError as e:
            raise ValueError(f"Ошибка при изменении формы raw_data: {e}")

        # Применяем множители к данным каналов
        df_data = pd.DataFrame(data_reshaped * multipliers, columns=channel_names)
        dataframe = pd.concat([time_axis, df_data], axis=1)
        return dataframe

    @staticmethod
    def _extract_channels_data(data_head: KukaDataHead) -> Tuple[np.ndarray, List[str]]:
        """
        Извлекает данные по каналам, сортируя их по номеру.
        
        :param data_head: Заголовочная информация.
        :return: Кортеж из массива множителей и списка имён каналов.
        """
        sorted_channels = sorted(
            (item for item in data_head.channels.items() if item[0] != "Zeit"),
            key=lambda item: item[1]['num']
        )
        multipliers = np.array([info['multiplier'] for _, info in sorted_channels])
        channel_names = [key for key, _ in sorted_channels]
        return multipliers, channel_names

    @staticmethod
    def _build_time_axis(data_head: KukaDataHead) -> pd.Series:
        """
        Строит временную ось на основе информации из канала 'Zeit'.
        
        :param data_head: Заголовочная информация.
        :return: Серия pandas с временными метками.
        """
        num_timestamps = data_head.channels['Zeit']['len'] - 1
        time_step = data_head.channels['Zeit']['multiplier']
        time_axis = pd.Series(np.arange(0, num_timestamps * time_step, time_step))
        time_axis.name = 'time'
        return time_axis

    @staticmethod
    def _read_r64_file(filepath: str) -> np.ndarray:
        """
        Считывает бинарный файл (.r64) и возвращает массив чисел.
        
        :param filepath: Путь к файлу.
        :return: Массив numpy с типом float.
        """
        with open(filepath, 'rb') as file:
            data = file.read()
            numbers = np.frombuffer(data, dtype='<d')
        return numbers


class KukaTextParser(BaseKukaTextParser):
    """
    Класс для парсинга текстовых данных Кука.
    
    Извлекает сообщения из файла.
    """
    def __init__(self):
        super().__init__()
        self._in_message = False
        self._data_packs: List[KukaTXT] = []

    def parse(self, filepath: str) -> List[KukaTXT]:
        """
        Парсит текстовый файл и возвращает список объектов KukaTXT.
        
        :param filepath: Путь к текстовому файлу.
        :return: Список объектов KukaTXT.
        """
        self._data_packs = []
        with open(filepath, 'r') as file:
            current_pack = KukaTXT()
            for line in file:
                line = line.strip()
                current_pack = self._process_message_flow(line, current_pack)
        return self._data_packs
                
    def _process_message_flow(self, line: str, current_pack: KukaTXT) -> KukaTXT:
        """
        Обрабатывает поток сообщений, определяя начало и конец блока.
        
        :param line: Строка из файла.
        :param current_pack: Текущий объект KukaTXT.
        :return: Обновлённый объект KukaTXT.
        """
        if line == "#BEGINMOTIONINFO":
            self._in_message = True
            current_pack = KukaTXT()
        elif line == "#ENDMOTIONINFO":
            self._in_message = False
            self._data_packs.append(current_pack)
            current_pack = KukaTXT()
        elif self._in_message:
            current_pack = self._process_line(line, current_pack)
        return current_pack

    def _process_line(self, line: str, current_pack: KukaTXT) -> KukaTXT:
        """
        Обрабатывает отдельную строку внутри блока сообщения.
        
        :param line: Строка с данными.
        :param current_pack: Текущий объект KukaTXT.
        :return: Обновлённый объект KukaTXT.
        """
        tag, value = line.split(":")
        match tag:
            case "TIME":
                current_pack.time = float(value)
            case "ENDTIME":
                current_pack.endtime = float(value)
            case "MODULE":
                pass
            case "FUNCTION/PROCEDURE":
                current_pack.func = value
            case "TYPE":
                current_pack.type_ = value
            case "SIGNAL":
                current_pack.signal = value
            case "LINE":
                pass
            case "POINT NAME":
                pass
            case "POINT COORDINATES":
                pass
            case "BLENDING":
                pass
            case "BLENDING PARAMETER":
                pass
            case "VELOCITIES":
                pass
            case "ACCELERATIONS":
                pass
            case "BASE":
                pass
            case "TOOL":
                pass
            case "IPO MODE":
                pass
            case "MOTION MODE":
                pass
            case "LOAD":
                pass
            case "LOAD A3":
                pass
        return current_pack


class TraceStageDetector(BaseTraceStageDetector):
    """
    Класс для детекции этапов (стадий) на основе данных трассировки.
    
    Определяет переходы между различными стадиями процесса по пороговым значениям.
    """
    def __init__(self, parent: TraceProcessor = None):
        super().__init__(parent)

    @staticmethod
    def is_closing(robot_velocity: float, actual_velocity: float, actual_force: float, thresholds: dict) -> bool:
        """
        Определяет, находится ли робот в стадии закрытия.
        
        :param robot_velocity: Скорость робота.
        :param actual_velocity: Фактическая скорость.
        :param actual_force: Фактическое значение силы.
        :param thresholds: Словарь пороговых значений.
        :return: True, если условие стадии закрытия выполнено.
        """
        return actual_velocity > thresholds['act_vel_min'] and abs(robot_velocity) < thresholds['rob_vel_thresh'] and actual_force < thresholds['act_force_close']

    @staticmethod
    def is_squeeze(robot_velocity: float, actual_velocity: float, actual_force: float, force_rate: float, thresholds: dict) -> bool:
        """
        Определяет, находится ли робот в стадии сжатия.
        
        :param robot_velocity: Скорость робота.
        :param actual_velocity: Фактическая скорость.
        :param actual_force: Фактическое значение силы.
        :param force_rate: Темп изменения силы.
        :param thresholds: Словарь пороговых значений.
        :return: True, если условие стадии сжатия выполнено.
        """
        return abs(robot_velocity) < thresholds['rob_vel_thresh'] and actual_velocity < thresholds['act_vel_close'] and force_rate > thresholds['force_increase']

    @staticmethod
    def is_welding(robot_velocity: float, actual_velocity: float, actual_force: float, thresholds: dict) -> bool:
        """
        Определяет, находится ли робот в стадии сварки.
        
        :param robot_velocity: Скорость робота.
        :param actual_velocity: Фактическая скорость.
        :param actual_force: Фактическое значение силы.
        :param thresholds: Словарь пороговых значений.
        :return: True, если условие стадии сварки выполнено.
        """
        return abs(robot_velocity) < thresholds['rob_vel_thresh'] and abs(actual_velocity) < thresholds['act_vel_thresh'] and actual_force > thresholds['act_force_weld']

    @staticmethod
    def is_relief(actual_velocity: float, actual_position_diff: float, force_rate: float, thresholds: dict) -> bool:
        """
        Определяет, находится ли робот в стадии снятия усилия (relief).
        
        :param actual_velocity: Фактическая скорость.
        :param actual_position_diff: Разница в позициях.
        :param force_rate: Темп изменения силы.
        :param thresholds: Словарь пороговых значений.
        :return: True, если условие стадии снятия усилия выполнено.
        """
        return force_rate < -thresholds['force_decrease'] and abs(actual_position_diff) < thresholds['act_pos_decrease'] and actual_velocity < -thresholds['act_vel_negative']

    def detect_stages(self, df: pd.DataFrame) -> List[Tuple[str, float, float]]:
        """
        Детектирует стадии процесса по данным трассировки.
        
        :param df: DataFrame с данными трассировки.
        :return: Список кортежей (название стадии, время начала, время окончания).
        """
        timestamps = df['time'].to_list()
        n = len(df)
        # Извлекаем пороговые значения из настроек
        thresholds = {key: item[0] for key, item in self._parent._settings.filter.items()}
        
        # Вычисляем производную силы и разницу позиций
        actual_force = df["DriveMotorTorq_Act7"].values
        force_diff = np.diff(actual_force, prepend=actual_force[0])
        actual_position = df["DriveMotorPos_Act7"].values
        position_diff = np.diff(actual_position, prepend=actual_position[0])

        stages = []
        current_state = "Oncomming"
        state_start = timestamps[0]
        
        # Проходим по всем записям DataFrame
        for i in range(n):
            robot_velocity = df.loc[i, "CartVel_Act"]
            actual_velocity = df.loc[i, "DriveMotorVel_Act7"]
            force_value = df.loc[i, "DriveMotorTorq_Act7"]
            current_force_rate = force_diff[i]
            current_position_diff = position_diff[i]
            
            if current_state == "Oncomming":
                if self.is_closing(robot_velocity, actual_velocity, force_value, thresholds):
                    state_end = timestamps[i]
                    stages.append(("Oncomming", state_start, state_end))
                    current_state = "Closing"
                    state_start = timestamps[i]
                    
            elif current_state == "Closing":
                if self.is_squeeze(robot_velocity, actual_velocity, force_value, current_force_rate, thresholds):
                    state_end = timestamps[i]
                    stages.append(("Closing", state_start, state_end))
                    current_state = "Squeeze"
                    state_start = timestamps[i]
                    
            elif current_state == "Squeeze":
                if self.is_welding(robot_velocity, actual_velocity, force_value, thresholds):
                    state_end = timestamps[i]
                    stages.append(("Squeeze", state_start, state_end))
                    current_state = "Welding"
                    state_start = timestamps[i]
                    
            elif current_state == "Welding":
                if self.is_relief(actual_velocity, current_position_diff, current_force_rate, thresholds):
                    state_end = timestamps[i]
                    stages.append(("Welding", state_start, state_end))
                    current_state = "Relief"
                    state_start = timestamps[i]
                    
            elif current_state == "Relief":
                # Если признаки снятия усилия не соблюдаются, завершаем цикл и переходим в Oncomming
                if not self.is_relief(actual_velocity, current_position_diff, current_force_rate, thresholds):
                    state_end = timestamps[i]
                    stages.append(("Relief", state_start, state_end))
                    current_state = "Oncomming"
                    state_start = timestamps[i]
        
        # Фиксируем последний сегмент
        stages.append((current_state, state_start, timestamps[-1]))
        return stages


class TextStageDetector(BaseTextStageDetector):
    """
    Класс для детекции сварочных стадий на основе текстовых данных.
    """
    def __init__(self, parent: TraceProcessor = None):
        super().__init__(parent)
    
    @staticmethod
    def detect_welding(data: List[KukaTXT]) -> List[dict]:
        """
        Детектирует этап сварки по текстовым сообщениям.
        
        :param data: Список объектов KukaTXT.
        :return: Список словарей с информацией о сварке (стадия, время начала и окончания).
        """
        stages = []
        for i in range(len(data) - 1):  # Предотвращаем выход за пределы списка
            if data[i].func == " SPOT" and data[i].signal == " END":
                stages.append({
                    "stage": "welding",
                    "start_time": data[i].time,
                    "end_time": data[i + 1].time
                })
        return stages


class TraceProcessor(BaseRawTraceProcessor):
    """
    Основной класс для обработки трассировок.
    
    Объединяет данные, детектирует стадии и инициирует рендеринг через медиатор.
    """
    def __init__(self):
        self._settings = Settings()
        data_parser = KukaDataParser()
        text_parser = KukaTextParser()
        trace_detector = TraceStageDetector(self)
        text_detector = TextStageDetector(self)

        super().__init__(data_parser, text_parser, trace_detector, text_detector)

    def prerender(self, data: List[str]) -> None:
        """
        Препроцессинг данных для отрисовки трейсов и обнаруженных событий.
        
        :param data: Список путей к файлам данных.
        """
        rendered_data = self._render_data(data)
        self._mediator.prerender_TCW(self, rendered_data)

    def final_render(self) -> None:
        """
        Финальный рендеринг данных для отрисовки 
        трейсов клиента и расчета трейсов ТСК.
        """
        events = self._detect_stages(self._trace_df, self._text_data)
        renamed_df = self._rename_df_columns(self._trace_df)
        part_position, max_open = self._detect_coords(renamed_df, events)
        self._mediator.render_TCW(self, [renamed_df, events, [part_position, max_open]])

    def _detect_coords(self, dataframe: pd.DataFrame, events: dict) -> Tuple[List[float], List[float]]:
        """
        Детектирует координаты на основе событий и данных.
        
        :param dataframe: DataFrame с данными.
        :param events: Словарь с событиями.
        :return: Кортеж из списков: координаты детали и максимальное открытие электрода.
        """
        weld_timings = events["Welding"]
        oncommings = events["Oncomming"]
        open_positions = []
        part_positions = []
        for i in range(len(weld_timings[0])):
            weld_start = weld_timings[0][i]
            weld_end = weld_timings[1][i]
            onc_start = oncommings[0][i]
            onc_end = oncommings[1][i]
            pos_part = dataframe[(dataframe["time"] > weld_start) & (dataframe["time"] < weld_end)]["Electrode Position, mm ME"].mean()
            part_positions.append(float(pos_part) / 1000)
            pos_open = dataframe[(dataframe["time"] > onc_start) & (dataframe["time"] < onc_end)]["Electrode Position, mm ME"].abs().max()
            open_positions.append(float(pos_open) / 1000)
        return (part_positions, open_positions)

    def _render_data(self, data: List[str]) -> List:
        """
        Обрабатывает входные файлы данных и возвращает предварительно обработанные данные.
        
        :param data: Список путей к файлам (первый для трассировки, второй для текстовых данных).
        :return: Список, содержащий DataFrame и события.
        """
        if data and len(data) == 2:
            trace_filepath = data[0]
            text_filepath = data[1]
        elif data:
            trace_filepath = data[0]
            text_filepath = None
        else:
            trace_filepath = None
            text_filepath = None

        self._trace_df = self._unpack_trace(trace_filepath)
        self._text_data = self._unpack_text(text_filepath)
        events = self._detect_stages(self._trace_df, self._text_data)
        renamed_df = self._rename_df_columns(self._trace_df)
        return [renamed_df, events]

    def update_settings(self, settings: Settings) -> None:
        """
        Обновляет настройки обработки.
        
        :param settings: Объект настроек.
        """
        self._settings = settings

    def _unpack_trace(self, trace_filepath: str = None) -> Optional[pd.DataFrame]:
        """
        Распаковывает трассировочные данные из файла.
        
        :param trace_filepath: Путь к файлу трассировки.
        :return: DataFrame с данными или None.
        """
        if trace_filepath:
            return self._dataparser.parse(trace_filepath)
        return None

    def _unpack_text(self, text_filepath: str = None) -> Optional[List[KukaTXT]]:
        """
        Распаковывает текстовые данные из файла.
        
        :param text_filepath: Путь к текстовому файлу.
        :return: Список объектов KukaTXT или None.
        """
        if text_filepath:
            return self._textparser.parse(text_filepath)
        return None
    
    def _detect_stages(self, trace_df: pd.DataFrame, text_data: List) -> Optional[dict]:
        """
        Детектирует события на основе трассировочных и текстовых данных.
        
        :param trace_df: DataFrame с трассировочными данными.
        :param text_data: Список текстовых данных.
        :return: Словарь с событиями или None.
        """
        if trace_df is not None and text_data is not None:
            trace_stages = self._data_detector.detect_stages(trace_df)
            events = self._form_events(trace_stages)
            return events
        return None
    
    @staticmethod
    def _form_events(trace_stages: list) -> dict:
        """
        Формирует словарь событий на основе списка этапов трассировки.
        
        :param trace_stages: Список этапов (кортежи: название, время начала, время окончания).
        :return: Словарь с событиями.
        """
        events = {
            "Closing": [[], []],
            "Squeeze": [[], []],
            "Welding": [[], []],
            "Relief": [[], []],
            "Oncomming": [[], []]
        }
        for stage in trace_stages:
            name, start_time, end_time = stage
            events[name][0].append(start_time)
            events[name][1].append(end_time)
        return events

    @staticmethod
    def _rename_df_columns(dataframe: pd.DataFrame) -> Optional[pd.DataFrame]:
        """
        Переименовывает столбцы DataFrame на основе корректной карты соответствия.
        
        :param dataframe: Исходный DataFrame.
        :return: DataFrame с переименованными столбцами или None в случае ошибки.
        """
        correct_mapping = {
            "time": ["Time", "Timestamp"],
            "Tool Coordinate, mm X": ["X_Act"],
            "Tool Coordinate, mm Y": ["Y_Act"],
            "Tool Coordinate, mm Z": ["Z_Act"], 
            "Electrode Force, N ME": ["DriveMotorTorq_Act7"],
            "Electrode Position, mm ME": ["AxisPos_Act7"],
            "Electrode Speed, mm ME": ["AxisVel_Act7"],
            "Rotor Position, deg FE": ["DriveMotorPos_Act7"], 
            "Rotor Speed, deg/s ME": ["DriveMotorVel_Act7"], 
            "Rotor Current, A ME": ["DriveMotorCurr_Act7"],
            "Cartesian Tool Speed, mm/s ME": ["CartVel_Act"], 
        }
        try:
            df_copy = dataframe.copy(deep=True)
            working_mapping = {key: [item.lower() for item in items] for key, items in correct_mapping.items()}
            
            new_columns = {}
            for col in df_copy.columns:
                col_lower = col.lower()
                for key, values in working_mapping.items():
                    if col_lower in values:
                        new_columns[col] = key
                        break
                else:
                    new_columns[col] = col
            df_copy.rename(columns=new_columns, inplace=True)
            df_copy = df_copy.loc[:, ~df_copy.columns.duplicated()]
            return df_copy
        except AttributeError as e:
            logger.error(f"_rename_df_columns - AttributeError: Проверьте, что переданный объект является DataFrame. {e}")
            return None
        except Exception as e:
            logger.error(f"_rename_df_columns - Непредвиденная ошибка: {e}")
            return None


"""
Примеры комментариев для этапов процесса:

Перемещение:
# FUNCTION/PROCEDURE: SW_RSP030TL01_SN - какое-то перемещение
#
# ИЛИ
#
# FUNCTION/PROCEDURE: SGL_MoveToPos - перемещение между точками
# SIGNAL: BLENDING 

Смыкание и набор усилия:
# FUNCTION/PROCEDURE: SGM_MOVE_TO_FORCE - перемещение электрода с движением робота
# ...
# FUNCTION/PROCEDURE: SGM_MOVE_TO_FORCE - перемещение 0.5 мм роботом с движением электрода
# ...
# FUNCTION/PROCEDURE: SGM_MOVE_TO_FORCE - остановка в позиции (может быть набор усилия?)

Сварка:
# FUNCTION/PROCEDURE: SPOT - Начало сварочного процесса
# SIGNAL: START
#
# FUNCTION/PROCEDURE: SPOT - Конец сварочного процесса
# SIGNAL: END

Снятие усилия и разъезд:
# FUNCTION/PROCEDURE: SGL_MoveToPos - Выход из контакта с точкой движением робота и электрода
# SIGNAL: START
"""
