import os
import shutil
import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state
from flax import jax_utils
import orbax.checkpoint as ocp
import numpy as np
from functools import partial
from datasets import load_dataset
from transformers import AutoTokenizer
from huggingface_hub import HfApi, login
import traceback # Error details dekhne ke liye

# ... (Upar ka Config, Attention, Model aur Data Loaders wala code same rahega) ...
# ==========================================
# 1. 200M MODEL CONFIGURATION (SCALED UP)
# ==========================================
class ProdModelConfig:
    vocab_size = 32000
    hidden_size = 896        # Scaled up for ~200M
    intermediate_size = 3584 # 4x hidden_size
    num_attention_heads = 14 # 896 / 14 = 64 dimension per head
    num_layers = 18          # Increased depth
    max_thought_loops = 6    
    num_memory_tokens = 16   
    context_length = 496     
    dtype = jnp.bfloat16     

# ==========================================
# 2. KV CACHE & SELF-ATTENTION WITH ROPE
# ==========================================
def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return jnp.concatenate([-x2, x1], axis=-1)

def apply_rope(xq, xk, position, dtype):
    dim = xq.shape[-1]
    inv_freq = 1.0 / (10000 ** (jnp.arange(0, dim, 2, dtype=jnp.float32) / dim))
    freqs = jnp.einsum("i,j->ij", position, inv_freq)
    emb = jnp.concatenate([freqs, freqs], axis=-1)
    cos = jnp.cos(emb)[None, :, None, :].astype(dtype)
    sin = jnp.sin(emb)[None, :, None, :].astype(dtype)
    return (xq * cos) + (rotate_half(xq) * sin), (xk * cos) + (rotate_half(xk) * sin)

class CausalAttentionWithRoPE(nn.Module):
    config: ProdModelConfig

    @nn.compact
    def __call__(self, x, mask=None, decode=False):
        B, S, C = x.shape
        H = self.config.num_attention_heads
        D = C // H
        
        q = nn.Dense(C, use_bias=False, dtype=self.config.dtype)(x)
        k = nn.Dense(C, use_bias=False, dtype=self.config.dtype)(x)
        v = nn.Dense(C, use_bias=False, dtype=self.config.dtype)(x)
        
        q = q.reshape(B, S, H, D)
        k = k.reshape(B, S, H, D)
        v = v.reshape(B, S, H, D)
        
        if decode:
            is_initialized = self.has_variable('cache', 'cached_key')
            cached_key = self.variable('cache', 'cached_key', jnp.zeros, (B, self.config.context_length, H, D), k.dtype)
            cached_value = self.variable('cache', 'cached_value', jnp.zeros, (B, self.config.context_length, H, D), v.dtype)
            cache_index = self.variable('cache', 'cache_index', lambda: jnp.array(0, dtype=jnp.int32))
            
            if is_initialized:
                idx = cache_index.value
                k = cached_key.value.at[:, idx:idx+S].set(k)
                v = cached_value.value.at[:, idx:idx+S].set(v)
                cache_index.value = idx + S
                positions = jnp.arange(idx, idx + S)
            else:
                positions = jnp.arange(S)
        else:
            positions = jnp.arange(S)
            
        q, k = apply_rope(q, k, positions, self.config.dtype)
        
        scores = jnp.einsum('bqhd,bkhd->bhqk', q, k) / jnp.sqrt(D).astype(self.config.dtype)
        if mask is not None:
            scores = jnp.where(mask, scores, -10000.0)
            
        attn_weights = jax.nn.softmax(scores, axis=-1)
        out = jnp.einsum('bhqk,bkhd->bqhd', attn_weights, v).reshape(B, S, C)
        return nn.Dense(C, use_bias=False, dtype=self.config.dtype)(out)

class MLPBlock(nn.Module):
    config: ProdModelConfig
    @nn.compact
    def __call__(self, x):
        h = nn.Dense(self.config.intermediate_size, use_bias=False, dtype=self.config.dtype)(x)
        h = jax.nn.silu(h)
        return nn.Dense(self.config.hidden_size, use_bias=False, dtype=self.config.dtype)(h)

class TransformerBlock(nn.Module):
    config: ProdModelConfig
    @nn.compact
    def __call__(self, x, mask=None, decode=False):
        attn_in = nn.RMSNorm(dtype=self.config.dtype)(x)
        x = x + CausalAttentionWithRoPE(self.config)(attn_in, mask=mask, decode=decode)
        
        mlp_in = nn.RMSNorm(dtype=self.config.dtype)(x)
        x = x + MLPBlock(self.config)(mlp_in)
        return x

