import numpy as np
from keras.callbacks import Callback
import os

class WeightedAccuracy(Callback):
  def __init__(self, tag, input, output, sample_weights, batch_size=32, do_gpu_status=False):
    self.tag = tag
    self.input = input
    self.output = output
    self.sample_weights = sample_weights
    self.w_acc = []
    self.batch_size = batch_size
    self.do_gpu_status = do_gpu_status
  
  def report_gpu_usage(self, msg):
    print(msg)
    if self.do_gpu_status:  
      os.system("nvidia-smi")     
      
  def on_train_begin(self, logs={}):
    self.losses = []
    self.w_acc = []
    self.report_gpu_usage('Train begin')

  def on_batch_end(self, batch, logs={}):
    self.losses.append(logs.get('loss'))
    
    
  def on_epoch_end(self, epoch, logs=None):
    self.report_gpu_usage('Epoch complete')
      
  def on_train_end(self, logs=None):
      self.calc_accuracy()
      self.print_info()
        
  def calc_accuracy(self):
    p=self.model.predict(self.input, batch_size=self.batch_size)
    cnt=np.sum(self.sample_weights)    
    err=np.sum(np.not_equal(np.argmax(p,axis=-1)*self.sample_weights,np.argmax(self.output,axis=-1)*self.sample_weights))
    self.w_acc=1.0-(err/cnt)
  
  def print_info(self):
    print(self.tag + ' weighted accuracy :{:5.2f}%'.format(self.w_acc*100))
    
#   def calc_confusion_matrix
#     classes = np.argmax(self.output,axis=-1)
#     print('Predicted')
#     for i in range(self.n_labels)
#       print(self.label[i] + '\t')
#     for i in range(self.n_labels)
#       for j in range(self.n_labels)
#         cnt=np.sum(self.sample_weights*np.where(self.output[j]==1))
#         cnt=np.sum(np.equal(self.output, axis=-1), )
#             
#         err=np.sum(np.not_equal(np.argmax(p,axis=-1)*self.sample_weights,np.argmax(self.output,axis=-1)*self.sample_weights))
#         self.w_acc=1.0-(err/cnt)
        
      
        