"""Small timeout-risk model for Huaxiang ChaPanda orders.

The first dataset only contains delayed orders, so the model supports two modes:

* labeled_empirical_bayes: used when both timeout and non-timeout samples exist.
* positive_only_time_density: used when only timeout samples exist. This is a
  relative risk estimate by order time, not a calibrated probability.
"""

from __future__ import annotations

import csv
import json
import math
from collections import Counter, defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Iterable

MODEL_VERSION = 1
CSV_FIELDS = [
    "product_name",
    "order_time",
    "timeout",
    "is_timeout",
    "expected_latest_time",
    "actual_arrival_time",
    "submitted_at",
    "notes",
]


@dataclass(frozen=True)
class TrainingRecord:
    product_name: str
    order_time: datetime
    is_timeout: bool
    source: str


def parse_datetime(value: Any) -> datetime | None:
    if value is None:
        return None
    text = str(value).strip()
    if not text:
        return None

    normalized = text.replace("/", "-").replace("T", " ")
    if len(normalized) == 16:
        normalized = f"{normalized}:00"

    formats = (
        "%Y-%m-%d %H:%M:%S",
        "%Y-%m-%d %H:%M",
        "%Y-%m-%d",
    )
    for fmt in formats:
        try:
            return datetime.strptime(normalized, fmt)
        except ValueError:
            pass

    try:
        return datetime.fromisoformat(normalized)
    except ValueError:
        return None


def normalize_row(row: dict[str, Any]) -> dict[str, str]:
    return {str(key).strip(): str(value).strip() for key, value in row.items() if key is not None}


def coerce_bool(value: Any) -> bool | None:
    if value is None:
        return None
    text = str(value).strip().lower()
    if not text:
        return None
    true_values = {"1", "true", "yes", "y", "timeout", "late", "超时", "是", "yes_timeout"}
    false_values = {"0", "false", "no", "n", "on_time", "not_timeout", "未超时", "否"}
    if text in true_values:
        return True
    if text in false_values:
        return False
    if "未超时" in text or "准时" in text:
        return False
    if "超时" in text or "延迟" in text:
        return True
    return None


def infer_timeout_label(row: dict[str, str], forced_timeout: bool | None = None) -> bool | None:
    if forced_timeout is not None:
        return forced_timeout

    for key in ("is_timeout", "timeout", "timeout_status", "status"):
        parsed = coerce_bool(row.get(key))
        if parsed is not None:
            return parsed

    expected = parse_datetime(row.get("expected_latest_time"))
    actual = parse_datetime(row.get("actual_arrival_time"))
    if expected and actual:
        return actual > expected
    return None


def read_csv_records(path: Path, forced_timeout: bool | None = None) -> list[TrainingRecord]:
    if not path.exists():
        return []

    records: list[TrainingRecord] = []
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle)
        for raw_row in reader:
            row = normalize_row(raw_row)
            order_time = parse_datetime(row.get("order_time"))
            is_timeout = infer_timeout_label(row, forced_timeout=forced_timeout)
            if not order_time or is_timeout is None:
                continue

            product_name = row.get("product_name", "").strip() or "unknown"
            records.append(
                TrainingRecord(
                    product_name=product_name,
                    order_time=order_time,
                    is_timeout=is_timeout,
                    source=str(path),
                )
            )
    return records


def load_dataset_records(dataset_dir: str | Path) -> list[TrainingRecord]:
    dataset_path = Path(dataset_dir)
    records_path = dataset_path / "records.csv"
    if records_path.exists():
        records = read_csv_records(records_path, forced_timeout=None)
    else:
        records = read_csv_records(dataset_path / "timeout_records.csv", forced_timeout=True)

    tmp_dir = dataset_path / "tmp"
    if tmp_dir.exists():
        for tmp_file in sorted(tmp_dir.glob("tmp_records_*.csv")):
            records.extend(read_csv_records(tmp_file, forced_timeout=None))

    return records


def _time_bucket(dt: datetime) -> str:
    hour = dt.hour
    if 6 <= hour < 11:
        return "morning"
    if 11 <= hour < 14:
        return "lunch"
    if 14 <= hour < 17:
        return "afternoon"
    if 17 <= hour < 20:
        return "dinner"
    if 20 <= hour < 24:
        return "night"
    return "late_night"


