server/data_cleaning_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.

"""
Data Cleaning Environment Implementation.

Simulates real-world tabular data cleaning tasks where an AI agent must
identify and fix data quality issues: missing values, duplicates, type
errors, and outliers.

Tasks (easy → medium → hard):
  1. easy   – Fix missing values in a small CSV (impute/drop)
  2. medium – Remove duplicates AND fix type-cast errors
  3. hard   – Full pipeline: missing values + duplicates + outliers + normalise

API:
    POST /reset          – start a new episode (optionally pass {"task": "easy"|"medium"|"hard"})
    POST /step           – apply one cleaning operation
    GET  /state          – internal episode state
    GET  /tasks          – list tasks + action schema
    POST /grader         – score a completed episode
    POST /baseline       – run a deterministic baseline agent on all 3 tasks
"""

import copy
import random
import re
from typing import Any, Optional
from uuid import uuid4

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

try:
    from ..models import DataCleaningAction, DataCleaningObservation
except ImportError:
    from models import DataCleaningAction, DataCleaningObservation

# ---------------------------------------------------------------------------
# Synthetic datasets (deterministic, reproducible)
# ---------------------------------------------------------------------------

def _make_easy_dataset(seed: int = 0) -> list[dict]:
    """5-row CSV with 3 missing values to impute or drop."""
    rng = random.Random(seed)
    names = ["Alice", None, "Charlie", "Diana", None]
    ages  = [25, 30, None, 22, 28]
    rows  = []
    for n, a in zip(names, ages):
        rows.append({"name": n, "age": a, "score": round(rng.uniform(50, 100), 1)})
    return rows

def _make_medium_dataset(seed: int = 0) -> list[dict]:
    """8-row dataset with 2 duplicate rows and 2 type errors (age stored as str)."""
    base = [
        {"id": 1, "name": "Alice",   "age": 25,    "city": "Delhi"},
        {"id": 2, "name": "Bob",     "age": "30x",  "city": "Mumbai"},   # bad type
        {"id": 3, "name": "Charlie", "age": 22,    "city": "Pune"},
        {"id": 2, "name": "Bob",     "age": "30x",  "city": "Mumbai"},   # duplicate of row 2
        {"id": 4, "name": "Diana",   "age": "abc",  "city": "Chennai"},  # bad type
        {"id": 5, "name": "Eve",     "age": 27,    "city": "Kolkata"},
        {"id": 3, "name": "Charlie", "age": 22,    "city": "Pune"},      # duplicate of row 3
        {"id": 6, "name": "Frank",   "age": 33,    "city": "Hyderabad"},
    ]
    return copy.deepcopy(base)

def _make_hard_dataset(seed: int = 0) -> list[dict]:
    """12-row dataset: missing values + duplicates + outliers + unnormalised strings."""
    rng = random.Random(seed)
    rows = [
        {"id": 1,  "product": " Widget A ",  "price": 19.99,   "quantity": None, "category": "Electronics"},
        {"id": 2,  "product": "Gadget B",    "price": 9999.0,  "quantity": 50,   "category": "electronics"},  # outlier price + case
        {"id": 3,  "product": "Widget A ",   "price": 19.99,   "quantity": None, "category": "Electronics"},  # dup of 1
        {"id": 4,  "product": "Doohickey C", "price": 5.49,    "quantity": 200,  "category": "HOME"},
        {"id": 5,  "product": None,          "price": 12.00,   "quantity": 75,   "category": "Electronics"},
        {"id": 6,  "product": "Thingamajig", "price": -50.0,   "quantity": 10,   "category": "Home"},        # neg price outlier
        {"id": 7,  "product": "Widget A",    "price": 19.99,   "quantity": 30,   "category": "Electronics"}, # near-dup of 1
        {"id": 8,  "product": "Gizmo D",     "price": 7.25,    "quantity": None, "category": "Home"},
        {"id": 9,  "product": "Gadget B",    "price": 9.99,    "quantity": 50,   "category": "Electronics"},
        {"id": 10, "product": "Sprocket E",  "price": 3.15,    "quantity": 500,  "category": " Electronics"},
        {"id": 11, "product": "Cog F",       "price": 1.99,    "quantity": 1000, "category": "Home"},
        {"id": 12, "product": "Bolt G",      "price": 0.50,    "quantity": None, "category": "home"},
    ]
    return copy.deepcopy(rows)

