import gradio as gr
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt

def extract_frames(gif_path):
    """Extract frames from a GIF and return as a list of numpy arrays."""
    try:
        img = Image.open(gif_path)
        frames = []
        while True:
            frame = img.convert('L')  # Convert to grayscale
            frames.append(np.array(frame))
            try:
                img.seek(img.tell() + 1)
            except EOFError:
                break
        return frames, None
    except Exception as e:
        return None, f"Error loading GIF: {str(e)}"

def preprocess_frame(frame):
    """Preprocess a frame: apply Gaussian blur and adaptive thresholding."""
    # Apply Gaussian blur to reduce noise
    blurred = cv2.GaussianBlur(frame, (9, 9), 0)
    # Adaptive thresholding to enhance lighter regions (CME features)
    thresh = cv2.adaptiveThreshold(
        blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
    )
    return thresh

def detect_circles(frame_diff, image_center, min_radius=20, max_radius=200):
    """Detect circles in a frame difference image, centered at the Sun."""
    circles = cv2.HoughCircles(
        frame_diff,
        cv2.HOUGH_GRADIENT,
        dp=1.5,  # Increase resolution for better detection
        minDist=100,  # Prevent overlapping circles
        param1=100,  # Higher edge threshold to reduce noise
        param2=20,  # Lower accumulator threshold to detect faint circles
        minRadius=min_radius,
        maxRadius=max_radius
    )

    if circles is not None:
        circles = np.round(circles[0, :]).astype("int")
        # Filter circles: only keep those centered near the image center
        filtered_circles = []
        center_tolerance = 30  # Allow 30-pixel deviation from the center
        for (x, y, r) in circles:
            if (abs(x - image_center[0]) < center_tolerance and 
                abs(y - image_center[1]) < center_tolerance):
                filtered_circles.append((x, y, r))
        return filtered_circles if filtered_circles else None
    return None

def analyze_gif(gif_file):
    """Analyze a GIF for growing concentric circles centered at the Sun."""
    try:
        # Handle Gradio file input
        gif_path = gif_file.name if hasattr(gif_file, 'name') else gif_file

        # Extract frames
        frames, error = extract_frames(gif_path)
        if error:
            return error, []

        if len(frames) < 2:
            return "GIF must have at least 2 frames for analysis.", []

        # Determine the image center (Sun's position)
        height, width = frames[0].shape
        image_center = (width // 2, height // 2)  # Assume Sun is at the center

        # Initialize results
        results = []
        circle_data = []
        min_radius = 20
        max_radius = min(height, width) // 2  # Limit max radius to half the image size

        # Process frames
        for i in range(len(frames) - 1):
            frame1 = preprocess_frame(frames[i])
            frame2 = preprocess_frame(frames[i + 1])

            # Compute absolute difference between consecutive frames
            frame_diff = cv2.absdiff(frame2, frame1)
            # Enhance contrast for lighter pixels
            frame_diff = cv2.convertScaleAbs(frame_diff, alpha=3.0, beta=0)

            # Detect circles centered at the Sun
            circles = detect_circles(frame_diff, image_center, min_radius, max_radius)

            if circles:
                # Take the largest circle (most prominent CME feature)
                largest_circle = max(circles, key=lambda c: c[2])  # Sort by radius
                x, y, r = largest_circle
                circle_data.append({
                    "frame": i + 1,
                    "center": (x, y),
                    "radius": r
                })

            # Visualize the frame with detected circle
            output_frame = cv2.cvtColor(frames[i + 1], cv2.COLOR_GRAY2RGB)
            if circles:
                x, y, r = largest_circle
                cv2.circle(output_frame, (x, y), r, (0, 255, 0), 2)
            # Convert to PIL Image for Gradio
            output_frame = Image.fromarray(output_frame)
            results.append(output_frame)

        # Analyze circle data for growth
        report = "Analysis Report:\n"
        if circle_data:
            radii = [c["radius"] for c in circle_data]
            centers = [c["center"] for c in circle_data]
            frames_with_circles = [c["frame"] for c in circle_data]

            # Check if radii are increasing over frames
            is_growing = all(radii[i] < radii[i + 1] for i in range(len(radii) - 1)) if len(radii) > 1 else False
            center_consistent = all(
                abs(centers[i][0] - centers[0][0]) < 20 and
                abs(centers[i][1] - centers[0][1]) < 20
                for i in range(1, len(centers))
            ) if len(centers) > 1 else False

            report += f"Detected {len(circle_data)} circles across frames.\n"
            for c in circle_data:
                report += f"Frame {c['frame']}: Center at {c['center']}, Radius {c['radius']} pixels\n"
            if is_growing and center_consistent:
                report += "\nConclusion: Growing concentric circles detected, indicative of a potential Earth-directed CME."
            else:
                report += "\nConclusion: Detected circles, but growth pattern or center consistency does not confirm a clear CME."
        else:
            report += "No concentric circles detected."

        return report, results
    except Exception as e:
        return f"Error during analysis: {str(e)}", []

# Gradio interface
iface = gr.Interface(
    fn=analyze_gif,
    inputs=gr.File(label="Upload Solar GIF", file_types=[".gif"]),
    outputs=[
        gr.Textbox(label="Analysis Report"),
        gr.Gallery(label="Frames with Detected Circles")
    ],
    title="Solar CME Detection",
    description="Upload a GIF of solar images to detect growing concentric circles indicative of Earth-directed coronal mass ejections (CMEs)."
)

if __name__ == "__main__":
    iface.launch()