Grapher

def encode(
    self,
    x: Optional[str] = None,
    y: Optional[str] = None,
    y_lower: Optional[str] = None,
    y_upper: Optional[str] = None,
    entity: Optional[str] = None,
    color: Optional[str] = None,
    size: Optional[str] = None,
) -> "Chart":
    """Map DataFrame columns to visual properties.

    This method establishes the visual encoding by mapping DataFrame columns to
    chart dimensions. The behavior varies by chart type:

    - **Line/Bar charts**: `x` is time, `y` is values, `entity` groups lines/bars
    - **Scatter plots**: `x` and `y` are both numeric values, `entity` groups points

    Args:
        x: Column name for x-axis. For line/bar charts, typically a time column
            ('year', 'date'). For scatter plots, a numeric value column.
        y: Column name for y-axis values to plot. For bar charts, can be the entity
            column if you want entities on the y-axis.
        y_lower: Column name for lower bound of confidence interval. When specified
            along with y_upper, renders a shaded confidence band around the main line.
        y_upper: Column name for upper bound of confidence interval. When specified
            along with y_lower, renders a shaded confidence band around the main line.
        entity: Column name for grouping data (e.g., 'country', 'region'). Each unique
            value becomes a separate line/series. Optional for single-series charts.
        color: Column name for color encoding in scatter plots. Values map to colors.
        size: Column name for size encoding in scatter plots. Values map to point sizes.

    Returns:
        Self for method chaining.

    Raises:
        ValueError: If a specified column name is not found in the DataFrame.

    Example:
        ```python
        # Line chart with multiple countries
        Chart(df).mark_line().encode(
            x='year',
            y='population',
            entity='country'
        )

        # Line chart with confidence intervals
        Chart(df).mark_line().encode(
            x='year',
            y='temperature',
            y_lower='temperature_low',
            y_upper='temperature_high',
            entity='region'
        )

        # Scatter plot with color and size
        Chart(df).mark_scatter().encode(
            x='gdp_per_capita',
            y='life_expectancy',
            entity='country',
            color='continent',
            size='population'
        )
        ```
    """
    self.x = x
    self.y = y
    self.y_lower = y_lower
    self.y_upper = y_upper
    self.entity = entity
    self.color = color
    self.size = size

    # fail early if there's been a typo
    for col in [x, y, y_lower, y_upper, entity, color, size]:
        if col and col not in self.data.columns:
            raise ValueError(f"no such column: {col}")

    if x == "date":
        self.time_type = TimeType.DAY

    self._state.hideLegend = not entity

    return self

def label(
    self, title: str = "", subtitle: str = "", source_desc: str = "", note: str = ""
) -> "Chart":
    """Add labels and text to the chart.

    Args:
        title: Chart title.
        subtitle: Chart subtitle.
        source_desc: Data source attribution.
        note: Additional notes or footnotes.

    Returns:
        Self for method chaining.
    """
    self._state.title = title if title else None
    self._state.subtitle = subtitle if subtitle else None
    self._state.sourceDesc = source_desc if source_desc else None
    self._state.note = note if note else None
    return self

def xaxis(
    self,
    label: Optional[str] = None,
    unit: Optional[str] = None,
    scale: Optional[Literal["linear", "log"]] = None,
    scale_control: Optional[bool] = None,
) -> "Chart":
    """Configure the x-axis.

    Args:
        label: Axis label text.
        unit: Unit of measurement (e.g., '$', 'kg').
        scale: Scale type ('linear' or 'log').
        scale_control: Allow users to toggle scale.

    Returns:
        Self for method chaining.
    """
    assert self._state.xAxis is not None  # Initialized in __init__
    if label is not None:
        self._state.xAxis.label = label
    if unit is not None:
        self.x_unit = unit
    if scale is not None:
        self._state.xAxis.scaleType = scale
    if scale_control is not None:
        self._state.xAxis.canChangeScaleType = scale_control
    return self

def yaxis(
    self,
    label: Optional[str] = None,
    unit: Optional[str] = None,
    scale: Optional[Literal["linear", "log"]] = None,
    scale_control: Optional[bool] = None,
) -> "Chart":
    """Configure the y-axis.

    Args:
        label: Axis label text.
        unit: Unit of measurement (e.g., '$', 'kg').
        scale: Scale type ('linear' or 'log').
        scale_control: Allow users to toggle scale.

    Returns:
        Self for method chaining.
    """
    assert self._state.yAxis is not None  # Initialized in __init__
    if label is not None:
        self._state.yAxis.label = label
    if unit is not None:
        self.y_unit = unit
    if scale is not None:
        self._state.yAxis.scaleType = scale
    if scale_control is not None:
        self._state.yAxis.canChangeScaleType = scale_control
    return self

