import gradio as gr
import json
import os
import numpy as np
import cv2 as cv
from google.cloud import storage

GCS_BUCKET_NAME = "veytel-cloud-store"
GCS_FOLDER_PATH = "density_mapper"
service_account_json_str = os.getenv('serviceKey')
service_account_json = json.loads(service_account_json_str)

if service_account_json:
    print("Secret Value Retrieved Successfully")
    #print(service_account_json)
else:
    print("Failed to Retrieve Secret Value")

count = 0
fresh_start = False

def set_min_dense_1(max_dense_0):
    global scaled_thresh1
    print("max_dense_0", max_dense_0)
    scaled_thresh1 = int(max_dense_0) * max_val / 255
    dense_0 = np.where(textured_cxr < scaled_thresh1, textured_cxr, 0)
    dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
    return max_dense_0, dense_0, dense_1


def set_min_dense_2(max_dense_1):
    global scaled_thresh2
    print("max_dense_1", max_dense_1)
    scaled_thresh2 = int(max_dense_1) * max_val / 255
    dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
    dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
    return max_dense_1, dense_1, dense_2


def set_min_dense_3(max_dense_2):
    global scaled_thresh3
    print("max_dense_2", max_dense_2)
    scaled_thresh3 = int(max_dense_2) * max_val / 255
    dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
    dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
    return max_dense_2, dense_2, dense_3

def authenticate_gcs():
    return storage.Client.from_service_account_info(service_account_json)


def download_csv_from_gcs(filename):
    client = authenticate_gcs()
    bucket = client.get_bucket(GCS_BUCKET_NAME)
    blob = bucket.blob(os.path.join(GCS_FOLDER_PATH, filename))
    csv_content = blob.download_as_text()
    return csv_content

def read_csv_from_gcs(user):
    filename = f"density_{user}.csv"
    try:
        csv_content = download_csv_from_gcs(filename)
        rows = csv_content.strip().split("\n")
        csv_data = [row.split(",") for row in rows]
        return csv_data
    except Exception as e:
        print("Error reading CSV from GCS:", e)
        return []

def upload_csv_to_gcs(csv_content, filename):
    client = authenticate_gcs()
    bucket = client.get_bucket(GCS_BUCKET_NAME)
    blob = bucket.blob(os.path.join(GCS_FOLDER_PATH, filename))
    blob.upload_from_string(csv_content)

def new__cxr(max_dense_0, max_dense_1, max_dense_2):
    global textured_cxr, lung_noised, max_val, cxr, mask, scaled_thresh1, scaled_thresh2, scaled_thresh3, image_id, index, count, label1
    csv_data = read_csv_from_gcs(user)
    if count < 30:
        csv_content = ""
        for row in csv_data:
            csv_content += ",".join(row) + "\n"
        if count > 0:
            csv_content += f"{count},{max_dense_0},{max_dense_1},{max_dense_2}\n"
        filename = f"density_{user}.csv"
        upload_csv_to_gcs(csv_content, filename)

    if (count >= 30):
        csv_content = ""
        for row in csv_data:
            csv_content += ",".join(row) + "\n"
        if count > 0:
            csv_content += f"{count},{max_dense_0},{max_dense_1},{max_dense_2}\n"
        filename = f"density_{user}.csv"
        upload_csv_to_gcs(csv_content, filename)
        empty_image = np.zeros((256, 256), dtype=np.uint8)
        label1.update(visible=False)
        count = 0
        image_id = 1
        index = 1
        fieldnames = ['count', 'thresh_1', 'thresh_2', 'thresh_3']
        csv_content += ",".join(fieldnames) + "\n"
        filename = f"density_{user}.csv"
        upload_csv_to_gcs(csv_content, filename)
        return empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, 0


    if count >= 0:
        index += 1
    if index > 3:
        index = 1
        image_id += 1

    count += 1
    # write count, thresh1, thresh2, thresh3 to csv file
    cxr_file = "cxr" + str(image_id) + "_cxr.png"
    mask_file = "cxr" + str(image_id) + "_mask.png"
    textured_cxr_file = "cxr" + str(image_id) + "_textured_" + str(index) + ".png"
    lung_noised_file = "cxr" + str(image_id) + "_lung_noised_" + str(index) + ".png"
    cxr_path = os.path.join(cxr_dir, cxr_file)
    mask_path = os.path.join(mask_dir, mask_file)
    textured_cxr_path = os.path.join(textured_cxr_dir, textured_cxr_file)
    lung_noised_path = os.path.join(lung_noised_dir, lung_noised_file)
    cxr = cv.imread(cxr_path, cv.IMREAD_GRAYSCALE)
    mask = cv.imread(mask_path, cv.IMREAD_GRAYSCALE)
    textured_cxr = cv.imread(textured_cxr_path, cv.IMREAD_GRAYSCALE)
    lung_noised = cv.imread(lung_noised_path, cv.IMREAD_GRAYSCALE)
    max_val = np.percentile(cxr, 97)  # To optimize later
    thresh1 = 50
    thresh2 = 100
    thresh3 = 150
    scaled_thresh1 = thresh1 * max_val / 255
    scaled_thresh2 = thresh2 * max_val / 255
    scaled_thresh3 = thresh3 * max_val / 255
    dense_0 = np.where(textured_cxr < scaled_thresh1, textured_cxr, 0)
    dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
    dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
    dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
    return cxr, textured_cxr, lung_noised, lung_noised, dense_0, dense_1, dense_2, dense_3, count


