from __future__ import division, print_function
import logging
import numpy as np
import scipy as sc
import skimage
from skimage import transform
import theano
import theano.tensor as T
import lasagne
import sys
import cPickle as pickle
import gc
from datetime import datetime
from collections import OrderedDict
from mscoco_threaded_iter import COCOCaptionDataset
sys.path.append('/home/noury/codevault/Recipes/modelzoo/')
sys.path.append('/home/noury/codevault/seq2seq-lasagne/')
from resnet50 import build_model

if __name__ == '__main__':
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s %(message)s', '%m/%d/%Y %I:%M:%S %p')
    fh = logging.FileHandler('mscoco_captions.log')
    fh.setLevel(logging.DEBUG)
    fh.setFormatter(formatter)

    sh = logging.StreamHandler(sys.stdout)
    sh.setLevel(logging.INFO)
    sh.setFormatter(formatter)

    logger.addHandler(fh)
    logger.addHandler(sh)

    logger.info('Loading the ResNet50 model.')
    # First get the ResNet50 model
    resnet_weights_path = '/home/noury/modelzoo/resnet50.pkl'
    resnet = build_model()
    model_params = pickle.load(open(resnet_weights_path, 'rb'))
    lasagne.layers.set_all_param_values(resnet['prob'], model_params['values'])
    mean_im = model_params['mean_image'].reshape((1, 3, 224, 224)).astype(np.float32)

    # Load the files needed for the MS COCO Captions dataset
    train_images_path = '/home/noury/datasets/mscoco/train2014'
    valid_images_path = '/home/noury/datasets/mscoco/val2014'
    train_annotations_filepath = '/home/noury/datasets/mscoco/annotations/captions_train2014.json'
    valid_annotations_filepath = '/home/noury/datasets/mscoco/annotations/captions_val2014.json'
    coco_captions = pickle.load(open('coco_captions_trainval2014.pkl', 'rb'))
    train_buckets = coco_captions['train buckets']
    valid_buckets = coco_captions['valid buckets']
    wordset = coco_captions['raw wordset']
    word2idx = {}
    word2idx['<PAD>'] = 0
    word2idx['<GO>'] = 1
    word2idx['<EOS>'] = 2
    for i, w in enumerate(wordset):
        word2idx[w] = i+3
    idx2word = map(lambda x: x[0], sorted(word2idx.items(), key=lambda x: x[1]))

    bucket_minibatch_sizes = {16:64, 32:32, 64:16}

    logger.info('Creating global variables')
    CONTINUE = True
    HIDDEN_SIZE = 2048
    EMBEDDING_SIZE = 256
    WORD_SIZE = len(idx2word)
    DENSE_SIZE = 1024
    ORDER_VIOLATION_COEFF = 0.2
    RNN_GRAD_CLIP = 32
    TOTAL_GRAD_CLIP = 64
    TOTAL_MAX_NORM = 128
    RESNET_ADAM_LR = theano.shared(np.float32(0.005 * 2), 'resnet adam lr')
    RECURR_ADAM_LR = theano.shared(np.float32(0.05 * 2), 'recurrent adam lr')
    RESNET_SGDM_LR = theano.shared(np.float32(0.0005 * 2), 'resnet sgdm lr')
    RECURR_SGDM_LR = theano.shared(np.float32(0.005 * 2), 'recurrent sgdm lr')
    EPOCH_LR_COEFF = np.float32(0.5)
    NUM_EPOCHS = 20
    ADAM_EPOCHS = 5

    if CONTINUE:
        import glob
        param_values = glob.glob('param_values_*.pkl')
        max_epoch = max(map(lambda x: int(x[len('param_values_'):-len('.pkl')]), param_values))
        logger.info('Continue training from epoch {}'.format(max_epoch + 1))
        logger.info('Setting previous parameter values from epoch {}'.format(max_epoch))
        if max_epoch >= ADAM_EPOCHS:
            for _ in xrange(max_epoch - ADAM_EPOCHS):
                RESNET_SGDM_LR.set_value(RESNET_SGDM_LR.get_value() * EPOCH_LR_COEFF)
                RECURR_SGDM_LR.set_value(RECURR_SGDM_LR.get_value() * EPOCH_LR_COEFF)
            ADAM_EPOCHS = 0
        else:
            for _ in xrange(max_epoch):
                RESNET_ADAM_LR.set_value(RESNET_ADAM_LR.get_value() * EPOCH_LR_COEFF)
                RECURR_ADAM_LR.set_value(RECURR_ADAM_LR.get_value() * EPOCH_LR_COEFF)
        param_values_file = 'param_values_{}.pkl'.format(max_epoch)

    logger.info('Building the network.')
    im_features = lasagne.layers.get_output(resnet['pool5'])
    im_features = T.flatten(im_features, outdim=2) # batch size, number of features
    cap_out_var = T.imatrix('cap_out')  # batch size, seq len
    cap_in_var = T.imatrix('cap_in')    # batch size, seq len
    mask_var = T.bmatrix('mask_var')    # batch size, seq len
    gate = lasagne.layers.Gate(W_in=lasagne.init.Orthogonal(), W_hid=lasagne.init.Orthogonal(),
                               W_cell=lasagne.init.Normal(), b=lasagne.init.Constant(0.0))
    cell_gate = lasagne.layers.Gate(W_in=lasagne.init.Orthogonal(), W_hid=lasagne.init.Orthogonal(),
                                    W_cell=None, b=lasagne.init.Constant(0.0),
                                    nonlinearity=lasagne.nonlinearities.tanh)
    forget_gate = lasagne.layers.Gate(W_in=lasagne.init.Orthogonal(), W_hid=lasagne.init.Orthogonal(),
                                      W_cell=lasagne.init.Normal(), b=lasagne.init.Constant(5.0))
    l_in = lasagne.layers.InputLayer((None, None), cap_in_var, name="l_in")
    l_mask = lasagne.layers.InputLayer((None, None), mask_var, name="l_mask")
    l_hid = lasagne.layers.InputLayer((None, HIDDEN_SIZE), input_var=im_features, name="l_hid")
    l_emb = lasagne.layers.EmbeddingLayer(l_in, input_size=WORD_SIZE, output_size=EMBEDDING_SIZE, name="l_emb")
    l_lstm = lasagne.layers.LSTMLayer(l_emb, HIDDEN_SIZE, ingate=gate, forgetgate=forget_gate, cell=cell_gate,
                                    outgate=gate, hid_init=l_hid, peepholes=True, grad_clipping=RNN_GRAD_CLIP,
                                    mask_input=l_mask, precompute_input=False, name="l_lstm") # batch size, seq len, hidden size
    l_reshape = lasagne.layers.ReshapeLayer(l_lstm, (-1, [2]), name="l_reshape") # batch size * seq len, hidden size
    l_fc = lasagne.layers.DenseLayer(l_reshape, DENSE_SIZE, b=lasagne.init.Constant(5.0),
                                    nonlinearity=lasagne.nonlinearities.rectify, name="l_fc")
    l_drp = lasagne.layers.DropoutLayer(l_fc, 0.2, name="l_drp")
    l_out = lasagne.layers.DenseLayer(l_drp, WORD_SIZE, nonlinearity=lasagne.nonlinearities.softmax,
                                       name='l_out') # batch size * seq len, WORD SIZE
    l_slice = lasagne.layers.SliceLayer(l_lstm, -1, axis=1, name="l_slice")

    if CONTINUE:
        logger.info('Setting model weights from epoch {}'.format(max_epoch))
        param_values = pickle.load(open(param_values_file, 'rb'))
        lasagne.layers.set_all_param_values(l_out, param_values['recurrent'])
        lasagne.layers.set_all_param_values(resnet['pool5'], param_values['resnet'])

    logger.info('Creating output and loss variables')
    prediction = lasagne.layers.get_output(l_out, deterministic=False)
    flat_cap_out_var = T.flatten(cap_out_var, outdim=1)
    loss = T.sum(lasagne.objectives.categorical_crossentropy(prediction, flat_cap_out_var))
    caption_features = lasagne.layers.get_output(l_slice, deterministic=False)
    order_embedding_loss = T.pow(T.maximum(0, caption_features - im_features), 2).sum()
    total_loss = loss + ORDER_VIOLATION_COEFF * order_embedding_loss

    deterministic_prediction = lasagne.layers.get_output(l_out, deterministic=True)
    deterministic_captions = lasagne.layers.get_output(l_slice, deterministic=True)
    deterministic_loss = T.sum(lasagne.objectives.categorical_crossentropy(deterministic_prediction, flat_cap_out_var))
    deterministic_order_embedding_loss = T.pow(T.maximum(0, deterministic_captions - im_features), 2).sum()
    deterministic_total_loss = deterministic_loss + ORDER_VIOLATION_COEFF * deterministic_order_embedding_loss

    logger.info('Getting all parameters and creating update rules.')
    resnet_params = lasagne.layers.get_all_params(resnet['pool5'], trainable=True)
    recurrent_params = lasagne.layers.get_all_params(l_out, trainable=True)
    resnet_grads = T.grad(total_loss, resnet_params)
    recurrent_grads = T.grad(total_loss, recurrent_params)
    resnet_grads = [T.clip(g, -TOTAL_GRAD_CLIP, TOTAL_GRAD_CLIP) for g in resnet_grads]
    recurrent_grads = [T.clip(g, -TOTAL_GRAD_CLIP, TOTAL_GRAD_CLIP) for g in recurrent_grads]
    resnet_grads, resnet_norm = lasagne.updates.total_norm_constraint(resnet_grads, TOTAL_MAX_NORM, return_norm=True)
    recurrent_grads, recurrent_norm = lasagne.updates.total_norm_constraint(recurrent_grads, TOTAL_MAX_NORM, return_norm=True)
    resnet_adam = lasagne.updates.adam(resnet_grads, resnet_params, learning_rate=RESNET_ADAM_LR)
    recurrent_adam = lasagne.updates.adam(recurrent_grads, recurrent_params, learning_rate=RECURR_ADAM_LR)
    resnet_adam_items = resnet_adam.items()
    recurrent_adam_items = recurrent_adam.items()
    resnet_adam_items.extend(recurrent_adam_items)
    adam_updates = OrderedDict(resnet_adam_items)

    logger.info("Creating the Adam update Theano function")
    adam_train_fun = theano.function([resnet['input'].input_var, cap_in_var, mask_var, cap_out_var],
                                     [total_loss, order_embedding_loss, resnet_norm, recurrent_norm],
                                     updates=adam_updates)

    resnet_sgdm = lasagne.updates.nesterov_momentum(resnet_grads, resnet_params, learning_rate=RESNET_SGDM_LR, momentum=0.8)
    recurrent_sgdm = lasagne.updates.nesterov_momentum(recurrent_grads, recurrent_params, learning_rate=RECURR_SGDM_LR, momentum=0.9)
    resnet_sgdm_items = resnet_sgdm.items()
    recurrent_sgdm_items = recurrent_sgdm.items()
    resnet_sgdm_items.extend(recurrent_sgdm_items)
    sgdm_updates = OrderedDict(resnet_sgdm_items)

    logger.info("Creating the SGDM update Theano function")
    sgdm_train_fun = theano.function([resnet['input'].input_var, cap_in_var, mask_var, cap_out_var],
                                     [total_loss, order_embedding_loss, resnet_norm, recurrent_norm],
                                     updates=sgdm_updates)
    logger.info("Creating the evaluation Theano function")
    eval_fun = theano.function([resnet['input'].input_var, cap_in_var, mask_var, cap_out_var],
                               [deterministic_total_loss, deterministic_order_embedding_loss])

    logger.info('Loading the COCO Captions training and validation sets.')
    coco_train = COCOCaptionDataset(train_images_path, train_annotations_filepath, train_buckets,
                                    bucket_minibatch_sizes, word2idx, mean_im, True)
    coco_valid = COCOCaptionDataset(valid_images_path, valid_annotations_filepath, valid_buckets,
                                    bucket_minibatch_sizes, word2idx, mean_im, False)

    deleted_adam_resources = False
    total_loss_values = {}
    order_embedding_loss_values = {}
    resnet_norm_values = {}
    recurrent_norm_values = {}
    det_total_loss_values = {}
    det_order_embedding_loss_values = {}

    logger.info("Starting the training process...")
    START = 1
    if CONTINUE:
        START = max_epoch + 1
    for e in xrange(START, NUM_EPOCHS + 1):
        logger.info("Starting epoch".format(e))

        total_loss_values[e] = []
        order_embedding_loss_values[e] = []
        resnet_norm_values[e] = []
        recurrent_norm_values[e] = []
        det_total_loss_values[e] = []
        det_order_embedding_loss_values[e] = []

        if e <= ADAM_EPOCHS:
            train_fun = adam_train_fun
            RESNET_ADAM_LR.set_value(np.float32(RESNET_ADAM_LR.get_value() * EPOCH_LR_COEFF))
            RECURR_ADAM_LR.set_value(np.float32(RECURR_ADAM_LR.get_value() * EPOCH_LR_COEFF))
            logger.info("Training with the Adam update function")
            logger.info("ResNet LR: {}, Recurrent LR: {}".format(RESNET_ADAM_LR.get_value(), RECURR_ADAM_LR.get_value()))
        else:
            if not deleted_adam_resources:
                deleted_adam_resources = True
                logger.info("Deleting Adam training function resources")
                del resnet_adam_items
                del adam_updates
                del adam_train_fun
                logger.info("Calling gc.collect() to remove resource from GC")
                while gc.collect() > 0:
                    pass
            train_fun = sgdm_train_fun
            RESNET_SGDM_LR.set_value(np.float32(RESNET_SGDM_LR.get_value() * EPOCH_LR_COEFF))
            RECURR_SGDM_LR.set_value(np.float32(RECURR_SGDM_LR.get_value() * EPOCH_LR_COEFF))
            logger.info("Training with the SGDM update function")
            logger.info("ResNet LR: {}, Recurrent LR: {}".format(RESNET_SGDM_LR.get_value(), RECURR_SGDM_LR.get_value()))
        mb = 0
        now = datetime.now()
        for im, cap_in, cap_out in coco_train:
            tl, oe, resn, recn = train_fun(im, cap_in, (cap_in > 0).astype(np.int8), cap_out)
            logger.debug("Epoch: {}, Minibatch: {}, Total Loss: {}, Order-embedding loss: {}, ResNet norm: {}, Recurrent norm: {}".format(e, mb, tl, oe, resn, recn))
            total_loss_values[e].append(tl)
            order_embedding_loss_values[e].append(oe)
            resnet_norm_values[e].append(resn)
            recurrent_norm_values[e].append(recn)
            mb += 1
        logger.info("Training epoch {} took {}.".format(e, datetime.now() - now))
        logger.info("Epoch {} results".format(e))
        logger.info("Mean total loss: {}".format(np.mean(total_loss_values[e])))
        logger.info("Mean order embedding loss: {}".format(np.mean(order_embedding_loss_values[e])))
        logger.info("Mean ResNet norm: {}".format(np.mean(resnet_norm_values[e])))
        logger.info("Mean Recurrent norm: {}".format(np.mean(recurrent_norm_values[e])))

        logger.info("Saving model parameters for epoch {}".format(e))
        pickle.dump({'resnet':lasagne.layers.get_all_param_values(resnet['pool5']),
                     'recurrent':lasagne.layers.get_all_param_values(l_out)},
                      open('param_values_{}.pkl'.format(e), 'wb'), protocol=-1)
        logger.info("Saving loss values for epoch {}".format(e))
        pickle.dump({'total loss':total_loss_values, 'oe loss': order_embedding_loss_values,
                     'resnet norm': resnet_norm_values, 'recurrent norm': recurrent_norm_values},
                     open('training_losses.pkl', 'wb'), protocol=-1)

        if e % 2 == 0:
            mb = 0
            logger.info("Evaluating the model on the validation set.")
            for im, cap_in, cap_out in coco_valid:
                tl, oe = eval_fun(im, cap_in, (cap_in > 0).astype(np.int8), cap_out)
                logger.debug("Epoch: {}, Minibatch: {}, Validation total loss: {}, Validation order-embedding loss: {}".format(e, mb, tl, oe))
                det_total_loss_values[e].append(tl)
                det_order_embedding_loss_values[e].append(oe)
                mb += 1
            logger.info("Epoch {} validation results".format(e))
            logger.info("Mean validation total loss: {}".format(np.mean(det_total_loss_values[e])))
            logger.info("Mean validation order-embedding loss: {}".format(np.mean(det_order_embedding_loss_values[e])))

            logger.info("Saving validation loss values for epoch {}".format(e))
            pickle.dump({'total loss': det_total_loss_values, 'oe loss': det_order_embedding_loss_values},
                        open('validation_losses.pkl', 'wb'), protocol=-1)


