import argparse
import math
import os
import random
import subprocess
import sys
import time
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path

try:
    import comet_ml  # must be imported before torch (if installed)
except ImportError:
    comet_ml = None

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

import val as validate  # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.downloads import attempt_download, is_url
from utils.general import (
    LOGGER,
    TQDM_BAR_FORMAT,
    check_amp,
    check_dataset,
    check_file,
    check_git_info,
    check_git_status,
    check_img_size,
    check_requirements,
    check_suffix,
    check_yaml,
    colorstr,
    get_latest_run,
    increment_path,
    init_seeds,
    intersect_dicts,
    labels_to_class_weights,
    labels_to_image_weights,
    methods,
    one_cycle,
    print_args,
    print_mutation,
    strip_optimizer,
    yaml_save,
)
from utils.loggers import LOGGERS, Loggers
from utils.loggers.comet.comet_utils import check_comet_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve
from utils.torch_utils import (
    EarlyStopping,
    ModelEMA,
    de_parallel,
    select_device,
    smart_DDP,
    smart_optimizer,
    smart_resume,
    torch_distributed_zero_first,
)

LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))
RANK = int(os.getenv("RANK", -1))
WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
GIT_INFO = check_git_info()

def train(hyp, opt, device, callbacks):
    """
    Train a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets,
    model architecture, loss computation, and optimizer steps.

    Args:
        hyp (str | dict): Path to the hyperparameters YAML file or a dictionary of hyperparameters.
        opt (argparse.Namespace): Parsed command-line arguments containing training options.
        device (torch.device): Device on which training occurs, e.g., 'cuda' or 'cpu'.
        callbacks (Callbacks): Callback functions for various training events.

    Returns:
        None

    Models and datasets download automatically from the latest YOLOv5 release.

    Example:
        Single-GPU training:
        ```bash
        $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640  # from pretrained (recommended)
        $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640  # from scratch
        ```

        Multi-GPU DDP training:
        ```bash
        $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights
        yolov5s.pt --img 640 --device 0,1,2,3
        ```

        For more usage details, refer to:
        - Models: https://github.com/ultralytics/yolov5/tree/master/models
        - Datasets: https://github.com/ultralytics/yolov5/tree/master/data
        - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data
    """

# Trainloader
train_loader, dataset = create_dataloader(
    train_path,
    imgsz,
    batch_size // WORLD_SIZE,
    gs,
    single_cls,
    hyp=hyp,
    augment=True, # 数据增强在训练中默认开启
    cache=None if opt.cache == "val" else opt.cache,
    rect=opt.rect,
    rank=LOCAL_RANK,
    workers=workers,
    image_weights=opt.image_weights,
    quad=opt.quad,
    prefix=colorstr("train: "),
    shuffle=True,
    seed=opt.seed,
)

# Process 0
# 验证数据加载器中无数据增强：
# 相比之下，验证数据加载器中没有启用数据增强，这是合理的，因为验证应该在不增强的原始数据上进行：
# 注意这里没有指定augment参数，它会使用create_dataloader函数的默认值False。
if RANK in {-1, 0}:
    val_loader = create_dataloader(
        val_path,
        imgsz,
        batch_size // WORLD_SIZE * 2,
        gs,
        single_cls,
        hyp=hyp,
        cache=None if noval else opt.cache,
        rect=True,
        rank=-1,
        workers=workers * 2,
        pad=0.5,
        prefix=colorstr("val: "),
    )[0]

def create_dataloader(
    path,
    imgsz,
    batch_size,
    stride,
    single_cls=False,
    hyp=None,
    augment=False,
    cache=False,
    pad=0.0,
    rect=False,
    rank=-1,
    workers=8,
    image_weights=False,
    quad=False,
    prefix="",
    shuffle=False,
    seed=0,
):
    # 创建数据集实例
    dataset = LoadImagesAndLabels(
        path,
        imgsz,
        batch_size,
        # 其他参数...
    )
  
    # 配置批次大小和采样器
    batch_size = min(batch_size, len(dataset))
    sampler = None if rank == -1 else distributed.DistributedSampler(...)
  
    # 选择数据加载器类型
    loader = InfiniteDataLoader if image_weights else DataLoader
  
    # 创建并返回数据加载器
    return loader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle and sampler is None,
        # 其他参数...
    ), dataset

