import arcpy
from arcpy import env
from arcpy.sa import *

# Set the environment
env.workspace = r"C:\path\to\your\shapefiles"  # Replace with your shapefile folder
shapefiles = arcpy.ListFeatureClasses()
env.outputCoordinateSystem = arcpy.SpatialReference("WGS 1984")

# Ensure the Spatial Analyst extension is enabled
arcpy.CheckOutExtension("Spatial")

# Set mask raster file path
mask_raster = r"C:\path\to\your\mask_raster.tif"  # Replace with your mask raster file path

# Get properties of the mask raster
mask_desc = arcpy.Describe(mask_raster)
cell_size = 0.0024998888

# Get extent of the mask raster and set it as the output extent
desc = arcpy.Describe(mask_raster)
env.extent = desc.extent

# Perform interpolation and masking for each shapefile
for shapefile in shapefiles:
    # Set the interpolation field, this should be a field name in your shapefile
    z_field = "pdsi"
    
    # Set the output raster paths
    idw_output_raster = r"C:\path\to\your\output\idw\{}_idw.tif".format(shapefile[:-4])  # Save IDW interpolation result
    mask_output_raster = r"C:\path\to\your\output\masked\{}.tif".format(shapefile[:-4])  # Save masked result
    
    # Perform IDW interpolation, setting the cell size
    outIDW = Idw(shapefile, z_field, cell_size=cell_size)
    
    # Save the IDW interpolation result
    outIDW.save(idw_output_raster)
    
    # Perform masking operation to match the properties of the mask raster
    outExtractByMask = ExtractByMask(outIDW, mask_raster)
    
    # Save the masked result
    outExtractByMask.save(mask_output_raster)

    print("IDW interpolation and masking completed for {}. IDW result saved as {}, masked result saved as {}.".format(shapefile, idw_output_raster, mask_output_raster))

print("Interpolation and masking completed for all shapefiles.")
