~~11111111

import numpy as np
import geopandas as gpd
import os
import glob
from tqdm import tqdm
from multiprocessing import Pool

# Load the RGI.shp file
rgi_shp = gpd.read_file("/media/cjz/WD_BLACK/Modis_20236/fishnet.shp")

# Create a grid with 0.25x0.25 degree cells based on RGI.shp
cell_size = 0.25
grid = rgi_shp.copy()
grid['geometry'] = grid['geometry'].centroid
grid['geometry'] = grid['geometry'].apply(lambda x: x.buffer(cell_size / 2))
grid = grid.set_index('FID')

modis_years_folders = glob.glob('/media/cjz/WD_BLACK/project/shp/modis/*')
omi_path = '/media/cjz/WD_BLACK/project/shp/omi'

def process_month(year, modis_month_folder, grid):
    try:
        month = os.path.basename(modis_month_folder)
        omi_year_path = os.path.join(omi_path, year)
        omi_month_path = os.path.join(omi_year_path, month)

        omi_shp_file = glob.glob(os.path.join(omi_month_path, '*.shp'))
        modis_shp_file = glob.glob(os.path.join(modis_month_folder, '*.shp'))

        modis_shp = gpd.read_file(modis_shp_file[0])
        uva_shp = gpd.read_file(omi_shp_file[0])

        joined_data_modis = gpd.sjoin(modis_shp, grid, op='within')

        aggregated_modis = joined_data_modis.groupby('index_right').agg({
            'AOD': 'mean',
            'FMF': 'mean',
            'AE': 'mean'
        }).reset_index()

        joined_data_uva = gpd.sjoin(uva_shp, grid, op='within')
        aggregated_uva = joined_data_uva.groupby('index_right').agg({
            'UVAI': 'mean'
        }).reset_index()

        merged_data = aggregated_modis.merge(aggregated_uva, on='index_right', how='inner')

        grid = grid.merge(merged_data[['index_right', 'AOD', 'FMF', 'AE', 'UVAI']], left_index=True,
                          right_on='index_right', how='left')
        grid = grid[grid['AOD'] != np.nan]
        gdf = grid.dropna(subset=['AOD'])

        pbar = tqdm(range(gdf.shape[0]), desc=f"Year {year}, Month {month}")
        gdf['class'] = 0

        for k in pbar:
            if gdf['AE'].iloc[k] > 0.95:
                if gdf['UVAI'].iloc[k] > 0.8:
                    gdf.at[k, 'class'] = 'BC'
            elif gdf['AE'].iloc[k] < 0.7:
                gdf.at[k, 'class'] = 'MIX'
            else:
                if gdf['UVAI'].iloc[k] > 0.8:
                    gdf.at[k, 'class'] = 'DUST'
                else:
                    gdf.at[k, 'class'] = 'others'

        output_folder = os.path.join('/media/cjz/WD_BLACK/project/shp/merge/', year, month)
        os.makedirs(output_folder, exist_ok=True)
        output_file = os.path.join(output_folder, month + '.shp')
        gdf.to_file(output_file, driver='ESRI Shapefile')

    except Exception as e:
        print(f"Error processing Year {year}, Month {month}: {e}")

def main():
    modis_years_folders = glob.glob('/media/cjz/WD_BLACK/project/shp/modis/*')

    with Pool(processes=os.cpu_count()) as pool:
        for modis_year_folder in tqdm(modis_years_folders, desc="Processing Years"):
            year = os.path.basename(modis_year_folder)
            omi_year_path = os.path.join(omi_path, year)
            modis_month_folders = glob.glob(os.path.join(modis_year_folder, '*'))
            pool.starmap(process_month, [(year, modis_month_folder, grid) for modis_month_folder in modis_month_folders])

if __name__ == "__main__":
    main()