import tensorflow as tf
from keras import Input, Model
from keras.layers import Conv2D, MaxPooling2D, concatenate, BatchNormalization, Activation, Dropout
from keras.layers import ConvLSTM2D, TimeDistributed, UpSampling2D
from keras.regularizers import l2
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.preprocessing.image import ImageDataGenerator
import os
import numpy as np
from keras_preprocessing.image import load_img, img_to_array
from sklearn.model_selection import train_test_split
import random

def convlstm_unet(input_size=(None, None, None, 3), timesteps=3, dropout_rate=0.5, l2_reg=1e-4):
    # 输入：时间序列的图像，形状为 (batch_size, timesteps, height, width, channels)
    inputs = Input(shape=input_size)
    
    # 下采样路径
    # 第一层
    c1 = ConvLSTM2D(filters=64, kernel_size=(3, 3), padding='same', return_sequences=True,
                   kernel_regularizer=l2(l2_reg))(inputs)
    c1 = BatchNormalization()(c1)
    c1 = Activation('relu')(c1)
    c1 = TimeDistributed(Dropout(dropout_rate))(c1)  # 添加Dropout
    p1 = TimeDistributed(MaxPooling2D(pool_size=(2, 2)))(c1)

    # 第二层
    c2 = ConvLSTM2D(filters=128, kernel_size=(3, 3), padding='same', return_sequences=True,
                   kernel_regularizer=l2(l2_reg))(p1)
    c2 = BatchNormalization()(c2)
    c2 = Activation('relu')(c2)
    c2 = TimeDistributed(Dropout(dropout_rate))(c2)  # 添加Dropout
    p2 = TimeDistributed(MaxPooling2D(pool_size=(2, 2)))(c2)

    # 第三层
    c3 = ConvLSTM2D(filters=256, kernel_size=(3, 3), padding='same', return_sequences=False,
                   kernel_regularizer=l2(l2_reg))(p2)
    c3 = BatchNormalization()(c3)
    c3 = Activation('relu')(c3)
    c3 = Dropout(dropout_rate)(c3)  # 添加Dropout

    # 上采样路径
    u1 = UpSampling2D(size=(2, 2))(c3)
    # 将 c2 的最后一个时间步与上采样结果进行拼接
    merge1 = concatenate([u1, c2[:, -1, :, :, :]], axis=-1)
    c4 = Conv2D(filters=128, kernel_size=(3, 3), padding='same', kernel_regularizer=l2(l2_reg))(merge1)
    c4 = BatchNormalization()(c4)
    c4 = Activation('relu')(c4)
    c4 = Dropout(dropout_rate)(c4)  # 添加Dropout

    u2 = UpSampling2D(size=(2, 2))(c4)
    merge2 = concatenate([u2, c1[:, -1, :, :, :]], axis=-1)
    c5 = Conv2D(filters=64, kernel_size=(3, 3), padding='same', kernel_regularizer=l2(l2_reg))(merge2)
    c5 = BatchNormalization()(c5)
    c5 = Activation('relu')(c5)
    c5 = Dropout(dropout_rate)(c5)  # 添加Dropout

    outputs = Conv2D(filters=1, kernel_size=(1, 1), activation='sigmoid')(c5)
    
    model = Model(inputs=inputs, outputs=outputs)
    return model

# 创建模型
model = convlstm_unet(input_size=(3, 128, 128, 3))
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 打印模型结构
model.summary()

# 加载图像序列和对应的掩膜序列
def load_sequence_data(image_folder, mask_folder, timesteps):
    image_files = sorted(os.listdir(image_folder))
    mask_files = sorted(os.listdir(mask_folder))
    
    images = []
    masks = []

    # 确保图像和掩膜数量一致
    assert len(image_files) == len(mask_files), "图像和掩膜数量不一致"

    num_samples = len(image_files) - timesteps + 1

    for i in range(num_samples):
        img_sequence = []
        mask_sequence = []

        for t in range(timesteps):
            img_path = os.path.join(image_folder, image_files[i + t])
            mask_path = os.path.join(mask_folder, mask_files[i + t])
            
            img = load_img(img_path, target_size=(128, 128))
            img = img_to_array(img) / 255.0
            img_sequence.append(img)
            
            mask = load_img(mask_path, target_size=(128, 128), color_mode='grayscale')
            mask = img_to_array(mask) / 255.0
            mask_sequence.append(mask)
        
        images.append(np.array(img_sequence))
        # 对于掩膜，我们只需要最后一帧的掩膜
        masks.append(mask_sequence[-1])
        
    return np.array(images), np.array(masks)

