server/my_env_environment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
MedCodeRL Environment Implementation.

Medical Coding & Billing Compliance environment where agents must:
1. Assign correct ICD-10 diagnosis codes
2. Assign correct CPT procedure codes
3. Make billing compliance decisions (approve/reject/review)
4. Provide clinical reasoning
5. Identify compliance risks

Contains: environment logic, deterministic grader, shaped rewards, action validation.
"""

import json
import os
import re
from typing import Any, Dict, List, Optional, Set, Tuple
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..models import MedAction, MedObservation
except ImportError:
    from models import MedAction, MedObservation


# ──────────────────────────────────────────────
#  Task loader
# ──────────────────────────────────────────────

TASKS_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tasks")


def _load_task_cases(difficulty: str) -> List[Dict[str, Any]]:
    """Load clinical cases for a given difficulty level."""
    filepath = os.path.join(TASKS_DIR, f"{difficulty}.json")
    if not os.path.exists(filepath):
        return []
    with open(filepath, "r", encoding="utf-8") as f:
        data = json.load(f)
    return data.get("cases", [])


# ──────────────────────────────────────────────
#  Action validation
# ──────────────────────────────────────────────

ICD10_PATTERN = re.compile(r"^[A-Z]\d{2}(\.\d{1,4})?$", re.IGNORECASE)
CPT_PATTERN = re.compile(r"^\d{5}$")
HCPCS_PATTERN = re.compile(r"^[A-Z]\d{4}$", re.IGNORECASE)


def _validate_action(action_dict: dict) -> Tuple[bool, List[str]]:
    """Validate an action dict. Returns (is_valid, error_list)."""
    errors: List[str] = []

    diag = action_dict.get("diagnosis_codes", [])
    if not isinstance(diag, list) or len(diag) == 0:
        errors.append("At least one diagnosis code is required")
    elif len(diag) > 5:
        errors.append("Maximum 5 diagnosis codes allowed")
    else:
        for code in diag:
            if not ICD10_PATTERN.match(str(code).strip()):
                errors.append(f"Invalid ICD-10 format: {code}")

    proc = action_dict.get("procedure_codes", [])
    if isinstance(proc, list):
        for code in proc:
            c = str(code).strip()
            if not CPT_PATTERN.match(c) and not HCPCS_PATTERN.match(c):
                errors.append(f"Invalid CPT/HCPCS format: {code}")

    decision = action_dict.get("decision", "")
    if str(decision).lower() not in ("approve", "reject", "review"):
        errors.append(f"Invalid decision: {decision}")

    confidence = action_dict.get("confidence", -1)
    try:
        conf_val = float(confidence)
        if conf_val < 0.0 or conf_val > 1.0:
            errors.append("confidence must be 0.0-1.0")
    except (TypeError, ValueError):
        errors.append("confidence must be a number")

    reasoning = action_dict.get("reasoning", "")
    if not isinstance(reasoning, str) or len(reasoning) < 15:
        errors.append("reasoning must be at least 15 characters")

    return len(errors) == 0, errors


# ──────────────────────────────────────────────
#  Deterministic grader
# ──────────────────────────────────────────────

def _set_similarity(predicted: List[str], ground_truth: List[str]) -> float:
    """Jaccard similarity between two code sets."""
    if not ground_truth and not predicted:
        return 1.0
    pred_set: Set[str] = set(c.strip().upper() for c in predicted if c)
    gt_set: Set[str] = set(c.strip().upper() for c in ground_truth if c)
    if not gt_set:
        return 1.0 if not pred_set else 0.0
    if not pred_set:
        return 0.0
    intersection = pred_set & gt_set
    union = pred_set | gt_set
    return len(intersection) / len(union) if union else 1.0


def _partial_code_match(predicted: List[str], ground_truth: List[str]) -> float:
    """Partial-credit matching for medical codes (prefix similarity)."""
    if not ground_truth:
        return 1.0 if not predicted else 0.0
    if not predicted:
        return 0.0
    pred_list = [c.strip().upper() for c in predicted if c]
    gt_list = [c.strip().upper() for c in ground_truth if c]
    total = 0.0
    for gt_code in gt_list:
        best = 0.0
        for pred_code in pred_list:
            if pred_code == gt_code:
                best = 1.0
                break
            gt_base = gt_code.split(".")[0]
            pred_base = pred_code.split(".")[0]
            if gt_base == pred_base:
                best = max(best, 0.5)
            elif len(gt_base) >= 3 and gt_base[:3] == pred_base[:3]:
                best = max(best, 0.25)
        total += best
    return total / len(gt_list)


def _grade(action_dict: dict, ground_truth: dict) -> Dict[str, float]:
    """
    Deterministic grading — 6 weighted components → score in [0.0, 1.0].

    Weights: diagnosis 35%, procedure 20%, decision 25%, reasoning 10%,
    risk flags 5%, confidence calibration 5%.
    """
    pred_diag = action_dict.get("diagnosis_codes", [])
    gt_diag = ground_truth.get("diagnosis_codes", [])
    pred_proc = action_dict.get("procedure_codes", [])
    gt_proc = ground_truth.get("procedure_codes", [])
    pred_dec = str(action_dict.get("decision", "")).lower()
    gt_dec = str(ground_truth.get("decision", "")).lower()
    pred_reasoning = str(action_dict.get("reasoning", ""))
    pred_risk = action_dict.get("risk_flags", [])
    gt_risk = ground_truth.get("risk_flags", [])
    pred_conf = float(action_dict.get("confidence", 0.5))

    # 1. Diagnosis codes (35%)
    diag_exact = _set_similarity(pred_diag, gt_diag)
    diag_partial = _partial_code_match(pred_diag, gt_diag)
    diag_score = 0.6 * diag_exact + 0.4 * diag_partial

    # 2. Procedure codes (20%)
    proc_exact = _set_similarity(pred_proc, gt_proc)
    proc_partial = _partial_code_match(pred_proc, gt_proc)
    proc_score = 0.6 * proc_exact + 0.4 * proc_partial

    # 3. Decision (25%)
    if pred_dec == gt_dec:
        dec_score = 1.0
    elif pred_dec == "review" and gt_dec in ("approve", "reject"):
        dec_score = 0.3
    elif pred_dec in ("approve", "reject") and gt_dec == "review":
        dec_score = 0.2
    else:
        dec_score = 0.0

    # 4. Reasoning quality (10%)
    reasoning_lower = pred_reasoning.lower()
    r_score = 0.0
    if len(pred_reasoning) >= 20:
        r_score += 0.3
    if len(pred_reasoning) >= 50:
        r_score += 0.2
    med_terms = [
        "icd", "cpt", "diagnosis", "procedure", "coding", "compliance",
        "medical", "clinical", "treatment", "patient", "billing",
        "authorization", "insurance", "modifier", "documentation",
        "justified", "appropriate", "medically necessary", "guideline",
    ]
    r_score += min(0.5, sum(1 for t in med_terms if t in reasoning_lower) * 0.1)
    r_score = min(1.0, r_score)

    # 5. Risk flags (5%)
    risk_score = _set_similarity(pred_risk, gt_risk)

    # 6. Confidence calibration (5%)
    correctness = diag_score * 0.5 + proc_score * 0.3 + dec_score * 0.2
    conf_score = max(0.0, 1.0 - abs(pred_conf - correctness) * 2.0)

    total = (
        diag_score * 0.35
        + proc_score * 0.20
        + dec_score * 0.25
        + r_score * 0.10
        + risk_score * 0.05
        + conf_score * 0.05
    )

    return {
        "score": round(min(1.0, max(0.0, total)), 4),
        "diagnosis_accuracy": round(diag_score, 4),
        "procedure_accuracy": round(proc_score, 4),
        "decision_accuracy": round(dec_score, 4),
        "reasoning_quality": round(r_score, 4),
        "risk_identification": round(risk_score, 4),
        "confidence_calibration": round(conf_score, 4),
    }


# ──────────────────────────────────────────────
#  Shaped reward engine
# ──────────────────────────────────────────────

def _compute_reward(action_dict: dict, ground_truth: dict, difficulty: str = "easy") -> Dict:
    """
    Shaped reward = base_grade + bonuses − penalties.

    Penalties: upcoding, undercoding, wrong denial/approval, unnecessary procedure,
    missing primary code, low confidence.
    Bonuses: perfect diagnosis, good reasoning, all risk flags.
    """
    grade_result = _grade(action_dict, ground_truth)
    base = grade_result["score"]

    pred_diag = set(c.strip().upper() for c in action_dict.get("diagnosis_codes", []) if c)
    gt_diag = set(c.strip().upper() for c in ground_truth.get("diagnosis_codes", []) if c)
    gt_diag_list = [c.strip().upper() for c in ground_truth.get("diagnosis_codes", []) if c]
    pred_proc = set(c.strip().upper() for c in action_dict.get("procedure_codes", []) if c)
    gt_proc = set(c.strip().upper() for c in ground_truth.get("procedure_codes", []) if c)
    pred_dec = str(action_dict.get("decision", "")).lower()
    gt_dec = str(ground_truth.get("decision", "")).lower()
    pred_conf = float(action_dict.get("confidence", 0.5))

    penalties: Dict[str, float] = {}
    bonuses: Dict[str, float] = {}

    # Penalties
    if len(pred_proc) > len(gt_proc) + 1:
        penalties["upcoding"] = -0.15
    if gt_diag and len(pred_diag & gt_diag) < len(gt_diag) * 0.5:
        penalties["undercoding"] = -0.10
    if pred_dec == "reject" and gt_dec == "approve":
        penalties["wrong_denial"] = -0.20
    if pred_dec == "approve" and gt_dec == "reject":
        penalties["wrong_approval"] = -0.25
    if pred_proc - gt_proc:
        penalties["unnecessary_procedure"] = -0.10
    if gt_diag_list and gt_diag_list[0] not in pred_diag:
        penalties["missing_primary_code"] = -0.15
    if pred_conf < 0.2:
        penalties["low_confidence"] = -0.05

    # Bonuses
    if grade_result["diagnosis_accuracy"] >= 0.99:
        bonuses["perfect_diagnosis"] = 0.05
    if grade_result["reasoning_quality"] >= 0.8:
        bonuses["good_reasoning"] = 0.03
    if grade_result["risk_identification"] >= 0.99:
        bonuses["all_risk_flags"] = 0.05

    diff_mult = {"easy": 0.8, "medium": 1.0, "hard": 1.2}.get(difficulty, 1.0)
    total_penalty = sum(penalties.values()) * diff_mult
    total_bonus = sum(bonuses.values())
    final = max(0.0, min(1.0, base + total_penalty + total_bonus))

    feedback_parts = []
    if penalties:
        feedback_parts.append(f"Penalties: {', '.join(penalties.keys())}")
    if bonuses:
        feedback_parts.append(f"Bonuses: {', '.join(bonuses.keys())}")
    if not penalties and not bonuses:
        feedback_parts.append("Clean submission.")

    return {
        "score": round(final, 4),
        "breakdown": {
            "base_grade": round(base, 4),
            "total_penalty": round(total_penalty, 4),
            "total_bonus": round(total_bonus, 4),
            "penalties": {k: round(v, 4) for k, v in penalties.items()},
            "bonuses": {k: round(v, 4) for k, v in bonuses.items()},
            "grade_components": grade_result,
        },
        "feedback": " | ".join(feedback_parts),
    }


# ──────────────────────────────────────────────
#  Core Environment
# ──────────────────────────────────────────────

class MyEnvironment(Environment):
    """
    MedCodeRL — Medical Coding & Billing Compliance Environment.

    90 realistic clinical cases (30 easy, 30 medium, 30 hard) covering:
    - Straightforward coding (easy)
    - Multi-diagnosis with comorbidities (medium)
    - Compliance dilemmas: upcoding, unbundling, fraud, ethical edge cases (hard)

    OpenEnv-compliant: reset() / step() / state property.
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        """Initialize the MedCodeRL environment."""
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._reset_count = 0

        # Load all task cases
        self._task_cases: Dict[str, List[Dict]] = {}
        self._case_index: Dict[str, int] = {"easy": 0, "medium": 0, "hard": 0}
        for diff in ("easy", "medium", "hard"):
            self._task_cases[diff] = _load_task_cases(diff)

        self._current_case: Optional[Dict] = None
        self._current_difficulty: str = "easy"
        self._done: bool = True
        self._action_history: List[dict] = []

    def _pick_case(self, task_id: Optional[str] = None) -> Dict:
        """Select a case by difficulty or specific case_id."""
        if task_id in ("easy", "medium", "hard"):
            difficulty = task_id
        elif task_id:
            # Search for specific case_id
            for diff in ("easy", "medium", "hard"):
                for case in self._task_cases.get(diff, []):
                    if case["id"] == task_id:
                        self._current_difficulty = diff
                        return case
            raise ValueError(f"Case not found: {task_id}")
        else:
            # Cycle through difficulties
            total = sum(self._case_index.values())
            difficulty = ["easy", "medium", "hard"][total % 3]

        cases = self._task_cases.get(difficulty, [])
        if not cases:
            raise ValueError(f"No cases for difficulty: {difficulty}")
        idx = self._case_index[difficulty] % len(cases)
        self._case_index[difficulty] = idx + 1
        self._current_difficulty = difficulty
        return cases[idx]

    def _build_observation(self, case: Dict, done: bool = False,
                           reward: Optional[float] = None,
                           reward_breakdown: Optional[Dict] = None,
                           feedback: str = "") -> MedObservation:
        """Build a MedObservation from a case dict."""
        inp = case.get("input", case)
        return MedObservation(
            case_id=case.get("id", ""),
            difficulty=case.get("difficulty", self._current_difficulty),
            clinical_note=inp.get("clinical_note", ""),
            symptoms=inp.get("symptoms", []),
            treatments=inp.get("treatments", []),
            insurance_type=inp.get("insurance_type", "Private"),
            prior_auth_required=inp.get("prior_auth_required", False),
            treatment_cost=inp.get("treatment_cost", "low"),
            patient_age=inp.get("patient_age", 0),
            patient_sex=inp.get("patient_sex", "M"),
            provider_specialty=inp.get("provider_specialty", ""),
            visit_type=inp.get("visit_type", "outpatient"),
            comorbidities=inp.get("comorbidities", []),
            lab_results=inp.get("lab_results"),
            medications=inp.get("medications", []),
            reward_breakdown=reward_breakdown,
            feedback=feedback,
            done=done,
            reward=reward,
            metadata={
                "step_count": self._state.step_count,
                "difficulty": self._current_difficulty,
            },
        )

    def reset(self, task_id: Optional[str] = None, **kwargs) -> MedObservation:
        """
        Reset the environment to a new episode.

        Args:
            task_id: 'easy', 'medium', 'hard', or a specific case_id.

        Returns:
            MedObservation with the clinical case to code.
        """
        # Accept task_id from kwargs if not provided directly
        if task_id is None:
            task_id = kwargs.get("task_id")

        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._reset_count += 1
        self._done = False
        self._action_history = []

        self._current_case = self._pick_case(task_id)

        return self._build_observation(self._current_case, done=False, reward=0.0)

    def step(self, action: MedAction) -> MedObservation:  # type: ignore[override]
        """
        Execute a step: grade the agent's medical coding action.

        Args:
            action: MedAction with diagnosis codes, procedure codes, decision, etc.

        Returns:
            MedObservation with reward and grading breakdown.
        """
        if self._done:
            return self._build_observation(
                self._current_case or {},
                done=True,
                reward=0.0,
                feedback="Episode already done. Call reset().",
            )

        self._state.step_count += 1

        # Convert action to dict
        action_dict = {
            "diagnosis_codes": action.diagnosis_codes,
            "procedure_codes": action.procedure_codes,
            "decision": action.decision,
            "confidence": action.confidence,
            "reasoning": action.reasoning,
            "modifier_codes": action.modifier_codes,
            "risk_flags": action.risk_flags,
        }

        # Validate
        is_valid, errors = _validate_action(action_dict)

        if not is_valid:
            self._action_history.append({"action": action_dict, "valid": False})
            if self._state.step_count >= 3:
                self._done = True
            return self._build_observation(
                self._current_case or {},
                done=self._done,
                reward=0.0,
                feedback=f"Invalid action: {'; '.join(errors)}",
            )

        # Grade against ground truth
        ground_truth = self._current_case.get("ground_truth", {})
        reward_result = _compute_reward(action_dict, ground_truth, self._current_difficulty)

        self._action_history.append({"action": action_dict, "valid": True})
        self._done = True  # single-step episode for valid actions

        score = reward_result["score"]

        return self._build_observation(
            self._current_case or {},
            done=True,
            reward=score,
            reward_breakdown=reward_result.get("breakdown"),
            feedback=reward_result.get("feedback", ""),
        )

    @property
    def state(self) -> State:
        """
        Get the current environment state.

        Returns:
            State with episode_id and step_count
        """
        return self._state