"""Parser utils for log file data."""

from pathlib import Path

import numpy as np
import pandas as pd
from loguru import logger
from matplotlib import pyplot as plt

from adaptive_oscillator.base_classes import (
    AngleXYZ,
    Body,
    Joint,
    LeftRight,
    Limb,
    Quaternion,
    SensorFile,
    VectorXYZ,
)
from adaptive_oscillator.definitions import (
    ALPHA,
    ANGLES_SEGMENT_FIELDS,
    FIG_SIZE,
    IMU_SEGMENT_FIELDS,
    QUATERNION_SEGMENT_FIELDS,
    AnglesHeader,
    IMUHeader,
    Joints,
    LogFileKeys,
    QuaternionHeader,
    Segments,
)
from adaptive_oscillator.utils import time_str_to_seconds


class LogFiles:
    """Main entry point for accessing all sensor log files."""

    def __init__(self, base_path: str | Path) -> None:
        self._path = Path(base_path)
        if not self._path.is_dir():
            msg = f"Path '{self._path}' does not exist."
            logger.error(msg)
            raise FileNotFoundError(msg)
        self.accel = SensorFile(LogFileKeys.ACCEL, self._path)
        self.angle = SensorFile(LogFileKeys.ANGLE, self._path)
        self.gravity = SensorFile(LogFileKeys.GRAVITY, self._path)
        self.gyro = SensorFile(LogFileKeys.GYRO, self._path)
        self.quat = SensorFile(LogFileKeys.QUAT, self._path)

    def __repr__(self) -> str:  # pragma: no cover
        """Return a string representation of the LogFiles object."""
        return (
            f"Log files for dir: '{self._path}'"
            f"\n\t{self.accel.left}"
            f"\n\t{self.accel.right}"
            f"\n\t{self.angle.left}"
            f"\n\t{self.angle.right}"
            f"\n\t{self.gravity.left}"
            f"\n\t{self.gravity.right}"
            f"\n\t{self.gyro.left}"
            f"\n\t{self.gyro.right}"
            f"\n\t{self.quat.left}"
            f"\n\t{self.quat.right})"
        )

    def plot(self, euler_only: bool = False, add_offset: bool = False) -> None:
        """Plot log files.

        :param euler_only: Plot only euler angles.
        :param add_offset: Add offset to angles.
        :return: None
        """
        logger.info("Plotting log file data.")

        for side in ["left", "right"]:
            if not euler_only:
                accel_data = IMUParser(getattr(self.accel, side))
                accel_data.parse()
                accel_data.plot(y_label="Acceleration (m/s2)")

                gyro_data = IMUParser(getattr(self.gyro, side))
                gyro_data.parse()
                gyro_data.plot(y_label="Angular Velocity (deg/s)")

                quat_data = QuaternionParser(getattr(self.quat, side))
                quat_data.parse()
                quat_data.plot()

            angle_data = AngleParser(getattr(self.angle, side))
            angle_data.parse(add_offset=add_offset)
            angle_data.plot(y_label="Euler Angle (deg)")


class IMUParser:
    """Parser for log files with limb information."""

    def __init__(self, filepath: Path):
        self.filepath = filepath
        self.time = np.array([])
        self.pelvis = VectorXYZ()
        self.upper_leg = VectorXYZ()
        self.lower_leg = VectorXYZ()
        self.foot = VectorXYZ()

    def parse(self):
        """Parse the log file and return a DataFrame."""
        raw_data = pd.read_csv(self.filepath, sep="\t+", engine="python")
        logger.debug(f"Parsing {self.filepath}")
        logger.debug(f"Columns: {raw_data.shape}")
        time_str = raw_data[IMUHeader.TIME]
        self.time = np.array([time_str_to_seconds(t) for t in time_str])

        for segment_name, fields in IMU_SEGMENT_FIELDS.items():
            x = raw_data[fields[0]].to_numpy()
            y = raw_data[fields[1]].to_numpy()
            z = raw_data[fields[2]].to_numpy()
            setattr(self, segment_name, VectorXYZ(x, y, z))

    def plot(self, y_label: str):  # pragma: no cover
        """Plot the x, y, z data."""
        _, ax = plt.subplots(figsize=FIG_SIZE, sharex=True, nrows=4, ncols=1)

        for ii, (name, segment) in enumerate(
            zip(
                [
                    Segments.PELVIS,
                    Segments.UPPER_LEG,
                    Segments.LOWER_LEG,
                    Segments.FOOT,
                ],
                [self.pelvis, self.upper_leg, self.lower_leg, self.foot],
            )
        ):
            time = self.time - self.time[0]
            for axis in ["x", "y", "z"]:
                imu_signal = getattr(segment, axis)
                ax[ii].plot(time, imu_signal, label=f"{name}-{axis}", alpha=ALPHA)
            ax[ii].set_title(f"{name} - {self.filepath.stem}")
            ax[ii].set_xlabel("Time (s)")
            ax[ii].set_ylabel(y_label)
            ax[ii].legend(loc="upper right")
            ax[ii].grid(True)
            plt.tight_layout()


