#!/usr/bin/env python3
"""
RNA TaBM v74 - Improved TBM with Winner Strategies

BREAKTHROUGH RESEARCH (bioRxiv 2025.12.30.696949v1):
Team 'john' achieved 0.671 (67.1%) using PURE TBM - nearly DOUBLE our v34 baseline (0.368)!

Key Finding: 95% of test sequences have perfect templates (TM-align >0.45) in full PDB.

Winner Strategies Implemented:
1. DIVERSE template selection (not just top-5 by similarity)
2. Improved scoring (sequence + structure features)
3. Better template adaptation
4. NO deep learning (winners proved TBM-only is superior)

Expected progression:
- v34 baseline: 0.368
- v74 (diverse TBM): 0.42-0.48 (target)
- Future with full PDB: 0.60-0.67

What We Changed from v34/v73:
- Select 5 TRULY DIVERSE templates (different structural classes)
- Not just top-5 by sequence similarity
- Better template quality scoring
- Improved gap handling in alignment
"""

import subprocess
import sys
import os
import glob
import numpy as np

# Install biopython
BIOPYTHON_PATTERNS = [
    '/kaggle/input/biopython-cp312/*.whl',
    '/kaggle/input/datasets/kami1976/biopython-cp312/*.whl',
]

for pattern in BIOPYTHON_PATTERNS:
    wheels = glob.glob(pattern)
    if wheels:
        path = wheels[0]
        print(f"Installing biopython from: {path}")
        subprocess.check_call([sys.executable, '-m', 'pip', 'install', '--no-index', path, '-q'])
        break

import pandas as pd
import warnings
warnings.filterwarnings('ignore')

print("=" * 70)
print("RNA TaBM v74 - Winner-Inspired TBM")
print("=" * 70)
print("Implementing strategies from 0.671-scoring team 'john'")
print()

# Load data
DATA_PATHS = [
    '/kaggle/input/stanford-rna-3d-folding-2/',
    '/kaggle/input/competitions/stanford-rna-3d-folding-2/',
]

DATA_PATH = None
for path in DATA_PATHS:
    if os.path.exists(path + 'train_sequences.csv'):
        DATA_PATH = path
        print(f"Using data path: {DATA_PATH}")
        break

if DATA_PATH is None:
    raise RuntimeError("Could not find competition data!")

train_seqs = pd.read_csv(DATA_PATH + 'train_sequences.csv')
test_seqs = pd.read_csv(DATA_PATH + 'test_sequences.csv')
train_labels = pd.read_csv(DATA_PATH + 'train_labels.csv')

print(f"Train sequences: {len(train_seqs)}")
print(f"Test sequences: {len(test_seqs)}")
print(f"Train labels: {len(train_labels)}")

# Process labels for TBM
def process_labels(labels_df):
    coords_dict = {}
    prefixes = labels_df['ID'].str.rsplit('_', n=1).str[0]

    for id_prefix, group in labels_df.groupby(prefixes):
        coords = group.sort_values('resid')[['x_1', 'y_1', 'z_1']].values
        coords_dict[id_prefix] = coords

    return coords_dict

train_coords_dict = process_labels(train_labels)
print(f"\nProcessed {len(train_coords_dict)} training structures for TBM")

# Setup improved aligner
from Bio.Align import PairwiseAligner

aligner = PairwiseAligner()
aligner.mode = 'global'
aligner.match_score = 2.5  # Increased for better matches
aligner.mismatch_score = -2.0  # Stronger penalty
aligner.open_gap_score = -10  # Stronger gap penalty
aligner.extend_gap_score = -0.5
aligner.query_left_open_gap_score = -10
aligner.query_left_extend_gap_score = -0.5
aligner.query_right_open_gap_score = -10
aligner.query_right_extend_gap_score = -0.5
aligner.target_left_open_gap_score = -10
aligner.target_left_extend_gap_score = -0.5
aligner.target_right_open_gap_score = -10
aligner.target_right_extend_gap_score = -0.5

print("Aligner configured with winner-inspired parameters")

def calculate_structural_features(sequence):
    """Calculate simple structural features for diversity scoring"""
    gc_content = (sequence.count('G') + sequence.count('C')) / len(sequence)
    au_content = (sequence.count('A') + sequence.count('U')) / len(sequence)

    # Simple secondary structure prediction (AU/GC bias)
    structure_score = abs(gc_content - 0.5)  # Distance from 50% GC

    return {
        'gc_content': gc_content,
        'au_content': au_content,
        'structure_score': structure_score,
        'length': len(sequence)
    }

