# -*- coding: utf-8 -*-
from __future__ import print_function

import os
import sys
import random
import jieba
from time import strftime, gmtime

import pickle
import codecs
from keras.optimizers import RMSprop, Adam
from scipy.stats import rankdata
from heapq import nlargest
from gensim.models import Word2Vec

from keras_models_v3 import *

random.seed(42)

class Evaluator:
    def __init__(self, conf=None):
        serialization_data_path = 'serialization_data'
        self.path = serialization_data_path
        self.conf = dict() if conf is None else conf
        self.params = conf.get('training_params', dict())
        self.type = conf.get('type', None)
        self.answers = self.load(self.type + '_answers_v3.pkl')
        self._vocab = None
        self._reverse_vocab = None
        self._eval_sets = None

    ##### Resources #####

    def load(self, name):
        return pickle.load(open(os.path.join(self.path, name), 'rb'))

    def vocab(self):
        if self._vocab is None:
            self._vocab = self.load(self.type + '_vocabulary_v3.pkl')
        return self._vocab

    def reverse_vocab(self):
        if self._reverse_vocab is None:
            vocab = self.vocab()
            self._reverse_vocab = dict((v.lower(), k) for k, v in vocab.items())
        return self._reverse_vocab

    ##### Loading / saving #####

    def save_epoch(self, model, epoch):
        if not os.path.exists('models/'):
            os.makedirs('models/')
        model.save_weights('models/' + self.type + '_weights_epoch_%d_v3.h5' % epoch, overwrite=True)

    def load_epoch(self, model, epoch):
        assert os.path.exists('models/' + self.type + '_weights_epoch_%d_v3.h5' % epoch), 'Weights at epoch %d not found' % epoch
        model.load_weights('models/' + self.type + '_weights_epoch_%d_v3.h5' % epoch)

    ##### Converting / reverting #####

    def convert(self, words):
        rvocab = self.reverse_vocab()
        if type(words) == str:
            words = words.strip().lower().split(' ')
        return [rvocab.get(w, 0) for w in words]

    def revert(self, indices):
        vocab = self.vocab()
        return [vocab.get(i, 'X') for i in indices]

    ##### Padding #####

    def padq(self, data):
        return self.pad(data, self.conf.get('question_len', None))

    def pade(self, data):
        return self.pad(data, self.conf.get('entity_len', None))

    def padd(self, data):
        return self.pad(data, self.conf.get('des_len', None))

    def pad(self, data, len=None):
        from keras.preprocessing.sequence import pad_sequences
        return pad_sequences(data, maxlen=len, padding='post', truncating='post', value=0)

    ##### Training #####

    def print_time(self):
        print(strftime('%Y-%m-%d %H:%M:%S :: ', gmtime()), end='')

    def train(self, model):
        eval_every = self.params.get('eval_every', None)
        save_every = self.params.get('save_every', None)
        batch_size = self.params.get('batch_size', 128)
        nb_epoch = self.params.get('nb_epoch', 10)
        split = self.params.get('validation_split', 0)

        training_set = self.load(self.type + '_train_v3.pkl')

        questions = list()
        good_entities = list()
        good_descriptions = list()
        bad_answer_candidates = list()

        for q in training_set:
            questions += [q['question']] * len(q['good_answers'])
            good_entities += [self.answers[i]['entity'] for i in q['good_answers']]
            good_descriptions += [self.answers[i]['des'] for i in q['good_answers']]
            bad_answer_candidates += [q['bad_answers']] * len(q['good_answers'])

        questions = self.padq(questions)
        good_entities = self.pade(good_entities)
        good_descriptions = self.padd(good_descriptions)
        print("Questions: ", len(questions))
        print("Good answers: ", len(good_entities))

        val_loss = {'loss': 1., 'epoch': 0}

        for i in range(1, nb_epoch+1):
            # bad_answers = np.roll(good_answers, random.randint(10, len(questions) - 10))
            # bad_answers = good_answers.copy()
            # random.shuffle(bad_answers)
            # bad_answers = self.pada(random.sample(self.answers.values(), len(good_answers)))
            neg_samples = [random.choice(candidate) for candidate in bad_answer_candidates]
            bad_entities = self.pade([self.answers[neg_sample]['entity'] for neg_sample in neg_samples])
            bad_descriptions = self.padd([self.answers[neg_sample]['des'] for neg_sample in neg_samples])
            # print("Bad answers: ", len(bad_entity_candidates))

            # shuffle questions
            # zipped = zip(questions, good_answers)
            # random.shuffle(zipped)
            # questions[:], good_answers[:] = zip(*zipped)

            print('Epoch %d :: ' % i, end='')
            self.print_time()
            model.fit([good_entities, bad_entities, questions, good_descriptions, bad_descriptions],
                      nb_epoch=1, batch_size=batch_size, validation_split=split)

            # if hist.history['val_loss'][0] < val_loss['loss']:
            #     val_loss = {'loss': hist.history['val_loss'][0], 'epoch': i}
            # print('Best: Loss = {}, Epoch = {}'.format(val_loss['loss'], val_loss['epoch']))

            if eval_every is not None and i % eval_every == 0:
                self.get_map(model)

            if save_every is not None and i % save_every == 0:
                self.save_epoch(model, i)

    ##### Evaluation #####
    def get_map(self, model, evaluate_all=False):
        top1s = list()
        top5s = list()
        top10s = list()
        maps = list()
        for name, data in self.eval_sets().items():
            if evaluate_all:
                self.print_time()
                print('----- %s -----' % name)

            random.shuffle(data)

            if not evaluate_all and 'n_eval' in self.params:
                data = data[:self.params['n_eval']]

            ap, h1, h5, h10 = 0, 0, 0, 0

            for i, d in enumerate(data):
                if evaluate_all:
                    self.prog_bar(i, len(data))

                indices = d['good_answers'] + d['bad_answers']
                entities = self.pade([self.answers[index]['entity'] for index in indices])
                descriptions = self.padd([self.answers[index]['des'] for index in indices])
                question = self.padq([d['question']] * len(indices))

                n_good = len(d['good_answers'])
                sims = model.predict([entities, question, descriptions], batch_size=len(indices)).flatten()
                r = rankdata(sims, method='ordinal')

                target_rank = np.asarray(r[:n_good])
                num_candidate = len(sims)
                ground_rank = num_candidate - target_rank + 1
                ground_rank.sort()
                one_ap = 0
                for rank in xrange(n_good):
                    one_ap += (rank + 1) / float(ground_rank[rank])
                one_ap /= n_good

                ap += one_ap
                h1 += 1 if np.argmax(sims) < n_good else 0
                h5 += 1 if set(list(ground_rank - 1)).intersection(set(range(5))) else 0
                h10 += 1 if set(list(ground_rank - 1)).intersection(set(range(10))) else 0

                # max_r = np.argmax(sims)
                # max_n = np.argmax(sims[:n_good])

                # print(' '.join(self.revert(d['question'])))
                # print(' '.join(self.revert(self.answers[indices[max_r]])))
                # print(' '.join(self.revert(self.answers[indices[max_n]])))

                # c_1 += 1 if max_r == max_n else 0
                # c_2 += 1 / float(r[max_r] - r[max_n] + 1)

            top1 = h1 / float(len(data))
            top5 = h5 / float(len(data))
            top10 = h10 / float(len(data))
            mean_ap = ap / float(len(data))

            del data

            if evaluate_all:
                print('Top-1 Precision: %f' % top1)
                print('Hit@5 Precision: %f' % top5)
                print('Hit@10 Precision: %f' % top10)
                print('MAP: %f' % mean_ap)

            top1s.append(top1)
            top5s.append(top5)
            top10s.append(top10)
            maps.append(mean_ap)

        # rerun the evaluation if above some threshold
        if not evaluate_all:
            print('Top-1 Precision: {}'.format(top1s))
            print('Hit@5 Precision: {}'.format(top5s))
            print('Hit@10 Precision: {}'.format(top10s))
            print('MAP: {}'.format(maps))
            evaluate_all_threshold = self.params.get('evaluate_all_threshold', dict())
            evaluate_mode = evaluate_all_threshold.get('mode', 'all')
            map_threshold = evaluate_all_threshold.get('map', 1)
            top1_threshold = evaluate_all_threshold.get('top1', 1)

            if evaluate_mode == 'any':
                evaluate_all = evaluate_all or any([x >= top1_threshold for x in top1s])
                evaluate_all = evaluate_all or any([x >= map_threshold for x in maps])
            else:
                evaluate_all = evaluate_all or all([x >= top1_threshold for x in top1s])
                evaluate_all = evaluate_all or all([x >= map_threshold for x in maps])

            if evaluate_all:
                return self.get_map(model, evaluate_all=True)

        return top1s, top5s, top10s, maps

    def prog_bar(self, so_far, total, n_bars=20):
        n_complete = int(so_far * n_bars / total)
        if n_complete >= n_bars - 1:
            print('\r[' + '=' * n_bars + ']', end='')
        else:
            s = '\r[' + '=' * (n_complete - 1) + '>' + '.' * (n_bars - n_complete) + ']'
            print(s, end='')

    def eval_sets(self):
        if self._eval_sets is None:
            self._eval_sets = dict([(s, self.load(s)) for s in [self.type + '_test_v3.pkl']])
        return self._eval_sets

    def max_sim(self, left_list, right_list, w2v_model):
        max_cos = -1
        for i in range(len(left_list)):
            for j in range(len(right_list)):
                try:
                    local_sim = w2v_model.similarity(left_list[i], right_list[j])
                    max_cos = local_sim if local_sim > max_cos else max_cos
                except KeyError:
                    pass
        return max_cos

    def make_submit(self, model, dev, submit_file):
        lines = dev.readlines()
        target_lines = list()
        data = self.eval_sets().values()[0]
        for i, d in enumerate(data):
            terms = lines[i].strip('\n').split('\t')
            query = terms[0]
            entities = terms[1:]
            num_candidate = len(entities)
            index_entities = xrange(num_candidate)

            indices = d['answers']
            answers = self.pada([self.answers[i] for i in indices])
            question = self.padq([d['question']] * len(indices))

            sims = model.predict([question, answers], batch_size=100).flatten()
            print(len(sims))
            r = rankdata(sims, method='ordinal')
            index_candidates = nlargest(num_candidate, index_entities, key=lambda j: r[j])
            for index_candidate in index_candidates:
                query = query + '\t' + entities[index_candidate]
            query += '\n'
            target_lines.append(query)
        submit_file.writelines(target_lines)

        del data

    def make_submit_v2(self, model, dev, submit_file, w2v_model):
        lines = dev.readlines()
        target_lines = list()
        data = self.eval_sets().values()[0]
        for i, d in enumerate(data):
            terms = lines[i].strip('\n').split('\t')
            query = terms[0]
            query_word_list = [x for x in jieba.cut(query)]
            entities = terms[1:]
            num_candidate = len(entities)
            sim_candidate = list()

            for entity in entities:
                sim_candidate.append(self.max_sim(query_word_list, [y for y in jieba.cut(entity)], w2v_model))

            sim_candidate = np.asarray(sim_candidate)

            index_entities = xrange(num_candidate)

            indices = d['answers']
            answers = self.pada([self.answers[i] for i in indices])
            question = self.padq([d['question']] * len(indices))

            sims = model.predict([question, answers], batch_size=100).flatten()
            sims = sim_candidate * 0.6 + sims * 0.4
            print(len(sims))
            r = rankdata(sims, method='ordinal')
            index_candidates = nlargest(num_candidate, index_entities, key=lambda j: r[j])
            for index_candidate in index_candidates:
                query = query + '\t' + entities[index_candidate]
            query += '\n'
            target_lines.append(query)
        submit_file.writelines(target_lines)

        del data

