server/cleaning_environment.py

"""FSDS Cleaning Environment.

This is an MCP/OpenEnv environment designed for the Silver-layer of the user's
Full-Stack Data Science thesis: a DS Agent cleans and validates messy tabular
business data while a QA/PCS layer enforces deterministic quality gates.

The environment exposes tools rather than a fixed action schema so that LLM
agents can inspect the dataset, choose cleaning operations, run gates, and
submit a final cleaned table.
"""

from __future__ import annotations

from copy import deepcopy
from dataclasses import dataclass, field
from typing import Any, Optional
from uuid import uuid4

import numpy as np
import pandas as pd
from fastmcp import FastMCP
from sklearn.model_selection import KFold
from sklearn.preprocessing import LabelEncoder
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor

from fsds_cleaning_env.dataset_generators import (
    SIZE_MEDIUM,
    make_dataset_factory,
)
from fsds_cleaning_env.reward import (
    FinalRewardInput,
    StepRewardInput,
    TOOL_ERROR_REWARD,
    compute_final_reward,
    compute_quality_gate_bonus,
    compute_step_reward,
)


try:
    from openenv.core.env_server.mcp_environment import MCPEnvironment
    from openenv.core.env_server.types import Action, Observation, State
except ImportError:
    from openenv.core.env_server.mcp_environment import MCPEnvironment
    from openenv.core.env_server.types import Action, Observation, State


INVALID_TOKENS = {"", " ", "unknown", "UNKNOWN", "n/a", "N/A", "null", "NULL", "?", "--"}

# Observation metadata schema version for downstream agents.
OBSERVATION_SCHEMA_VERSION = "1.0"

# Allowed operations for the apply_cleaning_operation tool. This defines the
# logical action space that agents should use.
AVAILABLE_OPERATIONS = [
    "drop_duplicates",
    "replace_invalid_with_null",
    "cast_numeric",
    "cast_datetime",
    "impute_numeric",
    "impute_categorical",
    "normalize_categories",
    "clip_outliers_iqr",
]


@dataclass
class TaskSpec:
    task_id: str
    title: str
    objective: str
    business_context: str
    target_column: str
    task_type: str  # classification | regression
    dataset_factory: Any
    expected_types: dict[str, str]
    required_ops: list[dict[str, Any]]
    allowed_columns_to_drop: list[str] = field(default_factory=list)
    min_retention_ratio: float = 0.85
    notes: list[str] = field(default_factory=list)


@dataclass
class EpisodeData:
    spec: TaskSpec
    raw_df: pd.DataFrame
    working_df: pd.DataFrame
    operation_log: list[dict[str, Any]] = field(default_factory=list)
    last_gate_report: dict[str, Any] = field(default_factory=dict)
    total_reward: float = 0.0
    submitted: bool = False
    max_steps: int = 18


