#!/usr/bin/env python

"""
Function to convert VCF file with RNA BAM file to FASTA protein sequences.

This module provides functionality to:
1. Read variants from VCF file
2. Process variants with RNA-seq BAM file using isovar
3. Extract mutant protein sequences
4. Write sequences to FASTA format
"""

import logging
import argparse
import os
from collections import OrderedDict

from isovar.cli import make_isovar_arg_parser, run_isovar_from_parsed_args

logger = logging.getLogger(__name__)


def vcf_bam_to_fasta(
        vcf_paths,
        bam_path,
        output_fasta_path,
        protein_sequence_length=None,
        centered_sequence_length=None,
        min_alt_rna_reads=2,
        min_alt_rna_fragments=2,
        min_variant_sequence_coverage=1,
        min_mapping_quality=1,
        use_duplicate_reads=False,
        use_secondary_alignments=False,
        min_base_quality=1,
        genome=None):
    """
    Convert VCF file(s) with RNA BAM file to FASTA protein sequences.

    This function reads variants from one or more VCF files, processes them with
    RNA-seq data from a BAM file using isovar, and writes the resulting mutant
    protein sequences to a FASTA file. Multiple VCF files from different variant
    callers (e.g., Mutect, Strelka) can be combined.

    Parameters
    ----------
    vcf_paths : str or list of str
        Path(s) to input VCF file(s) containing somatic variants.
        Can be a single path string or a list of paths for multiple VCF files.

    bam_path : str
        Path to RNA-seq BAM file aligned to the same reference genome

    output_fasta_path : str
        Path to output FASTA file where protein sequences will be written

    protein_sequence_length : int, optional
        Length of protein sequence to extract around each mutation.
        If None, uses the full protein sequence from isovar.
        Note: This parameter is passed to isovar and controls the initial
        sequence extraction length.

    centered_sequence_length : int, optional
        Fixed length of protein sequence centered on the mutation site (SNV/CNV).
        If specified, extracts a sequence of this length with the mutation
        positioned as close to the center as possible. If the mutation is too
        close to the sequence edge, the sequence will be shifted to include
        the mutation while maintaining the specified length.
        If None, uses the full protein sequence from isovar.
        This parameter takes precedence over protein_sequence_length for
        the final output.

    min_alt_rna_reads : int, optional
        Minimum number of RNA reads supporting the alternate allele (default: 2)

    min_alt_rna_fragments : int, optional
        Minimum number of RNA fragments supporting the alternate allele (default: 2)

    min_variant_sequence_coverage : int, optional
        Minimum coverage at variant position (default: 1)

    min_mapping_quality : int, optional
        Minimum mapping quality for reads (default: 1)

    use_duplicate_reads : bool, optional
        Whether to use duplicate reads (default: False)

    use_secondary_alignments : bool, optional
        Whether to use secondary alignments (default: False)

    min_base_quality : int, optional
        Minimum base quality (default: 1)

    genome : str, optional
        Genome reference (e.g., 'GRCh37', 'GRCh38'). If None, will try to infer
        from VCF or use default.

    Returns
    -------
    int
        Number of protein sequences written to FASTA file
    """
    # Normalize vcf_paths to a list
    if isinstance(vcf_paths, str):
        vcf_paths = [vcf_paths]

    if not vcf_paths:
        raise ValueError("At least one VCF file path must be provided")

    # Create argument parser to build isovar arguments
    arg_parser = make_isovar_arg_parser()

    # Build arguments for isovar
    # Add --vcf argument for each VCF file (isovar supports multiple --vcf)
    isovar_args = []
    for vcf_path in vcf_paths:
        isovar_args.extend(['--vcf', vcf_path])

    isovar_args.extend([
        '--bam', bam_path,
        '--min-alt-rna-reads', str(min_alt_rna_reads),
        '--min-alt-rna-fragments', str(min_alt_rna_fragments),
        '--min-variant-sequence-coverage', str(min_variant_sequence_coverage),
        '--min-mapping-quality', str(min_mapping_quality),
        '--min-base-quality', str(min_base_quality),
    ])

    if use_duplicate_reads:
        isovar_args.append('--use-duplicate-reads')

    if use_secondary_alignments:
        isovar_args.append('--use-secondary-alignments')

    if genome:
        isovar_args.extend(['--genome', genome])

    if protein_sequence_length:
        isovar_args.extend(['--protein-sequence-length', str(protein_sequence_length)])

    # Parse arguments
    args = arg_parser.parse_args(isovar_args)

    # Run isovar to get protein sequences
    vcf_files_str = ', '.join(vcf_paths) if len(vcf_paths) > 1 else vcf_paths[0]
    logger.info("Processing VCF file(s) %s with BAM file %s", vcf_files_str, bam_path)
    isovar_results = run_isovar_from_parsed_args(args)

    # Write protein sequences to FASTA file
    sequences_written = write_protein_sequences_to_fasta(
        isovar_results=isovar_results,
        output_fasta_path=output_fasta_path,
        vcf_paths=vcf_paths,
        centered_sequence_length=centered_sequence_length)

    logger.info(
        "Wrote %d protein sequences to %s",
        sequences_written,
        output_fasta_path)

    return sequences_written