if __name__ == '__main__':
    conf = {
        'type': 'tvShow',
        'question_len': 8,
        'entity_len': 4,
        'des_len': 2,
        'n_words': 9843,  # len(vocabulary) + 1
        'margin': 0.02,

        'training_params': {
            'save_every': 1000,
            'eval_every': 10,
            'batch_size': 32,
            'nb_epoch': 3000,
            'validation_split': 0,
            'optimizer': 'adam',
            # 'optimizer': Adam(clip_norm=0.1),
            # 'n_eval': 100,

            'evaluate_all_threshold': {
                'mode': 'all',
                'top1': 0.5,
            },
        },

        'model_params': {
            'n_embed_dims': 300,
            'n_hidden': 200,

            # convolution
            'nb_filters': 1000, # * 4
            'conv_activation': 'relu',

            # recurrent
            'n_lstm_dims': 141, # * 2

            'initial_embed_weights': np.load('embeddings/tvShow_300_dim_v3.embeddings'),
        },

        'similarity_params': {
            'mode': 'cosine',
            'gamma': 1,
            'c': 1,
            'd': 2,
        }
    }

    evaluator = Evaluator(conf)

    ##### Define model ######
    model = EmbeddingModel(conf)
    optimizer = conf.get('training_params', dict()).get('optimizer', 'adam')
    model.compile(optimizer=optimizer)

    import numpy as np

    # save embedding layer
    # evaluator.load_epoch(model, 33)
    # embedding_layer = model.prediction_model.layers[2].layers[2]
    # evaluator.load_epoch(model, 100)
    # evaluator.train(model)
    # weights = embedding_layer.get_weights()[0]
    # np.save(open('models/embedding_1000_dim.h5', 'wb'), weights)

    # train the model
    evaluator.train(model)
    # evaluate mrr for a particular epoch
    evaluator.load_epoch(model, 3000)
    evaluator.get_map(model, evaluate_all=True)

    # celebrity
    # evaluator.load_epoch(model, 3000)
    # dev_set = codecs.open('data/celebrity.TESTSET.txt', 'rb', 'gb18030')
    # submit = codecs.open('data/celebrity-final-v1.txt', 'wb', 'gb18030')
    # wor2vec_model = Word2Vec.load_word2vec_format('data/baike_vector.bin', binary=True)
    # evaluator.make_submit_v2(model, dev_set, submit, wor2vec_model)

    # movie
    # evaluator.load_epoch(model, 3000)
    # dev_set = codecs.open('data/movie.TESTSET.txt', 'rb', 'gb18030')
    # submit = codecs.open('data/movie-final-v1.txt', 'wb', 'gb18030')
    # evaluator.make_submit(model, dev_set, submit)

    # restaurant
    # evaluator.load_epoch(model, 3000)
    # dev_set = codecs.open('data/restaurant.TESTSET.txt', 'rb', 'gb18030')
    # submit = codecs.open('data/restaurant-final-v1.txt', 'wb', 'gb18030')
    # evaluator.make_submit(model, dev_set, submit)

    # tvShow
    # evaluator.load_epoch(model, 3000)
    # dev_set = codecs.open('data/tvShow.TESTSET.txt', 'rb', 'gb18030')
    # submit = codecs.open('data/tvShow-final-v1.txt', 'wb', 'gb18030')
    # evaluator.make_submit(model, dev_set, submit)


import tensorflow as tf
from keras.models import Model
from keras.layers import Dense, Input, concatenate
from keras.optimizers import Adam, RMSprop
from keras import backend as K
from keras.utils import to_categorical

