import os
import math
import torch

from safetensors import safe_open
from safetensors.torch import load_file, save_file


# =========================================================
# PATHS
# =========================================================

H3_EMB = r"E:\Models SD XL\UNET\Minimax_H3\h3_qwen3vl32b_embeddings_bf16.safetensors"

SENSENOVA = r"D:\ComfyUI_Python312\ComfyUI\models\unet\SenseNova-U1.5-8B-MoT-pruned-bf16.safetensors"

OUT = r"E:\Models SD XL\UNET\Minimax_H3\sensenova_to_h3_embedding_bridge_rank512.safetensors"


# =========================================================
# SETTINGS
# =========================================================

INPUT_DIM = 4096
OUTPUT_DIM = 5120
RANK = 512

BATCH_SIZE = 256
EPOCHS = 5
LR = 1e-3

VAL_SIZE = 12000

SEED = 1234


# =========================================================
# DEVICE
# =========================================================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("=" * 80)
print("SenseNova -> MiniMax H3 embedding bridge test")
print("=" * 80)

print("Device:", device)

if device.type == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))
    print(
        "VRAM:",
        round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2),
        "GB"
    )


# =========================================================
# LOAD H3 EMBEDDINGS
# =========================================================

print()
print("Loading H3 embeddings...")

h3_file = load_file(H3_EMB, device="cpu")

h3 = h3_file["model.embed_tokens.weight"]

print("H3 shape:", tuple(h3.shape))
print("H3 dtype:", h3.dtype)


# =========================================================
# LOAD SENSENOVA EMBEDDINGS ONLY
# =========================================================

print()
print("Loading SenseNova token embeddings...")

with safe_open(SENSENOVA, framework="pt", device="cpu") as f:

    key = "language_model.model.embed_tokens.weight"

    if key not in f.keys():
        raise RuntimeError(
            f"SenseNova embedding tensor not found: {key}"
        )

    sense = f.get_tensor(key)


print("SenseNova shape:", tuple(sense.shape))
print("SenseNova dtype:", sense.dtype)


# =========================================================
# VALIDATE
# =========================================================

if h3.shape[0] != sense.shape[0]:
    raise RuntimeError(
        f"Vocabulary mismatch: H3={h3.shape[0]}, SenseNova={sense.shape[0]}"
    )

if sense.shape[1] != INPUT_DIM:
    raise RuntimeError(
        f"Unexpected SenseNova hidden size: {sense.shape[1]}"
    )

if h3.shape[1] != OUTPUT_DIM:
    raise RuntimeError(
        f"Unexpected H3 hidden size: {h3.shape[1]}"
    )


VOCAB = h3.shape[0]

print()
print("Vocabulary rows:", VOCAB)


# =========================================================
# NORMALIZATION
# =========================================================

# Convert storage to float32 only per batch later.
# Calculate global average norms using manageable chunks.

def average_norm(x, chunk=4096):

    total = 0.0
    count = 0

    for start in range(0, x.shape[0], chunk):

        end = min(start + chunk, x.shape[0])

        b = x[start:end].float()

        total += b.norm(dim=1).sum().item()
        count += b.shape[0]

    return total / count


print()
print("Calculating embedding norm statistics...")

sense_norm = average_norm(sense)
h3_norm = average_norm(h3)

print("SenseNova avg norm:", sense_norm)
print("H3 avg norm:       ", h3_norm)


# =========================================================
# TRAIN / VALIDATION SPLIT
# =========================================================

g = torch.Generator()
g.manual_seed(SEED)

perm = torch.randperm(VOCAB, generator=g)

val_ids = perm[:VAL_SIZE]
train_ids = perm[VAL_SIZE:]

print()
print("Train rows:", len(train_ids))
print("Validation rows:", len(val_ids))


# =========================================================
# LOW-RANK PROJECTOR
#
# 4096 -> 512 -> 5120
# no activation: this is a low-rank linear mapping
# =========================================================

A = torch.nn.Linear(
    INPUT_DIM,
    RANK,
    bias=False,
    device=device,
    dtype=torch.float32,
)

B = torch.nn.Linear(
    RANK,
    OUTPUT_DIM,
    bias=False,
    device=device,
    dtype=torch.float32,
)


