#!/usr/bin/env python3
"""
Generate synthetic SAM alignments for teloclip integration testing.

This script creates realistic SAM alignment records that correspond to the
synthetic contigs generated by test_generate_data.py. The alignments simulate
various scenarios that teloclip extend command needs to handle.

SAM Record Types Generated:
1. Normal alignments within contig boundaries
2. Overhang alignments extending past contig ends
3. High-coverage regions (multiple reads per position)
4. One-sided overhangs (left or right end only)
5. Homopolymer sequences in overhangs
6. Break alignments with complex CIGAR strings
7. Control alignments with no overhangs

Each alignment includes:
- Proper SAM headers with contig references
- Realistic CIGAR strings matching alignment patterns
- Quality scores and mapping information
- Coordinated read names for tracking
"""

import argparse
from pathlib import Path
import random
from typing import Dict, List, Optional


class SAMAlignment:
    """Represent a SAM alignment record with all required fields."""

    def __init__(
        self,
        qname: str,
        flag: int,
        rname: str,
        pos: int,
        mapq: int,
        cigar: str,
        rnext: str,
        pnext: int,
        tlen: int,
        seq: str,
        qual: str,
        tags: Optional[Dict[str, str]] = None,
    ):
        self.qname = qname
        self.flag = flag
        self.rname = rname
        self.pos = pos
        self.mapq = mapq
        self.cigar = cigar
        self.rnext = rnext
        self.pnext = pnext
        self.tlen = tlen
        self.seq = seq
        self.qual = qual
        self.tags = tags or {}

    def to_sam_line(self) -> str:
        """Convert alignment to SAM format line."""
        tag_str = '\t'.join(f'{key}:{value}' for key, value in self.tags.items())
        base_fields = [
            self.qname,
            str(self.flag),
            self.rname,
            str(self.pos),
            str(self.mapq),
            self.cigar,
            self.rnext,
            str(self.pnext),
            str(self.tlen),
            self.seq,
            self.qual,
        ]

        if tag_str:
            return '\t'.join(base_fields) + '\t' + tag_str
        else:
            return '\t'.join(base_fields)


