"""Evaluation script for ResNet phenology classifier."""

import torch
import torch.nn as nn
from sklearn.metrics import accuracy_score, recall_score, f1_score, confusion_matrix, classification_report
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
import json
import os
from tqdm import tqdm

from src.model import create_model
from src.data_loader import get_data_loaders
from src.utils import load_checkpoint


def plot_confusion_matrix(
    cm: np.ndarray,
    class_names: list,
    save_path: str = None,
    normalize: bool = True
):
    """
    Plot confusion matrix.
    
    Args:
        cm: Confusion matrix
        class_names: List of class names
        save_path: Path to save plot
        normalize: Whether to normalize the confusion matrix
    """
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        fmt = '.2f'
    else:
        fmt = 'd'
    
    plt.figure(figsize=(10, 8))
    sns.heatmap(
        cm,
        annot=True,
        fmt=fmt,
        cmap='Blues',
        xticklabels=class_names,
        yticklabels=class_names,
        cbar_kws={'label': 'Proportion' if normalize else 'Count'}
    )
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.title('Confusion Matrix' + (' (Normalized)' if normalize else ''))
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Confusion matrix saved to {save_path}")
    
    plt.show()


def evaluate_model(
    model: nn.Module,
    dataloader: torch.utils.data.DataLoader,
    device: str,
    class_names: list
) -> dict:
    """
    Evaluate model and compute metrics.
    
    Args:
        model: Model to evaluate
        dataloader: Data loader
        device: Device to use
        class_names: List of class names
        
    Returns:
        Dictionary containing evaluation metrics
    """
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        pbar = tqdm(dataloader, desc='Evaluating')
        for images, labels in pbar:
            images = images.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.numpy())
    
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)
    
    # Compute metrics
    accuracy = accuracy_score(all_labels, all_preds)
    recall = recall_score(all_labels, all_preds, average='macro')
    macro_f1 = f1_score(all_labels, all_preds, average='macro')
    
    # Per-class metrics
    per_class_recall = recall_score(all_labels, all_preds, average=None)
    per_class_f1 = f1_score(all_labels, all_preds, average=None)
    
    # Confusion matrix
    cm = confusion_matrix(all_labels, all_preds)
    
    # Classification report
    report = classification_report(
        all_labels, all_preds,
        target_names=class_names,
        digits=4
    )
    
    metrics = {
        'accuracy': float(accuracy),
        'recall_macro': float(recall),
        'f1_macro': float(macro_f1),
        'per_class_recall': {class_names[i]: float(per_class_recall[i]) for i in range(len(class_names))},
        'per_class_f1': {class_names[i]: float(per_class_f1[i]) for i in range(len(class_names))},
        'confusion_matrix': cm.tolist(),
        'classification_report': report
    }
    
    return metrics, cm


def evaluate(
    model_path: str,
    data_dir: str,
    csv_file: str,
    class_mapping_path: str,
    output_dir: str = 'evaluation',
    batch_size: int = 32,
    num_workers: int = 4,
    split: str = 'test'
):
    """
    Main evaluation function.
    
    Args:
        model_path: Path to trained model checkpoint
        data_dir: Directory containing images
        csv_file: Path to CSV file with labels
        class_mapping_path: Path to class mapping JSON file
        output_dir: Directory to save evaluation results
        batch_size: Batch size
        num_workers: Number of data loader workers
        split: Which split to evaluate ('val' or 'test')
    """
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    # Device configuration
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    # Load class mapping
    print("\nLoading class mapping...")
    with open(class_mapping_path, 'r') as f:
        class_mapping = json.load(f)
    
    classes = class_mapping['classes']
    num_classes = len(classes)
    print(f"Number of classes: {num_classes}")
    print(f"Classes: {classes}")
    
    # Load data
    print(f"\nLoading {split} data...")
    data_loaders, _, _ = get_data_loaders(
        csv_file=csv_file,
        root_dir=data_dir,
        batch_size=batch_size,
        num_workers=num_workers
    )
    
    dataloader = data_loaders[split]
    print(f"Number of {split} samples: {len(dataloader.dataset)}")
    
    # Load model
    print("\nLoading model...")
    model = create_model(num_classes, pretrained=False, device=device)
    load_checkpoint(model_path, model, device=device)
    
    # Evaluate
    print(f"\nEvaluating on {split} set...")
    metrics, cm = evaluate_model(model, dataloader, device, classes)
    
    # Print results
    print(f"\n{'='*50}")
    print(f"EVALUATION RESULTS ({split.upper()} SET)")
    print(f"{'='*50}")
    print(f"Accuracy:     {metrics['accuracy']:.4f} ({metrics['accuracy']*100:.2f}%)")
    print(f"Recall (Macro): {metrics['recall_macro']:.4f}")
    print(f"F1 (Macro):   {metrics['f1_macro']:.4f}")
    print(f"\nPer-Class Metrics:")
    print("-" * 50)
    for cls in classes:
        print(f"  {cls:20s} - Recall: {metrics['per_class_recall'][cls]:.4f}, F1: {metrics['per_class_f1'][cls]:.4f}")
    print(f"{'='*50}")
    
    print(f"\nClassification Report:")
    print(metrics['classification_report'])
    
    # Save metrics
    metrics_path = os.path.join(output_dir, f'{split}_metrics.json')
    with open(metrics_path, 'w') as f:
        # Don't save classification report in JSON (it's a string)
        metrics_to_save = {k: v for k, v in metrics.items() if k != 'classification_report'}
        json.dump(metrics_to_save, f, indent=2)
    print(f"\nMetrics saved to: {metrics_path}")
    
    # Save classification report
    report_path = os.path.join(output_dir, f'{split}_classification_report.txt')
    with open(report_path, 'w') as f:
        f.write(metrics['classification_report'])
    print(f"Classification report saved to: {report_path}")
    
    # Plot confusion matrix
    cm_path = os.path.join(output_dir, f'{split}_confusion_matrix.png')
    plot_confusion_matrix(cm, classes, save_path=cm_path, normalize=True)
    
    # Also save non-normalized version
    cm_raw_path = os.path.join(output_dir, f'{split}_confusion_matrix_raw.png')
    plot_confusion_matrix(cm, classes, save_path=cm_raw_path, normalize=False)
    
    # Check if model meets success criteria
    if metrics['accuracy'] >= 0.9:
        print(f"\n✓ Model meets success criteria (accuracy >= 90%)")
    else:
        print(f"\n✗ Model does not meet success criteria (accuracy < 90%)")
        print(f"  Current accuracy: {metrics['accuracy']*100:.2f}%")
        print(f"  Model should be flagged for retraining")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Evaluate ResNet phenology classifier')
    parser.add_argument('--model_path', type=str, required=True, help='Path to model checkpoint')
    parser.add_argument('--data_dir', type=str, required=True, help='Directory containing images')
    parser.add_argument('--csv_file', type=str, required=True, help='Path to CSV file with labels')
    parser.add_argument('--class_mapping', type=str, required=True, help='Path to class mapping JSON')
    parser.add_argument('--output_dir', type=str, default='evaluation', help='Directory to save results')
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size')
    parser.add_argument('--num_workers', type=int, default=4, help='Number of data loader workers')
    parser.add_argument('--split', type=str, default='test', choices=['val', 'test'], help='Split to evaluate')
    
    args = parser.parse_args()
    
    evaluate(
        model_path=args.model_path,
        data_dir=args.data_dir,
        csv_file=args.csv_file,
        class_mapping_path=args.class_mapping,
        output_dir=args.output_dir,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        split=args.split
    )