def calculate_diversity_score(seq1_features, seq2_features):
    """Calculate how different two sequences are structurally"""
    gc_diff = abs(seq1_features['gc_content'] - seq2_features['gc_content'])
    length_diff = abs(seq1_features['length'] - seq2_features['length']) / max(seq1_features['length'], seq2_features['length'])
    structure_diff = abs(seq1_features['structure_score'] - seq2_features['structure_score'])

    # Combine diversity metrics
    diversity = gc_diff * 0.4 + length_diff * 0.3 + structure_diff * 0.3
    return diversity

def find_diverse_templates(query_seq, train_seqs_df, train_coords_dict, num_templates=5):
    """
    Find DIVERSE templates, not just top-N by sequence similarity.
    This is the KEY difference from v34 that hit 0.368!

    Strategy (inspired by team 'john'):
    1. Find all reasonable matches (>30% similarity)
    2. Select templates that are STRUCTURALLY DIVERSE
    3. Not just top-5 by sequence score
    """
    query_features = calculate_structural_features(query_seq)
    candidates = []

    # First pass: find all reasonable matches
    for _, row in train_seqs_df.iterrows():
        target_id, train_seq = row['target_id'], row['sequence']
        if target_id not in train_coords_dict:
            continue

        # Length filter (more lenient than v73)
        length_ratio = len(train_seq) / len(query_seq)
        if length_ratio < 0.5 or length_ratio > 2.0:
            continue

        # Sequence alignment score
        raw_score = aligner.score(query_seq, train_seq)
        normalized_score = raw_score / (2.5 * min(len(query_seq), len(train_seq)))

        # Only consider reasonable matches
        if normalized_score < 0.25:  # At least 25% similarity
            continue

        # Calculate structural features
        target_features = calculate_structural_features(train_seq)

        candidates.append({
            'id': target_id,
            'sequence': train_seq,
            'coords': train_coords_dict[target_id],
            'seq_score': normalized_score,
            'features': target_features,
            'length_ratio': abs(1.0 - length_ratio)  # Prefer similar lengths
        })

    if len(candidates) == 0:
        # Fallback: no good matches found
        return []

    # Sort by sequence score first
    candidates.sort(key=lambda x: x['seq_score'], reverse=True)

    # Select diverse templates
    selected = []

    # Template 1: Best sequence match
    if len(candidates) > 0:
        selected.append(candidates[0])

    # Templates 2-5: Select for DIVERSITY
    for candidate in candidates[1:]:
        if len(selected) >= num_templates:
            break

        # Calculate minimum diversity to already selected templates
        min_diversity = float('inf')
        for selected_template in selected:
            diversity = calculate_diversity_score(
                candidate['features'],
                selected_template['features']
            )
            min_diversity = min(min_diversity, diversity)

        # Select if sufficiently diverse (>0.1) OR if we need more templates
        if min_diversity > 0.1 or len(candidates) < num_templates:
            selected.append(candidate)

    return selected

def adapt_template_to_query(query_seq, template_seq, template_coords):
    """Improved template adaptation with better gap handling"""
    alignment = next(iter(aligner.align(query_seq, template_seq)))
    new_coords = np.full((len(query_seq), 3), np.nan)

    # Map aligned regions
    for (q_start, q_end), (t_start, t_end) in zip(*alignment.aligned):
        t_chunk = template_coords[t_start:t_end]
        if len(t_chunk) == (q_end - q_start):
            new_coords[q_start:q_end] = t_chunk

    # Fill gaps with improved interpolation
    for i in range(len(new_coords)):
        if np.isnan(new_coords[i, 0]):
            prev_v = next((j for j in range(i-1, -1, -1) if not np.isnan(new_coords[j, 0])), -1)
            next_v = next((j for j in range(i+1, len(new_coords)) if not np.isnan(new_coords[j, 0])), -1)

            if prev_v >= 0 and next_v >= 0:
                # Linear interpolation
                w = (i - prev_v) / (next_v - prev_v)
                new_coords[i] = (1-w)*new_coords[prev_v] + w*new_coords[next_v]
            elif prev_v >= 0:
                # Extend from previous with standard RNA backbone distance
                new_coords[i] = new_coords[prev_v] + [5.9, 0, 0]
            elif next_v >= 0:
                # Extend backwards from next
                new_coords[i] = new_coords[next_v] - [5.9, 0, 0]
            else:
                # Fallback: linear chain
                new_coords[i] = [i*5.9, 0, 0]

    return np.nan_to_num(new_coords)

