import shutil
import time
from datetime import datetime
from pathlib import Path

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from torch.optim import Adam
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

import old_datasets as dataset_F
import utils
from models import DPNet


# 训练一轮
def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    total_samples = 0

    for images, labels, image_names in tqdm(
        loader,
        desc="Train:",
        bar_format="{l_bar}{bar:20}{r_bar}",
        leave=False,
    ):
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True).float().view(-1)

        optimizer.zero_grad(set_to_none=True)

        outputs = model(images).view(-1)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        this_batch_size = images.size(0)
        running_loss += loss.item() * this_batch_size
        total_samples += this_batch_size

    epoch_loss = running_loss / total_samples
    return epoch_loss


# 验证一轮
@torch.no_grad()
def valid_epoch(model, loader, criterion, device):
    model.eval()
    running_loss = 0.0
    total_samples = 0

    for images, labels, image_names in tqdm(
        loader,
        desc="Valid:",
        bar_format="{l_bar}{bar:20}{r_bar}",
        leave=False,
    ):
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True).float().view(-1)

        outputs = model(images).view(-1)
        loss = criterion(outputs, labels)

        this_batch_size = images.size(0)
        running_loss += loss.item() * this_batch_size
        total_samples += this_batch_size

    epoch_loss = running_loss / total_samples
    return epoch_loss


# 主训练函数
def main():
    # ========== 1 配置文件与超参数 ==========
    config, config_path = utils.get_hyperparams()

    XLSX_FILES = config["xlsx_files"]
    BATCH_SIZE = config["batch_size"]
    NUM_WORKERS = config["num_workers"]
    LEARNING_RATE = config["learning_rate"]
    NUM_EPOCHS = config["epochs"]
    SEED = config["seed"]
    INIT_WEIGHT_PATH = config["init_weight"]

    # ========== 2 创建输出文件目录 ==========
    run_name = datetime.now().strftime("%Y_%m_%d_%H_%M_%S_dpn")
    run_dir = Path.cwd() / run_name
    run_dir.mkdir(parents=True, exist_ok=False)
    shutil.copy2(config_path, run_dir / config_path.name)

    # ========== 3 日志、TensorBoard、随机种子与设备 ==========
    logger = utils.get_logger(__name__, run_dir / "train.log")
    writer = SummaryWriter(str(run_dir / "run"))
    utils.set_seeds(SEED)
    device = torch.device("cuda:0")

    logger.info(f"Config path: {config_path}")
    logger.info(f"Loaded config: {str(config)}")
    logger.info(f"Run directory: {run_dir}")
    logger.info(f"Using device: {device}")

    # ========== 4 数据与 loader ==========
    (
        train_image_path_list,
        train_defocus_distance_list,
        valid_image_path_list,
        valid_defocus_distance_list,
    ) = dataset_F.get_DPNet_train_data_and_label(root_path_list=XLSX_FILES)

    train_transform = transforms.Compose(
        [
            transforms.ColorJitter(
                brightness=(0.9, 1.4),
                contrast=(0.8, 1.5),
                saturation=(0.8, 1.5),
            ),
            transforms.ToTensor(),
        ]
    )
    valid_transform = transforms.Compose(
        [
            transforms.ToTensor(),
        ]
    )

    train_set = dataset_F.MyDataset(
        train_image_path_list,
        train_defocus_distance_list,
        train_transform,
    )
    valid_set = dataset_F.MyDataset(
        valid_image_path_list,
        valid_defocus_distance_list,
        valid_transform,
    )

    train_loader = DataLoader(
        dataset=train_set,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=NUM_WORKERS,
        pin_memory=True,
        persistent_workers=(NUM_WORKERS > 0),
    )
    valid_loader = DataLoader(
        dataset=valid_set,
        batch_size=BATCH_SIZE,
        shuffle=False,
        num_workers=NUM_WORKERS,
        pin_memory=True,
        persistent_workers=(NUM_WORKERS > 0),
    )

    logger.info(f"Train dataset size: {len(train_set)}")
    logger.info(f"Val dataset size: {len(valid_set)}")
    logger.info(f"Train steps per epoch: {len(train_loader)}")

    # ========== 5 模型、损失、优化器、调度器 ==========
    model = DPNet().to(device)
    if INIT_WEIGHT_PATH:
        state_dict = torch.load(INIT_WEIGHT_PATH, map_location="cpu")
        model.load_state_dict(state_dict, strict=True)
        logger.info(f"Loaded init weight from: {INIT_WEIGHT_PATH}")
    else:
        logger.info("Training from scratch")

    criterion = nn.MSELoss(reduction="mean")
    optimizer = Adam(model.parameters(), lr=LEARNING_RATE, betas=(0.9, 0.999))
    scheduler = utils.get_warmup_cosine_scheduler(optimizer, NUM_EPOCHS)

    logger.info("Loss: MSELoss(reduction='mean')")
    logger.info(f"Optimizer: Adam(lr={LEARNING_RATE}, betas=(0.9, 0.999))")
    logger.info("Scheduler: epoch-based warmup + cosine annealing")

    # ========== 6 开始训练 ==========
    logger.info("START TRAINING")
    best_valid_loss = float("inf")

    try:
        for epoch in range(1, NUM_EPOCHS + 1):
            epoch_start_time = time.time()
            train_loss = train_epoch(model, train_loader, criterion, optimizer, device)
            valid_loss = valid_epoch(model, valid_loader, criterion, device)
            epoch_lr = optimizer.param_groups[0]["lr"]
            scheduler.step()
            epoch_time_cost = time.time() - epoch_start_time

            if valid_loss < best_valid_loss:
                best_valid_loss = valid_loss
                torch.save(model.state_dict(), run_dir / "best_model.pt")
                logger.info(f"Best model saved, valid_loss = {best_valid_loss:.4f}")

            logger.info(
                f"Epoch [{epoch}/{NUM_EPOCHS}] "
                f"Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | "
                f"Best Valid Loss: {best_valid_loss:.4f} | "
                f"Epoch Time Cost: {epoch_time_cost:.2f} s | "
                f"Epoch Learning Rate: {epoch_lr:.6e}"
            )

            writer.add_scalar("Loss/train", train_loss, epoch)
            writer.add_scalar("Loss/valid", valid_loss, epoch)
            writer.add_scalar("Loss/best_valid", best_valid_loss, epoch)
            writer.add_scalar("Time/epoch", epoch_time_cost, epoch)
            writer.add_scalar("Time/learning_rate", epoch_lr, epoch)

    except KeyboardInterrupt:
        logger.info("Training interrupted by user")

    finally:
        torch.save(model.state_dict(), run_dir / "last_model.pt")
        logger.info("Last model saved")

        writer.close()
        logger.info("TensorBoard writer closed")

        logger.info(f"Training finished, best validation loss: {best_valid_loss:.8f}")


if __name__ == "__main__":
    main()