#!/usr/bin/env python
import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import cartopy.crs as ccrs
import cartopy.feature as cfeature

# ==== Configuration ====
filename = 'avhrr3_n18.20250501T090000Z.nc4'  # Replace with your file
selected_channel = 1  # Index of the channel to show on map
filename = 'atms_n21.20250501T090000Z.nc4'
selected_channel = 7  # Index of the channel to show on map
filename = 'IODA_nnr_042.MOD04_L2a.ocean.20190601_1200z.nc4'
snoO3elected_channel = 2  # Index of the channel to show on map
filename = 'swell-3dvar_atmos.iasi_metop-c.20231009T210000Z.nc4'
selected_channel = 261 #1643  # Index of the channel to show on map
selected_channel = 254 #1579  # Index of the channel to show on map
selected_channel = 253 #1574  # Index of the channel to show on map
selected_channel = 218
selected_channel = 230
selected_channel = 245
selected_channel = 256
selected_channel = 259
selected_channel = 266
selected_channel = 261

filename = 'atms_n21.20260125T150000Z.nc4'
selected_channel = 15

varname = 'aerosolOpticalDepth'
varname = 'brightnessTemperature'
# =======================
selected_channel = selected_channel - 1

# Open NetCDF file
nc = Dataset(filename, 'r')

# Read core data
lat = nc.groups['MetaData'].variables['latitude'][:]
lon = nc.groups['MetaData'].variables['longitude'][:]
qc = nc.groups['EffectiveQC0'].variables[varname][:]
channel_numbers = nc.groups['MetaData'].variables['sensorChannelNumber'][:]  # shape: (Channel,)
#channel_numbers = nc.groups['MetaData'].variables['obs_wavelength'][:]

# Fill values
fill_value_latlon = -3.368795e+38
fill_value_qc = -2147483643
fill_value_ch = -3.368795e+38

# Clean channel numbers
channel_mask = channel_numbers != fill_value_ch
channel_numbers = channel_numbers[channel_mask]
channels = len(channel_numbers)
#print (channel_numbers)

if selected_channel >= channels:
    raise ValueError(f"Invalid selected_channel={selected_channel}, must be between 0 and {channels - 1}")

# Apply lat/lon validity mask
valid_mask = (lat != fill_value_latlon) & (lon != fill_value_latlon)
lat = lat[valid_mask]
lon = lon[valid_mask]
qc = qc[valid_mask, :]  # apply same mask to QC

# ---- For map plot ----
zero_qc_lats = []
zero_qc_lons = []
nonzero_qc_lats = []
nonzero_qc_lons = []

ch_qc = qc[:, selected_channel]
ch_valid_mask = ch_qc != fill_value_qc
ch_qc_valid = ch_qc[ch_valid_mask]
lat_valid = lat[ch_valid_mask]
lon_valid = lon[ch_valid_mask]

is_zero = ch_qc_valid == 0
is_nonzero = ~is_zero

zero_qc_lats = lat_valid[is_zero]
zero_qc_lons = lon_valid[is_zero]
nonzero_qc_lats = lat_valid[is_nonzero]
nonzero_qc_lons = lon_valid[is_nonzero]

# ---- Bar chart counts for all channels ----
zero_qc_counts = []
nonzero_qc_counts = []

for ch in range(channels):
    ch_qc = qc[:, ch]
    ch_valid_mask = ch_qc != fill_value_qc
    ch_qc_valid = ch_qc[ch_valid_mask]

    zero_qc_counts.append(np.sum(ch_qc_valid == 0))
    nonzero_qc_counts.append(np.sum(ch_qc_valid != 0))

# ---- Plotting ----
fig = plt.figure(figsize=(12, 10))
gs = fig.add_gridspec(2, 1, height_ratios=[4, 1])
ax_map = fig.add_subplot(gs[0], projection=ccrs.PlateCarree())
ax_bar = fig.add_subplot(gs[1])

# Map
actual_channel_number = int(channel_numbers[selected_channel])
ax_map.set_title(f"Observation Locations - EffectiveQC0 (Channel {selected_channel+1} → #{actual_channel_number})")
ax_map.coastlines()
#ax_map.add_feature(cfeature.BORDERS, linestyle=':')
ax_map.add_feature(cfeature.LAND, edgecolor='black', facecolor='lightgray')
ax_map.add_feature(cfeature.OCEAN, facecolor='lightblue')
ax_map.gridlines(draw_labels=True, linestyle='--', alpha=0.5)

# Plot red first (QC≠0), then green (QC=0)
ax_map.scatter(nonzero_qc_lons, nonzero_qc_lats, s=5, color='red', alpha=0.3, label='QC≠0')
ax_map.scatter(zero_qc_lons, zero_qc_lats, s=5, color='green', alpha=0.7, label='QC=0', zorder=3)
ax_map.legend(loc='upper right')

# Bar chart with actual channel numbers
x = np.arange(channels)
width = 0.35
ax_bar.bar(x + width/2, nonzero_qc_counts, width, label='QC≠0', color='red')
ax_bar.bar(x - width/2, zero_qc_counts, width, label='QC=0', color='green')
ax_bar.set_xticks(x)
ax_bar.set_xticklabels([f"#{int(ch)}" for ch in channel_numbers])
ax_bar.set_ylabel("Count")
ax_bar.set_title("Count of QC=0 and QC≠0 per Channel Number")
ax_bar.legend()
ax_bar.grid(True, axis='y')

plt.tight_layout()
plt.show()
#plt.savefig('myplot.png',dpi=300,orientation='landscape',format='png')