#!/usr/bin/env python3
"""
CAFA 6 - ESM2 Fine-Tuning for GO Term Prediction
================================================
Fine-tunes ESM2-650M on CAFA training data for multi-label GO term classification.

Run on Kaggle with GPU (P100 or T4) enabled.
Estimated training time: 4-6 hours.

Data sources required:
- cafa-6-protein-function-prediction (competition data)
- facebook/esm2_t33_650M_UR50D (ESM2 model from Kaggle datasets)
"""

import os
import gc
import random
import numpy as np
import pandas as pd
from pathlib import Path
from collections import defaultdict

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.cuda.amp import autocast, GradScaler

from transformers import AutoTokenizer, EsmModel, get_linear_schedule_with_warmup
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer
from tqdm.auto import tqdm

# =====================================================
# CONFIG
# =====================================================

class CFG:
    # Paths
    input_dir = '/kaggle/input/cafa-6-protein-function-prediction'
    esm_model_path = '/kaggle/input/cafa6-facebook-esm2-t33-650m-ur50d-train/esm2_t33_650M_UR50D'
    output_dir = '/kaggle/working'
    
    # Model
    model_name = 'facebook/esm2_t33_650M_UR50D'
    max_length = 1024  # Max protein sequence length
    hidden_dim = 512
    dropout = 0.3
    
    # Training
    epochs = 3
    batch_size = 4
    gradient_accumulation = 8
    learning_rate = 2e-5
    weight_decay = 0.01
    warmup_ratio = 0.1
    
    # GO terms
    min_term_count = 50  # Only train on GO terms with >= 50 examples
    
    # Device
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    seed = 42
    
    # Mixed precision
    use_amp = True


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


# =====================================================
# DATA LOADING
# =====================================================

def load_fasta(fasta_path):
    """Load sequences from FASTA file."""
    sequences = {}
    current_id = None
    current_seq = []
    
    with open(fasta_path, 'r') as f:
        for line in f:
            line = line.strip()
            if line.startswith('>'):
                if current_id:
                    sequences[current_id] = ''.join(current_seq)
                # Extract UniProt ID from header (e.g., "sp|A0A0C5B5G6|..." -> "A0A0C5B5G6")
                header = line[1:].split()[0]
                if '|' in header:
                    current_id = header.split('|')[1]  # Get ID from sp|ID|NAME format
                else:
                    current_id = header
                current_seq = []
            else:
                current_seq.append(line)
        if current_id:
            sequences[current_id] = ''.join(current_seq)
    
    return sequences


def load_annotations(terms_path, min_count=50):
    """Load GO annotations and filter rare terms."""
    df = pd.read_csv(terms_path, sep='\t')
    
    # Count term frequencies
    term_counts = df['term'].value_counts()
    valid_terms = term_counts[term_counts >= min_count].index.tolist()
    
    # Filter to valid terms
    df = df[df['term'].isin(valid_terms)]
    
    # Group by protein
    annotations = df.groupby('EntryID')['term'].apply(list).to_dict()
    
    print(f"Loaded {len(annotations)} proteins with {len(valid_terms)} GO terms")
    return annotations, valid_terms


class ProteinDataset(Dataset):
    """Dataset for protein sequences with GO annotations."""
    
    def __init__(self, protein_ids, sequences, annotations, tokenizer, mlb, max_length):
        self.protein_ids = protein_ids
        self.sequences = sequences
        self.annotations = annotations
        self.tokenizer = tokenizer
        self.mlb = mlb
        self.max_length = max_length
    
    def __len__(self):
        return len(self.protein_ids)
    
    def __getitem__(self, idx):
        protein_id = self.protein_ids[idx]
        sequence = self.sequences.get(protein_id, '')
        
        # Tokenize (add spaces between amino acids for ESM)
        sequence_spaced = ' '.join(list(sequence[:self.max_length]))
        
        encoding = self.tokenizer(
            sequence_spaced,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        
        # Get labels
        terms = self.annotations.get(protein_id, [])
        labels = self.mlb.transform([terms])[0]
        
        return {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0),
            'labels': torch.tensor(labels, dtype=torch.float32)
        }


