"""
pandas_compat.py - Pandas/NumPy compatibility utilities for OCG

Provides column-oriented result conversion for efficient pandas DataFrame creation
and columnar operations, while maintaining row-oriented as the default OpenCypher standard.
"""

from typing import Dict, List, Any, Optional


def to_columns(rows: List[Dict[str, Any]]) -> Dict[str, List[Any]]:
    """
    Convert row-oriented OpenCypher results to column-oriented format.

    This is useful for pandas DataFrame creation and columnar operations:
    - Efficient vectorized operations (NumPy)
    - Direct DataFrame construction
    - Memory locality for column-based analytics

    Args:
        rows: Row-oriented results from Graph.execute()
              Example: [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]

    Returns:
        Column-oriented dictionary
        Example: {"name": ["Alice", "Bob"], "age": [30, 25]}

    Performance:
        O(rows * columns) — single pass through all data
        Memory: 2x (both formats in memory during conversion)

    Example:
        >>> from ocg import Graph
        >>> from ocg.pandas_compat import to_columns
        >>>
        >>> graph = Graph()
        >>> graph.execute("CREATE (a:Person {name: 'Alice', age: 30})")
        >>> graph.execute("CREATE (b:Person {name: 'Bob', age: 25})")
        >>>
        >>> # Standard row-oriented (OpenCypher standard)
        >>> rows = graph.execute("MATCH (p:Person) RETURN p.name AS name, p.age AS age")
        >>> for row in rows:
        ...     print(f"{row['name']} is {row['age']}")
        Alice is 30
        Bob is 25
        >>>
        >>> # Column-oriented for pandas
        >>> cols = to_columns(rows)
        >>> import pandas as pd
        >>> df = pd.DataFrame(cols)
        >>> print(df)
             name  age
        0  Alice   30
        1    Bob   25
    """
    if not rows:
        return {}

    # Get all unique keys (union of keys from all rows, in case some rows have missing columns)
    all_keys = set()
    for row in rows:
        all_keys.update(row.keys())

    # Build column arrays
    columns = {key: [] for key in all_keys}
    for row in rows:
        for key in all_keys:
            columns[key].append(row.get(key, None))  # None for missing values

    return columns


def to_rows(columns: Dict[str, List[Any]]) -> List[Dict[str, Any]]:
    """
    Convert column-oriented format back to row-oriented OpenCypher results.

    Useful when you've processed data in columnar format (pandas/NumPy)
    and need to return to row-oriented for further OpenCypher operations.

    Args:
        columns: Column-oriented dictionary
                 Example: {"name": ["Alice", "Bob"], "age": [30, 25]}

    Returns:
        Row-oriented list of dictionaries
        Example: [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]

    Raises:
        ValueError: If column lengths are inconsistent

    Example:
        >>> cols = {"name": ["Alice", "Bob"], "age": [30, 25]}
        >>> rows = to_rows(cols)
        >>> for row in rows:
        ...     print(row)
        {'name': 'Alice', 'age': 30}
        {'name': 'Bob', 'age': 25}
    """
    if not columns:
        return []

    # Validate all columns have the same length
    lengths = {key: len(values) for key, values in columns.items()}
    unique_lengths = set(lengths.values())
    if len(unique_lengths) > 1:
        raise ValueError(
            f"Inconsistent column lengths: {lengths}. "
            "All columns must have the same number of rows."
        )

    num_rows = lengths[next(iter(lengths))] if lengths else 0

    # Build row dictionaries
    rows = []
    for i in range(num_rows):
        row = {key: values[i] for key, values in columns.items()}
        rows.append(row)

    return rows


def to_dataframe(rows: List[Dict[str, Any]], copy: bool = False):
    """
    Convert OpenCypher results directly to pandas DataFrame.

    This is a convenience wrapper around to_columns() that creates
    a DataFrame in one step. Requires pandas to be installed.

    Args:
        rows: Row-oriented OpenCypher results
        copy: If True, make a copy of the data (default False for efficiency)

    Returns:
        pandas.DataFrame with columns from the Cypher result

    Raises:
        ImportError: If pandas is not installed

    Performance:
        - No copy (default): Zero-copy construction where possible
        - With copy: Ensures data independence from original results

    Example:
        >>> from ocg import Graph
        >>> from ocg.pandas_compat import to_dataframe
        >>>
        >>> graph = Graph()
        >>> graph.execute("CREATE (a:Person {name: 'Alice', age: 30})")
        >>> graph.execute("CREATE (b:Person {name: 'Bob', age: 25})")
        >>>
        >>> rows = graph.execute("MATCH (p:Person) RETURN p.name, p.age ORDER BY p.age")
        >>> df = to_dataframe(rows)
        >>> print(df)
          p.name  p.age
        0    Bob     25
        1  Alice     30
        >>>
        >>> # Use pandas operations
        >>> print(df['p.age'].mean())
        27.5
    """
    try:
        import pandas as pd
    except ImportError as e:
        raise ImportError(
            "pandas is required for to_dataframe(). "
            "Install with: pip install pandas"
        ) from e

    if not rows:
        return pd.DataFrame()

    # Convert to columns first for efficient DataFrame construction
    columns = to_columns(rows)
    return pd.DataFrame(columns, copy=copy)


