#!/usr/bin/env python3
import os
import sys
import gzip
import time
import logging
import argparse
import resource
import subprocess
import networkx
import parallel_tools
import swalign
import shims
# There can be problems with the submodules, but none are essential.
# Try to load these modules, but if there's a problem, load a harmless dummy and continue.
version = shims.get_module_or_shim('utillib.version')
phone = shims.get_module_or_shim('ET.phone')

VERBOSE = (logging.DEBUG+logging.INFO)//2
USAGE = '$ %(prog)s [options] families.tsv barcodes.fa barcodes.sam > families.corrected.tsv'
DESCRIPTION = """Correct barcodes using an alignment of all barcodes to themselves. Reads the
alignment in SAM format and corrects the barcodes in an input "families" file (the output of
make-barcodes.awk). It will print the "families" file to stdout with barcodes (and orders)
corrected."""


def make_argparser():
  parser = argparse.ArgumentParser(usage=USAGE, description=DESCRIPTION)
  parser.add_argument('families', type=open_as_text_or_gzip,
    help='The sorted output of make-barcodes.awk. The important part is that it\'s a tab-delimited '
         'file with at least 2 columns: the barcode sequence and order, and it must be sorted in '
         'the same order as the "reads" in the SAM file.')
  parser.add_argument('reads', type=open_as_text_or_gzip,
    help='The fasta/q file given to the aligner. Used to get barcode sequences from read names.')
  parser.add_argument('sam', type=argparse.FileType('r'), nargs='?', default=sys.stdin,
    help='Barcode alignment, in SAM format. Omit to read from stdin. The read names must be '
         'integers, representing the (1-based) order they appear in the families file.')
  parser.add_argument('-P', '--prepend', action='store_true',
    help='Prepend the corrected barcodes and orders to the original columns.')
  parser.add_argument('-d', '--dist', type=int, default=1,
    help='NM edit distance threshold. Default: %(default)s')
  parser.add_argument('-m', '--mapq', type=int, default=20,
    help='MAPQ threshold. Default: %(default)s')
  parser.add_argument('-p', '--pos', type=int, default=2,
    help='POS tolerance. Alignments will be ignored if abs(POS - 1) is greater than this value. '
         'Set to greater than the barcode length for no threshold. Default: %(default)s')
  parser.add_argument('-c', '--choose-by', choices=('count', 'connect'), default='count',
    help='Choose the "correct" barcode in a network of related barcodes by either the count of how '
         'many times the barcode was observed ("freq") or how connected the barcode is to the '
         'others in the network ("connect").')
  parser.add_argument('-I', '--no-check-ids', dest='check_ids', action='store_false', default=True,
    help='Don\'t check to make sure read pairs have identical ids. By default, if this '
         'encounters a pair of reads in families.tsv with ids that aren\'t identical (minus an '
         'ending /1 or /2), it will throw an error.')
  parser.add_argument('--limit', type=int,
    help='Limit the number of lines that will be read from each input file, for testing purposes.')
  parser.add_argument('-S', '--structures', action='store_true',
    help='Print a list of the unique isoforms')
  parser.add_argument('--struct-human', action='store_true')
  parser.add_argument('-V', '--visualize', metavar='networks.png', nargs='?', default=0,
    help='Produce a visualization of the unique structures and write the image to this file. '
         'If you omit a filename, it will be displayed in a window.')
  parser.add_argument('-F', '--viz-format', choices=('dot', 'graphviz', 'png'), default='png')
  parser.add_argument('-n', '--no-output', dest='output', action='store_false', default=True)
  parser.add_argument('-l', '--log', type=argparse.FileType('w'), default=sys.stderr,
    help='Print log messages to this file instead of to stderr. Warning: Will overwrite the file.')
  parser.add_argument('-q', '--quiet', dest='volume', action='store_const', const=logging.CRITICAL,
                      default=logging.WARNING)
  parser.add_argument('-i', '--info', dest='volume', action='store_const', const=logging.INFO)
  parser.add_argument('-v', '--verbose', dest='volume', action='store_const', const=VERBOSE)
  parser.add_argument('-D', '--debug', dest='volume', action='store_const', const=logging.DEBUG,
    help='Print debug messages (very verbose).')
  parser.add_argument('--phone-home', action='store_true',
    help='Report helpful usage data to the developer, to better understand the use cases and '
         'performance of the tool. The only data which will be recorded is the name and version of '
         'the tool, the size of the input data, the time and memory taken to process it, and the '
         'IP address of the machine running it. Also, if the script fails, it will report the name '
         'of the exception thrown and the line of code it occurred in. No parameters or filenames '
         'are sent. All the reporting and recording code is available at '
         'https://github.com/NickSto/ET.')
  parser.add_argument('--galaxy', dest='platform', action='store_const', const='galaxy',
    help='Tell the script it\'s running on Galaxy. Currently this only affects data reported '
         'when phoning home.')
  parser.add_argument('--test', action='store_true',
    help='If reporting usage data, mark this as a test run.')
  parser.add_argument('--version', action='version', version=str(version.get_version()),
    help='Print the version number and exit.')
  return parser


