import os
import random
import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import time


# --- 1. 配置参数 ---
class Config:
    IMAGE_DIR = "data_up/images"
    MASK_DIR = "data_up/masks_line1"
    IMAGE_SIZE = 256  # 将所有图片缩放到 256x256
    BATCH_SIZE = 4
    EPOCHS = 50  # 训练轮数
    LEARNING_RATE = 1e-4
    TEST_SPLIT = 0.1  # 20% 的数据用作验证集


# --- 2. 数据集加载和预处理 ---
class WeldSeamDataset(Dataset):
    def __init__(self, image_paths, mask_paths, size):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.size = size

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        # 读取图像
        img = cv2.imread(self.image_paths[idx], cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, (self.size, self.size))
        img = img / 255.0  # 归一化到 [0, 1]
        img = np.expand_dims(img, axis=0)  # 增加通道维度 (H, W) -> (C, H, W)
        img_tensor = torch.from_numpy(img).float()

        # 读取掩码
        mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE)
        mask = cv2.resize(mask, (self.size, self.size))
        mask = mask / 255.0  # 归一化到 {0, 1}
        mask[mask > 0.5] = 1.0
        mask[mask <= 0.5] = 0.0
        mask = np.expand_dims(mask, axis=0)
        mask_tensor = torch.from_numpy(mask).float()

        return img_tensor, mask_tensor


# --- 3. LinkNet模型定义 ---
class DecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, in_channels // 4, kernel_size=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(in_channels // 4, in_channels // 4, kernel_size=2, stride=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels // 4, out_channels, kernel_size=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.block(x)


class LinkNet(nn.Module):
    def __init__(self, num_classes=1):
        super().__init__()
        # 使用预训练的ResNet18作为编码器
        resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

        self.firstconv = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.firstbn = resnet.bn1
        self.firstrelu = resnet.relu
        self.firstmaxpool = resnet.maxpool

        # 编码器层
        self.encoder1 = resnet.layer1
        self.encoder2 = resnet.layer2
        self.encoder3 = resnet.layer3
        self.encoder4 = resnet.layer4

        # 解码器层
        self.decoder4 = DecoderBlock(512, 256)
        self.decoder3 = DecoderBlock(256, 128)
        self.decoder2 = DecoderBlock(128, 64)
        self.decoder1 = DecoderBlock(64, 64)

        # 最终输出层
        self.final_deconv = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.final_relu = nn.ReLU(inplace=True)
        self.final_conv = nn.Conv2d(32, num_classes, kernel_size=1)

    def forward(self, x):
        # 编码器
        x = self.firstconv(x)
        x = self.firstbn(x)
        x = self.firstrelu(x)
        x = self.firstmaxpool(x)  # -> 64x64
        e1 = self.encoder1(x)  # -> 64x64
        e2 = self.encoder2(e1)  # -> 32x32
        e3 = self.encoder3(e2)  # -> 16x16
        e4 = self.encoder4(e3)  # -> 8x8

        # 解码器
        d4 = self.decoder4(e4) + e3  # -> 16x16
        d3 = self.decoder3(d4) + e2  # -> 32x32
        d2 = self.decoder2(d3) + e1  # -> 64x64
        d1 = self.decoder1(d2)  # -> 128x128

        f = self.final_deconv(d1)  # -> 256x256
        f = self.final_relu(f)
        f = self.final_conv(f)

        return torch.sigmoid(f)  # 使用Sigmoid输出概率图


# --- 4. 损失函数 (Dice Loss + BCE Loss) ---
def dice_loss(pred, target, smooth=1.):
    pred = pred.contiguous()
    target = target.contiguous()
    intersection = (pred * target).sum(dim=2).sum(dim=2)
    loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))
    return loss.mean()


def bce_dice_loss(pred, target):
    bce = nn.BCELoss()(pred, target)
    dice = dice_loss(pred, target)
    return bce + dice


