from sim_reveal.procedural_envs import BAG_PROXY, CLOTH_PROXY, FOLIAGE_PROXY, make_proxy_env


def test_proxy_scripted_bench():
    bag_wins = 0
    foliage_wins = 0
    cloth_wins = 0
    for seed in range(10):
        bag_env = make_proxy_env(proxy_name=BAG_PROXY.name, resolution=32, seed=seed, rollout_horizon=4)
        bag_env.reset(seed=seed)
        maintain = bag_env.evaluate_action_chunk(bag_env.macro_action_chunk("maintain_mouth", chunk_horizon=4), rollout_horizon=4)
        retrieve = bag_env.evaluate_action_chunk(bag_env.macro_action_chunk("premature_retrieve", chunk_horizon=4), rollout_horizon=4)
        if (maintain["hold_persistence"] >= retrieve["hold_persistence"]) and (maintain["reocclusion_rate"] <= retrieve["reocclusion_rate"]):
            bag_wins += 1

        foliage_env = make_proxy_env(proxy_name=FOLIAGE_PROXY.name, resolution=32, seed=seed, rollout_horizon=4)
        foliage_env.reset(seed=seed)
        pin = foliage_env.evaluate_action_chunk(foliage_env.macro_action_chunk("pin_canopy", chunk_horizon=4), rollout_horizon=4)
        swipe = foliage_env.evaluate_action_chunk(foliage_env.macro_action_chunk("foliage_immediate_reocclusion", chunk_horizon=4), rollout_horizon=4)
        if (pin["reocclusion_rate"] <= swipe["reocclusion_rate"]) and (pin["visibility_integral"] >= swipe["visibility_integral"]):
            foliage_wins += 1

        cloth_env = make_proxy_env(proxy_name=CLOTH_PROXY.name, resolution=32, seed=seed, rollout_horizon=4)
        cloth_env.reset(seed=seed)
        stabilize = cloth_env.evaluate_action_chunk(cloth_env.macro_action_chunk("stabilize_fold", chunk_horizon=4), rollout_horizon=4)
        lift_high = cloth_env.evaluate_action_chunk(cloth_env.macro_action_chunk("cloth_lift_high", chunk_horizon=4), rollout_horizon=4)
        if (stabilize["candidate_fold_preservation"] >= lift_high["candidate_fold_preservation"]) and (
            stabilize["final_disturbance_cost"] <= lift_high["final_disturbance_cost"]
        ):
            cloth_wins += 1

    assert bag_wins >= 8
    assert foliage_wins >= 8
    assert cloth_wins >= 8