source: trunk/LMDZ.COMMON/libf/evolution/deftank/visu_evol_layering.py @ 3958

#!/usr/bin/env python3
#######################################################################################################
### Python script to output stratification data over time from "restartpem#.nc" files               ###
### and to plot orbital parameters from "obl_ecc_lsp.asc"                                           ###
#######################################################################################################

import os
import sys
import numpy as np
from glob import glob
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from matplotlib.colors import LinearSegmentedColormap, LogNorm
from matplotlib.ticker import FuncFormatter
from scipy.interpolate import interp1d


def get_user_inputs():
    """
    Prompt the user for:
      - folder_path: directory containing NetCDF files (default: "starts")
      - base_name:   base filename (default: "restartpem")
      - infofile:    name of the PEM info file (default: "info_PEM.txt")
    Validates existence of folder and infofile before returning.
    """
    folder_path = input(
        "Enter the folder path containing the NetCDF files "
        "(press Enter for default [starts]): "
    ).strip() or "starts"
    while not os.path.isdir(folder_path):
        print(f"  » \"{folder_path}\" does not exist or is not a directory.")
        folder_path = input(
            "Enter a valid folder path (press Enter for default [starts]): "
        ).strip() or "starts"

    base_name = input(
        "Enter the base name of the NetCDF files "
        "(press Enter for default [restartpem]): "
    ).strip() or "restartpem"

    infofile = input(
        "Enter the name of the PEM info file "
        "(press Enter for default [info_PEM.txt]): "
    ).strip() or "info_PEM.txt"
    while not os.path.isfile(infofile):
        print(f"  » \"{infofile}\" does not exist or is not a file.")
        infofile = input(
            "Enter a valid PEM info filename (press Enter for default [info_PEM.txt]): "
        ).strip() or "info_PEM.txt"

    orbfile = input(
        "Enter the name of the orbital parameters ASCII file "
        "(press Enter for default [obl_ecc_lsp.asc]): "
    ).strip() or "obl_ecc_lsp.asc"
    while not os.path.isfile(orbfile):
        print(f"  » \"{orbfile}\" does not exist or is not a file.")
        orbfile = input(
            "Enter a valid orbital parameters ASCII filename (press Enter for default [obl_ecc_lsp.asc]): "
        ).strip() or "info_PEM.txt"

    return folder_path, base_name, infofile, orbfile


def list_netcdf_files(folder_path, base_name):
    """
    List and sort all NetCDF files matching the pattern {base_name}#.nc
    in folder_path. Returns a sorted list of full file paths.
    """
    pattern = os.path.join(folder_path, f"{base_name}[0-9]*.nc")
    all_files = glob(pattern)
    if not all_files:
        return []

    def extract_index(pathname):
        fname = os.path.basename(pathname)
        idx_str = fname[len(base_name):-3]
        return int(idx_str) if idx_str.isdigit() else float('inf')

    sorted_files = sorted(all_files, key=extract_index)
    return sorted_files


def open_sample_dataset(file_path):
    """
    Open a single NetCDF file and extract:
      - ngrid, nslope
      - longitude, latitude
    Returns (ngrid, nslope, longitude_array, latitude_array).
    """
    with Dataset(file_path, 'r') as ds:
        ngrid = ds.dimensions['physical_points'].size
        nslope = ds.dimensions['nslope'].size
        longitude = ds.variables['longitude'][:].copy()
        latitude = ds.variables['latitude'][:].copy()
    return ngrid, nslope, longitude, latitude


def collect_stratification_variables(files, base_name):
    """
    Scan all files to collect:
      - variable names for each stratification property
      - max number of strata (max_nb_str)
      - global min base elevation and max top elevation
    Returns:
      - var_info: dict mapping each property_name -> sorted list of var names
      - max_nb_str: int
      - min_base_elev: float
      - max_top_elev: float
    """
    max_nb_str = 0
    min_base_elev = np.inf
    max_top_elev = -np.inf

    property_markers = {
        'heights':   'stratif_slope',    # "..._top_elevation"
        'co2_ice':   'h_co2ice',
        'h2o_ice':   'h_h2oice',
        'dust':      'h_dust',
        'pore':      'h_pore',
        'pore_ice':  'poreice_volfrac'
    }
    var_info = {prop: set() for prop in property_markers}

    for file_path in files:
        with Dataset(file_path, 'r') as ds:
            if 'nb_str_max' in ds.dimensions:
                max_nb_str = max(max_nb_str, ds.dimensions['nb_str_max'].size)

            nslope = ds.dimensions['nslope'].size
            for k in range(1, nslope + 1):
                var_name = f"stratif_slope{k:02d}_top_elevation"
                if var_name in ds.variables:
                    arr = ds.variables[var_name][:]
                    min_base_elev = min(min_base_elev, np.min(arr))
                    max_top_elev = max(max_top_elev, np.max(arr))
                    var_info['heights'].add(var_name)

            for full_var in ds.variables:
                for prop, marker in property_markers.items():
                    if (marker in full_var) and prop != 'heights':
                        var_info[prop].add(full_var)

    for prop in var_info:
        var_info[prop] = sorted(var_info[prop])

    return var_info, max_nb_str, min_base_elev, max_top_elev


def load_full_datasets(files):
    """
    Open all NetCDF files and return a list of Dataset objects.
    (They should be closed by the caller after use.)
    """
    return [Dataset(fp, 'r') for fp in files]


