#!/usr/bin/env python3
"""
Generate PNG previews for existing GeoTIFF prediction files
"""

import numpy as np
import rasterio
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from pathlib import Path
import sys

def generate_png_preview(tif_file, output_png=None):
    """Generate PNG preview from GeoTIFF file"""
    tif_path = Path(tif_file)
    
    if not tif_path.exists():
        print(f"❌ File not found: {tif_file}")
        return False
    
    # Determine output PNG path
    if output_png is None:
        output_png = tif_path.with_suffix('.png')
    else:
        output_png = Path(output_png)
    
    try:
        # Read GeoTIFF
        with rasterio.open(tif_path) as src:
            data = src.read(1)
            
        print(f"📊 Data shape: {data.shape}, range: [{np.nanmin(data):.3f}, {np.nanmax(data):.3f}]")
        
        # Determine if it's classification or NDVI based on filename
        is_classification = 'classification' in tif_path.name.lower() or 'prediction' in tif_path.name.lower()
        is_ndvi = 'ndvi' in tif_path.name.lower()
        
        # Create figure
        fig, ax = plt.subplots(figsize=(12, 10), dpi=150)
        
        if is_ndvi:
            # NDVI: use RdYlGn colormap, range -1 to 1
            im = ax.imshow(data, cmap='RdYlGn', vmin=-1, vmax=1, interpolation='nearest')
            ax.set_title(f'NDVI - {tif_path.stem}', fontsize=14, fontweight='bold')
            cbar_label = 'NDVI'
        elif is_classification:
            # Classification: use tab20 colormap
            im = ax.imshow(data, cmap='tab20', interpolation='nearest')
            ax.set_title(f'Land Classification - {tif_path.stem}', fontsize=14, fontweight='bold')
            cbar_label = 'Class'
        else:
            # Generic: use viridis
            im = ax.imshow(data, cmap='viridis', interpolation='nearest')
            ax.set_title(f'{tif_path.stem}', fontsize=14, fontweight='bold')
            cbar_label = 'Value'
        
        ax.set_xlabel('X (pixels)', fontsize=10)
        ax.set_ylabel('Y (pixels)', fontsize=10)
        
        # Add colorbar
        cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label(cbar_label, rotation=270, labelpad=15)
        
        # For classification, try to set integer ticks
        if is_classification:
            try:
                unique_vals = np.unique(data[~np.isnan(data)])
                if len(unique_vals) < 20:  # Only if not too many classes
                    cbar.set_ticks(unique_vals)
                    cbar.set_ticklabels([str(int(v)) for v in unique_vals])
            except:
                pass
        
        # Add grid
        ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        
        # Save PNG
        plt.tight_layout()
        plt.savefig(str(output_png), dpi=150, bbox_inches='tight')
        plt.close(fig)
        
        print(f"✅ Created PNG: {output_png}")
        return True
        
    except Exception as e:
        print(f"❌ Error creating PNG: {e}")
        import traceback
        traceback.print_exc()
        return False


def generate_all_previews(predictions_dir="predictions"):
    """Generate PNG previews for all GeoTIFF files without PNGs"""
    pred_path = Path(predictions_dir)
    
    if not pred_path.exists():
        print(f"❌ Directory not found: {predictions_dir}")
        return
    
    tif_files = list(pred_path.glob("*.tif"))
    print(f"🔍 Found {len(tif_files)} GeoTIFF files")
    
    generated = 0
    skipped = 0
    
    for tif_file in tif_files:
        png_file = tif_file.with_suffix('.png')
        
        if png_file.exists():
            print(f"⏭️  Skipping {tif_file.name} (PNG already exists)")
            skipped += 1
            continue
        
        print(f"\n🎨 Processing {tif_file.name}...")
        if generate_png_preview(tif_file):
            generated += 1
    
    print(f"\n{'='*60}")
    print(f"✅ Generated {generated} new PNG previews")
    print(f"⏭️  Skipped {skipped} files (already have PNGs)")
    print(f"{'='*60}")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        # Process specific file
        tif_file = sys.argv[1]
        generate_png_preview(tif_file)
    else:
        # Process all files in predictions directory
        generate_all_previews()