class ResultAdapter:
    """
    Adapter that provides both row and column-oriented views of OpenCypher results.

    This allows flexible access patterns without storing data twice:
    - Access as rows (lazy, no conversion)
    - Access as columns (converts on first access, then cached)
    - Convert to DataFrame (converts on first access, then cached)

    Memory-efficient: Only one conversion is performed and cached.

    Example:
        >>> from ocg import Graph
        >>> from ocg.pandas_compat import ResultAdapter
        >>>
        >>> graph = Graph()
        >>> graph.execute("CREATE (a:Person {name: 'Alice', age: 30})")
        >>>
        >>> rows = graph.execute("MATCH (p:Person) RETURN p.name, p.age")
        >>> result = ResultAdapter(rows)
        >>>
        >>> # Iterate as rows (no conversion)
        >>> for row in result.rows:
        ...     print(row['p.name'])
        Alice
        >>>
        >>> # Access as columns (converts and caches)
        >>> print(result.columns['p.age'])
        [30]
        >>>
        >>> # Get DataFrame (uses cached columns)
        >>> df = result.to_dataframe()
        >>> print(df)
          p.name  p.age
        0  Alice     30
    """

    def __init__(self, rows: List[Dict[str, Any]]):
        """
        Initialize with row-oriented OpenCypher results.

        Args:
            rows: Row-oriented results from Graph.execute()
        """
        self._rows = rows
        self._columns: Optional[Dict[str, List[Any]]] = None
        self._dataframe = None

    @property
    def rows(self) -> List[Dict[str, Any]]:
        """Get row-oriented view (original format, no conversion)."""
        return self._rows

    @property
    def columns(self) -> Dict[str, List[Any]]:
        """Get column-oriented view (lazy conversion, cached)."""
        if self._columns is None:
            self._columns = to_columns(self._rows)
        return self._columns

    def to_dataframe(self, copy: bool = False):
        """
        Convert to pandas DataFrame (lazy conversion, cached).

        Args:
            copy: If True, make a copy of the data (default False)

        Returns:
            pandas.DataFrame
        """
        if self._dataframe is None:
            self._dataframe = to_dataframe(self._rows, copy=copy)
        return self._dataframe if not copy else self._dataframe.copy()

    def __len__(self) -> int:
        """Return number of rows."""
        return len(self._rows)

    def __iter__(self):
        """Iterate over rows."""
        return iter(self._rows)

    def __getitem__(self, index: int) -> Dict[str, Any]:
        """Get row by index."""
        return self._rows[index]


# Convenience functions for Graph class extension
def add_pandas_methods(graph_class):
    """
    Add pandas-compatible methods to the Graph class.

    This is a decorator that extends the Graph class with convenience methods:
    - graph.execute_df() -> returns DataFrame directly
    - graph.execute_columns() -> returns column-oriented dict

    Example:
        >>> from ocg import Graph
        >>> from ocg.pandas_compat import add_pandas_methods
        >>>
        >>> # Extend Graph class
        >>> Graph = add_pandas_methods(Graph)
        >>>
        >>> graph = Graph()
        >>> graph.execute("CREATE (a:Person {name: 'Alice', age: 30})")
        >>>
        >>> # Now you can use the new methods
        >>> df = graph.execute_df("MATCH (p:Person) RETURN p.name, p.age")
        >>> print(df)
          p.name  p.age
        0  Alice     30
    """

    def execute_df(self, query: str, copy: bool = False):
        """
        Execute OpenCypher query and return pandas DataFrame.

        Args:
            query: OpenCypher query string
            copy: If True, make a copy of the data

        Returns:
            pandas.DataFrame
        """
        rows = self.execute(query)
        return to_dataframe(rows, copy=copy)

    def execute_columns(self, query: str) -> Dict[str, List[Any]]:
        """
        Execute OpenCypher query and return column-oriented results.

        Args:
            query: OpenCypher query string

        Returns:
            Dictionary mapping column names to lists of values
        """
        rows = self.execute(query)
        return to_columns(rows)

    # Add methods to the class
    graph_class.execute_df = execute_df
    graph_class.execute_columns = execute_columns

    return graph_class


__all__ = [
    "to_columns",
    "to_rows",
    "to_dataframe",
    "ResultAdapter",
    "add_pandas_methods",
]
