File size: 5,821 Bytes
4052d84
 
 
 
 
 
 
 
2bc545f
4052d84
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2bc545f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4052d84
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
"""
UndertriAI — Schema Drift Injector (Patronus AI Bonus Track)
Applies realistic schema drift to simulate IPC→BNSS transition and
cross-state FIR format variation, testing agent robustness.
"""

import random
import re
from typing import Any, Dict, Optional


# ---------------------------------------------------------------------------
# IPC → BNSS Section Remapping (Bharatiya Nyaya Sanhita 2023)
# ---------------------------------------------------------------------------

IPC_TO_BNSS: Dict[str, str] = {
    "120B": "61",   "121": "147",  "121A": "148",
    "302":  "103",  "304": "105",  "304B": "80",
    "306":  "108",  "307": "109",  "323": "115",
    "324":  "117",  "325": "118",  "326A": "124",
    "341":  "126",  "354": "74",   "354A": "75",
    "376":  "64",   "377": "377",  "379": "303",
    "380":  "305",  "392": "309",  "395": "310",
    "396":  "311",  "406": "316",  "420": "318",
    "465":  "336",  "467": "337",  "468": "338",
    "471":  "340",  "474": "341",  "498A": "85",
    "504":  "352",  "506": "351",  "509": "79",
    "34":   "3(5)", "149": "190",
}

BNSS_TO_IPC = {v: k for k, v in IPC_TO_BNSS.items()}


def remap_sections_to_bnss(sections: list) -> list:
    """Remap IPC section numbers to BNSS equivalents."""
    result = []
    for sec in sections:
        clean = sec.strip()
        mapped = IPC_TO_BNSS.get(clean, clean)
        result.append(f"BNS {mapped}")
    return result


# ---------------------------------------------------------------------------
# Regional FIR Format Templates
# ---------------------------------------------------------------------------

REGIONAL_HEADERS = {
    "Tamil Nadu": (
        "FIRST INFORMATION REPORT\n"
        "நம்பகமான தகவல் அறிக்கை (F.I.R)\n"
        "Tamil Nadu Police | State Crime Records Bureau\n"
        "---\n"
    ),
    "Kerala": (
        "FIRST INFORMATION REPORT\n"
        "Kerala Police | District Crime Records Bureau\n"
        "Registered under CrPC Section 154 / BNSS Section 173\n"
        "---\n"
    ),
    "Punjab": (
        "FIRST INFORMATION REPORT\n"
        "Punjab Police | District: __\n"
        "Thana No. __ | FIR No. __ | Year: __\n"
        "Under Sections: Bharatiya Nyaya Sanhita, 2023\n"
        "---\n"
    ),
    "Maharashtra": (
        "प्रथम सूचना अहवाल (F.I.R)\n"
        "FIRST INFORMATION REPORT\n"
        "Maharashtra Police | Taluka: __ | District: __\n"
        "---\n"
    ),
    "Assam": (
        "FIRST INFORMATION REPORT\n"
        "Assam Police | Registration No.: __\n"
        "Under Bharatiya Nagarik Suraksha Sanhita (BNSS) 2023\n"
        "---\n"
    ),
}

BNSS_PROCEDURAL_INSERTS = [
    "This matter is registered under the Bharatiya Nyaya Sanhita (BNS), 2023, "
    "which supersedes the Indian Penal Code (IPC), 1860 with effect from July 1, 2024.",

    "Bail conditions are governed by Chapter XXXV of the Bharatiya Nagarik Suraksha "
    "Sanhita (BNSS), 2023 (formerly CrPC), specifically Sections 479–483.",

    "Default bail under Section 479 BNSS: accused is eligible where custody exceeds "
    "one-half of the maximum sentence for the offence, subject to no special law restriction.",

    "The Bharatiya Sakshya Adhiniyam (BSA) 2023 governs evidentiary standards in this matter.",
]


# ---------------------------------------------------------------------------
# Main drift injector
# ---------------------------------------------------------------------------

def apply_schema_drift(episode: Dict[str, Any], drift_type: str = "auto") -> Dict[str, Any]:
    """
    Apply schema drift to an episode dict. Returns a modified copy.
    drift_type: "bnss" | "regional" | "auto" (chooses randomly)
    """
    import copy
    ep = copy.deepcopy(episode)
    region = ep.get("region", "")

    if drift_type == "auto":
        if region in REGIONAL_HEADERS:
            drift_type = random.choice(["bnss", "regional", "combined"])
        else:
            drift_type = "bnss"

    # --- Apply BNSS section remapping ---
    if drift_type in ("bnss", "combined"):
        ep["ipc_sections"] = remap_sections_to_bnss(ep.get("ipc_sections", []))
        ep["schema_variant"] = "bnss"

        # Insert BNSS procedural text into charge sheet
        insert = random.choice(BNSS_PROCEDURAL_INSERTS)
        ep["charge_sheet"] = insert + "\n\n" + ep.get("charge_sheet", "")

        # Replace "IPC Section" mentions in text with "BNS Section"
        ep["charge_sheet"] = re.sub(
            r"Section\s+(\d+[A-Z]?)\s+(?:of\s+)?(?:the\s+)?I\.P\.C\.?|IPC\s+[Ss]ection\s+(\d+[A-Z]?)",
            lambda m: f"BNS Section {IPC_TO_BNSS.get(m.group(1) or m.group(2), m.group(1) or m.group(2))}",
            ep["charge_sheet"],
        )

    # --- Apply regional FIR header ---
    if drift_type in ("regional", "combined") and region in REGIONAL_HEADERS:
        header = REGIONAL_HEADERS[region]
        ep["charge_sheet"] = header + ep.get("charge_sheet", "")
        ep["schema_variant"] = f"regional_{region.lower().replace(' ', '_')}"

    # Mark as drifted
    ep["schema_drifted"] = True
    return ep


def maybe_apply_drift(
    episode: Dict[str, Any],
    probability: float = 0.25,
    seed: Optional[int] = None,
) -> Dict[str, Any]:
    """Apply drift with the given probability (used during Stage 4 training).

    Args:
        probability: Chance of applying drift (0.0–1.0).
        seed: If provided, uses a local RNG so drift is fully deterministic.
              Pass the same seed that was used for episode selection.
    """
    if not episode.get("schema_drift_eligible"):
        return episode
    rng = random.Random(seed) if seed is not None else random
    if rng.random() < probability:
        return apply_schema_drift(episode)
    return episode