# 定义时间步长
timesteps = 3

# 调用函数加载数据
train_images, train_masks = load_sequence_data(
    r"C:\Users\jack8\OneDrive\113-1_school and company\main file\data\Mix train\train\img",
    r"C:\Users\jack8\OneDrive\113-1_school and company\main file\data\Mix train\train\mask",
    timesteps
)

print("训练数据形状：", train_images.shape)  # (num_samples, 3, 128, 128, 3)
print("训练掩膜形状：", train_masks.shape)  # (num_samples, 128, 128, 1)

# 数据集划分为训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(
    train_images, train_masks, test_size=0.2, random_state=42
)

# 数据增强参数
data_gen_args = dict(
    rotation_range=10,
    width_shift_range=0.1,
    height_shift_range=0.1,
    shear_range=0.1,
    zoom_range=0.1,
    horizontal_flip=True,
    fill_mode='nearest'
)

image_datagen = ImageDataGenerator(**data_gen_args)
mask_datagen = ImageDataGenerator(**data_gen_args)

# 不需要对整个序列进行 fit，只对单帧图像进行 fit
# 重塑 X_train 和 y_train 为单帧图像
X_train_reshaped = X_train.reshape(-1, 128, 128, 3)
y_train_reshaped = y_train.reshape(-1, 128, 128, 1)

image_datagen.fit(X_train_reshaped)
mask_datagen.fit(y_train_reshaped)

# 自定义生成器，确保相同的随机变换应用于所有时间步
def create_sequence_generator(X, y, batch_size, image_datagen, mask_datagen):
    while True:
        # 打乱数据
        indices = np.arange(len(X))
        np.random.shuffle(indices)
        X = X[indices]
        y = y[indices]
        
        for start in range(0, len(X), batch_size):
            end = min(start + batch_size, len(X))
            batch_X = X[start:end]
            batch_y = y[start:end]
            
            augmented_X = []
            augmented_y = []
            
            for seq, mask in zip(batch_X, batch_y):
                # 随机选择一个变换
                seed = random.randint(0, 10000)
                
                # 对序列中的每一帧应用相同的变换
                augmented_seq = []
                for img in seq:
                    augmented_img = image_datagen.random_transform(img, seed=seed)
                    augmented_img = image_datagen.standardize(augmented_img)
                    augmented_seq.append(augmented_img)
                augmented_seq = np.array(augmented_seq)
                augmented_X.append(augmented_seq)
                
                # 对掩膜应用相同的变换
                augmented_mask = mask_datagen.random_transform(mask, seed=seed)
                augmented_mask = mask_datagen.standardize(augmented_mask)
                augmented_y.append(augmented_mask)
            
            augmented_X = np.array(augmented_X)
            augmented_y = np.array(augmented_y)
            
            yield (augmented_X, augmented_y)

# 创建生成器
batch_size = 8
train_generator = create_sequence_generator(
    X_train, y_train, batch_size, image_datagen, mask_datagen
)

steps_per_epoch = len(X_train) // batch_size

# 定义回调函数
callbacks = [
    EarlyStopping(monitor='val_loss', patience=5, verbose=1, restore_best_weights=True),
    ModelCheckpoint(
        filepath=r"C:\Users\jack8\OneDrive\113-1_school and company\main file\Model\CNN-LSTM_model.h5",
        monitor='val_loss',
        save_best_only=True,
        verbose=1
    )
]

# 训练模型
history = model.fit(
    train_generator,
    steps_per_epoch=steps_per_epoch,
    epochs=10,  # 设置为10个epoch
    validation_data=(X_val, y_val),
    callbacks=callbacks
)


import numpy as np
import cv2
import matplotlib.pyplot as plt
from sklearn.metrics import (
    mean_absolute_error,
    mean_squared_error,
    precision_score,
    recall_score,
    accuracy_score,
    f1_score,
    roc_curve,
    auc
)
from skimage.transform import resize
from keras.models import load_model
from keras_preprocessing.image import img_to_array, load_img

# 1. Load your LSTM model
model = load_model(r"C:\Users\jack8\OneDrive\113-1_school and company\main file\Model\CNN-LSTM_model.h5")

# 2. Define helper functions with overlap and blending
def image_to_patches(image, patch_size, overlap):
    patches = []
    positions = []
    img_height, img_width = image.shape[:2]
    patch_height, patch_width = patch_size
    step_y = patch_height - overlap
    step_x = patch_width - overlap
    for y in range(0, img_height - overlap, step_y):
        for x in range(0, img_width - overlap, step_x):
            patch = image[y:y+patch_height, x:x+patch_width]
            # If the patch is smaller than the patch size, pad it
            if patch.shape[0] != patch_height or patch.shape[1] != patch_width:
                pad_height = patch_height - patch.shape[0]
                pad_width = patch_width - patch.shape[1]
                patch = np.pad(
                    patch,
                    ((0, pad_height), (0, pad_width), (0, 0)),
                    mode='constant'
                )
            patches.append(patch)
            positions.append((y, x))
    return np.array(patches), positions

