import torch
from models import clip_model, clip_processor, segformer_model, segformer_processor
from constants import fashion_categories, fashion_clip_to_segformer, class_names, category_to_segment_mapping, garment_to_segments
from segmentation import identify_garment_segformer

def identify_garment_clip(image):
    """Identify the garment type using Fashion-CLIP model"""
    # Prepare text prompts
    texts = [f"a photo of a {category}" for category in fashion_categories]
    
    # Process inputs
    inputs = clip_processor(text=texts, images=image, return_tensors="pt", padding=True)
    
    # Get predictions
    with torch.no_grad():
        outputs = clip_model(**inputs)
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=1)
    
    # Get the top prediction
    top_idx = torch.argmax(probs[0]).item()
    top_category = fashion_categories[top_idx]
    confidence = probs[0][top_idx].item() * 100
    
    # Map to SegFormer class if possible
    if top_category in fashion_clip_to_segformer:
        segformer_idx = fashion_clip_to_segformer[top_category]
        segformer_class = class_names[segformer_idx]
        return top_category, segformer_idx, confidence
    else:
        # Fallback to using SegFormer directly
        return top_category, None, confidence

def get_segments_for_garment(garment_image):
    """Get the segments that should be included for a given garment image"""
    # First try to identify the garment using Fashion-CLIP
    clip_category, segformer_idx, confidence = identify_garment_clip(garment_image)
    
    # If CLIP couldn't map to a SegFormer class, fall back to SegFormer
    if segformer_idx is None:
        garment_name, segformer_idx = identify_garment_segformer(garment_image)
        method = "SegFormer (fallback)"
        confidence_text = ""
    else:
        garment_name = class_names[segformer_idx]
        method = "Fashion-CLIP"
        confidence_text = f" with {confidence:.2f}% confidence"
    
    if segformer_idx is None:
        return None, None, "No clear garment detected in the garment image"
    
    # Get all segments that should be included based on the detected garment
    if method == "Fashion-CLIP" and clip_category in category_to_segment_mapping:
        # Use the detailed mapping for CLIP categories
        selected_class = category_to_segment_mapping[clip_category]
    elif segformer_idx in garment_to_segments:
        # Fall back to the SegFormer class-based mapping
        segment_indices = garment_to_segments[segformer_idx]
        selected_class = [class_names[idx] for idx in segment_indices]
    else:
        # Fallback to just the detected garment if no mapping exists
        selected_class = [class_names[segformer_idx]]
    
    # Prepare a more descriptive result text
    included_segments = ", ".join(selected_class)
    if method == "Fashion-CLIP":
        result_text = f"Detected garment: {clip_category} (mapped to {garment_name})\nUsing {method}{confidence_text}\nSegmented parts: {included_segments}"
    else:
        result_text = f"Detected garment: {garment_name}\nUsing {method}\nSegmented parts: {included_segments}"
    
    return selected_class, segformer_idx, result_text