"""Sequential Monte Carlo sampler also known as
Adaptive Transitional Marcov Chain Monte Carlo sampler.

Runs on any pymc3 model.

Created on March, 2016

Various significant updates July, August 2016

Made pymc3 compatible November 2016
Renamed to SMC and further improvements March 2017

@author: Hannes Vasyura-Bathke
"""
import numpy as np
import pymc3 as pm
from tqdm import tqdm

import theano
import copy
import warnings

from ..model import modelcontext
from ..vartypes import discrete_types
from ..theanof import inputvars, make_shared_replacements, join_nonshared_inputs
import numpy.random as nr

from .arraystep import metrop_select
from ..backends import smc_text as atext

__all__ = ['SMC', 'ATMIP_sample']

EXPERIMENTAL_WARNING = "Warning: SMC is an experimental step method, and not yet"\
    " recommended for use in PyMC3!"

class Proposal(object):
    """Proposal distributions modified from pymc3 to initially create all the
    Proposal steps without repeated execution of the RNG - significant speedup!

    Parameters
    ----------
    s : :class:`numpy.ndarray`
    """
    def __init__(self, s):
        self.s = np.atleast_1d(s)

class MultivariateNormalProposal(Proposal):
    def __call__(self, num_draws=None):
        return np.random.multivariate_normal(
                mean=np.zeros(self.s.shape[0]), cov=self.s, size=num_draws)

proposal_dists = {
    'MultivariateNormal': MultivariateNormalProposal,
        }

def choose_proposal(proposal_name, scale=1.):
    """Initialise and select proposal distribution.

    Parameters
    ----------
    proposal_name : string
        Name of the proposal distribution to initialise
    scale : float or :class:`numpy.ndarray`

    Returns
    -------
    class:`pymc3.Proposal` Object
    """
    return proposal_dists[proposal_name](scale)

class SMC(atext.ArrayStepSharedLLK):
    """Adaptive Transitional Markov-Chain Monte-Carlo sampler class.

    Creates initial samples and framework around the (C)ATMIP parameters

    Parameters
    ----------
    vars : list
        List of variables for sampler
    out_vars : list
        List of output variables for trace recording. If empty unobserved_RVs
        are taken.
    n_chains : int
        Number of chains per stage has to be a large number
        of number of n_jobs (processors to be used) on the machine.
    scaling : float
        Factor applied to the proposal distribution i.e. the step size of the
        Markov Chain
    covariance : :class:`numpy.ndarray`
        (n_chains x n_chains)
        Initial Covariance matrix for proposal distribution,
        if None - identity matrix taken
    likelihood_name : string
        name of the :class:`pymc3.determinsitic` variable that contains the
        model likelihood - defaults to 'like'
    proposal_name :
        Type of proposal distribution, see
        smc.proposal_dists.keys() for options
    tune : boolean
        Flag for adaptive scaling based on the acceptance rate
    tune_interval : int
        Number of steps to tune for
    coef_variation : scalar, float
        Coefficient of variation, determines the change of beta
        from stage to stage, i.e.indirectly the number of stages,
        low coef_variation --> slow beta change,
        results in many stages and vice verca (default: 1.)
    check_bound : boolean
        Check if current sample lies outside of variable definition
        speeds up computation as the forward model wont be executed
        default: True
    model : :class:`pymc3.Model`
        Optional model for sampling step.
        Defaults to None (taken from context).
    random_seed : int
        Optional to set the random seed.  Necessary for initial population.

    References
    ----------
    .. [Ching2007] Ching, J. and Chen, Y. (2007).
        Transitional Markov Chain Monte Carlo Method for Bayesian Model
        Updating, Model Class Selection, and Model Averaging.
        J. Eng. Mech., 10.1061/(ASCE)0733-9399(2007)133:7(816), 816-832.
        `link <http://ascelibrary.org/doi/abs/10.1061/%28ASCE%290733-9399
        %282007%29133:7%28816%29>`__
    """
    default_blocked = True

    def __init__(self, vars=None, out_vars=None, n_chains=100, scaling=1., covariance=None,
                 likelihood_name='like', proposal_name='MultivariateNormal', tune=True,
                 tune_interval=100, coef_variation=1., check_bound=True, model=None,
                 random_seed=-1):
        warnings.warn(EXPERIMENTAL_WARNING)
        if random_seed != -1:
            nr.seed(random_seed)

        model = modelcontext(model)

        if vars is None:
            vars = model.vars

        vars = inputvars(vars)

        if out_vars is None:
            out_vars = model.unobserved_RVs

        out_varnames = [out_var.name for out_var in out_vars]

        if covariance is None and proposal_name == 'MultivariateNormal':
            self.covariance = np.eye(sum(v.dsize for v in vars))
            scale = self.covariance
        elif covariance is None:
            scale = np.ones(sum(v.dsize for v in vars))
        else:
            scale = covariance

        self.proposal_name = proposal_name
        self.proposal_dist = choose_proposal(self.proposal_name, scale=scale)

        self.scaling = np.atleast_1d(scaling)
        self.tune = tune
        self.check_bnd = check_bound
        self.tune_interval = tune_interval
        self.steps_until_tune = tune_interval

        self.proposal_samples_array = self.proposal_dist(n_chains)

        self.stage_sample = 0
        self.accepted = 0

        self.beta = 0
        self.stage = 0
        self.chain_index = 0
        self.resampling_indexes = np.arange(n_chains)

        self.coef_variation = coef_variation
        self.n_chains = n_chains
        self.likelihoods = np.zeros(n_chains)

        self.likelihood_name = likelihood_name
        self._llk_index = out_varnames.index(likelihood_name)
        self.discrete = np.concatenate([[v.dtype in discrete_types] * (v.dsize or 1) for v in vars])
        self.any_discrete = self.discrete.any()
        self.all_discrete = self.discrete.all()

        # create initial population
        self.population = []
        self.array_population = np.zeros(n_chains)
        for _ in range(self.n_chains):
            self.population.append(pm.Point({v.name: v.random() for v in vars}, model=model))

        self.chain_previous_lpoint = copy.deepcopy(self.population)

        shared = make_shared_replacements(vars, model)
        self.logp_forw = logp_forw(out_vars, vars, shared)
        self.check_bnd = logp_forw([model.varlogpt], vars, shared)

        super(SMC, self).__init__(vars, out_vars, shared)

    def astep(self, q0):
        if self.stage == 0:
            l_new = self.logp_forw(q0)
            q_new = q0

        else:
            if not self.stage_sample:
                self.proposal_samples_array = self.proposal_dist(self.n_steps)

            if not self.steps_until_tune and self.tune:
                # Tune scaling parameter
                self.scaling = tune(self.accepted / float(self.tune_interval))
                # Reset counter
                self.steps_until_tune = self.tune_interval
                self.accepted = 0

            delta = self.proposal_samples_array[self.stage_sample, :] * self.scaling

            if self.any_discrete:
                if self.all_discrete:
                    delta = np.round(delta, 0)
                    q0 = q0.astype(int)
                    q = (q0 + delta).astype(int)
                else:
                    delta[self.discrete] = np.round(delta[self.discrete], 0).astype(int)
                    q = q0 + delta
                    q = q[self.discrete].astype(int)
            else:
                q = q0 + delta

            l0 = self.chain_previous_lpoint[self.chain_index]

            if self.check_bnd:
                varlogp = self.check_bnd(q)

                if np.isfinite(varlogp):
                    logp = self.logp_forw(q)
                    q_new, accepted = metrop_select(
                        self.beta * (logp[self._llk_index] - l0[self._llk_index]), q, q0)

                    if accepted:
                        self.accepted += 1
                        l_new = logp
                        self.chain_previous_lpoint[self.chain_index] = l_new
                    else:
                        l_new = l0
                else:
                    q_new = q0
                    l_new = l0

            else:
                logp = self.logp_forw(q)
                q_new, accepted = metrop_select(
                    self.beta * (logp[self._llk_index] - l0[self._llk_index]), q, q0)

                if accepted:
                    self.accepted += 1
                    l_new = logp
                    self.chain_previous_lpoint[self.chain_index] = l_new
                else:
                    l_new = l0

            self.steps_until_tune -= 1
            self.stage_sample += 1

            # reset sample counter
            if self.stage_sample == self.n_steps:
                self.stage_sample = 0

        return q_new, l_new

    def calc_beta(self):
        """Calculate next tempering beta and importance weights based on
        current beta and sample likelihoods.

        Returns
        -------
        beta(m+1) : scalar, float
            tempering parameter of the next stage
        beta(m) : scalar, float
            tempering parameter of the current stage
        weights : :class:`numpy.ndarray`
            Importance weights (floats)
        """
        low_beta = self.beta
        up_beta = 2.
        old_beta = self.beta

        while up_beta - low_beta > 1e-6:
            current_beta = (low_beta + up_beta) / 2.
            temp = np.exp((current_beta - self.beta) * (self.likelihoods - self.likelihoods.max()))
            cov_temp = np.std(temp) / np.mean(temp)
            if cov_temp > self.coef_variation:
                up_beta = current_beta
            else:
                low_beta = current_beta

        beta = current_beta
        weights = temp / np.sum(temp)
        return beta, old_beta, weights

    def calc_covariance(self):
        """Calculate trace covariance matrix based on importance weights.

        Returns
        -------
        cov : :class:`numpy.ndarray`
            weighted covariances (NumPy > 1.10. required)
        """
        cov = np.cov(self.array_population, aweights=self.weights.ravel(), bias=False, rowvar=0)
        if np.isnan(cov).any() or np.isinf(cov).any():
            raise ValueError('Sample covariances not valid! Likely "n_chains" is too small!')
        return np.atleast_2d(cov)

    def select_end_points(self, mtrace):
        """Read trace results (variables and model likelihood) and take end points
        for each chain and set as start population for the next stage.

        Parameters
        ----------
        mtrace : :class:`.base.MultiTrace`

        Returns
        -------
        population : list
            of :func:`pymc3.Point` dictionaries
        array_population : :class:`numpy.ndarray`
            Array of trace end-points
        likelihoods : :class:`numpy.ndarray`
            Array of likelihoods of the trace end-points
        """
        array_population = np.zeros((self.n_chains, self.ordering.dimensions))
        n_steps = len(mtrace)

        # collect end points of each chain and put into array
        for var, slc, shp, _ in self.ordering.vmap:
            slc_population = mtrace.get_values(varname=var, burn=n_steps - 1, combine=True)
            if len(shp) == 0:
                array_population[:, slc] = np.atleast_2d(slc_population).T
            else:
                array_population[:, slc] = slc_population
        # get likelihoods
        likelihoods = mtrace.get_values(varname=self.likelihood_name,
                                        burn=n_steps - 1,
                                        combine=True)

        # map end array_endpoints to dict points
        population = [self.bij.rmap(row) for row in array_population]

        return population, array_population, likelihoods

    def get_chain_previous_lpoint(self, mtrace):
        """Read trace results and take end points for each chain and set as
        previous chain result for comparison of metropolis select.

        Parameters
        ----------
        mtrace : :class:`.base.MultiTrace`

        Returns
        -------
        chain_previous_lpoint : list
            all unobservedRV values, including dataset likelihoods
        """
        array_population = np.zeros((self.n_chains, self.lordering.dimensions))
        n_steps = len(mtrace)
        for var, (_, slc, shp, _) in zip(mtrace.varnames, self.lordering.vmap):
            slc_population = mtrace.get_values(varname=var, burn=n_steps - 1, combine=True)
            if len(shp) == 0:
                array_population[:, slc] = np.atleast_2d(slc_population).T
            else:
                array_population[:, slc] = slc_population

        return [self.lij.rmap(row) for row in array_population[self.resampling_indexes, :]]

    def mean_end_points(self):
        """Calculate mean of the end-points and return point.

        Returns
        -------
        Dictionary of trace variables
        """
        return self.bij.rmap(self.array_population.mean(axis=0))

    def resample(self):
        """Resample pdf based on importance weights.
        based on Kitagawas deterministic resampling algorithm.

        Returns
        -------
        outindex : :class:`numpy.ndarray`
            Array of resampled trace indexes
        """
        parents = np.arange(self.n_chains)
        N_childs = np.zeros(self.n_chains, dtype=int)

        cum_dist = np.cumsum(self.weights)
        u = (parents + np.random.rand()) / self.n_chains
        j = 0
        for i in parents:
            while u[i] > cum_dist[j]:
                j += 1

            N_childs[j] += 1

        indx = 0
        outindx = np.zeros(self.n_chains, dtype=int)
        for i in parents:
            if N_childs[i] > 0:
                for j in range(indx, (indx + N_childs[i])):
                    outindx[j] = parents[i]

            indx += N_childs[i]

        return outindx

