import pandas as pd
import argparse

parser = argparse.ArgumentParser(description='Remove chromosome 5 and add a reference population', usage="proj_path")
parser.add_argument('proj_path', type=str, help='path to project root name')

arg = parser.parse_args()

snps = pd.read_csv(arg.proj_path+'.pileupCaller.snp',sep='\t',names=['rsID','chr','unk','pos','major','minor'])

idx1 = snps[(snps['chr'] == 5) & (snps['pos'] > 32000000) & (snps['pos'] < 85000000)].index

geno = pd.read_csv(arg.proj_path+'.pileupCaller.geno',names=['hap'])

geno['ref'] = '2'
geno['out'] = geno['hap'] + geno['ref']

geno.loc[idx1,'out'].to_csv(arg.proj_path+'.pileupCaller.ref.geno',header=False,index=False)
snps.loc[idx1].to_csv(arg.proj_path+'.pileupCaller.ref.snp',header=False,index=False,sep='\t')

# modify pops file, add ref to it

pops = pd.read_csv(arg.proj_path+'.pileupCaller.ind',sep='\t',names=['samp','sex','group'])
pops.loc[len(pops.index)] = ['REF00000000', 'U', 'temp']
pops.to_csv(arg.proj_path+'.pileupCaller.ref.ind',sep='\t',header=False,index=False)
