#!/usr/bin/env python3
"""Bab 13 AI-first capstone playground.

A small, offline, standard-library-only capstone demo:
- generate synthetic UMKM promo data
- compare rule baseline vs learned linear model
- produce recommendations, error analysis, and risk register
"""

from __future__ import annotations

import argparse
import json
import math
import os
import random
from dataclasses import dataclass, asdict
from typing import Dict, List, Sequence, Tuple

PROMOS = ["tanpa_promo", "diskon_10", "bundling", "diskon_20"]
WEATHERS = ["cerah", "hujan"]
DAY_TYPES = ["weekday", "weekend"]


@dataclass
class PromoRow:
    day_type: str
    weather: str
    stock_level: int
    promo: str
    profit: float
    complaints: int
    stockout: int


def true_profit(day_type: str, weather: str, stock_level: int, promo: str, rng: random.Random) -> Tuple[float, int, int]:
    """Synthetic business process for educational use only."""
    base = 90.0
    if day_type == "weekend":
        base += 22
    if weather == "hujan":
        base -= 8
    base += min(stock_level, 80) * 0.35

    promo_effect = {
        "tanpa_promo": 0,
        "diskon_10": 18 if weather == "hujan" else 9,
        "bundling": 24 if day_type == "weekend" else 12,
        "diskon_20": 28,
    }[promo]
    margin_penalty = {
        "tanpa_promo": 0,
        "diskon_10": 7,
        "bundling": 10,
        "diskon_20": 24,
    }[promo]
    demand = base + promo_effect - margin_penalty + rng.gauss(0, 6)
    stockout = int(demand > stock_level + 75)
    complaints = int(promo == "diskon_20" and stockout) + int(rng.random() < 0.04)
    profit = demand - 18 * stockout - 12 * complaints
    return round(profit, 3), complaints, stockout


def generate_dataset(seed: int = 13, n: int = 480) -> List[PromoRow]:
    rng = random.Random(seed)
    rows: List[PromoRow] = []
    for _ in range(n):
        day_type = rng.choice(DAY_TYPES)
        weather = rng.choice(WEATHERS)
        stock_level = rng.randint(25, 120)
        promo = rng.choice(PROMOS)
        profit, complaints, stockout = true_profit(day_type, weather, stock_level, promo, rng)
        rows.append(PromoRow(day_type, weather, stock_level, promo, profit, complaints, stockout))
    return rows


def features(row: PromoRow, promo_override: str | None = None) -> List[float]:
    promo = promo_override or row.promo
    return [
        1.0,
        1.0 if row.day_type == "weekend" else 0.0,
        1.0 if row.weather == "hujan" else 0.0,
        row.stock_level / 120.0,
        1.0 if promo == "diskon_10" else 0.0,
        1.0 if promo == "bundling" else 0.0,
        1.0 if promo == "diskon_20" else 0.0,
    ]