# =====================================================
# MODEL
# =====================================================

class ESM2Classifier(nn.Module):
    """ESM2 with multi-label classification head."""
    
    def __init__(self, esm_model_path, num_labels, hidden_dim=512, dropout=0.3):
        super().__init__()
        
        # Load ESM2 (try local path first, then HuggingFace)
        if os.path.exists(esm_model_path):
            self.esm = EsmModel.from_pretrained(esm_model_path)
        else:
            self.esm = EsmModel.from_pretrained('facebook/esm2_t33_650M_UR50D')
        
        # Freeze early layers (keep last 3 trainable)
        for param in self.esm.parameters():
            param.requires_grad = False
        for layer in self.esm.encoder.layer[-3:]:
            for param in layer.parameters():
                param.requires_grad = True
        
        # Classification head
        esm_dim = self.esm.config.hidden_size  # 1280 for ESM2-650M
        self.classifier = nn.Sequential(
            nn.Linear(esm_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_labels)
        )
    
    def forward(self, input_ids, attention_mask):
        outputs = self.esm(input_ids=input_ids, attention_mask=attention_mask)
        
        # Mean pooling over sequence
        hidden_states = outputs.last_hidden_state
        mask = attention_mask.unsqueeze(-1).float()
        pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
        
        logits = self.classifier(pooled)
        return logits


# =====================================================
# TRAINING
# =====================================================

def train_epoch(model, dataloader, optimizer, scheduler, scaler, device, accumulation_steps):
    model.train()
    total_loss = 0
    optimizer.zero_grad()
    
    pbar = tqdm(dataloader, desc='Training')
    for step, batch in enumerate(pbar):
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        
        with autocast(enabled=CFG.use_amp):
            logits = model(input_ids, attention_mask)
            
            # Weighted BCE loss (upweight positive samples)
            pos_weight = torch.ones(labels.shape[1], device=device) * 10
            loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
            loss = loss_fn(logits, labels) / accumulation_steps
        
        scaler.scale(loss).backward()
        total_loss += loss.item() * accumulation_steps
        
        if (step + 1) % accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            optimizer.zero_grad()
        
        pbar.set_postfix({'loss': total_loss / (step + 1)})
    
    return total_loss / len(dataloader)


def validate(model, dataloader, device):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for batch in tqdm(dataloader, desc='Validating'):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels']
            
            with autocast(enabled=CFG.use_amp):
                logits = model(input_ids, attention_mask)
            
            probs = torch.sigmoid(logits).cpu().numpy()
            all_preds.append(probs)
            all_labels.append(labels.numpy())
    
    preds = np.vstack(all_preds)
    labels = np.vstack(all_labels)
    
    # Calculate F-max (simplified)
    best_f1 = 0
    for thresh in np.arange(0.1, 0.9, 0.05):
        binary_preds = (preds > thresh).astype(int)
        
        # Micro F1
        tp = (binary_preds * labels).sum()
        fp = (binary_preds * (1 - labels)).sum()
        fn = ((1 - binary_preds) * labels).sum()
        
        precision = tp / (tp + fp + 1e-9)
        recall = tp / (tp + fn + 1e-9)
        f1 = 2 * precision * recall / (precision + recall + 1e-9)
        
        if f1 > best_f1:
            best_f1 = f1
    
    return best_f1


# =====================================================
# MAIN
# =====================================================