def write_protein_sequences_to_fasta(
        isovar_results,
        output_fasta_path,
        vcf_paths=None,
        centered_sequence_length=None):
    """
    Write protein sequences from isovar results to FASTA file.

    Parameters
    ----------
    isovar_results : list of isovar.IsovarResult
        List of isovar results containing protein sequence information

    output_fasta_path : str
        Path to output FASTA file

    vcf_paths : list of str, optional
        List of VCF file paths used as input. If provided, will be included
        in FASTA header information.

    centered_sequence_length : int, optional
        Fixed length of protein sequence centered on the mutation site.
        If specified, extracts a sequence of this length with the mutation
        positioned as close to the center as possible.

    Returns
    -------
    int
        Number of sequences written
    """
    sequences_written = 0

    with open(output_fasta_path, 'w') as fasta_file:
        for isovar_result in isovar_results:
            # Skip variants that don't pass filters or don't have protein sequences
            if not isovar_result.passes_all_filters:
                continue

            protein_sequence = isovar_result.top_protein_sequence
            if protein_sequence is None:
                continue

            variant = isovar_result.variant

            # Extract centered sequence if requested
            amino_acids = protein_sequence.amino_acids
            mutation_start = protein_sequence.mutation_start_idx
            mutation_end = protein_sequence.mutation_end_idx

            if centered_sequence_length and centered_sequence_length > 0:
                amino_acids, mutation_start, mutation_end = extract_centered_sequence(
                    amino_acids=amino_acids,
                    mutation_start=mutation_start,
                    mutation_end=mutation_end,
                    target_length=centered_sequence_length)

            # Generate FASTA header with variant information
            header = generate_fasta_header(
                variant=variant,
                isovar_result=isovar_result,
                protein_sequence=protein_sequence,
                vcf_paths=vcf_paths,
                extracted_sequence_length=len(amino_acids) if centered_sequence_length else None)

            # Write FASTA entry
            fasta_file.write('>%s\n' % header)

            # Write protein sequence (wrap at 80 characters per line)
            for i in range(0, len(amino_acids), 80):
                fasta_file.write('%s\n' % amino_acids[i:i + 80])

            sequences_written += 1

    return sequences_written


