"""
U1FieldNode (Electromagnetism Metaphor)

Simulates a U(1) gauge force, like electromagnetism.
It takes a grayscale "charge density" map and calculates
the resulting force field (like an E-field).

[FIXED] Initialized self.potential in __init__ and
saved potential to self.potential in step().
"""

import numpy as np
import cv2

# --- Magic import block ---
import __main__
BaseNode = __main__.BaseNode
QtGui = __main__.QtGui
# --------------------------

class U1FieldNode(BaseNode):
    """
    Generates a U(1) force field from a charge density map.
    """
    NODE_CATEGORY = "Generator"
    NODE_COLOR = QtGui.QColor(100, 150, 220) # Blue

    def __init__(self, size=128):
        super().__init__()
        self.node_title = "U(1) Field (E/M)"
        
        self.inputs = {
            'charge_in': 'image',    # Grayscale image (0-1)
            'strength': 'signal'     # 0-1, force strength
        }
        self.outputs = {
            'potential_out': 'image', # The scalar potential (blurred charge)
            'field_viz': 'image'      # Vector field visualization
        }
        
        self.size = int(size)
        
        # --- START FIX ---
        # Initialize the output variables to prevent AttributeError
        self.viz = np.zeros((self.size, self.size, 3), dtype=np.float32)
        self.potential = np.zeros((self.size, self.size), dtype=np.float32)
        # --- END FIX ---

    def _prepare_image(self, img):
        if img is None:
            return np.full((self.size, self.size), 0.5, dtype=np.float32)
        
        if img.dtype != np.float32: img = img.astype(np.float32)
        if img.max() > 1.0: img /= 255.0
            
        img_resized = cv2.resize(img, (self.size, self.size), 
                                 interpolation=cv2.INTER_LINEAR)
        
        if img_resized.ndim == 3:
            return cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
        return img_resized

    def step(self):
        # --- 1. Get Charge Density ---
        # Map input [0, 1] to charge [-1, 1]
        charge_density = (self._prepare_image(
            self.get_blended_input('charge_in', 'first')
        ) * 2.0) - 1.0
        
        strength = self.get_blended_input('strength', 'sum') or 1.0
        
        # --- 2. Calculate Potential ---
        # Simulate long-range 1/r potential by blurring
        # A large blur kernel simulates the 1/r falloff
        ksize = self.size // 4 * 2 + 1 # Must be odd
        
        self.potential = cv2.GaussianBlur(charge_density, (ksize, ksize), 0)
        
        # --- 3. Calculate Force Field (E-Field) ---
        # E = -∇V (Force is the negative gradient of potential)
        grad_x = -cv2.Sobel(self.potential, cv2.CV_32F, 1, 0, ksize=3) * strength
        grad_y = -cv2.Sobel(self.potential, cv2.CV_32F, 0, 1, ksize=3) * strength
        
        # --- 4. Create Visualization ---
        self.viz = np.zeros((self.size, self.size, 3), dtype=np.float32)
        step = 8 # Draw an arrow every 8 pixels
        for y in range(0, self.size, step):
            for x in range(0, self.size, step):
                vx = grad_x[y, x] * 20 # Scale for viz
                vy = grad_y[y, x] * 20
                
                pt1 = (x, y)
                pt2 = (int(np.clip(x + vx, 0, self.size-1)), 
                       int(np.clip(y + vy, 0, self.size-1)))
                
                # Color based on direction
                angle = np.arctan2(vy, vx) + np.pi
                hue = int(angle / (2 * np.pi) * 179) # 0-179 for OpenCV HSV
                color_hsv = np.uint8([[[hue, 255, 255]]])
                color_rgb = cv2.cvtColor(color_hsv, cv2.COLOR_HSV2RGB)[0][0]
                color_float = color_rgb.astype(np.float32) / 255.0
                
                # --- START FIX ---
                # Convert numpy.float32 to standard Python floats for OpenCV
                color_tuple = (float(color_float[0]), float(color_float[1]), float(color_float[2]))
                cv2.arrowedLine(self.viz, pt1, pt2, color_tuple, 1, cv2.LINE_AA)
                # --- END FIX ---

    def get_output(self, port_name):
        if port_name == 'potential_out':
            # Normalize potential [-max, +max] to [0, 1]
            p_max = np.max(np.abs(self.potential))
            if p_max == 0: return np.full((self.size, self.size), 0.5, dtype=np.float32)
            return (self.potential / (2 * p_max)) + 0.5
            
        elif port_name == 'field_viz':
            return self.viz
        return None

    def get_display_image(self):
        return self.viz