import torch
# PyTorch 2.6+ fix: Allow YOLO weights to load safely
try:
    import ultralytics
    # Add common YOLO classes to safe globals
    torch.serialization.add_safe_globals(["ultralytics.nn.tasks.DetectionModel", "ultralytics.nn.tasks.SegmentationModel", "ultralytics.nn.tasks.PoseModel"])
except:
    pass
from transformers import T5Tokenizer, T5ForConditionalGeneration
from models.blip_model import processor_cap, model_cap, processor_vqa, model_vqa
from config import DEVICE

try:
    from ultralytics import YOLO
    yolo_model = YOLO("yolov8n.pt")
except Exception as e:
    print("YOLO initialization failed:", e)
    yolo_model = None

# Initialize FLAN-T5 specifically to synthesize intelligent explanations locally
try:
    tokenizer_t5 = T5Tokenizer.from_pretrained("google/flan-t5-base")
    model_t5 = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base").to(DEVICE)
    model_t5.eval()
except Exception as e:
    tokenizer_t5 = None
    model_t5 = None

def blip_yolo_answer(image, question, lang="en"):
    objects_detected = []
    if yolo_model is not None:
        try:
            results = yolo_model(image, verbose=False)
            for r in results:
                for c in r.boxes.cls:
                    objects_detected.append(yolo_model.names[int(c)])
        except:
            pass
    
    unique_objs = list(set(objects_detected))
    counts = {obj: objects_detected.count(obj) for obj in unique_objs}
    det_str = ", ".join([f"{count} {obj}(s)" for obj, count in counts.items()]) if objects_detected else "no distinct YOLO objects"

    if image.mode != "RGB":
        image = image.convert("RGB")
        
    cap_inputs = processor_cap(images=image, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        cap_output = model_cap.generate(**cap_inputs, max_new_tokens=30)
    caption = processor_cap.decode(cap_output[0], skip_special_tokens=True).strip()

    context_question = f"The image features {caption} and {det_str}. {question}"
    vqa_inputs = processor_vqa(images=image, text=context_question, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        vqa_output = model_vqa.generate(**vqa_inputs, max_new_tokens=20)
    answer = processor_vqa.decode(vqa_output[0], skip_special_tokens=True).strip()

    if model_t5 is not None:
        prompt = f"Explain logically why the answer to the question '{question}' is '{answer}', knowing the scene is {caption} and precisely contains {det_str}. Be concise and logical."
        input_ids = tokenizer_t5(prompt, return_tensors="pt").input_ids.to(DEVICE)
        with torch.no_grad():
            out = model_t5.generate(input_ids, max_new_tokens=50)
        explanation = tokenizer_t5.decode(out[0], skip_special_tokens=True)
        # Capitalize first letter strictly
        explanation = explanation[0].upper() + explanation[1:] if len(explanation) > 0 else explanation
    else:
        explanation = f"The visual structure shows {caption} alongside {det_str}, directly leading to the conclusion that the answer is {answer}."

    if objects_detected:
        caption = f"YOLO found: {det_str}. Scene: {caption}"
    
    return caption, answer, explanation