# %% [code]
# %% [code]
# %% [code]
# 1. 环境配置与库导入（第一步：基础准备）
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
from sklearn.preprocessing import StandardScaler, OneHotEncoder, RobustScaler
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split
from sklearn.impute import SimpleImputer
import tensorflow as tf
from tensorflow.keras import layers, models, optimizers
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import warnings
warnings.filterwarnings('ignore')

# 固定随机种子（保证结果可复现）
SEED = 42
np.random.seed(SEED)
tf.random.set_seed(SEED)

# 关键参数配置（核心修改：测试图像目录改为./test）
TRAIN_IMG_DIR = "/kaggle/input/csiro-biomass/train"
TEST_IMG_DIR = "/kaggle/input/csiro-biomass/test"
IMG_EXT = [".jpg", ".png"]
IMG_SIZE = (512, 256)
BATCH_SIZE = 16
EPOCHS = 20

# 适配train.csv的target_name列实际值
TARGET_COLS = ['Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g', 'Dry_Total_g', 'GDM_g']
NUM_TARGETS = len(TARGET_COLS)


# 2. 数据探索
print("="*50)
print("2. 数据探索开始")
print("="*50)

# 加载数据
train_df = pd.read_csv('/kaggle/input/csiro-biomass/train.csv')
test_df = pd.read_csv('/kaggle/input/csiro-biomass/test.csv')

# 2.1 基础信息分析
print(f"【训练集基础信息】")
print(f"  行数：{train_df.shape[0]}, 列数：{train_df.shape[1]}")
print(f"  列名：{train_df.columns.tolist()}")
print(f"\n【测试集基础信息】")
print(f"  行数：{test_df.shape[0]}, 列数：{test_df.shape[1]}")
print(f"  列名：{test_df.columns.tolist()}")

# 2.2 图像ID列确认
IMG_ID_COL = 'sample_id'
print(f"\n【图像ID确认】")
print(f"  ID列名：{IMG_ID_COL}")
print(f"  训练集唯一图像数：{train_df[IMG_ID_COL].nunique()}")
print(f"  测试集唯一图像数：{test_df[IMG_ID_COL].nunique()}")
print(f"  训练集ID前5个：{train_df[IMG_ID_COL].head().tolist()}")
print(f"  测试集ID前5个：{test_df[IMG_ID_COL].head().tolist()}")

# 2.3 目标列确认与数据清洗
train_wide = train_df.pivot(
    index=IMG_ID_COL,
    columns='target_name',
    values='target'
).reset_index()
train_wide.columns.name = None

# 合并原表中的特征列（同一图像的特征值去重后合并）
feature_cols = [col for col in train_df.columns if col not in ['target_name', 'target']]
train_wide = train_wide.merge(
    train_df[feature_cols].drop_duplicates(subset=IMG_ID_COL),
    on=IMG_ID_COL,
    how='left'
)

# 目标值清洗（核心：填充NaN+裁剪极值）
print(f"\n【目标值清洗】")
# 填充NaN（用中位数）
train_wide[TARGET_COLS] = train_wide[TARGET_COLS].fillna(train_wide[TARGET_COLS].median())
# 裁剪极值（0~1000，可根据实际业务调整）
train_wide[TARGET_COLS] = train_wide[TARGET_COLS].clip(lower=0, upper=1000)
# 检查清洗后结果
print(f"  清洗后目标列NaN数量：{train_wide[TARGET_COLS].isnull().sum().to_dict()}")
print(f"  清洗后目标列范围：")
for col in TARGET_COLS:
    print(f"    {col}: [{train_wide[col].min():.2f}, {train_wide[col].max():.2f}]")

# 2.4 特征分类（基于原始train_df，且对齐test_df列名）
train_candidate_features = [col for col in train_df.columns if col not in [IMG_ID_COL, 'target_name', 'target']]
common_features = [col for col in train_candidate_features if col in test_df.columns]
if not common_features:
    print("警告：训练集和测试集无共同特征列，将仅使用图像特征训练")

# 自动识别类别/数值特征
CAT_FEATURES = [col for col in common_features if train_df[col].dtype == 'object']
NUM_FEATURES = [col for col in common_features if train_df[col].dtype in ['float64', 'int64']]