def main():
    print("=" * 60)
    print("🧬 CAFA 6 - ESM2 Fine-Tuning")
    print("=" * 60)
    
    set_seed(CFG.seed)
    
    # Check GPU
    print(f"\n📍 Device: {CFG.device}")
    if torch.cuda.is_available():
        print(f"   GPU: {torch.cuda.get_device_name(0)}")
        print(f"   Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
    
    # Load data
    print("\n📂 Loading data...")
    sequences = load_fasta(f'{CFG.input_dir}/Train/train_sequences.fasta')
    annotations, go_terms = load_annotations(
        f'{CFG.input_dir}/Train/train_terms.tsv', 
        min_count=CFG.min_term_count
    )
    
    # Filter to proteins with both sequence and annotations
    valid_proteins = [p for p in sequences.keys() if p in annotations]
    print(f"Valid proteins: {len(valid_proteins)}")
    
    # Train/val split
    train_ids, val_ids = train_test_split(
        valid_proteins, test_size=0.1, random_state=CFG.seed
    )
    print(f"Train: {len(train_ids)}, Val: {len(val_ids)}")
    
    # Setup tokenizer and label encoder
    print("\n🔧 Setting up model...")
    # Tokenizer is small (~200KB) - download from HF
    # Model is large (2.6GB) - use local dataset
    tokenizer = AutoTokenizer.from_pretrained('facebook/esm2_t33_650M_UR50D')
    
    mlb = MultiLabelBinarizer(classes=go_terms)
    mlb.fit([go_terms])  # Fit on all terms
    
    num_labels = len(go_terms)
    print(f"Number of GO terms: {num_labels}")
    
    # Create datasets
    train_dataset = ProteinDataset(
        train_ids, sequences, annotations, tokenizer, mlb, CFG.max_length
    )
    val_dataset = ProteinDataset(
        val_ids, sequences, annotations, tokenizer, mlb, CFG.max_length
    )
    
    train_loader = DataLoader(
        train_dataset, batch_size=CFG.batch_size, shuffle=True, num_workers=0
    )
    val_loader = DataLoader(
        val_dataset, batch_size=CFG.batch_size * 2, shuffle=False, num_workers=0
    )
    
    # Create model
    model = ESM2Classifier(
        CFG.esm_model_path, num_labels, CFG.hidden_dim, CFG.dropout
    ).to(CFG.device)
    
    # Count trainable parameters
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.1f}%)")
    
    # Optimizer and scheduler
    optimizer = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=CFG.learning_rate,
        weight_decay=CFG.weight_decay
    )
    
    total_steps = len(train_loader) * CFG.epochs // CFG.gradient_accumulation
    warmup_steps = int(total_steps * CFG.warmup_ratio)
    scheduler = get_linear_schedule_with_warmup(
        optimizer, warmup_steps, total_steps
    )
    
    scaler = GradScaler(enabled=CFG.use_amp)
    
    # Training loop
    print("\n🚀 Starting training...")
    best_f1 = 0
    
    for epoch in range(CFG.epochs):
        print(f"\n{'='*40}")
        print(f"Epoch {epoch + 1}/{CFG.epochs}")
        print('='*40)
        
        train_loss = train_epoch(
            model, train_loader, optimizer, scheduler, scaler, 
            CFG.device, CFG.gradient_accumulation
        )
        print(f"Train loss: {train_loss:.4f}")
        
        val_f1 = validate(model, val_loader, CFG.device)
        print(f"Val F-max: {val_f1:.4f}")
        
        if val_f1 > best_f1:
            best_f1 = val_f1
            print(f"✅ New best! Saving model...")
            torch.save({
                'model_state_dict': model.state_dict(),
                'mlb_classes': mlb.classes_.tolist(),
                'config': {
                    'num_labels': num_labels,
                    'hidden_dim': CFG.hidden_dim,
                    'dropout': CFG.dropout,
                }
            }, f'{CFG.output_dir}/esm2_best.pt')
        
        gc.collect()
        torch.cuda.empty_cache()
    
    print(f"\n{'='*60}")
    print(f"🎉 Training complete! Best F-max: {best_f1:.4f}")
    print(f"Model saved to: {CFG.output_dir}/esm2_best.pt")
    print('='*60)


if __name__ == "__main__":
    main()