import numpy as np
import netCDF4 as nc
from tqdm import tqdm
#Convert NDVI data to a shape of (latitude, longitude, years, 24 months) 
def prepare_heatmap_data(ndvi_data):
    n_years = ndvi_data.shape[0] // 24
    n_lat, n_lon = ndvi_data.shape[1], ndvi_data.shape[2]

    heatmap_data = np.zeros((n_lat, n_lon, n_years, 24), dtype=np.float32)

    print("Preparing heatmap data...")
    for i in tqdm(range(n_years)):
        yearly_data = ndvi_data[i*24:(i+1)*24, :, :].transpose(1, 2, 0).astype(np.float32)
        heatmap_data[:, :, i, :] = yearly_data

    return heatmap_data
# Compute NDVI histograms adjusted to 20 bins 
def compute_histogram(heatmap_data, bins=20):
    n_lat, n_lon, n_years, _ = heatmap_data.shape
    histogram_data = np.zeros((n_lat, n_lon, n_years, bins), dtype=np.float32)

    print("Computing histograms with 20 bins...")
    for lat in tqdm(range(n_lat)):
        for lon in range(n_lon):
            for year in range(n_years):
                # Split the NDVI value range into 20 bins
                hist, bin_edges = np.histogram(heatmap_data[lat, lon, year, :], bins=bins, range=(0, 10000))
                histogram_data[lat, lon, year, :] = hist  # Store the count for the 20 bins

    return histogram_data
# Save the histogram data to a .npy file 
def save_histogram_data(histogram_data, filename='/NDVI_20bins_data.npy'):
    np.save(filename, histogram_data)
    print(f"Histogram data with 20 bins saved to {filename}")
# Main function to execute the entire process 
def main():
    file_path = '/NDVI_analy_raw.nc'
    dataset = nc.Dataset(file_path, 'r')
    ndvi_data = dataset.variables['ndvi'][:]

    heatmap_data = prepare_heatmap_data(ndvi_data)
    print("Heatmap data prepared. Shape:", heatmap_data.shape)

    histogram_data = compute_histogram(heatmap_data, bins=20)
    print("Histogram data computed with 20 bins.")

    save_histogram_data(histogram_data)

if __name__ == "__main__":
    main()