"""训练/部署观测和动作的一致性，以及RL改进的目标口径。"""
import importlib.util
import math
from pathlib import Path

import pytest

np = pytest.importorskip('numpy')
pytest.importorskip('gymnasium')

from mcm26b.rl.env import LocalizationEnv
from mcm26b.rl.features import action_mask, build_observation, candidate_actions


def test_training_observation_matches_runtime_channel_belief():
    env = LocalizationEnv(problem=3, handoff_probability=1)
    obs, _ = env.reset(seed=1)
    belief = env._belief.channels[env._channel]
    expected = build_observation(belief.region, env._client.state.position, env._channel,
                                 env.config, env.problem, env._blocked, belief,
                                 virtual_time_s=env._client.state.virtual_time_s,
                                 receiver_channel=env._client.state.channel)
    np.testing.assert_array_equal(obs, expected)
    assert obs[7] == np.float32(1/12)
    assert tuple(obs[14:17]) == (0,1,0)
    assert env._client.state.channel != env._channel
    env.close()


def test_mask_rejects_measured_and_duplicate_candidates():
    env = LocalizationEnv(problem=3)
    env.reset(seed=17)
    belief = env._belief.channels[env._channel]
    region = belief.region
    position = env._client.state.position
    candidates = candidate_actions(region, position, env.config)
    belief.measurements.append(dict(position=candidates[0], result='no_signal', svd_deg=None))
    mask = env.action_masks()
    assert not mask[0]
    keys = [(round(p[0]*2),round(p[1]*2)) for p,v in zip(candidates,mask) if v]
    assert len(keys) == len(set(keys))
    assert any(mask)
    assert all(math.dist(p,region.mec()[0]) <= 1000-region.mec()[1]+1e-6
               for p,v in zip(candidates,mask) if v)
    env.close()


def test_q4_blocked_sequence_resets_on_direction():
    env = LocalizationEnv(problem=4,directional_probability=1,handoff_probability=0)
    env.reset(seed=7)
    for expected in ('no_signal','direction'):
        region = env._belief.candidate_set(env._channel)
        candidates = candidate_actions(region,env._client.state.position,env.config,env._blocked)
        # 只有测试端可用真值选择一个指定分支，策略从不接触该接口。
        index = next(i for i,p in enumerate(candidates)
                     if env._simulator._observe(p,env._channel)[0] == expected)
        _, _, terminated, truncated, _ = env.step(index)
        assert len(env._blocked) == (1 if expected=='no_signal' else 0)
        if expected=='no_signal':
            assert not terminated and not truncated
    env.close()


def test_success_reward_is_constant_minus_actual_virtual_cost():
    env = LocalizationEnv(problem=3,handoff_probability=0)
    env.reset(seed=3)
    total = 0
    for _ in range(16):
        legal = np.flatnonzero(env.action_masks())
        action = int(0 if 0 in legal else legal[0])
        _, reward, terminated, truncated, info = env.step(action)
        total += reward
        if terminated or truncated:
            break
    assert info['cleared']
    assert total == pytest.approx(20-info['virtual_cost_s']/100)
    assert env._client.state.exited


def test_bc_and_ppo_forward_mapping_agree():
    torch = pytest.importorskip('torch')
    from stable_baselines3 import PPO
    spec = importlib.util.spec_from_file_location('pretrain',Path('scripts/pretrain_rl_expert.py'))
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    torch.set_num_threads(1)
    env = LocalizationEnv()
    ppo = PPO('MlpPolicy',env,policy_kwargs={'net_arch':[128,128]},device='cpu',n_steps=16,batch_size=8)
    bc = module.BCMlp()
    state = bc.net.state_dict()
    ppo.policy.mlp_extractor.policy_net.load_state_dict({k:v for k,v in state.items() if not k.startswith('4.')})
    ppo.policy.action_net.load_state_dict({'weight':state['4.weight'],'bias':state['4.bias']})
    x = torch.randn(10,401)
    torch.testing.assert_close(bc(x),ppo.policy.action_net(ppo.policy.mlp_extractor.policy_net(x)))
    env.close()