import os
import glob
import rasterio
from rasterio.warp import reproject, Resampling
import numpy as np
import geopandas as gpd
from rasterio.features import rasterize

INPUT_DIR = "mixed_res_dataset"
SHAPEFILE_PATH = "Final_merged01/Final_merged01.shp"
OUTPUT_FEATURE = "mixed_res_dataset/stacked_features.tif"
OUTPUT_MASK = "mixed_res_dataset/stacked_mask.tif"

def main():
    # Find all downloaded bands
    band_files = sorted(glob.glob(os.path.join(INPUT_DIR, "Band*.tif")))
    if not band_files:
        print(f"No band files found in {INPUT_DIR}")
        return

    print(f"Found {len(band_files)} band files.")

    # We will use Band01 as the reference geometry (10m resolution)
    ref_file = [f for f in band_files if "Band01" in f][0]
    
    with rasterio.open(ref_file) as src_ref:
        ref_meta = src_ref.meta.copy()
        ref_transform = src_ref.transform
        ref_crs = src_ref.crs
        ref_height = src_ref.height
        ref_width = src_ref.width

    print(f"Reference geometry based on {os.path.basename(ref_file)}:")
    print(f"  Shape: ({ref_height}, {ref_width})")
    print(f"  Transform: {ref_transform}")
    print(f"  CRS: {ref_crs}")

    # Prepare output metadata for features
    ref_meta.update({
        "count": len(band_files),
        "dtype": "float32",
        "nodata": 0.0,
        "compress": "lzw",
        "interleave": "band", # band sequential for easy reading
        "BIGTIFF": "YES"
    })

    print(f"\n1. Stacking and aligning features to {OUTPUT_FEATURE}")
    if not os.path.exists(OUTPUT_FEATURE):
        with rasterio.open(OUTPUT_FEATURE, "w", **ref_meta) as dst:
            for i, b_path in enumerate(band_files, start=1):
                b_name = os.path.basename(b_path)
                print(f"  Processing [{i}/{len(band_files)}] {b_name}...")
                
                with rasterio.open(b_path) as src_b:
                    # Read data
                    data = src_b.read(1)
                    
                    # If shapes match exactly, just cast and write
                    if src_b.shape == (ref_height, ref_width) and src_b.transform == ref_transform:
                        data = data.astype(np.float32)
                        dst.write(data, i)
                    else:
                        # Resample needed
                        print(f"    Resampling {b_name} from {src_b.shape} to ({ref_height}, {ref_width})")
                        resampled_data = np.empty((ref_height, ref_width), dtype=np.float32)
                        reproject(
                            source=data,
                            destination=resampled_data,
                            src_transform=src_b.transform,
                            src_crs=src_b.crs,
                            dst_transform=ref_transform,
                            dst_crs=ref_crs,
                            resampling=Resampling.bilinear
                        )
                        dst.write(resampled_data, i)
    else:
        print(f"  {OUTPUT_FEATURE} already exists. Skipping feature stacking.")

    print("\n2. Generating mask from shapefile")
    if os.path.exists(SHAPEFILE_PATH):
        try:
            gdf = gpd.read_file(SHAPEFILE_PATH)
            # Ensure the CRS matches the reference raster
            if gdf.crs != ref_crs:
                print(f"  Reprojecting shapefile from {gdf.crs} to {ref_crs}")
                gdf = gdf.to_crs(ref_crs)
            
            # Rasterize
            mask_meta = ref_meta.copy()
            mask_meta.update({
                "count": 1,
                "dtype": "uint8",
                "nodata": 0
            })
            
            print(f"  Rasterizing to {OUTPUT_MASK}")
            mask_array = rasterize(
                [(geom, 1) for geom in gdf.geometry if geom is not None and geom.is_valid],
                out_shape=(ref_height, ref_width),
                transform=ref_transform,
                fill=0,
                dtype="uint8"
            )
            
            with rasterio.open(OUTPUT_MASK, "w", **mask_meta) as dst_mask:
                dst_mask.write(mask_array, 1)
            print("  Mask generation completed.")
            
        except Exception as e:
            print(f"  Failed to rasterize mask: {e}")
    else:
        print(f"  Shapefile not found at {SHAPEFILE_PATH}. Skipping mask generation.")

    print("\nAll done!")

if __name__ == '__main__':
    main()