def extract_stratification_data(datasets, var_info, ngrid, nslope, max_nb_str):
    """
    Build:
      - heights_data[t_idx][isl] = 2D array (ngrid, n_strata_current) of top_elevations.
      - raw_prop_arrays[prop] = 4D array (ngrid, ntime, nslope, max_nb_str) of per-strata values.
    Returns:
      - heights_data: list (ntime) of lists (nslope) of 2D arrays
      - raw_prop_arrays: dict mapping each property_name -> 4D array
      - ntime: number of time steps (files)
    """
    ntime = len(datasets)

    heights_data = [
        [None for _ in range(nslope)]
        for _ in range(ntime)
    ]
    for t_idx, ds in enumerate(datasets):
        for var_name in var_info['heights']:
            slope_idx = int(var_name.split("slope")[1].split("_")[0]) - 1
            if 0 <= slope_idx < nslope:
                raw = ds.variables[var_name][0, :, :]  # (n_strata, ngrid)
                heights_data[t_idx][slope_idx] = raw.T  # (ngrid, n_strata)

    raw_prop_arrays = {}
    for prop in var_info:
        if prop == 'heights':
            continue
        raw_prop_arrays[prop] = np.zeros((ngrid, ntime, nslope, max_nb_str), dtype=np.float32)

    def slope_index_from_var(vname):
        return int(vname.split("slope")[1].split("_")[0]) - 1

    for prop in raw_prop_arrays:
        slope_map = {}
        for vname in var_info[prop]:
            isl = slope_index_from_var(vname)
            if 0 <= isl < nslope:
                slope_map[isl] = vname

        arr = raw_prop_arrays[prop]
        for t_idx, ds in enumerate(datasets):
            for isl, var_name in slope_map.items():
                raw = ds.variables[var_name][0, :, :]  # (n_strata, ngrid)
                n_strata_current = raw.shape[0]
                arr[:, t_idx, isl, :n_strata_current] = raw.T

    return heights_data, raw_prop_arrays, ntime


def normalize_to_fractions(raw_prop_arrays):
    """
    Given raw_prop_arrays for 'co2_ice', 'h2o_ice', 'dust', 'pore' (in meters),
    normalize each set of strata so that the sum of those four = 1 per cell.
    Returns:
      - frac_arrays: dict mapping same keys -> 4D arrays of fractions (0..1).
    """
    co2 = raw_prop_arrays['co2_ice']
    h2o = raw_prop_arrays['h2o_ice']
    dust = raw_prop_arrays['dust']
    pore = raw_prop_arrays['pore']

    total = co2 + h2o + dust + pore
    mask = total > 0.0

    frac_co2 = np.zeros_like(co2, dtype=np.float32)
    frac_h2o = np.zeros_like(h2o, dtype=np.float32)
    frac_dust = np.zeros_like(dust, dtype=np.float32)
    frac_pore = np.zeros_like(pore, dtype=np.float32)

    frac_co2[mask] = co2[mask] / total[mask]
    frac_h2o[mask] = h2o[mask] / total[mask]
    frac_dust[mask] = dust[mask] / total[mask]
    frac_pore[mask] = pore[mask] / total[mask]

    return {
        'co2_ice': frac_co2,
        'h2o_ice': frac_h2o,
        'dust':     frac_dust,
        'pore':     frac_pore
    }


def read_infofile(file_name):
    """
    Reads "info_PEM.txt". Expects:
      - First line: parameters where the 3rd value is martian_to_earth conversion factor.
      - Each subsequent line: floats where first value is simulation timestamp (in Mars years).
    Returns:
      - date_time: 1D numpy array of timestamps (Mars years)
      - martian_to_earth: float conversion factor
    """
    date_time = []
    with open(file_name, 'r') as fp:
        first = fp.readline().split()
        martian_to_earth = float(first[2])
        for line in fp:
            parts = line.strip().split()
            if not parts:
                continue
            try:
                date_time.append(float(parts[0]))
            except ValueError:
                continue
    return np.array(date_time, dtype=np.float64), martian_to_earth


def get_yes_no_input(prompt: str) -> bool:
    """
    Prompt the user with a yes/no question. Returns True for yes, False for no.
    """
    while True:
        choice = input(f"{prompt} (y/n): ").strip().lower()
        if choice in ['y', 'yes']:
            return True
        elif choice in ['n', 'no']:
            return False
        else:
            print("Please respond with y or n.")


def prompt_discretization_step(max_top_elev):
    """
    Prompt for a positive float dz such that 0 < dz <= max_top_elev.
    """
    while True:
        entry = input(
            "Enter the discretization step of the reference grid for the elevation [m]: "
        ).strip()
        try:
            dz = float(entry)
            if dz <= 0:
                print("  » Discretization step must be strictly positive!")
                continue
            if dz > max_top_elev:
                print(
                    f"  » {dz:.3e} m is greater than the maximum top elevation "
                    f"({max_top_elev:.3e} m). Please enter a smaller value."
                )
                continue
            return dz
        except ValueError:
            print("  » Invalid numeric value. Please try again.")


def interpolate_data_on_refgrid(
    heights_data,
    prop_arrays,
    min_base_for_interp,
    max_top_elev,
    dz,
    exclude_sub=False
):
    """
    Build a reference elevation grid and interpolate strata fractions onto it.

    Returns:
      - ref_grid: 1D array of elevations (nz,)
      - gridded_data: dict mapping each property_name to 4D array
        (ngrid, ntime, nslope, nz) with interpolated fractions.
      - top_index: 3D array (ngrid, ntime, nslope) of ints:
        number of levels covered by the topmost stratum.
    """
    if exclude_sub and (dz > max_top_elev):
        ref_grid = np.array([0.0, max_top_elev], dtype=np.float32)
    else:
        ref_grid = np.arange(min_base_for_interp, max_top_elev + dz/2, dz)
    nz = len(ref_grid)
    print(f"> Number of reference grid points = {nz}")

    sample_prop = next(iter(prop_arrays.values()))
    ngrid, ntime, nslope, max_nb_str = sample_prop.shape

    gridded_data = {
        prop: np.full((ngrid, ntime, nslope, nz), -1.0, dtype=np.float32)
        for prop in prop_arrays
    }
    top_index = np.zeros((ngrid, ntime, nslope), dtype=np.int32)

    for ig in range(ngrid):
        for t_idx in range(ntime):
            for isl in range(nslope):
                h_mat = heights_data[t_idx][isl]
                if h_mat is None:
                    continue

                raw_h = h_mat[ig, :]
                h_all = np.full((max_nb_str,), np.nan, dtype=np.float32)
                n_strata_current = raw_h.shape[0]
                h_all[:n_strata_current] = raw_h

                if exclude_sub:
                    epsilon = 1e-6
                    valid_mask = (h_all >= -epsilon)
                else:
                    valid_mask = (~np.isnan(h_all)) & (h_all != 0.0)

                if not np.any(valid_mask):
                    continue

                h_valid = h_all[valid_mask]
                top_h = np.max(h_valid)
                i_zmax = np.searchsorted(ref_grid, top_h, side='right')
                top_index[ig, t_idx, isl] = i_zmax
                if i_zmax == 0:
                    continue

                for prop, arr in prop_arrays.items():
                    prop_profile_all = arr[ig, t_idx, isl, :]
                    prop_profile = prop_profile_all[valid_mask]
                    if prop_profile.size == 0:
                        continue

                    f_interp = interp1d(
                        h_valid,
                        prop_profile,
                        kind='next',
                        bounds_error=False,
                        fill_value=-1.0
                    )
                    gridded_data[prop][ig, t_idx, isl, :i_zmax] = f_interp(ref_grid[:i_zmax])

    return ref_grid, gridded_data, top_index


