#!/usr/bin/env python
from optparse import OptionParser
import json
import os
import pdb
import random
import sys

from intervaltree import IntervalTree
import networkx as nx
import numpy as np
import pysam
import pandas as pd
import tensorflow as tf

from basenji.dna_io import dna_1hot
import gff

"""
basenji_data_gene.py

Write TF Records for gene TSS with expression measurements.
"""

################################################################################
# main
################################################################################
def main():
  usage = 'usage: %prog [options] <fasta> <tss_gff> <expr_file>'
  parser = OptionParser(usage)
  parser.add_option('-b', dest='split_bed',
      default=None, help='BED file to restrict data splitting.')
  parser.add_option('-c', dest='cluster_gene_distance',
      default=2000, type='int',
      help='Cluster genes into the same split within this distance [Default: %default]')
  parser.add_option('-f', dest='folds',
      default=None, type='int',
      help='Generate cross fold split [Default: %default]')
  parser.add_option('-g', dest='gene_index',
      default='gene_id',
      help='Key to match TSS GFF to expression table [Default: %default]')
  parser.add_option('-l', dest='seq_length',
      default=65536, type='int',
      help='Sequence length [Default: %default]')
  parser.add_option('-o', dest='out_dir',
      default='genes_out')
  parser.add_option('-n', dest='n_allowed_pct',
      default=0.25, type='float',
      help='Proportion of sequence allowed to be Ns on one side [Default: %default]')
  parser.add_option('-r', dest='seqs_per_tfr',
      default=256, type='int',
      help='Sequences per TFRecord file [Default: %default]')
  parser.add_option('-t', dest='test_pct_or_chr',
      default=0.1, type='str',
      help='Proportion of the data for testing [Default: %default]')
  parser.add_option('-v', dest='valid_pct_or_chr',
      default=0.1, type='str',
      help='Proportion of the data for validation [Default: %default]')
  (options, args) = parser.parse_args()

  if len(args) != 3:
    parser.error('')
  else:
    fasta_file = args[0]
    tss_gff_file = args[1]
    expr_file = args[2]

  if os.path.isdir(options.out_dir):
    print('Remove output directory %s.' % options.out_dir)
    exit(1)
  else:
    os.mkdir(options.out_dir)

  ################################################################
  # read genes and targets

  genes_raw_df = gff_df(tss_gff_file, options.gene_index)
  expr_raw_df = pd.read_csv(expr_file, index_col=0)

  # filter for shared genes
  shared_genes = set(genes_raw_df.index) & set(expr_raw_df.index)
  shared_genes = sorted(shared_genes)
  print('Shared %d genes of %d described and %d quantified' %  \
    (len(shared_genes), genes_raw_df.shape[0], expr_raw_df.shape[0]))

  # align gene info and expression
  genes_df = genes_raw_df.loc[shared_genes]
  expr_df = expr_raw_df.loc[shared_genes]
  assert(genes_df.shape[0] == expr_df.shape[0])

  ################################################################
  # filter genes from chromosome ends

  gene_valid_mask = sufficient_sequence(fasta_file, genes_df, 
      options.seq_length, options.n_allowed_pct)
  genes_df = genes_df.loc[gene_valid_mask]
  expr_df = expr_df.loc[gene_valid_mask]

  ################################################################
  # divide between train/valid/test

  # permute genes
  np.random.seed(44)
  permute_order = np.random.permutation(genes_df.shape[0])
  genes_df = genes_df.iloc[permute_order]
  expr_df = expr_df.iloc[permute_order]
  assert((genes_df.index == expr_df.index).all())

  if options.folds is not None:
    if options.split_bed is None:
      fold_indexes = divide_genes_folds(genes_df, options.folds,
                                        options.cluster_gene_distance)
    else:
      fold_indexes = divide_genes_folds_bed(genes_df, options.split_bed,
                                            options.cluster_gene_distance)

  else:
    try:
      # convert to float pct
      valid_pct = float(options.valid_pct_or_chr)
      test_pct = float(options.test_pct_or_chr)
      assert(0 <= valid_pct <= 1)
      assert(0 <= test_pct <= 1)

      # divide by pct
      fold_indexes = divide_genes_pct(genes_df, test_pct, valid_pct, options.cluster_gene_distance)

    except (ValueError, AssertionError):
      # divide by chr
      valid_chrs = options.valid_pct_or_chr.split(',')
      test_chrs = options.test_pct_or_chr.split(',')
      fold_indexes = divide_genes_chr(genes_df, test_chrs, valid_chrs)

  # label folds
  if options.folds is not None:
    fold_labels = ['fold%d' % fi for fi in range(options.folds)]
    num_folds = options.folds
  else:
    fold_labels = ['train', 'valid', 'test']
    num_folds = 3

  # write genes BED
  genes_bed_file = '%s/genes.bed' % options.out_dir
  genes_bed_open = open(genes_bed_file, 'w')
  for fi in range(num_folds):
    for gi in fold_indexes[fi]:
      gene = genes_df.iloc[gi]
      name_col = '%s;%s' % (gene.name, fold_labels[fi])
      cols = [gene.chr, str(gene.start), str(gene.end), name_col, '.', gene.strand]
      print('\t'.join(cols), file=genes_bed_open)
  genes_bed_open.close()

  # write targets
  descriptions = ['RNA:%s' % label for label in expr_df.columns]
  targets_df = pd.DataFrame({'identifier':expr_df.columns,
                             'description':descriptions})
  targets_df.index.name = 'index'
  targets_df.to_csv('%s/targets.txt' % options.out_dir, sep='\t')

  ################################################################
  # write TFRecords

  tfr_dir = '%s/tfrecords' % options.out_dir 
  os.mkdir(tfr_dir)

  # open FASTA
  fasta_open = pysam.Fastafile(fasta_file)

  # define options
  tf_opts = tf.io.TFRecordOptions(compression_type='ZLIB')

  for fi in range(num_folds):
    genes_set_df = genes_df.iloc[fold_indexes[fi]]
    expr_set_df = expr_df.iloc[fold_indexes[fi]]

    num_set = genes_set_df.shape[0]
    num_set_tfrs = int(np.ceil(num_set / options.seqs_per_tfr))

    # gene sequence index
    si = 0

    for tfr_i in range(num_set_tfrs):
      tfr_file = '%s/%s-%d.tfr' % (tfr_dir, fold_labels[fi], tfr_i)
      print(tfr_file)
      with tf.io.TFRecordWriter(tfr_file, tf_opts) as writer:
        # TFR index
        ti = 0
        while ti < options.seqs_per_tfr and si < num_set:
          gene = genes_set_df.iloc[si]
          seq_chrm = gene.chr
          mid_pos = (gene.start + gene.end) // 2
          seq_start = mid_pos - options.seq_length//2
          seq_end = seq_start + options.seq_length

          if seq_start < 0:
            # fill left side first
            n_requested = -seq_start
            seq_dna = ''.join([random.choice('ACGT') for i in range(n_requested)])
            seq_dna += fasta_open.fetch(seq_chrm, 0, seq_end)
          else:
            seq_dna = fasta_open.fetch(seq_chrm, seq_start, seq_end)

          # fill out right side          
          if len(seq_dna) > 0:
            n_requested = options.seq_length - len(seq_dna)
            seq_dna += ''.join([random.choice('ACGT') for i in range(n_requested)])

          # verify length
          assert(len(seq_dna) == options.seq_length)

          # orient
          if gene.strand == '-':
            seq_dna = rc(seq_dna)

          # one hot code
          seq_1hot = dna_1hot(seq_dna)

          # get targets
          targets = expr_set_df.iloc[si].values
          targets = targets.reshape((1,-1)).astype('float16')
          
          # make example
          example = tf.train.Example(features=tf.train.Features(feature={
            'sequence': _bytes_feature(seq_1hot.flatten().tostring()),
            'target': _bytes_feature(targets.flatten().tostring())}))

          # write
          writer.write(example.SerializeToString())

          # advance indexes
          ti += 1
          si += 1

  fasta_open.close()

  ################################################################
  # stats

  stats_dict = {}
  stats_dict['num_targets'] = targets_df.shape[0]
  stats_dict['seq_length'] = options.seq_length
  stats_dict['target_length'] = 1

  for fi in range(num_folds):
    stats_dict['%s_seqs' % fold_labels[fi]] = len(fold_indexes[fi])

  with open('%s/statistics.json' % options.out_dir, 'w') as stats_json_out:
    json.dump(stats_dict, stats_json_out, indent=4)


