import os
import sys
import numpy as np
import torch
from PIL import Image

# Add parent directory to sys.path to enable loading of local package
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from gemmasight.config import DIM_FUSED
from gemmasight.utils.preprocess import load_and_preprocess_image
from gemmasight.models.feature_extractor import DualEncoderFeatureExtractor
from gemmasight.models.classifier import MSIClassifier
from gemmasight.models.retriever import CaseRetriever
from gemmasight.utils.occlusion import generate_occlusion_heatmap
from gemmasight.inference import GemmaSightInferencePipeline

def run_pipeline_tests():
    print("=== GemmaSight Verification Test Suite ===")
    
    # Create random synthetic image for test
    test_img_np = (np.random.rand(224, 224, 3) * 255).astype(np.uint8)
    test_img = Image.fromarray(test_img_np)
    print("✓ Created synthetic H&E image.")
    
    # 1. Preprocessing Test
    print("\nRunning Test 1: Preprocessing...")
    tensor = load_and_preprocess_image(test_img)
    assert tensor.shape == (1, 3, 224, 224), f"Unexpected tensor shape {tensor.shape}"
    print(f"✓ Preprocessed image successfully. Output shape: {tensor.shape}")
    
    # 2. Feature Extraction Test
    print("\nRunning Test 2: Dual-Encoder Feature Extraction...")
    extractor = DualEncoderFeatureExtractor(force_simulation=True)
    embeddings = extractor(test_img)
    assert embeddings.shape == (1, DIM_FUSED), f"Unexpected embedding shape {embeddings.shape}"
    print(f"✓ Features extracted successfully. Dimensions: {embeddings.shape}")
    
    # 3. Classifier Test
    print("\nRunning Test 3: Classification MLP...")
    classifier = MSIClassifier(input_dim=DIM_FUSED)
    classifier.eval()
    prob_tensor = classifier(embeddings)

    assert prob_tensor.shape == (1, 1), f"Unexpected prob shape {prob_tensor.shape}"
    prob = prob_tensor.item()
    assert 0.0 <= prob <= 1.0, f"Probability {prob} out of bounds"
    print(f"✓ Classifier forward pass successful. Predicted MSI-High Probability: {prob:.4%}")
    
    # 4. FAISS Index & Case Retrieval Test
    print("\nRunning Test 4: Case Retrieval (FAISS)...")
    retriever = CaseRetriever(index_dim=DIM_FUSED)
    
    # Build mini-index
    mock_emb = np.random.randn(5, DIM_FUSED).astype(np.float32)
    # L2 normalize
    mock_emb /= np.linalg.norm(mock_emb, axis=1, keepdims=True)
    
    mock_labels = [0, 1, 0, 1, 0]
    mock_descriptions = ["Desc A", "Desc B", "Desc C", "Desc D", "Desc E"]
    mock_patient_ids = ["P-1", "P-2", "P-3", "P-4", "P-5"]
    
    retriever.build_index(mock_emb, mock_labels, mock_descriptions, mock_patient_ids)
    
    # Query index
    matches = retriever.retrieve_top_k(embeddings.cpu().numpy(), k=3)
    assert len(matches) == 3, f"Expected 3 matches, got {len(matches)}"
    for match in matches:
        assert "similarity" in match, "Matching similarity score missing"
        assert "patient_id" in match, "Patient metadata missing"
        assert 0.0 <= match["similarity"] <= 1.0, "Similarity out of bounds"
    print("✓ FAISS build, serialization, and top-3 retrieval verified successfully.")
    
    # 5. Occlusion Heatmap Test
    print("\nRunning Test 5: Sliding-Window Occlusion Saliency...")
    heatmap_overlay = generate_occlusion_heatmap(test_img, extractor, classifier, window_size=40, stride=20)
    assert isinstance(heatmap_overlay, Image.Image), "Heatmap output is not PIL Image"
    assert heatmap_overlay.size == (224, 224), f"Unexpected heatmap size {heatmap_overlay.size}"
    print(f"✓ Occlusion heatmap visual overlay generated successfully. Resolution: {heatmap_overlay.size}")
    
    # 6. End-to-End Inference Pipeline Orchestrator Test
    print("\nRunning Test 6: End-to-End Pipeline Orchestration...")
    pipeline = GemmaSightInferencePipeline(force_simulation=True)
    prob_p, label_p, heatmap_p, cases_p, report_p = pipeline.run_diagnosis(test_img)
    
    assert 0.0 <= prob_p <= 1.0, "Pipeline probability out of bounds"
    assert label_p in ["MSS", "MSI-High"], f"Unexpected label {label_p}"
    assert isinstance(heatmap_p, Image.Image), "Pipeline heatmap is not PIL Image"
    assert len(cases_p) == 3, f"Expected 3 pipeline cases, got {len(cases_p)}"
    assert isinstance(report_p, str) and len(report_p) > 0, "Pipeline report is empty or invalid"
    print("✓ End-to-end pipeline run executed successfully!")
    print(f"  - Calculated Probability: {prob_p:.2%}")
    print(f"  - Status Classification: {label_p}")
    print(f"  - Retrieved Case Matches: {', '.join([c['patient_id'] for c in cases_p])}")
    print(f"  - Clinical Report Word Count: {len(report_p.split())} words")
    
    print("\n===========================================")
    print("🎉 ALL GEMMASIGHT TESTS PASSED SUCCESSFULLY! 🎉")
    print("===========================================")

if __name__ == "__main__":
    run_pipeline_tests()