server.py

from __future__ import annotations

import tempfile
import uuid
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from sklearn.metrics import confusion_matrix

from hyperparameters import HYPERPARAMETER_SCHEMAS
from ml_core import (
    ALGORITHMS,
    DataFormatError,
    dataset_summary,
    export_model,
    infer_task,
    parse_feature_text,
    parse_training_text,
    predict_with_confidence,
    train_model,
)


app = FastAPI(title="AnyToAny ML API", version="2.0")

DATASETS: dict[str, dict[str, Any]] = {}
MODELS: dict[str, dict[str, Any]] = {}
PREDICTIONS: dict[str, Path] = {}


class TrainRequest(BaseModel):
    dataset_id: str
    task: str = "auto"
    algorithm: str
    test_size: float = Field(default=0.2, ge=0.1, le=0.4)
    random_state: int = 42
    standardize: bool = True
    hyperparameters: dict[str, Any]


class PredictRequest(BaseModel):
    model_id: str
    text: str


def _decode(contents: bytes) -> str:
    for encoding in ("utf-8-sig", "gb18030"):
        try:
            return contents.decode(encoding)
        except UnicodeDecodeError:
            continue
    raise DataFormatError("文件编码需要是 UTF-8 或 GB18030")


def _json_value(value: Any) -> Any:
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, np.ndarray):
        return value.tolist()
    if isinstance(value, tuple):
        return list(value)
    return value


def _records(frame: pd.DataFrame) -> list[dict[str, Any]]:
    return [
        {str(key): _json_value(value) for key, value in row.items()}
        for row in frame.to_dict(orient="records")
    ]


def _finite_float(value: float) -> float:
    number = float(value)
    return number if np.isfinite(number) else 0.0


def _dataset_profile(features: np.ndarray, targets: np.ndarray) -> dict[str, Any]:
    x = np.asarray(features, dtype=float)
    y = np.asarray(targets, dtype=float)
    task = infer_task(y)
    feature_items = []
    for index in range(x.shape[1]):
        column = x[:, index]
        feature_items.append(
            {
                "feature": f"x{index + 1}",
                "mean": _finite_float(np.mean(column)),
                "std": _finite_float(np.std(column)),
                "min": _finite_float(np.min(column)),
                "max": _finite_float(np.max(column)),
                "missing": int(np.isnan(column).sum()),
            }
        )

    first_target = y[:, 0]
    if task == "classification":
        values, counts = np.unique(np.round(first_target).astype(int), return_counts=True)
        target = {
            "kind": "classes",
            "items": [
                {"label": str(_json_value(value)), "count": int(count)}
                for value, count in zip(values, counts)
            ],
        }
        balance = float(np.min(counts) / np.max(counts)) if counts.size else 0.0
    else:
        counts, edges = np.histogram(first_target, bins=min(10, max(4, int(np.sqrt(len(first_target))))))
        target = {
            "kind": "bins",
            "items": [
                {
                    "label": f"{edges[index]:.2f}–{edges[index + 1]:.2f}",
                    "count": int(count),
                }
                for index, count in enumerate(counts)
            ],
        }
        balance = 1.0

    rows_per_feature = round(float(x.shape[0] / max(x.shape[1], 1)), 2)
    quality = {
        "rows_per_feature": rows_per_feature,
        "target_balance": round(balance, 3),
        "missing_cells": int(np.isnan(x).sum() + np.isnan(y).sum()),
        "recommendation": (
            "样本较少，优先使用树模型或较强正则化"
            if rows_per_feature < 10
            else "数据规模适合快速比较多种模型"
        ),
    }
    return {"features": feature_items, "target": target, "quality": quality}


def _primary_metric(metrics: dict[str, float], task: str) -> dict[str, Any]:
    preferred = "Accuracy" if task == "classification" else "R²"
    name = preferred if preferred in metrics else next(iter(metrics))
    return {"name": name, "value": float(metrics[name])}


def _train_diagnostics(
    y_true: np.ndarray,
    y_pred: np.ndarray,
    metrics: dict[str, float],
    task: str,
) -> dict[str, Any]:
    primary = _primary_metric(metrics, task)
    true_1d = np.asarray(y_true).reshape(len(y_true), -1)[:, 0]
    pred_1d = np.asarray(y_pred).reshape(len(y_pred), -1)[:, 0]

    if task == "classification":
        labels = sorted({*_json_value(np.unique(true_1d)), *_json_value(np.unique(pred_1d))})
        matrix = confusion_matrix(true_1d, pred_1d, labels=labels).tolist()
        accuracy = float(metrics.get("Accuracy", 0.0))
        insight = (
            "分类效果稳定，可进入预测验证"
            if accuracy >= 0.85
            else "分类边界仍有混淆，建议调整模型或扩大数据量"
        )
        return {
            "primary_metric": primary,
            "confusion_matrix": {"labels": [str(_json_value(item)) for item in labels], "matrix": matrix},
            "residual_histogram": None,
            "insight": insight,
        }

    residuals = true_1d - pred_1d
    counts, edges = np.histogram(residuals, bins=12)
    r2 = float(metrics.get("R²", 0.0))
    insight = (
        "残差集中，当前模型拟合质量较好"
        if r2 >= 0.75
        else "残差离散，建议检查特征尺度或尝试非线性模型"
    )
    return {
        "primary_metric": primary,
        "confusion_matrix": None,
        "residual_histogram": [
            {
                "bin": f"{edges[index]:.2f}–{edges[index + 1]:.2f}",
                "count": int(count),
            }
            for index, count in enumerate(counts)
        ],
        "insight": insight,
    }