def generate_random_sequence(length: int, gc_content: float = 0.5) -> str:
    """Generate random DNA sequence with specified GC content."""
    bases = ['A', 'T', 'G', 'C']

    # Calculate base frequencies
    gc_bases = int(length * gc_content)
    at_bases = length - gc_bases

    # Create base pool
    sequence_bases = (
        ['G'] * (gc_bases // 2)
        + ['C'] * (gc_bases // 2)
        + ['A'] * (at_bases // 2)
        + ['T'] * (at_bases // 2)
    )

    # Add remaining bases if odd numbers
    while len(sequence_bases) < length:
        sequence_bases.append(random.choice(bases))

    # Shuffle to randomize order
    random.shuffle(sequence_bases)

    return ''.join(sequence_bases)


def generate_quality_string(length: int, avg_quality: int = 30) -> str:
    """Generate quality scores string for SAM record."""
    # Quality scores range from 0-40, convert to ASCII offset 33
    qualities = []
    for _ in range(length):
        # Add some variation around average quality
        qual_score = max(0, min(40, avg_quality + random.randint(-5, 5)))
        qualities.append(chr(qual_score + 33))

    return ''.join(qualities)


def get_contig_lengths() -> Dict[str, int]:
    """Return the lengths of synthetic contigs (must match test_generate_data.py)."""
    return {
        'contig_1': 1102,  # 1000 + motifs + buffers
        'contig_2': 2136,  # 2000 + motifs + buffers
        'contig_3': 1624,  # 1500 + left motif + buffers
        'contig_4': 916,  # 800 + motifs + buffers
        'contig_5': 1354,  # 1200 + motifs + buffers
        'contig_6': 900,  # Exactly 900 (no motifs)
        'contig_7': 1216,  # 1100 + motifs + buffers
    }


def generate_normal_alignments(
    contig: str, contig_length: int, read_count: int = 5
) -> List[SAMAlignment]:
    """Generate normal alignments completely within contig boundaries."""
    alignments = []
    read_length = 150

    for i in range(read_count):
        # Position read within contig boundaries with safe margins
        start_pos = random.randint(100, max(100, contig_length - read_length - 100))

        # Generate read sequence and quality
        read_seq = generate_random_sequence(read_length)
        quality = generate_quality_string(read_length)

        # Create alignment
        alignment = SAMAlignment(
            qname=f'normal_read_{contig}_{i + 1}',
            flag=0,  # Primary alignment, forward strand
            rname=contig,
            pos=start_pos,
            mapq=60,  # High mapping quality
            cigar=f'{read_length}M',  # Perfect match
            rnext='*',
            pnext=0,
            tlen=0,
            seq=read_seq,
            qual=quality,
            tags={
                'NM:i': '0',
                'AS:i': str(read_length * 2),
            },  # Perfect alignment scores
        )
        alignments.append(alignment)

    return alignments


def generate_overhang_alignments(
    contig: str, contig_length: int, side: str = 'both'
) -> List[SAMAlignment]:
    """Generate alignments that extend past contig ends (overhangs)."""
    alignments = []
    read_length = 150
    overhang_length = 50  # How much extends past contig

    sides_to_generate = []
    if side == 'both':
        sides_to_generate = ['left', 'right']
    elif side in ['left', 'right']:
        sides_to_generate = [side]

    for side_name in sides_to_generate:
        for i in range(3):  # Generate 3 reads per side
            if side_name == 'left':
                # Read starts before contig and extends into it
                start_pos = (
                    -overhang_length + 1
                )  # SAM is 1-based, negative means before start
                # CIGAR: soft clip for overhang + match for contig part
                cigar = f'{overhang_length}S{read_length - overhang_length}M'

                # Generate sequence: overhang part + contig part
                overhang_seq = 'TTAGGG' * (overhang_length // 6 + 1)  # Telomeric motif
                overhang_seq = overhang_seq[:overhang_length]
                contig_seq = generate_random_sequence(read_length - overhang_length)
                read_seq = overhang_seq + contig_seq

            else:  # right side
                # Read starts within contig and extends past end
                start_pos = contig_length - (read_length - overhang_length) + 1
                # CIGAR: match for contig part + soft clip for overhang
                cigar = f'{read_length - overhang_length}M{overhang_length}S'

                # Generate sequence: contig part + overhang part
                contig_seq = generate_random_sequence(read_length - overhang_length)
                overhang_seq = 'CCCTAA' * (overhang_length // 6 + 1)  # Telomeric motif
                overhang_seq = overhang_seq[:overhang_length]
                read_seq = contig_seq + overhang_seq

            quality = generate_quality_string(read_length)

            alignment = SAMAlignment(
                qname=f'overhang_{side_name}_{contig}_{i + 1}',
                flag=0,
                rname=contig,
                pos=max(1, start_pos),  # SAM positions must be >= 1
                mapq=50,
                cigar=cigar,
                rnext='*',
                pnext=0,
                tlen=0,
                seq=read_seq,
                qual=quality,
                tags={'NM:i': '1', 'AS:i': str((read_length - overhang_length) * 2)},
            )
            alignments.append(alignment)

    return alignments


def generate_high_coverage_alignments(
    contig: str, contig_length: int, coverage_multiplier: int = 10
) -> List[SAMAlignment]:
    """Generate high coverage alignments for specific regions."""
    alignments = []
    read_length = 150

    # Focus on both ends of the contig for high coverage
    target_regions = [
        (50, 200),  # Near start
        (max(200, contig_length - 200), contig_length - 50),  # Near end
    ]

    for region_idx, (start_region, end_region) in enumerate(target_regions):
        for i in range(coverage_multiplier):
            # Random position within target region
            start_pos = random.randint(
                start_region, max(start_region, end_region - read_length)
            )

            read_seq = generate_random_sequence(read_length)
            quality = generate_quality_string(read_length)

            alignment = SAMAlignment(
                qname=f'high_cov_region{region_idx + 1}_{contig}_{i + 1}',
                flag=0,
                rname=contig,
                pos=start_pos,
                mapq=55,
                cigar=f'{read_length}M',
                rnext='*',
                pnext=0,
                tlen=0,
                seq=read_seq,
                qual=quality,
                tags={'NM:i': '0', 'AS:i': str(read_length * 2)},
            )
            alignments.append(alignment)

    return alignments


def generate_homopolymer_alignments(contig: str) -> List[SAMAlignment]:
    """Generate alignments with homopolymer sequences in overhangs."""
    alignments = []
    read_length = 150

    for i in range(2):  # Two homopolymer reads
        # Create read with homopolymer overhang
        overhang_length = 40
        homopolymer_base = random.choice(['A', 'T'])  # Common in telomeres

        # Read extending past left end with homopolymer
        start_pos = 1  # At contig start

        # CIGAR: soft clip for homopolymer overhang + match
        cigar = f'{overhang_length}S{read_length - overhang_length}M'

        # Sequence: homopolymer + random contig sequence
        homopolymer_seq = homopolymer_base * overhang_length
        contig_seq = generate_random_sequence(read_length - overhang_length)
        read_seq = homopolymer_seq + contig_seq

        quality = generate_quality_string(
            read_length, avg_quality=25
        )  # Lower quality for homopolymers

        alignment = SAMAlignment(
            qname=f'homopolymer_{homopolymer_base}_{contig}_{i + 1}',
            flag=0,
            rname=contig,
            pos=start_pos,
            mapq=40,  # Lower mapping quality due to repetitive sequence
            cigar=cigar,
            rnext='*',
            pnext=0,
            tlen=0,
            seq=read_seq,
            qual=quality,
            tags={'NM:i': '2', 'AS:i': str((read_length - overhang_length) * 2 - 4)},
        )
        alignments.append(alignment)

    return alignments


def generate_break_alignments(contig: str, contig_length: int) -> List[SAMAlignment]:
    """Generate alignments with complex CIGAR strings (insertions, deletions)."""
    alignments = []
    read_length = 150

    for i in range(2):  # Two complex alignments
        start_pos = random.randint(50, contig_length - read_length)

        # Create complex CIGAR with matches, insertions, deletions
        if i == 0:
            # Match + Insertion + Match + Soft clip
            cigar = '75M5I50M20S'
            match_len = 75 + 50  # Total matching bases
            insertion_len = 5
            total_seq_len = match_len + insertion_len + 20  # +20 for soft clip
        else:
            # Soft clip + Match + Deletion + Match
            cigar = '25S60M10D40M'
            match_len = 60 + 40
            total_seq_len = 25 + match_len  # No sequence for deletion

        read_seq = generate_random_sequence(total_seq_len)
        quality = generate_quality_string(total_seq_len, avg_quality=28)

        alignment = SAMAlignment(
            qname=f'complex_cigar_{contig}_{i + 1}',
            flag=0,
            rname=contig,
            pos=start_pos,
            mapq=45,
            cigar=cigar,
            rnext='*',
            pnext=0,
            tlen=0,
            seq=read_seq,
            qual=quality,
            tags={'NM:i': '8', 'AS:i': str(match_len * 2 - 16)},  # Penalize for indels
        )
        alignments.append(alignment)

    return alignments


def generate_all_alignments() -> List[SAMAlignment]:
    """Generate comprehensive set of test alignments for all contigs."""
    all_alignments = []
    contig_lengths = get_contig_lengths()

    print('Generating alignments for each contig...')

    for contig_name, contig_length in contig_lengths.items():
        print(f'  Processing {contig_name} (length: {contig_length})')

        # contig_1: Normal contig with moderate overhangs
        if contig_name == 'contig_1':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 5)
            )
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'both')
            )

        # contig_2: High-overhang contig (10x more alignments)
        elif contig_name == 'contig_2':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 3)
            )
            all_alignments.extend(
                generate_high_coverage_alignments(contig_name, contig_length, 10)
            )
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'both')
            )

        # contig_3: One-sided high overhangs (left end only)
        elif contig_name == 'contig_3':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 3)
            )
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'left')
            )
            all_alignments.extend(
                generate_high_coverage_alignments(contig_name, contig_length, 5)
            )

        # contig_4: Homopolymer test contig
        elif contig_name == 'contig_4':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 3)
            )
            all_alignments.extend(generate_homopolymer_alignments(contig_name))
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'both')
            )

        # contig_5: Motif-rich contig
        elif contig_name == 'contig_5':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 4)
            )
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'both')
            )
            # Extra overhangs to test motif detection
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'both')
            )

        # contig_6: Control (no overhangs)
        elif contig_name == 'contig_6':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 8)
            )
            # No overhang alignments for control

        # contig_7: Break alignment test contig
        elif contig_name == 'contig_7':
            all_alignments.extend(
                generate_normal_alignments(contig_name, contig_length, 3)
            )
            all_alignments.extend(generate_break_alignments(contig_name, contig_length))
            all_alignments.extend(
                generate_overhang_alignments(contig_name, contig_length, 'both')
            )

    print(f'Generated {len(all_alignments)} total alignments')
    return all_alignments