def ATMIP_sample(n_steps, step=None, start=None, homepath=None, chain=0, stage=0, n_jobs=1,
                 tune=None, progressbar=False, model=None, random_seed=-1, rm_flag=False):
    """(C)ATMIP sampling algorithm (Cascading - (C) not always relevant)

    Samples the solution space with n_chains of Metropolis chains, where each
    chain has n_steps iterations. Once finished, the sampled traces are
    evaluated:

    (1) Based on the likelihoods of the final samples, chains are weighted
    (2) the weighted covariance of the ensemble is calculated and set as new
        proposal distribution
    (3) the variation in the ensemble is calculated and the next tempering
        parameter (beta) calculated
    (4) New n_chains Metropolis chains are seeded on the traces with high
        weight for n_steps iterations
    (5) Repeat until beta > 1.

    Parameters
    ----------
    n_steps : int
        The number of samples to draw for each Markov-chain per stage
    step : :class:`SMC`
        SMC initialisation object
    start : List of dictionaries
        with length of (n_chains)
        Starting points in parameter space (or partial point)
        Defaults to random draws from variables (defaults to empty dict)
    homepath : string
        Result_folder for storing stages, will be created if not existing.
    chain : int
        Chain number used to store sample in backend. If `n_jobs` is
        greater than one, chain numbers will start here.
    stage : int
        Stage where to start or continue the calculation. It is possible to
        continue after completed stages (stage should be the number of the
        completed stage + 1). If None the start will be at stage = 0.
    n_jobs : int
        The number of cores to be used in parallel. Be aware that theano has
        internal parallelisation. Sometimes this is more efficient especially
        for simple models.
        step.n_chains / n_jobs has to be an integer number!
    tune : int
        Number of iterations to tune, if applicable (defaults to None)
    progressbar : bool
        Flag for displaying a progress bar
    model : :class:`pymc3.Model`
        (optional if in `with` context) has to contain deterministic
        variable name defined under step.likelihood_name' that contains the
        model likelihood
    random_seed : int or list of ints
        A list is accepted, more if `n_jobs` is greater than one.
    rm_flag : bool
        If True existing stage result folders are being deleted prior to
        sampling.

    References
    ----------
    .. [Minson2013] Minson, S. E. and Simons, M. and Beck, J. L., (2013),
        Bayesian inversion for finite fault earthquake source models
        I- Theory and algorithm. Geophysical Journal International, 2013,
        194(3), pp.1701-1726,
        `link <https://gji.oxfordjournals.org/content/194/3/1701.full>`__
    """
    warnings.warn(EXPERIMENTAL_WARNING)

    model = modelcontext(model)
    step.n_steps = int(n_steps)
    if random_seed != -1:
        nr.seed(random_seed)

    if step is None:
        raise TypeError('Argument `step` has to be a SMC step object.')

    if homepath is None:
        raise TypeError('Argument `homepath` should be path to result_directory.')

    if n_jobs > 1:
        if not (step.n_chains / float(n_jobs)).is_integer():
            raise TypeError('n_chains / n_jobs has to be a whole number!')

    if start is not None:
        if len(start) != step.n_chains:
            raise TypeError('Argument `start` should have dicts equal the '
                            'number of chains (`step.n_chains`)')
        else:
            step.population = start

    if not any(step.likelihood_name in var.name for var in model.deterministics):
        raise TypeError('Model (deterministic) variables need to contain a variable %s '
                        'as defined in `step`.' % step.likelihood_name)

    stage_handler = atext.TextStage(homepath)

    if progressbar and n_jobs > 1:
        progressbar = False

    if stage == 0:
        # continue or start initial stage
        step.stage = stage
        draws = 1
    else:
        step = stage_handler.load_atmip_params(stage)
        draws = step.n_steps

    stage_handler.clean_directory(stage, None, rm_flag)
    with model:
        chains = stage_handler.recover_existing_results(stage, draws, step, n_jobs)
        if chains is not None:
            rest = len(chains) % n_jobs
            if rest > 0:
                pm._log.info('Fixing %i chains ...' % rest)
                chains, rest_chains = chains[:-rest], chains[-rest:]
                # process traces that are not a multiple of n_jobs
                sample_args = {
                    'draws': draws,
                    'step': step,
                    'stage_path': stage_handler.stage_path(stage),
                    'progressbar': progressbar,
                    'model': model,
                    'n_jobs': rest,
                    'chains': rest_chains}

                _iter_parallel_chains(**sample_args)
                pm._log.info('Back to normal!')

    with model:
        while step.beta < 1:
            if step.stage == 0:
                # Initial stage
                pm._log.info('Sample initial stage: ...')
                draws = 1
            else:
                draws = n_steps

            pm._log.info('Beta: %f Stage: %i' % (step.beta, step.stage))

            # Metropolis sampling intermediate stages
            chains = stage_handler.clean_directory(stage, chains, rm_flag)
            sample_args = {
                    'draws': draws,
                    'step': step,
                    'stage_path': stage_handler.stage_path(step.stage),
                    'progressbar': progressbar,
                    'model': model,
                    'n_jobs': n_jobs,
                    'chains': chains}

            _iter_parallel_chains(**sample_args)

            mtrace = stage_handler.load_multitrace(step.stage)

            step.population, step.array_population, step.likelihoods = step.select_end_points(
                mtrace)
            step.beta, step.old_beta, step.weights = step.calc_beta()

            if step.beta > 1.:
                pm._log.info('Beta > 1.: %f' % step.beta)
                step.beta = 1.
                stage_handler.dump_atmip_params(step)
                if stage == -1:
                    chains = []
                else:
                    chains = None
            else:
                step.covariance = step.calc_covariance()
                step.proposal_dist = choose_proposal(step.proposal_name, scale=step.covariance)
                step.resampling_indexes = step.resample()
                step.chain_previous_lpoint = step.get_chain_previous_lpoint(mtrace)

                stage_handler.dump_atmip_params(step)

                step.stage += 1
                del(mtrace)

        # Metropolis sampling final stage
        pm._log.info('Sample final stage')
        step.stage = -1
        temp = np.exp((1 - step.old_beta) * (step.likelihoods - step.likelihoods.max()))
        step.weights = temp / np.sum(temp)
        step.covariance = step.calc_covariance()
        step.proposal_dist = choose_proposal(step.proposal_name, scale=step.covariance)
        step.resampling_indexes = step.resample()
        step.chain_previous_lpoint = step.get_chain_previous_lpoint(mtrace)

        sample_args['step'] = step
        sample_args['stage_path'] = stage_handler.stage_path(step.stage)
        sample_args['chains'] = chains
        _iter_parallel_chains(**sample_args)

        stage_handler.dump_atmip_params(step)
        return stage_handler.load_multitrace(step.stage, model=model)

def _sample(draws, step=None, start=None, trace=None, chain=0, tune=None,
            progressbar=True, model=None, random_seed=-1):

    sampling = _iter_sample(draws, step, start, trace, chain,
                            tune, model, random_seed)

    if progressbar:
        sampling = tqdm(sampling, total=draws)

    try:
        for strace in sampling:
            pass

    except KeyboardInterrupt:
        pass

    return chain

def _iter_sample(draws, step, start=None, trace=None, chain=0, tune=None,
                 model=None, random_seed=-1):
    """Modified from :func:`pymc3.sampling._iter_sample` to be more efficient with SMC algorithm."""
    model = modelcontext(model)
    draws = int(draws)
    if draws < 1:
        raise ValueError('Argument `draws` should be above 0.')

    if start is None:
        start = {}

    if random_seed != -1:
        nr.seed(random_seed)

    try:
        step = pm.step_methods.CompoundStep(step)
    except TypeError:
        pass

    point = pm.Point(start, model=model)
    step.chain_index = chain
    trace.setup(draws, chain)
    for i in range(draws):
        if i == tune:
            step = pm.sampling.stop_tuning(step)

        point, out_list = step.step(point)
        trace.record(out_list)
        yield trace

def _work_chain(work):
    """Wrapper function for parallel execution of _sample i.e. the Markov Chains.

    Parameters
    ----------
    work : List
        Containing all the information that is unique for each Markov Chain
        i.e. [:class:'SMC', chain_number(int),
        sampling index(int), start_point(dictionary)]

    Returns
    -------
    chain : int
        Index of chain that has been sampled
    """
    return _sample(*work)

def _iter_parallel_chains(draws, step, stage_path, progressbar, model, n_jobs, chains=None):
    """Do Metropolis sampling over all the chains with each chain being
    sampled 'draws' times. Parallel execution according to n_jobs.
    """
    if chains is None:
        chains = range(step.n_chains)

    pm._log.info('Initialising chain traces ...')

    max_int = np.iinfo(np.int32).max
    random_seeds = nr.randint(1, max_int, size=len(chains))

    pm._log.info('Sampling ...')
    work = [(draws,
             step,
             step.population[step.resampling_indexes[chain]],
             atext.TextChain(stage_path, model=model),
             chain,
             None,
             False,
             model,
             rseed) for chain, rseed in
            zip(chains, random_seeds)]

    if draws < 10:
        chunksize = n_jobs
    else:
        chunksize = 1

    p = atext.paripool(_work_chain, work, chunksize=chunksize, nprocs=n_jobs)

    if n_jobs == 1 and progressbar:
        p = tqdm(p, total=len(chains))

    for _ in p:
        pass

def tune(acc_rate):
    """Tune adaptively based on the acceptance rate.

    Parameters
    ----------
    acc_rate: scalar, float
        Acceptance rate of the Metropolis sampling

    Returns
    -------
    scaling: scalar float
    """
    # a and b after Muto & Beck 2008 .
    a = 1. / 9
    b = 8. / 9
    return np.power((a + (b * acc_rate)), 2)

def logp_forw(out_vars, vars, shared):
    """Compile Theano function of the model and the input and output variables.

    Parameters
    ----------
    out_vars : List
        containing :class:`pymc3.Distribution` for the output variables
    vars : List
        containing :class:`pymc3.Distribution` for the input variables
    shared : List
        containing :class:`theano.tensor.Tensor` for dependend shared data
    """
    out_list, inarray0 = join_nonshared_inputs(out_vars, vars, shared)
    f = theano.function([inarray0], out_list)
    f.trust_input = True
    return f

import numpy,import theano,import pymc3,from pymc3


""" Conditional Similarity Network
During training TripletDataset load anchor, positive, negative images, and condition
A Test Dataset load pair images, target and condition, conditions are:

    Conditions
    upper_bottom
    upper_shoe
    upper_bag
    upper_accessory
    bottom_shoe
    bottom_bag
    bottom_accessory
    shoe_bag
    shoe_accessory
    bag_accessory
"""

import sys
import torch
import logging
import torchvision
import numpy as np
import torch.nn as nn
from sklearn import metrics
from torch.utils.data import DataLoader, Dataset
from polyvore_dataset import TripletDataset, CategoryDataset
from evaluate import test_compatibility_auc, test_fitb_quesitons
from utils import AverageMeter, prepare_dataloaders, config_logging
from model import CompatModel