from pdb import set_trace as T
import abc

import numpy as np
from nmmo.systems import experience, combat, ai

from nmmo.lib import material

### Infrastructure ###
class SkillGroup:
   def __init__(self, realm):
      self.expCalc = experience.ExperienceCalculator()
      self.config  = realm.dataframe.config
      self.skills  = set()

   def update(self, realm, entity, actions):
      for skill in self.skills:
         skill.update(realm, entity)

   def packet(self):
      data = {}
      for skill in self.skills:
         data[skill.__class__.__name__.lower()] = skill.packet()

      return data

class Skill:
   skillItems = abc.ABCMeta

   def __init__(self, skillGroup):
      self.config  = skillGroup.config
      self.expCalc = skillGroup.expCalc
      self.exp     = 0

      skillGroup.skills.add(self)

   def packet(self):
      data = {}

      data['exp']   = self.exp
      data['level'] = self.level

      return data

   def update(self, realm, entity):
      pass

   def setExpByLevel(self, level):
      self.exp = self.expCalc.expAtLevel(level)

   @property
   def level(self):
      lvl = self.expCalc.levelAtExp(self.exp)
      assert lvl == int(lvl)
      return int(lvl)

### Skill Subsets ###
class Harvesting(SkillGroup):
   def __init__(self, realm):
      super().__init__(realm)

      self.fishing      = Fishing(self)
      self.hunting      = Hunting(self)

class Combat(SkillGroup):
   def __init__(self, realm):
      super().__init__(realm)

      self.constitution = Constitution(self)
      self.defense      = Defense(self)
      self.melee        = Melee(self)
      self.range        = Range(self)
      self.mage         = Mage(self)

   def packet(self):
      data          = super().packet()
      data['level'] = combat.level(self)

      return data

   def applyDamage(self, dmg, style):
      if not self.config.game_system_enabled('Progression'):
         return

      config    = self.config
      baseScale = config.PROGRESSION_BASE_XP_SCALE
      combScale = config.PROGRESSION_COMBAT_XP_SCALE
      conScale  = config.PROGRESSION_CONSTITUTION_XP_SCALE

      self.constitution.exp += dmg * baseScale * conScale

      skill      = self.__dict__[style]
      skill.exp += dmg * baseScale * combScale

   def receiveDamage(self, dmg):
      if not self.config.game_system_enabled('Progression'):
         return

      config    = self.config
      baseScale = config.PROGRESSION_BASE_XP_SCALE
      combScale = config.PROGRESSION_COMBAT_XP_SCALE
      conScale  = config.PROGRESSION_CONSTITUTION_XP_SCALE

      self.constitution.exp += dmg * baseScale * conScale
      self.defense.exp      += dmg * baseScale * combScale

class Skills(Harvesting, Combat):
   pass

### Individual Skills ###
class CombatSkill(Skill): pass

