import pandas as pd
import numpy as np
import argparse
import re 
import os

#------------------------------PARSE PARAMETERS--------------------------------
parser = argparse.ArgumentParser(description='test')
parser.add_argument('-b', '--basename', dest='basename',  help='Full path for basename.forward.read_start_counts.tsv')
parser.add_argument('-s', '--strand',   dest='strand',    help='strand [forward or reverse]')
parser.add_argument('-v', '--version', action='version', version='Anne de Jong, version 1.0, May 2025')
args = parser.parse_args()

max_gap = 5                  # Maximum gap (in bp) between positions to be considered part of the same peak
min_peak_value = 10          # Minimum sum of values in a peak to be considered

df = pd.read_csv(args.basename+"."+args.strand+".read_start_counts.tsv", sep="\t", header=None, names=["name", "pos", "strand", "value"])   # Load the data
df.sort_values(by=["name", "strand", "pos"], inplace=True)  # Ensure data is sorted



def peakcalling():
	peaks = []
	# Group by name 
	for name, group in df.groupby("name"):
		group = group.reset_index(drop=True)
		current_peak = [group.loc[0, "pos"], group.loc[0, "pos"]]
		values = [group.loc[0, "value"]]
		positions = [group.loc[0, "pos"]]
		for i in range(1, len(group)):
			pos = group.loc[i, "pos"]
			val = group.loc[i, "value"]
			if pos - group.loc[i - 1, "pos"] <= max_gap:
				# Extend current peak
				current_peak[1] = pos
				values.append(val)
				positions.append(pos)
			else:
				# Filter using max value instead of sum
				max_val = max(values)
				if max_val >= min_peak_value:
					max_idx = values.index(max_val)
					max_pos = positions[max_idx]
					peaks.append([name, current_peak[0], current_peak[1], max_val, max_pos])

				# Start new peak
				current_peak = [pos, pos]
				values = [val]
				positions = [pos]

		# Final peak
		max_val = max(values)
		if max_val >= min_peak_value:
			max_idx = values.index(max_val)
			max_pos = positions[max_idx]
			peaks.append([name, current_peak[0], current_peak[1], max_val, max_pos])
	return peaks

peaks = peakcalling()

peak_records = [[name, max_pos, max_pos + 1, max_val] for name, start, end, max_val, max_pos in peaks]
peaks_df = pd.DataFrame(peak_records, columns=["name", "max_position", "max_position_end", "max_value"])
peaks_df['strand']=args.strand
peaks_df.to_csv(args.basename+"."+args.strand+".peaks.tsv", sep="\t", index=False, header=False)

PromLoci = [[name, max_pos - 40, max_pos + 10, max_val]  for name, start, end, max_val, max_pos in peaks  if args.strand == 'forward' ]
PromLoci = [[name, max_pos - 10, max_pos + 40, max_val]  for name, start, end, max_val, max_pos in peaks  if args.strand == 'reverse' ]

PromLoci_df = pd.DataFrame(PromLoci, columns=["name", "max_position", "max_position_end", "max_value"])
PromLoci_df.to_csv(args.basename+"."+args.strand+".promloci.tsv", sep="\t", index=False, header=False)


print(f"Done. {len(peaks)} peaks written to {args.basename}")