#!/usr/bin/env python3

import os
import sys
import logging
import resource
import platform
import shutil
import multiprocessing as mp
from pyfastatools import Parser, Record, Header
from collections import defaultdict
from functools import partial
from tqdm import tqdm

def set_memory_limit(limit_in_gb):
    limit_in_bytes = limit_in_gb * 1024 * 1024 * 1024
    try:
        resource.setrlimit(resource.RLIMIT_AS, (limit_in_bytes, limit_in_bytes))
    except (ValueError, OSError, AttributeError) as e:
        logger.warning(f"Unable to set memory limit. Error: {e}")

log_level = logging.DEBUG if snakemake.params.debug else logging.INFO
log_file = snakemake.params.log
logging.basicConfig(
    level=log_level,
    format='%(asctime)s | %(levelname)s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    handlers=[
        logging.FileHandler(log_file, mode='a'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger()

print("========================================================================\n        Step 4/11: Filter the input sequences by number of ORFs         \n========================================================================")
with open(log_file, "a") as log:
    log.write("========================================================================\n        Step 4/11: Filter the input sequences by number of ORFs         \n========================================================================\n")

def count_cds_and_filter_by_contig(file_path, min_cds):
    """
    Groups amino-acid sequences by contig and filters contigs based on a minimum CDS threshold.
    Returns a dictionary of contigs with CDS counts meeting the threshold.
    """
    contig_cds_counts = {}
    for record in Parser(file_path):
        contig_id = record.header.name.rsplit("_", 1)[0]
        contig_cds_counts.setdefault(contig_id, []).append(record)
    
    # Filter contigs based on the min_cds threshold
    return {contig_id: cds for contig_id, cds in contig_cds_counts.items() if len(cds) >= min_cds}

def process_nucl_input(input_prot_subdir):
    """
    Processes input files when the input type is 'nucl' (nucleotide-based inputs translated to proteins).
    Checks if single-contig and/or vMAG protein files exist and organizes input paths.

    Parameters:
    - input_prot_subdir (str): The directory containing amino-acid FASTAs generated by pyrodigal-GV in the CheckAMG workflow.

    Returns:
    - input_files (list): List of valid protein file paths.
    - has_single_contig (bool): Whether single-contig protein file exists.
    - has_vmag_files (bool): Whether vMAG protein files exist.
    """
    single_contig_file = os.path.join(input_prot_subdir, "single_contig_proteins.faa")
    vmag_dir = os.path.join(input_prot_subdir, "vMAG_proteins")

    has_single_contig = os.path.exists(single_contig_file)
    has_vmag_files = os.path.exists(vmag_dir) and any(
        f.endswith('.faa') for f in os.listdir(vmag_dir) if os.path.isfile(os.path.join(vmag_dir, f))
    )

    if not has_single_contig and not has_vmag_files:
        logger.error("No valid nucleotide-based input files found.")
        raise FileNotFoundError("No valid nucleotide-based input files found for amino-acid sequence filtering.")

    input_files = []
    if has_single_contig:
        input_files.append(single_contig_file)
    if has_vmag_files:
        input_files.extend([
            os.path.join(vmag_dir, faa) for faa in os.listdir(vmag_dir) if faa.endswith(".faa")
        ])

    return input_files, has_single_contig, has_vmag_files

def process_prot_input(input_single_contig_prots, input_vmag_prots):
    """
    Processes input files when the input type is 'prot' (already translated amino acid sequences).
    It ensures both single-contig and vMAG protein files exist and organizes input paths.
    """
    has_single_contig = os.path.exists(input_single_contig_prots)
    has_vmag_files = True if input_vmag_prots else False

    if not has_single_contig and not has_vmag_files:
        logger.error("No valid input protein files found.")
        raise FileNotFoundError("No valid input protein files found.")

    input_files = []
    if has_single_contig:
        input_files.append(input_single_contig_prots)
    if has_vmag_files:
        input_files.extend([
            faa for faa in input_vmag_prots
        ])
    
    return input_files, has_single_contig, has_vmag_files

def get_totals_before_filtering(input_files, single_contig_file):
    """
    Calculate total contigs and amino-acid sequences before filtering.

    Parameters:
    - input_files (list): List of protein file paths.
    - single_contig_file (str): Path to the single-contig protein file, if applicable.

    Returns:
    - total_contigs (int): Number of contigs before filtering.
    - total_amino_acid_sequences (int): Number of amino-acid sequences before filtering.
    """
    total_genomes = 0
    total_contigs = 0
    total_amino_acid_sequences = 0

    for file_path in input_files:
        contig_counts = {}
        for record in Parser(file_path):
            contig_id = record.header.name.rsplit("_", 1)[0]
            contig_counts.setdefault(contig_id, []).append(record)
            
        if file_path == single_contig_file:
            total_genomes += len(contig_counts)
        else:
            total_genomes += 1
        total_contigs += len(contig_counts)
        total_amino_acid_sequences += sum(len(records) for records in contig_counts.values())

    return total_genomes, total_contigs, total_amino_acid_sequences

def process_vmag_file(vmag_file, output_folder, min_num_sequences):
    filtered_contigs = count_cds_and_filter_by_contig(vmag_file, min_num_sequences)
    output_file = os.path.join(output_folder, os.path.basename(vmag_file))
    buffer = []
    chunk_size = 10000
    with open(output_file, "w", buffering=1024*1024) as output_handle:
        for records in filtered_contigs.values():
            for record in records:
                buffer.append(f">{record.header.name} {record.header.desc}\n{record.seq}\n")
                if len(buffer) >= chunk_size:
                    output_handle.write(''.join(buffer))
                    buffer = []
        if buffer:
            output_handle.write(''.join(buffer))
    return len(filtered_contigs), sum(len(v) for v in filtered_contigs.values())

def filter_and_save_vmag_proteins(input_files, output_folder, min_num_sequences):
    vmag_files = [f for f in input_files if "vMAG_proteins" in f]
    if not vmag_files:
        return

    vmag_output_folder = os.path.join(output_folder, "vMAG_proteins")
    os.makedirs(vmag_output_folder, exist_ok=True)

    logger.info(f"Processing {len(vmag_files)} vMAG files...")

    with mp.Pool(processes=mp.cpu_count()) as pool:
        func = partial(process_vmag_file, output_folder=vmag_output_folder, min_num_sequences=min_num_sequences)
        results = list(tqdm(pool.imap_unordered(func, vmag_files), total=len(vmag_files), desc="Filtering vMAGs", unit="file"))

    total_contigs = sum(r[0] for r in results)
    total_orfs = sum(r[1] for r in results)
    logger.info(f"Filtered vMAGs: {total_contigs:,} contigs and {total_orfs:,} ORFs retained")

def extract_and_filter_contig_group(args):
    contig_id, records, min_num_sequences = args
    if len(records) >= min_num_sequences:
        return contig_id, records
    return None

def filter_and_save_single_contig_proteins(input_file, output_folder, min_num_sequences):
    contig_cds = defaultdict(list)
    for record in Parser(input_file):
        contig_id = record.header.name.rsplit("_", 1)[0]
        contig_cds[contig_id].append(record)

    args = [
        (contig_id, [(rec.header.name, rec.header.desc, rec.seq) for rec in records], min_num_sequences)
        for contig_id, records in contig_cds.items()
    ]

    with mp.Pool(processes=mp.cpu_count()) as pool:
        filtered_results = list(tqdm(pool.imap_unordered(extract_and_filter_contig_group, args),
                                     total=len(args), desc="Filtering contigs", unit="contig"))

    total_orfs = 0
    single_contig_output_file = os.path.join(output_folder, "single_contig_proteins.faa")
    buffer = []
    chunk_size = 10000
    with open(single_contig_output_file, "w", buffering=1024*1024) as output_handle:
        for result in filtered_results:
            if result is None:
                continue
            _, records_data = result
            for name, desc, seq in records_data:
                total_orfs += 1
                buffer.append(f">{name} {desc}\n{seq}\n")
                if len(buffer) >= chunk_size:
                    output_handle.write(''.join(buffer))
                    buffer = []
        if buffer:
            output_handle.write(''.join(buffer))
    logger.info(f"Filtered single-contig genomes: {len(contig_cds):,} contigs and {total_orfs:,} ORFs retained")

def get_totals_after_filtering(input_files, min_num_sequences, single_contig_file):
    """
    Calculate total contigs and amino-acid sequences after filtering.

    Parameters:
    - input_files (list): List of protein file paths.
    - min_num_sequences (int): Minimum CDS count threshold for filtering.

    Returns:
    - total_filtered_contigs (int): Number of contigs after filtering.
    - total_filtered_amino_acid_sequences (int): Number of amino-acid sequences after filtering.
    """
    total_filtered_genomes = 0
    total_filtered_contigs = 0
    total_filtered_amino_acid_sequences = 0

    for file_path in input_files:
        filtered_contigs = count_cds_and_filter_by_contig(file_path, min_num_sequences)
        total_filtered_contigs += len(filtered_contigs)
        if file_path == single_contig_file:
            total_filtered_genomes += len(filtered_contigs)
        else:
            total_filtered_genomes += 1
        total_filtered_amino_acid_sequences += sum(len(cds) for cds in filtered_contigs.values())

    return total_filtered_genomes, total_filtered_contigs, total_filtered_amino_acid_sequences

def main():
    input_type = snakemake.params.input_type
    output_folder = snakemake.params.output
    min_num_sequences = snakemake.params.min_cds
    mem_limit = snakemake.resources.mem
    set_memory_limit(mem_limit)

    logger.info("Amino-acid sequence filtering starting...")

    input_files, has_single_contig, has_vmag_files = [], False, False

    if input_type == "nucl":
        input_prot_subdir = snakemake.params.input_prot_subdir
        logger.debug(f"Input pyrodigal-gv translated protein subdirectory: {input_prot_subdir}")
        input_files, has_single_contig, has_vmag_files = process_nucl_input(input_prot_subdir)

    elif input_type == "prot":
        input_single_contig_prots = snakemake.params.input_single_contig_prots
        logger.debug(f"Input single-contig proteins: {input_single_contig_prots}")
        input_vmag_prots = snakemake.params.input_vmag_prots
        logger.debug(f"Input vMAG proteins: {input_vmag_prots}")
        input_files, has_single_contig, has_vmag_files = process_prot_input(input_single_contig_prots, input_vmag_prots)

    else:
        logger.error(f"Invalid input type specified: {input_type}")
        raise ValueError(f"Invalid input type specified: {input_type}")

    if has_single_contig:
        single_contig_file = input_files[0]
    else:
        single_contig_file = None
        
    # Totals before filtering
    total_genomes, total_contigs, total_amino_acid_sequences = get_totals_before_filtering(input_files, single_contig_file)
    logger.info(f"Total contigs before filtering: {total_contigs:,} ({total_genomes:,} genomes)")
    logger.info(f"Total amino-acid sequences before filtering: {total_amino_acid_sequences:,}")

    # Ensure output folder exists
    logger.debug(f"Output folder: {output_folder}")
    os.makedirs(output_folder, exist_ok=True)

    # Process vMAG files
    if has_vmag_files:
        filter_and_save_vmag_proteins(input_files, output_folder, min_num_sequences)

    # Process single-contig proteins
    if has_single_contig:
        filter_and_save_single_contig_proteins(input_files[0], output_folder, min_num_sequences)

    # Totals after filtering
    total_filtered_genomes, total_filtered_contigs, total_filtered_amino_acid_sequences = get_totals_after_filtering(input_files, min_num_sequences, single_contig_file)

    logger.info(f"Total contigs after filtering: {total_filtered_contigs:,} ({total_filtered_genomes:,} genomes)")
    logger.info(f"Total amino-acid sequences after filtering: {total_filtered_amino_acid_sequences:,}")
    
    # Remove original pyrodigal-gv directory to save space
    if input_type == "nucl":
        logger.debug(f"Removing original pyrodigal-gv directory: {input_prot_subdir}")
        shutil.rmtree(input_prot_subdir)
        
    logger.info("Amino-acid sequence filtering completed.")

if __name__ == "__main__":
    main()