import numpy as np
import pymc3 as pm
import theano as th
import theano.tensor as tt
import inspect
import io

import logging
import time
import tqdm
import argparse
from .. models.fancy_model import GeneralizedContinuousModel
from typing import List, Callable, Optional, Set, Tuple, Any, Dict
from abc import abstractmethod
from ..inference.convergence_tracker import NoisyELBOConvergenceTracker
from ..inference.param_tracker import VariationalParameterTrackerConfig, VariationalParameterTracker
from ..inference import fancy_optimizers
from ..inference.deterministic_annealing import ADVIDeterministicAnnealing
from .. import types

_logger = logging.getLogger(__name__)

class Sampler:
    """Base class for log emission posterior probability samplers to be used in the hybrid ADVI scheme."""
    def __init__(self, hybrid_inference_params: 'HybridInferenceParameters'):
        self.hybrid_inference_params = hybrid_inference_params

    @abstractmethod
    def update_approximation(self, approx: pm.approximations.MeanField) -> None:
        """Take a new mean-field approximation and update the sampler routine accordingly.

        Args:
            approx: an instance of PyMC3 mean-field posterior approximation

        Returns:
            None
        """
        raise NotImplementedError

    @abstractmethod
    def draw(self) -> np.ndarray:
        """Draw one sample (or average of several samples) from log emission posterior probability."""
        raise NotImplementedError

    @abstractmethod
    def reset(self):
        """Reset the internal state of the sampler (e.g. previously accumulated samples)."""
        raise NotImplementedError

    @abstractmethod
    def increment(self, update: np.ndarray):
        """Add an incremental update to the current estimate of the log emission posterior mean."""
        raise NotImplementedError

    @abstractmethod
    def get_latest_log_emission_posterior_mean_estimator(self) -> np.ndarray:
        """Returns the latest estimate of the log emission posterior mean."""
        raise NotImplementedError

class Caller:
    """Base class for callers, i.e. routines that update the posterior of discrete RVs, to be used in the
    hybrid ADVI scheme."""
    @abstractmethod
    def call(self) -> 'CallerUpdateSummary':
        """Update the posterior of discrete RVs and return a summary"""
        raise NotImplementedError

    @abstractmethod
    def update_auxiliary_vars(self) -> None:
        """Update auxiliary variables in workspaces, if any. This routine is called after one
        (or more) round(s) of invoking `Caller.call`."""
        raise NotImplemented

class CallerUpdateSummary:
    """Represents a summary of updates made to discrete RV posteriors by a `Caller`."""
    @abstractmethod
    def reduce_to_scalar(self) -> float:
        """A function that reduces arrays to scalars. It is used to summarize tensor-valued updates."""
        raise NotImplementedError

    @abstractmethod
    def __repr__(self):
        """Represents the caller update summary in a human readable format (used in logging)."""
        raise NotImplementedError

class LoggerTQDMAdapter(io.StringIO):
    """Output stream for `tqdm` which will output to logger module instead of `stderr`."""
    logger = None
    level = None
    buf = ''

    def __init__(self, logger, level=logging.INFO):
        super().__init__()
        self.logger = logger
        self.level = level

    def write(self, buf):
        self.buf = buf.strip('\r\n\t ')

    def flush(self):
        self.logger.log(self.level, self.buf)

class InferenceTask:
    """Base class of all inference tasks."""

    # Lay in a course for starbase one two warp nine point five...
    @abstractmethod
    def engage(self):
        """Initiate inference algorithm."""
        raise NotImplementedError("Core breach imminent!")

    @abstractmethod
    def disengage(self):
        """Wrap up the inference algorithm (clean up workspaces, etc.)"""
        raise NotImplementedError

