from concurrent.futures import ProcessPoolExecutor
import pyfaidx
import ganz
from typing import Tuple, Dict
from pathlib import Path
import os

GENOME_PATH = Path("/Users/danila/pd/hematopoiesis/data/ref/GRCh37 genome.fa")

def load_chromosome(chrom: int) -> Tuple[str, str]:
    """Load a single chromosome sequence.
    
    Returns raw string sequence instead of DnaSequence object
    """
    genome = pyfaidx.Fasta(str(GENOME_PATH))
    chrom_name = f'chr{chrom}'
    return chrom_name, genome[chrom_name][:].seq

def load_genome_parallel(n_processes: int = None) -> Dict[str, ganz.DnaSequence]:
    """Load genome sequences in parallel."""
    if n_processes is None:
        n_processes = os.cpu_count()
        
    # First load raw sequences in parallel
    with ProcessPoolExecutor(max_workers=n_processes) as executor:
        futures = [executor.submit(load_chromosome, i) for i in range(1, 23)]
        raw_sequences = [future.result() for future in futures]
    
    # Then convert to DnaSequence objects in the main process
    return {
        chrom: ganz.DnaSequence(seq) 
        for chrom, seq in raw_sequences
    }

if __name__ == '__main__':
    sequence = load_genome_parallel()
    print(f"Loaded {len(sequence)} chromosomes")
