from concurrent.futures import ProcessPoolExecutor
import pyfaidx
import ganz
from typing import Tuple, Dict
from pathlib import Path
import os
import time
import numpy as np
from multiprocessing import shared_memory

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

def load_sequences_to_shared():
    """Load all chromosome sequences into shared memory."""
    genome = pyfaidx.Fasta(str(GENOME_PATH), rebuild=False, mutable=False, 
                          sequence_always_upper=True)
    
    # Create dict to store metadata about shared sequences
    shared_data = {}
    
    for i in range(1, 23):
        chrom = f'chr{i}'
        seq = genome[chrom][:].seq
        # Convert sequence to bytes
        seq_bytes = seq.encode('ascii')
        # Create shared memory block
        shm = shared_memory.SharedMemory(create=True, size=len(seq_bytes))
        # Copy data into shared memory
        np.frombuffer(shm.buf, dtype='b')[0:len(seq_bytes)] = np.frombuffer(seq_bytes, dtype='b')
        # Store metadata
        shared_data[chrom] = {
            'name': shm.name,
            'size': len(seq_bytes)
        }
    
    return shared_data

def convert_to_dna_sequence(chrom_data: Tuple[str, dict]) -> Tuple[str, ganz.DnaSequence]:
    """Convert shared memory sequence to DnaSequence object."""
    chrom, metadata = chrom_data
    # Access shared memory
    shm = shared_memory.SharedMemory(name=metadata['name'])
    # Get sequence bytes
    seq_bytes = bytes(np.frombuffer(shm.buf, dtype='b')[0:metadata['size']])
    # Convert to string and then DnaSequence
    seq = seq_bytes.decode('ascii')
    dna_seq = ganz.DnaSequence(seq)
    # Clean up
    shm.close()
    return chrom, dna_seq

def load_genome_parallel(n_processes: int = None) -> Dict[str, ganz.DnaSequence]:
    """Load genome sequences using shared memory and parallel conversion."""
    if n_processes is None:
        n_processes = os.cpu_count()
    n_processes = min(n_processes, 8)
    
    try:
        # First load all sequences into shared memory
        shared_data = load_sequences_to_shared()
        
        # Convert sequences to DnaSequence objects in parallel
        with ProcessPoolExecutor(max_workers=n_processes) as executor:
            futures = [
                executor.submit(convert_to_dna_sequence, (chrom, metadata))
                for chrom, metadata in shared_data.items()
            ]
            
            # Collect results
            sequences = {}
            for future in futures:
                chrom, dna_seq = future.result()
                sequences[chrom] = dna_seq
                
        return sequences
    
    finally:
        # Clean up shared memory
        for metadata in shared_data.values():
            try:
                shm = shared_memory.SharedMemory(name=metadata['name'])
                shm.close()
                shm.unlink()
            except:
                pass

if __name__ == '__main__':
    start = time.time()
    sequence = load_genome_parallel()
    print(f"Loaded {len(sequence)} chromosomes in {time.time() - start:.2f} seconds")
