app/environment/core.py

import json
from pathlib import Path
from typing import Any, Literal, cast

from app.environment.graders import grade_task
from app.environment.scenarios.accident import generate_accident_case
from app.environment.scenarios.fire import generate_fire_case
from app.environment.scenarios.medical import generate_medical_case
from app.environment.validation import DifficultyModifier, HospitalValidator
from app.models.action import Action
from app.models.observation import ArrivalOutcomeObservation, HospitalObservation, Observation
from app.models.reward import RewardBreakdown, RewardModel, StepInfo
from app.models.state import ArrivalOutcome, EnvState, HospitalState, HospitalValidationDetails, LearningEntry
from app.utils.calculations import compute_speed_kmh, compute_travel_time_minutes
from app.utils.randomizer import SeededRandomizer


TASKS = {
    "acde_easy": {
        "difficulty": "easy",
        "objective": "Stabilize quickly while information is mostly reliable.",
    },
    "acde_medium": {
        "difficulty": "medium",
        "objective": "Balance speed, uncertainty, and specialization constraints.",
    },
    "acde_hard": {
        "difficulty": "hard",
        "objective": "Make least-bad decisions when every hospital has trade-offs.",
    },
}

MIN_REWARD = 0.001
MAX_REWARD = 0.999

OUTCOME_SCORE = {"accepted": 3, "partial": 2, "rejected": 1}
CONDITION_ORDER = ["stable", "serious", "unstable", "critical"]


