import os
# os.environ['TORCH_LOGS'] = '+dynamic'
# os.environ['TORCH_LOGS'] = '+export'
# os.environ['TORCHDYNAMO_EXTENDED_DEBUG_GUARD_ADDED']="u0 >= 0"
# os.environ['TORCHDYNAMO_EXTENDED_DEBUG_CPP']="1"
# os.environ['TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL']="u0"


from kokoro import phonemize, tokenize, length_to_mask
import torch.nn.functional as F
from models_scripting import build_model
import torch
from typing import Dict

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

model = build_model('kokoro-v0_19.pth', device)

voicepack = torch.load('voices/af.pt', weights_only=True).to(device)

speed = 1.

text = "How could I know? It's an unanswerable question. Like asking an unborn child if they'll lead a good life. They haven't even been born."

ps = phonemize(text, "a")
tokens = tokenize(ps)

tokens = torch.LongTensor([[0, *tokens, 0]]).to(device)

class StyleTTS2(torch.nn.Module):
    def __init__(self, model, voicepack):
        super().__init__()
        # self.model = model
        self.bert = model.bert
        self.bert_encoder = model.bert_encoder
        self.predictor = model.predictor
        self.decoder = model.decoder
        self.text_encoder = model.text_encoder
        self.voicepack = voicepack
    
    def forward(self, tokens : torch.Tensor):
        speed = 1.
        # tokens = torch.nn.functional.pad(tokens, (0, 510 - tokens.shape[-1]))
        device = tokens.device
        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)

        text_mask = length_to_mask(input_lengths).to(device)
        bert_dur = self.bert(tokens)

        d_en = self.bert_encoder(bert_dur).transpose(-1, -2)

        ref_s = self.voicepack[tokens.shape[1]]
        s = ref_s[:, 128:]

        d = self.predictor.text_encoder.inference(d_en, s)
        x, _ = self.predictor.lstm(d)

        duration = self.predictor.duration_proj(x)
        duration = torch.sigmoid(duration).sum(axis=-1) / speed
        pred_dur = torch.round(duration).clamp(min=1).long()
        
        c_start = F.pad(pred_dur,(1,0), "constant").cumsum(dim=1)[0,0:-1]
        c_end = c_start + pred_dur[0,:]

        # torch._check(pred_dur.sum().item()>0, lambda: print(f"Got {pred_dur.sum().item()}"))
        indices = torch.arange(0, pred_dur.sum().item()).long().to(device)

        pred_aln_trg_list=[]
        for cs, ce in zip(c_start, c_end):
            row = torch.where((indices>=cs) & (indices<ce), 1., 0.)
            pred_aln_trg_list.append(row)
        pred_aln_trg=torch.vstack(pred_aln_trg_list)
            
        en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device)
        
        F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
        t_en = self.text_encoder.inference(tokens)
        asr = t_en @ pred_aln_trg.unsqueeze(0).to(device)
        return (asr, F0_pred, N_pred, ref_s[:, :128])
        # output = self.model.decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze().detach().cpu().numpy()


# bert = torch.jit.script(model.bert)
# bert_encoder = torch.jit.script(model.bert_encoder)
# predictor = torch.jit.script(model.predictor)
# text_encoder = torch.jit.script(model.text_encoder)

# model["bert"] = torch.jit.trace(model["bert"], (tokens, ))
# # model["decoder"] = torch.jit.script(model["decoder"])
# bert_dur = model["bert"](tokens)
# model["bert_encoder"] = torch.jit.trace(model["bert_encoder"], (bert_dur,))
# model["predictor"] = torch.compile(model["predictor"], backend=backend)
# model["text_encoder"] = torch.compile(model["text_encoder"], backend=backend)

style_model = StyleTTS2(model=model, voicepack=voicepack)
style_model.eval()
# style_model = torch.jit.trace_module(style_model.eval(), inputs={'forward': (tokens, )}) 
# style_model.model["predictor"].F0Ntrain = torch.jit.script(style_model.model["predictor"].F0Ntrain)
(asr, F0_pred, N_pred, ref_s) = style_model(tokens)
print(asr.shape, F0_pred.shape, N_pred.shape, ref_s.shape)

# scripted_style_model = torch.jit.script(style_model)

# (asr, F0_pred, N_pred, ref_s) = scripted_style_model(tokens)
# print(asr.shape, F0_pred.shape, N_pred.shape, ref_s.shape)

# torch.onnx.export(scripted_style_model, ( tokens, ), "style_model.onnx", verbose=True, opset_version=17, input_names=["tokens"], output_names=["asr", "F0_pred", "N_pred", "ref_s"])
# token_len = torch.export.Dim("token_len", min=2, max=510)
# batch = torch.export.Dim("batch")
# dynamic_shapes = {"tokens":{ 1:token_len}}
dynamic_shapes = {"tokens":{ 1:"token_len"}}
print(f"{tokens.shape=}")
torch.onnx.export(model=style_model, args=( tokens, ), dynamic_axes=dynamic_shapes, input_names=["tokens"], f="style_model.onnx", 
                output_names=["asr", "F0_pred", "N_pred", "ref_s"], opset_version=13, verbose=False, dynamo=False)


# with torch.no_grad():
# torch.export.export(style_model, args=( tokens, ), dynamic_shapes=dynamic_shapes, strict=False)

# export_mod = torch.export.export(style_model, args=( tokens, ), strict=False)