import torch
import cv2
import easyocr
import numpy as np
from PIL import Image
from torchvision import transforms

from .model import (
    load_classifier,
    load_hbar_model,
    load_vbar_model,
    load_pie_model
)

class GraphAnalyzer:
    def __init__(self):
        # Automatically use Mac MPS, Nvidia CUDA, or CPU
        self.device = torch.device("mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu"))

        # 1. Load Models 
        self.classifier = load_classifier(self.device)
        self.hbar_model = load_hbar_model(self.device)
        self.vbar_model = load_vbar_model(self.device)
        self.pie_model = load_pie_model(self.device)

        # 2. Setup OCR Engine
        self.reader = easyocr.Reader(['en'])

        # 3. CRITICAL: Setup strict ResNet Normalization
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

        # Adjust these indices to match exactly how the notebook classifier sorted the folders
        self.class_map = {
            0: "dot_line",
            1: "hbar_categorical",
            2: "line",
            3: "pie",
            4: "vbar_categorical"
        }

    def analyze(self, image):
        """Main routing function. Classifies the image and extracts data."""
        
        # --- ROBUST IMAGE LOADING ---
        if isinstance(image, str):
            img_pil = Image.open(image).convert('RGB')
            cv_img = cv2.imread(image)
        elif isinstance(image, np.ndarray):
            cv_img = image
            img_pil = Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
        elif isinstance(image, Image.Image):
            img_pil = image.convert('RGB')
            cv_img = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
        else:
            raise TypeError("Input must be a file path (str), NumPy array (cv2), or PIL Image.")

        # 1. Classify the graph type
        image_tensor = self.data_transforms(img_pil).unsqueeze(0).to(self.device)

        with torch.no_grad():
            outputs = self.classifier(image_tensor)
            predicted_class = torch.argmax(outputs, dim=1).item()

        graph_type = self.class_map.get(predicted_class, "unknown")

        # 2. Route to the correct OCR + Prediction logic
        if graph_type == "pie":
            return self._extract_pie_data(cv_img, image_tensor)
        elif graph_type == "vbar_categorical":
            return self._extract_vbar_data(cv_img, image_tensor)
        elif graph_type == "hbar_categorical":
            return self._extract_hbar_data(cv_img, image_tensor)
        else:
            return {"type": graph_type, "status": f"Extraction logic for {graph_type} is not completely implemented in this snippet"}

    def _extract_pie_data(self, cv_img, image_tensor):
        """The dedicated Pie Chart extraction logic"""
        with torch.no_grad():
            preds = self.pie_model(image_tensor).squeeze().cpu().numpy() * 100.0
            
        h, w, _ = cv_img.shape
        gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)
        all_text_results = self.reader.readtext(gray, mag_ratio=2.5)
        
        title = "Untitled"
        raw_legend_names = [] 
        
        for bbox, text, conf in all_text_results:
            clean_text = text.strip()
            if not clean_text: continue
                
            x_center = (bbox[0][0] + bbox[2][0]) / 2
            y_center = (bbox[0][1] + bbox[2][1]) / 2
            x_pct = x_center / w
            y_pct = y_center / h
            
            if y_pct < 0.15:
                title = clean_text
            elif y_pct > 0.15:
                if len(clean_text) > 2 and not clean_text.replace('.', '', 1).isdigit():
                    if clean_text.lower() == "grav": clean_text = "Gray"
                    raw_legend_names.append((y_pct, clean_text)) 

        legend_names = [item[1] for item in sorted(raw_legend_names, key=lambda i: i[0])]
        
        num_slices = len(legend_names)
        if num_slices == 0:
            valid_preds = [v for v in preds if v > 1.5]
            num_slices = len(valid_preds)
            
        num_slices = min(num_slices, 10)
        slice_preds = preds[:num_slices]
        total_pred = sum(slice_preds)
        
        if total_pred > 0:
            normalized_preds = [(v / total_pred) * 100.0 for v in slice_preds]
        else:
            normalized_preds = slice_preds
        
        final_slices = {}
        for i in range(num_slices):
            slice_name = legend_names[i] if i < len(legend_names) else f"Unknown_Slice_{i+1}"
            slice_value = round(normalized_preds[i], 2) if i < len(normalized_preds) else 0.0
            final_slices[slice_name] = slice_value

        return {
            "type": "pie",
            "title": title,
            "data": final_slices
        }

    def _extract_vbar_data(self, cv_img, image_tensor):
        """The dedicated Vertical Bar Chart extraction logic"""
        with torch.no_grad():
            preds = self.vbar_model(image_tensor).squeeze().cpu().numpy() * 100.0
            
        h, w, _ = cv_img.shape
        
        # 1. Title
        title_res = self.reader.readtext(cv_img[0:int(h*0.15), :])
        title = title_res[0][1] if title_res else "Untitled"
        
        # 2. X-Axis Labels (Rotated crop to catch vertical text)
        label_area = cv_img[int(h*0.70):h, :] 
        label_rot = cv2.rotate(label_area, cv2.ROTATE_90_COUNTERCLOCKWISE)
        label_res = self.reader.readtext(label_rot)
        
        x_labels = [res[1] for res in label_res if len(res[1]) > 1]
        x_labels = x_labels[::-1]

        # 3. Intelligent Matching (The dynamic fix)
        max_val = np.max(preds)
        noise_threshold = max_val * 0.05 
        
        valid_preds = [v for v in preds if v > noise_threshold]
        
        count = max(len(x_labels), len(valid_preds))
        count = min(count, 10) 
        
        final_values = [round(float(v), 2) for v in preds[:count]]
        
        while len(x_labels) < len(final_values):
            x_labels.append(f"Unknown_Label_{len(x_labels)+1}")

        return {
            "type": "vbar_categorical",
            "title": title,
            "x_axis_labels": x_labels,
            "values": final_values
        }

    def _extract_hbar_data(self, cv_img, image_tensor):
        """The dedicated Horizontal Bar Chart extraction logic"""
        with torch.no_grad():
            preds = self.hbar_model(image_tensor).squeeze().cpu().numpy() * 100.0
            
        h, w, _ = cv_img.shape
        
        # 1. Title
        title_res = self.reader.readtext(cv_img[0:int(h*0.15), :])
        title = title_res[0][1] if title_res else "Untitled"
        
        # 2. Labels (Smooth Contrast applied)
        label_area = cv_img[int(h*0.10):int(h*0.95), 0:int(w*0.45)] 
        gray = cv2.cvtColor(label_area, cv2.COLOR_BGR2GRAY)
        contrast_img = cv2.convertScaleAbs(gray, alpha=1.5, beta=-50)
        
        label_res = self.reader.readtext(contrast_img, mag_ratio=2.0)
        y_labels = [res[1] for res in label_res if len(res[1]) > 1]
        y_labels = y_labels[::-1] 

        # 3. Intelligent Matching
        valid_preds = [v for v in preds if v > 2.0]
        
        count = len(y_labels) if len(y_labels) > 0 else len(valid_preds)
        count = min(count, 10) 
        
        final_values = [round(float(v), 2) for v in preds[:count]]
        
        while len(y_labels) < len(final_values):
            y_labels.append(f"Unknown_Label_{len(y_labels)+1}")

        return {
            "type": "hbar_categorical",
            "title": title,
            "y_axis_labels": y_labels,
            "values": final_values
        }