from sage.base_app import BaseApp
from .Rotation import Rotation as R
from .RelativeAngleCalc import TrunkAngle, IntrinsicZYXEuler
import sage
import time
import numpy as np


def get_rotation(data, node_num):
    return R.from_quat(
        [
            data[node_num]["Quat1"],
            data[node_num]["Quat2"],
            data[node_num]["Quat3"],
            data[node_num]["Quat4"],
        ],
        scalar_first=True,
    )


class Core(BaseApp):
    def __init__(self, my_sage):
        BaseApp.__init__(self, my_sage, __file__)
        self.iteration = 0

        self.l1_sensor = self.info["sensors"].index("l1_sensor")
        self.s1_sensor = self.info["sensors"].index("s1_sensor")
        self.feedback_node = self.info["feedback"].index("feedback")
        self.initial_angle = 0

        self.feedback_threshold = self.config.get("feedback_threshold", 0)
        self.feedback_enabled = self.config.get("feedback_enabled", True)
        self.trunk_angle = TrunkAngle()

        # Variables to track button press logic
        self.button_press_count = 0
        self.button_event_flag = False
        self.max_flex_angle = 0.0
        self.max_extend_angle = 0.0
        self.trunk_flex_percentage = 0.0
        self.sample_angle_data = []
        # get these values from config # Constants
        self.SAMPLE_LIMIT = (
            self.config.get("angle_capture_duration") * self.info["datarate"]
        )
        self.REST_PERIOD_DURATION = self.config.get("rest_period_duration")

        # User prompt message
        self.user_defined_status = "Extend back to max degree and press the button."

        # Register the callback
        my_sage.register_function_button_callback(self.function_button_callback)

        self.flex_angle_range = 0.0
        self.is_rest_period = False
        self.is_calibrated = False

    def function_button_callback(self, event):
        self.button_press_count += int(event == self.my_sage.BUTTON_DOWN)
        self.button_event_flag = True

    def run_in_loop(self):
        data = self.my_sage.get_next_data()
        s1_quat = get_rotation(data, self.s1_sensor)
        l1_quat = get_rotation(data, self.l1_sensor)
        s1_intrinsic = IntrinsicZYXEuler(s1_quat)
        l1_intrinsic = IntrinsicZYXEuler(l1_quat)

        s1_yaw = s1_intrinsic.yaw
        l1_yaw = l1_intrinsic.yaw
        s1_pitch = s1_intrinsic.pitch
        l1_pitch = l1_intrinsic.pitch
        s1_roll = s1_intrinsic.roll
        l1_roll = l1_intrinsic.roll

        if self.iteration == 0:
            self.trunk_angle.calibrate(s1_quat, l1_quat)

        self.trunk_angle.update(s1_quat, l1_quat)

        trunk_flex_angle = self.trunk_angle.flex

        time_now = self.iteration / self.info["datarate"]  # time in seconds
        feedback_state = False

        # Handle button press logic in the loop
        if self.button_event_flag:
            if self.button_press_count == 1:
                if len(self.sample_angle_data) < self.SAMPLE_LIMIT:
                    self.sample_angle_data.append(trunk_flex_angle)
                    self.user_defined_status = (
                        "Button Pressed. Wait for sample to complete."
                    )
                else:
                    self.max_extend_angle = np.mean(self.sample_angle_data)
                    self.sample_angle_data = []
                    self.user_defined_status = (
                        "Flex to max degree and press the button."
                    )
                    self.button_event_flag = False

            elif self.button_press_count == 2:
                if len(self.sample_angle_data) < self.SAMPLE_LIMIT:
                    self.sample_angle_data.append(trunk_flex_angle)
                    self.user_defined_status = (
                        "Button Pressed. Wait for sample to complete."
                    )
                else:
                    self.max_flex_angle = np.mean(self.sample_angle_data)
                    self.sample_angle_data = []
                    self.button_event_flag = False
                    self.flex_angle_range = abs(
                        self.max_flex_angle - self.max_extend_angle
                    )
                    self.user_defined_status = "Rest period. Adjust to normal position."
                    self.is_rest_period = True
                    self.is_calibrated = True
                    self.rest_period_end_time = time.time() + self.REST_PERIOD_DURATION

            else:
                self.button_event_flag = False
        if self.is_calibrated:
            if self.is_rest_period:
                if time.time() >= self.rest_period_end_time:
                    self.is_rest_period = False
                    self.user_defined_status = "Continue with the trial."

            else:
                self.trunk_flex_percentage = (
                    (trunk_flex_angle - self.max_extend_angle)
                    / self.flex_angle_range
                    * 100
                )
                if self.trunk_flex_percentage > self.feedback_threshold:
                    feedback_state = True

                if self.feedback_enabled:
                    self.set_feedback_vibration(
                        feedbackNode=self.feedback_node,
                        duration=1,
                        feedback_state=feedback_state,
                    )

        my_data = {
            "time": [time_now],
            "trunk_flex_angle": [float(trunk_flex_angle)],
            "trunk_flex_percentage": [float(self.trunk_flex_percentage)],
            "feedback_threshold": [float(self.feedback_threshold)],
            "feedback_state": [int(feedback_state)],
            "feedback_enabled": [self.feedback_enabled],
            "s1_yaw": [float(s1_yaw)],
            "l1_yaw": [float(l1_yaw)],
            "s1_pitch": [float(s1_pitch)],
            "l1_pitch": [float(l1_pitch)],
            "s1_roll": [float(s1_roll)],
            "l1_roll": [float(l1_roll)],
            "user_defined_status": [self.user_defined_status],
            "max_flex_angle": [self.max_flex_angle],
            "max_extend_angle": [self.max_extend_angle],
            "flex_angle_range": [self.flex_angle_range],
        }

        self.my_sage.save_data(data, my_data)
        self.my_sage.send_stream_data(data, my_data)

        self.iteration += 1

        return True

    def set_feedback_vibration(self, feedbackNode=0, duration=1, feedback_state=False):
        if feedback_state:
            self.my_sage.feedback_vibration_on(feedbackNode, duration)
        else:
            self.my_sage.feedback_vibration_off(feedbackNode)