def attach_format_coord(ax, mat, x, y, is_pcolormesh=True):
    """
    Attach a format_coord function to the axes to display x, y, and value at cursor.
    Works for both pcolormesh and imshow style grids.
    """
    # Determine dimensions
    if mat.ndim == 2:
        ny, nx = mat.shape
    elif mat.ndim == 3 and mat.shape[2] in (3, 4):
        ny, nx, nc = mat.shape
    else:
        raise ValueError(f"Unsupported mat shape {mat.shape}")
    # Edges or extents
    if is_pcolormesh:
        xedges, yedges = x, y
    else:
        x0, x1 = x.min(), x.max()
        y0, y1 = y.min(), y.max()

    def format_coord(xp, yp):
        # Map to indices
        if is_pcolormesh:
            col = np.searchsorted(xedges, xp) - 1
            row = np.searchsorted(yedges, yp) - 1
        else:
            col = int((xp - x0) / (x1 - x0) * nx)
            row = int((yp - y0) / (y1 - y0) * ny)
        # Within bounds?
        if 0 <= row < ny and 0 <= col < nx:
            if mat.ndim == 2:
                v = mat[row, col]
                return f"x={xp:.3g}, y={yp:.3g}, val={v:.3g}"
            else:
                vals = mat[row, col]
                txt = ", ".join(f"{vv:.3g}" for vv in vals[:3])
                return f"x={xp:.3g}, y={yp:.3g}, val=({txt})"
        return f"x={xp:.3g}, y={yp:.3g}"

    ax.format_coord = format_coord


def plot_stratification_over_time(
    gridded_data,
    ref_grid,
    top_index,
    heights_data,
    date_time,
    exclude_sub=False,
    output_folder="."
):
    """
    For each grid point and slope, generate a 2×2 figure of:
      - CO2 ice fraction
      - H2O ice fraction
      - Dust fraction
      - Pore fraction
    """
    prop_names = ['co2_ice', 'h2o_ice', 'dust', 'pore']
    titles = ["CO2 ice", "H2O ice", "Dust", "Pore"]
    cmap = plt.get_cmap('turbo').copy()
    cmap.set_under('white')
    vmin, vmax = 0.0, 1.0

    sample_prop = next(iter(gridded_data.values()))
    ngrid, ntime, nslope, nz = sample_prop.shape

    if exclude_sub:
        positive_indices = np.where(ref_grid >= 0.0)[0]
        sub_ref_grid = ref_grid[positive_indices]
    else:
        positive_indices = np.arange(nz)
        sub_ref_grid = ref_grid

    for ig in range(ngrid):
        for isl in range(nslope):
            fig, axes = plt.subplots(2, 2, figsize=(10, 8))
            fig.suptitle(
                f"Content variation over time for (Grid point {ig+1}, Slope {isl+1})",
                fontsize=14,
                fontweight='bold'
            )

            # Precompute valid stratum tops per time
            valid_tops_per_time = []
            for t_idx in range(ntime):
                raw_h = heights_data[t_idx][isl][ig, :]
                h_all = raw_h[~np.isnan(raw_h)]
                if exclude_sub:
                    h_all = h_all[h_all >= 0.0]
                valid_tops_per_time.append(np.unique(h_all))

            for idx, prop in enumerate(prop_names):
                ax = axes.flat[idx]
                data_3d = gridded_data[prop][ig, :, isl, :]
                mat_full = data_3d.T
                mat = mat_full[positive_indices, :].copy()
                mat[mat < 0.0] = np.nan

                # Mask above top stratum
                for t_idx in range(ntime):
                    i_zmax = top_index[ig, t_idx, isl]
                    if i_zmax <= positive_indices[0]:
                        mat[:, t_idx] = np.nan
                    else:
                        count_z = np.count_nonzero(positive_indices < i_zmax)
                        mat[count_z:, t_idx] = np.nan

                im = ax.pcolormesh(
                    date_time,
                    sub_ref_grid,
                    mat,
                    cmap=cmap,
                    shading='auto',
                    vmin=vmin,
                    vmax=vmax
                )
                x_edges = np.concatenate([date_time, [date_time[-1] + (date_time[-1]-date_time[-2])]])
                attach_format_coord(ax, mat, x_edges, np.concatenate([sub_ref_grid, [sub_ref_grid[-1] + (sub_ref_grid[-1]-sub_ref_grid[-2])]]), is_pcolormesh=True)
                ax.set_title(titles[idx], fontsize=12)
                ax.set_xlabel("Time (Mars years)")
                ax.set_ylabel("Elevation (m)")

            fig.subplots_adjust(right=0.88)
            fig.tight_layout(rect=[0, 0, 0.88, 1.0])
            cbar_ax = fig.add_axes([0.90, 0.15, 0.02, 0.7])
            fig.colorbar(im, cax=cbar_ax, orientation='vertical', label="Content")

            fname = os.path.join(
                output_folder, f"layering_evolution_ig{ig+1}_is{isl+1}.png"
            )
            fig.savefig(fname, dpi=1200, bbox_inches='tight')