print(f"\n【特征分类】")
print(f"  类别特征（需编码）：{CAT_FEATURES}")
print(f"  数值特征（需标准化）：{NUM_FEATURES}")

# 2.5 缺失值检查
print(f"\n【特征缺失值检查】")
if common_features:
    feat_missing = train_wide[common_features].isnull().sum()
    print(f"  特征NaN数量：{feat_missing[feat_missing>0].to_dict()}")

print("="*50)
print("2. 数据探索结束")
print("="*50)


# 3. 数据预处理
print("\n" + "="*50)
print("3. 数据预处理开始")
print("="*50)

# 3.1 数据划分
train_unique_ids = train_wide[IMG_ID_COL].unique()
train_ids, val_ids = train_test_split(train_unique_ids, test_size=0.2, random_state=SEED)
train_subset = train_wide[train_wide[IMG_ID_COL].isin(train_ids)].reset_index(drop=True)
val_subset = train_wide[train_wide[IMG_ID_COL].isin(val_ids)].reset_index(drop=True)

print(f"【数据划分结果】")
print(f"  训练集：{len(train_ids)}张图 → {len(train_subset)}行样本")
print(f"  验证集：{len(val_ids)}张图 → {len(val_subset)}行样本")

# 3.2 特征预处理流水线（核心：用RobustScaler提升稳定性）
transformers = []
# 数值特征：填充中位数 + RobustScaler（对极值鲁棒）
if NUM_FEATURES:
    num_preprocessor = Pipeline(steps=[
        ('imputer', SimpleImputer(strategy='median')),
        ('scaler', RobustScaler())
    ])
    transformers.append(('num', num_preprocessor, NUM_FEATURES))

# 类别特征：填充众数 + OneHot编码
if CAT_FEATURES:
    cat_preprocessor = Pipeline(steps=[
        ('imputer', SimpleImputer(strategy='most_frequent')),
        ('onehot', OneHotEncoder(handle_unknown='ignore', sparse_output=False))
    ])
    transformers.append(('cat', cat_preprocessor, CAT_FEATURES))

preprocessor = ColumnTransformer(transformers=transformers, remainder='drop')

# 3.3 拟合预处理流水线（仅用训练集）
print(f"\n【特征预处理】")
X_train_processed = np.zeros((len(train_subset), 0))
X_val_processed = np.zeros((len(val_subset), 0))
X_test_processed = np.zeros((len(test_df), 0))

if transformers:
    X_train_num = train_subset[common_features]
    X_train_processed = preprocessor.fit_transform(X_train_num)
    X_val_num = val_subset[common_features]
    X_val_processed = preprocessor.transform(X_val_num)
    X_test_num = test_df[common_features]
    X_test_processed = preprocessor.transform(X_test_num)

# 特征异常值清洗（替换NaN/无穷大）
X_train_processed = np.nan_to_num(X_train_processed, nan=0.0, posinf=1e6, neginf=-1e6)
X_val_processed = np.nan_to_num(X_val_processed, nan=0.0, posinf=1e6, neginf=-1e6)
X_test_processed = np.nan_to_num(X_test_processed, nan=0.0, posinf=1e6, neginf=-1e6)

# 提取目标值
y_train = train_subset[TARGET_COLS].values
y_val = val_subset[TARGET_COLS].values

# 记录特征维度
FEATURE_DIM = X_train_processed.shape[1] if X_train_processed.size > 0 else 32
print(f"  处理后训练集特征维度：{X_train_processed.shape}")
print(f"  处理后验证集特征维度：{X_val_processed.shape}")
print(f"  处理后测试集特征维度：{X_test_processed.shape}")
print(f"  特征维度（用于模型）：{FEATURE_DIM}")

