import mlflow
import torch
from sklearn.metrics import average_precision_score, precision_score, recall_score
from src.models.evaluate import get_probs, metrics_at, best_threshold


def train_one_epoch(model, dataloader, loss_fn, optimizer, device):
    model.train()
    train_loss, correct, n_samples, n_batches = 0.0, 0, 0, 0
    all_probs, all_y = [], []

    for X, y in dataloader:
        X, y = X.to(device), y.to(device)

        logits = model(X)
        loss = loss_fn(logits, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        n_batches += 1

        preds = (logits > 0).float()
        correct += (preds == y).sum().item()
        n_samples += y.numel()

        all_probs.append(torch.sigmoid(logits).detach().cpu())
        all_y.append(y.detach().cpu())

    probs = torch.cat(all_probs).numpy().ravel()
    y_true = torch.cat(all_y).numpy().ravel()
    preds_np = (probs > 0.5).astype(int)

    return {
        "loss": train_loss / n_batches,
        "acc": correct / n_samples,
        "pr_auc": average_precision_score(y_true, probs),
        "precision": precision_score(y_true, preds_np, zero_division=0),
        "recall": recall_score(y_true, preds_np, zero_division=0),
    }


def run_training(model, train_dl, test_dl, epochs, lr):
    torch.manual_seed(42)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)

    # pos_weight = negatives / positives, from the TRAIN split
    y_all = torch.cat([y for _, y in train_dl])
    n_pos = y_all.sum().item()
    n_neg = y_all.numel() - n_pos
    pos_weight = torch.tensor([n_neg / max(n_pos, 1)], device=device)
    print(f"pos_weight = {pos_weight.item():.1f}  ({int(n_pos)} pos / {int(n_neg)} neg)")

    loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    with mlflow.start_run():
        for epoch in range(epochs):
            m = train_one_epoch(model, train_dl, loss_fn, optimizer, device)
            for k, v in m.items():
                mlflow.log_metric(f"train_{k}", v, step=epoch)
            print(f"epoch {epoch}: loss={m['loss']:.4f} acc={m['acc']:.4f} "
                  f"pr_auc={m['pr_auc']:.4f} P={m['precision']:.3f} R={m['recall']:.3f}")
                # ---------- evaluate on the held-out TEST set ----------
        probs, y_true = get_probs(model, test_dl)

        # (a) the naive 0.5 cutoff 
        base = metrics_at(probs, y_true, 0.5)
        print(f"\nTEST @0.50 : P={base['precision']:.3f} R={base['recall']:.3f} "
              f"F1={base['f1']:.3f}  PR-AUC={base['pr_auc']:.3f}")

        # (b) the F1-best cutoff — same model, better trade-off
        thr = best_threshold(probs, y_true)
        tuned = metrics_at(probs, y_true, thr)
        print(f"TEST @{thr:.2f} : P={tuned['precision']:.3f} R={tuned['recall']:.3f} "
              f"F1={tuned['f1']:.3f}")

        # log BOTH so you can compare in MLflow
        for k, v in base.items():
            mlflow.log_metric(f"test_{k}", v)
        for k, v in tuned.items():
            mlflow.log_metric(f"test_tuned_{k}", v)

    return model