def axis(
    self,
    x_label: Optional[str] = None,
    y_label: Optional[str] = None,
    x_unit: Optional[str] = None,
    y_unit: Optional[str] = None,
    x_scale: Optional[Literal["linear", "log"]] = None,
    y_scale: Optional[Literal["linear", "log"]] = None,
    x_scale_control: Optional[bool] = None,
    y_scale_control: Optional[bool] = None,
) -> "Chart":
    """Configure both axes at once.

    Convenience method for setting properties on both axes. For single-axis
    configuration, use xaxis() or yaxis() instead.

    Args:
        x_label: Label text for x-axis.
        y_label: Label text for y-axis.
        x_unit: Unit suffix for x-axis values (e.g., '$', '%', 'kg').
        y_unit: Unit suffix for y-axis values.
        x_scale: Scale type for x-axis ('linear' or 'log').
        y_scale: Scale type for y-axis ('linear' or 'log').
        x_scale_control: If True, adds UI control to toggle x-axis scale.
        y_scale_control: If True, adds UI control to toggle y-axis scale.

    Returns:
        Self for method chaining.
    """
    assert self._state.xAxis is not None  # Initialized in __init__
    assert self._state.yAxis is not None  # Initialized in __init__
    if x_label is not None:
        self._state.xAxis.label = x_label
    if y_label is not None:
        self._state.yAxis.label = y_label
    if x_unit is not None:
        self.x_unit = x_unit
    if y_unit is not None:
        self.y_unit = y_unit
    if x_scale is not None:
        self._state.xAxis.scaleType = x_scale
    if y_scale is not None:
        self._state.yAxis.scaleType = y_scale
    if x_scale_control is not None:
        self._state.xAxis.canChangeScaleType = x_scale_control
    if y_scale_control is not None:
        self._state.yAxis.canChangeScaleType = y_scale_control
    return self

def mark_scatter(self) -> "Chart":
    """Add scatter plot to available chart types.

    Scatter plots display individual data points with x and y positions. Use `color`
    and `size` encodings for additional dimensions. Best for showing relationships
    between two numeric variables.

    Can be combined with other mark_*() methods to enable multiple chart types.

    Returns:
        Self for method chaining.
    """
    self._add_chart_type("ScatterPlot")
    return self

def mark_line(self) -> "Chart":
    """Add line chart to available chart types.

    Line charts connect data points with lines, ideal for showing trends over time.
    Multiple entities create multiple lines.

    Can be combined with other mark_*() methods to enable multiple chart types.

    Returns:
        Self for method chaining.
    """
    self._add_chart_type("LineChart")
    return self

def mark_bar(self, stacked: bool = False) -> "Chart":
    """Add bar chart to available chart types.

    Bar charts display categorical data with rectangular bars. Bars can be shown
    side-by-side (default) or stacked on top of each other.

    Can be combined with other mark_*() methods to enable multiple chart types.

    Args:
        stacked: If True, creates a stacked bar chart where bars for different
            entities are stacked vertically. If False (default), bars are shown
            side-by-side.

    Returns:
        Self for method chaining.
    """
    chart_type = "StackedDiscreteBar" if stacked else "DiscreteBar"
    self._add_chart_type(chart_type)
    return self

def mark_slope(self) -> "Chart":
    """Add slope chart to available chart types.

    Slope charts compare values between two time points, showing the change
    as connecting lines between start and end values. Ideal for highlighting
    increases and decreases across entities.

    Can be combined with other mark_*() methods to enable multiple chart types.

    Returns:
        Self for method chaining.
    """
    self._add_chart_type("SlopeChart")
    return self

def mark_marimekko(self) -> "Chart":
    """Add Marimekko chart to available chart types.

    Marimekko charts (also called mosaic charts) show part-to-whole relationships
    where both width and height of segments are meaningful. Width represents one
    dimension (e.g., population) and height represents another (e.g., percentage).

    Can be combined with other mark_*() methods to enable multiple chart types.

    Returns:
        Self for method chaining.
    """
    self._add_chart_type("Marimekko")
    return self

