import numpy as np

from typing import Literal
from numpy.typing import NDArray

from spatz.simulations.csv_source import CSVSource

class TeleMega(CSVSource):
    def __init__(self, path: str, interpolation: Literal['linear'] = 'linear') -> None:
        super().__init__(path, 'time', interpolation)

    def get_altitude(self) -> float:
        return self.fetch_value('altitude')

    def get_acceleration(self, frame: Literal['global', 'local']) -> NDArray:
        acc_local = self.fetch_values(['accel_x', 'accel_y', 'accel_z'])

        if frame == 'global':
            raise NotImplementedError()
        
        return acc_local
    
    def get_angular_velocity(self) -> NDArray:
        return self.fetch_values(['gyro_roll', 'gyro_pitch', 'gyro_yaw'])
    
    def get_static_pressure(self) -> float:
        return self.fetch_value('pressure')