"""
================================================================================
   CAFA-6 V5 - LIGHTWEIGHT WITH GO PROPAGATION

   Based on 0.230 kernel but with GO propagation added.
   Target: 0.35+ (current 0.231, baseline 0.230)
================================================================================
"""

# Install biopython
import subprocess
subprocess.run(['pip', 'install', 'biopython', '-q'], check=True)

import os
import gc
import numpy as np
import pandas as pd
from collections import defaultdict
from tqdm.auto import tqdm
import warnings
warnings.filterwarnings('ignore')

from sklearn.decomposition import PCA
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier

# ============================================================================
# CONFIG
# ============================================================================
CONFIG = {
    'PCA_COMPONENTS': 150,      # Slightly more than baseline 100
    'MIN_CONFIDENCE': 0.01,    # Lower threshold for more predictions
    'MAX_PREDS_PER_PROTEIN': 1500,
    'BASE_PATH': '/kaggle/input/cafa-6-protein-function-prediction',
    'EMB_PATH': '/kaggle/input/esm2-t6-8m-ur50d-cafa6',
}

print(f"CAFA-6 V5 - Lightweight with GO Propagation")
print(f"PCA components: {CONFIG['PCA_COMPONENTS']}")

# ============================================================================
# GO ONTOLOGY
# ============================================================================
def parse_go_obo(obo_path):
    """Parse GO ontology for parent relationships."""
    go_parents = defaultdict(set)
    current_id = None

    with open(obo_path, 'r') as f:
        for line in f:
            line = line.strip()
            if line == '[Term]':
                current_id = None
            elif line.startswith('id: GO:'):
                current_id = line[4:]
            elif line.startswith('is_a:') and current_id:
                parent = line.split()[1]
                go_parents[current_id].add(parent)
            elif line.startswith('relationship: part_of') and current_id:
                parent = line.split()[2]
                go_parents[current_id].add(parent)

    return go_parents

def get_all_ancestors(go_term, go_parents, cache=None):
    """Get all ancestors of a GO term."""
    if cache is None:
        cache = {}
    if go_term in cache:
        return cache[go_term]

    ancestors = set()
    for parent in go_parents.get(go_term, set()):
        ancestors.add(parent)
        ancestors.update(get_all_ancestors(parent, go_parents, cache))

    cache[go_term] = ancestors
    return ancestors

def propagate_predictions(predictions_df, go_parents):
    """Apply GO propagation - if child has score, parent gets max(parent, child)."""
    print("Applying GO propagation...")

    # Build ancestors cache
    cache = {}
    all_terms = predictions_df['term'].unique()
    for term in tqdm(all_terms, desc="Building ancestors"):
        get_all_ancestors(term, go_parents, cache)

    # Propagate
    protein_preds = defaultdict(dict)

    for _, row in tqdm(predictions_df.iterrows(), total=len(predictions_df), desc="Propagating"):
        protein = row['protein']
        term = row['term']
        score = row['score']

        # Set score for term
        if term not in protein_preds[protein] or score > protein_preds[protein][term]:
            protein_preds[protein][term] = score

        # Propagate to ancestors
        for ancestor in cache.get(term, set()):
            if ancestor not in protein_preds[protein] or score > protein_preds[protein][ancestor]:
                protein_preds[protein][ancestor] = score

    # Convert to DataFrame
    rows = []
    for protein, term_scores in protein_preds.items():
        for term, score in term_scores.items():
            if score >= CONFIG['MIN_CONFIDENCE']:
                rows.append({'protein': protein, 'term': term, 'score': score})

    result = pd.DataFrame(rows)
    print(f"After propagation: {len(result):,} predictions")
    return result

