import subprocess
import sys

def has_numbers(inputString):
    return any(char.isdigit() for char in inputString)

def getChrName(chromosome_field):
    chromosome_id=""
    if has_numbers(chromosome_field):
        for char in reversed(chromosome_field): # take the last digits of the field
            if not char.isdigit():
                break
            else:
                chromosome_id+=char
        chromosome_id="Chr"+chromosome_id[::-1]
    else:
        for char in reversed(chromosome_field): # take the last uppercase chars of the fied
            if not char.isupper():
                break
            else:
                chromosome_id+=char
        chromosome_id="Chr"+chromosome_id[::-1]
    return chromosome_id

if not(len(sys.argv)==2) :
        print("expected input : gfa file with walks.")
        print("output : bed files giving the coordinates of the segments on the genomes (or on minigraph segments).")
        sys.exit(1)
elif (sys.argv[1]=="-h") :
    print("expected input : gfa file with walks.")
    print("output : bed files giving the coordinates of the segments on the genomes (or on minigraph segments).")
    sys.exit(1)

gfa=sys.argv[1]

# get the lines that start with "S"
command="grep ^S "+gfa+" > segments.txt"
subprocess.run(command,shell=True)
segments = open('segments.txt','r')
lines=segments.readlines()
segments.close()

# build a dictionnary with the segment sizes
segments_size={}
for line in lines:
    line=line.split()
    seg_id='s'+line[1]
    seg_size=len(line[2])
    segments_size[seg_id]=seg_size

# get the lines that start with "W"
command="grep ^W "+gfa+" | sed 's/>/,>/g' | sed 's/</,</g' > walks.txt"
subprocess.run(command,shell=True)
walks=open('walks.txt','r')
lines=walks.readlines()
walks.close()

# on these lines, get the name of the genome to name the output bed file
file_names=list()
for line in lines :
    line=line.split()
    name=line[3]
    path_start=int(line[4])
    chromosome_field=line[3]
    chromosome_id=getChrName(chromosome_field)

    file_name=name+'.bed'

    # if we are writing in the file for the first time, overwrite it. else, append it
    # this is because chromosomes can be fragmented. the coordinates of all the fragments from the same chromosome will be written in the same bed file.
    if file_name not in file_names:
        file_names.append(file_name)
        out_bed = open(file_name, 'w')
    else :
        out_bed = open(file_name, 'a')
        
    path=line[6].split(',')
    position=path_start
    
    for i in range(1, len(path)): # for each segment in the path, write the position of the segment in the output bed file
        # coordinates calculation : start=position, stop=position+segment_size-1, then position+=segment_size
        
        chr='Chr'+chromosome_id[len(chromosome_id)-2:]
        
        seg_start=position
        seg_name='s'+path[i][1:]
        seg_stop=position+segments_size[seg_name]

        if path[i][0:1]=='>':
            orientation='+'
        elif path[i][0:1]=='<':
            orientation='-'
        else:
            orientation='.'
    
        out_line=chr+'\t'+str(seg_start)+'\t'+str(seg_stop)+'\t'+path[i][0:1]+seg_name+'\n'
        out_bed.write(out_line)
        
        position+=segments_size[seg_name]
    out_bed.close()
 
command="rm segments.txt && rm walks.txt"
subprocess.run(command,shell=True)


command="if [ 0 -lt $(ls _MINIGRAPH_* 2>/dev/null | wc -w) ]; then mkdir minigraph_segments && mv _MINIGRAPH_* minigraph_segments; fi"
subprocess.run(command,shell=True)