from netCDF4 import Dataset
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.mpl.ticker as cticker
import cartopy.feature as cfeature
import subprocess

# Open NetCDF
nclim = Dataset("CFSv2.prate.ENSMEAN.fcst.1991-2020_09.clim.nc")
ncurr = Dataset("CFSv2.prate.202509.ENSMEAN.fcst.varb.nc")

# Lon/Lat grid
Lon = np.linspace(0, 359, 360)
Lat = np.linspace(90, -90, 181)
Lon, Lat = np.meshgrid(Lon, Lat)

# Levels & colors
levels = [-15,-10, -6, -4, -2, -1, -0.5, -0.25, 0.25, 0.5, 1, 2, 4, 6, 10,15]

colors = [(0.47, 0.31, 0.25),
          (0.75, 0.0, 0.0),
         (1.0, 0.08, 0.00),
         (1.0, 0.38, 0.00),
         (1.0, 0.63, 0.00),
         (1.0, 0.75, 0.24),
         (1.0, 0.91, 0.48),
         (1.0, 1.00, 1.00),
         (0.78, 1.0, 0.75),
         (0.59, 0.96, 0.55),
         (0.47, 0.96, 0.45),
         (0.22, 0.82, 0.24),
         (0.12, 0.71, 0.12),
         (0.02, 0.59, 0.02),
         (0.12, 0.43, 0.92),
         (0.0, 0.0, 1.0)
       ]

#

#-------------------------------------------------------------------------
# Read lines from file
with open("./seed/months_file", "r") as f:
    lines = f.readlines()

# Store words in an array (list of lists)
words_array = [line.strip().split() for line in lines]
monic = words_array[0][0]
#-------------------------------------------------------------------------

# Loop over NetCDF indices (0 to 9)
for lead in range(0,9):   # 0 → 9
    fcm = words_array[0][lead+1]
    # Compute anomaly (scale by 86400)
    anm = (ncurr.variables["fcst"][lead] - nclim.variables["fcst"][lead]) * 86400

    # Plot with Cartopy
    fig, ax = plt.subplots(
        figsize=(10, 6),
        subplot_kw={"projection": ccrs.PlateCarree()}
    )

    # Map features
    ax.coastlines(resolution="110m")
    ax.add_feature(cfeature.BORDERS, linewidth=0.5)
    ax.add_feature(cfeature.LAND, facecolor="lightgray")
    ax.add_feature(cfeature.OCEAN, facecolor="lightblue")

    # Ticks
    ax.set_xticks(np.arange(-180, 181, 30), crs=ccrs.PlateCarree())
    ax.set_yticks(np.arange(-90, 91, 30), crs=ccrs.PlateCarree())
    lon_formatter = cticker.LongitudeFormatter()
    lat_formatter = cticker.LatitudeFormatter()
    ax.xaxis.set_major_formatter(lon_formatter)
    ax.yaxis.set_major_formatter(lat_formatter)
    ax.tick_params(axis='both', which='major', labelsize=10, direction='out')
#-------------------------------------
    gl = ax.gridlines(
        crs=ccrs.PlateCarree(),
        draw_labels=True,    # show tick labels
        linewidth=0.5,
        color='gray',
        alpha=0.5,
        linestyle='--'
    )

# Control which labels show up
    gl.top_labels = False
    gl.right_labels = False
    gl.left_labels = False
    gl.bottom_labels = False
#-------------------------------------

    # Filled contour
    contour = ax.contourf(
        Lon, Lat, anm,
        levels=levels,
        colors=colors,
        transform=ccrs.PlateCarree()
    )

    # Leadtime (index + 1)
    file_name  = "contour_lead" + str(lead) + "_global.png"

    # Title
    plt.title(f"CFSv2 Precip Rate Anom. (mm per day)  {monic}2025   {lead} Month Lead   {fcm} Fcst", fontsize=8)

    # Colorbar
#    cbar = fig.colorbar(contour, extend="both", shrink=0.9, aspect=40,
#                        orientation="horizontal", pad=0.1,
#                        label="Precipitation (mm/day)", ax=ax, ticks=levels)
#    
    cbar = fig.colorbar(contour, extend="both", shrink=0.9, aspect=40,
                        orientation="horizontal", pad=0.1,
                        ax=ax, ticks=levels)
    cbar.ax.xaxis.set_major_formatter(plt.FuncFormatter(lambda val, pos: f'{int(val)}'))
    # Suppress first and last tick labels

#
# --- Remove end ticks entirely ---
    ticks = cbar.get_ticks()
    cbar.set_ticks(ticks[1:-1])
# --- Custom tick formatter: integer if x < -0.5 or  x > 0.5 , else float
    def custom_format(x, pos):
        return f'{int(x)}' if x < -0.5 or  x > 0.5 else f'{x:.2f}'
    cbar.formatter = plt.FuncFormatter(custom_format)
    cbar.update_ticks()  # apply form
#
    # Save with unique filename
    plt.savefig(f'contour_lead{lead}_global.png')
    plt.close(fig)
    # Run a Linux command after saving
# Example: open the image using `xdg-open` (Linux)
    subprocess.run(["convert", file_name, "-trim", file_name ])

