MLflow

import random
from dataclasses import dataclass, field
from typing import Dict, List, Any

# ─── MLflow Simulation ─────────────────────────────────────
@dataclass
class MLRun:
    run_id:     str
    experiment: str
    params:     Dict[str, Any]   = field(default_factory=dict)
    metrics:    Dict[str, float] = field(default_factory=dict)
    tags:       Dict[str, str]   = field(default_factory=dict)
    status:     str = "FINISHED"

@dataclass
class RegisteredModel:
    name:    str
    version: str
    run_id:  str
    stage:   str = "None"  # None, Staging, Production, Archived

class MLflowClient:
    """محاكاة MLflow Tracking + Registry"""

    def __init__(self):
        self._experiments: Dict[str, List[MLRun]] = {}
        self._registry:    Dict[str, List[RegisteredModel]] = {}
        self._counter      = 0

    def set_experiment(self, name: str):
        if name not in self._experiments:
            self._experiments[name] = []
        print(f"  🧪 Experiment: {name}")

    def start_run(self, experiment: str, **params) -> MLRun:
        self._counter += 1
        run = MLRun(
            run_id=f"run_{self._counter:04d}",
            experiment=experiment,
            params=params,
        )
        self._experiments[experiment].append(run)
        return run

    def log_metric(self, run: MLRun, key: str, value: float):
        run.metrics[key] = value

    def log_tag(self, run: MLRun, **tags):
        run.tags.update(tags)

    def register_model(self, run: MLRun, name: str, version: str) -> RegisteredModel:
        if name not in self._registry:
            self._registry[name] = []
        rm = RegisteredModel(name=name, version=version, run_id=run.run_id)
        self._registry[name].append(rm)
        print(f"  📦 Registered: {name} v{version} (run={run.run_id})")
        return rm

    def transition_stage(self, name: str, version: str, stage: str):
        for rm in self._registry.get(name, []):
            if rm.version == version:
                rm.stage = stage
                print(f"  🔄 {name} v{version}: None → {stage}")

    def compare_runs(self, experiment: str):
        runs = self._experiments.get(experiment, [])
        print(f"\n{'─'*60}")
        print(f"📊 مقارنة تجارب: {experiment}")
        print(f"{'─'*60}")
        header = f"  {'Run ID':<12} {'LR':>8} {'Epochs':>7} {'Accuracy':>10} {'F1':>8}"
        print(header)
        print("  " + "-"*50)
        for r in sorted(runs, key=lambda x: x.metrics.get("accuracy", 0), reverse=True):
            lr_s  = str(r.params.get("learning_rate", "-"))
            ep_s  = str(r.params.get("epochs", "-"))
            acc   = r.metrics.get("accuracy", 0)
            f1    = r.metrics.get("f1_score",  0)
            print(f"  {r.run_id:<12} {lr_s:>8} {ep_s:>7} {acc:>10.4f} {f1:>8.4f}")

    def best_run(self, experiment: str, metric: str = "accuracy") -> MLRun:
        runs = self._experiments.get(experiment, [])
        return max(runs, key=lambda r: r.metrics.get(metric, 0))

# ─── Hyperparameter Tuning ─────────────────────────────────
mlflow = MLflowClient()
exp    = "sentiment-classifier-v2"

print("🔬 MLflow — Hyperparameter Tuning:")
print("=" * 55)
mlflow.set_experiment(exp)

configs = [
    {"learning_rate": 0.001,  "epochs": 10, "model": "DistilBERT"},
    {"learning_rate": 0.0005, "epochs": 20, "model": "BERT-base"},
    {"learning_rate": 0.0001, "epochs": 30, "model": "RoBERTa"},
]

random.seed(42)
for cfg in configs:
    run = mlflow.start_run(exp, **cfg)
    acc = round(0.82 + random.uniform(0, 0.13), 4)
    f1  = round(acc - random.uniform(0.01, 0.05), 4)
    mlflow.log_metric(run, "accuracy",  acc)
    mlflow.log_metric(run, "f1_score",  f1)
    mlflow.log_metric(run, "val_loss",  round(random.uniform(0.05, 0.3), 4))
    mlflow.log_tag(run, framework="PyTorch", dataset="reviews-v3")
    print(f"  ✅ {run.run_id} — acc={acc:.4f}, f1={f1:.4f} ({cfg['model']})")

# مقارنة النتائج
mlflow.compare_runs(exp)

# أفضل نموذج → Registry
best = mlflow.best_run(exp)
print(f"\n🏆 أفضل نموذج: {best.run_id}")
best_acc_str = f"{best.metrics['accuracy']:.4f}"
best_lr_str  = str(best.params.get('learning_rate'))
print(f"   Accuracy: {best_acc_str}  LR: {best_lr_str}")

rm = mlflow.register_model(best, "SentimentClassifier", "1.0.0")

# Staging → Production
print(f"\n🔄 ترقية النموذج:")
mlflow.transition_stage("SentimentClassifier", "1.0.0", "Staging")
mlflow.transition_stage("SentimentClassifier", "1.0.0", "Production")
print(f"\n✅ MLflow Tracking & Registry مكتمل!")