def style_scan(df):
    def hide_nan_css(v):
        return "visibility: hidden;" if pd.isna(v) else ""

    # Titleize column names
    df.columns = pd.MultiIndex.from_arrays(
        [
            [titleize(col[0]) for col in df.columns],
            [titleize(col[1]) for col in df.columns],
        ]
    )

    # Create styler with background gradient per column
    styler = df.style.background_gradient(axis=0)

    # Hide NaNs
    # styler = styler.format(lambda x: '' if pd.isna(x) else '{:.2f}'.format(x))
    styler = styler.map(hide_nan_css).format(na_rep="", precision=2)

    return styler


def get_cai_series(geographyId):
    """
    Get and process CAI series data for a given country.

    Parameters:
    -----------
    ctry_iso2 : str
        Two-letter ISO country code

    Returns:
    --------
    pandas.DataFrame
        DataFrame containing Hard, Soft, and Headline CAI metrics
    """
    cai_series_soft_vs_hard = gs.get_CAI_series(
        geographyId=geographyId,
        metricName=[
            "CAI_HEADLINE",
            "CAI_CONTRIBUTION_TYPE_HARD",
            "CAI_CONTRIBUTION_TYPE_SOFT",
        ],
        startDate="1980-01-01",
    )
    cai_series_soft_vs_hard = cai_series_soft_vs_hard.set_index(
        "metricName", append=True
    )["metricValue"].unstack("metricName")
    cai_series_soft_vs_hard = cai_series_soft_vs_hard.rename(
        columns={
            "CAI_CONTRIBUTION_TYPE_HARD": "Hard",
            "CAI_CONTRIBUTION_TYPE_SOFT": "Soft",
            "CAI_HEADLINE": "Headline",
        }
    )
    return cai_series_soft_vs_hard


def get_cfi_series(geographyId):
    """
    Get and process CFI series data for a given country.

    Parameters:
    -----------
    ctry_iso2 : str
        Two-letter ISO country code

    Returns:
    --------
    pandas.DataFrame
        DataFrame containing Hard, Soft, and Headline CAI metrics
    """
    data = gs.get_FCI_series(geographyId=geographyId, startDate="1980-01-01")

    return data


def summarize_gs_eco_fct(raw_data):
    """
    Summarize GS economic forecast data by filtering and organizing key metrics.

    This function processes raw economic forecast data to extract and organize
    specific key metrics into a standardized format. It filters for important
    economic indicators like CPI, GDP, current account, and output gap metrics,
    ensuring they have the correct periodicity, value types, and adjustments.

    Parameters
    ----------
    raw_data : pandas.DataFrame
        Raw economic forecast data with multi-level index containing various
        economic metrics and their metadata

    Returns
    -------
    pandas.Series
        Filtered and organized economic forecast values with standardized index.
        Returns a Series indexed by metric shorthand names (e.g., 'core_cpi',
        'rgdp_yoy') containing the latest forecast values.

    Raises
    ------
    ValueError
        If no valid forecast data is found after filtering
    AssertionError
        If the input data contains more than one geography or geographyId

    Examples
    --------
    >>> gs_data = gs.get_eco_forecast(geographyId='US')
    >>> summary = summarize_gs_eco_fct(gs_data)
    >>> summary['rgdp_yoy']  # Real GDP YoY growth forecast
    """

    # Define grouping columns for clarity
    GROUPBY_COLS = [
        "geographyName",
        "geographyId",
        "metricName",
        "periodType",
        "valueType",
        "isSeasonallyAdjusted",
        "isAnnualized",
        "isReal",
        "valueCurrency",
        "weightingSchema",
    ]

    # Aggregate to get the latest forecast for each metric specification
    eco_fct = (
        raw_data.reset_index()
        .dropna(subset=["forecastValue"])  # Only drop rows with missing forecast values
        .sort_index()
        .groupby(GROUPBY_COLS)
        .last()  # Get most recent forecast for each metric
    )

    # Validate single geography
    if len(eco_fct) == 0:
        raise ValueError("No valid forecast data found after filtering")

    assert len(eco_fct.index.levels[0]) == 1, (
        f"Expected one geography, found {len(eco_fct.index.levels[0])}: "
        f"{eco_fct.index.levels[0].tolist()}"
    )
    assert len(eco_fct.index.levels[1]) == 1, (
        f"Expected one geographyId, found {len(eco_fct.index.levels[1])}: "
        f"{eco_fct.index.levels[1].tolist()}"
    )

    # Drop geography levels now that we've validated
    eco_fct = eco_fct.droplevel(["geographyName", "geographyId"])

    # Define key metrics to extract with their specifications
    # Format: (metricName, periodType, valueType, isSeasonallyAdjusted, isAnnualized, isReal, valueCurrency, weightingSchema)
    KEY_METRICS = {
        # Inflation metrics (Yearly YoY)
        "core_cpi": ("core_cpi", "Yearly", "YoY", 1, 0, 0, "N/A", "N/A"),
        "cpi_avg": ("cpi_avg", "Yearly", "YoY", 1, 0, 0, "N/A", "N/A"),
        # External balance (Quarterly)
        "current_account": (
            "current_account",
            "Quarterly",
            "% of GDP",
            1,
            0,
            0,
            "N/A",
            "N/A",
        ),
        # GDP metrics
        "ngdp_lcl": (
            "gdp",
            "Quarterly",
            "Level",
            1,
            0,
            0,
            "Local",
            "N/A",
        ),  # Nominal GDP in local currency
        "ngdp_usd": (
            "gdp",
            "Quarterly",
            "Level",
            1,
            0,
            0,
            "USD",
            "N/A",
        ),  # Nominal GDP in USD
        "rgdp_qoq": (
            "gdp",
            "Quarterly",
            "QoQ",
            1,
            1,
            1,
            "N/A",
            "N/A",
        ),  # Real GDP QoQ (annualized)
        "rgdp_yoy": ("gdp", "Yearly", "YoY", 1, 0, 1, "N/A", "N/A"),  # Real GDP YoY
        # Slack measure
        "output_gap": ("output_gap", "Quarterly", "% of GDP", 0, 0, 0, "N/A", "N/A"),
    }

    # Create MultiIndex for the metrics we want
    metric_specs = pd.MultiIndex.from_tuples(
        KEY_METRICS.values(),
        names=[
            "metricName",
            "periodType",
            "valueType",
            "isSeasonallyAdjusted",
            "isAnnualized",
            "isReal",
            "valueCurrency",
            "weightingSchema",
        ],
    )

    # Reindex to get only the metrics we want (fills with NaN if missing)
    eco_filtered = eco_fct.reindex(metric_specs)

    # Rename index to short metric names
    eco_filtered.index = list(KEY_METRICS.keys())
    eco_filtered.index.name = "metric"

    # Extract just the forecast values as a Series
    result = eco_filtered["forecastValue"]

    return result

