ResInsight/GrpcInterface/Python/rips/well_events.py

"""
Well Events module for managing time-dependent well events.

This module provides functionality for creating and managing well events
in a timeline-based event system. Events can be perforation events, valve events,
tubing changes, well state changes, and production/injection control changes.
"""

from typing import Any, Dict, List
from datetime import date, datetime

from .pdmobject import add_method
from .resinsight_classes import EclipseCase
from .generated.generated_classes import (
    KeywordEvent,
    WellEventKeyword,
    WellEventTimeline,
    WellPath,
)


def _format_date(event_date: str | date | datetime) -> str:
    """Convert date to ISO format string (YYYY-MM-DD)."""
    if isinstance(event_date, str):
        return event_date
    elif isinstance(event_date, datetime):
        return event_date.strftime("%Y-%m-%d")
    elif isinstance(event_date, date):
        return event_date.strftime("%Y-%m-%d")
    else:
        raise TypeError(
            f"event_date must be a string, date, or datetime, not {type(event_date)}"
        )


@add_method(WellEventTimeline)
def add_well_keyword_event(
    self: WellEventTimeline,
    event_date: str | date | datetime,
    well_path: Any,
    keyword_name: str,
    keyword_data: Dict[str, Any],
) -> WellEventKeyword:
    """Add a well keyword event with arbitrary keyword data.

    This is a convenience method that automatically infers types from Python
    values and calls the underlying GRPC method with parallel arrays.

    Type inference rules:
    - str → STRING
    - int → INT
    - float → DOUBLE
    - bool → INT (1 or 0)

    Arguments:
        event_date: Date string in YYYY-MM-DD format, date, or datetime object
        well_path (WellPath): The well path object
        keyword_name (str): Keyword name (e.g., "WCONHIST", "WELTARG", "WRFTPLT")
        keyword_data (dict): Dictionary mapping keyword item names to values

    Returns:
        WellEventKeyword: The created keyword event object

    Raises:
        TypeError: If keyword_data contains unsupported value types

    Example:
        ```python
        # Get the timeline
        well_path = project.well_paths()[0]
        timeline = well_path.event_timeline()

        # Add WCONHIST - Historical production data
        timeline.add_well_keyword_event(
            event_date="2018-04-01",
            well_path=well_path,
            keyword_name="WCONHIST",
            keyword_data={
                "WELL": well_path.name,
                "STATUS": "OPEN",
                "CMODE": "RESV",
                "ORAT": 3999.98999,
                "WRAT": 0.01,
                "GRAT": 550678.438,
                "VFP_TABLE": 1
            }
        )

        # Add WELTARG - Change target
        timeline.add_well_keyword_event(
            event_date="2018-05-01",
            well_path=well_path,
            keyword_name="WELTARG",
            keyword_data={
                "WELL": well_path.name,
                "CMODE": "ORAT",
                "NEW_VALUE": 5000.0
            }
        )

        # Add WRFTPLT - Enable RFT output
        timeline.add_well_keyword_event(
            event_date="2018-06-01",
            well_path=well_path,
            keyword_name="WRFTPLT",
            keyword_data={
                "WELL": well_path.name,
                "OUTPUT_RFT": "YES",
                "OUTPUT_PLT": "NO"
            }
        )

        # Generate schedule
        case = project.cases()[0]
        schedule_text = timeline.generate_schedule_text(
            eclipse_case=case, export_msw_for_wells=[well_path]
        )
        print(schedule_text)
        ```
    """
    # Type inference and conversion
    item_names = []
    item_types = []
    item_values = []

    for name, value in keyword_data.items():
        # Handle bool before int (bool is subclass of int in Python).
        # bool semantics: True -> flag (emitted as bare KEY in mnemonic-list keywords,
        # ignored elsewhere); False -> the entry is dropped entirely.
        if isinstance(value, bool):
            if not value:
                continue
            item_names.append(name)
            item_types.append("FLAG")
            # The value is ignored for FLAG items but must be a non-empty string;
            # empty strings get dropped by the GRPC vector<string> serialization.
            item_values.append("1")
            continue

        item_names.append(name)

        if isinstance(value, str):
            item_types.append("STRING")
            item_values.append(value)
        elif isinstance(value, int):
            item_types.append("INT")
            item_values.append(str(value))
        elif isinstance(value, float):
            item_types.append("DOUBLE")
            item_values.append(str(value))
        else:
            raise TypeError(
                f"Unsupported type for keyword item '{name}': {type(value).__name__}. "
                "Supported types: str, int, float, bool"
            )

    # Format date
    date_str = _format_date(event_date)

    # Call internal GRPC method with parallel arrays
    return self.add_well_keyword_event_internal(
        event_date=date_str,
        well_path=well_path,
        keyword_name=keyword_name,
        item_names=item_names,
        item_types=item_types,
        item_values=item_values,
    )