class Actor:
    '''
    Policy Function Approximator - "actor" for SAC algorithm.

    Stochastic actor - augments classic reinforcement learning objective of maximising
    the expected sum of rewards with the expected entropy of the policy π.

    Network Input: state.
    Network Output: action under a stochastic policy.
    '''
    # Khev
    def __init__(self,input_dim, output_dim, lr, gamma, tau, alpha, clipnorm, clipnorm_val, verbose = False):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.model = self._make_network()
        self.target_model = self._make_network()    #we have target networks to stabilize learning.
        self.target_model.set_weights(self.model.get_weights())         #clone the network
        self.lr = lr  #learning rate for optimizer
        self.gamma = gamma
        self.tau = tau
        self.alpha = alpha
        self.clipnorm = clipnorm
        self.clipnorm_val = clipnorm_val
        self.verbose = verbose
        self.opt = self.optimizer()

    # JP
    def __init__(self, config, state_dim, action_dim, action_bound):
        '''
        :param state_dim: dimensionality of environment state space
        :param action_dim: dimensionality of action space
        :param action_bound: bound on action magnitude
        :param learning_rate: learning rate of the critic network
        :num_actor_vars: number of trainable variables in the actor to be given to the critic
        '''

        # Build neural network function approximator for the actor (policy)
        self.state, _, self.scaled_actions = self.build_actor_model(
            state_dim, action_dim, action_bound)

        # Define parameters for current (+stochastic noise) actor network
        # N.B. order here is important: this call must be before target network is set up.
        network_params = tf.compat.v1.trainable_variables()

        # Target Network
        self.target_state, _, self.target_scaled_actions = self.build_actor_model(
            state_dim, action_dim, action_bound)

        # Define parameters for target (deterministic) actor network
        target_network_params = tf.compat.v1.trainable_variables()[len(network_params):]

        # Mixing factor for periodicity of frozen target network parameter updates (for stability)
        tau = config['agent']['tau']

        # Op for periodically updating target network with online network
        # weights. Copies the parameters of the online network with a mixing factor tau.
        self.update_target_network_params_op = \
            [target_network_params[i].assign(tf.multiply(network_params[i], tau) +
                                             tf.multiply(target_network_params[i], 1.0 - tau))
                for i in range(len(target_network_params))]

        # This gradient will be provided by the critic network's feedback
        self.action_gradient = tf.compat.v1.placeholder(tf.float32, [None, action_dim])

        # Combine gradients.
        # -ve action gradient --> make actor follow action_gradients by gradient ascent.
        unnormalized_actor_gradients = tf.gradients(
            self.scaled_actions, network_params, -self.action_gradient)

        # Combine the gradients - divide by batch_size to account
        # for gradients being summed over the batch by tf.gradients
        actor_gradients = list(map(lambda x: tf.math.divide(
            x, config['train']['batch_size']), unnormalized_actor_gradients))

        # Optimization Op
        self.train_op = tf.compat.v1.train.AdamOptimizer(config['agent']['actor_lr']).\
            apply_gradients(zip(actor_gradients, network_params))

        self.num_trainable_vars = len(network_params) + len(target_network_params)

    # Khev
    def _make_network(self):

        #Hyperpars chosen according to the paper (see README)
        S = Input(shape=(self.input_dim,))
        x = Dense(256, activation = 'relu')(S)
        x = Dense(256, activation = 'relu')(x)
        out = Dense(self.output_dim, activation = 'softmax')(x)
        model = Model(inputs = S, outputs = out)
        return model

    # JP
    def build_actor_model(self, state_dim, action_dim, action_bound, is_training=True):
        '''
        Builds neural network approximator for SAC actor model.
        Architecture built according to Appendix D https://arxiv.org/pdf/1801.01290.pdf

        :param state_dim: dimensionality of environment state space
        :param action_dim: dimensionality of action space
        :param action_bound: bound on action magnitude
        :param is_training: mode of execution, for BN

        :return state: pass state on to critic
        :return action: action only bounded by tanh
        :return scaled_action: action bounded by tanh and [-action_bound, +action_bound]
        '''
        with tf.compat.v1.variable_scope('build_actor'):
            # Keras input --> placeholder of [None, state_dim]
            state = tf.keras.Input(shape=(state_dim,))

            with tf.compat.v1.variable_scope('actor_h1'):
                net = tf.keras.layers.Dense(units=256, activation=None)(state)
                net = tf.keras.layers.BatchNormalization()(net, training=is_training)
                net = tf.keras.activations.relu(net)

            with tf.compat.v1.variable_scope('actor_h2'):
                net = tf.keras.layers.Dense(units=256, activation=None)(net)
                net = tf.keras.layers.BatchNormalization()(net, training=is_training)
                net = tf.keras.activations.relu(net)

            with tf.compat.v1.variable_scope('actor_out'):
                actions = tf.keras.layers.Dense(
                    units=action_dim, activation=tf.keras.activations.softmax)(net)

                scaled_actions = tf.multiply(actions, action_bound)

        return state, actions, scaled_actions

    # Khev init - self.opt = self.optimizer()
    def learn(self,S,A,Q,V):
        """
        S = batch of states
        A = batch of actions
        Q = batch of Q-vals
        V = batch of values
        """

        self.opt([S,A,Q,V])

    def optimizer(self):

        """ grad_loss =  \grad_{theta} \log(\pi(a_t,s_t))* (  \alpha \log(\pi(s_t, a_t) ) - Q(s_t, a_t)  + V(s_t)  )

            where
            s_t = state at time t (t is discrete)
            a_t = action
            \pi(s,a) = the policy (actor)
            Q(s,a) = the Q-value, from the critic
            V(s) = the value  (gets its own neural net -- see paper)
            \alpha = parameter controlling the strength of the entropy

        """
        #Inputs
        S_pl = self.model.input
        A_pl = K.placeholder(shape=(None, self.output_dim))  #onehot
        Q_pl = K.placeholder(shape=(None, 1))
        V_pl = K.placeholder(shape=(None, 1))

        #Find terms in bracket
        pi_vec = self.model.output
        pi = K.sum(pi_vec * A_pl, axis=1)    # get \pi(s_t, a_t) -- prob for specific action
        entropy = self.alpha * K.log(pi)
        temp = entropy - K.transpose(Q_pl) + K.transpose(V_pl)  #this is a row vec
        temp = K.transpose(temp)      #turn it into col vec

        # TODO (JP): you could give a lot of this to the agent class train method.
        #Find grad log(pi)
        pi_pl = self.model.output
        pars = self.model.trainable_weights
        grads = tf.gradients( K.log(pi_pl), pars, temp)   #scalar multiply by temp

        #Clip gradients
        if self.clipnorm == True:
            grads = tf.clip_by_global_norm(grads, self.clipnorm_val)[0]

        #Do learning
        #To get keras to apply updates given a custom gradients (i.e. run the above line) I had to alter the source
        #Code. It was easy to do. See line X in the get_updates function.
        opt = Adam(self.lr)
        loss = grads  #placeholder, I won't use it
        updates = opt.get_updates(loss = grads, params = pars, grads = grads)

        #This function will apply updates when called
        func = K.function(inputs = [S_pl, A_pl, Q_pl, V_pl], outputs = [], updates = updates)
        return func


import datetime
import dateutil.parser
import xml.etree.ElementTree as ET

class TravelDecorator:
    STATUS_CLASSES = {
        'NIET-MOGELIJK': 'impossible',
        'VOLGENS-PLAN':  'planned',
        'NIEUW':         'new',
        'VERTRAAGD':     'delayed',
        'NIET-OPTIMAAL': 'not-optimal'
    }

    def __init__(self, xml):
        self.xml = xml
        self.root = ET.fromstring(self.xml)

    def decorate(self):
        results = []

        for travel_possibility in self.root.findall('.//ReisMogelijkheid'):
            actual_travel_time      = travel_possibility.find('ActueleReisTijd')
            actual_departure_time   = travel_possibility.find('ActueleVertrekTijd')
            actual_arrival_time     = travel_possibility.find('ActueleAankomstTijd')
            status                  = travel_possibility.find('Status')
            delay                   = travel_possibility.find('VertrekVertraging')

            possibility = {
                'actual_travel_time':     actual_travel_time.text,
                'actual_departure_time':  self.get_time(actual_departure_time.text),
                'actual_arrival_time':    self.get_time(actual_arrival_time.text),
                'status':                 status.text,
                'status_class':           self.STATUS_CLASSES[status.text.strip()]
             }

            if delay is not None:
                possibility['delay'] = delay.text

            results.append(possibility)

        return results

    def get_time(self, time):
        return str(dateutil.parser.parse(time).time())


from sqlalchemy.orm import class_mapper, ColumnProperty
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base
import models as mf
from database import Base as Base_from
from database import db_session as db_session_from
from sqlalchemy.sql.schema import Table
import settings

engine_to = create_engine(f"sqlite:///{settings.work_dir}/.report/db.db")
db_session_to = scoped_session(sessionmaker(autocommit=False,
                                         autoflush=False,
                                         bind=engine_to))

def get_list_classes(module):
    """Retorna a lista de models de um modulo"""
    list_ = []
    for item in dir(module):
        obj = getattr(module, item)
        if hasattr(obj, "__tablename__"):
            list_.append(item)
    return list_

def attribute_names(cls):
    """Retorna os atributos da classe que são do tipo Column"""
    return [prop.key for prop in class_mapper(cls).iterate_properties
        if isinstance(prop, ColumnProperty)]

def run_migrate():
    Base_from.metadata.create_all(bind=engine_to)

    for class_name in get_list_classes(mf):
        class_ = getattr(mf, class_name)
        props = attribute_names(class_)
        items = db_session_from.query(class_).all()
        for item in items:
            new_item = class_()
            for prop in props:
                setattr(new_item, prop, getattr(item, prop))
            db_session_to.add(new_item)
    db_session_to.commit()

    #Copiar dados das tabelas ponte
    conn_from = db_session_from.connection()
    conn_to = db_session_to.connection()
    for item in dir(mf):
        obj = getattr(mf, item)
        if obj.__class__ == Table:
            stm = obj.select()
            res = conn_from.execute(stm).fetchall()
            for item in res:
                cols = [c.name for c in obj.columns]
                stm = obj.insert().values(dict(zip(cols, item)))
                conn_to.execute(stm)
    db_session_to.commit()


