#!/usr/bin/env python

__version__ = "3.29"

import os
import io
import shutil
import re
import glob
import pandas as pd
from collections import OrderedDict
import argparse
import textwrap
import subprocess  # Added subprocess for more secure command execution


class Bruc_MLST:

    def __init__(self, read1, read2):
        self.fastq_list = []
        for read in [read1, read2]:
            if read:  # append if not None
                self.fastq_list.append(read)
        self.paired = False
        if read2:
            self.paired = True
        self.sample_name = re.sub('[_.].*', '', os.path.basename(read1))

        # https://bmcmicrobiol.biomedcentral.com/articles/10.1186/1471-2180-7-34
        with open("ST1-MLST.fasta", 'w') as write_ref:
            print(">ST1-MLST", file=write_ref)
            print(" ", file=write_ref)

    def filter_vcf(self, input_file, output_file, quality_threshold=20):
        """
        Filter a VCF file based on quality threshold
        
        Args:
            input_file (str): Path to input VCF file
            output_file (str): Path to output filtered VCF file
            quality_threshold (float): Minimum quality score to retain variant
        """
        with open(input_file, 'r') as infile, open(output_file, 'w') as outfile:
            for line in infile:
                if line.startswith('#'):
                    outfile.write(line)
                else:
                    fields = line.strip().split('\t')
                    if len(fields) >= 6:
                        try:
                            qual = float(fields[5])
                            if qual > quality_threshold:
                                outfile.write(line)
                        except ValueError:
                            continue

    def read_vcf(self, path):
        """
        Read a VCF file into a pandas DataFrame
        
        Args:
            path (str): Path to VCF file
            
        Returns:
            pandas.DataFrame: DataFrame containing VCF data
        """
        with open(path, 'r') as f:
            lines = [l for l in f if not l.startswith('##')]
        
        df = pd.read_csv(
            io.StringIO(''.join(lines)),
            dtype={'#CHROM': str, 'POS': str, 'ID': str, 'REF': str, 'ALT': str,
                'QUAL': str, 'FILTER': str, 'INFO': str},
            sep='\t'
        ).rename(columns={'#CHROM': 'CHROM'})
        
        # Convert numeric columns
        df['POS'] = pd.to_numeric(df['POS'], errors='coerce').fillna(0).astype(int)
        df['QUAL'] = pd.to_numeric(df['QUAL'], errors='coerce').fillna(0).astype(int)
        
        # Split the INFO column and extract the AC, DP and MQ fields
        df['AC'] = df['INFO'].apply(lambda x: dict(item.split("=") for item in x.split(";") if "=" in item).get('AC', None))
        df['DP'] = df['INFO'].apply(lambda x: dict(item.split("=") for item in x.split(";") if "=" in item).get('DP', None))
        df['MQ'] = df['INFO'].apply(lambda x: dict(item.split("=") for item in x.split(";") if "=" in item).get('MQ', None))
        df['AC'] = pd.to_numeric(df['AC'], errors='coerce').fillna(0).astype(int)
        df['DP'] = pd.to_numeric(df['DP'], errors='coerce').fillna(0).astype(int)
        df['MQ'] = pd.to_numeric(df['MQ'], errors='coerce').fillna(0).astype(int)
        df = df.drop(columns=['INFO', 'ID', 'FILTER', 'FORMAT'])

        return df

    def run_command(self, command):
        """
        Execute a command and return the output
        
        Args:
            command (str): Command to execute
            
        Returns:
            str: Command output (stdout)
        """
        try:
            result = subprocess.run(command, shell=True, capture_output=True, text=True, check=False)
            if result.returncode != 0:
                print(f"Warning: Command '{command}' returned non-zero exit status {result.returncode}")
                if result.stderr:
                    print(f"Error: {result.stderr.strip()}")
            return result.stdout
        except Exception as e:
            print(f"Error executing command '{command}': {e}")
            return ""

    def find_mlst(self):
        """
        Identify the MLST type for Brucella from the input sequence data
        """
        sample_name = self.sample_name + "_mlst"
        fastq_list = self.fastq_list
        mlst_reference = "ST1-MLST.fasta"
        ref = re.sub('.fasta', '', os.path.basename(mlst_reference))
        
        # Index reference sequence
        self.run_command(f"samtools faidx {mlst_reference} 2> /dev/null")
        self.run_command(f"picard CreateSequenceDictionary REFERENCE={mlst_reference} OUTPUT={ref}.dict 2> /dev/null")
        self.run_command(f"bwa index {mlst_reference} 2> /dev/null")
        
        samfile_mlst = sample_name + ".sam"
        print(f'### {sample_name} BWA Aligning reads for Brucella MLST...')
        
        # Align reads to reference
        if self.paired:
            self.run_command(f'bwa mem -M -R "@RG\\tID:{sample_name}\\tSM:{sample_name}\\tPL:ILLUMINA\\tPI:250" -t 8 {mlst_reference} {fastq_list[0]} {fastq_list[1]} > {samfile_mlst}')
        else:
            self.run_command(f'bwa mem -M -R "@RG\\tID:{sample_name}\\tSM:{sample_name}\\tPL:ILLUMINA\\tPI:250" -t 8 {mlst_reference} {fastq_list[0]} > {samfile_mlst}')
        
        print(f'{sample_name} finding MLST type...')
        all_bam_mlst = sample_name + "_all.bam"
        self.run_command(f"picard AddOrReplaceReadGroups INPUT={samfile_mlst} OUTPUT={all_bam_mlst} RGLB=lib1 RGPU=unit1 RGSM={sample_name} RGPL=illumina 2> /dev/null")
        
        # Process BAM file
        mapbam = sample_name + "_mapped.bam"
        self.run_command(f"samtools view -h -F4 -b -T {mlst_reference} {all_bam_mlst} -o {mapbam} 2> /dev/null")
        
        sortedbam = sample_name + "_sorted.bam"
        self.run_command(f"samtools sort {mapbam} -o {sortedbam} 2> /dev/null")
        self.run_command(f"samtools index {sortedbam} 2> /dev/null")
        
        # Call variants
        unfiltered_vcf_mlst = sample_name + "_unfiltered_mlst" + ".vcf"
        mapq_fix = sample_name + "_mapq_fix_mlst.vcf"
        vcf_mlst = sample_name + ".vcf"
        
        self.run_command(f"freebayes -E -1 -f {mlst_reference} {sortedbam} > {unfiltered_vcf_mlst}")
        
        # "fix" MQ notation in VCF to match GATK output
        with open(mapq_fix, 'w') as write_fix:
            with open(unfiltered_vcf_mlst, 'r') as unfiltered:
                for line in unfiltered:
                    line = line.strip()
                    new_line = re.sub(r';MQM=', r';MQ=', line)
                    print(new_line, file=write_fix)
        
        # Remove clearly poor positions
        self.filter_vcf(mapq_fix, vcf_mlst, quality_threshold=20)
        
        # Parse VCF file
        df = self.read_vcf(vcf_mlst)
        pos_call_dict = dict(zip(df['POS'], df['ALT']))
        
        # Position 1629 was too close to the end of glk sequence. Reads would not assemble properly to call possible SNP, 
        # therefore 100 bases of the gene were added. Because of this all positions beyond this point are 100 more.  
        # Same with position 1645 and 2693.

        # Define reference positions for MLST typing
        target_pos_ref = {231: 'C', 297: 'T', 363: 'C', 398: 'C', 429: 'C', 523: 'G', 631: 'G', 730: 'G', 1247: 'G', 1296: 'C', 1342: 'G', 1381: 'A', 1648: 'C', 1685: 'C', 1741: 'C', 1754: 'G', 2165: 'A', 2224: 'T', 2227: 'C', 2297: 'G', 2300: 'A', 2344: 'A', 2352: 'G', 2403: 'C', 2530: 'G', 2557: 'G', 2578: 'G', 2629: 'A', 3045: 'A', 3054: 'G', 3118: 'G', 3295: 'C', 3328: 'C', 3388: 'A', 3966: 'C', 3969: 'G', 4167: 'G', 4271: 'C', 4296: 'G', 4893: 'C', 4996: 'G', 4998: 'T', 5058: 'G', 5248: 'A', 5672: 'G', 5737: 'C', 5928: 'A', 5963: 'G', 5984: 'C', 5987: 'C', 6025: 'G', 6045: 'G', 6498: 'G', 6499: 'C', 6572: 'A', 6627: 'T', 6715: 'C', 6735: 'T', 6745: 'G', 6785: 'T', 6810: 'C', 6828: 'C', 6845: 'C', 6864: 'G', 6875: 'C', 7382: 'G', 7432: 'G', 7464: 'G', 7594: 'G', 7660: 'T', 7756: 'A'}

        # Update reference positions with variant calls
        for key, value in pos_call_dict.items():
            if key in target_pos_ref:
                target_pos_ref[key] = value
        
        # Python 3.7+ dictionaries maintain insertion order, but using OrderedDict for clarity
        ordered_combined_dict = OrderedDict(sorted(target_pos_ref.items()))
        combined_value_list = list(ordered_combined_dict.values())
        mlst_join = ''.join(combined_value_list)

        # Define MLST type dictionary
        mlst_dictionary = {
            "CTCCCGGGGCGACCCGATCGAAGCGGGAAGGCCACGGCGCGTGAGCAGCCGGGCATCTGTCCCGCGGGGTA": "MLST type 01",
            "CTCCCGGGGCGACCCGAGCGAAGCGGGAAGGCCACGGCGCGTGAGCAGCCGGGCATCTGTCCCGCGGGGTA": "MLST type 02",
            "CTCCCGTGGCGACCCGAGCGAAGCGGGAAGGCCACGGCGCGTGAGCAGCCGGGCATCTGTCCCGCGGGGTA": "MLST type 03",
            "CTCCCGGGGCGACCCGAGCGAAGCGGGAAGGCCAAGGCGCGTGAGCAGCCGGGCATCTGTCCCGCGGGGTA": "MLST type 04",
            "CTCCCGGGGCGACCCGATCGAAGCGGGAAGGCCACGGCGAGTGAGCAGCCGGGCATCTGTCCCGCGGGGTA": "MLST type 05",
            "TTCCTGGGGCAACCCGAGCGAGGCAGGGAGGCCGCGGCTCGTGAGCGGTCGGGCATCTGTCCCGCGGGGTA": "MLST type 06",
            "CTTCCTGGCCGAGCCGAGTGAAGGGGGGAGGCCACGGCGCGTGCTCGGCTGGGTACCTGTCTCGCGGTGCT": "MLST type 07",
            "CTTCCTGGCCGACCCGAGTGAAGGGGGGAGGCCACGGCGCGTGCGCGGCTGGGTACCCGTCTCGCGGTGCT": "MLST type 08",
            "CTTCCTGGCCGACCCGAGTGAAGGGGGGAGGCCACGGCGCGTGCGCGGCTGGGTACCTGTCTCGTGGTGCT": "MLST type 09",
            "CTTCCTGGCCGACCCGAGTGAAGGGGGGGGGCCACGGCGCGTGCTCGGCTGGGTACCTGTCTCGCGGTGCT": "MLST type 10",
            "CTTCCTGGCCGACCCGAGTGAAGGGGGGAGGCCACGGCGCGTGCGCGGCTGGGTACCTGTCTCGCGGTGCT": "MLST type 11",
            "CTTCCTGGCCGACCCGAGTGAAGGGGGGAGGCCACGGCGCGTGCTCGGCTGGGTACCTGTCTCGCGGTGCT": "MLST type 12",
            "CCCCCGGGCCGACTCGAGCGAAGCGAAGAGGCCACGGCGCGTGAGTGACCAGGCACCTATCCCACGGGGTA": "MLST type 13",
            "CCCCCGGGCCGGCCCAAGCGAAGCGGGGAGGCTACAGTGCGTGAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 14",
            "CCCCCGGGCCGACCCGGGCGAAGCGGGGAGGCTACGGTGCGTGAGTGGCCAGGCACCTGTCCCGCGAGGTA": "MLST type 15",
            "CCCCCGGCCCGACCCGGGCGAAGCGGGGAGGCTACGGTGCGTGAGTGGCCAGGCACCTGTCCCGCGAGGTA": "MLST type 16",
            "CCCCCGGGCCGGCCCAAGCGAAGCGGGGAGGCTACAATGCGTGAGTGGCCAGGCACCTGTCCCGCAGGGTA": "MLST type 17",
            "CCCCCGGGCCGGCCCAAGCGAAGCGGGGAGGCTACAATGCGTGAGTGGCCAGGCACCTGTCCCGCAGGCTA": "MLST type 18",
            "CCCCCGGGCCGACCCGAGCGAAGCGGGGAGGACACGGCGCGTGAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 19",
            "CCCCCGGGCCGGCCCAAGCGAAGCGGGGAGGCTACAATGCGTGAGTGGCCAGGCACATGTCCCGCAGGGTA": "MLST type 20",
            "CCCCCGGGCCGGCCCAAGCGAAGCGGGGAGGCTACAATGCGTGAGTGGCCAGGCACATGCCCCGCAGGGTA": "MLST type 21",
            "CCCCCGGGCCGACCCGAGCGAGGCGGGGAGGCCACGGCGCGGGAGTGGCCAGACACCTGTCCTGCGGGGTA": "MLST type 22",
            "CCCCCGGGCTGACCCGAGCGAAACGGGGAAGCCACGGCGCGTAAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 23",
            "CCCCCGGGCTGACCCGAGCGGAACGGGGAAGCCACGGCGCGTAAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 23x",
            "CCCCCGGGCCGACCCGAGCAAAGCGGGGAGGCCACGGCGCGTAAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 24",
            "CCCCCGGGCCGACCCGAGCGAAGCGGGGAGGCCACGGCGCGTAAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 25",
            "CCCCCGGGCCGACCCGAGCGAAGCGGGGAGGCCACGGCGCGTAAGTGGCCAAGCACCTGTTCCGCGGGGTA": "MLST type 26",
            "CCCCCGGGCCGACCCGAGCGAAGCGGGGAGACCACGGCGCATAAGTGGCCAGGCACCTGTCCCGCGGGGTA": "MLST type 27",
            "CCCTCGGGCCGACCTGAGCGAAGCGGGGAGACCACGGCGCATAAGTGGCCAGGCTCCTGTCCCGCGGGGTA": "MLST type 28" }
        # Clean up temporary files
        try:
            # Use pathlib for more modern file operations in Python 3.12
            remove_files = glob.glob('ST1-MLST*')
            for i in remove_files:
                os.remove(i)
            remove_files = glob.glob('*-mlst*')
            for i in remove_files:
                os.remove(i)
            remove_files = glob.glob('*_mlst.vcf.idx')
            for i in remove_files:
                os.remove(i)
            if os.path.exists(unfiltered_vcf_mlst):  # Check before removing
                os.remove(unfiltered_vcf_mlst)

            # Use context manager for file operations
            with open("mlst.txt", 'w') as write_out:
                if mlst_join in mlst_dictionary:
                    mlst_type = mlst_dictionary[mlst_join]
                    print(mlst_type)
                    print(mlst_type, file=write_out)
                else:
                    print("NO MLST MATCH FOUND")
                    print("NO MLST MATCH FOUND", file=write_out)
                    mlst_type = "No Brucella MLST type found"
            
            self.mlst_type = mlst_type

            # Create directory if it doesn't exist
            os.makedirs("mlst", exist_ok=True)
            
            # Move and remove files with existence checks
            if os.path.exists(vcf_mlst):
                shutil.move(vcf_mlst, "mlst")
            if os.path.exists("mlst.txt"):
                shutil.move("mlst.txt", "mlst")
            
            # Clean up with existence checks
            for file_to_remove in [samfile_mlst, all_bam_mlst, mapbam, sortedbam, f'{sortedbam}.bai', mapq_fix]:
                if os.path.exists(file_to_remove):
                    os.remove(file_to_remove)
                    
            return self.mlst_type
            
        except Exception as e:
            print(f"Error during cleanup: {e}")
            self.mlst_type = "Error determining MLST type"
            return self.mlst_type


if __name__ == "__main__": # execute if directly access by the interpreter
    parser = argparse.ArgumentParser(prog='PROG', formatter_class=argparse.RawDescriptionHelpFormatter, description=textwrap.dedent('''\

    ---------------------------------------------------------
    Based on https://bmcmicrobiol.biomedcentral.com/articles/10.1186/1471-2180-7-34

    Brucella MLST from WGS sequence
    
    '''), epilog='''---------------------------------------------------------''')
    
    parser.add_argument('-r1', '--read1', action='store', dest='read1', required=True, help='Required: single read')
    parser.add_argument('-r2', '--read2', action='store', dest='read2', required=False, default=None, help='Optional: paired read')
    parser.add_argument('-v', '--version', action='version', version=f'{os.path.abspath(__file__)}: version {__version__}')

    args = parser.parse_args()
    read1 = args.read1
    read2 = args.read2

    mlst = Bruc_MLST(read1, read2)
    mlst.find_mlst()