def _sigmoid(value: float) -> float:
    return 1.0 / (1.0 + math.exp(-value))


def _logit(probability: float) -> float:
    clipped = min(max(probability, 0.02), 0.98)
    return math.log(clipped / (1.0 - clipped))


def _clip_probability(value: float) -> float:
    return round(min(max(value, 0.02), 0.98), 4)


def _empirical_probabilities(
    records: list[TrainingRecord],
    key_fn: Callable[[TrainingRecord], str],
    base_probability: float,
    prior_strength: float,
) -> dict[str, dict[str, float | int]]:
    grouped: dict[str, list[TrainingRecord]] = defaultdict(list)
    for record in records:
        grouped[key_fn(record)].append(record)

    stats: dict[str, dict[str, float | int]] = {}
    for key, items in grouped.items():
        total = len(items)
        timeouts = sum(1 for item in items if item.is_timeout)
        probability = (timeouts + prior_strength * base_probability) / (total + prior_strength)
        stats[str(key)] = {
            "total": total,
            "timeouts": timeouts,
            "probability": _clip_probability(probability),
        }
    return stats


def _positive_density_probabilities(
    records: list[TrainingRecord],
    key_fn: Callable[[TrainingRecord], str],
    keys: Iterable[str],
    low: float,
    high: float,
) -> dict[str, dict[str, float | int]]:
    counts = Counter(key_fn(record) for record in records if record.is_timeout)
    max_count = max(counts.values(), default=0)
    output: dict[str, dict[str, float | int]] = {}
    for key in keys:
        count = counts.get(key, 0)
        density = count / max_count if max_count else 0.0
        probability = low + (high - low) * density
        output[str(key)] = {
            "positive_count": count,
            "probability": _clip_probability(probability),
        }
    return output


def _smoothed_hour_density(records: list[TrainingRecord]) -> dict[str, dict[str, float | int]]:
    raw_counts = Counter(record.order_time.hour for record in records if record.is_timeout)
    smoothed: dict[int, float] = {}
    for hour in range(24):
        smoothed[hour] = (
            raw_counts.get(hour, 0)
            + 0.45 * raw_counts.get((hour - 1) % 24, 0)
            + 0.45 * raw_counts.get((hour + 1) % 24, 0)
            + 0.15 * raw_counts.get((hour - 2) % 24, 0)
            + 0.15 * raw_counts.get((hour + 2) % 24, 0)
        )
    max_score = max(smoothed.values(), default=0.0)

    output: dict[str, dict[str, float | int]] = {}
    for hour in range(24):
        density = smoothed[hour] / max_score if max_score else 0.0
        probability = 0.12 + 0.68 * density
        output[str(hour)] = {
            "positive_count": raw_counts.get(hour, 0),
            "smoothed_density": round(density, 4),
            "probability": _clip_probability(probability),
        }
    return output


def train_timeout_model(records: list[TrainingRecord]) -> dict[str, Any]:
    total = len(records)
    positives = sum(1 for record in records if record.is_timeout)
    negatives = total - positives
    generated_at = datetime.now(timezone.utc).isoformat()

    if total == 0:
        return {
            "model_version": MODEL_VERSION,
            "mode": "empty_default",
            "generated_at": generated_at,
            "base_probability": 0.35,
            "hour_probabilities": {str(hour): {"probability": 0.35, "total": 0, "timeouts": 0} for hour in range(24)},
            "weekday_probabilities": {str(day): {"probability": 0.35, "total": 0, "timeouts": 0} for day in range(7)},
            "time_bucket_probabilities": {},
            "product_probabilities": {},
            "metadata": {
                "training_records": 0,
                "timeout_records": 0,
                "non_timeout_records": 0,
                "calibration": "No valid records were found. Default probability is a placeholder.",
            },
        }

    if positives > 0 and negatives > 0:
        base_probability = (positives + 2.0) / (total + 4.0)
        mode = "labeled_empirical_bayes"
        calibration = "Calibrated from user-labeled timeout and non-timeout records with Bayesian smoothing."
        hour_probabilities = _empirical_probabilities(records, lambda r: str(r.order_time.hour), base_probability, 6.0)
        weekday_probabilities = _empirical_probabilities(records, lambda r: str(r.order_time.weekday()), base_probability, 8.0)
        time_bucket_probabilities = _empirical_probabilities(records, lambda r: _time_bucket(r.order_time), base_probability, 5.0)
        product_probabilities = _empirical_probabilities(records, lambda r: r.product_name, base_probability, 8.0)
    else:
        base_probability = 0.35
        mode = "positive_only_time_density"
        calibration = (
            "Only timeout records were available. Output is a relative time-risk estimate, "
            "not a fully calibrated timeout probability."
        )
        positive_records = [record for record in records if record.is_timeout]
        hour_probabilities = _smoothed_hour_density(positive_records)
        weekday_probabilities = _positive_density_probabilities(
            positive_records,
            lambda r: str(r.order_time.weekday()),
            [str(day) for day in range(7)],
            low=0.18,
            high=0.62,
        )
        time_bucket_probabilities = _positive_density_probabilities(
            positive_records,
            lambda r: _time_bucket(r.order_time),
            ["morning", "lunch", "afternoon", "dinner", "night", "late_night"],
            low=0.16,
            high=0.70,
        )
        product_probabilities = {}

    return {
        "model_version": MODEL_VERSION,
        "mode": mode,
        "generated_at": generated_at,
        "base_probability": _clip_probability(base_probability),
        "hour_probabilities": hour_probabilities,
        "weekday_probabilities": weekday_probabilities,
        "time_bucket_probabilities": time_bucket_probabilities,
        "product_probabilities": product_probabilities,
        "metadata": {
            "training_records": total,
            "timeout_records": positives,
            "non_timeout_records": negatives,
            "calibration": calibration,
            "features": ["order hour", "weekday", "time bucket", "optional product name"],
        },
    }


