import math
import random

from sage.base_app import BaseApp


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

        self.iteration = 0
        self.inside_node = self.info["sensors"].index("inside_wrist")
        self.outside_node = self.info["sensors"].index("outside_wrist")

        self.cue_interval = max(float(self.config["cue_interval"]), 0.25)
        self.game_time = max(float(self.config["game_time"]), self.cue_interval)
        self.total_intervals = max(1, int(math.ceil(self.game_time / self.cue_interval)))
        self.actual_game_time = self.total_intervals * self.cue_interval

        self.decision_window = min(
            max(float(self.config["decision_window"]), 0.05), self.cue_interval
        )
        self.points_per_correct = int(self.config["points_per_correct"])
        self.penalty_same_target_miss = int(self.config["penalty_same_target_miss"])
        self.penalty_changed_target_miss = int(
            self.config["penalty_changed_target_miss"]
        )
        self.penalty_repeat_rotation = int(self.config["penalty_repeat_rotation"])
        self.unknown_counts_as_miss = self.config["unknown_counts_as_miss"]
        self.repeat_probability = min(max(float(self.config["repeat_probability"]), 0), 1)

        self.vibration_duration = max(float(self.config["vibration_duration"]), 0.05)
        self.top_margin_threshold = max(float(self.config["top_margin_threshold"]), 0)
        self.repeat_rotation_margin_threshold = max(
            float(self.config["repeat_rotation_margin_threshold"]), 0
        )
        self.orientation_confidence_threshold = max(
            float(self.config["orientation_confidence_threshold"]), 0
        )

        self.inside_axis_key = "AccelZ"
        self.outside_axis_key = "AccelZ"
        self.inside_axis_sign = 1
        self.outside_axis_sign = 1

        start_target = self.config["start_target"]
        self.current_target = (
            random.choice(["inside_wrist", "outside_wrist"])
            if start_target == "random"
            else start_target
        )
        self.previous_target = None
        self.current_interval = 0
        self.running_score = 0
        self.correct_count = 0
        self.same_target_penalty_count = 0
        self.changed_target_miss_count = 0
        self.repeat_rotation_penalty_count = 0
        self.unknown_count = 0

        self.window_sum_inside = 0.0
        self.window_sum_outside = 0.0
        self.window_count = 0
        self.interval_margin_min = None
        self.interval_margin_max = None

        self.last_points_awarded = 0
        self.last_actual_top = "unknown"
        self.last_correct = 0
        self.last_repeated_target = 0
        self.last_cue_changed = 1
        self.last_inside_score = 0.0
        self.last_outside_score = 0.0
        self.last_orientation_margin = 0.0
        self.last_orientation_margin_swing = 0.0
        self.pending_stop = False
        self.first_cue_sent = False

    def run_in_loop(self):
        data = self.my_sage.get_next_data()
        time_now = self.iteration / self.info["datarate"]
        inside_score = self._top_score(
            data[self.inside_node], self.inside_axis_key, self.inside_axis_sign
        )
        outside_score = self._top_score(
            data[self.outside_node], self.outside_axis_key, self.outside_axis_sign
        )
        margin = inside_score - outside_score
        if self.interval_margin_min is None or margin < self.interval_margin_min:
            self.interval_margin_min = margin
        if self.interval_margin_max is None or margin > self.interval_margin_max:
            self.interval_margin_max = margin

        if not self.first_cue_sent:
            self._trigger_cue()
            self.first_cue_sent = True

        interval_start = self.current_interval * self.cue_interval
        interval_elapsed = time_now - interval_start

        if interval_elapsed >= self.cue_interval - self.decision_window:
            self.window_sum_inside += inside_score
            self.window_sum_outside += outside_score
            self.window_count += 1
            self.last_inside_score = inside_score
            self.last_outside_score = outside_score

        if not self.pending_stop and time_now >= (self.current_interval + 1) * self.cue_interval:
            self._finish_interval()
            self.current_interval += 1
            if self.current_interval >= self.total_intervals:
                self.pending_stop = True
                self._stop_all_vibration()
            else:
                self._choose_next_target()
                self._trigger_cue()

        time_remaining = max(self.actual_game_time - time_now, 0)
        user_defined_status = (
            f"Score: {self.running_score} | "
            f"Cue {min(self.current_interval + 1, self.total_intervals)}/{self.total_intervals} | "
            f"Target: {self.current_target} | "
            f"Top: {self.last_actual_top} | "
            f"Time Left: {time_remaining:.1f}s"
        )

        my_data = {
            "time": [time_now],
            "interval_index": [self.current_interval],
            "time_remaining": [time_remaining],
            "running_score": [self.running_score],
            "correct_count": [self.correct_count],
            "same_target_penalty_count": [self.same_target_penalty_count],
            "changed_target_miss_count": [self.changed_target_miss_count],
            "repeat_rotation_penalty_count": [self.repeat_rotation_penalty_count],
            "unknown_count": [self.unknown_count],
            "target_top": [self._target_code(self.current_target)],
            "actual_top": [self._target_code(self.last_actual_top)],
            "cue_changed": [self.last_cue_changed],
            "repeated_target": [self.last_repeated_target],
            "points_awarded": [self.last_points_awarded],
            "correct_this_interval": [self.last_correct],
            "inside_top_score": [self.last_inside_score],
            "outside_top_score": [self.last_outside_score],
            "orientation_margin": [self.last_orientation_margin],
            "orientation_margin_swing": [self.last_orientation_margin_swing],
            "user_defined_status": [user_defined_status],
        }

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

        self.iteration += 1
        if self.pending_stop:
            return False
        return True

    def on_stop_event(self, _):
        self._stop_all_vibration()
        self.logger.info(
            (
                f"Final score: {self.running_score}. "
                f"Correct: {self.correct_count}, "
                f"same-target penalties: {self.same_target_penalty_count}, "
                f"changed-target misses: {self.changed_target_miss_count}, "
                f"repeat-rotation penalties: {self.repeat_rotation_penalty_count}, "
                f"unknown intervals: {self.unknown_count}."
            )
        )

    def _choose_next_target(self):
        self.previous_target = self.current_target
        if random.random() < self.repeat_probability:
            next_target = self.current_target
        else:
            next_target = (
                "outside_wrist"
                if self.current_target == "inside_wrist"
                else "inside_wrist"
            )
        self.current_target = next_target
        self.last_repeated_target = int(self.previous_target == self.current_target)
        self.last_cue_changed = int(not self.last_repeated_target)

    def _trigger_cue(self):
        self._stop_all_vibration()
        target_node = (
            self.inside_node
            if self.current_target == "inside_wrist"
            else self.outside_node
        )
        self.my_sage.sensor_vibration_on(
            target_node,
            self.vibration_duration,
        )

    def _finish_interval(self):
        actual_top, inside_mean, outside_mean, margin = self._decide_top_node()
        margin_swing = self._interval_margin_swing()
        self.last_actual_top = actual_top
        self.last_inside_score = inside_mean
        self.last_outside_score = outside_mean
        self.last_orientation_margin = margin
        self.last_orientation_margin_swing = margin_swing
        self.last_points_awarded = 0
        self.last_correct = 0

        is_correct = actual_top == self.current_target
        repeated_target = self.previous_target == self.current_target
        excessive_repeat_rotation = (
            repeated_target
            and margin_swing > self.repeat_rotation_margin_threshold
        )

        if is_correct:
            points_awarded = self.points_per_correct
            if excessive_repeat_rotation:
                points_awarded += self.penalty_repeat_rotation
                self.repeat_rotation_penalty_count += 1
            self.running_score += points_awarded
            self.correct_count += 1
            self.last_points_awarded = points_awarded
            self.last_correct = 1
        else:
            if actual_top == "unknown":
                self.unknown_count += 1
            if repeated_target and (actual_top != "unknown" or self.unknown_counts_as_miss):
                self.running_score += self.penalty_same_target_miss
                self.same_target_penalty_count += 1
                self.last_points_awarded = self.penalty_same_target_miss
            elif not repeated_target and (
                actual_top != "unknown" or self.unknown_counts_as_miss
            ):
                self.running_score += self.penalty_changed_target_miss
                self.changed_target_miss_count += 1
                self.last_points_awarded = self.penalty_changed_target_miss

        self.window_sum_inside = 0.0
        self.window_sum_outside = 0.0
        self.window_count = 0
        self.interval_margin_min = None
        self.interval_margin_max = None

    def _decide_top_node(self):
        if self.window_count == 0:
            return "unknown", 0.0, 0.0, 0.0

        inside_mean = self.window_sum_inside / self.window_count
        outside_mean = self.window_sum_outside / self.window_count
        margin = inside_mean - outside_mean

        confidence = max(abs(inside_mean), abs(outside_mean))
        if confidence < self.orientation_confidence_threshold:
            return "unknown", inside_mean, outside_mean, margin
        if abs(margin) < self.top_margin_threshold:
            return "unknown", inside_mean, outside_mean, margin
        if margin > 0:
            return "inside_wrist", inside_mean, outside_mean, margin
        return "outside_wrist", inside_mean, outside_mean, margin

    def _interval_margin_swing(self):
        if self.interval_margin_min is None or self.interval_margin_max is None:
            return 0.0
        return self.interval_margin_max - self.interval_margin_min

    def _top_score(self, node_data, axis_key, sign):
        accel_norm = (
            node_data["AccelX"] ** 2
            + node_data["AccelY"] ** 2
            + node_data["AccelZ"] ** 2
        ) ** 0.5
        if accel_norm <= 1e-6:
            return 0.0
        return sign * node_data[axis_key] / accel_norm

    def _stop_all_vibration(self):
        self.my_sage.sensor_vibration_off(self.inside_node)
        self.my_sage.sensor_vibration_off(self.outside_node)

    def _target_code(self, label):
        if label == "inside_wrist":
            return 1
        if label == "outside_wrist":
            return 2
        return 0