class HybridInferenceTask(InferenceTask):
    """The hybrid inference framework is applicable to PGMs with the following general structure:

        +--------------+           +----------------+
        | discrete RVs + --------► + continuous RVs |
        +--------------+           +----------------+

    Note that discrete RVs do not have any continuous parents. The inference is approximately
    performed by factorizing the true posterior into an uncorrelated product of discrete RVs (DRVs)
    and continuous RVs (CRVs):

        p(CRVs, DRVs | observed) ~ q(CRVs) q(DRVs)

    q(CRVs) is updated via deterministic annealing mean-field ADVI.
    q(DRV) is updated by the provided "caller" and is out the scope of this class.

    Usage:
    ------

    Preliminaries. Let us decompose the log joint as follows:

        -log_P(CRVs, DRVs, observed) = F_c(CRVs, observed)
                                        + F_d(DRVs, observed)
                                            + F_cd(CRVs, DRVs, observed)

    The last term in the free energy (negative log joint) is the only term with cross terms between
    the discrete and continuous sectors.

    The user must supply the following components:

        (1) a pm.Model that yields the DRV-posterior-expectation of the free energy,

            F_c^{eff}(CRVs, observed) = E_{DRVs ~ q(DRVs)} [-log_P(CRVs, DRVs, observed)]
                                        = F_c(CRVs, observed)
                                            + E_{DRVs ~ q(DRVs)} [F_cd(CRVs, DRVs, observed)]
                                                + E_{DRVs ~ q(DRVs)} [F_d(DRVs, observed)]

            Note: the last term is fully determined by q(DRVs) and can be dropped while performing
            ADVI updates in the continuous sector.

        (2) a "sampler" that provides samples from the cross term, which we call "log emission",
            defined as:

            log_emission(DRVs) = E_{CRVs ~ q(CRVs)} [-F_cd (CRVs, DRV, observed)]

        (3) a "caller" that updates q(DRVs) given log_emission(DRV), i.e.:

            q(DRVs) \propto \exp[log_emission(DRVs) - F_d(DRVs, observed)]

            In practice, one does not need the complete joint posterior of DRVs: only sufficient
            statistics, to the extent required for calculating F_c^{eff} is needed. Calculating such
            sufficient statistics could be as simple as using the Bayes rule, or more complicated if
            the DRVs are strongly correlated.

    The general implementation motif is:

        (a) to store sufficient statistics from q(DRVs) as a shared theano tensor such that the the
            model can access it,
        (b) to store log_emission(DRVs) as a shared theano tensor (or ndarray) such that the caller
            can access it, and:
        (c) let the caller directly update the shared sufficient statistics.
    """

    # if the temperature is within the following tolerance of 1.0, it is assumed that annealing
    # has finished
    temperature_tolerance = 1e-6

    def __init__(self,
                 hybrid_inference_params: 'HybridInferenceParameters',
                 continuous_model: GeneralizedContinuousModel,
                 sampler: Optional[Sampler],
                 caller: Optional[Caller],
                 **kwargs):
        """Initializer.

        Args:
            hybrid_inference_params: inference configuration
            continuous_model: a PyMC3 model representing the continuous sector of the PGM
            sampler: log emission probability sampler
            caller: discrete RV posterior updater
            **kwargs: extra keywords

        Keyword Args:
            custom_optimizer: a custom stochastic optimizer to be used in place of the default optimizer (adamax
            elbo_normalization_factor: normalization factor of the full model ELBO (for logging)
            advi_task_name: name of the ADVI step (for logging)
            sampling_task_name: name of the sampling step (for logging)
            calling_task_name: name of the calling step (for logging)
        """
        assert hybrid_inference_params is not None
        self.hybrid_inference_params = hybrid_inference_params

        assert continuous_model is not None
        self.continuous_model = continuous_model
        self.continuous_model.verify_var_registry()

        if sampler is None:
            _logger.warning("No log emission sampler given; skipping the sampling step")
        self.sampler = sampler

        if caller is None:
            _logger.warning("No caller given; skipping the calling step")
        self.caller = caller

        if self.hybrid_inference_params.track_model_params:
            _logger.info("Instantiating the parameter tracker...")
            self.param_tracker = self._create_param_tracker()
        else:
            self.param_tracker = None

        _logger.info("Instantiating the convergence tracker...")
        self.advi_convergence_tracker = NoisyELBOConvergenceTracker(
            self.hybrid_inference_params.convergence_snr_averaging_window,
            self.hybrid_inference_params.convergence_snr_trigger_threshold,
            self.hybrid_inference_params.convergence_snr_countdown_window)

        _logger.info("Setting up DA-ADVI...")
        with self.continuous_model:
            if not hasattr(self, 'temperature'):
                initial_temperature = self.hybrid_inference_params.initial_temperature
                self.temperature: types.TensorSharedVariable = th.shared(
                    np.asarray([initial_temperature], dtype=types.floatX))
            initial_temperature = self.temperature.get_value()[0]
            if (np.abs(initial_temperature - 1.0) < self.temperature_tolerance or
                    hybrid_inference_params.disable_annealing):
                # no annealing
                temperature_update = None
            else:
                # linear annealing
                max_thermal_advi_iterations = self.hybrid_inference_params.max_advi_iter_first_epoch
                max_thermal_advi_iterations += ((self.hybrid_inference_params.num_thermal_epochs - 1)
                                                * self.hybrid_inference_params.max_advi_iter_first_epoch)
                temperature_drop_per_iter = (initial_temperature - 1.0) / max_thermal_advi_iterations
                temperature_update = [(self.temperature,
                                       tt.maximum(1.0, self.temperature - temperature_drop_per_iter))]

            self.continuous_model_advi = ADVIDeterministicAnnealing(
                random_seed=self.hybrid_inference_params.random_seed,
                temperature=self.temperature)
            self.continuous_model_approx: pm.MeanField = self.continuous_model_advi.approx

            if 'custom_optimizer' in kwargs.keys():
                opt = kwargs['custom_optimizer']
                assert issubclass(type(opt), fancy_optimizers.FancyStochasticOptimizer)
                self.fancy_opt = opt
            else:
                self.fancy_opt = fancy_optimizers.FancyAdamax(
                    learning_rate=hybrid_inference_params.learning_rate,
                    beta1=hybrid_inference_params.adamax_beta1,
                    beta2=hybrid_inference_params.adamax_beta2,
                    sample_specific_only=False)

            self.continuous_model_step_func = self.continuous_model_advi.objective.step_function(
                score=True,
                obj_optimizer=self.fancy_opt.get_optimizer(self.continuous_model, self.continuous_model_approx),
                total_grad_norm_constraint=self.hybrid_inference_params.total_grad_norm_constraint,
                obj_n_mc=self.hybrid_inference_params.obj_n_mc,
                more_updates=temperature_update)

        if self.sampler is not None:
            self.sampler.update_approximation(self.continuous_model_approx)

        if 'elbo_normalization_factor' in kwargs.keys():
            self.elbo_normalization_factor = kwargs['elbo_normalization_factor']
        else:
            self.elbo_normalization_factor = 1.0

        if 'advi_task_name' in kwargs.keys():
            self.advi_task_name = kwargs['advi_task_name']
        else:
            self.advi_task_name = "ADVI"

        if 'sampling_task_name' in kwargs.keys():
            self.sampling_task_name = kwargs['sampling_task_name']
        else:
            self.sampling_task_name = "sampling"

        if 'calling_task_name' in kwargs.keys():
            self.calling_task_name = kwargs['calling_task_name']
        else:
            self.calling_task_name = "calling_task_name"

        self._t0 = None
        self._t1 = None
        self.elbo_hist: List[float] = []
        self.rls_elbo_hist: List[float] = []
        self.snr_hist: List[float] = []
        self.i_epoch = 1
        self.i_advi = 1
        self.calling_hist: List[Tuple[int, bool, bool]] = []
        self.previous_sampling_rounds = 0
        self.latest_caller_update_summary: Optional[CallerUpdateSummary] = None
        self.tqdm_out = LoggerTQDMAdapter(_logger)

    @abstractmethod
    def disengage(self):
        raise NotImplementedError

    def engage(self):
        try:
            all_converged = False
            while self.i_epoch <= self.hybrid_inference_params.max_training_epochs:
                _logger.debug("Starting epoch {0}...".format(self.i_epoch))
                converged_continuous = self._update_continuous_posteriors()
                all_converged = converged_continuous
                if self.sampler is not None:
                    converged_sampling = self._update_log_emission_posterior_expectation()
                    all_converged = all_converged and converged_sampling
                if self.caller is not None:
                    converged_discrete = self._update_discrete_posteriors()
                    all_converged = all_converged and converged_discrete
                self.i_epoch += 1
                if all_converged and not self._premature_convergence():
                    break
            if all_converged:
                _logger.info("Inference task completed successfully and convergence achieved.")
            else:
                _logger.warning("Inference task completed successfully but convergence not achieved.")
        except KeyboardInterrupt:
            pass

    def _log_start(self, task_name: str, i_epoch: int):
        self._t0 = time.time()
        _logger.debug("Starting {0} for epoch {1}...".format(task_name, i_epoch))

    def _log_stop(self, task_name: str, i_epoch: int):
        self._t1 = time.time()
        _logger.debug('The {0} for epoch {1} successfully finished in {2:.2f}s'.format(
            task_name, i_epoch, self._t1 - self._t0))

    def _log_interrupt(self, task_name: str, i_epoch: int):
        _logger.warning('The {0} for epoch {1} was interrupted'.format(task_name, i_epoch))

    def _premature_convergence(self):
        too_few_epochs = self.i_epoch < self.hybrid_inference_params.min_training_epochs
        temperature_still_high = np.abs(self.temperature.get_value()[0] - 1) > self.temperature_tolerance
        still_in_annealing = temperature_still_high and not self.hybrid_inference_params.disable_annealing
        return too_few_epochs or still_in_annealing

    def _create_param_tracker(self):
        assert all([param_name in self.continuous_model.vars or
                    param_name in self.continuous_model.deterministics
                    for param_name in self.hybrid_inference_params.param_tracker_config.var_names]),\
            "Some of the parameters chosen to be tracker are not present in the model"
        return VariationalParameterTracker(self.hybrid_inference_params.param_tracker_config)

    def _update_continuous_posteriors(self) -> bool:
        self._log_start(self.advi_task_name, self.i_epoch)
        max_advi_iters = self.hybrid_inference_params.max_advi_iter_subsequent_epochs if self.i_epoch > 1 \
            else self.hybrid_inference_params.max_advi_iter_first_epoch
        converged = False
        with tqdm.trange(max_advi_iters,
                         desc="({0}) starting...".format(self.advi_task_name),
                         file=self.tqdm_out) as progress_bar:
            try:
                for _ in progress_bar:
                    loss = self.continuous_model_step_func() / self.elbo_normalization_factor
                    self.i_advi += 1
                    try:
                        self.advi_convergence_tracker(self.continuous_model_advi.approx, loss, self.i_advi)
                    except StopIteration:
                        if not self._premature_convergence():  # suppress signal if deemed premature
                            raise StopIteration
                    snr = self.advi_convergence_tracker.snr
                    elbo_mean = self.advi_convergence_tracker.mean
                    elbo_variance = self.advi_convergence_tracker.variance
                    if snr is not None:
                        self.snr_hist.append(snr)
                    self.elbo_hist.append(-loss)
                    self.rls_elbo_hist.append(elbo_mean)
                    progress_bar.set_description("({0} epoch {1}) ELBO: {2}, SNR: {3}, T: {4:.2f}".format(
                        self.advi_task_name,
                        self.i_epoch,
                        "{0:.3f} +/- {1:.3f}".format(-elbo_mean, np.sqrt(elbo_variance))
                        if elbo_mean is not None and elbo_variance is not None else "N/A",
                        "{0:.1f}".format(snr) if snr is not None else "N/A",
                        self.temperature.get_value()[0]),
                        refresh=False)
                    if self.param_tracker is not None \
                            and self.i_advi % self.hybrid_inference_params.track_model_params_every == 0:
                        self.param_tracker(self.continuous_model_advi.approx, loss, self.i_advi)

            except StopIteration:
                converged = True
                progress_bar.refresh()
                progress_bar.close()
                self._log_stop(self.advi_task_name, self.i_epoch)

            except KeyboardInterrupt:
                progress_bar.refresh()
                progress_bar.close()
                self._log_interrupt(self.advi_task_name, self.i_epoch)
                raise KeyboardInterrupt

        return converged

    def _update_log_emission_posterior_expectation(self):
        self._log_start(self.sampling_task_name, self.i_epoch)
        if self.i_epoch == 1:
            self.sampler.reset()  # clear out log emission
        lag = min(self.previous_sampling_rounds, self.hybrid_inference_params.sampler_smoothing_window)
        converged = False
        median_rel_err = np.nan
        with tqdm.trange(self.hybrid_inference_params.log_emission_sampling_rounds,
                         desc="({0} epoch {1})".format(self.sampling_task_name, self.i_epoch),
                         file=self.tqdm_out) as progress_bar:
            try:
                for i_round in progress_bar:
                    update_to_estimator = self.sampler.draw()
                    latest_estimator = self.sampler.get_latest_log_emission_posterior_mean_estimator()
                    update_to_estimator = (update_to_estimator - latest_estimator) / (i_round + 1 + lag)
                    self.sampler.increment(update_to_estimator)
                    latest_estimator = self.sampler.get_latest_log_emission_posterior_mean_estimator()
                    median_rel_err = np.median(np.abs(update_to_estimator / latest_estimator).flatten())
                    std_rel_err = np.std(np.abs(update_to_estimator / latest_estimator).flatten())
                    del update_to_estimator
                    progress_bar.set_description("({0} epoch {1}) relative error: {2:2.4f} +/- {3:2.4f}".format(
                        self.sampling_task_name, self.i_epoch, median_rel_err, std_rel_err),
                        refresh=False)
                    if median_rel_err < self.hybrid_inference_params.log_emission_sampling_median_rel_error:
                        _logger.debug('{0} converged after {1} rounds with final '
                                      'median relative error {2:.3}.'.format(self.sampling_task_name, i_round + 1,
                                                                             median_rel_err))
                        raise StopIteration

            except StopIteration:
                converged = True
                progress_bar.refresh()
                progress_bar.close()
                self._log_stop(self.sampling_task_name, self.i_epoch)

            except KeyboardInterrupt:
                progress_bar.refresh()
                progress_bar.close()
                raise KeyboardInterrupt

            finally:
                self.previous_sampling_rounds = i_round + 1
                if not converged:
                    _logger.warning('{0} did not converge (median relative error '
                                    '= {1:.3}). Increase sampling rounds (current: {2}) if this behavior persists.'
                                    .format(self.sampling_task_name, median_rel_err,
                                            self.hybrid_inference_params.log_emission_sampling_rounds))

        return converged

    def _update_discrete_posteriors(self):
        self._log_start(self.calling_task_name, self.i_epoch)
        first_call_converged = False  # if convergence is achieved on the first call (stronger)
        iters_converged = False  # if internal loop is converged (weaker, does not imply global convergence)
        with tqdm.trange(self.hybrid_inference_params.max_calling_iters,
                         desc="({0} epoch {1})".format(self.calling_task_name, self.i_epoch),
                         file=self.tqdm_out) as progress_bar:
            try:
                for i_calling_iter in progress_bar:
                    caller_summary = self.caller.call()
                    self.latest_caller_update_summary = caller_summary
                    progress_bar.set_description("({0} epoch {1}) {2}".format(
                        self.calling_task_name, self.i_epoch, repr(caller_summary)), refresh=False)
                    caller_update_size_scalar = caller_summary.reduce_to_scalar()
                    if caller_update_size_scalar < self.hybrid_inference_params.caller_update_convergence_threshold:
                        iters_converged = True
                        if i_calling_iter == 0:
                            first_call_converged = True
                        raise StopIteration

            except StopIteration:
                progress_bar.refresh()
                progress_bar.close()
                self._log_stop(self.calling_task_name, self.i_epoch)

            except KeyboardInterrupt:
                progress_bar.refresh()
                progress_bar.close()
                self._log_interrupt(self.calling_task_name, self.i_epoch)
                raise KeyboardInterrupt

            finally:
                self.caller.update_auxiliary_vars()
                self.calling_hist.append((self.i_advi, iters_converged, first_call_converged))
                # if there is a self-consistency loop and not converged ...
                if not iters_converged and self.hybrid_inference_params.max_calling_iters > 1:
                    _logger.warning('{0} did not converge. Increase maximum calling rounds (current: {1}) '
                                    'if this behavior persists.'.format(
                        self.calling_task_name, self.hybrid_inference_params.max_calling_iters))

        return first_call_converged