# --- 5. 训练和评估 ---
def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    print(f"Training on {device}")

    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        start_time = time.time()
        model.train()
        running_loss = 0.0

        for images, masks in train_loader:
            images = images.to(device)
            masks = masks.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()

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

        epoch_loss = running_loss / len(train_loader.dataset)

        # 验证
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for images, masks in val_loader:
                images = images.to(device)
                masks = masks.to(device)
                outputs = model(images)
                loss = criterion(outputs, masks)
                val_loss += loss.item() * images.size(0)

        val_loss /= len(val_loader.dataset)

        duration = time.time() - start_time
        print(f"Epoch {epoch + 1}/{num_epochs}.. "
              f"Train Loss: {epoch_loss:.4f}.. "
              f"Val Loss: {val_loss:.4f}.. "
              f"Time: {duration:.2f}s")

        # 保存最佳模型
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_linknet_model.pth')
            print("Model saved!")

    print("Training complete.")


def predict_and_visualize(model, image_path, model_path, size):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)
    model.eval()

    # 加载和预处理单张图片
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    original_size = (img.shape[1], img.shape[0])  # (width, height)
    img_resized = cv2.resize(img, (size, size))
    img_normalized = img_resized / 255.0
    img_tensor = torch.from_numpy(np.expand_dims(np.expand_dims(img_normalized, axis=0), axis=0)).float()
    img_tensor = img_tensor.to(device)

    with torch.no_grad():
        output = model(img_tensor)

    # 后处理
    pred_mask = output.cpu().numpy()[0, 0]  # 从 (B, C, H, W) -> (H, W)
    pred_mask = (pred_mask > 0.5).astype(np.uint8) * 255  # 二值化
    pred_mask = cv2.resize(pred_mask, original_size)  # 恢复到原始尺寸

    # 可视化
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 3, 1)
    plt.title("Original Image")
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_GRAY2RGB))

    plt.subplot(1, 3, 2)
    plt.title("Predicted Mask")
    plt.imshow(pred_mask, cmap='gray')

    # 将掩码叠加到原图
    overlay = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    overlay[pred_mask == 255] = [255, 0, 0]  # 红色
    plt.subplot(1, 3, 3)
    plt.title("Overlay")
    plt.imshow(overlay)

    plt.show()


# --- 6. 主程序入口 ---
if __name__ == '__main__':
    cfg = Config()

    # 准备数据集
    image_files = sorted([os.path.join(cfg.IMAGE_DIR, f) for f in os.listdir(cfg.IMAGE_DIR)])
    mask_files = sorted([os.path.join(cfg.MASK_DIR, f) for f in os.listdir(cfg.MASK_DIR)])

    # 划分训练集和验证集
    train_imgs, val_imgs, train_masks, val_masks = train_test_split(
        image_files, mask_files, test_size=cfg.TEST_SPLIT, random_state=42
    )

    train_dataset = WeldSeamDataset(train_imgs, train_masks, cfg.IMAGE_SIZE)
    val_dataset = WeldSeamDataset(val_imgs, val_masks, cfg.IMAGE_SIZE)

    train_loader = DataLoader(train_dataset, batch_size=cfg.BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=cfg.BATCH_SIZE, shuffle=False)

    # 初始化模型、损失函数和优化器
    model = LinkNet(num_classes=1)
    criterion = bce_dice_loss
    optimizer = optim.Adam(model.parameters(), lr=cfg.LEARNING_RATE)

    # --- 训练模型 ---
    # 如果你想开始训练，取消下面这行的注释
    train_model(model, train_loader, val_loader, criterion, optimizer, cfg.EPOCHS)

    # --- 使用训练好的模型进行预测 ---
    # 训练完成后，使用这个函数来测试
    # 确保 'best_linknet_bottom_model_line1.pth' 文件存在
    # print("\n--- Running Prediction ---")
    # # 随机选择一张验证集图片进行测试
    # test_image_path = random.choice(val_imgs)
    # print(f"Predicting on image: {test_image_path}")
    # predict_and_visualize(model, test_image_path, 'best_linknet_up_model_line1.pth', cfg.IMAGE_SIZE)