def plot_stratification_rgb_over_time(
    gridded_data,
    ref_grid,
    top_index,
    heights_data,
    date_time,
    exclude_sub=False,
    output_folder="."
):
    """
    Plot stratification over time colored using RGB ternary mix of H2O ice (blue), CO2 ice (violet), and dust (orange).
    Includes a triangular legend showing the mix proportions.
    """
    # Define constant colors
    violet = np.array([255,   0, 255], dtype=float) / 255
    blue   = np.array([  0,   0, 255], dtype=float) / 255
    orange = np.array([255, 165,   0], dtype=float) / 255

    # Elevation mask and array
    if exclude_sub:
        elevation_mask = (ref_grid >= 0.0)
        elev = ref_grid[elevation_mask]
    else:
        elevation_mask = np.ones_like(ref_grid, dtype=bool)
        elev = ref_grid

    # Pre-compute legend triangle
    res = 300
    u = np.linspace(0, 1, res)
    v = np.linspace(0, np.sqrt(3)/2, res)
    X, Y = np.meshgrid(u, v)
    V_bary = 2 * Y / np.sqrt(3)
    U_bary = X - 0.5 * V_bary
    W_bary = 1 - U_bary - V_bary
    mask_triangle = (U_bary >= 0) & (V_bary >= 0) & (W_bary >= 0)
    legend_rgb = (
        U_bary[..., None] * violet
        + V_bary[..., None] * orange
        + W_bary[..., None] * blue
    )
    legend_rgb = np.clip(legend_rgb, 0.0, 1.0)
    legend_rgba = np.zeros((res, res, 4))
    legend_rgba[..., :3] = legend_rgb
    legend_rgba[..., 3] = mask_triangle.astype(float)

    # Extract data arrays
    h2o = gridded_data['h2o_ice']
    co2 = gridded_data['co2_ice']
    dust = gridded_data['dust']
    ngrid, ntime, nslope, nz = h2o.shape

    # Fill missing depths
    ti = top_index.copy().astype(int)
    for ig in range(ngrid):
        for isl in range(nslope):
            for t in range(1, ntime):
                if ti[ig, t, isl] <= 0:
                    ti[ig, t, isl] = ti[ig, t-1, isl]

    # Loop over grid and slope
    for ig in range(ngrid):
        for isl in range(nslope):
            # Compute RGB stratification over time
            rgb = np.ones((nz, ntime, 3), dtype=float)
            frac_all = np.full((nz, ntime, 3), np.nan, dtype=float)  # store fH2O, fCO2, fDust
            for t in range(ntime):
                depth = ti[ig, t, isl]
                if depth <= 0:
                    continue
                cH2O = np.clip(h2o[ig, t, isl, :depth], 0, None)
                cCO2 = np.clip(co2[ig, t, isl, :depth], 0, None)
                cDust = np.clip(dust[ig, t, isl, :depth], 0, None)
                total = cH2O + cCO2 + cDust
                total[total == 0] = 1.0
                fH2O = cH2O / total
                fCO2 = cCO2 / total
                fDust = cDust / total
                frac_all[:depth, t, 0] = fH2O
                frac_all[:depth, t, 1] = fCO2
                frac_all[:depth, t, 2] = fDust
                rgb[:depth, t, 0] = fH2O * blue[0] + fCO2 * violet[0] + fDust * orange[0]
                rgb[:depth, t, 1] = fH2O * blue[1] + fCO2 * violet[1] + fDust * orange[1]
                rgb[:depth, t, 2] = fH2O * blue[2] + fCO2 * violet[2] + fDust * orange[2]

            # Mask elevation
            display_rgb = rgb[elevation_mask, :, :]
            display_frac = frac_all[elevation_mask, :, :]

            # Compute edges for pcolormesh
            dt = date_time[1] - date_time[0] if len(date_time) > 1 else 1
            x_edges = np.concatenate([date_time, [date_time[-1] + dt]])
            d_e = np.diff(elev)
            last_e = elev[-1] + (d_e[-1] if len(d_e)>0 else 1)
            y_edges = np.concatenate([elev, [last_e]])

            # Create figure with legend
            fig, (ax_main, ax_leg) = plt.subplots(
                1, 2, figsize=(8, 4), dpi=150,
                gridspec_kw={'width_ratios': [5, 1]}
            )

            # Main stratification panel
            mesh = ax_main.pcolormesh(
                x_edges,
                y_edges,
                display_rgb,
                shading='auto',
                edgecolors='none'
            )

            # Custom coordinate formatter: show time, elevation, and mixture fractions
            def main_format(x, y):
                # check bounds
                if x < x_edges[0] or x > x_edges[-1] or y < y_edges[0] or y > y_edges[-1]:
                    return ''
                # locate cell
                i = np.searchsorted(x_edges, x) - 1
                j = np.searchsorted(y_edges, y) - 1
                i = np.clip(i, 0, display_rgb.shape[1] - 1)
                j = np.clip(j, 0, display_rgb.shape[0] - 1)
                # get fractions
                fH2O, fCO2, fDust = display_frac[j, i]
                return f"Time={x:.2f}, Elev={y:.2f}, H2O={fH2O:.4f}, CO2={fCO2:.4f}, Dust={fDust:.4f}"
            ax_main.format_coord = main_format
            ax_main.set_facecolor('white')
            ax_main.set_title(f"Ternary mix over time (Grid point {ig+1}, Slope {isl+1})", fontweight='bold')
            ax_main.set_xlabel("Time (Mars years)")
            ax_main.set_ylabel("Elevation (m)")

            # Legend panel using proper edges
            u_edges = np.linspace(0, 1, res+1)
            v_edges = np.linspace(0, np.sqrt(3)/2, res+1)
            ax_leg.pcolormesh(
                u_edges,
                v_edges,
                legend_rgba,
                shading='auto',
                edgecolors='none'
            )
            ax_leg.set_aspect('equal')

            # Custom coordinate formatter for legend: show barycentric fractions
            def legend_format(x, y):
                # compute barycentric coords from cartesian (x,y)
                V = 2 * y / np.sqrt(3)
                U = x - 0.5 * V
                W = 1 - U - V
                if U >= 0 and V >= 0 and W >= 0:
                    return f"H2O: {W:.2f}, Dust: {V:.2f}, CO2: {U:.2f}"
                else:
                    return ''
            ax_leg.format_coord = legend_format

            # Draw triangle border and gridlines
            triangle = np.array([[0, 0], [1, 0], [0.5, np.sqrt(3)/2], [0, 0]])
            ax_leg.plot(triangle[:, 0], triangle[:, 1], 'k-', linewidth=1, clip_on=False, zorder=10)
            ticks = np.linspace(0.25, 0.75, 3)
            for f in ticks:
                ax_leg.plot([1 - f, 0.5 * (1 - f)], [0, (1 - f)*np.sqrt(3)/2], '--', color='k', linewidth=0.5, clip_on=False, zorder=9)
                ax_leg.plot([f, f + 0.5 * (1 - f)], [0, (1 - f)*np.sqrt(3)/2], '--', color='k', linewidth=0.5, clip_on=False, zorder=9)
                y = (np.sqrt(3)/2) * f
                ax_leg.plot([0.5 * f, 1 - 0.5 * f], [y, y], '--', color='k', linewidth=0.5, clip_on=False, zorder=9)

            # Legend labels
            ax_leg.text(0, -0.05, 'H2O ice', ha='center', va='top', fontsize=8)
            ax_leg.text(1, -0.05, 'CO2 ice', ha='center', va='top', fontsize=8)
            ax_leg.text(0.5, np.sqrt(3)/2 + 0.05, 'Dust', ha='center', va='bottom', fontsize=8)
            ax_leg.axis('off')

            # Save figure
            plt.tight_layout()
            fname = os.path.join(output_folder, f"layering_rgb_evolution_ig{ig+1}_is{isl+1}.png")
            fig.savefig(fname, dpi=1200, bbox_inches='tight')