def extract_centered_sequence(amino_acids, mutation_start, mutation_end, target_length):
    """
    Extract a fixed-length protein sequence centered on the mutation site.

    Parameters
    ----------
    amino_acids : str
        Full protein amino acid sequence

    mutation_start : int
        Start index of mutation in the protein sequence (0-based)

    mutation_end : int
        End index of mutation in the protein sequence (0-based, exclusive)

    target_length : int
        Desired length of the extracted sequence

    Returns
    -------
    tuple of (str, int, int)
        Extracted amino acid sequence, adjusted mutation_start, adjusted mutation_end
        (mutation indices are relative to the extracted sequence)
    """
    if mutation_start is None or mutation_end is None:
        # If mutation position is unknown, return original sequence or truncate
        if len(amino_acids) <= target_length:
            return amino_acids, mutation_start, mutation_end
        else:
            # Take the first target_length amino acids
            return amino_acids[:target_length], mutation_start, mutation_end

    full_length = len(amino_acids)

    # If the full sequence is shorter than target, return as is
    if full_length <= target_length:
        return amino_acids, mutation_start, mutation_end

    # Calculate the center of the mutation
    mutation_center = (mutation_start + mutation_end) // 2

    # Calculate desired start position (centered on mutation)
    desired_start = mutation_center - target_length // 2

    # Adjust if we go beyond the sequence boundaries
    if desired_start < 0:
        # Mutation is too close to the start, align to start
        sequence_start = 0
        sequence_end = target_length
    elif desired_start + target_length > full_length:
        # Mutation is too close to the end, align to end
        sequence_end = full_length
        sequence_start = full_length - target_length
    else:
        # Normal case: can center on mutation
        sequence_start = desired_start
        sequence_end = sequence_start + target_length

    # Extract the sequence
    extracted_sequence = amino_acids[sequence_start:sequence_end]

    # Adjust mutation indices relative to the extracted sequence
    adjusted_mutation_start = mutation_start - sequence_start
    adjusted_mutation_end = mutation_end - sequence_start

    # Ensure mutation indices are within bounds
    adjusted_mutation_start = max(0, min(adjusted_mutation_start, len(extracted_sequence)))
    adjusted_mutation_end = max(0, min(adjusted_mutation_end, len(extracted_sequence)))

    return extracted_sequence, adjusted_mutation_start, adjusted_mutation_end


def generate_fasta_header(
        variant,
        isovar_result,
        protein_sequence,
        vcf_paths=None,
        extracted_sequence_length=None):
    """
    Generate FASTA header line from variant and protein sequence information.

    Parameters
    ----------
    variant : varcode.Variant
        Variant object

    isovar_result : isovar.IsovarResult
        Isovar result object

    protein_sequence : isovar.ProteinSequence
        Protein sequence object

    vcf_paths : list of str, optional
        List of VCF file paths used as input. If provided, will be included
        in FASTA header to track variant source.

    extracted_sequence_length : int, optional
        Length of the extracted sequence if centered extraction was used.

    Returns
    -------
    str
        FASTA header string
    """
    # Build header components
    header_parts = OrderedDict()

    # Variant identifier
    header_parts['variant'] = variant.short_description

    # Gene name
    if protein_sequence.gene_name:
        header_parts['gene'] = protein_sequence.gene_name

    # Variant location
    header_parts['contig'] = variant.contig
    header_parts['start'] = str(variant.start)
    header_parts['ref'] = variant.ref
    header_parts['alt'] = variant.alt

    # VCF source information (if multiple VCFs provided)
    if vcf_paths and len(vcf_paths) > 1:
        # Note: When multiple VCFs are merged by varcode, individual variants
        # may not retain their source VCF information. We include the list
        # of all VCF files used as input for reference.
        # If variant has source information, use it; otherwise, note that
        # multiple VCFs were used
        if hasattr(variant, 'sources') and variant.sources:
            vcf_source = ';'.join(variant.sources)
            header_parts['vcf_source'] = vcf_source
        else:
            # Indicate that multiple VCFs were used (simplified paths)
            vcf_names = [os.path.basename(vcf) for vcf in vcf_paths]
            header_parts['vcf_inputs'] = ';'.join(vcf_names)

    # Mutation position in protein
    if protein_sequence.mutation_start_idx is not None:
        header_parts['mutation_start'] = str(protein_sequence.mutation_start_idx)
    if protein_sequence.mutation_end_idx is not None:
        header_parts['mutation_end'] = str(protein_sequence.mutation_end_idx)

    # RNA support information
    header_parts['num_alt_reads'] = str(isovar_result.num_alt_fragments)
    header_parts['num_ref_reads'] = str(isovar_result.num_ref_fragments)
    header_parts['num_supporting_fragments'] = str(
        protein_sequence.num_supporting_fragments)

    # Sequence length
    if extracted_sequence_length is not None:
        header_parts['length'] = str(extracted_sequence_length)
        header_parts['centered_extraction'] = 'true'
    else:
        header_parts['length'] = str(len(protein_sequence.amino_acids))

    # Format as key=value pairs
    header = '|'.join('%s=%s' % (k, v) for k, v in header_parts.items())

    return header