def load_model(path: str | Path) -> dict[str, Any]:
    with Path(path).open("r", encoding="utf-8") as handle:
        return json.load(handle)


def save_model(model: dict[str, Any], path: str | Path) -> None:
    output_path = Path(path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", encoding="utf-8") as handle:
        json.dump(model, handle, ensure_ascii=False, indent=2, sort_keys=True)
        handle.write("\n")


def predict_timeout_probability(
    model: dict[str, Any],
    order_time: str | datetime,
    product_name: str = "",
) -> dict[str, Any]:
    dt = order_time if isinstance(order_time, datetime) else parse_datetime(order_time)
    if dt is None:
        raise ValueError("order_time must look like '2026-05-24 18:30'.")

    base = float(model.get("base_probability", 0.35))
    mode = model.get("mode", "unknown")

    def component_probability(section: str, key: str, fallback: float) -> float:
        value = model.get(section, {}).get(str(key), {})
        try:
            return float(value.get("probability", fallback))
        except (TypeError, ValueError):
            return fallback

    hour_probability = component_probability("hour_probabilities", str(dt.hour), base)
    weekday_probability = component_probability("weekday_probabilities", str(dt.weekday()), base)
    bucket_probability = component_probability("time_bucket_probabilities", _time_bucket(dt), base)
    product_probability = None
    if product_name:
        product_probability = component_probability("product_probabilities", product_name.strip(), base)

    if mode == "labeled_empirical_bayes":
        combined_logit = _logit(base)
        combined_logit += 0.58 * (_logit(hour_probability) - _logit(base))
        combined_logit += 0.18 * (_logit(weekday_probability) - _logit(base))
        combined_logit += 0.16 * (_logit(bucket_probability) - _logit(base))
        if product_probability is not None:
            combined_logit += 0.08 * (_logit(product_probability) - _logit(base))
        probability = _sigmoid(combined_logit)
    else:
        combined_logit = _logit(base)
        combined_logit += 0.70 * (_logit(hour_probability) - _logit(base))
        combined_logit += 0.18 * (_logit(bucket_probability) - _logit(base))
        combined_logit += 0.12 * (_logit(weekday_probability) - _logit(base))
        probability = _sigmoid(combined_logit)

    clipped = _clip_probability(probability)
    if clipped >= 0.65:
        level = "high"
    elif clipped >= 0.4:
        level = "medium"
    else:
        level = "low"

    return {
        "probability": clipped,
        "risk_level": level,
        "mode": mode,
        "order_time": dt.strftime("%Y-%m-%d %H:%M"),
        "components": {
            "base": base,
            "hour": hour_probability,
            "weekday": weekday_probability,
            "time_bucket": bucket_probability,
            "product": product_probability,
        },
        "calibration": model.get("metadata", {}).get("calibration", ""),
    }