# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import sys
from unittest import mock

import numpy as np
import torch

from tests import MODEL, SOURCE
from ultralytics import YOLO
from ultralytics.cfg import get_cfg
from ultralytics.engine.exporter import Exporter
from ultralytics.models.yolo import classify, detect, segment
from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR


def test_func(*args, **kwargs):
    """Test function used as a callback stub to verify callback registration."""
    print("callback test passed")


def test_export():
    """Test model exporting functionality by adding a callback and verifying its execution."""
    exporter = Exporter()
    exporter.add_callback("on_export_start", test_func)
    assert test_func in exporter.callbacks["on_export_start"], "callback test failed"
    f = exporter(model=YOLO("yolo26n.yaml").model)
    YOLO(f)(SOURCE)  # exported model inference


def test_detect():
    """Test YOLO object detection training, validation, and prediction functionality."""
    overrides = {"data": "coco8.yaml", "model": "yolo26n.yaml", "imgsz": 32, "epochs": 1, "save": False}
    cfg = get_cfg(DEFAULT_CFG)
    cfg.data = "coco8.yaml"
    cfg.imgsz = 32

    # Trainer
    trainer = detect.DetectionTrainer(overrides=overrides)
    trainer.add_callback("on_train_start", test_func)
    assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
    trainer.train()

    # Validator
    val = detect.DetectionValidator(args=cfg)
    val.add_callback("on_val_start", test_func)
    assert test_func in val.callbacks["on_val_start"], "callback test failed"
    val(model=trainer.best)  # validate best.pt

    # Predictor
    pred = detect.DetectionPredictor(overrides={"imgsz": [64, 64]})
    pred.add_callback("on_predict_start", test_func)
    assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
    # Confirm there is no issue with sys.argv being empty
    with mock.patch.object(sys, "argv", []):
        result = pred(source=ASSETS, model=MODEL)
        assert len(result), "predictor test failed"

    # Test resume functionality
    overrides["resume"] = trainer.last
    trainer = detect.DetectionTrainer(overrides=overrides)
    try:
        trainer.train()
    except Exception as e:
        print(f"Expected exception caught: {e}")
        return

    raise Exception("Resume test failed!")


def test_segment():
    """Test image segmentation training, validation, and prediction pipelines using YOLO models."""
    overrides = {
        "data": "coco8-seg.yaml",
        "model": "yolo26n-seg.yaml",
        "imgsz": 32,
        "epochs": 1,
        "save": False,
        "mask_ratio": 1,
        "overlap_mask": False,
    }
    cfg = get_cfg(DEFAULT_CFG)
    cfg.data = "coco8-seg.yaml"
    cfg.imgsz = 32

    # Trainer
    trainer = segment.SegmentationTrainer(overrides=overrides)
    trainer.add_callback("on_train_start", test_func)
    assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
    trainer.train()

    # Validator
    val = segment.SegmentationValidator(args=cfg)
    val.add_callback("on_val_start", test_func)
    assert test_func in val.callbacks["on_val_start"], "callback test failed"
    val(model=trainer.best)  # validate best.pt

    # Predictor
    pred = segment.SegmentationPredictor(overrides={"imgsz": [64, 64]})
    pred.add_callback("on_predict_start", test_func)
    assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
    result = pred(source=ASSETS, model=WEIGHTS_DIR / "yolo26n-seg.pt")
    assert len(result), "predictor test failed"

    # Test resume functionality
    overrides["resume"] = trainer.last
    trainer = segment.SegmentationTrainer(overrides=overrides)
    try:
        trainer.train()
    except Exception as e:
        print(f"Expected exception caught: {e}")
        return

    raise Exception("Resume test failed!")


def test_classify():
    """Test image classification including training, validation, and prediction phases."""
    overrides = {"data": "imagenet10", "model": "yolo26n-cls.yaml", "imgsz": 32, "epochs": 1, "save": False}
    cfg = get_cfg(DEFAULT_CFG)
    cfg.data = "imagenet10"
    cfg.imgsz = 32

    # Trainer
    trainer = classify.ClassificationTrainer(overrides=overrides)
    trainer.add_callback("on_train_start", test_func)
    assert test_func in trainer.callbacks["on_train_start"], "callback test failed"
    trainer.train()

    # Validator
    val = classify.ClassificationValidator(args=cfg)
    val.add_callback("on_val_start", test_func)
    assert test_func in val.callbacks["on_val_start"], "callback test failed"
    val(model=trainer.best)

    # Predictor
    pred = classify.ClassificationPredictor(overrides={"imgsz": [64, 64]})
    pred.add_callback("on_predict_start", test_func)
    assert test_func in pred.callbacks["on_predict_start"], "callback test failed"
    result = pred(source=ASSETS, model=trainer.best)
    assert len(result), "predictor test failed"


def test_nan_recovery():
    """Test NaN loss detection and recovery during training."""
    nan_injected = [False]

    def inject_nan(trainer):
        """Inject NaN into loss during batch processing to test recovery mechanism."""
        if trainer.epoch == 1 and trainer.tloss is not None and not nan_injected[0]:
            trainer.tloss *= torch.tensor(float("nan"))
            nan_injected[0] = True

    overrides = {"data": "coco8.yaml", "model": "yolo26n.yaml", "imgsz": 32, "epochs": 3}
    trainer = detect.DetectionTrainer(overrides=overrides)
    trainer.add_callback("on_train_batch_end", inject_nan)
    trainer.train()
    assert nan_injected[0], "NaN injection failed"


