import re

with open("env/environment.py", "r") as f:
    code = f.read()

old_block = """        reward = self.reward_calculator.calculate_step_reward(
            state={
                "step_count": self._state.step_count,
                "previous_config": self._state.previous_config,
                "expected_config": self._state.task.expected_config,
                "original_config": self._state.task.broken_config,
                "error": self._state.last_error,
                "changed_files_count": self._state.file_modification_count,
                "changed_lines_count": self._state.total_changed_lines,
                "consecutive_edit_actions": self._state.consecutive_edit_actions,
                "failed_validations": self._state.failed_validations,
            },
            action=tool,
            result=result,
            original_config=self._state.task.broken_config,
            fixed_config=self._state.current_config,
            error_message=self._state.last_error,
            expected_config=self._state.task.expected_config,
            metadata=self._state.task.metadata,
        )

        reward_model = Reward(value=float(reward), components={"total": float(reward)})"""

new_block = """        reward, comps = self.reward_calculator.calculate_step_reward(
            state={
                "step_count": self._state.step_count,
                "previous_config": self._state.previous_config,
                "expected_config": self._state.task.expected_config,
                "original_config": self._state.task.broken_config,
                "error": self._state.last_error,
                "changed_files_count": self._state.file_modification_count,
                "changed_lines_count": self._state.total_changed_lines,
                "consecutive_edit_actions": self._state.consecutive_edit_actions,
                "failed_validations": self._state.failed_validations,
            },
            action=tool,
            result=result,
            original_config=self._state.task.broken_config,
            fixed_config=self._state.current_config,
            error_message=self._state.last_error,
            expected_config=self._state.task.expected_config,
            metadata=self._state.task.metadata,
        )

        reward_model = Reward(value=float(reward), components=comps)"""

code = code.replace(old_block, new_block)

with open("env/environment.py", "w") as f:
    f.write(code)