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,
segment changes, well state changes, and production/injection control changes.
"""

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

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


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:
        .. code-block:: 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:
        .. code-block:: 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_GROUP": "OP",
                    "PARENT_GROUP": "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 add_raw_text_event(
    self: WellEventTimeline,
    event_date: str | date | datetime,
    text: str,
    placement: str = "AFTER_DATE",
    anchor_keyword: str = "",
    priority: int = 0,
) -> Any:
    """Add raw text at a specific position in a dated schedule section.

    ``placement`` is one of ``AFTER_DATE``, ``BEFORE_KEYWORD``,
    ``AFTER_KEYWORD``, or ``END_OF_DATE``. ``anchor_keyword`` is required for
    before/after-keyword placement and must be empty for the other placements.
    Lower priority values are emitted first; source order breaks ties.
    """
    return self.add_raw_text_event_internal(
        event_date=_format_date(event_date),
        text=text,
        placement=Placement(placement),
        anchor_keyword=anchor_keyword,
        priority=priority,
    )


@add_method(WellEventTimeline)
def add_insert_date_event(
    self: WellEventTimeline,
    event_date: str | date | datetime,
    comment: str = "",
) -> Any:
    """Add a date that is emitted as a DATES keyword even without other events.

    In Eclipse/Flow a DATES entry ensures a summary report at that date. The
    event creates no keyword of its own; ``comment`` (if any) is written as
    "--"-prefixed lines directly below the date.
    """
    return self.add_insert_date_event_internal(
        event_date=_format_date(event_date),
        comment=comment,
    )


@add_method(WellEventTimeline)
def generate_schedule_text(
    self: WellEventTimeline,
    eclipse_case: Optional[EclipseCase] = None,
    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 (Optional[EclipseCase]): Grid case used for the
            grid-dependent keywords. Any Eclipse-type case works, including
            Roff and corner point cases. When None (the default), the
            keywords that need grid geometry (WELSPECS, COMPORD, COMPDAT,
            COMPLUMP, WELSEGS, COMPSEGS, WSEGVALV, WSEGAICD) are skipped. Each
            well with skipped keywords gets a "-- WARNING:" comment line in the
            schedule header. DATES, well controls (WCONHIST, WCONINJH, WELTARG,
            ...), schedule keywords and raw text are still emitted.
        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.

    Dates added with add_insert_date_event() are ordinary timeline events, but
    since they only force a DATES keyword (e.g. to trigger a summary report)
    they are not filtered by set_timestamp(). If such a date precedes all other
    event dates it becomes the earliest date and is therefore emitted as a
    comment when first_date_as_comment is True.

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

    Example:
        .. code-block:: 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 ""


def _format_query_date(query_date: Optional[str | date | datetime]) -> str:
    """Convert an optional query bound to an ISO string. None gives an empty
    string (no limit). A datetime keeps its time of day."""
    if query_date is None:
        return ""
    if isinstance(query_date, datetime):
        return query_date.isoformat(timespec="seconds")
    return _format_date(query_date)


@add_method(WellEventTimeline)
def event_dates(
    self: WellEventTimeline,
    start_date: Optional[str | date | datetime] = None,
    end_date: Optional[str | date | datetime] = None,
    wells: Optional[List[WellPath]] = None,
    event_types: Optional[List[str]] = None,
) -> List[datetime]:
    """Return the sorted, unique dates of the timeline events matching all filters.

    The full timeline is queried; set_timestamp() does not limit the result.

    Arguments:
        start_date: Earliest event date to include (inclusive). None: no limit.
        end_date: Latest event date to include (inclusive). None: no limit. A
            date without time of day means midnight, so events later that day
            are excluded.
        wells (Optional[List[WellPath]]): Only include events for these wells.
            When given, schedule keyword, raw text and insert-date events are
            excluded. None or empty: no well filter.
        event_types (Optional[List[str]]): Timeline event types (SEGMENT, PERF,
            VALVE, WSTATE, WTYPE, WELLSPEC, WCONTROL, KEYWORD,
            SCHEDULE_KEYWORD, RAW_TEXT, INSERT_DATE) or keyword names
            (WCONHIST, RPTRST, ...), case-insensitive. None or empty: all
            events.

    Returns:
        List[datetime]: The event dates.

    Example:
        .. code-block:: python

            last_date = timeline.latest_event_date()
            timeline.set_timestamp(timestamp=last_date.isoformat())
    """
    container = self.query_event_dates(
        start_date=_format_query_date(start_date),
        end_date=_format_query_date(end_date),
        well_paths=wells or [],
        event_types=event_types or [],
    )
    values = container.values if container and container.values else []
    return [datetime.fromisoformat(value) for value in values]


@add_method(WellEventTimeline)
def earliest_event_date(
    self: WellEventTimeline,
    start_date: Optional[str | date | datetime] = None,
    end_date: Optional[str | date | datetime] = None,
    wells: Optional[List[WellPath]] = None,
    event_types: Optional[List[str]] = None,
) -> Optional[datetime]:
    """Return the earliest date of the timeline events matching all filters,
    or None when no event matches. See event_dates() for the arguments."""
    dates = event_dates(self, start_date, end_date, wells, event_types)
    return dates[0] if dates else None


@add_method(WellEventTimeline)
def latest_event_date(
    self: WellEventTimeline,
    start_date: Optional[str | date | datetime] = None,
    end_date: Optional[str | date | datetime] = None,
    wells: Optional[List[WellPath]] = None,
    event_types: Optional[List[str]] = None,
) -> Optional[datetime]:
    """Return the latest date of the timeline events matching all filters, or
    None when no event matches. See event_dates() for the arguments."""
    dates = event_dates(self, start_date, end_date, wells, event_types)
    return dates[-1] if dates else None


@add_method(WellEventTimeline)
def wells_with_events(
    self: WellEventTimeline,
    start_date: Optional[str | date | datetime] = None,
    end_date: Optional[str | date | datetime] = None,
    event_types: Optional[List[str]] = None,
) -> List[WellPath]:
    """Return the well paths with timeline events matching all filters, sorted
    by name. See event_dates() for the arguments.

    Matching well names must be unique within the well path collection.
    Raises ValueError if a name is ambiguous or cannot be resolved.

    Returns:
        List[WellPath]: The well paths.
    """
    container = self.query_wells_with_events(
        start_date=_format_query_date(start_date),
        end_date=_format_query_date(end_date),
        event_types=event_types or [],
    )
    names = container.values if container and container.values else []
    if not names:
        return []

    collection = self.ancestor(WellPathCollection)
    if collection is None:
        raise ValueError("Cannot resolve event wells without a well path collection")
    matching_names = set(names)
    well_paths_by_name: Dict[str, WellPath] = {}
    for well_path in collection.descendants(WellPath):
        if well_path.name not in matching_names:
            continue
        if well_path.name in well_paths_by_name:
            raise ValueError(f"Ambiguous well path name: {well_path.name!r}")
        well_paths_by_name[well_path.name] = well_path
    missing_names = matching_names - well_paths_by_name.keys()
    if missing_names:
        raise ValueError(f"Cannot resolve event well names: {sorted(missing_names)!r}")
    return [well_paths_by_name[name] for name in names]


@add_method(WellEventTimeline)
def wells_with_msw_events(self: WellEventTimeline) -> List[WellPath]:
    """Return the well paths with SEGMENT or VALVE events, sorted by name.

    Wells with only PERF events are not included, even if segments were
    configured directly on their well paths; add such MSW wells explicitly
    to export_msw_for_wells.

    Example:
        .. code-block:: python

            schedule_text = timeline.generate_schedule_text(
                eclipse_case=case,
                export_msw_for_wells=timeline.wells_with_msw_events(),
            )
    """
    return wells_with_events(self, event_types=["SEGMENT", "VALVE"])