class Constitution(CombatSkill):
   def __init__(self, skillGroup):
      super().__init__(skillGroup)
      self.setExpByLevel(self.config.BASE_HEALTH)

   def update(self, realm, entity):
      health = entity.resources.health
      food   = entity.resources.food
      water  = entity.resources.water
      config = self.config

      if not config.game_system_enabled('Resource'):
         health.increment(1)
         return

      # Heal if above fractional resource threshold
      regen       = config.RESOURCE_HEALTH_REGEN_THRESHOLD
      foodThresh  = food > regen * entity.skills.hunting.level
      waterThresh = water > regen * entity.skills.fishing.level

      if foodThresh and waterThresh:
         restore = config.RESOURCE_HEALTH_RESTORE_FRACTION
         restore = np.floor(restore * self.level)
         health.increment(restore)

      if food.empty:
         health.decrement(1)

      if water.empty:
         health.decrement(1)

class Melee(CombatSkill): pass
class Range(CombatSkill): pass
class Mage(CombatSkill): pass
class Defense(CombatSkill): pass

class Fishing(Skill):
   def __init__(self, skillGroup):
      super().__init__(skillGroup)
      config, level = self.config, 1
      if config.game_system_enabled('Progression'):
         level = config.PROGRESSION_BASE_RESOURCE
      elif config.game_system_enabled('Resource'):
         level = config.RESOURCE_BASE_RESOURCE

      self.setExpByLevel(level)

   def update(self, realm, entity):
      if not self.config.game_system_enabled('Resource'):
         return

      water = entity.resources.water
      water.decrement(1)

      if material.Water not in ai.utils.adjacentMats(
            realm.map.tiles, entity.pos):
         return

      restore = self.config.RESOURCE_HARVEST_RESTORE_FRACTION
      restore = np.floor(restore * self.level)
      water.increment(restore)

      if self.config.game_system_enabled('Progression'):
         self.exp += self.config.PROGRESSION_BASE_XP_SCALE * restore

class Hunting(Skill):
   def __init__(self, skillGroup):
      super().__init__(skillGroup)
      config, level = self.config, 1
      if config.game_system_enabled('Progression'):
         level = config.PROGRESSION_BASE_RESOURCE
      elif config.game_system_enabled('Resource'):
         level = config.RESOURCE_BASE_RESOURCE

      self.setExpByLevel(level)

   def update(self, realm, entity):
      if not self.config.game_system_enabled('Resource'):
         return

      food = entity.resources.food
      food.decrement(1)

      r, c = entity.pos
      if (type(realm.map.tiles[r, c].mat) not in [material.Forest] or
            not realm.map.harvest(r, c)):
         return

      restore = self.config.RESOURCE_HARVEST_RESTORE_FRACTION
      restore = np.floor(restore * self.level)
      food.increment(restore)

      if self.config.game_system_enabled('Progression'):
         self.exp += self.config.PROGRESSION_BASE_XP_SCALE * restore


from math import ceil
from os import path

import numpy as np

from decibel.import_export.chord_annotation_io import import_chord_annotation
from decibel.music_objects.chord import Chord
from decibel.music_objects.chord_alphabet import ChordAlphabet
from decibel.music_objects.chord_vocabulary import ChordVocabulary
from decibel.music_objects.song import Song
from decibel.import_export import filehandler
from decibel.import_export.midi_alignment_score_io import read_chord_alignment_score

def get_well_aligned_midis(song: Song) -> [str]:
    """
    Return names of only the well-aligned MIDIs for this Song (excluding duplicates)

    :param song: Song in our data set
    """
    # Find duplicate MIDIs in this song; we will exclude them
    duplicate_midis = filehandler.find_duplicate_midis(song)

    well_aligned_midis = []
    for full_midi_path in song.full_midi_paths:
        midi_name = filehandler.get_file_name_from_full_path(full_midi_path)
        if midi_name not in duplicate_midis:
            alignment_score = read_chord_alignment_score(midi_name)
            if alignment_score.is_well_aligned:
                well_aligned_midis.append(midi_name)

    return well_aligned_midis

def get_expected_best_midi(song: Song) -> (str, str):
    """
    Find name of the expected best well-aligned MIDI and segmentation type for this Song
    (based on MIDI chord probability)

    :param song: Song in our data set
    """
    # We only consider well-aligned MIDIs
    well_aligned_midis = get_well_aligned_midis(song)

    # Return path to the best MIDI of the song
    best_midi_name, best_midi_quality, best_segmentation = '', -9999999999, ''
    for segmentation_type in 'bar', 'beat':
        for full_midi_path in well_aligned_midis:
            midi_name = filehandler.get_file_name_from_full_path(full_midi_path)
            midi_chord_probability = filehandler.read_midi_chord_probability(segmentation_type, midi_name)
            if midi_chord_probability > best_midi_quality:
                # This is the best MIDI & segmentation type we have seen until now
                best_midi_name, best_midi_quality, best_segmentation = \
                    midi_name, midi_chord_probability, segmentation_type

    return best_midi_name, best_segmentation

def get_expected_best_tab_lab(song: Song) -> str:
    """
    Find the lab file of the expected best tab for this Song (based on log-likelihood returned by Jump Alignment)

    :param song: Song in our data set
    """
    best_tab_lab, best_tab_quality = '', 0

    for tab_path in song.full_tab_paths:
        tab_write_path = filehandler.get_full_tab_chord_labs_path(tab_path)
        if filehandler.file_exists(tab_write_path):
            tab_quality, _ = filehandler.read_log_likelihood(song.key, tab_path)
            if tab_quality > best_tab_quality:
                best_tab_lab, best_tab_quality = tab_write_path, tab_quality

    return best_tab_lab

def get_actual_best_midi_lab(song: Song) -> str:
    """
    Find path to the actual best MIDI and segmentation type for this Song (based on CSR). Only use after evaluation.

    :param song: Song in our data set
    """
    best_beat_midi_name, best_beat_midi_csr = filehandler.get_actual_best_midi_for_song('beat', song.key)
    best_bar_midi_name, best_bar_midi_csr = filehandler.get_actual_best_midi_for_song('bar', song.key)
    if best_beat_midi_csr > best_bar_midi_csr:
        return filehandler.get_full_midi_chord_labs_path(best_beat_midi_name, 'beat')
    else:
        return filehandler.get_full_midi_chord_labs_path(best_bar_midi_name, 'bar')

def get_actual_best_tab_lab(song: Song) -> str:
    best_tab_name, _ = filehandler.get_actual_best_tab_for_song(song.key)
    #if(best_tab_name != -1000):
    return filehandler.get_full_tab_chord_labs_path(best_tab_name)
    #return filehandler.get_full_tab_chord_labs_path('notapath.txt')

def _write_final_labels(final_labels, lab_path, alphabet):
    """
    Write the row of final labels (after data fusion) from the chord matrix to a .lab file at lab_path

    :param final_labels: The row of final labels after data fusion
    :param lab_path: The path to write the final lab file to
    :param alphabet: The chord vocabulary (used to translate chord indices to chord strings)
    """

    # First, we fill the write_labels list with (start_time, end_time, chord_string) 3-tuples
    write_labels = []
    start_time = -1
    last_chord = ''
    last_added = True
    for i in range(len(final_labels)):
        if final_labels[i] == last_chord:
            last_added = False
        else:
            if not last_added:
                # Write previous chord
                write_labels.append((start_time, float(i) / 100, last_chord))
            # Set parameters for next
            start_time = float(i) / 100
            last_chord = final_labels[i]
    if not last_added:
        # The last chord has not been added yet, so we do it now
        write_labels.append((start_time, float(len(final_labels) - 1) / 100, last_chord))

    # Now write the chords to the .lab file in Harte format
    with open(lab_path, 'w') as write_file:
        for write_label in write_labels:
            chord_string = str(Chord.from_common_tab_notation_string(alphabet[write_label[2]]))
            if chord_string == 'None':
                chord_string = 'N'
            write_file.write('{0} {1} {2}\n'.format(str(write_label[0]), str(write_label[1]), chord_string))

def _random_chord_label_combination(chord_matrix, nr_of_samples):
    """
    For each 10ms segment, pick the chord label from a random source

    :param chord_matrix: The chord label matrix
    :param nr_of_samples: Total number of 10ms samples
    :return: Array with the final chord index for each segment
    """
    final_labels = np.empty(nr_of_samples, dtype=int)
    for sample_nr in range(nr_of_samples):
        # Select at every sample the chord from a random source
        source_nr = filehandler.random.randint(0, chord_matrix.shape[0] - 1)
        final_labels[sample_nr] = chord_matrix[source_nr, sample_nr]
    return final_labels

def _majority_vote_chord_label_combination(chord_matrix, nr_of_samples, alphabet):
    """
    For each 10ms segment, pick the chord label that occurs in most sources (or randomly pick one of the chord labels
    that occur in most sources)

    :param chord_matrix: The chord label matrix
    :param nr_of_samples: Total number of 10ms samples
    :param alphabet: The chord vocabulary
    :return: Array with the final chord index for each segment
    """
    final_labels = np.empty(nr_of_samples, dtype=int)
    for sample_nr in range(nr_of_samples):
        # Select at every sample the most frequent chord (or a random one of multiple most frequent chords)
        frequencies = np.zeros(len(alphabet))
        for lab_nr in range(chord_matrix.shape[0]):
            frequencies[chord_matrix[lab_nr, sample_nr]] += 1
        most_occurring_chord_indices = np.argwhere(frequencies == np.amax(frequencies))
        chosen_chord_index = most_occurring_chord_indices[
            filehandler.random.randint(0, len(most_occurring_chord_indices) - 1)]
        final_labels[sample_nr] = chosen_chord_index
    return final_labels