################################################################################
def divide_genes_chr(genes_df, test_chrs, valid_chrs):
  """Divide genes into train/valid/test lists by chromosome."""

  # initialize train/valid/test genes lists
  train_index = []
  valid_index = []
  test_index = []

  # process contigs
  gi = 0
  for gene in genes_df.itertuples(index=False):
    if gene.chr in test_chrs:
      test_index.append(gi)      
    elif gene.chr in valid_chrs:
      valid_index.append(gi)
    else:
      train_index.append(gi)
    gi += 1

  train_n = len(train_index)
  valid_n = len(valid_index)
  test_n = len(test_index)
  genes_n = genes_df.shape[0]

  print('Genes divided into')
  print(' Train: %5d genes, (%.4f)' % \
      (train_n, train_n/genes_n))
  print(' Valid: %5d genes, (%.4f)' % \
      (valid_n, valid_n/genes_n))
  print(' Test:  %5d genes, (%.4f)' % \
      (test_n, test_n/genes_n))

  return train_index, valid_index, test_index

################################################################################
def cluster_genes(genes_df, cluster_gene_distance):
  """Cluster genes within a specific distance into a graph."""

  # add numeric index
  genes_df['index_num'] = np.arange(genes_df.shape[0])

  # initialize undirected graph
  genes_graph = nx.Graph()

  # add all genes
  for gi in range(genes_df.shape[0]):
    genes_graph.add_node(gi)

  # for each chromosome
  for chrm in set(genes_df.chr):

    # sort chromosome genes
    genes_chr_df = genes_df[genes_df.chr==chrm].sort_values('start')

    # maintain list of open genes that may still be within cluster distance
    open_genes = []

    # for each chromosome gene
    for gene in genes_chr_df.itertuples():

      # initialize a list of still open genes
      still_open_genes = [gene]

      for ogene in open_genes:
        if gene.start - ogene.start < cluster_gene_distance:
          # add edge between genes
          genes_graph.add_edge(ogene.index_num, gene.index_num)

          # ogene remains open
          still_open_genes.append(ogene)

      # set new open gene list
      open_genes = still_open_genes

  # remove numeric index
  del genes_df['index_num']

  return genes_graph

