import os
import argparse
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from TrainModel import SimpleAE  # Ensure TrainModel.py is in PYTHONPATH

def load_checkpoint(path, device='cpu'):
    """
    Load autoencoder model and threshold from checkpoint file.
    """
    # Load full checkpoint
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    model = SimpleAE().to(device)
    model.load_state_dict(checkpoint['model_state'])
    threshold = checkpoint['threshold']
    model.eval()
    return model, threshold


def preprocess_image(image_path, img_size):
    """
    Load and preprocess image: resize, grayscale, to tensor.
    """
    img = Image.open(image_path).convert('L')
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor()
    ])
    tensor = transform(img).unsqueeze(0)  # shape: 1x1xHxW
    return tensor


def compute_reconstruction_error(model, x_tensor, device='cpu'):
    """
    Compute mean squared reconstruction error for input tensor.
    """
    x = x_tensor.to(device)
    with torch.no_grad():
        recon = model(x)
    mse = torch.mean((recon - x) ** 2).item()
    return mse


def map_error_to_quality(mse, threshold):
    """
    Convert reconstruction error to a quality percentage in [0,100].
    - If mse <= threshold: quality = 100 * (1 - mse/threshold)
    - Else: quality = 0
    Confidence equals quality.
    """
    if mse <= threshold:
        quality = (1 - mse / threshold) * 100
    else:
        quality = 0.0
    # Clamp to [0,100]
    quality = max(0.0, min(100.0, quality))
    confidence = quality
    return round(quality, 2), round(confidence, 2)


def main(args):
    # Validate inputs
    if not os.path.isfile(args.model_path):
        raise FileNotFoundError(f"Model checkpoint not found: {args.model_path}")
    if not os.path.isfile(args.image_path):
        raise FileNotFoundError(f"Image file not found: {args.image_path}")

    # Load model and threshold
    model, threshold = load_checkpoint(args.model_path, device='cpu')

    # Preprocess image
    x = preprocess_image(args.image_path, args.img_size)

    # Compute error
    mse_error = compute_reconstruction_error(model, x, device='cpu')

    # Map to quality and confidence
    quality_pct, confidence_pct = map_error_to_quality(mse_error, threshold)

    # Output detailed results
    print(f"Image: {os.path.basename(args.image_path)}")
    print(f"Reconstruction MSE Error: {mse_error:.6f}")
    print(f"Threshold (95th percentile): {threshold:.6f}")
    print(f"Quality Score: {quality_pct:.2f}%")
    print(f"Confidence: {confidence_pct:.2f}%")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Assess quality of a radiology image using trained QA model')
    parser.add_argument('--model_path', type=str, default='./models/autoencoder_qc.pth', help='Path to QA model checkpoint')
    parser.add_argument('--image_path', type=str, required=True, help='Path to radiology image (PNG/JPG)')
    parser.add_argument('--img_size', type=int, default=128, help='Image resize dimension (square)')
    args = parser.parse_args()
    main(args)