import os
import numpy as np
import xarray as xr
import matplotlib as mpl
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.patches import Rectangle
from matplotlib.colors import ListedColormap

# --- Configuration ---
SCRIPTDIR = os.environ['SCRIPTDIR']
PREFIX = os.environ['PREFIX']
EXPDIR = os.environ['EXPDIR']
ntimestep, nlevel, num_levels = 18, 32, 40
show_nest_outline = True

# --- Filepaths ---
file_dom01_ml = os.path.join(EXPDIR, f"{PREFIX}_JW_DOM01_ML_0001.nc")
file_dom01_pl = os.path.join(EXPDIR, f"{PREFIX}_JW_DOM01_PL_0001.nc")
grid_file_n1 = os.path.join(EXPDIR, "icon_grid_0055_R02B05_N.nc")
grid_file_n2 = os.path.join(EXPDIR, "icon_grid_0056_R02B05_N.nc")

# --- Helper functions ---
def shift_lon(lon):
    lon = ((np.asarray(lon) + 180) % 360)
    return ((lon + 180) % 360) - 180

def read_ncl_colormap(filename, start=None, end=None):
    rgb = []
    with open(filename) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith((';', '#')) or 'ncolors=' in line: continue
            parts = line.split()
            if len(parts) >= 3 and all(p.isdigit() for p in parts[:3]):
                rgb.append([int(parts[0]), int(parts[1]), int(parts[2])])
    rgb = np.array(rgb) / 255.0
    if start is not None and end is not None: rgb = rgb[start:end+1]
    return ListedColormap(rgb)

def plot_nest(ax, lon_min, lon_max, lat_min, lat_max, **kw):
    if lon_max < lon_min:
        ax.add_patch(Rectangle((lon_min, lat_min), 180-lon_min, lat_max-lat_min, **kw))
        ax.add_patch(Rectangle((-180, lat_min), lon_max+180, lat_max-lat_min, **kw))
    else:
        ax.add_patch(Rectangle((lon_min, lat_min), lon_max-lon_min, lat_max-lat_min, **kw))

def shift_degrees(ds):
    for c in ['clon', 'clat']:
        ds[c] = shift_lon(np.rad2deg(ds[c])) if c == 'clon' else np.rad2deg(ds[c])
    return ds

# --- Data loading and preprocessing ---
ds_ml = xr.open_dataset(file_dom01_ml)
ds_pl = xr.open_dataset(file_dom01_pl)
ds_grid_n1 = xr.open_dataset(grid_file_n1)
ds_grid_n2 = xr.open_dataset(grid_file_n2)

for ds in [ds_ml, ds_pl, ds_grid_n1, ds_grid_n2]:
    shift_degrees(ds)

p = ds_ml.pres_sfc.isel(time=ntimestep) * 1e-2
temp = ds_pl.temp.isel(time=ntimestep, plev=2)
zeta_c = ds_pl.vor.isel(time=ntimestep, plev=2) * 1e5

elon_n1, elat_n1 = ds_grid_n1.clon, ds_grid_n1.clat
elon_n2, elat_n2 = ds_grid_n2.clon, ds_grid_n2.clat
max_elat_n1, min_elat_n1 = elat_n1.max(), elat_n1.min()
max_elon_n1, min_elon_n1 = elon_n1.max(), elon_n1.min()
max_elat_n2, min_elat_n2 = elat_n2.max(), elat_n2.min()
max_elon_n2, min_elon_n2 = elon_n2.max(), elon_n2.min()

# --- Plot setup ---
output_width_px, dpi = 700, 100
width_in, height_in = output_width_px / dpi, (output_width_px / dpi) * 0.7
mpl.rcParams['font.size'] = 6
projection = ccrs.PlateCarree(central_longitude=180)

