import os
import pandas as pd
import numpy as np
from scipy.signal import butter, lfilter
from scipy.fftpack import fft
from scipy.signal import welch
from tensorflow import keras
import time
import serial
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QFrame, QPushButton
from PyQt5.QtCore import QTimer
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
from PyQt5.QtGui import QFont, QColor, QPalette, QPixmap
from PyQt5.QtCore import Qt
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['Arial']
matplotlib.rcParams['axes.unicode_minus'] = False

# Ignore TensorFlow warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# Pandas display options
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)


def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a


def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y


def create_sequences(data, time_steps=5):
    X = []
    for i in range(len(data) - time_steps + 1):
        X.append(data[i:i + time_steps])
    return np.array(X)


# Reference columns for model prediction
reference_columns = [
    'RawData', 'Attention', 'Delta', 'HighAlpha', 'HighBeta',
    'new_column_25', 'new_column_26', 'new_column_27', 'new_column_28',
    'new_column_29', 'new_column_30', 'new_column_31', 'new_column_32',
    'mean', 'std', 'max', 'min', 'rms', 'dominant_freq', 'psd_mean', 'psd_max'
]

# Read 256 rows each time
num_rows_to_read = 256
num_last_rows = 52

# Label names mapping
label_names = {0: 'Decelerate', 1: 'LaneChangeLeft', 2: 'LaneChangeRight', 3: 'InvalidCommand'}
# Command mapping
command_mapping = {
    'Decelerate': 'slow',
    'LaneChangeLeft': 'left',
    'LaneChangeRight': 'righ',
    'Advance': 'adva',
    'Stop': 'stop'
}

# Simulated training distribution
training_distribution = np.array([0.2, 0.2, 0.3, 0.3])

# Thresholds
threshold = 0.9
right_lane_change_threshold = threshold * 0.000001
slow_down_threshold = threshold * 1000000
left_lane_change_threshold = threshold * 0.5

confidence_threshold = 0.5
port = 'COM3'

last_slow_down_time = 0
last_lane_change_time = 0
recent_command_5s = None
current_lane = 3
last_sent_command_name = None
last_sent_command_time = 0
send_commands = False


def send_data_to_bluetooth(ser, command):
    try:
        print(f"Trying to send data to Bluetooth device: {command}")
        if not ser.is_open:
            print("Serial port not open, cannot send data.")
            return
        ser.write(command.encode())
        print(f"Data sent to Bluetooth device: {command}")
        time.sleep(1)

        print("Trying to receive data from Bluetooth device...")
        received_bytes = ser.readline()
        if received_bytes:
            received_data = received_bytes.decode().strip()
            print(f"Received data from Bluetooth: {received_data}")
        else:
            print("No data received from Bluetooth device")
        return True
    except serial.SerialException as e:
        print(f"Serial communication error: {e}")
    except Exception as e:
        print(f"Unknown error: {e}")
    return False


def extract_time_features(signal):
    if len(signal) == 0:
        return {'mean': 0, 'std': 0, 'max': 0, 'min': 0, 'rms': 0}
    return {
        'mean': np.mean(signal),
        'std': np.std(signal),
        'max': np.max(signal),
        'min': np.min(signal),
        'rms': np.sqrt(np.mean(signal ** 2))
    }


