"""
Phase 3: Interpretability & Analysis
Generates Grad-CAM heatmaps and runs NDVI correlation analysis
using real DeepGlobe images through the EcoPulse pipeline.
"""
import argparse
import os
import yaml
import numpy as np
import cv2
import torch
from PIL import Image
from tqdm import tqdm

from src.cnn_model import load_model
from src.transforms import EUROSAT_TRANSFORM
from src.visualization import apply_grad_cam, save_grad_cam_overlay, plot_greenery_overlay
from src.pipeline import EcoPulsePipeline
from src.data_loader import DeepGlobeDataset
from src.metrics import ndvi_correlation


def generate_grad_cam_samples(config_path, num_samples=5):
    """
    Generates Grad-CAM heatmap overlays for a set of DeepGlobe images
    using the trained ResNet-50 classifier.
    """
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    output_dir = os.path.join(config['paths']['output_figures'], 'grad_cam')
    os.makedirs(output_dir, exist_ok=True)

    classes = config['classes']
    greenery_classes = config['greenery_classes']

    # Load the trained CNN
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    model = load_model(
        weights_path=os.path.join(config['paths']['output_models'], 'resnet50_eurosat.pth'),
        num_classes=config['model']['num_classes'],
        device=device
    )
    model.eval()

    transform = EUROSAT_TRANSFORM

    # Load DeepGlobe images
    deepglobe_dir = config['paths']['deepglobe_dir']
    dataset = DeepGlobeDataset(deepglobe_dir)
    num_samples = min(num_samples, len(dataset))

    print(f"\n{'='*60}")
    print(f"  GRAD-CAM HEATMAP GENERATION")
    print(f"  Processing {num_samples} DeepGlobe images...")
    print(f"{'='*60}\n")

    for i in tqdm(range(num_samples), desc="Generating Grad-CAMs"):
        img_name = dataset.image_files[i]
        img_path = os.path.join(dataset.images_dir, img_name)

        # Load the raw image for the overlay
        raw_img = cv2.imread(img_path)
        if raw_img is None:
            raise FileNotFoundError(f"Could not read image at {img_path} — file may be missing or corrupt")
        raw_img_rgb = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)

        # Create a center crop (64x64-like patch) for the CNN input
        h, w = raw_img_rgb.shape[:2]
        cx, cy = w // 2, h // 2
        patch_size = min(h, w) // 3
        crop = raw_img_rgb[cy - patch_size:cy + patch_size, cx - patch_size:cx + patch_size]

        crop_pil = Image.fromarray(crop)
        input_tensor = transform(crop_pil).unsqueeze(0)

        # Generate Grad-CAM
        heatmap, pred_idx = apply_grad_cam(model, input_tensor, target_class=None)
        pred_class = classes[pred_idx]
        is_green = pred_class in greenery_classes

        # Save the overlay
        base_name = os.path.splitext(img_name)[0]
        out_path = os.path.join(output_dir, f"{base_name}_gradcam.png")
        save_grad_cam_overlay(crop, heatmap, out_path)

        label = "[GREENERY]" if is_green else "[NON-GREEN]"
        print(f"  [{i+1}/{num_samples}] {img_name} -> Predicted: {pred_class} ({label})")

    print(f"\nAll Grad-CAM overlays saved to: {output_dir}")
    return output_dir