def plot_dust_to_ice_ratio_over_time(
    gridded_data,
    ref_grid,
    top_index,
    heights_data,
    date_time,
    exclude_sub=False,
    output_folder="."
):
    """
    Plot the dust-to-ice ratio in the stratification over time,
    using a blue-to-orange colormap:
    - blue: ice-dominated (low dust-to-ice ratio)
    - orange: dust-dominated (high dust-to-ice ratio)
    """
    h2o = gridded_data['h2o_ice']
    co2 = gridded_data['co2_ice']
    dust = gridded_data['dust']
    ngrid, ntime, nslope, nz = h2o.shape

    # Define custom blue-to-orange colormap
    blue = np.array([0, 0, 255], dtype=float) / 255
    orange = np.array([255, 165, 0], dtype=float) / 255
    custom_cmap = LinearSegmentedColormap.from_list('BlueOrange', [blue, orange], N=256)

    # Log‑ratio bounds and small epsilon to avoid log(0)
    vmin, vmax = -2, 1
    epsilon = 1e-6

    # Loop over grids and slopes
    for ig in range(ngrid):
        for isl in range(nslope):
            ti = top_index[ig, :, isl].copy().astype(int)

            # Compute log10(dust/ice) profile at each time step
            log_ratio_array = np.full((nz, ntime), np.nan, dtype=np.float32)
            for t in range(ntime):
                if t > 0 and ti[t] <= 0:
                    ti[t] = ti[t-1]
                elif ti[t] <= 0:
                    continue
                zmax = ti[t]
                if zmax <= 0:
                    continue

                cH2O = np.clip(h2o[ig, t, isl, :zmax], 0, None)
                cCO2 = np.clip(co2[ig, t, isl, :zmax], 0, None)
                cDust = np.clip(dust[ig, t, isl, :zmax], 0, None)

                with np.errstate(divide='ignore', invalid='ignore'):
                    ratio = np.where(
                        cH2O > 0,
                        cDust / cH2O,
                        10**(vmax + 1)
                    )
                    log_ratio = np.log10(ratio + epsilon)
                    log_ratio = np.clip(log_ratio, vmin, vmax)

                log_ratio_array[:zmax, t] = log_ratio

            ratio_array = 10**log_ratio_array

            # Compute edges for pcolormesh
            x_edges = np.concatenate([date_time, [date_time[-1] + (date_time[-1] - date_time[-2])]]) * martian_to_earth
            y_edges = np.concatenate([ref_grid, [ref_grid[-1] + (ref_grid[-1] - ref_grid[-2])]])

            # Plot
            fig, ax = plt.subplots(figsize=(8, 6), dpi=150)
            im = ax.pcolormesh(
                date_time,
                elev,
                ratio_array,
                shading='auto',
                cmap='managua_r',
                norm=LogNorm(vmin=10**vmin, vmax=10**vmax),
            )
            attach_format_coord(ax, ratio_array, x_edges, y_edges, is_pcolormesh=True)
            ax.set_title(f"Dust-to-Ice ratio over time (Grid point {ig+1}, Slope {isl+1})", fontweight='bold')
            ax.set_xlabel('Time (Mars years)')
            ax.set_ylabel('Elevation (m)')

            # Add colorbar
            cbar = fig.colorbar(im, ax=ax, orientation='vertical', pad=0.15)
            cbar.set_label('Dust / H₂O ice (ratio)')
            cbar.set_ticks([1e-2, 1e-1, 1, 1e1])
            cbar.set_ticklabels(['1:100', '1:10', '1:1', '10:1'])

            # Save figure
            plt.tight_layout()
            outname = os.path.join(
                output_folder,
                f"dust_to_ice_ratio_grid{ig+1}_slope{isl+1}.png"
            )
            fig.savefig(outname, dpi=1200, bbox_inches='tight')


def plot_strata_count_and_total_height(heights_data, date_time, output_folder="."):
    """
    For each grid point and slope, plot:
      - Number of strata vs time
      - Total deposit height vs time
    """
    ntime = len(heights_data)
    nslope = len(heights_data[0])
    ngrid = heights_data[0][0].shape[0]

    for ig in range(ngrid):
        for isl in range(nslope):
            n_strata_t = np.zeros(ntime, dtype=int)
            total_height_t = np.zeros(ntime, dtype=float)

            for t_idx in range(ntime):
                h_mat = heights_data[t_idx][isl]
                raw_h = h_mat[ig, :]
                valid_mask = (~np.isnan(raw_h)) & (raw_h != 0.0)
                if np.any(valid_mask):
                    h_valid = raw_h[valid_mask]
                    n_strata_t[t_idx] = h_valid.size
                    total_height_t[t_idx] = np.max(h_valid)

            fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
            fig.suptitle(
                f"Strata count & total height over time for (Grid point {ig+1}, Slope {isl+1})",
                fontsize=14,
                fontweight='bold'
            )

            axes[0].plot(date_time, n_strata_t, marker='+', linestyle='-')
            axes[0].set_ylabel("Number of strata")
            axes[0].grid(True)

            axes[1].plot(date_time, total_height_t, marker='+', linestyle='-')
            axes[1].set_xlabel("Time (Mars years)")
            axes[1].set_ylabel("Total height (m)")
            axes[1].grid(True)

            fig.tight_layout(rect=[0, 0, 1, 0.95])
            fname = os.path.join(
                output_folder, f"strata_count_height_ig{ig+1}_is{isl+1}.png"
            )
            fig.savefig(fname, dpi=150)