# 3.4 图像加载函数（核心：支持训练/测试不同目录+自动检测后缀+容错）
def load_image(img_id, is_test=False):
    """
    根据图像ID加载并预处理图像
    参数：
        img_id: 图像ID（如ID1001187975）
        is_test: 是否为测试集图像（决定读取目录）
    返回：
        预处理后的图像数组（512,256,3）
    """
    # 1. 确定图像目录
    img_dir = TEST_IMG_DIR if is_test else TRAIN_IMG_DIR
    
    # 2. 提取纯图像ID（去掉__分隔的目标名）
    pure_img_id = img_id.split('__')[0] if '__' in img_id else img_id
    
    # 3. 自动检测图像后缀（.jpg/.png）
    img_path = None
    for ext in IMG_EXT:
        temp_path = os.path.join(img_dir, f"{pure_img_id}{ext}")
        if os.path.exists(temp_path):
            img_path = temp_path
            break
    
    # 4. 容错处理：文件不存在时返回默认图像（避免程序崩溃）
    if img_path is None:
        print(f"⚠️ 警告：图像 {pure_img_id} 在 {img_dir} 中不存在，返回全0图像")
        return np.zeros(IMG_SIZE + (3,), dtype=np.float32)
    
    # 5. 读取并预处理图像
    img = cv2.imread(img_path)
    if img is None:
        print(f"⚠️ 警告：图像 {img_path} 损坏，返回全0图像")
        return np.zeros(IMG_SIZE + (3,), dtype=np.float32)
    
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (IMG_SIZE[1], IMG_SIZE[0]))
    img = img / 255.0  # 归一化到0-1
    return img.astype(np.float32)

# 3.5 图像数据增强（仅训练集）
train_img_datagen = ImageDataGenerator(
    horizontal_flip=True,
    brightness_range=[0.8, 1.2],
    zoom_range=0.1,
    fill_mode='nearest'
)
val_test_img_datagen = ImageDataGenerator()

# 3.6 数据生成器函数
def data_generator_fn(df, processed_features, img_datagen, is_train=True, is_test=False):
    """
    数据生成器：返回({"image_input": 图像, "numeric_input": 特征}, 标签)
    """
    unique_ids = df[IMG_ID_COL].unique()
    id_to_idx = {img_id: idx for idx, img_id in enumerate(df[IMG_ID_COL].values)}
    
    while True:
        if is_train:
            np.random.shuffle(unique_ids)
        for img_id in unique_ids:
            img = load_image(img_id, is_test=is_test)
            if is_train:
                img = img_datagen.random_transform(img)
            
            if processed_features.size > 0:
                feat = processed_features[id_to_idx[img_id]].astype(np.float32)
            else:
                feat = np.zeros((FEATURE_DIM,), dtype=np.float32)
            
            if is_test:
                target = np.zeros((NUM_TARGETS,), dtype=np.float32)
            else:
                target = df[df[IMG_ID_COL] == img_id][TARGET_COLS].values[0].astype(np.float32)
            
            yield ({"image_input": img, "numeric_input": feat}, target)

# 3.7 创建tf.data.Dataset
def create_tf_dataset(df, processed_features, img_datagen, batch_size, is_train=True, is_test=False):
    """创建TensorFlow Dataset，指定输出签名"""
    output_signature = (
        {
            "image_input": tf.TensorSpec(shape=IMG_SIZE + (3,), dtype=tf.float32),
            "numeric_input": tf.TensorSpec(shape=(FEATURE_DIM,), dtype=tf.float32)
        },
        tf.TensorSpec(shape=(NUM_TARGETS,), dtype=tf.float32)
    )
    
    dataset = tf.data.Dataset.from_generator(
        lambda: data_generator_fn(df, processed_features, img_datagen, is_train, is_test),
        output_signature=output_signature
    )
    
    if is_train:
        dataset = dataset.shuffle(buffer_size=batch_size*10)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    return dataset

# 3.8 构建训练/验证Dataset
print(f"\n【创建Dataset】")
train_dataset = create_tf_dataset(
    df=train_subset,
    processed_features=X_train_processed,
    img_datagen=train_img_datagen,
    batch_size=BATCH_SIZE,
    is_train=True,
    is_test=False
)
val_dataset = create_tf_dataset(
    df=val_subset,
    processed_features=X_val_processed,
    img_datagen=val_test_img_datagen,
    batch_size=BATCH_SIZE,
    is_train=False,
    is_test=False
)

# 计算步数
train_steps = len(train_ids) // BATCH_SIZE
val_steps = len(val_ids) // BATCH_SIZE
print(f"  训练集步数：{train_steps}")
print(f"  验证集步数：{val_steps}")