################################################################################
def divide_genes_folds(genes_df, folds, cluster_gene_distance):
  """Divide genes uniformly into folds."""

  # make gene graph
  genes_graph = cluster_genes(genes_df, cluster_gene_distance)

  # initialize fold gene lists
  fold_size = np.zeros(folds)
  fold_genes = []
  for fi in range(folds):
    fold_genes.append([])

  # determine aimed genes/fold
  num_genes = genes_df.shape[0]
  fold_gene_aim = int(np.ceil(num_genes/folds))

  # process connected componenets
  for genes_cc in nx.connected_components(genes_graph):
    # compute gap between current and aim
    fold_gene_gap = fold_gene_aim - fold_size
    fold_gene_gap = np.clip(fold_gene_gap, 0, np.inf)

    # compute sample probability
    fold_prob = fold_gene_gap / fold_gene_gap.sum()

    # sample train/valid/test
    fi = np.random.choice(folds, p=fold_prob)
    fold_genes[fi] += genes_cc
    fold_size[fi] += len(genes_cc)

  print('Genes divided into')
  for fi in range(folds):
    print(' Fold%d: %5d genes, (%.4f)' % \
      (fi, fold_size[fi], fold_size[fi]/num_genes))

  return fold_genes

################################################################################
def divide_genes_folds_bed(genes_df, split_bed_file, cluster_gene_distance):
  """Divide genes into folds according to an existing split."""

  # make gene graph
  genes_graph = cluster_genes(genes_df, cluster_gene_distance)

  # create interval trees for existing splits
  split_trees = {}
  fold_labels = set()
  for line in open(split_bed_file):
    a = line.split()
    chrm = a[0]
    start = int(a[1])
    end = int(a[2])
    fold_label = a[3]
    fold_labels.add(fold_label)
    fold_index = int(fold_label.replace('fold',''))
    if chrm not in split_trees:
      split_trees[chrm] = IntervalTree()
    split_trees[chrm][start:end] = fold_index

  # initialize fold gene lists
  folds = len(fold_labels)
  fold_size = np.zeros(folds)
  fold_genes = []
  for fi in range(folds):
    fold_genes.append([])

  # determine aimed genes/fold
  num_genes = genes_df.shape[0]
  fold_gene_aim = int(np.ceil(num_genes/folds))

  # process connected componenets
  for genes_cc in nx.connected_components(genes_graph):
    # maintain order with list
    genes_cc = list(genes_cc)

    # map genes to folds
    genes_cc_splits = []
    for gi in genes_cc:
      gene = genes_df.iloc[gi]
      if gene.chr not in split_trees:
        genes_cc_splits.append(-1)
      else:
        ctg_intervals = list(split_trees[gene.chr][gene.start])
        split_intervals = list(set([ctg.data for ctg in ctg_intervals]))
        if len(split_intervals) == 0:
          genes_cc_splits.append(-1)
        elif len(split_intervals) == 1:
          genes_cc_splits.append(split_intervals[0])
        else:
          genes_cc_splits.append(split_intervals[0])
          print('WARNING: Multiple overlapping contigs for %s.' % gene.name, file=sys.stderr)

    # if component is unmapped    
    genes_cc_splits_set = sorted(set(genes_cc_splits))
    if len(genes_cc_splits_set) == 1 and genes_cc_splits_set[0] == -1:
      # compute gap between current and aim
      fold_gene_gap = fold_gene_aim - fold_size
      fold_gene_gap = np.clip(fold_gene_gap, 0, np.inf)
      
      # sample split
      # fi = np.random.choice(folds, p=fold_prob)
      fi = np.argmax(fold_gene_gap)
      fold_genes[fi] += genes_cc
      fold_size[fi] += len(genes_cc)
      # print('Unmapped to fold%d' % fi)

    else:
      # map according to overlap
      for ci, gi in enumerate(genes_cc):
        fi = genes_cc_splits[ci]

        # set unmapped to next split
        if fi == -1:
          fi = genes_cc_splits_set[1]

        fold_genes[fi].append(gi)
        fold_size[fi] += 1

  print('Genes divided into')
  for fi in range(folds):
    print(' Fold%d: %5d genes, (%.4f)' % \
      (fi, fold_size[fi], fold_size[fi]/num_genes))

  return fold_genes