def _data_fusion_chord_label_combination(chord_matrix, nr_of_samples, alphabet):
    """
    For each 10ms segment, pick the chord label using data fusion (i.e. based on chord label probabilities and source
    accuracies)

    :param chord_matrix: The chord label matrix
    :param nr_of_samples: Total number of 10ms samples
    :param alphabet: The chord vocabulary
    :return: Array with the final chord index for each segment
    """
    # Calculate initial chord label probabilities
    chord_label_probabilities = np.empty((nr_of_samples, len(alphabet)))
    for sample_nr in range(nr_of_samples):
        for chord_nr in range(len(alphabet)):
            chord_label_probabilities[sample_nr, chord_nr] = \
                np.mean([chord_matrix[lab_nr, sample_nr] == chord_nr for lab_nr in range(chord_matrix.shape[0])])

    # Calculate source accuracies
    for i in range(5):
        source_accuracies = np.empty(chord_matrix.shape[0])
        for lab_nr in range(chord_matrix.shape[0]):
            source_accuracies[lab_nr] = np.mean(
                [chord_label_probabilities[sample_nr, chord_matrix[lab_nr, sample_nr]]
                 for sample_nr in range(nr_of_samples)])

        chord_label_probabilities = np.empty((nr_of_samples, len(alphabet)))
        for sample_nr in range(nr_of_samples):
            vote_counts = np.zeros(len(alphabet))
            for lab_nr in range(chord_matrix.shape[0]):
                vote_counts[chord_matrix[lab_nr, sample_nr]] += source_accuracies[lab_nr]
            sum_exp_vote_counts_all_labels = sum([np.exp(vote_counts[chord_nr]) for chord_nr in range(len(alphabet))])
            for chord_nr in range(len(alphabet)):
                chord_label_probabilities[sample_nr, chord_nr] = \
                    np.exp(vote_counts[chord_nr]) / sum_exp_vote_counts_all_labels

    final_labels = np.empty(nr_of_samples, dtype=int)
    for sample_nr in range(nr_of_samples):
        final_labels[sample_nr] = np.argmax(chord_label_probabilities[sample_nr])

    return final_labels

def load_lab_file_into_chord_matrix(lab_path, i, chord_matrix, alphabet, nr_of_samples):
    """
    Load a chord file (in Harte's chord annotation format) into a chord matrix, having the chord label per 10ms sample.
    :param lab_path: Path to the .lab file with chord annotation/estimation
    :param i: Index to the row in the chord_matrix to fill
    :param chord_matrix: Chord matrix to which we add samples from this .lab file
    :param alphabet: The chord vocabulary (used to get the index belonging to a chord string)
    :param nr_of_samples: Total number of samples in the song
    :return: Chord matrix to which we just added samples from this .lab file to row i
    """
    chord_annotation = import_chord_annotation(lab_path)
    for chord_annotation_item in chord_annotation.chord_annotation_items:
        chord_label_alphabet_index = alphabet.get_index_of_chord_in_alphabet(chord_annotation_item.chord)
        for s in range(int(chord_annotation_item.from_time * 100),
                       min(int(chord_annotation_item.to_time * 100), nr_of_samples - 1)):
            chord_matrix[i, s] = chord_label_alphabet_index

def data_fuse_song(song: Song, chord_vocabulary: ChordVocabulary):
    """
    Data fuse a song using all combinations of selection and combination methods, write the final labels to .lab files

    :param song: The song on which we want to apply data fusion
    :param chord_vocabulary: The chord vocabulary
    """
    # Check if data fusion has already been calculated  TODO: make this check more robust
    if path.isfile(filehandler.get_data_fusion_path(song.key, 'df', 'best', 'CHF_2017')):
        return

    # Get list of symbolic lab files (all / expected best)
    well_aligned_midis = get_well_aligned_midis(song)
    all_symbolic_lab_paths = \
        [filehandler.get_full_midi_chord_labs_path(wam, 'bar') for wam in well_aligned_midis] + \
        [filehandler.get_full_midi_chord_labs_path(wam, 'beat') for wam in well_aligned_midis] + \
        [filehandler.get_full_tab_chord_labs_path(t) for t in song.full_tab_paths]
    expected_best_symbolic_lab_paths = []
    if well_aligned_midis:
        expected_best_symbolic_lab_paths.append(
            filehandler.get_full_midi_chord_labs_path(*get_expected_best_midi(song)))
    if [filehandler.get_full_tab_chord_labs_path(t) for t in song.full_tab_paths]:
        expected_best_symbolic_lab_paths.append(
            filehandler.get_full_tab_chord_labs_path(get_expected_best_tab_lab(song)))

    # Remove non-existing files (e.g. tab files in which too little chords were observed)
    all_symbolic_lab_paths = [lab for lab in all_symbolic_lab_paths if filehandler.file_exists(lab)]
    expected_best_symbolic_lab_paths = [lab for lab in expected_best_symbolic_lab_paths if filehandler.file_exists(lab)]

    # Get list of audio lab files
    audio_labs = song.full_mirex_chord_lab_paths
    audio_labs['CHF_2017'] = song.full_chordify_chord_labs_path

    # Sample every 10ms, so 100 samples per second
    song_duration = song.duration
    nr_of_samples = int(ceil(song_duration * 100))

    # Turn the chords list (a list of (key, mode-str, chroma-list) tuples) into an chord_vocabulary (a list of strings)
    # alphabet = _chords_list_to_alphabet(chords_list) TODO remove if possible
    alphabet = ChordAlphabet(chord_vocabulary)

    # Iterate over the two types of selection (all / best)
    for lab_list_i in [0, 1]:
        lab_list = [all_symbolic_lab_paths, expected_best_symbolic_lab_paths][lab_list_i]
        lab_list = [i for i in lab_list if i != '']
        selection_name = ['all', 'best'][lab_list_i]

        # Fill a numpy array with chord labels for each of the lab files
        chord_matrix = np.zeros((len(lab_list) + 1, nr_of_samples), dtype=int)
        for lab_nr in range(len(lab_list)):
            load_lab_file_into_chord_matrix(lab_list[lab_nr], lab_nr, chord_matrix, alphabet, nr_of_samples)

        # Iterate over the audio types:
        for audio_name, audio_lab in audio_labs.items():
            if filehandler.file_exists(audio_lab):
                # Add the lab file to our chord matrix
                load_lab_file_into_chord_matrix(audio_lab, len(lab_list), chord_matrix, alphabet, nr_of_samples)
                # Iterate over the three combination types; calculate labels and write them:
                final_labels_random = _random_chord_label_combination(chord_matrix, nr_of_samples)
                final_labels_majority = _majority_vote_chord_label_combination(chord_matrix, nr_of_samples, alphabet)
                final_labels_data_fusion = _data_fusion_chord_label_combination(chord_matrix, nr_of_samples, alphabet)
                _write_final_labels(final_labels_random,
                                    filehandler.get_data_fusion_path(song.key, 'rnd', selection_name, audio_name),
                                    alphabet)
                _write_final_labels(final_labels_majority,
                                    filehandler.get_data_fusion_path(song.key, 'mv', selection_name, audio_name),
                                    alphabet)
                _write_final_labels(final_labels_data_fusion,
                                    filehandler.get_data_fusion_path(song.key, 'df', selection_name, audio_name),
                                    alphabet)

def data_fuse_song_with_actual_best_midi_and_tab(song: Song, chord_vocabulary: ChordVocabulary):
    """
    Data fuse a song using all combinations of selection and combination methods, write the final labels to .lab files

    :param song: The song on which we want to apply data fusion
    :param chord_vocabulary: The chord vocabulary
    """
    # Check if data fusion has already been calculated  TODO: make this check more robust
    if path.isfile(filehandler.get_data_fusion_path(song.key, 'df', 'actual-best', 'CHF_2017')):
        return

    # Get list of audio lab files
    audio_labs = song.full_mirex_chord_lab_paths
    audio_labs['CHF_2017'] = song.full_chordify_chord_labs_path

    # Sample every 10ms, so 100 samples per second
    song_duration = song.duration
    nr_of_samples = int(ceil(song_duration * 100))

    # Turn the chords list (a list of (key, mode-str, chroma-list) tuples) into an chord_vocabulary (a list of strings)
    alphabet = ChordAlphabet(chord_vocabulary)

    selection_name = 'actual-best'
    lab_list = [get_actual_best_tab_lab(song)]
    # lab_list = [get_actual_best_midi_lab(song), get_actual_best_tab_lab(song)]

    # Fill a numpy array with chord labels for each of the lab files
    chord_matrix = np.zeros((len(lab_list) + 1, nr_of_samples), dtype=int)
    for lab_nr in range(len(lab_list)):
        load_lab_file_into_chord_matrix(lab_list[lab_nr], lab_nr, chord_matrix, alphabet, nr_of_samples)

    # Iterate over the audio types:
    for audio_name, audio_lab in audio_labs.items():
        if filehandler.file_exists(audio_lab):
            # Add the lab file to our chord matrix
            load_lab_file_into_chord_matrix(audio_lab, len(lab_list), chord_matrix, alphabet, nr_of_samples)
            final_labels_data_fusion = _data_fusion_chord_label_combination(chord_matrix, nr_of_samples, alphabet)
            _write_final_labels(final_labels_data_fusion,
                                filehandler.get_data_fusion_path(song.key, 'df', selection_name, audio_name),
                                alphabet)


import numpy as np
from joblib import load

NUM_RECOMM_VIDEOS = 12