print("="*50)
print("3. 数据预处理结束")
print("="*50)


# 4. 模型构建与训练（第四步：多模态模型）
print("\n" + "="*50)
print("4. 模型构建与训练开始")
print("="*50)

# 4.1 构建图像分支（MobileNetV2）
def build_image_branch(input_shape):
    inputs = layers.Input(shape=input_shape, name="image_input")
    x = layers.Rescaling(scale=1./127.5, offset=-1)(inputs)
    base_model = tf.keras.applications.MobileNetV2(
        input_shape=input_shape,
        include_top=False,
        weights=None
    )
    base_model.load_weights(
    "/kaggle/input/mobilenetv2/tensorflow2/tf2_mobilenetv2_imagebase/1/mobilenet_v2_weights_tf_dim_ordering_tf_kernels_1.0_224_no_top.h5"
    )

    base_model.trainable = True
    for layer in base_model.layers[:-20]:
        layer.trainable = False
    
    x = base_model(x)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(128, activation='relu')(x)
    x = layers.Dropout(0.3)(x)
    outputs = layers.Dense(64, activation='relu')(x)
    return models.Model(inputs, outputs, name='image_branch')

# 4.2 构建数值特征分支
def build_numeric_branch(input_dim):
    inputs = layers.Input(shape=(input_dim,), name="numeric_input")
    if input_dim == 0:
        x = layers.Lambda(lambda x: tf.zeros((tf.shape(x)[0], 32)))(inputs)
    else:
        x = layers.Dense(64, activation='relu')(inputs)
        x = layers.Dropout(0.2)(x)
        x = layers.Dense(32, activation='relu')(x)
    outputs = layers.Dense(32, activation='relu')(x)
    return models.Model(inputs, outputs, name='numeric_branch')

# 4.3 融合多模态模型
def build_combined_model(image_shape, numeric_dim):
    image_branch = build_image_branch(image_shape)
    numeric_branch = build_numeric_branch(numeric_dim)
    
    combined = layers.concatenate([image_branch.output, numeric_branch.output])
    x = layers.Dense(128, activation='relu')(combined)
    x = layers.Dropout(0.3)(x)
    outputs = layers.Dense(NUM_TARGETS, activation='linear')(x)
    
    model = models.Model(
        inputs=[image_branch.input, numeric_branch.input],
        outputs=outputs,
        name='biomass_combined_model'
    )
    return model

# 4.4 初始化并编译模型（降低学习率提升稳定性）
model = build_combined_model(image_shape=IMG_SIZE + (3,), numeric_dim=FEATURE_DIM)

# 稳定版加权MSE损失
def weighted_mse_loss(y_true, y_pred):
    weights = tf.constant([0.1, 0.1, 0.1, 0.5, 0.2], dtype=tf.float32)
    y_true = tf.clip_by_value(y_true, 1e-6, 1e12)
    y_pred = tf.clip_by_value(y_pred, 1e-6, 1e12)
    return tf.reduce_mean(weights * tf.square(y_true - y_pred))

model.compile(
    optimizer=optimizers.AdamW(learning_rate=1e-2, weight_decay=1e-5),
    loss=weighted_mse_loss,
    metrics=[tf.keras.metrics.R2Score(name='r2')]
)

# 打印模型结构
print(f"\n【模型结构】")
model.summary()

# 4.5 训练回调函数
callbacks = [
    EarlyStopping(monitor='val_r2', patience=5, mode='max', restore_best_weights=True, verbose=1),
    ReduceLROnPlateau(monitor='val_loss', factor=0.75, patience=3, min_lr=1e-6, verbose=1)
]

# 4.6 开始训练
print(f"\n【模型训练】")
history = model.fit(
    train_dataset,
    validation_data=val_dataset,
    steps_per_epoch=train_steps,
    validation_steps=val_steps,
    epochs=EPOCHS,
    callbacks=callbacks,
    verbose=1
)

# 4.7 训练曲线可视化
print(f"\n【训练曲线保存】")
plt.figure(figsize=(12, 4))

# 损失曲线
plt.subplot(1, 2, 1)
plt.plot(history.history['loss'], label='Train Loss')
plt.plot(history.history['val_loss'], label='Val Loss')
plt.xlabel('Epoch')
plt.ylabel('Weighted MSE')
plt.legend()
plt.title('Loss Curve')

