#!/usr/bin/env python3

import argparse
from fastaUtils.fasta import parse_fasta, parse_header

if __name__=="__main__":
  parser = argparse.ArgumentParser(prog='fst-sort',description="Sort fasta file",formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument('infile', nargs='?', default=None, help='Input file in fasta format')
  parser.add_argument('-k', dest="key", default="uid", help='Key used to sort the sequences')
  parser.add_argument('-r', dest="reverse", action="store_true", default=False, help='Reverse the ordering')
  parser.add_argument('-F', dest="orderfile", type=str,default=None, help='Use same order as in this file. Sequences that have no correspondence in this file will be excluded')
  parser.add_argument('-R', dest="random", action="store_true", default=False, help='Random sort/shuffle')
  parser.add_argument('--seed', dest="seed", type=int, default=None, help='Seed used during random sort')
  args=parser.parse_args()

  if args.random:
    msa=[seq for seq in parse_fasta(args.infile)]
    import random
    random.seed(args.seed)
    random.shuffle(msa)
    for seq in msa:
      print(seq)
  else:
    msa={}
    sequences=parse_fasta(args.infile)
    if args.orderfile is not None:
      order=[]
      with open(args.orderfile,'r') as of:
        for val in of:
          try:
            order.append(int(val.strip()))
          except:
            order.append(val.strip())
      if args.reverse:
        order=order[::-1]
      fsorder=frozenset(order)
      for seq in sequences:
        db,uid,name,descr,os,ox,gn,pe,sv,beg,end=parse_header(seq.header.strip())
        if locals()[args.key] in fsorder:
          msa[locals()[args.key]]=seq
    else:
      for seq in sequences:
        db,uid,name,descr,os,ox,gn,pe,sv,beg,end=parse_header(seq.header.strip())
        msa[locals()[args.key]]=seq
      order=sorted(msa.keys(),reverse=args.reverse)

    for val in order:
      try:
        print(msa[val])
      except:
        continue
