import subprocess
import sys

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[1]
    chr_field=line[3].split('_')
    chromosome=chr_field[len(chr_field)-1]
    path_start=int(line[4])

    file_name=name+'_'+chromosome+'.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[len(chromosome)-2:]
        
        seg_start=position
        seg_name='s'+path[i][1:]
        seg_stop=position+segments_size[seg_name]-1

        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="mkdir minigraph_segments && mv _MINIGRAPH_* minigraph_segments"
subprocess.run(command,shell=True)