generate_test_data.py

#!/usr/bin/env python3
"""
Generate synthetic test data for kmer-diff development and testing.

Creates:
- A synthetic bacterial reference genome (100kb)
- Core genes FASTA for normalization
- Paired-end FASTQ files for multiple samples with known variants
- Ground truth variants file

Usage:
    python generate_test_data.py -o test_data/
"""

import argparse
import gzip
import math
import random
from dataclasses import dataclass
from pathlib import Path
from typing import List, Tuple, Optional


# =============================================================================
# Configuration
# =============================================================================

@dataclass
class Config:
    """Test data generation parameters."""
    # Genome
    genome_length: int = 100_000
    gc_content: float = 0.50
    
    # Core genes
    num_core_genes: int = 20
    core_gene_min_length: int = 500
    core_gene_max_length: int = 1000
    
    # Sequencing
    read_length: int = 150
    insert_size_mean: int = 400
    insert_size_std: int = 50
    insert_size_min: int = 300
    insert_size_max: int = 500
    coverage: int = 100
    
    # Error model
    error_rate_start: float = 0.001  # 0.1% at read start
    error_rate_end: float = 0.005    # 0.5% at read end
    
    # Random seeds
    seed_reference: int = 42
    seed_variants: int = 43
    seed_sample_a: int = 100
    seed_sample_b: int = 101
    seed_sample_c: int = 102


# =============================================================================
# Data structures
# =============================================================================

@dataclass
class Gene:
    """A core gene region."""
    name: str
    start: int
    end: int
    
    @property
    def length(self) -> int:
        return self.end - self.start


@dataclass
class SNP:
    """A single nucleotide polymorphism."""
    position: int
    ref: str
    alt: str


@dataclass
class StructuralVariant:
    """A structural variant (duplication or deletion)."""
    sv_type: str  # 'DUP' or 'DEL'
    start: int
    end: int
    
    @property
    def length(self) -> int:
        return self.end - self.start


@dataclass
class SampleVariants:
    """Variants for a sample."""
    name: str
    snps: List[SNP]
    structural_variants: List[StructuralVariant]


# =============================================================================
# Reference genome generation
# =============================================================================

def generate_reference(length: int, gc_content: float, seed: int) -> str:
    """Generate a random reference genome with specified GC content."""
    random.seed(seed)
    
    gc_bases = ['G', 'C']
    at_bases = ['A', 'T']
    
    sequence = []
    for _ in range(length):
        if random.random() < gc_content:
            sequence.append(random.choice(gc_bases))
        else:
            sequence.append(random.choice(at_bases))
    
    return ''.join(sequence)


def generate_core_genes(genome_length: int, num_genes: int, 
                        min_length: int, max_length: int, seed: int) -> List[Gene]:
    """Generate core gene positions distributed across the genome."""
    random.seed(seed)
    
    # Calculate spacing between genes
    # Leave buffer at start and end
    buffer = 1000
    available_length = genome_length - 2 * buffer
    spacing = available_length // num_genes
    
    genes = []
    for i in range(num_genes):
        # Calculate start position with some jitter
        base_start = buffer + i * spacing
        jitter = random.randint(-100, 100)
        start = max(buffer, base_start + jitter)
        
        # Random length
        length = random.randint(min_length, max_length)
        end = start + length
        
        # Ensure we don't exceed genome
        if end > genome_length - buffer:
            end = genome_length - buffer
            start = end - length
        
        genes.append(Gene(
            name=f"core_gene_{i+1:02d}",
            start=start,
            end=end
        ))
    
    return genes


# =============================================================================
# Variant generation
# =============================================================================

def generate_snp_positions(genome_length: int, num_snps: int, 
                           excluded_regions: List[Tuple[int, int]], seed: int) -> List[int]:
    """Generate random SNP positions avoiding excluded regions."""
    random.seed(seed)
    
    def is_excluded(pos: int) -> bool:
        for start, end in excluded_regions:
            if start <= pos < end:
                return True
        return False
    
    positions = []
    attempts = 0
    max_attempts = num_snps * 100
    
    while len(positions) < num_snps and attempts < max_attempts:
        pos = random.randint(0, genome_length - 1)
        if not is_excluded(pos) and pos not in positions:
            positions.append(pos)
        attempts += 1
    
    if len(positions) < num_snps:
        raise ValueError(f"Could only place {len(positions)} SNPs, requested {num_snps}")
    
    return sorted(positions)