@app.get("/api/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/api/algorithms")
def algorithms() -> dict[str, Any]:
    items = []
    for name, description in ALGORITHMS.items():
        items.append(
            {
                "name": name,
                "description": description,
                "supports": {
                    "classification": name != "Linear Regression",
                    "regression": name != "Logistic Regression",
                },
                "parameters": [
                    {
                        "name": parameter.name,
                        "default": _json_value(parameter.default),
                        "kind": parameter.kind,
                        "description": parameter.description,
                    }
                    for parameter in HYPERPARAMETER_SCHEMAS[name]
                ],
            }
        )
    return {"algorithms": items}


@app.post("/api/datasets")
async def upload_dataset(file: UploadFile = File(...)) -> dict[str, Any]:
    try:
        text = _decode(await file.read())
        features, targets = parse_training_text(text)
    except DataFormatError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc

    dataset_id = uuid.uuid4().hex
    summary = dataset_summary(features, targets)
    preview = pd.DataFrame(
        features[:5],
        columns=[f"x{index + 1}" for index in range(features.shape[1])],
    )
    for index in range(targets.shape[1]):
        preview[f"y{index + 1}"] = targets[:5, index]
    DATASETS[dataset_id] = {
        "features": features,
        "targets": targets,
        "filename": file.filename or "dataset.txt",
    }
    return {
        "dataset_id": dataset_id,
        "filename": file.filename or "dataset.txt",
        "summary": summary,
        "profile": _dataset_profile(features, targets),
        "preview": _records(preview),
        "columns": list(preview.columns),
    }


@app.post("/api/train")
def train(request: TrainRequest) -> dict[str, Any]:
    dataset = DATASETS.get(request.dataset_id)
    if dataset is None:
        raise HTTPException(status_code=404, detail="数据集会话已失效，请重新上传")
    if request.algorithm not in ALGORITHMS:
        raise HTTPException(status_code=422, detail="未知算法")

    task = infer_task(dataset["targets"]) if request.task == "auto" else request.task
    try:
        result = train_model(
            dataset["features"],
            dataset["targets"],
            task=task,
            algorithm=request.algorithm,
            test_size=request.test_size,
            random_state=request.random_state,
            standardize=request.standardize,
            hyperparameters=request.hyperparameters,
        )
    except (TypeError, ValueError) as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc

    model_id = uuid.uuid4().hex
    artifact_dir = Path(tempfile.mkdtemp(prefix="anytoany-api-model-"))
    model_path = Path(export_model(result.bundle, artifact_dir))
    MODELS[model_id] = {"bundle": result.bundle, "path": model_path}
    return {
        "model_id": model_id,
        "metrics": {name: float(value) for name, value in result.metrics.items()},
        "evaluation": {
            "y_true": result.y_true.tolist(),
            "y_pred": result.y_pred.tolist(),
            "comparison": _records(result.preview),
        },
        "feature_importance": _records(result.feature_importance),
        "diagnostics": _train_diagnostics(result.y_true, result.y_pred, result.metrics, task),
        "summary": {
            "algorithm": request.algorithm,
            "task": task,
            "test_size": request.test_size,
            "random_state": request.random_state,
            "standardize": request.standardize,
            "elapsed_seconds": result.elapsed_seconds,
            "feature_count": result.bundle["feature_count"],
            "target_count": result.bundle["target_count"],
            "hyperparameters": request.hyperparameters,
        },
        "downloads": {"model": f"/api/models/{model_id}/download"},
    }


@app.post("/api/predict")
def predict(request: PredictRequest) -> dict[str, Any]:
    model = MODELS.get(request.model_id)
    if model is None:
        raise HTTPException(status_code=404, detail="模型会话已失效，请重新训练")
    try:
        features = parse_feature_text(
            request.text, expected_features=model["bundle"]["feature_count"]
        )
        predictions, confidence = predict_with_confidence(model["bundle"], features)
    except (DataFormatError, ValueError) as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc

    frame = pd.DataFrame(
        predictions,
        columns=[f"prediction_{index + 1}" for index in range(predictions.shape[1])],
    )
    if confidence is not None:
        for index in range(confidence.shape[1]):
            frame[f"confidence_{index + 1}"] = confidence[:, index]
    prediction_id = uuid.uuid4().hex
    directory = Path(tempfile.mkdtemp(prefix="anytoany-api-prediction-"))
    path = directory / "predictions.csv"
    frame.to_csv(path, index=False)
    PREDICTIONS[prediction_id] = path
    return {
        "rows": _records(frame),
        "columns": list(frame.columns),
        "download": f"/api/predictions/{prediction_id}/download",
    }


@app.get("/api/models/{model_id}/download")
def download_model(model_id: str):
    model = MODELS.get(model_id)
    if model is None:
        raise HTTPException(status_code=404, detail="模型不存在")
    return FileResponse(
        model["path"],
        media_type="application/octet-stream",
        filename="anytoany_model.joblib",
    )


@app.get("/api/predictions/{prediction_id}/download")
def download_predictions(prediction_id: str):
    path = PREDICTIONS.get(prediction_id)
    if path is None:
        raise HTTPException(status_code=404, detail="预测结果不存在")
    return FileResponse(path, media_type="text/csv", filename="predictions.csv")


FRONTEND_DIST = Path(__file__).parent / "frontend" / "dist"
if FRONTEND_DIST.exists():
    app.mount("/", StaticFiles(directory=FRONTEND_DIST, html=True), name="frontend")