def patches_to_image(patches, positions, image_shape, overlap):
    """
    Reconstruct the original image from patches with overlapping and blending.

    Args:
        patches (numpy.ndarray): Array of patches with shape (num_patches, patch_height, patch_width, channels).
        positions (list): List of (y, x) positions where each patch is placed.
        image_shape (tuple): Shape of the original image (height, width, channels).
        overlap (int): Overlap size between patches.

    Returns:
        numpy.ndarray: Reconstructed image with shape `image_shape`.
    """
    reconstructed_image = np.zeros(image_shape, dtype=np.float32)  # Shape: (H, W, C)
    weight_mask = np.zeros(image_shape, dtype=np.float32)        # Shape: (H, W, C)
    patch_height, patch_width = patches.shape[1:3]

    for patch, (y, x) in zip(patches, positions):
        y_end = min(y + patch_height, image_shape[0])
        x_end = min(x + patch_width, image_shape[1])

        h = y_end - y
        w = x_end - x

        # 确保 patch 的维度为 (h, w, C)
        patch_crop = patch[:h, :w, 0]  # 取第一个通道，形状: (h, w)
        patch_crop = patch_crop[..., np.newaxis]  # 扩展为 (h, w, 1)

        # 将 patch 添加到重建图像
        reconstructed_image[y:y_end, x:x_end, 0] += patch_crop[:, :, 0]

        # 更新权重掩膜
        weight_mask[y:y_end, x:x_end, 0] += 1

    # 避免除以零
    weight_mask[weight_mask == 0] = 1

    # 平均重叠区域
    reconstructed_image /= weight_mask

    # 如果只处理单通道，可以移除多余的维度
    return reconstructed_image

# 3. Load and preprocess the input image
input_img_path = r"C:\Users\jack8\OneDrive\113-1_school and company\main file\data\north\Finish\segmentation-vidas\2024.png"
new_img = load_img(input_img_path)
new_img_array = img_to_array(new_img) / 255.0

# Get original image dimensions
original_height, original_width = new_img_array.shape[:2]

# Define model input dimensions
patch_size = (128, 128)  # Adjust based on your model's input size
overlap = 32  # You can adjust this value based on your needs

# Extract overlapping patches
patches, positions = image_to_patches(new_img_array, patch_size, overlap)

# 4. Predict on each patch
timesteps = 3  # Ensure this matches your model's expected timesteps

predicted_patches = []
for i, patch in enumerate(patches):
    # Create a time sequence by repeating the same patch
    patch_sequence = np.array([patch] * timesteps)
    # Expand dimensions to match the model input shape
    patch_input = np.expand_dims(patch_sequence, axis=0)  # Shape: (1, timesteps, height, width, channels)
    # Predict
    prediction = model.predict(patch_input)
    predicted_patches.append(prediction[0])

predicted_patches = np.array(predicted_patches)

# Reconstruct the image with blending
prediction_mask = patches_to_image(predicted_patches, positions, (original_height, original_width, 1), overlap)

# 5. Load and preprocess the ground truth mask
mask_img_path = r"C:\Users\jack8\OneDrive\113-1_school and company\main file\data\north\Finish\cover\2024.png"
mask_img = load_img(mask_img_path, color_mode='grayscale')
mask_img_array = img_to_array(mask_img) / 255.0
mask_img_array = mask_img_array[:, :, 0]  # Keep only one channel

# Ensure the ground truth mask size matches the prediction mask
if mask_img_array.shape != prediction_mask.shape[:2]:
    mask_img_array_resized = resize(mask_img_array, prediction_mask.shape[:2], preserve_range=True)
    mask_img_array_resized = mask_img_array_resized.astype(np.float32)  # Ensure float type for consistency
else:
    mask_img_array_resized = mask_img_array

# Threshold the prediction mask to get binary mask
prediction_mask_binary = (prediction_mask > 0.3).astype(np.uint8)

# Apply morphological closing to the prediction mask
kernel_size = 20
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
closed_mask = cv2.morphologyEx(prediction_mask_binary[:, :, 0], cv2.MORPH_CLOSE, kernel)

# Convert to binary mask
closed_mask_binary = (closed_mask > 0).astype(np.uint8)