class AngleParser:
    """Parser for log files with angle."""

    def __init__(self, filepath: Path):
        self.filepath = filepath
        self.time = np.array([])
        self.hip = AngleXYZ()
        self.knee = AngleXYZ()
        self.ankle = AngleXYZ()
        self.side: str = ""

        if "left" in self.filepath.stem:
            self.side = "left"
        elif "right" in self.filepath.stem:
            self.side = "right"
        else:
            logger.warning(f"Words 'left' or 'right' not found in {self.filepath}")

    def parse(self, add_offset: bool = False):
        """Parse the log file and return a DataFrame."""
        raw_data = pd.read_csv(self.filepath, sep="\t+", engine="python")
        logger.debug(f"Parsing {self.filepath}")
        logger.debug(f"Columns: {raw_data.shape}")

        time_str = raw_data[AnglesHeader.TIME]
        self.time = np.array([time_str_to_seconds(t) for t in time_str])

        for segment_name, fields in ANGLES_SEGMENT_FIELDS.items():
            x_deg = raw_data[fields[0]].to_numpy()
            y_deg = raw_data[fields[1]].to_numpy()
            z_deg = raw_data[fields[2]].to_numpy()
            setattr(self, segment_name, AngleXYZ(x_deg, y_deg, z_deg))

        if add_offset:
            logger.info(f"Adding offset to {self.filepath}")
            self.hip.add_offset()
            self.knee.add_offset()
            self.ankle.add_offset()

    def plot(self, y_label: str) -> None:
        """Plot the x, y, z data.

        :param y_label: label for the y-axis
        :return: None
        """
        _, ax = plt.subplots(figsize=FIG_SIZE, sharex=True, nrows=3, ncols=1)

        for ii, (name, segment) in enumerate(
            zip(
                [
                    Joints.HIP,
                    Joints.KNEE,
                    Joints.ANKLE,
                ],
                [self.hip, self.knee, self.ankle],
            )
        ):
            time = self.time - self.time[0]
            for axis in ["x", "y", "z"]:
                angle = getattr(segment, axis)
                ax[ii].plot(time, angle, label=axis, alpha=ALPHA)
            ax[ii].set_title(f"{name} - {self.filepath.stem}")
            ax[ii].set_xlabel("Time (s)")
            ax[ii].set_ylabel(y_label)
            ax[ii].legend(loc="upper right")
            ax[ii].grid(True)
            plt.tight_layout()


class QuaternionParser:
    """Parser for log files with quaternion information."""

    def __init__(self, filepath: Path):
        self.filepath = filepath
        self.time = np.array([])
        self.pelvis = Quaternion()
        self.upper_leg = Quaternion()
        self.lower_leg = Quaternion()
        self.foot = Quaternion()

    def parse(self):
        """Parse the log file and return a DataFrame."""
        try:
            raw_data = pd.read_csv(self.filepath, sep="\t")
            logger.debug(f"Parsing {self.filepath}")
            logger.debug(f"Columns: {raw_data.shape}")

            time_str = raw_data[QuaternionHeader.TIME]
            self.time = np.array([time_str_to_seconds(t) for t in time_str])

            for segment_name, fields in QUATERNION_SEGMENT_FIELDS.items():
                w = raw_data[fields[0]].to_numpy()
                x = raw_data[fields[1]].to_numpy()
                y = raw_data[fields[2]].to_numpy()
                z = raw_data[fields[3]].to_numpy()
                setattr(self, segment_name, Quaternion(w, x, y, z))
        except FileNotFoundError:
            logger.error(f"File not found: {self.filepath}")

    def plot(self):  # pragma: no cover
        """Plot the Quaternion data."""
        fig, ax = plt.subplots(figsize=FIG_SIZE, sharex=True, nrows=4, ncols=1)

        try:
            for ii, (name, segment) in enumerate(
                zip(
                    [
                        Segments.PELVIS,
                        Segments.UPPER_LEG,
                        Segments.LOWER_LEG,
                        Segments.FOOT,
                    ],
                    [self.pelvis, self.upper_leg, self.lower_leg, self.foot],
                )
            ):
                time = self.time - self.time[0]
                for axis in ["w", "x", "y", "z"]:
                    quat_component = getattr(segment, axis)
                    ax[ii].plot(
                        time, quat_component, label=f"{name}-{axis}", alpha=ALPHA
                    )
                ax[ii].set_title(f"{name} Orientation")
                ax[ii].set_xlabel("Time (s)")
                ax[ii].set_ylabel("Quaternion")
                ax[ii].legend(loc="upper right")
                ax[ii].grid(True)
                plt.tight_layout()
            plt.show()
        except Exception as err:
            logger.error(f"Exception: '{err}' for '{self.filepath}'.")


