# 添加起始和终止位置
import pandas as pd
def Add_location(input_file, blocks_zhong, output_location_file):
    # 添加位置信息
    df = pd.read_csv(input_file)
    df2 = df.reset_index(drop=False)
    file_path = blocks_zhong
    data_dict = {}
    with open(file_path, 'r') as file:
        for lineno, line in enumerate(file, start=1):
            parts = line.split()
            begin_str = parts[1]  # 假设起始位置在第二列，也有可能在第一列，看blocks文件中是否有*
            begin_parts = begin_str.split(':')
            position1 = begin_parts[1]  # 提取位置数字
            end_str = parts[-1]  # 假设终止位置在最后一列
            end_parts = end_str.split(':')
            position2 = end_parts[1]  # 提取位置数字
            chr = end_parts[0]
            data_dict[lineno] = {'start': position1, 'end': position2, 'chr': chr}
    # 步骤4和5: 更新df_original中的Begin和End列
    for index, row in df2.iterrows():
        locus_value = int(row['LOCUS'].replace('H', ''))
        if locus_value in data_dict:
            df2.at[index, 'chr'] = data_dict[locus_value]['chr']
            df2.at[index, 'start'] = data_dict[locus_value]['start']
            df2.at[index, 'end'] = data_dict[locus_value]['end']
    df2 = df2.loc[:, df2.columns[1:]]
    columns = df2.columns.tolist()
    # 将最后三列移动到列表的开始位置
    new_columns = columns[-3:] + columns[:-3]
    # 重新排序DataFrame的列
    df2 = df2.reindex(columns=new_columns)
    # 添加SNP_ID号
    # 步骤1: 提取LOCUS列中的数字并排序
    df2['LOCUS_num'] = df2['LOCUS'].str.extract(r'(?<=H)(\d+)').astype(int)
    # 将SNP的ID号整理进来
    locus_values = df2['LOCUS_num'].astype(str)
    file_path = blocks_zhong
    with open(file_path, 'r') as file:
        lines = file.readlines()
    df2['SNP_ID'] = ''
    for index, row in df2.iterrows():
        locus_value = locus_values[index]  # 获取LOCUS列的值
        # 检查LOCUS值是否为数字，且在文件行数范围内
        if locus_value.isdigit() and int(locus_value) - 1 < len(lines):
            # 根据LOCUS值找到另一个文件的相应行
            line = lines[int(locus_value) - 1].strip()  # 行号从1开始，所以减1
            cleaned_line = line[2:]
            # 将分隔后的部分添加到新的列 'SNP_ID'
            df2.at[index, 'SNP_ID'] = cleaned_line  # 假设SNP_ID是每行的第一个值
    df2 = df2.drop('LOCUS_num', axis=1)
    df2.to_csv(output_location_file, sep=',', index=False, header=True)




