#!/usr/bin/env python3
"""
Optimized version of dsRNAscan with improved parallelization
IMMEDIATE FIX: Change ThreadPoolExecutor to ProcessPoolExecutor for RNAduplex
"""

import RNA
import pandas as pd
from concurrent.futures import ProcessPoolExecutor, as_completed
import numpy as np

def process_rnaduplex_row_parallel(args):
    """
    Process a single einverted result through RNAduplex.
    Moved outside class for pickling with ProcessPoolExecutor.
    """
    row, temperature = args

    # Set temperature for this process
    RNA.cvar.temperature = temperature

    try:
        result = RNA.duplexfold(row['i_seq'], row['j_seq'])

        structure_parts = result.structure.split('&')
        if len(structure_parts) == 2:
            len1 = len(structure_parts[0])
            len2 = len(structure_parts[1])

            return {
                'idx': row.name,  # Preserve index
                'structure': result.structure,
                'energy': result.energy,
                'i_trim_start': result.i - len1 + 1,
                'i_trim_end': result.i,
                'j_trim_start': result.j,
                'j_trim_end': result.j + len2 - 1
            }
        else:
            return {
                'idx': row.name,
                'structure': None,
                'energy': None,
                'i_trim_start': None,
                'i_trim_end': None,
                'j_trim_start': None,
                'j_trim_end': None
            }
    except Exception as e:
        return {
            'idx': row.name,
            'structure': None,
            'energy': None,
            'i_trim_start': None,
            'i_trim_end': None,
            'j_trim_start': None,
            'j_trim_end': None
        }

def run_rnaduplex_batch_optimized(einverted_df, temperature=37, max_workers=None, verbose=False):
    """
    Optimized RNAduplex processing using ProcessPoolExecutor instead of ThreadPoolExecutor
    """
    if einverted_df.empty:
        return einverted_df

    # Deduplicate identical einverted results before RNAduplex
    total_before = len(einverted_df)
    einverted_df = einverted_df.drop_duplicates(
        subset=['seq_hash', 'strand', 'i_start_local', 'i_end_local',
                'j_start_local', 'j_end_local', 'i_seq', 'j_seq'],
        keep='first'
    )
    total_after = len(einverted_df)

    if verbose:
        if total_before > total_after:
            print(f"    Deduplicated {total_before} → {total_after} unique inverted repeats")
        print(f"    Processing {total_after} inverted repeats through RNAduplex")

    # Apply RNAduplex processing with ProcessPoolExecutor for true parallelization
    if max_workers and max_workers > 1 and len(einverted_df) > 10:
        if verbose:
            print(f"    Using ProcessPoolExecutor with {max_workers} workers")

        # Prepare data for parallel processing
        rows_with_temp = [(row, temperature) for _, row in einverted_df.iterrows()]

        # Process with ProcessPoolExecutor for true parallelization (no GIL)
        results = []
        with ProcessPoolExecutor(max_workers=max_workers) as executor:
            futures = {executor.submit(process_rnaduplex_row_parallel, row_data): i
                      for i, row_data in enumerate(rows_with_temp)}

            for future in as_completed(futures):
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    print(f"Error in RNAduplex processing: {e}")

        # Convert results to DataFrame
        duplex_df = pd.DataFrame(results)
        duplex_df = duplex_df.set_index('idx').sort_index()

        # Merge results back
        for col in ['structure', 'energy', 'i_trim_start', 'i_trim_end', 'j_trim_start', 'j_trim_end']:
            einverted_df[col] = duplex_df[col]

    else:
        # For small datasets or single CPU, use the original approach
        if verbose:
            print(f"    Using serial processing (dataset too small or single CPU)")

        # Set temperature
        RNA.cvar.temperature = temperature

        for idx, row in einverted_df.iterrows():
            try:
                result = RNA.duplexfold(row['i_seq'], row['j_seq'])

                structure_parts = result.structure.split('&')
                if len(structure_parts) == 2:
                    len1 = len(structure_parts[0])
                    len2 = len(structure_parts[1])

                    einverted_df.at[idx, 'structure'] = result.structure
                    einverted_df.at[idx, 'energy'] = result.energy
                    einverted_df.at[idx, 'i_trim_start'] = result.i - len1 + 1
                    einverted_df.at[idx, 'i_trim_end'] = result.i
                    einverted_df.at[idx, 'j_trim_start'] = result.j
                    einverted_df.at[idx, 'j_trim_end'] = result.j + len2 - 1
            except:
                einverted_df.at[idx, 'structure'] = None
                einverted_df.at[idx, 'energy'] = None

    # Drop entries where RNAduplex failed
    einverted_df = einverted_df.dropna(subset=['structure'])

    return einverted_df

# Additional optimization: Vectorized structure analysis
def vectorized_structure_analysis(structures):
    """
    Vectorized calculation of structure properties
    Faster than using .apply()
    """
    if len(structures) == 0:
        return pd.DataFrame()

    # Convert to numpy array for faster processing
    struct_array = structures.values

    # Vectorized calculations
    base_pairs = np.array([s.count('(') if s else 0 for s in struct_array])
    total_lengths = np.array([len(s) - 1 if s else 0 for s in struct_array])  # -1 for '&'

    # Avoid division by zero
    total_lengths[total_lengths == 0] = 1

    percent_paired = np.round(base_pairs * 2 / total_lengths * 100, 2)

    return pd.DataFrame({
        'base_pairs': base_pairs,
        'percent_paired': percent_paired
    }, index=structures.index)

# Test function to compare performance
def test_rnaduplex_optimization():
    """
    Test function to compare ThreadPoolExecutor vs ProcessPoolExecutor performance
    """
    import time
    import random

    # Create test data
    test_data = []
    for i in range(100):
        seq1 = ''.join(random.choice('ACGU') for _ in range(50))
        seq2 = ''.join(random.choice('ACGU') for _ in range(50))
        test_data.append({
            'seq_hash': f'hash_{i}',
            'strand': '+',
            'i_start_local': i * 10,
            'i_end_local': i * 10 + 50,
            'j_start_local': i * 10 + 100,
            'j_end_local': i * 10 + 150,
            'i_seq': seq1,
            'j_seq': seq2
        })

    test_df = pd.DataFrame(test_data)

    print("Testing RNAduplex parallelization optimization...")
    print("=" * 60)

    # Test with different CPU counts
    for cpus in [1, 2, 4, 8]:
        print(f"\nTesting with {cpus} CPUs:")

        start = time.time()
        result = run_rnaduplex_batch_optimized(test_df.copy(), max_workers=cpus, verbose=True)
        runtime = time.time() - start

        valid_results = result.dropna(subset=['structure'])
        print(f"  Runtime: {runtime:.2f}s")
        print(f"  Valid results: {len(valid_results)}/{len(test_df)}")

    print("\n" + "=" * 60)
    print("ProcessPoolExecutor should show better scaling than ThreadPoolExecutor")

if __name__ == "__main__":
    test_rnaduplex_optimization()