data_gen.py

"""Synthetic data generation and CSV loading for claims and adjusters.

Generates data/claims.csv and data/adjusters.csv. To use real data, produce
CSVs with the same columns (or adapt load_claims / load_adjusters).

Time columns are stored as "HH:MM" strings for readability and parsed to
integer minutes since midnight on load.
"""

from __future__ import annotations

import csv
import os
import random
from dataclasses import dataclass, field

import config


def hhmm_to_min(s: str) -> int:
    h, m = s.split(":")
    return int(h) * 60 + int(m)


def min_to_hhmm(m: int) -> str:
    return f"{m // 60:02d}:{m % 60:02d}"


@dataclass
class Claim:
    claim_id: str
    lat: float
    lon: float
    peril: str
    priority: int  # 1 = must inspect today, 2 = high, 3 = normal
    window_start: int  # minutes since midnight
    window_end: int
    service_minutes: int
    age_days: int = 0  # days since FNOL; escalates the drop penalty
    assigned_to: str | None = None  # upstream pre-assignment (sequencer)
    # Split availability (VRP with multiple time windows): optional
    # extra windows beyond [window_start, window_end]. None or empty =
    # the classic single window. Service must start AND finish inside
    # one chosen window, same semantics as the primary pair.
    extra_windows: list[tuple[int, int]] | None = None


@dataclass
class Adjuster:
    adjuster_id: str
    name: str
    home_lat: float
    home_lon: float
    skills: list[str] = field(default_factory=list)
    shift_start: int = 8 * 60
    shift_end: int = 17 * 60
    # Service territory: only take claims within this road-mile radius of
    # home (optional adjusters.csv column; blank/absent = no limit).
    max_radius_miles: float | None = None


# ---------------------------------------------------------------------------
# Generation
# ---------------------------------------------------------------------------

ADJUSTER_NAMES = [
    "Alice Chen", "Bob Rivera", "Carol Nguyen", "David Okafor", "Erin Walsh",
    "Frank Kim", "Grace Patel", "Hank Torres", "Ivy Johnson", "Jack Murphy",
]

# (window label, start, end, weight)
WINDOW_CHOICES = [
    ("all-day", "08:00", "17:00", 0.4),
    ("morning", "08:00", "12:00", 0.2),
    ("afternoon", "12:00", "17:00", 0.2),
    ("narrow-am", "09:00", "11:00", 0.1),
    ("narrow-pm", "13:00", "15:00", 0.1),
]


def _rand_point(rng: random.Random) -> tuple[float, float]:
    lat = config.REGION_CENTER[0] + rng.uniform(-1, 1) * config.REGION_SPREAD_DEG
    lon = config.REGION_CENTER[1] + rng.uniform(-1, 1) * config.REGION_SPREAD_DEG
    return round(lat, 5), round(lon, 5)