# Leave a comment for this training, and it will be used for name suffix of log and saved model
comment = '_'.join(sys.argv[1:])

# Logger
config_logging(comment)

# Hyperparameters
img_size = 112
emb_size = 64
device = torch.device("cuda")

# Dataloader
transform = torchvision.transforms.Compose(
    [
        torchvision.transforms.Scale((img_size, img_size)),
        torchvision.transforms.CenterCrop(112),
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        ),
    ]
)
train_dataset = TripletDataset(
    root_dir="../../data/images/",
    data_dir="../../data/",
    transform=transform,
    data_file="train_no_dup_with_category_3more_name.json",
)
train_loader = DataLoader(train_dataset, 32, shuffle=True, num_workers=4)
val_dataset = TripletDataset(
    root_dir="../../data/images/",
    data_dir="../../data/",
    transform=transform,
    data_file="valid_no_dup_with_category_3more_name.json",
    is_train=True,
)
val_loader = DataLoader(val_dataset, 32, shuffle=False, num_workers=4)
test_dataset = TripletDataset(
    root_dir="../../data/images/",
    data_dir="../../data/",
    transform=transform,
    data_file="test_no_dup_with_category_3more_name.json",
    is_train=True,
)
test_loader = DataLoader(test_dataset, 32, shuffle=False, num_workers=4)

val_auc_dataset = CategoryDataset(
    root_dir="../../data/images/",
    data_dir="../../data/",
    transform=transform,
    use_mean_img=True,
    data_file="valid_no_dup_with_category_3more_name.json",
    neg_samples=True,
)

# Model
tnet = CompatModel(
    emb_size,
    n_conditions=len(train_dataset.conditions) // 2,
    learnedmask=True,
    prein=False,
)
tnet = tnet.to(device)

def accuracy(dista, distb):
    margin = 0
    pred = (dista - distb - margin).cpu().data
    return (pred > 0).sum().item() / dista.size(0)

# Hyperparameters
criterion = torch.nn.MarginRankingLoss(margin=0.2)
parameters = filter(lambda p: p.requires_grad, tnet.parameters())
optimizer = torch.optim.Adam(parameters, lr=5e-5)
n_parameters = sum([p.data.nelement() for p in tnet.parameters()])
logging.info(" + Number of params: {}".format(n_parameters))

# Train process
best_acc = -1
for epoch in range(1, 50 + 1):
    logging.info("** Epoch: {} **".format(epoch))
    # Train phase
    tnet.train()
    losses = AverageMeter()
    accs = AverageMeter()
    for batch_num, (a_img, p_img, n_img, c) in enumerate(train_loader, 1):
        a_img, p_img, n_img, c = (
            a_img.to(device),
            p_img.to(device),
            n_img.to(device),
            c.to(device),
        )
        # Original code need input triplet like: anchor, far, close
        dista, distb, mask_norm, embed_norm, mask_embed_norm = tnet(a_img, n_img, p_img, c)
        target = torch.FloatTensor(dista.size()).fill_(1).to(device)
        loss_triplet = criterion(dista, distb, target)
        loss_embed = embed_norm / np.sqrt(a_img.size(0))
        loss_mask = mask_norm / a_img.size(0)
        loss = loss_triplet + 5e-3 * loss_embed + 5e-4 * loss_mask
        losses.update(loss_triplet.item(), a_img.size(0))
        accs.update(accuracy(dista, distb), a_img.size(0))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch_num % 50 == 0:
            logging.info("#{} Loss: {:.4f} (avg: {:.4f}), Accuracy: {:.4f}".format(
                batch_num, losses.val, losses.avg, accs.avg))
    logging.info("Train Loss: {:.4f}, Accuracy: {:.4f}".format(losses.avg, accs.avg))

    # Evaluation phase
    tnet.eval()
    losses = AverageMeter()
    accs = AverageMeter()
    for batch_num, (a_img, p_img, n_img, c) in enumerate(val_loader, 1):
        a_img, p_img, n_img, c = (
            a_img.to(device),
            p_img.to(device),
            n_img.to(device),
            c.to(device),
        )
        # Original code need input triplet like: anchor, far, close
        with torch.no_grad():
            dista, distb, mask_norm, embed_norm, mask_embed_norm = tnet(a_img, n_img, p_img, c)

        target = torch.FloatTensor(dista.size()).fill_(1).to(device)
        loss_triplet = criterion(dista, distb, target)
        losses.update(loss_triplet.item(), a_img.size(0))
        accs.update(accuracy(dista, distb), a_img.size(0))

    # Valid AUC
    auc = test_compatibility_auc(val_auc_dataset, tnet.embeddingnet)
    acc = test_fitb_quesitons(val_auc_dataset, tnet.embeddingnet)
    logging.info("Valid Loss: {:.4f}, Accuracy: {:.4f}, AUC: {:.4f}, FitbACC: {:.4f}".format(
        losses.avg, accs.avg, auc, acc))

    # Save Model
    if accs.avg > best_acc:
        best_acc = accs.avg
        torch.save(tnet.state_dict(), "csn_model_best.pth")
        logging.info("Found Best Accuracy {:.4f}, saved model to csn_model_best.pth".format(accs.avg))


#!/usr/bin/env python

import getopt
import math
import copy
import cv2
import numpy
import os
import time
import argparse
import torch
#import torch.utils.serialization
from tqdm import tqdm
import random

# from sepconv.model import Network           # model
from sepconv.model import MetaNetwork as Network           # model
#from combined240 import Combined240 # data
from data.vimeo_septuplet import VimeoSeptuplet # data
import utils

import warnings
warnings.filterwarnings("ignore", category=UserWarning)

VIZ = True

torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance

##########################################################

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--exp_name', type=str, default='temp', help='Name of the experiment')
    parser.add_argument('--meta_algorithm', type=str, default='maml', help='Meta-learning algorithm to use')
    parser.add_argument('--batch_size', '--bs', type=int, default=4)
    parser.add_argument('--val_batch_size', type=int, default=4)
    parser.add_argument('--test_batch_size', type=int, default=4)
    parser.add_argument('--inner_lr', type=float, default=0.00001)
    parser.add_argument('--outer_lr', type=float, default=0.00001)
    parser.add_argument('--max_epoch', type=int, default=100)
    parser.add_argument('--val_iter', type=int, default=200)
    parser.add_argument('--resume', action='store_true', help='True if resuming from a pretrained model')
    parser.add_argument('--resume_model', type=str, default='l1')
    parser.add_argument('--resume_ckpt', action='store_true')
    parser.add_argument('--resume_exp', type=str, default=None)
    parser.add_argument('--start_epoch', type=int, default=0)
    parser.add_argument('--logfreq', type=int, default=20)
    parser.add_argument('--num_inner_update', type=int, default=1)
    parser.add_argument('--mode', type=str, default='')
    args = parser.parse_args()

    # Train args
    args.step = 3
    args.crop_size = [128, 128]
    args.rotation = [-10, 10]
    args.crop_policy = 'random'
    args.flip = True
    args.scale_factor = [1.07, 1.5]

    return args

##########################################################
cfg = parse_args()
VAL_ITER_CUT = 1e6 if cfg.val_iter==-1 else cfg.val_iter

net = Network(cfg.resume, cfg.resume_model).cuda() #.eval()

criterion = torch.nn.L1Loss()
optimizer = torch.optim.Adamax(net.parameters(), lr=cfg.outer_lr)

if cfg.resume_ckpt:
    utils.load_checkpoint(cfg, net, optimizer)
    '''
    print('loading pretrained checkpoint.. %s' % cfg.exp_name)
    #checkpoint = torch.load(os.path.join('checkpoint', cfg.exp_name, 'checkpoint.pth'))
    checkpoint = torch.load(os.path.join('checkpoint', cfg.exp_name, 'model_best.pth'))
    cfg.start_epoch = checkpoint['epoch']
    net.load_state_dict(checkpoint['state_dict'])
    #optimizer.load_state_dict(checkpoint['grad_dict'])
    del checkpoint'''

##########################################################

train_loader = torch.utils.data.DataLoader(#Combined240('data/combined', cfg),
    #VimeoSeptuplet('data/vimeo_septuplet', [1, 2, 3, 4, 5, 6, 7], is_training=True),
    VimeoSeptuplet(cfg),
    batch_size=cfg.batch_size,
    shuffle=True,
    num_workers=16,
    pin_memory=True,
    drop_last=False
)

def train(epoch):
    print("Training Epoch: %d" % epoch)
    net.train()

    losses = utils.AverageMeter()
    psnrs = utils.AverageMeter()
    ssims = utils.AverageMeter()
    batch_time = utils.AverageMeter()
    t = time.time()
    for i, (images, _) in enumerate(train_loader):
        if i == 4000:
            break
        images = [im.cuda() for im in images]

        # Meta training
        k = cfg.num_inner_update   # inner update iteration

        if cfg.meta_algorithm == 'reptile':
            weights_before = copy.deepcopy(net.state_dict())
            inner_optimizer = torch.optim.Adamax(net.parameters(), lr=cfg.inner_lr)

            # inner loop
            for _k in range(k):
                indices = [ [0, 2, 4], [2, 4, 6] ]
                total_loss = 0
                for ind in indices:
                    output = net(images[ind[0]].clone(), images[ind[2]].clone())
                    loss = criterion(output, images[ind[1]])
                    total_loss = total_loss + loss

                inner_optimizer.zero_grad()
                total_loss.backward()
                inner_optimizer.step()

            # Reptile - outer update
            outerstepsize = cfg.outer_lr
            weights_after = net.state_dict()
            net.load_state_dict({name:
                weights_before[name] + (weights_after[name] - weights_before[name]) * outerstepsize
                for name in weights_before})

            # calculate loss w/ updated model
            input0, input1, target = images[2], images[4], images[3]

            with torch.no_grad():
                output = net(input0, input1)
                loss = criterion(output, target)

        elif cfg.meta_algorithm == 'maml':

            base_net = copy.deepcopy(net)
            inner_optimizer = torch.optim.Adamax(net.parameters(), lr=cfg.inner_lr)

            # inner loop
            for _k in range(k):
                indices = [ [0, 2, 4], [2, 4, 6] ]
                total_loss = 0
                for ind in indices:
                    output = net(images[ind[0]].clone(), images[ind[2]].clone())
                    loss = criterion(output, images[ind[1]])
                    total_loss = total_loss + loss

                inner_optimizer.zero_grad()
                total_loss.backward()
                inner_optimizer.step()

            # Forward on query data
            outerstepsize = cfg.outer_lr
            input0, input1, target = images[2], images[4], images[3]
            output = net(input0, input1)
            loss = criterion(output, target)

            # Copy base parameters to 'net' to connect the computation graph
            for param, base_param in zip(net.parameters(), base_net.parameters()):
                param.data = base_param.data

            # Calculate gradient & update meta-learner
            optimizer.zero_grad()
            grads = torch.autograd.grad(loss, net.parameters())
            for j, param in enumerate(net.parameters()):
                #param = param - grads[j] * outerstepsize
                param.grad = grads[j]
            optimizer.step()

        losses.update(loss.item(), images[0].size(0))
        batch_time.update(time.time() - t)
        t = time.time()

        # Logging
        if i % cfg.logfreq == 0:
            utils.eval_metrics(output, target, psnrs, ssims, None)
            print(('Epoch: [%d][%d/%d],\tTime %.3f (%.3f)\tLoss %.4f (%.4f)\tPSNR %.2f\t'
                   % (epoch, i, len(train_loader), batch_time.val, batch_time.avg,
                      losses.val, losses.avg, psnrs.val)))
            losses.reset()
            psnrs.reset()
            ssims.reset()
            batch_time.reset()