TASKS: dict[str, TaskSpec] = {
    "ecommerce_mobile": TaskSpec(
        task_id="ecommerce_mobile",
        title="Mobile conversion cleaning",
        objective="Prepare a mobile conversion table for downstream churn/conversion modeling.",
        business_context=(
            "Marketing needs a trustworthy Silver table to investigate why mobile conversions fell. "
            "Clean the dataset without destroying row retention."
        ),
        target_column="converted",
        task_type="classification",
        dataset_factory=make_dataset_factory("ecommerce_mobile", n_rows=SIZE_MEDIUM),
        expected_types={
            "session_id": "int64",
            "device_os": "str",
            "customer_id": "str",
            "country": "str",
            "items_in_cart": "float64",
            "order_value": "float64",
            "event_date": "datetime64[us]",
            "converted": "int64",
        },
        required_ops=[
            {"operation": "replace_invalid_with_null", "column": "country"},
            {"operation": "replace_invalid_with_null", "column": "items_in_cart"},
            {"operation": "replace_invalid_with_null", "column": "device_os"},
            {"operation": "cast_numeric", "column": "items_in_cart"},
            {"operation": "cast_numeric", "column": "order_value"},
            {"operation": "impute_numeric", "column": "items_in_cart"},
            {"operation": "impute_numeric", "column": "order_value"},
            {"operation": "clip_outliers_iqr", "column": "items_in_cart"},
            {"operation": "clip_outliers_iqr", "column": "order_value"},
            {"operation": "normalize_categories", "column": "device_os"},
            {"operation": "normalize_categories", "column": "country"},
            {"operation": "impute_categorical", "column": "device_os"},
            {"operation": "impute_categorical", "column": "country"},
            {"operation": "cast_datetime", "column": "event_date"},
            {"operation": "drop_duplicates"},
        ],
        notes=[
            "Preserve the target column.",
            "Do not drop high-value rows unless a quality gate demands it.",
        ],
    ),
    "subscription_churn": TaskSpec(
        task_id="subscription_churn",
        title="Subscription churn cleaning",
        objective="Create a clean churn-ready subscriber table.",
        business_context=(
            "Retention wants to model churn risk, but the raw CRM extract has duplicated customers, "
            "invalid tokens, and numeric columns stored as strings."
        ),
        target_column="churned",
        task_type="classification",
        dataset_factory=make_dataset_factory("subscription_churn", n_rows=SIZE_MEDIUM),
        expected_types={
            "customer_key": "str",
            "age": "float64",
            "monthly_spend": "float64",
            "plan_type": "str",
            "tenure_months": "float64",
            "payment_method": "str",
            "churned": "int64",
        },
        required_ops=[
            {"operation": "replace_invalid_with_null", "column": "monthly_spend"},
            {"operation": "replace_invalid_with_null", "column": "age"},
            {"operation": "replace_invalid_with_null", "column": "tenure_months"},
            {"operation": "replace_invalid_with_null", "column": "payment_method"},
            {"operation": "cast_numeric", "column": "age"},
            {"operation": "cast_numeric", "column": "monthly_spend"},
            {"operation": "cast_numeric", "column": "tenure_months"},
            {"operation": "impute_numeric", "column": "age"},
            {"operation": "impute_numeric", "column": "monthly_spend"},
            {"operation": "impute_numeric", "column": "tenure_months"},
            {"operation": "clip_outliers_iqr", "column": "monthly_spend"},
            {"operation": "normalize_categories", "column": "plan_type"},
            {"operation": "normalize_categories", "column": "payment_method"},
            {"operation": "impute_categorical", "column": "plan_type"},
            {"operation": "impute_categorical", "column": "payment_method"},
            {"operation": "drop_duplicates"},
        ],
        notes=["Monthly spend contains a severe outlier that should be handled, not ignored."],
    ),
    "delivery_eta": TaskSpec(
        task_id="delivery_eta",
        title="Delivery ETA cleaning",
        objective="Clean a last-mile delivery table before ETA modeling.",
        business_context=(
            "Operations needs a reliable table to predict ETA. The raw export mixes city aliases, "
            "rating gaps, duplicated rows and suspicious distance values."
        ),
        target_column="delivery_time_minutes",
        task_type="regression",
        dataset_factory=make_dataset_factory("delivery_eta", n_rows=SIZE_MEDIUM),
        expected_types={
            "route_id": "str",
            "city": "str",
            "driver_rating": "float64",
            "delivery_distance_km": "float64",
            "late_deliveries_last_30d": "float64",
            "vehicle_type": "str",
            "delivery_time_minutes": "float64",
        },
        required_ops=[
            {"operation": "replace_invalid_with_null", "column": "driver_rating"},
            {"operation": "replace_invalid_with_null", "column": "late_deliveries_last_30d"},
            {"operation": "replace_invalid_with_null", "column": "city"},
            {"operation": "replace_invalid_with_null", "column": "vehicle_type"},
            {"operation": "cast_numeric", "column": "driver_rating"},
            {"operation": "cast_numeric", "column": "delivery_distance_km"},
            {"operation": "cast_numeric", "column": "late_deliveries_last_30d"},
            {"operation": "impute_numeric", "column": "driver_rating"},
            {"operation": "impute_numeric", "column": "late_deliveries_last_30d"},
            {"operation": "impute_numeric", "column": "delivery_distance_km"},
            {"operation": "clip_outliers_iqr", "column": "delivery_distance_km"},
            {"operation": "normalize_categories", "column": "city"},
            {"operation": "normalize_categories", "column": "vehicle_type"},
            {"operation": "impute_categorical", "column": "city"},
            {"operation": "impute_categorical", "column": "vehicle_type"},
            {"operation": "drop_duplicates"},
        ],
        notes=["City aliases should be standardized before downstream feature engineering."],
    ),
}