# ---------------------------------------------------------------------------
# Grader helpers
# ---------------------------------------------------------------------------

def _grade_easy(rows: list[dict]) -> float:
    """Score 0-1: reward for each missing value resolved (imputed or dropped)."""
    if not rows:
        return 0.0
    missing = sum(1 for r in rows for v in r.values() if v is None)
    total_possible = 3  # 2 missing names + 1 missing age in original
    resolved = max(0, total_possible - missing)
    return round(resolved / total_possible, 4)

def _grade_medium(rows: list[dict]) -> float:
    """Score 0-1: half for no duplicates, half for no type errors in 'age'."""
    if not rows:
        return 0.0
    # Dedup check: unique (id, name) combos
    seen = set()
    dupes = 0
    for r in rows:
        key = (r.get("id"), r.get("name"))
        if key in seen:
            dupes += 1
        seen.add(key)
    dedup_score = 1.0 if dupes == 0 else max(0.0, 1.0 - dupes * 0.5)
    # Type-fix check: all 'age' values must be int or float
    type_errors = sum(1 for r in rows if not isinstance(r.get("age"), (int, float)))
    type_score = 1.0 if type_errors == 0 else max(0.0, 1.0 - type_errors * 0.5)

    return round((dedup_score + type_score) / 2, 4)

def _grade_hard(rows: list[dict]) -> float:
    """Score 0-1 across 4 sub-criteria."""
    if not rows:
        return 0.0
    scores = []

    # 1. No missing values
    missing = sum(1 for r in rows for v in r.values() if v is None)
    scores.append(1.0 if missing == 0 else max(0.0, 1.0 - missing * 0.1))

    # 2. No duplicate products (after stripping)
    products = [str(r.get("product", "")).strip().lower() for r in rows if r.get("product")]
    dupes = len(products) - len(set(products))
    scores.append(1.0 if dupes == 0 else max(0.0, 1.0 - dupes * 0.2))

    # 3. No price outliers (price must be 0 < price < 500)
    outliers = sum(1 for r in rows if isinstance(r.get("price"), (int, float))
                   and not (0 < r["price"] < 500))
    scores.append(1.0 if outliers == 0 else max(0.0, 1.0 - outliers * 0.3))

    # 4. Normalised category strings (no leading/trailing spaces, title-case)
    bad_cats = sum(
        1 for r in rows
        if isinstance(r.get("category"), str)
        and (r["category"] != r["category"].strip()
             or r["category"] != r["category"].title())
    )
    scores.append(1.0 if bad_cats == 0 else max(0.0, 1.0 - bad_cats * 0.15))

    return round(sum(scores) / len(scores), 4)

GRADERS = {
    "easy":   _grade_easy,
    "medium": _grade_medium,
    "hard":   _grade_hard,
}

DATASETS = {
    "easy":   _make_easy_dataset,
    "medium": _make_medium_dataset,
    "hard":   _make_hard_dataset,
}

# ---------------------------------------------------------------------------
# Environment
# ---------------------------------------------------------------------------

MAX_STEPS = {
    "easy": 10,
    "medium": 15,
    "hard": 25,
}

VALID_OPERATIONS_PER_TASK = {
    "easy": {
        "impute_mean",
        "impute_mode",
        "drop_missing_rows",
    },
    "medium": {
        "remove_duplicates",
        "fix_type_errors",
        "drop_missing_rows",
    },
    "hard": {
        "drop_missing_rows",
        "fill_quantity_mean",
        "remove_duplicates",
        "fix_type_errors",
        "remove_outliers",
        "normalize_text",
    },
}

