import os, sys
import tensorflow as tf
import numpy as np
from tensorflow.keras.losses import Loss, MeanSquaredError

seed = 42
tf.random.set_seed(seed)
np.random.seed(seed)

def sequence_cross_entropy(speech_label, text_label, logits, reduction='sum'):
    """
    args
    speech_label        : [B, Ls]
    text_label          : [B, Lt]
    logits              : [B, Lt]
    logits._keras_mask  : [B, Lt]
    """
    # Data pre-processing
    if tf.shape(text_label)[1] > tf.shape(speech_label)[1]:
        speech_label =  tf.pad(speech_label, [[0, 0],[0, tf.shape(text_label)[1] - tf.shape(speech_label)[1]]], 'CONSTANT', constant_values=0)
    elif tf.shape(text_label)[1] < tf.shape(speech_label)[1]:
        speech_label = speech_label[:, :text_label.shape[1]]
    
    # Make paired data between text and speech phonemes
    paired_label = tf.math.equal(text_label, speech_label)
    paired_label = tf.cast(tf.math.logical_and(tf.cast(paired_label, tf.bool), tf.cast(logits._keras_mask, tf.bool)), tf.float32)
    paired_label = tf.reshape(tf.ragged.boolean_mask(paired_label, tf.cast(logits._keras_mask, tf.bool)).flat_values, [-1,1])
    logits = tf.reshape(tf.ragged.boolean_mask(logits, tf.cast(logits._keras_mask, tf.bool)).flat_values, [-1,1])
    
    # Get BinaryCrossEntropy loss
    BCE = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.SUM)
    loss = BCE(paired_label, logits)
    
    if reduction == 'sum':
        loss = tf.math.divide_no_nan(loss, tf.cast(tf.shape(logits)[0], loss.dtype))
        loss = tf.math.multiply_no_nan(loss, tf.cast(tf.shape(speech_label)[0], loss.dtype))

    return loss

def detection_loss(y_true, y_pred):
    BFC = tf.keras.losses.BinaryCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.SUM)
    return(BFC(y_true, y_pred))

def ctc_loss(affinity_matrix, speech_labels, text_labels,n_speech):
    #logit_length
    # n_speech = tf.math.reduce_sum(tf.cast(affinity_matrix._keras_mask, tf.float32), -1)
    
    #logit
    transposed_logits = tf.transpose(affinity_matrix, perm=[0, 2, 1])
    # log_probs = tf.math.log(transposed_logits+ 1e-8)
    # logits_approx = log_probs - tf.reduce_max(log_probs, axis=-1, keepdims=True)
    
    #label
    matches = tf.equal(speech_labels, text_labels)
    indices = tf.range(text_labels.shape[1], dtype=tf.int32)
    selected_indices = tf.where(matches, indices, tf.fill(tf.shape(text_labels), 0))
    labels = tf.where(tf.equal(text_labels, 0), text_labels, selected_indices)
    
    #label_length
    label_length = tf.math.count_nonzero(labels, axis=1)
    
    ctc_loss = tf.nn.ctc_loss(labels,transposed_logits,label_length,n_speech,
                   logits_time_major=False,
                    unique=None,
                    blank_index=0,
                    name=None)
    
    return ctc_loss

class TotalLoss(Loss):
    def __init__(self, weight=1.0):
        super().__init__()
        self.weight = weight
    
    def __call__(self, y_true, y_pred, reduction='sum'):
        LD = detection_loss(y_true, y_pred)

        return self.weight * LD, LD


class TotalLoss_SCE(Loss):
    def __init__(self, weight=[1.0, 1.0, 0.2]):
        super().__init__()
        self.weight = weight
    
    def __call__(self, y_true, y_pred, speech_label, text_label, logit,affinity_matrix,n_speech, reduction='sum'):
        ctc = ctc_loss(affinity_matrix, speech_label, text_label,n_speech)
        
        if self.weight[0] != 0.0:   
            LD = detection_loss(y_true, y_pred)
        else:
            LD = 0
        if self.weight[1] != 0.0:
            LC = sequence_cross_entropy(speech_label, text_label, logit, reduction=reduction)
        else:
            LC = 0
        return self.weight[0] * LD + self.weight[1] * LC + self.weight[2]*ctc, LD, LC