import sys
import numpy as np
import random
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QTextEdit, QLabel, 
                             QTabWidget, QTableWidget, QTableWidgetItem, 
                             QHeaderView, QGroupBox, QSpinBox, QDoubleSpinBox,
                             QFormLayout)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure

class Particle:
    def __init__(self, dim, bounds):
        self.position = np.array([random.uniform(bounds[i][0], bounds[i][1]) for i in range(dim)])
        self.velocity = np.array([random.uniform(-1, 1) for _ in range(dim)])
        self.best_position = self.position.copy()
        self.best_value = float('inf')
        self.value = float('inf')
        
    def update_velocity(self, global_best_position, w=0.5, c1=1.5, c2=1.5):
        r1, r2 = random.random(), random.random()
        cognitive = c1 * r1 * (self.best_position - self.position)
        social = c2 * r2 * (global_best_position - self.position)
        self.velocity = w * self.velocity + cognitive + social
        
    def update_position(self, bounds):
        self.position += self.velocity
        # Apply bounds
        for i in range(len(self.position)):
            if self.position[i] < bounds[i][0]:
                self.position[i] = bounds[i][0]
            elif self.position[i] > bounds[i][1]:
                self.position[i] = bounds[i][1]
                
    def evaluate(self, cost_function):
        self.value = cost_function(self.position)
        if self.value < self.best_value:
            self.best_value = self.value
            self.best_position = self.position.copy()

class PSOThread(QThread):
    update_signal = pyqtSignal(str, int, float, list)
    finished_signal = pyqtSignal(list, list)
    
    def __init__(self, cost_function, bounds, num_particles=30, max_iter=100):
        super().__init__()
        self.cost_function = cost_function
        self.bounds = bounds
        self.num_particles = num_particles
        self.max_iter = max_iter
        self.dim = len(bounds)
        self.running = True
        
    def run(self):
        # Initialize particles
        particles = [Particle(self.dim, self.bounds) for _ in range(self.num_particles)]
        global_best_position = particles[0].position.copy()
        global_best_value = float('inf')
        
        # Find initial global best
        for particle in particles:
            particle.evaluate(self.cost_function)
            if particle.best_value < global_best_value:
                global_best_value = particle.best_value
                global_best_position = particle.best_position.copy()
        
        # PSO main loop
        iteration_data = []
        position_data = []
        
        for iteration in range(self.max_iter):
            if not self.running:
                break
                
            for particle in particles:
                particle.update_velocity(global_best_position)
                particle.update_position(self.bounds)
                particle.evaluate(self.cost_function)
                
                if particle.best_value < global_best_value:
                    global_best_value = particle.best_value
                    global_best_position = particle.best_position.copy()
            
            # Store data for plotting
            iteration_data.append(iteration + 1)
            position_data.append(global_best_position.copy())
            
            # Emit update signal
            self.update_signal.emit(
                f"Iteration {iteration+1}/{self.max_iter}", 
                iteration+1, 
                global_best_value, 
                global_best_position.tolist()
            )
            
        self.finished_signal.emit(iteration_data, position_data)
        
    def stop(self):
        self.running = False

class CircuitExample:
    def __init__(self, example_num):
        self.example_num = example_num
        self.R1 = example_num  # Ohms
        self.V_out = example_num  # Ohms
        self.C1 = self.C2 = 1/(example_num + 1)  # 1/s Ohms
        self.L1 = self.L2 = 0.5 + 0.1 * example_num  # s Ohms
        self.V1 = self.V2 = example_num  # Volts
        self.alpha = self.R1
        
    def get_description(self):
        return f"""
Example {self.example_num}:
- R1 = {self.R1} Ω
- V_out(s) = {self.V_out} Ω
- C1 = C2 = {self.C1:.3f}/s Ω
- L1 = L2 = {self.L2:.3f}s Ω
- V1 = V2 = {self.V1} V
- α = {self.alpha}
        """
        
    def theoretical_impedance(self, s):
        # Theoretical impedance: Z(s) = αs
        return self.alpha * s
        
    def cost_function(self, x):
        # x[0] is the estimated alpha
        # We'll evaluate at multiple s values to get a better estimate
        s_values = [0.1, 0.5, 1.0, 2.0, 5.0]
        error = 0
        for s in s_values:
            theoretical = self.theoretical_impedance(s)
            estimated = x[0] * s
            error += (theoretical - estimated) ** 2
        return error

