import pandas as pd
import sys

input_file = sys.argv[1]

output_file = input_file.replace('.tsv', '_VAF.tsv')

# Load the data
data = pd.read_csv(input_file, sep='\t')

# Create the "VAF" column, which is the "Alt_Count" divided by the "Total_Coverage" or 0 if "Total_Coverage" is 0
data['VAF'] = data['Alt_Count'] / data['Total_Coverage']
data['VAF'] = data['VAF'].fillna(0)
data['VAF'] = data['VAF'].round(2)

# sort data by "Variant"
data['chr'] = data['Variant'].str.extract(r'(chr\d+|chrX|chrY)')

# If chrX is present in chr, replace it with 23 
data['chr'] = data['chr'].str.replace('chrX', '23')

data['chr'] = data['chr'].str.extract(r'(\d+)')
data['chr'] = data['chr'].astype(int)
data = data.sort_values(by=['chr', 'Variant'])
data.drop(columns='chr', inplace=True)

# Save the updated dataframe to a new TSV file
data.to_csv(output_file, sep='\t', index=False)

print(f'Updated file saved as {output_file}')