def generate_snps(reference: str, positions: List[int], seed: int) -> List[SNP]:
    """Generate SNPs at specified positions."""
    random.seed(seed)
    
    bases = ['A', 'C', 'G', 'T']
    snps = []
    
    for pos in positions:
        ref = reference[pos]
        alt_choices = [b for b in bases if b != ref]
        alt = random.choice(alt_choices)
        snps.append(SNP(position=pos, ref=ref, alt=alt))
    
    return snps


def create_sample_b_variants(reference: str, core_genes: List[Gene], seed: int) -> SampleVariants:
    """Create variants for sample B."""
    random.seed(seed)
    
    # Excluded regions: core genes + buffer
    excluded = [(g.start - 100, g.end + 100) for g in core_genes]
    
    # Also exclude regions we'll use for structural variants
    dup_region = (40000, 42000)
    del_region = (70000, 71000)
    excluded.append((dup_region[0] - 500, dup_region[1] + 500))
    excluded.append((del_region[0] - 500, del_region[1] + 500))
    
    # Generate 10 SNPs
    snp_positions = generate_snp_positions(len(reference), 10, excluded, seed)
    snps = generate_snps(reference, snp_positions, seed + 1)
    
    # Structural variants
    structural_variants = [
        StructuralVariant('DUP', dup_region[0], dup_region[1]),
        StructuralVariant('DEL', del_region[0], del_region[1]),
    ]
    
    return SampleVariants(name='B', snps=snps, structural_variants=structural_variants)


def create_sample_c_variants(reference: str, core_genes: List[Gene], 
                             sample_b: SampleVariants, seed: int) -> SampleVariants:
    """Create variants for sample C (shares some SNPs with B)."""
    random.seed(seed)
    
    # Share first 2 SNPs with B
    shared_snps = sample_b.snps[:2]
    
    # Excluded regions
    excluded = [(g.start - 100, g.end + 100) for g in core_genes]
    dup_region = (50000, 51500)
    del_region = (82000, 82500)
    excluded.append((dup_region[0] - 500, dup_region[1] + 500))
    excluded.append((del_region[0] - 500, del_region[1] + 500))
    
    # Also exclude positions of shared SNPs and B's unique SNPs to avoid overlap
    for snp in sample_b.snps:
        excluded.append((snp.position - 50, snp.position + 50))
    
    # Generate 6 unique SNPs for C
    unique_positions = generate_snp_positions(len(reference), 6, excluded, seed)
    unique_snps = generate_snps(reference, unique_positions, seed + 1)
    
    # Combine shared and unique
    all_snps = sorted(shared_snps + unique_snps, key=lambda s: s.position)
    
    structural_variants = [
        StructuralVariant('DUP', dup_region[0], dup_region[1]),
        StructuralVariant('DEL', del_region[0], del_region[1]),
    ]
    
    return SampleVariants(name='C', snps=all_snps, structural_variants=structural_variants)


def apply_variants(reference: str, variants: SampleVariants) -> str:
    """Apply variants to reference to create sample genome."""
    # Start with reference
    genome = list(reference)
    
    # Apply SNPs
    for snp in variants.snps:
        genome[snp.position] = snp.alt
    
    # Convert back to string
    genome_str = ''.join(genome)
    
    # Apply structural variants
    # Sort by position descending so deletions don't shift coordinates
    svs = sorted(variants.structural_variants, key=lambda sv: sv.start, reverse=True)
    
    for sv in svs:
        if sv.sv_type == 'DEL':
            # Delete the region
            genome_str = genome_str[:sv.start] + genome_str[sv.end:]
        elif sv.sv_type == 'DUP':
            # Insert a copy after the region (tandem duplication)
            duplicated = genome_str[sv.start:sv.end]
            genome_str = genome_str[:sv.end] + duplicated + genome_str[sv.end:]
    
    return genome_str


# =============================================================================
# Read simulation
# =============================================================================

def error_rate_at_position(position: int, read_length: int, config: Config) -> float:
    """Calculate error rate at a given position in the read."""
    fraction = position / read_length
    return config.error_rate_start + (config.error_rate_end - config.error_rate_start) * fraction


def quality_score(error_prob: float) -> int:
    """Convert error probability to Phred quality score."""
    if error_prob <= 0:
        return 40
    score = int(-10 * math.log10(error_prob))
    return min(40, max(0, score))


def introduce_errors(sequence: str, config: Config, seed: int) -> Tuple[str, str]:
    """Introduce sequencing errors and generate quality scores."""
    random.seed(seed)
    
    bases = ['A', 'C', 'G', 'T']
    result = list(sequence)
    qualities = []
    
    for i, base in enumerate(result):
        error_rate = error_rate_at_position(i, len(sequence), config)
        q = quality_score(error_rate)
        
        # Introduce error with probability = error_rate
        if random.random() < error_rate:
            # Substitution error
            alt_bases = [b for b in bases if b != base]
            result[i] = random.choice(alt_bases)
            # Lower quality for error positions (though in real data we wouldn't know)
            q = max(2, q - 10)
        
        qualities.append(chr(q + 33))  # Phred+33 encoding
    
    return ''.join(result), ''.join(qualities)