def mark_map(
    self,
    time_tolerance: Optional[int] = None,
    color_scheme: Optional[ColorSchemeName] = None,
    binning_strategy: Optional[BinningStrategy] = None,
    custom_numeric_values: Optional[List[float]] = None,
) -> "Chart":
    """Enable the map tab with optional configuration.

    Adds a world map visualization showing data geographically. Can be combined
    with other mark_*() methods.

    Args:
        time_tolerance: How many years to look back/forward for data
        color_scheme: Color scheme name (e.g., "Reds", "Blues", "YlOrRd")
        binning_strategy: How to bin values ("auto", "manual", "equalInterval", "quantiles")
        custom_numeric_values: Custom bin boundaries when using manual binning

    Returns:
        Self for method chaining.

    Example:
        ```python
        # Line chart with map, defaulting to map view
        Chart(df).mark_line().mark_map().show("map").encode(...)

        # Line chart with customized map
        Chart(df).mark_line().mark_map(
            color_scheme='Reds',
            binning_strategy='manual',
            custom_numeric_values=[0, 1000, 10000, 100000]
        ).encode(...)
        ```
    """
    self._state.hasMapTab = True

    # Set default tab to map if this is the first mark_*() call
    if self._state.tab == "chart":  # Default value means not yet set
        self._state.tab = "map"

    # Configure map options if any provided
    if any([time_tolerance, color_scheme, binning_strategy, custom_numeric_values]):
        color_scale = ColorScaleConfig(
            baseColorScheme=color_scheme,
            binningStrategy=binning_strategy,
            customNumericValues=custom_numeric_values,
        )
        self._state.map = MapConfig(
            timeTolerance=time_tolerance,
            colorScale=color_scale
            if any([color_scheme, binning_strategy, custom_numeric_values])
            else None,
        )

    return self

def show(
    self,
    tab: Literal[
        "line",
        "discrete-bar",
        "stacked-discrete-bar",
        "scatter",
        "stacked-area",
        "slope",
        "stacked-bar",
        "marimekko",
        "map",
        "table",
    ],
) -> "Chart":
    """Set which tab to display by default.

    Use this to control which visualization is shown when the chart first loads.
    The tab must correspond to an enabled chart type (via mark_*() methods).

    Args:
        tab: The tab to show by default. Options:
            - "line": Line chart
            - "discrete-bar": Bar chart
            - "stacked-discrete-bar": Stacked bar chart
            - "scatter": Scatter plot
            - "map": World map
            - "table": Data table

    Returns:
        Self for method chaining.

    Example:
        ```python
        # Enable line and bar, but show bar by default
        Chart(df).mark_line().mark_bar().show("discrete-bar").encode(...)
        ```
    """
    self._state.tab = tab
    return self

def interact(
    self,
    allow_relative: Optional[bool] = None,
    scale_control: Optional[bool] = None,
    entity_control: Optional[bool] = None,
    entity_mode: Optional[EntitySelectionMode] = None,
) -> "Chart":
    """Add interactive controls to the chart.

    Args:
        allow_relative: Show relative/absolute toggle.
        scale_control: Show log/linear scale toggle.
        entity_control: Show entity/country picker (shorthand for entity_mode).
        entity_mode: Entity selection mode. Options:
            - "add-country": Allow adding multiple entities (default when entity_control=True)
            - "change-country": Allow only single entity selection (useful for
              charts with multiple lines per entity, e.g., confidence intervals)
            - "disabled": Disable entity selection

    Returns:
        Self for method chaining.
    """
    if allow_relative is not None:
        self._state.hideRelativeToggle = False

    if scale_control is not None:
        # Update yAxis without overwriting existing settings
        assert self._state.yAxis is not None  # Initialized in __init__
        self._state.yAxis.scaleType = "linear"
        self._state.yAxis.canChangeScaleType = scale_control

    if entity_mode is not None:
        self._state.addCountryMode = entity_mode
    elif entity_control is not None:
        self._state.addCountryMode = "add-country" if entity_control else "disabled"

    return self

def select(
    self,
    entities: Optional[List[str]] = None,
    timespan: Optional[Tuple[Any, Any]] = None,
) -> "Chart":
    """Pre-select entities and time range.

    Args:
        entities: List of entity names to display.
        timespan: Tuple of (start, end) for time range.

    Returns:
        Self for method chaining.
    """
    if entities:
        self.selection = entities

    if timespan:
        if isinstance(timespan, (str, int)):
            timespan = (timespan, None)
        self.timespan = timespan

    return self

def transform(self, relative: bool) -> "Chart":
    """Transform data to relative or absolute values.

    Display values as percentage change from a baseline (relative mode) or as
    absolute values (default). In relative mode, the first time period serves
    as the baseline (100%).

    Args:
        relative: If True, show values as percentage change from baseline.
            If False, show absolute values.

    Returns:
        Self for method chaining.
    """
    self._state.stackMode = "relative" if relative else "absolute"
    return self