def extract_frequency_features(signal, sampling_rate=128):
    if len(signal) < 2:
        return {'dominant_freq': 0, 'psd_mean': 0, 'psd_max': 0}

    fft_values = fft(signal)
    fft_magnitude = np.abs(fft_values[:len(fft_values) // 2])
    freqs = np.fft.fftfreq(len(signal), d=1 / sampling_rate)[:len(fft_values) // 2]

    dominant_freq = freqs[np.argmax(fft_magnitude)] if len(fft_magnitude) > 0 and not np.all(fft_magnitude == 0) else 0

    nperseg = len(signal) if len(signal) < 256 else 256
    freqs_psd, psd = welch(signal, fs=sampling_rate, nperseg=nperseg)
    psd = np.abs(psd)

    return {
        'dominant_freq': dominant_freq,
        'psd_mean': np.mean(psd) if len(psd) > 0 else 0,
        'psd_max': np.max(psd) if len(psd) > 0 else 0
    }


def handle_command(ser, command, current_time, current_lane, last_slow_down_time, last_lane_change_time):
    global last_sent_command_name, last_sent_command_time
    if not send_commands:
        return last_slow_down_time, current_lane, last_lane_change_time
    if command == 'LaneChangeLeft':
        if current_lane > 1 and current_time - last_lane_change_time > 5:
            if send_data_to_bluetooth(ser, command_mapping[command]):
                last_sent_command_name = command
                last_sent_command_time = current_time
                return current_time, current_lane - 1, current_time
        else:
            print("Invalid left lane change request.")
            return last_slow_down_time, current_lane, last_lane_change_time
    elif command == 'LaneChangeRight':
        if current_lane < 5 and current_time - last_lane_change_time > 5:
            if send_data_to_bluetooth(ser, command_mapping[command]):
                last_sent_command_name = command
                last_sent_command_time = current_time
                return last_slow_down_time, current_lane + 1, current_time
        else:
            print("Invalid right lane change request.")
            return last_slow_down_time, current_lane, last_lane_change_time
    elif command == 'Decelerate':
        if current_time - last_slow_down_time > 2:
            if send_data_to_bluetooth(ser, command_mapping[command]):
                last_sent_command_name = command
                last_sent_command_time = current_time
                return current_time, current_lane, last_lane_change_time
        else:
            print("Decelerate command too frequent, ignored.")
            return last_slow_down_time, current_lane, last_lane_change_time
    elif command != 'InvalidCommand':
        if send_data_to_bluetooth(ser, command_mapping[command]):
            last_sent_command_name = command
            last_sent_command_time = current_time
            return last_slow_down_time, current_lane, last_lane_change_time
    return last_slow_down_time, current_lane, last_lane_change_time


custom_objects = {'LeakyReLU': keras.layers.LeakyReLU}
model = keras.models.load_model('EEG_4class_trained_model_newUI.h5', custom_objects=custom_objects)
model.compile(optimizer=keras.optimizers.Adam(),
              loss='categorical_crossentropy',
              metrics=['accuracy'])


class MainWindow(QWidget):
    def __init__(self):
        self.eeg_bands = ['Delta', 'Theta', 'Low Alpha', 'High Alpha', 'Low Beta', 'High Beta', 'Low Gamma', 'High Gamma']
        self.eeg_band_values = [0] * 8

        super().__init__()
        self.setWindowTitle("Brain-Controlled Emergency Takeover System")
        self.setFixedSize(1400, 800)

        self.setAutoFillBackground(True)
        palette = self.palette()
        palette.setColor(QPalette.Window, Qt.white)
        self.setPalette(palette)

        self.signal_values = []
        self.attention_values = []
        self.relaxation_values = []
        self.confidence_values = []
        self.time_points = []

        self.init_ui()
        self.update_interface()

        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_interface)
        self.timer.start(1000)

    def bordered_widget(self, inner_widget, border_color):
        frame = QFrame()
        frame.setStyleSheet(f"border: 3px solid {border_color}; border-radius: 6px;")
        layout = QVBoxLayout(frame)
        layout.setContentsMargins(2, 2, 2, 2)
        layout.addWidget(inner_widget)
        return frame

    def neon_frame(self, widget, border_color):
        widget.setStyleSheet(f"""
            background-color: #121625;
            border: 2px solid {border_color};
            border-radius: 10px;
        """)

    def init_ui(self):
        main_layout = QHBoxLayout()
        CYBER_BORDER = '#0099FF'

        # Left layout: car video + EEG signal plot
        left_layout = QVBoxLayout()

        self.car_image_label = QLabel()
        self.car_image_label.setAlignment(Qt.AlignCenter)
        self.car_image_label.setStyleSheet("background-color: #FFFFFF;")
        left_layout.addWidget(self.bordered_widget(self.car_image_label, CYBER_BORDER))

        self.fig1 = Figure(figsize=(6, 3), facecolor='white')
        self.ax1 = self.fig1.add_subplot(111)
        self.line1, = self.ax1.plot([], [], label='Signal', color='#FF4B91')
        self.line2, = self.ax1.plot([], [], label='Attention', color='#00E0FF')
        self.line3, = self.ax1.plot([], [], label='Relaxation', color='#7CFFCB')
        self.ax1.set_facecolor('#FFFFFF')
        self.ax1.tick_params(colors='black')
        self.ax1.set_title('EEG Signal', color='black')
        self.ax1.set_xlabel('Time', color='black')
        self.ax1.set_ylabel('Value', color='black')
        self.ax1.legend(facecolor='#FFFFFF', edgecolor='black', labelcolor='black')
        self.canvas1 = FigureCanvas(self.fig1)
        self.neon_frame(self.canvas1, '#00FFFF')
        left_layout.addWidget(self.bordered_widget(self.canvas1, CYBER_BORDER))

        # Center layout: confidence chart + lane chart + prediction
        center_layout = QVBoxLayout()
        self.fig2 = Figure(figsize=(6, 2.5), facecolor='white')
        self.ax2 = self.fig2.add_subplot(111)
        self.fig2.subplots_adjust(bottom=0.25)
        self.confidence_line, = self.ax2.plot([], [], label='Confidence', color='#FFA500')
        self.ax2.set_facecolor('#FFFFFF')
        self.ax2.tick_params(colors='black')
        self.ax2.set_title('Confidence Trend', color='black')
        self.ax2.set_xlabel('Time', color='black')
        self.ax2.set_ylabel('Confidence', color='black')
        self.ax2.legend(facecolor='#FFFFFF', edgecolor='black', labelcolor='black')
        self.canvas2 = FigureCanvas(self.fig2)
        self.neon_frame(self.canvas2, '#FF00FF')
        center_layout.addWidget(self.bordered_widget(self.canvas2, CYBER_BORDER))

        self.canvas3 = FigureCanvas(Figure(figsize=(6, 1.5)))
        self.neon_frame(self.canvas3, '#7CFFCB')
        center_layout.addWidget(self.bordered_widget(self.canvas3, CYBER_BORDER))

        self.prediction_label = QLabel("Prediction")
        self.prediction_label.setFont(QFont("Arial", 28, QFont.Bold))
        self.prediction_label.setAlignment(Qt.AlignCenter)
        self.prediction_label.setStyleSheet("color: #00FFFF; background-color: #FFFFFF; padding: 10px;")
        center_layout.addWidget(self.bordered_widget(self.prediction_label, CYBER_BORDER))

        # Right layout: equipment image + radar chart + start/stop buttons
        right_layout = QVBoxLayout()

        self.equipment_image_label = QLabel()
        self.equipment_image_label.setAlignment(Qt.AlignCenter)
        self.equipment_image_label.setStyleSheet("background-color: #FFFFFF;")
        right_layout.addWidget(self.bordered_widget(self.equipment_image_label, CYBER_BORDER))

        self.radar_fig = Figure(figsize=(3.5, 3.4), facecolor='white')
        self.radar_ax = self.radar_fig.add_subplot(111)
        self.radar_ax.set_facecolor('#FFFFFF')
        self.radar_ax.tick_params(colors='black')
        self.radar_ax.set_title('EEG Band Power', color='black')
        self.radar_ax.set_ylabel('Power Value', color='black')
        self.radar_ax.set_xticks(range(8))
        self.radar_ax.set_xticklabels(self.eeg_bands, rotation=30, ha='right', fontsize=8, color='black')

        self.radar_bars = self.radar_ax.bar(range(8), self.eeg_band_values, color='#7CFFCB')
        self.radar_canvas = FigureCanvas(self.radar_fig)
        self.neon_frame(self.radar_canvas, '#ADFF2F')
        right_layout.addWidget(self.bordered_widget(self.radar_canvas, CYBER_BORDER))

        self.start_button = QPushButton("▶ Start Takeover")
        self.start_button.setStyleSheet("""
            QPushButton {
                background-color: #00E0FF;
                color: black;
                font-size: 24px;
                padding: 20px;
                border-radius: 10px;
            }
            QPushButton:hover {
                background-color: #3CFFFB;
            }
        """)
        self.start_button.clicked.connect(self.start_sending)

        self.stop_button = QPushButton("■ Stop")
        self.stop_button.setStyleSheet("""
            QPushButton {
                background-color: #FF4B91;
                color: white;
                font-size: 24px;
                padding: 20px;
                border-radius: 10px;
            }
            QPushButton:hover {
                background-color: #FF80B3;
            }
        """)
        self.stop_button.clicked.connect(self.stop_sending)

        right_layout.addWidget(self.start_button)
        right_layout.addWidget(self.stop_button)

        main_layout.addLayout(left_layout, 2)
        main_layout.addLayout(center_layout, 2)
        main_layout.addLayout(right_layout, 1)
        self.setLayout(main_layout)

    def start_sending(self):
        global send_commands
        send_commands = True
        print("Start sending commands")

    def stop_sending(self):
        global send_commands, ser
        send_commands = False
        try:
            if ser.is_open:
                send_data_to_bluetooth(ser, command_mapping['Stop'])
        except NameError:
            print("Serial port not initialized, cannot send Stop command.")
        print("Stop sending commands")

    def update_interface(self):
        try:
            ser = serial.Serial(port, 9600, timeout=1)
            csv_data = pd.read_csv("C:/Users/asus/Desktop/naobo.csv", encoding='iso-8859-1')
            if len(csv_data) < num_rows_to_read:
                print("Not enough rows, waiting for update...")
            else:
                # Data processing steps omitted (same structure as original)
                pass
        except FileNotFoundError:
            print("CSV file not found, please check the path.")
        except pd.errors.ParserError:
            print("Error parsing CSV file, please check the format.")
        except Exception as e:
            print(f"Unknown error: {e}")


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