def reverse_complement(sequence: str) -> str:
    """Return reverse complement of sequence."""
    complement = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C', 'N': 'N'}
    return ''.join(complement.get(b, 'N') for b in reversed(sequence))


def simulate_read_pair(genome: str, config: Config, read_num: int, seed: int) -> Optional[Tuple[str, str, str, str]]:
    """Simulate a single paired-end read pair."""
    random.seed(seed)
    
    genome_len = len(genome)
    
    # Sample insert size
    insert_size = int(random.gauss(config.insert_size_mean, config.insert_size_std))
    insert_size = max(config.insert_size_min, min(config.insert_size_max, insert_size))
    
    # Ensure insert fits in genome
    if insert_size > genome_len:
        insert_size = genome_len
    
    # Random start position
    max_start = genome_len - insert_size
    if max_start < 0:
        return None
    
    start = random.randint(0, max_start)
    
    # Extract insert
    insert = genome[start:start + insert_size]
    
    # R1: forward read from start of insert
    r1_seq = insert[:config.read_length]
    if len(r1_seq) < config.read_length:
        r1_seq = r1_seq + 'N' * (config.read_length - len(r1_seq))
    
    # R2: reverse read from end of insert (reverse complement)
    r2_region = insert[-config.read_length:] if len(insert) >= config.read_length else insert
    r2_seq = reverse_complement(r2_region)
    if len(r2_seq) < config.read_length:
        r2_seq = r2_seq + 'N' * (config.read_length - len(r2_seq))
    
    # Introduce errors
    r1_seq, r1_qual = introduce_errors(r1_seq, config, seed + 1)
    r2_seq, r2_qual = introduce_errors(r2_seq, config, seed + 2)
    
    return r1_seq, r1_qual, r2_seq, r2_qual


def simulate_reads(genome: str, config: Config, sample_name: str, 
                   seed: int) -> Tuple[List[Tuple[str, str, str]], List[Tuple[str, str, str]]]:
    """Simulate all reads for a sample."""
    random.seed(seed)
    
    # Calculate number of read pairs needed
    total_bases_needed = len(genome) * config.coverage
    bases_per_pair = config.read_length * 2
    num_pairs = int(total_bases_needed / bases_per_pair) + 100  # slight oversampling
    
    r1_reads = []
    r2_reads = []
    
    for i in range(num_pairs):
        read_seed = seed + i * 10
        result = simulate_read_pair(genome, config, i, read_seed)
        
        if result is None:
            continue
        
        r1_seq, r1_qual, r2_seq, r2_qual = result
        read_name = f"read_{i+1:06d}"
        
        r1_reads.append((f"{read_name}/1", r1_seq, r1_qual))
        r2_reads.append((f"{read_name}/2", r2_seq, r2_qual))
    
    return r1_reads, r2_reads


# =============================================================================
# File output
# =============================================================================

def write_fasta(path: Path, sequences: List[Tuple[str, str]], line_width: int = 80):
    """Write sequences to FASTA file."""
    with open(path, 'w') as f:
        for name, seq in sequences:
            f.write(f">{name}\n")
            for i in range(0, len(seq), line_width):
                f.write(seq[i:i+line_width] + "\n")


def write_fastq_gz(path: Path, reads: List[Tuple[str, str, str]]):
    """Write reads to gzipped FASTQ file."""
    with gzip.open(path, 'wt') as f:
        for name, seq, qual in reads:
            f.write(f"@{name}\n")
            f.write(f"{seq}\n")
            f.write("+\n")
            f.write(f"{qual}\n")


def write_variants_tsv(path: Path, samples: List[SampleVariants]):
    """Write ground truth variants to TSV file."""
    with open(path, 'w') as f:
        # Header
        f.write("sample\ttype\tchrom\tstart\tend\tref\talt\tnotes\n")
        
        for sample in samples:
            # SNPs
            for snp in sample.snps:
                f.write(f"{sample.name}\tSNP\tchr1\t{snp.position}\t{snp.position}\t{snp.ref}\t{snp.alt}\t\n")
            
            # Structural variants
            for sv in sample.structural_variants:
                notes = f"length={sv.length}"
                f.write(f"{sample.name}\t{sv.sv_type}\tchr1\t{sv.start}\t{sv.end}\t.\t.\t{notes}\n")