class DynamicThinkingLM(nn.Module):
    config: ProdModelConfig
    @nn.compact
    def __call__(self, input_ids, decode=False):
        batch_size, seq_len = input_ids.shape
        total_len = seq_len + self.config.num_memory_tokens
        
        x = nn.Embed(self.config.vocab_size, self.config.hidden_size, dtype=self.config.dtype)(input_ids)
        memory_tokens = self.param('memory_tokens', jax.nn.initializers.normal(stddev=0.02), (1, self.config.num_memory_tokens, self.config.hidden_size))
        memory_tokens = jnp.broadcast_to(memory_tokens, (batch_size, self.config.num_memory_tokens, self.config.hidden_size)).astype(self.config.dtype)
        
        if not decode:
            x = jnp.concatenate([memory_tokens, x], axis=1)
            causal_mask = jnp.tril(jnp.ones((1, 1, total_len, total_len), dtype=bool))
        else:
            causal_mask = None 

        for i in range(self.config.num_layers):
            x = TransformerBlock(self.config, name=f"perception_{i}")(x, mask=causal_mask, decode=decode)
            
        reasoning_core = TransformerBlock(self.config, name="reasoning_core")
        halt_classifier = nn.Dense(1, use_bias=False, dtype=jnp.float32)
        
        accumulated_state = jnp.zeros_like(x)
        remainders = jnp.ones((batch_size, 1, 1))
        halt_probs_sum = jnp.zeros((batch_size, 1, 1))
        ponder_steps = jnp.zeros((batch_size, 1, 1))

        for _ in range(self.config.max_thought_loops):
            x = reasoning_core(x, mask=causal_mask, decode=decode)
            p_halt = jax.nn.sigmoid(halt_classifier(jnp.mean(x, axis=1, keepdims=True)))
            still_thinking = (halt_probs_sum < 1.0).astype(jnp.float32)
            p_step = jnp.where(halt_probs_sum + p_halt > 1.0, remainders, p_halt) * still_thinking
            
            halt_probs_sum += p_step
            remainders -= p_step
            ponder_steps += still_thinking
            accumulated_state += p_step.astype(self.config.dtype) * x

        x = nn.RMSNorm(dtype=self.config.dtype)(accumulated_state)
        
        if not decode:
            x = x[:, self.config.num_memory_tokens:, :] 
            
        logits = nn.Dense(self.config.vocab_size, use_bias=False, dtype=self.config.dtype)(x)
        return logits, ponder_steps

# ==========================================
# 3. DATA LOADERS (Streaming - Web, Chat, Math)
# ==========================================
tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
tokenizer.pad_token = tokenizer.eos_token 

def create_fineweb_stream(global_batch_size, seq_len):
    dataset = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)
    buffer = []
    for example in dataset:
        buffer.extend(tokenizer(example['text'])['input_ids'])
        while len(buffer) >= global_batch_size * seq_len:
            batch = buffer[:global_batch_size * seq_len]
            buffer = buffer[global_batch_size * seq_len:]
            yield np.array(batch, dtype=np.int32).reshape(global_batch_size, seq_len)

def create_chat_stream(global_batch_size, seq_len):
    dataset = load_dataset("yahma/alpaca-cleaned", split="train", streaming=True)
    buffer = []
    for example in dataset:
        chat_text = f"[INST] {example['instruction']} [/INST] {example['output']} {tokenizer.eos_token}"
        buffer.extend(tokenizer(chat_text)['input_ids'])
        while len(buffer) >= global_batch_size * seq_len:
            batch = buffer[:global_batch_size * seq_len]
            buffer = buffer[global_batch_size * seq_len:]
            yield np.array(batch, dtype=np.int32).reshape(global_batch_size, seq_len)

def create_math_stream(global_batch_size, seq_len):
    # Math dataset for logical reasoning
    dataset = load_dataset("meta-math/MetaMathQA", split="train", streaming=True)
    buffer = []
    for example in dataset:
        math_text = f"[INST] Solve this math problem: {example['query']} [/INST] {example['response']} {tokenizer.eos_token}"
        buffer.extend(tokenizer(math_text)['input_ids'])
        while len(buffer) >= global_batch_size * seq_len:
            batch = buffer[:global_batch_size * seq_len]
            buffer = buffer[global_batch_size * seq_len:]
            yield np.array(batch, dtype=np.int32).reshape(global_batch_size, seq_len)
