import argparse
import re
from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd


def find_col(df, name):
    """Find a column name in df using case-insensitive exact match."""
    for col in df.columns:
        if str(col).strip().lower() == name.lower():
            return col
    return None


def main():
    parser = argparse.ArgumentParser(
        description="Plot all CO2 concentration trajectories from an RCMIP concentrations CSV."
    )
    parser.add_argument(
        "--input",
        default="../RCMIP3_input_datafiles/rcmip_phase3_concentrations_v2.0.0.csv",
        help="Path to the concentrations CSV file.",
    )
    parser.add_argument(
        "--output",
        default="co2_concentration_trajectories.png",
        help="Output image file path.",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Show the figure interactively in addition to saving it.",
    )
    args = parser.parse_args()

    df = pd.read_csv(args.input)

    scenario_col = find_col(df, "Scenario")
    variable_col = find_col(df, "Variable")
    region_col = find_col(df, "Region")

    if scenario_col is None or variable_col is None:
        raise ValueError("Input CSV must contain Scenario and Variable columns.")

    year_cols = [c for c in df.columns if re.fullmatch(r"\d{4}", str(c))]
    if not year_cols:
        raise ValueError("No 4-digit year columns found in the input file.")
    year_cols = sorted(year_cols, key=int)

    # Keep only global rows if a Region column exists.
    if region_col is not None:
        df = df[df[region_col] == "World"]

    # Prefer exact variable name, fallback to any row containing CO2.
    co2 = df[df[variable_col] == "Atmospheric Concentrations|CO2"].copy()
    if co2.empty:
        co2 = df[df[variable_col].astype(str).str.contains("CO2", regex=False)].copy()

    if co2.empty:
        raise ValueError("No CO2 concentration rows found.")

    co2[year_cols] = co2[year_cols].apply(pd.to_numeric, errors="coerce")

    # If there are duplicate rows per scenario, combine by mean trajectory.
    co2 = co2.groupby(scenario_col, as_index=False)[year_cols].mean(numeric_only=True)

    years = [int(y) for y in year_cols]

    plt.figure(figsize=(12, 7))
    for _, row in co2.iterrows():
        plt.plot(
            years,
            row[year_cols].values,
            alpha=0.7,
            linewidth=1.0,
            label=str(row[scenario_col]),
        )

    plt.title("CO2 Concentration Trajectories (All Scenarios)")
    plt.xlabel("Year")
    plt.ylabel("CO2 concentration")
    plt.grid(alpha=0.3)
    # Put labels in a legend outside the plot area so lines remain readable.
    plt.legend(
        title="Scenario",
        bbox_to_anchor=(1.02, 1),
        loc="upper left",
        borderaxespad=0,
        fontsize=8,
    )
    plt.tight_layout()

    out_path = Path(args.output)
    plt.savefig(out_path, dpi=200)
    print(f"Saved plot to: {out_path.resolve()}")

    if args.show:
        plt.show()


if __name__ == "__main__":
    main()
