import pandas as pd
import ast
from pathlib import Path
from typing import Union
import random

class CabinDataset:
    """
    Iterable dataset for the Aircraft Cabin Localization data.
    
    - Supports 'pickle' (fast) and 'csv' (requires parsing) backends.
    - Can load specific subsets: 'train', 'test', or 'all' (no suffix).
    - Can load specific sensors: 'ble', 'uwb', or 'fused'.
    - Automatically parses stringified lists in CSV mode.
    """

    SUPPORTED_BACKENDS = {"pickle", "csv"}
    LIST_COLUMNS = ["anchors", "measured_range", "true_range", "residual_range"]

    def __init__(
        self,
        folder: Union[str, Path] = "../data",  # Default assumes running from scripts/
        subset: str = "train",                 # Options: 'train', 'test', '' (for all)
        sensor: str = "uwb",                   # Options: 'uwb', 'ble', 'fused'
        backend: str = "pickle",               # Options: 'pickle', 'csv'
    ):
        self.folder = Path(folder)
        self.subset = subset
        self.sensor = sensor
        self.backend = backend.lower()

        if self.backend not in self.SUPPORTED_BACKENDS:
            raise ValueError(f"Unsupported backend '{backend}'. Choose {self.SUPPORTED_BACKENDS}")

        # Construct file name (e.g., 'uwb_train.pkl' or 'ble.csv')
        suffix = f"_{subset}" if subset else ""
        ext = "pkl" if self.backend == "pickle" else "csv"
        filename = f"{sensor}{suffix}.{ext}"
        
        # Path: data/pickle/uwb_train.pkl
        self.file_path = self.folder / self.backend / filename

        if not self.file_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.file_path}")

        print(f"Loading {self.sensor.upper()} data from: {self.file_path} ...")
        self.data = self._load_data()
        print(f"Loaded {len(self.data)} measurements.")

    def _load_data(self):
        """Loads data and handles CSV list parsing if necessary."""
        if self.backend == "pickle":
            return pd.read_pickle(self.file_path)
        
        # CSV handling: parse stringified lists into actual lists
        elif self.backend == "csv":
            # helper to eval strings
            def parse_list(x):
                try:
                    return ast.literal_eval(x)
                except (ValueError, SyntaxError):
                    return x # Return raw if parsing fails

            converters = {col: parse_list for col in self.LIST_COLUMNS}
            return pd.read_csv(self.file_path, converters=converters)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx: int):
        """Returns a dictionary representation of the row."""
        row = self.data.iloc[idx]
        return row.to_dict()

    def __iter__(self):
        """Iterates through the dataset rows."""
        for i in range(len(self)):
            yield self[i]

if __name__ == "__main__":
    
    sensors = ["uwb", "ble"]
    subset = "train"
    backend = "pickle"  # Change to 'csv' to test CSV loading

    for sensor in sensors:
        print(f"\n{'='*50}")
        print(f"  READING SENSOR: {sensor.upper()}")
        print(f"{'='*50}")

        try:
            # Initialize Dataset
            ds = CabinDataset(folder="data", subset=subset, sensor=sensor, backend=backend)

            # Get 3 random samples from the dataset
            random_indices = random.sample(range(len(ds)), 3)
            
            # Iterate through a few examples
            print(f"\n--- Showing 3 random samples from {sensor.upper()} ---")
            for i in random_indices:
                item = ds[i]
                print(f"\nMeasurement ID: {item['mess_id']}")
                print(f"Timestamp:      {item['timestamp']}")
                print(f"Anchors:        {item['anchors']}")
                print(f"Ranges:         {item['measured_range']}")
                print(f"Ground Truth:   ({item['ref_x']}, {item['ref_y']})")
                print(f"True Ranges:         {item['true_range']}")
                print(f"Residual Ranges:         {item['residual_range']}")
                
        
        except FileNotFoundError as e:
            print(f"\n[!] Error loading {sensor}: {e}")
        except Exception as e:
            print(f"\n[!] Unexpected error: {e}")

    print(f"\n{'='*50}")
    print("Done.")