#!/usr/bin/env python3

import argparse
import sys

from grimoire.sequence import DNA
from grimoire.feature import Feature, mRNA, Gene
from grimoire.genome import Reader

## Command line stuff ##

extended_help = """
%(prog)s is used for converting chromosome-based genomic data into
smaller pieces. This is useful in a training environment when you want
to create cross-validation training and testing sets.
"""

parser = argparse.ArgumentParser(
	description='Segments chromosome sequences and features.',
	formatter_class=argparse.RawDescriptionHelpFormatter,
	epilog=extended_help)
parser.add_argument('--fasta', required=True, type=str,
	metavar='<path>', help='path to input fasta file')
parser.add_argument('--gff', required=True, type=str,
	metavar='<path>', help='path to input GFF3 (or similar) file')
parser.add_argument('--out', required=True, type=str,
	metavar='<str>', help='output name for fasta and gff3')
parser.add_argument('--segment', required=True, type=str,
	metavar='<str>', help='feature type (e.g. gene)')
parser.add_argument('--split', required=False, type=int, default=1,
	metavar='<int>', help='split into <int> parts [%(default)d]')
parser.add_argument('--padding', required=False, type=int, default=100,
	metavar='<int>', help='length of flanking sequence [%(default)d]')
arg = parser.parse_args()

class HamanError(Exception):
	pass

def chromify(g, beg, end):
	dna = DNA(seq=g.dna.seq[beg:end], name=g.id + '-remapped')
	dna.ftable.add_feature(Feature(dna, g.beg -beg, g.end -beg, g.strand,
		'gene', id=g.id))
	for t in g.transcripts():
		dna.ftable.add_feature(Feature(dna, t.beg -beg, t.end -beg,
			t.strand, t.type, id=t.id, pid=t.pid))
		for f in t.children:
			dna.ftable.add_feature(Feature(dna, f.beg -beg, f.end -beg,
				f.strand, f.type, pid=f.pid))
	if g.strand == '-':
		dna.revcomp()
	return dna

def segment_by_genes():

	fp = []
	if arg.split:
		for i in range(arg.split):
			fp.append({
				'gf':open(arg.out + '-' + str(i) + '.gff3', 'w+'),
				'ff':open(arg.out + '-' + str(i) + '.fa', 'w+'),
			})

	count = 0
	genome = Reader(gff=arg.gff, fasta=arg.fasta)
	for chrom in genome:
		for gene in chrom.ftable.build_genes():
			if gene.issues: continue
			if len(gene.transcripts()) == 0: continue # RNA genes?
			rbeg = gene.beg - arg.padding -1
			rend = gene.end + arg.padding +1
			if rbeg < 1 or rend > len(chrom.seq): continue
			dna = chromify(gene, rbeg, rend)
			set = count % arg.split
			count += 1
			fp[set]['gf'].write(dna.ftable.gff() + '\n')
			fp[set]['ff'].write(dna.fasta())

if __name__ == '__main__':
	if arg.segment == 'gene':
		segment_by_genes()
	elif arg.segment == 'region':
		raise NotImplemented('planned but not yet done')
	else:
		raise HamanError('unknown segmentation type: ' + arg.segment)

# There might be a better way to do this by simply taking all the gff
# in any region. Why build genes at all?
