server/environment.py

"""
PharmaAgent — environment.py
Clinical Decision RL Environment

Design principles:
  1. Every reward signal is derived exclusively from DrugBank data.
     No hardcoded drug lists, no opinion-based scoring.
  2. Patient cases are generated dynamically from the live database.
     The agent cannot memorise them — it must reason.
  3. The LLM (Groq) is called once per episode, at the very end,
     only to format a human-readable clinical summary of the final
     regimen. It has zero influence on any reward value.
  4. If DrugBank has no data on a drug, the reward is 0.0 — not a
     guess. Unknown = unverified = no credit.
"""

import random
import re
import sqlite3
import os
from typing import Tuple, Optional
from .models import Action, Observation, State

DB_PATH = os.environ.get("DB_PATH", "drugbank.db")

# ── Reward ceiling (used by inference.py for normalisation) ───────────────
# Per episode maximum achievable:
#   0.30  diagnosis   (DrugBank indication match, full score)
#   0.60  drugs       (up to 3 drugs x 0.20 each from DrugBank)
#   0.30  DDI safety  (catching a real contraindicated/major pair)
#   0.30  finalize    (clean regimen, safety checks done, no banned drugs)
MAX_EPISODE_REWARD = 1.5

# Cap DDI checks that earn reward — prevents reward farming
MAX_REWARDED_DDI_CHECKS = 3

# Severity keywords parsed from DrugBank interaction description text.
# Ordered most-to-least severe so the first match wins.
_SEVERITY_PATTERNS = [
    ("contraindicated", re.compile(r"\bcontraindicated\b", re.I)),
    ("major",           re.compile(r"\b(major|serious|severe|life.?threatening|fatal)\b", re.I)),
    ("moderate",        re.compile(r"\b(moderate|caution|monitor)\b", re.I)),
    ("minor",           re.compile(r"\b(minor|mild|small)\b", re.I)),
]


# ── DB helpers ────────────────────────────────────────────────────────────

def _get_db() -> Optional[sqlite3.Connection]:
    if not os.path.exists(DB_PATH) or os.path.getsize(DB_PATH) == 0:
        return None
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def db_get_drug(drug_name: str) -> Optional[dict]:
    """Fetch a drug record from DrugBank by exact name (case-insensitive)."""
    conn = _get_db()
    if not conn:
        return None
    try:
        row = conn.execute(
            """SELECT name, indication, mechanism_of_action,
                      pharmacodynamics, metabolism, status, type
               FROM drugs WHERE LOWER(name) = LOWER(?)""",
            (drug_name,)
        ).fetchone()
        return dict(row) if row else None
    finally:
        conn.close()


def db_check_interaction(drug1: str, drug2: str) -> Optional[dict]:
    """
    Look up a DDI pair in DrugBank.
    Returns dict with keys: description, severity_label.
    severity_label is derived by parsing DrugBank's own description text.
    """
    conn = _get_db()
    if not conn:
        return None
    try:
        row = conn.execute(
            """SELECT description FROM interactions
               WHERE (LOWER(drug1_name) = LOWER(?) AND LOWER(drug2_name) = LOWER(?))
                  OR (LOWER(drug1_name) = LOWER(?) AND LOWER(drug2_name) = LOWER(?))
               LIMIT 1""",
            (drug1, drug2, drug2, drug1)
        ).fetchone()
        if not row:
            return None
        desc = row["description"] or ""
        severity = "unknown"
        for label, pattern in _SEVERITY_PATTERNS:
            if pattern.search(desc):
                severity = label
                break
        return {"description": desc, "severity_label": severity}
    finally:
        conn.close()


def db_get_interactions_for_drug(drug_name: str, limit: int = 10) -> list:
    """Return known interaction partners for a drug (used in case generation)."""
    conn = _get_db()
    if not conn:
        return []
    try:
        rows = conn.execute(
            """SELECT drug2_name AS partner, description
               FROM interactions
               WHERE LOWER(drug1_name) = LOWER(?)
               AND description != ''
               LIMIT ?""",
            (drug_name, limit)
        ).fetchall()
        return [dict(r) for r in rows]
    finally:
        conn.close()