def run_ndvi_correlation(config_path, num_samples=5):
    """
    Runs the EcoPulse pipeline on DeepGlobe images, then computes a
    pseudo-NDVI from the RGB channels and correlates it with the
    pipeline's predicted greenery percentage.

    Note: DeepGlobe provides standard RGB imagery, not multispectral.
    We approximate NDVI using the Normalized Green-Red Difference Index (NGRDI):
        NGRDI = (Green - Red) / (Green + Red)
    This is a well-known vegetation proxy for RGB-only data.
    """
    print(f"\n{'='*60}")
    print(f"  NDVI CORRELATION ANALYSIS")
    print(f"{'='*60}\n")

    pipeline = EcoPulsePipeline(config_path)
    config = pipeline.config

    deepglobe_dir = config['paths']['deepglobe_dir']
    dataset = DeepGlobeDataset(deepglobe_dir)
    num_samples = min(num_samples, len(dataset))

    predicted_greenery = []
    ngrdi_values = []

    print(f"Processing {num_samples} images through the full pipeline...\n")

    for i in tqdm(range(num_samples), desc="Pipeline + NGRDI"):
        img_name = dataset.image_files[i]
        img_path = os.path.join(dataset.images_dir, img_name)

        # Run the full EcoPulse pipeline
        image_np, results = pipeline.process_image(img_path)
        predicted_greenery.append(results['greenery_percentage'])

        # Compute NGRDI (RGB-based vegetation proxy)
        green_band = image_np[:, :, 1].astype(float)
        red_band = image_np[:, :, 0].astype(float)
        denominator = green_band + red_band
        denominator[denominator == 0] = 1e-6
        ngrdi = (green_band - red_band) / denominator
        mean_ngrdi = np.mean(ngrdi)
        ngrdi_values.append(mean_ngrdi)

        print(f"  [{i+1}/{num_samples}] {img_name}: "
              f"EcoPulse Greenery = {results['greenery_percentage']:.1f}% | "
              f"Mean NGRDI = {mean_ngrdi:.4f}")

    # Compute correlation
    corr = ndvi_correlation(predicted_greenery, ngrdi_values)

    print(f"\n{'='*60}")
    print(f"  RESULTS")
    print(f"{'='*60}")
    print(f"  Pearson Correlation (EcoPulse vs NGRDI): {corr:.4f}")

    if corr > 0.7:
        print("  [OK] Strong positive correlation -- EcoPulse's deep learning approach")
        print("       closely aligns with traditional vegetation indices.")
    elif corr > 0.4:
        print("  [!!] Moderate correlation -- partial alignment with vegetation indices.")
    else:
        print("  [ii] Weak correlation -- expected when comparing learned features")
        print("       against a simple spectral proxy on RGB-only imagery.")

    # Save results
    output_dir = config['paths']['output_figures']
    os.makedirs(output_dir, exist_ok=True)

    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(8, 6))
    ax.scatter(ngrdi_values, predicted_greenery, c='forestgreen', s=100, edgecolors='black', zorder=5)

    # Trend line
    if len(ngrdi_values) >= 2:
        z = np.polyfit(ngrdi_values, predicted_greenery, 1)
        p = np.poly1d(z)
        x_line = np.linspace(min(ngrdi_values), max(ngrdi_values), 100)
        ax.plot(x_line, p(x_line), '--', color='darkgreen', alpha=0.7, label=f'Trend (r={corr:.3f})')

    ax.set_xlabel('Mean NGRDI (RGB Vegetation Proxy)', fontsize=12)
    ax.set_ylabel('EcoPulse Predicted Greenery (%)', fontsize=12)
    ax.set_title('NDVI Correlation: EcoPulse vs NGRDI', fontsize=14)
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()

    scatter_path = os.path.join(output_dir, 'ndvi_correlation_scatter.png')
    plt.savefig(scatter_path, dpi=150)
    plt.close()
    print(f"\n  Scatter plot saved to: {scatter_path}")

    return corr


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="EcoPulse Phase 3: Interpretability & Analysis")
    parser.add_argument("--config", default="config/config.yaml", help="Path to config file")
    parser.add_argument("--samples", type=int, default=5, help="Number of images to process")
    args = parser.parse_args()

    print("=" * 60)
    print("  EcoPulse: Interpretability & Analysis")
    print("=" * 60)

    # Part 1: Grad-CAM Heatmaps
    generate_grad_cam_samples(args.config, num_samples=args.samples)

    # Part 2: NDVI Correlation
    corr = run_ndvi_correlation(args.config, num_samples=args.samples)

    print(f"\n{'='*60}")
    print(f"  INTERPRETABILITY ANALYSIS COMPLETE")
    print(f"  Grad-CAM heatmaps: outputs/figures/grad_cam/")
    print(f"  NDVI scatter plot:  outputs/figures/ndvi_correlation_scatter.png")
    print(f"  Pearson Correlation: {corr:.4f}")
    print(f"{'='*60}")