import pandas as pd
from sklearn.linear_model import LassoCV, Lasso
from sklearn.preprocessing import StandardScaler
import json
import sys

# 重新配置标准输出流的编码为 UTF-8
sys.stdout.reconfigure(encoding='utf-8')


def load_data(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        json_str = file.read()

    # 解析 JSON 字符串
    data_dict = json.loads(json_str)
    # 转为 Pandas DataFrame
    data = pd.DataFrame(data_dict)

    # # """加载数据并重命名列"""
    # data = pd.read_csv(file_path)
    new_names = [f'q{i}' for i in range(7, 189)]
    data.columns.values[6:189] = new_names
    return data


def prepare_data(data):
    data['month'] = data['month'].astype(int)
    data['year'] = data['year'].astype(int)
    # """选择所需的数据，计算平均值并合并数据"""
    new_names = [f'q{i}' for i in range(7, 189)]  # 确保此列名称与选择特征匹配
    selected_data = (
        data[(data['year'] == 2021) | (data['year'] == 2022) | ((data['year'] == 2023) & (data['month'] <= 5))]
        .groupby('area_name')[new_names].mean().reset_index()
    )
    all_score_2023_06 = data[(data['year'] == 2023) & (data['month'] == 6)].dropna(subset=['all_score'])[
        ['area_name', 'all_score']]
    model_data = pd.merge(selected_data, all_score_2023_06, on='area_name')
    return model_data


def train_lasso(X, y):
    """训练LASSO回归模型并返回最佳模型和预测值"""
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    lasso_cv = LassoCV(cv=16, random_state=42).fit(X_scaled, y)
    best_lambda = lasso_cv.alpha_
    lasso_model = Lasso(alpha=best_lambda).fit(X_scaled, y)
    predictions = lasso_model.predict(X_scaled)
    return lasso_model, predictions, best_lambda, lasso_cv


if __name__ == "__main__":

    # 1. 加载数据
    file_path = sys.argv[1]
    data = load_data(file_path)
    # data = load_data(
    #     "H:/develop/dama/java/buliangfanying/src/main/resources/python/二级指标加权分数.csv")  # 替换为您的实际文件路径

    # 2. 准备数据
    model_data = prepare_data(data)
    # 2.1移除所有值相同的列
    model_data = model_data.loc[:, model_data.nunique() > 1]

    # 3. 准备特征和目标变量
    feature_names = model_data.drop(columns=['area_name', 'all_score']).columns.tolist()
    X = model_data[feature_names].values
    y = model_data['all_score'].values

    # 保留但不参与计算的列
    columns_to_exclude = ['area_name']
    # 移除不需要参与计算的列
    X_for_model = model_data.drop(columns=columns_to_exclude)

    # 4. 训练模型
    lasso_model, lasso_predictions, best_lambda, lasso_cv = train_lasso(X, y)

    result_df = pd.DataFrame({
        'area_name': model_data['area_name'],
        'predicted_score': lasso_predictions
    })

    # 将 DataFrame 转换为 JSON 字符串并输出
    print(json.dumps(json.loads(result_df.to_json(orient='records')), ensure_ascii=False, indent=4))