import sys
import pandas as pd
import numpy as np  
import re
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

import matplotlib
matplotlib.use("Agg")  

datadir = "../data/constraining_data/"
rcmip_output_dir = "../RCMIP3_input_datafiles/"


def define_and_print_uncertainty_ranges(version="v2.0.0"):
    constraint_w_unc = {
        "Global Mean Surface Temperature (GMST)": ["IGCC2024", "degC", 1850, 1900, 2014, 2023, 1.0, 1.19, 1.4, "90th percentile"],
        "Ocean Heat Content|Global|Total": ["IGCC2024", "ZJ", 1971, 1971, 2020, 2020, 435.1049033  - 72.5,  435.1049033 , 438.804903 + 72.5, "1 sigma"],
        "Atmospheric Concentrations|CO2": ["NOA_GML", "ppm", np.nan, np.nan, 2014, 2023, 407.049, 408.649, 410.249, "expert judgment 1.26"],
        "Carbon Flux to Oceans": ["GCB2024_updated","Pg/y", np.nan, np.nan, 2014, 2023, 2.6, 3.1, 3.6, "1 sigma"],
        "Carbon Flux to Land": ["GCB2024_updated", "Pg/y", np.nan, np.nan, 2014, 2023, 1.8, 2.7, 3.6, "1 sigma"],
        "Effective Radiative Forcing|Aerosols": ["AR6", "W/m2", 1750, 1750, 2005, 2014, -2.0, -1.3, -0.6, "90th percentile"]
    }
    rows = []
    for key, item in constraint_w_unc.items():
        print(f"{key}: Source={item[0]}, Baseline years=({item[1]}, {item[2]}), Recent years=({item[3]}, {item[4]}), Value={item[5]}, Lower bound={item[6]}, Upper bound={item[7]}")
        row = pd.DataFrame(
            {
                "Model": "Global Carbon Budget",
                "Scenario": "historical",
                "Region": "World",
                "Variable": key,
                "Unit": item[1],
                "Activity_Id": "Constraining Data",
                "Source": item[0],
                "Version": f"RCMIP Phase 3 {version}",
                "Baseline_period": f"{item[2]}-{item[3]}",
                "Constraint_period": f"{item[4]}-{item[5]}",
                "Lower": item[6],
                "Central_estimate": item[7],
                "Upper": item[8],
                "Uncertainty_type": item[9]
            },
            index=[0],
        )
        rows.append(row)
    constraints_df = pd.concat(rows, ignore_index=True)
    for idx, row in constraints_df.iterrows():
        if row['Uncertainty_type'].lower() == '1 sigma':
            constraints_df.loc[idx, 'Lower_bound'] = row['Central_estimate'] - 1.645 * (row['Central_estimate'] - row['Lower'])
            constraints_df.loc[idx, 'Upper_bound'] = row['Central_estimate'] + 1.645 * (row['Upper'] - row['Central_estimate'])
        else:
            constraints_df.loc[idx, 'Lower_bound'] = row['Lower']
            constraints_df.loc[idx, 'Upper_bound'] = row['Upper']
    constraints_df.drop(columns=['Uncertainty_type', "Lower", "Upper"], inplace=True)
    return constraints_df