def read_orbital_data(orb_file, martian_to_earth):
    """
    Read the .asc file containing obliquity, eccentricity and Ls p.
    Columns:
      0 = time in thousand Martian years
      1 = obliquity (deg)
      2 = eccentricity
      3 = Ls p (deg)
    Converts times to Earth years.
    """
    data = np.loadtxt(orb_file)
    dates_mka = data[:, 0]
    dates_yr = dates_mka * 1e3 / martian_to_earth
    obliquity = data[:, 1]
    eccentricity = data[:, 2]
    lsp = data[:, 3]
    return dates_yr, obliquity, eccentricity, lsp


def plot_orbital_parameters(infofile, orb_file, date_time, output_folder="."):
    """
    Plot the evolution of obliquity, eccentricity and Ls p
    versus simulated time.
    """
    # Read conversion factor from infofile
    _, martian_to_earth = read_infofile(infofile)

    # Read orbital data
    dates_yr, obl, ecc, lsp = read_orbital_data(orb_file, martian_to_earth)

    # Interpolate orbital parameters at simulation dates (date_time)
    obl_interp = interp1d(dates_yr, obl, kind='linear', bounds_error=False, fill_value="extrapolate")(date_time)
    ecc_interp = interp1d(dates_yr, ecc, kind='linear', bounds_error=False, fill_value="extrapolate")(date_time)
    lsp_interp = interp1d(dates_yr, lsp, kind='linear', bounds_error=False, fill_value="extrapolate")(date_time)

    # Plot
    fig, axes = plt.subplots(3, 1, figsize=(8, 10), sharex=True)
    fig.suptitle("Orbital parameters vs simulated time", fontsize=14, fontweight='bold')

    axes[0].plot(date_time, obl_interp, 'r-', marker='+')
    axes[0].set_ylabel("Obliquity (°)")
    axes[0].grid(True)

    axes[1].plot(date_time, ecc_interp, 'b-', marker='+')
    axes[1].set_ylabel("Eccentricity")
    axes[1].grid(True)

    axes[2].plot(date_time, lsp_interp, 'g-', marker='+')
    axes[2].set_ylabel("Ls of perihelion  (°)")
    axes[2].set_xlabel("Time (Mars years)")
    axes[2].grid(True)

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    fname = os.path.join(output_folder, "orbital_parameters_laskar.png")
    fig.savefig(fname, dpi=150)


def mars_ls(pday, peri_day, e_elips, year_day, lsperi=0.0):
    """
    Compute solar longitude (Ls) in radians for a given Mars date array 'pday'.
    Returns Ls in degrees [0, 360).
    """
    zz = (pday - peri_day) / year_day
    zanom = 2 * np.pi * (zz - np.round(zz))
    xref = np.abs(zanom)

    # Solve Kepler's equation via Newton–Raphson
    zx0 = xref + e_elips * np.sin(xref)
    for _ in range(10):
        f  = zx0 - e_elips * np.sin(zx0) - xref
        fp = 1 - e_elips * np.cos(zx0)
        dz = -f / fp
        zx0 += dz
        if np.all(np.abs(dz) <= 1e-7):
            break

    zx0 = np.where(zanom < 0, -zx0, zx0)
    zteta = 2 * np.arctan(
        np.sqrt((1 + e_elips) / (1 - e_elips)) * np.tan(zx0 / 2)
    )
    psollong = np.mod(zteta + lsperi, 2 * np.pi)

    return np.degrees(psollong)


def read_orbital_data_nc(starts_folder, infofile=None):
    """
    Read orbital parameters from restartfi_postPEM*.nc files in starts_folder.
    """
    if not os.path.isdir(starts_folder):
        raise ValueError(f"Invalid starts_folder '{starts_folder}': not a directory.")

    # Read simulation time mapping if provided
    if infofile:
        dates_yr, martian_to_earth = read_infofile(infofile)
    else:
        dates_yr = None

    pattern = os.path.join(starts_folder, "restartfi_postPEM*.nc")
    files = glob(pattern)
    if not files:
        raise FileNotFoundError(f"No NetCDF restart files found matching {pattern}")

    def extract_number(path):
        name = os.path.basename(path)
        prefix = 'restartfi_postPEM'
        if name.startswith(prefix) and name.endswith('.nc'):
            num_str = name[len(prefix):-3]
            if num_str.isdigit():
                return int(num_str)
        return float('inf')

    files = sorted(files, key=extract_number)

    all_year_day, all_peri, all_aphe, all_date_peri, all_obl = [], [], [], [], []
    for nc_path in files:
        with Dataset(nc_path, 'r') as nc:
            ctrl = nc.variables['controle'][:]
            all_year_day.append(ctrl[13])
            all_peri.append(ctrl[14])
            all_aphe.append(ctrl[15])
            all_date_peri.append(ctrl[16])
            all_obl.append(ctrl[17])

    year_day      = np.array(all_year_day)
    perihelion    = np.array(all_peri)
    aphelion      = np.array(all_aphe)
    date_peri_day = np.array(all_date_peri)
    obliquity     = np.array(all_obl)

    eccentricity  = (aphelion - perihelion) / (aphelion + perihelion)
    ls_perihelion = mars_ls(date_peri_day,0.,eccentricity,year_day)

    return dates_yr, obliquity, eccentricity, ls_perihelion, martian_to_earth


