import numpy as np
import gradio as gr
from PIL import Image
from scipy import ndimage
import matplotlib.pyplot as plt
from bulk_bulge_generation import definitions, smooth
# from transformers import pipeline
import fastai
from fastcore.all import *
from fastai.vision.all import *

def apply_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1, center_smoothness=0.20):
    # 0.106 strength = .50
    # 0.106 strength = 1
    rows, cols = image.shape[:2]
    max_dim = max(rows, cols)
    
    #Normalize the positions
    # Y Needs to be flipped
    center_y = int(center[1] * rows)
    center_x = int(center[0] * cols)

    # Inverts the Y axis (Numpy is 0 index at top of image)
    center_y = abs(rows - center_y)

    print()
    print(rows, cols)
    print("y =", center_y, "/", rows)
    print("x =", center_x, "/", cols)
    print()
    
    pixel_radius = int(max_dim * radius)
    
    y, x = np.ogrid[:rows, :cols]
    y = (y - center_y) / max_dim
    x = (x - center_x) / max_dim
    
    # Calculate distance from center
    dist_from_center = np.sqrt(x**2 + y**2)
    
    # Calculate function values
    z = func(x, y)
    
    # Calculate gradients
    gy, gx = np.gradient(z)

    # Creating a sigmoid function to apply to masks
    def sigmoid(x, center, steepness):
        return 1 / (1 + np.exp(-steepness * (x - center)))
    
    print(radius)
    print(strength)
    print(edge_smoothness)
    print(center_smoothness)

    # Masking
    edge_mask = np.clip((radius - dist_from_center) / (radius * edge_smoothness), 0, 1)

    center_mask = np.clip((dist_from_center - radius * center_smoothness) / (radius * center_smoothness), 0, 1)

    mask = edge_mask * center_mask
    
    # Apply mask to gradients
    gx = gx * mask
    gy = gy * mask
    
    # Normalize gradient vectors
    magnitude = np.sqrt(gx**2 + gy**2)
    magnitude[magnitude == 0] = 1  # Avoid division by zero
    gx = gx / magnitude
    gy = gy / magnitude
    
    # Scale the effect (Play with the number 5)
    scale_factor = strength * np.log(max_dim) / 100  # Adjust strength based on image size
    gx = gx * scale_factor * mask
    gy = gy * scale_factor * mask
    
    # Create the mapping
    x_new = x + gx
    y_new = y + gy
    
    # Convert back to pixel coordinates
    x_new = x_new * max_dim + center_x
    y_new = y_new * max_dim + center_y
    
    # Ensure the new coordinates are within the image boundaries
    x_new = np.clip(x_new, 0, cols - 1)
    y_new = np.clip(y_new, 0, rows - 1)
    
    # Apply the transformation to each channel
    channels = [ndimage.map_coordinates(image[..., i], [y_new, x_new], order=1, mode='reflect') 
                for i in range(image.shape[2])]
    
    transformed_image = np.dstack(channels).astype(image.dtype)
    
    return transformed_image, (gx, gy)


def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
    """
    Create a gradient vector field visualization with option to reverse direction.
    
    :param gx: X-component of the gradient
    :param gy: Y-component of the gradient
    :param image_shape: Shape of the original image (height, width)
    :param step: Spacing between arrows
    :param reverse: If True, reverse the direction of the arrows
    :return: Gradient vector field as a numpy array (RGB image)
    """
    rows, cols = image_shape
    y, x = np.mgrid[step/2:rows:step, step/2:cols:step].reshape(2, -1).astype(int)
    
    # Calculate the scale based on image size
    max_dim = max(rows, cols)
    scale = max_dim / 1000  # Adjusted for longer arrows
    
    # Reverse direction if specified
    direction = -1 if reverse else 1
    
    fig, ax = plt.subplots(figsize=(cols/50, rows/50), dpi=100)
    ax.quiver(x, y, direction * gx[y, x], direction * -gy[y, x], 
              scale=scale, 
              scale_units='width', 
              width=0.002 * max_dim / 500,
              headwidth=8, 
              headlength=12, 
              headaxislength=0, 
              color='black',
              minshaft=2,
              minlength=0,
              pivot='tail')
    ax.set_xlim(0, cols)
    ax.set_ylim(rows, 0)
    ax.set_aspect('equal')
    ax.axis('off')
    
    fig.tight_layout(pad=0)
    fig.canvas.draw()
    vector_field = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    vector_field = vector_field.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close(fig)
    
    return vector_field