def validate(model, epoch):
    val_dataset = VimeoSeptuplet(cfg) #('data/vimeo_septuplet', [1, 2, 3, 4, 5, 6, 7], is_training=False)
    val_loader = torch.utils.data.DataLoader(val_dataset,
        batch_size=cfg.batch_size, shuffle=False, num_workers=8, pin_memory=True)

    #model.eval()
    if True: #with torch.no_grad():
        losses = utils.AverageMeter()
        psnrs = utils.AverageMeter()
        ssims = utils.AverageMeter()
        batch_time = utils.AverageMeter()
        t = time.time()

        for i, (images, metadata) in enumerate(tqdm(val_loader)):
            if i == VAL_ITER_CUT:
                break
            t = time.time()
            images = [im.cuda() for im in images]
            target = images[3]

            weights_before = copy.deepcopy(net.state_dict())
            val_optimizer = torch.optim.Adamax(net.parameters(), lr=cfg.inner_lr)

            # inner loop
            k = cfg.num_inner_update
            for _k in range(k):
                indices = [ [0, 2, 4], [2, 4, 6] ]
                total_loss = 0
                for ind in indices:
                    output = net(images[ind[0]].clone(), images[ind[2]].clone())
                    loss = criterion(output, images[ind[1]])
                    total_loss = total_loss + loss

                val_optimizer.zero_grad()
                total_loss.backward()
                val_optimizer.step()

            with torch.no_grad():
                input0, input1, target = images[2], images[4], images[3]
                output = model(input0, input1)

                batch_time.update(time.time() - t)
                loss = criterion(output, target)
                losses.update(loss.item(), input0.size(0))

                utils.eval_metrics(output, target, psnrs, ssims, None)
                print(psnrs.val, ssims.val)

                if VIZ:
                    for b in range(images[0].size(0)):
                        imgpath = metadata['imgpaths'][0][b]
                        #print(imgpath)
                        savepath = os.path.join('checkpoint', cfg.exp_name, 'vimeoSeptuplet', imgpath.split('/')[-3], imgpath.split('/')[-2])
                        if not os.path.exists(savepath):
                            #print('make dirs... %s' % savepath)
                            os.makedirs(savepath)
                        img_pred = (output[b].data.permute(1, 2, 0).clamp_(0, 1).cpu().numpy()[..., ::-1] * 255).astype(numpy.uint8)
                        cv2.imwrite(os.path.join(savepath, 'im2_pred.png'), img_pred)

            #batch_time.update(time.time() - t)
            #t = time.time()

            # restore the original base weight
            net.load_state_dict(weights_before)

    print("val_losses: %f" % losses.avg)
    print("val_PSNR: %f, val_SSIM: %f" % (psnrs.avg, ssims.avg))
    print("Time per batch: %.3f" % batch_time.avg)
    return psnrs.avg
# end

##########################################################
if __name__ == '__main__':
    best_PSNR = 0.0
    if cfg.mode == 'test':
        PSNR = validate(net, cfg.start_epoch)
        exit(0)

    for epoch in range(cfg.start_epoch, cfg.max_epoch):
        train(epoch)

        PSNR = validate(net, epoch)
        is_best = PSNR > best_PSNR
        best_PSNR = max(PSNR, best_PSNR)
        utils.save_checkpoint({
            'epoch': epoch + 1,
            'arch': cfg,
            'state_dict': net.state_dict(),
            'optimizer': optimizer.state_dict(),
            'best_PSNR': best_PSNR
        }, is_best, cfg.exp_name)


"""Trainer for OCR CTC model."""
import paddle.fluid as fluid
import paddle.fluid.profiler as profiler
from utility import add_arguments, print_arguments, to_lodtensor, get_feeder_data
from crnn_ctc_model import ctc_train_net
import ctc_reader
import argparse
import functools
import sys
import time
import os
import numpy as np

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
# yapf: disable
add_arg('batch_size',        int,   32,         "Minibatch size.")
add_arg('total_step',        int,   720000,    "The number of iterations. Zero or less means whole training set. More than 0 means the training set might be looped until # of iterations is reached.")
add_arg('log_period',        int,   1000,       "Log period.")
add_arg('save_model_period', int,   15000,      "Save model period. '-1' means never saving the model.")
add_arg('eval_period',       int,   15000,      "Evaluate period. '-1' means never evaluating the model.")
add_arg('save_model_dir',    str,   "./models", "The directory the model to be saved to.")
add_arg('init_model',        str,   None,       "The init model file of directory.")
add_arg('use_gpu',           bool,  True,      "Whether use GPU to train.")
add_arg('min_average_window',int,   10000,     "Min average window.")
add_arg('max_average_window',int,   12500,     "Max average window. It is proposed to be set as the number of minibatch in a pass.")
add_arg('average_window',    float, 0.15,      "Average window.")
add_arg('parallel',          bool,  False,     "Whether use parallel training.")
add_arg('profile',           bool,  False,      "Whether to use profiling.")
add_arg('skip_batch_num',    int,   0,          "The number of first minibatches to skip as warm-up for better performance test.")
add_arg('skip_test',         bool,  False,      "Whether to skip test phase.")
# yapf: enable

def train(args, data_reader=ctc_reader):
    """OCR CTC training"""
    num_classes = None
    train_images = None
    train_list = None
    test_images = None
    test_list = None
    num_classes = data_reader.num_classes(
    ) if num_classes is None else num_classes
    data_shape = data_reader.data_shape()
    # define network
    images = fluid.layers.data(name='pixel', shape=data_shape, dtype='float32')
    label = fluid.layers.data(
        name='label', shape=[1], dtype='int32', lod_level=1)
    sum_cost, error_evaluator, inference_program, model_average = ctc_train_net(
        images, label, args, num_classes)

    # data reader
    train_reader = data_reader.train(
        args.batch_size,
        train_images_dir=train_images,
        train_list_file=train_list,
        cycle=args.total_step > 0)
    test_reader = data_reader.test(
        test_images_dir=test_images, test_list_file=test_list)

    # prepare environment
    place = fluid.CPUPlace()
    if args.use_gpu:
        place = fluid.CUDAPlace(0)
    exe = fluid.Executor(place)
    exe.run(fluid.default_startup_program())

    # load init model
    if args.init_model is not None:
        model_dir = args.init_model
        model_file_name = None
        if not os.path.isdir(args.init_model):
            model_dir = os.path.dirname(args.init_model)
            model_file_name = os.path.basename(args.init_model)
        fluid.io.load_params(exe, dirname=model_dir, filename=model_file_name)
        print "Init model from: %s." % args.init_model

    train_exe = exe
    error_evaluator.reset(exe)
    if args.parallel:
        train_exe = fluid.ParallelExecutor(
            use_cuda=True if args.use_gpu else False, loss_name=sum_cost.name)

    fetch_vars = [sum_cost] + error_evaluator.metrics

    def train_one_batch(data):
        var_names = [var.name for var in fetch_vars]
        if args.parallel:
            results = train_exe.run(var_names,
                                    feed=get_feeder_data(data, place))
            results = [np.array(result).sum() for result in results]
        else:
            results = train_exe.run(feed=get_feeder_data(data, place),
                                    fetch_list=fetch_vars)
            results = [result[0] for result in results]
        return results

    def test(iter_num):
        error_evaluator.reset(exe)
        for data in test_reader():
            exe.run(inference_program, feed=get_feeder_data(data, place))
        _, test_seq_error = error_evaluator.eval(exe)
        print "\nTime: %s; Iter[%d]; Test seq error: %s.\n" % (
            time.time(), iter_num, str(test_seq_error[0]))

    def save_model(args, exe, iter_num):
        filename = "model_%05d" % iter_num
        fluid.io.save_params(
            exe, dirname=args.save_model_dir, filename=filename)
        print "Saved model to: %s/%s." % (args.save_model_dir, filename)

    iter_num = 0
    stop = False
    while not stop:
        total_loss = 0.0
        total_seq_error = 0.0
        batch_times = []
        # train a pass
        for data in train_reader():
            if args.total_step > 0 and iter_num == args.total_step + args.skip_batch_num:
                stop = True
                break
            if iter_num < args.skip_batch_num:
                print("Warm-up iteration")
            if iter_num == args.skip_batch_num:
                profiler.reset_profiler()
            start = time.time()
            results = train_one_batch(data)
            batch_time = time.time() - start
            fps = args.batch_size / batch_time
            batch_times.append(batch_time)
            total_loss += results[0]
            total_seq_error += results[2]

            iter_num += 1
            # training log
            if iter_num % args.log_period == 0:
                print "\nTime: %s; Iter[%d]; Avg Warp-CTC loss: %.3f; Avg seq err: %.3f" % (
                    time.time(), iter_num,
                    total_loss / (args.log_period * args.batch_size),
                    total_seq_error / (args.log_period * args.batch_size))
                sys.stdout.flush()
                total_loss = 0.0
                total_seq_error = 0.0

            # evaluate
            if not args.skip_test and iter_num % args.eval_period == 0:
                if model_average:
                    with model_average.apply(exe):
                        test(iter_num)
                else:
                    test(iter_num)

            # save model
            if iter_num % args.save_model_period == 0:
                if model_average:
                    with model_average.apply(exe):
                        save_model(args, exe, iter_num)
                else:
                    save_model(args, exe, iter_num)
        # Postprocess benchmark data
        latencies = batch_times[args.skip_batch_num:]
        latency_avg = np.average(latencies)
        latency_pc99 = np.percentile(latencies, 99)
        fpses = np.divide(args.batch_size, latencies)
        fps_avg = np.average(fpses)
        fps_pc99 = np.percentile(fpses, 1)

        # Benchmark output
        print('\nTotal examples (incl. warm-up): %d' %
              (iter_num * args.batch_size))
        print('average latency: %.5f s, 99pc latency: %.5f s' % (latency_avg,
                                                                 latency_pc99))
        print('average fps: %.5f, fps for 99pc latency: %.5f' % (fps_avg,
                                                                 fps_pc99))

def main():
    args = parser.parse_args()
    print_arguments(args)
    if args.profile:
        if args.use_gpu:
            with profiler.cuda_profiler("cuda_profiler.txt", 'csv') as nvprof:
                train(args, data_reader=ctc_reader)
        else:
            with profiler.profiler("CPU", sorted_key='total') as cpuprof:
                train(args, data_reader=ctc_reader)
    else:
        train(args, data_reader=ctc_reader)

if __name__ == "__main__":
    main()


import os.path
from data.animation_dataset import AnimationDataset
from data.keypoint_dataset import KeypointDataset
from data.image_folder import make_grouped_dataset, check_path_valid
from data.keypoint2img import interpPoints, drawEdge
import pandas as pd
import numpy as np
import torch
import torchvision.transforms as transforms
import torchvision.transforms.functional as F
from PIL import Image
import cv2
from skimage import feature
from util import openpose_utils
import json
import glob