def test_ground3d_validator_prints_combined_metrics(caplog):
    """Test Ground3DDetectionValidator prints whole and face 3D metrics."""
    from ultralytics.models.yolo.detect.train import Ground3DDetectionValidator

    validator = Ground3DDetectionValidator(args=get_cfg(DEFAULT_CFG))
    validator.names = {0: "car"}
    validator.args.task = "detect"
    validator.args.verbose = False
    validator.training = True
    validator.seen = 3531
    validator.nc = 1
    validator.stats = []
    validator.metrics.nt_per_class = torch.tensor([59691])
    validator.metrics.mean_results = lambda: [0.73, 0.585, 0.657, 0.499]
    validator.metrics.keys = ["metrics/precision(B)", "metrics/recall(B)", "metrics/mAP50(B)", "metrics/mAP50-95(B)"]
    validator.metrics_3d_results = {
        "whole": {
            "depth_abs": 48.72,
            "depth_rel": 1.915,
            "depth_rmse": 54.22,
            "center": 52.01,
            "uv": 17.35,
            "orient": 85.47,
            "size": 2.011,
            "matched": 5738,
        },
        "face": {
            "depth_abs": 12.31,
            "depth_rel": 0.441,
            "depth_rmse": 18.22,
            "center": 14.07,
            "uv": 9.56,
            "matched": 2814,
        },
    }

    with caplog.at_level("INFO", logger="ultralytics"):
        validator.print_results()

    all_lines = [record.message for record in caplog.records if "all" in record.message]
    assert len(all_lines) == 1
    all_line = all_lines[0]
    assert "0.73" in all_line
    assert "48.7" in all_line
    assert "1.92" in all_line
    assert "54.2" in all_line
    assert "52" in all_line
    assert "17.4" in all_line
    assert "85.5" in all_line
    assert "2.01" in all_line
    assert "5738" in all_line

    face_lines = [record.message for record in caplog.records if "3d-face" in record.message]
    assert len(face_lines) == 1
    face_line = face_lines[0]
    assert "12.3" in face_line
    assert "0.441" in face_line
    assert "18.2" in face_line
    assert "14.1" in face_line
    assert "9.56" in face_line
    assert "2814" in face_line


def test_ground3d_validator_prints_invalid_visible_yaw_as_dash(caplog):
    """Test Ground3DDetectionValidator renders invalid visible-yaw metrics as '-'."""
    from ultralytics.models.yolo.detect.train import Ground3DDetectionValidator

    validator = Ground3DDetectionValidator(args=get_cfg(DEFAULT_CFG))
    validator.names = {0: "car"}
    validator.args.task = "detect"
    validator.args.verbose = False
    validator.training = True
    validator.seen = 1
    validator.nc = 1
    validator.stats = []
    validator.metrics.nt_per_class = torch.tensor([1])
    validator.metrics.mean_results = lambda: [0.73, 0.585, 0.657, 0.499]
    validator.metrics.keys = ["metrics/precision(B)", "metrics/recall(B)", "metrics/mAP50(B)", "metrics/mAP50-95(B)"]
    validator.metrics_3d_results = {
        "whole": {
            "depth_abs": 1.0,
            "depth_rel": 0.1,
            "depth_rmse": 1.5,
            "center": 2.0,
            "uv": 3.0,
            "orient": 4.0,
            "size": 0.5,
            "matched": 1,
        },
        "face": {
            "depth_abs": 2.0,
            "depth_rel": 0.2,
            "depth_rmse": 2.5,
            "center": 4.0,
            "uv": 5.0,
            "size": 0.6,
            "direct_orient_visible": np.nan,
            "edge_orient_visible": np.nan,
            "matched": 1,
        },
    }

    with caplog.at_level("INFO", logger="ultralytics"):
        validator.print_results()

    all_lines = [record.message for record in caplog.records if "all-3d" in record.message]
    assert len(all_lines) == 1
    assert "      nan" not in all_lines[0]
    assert all_lines[0].rstrip().endswith("-")


def test_detect3d_select_topk_3d_metadata_keeps_batch_alignment():
    """Test Detect3D keeps selected 3D predictions, anchors, and strides aligned per sample."""
    from ultralytics.nn.modules.head import Detect3D

    detect3d = Detect3D(nc=2, reg_max=1, end2end=True, ch=(16,))
    detect3d.anchors = torch.tensor([[0.5, 1.5, 2.5, 3.5], [10.5, 11.5, 12.5, 13.5]], dtype=torch.float32)
    detect3d.strides = torch.tensor([[8.0, 16.0, 32.0, 64.0]], dtype=torch.float32)

    preds_3d = torch.arange(2 * detect3d.no_3d * 4, dtype=torch.float32).reshape(2, detect3d.no_3d, 4)
    idx = torch.tensor([[[0], [2]], [[1], [3]]], dtype=torch.long)

    preds_sel, anchors_sel, strides_sel = detect3d._select_topk_3d_metadata(preds_3d, idx)

    assert preds_sel.shape == (2, 2, detect3d.no_3d)
    assert anchors_sel.shape == (2, 2, 2)
    assert strides_sel.shape == (2, 2)

    torch.testing.assert_close(preds_sel[0, 0], preds_3d[0, :, 0])
    torch.testing.assert_close(preds_sel[0, 1], preds_3d[0, :, 2])
    torch.testing.assert_close(preds_sel[1, 0], preds_3d[1, :, 1])
    torch.testing.assert_close(preds_sel[1, 1], preds_3d[1, :, 3])
    torch.testing.assert_close(anchors_sel[0], detect3d.anchors[:, [0, 2]])
    torch.testing.assert_close(anchors_sel[1], detect3d.anchors[:, [1, 3]])
    torch.testing.assert_close(strides_sel[0], detect3d.strides[0, [0, 2]])
    torch.testing.assert_close(strides_sel[1], detect3d.strides[0, [1, 3]])