import argparse
import math
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import torch
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer
from torch.utils.data import DataLoader, TensorDataset

from b1_tensor_core import B1Linear
from src.r4t.b1_diffusion import B1EDMDenoiser
from src.r4t.config import DiffusionConfig
from src.r4t.diffusion import sample_edm
from src.r4t.journal import ExperimentJournal
from scripts.train_diffusion import ExponentialMovingAverage, diffusion_loss

DATA_PATH = ROOT / "data" / "diffusion_dataset_540k.pt"
TAXONOMY_PATH = ROOT / "data" / "taxonomy_embeddings.pt"
CHECKPOINT_DIR = ROOT / "checkpoints"
CHECKPOINT_DIR.mkdir(exist_ok=True)

TEST_QUERIES = [
    "running shoes",
    "wireless noise cancelling headphones",
    "organic gardening tools",
]


def load_taxonomy_bank(device):
    if not TAXONOMY_PATH.exists():
        return None, None
    tax_data = torch.load(TAXONOMY_PATH, map_location="cpu", weights_only=False)
    emb = F.normalize(tax_data["embeddings"].float(), dim=-1).to(device)
    names = tax_data["names"]
    return emb, names


def evaluate_qualitative(model, embedder, tax_emb, tax_names, device, tracker, run_name, sampling_steps=8):
    model.eval()
    all_alignments = []
    all_diversities = []

    for q_idx, q_text in enumerate(TEST_QUERIES):
        with torch.no_grad():
            z_q = embedder.encode([q_text], convert_to_tensor=True, normalize_embeddings=True, device=device).float()
            fanout = sample_edm(model, z_q, sampling_steps=sampling_steps, cfg_strength=0.1)[0]
            fanout = F.normalize(fanout, dim=-1)

            align_scores = torch.mm(fanout, z_q.T).squeeze(-1)
            mean_align = align_scores.mean().item()
            all_alignments.append(mean_align)

            sim_mat = torch.mm(fanout, fanout.T)
            mask = ~torch.eye(10, dtype=torch.bool, device=device)
            pairwise_sim = sim_mat[mask].mean().item()
            pairwise_div = 1.0 - pairwise_sim
            all_diversities.append(pairwise_div)

            decoded_subqueries = []
            if tax_emb is not None:
                top_matches = torch.mm(fanout, tax_emb.T)
                best_indices = top_matches.argmax(dim=-1).tolist()
                for slot_idx, tax_idx in enumerate(best_indices):
                    cat_name = tax_names[tax_idx]
                    slot_sim = top_matches[slot_idx, tax_idx].item()
                    decoded_subqueries.append(f"{cat_name} (sim: {slot_sim:.3f})")
            else:
                decoded_subqueries = [f"Vector #{i+1} [768-d]" for i in range(10)]

            tracker.log_sample(
                step=100 + q_idx,
                sample_type="fanout_query_expansion",
                input_payload={"query": q_text, "cfg_strength": 0.1, "sampling_steps": sampling_steps},
                output_payload={
                    "subqueries": decoded_subqueries,
                    "mean_prompt_align": mean_align,
                    "pairwise_diversity": pairwise_div,
                },
                target_payload=None,
                metadata={"run": run_name, "test_query_index": q_idx + 1, "sampling_steps": sampling_steps, "backend": "b1_tc"},
            )

    overall_align = float(sum(all_alignments) / max(1, len(all_alignments)))
    overall_div = float(sum(all_diversities) / max(1, len(all_diversities)))
    return overall_align, overall_div


def benchmark_latency(model, device, sampling_steps=8):
    model.eval()
    dummy_q = torch.randn(1, 768, device=device)
    dummy_q = F.normalize(dummy_q, dim=-1)

    # Warmup
    with torch.no_grad():
        sample_edm(model, dummy_q, sampling_steps=sampling_steps, cfg_strength=0.1)
    torch.cuda.synchronize()

    times = []
    for _ in range(10):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.no_grad():
            sample_edm(model, dummy_q, sampling_steps=sampling_steps, cfg_strength=0.1)
        torch.cuda.synchronize()
        times.append((time.perf_counter() - t0) * 1000.0)

    return float(sum(times) / len(times))


