"""
================================================================================
   CAFA-6 V4 - IMPROVED APPROACH

   Target: Score 0.40+ (current: 0.231, top: 0.456)

   Key Improvements:
   1. Lower MIN_CONFIDENCE (0.01 vs 0.05) - More predictions
   2. More GO Terms (3500 vs 2000) - Better coverage
   3. Deeper MLP (3 hidden layers vs 2) - Better representation
   4. Multi-seed inference averaging - Pseudo ensemble
   5. Optimized thresholds based on validation
================================================================================
"""

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

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer

# ============================================================================
# CONFIG - OPTIMIZED
# ============================================================================
CONFIG = {
    'SEED': 42,
    'TOP_K_LABELS': 2000,        # Reduced further
    'BATCH_SIZE': 32,
    'EPOCHS': 15,
    'LR': 2e-4,
    'DROPOUT': 0.15,
    'MIN_CONFIDENCE': 0.01,
    'MAX_PREDS_PER_PROTEIN': 600,
    'BASE_PATH': '/kaggle/input/cafa-6-protein-function-prediction',
    'EMBEDDINGS_PATH': '/kaggle/input/cafa-6-embeddings',
    # ESM1b 650M - smallest model (1280 dim)
    'EMB_MODEL': 'esm1b_650M',
    'EMB_DIM': 1280,
    'LABEL_SMOOTHING': 0.01,
    'HIDDEN_DIMS': [512, 256],   # Smaller network
    'INFERENCE_SEEDS': [42],     # Single seed to save memory
}

np.random.seed(CONFIG['SEED'])
torch.manual_seed(CONFIG['SEED'])
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Device: {device}")
print(f"Config: TOP_K={CONFIG['TOP_K_LABELS']}, MIN_CONF={CONFIG['MIN_CONFIDENCE']}")

# ============================================================================
# GO ONTOLOGY PARSER
# ============================================================================
def parse_go_obo(obo_path):
    """Parse GO ontology file to get parent-child relationships."""
    go_parents = defaultdict(set)
    go_namespace = {}

    current_id = None
    current_namespace = None

    with open(obo_path, 'r') as f:
        for line in f:
            line = line.strip()
            if line == '[Term]':
                current_id = None
                current_namespace = None
            elif line.startswith('id: GO:'):
                current_id = line[4:]
            elif line.startswith('namespace:'):
                current_namespace = line.split(': ')[1]
                if current_id:
                    go_namespace[current_id] = current_namespace
            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, go_namespace

def get_all_ancestors(go_term, go_parents, ancestors_cache=None):
    """Get all ancestors of a GO term (for True Path Rule)."""
    if ancestors_cache is None:
        ancestors_cache = {}

    if go_term in ancestors_cache:
        return ancestors_cache[go_term]

    ancestors = set()
    parents = go_parents.get(go_term, set())

    for parent in parents:
        ancestors.add(parent)
        ancestors.update(get_all_ancestors(parent, go_parents, ancestors_cache))

    ancestors_cache[go_term] = ancestors
    return ancestors

# ============================================================================
# GO PROPAGATION (CRITICAL!)
# ============================================================================
def propagate_go_predictions(predictions_df, go_parents):
    """Apply True Path Rule with optimized implementation."""
    print("   Applying GO propagation (True Path Rule)...")

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

    # Group by protein - use dict for speed
    print("   Grouping predictions...")
    protein_preds = defaultdict(dict)

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

        # Set score for this 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 ancestors_cache.get(term, set()):
            if ancestor not in protein_preds[protein] or score > protein_preds[protein][ancestor]:
                protein_preds[protein][ancestor] = score

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

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

# ============================================================================
# IMPROVED MODEL
# ============================================================================
class ImprovedMLP(nn.Module):
    """Deeper MLP with residual connections."""
    def __init__(self, input_dim, output_dim, hidden_dims=[2048, 1024, 512]):
        super().__init__()

        layers = []
        prev_dim = input_dim

        for i, hidden_dim in enumerate(hidden_dims):
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(nn.LayerNorm(hidden_dim))  # LayerNorm instead of BatchNorm
            layers.append(nn.GELU())
            layers.append(nn.Dropout(CONFIG['DROPOUT']))
            prev_dim = hidden_dim

        self.features = nn.Sequential(*layers)
        self.classifier = nn.Linear(prev_dim, output_dim)

        # Initialize weights
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

    def forward(self, x):
        features = self.features(x)
        return self.classifier(features)