def create_csv():
    global count, fieldnames, image_id, index, csv_path, fresh_start
    fieldnames = ['count', 'thresh_1', 'thresh_2', 'thresh_3']
    csv_data = read_csv_from_gcs(user)
    csv_content = ""
    last_row_count = 0
    if not csv_data:
        csv_content = ",".join(fieldnames) + "\n"
    else:
        for i, row in enumerate(csv_data):
            csv_content += ",".join(row) + "\n"
            if i != 0:
                if isinstance(row[0], str):
                    last_row_count = 0
                else:
                    last_row_count = int(row[0])

    #csv_content += f"{count+1},{thresh1},{thresh2},{thresh3}\n"
    filename = f"density_{user}.csv"
    upload_csv_to_gcs(csv_content, filename)
    count = last_row_count
    if(count>0):
        image_id, index = get_image_id_index(count)
    fresh_start = True


def check_auth(username, password):
    global user, fresh_start
    user = username
    if (user == 'gk' and password == 'upmc2023'):
        create_csv()
        return True
    elif (user == 'veytel' and password == 'pittsburgh'):
        create_csv()
        return True
    elif (user == 'cathy' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'ellen' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'kevin' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'swathi' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'mike' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'test' and password == 'test'):
        create_csv()
        return True
    elif (user == 'nischal' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'vijayakumar' and password == 'veytel'):
        create_csv()
        return True
    elif (user == 'konstantine' and password == 'upmc2023'):
        create_csv()
        return True
    elif (user == 'taaha' and password == 'upmc2023'):
        create_csv()
        return True
    elif (user == 'nameer' and password == 'upmc2023'):
        create_csv()
        return True
    elif (user == 'siddique' and password == 'upmc2023'):
        create_csv()
        return True


def change_vis():
    global count, fresh_start, button1, button2
    if fresh_start:
        fresh_start = not fresh_start
    if count >= 30:
        return gr.Label(visible=True), gr.Button(visible=False), gr.Button(visible=False)
    else:
        return gr.Label(visible=False), gr.Button(visible=True), gr.Button(visible=False)


image_id = 1
index = 1

def get_image_id_index(count):
    #global image_id, index
    if count ==0:
        return 1, 1
    image_id = (count - 1) // 3 + 1
    index = (count - 1) % 3 + 1
    return image_id, index

