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 *
from ultralytics import YOLO

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')

# Loads the YOLO Model
model = YOLO("bulge_yolo_model.pt")


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 bulge_inverse(x, y, f=bulge, a=1, b=1, c=1, d=0, e=0):
        t = np.arctan2(y, x)
        term = ((f - e) / (-a))**2 - d
        if term < 0:
            return None, None

        x = (1/np.sqrt(b)) * np.sqrt(term) * np.cos(t)
        y = (1/np.sqrt(c)) * np.sqrt(term) * np.sin(t)

        return x, y

    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
        func2 = bulge_inverse
        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)


    # Original Image Transformation
    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)

    reverted, (gx_inverse, gy_inverse) = apply_vector_field_transform(I, func2, radius, (center_x, center_y), strength, edge_smoothness, center_smoothness)
    vector_field_reverted = create_gradient_vector_field(gx_inverse, gy_inverse, I.shape[:2], reverse=reverse_gradient)
    

    # GRADIO CHANGE HERE
    # predictions = pipeline(transformed) 

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

    categories = ['Distorted', 'Maze']

    def clean_output(result_values):
        pred, idx, probs = result_values[0], result_values[1], result_values[2]
        return dict(zip(categories, map(float, probs)))

    result_bias = learn_bias.predict(result)
    result_fresh = learn_fresh.predict(result)
    print("Results")
    result_bias_final = clean_output(result_bias)
    result_fresh_final = clean_output(result_fresh)

    print("saving?")
    result_localization = model.predict(transformed, save=True)
    print(result_localization)

    return transformed, result_bias_final, result_fresh_final, vector_field, vector_field_reverted

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"),
    ],
    examples=[
        [np.asarray(Image.open("examples/1500_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
        [np.asarray(Image.open("examples/2048_maze.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
        [np.asarray(Image.open("examples/2300_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5],
        [np.asarray(Image.open("examples/50_fresh.jpg")), "Bulge", True, 0.25, 0.5, 0.5, 0.5]
    ],
    outputs=[
        gr.Image(label="Transformed Image"),
        # gr.Image(label="Result", num_top_classes=2)
        gr.Label(),
        gr.Label(),
        gr.Image(label="Gradient Vector Field"),
        gr.Image(label="Gradient Vector Field Reverted")
    ],
    title="Image Transformation Demo!",
    article="If you like this demo, please star the github repository for the project! Located [here!](https://github.com/nick-leland/DistortionML)",
    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!\nYou 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'/'Volcano' function, I can't really guarantee that the others work as well!\nI have just added the first baseline ML model to detect if a distortion has taken place! It was only trained on mazes though ([Dataset Here](https://www.kaggle.com/datasets/nickleland/distorted-mazes)) so in order for it to detect a distortion you have to use one of the images provided in the examples! Feel free to mess around wtih other images in the meantime though!"
)

demo.launch(share=True)