import os
import rasterio
from rasterio.windows import Window
import numpy as np
import random
import glob

FEATURE_STACK = "mixed_res_dataset/stacked_features.tif"
MASK_STACK = "mixed_res_dataset/stacked_mask.tif"
OUTPUT_DIR = "dataset/extracted_patches"
PATCH_SIZE = 128
STRIDE = 64  # 50% overlap

def main():
    if not os.path.exists(FEATURE_STACK) or not os.path.exists(MASK_STACK):
        print("Stacked feature or mask not found. Please run align_and_stack.py first.")
        return

    os.makedirs(os.path.join(OUTPUT_DIR, "features"), exist_ok=True)
    os.makedirs(os.path.join(OUTPUT_DIR, "masks"), exist_ok=True)

    with rasterio.open(FEATURE_STACK) as src_feat, rasterio.open(MASK_STACK) as src_mask:
        height = src_feat.height
        width = src_feat.width
        num_bands = src_feat.count
        
        print(f"Opened stacked files. Size: {width}x{height}, Bands: {num_bands}")
        
        patch_idx = 0
        valid_patches = 0
        
        # Read the whole mask into memory to fast-check if a patch has any glacier
        # It's uint8, so for 11477x8969 it's ~100MB, easily fits in RAM
        print("Loading full mask into memory for filtering...")
        full_mask = src_mask.read(1)
        
        print("Extracting patches...")
        for y in range(0, height - PATCH_SIZE + 1, STRIDE):
            for x in range(0, width - PATCH_SIZE + 1, STRIDE):
                # Check mask first
                patch_mask = full_mask[y:y+PATCH_SIZE, x:x+PATCH_SIZE]
                
                # Filter: Keep if it contains any glacier OR 10% chance of background
                is_glacier = np.any(patch_mask > 0)
                if not is_glacier and random.random() > 0.1:
                    continue
                    
                # Read feature patch
                window = Window(x, y, PATCH_SIZE, PATCH_SIZE)
                patch_feat = src_feat.read(window=window)
                
                # patch_feat shape is (Bands, H, W) -> we need (H, W, Bands) for generator
                patch_feat = np.moveaxis(patch_feat, 0, -1)
                patch_mask_expanded = np.expand_dims(patch_mask, axis=-1)
                
                # Save as .npy
                feat_out = os.path.join(OUTPUT_DIR, "features", f"patch_{patch_idx:06d}.npy")
                mask_out = os.path.join(OUTPUT_DIR, "masks", f"patch_{patch_idx:06d}.npy")
                
                np.save(feat_out, patch_feat)
                np.save(mask_out, patch_mask_expanded)
                
                patch_idx += 1
                if is_glacier:
                    valid_patches += 1
                    
                if patch_idx % 100 == 0:
                    print(f"  Extracted {patch_idx} patches (Glacier: {valid_patches})")

        print(f"\nDone! Extracted total {patch_idx} patches.")
        print(f"Patches containing glacier: {valid_patches}")
        print(f"Background patches (kept): {patch_idx - valid_patches}")
        print(f"Saved to {OUTPUT_DIR}")

if __name__ == '__main__':
    main()