class HybridInferenceParameters:
    """Hybrid ADVI inference parameters."""
    def __init__(self,
                 learning_rate: float = 0.05,
                 adamax_beta1: float = 0.9,
                 adamax_beta2: float = 0.99,
                 obj_n_mc: int = 1,
                 random_seed: int = 1984,
                 total_grad_norm_constraint: Optional[float] = None,
                 log_emission_samples_per_round: int = 50,
                 log_emission_sampling_median_rel_error: float = 5e-3,
                 log_emission_sampling_rounds: int = 10,
                 max_advi_iter_first_epoch: int = 100,
                 max_advi_iter_subsequent_epochs: int = 100,
                 min_training_epochs: int = 5,
                 max_training_epochs: int = 50,
                 initial_temperature: float = 2.0,
                 num_thermal_epochs: int = 20,
                 track_model_params: bool = False,
                 track_model_params_every: int = 10,
                 param_tracker_config: Optional['VariationalParameterTrackerConfig'] = None,
                 convergence_snr_averaging_window: int = 500,
                 convergence_snr_trigger_threshold: float = 0.1,
                 convergence_snr_countdown_window: int = 10,
                 max_calling_iters: int = 10,
                 caller_update_convergence_threshold: float = 1e-3,
                 caller_admixing_rate: float = 0.75,
                 sampler_smoothing_window: int = 0,
                 caller_summary_statistics_reducer: Callable[[np.ndarray], float] = np.mean,
                 disable_sampler: bool = False,
                 disable_caller: bool = False,
                 disable_annealing: bool = False):
        """See `expose_args` for the description of arguments."""
        self.learning_rate = learning_rate
        self.adamax_beta1 = adamax_beta1
        self.adamax_beta2 = adamax_beta2
        self.obj_n_mc = obj_n_mc
        self.random_seed = random_seed
        self.total_grad_norm_constraint = total_grad_norm_constraint
        self.log_emission_samples_per_round = log_emission_samples_per_round
        self.log_emission_sampling_median_rel_error = log_emission_sampling_median_rel_error
        self.log_emission_sampling_rounds = log_emission_sampling_rounds
        self.max_advi_iter_first_epoch = max_advi_iter_first_epoch
        self.max_advi_iter_subsequent_epochs = max_advi_iter_subsequent_epochs
        self.min_training_epochs = min_training_epochs
        self.max_training_epochs = max_training_epochs
        self.initial_temperature = initial_temperature
        self.num_thermal_epochs = num_thermal_epochs
        self.track_model_params = track_model_params
        self.track_model_params_every = track_model_params_every
        self.param_tracker_config = param_tracker_config
        self.convergence_snr_averaging_window = convergence_snr_averaging_window
        self.convergence_snr_trigger_threshold = convergence_snr_trigger_threshold
        self.convergence_snr_countdown_window = convergence_snr_countdown_window
        self.max_calling_iters = max_calling_iters
        self.caller_update_convergence_threshold = caller_update_convergence_threshold
        self.caller_admixing_rate = caller_admixing_rate
        self.sampler_smoothing_window = sampler_smoothing_window
        self.caller_summary_statistics_reducer = caller_summary_statistics_reducer
        self.disable_sampler = disable_sampler
        self.disable_caller = disable_caller
        self.disable_annealing = disable_annealing

        self._assert_params()

    def _assert_params(self):
        assert self.learning_rate > 0
        assert self.adamax_beta1 > 0
        assert self.adamax_beta2 > 0
        assert self.obj_n_mc > 0
        assert self.log_emission_samples_per_round > 0
        assert 0.0 < self.log_emission_sampling_median_rel_error < 1.0
        assert self.log_emission_sampling_rounds > 0
        assert self.max_advi_iter_first_epoch > 0
        assert self.max_advi_iter_subsequent_epochs > 0
        assert self.min_training_epochs > 0
        assert self.max_training_epochs > 0
        assert self.max_training_epochs >= self.min_training_epochs
        assert self.num_thermal_epochs <= self.max_training_epochs
        assert self.initial_temperature >= 1.0
        assert self.disable_annealing or self.num_thermal_epochs > 0
        assert self.track_model_params_every > 0
        assert self.convergence_snr_averaging_window > 0
        assert self.convergence_snr_trigger_threshold > 0
        assert self.max_calling_iters > 0
        assert self.caller_update_convergence_threshold > 0
        assert self.caller_admixing_rate > 0
        assert self.sampler_smoothing_window >= 0

        if self.track_model_params:
            assert self.param_tracker_config is not None

        if self.disable_annealing and self.initial_temperature > 1.0:
            _logger.warning("The initial temperature ({0}) is above 1.0 and annealing is disabled. This run "
                            "makes inferences in a thermal state. Ignore this warning if "
                            "this setting is intended.".format(self.initial_temperature))

    @staticmethod
    def expose_args(args: argparse.ArgumentParser, override_default: Dict[str, Any] = None, hide: Set[str] = None):
        """Exposes arguments of `__init__` to a given instance of `ArgumentParser`.

        Args:
            args: an instance of `ArgumentParser`
            override_default: a dictionary containing arguments the default values of which
                are to be overridden before passing to the argument parser
            hide: a set of arguments not to expose

        Returns:
            None
        """
        group = args.add_argument_group(title="Inference parameters")
        if override_default is None:
            override_default = dict()
        if hide is None:
            hide = set()

        initializer_params = inspect.signature(HybridInferenceParameters.__init__).parameters
        valid_args = {"--" + arg for arg in initializer_params.keys()}
        for hidden_arg in hide:
            assert hidden_arg in valid_args, \
                "Initializer argument to be hidden {0} is not a valid initializer arguments; possible " \
                "choices are: {1}".format(hidden_arg, valid_args)
        for override_default_arg in override_default.keys():
            assert override_default_arg in valid_args, \
                "Initializer argument of which the default is to be overridden {0} is not a valid initializer " \
                "arguments; possible choices are: {1}".format(override_default_arg, valid_args)

        def process_and_maybe_add(arg, **kwargs):
            full_arg = "--" + arg
            if full_arg in hide:
                return
            if full_arg in override_default:
                kwargs['default'] = override_default[full_arg]
            else:
                kwargs['default'] = initializer_params[arg].default
            group.add_argument(full_arg, **kwargs)

        def str_to_bool(value: str):
            if value.lower() in ('yes', 'true', 't', 'y', '1'):
                return True
            elif value.lower() in ('no', 'false', 'f', 'n', '0'):
                return False
            else:
                raise argparse.ArgumentTypeError('Boolean value expected.')

        process_and_maybe_add("learning_rate",
                              type=float,
                              help="Adamax optimizer learning rate")

        process_and_maybe_add("adamax_beta1",
                              type=float,
                              help="Adamax first moment estimator forgetting factor")

        process_and_maybe_add("adamax_beta2",
                              type=float,
                              help="Adamax second moment estimator forgetting factor")

        process_and_maybe_add("log_emission_samples_per_round",
                              type=int,
                              help="Number of log emission posterior samples per sampling round")

        process_and_maybe_add("log_emission_sampling_median_rel_error",
                              type=float,
                              help="Maximum tolerated median relative error in log emission posterior sampling")

        process_and_maybe_add("log_emission_sampling_rounds",
                              type=int,
                              help="Maximum log emission posterior sampling rounds")

        process_and_maybe_add("max_advi_iter_first_epoch",
                              type=int,
                              help="Maximum ADVI iterations in the first epoch (before sampling and calling)")

        process_and_maybe_add("max_advi_iter_subsequent_epochs",
                              type=int,
                              help="Maximum ADVI iterations after the first epoch")

        process_and_maybe_add("min_training_epochs",
                              type=int,
                              help="Minimum number of training epochs before evaluating convergence")

        process_and_maybe_add("max_training_epochs",
                              type=int,
                              help="Maximum number of training epochs before stopping")

        process_and_maybe_add("initial_temperature",
                              type=float,
                              help="Initial temperature for deterministic annealing (must be >= 1.0)")

        process_and_maybe_add("num_thermal_epochs",
                              type=int,
                              help="Annealing duration (in the units of training epochs)")

        process_and_maybe_add("convergence_snr_averaging_window",
                              type=int,
                              help="Averaging window for calculating training SNR for evaluating convergence")

        process_and_maybe_add("convergence_snr_trigger_threshold",
                              type=float,
                              help="The SNR threshold to be reached for triggering convergence")

        process_and_maybe_add("convergence_snr_countdown_window",
                              type=int,
                              help="The number of ADVI iterations during which the SNR is required to stay below the "
                                   "set threshold for convergence")

        process_and_maybe_add("max_calling_iters",
                              type=int,
                              help="Maximum number of calling internal self-consistency iterations")

        process_and_maybe_add("caller_update_convergence_threshold",
                              type=float,
                              help="Maximum tolerated calling update size for convergence")

        process_and_maybe_add("caller_admixing_rate",
                              type=float,
                              help="Admixing ratio of new and old caller posteriors (between 0 and 1; higher means "
                                   "using more of the new posterior)")

        process_and_maybe_add("disable_sampler",
                              type=str_to_bool,
                              help="(advanced) Disable sampler")

        process_and_maybe_add("disable_caller",
                              type=str_to_bool,
                              help="(advanced) Disable caller")

        process_and_maybe_add("disable_annealing",
                              type=str_to_bool,
                              help="(advanced) Keep the temperature pinned at its initial value and disable annealing")

    @staticmethod
    def from_args_dict(args_dict: Dict):
        """Initialize an instance of `HybridInferenceParameters` from a dictionary of arguments.

        Args:
            args_dict: a dictionary of arguments; the keys must match argument names in
                `HybridInferenceParameters.__init__`

        Returns:
            an instance of `HybridInferenceParameters`
        """
        relevant_keys = set(inspect.getfullargspec(HybridInferenceParameters.__init__).args)
        relevant_kwargs = {k: v for k, v in args_dict.items() if k in relevant_keys}
        return HybridInferenceParameters(**relevant_kwargs)


#!/usr/bin/env python3

"""
Script that takes the output of the universe machine.

This output is a large text file, each
row being one object. We convert this into a numpy array and
1) throw away the columns we don't need
2) throw away central galaxies (and their satellites) smaller than a certain
    mass

The output is a .npz file with two keys, "centrals" and "satellites".
Satellites are sorted (ascending) by their parent ID and then by their total
stellar mass

Step 1 runs at ~25M rows per hour
Step 2 is very fast

We will not re-run steps if the intermediate files already exist
"""
import os
import argparse

from io import StringIO

import numpy as np

full_catalog_dtype = [
    # Halo stuff
    ("id", "int"),           # Unique halo ID
    ("descid", "int"),       # ID of descendant halo (or -1 at z=0).
    ("upid", "int"),         # Parent halo ID (or -1 if this is a central)
    ("flags", "int"),        # Ignore
    ("updist", "float64"),   # Ignore
    ("x", "float64"),
    ("y", "float64"),
    ("z", "float64"),        # halo position (comoving Mpc/h)
    ("vx", "float64"),
    ("vy", "float64"),
    ("vz", "float64"),       # halo velocity (physical peculiar km/s)
    ("m", "float64"),
    # Halo mass (Bryan & Norman 1998 virial mass, Msun)
    ("v", "float64"),        # Halo vmax (physical km/s)
    ("mp", "float64"),       # Halo peak historical mass (BN98 vir, Msun)
    ("vmp", "float64"),
    # VMP: Halo vmax at the time when peak mass was reached
    ("r", "float64"),        # Halo radius (BN98 vir, comoving kpc/h)
    ("rank1", "float64"),
    # halo rank in Delta_vmax (see UniverseMachine paper)
    ("rank2", "float64"),    # Ignore
    ("ra", "float64"),       # Ignore
    ("ra_rank", "float64"),  # Ignore
    # Stellar stuff
    ("sm", "float64"),       # True stellar mass (Msun)
    ("icl", "float64"),      # True intracluster stellar mass (Msun)
    ("sfr", "float64"),      # True star formation rate (Msun/yr)
    ("obs_sm", "float64"),
    # observed stellar mass, including random & systematic errors (Msun)
    ("obs_sfr", "float64"),
    # observed SFR, including random & systematic errors (Msun/yr)
    ("ssfr", "float64"),     # observed SSFR
    ("sm_hm", "float64"),    # SMHM ratio
    ("obs_uv", "float64"),   # Observed UV Magnitude (M_1500 AB)
]

