import mlflow
import pandas as pd
from mlflow.tracking import MlflowClient


class MLflowManager:

    def __init__(self, experiment_name="chandanops"):
        mlflow.set_experiment(experiment_name)
        self.client = MlflowClient()

    def start_run(self, run_name=None):
        return mlflow.start_run(run_name=run_name)

    def log_params(self, params: dict):
        mlflow.log_params(params)

    def log_metrics(self, metrics: dict):
        mlflow.log_metrics(metrics)

    def log_model(self, model, registered_model_name=None):
        mlflow.sklearn.log_model(
            model,
            artifact_path="model",
            registered_model_name=registered_model_name
        )

    def compare_runs(self, metric_name, top_n=5):
        runs = mlflow.search_runs(order_by=[f"metrics.{metric_name} DESC"])
        if runs.empty:
            return None
        return runs[["run_id", f"metrics.{metric_name}"]].head(top_n)
