import argparse
import numpy as np
from pathlib import Path
from PIL import Image
import tensorflow as tf
import numpy as np
import torch
from torchvision import transforms
from model_pytorch import SaraCNN





CLASS_NAMES = ["buildings", "forest", "glacier", "mountain", "sea", "street"]
IMAGE_SIZE = (150, 150)
MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]



#  PyTorch Prediction
def predict_pytorch(image_path: str, model_path: str = "sara_model.pth") -> dict:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    checkpoint = torch.load(model_path, map_location=device)
    class_names = checkpoint.get("class_names", CLASS_NAMES)

    model = SaraCNN(num_classes=len(class_names))
    model.load_state_dict(checkpoint["model_state"])
    model.to(device)
    model.eval()

    transform = transforms.Compose([
        transforms.Resize(IMAGE_SIZE),
        transforms.ToTensor(),
        transforms.Normalize(MEAN, STD),
    ])

    img = Image.open(image_path).convert("RGB")
    tensor = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        logits = model(tensor)
        probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()

    idx = int(np.argmax(probs))
    pred_class = class_names[idx]
    confidence = float(probs[idx])

    return {
        "predicted_class":    pred_class,
        "confidence":         round(confidence * 100, 2),
        "all_probabilities":  {
            cls: round(float(p) * 100, 2)
            for cls, p in zip(class_names, probs)
        },
    }



#  TensorFlow Prediction
def predict_tensorflow(image_path: str, model_path: str = "sara_model.keras") -> dict:
    model = tf.keras.models.load_model(model_path, compile=False)

    img = tf.keras.utils.load_img(image_path, target_size=(150, 150))
    arr = tf.keras.utils.img_to_array(img) / 255.0
    arr = np.expand_dims(arr, axis=0)

   
    probs = model.predict(arr, verbose=0)[0]
    idx = int(np.argmax(probs))
    pred_class = CLASS_NAMES[idx]
    confidence = float(probs[idx])

    return {
        "predicted_class":   pred_class,
        "confidence":        round(confidence * 100, 2),
        "all_probabilities": {
            cls: round(float(p) * 100, 2)
            for cls, p in zip(CLASS_NAMES, probs)
        },
    }




if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model",      required=True, choices=["pytorch", "tensorflow"])
    parser.add_argument("--image",      required=True)
    parser.add_argument("--model_path", default=None)
    args = parser.parse_args()

    if args.model == "pytorch":
        path = args.model_path or "sara_model.pth"
        result = predict_pytorch(args.image, model_path=path)
    else:
        path = args.model_path or "sara_model.keras"
        result = predict_tensorflow(args.image, model_path=path)

    print(f"\n Classe     : {result['predicted_class']}")
    print(f" Confidence  : {result['confidence']}%")
    print("\n All probabilities :")
    for cls, prob in sorted(result["all_probabilities"].items(), key=lambda x: -x[1]):
        bar = " " * int(prob / 5)
        print(f"   {cls:<12} {prob:6.2f}%  {bar}")