def set_layout():
    global textured_cxr, lung_noised, max_val, cxr, mask, scaled_thresh1, scaled_thresh2, scaled_thresh3, image_id, index, count
    csv_data = read_csv_from_gcs(user)
    csv_content = ""
    last_row_count = 0
    if csv_data:
        for i, row in enumerate(csv_data):
            if i != 0:
                if str(row[0]) == 'count':
                    last_row_count = 0
                else:
                    last_row_count = int(row[0])


    count = last_row_count
    if (count > 0):
        image_id, index = get_image_id_index(count)

    if count >= 30:
        empty_image = np.zeros((256, 256), dtype=np.uint8)
        return empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, empty_image, 30
    image_id, index = get_image_id_index(count)
    if count == 0:
        count = 1
        image_id = 1
        index = 1
    else:
        count +=1
        index += 1
        if index > 3:
            index = 1
            image_id += 1

    #new__cxr(max_dense_0, max_dense_1, max_dense_2)
    cxr_file = "cxr" + str(image_id) + "_cxr.png"
    mask_file = "cxr" + str(image_id) + "_mask.png"
    textured_cxr_file = "cxr" + str(image_id) + "_textured_" + str(index) + ".png"
    lung_noised_file = "cxr" + str(image_id) + "_lung_noised_" + str(index) + ".png"
    cxr_path = os.path.join(cxr_dir, cxr_file)
    mask_path = os.path.join(mask_dir, mask_file)
    textured_cxr_path = os.path.join(textured_cxr_dir, textured_cxr_file)
    lung_noised_path = os.path.join(lung_noised_dir, lung_noised_file)
    cxr = cv.imread(cxr_path, cv.IMREAD_GRAYSCALE)
    mask = cv.imread(mask_path, cv.IMREAD_GRAYSCALE)
    textured_cxr = cv.imread(textured_cxr_path, cv.IMREAD_GRAYSCALE)
    lung_noised = cv.imread(lung_noised_path, cv.IMREAD_GRAYSCALE)
    max_val = np.percentile(cxr, 97)  # To optimize later
    thresh1 = 50
    thresh2 = 100
    thresh3 = 150
    scaled_thresh1 = thresh1 * max_val / 255
    scaled_thresh2 = thresh2 * max_val / 255
    scaled_thresh3 = thresh3 * max_val / 255
    dense_0 = np.where(textured_cxr < scaled_thresh1, textured_cxr, 0)
    dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
    dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
    dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
    return cxr, textured_cxr, lung_noised, lung_noised, dense_0, dense_1, dense_2, dense_3, count



