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, Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt

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

# Pandas display settings
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)


# Model reference columns
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 at a 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.8
right_lane_change_threshold = threshold * 0.5
slow_down_threshold = threshold * 0.5
left_lane_change_threshold = threshold * 0.5
confidence_threshold = 0.5

# Bluetooth port
port = 'COM3'

# State tracking
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 (too soon or already at leftmost lane).")
            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 (too soon or already at rightmost lane).")
            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 sent within 2s, 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_LSTM.h5', custom_objects=custom_objects)
model.compile(optimizer=keras.optimizers.Adam(),
              loss='categorical_crossentropy',
              metrics=['accuracy'])


class MainWindow(QWidget):
    def __init__(self):
        super().__init__()

        self.setWindowTitle("Driver Brain-Controlled Emergency Takeover System")

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

        self.frame1 = QFrame()
        self.frame2 = QFrame()
        self.frame3 = QFrame()
        self.frame4 = QFrame()

        self.init_ui()
        self.update_interface()

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

    def init_ui(self):
        main_layout = QHBoxLayout()
        left_layout = QVBoxLayout()
        right_layout = QVBoxLayout()

        # First chart: Signal, Attention, Relaxation
        self.fig1 = Figure(figsize=(6, 4))
        self.ax1 = self.fig1.add_subplot(111)
        self.line1, = self.ax1.plot([], [], label='Signal', color='red')
        self.line2, = self.ax1.plot([], [], label='Attention', color='blue')
        self.line3, = self.ax1.plot([], [], label='Relaxation', color='green')
        self.ax1.set_xlabel('Time')
        self.ax1.set_ylabel('Value')
        self.ax1.set_title('Signal, Attention, Relaxation')
        self.ax1.legend()
        self.canvas1 = FigureCanvas(self.fig1)
        left_layout.addWidget(self.canvas1)

        # Prediction label
        self.prediction_label = QLabel("None")
        self.prediction_label.setStyleSheet("font-size: 30px;")
        self.prediction_label.setAlignment(Qt.AlignCenter)
        left_layout.addWidget(self.prediction_label)

        # Confidence chart
        self.fig2 = Figure(figsize=(6, 4))
        self.ax2 = self.fig2.add_subplot(111)
        self.confidence_line, = self.ax2.plot([], [], label='Confidence', color='orange')
        self.ax2.set_xlabel('Time')
        self.ax2.set_ylabel('Confidence')
        self.ax2.set_title('Prediction Confidence')
        self.ax2.legend()
        self.canvas2 = FigureCanvas(self.fig2)
        right_layout.addWidget(self.canvas2)

        # Lane display
        self.canvas3 = FigureCanvas(Figure(figsize=(6, 2)))
        right_layout.addWidget(self.canvas3)

        # Buttons
        button_layout = QHBoxLayout()
        self.start_button = QPushButton("Start Sending")
        self.start_button.clicked.connect(self.start_sending)
        self.stop_button = QPushButton("Stop")
        self.stop_button.clicked.connect(self.stop_sending)
        button_layout.addWidget(self.start_button)
        button_layout.addWidget(self.stop_button)
        right_layout.addLayout(button_layout)

        main_layout.addLayout(left_layout)
        main_layout.addLayout(right_layout)
        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):
        global last_slow_down_time, last_lane_change_time, last_sent_command_time, last_sent_command_name, current_lane, send_commands
        try:
            ser = serial.Serial(port, 9600, timeout=1)
            csv_data = pd.read_csv("C:/Users/26765/Desktop/naobo.csv", encoding='iso-8859-1')
            if len(csv_data) < num_rows_to_read:
                print("Not enough rows, waiting for more data...")
            else:
                # Processing omitted for brevity (same structure as original)
                pass

        except FileNotFoundError:
            print("CSV file not found, please check file path.")
        except pd.errors.ParserError:
            print("Error parsing CSV file, please check file format.")
        except Exception as e:
            print(f"Unknown error while reading file: {e}")


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