def dot(a: Sequence[float], b: Sequence[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def train_linear_model(rows: Sequence[PromoRow], epochs: int = 140, lr: float = 0.045) -> List[float]:
    """Tiny gradient descent linear regressor for profit prediction."""
    weights = [0.0] * 7
    for _ in range(epochs):
        grad = [0.0] * len(weights)
        for row in rows:
            x = features(row)
            pred = dot(weights, x)
            err = pred - row.profit
            for i, value in enumerate(x):
                grad[i] += err * value / len(rows)
        for i in range(len(weights)):
            weights[i] -= lr * grad[i]
    return weights


def baseline_policy(day_type: str, weather: str, stock_level: int) -> str:
    if stock_level < 45:
        return "tanpa_promo"
    if day_type == "weekend":
        return "bundling"
    if weather == "hujan":
        return "diskon_10"
    return "tanpa_promo"


def recommend_with_model(weights: Sequence[float], context: PromoRow) -> str:
    scores = {promo: dot(weights, features(context, promo_override=promo)) for promo in PROMOS}
    # Guardrail: avoid aggressive discount when stock is low.
    if context.stock_level < 45:
        scores["diskon_20"] -= 30
    return max(scores, key=scores.get)


def evaluate_policy(rows: Sequence[PromoRow], policy_name: str, weights: Sequence[float] | None = None, seed: int = 99) -> Dict[str, float]:
    rng = random.Random(seed)
    total_profit = 0.0
    complaints = 0
    stockouts = 0
    for row in rows:
        if policy_name == "baseline":
            promo = baseline_policy(row.day_type, row.weather, row.stock_level)
        elif policy_name == "model":
            assert weights is not None
            promo = recommend_with_model(weights, row)
        else:
            raise ValueError(policy_name)
        profit, c, s = true_profit(row.day_type, row.weather, row.stock_level, promo, rng)
        total_profit += profit
        complaints += c
        stockouts += s
    n = len(rows)
    return {
        "average_profit": round(total_profit / n, 3),
        "complaint_rate": round(complaints / n, 3),
        "stockout_rate": round(stockouts / n, 3),
    }


def error_analysis(rows: Sequence[PromoRow], weights: Sequence[float], limit: int = 5) -> List[Dict[str, object]]:
    errors = []
    for row in rows:
        pred = dot(weights, features(row))
        errors.append((abs(pred - row.profit), row, pred))
    errors.sort(reverse=True, key=lambda x: x[0])
    out = []
    for err, row, pred in errors[:limit]:
        out.append({
            "absolute_error": round(err, 3),
            "actual_profit": row.profit,
            "predicted_profit": round(pred, 3),
            "context": asdict(row),
            "possible_reason": "noise, stockout, or promo interaction not fully captured by linear model",
        })
    return out


def risk_register() -> List[Dict[str, object]]:
    return [
        {"risk": "rekomendasi diskon berlebihan", "impact": 4, "probability": 3, "mitigation": "batas diskon dan approval manual"},
        {"risk": "data sintetis tidak mewakili toko nyata", "impact": 4, "probability": 4, "mitigation": "validasi dengan data riil sebelum klaim"},
        {"risk": "stockout karena promo agresif", "impact": 5, "probability": 2, "mitigation": "guardrail stok minimum"},
        {"risk": "overclaim hasil capstone", "impact": 3, "probability": 4, "mitigation": "tulis batasan pada model card"},
        {"risk": "bias terhadap pola weekend tertentu", "impact": 3, "probability": 3, "mitigation": "audit per segmen dan periode"},
    ]


def run_demo(seed: int = 13) -> Dict[str, object]:
    rows = generate_dataset(seed=seed, n=520)
    split = int(0.75 * len(rows))
    train_rows, test_rows = rows[:split], rows[split:]
    weights = train_linear_model(train_rows)
    baseline = evaluate_policy(test_rows, "baseline", seed=seed + 1)
    model = evaluate_policy(test_rows, "model", weights=weights, seed=seed + 1)
    improvement = (model["average_profit"] - baseline["average_profit"]) / abs(baseline["average_profit"])

    sample_contexts = [
        PromoRow("weekday", "cerah", 80, "tanpa_promo", 0, 0, 0),
        PromoRow("weekday", "hujan", 70, "tanpa_promo", 0, 0, 0),
        PromoRow("weekend", "cerah", 95, "tanpa_promo", 0, 0, 0),
        PromoRow("weekend", "hujan", 35, "tanpa_promo", 0, 0, 0),
    ]
    recommendations = [
        {
            "context": {"day_type": r.day_type, "weather": r.weather, "stock_level": r.stock_level},
            "baseline_promo": baseline_policy(r.day_type, r.weather, r.stock_level),
            "model_promo": recommend_with_model(weights, r),
        }
        for r in sample_contexts
    ]

    return {
        "metadata": {"chapter": "Bab 13 — Capstone AI-First", "seed": seed, "dataset": "synthetic_umkm_promo"},
        "data_card_summary": {
            "rows": len(rows),
            "features": ["day_type", "weather", "stock_level", "promo"],
            "target": "profit",
            "license": "synthetic educational data",
            "warning": "not representative of a real store without validation",
        },
        "model_card_summary": {
            "model": "linear regression trained from scratch",
            "baseline": "rule-based promo policy",
            "intended_use": "educational capstone demo",
            "not_intended_use": "production pricing or automated discounting",
        },
        "evaluation": {"baseline": baseline, "model": model, "relative_profit_improvement": round(improvement, 3)},
        "learned_weights": [round(w, 3) for w in weights],
        "recommendations": recommendations,
        "error_analysis_top_cases": error_analysis(test_rows, weights),
        "risk_register": risk_register(),
    }


def save_results(results: Dict[str, object], output_path: str) -> None:
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
        f.write("\n")


def self_test() -> None:
    rows = generate_dataset(seed=1, n=40)
    assert len(rows) == 40
    assert all(r.promo in PROMOS for r in rows)
    weights = train_linear_model(rows, epochs=20)
    assert len(weights) == 7
    demo = run_demo(seed=3)
    assert demo["evaluation"]["model"]["average_profit"] > 0
    assert "risk_register" in demo and len(demo["risk_register"]) >= 5
    assert demo["data_card_summary"]["license"] == "synthetic educational data"


def main(argv: Sequence[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Bab 13 capstone AI-first playground")
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--seed", type=int, default=13)
    parser.add_argument("--output", default=os.path.join("outputs", "bab13_capstone_results.json"))
    args = parser.parse_args(argv)
    if args.self_test:
        self_test()
        print("self-test passed")
        return
    results = run_demo(args.seed)
    save_results(results, args.output)
    print(json.dumps(results, ensure_ascii=False, indent=2))
    print(f"\nSaved results to {args.output}")


if __name__ == "__main__":
    main()