def db_drugs_for_indication(keywords: list, limit: int = 40) -> list:
    """
    Return approved small-molecule drugs whose indication text contains
    at least one of the given keywords.
    """
    conn = _get_db()
    if not conn:
        return []
    try:
        results = []
        seen = set()
        for kw in keywords:
            rows = conn.execute(
                """SELECT name, indication, mechanism_of_action, status
                   FROM drugs
                   WHERE LOWER(indication) LIKE LOWER(?)
                   AND status LIKE '%approved%'
                   AND type = 'small molecule'
                   AND indication != ''
                   LIMIT ?""",
                (f"%{kw}%", limit)
            ).fetchall()
            for r in rows:
                if r["name"] not in seen:
                    seen.add(r["name"])
                    results.append(dict(r))
        return results
    finally:
        conn.close()


# ── Condition seeds ───────────────────────────────────────────────────────
# Plain-English conditions + keywords used to query DrugBank.
# These are established clinical facts, not scoring opinions.

_CONDITION_SEEDS = [
    {
        "condition": "Hypertension",
        "symptoms": ["persistent headache", "elevated blood pressure", "dizziness", "blurred vision"],
        "indication_keywords": ["hypertension", "high blood pressure", "antihypertensive"],
        "diagnosis_keywords": ["hypertension", "blood pressure", "antihypertensive", "vascular"],
    },
    {
        "condition": "Type 2 Diabetes Mellitus",
        "symptoms": ["excessive thirst", "frequent urination", "fatigue", "blurred vision"],
        "indication_keywords": ["type 2 diabetes", "diabetes mellitus", "hyperglycemia", "glycemic"],
        "diagnosis_keywords": ["diabetes", "hyperglycemia", "insulin resistance", "glycemic", "glucose"],
    },
    {
        "condition": "Chronic Heart Failure",
        "symptoms": ["shortness of breath on exertion", "ankle swelling", "fatigue", "orthopnoea"],
        "indication_keywords": ["heart failure", "cardiac failure", "congestive heart"],
        "diagnosis_keywords": ["heart failure", "cardiac", "ejection fraction", "congestive"],
    },
    {
        "condition": "Rheumatoid Arthritis",
        "symptoms": ["symmetric joint pain", "morning stiffness", "joint swelling", "fatigue"],
        "indication_keywords": ["rheumatoid arthritis", "rheumatoid", "autoimmune arthritis"],
        "diagnosis_keywords": ["rheumatoid", "arthritis", "autoimmune", "synovitis", "joint inflammation"],
    },
    {
        "condition": "Bronchial Asthma",
        "symptoms": ["recurrent wheeze", "chest tightness", "shortness of breath", "nocturnal cough"],
        "indication_keywords": ["asthma", "bronchial asthma", "bronchospasm", "airway obstruction"],
        "diagnosis_keywords": ["asthma", "bronchospasm", "airway", "bronchial", "respiratory"],
    },
    {
        "condition": "Epilepsy",
        "symptoms": ["recurrent seizures", "transient loss of consciousness", "post-ictal confusion"],
        "indication_keywords": ["epilepsy", "seizure", "anticonvulsant", "antiepileptic"],
        "diagnosis_keywords": ["epilepsy", "seizure", "anticonvulsant", "antiepileptic", "ictal"],
    },
    {
        "condition": "Hypothyroidism",
        "symptoms": ["fatigue", "weight gain", "cold intolerance", "constipation", "dry skin"],
        "indication_keywords": ["hypothyroidism", "thyroid deficiency", "levothyroxine"],
        "diagnosis_keywords": ["hypothyroidism", "thyroid", "TSH", "levothyroxine", "thyroid hormone"],
    },
    {
        "condition": "Major Depressive Disorder",
        "symptoms": ["persistent low mood", "anhedonia", "insomnia", "fatigue", "poor concentration"],
        "indication_keywords": ["depression", "major depressive", "antidepressant"],
        "diagnosis_keywords": ["depression", "depressive", "antidepressant", "mood", "serotonin"],
    },
    {
        "condition": "Peptic Ulcer Disease",
        "symptoms": ["epigastric pain", "nausea", "bloating", "pain relieved by food"],
        "indication_keywords": ["peptic ulcer", "gastric ulcer", "duodenal ulcer", "H. pylori"],
        "diagnosis_keywords": ["peptic ulcer", "gastric", "H. pylori", "proton pump", "acid"],
    },
    {
        "condition": "Atrial Fibrillation",
        "symptoms": ["palpitations", "irregular heartbeat", "dyspnoea on exertion", "fatigue"],
        "indication_keywords": ["atrial fibrillation", "AF", "anticoagulation", "rate control"],
        "diagnosis_keywords": ["atrial fibrillation", "arrhythmia", "anticoagul", "rate control", "AF"],
    },
]


