"""
End-to-end validation of the complete Fusion-LLM pipeline.

Tests:
1. FusionModel training - loss decreases over 50 steps
2. ThinkingDialModel generate - different depths produce different logits
3. GRPO training pipeline - train_step completes without errors
4. Full generate loop - model produces valid token sequences

Run: python train/e2e_validation.py
"""
import sys
import torch
import torch.nn as nn
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from models.fusion_model import FusionModel, FusionConfig
from models.thinking_dial import ThinkingDialModel, ThinkingConfig, GRPOTrainer, GRPOConfig
from train.model_utils import create_local_model

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[E2E] Device: {DEVICE}")
print()

PASS_COUNT = 0
FAIL_COUNT = 0

def check(name, condition, detail=""):
    global PASS_COUNT, FAIL_COUNT
    if condition:
        PASS_COUNT += 1
        print(f"  [PASS] {name}")
    else:
        FAIL_COUNT += 1
        print(f"  [FAIL] {name} {detail}")


def test_training_loss_decrease():
    """Test 1: FusionModel training - loss should decrease"""
    print("\n=== Test 1: FusionModel Training Loss ===")
    
    config = FusionConfig(
        vocab_size=500,
        hidden_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=256,
        block_size=8,
        latent_dim=16,
        window_size=64,
    )
    model = FusionModel(config)
    model.train()
    model.to(DEVICE)
    
    # Count params
    param_count = sum(p.numel() for p in model.parameters())
    print(f"  Parameters: {param_count:,}")
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    batch_size, seq_len = 4, 32
    
    # Synthetic data
    input_ids = torch.randint(1, config.vocab_size, (batch_size, seq_len), device=DEVICE)
    
    losses = []
    for step in range(50):
        optimizer.zero_grad()
        outputs = model(input_ids=input_ids, labels=input_ids)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
        
        if (step + 1) % 25 == 0:
            print(f"    Step {step+1:3d}: Loss = {loss.item():.4f}")
    
    initial = losses[0]
    final = losses[-1]
    decreased = final < initial
    print(f"    Initial: {initial:.4f}, Final: {final:.4f}, Delta: {final - initial:+.4f}")
    check("Loss decreased over 50 steps", decreased, f"(final {final:.4f} >= initial {initial:.4f})")
    
    del model, optimizer
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return decreased


def test_thinking_dial_different_depths():
    """Test 2: Different thinking_depths produce different logits"""
    print("\n=== Test 2: ThinkingDialModel Depth Sensitivity ===")
    
    config = FusionConfig(
        vocab_size=500,
        hidden_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=256,
        block_size=8,
        latent_dim=16,
        window_size=64,
    )
    base_model = FusionModel(config)
    base_model.eval()
    base_model.to(DEVICE)
    
    thinking_config = ThinkingConfig(num_thinking_depths=4)
    td_model = ThinkingDialModel(base_model, thinking_config)
    
    input_ids = torch.randint(1, config.vocab_size, (1, 8), device=DEVICE)
    
    with torch.no_grad():
        outputs_depth0 = td_model.generate(input_ids, max_new_tokens=4, thinking_depth=0, do_sample=False)
        outputs_depth3 = td_model.generate(input_ids, max_new_tokens=4, thinking_depth=3, do_sample=False)
    
    same = torch.equal(outputs_depth0, outputs_depth3)
    check("Different depths produce different outputs", not same, "(outputs identical)")
    
    print(f"    Depth 0 output: {outputs_depth0[0, 8:].tolist()}")
    print(f"    Depth 3 output: {outputs_depth3[0, 8:].tolist()}")
    
    del td_model, base_model
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return not same


def test_grpo_train_step():
    """Test 3: GRPO train_step completes without errors"""
    print("\n=== Test 3: GRPO Train Step ===")
    
    config = FusionConfig(
        vocab_size=500,
        hidden_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=256,
        block_size=8,
        latent_dim=16,
        window_size=64,
    )
    base_model = FusionModel(config)
    base_model.train()
    base_model.to(DEVICE)
    
    thinking_config = ThinkingConfig(num_thinking_depths=4)
    td_model = ThinkingDialModel(base_model, thinking_config)
    td_model.train()
    
    grpo_config = GRPOConfig(
        grpo_sample_size=2,
        kl_coef=0.1,
    )
    trainer = GRPOTrainer(td_model, grpo_config=grpo_config, thinking_config=thinking_config)
    
    input_ids = torch.randint(1, config.vocab_size, (2, 8), device=DEVICE)
    
    try:
        result = trainer.train_step(input_ids, thinking_depth=2)
        check("train_step completed", True)
        check("train_step returned loss", "loss" in result, f"loss={result.get('loss')}")
        check("train_step returned rewards", "mean_reward" in result)
        check("step_count incremented", trainer.step_count == 1, f"step_count={trainer.step_count}")
        print(f"    Loss: {result['loss']:.4f}, Mean Reward: {result['mean_reward']:.4f}")
        # Note: loss can be 0 if all rewards are equal (identical dummy text)
    except Exception as e:
        check("train_step completed", False, f"Exception: {e}")
        import traceback
        traceback.print_exc()
    
    del td_model, base_model, trainer
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return True