torch.nn.init.normal_(
    A.weight,
    std=1.0 / math.sqrt(INPUT_DIM)
)

torch.nn.init.zeros_(B.weight)


params = list(A.parameters()) + list(B.parameters())

optimizer = torch.optim.AdamW(
    params,
    lr=LR,
    weight_decay=1e-4,
)


# =========================================================
# LOSS
# =========================================================

def bridge_forward(x):

    return B(A(x))


def loss_fn(pred, target):

    # Directional similarity.
    cosine = 1.0 - torch.nn.functional.cosine_similarity(
        pred,
        target,
        dim=-1
    ).mean()

    # Normalize magnitude sensitivity.
    pred_n = torch.nn.functional.normalize(pred, dim=-1)
    target_n = torch.nn.functional.normalize(target, dim=-1)

    mse = torch.nn.functional.mse_loss(
        pred_n,
        target_n
    )

    return cosine + mse


# =========================================================
# EVALUATION
# =========================================================

@torch.no_grad()
def evaluate(ids):

    A.eval()
    B.eval()

    cos_sum = 0.0
    mse_sum = 0.0
    count = 0

    for start in range(0, len(ids), BATCH_SIZE):

        idx = ids[start:start + BATCH_SIZE]

        x = sense[idx].float().to(device)
        y = h3[idx].float().to(device)

        pred = bridge_forward(x)

        cos = torch.nn.functional.cosine_similarity(
            pred,
            y,
            dim=-1
        )

        mse = torch.nn.functional.mse_loss(
            pred,
            y,
            reduction="none"
        ).mean(dim=1)

        cos_sum += cos.sum().item()
        mse_sum += mse.sum().item()

        count += len(idx)

    return cos_sum / count, mse_sum / count


# =========================================================
# BASELINE
# =========================================================

print()
print("=" * 80)
print("TRAINING")
print("=" * 80)


# =========================================================
# TRAIN
# =========================================================

for epoch in range(EPOCHS):

    A.train()
    B.train()

    shuffled = train_ids[
        torch.randperm(len(train_ids), generator=g)
    ]

    running = 0.0
    batches = 0

    for start in range(0, len(shuffled), BATCH_SIZE):

        idx = shuffled[start:start + BATCH_SIZE]

        x = sense[idx].float().to(device)
        y = h3[idx].float().to(device)

        optimizer.zero_grad(set_to_none=True)

        pred = bridge_forward(x)

        loss = loss_fn(pred, y)

        loss.backward()

        torch.nn.utils.clip_grad_norm_(
            params,
            1.0
        )

        optimizer.step()

        running += loss.item()
        batches += 1

        if batches % 100 == 0:

            print(
                f"Epoch {epoch + 1}/{EPOCHS} | "
                f"batch {batches} | "
                f"loss {running / batches:.6f}"
            )

    val_cos, val_mse = evaluate(val_ids)

    print()
    print(
        f"EPOCH {epoch + 1} COMPLETE | "
        f"train_loss={running / batches:.6f} | "
        f"val_cos={val_cos:.6f} | "
        f"val_mse={val_mse:.8f}"
    )
    print()


# =========================================================
# FINAL EVALUATION
# =========================================================

val_cos, val_mse = evaluate(val_ids)

print("=" * 80)
print("FINAL")
print("=" * 80)

print("Validation cosine similarity:", val_cos)
print("Validation MSE:", val_mse)


# =========================================================
# SAVE
# =========================================================

save_file(
    {
        "proj_in.weight": A.weight.detach().cpu().to(torch.float16),
        "proj_out.weight": B.weight.detach().cpu().to(torch.float16),
    },
    OUT,
    metadata={
        "source": "SenseNova-U1.5-8B-MoT",
        "target": "MiniMax-H3-Qwen3VL32B",
        "input_dim": str(INPUT_DIM),
        "rank": str(RANK),
        "output_dim": str(OUTPUT_DIM),
        "validation_cosine": str(val_cos),
        "validation_mse": str(val_mse),
    }
)

print()
print("Saved bridge:")
print(OUT)
print()
print("DONE")