################################################################################
def divide_genes_pct(genes_df, test_pct, valid_pct, cluster_gene_distance):
  """Divide genes into train/valid/test lists by percentage."""

  # make gene graph
  genes_graph = cluster_genes(genes_df, cluster_gene_distance)

  # initialize train/valid/test genes lists
  train_index = []
  valid_index = []
  test_index = []

  # process connected componenets
  for genes_cc in nx.connected_components(genes_graph):
    r = random.random()
    if r < test_pct:
      test_index += genes_cc
    elif r < test_pct+valid_pct:
      valid_index += genes_cc
    else:
      train_index += genes_cc

  train_n = len(train_index)
  valid_n = len(valid_index)
  test_n = len(test_index)
  genes_n = genes_df.shape[0]

  print('Genes divided into')
  print(' Train: %5d genes, (%.4f)' % \
      (train_n, train_n/genes_n))
  print(' Valid: %5d genes, (%.4f)' % \
      (valid_n, valid_n/genes_n))
  print(' Test:  %5d genes, (%.4f)' % \
      (test_n, test_n/genes_n))

  return [train_index, valid_index, test_index]

################################################################################
def genes_bed(genes_df, bed_file):
  """Write BED file representing gene sequencs."""
  bed_open = open(bed_file, 'w')
  for gene in genes_df.itertuples():
    cols = [gene.chr, gene.start-1, gene.end]

    # ...
  bed_open.close()

################################################################################
def gff_df(gff_file, gene_index):
  """Read GFF w/ keys into DataFrame."""

  chrms = []
  starts = []
  ends = []
  strands = []
  gtf_lists = {}
  for line in open(gff_file):
    a = line.split('\t')
    chrms.append(a[0])
    starts.append(int(a[3]))
    ends.append(int(a[3]))
    strands.append(a[6])
    for kv in gff.gtf_kv(a[-1]).items():
      gtf_lists.setdefault(kv[0],[]).append(kv[1])

  df = pd.DataFrame({
    'chr':chrms,
    'start':starts,
    'end':ends,
    'strand':strands
    })

  for k, kl in gtf_lists.items():
    df[k] = kl

  df.set_index(gene_index, inplace=True)

  return df

################################################################################
def sufficient_sequence(fasta_file, genes_df, seq_length, n_allowed_pct):
  """Return boolean mask specifying genes with sufficient sequence."""

  # open FASTA
  fasta_open = pysam.Fastafile(fasta_file)

  # initialize gene boolean
  gene_valid = np.ones(genes_df.shape[0], dtype='bool')

  gi = 0
  for gene in genes_df.itertuples():
    chr_len = fasta_open.get_reference_length(gene.chr)
    mid_pos = (gene.start + gene.end) // 2
    seq_start = mid_pos - seq_length//2
    seq_end = seq_start + seq_length

    # count requested N's
    n_requested = 0
    if seq_start < 0:
      n_requested += -seq_start
    if seq_end > chr_len:
      n_requested += seq_end - chr_len

    if n_requested > 0:
      if n_requested/seq_length < n_allowed_pct:                
        print('Allowing %s with %d Ns' % (gene.Index, n_requested), file=sys.stderr)
      else:
        print('Skipping %s with %d Ns' % (gene.Index, n_requested), file=sys.stderr)
        gene_valid[gi] = False

    gi += 1

  fasta_open.close()

  return gene_valid


def rc(seq):
  return seq.translate(str.maketrans("ATCGatcg","TAGCtagc"))[::-1]

def _bytes_feature(value):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


################################################################################
# __main__
################################################################################
if __name__ == '__main__':
  main()

#!/usr/bin/env python
from optparse import OptionParser
import json
import os
import pdb
import random
import sys

from intervaltree import IntervalTree
import networkx as nx
import numpy as np
import pysam
import pandas as pd
import tensorflow as tf

from basenji.dna_io import dna_1hot
import gff

"""
basenji_data_gene.py

Write TF Records for gene TSS with expression measurements.
"""