# ── Dynamic case generation ───────────────────────────────────────────────

def generate_case_from_db() -> dict:
    """
    Build a patient case dynamically from DrugBank.

    Steps:
      1. Pick a random condition seed.
      2. Query DrugBank for all approved drugs indicated for that condition.
      3. Pick one of those drugs as the patient's existing medication.
      4. Fetch that drug's known interaction partners from DrugBank.
      5. Mark partners with contraindicated/major severity as avoid_drugs.

    The result is a case where:
      - indicated_drug_names comes from DrugBank (19,842 drugs)
      - avoid_drugs comes from DrugBank (2.9M interaction pairs)
      - No hardcoded drug names anywhere

    Falls back to a minimal case if the DB is unavailable.
    """
    seed = random.choice(_CONDITION_SEEDS)
    indicated_drugs = db_drugs_for_indication(seed["indication_keywords"], limit=40)

    if not indicated_drugs:
        # DB unavailable — minimal fallback
        return {
            "id": f"fallback_{seed['condition'].replace(' ', '_')}",
            "condition": seed["condition"],
            "symptoms": seed["symptoms"],
            "existing_medications": [],
            "indication_keywords": seed["indication_keywords"],
            "diagnosis_keywords": seed["diagnosis_keywords"],
            "indicated_drug_names": [],
            "avoid_drugs": [],
            "existing_med_interactions": {},
        }

    existing_med = random.choice(indicated_drugs)
    existing_med_name = existing_med["name"]

    raw_interactions = db_get_interactions_for_drug(existing_med_name, limit=10)

    avoid_drugs = []
    existing_med_interactions = {}
    for ix in raw_interactions:
        partner = ix["partner"]
        desc = ix["description"] or ""
        severity = "unknown"
        for label, pattern in _SEVERITY_PATTERNS:
            if pattern.search(desc):
                severity = label
                break
        existing_med_interactions[partner] = {"description": desc, "severity": severity}
        if severity in ("contraindicated", "major"):
            avoid_drugs.append(partner)

    return {
        "id": f"dynamic_{seed['condition'].replace(' ', '_')}_{random.randint(1000, 9999)}",
        "condition": seed["condition"],
        "symptoms": seed["symptoms"],
        "existing_medications": [existing_med_name],
        "indication_keywords": seed["indication_keywords"],
        "diagnosis_keywords": seed["diagnosis_keywords"],
        "indicated_drug_names": [d["name"] for d in indicated_drugs],
        "avoid_drugs": avoid_drugs,
        "existing_med_interactions": existing_med_interactions,
    }


# ── Scoring ───────────────────────────────────────────────────────────────

def score_diagnosis(proposed: str, case: dict) -> Tuple[float, str]:
    """
    Score against DrugBank-derived clinical keywords.
    Full 0.30 for >= 2 keyword matches, 0.15 for 1, 0.0 for none.
    """
    proposed_lower = proposed.lower()
    keywords = case.get("diagnosis_keywords", [])
    matches = [kw for kw in keywords if kw.lower() in proposed_lower]

    if len(matches) >= 2:
        return 0.30, (
            f"Diagnosis supported — matched: {', '.join(matches[:3])}."
        )
    elif len(matches) == 1:
        return 0.15, (
            f"Partial diagnosis — matched '{matches[0]}'. Consider the full clinical picture."
        )
    else:
        return 0.00, (
            "Diagnosis does not align with the presenting symptoms."
        )


def score_drug_selection(drug: str, case: dict) -> Tuple[float, str]:
    """
    Score using DrugBank as sole arbiter.

    -0.25  Drug is in avoid_drugs (derived from real DrugBank DDI severity)
     0.20  DrugBank confirms indication for this condition
     0.05  Drug exists in DrugBank but indication doesn't match
     0.00  Drug not found in DrugBank (hallucinated names earn nothing)
    """
    drug_lower = drug.lower()
    avoid = [d.lower() for d in case.get("avoid_drugs", [])]

    if drug_lower in avoid:
        existing = case.get("existing_medications", ["existing medication"])
        return -0.25, (
            f"SAFETY: {drug} has a contraindicated or major interaction "
            f"with {', '.join(existing)} per DrugBank. Not added to regimen."
        )

    db_record = db_get_drug(drug)

    if db_record is None:
        return 0.00, (
            f"'{drug}' not found in DrugBank. "
            f"Only verified drugs earn reward. Check spelling or use the full approved name."
        )

    indication_text = (db_record.get("indication") or "").lower()
    keywords = case.get("indication_keywords", [])
    matched_kw = [kw for kw in keywords if kw.lower() in indication_text]
    in_indicated_list = drug in case.get("indicated_drug_names", [])

    if matched_kw or in_indicated_list:
        return 0.20, (
            f"{drug} is indicated for this condition per DrugBank "
            f"({db_record.get('status', 'approved')})."
        )
    else:
        return 0.05, (
            f"{drug} is a known approved drug in DrugBank, but its primary "
            f"indication does not clearly match this condition."
        )


