import os
import random
import numpy as np
import rasterio
import matplotlib.pyplot as plt
from pathlib import Path

# Setup paths based on the script location
SCRIPT_DIR = Path(__file__).resolve().parent
BASE_DIR = SCRIPT_DIR.parent
TRAIN_FEAT_DIR = BASE_DIR / "patches_128" / "train" / "features"
TRAIN_MASK_DIR = BASE_DIR / "patches_128" / "train" / "masks"
OUT_DIR = SCRIPT_DIR / "plots"

def normalize_band(band, pmin=2, pmax=98):
    """Normalize a band for visualization using percentiles (robust to outliers)."""
    # Exclude NaNs and extreme values when calculating percentiles
    valid_pixels = band[~np.isnan(band)]
    if len(valid_pixels) == 0:
        return np.zeros_like(band)
        
    b_min, b_max = np.percentile(valid_pixels, (pmin, pmax))
    if b_max - b_min == 0:
        return np.zeros_like(band)
        
    band_norm = (band - b_min) / (b_max - b_min)
    return np.clip(band_norm, 0, 1)

def main(num_samples=10):
    # Ensure the output directory exists
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    # Get a list of all feature files in the train split
    feature_files = list(TRAIN_FEAT_DIR.glob("feature_*.tif"))
    
    if not feature_files:
        print(f"No feature files found in {TRAIN_FEAT_DIR}")
        print("Please make sure you have run the patch building script.")
        return

    # Randomly select a few files to visualize
    samples = random.sample(feature_files, min(num_samples, len(feature_files)))
    
    # Optional: we can force it to pick at least some patches that contain a glacier
    # by reading the metadata, but for a quick script we will just randomly sample.
    
    print(f"Saving {len(samples)} visualization plots to {OUT_DIR} ...")
    
    for feat_path in samples:
        # Construct the matching mask path
        mask_name = feat_path.name.replace("feature_", "mask_")
        mask_path = TRAIN_MASK_DIR / mask_name

        if not mask_path.exists():
            print(f"Mask not found for {feat_path.name}, skipping.")
            continue

        # Read feature data
        with rasterio.open(feat_path) as src:
            # According to config.py, the bands are:
            # 1: B2(Blue), 2: B3(Green), 3: B4(Red), 4: B8(NIR), 5: B11, 6: B12
            # 7: NDVI, 8: NDWI, 9: NDSI, 10: Elev, 11: Slope, 12: Aspect
            b = src.read(1)
            g = src.read(2)
            r = src.read(3)
            
            ndvi = src.read(7)
            elev = src.read(10)

        # Read mask data
        with rasterio.open(mask_path) as src:
            mask = src.read(1)

        # Normalize RGB for display
        r_norm = normalize_band(r)
        g_norm = normalize_band(g)
        b_norm = normalize_band(b)
        rgb_img = np.dstack((r_norm, g_norm, b_norm))

        # Create the plot
        fig, axes = plt.subplots(1, 4, figsize=(20, 5))
        
        # 1. True Color (RGB)
        axes[0].imshow(rgb_img)
        axes[0].set_title("True Color (B4-B3-B2)")
        axes[0].axis('off')
        
        # 2. NDVI
        im1 = axes[1].imshow(ndvi, cmap='RdYlGn', vmin=-1, vmax=1)
        axes[1].set_title("NDVI")
        axes[1].axis('off')
        plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)
        
        # 3. Elevation
        im2 = axes[2].imshow(elev, cmap='terrain')
        axes[2].set_title("Elevation (m)")
        axes[2].axis('off')
        plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04)

        # 4. Ground Truth Mask
        axes[3].imshow(mask, cmap='gray', interpolation='nearest')
        axes[3].set_title("Ground Truth Mask")
        axes[3].axis('off')

        # Add a super title with the patch ID
        patch_id = feat_path.stem.replace('feature_', '')
        plt.suptitle(f"Patch: {patch_id}", fontsize=16)
        plt.tight_layout()
        
        # Save the figure
        out_path = OUT_DIR / f"{patch_id}.png"
        plt.savefig(out_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        print(f"Saved {out_path.name}")

if __name__ == "__main__":
    main(10)