################################################################################
# main
################################################################################
def main():
  usage = 'usage: %prog [options] <fasta> <tss_gff> <expr_file>'
  parser = OptionParser(usage)
  parser.add_option('-b', dest='split_bed',
      default=None, help='BED file to restrict data splitting.')
  parser.add_option('-c', dest='cluster_gene_distance',
      default=2000, type='int',
      help='Cluster genes into the same split within this distance [Default: %default]')
  parser.add_option('-f', dest='folds',
      default=None, type='int',
      help='Generate cross fold split [Default: %default]')
  parser.add_option('-g', dest='gene_index',
      default='gene_id',
      help='Key to match TSS GFF to expression table [Default: %default]')
  parser.add_option('-l', dest='seq_length',
      default=65536, type='int',
      help='Sequence length [Default: %default]')
  parser.add_option('-o', dest='out_dir',
      default='genes_out')
  parser.add_option('-n', dest='n_allowed_pct',
      default=0.25, type='float',
      help='Proportion of sequence allowed to be Ns on one side [Default: %default]')
  parser.add_option('-r', dest='seqs_per_tfr',
      default=256, type='int',
      help='Sequences per TFRecord file [Default: %default]')
  parser.add_option('-t', dest='test_pct_or_chr',
      default=0.1, type='str',
      help='Proportion of the data for testing [Default: %default]')
  parser.add_option('-v', dest='valid_pct_or_chr',
      default=0.1, type='str',
      help='Proportion of the data for validation [Default: %default]')
  (options, args) = parser.parse_args()

  if len(args) != 3:
    parser.error('')
  else:
    fasta_file = args[0]
    tss_gff_file = args[1]
    expr_file = args[2]

  if os.path.isdir(options.out_dir):
    print('Remove output directory %s.' % options.out_dir)
    exit(1)
  else:
    os.mkdir(options.out_dir)

  ################################################################
  # read genes and targets

  genes_raw_df = gff_df(tss_gff_file, options.gene_index)
  expr_raw_df = pd.read_csv(expr_file, index_col=0)

  # filter for shared genes
  shared_genes = set(genes_raw_df.index) & set(expr_raw_df.index)
  shared_genes = sorted(shared_genes)
  print('Shared %d genes of %d described and %d quantified' %  \
    (len(shared_genes), genes_raw_df.shape[0], expr_raw_df.shape[0]))

  # align gene info and expression
  genes_df = genes_raw_df.loc[shared_genes]
  expr_df = expr_raw_df.loc[shared_genes]
  assert(genes_df.shape[0] == expr_df.shape[0])

  ################################################################
  # filter genes from chromosome ends

  gene_valid_mask = sufficient_sequence(fasta_file, genes_df, 
      options.seq_length, options.n_allowed_pct)
  genes_df = genes_df.loc[gene_valid_mask]
  expr_df = expr_df.loc[gene_valid_mask]

  ################################################################
  # divide between train/valid/test

  # permute genes
  np.random.seed(44)
  permute_order = np.random.permutation(genes_df.shape[0])
  genes_df = genes_df.iloc[permute_order]
  expr_df = expr_df.iloc[permute_order]
  assert((genes_df.index == expr_df.index).all())

  if options.folds is not None:
    if options.split_bed is None:
      fold_indexes = divide_genes_folds(genes_df, options.folds,
                                        options.cluster_gene_distance)
    else:
      fold_indexes = divide_genes_folds_bed(genes_df, options.split_bed,
                                            options.cluster_gene_distance)

  else:
    try:
      # convert to float pct
      valid_pct = float(options.valid_pct_or_chr)
      test_pct = float(options.test_pct_or_chr)
      assert(0 <= valid_pct <= 1)
      assert(0 <= test_pct <= 1)

      # divide by pct
      fold_indexes = divide_genes_pct(genes_df, test_pct, valid_pct, options.cluster_gene_distance)

    except (ValueError, AssertionError):
      # divide by chr
      valid_chrs = options.valid_pct_or_chr.split(',')
      test_chrs = options.test_pct_or_chr.split(',')
      fold_indexes = divide_genes_chr(genes_df, test_chrs, valid_chrs)

  # label folds
  if options.folds is not None:
    fold_labels = ['fold%d' % fi for fi in range(options.folds)]
    num_folds = options.folds
  else:
    fold_labels = ['train', 'valid', 'test']
    num_folds = 3

  # write genes BED
  genes_bed_file = '%s/genes.bed' % options.out_dir
  genes_bed_open = open(genes_bed_file, 'w')
  for fi in range(num_folds):
    for gi in fold_indexes[fi]:
      gene = genes_df.iloc[gi]
      name_col = '%s;%s' % (gene.name, fold_labels[fi])
      cols = [gene.chr, str(gene.start), str(gene.end), name_col, '.', gene.strand]
      print('\t'.join(cols), file=genes_bed_open)
  genes_bed_open.close()

  # write targets
  descriptions = ['RNA:%s' % label for label in expr_df.columns]
  targets_df = pd.DataFrame({'identifier':expr_df.columns,
                             'description':descriptions})
  targets_df.index.name = 'index'
  targets_df.to_csv('%s/targets.txt' % options.out_dir, sep='\t')

  ################################################################
  # write TFRecords

  tfr_dir = '%s/tfrecords' % options.out_dir 
  os.mkdir(tfr_dir)

  # open FASTA
  fasta_open = pysam.Fastafile(fasta_file)

  # define options
  tf_opts = tf.io.TFRecordOptions(compression_type='ZLIB')

  for fi in range(num_folds):
    genes_set_df = genes_df.iloc[fold_indexes[fi]]
    expr_set_df = expr_df.iloc[fold_indexes[fi]]

    num_set = genes_set_df.shape[0]
    num_set_tfrs = int(np.ceil(num_set / options.seqs_per_tfr))

    # gene sequence index
    si = 0

    for tfr_i in range(num_set_tfrs):
      tfr_file = '%s/%s-%d.tfr' % (tfr_dir, fold_labels[fi], tfr_i)
      print(tfr_file)
      with tf.io.TFRecordWriter(tfr_file, tf_opts) as writer:
        # TFR index
        ti = 0
        while ti < options.seqs_per_tfr and si < num_set:
          gene = genes_set_df.iloc[si]
          seq_chrm = gene.chr
          mid_pos = (gene.start + gene.end) // 2
          seq_start = mid_pos - options.seq_length//2
          seq_end = seq_start + options.seq_length

          if seq_start < 0:
            # fill left side first
            n_requested = -seq_start
            seq_dna = ''.join([random.choice('ACGT') for i in range(n_requested)])
            seq_dna += fasta_open.fetch(seq_chrm, 0, seq_end)
          else:
            seq_dna = fasta_open.fetch(seq_chrm, seq_start, seq_end)

          # fill out right side          
          if len(seq_dna) > 0:
            n_requested = options.seq_length - len(seq_dna)
            seq_dna += ''.join([random.choice('ACGT') for i in range(n_requested)])

          # verify length
          assert(len(seq_dna) == options.seq_length)

          # orient
          if gene.strand == '-':
            seq_dna = rc(seq_dna)

          # one hot code
          seq_1hot = dna_1hot(seq_dna)

          # get targets
          targets = expr_set_df.iloc[si].values
          targets = targets.reshape((1,-1)).astype('float16')
          
          # make example
          example = tf.train.Example(features=tf.train.Features(feature={
            'sequence': _bytes_feature(seq_1hot.flatten().tostring()),
            'target': _bytes_feature(targets.flatten().tostring())}))

          # write
          writer.write(example.SerializeToString())

          # advance indexes
          ti += 1
          si += 1

  fasta_open.close()

  ################################################################
  # stats

  stats_dict = {}
  stats_dict['num_targets'] = targets_df.shape[0]
  stats_dict['seq_length'] = options.seq_length
  stats_dict['target_length'] = 1

  for fi in range(num_folds):
    stats_dict['%s_seqs' % fold_labels[fi]] = len(fold_indexes[fi])

  with open('%s/statistics.json' % options.out_dir, 'w') as stats_json_out:
    json.dump(stats_dict, stats_json_out, indent=4)