def test_generate_loop():
    """Test 4: Full generate loop produces valid sequences"""
    print("\n=== Test 4: Full Generate Loop ===")
    
    config = FusionConfig(
        vocab_size=500,
        hidden_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=256,
        block_size=8,
        latent_dim=16,
        window_size=64,
    )
    model = FusionModel(config)
    model.eval()
    model.to(DEVICE)
    
    input_ids = torch.randint(1, config.vocab_size, (2, 8), device=DEVICE)
    
    with torch.no_grad():
        # Greedy
        out_greedy = model.generate(input_ids, max_new_tokens=16, do_sample=False)
        check("Greedy generate shape correct", out_greedy.shape == (2, 24), f"shape={out_greedy.shape}")
        
        # Sampling
        out_sample = model.generate(input_ids, max_new_tokens=16, do_sample=True, temperature=1.0)
        check("Sampled generate shape correct", out_sample.shape == (2, 24), f"shape={out_sample.shape}")
        
        # Greedy == Greedy (deterministic)
        out_greedy2 = model.generate(input_ids, max_new_tokens=16, do_sample=False)
        check("Greedy is deterministic", torch.equal(out_greedy, out_greedy2))
        
        # All tokens valid
        valid = (out_greedy >= 0).all() and (out_greedy < config.vocab_size).all()
        check("All output tokens valid", valid)
        
        # Prefix preserved
        prefix_match = torch.equal(out_greedy[:, :8], input_ids)
        check("Input prefix preserved", prefix_match)
    
    print(f"    Greedy output[0]: {out_greedy[0].tolist()}")
    print(f"    Sampled output[0]: {out_sample[0].tolist()}")
    
    del model
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return True


def test_thinking_dial_with_thinking_depth():
    """Test 5: ThinkingDial generate_with_thinking produces coherent results"""
    print("\n=== Test 5: ThinkingDial Thinking Depth Integration ===")
    torch.manual_seed(42)
    
    config = FusionConfig(
        vocab_size=500,
        hidden_size=128,
        num_hidden_layers=2,
        num_attention_heads=4,
        num_key_value_heads=2,
        intermediate_size=256,
        block_size=8,
        latent_dim=16,
        window_size=64,
    )
    base_model = FusionModel(config)
    base_model.eval()
    base_model.to(DEVICE)
    
    thinking_config = ThinkingConfig(num_thinking_depths=4)
    td_model = ThinkingDialModel(base_model, thinking_config)
    
    input_ids = torch.randint(1, config.vocab_size, (1, 8), device=DEVICE)
    
    results = {}
    for depth in range(4):
        with torch.no_grad():
            out = td_model.generate(input_ids, max_new_tokens=8, thinking_depth=depth, do_sample=False)
            results[depth] = out[0, 8:].tolist()
            print(f"    Depth {depth}: {results[depth]}")
    
    # All depths should produce valid sequences
    all_valid = all(
        all(0 <= t < config.vocab_size for t in results[d])
        for d in range(4)
    )
    check("All depths produce valid tokens", all_valid)
    
    # Note: with random weights, depth=0 vs depth=3 may sometimes produce
    # identical outputs. This is not a bug. Test 2 already verified depth
    # sensitivity in isolation. Here we just verify no crashes.
    check("All 4 depths generated without errors", True)
    
    del td_model, base_model
    torch.cuda.empty_cache() if torch.cuda.is_available() else None
    return True


if __name__ == "__main__":
    print("=" * 60)
    print("Fusion-LLM End-to-End Pipeline Validation")
    print("=" * 60)
    
    test_training_loss_decrease()
    test_thinking_dial_different_depths()
    test_grpo_train_step()
    test_generate_loop()
    test_thinking_dial_with_thinking_depth()
    
    print()
    print("=" * 60)
    print(f"Results: {PASS_COUNT} PASSED, {FAIL_COUNT} FAILED out of {PASS_COUNT + FAIL_COUNT}")
    if FAIL_COUNT == 0:
        print("ALL TESTS PASSED - Pipeline is runtime-verified!")
    else:
        print(f"FAILURES DETECTED - {FAIL_COUNT} test(s) need attention")
    print("=" * 60)
    
    sys.exit(1 if FAIL_COUNT > 0 else 0)