# -*- coding: utf-8 -*-
"""

Anne de Jong, Januari 2025

"""

import sys
import pandas as pd
import argparse 
import subprocess
import os
import re
import argparse 

parser = argparse.ArgumentParser(description='FACoP upstream and downstream Intergenic regions')
parser.add_argument('-f', '--fna',           dest='fnaFilename',   help='full path of the reference genome.fna file')
parser.add_argument('-t', '--table',         dest='table',         help='full path of a tab delimited input table with at first 3 column: [chrom, start, end, strand]')
parser.add_argument('-u', '--upstreamAug',   dest='upstreamAug',   help='Length of upstream augmentation',   nargs='?', default=1, type=int)
parser.add_argument('-d', '--downstreamAug', dest='downstreamAug', help='Length of downstream augmentation', nargs='?', default=1, type=int)
parser.add_argument('-o', '--out',           dest='outBasename',   help='Full path for outputfilename without extension')
parser.add_argument('--version', action='version', version='version 1.0')
args = parser.parse_args()

complement = {'A': 'T', 'C': 'G', 'G': 'C', 'T': 'A'}


def readMultiFasta(fna_file):
	# Note: string adding to dict is slow, here we use an intermediate string; seq
	fasta = {}
	key= ''
	seq=''
	with open(fna_file) as f:
		for line in f:
			if line.startswith(">"):
				if (key != ''):   # chk previous key
					fasta[key] = seq.strip().upper()
					seq = ''
				items = re.match(">(.*?)\s", line)
				if items: key=items.group(1)
			else: seq += line.rstrip()
		if (key != ''): fasta[key] = seq.strip().upper()	# add last record	
	return fasta	


def reverse_complement(dna):
	dna = re.sub('[BDEFHIJKLMNOPQRSUVWXYZ]','G', dna)  # only keep GATC
	return ''.join([complement[base] for base in dna[::-1]])


def get_sequence(FASTA, gene):
	seq = FASTA[gene['chrom']][gene['start']-1:gene['end']]
	if gene['strand'] == 'reverse': seq = reverse_complement(seq)
	return seq

def clean_seq(seq):
	unwanted = "'[]/+.!@#$;:!*%)(&^~="
	return ''.join( c for c in seq.strip() if c not in unwanted )

def adjust_coordinates(row, upstream, downstream):
    if row['strand'] in ('forward', '+'):
        new_start = max(1, row['start'] - upstream)
        new_end = row['end'] + downstream
    else:
        new_start = max(1, row['start'] - downstream)
        new_end = row['end'] + upstream
    return pd.Series({'aug_start': new_start, 'aug_end': new_end})


def get_augmented_sequence(FASTA, row):# Fetch augmented sequences
    seq = FASTA[row['chrom']][row['aug_start']-1:row['aug_end']]
    if row['strand'] in ('reverse', '-'):  seq = reverse_complement(seq)
    return clean_seq(seq)


FNA = readMultiFasta(args.fnaFilename)

regions_df = pd.read_csv(args.table,comment='#',sep='\t',header=None,  # file has no header 
    usecols=[0, 1, 2, 3, 4],  # first 4 columns
    names=["chrom", "start", "end", "count", "strand"]).astype({'start': int, 'end': int})

regions_df[['aug_start', 'aug_end']] = regions_df.apply( lambda row: adjust_coordinates(row, args.upstreamAug, args.downstreamAug),  axis=1 ) # Augment region
regions_df["sequence"] = regions_df.apply(lambda row: get_augmented_sequence(FNA, row), axis=1)


# Export as table
regions_df.to_csv(args.outBasename+".regions.tsv", sep="\t", index=False)

# Export as fasta
fasta = ''
for index, row in regions_df.iterrows(): fasta += f">prom_{index:04d}\n{row['sequence']}\n"
with open(args.outBasename+".regions.fnn", 'w')   as f: f.write(fasta)

# Export as GFF