# -*- coding: utf-8 -*-
#
# Copyright © 2013 Red Hat, Inc.
#
# This software is licensed to you under the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the License
# (GPLv2) or (at your option) any later version.
# There is NO WARRANTY for this software, express or implied, including the
# implied warranties of MERCHANTABILITY, NON-INFRINGEMENT, or FITNESS FOR A
# PARTICULAR PURPOSE.
# You should have received a copy of GPLv2 along with this software; if not,
# see http://www.gnu.org/licenses/old-licenses/gpl-2.0.txt

from copy import deepcopy
import gdbm
import gzip
import hashlib
import logging
import lzma
import os
from urlparse import urljoin
from xml.etree import ElementTree
from xml.etree.cElementTree import iterparse

from nectar.listener import AggregatingEventListener
from nectar.request import DownloadRequest
from pulp_rpm.plugins.importers.yum import utils

from pulp_rpm.plugins.importers.yum.repomd import filelists, nectar_factory, other, packages
from pulp_rpm.plugins.importers.yum.repomd.packages import package_list_generator

_LOGGER = logging.getLogger(__name__)

# repomd.xml element tags ------------------------------------------------------

REPOMD_FILE_NAME = 'repomd.xml'
REPOMD_URL_RELATIVE_PATH = 'repodata/%s' % REPOMD_FILE_NAME

SPEC_URL = 'http://linux.duke.edu/metadata/repo'

REVISION_TAG = '{%s}revision' % SPEC_URL

DATA_TAG = '{%s}data' % SPEC_URL

LOCATION_TAG = '{%s}location' % SPEC_URL
CHECKSUM_TAG = '{%s}checksum' % SPEC_URL
SIZE_TAG = '{%s}size' % SPEC_URL
TIMESTAMP_TAG = '{%s}timestamp' % SPEC_URL
OPEN_CHECKSUM_TAG = '{%s}open-checksum' % SPEC_URL
OPEN_SIZE_TAG = '{%s}open-size' % SPEC_URL

# metadata file information skeleton -------------------------------------------

FILE_INFO_SKEL = {'name': None,
                  'relative_path': None,
                  'checksum': {'algorithm': None, 'hex_digest': None},
                  'size': None,
                  'timestamp': None,
                  'open_checksum': {'algorithm': None, 'hex_digest': None},
                  'open_size': None}

# metadata files downloader, parser, and validator -----------------------------

class MetadataFiles(object):
    """
    Stateful downloader, parser, and validator of the metadata files of a Yum
    repository.

    Given a Yum repository URL, this class presents a clean work flow for
    fetching and validating the metadata files of that repo. The workflow is as
    follows:

    1. instantiate MetadataFiles instance with repository URL
    2. call `download_repomd` method
    3. call `parse_repomd` method
    4. call `download_metadata_files` method
    5. optionally call `validate_metadata_files` method

    If all goes well, the instance will have have populated its `metadata` dict
    with `key` -> file path information

    Keys of interest:

     * `primary`: path the primary.xml file containing the metadata of all packages in the repository
     * `filelists`: path the filelists.xml file containing the files provided by all of the packages in the repository
     * `other`
     * `group`
     * `group_gz`
     * `updateinfo`

    :ivar repo_url: Yum repository URL
    :ivar dst_dir: Directory to store downloaded metadata files in
    :ivar event_listener: nectar.listener.DownloadEventListener instance
    :ivar downloader: nectar.downloaders.base.DownloaderBackend instance
    :ivar revision: revision number of the metadata, set during the `parse_repomd` call
    :ivar metadata: dictionary of the main metadata type keys to the corresponding file paths
    """

    # these are metadata file types listed in "repomd" that we do not want to
    # store as units
    KNOWN_TYPES = set(['group', 'group_gz',
                   'filelists', 'filelists_db',
                   'other', 'other_db',
                   'primary', 'primary_db',
                   'updateinfo', 'updateinfo_db'])

    def __init__(self, repo_url, dst_dir, nectar_config):
        """
        :param repo_url:        URL for the base of a yum repository
        :type  repo_url:        basestring
        :param dst_dir:         full path to a destination to which files
                                should be downloaded
        :type  dst_dir:         basestring
        :param nectar_config:   download config for nectar
        :type  nectar_config:   nectar.config.DownloaderConfig
        """
        super(MetadataFiles, self).__init__()
        self.repo_url = repo_url
        self.dst_dir = dst_dir
        self.event_listener = AggregatingEventListener()

        self.downloader = nectar_factory.create_downloader(repo_url, nectar_config,
                                                           self.event_listener)

        self.revision = None
        self.metadata = {}
        self.dbs = {}

    def download_repomd(self):
        """
        Download the main repomd.xml file.
        """
        repomd_dst_path = os.path.join(self.dst_dir, REPOMD_FILE_NAME)
        repomd_url = urljoin(self.repo_url, REPOMD_URL_RELATIVE_PATH)
        repomd_request = DownloadRequest(repomd_url, repomd_dst_path)
        self.downloader.download([repomd_request])
        if self.event_listener.failed_reports:
            error_report = self.event_listener.failed_reports[0]
            raise IOError(error_report.error_msg)

    # TODO (jconnonr 2013-03-07) add a method to validate/verify the repomd.xml file

    def parse_repomd(self):
        """
        Parse the downloaded repomd.xml file and populate the metadata dictionary.
        """
        repomd_file_path = os.path.join(self.dst_dir, REPOMD_FILE_NAME)

        if not os.access(repomd_file_path, os.F_OK | os.R_OK):
            raise RuntimeError('%s has not been downloaded' % REPOMD_FILE_NAME)

        parser = iterparse(repomd_file_path, events=('start', 'end'))
        xml_iterator = iter(parser)

        # get a hold of the root element so that we can clear it
        # this prevents the entire parsed document from building up in memory
        try:
            root_element = xml_iterator.next()[1]
        except SyntaxError:
            raise ValueError('could not parse repo metadata')

        for event, element in xml_iterator:
            if event != 'end':
                continue

            root_element.clear()

            if element.tag == REVISION_TAG:
                self.revision = element.text

            if element.tag == DATA_TAG:
                file_info = process_repomd_data_element(element)
                self.metadata[file_info['name']] = file_info

    def download_metadata_files(self):
        """
        Download the remaining metadata files.
        """
        if not self.metadata:
            raise RuntimeError('%s has not been parsed' % REPOMD_FILE_NAME)

        download_request_list = []

        for file_name, file_info in self.metadata.iteritems():
            # we don't care about the sqlite files
            if file_name.endswith('_db') and file_name in self.KNOWN_TYPES:
                continue
            url = urljoin(self.repo_url, file_info['relative_path'])
            dst = os.path.join(self.dst_dir, file_info['relative_path'].rsplit('/', 1)[-1])

            file_info['local_path'] = dst

            request = DownloadRequest(url, dst)
            download_request_list.append(request)

        self.downloader.download(download_request_list)

    def verify_metadata_files(self):
        """
        Optionally verify the metadata files using both reported size and checksum.
        """
        # TODO: vet this method and determine if it should be used
        for md in self.metadata.values():
            if 'local_path' not in md:
                raise RuntimeError('%s has not been downloaded' % md['relative_path'].rsplit('/', 1)[-1])

            if md['size'] is None:
                raise RuntimeError('%s cannot be verified, no file size' % md['local_path'])

            local_file_size = os.path.getsize(md['local_path'])
            # prevents the rounding errors better than: md['size'] * 1024
            if local_file_size / 1024 != md['size']:
                raise RuntimeError('%s failed verification, file size mismatch' % md['local_path'])

            if md['checksum']['algorithm'] is None:
                raise RuntimeError('%s cannot be verified, no checksum' % md['local_path'])

            hash_constructor = getattr(hashlib, md['checksum']['algorithm'], None)
            if hash_constructor is None:
                raise RuntimeError('%s failed verification, unsupported hash algorithm: %s' % (md['local_path'], md['checksum']['algorithm']))

            hash_obj = hash_constructor()
            with open(md['local_path'], 'rb') as file_handle:
                hash_obj.update(file_handle.read())
            if hash_obj.hexdigest() != md['checksum']['hex_digest']:
                raise RuntimeError('%s failed verification, checksum mismatch' % md['local_path'])

    def get_metadata_file_handle(self, name):
        """
        Given a standard name for a metadata file, as appears in a repomd.xml file
        as a "data" element's "type", return an open file handle in read mode for
        that file.

        :param name:    name of a metadata file as would be found in the
                        repomd.xml file as a "type" attribute of a "data" block.
        :type  name:    basestring

        :return: open file handle to a file containing XML
        :rtype:  file
        """
        try:
            file_path = self.metadata[name]['local_path']
        except KeyError:
            return

        if file_path.endswith('.gz'):
            file_handle = gzip.open(file_path, 'r')
        elif file_path.endswith('.xz'):
            file_handle = lzma.LZMAFile(file_path, 'r')
        else:
            file_handle = open(file_path, 'r')
        return file_handle

    def get_group_file_handle(self):
        """
        return an open file handle from which the group XML can be read.

        :return:    open file handle to the file containing group XML
        :rtype:     file
        """
        group_file_handle = self.get_metadata_file_handle('group_gz')
        if group_file_handle is None:
            group_file_handle = self.get_metadata_file_handle('group')
        return group_file_handle

    def generate_dbs(self):
        """
        For repo data files that contain data we need to access later for each
        unit in the repo, generate a local db file that gives us quick read
        access to each unit's data.
        """
        for filename, tag, process_func in (
            (filelists.METADATA_FILE_NAME, filelists.PACKAGE_TAG, filelists.process_package_element),
            (other.METADATA_FILE_NAME, other.PACKAGE_TAG, other.process_package_element),
        ):
            xml_file_handle = self.get_metadata_file_handle(filename)
            try:
                generator = package_list_generator(xml_file_handle, tag)
                db_filename = os.path.join(self.dst_dir, '%s.db' % filename)
                # always a New file, and open with Fast writing mode.
                db_file_handle = gdbm.open(db_filename, 'nf')
                try:
                    for element in generator:
                        utils.strip_ns(element)
                        raw_xml = utils.element_to_raw_xml(element)
                        unit_key, _ = process_func(element)
                        db_key = self.generate_db_key(unit_key)
                        db_file_handle[db_key] = raw_xml
                    db_file_handle.sync()
                finally:
                    db_file_handle.close()
            finally:
                xml_file_handle.close()
            self.dbs[filename] = db_filename

    @staticmethod
    def generate_db_key(unit_key):
        """
        :param unit_key:    dictionary of key:value pairs that make a unique
                            entry in the given database.
        :type  unit_key:    dict

        :return:    a string that is a suitable unit key for the database
        :rtype:     basestring
        """
        unit_key = unit_key.copy()
        # clean out these entries if they exist, because they won't be in the
        # XML files we're indexing.
        unit_key.pop('checksum', None)
        unit_key.pop('checksumtype', None)
        sorted_key_names = sorted(unit_key.keys())
        return '::'.join('%s:%s' % (name, unit_key[name]) for name in sorted_key_names)

    def add_repodata(self, model):
        """
        Given a model, add the "repodata" attribute to it (which includes raw
        XML used for publishing), and add the "files" and "changelog" attributes
        based on data obtained in the raw XML snippets.

        :param model:   model instance to manipulate
        :type  model:   pulp_rpm.common.models.RPM
        """
        repodata = model.metadata.setdefault('repodata',{})
        db_key = self.generate_db_key(model.unit_key)
        for filename, metadata_key, process_func in (
            (filelists.METADATA_FILE_NAME, 'files', filelists.process_package_element),
            (other.METADATA_FILE_NAME, 'changelog', other.process_package_element)
        ):
            try:
                db_file = gdbm.open(self.dbs[filename], 'r')
                raw_xml = db_file[db_key]
            finally:
                db_file.close()
            repodata[filename] = raw_xml
            element = ElementTree.fromstring(raw_xml)
            unit_key, items = process_func(element)
            model.metadata[metadata_key] = items

        repodata['primary'] = model.raw_xml