MAIN_COUNTRIES = {
    "ARG": {"iso_2": "AR", "iso_3": "032", "name": "Argentina"},
    "AUS": {"iso_2": "AU", "iso_3": "036", "name": "Australia"},
    "BRA": {"iso_2": "BR", "iso_3": "986", "name": "Brazil"},
    "CAN": {"iso_2": "CA", "iso_3": "124", "name": "Canada"},
    "CHN": {"iso_2": "CN", "iso_3": "156", "name": "China"},
    "FRA": {"iso_2": "FR", "iso_3": "978", "name": "France"},
    "DEU": {"iso_2": "DE", "iso_3": "978", "name": "Germany"},
    "IND": {"iso_2": "IN", "iso_3": "356", "name": "India"},
    "IDN": {"iso_2": "ID", "iso_3": "360", "name": "Indonesia"},
    "ITA": {"iso_2": "IT", "iso_3": "978", "name": "Italy"},
    "JPN": {"iso_2": "JP", "iso_3": "392", "name": "Japan"},
    "KOR": {"iso_2": "KR", "iso_3": "410", "name": "Suth Korea"},
    "MEX": {"iso_2": "MX", "iso_3": "484", "name": "Mexico"},
    "RUS": {"iso_2": "RU", "iso_3": "643", "name": "Russia"},
    "SAU": {"iso_2": "SA", "iso_3": "682", "name": "Saudi Arabia"},
    "ZAF": {"iso_2": "ZA", "iso_3": "710", "name": "South Africa"},
    "TUR": {"iso_2": "TR", "iso_3": "949", "name": "Turkey"},
    "GBR": {"iso_2": "GB", "iso_3": "826", "name": "United Kingdom"},
    "USA": {"iso_2": "US", "iso_3": "840", "name": "United States"},
    "NLD": {"iso_2": "NL", "iso_3": "978", "name": "Netherlands"},
    "POL": {"iso_2": "PL", "iso_3": "985", "name": "Poland"},
    "DNK": {"iso_2": "DK", "iso_3": "208", "name": "Denmark"},
    "FIN": {"iso_2": "FI", "iso_3": "978", "name": "Finland"},
    "ISL": {"iso_2": "IS", "iso_3": "352", "name": "Iceland"},
    "NOR": {"iso_2": "NO", "iso_3": "578", "name": "Norway"},
    "SWE": {"iso_2": "SE", "iso_3": "752", "name": "Sweden"},
    "EUR": {"iso_2": "EAagg", "iso_3": None, "name": "Eurozone"},
    "WLD": {"iso_2": "GLagg", "iso_3": None, "name": "World"},
    "DEV": {"iso_2": "DMagg", "iso_3": None, "name": "Developed World"},
    "EME": {"iso_2": "EMagg", "iso_3": None, "name": "Emerging World"},
}