def plot_orbital_parameters_nc(starts_folder, infofile, date_time, output_folder="."):
    """
    Plot the evolution of obliquity, eccentricity and Ls p coming from simulation data
    versus simulated time, plus an additional figure of sin(eccentricity)*Lsp.
    versus simulated time.
    """
    # Read orbital data
    times_yr, obl, ecc, lsp, martian_to_earth = read_orbital_data_nc(starts_folder, infofile)

    fargs = dict(kind='linear', bounds_error=False, fill_value='extrapolate')
    obl_i = interp1d(times_yr, obl, **fargs)(date_time)
    ecc_i = interp1d(times_yr, ecc, **fargs)(date_time)
    lsp_i = interp1d(times_yr, lsp, **fargs)(date_time)

    date_time = date_time * martian_to_earth / 1e6

    fig, axes = plt.subplots(3,1, figsize=(8,10), sharex=True)
    fig.suptitle("Orbital parameters vs simulated time", fontsize=14, fontweight='bold')

    # Plot
    axes[0].plot(date_time, obl_i, 'r-', marker='+')
    axes[0].set_ylabel("Obliquity (°)")
    axes[0].grid(True)

    axes[1].plot(date_time, ecc_i, 'b-', marker='+')
    axes[1].set_ylabel("Eccentricity")
    axes[1].grid(True)

    axes[2].plot(date_time, lsp_i, 'g-', marker='+')
    axes[2].set_ylabel("Ls of perihelion (°)")
    axes[2].set_xlabel("Time (Myr)")
    axes[2].grid(True)

    plt.tight_layout(rect=[0,0,1,0.96])
    outname = os.path.join(output_folder, "orbital_parameters_simu.png")
    fig.savefig(outname, dpi=150)

    eps_sin_lsp = ecc_i * np.sin(np.radians(lsp_i)) 

    fig2, ax2 = plt.subplots(figsize=(8,5))
    fig2.suptitle(r"$\epsilon \times \sin(L_{sp})$", fontweight='bold')

    ax2.plot(date_time, eps_sin_lsp, 'm-', marker='+')
    ax2.set_ylabel(r"$\epsilon \cdot \sin(L_{sp})$")
    ax2.set_xlabel("Time (Myr)")
    ax2.grid(True)

    plt.tight_layout(rect=[0,0,1,0.95])
    outname2 = os.path.join(output_folder, "sin_ecc_times_Lsp.png")
    fig2.savefig(outname2, dpi=150)


def plot_dust_to_ice_ratio_with_obliquity(
    starts_folder,
    infofile,
    gridded_data,
    ref_grid,
    top_index,
    heights_data,
    date_time,
    exclude_sub=False,
    output_folder="."
):
    """
    Plot the dust-to-ice ratio over time as a heatmap, and overlay the evolution of
    obliquity on a secondary y-axis.
    """
    h2o = gridded_data['h2o_ice']
    co2 = gridded_data['co2_ice']
    dust = gridded_data['dust']
    ngrid, ntime, nslope, nz = h2o.shape

    # Read orbital data
    times_yr, obl, _, _, martian_to_earth = read_orbital_data_nc(starts_folder, infofile)
    fargs = dict(kind='linear', bounds_error=False, fill_value='extrapolate')
    obliquity = interp1d(times_yr, obl, **fargs)(date_time)

    # Define custom blue-to-orange colormap
    blue = np.array([0, 0, 255], dtype=float) / 255
    orange = np.array([255, 165, 0], dtype=float) / 255
    custom_cmap = LinearSegmentedColormap.from_list('BlueOrange', [blue, orange], N=256)
    color_map = { 1: 'green', -1: 'red', 0: 'orange' }

    # Log‑ratio bounds and small epsilon to avoid log(0)
    vmin, vmax = -2, 1
    epsilon = 1e-6

    # Loop over grids and slopes
    for ig in range(ngrid):
        for isl in range(nslope):
            # Compute total height time series
            total_height_t = np.zeros(ntime, dtype=float)
            for t_idx in range(ntime):
                h_mat = heights_data[t_idx][isl]
                raw_h = h_mat[ig, :]
                valid_mask = (~np.isnan(raw_h)) & (raw_h != 0.0)
                if np.any(valid_mask):
                    h_valid = raw_h[valid_mask]
                    total_height_t[t_idx] = np.max(h_valid)

            # Compute the per-interval sign of height change
            if ntime > 1:
                dh = np.diff(total_height_t)
                signs = np.sign(dh).astype(int)
            else:
                dh = np.array([], dtype=float)
                signs = np.array([], dtype=int)

            # Prepare fraction and ratio arrays
            ti = top_index[ig, :, isl].copy().astype(int)
            log_ratio_array = np.full((nz, ntime), np.nan, dtype=np.float32)
            frac_all = np.full((nz, ntime, 3), np.nan, dtype=float)  # store fH2O, fCO2, fDust
            for t in range(ntime):
                if t > 0 and ti[t] <= 0:
                    ti[t] = ti[t-1]
                elif ti[t] <= 0:
                    continue
                zmax = ti[t]
                if zmax <= 0:
                    continue

                cH2O = np.clip(h2o[ig, t, isl, :zmax], 0, None)
                cCO2 = np.clip(co2[ig, t, isl, :zmax], 0, None)
                cDust = np.clip(dust[ig, t, isl, :zmax], 0, None)
                total = cH2O + cCO2 + cDust
                total[total == 0] = 1.0
                fH2O = cH2O / total
                fCO2 = cCO2 / total
                fDust = cDust / total
                frac_all[:zmax, t, 0] = fH2O
                frac_all[:zmax, t, 1] = fCO2
                frac_all[:zmax, t, 2] = fDust

                with np.errstate(divide='ignore', invalid='ignore'):
                    ratio = np.where(cH2O > 0, cDust / cH2O, 10**(vmax + 1)
                    )
                    log_ratio = np.log10(ratio + epsilon)
                    log_ratio = np.clip(log_ratio, vmin, vmax)

                log_ratio_array[:zmax, t] = log_ratio

            ratio_array = 10**log_ratio_array

            # Compute edges for pcolormesh
            dt = date_time[1] - date_time[0] if len(date_time) > 1 else 1
            x_edges = np.concatenate([date_time, [date_time[-1] + (date_time[-1] - date_time[-2])]]) * martian_to_earth
            y_edges = np.concatenate([ref_grid, [ref_grid[-1] + (ref_grid[-1] - ref_grid[-2])]])

            # Plot
            fig, ax = plt.subplots(figsize=(8, 6), dpi=150)
            im = ax.pcolormesh(
                x_edges,
                y_edges,
                ratio_array,
                shading='auto',
                cmap='managua_r',
                norm=LogNorm(vmin=10**vmin, vmax=10**vmax),
            )

            # Custom formatter for millions of Earth years
            def millions_formatter(x, pos):
                return f"{x/1e6:.1f}"

            def format_coord_custom(x_input, y_input):
                # map onto the main axis
                if plt.gca() is ax2:
                    x_pix, y_pix = ax2.transData.transform((x_input, y_input))
                    x, y = ax.transData.inverted().transform((x_pix, y_pix))
                else:
                    x, y = x_input, y_input
                # check bounds
                if x < x_edges[0] or x > x_edges[-1] or y < y_edges[0] or y > y_edges[-1]:
                    return ''
                # locate cell
                i = np.searchsorted(x_edges, x) - 1
                j = np.searchsorted(y_edges, y) - 1
                i = np.clip(i, 0, ratio_array.shape[1] - 1)
                j = np.clip(j, 0, ratio_array.shape[0] - 1)
                # get fractions and obliquity
                fH2O, fCO2, fDust = frac_all[j, i]
                obl   = np.interp(x / martian_to_earth, date_time, obliquity)
                return f"Time={x:.2f}, Elev={y:.2f}, H2O={fH2O:.4f}, Dust={fDust:.4f}, Obl={obl:.2f}°"

            ax.set_title(f"Dust-to-Ice ratio over time (Grid point {ig+1}, Slope {isl+1})", fontweight='bold')
            ax.xaxis.set_major_formatter(FuncFormatter(millions_formatter))
            ax.set_xlabel('Time (Myr)')
            ax.set_ylabel('Elevation (m)')

            # Add colorbar
            cbar = fig.colorbar(im, ax=ax, orientation='vertical', pad=0.15)
            cbar.set_label('Dust / H₂O ice (ratio)')
            cbar.set_ticks([1e-2, 1e-1, 1, 1e1])
            cbar.set_ticklabels(['1:100', '1:10', '1:1', '10:1'])

            # Overlay obliquity on secondary y-axis
            ax2 = ax.twinx()
            for i in range(len(dh)):
                ax2.plot(
                    [date_time[i] * martian_to_earth, date_time[i+1] * martian_to_earth],
                    [obliquity[i], obliquity[i+1]],
                    color=color_map[signs[i]],
                    marker='+',
                    linewidth=1.5
                )
            ax2.format_coord = format_coord_custom
            ax2.set_ylabel('Obliquity (°)')
            ax2.tick_params(axis='y')
            ax2.grid(False)

            # Save figure
            plt.tight_layout()
            outname = os.path.join(
                output_folder,
                f'dust_ice_obliquity_grid{ig+1}_slope{isl+1}.png'
            )
            fig.savefig(outname, dpi=1200, bbox_inches='tight')