reduced_catalog_dtype = [
    # Halo stuff
    ("id", "int"),           # Unique halo ID
    ("upid", "int"),         # Parent halo ID (or -1 if this is a central)
    ("x", "float64"),
    ("y", "float64"),
    ("z", "float64"),        # halo position (comoving Mpc/h)
    ("vx", "float64"),
    ("vy", "float64"),
    ("vz", "float64"),       # halo position (comoving Mpc/h)
    ("m", "float64"),
    # Halo mass (Bryan & Norman 1998 virial mass, Msun)
    ("mp", "float64"),       # Halo peak historical mass (BN98 vir, Msun)
    ("v", "float64"),        # Halo vmax (physical km/s)
    # Stellar stuff
    ("sm", "float64"),       # True stellar mass (Msun)
    ("icl", "float64"),      # True intracluster stellar mass (Msun)
    ("sfr", "float64"),      # True star formation rate (Msun/yr)
]

def main(data_file, mpeak_lim=11.0):
    """Update galaxy atalogs.

    With current data:
    # The intermediate data contains 68448380 centrals (80%)
    # The final data contains 385125 centrals (with a cut at 12)
      and 10809469 satellites (with no cut)
    """
    # Input catalog
    if (not os.path.isfile(data_file)) and (not os.path.islink(data_file)):
        raise IOError("# Can not find the UM results: %s" % data_file)

    um_pre, _ = os.path.splitext(data_file)

    # Reduce the number of colunms and save as a numpy array
    inter_file = um_pre + '_reduced_cols.npy'
    print("# The catalog with reduced columns will be saved as: %s" % inter_file)

    # Remove small parent halos and satellites associated with these
    # small things
    final_file = um_pre + "_logmp_%4.1f.npz" % mpeak_lim
    print("# The final catalog after the Mpeak cut will be saved as: %s" % final_file)

    # If we have already generated the inter_file, don't do it again...
    blockSize = 100000

    if not os.path.isfile(inter_file):
        reduced_catalog = np.zeros(blockSize, dtype=reduced_catalog_dtype)
        with open(data_file) as f:
            count = 0
            for line in f:
                if line.startswith("#"):
                    continue
                row = np.loadtxt(StringIO(line), dtype=full_catalog_dtype)
                reduced_catalog[count] = row[
                    list(reduced_catalog.dtype.names)].copy()
                count += 1
                if count % blockSize == 0:
                    reduced_catalog.resize(len(reduced_catalog) + blockSize)
                    print(count)
        np.save(inter_file, reduced_catalog[:count])
        del reduced_catalog
    else:
        print("# Skipping reducing cols, file already exists")

    if not os.path.isfile(final_file):
        reduced_catalog = np.load(inter_file)

        # All centrals greater than a given mass
        central_catalog = np.sort(
            reduced_catalog[
                (reduced_catalog["upid"] == -1) &
                (reduced_catalog["mp"] > (10.0 ** mpeak_lim))],
            order="id")
        central_ids = frozenset(central_catalog["id"])

        # Remove all halos not associated with one of those centrals
        satellite_catalog = reduced_catalog[
            np.array(
                [parent in central_ids for parent in reduced_catalog["upid"]]
                )]
        np.savez(final_file,
                 centrals=central_catalog,
                 satellites=satellite_catalog)

        del reduced_catalog, central_catalog, central_ids, satellite_catalog
    else:
        print("# File already exists !")

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument(
        'um_file', type=str,
        help=('UniverseMachine snapshot file in .txt format'))

    parser.add_argument(
        '-m', '--mpeak', dest='mpeak_lim',
        help='Lower limit of Mpeak',
        type=float, default=11.0)

    args = parser.parse_args()

    main(args.um_file, mpeak_lim=args.mpeak_lim)


#!/usr/bin/env python

import argparse
import os
import os.path as osp
import glob
import numpy as np
import cv2
import sys
this_dir=osp.dirname(__file__)
sys.path.insert(0, osp.join(this_dir, '../../external/py-faster-rcnn/lib'))
sys.path.insert(0, osp.join(this_dir, '../../src'))
from fast_rcnn.nms_wrapper import nms
import cPickle
from time import time
from tpn.data_io import tpn_test_iterator
sys.path.insert(0, osp.join(this_dir, '../../external'))
from vdetlib.utils.visual import unique_colors, add_bbox
from vdetlib.utils.common import imread, imwrite
from vdetlib.utils.protocol import frame_path_at, proto_load
from vdetlib.vdet.dataset import imagenet_vdet_classes

def _frame_dets(tracks, frame_idx, score_key, box_key):
    scores = []
    boxes = []
    track_ids = []
    for track_idx, track in enumerate(tracks):
        if frame_idx not in track['frame']: continue
        assert score_key in track
        assert box_key in track
        ind = track['frame'] == frame_idx
        cur_scores = track[score_key][ind]
        cur_boxes = track[box_key][ind,:]
        num_cls = cur_scores.shape[1]
        # repeat boxes if not class specific
        if cur_boxes.shape[1] != num_cls:
            cur_boxes = np.repeat(cur_boxes[:,np.newaxis,:], num_cls, axis=1)
        track_ids.append(track_idx)
        scores.append(cur_scores)
        boxes.append(cur_boxes)
    scores = np.concatenate(scores, 0)
    boxes = np.concatenate(boxes, 0)
    return track_ids, scores, boxes

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('vid_file')
    parser.add_argument('track_file',
        help='Track file.')
    parser.add_argument('score_key')
    parser.add_argument('box_key')
    parser.add_argument('--thres', type=float, default=0.1,
        help='Detection score threshold. [0.1]')
    parser.add_argument('--num_classes', type=int, default=31,
        help='Number of classes. [31]')
    parser.add_argument('--max_per_image', type=int, default=100,
        help='Maximum number of detections per image. [100]')
    parser.add_argument('--save_dir',
        help='Save directory.')
    args = parser.parse_args()

    num_classes = args.num_classes

    # process vid detections
    vid_proto = proto_load(args.vid_file)
    tracks = tpn_test_iterator(args.track_file)
    end_frame_idx = set([track['frame'][-1] for track in tracks])
    end_frames = [frame for frame in vid_proto['frames'] \
        if frame['frame'] in end_frame_idx]
    vid_name = vid_proto['video']
    kept_track_ids = []
    kept_class = []
    for frame in end_frames:
        frame_name = osp.join(vid_name, osp.splitext(frame['path'])[0])

        frame_idx = frame['frame']
        start_time = time()
        track_ids, scores, boxes = _frame_dets(tracks, frame_idx, args.score_key, args.box_key)
        boxes = boxes.reshape((boxes.shape[0], -1))

        for j in xrange(1, num_classes):
            inds = np.where(scores[:, j] > args.thres)[0]
            cls_scores = scores[inds, j]
            cls_boxes = boxes[inds, j*4:(j+1)*4]
            cls_dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])) \
                .astype(np.float32, copy=False)
            keep = nms(cls_dets, 0.3, force_cpu=True)
            for keep_id in keep:
                kept_track_ids.append(track_ids[inds[keep_id]])
                kept_class.append(j)

    colors = unique_colors(len(kept_track_ids))
    kept_tracks = [tracks[i] for i in kept_track_ids]
    idx = 0
    while True:
        frame = vid_proto['frames'][idx]
        frame_id = frame['frame']
        print "Frame id: {}".format(frame_id)
        img = imread(frame_path_at(vid_proto, frame['frame']))
        boxes = []
        scores = []
        show_track_ids = []
        cur_colors = []
        cur_classes = []
        for track_id, (class_id, track) in enumerate(zip(kept_class, kept_tracks)):
            if frame_id in track['frame']:
                boxes.append(track[args.box_key][track['frame'] == frame_id][0,class_id,:].tolist())
                scores.append(track[args.score_key][track['frame'] == frame_id][0,class_id].tolist())
                cur_colors.append(colors[track_id])
                cur_classes.append(imagenet_vdet_classes[class_id])
                show_track_ids.append(track_id)
        tracked = add_bbox(img, boxes, classes=cur_classes,
                           scores=scores, line_width=10)
        if args.save_dir:
            imwrite(os.path.join(args.save_dir, "{:04d}.jpg".format(frame_id)),
                tracked)
            idx += 1
            if idx >= len(vid_proto['frames']):
                break
            continue
        cv2.imshow('tracks', tracked)
        key = cv2.waitKey(0)
        if key == ord('q'):
            cv2.destroyAllWindows()
            sys.exit(0)
        elif key == ord('a'):
            if idx > 0:
                idx -= 1
        elif key == ord('d'):
            if idx < len(vid_proto['frames']) - 1:
                idx += 1
    if not args.save_dir:
        cv2.destroyAllWindows()


#!/usr/bin/python
#
# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import functools
import os
import json
import math
from collections import defaultdict
import random
import numpy as np
import torch
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import torchvision.models as models
import pdb

from sg2im.data.coco import CocoSceneGraphDataset, coco_collate_fn
from sg2im.discriminators import PatchDiscriminator, AcCropDiscriminator, CondGANPatchDiscriminator, CondGANDiscriminator
from sg2im.losses import get_gan_losses
from sg2im.metrics import jaccard, relation_score

#from sg2im.model_layout_8ptembed_mask import Sg2ImModel
#from sg2im.model_layout_8ptembed_maskembed import Sg2ImModel
#from sg2im.model_layout_embed import Sg2ImModel
from sg2im.model_layout_mask import Sg2ImModel
#from sg2im.model import Sg2ImModel
from sg2im.utils import int_tuple, float_tuple, str_tuple
from sg2im.utils import timeit, bool_flag, LossManager

####
from sg2im.perceptual_loss import FeatureExtractor
from sg2im.data.utils import imagenet_deprocess_batch
from sg2im.data.utils import perc_process_batch
#from sg2im.logger import Logger

