#import library
import gradio as gr
from roboflow import Roboflow
import numpy as np
from PIL import Image
import requests
from io import BytesIO
import pandas as pd
import os

# Initialize Roboflow with your API key
rf = Roboflow(api_key="kKDoCn3ABT9AKeFQDCB4")

# Function to calculate the area of a polygon using the shoelace formula
def calculate_polygon_area(points):
    n = len(points)
    area = 0.0
    for i in range(n):
        x1, y1 = points[i]
        x2, y2 = points[(i + 1) % n]
        area += (x1 * y2 - x2 * y1)
    return abs(area) / 2.0

# Function to process Roboflow prediction JSON and calculate corrosion areas
def calculate_corrosion_areas(json_data, unit="pixels", conversion_factor=1):
    corrosion_areas = []
    for prediction in json_data["predictions"]:
        if prediction["class"] == "Corrosion":
            points = [(point["x"], point["y"]) for point in prediction["points"]]
            area = calculate_polygon_area(points)
            if unit == "cm??":
                area *= conversion_factor  # Convert area from pixels to cm??
            corrosion_areas.append(area)

    total_corrosion_area = sum(corrosion_areas)

# Prepare output
    result = {
        "individual_areas": [f"{area} {unit}" for area in corrosion_areas],
        "total_area": f"{total_corrosion_area} {unit}",
        "recommendation": get_inspection_recommendation(total_corrosion_area)
    }
    return result

# Function to provide inspection recommendation based on total corrosion area
def get_inspection_recommendation(total_area):
    if total_area < 1000:
        return "No immediate inspection needed."
    elif total_area < 5000:
        return "Schedule an inspection in the next 6 months."
    else:
        return "Immediate inspection required."

# Define a Gradio interface to input a URL, run inference, and calculate corrosion areas
def url_infer_and_calculate(url, location, unit="pixels", conversion_factor=1, corrosion_type="", inspection_standards=[], ndt_methods=[], manual_recommendation="", supporting_data=""):
    try:
        # Run inference using the Roboflow script
        rf_project = rf.workspace().project("corrosion-instance-segmentation-sfcpc")
        model = rf_project.version(3).model
        prediction = model.predict(url)

        # Ensure the response is properly formatted as JSON
        prediction_json = prediction.json()

        # Calculate corrosion areas from the Roboflow prediction
        corrosion_areas = calculate_corrosion_areas(prediction_json, unit, float(conversion_factor))

        # Download the image from the URL and convert it to a PIL Image
        response = requests.get(url)
        img = Image.open(BytesIO(response.content))

        # Create a pandas DataFrame for reporting
        df = pd.DataFrame([{'Number': index+1, 'URL': url, 'Location': location, 'corrosion_areas': corrosion_areas, 'Recommendation': corrosion_areas['recommendation']} for index in range(len(corrosion_areas))])

        # Write DataFrame to local CSV file with index included immediately after creating it.
        df.to_csv('Corrosion_Report.csv', index=False)

        # Write DataFrame to a string in CSV format
        csv_string = df.to_csv(index=False)

        return img, corrosion_areas, prediction_json, csv_string

    except Exception as e:
        return {"error": str(e)}

# Create a Gradio interface for URL input, inference, and corrosion area calculation
iface = gr.Interface(
    fn=url_infer_and_calculate,
    inputs=[
        gr.inputs.Textbox(label="Enter the URL of an image"),
        gr.inputs.Textbox(label="Enter the Location"),
        gr.inputs.Dropdown(choices=["pixels", "cm"], label="Area Unit"),
        gr.inputs.Textbox(label="Conversion Factor"),
        gr.inputs.Textbox(label="Enter the Corrosion Type"),
        gr.inputs.Textbox(label="Inspection Standards"),
        gr.inputs.CheckboxGroup(choices=["UT thickness", "UT scan", "Phased Array UT", "Short range UT", "Long range UT", "MT", "PT", "Acfm", "Pulse eddy current", "magnetic flux leakage", "positive material identification (PMI)"], label="NDT Inspection Methods"),
        gr.inputs.Textbox(label="Enter Manual Recommendation"),
        gr.inputs.Textbox(lines=5, label="Enter Supporting Data URLs (separated by commas)")
    ],
    outputs=[
        gr.outputs.Image(type="pil"), 
        "json",  # JSON output
        gr.outputs.Textbox(label="CSV Data", type="text"),  # CSV data as a plain text
        gr.outputs.Textbox(label="Corrosion Data"),  # Display CSV data as a table
    ],
    title="Tim CCG",
    description="Enter the URL of an image to perform rust detection and calculate corrosion areas.",
)

# Launch the Gradio interface
iface.launch(debug=False, share=False)