class LogParser:
    """Parser for log files with limb information."""

    def __init__(self, log_files: LogFiles, add_offset: bool = False):
        """Parse the log files.

        :param log_files: LogFiles
        :param add_offset: bool
        :return: None
        """
        logger.info(f"Parsing {log_files}")
        accel_data_right = IMUParser(log_files.accel.right)
        accel_data_right.parse()
        accel_data_left = IMUParser(log_files.accel.left)
        accel_data_left.parse()

        gyro_data_right = IMUParser(log_files.gyro.right)
        gyro_data_right.parse()
        gyro_data_left = IMUParser(log_files.gyro.left)
        gyro_data_left.parse()

        quat_data_right = QuaternionParser(log_files.quat.right)
        quat_data_right.parse()
        quat_data_left = QuaternionParser(log_files.quat.left)
        quat_data_left.parse()

        angles_right = AngleParser(log_files.angle.right)
        angles_right.parse(add_offset=add_offset)
        angles_left = AngleParser(log_files.angle.left)
        angles_left.parse(add_offset=add_offset)

        time = accel_data_right.time

        pelvis_right = Limb(
            time=time,
            accel=accel_data_right.pelvis,
            gyro=gyro_data_right.pelvis,
            quat=quat_data_right.pelvis,
        )
        upper_leg_right = Limb(
            time=time,
            accel=accel_data_right.upper_leg,
            gyro=gyro_data_right.upper_leg,
            quat=quat_data_right.upper_leg,
        )
        lower_leg_right = Limb(
            time=time,
            accel=accel_data_right.lower_leg,
            gyro=gyro_data_right.lower_leg,
            quat=quat_data_right.lower_leg,
        )
        foot_right = Limb(
            time=time,
            accel=accel_data_right.foot,
            gyro=gyro_data_right.foot,
            quat=quat_data_right.foot,
        )
        hip_right = Joint(time=time, angles=angles_right.hip)
        knee_right = Joint(time=time, angles=angles_right.knee)
        ankle_right = Joint(time=time, angles=angles_right.ankle)

        pelvis_left = Limb(
            time=time,
            accel=accel_data_left.pelvis,
            gyro=gyro_data_left.pelvis,
            quat=quat_data_left.pelvis,
        )
        upper_leg_left = Limb(
            time=time,
            accel=accel_data_left.upper_leg,
            gyro=gyro_data_left.upper_leg,
            quat=quat_data_left.upper_leg,
        )
        lower_leg_left = Limb(
            time=time,
            accel=accel_data_left.lower_leg,
            gyro=gyro_data_left.lower_leg,
            quat=quat_data_left.lower_leg,
        )
        foot_left = Limb(
            time=time,
            accel=accel_data_left.foot,
            gyro=gyro_data_left.foot,
            quat=quat_data_left.foot,
        )
        hip_left = Joint(time=time, angles=angles_left.hip)
        knee_left = Joint(time=time, angles=angles_left.knee)
        ankle_left = Joint(time=time, angles=angles_left.ankle)

        self.log_files = log_files
        self.time = accel_data_right.time
        self.data = LeftRight(
            left=Body(
                pelvis=pelvis_left,
                upper_leg=upper_leg_left,
                lower_leg=lower_leg_left,
                foot=foot_left,
                hip=hip_left,
                knee=knee_left,
                ankle=ankle_left,
            ),
            right=Body(
                pelvis=pelvis_right,
                upper_leg=upper_leg_right,
                lower_leg=lower_leg_right,
                foot=foot_right,
                hip=hip_right,
                knee=knee_right,
                ankle=ankle_right,
            ),
        )
