#!/usr/bin/env python3

from collections import Counter, defaultdict
import argparse

# Target sequences
ANCHOR_SEQ = "CACTCTTTCCCTACACGACGCTCTTCCGATCT"
DOWNSTREAM_SEQ = "GATCGGA"

# Barcode settings
BARCODE_OFFSET = 8
BARCODE_LENGTH = 4

# Barcode name mapping
BARCODE_NAMES = {
    "agtc": "Barcode 1",
    "cttg": "Barcode 3",
    "aagg": "Barcode 5",
    "ttcc": "Barcode 6",
    "gtgc": "Barcode 7",
    "gcca": "Barcode 8",
}


def reverse_complement(seq):
    """Generate reverse complement of DNA sequence."""
    complement = str.maketrans("ACGTNacgtn", "TGCANtgcan")
    return seq.translate(complement)[::-1]


def read_fastq_sequences(fastq_file):
    """
    Read FASTQ file and extract sequence lines
    (line 2 and every 4th line thereafter).
    """
    sequences = []

    with open(fastq_file, "r") as f:
        for i, line in enumerate(f):
            if i % 4 == 1:
                sequences.append(line.strip())

    return sequences


def generate_all_sequences(sequences):
    """Add reverse complements to sequence list."""
    revcomp_sequences = [reverse_complement(seq) for seq in sequences]
    return sequences + revcomp_sequences


def extract_barcode_and_position(seq):
    """
    Find anchor sequence.
    Extract 4-nt barcode beginning 9 nt after anchor end.
    Find position of GATCGGA relative to barcode end.
    """

    anchor_pos = seq.find(ANCHOR_SEQ)

    if anchor_pos == -1:
        return None, None

    anchor_end = anchor_pos + len(ANCHOR_SEQ)

    barcode_start = anchor_end + BARCODE_OFFSET
    barcode_end = barcode_start + BARCODE_LENGTH

    if barcode_end > len(seq):
        return None, None

    barcode = seq[barcode_start:barcode_end].lower()

    downstream_pos = seq.find(DOWNSTREAM_SEQ, barcode_end)

    if downstream_pos == -1:
        return barcode, None

    relative_position = downstream_pos - barcode_end

    return barcode, relative_position


def barcode_label(barcode):
    """
    Return named barcode if present in lookup table,
    otherwise return barcode sequence itself.
    """
    return BARCODE_NAMES.get(barcode.lower(), barcode)


def main():

    parser = argparse.ArgumentParser(
        description="Extract 4-nt barcodes from FASTQ reads and analyze GATCGGA positions."
    )

    parser.add_argument(
        "fastq",
        help="Input FASTQ file"
    )

    args = parser.parse_args()

    print("Reading FASTQ sequences...")
    sequences = read_fastq_sequences(args.fastq)

    print(f"Loaded {len(sequences)} sequences")

    print("Generating reverse complements...")
    all_sequences = generate_all_sequences(sequences)

    print(f"Total sequences including reverse complements: {len(all_sequences)}")

    barcode_counter = Counter()
    barcode_positions = defaultdict(list)

    print("Searching for anchor sequence and extracting barcodes...")

    for seq in all_sequences:

        barcode, rel_pos = extract_barcode_and_position(seq)

        if barcode is not None:

            barcode_counter[barcode] += 1

            if rel_pos is not None:
                barcode_positions[barcode].append(rel_pos)

    # Top 10 barcodes
    top10 = barcode_counter.most_common(10)

    # Total reads among top 10 barcodes
    total_top10_reads = sum(count for barcode, count in top10)

    print()
    print(f"Total reads containing top 10 barcodes:\t{total_top10_reads}")
    print()

    # Output table header
    print("Barcode\tFrequency (%)\tPosition of ligated sequence\tFrequency of Position (%)")

    for barcode, count in top10:

        barcode_name = barcode_label(barcode)

        # Barcode percentage among top 10 reads
        barcode_percent = (count / total_top10_reads) * 100

        positions = barcode_positions[barcode]

        if len(positions) > 0:

            pos_counter = Counter(positions)

            most_common_position, pos_freq = pos_counter.most_common(1)[0]

            # Percentage within reads containing this barcode
            pos_percent = (pos_freq / count) * 100

            print(
                f"{barcode_name}\t"
                f"{barcode_percent:.2f}\t"
                f"{most_common_position}\t"
                f"{pos_percent:.2f}"
            )

        else:

            print(
                f"{barcode_name}\t"
                f"{barcode_percent:.2f}\t"
                f"NA\t"
                f"0.00"
            )


if __name__ == "__main__":
    main()