#!/usr/bin/env python
# coding: utf-8
"""
True regret loss
"""

import numpy as np
import torch

from pyepo import EPO
import torch
from torch import nn
from torch.utils.data import DataLoader
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from utils.tools import adjust_learning_rate, visual
from time import *

def regret(predmodel, optmodel, dataloader):
    """
    A function to evaluate model performance with normalized true regret

    Args:
        predmodel (nn): a regression neural network for cost prediction
        optmodel (optModel): an PyEPO optimization model
        dataloader (DataLoader): Torch dataloader from optDataSet

    Returns:
        float: true regret loss
    """
    # evaluate
    predmodel.eval()
    loss = 0
    optsum = 0
    # load data
    for data in dataloader:
        x, c, w, z, b = data
        # cuda
        if next(predmodel.parameters()).is_cuda:
            x, c, w, z, b = x.cuda(), c.cuda(), w.cuda(), z.cuda(), b.cuda()
        # predict
        with torch.no_grad(): # no grad
            cp = predmodel(x).to("cpu").detach().numpy()
        # solve
        for j in range(cp.shape[0]):
            # accumulate loss
            loss += calRegret(optmodel, cp[j], c[j].to("cpu").detach().numpy(),
                              z[j].item())
        optsum += abs(z).sum().item()
    # turn back train mode
    predmodel.train()
    # normalized
    return loss / (optsum + 1e-7)

def calRegret_PDG(optmodel, optmodel_d, pred_cost, true_cost, true_obj, para_constr):
    """
    A function to calculate normalized true regret for a batch

    Args:
        optmodel (optModel): optimization model
        pred_cost (torch.tensor): predicted costs
        true_cost (torch.tensor): true costs
        true_obj (torch.tensor): true optimal objective values

    Returns:predmodel
        float: true regret losses
    """
    # opt sol for pred cost
    optmodel.setObj(pred_cost)
    optmodel.updateConstr(para_constr)
    sol, _ = optmodel.solve()
    # obj with true cost
    obj = np.dot(sol, true_cost)
    # loss
    if optmodel.modelSense == EPO.MINIMIZE:
        loss = obj - true_obj
    if optmodel.modelSense == EPO.MAXIMIZE:
        loss = true_obj - obj
    return loss


def calRegret(optmodel, pred_cost, true_cost, true_obj):
    """
    A function to calculate normalized true regret for a batch

    Args:
        optmodel (optModel): optimization model
        pred_cost (torch.tensor): predicted costs
        true_cost (torch.tensor): true costs
        true_obj (torch.tensor): true optimal objective values

    Returns:predmodel
        float: true regret losses
    """
    # pred_cost = pred_cost.detach().to("cpu").numpy()
    # true_cost = true_cost.detach().to("cpu").numpy()
    # opt sol for pred cost
    optmodel.setObj(pred_cost)
    # optmodel.updateConstr(para_constr)
    sol, _ = optmodel.solve()  # sol = sol_pred
    # print('sol shape', np.array(sol).shape)
    # print('cost shape', true_cost.shape, pred_cost.shape)
    # obj with true cost
    obj = np.dot(sol, true_cost)
    obj_list = sol * true_cost
    # loss
    if optmodel.modelSense == EPO.MINIMIZE:
        loss = obj - true_obj
    if optmodel.modelSense == EPO.MAXIMIZE:
        loss = true_obj - obj
    return abs(loss), sol, obj_list