# 读取本地测试结果的第一条数据（手动粘贴一下吧！）
input_title = "《黑暗之魂3 》艾利德修女boss战 ——来自高端玩家的教学视频"
input_cateName = "视频-游戏视频"
input_videoId = "c07bd7d0e4784013a234a62f653be295"
input_results = '''
[['《GTA5里等红绿灯的女司机》' '视频-游戏视频' '43d08001f09d474fb68e51d49b0a270b']
 ['《月入百万公交坐过站，隔桌明星一起吃过饭》' '视频-直播剪辑' '198004c6eb1d48f390617e6c405da385']
 ['《一封特别的信》' '视频-直播剪辑' '0b2f8dea47fa418780cd9e68d8f8f01c']
 ['《受邀参加英女王生日会！》' '视频-直播剪辑' 'd975f997a5f44961ae92195d00f4659e']
 ['《洗面奶的故事》' '视频-直播剪辑' '8669c160fc1046aebfd3cd29353fee2e']
 ['陈一发儿：油管直播终极计划——《格局》' '视频-直播剪辑' 'ea477760e11641049e7ca0f139aefd75']
 ['之前一曲不插电的《我们》， 送给各位毕业生' '视频-直播剪辑' '154aabf4d9dd49f6a3df2187ba5c6e6b']
 ['《一个鸡蛋的暴走》——上海徒步公益活动' '视频-日常' '8572790d11a3481a95e3b3634e4b7504']
 ['《白車过隙听过没？》' '视频-游戏视频' '4539d1ccb0244d1c808137077d7998d0']
 ['潜水艇小游戏' '视频-日常' '7e912121d1ac4e9b954c89a6697154ed']
 ['你比特币赚了多少？' '视频-直播剪辑' 'cf0d8273b34342269ea6d898a7d982dc']
 ['分享一些男性健康方面的小知识，不用谢。' '视频-直播剪辑' '5a342d4f3ccd4955a06d64501e4f8bbd']]
'''

def load_model(file_name):
    return load('models/{}'.format(file_name))

def load_dataset():
    f = open("dataset/dataset.txt", "r")
    lines = f.read().splitlines()
    dataset = []
    for line in lines:
        arr = line.split(',')
        dataset.append(arr)
    dataset = np.array(dataset)
    f.close()
    return dataset

def BoW(inputs, codebook):
    vec = [0] * len(codebook)
    indices = []
    if (type(inputs) == list):
        inputs = ''.join(inputs)
    for char in list(inputs):
        index = np.argwhere(codebook == char)
        if (index.shape == (1, 1)):
            vec[index[0][0]] = 1
            indices.append(index[0][0])
    return np.array(vec), np.array(indices)

dataset = load_dataset()
codebook = load_model('model_codebook')
som = load_model('model_som')
knn = load_model('model_knn')

vec = np.array(BoW([input_title, input_cateName], codebook)[0])
knn_results = knn.kneighbors([som.winner(vec)])[1][0]

recomms = [data[[0, 1, 5]] for data in dataset[knn_results]]
filters = np.array(
    list(filter(lambda data: data[2] != input_videoId, recomms)))

print(str(filters))
print(input_results)


# -*- coding: utf-8 -*-
# @Brief: 获取预测框的数据，并转mAP的txt检测格式
from predict import FasterRCNN
from PIL import Image
import config.config as cfg
import os

import tensorflow as tf
from core.anchorGenerate import get_anchors
from core.boxParse import BoundingBox
import numpy as np
from tqdm import tqdm

class Frcnn(FasterRCNN):
    def __init__(self, weight_path):
        self.weight_path = weight_path
        super().__init__(weight_path)

    def detect_single_image(self, image, image_id):
        f = open("./input/detection-results/" + image_id + ".txt", "w")

        resize_image = self.process_image(image)
        old_w, old_h = image.size
        # 计算特征层下的宽度、高度
        new_h, new_w = np.shape(resize_image)[1:-1]
        # 计算图片输入到rpn的输出shape，[-1]是因为rpn的输出list，最后一个是共享特征层
        rpn_height, rpn_width = self.model_rpn.compute_output_shape((1, new_h, new_w, 3))[-1][1:-1]

        predict_rpn = self.model_rpn(resize_image, training=False)

        share_layer = predict_rpn[2]

        anchors = get_anchors(share_layer_shape=(rpn_width, rpn_height), image_shape=(new_w, new_h))
        box_parse = BoundingBox(anchors)
        predict_boxes = box_parse.detection_out(predict_rpn, anchors,  confidence_threshold=0)
        predict_boxes = predict_boxes[0]

        predict_boxes[:, 0] = np.array(np.round(predict_boxes[:, 0]*new_w/cfg.rpn_stride), dtype=np.int32)
        predict_boxes[:, 1] = np.array(np.round(predict_boxes[:, 1]*new_h/cfg.rpn_stride), dtype=np.int32)
        predict_boxes[:, 2] = np.array(np.round(predict_boxes[:, 2]*new_w/cfg.rpn_stride), dtype=np.int32)
        predict_boxes[:, 3] = np.array(np.round(predict_boxes[:, 3]*new_h/cfg.rpn_stride), dtype=np.int32)

        predict_boxes[:, 2] -= predict_boxes[:, 0]
        predict_boxes[:, 3] -= predict_boxes[:, 1]

        delete_line = []
        for i, r in enumerate(predict_boxes):
            if r[2] < 1 or r[3] < 1:
                delete_line.append(i)

        predict_boxes = np.delete(predict_boxes, delete_line, axis=0)

        boxes = []
        probs = []
        labels = []

        for i in range(predict_boxes.shape[0] // cfg.num_rois + 1):
            # 把predict_boxes分成多个(128，4）
            rois = predict_boxes[cfg.num_rois * i: cfg.num_rois * (i+1), :]
            rois = np.expand_dims(rois, axis=0)

            if rois.shape[1] == 0:
                print("Your model can't detect any bounding boxes.")
                exit(0)

            # 把不够128框的部分，补足128个
            if i == predict_boxes.shape[0] // cfg.num_rois:

                curr_shape = rois.shape

                target_shape = (curr_shape[0], cfg.num_rois, curr_shape[2])

                rois_padded = np.zeros(target_shape).astype(rois.dtype)
                # 前面框足够的部分不变
                rois_padded[:, :curr_shape[1], :] = rois
                # 后面不够的地方用同一个数据补足
                rois_padded[0, curr_shape[1]:, :] = rois[0, 0, :]

                rois = rois_padded

            p_cls, p_regr = self.model_classifier([share_layer, rois], training=False)

            for j in range(p_cls.shape[1]):
                # 如果这个框置信度都小于阈值，那就直接跳过了
                conf = np.max(p_cls[0, j, :-1])
                if conf < self.confidence:
                    continue

                x, y, w, h = rois[0, j, :]

                # 获取最大分类值的索引
                cls_num = np.argmax(p_cls[0, j, :-1])

                tx, ty, tw, th = p_regr[0, j, 4 * cls_num: 4 * (cls_num + 1)]

                # 在当时有乘上一个系数，现在就要除回去
                tx /= cfg.classifier_regr_std[0]
                ty /= cfg.classifier_regr_std[1]
                tw /= cfg.classifier_regr_std[2]
                th /= cfg.classifier_regr_std[3]

                # 左上角坐标变换为中心坐标
                cx = x + w / 2.0
                cy = y + h / 2.0
                # 论文公式的逆变换
                cx1 = tx * w + cx
                cy1 = ty * h + cy
                w1 = np.exp(tw) * w
                h1 = np.exp(th) * h

                # 再转换回左上角和右下角的坐标
                x1 = cx1 - w1 / 2.0
                y1 = cy1 - h1 / 2.0
                x2 = cx1 + w1 / 2.0
                y2 = cy1 + h1 / 2.0

                x1 = int(tf.round(x1))
                y1 = int(tf.round(y1))
                x2 = int(tf.round(x2))
                y2 = int(tf.round(y2))

                boxes.append([x1, y1, x2, y2])
                probs.append(conf)
                labels.append(cls_num)

        if len(boxes) == 0:
            return image

        # 筛选出其中得分高于confidence的框
        labels = np.array(labels)
        probs = np.array(probs)
        boxes = np.array(boxes, dtype=np.float32)

        # 变换为标准尺度的框
        boxes[:, 0] = boxes[:, 0] * cfg.rpn_stride / new_w
        boxes[:, 1] = boxes[:, 1] * cfg.rpn_stride / new_h
        boxes[:, 2] = boxes[:, 2] * cfg.rpn_stride / new_w
        boxes[:, 3] = boxes[:, 3] * cfg.rpn_stride / new_h

        result = None
        for c in range(cfg.num_classes - 1):

            mask = labels == c

            if len(probs[mask]) > 0:
                # 取出得分高于confidence_threshold的框
                boxes_to_process = boxes[mask]
                score_to_process = probs[mask]

                nms_index = tf.image.non_max_suppression(boxes_to_process, score_to_process, 300,
                                                         iou_threshold=0.4, score_threshold=self.confidence)

                nms_boxes = tf.gather(boxes_to_process, nms_index).numpy()
                nms_score = tf.gather(score_to_process, nms_index).numpy()
                nms_label = c * np.ones((len(nms_index), 1))

                if result is None:
                    result = np.concatenate((nms_label, np.expand_dims(nms_score, axis=-1), nms_boxes), axis=1)
                else:
                    result = np.vstack((result,
                                        np.concatenate((nms_label, np.expand_dims(nms_score, axis=-1), nms_boxes), axis=1)))

        top_label = result[:, 0]
        top_conf = result[:, 1]
        top_boxes = result[:, 2:]

        top_boxes[:, 0] = top_boxes[:, 0] * old_w
        top_boxes[:, 1] = top_boxes[:, 1] * old_h
        top_boxes[:, 2] = top_boxes[:, 2] * old_w
        top_boxes[:, 3] = top_boxes[:, 3] * old_h

        for i, c in enumerate(top_label):

            predicted_class = self.class_names[int(c)]
            score = top_conf[i]

            left, top, right, bottom = top_boxes[i]

            f.write("{} {:.6} {} {} {} {}\n".format(
                predicted_class, score, str(int(left)), str(int(top)), str(int(right)), str(int(bottom))))

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    gpus = tf.config.experimental.list_physical_devices("GPU")
    if gpus:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)

    frcnn = Frcnn("../model/voc_1.7687.h5")

    image_infos = open("../config/test.txt").read().strip().split('\n')

    if not os.path.exists("./input"):
        os.makedirs("./input")
    if not os.path.exists("./input/detection-results"):
        os.makedirs("./input/detection-results")
    if not os.path.exists("./input/images-optional"):
        os.makedirs("./input/images-optional")

    process_bar = tqdm(range(len(image_infos)), ncols=100, unit="step")
    for i in process_bar:
        image_boxes = image_infos[i].split(' ')
        image_path = image_boxes[0]
        image_id = os.path.basename(image_path)[:-4]

        image = Image.open(image_path)
        # image = Image.open(image_path)
        # 开启后在之后计算mAP可以可视化
        # image.save("./input/images-optional/"+image_id+".jpg")
        frcnn.detect_single_image(image, image_id)

    print("Conversion completed!")