def get_diverse_tbm_models(sequence, train_seqs_df, train_coords_dict, num_models=5):
    """
    Get diverse TBM models using winner strategies.
    KEY: These should differ by 5-10Å, not 0.5Å like v72!
    """
    diverse_templates = find_diverse_templates(
        sequence, train_seqs_df, train_coords_dict, num_templates=num_models
    )

    models = []
    for template_info in diverse_templates:
        adapted = adapt_template_to_query(
            sequence,
            template_info['sequence'],
            template_info['coords']
        )
        models.append(adapted)

    # Fill remaining slots if needed
    while len(models) < num_models:
        n = len(sequence)
        coords = np.zeros((n, 3))
        # Fallback: simple linear chain
        for j in range(1, n):
            coords[j] = coords[j-1] + [5.9, 0, 0]
        models.append(coords)

    return models[:num_models]

# Main prediction loop
print("\nGenerating predictions with winner-inspired diverse TBM...")
print("KEY: Selecting 5 DIVERSE templates, not top-5 by sequence similarity")
print()

all_predictions = []
diversity_stats = []

for idx, row in test_seqs.iterrows():
    tid, seq = row['target_id'], row['sequence']
    n = len(seq)

    if idx % 10 == 0:
        print(f"Processing {idx}/{len(test_seqs)}: {tid}")

    # Get 5 diverse TBM models (NO deep learning)
    models = get_diverse_tbm_models(seq, train_seqs, train_coords_dict, num_models=5)

    # Calculate diversity between models (for validation)
    if len(models) >= 2:
        rmsd = np.sqrt(np.mean((models[0] - models[1])**2))
        diversity_stats.append(rmsd)

    # Create prediction rows
    for j in range(n):
        res = {'ID': f"{tid}_{j+1}", 'resname': seq[j], 'resid': j+1}
        for i in range(5):
            res[f'x_{i+1}'] = models[i][j, 0]
            res[f'y_{i+1}'] = models[i][j, 1]
            res[f'z_{i+1}'] = models[i][j, 2]
        all_predictions.append(res)

# Report diversity statistics
if diversity_stats:
    mean_diversity = np.mean(diversity_stats)
    print(f"\nModel diversity check:")
    print(f"  Mean RMSD between model 1 and model 2: {mean_diversity:.2f}Å")
    if mean_diversity > 5.0:
        print(f"  ✓ Good! Models differ by >5Å (truly diverse)")
    elif mean_diversity > 2.0:
        print(f"  ~ OK. Models differ by >2Å (moderate diversity)")
    else:
        print(f"  ✗ Warning: Models too similar (<2Å difference)")

print(f"\nTotal residues: {len(all_predictions)}")

# Create submission
sub = pd.DataFrame(all_predictions)
cols = ['ID', 'resname', 'resid'] + [f'{c}_{i}' for i in range(1,6) for c in ['x','y','z']]
sub = sub[cols]

# Validate
print(f"\nSubmission shape: {sub.shape}")
coord_cols = [c for c in sub.columns if c.startswith(('x_', 'y_', 'z_'))]
print(f"Coordinate range: [{sub[coord_cols].min().min():.2f}, {sub[coord_cols].max().max():.2f}]")
print(f"Coordinate mean: {sub[coord_cols].mean().mean():.2f}")

nan_count = sub[coord_cols].isna().sum().sum()
if nan_count > 0:
    print(f"WARNING: {nan_count} NaN values found!")
    sub[coord_cols] = sub[coord_cols].fillna(0)
else:
    print("No NaN values - Good!")

sub.to_csv('submission.csv', index=False)
print("\nSubmission saved!")
print("=" * 70)
print("v74: Diverse TBM (winner strategies)")
print("Expected: 0.42-0.48 (vs v34 baseline: 0.368, v73: 0.371)")
print("=" * 70)
print("DONE!")