from imageio import imwrite

#torch.backends.cudnn.benchmark = True

COCO_DIR = os.path.expanduser('/dataset/coco_stuff')
#COCO_DIR = os.path.expanduser('/Users/brigitsc/sandbox/sg2im/datasets/coco')

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='coco', choices=['vg', 'coco'])

# Optimization hyperparameters
parser.add_argument('--batch_size', default=32, type=int)
parser.add_argument('--num_iterations', default=1000000, type=int)
parser.add_argument('--learning_rate', default=1e-4, type=float)

# Switch the generator to eval mode after this many iterations
parser.add_argument('--eval_mode_after', default=100000, type=int)

# Dataset options common to both VG and COCO
parser.add_argument('--image_size', default='64,64', type=int_tuple)
parser.add_argument('--num_train_samples', default=None, type=int)
parser.add_argument('--num_val_samples', default=1024, type=int)
parser.add_argument('--shuffle_val', default=False, type=bool_flag)
parser.add_argument('--loader_num_workers', default=4, type=int)
parser.add_argument('--include_relationships', default=True, type=bool_flag)

# COCO-specific options
parser.add_argument('--coco_train_image_dir',
         default=os.path.join(COCO_DIR, 'images/train2017'))
parser.add_argument('--coco_val_image_dir',
         default=os.path.join(COCO_DIR, 'images/val2017'))
parser.add_argument('--coco_train_instances_json',
         default=os.path.join(COCO_DIR, 'annotations/instances_train2017.json'))
parser.add_argument('--coco_train_stuff_json',
         default=os.path.join(COCO_DIR, 'annotations/stuff_train2017.json'))
parser.add_argument('--coco_val_instances_json',
         default=os.path.join(COCO_DIR, 'annotations/instances_val2017.json'))
parser.add_argument('--coco_val_stuff_json',
         default=os.path.join(COCO_DIR, 'annotations/stuff_val2017.json'))
parser.add_argument('--instance_whitelist', default=None, type=str_tuple)
parser.add_argument('--stuff_whitelist', default=None, type=str_tuple)
parser.add_argument('--coco_include_other', default=False, type=bool_flag)
parser.add_argument('--min_object_size', default=0.02, type=float)
parser.add_argument('--min_objects_per_image', default=3, type=int)
parser.add_argument('--coco_stuff_only', default=True, type=bool_flag)
parser.add_argument('--random_seed', default=42, type=int)  #For comparative results/debug

# Generator options
parser.add_argument('--mask_size', default=16, type=int) # Set this to 0 to use no masks
parser.add_argument('--embedding_dim', default=128, type=int)
parser.add_argument('--gconv_dim', default=128, type=int)
parser.add_argument('--gconv_hidden_dim', default=512, type=int)
parser.add_argument('--gconv_num_layers', default=5, type=int)
parser.add_argument('--mlp_normalization', default='none', type=str)
parser.add_argument('--refinement_network_dims', default='1024,512,256,128,64', type=int_tuple)
parser.add_argument('--normalization', default='batch')
parser.add_argument('--activation', default='leakyrelu-0.2')
parser.add_argument('--layout_noise_dim', default=32, type=int)
parser.add_argument('--use_boxes_pred_after', default=-1, type=int)

## scene_graph conditioning on GAN
parser.add_argument('--sg_context_dim', default=0, type=int)
parser.add_argument('--sg_context_dim_d', default=0, type=int)
parser.add_argument('--image_patch_discr', default=True, type=bool_flag)
parser.add_argument('--gcnn_pooling', default='sum')
parser.add_argument('--layout_for_discrim', default=False, type=bool_flag)

parser.add_argument('--matching_aware_loss', default=False, type=bool_flag)

# Generator losses
parser.add_argument('--mask_loss_weight', default=0, type=float)
parser.add_argument('--l1_pixel_loss_weight', default=1.0, type=float)
parser.add_argument('--bbox_pred_loss_weight', default=10, type=float)
parser.add_argument('--predicate_pred_loss_weight', default=0, type=float) # DEPRECATED
parser.add_argument('--perceptual_loss_weight', default=0, type=float)
parser.add_argument('--grayscale_perceptual', action='store_true', help='Calculate perceptual loss with grayscale images')
parser.add_argument('--log_perceptual', action='store_true', help='Take logarithm of perceptual loss')

# Generic discriminator options
parser.add_argument('--discriminator_loss_weight', default=0.01, type=float)
parser.add_argument('--gan_loss_type', default='gan')
parser.add_argument('--d_clip', default=None, type=float)
parser.add_argument('--d_normalization', default='batch')
parser.add_argument('--d_padding', default='valid')
parser.add_argument('--d_activation', default='leakyrelu-0.2')

# Object discriminator
parser.add_argument('--d_obj_arch',
    default='C4-64-2,C4-128-2,C4-256-2')
parser.add_argument('--crop_size', default=32, type=int)
parser.add_argument('--d_obj_weight', default=1.0, type=float) # multiplied by d_loss_weight
parser.add_argument('--ac_loss_weight', default=0.1, type=float)

# Image discriminator
parser.add_argument('--d_img_arch',
    default='C4-64-2,C4-128-2,C4-256-2')
parser.add_argument('--d_img_weight', default=1.0, type=float) # multiplied by d_loss_weight

# Output options
parser.add_argument('--print_every', default=10, type=int)
parser.add_argument('--timing', default=False, type=bool_flag)
parser.add_argument('--output_dir', default=os.getcwd())

parser.add_argument('--checkpoint', default='sg2im-models/coco64.pt')
parser.add_argument('--device', default='gpu', choices=['cpu', 'gpu'])

def add_loss(total_loss, curr_loss, loss_dict, loss_name, weight=1, logarithm=False):
  curr_loss = curr_loss * weight
  if logarithm:
    curr_loss = torch.log(curr_loss)
  loss_dict[loss_name] = curr_loss.item()
  if total_loss is not None:
    total_loss += curr_loss
  else:
    total_loss = curr_loss
  return total_loss

def check_args(args):
  H, W = args.image_size
  for _ in args.refinement_network_dims[1:]:
    H = H // 2
  if H == 0:
    raise ValueError("Too many layers in refinement network")

def build_model(args, vocab):
  kwargs = {
    'vocab': vocab,
    'image_size': args.image_size,
    'embedding_dim': args.embedding_dim,
    'gconv_dim': args.gconv_dim,
    'gconv_hidden_dim': args.gconv_hidden_dim,
    'gconv_num_layers': args.gconv_num_layers,
    'mlp_normalization': args.mlp_normalization,
    'refinement_dims': args.refinement_network_dims,
    'normalization': args.normalization,
    'activation': args.activation,
    'mask_size': args.mask_size,
    'layout_noise_dim': args.layout_noise_dim,
  }
  model = Sg2ImModel(**kwargs)
  return model, kwargs

def build_coco_dsets(args):
  dset_kwargs = {
    'image_dir': args.coco_train_image_dir,
    'instances_json': args.coco_train_instances_json,
    'stuff_json': args.coco_train_stuff_json,
    'stuff_only': args.coco_stuff_only,
    'image_size': args.image_size,
    'mask_size': args.mask_size,
    'max_samples': args.num_train_samples,
    'min_object_size': args.min_object_size,
    'min_objects_per_image': args.min_objects_per_image,
    'instance_whitelist': args.instance_whitelist,
    'stuff_whitelist': args.stuff_whitelist,
    'include_other': args.coco_include_other,
    'include_relationships': args.include_relationships,
    'seed': args.random_seed
  }
  train_dset = None
  #train_dset = CocoSceneGraphDataset(**dset_kwargs)
  #num_objs = train_dset.total_objects()
  #num_imgs = len(train_dset)
  #print('Training dataset has %d images and %d objects' % (num_imgs, num_objs))
  #print('(%.2f objects per image)' % (float(num_objs) / num_imgs))

  dset_kwargs['image_dir'] = args.coco_val_image_dir
  dset_kwargs['instances_json'] = args.coco_val_instances_json
  dset_kwargs['stuff_json'] = args.coco_val_stuff_json
  dset_kwargs['max_samples'] = args.num_val_samples
  val_dset = CocoSceneGraphDataset(**dset_kwargs)

  #assert train_dset.vocab == val_dset.vocab
  #vocab = json.loads(json.dumps(train_dset.vocab))
  vocab = json.loads(json.dumps(val_dset.vocab))

  return vocab, train_dset, val_dset

def build_loaders(args):
  if args.dataset == 'coco':
    vocab, train_dset, val_dset = build_coco_dsets(args)
    collate_fn = coco_collate_fn

  loader_kwargs = {
    'batch_size': args.batch_size,
    'num_workers': args.loader_num_workers,
    'shuffle': True,
    'collate_fn': collate_fn,
  }
  #train_loader = DataLoader(train_dset, **loader_kwargs)
  train_loader = None

  loader_kwargs['shuffle'] = args.shuffle_val
  val_loader = DataLoader(val_dset, **loader_kwargs)
  return vocab, train_loader, val_loader

def get_rel_score(args, t, loader, model):
  float_dtype = torch.FloatTensor
  long_dtype = torch.LongTensor
  num_samples = 0
  all_losses = defaultdict(list)
  total_iou = 0
  total_boxes = 0

  ###################
  if not os.path.isdir(args.output_dir):
    os.mkdir(args.output_dir)
    print('Created %s' %args.output_dir)

  img_dir = args.output_dir+'/img_dir'

  if not os.path.isdir(img_dir):
    os.mkdir(img_dir)
    print('Created %s' %img_dir)
  ##################
  rel_score = 0
  with torch.no_grad():
    o_start = o_end = 0
    t_start = t_end = 0
    last_o_idx = last_t_idx = 0

    b = 0
    total_boxes = 0
    total_iou = 0
    for batch in loader:
      batch = [tensor.cuda() for tensor in batch]
      #batch = [tensor for tensor in batch]
      masks = None
      if len(batch) == 6:
        imgs, objs, boxes, triples, obj_to_img, triple_to_img = batch
      elif len(batch) == 8:
        imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, triplet_masks = batch
      #elif len(batch) == 7:
      #  imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img = batch
      predicates = triples[:, 1]

      objs = objs.detach()
      triples = triples.detach()
      # Run the model as it has been run during training
      model_masks = masks
      model_out = model(objs, triples, obj_to_img, boxes_gt=boxes, masks_gt=model_masks)
      # imgs_pred, boxes_pred, masks_pred, predicate_scores = model_out
      #imgs_pred, boxes_pred, masks_pred, objs_vec, layout, layout_boxes, layout_masks, obj_to_img, sg_context_pred, sg_context_pred_d, predicate_scores, obj_embeddings, pred_embeddings, triple_boxes_pred, triple_boxes_gt = model_out
      imgs_pred, boxes_pred, masks_pred, objs_vec, layout, layout_boxes, layout_masks, obj_to_img, sg_context_pred, sg_context_pred_d, predicate_scores, obj_embeddings, pred_embeddings, triple_boxes_pred, triple_boxes_gt, triplet_masks_pred = model_out
      num_samples += imgs.size(0)
      if num_samples >= args.num_val_samples:
        break

      rel_score += relation_score(boxes_pred, boxes, masks_pred, masks, model.vocab)
      b += 1

      total_iou += jaccard(boxes_pred, boxes)
      total_boxes += boxes_pred.size(0)

    print (b)
    print (total_boxes)

    rel_score = rel_score/b
    avg_iou = total_iou / total_boxes

    # print('rel score:' rel_score)
    # print('avg iou:' avg_iou)
    return rel_score, avg_iou