def score_ddi_check(drug1: str, drug2: str, case: dict,
                    checked_interactions: list) -> Tuple[float, str]:
    """
    Score using DrugBank's 2.9M interaction table exclusively.
    Severity is parsed from DrugBank description text, never hardcoded.
    """
    if drug1.lower() == drug2.lower():
        return 0.0, f"Both drugs are the same ('{drug1}'). Provide two different drugs."

    already_checked = {
        frozenset([i["drug1"].lower(), i["drug2"].lower()])
        for i in checked_interactions
    }
    this_pair = frozenset([drug1.lower(), drug2.lower()])

    if this_pair in already_checked:
        return 0.0, f"{drug1} x {drug2} already checked. No additional reward."

    if len(already_checked) >= MAX_REWARDED_DDI_CHECKS:
        return 0.0, f"Reward cap ({MAX_REWARDED_DDI_CHECKS} checks) reached."

    avoid = [d.lower() for d in case.get("avoid_drugs", [])]
    is_flagged = drug1.lower() in avoid or drug2.lower() in avoid
    db_result = db_check_interaction(drug1, drug2)

    if is_flagged and db_result:
        sev = db_result["severity_label"].upper()
        return 0.30, (
            f"CRITICAL DDI [{sev}]: {drug1} x {drug2}\n"
            f"DrugBank: {db_result['description'][:250]}\n"
            f"This combination must be avoided for this patient."
        )
    elif is_flagged:
        return 0.15, (
            f"{drug1} or {drug2} is flagged as potentially dangerous "
            f"with this patient's existing medications. Caution warranted."
        )
    elif db_result:
        sev = db_result["severity_label"].upper()
        return 0.15, (
            f"Interaction found [{sev}]: {drug1} x {drug2}\n"
            f"DrugBank: {db_result['description'][:200]}"
        )
    else:
        return 0.05, (
            f"No interaction found in DrugBank between {drug1} and {drug2}."
        )


def score_finalize(state: State) -> Tuple[float, str]:
    """
    Score the completed regimen. All checks against DrugBank-derived case data.
    """
    case = state.patient_case
    selected_lower = [d.lower() for d in state.selected_drugs]
    indicated = [d.lower() for d in case.get("indicated_drug_names", [])]
    avoid = [d.lower() for d in case.get("avoid_drugs", [])]
    existing_meds = case.get("existing_medications", [])

    indicated_hits = [d for d in selected_lower if d in indicated]
    bad_drugs = [d for d in selected_lower if d in avoid]
    has_diagnosis = state.proposed_diagnosis is not None
    checked_safety = len(state.checked_interactions) > 0
    skipped_safety = (
        len(existing_meds) > 0
        and len(state.selected_drugs) > 0
        and not checked_safety
    )

    reward = 0.0
    parts = []

    if indicated_hits:
        reward += 0.10
        parts.append(f"Regimen contains DrugBank-indicated drug(s): {', '.join(d.title() for d in indicated_hits)}")
    else:
        parts.append("No DrugBank-indicated drugs found in the final regimen.")

    if not bad_drugs:
        reward += 0.10
        parts.append("No contraindicated drugs in regimen.")
    else:
        penalty = 0.20 * len(bad_drugs)
        reward -= penalty
        parts.append(f"Contraindicated drug(s) in regimen: {', '.join(d.title() for d in bad_drugs)} (penalty -{penalty:.2f})")

    if has_diagnosis:
        reward += 0.05
        parts.append(f"Diagnosis established: '{state.proposed_diagnosis}'.")

    if checked_safety:
        reward += 0.05
        parts.append(f"{len(state.checked_interactions)} DDI check(s) performed.")
    elif skipped_safety:
        reward -= 0.10
        parts.append(f"Safety penalty: patient on {', '.join(existing_meds)} but no DDI checks performed.")

    feedback = "Final Regimen Evaluation\n" + "\n".join(f"  {p}" for p in parts)
    return round(reward, 3), feedback