# Flatten arrays for evaluation
true_values = mask_img_array_resized.flatten()
predicted_values = closed_mask_binary.flatten()

# Binarize the true values
true_values = (true_values > 0.5).astype(np.float32)

# ----------------------- 添加 ROC 曲线的绘制部分 -----------------------

# 要绘制 ROC 曲线，我们需要真实标签和模型预测的概率分数
# 真实标签 y_true 应该是二进制的 (0 和 1)
# 预测分数 y_scores 应该是连续的概率值

# 确保 mask_img_array_resized 是二值的
y_true = (mask_img_array_resized > 0.5).astype(int).flatten()  # 使用 int 替代 np.int

# 获取预测的概率分数
# prediction_mask 是通过 patches_to_image 重建的概率掩膜
# 确保 prediction_mask 的形状正确
prediction_prob_mask = prediction_mask[:, :, 0].flatten()

# 检查 y_true 是否仅包含 0 和 1
unique_values = np.unique(y_true)
print("Unique values in y_true:", unique_values)

# 如果 y_true 中有非二进制值，强制二值化
if not set(unique_values).issubset({0, 1}):
    print("Warning: y_true contains non-binary values. Binarizing y_true.")
    y_true = (y_true > 0.5).astype(int)

# 计算 ROC 曲线和 AUC
fpr, tpr, thresholds = roc_curve(y_true, prediction_prob_mask)
roc_auc = auc(fpr, tpr)
print(f"ROC AUC: {roc_auc:.4f}")

# 绘制 ROC 曲线
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (AUC = {roc_auc:.4f})')
plt.plot([0, 1], [0, 1], color='red', lw=2, linestyle='--', label='Random Guess')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('False Positive Rate', fontsize=14)
plt.ylabel('True Positive Rate', fontsize=14)
plt.title('Receiver Operating Characteristic (ROC) Curve', fontsize=16)
plt.legend(loc="lower right", fontsize=12)
plt.grid(True)
plt.show()

# ----------------------- 评估指标计算部分 -----------------------

# Mean Absolute Percentage Error
def mean_absolute_percentage_error(y_true, y_pred):
    non_zero_indices = y_true != 0
    y_true_non_zero = y_true[non_zero_indices]
    y_pred_non_zero = y_pred[non_zero_indices]
    if len(y_true_non_zero) == 0:
        print("Warning: No non-zero elements in y_true. MAPE is undefined.")
        return np.nan
    else:
        mape = np.mean(np.abs((y_true_non_zero - y_pred_non_zero) / y_true_non_zero)) * 100
        return mape

mape = mean_absolute_percentage_error(true_values, predicted_values)
print(f"MAPE: {mape}%")

# Mean Absolute Error
mae = mean_absolute_error(true_values, predicted_values)
print(f"MAE: {mae}")

# Root Mean Squared Error
rmse = np.sqrt(mean_squared_error(true_values, predicted_values))
print(f"RMSE: {rmse}")

# Intersection over Union (IoU)
intersection = np.logical_and(true_values, predicted_values).sum()
union = np.logical_or(true_values, predicted_values).sum()
iou_score = intersection / union if union != 0 else 1
print(f"IoU (Jaccard Index): {iou_score}")

# Accuracy
accuracy = accuracy_score(true_values, predicted_values)
print(f"Accuracy: {accuracy}")

# F1 Score
f1 = f1_score(true_values, predicted_values, zero_division=0)
print(f"F1 Score: {f1}")

# Dice Coefficient
dice_score = 2 * intersection / (true_values.sum() + predicted_values.sum()) if (true_values.sum() + predicted_values.sum()) != 0 else 1
print(f"Dice Coefficient: {dice_score}")

# Precision and Recall
precision = precision_score(true_values, predicted_values, zero_division=0)
recall = recall_score(true_values, predicted_values, zero_division=0)
print(f"Precision: {precision}")
print(f"Recall: {recall}")

# ----------------------- 结果可视化部分 -----------------------

# Visualize results
plt.figure(figsize=(18, 6))

plt.subplot(1, 3, 1)
plt.title('Ground Truth Mask')
plt.imshow(mask_img_array_resized, cmap='gray')
plt.axis('off')

plt.subplot(1, 3, 2)
plt.title('Prediction Mask')
plt.imshow(prediction_mask[:, :, 0], cmap='gray')
plt.axis('off')

plt.subplot(1, 3, 3)
plt.title('Closed Prediction Mask (Binary)')
plt.imshow(closed_mask_binary, cmap='gray')
plt.axis('off')

plt.show()

###幫我改成只需要cnn就好