class EmergencyEnv:
    """Stateful local RL environment for emergency routing under uncertainty."""

    def __init__(self, memory_file: str):
        self.memory_path = Path(memory_file)
        self.memory_path.parent.mkdir(parents=True, exist_ok=True)
        if not self.memory_path.exists():
            self.memory_path.write_text("{}", encoding="utf-8")

        self.episode_counter = 0
        self._rng = SeededRandomizer(seed=42)
        self.state_data: EnvState | None = None
        self.validator = HospitalValidator(self._rng)
        self.trajectory: list[dict[str, Any]] = []
        self.last_info: StepInfo | None = None
        self.last_outcome_status: str | None = None
        self.base_speed_kmh = 60.0

    def reset(self, seed: int | None = None, task_id: str | None = None) -> Observation:
        if seed is None:
            seed = self._rng.randint(1, 10**9)

        resolved_task_id = cast(Literal["acde_easy", "acde_medium", "acde_hard"], task_id if task_id in TASKS else self._rng.choice(list(TASKS.keys())))
        difficulty = TASKS[resolved_task_id]["difficulty"]

        self._rng = SeededRandomizer(seed)
        self.validator = HospitalValidator(self._rng)
        self.episode_counter += 1
        self.trajectory = []
        self.last_outcome_status = None

        scenario, scenario_type = self._sample_scenario_for_difficulty(difficulty)
        hospitals = self._build_hospital_states(scenario)
        hospitals = self._augment_hospital_options(
            hospitals,
            difficulty,
            required_specialization=scenario["required_specialization"],
        )
        hospitals = self._inject_no_perfect_option(hospitals, difficulty)

        max_steps = {"easy": 3, "medium": 4, "hard": 4}.get(difficulty, 4)

        self.state_data = EnvState(
            episode_id=self.episode_counter,
            seed=seed,
            task_id=cast(Literal["acde_easy", "acde_medium", "acde_hard"], resolved_task_id),
            task_objective=TASKS[resolved_task_id]["objective"],
            scenario_type=cast(Literal["medical", "accident", "fire"], scenario_type),
            scenario_name=scenario["scenario_name"],
            scenario_difficulty=cast(Literal["easy", "medium", "hard"], difficulty),
            patient_condition=scenario["patient_condition"],
            required_specialization=scenario["required_specialization"],
            initial_critical_time_limit_minutes=scenario["critical_time_limit_minutes"],
            critical_time_limit_minutes=scenario["critical_time_limit_minutes"],
            step=1,
            max_steps=max_steps,
            hospitals=hospitals,
            selected_hospital_id=None,
            done=False,
            final_outcome=None,
            reward=MIN_REWARD,
            final_score=MIN_REWARD,
            ambulance_status="en_route",
            current_location_context="incident_site",
            visited_hospitals=[],
            failed_hospitals=[],
            recent_failed_hospitals=[],
            failed_reasons={},
            total_time_spent_minutes=0.0,
            rerouting_reason=None,
            last_arrival_outcome=None,
            accepted_hospital_id=None,
            explanation=[
                "Episode initialized with seeded uncertainty.",
                f"Difficulty: {difficulty}. Hidden hospital state can change during transit.",
                f"Patient condition: {scenario['patient_condition']}.",
                f"Required specialization: {scenario['required_specialization']}.",
                "Primary objective: admit patient successfully under uncertainty.",
            ],
            memory=self._load_memory(),
        )

        self.last_info = StepInfo(
            task_id=cast(Literal["acde_easy", "acde_medium", "acde_hard"], resolved_task_id),
            difficulty=cast(Literal["easy", "medium", "hard"], difficulty),
            objective=TASKS[resolved_task_id]["objective"],
            progress_score=MIN_REWARD,
            reward_model=RewardModel(
                value=MIN_REWARD,
                breakdown=RewardBreakdown(
                    survival_component=MIN_REWARD,
                    time_efficiency_component=MIN_REWARD,
                    specialization_component=MIN_REWARD,
                    delay_penalty=MIN_REWARD,
                ),
            ),
            grader=grade_task(
                task_id=resolved_task_id,
                difficulty=difficulty,
                objective=TASKS[resolved_task_id]["objective"],
                trajectory=[],
            ),
            last_action_error=None,
            outcome=None,
        )

        return self._build_observation()

    def state(self) -> EnvState:
        if self.state_data is None:
            self.reset(seed=42, task_id="acde_medium")
        assert self.state_data is not None
        return self.state_data

    def step(self, action: Action | str | dict[str, Any]) -> dict[str, Any]:
        if self.state_data is None:
            self.reset(seed=42, task_id="acde_medium")
        assert self.state_data is not None

        if self.state_data.done:
            info = self.last_info.model_dump() if self.last_info else {}
            return {
                "observation": self._build_observation(),
                "reward": MIN_REWARD,
                "done": True,
                "info": info,
            }

        normalized_action = self._normalize_action(action)
        if normalized_action.step != self.state_data.step:
            raise ValueError(
                f"Action step {normalized_action.step} does not match environment step {self.state_data.step}."
            )

        selected = self._find_hospital(normalized_action.hospital_id)
        if selected is None:
            raise ValueError(f"Unknown hospital id: {normalized_action.hospital_id}")

        was_visited_before = selected.hospital_id in self.state_data.visited_hospitals
        was_failed_before = selected.hospital_id in self.state_data.failed_hospitals

        original_traffic = selected.traffic
        selected.traffic = cast(Literal["low", "medium", "high"], self._traffic_shift(selected.traffic, self.state_data.scenario_difficulty))

        speed = compute_speed_kmh(self.base_speed_kmh, selected.traffic)
        travel_time = compute_travel_time_minutes(selected.distance_km, speed)

        delay_probability = {
            "easy": 0.10,
            "medium": 0.25,
            "hard": 0.45,
        }.get(self.state_data.scenario_difficulty, 0.25)
        dynamic_delay = self._rng.uniform(0.5, 2.5) if self._rng.random() < delay_probability else 0.0
        travel_time += dynamic_delay

        selected, travel_time, enroute_note = self._apply_enroute_diversion(selected, travel_time)

        self.state_data.total_time_spent_minutes += travel_time

        if selected.hospital_id not in self.state_data.visited_hospitals:
            self.state_data.visited_hospitals.append(selected.hospital_id)

        self.state_data.ambulance_status = "arrived"
        self.state_data.current_location_context = f"arrived_at_{selected.hospital_id}"

        arrival_outcome = self.validator.validate_arrival(
            hospital=selected,
            difficulty=self.state_data.scenario_difficulty,
            patient_condition=self.state_data.patient_condition,
            required_specialization=self.state_data.required_specialization,
            total_time_spent=self.state_data.total_time_spent_minutes,
            critical_time_limit=self.state_data.critical_time_limit_minutes,
            step_number=self.state_data.step,
        )

        # Hidden-case guess: selecting uncertain ICU may lead to wrong guess at arrival.
        arrival_outcome, hidden_case_penalty, hidden_case_note = self._apply_hidden_guess_case(
            selected,
            arrival_outcome,
        )

        # Late-arrival shocks: on arrival, resources may suddenly become unavailable.
        arrival_outcome, shock_penalty, shock_note = self._apply_arrival_hidden_shock(
            arrival_outcome,
            difficulty=self.state_data.scenario_difficulty,
        )

        # Fix 1: cap partial chains so they resolve after repeated delays.
        arrival_outcome, partial_cap_note = self._apply_partial_chain_cap(arrival_outcome)

        # Critical polish: early hard rejections can degrade to partial to preserve recoverability.
        arrival_outcome, early_reject_note = self._apply_early_reject_protection(arrival_outcome)

        # Critical polish: partial outcomes after step 2 can recover into acceptance.
        arrival_outcome, late_partial_note = self._apply_late_partial_recovery(arrival_outcome)

        # Fix 3: final-attempt pressure can produce emergency stabilization.
        arrival_outcome, last_chance_note = self._apply_last_chance_outcome(arrival_outcome)

        reward, breakdown = self._calculate_reward(
            selected=selected,
            arrival_outcome=arrival_outcome,
            travel_time=travel_time,
            was_visited_before=was_visited_before,
            was_failed_before=was_failed_before,
            hidden_case_penalty=hidden_case_penalty + shock_penalty,
        )

        success = arrival_outcome.status in {"accepted", "partial"}
        self._update_learning_memory(selected.hospital_id, success, reward)
        self.state_data.memory = self._load_memory()

        self._record_trajectory(
            selected=selected,
            arrival_outcome=arrival_outcome,
            reward=reward,
            travel_time=travel_time,
            dynamic_delay=dynamic_delay,
            original_traffic=original_traffic,
        )

        self.state_data.selected_hospital_id = selected.hospital_id
        self.state_data.reward = reward
        self.state_data.last_arrival_outcome = arrival_outcome

        self._advance_patient_state(arrival_outcome.status, travel_time, dynamic_delay)

        self._resolve_transition(selected, arrival_outcome)

        self._build_last_info(reward, breakdown, arrival_outcome)

        if not self.state_data.done:
            self._evolve_hospital_uncertainty()

        self._set_explanation(
            selected,
            arrival_outcome,
            travel_time,
            dynamic_delay,
            original_traffic,
            [
                note
                for note in [
                    enroute_note,
                    hidden_case_note,
                    shock_note,
                    partial_cap_note,
                    early_reject_note,
                    late_partial_note,
                    last_chance_note,
                ]
                if note
            ],
        )

        info = self.last_info.model_dump() if self.last_info else {}
        # Clamp reward into the strict open interval (0, 1) for the external validator.
        clamped_reward = max(MIN_REWARD, min(MAX_REWARD, reward))
        return {
            "observation": self._build_observation(),
            "reward": clamped_reward,
            "done": self.state_data.done,
            "info": info,
        }

    def _normalize_action(self, action: Action | str | dict[str, Any]) -> Action:
        if isinstance(action, Action):
            return action
        if isinstance(action, str):
            assert self.state_data is not None
            return Action(step=self.state_data.step, hospital_id=action, rationale="policy selection")
        if isinstance(action, dict):
            assert self.state_data is not None
            return Action(
                step=action.get("step", self.state_data.step),
                hospital_id=str(action.get("hospital_id", "")),
                rationale=action.get("rationale"),
            )
        raise ValueError("Action must be Action, hospital_id string, or action dict.")

    def _build_hospital_states(self, scenario: dict[str, Any]) -> list[HospitalState]:
        hospitals: list[HospitalState] = []
        for template in scenario["hospitals"]:
            distance = round(
                self._rng.uniform(template["distance_range"][0], template["distance_range"][1]),
                1,
            )
            traffic = self._rng.choice(template["traffic_options"])
            icu_actual = self._rng.random() < template["icu_true_probability"]

            if icu_actual:
                icu_display = "available" if self._rng.random() < 0.85 else "unknown"
            else:
                icu_display = "available" if self._rng.random() < 0.2 else "unknown"

            hospitals.append(
                HospitalState(
                    hospital_id=template["hospital_id"],
                    distance_km=distance,
                    icu_display=icu_display,
                    icu_actual=icu_actual,
                    specialization=template["specialization"],
                    traffic=traffic,
                )
            )
        return hospitals

    def _inject_no_perfect_option(self, hospitals: list[HospitalState], difficulty: str) -> list[HospitalState]:
        trigger = {"easy": 0.06, "medium": 0.30, "hard": 0.42}.get(difficulty, 0.30)
        if self._rng.random() >= trigger:
            return hospitals

        if len(hospitals) < 3:
            return hospitals

        hospitals[0].traffic = "high"
        hospitals[1].icu_display = "unknown"
        hospitals[2].specialization = "general" if hospitals[2].specialization != "general" else "trauma"
        hospitals[2].icu_display = "unknown"
        return hospitals

    def _augment_hospital_options(
        self,
        hospitals: list[HospitalState],
        difficulty: str,
        required_specialization: str,
    ) -> list[HospitalState]:
        """Add extra decoy/alternative hospitals to increase decision ambiguity."""
        target_extra = {"easy": 1, "medium": 1, "hard": 2}.get(difficulty, 1)
        extra_count = 0
        while extra_count < target_extra:
            new_id = f"H{len(hospitals) + 1}"
            # Keep options plausible but uncertain: mixed specialization and variable traffic.
            spec_roll = self._rng.random()
            if spec_roll < 0.45:
                specialization = required_specialization
            elif spec_roll < 0.75:
                specialization = "general"
            else:
                specialization = "trauma" if required_specialization != "trauma" else "cardiac"

            distance = round(self._rng.uniform(4.0, 13.5), 1)
            traffic = self._rng.choice(["low", "medium", "high"])

            icu_prob = {"easy": 0.62, "medium": 0.52, "hard": 0.42}.get(difficulty, 0.52)
            icu_actual = self._rng.random() < icu_prob
            if icu_actual:
                icu_display = "available" if self._rng.random() < 0.74 else "unknown"
            else:
                icu_display = "available" if self._rng.random() < 0.18 else "unknown"

            hospitals.append(
                HospitalState(
                    hospital_id=new_id,
                    distance_km=distance,
                    icu_display=icu_display,
                    icu_actual=icu_actual,
                    specialization=cast(Literal["cardiac", "trauma", "general"], specialization),
                    traffic=cast(Literal["low", "medium", "high"], traffic),
                )
            )
            extra_count += 1
        return hospitals

    def _calculate_reward(
        self,
        selected: HospitalState,
        arrival_outcome: ArrivalOutcome,
        travel_time: float,
        was_visited_before: bool,
        was_failed_before: bool,
        hidden_case_penalty: float,
    ) -> tuple[float, RewardBreakdown]:
        assert self.state_data is not None

        base_status_reward = {
            "accepted": 0.92,
            "partial": 0.55,
            "rejected": 0.08,
        }[arrival_outcome.status]

        if arrival_outcome.status == "rejected":
            status_reward = base_status_reward
        else:
            outcome_modifier = max(0.5, min(1.2, float(arrival_outcome.reward_modifier)))
            status_reward = base_status_reward * outcome_modifier

        critical_patient = self.state_data.patient_condition in {"critical", "unstable"}
        unknown_critical_penalty = (
            0.12
            if critical_patient and selected.icu_display == "unknown"
            else 0.0
        )
        repeat_penalty = 0.15 if was_visited_before else 0.0
        failed_repeat_penalty = 0.20 if was_failed_before else 0.0
        traffic_penalty = 0.10 if critical_patient and selected.traffic == "high" else 0.04 if critical_patient and selected.traffic == "medium" else 0.0

        time_bonus = 0.06 if travel_time <= 8.0 else (0.03 if travel_time <= 14.0 else 0.0)

        improvement_bonus = self._improvement_bonus(arrival_outcome.status)

        reward = (
            status_reward
            + time_bonus
            + improvement_bonus
            - unknown_critical_penalty
            - repeat_penalty
            - failed_repeat_penalty
            - traffic_penalty
            - hidden_case_penalty
        )
        reward = max(MIN_REWARD, min(MAX_REWARD, reward))

        raw_delay = (
            unknown_critical_penalty
            + repeat_penalty
            + failed_repeat_penalty
            + traffic_penalty
            + hidden_case_penalty
        )
        breakdown = RewardBreakdown(
            survival_component=max(MIN_REWARD, min(MAX_REWARD, (status_reward + 0.5) / 1.5)),
            time_efficiency_component=max(MIN_REWARD, min(MAX_REWARD, 1.0 - (travel_time / 25.0))),
            specialization_component=max(MIN_REWARD, min(MAX_REWARD, MAX_REWARD if self._specialization_match(selected) else 0.4)),
            delay_penalty=max(MIN_REWARD, min(MAX_REWARD, raw_delay)),
        )

        return reward, breakdown

    def _improvement_bonus(self, status: str) -> float:
        if self.last_outcome_status is None:
            self.last_outcome_status = status
            return MIN_REWARD

        delta = OUTCOME_SCORE[status] - OUTCOME_SCORE[self.last_outcome_status]
        self.last_outcome_status = status
        if delta > 0:
            return 0.04
        return MIN_REWARD

    def _specialization_match(self, hospital: HospitalState) -> bool:
        assert self.state_data is not None
        return (
            hospital.specialization == self.state_data.required_specialization
            or hospital.specialization == "general"
            or self.state_data.required_specialization == "general"
        )

    def _advance_patient_state(self, outcome_status: str, travel_time: float, dynamic_delay: float) -> None:
        assert self.state_data is not None

        condition = self.state_data.patient_condition
        idx = CONDITION_ORDER.index(condition) if condition in CONDITION_ORDER else 2

        deterioration_risk = 0.0
        if travel_time > 12.0:
            deterioration_risk += 0.20
        if dynamic_delay > 0:
            deterioration_risk += 0.15
        if outcome_status == "rejected":
            deterioration_risk += 0.20

        if self._rng.random() < min(0.95, deterioration_risk):
            idx = min(len(CONDITION_ORDER) - 1, idx + 1)

        if outcome_status == "partial":
            stabilize_prob = {"easy": 0.35, "medium": 0.22, "hard": 0.12}.get(
                self.state_data.scenario_difficulty,
                0.22,
            )
            if self._rng.random() < stabilize_prob:
                idx = max(0, idx - 1)

        self.state_data.patient_condition = CONDITION_ORDER[idx]

    def _resolve_transition(self, selected: HospitalState, arrival_outcome: ArrivalOutcome) -> None:
        assert self.state_data is not None

        if arrival_outcome.status == "accepted":
            self.state_data.accepted_hospital_id = selected.hospital_id
            self.state_data.ambulance_status = "admitted"
            self.state_data.current_location_context = selected.hospital_id
            self.state_data.done = True
            self.state_data.final_outcome = "SUCCESS"
            self.state_data.final_score = self._success_score()
            return

        if arrival_outcome.status == "rejected":
            if selected.hospital_id not in self.state_data.failed_hospitals:
                self.state_data.failed_hospitals.append(selected.hospital_id)

            # Cooldown memory: block immediate retries, but allow later reconsideration.
            self.state_data.recent_failed_hospitals.append(selected.hospital_id)
            if len(self.state_data.recent_failed_hospitals) > 3:
                self.state_data.recent_failed_hospitals.pop(0)

            self.state_data.failed_reasons[selected.hospital_id] = arrival_outcome.reason

            if arrival_outcome.terminal:
                self.state_data.done = True
                self.state_data.final_outcome = "FAILURE"
                self.state_data.final_score = self._failure_score()
                self.state_data.rerouting_reason = arrival_outcome.reason
                self.state_data.ambulance_status = "arrived"
                self.state_data.current_location_context = f"terminal_failure_at_{selected.hospital_id}"
                return

            self.state_data.rerouting_reason = arrival_outcome.reason
            self.state_data.ambulance_status = "rerouting"
            self.state_data.current_location_context = f"rejected_at_{selected.hospital_id}"
        else:
            self.state_data.ambulance_status = "in_transit"
            self.state_data.current_location_context = "post_partial_treatment"

        if self._critical_failure():
            self.state_data.done = True
            self.state_data.final_outcome = "FAILURE"
            self.state_data.final_score = self._failure_score()
            return

        if self.state_data.step >= self.state_data.max_steps:
            self.state_data.done = True
            self.state_data.final_outcome = "FAILURE"
            self.state_data.final_score = self._failure_score()
            return

        self.state_data.step += 1
        self.state_data.done = False
        self.state_data.final_outcome = None

    def _critical_failure(self) -> bool:
        # Time-window based failure is disabled. Episodes end by acceptance or max steps.
        return False

    def _set_explanation(
        self,
        selected: HospitalState,
        arrival_outcome: ArrivalOutcome,
        travel_time: float,
        dynamic_delay: float,
        original_traffic: str,
        hidden_case_notes: list[str],
    ) -> None:
        assert self.state_data is not None
        v = arrival_outcome.validation_details
        assert v is not None
        self.state_data.explanation = [
            f"Step {self.state_data.step}: selected {selected.hospital_id}.",
            f"Traffic changed {original_traffic} -> {selected.traffic} before arrival.",
            f"Travel time: {travel_time:.2f} min (delay {dynamic_delay:.2f} min).",
            f"Validation checks: ICU={v.icu_available}, doctor={v.doctor_available}, equipment={v.equipment_functional}, overload={v.overload_status}",
            f"Patient suitability score = {v.patient_suitability:.2f}",
            f"Arrival outcome = {arrival_outcome.status.upper()}",
            f"Arrival reason = {arrival_outcome.reason}",
            f"Patient condition now = {self.state_data.patient_condition}",
            f"Total time spent = {self.state_data.total_time_spent_minutes:.2f} min",
        ]
        for note in hidden_case_notes:
            self.state_data.explanation.append(note)

    def _apply_enroute_diversion(
        self,
        selected: HospitalState,
        travel_time: float,
    ) -> tuple[HospitalState, float, str | None]:
        """Sometimes traffic collapses mid-route and ambulance diverts before arrival."""
        assert self.state_data is not None

        base_diversion_prob = {
            "easy": 0.04,
            "medium": 0.12,
            "hard": 0.18,
        }.get(self.state_data.scenario_difficulty, 0.20)

        if selected.traffic == "high":
            base_diversion_prob += 0.08
        elif selected.traffic == "medium":
            base_diversion_prob += 0.04

        if self._rng.random() >= min(0.85, base_diversion_prob):
            return selected, travel_time, None

        alternatives = [
            h
            for h in self.state_data.hospitals
            if h.hospital_id != selected.hospital_id and h.hospital_id not in self.state_data.failed_hospitals
        ]
        if not alternatives:
            return selected, travel_time, None

        def _rank(h: HospitalState) -> tuple[int, float]:
            traffic_rank = {"low": 0, "medium": 1, "high": 2}.get(h.traffic, 1)
            return (traffic_rank, h.distance_km)

        diverted = sorted(alternatives, key=_rank)[0]
        diverted_speed = compute_speed_kmh(self.base_speed_kmh, diverted.traffic)
        diverted_time = compute_travel_time_minutes(diverted.distance_km, diverted_speed)
        diversion_overhead = {
            "easy": self._rng.uniform(0.4, 1.1),
            "medium": self._rng.uniform(0.8, 1.8),
            "hard": self._rng.uniform(1.2, 2.6),
        }.get(self.state_data.scenario_difficulty, self._rng.uniform(1.0, 2.2))

        note = (
            f"Hidden case: severe traffic lock en-route to {selected.hospital_id}; "
            f"ambulance diverted to {diverted.hospital_id}."
        )
        return diverted, diverted_time + diversion_overhead, note

    def _apply_hidden_guess_case(
        self,
        selected: HospitalState,
        arrival_outcome: ArrivalOutcome,
    ) -> tuple[ArrivalOutcome, float, str | None]:
        """Resolve hidden guess cases for uncertain hospitals.

        If ICU is shown as unknown, the agent is effectively guessing.
        Wrong guess triggers stronger penalty and forced reroute.
        """
        assert self.state_data is not None

        if selected.icu_display != "unknown":
            return arrival_outcome, MIN_REWARD, None

        difficulty = self.state_data.scenario_difficulty
        guess_success_prob = {
            "easy": 0.82,
            "medium": 0.72,
            "hard": 0.58,
        }.get(difficulty, 0.52)
        guess_correct = self._rng.random() < guess_success_prob

        if guess_correct:
            return (
                arrival_outcome,
                MIN_REWARD,
                "Hidden case: risky ICU-unknown guess was correct this time.",
            )

        # Wrong hidden guess: downgrade to rejected and enforce rerouting signal.
        forced_reject = ArrivalOutcome(
            status="rejected",
            reason="Hidden mismatch at arrival (wrong risky guess). Rerouting required.",
            validation_details=arrival_outcome.validation_details,
            reward_modifier=0.001,
        )
        return (
            forced_reject,
            0.14,
            "Hidden case: risky ICU-unknown guess failed; penalty applied.",
        )

    def _apply_arrival_hidden_shock(
        self,
        arrival_outcome: ArrivalOutcome,
        difficulty: str,
    ) -> tuple[ArrivalOutcome, float, str | None]:
        """Late-arrival operational shocks: ICU/doctor/bed/equipment can fail at handover."""
        if arrival_outcome.status == "rejected":
            return arrival_outcome, MIN_REWARD, None

        shock_prob = {
            "easy": 0.03,
            "medium": 0.05,
            "hard": 0.10,
        }.get(difficulty, 0.14)
        if self._rng.random() >= shock_prob:
            return arrival_outcome, MIN_REWARD, None

        v = arrival_outcome.validation_details
        if v is None:
            return arrival_outcome, MIN_REWARD, None

        shock = self._rng.choice([
            "doctor_unavailable",
            "icu_full",
            "beds_full",
            "machine_failed",
        ])

        if shock == "doctor_unavailable":
            reason = "Doctor was reassigned to another emergency at arrival"
            new_validation = HospitalValidationDetails(
                icu_available=v.icu_available,
                doctor_available=False,
                equipment_functional=v.equipment_functional,
                overload_status=v.overload_status,
                patient_suitability=v.patient_suitability,
            )
        elif shock == "icu_full":
            reason = "ICU got full moments before handover"
            new_validation = HospitalValidationDetails(
                icu_available=False,
                doctor_available=v.doctor_available,
                equipment_functional=v.equipment_functional,
                overload_status=v.overload_status,
                patient_suitability=v.patient_suitability,
            )
        elif shock == "beds_full":
            reason = "Emergency beds became unavailable during arrival"
            new_validation = HospitalValidationDetails(
                icu_available=v.icu_available,
                doctor_available=v.doctor_available,
                equipment_functional=v.equipment_functional,
                overload_status="severe",
                patient_suitability=v.patient_suitability,
            )
        else:
            reason = "Critical treatment machine failed at admission"
            new_validation = HospitalValidationDetails(
                icu_available=v.icu_available,
                doctor_available=v.doctor_available,
                equipment_functional=False,
                overload_status=v.overload_status,
                patient_suitability=v.patient_suitability,
            )

        return (
            ArrivalOutcome(
                status="rejected",
                reason=reason,
                validation_details=new_validation,
                reward_modifier=0.001,
            ),
            0.12,
            f"Hidden case: {reason}. Rerouting required.",
        )

    def _apply_partial_chain_cap(self, arrival_outcome: ArrivalOutcome) -> tuple[ArrivalOutcome, str | None]:
        """Fix 1: after repeated partials, force resolution to accepted or rejected."""
        assert self.state_data is not None
        if arrival_outcome.status != "partial":
            return arrival_outcome, None

        prior_partials = sum(1 for t in self.trajectory if t.get("outcome_status") == "partial")
        partial_count = prior_partials + 1
        if partial_count < 2:
            return arrival_outcome, None

        stabilize_chance = {
            "easy": 0.45,
            "medium": 0.28,
            "hard": 0.05,
        }.get(self.state_data.scenario_difficulty, 0.28)

        if self._rng.random() < stabilize_chance:
            return (
                ArrivalOutcome(
                    status="accepted",
                    reason="Patient stabilized after critical delay",
                    validation_details=arrival_outcome.validation_details,
                    reward_modifier=0.78 if self.state_data.scenario_difficulty == "easy" else 0.68,
                ),
                "Partial chain cap: resolved as emergency stabilization.",
            )

        carry_partial_chance = 0.3 if self.state_data.scenario_difficulty != "hard" else 0.15
        if self._rng.random() < carry_partial_chance:
            return (
                ArrivalOutcome(
                    status="partial",
                    reason="Condition worsened but remains temporarily transferable",
                    validation_details=arrival_outcome.validation_details,
                    reward_modifier=0.44,
                ),
                "Partial chain cap: temporary recovery preserved rerouting chance.",
            )

        return (
            ArrivalOutcome(
                status="rejected",
                reason="Condition became irreversible after delays",
                validation_details=arrival_outcome.validation_details,
                reward_modifier=0.001,
            ),
            "Partial chain cap: condition became irreversible.",
        )

    def _apply_last_chance_outcome(self, arrival_outcome: ArrivalOutcome) -> tuple[ArrivalOutcome, str | None]:
        """Fix 3: near final attempt, allow emergency stabilization chance."""
        assert self.state_data is not None
        if arrival_outcome.status == "accepted":
            return arrival_outcome, None
        # Apply only on the literal final step, not one step earlier.
        if self.state_data.step != self.state_data.max_steps:
            return arrival_outcome, None

        chance = {
            "easy": 0.35,
            "medium": 0.18,
            "hard": 0.02,
        }.get(self.state_data.scenario_difficulty, 0.18)

        reward_modifier = {
            "easy": 0.82,
            "medium": 0.70,
            "hard": 0.58,
        }.get(self.state_data.scenario_difficulty, 0.70)

        if self._rng.random() < chance:
            return (
                ArrivalOutcome(
                    status="accepted",
                    reason="Emergency stabilization at last attempt",
                    validation_details=arrival_outcome.validation_details,
                    reward_modifier=reward_modifier,
                ),
                "Last-chance rule: emergency stabilization triggered.",
            )
        return arrival_outcome, None

    def _apply_early_reject_protection(self, arrival_outcome: ArrivalOutcome) -> tuple[ArrivalOutcome, str | None]:
        """Avoid excessive instant dead-ends by softening some step-1 rejections."""
        assert self.state_data is not None
        if arrival_outcome.status != "rejected":
            return arrival_outcome, None
        if self.state_data.step >= 2:
            return arrival_outcome, None
        soften_reject_chance = 0.3 if self.state_data.scenario_difficulty != "hard" else 0.05
        if self._rng.random() >= soften_reject_chance:
            return arrival_outcome, None

        return (
            ArrivalOutcome(
                status="partial",
                reason="Early rejection mitigated by emergency field stabilization",
                validation_details=arrival_outcome.validation_details,
                reward_modifier=0.50,
                terminal=False,
            ),
            "Recovery guard: early rejection softened to partial.",
        )

    def _apply_late_partial_recovery(self, arrival_outcome: ArrivalOutcome) -> tuple[ArrivalOutcome, str | None]:
        """Allow realistic comeback from partial outcomes after initial stabilization attempts."""
        assert self.state_data is not None
        if arrival_outcome.status != "partial":
            return arrival_outcome, None
        if self.state_data.step < 2:
            return arrival_outcome, None
        recovery_trigger = 0.5 if self.state_data.scenario_difficulty != "hard" else 0.25
        if self._rng.random() >= recovery_trigger:
            return arrival_outcome, None

        reject_from_partial = 0.5 if self.state_data.scenario_difficulty != "hard" else 0.8
        if self._rng.random() < reject_from_partial:
            return (
                ArrivalOutcome(
                    status="rejected",
                    reason="Condition relapsed after temporary stabilization",
                    validation_details=arrival_outcome.validation_details,
                    reward_modifier=0.001,
                    terminal=False,
                ),
                "Recovery guard: partial relapsed to rejected.",
            )

        return (
            ArrivalOutcome(
                status="accepted",
                reason="Condition stabilized after progressive treatment",
                validation_details=arrival_outcome.validation_details,
                reward_modifier=max(0.7, float(arrival_outcome.reward_modifier)),
                terminal=False,
            ),
            "Recovery guard: partial upgraded to accepted after continued care.",
        )

    def _build_last_info(
        self,
        reward: float,
        breakdown: RewardBreakdown,
        arrival_outcome: ArrivalOutcome,
    ) -> None:
        assert self.state_data is not None

        # Always expose a bounded task score snapshot so external validators
        # never see a missing grader and fallback to 0.0/1.0.
        grader_result = grade_task(
            task_id=self.state_data.task_id,
            difficulty=self.state_data.scenario_difficulty,
            objective=self.state_data.task_objective,
            trajectory=self.trajectory,
        )

        self.last_info = StepInfo(
            task_id=self.state_data.task_id,
            difficulty=self.state_data.scenario_difficulty,
            objective=self.state_data.task_objective,
            progress_score=self._progress_score(),
            reward_model=RewardModel(value=reward, breakdown=breakdown),
            grader=grader_result,
            last_action_error=None,
            outcome={
                "status": arrival_outcome.status,
                "reason": arrival_outcome.reason,
            },
        )

    def _record_trajectory(
        self,
        selected: HospitalState,
        arrival_outcome: ArrivalOutcome,
        reward: float,
        travel_time: float,
        dynamic_delay: float,
        original_traffic: str,
    ) -> None:
        assert self.state_data is not None
        self.trajectory.append(
            {
                "step": self.state_data.step,
                "state": {
                    "patient_condition": self.state_data.patient_condition,
                    "remaining_time_minutes": self.state_data.critical_time_limit_minutes,
                    "visited_hospitals": list(self.state_data.visited_hospitals),
                    "failed_hospitals": list(self.state_data.failed_hospitals),
                },
                "action": {
                    "hospital_id": selected.hospital_id,
                    "traffic_before": original_traffic,
                    "traffic_at_arrival": selected.traffic,
                },
                "outcome_status": arrival_outcome.status,
                "outcome_reason": arrival_outcome.reason,
                "reward": reward,
                "travel_time": travel_time,
                "dynamic_delay": dynamic_delay,
                "critical_limit": self.state_data.critical_time_limit_minutes,
                "specialization_match": self._specialization_match(selected),
                "suitability_score": max(MIN_REWARD, min(MAX_REWARD, arrival_outcome.validation_details.patient_suitability if arrival_outcome.validation_details else 0.5)),
            }
        )

    def _build_observation(self) -> Observation:
        assert self.state_data is not None

        last_outcome_obs = None
        if self.state_data.last_arrival_outcome and self.state_data.last_arrival_outcome.validation_details:
            last_outcome_obs = ArrivalOutcomeObservation(
                status=self.state_data.last_arrival_outcome.status,
                reason=self.state_data.last_arrival_outcome.reason,
                suitability_score=self.state_data.last_arrival_outcome.validation_details.patient_suitability,
            )

        return Observation(
            episode_id=self.state_data.episode_id,
            seed=self.state_data.seed,
            task_id=self.state_data.task_id,
            task_objective=self.state_data.task_objective,
            scenario_type=self.state_data.scenario_type,
            scenario_name=self.state_data.scenario_name,
            scenario_difficulty=self.state_data.scenario_difficulty,
            patient_condition=self.state_data.patient_condition,
            required_specialization=self.state_data.required_specialization,
            initial_critical_time_limit_minutes=self.state_data.initial_critical_time_limit_minutes,
            critical_time_limit_minutes=self.state_data.critical_time_limit_minutes,
            remaining_time_minutes=self.state_data.critical_time_limit_minutes,
            step=self.state_data.step,
            max_steps=self.state_data.max_steps,
            hospitals=[
                HospitalObservation(
                    hospital_id=h.hospital_id,
                    distance_km=h.distance_km,
                    icu=h.icu_display,
                    specialization=h.specialization,
                    traffic=h.traffic,
                )
                for h in self.state_data.hospitals
            ],
            previous_action=self.state_data.selected_hospital_id,
            ambulance_status=self.state_data.ambulance_status,
            current_location_context=self.state_data.current_location_context,
            visited_hospitals=list(self.state_data.visited_hospitals),
            failed_hospitals=list(self.state_data.failed_hospitals),
            recent_failed_hospitals=list(self.state_data.recent_failed_hospitals),
            failed_reasons=dict(self.state_data.failed_reasons),
            total_time_spent_minutes=self.state_data.total_time_spent_minutes,
            rerouting_reason=self.state_data.rerouting_reason,
            last_arrival_outcome=last_outcome_obs,
            explanation=list(self.state_data.explanation),
            memory_snapshot={k: v.model_dump() for k, v in self.state_data.memory.items()},
        )

    def _evolve_hospital_uncertainty(self) -> None:
        assert self.state_data is not None
        for hospital in self.state_data.hospitals:
            if self._rng.random() < 0.40:
                hospital.traffic = cast(Literal["low", "medium", "high"], self._traffic_shift(hospital.traffic, self.state_data.scenario_difficulty))

            if self._rng.random() < DifficultyModifier.get_icu_mismatch_probability(self.state_data.scenario_difficulty):
                hospital.icu_actual = not hospital.icu_actual

            if hospital.icu_actual:
                hospital.icu_display = "available" if self._rng.random() < 0.80 else "unknown"
            else:
                hospital.icu_display = "available" if self._rng.random() < 0.2 else "unknown"

    def _traffic_shift(self, current: str, difficulty: str) -> str:
        worsening_prob = {"easy": 0.12, "medium": 0.25, "hard": 0.38}.get(difficulty, 0.25)
        improving_prob = {"easy": 0.18, "medium": 0.10, "hard": 0.06}.get(difficulty, 0.10)

        if current == "low":
            if self._rng.random() < worsening_prob:
                return "medium"
            return "low"

        if current == "medium":
            roll = self._rng.random()
            if roll < worsening_prob:
                return "high"
            if roll < worsening_prob + improving_prob:
                return "low"
            return "medium"

        if self._rng.random() < improving_prob:
            return "medium"
        return "high"

    def _sample_scenario_for_difficulty(self, difficulty: str) -> tuple[dict[str, Any], str]:
        generators = [
            (generate_medical_case, "medical"),
            (generate_accident_case, "accident"),
            (generate_fire_case, "fire"),
        ]
        scenario: dict[str, Any] = {}
        scenario_type = "medical"
        for _ in range(64):
            generator, scenario_type = self._rng.choice(generators)
            scenario = generator(self._rng)
            if scenario["difficulty"] == difficulty:
                return scenario, scenario_type

        for generator, scenario_type in generators:
            scenario = generator(self._rng)
            if scenario["difficulty"] == difficulty:
                return scenario, scenario_type
        return scenario, scenario_type

    def _find_hospital(self, hospital_id: str) -> HospitalState | None:
        assert self.state_data is not None
        for hospital in self.state_data.hospitals:
            if hospital.hospital_id == hospital_id:
                return hospital
        return None

    def _load_memory(self) -> dict[str, LearningEntry]:
        text = self.memory_path.read_text(encoding="utf-8-sig").strip()
        raw = json.loads(text) if text else {}
        return {k: LearningEntry(**v) for k, v in raw.items()}

    def _update_learning_memory(self, hospital_id: str, success: bool, reward: float) -> None:
        memory = self._load_memory()
        entry = memory.get(hospital_id)
        if entry is None:
            entry = LearningEntry()

        if success:
            entry.success += 1
            entry.accepted += 1
        else:
            entry.fail += 1
            entry.rejected += 1

        total = entry.success + entry.fail
        if total == 1:
            entry.avg = max(0.0, min(1.0, reward))
        else:
            normalized_reward = max(0.0, min(1.0, reward))
            entry.avg = ((entry.avg * (total - 1)) + normalized_reward) / total

        memory[hospital_id] = entry
        serialized = {k: v.model_dump() for k, v in memory.items()}
        self.memory_path.write_text(json.dumps(serialized, indent=2), encoding="utf-8")

    def _progress_score(self) -> float:
        if not self.trajectory:
            return MIN_REWARD
        raw = sum(float(t["reward"]) for t in self.trajectory) / len(self.trajectory)
        return max(MIN_REWARD, min(MAX_REWARD, raw))

    def _failure_score(self) -> float:
        assert self.state_data is not None
        progress_component = self._progress_score()
        reward_component = max(MIN_REWARD, min(MAX_REWARD, self.state_data.reward))
        score = 0.15 + (0.35 * reward_component) + (0.25 * progress_component)
        return max(MIN_REWARD, min(MAX_REWARD, max(0.1, min(0.85, score))))

    def _success_score(self) -> float:
        assert self.state_data is not None
        progress_component = self._progress_score()
        reward_component = max(MIN_REWARD, min(MAX_REWARD, self.state_data.reward))
        total_steps = max(1, len(self.trajectory))
        rejected_steps = sum(1 for item in self.trajectory if item.get("outcome_status") == "rejected")
        route_quality = max(0.0, 1.0 - (rejected_steps / total_steps))
        score = (0.45 * reward_component) + (0.40 * progress_component) + (0.15 * route_quality)
        return max(MIN_REWARD, min(MAX_REWARD, max(0.25, min(0.99, score))))


ACDEEnvironment = EmergencyEnv