# ── Environment ───────────────────────────────────────────────────────────

class PharmaAgentEnvironment:

    def reset(self) -> Tuple[State, Observation]:
        case = generate_case_from_db()
        state = State(patient_case=case)
        obs = Observation(
            step=0,
            phase="triage",
            patient_case={
                "condition": case["condition"],
                "symptoms": case["symptoms"],
                "existing_medications": case["existing_medications"],
                "condition_hint": "Unknown — diagnose from symptoms only.",
            },
            feedback=(
                f"New patient case.\n"
                f"Symptoms: {', '.join(case['symptoms'])}\n"
                f"Current medications: {', '.join(case['existing_medications']) or 'None'}\n\n"
                f"Begin with action_type='diagnose'."
            ),
            valid_options=["diagnose"],
            reward_so_far=0.0,
            done=False,
        )
        return state, obs

    def step(self, state: State, action: Action) -> Tuple[State, Observation, float, bool]:
        state.step += 1
        reward = 0.0
        feedback = ""
        done = False
        next_phase = state.current_phase
        valid_options = ["diagnose", "select_drug", "check_ddi", "finalize"]

        if action.action_type == "diagnose":
            reward, feedback = score_diagnosis(action.value, state.patient_case)
            state.proposed_diagnosis = action.value
            state.cumulative_reward += reward
            state.phase_rewards["triage"] = reward
            next_phase = "selection"
            valid_options = ["select_drug", "finalize"]
            feedback += (
                "\n\nNext: action_type='select_drug'. "
                "Only drugs found in DrugBank earn reward. "
                "Add 2-3 drugs, then use 'check_ddi' before finalising."
            )

        elif action.action_type == "select_drug":
            reward, feedback = score_drug_selection(action.value, state.patient_case)
            avoid = [d.lower() for d in state.patient_case.get("avoid_drugs", [])]
            if action.value.lower() not in avoid and action.value not in state.selected_drugs:
                state.selected_drugs.append(action.value)
            state.cumulative_reward += reward
            next_phase = "safety"
            valid_options = ["select_drug", "check_ddi", "finalize"]
            feedback += (
                f"\n\nCurrent regimen: {', '.join(state.selected_drugs) or 'None'}"
            )

        elif action.action_type == "check_ddi":
            parts = action.value.replace(" vs ", ",").replace(" and ", ",").split(",")
            if len(parts) >= 2:
                d1, d2 = parts[0].strip(), parts[1].strip()
                reward, feedback = score_ddi_check(
                    d1, d2, state.patient_case, state.checked_interactions
                )
                state.checked_interactions.append({"drug1": d1, "drug2": d2, "reward": reward})
                state.cumulative_reward += reward
                next_phase = "safety"
                valid_options = ["select_drug", "check_ddi", "finalize"]
                feedback += f"\n\nCurrent regimen: {', '.join(state.selected_drugs) or 'None'}"
            else:
                feedback = "Provide two drug names separated by comma, 'vs', or 'and'."
                reward = 0.0
                valid_options = ["check_ddi"]

        elif action.action_type == "finalize":
            reward, feedback = score_finalize(state)
            state.cumulative_reward += reward
            state.phase_rewards["finalize"] = reward
            done = True
            next_phase = "done"
            valid_options = []
            feedback += (
                f"\n\nEpisode complete.\n"
                f"Final regimen: {', '.join(state.selected_drugs) or 'None'}\n"
                f"Total reward: {round(state.cumulative_reward, 3)} / {MAX_EPISODE_REWARD}"
            )

        else:
            feedback = (
                f"Unknown action_type '{action.action_type}'. "
                f"Valid: diagnose, select_drug, check_ddi, finalize."
            )

        if state.step >= state.max_steps and not done:
            done = True
            next_phase = "done"
            valid_options = []
            feedback += f"\n\nStep limit ({state.max_steps}) reached."

        state.current_phase = next_phase
        state.done = done

        obs = Observation(
            step=state.step,
            phase=next_phase,
            patient_case={
                "condition": state.patient_case.get("condition"),
                "symptoms": state.patient_case["symptoms"],
                "existing_medications": state.patient_case["existing_medications"],
                "proposed_diagnosis": state.proposed_diagnosis,
                "current_regimen": state.selected_drugs,
            },
            feedback=feedback,
            valid_options=valid_options,
            reward_so_far=round(state.cumulative_reward, 3),
            done=done,
        )

        return state, obs, round(reward, 3), done