def main():
    # 1) Get user inputs
    folder_path, base_name, infofile, orbfile = get_user_inputs()

    # 2) List and verify NetCDF files
    files = list_netcdf_files(folder_path, base_name)
    if not files:
        print(f"No NetCDF files named \"{base_name}#.nc\" found in \"{folder_path}\".")
        sys.exit(1)
    print(f"> Found {len(files)} NetCDF file(s).")

    # 3) Open one sample to get grid dimensions & coordinates
    sample_file = files[0]
    ngrid, nslope, longitude, latitude = open_sample_dataset(sample_file)
    print(f"> ngrid  = {ngrid}, nslope = {nslope}")

    # 4) Collect variable info + global min/max elevations
    var_info, max_nb_str, min_base_elev, max_top_elev = collect_stratification_variables(files, base_name)
    print(f"> max strata per slope = {max_nb_str}")
    print(f"> min base elev = {min_base_elev:.3f} m, max top elev = {max_top_elev:.3f} m")

    # 5) Load full datasets
    datasets = load_full_datasets(files)

    # 6) Extract stratification data
    heights_data, raw_prop_arrays, ntime = extract_stratification_data(datasets, var_info, ngrid, nslope, max_nb_str)

    # 7) Close datasets
    for ds in datasets:
        ds.close()

    # 8) Normalize to fractions
    frac_arrays = normalize_to_fractions(raw_prop_arrays)

    # 9) Ask whether to include subsurface
    show_subsurface = get_yes_no_input("Show subsurface layers?")
    exclude_sub = not show_subsurface
    if exclude_sub:
        min_base_for_interp = 0.0
        print("> Interpolating only elevations >= 0 m (surface strata).")
    else:
        min_base_for_interp = min_base_elev
        print(f"> Interpolating full depth down to {min_base_elev:.3f} m.")

    # 10) Prompt discretization step
    dz = prompt_discretization_step(max_top_elev)

    # 11) Build reference grid and interpolate
    ref_grid, gridded_data, top_index = interpolate_data_on_refgrid(
        heights_data, frac_arrays, min_base_for_interp, max_top_elev, dz, exclude_sub=exclude_sub
    )

    # 12) Read timestamps and conversion factor from infofile
    date_time, martian_to_earth = read_infofile(infofile)
    if date_time.size != ntime:
        print(f"Warning: {date_time.size} timestamps vs {ntime} NetCDF files.")

    # 13) Plot stratification data over time
    plot_stratification_over_time(
        gridded_data, ref_grid, top_index, heights_data, date_time,
        exclude_sub=exclude_sub, output_folder="."
    )
    plot_stratification_rgb_over_time(
        gridded_data, ref_grid, top_index, heights_data, date_time,
        exclude_sub=exclude_sub, output_folder="."
    )
    #plot_dust_to_ice_ratio_over_time(
    #    gridded_data, ref_grid, top_index, heights_data, date_time,
    #    exclude_sub=exclude_sub, output_folder="."
    #)
    plot_dust_to_ice_ratio_with_obliquity(
        folder_path, infofile,
        gridded_data, ref_grid, top_index, heights_data, date_time,
        exclude_sub=exclude_sub, output_folder="."
    )
    #plot_strata_count_and_total_height(heights_data, date_time, output_folder=".")

    # 14) Plot orbital parameters
    #plot_orbital_parameters(infofile, orbfile, date_time, output_folder=".")
    plot_orbital_parameters_nc(folder_path, infofile, date_time, output_folder=".")

    # 15) Show all figures
    plt.show()


if __name__ == "__main__":
    main()