"""DataFrame export, NumPy vectorized ops, CSV/Parquet/Arrow export.

Adds to_numpy(), to_dataframe(), to_csv(), to_parquet(), to_arrow()
to Trace and SignalRef objects.
"""

from __future__ import annotations
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from ewal.trace import Trace
    from ewal.signal import SignalRef


def signal_to_numpy(trace: "Trace", signal_path: str):
    """Get a signal's values across all timesteps as a numpy array."""
    import numpy as np
    n = trace.max_index + 1
    vals = [trace._get_signal_value(signal_path, i) for i in range(n)]
    # Convert to numeric where possible
    numeric = []
    for v in vals:
        if v is None:
            numeric.append(0)
        elif isinstance(v, (int, float)):
            numeric.append(v)
        else:
            try:
                numeric.append(int(v))
            except (ValueError, TypeError):
                numeric.append(0)
    return np.array(numeric, dtype=np.uint64)


def trace_to_dataframe(trace: "Trace", signals: Optional[list[str]] = None):
    """Export trace data as a pandas DataFrame.

    Args:
        trace: The Trace object
        signals: Optional list of signal names. If None, exports all.

    Returns:
        pandas.DataFrame with columns: index, timestamp, <signal_names>...
    """
    import pandas as pd
    import numpy as np

    n = trace.max_index + 1

    if signals is None:
        signals = [s for s in trace._core.signal_names()
                   if s not in trace._virtual_signals]

    data = {
        "index": list(range(n)),
        "timestamp": [trace._core.timestamp(i) for i in range(n)],
    }

    for sig in signals:
        vals = []
        for i in range(n):
            v = trace._get_signal_value(sig, i)
            vals.append(v if v is not None else 0)
        data[sig] = vals

    return pd.DataFrame(data)


def trace_to_csv(trace: "Trace", path: str, signals: Optional[list[str]] = None):
    """Export trace data to a CSV file."""
    df = trace_to_dataframe(trace, signals)
    df.to_csv(path, index=False)


def trace_to_parquet(trace: "Trace", path: str, signals: Optional[list[str]] = None):
    """Export trace data to a Parquet file."""
    df = trace_to_dataframe(trace, signals)
    df.to_parquet(path, index=False)


def trace_to_arrow(trace: "Trace", signals: Optional[list[str]] = None):
    """Export trace data as a PyArrow Table."""
    import pyarrow as pa
    df = trace_to_dataframe(trace, signals)
    return pa.Table.from_pandas(df)


# ── Monkey-patch these onto Trace and SignalRef ──

def _patch_trace_exports():
    """Add export methods to Trace class."""
    from ewal.trace import Trace
    from ewal.signal import SignalRef

    def _to_dataframe(self, signals=None):
        return trace_to_dataframe(self, signals)
    Trace.to_dataframe = _to_dataframe

    def _to_csv(self, path, signals=None):
        return trace_to_csv(self, path, signals)
    Trace.to_csv = _to_csv

    def _to_parquet(self, path, signals=None):
        return trace_to_parquet(self, path, signals)
    Trace.to_parquet = _to_parquet

    def _to_arrow(self, signals=None):
        return trace_to_arrow(self, signals)
    Trace.to_arrow = _to_arrow

    def _sig_to_numpy(self):
        return signal_to_numpy(self._trace, self._path)
    SignalRef.to_numpy = _sig_to_numpy


# Auto-patch when imported
_patch_trace_exports()
