"""
Training script for the Shared Encoder

This script trains the encoder on a mixed corpus from all domains.
The encoder is trained once and then frozen for expert training.

Usage:
    python train.py --data_dir ../../data --output_dir ../../experts --epochs 10
"""

import os
import argparse
import json
from pathlib import Path
from typing import Dict, List
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import AutoTokenizer
from tqdm import tqdm
import wandb

from model import create_encoder, ContrastiveLoss


class MixedCorpusDataset(Dataset):
    """Dataset that combines samples from all expert domains."""
    
    def __init__(
        self,
        data_dir: Path,
        tokenizer,
        max_length: int = 512,
        samples_per_expert: int = 5000
    ):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = []
        
        print("Loading datasets from all expert domains...")
        
        # Load samples from each expert's data
        expert_dirs = sorted([d for d in data_dir.iterdir() if d.is_dir()])
        
        for expert_dir in tqdm(expert_dirs, desc="Loading experts"):
            metadata_path = expert_dir / "metadata.json"
            
            if not metadata_path.exists():
                continue
            
            with open(metadata_path) as f:
                metadata = json.load(f)
            
            # Skip synthetic datasets
            if metadata.get("synthetic", False):
                continue
            
            # Load raw datasets
            expert_samples = self._load_expert_data(expert_dir, samples_per_expert)
            
            # Add domain label
            for sample in expert_samples:
                sample["domain"] = metadata["expert_id"]
            
            self.samples.extend(expert_samples)
        
        print(f"Loaded {len(self.samples)} total samples from {len(expert_dirs)} domains")
    
    def _load_expert_data(self, expert_dir: Path, max_samples: int) -> List[Dict]:
        """Load samples from an expert's raw datasets."""
        from datasets import load_from_disk
        
        samples = []
        raw_datasets = sorted(expert_dir.glob("raw_*"))
        
        for raw_path in raw_datasets:
            try:
                dataset = load_from_disk(str(raw_path))
                
                # Extract text from various dataset formats
                for item in dataset:
                    text = self._extract_text(item)
                    if text and len(text) > 20:  # Minimum length filter
                        samples.append({"text": text})
                    
                    if len(samples) >= max_samples:
                        break
                
                if len(samples) >= max_samples:
                    break
            except Exception as e:
                print(f"  Warning: Could not load {raw_path}: {e}")
                continue
        
        return samples[:max_samples]
    
    def _extract_text(self, item: Dict) -> str:
        """Extract text from various dataset formats."""
        # Try common text fields
        for field in ["text", "sentence", "question", "document", "article", "code"]:
            if field in item and item[field]:
                return str(item[field])
        
        # For QA datasets
        if "question" in item and "context" in item:
            return f"{item['question']} {item['context']}"
        
        # For summarization
        if "article" in item:
            return item["article"]
        
        # Fallback: convert entire item to string
        return str(item)
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample = self.samples[idx]
        
        # Tokenize
        encoding = self.tokenizer(
            sample["text"],
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        
        return {
            "input_ids": encoding["input_ids"].squeeze(0),
            "attention_mask": encoding["attention_mask"].squeeze(0),
            "domain": hash(sample["domain"]) % 1000  # Simple domain label
        }


def train_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    epoch: int
) -> float:
    """Train for one epoch."""
    model.train()
    total_loss = 0
    
    pbar = tqdm(dataloader, desc=f"Epoch {epoch}")
    
    for batch_idx, batch in enumerate(pbar):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["domain"].to(device)
        
        # Forward pass
        embeddings = model(input_ids, attention_mask)
        
        # Compute loss
        loss = criterion(embeddings, labels)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        
        # Track metrics
        total_loss += loss.item()
        avg_loss = total_loss / (batch_idx + 1)
        
        pbar.set_postfix({"loss": f"{avg_loss:.4f}"})
        
        # Log to wandb
        if batch_idx % 100 == 0:
            wandb.log({
                "train/loss": loss.item(),
                "train/avg_loss": avg_loss,
                "train/epoch": epoch
            })
    
    return total_loss / len(dataloader)


@torch.no_grad()
def validate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device
) -> float:
    """Validate the model."""
    model.eval()
    total_loss = 0
    
    for batch in tqdm(dataloader, desc="Validating"):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["domain"].to(device)
        
        embeddings = model(input_ids, attention_mask)
        loss = criterion(embeddings, labels)
        
        total_loss += loss.item()
    
    return total_loss / len(dataloader)


