import pysam
import numpy as np
import argparse
import sys

def create_strand_coverage_maps(bam_file, reference_length=None):
    """
    Create coverage maps for positive and negative strands from a BAM file.
    
    Args:
        bam_file: Path to the BAM file. Note: The BAM file must be indexed
                  if reference_length is not provided.
        reference_length: Length of reference sequence (if None, will be inferred
                          from the BAM header).
    
    Returns:
        pos_coverage: Coverage array for positive strand
        neg_coverage: Coverage array for negative strand
    """
    # Open the BAM file in binary read mode ("rb")
    samfile = pysam.AlignmentFile(bam_file, "rb")
    
    # Get reference length if not provided
    if reference_length is None:
        if samfile.header and 'SQ' in samfile.header and samfile.header['SQ']:
            # Use the length of the first reference sequence
            reference_length = samfile.header['SQ'][0]['LN']
        else:
            # This fallback is less likely for BAM but included for robustness
            print("Warning: No header or 'SQ' tag found, scanning file to determine reference length...", file=sys.stderr)
            max_pos = 0
            # Re-iterate from the beginning
            for read in samfile.fetch(until_eof=True):
                if not read.is_unmapped and read.reference_end is not None:
                    max_pos = max(max_pos, read.reference_end)
            reference_length = max_pos
            # Reset the iterator to the beginning of the file for the main loop
            samfile.close()
            samfile = pysam.AlignmentFile(bam_file, "rb")
    
    # Initialize coverage arrays
    pos_coverage = np.zeros(reference_length, dtype=np.int32)
    neg_coverage = np.zeros(reference_length, dtype=np.int32)
    
    # Process each read using fetch(until_eof=True) for efficiency
    for read in samfile.fetch(until_eof=True):
        if read.is_unmapped:
            continue
        
        positions = read.get_reference_positions(full_length=True)
        
        if read.is_reverse:
            for pos in positions:
                if pos is not None and pos < reference_length:
                    neg_coverage[pos] += 1
        else:
            for pos in positions:
                if pos is not None and pos < reference_length:
                    pos_coverage[pos] += 1
    
    samfile.close()
    
    return pos_coverage, neg_coverage

# Main execution block
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Create strand-specific coverage maps from a BAM file and save them as NumPy arrays."
    )
    parser.add_argument(
        "bam_file", 
        help="Path to the input BAM file. The file should be indexed."
    )
    args = parser.parse_args()
    
    input_file = args.bam_file

    try:
        # Process the BAM file provided as an argument
        print(f"Processing {input_file}...")
        pos_coverage, neg_coverage = create_strand_coverage_maps(input_file)
        
        # Save the coverage maps using the original, fixed filenames
        np.save("positive_strand_coverage.npy", pos_coverage)
        np.save("negative_strand_coverage.npy", neg_coverage)
        
        print(f"Coverage maps created successfully:")
        print(f"  - positive_strand_coverage.npy: shape {pos_coverage.shape}, max coverage {pos_coverage.max()}")
        print(f"  - negative_strand_coverage.npy: shape {neg_coverage.shape}, max coverage {neg_coverage.max()}")
        
        # Optional: Print some basic statistics
        print(f"\nStatistics:")
        print(f"  Reference length inferred: {len(pos_coverage)}")
        print(f"  Positive strand: mean coverage = {pos_coverage.mean():.2f}")
        print(f"  Negative strand: mean coverage = {neg_coverage.mean():.2f}")
        print(f"  Total bases covered (positive): {np.sum(pos_coverage > 0)}")
        print(f"  Total bases covered (negative): {np.sum(neg_coverage > 0)}")

    except FileNotFoundError:
        print(f"\nError: Input file '{input_file}' not found.", file=sys.stderr)
        sys.exit(1)
    except (ValueError, pysam.utils.SamtoolsError) as e:
        print(f"\nAn error occurred while processing the BAM file: {e}", file=sys.stderr)
        print("Please ensure the input file is a valid, sorted, and indexed BAM file.", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"\nAn unexpected error occurred: {e}", file=sys.stderr)
        sys.exit(1)

Modify so I can choose prefix for output