#############################
#    MAIN FUNCTION HERE
#############################

# pipeline = pipeline(task="image-classification", model="nick-leland/distortionml")

# Version Check 
print(f"NumPy version: {np.__version__}")
print(f"PyTorch version: {torch.__version__}")
print(f"FastAI version: {fastai.__version__}")

learn_bias = load_learner('model_bias.pkl')
learn_fresh = load_learner('model_fresh.pkl')


def transform_image(image, func_choice, randomization_check, radius, center_x, center_y, strength, reverse_gradient=True, spiral_frequency=1):
    I = np.asarray(Image.open(image))    

    def pinch(x, y):
        return x**2 + y**2

    def shift(x, y):
        return np.arctan2(y, x)

    def bulge(x, y):
        r = -np.sqrt(x**2 + y**2)
        return r 

    def spiral(x, y, frequency=1):
        r = np.sqrt(x**2 + y**2)
        theta = np.arctan2(y, x)
        return r * np.sin(theta - frequency * r)

    rng = np.random.default_rng()
    if randomization_check == True:
        radius, location, strength, edge_smoothness= definitions(rng)
        center_x = location[0]
        center_y = location[1]
            
    # Temporarily disabling and using these values.
    # edge_smoothness = 0.25 * strength
    # center_smoothness = 0.25 * strength
    edge_smoothness, center_smoothness = smooth(rng, strength)


    if func_choice == "Pinch":
        func = pinch
    elif func_choice == "Spiral":
        func = shift 
    elif func_choice == "Bulge":
        func = bulge
        edge_smoothness = 0
        center_smoothness = 0
    elif func_choice == "Volcano":
        func = bulge
    elif func_choice == "Shift Up":
        func = lambda x, y: spiral(x, y, frequency=spiral_frequency)


    transformed, (gx, gy) = apply_vector_field_transform(I, func, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
    vector_field = create_gradient_vector_field(gx, gy, I.shape[:2], reverse=reverse_gradient)

    # GRADIO CHANGE HERE
    # predictions = pipeline(transformed) 

    # Have to convert to image first
    result = Image.fromarray(transformed)

    result_bias = str(learn_bias.predict(result))
    result_fresh = str(learn_fresh.predict(result))
    print("Results")
    print(result_bias)
    print(result_fresh)

    return transformed, result_bias, result_fresh, vector_field

demo = gr.Interface(
    fn=transform_image,
    inputs=[
        gr.Image(type="filepath"),
        gr.Dropdown(["Pinch", "Spiral", "Shift Up", "Bulge", "Volcano"], value="Volcano", label="Function"), 
        gr.Checkbox(label="Randomize inputs?"),
        gr.Slider(0, 0.5, value=0.25, label="Radius (as fraction of image size)"),
        gr.Slider(0, 1, value=0.5, label="Center X"),
        gr.Slider(0, 1, value=0.5, label="Center Y"),
        gr.Slider(0, 1, value=0.5, label="Strength"),
        # gr.Slider(0, 1, value=0.5, label="Edge Smoothness"),
        # gr.Slider(0, 0.5, value=0.1, label="Center Smoothness")
        # gr.Checkbox(label="Reverse Gradient Direction"),
    ],
    outputs=[
        gr.Image(label="Transformed Image"),
        # gr.Image(label="Result", num_top_classes=2)
        gr.Textbox(label='Result Bias'),
        gr.Textbox(label='Result Fresh'),
        gr.Image(label="Gradient Vector Field")
    ],
    title="Image Transformation Demo!",
    description="This is the baseline function that will be used to generate the database for a machine learning model I am working on called 'DistortionMl'! The goal of this model is to detect and then reverse image transformations that can be generated here! You can read more about the project at this repository link : https://github.com/nick-leland/DistortionML. The main function that I was working on is the 'Bulge' function, I can't really guarantee that the others work well (;"
)

demo.launch(share=True)