import argparse,math,time,warnings,copy, numpy as np, os.path as path 
import utils.evals as evals
import utils.utils as utils
from utils.data_loader import process_data
import torch, torch.nn as nn, torch.nn.functional as F
import lamp.Constants as Constants
from lamp.Models import LAMP
from lamp.Translator import translate
from config_args import config_args,get_args
from pdb import set_trace as stop
from tqdm import tqdm
from lamp.SubLayers import PLoss


def train_epoch(model,train_data, crit, optimizer,adv_optimizer,epoch,data_dict,opt,logger):
    ##added attention_decay
    if opt.ploss:
        ploss=PLoss(opt.ploss).apply
    model.train()
    out_len = (opt.tgt_vocab_size) if opt.binary_relevance else (opt.tgt_vocab_size-1)
    all_predictions = torch.zeros(len(train_data._src_insts),out_len)
    all_targets = torch.zeros(len(train_data._src_insts),out_len)

    batch_idx,batch_size = 0,train_data._batch_size
    bce_total,d_total,d_fake_total,g_total = 0,0,0,0

    
    for batch in tqdm(train_data, mininterval=0.5,desc='(Training)', leave=False):
            src,adj,tgt = batch
            loss,d_loss = 0,0
            gold = tgt[:, 1:]
    
            if opt.binary_relevance:
                gold_binary = utils.get_gold_binary(gold.data.cpu(),opt.tgt_vocab_size).cuda()
                optimizer.zero_grad()
                pred,enc_output,_,dec_output2= model(src,adj,None,gold_binary,return_attns=True,int_preds=opt.int_preds)
                dec_self_attn,dec_enc_attn=dec_output2
                                
                logger.push_attentions(dec_self_attn[1])
                
                norm_pred = F.sigmoid(pred)
                bce_loss =  F.binary_cross_entropy_with_logits(pred, gold_binary,reduction='mean')
                loss+= bce_loss
                bce_total += bce_loss.item()
                if opt.int_preds and not opt.matching_mlp:
                    for i in range(len(results[0])):
                        bce_loss =  F.binary_cross_entropy_with_logits(results[0][i], gold_binary,reduction='mean')
                        loss += (opt.int_pred_weight)*bce_loss
                
                if opt.ploss:
                    loss+=opt.p_coef*ploss(dec_self_attn[1].sum(axis=0).view(-1))
                    
                if epoch == opt.thresh1:
                    opt.init_model = copy.deepcopy(model)
                loss.backward()
                optimizer.step()
                tgt_out = gold_binary.data
                pred_out = norm_pred.data
    
            else: 
                # Non Binary Outputs
                optimizer.zero_grad()
                pred,enc_output,_,dec_output2 = model(src,adj,tgt,None,int_preds=opt.int_preds,return_attns=True)                                               
                dec_self_attn,dec_enc_attn=dec_output2
                                
                logger.push_attentions(dec_self_attn)
                
                loss ,_= crit(F.log_softmax(pred), gold.contiguous().view(-1)),logger.push_loss(loss)
                
                if opt.ploss:
                    loss+=opt.p_coef*PLoss(opt.p,dec_self_attn.sum(axis=0).view(-1))
                pred = F.softmax(pred,dim=1)
                pred_vals,pred_idxs = pred.max(1)
                pred_vals = pred_vals.view(gold.size()).data.cpu()
                pred_idxs = pred_idxs.view(gold.size()).data.cpu()
                pred_out = torch.zeros(pred_vals.size(0),pred.size(1)).scatter_(1,pred_idxs.long(),pred_vals)
                tgt_out = torch.zeros(pred_vals.size(0),pred.size(1)).scatter_(1,gold.data.cpu().long(),torch.ones(pred_vals.size()))
                pred_out = pred_out[:,1:]
                tgt_out = tgt_out[:,1:]
                loss.backward()
                optimizer.step()
    
            
            ## Updates ##
            start_idx, end_idx = (batch_idx*batch_size),((batch_idx+1)*batch_size)
            all_predictions[start_idx:end_idx] = pred_out
            all_targets[start_idx:end_idx] = tgt_out
            batch_idx +=1
            
        
    return all_predictions, all_targets, bce_total
            
            
            
