import torch
import numpy as np
import torch.nn.functional as F
from utils import *
class SLMAdversarialLoss(torch.nn.Module):

    def __init__(self, model, wl, sampler, min_len, max_len, batch_percentage=0.5, skip_update=10, sig=1.5):
        super(SLMAdversarialLoss, self).__init__()
        self.model = model
        self.wl = wl
        self.sampler = sampler
        
        self.min_len = min_len
        self.max_len = max_len
        self.batch_percentage = batch_percentage
        self.sig = sig
        self.skip_update = skip_update
        
    def forward(self,  iters, raw_texts, mean, gemma_mrl, gemma_tokenizer, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_text, ref_lengths, use_ind, s_trg, ref_s=None):
        text_mask = length_to_mask(ref_lengths).to(ref_text.device)
        #bert_dur = self.model.bert(ref_text, attention_mask=(~text_mask).int())
        #d_en = self.model.bert_encoder(bert_dur).transpose(-1, -2) 
        token_id_start_generation = 2
        token_id_end_generation = 1
        token_id_unused_0_separator =6
        token_id_mask  =4
        exclude_token_ids = {token_id_unused_0_separator, token_id_start_generation, token_id_end_generation}
        concat_texts = []
        for ctx in raw_texts:
            parts = []
            idxs = [0, 1, 2, 3]  # prev2, prev1, cur, next1
            for idx in idxs:
                txt = ""
                try:
                    txt = ctx[idx] if ctx[idx] is not None else ""
                except Exception:
                    txt = ""
                parts.append(txt.strip().lower())
            concat = "<unused0>".join([p for p in parts if p != ""])
            concat_texts.append(concat)

        tokens = gemma_tokenizer(
            concat_texts,
            padding=True,
            truncation=True,
            return_tensors="pt",
            max_length=512,
        )
        input_ids = tokens["input_ids"].to("cuda")          # original ids (used to locate seps, bos, eos)
        attention_mask_g = tokens["attention_mask"].to("cuda")
        
        p_mask = 0.15
        p_shuffle = 0.05
        aug_input_ids = input_ids.clone()
        aug_input_ids = mask_input_ids_random(
            aug_input_ids,
            attention_mask_g,
            mask_token_id=token_id_mask,
            p_mask=p_mask,
            exclude_token_ids=exclude_token_ids,
            seed=None,
        )
        aug_input_ids = shuffle_segments_by_separator(
            aug_input_ids,
            attention_mask_g,
            sep_token_id=token_id_unused_0_separator,
            p_shuffle=p_shuffle,
            bos_token_id=token_id_start_generation,
            eos_token_id=token_id_end_generation,
            seed=None,
        )
        raw_hidden = self.model.gemma_encoder(
            input_ids=aug_input_ids,
            attention_mask=attention_mask_g,
        ).last_hidden_state  # [B, L, D]
        raw_hidden = gemma_mrl(raw_hidden, attention_mask=attention_mask_g)  # (B, T, D)

        
        gated_hidden, context_weights = self.model.conditioner(
            raw_hidden,
            attention_mask_g,
            input_ids=input_ids,  # original ids used to find separators
            sep_token_id=token_id_unused_0_separator,
            current_context_idx=2,
            return_weights=True,
        )
    
        
        if use_ind and np.random.rand() < 0.5:
            s_preds = s_trg
        else:
            loss_diff, s_preds = self.model.flowMatching(
                    s_trg.unsqueeze(1),
                    embedding=gated_hidden,
            )
            s_preds = s_preds.squeeze(1) + mean
            '''    
            with torch.no_grad():
                num_steps = np.random.randint(10, 16 + 1)
                s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
                  embedding=gated_hidden,
                  embedding_scale=1,
                     embedding_mask_proba=0.1,
                     num_steps=num_steps).squeeze(1)
                s_preds = s_preds+ mean
            '''
     
            
        s = s_preds[:, :128]
        s_dur = s_preds[:, 128:256]
        s_pitch = s_preds[:, 256:]
        #print(s_pitch.shape)
        #print(s_dur.shape)
        
        d, _ = self.model.predictor(ref_text, s_dur,s_pitch, 
                                                ref_lengths, 
                                                torch.randn(ref_lengths.shape[0], ref_lengths.max(), 2).to(ref_text.device), 
                                                text_mask)
        
        bib = 0

        output_lengths = []
        attn_preds = []
        
        # differentiable duration modeling
        for _s2s_pred, _text_length in zip(d, ref_lengths):

            logits = _s2s_pred[:_text_length, :]
            probs  = torch.sigmoid(logits)  # probabilities for duration stats / kernel

            # per-token expected durations (sum over frames)
            dur_tok = probs.sum(dim=-1)  # [T]

            # predicted total frames l (clamped to safe range)
            l_float = dur_tok.sum()  # scalar
            l_min = max(1, int(self.min_len // 2))   # ensure nonzero and not too tiny
            l_max = max(l_min, int(self.max_len // 2))
            l = int(torch.clamp(l_float.round(), min=l_min, max=l_max).item())

            # time grid and token centers
            t = torch.arange(l, device=ref_text.device, dtype=logits.dtype)
            t = t.unsqueeze(0).expand(int(_text_length), l)  # [T, l]
            loc = torch.cumsum(dur_tok, dim=0) - 0.5 * dur_tok  # [T]
            # keep centers inside [0, l-1] to avoid extreme kernels
            loc = loc.clamp(min=0, max=max(0, l - 1))

            # Gaussian kernel per token built from probabilities
            h = torch.exp(-0.5 * ((t - loc.unsqueeze(-1)) / self.sig) ** 2)  # [T, l]
            h = h / (h.sum(dim=-1, keepdim=True) + 1e-8)
            # optional: normalize per-token kernels (not required if softmax later)
            # h = h / (h.sum(dim=-1, keepdim=True) + 1e-6)

            # depthwise conv over frames: one kernel per token/channel
            logits = logits.float()
            h = h.to(logits.dtype)
            out = F.conv1d(logits.unsqueeze(0), h.unsqueeze(1), padding=l-1, groups=int(_text_length))[..., :l]
            attn = F.softmax(out.squeeze(0).float(), dim=0)  # softmax across tokens
            attn_preds.append(attn)
            output_lengths.append(l)

        max_len = max(output_lengths)
        
        with torch.no_grad():
            t_en = self.model.text_encoder(ref_text, ref_lengths, text_mask)
            
        s2s_attn = torch.zeros(len(ref_lengths), int(ref_lengths.max()), max_len).to(ref_text.device)
        for bib in range(len(output_lengths)):
            s2s_attn[bib, :ref_lengths[bib], :output_lengths[bib]] = attn_preds[bib]

        asr_pred = t_en @ s2s_attn

        _, p_pred = self.model.predictor(ref_text, s_dur, s_pitch,
                                                ref_lengths, 
                                                s2s_attn, 
                                                text_mask)
        
        mel_len = max(int(min(output_lengths) / 2 - 1), self.min_len // 2)
        mel_len = min(mel_len, self.max_len // 2)
        
        # get clips
        
        en = []
        p_en = []
        sp = []
        
        F0_fakes = []
        N_fakes = []
        
        wav = []

        for bib in range(len(output_lengths)):
            mel_length_pred = output_lengths[bib]
            mel_length_gt = int(mel_input_length[bib].item() / 2)
            if mel_length_gt <= mel_len or mel_length_pred <= mel_len:
                continue

            sp.append(s_preds[bib])

            random_start = np.random.randint(0, mel_length_pred - mel_len)
            en.append(asr_pred[bib, :, random_start:random_start+mel_len])
            p_en.append(p_pred[bib, :, random_start:random_start+mel_len])

            # get ground truth clips
            random_start = np.random.randint(0, mel_length_gt - mel_len)
            y = waves[bib][(random_start * 2) * 600:((random_start+mel_len) * 2) * 600]
            wav.append(torch.from_numpy(y).to(ref_text.device))
            
            if len(wav) >= self.batch_percentage * len(waves): # prevent OOM due to longer lengths
                break

        if len(sp) <= 1:
            return None
            
        sp = torch.stack(sp)
        wav = torch.stack(wav).float()
        en = torch.stack(en)
        p_en = torch.stack(p_en)
        
        F0_fake, N_fake = self.model.predictor.F0Ntrain(p_en, sp[:, 256:])
        y_pred, _ , _ = self.model.decoder(en, F0_fake, N_fake, sp[:, :128])
        
        # discriminator loss
        if (iters + 1) % self.skip_update == 0:
            if np.random.randint(0, 2) == 0:
                wav = y_rec_gt_pred
                use_rec = True
            else:
                use_rec = False

            crop_size = min(wav.size(-1), y_pred.size(-1))
            if use_rec: # use reconstructed (shorter lengths), do length invariant regularization
                if wav.size(-1) > y_pred.size(-1):
                    real_GP = wav[:, : , :crop_size]
                    out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
                    out_org = self.wl.discriminator_forward(wav.detach().squeeze())
                    loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])

                    if np.random.randint(0, 2) == 0:
                        d_loss = self.wl.discriminator(real_GP.detach().squeeze(), y_pred.detach().squeeze()).mean()
                    else:
                        d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
                else:
                    real_GP = y_pred[:, : , :crop_size]
                    out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
                    out_org = self.wl.discriminator_forward(y_pred.detach().squeeze())
                    loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])

                    if np.random.randint(0, 2) == 0:
                        d_loss = self.wl.discriminator(wav.detach().squeeze(), real_GP.detach().squeeze()).mean()
                    else:
                        d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
                
                # regularization (ignore length variation)
                d_loss += loss_reg

                out_gt = self.wl.discriminator_forward(y_rec_gt.detach().squeeze())
                out_rec = self.wl.discriminator_forward(y_rec_gt_pred.detach().squeeze())

                # regularization (ignore reconstruction artifacts)
                d_loss += F.l1_loss(out_gt, out_rec)

            else:
                d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
        else:
            d_loss = 0
        gen_loss = self.wl.generator(y_pred.squeeze())
        
        gen_loss = gen_loss.mean()
        
        return d_loss, gen_loss, y_pred.detach().cpu().numpy()    
    
def length_to_mask(lengths):
    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
    mask = torch.gt(mask+1, lengths.unsqueeze(1))
    return mask