################################################################################
def divide_genes_chr(genes_df, test_chrs, valid_chrs):
  """Divide genes into train/valid/test lists by chromosome."""

  # initialize train/valid/test genes lists
  train_index = []
  valid_index = []
  test_index = []

  # process contigs
  gi = 0
  for gene in genes_df.itertuples(index=False):
    if gene.chr in test_chrs:
      test_index.append(gi)      
    elif gene.chr in valid_chrs:
      valid_index.append(gi)
    else:
      train_index.append(gi)
    gi += 1

  train_n = len(train_index)
  valid_n = len(valid_index)
  test_n = len(test_index)
  genes_n = genes_df.shape[0]

  print('Genes divided into')
  print(' Train: %5d genes, (%.4f)' % \
      (train_n, train_n/genes_n))
  print(' Valid: %5d genes, (%.4f)' % \
      (valid_n, valid_n/genes_n))
  print(' Test:  %5d genes, (%.4f)' % \
      (test_n, test_n/genes_n))

  return train_index, valid_index, test_index

################################################################################
def cluster_genes(genes_df, cluster_gene_distance):
  """Cluster genes within a specific distance into a graph."""

  # add numeric index
  genes_df['index_num'] = np.arange(genes_df.shape[0])

  # initialize undirected graph
  genes_graph = nx.Graph()

  # add all genes
  for gi in range(genes_df.shape[0]):
    genes_graph.add_node(gi)

  # for each chromosome
  for chrm in set(genes_df.chr):

    # sort chromosome genes
    genes_chr_df = genes_df[genes_df.chr==chrm].sort_values('start')

    # maintain list of open genes that may still be within cluster distance
    open_genes = []

    # for each chromosome gene
    for gene in genes_chr_df.itertuples():

      # initialize a list of still open genes
      still_open_genes = [gene]

      for ogene in open_genes:
        if gene.start - ogene.start < cluster_gene_distance:
          # add edge between genes
          genes_graph.add_edge(ogene.index_num, gene.index_num)

          # ogene remains open
          still_open_genes.append(ogene)

      # set new open gene list
      open_genes = still_open_genes

  # remove numeric index
  del genes_df['index_num']

  return genes_graph

