from __future__ import annotations

import unittest

import numpy as np

from pong_policy import (
    DOWN,
    NOOP,
    TARGET_MAX_Y_RAM,
    TARGET_MIN_Y_RAM,
    UP,
    INACTIVE_CENTER_Y,
    LOW_SERVE_TARGET_Y,
    LATE_LOW_LOOP_CORRECTION_LIMIT,
    LOW_MID_SCORE00_CORRECTION_LIMIT,
    LOW_UP_Y111_CORRECTION_LIMIT,
    RamPolicy,
    is_high_down_score40_bounce_down_correction_state,
    is_high_down_score80_bounce_down_correction_state,
    is_high_down_score120_bounce_down_correction_state,
    is_high_down_score160_bounce_down_correction_state,
    is_high_down_setup_noop_correction_state,
    is_late_low_loop_pre_up_correction_state,
    is_low_mid_score00_down_correction_state,
    is_low_serve_late_noop_correction_state,
    is_low_up_score01_down_correction_state,
    is_low_up_score10_final_down_correction_state,
    is_low_up_y111_pre_up_correction_state,
    is_mid_down_precontact_up_correction_state,
    is_mid_up_precontact_up_correction_state,
    is_opening_low_dy8_correction_state,
    is_score00_mid_down_y154_final_up_correction_state,
    is_low_serve_pattern,
    is_chain2_score00_initial_down_correction_state,
    is_y86_score22_setup_up_correction_state,
    is_y86_score42_setup_up_correction_state,
    is_y86_score62_setup_up_correction_state,
    is_y86_score82_setup_up_correction_state,
    is_y86_score20_serve_initial_down_correction_state,
    is_y86_score60_serve_initial_down_correction_state,
    is_y86_score100_serve_initial_down_correction_state,
    is_y86_score140_serve_initial_down_correction_state,
    is_y86_bounce_score22_x166_down_correction_state,
    is_y86_bounce_x166_down_correction_state,
    reflected_y,
    select_action,
    select_action_from_state,
    target_y_from_state,
)
from pong_state import PongState, detect_state, ram_ball_to_screen, ram_paddle_y_to_screen


def info_with_ram(**values: int) -> dict:
    ram = np.zeros((1, 128), dtype=np.uint8)
    for index, value in values.items():
        ram[0, int(index)] = value
    return {"ram": ram}