def main():
    parser = argparse.ArgumentParser(description="Train 1-Bit Tensor Core Diffusion Denoiser")
    parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs")
    parser.add_argument("--lr", type=float, default=1.5e-4, help="Learning rate")
    parser.add_argument("--batch_size", type=int, default=128, help="Batch size")
    parser.add_argument("--run_name", type=str, default="Sweep B1-TC: 1-Bit Tensor Core (2L, 10ep)", help="Journal run name")
    parser.add_argument("--checkpoint_name", type=str, default="b1_tc_10ep_champion.pt", help="Checkpoint filename")
    parser.add_argument("--pure_1bit", action="store_true", help="Binarize all outer projections and time MLP to 1-bit")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Compute device: {device} ({torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'})")

    print(f"Loading dataset from {DATA_PATH}...")
    dataset_dict = torch.load(DATA_PATH, map_location="cpu", weights_only=False)
    queries = dataset_dict["query_embeddings"].float()
    targets = dataset_dict["targets"].float()
    sigma_data = float(dataset_dict.get("sigma_data", 0.0361))
    N, L, D = targets.shape
    print(f"Loaded {N:,} samples: L={L}, D={D}, sigma_data={sigma_data:.4f}")

    # 90/10 split
    n_train = int(0.9 * N)
    train_queries, val_queries = queries[:n_train], queries[n_train:]
    train_targets, val_targets = targets[:n_train], targets[n_train:]

    train_dataset = TensorDataset(train_queries, train_targets)
    val_dataset = TensorDataset(val_queries, val_targets)

    batch_size = args.batch_size
    epochs = args.epochs
    lr = args.lr
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)

    config = DiffusionConfig(
        sequence_length=L,
        embedding_dim=D,
        hidden_dim=512,
        mlp_dim=1024,
        heads=8,
        layers=2,
        dropout=0.1,
        sigma_min=0.0001,
        sigma_max=80.0,
        sigma_data=sigma_data,
        sampling_steps=8,
    )

    mode_str = "Pure 1-Bit (End-to-End)" if args.pure_1bit else "Hybrid 1-Bit Core"
    print(f"\nInitializing B1EDMDenoiser with Hardware 1-Bit Tensor Cores ({mode_str}, backend='tc')...")
    model = B1EDMDenoiser(config, backend="tc", pure_1bit=args.pure_1bit).to(device)
    ema = ExponentialMovingAverage(model, decay=0.999)

    param_count = sum(p.numel() for p in model.parameters())
    print(f"Initialized B1EDMDenoiser ({param_count / 1e6:.2f}M parameters).")

    run_name = args.run_name
    journal = ExperimentJournal()
    tracker = journal.start_run(
        name=run_name,
        experiment_name="1-Bit Tensor Core Innovation",
        task_type="diffusion",
        config={
            "layers": 2,
            "hidden_dim": 512,
            "mlp_dim": 1024,
            "heads": 8,
            "lr": lr,
            "epochs": epochs,
            "batch_size": batch_size,
            "target_ordering": "cosine",
            "backend": "b1_tc",
            "precision": "1-bit_tc_ste",
            "dataset": "diffusion_dataset_540k.pt",
        },
        tags=["1bit", "tensor_core", "ptx_mma", "2layer", "ste"],
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    total_steps = len(train_loader) * epochs
    warmup_steps = len(train_loader) * 1

    warmup_sched = torch.optim.lr_scheduler.LinearLR(optimizer, start_factor=0.05, total_iters=warmup_steps)
    cosine_sched = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1, total_steps - warmup_steps), eta_min=lr * 0.05)
    scheduler = torch.optim.lr_scheduler.SequentialLR(optimizer, schedulers=[warmup_sched, cosine_sched], milestones=[warmup_steps])

    best_val_loss = float("inf")
    best_checkpoint_path = CHECKPOINT_DIR / args.checkpoint_name

    use_amp = torch.cuda.is_available()
    amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16

    print("\nStarting 1-Bit Tensor Core STE Diffusion Training:")
    print(f"  Epochs: {epochs} | Batch size: {batch_size} | Batches/epoch: {len(train_loader):,}")
    print(f"  Peak LR: {lr:.2e} | Target checkpoint: {best_checkpoint_path}")
    print("=" * 75)

    t0 = time.time()
    for epoch in range(1, epochs + 1):
        ep_t0 = time.time()
        model.train()
        train_loss_total = 0.0

        for b_queries, b_targets in train_loader:
            b_queries = b_queries.to(device, non_blocking=True)
            b_targets = b_targets.to(device, non_blocking=True)

            # Cosine sorting targets
            sims = torch.einsum("bd,bld->bl", F.normalize(b_queries, dim=-1), F.normalize(b_targets, dim=-1))
            sorted_idx = torch.argsort(sims, dim=1, descending=True)
            b_targets = torch.gather(b_targets, 1, sorted_idx.unsqueeze(-1).expand(-1, -1, D))

            optimizer.zero_grad()
            with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=use_amp):
                loss = diffusion_loss(model, b_targets, b_queries)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            # BNN STE Best Practice: clamp latent weights to [-1.0, 1.0]
            with torch.no_grad():
                for m in model.modules():
                    if isinstance(m, B1Linear):
                        m.weight.clamp_(-1.0, 1.0)

            scheduler.step()
            ema.update(model)

            train_loss_total += loss.item() * len(b_queries)

        train_loss = train_loss_total / len(train_dataset)

        # Validation
        model.eval()
        val_loss_total = 0.0
        val_gen = torch.Generator(device=device).manual_seed(1337)
        with torch.no_grad():
            for b_queries, b_targets in val_loader:
                b_queries = b_queries.to(device, non_blocking=True)
                b_targets = b_targets.to(device, non_blocking=True)

                sims = torch.einsum("bd,bld->bl", F.normalize(b_queries, dim=-1), F.normalize(b_targets, dim=-1))
                sorted_idx = torch.argsort(sims, dim=1, descending=True)
                b_targets = torch.gather(b_targets, 1, sorted_idx.unsqueeze(-1).expand(-1, -1, D))

                with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=use_amp):
                    v_loss = diffusion_loss(model, b_targets, b_queries, generator=val_gen)
                val_loss_total += v_loss.item() * len(b_queries)

        val_loss = val_loss_total / len(val_dataset)
        ep_time = time.time() - ep_t0
        elapsed = time.time() - t0

        is_best = val_loss < best_val_loss
        best_marker = " [BEST SAVED]" if is_best else ""
        print(f"Epoch [{epoch:2d}/{epochs:2d}] | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}{best_marker} | LR: {scheduler.get_last_lr()[0]:.2e} | Ep Time: {ep_time:.1f}s | Elapsed: {elapsed:.1f}s")

        if is_best:
            best_val_loss = val_loss
            torch.save({
                "epoch": epoch,
                "config": config,
                "model_state_dict": model.state_dict(),
                "ema_state_dict": ema.state_dict(),
                "val_loss": val_loss,
            }, best_checkpoint_path)

        tracker.log_metrics(
            step=epoch * len(train_loader),
            epoch=epoch,
            train_loss=train_loss,
            val_loss=val_loss,
            lr=scheduler.get_last_lr()[0],
            is_best=is_best,
        )

    print(f"\nTraining Complete! Best Val Loss: {best_val_loss:.4f}")
    
    # Load best EMA weights for evaluation
    ckpt = torch.load(best_checkpoint_path, map_location=device, weights_only=False)
    if "ema_state_dict" in ckpt and "shadow" in ckpt["ema_state_dict"]:
        shadow = ckpt["ema_state_dict"]["shadow"]
        model.load_state_dict({k: shadow[k].to(device) for k in shadow})
    else:
        model.load_state_dict(ckpt["model_state_dict"])
    model.eval()

    # Freeze weights for zero-overhead inference
    print("Freezing B1Linear weights for zero-overhead Tensor Core inference...")
    model.freeze_for_inference()

    if args.pure_1bit:
        export_dict = {"config": config, "weights": {}}
        for k, v in model.state_dict().items():
            if "packed_weight" in k:
                export_dict["weights"][k] = v.cpu()
            elif "bias" in k or "norm" in k or "position" in k:
                export_dict["weights"][k] = v.to(torch.float16).cpu()
        out_sub1mb = CHECKPOINT_DIR / "champion_pure_1bit_sub1mb.pt"
        torch.save(export_dict, out_sub1mb)
        print(f"Exported Pure 1-Bit sub-1MB checkpoint: {out_sub1mb} ({out_sub1mb.stat().st_size / 1024:.1f} KB / {out_sub1mb.stat().st_size / (1024*1024):.2f} MB)")

    # Benchmarking latency
    lat_8 = benchmark_latency(model, device, sampling_steps=8)
    lat_16 = benchmark_latency(model, device, sampling_steps=16)
    print(f"Latency benchmarks: 8-Step Heun = {lat_8:.2f} ms | 16-Step Heun = {lat_16:.2f} ms")

    # Qualitative decoding
    tax_emb, tax_names = load_taxonomy_bank(device)
    print("Loading Google EmbeddingGemma embedder...")
    embedder = SentenceTransformer("google/embeddinggemma-300m", model_kwargs={"torch_dtype": torch.bfloat16}, device=device)

    align, div = evaluate_qualitative(
        model=model,
        embedder=embedder,
        tax_emb=tax_emb,
        tax_names=tax_names,
        device=device,
        tracker=tracker,
        run_name=run_name,
        sampling_steps=8,
    )

    tracker.log_benchmark(
        latency_us=int(lat_8 * 1000),
        throughput_items_per_sec=1000.0 / lat_8 if lat_8 > 0 else 0.0,
        device_name=torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU",
        notes=f"1-Bit Tensor Core Denoiser (backend='tc', 8-step Heun); 16-step latency: {lat_16:.2f}ms",
    )

    tracker.finish(
        status="completed",
        summary_metrics={
            "best_val_loss": best_val_loss,
            "prompt_alignment": align,
            "pairwise_diversity": div,
            "sampling_latency_ms": lat_8,
            "sampling_latency_16step_ms": lat_16,
            "backend": "tc_1bit",
            "layers": 2,
        },
    )
    print(f"\nRun Complete: Val Loss={best_val_loss:.4f}, Align={align:.3f}, Div={div:.3f}, Latency(8-step)={lat_8:.2f}ms")


if __name__ == "__main__":
    main()