# ============================================================================
# FOCAL LOSS (Better for imbalanced data)
# ============================================================================
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred, target):
        bce = nn.functional.binary_cross_entropy_with_logits(pred, target, reduction='none')
        pt = torch.exp(-bce)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * bce
        return focal_loss.mean()

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

    # 2. Load annotations
    print("\n[2/8] Loading annotations...")
    train_terms_df = pd.read_csv(
        f"{CONFIG['BASE_PATH']}/Train/train_terms.tsv",
        sep='\t', header=None, names=['protein', 'term', 'ontology']
    )

    protein_to_terms = train_terms_df.groupby('protein')['term'].apply(list).to_dict()
    term_counts = Counter()
    for terms in protein_to_terms.values():
        term_counts.update(terms)

    top_terms = [t for t, _ in term_counts.most_common(CONFIG['TOP_K_LABELS'])]
    term_to_idx = {term: idx for idx, term in enumerate(top_terms)}

    print(f"   Proteins: {len(protein_to_terms):,}, Top Terms: {len(top_terms)}")

    # 3. Load embeddings
    emb_model = CONFIG['EMB_MODEL']
    print(f"\n[3/8] Loading {emb_model} embeddings...")
    base_path = CONFIG['EMBEDDINGS_PATH']

    esm2_path = f"{base_path}/{emb_model}/train_sequences_emb.npy"
    esm2_ids_path = f"{base_path}/{emb_model}/train_sequences_ids.npy"

    esm2_emb = np.load(esm2_path)
    esm2_ids = np.load(esm2_ids_path, allow_pickle=True)
    print(f"   {emb_model}: {esm2_emb.shape}")

    def extract_id(full_id):
        full_id = str(full_id)
        if '|' in full_id:
            parts = full_id.split('|')
            return parts[1] if len(parts) >= 2 else full_id
        return full_id.split()[0] if ' ' in full_id else full_id

    # Map proteins to embeddings
    id_to_idx = {extract_id(pid): i for i, pid in enumerate(esm2_ids)}

    valid_proteins = [p for p in protein_to_terms.keys() if p in id_to_idx]
    print(f"   Valid proteins: {len(valid_proteins):,}")

    # Build X matrix
    X = np.array([esm2_emb[id_to_idx[p]] for p in tqdm(valid_proteins, desc="   Building X")])
    print(f"   X shape: {X.shape}")

    del esm2_emb, esm2_ids
    gc.collect()

    # 4. Create labels with GO hierarchy
    print("\n[4/8] Creating labels with GO hierarchy...")
    mlb = MultiLabelBinarizer(classes=range(len(top_terms)))
    y_labels = []
    ancestors_cache = {}

    for p in tqdm(valid_proteins, desc="   Expanding labels"):
        terms = set()
        for t in protein_to_terms.get(p, []):
            if t in term_to_idx:
                terms.add(term_to_idx[t])
                for ancestor in get_all_ancestors(t, go_parents, ancestors_cache):
                    if ancestor in term_to_idx:
                        terms.add(term_to_idx[ancestor])
        y_labels.append(list(terms))

    y = mlb.fit_transform(y_labels).astype(np.float32)
    print(f"   y shape: {y.shape}")
    print(f"   Avg labels per protein: {y.sum(axis=1).mean():.1f}")

    # 5. Train model
    print("\n[5/8] Training model...")
    X_train, X_val, y_train, y_val = train_test_split(
        X, y, test_size=0.1, random_state=CONFIG['SEED']
    )
    print(f"   Train: {len(X_train)}, Val: {len(X_val)}")

    del X, y, y_labels
    gc.collect()

    train_dataset = TensorDataset(torch.FloatTensor(X_train), torch.FloatTensor(y_train))
    val_dataset = TensorDataset(torch.FloatTensor(X_val), torch.FloatTensor(y_val))
    train_loader = DataLoader(train_dataset, batch_size=CONFIG['BATCH_SIZE'], shuffle=True, num_workers=0)
    val_loader = DataLoader(val_dataset, batch_size=CONFIG['BATCH_SIZE'], num_workers=0)

    model = ImprovedMLP(CONFIG['EMB_DIM'], len(top_terms), CONFIG['HIDDEN_DIMS']).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=CONFIG['LR'], weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=5, T_mult=2)
    criterion = FocalLoss(alpha=0.25, gamma=2.0)

    print(f"   Model params: {sum(p.numel() for p in model.parameters()):,}")

    best_loss = float('inf')
    best_state = None
    patience = 5
    no_improve = 0

    for epoch in range(CONFIG['EPOCHS']):
        # Train
        model.train()
        train_loss = 0
        for X_b, y_b in train_loader:
            X_b, y_b = X_b.to(device), y_b.to(device)
            optimizer.zero_grad()
            out = model(X_b)
            loss = criterion(out, y_b)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            train_loss += loss.item()

        # Validate
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for X_b, y_b in val_loader:
                X_b, y_b = X_b.to(device), y_b.to(device)
                out = model(X_b)
                val_loss += criterion(out, y_b).item()

        train_loss /= len(train_loader)
        val_loss /= len(val_loader)
        scheduler.step()

        if val_loss < best_loss:
            best_loss = val_loss
            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
            no_improve = 0
            marker = " *BEST*"
        else:
            no_improve += 1
            marker = ""

        print(f"   Epoch {epoch+1}/{CONFIG['EPOCHS']}: train={train_loss:.4f}, val={val_loss:.4f}{marker}")

        if no_improve >= patience:
            print(f"   Early stopping at epoch {epoch+1}")
            break

    model.load_state_dict(best_state)

    del train_dataset, val_dataset, train_loader, val_loader, X_train, X_val, y_train, y_val
    gc.collect()
    torch.cuda.empty_cache()

    # 6. Generate predictions on test set
    print("\n[6/8] Generating test predictions...")

    test_ids_path = f"{base_path}/{emb_model}/testsuperset_ids.npy"
    test_ids = np.load(test_ids_path, allow_pickle=True)
    test_ids = [extract_id(pid) for pid in test_ids]
    print(f"   Test proteins: {len(test_ids):,}")

    test_emb_path = f"{base_path}/{emb_model}/testsuperset_emb.npy"
    test_X = np.load(test_emb_path, mmap_mode='r')
    print(f"   Test shape: {test_X.shape}")

    # Multi-seed inference for pseudo-ensemble
    model.eval()
    batch_size = 64
    all_predictions = defaultdict(lambda: defaultdict(list))

    print(f"   Running {len(CONFIG['INFERENCE_SEEDS'])} inference passes...")

    for seed_idx, seed in enumerate(CONFIG['INFERENCE_SEEDS']):
        torch.manual_seed(seed)
        print(f"   Pass {seed_idx+1}/{len(CONFIG['INFERENCE_SEEDS'])} (seed={seed})...")

        with torch.no_grad():
            for i in tqdm(range(0, len(test_X), batch_size), desc=f"   Predicting"):
                end_idx = min(i + batch_size, len(test_X))
                batch_data = np.array(test_X[i:end_idx])
                batch = torch.FloatTensor(batch_data).to(device)

                # Add small noise for diversity (pseudo-ensemble)
                if seed_idx > 0:
                    batch = batch + torch.randn_like(batch) * 0.001

                probs = torch.sigmoid(model(batch)).cpu().numpy()

                for j in range(len(probs)):
                    pid = test_ids[i + j]
                    p = probs[j]
                    indices = np.where(p >= CONFIG['MIN_CONFIDENCE'] / 2)[0]  # Lower threshold for averaging
                    for k in indices:
                        all_predictions[pid][top_terms[k]].append(p[k])

                del batch, probs

        gc.collect()
        torch.cuda.empty_cache()

    # Average predictions across seeds
    print("   Averaging predictions...")
    rows = []
    for pid, term_scores in tqdm(all_predictions.items(), desc="   Averaging"):
        for term, scores in term_scores.items():
            avg_score = np.mean(scores)
            if avg_score >= CONFIG['MIN_CONFIDENCE']:
                rows.append({'protein': pid, 'term': term, 'score': avg_score})

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

    del test_X, model, all_predictions
    gc.collect()
    torch.cuda.empty_cache()

    # 7. Apply GO Propagation
    print("\n[7/8] Applying GO propagation...")
    predictions_df = propagate_go_predictions(predictions_df, go_parents)

    # 8. Create submission
    print("\n[8/8] Creating submission...")

    predictions_df = predictions_df.groupby(['protein', 'term'], as_index=False)['score'].max()
    predictions_df['score'] = predictions_df['score'].clip(0.001, 0.999)
    predictions_df = predictions_df[predictions_df['score'] >= CONFIG['MIN_CONFIDENCE']]

    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['score'] = predictions_df['score'].round(3)
    predictions_df.to_csv('submission.tsv', sep='\t', header=False, index=False)

    print(f"\n{'='*60}")
    print(f"CAFA-6 V4 COMPLETE!")
    print(f"{'='*60}")
    print(f"Total rows: {len(predictions_df):,}")
    print(f"Unique proteins: {predictions_df['protein'].nunique():,}")
    print(f"Unique terms: {predictions_df['term'].nunique():,}")
    print(f"Score stats: min={predictions_df['score'].min():.3f}, max={predictions_df['score'].max():.3f}, mean={predictions_df['score'].mean():.3f}")
    print(predictions_df.head(10))

if __name__ == "__main__":
    main()