def __getitem__(self, index):
    # 1. 获取索引
    index = self.indices[index]
  
    # 2. 决定是否使用Mosaic增强
    if self.mosaic and random.random() < self.hyp["mosaic"]:
        # 加载Mosaic增强图像
        img, labels = self.load_mosaic(index)
      
        # 可能应用MixUp增强
        if random.random() < self.hyp["mixup"]:
            img, labels = mixup(...)
    else:
        # 常规加载图像
        img, (h0, w0), (h, w) = self.load_image(index)
        # 应用Letterbox
        img, ratio, pad = letterbox(...)
        # 处理标签
        labels = self.labels[index].copy()
      
    # 3. 应用更多增强操作
    if self.augment:
        # 随机透视变换
        img, labels = random_perspective(...)
      
        # Albumentations库增强
        img, labels = self.albumentations(img, labels)
      
        # HSV颜色空间增强
        augment_hsv(...)
      
        # 随机翻转
        if random.random() < self.hyp["flipud"]:
            img = np.flipud(img)
          
        if random.random() < self.hyp["fliplr"]:
            img = np.fliplr(img)
  
    # 4. 最终处理
    # 标签格式转换
    labels_out = torch.zeros((len(labels), 6))
    # 图像格式转换
    img = img.transpose((2, 0, 1))[::-1]
  
    return torch.from_numpy(img), labels_out, self.im_files[index], shapes

pbar = enumerate(train_loader)
...
for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
    callbacks.run("on_train_batch_start")
    ni = i + nb * epoch  # number integrated batches (since train start)
    imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0
    ...

def __getitem__(self, index):
    """获取数据集中指定索引的样本，考虑线性、随机或加权采样。"""
    index = self.indices[index]  # 线性、随机或图像权重
  
    hyp = self.hyp
    if mosaic := self.mosaic and random.random() < hyp["mosaic"]:
        # 加载Mosaic增强
        img, labels = self.load_mosaic(index)
        shapes = None
      
        # MixUp增强
        if random.random() < hyp["mixup"]:
            img, labels = mixup(img, labels, *self.load_mosaic(random.choice(self.indices)))
    else:
        # 常规加载图像
        img, (h0, w0), (h, w) = self.load_image(index)
      
        # Letterbox
        shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size
        img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
        shapes = (h0, w0), ((h / h0, w / w0), pad)
      
        # 处理标签
        labels = self.labels[index].copy()
        if labels.size:
            labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
      
        # 随机透视变换
        if self.augment:
            img, labels = random_perspective(
                img,
                labels,
                degrees=hyp["degrees"],
                translate=hyp["translate"],
                scale=hyp["scale"],
                shear=hyp["shear"],
                perspective=hyp["perspective"],
            )
  
    nl = len(labels)  # 标签数量
    if nl:
        labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
  
    # 更多数据增强操作
    if self.augment:
        # Albumentations库增强
        img, labels = self.albumentations(img, labels)
        nl = len(labels)  # 更新标签数量
      
        # HSV颜色空间增强
        augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
      
        # 上下翻转
        if random.random() < hyp["flipud"]:
            img = np.flipud(img)
            if nl:
                labels[:, 2] = 1 - labels[:, 2]
      
        # 左右翻转
        if random.random() < hyp["fliplr"]:
            img = np.fliplr(img)
            if nl:
                labels[:, 1] = 1 - labels[:, 1]
      
        # Cutout(被注释掉了)
        # labels = cutout(img, labels, p=0.5)
  
    # 格式转换
    labels_out = torch.zeros((nl, 6))
    if nl:
        labels_out[:, 1:] = torch.from_numpy(labels)
  
    # 图像格式转换：HWC到CHW，BGR到RGB
    img = img.transpose((2, 0, 1))[::-1]
    img = np.ascontiguousarray(img)
  
    return torch.from_numpy(img), labels_out, self.im_files[index], shapes

if mosaic := self.mosaic and random.random() < hyp["mosaic"]:
    img, labels = self.load_mosaic(index)

if random.random() < hyp["mixup"]:
    img, labels = mixup(img, labels, *self.load_mosaic(random.choice(self.indices)))

if self.augment:
    img, labels = random_perspective(
        img,
        labels,
        degrees=hyp["degrees"],
        translate=hyp["translate"],
        scale=hyp["scale"],
        shear=hyp["shear"],
        perspective=hyp["perspective"],
    )

if self.augment:
    img, labels = self.albumentations(img, labels)

if self.augment:
    augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])

if self.augment:
    if random.random() < hyp["flipud"]:
        img = np.flipud(img)
        if nl:
            labels[:, 2] = 1 - labels[:, 2]
  
    if random.random() < hyp["fliplr"]:
        img = np.fliplr(img)
        if nl:
            labels[:, 1] = 1 - labels[:, 1]