import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torch.nn.functional as F
from torch.cuda.amp import GradScaler, autocast
import os
import json
import argparse
import time
import math
import glob
from typing import Dict, List, Optional
from tqdm import tqdm
import numpy as np
import gc
import logging
from collections import defaultdict
import multiprocessing

# Import custom modules
try:
    from model_slm import MixtureOfRecursions, count_parameters, TextGenerator
    from custom_tokenizer import TechnicalTokenizer
except ImportError as e:
    raise ImportError(f"Failed to import custom modules: {e}")

# Constants for configuration
DEFAULT_MAX_LENGTH = 128
DEFAULT_MAX_EXAMPLES = 50000
DEFAULT_D_MODEL = 384
DEFAULT_N_LAYERS = 6
DEFAULT_N_HEADS = 6
DEFAULT_EPOCHS = 15
DEFAULT_BATCH_SIZE = 16
DEFAULT_LEARNING_RATE = 5e-4
DEFAULT_GRAD_ACCUM_STEPS = 1
DEFAULT_EVAL_EVERY = 500
DEFAULT_WARMUP_RATIO = 0.05
DEFAULT_CHECKPOINT_DIR = "checkpoints"
DEFAULT_LOG_LEVEL = "INFO"

# Set up logging
logging.basicConfig(
    level=DEFAULT_LOG_LEVEL,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("training.log")
    ]
)
logger = logging.getLogger(__name__)