def plot_timeseries_constraints(df, constraints_df=None, figsize=None, save_path=None, show=False):
    """Plot one timeseries subplot per row in a constraints DataFrame.

    Each row is expected to contain metadata columns including ``Variable``,
    ``Source`` and ``Unit`` and year columns named as four-digit strings.
    The function plots only the contiguous non-NaN section of each row's
    timeseries (years with any non-NaN data), uses the ``Variable`` value as
    the subplot title with the ``Source`` in parentheses, and places the
    ``Unit`` and a short variable shorthand on the y-axis label. Years are
    shown on the x-axis.

    Parameters
    ----------
    df : pandas.DataFrame
        DataFrame where each row is a timeseries with year columns and
        metadata columns ``Variable``, ``Source``, ``Unit``. Year columns
        should be string names like '1850', '2020', etc.
    constraints_df : pandas.DataFrame or None, optional
        Optional DataFrame containing constraint metadata. Expected columns
        include ``Variable``, ``Central_estimate``, ``Lower_bound``,
        ``Upper_bound``, ``Uncertainty_type`` and ``Constraint_period``
        (e.g. '2014-2023'). When provided the function will draw a
        semi-transparent rectangle spanning the constraint period on the x-\
        axis and the (possibly scaled) lower/upper bounds on the y-axis.
    figsize : tuple or None, optional
        Matplotlib figure size (width, height). If ``None`` a reasonable
        default is chosen based on number of rows.
    save_path : str or None, optional
        If provided, the plot is saved to this path.
    show : bool, optional
        If True, call ``plt.show()`` after drawing.

    Returns
    -------
    matplotlib.figure.Figure
        The created figure object.
    """
    if df is None or df.empty:
        raise ValueError("Input dataframe is empty or None")
    # Find year columns: accept ints (e.g. 1850) or 4-digit string names
    year_cols = [c for c in df.columns if (isinstance(c, np.int64) and 1000 <= c <= 9999) or (isinstance(c, str) and re.fullmatch(r"\d{4}", c))]
    if not year_cols:
        raise ValueError("No 4-digit year columns found in dataframe")

    df_tothc = df[df['Variable'] == 'Total Heat Content'].copy()
    df = df.drop(df[df['Variable'] == 'Total Heat Content'].index)
    # Determine figure size
    n = len(df)
    if figsize is None:
        # height per subplot ~1.5 inches, width default 8
        figsize = (30, max(3, int(1.5 * n //2)))

    fig, axes = plt.subplots(nrows=n//2, ncols=2, figsize=figsize, constrained_layout=True)
    if n == 1:
        axes = [axes]
    else:
        axes = axes.flatten()

    # Add subplot letters a), b), c), ...
    labels = [f"{chr(97+i)})" for i in range(len(axes))]
    for ax, label in zip(axes, labels):
        ax.text(
            0.02, 0.95, label, transform=ax.transAxes,
            fontsize=14, fontweight='bold', va='top', ha='left'
        )

    shorthand = {
        "Global Mean Surface Temperature (GMST)": "GMST",
        "Ocean Heat Content|Total": "Ocean Heat Content",
        "Atmospheric Concentrations|CO2": "CO$_2$ atmospheric concentration",
        "Carbon Flux to Oceans": "Carbon Flux To Oceans",
        "Carbon Flux to Land": "Carbon Flux to Land",
        "Effective Radiative Forcing|Aerosols": "Aerosol ERF",
    }

    for i, (_, row) in enumerate(df.iterrows()):
        if i < n:
            ax = axes[i]
        else:
            ax = axes[-1]
        # extract metadata
        var = row.get('Variable', '')
        unit = row.get('Unit', '')

        # select year values and drop NaNs
        ts = row[[y for y in year_cols]].dropna()

        # ensure index are ints for plotting order
        years = [int(y) for y in ts.index]
        values = ts.values.astype(float)
        ax.plot(years, values, '-o', markersize=3, label=shorthand.get(var, var.split('|')[-1]))
        # If constraints provided, try to draw them for this variable
        # exact match on Variable column (fall back to substring match)
     
        cons_row = constraints_df[constraints_df['Variable'] == var]
        if cons_row.empty:
            cons_row = constraints_df[constraints_df['Variable'].str.contains(var.split('|')[-1], na=False)]
        if not cons_row.empty:
            cons = cons_row.iloc[0]
            central = cons.get('Central_estimate')
            lower = cons.get('Lower_bound')
            upper = cons.get('Upper_bound')
            unc_type = str(cons.get('Uncertainty_type', '')).strip()
            period = cons.get('Constraint_period')
            # Apply 1-sigma scaling if requested
            try:
                central = float(central)
                lower = float(lower)
                upper = float(upper)
                if unc_type.lower() == '1 sigma' or '1 sigma' in unc_type.lower():
                    # scale distances from central by 1.645
                    d_low = central - lower
                    d_high = upper - central
                    lower = central - 1.645 * d_low
                    upper = central + 1.645 * d_high
            except Exception:
                # if conversion fails, skip drawing
                central = lower = upper = None

            if period and isinstance(period, str) and '-' in period and central is not None:
                try:
                    p0, p1 = period.split('-', 1)
                    x0 = int(p0)
                    x1 = int(p1)
                    height = upper - lower
                    if x0 < x1:
                        # draw semi-transparent rectangle spanning x0..x1 and lower..upper

                        rect = Rectangle((x0, lower), x1 - x0, height, color='gray', alpha=0.25, zorder=0)
                        ax.add_patch(rect)
                        # draw central estimate as a thicker line
                        ax.hlines(central, x0, x1, colors='k', linewidth=1.2, zorder=5)
                    else:
                        # single-year constraint: draw a vertical error bar at x0 with central point
                        lower_f = float(lower)
                        upper_f = float(upper)
                        central_f = float(central)
                        # vertical line spanning lower..upper
                        ax.vlines(x0, lower_f, upper_f, colors='gray', linewidth=2, alpha=0.5, zorder=1)
                        # central point with caps using errorbar (yerr expects [[low],[high]])
                        low_err = central_f - lower_f
                        high_err = upper_f - central_f
                        ax.errorbar([x0], [central_f], yerr=np.array([[low_err], [high_err]]),
                                    fmt='o', color='k', ecolor='k', elinewidth=1.5, capsize=6, zorder=5)
                except Exception:
                    pass
        if var == "Ocean Heat Content|Total":

            # select year values and drop NaNs
            ts_thc = df_tothc.iloc[0,:][[y for y in year_cols]].dropna()
            # ensure index are ints for plotting order
            years_thc = [int(y) for y in ts_thc.index]
            values = ts_thc.values.astype(float)
            ax.plot(years_thc, values, '-o', markersize=3, label="Total Heat Content")
            ax.legend(fontsize=12)
            ax.set_title(f"Heat Content", fontsize=16)
        else:
            ax.set_title(f"{var}", fontsize=16)
        # ax.set_title(f"{chr(97 + i)})", loc='left')
        ax.set_ylabel(f"{shorthand.get(var, var.split('|')[-1])} ({unit})", fontsize=15)
        ax.set_xlabel('Year', fontsize=15)
        ax.tick_params(axis='both', labelsize=15)
        ax.grid(True, linestyle=':', linewidth=0.5)
        if years[0] < 1950:
            ax.set_xlim(left=1950)
        ax.set_xlim(right=2025)
    if save_path:
        fig.savefig(save_path, dpi=150)
    plt.close(fig)
    return fig

def read_noaa_gml_ml_means(timeres, version="v2.0.0"):
    """
    Reads NOAA GML mean CO2 data for a given time resolution.

    Parameters
    ----------
    timeres : str
        Time resolution of the data to read. Must be either "year" or "month".

    Returns
    -------
    np.ndarray
        Transposed array of selected columns containing time and mean CO2 values.

    Raises
    ------
    KeyError
        If `timeres` is not one of the supported keys ("year", "month").
    FileNotFoundError
        If the corresponding CSV file does not exist.
    """
    timeres_dict = {
        "year": ["annmean", 37, ["year", "mean"]],
        "month": ["mm", 38, ["decimal", "average"]],
    }
    data = pd.read_csv(
        f"{datadir}/noa_gml_co2.csv",
        skiprows=timeres_dict[timeres][1],
    )
    # print(data.tail())
    # print(data["mean"].values[-10:])
    # print(data["unc"].values[-10:])
    # print(np.mean(data["mean"].values[-10:]))
    #sys.exit(4)
    conc_row = pd.DataFrame(
        {
            "Model": "Observed",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Atmospheric Concentrations|CO2",
            "Unit": "ppm",
            "Activity_Id": "Constraining Data",
            "Source": "NOAA GML",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                data["year"].values[i] : data["mean"].values[i]
                for i in range(len(data["year"].values))
            },
        },
        index=[0],
    )
    return conc_row


def read_gcb_data(version="v2.0.0", data_type = "csv"):
    """
    Reads and processes the Global Carbon Budget data from an Excel file.
    Loads the "Historical Budget" sheet from the Global_Carbon_Budget_2024_v1.02.xlsx file,
    starting from row 16, fills missing values with 0.0, and computes additional columns:
    "emissions_tot" (total emissions) and "fossil emissions" (fossil emissions including cement carbonation sink).

    Returns
    -------
    pandas.DataFrame
        A DataFrame containing the processed Global Carbon Budget data with additional columns:
        "emissions_tot" and "fossil emissions".
    """
    if data_type == "csv":
        data_csv = pd.read_csv(f"{datadir}/Global_Carbon_Budget_2024_updated_methodology.csv")
        idx_to_drop = []
        for idx, row in data_csv.iterrows():
            if pd.isna(row["corrected land sink"]):
                idx_to_drop.append(idx)
            if not re.fullmatch(r"\d{4}", str(row["Year"])):
                idx_to_drop.append(idx)
        data_csv.drop(index=idx_to_drop, inplace=True)
        fco2_land = pd.DataFrame(
            {
                "Model": "Global Carbon Budget",
                "Scenario": "historical",
                "Region": "World",
                "Variable": "Carbon Flux to Land",
                "Unit": "PgC/yr",
                "Activity_Id": "Constraining Data",
                "Source": "GCB2024_updated",
                "Version": f"RCMIP Phase 3 {version}",
                **{
                    data_csv["Year"].values[i] : data_csv["corrected land sink"].values[i]
                    for i in range(len(data_csv["Year"].values))
                },
            },
            index=[0],
        )
        return fco2_land
    data_xls = pd.read_excel(
        f"{datadir}/Global_Carbon_Budget_2024_v1.02.xlsx",
        sheet_name="Historical Budget",
        skiprows=15,
    ).fillna(0.0)
    fco2_land = pd.DataFrame(
        {
            "Model": "Global Carbon Budget",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Carbon Flux to Land",
            "Unit": "PgC/yr",
            "Activity_Id": "Constraining Data",
            "Source": "GCB2024_updated",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                data_xls["Year"].values[i] : data_xls["land sink"].values[i]
                for i in range(len(data_xls["Year"].values))
            },
        },
        index=[0],
    )
    return fco2_land


def read_gcb_ocean_carbon_data(version="v2.0.0", data_type = "csv"):
    """
    Reads ocean carbon sink data from the Global Carbon Budget Excel file.
    This function loads data from the "Ocean Sink" sheet of the
    'Global_Carbon_Budget_2024_v1.02.xlsx' Excel file, skipping the first 31 rows
    and reading the next 65 rows. It fills missing values with 0.0 and removes
    columns with names starting with 'Unnamed'.

    Returns
    -------
    pandas.DataFrame
        A DataFrame containing the processed ocean carbon sink data.
    """
    if data_type == "csv":
        data_csv = pd.read_csv(f"{datadir}/Global_Carbon_Budget_2024_updated_methodology.csv")
        idx_to_drop = []
        for idx, row in data_csv.iterrows():
            if pd.isna(row["corrected ocean sink"]):
                idx_to_drop.append(idx)
            if not re.fullmatch(r"\d{4}", str(row["Year"])):
                idx_to_drop.append(idx)
        data_csv.drop(index=idx_to_drop, inplace=True)
        fco2_ocean = pd.DataFrame(
            {
                "Model": "Global Carbon Budget",
                "Scenario": "historical",
                "Region": "World",
                "Variable": "Carbon Flux to Oceans",
                "Unit": "PgC/yr",
                "Activity_Id": "Constraining Data",
                "Source": "GCB2024_updated",
                "Version": f"RCMIP Phase 3 {version}",
                **{
                    data_csv["Year"].values[i] : data_csv["corrected ocean sink"].values[i]
                    for i in range(len(data_csv["Year"].values))
                },
            },
            index=[0],
        )
        return fco2_ocean
    data_xls = pd.read_excel(
        f"{datadir}/Global_Carbon_Budget_2024_v1.02.xlsx",
        sheet_name="Ocean Sink",
        skiprows=31,
        nrows=65,
    ).fillna(0.0)
    data_csv = pd.read_csv(f"{datadir}/Global_Carbon_Budget_2024_updated_methodology.csv")
    fco2_ocean = pd.DataFrame(
        {
            "Model": "Global Carbon Budget",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Carbon Flux to Oceans",
            "Unit": "PgC/yr",
            "Activity_Id": "Constraining Data",
            "Source": "GCB2024_updated",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                data_xls["year"].values[i] : data_xls["GCB"].values[i]
                for i in range(len(data_xls["year"].values))
            },
        },
        index=[0],
    )   
    return fco2_ocean

def read_temperature_data(version="v2.0.0"):
    """
    Reads annual average temperature data from a CSV file.

    Returns
    -------
    pandas.DataFrame
        A DataFrame containing the annual average temperature data.
    """
    temp_data = pd.read_csv(f"{datadir}igcc_temperature_timeseries.csv")
    temperature_row = pd.DataFrame(
        {
            "Model": "Observed",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Global Mean Surface Temperature (GMST)",
            "Unit": "°C",
            "Activity_Id": "Constraining Data",
            "Source": "IGCC2024",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                temp_data["timebound_lower"].values[i] : temp_data["GMST"].values[i]
                for i in range(len(temp_data["timebound_lower"].values))
            },
        },
        index=[0],
    )
    #print(temperature_row)
    return temperature_row