from basicsr.utils.registry import ARCH_REGISTRY
from torch import nn as nn
from torch.nn import functional as F
from torch.nn.utils import spectral_norm
import torch

class add_attn(nn.Module):

    def __init__(self, x_channels, g_channels=256):
        super(add_attn, self).__init__()
        self.W = nn.Sequential(
            nn.Conv2d(x_channels, x_channels, kernel_size=1, stride=1, padding=0), nn.BatchNorm2d(x_channels))
        self.theta = nn.Conv2d(x_channels, x_channels, kernel_size=2, stride=2, padding=0, bias=False)

        self.phi = nn.Conv2d(g_channels, x_channels, kernel_size=1, stride=1, padding=0, bias=True)
        self.psi = nn.Conv2d(x_channels, out_channels=1, kernel_size=1, stride=1, padding=0, bias=True)

    def forward(self, x, g):
        input_size = x.size()
        batch_size = input_size[0]
        assert batch_size == g.size(0)

        theta_x = self.theta(x)
        theta_x_size = theta_x.size()
        phi_g = F.interpolate(self.phi(g), size=theta_x_size[2:], mode='bilinear', align_corners=False)
        f = F.relu(theta_x + phi_g, inplace=True)

        sigm_psi_f = torch.sigmoid(self.psi(f))
        sigm_psi_f = F.interpolate(sigm_psi_f, size=input_size[2:], mode='bilinear', align_corners=False)

        y = sigm_psi_f.expand_as(x) * x
        W_y = self.W(y)
        return W_y, sigm_psi_f

class unetCat(nn.Module):

    def __init__(self, dim_in, dim_out):
        super(unetCat, self).__init__()
        norm = spectral_norm
        self.convU = norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1, bias=False))

    def forward(self, input_1, input_2):
        # Upsampling
        input_2 = F.interpolate(input_2, scale_factor=2, mode='bilinear', align_corners=False)

        output_2 = F.leaky_relu(self.convU(input_2), negative_slope=0.2, inplace=True)

        offset = output_2.size()[2] - input_1.size()[2]
        padding = 2 * [offset // 2, offset // 2]
        output_1 = F.pad(input_1, padding)
        y = torch.cat([output_1, output_2], 1)
        return y

class UNetDiscriminatorAesrgan(nn.Module):
    """Defines a U-Net discriminator with spectral normalization (SN)"""

    def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
        super(UNetDiscriminatorAesrgan, self).__init__()
        norm = spectral_norm

        self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)

        self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 3, 2, 1, bias=False))
        self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 3, 2, 1, bias=False))

        # Center
        self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 3, 2, 1, bias=False))

        self.gating = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 1, 1, 1, bias=False))

        # attention Blocks
        self.attn_1 = add_attn(x_channels=num_feat * 4, g_channels=num_feat * 4)
        self.attn_2 = add_attn(x_channels=num_feat * 2, g_channels=num_feat * 4)
        self.attn_3 = add_attn(x_channels=num_feat, g_channels=num_feat * 4)

        # Cat
        self.cat_1 = unetCat(dim_in=num_feat * 8, dim_out=num_feat * 4)
        self.cat_2 = unetCat(dim_in=num_feat * 4, dim_out=num_feat * 2)
        self.cat_3 = unetCat(dim_in=num_feat * 2, dim_out=num_feat)

        # upsample
        self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
        self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
        self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))

        # extra
        self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
        self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
        self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)

    def forward(self, x):
        x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
        x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
        x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
        x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)

        gated = F.leaky_relu(self.gating(x3), negative_slope=0.2, inplace=True)

        # Attention
        attn1, ly1 = self.attn_1(x2, gated)
        attn2, ly2 = self.attn_2(x1, gated)
        attn3, ly3 = self.attn_3(x0, gated)

        return ly1, ly2, ly3

if __name__ == "__main__":
    from torchsummary import summary
    from PIL import Image
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--img_path', default=r"..\inputs\img_015_SRF_4_HR.png", help='image path')
    parser.add_argument(
        '--model_path', default=r"..\experiments\pretrained_models\single", help='single model list path')
    parser.add_argument('--save_path', default=r".\Visual", help='path to save the heat map')

    args = parser.parse_args()

    uNet = UNetDiscriminatorAesrgan(3, num_feat=64, skip_connection=True)
    import numpy as np

    imgpath = args.img_path
    modelpath = args.model_path
    save_dir = args.save_path

    import cv2
    import torchvision.transforms as transforms
    img = cv2.imread(imgpath)

    import os
    import shutil
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    else:
        shutil.rmtree(save_dir)
        os.mkdir(save_dir)

    for i in range(5000, 400000, 5000):
        path = modelpath + r"\net_d_" + str(i) + ".pth"
        l = torch.load(path)

        p = uNet.state_dict()
        n = {}
        n["params"] = {}
        for k in l["params"]:
            k1 = k.replace("layer0.", "")
            k2 = k1.replace("layer1.", "")
            n["params"][k2] = l["params"][k]
        uNet.load_state_dict(n["params"], strict=True)

        input = transforms.ToTensor()(img)
        input = input.unsqueeze(0)

        AttentionLayer1, AttentionLayer2, AttentionLayer3 = uNet(input)

        A1 = AttentionLayer1.detach().numpy()
        A1 = np.squeeze(A1)
        A1 = A1 * 255
        A1 = cv2.applyColorMap(np.uint8(A1), cv2.COLORMAP_JET)

        save_path = save_dir + "\A1_" + str(i) + ".png"
        cv2.imwrite(save_path, A1)

        A2 = AttentionLayer2.detach().numpy()
        A2 = np.squeeze(A2)
        A2 = A2 * 255
        A2 = cv2.applyColorMap(np.uint8(A2), cv2.COLORMAP_JET)
        save_path = save_dir + "\A2_" + str(i) + ".png"
        cv2.imwrite(save_path, A2)

        A3 = AttentionLayer3.detach().numpy()
        A3 = np.squeeze(A3)
        A3 = A3 * 255
        A3 = cv2.applyColorMap(np.uint8(A3), cv2.COLORMAP_JET)
        save_path = save_dir + "\A3_" + str(i) + ".png"
        cv2.imwrite(save_path, A3)


import pyqtgraph as pg
import numpy as np

Terrains = {
    2213004848: {"TerrainType": "grassland",                "color": pg.mkBrush(pg.mkColor(np.array([103, 125, 48]))),  "searchTag": b'\x30\xc6\xe7\x83'},
    1855786096: {"TerrainType": "grassland (hills)",        "color": pg.mkBrush(pg.mkColor(np.array([103, 125, 48]))),  "searchTag": b'\x70\x0C\x9D\x6E'},
    1602466867: {"TerrainType": "grassland (mountains)",    "color": pg.mkBrush(pg.mkColor(np.array([132, 133, 134]))), "searchTag": b'\x33\xB4\x83\x5F'},
    4226188894: {"TerrainType": "plains",                   "color": pg.mkBrush(pg.mkColor(np.array([159, 159, 53]))),  "searchTag": b'\x5E\x86\xE6\xFB'},
    3872285854: {"TerrainType": "plains (hills)",           "color": pg.mkBrush(pg.mkColor(np.array([159, 159, 53]))),  "searchTag": b'\x9E\x64\xCE\xE6'},
    2746853616: {"TerrainType": "plains (mountains)",       "color": pg.mkBrush(pg.mkColor(np.array([132, 133, 134]))), "searchTag": b'\xF0\xA8\xB9\xA3'},
    3852995116: {"TerrainType": "desert",                   "color": pg.mkBrush(pg.mkColor(np.array([236, 196, 111]))), "searchTag": b'\x2C\x0A\xA8\xE5'},
    3108058291: {"TerrainType": "desert (hills)",           "color": pg.mkBrush(pg.mkColor(np.array([236, 196, 111]))), "searchTag": b'\xB3\x34\x41\xB9'},
    1418772217: {"TerrainType": "desert (mountains)",       "color": pg.mkBrush(pg.mkColor(np.array([132, 133, 134]))), "searchTag": b'\xF9\xBE\x90\x54'},
    1223859883: {"TerrainType": "tundra",                   "color": pg.mkBrush(pg.mkColor(np.array([171, 170, 139]))), "searchTag": b'\xAB\x9E\xF2\x48'},
    3949113590: {"TerrainType": "tundra (hills)",           "color": pg.mkBrush(pg.mkColor(np.array([171, 170, 139]))), "searchTag": b'\xF6\xB0\x62\xEB'},
    3746160061: {"TerrainType": "tundra (mountains)",       "color": pg.mkBrush(pg.mkColor(np.array([132, 133, 134]))), "searchTag": b'\xBD\xDD\x49\xDF'},
    1743422479: {"TerrainType": "snow",                     "color": pg.mkBrush(pg.mkColor(np.array([204, 223, 243]))), "searchTag": b'\x0F\x84\xEA\x67'},
    3842183808: {"TerrainType": "snow (hills)",             "color": pg.mkBrush(pg.mkColor(np.array([204, 223, 243]))), "searchTag": b'\x80\x12\x03\xE5'},
    699483892:  {"TerrainType": "snow (mountains)",         "color": pg.mkBrush(pg.mkColor(np.array([132, 133, 134]))), "searchTag": b'\xF4\x46\xB1\x29'},
    1204357597: {"TerrainType": "Ocean",                    "color": pg.mkBrush(pg.mkColor(np.array([45, 49, 86]))),    "searchTag": b'\xDD\x09\xC9\x47'},
    1248885265: {"TerrainType": "Coast",                    "color": pg.mkBrush(pg.mkColor(np.array([45, 89, 120]))),   "searchTag": b'\x11\x7A\x70\x4A'},
}

