import sys
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
from config import config
from text.japanese import text2sep_kata

MODEL_ID = "ku-nlp/deberta-v2-large-japanese-char-wwm"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

models = dict()


def get_bert_feature(
    text,
    word2ph,
    device=config.bert_gen_config.device,
    style_text=None,
    style_weight=0.7,
):
    text = "".join(text2sep_kata(text)[0])
    if style_text:
        style_text = "".join(text2sep_kata(style_text)[0])

    if sys.platform == "darwin" and torch.backends.mps.is_available() and device == "cpu":
        device = "mps"
    if not device:
        device = "cuda"

    if device not in models:
        if config.webui_config.fp16_run:
            models[device] = AutoModelForMaskedLM.from_pretrained(
                MODEL_ID, torch_dtype=torch.float16
            ).to(device)
        else:
            models[device] = AutoModelForMaskedLM.from_pretrained(MODEL_ID).to(device)

    with torch.no_grad():
        inputs = tokenizer(text, return_tensors="pt", add_special_tokens=True)
        for k in inputs:
            inputs[k] = inputs[k].to(device)
        res = models[device](**inputs, output_hidden_states=True)
        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].float().cpu()

        if style_text:
            style_inputs = tokenizer(style_text, return_tensors="pt", add_special_tokens=True)
            for k in style_inputs:
                style_inputs[k] = style_inputs[k].to(device)
            style_res = models[device](**style_inputs, output_hidden_states=True)
            style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].float().cpu()
            style_res_mean = style_res.mean(0)

    # ✅ Force truncate ให้ความยาวตรงกับ word2ph
    min_len = min(len(word2ph), res.shape[0])
    word2phone = word2ph[:min_len]
    res = res[:min_len]

    phone_level_feature = []
    for i in range(len(word2phone)):
        if style_text:
            repeat_feature = (
                res[i].repeat(word2phone[i], 1) * (1 - style_weight)
                + style_res_mean.repeat(word2phone[i], 1) * style_weight
            )
        else:
            repeat_feature = res[i].repeat(word2phone[i], 1)
        phone_level_feature.append(repeat_feature)

    phone_level_feature = torch.cat(phone_level_feature, dim=0)
    return phone_level_feature.T