titles = [
    "Surface Pressure (Northern Hemisphere)", "Surface Pressure (Southern Hemisphere)",
    "Vorticity (Northern Hemisphere)", "Vorticity (Southern Hemisphere)",
    "Temperature (Northern Hemisphere)", "Temperature (Southern Hemisphere)"
]
hemisphere_limits = [
    {'lon_min': 60, 'lon_max': 270, 'lat_min': 25, 'lat_max': 77},
    {'lon_min': 60, 'lon_max': 270, 'lat_min': -77, 'lat_max': -25},
]
p_levels = [[940, 1020, 10], [995, 1002, 0.5]]
vort_levels = [[-10, 35, 5], [-0.5, 1.75, 0.25]]
temp_levels = [220, 320, 2]
cmap_pressure = read_ncl_colormap(f"{SCRIPTDIR}/colormaps/JWw_PS_colormap.rgb")
cmap_vorticity = read_ncl_colormap(f"{SCRIPTDIR}/colormaps/BlWhRe.rgb", 37, 97)
cmap_temperature = read_ncl_colormap(f"{SCRIPTDIR}/colormaps/pressurecolors.rgb", 44, 209)
data_p, data_vort, data_temp = [p, p], [zeta_c, zeta_c], [temp, temp]

# --- Plotting ---
fig, axes = plt.subplots(3, 2, figsize=(width_in, height_in), subplot_kw={'projection': projection}, layout="tight")
for ax_row in axes:
    for ax in ax_row:
        ax.set_aspect('auto')

for n in [0, 1]:
    for i, (data, levels, cmap, varname) in enumerate(zip(
        [data_p[n], data_vort[n], data_temp[n]],
        [p_levels[n], vort_levels[n], temp_levels],
        [cmap_pressure, cmap_vorticity, cmap_temperature],
        [r'Surface Pressure (hPa)', r'Vorticity ($10^{-5} s^{-1}$)', 'Temperature (K)']
    )):
        ax = axes[i, n]
        ax.set_extent([hemisphere_limits[n]['lon_min'], hemisphere_limits[n]['lon_max'],
                       hemisphere_limits[n]['lat_min'], hemisphere_limits[n]['lat_max']], crs=projection)
        ax.coastlines(resolution='50m', linewidth=0.3)
        ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.3)
        clevs = np.arange(levels[0], levels[1] + levels[2], levels[2])
        cf = ax.tricontourf(ds_ml.clon, ds_ml.clat, data.values, levels=clevs, cmap=cmap, extend='both')
        if varname == "Surface Pressure (hPa)":
            ax.tricontour(ds_ml.clon, ds_ml.clat, data.values, levels=clevs, colors='k', linewidths=0.2)
        cbar = fig.colorbar(cf, ax=ax, orientation='horizontal', pad=0.15, aspect=40)
        cbar.set_label(varname)
        if show_nest_outline:
            plot_nest(ax, min_elon_n1, max_elon_n1, min_elat_n1, max_elat_n1, edgecolor='black', facecolor='none', linewidth=0.5, linestyle='--')
            plot_nest(ax, min_elon_n2, max_elon_n2, min_elat_n2, max_elat_n2, edgecolor='black', facecolor='none', linewidth=0.5, linestyle='--')
        ax.set_title(titles[i * 2 + n])

labels = ['a)', 'b)', 'c)', 'd)', 'e)', 'f)']
for ax, label in zip(axes.flat, labels):
    ax.annotate(label, xy=(0.01, 0.99), xycoords='axes fraction',
                fontsize='large', fontweight='bold', va='top', ha='left',
                bbox=dict(boxstyle='round,pad=0.2', facecolor='white', edgecolor='none'))


plt.tight_layout(pad=0.1)
fig.subplots_adjust(hspace=0.05, wspace=0.1)
output_file = os.path.join(SCRIPTDIR, f"../{PREFIX}_JW_R2B4N5_day9_pres_vort_temp.png")
plt.savefig(output_file, dpi=300, bbox_inches='tight')
print(f"Saved figure to {output_file}")
plt.show()