import torch
import torch.nn.functional as F

def warmup_lr(step):
    return min(step, 5000) / 5000


def training_with_esm(x_0_con, x_0_dis, trainer, trainer_dis, ns_con, ns_dis, trans, FLAGS, cls_head, class_weights):

    t = torch.randint(FLAGS.T, size=(x_0_con['input_ids'].shape[0], ), device=x_0_con['input_ids'].device)
    pt = torch.ones_like(t).float() / FLAGS.T

    embeddings = trainer.model.get_embeddings(x_0_con)

    noise = torch.randn_like(embeddings)
    x_t_con = trainer.make_x_t(embeddings, t, noise)
    torch.cuda.empty_cache()
    log_x_start = torch.log(x_0_dis.float().clamp(min=1e-30))
    x_t_dis = trainer_dis.q_sample(log_x_start=log_x_start, t=t)

    eps = trainer.model(x_t_con, t, x_t_dis.to(x_t_con.device))
    ps_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps)
    con_loss = F.mse_loss(eps, noise, reduction='none')
    con_loss = con_loss.mean()
    kl, ps_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, x_t_con)
    ps_0_dis = torch.exp(ps_0_dis)
    kl_prior = trainer_dis.kl_prior(log_x_start)
    dis_loss = (kl / pt + kl_prior).mean()
    torch.cuda.empty_cache()
    noise_ns = torch.randn_like(ns_con)
    ns_t_con = trainer.make_x_t(ns_con, t, noise_ns)
    log_ns_start = torch.log(ns_dis.float().clamp(min=1e-30))
    ns_t_dis = trainer_dis.q_sample(log_x_start=log_ns_start, t=t)
    torch.cuda.empty_cache()
    eps_ns = trainer.model(x_t_con, t, ns_t_dis.to(ns_t_dis.device))
    ns_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps_ns)
    _, ns_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, ns_t_con)
    ns_0_dis = torch.exp(ns_0_dis)


    torch.cuda.empty_cache()
    triplet_loss = torch.nn.TripletMarginLoss(margin=0.5, p=2)
    triplet_con = triplet_loss(embeddings, ps_0_con, ns_0_con)
    total_loss = con_loss
    st=0
    triplet_dis = []
    for item in trans.output_info:
        ed = st + item[0]
        ps_dis = F.cross_entropy(ps_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none')
        ns_dis = F.cross_entropy(ns_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none')

        triplet_dis.append(max((ps_dis-ns_dis).mean()+1,0))
        st = ed
    triplet_dis = sum(triplet_dis)/len(triplet_dis)


    return total_loss, triplet_con, dis_loss, triplet_dis


def make_negative_condition_esm(x_0_con, x_0_dis, base_model):
    device = x_0_dis.device

    with torch.no_grad():
        _, _, embeddings, _ = base_model(
            input_ids=x_0_con['input_ids'],
            attention_masks=x_0_con['attention_mask']
        )

    B = embeddings.shape[0]

    indices = torch.randperm(B)

    ns_con = embeddings[indices].clone()
    ns_dis = x_0_dis[indices].clone()

    return ns_con, ns_dis