import numpy as np
import torch
import random

# 正规包导入 - 使用已安装的Python包
from mfdnn import MFDNN, MFDNN_predict, SimpsonIntegral2D

# 设置随机种子
torch.manual_seed(756)
np.random.seed(756)
random.seed(756)

# 定义函数 (保持不变)
def B1(s, t):
    return np.sin(s) * np.sin(t)

def B2(s, t):
    return np.cos(s) * np.cos(t)

def B3(s, t):
    return s + t

def B4(s, t):
    return np.sin(s * t)

def B5(s, t):
    return np.cos(s * t)

def B6(s, t):
    return s * t

def x1(s, t):
    return np.random.normal(1) + np.random.normal(1) * np.sin(s) + np.random.normal(1) * np.sin(t)

def x2(s, t):
    return np.random.normal(1) + np.random.normal(1) * np.cos(s) + np.random.normal(1) * np.cos(t)

def x3(s, t):
    return np.random.normal(1) * np.square(s) + np.random.normal(1) * np.square(t) + np.random.normal(1) * s * t

def x4(s, t):
    return np.random.normal(1) * s * np.sin(t)

def x5(s, t):
    return np.random.normal(1) * s * np.cos(t)

def x6(s, t):
    return np.random.normal(1) * np.exp(- s * t)

def generate_single_dataset(T=32, n=400):
    """生成单个数据集"""
    t1 = np.linspace(0, 1, T)
    t2 = np.linspace(0, 1, T)
    t1_grid, t2_grid = np.meshgrid(t1, t2)
    grid = np.vstack((t1_grid.ravel(), t2_grid.ravel())).T
    
    X_array = np.zeros((6, 2 * n, T, T))
    mX1 = []
    
    for k in range(2 * n):
        X1 = np.array([[x1(s, t) for s, t in grid]]).reshape((T, T)).T
        X2 = np.array([[x2(s, t) for s, t in grid]]).reshape((T, T)).T
        X3 = np.array([[x3(s, t) for s, t in grid]]).reshape((T, T)).T
        X4 = np.array([[x4(s, t) for s, t in grid]]).reshape((T, T)).T
        X5 = np.array([[x5(s, t) for s, t in grid]]).reshape((T, T)).T
        X6 = np.array([[x6(s, t) for s, t in grid]]).reshape((T, T)).T

        B1_grid = np.array([[B1(s, t) for s, t in grid]]).reshape((T, T)).T
        B2_grid = np.array([[B2(s, t) for s, t in grid]]).reshape((T, T)).T

        mX1.append(SimpsonIntegral2D(X1 * B1_grid, (0, 0), (1, 1)) + 
                   SimpsonIntegral2D(X2 * B2_grid, (0, 0), (1, 1)))

        X_array[0, k, :, :] = X1
        X_array[1, k, :, :] = X2
        X_array[2, k, :, :] = X3
        X_array[3, k, :, :] = X4
        X_array[4, k, :, :] = X5
        X_array[5, k, :, :] = X6
        
    mX1 = np.array(mX1)  
    y1 = mX1 + np.random.normal(0, 0.3 * np.std(mX1), 2 * n)
    
    return X_array, y1

def calculate_selection_metrics(l21_norms, true_vars, epsilon=0.01, p=6):
    """计算变量选择指标"""
    selected_vars = set(i for i, norm in enumerate(l21_norms) if norm > epsilon)
    
    true_positive = len(selected_vars & true_vars)
    false_positive = len(selected_vars - true_vars)
    false_negative = len(true_vars - selected_vars)
    
    f1_score = (2 * true_positive) / (2 * true_positive + false_positive + false_negative) if (2 * true_positive + false_positive + false_negative) > 0 else 0
    perfect_selection = 1.0 if selected_vars == true_vars else 0.0
    
    return f1_score, perfect_selection, selected_vars