@add_method(WellEventTimeline)
def add_keyword_event(
    self: WellEventTimeline,
    event_date: str | date | datetime,
    keyword_name: str,
    keyword_data: Dict[str, Any],
) -> KeywordEvent:
    """Add a schedule-level keyword event (not tied to a specific well path).

    This is for global schedule keywords like RPTRST, GRUPTREE, RPTSCHED, etc.
    that apply to the entire simulation rather than a specific well.

    Type inference rules:
    - str → STRING
    - int → INT
    - float → DOUBLE
    - bool → INT (1 or 0)

    Arguments:
        event_date: Date string in YYYY-MM-DD format, date, or datetime object
        keyword_name (str): Keyword name (e.g., "RPTRST", "GRUPTREE", "RPTSCHED")
        keyword_data (dict): Dictionary mapping keyword item names to values

    Returns:
        KeywordEvent: The created keyword event object

    Raises:
        TypeError: If keyword_data contains unsupported value types

    Example:
        ```python
        # Get the timeline
        well_path_coll = project.descendants(rips.WellPathCollection)[0]
        timeline = well_path_coll.event_timeline()

        # Add RPTRST - Report restart settings (schedule-level, not well-specific)
        timeline.add_keyword_event(
            event_date="2024-01-01",
            keyword_name="RPTRST",
            keyword_data={
                "BASIC": 2,
                "FREQ": 1,
            }
        )

        # Add GRUPTREE - Group tree definition
        timeline.add_keyword_event(
            event_date="2024-01-01",
            keyword_name="GRUPTREE",
            keyword_data={
                "CHILD": "OP",
                "PARENT": "FIELD",
            }
        )

        # Add RPTSCHED - Report schedule settings
        timeline.add_keyword_event(
            event_date="2024-01-01",
            keyword_name="RPTSCHED",
            keyword_data={
                "FIP": 1,
                "WELLS": 2,
            }
        )
        ```
    """
    # Type inference and conversion
    item_names = []
    item_types = []
    item_values = []

    for name, value in keyword_data.items():
        # Handle bool before int (bool is subclass of int in Python).
        # bool semantics: True -> flag (emitted as bare KEY in mnemonic-list keywords,
        # ignored elsewhere); False -> the entry is dropped entirely.
        if isinstance(value, bool):
            if not value:
                continue
            item_names.append(name)
            item_types.append("FLAG")
            # The value is ignored for FLAG items but must be a non-empty string;
            # empty strings get dropped by the GRPC vector<string> serialization.
            item_values.append("1")
            continue

        item_names.append(name)

        if isinstance(value, str):
            item_types.append("STRING")
            item_values.append(value)
        elif isinstance(value, int):
            item_types.append("INT")
            item_values.append(str(value))
        elif isinstance(value, float):
            item_types.append("DOUBLE")
            item_values.append(str(value))
        else:
            raise TypeError(
                f"Unsupported type for keyword item '{name}': {type(value).__name__}. "
                "Supported types: str, int, float, bool"
            )

    # Format date
    date_str = _format_date(event_date)

    # Call internal GRPC method with parallel arrays
    return self.add_keyword_event_internal(
        event_date=date_str,
        keyword_name=keyword_name,
        item_names=item_names,
        item_types=item_types,
        item_values=item_values,
    )


@add_method(WellEventTimeline)
def generate_schedule_text(
    self: WellEventTimeline,
    eclipse_case: EclipseCase,
    export_msw_for_wells: List[WellPath] = [],
    first_date_as_comment: bool = True,
    align_columns: bool = False,
) -> str:
    """Generate Eclipse schedule text for all wells in the collection.

    The timeline is shared across all wells in the well path collection.
    This method generates schedule data for all wells that have events in
    the timeline.

    This is a convenience wrapper around generate_schedule() that returns the
    text directly instead of a DataContainerString.

    Arguments:
        eclipse_case (EclipseCase): Eclipse case to use for schedule generation.
        export_msw_for_wells (List[WellPath]): Wells for which the
            multi-segment-well keywords (WELSEGS, COMPSEGS, WSEGVALV, WSEGAICD)
            are exported. Wells not in the list get no MSW keywords. An empty
            list (the default) suppresses MSW export for all wells.
        first_date_as_comment (bool): When True (the default), the first
            (earliest) date is written as a comment line (e.g. "-- Date: 1 JAN
            2024") instead of a DATES keyword. This avoids a DATES entry equal
            to the simulation start date, which some commercial simulators
            reject. Later dates are always emitted as DATES keywords.
        align_columns (bool): When True, emit each keyword with a "--"-prefixed
            column-header comment and right-aligned, fixed-width columns instead
            of the compact default form. Defaults to False.

    Returns:
        str: Eclipse schedule text containing DATES, COMPDAT, WELSEGS, WCONPROD, etc.
             for all wells in the collection.

    Example:
        ```python
        # Get the timeline (shared across all wells)
        well_path = project.well_paths()[0]
        timeline = well_path.event_timeline()

        # Add events for multiple wells
        timeline.add_perf_event(
            event_date="2024-01-01",
            well_name="WELL-1",
            start_md=1000,
            end_md=1500,
            diameter=0.1,
            skin_factor=0.5,
            state="OPEN"
        )

        timeline.add_perf_event(
            event_date="2024-02-01",
            well_name="WELL-2",
            start_md=2000,
            end_md=2500,
            diameter=0.1,
            state="OPEN"
        )

        # Generate schedule text, exporting MSW keywords for all wells
        case = project.cases()[0]
        schedule_text = timeline.generate_schedule_text(
            eclipse_case=case, export_msw_for_wells=project.well_paths()
        )
        print(schedule_text)
        ```
    """
    container = self.generate_schedule(
        eclipse_case=eclipse_case,
        export_msw_for_wells=export_msw_for_wells,
        first_date_as_comment=first_date_as_comment,
        align_columns=align_columns,
    )
    if container and container.values:
        return "".join(container.values)
    return ""