import os
import time
import argparse
import torch
import numpy as np
from glob import glob
from PIL import Image, ImageFilter
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, accuracy_score, roc_curve, auc
import torchvision.transforms as transforms
from TrainModel import SimpleAE  # ensure TrainModel.py in PYTHONPATH

# -------- Dataset for Testing --------
def load_test_set(data_dir, img_size, noise_std):
    files = glob(os.path.join(data_dir, '*.png')) + glob(os.path.join(data_dir, '*.jpg'))
    X, y = [], []
    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.Grayscale(num_output_channels=1),
        transforms.ToTensor()
    ])
    for p in files:
        img = Image.open(p).convert('L')
        x = transform(img)
        # clean
        X.append(x); y.append(1)
        # noisy
        arr = np.array(img).astype(np.float32)
        arr_n = np.clip(arr + np.random.normal(0, noise_std, arr.shape), 0, 255).astype(np.uint8)
        img_n = Image.fromarray(arr_n)
        X.append(transform(img_n)); y.append(0)
        # blurred
        img_b = img.filter(ImageFilter.GaussianBlur(radius=2))
        X.append(transform(img_b)); y.append(0)
    return torch.stack(X), np.array(y)

# -------- Testing Script --------
def main(args):
    os.makedirs(args.output_dir, exist_ok=True)

    # Load checkpoint with weights_only=False to allow loading full dict
    checkpoint = torch.load(args.model_path, map_location='cpu', weights_only=False)
    model = SimpleAE()
    model.load_state_dict(checkpoint['model_state'])
    threshold = checkpoint['threshold']
    model.eval()

    # Prepare test set
    X, y_true = load_test_set(args.data_dir, args.img_size, args.noise_std)

    # Inference metrics
    errors, times = [], []
    with torch.no_grad():
        for x in X:
            start = time.time()
            out = model(x.unsqueeze(0))
            times.append(time.time() - start)
            err = ((out - x.unsqueeze(0))**2).mean().item()
            errors.append(err)
    y_pred = np.array([1 if e <= threshold else 0 for e in errors])

    # Calculate metrics
    acc = accuracy_score(y_true, y_pred)
    prec = precision_score(y_true, y_pred)
    rec = recall_score(y_true, y_pred)
    fpr, tpr, _ = roc_curve(y_true, [-e for e in errors])
    roc_auc = auc(fpr, tpr)
    avg_time = np.mean(times)
    model_size = os.path.getsize(args.model_path) / (1024*1024)

    # Save summary
    with open(os.path.join(args.output_dir, 'metrics.txt'), 'w') as f:
        f.write(f"Accuracy: {acc:.4f}\nPrecision: {prec:.4f}\nRecall: {rec:.4f}\nROC AUC: {roc_auc:.4f}\n")
        f.write(f"Avg inference time (s): {avg_time:.6f}\nModel size (MB): {model_size:.2f}\n")

    # Plot and save
    plt.figure()
    plt.hist(errors, bins=50)
    plt.title('Reconstruction Error Distribution')
    plt.xlabel('MSE Error')
    plt.ylabel('Count')
    plt.savefig(os.path.join(args.output_dir, 'error_distribution.png'))
    plt.close()

    plt.figure()
    plt.plot(fpr, tpr, label=f'AUC = {roc_auc:.2f}')
    plt.plot([0,1],[0,1],'--')
    plt.title('ROC Curve')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.legend(loc='lower right')
    plt.savefig(os.path.join(args.output_dir, 'roc_curve.png'))
    plt.close()

    plt.figure()
    plt.plot(sorted(times), marker='.')
    plt.title('Inference Time per Sample')
    plt.xlabel('Sample index')
    plt.ylabel('Time (s)')
    plt.savefig(os.path.join(args.output_dir, 'inference_time.png'))
    plt.close()

    print(f"Test results saved in {args.output_dir}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Evaluate QA model')
    parser.add_argument('--data_dir', type=str, default='./data', help='Test images directory')
    parser.add_argument('--model_path', type=str, default='./models/autoencoder_qc.pth', help='Path to QC model')
    parser.add_argument('--output_dir', type=str, default='./TestModel', help='Directory to save results')
    parser.add_argument('--img_size', type=int, default=128, help='Image resize')
    parser.add_argument('--noise_std', type=float, default=25.0, help='Synthetic noise std')
    args = parser.parse_args()
    main(args)