def main(argv):

  # Allow using -v for --version if it's the only argument, and --verbose if there are more.
  if len(argv) == 2 and argv[1] == '-v':
    print(version.get_version())
    return

  parser = make_argparser()
  args = parser.parse_args(argv[1:])

  logging.basicConfig(stream=args.log, level=args.volume, format='%(message)s')
  tone_down_logger()

  start_time = time.time()
  # If the user requested, report back some data about the start of the run.
  if args.phone_home:
    call = phone.Call(__file__, version.get_version(), platform=args.platform, test=args.test,
                      fail='warn')
    call.send_data('start')
    call.send_data('prelim', run_data=gather_prelim_data(args.families, args.reads, args.sam))

  # Execute as much of the script as possible in a try/except to catch any exception that occurs
  # and report it via ET.phone.
  try:
    logging.info('Reading the fasta/q to map read names to barcodes..')
    names_to_barcodes = map_names_to_barcodes(args.reads, args.limit)

    logging.info('Reading the SAM to build the graph of barcode relationships..')
    graph, reversed_barcodes, num_good_alignments = read_alignments(args.sam, names_to_barcodes,
                                                                    args.pos, args.mapq,
                                                                    args.dist, args.limit)

    logging.info('Reading the families.tsv to get the counts of each family..')
    family_counts, read_pairs = get_family_counts(args.families, limit=args.limit,
                                                  check_ids=args.check_ids)

    if args.structures or args.visualize != 0:
      logging.info('Counting the unique barcode networks..')
      structures = count_structures(graph, family_counts)
      if args.structures:
        print_structures(structures, args.struct_human)
      if args.visualize != 0:
        logging.info('Generating a visualization of barcode networks..')
        visualize([s['graph'] for s in structures], args.visualize, args.viz_format)

    logging.info('Building the correction table from the graph..')
    corrections = make_correction_table(graph, family_counts, args.choose_by)

    logging.info('Reading the families.tsv again to print corrected output..')
    with open_as_text_or_gzip(args.families.name) as families:
      print_corrected_output(families, corrections, reversed_barcodes, args.prepend, args.limit,
                             args.output)

    run_time = int(time.time() - start_time)
    max_mem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
    logging.info('Max memory usage: {:0.2f}MB'.format(max_mem))
    logging.info('Wall clock time:  {} seconds'.format(run_time))

  except (Exception, KeyboardInterrupt) as exception:
    if args.phone_home and call:
      try:
        exception_data = getattr(exception, 'child_context', parallel_tools.get_exception_data())
        logging.critical(parallel_tools.format_traceback(exception_data))
        exception_data = parallel_tools.scrub_tb_paths(exception_data, script_path=__file__)
      except Exception:
        exception_data = {}
        raise
      run_time = int(time.time() - start_time)
      try:
        run_data = {'barcodes':len(names_to_barcodes), 'good_alignments':num_good_alignments,
                    'read_pairs':read_pairs, 'max_mem':int(max_mem)}
      except Exception:
        run_data = {}
      run_data['failed'] = True
      if exception_data:
        run_data['exception'] = exception_data
      call.send_data('end', run_time=run_time, run_data=run_data)
      raise exception
    else:
      raise

  if args.phone_home:
    run_data = {'barcodes':len(names_to_barcodes), 'good_alignments':num_good_alignments,
                'read_pairs':read_pairs, 'max_mem':int(max_mem)}
    call.send_data('end', run_time=run_time, run_data=run_data)