#def check_model(args, t, loader, model, logger=None, log_tag='', write_images=False):
def check_model(args, t, loader, model, log_tag='', write_images=False):
  float_dtype = torch.cuda.FloatTensor
  long_dtype = torch.cuda.LongTensor
  #float_dtype = torch.FloatTensor
  #long_dtype = torch.LongTensor
  num_samples = 0
  all_losses = defaultdict(list)
  total_iou = 0
  total_boxes = 0

  ###################
  if not os.path.isdir(args.output_dir):
    os.mkdir(args.output_dir)
    print('Created %s' %args.output_dir)

  img_dir = args.output_dir+'/img_dir'

  if not os.path.isdir(img_dir):
    os.mkdir(img_dir)
    print('Created %s' %img_dir)
  ##################

  t = 0
  with torch.no_grad():
    for batch in loader:
      batch = [tensor.cuda() for tensor in batch]
      #batch = [tensor for tensor in batch]
      masks = None
      if len(batch) == 6:
        imgs, objs, boxes, triples, obj_to_img, triple_to_img = batch
      elif len(batch) == 8:
        imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, triplet_masks = batch
      #elif len(batch) == 7:
      #  imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img = batch
      predicates = triples[:, 1]

      # Run the model as it has been run during training
      model_masks = masks
      model_out = model(objs, triples, obj_to_img, boxes_gt=boxes, masks_gt=model_masks)
      # imgs_pred, boxes_pred, masks_pred, predicate_scores = model_out
      #imgs_pred, boxes_pred, masks_pred, objs_vec, layout, layout_boxes, layout_masks, obj_to_img, sg_context_pred, sg_context_pred_d, predicate_scores, obj_embeddings, pred_embeddings, triple_boxes_pred, triple_boxes_gt = model_out
      imgs_pred, boxes_pred, masks_pred, objs_vec, layout, layout_boxes, layout_masks, obj_to_img, sg_context_pred, sg_context_pred_d, predicate_scores, obj_embeddings, pred_embeddings, triple_boxes_pred, triple_boxes_gt, triplet_masks_pred = model_out

      skip_pixel_loss = True
      skip_perceptual_loss = True
      total_loss, losses =  calculate_model_losses(
                                args, skip_pixel_loss, model, imgs, imgs_pred,
                                boxes, boxes_pred, masks, masks_pred,
                                predicates, predicate_scores,
                                skip_perceptual_loss)

      total_iou += jaccard(boxes_pred, boxes)
      total_boxes += boxes_pred.size(0)

      for loss_name, loss_val in losses.items():
        all_losses[loss_name].append(loss_val)
      num_samples += imgs.size(0)
      if num_samples >= args.num_val_samples:
        break

      ###################################################
      # can't use this since image generation is turned off.
      samples = {}
      samples['gt_img'] = imgs

      #model_out = model(objs, triples, obj_to_img, boxes_gt=boxes, masks_gt=masks)
      #samples['gt_box_gt_mask'] = model_out[0]

      #model_out = model(objs, triples, obj_to_img, boxes_gt=boxes)
      #samples['gt_box_pred_mask'] = model_out[0]

      #model_out = model(objs, triples, obj_to_img)
      #samples['pred_box_pred_mask'] = model_out[0]

      #for k, v in samples.items():
      #  samples[k] = imagenet_deprocess_batch(v)

      #if write_images:
        #3. Log ground truth and predicted images
      #  with torch.no_grad():
      #    p_imgs = samples['gt_box_gt_mask'].detach()
      #    gt_imgs = samples['gt_img'].detach()
      #    p_gbox_pmsk_img = samples['gt_box_pred_mask']
      #    p_test_imgs = samples['pred_box_pred_mask']

      #  np_gt_imgs = [gt.cpu().numpy().transpose(1,2,0) for gt in gt_imgs]
      #  np_pred_imgs = [pred.cpu().numpy().transpose(1,2,0) for pred in p_imgs]
      #  np_gbox_pmsk_imgs = [pred.cpu().numpy().transpose(1,2,0) for pred in p_gbox_pmsk_img]
      #  np_test_pred_imgs = [pred.cpu().numpy().transpose(1,2,0) for pred in p_test_imgs]
      #  np_all_imgs = []

      #  for gt_img, gtb_gtm_img, gtb_pm_img, pred_img in zip(np_gt_imgs, np_pred_imgs, np_gbox_pmsk_imgs, np_test_pred_imgs):
      #    img_path = os.path.join(img_dir, '%06d_gt_img.png' % t)
      #    imwrite(img_path, gt_img)

      #    img_path = os.path.join(img_dir, '%06d_gtb_gtm_img.png' % t)
      #    imwrite(img_path, gtb_gtm_img)

      #    img_path = os.path.join(img_dir, '%06d_gtb_pm_img.png' % t)
      #    imwrite(img_path, gtb_pm_img)

      #    img_path = os.path.join(img_dir, '%06d_pred_img.png' % t)
      #    imwrite(img_path, pred_img)

      #    t=t+1

        # for gt_img, gtb_gtm_img, gtb_pm_img, pred_img in zip(np_gt_imgs, np_pred_imgs, np_gbox_pmsk_imgs, np_test_pred_imgs):
        #   np_all_imgs.append((gt_img * 255.0).astype(np.uint8))
        #   np_all_imgs.append((gtb_gtm_img * 255.0).astype(np.uint8))
        #   np_all_imgs.append((gtb_pm_img * 255.0).astype(np.uint8))
        #   np_all_imgs.append((pred_img * 255.0).astype(np.uint8))

        # logger.image_summary(log_tag, np_all_imgs, t)
      #########################################################################

    mean_losses = {k: np.mean(v) for k, v in all_losses.items()}
    avg_iou = total_iou / total_boxes

    masks_to_store = masks
    if masks_to_store is not None:
      masks_to_store = masks_to_store.data.cpu().clone()

    masks_pred_to_store = masks_pred
    if masks_pred_to_store is not None:
      masks_pred_to_store = masks_pred_to_store.data.cpu().clone()

  batch_data = {
    'objs': objs.detach().cpu().clone(),
    'boxes_gt': boxes.detach().cpu().clone(),
    'masks_gt': masks_to_store,
    'triples': triples.detach().cpu().clone(),
    'obj_to_img': obj_to_img.detach().cpu().clone(),
    'triple_to_img': triple_to_img.detach().cpu().clone(),
    'boxes_pred': boxes_pred.detach().cpu().clone(),
    'masks_pred': masks_pred_to_store
  }
  #out = [mean_losses, samples, batch_data, avg_iou]
  out = [mean_losses, avg_iou]

  return tuple(out)

def calculate_model_losses(args, skip_pixel_loss, model, img, img_pred,
                           bbox, bbox_pred, masks, masks_pred,
                           predicates, predicate_scores,
                           skip_perceptual_loss,
                           perceptual_extractor=None):  #FeatureExtractor(requires_grad=False).cuda()):
  total_loss = torch.zeros(1).to(img)
  losses = {}

  l1_pixel_weight = args.l1_pixel_loss_weight
  if skip_pixel_loss:
    l1_pixel_weight = 0

#  l1_pixel_loss = F.l1_loss(img_pred, img)
#  total_loss = add_loss(total_loss, l1_pixel_loss, losses, 'L1_pixel_loss',
#                        l1_pixel_weight)

  loss_bbox = F.mse_loss(bbox_pred, bbox)
  total_loss = add_loss(total_loss, loss_bbox, losses, 'bbox_pred',
                        args.bbox_pred_loss_weight)

  if args.predicate_pred_loss_weight > 0:
    loss_predicate = F.cross_entropy(predicate_scores, predicates)
    total_loss = add_loss(total_loss, loss_predicate, losses, 'predicate_pred',
                          args.predicate_pred_loss_weight)

  if args.mask_loss_weight > 0 and masks is not None and masks_pred is not None:
    mask_loss = F.binary_cross_entropy(masks_pred, masks.float())
    total_loss = add_loss(total_loss, mask_loss, losses, 'mask_loss',
                          args.mask_loss_weight)

  #########################
  perceptual_weight = args.perceptual_loss_weight
  if skip_perceptual_loss:
    perceptual_weight = 0

  if perceptual_weight > 0:
    with torch.no_grad():
      img_p = perc_process_batch(img.detach(), rescale=True, grayscale=args.grayscale_perceptual)
      img_pred_p = perc_process_batch(img_pred.detach(), rescale=True, grayscale=args.grayscale_perceptual)

      real = perceptual_extractor(img_p)
      fake = perceptual_extractor(img_pred_p)

      vgg_perceptual_loss = F.mse_loss(fake.relu1_2, real.relu1_2)
      vgg_perceptual_loss += F.mse_loss(fake.relu2_2, real.relu2_2)
      vgg_perceptual_loss += F.mse_loss(fake.relu3_3, real.relu3_3)
      vgg_perceptual_loss += F.mse_loss(fake.relu4_3, real.relu4_3)

    total_loss = add_loss(total_loss, vgg_perceptual_loss, losses, 'perceptual_loss',
                          perceptual_weight, logarithm=args.log_perceptual)
  #########################
  return total_loss, losses

