"""Inference script for ResNet phenology classifier."""

import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import argparse
import json
import os
import matplotlib.pyplot as plt
import numpy as np

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


def load_inference_model(model_path: str, class_mapping_path: str, device: str = None):
    """
    Load model for inference.
    
    Args:
        model_path: Path to model checkpoint
        class_mapping_path: Path to class mapping JSON
        device: Device to use
        
    Returns:
        Tuple of (model, classes, device)
    """
    if device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # Load class mapping
    with open(class_mapping_path, 'r') as f:
        class_mapping = json.load(f)
    
    classes = class_mapping['classes']
    num_classes = len(classes)
    
    # Load model
    model = create_model(num_classes, pretrained=False, device=device)
    load_checkpoint(model_path, model, device=device)
    model.eval()
    
    return model, classes, device


def preprocess_image(image_path: str) -> torch.Tensor:
    """
    Preprocess image for inference.
    
    Args:
        image_path: Path to image file
        
    Returns:
        Preprocessed image tensor
    """
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    image = Image.open(image_path).convert('RGB')
    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
    
    return image_tensor, image


def predict(
    model: torch.nn.Module,
    image_tensor: torch.Tensor,
    classes: list,
    device: str
) -> dict:
    """
    Make prediction on image.
    
    Args:
        model: Trained model
        image_tensor: Preprocessed image tensor
        classes: List of class names
        device: Device to use
        
    Returns:
        Dictionary with prediction results
    """
    image_tensor = image_tensor.to(device)
    
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = F.softmax(outputs, dim=1)
        confidence, predicted = torch.max(probabilities, 1)
    
    predicted_class = classes[predicted.item()]
    confidence_score = confidence.item()
    
    # Get probabilities for all classes
    all_probs = {classes[i]: probabilities[0][i].item() for i in range(len(classes))}
    
    result = {
        'phase': predicted_class,
        'confidence': confidence_score,
        'probabilities': all_probs
    }
    
    return result


def visualize_prediction(
    image: Image.Image,
    result: dict,
    save_path: str = None
):
    """
    Visualize prediction result.
    
    Args:
        image: Original PIL image
        result: Prediction result dictionary
        save_path: Optional path to save visualization
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
    
    # Display image
    ax1.imshow(image)
    ax1.axis('off')
    ax1.set_title(f"Predicted: {result['phase']}\nConfidence: {result['confidence']:.2%}", 
                  fontsize=14, fontweight='bold')
    
    # Display probabilities
    classes = list(result['probabilities'].keys())
    probs = list(result['probabilities'].values())
    
    colors = ['green' if cls == result['phase'] else 'lightblue' for cls in classes]
    
    ax2.barh(classes, probs, color=colors)
    ax2.set_xlabel('Probability', fontsize=12)
    ax2.set_title('Class Probabilities', fontsize=14, fontweight='bold')
    ax2.set_xlim(0, 1)
    
    # Add probability values on bars
    for i, (cls, prob) in enumerate(zip(classes, probs)):
        ax2.text(prob + 0.02, i, f'{prob:.2%}', va='center', fontsize=10)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Visualization saved to {save_path}")
    
    plt.show()


def classify_image(
    image_path: str,
    model_path: str,
    class_mapping_path: str,
    output_dir: str = None,
    visualize: bool = True
):
    """
    Classify a single image.
    
    Args:
        image_path: Path to image file
        model_path: Path to model checkpoint
        class_mapping_path: Path to class mapping JSON
        output_dir: Optional directory to save results
        visualize: Whether to create visualization
    """
    print(f"Loading model from {model_path}...")
    model, classes, device = load_inference_model(model_path, class_mapping_path)
    
    print(f"Processing image: {image_path}")
    image_tensor, original_image = preprocess_image(image_path)
    
    print("Making prediction...")
    result = predict(model, image_tensor, classes, device)
    
    # Print results
    print(f"\n{'='*50}")
    print(f"CLASSIFICATION RESULT")
    print(f"{'='*50}")
    print(f"Predicted Phase: {result['phase']}")
    print(f"Confidence: {result['confidence']:.2%}")
    print(f"\nAll Probabilities:")
    for cls, prob in sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True):
        print(f"  {cls:20s}: {prob:.2%}")
    print(f"{'='*50}")
    
    # Save results
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)
        
        # Save JSON result
        result_path = os.path.join(output_dir, 'prediction.json')
        with open(result_path, 'w') as f:
            json.dump(result, f, indent=2)
        print(f"\nPrediction saved to: {result_path}")
    
    # Visualize
    if visualize:
        save_path = os.path.join(output_dir, 'prediction_visualization.png') if output_dir else None
        visualize_prediction(original_image, result, save_path)
    
    return result


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Classify plant image by phenological phase')
    parser.add_argument('--image_path', type=str, required=True, help='Path to image file')
    parser.add_argument('--model_path', type=str, required=True, help='Path to model checkpoint')
    parser.add_argument('--class_mapping', type=str, required=True, help='Path to class mapping JSON')
    parser.add_argument('--output_dir', type=str, default=None, help='Directory to save results')
    parser.add_argument('--no_visualize', action='store_true', help='Do not create visualization')
    
    args = parser.parse_args()
    
    classify_image(
        image_path=args.image_path,
        model_path=args.model_path,
        class_mapping_path=args.class_mapping,
        output_dir=args.output_dir,
        visualize=not args.no_visualize
    )