def main():
    parser = argparse.ArgumentParser(description="Train shared encoder")
    parser.add_argument("--data_dir", type=str, default="../../data", help="Data directory")
    parser.add_argument("--output_dir", type=str, default="../../experts", help="Output directory")
    parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
    parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
    parser.add_argument("--samples_per_expert", type=int, default=5000, help="Samples per expert")
    parser.add_argument("--max_length", type=int, default=512, help="Max sequence length")
    parser.add_argument("--wandb_project", type=str, default="lego-moe", help="W&B project name")
    parser.add_argument("--no_wandb", action="store_true", help="Disable W&B logging")
    
    args = parser.parse_args()
    
    # Setup
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    data_dir = Path(args.data_dir)
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Initialize wandb
    if not args.no_wandb:
        wandb.init(
            project=args.wandb_project,
            name="shared-encoder-training",
            config=vars(args)
        )
    
    # Load tokenizer
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    
    # Create datasets
    print("Creating datasets...")
    train_dataset = MixedCorpusDataset(
        data_dir,
        tokenizer,
        max_length=args.max_length,
        samples_per_expert=args.samples_per_expert
    )
    
    # Split into train/val
    train_size = int(0.95 * len(train_dataset))
    val_size = len(train_dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(
        train_dataset,
        [train_size, val_size]
    )
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=True
    )
    
    # Create model
    print("Creating model...")
    model = create_encoder({
        "vocab_size": tokenizer.vocab_size,
        "embed_dim": 256,
        "num_layers": 4,
        "num_heads": 8,
        "ff_dim": 1024,
        "dropout": 0.1,
        "max_seq_len": args.max_length
    })
    
    model = model.to(device)
    
    num_params, size_mb = model.get_model_size()
    print(f"Model parameters: {num_params:,}")
    print(f"Model size: {size_mb:.2f} MB")
    
    # Create optimizer and scheduler
    optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
    scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs * len(train_loader))
    
    # Create loss function
    criterion = ContrastiveLoss(temperature=0.07)
    
    # Training loop
    print(f"\nStarting training for {args.epochs} epochs...")
    best_val_loss = float("inf")
    
    for epoch in range(1, args.epochs + 1):
        print(f"\n{'='*60}")
        print(f"Epoch {epoch}/{args.epochs}")
        print(f"{'='*60}")
        
        # Train
        train_loss = train_epoch(model, train_loader, optimizer, criterion, device, epoch)
        
        # Validate
        val_loss = validate(model, val_loader, criterion, device)
        
        # Update scheduler
        scheduler.step()
        
        print(f"Train Loss: {train_loss:.4f}")
        print(f"Val Loss: {val_loss:.4f}")
        
        # Log to wandb
        if not args.no_wandb:
            wandb.log({
                "epoch": epoch,
                "train/epoch_loss": train_loss,
                "val/loss": val_loss,
                "lr": optimizer.param_groups[0]["lr"]
            })
        
        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            checkpoint_path = output_dir / "encoder_best.pt"
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "val_loss": val_loss,
                "config": model.__dict__
            }, checkpoint_path)
            print(f"✓ Saved best model to {checkpoint_path}")
        
        # Save checkpoint
        checkpoint_path = output_dir / f"encoder_epoch_{epoch}.pt"
        torch.save({
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "val_loss": val_loss
        }, checkpoint_path)
    
    # Export final model
    print("\nExporting model to ONNX...")
    model.eval()
    onnx_path = output_dir / "encoder.onnx"
    model.export_to_onnx(str(onnx_path))
    
    # Save tokenizer
    tokenizer.save_pretrained(str(output_dir / "tokenizer"))
    
    print("\n✅ Training complete!")
    print(f"Best validation loss: {best_val_loss:.4f}")
    print(f"Model saved to: {output_dir}")
    print(f"\nNext steps:")
    print(f"  1. Generate embeddings: python ../datasets/generate_embeddings.py")
    print(f"  2. Train experts: python ../distillation/train_expert.py")
    
    if not args.no_wandb:
        wandb.finish()


if __name__ == "__main__":
    main()