def gather_prelim_data(families, reads, sam):
  data = {}
  data['families_gzipped'] = isinstance(families, gzip.GzipFile)
  data['families_size'] = os.path.getsize(families.name)
  data['reads_gzipped'] = isinstance(reads, gzip.GzipFile)
  data['reads_size'] = os.path.getsize(reads.name)
  data['sam_stdin'] = sam is sys.stdin
  if data['sam_stdin']:
    data['sam_size'] = None
  else:
    data['sam_size'] = os.path.getsize(sam.name)
  return data


def detect_format(reads_file, max_lines=7):
  """Detect whether a file is a fastq or a fasta, based on its content."""
  fasta_votes = 0
  fastq_votes = 0
  line_num = 0
  for line in reads_file:
    line_num += 1
    if line_num % 4 == 1:
      if line.startswith('@'):
        fastq_votes += 1
      elif line.startswith('>'):
        fasta_votes += 1
    elif line_num % 4 == 3:
      if line.startswith('+'):
        fastq_votes += 1
      elif line.startswith('>'):
        fasta_votes += 1
    if line_num >= max_lines:
      break
  reads_file.seek(0)
  if fasta_votes > fastq_votes:
    return 'fasta'
  elif fastq_votes > fasta_votes:
    return 'fastq'
  else:
    return None


def read_fastaq(reads_file):
  filename = reads_file.name
  if filename.endswith('.fa') or filename.endswith('.fasta'):
    format = 'fasta'
  elif filename.endswith('.fq') or filename.endswith('.fastq'):
    format = 'fastq'
  else:
    format = detect_format(reads_file)
  if format == 'fasta':
    return read_fasta(reads_file)
  elif format == 'fastq':
    return read_fastq(reads_file)


def read_fasta(reads_file):
  """Read a FASTA file, yielding read names and sequences.
  NOTE: This assumes sequences are only one line!"""
  line_num = 0
  for line_raw in reads_file:
    line = line_raw.rstrip('\r\n')
    line_num += 1
    if line_num % 2 == 1:
      assert line.startswith('>'), line
      read_name = line[1:]
    elif line_num % 2 == 0:
      read_seq = line
      yield read_name, read_seq


def read_fastq(reads_file):
  """Read a FASTQ file, yielding read names and sequences.
  NOTE: This assumes sequences are only one line!"""
  line_num = 0
  for line in reads_file:
    line_num += 1
    if line_num % 4 == 1:
      assert line.startswith('@'), line
      read_name = line[1:].rstrip('\r\n')
    elif line_num % 4 == 2:
      read_seq = line.rstrip('\r\n')
      yield read_name, read_seq


def map_names_to_barcodes(reads_file, limit=None):
  """Map barcode names to their sequences."""
  names_to_barcodes = {}
  read_num = 0
  for read_name, read_seq in read_fastaq(reads_file):
    read_num += 1
    if limit is not None and read_num > limit:
      break
    try:
      name = int(read_name)
    except ValueError:
      logging.critical('Non-int read name "{}"'.format(read_name))
      raise
    names_to_barcodes[name] = read_seq
  reads_file.close()
  return names_to_barcodes


