import streamlit as st
import cv2
import numpy as np
from pathlib import Path
import tempfile
import os
import plotly.express as px
import pytesseract
import re
from ultralytics import YOLO


# Colors for labels
LABEL_COLORS = {
    "Circuit": (37, 255, 76),        # Green
    "Fuse": (255, 76, 76),           # Red
    "Missing Fuse": (255, 204, 0),   # Yellow
    "Missing Circuit": (0, 128, 255) # Blue
}

def test_trained_model(model_path, test_image_path, conf_threshold=0.5):
    model = YOLO(model_path)
    test_image_path = Path(test_image_path)

    if not test_image_path.exists():
        st.error(f"No image found at {test_image_path}")
        return None

    results = model(str(test_image_path), conf=conf_threshold, verbose=False)
    result = results[0]
    image = cv2.imread(str(test_image_path))

    if result.boxes is not None:
        boxes = result.boxes.xyxy.cpu().numpy()
        classes = result.boxes.cls.cpu().numpy().astype(int)

        for box, cls_id in zip(boxes, classes):
            x1, y1, x2, y2 = map(int, box)
            label = model.names[cls_id]
            color = LABEL_COLORS.get(label, (255, 255, 255))  # fallback white
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    st.image(image_rgb, caption=f"{test_image_path.name} - Detection with Labels", use_column_width=True)

    # Legend
    st.markdown("### 🏷 Legend (Bounding Box Colors)")
    for label, color in LABEL_COLORS.items():
        r, g, b = color
        hex_color = f'#{r:02x}{g:02x}{b:02x}'
        st.markdown(f"<span style='color:{hex_color}; font-weight:bold;'>&#9632;</span> {label}", unsafe_allow_html=True)

    detections = [model.names[int(cls)] for cls in result.boxes.cls]
    return detections

def count_circuit_components(detections):
    target_labels = ['Circuit', 'Fuse', 'Missing Fuse', 'Missing Circuit']
    component_counts = {}

    for label in target_labels:
        count = sum(1 for detection in detections if detection.lower().replace('_', ' ') == label.lower())
        component_counts[label] = count

    return component_counts

def display_component_summary(component_counts, image_name):
    st.write(f"\n{'=' * 50}")
    st.write(f"COMPONENT SUMMARY FOR: {image_name}")
    st.write(f"{'=' * 50}")

    for component, count in component_counts.items():
        st.write(f"{component:15}: {count:3d}")

    st.write(f"{'-' * 25}")

def extract_temperatures_from_image(image_path):
    st.subheader("📋 Temperature Report")

    image = cv2.imread(image_path)
    if image is None:
        st.error("Error loading image for OCR.")
        return None

    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    lower_white = np.array([0, 0, 200])
    upper_white = np.array([180, 50, 255])
    white_mask = cv2.inRange(hsv_image, lower_white, upper_white)
    white_text_image = cv2.bitwise_and(image, image, mask=white_mask)
    white_text_gray = cv2.cvtColor(white_text_image, cv2.COLOR_BGR2GRAY)

    ocr_text = pytesseract.image_to_string(white_text_gray)
    ocr_text = re.sub(r'(\d+)\s(\d)°C', r'\1.\2°C', ocr_text)

    temps_found = re.findall(r'(\d{1,3}\.\d)\s*°C', ocr_text)
    temps_found = sorted([float(t) for t in temps_found])

    if len(temps_found) == 0:
        st.warning("No temperature values detected.")
        return None

    highest_temp = max(temps_found)
    lowest_temp = min(temps_found)

    if len(temps_found) >= 3:
        ambient_temp = temps_found[1]
    elif len(temps_found) == 2:
        ambient_temp = sum(temps_found) / 2
    else:
        ambient_temp = temps_found[0]

    st.write("Detected Temperatures:")
    st.write(f"- Highest Temp: {highest_temp:.2f} °C")
    st.write(f"- Lowest Temp: {lowest_temp:.2f} °C")
    st.write(f"- Ambient Temp: {ambient_temp:.2f} °C")

    return {
        "highest": highest_temp,
        "lowest": lowest_temp,
        "ambient": ambient_temp,
    }

