"""Synthetic case generator (Phase 2). Generates Case objects between start_date and end_date using: - CASE_TYPE_DISTRIBUTION - Monthly seasonality factors - Urgent case percentage - Court working days (CourtCalendar) Also provides CSV export/import helpers compatible with scripts. """ from __future__ import annotations import csv import random from dataclasses import dataclass from datetime import date, timedelta from pathlib import Path from typing import Iterable, List, Tuple from src.core.case import Case from src.data.config import ( CASE_TYPE_DISTRIBUTION, MONTHLY_SEASONALITY, URGENT_CASE_PERCENTAGE, ) from src.data.param_loader import load_parameters from src.utils.calendar import CourtCalendar def _month_iter(start: date, end: date) -> Iterable[Tuple[int, int]]: y, m = start.year, start.month while (y, m) <= (end.year, end.month): yield (y, m) if m == 12: y += 1 m = 1 else: m += 1 @dataclass class CaseGenerator: start: date end: date seed: int = 42 def generate( self, n_cases: int, stage_mix: dict | None = None, stage_mix_auto: bool = False, case_type_distribution: dict | None = None, ) -> List[Case]: random.seed(self.seed) cal = CourtCalendar() if stage_mix_auto: params = load_parameters() stage_mix = params.get_stage_stationary_distribution() stage_mix = stage_mix or {"ADMISSION": 1.0} # normalize explicitly total_mix = sum(stage_mix.values()) or 1.0 stage_mix = {k: v / total_mix for k, v in stage_mix.items()} # precompute cumulative for stage sampling stage_items = list(stage_mix.items()) scum = [] accs = 0.0 for _, p in stage_items: accs += p scum.append(accs) if scum: scum[-1] = 1.0 def sample_stage() -> str: if not stage_items: return "ADMISSION" r = random.random() for i, (st, _) in enumerate(stage_items): if r <= scum[i]: return st return stage_items[-1][0] # duration sampling helpers (lognormal via median & p90) def sample_stage_duration(stage: str) -> float: params = getattr(sample_stage_duration, "_params", None) if params is None: sample_stage_duration._params = load_parameters() params = sample_stage_duration._params med = params.get_stage_duration(stage, "median") p90 = params.get_stage_duration(stage, "p90") import math med = max(med, 1e-3) p90 = max(p90, med + 1e-6) z = 1.2815515655446004 sigma = max(1e-6, math.log(p90) - math.log(med)) / z mu = math.log(med) # Box-Muller normal sample u1 = max(random.random(), 1e-9) u2 = max(random.random(), 1e-9) z0 = ((-2.0 * math.log(u1)) ** 0.5) * math.cos(2.0 * math.pi * u2) val = math.exp(mu + sigma * z0) return max(1.0, val) # 1) Build monthly working-day lists and weights (seasonality * working days) month_days = {} month_weight = {} for y, m in _month_iter(self.start, self.end): days = cal.get_working_days_in_month(y, m) # restrict to [start, end] days = [d for d in days if self.start <= d <= self.end] if not days: continue month_days[(y, m)] = days month_weight[(y, m)] = MONTHLY_SEASONALITY.get(m, 1.0) * len(days) # normalize weights total_w = sum(month_weight.values()) if total_w == 0: return [] # 2) Allocate case counts per month (round, then adjust) alloc = {} for key, w in month_weight.items(): cnt = int(round(n_cases * (w / total_w))) alloc[key] = cnt # adjust rounding to total n_cases diff = n_cases - sum(alloc.values()) if diff != 0: # distribute the difference across months deterministically by key order keys = sorted(alloc.keys()) idx = 0 step = 1 if diff > 0 else -1 for _ in range(abs(diff)): alloc[keys[idx]] += step idx = (idx + 1) % len(keys) # 3) Sampling helpers (case type distribution) # Allow custom distribution override; default to historical (from config/EDA) if case_type_distribution is None: type_dist = dict(CASE_TYPE_DISTRIBUTION) else: # Validate and normalize user-provided distribution # Filter out zero/negative and None values valid_items = { str(k): float(v) for k, v in case_type_distribution.items() if v is not None and float(v) > 0.0 and str(k) } # Fallback to defaults if invalid or empty after filtering if not valid_items: type_dist = dict(CASE_TYPE_DISTRIBUTION) else: total = sum(valid_items.values()) # Normalize to 1.0 type_dist = {k: v / total for k, v in valid_items.items()} type_items = list(type_dist.items()) type_acc = [] cum = 0.0 for _, p in type_items: cum += p type_acc.append(cum) # ensure last is exactly 1.0 in case of rounding issues if type_acc: type_acc[-1] = 1.0 def sample_case_type() -> str: r = random.random() for i, (ct, _) in enumerate(type_items): if r <= type_acc[i]: return ct return type_items[-1][0] cases: List[Case] = [] seq = 0 for key in sorted(alloc.keys()): y, m = key days = month_days[key] if not days or alloc[key] <= 0: continue # simple distribution across working days of the month for _ in range(alloc[key]): filed = days[seq % len(days)] seq += 1 ct = sample_case_type() urgent = random.random() < URGENT_CASE_PERCENTAGE cid = f"{ct}/{filed.year}/{len(cases) + 1:05d}" init_stage = sample_stage() # For initial cases: they're filed on 'filed' date, started current stage on filed date # days_in_stage represents how long they've been in this stage as of simulation start # We sample a duration but cap it to not go before filed_date int(sample_stage_duration(init_stage)) # stage_start should be between filed_date and some time after # For simplicity: set stage_start = filed_date, case just entered this stage c = Case( case_id=cid, case_type=ct, filed_date=filed, current_stage=init_stage, is_urgent=urgent, ) c.stage_start_date = filed c.days_in_stage = 0 # Initialize realistic hearing history # Spread last hearings across past 7-30 days to simulate realistic court flow # This ensures constant stream of cases becoming eligible, not all at once days_since_filed = (self.end - filed).days if days_since_filed > 30: # Only if filed at least 30 days before end # Determine number of historical hearings based on age (roughly monthly) c.hearing_count = max(1, days_since_filed // 30) # Define pools of purposes bottleneck_purposes = [ "ISSUE SUMMONS", "FOR NOTICE", "AWAIT SERVICE OF NOTICE", "STAY APPLICATION PENDING", "FOR ORDERS", ] ripe_purposes = [ "ARGUMENTS", "HEARING", "FINAL ARGUMENTS", "FOR JUDGMENT", "EVIDENCE", ] # Build a small hearing history list on the Case.history c.history = [] # Generate hearing dates spaced across the case lifetime, ending 7-30 days before end days_before_end = random.randint(7, 30) last_hearing_date = self.end - timedelta(days=days_before_end) # approximate spacing if c.hearing_count == 1: hearing_dates = [last_hearing_date] else: span_days = max(days_since_filed - days_before_end, 30) step = max(1, span_days // c.hearing_count) hearing_dates = [ last_hearing_date - timedelta(days=step * i) for i in range(c.hearing_count - 1) ] hearing_dates = sorted(hearing_dates) + [last_hearing_date] # Assign purposes: earlier ones mixed; final one stage-dependent for i, hdt in enumerate(hearing_dates): if i == len(hearing_dates) - 1: # Final hearing purpose depends on stage and random bottleneck share if init_stage == "ADMISSION" and c.hearing_count < 3: purpose = ( random.choice(bottleneck_purposes) if random.random() < 0.4 else random.choice(ripe_purposes) ) elif init_stage in [ "ARGUMENTS", "ORDERS / JUDGMENT", "FINAL DISPOSAL", ]: purpose = random.choice(ripe_purposes) else: purpose = ( random.choice(bottleneck_purposes) if random.random() < 0.2 else random.choice(ripe_purposes) ) else: purpose = random.choice(bottleneck_purposes + ripe_purposes) was_heard = purpose not in ( "ISSUE SUMMONS", "FOR NOTICE", "AWAIT SERVICE OF NOTICE", ) c.history.append( { "date": hdt, "event": "hearing", "was_heard": was_heard, "outcome": "", "stage": init_stage, "purpose": purpose, } ) # Update aggregates from generated history c.last_hearing_date = last_hearing_date c.days_since_last_hearing = days_before_end c.last_hearing_purpose = ( c.history[-1]["purpose"] if c.history else None ) cases.append(c) return cases # CSV helpers ----------------------------------------------------------- @staticmethod def to_csv(cases: List[Case], out_path: Path) -> None: out_path.parent.mkdir(parents=True, exist_ok=True) with out_path.open("w", newline="") as f: w = csv.writer(f) w.writerow( [ "case_id", "case_type", "filed_date", "current_stage", "is_urgent", "hearing_count", "last_hearing_date", "days_since_last_hearing", "last_hearing_purpose", ] ) for c in cases: w.writerow( [ c.case_id, c.case_type, c.filed_date.isoformat(), c.current_stage, 1 if c.is_urgent else 0, c.hearing_count, c.last_hearing_date.isoformat() if c.last_hearing_date else "", c.days_since_last_hearing, c.last_hearing_purpose or "", ] ) @staticmethod def to_hearings_csv(cases: List[Case], out_path: Path) -> None: """Write flattened hearing histories for generated cases. Schema: case_id,date,stage,purpose,was_heard,event """ out_path.parent.mkdir(parents=True, exist_ok=True) with out_path.open("w", newline="") as f: w = csv.writer(f) w.writerow(["case_id", "date", "stage", "purpose", "was_heard", "event"]) for c in cases: for ev in getattr(c, "history", []) or []: if ev.get("event") == "hearing": w.writerow( [ c.case_id, (ev.get("date") or c.filed_date).isoformat(), ev.get("stage") or c.current_stage, ev.get("purpose", ""), 1 if ev.get("was_heard", False) else 0, ev.get("event"), ] ) @staticmethod def from_csv(path: Path) -> List[Case]: cases: List[Case] = [] with path.open("r", newline="") as f: r = csv.DictReader(f) for row in r: c = Case( case_id=row["case_id"], case_type=row["case_type"], filed_date=date.fromisoformat(row["filed_date"]), current_stage=row.get("current_stage", "ADMISSION"), is_urgent=(str(row.get("is_urgent", "0")) in ("1", "true", "True")), ) # Load hearing history if available if "hearing_count" in row and row["hearing_count"]: c.hearing_count = int(row["hearing_count"]) if "last_hearing_date" in row and row["last_hearing_date"]: c.last_hearing_date = date.fromisoformat(row["last_hearing_date"]) if "days_since_last_hearing" in row and row["days_since_last_hearing"]: c.days_since_last_hearing = int(row["days_since_last_hearing"]) if "last_hearing_purpose" in row and row["last_hearing_purpose"]: c.last_hearing_purpose = row["last_hearing_purpose"] cases.append(c) return cases