class MplCanvas(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)
        
    def plot_convergence(self, iterations, best_values):
        self.fig.clear()
        ax = self.fig.add_subplot(111)
        ax.plot(iterations, best_values, 'b-', linewidth=2)
        ax.set_xlabel('Iteration')
        ax.set_ylabel('Best Cost Value')
        ax.set_title('PSO Convergence')
        ax.grid(True)
        self.draw()
        
    def plot_parameter_evolution(self, iterations, parameters):
        self.fig.clear()
        ax = self.fig.add_subplot(111)
        for i in range(len(parameters[0])):
            param_values = [p[i] for p in parameters]
            ax.plot(iterations, param_values, label=f'Parameter {i+1}')
        ax.set_xlabel('Iteration')
        ax.set_ylabel('Parameter Value')
        ax.set_title('Parameter Evolution')
        ax.legend()
        ax.grid(True)
        self.draw()

class PSOCircuitApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.examples = [CircuitExample(i+1) for i in range(10)]
        self.current_example = 0
        self.pso_thread = None
        self.init_ui()
        
    def init_ui(self):
        self.setWindowTitle("PSO Circuit Analysis")
        self.setGeometry(100, 100, 1200, 800)
        
        # Central widget and main layout
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        
        # Left panel for controls and info
        left_panel = QVBoxLayout()
        
        # Example selection
        example_group = QGroupBox("Circuit Examples")
        example_layout = QVBoxLayout()
        
        self.example_combo = QSpinBox()
        self.example_combo.setMinimum(1)
        self.example_combo.setMaximum(10)
        self.example_combo.valueChanged.connect(self.change_example)
        
        self.example_info = QTextEdit()
        self.example_info.setMaximumHeight(200)
        self.example_info.setReadOnly(True)
        
        example_layout.addWidget(QLabel("Select Example:"))
        example_layout.addWidget(self.example_combo)
        example_layout.addWidget(QLabel("Circuit Parameters:"))
        example_layout.addWidget(self.example_info)
        example_group.setLayout(example_layout)
        left_panel.addWidget(example_group)
        
        # PSO controls
        pso_group = QGroupBox("PSO Parameters")
        pso_layout = QFormLayout()
        
        self.num_particles_spin = QSpinBox()
        self.num_particles_spin.setMinimum(10)
        self.num_particles_spin.setMaximum(100)
        self.num_particles_spin.setValue(30)
        
        self.max_iter_spin = QSpinBox()
        self.max_iter_spin.setMinimum(10)
        self.max_iter_spin.setMaximum(500)
        self.max_iter_spin.setValue(100)
        
        self.w_spin = QDoubleSpinBox()
        self.w_spin.setMinimum(0.1)
        self.w_spin.setMaximum(2.0)
        self.w_spin.setValue(0.5)
        self.w_spin.setSingleStep(0.1)
        
        self.c1_spin = QDoubleSpinBox()
        self.c1_spin.setMinimum(0.1)
        self.c1_spin.setMaximum(3.0)
        self.c1_spin.setValue(1.5)
        self.c1_spin.setSingleStep(0.1)
        
        self.c2_spin = QDoubleSpinBox()
        self.c2_spin.setMinimum(0.1)
        self.c2_spin.setMaximum(3.0)
        self.c2_spin.setValue(1.5)
        self.c2_spin.setSingleStep(0.1)
        
        pso_layout.addRow("Number of Particles:", self.num_particles_spin)
        pso_layout.addRow("Maximum Iterations:", self.max_iter_spin)
        pso_layout.addRow("Inertia Weight (w):", self.w_spin)
        pso_layout.addRow("Cognitive Parameter (c1):", self.c1_spin)
        pso_layout.addRow("Social Parameter (c2):", self.c2_spin)
        
        pso_group.setLayout(pso_layout)
        left_panel.addWidget(pso_group)
        
        # Control buttons
        self.run_button = QPushButton("Run PSO")
        self.run_button.clicked.connect(self.run_pso)
        
        self.stop_button = QPushButton("Stop PSO")
        self.stop_button.clicked.connect(self.stop_pso)
        self.stop_button.setEnabled(False)
        
        left_panel.addWidget(self.run_button)
        left_panel.addWidget(self.stop_button)
        
        # Results display
        results_group = QGroupBox("Results")
        results_layout = QVBoxLayout()
        
        self.results_text = QTextEdit()
        self.results_text.setMaximumHeight(150)
        self.results_text.setReadOnly(True)
        
        results_layout.addWidget(self.results_text)
        results_group.setLayout(results_layout)
        left_panel.addWidget(results_group)
        
        # Add left panel to main layout
        main_layout.addLayout(left_panel, 1)
        
        # Right panel for plots
        right_panel = QVBoxLayout()
        
        # Tab widget for different plots
        self.plot_tabs = QTabWidget()
        
        # Convergence plot
        self.convergence_canvas = MplCanvas(self, width=5, height=4, dpi=100)
        self.plot_tabs.addTab(self.convergence_canvas, "Convergence")
        
        # Parameter evolution plot
        self.param_canvas = MplCanvas(self, width=5, height=4, dpi=100)
        self.plot_tabs.addTab(self.param_canvas, "Parameter Evolution")
        
        right_panel.addWidget(self.plot_tabs)
        
        # Add right panel to main layout
        main_layout.addLayout(right_panel, 2)
        
        # Initialize with first example
        self.change_example(1)
        
    def change_example(self, value):
        self.current_example = value - 1
        example = self.examples[self.current_example]
        self.example_info.setText(example.get_description())
        self.results_text.clear()
        
    def run_pso(self):
        example = self.examples[self.current_example]
        
        # Define bounds for alpha (0 to 2*expected alpha)
        bounds = [(0.1, 2 * example.alpha)]
        
        # Create and configure PSO thread
        self.pso_thread = PSOThread(
            example.cost_function,
            bounds,
            self.num_particles_spin.value(),
            self.max_iter_spin.value()
        )
        
        # Connect signals
        self.pso_thread.update_signal.connect(self.update_progress)
        self.pso_thread.finished_signal.connect(self.pso_finished)
        
        # Update UI
        self.run_button.setEnabled(False)
        self.stop_button.setEnabled(True)
        self.results_text.clear()
        self.results_text.append("Running PSO...")
        
        # Start PSO
        self.pso_thread.start()
        
    def stop_pso(self):
        if self.pso_thread and self.pso_thread.isRunning():
            self.pso_thread.stop()
            self.pso_thread.wait()
            self.results_text.append("PSO stopped by user.")
            self.run_button.setEnabled(True)
            self.stop_button.setEnabled(False)
            
    def update_progress(self, status, iteration, best_value, best_position):
        example = self.examples[self.current_example]
        self.results_text.clear()
        self.results_text.append(f"Status: {status}")
        self.results_text.append(f"Best Cost: {best_value:.6f}")
        self.results_text.append(f"Estimated α: {best_position[0]:.4f}")
        self.results_text.append(f"Theoretical α: {example.alpha}")
        self.results_text.append(f"Error: {abs(best_position[0] - example.alpha):.4f}")
        
    def pso_finished(self, iterations, positions):
        example = self.examples[self.current_example]
        best_alpha = positions[-1][0]
        
        self.results_text.append("\n--- PSO Completed ---")
        self.results_text.append(f"Final Estimated α: {best_alpha:.4f}")
        self.results_text.append(f"Theoretical α: {example.alpha}")
        self.results_text.append(f"Absolute Error: {abs(best_alpha - example.alpha):.4f}")
        self.results_text.append(f"Relative Error: {abs(best_alpha - example.alpha)/example.alpha*100:.2f}%")
        
        # Plot convergence
        best_values = [example.cost_function(p) for p in positions]
        self.convergence_canvas.plot_convergence(iterations, best_values)
        
        # Plot parameter evolution
        self.param_canvas.plot_parameter_evolution(iterations, positions)
        
        # Update UI
        self.run_button.setEnabled(True)
        self.stop_button.setEnabled(False)

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