AGGREGATES = ["WLD", "DEV", "EME"]
COUNTRIES = list(set(MAIN_COUNTRIES.keys()) - set(AGGREGATES))

gs_cov = GS_COVERAGE.groupby("Geography Id").count()
gs_cov = gs_cov.loc[gs_cov["Dataset"] > 1].index.sort_values().tolist()

DATA = {k: v for k, v in MAIN_COUNTRIES.items() if v["iso_2"] in gs_cov}

for ctry, value in tqdm(DATA.items(), leave=False):
    iso2 = value["iso_2"]
    DATA[ctry]["eco"] = gs.get_eco_forecast(geographyId=iso2)
    DATA[ctry]["cai"] = get_cai_series(geographyId=iso2)

cai_summary = {ctry: DATA[ctry]["cai"].iloc[-1].rename(ctry) for ctry in DATA}
cai_summary = pd.concat(
    cai_summary, axis=1
).T  # .rename(columns={'metricName': 'Country'})
cai_summary.columns.name = "Component"
cai_summary.index.name = "Country"
cai_summary = cai_summary.loc[:, ["Headline", "Hard", "Soft"]]

eco_summary = {ctry: summarize_gs_eco_fct(DATA[ctry]["eco"]) for ctry in DATA}
eco_summary = pd.concat(eco_summary, axis=1).T.sort_values(
    by="ngdp_usd", ascending=False
)
eco_summary.index.name = "Country"
eco_summary.columns.name = "Metric"

scan = (
    pd.concat([cai_summary, eco_summary], keys=["CAI", "Economic Forecasts"], axis=1)
    .sort_values(by=("Economic Forecasts", "ngdp_usd"), ascending=False)
    .drop(["ngdp_usd", "ngdp_lcl"], axis=1, level=1)
)

WLD_CAI = pd.DataFrame(
    {
        "World": DATA["WLD"]["cai"]["Headline"],
        "Developed Economies": DATA["DEV"]["cai"]["Headline"],
        "Emerging Economies ": DATA["EME"]["cai"]["Headline"],
    }
)

fig = plot_lines(
    WLD_CAI,
    show_0=True,
    title=f"<b>Global Current Activity Indicator</b>",
    years_limit=4,
    source="Goldman Sachs",
)

widths = [4, 2, 2]  # Example: customize per series
for trace, w in zip(fig.data, widths):
    trace.update(line=dict(width=w))

display_table_and_chart(
    table=style_scan(scan.loc[AGGREGATES].dropna(axis=1, how="all")), chart=fig
)

months = 12
contribution_mapping = {
    "realLongRatesContribution": "Real Long Rates",
    "longRatesContribution": "Long Rates",
    "equitiesContribution": "Equities",
    "shortRatesContribution": "Short Rates",
    "realShortRatesContribution": "Real Short Rates",
}

global_agg_cfi = get_cfi_series("GLagg")
topline = global_agg_cfi[["fci", "realFCI"]].rename(
    columns={"fci": "FCI", "realFCI": "Real FCI"}
)
contributor_line = global_agg_cfi[list(contribution_mapping.keys())].rename(
    columns=contribution_mapping
)
topline_fig = plot_lines(topline, title="Overall Index", figsize=(500, 500))
topline_chg = plot_lines(
    topline.diff(months),
    title="12M Change of Index",
    years_limit=3,
    figsize=(500, 500),
    show_0=True,
)
contributor_chg = plot_lines(
    contributor_line.diff(months),
    title="",
    years_limit=3,
    figsize=(500, 500),
    show_0=True,
)  # Contributors

display_by_columns([topline_fig, topline_chg, contributor_chg])

country_scan = scan.drop(AGGREGATES).dropna(axis=1, how="all")
scan_styled = style_scan(country_scan)

scan_styled

cai_plots = {}
for ctry in country_scan.index:
    cai_plots[ctry] = plot_lines(
        DATA[ctry]["cai"],
        show_0=True,
        title=f"<b>{ctry} Current Activity Indicator</b>",
        years_limit=4,
    )

for ctry, plot in cai_plots.items():
    plot.show()

from tulip_mania.notebook_related import notebook_updated

notebook_updated()