import pywifi
from pywifi import const
import pygame
import threading
import time
import numpy as np
from scipy.optimize import least_squares

# Initialize pygame
pygame.init()

# Set up display
WIDTH, HEIGHT = 800, 600
WINDOW = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Real-Time Wi-Fi Signal Mapping")

# Colors
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
AP_COLOR = (0, 255, 0)
DEVICE_COLOR = (255, 0, 0)
TEXT_COLOR = (255, 255, 255)

# Initialize font
FONT = pygame.font.SysFont('Arial', 16)

# Wi-Fi scanner class
class WiFiScanner(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.daemon = True  # Allow thread to exit when main program exits
        self.wifi = pywifi.PyWiFi()
        self.iface = self.wifi.interfaces()[0]
        self.scan_results = []
        self.running = True

    def run(self):
        while self.running:
            self.iface.scan()
            time.sleep(1.5)  # Wait for scan to complete
            self.scan_results = self.iface.scan_results()
            time.sleep(1)

    def stop(self):
        self.running = False

# Function to estimate distance from RSSI
def rssi_to_distance(rssi, tx_power=-40, n=2):
    """
    Estimate distance based on RSSI.
    :param rssi: Received signal strength (dBm)
    :param tx_power: Transmit power (dBm)
    :param n: Path-loss exponent (environmental factor)
    :return: Estimated distance in meters
    """
    return 10 ** ((tx_power - rssi) / (10 * n))

# Function to perform trilateration
def trilaterate(positions, distances):
    """
    Estimate position based on distances to known points.
    :param positions: List of tuples [(x1,y1), (x2,y2), (x3,y3)]
    :param distances: List of distances [d1, d2, d3]
    :return: Estimated position (x, y)
    """
    def residuals(point, positions, distances):
        return [np.linalg.norm(np.array(point) - np.array(pos)) - d for pos, d in zip(positions, distances)]

    # Initial guess
    x0 = np.mean(positions, axis=0)
    res = least_squares(residuals, x0, args=(positions, distances))
    return res.x

# Main function
def main():
    scanner = WiFiScanner()
    scanner.start()

    running = True
    clock = pygame.time.Clock()

    # Simulated positions of known Wi-Fi APs (for demonstration purposes)
    wifi_aps = {}

    # Main loop
    while running:
        WINDOW.fill(BLACK)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # Draw device at center
        device_pos = (WIDTH // 2, HEIGHT // 2)
        pygame.draw.circle(WINDOW, DEVICE_COLOR, device_pos, 5)
        device_text = FONT.render('Your Device (0,0)', True, TEXT_COLOR)
        WINDOW.blit(device_text, (device_pos[0] + 10, device_pos[1]))

        # Process scan results
        aps = scanner.scan_results
        positions = []
        distances = []

        for ap in aps:
            ssid = ap.ssid
            bssid = ap.bssid
            rssi = ap.signal  # RSSI value
            # For this demo, assign random positions to APs if not already assigned
            if bssid not in wifi_aps:
                wifi_aps[bssid] = (np.random.randint(50, WIDTH - 50), np.random.randint(50, HEIGHT - 50))

            # Estimate distance
            distance = rssi_to_distance(rssi)
            positions.append(wifi_aps[bssid])
            distances.append(distance)

            # Draw AP
            ap_pos = wifi_aps[bssid]
            pygame.draw.circle(WINDOW, AP_COLOR, ap_pos, 5)
            ap_text = FONT.render(f'{ssid} ({rssi}dBm)', True, TEXT_COLOR)
            WINDOW.blit(ap_text, (ap_pos[0] + 10, ap_pos[1]))

            # Draw line representing distance
            pygame.draw.line(WINDOW, WHITE, device_pos, ap_pos, 1)
            # Draw circle representing estimated distance
            pygame.draw.circle(WINDOW, WHITE, ap_pos, int(distance), 1)

        # If we have at least 3 APs, perform trilateration
        if len(positions) >= 3:
            estimated_pos = trilaterate(positions[:3], distances[:3])
            pygame.draw.circle(WINDOW, (0, 0, 255), (int(estimated_pos[0]), int(estimated_pos[1])), 5)
            est_text = FONT.render('Estimated Position', True, TEXT_COLOR)
            WINDOW.blit(est_text, (int(estimated_pos[0]) + 10, int(estimated_pos[1])))

        pygame.display.flip()
        clock.tick(30)  # Limit to 30 FPS

    scanner.stop()
    pygame.quit()

if __name__ == "__main__":
    main()