def main(args):

  if args.device == 'cpu':
    device = torch.device('cpu')
  elif args.device == 'gpu':
    device = torch.device('cuda:0')
    if not torch.cuda.is_available():
      print('WARNING: CUDA not available; falling back to CPU')
      device = torch.device('cpu')

  # Load the model, with a bit of care in case there are no GPUs
  map_location = 'cpu' if device == torch.device('cpu') else None
  checkpoint = torch.load(args.checkpoint, map_location=map_location)
  model = Sg2ImModel(**checkpoint['model_kwargs'])
  model.load_state_dict(checkpoint['model_state'], strict=False)
  model.eval()
  model.to(device)

  vocab, train_loader, val_loader = build_loaders(args)

  if not os.path.isdir(args.output_dir):
    os.mkdir(args.output_dir)
    print('Created %s' %args.output_dir)

  ## add code for validation visualization
  #logger = Logger(args.output_dir)
  logger = None

  t = 1
  with timeit('forward', args.timing):
    print('checking on val')
    #val_results = check_model(args, t, val_loader, model, logger=logger, log_tag='Validation', write_images=True)
    #val_results = check_model(args, t, val_loader, model, log_tag='Validation', write_images=True)

    rel_score, avg_iou = get_rel_score(args, t, val_loader, model)
    print ('relation score: ', rel_score)
    print ('get_rel_score average iou: ', avg_iou)
    #val_losses, val_avg_iou = val_results
    #print('val iou: ', val_avg_iou)

if __name__ == '__main__':
  args = parser.parse_args()
  main(args)


import numpy as np
from torch.utils import data
from data.gta5_dataset import GTA5DataSet
from data.cityscapes_dataset import cityscapesDataSet
from data.cityscapes_dataset_label import cityscapesDataSetLabel
from data.cityscapes_dataset_SSL import cityscapesDataSetSSL
from data.synthia_dataset import SYNDataSet

IMG_MEAN = np.array((0.0, 0.0, 0.0), dtype=np.float32)
image_sizes = {'cityscapes': (1024, 512), 'gta5': (
    1280, 720), 'synthia': (1280, 760)}
cs_size_test = {'cityscapes': (1344, 576)}

def CreateSrcDataLoader(args):
    """
    Creates the source data loader

    Args:
        commandline arguments

    Raises:
        ValueError: The value arg.source must be either gta5 or synthia

    Returns:
        torch.utils.data.DataLoader
    """
    if args.source == 'gta5':
        source_dataset = GTA5DataSet(args.data_dir, args.data_list, crop_size=image_sizes['cityscapes'],
                                     resize=image_sizes['gta5'], mean=IMG_MEAN,
                                     max_iters=args.num_steps * args.batch_size)
    elif args.source == 'synthia':
        source_dataset = SYNDataSet(args.data_dir, args.data_list, crop_size=image_sizes['cityscapes'],
                                    resize=image_sizes['synthia'], mean=IMG_MEAN,
                                    max_iters=args.num_steps * args.batch_size)
    else:
        raise ValueError('The source dataset mush be either gta5 or synthia')

    source_dataloader = data.DataLoader(source_dataset,
                                        batch_size=args.batch_size,
                                        shuffle=True,
                                        num_workers=args.num_workers,
                                        pin_memory=True)
    return source_dataloader

def CreateTrgDataLoader(args):
    """
    Creates Target data loader

    Args:
        commandline arguments

    Returns:
        torch.utils.data.DataLoader
    """
    if args.set == 'train' or args.set == 'trainval':
        target_dataset = cityscapesDataSetLabel(args.data_dir_target,
                                                args.data_list_target,
                                                crop_size=image_sizes['cityscapes'],
                                                mean=IMG_MEAN,
                                                max_iters=args.num_steps * args.batch_size,
                                                set=args.set)
    else:
        target_dataset = cityscapesDataSet(args.data_dir_target,
                                           args.data_list_target,
                                           crop_size=cs_size_test['cityscapes'],
                                           mean=IMG_MEAN,
                                           set=args.set)

    if args.set == 'train' or args.set == 'trainval':
        target_dataloader = data.DataLoader(target_dataset,
                                            batch_size=args.batch_size,
                                            shuffle=True,
                                            num_workers=args.num_workers,
                                            pin_memory=True)
    else:
        target_dataloader = data.DataLoader(target_dataset,
                                            batch_size=1,
                                            shuffle=False,
                                            pin_memory=True)

    return target_dataloader

def CreateTrgDataSSLLoader(args):
    """
    Create the Target data loader for SSL training

    Args:
        commandline arguments

    Returns:
        torch.utils.data.DataLoader
    """
    target_dataset = cityscapesDataSet(args.data_dir_target,
                                       args.data_list_target,
                                       crop_size=image_sizes['cityscapes'],
                                       mean=IMG_MEAN,
                                       set=args.set)
    target_dataloader = data.DataLoader(target_dataset,
                                        batch_size=1,
                                        shuffle=False,
                                        pin_memory=True)
    return target_dataloader

def CreatePseudoTrgLoader(args):
    """
    Create the Target data loader for pseudo label generation

    Args:
        commandline arguments

    Returns:
        torch.utils.data.DataLoader
    """
    target_dataset = cityscapesDataSetSSL(args.data_dir_target,
                                          args.data_list_target,
                                          crop_size=image_sizes['cityscapes'],
                                          mean=IMG_MEAN,
                                          max_iters=args.num_steps * args.batch_size,
                                          set=args.set,
                                          label_folder=args.label_folder)

    target_dataloader = data.DataLoader(target_dataset,
                                        batch_size=args.batch_size,
                                        shuffle=True,
                                        num_workers=args.num_workers,
                                        pin_memory=True)

    return target_dataloader


import numpy as np
import cv2
import glob
import os
import json
from PIL import Image

class Cam:
    """
    Class for the camera object
    """
    def __init__(self, calib_file=None, picDir=None, firstPic=None, pic_to_cm=None,
                 framerate=None, camRes=None, res=None, cropsize=0, cam_thres=20.):
        """Camera object initialisation

        :param calib_file: calibration file for the camera
        :param picDir: shot pictures directory
        :param firstPic: first picture index
        :param pic_to_cm: pixel to cm ratio
        :param framerate: camera framerate
        :param camRes: camera resolution
        :param res: picture resolution (W, H). Warning, it is the picture resolution written on each picture e.g. not accounting for the banner.
        :param cropsize: size of the screen to crop (usefull when pictures information was written on each pictures)
        :param cam_thres: shot detection trheshold
        """
        if not(calib_file is None):
            self.load_calibration_file(calib_file)
        self.dir = picDir
        self.firstPic = firstPic
        self.pic_to_cm = pic_to_cm
        self.framerate = framerate
        self.cropSize = cropsize
        self.res = res
        self.camRes = camRes
        if not(res is None) and not(camRes is None):
            self.set_crop_size()
        self.cam_thres = cam_thres
        self.mask_w = 0
        self.mask_h = 0

    def undistort(self):
        """Undistort the camera pictures and change the picture file to the undistorted one"""

        if "corrected" in self.dir.split("/"):
            print("The camera was already calibrated.. exiting")
        else:
            images = glob.glob(self.dir+'/*.jpg')
            if not("corrected" in os.listdir(self.dir)):
                print('mkdir "'+self.dir+'/corrected"')
                os.system('mkdir "'+self.dir+'/corrected"')
            for elem in images:
                img = cv2.imread(elem)
                h, w = img.shape[:2]
                newcameramtx, roi = cv2.getOptimalNewCameraMatrix(self.mtx, self.dist, (w, h), 0, (w, h))
                dst = cv2.undistort(img, self.mtx, self.dist, None, newcameramtx)
                dst = cv2.cvtColor(dst, cv2.COLOR_BGR2GRAY)
                cv2.imwrite(self.dir+'/corrected/'+elem.split('\\')[1], dst)

            self.dir = self.dir+"/corrected"

    def set_mask(self, mask_w, mask_h):
        """Set black mask to apply on each picture to remove reflexive surfaces

        :param mask_w, mask_h: mask width and height in pixels
        """
        self.mask_w = mask_w
        self.mask_h = mask_h

    def set_mtx(self, mtx):
        """Set camera intrinsic matrix for a given path name (deprecated)

        :param mtx: camera intrinsic matrix file name
        """
        self.mtx = np.loadtxt(mtx)

    def set_dist(self, dist):
        """Set camera distorsion matrix for a given path name (deprecated)

        :param dist: camera distorsion matrix file name
        """
        self.dist = np.loadtxt(dist)

    def set_R(self, R):
        """Set camera rotation matrix for a given Rodrigues vector

        :param R: camera rotation Rodrigues vector
        """
        self.R = np.zeros((3, 3))
        cv2.Rodrigues(R, self.R)

    def set_R_by_matrix(self, R):
        """Set camera rotation matrix for a given matrix

        :param R: camera rotation matrix
        """
        self.R = R

    def set_T(self, T):
        """Set camera Translation vector

        :param T: camera Translation vector
        """
        self.T = np.loadtxt(T)

    def write_cam_data(self):
        """Writes all camera data into a formated tring

        :return: String containing all relevant camera data
        """
        out_str = ''
        out_str += "Screen resolution:\n" + json.dumps(self.camRes) + '\n'
        out_str += "Acquisition resolution:\n" + json.dumps(self.res) + '\n'
        out_str += "Crop array:\n" + json.dumps(self.cropSize) + '\n'
        out_str += "Intrisinc matrix:\n" + json.dumps(self.mtx.tolist())+'\n'
        out_str += "Distorsion matrix:\n" + json.dumps(self.dist.tolist())+'\n'
        out_str += "Rotation matrix:\n" + json.dumps(self.R.tolist()) + '\n'
        out_str += "Translation vector:\n" + json.dumps(self.T.tolist()) + '\n'
        out_str += "Picture directory: \n"+self.dir + '\n'
        out_str += "First picture: \n"+str(self.firstPic) + '\n'
        return out_str

    def load_calibration_file(self, f_name):
        """Load camera intrinsic, distorsion and transformation matrices from a calibration file

        :param f_name:calibration file path
        """
        fichier = open(f_name)
        lines = fichier.read().split('\n')
        fichier.close()

        self.mtx = np.matrix(json.loads(lines[1]))
        self.dist = np.matrix(json.loads(lines[3]))
        R = np.matrix(json.loads(lines[5]))
        self.set_R(R)
        self.T = np.array(json.loads(lines[7])[0])

    def load_from_string(self, data):
        """Initialize a camera object from a formatted string such as produced by write_cam_data

        :param data: formated string to parse
        """
        lines = data.split('\n')
        self.camRes = tuple(json.loads(lines[1]))
        self.res = tuple(json.loads(lines[3]))
        self.cropSize = json.loads(lines[5])
        self.mtx = np.matrix(json.loads(lines[7]))
        self.dist = np.matrix(json.loads(lines[9]))
        self.R = np.matrix(json.loads(lines[11]))
        self.T = np.matrix(json.loads(lines[13]))
        self.dir = lines[15]
        self.firstPic = int(lines[17])

    def set_crop_size(self):
        """Set camera crop size to remove the banner according to the picture effective resolution and the picture target resolution (cam.res)"""
        picList = glob.glob(self.dir + "/*.tif")
        picList += glob.glob(self.dir + "/*.jpg")
        if len(picList) == 0:
            raise NameError("Empty camera picture file :"+self.dir)
        else:
            testpic = np.array(Image.open(picList[0]))
            hor_crop = int(np.abs(testpic.shape[1] - self.res[0])/2.)
            self.cropSize = [hor_crop, hor_crop, 0, testpic.shape[0] - self.res[1]]