# utilities --------------------------------------------------------------------

def process_repomd_data_element(data_element):
    """
    Process the data elements of the repomd.xml file.

    This returns a file information dictionary with the following keys:

     * `name`: name of the element
     * `relative_path`: the path of the metadata file, relative to the repository URL
     * `checksum`: dictionary of `algorithm` and `hex_digest` keys and values
     * `size`: size of the metadata file, in bytes
     * `timestamp`: unix timestamp of the file's creation, as a float
     * `open_checksum`: optional checksum dictionary of uncompressed metadata file
     * `open_size`: optional size of the uncompressed metadata file, in bytes

    :param data_element: XML data element parsed from the repomd.xml file
    :return: file_info dictionary
    :rtype: dict
    """

    file_info = deepcopy(FILE_INFO_SKEL)

    file_info['name'] = data_element.attrib['type']

    location_element = data_element.find(LOCATION_TAG)
    if location_element is not None:
        file_info['relative_path'] = location_element.attrib['href']

    checksum_element = data_element.find(CHECKSUM_TAG)
    if checksum_element is not None:
        file_info['checksum']['algorithm'] = checksum_element.attrib['type']
        file_info['checksum']['hex_digest'] = checksum_element.text

    size_element = data_element.find(SIZE_TAG)
    if size_element is not None:
        file_info['size'] = int(size_element.text)

    timestamp_element = data_element.find(TIMESTAMP_TAG)
    if timestamp_element is not None:
        file_info['timestamp'] = float(timestamp_element.text)

    open_checksum_element = data_element.find(OPEN_CHECKSUM_TAG)
    if open_checksum_element is not None:
        file_info['open_checksum']['algorithm'] = open_checksum_element.attrib['type']
        file_info['open_checksum']['hex_digest'] = open_checksum_element.text

    open_size_element = data_element.find(OPEN_SIZE_TAG)
    if open_size_element is not None:
        file_info['open_size'] = int(open_size_element.text)

    for child in data_element.getchildren():
        child.clear()
    data_element.clear()

    return file_info


# -*- coding: utf-8 -*-
import scrapy
from ..items import AmazonItem

class AmazonSpider(scrapy.Spider):
    name = 'Amazon'
    page_number = 2
    allowed_domains = ['amazon.com']
    start_urls = ['https://www.amazon.com/b?ie=UTF8&node=17143709011']

    def parse(self, response):
        items = AmazonItem()

        product_name = response.css('.s-access-title::text').extract()
        product_author = response.css('.a-color-secondary .a-text-normal').css('::text').extract()
        product_price = response.css('.a-price-whole::text').extract()
        product_imagelink = response.css('.cfMarker::attr(src)').extract()

        items['product_name'] = product_name
        items['product_author'] = product_author
        items['product_price'] = product_price
        items['product_imagelink'] = product_imagelink

        yield items

        next_page = 'https://www.amazon.com/s?i=specialty-aps&srs=17143709011&page= ' + str(AmazonSpider.page_number) + ' &qid=1575368837&ref=sr_pg_2'
        if AmazonSpider.page_number <= 10:
            AmazonSpider.page_number +=1
            yield response.follow(next_page, callback = self.parse)


import gym
import numpy as np
import matplotlib.pyplot as plt
from gym.envs.registration import register
import random as pr

env = gym.make('FrozenLake-v0') #미끄러지는 부분

Q = np.zeros([env.observation_space.n, env.action_space.n])

#learning parameter 설정
learning_rate = 0.85
dis = 0.99
num_episodes = 2000