class KeypointTestDataset(AnimationDataset):

    @staticmethod
    def modify_commandline_options(parser, is_train):
        parser = KeypointDataset.modify_commandline_options(parser, is_train)
        parser.add_argument('--pose_extractor', type=str, default='alphapose', help='openpose | alphapose' )
        parser.add_argument('--sub_dataset', type=str, default='iper', help='iper | fashon')
        parser.add_argument('--eval_set', type=str, default='train', help='train | test')
        parser.add_argument('--start_frame', type=int, default=0)
        parser.add_argument('--total_test_frames', default=None)
        parser.set_defaults(load_size=256)
        parser.set_defaults(nThreads=1)

        return parser

    def initialize(self, opt):
        assert opt.isTrain == False
        self.opt = opt
        self.A_paths, self.B_paths_noise = self.get_paths(opt)
        self.init_frame_idx(self.A_paths)
        self.load_size = (self.opt.load_size, self.opt.load_size)

    def get_paths(self, opt):
        root = opt.dataroot
        phase = opt.eval_set +'_256'
        dir_A = os.path.join(opt.dataroot, phase, 'train_A')
        dir_B_noise = os.path.join(opt.dataroot, phase, 'train_'+self.opt.pose_extractor)
        A_paths = sorted(make_grouped_dataset(dir_A))
        B_paths_noise = sorted(make_grouped_dataset(dir_B_noise))
        check_path_valid(A_paths, B_paths_noise)
        return A_paths, B_paths_noise

    def __getitem__(self, index):
        _, _, _, seq_idx = self.update_seq_idx(self.B_paths_noise, index)
        A_paths = self.A_paths[seq_idx]
        B_paths_noise = self.B_paths_noise[seq_idx]
        gen_images = None

        n_frames_total, start_idx, t_step, org_size = self.get_video_params(self.opt, self.n_frames_total, len(B_paths_noise), self.frame_idx, A_paths)
        left_pad, current, right_pad = self.get_keypoint_param(n_frames_total, start_idx, t_step, len(B_paths_noise))
        self.org_size = (org_size[1], org_size[0])

        left,_ = self.load_keypoints(left_pad, B_paths_noise)
        current, gen_path = self.load_keypoints(current, B_paths_noise)
        right,_ = self.load_keypoints(right_pad, B_paths_noise)
        input_kp = torch.cat([left,current,right],1)
        self.frame_idx += self.opt.n_frames_pre_load_test

        for i in range(n_frames_total):
            A_path = A_paths[start_idx + i * t_step]
            image = self.load_image(A_path)
            gen_images = self.concat_frame(gen_images, image)

        return_list = {'gen_keypoints': input_kp,  'current_keypoints':current, 'gen_images':gen_images,
                       'shape':org_size, 'gen_paths':gen_path, 'frame_idx':self.frame_idx, 'change_seq':self.change_seq
                       }

        return return_list

    def load_image(self, A_path):
        A_img = Image.open(A_path)
        Ai = self.transform_image(A_img, self.load_size)
        return Ai

    def load_keypoints(self, keypoint_list, B_paths):
        pose_list=[]
        B_path_recoder=[]
        for index in keypoint_list:
            B_path = B_paths[index]
            B_path_recoder.append(B_path)
            B_coor = json.load(open(B_path))["people"]
            if len(B_coor)==0:
                pose = torch.zeros(17*2, 1)
            else:
                B_coor = B_coor[0]
                pose_dict = openpose_utils.obtain_2d_cords(B_coor, resize_param=self.load_size, org_size=self.org_size)
                pose = openpose_utils.openpose18_to_coco17(pose_dict['body'])
                pose = torch.Tensor(pose).float()
                pose = pose.view(17*2, 1)
                # normalize pose
                pose = 2*pose/self.load_size[1]-1

            pose_list.append(pose)
        B = torch.cat(pose_list, 1)
        return B, B_path_recoder

    def get_keypoint_param(self, n_frames_total, start_idx, t_step, cur_seq_len):
        end_id = start_idx + n_frames_total * t_step
        current = [i for i in range(start_idx, end_id, t_step)]

        pad_range = self.opt.n_receptive//2
        left_most = start_idx - pad_range*t_step

        left_pad = [i for i in range(left_most, start_idx, t_step)]
        left_pad = [0 if i<0 else i for i in left_pad]

        right_most = start_idx + (n_frames_total+pad_range)*t_step
        right_pad = [i for i in range(end_id, right_most, t_step)]
        right_pad = [cur_seq_len-1 if i>cur_seq_len-1 else i for i in right_pad]
        return left_pad, current, right_pad

    def update_seq_idx(self, A_paths, index):
        self.change_seq = self.frame_idx >= self.frames_count[self.seq_idx]
        if not self.change_seq and self.frame_idx + self.opt.n_frames_pre_load_test > self.frames_count[self.seq_idx]:
            self.frame_idx = self.frames_count[self.seq_idx]-self.opt.n_frames_pre_load_test
        if self.change_seq:
            self.seq_idx += 1
            self.frame_idx = self.opt.start_frame
        return None, None, None, self.seq_idx

    def __len__(self):
        return sum(self.frames_count)

    def name(self):
        return 'keypointDataset'


// Copyright
2019
Amazon.com, Inc. or its
affiliates.All
Rights
Reserved.
// SPDX - License - Identifier: LicenseRef -.amazon.com. - AmznSL - 1.0
// Licensed
under
the
Amazon
Software
License
http: // aws.amazon.com / asl /

import pdb
import math
import logging
import os.path
import pickle
import random
from typing import Any, Dict, List
from overrides import overrides
import numpy as np

import torch
import torch.nn.functional as F
from torch.nn.functional import nll_loss
from torch.nn import BCEWithLogitsLoss
from torch.nn import CrossEntropyLoss

from allennlp.data.token_indexers import SingleIdTokenIndexer, TokenIndexer
from allennlp.data.tokenizers import Token, Tokenizer, WordTokenizer
from allennlp.data.fields import Field, TextField, ArrayField
from allennlp.common.checks import check_dimensions_match
from allennlp.data import Token, Vocabulary, Instance
from allennlp.data.dataset import Batch
from allennlp.models.model import Model
from allennlp.modules import Seq2SeqEncoder, TimeDistributed, TokenEmbedder, TextFieldEmbedder, FeedForward, ScalarMix
from allennlp.modules.input_variational_dropout import InputVariationalDropout
from allennlp.modules.matrix_attention.linear_matrix_attention import LinearMatrixAttention
from allennlp.modules.seq2seq_encoders.pytorch_seq2seq_wrapper import PytorchSeq2SeqWrapper
from allennlp.modules.layer_norm import LayerNorm
from allennlp.nn import Activation
from allennlp.nn import InitializerApplicator, util
from allennlp.nn.util import logsumexp
from allennlp.tools import squad_eval
from allennlp.training.metrics import Average, BooleanAccuracy, CategoricalAccuracy
from allennlp.modules.elmo import batch_to_ids as elmo_batch_to_ids
from allennlp.modules.elmo import Elmo

from .accuracy import Accuracy
from . import dstqa_util

logger = logging.getLogger(__name__)

