import os
import gradio as gr
import numpy as np
from PIL import Image
from tqdm import tqdm
import matplotlib.pyplot as plt
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)

# Disable GPU usage by default
os.environ['CUDA_VISIBLE_DEVICES'] = ''

# HDC Encoding and Decoding Functions
DIM = 1000  # Hypervector dimensionality

def pixel_to_hypervector(pixel):
    """Convert a pixel intensity to a high-dimensional binary hypervector."""
    np.random.seed(int(pixel))
    return np.random.choice([1, -1], size=(DIM,))

def image_to_hdc(image):
    """Encode the entire image into hypervectors (by pixel intensity)."""
    return np.array([pixel_to_hypervector(p) for p in image.flatten()])

def hdc_to_image(hdc_vectors, shape):
    """Decode hypervectors back into an image."""
    decoded_pixels = np.mean(hdc_vectors, axis=1)  # Aggregate hypervector values
    decoded_pixels = np.clip((decoded_pixels + 1) / 2 * 255, 0, 255)  # Rescale to [0, 255]
    return decoded_pixels.reshape(shape).astype(np.uint8)

class SwarmAgent:
    def __init__(self, position, velocity):
        self.position = position
        self.velocity = velocity

class SwarmNeuralNetwork:
    def __init__(self, num_agents, image_shape, target_image):
        self.image_shape = image_shape
        self.agents = [SwarmAgent(self.random_position(), self.random_velocity()) for _ in range(num_agents)]
        self.target_image = target_image

    def random_position(self):
        return np.random.randn(*self.image_shape)

    def random_velocity(self):
        return np.random.randn(*self.image_shape) * 0.01

    def update_agents(self, timestep):
        for agent in self.agents:
            # Convert agent's position and target image into HDC space
            agent_hdc = image_to_hdc(agent.position)
            target_hdc = image_to_hdc(self.target_image)

            # Compute similarity between the agent's position and the target image
            similarity = np.mean(agent_hdc * target_hdc, axis=1)  # Cosine-like similarity
            attention = similarity / np.sum(similarity)

            # Adjust the agent's position based on HDC-guided noise reduction
            noise = np.random.randn(*self.image_shape) * 0.1
            agent.position += attention.reshape(self.image_shape) * noise

            # Clip values to ensure valid range
            agent.position = np.clip(agent.position, -1, 1)

    def generate_image(self):
        generated_image = np.mean([agent.position for agent in self.agents], axis=0)
        return (generated_image + 1) / 2

    def train(self, epochs):
        for epoch in tqdm(range(epochs), desc="Training"):
            self.update_agents(epoch)

        return self.generate_image()

def preprocess_image(image):
    """Preprocess the input image."""
    if image is None:
        raise ValueError("No image provided")
    
    if isinstance(image, np.ndarray):
        # If it's already a numpy array, just resize and normalize
        image = Image.fromarray(image)
    elif isinstance(image, str):
        # If it's a file path, open the image
        image = Image.open(image)
    else:
        raise ValueError("Unsupported image type")
    
    image = image.convert('RGB').resize((128, 128))
    return np.array(image) / 127.5 - 1

def train_snn(image, num_agents, epochs):
    try:
        logging.info(f"Received image type: {type(image)}")
        preprocessed_image = preprocess_image(image)
        logging.info(f"Preprocessed image shape: {preprocessed_image.shape}")
        
        snn = SwarmNeuralNetwork(num_agents=num_agents, image_shape=(128, 128, 3), target_image=preprocessed_image)
        generated_image = snn.train(epochs=epochs)
        return (generated_image * 255).astype(np.uint8)
    except Exception as e:
        logging.error(f"Error in train_snn: {str(e)}")
        return None

interface = gr.Interface(
    fn=train_snn,
    inputs=[
        gr.Image(type="numpy", label="Upload Target Image"),
        gr.Slider(minimum=100, maximum=1000, value=500, step=50, label="Number of Agents"),
        gr.Slider(minimum=5, maximum=20, value=10, step=1, label="Number of Epochs")
    ],
    outputs=gr.Image(type="numpy", label="Generated Image"),
    title="HDC Swarm Neural Network Image Generation"
)

interface.launch()