import pandas as pd
import argparse
import pdb
import re

parser = argparse.ArgumentParser(description='Script to process the allele frequency matrix')

parser.add_argument('--region',help='Region to analyse. i.e. chr20:1-1000',required=True)
parser.add_argument('--pops',help='Populations or superpopulations to analyse. i.e. EAS,EUR,AFR,AMR,SAS',required=True)
parser.add_argument('--outfile',help='Output file name',required=True)
parser.add_argument('--af_file',help='Allele frequency file generated by Perl calculate_allele_frq_from_vcf.pl',required=True)
parser.add_argument('--depth_f',help='Path to file containing the depths per position created by function \'get_annotation\'',required=True)

args = parser.parse_args()

def process_AFs(AF_f,depth_f):
    '''
    Function to create a table containing the AFs and the depths for each position

    Args
    ----
    AF_f: string
          Path to file containing the allele frequencies created by 'get_allele_frequencies'
    depth_f: string
             Path to file containing the depths per position created by function 'get_annotation'

    Returns
    -------
    Nothing
    '''
    # parse first row in order to get the populations in the file
    p=re.compile('.*_FRQ')
    pops=[]
    with open(AF_f) as f:
        header = f.readline().split('\t')
        for i in header:
            i=i.rstrip('\n')
            if p.match(i):
                if i!='ALL_FRQ': pops.append(i)

    #Process the AFs
    cols=['#CHROM','POS','REF','ALT','ALL_TOTAL_CNT','ALL_ALT_CNT','ALL_FRQ']+pops
    names=['#CHR','POS', 'REF','ALT','AN','AC','AF']+[re.sub('FRQ','AF',i) for i in pops]
    final=['#CHR','FROM','TO','REF','ALT','DP','AN','AC','AF']+[re.sub('FRQ','AF',i) for i in pops]
    AFs_df=pd.read_csv(AF_f, sep="\t", usecols=cols, index_col=False)[cols]
    AFs_df.columns =names
    AFs_df=AFs_df.round(decimals=2).astype(object)

    #Proces the depths
    depths_DF=pd.read_csv(depth_f, sep="\t",  header=None)
    depths_DF.columns=["#CHR","POS","DP"]

    #Merge by chr-pos
    merged=pd.merge(AFs_df, depths_DF, on=['#CHR','POS'], how='inner')
    mergedA=merged.assign(TO = lambda x: x.POS)
    mergedA=mergedA.rename(columns={'POS': 'FROM'})
    mergedFinal=mergedA[final]

    #output filename
    region=args.region.replace(':','.')
    pops=args.pops.replace(',','_')

    mergedFinal.to_csv(args.outfile,sep='\t',index=False,float_format='%g')


process_AFs(args.af_file, args.depth_f)






