import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.layers import Conv2D, BatchNormalization, LeakyReLU, Flatten, Dense, Reshape, Dropout, Add

IMG_SIZE = 512

def residual_block(X,filters):
    # Retrieve Filters
    F1, F2 = filters
    # Saving the input value.we need this later to add to the output. 
    X_shortcut = X
    
    # First component of main path
    X = Conv2D(filters = F1, kernel_size = (3, 3), strides = (1,1), padding = 'same')(X)
    X = BatchNormalization()(X)
    X = LeakyReLU(alpha=0.1)(X)

    # Second component of main path 
    X = Conv2D(filters = F2, kernel_size = (3, 3), strides = (1,1), padding = 'same')(X)
    X = BatchNormalization()(X)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation 
    X = Add()([X, X_shortcut])
    X = LeakyReLU(alpha=0.1)(X)
    return X

def build_model():
    # Inputs to the model
    base_model = ResNet50(
        weights='imagenet',
        input_shape=(512, 512, 3),  # Input shape of the images (height, width, channels)
        include_top=False  # Exclude the top classification layers
    )

    # Freeze the base model's layers to prevent them from being trained
    base_model.trainable = False
    x = base_model.output
    
    # First conv block
    x = Conv2D(32,(3, 3),kernel_initializer="he_normal",padding="same")(x)
    x = LeakyReLU(alpha=0.1)(x)
    
    # Add a residual block
    x = residual_block(x, [64, 32]) 
    x = residual_block(x, [64, 32]) 
    x = Flatten()(x)
    
    x = Dense(128, kernel_initializer='he_normal')(x)
    x = Dropout(0.2)(x)
    x = Dense(64, kernel_initializer='he_normal')(x)
    x = Dropout(0.2)(x)
    
    x = Dense(1, activation='sigmoid')(x)
    model = Model(inputs=base_model.input, outputs=x)
    
    return model


stage1_cc = build_model()
stage1_cc.load_weights("weights/stage1_cc_weights.weights.h5")

stage1_mlo = build_model()
stage1_mlo.load_weights("weights/stage1_mlo_weight.weights.h5")

stage2_cc = build_model()
stage2_cc.load_weights("weights/stage2_cc_weights.weights.h5")

stage2_mlo = build_model()
stage2_mlo.load_weights("weights/stage2_mlo_weight.weights.h5")


def get_diff(prior_image, current_image):
    print(prior_image.shape)
    print(current_image.shape)

    prior_image = np.array(prior_image, dtype=np.float32)
    current_image = np.array(current_image, dtype=np.float32)

    prior_image = prior_image.astype(np.float32)
    current_image = current_image.astype(np.float32)

    avg_width = int((prior_image.shape[0] + current_image.shape[0])/2)
    avg_height = int((prior_image.shape[1] + current_image.shape[1])/2)

    # print(avg_width, avg_width)

    prior_image = np.resize(prior_image, [avg_width, avg_height, 3])
    current_image = np.resize(current_image, [avg_width, avg_height, 3])

    subtract_image =current_image - prior_image
    subtract_image[subtract_image<0] = 0
    return subtract_image

def stage1_run(cc_diff_img, mlo_diff_img):
    cc_diff_img = np.resize(cc_diff_img, [IMG_SIZE, IMG_SIZE, 3])
    cc_diff_img = np.expand_dims(cc_diff_img, axis=0)
    cc_diff_img = tf.constant(cc_diff_img, dtype=tf.float32)
    # print(cc_diff_img.shape)
    cc_res = stage1_cc.predict(cc_diff_img)
    # print(cc_res)
    # mlo_res = stage1_mlo.predict(mlo_diff_img)
    mlo_diff_img = np.resize(mlo_diff_img, [IMG_SIZE, IMG_SIZE, 3])
    mlo_diff_img = np.expand_dims(mlo_diff_img, axis=0)
    mlo_diff_img = tf.constant(mlo_diff_img, dtype=tf.float32)

    mlo_res = stage1_mlo.predict(mlo_diff_img)
    # print(mlo_res)
    
    return (cc_res + mlo_res)/2

def stage2_run(cc_diff_img, mlo_diff_img):
    cc_diff_img = np.resize(cc_diff_img, [IMG_SIZE, IMG_SIZE, 3])
    cc_diff_img = np.expand_dims(cc_diff_img, axis=0)
    cc_diff_img = tf.constant(cc_diff_img, dtype=tf.float32)
    # print(cc_diff_img.shape)
    cc_res = stage2_cc.predict(cc_diff_img)
    # print(cc_res)
    # mlo_res = stage1_mlo.predict(mlo_diff_img)
    mlo_diff_img = np.resize(mlo_diff_img, [IMG_SIZE, IMG_SIZE, 3])
    mlo_diff_img = np.expand_dims(mlo_diff_img, axis=0)
    mlo_diff_img = tf.constant(mlo_diff_img, dtype=tf.float32)

    mlo_res = stage2_mlo.predict(mlo_diff_img)
    # print(mlo_res)
    
    return (cc_res + mlo_res)/2

def give_result(cc_prior_image, mlo_prior_image, cc_recent_image, mlo_recent_image):
    cc_prior_img = np.array(cc_prior_image)
    mlo_prior_img = np.array(mlo_prior_image)
    cc_recent_img = np.array(cc_recent_image)
    mlo_recent_img = np.array(mlo_recent_image)

    cc_diff_img = get_diff(cc_prior_img, cc_recent_img)
    mlo_diff_img = get_diff(mlo_prior_img, mlo_recent_img)

    stage1_res = stage1_run(cc_diff_img, mlo_diff_img)

    if(stage1_res<0.4):
        res = stage1_res.numpy()
        return f"Normal, you have {1-res} chance of being suspecious."
    
    stage2_res = stage2_run(cc_diff_img, mlo_diff_img)
    if(stage2_res<0.4):
        return "Benign"
    else:
        return "Suspecious, You need to go for biopsy immediately."

with gr.Blocks(title="Breast Cancer detection", css=".gradio-container {background:lightyellow;}") as demo:
    gr.HTML("<h1>Breast Cancer Detection</h1>")
    with gr.Row():
        cc_prior_image = gr.Image(label="CC Prior image", type="pil")
        mlo_prior_image = gr.Image(label="MLO Prior image", type="pil")

    with gr.Row():
        cc_recent_image = gr.Image(label="CC Recent image", type="pil")
        mlo_recent_image = gr.Image(label="MLO Recent image", type="pil")
    gr.HTML("<br/>")
    output_label = gr.TextArea(placeholder="Result")
    send_btn = gr.Button("Detect")
    send_btn.click(fn=give_result, inputs=[cc_prior_image, mlo_prior_image, cc_recent_image, mlo_recent_image], outputs=[output_label])

demo.launch(debug=True)