def generate(n_claims: int = 25, n_adjusters: int = 5, seed: int = 42,
             data_dir: str = config.DATA_DIR) -> None:
    rng = random.Random(seed)
    os.makedirs(data_dir, exist_ok=True)

    # --- Adjusters: 2-3 skills each, and every peril covered by >= 2 people
    adjusters = []
    for i in range(n_adjusters):
        lat, lon = _rand_point(rng)
        n_skills = rng.choice([2, 2, 3])
        skills = rng.sample(config.PERILS, n_skills)
        shift_start, shift_end = rng.choice([("08:00", "17:00"),
                                             ("07:00", "16:00"),
                                             ("08:00", "17:00")])
        adjusters.append({
            "adjuster_id": f"ADJ-{i + 1:02d}",
            "name": ADJUSTER_NAMES[i % len(ADJUSTER_NAMES)],
            "home_lat": lat,
            "home_lon": lon,
            "skills": "|".join(sorted(skills)),
            "shift_start": shift_start,
            "shift_end": shift_end,
        })
    # Patch coverage: make sure each peril appears in at least 2 skill sets.
    for peril in config.PERILS:
        holders = [a for a in adjusters if peril in a["skills"].split("|")]
        while len(holders) < 2:
            victim = rng.choice([a for a in adjusters if a not in holders])
            victim["skills"] = "|".join(sorted(victim["skills"].split("|") + [peril]))
            holders.append(victim)

    # --- Claims
    claims = []
    for i in range(n_claims):
        lat, lon = _rand_point(rng)
        peril = rng.choice(config.PERILS)
        r = rng.random()
        if r < 0.12:
            priority = config.PRIORITY_MUST_TODAY
        elif r < 0.40:
            priority = config.PRIORITY_HIGH
        else:
            priority = config.PRIORITY_NORMAL
        _, ws, we, _ = rng.choices(
            WINDOW_CHOICES, weights=[w for *_, w in WINDOW_CHOICES])[0]
        # Serious losses take longer to inspect.
        if priority == config.PRIORITY_MUST_TODAY:
            service = rng.choice([120, 150, 180])
        else:
            service = rng.choice([60, 75, 90, 120])
        claims.append({
            "claim_id": f"CLM-{i + 1:03d}",
            "lat": lat,
            "lon": lon,
            "peril": peril,
            "priority": priority,
            "window_start": ws,
            "window_end": we,
            "service_minutes": service,
        })

    with open(os.path.join(data_dir, "adjusters.csv"), "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(adjusters[0].keys()))
        w.writeheader()
        w.writerows(adjusters)
    with open(os.path.join(data_dir, "claims.csv"), "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(claims[0].keys()))
        w.writeheader()
        w.writerows(claims)


# ---------------------------------------------------------------------------
# Loading
# ---------------------------------------------------------------------------

def claim_windows(c: Claim) -> list[tuple[int, int]]:
    """All availability windows of a claim, primary first, sorted."""
    wins = [(c.window_start, c.window_end)]
    for lo, hi in (c.extra_windows or []):
        wins.append((int(lo), int(hi)))
    return sorted(set(wins))


def arrival_ranges(c: Claim) -> list[tuple[int, int]]:
    """Allowed SERVICE-START ranges: inside a window, finishing by its
    end (the same tightening every backend applies), merged/sorted."""
    ranges = []
    for lo, hi in claim_windows(c):
        ranges.append((lo, max(lo, hi - c.service_minutes)))
    ranges.sort()
    merged = [list(ranges[0])]
    for lo, hi in ranges[1:]:
        if lo <= merged[-1][1] + 1:
            merged[-1][1] = max(merged[-1][1], hi)
        else:
            merged.append([lo, hi])
    return [(lo, hi) for lo, hi in merged]


def win_str(c: Claim) -> str:
    """Display text for all of a claim's availability windows."""
    return " | ".join(f"{min_to_hhmm(lo)}-{min_to_hhmm(hi)}"
                      for lo, hi in claim_windows(c))


def _parse_windows(text):
    """'08:00-10:30|16:00-17:00' -> list of (start_min, end_min)."""
    wins = []
    for part in text.split("|"):
        lo, hi = part.strip().split("-")
        wins.append((hhmm_to_min(lo.strip()), hhmm_to_min(hi.strip())))
    return wins


def load_claims(data_dir: str = config.DATA_DIR) -> list[Claim]:
    claims = []
    with open(os.path.join(data_dir, "claims.csv"), newline="") as f:
        for row in csv.DictReader(f):
            claims.append(Claim(
                claim_id=row["claim_id"],
                lat=float(row["lat"]),
                lon=float(row["lon"]),
                peril=row["peril"],
                priority=int(row["priority"]),
                window_start=hhmm_to_min(row["window_start"]),
                window_end=hhmm_to_min(row["window_end"]),
                service_minutes=int(row["service_minutes"]),
                age_days=int(row.get("age_days") or 0),
                assigned_to=(row.get("assigned_to") or "").strip() or None,
            ))
            wtext = (row.get("windows") or "").strip()
            if wtext:
                wins = _parse_windows(wtext)
                c = claims[-1]
                c.window_start, c.window_end = wins[0]
                c.extra_windows = wins[1:] or None
    return claims


def load_adjusters(data_dir: str = config.DATA_DIR) -> list[Adjuster]:
    adjusters = []
    with open(os.path.join(data_dir, "adjusters.csv"), newline="") as f:
        for row in csv.DictReader(f):
            adjusters.append(Adjuster(
                adjuster_id=row["adjuster_id"],
                name=row["name"],
                home_lat=float(row["home_lat"]),
                home_lon=float(row["home_lon"]),
                skills=row["skills"].split("|"),
                shift_start=hhmm_to_min(row["shift_start"]),
                shift_end=hhmm_to_min(row["shift_end"]),
                max_radius_miles=(float(row["max_radius_miles"])
                                  if row.get("max_radius_miles")
                                  else None),
            ))
    return adjusters


# ---------------------------------------------------------------------------
# Writing (round-trip of the loaders, optional columns included)
# ---------------------------------------------------------------------------

def write_claims_csv(claims: list[Claim], path: str) -> str:
    """Write claims in the exact upload schema, optional columns included."""
    with_assign = any(c.assigned_to for c in claims)
    with_windows = any(c.extra_windows for c in claims)
    with open(path, "w", newline="") as f:
        w = csv.writer(f)
        header = ["claim_id", "lat", "lon", "peril", "priority",
                  "window_start", "window_end", "service_minutes",
                  "age_days"]
        if with_assign:
            header.append("assigned_to")
        if with_windows:
            header.append("windows")
        w.writerow(header)
        for c in claims:
            row = [c.claim_id, c.lat, c.lon, c.peril, c.priority,
                   min_to_hhmm(c.window_start),
                   min_to_hhmm(c.window_end),
                   c.service_minutes, c.age_days]
            if with_assign:
                row.append(c.assigned_to or "")
            if with_windows:
                row.append("|".join(f"{min_to_hhmm(lo)}-{min_to_hhmm(hi)}"
                                    for lo, hi in claim_windows(c))
                           if c.extra_windows else "")
            w.writerow(row)
    return path


def write_adjusters_csv(adjusters: list[Adjuster], path: str) -> str:
    """Write adjusters in the exact upload schema."""
    with open(path, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["adjuster_id", "name", "home_lat", "home_lon",
                    "skills", "shift_start", "shift_end",
                    "max_radius_miles"])
        for a in adjusters:
            w.writerow([a.adjuster_id, a.name, a.home_lat, a.home_lon,
                        "|".join(a.skills),
                        min_to_hhmm(a.shift_start),
                        min_to_hhmm(a.shift_end),
                        "" if a.max_radius_miles is None
                        else a.max_radius_miles])
    return path