#!/usr/bin/env python3
"""
Chapter 13: Python for Next Gen Biologist
Script: ch13_hpc_parallel_genomics.py
Description: High-Performance Computing & Population Genetics Simulation:
             Parallel multi-core processing of sequence chunks using multiprocessing,
             and Wright-Fisher population genetics simulation of genetic drift.
Usage:
    python ch13_hpc_parallel_genomics.py
"""

import multiprocessing as mp
import random
import time

def process_sequence_chunk(chunk_id, sequences):
    """Worker function: processes batch of sequences in parallel."""
    gc_results = []
    for s in sequences:
        s_up = s.upper()
        gc = (s_up.count('G') + s_up.count('C')) / len(s_up) if s_up else 0
        gc_results.append(round(gc * 100, 2))
    return chunk_id, len(sequences), sum(gc_results) / len(gc_results)

def wright_fisher_drift(pop_size=100, initial_p=0.5, generations=50):
    """
    Simulates genetic drift of a neutral biallelic locus over generations.
    p = frequency of allele A.
    """
    p = initial_p
    trajectory = [p]
    
    for gen in range(generations):
        # Sample 2N alleles binomially
        successes = 0
        for _ in range(2 * pop_size):
            if random.random() < p:
                successes += 1
        p = successes / (2 * pop_size)
        trajectory.append(p)
        if p == 0.0 or p == 1.0:
            # Fixation or loss
            break
    return trajectory

def main():
    print("=== Chapter 13: HPC & Simulation in Biology ===")
    
    # 1. Parallel Computing Demo
    cores = min(mp.cpu_count(), 4)
    print(f"\n[1] Parallel Multi-core Batch Processing (Detected {mp.cpu_count()} CPU cores, using {cores}):")
    
    mock_seqs = [
        "ATGCGATCGATCGATCGATC" * 50 for _ in range(1000)
    ]
    chunk_size = len(mock_seqs) // cores
    chunks = [mock_seqs[i:i+chunk_size] for i in range(0, len(mock_seqs), chunk_size)]
    
    start_t = time.time()
    with mp.Pool(processes=cores) as pool:
        args = [(i, chunks[i]) for i in range(len(chunks))]
        results = pool.starmap(process_sequence_chunk, args)
        
    elapsed = time.time() - start_t
    print(f"    - Processed {len(mock_seqs)} sequences in {elapsed:.4f} seconds.")
    for cid, n, mean_gc in results:
        print(f"    - Worker {cid}: {n} seqs | Mean GC: {mean_gc:.2f}%")

    # 2. Wright-Fisher Simulation
    print(f"\n[2] Wright-Fisher Genetic Drift Simulation (N=50 individuals, initial freq=0.5):")
    traj = wright_fisher_drift(pop_size=50, initial_p=0.5, generations=15)
    for g, freq in enumerate(traj[:10]):
        print(f"    Gen {g:02d}: Allele 'A' Frequency = {freq:.3f} [{'#' * int(freq * 30)}]")

if __name__ == "__main__":
    main()