# R²曲线
plt.subplot(1, 2, 2)
plt.plot(history.history['r2'], label='Train R²')
plt.plot(history.history['val_r2'], label='Val R²')
plt.xlabel('Epoch')
plt.ylabel('R² Score')
plt.legend()
plt.title('R² Curve')

# 保存
plt.tight_layout()
#plt.savefig('./training_history.png')
plt.close()
#print(f"  训练曲线已保存至：./training_history.png")

print("="*50)
print("4. 模型构建与训练结束")
print("="*50)


# 5. 测试集预测
print("\n" + "="*50)
print("5. 测试集预测开始")
print("="*50)

# 5.1 预处理测试集：从sample_id中提取【纯图像ID】（核心修复）
def extract_pure_img_id(sample_id_str):
    """从“图像ID__指标名”中提取纯图像ID"""
    return sample_id_str.split('__')[0]

test_df['pure_img_id'] = test_df[IMG_ID_COL].apply(extract_pure_img_id)
pure_img_ids = test_df['pure_img_id'].unique()
print(f"【测试集纯图像ID数量】：{len(pure_img_ids)}")


# 5.2 测试集预测函数（彻底修复格式问题）
def predict_test_set(model, test_df, processed_features, batch_size):
    """测试集预测，生成正确格式的提交结果"""
    predictions = []
    pure_id_to_idx = {}
    for pure_id in pure_img_ids:
        first_row_idx = test_df[test_df['pure_img_id'] == pure_id].index[0]
        pure_id_to_idx[pure_id] = first_row_idx
    
    # 分批处理纯图像ID
    for batch_start in range(0, len(pure_img_ids), batch_size):
        batch_pure_ids = pure_img_ids[batch_start:batch_start+batch_size]
        batch_imgs = []
        batch_feats = []
        
        for pure_id in batch_pure_ids:
            img = load_image(pure_id, is_test=True)
            batch_imgs.append(img)
            
            if processed_features.size > 0:
                feat_idx = pure_id_to_idx[pure_id]
                feat = processed_features[feat_idx].astype(np.float32)
            else:
                feat = np.zeros((FEATURE_DIM,), dtype=np.float32)
            batch_feats.append(feat)
        
        batch_imgs_np = np.array(batch_imgs)
        batch_feats_np = np.array(batch_feats)
        batch_preds = model.predict(
            {"image_input": batch_imgs_np, "numeric_input": batch_feats_np},
            verbose=0
        )
        
        # 4. 整理结果（严格匹配正确格式）
        for idx_in_batch, pure_id in enumerate(batch_pure_ids):
            img_preds = batch_preds[idx_in_batch]
            for target_idx, target_name in enumerate(TARGET_COLS):
                sample_id = f"{pure_id}__{target_name}"
                pred_value = max(float(img_preds[target_idx]), 0.0)
                predictions.append([sample_id, pred_value])
    
    return predictions


# 5.3 执行预测
print(f"【执行预测】")
test_predictions = predict_test_set(
    model=model,
    test_df=test_df,
    processed_features=X_test_processed,
    batch_size=BATCH_SIZE
)


# 5.4 生成提交文件
submission_df = pd.DataFrame(test_predictions, columns=['sample_id', 'target'])


# 5.5 验证提交文件格式（确保完全正确）
print(f"\n【提交文件验证】")
expected_rows = len(pure_img_ids) * len(TARGET_COLS)
actual_rows = len(submission_df)
print(f"  预期行数：{expected_rows}，实际行数：{actual_rows}")
correct_sample_id_example = f"{pure_img_ids[0]}__{TARGET_COLS[0]}"
print(f"  正确sample_id格式示例：{correct_sample_id_example}")
print(f"  实际生成的sample_id示例：{submission_df['sample_id'].iloc[0]}")
print(f"  前5行预测结果：")
print(submission_df.head())


# 5.6 保存提交文件
submission_df.to_csv('submission.csv', index=False)
print(f"\n【提交文件保存】")
print(f"  文件已保存至：./submission.csv")


print("="*50)
print("5. 测试集预测结束")
print("="*50)
print("\n所有步骤执行完成！")