import numpy as np
import os
from PIL import Image
import torch
import scipy.io as scio
import sys
from graspnetAPI import GraspGroup
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(ROOT_DIR)
from utils.loss_utils import batch_viewpoint_params_to_matrix
from utils.data_utils import CameraInfo
from utils.collision_detector import ModelFreeCollisionDetector
from rgbmatter.util import create_point_cloud_from_depth_image, heatmap_to_xyz_ori, collision_detection_with_full_models
from rgbmatter.configs import get_config_rgbmatter

if __name__ == '__main__':
    config = get_config_rgbmatter()

    if not os.path.exists(config['scene_res_6Dpose_path']):
        os.makedirs(config['scene_res_6Dpose_path'])

    for i in range(3, 256):   # 256 images
        print('processing scene {}, image {}\n'.format(config['scene_id'], i))
        grasp_gt = np.load(os.path.join(config['grasp_gt_path'], config['scene_id_str'], config['camera'],
                                        str(i).zfill(4) + '.npy'))
        # for this format, rotation matrix's first column is approaching vector because gripper's orientation is x-axis
        up_inds = grasp_gt[:, 10] > 0
        grasp_gt = grasp_gt[up_inds]

        heatmap = np.load(os.path.join(config['scene_heatmap_path'], str(i).zfill(4) + '.npy'))
        depth = np.array(Image.open(os.path.join(config['dataset_path'], 'scenes', config['scene_id_str'],
                                                 config['camera'], 'depth', str(i).zfill(4) + '.png')))   # mm
        meta = scio.loadmat(os.path.join(config['dataset_path'], 'scenes', config['scene_id_str'], config['camera'],
                                         'meta', str(i).zfill(4) + '.mat'))
        intrinsic = meta['intrinsic_matrix']
        factor_depth = meta['factor_depth']

        # generate cloud
        camera = CameraInfo(1280.0, 720.0, intrinsic[0][0], intrinsic[1][1], intrinsic[0][2], intrinsic[1][2],
                            factor_depth)
        cloud = create_point_cloud_from_depth_image(depth, camera, organized=True)  # 720 * 1280 * 3 numpy.ndarray, m
        incomplete_pose = heatmap_to_xyz_ori(heatmap, cloud, grasp_gt)
        xyz = incomplete_pose[:, :3]
        print('grasp num before depth and width sampling: ', xyz.shape[0])

        approaching_vector = incomplete_pose[:, 3:6]
        angle = incomplete_pose[:, 6]

        # gt_width = incomplete_pose[:, 7]  # use gt width and depth
        # gt_depth = incomplete_pose[:, 8]
        # grasp_width = torch.from_numpy(gt_width).unsqueeze(-1)
        # grasp_depth = torch.from_numpy(gt_depth).unsqueeze(-1)
        grasp_width = torch.arange(3, 11) * 0.01  # sample grasp width , 10
        grasp_dist = torch.arange(-2, 3) * 0.01
        # grasp_depth = torch.arange(-2, 3) * 0.01  # sample grasp depth, 5

        grasp_num = incomplete_pose.shape[0]  # grasp num in heatmap
        grasp_score = torch.ones((grasp_num, 1))
        ori = batch_viewpoint_params_to_matrix(torch.from_numpy(approaching_vector), torch.from_numpy(angle))
        ori = ori.view(grasp_num, 9)
        xyz = torch.from_numpy(xyz)

        grasp_width = grasp_width.unsqueeze(-1).repeat(grasp_num, 5).view(-1, 1)  # 50 combinations of width and depth
        # grasp_depth = grasp_depth.unsqueeze(-1).repeat(grasp_num * 8, 1).view(-1, 1)
        grasp_dist = grasp_dist.unsqueeze(-1).repeat(grasp_num * 8, 1).view(-1, 1)
        ori = ori.repeat(1, 40).view(-1, 9)  # each one has 50 combinations of width and depth
        xyz = xyz.repeat(1, 40).view(-1, 3)  # each one 50 combination of width and depth
        # grasp_dist = torch.cat([grasp_dist, torch.zeros_like(grasp_dist), torch.zeros_like(grasp_dist)],
        #                        axis=-1).unsqueeze(-1).double()
        # grasp_dist = torch.bmm(ori.view(-1, 3, 3).double(), grasp_dist)
        # xyz = xyz + grasp_dist.squeeze()
        xyz[:, 2] = xyz[:, 2] + grasp_dist.squeeze()
        grasp_score = grasp_score.repeat(1, 40).view(-1, 1)  # each one 50 combination of width and depth

        obj_ids = -1 * torch.ones_like(grasp_score)
        grasp_height = 0.02 * torch.ones_like(grasp_score)
        grasp_depth = 0.03 * torch.ones_like(grasp_score)
        grasp_decode = torch.cat([grasp_score, grasp_width, grasp_height, grasp_depth, ori, xyz, obj_ids],
                                 axis=-1).numpy()
        print('grasp after sampling: ', grasp_decode.shape)

        if config['collision_detection_choice'] == 'point_cloud':
            gg = GraspGroup(grasp_decode)
            mfcdetector = ModelFreeCollisionDetector(cloud.reshape([-1, 3]),  # collision detection use point cloud
                                                     voxel_size=0.08)  # filter the collision and empty grasps
            collision_mask = mfcdetector.detect(gg, approach_dist=0.03, collision_thresh=0.0, empty_thresh=0.3,
                                                return_empty_grasp=True)
            final_gg = gg[np.logical_and(~collision_mask[0], ~collision_mask[1])]
        else:
            final_gg = GraspGroup(np.empty([0, 17]))
            for j in range(int(np.ceil(grasp_decode.shape[0] / 6000))):
                start = j * 6000
                print('processing grasps {} ~ {}'.format(start, start + 6000))
                gg = GraspGroup(grasp_decode[start:start + 6000])
                gg = collision_detection_with_full_models(gg, i)  # collision detection use full object models
                final_gg.grasp_group_array = np.concatenate([final_gg.grasp_group_array, gg.grasp_group_array])

        print('grasp shape after filter: ', final_gg.grasp_group_array.shape)  # (39999, 17)  why
        save_path = os.path.join(config['scene_res_6Dpose_path'], str(i).zfill(4) + '.npy')
        final_gg.save_npy(save_path)
        print('save {} successfully!!!'.format(save_path))
        break


# -*- coding: utf-8 -*-

# Contains the "Note" class used in our program, that represent a note
# and keeps various attributes assigned to it. It also contains static
# methods to get a random note, or convert a whole list.

import random
import re
from time import sleep

import numpy as np
import simpleaudio as sa

# List of notes and figures
NOTE_NAMES = ['DO', 'RE', 'MI', 'FA', 'SOL', 'LA', 'SI']
NOTE_FIGURES = ['r', 'b', 'n', 'c']

class Note:
    def __init__(self, raw_input):
        self.raw_input = raw_input

        self.parsed_data = self.parse()

        # Name of the note (one of NOTE_NAMES)
        self.name = self.parsed_data[0]
        # Figure of the note (one of NOTE_FIGURES)
        self.figure = self.parsed_data[1]
        # Wether the note has its duration extended with a point
        self.has_point = bool(self.parsed_data[2])
        # The frequency of the note
        self.frequency = self.get_frequency()
        # The duration of the note in seconds
        self.duration = self.get_duration()
        # Wether the note is a pause
        self.is_pause = self.name == "Z"

    @staticmethod
    def create_random_note():
        """ Create a random note. """
        raw_note = random.choice(NOTE_NAMES)
        raw_note += random.choice(NOTE_FIGURES)
        return Note(raw_note)

    @staticmethod
    def to_raw(parsed_notes):
        """ Transform a list of Note instances into a list of raw notes, as in the partition files """
        output = []
        for note in parsed_notes:
            new_note = f'{note.name}{note.figure}'
            new_note += 'p' if note.has_point else ''
            output.append(new_note)
        return output

    def parse(self):
        """
        Parse notes in the partitions, and returns a tuple containing the name,
        the duration and whether it contains a point.
        """
        groups = re.findall("([A-Z]*)(r|b|n|c)(p)?", self.raw_input)
        return groups[0]

    def get_frequency(self):
        """ Assign a frequency for each available notes """
        if self.name == "DO":
            return 264
        elif self.name == "RE":
            return 297
        elif self.name == "MI":
            return 330
        elif self.name == "FA":
            return 352
        elif self.name == "SOL":
            return 396
        elif self.name == "LA":
            return 440
        elif self.name == "SI":
            return 495
        return 0

    def get_duration(self):
        """ Get the duration depending on the figure, and if there is a point """
        duration = 0
        if self.figure == 'r':    # "Ronde"
            duration = 1000
        elif self.figure == 'b':  # "Blanche"
            duration = 500
        elif self.figure == 'n':  # "Noire"
            duration = 250
        elif self.figure == 'c':  # "Croche"
            duration = 125
        # If there is a point, set it to half the duration, otherwise 0
        point_duration = duration / 2 if self.has_point else 0
        # Add the duration and the extended duration (point), and normalize it
        return (duration + point_duration) / 1000

    def play(self):
        """ If there is a frequency, play the note, otherwise it is a pause """
        if self.is_pause:
            # Sleep for the duration of the note when the note is "Z"
            sleep(self.duration)
        else:
            # Get timesteps for each sample, "duration" is note duration in seconds
            sample_rate = 44100
            t = np.linspace(0, self.duration, int(self.duration * sample_rate), False)
            # Generate sine wave tone
            tone = np.sin(self.frequency * t * 6 * np.pi)
            # Normalize to 24−bit range
            tone *= 8388607 / np.max(np.abs(tone))
            # Convert to 32−bit data
            tone = tone.astype(np.int32)

            # Convert from 32−bit to 24−bit by building a new byte buffer,
            # skipping every fourth bit
            # Note: this also works for 2−channel audio
            i = 0
            byte_array = []
            for b in tone.tobytes():
                if i % 4 != 3:
                    byte_array.append(b)
                i += 1

            audio = bytearray(byte_array)
            # Start playback
            play_obj = sa.play_buffer(audio, 1, 3, sample_rate)
            # Wait for playback to finish before exiting
            play_obj.wait_done()