import torch
import numpy as np
import cv2
import json
import sys
from pathlib import Path
from ensemble_models import TBEnsemble
from preprocessing import LungPreprocessor, get_val_transforms

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

_model = None
_preprocessor = None
_transforms = None


def _load_model():
    global _model, _preprocessor, _transforms
    if _model is None:
        _model = TBEnsemble()
        state = torch.load("models/ensemble_best.pth", map_location=DEVICE)
        _model.load_state_dict(state)
        _model.to(DEVICE)
        _model.eval()
        _preprocessor = LungPreprocessor()
        _transforms = get_val_transforms(224)


def predict(image_path, threshold=0.52):
    _load_model()
    img = _preprocessor.preprocess(str(image_path), segment_lung=True)
    if img is None:
        return None
    augmented = _transforms(image=img)
    tensor = augmented['image'].unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        prob = _model(tensor).item()
    return {"file": Path(image_path).name, "tb_probability": round(prob, 4), "prediction": "TB" if prob > threshold else "Normal"}


def evaluate_dir(dir_path, threshold=0.52):
    _load_model()
    paths = [p for p in Path(dir_path).rglob("*") if p.suffix.lower() in (".png", ".jpg", ".jpeg")]
    results = []
    for p in paths:
        r = predict(p, threshold)
        if r:
            results.append(r)
            print(f"  {r['file']:40s} {r['prediction']:8s} ({r['tb_probability']:.4f})")
    return results


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python predict.py <image_or_dir> [threshold]")
        sys.exit(1)
    path = Path(sys.argv[1])
    thresh = float(sys.argv[2]) if len(sys.argv) > 2 else 0.52
    if path.is_dir():
        results = evaluate_dir(path, thresh)
        tbs = sum(1 for r in results if r["prediction"] == "TB")
        print(f"\n{tbs}/{len(results)} TB positive")
    else:
        r = predict(path, thresh)
        if r:
            print(f"{r['file']}: {r['prediction']} (TB prob: {r['tb_probability']:.4f})")
        else:
            print("Failed to load image")