def read_aer_ts(version="v2.0.0"):
    """
    Reads annual average temperature data from a CSV file.

    Returns
    -------
    pandas.DataFrame
        A DataFrame containing the annual average temperature data.
    """
    aer_data = pd.read_csv(f"{datadir}igcc_erf_best.csv")
    #print(aer_data.columns)
    aer_row = pd.DataFrame(
        {
            "Model": "Observed",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Effective Radiative Forcing|Aerosols",
            "Unit": "W/m2",
            "Activity_Id": "Constraining Data",
            "Source": "IGCC2024",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                aer_data["timebound_lower"].values[i] : aer_data["aerosol-radiation_interactions"].values[i] + aer_data["aerosol-cloud_interactions"].values[i]
                for i in range(100,len(aer_data["timebound_lower"].values))
            },
        },
        index=[0],
    )
    return aer_row

def read_ohc(version="v2.0.0"):
    """
    Reads annual average temperature data from a CSV file.

    Returns
    -------
    pandas.DataFrame
        A DataFrame containing the annual average temperature data.
    """
    ohc_data = pd.read_csv(
        f"{datadir}earth_energy_imbalance.csv"
    )
    # print(ohc_data.head())
    # print(ohc_data.tail())
    ohc_row = pd.DataFrame(
        {
            "Model": "Observed",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Ocean Heat Content|Total",
            "Unit": "ZJ",
            "Activity_Id": "Constraining Data",
            "Source": "IGCC2024",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                ohc_data["timebound_lower"].values[i] : ohc_data["ocean_full-depth"].values[i]
                for i in range(len(ohc_data["timebound_lower"].values))
            },
        },
        index=[0],
    )
    tothc_row = pd.DataFrame(
        {
            "Model": "Observed",
            "Scenario": "historical",
            "Region": "World",
            "Variable": "Total Heat Content",
            "Unit": "ZJ",
            "Activity_Id": "Constraining Data",
            "Source": "IGCC2024",
            "Version": f"RCMIP Phase 3 {version}",
            **{
                ohc_data["timebound_lower"].values[i] : ohc_data["total"].values[i]
                for i in range(len(ohc_data["timebound_lower"].values))
            },
        },
        index=[0],
    )
    return ohc_row, tothc_row

