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
from tqdm import tqdm
import numpy as np
import gc
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:
    print(f"Import error: {e}")
    exit(1)
class FastTechnicalTextDataset(Dataset):
    """Ultra-fast dataset with aggressive optimizations for 4-5hr training"""    
    def __init__(self, data_file: str, tokenizer: TechnicalTokenizer, max_length: int = 128, max_examples: int = 50000):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.pad_token_id = tokenizer.vocab.get('<pad>', 0)
        self.max_examples = max_examples        
        print(f"FAST DATASET LOADING")
        print(f"Data file: {data_file}")
        print(f"Max sequence length: {max_length}")
        print(f"Max examples: {max_examples}")        
        start_time = time.time()
        self.examples = []
        self._fast_load_data(data_file)        
        load_time = time.time() - start_time
        print(f" Loaded {len(self.examples)} examples in {load_time:.1f}s")        
        self._tensorize_data()
        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None    
    def _fast_load_data(self, data_file: str):
        print("🔍 Fast reading file...")        
        with open(data_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()        
        print(f"File has {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        
        print(f"After fast filtering: {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 = tokens + [self.pad_token_id] * (self.max_length - len(tokens))
                        self.examples.append(tokens)
                except:
                    continue
            if i % 5000 == 0:
                print(f"Processed {len(self.examples)} examples...")        
        print(f"Final dataset: {len(self.examples)} examples")    
    def _tensorize_data(self):
        print("Pre-tensorizing data for maximum speed...")
        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
        print("All data pre-tensorized")    
    def __len__(self):
        return len(self.examples)    
    def __getitem__(self, idx):
        return self.examples[idx]
class FastCosineScheduler:
    def __init__(self, optimizer, total_steps: int, warmup_ratio: float = 0.05):
        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):
        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:
    def __init__(self, model, tokenizer, train_dataset, val_dataset=None, config=None):
        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._fast_init_weights()
        self._setup_fast_optimizer()        
        epochs = self.config.get('epochs', 15)
        batch_size = self.config.get('batch_size', 16)
        total_steps = len(train_dataset) // batch_size * epochs
        self.scheduler = FastCosineScheduler(self.optimizer, total_steps)
        self.scaler = GradScaler()        
        self.global_step = 0
        self.best_loss = float('inf')
        self.grad_accum_steps = self.config.get('gradient_accumulation_steps', 1)
        self.eval_every = self.config.get('eval_every', 500)    
    def _fast_init_weights(self):
        def fast_init(module):
            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)    
    def _setup_fast_optimizer(self):
        lr = self.config.get('learning_rate', 5e-4)
        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)    
    def compute_fast_loss(self, logits, targets, mask):
        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)
        loss = F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat])
        return loss    
    def train_epoch_fast(self, epoch: int, dataloader: DataLoader) -> Dict[str, float]:
        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():
                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            
            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            
            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:
                torch.cuda.empty_cache()       
        epoch_time = time.time() - start_time
        avg_loss = total_loss / max(num_batches, 1)
        return {'loss': avg_loss, 'perplexity': math.exp(min(avg_loss, 10)), 'epoch_time_min': epoch_time / 60}    
    def validate_fast(self, dataloader: DataLoader) -> Dict[str, float]:
        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():
                    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 = "checkpoints"):
        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()
            }
            best_path = os.path.join(save_dir, "best_model.pt")
            torch.save(checkpoint, best_path)
            print(f"New best! Loss: {val_loss:.4f}")
            return best_path
        return None    
    def train_ultra_fast(self, num_epochs: int = 15, batch_size: int = 16):
        print(f"\n ULTRA-FAST TRAINING")
        print(f" Target: Loss < 2.0, PPL < 12")
        print(f" Time target: 4-5 hours")
        print(f" Epochs: {num_epochs}")
        print(f" Batch size: {batch_size}")
        print("-" * 60)        
        train_loader = DataLoader(
            self.train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=4,
            pin_memory=True,
            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=2,
                pin_memory=True
            )        
        total_start_time = time.time()
        history = []        
        for epoch in range(1, num_epochs + 1):
            epoch_start = time.time()
            print(f"\n 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 = time.time() - epoch_start
            epoch_info = {
                'epoch': epoch,
                'train_loss': train_metrics['loss'],
                'train_ppl': train_metrics['perplexity'],
                'epoch_time_min': epoch_time / 60
            }
            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) / epoch            
            print(f"\n EPOCH {epoch} RESULTS:")
            print(f" Epoch time: {epoch_time/60:.1f} min")
            print(f" Total elapsed: {elapsed_hours:.1f}h")
            print(f" Est. remaining: {remaining_hours:.1f}h")
            print(f" Train Loss: {train_metrics['loss']:.4f}")
            print(f" Train PPL: {train_metrics['perplexity']:.1f}")
            if val_metrics:
                print(f" Val Loss: {val_metrics['loss']:.4f}")
                print(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:
                print(f" TARGETS ACHIEVED!")
                print(f" Loss: {current_loss:.4f} < 2.0")
                print(f" 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)           
            torch.cuda.empty_cache()
            gc.collect()            
            if current_loss < 1.8 and current_ppl < 10:
                print(f"EARLY STOPPING - Excellent performance achieved!")
                break        
        total_time = time.time() - total_start_time
        print(f"\n TRAINING COMPLETED!")
        print(f"Total time: {total_time/3600:.1f} hours")
        print(f" Best loss: {self.best_loss:.4f}")
        return history
def run_ultra_fast_training():
    parser = argparse.ArgumentParser(description="Ultra-Fast Training for 4-5 Hours")
    parser.add_argument("--train_file", default=None)
    parser.add_argument("--val_file", default=None)
    parser.add_argument("--tokenizer_dir", default="tokenizer")
    parser.add_argument("--max_examples", type=int, default=50000)
    parser.add_argument("--d_model", type=int, default=384)
    parser.add_argument("--n_layers", type=int, default=6)
    parser.add_argument("--n_heads", type=int, default=6)
    parser.add_argument("--max_seq_len", type=int, default=128)
    parser.add_argument("--epochs", type=int, default=15)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--learning_rate", type=float, default=5e-4)
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument("--eval_every", type=int, default=500)    
    args = parser.parse_args()    
    torch.manual_seed(42)
    np.random.seed(42)    
    print("Training My Model")
    print("-" * 50)    
    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(f"split_data/{pattern}"))
            files.extend(glob.glob(f"data/{pattern}"))
        if files:
            args.train_file = files[0]
            print(f"Found: {args.train_file}")
        else:
            print(" No training files found!")
            return 1    
    tokenizer = TechnicalTokenizer()
    try:
        tokenizer.load(args.tokenizer_dir)
        print(f"Tokenizer loaded. Vocab size: {tokenizer.get_vocab_size()}")
    except Exception as e:
        print(f" Tokenizer error: {e}")
        return 1    
    print(" Creating ultra-fast dataset...")
    train_dataset = FastTechnicalTextDataset(
        args.train_file, tokenizer, args.max_seq_len, args.max_examples
    )    
    val_dataset = None
    if args.val_file and os.path.exists(args.val_file):
        val_dataset = FastTechnicalTextDataset(
            args.val_file, tokenizer, args.max_seq_len, max_examples=5000
        )    
    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, # Pass the actual sequence length to the model
        padding_idx=tokenizer.vocab.get('<pad>', 0)
    )    
    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
    }    
    trainer = UltraFastTrainer(model, tokenizer, train_dataset, val_dataset, config)
    print(f"\n START TRAINING")
    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)    
    print("\n Training Completed!")
    print(" Results saved to: ultra_fast_results.json")
    return 0
if __name__ == "__main__":
    exit(run_ultra_fast_training())