import torch
import os
from pathlib import Path

checkpoint_path = "files/checkpoint.pth"

print("=" * 60)
print("CHECKPOINT QUALITY ANALYSIS")
print("=" * 60)

try:
    checkpoint = torch.load(checkpoint_path, map_location='cpu')
    
    print("\n1. FILE INFORMATION:")
    print(f"   File size: {os.path.getsize(checkpoint_path) / (1024*1024):.2f} MB")
    print(f"   File exists: YES")
    
    print("\n2. CHECKPOINT STRUCTURE:")
    if isinstance(checkpoint, dict):
        print(f"   Type: Dictionary")
        print(f"   Keys: {list(checkpoint.keys())}")
    else:
        print(f"   Type: Direct state dict")
    
    print("\n3. TRAINING METADATA:")
    if isinstance(checkpoint, dict):
        if 'epoch' in checkpoint:
            print(f"   Epoch: {checkpoint['epoch']}")
        if 'loss' in checkpoint:
            print(f"   Final Loss: {checkpoint['loss']:.6f}")
        if 'best_loss' in checkpoint:
            print(f"   Best Loss: {checkpoint['best_loss']:.6f}")
    
    print("\n4. MODEL STATE:")
    state_dict = checkpoint.get('model_state_dict', checkpoint) if isinstance(checkpoint, dict) else checkpoint
    print(f"   Number of parameters: {len(state_dict)}")
    print(f"   First 5 layer names:")
    for i, key in enumerate(list(state_dict.keys())[:5]):
        print(f"      - {key}")
    
    total_params = 0
    for key, tensor in state_dict.items():
        total_params += tensor.numel()
    print(f"   Total trainable parameters: {total_params:,}")
    
    print("\n5. LOADING TEST:")
    from app.models.unet_model import BuildUNet
    model = BuildUNet(num_classes=4)
    model.load_state_dict(state_dict)
    model.eval()
    print(f"   Model loaded: YES")
    print(f"   Model in eval mode: YES")
    
    print("\n6. INFERENCE TEST:")
    dummy_input = torch.randn(1, 3, 256, 512)
    with torch.no_grad():
        output = model(dummy_input)
    print(f"   Input shape: {dummy_input.shape}")
    print(f"   Output shape: {output.shape}")
    print(f"   Output classes: {output.shape[1]}")
    print(f"   Inference successful: YES")
    
    print("\n7. OUTPUT ANALYSIS:")
    probs = torch.softmax(output, dim=1)
    pred_mask = torch.argmax(probs, dim=1)
    unique_classes = torch.unique(pred_mask)
    print(f"   Predicted classes in dummy test: {unique_classes.tolist()}")
    
    class_percentages = []
    for i in range(4):
        percentage = (pred_mask == i).sum().item() / pred_mask.numel() * 100
        class_percentages.append(percentage)
        print(f"   Class {i}: {percentage:.2f}%")
    
    print("\n8. QUALITY INDICATORS:")
    if class_percentages[0] > 90:
        print(f"   WARNING: Dummy test predicts {class_percentages[0]:.1f}% background")
        print(f"   This MAY indicate class imbalance")
    
    if checkpoint.get('loss', 1.0) > 0.5:
        print(f"   WARNING: Training loss is high ({checkpoint.get('loss', 'N/A')})")
        print(f"   Model may be undertrained")
    
    print("\n9. TRAINING SCORES (from score.csv):")
    mean_f1 = 0.0
    try:
        with open('files/score.csv', 'r') as f:
            scores = f.read()
            print(scores)
            for line in scores.split('\n'):
                if line.startswith('Mean'):
                    parts = line.split(',')
                    if len(parts) >= 2:
                        mean_f1 = float(parts[1])
    except:
        print("   score.csv not found")
    
    print("\n" + "=" * 60)
    print("SUMMARY:")
    print("=" * 60)
    print(f"Mean F1 Score: {mean_f1:.4f} ({mean_f1*100:.2f}%)")
    
    if mean_f1 < 0.60:
        print("\nCHECKPOINT QUALITY: POOR")
        print("   Reason: Mean F1 score is {:.2f}% (need 80%+ for production)".format(mean_f1*100))
        print("   - Model is undertrained")
        print("   - Training dataset too small (997 images)")
        print("   - Needs retraining with 10,000+ images")
        print("   - Class-specific performance very low")
    elif mean_f1 < 0.80:
        print("\nCHECKPOINT QUALITY: MODERATE")
        print("   Reason: Mean F1 score is {:.2f}% (acceptable but not great)".format(mean_f1*100))
        print("   - Model works but could be better")
        print("   - More training data recommended")
    else:
        print("\nCHECKPOINT QUALITY: GOOD")
        print("   Reason: Mean F1 score is {:.2f}% (production ready)".format(mean_f1*100))
        print("   - Model is properly trained")
        print("   - Ready for deployment")
    print("=" * 60)
    
except Exception as e:
    print(f"\nERROR: {str(e)}")
    import traceback
    traceback.print_exc()