import array
import inlinecpp
import numpy as np

def output_planes_to_cook(cop_node):
    # This sample only modifies the color plane.
    return ("C",)

def required_input_planes(cop_node, output_plane):
    # This sample requires the color and alpha planes from the first input.
    if output_plane == "C":
        return ("0", "C", "0", "A")
    return ()

def cook(cop_node, plane, resolution):
    assert(plane == "C")
    input = cop_node.inputs()[0]
    print(input.type())
    xres = input.xRes()
    yres = input.yRes()
    print(xres,yres)
    
    # Get the color plane as an array.
    color = array.array("f", input.allPixelsAsString("C"))

    # Convert the color array into a NumPy array and reshape it into an image.
    np_color = np.array(color).reshape((xres, yres, 3))
    # Separate the R, G, and B channels.
    r_channel = np_color[:, :, 0]
    g_channel = np_color[:, :, 1]
    b_channel = np_color[:, :, 2]

    # Perform an FFT on each channel.
    r_fft = np.fft.fft2(r_channel)
    g_fft = np.fft.fft2(g_channel)
    b_fft = np.fft.fft2(b_channel)

    # Combine the channels back into one array.
    fft_color = np.stack((r_fft, g_fft, b_fft), axis=-1)

    # Flatten the array and convert it back to a Python array.
    fft_color_flat = fft_color.flatten()
    color = array.array("f", fft_color_flat.real)  # Use .real to get the real part of the complex numbers.

    # Set the pixels of the cooking plane with the FFT data.
    cop_node.setPixelsOfCookingPlaneFromString(color.tostring())