Features = {  # + goodyhut codes
    4294967295: {"FeatureType": "NoFeature",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    662085235:  {"FeatureType": "Aerodome",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    607694793:  {"FeatureType": "Airstrip",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2588244546: {"FeatureType": "Alcazar",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    170002756:  {"FeatureType": "AntiquitySite",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3727362748: {"FeatureType": "BarbCamp",             "color": pg.mkBrush(pg.mkColor(np.array([183, 20, 20])))},
    1765528555: {"FeatureType": "Barricade",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2534355737: {"FeatureType": "Batey",                "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2593014542: {"FeatureType": "BurningWoods",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3000670711: {"FeatureType": "BurntWoods",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2042437235: {"FeatureType": "BurningRainforest",    "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1374274698: {"FeatureType": "BurntRainforest",      "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1957694686: {"FeatureType": "CahokiaMounds",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2475408324: {"FeatureType": "Camp",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    248721113:  {"FeatureType": "CityPark",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    38972510:   {"FeatureType": "Chateau",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2176791945: {"FeatureType": "ColossalHead",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    119268998:  {"FeatureType": "Corporation",          "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1387624230: {"FeatureType": "DesertFloodplains",    "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1555056577: {"FeatureType": "GeothermalFissure",    "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2686085371: {"FeatureType": "GolfCourse",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1434118760: {"FeatureType": "GoodyHut",             "color": pg.mkBrush(pg.mkColor(np.array([241, 209, 100])))},
    2414989200: {"FeatureType": "GreatWall",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2973166745: {"FeatureType": "GrasslandFloodplains", "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    168372657:  {"FeatureType": "Farm",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    578093457:  {"FeatureType": "FishingBoats",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1694280827: {"FeatureType": "Fort",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1542194068: {"FeatureType": "Ice",                  "color": pg.mkBrush(pg.mkColor(np.array([171, 188, 219])))},
    168256235:  {"FeatureType": "IceHockeyRink",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    554019297:  {"FeatureType": "Industry",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    874973008:  {"FeatureType": "Kampung",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3808671749: {"FeatureType": "Kurgan",               "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2048582848: {"FeatureType": "LumberMill",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    592796099:  {"FeatureType": "March",                "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    590245655:  {"FeatureType": "Mekewap",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    933499311:  {"FeatureType": "MeteorSite?",          "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2929403647: {"FeatureType": "MeteorSite",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1001859687: {"FeatureType": "Mine",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2970879537: {"FeatureType": "Moai",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2457279608: {"FeatureType": "Monastery",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3108964764: {"FeatureType": "MountainTunnel",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1163354216: {"FeatureType": "NazcaLine",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1646220195: {"FeatureType": "Oasis",                "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3898338829: {"FeatureType": "OffshoreOilRig",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1058336991: {"FeatureType": "OffshoreWindFarm",     "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2939453696: {"FeatureType": "OilWell",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1151932567: {"FeatureType": "OutbackStation",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1719494282: {"FeatureType": "Pairidaeza",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    154488225:  {"FeatureType": "Pasture",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2265104033: {"FeatureType": "PlainsFloodplains",    "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1523996587: {"FeatureType": "Plantation",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1655685824: {"FeatureType": "Ponder",               "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    4214473799: {"FeatureType": "Quarry",               "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3910227001: {"FeatureType": "Rainforest",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3943953367: {"FeatureType": "Reef",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2334474827: {"FeatureType": "RockHewnChurch",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1195876395: {"FeatureType": "RomanFort",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    570930386:  {"FeatureType": "SeasideResort",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2182120684: {"FeatureType": "Sphinx",               "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2135005470: {"FeatureType": "SkiResort",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2127465633: {"FeatureType": "SolarFarm",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1021479658: {"FeatureType": "TradingDome",          "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    4092364633: {"FeatureType": "Trap",                 "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2005950930: {"FeatureType": "VolcanicSoil",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2771557557: {"FeatureType": "Volcano",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    520634903:  {"FeatureType": "WindFarm",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3029993898: {"FeatureType": "Ziggurat",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},

    # Wonders
    2773098215: {"FeatureType": "ChocolateHills",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2072857400: {"FeatureType": "CliffsOfDover",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1253249868: {"FeatureType": "CraterLake",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1894931018: {"FeatureType": "DeadSea",              "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1781837174: {"FeatureType": "DelicateArch",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    656656409:  {"FeatureType": "EyeOfTheSahara",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    853279971:  {"FeatureType": "Eyjafjallajokull",     "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    226585075:  {"FeatureType": "Galapagos",            "color": pg.mkBrush(pg.mkColor(np.array([243, 212, 1])))},
    1653648472: {"FeatureType": "GiantsCauseway",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    978567123:  {"FeatureType": "Gobustan",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3583470385: {"FeatureType": "GreatBarrierReef",     "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    172143503:  {"FeatureType": "FountainOfYouth",      "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1084731038: {"FeatureType": "HalongBay",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    501201242:  {"FeatureType": "IkKil",                "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    482180537:  {"FeatureType": "LakeRetba",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    308108439:  {"FeatureType": "Lysefjord",            "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3082807326: {"FeatureType": "MatoTipila",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    224663563:  {"FeatureType": "Matterhorn",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1080563218: {"FeatureType": "MountEverest",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2354893825: {"FeatureType": "MountKilimanjaro",     "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3861018290: {"FeatureType": "MountRoraima",         "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2284878731: {"FeatureType": "MountVesuvius",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1491718565: {"FeatureType": "Pamukkale",            "color": pg.mkBrush(pg.mkColor(np.array([132, 133, 134])))},
    1394593257: {"FeatureType": "Pantanal",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    1624801336: {"FeatureType": "Piopiotahi",           "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    3738443622: {"FeatureType": "SaharaElBeyda",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    2666510884: {"FeatureType": "TorresDelPaine",       "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    981353664:  {"FeatureType": "TsingyDeBemaraha",     "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    21: {"FeatureType": "UbsunurHollow",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    4277999168: {"FeatureType": "Uluru",                "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
    4283954533: {"FeatureType": "Yosemite",             "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},  # 1253249868
    23: {"FeatureType": "ZhangyeDanxia",        "color": pg.mkBrush(pg.mkColor(np.array([0, 0, 0])))},
}


import argparse
import multiprocessing as mp
import os
from typing import Any, List

from loguru import logger
import numpy as np
from sklearn.svm import LinearSVC
from sklearn.metrics import average_precision_score
from sklearn.model_selection import cross_val_score
import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from virtex.config import Config
from virtex.factories import PretrainingModelFactory, DownstreamDatasetFactory
from virtex.utils.checkpointing import CheckpointManager
from virtex.utils.common import common_parser, common_setup

parser = common_parser(
    description="Train SVMs for VOC2007 classification on a pretrained model."
)
group = parser.add_argument_group("Downstream config arguments.")
group.add_argument(
    "--down-config", metavar="FILE", help="Path to a downstream config file."
)
group.add_argument(
    "--down-config-override",
    nargs="*",
    default=[],
    help="A list of key-value pairs to modify downstream config params.",
)

# fmt: off
parser.add_argument_group("Checkpointing")
parser.add_argument(
    "--weight-init", choices=["random", "imagenet", "torchvision", "virtex"],
    default="virtex", help="""How to initialize weights:
        1. 'random' initializes all weights randomly
        2. 'imagenet' initializes backbone weights from torchvision model zoo
        3. {'torchvision', 'virtex'} load state dict from --checkpoint-path
            - with 'torchvision', state dict would be from PyTorch's training
              script.
            - with 'virtex' it should be for our full pretrained model."""
)
parser.add_argument(
    "--checkpoint-path",
    help="Path to load checkpoint and run downstream task evaluation."
)
# fmt: on

def train_test_single_svm(args):

    feats_train, tgts_train, feats_test, tgts_test, cls_name = args
    SVM_COSTS = [0.01, 0.1, 1.0, 10.0]

    cls_labels = np.copy(tgts_train)
    # Meaning of labels in VOC/COCO original loaded target files:
    # label 0 = not present, set it to -1 as svm train target
    # label 1 = present. Make the svm train target labels as -1, 1.
    cls_labels[np.where(cls_labels == 0)] = -1

    # See which cost maximizes the AP for this class.
    best_crossval_ap: float = 0.0
    best_crossval_clf = None
    best_cost: float = 0.0

    # fmt: off
    for cost in SVM_COSTS:
        clf = LinearSVC(
            C=cost, class_weight={1: 2, -1: 1}, penalty="l2",
            loss="squared_hinge", max_iter=2000,
        )
        ap_scores = cross_val_score(
            clf, feats_train, cls_labels, cv=3, scoring="average_precision",
        )
        clf.fit(feats_train, cls_labels)

        # Keep track of best SVM (based on cost) for each class.
        if ap_scores.mean() > best_crossval_ap:
            best_crossval_ap = ap_scores.mean()
            best_crossval_clf = clf
            best_cost = cost

    logger.info(f"Best SVM {cls_name}: cost {best_cost}, mAP {best_crossval_ap * 100}")
    # fmt: on

    # -------------------------------------------------------------------------
    #   TEST THE TRAINED SVM (PER CLASS)
    # -------------------------------------------------------------------------
    predictions = best_crossval_clf.decision_function(feats_test)
    evaluate_data_inds = tgts_test != -1
    eval_preds = predictions[evaluate_data_inds]

    cls_labels = np.copy(tgts_test)
    eval_cls_labels = cls_labels[evaluate_data_inds]
    eval_cls_labels[np.where(eval_cls_labels == 0)] = -1

    # Binarize class labels to make AP targets.
    targets = eval_cls_labels > 0
    return average_precision_score(targets, eval_preds)

def main(_A: argparse.Namespace):

    if _A.num_gpus_per_machine == 0:
        # Set device as CPU if num_gpus_per_machine = 0.
        device = torch.device("cpu")
    else:
        # Get the current device (this will be zero here by default).
        device = torch.cuda.current_device()

    # Create a downstream config object (this will be immutable) and perform
    # common setup such as logging and setting up serialization directory.
    _DOWNC = Config(_A.down_config, _A.down_config_override)
    common_setup(_DOWNC, _A, job_type="downstream")

    # Create a (pretraining) config object and backup in serialization directory.
    _C = Config(_A.config, _A.config_override)
    _C.dump(os.path.join(_A.serialization_dir, "pretrain_config.yaml"))

    # -------------------------------------------------------------------------
    #   INSTANTIATE DATALOADER, MODEL, AND FEATURE EXTRACTOR
    # -------------------------------------------------------------------------

    train_dataset = DownstreamDatasetFactory.from_config(_DOWNC, split="trainval")
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=_DOWNC.OPTIM.BATCH_SIZE,
        num_workers=_A.cpu_workers,
        pin_memory=True,
    )
    test_dataset = DownstreamDatasetFactory.from_config(_DOWNC, split="test")
    test_dataloader = DataLoader(
        test_dataset,
        batch_size=_DOWNC.OPTIM.BATCH_SIZE,
        num_workers=_A.cpu_workers,
        pin_memory=True,
    )
    NUM_CLASSES = len(train_dataset.class_names)

    # Initialize from a checkpoint, but only keep the visual module.
    model = PretrainingModelFactory.from_config(_C)

    # Load weights according to the init method, do nothing for `random`, and
    # `imagenet` is already taken care of.
    if _A.weight_init == "virtex":
        ITERATION = CheckpointManager(model=model).load(_A.checkpoint_path)
    elif _A.weight_init == "torchvision":
        # Keep strict=False because this state dict may have weights for
        # last fc layer.
        model.visual.cnn.load_state_dict(
            torch.load(_A.checkpoint_path, map_location="cpu")["state_dict"],
            strict=False,
        )
        # Set ``ITERATION`` to a dummy value.
        ITERATION = 0

    # Transfer model to GPU and set to eval mode. This is a torchvision model
    # and it returns features as ``(batch_size, 2048, 7, 7)``.
    model = model.visual.cnn.to(device).eval()

    # -------------------------------------------------------------------------
    #   EXTRACT FEATURES FOR TRAINING SVMs
    # -------------------------------------------------------------------------

    features_train: List[torch.Tensor] = []
    targets_train: List[torch.Tensor] = []

    features_test: List[torch.Tensor] = []
    targets_test: List[torch.Tensor] = []

    # VOC07 is small, extract all features and keep them in memory.
    with torch.no_grad():
        for batch in tqdm(train_dataloader, desc="Extracting train features:"):
            features = model(batch["image"].to(device))

            # Global average pool features. Assume the tensor is in NCHW format.
            if len(features.size()) > 2:
                features = features.view(features.size(0), features.size(1), -1)

                # shape: (batch_size, visual_feature_size)
                features = features.mean(dim=-1)

            # shape: (batch_size, visual_feature_size)
            features = features.view(features.size(0), -1)

            # L2-normalize the global average pooled features.
            features = features / torch.norm(features, dim=-1).unsqueeze(-1)

            features_train.append(features.cpu())
            targets_train.append(batch["label"])

        # Similarly extract test features.
        for batch in tqdm(test_dataloader, desc="Extracting test features:"):
            features = model(batch["image"].to(device))

            if len(features.size()) > 2:
                features = features.view(features.size(0), features.size(1), -1)
                features = features.mean(dim=-1)

            features = features.view(features.size(0), -1)
            features = features / torch.norm(features, dim=-1).unsqueeze(-1)

            features_test.append(features.cpu())
            targets_test.append(batch["label"])

    # Convert batches of features/targets to one large numpy array
    features_train = torch.cat(features_train, dim=0).numpy()
    targets_train = torch.cat(targets_train, dim=0).numpy().astype(np.int32)

    features_test = torch.cat(features_test, dim=0).numpy()
    targets_test = torch.cat(targets_test, dim=0).numpy().astype(np.int32)

    # -------------------------------------------------------------------------
    #   TRAIN AND TEST SVMs WITH EXTRACTED FEATURES
    # -------------------------------------------------------------------------

    input_args: List[Any] = []

    # Iterate over all VOC07 classes and train one-vs-all linear SVMs.
    for cls_idx in range(NUM_CLASSES):
        # fmt: off
        input_args.append((
            features_train, targets_train[:, cls_idx],
            features_test, targets_test[:, cls_idx],
            train_dataset.class_names[cls_idx],
        ))
        # fmt: on

    pool = mp.Pool(processes=_A.cpu_workers)
    pool_output = pool.map(train_test_single_svm, input_args)

    # -------------------------------------------------------------------------
    #   TENSORBOARD LOGGING (RELEVANT MAINLY FOR weight_init=checkpoint)
    # -------------------------------------------------------------------------

    # Tensorboard writer for logging mAP scores. This is useful especially
    # when weight_init=checkpoint (which maybe be coming from a training job).
    tensorboard_writer = SummaryWriter(log_dir=_A.serialization_dir)

    # Test set mAP for each class, for features from every layer.
    test_map = torch.tensor(pool_output).mean()
    logger.info(f"Iteration: {ITERATION}, mAP: {test_map * 100}")
    tensorboard_writer.add_scalars(
        "metrics/voc07_clf", {f"voc07_mAP": test_map * 100}, ITERATION
    )

if __name__ == "__main__":
    _A = parser.parse_args()

    if _A.num_gpus_per_machine > 1:
        raise ValueError("Using multiple GPUs is not supported for this script.")

    # Add an arg in config override if `--weight-init` is imagenet.
    if _A.weight_init == "imagenet":
        _A.config_override.extend(["MODEL.VISUAL.PRETRAINED", True])

    # No distributed training here, just a single process.
    main(_A)


import numpy as np
import ipywidgets
from IPython.display import display
import PIL.Image as Image

_terrain_attrs = (
    "internal_wall_edges",
    "floor",
    "floor_special_tile",
    "wall_edges",
    "keys_switches",
    "door",
    "unk7",
    "unk8",
    "unk9",
    "unk12",
    "big_wooden_thing",
    "big_boulder",
)

class TerrainSprites(dict):
    def __init__(self, game, level_asset):
        super(dict, self).__init__()
        self.game = game
        self.level_asset = level_asset
        offset = self.game.assets["INIT"].sprite_offsets.terrain
        sprites = self.game.resources.sprites
        props = self.level_asset.properties

        for attr in _terrain_attrs:
            self[attr] = sprites[offset + getattr(props, attr)]
        offset = self.game.assets["INIT"].sprite_offsets.wall_decoration
        self["wall_overlay_tiles"] = sprites[offset + props.wall_overlay_tiles]

    def _ipython_display_(self):
        children = []
        titles = []
        for attr in _terrain_attrs + ("wall_overlay_tiles",):
            titles.append(attr.replace("_", " ").capitalize())
            o = ipywidgets.Output()
            with o:
                display(self[attr])
            children.append(ipywidgets.HBox([ipywidgets.Label(attr), o]))
        t = ipywidgets.Tab(children)
        for i, title in enumerate(titles):
            t.set_title(i, title)
        display(t)

class LevelMap:
    def __init__(self, game, level_number):
        self.game = game
        self.level_asset = game.assets["LEVELS"].levels[level_number]
        self.tiles = np.frombuffer(self.level_asset.map, dtype="u1").reshape(
            (self.level_asset.height, self.level_asset.width)
        )

        self.items = {}
        self.info = {}
        # Always one extra to signal terminus of item list
        for item in self.level_asset.items[:-1]:
            self.items[item.x, item.y] = [game.objects[_] for _ in item.info.items]
            self.info[item.x, item.y] = item.info

        # Setup terrain sprites
        self.terrain_sprites = TerrainSprites(self.game, self.level_asset)

    def create_map(self):
        # OK!  This will be our really hard part.  We have to do an isometric draw.
        # The way this works is that we have a uint8 for each tile.  The values they take on range from 0 .. 53,
        # but with a few missing. (18, 19, 20, 22)
        #
        # Our walls have 15 "upper" sprites and 15 "lower" sprites (in that order)
        # We have 8 floor tiles, 8 "special" floor tiles, and 8 "wall edges". 4 keys.
        # Doors I haven't totally worked out.
        # For a wall edge, the sprites are 32 across, 16 high. But, we end up
        # with 2 to make a wall.  So 32x32 net.  Floor tiles are 16 high, 32 across.
        #
        # We're going to roughly follow
        # https://stackoverflow.com/questions/892811/drawing-isometric-game-worlds
        # and note that our tiles are oriented with our origin in what would be
        # the top of the diamond.
        # How big does our image need to be?
        # Add one on to the height for the top/bottom
        w, h = 32, 16
        image_shape = (
            (self.tiles.shape[0] + self.tiles.shape[1] - 1) * w // 2,
            (self.tiles.shape[0] + self.tiles.shape[1] - 1) * h // 2,
        )
        image = Image.new("RGBA", image_shape)
        tile_frame = self.terrain_sprites["floor"].frames[0]
        offset = (self.tiles.shape[0] - 1) * w // 2
        # i is for y, j for x
        for i in range(self.tiles.shape[0]):
            for j in range(self.tiles.shape[1]):
                if self.tiles[i, j] == 255:
                    continue
                start_x = j * w // 2 - i * w // 2 + offset
                start_y = i * h // 2 + j * h // 2
                image.paste(tile_frame, (start_x, start_y))
        return image