rList=[]
for i in range(num_episodes):
    #Reset environment and 첫번째 관찰
    state = env.reset()
    rAll = 0
    done = False
    # decaying E-greedy
    e = 1. / ((i // 100) + 1)
    #Q-table learning algorithm
    while not done:
        #액션의 greedy 선택
        action = np.argmax(Q[state, :] + np.random.randn(1, env.action_space.n) / (i + 1))

        # S'의 상태와 보상
        new_state, reward, done,_ = env.step(action)

        #Update
        Q[state,action] = (1-learning_rate)*Q[state,action] + learning_rate*(reward + dis*np.max(Q[new_state,:]))

        rAll += reward
        state = new_state

    rList.append(rAll)

print("Success rate : "+str(sum(rList) / num_episodes))#승률
print("Final Q-Table Values")
print(Q)

plt.bar(range(len(rList)), rList, color="blue")
plt.show()


'''
Created on Feb 11, 2013

@package: ally base
@copyright: 2012 Sourcefabric o.p.s.
@license: http://www.gnu.org/licenses/gpl-3.0.txt
@author: Gabriel Nistor

Provides the attributes support.
'''

from .spec import IAttribute, AttrError, IResolver, ResolverError, \
    ContextMetaClass, LIST_UNAVAILABLE, LIST_UNUSED, CREATE_DEFINITION
from ally.design.processor.spec import LIST_CLASSES
from ally.support.util_spec import IGet, ISet
from ally.support.util_sys import locationStack
from collections import Iterable
from inspect import isclass
import logging
from ally.support.util import FlagSet

# --------------------------------------------------------------------

log = logging.getLogger(__name__)

# --------------------------------------------------------------------

def defines(*types, doc=None):
    '''
    Construct a defining attribute for the context. The defines attribute means that the context can provide a value
    for the attribute, but is not mandatory also whenever managing an attribute if this type is a good idea to check
    if there aren't already values provided.

    @param types: arguments[class]
        The types of the defined attribute.
    @param doc: string|None
        The documentation associated with the attribute.
    '''
    return Attribute(Specification(DEFINED, types, doc=doc))

def definesIf(*types, doc=None):
    '''
    Construct a defining attribute for the context. The defines attribute means that the context can provide a value
    for the attribute, but is not mandatory also whenever managing an attribute if this type is a good idea to check
    if there aren't already values provided. Whenever using this type of attributes always check if the context has them since
    if they are optional they might not event be populated if there is no definition for them, so always to a check like:
        MyContext.myAttribute in myInstance
    , otherwise you might get attribute error.

    @param types: arguments[class]
        The types of the defined attribute.
    @param doc: string|None
        The documentation associated with the attribute.
    '''
    return Attribute(Specification(DEFINED | OPTIONAL, types, doc=doc))

def optional(*types, doc=None):
    '''
    Construct an optional attribute for the context. The optional attribute means that the context is valid even if
    there is no value for the attribute. Whenever using this type of attributes always check if the context has them since
    if they are optional they might not event be populated if there is no definition for them, so always to a check like:
        MyContext.myAttribute in myInstance
    , otherwise you might get attribute error.

    @param types: arguments[class]
        The types of the optional attribute, the attribute value can be any one of the provided attributes.
    @param doc: string|None
        The documentation associated with the attribute.
    '''
    return Attribute(Specification(OPTIONAL, types, doc=doc))

def requires(*types, doc=None):
    '''
    Construct a required attribute for the context. The requires attribute means that the context is valid only if
    there is a value for the attribute.

    @param types: arguments[class]
        The types of the required attribute, the attribute value can be any one of the provided attributes.
    @param doc: string|None
        The documentation associated with the attribute.
    '''
    return Attribute(Specification(REQUIRED, types, doc=doc))

def attribute(*types, doc=None):
    '''
    Construct a simple attribute for the context that will be used by the processor that define it and is also available
    as a defined attribute.

    @param types: arguments[class]
        The types of the defined attribute.
    @param doc: string|None
        The documentation associated with the attribute.
    '''
    return Attribute(Specification(DEFINED | USED, types, doc=doc))

# --------------------------------------------------------------------

DEFINED = 1 << 1
# Status flag for defined attributes.
REQUIRED = 1 << 2
# Status flag for required attributes.
OPTIONAL = 1 << 3
# Status flag for optional attributes.
USED = 1 << 4
# Status flag for attributes that should be considered always used.

# --------------------------------------------------------------------

class Specification:
    '''
    Provides attribute specifications.
    '''

    def __init__(self, status, types, definedIn=None, doc=None, defined=None):
        '''
        Construct the attribute specification.

        @param status: integer
            The status of the attribute specification.
        @param types: Iterable(class)
            The type(s) for the attribute specification.
        @param definedIn: class|None
            The class that defines the specification, this is just in case the specification was created based on a definer.
        @param doc: string|None
            The documentation associated with the attribute specification.
        @param defined: Iterable(class)|None
            The defined classes.
        '''
        assert isinstance(status, int), 'Invalid status %s' % status
        types = tuple(reduce(types))
        assert types, 'At least a type is required'
        assert definedIn is None or isclass(definedIn), 'Invalid defined in %s' % definedIn

        if defined is None:
            if status & DEFINED: defined = types
            else: defined = ()
        else: assert isinstance(defined, Iterable), 'Invalid defined classes %s' % defined

        self.status = status
        self.types = types
        self.definedIn = definedIn
        self.doc = doc
        self.defined = frozenset(defined)

        self.usedIn = {}

    def __str__(self):
        status = []
        if self.status & DEFINED: status.append('DEFINES')
        if self.status & REQUIRED: status.append('REQUIRED')
        if self.status & OPTIONAL: status.append('OPTIONAL')

        return ''.join(('|'.join(status), '[', ','.join(t.__name__ for t in self.types), ']'))

class Resolver(IResolver):
    '''
    Implementation for a @see: IResolver that manages contexts with @see: Attribute.
    '''

    def __init__(self, context):
        '''
        Construct the attribute resolver.

        @param context: ContextMetaClass|dictionary{string: Specification}
            The context class to construct the resolver based on, or the specifications dictionary.
        '''
        if isinstance(context, ContextMetaClass):
            assert isinstance(context, ContextMetaClass)
            specifications = {}
            for name, attribute in context.__attributes__.items():
                assert isinstance(name, str), 'Invalid name %s' % name
                assert isinstance(attribute, Attribute), 'Invalid attribute %s' % attribute
                specifications[name] = attribute.specification

        else:
            assert isinstance(context, dict), 'Invalid context %s' % context
            if __debug__:
                for name, spec in context.items():
                    assert isinstance(name, str), 'Invalid name %s' % name
                    assert isinstance(spec, Specification), 'Invalid specification %s' % spec
            specifications = context

        self.specifications = specifications

    def copy(self, names=None):
        '''
        @see: IResolver.copy
        '''
        specifications = {}
        if names is None:
            specifications.update(self.specifications)

        else:
            assert isinstance(names, Iterable), 'Invalid names %s' % names
            for name in names:
                assert isinstance(name, str), 'Invalid name %s' % name
                spec = self.specifications.get(name)
                if spec: specifications[name] = spec

        return self.__class__(specifications)

    def merge(self, other, isFirst=True):
        '''
        @see: IResolver.merge
        '''
        assert isinstance(other, IResolver), 'Invalid other resolver %s' % other
        assert isinstance(isFirst, bool), 'Invalid is first flag %s' % isFirst

        if self is other: return self

        if not issubclass(self.__class__, other.__class__):
            if isFirst: return other.merge(self, False)
            raise ResolverError('Cannot merge %s with %s' % (self, other))
        assert isinstance(other, Resolver), 'Invalid other resolver %s' % other

        return self.__class__(self.mergeSpecifications(self.specifications, other.specifications))

    def solve(self, other):
        '''
        @see: IResolver.solve
        '''
        assert isinstance(other, IResolver), 'Invalid other resolver %s' % other
        if self is other: return self

        if not issubclass(self.__class__, other.__class__): return other.solve(self)
        assert isinstance(other, Resolver), 'Invalid other resolver %s' % other

        return self.__class__(self.solveSpecifications(self.specifications, other.specifications))

    def list(self, *flags):
        '''
        @see: IResolver.list
        '''
        flags, attributes = FlagSet(flags), {}

        listed = False
        if flags.checkOnce(LIST_UNAVAILABLE):
            listed = True
            for name, spec in self.specifications.items():
                assert isinstance(spec, Specification), 'Invalid specification %s' % spec
                if spec.status == REQUIRED: attributes[name] = None

        listClasses = flags.checkOnce(LIST_CLASSES)

        if flags.checkOnce(LIST_UNUSED):
            listed = True
            for name, spec in self.specifications.items():
                assert isinstance(spec, Specification), 'Invalid specification %s' % spec
                if spec.status & USED: continue
                if spec.status & OPTIONAL: continue
                for status in spec.usedIn.values():
                    if not status & DEFINED:
                        if spec.definedIn is not None: attributes[name] = (spec.definedIn,) if listClasses else None
                        break
                else: attributes[name] = spec.usedIn.keys() if listClasses else None
            listClasses = False

        if not listed:
            for name in self.specifications: attributes[name] = None

        if listClasses:
            for name in attributes: attributes[name] = self.specifications[name].usedIn.keys()

        assert not flags, 'Unknown flags: %s' % ', '.join(flags)
        return attributes

    def create(self, *flags):
        '''
        @see: IResolver.create
        '''
        flags = FlagSet(flags)

        attributes = {}
        if flags.checkOnce(CREATE_DEFINITION): attributes = self.createDefinitions(self.specifications)
        else: attributes = self.createDescriptors(self.specifications)

        assert not flags, 'Unknown flags: %s' % ', '.join(flags)
        return attributes

    # ----------------------------------------------------------------

    def mergeSpecification(self, mergeSpec, withSpec, **keyargs):
        '''
        Merges the provided specifications.

        @param mergeSpec: Specification
            The specification to be merged.
        @param withSpec: Specification
            The specification to merge with.
        @param keyargs: key arguments
            Additional key arguments to be used in constructing the merged specification.
        @return: Specification
            The merged specification.
        '''
        assert isinstance(mergeSpec, Specification), 'Invalid merge specification %s' % mergeSpec
        assert isinstance(withSpec, Specification), 'Invalid with specification %s' % withSpec

        if mergeSpec is withSpec: return mergeSpec

        if mergeSpec.status == REQUIRED:
            if withSpec.status & DEFINED:
                raise AttrError('Improper order for %s, it should be before %s' % (withSpec, mergeSpec))
            status = REQUIRED
            types = intersect(mergeSpec.types, withSpec.types)
            if not types:
                raise AttrError('Incompatible required types of %s, with required types of %s' % (mergeSpec, withSpec))

        elif mergeSpec.status == OPTIONAL:
            if withSpec.status & DEFINED:
                status = DEFINED
                types = set(mergeSpec.types)
            else:
                status = withSpec.status
                types = intersect(mergeSpec.types, withSpec.types)
                if not types:
                    raise AttrError('Incompatible required types of %s, with required types of %s' % (mergeSpec, withSpec))

        elif mergeSpec.status & DEFINED:
            status = DEFINED
            if mergeSpec.status & USED: status |= USED
            if withSpec.status & DEFINED:
                types = set(mergeSpec.types)
                types.update(withSpec.types)
                # In case both definitions are optional we need to relate that to the merged attribute
                if withSpec.status & USED: status |= USED
                if mergeSpec.status & OPTIONAL and withSpec.status & OPTIONAL: status |= OPTIONAL
            else:
                types = set(withSpec.types)

        keyargs['status'] = status

        defined = set(mergeSpec.defined)
        defined.update(withSpec.defined)
        defined, types = reduce(defined), reduce(types)
        if defined != types and not types.issuperset(defined):
            raise ResolverError('Invalid types %s and defined types %s, they cannot be joined, for %s, and %s, '
                                'from merged:%s\n, with:%s' % (', '.join('\'%s\'' % typ.__name__ for typ in types),
                                 ', '.join('\'%s\'' % typ.__name__ for typ in defined), mergeSpec, withSpec,
                                 ''.join(locationStack(clazz) for clazz in mergeSpec.usedIn),
                                 ''.join(locationStack(clazz) for clazz in withSpec.usedIn)))

        keyargs['types'] = types
        keyargs['defined'] = defined

        docs = []
        if mergeSpec.doc is not None: docs.append(mergeSpec.doc)
        if withSpec.doc is not None: docs.append(withSpec.doc)
        keyargs['doc'] = '\n'.join(docs) if docs else None

        spec = mergeSpec.__class__(**keyargs)
        assert isinstance(spec, Specification), 'Invalid specification %s' % spec
        spec.usedIn.update(mergeSpec.usedIn)
        spec.usedIn.update(withSpec.usedIn)

        return spec

    def mergeSpecifications(self, mergeSpecs, withSpecs):
        '''
        Merges the provided specifications.

        @param mergeSpecs: dictionary{string: Specification}
            The specifications to be merged.
        @param withSpecs: dictionary{string: Specification}
            The specifications to merge with.
        @return: dictionary{string: Specification}
            The merged specifications.
        '''
        assert isinstance(mergeSpecs, dict), 'Invalid specifications %s' % mergeSpecs
        assert isinstance(withSpecs, dict), 'Invalid specifications %s' % withSpecs

        specifications = dict(mergeSpecs)
        for name, spec in withSpecs.items():
            ownSpec = specifications.get(name)
            if ownSpec is None: specifications[name] = spec
            else:
                assert isinstance(spec, Specification), 'Invalid specification %s' % spec
                try: specifications[name] = self.mergeSpecification(ownSpec, spec, definedIn=spec.definedIn)
                except AttrError:
                    raise AttrError('Cannot merge attribute \'%s\', from:%s\n, with:%s' %
                                    (name, ''.join(locationStack(clazz) for clazz in ownSpec.usedIn),
                                     ''.join(locationStack(clazz) for clazz in spec.usedIn)))

        return specifications

    def solveSpecifications(self, mergeSpecs, withSpecs):
        '''
        Solve the provided specifications.

        @param mergeSpecs: dictionary{string: Specification}
            The specifications to be solved.
        @param withSpecs: dictionary{string: Specification}
            The specifications to solve with.
        @return: dictionary{string: Specification}
            The solved specifications.
        '''
        assert isinstance(mergeSpecs, dict), 'Invalid specifications %s' % mergeSpecs
        assert isinstance(withSpecs, dict), 'Invalid specifications %s' % withSpecs

        specifications = dict(self.specifications)
        for name, spec in withSpecs.items():
            assert isinstance(spec, Specification), 'Invalid specification %s' % spec
            ownSpec = specifications.get(name)
            if ownSpec is None: specifications[name] = spec
            elif ownSpec.status & DEFINED:
                assert isinstance(ownSpec, Specification), 'Invalid specification %s' % ownSpec
                specifications[name] = self.mergeSpecification(ownSpec, spec)
            else:
                if spec.definedIn is not None and ownSpec.definedIn is not None:
                    specifications[name] = self.mergeSpecification(spec, ownSpec)
                else:
                    specifications[name] = self.mergeSpecification(spec, ownSpec)

        return specifications

    def createDefinitions(self, specifications):
        '''
        Create the definitions attributes.

        @param specifications: dictionary{string: Specification}
            The specifications to create the definitions for.
        @return: dictionary{string: IAttribute}
            The created attributes.
        '''
        assert isinstance(specifications, dict), 'Invalid specifications %s' % specifications
        attributes = {}
        for name, spec in specifications.items():
            assert isinstance(name, str), 'Invalid name %s' % name
            assert isinstance(spec, Specification), 'Invalid specification %s' % spec
            attributes[name] = Attribute(spec)

        return attributes

    def createDescriptors(self, specifications):
        '''
        Create the descriptors attribute.

        @param specifications: dictionary{string: Specification}
            The specifications to create the descriptors for.
        @return: dictionary{string: IAttribute}
            The created attributes.
        '''
        assert isinstance(specifications, dict), 'Invalid specifications %s' % specifications
        attributes = {}
        for name, spec in specifications.items():
            assert isinstance(name, str), 'Invalid name %s' % name
            assert isinstance(spec, Specification), 'Invalid specification %s' % spec
            if spec.status == REQUIRED:
                raise AttrError('Cannot generate attribute %s=%s, used in:%s' %
                                (name, spec, ''.join(locationStack(clazz) for clazz in spec.usedIn)))
            if spec.status & OPTIONAL: continue  # If is optional then no need to create it
            attributes[name] = AttributeObject(spec)

        return attributes

    def __str__(self):
        if not self.specifications: return '%s empty' % self.__class__.__name__
        return '%s[%s]' % (self.__class__.__name__, ', '.join('%s=%s' % (name, self.specifications.get(name))
                                                              for name in sorted(self.specifications)))

class Attribute(IAttribute):
    '''
    Base attribute implementation for a @see: IAttribute that manages a attributes by status.
    '''

    def __init__(self, specification, Resolver=Resolver):
        '''
        Construct the attribute.

        @param specification: Specification
            The attribute specification.
        @param Resolver: class
            The resolver class for the attribute.
        '''
        assert isinstance(specification, Specification), 'Invalid specification %s' % specification
        assert isclass(Resolver), 'Invalid resolver class %s' % Resolver

        self.specification = specification
        self.Resolver = Resolver

        self.isPlaced = False

    def resolver(self):
        '''
        @see: IAttribute.resolver
        '''
        return self.Resolver

    def place(self, clazz, name):
        '''
        @see: IAttribute.place
        '''
        if not self.isPlaced:
            assert isinstance(clazz, ContextMetaClass), 'Invalid class %s' % clazz
            assert isinstance(name, str), 'Invalid name %s' % name
            self.isPlaced, self.__objclass__, self.__name__ = True, clazz, name
            self.specification.usedIn[clazz] = self.specification.status
            if self.specification.definedIn is None and self.specification.status == DEFINED:
                self.specification.definedIn = clazz
        elif not issubclass(clazz, self.__objclass__) or self.__name__ != name:
            raise AttrError('%s\n, is already placed in:%s as attribute %s' %
                            (self, locationStack(self.__objclass__), self.__name__))

    def isValid(self, clazz):
        '''
        @see: IAttribute.isValid
        '''
        if not self.isPlaced: return False
        assert isclass(clazz), 'Invalid class %s' % clazz

        if not isinstance(clazz, ContextMetaClass): return False
        assert isinstance(clazz, ContextMetaClass)

        other = clazz.__attributes__.get(self.__name__)
        if other is None:
            if self.specification.status != REQUIRED: return True
            return False

        if not isinstance(other, Attribute): return False
        assert isinstance(other, Attribute)

        for typ in self.specification.types:
            if typ in other.specification.types: break
        else: return False

        return True

    def isIn(self, clazz):
        '''
        @see: IAttribute.isIn
        '''
        if not self.isPlaced: return False
        assert isclass(clazz), 'Invalid class %s' % clazz

        if not isinstance(clazz, ContextMetaClass): return False
        assert isinstance(clazz, ContextMetaClass)

        other = clazz.__attributes__.get(self.__name__)
        if other is None: return False

        if not isinstance(other, Attribute): return False
        assert isinstance(other, Attribute)

        for typ in self.specification.types:
            if typ in other.specification.types: break
        else: return False

        return True

    # ----------------------------------------------------------------

    def __get__(self, obj, owner=None):
        '''
        Descriptor get.
        '''
        if obj is not None: raise TypeError('Operation not allowed')
        assert self.isPlaced, 'Attribute %s, is not placed in a class' % self
        return self

    def __set__(self, obj, value):
        '''
        Descriptor set.
        '''
        raise TypeError('Operation not allowed')

    def __str__(self):
        st = ''.join((self.__class__.__name__, '.', str(self.specification)))
        if self.isPlaced:
            return ''.join((st, ' in:', locationStack(self.__objclass__), ' as attribute ', self.__name__))
        return ''.join((st, ' unplaced'))

class AttributeObject(Attribute):
    '''
    Object descriptor implementation for a @see: Attribute.
    '''

    def __init__(self, specification, Resolver=Resolver):
        '''
        @see: Attribute.__init__
        '''
        super().__init__(specification, Resolver)
        self.descriptor = None

    def place(self, clazz, name):
        '''
        @see: IAttribute.place
        '''
        if not self.isPlaced:
            assert isinstance(clazz, ContextMetaClass), 'Invalid class %s' % clazz
            assert isinstance(name, str), 'Invalid name %s' % name

            if __debug__:
                assert hasattr(clazz, name), 'Invalid class %s has no descriptor for %s' % (clazz, name)
                self.descriptor = getattr(clazz, name)
                assert isinstance(self.descriptor, IGet), 'Invalid descriptor %s' % self.descriptor
                assert isinstance(self.descriptor, ISet), 'Invalid descriptor %s' % self.descriptor
                setattr(clazz, name, self)
            self.isPlaced, self.__objclass__, self.__name__ = True, clazz, name
        elif not issubclass(clazz, self.__objclass__) or self.__name__ != name:
            raise AttrError('%s\n, is already placed in:%s as attribute %s' %
                            (self, locationStack(self.__objclass__), self.__name__))

    # ----------------------------------------------------------------

    def __get__(self, obj, owner=None):
        '''
        @see: Attribute.__get__
        '''
        if obj is None: return self
        assert self.descriptor, 'Attribute %s, is not placed in a class' % self
        try: return self.descriptor.__get__(obj, owner)
        except AttributeError: return None

    def __set__(self, obj, value):
        '''
        @see: Attribute.__set__
        '''
        assert value is None or isinstance(value, self.specification.types), \
        'Invalid value \'%s\' for %s' % (value, self.specification.types)
        assert self.descriptor, 'Attribute %s, is not placed in a class' % self
        self.descriptor.__set__(obj, value)

# --------------------------------------------------------------------

def reduce(types):
    '''
    Reduces the provided types to only the classes that are the top classes.

    @param types: Iterable(class)
        The types to reduce.
    @return: set(class)
        The reduced types.
    '''
    assert isinstance(types, Iterable), 'Invalid types %s' % types
    reduced = []
    for clazz in types:
        assert isclass(clazz), 'Invalid class %s' % clazz
        k, solved = 0, False
        while k < len(reduced):
            rclazz = reduced[k]

            if rclazz == clazz:
                solved = True
                break
            elif issubclass(clazz, rclazz):
                solved = True
                reduced[k] = clazz
            elif issubclass(rclazz, clazz):
                solved = True
                break

            k += 1

        if not solved: reduced.append(clazz)
    return set(reduced)

def intersect(first, second):
    '''
    Reduces the provided types to only the classes that are the top classes.

    @param first: Iterable(class)
        The first types to reduce.
    @param second: Iterable(class)
        The second types to reduce.
    @return: set(class)
        The intersected types.
    '''
    assert isinstance(first, Iterable), 'Invalid first types %s' % first
    assert isinstance(second, Iterable), 'Invalid types %s' % second
    if not isinstance(second, (list, tuple, set)): second = list(second)
    intersect = []
    for fclazz in first:
        for sclazz in second:
            if issubclass(fclazz, sclazz): intersect.append(sclazz)
            elif issubclass(sclazz, fclazz): intersect.append(fclazz)
    return reduce(intersect)


import torch
import torch.nn as nn
import torch.nn.functional as F

from networks.autoencoder import Autoencoder, Encoder

def conv2d_size_out(size, kernel_size=5, stride=2):
    return (size - (kernel_size - 1) - 1) // stride  + 1

class RecogNet(nn.Module):
    def __init__(self, h, w, outputs):
        super(RecogNet, self).__init__()

        self.conv1 = nn.Conv2d(3, 16, kernel_size=5, stride=2)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=5, stride=2)
        self.bn2 = nn.BatchNorm2d(16)
        self.conv3 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
        self.bn3 = nn.BatchNorm2d(32)

        convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h)))
        convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w)))
        linear_input_size = convw * convh * 32

        self.head = nn.Linear(linear_input_size, outputs)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = x.view(x.size(0), -1)
        x = self.head(x)
        return x