import argparse,math,time,warnings,copy, numpy as np, os.path as path 
import utils.evals as evals
import utils.utils as utils
from utils.data_loader import process_data
import torch, torch.nn as nn, torch.nn.functional as F
import lamp.Constants as Constants
from lamp.Models import LAMP
from lamp.Translator import translate
from config_args import config_args,get_args
from pdb import set_trace as stop
from tqdm import tqdm
from lamp.SubLayers import PLoss


def train_epoch(model,train_data, crit, optimizer,adv_optimizer,epoch,data_dict,opt,logger):
    ##added attention_decay
    if opt.ploss:
        ploss=PLoss(opt.ploss).apply
    model.train()
    out_len = (opt.tgt_vocab_size) if opt.binary_relevance else (opt.tgt_vocab_size-1)
    all_predictions = torch.zeros(len(train_data._src_insts),out_len)
    all_targets = torch.zeros(len(train_data._src_insts),out_len)

    batch_idx,batch_size = 0,train_data._batch_size
    bce_total,d_total,d_fake_total,g_total = 0,0,0,0

    
    for batch in tqdm(train_data, mininterval=0.5,desc='(Training)', leave=False):
            src,adj,tgt = batch
            loss,d_loss = 0,0
            gold = tgt[:, 1:]
    
            if opt.binary_relevance:
                gold_binary = utils.get_gold_binary(gold.data.cpu(),opt.tgt_vocab_size).cuda()
                optimizer.zero_grad()
                pred,enc_output,_,dec_output2= model(src,adj,None,gold_binary,return_attns=True,int_preds=opt.int_preds)
                dec_self_attn,dec_enc_attn=dec_output2
                                
                logger.push_attentions(dec_self_attn[1])
                
                norm_pred = F.sigmoid(pred)
                bce_loss =  F.binary_cross_entropy_with_logits(pred, gold_binary,reduction='mean')
                loss+= bce_loss
                bce_total += bce_loss.item()
                if opt.int_preds and not opt.matching_mlp:
                    for i in range(len(results[0])):
                        bce_loss =  F.binary_cross_entropy_with_logits(results[0][i], gold_binary,reduction='mean')
                        loss += (opt.int_pred_weight)*bce_loss
                
                if opt.ploss:
                    loss+=opt.p_coef*ploss(dec_self_attn[1].sum(axis=0).view(-1))
                    
                if epoch == opt.thresh1:
                    opt.init_model = copy.deepcopy(model)
                loss.backward()
                optimizer.step()
                tgt_out = gold_binary.data
                pred_out = norm_pred.data
    
            else: 
                # Non Binary Outputs
                optimizer.zero_grad()
                pred,enc_output,_,dec_output2 = model(src,adj,tgt,None,int_preds=opt.int_preds,return_attns=True)                                               
                dec_self_attn,dec_enc_attn=dec_output2
                                
                logger.push_attentions(dec_self_attn)
                
                loss ,_= crit(F.log_softmax(pred), gold.contiguous().view(-1)),logger.push_loss(loss)
                
                if opt.ploss:
                    loss+=opt.p_coef*PLoss(opt.p,dec_self_attn.sum(axis=0).view(-1))
                pred = F.softmax(pred,dim=1)
                pred_vals,pred_idxs = pred.max(1)
                pred_vals = pred_vals.view(gold.size()).data.cpu()
                pred_idxs = pred_idxs.view(gold.size()).data.cpu()
                pred_out = torch.zeros(pred_vals.size(0),pred.size(1)).scatter_(1,pred_idxs.long(),pred_vals)
                tgt_out = torch.zeros(pred_vals.size(0),pred.size(1)).scatter_(1,gold.data.cpu().long(),torch.ones(pred_vals.size()))
                pred_out = pred_out[:,1:]
                tgt_out = tgt_out[:,1:]
                loss.backward()
                optimizer.step()
    
            
            ## Updates ##
            start_idx, end_idx = (batch_idx*batch_size),((batch_idx+1)*batch_size)
            all_predictions[start_idx:end_idx] = pred_out
            all_targets[start_idx:end_idx] = tgt_out
            batch_idx +=1
            
        
    return all_predictions, all_targets, bce_total
            
            
            
##added attention_decay # Non Binary Outputs ## Updates ##