################################################################################
def divide_genes_folds(genes_df, folds, cluster_gene_distance):
  """Divide genes uniformly into folds."""

  # make gene graph
  genes_graph = cluster_genes(genes_df, cluster_gene_distance)

  # initialize fold gene lists
  fold_size = np.zeros(folds)
  fold_genes = []
  for fi in range(folds):
    fold_genes.append([])

  # determine aimed genes/fold
  num_genes = genes_df.shape[0]
  fold_gene_aim = int(np.ceil(num_genes/folds))

  # process connected componenets
  for genes_cc in nx.connected_components(genes_graph):
    # compute gap between current and aim
    fold_gene_gap = fold_gene_aim - fold_size
    fold_gene_gap = np.clip(fold_gene_gap, 0, np.inf)

    # compute sample probability
    fold_prob = fold_gene_gap / fold_gene_gap.sum()

    # sample train/valid/test
    fi = np.random.choice(folds, p=fold_prob)
    fold_genes[fi] += genes_cc
    fold_size[fi] += len(genes_cc)

  print('Genes divided into')
  for fi in range(folds):
    print(' Fold%d: %5d genes, (%.4f)' % \
      (fi, fold_size[fi], fold_size[fi]/num_genes))

  return fold_genes

################################################################################
def divide_genes_folds_bed(genes_df, split_bed_file, cluster_gene_distance):
  """Divide genes into folds according to an existing split."""

  # make gene graph
  genes_graph = cluster_genes(genes_df, cluster_gene_distance)

  # create interval trees for existing splits
  split_trees = {}
  fold_labels = set()
  for line in open(split_bed_file):
    a = line.split()
    chrm = a[0]
    start = int(a[1])
    end = int(a[2])
    fold_label = a[3]
    fold_labels.add(fold_label)
    fold_index = int(fold_label.replace('fold',''))
    if chrm not in split_trees:
      split_trees[chrm] = IntervalTree()
    split_trees[chrm][start:end] = fold_index

  # initialize fold gene lists
  folds = len(fold_labels)
  fold_size = np.zeros(folds)
  fold_genes = []
  for fi in range(folds):
    fold_genes.append([])

  # determine aimed genes/fold
  num_genes = genes_df.shape[0]
  fold_gene_aim = int(np.ceil(num_genes/folds))

  # process connected componenets
  for genes_cc in nx.connected_components(genes_graph):
    # maintain order with list
    genes_cc = list(genes_cc)

    # map genes to folds
    genes_cc_splits = []
    for gi in genes_cc:
      gene = genes_df.iloc[gi]
      if gene.chr not in split_trees:
        genes_cc_splits.append(-1)
      else:
        ctg_intervals = list(split_trees[gene.chr][gene.start])
        split_intervals = list(set([ctg.data for ctg in ctg_intervals]))
        if len(split_intervals) == 0:
          genes_cc_splits.append(-1)
        elif len(split_intervals) == 1:
          genes_cc_splits.append(split_intervals[0])
        else:
          genes_cc_splits.append(split_intervals[0])
          print('WARNING: Multiple overlapping contigs for %s.' % gene.name, file=sys.stderr)

    # if component is unmapped    
    genes_cc_splits_set = sorted(set(genes_cc_splits))
    if len(genes_cc_splits_set) == 1 and genes_cc_splits_set[0] == -1:
      # compute gap between current and aim
      fold_gene_gap = fold_gene_aim - fold_size
      fold_gene_gap = np.clip(fold_gene_gap, 0, np.inf)
      
      # sample split
      # fi = np.random.choice(folds, p=fold_prob)
      fi = np.argmax(fold_gene_gap)
      fold_genes[fi] += genes_cc
      fold_size[fi] += len(genes_cc)
      # print('Unmapped to fold%d' % fi)

    else:
      # map according to overlap
      for ci, gi in enumerate(genes_cc):
        fi = genes_cc_splits[ci]

        # set unmapped to next split
        if fi == -1:
          fi = genes_cc_splits_set[1]

        fold_genes[fi].append(gi)
        fold_size[fi] += 1

  print('Genes divided into')
  for fi in range(folds):
    print(' Fold%d: %5d genes, (%.4f)' % \
      (fi, fold_size[fi], fold_size[fi]/num_genes))

  return fold_genes

################################################################################
def divide_genes_pct(genes_df, test_pct, valid_pct, cluster_gene_distance):
  """Divide genes into train/valid/test lists by percentage."""

  # make gene graph
  genes_graph = cluster_genes(genes_df, cluster_gene_distance)

  # initialize train/valid/test genes lists
  train_index = []
  valid_index = []
  test_index = []

  # process connected componenets
  for genes_cc in nx.connected_components(genes_graph):
    r = random.random()
    if r < test_pct:
      test_index += genes_cc
    elif r < test_pct+valid_pct:
      valid_index += genes_cc
    else:
      train_index += genes_cc

  train_n = len(train_index)
  valid_n = len(valid_index)
  test_n = len(test_index)
  genes_n = genes_df.shape[0]

  print('Genes divided into')
  print(' Train: %5d genes, (%.4f)' % \
      (train_n, train_n/genes_n))
  print(' Valid: %5d genes, (%.4f)' % \
      (valid_n, valid_n/genes_n))
  print(' Test:  %5d genes, (%.4f)' % \
      (test_n, test_n/genes_n))

  return [train_index, valid_index, test_index]