def parse_alignment(sam_file, pos_thres, mapq_thres, dist_thres):
  """Parse the SAM file and yield reads that pass the filters.
  Returns (qname, rname, reversed)."""
  line_num = 0
  for line in sam_file:
    line_num += 1
    if line.startswith('@'):
      logging.debug('Header line ({})'.format(line_num))
      continue
    fields = line.split('\t')
    logging.debug('read {} -> ref {} (read seq {}):'.format(fields[2], fields[0], fields[9]))
    qname_str = fields[0]
    rname_str = fields[2]
    if qname_str == rname_str:
      logging.debug('\tRead aligned to itself.')
      continue
    rname_fields = rname_str.split(':')
    if len(rname_fields) == 2 and rname_fields[1] == 'rev':
      reversed = True
      rname_str = rname_fields[0]
    else:
      reversed = False
    try:
      qname = int(qname_str)
      rname = int(rname_str)
    except ValueError:
      if fields[2] == '*':
        logging.debug('\tRead unmapped (reference == "*")')
        continue
      else:
        logging.error('Non-integer read name(s) on line {}: "{}", "{}".'
                      .format(line_num, qname, rname))
        raise
    # Apply alignment quality filters.
    try:
      flags = int(fields[1])
      pos = int(fields[3])
      mapq = int(fields[4])
    except ValueError:
      logging.warn('\tNon-integer flag ({}), pos ({}), or mapq ({})'
                   .format(fields[1], fields[3], fields[4]))
      continue
    if flags & 4:
      logging.debug('\tRead unmapped (flag & 4 == True)')
      continue
    if abs(pos - 1) > pos_thres:
      logging.debug('\tAlignment failed pos filter: abs({} - 1) > {}'.format(pos, pos_thres))
      continue
    if mapq < mapq_thres:
      logging.debug('\tAlignment failed mapq filter: {} > {}'.format(mapq, mapq_thres))
      continue
    nm = None
    for tag in fields[11:]:
      if tag.startswith('NM:i:'):
        try:
          nm = int(tag[5:])
        except ValueError:
          logging.error('Invalid NM tag "{}" on line {}.'.format(tag, line_num))
          raise
        break
    assert nm is not None, line_num
    if nm > dist_thres:
      logging.debug('\tAlignment failed NM distance filter: {} > {}'.format(nm, dist_thres))
      continue
    yield qname, rname, reversed
  sam_file.close()


def read_alignments(sam_file, names_to_barcodes, pos_thres, mapq_thres, dist_thres, limit=None):
  """Read the alignments from the SAM file.
  Returns (graph, reversed_barcodes, num_good_alignments):
  graph: A networkx.Graph() containing a node per barcode (the sequence as a str), and an edge
    between every pair of barcodes that align to each other (with a threshold-passing alignment).
  reversed_barcodes: The set() of all barcode sequences that are involved in an alignment where the
    target is reversed (swapped halves, like alpha+beta -> beta+alpha). Both the query and reference
    sequence in each alignment are marked here.
  num_good_alignments: The raw number of alignments processed that passed the filters."""
  graph = networkx.Graph()
  reversed_barcodes = set()
  # Maps correct barcode numbers to sets of original barcodes (includes correct ones).
  num_good_alignments = 0
  for qname, rname, reversed in parse_alignment(sam_file, pos_thres, mapq_thres, dist_thres):
    num_good_alignments += 1
    if limit is not None and num_good_alignments > limit:
      break
    # Skip self-alignments.
    if rname == qname:
      continue
    rseq = names_to_barcodes[rname]
    qseq = names_to_barcodes[qname]
    # Is this an alignment to a reversed barcode?
    if reversed:
      reversed_barcodes.add(rseq)
      reversed_barcodes.add(qseq)
    graph.add_node(rseq)
    graph.add_node(qseq)
    graph.add_edge(rseq, qseq)
  return graph, reversed_barcodes, num_good_alignments


def get_family_counts(families_file, limit=None, check_ids=True):
  """For each family (barcode), count how many read pairs exist for each strand (order)."""
  family_counts = {}
  last_barcode = None
  this_family_counts = None
  read_pairs = 0
  for line in families_file:
    read_pairs += 1
    if limit is not None and read_pairs > limit:
      break
    fields = line.rstrip('\r\n').split('\t')
    if check_ids:
      assert_read_ids_match(fields[2], fields[5])
    barcode = fields[0]
    order = fields[1]
    if barcode != last_barcode:
      if this_family_counts:
        this_family_counts['all'] = this_family_counts['ab'] + this_family_counts['ba']
      family_counts[last_barcode] = this_family_counts
      this_family_counts = {'ab':0, 'ba':0}
      last_barcode = barcode
    this_family_counts[order] += 1
  this_family_counts['all'] = this_family_counts['ab'] + this_family_counts['ba']
  family_counts[last_barcode] = this_family_counts
  families_file.close()
  return family_counts, read_pairs