@Model.register("dstqa")
class DSTQA(Model):
    def __init__(self, vocab: Vocabulary,
                 base_dim,
                 loss_scale_by_num_values,
                 use_pre_calc_elmo_embeddings,
                 elmo_embedding_path,
                 domain_slot_list_path,
                 word_embeddings,
                 token_indexers: Dict[str, TokenIndexer],
                 text_field_embedder: TextFieldEmbedder,
                 text_field_char_embedder: TextFieldEmbedder,
                 symbol_embedder: TextFieldEmbedder,
                 phrase_layer: Seq2SeqEncoder,
                 class_prediction_layer: FeedForward,
                 span_prediction_layer: FeedForward,
                 span_start_encoder: FeedForward,
                 span_end_encoder: FeedForward,
                 span_label_predictor: FeedForward,
                 initializer: InitializerApplicator,
                 use_graph,
                 bi_dropout: float = 0.2,
                 dropout: float = 0.2) -> None:
        super().__init__(vocab)
        self._is_in_training_mode = False
        self._loss_scale_by_num_values = loss_scale_by_num_values
        self._use_pre_calc_elmo_embeddings = use_pre_calc_elmo_embeddings
        self._word_embeddings = word_embeddings
        self._is_use_elmo = True if self._word_embeddings == "elmo" else False
        self._is_use_graph = use_graph
        if self._is_use_elmo and use_pre_calc_elmo_embeddings:
            self._dialog_elmo_embeddings = self.load_elmo_embeddings(elmo_embedding_path)
            self._dialog_scalar_mix = ScalarMix(mixture_size=3, trainable=True)

        self._domains, self._ds_id2text, self._ds_text2id, self.value_file_path, \
        self._ds_type, self._ds_use_value_list, num_ds_use_value, self._ds_masked \
            = self.read_domain_slot_list(domain_slot_list_path)
        self._value_id2text, self._value_text2id = self.load_value_list(domain_slot_list_path)
        self._span_id2text, self._class_id2text = dstqa_util.gen_id2text(self._ds_id2text, self._ds_type)
        self._token_indexers = token_indexers

        self._text_field_embedder = text_field_embedder
        self._text_field_char_embedder = text_field_char_embedder
        self._symbol_embedder = symbol_embedder

        self._ds_dialog_attention = LinearMatrixAttention(base_dim, base_dim, 'x,y,x*y')
        self._dialog_dsv_attention = LinearMatrixAttention(base_dim, base_dim, 'x,y,x*y')
        self._dsv_dialog_attention = LinearMatrixAttention(base_dim, base_dim, 'x,y,x*y')
        self._ds_attention = LinearMatrixAttention(base_dim, base_dim, 'x,y,x*y')
        self._dsv_attention = LinearMatrixAttention(base_dim, base_dim, 'x,y,x*y')
        self._agg_value = torch.nn.Linear(base_dim, base_dim)
        self._agg_nodes = torch.nn.Linear(base_dim, base_dim)
        self._graph_gamma = torch.nn.Linear(base_dim, 1)
        self._class_prediction_layer = class_prediction_layer
        self._span_prediction_layer = span_prediction_layer
        self._span_label_predictor = span_label_predictor
        self._span_start_encoder = span_start_encoder
        self._span_end_encoder = span_end_encoder
        self._phrase_layer = phrase_layer
        self._cross_entropy = CrossEntropyLoss(ignore_index=-1)
        self._accuracy = Accuracy(self._ds_id2text, self._ds_type)
        self._dropout = torch.nn.Dropout(dropout)
        self._bi_dropout = torch.nn.Dropout(bi_dropout)
        self._dropout2 = torch.nn.Dropout(0.1)
        self._sigmoid = torch.nn.Sigmoid()
        initializer(self)

    def load_elmo_embeddings(self, elmo_embedding_path):
        elmo_embeddings = {}
        for suffix in ["train", "dev", "test"]:
            with open(elmo_embedding_path + suffix, "rb") as fp:
                elmo_embeddings.update(pickle.load(fp))
        return elmo_embeddings

    def gen_utt_masks(self, turn_offset, batch_size, max_turn_count, max_dialog_len):
        masks = torch.arange(0, max_dialog_len).unsqueeze(0).unsqueeze(0).cuda()
        masks = masks.repeat(batch_size, max_turn_count, 1)
        repeated_turn_offset = turn_offset.unsqueeze(2).repeat(1, 1, max_dialog_len)
        masks = masks < repeated_turn_offset
        # two types of masks: (1) all previous and current utt are marked as 1, (2) only current utt are marked as 1
        bmasks = masks.clone().detach()
        bmasks = (~bmasks)[:, :-1, :]
        cmasks = masks.clone().detach()
        cmasks[:, 1:, :] = cmasks[:, 1:, :] & bmasks
        return masks, cmasks

    def mix_dialog_embeddings(self, dialog_indices):
        dialog_embeddings = []
        max_dialog_len = 0
        for idx in dialog_indices:
            elmo_embeddings_cuda = []
            for v in self._dialog_elmo_embeddings[idx]:
                elmo_embeddings_cuda.append(v.cuda())
            dialog_embeddings.append(self._dialog_scalar_mix(elmo_embeddings_cuda))
            if max_dialog_len < dialog_embeddings[-1].size(0):
                max_dialog_len = dialog_embeddings[-1].size(0)
        for i, e in enumerate(dialog_embeddings):
            pad = torch.zeros(max_dialog_len - e.size(0), e.size(1)).cuda()
            dialog_embeddings[i] = torch.cat((e, pad), dim=0)
        dialog_embeddings = torch.stack(dialog_embeddings, dim=0)
        return dialog_embeddings

    def mask_time_step(self, dialogs, dialog_masks):
        batch_size, max_dialog_len, max_char_len = dialogs['token_characters'].size()
        masks = self._dropout2(torch.ones(batch_size, max_dialog_len))
        masks = masks < 0.5
        char_masked = torch.tensor([259, 260] + [0] * (max_char_len - 2)).cuda()
        char_padded = torch.tensor([0] * max_char_len).cuda()
        dialogs["token_characters"][masks] = char_masked
        dialogs["token_characters"][dialog_masks == 0] = char_padded

        if "tokens" in dialogs:
            dialogs["tokens"][masks] = 1  # 1 is the index for unknown
            dialogs["tokens"][dialog_masks == 0] = 0
        if "elmo" in dialogs:
            elmo_masked = torch.tensor([259, 260] + [261] * (50 - 2)).cuda()
            elmo_padded = torch.tensor([0] * 50).cuda()
            dialogs["elmo"][masks] = elmo_masked
            dialogs["elmo"][dialog_masks == 0] = elmo_padded

    def forward(self, dialogs, tags, utt_lens, exact_match, dialog_indices, epoch_num=None,
                labels=None, spans_start=None, spans_end=None, metadata=None, span_labels=None):
        self._is_in_training_mode = self.training
        # dialog embeddings
        batch_size, max_dialog_len, _ = dialogs['token_characters'].size()
        dialog_masks = util.get_text_field_mask(dialogs, num_wrapping_dims=0)
        self.mask_time_step(dialogs, dialog_masks)
        char_embedder_input = {'token_characters': dialogs['token_characters']}
        dialog_char_embeddings = self._text_field_char_embedder(char_embedder_input, num_wrapping_dims=0)
        if self._is_use_elmo:
            if self._use_pre_calc_elmo_embeddings == False:
                elmo_embedder_input = {'elmo': dialogs['elmo']}
                dialog_elmo_embeddings = self._text_field_embedder(elmo_embedder_input, num_wrapping_dims=0)
                dialog_embeddings = torch.cat((dialog_elmo_embeddings, dialog_char_embeddings), dim=2)
            else:
                dialog_elmo_embeddings = self.mix_dialog_embeddings(dialog_indices)
                dialog_embeddings = torch.cat((dialog_elmo_embeddings, dialog_char_embeddings), dim=2)
        else:
            embedder_input = {'tokens': dialogs['tokens']}
            dialog_elmo_embeddings = self._text_field_embedder(embedder_input, num_wrapping_dims=0)
            dialog_embeddings = torch.cat((dialog_elmo_embeddings, dialog_char_embeddings), dim=2)
        tag_embeddings = self._symbol_embedder(tags, num_wrapping_dims=0)
        turn_offset = torch.cumsum(utt_lens, dim=1)
        max_turn_count = utt_lens.size(1)
        context_masks, utt_masks = self.gen_utt_masks(turn_offset, batch_size, max_turn_count, max_dialog_len)

        # dsv embeddings
        ds_embeddings, v_embeddings = self.get_dsv_embeddings()

        # phrase layer
        merged_dialog_embeddings = torch.cat((dialog_embeddings, tag_embeddings, exact_match), dim=2)

        total_loss = 0.0
        predictions = []
        if self._is_in_training_mode == True:  # # only train one domain per turn for GPU memory limits
            sampled_turn = random.choice(list(range(max_turn_count)))
        for turn_i in range(max_turn_count):
            predictions.append(({}, {}))
            if self._is_in_training_mode == True and self._is_use_graph == False:
                if turn_i != sampled_turn:
                    continue
            if self._is_in_training_mode == True:
                if turn_i < sampled_turn:
                    self.set_module_to_eval()
                if turn_i > sampled_turn: break
            # compute new domain slot embeddings
            attention_ds_embeddings = None
            if turn_i > 0 and self._is_use_graph:
                attention_ds_embeddings = self.ds_graph_embeddings(batch_size, predictions[turn_i - 1], ds_embeddings,
                                                                   v_embeddings)
            repeated_ds_embeddings = ds_embeddings.unsqueeze(0).repeat(batch_size, 1, 1)
            reduced_dialog_masks = self._phrase_layer(self._dropout(merged_dialog_embeddings),
                                                      context_masks[:, turn_i, :])
            span_ds_i = 0
            for ds_i, ds_name in enumerate(self._ds_id2text):
                cur_repeated_ds_embeddings = repeated_ds_embeddings[:, ds_i, :].unsqueeze(1)
                cur_context_masks = context_masks[:, turn_i, :]
                if self._ds_type[ds_name] == "classification":
                    cur_labels = labels[:, turn_i, ds_i]
                    cur_v_embeddings = v_embeddings[ds_name]
                    loss, prediction = self.forward_classification(ds_name, reduced_dialog_masks,
                                                                   cur_repeated_ds_embeddings, cur_v_embeddings,
                                                                   cur_context_masks, cur_labels,
                                                                   attention_ds_embeddings)
                    predictions[turn_i][0][ds_name] = prediction
                    if self._loss_scale_by_num_values:
                        loss = loss * max(1.0, math.log(cur_v_embeddings.size(0)))
                elif self._ds_type[ds_name] == "span":
                    cur_span_labels = span_labels[:, turn_i, span_ds_i]
                    cur_spans_start = spans_start[:, turn_i, span_ds_i]
                    cur_spans_end = spans_end[:, turn_i, span_ds_i]
                    loss, prediction = self.forward_span(ds_name, reduced_dialog_masks, cur_repeated_ds_embeddings,
                                                         cur_context_masks, cur_span_labels, cur_spans_start,
                                                         cur_spans_end)
                    predictions[turn_i][1][ds_name] = prediction
                    span_ds_i += 1
                if self._is_in_training_mode == True and turn_i == sampled_turn:
                    if not self._ds_masked[ds_name]:
                        total_loss += loss
            if self._is_in_training_mode == True:
                if turn_i < sampled_turn:
                    self.set_module_to_train()
        output = {}
        if self._is_in_training_mode == True:
            output["loss"] = total_loss
        output["predictions"] = predictions
        output["metadata"] = metadata
        return output

    def set_module_to_eval(self):
        self.eval()
        self._phrase_layer.eval()
        self._class_prediction_layer.eval()
        self._span_prediction_layer.eval()
        self._span_start_encoder.eval()
        self._span_end_encoder.eval()
        self._span_label_predictor.eval()
        torch.set_grad_enabled(False)

    def set_module_to_train(self):
        self.train()
        self._phrase_layer.train()
        self._class_prediction_layer.train()
        self._span_prediction_layer.train()
        self._span_start_encoder.train()
        self._span_end_encoder.train()
        self._span_label_predictor.train()
        torch.set_grad_enabled(True)

    def bi_att(self, dialog_embeddings, dsv_embeddings, context_masks):
        batch_size, max_dialog_len = context_masks.size()
        num_values = dsv_embeddings.size(1)
        dialog_dsv_similarity = self._dialog_dsv_attention(self._bi_dropout(dialog_embeddings),
                                                           self._bi_dropout(dsv_embeddings))
        # attention on dsv
        dialog_dsv_att = util.masked_softmax(dialog_dsv_similarity.view(-1, num_values), None)
        dialog_dsv_att = dialog_dsv_att.view(batch_size, max_dialog_len, num_values)
        dialog_dsv = util.weighted_sum(dsv_embeddings, dialog_dsv_att)
        new_dialog_embeddings = dialog_embeddings + dialog_dsv

        # attention on dialog
        dsv_dialog_att = util.masked_softmax(
            dialog_dsv_similarity.transpose(1, 2).contiguous().view(-1, max_dialog_len),
            context_masks.unsqueeze(1).repeat(1, num_values, 1).view(-1, max_dialog_len))
        dsv_dialog_att = dsv_dialog_att.view(batch_size, num_values, max_dialog_len)
        dsv_dialog = util.weighted_sum(dialog_embeddings, dsv_dialog_att)
        new_dsv_embeddings = dsv_embeddings + dsv_dialog
        return new_dialog_embeddings, new_dsv_embeddings

    def forward_classification(self, ds_name, dialog_repr, ds_embeddings, value_embeddings, context_masks, labels=None,
                               attention_ds_embeddings=None):
        batch_size, max_dialog_len = context_masks.size()
        num_values = value_embeddings.size(0)
        repeated_dsv_embeddings = ds_embeddings.repeat(1, num_values, 1)
        repeated_dsv_embeddings += value_embeddings.unsqueeze(0).repeat(batch_size, 1, 1)
        dialog_repr, repeated_dsv_embeddings = self.bi_att(dialog_repr, repeated_dsv_embeddings, context_masks)

        ds_dialog_sim = self._ds_dialog_attention(self._bi_dropout(ds_embeddings), self._bi_dropout(dialog_repr))
        ds_dialog_att = util.masked_softmax(ds_dialog_sim.view(-1, max_dialog_len),
                                            context_masks.view(-1, max_dialog_len))
        ds_dialog_att = ds_dialog_att.view(batch_size, max_dialog_len)
        ds_dialog_repr = util.weighted_sum(dialog_repr, ds_dialog_att)

        if attention_ds_embeddings is not None:
            self_att_matrix = self._ds_attention(self._bi_dropout(ds_dialog_repr.unsqueeze(1)), attention_ds_embeddings)
            self_probs = util.masked_softmax(self_att_matrix, None)
            ret = util.weighted_sum(attention_ds_embeddings, self_probs).squeeze(1)
            gamma = torch.sigmoid(self._graph_gamma(ds_dialog_repr + ret))
            ds_dialog_repr = (1 - gamma) * ds_dialog_repr + gamma * ret

        w = self._class_prediction_layer(self._bi_dropout(ds_dialog_repr)).unsqueeze(1)
        logits = torch.bmm(w, repeated_dsv_embeddings.transpose(1, 2)).squeeze(1)
        prediction = torch.argmax(logits, dim=1)
        loss = self._cross_entropy(logits.view(-1, num_values), labels.view(-1))
        if labels is not None:
            self._accuracy.value_acc(ds_name, logits, labels, labels != -1)
        return loss, prediction

    def forward_span(self, ds_name, dialog_repr, repeated_ds_embeddings, context_masks, span_labels=None,
                     spans_start=None, spans_end=None):
        batch_size, max_dialog_len = context_masks.size()
        ds_dialog_sim = self._ds_dialog_attention(self._dropout(repeated_ds_embeddings), self._dropout(dialog_repr))
        ds_dialog_att = util.masked_softmax(ds_dialog_sim.view(-1, max_dialog_len),
                                            context_masks.view(-1, max_dialog_len))
        ds_dialog_att = ds_dialog_att.view(batch_size, max_dialog_len)
        ds_dialog_repr = util.weighted_sum(dialog_repr, ds_dialog_att)
        ds_dialog_repr = ds_dialog_repr + repeated_ds_embeddings.squeeze(1)
        span_label_logits = self._span_label_predictor(F.relu(self._dropout(ds_dialog_repr)))
        span_label_prediction = torch.argmax(span_label_logits, dim=1)
        span_label_loss = 0.0
        if span_labels is not None:
            span_label_loss = self._cross_entropy(span_label_logits, span_labels)  # loss averaged by #turn
            self._accuracy.span_label_acc(ds_name, span_label_logits, span_labels, span_labels != -1)
        loss = span_label_loss

        w = self._span_prediction_layer(self._dropout(ds_dialog_repr)).unsqueeze(1)
        span_start_repr = self._span_start_encoder(self._dropout(dialog_repr))
        span_start_logits = torch.bmm(w, span_start_repr.transpose(1, 2)).squeeze(1)
        span_start_probs = util.masked_softmax(span_start_logits, context_masks)
        span_start_logits = util.replace_masked_values(span_start_logits, context_masks.to(dtype=torch.int8), -1e7)

        span_end_repr = self._span_end_encoder(self._dropout(span_start_repr))
        span_end_logits = torch.bmm(w, span_end_repr.transpose(1, 2)).squeeze(1)
        span_end_probs = util.masked_softmax(span_end_logits, context_masks)
        span_end_logits = util.replace_masked_values(span_end_logits, context_masks.to(dtype=torch.int8), -1e7)

        best_span = self.get_best_span(span_start_logits, span_end_logits)
        best_span = best_span.view(batch_size, -1)

        spans_loss = 0.0
        if spans_start is not None:
            spans_loss = self._cross_entropy(span_start_logits, spans_start)
            self._accuracy.span_start_acc(ds_name, span_start_logits, spans_start, spans_start != -1)
            spans_loss += self._cross_entropy(span_end_logits, spans_end)
            self._accuracy.span_end_acc(ds_name, span_end_logits, spans_end, spans_end != -1)
        loss += spans_loss

        return loss, (span_label_prediction, best_span)

    @overrides
    def decode(self, output_dict):
        num_turns = len(output_dict["predictions"])
        class_output = []
        for t in range(num_turns):
            class_predictions = output_dict["predictions"][t][0]
            res = []
            for ds_name, pred in class_predictions.items():
                value = self._value_id2text[ds_name][pred.item()]
                res.append(ds_name + ":" + value)
            class_output.append(res)

        span_output = []
        for t in range(num_turns):
            span_predictions = output_dict["predictions"][t][1]
            res = []
            for ds_name, pred in span_predictions.items():
                span_label = pred[0]
                if span_label == 0: value = "none"
                if span_label == 1: value = "dont care"
                if span_label == 2:
                    start, end = pred[1][0][0], pred[1][0][1]
                    value = " ".join([output_dict["metadata"][0][i].text for i in range(start, end + 1)])
                    value = value.lower()
                res.append(ds_name + ":" + value)
            span_output.append(res)
        # merge class output and span output
        output = []
        if len(span_output) != 0 and len(class_output) != 0:
            for x, y in zip(class_output, span_output):
                output.append(x + y)
        elif len(span_output) == 0:
            output = class_output
        elif len(class_output) == 0:
            output = span_output
        else:
            assert (False)
        output_dict["predicted_labels"] = [output]
        del output_dict["metadata"]
        del output_dict["predictions"]
        return output_dict

    def get_metrics(self, reset=False):
        acc = self._accuracy.get_metrics(reset)
        return acc

    def get_dsv_embeddings(self):
        def batch_to_id(batch: List[List[str]]):
            instances = []
            for b in batch:
                tokens = [Token(w) for w in b.split(" ")]
                field = TextField(tokens, self._token_indexers)
                instance = Instance({"b": field})
                instances.append(instance)
            dataset = Batch(instances)
            vocab = self.vocab
            dataset.index_instances(vocab)
            res = {}
            for k, v in dataset.as_tensor_dict()['b'].items():
                res[k] = v.cuda()
            return res

        ds_ids = batch_to_id(self._ds_id2text)
        if 'tokens' in ds_ids:
            elmo_embedder_input = {'tokens': ds_ids['tokens']}
        elif 'elmo' in ds_ids:
            elmo_embedder_input = {'elmo': ds_ids['elmo']}
        ds_elmo_embeddings = self._text_field_embedder(elmo_embedder_input, num_wrapping_dims=0).sum(1)
        char_embedder_input = {'token_characters': ds_ids['token_characters']}
        ds_char_embeddings = self._text_field_char_embedder(char_embedder_input, num_wrapping_dims=0).sum(1)
        ds_embeddings = torch.cat((ds_elmo_embeddings, ds_char_embeddings), dim=1)
        ds_masks = util.get_text_field_mask(ds_ids, num_wrapping_dims=0).sum(1).float()
        ds_embeddings = ds_embeddings / ds_masks.unsqueeze(1).repeat(1, ds_embeddings.size(1))
        v_embeddings = {}
        for v, v_list in self._value_id2text.items():
            v_ids = batch_to_id(v_list)
            if 'tokens' in v_ids:
                elmo_embedder_input = {'tokens': v_ids['tokens']}
            elif 'elmo' in v_ids:
                elmo_embedder_input = {'elmo': v_ids['elmo']}
            v_elmo_embeddings = self._text_field_embedder(elmo_embedder_input, num_wrapping_dims=0).sum(1)
            char_embedder_input = {'token_characters': v_ids['token_characters']}
            v_char_embeddings = self._text_field_char_embedder(char_embedder_input, num_wrapping_dims=0).sum(1)
            v_embeddings[v] = torch.cat((v_elmo_embeddings, v_char_embeddings), dim=1)
            v_masks = util.get_text_field_mask(v_ids, num_wrapping_dims=0).sum(1).float()
            v_embeddings[v] = v_embeddings[v] / v_masks.unsqueeze(1).repeat(1, v_embeddings[v].size(1))
        return ds_embeddings, v_embeddings

    def read_domain_slot_list(self, filename):
        with open(filename) as fp:
            lines = fp.readlines()
        domains = []
        domain_slots = []
        value_file_path = {}
        domain_slots_type = {}
        domain_slots_use_value_list = {}
        ds_masked = {}
        num_ds_use_value = 0
        for line in lines:
            line = line.strip("\n ")
            if line.startswith("#"):
                continue
            if len(line.strip("\n ")) == 0:
                continue
            line_arr = line.split("\t")
            ds = line_arr[0] + " " + line_arr[1]
            if line_arr[3] == "n":
                domains.append(line_arr[0])
                domain_slots.append(ds)
                value_file_path[ds] = line_arr[4].strip(" \n")
                domain_slots_type[ds] = line_arr[2]
                domain_slots_use_value_list[ds] = True if line_arr[5] == "y" else False
                num_ds_use_value += 1 if line_arr[5] == "y" else 0
                ds_masked[ds] = True if line_arr[6] == "y" else False
        ds_text2id = {}
        for i, s in enumerate(domain_slots):
            ds_text2id[s] = i
        return domains, domain_slots, ds_text2id, value_file_path, domain_slots_type, domain_slots_use_value_list, num_ds_use_value, ds_masked

    def load_value_list(self, ds_path):
        def read_value_list(ds_path, ds, value_path_list):
            dir_path = os.path.dirname(ds_path)
            filename = dir_path + "/" + value_path_list[ds]
            with open(filename) as fp:
                lines = fp.readlines()
            values = []
            for line_i, line in enumerate(lines):
                if len(line.strip("\n ")) == 0:
                    continue
                values.append(line.strip("\n "))
            value2id = {}
            for i, v in enumerate(values):
                value2id[v] = i
            return values, value2id

        value_text2id = {}
        value_id2text = {}
        for ds in self._ds_text2id.keys():
            if not self._ds_use_value_list[ds]: continue
            id2v, v2id = read_value_list(ds_path, ds, self.value_file_path)
            value_text2id[ds] = v2id
            value_id2text[ds] = id2v
        return value_id2text, value_text2id

    # code from https://github.com/allenai/allennlp/blob/master/allennlp/models/reading_comprehension/bidaf.py
    def get_best_span(self, span_start_logits, span_end_logits):
        # We call the inputs "logits" - they could either be unnormalized logits or normalized log
        # probabilities.  A log_softmax operation is a constant shifting of the entire logit
        # vector, so taking an argmax over either one gives the same result.
        if span_start_logits.dim() != 2 or span_end_logits.dim() != 2:
            raise ValueError("Input shapes must be (batch_size, passage_length)")
        batch_size, passage_length = span_start_logits.size()
        device = span_start_logits.device
        # (batch_size, passage_length, passage_length)
        span_log_probs = span_start_logits.unsqueeze(2) + span_end_logits.unsqueeze(1)
        # Only the upper triangle of the span matrix is valid; the lower triangle has entries where
        # the span ends before it starts.
        span_log_mask = torch.triu(torch.ones((passage_length, passage_length),
                                              device=device)).log().unsqueeze(0)
        valid_span_log_probs = span_log_probs + span_log_mask

        # Here we take the span matrix and flatten it, then find the best span using argmax.  We
        # can recover the start and end indices from this flattened list using simple modular
        # arithmetic.
        # (batch_size, passage_length * passage_length)
        best_spans = valid_span_log_probs.view(batch_size, -1).argmax(-1)
        span_start_indices = best_spans // passage_length
        span_end_indices = best_spans % passage_length
        return torch.stack([span_start_indices, span_end_indices], dim=-1)

    def ds_graph_embeddings(self, batch_size, predictions, ds_embeddings, v_embeddings):
        repeated_ds_embeddings = ds_embeddings.unsqueeze(0).repeat(batch_size, 1, 1)
        for node_i, node in enumerate(self._ds_id2text):
            if not self._ds_use_value_list[node]: continue
            val_node = v_embeddings[node][predictions[0][node]]
            ds_node = repeated_ds_embeddings[:, self._ds_text2id[node], :]
            ds_node = ds_node + val_node
            repeated_ds_embeddings = repeated_ds_embeddings.clone()
            repeated_ds_embeddings[:, self._ds_text2id[node], :] = ds_node
        return repeated_ds_embeddings