################################################################################
def genes_bed(genes_df, bed_file):
  """Write BED file representing gene sequencs."""
  bed_open = open(bed_file, 'w')
  for gene in genes_df.itertuples():
    cols = [gene.chr, gene.start-1, gene.end]

    # ...
  bed_open.close()

################################################################################
def gff_df(gff_file, gene_index):
  """Read GFF w/ keys into DataFrame."""

  chrms = []
  starts = []
  ends = []
  strands = []
  gtf_lists = {}
  for line in open(gff_file):
    a = line.split('\t')
    chrms.append(a[0])
    starts.append(int(a[3]))
    ends.append(int(a[3]))
    strands.append(a[6])
    for kv in gff.gtf_kv(a[-1]).items():
      gtf_lists.setdefault(kv[0],[]).append(kv[1])

  df = pd.DataFrame({
    'chr':chrms,
    'start':starts,
    'end':ends,
    'strand':strands
    })

  for k, kl in gtf_lists.items():
    df[k] = kl

  df.set_index(gene_index, inplace=True)

  return df

################################################################################
def sufficient_sequence(fasta_file, genes_df, seq_length, n_allowed_pct):
  """Return boolean mask specifying genes with sufficient sequence."""

  # open FASTA
  fasta_open = pysam.Fastafile(fasta_file)

  # initialize gene boolean
  gene_valid = np.ones(genes_df.shape[0], dtype='bool')

  gi = 0
  for gene in genes_df.itertuples():
    chr_len = fasta_open.get_reference_length(gene.chr)
    mid_pos = (gene.start + gene.end) // 2
    seq_start = mid_pos - seq_length//2
    seq_end = seq_start + seq_length

    # count requested N's
    n_requested = 0
    if seq_start < 0:
      n_requested += -seq_start
    if seq_end > chr_len:
      n_requested += seq_end - chr_len

    if n_requested > 0:
      if n_requested/seq_length < n_allowed_pct:                
        print('Allowing %s with %d Ns' % (gene.Index, n_requested), file=sys.stderr)
      else:
        print('Skipping %s with %d Ns' % (gene.Index, n_requested), file=sys.stderr)
        gene_valid[gi] = False

    gi += 1

  fasta_open.close()

  return gene_valid


def rc(seq):
  return seq.translate(str.maketrans("ATCGatcg","TAGCtagc"))[::-1]

def _bytes_feature(value):
  return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


################################################################################
# __main__
################################################################################
if __name__ == '__main__':
  main()

#!/usr/bin/env python basenji_data_gene.py  Write TF Records for gene TSS with expression measurements. ################################################################################ # main ################################################################################ ################################################################ # read genes and targets # filter for shared genes # align gene info and expression ################################################################ # filter genes from chromosome ends ################################################################ # divide between train/valid/test # permute genes # convert to float pct # divide by pct # divide by chr # label folds # write genes BED # write targets ################################################################ # write TFRecords # open FASTA # define options # gene sequence index # TFR index # fill left side first # fill out right side # verify length # orient # one hot code # get targets # make example # write # advance indexes ################################################################ # stats ################################################################################ Divide genes into train/valid/test lists by chromosome. # initialize train/valid/test genes lists # process contigs ################################################################################ Cluster genes within a specific distance into a graph. # add numeric index # initialize undirected graph # add all genes # for each chromosome # sort chromosome genes # maintain list of open genes that may still be within cluster distance # for each chromosome gene # initialize a list of still open genes # add edge between genes # ogene remains open # set new open gene list # remove numeric index ################################################################################ Divide genes uniformly into folds. # make gene graph # initialize fold gene lists # determine aimed genes/fold # process connected componenets # compute gap between current and aim # compute sample probability # sample train/valid/test ################################################################################ Divide genes into folds according to an existing split. # make gene graph # create interval trees for existing splits # initialize fold gene lists # determine aimed genes/fold # process connected componenets # maintain order with list # map genes to folds # if component is unmapped # compute gap between current and aim # sample split # fi = np.random.choice(folds, p=fold_prob) # print('Unmapped to fold%d' % fi) # map according to overlap # set unmapped to next split ################################################################################ Divide genes into train/valid/test lists by percentage. # make gene graph # initialize train/valid/test genes lists # process connected componenets ################################################################################ Write BED file representing gene sequencs. # ... ################################################################################ Read GFF w/ keys into DataFrame. ################################################################################ Return boolean mask specifying genes with sufficient sequence. # open FASTA # initialize gene boolean # count requested N's ################################################################################ # __main__ ################################################################################