def main():
    """
    Command-line interface for vcf_bam_to_fasta function.
    """
    parser = argparse.ArgumentParser(
        description='Convert VCF file with RNA BAM file to FASTA protein sequences',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Single VCF file
  %(prog)s --vcf variants.vcf --bam rna.bam --output proteins.fasta

  # Multiple VCF files from different callers
  %(prog)s --vcf mutect.vcf --vcf strelka.vcf --vcf vardict.vcf \\
      --bam rna.bam --output proteins.fasta

  # With custom parameters
  %(prog)s --vcf variants.vcf --bam rna.bam --output proteins.fasta \\
      --protein-sequence-length 50 --min-alt-rna-reads 3

  # Extract fixed-length sequences centered on mutation
  %(prog)s --vcf variants.vcf --bam rna.bam --output proteins.fasta \\
      --centered-sequence-length 50
        """)

    parser.add_argument(
        '--vcf',
        action='append',
        required=True,
        dest='vcf_paths',
        help='Path to input VCF file. Can be specified multiple times for multiple VCF files (e.g., --vcf mutect.vcf --vcf strelka.vcf)')

    parser.add_argument(
        '--bam',
        required=True,
        help='Path to RNA-seq BAM file')

    parser.add_argument(
        '--output',
        required=True,
        dest='output_fasta_path',
        help='Path to output FASTA file')

    parser.add_argument(
        '--protein-sequence-length',
        type=int,
        default=None,
        help='Length of protein sequence to extract (default: full sequence). '
             'This is passed to isovar for initial sequence extraction.')

    parser.add_argument(
        '--centered-sequence-length',
        type=int,
        default=None,
        dest='centered_sequence_length',
        help='Fixed length of protein sequence centered on the mutation site (SNV/CNV). '
             'If specified, extracts a sequence of this length with the mutation '
             'positioned as close to the center as possible. This parameter takes '
             'precedence over --protein-sequence-length for the final output.')

    parser.add_argument(
        '--min-alt-rna-reads',
        type=int,
        default=2,
        help='Minimum number of RNA reads supporting alternate allele (default: 2)')

    parser.add_argument(
        '--min-alt-rna-fragments',
        type=int,
        default=2,
        help='Minimum number of RNA fragments supporting alternate allele (default: 2)')

    parser.add_argument(
        '--min-variant-sequence-coverage',
        type=int,
        default=1,
        help='Minimum coverage at variant position (default: 1)')

    parser.add_argument(
        '--min-mapping-quality',
        type=int,
        default=1,
        help='Minimum mapping quality (default: 1)')

    parser.add_argument(
        '--min-base-quality',
        type=int,
        default=1,
        help='Minimum base quality (default: 1)')

    parser.add_argument(
        '--use-duplicate-reads',
        action='store_true',
        help='Use duplicate reads')

    parser.add_argument(
        '--use-secondary-alignments',
        action='store_true',
        help='Use secondary alignments')

    parser.add_argument(
        '--genome',
        default=None,
        help='Genome reference (e.g., GRCh37, GRCh38)')

    parser.add_argument(
        '--log-level',
        default='INFO',
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
        help='Logging level (default: INFO)')

    args = parser.parse_args()

    # Configure logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    # Run conversion
    sequences_written = vcf_bam_to_fasta(
        vcf_paths=args.vcf_paths,
        bam_path=args.bam,
        output_fasta_path=args.output_fasta_path,
        protein_sequence_length=args.protein_sequence_length,
        centered_sequence_length=args.centered_sequence_length,
        min_alt_rna_reads=args.min_alt_rna_reads,
        min_alt_rna_fragments=args.min_alt_rna_fragments,
        min_variant_sequence_coverage=args.min_variant_sequence_coverage,
        min_mapping_quality=args.min_mapping_quality,
        use_duplicate_reads=args.use_duplicate_reads,
        use_secondary_alignments=args.use_secondary_alignments,
        min_base_quality=args.min_base_quality,
        genome=args.genome)

    print("Successfully wrote %d protein sequences to %s" % (
        sequences_written, args.output_fasta_path))


if __name__ == '__main__':
    main()

