import numpy as np
import cv2
from scipy.fft import ifft2, ifftshift

import __main__
try:
    BaseNode = __main__.BaseNode
    QtGui = __main__.QtGui
except AttributeError:
    class BaseNode: 
        def get_blended_input(self, name, mode): return None
    import PyQt6.QtGui as QtGui

class HolographicIFFTNodeX(BaseNode):
    """
    Holographic IFFT - The Inverse Reality
    ======================================
    
    Treats the input image not as "Space", but as "Frequency" (k-space).
    It effectively asks: "If this image were a diffraction pattern, 
    what hidden object created it?"
    
    MECHANISM:
    1. Input (Gradient/Field) -> Treated as MAGNITUDE spectrum.
    2. Phase Synthesis -> We generate the missing phase (The "Light").
       - 'Void': Phase = 0 (Autocorrelation)
       - 'Chaos': Random Phase
       - 'Structure': Phase derived from image structure
    3. Inverse FFT -> Collapses the spectrum into a spatial image.
    
    This reveals the "Reciprocal Ghost" - the hidden symmetries that 
    exist in the frequency domain of your field.
    """
    NODE_CATEGORY = "Analysis"
    NODE_TITLE = "Holographic IFFT 2"
    NODE_COLOR = QtGui.QColor(100, 255, 200) # Spectral Cyan
    
    def __init__(self):
        super().__init__()
        
        self.inputs = {
            'diffraction_pattern': 'image', # The Gradient Field (treated as FFT Mag)
            'phase_mod': 'signal',          # Rotate the phase (The "Laser" angle)
            'frequency_scale': 'signal',    # Zoom in k-space
            'reset': 'signal'
        }
        
        self.outputs = {
            'hologram': 'image',            # The reconstructed "Ghost"
            'phase_mask': 'image'           # The phase we used
        }
        
        self.size = 128
        self.phase_mode = 'structure' # 'void', 'chaos', 'structure'
        
        # Internal state
        self.phase_map = None
        self.t = 0.0

    def step(self):
        # 1. Get Inputs
        pattern = self.get_blended_input('diffraction_pattern', 'first')
        p_mod = self.get_blended_input('phase_mod', 'sum')
        f_scale = self.get_blended_input('frequency_scale', 'sum')
        
        if pattern is None: return
        
        # Defaults
        if p_mod is None: p_mod = 0.0
        if f_scale is None: f_scale = 1.0
        
        # Resize/Format
        if pattern.shape[:2] != (self.size, self.size):
            pattern = cv2.resize(pattern, (self.size, self.size))
        if pattern.ndim == 3:
            pattern = np.mean(pattern, axis=2)
            
        # 2. Treat Input as MAGNITUDE (The Diffraction Pattern)
        # We shift it so the center of the image is DC (0 frequency)
        magnitude = np.abs(pattern)
        
        # Apply Frequency Scaling (Zooming in Reciprocal Space)
        if f_scale != 1.0 and f_scale > 0.1:
            # Zooming the spectrum changes the size of the object
            center = self.size // 2
            M = cv2.getRotationMatrix2D((center, center), 0, f_scale)
            magnitude = cv2.warpAffine(magnitude, M, (self.size, self.size))
        
        # 3. Synthesize the Missing Phase
        # This is where we "change the values" to catch the ghost
        if self.phase_map is None:
            self.phase_map = np.zeros_like(magnitude)
            
        # Create a phase ramp (spatial offset) + modulation
        y, x = np.ogrid[:self.size, :self.size]
        
        # 'Structure' Mode: Phase is related to position (creates coherence)
        # We animate the phase over time to "scan" the hologram
        self.t += 0.05
        phase_structure = (x * np.cos(self.t) + y * np.sin(self.t)) * (0.1 + p_mod * 0.1)
        
        # Combine Magnitude and Phase into Complex Spectrum
        # Z = Mag * e^(i * theta)
        complex_spectrum = magnitude * np.exp(1j * phase_structure)
        
        # 4. The Inverse FFT (The Reconstruction)
        # We assume the input image had (0,0) in top-left, so we don't shift first.
        # But usually spectral view has DC in center. Let's try shifting.
        complex_spectrum = ifftshift(complex_spectrum)
        
        reconstructed = ifft2(complex_spectrum)
        
        # 5. Extract Real Component (The Hologram)
        self.hologram = np.abs(reconstructed)
        self.current_phase = np.angle(complex_spectrum)

    def get_output(self, port_name):
        if port_name == 'hologram':
            if hasattr(self, 'hologram'):
                return self._normalize(self.hologram)
        elif port_name == 'phase_mask':
            if hasattr(self, 'current_phase'):
                return self._normalize(self.current_phase)
        return None

    def _normalize(self, img):
        img = np.nan_to_num(img)
        norm = (img - np.min(img)) / (np.max(img) - np.min(img) + 1e-10)
        return (norm * 255).astype(np.uint8)

    def get_display_image(self):
        if not hasattr(self, 'hologram'): return None
        
        # Display: The Hologram (Reconstructed Reality)
        img = self._normalize(self.hologram)
        display = cv2.applyColorMap(img, cv2.COLORMAP_OCEAN)
        
        cv2.putText(display, "Holographic IFFT", (5, 15), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255,255,255), 1)
        
        return QtGui.QImage(display.data, self.size, self.size, 
                           self.size*3, QtGui.QImage.Format.Format_RGB888)
    
    def get_config_options(self):
        return [
             ("Resolution", "size", 128, None),
             ("Scale", "scale", 1.0, None)
        ]