class ControlNet(nn.Module):
    def __init__(self, h, w, outputs):
        super(ControlNet, self).__init__()

        self.conv1 = nn.Conv2d(3, 16, kernel_size=5, stride=2)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 16, kernel_size=5, stride=2)
        self.bn2 = nn.BatchNorm2d(16)
        self.conv3 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
        self.bn3 = nn.BatchNorm2d(32)

        convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h)))
        convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w)))
        linear_input_size = convw * convh * 32

        self.dense = nn.Linear(linear_input_size, 256)
        self.bn4 = nn.BatchNorm1d(256)
        self.head = nn.Linear(64, outputs)

    def forward(self, x, m):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.bn4(self.dense(x)))

        x = x.view(-1, 4, 64)
        m = m.view(-1, 1, 4)
        x = torch.matmul(m, x)
        x = x.view(-1, 64)

        x = self.head(x)
        return x

class SeasonNet(nn.Module):
    def __init__(self, h, w, num_z, num_classes, num_actions):
        super(SeasonNet, self).__init__()

        self.div = int(h / 2)
        self.num_classes = num_classes
        self.num_actions = num_actions

        self.recog = RecogNet(self.div, w, num_classes)
        self.encoder = Encoder(self.div, w, num_z, False)
        self.head = nn.Linear(num_z, num_actions*num_classes)

    def forward(self, x):
        x1 = x[:, :, :self.div, :]
        x2 = x[:, :, self.div:, :]

        x1 = self.recog(x1)
        mask = F.one_hot(x1.argmax(1), num_classes=self.num_classes).to(torch.float)
        mask = mask.view(-1, 1, self.num_classes)

        x2 = self.encoder(x2)
        x2 = self.head(x2)
        x2 = x2.view(-1, self.num_classes, self.num_actions)

        x = torch.matmul(mask, x2)
        x = x.view(-1, self.num_actions)
        return x