def filter(self, matching_entities_only: bool = True) -> "Chart":
    """Filter entities to only show those with complete data.

    When enabled, only entities that have data for all time periods and dimensions
    will be shown. Useful for ensuring fair comparisons by excluding entities with
    incomplete data.

    Args:
        matching_entities_only: If True, only show entities with complete data across
            all dimensions and time periods. If False, show all entities even with gaps.

    Returns:
        Self for method chaining.
    """
    self._state.matchingEntitiesOnly = matching_entities_only
    return self

def variable(
    self,
    column: str,
    name: Optional[str] = None,
    description_short: Optional[str] = None,
    description_from_producer: Optional[str] = None,
    description_processing: Optional[str] = None,
    description_key: Optional[List[str]] = None,
    unit: Optional[str] = None,
    short_unit: Optional[str] = None,
    color: Optional[str] = None,
    source_name: Optional[str] = None,
    source_link: Optional[str] = None,
) -> "Chart":
    """Add rich metadata to a data column/variable.

    Configures display properties and documentation for a column that will
    appear in tooltips, the data table, and source information.

    Args:
        column: Name of the DataFrame column to configure.
        name: Display name (e.g., "Population" instead of "pop").
        description_short: Brief description shown in tooltips.
        description_from_producer: Original description from data source.
        description_processing: How the data was processed/transformed.
        description_key: List of key points about the variable.
        unit: Full unit name (e.g., "million people").
        short_unit: Abbreviated unit for compact display (e.g., "M").
        color: Hex color for the line/series (e.g., "#ca2628").
        source_name: Name of the data source.
        source_link: URL to the data source.

    Returns:
        Self for method chaining.

    Note:
        The timespan is computed automatically from the data's time column.

    Example:
        ```python
        Chart(df).mark_line().encode(
            x='year',
            y='population',
            entity='country'
        ).variable(
            'population',
            name='Population',
            description_short='Total population in millions',
            unit='million people',
            short_unit='M',
            source_name='World Bank',
            source_link='https://data.worldbank.org'
        )
        ```
    """
    self.variable_configs[column] = VariableConfig(
        name=name,
        description_short=description_short,
        description_from_producer=description_from_producer,
        description_processing=description_processing,
        description_key=description_key,
        unit=unit,
        short_unit=short_unit,
        color=color,
        source_name=source_name,
        source_link=source_link,
    )
    return self

def export(self) -> Dict[str, Any]:
    """Export the chart as the three components needed for rendering.

    Returns:
        Dictionary with keys:
            - csv_data: CSV string of the data
            - column_defs: List of column definition dicts for OwidTable
            - grapher_config: Dict of GrapherState configuration
    """
    (
        df,
        x_col,
        y_cols,
        entity_col,
        color_col,
        year_col,
        selected_entities,
        min_time,
        max_time,
    ) = self._prepare_data()

    return {
        "csv_data": self._build_csv(df),
        "column_defs": self._build_column_defs(y_cols, x_col, year_col),
        "grapher_config": self._build_grapher_config(
            y_cols, x_col, color_col, selected_entities, min_time, max_time
        ),
    }

def save_png(
    self, path: str, include_details: bool = False, timeout: int = 30000
) -> None:
    """Save the chart as a PNG image.

    This method uses Playwright to render the chart in a headless browser
    and export it using the Grapher's built-in rasterize function.

    Args:
        path: File path to save the PNG.
        include_details: Whether to include chart details/sources in export.
        timeout: Timeout in milliseconds for rendering.

    Raises:
        ImportError: If Playwright is not installed. Install with:
            pip install playwright && playwright install chromium

    Example:
        ```python
        chart = Chart(df).mark_line().encode(x='year', y='population')
        chart.save_png("my_chart.png")
        ```
    """
    from owid.grapher.export import save_png

    save_png(self, path, include_details=include_details, timeout=timeout)

def save_svg(
    self, path: str, include_details: bool = False, timeout: int = 30000
) -> None:
    """Save the chart as an SVG image.

    This method uses Playwright to render the chart in a headless browser
    and export it using the Grapher's built-in rasterize function.

    Args:
        path: File path to save the SVG.
        include_details: Whether to include chart details/sources in export.
        timeout: Timeout in milliseconds for rendering.

    Raises:
        ImportError: If Playwright is not installed. Install with:
            pip install playwright && playwright install chromium

    Example:
        ```python
        chart = Chart(df).mark_line().encode(x='year', y='population')
        chart.save_svg("my_chart.svg")
        ```
    """
    from owid.grapher.export import save_svg

    save_svg(self, path, include_details=include_details, timeout=timeout)