def assert_read_ids_match(name1, name2):
  id1 = name1.split()[0]
  id2 = name2.split()[0]
  if id1.endswith('/1'):
    id1 = id1[:-2]
  if id2.endswith('/2'):
    id2 = id2[:-2]
  if id1 == id2:
    return True
  elif id1.endswith('/2') and id2.endswith('/1'):
    raise ValueError('Read names not as expected. Mate 1 ends with /2 and mate 2 ends with /1:\n'
                     '  Mate 1: {!r}\n  Mate 2: {!r}'.format(name1, name2))
  else:
    raise ValueError('Read names "{}" and "{}" do not match.'.format(name1, name2))


def make_correction_table(meta_graph, family_counts, choose_by='count'):
  """Make a table mapping original barcode sequences to correct barcodes.
  Assumes the most connected node in the graph as the correct barcode."""
  corrections = {}
  for nodes in networkx.connected_components(meta_graph):
    graph = meta_graph.subgraph(nodes)
    if choose_by == 'count':
      def key(bar):
        return family_counts[bar]['all']
    elif choose_by == 'connect':
      degrees = graph.degree()
      def key(bar):
        return degrees[bar]
    barcodes = sorted(graph.nodes(), key=key, reverse=True)
    correct = barcodes[0]
    for barcode in barcodes:
      if barcode != correct:
        logging.debug('Correcting {} ->\n           {}\n'.format(barcode, correct))
        corrections[barcode] = correct
  return corrections


def print_corrected_output(families_file, corrections, reversed_barcodes, prepend=False, limit=None,
                           output=True):
  line_num = 0
  barcode_num = 0
  barcode_last = None
  corrected = {'reads':0, 'barcodes':0, 'reversed':0}
  reads = [0, 0]
  corrections_in_this_family = 0
  for line in families_file:
    line_num += 1
    if limit is not None and line_num > limit:
      break
    fields = line.rstrip('\r\n').split('\t')
    raw_barcode = fields[0]
    order = fields[1]
    if raw_barcode != barcode_last:
      # We just started a new family.
      barcode_num += 1
      family_info = '{}\t{}\t{}'.format(barcode_last, reads[0], reads[1])
      if corrections_in_this_family:
        corrected['reads'] += corrections_in_this_family
        corrected['barcodes'] += 1
        family_info += '\tCORRECTED!'
      else:
        family_info += '\tuncorrected'
      logging.log(VERBOSE, family_info)
      reads = [0, 0]
      corrections_in_this_family = 0
      barcode_last = raw_barcode
    if order == 'ab':
      reads[0] += 1
    elif order == 'ba':
      reads[1] += 1
    if raw_barcode in corrections:
      correct_barcode = corrections[raw_barcode]
      corrections_in_this_family += 1
      # Check if the order of the barcode reverses in the correct version.
      # First, we check in reversed_barcodes whether either barcode was involved in a reversed
      # alignment, to save time (is_alignment_reversed() does a full smith-waterman alignment).
      if ((raw_barcode in reversed_barcodes or correct_barcode in reversed_barcodes) and
          is_alignment_reversed(raw_barcode, correct_barcode)):
        # If so, then switch the order field.
        corrected['reversed'] += 1
        if order == 'ab':
          correct_order = 'ba'
        else:
          correct_order = 'ab'
      else:
        correct_order = order
    else:
      correct_barcode = raw_barcode
      correct_order = order
    # Add the corrected barcode and order to the output.
    if prepend:
      fields.insert(0, correct_barcode)
      fields.insert(1, correct_order)
    else:
      fields[0] = correct_barcode
      fields[1] = correct_order
    if output:
      print(*fields, sep='\t')
  if corrections_in_this_family:
    corrected['reads'] += corrections_in_this_family
    corrected['barcodes'] += 1
  logging.info('Corrected {barcodes} barcodes on {reads} read pairs, with {reversed} reversed.'
               .format(**corrected))


def is_alignment_reversed(barcode1, barcode2):
  """Return True if the barcodes are reversed with respect to each other, False otherwise.
  "reversed" in this case meaning the alpha + beta halves are swapped.
  Determine by aligning the two to each other, once in their original forms, and once with the