import random
from datetime import datetime, timedelta

# 生成人员数据
def generate_data(count=50):
    # 姓氏和名字
    surnames = ['张', '王', '李', '刘', '陈', '杨', '黄', '赵', '周', '吴', '徐', '孙', '马', '朱', '胡', '郭', '何', '林', '罗', '高']
    names = ['伟', '强', '磊', '军', '洋', '勇', '杰', '涛', '明', '超', '华', '刚', '辉', '鹏', '飞', '波', '斌', '宇', '浩', '凯']
    
    # 监区
    prisons = ['一监区', '二监区', '三监区', '四监区', '五监区']
    
    # 班级（使用指定的班级）
    classes = ['中级班1班', '初级班1班', '高级班1班', '攻坚转换班1班', '攻坚转换班2班']
    
    # 民族
    ethnicities = ['汉族', '回族', '满族', '蒙古族', '藏族', '维吾尔族', '苗族', '彝族', '壮族', '布依族']
    
    # 文化程度
    educations = ['小学', '初中', '高中', '中专', '大专', '本科', '硕士']
    
    # 罪名
    crimes = ['盗窃罪', '诈骗罪', '抢劫罪', '故意伤害罪', '贩卖毒品罪', '交通肇事罪', '寻衅滋事罪', '聚众斗殴罪', '非法拘禁罪', '敲诈勒索罪']
    
    data = []
    
    for i in range(count):
        # 信息编号（从201开始）
        info_id = 201 + i
        
        # 姓名
        name = random.choice(surnames) + random.choice(names) + (random.choice(names) if random.random() > 0.6 else '')
        
        # 监区
        prison = random.choice(prisons)
        
        # 班级
        class_name = random.choice(classes)
        
        # 性别
        gender = '男'
        
        # 民族
        ethnicity = random.choice(ethnicities)
        
        # 文化程度
        education = random.choice(educations)
        
        # 罪名
        crime = random.choice(crimes)
        
        # 刑期 (月数)
        sentence_months = random.choice([12, 18, 24, 36, 48, 60, 72, 84, 96, 120, 144, 180])
        
        # 刑期起日
        start_year = random.randint(2022, 2024)
        start_month = random.randint(1, 12)
        start_day = random.randint(1, 28)
        start_date = f"{start_year}-{str(start_month).zfill(2)}-{str(start_day).zfill(2)}"
        
        # 刑期止日
        start_dt = datetime(start_year, start_month, start_day)
        end_dt = start_dt + timedelta(days=sentence_months * 30)
        end_date = end_dt.strftime("%Y-%m-%d")
        
        # 入监时间
        entry_date = start_date
        
        # 状态
        status = '在押'
        
        data.append([
            info_id, name, prison, class_name, gender, ethnicity, education, crime,
            sentence_months, start_date, end_date, entry_date, status
        ])
    
    return data

# 生成Excel文件
def save_to_excel(data, filename='prison_data.xlsx'):
    from openpyxl import Workbook
    
    headers = ['信息编号', '姓名', '监区', '班级', '性别', '民族', '文化程度', '罪名', '刑期', '刑期起日', '刑期止日', '入监时间', '状态']
    
    wb = Workbook()
    ws = wb.active
    ws.title = '人员数据'
    
    # 写入表头
    ws.append(headers)
    
    # 写入数据
    for row in data:
        ws.append(row)
    
    wb.save(filename)
    print(f'已生成 {len(data)} 条数据，保存到 {filename}')

if __name__ == '__main__':
    data = generate_data(30)
    save_to_excel(data, 'prison_data.xlsx')