def season_recog_net(img_size, num_z, num_classes, num_actions, ae_model, recog_model):

    h, w = img_size

    ae = Autoencoder(int(h /2), w, num_z)
    ae.load_state_dict(torch.load(ae_model))

    model = SeasonNet(h, w, num_z, num_classes, num_actions)

    model.recog.load_state_dict(torch.load(recog_model))
    model.encoder.load_state_dict(ae.encoder.state_dict())

    for name, param in model.named_parameters():
        layer_name = name.split('.')[0]
        if layer_name in ["recog", "encoder"]:
            param.requires_grad = False

    return model

if __name__ == "__main__":

    h = 240
    w = 320
    num_classes = 4
    num_actions = 3
    num_z = 256

    img = torch.rand(4, 3, h, w)

    recog = RecogNet(h / 2, w, num_classes)

    ae = Autoencoder(h / 2, w, num_z)

    model = SeasonNet(h, w, num_z, num_classes, num_actions)

    model.recog.load_state_dict(recog.state_dict())
    model.encoder.load_state_dict(ae.encoder.state_dict())

    for name, param in model.named_parameters():
        layer_name = name.split('.')[0]
        if layer_name in ["recog", "encoder"]:
            param.requires_grad = False

    print(model)

    for name, param in model.named_parameters():
        print(name, param.requires_grad)

    y = model(img)

    print(y)