# ============================================================================
# MAIN
# ============================================================================
def main():
    # 1. Load GO ontology
    print("\n[1/6] Loading GO ontology...")
    go_parents = parse_go_obo(f"{CONFIG['BASE_PATH']}/Train/go-basic.obo")
    print(f"GO terms: {len(go_parents):,}")

    # 2. Load sequences and IDs
    print("\n[2/6] Loading sequence IDs...")
    from Bio import SeqIO

    train_seqs = {rec.id: str(rec.seq) for rec in SeqIO.parse(
        f"{CONFIG['BASE_PATH']}/Train/train_sequences.fasta", 'fasta')}
    test_seqs = {rec.id: str(rec.seq) for rec in SeqIO.parse(
        f"{CONFIG['BASE_PATH']}/Test/testsuperset.fasta", 'fasta')}

    train_ids = [i.split('|')[1] for i in train_seqs.keys()]
    test_ids = list(test_seqs.keys())

    print(f"Train: {len(train_ids)}, Test: {len(test_ids)}")

    # 3. Load embeddings and apply PCA
    print("\n[3/6] Loading embeddings and applying PCA...")
    X_train = np.load(f"{CONFIG['EMB_PATH']}/protein_embeddings_train.npy")
    X_test = np.load(f"{CONFIG['EMB_PATH']}/protein_embeddings_test.npy")
    print(f"Train shape: {X_train.shape}, Test shape: {X_test.shape}")

    pca = PCA(n_components=CONFIG['PCA_COMPONENTS'], random_state=42)
    X_train_pca = pca.fit_transform(X_train)
    X_test_pca = pca.transform(X_test)
    print(f"After PCA: {X_train_pca.shape}")

    del X_train, X_test
    gc.collect()

    # 4. Train models
    print("\n[4/6] Training models...")
    train_terms_df = pd.read_csv(
        f"{CONFIG['BASE_PATH']}/Train/train_terms.tsv", sep="\t")

    mlb_dict = {}
    models = {}

    for aspect in tqdm(['P', 'C', 'F'], desc="Training"):
        ont_terms = train_terms_df[train_terms_df['aspect'] == aspect]
        protein_terms = ont_terms.groupby('EntryID')['term'].apply(list).to_dict()
        labels = [protein_terms.get(eid, []) for eid in train_ids]

        mlb = MultiLabelBinarizer(sparse_output=True)
        y_train = mlb.fit_transform(labels)
        print(f"  {aspect}: {y_train.shape[1]} classes")

        model = OneVsRestClassifier(
            LogisticRegression(max_iter=800, solver='lbfgs', C=0.5, random_state=42),
            n_jobs=-1
        )
        model.fit(X_train_pca, y_train)

        mlb_dict[aspect] = mlb
        models[aspect] = model

    del X_train_pca
    gc.collect()

    # 5. Generate predictions
    print("\n[5/6] Generating predictions...")
    BATCH_SIZE = 5000
    submission_list = []

    for i in tqdm(range(0, len(test_ids), BATCH_SIZE), desc="Predicting"):
        batch_ids = test_ids[i:i+BATCH_SIZE]
        X_batch = X_test_pca[i:i+BATCH_SIZE]

        for aspect, model in models.items():
            mlb = mlb_dict[aspect]
            y_pred = model.predict_proba(X_batch)

            for j, eid in enumerate(batch_ids):
                probs = y_pred[j]
                # Lower threshold to get more predictions for propagation
                indices = np.where(probs > CONFIG['MIN_CONFIDENCE'])[0]
                for idx in indices:
                    submission_list.append({
                        'protein': eid,
                        'term': mlb.classes_[idx],
                        'score': probs[idx]
                    })

    del X_test_pca
    gc.collect()

    predictions_df = pd.DataFrame(submission_list)
    print(f"Raw predictions: {len(predictions_df):,}")

    # 6. Apply GO propagation and save
    print("\n[6/6] Applying GO propagation and saving...")
    predictions_df = propagate_predictions(predictions_df, go_parents)

    # Aggregate and clean
    predictions_df = predictions_df.groupby(['protein', 'term'], as_index=False)['score'].max()
    predictions_df['score'] = predictions_df['score'].clip(0.001, 0.999).round(3)

    # Keep top predictions per protein
    predictions_df = predictions_df.sort_values(['protein', 'score'], ascending=[True, False])
    predictions_df = predictions_df.groupby('protein').head(CONFIG['MAX_PREDS_PER_PROTEIN'])

    predictions_df.to_csv('submission.tsv', sep='\t', header=False, index=False)

    print(f"\n{'='*60}")
    print("CAFA-6 V5 COMPLETE!")
    print(f"{'='*60}")
    print(f"Total: {len(predictions_df):,} predictions")
    print(f"Proteins: {predictions_df['protein'].nunique():,}")
    print(f"Terms: {predictions_df['term'].nunique():,}")
    print(predictions_df.head(10))

if __name__ == "__main__":
    main()