def select_best_hyperparameters(X_train, y_train, true_vars, p, domain_range, lam1_values, lam2_values, model_params, epsilon=0.01):
    """选择最佳超参数"""
    mse_results = np.zeros((len(lam1_values), len(lam2_values)))
    f1_results = np.zeros((len(lam1_values), len(lam2_values)))
    selection_info = {}
    
    y_train_mean = np.mean(y_train)
    y_train_std = np.std(y_train)
    
    for i, lam1 in enumerate(lam1_values):
        for j, lam2 in enumerate(lam2_values):
            try:
                train_losses, val_losses, model, l21 = MFDNN(
                    p=p, resp=y_train, func_cov=X_train,
                    num_basis=model_params['num_basis'],
                    layer_sizes=model_params['layer_sizes'],
                    domain_range=domain_range,
                    epochs=model_params['epochs'],
                    val_ratio=model_params['val_ratio'],
                    patience=model_params['patience'],
                    lam1=lam1, lam2=lam2, std_resp=True
                )
                
                mse_results[i, j] = min(val_losses) if len(val_losses) > 0 else np.mean(train_losses[-10:])
                
                f1_score, perfect_selection, selected_vars = calculate_selection_metrics(l21, true_vars, epsilon, p)
                f1_results[i, j] = f1_score
                
                selection_info[f"{i}_{j}"] = {
                    'model': model, 'lam1': lam1, 'lam2': lam2,
                    'f1_score': f1_score, 'mse': mse_results[i, j],
                    'selected_vars': list(selected_vars),
                    'y_mean': y_train_mean, 'y_std': y_train_std,
                    'perfect_selection': perfect_selection
                }
                
            except Exception as e:
                mse_results[i, j] = np.inf
                f1_results[i, j] = 0
    
    best_f1 = np.max(f1_results)
    best_f1_indices = np.where(f1_results == best_f1)
    
    if len(best_f1_indices[0]) > 0:
        best_candidates = [selection_info[f"{i}_{j}"] for i, j in zip(best_f1_indices[0], best_f1_indices[1])]
        best_candidate = min(best_candidates, key=lambda x: x['mse'])
    else:
        default_lam1 = 1
        default_lam2 = 0.1
        print("使用默认超参数: lam1=1, lam2=0.1")
        best_candidate = {
            'model': None, 'lam1': default_lam1, 'lam2': default_lam2,
            'f1_score': 0, 'mse': np.inf, 'selected_vars': [],
            'y_mean': y_train_mean, 'y_std': y_train_std, 'perfect_selection': 0
        }
    
    return best_candidate['lam1'], best_candidate['lam2'], best_candidate

def evaluate_on_test_set(best_candidate, X_test, y_test, p, domain_range, model_params):
    """在测试集上评估性能 - 修改为RMSE"""
    try:
        y_mean = best_candidate['y_mean']
        y_std = best_candidate['y_std']
        
        test_predictions_normalized = MFDNN_predict(p, best_candidate['model'], X_test, model_params['num_basis'], domain_range)
        test_predictions_original = test_predictions_normalized.detach().numpy() * y_std + y_mean
        
        test_mse = np.mean((test_predictions_original.flatten() - y_test) ** 2)
        test_rmse = np.sqrt(test_mse)  # 只返回RMSE，不计算NRMSE
        
        return test_rmse, best_candidate['f1_score']
    except Exception as e:
        return np.inf, best_candidate['f1_score']

# 主程序
if __name__ == "__main__":
    print("生成测试数据...")
    X, y = generate_single_dataset(T=32, n=400)
    print(f"数据生成完成: X shape {X.shape}, y shape {y.shape}")
    
    # 超参数
    lam1_values = [0.5, 1, 1.5, 2, 2.5, 3]
    lam2_values = [0, 0.001, 0.01, 0.1, 0.5, 1]
    
    model_params = {
        'num_basis': (5, 5),
        'layer_sizes': [64, 64],
        'epochs': 100,
        'val_ratio': 0.25,
        'patience': 10
    }
    
    ground_truth = {0: {0, 1}}  # y1的真实变量
    epsilon = 0.01
    
    # 数据分割
    p, N, T1, T2 = X.shape
    split_idx = N // 2
    X_train = X[:, :split_idx, :, :]
    X_test = X[:, split_idx:, :, :]
    y_train = y[:split_idx]
    y_test = y[split_idx:]
    
    domain_range = [[[0, 0], [1, 1]] for _ in range(p)]
    true_vars = ground_truth[0]
    
    print(f"\n训练数据: X_train shape {X_train.shape}, y_train shape {y_train.shape}")
    print(f"测试数据: X_test shape {X_test.shape}, y_test shape {y_test.shape}")
    print(f"真实显著变量: {true_vars}")
    
    print("\n开始超参数选择...")
    lam1, lam2, best_candidate = select_best_hyperparameters(
        X_train, y_train, true_vars, p, domain_range, 
        lam1_values, lam2_values, model_params, epsilon
    )
    
    print(f"最佳超参数: lam1={lam1}, lam2={lam2}")
    print(f"训练F1分数: {best_candidate['f1_score']:.4f}")
    print(f"选择变量: {set(best_candidate['selected_vars'])}")
    
    print("\n在测试集上评估...")
    test_rmse, test_f1 = evaluate_on_test_set(
        best_candidate, X_test, y_test, p, domain_range, model_params
    )
    
    print(f"测试结果:")
    print(f"  RMSE: {test_rmse:.4f}")  # 修改为RMSE
    print(f"  F1分数: {test_f1:.4f}")
    print(f"  完美匹配: {'是' if set(best_candidate['selected_vars']) == true_vars else '否'}")