class FSDSCleaningEnvironment(MCPEnvironment):
    def __init__(self) -> None:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self._episode: Optional[EpisodeData] = None

        mcp = FastMCP("fsds_cleaning_env")

        @mcp.tool
        def list_tasks() -> dict[str, Any]:
            return {
                "tasks": [
                    {
                        "task_id": task.task_id,
                        "title": task.title,
                        "objective": task.objective,
                        "task_type": task.task_type,
                    }
                    for task in TASKS.values()
                ]
            }

        @mcp.tool
        def get_task_brief() -> dict[str, Any]:
            episode = self._require_episode()
            spec = episode.spec
            return {
                "task_id": spec.task_id,
                "title": spec.title,
                "objective": spec.objective,
                "business_context": spec.business_context,
                "target_column": spec.target_column,
                "success_criteria": {
                    "quality_gate": "pass",
                    "minimum_row_retention": spec.min_retention_ratio,
                    "no_missing_values_outside_allowed": True,
                    "target_preserved": True,
                },
                "notes": spec.notes,
            }

        @mcp.tool
        def preview_data(n: int = 5) -> dict[str, Any]:
            episode = self._require_episode()
            return {
                "rows": episode.working_df.head(n).to_dict(orient="records"),
                "shape": list(episode.working_df.shape),
            }

        @mcp.tool
        def profile_data() -> dict[str, Any]:
            episode = self._require_episode()
            df = episode.working_df
            numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
            numeric_summary = {}
            if numeric_cols:
                desc = df[numeric_cols].describe(include="all").fillna("NaN")
                numeric_summary = desc.to_dict()
            return {
                "columns": list(df.columns),
                "dtypes": {c: str(t) for c, t in df.dtypes.items()},
                "missing_counts": df.isna().sum().to_dict(),
                "duplicate_rows": int(df.duplicated().sum()),
                "invalid_token_counts": {c: int(df[c].astype(str).isin(INVALID_TOKENS).sum()) for c in df.columns},
                "numeric_summary": numeric_summary,
            }

        @mcp.tool
        def get_operation_history() -> dict[str, Any]:
            episode = self._require_episode()
            return {"operations": episode.operation_log, "total_reward": round(episode.total_reward, 4)}

        @mcp.tool
        def render_episode(n_preview_rows: int = 5) -> dict[str, Any]:
            """Human-friendly snapshot of the current episode state."""
            episode = self._require_episode()
            df = episode.working_df
            return {
                "task_id": episode.spec.task_id,
                "title": episode.spec.title,
                "step_count": self._state.step_count,
                "max_steps": episode.max_steps,
                "total_reward": round(episode.total_reward, 4),
                "shape": list(df.shape),
                "last_gate_report": episode.last_gate_report or None,
                "operations": episode.operation_log,
                "preview": df.head(n_preview_rows).to_dict(orient="records"),
            }

        @mcp.tool
        def apply_cleaning_operation(operation: str, column: Optional[str] = None, strategy: str = "median") -> dict[str, Any]:
            episode = self._require_episode()
            if episode.submitted:
                return self._tool_error("Episode already submitted.")
            if self._state.step_count >= episode.max_steps:
                return self._tool_error("Maximum steps reached. Submit or reset the episode.")

            before_score = self._quality_score(episode)
            before_df = episode.working_df.copy(deep=True)
            message = self._apply_operation(episode, operation=operation, column=column, strategy=strategy)
            after_score = self._quality_score(episode)
            delta = round(after_score - before_score, 4)
            step_reward = compute_step_reward(
                StepRewardInput(quality_before=before_score, quality_after=after_score)
            )
            episode.total_reward += step_reward
            episode.operation_log.append(
                {
                    "operation": operation,
                    "column": column,
                    "strategy": strategy,
                    "message": message,
                    "quality_delta": delta,
                    "reward": round(step_reward, 4),
                    "shape_before": list(before_df.shape),
                    "shape_after": list(episode.working_df.shape),
                }
            )
            return {
                "message": message,
                "reward": round(float(step_reward), 4),
                "quality_score": round(float(after_score), 4),
                "shape": [int(x) for x in episode.working_df.shape],
            }

        @mcp.tool
        def run_quality_gates() -> dict[str, Any]:
            episode = self._require_episode()
            report = self._evaluate_quality_gates(episode)
            episode.last_gate_report = report
            bonus = compute_quality_gate_bonus(bool(report["passed"]))
            episode.total_reward += bonus
            return {
                **report,
                "reward": bonus,
                "total_reward": round(episode.total_reward, 4),
            }

        @mcp.tool
        def submit_solution() -> dict[str, Any]:
            import json as _json
            episode = self._require_episode()
            report = self._evaluate_quality_gates(episode)
            match_score = self._required_operations_score(episode)
            final_reward = round(
                compute_final_reward(
                    FinalRewardInput(
                        quality_score=self._quality_score(episode),
                        gate_passed=bool(report["passed"]),
                        required_operation_coverage=match_score,
                    )
                ),
                4,
            )
            episode.total_reward += final_reward
            episode.submitted = True
            # Use pandas to_json to safely serialize numpy/pandas types (Timestamps, int64, float64).
            cleaned_preview = _json.loads(
                episode.working_df.head(5).to_json(orient="records", date_format="iso")
            )
            return {
                "done": True,
                "passed": bool(report["passed"]),
                "final_reward": final_reward,
                "cumulative_reward": round(float(episode.total_reward), 4),
                "quality_report": report,
                "required_operation_coverage": round(float(match_score), 4),
                "cleaned_preview": cleaned_preview,
            }

        super().__init__(mcp)

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        task_id: Optional[str] = None,
        **kwargs: Any,
    ) -> Observation:
        spec = TASKS[task_id] if task_id in TASKS else TASKS["ecommerce_mobile"]
        dataset_mode = kwargs.get("dataset_mode")
        dataset_n_rows = kwargs.get("dataset_n_rows")
        try:
            raw_df = spec.dataset_factory(
                seed=seed,
                dataset_mode=dataset_mode,
                n_rows_override=dataset_n_rows,
            )
        except TypeError:
            raw_df = spec.dataset_factory()
        episode = EpisodeData(spec=spec, raw_df=raw_df.copy(deep=True), working_df=raw_df.copy(deep=True))
        self._episode = episode
        self._state = State(episode_id=episode_id or str(uuid4()), step_count=0)
        return Observation(
            done=False,
            reward=0.0,
            metadata={
                "status": "ready",
                "schema_version": OBSERVATION_SCHEMA_VERSION,
                "task_id": spec.task_id,
                "task_type": spec.task_type,
                "target_column": spec.target_column,
                "episode_id": self._state.episode_id,
                "step_count": self._state.step_count,
                "max_steps": episode.max_steps,
                "message": f"FSDS cleaning environment ready for task '{spec.task_id}'.",
                "available_tools": [
                    "list_tasks",
                    "get_task_brief",
                    "preview_data",
                    "profile_data",
                    "get_operation_history",
                    "apply_cleaning_operation",
                    "run_quality_gates",
                    "submit_solution",
                    "render_episode",
                ],
                "available_operations": AVAILABLE_OPERATIONS,
            },
        )

    def _step_impl(self, action: Action, timeout_s: Optional[float] = None, **kwargs: Any) -> Observation:
        del timeout_s, kwargs
        return Observation(
            done=False,
            reward=0.0,
            metadata={
                "error": f"Unknown action type: {type(action).__name__}. Use MCP tools instead."
            },
        )

    def step(self, action: Action, timeout_s: Optional[float] = None, **kwargs: Any) -> Observation:
        self._state.step_count += 1
        return super().step(action, timeout_s=timeout_s, **kwargs)

    @property
    def state(self) -> State:
        return self._state

    def _require_episode(self) -> EpisodeData:
        if self._episode is None:
            raise RuntimeError("Reset the environment before using tools.")
        return self._episode

    def _tool_error(self, message: str) -> dict[str, Any]:
        return {"error": message, "reward": TOOL_ERROR_REWARD}

    def _apply_operation(self, episode: EpisodeData, operation: str, column: Optional[str], strategy: str) -> str:
        df = episode.working_df
        if operation == "drop_duplicates":
            before = len(df)
            episode.working_df = df.drop_duplicates().reset_index(drop=True)
            return f"Dropped {before - len(episode.working_df)} duplicate rows."

        if column is None:
            raise ValueError(f"Operation '{operation}' requires a column.")
        if column not in df.columns:
            raise ValueError(f"Unknown column '{column}'.")

        if operation == "replace_invalid_with_null":
            mask = df[column].astype(str).isin(INVALID_TOKENS)
            changed = int(mask.sum())
            episode.working_df.loc[mask, column] = np.nan
            return f"Replaced {changed} invalid tokens in '{column}' with nulls."

        if operation == "cast_numeric":
            episode.working_df[column] = pd.to_numeric(df[column], errors="coerce")
            return f"Cast '{column}' to numeric with coercion."

        if operation == "cast_datetime":
            episode.working_df[column] = pd.to_datetime(df[column], errors="coerce")
            return f"Cast '{column}' to datetime with coercion."

        if operation == "impute_numeric":
            series = pd.to_numeric(episode.working_df[column], errors="coerce")
            if strategy == "mean":
                fill_value = float(series.mean())
            else:
                fill_value = float(series.median())
            episode.working_df[column] = series.fillna(fill_value)
            return f"Imputed '{column}' using {strategy}={round(fill_value, 4)}."

        if operation == "impute_categorical":
            mode = episode.working_df[column].dropna().astype(str).mode()
            fill_value = mode.iloc[0] if not mode.empty else "unknown"
            episode.working_df[column] = episode.working_df[column].fillna(fill_value)
            return f"Imputed '{column}' with mode='{fill_value}'."

        if operation == "normalize_categories":
            null_mask = episode.working_df[column].isna()
            normalized = (
                episode.working_df[column]
                .astype(str)
                .str.strip()
                .str.lower()
                .replace({"ios": "ios", "android": "android", "android ": "android", "mty": "monterrey", "car": "car", "CAR": "car"})
            )
            normalized = normalized.replace({
                "ca": "ca",
                "mx": "mx",
                "us": "us",
                "monterrey": "monterrey",
                "cdmx": "cdmx",
                "gdl": "gdl",
                "motorbike": "motorbike",
                "bike": "bike",
            })
            normalized[null_mask] = np.nan
            episode.working_df[column] = normalized
            return f"Normalized categories in '{column}'."

        if operation == "clip_outliers_iqr":
            series = pd.to_numeric(episode.working_df[column], errors="coerce")
            q1 = series.quantile(0.25)
            q3 = series.quantile(0.75)
            iqr = q3 - q1
            lower = q1 - 1.5 * iqr
            upper = q3 + 1.5 * iqr
            clipped = series.clip(lower=lower, upper=upper)
            changed = int((series.fillna(clipped) != clipped).sum())
            episode.working_df[column] = clipped
            return f"Clipped {changed} outliers in '{column}' using IQR bounds [{round(lower, 4)}, {round(upper, 4)}]."

        raise ValueError(f"Unsupported operation '{operation}'.")

    def _evaluate_quality_gates(self, episode: EpisodeData) -> dict[str, Any]:
        df = episode.working_df
        raw_df = episode.raw_df
        spec = episode.spec
        missing_counts = df.isna().sum().to_dict()
        non_target_missing = {
            col: int(count)
            for col, count in missing_counts.items()
            if col != spec.target_column and count > 0
        }
        duplicate_rows = int(df.duplicated().sum())
        retention_ratio = len(df) / max(len(raw_df), 1)
        schema_same_columns = list(df.columns) == list(raw_df.columns)
        dtype_alignment = {
            col: str(df[col].dtype) == expected for col, expected in spec.expected_types.items() if col in df.columns
        }
        target_preserved = spec.target_column in df.columns
        stability = self._stability_probe(df, spec)
        tests = {
            "test_missing_values": {"passed": len(non_target_missing) == 0, "details": non_target_missing},
            "test_duplicated_rows": {"passed": duplicate_rows == 0, "details": duplicate_rows},
            "test_data_consistency": {"passed": schema_same_columns and target_preserved, "details": {"schema_same_columns": schema_same_columns, "target_preserved": target_preserved}},
            "test_data_retention": {"passed": retention_ratio >= spec.min_retention_ratio, "details": round(retention_ratio, 4)},
            "test_dtype_alignment": {"passed": all(dtype_alignment.values()), "details": dtype_alignment},
            "test_stability_probe": {"passed": stability["score_std"] <= 0.15, "details": stability},
        }
        passed = all(test["passed"] for test in tests.values())
        return {
            "passed": bool(passed),
            "tests": tests,
            "shape": [int(x) for x in df.shape],
            "retention_ratio": round(float(retention_ratio), 4),
        }

    def _quality_score(self, episode: EpisodeData) -> float:
        df = episode.working_df
        raw_df = episode.raw_df
        spec = episode.spec
        retention = len(df) / max(len(raw_df), 1)
        missing_penalty = df.drop(columns=[spec.target_column], errors="ignore").isna().sum().sum()
        duplicate_penalty = int(df.duplicated().sum())
        dtype_score = 0.0
        for col, expected in spec.expected_types.items():
            if col in df.columns and str(df[col].dtype) == expected:
                dtype_score += 1.0
        dtype_score /= max(len(spec.expected_types), 1)
        score = 1.2 * retention + 0.7 * dtype_score - 0.08 * missing_penalty - 0.2 * duplicate_penalty
        return max(0.0, min(2.0, float(score)))

    def _required_operations_score(self, episode: EpisodeData) -> float:
        executed = [
            {k: v for k, v in op.items() if k in {"operation", "column"} and v is not None}
            for op in episode.operation_log
        ]
        matched = 0
        for required in episode.spec.required_ops:
            req = {k: v for k, v in required.items() if k in {"operation", "column"}}
            if req in executed:
                matched += 1
        return matched / max(len(episode.spec.required_ops), 1)

    def _stability_probe(self, df: pd.DataFrame, spec: TaskSpec) -> dict[str, Any]:
        if spec.target_column not in df.columns or len(df) < 8:
            return {"score_mean": 0.0, "score_std": 1.0, "n_folds": 0}
        clean_df = df.copy()
        X = clean_df.drop(columns=[spec.target_column])
        y = clean_df[spec.target_column]
        X_enc = self._encode_features(X)
        if X_enc.empty:
            return {"score_mean": 0.0, "score_std": 1.0, "n_folds": 0}
        kf = KFold(n_splits=min(3, len(clean_df)), shuffle=True, random_state=42)
        scores: list[float] = []
        for train_idx, test_idx in kf.split(X_enc):
            X_train, X_test = X_enc.iloc[train_idx], X_enc.iloc[test_idx]
            y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]
            try:
                if spec.task_type == "classification":
                    model = RandomForestClassifier(n_estimators=25, random_state=42)
                    model.fit(X_train, y_train)
                    scores.append(float(model.score(X_test, y_test)))
                else:
                    model = RandomForestRegressor(n_estimators=25, random_state=42)
                    model.fit(X_train, y_train)
                    pred = model.predict(X_test)
                    mae = np.mean(np.abs(pred - y_test))
                    denom = max(float(np.mean(np.abs(y_test))), 1e-6)
                    scores.append(float(max(0.0, 1.0 - (mae / denom))))
            except Exception:
                scores.append(0.0)
        return {
            "score_mean": round(float(np.mean(scores)), 4) if scores else 0.0,
            "score_std": round(float(np.std(scores)), 4) if scores else 1.0,
            "n_folds": len(scores),
        }

    def _encode_features(self, X: pd.DataFrame) -> pd.DataFrame:
        out = pd.DataFrame(index=X.index)
        for col in X.columns:
            if pd.api.types.is_numeric_dtype(X[col]):
                out[col] = pd.to_numeric(X[col], errors="coerce").fillna(0.0)
            elif pd.api.types.is_datetime64_any_dtype(X[col]):
                out[col] = pd.to_datetime(X[col], errors="coerce").map(lambda x: x.toordinal() if not pd.isna(x) else 0)
            else:
                le = LabelEncoder()
                out[col] = le.fit_transform(X[col].astype(str).fillna("missing"))
        return out