class FastTechnicalTextDataset(Dataset):
    """Optimized dataset for fast loading and processing of technical text."""
    
    def __init__(
        self,
        data_file: str,
        tokenizer: TechnicalTokenizer,
        max_length: int = DEFAULT_MAX_LENGTH,
        max_examples: int = DEFAULT_MAX_EXAMPLES
    ):
        """
        Initialize the dataset with optimized loading.

        Args:
            data_file (str): Path to the training data file.
            tokenizer (TechnicalTokenizer): Tokenizer for encoding text.
            max_length (int): Maximum sequence length.
            max_examples (int): Maximum number of examples to load.

        Raises:
            FileNotFoundError: If the data file does not exist.
            ValueError: If max_length or max_examples is invalid.
        """
        if not os.path.exists(data_file):
            raise FileNotFoundError(f"Data file not found: {data_file}")
        if max_length <= 0 or max_examples <= 0:
            raise ValueError("max_length and max_examples must be positive")

        self.tokenizer = tokenizer
        self.max_length = max_length
        self.pad_token_id = tokenizer.vocab.get('<pad>', 0)
        self.max_examples = max_examples
        self.examples = []

        logger.info(f"Loading dataset from {data_file} with max_length={max_length}, max_examples={max_examples}")
        start_time = time.time()
        self._fast_load_data(data_file)
        self._tensorize_data()
        logger.info(f"Loaded {len(self.examples)} examples in {time.time() - start_time:.1f}s")

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()

    def _fast_load_data(self, data_file: str) -> None:
        """Load and filter data efficiently."""
        logger.info("Reading and filtering data...")
        with open(data_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        logger.info(f"File contains {len(lines)} lines")
        good_examples = []
        seen_hashes = set()

        for line in lines[:self.max_examples * 3]:
            line = line.strip()
            if (
                50 <= len(line) <= 400 and
                line.count(' ') >= 8 and
                not line.lower().startswith(('http', 'www', 'ftp')) and
                line.count('.') <= len(line) * 0.1
            ):
                line_hash = hash(line[:100])
                if line_hash not in seen_hashes:
                    seen_hashes.add(line_hash)
                    good_examples.append(line)
                    if len(good_examples) >= self.max_examples:
                        break

        logger.info(f"Filtered to {len(good_examples)} quality examples")

        batch_size = 1000
        for i in range(0, len(good_examples), batch_size):
            batch = good_examples[i:i + batch_size]
            for line in batch:
                try:
                    if not line.endswith('<|endoftext|>'):
                        line += ' <|endoftext|>'
                    tokens = self.tokenizer.encode_ids(line, add_special_tokens=True)
                    if 30 <= len(tokens) <= self.max_length:
                        if len(tokens) < self.max_length:
                            tokens.extend([self.pad_token_id] * (self.max_length - len(tokens)))
                        self.examples.append(tokens)
                except Exception as e:
                    logger.warning(f"Failed to process line: {e}")
                    continue
            if i % 5000 == 0:
                logger.info(f"Processed {len(self.examples)} examples...")

        logger.info(f"Final dataset size: {len(self.examples)} examples")

    def _tensorize_data(self) -> None:
        """Pre-tensorize data for faster training."""
        logger.info("Pre-tensorizing data...")
        seq_len = self.max_length - 1
        tensorized_examples = []

        for tokens in self.examples:
            if len(tokens) != self.max_length:
                continue
            input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
            targets = torch.tensor(tokens[1:], dtype=torch.long)
            original_len = next((i for i, x in enumerate(tokens) if x == self.pad_token_id), self.max_length)
            mask_len = min(original_len, seq_len)
            attention_mask = torch.zeros(seq_len, dtype=torch.long)
            attention_mask[:mask_len] = 1
            tensorized_examples.append({
                'input_ids': input_ids,
                'targets': targets,
                'attention_mask': attention_mask
            })

        self.examples = tensorized_examples
        logger.info("Data pre-tensorized successfully")

    def __len__(self) -> int:
        """Return the number of examples in the dataset."""
        return len(self.examples)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Return a single example from the dataset."""
        return self.examples[idx]

class FastCosineScheduler:
    """Cosine learning rate scheduler with warmup."""
    
    def __init__(self, optimizer: optim.Optimizer, total_steps: int, warmup_ratio: float = DEFAULT_WARMUP_RATIO):
        """
        Initialize the cosine scheduler.

        Args:
            optimizer (optim.Optimizer): Optimizer to schedule.
            total_steps (int): Total training steps.
            warmup_ratio (float): Ratio of steps for warmup phase.

        Raises:
            ValueError: If total_steps or warmup_ratio is invalid.
        """
        if total_steps <= 0 or not 0 <= warmup_ratio <= 1:
            raise ValueError("total_steps must be positive and warmup_ratio must be in [0, 1]")

        self.optimizer = optimizer
        self.total_steps = total_steps
        self.warmup_steps = int(total_steps * warmup_ratio)
        self.base_lr = optimizer.param_groups[0]['lr']
        self.step_count = 0

    def step(self) -> float:
        """
        Update the learning rate.

        Returns:
            float: Current learning rate.
        """
        self.step_count += 1
        if self.step_count <= self.warmup_steps:
            lr = self.base_lr * self.step_count / self.warmup_steps
        else:
            progress = (self.step_count - self.warmup_steps) / (self.total_steps - self.warmup_steps)
            lr = self.base_lr * 0.5 * (1 + math.cos(math.pi * progress))
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
        return lr

class UltraFastTrainer:
    """Trainer optimized for fast training of transformer models."""
    
    def __init__(
        self,
        model: nn.Module,
        tokenizer: TechnicalTokenizer,
        train_dataset: FastTechnicalTextDataset,
        val_dataset: Optional[FastTechnicalTextDataset] = None,
        config: Optional[Dict] = None
    ):
        """
        Initialize the trainer.

        Args:
            model (nn.Module): The transformer model to train.
            tokenizer (TechnicalTokenizer): Tokenizer for encoding/decoding.
            train_dataset (FastTechnicalTextDataset): Training dataset.
            val_dataset (Optional[FastTechnicalTextDataset]): Validation dataset.
            config (Optional[Dict]): Training configuration.

        Raises:
            ValueError: If config contains invalid parameters.
        """
        self.model = model
        self.tokenizer = tokenizer
        self.train_dataset = train_dataset
        self.val_dataset = val_dataset
        self.config = config or {}
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model.to(self.device)
        
        self._validate_config()
        self._fast_init_weights()
        self._setup_fast_optimizer()
        
        epochs = self.config.get('epochs', DEFAULT_EPOCHS)
        batch_size = self.config.get('batch_size', DEFAULT_BATCH_SIZE)
        total_steps = len(train_dataset) // batch_size * epochs
        self.scheduler = FastCosineScheduler(self.optimizer, total_steps)
        self.scaler = GradScaler() if self.device.type == 'cuda' else None
        self.global_step = 0
        self.best_loss = float('inf')
        self.grad_accum_steps = self.config.get('gradient_accumulation_steps', DEFAULT_GRAD_ACCUM_STEPS)
        self.eval_every = self.config.get('eval_every', DEFAULT_EVAL_EVERY)

    def _validate_config(self) -> None:
        """Validate training configuration."""
        if self.config.get('batch_size', DEFAULT_BATCH_SIZE) <= 0:
            raise ValueError("batch_size must be positive")
        if self.config.get('epochs', DEFAULT_EPOCHS) <= 0:
            raise ValueError("epochs must be positive")
        if self.config.get('learning_rate', DEFAULT_LEARNING_RATE) <= 0:
            raise ValueError("learning_rate must be positive")
        if self.config.get('gradient_accumulation_steps', DEFAULT_GRAD_ACCUM_STEPS) <= 0:
            raise ValueError("gradient_accumulation_steps must be positive")

    def _fast_init_weights(self) -> None:
        """Initialize model weights."""
        def fast_init(module: nn.Module) -> None:
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, std=0.02)
        self.model.apply(fast_init)
        logger.info("Model weights initialized")

    def _setup_fast_optimizer(self) -> None:
        """Set up AdamW optimizer."""
        lr = self.config.get('learning_rate', DEFAULT_LEARNING_RATE)
        params = [p for p in self.model.parameters() if p.requires_grad]
        self.optimizer = optim.AdamW(params, lr=lr, betas=(0.9, 0.99), weight_decay=0.01, eps=1e-6)
        logger.info(f"Optimizer initialized with learning rate: {lr}")

    def compute_fast_loss(self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """
        Compute masked cross-entropy loss.

        Args:
            logits (torch.Tensor): Model output logits of shape (batch_size, seq_len, vocab_size).
            targets (torch.Tensor): Target token IDs of shape (batch_size, seq_len).
            mask (torch.Tensor): Attention mask of shape (batch_size, seq_len).

        Returns:
            torch.Tensor: Computed loss.
        """
        logits_flat = logits.view(-1, logits.size(-1))
        targets_flat = targets.view(-1)
        mask_flat = mask.view(-1).bool()
        
        if not mask_flat.any():
            return torch.tensor(0.0, device=logits.device, requires_grad=True)
        
        return F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat])

    def train_epoch_fast(self, epoch: int, dataloader: DataLoader) -> Dict[str, float]:
        """
        Train for one epoch.

        Args:
            epoch (int): Current epoch number.
            dataloader (DataLoader): Training data loader.

        Returns:
            Dict[str, float]: Training metrics (loss, perplexity, epoch_time_min).
        """
        self.model.train()
        total_loss = 0
        num_batches = 0
        start_time = time.time()
        
        progress_bar = tqdm(dataloader, desc=f"Epoch {epoch}", leave=False, miniters=50)
        for batch_idx, batch in enumerate(progress_bar):
            input_ids = batch['input_ids'].to(self.device, non_blocking=True)
            targets = batch['targets'].to(self.device, non_blocking=True)
            mask = batch['attention_mask'].to(self.device, non_blocking=True)
            
            with autocast(enabled=self.device.type == 'cuda'):
                logits, comp_loss = self.model(input_ids, mask)
                lm_loss = self.compute_fast_loss(logits, targets, mask)
                total_loss_step = lm_loss + 0.0001 * comp_loss
                if self.grad_accum_steps > 1:
                    total_loss_step = total_loss_step / self.grad_accum_steps
            
            if self.scaler:
                self.scaler.scale(total_loss_step).backward()
                if (batch_idx + 1) % self.grad_accum_steps == 0:
                    self.scaler.unscale_(self.optimizer)
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                    self.optimizer.zero_grad(set_to_none=True)
                    self.scheduler.step()
                    self.global_step += 1
            else:
                total_loss_step.backward()
                if (batch_idx + 1) % self.grad_accum_steps == 0:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.optimizer.step()
                    self.optimizer.zero_grad(set_to_none=True)
                    self.scheduler.step()
                    self.global_step += 1
            
            total_loss += lm_loss.item()
            num_batches += 1
            
            if batch_idx % 100 == 0:
                current_loss = total_loss / num_batches
                progress_bar.set_postfix({'loss': f"{current_loss:.3f}", 'ppl': f"{math.exp(min(current_loss, 10)):.1f}"})
            
            if batch_idx % 200 == 0 and batch_idx > 0 and self.device.type == 'cuda':
                torch.cuda.empty_cache()
        
        avg_loss = total_loss / max(num_batches, 1)
        return {
            'loss': avg_loss,
            'perplexity': math.exp(min(avg_loss, 10)),
            'epoch_time_min': (time.time() - start_time) / 60
        }

    def validate_fast(self, dataloader: DataLoader) -> Dict[str, float]:
        """
        Validate the model on the validation dataset.

        Args:
            dataloader (DataLoader): Validation data loader.

        Returns:
            Dict[str, float]: Validation metrics (loss, perplexity).
        """
        self.model.eval()
        total_loss = 0
        num_batches = 0
        max_val_batches = min(100, len(dataloader))
        
        with torch.no_grad():
            for batch_idx, batch in enumerate(dataloader):
                if batch_idx >= max_val_batches:
                    break
                input_ids = batch['input_ids'].to(self.device, non_blocking=True)
                targets = batch['targets'].to(self.device, non_blocking=True)
                mask = batch['attention_mask'].to(self.device, non_blocking=True)
                
                with autocast(enabled=self.device.type == 'cuda'):
                    logits, _ = self.model(input_ids, mask)
                    loss = self.compute_fast_loss(logits, targets, mask)
                
                total_loss += loss.item()
                num_batches += 1
        
        avg_loss = total_loss / max(num_batches, 1)
        return {'loss': avg_loss, 'perplexity': math.exp(min(avg_loss, 10))}

    def save_checkpoint_fast(self, epoch: int, metrics: Dict, save_dir: str = DEFAULT_CHECKPOINT_DIR) -> Optional[str]:
        """
        Save a checkpoint if the loss improves.

        Args:
            epoch (int): Current epoch number.
            metrics (Dict): Training and validation metrics.
            save_dir (str): Directory to save checkpoints.

        Returns:
            Optional[str]: Path to the saved checkpoint or None.
        """
        os.makedirs(save_dir, exist_ok=True)
        val_loss = metrics.get('val_loss', metrics.get('loss', float('inf')))
        
        if val_loss < self.best_loss:
            self.best_loss = val_loss
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': self.model.state_dict(),
                'optimizer_state_dict': self.optimizer.state_dict(),
                'metrics': metrics,
                'scaler_state_dict': self.scaler.state_dict() if self.scaler else None
            }
            best_path = os.path.join(save_dir, "best_model.pt")
            torch.save(checkpoint, best_path)
            logger.info(f"New best checkpoint saved: {best_path}, Loss: {val_loss:.4f}")
            return best_path
        return None

    def train_ultra_fast(self, num_epochs: int = DEFAULT_EPOCHS, batch_size: int = DEFAULT_BATCH_SIZE) -> List[Dict]:
        """
        Train the model with optimized settings.

        Args:
            num_epochs (int): Number of training epochs.
            batch_size (int): Batch size for training.

        Returns:
            List[Dict]: Training history with metrics for each epoch.
        """
        logger.info(f"Starting ultra-fast training: {num_epochs} epochs, batch_size={batch_size}")
        logger.info("Target: Loss < 2.0, PPL < 12, Time: 4-5 hours")

        train_loader = DataLoader(
            self.train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=min(multiprocessing.cpu_count(), 4),
            pin_memory=self.device.type == 'cuda',
            persistent_workers=True,
            drop_last=True
        )

        val_loader = None
        if self.val_dataset:
            val_loader = DataLoader(
                self.val_dataset,
                batch_size=batch_size * 2,
                shuffle=False,
                num_workers=min(multiprocessing.cpu_count() // 2, 2),
                pin_memory=self.device.type == 'cuda'
            )

        total_start_time = time.time()
        history = []

        for epoch in range(1, num_epochs + 1):
            logger.info(f"Starting epoch {epoch}/{num_epochs}")
            train_metrics = self.train_epoch_fast(epoch, train_loader)
            
            val_metrics = {}
            if val_loader and (epoch % 2 == 0 or epoch == num_epochs):
                val_metrics = self.validate_fast(val_loader)
            
            epoch_time = train_metrics['epoch_time_min'] * 60
            epoch_info = {
                'epoch': epoch,
                'train_loss': train_metrics['loss'],
                'train_ppl': train_metrics['perplexity'],
                'epoch_time_min': train_metrics['epoch_time_min']
            }
            if val_metrics:
                epoch_info.update({'val_loss': val_metrics['loss'], 'val_ppl': val_metrics['perplexity']})
            
            history.append(epoch_info)
            
            elapsed_hours = (time.time() - total_start_time) / 3600
            remaining_hours = elapsed_hours * (num_epochs - epoch) / max(epoch, 1)
            
            logger.info(f"Epoch {epoch} results:")
            logger.info(f"  Epoch time: {epoch_time/60:.1f} min")
            logger.info(f"  Total elapsed: {elapsed_hours:.1f}h")
            logger.info(f"  Est. remaining: {remaining_hours:.1f}h")
            logger.info(f"  Train Loss: {train_metrics['loss']:.4f}")
            logger.info(f"  Train PPL: {train_metrics['perplexity']:.1f}")
            if val_metrics:
                logger.info(f"  Val Loss: {val_metrics['loss']:.4f}")
                logger.info(f"  Val PPL: {val_metrics['perplexity']:.1f}")
            
            current_loss = val_metrics.get('loss', train_metrics['loss'])
            current_ppl = val_metrics.get('perplexity', train_metrics['perplexity'])
            if current_loss < 2.0 and current_ppl < 12:
                logger.info(f"Targets achieved: Loss={current_loss:.4f} < 2.0, PPL={current_ppl:.1f} < 12")
            
            combined_metrics = {**train_metrics}
            if val_metrics:
                combined_metrics.update({f"val_{k}": v for k, v in val_metrics.items()})
            self.save_checkpoint_fast(epoch, combined_metrics)
            
            if self.device.type == 'cuda':
                torch.cuda.empty_cache()
            gc.collect()
            
            if current_loss < 1.8 and current_ppl < 10:
                logger.info("Early stopping: Excellent performance achieved!")
                break
        
        total_time = (time.time() - total_start_time) / 3600
        logger.info(f"Training completed in {total_time:.1f} hours")
        logger.info(f"Best loss: {self.best_loss:.4f}")
        return history

def run_ultra_fast_training() -> int:
    """
    Run the ultra-fast training pipeline.

    Returns:
        int: Exit code (0 for success, 1 for failure).
    """
    parser = argparse.ArgumentParser(description="Ultra-Fast Training for MixtureOfRecursions Model")
    parser.add_argument("--train_file", default=None, help="Path to training data file")
    parser.add_argument("--val_file", default=None, help="Path to validation data file")
    parser.add_argument("--tokenizer_dir", default="tokenizer", help="Directory for tokenizer files")
    parser.add_argument("--max_examples", type=int, default=DEFAULT_MAX_EXAMPLES, help="Maximum number of training examples")
    parser.add_argument("--d_model", type=int, default=DEFAULT_D_MODEL, help="Model embedding dimension")
    parser.add_argument("--n_layers", type=int, default=DEFAULT_N_LAYERS, help="Number of transformer layers")
    parser.add_argument("--n_heads", type=int, default=DEFAULT_N_HEADS, help="Number of attention heads")
    parser.add_argument("--max_seq_len", type=int, default=DEFAULT_MAX_LENGTH, help="Maximum sequence length")
    parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS, help="Number of training epochs")
    parser.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE, help="Batch size for training")
    parser.add_argument("--learning_rate", type=float, default=DEFAULT_LEARNING_RATE, help="Learning rate")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=DEFAULT_GRAD_ACCUM_STEPS, help="Gradient accumulation steps")
    parser.add_argument("--eval_every", type=int, default=DEFAULT_EVAL_EVERY, help="Evaluate every N steps")
    
    args = parser.parse_args()
    
    torch.manual_seed(42)
    np.random.seed(42)
    
    logger.info("Starting ultra-fast training pipeline")
    
    if args.train_file is None:
        patterns = ["*train*.txt", "*_train.txt"]
        files = []
        for pattern in patterns:
            files.extend(glob.glob(pattern))
            files.extend(glob.glob(os.path.join("split_data", pattern)))
            files.extend(glob.glob(os.path.join("data", pattern)))
        if files:
            args.train_file = files[0]
            logger.info(f"Found training file: {args.train_file}")
        else:
            logger.error("No training files found!")
            return 1
    
    try:
        tokenizer = TechnicalTokenizer()
        tokenizer.load(args.tokenizer_dir)
        logger.info(f"Tokenizer loaded with vocab size: {tokenizer.get_vocab_size()}")
    except Exception as e:
        logger.error(f"Failed to load tokenizer: {e}")
        return 1
    
    logger.info("Creating training dataset...")
    try:
        train_dataset = FastTechnicalTextDataset(
            args.train_file, tokenizer, args.max_seq_len, args.max_examples
        )
    except Exception as e:
        logger.error(f"Failed to create training dataset: {e}")
        return 1
    
    val_dataset = None
    if args.val_file and os.path.exists(args.val_file):
        try:
            val_dataset = FastTechnicalTextDataset(
                args.val_file, tokenizer, args.max_seq_len, max_examples=5000
            )
            logger.info("Validation dataset created")
        except Exception as e:
            logger.warning(f"Failed to create validation dataset: {e}")
    
    try:
        model = MixtureOfRecursions(
            vocab_size=tokenizer.get_vocab_size(),
            d_model=args.d_model,
            n_layers=args.n_layers,
            n_heads=args.n_heads,
            max_seq_len=args.max_seq_len - 1,
            padding_idx=tokenizer.vocab.get('<pad>', 0)
        )
        logger.info("Model initialized")
    except Exception as e:
        logger.error(f"Failed to initialize model: {e}")
        return 1
    
    config = {
        'learning_rate': args.learning_rate,
        'gradient_accumulation_steps': args.gradient_accumulation_steps,
        'eval_every': args.eval_every,
        'batch_size': args.batch_size,
        'epochs': args.epochs
    }
    
    try:
        trainer = UltraFastTrainer(model, tokenizer, train_dataset, val_dataset, config)
        results = trainer.train_ultra_fast(args.epochs, args.batch_size)
        
        with open('ultra_fast_results.json', 'w') as f:
            json.dump(results, f, indent=2)
        logger.info("Training results saved to ultra_fast_results.json")
        
        return 0
    except Exception as e:
        logger.error(f"Training failed: {e}")
        return 1

if __name__ == "__main__":
    exit(run_ultra_fast_training())