Grapher

def encode(
    self,
    x: Optional[str] = None,
    y: 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.
        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'
        )

        # 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.entity = entity
    self.color = color
    self.size = size

    # fail early if there's been a typo
    for col in [x, y, 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.config.hide_legend = 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.config.title = title
    self.config.subtitle = subtitle
    self.config.source_desc = source_desc
    self.config.note = note
    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.
    """
    if label is not None:
        self.config.x_axis["label"] = label
    if unit is not None:
        self.x_unit = unit
    if scale is not None:
        self.config.x_axis["scaleType"] = scale
    if scale_control is not None:
        self.config.x_axis["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.
    """
    if label is not None:
        self.config.y_axis["label"] = label
    if unit is not None:
        self.y_unit = unit
    if scale is not None:
        self.config.y_axis["scaleType"] = scale
    if scale_control is not None:
        self.config.y_axis["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.
    """
    if x_label is not None:
        self.config.x_axis["label"] = x_label
    if y_label is not None:
        self.config.y_axis["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.config.x_axis["scaleType"] = x_scale
    if y_scale is not None:
        self.config.y_axis["scaleType"] = y_scale
    if x_scale_control is not None:
        self.config.x_axis["canChangeScaleType"] = x_scale_control
    if y_scale_control is not None:
        self.config.y_axis["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_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.config.has_map_tab = True

    # Set default tab to map if this is the first mark_*() call
    if self.default_tab is None:
        self.default_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.config.map_config = 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.default_tab = tab
    return self

def interact(
    self,
    allow_relative: Optional[bool] = None,
    scale_control: Optional[bool] = None,
    entity_control: Optional[bool] = 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.

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

    if scale_control is not None:
        # Update y_axis without overwriting existing settings
        self.config.y_axis["scaleType"] = "linear"
        self.config.y_axis["canChangeScaleType"] = scale_control

    if entity_control is not None:
        self.config.hide_entity_controls = not entity_control

    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.config.stack_mode = "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.config.matching_entities_only = 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,
    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").
        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,
        source_name=source_name,
        source_link=source_link,
    )
    return self

def export(self, include_data: bool = True) -> Dict[str, Any]:
    """Export the chart configuration as a dictionary.

    Args:
        include_data: If True, includes the data in the export.

    Returns:
        Dictionary containing the full chart configuration.
    """
    config = self.config.to_dict()  # type: ignore

    # Add chart types and tab
    chart_types = self.chart_types if self.chart_types else ["LineChart"]
    config["chartTypes"] = chart_types

    # Set tab based on default_tab or first chart type
    if self.default_tab:
        config["tab"] = self.default_tab
    elif chart_types:
        config["tab"] = _CHART_TYPE_TO_TAB.get(chart_types[0], "line")

    # Auto-improve: show title annotation for line charts with titles
    if self.config.title and "LineChart" in chart_types:
        config["hideTitleAnnotation"] = False

    # Convert MapConfig to dict (dataclass_json doesn't handle it automatically)
    if self.config.map_config is not None:
        config["mapConfig"] = self.config.map_config.to_dict()

    primary_type = self._get_primary_chart_type()
    config.update(self.data_config().to_dict(primary_type))
    config = prune(config)
    if not include_data:
        del config["owidDataset"]["data"]
    return config