#load relevant modules
from stat import FILE_ATTRIBUTE_INTEGRITY_STREAM
import eyepy as ep
import matplotlib.pyplot as plt
import os
import pandas as pd
from PIL import Image
import numpy as np
import torch
import shutil
from skimage import measure, segmentation, morphology, exposure
from octolyzer.segment.sloseg import slo_inference, avo_inference, fov_inference
from octolyzer.utils import generate_imgmask, generate_zonal_masks
from octolyzer.measure.slo import feature_measurement, get_vessel_coords
import octolyzer.utils as utils
import gradio as gr
import tempfile

# Load models
# SLO segmentation models
slo_model = slo_inference.SLOSegmenter()
# FOV segmentation models
fov_model = fov_inference.FOVSegmenter()
# AVO segmentation models
avo_model = avo_inference.AVOSegmenter() 

def load_file(e2e_file):
    '''
    read e2e file and return slo image as float
    '''
    if e2e_file is None:
        return "**Error:** Please upload an image first."

    filetype = e2e_file.rsplit('.', 1)[1].lower()
    if filetype == "e2e":
        ep_obj = ep.import_heyex_e2e(e2e_file) 
        return ep_obj.localizer.data.astype(float)
    else:
        return "**Error:** Unsupported filetype"

def create_SLO_segmentation(slo):
    slo_vbinmap = slo_model.predict_img(slo)
    return slo_vbinmap

def create_avo_segmentation(slo):
    slo_avimout, od_centre = avo_model.predict_img(slo, location="macula")
    return slo_avimout, od_centre

def plot_SLO_segmentation(slo, slo_vbinmap):
    # Assuming you have slo and slo_vbinmap from your experiment
    fig, ax = plt.subplots(1, 1, figsize=(8, 8))

    # Display the SLO image as background
    ax.imshow(slo, cmap='gray')

    # Create a colored overlay for the binary vessel map
    # Generate RGBA mask with red channel (cmap=0) for vessels
    vessel_overlay = generate_imgmask(slo_vbinmap)  # Red vessels
    ax.imshow(vessel_overlay, alpha=0.6)

    ax.set_title('SLO with Binary Vessel Overlay')
    ax.axis('off')
    
    return fig

def plot_avo_segmentation(slo, slo_avimout):
    # Assuming you have slo and slo_vbinmap from your experiment
    fig, ax = plt.subplots(1, 1, figsize=(8, 8))

    # Display the SLO image as background
    ax.imshow(slo, cmap='gray')

    # Create a colored overlay for the binary vessel map
    # Generate RGBA mask with red channel (cmap=0) for vessels
    avoimout_save = 191*slo_avimout[...,0] + 127*slo_avimout[...,2] + 255*slo_avimout[...,1]
    #vessel_overlay = generate_imgmask(slo_avimout[...,2])  # Red vessels
    ax.imshow(avoimout_save, alpha=0.6)

    ax.set_title('SLO with Artery / Veins / Optic Nerve Overlay')
    ax.axis('off')

    return fig

def features(slo, od_centre, slo_vbinmap, slo_avimout):
    img_shape = slo.shape
    _, N = img_shape
    od_radius = None #macula centred map
    location = "Macula"
    scale = 1
    slo_dict = {}
    slo_keys = ["binary", "artery", "vein"]
    masks = generate_zonal_masks((N,N), od_radius, od_centre, location)
    artery_vbinmap, vein_vbinmap = slo_avimout[...,0], slo_avimout[...,2]
    od_mask = slo_avimout[...,1]
    for v_map, v_type in zip([slo_vbinmap, artery_vbinmap, vein_vbinmap], slo_keys):
        vcoords = get_vessel_coords.generate_vessel_skeleton(v_map, od_mask, od_centre, min_length=10)
        slo_dict[v_type] = feature_measurement.vessel_metrics(v_map, vcoords, masks, scale=scale, vessel_type=v_type)
    slo_df = utils.nested_dict_to_df(slo_dict).reset_index()
    slo_df = slo_df.rename({"level_0":"vessel_map", "level_1":"zone"}, axis=1, inplace=False)
    reorder_cols = ["vessel_map", "zone", "fractal_dimension", "vessel_density", "average_global_calibre", 
                "average_local_calibre", "tortuosity_density", "tortuosity_distance", "CRAE_Knudtson", "CRVE_Knudtson"]
    slo_df = slo_df[reorder_cols]
    return slo_df

def predict(e2e_file):
    slo = load_file(e2e_file)
    all_vessels = create_SLO_segmentation(slo)
    AVO, OD_centre = create_avo_segmentation(slo)
    all_vessel_plot = plot_SLO_segmentation(slo, all_vessels)
    AVO_plot = plot_avo_segmentation(slo, AVO)
    slo_df = features(slo, OD_centre, all_vessels, AVO)
    csv_file = tempfile.NamedTemporaryFile(delete=False, suffix='.csv')
    slo_df.to_csv(csv_file.name, index=False)
    csv_file.close()

    return all_vessel_plot, AVO_plot, slo_df, Image.fromarray(slo), csv_file.name

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # Automated Retinal Vascular Morphology Quantification from SLO images
        
        Upload a Heidelberg Spectralis .E2E file to automatically segment and assess vessel metrics
        
        **Accepted formats:** E2E

        **Scan type:** Only macular centred scans accepted.
        
        **Disclaimer:** This is a research tool and not intended for clinical use.
        """
    )
    
    with gr.Row():
        with gr.Column():
            image_input = gr.File(
                label="Upload e2e Image"
            )
            #predict_btn = gr.Button("🔍 Analyze Image", variant="primary")

    with gr.Row():   
        with gr.Column():
            slo_image = gr.Image(label = "SLO image to be analysed") 
        with gr.Column():
            plot_output_1 = gr.Plot(label = "segmentation map of all vessels")
        with gr.Column():
            plot_output_2 = gr.Plot(label = "segmentation map of arteries / veins / optic nerves")
    
    with gr.Row():
        with gr.Column():
            dataframe_output = gr.Dataframe(label="Vessel Metrics", wrap=True)
            csv_download = gr.File(label="Download CSV", file_count="single")   


    image_input.upload(
        fn=predict,
        inputs=image_input,
        outputs=[plot_output_1, plot_output_2, dataframe_output, slo_image, csv_download]
    )
    
    gr.Markdown(
        """
        ---
        ### About this tool
        
        This tool is adapted from Octolyzer, a fully automatic toolkit for segmentation and feature extracting in optical coherence tomography and scanning laser ophthalmoscopy data
        
        **Citation:** https://arxiv.org/abs/2407.14128 and https://github.com/jaburke166/OCTolyzer
        
        **License:** GPL-3.0 license
        """
    )

# Launch
if __name__ == "__main__":
    demo.launch()