import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision import models
from torch.autograd import Variable

from net import vgg16, vgg16_bn
from resnet_yolo import resnet50, resnet18
from dataset import yoloDataset
from yololoss import Loss

from visualize import Visualizer
import numpy as np
# import warnings
# warnings.filterwarnings("ignore", category=Warning)
# torch.cuda.current_device()
# torch.cuda._initialized = True

use_gpu = torch.cuda.is_available()

# file_root = '/Users/mr.xue/Desktop/code/pytorch-YOLO-v1/tiny_image/'

# train_dataset = yoloDataset(root=file_root,list_file=['voc2012_tiny.txt'],train=True,transform = [transforms.ToTensor()] )
# train_loader = DataLoader(train_dataset,batch_size=1,shuffle=True,num_workers=0)

learning_rate = 0.001
batch_size = 1
use_resnet = False
if use_resnet:
    net = resnet50()
else:
    net = vgg16_bn()

# print(net)

print('load pre-trined model')
if use_resnet:
    resnet = models.resnet50(pretrained=True)
    new_state_dict = resnet.state_dict()
    dd = net.state_dict()
    for k in new_state_dict.keys():
        print(k)
        if k in dd.keys() and not k.startswith('fc'):
            print('yes')
            dd[k] = new_state_dict[k]
    net.load_state_dict(dd)
else:
    vgg = models.vgg16_bn(pretrained=True)
    new_state_dict = vgg.state_dict()
    dd = net.state_dict()
    for k in new_state_dict.keys():
        # print(k)
        if k in dd.keys() and k.startswith('features'):
            # print('yes')
            dd[k] = new_state_dict[k]
    net.load_state_dict(dd)
if False:
    net.load_state_dict(torch.load('best.pth'))
print('cuda', torch.cuda.current_device(), torch.cuda.device_count())

if use_gpu:
    net.cuda()

net.train()
# different learning rate
params=[]
params_dict = dict(net.named_parameters())
for key,value in params_dict.items():
    if key.startswith('features'):
        params += [{'params':[value],'lr':learning_rate*1}]
    else:
        params += [{'params':[value],'lr':learning_rate}]

optimizer = torch.optim.SGD(params, lr=learning_rate, momentum=0.9, weight_decay=5e-4)
# optimizer = torch.optim.Adam(net.parameters(),lr=learning_rate,weight_decay=1e-4)

# file_root = '/Datas/yc/VOC_0712/VOCdevkit/VOC2012/JPEGImages/'
file_root = './tiny_database/train_images/'
# file_root = '/home/ubuntu/voc2012/VOCtrainval_11-May-2012/VOCdevkit/voc2012_tiny_test/images/'

# train_dataset = yoloDataset(root=file_root,list_file=['voc2012_tiny.txt'],train=True,transform = [transforms.ToTensor()] )
train_dataset = yoloDataset(root=file_root,list_file=['voc2012train.txt'],train=True,transform = [transforms.ToTensor()] )
train_loader = DataLoader(train_dataset,batch_size=1,shuffle=True,num_workers=0)
# test_dataset = yoloDataset(root=file_root,list_file='voc2007test.txt',train=False,transform = [transforms.ToTensor()] )
# test_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=False,num_workers=4)
print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
# logfile = open('log.txt', 'w')