# =============================================================================
# Main
# =============================================================================

def main():
    parser = argparse.ArgumentParser(description="Generate synthetic test data for kmer-diff")
    parser.add_argument('-o', '--output', type=Path, required=True, help="Output directory")
    parser.add_argument('--seed', type=int, default=42, help="Base random seed")
    args = parser.parse_args()
    
    # Create output directory
    args.output.mkdir(parents=True, exist_ok=True)
    
    # Configuration
    config = Config()
    
    print("Generating test data...")
    
    # Generate reference genome
    print("  Creating reference genome...")
    reference = generate_reference(config.genome_length, config.gc_content, config.seed_reference)
    
    # Generate core genes
    print("  Defining core genes...")
    core_genes = generate_core_genes(
        config.genome_length, 
        config.num_core_genes,
        config.core_gene_min_length,
        config.core_gene_max_length,
        config.seed_reference + 1
    )
    
    # Write reference FASTA
    write_fasta(
        args.output / "reference.fasta",
        [("reference length=100000", reference)]
    )
    
    # Write core genes FASTA
    core_gene_seqs = [
        (f"{g.name} start={g.start} end={g.end}", reference[g.start:g.end])
        for g in core_genes
    ]
    write_fasta(args.output / "core_genes.fasta", core_gene_seqs)
    
    # Generate variants
    print("  Generating variants...")
    sample_b_variants = create_sample_b_variants(reference, core_genes, config.seed_variants)
    sample_c_variants = create_sample_c_variants(reference, core_genes, sample_b_variants, config.seed_variants + 100)
    
    # Write variants file
    write_variants_tsv(args.output / "variants.tsv", [sample_b_variants, sample_c_variants])
    
    # Create sample genomes
    genome_a = reference  # No variants
    genome_b = apply_variants(reference, sample_b_variants)
    genome_c = apply_variants(reference, sample_c_variants)
    
    # Simulate reads for each sample
    print("  Simulating reads for sample A...")
    r1_a, r2_a = simulate_reads(genome_a, config, "A", config.seed_sample_a)
    write_fastq_gz(args.output / "sample_a_R1.fastq.gz", r1_a)
    write_fastq_gz(args.output / "sample_a_R2.fastq.gz", r2_a)
    print(f"    Generated {len(r1_a)} read pairs")
    
    print("  Simulating reads for sample B...")
    r1_b, r2_b = simulate_reads(genome_b, config, "B", config.seed_sample_b)
    write_fastq_gz(args.output / "sample_b_R1.fastq.gz", r1_b)
    write_fastq_gz(args.output / "sample_b_R2.fastq.gz", r2_b)
    print(f"    Generated {len(r1_b)} read pairs")
    
    print("  Simulating reads for sample C...")
    r1_c, r2_c = simulate_reads(genome_c, config, "C", config.seed_sample_c)
    write_fastq_gz(args.output / "sample_c_R1.fastq.gz", r1_c)
    write_fastq_gz(args.output / "sample_c_R2.fastq.gz", r2_c)
    print(f"    Generated {len(r1_c)} read pairs")
    
    print("\nTest data generation complete!")
    print(f"\nOutput files in {args.output}/:")
    print("  reference.fasta       - 100kb synthetic genome")
    print("  core_genes.fasta      - 20 single-copy genes for normalization")
    print("  sample_a_*.fastq.gz   - Baseline reads (no variants)")
    print("  sample_b_*.fastq.gz   - Reads with 10 SNPs, 1 dup, 1 del")
    print("  sample_c_*.fastq.gz   - Reads with 8 SNPs (2 shared), 1 dup, 1 del")
    print("  variants.tsv          - Ground truth variant positions")
    
    # Print summary of variants
    print("\nSample B variants:")
    print(f"  SNPs: {len(sample_b_variants.snps)}")
    for snp in sample_b_variants.snps:
        print(f"    pos {snp.position}: {snp.ref}>{snp.alt}")
    for sv in sample_b_variants.structural_variants:
        print(f"  {sv.sv_type}: {sv.start}-{sv.end} ({sv.length} bp)")
    
    print("\nSample C variants:")
    print(f"  SNPs: {len(sample_c_variants.snps)}")
    for snp in sample_c_variants.snps:
        print(f"    pos {snp.position}: {snp.ref}>{snp.alt}")
    for sv in sample_c_variants.structural_variants:
        print(f"  {sv.sv_type}: {sv.start}-{sv.end} ({sv.length} bp)")


if __name__ == "__main__":
    main()