class StatePolicyTests(unittest.TestCase):
    def test_detect_state_reads_known_ram_bytes(self) -> None:
        state = detect_state(
            info_with_ram(**{"13": 2, "14": 5, "49": 140, "50": 88, "51": 120, "54": 130})
        )

        self.assertEqual(state.ball_x, 140)
        self.assertEqual(state.ball_y, 130)
        self.assertEqual(state.opponent_paddle_y, 88)
        self.assertEqual(state.self_paddle_y, 120)
        self.assertTrue(state.ball_active)
        self.assertTrue(state.ball_in_right_half)
        self.assertEqual(state.opponent_score, 2)
        self.assertEqual(state.self_score, 5)

    def test_policy_tracks_active_ball_in_right_half(self) -> None:
        self.assertEqual(
            select_action_from_state(
                PongState(140, 80, 120, 100, ball_active=True, ball_in_right_half=True)
            ),
            UP,
        )
        self.assertEqual(
            select_action_from_state(
                PongState(140, 160, 120, 100, ball_active=True, ball_in_right_half=True)
            ),
            DOWN,
        )

    def test_policy_centers_when_ball_inactive_or_left(self) -> None:
        self.assertEqual(
            select_action_from_state(
                PongState(80, 180, 160, 100, ball_active=True, ball_in_right_half=False)
            ),
            UP,
        )
        self.assertEqual(
            select_action_from_state(
                PongState(205, 0, 100, 100, ball_active=False, ball_in_right_half=True)
            ),
            NOOP,
        )

    def test_inactive_ball_uses_serve_ready_center(self) -> None:
        state = PongState(205, 0, 140, 100, ball_active=False, ball_in_right_half=True)

        self.assertEqual(target_y_from_state(state), INACTIVE_CENTER_Y)

    def test_select_action_accepts_envpool_info(self) -> None:
        self.assertEqual(select_action(info_with_ram(**{"49": 140, "51": 121, "54": 123})), NOOP)

    def test_policy_predicts_right_moving_intercept(self) -> None:
        previous = PongState(100, 100, 140, 100, ball_active=True, ball_in_right_half=False)
        current = PongState(110, 110, 140, 100, ball_active=True, ball_in_right_half=False)

        self.assertGreater(target_y_from_state(current, previous), current.ball_y)
        self.assertEqual(select_action_from_state(current, previous), DOWN)

    def test_reflected_prediction_stays_inside_playfield(self) -> None:
        self.assertLessEqual(reflected_y(260), 210)
        self.assertGreaterEqual(reflected_y(20), 46)

    def test_target_prediction_stays_inside_paddle_target_band(self) -> None:
        previous = PongState(180, 190, 140, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(184, 202, 140, 100, ball_active=True, ball_in_right_half=True)

        target_y = target_y_from_state(current, previous)

        self.assertLessEqual(target_y, TARGET_MAX_Y_RAM)
        self.assertGreaterEqual(target_y, TARGET_MIN_Y_RAM)

    def test_visual_calibration_helpers_use_separate_ball_and_paddle_offsets(self) -> None:
        self.assertEqual(ram_ball_to_screen(100, 100), (51.5, 87.5))
        self.assertEqual(ram_paddle_y_to_screen(100), 102.75)

    def test_low_serve_pattern_detection(self) -> None:
        previous = PongState(132, 124, 120, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_low_serve_pattern(previous, dx=4, dy=-4))
        self.assertFalse(is_low_serve_pattern(previous, dx=-4, dy=-4))

    def test_ram_policy_applies_low_serve_target(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(132, 124, 120, 100, ball_active=True, ball_in_right_half=True)
        info = info_with_ram(**{"49": 136, "50": 100, "51": 70, "54": 120})

        action = policy.select_action(info)

        self.assertEqual(policy.low_serve_steps, 23)
        self.assertEqual(LOW_SERVE_TARGET_Y, 50)
        self.assertEqual(action, UP)

    def test_opening_low_dy8_correction_is_narrow(self) -> None:
        previous = PongState(182, 129, 147, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(
            186,
            137,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )

        self.assertTrue(is_opening_low_dy8_correction_state(current, previous, DOWN, -12))

        later_score = PongState(
            186,
            137,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=0,
        )
        self.assertFalse(is_opening_low_dy8_correction_state(later_score, previous, DOWN, -12))
        self.assertFalse(is_opening_low_dy8_correction_state(current, previous, NOOP, -12))

    def test_chain2_score00_initial_down_correction_is_narrow(self) -> None:
        state = PongState(
            125,
            131,
            85,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )

        self.assertTrue(
            is_chain2_score00_initial_down_correction_state(state, DOWN, "initial", 131, -21)
        )
        self.assertFalse(
            is_chain2_score00_initial_down_correction_state(state, NOOP, "initial", 131, -21)
        )
        self.assertFalse(
            is_chain2_score00_initial_down_correction_state(state, DOWN, "intercept", 131, -21)
        )
        self.assertFalse(
            is_chain2_score00_initial_down_correction_state(state, DOWN, "initial", 131, -10)
        )

    def test_chain2_score00_initial_policy_applies_up_once(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            0,
            0,
            106,
            100,
            ball_active=False,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        policy.previous_paddle_y = 106
        info = info_with_ram(**{"49": 125, "50": 100, "51": 85, "54": 131})

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertEqual(policy.last_mode, "chain2_score00_initial_down_correction")

    def test_low_serve_late_noop_correction_is_narrow(self) -> None:
        previous = PongState(193, 63, 39, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(197, 59, 47, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_low_serve_late_noop_correction_state(current, previous, NOOP, 50))
        self.assertFalse(is_low_serve_late_noop_correction_state(current, previous, DOWN, 50))
        self.assertFalse(is_low_serve_late_noop_correction_state(current, previous, NOOP, 51))

    def test_mid_down_precontact_up_correction_is_narrow(self) -> None:
        previous = PongState(186, 70, 77, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(190, 74, 86, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_mid_down_precontact_up_correction_state(current, previous, UP))
        self.assertFalse(is_mid_down_precontact_up_correction_state(current, previous, DOWN))

        too_late = PongState(198, 82, 69, 100, ball_active=True, ball_in_right_half=True)
        self.assertFalse(is_mid_down_precontact_up_correction_state(too_late, current, UP))

    def test_mid_up_precontact_up_correction_is_narrow(self) -> None:
        previous = PongState(183, 204, 199, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(187, 196, 202, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_mid_up_precontact_up_correction_state(current, previous, UP))
        self.assertFalse(is_mid_up_precontact_up_correction_state(current, previous, NOOP))

        too_low = PongState(195, 180, 180, 100, ball_active=True, ball_in_right_half=True)
        self.assertFalse(is_mid_up_precontact_up_correction_state(too_low, current, UP))

    def test_high_down_setup_noop_correction_is_narrow(self) -> None:
        previous = PongState(137, 135, 173, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(141, 143, 161, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_high_down_setup_noop_correction_state(current, previous, NOOP, 175))
        self.assertFalse(is_high_down_setup_noop_correction_state(current, previous, UP, 175))
        self.assertFalse(is_high_down_setup_noop_correction_state(current, previous, NOOP, 151))

    def test_high_down_score40_bounce_down_correction_is_narrow(self) -> None:
        previous = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=0,
        )
        current = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=0,
        )

        self.assertTrue(is_high_down_score40_bounce_down_correction_state(current, previous, DOWN, 156.5))
        self.assertFalse(is_high_down_score40_bounce_down_correction_state(current, previous, UP, 156.5))

        wrong_score = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        self.assertFalse(is_high_down_score40_bounce_down_correction_state(wrong_score, previous, DOWN, 156.5))

    def test_high_down_score40_bounce_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 4, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertEqual(policy.last_mode, "high_down_score40_bounce_down_correction")

    def test_high_down_score80_bounce_down_correction_is_narrow(self) -> None:
        previous = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=0,
        )
        current = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=0,
        )

        self.assertTrue(is_high_down_score80_bounce_down_correction_state(current, previous, DOWN, 156.5))
        self.assertFalse(is_high_down_score80_bounce_down_correction_state(current, previous, UP, 156.5))

        wrong_score = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=1,
        )
        self.assertFalse(is_high_down_score80_bounce_down_correction_state(wrong_score, previous, DOWN, 156.5))

    def test_high_down_score80_bounce_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.high_down_score80_bounce_correction_used = True
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 8, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.high_down_score80_bounce_correction_used)

    def test_high_down_score80_bounce_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 8, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.high_down_score80_bounce_correction_used)
        self.assertEqual(policy.last_mode, "high_down_score80_bounce_down_correction")

    def test_high_down_score120_bounce_down_correction_is_narrow(self) -> None:
        previous = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=12,
            opponent_score=0,
        )
        current = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=12,
            opponent_score=0,
        )

        self.assertTrue(is_high_down_score120_bounce_down_correction_state(current, previous, DOWN, 156.5))
        self.assertFalse(is_high_down_score120_bounce_down_correction_state(current, previous, UP, 156.5))

        wrong_score = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=12,
            opponent_score=1,
        )
        self.assertFalse(is_high_down_score120_bounce_down_correction_state(wrong_score, previous, DOWN, 156.5))

    def test_high_down_score120_bounce_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.high_down_score120_bounce_correction_used = True
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=12,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 12, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.high_down_score120_bounce_correction_used)

    def test_high_down_score120_bounce_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=12,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 12, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.high_down_score120_bounce_correction_used)
        self.assertEqual(policy.last_mode, "high_down_score120_bounce_down_correction")

    def test_high_down_score160_bounce_down_correction_is_narrow(self) -> None:
        previous = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=16,
            opponent_score=0,
        )
        current = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=16,
            opponent_score=0,
        )

        self.assertTrue(is_high_down_score160_bounce_down_correction_state(current, previous, DOWN, 156.5))
        self.assertFalse(is_high_down_score160_bounce_down_correction_state(current, previous, UP, 156.5))

        wrong_score = PongState(
            146,
            53,
            146,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=16,
            opponent_score=1,
        )
        self.assertFalse(is_high_down_score160_bounce_down_correction_state(wrong_score, previous, DOWN, 156.5))

    def test_high_down_score160_bounce_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.high_down_score160_bounce_correction_used = True
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=16,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 16, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.high_down_score160_bounce_correction_used)

    def test_high_down_score160_bounce_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            142,
            44,
            138,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=16,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 16, "49": 146, "50": 100, "51": 146, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.high_down_score160_bounce_correction_used)
        self.assertEqual(policy.last_mode, "high_down_score160_bounce_down_correction")

    def test_score00_mid_down_y154_final_up_correction_is_narrow(self) -> None:
        previous = PongState(
            195,
            146,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        current = PongState(
            199,
            150,
            155,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )

        self.assertTrue(
            is_score00_mid_down_y154_final_up_correction_state(current, previous, UP, 150.0, 20)
        )
        self.assertFalse(
            is_score00_mid_down_y154_final_up_correction_state(current, previous, DOWN, 150.0, 20)
        )
        self.assertFalse(
            is_score00_mid_down_y154_final_up_correction_state(current, previous, UP, 150.0, 8)
        )

        wrong_score = PongState(
            199,
            150,
            155,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=0,
        )
        self.assertFalse(
            is_score00_mid_down_y154_final_up_correction_state(wrong_score, previous, UP, 150.0, 20)
        )

    def test_score00_mid_down_y154_final_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.score00_mid_down_y154_final_correction_used = True
        policy.previous_state = PongState(
            195,
            146,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        policy.previous_paddle_y = 135
        info = info_with_ram(**{"13": 0, "14": 0, "49": 199, "50": 100, "51": 155, "54": 150})

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertTrue(policy.score00_mid_down_y154_final_correction_used)

    def test_score00_mid_down_y154_final_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            195,
            146,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        policy.previous_paddle_y = 135
        info = info_with_ram(**{"13": 0, "14": 0, "49": 199, "50": 100, "51": 155, "54": 150})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.score00_mid_down_y154_final_correction_used)
        self.assertEqual(policy.last_mode, "score00_mid_down_y154_final_up_correction")

    def test_late_low_loop_pre_up_correction_is_narrow(self) -> None:
        previous = PongState(169, 49, 86, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(173, 53, 98, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_late_low_loop_pre_up_correction_state(current, previous, UP, 72))
        self.assertFalse(is_late_low_loop_pre_up_correction_state(current, previous, DOWN, 72))
        self.assertFalse(is_late_low_loop_pre_up_correction_state(current, previous, UP, 90))

    def test_late_low_loop_correction_limit_matches_probe(self) -> None:
        policy = RamPolicy()
        policy.late_low_loop_corrections = LATE_LOW_LOOP_CORRECTION_LIMIT
        policy.previous_state = PongState(169, 49, 86, 100, ball_active=True, ball_in_right_half=True)
        info = info_with_ram(**{"49": 173, "50": 100, "51": 98, "54": 53})

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertEqual(policy.late_low_loop_corrections, LATE_LOW_LOOP_CORRECTION_LIMIT)

    def test_low_mid_score00_down_correction_is_narrow(self) -> None:
        previous = PongState(
            177,
            53,
            49,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        current = PongState(
            181,
            57,
            51,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )

        self.assertTrue(is_low_mid_score00_down_correction_state(current, previous, DOWN, 68))
        self.assertFalse(is_low_mid_score00_down_correction_state(current, previous, UP, 68))

        wrong_score = PongState(
            181,
            57,
            51,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=1,
        )
        self.assertFalse(is_low_mid_score00_down_correction_state(wrong_score, previous, DOWN, 68))

    def test_low_mid_score00_correction_limit_matches_probe(self) -> None:
        policy = RamPolicy()
        policy.low_mid_score00_corrections = LOW_MID_SCORE00_CORRECTION_LIMIT
        policy.previous_state = PongState(
            177,
            53,
            49,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 0, "49": 181, "50": 100, "51": 51, "54": 57})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertEqual(policy.low_mid_score00_corrections, LOW_MID_SCORE00_CORRECTION_LIMIT)

    def test_low_up_y111_pre_up_correction_is_narrow(self) -> None:
        previous = PongState(183, 151, 119, 100, ball_active=True, ball_in_right_half=True)
        current = PongState(187, 143, 137, 100, ball_active=True, ball_in_right_half=True)

        self.assertTrue(is_low_up_y111_pre_up_correction_state(current, previous, UP, 133))
        self.assertFalse(is_low_up_y111_pre_up_correction_state(current, previous, NOOP, 133))
        self.assertFalse(is_low_up_y111_pre_up_correction_state(current, previous, UP, 140))

    def test_low_up_y111_correction_limit_matches_probe(self) -> None:
        policy = RamPolicy()
        policy.low_up_y111_corrections = LOW_UP_Y111_CORRECTION_LIMIT
        policy.previous_state = PongState(183, 151, 119, 100, ball_active=True, ball_in_right_half=True)
        info = info_with_ram(**{"49": 187, "50": 100, "51": 137, "54": 143})

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertEqual(policy.low_up_y111_corrections, LOW_UP_Y111_CORRECTION_LIMIT)

    def test_y86_score22_setup_up_correction_is_narrow(self) -> None:
        previous = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=2,
        )
        current = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=2,
        )

        self.assertTrue(is_y86_score22_setup_up_correction_state(current, previous, UP, 76))
        self.assertFalse(is_y86_score22_setup_up_correction_state(current, previous, NOOP, 76))

        wrong_score = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=0,
        )
        self.assertFalse(is_y86_score22_setup_up_correction_state(wrong_score, previous, UP, 76))

    def test_y86_score22_setup_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.y86_score22_setup_correction_used = True
        policy.previous_state = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=2,
        )
        info = info_with_ram(**{"13": 2, "14": 2, "49": 174, "50": 100, "51": 87, "54": 58})

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertTrue(policy.y86_score22_setup_correction_used)

    def test_y86_score42_setup_up_correction_is_narrow(self) -> None:
        previous = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=2,
        )
        current = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=2,
        )

        self.assertTrue(is_y86_score42_setup_up_correction_state(current, previous, UP, 76))
        self.assertFalse(is_y86_score42_setup_up_correction_state(current, previous, NOOP, 76))

        wrong_score = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=3,
        )
        self.assertFalse(is_y86_score42_setup_up_correction_state(wrong_score, previous, UP, 76))

    def test_y86_score42_setup_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.y86_score42_setup_correction_used = True
        policy.previous_state = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=4,
            opponent_score=2,
        )
        info = info_with_ram(**{"13": 2, "14": 4, "49": 174, "50": 100, "51": 87, "54": 58})

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertTrue(policy.y86_score42_setup_correction_used)

    def test_y86_score62_setup_up_correction_is_narrow(self) -> None:
        previous = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=2,
        )
        current = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=2,
        )

        self.assertTrue(is_y86_score62_setup_up_correction_state(current, previous, UP, 76))
        self.assertFalse(is_y86_score62_setup_up_correction_state(current, previous, NOOP, 76))

        wrong_score = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=3,
        )
        self.assertFalse(is_y86_score62_setup_up_correction_state(wrong_score, previous, UP, 76))

    def test_y86_score62_setup_correction_uses_early_step_gate(self) -> None:
        policy = RamPolicy()
        policy.step_count = 1499
        policy.previous_state = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=2,
        )
        info = info_with_ram(**{"13": 2, "14": 6, "49": 174, "50": 100, "51": 87, "54": 58})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.y86_score62_early_setup_correction_used)

        policy = RamPolicy()
        policy.step_count = 1700
        policy.previous_state = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=2,
        )

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertFalse(policy.y86_score62_early_setup_correction_used)

    def test_y86_score82_setup_up_correction_is_narrow(self) -> None:
        previous = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=2,
        )
        current = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=2,
        )

        self.assertTrue(is_y86_score82_setup_up_correction_state(current, previous, UP, 76))
        self.assertFalse(is_y86_score82_setup_up_correction_state(current, previous, NOOP, 76))

        wrong_score = PongState(
            174,
            58,
            87,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=3,
        )
        self.assertFalse(is_y86_score82_setup_up_correction_state(wrong_score, previous, UP, 76))

    def test_y86_score82_setup_correction_uses_early_step_gate(self) -> None:
        policy = RamPolicy()
        policy.step_count = 2049
        policy.previous_state = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=2,
        )
        info = info_with_ram(**{"13": 2, "14": 8, "49": 174, "50": 100, "51": 87, "54": 58})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.y86_score82_early_setup_correction_used)

        policy = RamPolicy()
        policy.step_count = 2110
        policy.previous_state = PongState(
            170,
            54,
            79,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=2,
        )

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertFalse(policy.y86_score82_early_setup_correction_used)

    def test_y86_score20_serve_initial_down_correction_is_narrow(self) -> None:
        current = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=0,
        )

        self.assertTrue(is_y86_score20_serve_initial_down_correction_state(current, DOWN, "initial"))
        self.assertFalse(is_y86_score20_serve_initial_down_correction_state(current, NOOP, "initial"))
        self.assertFalse(is_y86_score20_serve_initial_down_correction_state(current, DOWN, "intercept"))

        wrong_score = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=1,
        )
        self.assertFalse(is_y86_score20_serve_initial_down_correction_state(wrong_score, DOWN, "initial"))

    def test_y86_score20_serve_initial_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.y86_score20_serve_initial_correction_used = True
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=2,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 2, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.y86_score20_serve_initial_correction_used)

    def test_y86_score60_serve_initial_down_correction_is_narrow(self) -> None:
        current = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=0,
        )

        self.assertTrue(is_y86_score60_serve_initial_down_correction_state(current, DOWN, "initial"))
        self.assertFalse(is_y86_score60_serve_initial_down_correction_state(current, NOOP, "initial"))
        self.assertFalse(is_y86_score60_serve_initial_down_correction_state(current, DOWN, "intercept"))

        wrong_score = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=1,
        )
        self.assertFalse(is_y86_score60_serve_initial_down_correction_state(wrong_score, DOWN, "initial"))

    def test_y86_score60_serve_initial_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.y86_score60_serve_initial_correction_used = True
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=6,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 6, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.y86_score60_serve_initial_correction_used)

    def test_y86_score60_serve_initial_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=6,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 6, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.y86_score60_serve_initial_correction_used)
        self.assertEqual(policy.last_mode, "y86_score60_serve_initial_down_correction")

    def test_y86_score100_serve_initial_down_correction_is_narrow(self) -> None:
        current = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=10,
            opponent_score=0,
        )

        self.assertTrue(is_y86_score100_serve_initial_down_correction_state(current, DOWN, "initial"))
        self.assertFalse(is_y86_score100_serve_initial_down_correction_state(current, NOOP, "initial"))
        self.assertFalse(is_y86_score100_serve_initial_down_correction_state(current, DOWN, "intercept"))

        wrong_score = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=10,
            opponent_score=1,
        )
        self.assertFalse(is_y86_score100_serve_initial_down_correction_state(wrong_score, DOWN, "initial"))

    def test_y86_score100_serve_initial_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.y86_score100_serve_initial_correction_used = True
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=10,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 10, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.y86_score100_serve_initial_correction_used)

    def test_y86_score100_serve_initial_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=10,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 10, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.y86_score100_serve_initial_correction_used)
        self.assertEqual(policy.last_mode, "y86_score100_serve_initial_down_correction")

    def test_y86_score140_serve_initial_down_correction_is_narrow(self) -> None:
        current = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=14,
            opponent_score=0,
        )

        self.assertTrue(is_y86_score140_serve_initial_down_correction_state(current, DOWN, "initial"))
        self.assertFalse(is_y86_score140_serve_initial_down_correction_state(current, NOOP, "initial"))
        self.assertFalse(is_y86_score140_serve_initial_down_correction_state(current, DOWN, "intercept"))

        wrong_score = PongState(
            126,
            126,
            101,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=14,
            opponent_score=1,
        )
        self.assertFalse(is_y86_score140_serve_initial_down_correction_state(wrong_score, DOWN, "initial"))

    def test_y86_score140_serve_initial_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.y86_score140_serve_initial_correction_used = True
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=14,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 14, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, DOWN)
        self.assertTrue(policy.y86_score140_serve_initial_correction_used)

    def test_y86_score140_serve_initial_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            52,
            0,
            101,
            100,
            ball_active=False,
            ball_in_right_half=False,
            self_score=14,
            opponent_score=0,
        )
        info = info_with_ram(**{"13": 0, "14": 14, "49": 126, "50": 100, "51": 101, "54": 126})

        action = policy.select_action(info)

        self.assertEqual(action, NOOP)
        self.assertTrue(policy.y86_score140_serve_initial_correction_used)
        self.assertEqual(policy.last_mode, "y86_score140_serve_initial_down_correction")

    def test_y86_bounce_x166_down_correction_is_narrow(self) -> None:
        previous = PongState(
            162,
            45,
            86,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=2,
        )
        current = PongState(
            166,
            50,
            78,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=2,
        )

        self.assertTrue(is_y86_bounce_x166_down_correction_state(current, previous, DOWN, 82.5))
        self.assertTrue(is_y86_bounce_score22_x166_down_correction_state(current, previous, DOWN, 82.5))
        self.assertFalse(is_y86_bounce_x166_down_correction_state(current, previous, NOOP, 82.5))

        wrong_score = PongState(
            166,
            50,
            78,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=6,
            opponent_score=0,
        )
        self.assertFalse(is_y86_bounce_score22_x166_down_correction_state(wrong_score, previous, DOWN, 82.5))

    def test_y86_bounce_after_score22_policy_uses_family_state(self) -> None:
        previous = PongState(
            162,
            45,
            86,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=2,
            opponent_score=2,
        )
        info = info_with_ram(**{"13": 2, "14": 2, "49": 166, "50": 100, "51": 78, "54": 50})
        policy = RamPolicy()
        policy.previous_state = previous

        action = policy.select_action(info)

        self.assertEqual(action, UP)
        self.assertTrue(policy.y86_score22_bounce_family_seen)
        self.assertEqual(policy.last_mode, "y86_bounce_after_score22_x166_down_correction")

        no_family_policy = RamPolicy()
        no_family_policy.previous_state = PongState(
            162,
            45,
            86,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=1,
        )
        later_info = info_with_ram(**{"13": 1, "14": 8, "49": 166, "50": 100, "51": 78, "54": 50})

        self.assertEqual(no_family_policy.select_action(later_info), DOWN)
        self.assertFalse(no_family_policy.y86_score22_bounce_family_seen)

        family_policy = RamPolicy()
        family_policy.y86_score22_bounce_family_seen = True
        family_policy.previous_state = PongState(
            162,
            45,
            86,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=8,
            opponent_score=1,
        )

        self.assertEqual(family_policy.select_action(later_info), UP)
        self.assertEqual(family_policy.last_mode, "y86_bounce_after_score22_x166_down_correction")

    def test_low_up_score01_down_correction_is_narrow(self) -> None:
        previous = PongState(
            173,
            182,
            156,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=1,
        )
        current = PongState(
            177,
            174,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=1,
        )

        self.assertTrue(is_low_up_score01_down_correction_state(current, previous, DOWN, 144, -21))
        self.assertFalse(is_low_up_score01_down_correction_state(current, previous, UP, 144, -21))

        wrong_score = PongState(
            177,
            174,
            135,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=1,
        )
        self.assertFalse(is_low_up_score01_down_correction_state(wrong_score, previous, DOWN, 144, -21))

    def test_low_up_score10_final_down_correction_is_narrow(self) -> None:
        previous = PongState(
            200,
            128,
            131,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=0,
        )
        current = PongState(
            204,
            120,
            115,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=0,
        )

        self.assertTrue(is_low_up_score10_final_down_correction_state(current, previous, DOWN, 120, -16))
        self.assertFalse(is_low_up_score10_final_down_correction_state(current, previous, UP, 120, -16))
        self.assertFalse(is_low_up_score10_final_down_correction_state(current, previous, DOWN, 120, -2))

        wrong_score = PongState(
            204,
            120,
            115,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=0,
            opponent_score=1,
        )
        self.assertFalse(is_low_up_score10_final_down_correction_state(wrong_score, previous, DOWN, 120, -16))

    def test_low_up_score10_final_correction_is_one_shot(self) -> None:
        policy = RamPolicy()
        policy.low_up_score10_final_correction_used = True
        policy.previous_state = PongState(
            200,
            128,
            131,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=0,
        )
        policy.previous_paddle_y = 131

        info = info_with_ram(**{"13": 0, "14": 1, "49": 204, "50": 100, "51": 115, "54": 120})

        self.assertEqual(policy.select_action(info), DOWN)
        self.assertEqual(policy.last_mode, "intercept")

    def test_low_up_score10_final_policy_applies_noop(self) -> None:
        policy = RamPolicy()
        policy.previous_state = PongState(
            200,
            128,
            131,
            100,
            ball_active=True,
            ball_in_right_half=True,
            self_score=1,
            opponent_score=0,
        )
        policy.previous_paddle_y = 131

        info = info_with_ram(**{"13": 0, "14": 1, "49": 204, "50": 100, "51": 115, "54": 120})

        self.assertEqual(policy.select_action(info), NOOP)
        self.assertTrue(policy.low_up_score10_final_correction_used)
        self.assertEqual(policy.last_mode, "low_up_score10_final_down_correction")


if __name__ == "__main__":
    unittest.main()
