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 import build_model
import torch
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)

model = MODEL
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.voicepack = voicepack
    
    def forward(self, tokens):
        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.model['bert'](tokens, attention_mask=(~text_mask).int())

        d_en = self.model["bert_encoder"](bert_dur).transpose(-1, -2)

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

        d = self.model["predictor"].text_encoder.inference(d_en, s)
        x, _ = self.model["predictor"].lstm(d)

        duration = self.model["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.model["predictor"].F0Ntrain(en, s)
        t_en = self.model["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()


style_model = StyleTTS2(model=model, voicepack=voicepack)
(asr, F0_pred, N_pred, ref_s) = style_model(tokens)

token_len = torch.export.Dim("token_len", min=2, max=510)
batch = torch.export.Dim("batch")
dynamic_shapes = {"tokens":{0:batch, 1:token_len}}

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

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