# ==========================================
# 4. HIGH-STEP 8B TOKENS TRAINING LOOP (WITH FAILSAFE)
# ==========================================
def main():
    # 1️⃣ HARDCODED HUGGING FACE TOKEN
    # Yahan apni asli token daalein "hf_..."
    hf_token = "hf_mpvCHCISxesjIaGXwaumyxJlyTzRRMuf" 
    print("🔑 Logging into Hugging Face...")
    login(token=hf_token)

    num_devices = jax.device_count()
    print(f"🚀 Initializing on {num_devices} TPU cores!")

    config = ProdModelConfig()
    
    per_device_batch = 16 
    global_batch_size = per_device_batch * num_devices 
    tokens_per_step = global_batch_size * config.context_length
    print(f"📦 Global Batch: {global_batch_size} | Tokens/Step: {tokens_per_step:,}")

    fineweb_loader = create_fineweb_stream(global_batch_size, config.context_length)
    chat_loader = create_chat_stream(global_batch_size, config.context_length)
    math_loader = create_math_stream(global_batch_size, config.context_length)

    model = DynamicThinkingLM(config)
    rng = jax.random.PRNGKey(42)
    dummy_input = jnp.ones((1, config.context_length), dtype=jnp.int32)
    params = model.init(rng, dummy_input)

    switch_step = 110250  
    total_steps = 126000 
    
    tx = optax.chain(
        optax.clip_by_global_norm(1.0),
        optax.adamw(learning_rate=optax.cosine_decay_schedule(4e-4, total_steps), weight_decay=0.1)
    )
    
    state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)
    state = jax_utils.replicate(state)

    def loss_fn(params, batch):
        logits, ponder_steps = model.apply(params, batch)
        shift_logits = logits[..., :-1, :]
        shift_labels = batch[..., 1:]
        ce_loss = optax.softmax_cross_entropy_with_integer_labels(shift_logits, shift_labels).mean()
        return ce_loss + (1e-3 * ponder_steps.mean()), ce_loss

    @partial(jax.pmap, axis_name='batch')
    def train_step(state, batch):
        grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
        (total_loss, ce_loss), grads = grad_fn(state.params, batch)
        grads = jax.lax.pmean(grads, axis_name='batch')
        total_loss = jax.lax.pmean(total_loss, axis_name='batch')
        ce_loss = jax.lax.pmean(ce_loss, axis_name='batch')
        state = state.apply_gradients(grads=grads)
        return state, total_loss, ce_loss

    ckpt_dir = '/kaggle/working/200M_Model_Checkpoints'
    ocp_options = ocp.CheckpointManagerOptions(max_to_keep=2, create=True)
    checkpoint_manager = ocp.CheckpointManager(ckpt_dir, options=ocp_options)

    print(f"\n🔥 PHASE 1: Pre-Training on 7 Billion Tokens ({switch_step} Steps)...")
    
    step = 0 # Step variable bahar declare kiya taaki exception block me use ho sake
    
    # 2️⃣ FAILSAFE TRY-EXCEPT BLOCK
    try:
        for step in range(1, total_steps + 1):
            
            if step == switch_step:
                print("\n==============================================")
                print("🚀 PHASE 2 INITIATED: 1 Billion Tokens Chat & Math Alignment!")
                print("==============================================\n")
                
            if step < switch_step:
                batch = next(fineweb_loader)
            else:
                batch = next(chat_loader) if step % 2 == 0 else next(math_loader)
                
            batch = batch.reshape(num_devices, per_device_batch, config.context_length)
                
            state, total_loss, ce_loss = train_step(state, batch)
            
            if step % 100 == 0:
                loss_v = float(total_loss[0].block_until_ready())
                ppl_v = np.exp(min(float(ce_loss[0]), 100)) 
                mode = "PRETRAIN" if step < switch_step else ("CHAT-TUNE" if step % 2 == 0 else "MATH-TUNE")
                print(f"[{mode}] Step {step}/{total_steps} | Loss: {loss_v:.4f} | PPL: {ppl_v:.2f}")
                
            if step % 3000 == 0 or step == total_steps:
                cpu_state = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state))
                checkpoint_manager.save(step, args=ocp.args.StandardSave(cpu_state))
                print(f"💾 Checkpoint Saved at Step {step}.")

    except (KeyboardInterrupt, Exception) as e:
        # 🚨 AGAR KUCH GADBAD HUI YA KISI NE ROKA TOH YAHAN AAYEGA
        print(f"\n⚠️ TRAINING RUK GAYI! Error: {e}")
        # Agar manual stop nahi hai toh error ki details print karega
        if not isinstance(e, KeyboardInterrupt):
            traceback.print_exc()
            
        print(f"\n🛑 Ghabrayein nahi bro, Step {step} par Emergency Checkpoint save ho raha hai...")
        try:
            cpu_state = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state))
            # Step ko save karenge taaki pichla checkpoint overwrite na ho, usme +1 kar diya taaki conflict na ho
            save_step = step if step % 3000 != 0 else step + 1 
            checkpoint_manager.save(save_step, args=ocp.args.StandardSave(cpu_state))
            print("✅ Emergency Checkpoint safely saved!")
        except Exception as save_err:
            print(f"❌ Emergency save fail ho gaya: {save_err}")

    finally:
        # 📦 YE WALA BLOCK HAMESHA CHALEGA (Chaahe error aaye ya puri training ho)
        print("\n📦 Hugging Face par push karne ki koshish kar raha hoon...")
        try:
            api = HfApi()
            repo_id = "Nebulixlabs/Dynamic-200M-Math"
            
            print(f"Creating private repository: {repo_id}")
            api.create_repo(repo_id=repo_id, private=True, exist_ok=True)
            
            print("Uploading checkpoint folder... This might take a few minutes.")
            api.upload_folder(
                folder_path=ckpt_dir,
                repo_id=repo_id,
                repo_type="model"
            )
            print(f"✅ Model successfully uploaded to https://huggingface.co/{repo_id}")
        except Exception as e:
            print(f"❌ Upload failed: {e}")
            print("Model is saved locally in Kaggle. Zipping...")
            shutil.make_archive('/kaggle/working/200M_Ready', 'zip', ckpt_dir)
            print("✅ Zip file ready to download manually!")

if __name__ == "__main__":
    main()