import torch
import pandas as pd
from base.data.base_dataset import BaseDataset


class Arpl(BaseDataset):
    """
    The Dataset class to handle the csv generated by acquisition by the ARPL lab,
    NYU. The original data is stored in RosBags and the .csv files are generated
    with the python script "arpl_stand_alone_preprocess.py". Setting up the
    script is possible to obtain a  csv with a lot of measurements:
    - CSV labels:
            - Linear Velocities: "v_x_t", "v_y_t", "v_z_t"
            - Quaternions: "q_w_t", "q_x_t", "q_y_t", "q_z_t"
            - Angular Velocities: "w_x_t", "w_y_t", "w_z_t",
            - Motors forces: "f_0_t", "f_1_t", "f_2_t", "f_3_t"
            - Linear Accelerations: "vdot_x", "vdot_y", "vdot_z"
            - Angular Accelerations: "wdot_x", "wdot_y", "wdot_z"
            - Residual linear acceleration: "res_vdot_x", "res_vdot_y", "res_vdot_z
            - Nominal linear, angular velocities: "vdot_nom", "wdot_nom"

    The regressor in this framework work just with single vectors for samples and labels, so you can just pass as
    argument the desired column labels.
    If Some special operations are needed, you can simply create a new class for the desired data:
     1) extending the ARPL class,
     2) implementing the _extra_operations method.
    """

    def __init__(self, alias: str, csv_path: str,
                 x_col_label: str, y_col_label: str, t_col_label: str, start_ix: int, end_ix: int):
        self.x_column = x_col_label
        self.y_column = y_col_label
        self.t_column = t_col_label
        self.csv_path = csv_path
        alias = ''.join(['ARPL_', alias])
        super().__init__(dataset_label=alias)
        self.select_subset(start_ix, end_ix)

    def _init_dataset(self):
        data_frame = pd.read_csv(self.csv_path)
        df_x = data_frame[self.x_column]
        df_y = data_frame[self.y_column]
        df_t = data_frame[self.t_column]
        self.X = torch.tensor(df_x.to_numpy(), dtype=torch.float)
        self.Y = torch.tensor(df_y.to_numpy(), dtype=torch.float)
        self.timestamp = torch.tensor(df_t.to_numpy(), dtype=torch.float)