num_iter = 0
num_epochs = 10
#vis = Visualizer(env='xiong')
# best_test_loss = np.inf

criterion = Loss(5, 0.5)  # lambda_coord, lambda_noobj
# f = open('result.txt','a',encoding = 'utf-8')

for epoch in range(num_epochs):
    net.train()
    # if epoch == 1:
    #     learning_rate = 0.0005
    # if epoch == 2:
    #     learning_rate = 0.00075
    # if epoch == 3:
    #     learning_rate = 0.001
    if epoch == 30:
        learning_rate=0.0001
    if epoch == 40:
        learning_rate=0.00001
    # optimizer = torch.optim.SGD(net.parameters(),lr=learning_rate*0.1,momentum=0.9,weight_decay=1e-4)
    for param_group in optimizer.param_groups:
        param_group['lr'] = learning_rate

    print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
    print('Learning Rate for this epoch: {}'.format(learning_rate))

    total_loss = 0.

    for i,(images,target) in enumerate(train_loader):
        # print("target", target.shape)
        # print("image", images.shape)
        images = Variable(images)
        target = Variable(target)
        if use_gpu:
            images,target = images.cuda(),target.cuda()

        pred = net(images)
        # print("pred is ok", pred.shape)
        loss = criterion(pred,target)
        # print("loss:", loss, loss.shape)
        total_loss += loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f'
        #       % (epoch + 1, num_epochs, i + 1, len(train_loader), loss, total_loss / (i + 1)))
        if (i+1) % 10 == 0:
            print ('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f'
            %(epoch+1, num_epochs, i+1, len(train_loader), loss, total_loss / (i+1)))
        #     num_iter += 1
        #     # vis.plot_train_val(loss_train=total_loss/(i+1))

    print('all: Epoch [%d/%d],  Loss: %.4f, average_loss: %.4f'
              % (epoch + 1, num_epochs, loss, total_loss / len(train_loader)))
    # f.write(str(epoch+1)+" " + str(total_loss / len(train_loader)) + "\n")

    # #validation
    # validation_loss = 0.0
    # net.eval()
    # for i,(images,target) in enumerate(test_loader):
    #     images = Variable(images,volatile=True)
    #     target = Variable(target,volatile=True)
    #     if use_gpu:
    #         images,target = images.cuda(),target.cuda()

    #     pred = net(images)
    #     loss = criterion(pred,target)
    #     validation_loss += loss.data[0]
    # validation_loss /= len(test_loader)
    # vis.plot_train_val(loss_val=validation_loss)

    # if best_test_loss > validation_loss:
    #     best_test_loss = validation_loss
    #     print('get best test loss %.5f' % best_test_loss)
    #     torch.save(net.state_dict(),'best.pth')
    # logfile.writelines(str(epoch) + '\t' + str(validation_loss) + '\n')
    # logfile.flush()
    # torch.save(net.state_dict(),'yolo.pth')

# f.close()


# vim: set fileencoding=<utf-8> :
# Copyright 2018-2020 John Lees and Nick Croucher

'''Functions for construction and processing of trees'''

import sys
import os
import subprocess
import numpy as np
import pandas as pd
import dendropy

from .utils import check_and_set_gpu

# Load GPU libraries
try:
    import cupyx
    import cugraph
    import cudf
    import cupy as cp
    from numba import cuda
    import rmm
    gpu_lib = True
except ImportError as e:
    gpu_lib = False

def buildRapidNJ(rapidnj, refList, coreMat, outPrefix, threads = 1):
    """Use rapidNJ for more rapid tree building

    Creates a phylip of core distances, system call to rapidnj executable, loads tree as
    dendropy object (cleaning quotes in node names), removes temporary files.

    Args:
        rapidnj (str)
            Location of rapidnj executable
        refList (list)
            Names of sequences in coreMat (same order)
        coreMat (numpy.array)
            NxN core distance matrix produced in :func:`~outputsForMicroreact`
        outPrefix (str)
            Prefix for all generated output files, which will be placed in `outPrefix` subdirectory
        threads (int)
            Number of threads to use

    Returns:
        tree (str)
            Newick-formatted NJ tree from core distances
    """
    # generate phylip matrix
    phylip_name = outPrefix + "/" + os.path.basename(outPrefix) + "_core_distances.phylip"
    with open(phylip_name, 'w') as pFile:
        pFile.write(str(len(refList))+"\n")
        for coreDist, ref in zip(coreMat, refList):
            pFile.write(ref)
            pFile.write(' '+' '.join(map(str, coreDist)))
            pFile.write("\n")

    # construct tree
    tree_filename = outPrefix + "/" + os.path.basename(outPrefix) + "_core_NJ.nwk"
    rapidnj_cmd = rapidnj + " " + phylip_name + " -n -i pd -o t -x " + tree_filename + ".raw -c " + str(threads)
    try:
        # run command
        subprocess.run(rapidnj_cmd, shell=True, check=True)

        # remove quotation marks for microreact
        with open(tree_filename + ".raw", 'r') as f, open(tree_filename, 'w') as fo:
            for line in f:
                fo.write(line.replace("'", ''))
        # tidy unnecessary files
        os.remove(tree_filename+".raw")
        os.remove(phylip_name)

    # record errors
    except subprocess.CalledProcessError as e:
        sys.stderr.write("Could not run command " + rapidnj_cmd + "; returned code: " + str(e.returncode) + "\n")
        sys.exit(1)

    # read tree and return
    tree = dendropy.Tree.get(path=tree_filename, schema="newick", preserve_underscores=True)
    os.remove(tree_filename)
    return tree

def write_tree(tree, prefix, suffix, overwrite):
    """Prints a Newick-formatted string to an output file

    Args:
        tree (str)
            Newick-formatted string representation of tree
        prefix (str)
            Prefix for output files
        suffix (str)
            Suffix for output files
        overwrite (bool)
            Whether to overwrite existing files
    """
    tree_filename = prefix + "/" + os.path.basename(prefix) + suffix
    if overwrite or not os.path.isfile(tree_filename):
        with open(tree_filename, 'w') as tree_file:
            tree_file.write(tree)
    else:
        sys.stderr.write("Unable to write phylogeny to " + tree_filename + "\n")

def load_tree(prefix, type, distances = 'core'):
    """Checks for existing trees from previous runs.

    Args:
        prefix (str)
            Output prefix used for search
        type (str)
            Type of tree (NJ or MST)

    Returns:
        tree_string (str)
            Newick-formatted string of NJ tree
    """
    tree_string = None
    tree_prefix = os.path.join(prefix,os.path.basename(prefix))
    for suffix in ['_' + distances + '_' + type + ".tree",'_' + distances + '_' + type + ".nwk"]:
            tree_fn = tree_prefix + suffix
            if os.path.isfile(tree_fn):
                sys.stderr.write("Reading existing tree from " + tree_fn + "\n")
                tree = dendropy.Tree.get(path=tree_fn, schema="newick", preserve_underscores=True)
                tree_string = tree.as_string(schema="newick",
                suppress_rooting=True,
                unquoted_underscores=True)
                break

    return tree_string

def generate_nj_tree(coreMat, seqLabels, outPrefix, rapidnj, threads):
    """Generate phylogeny using dendropy or RapidNJ

    Writes a neighbour joining tree (.nwk) from core distances.

    Args:
        coreMat (numpy.array)
            n x n array of core distances for n samples.
        seqLabels (list)
            Processed names of sequences being analysed.
        outPrefix (str)
            Output prefix for core distances file
        rapidnj (str)
            A string with the location of the rapidnj executable for tree-building. If None, will
            use dendropy by default
        threads (int)
            Number of threads to use with rapidnj

    Returns:
        tree_string (str)
            Newick-formatted string of NJ tree
    """
    # Save distances to file
    core_dist_file = outPrefix + "/" + os.path.basename(outPrefix) + "_core_dists.csv"
    np.savetxt(core_dist_file, coreMat, delimiter=",", header = ",".join(seqLabels), comments="")

    # calculate phylogeny
    sys.stderr.write("Building phylogeny\n")
    if rapidnj is not None:
        tree = buildRapidNJ(rapidnj, seqLabels, coreMat, outPrefix, threads = threads)
    else:
        pdm = dendropy.PhylogeneticDistanceMatrix.from_csv(src=open(core_dist_file),
                                                           delimiter=",",
                                                           is_first_row_column_names=True,
                                                           is_first_column_row_names=False)
        tree = pdm.nj_tree()

    # Midpoint root tree and write outout
    tree.reroot_at_midpoint(update_bipartitions=True, suppress_unifurcations=False)

    # remove file as it can be large
    os.remove(core_dist_file)

    # return Newick string
    tree_string = tree.as_string(schema="newick",
                                 suppress_rooting=True,
                                 unquoted_underscores=True)
    tree_string = tree_string.replace("'","")

    return tree_string

def mst_to_phylogeny(mst_network, names, use_gpu = False):
    """Convert a MST graph to a phylogeny

    Args:
       mst_network (network)
           Minimum spanning tree from
           :func:`~PopPUNK.network.generate_minimum_spanning_tree`
       names (list)
            Names of sequences in network
       use_gpu (bool)
            Whether to use GPU-specific libraries for processing

    Returns:
       mst_network (str)
           Minimum spanning tree (as graph-tool graph)
    """
    #
    # MST graph -> phylogeny
    #

    use_gpu = check_and_set_gpu(use_gpu, gpu_lib)

    # Define sequences names for tree
    taxon_namespace = dendropy.TaxonNamespace(names)
    # Initialise tree and create nodes
    tree = dendropy.Tree(taxon_namespace=taxon_namespace)

    # Identify edges
    if use_gpu:
        tree_nodes = {v:dendropy.Node(taxon=taxon_namespace[int(v)]) \
            for v in range(0,mst_network.number_of_vertices())}
        mst_edges_df = mst_network.view_edge_list()
    else:
        tree_nodes = {v:dendropy.Node(taxon=taxon_namespace[int(v)]) for v in mst_network.get_vertices()}
        mst_edges_df = pd.DataFrame(mst_network.get_edges(),
                                    columns = ['src', 'dst'])
        mst_edges_df['weights'] = list(mst_network.ep['weight'])
    seed_node_index = int(mst_edges_df[['src','dst']].stack().mode()[0])
    tree.seed_node = tree_nodes[seed_node_index]

    # Generate links of tree
    parent_node_indices = [seed_node_index]
    added_nodes = set(parent_node_indices)
    i = 0
    while i < len(parent_node_indices): # NB loop end will increase
        # select links from data frames
        mst_links_dst = mst_edges_df[['dst','weights']].loc[mst_edges_df['src']==parent_node_indices[i]]
        mst_links_src = mst_edges_df[['src','weights']].loc[mst_edges_df['dst']==parent_node_indices[i]]
        mst_links_src.columns = ['dst','weights']
        mst_links = mst_links_dst.append(mst_links_src)
        if use_gpu:
            mst_links = mst_links.to_pandas()
        for (child_node,edge_length) in mst_links.itertuples(index=False, name=None):
            if child_node not in added_nodes:
                tree_nodes[parent_node_indices[i]].add_child(tree_nodes[child_node])
                tree_nodes[child_node].edge_length = edge_length
                added_nodes.add(child_node)
                parent_node_indices.append(child_node)
        i = i + 1

    # Add zero length branches for internal nodes in MST
    for node in tree.preorder_node_iter():
        if not node.is_leaf():
            new_child = dendropy.Node(taxon=node.taxon, edge_length=0.0)
            node.taxon = None
            node.add_child(new_child)

    # Return tree as string
    tree_string = tree.as_string(schema="newick",
                                 suppress_rooting=True,
                                 unquoted_underscores=True)

    return tree_string