class DataCleaningEnvironment(Environment):
    """
    Data Cleaning RL Environment.

    The agent receives a dirty dataset as observation and must apply
    cleaning operations step-by-step to maximise data quality.

    Each task presents a different level of difficulty and requires
    different cleaning operations.
    """
    SUPPORTS_CONCURRENT_SESSIONS: bool = True
    # -----------------------------------------------------------------------
    # Lifecycle
    # -----------------------------------------------------------------------

    def __init__(self):
        super().__init__()
        self._task: str = "easy"
        self._rows: list[dict] = []
        self._original_rows: list[dict] = []
        self._step_count: int = 0
        self._episode_id: str = str(uuid4())
        self._done: bool = False
        self._last_reward: float = 0.0
        self._applied_ops: list[str] = []

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        task: str = "easy",
        **kwargs: Any,
    ) -> DataCleaningObservation:
        """
        Reset the environment for a new episode.

        Args:
            seed: Random seed for dataset generation.
            episode_id: Optional custom episode identifier.
            task: One of "easy", "medium", "hard".
        """
        if task not in DATASETS:
            task = "easy"
        self._task = task
        _seed = seed if seed is not None else random.randint(0, 9999)
        self._rows = DATASETS[task](_seed)
        self._original_rows = copy.deepcopy(self._rows)
        self._step_count = 0
        self._episode_id = episode_id or str(uuid4())
        self._done = False
        self._last_reward = 0.0
        self._applied_ops = []

        return self._make_observation(reward=0.0, done=False)

    def step(
        self,
        action: DataCleaningAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> DataCleaningObservation:
        """
        Apply one cleaning operation to the dataset.

        Args:
            action: DataCleaningAction with fields:
                - operation (str): cleaning op name
                - column (str, optional): target column
        """
        if self._done:
            return self._make_observation(reward=0.0, done=True)

        self._step_count += 1
        op = action.operation.strip().lower()
        col = getattr(action, "column", None)

        valid_ops = VALID_OPERATIONS_PER_TASK[self._task]
        if op not in valid_ops:
            # Invalid operation → small penalty, episode continues
            reward = -0.05
            self._last_reward = reward
            obs = self._make_observation(reward=reward, done=False)
            obs.metadata["error"] = f"Unknown operation '{op}'. Valid: {sorted(valid_ops)}"
            return obs

        before_score = GRADERS[self._task](self._rows)
        self._apply_operation(op, col)
        after_score = GRADERS[self._task](self._rows)

        # Reward = improvement in quality score (partial progress signal)
        improvement = after_score - before_score
        # 0.01
        if improvement > 0:
            reward = round(improvement + 0.02, 4)   # bonus for improvement
        elif improvement < 0:
            reward = round(improvement - 0.02, 4)   # penalty for harming dataset
        else:
            reward = 0.0                            # neutral if no change
        
        if after_score >= 1.0:
            reward += 1.0                       # big bonus for perfect score (episode complete)
        
        self._last_reward = reward
        self._applied_ops.append(op)

        # Episode ends when perfect score or step limit reached
        max_steps = MAX_STEPS[self._task]
        done = (after_score >= 1.0) or (self._step_count >= max_steps)
        self._done = done

        return self._make_observation(reward=reward, done=done)

    # -----------------------------------------------------------------------
    # State
    # -----------------------------------------------------------------------

    @property
    def state(self) -> State:
        return State(
            episode_id=self._episode_id,
            step_count=self._step_count,
            task=self._task,
            done=self._done,
            current_score=GRADERS[self._task](self._rows),
            applied_ops=self._applied_ops,
            rows_remaining=len(self._rows),
        )

    # -----------------------------------------------------------------------
    # Grader (called via POST /grader)
    # -----------------------------------------------------------------------

    def grade(self) -> dict:
        """Return final grader score for the current episode."""
        score = GRADERS[self._task](self._rows)
        return {
            "task": self._task,
            "score": score,
            "steps_taken": self._step_count,
            "ops_applied": self._applied_ops,
        }

    # -----------------------------------------------------------------------
    # Tasks manifest (called via GET /tasks)
    # -----------------------------------------------------------------------

    @staticmethod
    def tasks() -> list[dict]:
        return [
            {
                "task_id": "easy",
                "description": "Fix missing values in a 5-row name/age/score table.",
                "difficulty": "easy",
                "max_steps": MAX_STEPS["easy"],
                "action_schema": {
                    "operation": {
                        "type": "string",
                        "enum": ["impute_mean", "impute_mode", "drop_missing_rows"],
                    },
                    "column": {"type": "string", "description": "Target column (optional)"},
                },
            },
            {
                "task_id": "medium",
                "description": "Remove duplicate rows and fix type errors in the 'age' column.",
                "difficulty": "medium",
                "max_steps": MAX_STEPS["medium"],
                "action_schema": {
                    "operation": {
                        "type": "string",
                        "enum": ["remove_duplicates", "fix_type_errors", "drop_missing_rows"],
                    },
                    "column": {"type": "string", "description": "Target column (optional)"},
                },
            },
            {
                "task_id": "hard",
                "description": (
                    "Full pipeline: fix missing values, remove duplicates, "
                    "remove price outliers, and normalise text fields."
                ),
                "difficulty": "hard",
                "max_steps": MAX_STEPS["hard"],
                "action_schema": {
                    "operation": {
                        "type": "string",
                        "enum": list(VALID_OPERATIONS_PER_TASK),
                    },
                    "column": {"type": "string", "description": "Target column (optional)"},
                },
            },
        ]

    # -----------------------------------------------------------------------
    # Baseline (called via POST /baseline)
    # -----------------------------------------------------------------------

    def run_baseline(self) -> dict:
        """
        Run a deterministic rule-based baseline agent on all 3 tasks.
        Returns scores dict compatible with the hackathon /baseline endpoint.
        """
        results = {}
        for task_id in ["easy", "medium", "hard"]:
            self.reset(seed=42, task=task_id)
            ops = _BASELINE_POLICIES[task_id]
            for op in ops:
                if not self._done:
                    self.step(DataCleaningAction(operation=op))
            results[task_id] = GRADERS[task_id](self._rows)
        return {"baseline_scores": results}

    # -----------------------------------------------------------------------
    # Internal helpers
    # -----------------------------------------------------------------------

    def _make_observation(self, reward: float, done: bool) -> DataCleaningObservation:
        missing_count = sum(1 for r in self._rows for v in r.values() if v is None)
        has_dupes = self._has_duplicates()
        has_outliers = self._has_outliers()
        score = GRADERS[self._task](self._rows)
        valid_ops = VALID_OPERATIONS_PER_TASK[self._task]
        return DataCleaningObservation(
            current_text=self._rows_to_text(),
            is_normalized=not has_outliers and missing_count == 0 and not has_dupes,
            html_found=False,
            remaining_typos=missing_count + (2 if has_dupes else 0) + (1 if has_outliers else 0),
            done=done,
            reward=reward,
            metadata={
                "task": self._task,
                "step": self._step_count,
                "rows": copy.deepcopy(self._rows),
                "missing_count": missing_count,
                "has_duplicates": has_dupes,
                "has_outliers": has_outliers,
                "quality_score": score,
                "valid_operations": sorted(valid_ops),
                "ops_already_applied": list(self._applied_ops),
            },
        )

    def _rows_to_text(self) -> str:
        if not self._rows:
            return "[]"
        cols = list(self._rows[0].keys())
        header = " | ".join(cols)
        lines = [header, "-" * len(header)]
        for r in self._rows:
            lines.append(" | ".join(str(r.get(c, "")) for c in cols))
        return "\n".join(lines)

    def _has_duplicates(self) -> bool:
        seen = set()
        for r in self._rows:
            key = tuple(
                (k, str(v).strip().lower() if isinstance(v, str) else v)
                for k, v in sorted(r.items())
            )
            if key in seen:
                return True
            seen.add(key)
        return False

    def _has_outliers(self) -> bool:
        return any(
            isinstance(r.get("price"), (int, float)) and not (0 < r["price"] < 500)
            for r in self._rows
        )

    def _apply_operation(self, op: str, column: Optional[str]):
        """Mutate self._rows according to the chosen operation."""

        if op == "drop_missing_rows":
            self._rows = [r for r in self._rows if all(v is not None for v in r.values())]

        elif op == "impute_mean":
            cols = [column] if column else _numeric_cols(self._rows)
            for c in cols:
                vals = [r[c] for r in self._rows if isinstance(r.get(c), (int, float))]
                if vals:
                    mean = sum(vals) / len(vals)
                    for r in self._rows:
                        if r.get(c) is None:
                            r[c] = round(mean, 2)

        elif op == "impute_mode":
            cols = [column] if column else _string_cols(self._rows)
            for c in cols:
                vals = [r[c] for r in self._rows if r.get(c) is not None]
                if vals:
                    mode = max(set(vals), key=vals.count)
                    for r in self._rows:
                        if r.get(c) is None:
                            r[c] = mode

        elif op == "remove_duplicates":
            seen: set = set()
            unique = []
            for r in self._rows:
                key = tuple(sorted(r.items()))
                if key not in seen:
                    seen.add(key)
                    unique.append(r)
            self._rows = unique

        elif op == "fix_type_errors":
            cols = [column] if column else _numeric_cols(self._rows) + ["age"]
            for c in set(cols):
                for r in self._rows:
                    val = r.get(c)
                    if val is not None and not isinstance(val, (int, float)):
                        val_str = str(val).strip()
                        if re.fullmatch(r"-?\d+(\.\d+)?", val_str):
                            r[c] = float(val_str)
                        else:
                            r[c] = None

        elif op == "remove_outliers":
            self._rows = [
                r for r in self._rows
                if not isinstance(r.get("price"), (int, float))
                or (0 < r["price"] < 500)
            ]

        elif op == "normalize_text":
            for r in self._rows:
                for k, v in list(r.items()):
                    if isinstance(v, str):
                        r[k] = v.strip().title()

        elif op == "fill_quantity_mean":
            vals = [r["quantity"] for r in self._rows if isinstance(r.get("quantity"), (int, float))]
            if vals:
                mean = round(sum(vals) / len(vals), 1)
                for r in self._rows:
                    if r.get("quantity") is None:
                        r["quantity"] = mean

# ---------------------------------------------------------------------------
# Baseline policies (deterministic rule-based agents)
# ---------------------------------------------------------------------------

_BASELINE_POLICIES: dict[str, list[str]] = {
    "easy":   ["impute_mean", "impute_mode", "drop_missing_rows"],
    "medium": ["remove_duplicates", "fix_type_errors", "drop_missing_rows"],
    "hard":   [
        "drop_missing_rows",
        "fill_quantity_mean",
        "remove_duplicates",
        "fix_type_errors",
        "remove_outliers",
        "normalize_text",
    ],
}

# ---------------------------------------------------------------------------
# Recommendation helper (guides the LLM toward correct next op)
# ---------------------------------------------------------------------------

# def _recommend_next(
#     task: str,
#     missing_count: int,
#     has_dupes: bool,
#     has_outliers: bool,
#     applied_ops: list,
# ) -> str:
#     """Return a plain-English hint for the LLM about the best next operation."""
#     applied = set(applied_ops)

#     if task == "easy":
#         if missing_count > 0:
#             return "There are missing values. Use impute_mean (numeric) or impute_mode (text)."
#         return "No issues remain. Episode should be complete."

#     if task == "medium":
#         if has_dupes and "remove_duplicates" not in applied:
#             return "Duplicate rows exist. Use remove_duplicates."
#         if "fix_type_errors" not in applied:
#             return "Non-numeric values in numeric columns. Use fix_type_errors."
#         if missing_count > 0 and "drop_missing_rows" not in applied:
#             return "Some values still missing after type fix. Use drop_missing_rows."
#         return "No issues remain. Episode should be complete."

#     # hard
#     if missing_count > 0 and "fill_quantity_mean" not in applied:
#         return "Missing quantity values. Use fill_quantity_mean first."
#     if missing_count > 0 and "drop_missing_rows" not in applied:
#         return "Missing product/category values. Use drop_missing_rows."
#     if has_outliers and "remove_outliers" not in applied:
#         return "Price outliers detected (price<=0 or price>=500). Use remove_outliers."
#     if "normalize_text" not in applied:
#         return "String columns have inconsistent casing/whitespace. Use normalize_text."
#     if has_dupes and "remove_duplicates" not in applied:
#         return "Duplicate rows remain. Use remove_duplicates."
#     return "All issues fixed. Episode should be complete."

# ---------------------------------------------------------------------------
# Column utility helpers
# ---------------------------------------------------------------------------

def _numeric_cols(rows: list[dict]) -> list[str]:
    if not rows:
        return []
    return [
        c for c in rows[0]
        if any(isinstance(r.get(c), (int, float)) for r in rows)
    ]

def _string_cols(rows: list[dict]) -> list[str]:
    if not rows:
        return []
    return [
        c for c in rows[0]
        if any(isinstance(r.get(c), str) for r in rows)
    ]