def get_data_for_plots(version="v2.0.0"):
    """
    Loads and returns multiple datasets required for plotting climate-related distributions.

    Reads temperature, CO2 concentration, global carbon budget, aerosol forcing, and ocean heat content data
    from various CSV files and returns them as pandas DataFrames or appropriate objects.

    Returns
    -------
    temp_data : pandas.DataFrame
        Annual average temperature data.
    co2_conc : pandas.DataFrame or appropriate object
        CO2 concentration data, as returned by `read_noaa_gml_ml_means`.
    data_gcb : pandas.DataFrame or appropriate object
        Global Carbon Budget data, as returned by `read_gcb_data`.
    data_aer_best : pandas.DataFrame
        Best estimate of aerosol effective radiative forcing (ERF) data.
    data_aer_5 : pandas.DataFrame
        5th percentile aggregate aerosol ERF data.
    data_aer_95 : pandas.DataFrame
        95th percentile aggregate aerosol ERF data.
    data_ohc : pandas.DataFrame
        Ocean heat content (OHC) ensemble data.

    Notes
    -----
    The function expects certain file paths and directory variables (e.g., `datadir`) to be defined in the scope.
    """
    temp_data = read_temperature_data(version= version)
    co2_conc = read_noaa_gml_ml_means("year", version= version)
    data_gcb_ocean = read_gcb_ocean_carbon_data(version= version)
    data_gcb_land = read_gcb_data(version= version)
    # print(data_gcb_ocean.head())
    # print(data_gcb_land.head())
    data_aer_best = read_aer_ts(version= version)
    #data_aer_5 = pd.read_csv(f"{datadir}igcc_erf_p05_aggregates.csv")
    #data_aer_95 = pd.read_csv(f"{datadir}igcc_erf_p95_aggregates.csv")
    data_ohc, data_tothc = read_ohc(version= version)
    # print(data_ohc.head())
    full_dataframe = pd.concat([data_aer_best, temp_data, co2_conc, data_ohc, data_tothc, data_gcb_land, data_gcb_ocean], ignore_index=True)
    # print(full_dataframe)
    # print(full_dataframe.columns)
    cols_to_drop = [
        col
        for col in full_dataframe.columns
        if re.fullmatch(r"\d{4}", str(col)) and (int(str(col)) < 1850 or int(str(col)) > 2023)
    ]
    if cols_to_drop:
        full_dataframe.drop(columns=cols_to_drop, inplace=True)
        #print("Dropped columns:", cols_to_drop)
    return full_dataframe



if __name__ == "__main__":
    version = "v2.0.0"
    full_df = get_data_for_plots(version=version)
    full_df.to_csv(f"{rcmip_output_dir}/rcmip_phase3_processed_constraining_data_{version}.csv", index=False)
    constraints_df = define_and_print_uncertainty_ranges(version=version)
    constraints_df.to_csv(f"{rcmip_output_dir}/rcmip_phase3_constraint_targets_with_uncertainty_{version}.csv", index=False)
    constraints_df.to_latex(f"{rcmip_output_dir}/constraints_table_for_overleaf.tex", index=False)
    print(full_df)
    print(constraints_df)
    plot_timeseries_constraints(full_df, constraints_df=constraints_df, figsize=(12, 15), save_path=f"{datadir}/../figures/rcmip_phase3_emissions_processed_constraining_data_{version}.png")