def plot_hover_temperature(image_path, temp_min=0.0, temp_max=100.0):
    img = cv2.imread(image_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    red_channel = img_rgb[:, :, 0]

    def map_temp(val):
        return round((val / 255.0) * (temp_max - temp_min) + temp_min, 2)

    temp_map = np.vectorize(map_temp)(red_channel)

    fig = px.imshow(
        temp_map,
        color_continuous_scale="Hot",
        labels={'x': 'X', 'y': 'Y', 'color': 'Temperature (°C)'},
        title=f"Temperature Map (Hover) - Range: {temp_min}°C to {temp_max}°C"
    )

    fig.update_layout(dragmode=False, margin=dict(l=20, r=20, t=40, b=20), height=600)
    fig.update_traces(hovertemplate="X: %{x}<br>Y: %{y}<br>Temp: %{z:.2f} °C")
    st.plotly_chart(fig, use_container_width=True)

def analyze_thermal_circuit_image(model_path, test_image_path, conf_threshold=0.5):
    st.write("Starting thermal circuit analysis...")
    detections = test_trained_model(model_path, test_image_path, conf_threshold)

    if detections is None:
        return

    component_counts = count_circuit_components(detections)
    image_name = Path(test_image_path).name
    display_component_summary(component_counts, image_name)

    return component_counts

def extract_ambient_temperature(image_path):
    image = cv2.imread(image_path)
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    white_mask = cv2.inRange(hsv_image, np.array([0, 0, 200]), np.array([180, 50, 255]))
    text_image = cv2.bitwise_and(image, image, mask=white_mask)
    gray = cv2.cvtColor(text_image, cv2.COLOR_BGR2GRAY)
    text = pytesseract.image_to_string(gray)
    temps = re.findall(r'(\d{1,3}\.\d)\s*°C', text)

    if not temps:
        return None

    temps = sorted([float(t) for t in temps])
    return temps[1] if len(temps) >= 3 else sum(temps) / len(temps)

def detect_circuit_boxes(model_path, image_path, conf_threshold=0.5):
    model = YOLO(model_path)
    results = model(image_path, conf=conf_threshold)[0]
    boxes = []

    if results.boxes is not None:
        for i, cls in enumerate(results.boxes.cls):
            label = model.names[int(cls)]
            if label.lower() == "circuit":
                box = results.boxes.xyxy[i].cpu().numpy().astype(int)
                boxes.append(box)

    return boxes

def draw_hotspot_boxes(image_path, boxes, ambient_temp, min_temp=0, max_temp=100):
    image = cv2.imread(image_path)
    red_channel = image[:, :, 2]
    temp_map = (red_channel / 255.0) * (max_temp - min_temp) + min_temp
    threshold = ambient_temp + 5

    mask = np.zeros_like(red_channel, dtype=np.uint8)
    annotated = image.copy()

    for box in boxes:
        x1, y1, x2, y2 = box
        region_temp = temp_map[y1:y2, x1:x2]
        if np.mean(region_temp) > threshold:
            cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 0, 255), 2)
            cv2.putText(annotated, f"HOT > {threshold:.1f}°C", (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
            mask[y1:y2, x1:x2] = 255

    return annotated, mask, temp_map


def main():
    st.set_page_config(layout="wide")
    st.title("Thermal Circuit & Hotspot Analysis")

    uploaded_file = st.file_uploader("Upload Thermal Image (.jpg/.png)", type=["jpg", "jpeg", "png"], key="upload")
    if uploaded_file:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
            tmp_file.write(uploaded_file.getvalue())
            tmp_path = tmp_file.name

        st.image(cv2.cvtColor(cv2.imread(tmp_path), cv2.COLOR_BGR2RGB), caption="Uploaded Thermal Image", use_column_width=True)

        conf_threshold = st.slider("Detection Confidence Threshold", 0.0, 1.0, 0.5, 0.01, key="conf_thresh")
        if st.button("Analyze Image", key="analyze"):
            model_path = "best.pt"
            analyze_thermal_circuit_image(model_path, tmp_path, conf_threshold)
            temps = extract_temperatures_from_image(tmp_path)
            if temps:
                plot_hover_temperature(tmp_path, temp_min=temps["lowest"], temp_max=temps["highest"])
            else:
                st.info("Using default temperature range (0°C to 100°C) for hover map.")
                plot_hover_temperature(tmp_path)


            ambient = extract_ambient_temperature(tmp_path)

            boxes = detect_circuit_boxes(model_path, tmp_path)

            if not boxes:
                st.warning("No circuit components detected.")
            else:
                st.success(f"Detected {len(boxes)} Circuit Regions")

                st.subheader("Mapping Hotspots...")
                annotated, mask, temp_map = draw_hotspot_boxes(tmp_path, boxes, ambient)

                st.image(cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB), caption="Annotated Hot Circuits", use_column_width=True)



                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as mask_file:
                    cv2.imwrite(mask_file.name, mask)
                    st.download_button("Download Hotspot Mask", data=open(mask_file.name, "rb"), file_name="hotspot_mask.png")

        os.unlink(tmp_path)

if __name__ == "__main__":
    main()