
import polars as pl
from pathlib import Path
from typing import Any
from ..pipeline import PipelineComponent, register_component

from ..pipeline import PipelineComponent, register_component

@register_component("ChronologicalSplitter")
class ChronologicalSplitter(PipelineComponent):
    """
    Splits dataset into Train/Val/Test preserving chronological order.
    Respects trip boundaries (doesn't cut a trip in half).
    """
    def __init__(self, output_dir: str, train_ratio: float = 0.60, val_ratio: float = 0.20):
        self.output_dir = output_dir
        self.train_ratio = train_ratio
        self.val_ratio = val_ratio
        super().__init__(config={"out": output_dir, "ratios": [train_ratio, val_ratio]})

    def fit(self, data: Any) -> 'ChronologicalSplitter':
        return self

    def transform(self, df: pl.DataFrame) -> dict[str, str]:
        """
        Splits DataFrame and saves parts.
        Returns paths to saved files.
        """
        if df.is_empty():
            return {}
            
        if 'trip_id' not in df.columns:
            df = df.with_columns(pl.lit(0).alias("trip_id"))
            
        Path(self.output_dir).mkdir(parents=True, exist_ok=True)
        
        # Trip Stats: start_time, count
        trip_stats = (
            df.group_by("trip_id")
            .agg([
                pl.col("timestamp").min().alias("start_time"),
                pl.len().alias("count")
            ])
            .sort("start_time")
        )
        
        total_rows = trip_stats["count"].sum()
        target_train = total_rows * self.train_ratio
        target_val = total_rows * self.val_ratio
        
        train_trips = []
        val_trips = []
        test_trips = []
        
        current_count = 0
        
        # Iterate (small number of trips usually, loop is invalid overhead)
        for row in trip_stats.iter_rows(named=True):
            trip_id = row['trip_id']
            count = row['count']
            
            if current_count < target_train:
                train_trips.append(trip_id)
            elif current_count < (target_train + target_val):
                val_trips.append(trip_id)
            else:
                test_trips.append(trip_id)
            
            current_count += count
            
        # Save
        parts = {}
        for name, trips in [('train', train_trips), ('val', val_trips), ('test', test_trips)]:
            if not trips:
                continue
                
            # Filter
            subset = df.filter(pl.col("trip_id").is_in(trips))
            
            path = Path(self.output_dir) / f"{name}.parquet"
            subset.write_parquet(path)
            parts[name] = str(path)
            
        return parts