def write_sam_file(
    alignments: List[SAMAlignment], output_path: Path, contig_lengths: Dict[str, int]
) -> None:
    """Write alignments to SAM format file with proper headers."""
    with open(output_path, 'w') as f:
        # Write SAM header
        f.write('@HD\tVN:1.6\tSO:coordinate\n')

        # Write sequence headers for each contig
        for contig_name, length in sorted(contig_lengths.items()):
            f.write(f'@SQ\tSN:{contig_name}\tLN:{length}\n')

        # Write program header
        f.write(
            '@PG\tID:test_generate_alignments\tVN:1.0\tCL:test_generate_alignments.py\n'
        )

        # Sort alignments by reference name and position
        sorted_alignments = sorted(alignments, key=lambda x: (x.rname, x.pos))

        # Write alignment records
        for alignment in sorted_alignments:
            f.write(alignment.to_sam_line() + '\n')

    print(f'Wrote {len(alignments)} alignments to {output_path}')


def main():
    """Main function to generate synthetic SAM alignments."""
    parser = argparse.ArgumentParser(
        description='Generate synthetic SAM alignments for teloclip integration testing'
    )
    parser.add_argument(
        '--output-dir',
        type=Path,
        default=Path('./test_data'),
        help='Output directory for SAM file (default: ./test_data)',
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Random seed for reproducibility (default: 42)',
    )

    args = parser.parse_args()

    # Set random seed for reproducibility
    random.seed(args.seed)

    # Ensure output directory exists
    args.output_dir.mkdir(parents=True, exist_ok=True)

    # Generate alignments
    print('Generating synthetic SAM alignments...')
    alignments = generate_all_alignments()

    # Get contig lengths (must match the FASTA file)
    contig_lengths = get_contig_lengths()

    # Write SAM file
    sam_path = args.output_dir / 'synthetic_alignments.sam'
    write_sam_file(alignments, sam_path, contig_lengths)

    print('\nAlignment generation completed successfully!')
    print(f'SAM file: {sam_path}')
    print('\nAlignment Summary by Contig:')
    print('Contig\t\tAlignments\tTypes')
    print('-' * 60)

    # Count alignments per contig
    alignment_counts = {}
    alignment_types = {}

    for alignment in alignments:
        contig = alignment.rname
        alignment_counts[contig] = alignment_counts.get(contig, 0) + 1

        # Determine alignment type from read name
        if 'normal' in alignment.qname:
            type_name = 'normal'
        elif 'overhang' in alignment.qname:
            type_name = 'overhang'
        elif 'high_cov' in alignment.qname:
            type_name = 'high_cov'
        elif 'homopolymer' in alignment.qname:
            type_name = 'homopolymer'
        elif 'complex' in alignment.qname:
            type_name = 'complex'
        else:
            type_name = 'other'

        if contig not in alignment_types:
            alignment_types[contig] = {}
        alignment_types[contig][type_name] = (
            alignment_types[contig].get(type_name, 0) + 1
        )

    for contig in sorted(alignment_counts.keys()):
        count = alignment_counts[contig]
        types = ', '.join(
            f'{k}:{v}' for k, v in sorted(alignment_types[contig].items())
        )
        print(f'{contig}\t{count}\t\t{types}')

    print('\nNext steps:')
    print(
        '1. Convert to BAM: samtools view -bS synthetic_alignments.sam > synthetic_alignments.bam'
    )
    print(
        '2. Sort BAM: samtools sort synthetic_alignments.bam -o synthetic_alignments_sorted.bam'
    )
    print('3. Index BAM: samtools index synthetic_alignments_sorted.bam')
    print('4. Run integration tests with teloclip extend')


if __name__ == '__main__':
    main()