def update_layout(image_id, index):
    global cxr_dir,mask_dir, textured_cxr_dir,lung_noised_dir,  cxr, mask, textured_cxr, lung_noised, max_val, thresh1, thresh2, thresh3, scaled_thresh1, scaled_thresh2, scaled_thresh3, dense_0, dense_1, dense_2, dense_3, label_title, min_dense_0, max_dense_3, button1, button2, label1
    executable_path = os.path.dirname(os.path.realpath(__file__))
    cxr_dir = os.path.join(executable_path, "Images/cxr")
    mask_dir = os.path.join(executable_path, "Images/mask")
    textured_cxr_dir = os.path.join(executable_path, "Images/textured_cxr")
    lung_noised_dir = os.path.join(executable_path, "Images/lung_noised")

    empty_image = np.zeros((256, 256), dtype=np.uint8)
    cxr = empty_image
    mask = empty_image
    textured_cxr = empty_image
    lung_noised = empty_image
    max_val = np.percentile(cxr, 97)  # To optimize later
    thresh1 = 50
    thresh2 = 100
    thresh3 = 150
    scaled_thresh1 = thresh1 * max_val / 255
    scaled_thresh2 = thresh2 * max_val / 255
    scaled_thresh3 = thresh3 * max_val / 255
    dense_0 = np.where(textured_cxr < scaled_thresh1, textured_cxr, 0)
    dense_1 = np.where(((textured_cxr < scaled_thresh2) & (textured_cxr >= scaled_thresh1)), textured_cxr, 0)
    dense_2 = np.where(((textured_cxr < scaled_thresh3) & (textured_cxr >= scaled_thresh2)), textured_cxr, 0)
    dense_3 = np.where(((textured_cxr < 255) & (textured_cxr >= scaled_thresh3)), textured_cxr, 0)
    # open csv in append mode
    # add title
    with gr.Row():
        # label_title = gr.Label("Density Mapper", visible=True)
        gr.Markdown(
            """
        <center>
        <h1>Density Mapper</h1>
        </center>
        <h3>Instructions:</h3>
        1. Set the brightness of your display to maximum<br>
        2. Initiate the process by clicking the 'Start' button <br>
        3. Synthetic density(middle image in top row) is added to "Original CXR" to obtain "Synthetic CXR"<br>
        4. Adjust the brightness thresholds using the sliders provided \
            to obtain the correct density maps for each level of RALE density<br>
        5. If a density level has absence of pixels at the upper limit, please set the Max Value to 255<br>
        6. Click "Save & continue" to proceed to the next image. The progress is shown in the progress bar<br>
        7. You may close the window and resume the process later when you reopen the window
        """
        )

    with gr.Row():
        with gr.Column():
            im1 = gr.Image(cxr, label="Original CXR")
        with gr.Column():
            im2 = gr.Image(textured_cxr, label="Combined Synthetic Density 0-3")
        with gr.Column():
            im3 = gr.Image(lung_noised, label="Synthetic CXR")
        with gr.Column():
            label1 = gr.Label("Completed! Please close window", visible=False)

    with gr.Row():
        with gr.Column():
            dense0 = gr.Image(dense_0, label="Pixels @ Density 0")
            with gr.Row():
                min_dense_0 = gr.Textbox(value='0', label="Min")
                max_dense_0 = gr.Slider(0, 255, value=50, step=1, label="Max")

        with gr.Column():
            dense1 = gr.Image(dense_1, label="Pixels @ Density 1")
            with gr.Row():
                min_dense_1 = gr.Textbox(value='50', label="Min")
                max_dense_1 = gr.Slider(0, 255, value=100, step=1, label="Max")
                max_dense_0.change(set_min_dense_1, inputs=max_dense_0, outputs=[min_dense_1, dense0, dense1]).then(
                    set_min_dense_1, inputs=max_dense_0, outputs=[min_dense_1, dense0, dense1])

        with gr.Column():
            progress_log = gr.Slider(1, 30, value=0, step=1, label="progress")
            button1 = gr.Button(value="Save & continue", visible=fresh_start)
            button2 = gr.Button(value="Start", visible=not fresh_start)

    with gr.Row():
        with gr.Column():
            dense2 = gr.Image(dense_2, label="Pixels @ Density 2")
            with gr.Row():
                min_dense_2 = gr.Textbox(value='100', label="Min")
                max_dense_2 = gr.Slider(0, 255, value=150, step=1, label="Max")
                max_dense_1.change(set_min_dense_2, inputs=max_dense_1, outputs=[min_dense_2, dense1, dense2]).then(
                    set_min_dense_2, inputs=max_dense_1, outputs=[min_dense_2, dense1, dense2])

        with gr.Column():
            dense3 = gr.Image(dense_3, label="Pixels @ Density 3")
            with gr.Row():
                min_dense_3 = gr.Textbox(value='150', label="Min")
                max_dense_3 = gr.Textbox(value='255', label="Max")
                max_dense_2.change(set_min_dense_3, inputs=max_dense_2, outputs=[min_dense_3, dense2, dense3]).then(
                    set_min_dense_3, inputs=max_dense_2, outputs=[min_dense_3, dense2, dense3])

        with gr.Column():  # adding additional for better visualization
            im3_1 = gr.Image(lung_noised, label="Synthetic CXR")
    button1.click(new__cxr, inputs=[max_dense_0, max_dense_1, max_dense_2],
                  outputs=[im1, im2, im3, im3_1, dense0, dense1, dense2, dense3, progress_log])
    button1.click(change_vis, outputs=[label1, button1, button2])
    button2.click(set_layout,
                  outputs=[im1, im2, im3, im3_1, dense0, dense1, dense2, dense3, progress_log])
    button2.click(change_vis, outputs=[label1, button1, button2])

with gr.Blocks() as demo:
    update_layout(image_id, index)

if __name__ == "__main__":
    demo.launch(auth=check_auth)