import random

import pymysql

# 批量模拟数据并插入mysql数据库

# 建立数据库连接
connection = pymysql.connect(
    host='mysql.aimpdev.svc.k5.bigtree.zone',
    user='test_user',
    password='tmOxmmc+3jznq2cX',
    database='aimp_model',
    cursorclass=pymysql.cursors.DictCursor
)

try:
    with connection.cursor() as cursor:
        # 生成并执行插入语句
        for _ in range(1000):
            age = random.randint(18, 65)
            workclass = random.choice(['Private', 'Self-emp-not-inc', 'Self-emp-inc', 'Federal-gov', 'Local-gov', 'State-gov', 'Without-pay', 'Never-worked'])
            fnlwgt = random.randint(10000, 500000)
            education = random.choice(['Bachelors', 'Masters', 'Doctorate', 'Some-college', 'HS-grad', 'Assoc-voc', 'Assoc-acdm', 'Prof-school'])
            education_num = random.randint(1, 16)
            marital_status = random.choice(['Married-civ-spouse', 'Divorced', 'Never-married', 'Separated', 'Widowed', 'Married-spouse-absent', 'Married-AF-spouse'])
            occupation = random.choice(['Tech-support', 'Craft-repair', 'Other-service', 'Sales', 'Exec-managerial', 'Prof-specialty', 'Handlers-cleaners', 'Machine-op-inspct'])
            happy = random.randint(0, 1)
            birth_date = f"{random.randint(1950, 2005)}-{random.randint(1, 12):02d}-{random.randint(1, 28):02d}"

            insert_query = f"INSERT INTO model_train_set (age, workclass, fnlwgt, education, education_num, marital_status, occupation, happy, birth_date) VALUES ({age}, '{workclass}', {fnlwgt}, '{education}', {education_num}, '{marital_status}', '{occupation}', {happy}, '{birth_date}')"
            cursor.execute(insert_query)

    # 提交事务
    connection.commit()
    print("数据插入成功！")

except pymysql.Error as e:
    print("数据库错误：", e)

finally:
    # 关闭数据库连接
    connection.close()
    print("数据库连接已关闭！")