# coding: utf8
# cython: profile=True
# cython: infer_types=True
"""Coref resolution"""

from __future__ import unicode_literals
from __future__ import print_function

import sys
import os
import spacy
import numpy as np

from neuralcoref.utils import PACKAGE_DIRECTORY, SIZE_PAIR_IN, SIZE_SINGLE_IN
from neuralcoref.compat import unicode_
from neuralcoref.document import Document, MENTION_TYPE, NO_COREF_LIST

#######################
##### UTILITIES #######

MAX_FOLLOW_UP = 50

#######################
###### CLASSES ########

class Model(object):
    '''
    Coreference neural model
    '''
    def __init__(self, model_path):
        weights, biases = [], []
        for file in sorted(os.listdir(model_path)):
            if file.startswith("single_mention_weights"):
                w = np.load(os.path.join(model_path, file))
                weights.append(w)
            if file.startswith("single_mention_bias"):
                w = np.load(os.path.join(model_path, file))
                biases.append(w)
        self.single_mention_model = list(zip(weights, biases))
        weights, biases = [], []
        for file in sorted(os.listdir(model_path)):
            if file.startswith("pair_mentions_weights"):
                w = np.load(os.path.join(model_path, file))
                weights.append(w)
            if file.startswith("pair_mentions_bias"):
                w = np.load(os.path.join(model_path, file))
                biases.append(w)
        self.pair_mentions_model = list(zip(weights, biases))

    def _score(self, features, layers):
        for weights, bias in layers:
            #print("features", features.shape)
            features = np.matmul(weights, features) + bias
            if weights.shape[0] > 1:
                features = np.maximum(features, 0) # ReLU
        return np.sum(features, axis=0)

    def get_multiple_single_score(self, first_layer_input):
        return self._score(first_layer_input, self.single_mention_model)

    def get_multiple_pair_score(self, first_layer_input):
        return self._score(first_layer_input, self.pair_mentions_model)

class Coref(object):
    '''
    Main coreference resolution algorithm
    '''
    def __init__(self, nlp=None, greedyness=0.5, max_dist=50, max_dist_match=500, conll=None,
                 blacklist=True, debug=False):
        self.greedyness = greedyness
        self.max_dist = max_dist
        self.max_dist_match = max_dist_match
        self.debug = debug
        model_path = os.path.join(PACKAGE_DIRECTORY, "weights/conll/" if conll is not None else "weights/")
        model_path = os.path.join(PACKAGE_DIRECTORY, "weights/")
        print("Loading neuralcoref model from", model_path)
        self.coref_model = Model(model_path)
        if nlp is None:
            print("Loading spacy model")
            try:
                spacy.info('en_core_web_sm')
                model = 'en_core_web_sm'
            except IOError:
                print("No spacy 2 model detected, using spacy1 'en' model")
                spacy.info('en')
                model = 'en'
            nlp = spacy.load(model)
        self.data = Document(nlp, conll=conll, blacklist=blacklist, model_path=model_path)
        self.clusters = {}
        self.mention_to_cluster = []
        self.mentions_single_scores = {}
        self.mentions_pairs_scores = {}

    ###################################
    #### ENTITY CLUSTERS FUNCTIONS ####
    ###################################

    def _prepare_clusters(self):
        '''
        Clean up and prepare one cluster for each mention
        '''
        self.mention_to_cluster = list(range(len(self.data.mentions)))
        self.clusters = dict((i, [i]) for i in self.mention_to_cluster)
        self.mentions_single_scores = {}
        self.mentions_pairs_scores = {}
        for mention in self.mention_to_cluster:
            self.mentions_single_scores[mention] = None
            self.mentions_pairs_scores[mention] = {}

    def _merge_coreference_clusters(self, ant_idx, mention_idx):
        '''
        Merge two clusters together
        '''
        if self.mention_to_cluster[ant_idx] == self.mention_to_cluster[mention_idx]:
            return

        remove_id = self.mention_to_cluster[ant_idx]
        keep_id = self.mention_to_cluster[mention_idx]
        for idx in self.clusters[remove_id]:
            self.mention_to_cluster[idx] = keep_id
            self.clusters[keep_id].append(idx)

        del self.clusters[remove_id]

    def remove_singletons_clusters(self):
        remove_id = []
        for key, mentions in self.clusters.items():
            if len(mentions) == 1:
                remove_id.append(key)
                self.mention_to_cluster[key] = None
        for rem in remove_id:
            del self.clusters[rem]

    def display_clusters(self):
        '''
        Print clusters informations
        '''
        print(self.clusters)
        for key, mentions in self.clusters.items():
            print("cluster", key, "(", ", ".join(unicode_(self.data[m]) for m in mentions), ")")

    ###################################
    ####### MAIN COREF FUNCTIONS ######
    ###################################

    def run_coref_on_mentions(self, mentions):
        '''
        Run the coreference model on a mentions list
        '''
        best_ant = {}
        best_score = {}
        n_ant = 0
        inp = np.empty((SIZE_SINGLE_IN, len(mentions)))
        for i, mention_idx in enumerate(mentions):
            mention = self.data[mention_idx]
            print("mention.embedding", mention.embedding.shape)
            inp[:len(mention.embedding), i] = mention.embedding
            inp[:len(mention.embedding), i] = mention.features
            inp[:len(mention.embedding), i] = self.data.genre
        score = self.coref_model.get_multiple_single_score(inp.T)
        for mention_idx, s in zip(mentions, score):
            self.mentions_single_scores[mention_idx] = s
            best_score[mention_idx] = s - 50 * (self.greedyness - 0.5)

        for mention_idx, ant_list in self.data.get_candidate_pairs(mentions, self.max_dist, self.max_dist_match):
            if len(ant_list) == 0:
                continue
            inp_l = []
            for ant_idx in ant_list:
                mention = self.data[mention_idx]
                antecedent = self.data[ant_idx]
                feats_, pwf = self.data.get_pair_mentions_features(antecedent, mention)
                inp_l.append(pwf)
            inp = np.stack(inp_l, axis=0)
            score = self.coref_model.get_multiple_pair_score(inp.T)
            for ant_idx, s in zip(ant_list, score):
                self.mentions_pairs_scores[mention_idx][ant_idx] = s
                if s > best_score[mention_idx]:
                    best_score[mention_idx] = s
                    best_ant[mention_idx] = ant_idx
            if mention_idx in best_ant:
                n_ant += 1
                self._merge_coreference_clusters(best_ant[mention_idx], mention_idx)
        return (n_ant, best_ant)

    def run_coref_on_utterances(self, last_utterances_added=False, follow_chains=True, debug=False):
        ''' Run the coreference model on some utterances

        Arg:
            last_utterances_added: run the coreference model over the last utterances added to the data
            follow_chains: follow coreference chains over previous utterances
        '''
        if debug: print("== run_coref_on_utterances == start")
        self._prepare_clusters()
        if debug: self.display_clusters()
        mentions = list(self.data.get_candidate_mentions(last_utterances_added=last_utterances_added))
        n_ant, antecedents = self.run_coref_on_mentions(mentions)
        mentions = antecedents.values()
        if follow_chains and last_utterances_added and n_ant > 0:
            i = 0
            while i < MAX_FOLLOW_UP:
                i += 1
                n_ant, antecedents = self.run_coref_on_mentions(mentions)
                mentions = antecedents.values()
                if n_ant == 0:
                    break
        if debug: self.display_clusters()
        if debug: print("== run_coref_on_utterances == end")

    def one_shot_coref(self, utterances, utterances_speakers_id=None, context=None,
                       context_speakers_id=None, speakers_names=None):
        ''' Clear history, load a list of utterances and an optional context and run the coreference model on them

        Arg:
        - `utterances` : iterator or list of string corresponding to successive utterances (in a dialogue) or sentences.
            Can be a single string for non-dialogue text.
        - `utterances_speakers_id=None` : iterator or list of speaker id for each utterance (in the case of a dialogue).
            - if not provided, assume two speakers speaking alternatively.
            - if utterances and utterances_speaker are not of the same length padded with None
        - `context=None` : iterator or list of string corresponding to additionnal utterances/sentences sent prior to `utterances`. Coreferences are not computed for the mentions identified in `context`. The mentions in `context` are only used as possible antecedents to mentions in `uterrance`. Reduce the computations when we are only interested in resolving coreference in the last sentences/utterances.
        - `context_speakers_id=None` : same as `utterances_speakers_id` for `context`.
        - `speakers_names=None` : dictionnary of list of acceptable speaker names (strings) for speaker_id in `utterances_speakers_id` and `context_speakers_id`
        Return:
            clusters of entities with coreference resolved
        '''
        self.data.set_utterances(context, context_speakers_id, speakers_names)
        self.continuous_coref(utterances, utterances_speakers_id, speakers_names)
        return self.get_clusters()

    def continuous_coref(self, utterances, utterances_speakers_id=None, speakers_names=None):
        '''
        Only resolve coreferences for the mentions in the utterances
        (but use the mentions in previously loaded utterances as possible antecedents)
        Arg:
            utterances : iterator or list of string corresponding to successive utterances
            utterances_speaker : iterator or list of speaker id for each utterance.
                If not provided, assume two speakers speaking alternatively.
                if utterances and utterances_speaker are not of the same length padded with None
            speakers_names : dictionnary of list of acceptable speaker names for each speaker id
        Return:
            clusters of entities with coreference resolved
        '''
        self.data.add_utterances(utterances, utterances_speakers_id, speakers_names)
        self.run_coref_on_utterances(last_utterances_added=True, follow_chains=True)
        return self.get_clusters()

    ###################################
    ###### INFORMATION RETRIEVAL ######
    ###################################

    def get_utterances(self, last_utterances_added=True):
        ''' Retrieve the list of parsed uterrances'''
        if last_utterances_added and len(self.data.last_utterances_loaded):
            return [self.data.utterances[idx] for idx in self.data.last_utterances_loaded]
        else:
            return self.data.utterances

    def get_resolved_utterances(self, last_utterances_added=True, blacklist=True):
        ''' Return a list of utterrances text where the '''
        coreferences = self.get_most_representative(last_utterances_added, blacklist)
        resolved_utterances = []
        for utt in self.get_utterances(last_utterances_added=last_utterances_added):
            resolved_utt = ""
            in_coref = None
            for token in utt:
                if in_coref is None:
                    for coref_original, coref_replace in coreferences.items():
                        if coref_original[0] == token:
                            in_coref = coref_original
                            resolved_utt += coref_replace.text.lower()
                            break
                    if in_coref is None:
                        resolved_utt += token.text_with_ws
                if in_coref is not None and token == in_coref[-1]:
                    resolved_utt += ' ' if token.whitespace_ and resolved_utt[-1] is not ' ' else ''
                    in_coref = None
            resolved_utterances.append(resolved_utt)
        return resolved_utterances

    def get_mentions(self):
        ''' Retrieve the list of mentions'''
        return self.data.mentions

    def get_scores(self):
        ''' Retrieve scores for single mentions and pair of mentions'''
        return {"single_scores": self.mentions_single_scores,
                "pair_scores": self.mentions_pairs_scores}

    def get_clusters(self, remove_singletons=False, blacklist=False):
        ''' Retrieve cleaned clusters'''
        clusters = self.clusters
        mention_to_cluster = self.mention_to_cluster
        remove_id = []
        if blacklist:
            for key, mentions in clusters.items():
                cleaned_list = []
                for mention_idx in mentions:
                    mention = self.data.mentions[mention_idx]
                    if mention.lower_ not in NO_COREF_LIST:
                        cleaned_list.append(mention_idx)
                clusters[key] = cleaned_list
            # Also clean up keys so we can build coref chains in self.get_most_representative
            added = {}
            for key, mentions in clusters.items():
                if self.data.mentions[key].lower_ in NO_COREF_LIST:
                    remove_id.append(key)
                    mention_to_cluster[key] = None
                    if mentions:
                        added[mentions[0]] = mentions
            for rem in remove_id:
                del clusters[rem]
            clusters.update(added)

        if remove_singletons:
            remove_id = []
            for key, mentions in clusters.items():
                if len(mentions) == 1:
                    remove_id.append(key)
                    mention_to_cluster[key] = None
            for rem in remove_id:
                del clusters[rem]

        return clusters, mention_to_cluster

    def get_most_representative(self, last_utterances_added=True, blacklist=True):
        '''
        Find a most representative mention for each cluster

        Return:
            Dictionnary of {original_mention: most_representative_resolved_mention, ...}
        '''
        clusters, _ = self.get_clusters(remove_singletons=True, blacklist=blacklist)
        coreferences = {}
        for key in self.data.get_candidate_mentions(last_utterances_added=last_utterances_added):
            if self.mention_to_cluster[key] is None:
                continue
            mentions = clusters.get(self.mention_to_cluster[key], None)
            if mentions is None:
                continue
            representative = self.data.mentions[key]
            for mention_idx in mentions[1:]:
                mention = self.data.mentions[mention_idx]
                if mention.mention_type is not representative.mention_type:
                    if mention.mention_type == MENTION_TYPE["PROPER"] \
                        or (mention.mention_type == MENTION_TYPE["NOMINAL"] and
                                representative.mention_type == MENTION_TYPE["PRONOMINAL"]):
                        coreferences[self.data.mentions[key]] = mention
                        representative = mention

        return coreferences


# -*- mode: python; coding: utf-8 -*
# Copyright (c) 2018 Radio Astronomy Software Group
# Licensed under the 3-clause BSD License

"""
The primary machinery of pyuvsim.

The :class:`~UVTask` and :class:`~UVEngine` classes and the functions that actually run
the simulation.
"""

import numpy as np
import yaml
import warnings
import astropy.units as units
from astropy.coordinates import EarthLocation
from astropy.units import Quantity
from astropy.constants import c as speed_of_light
from pyuvdata import UVData

from . import mpi
from . import simsetup
from . import utils as simutils
from .antenna import Antenna
from .baseline import Baseline
from .telescope import Telescope
from .simsetup import SkyModelData

from .astropy_interface import MoonLocation, hasmoon, Time

__all__ = ['UVTask', 'UVEngine', 'uvdata_to_task_iter', 'run_uvsim', 'run_uvdata_uvsim']

class UVTask:
    """
    An object to hold all the information to calculate a visibility for a set of sources.

    Parameters
    ----------
    sources : :class:`pyuvsim.simsetup.SkyModelData`
        The sources to include in the visibility.
    time : :class:`astropy.time.Time` object or float
        Time at which to calculate the visibility, either an astropy Time object or a
        float in Julian Day (will be converted to a Time object using `format='jd'`).
    freq : :class:`astropy.units.Quantity` or float
        Frequency at which to calculate the visibility, either an astropy Quantity with
        units compatible with Hz or a float in Hz (will be converted to an astropy
        Quantity with units of Hz)
    baseline : :class:`pyuvsim.Baseline`
        The baseline to calculate the visibility for.
    telescope : :class:`pyuvsim.Telescope`
        The telescope object this visibility belongs to, which carries the telescope
        location and beam information.
    freq_i : int
        Frequency index for this visibility's source coherency.

    Attributes
    ----------
    sources : :class:`pyuvsim.simsetup.SkyModelData`
        The sources to include in the visibility.
    time : :class:`astropy.time.Time`
        Time at which to calculate the visibility.
    freq : :class:`astropy.units.Quantity`
        Frequency at which to calculate the visibility.
    baseline : :class:`pyuvsim.Baseline`
        The baseline to calculate the visibility for.
    telescope : :class:`pyuvsim.Telescope`
        The telescope object this visibility belongs to.
    freq_i : int
        Frequency index for this visibility's source coherency. Always `0` for flat
        spectrum sources, otherwise should be the index of `freq` in the full frequency
        array for this simulation.
    uvdata_index : tuple of int
        Indexes for this visibility location in the uvdata data_array, length 3 giving
        (blt_index, 0, frequency_index), where the zero index is for the old spw axis.
    visibility_vector : ndarray of complex
        The calculated visibility, shape (4,) ordered as [xx, yy, xy, yx].

    """

    def __init__(self, sources, time, freq, baseline, telescope, freq_i=0):
        self.time = time
        self.freq = freq
        self.sources = sources  # SkyModel object
        self.baseline = baseline
        self.telescope = telescope
        self.freq_i = freq_i
        self.visibility_vector = None
        self.uvdata_index = None  # Where to add the visibility in the uvdata object.

        if isinstance(self.time, float):
            self.time = Time(self.time, format='jd')
        if isinstance(self.freq, float):
            self.freq = self.freq * units.Hz
        if sources.spectral_type == 'flat':
            self.freq_i = 0

    def __eq__(self, other):
        """
        Test for equality between two UVTask objects.

        Parameters
        ----------
        other : :class:`UVTask`
            The other

        """
        return (np.isclose(self.time.jd, other.time.jd, atol=1e-4)
                and np.isclose(self.freq.value, other.freq.value, atol=1e-4)
                and (self.sources == other.sources)
                and (self.baseline == other.baseline)
                and (self.visibility_vector == other.visibility_vector)
                and (self.uvdata_index == other.uvdata_index)
                and (self.telescope == other.telescope))

    def __gt__(self, other):
        """
        Check if this UVTask has a larger `uvdata_index` than another one.

        Parameters
        ----------
        other : :class:`UVTask`
            The other

        """
        blti0, fi0 = self.uvdata_index
        blti1, fi1 = other.uvdata_index
        if self.baseline == other.baseline:
            if fi0 == fi1:
                return blti0 > blti1
            return fi0 > fi1
        return self.baseline > other.baseline

    def __ge__(self, other):
        """
        Check if this UVTask has a larger or equal `uvdata_index` than another one.

        Parameters
        ----------
        other : :class:`UVTask`
            The other

        """
        blti0, fi0 = self.uvdata_index
        blti1, fi1 = other.uvdata_index
        if self.baseline == other.baseline:
            if fi0 == fi1:
                return blti0 >= blti1
            return fi0 >= fi1
        return self.baseline >= other.baseline

    def __lt__(self, other):
        """
        Check if this UVTask has a smaller `uvdata_index` than another one.

        Parameters
        ----------
        other : :class:`UVTask`
            The other

        """
        return not self.__ge__(other)

    def __le__(self, other):
        """
        Check if this UVTask has a smaller or equal `uvdata_index` than another one.

        Parameters
        ----------
        other : :class:`UVTask`
            The other

        """
        return not self.__gt__(other)

class UVEngine:
    """
    The object where the visibility calculations actually take place.

    Parameters
    ----------
    task : :class:`~UVTask`
        The task to do the calculations for.
    reuse_spline : bool
        Option to reuse the spline in the beam interpolation to save time.
    update_positions : bool
        Flag indicating that positions need to be updated (when time or source changes).
    update beams : bool
        Flag indicating that beams need to be updated (when time, sources, frequency,
        or beam pair changes).

    Attributes
    ----------
    reuse_spline : bool
        Option to reuse the spline in the beam interpolation to save time.
    update_positions : bool
        Flag indicating that positions need to be updated (when time or source changes).
    update_local_coherency : bool
        Flag indicating that local coherencies need to be updated (when time or source changes).
    update beams : bool
        Flag indicating that beams need to be updated (when time, sources, frequency,
        or beam pair changes).
    sources : :class:`pyuvsim.simsetup.SkyModelData`
        The sources to include in the visibility.
    current_time : :class:`astropy.time.Time`
        Time for the current calculation.
    current_freq : :class:`astropy.units.Quantity`
        Frequency for the current calculation, units compatible with Hz.
    current_beam_pair : 2-tuple of ints
        Tuple containing the beam ids for the current calculation.
    beam1_jones : array_like of float
        Jones matricies for the first beam for each source location, shape
        (2,2, Ncomponents). The first axis is feed, the second axis is vector component
        on the sky in az/za.
    beam2_jones : array_like of float
        Jones matricies for the second beam for each source location, shape
        (2,2, Ncomponents). The first axis is feed, the second axis is vector component
        on the sky in az/za.
    local_coherency : array_like of float
        Source local coherencies in alt/az basis, shape (2, 2, Nfreqs, Ncomponents).
    apparent_coherency : array_like of float
        Source apparent coherencies (including the beam response), shape
        (2, 2, Nfreqs, Ncomponents).
    task :  :class:`~UVTask`
        The task currently being calculated.
    """

    def __init__(self, task=None, update_positions=True, update_beams=True, reuse_spline=True):
        self.reuse_spline = reuse_spline  # Reuse spline fits in beam interpolation
        self.update_positions = update_positions
        self.update_beams = update_beams

        self.sources = None
        self.current_time = None
        self.current_freq = None
        self.current_beam_pair = None
        self.beam1_jones = None
        self.beam2_jones = None
        self.local_coherency = None
        self.apparent_coherency = None

        if task is not None:
            self.set_task(task)

    def set_task(self, task):
        """
        Set the task for the Engine.

        Parameters
        ----------
        task : :class:`UVTask`
            The task to do the calculations for.

        """
        self.task = task

        # These flags control whether quantities can be re-used from the last task:
        #   Update local coherency for all frequencies when time or sources changes.
        #   Update source positions when time or sources changes.
        #   Update saved beam jones matrices when time, sources, frequency, or beam pair changes.

        baseline = self.task.baseline
        beam_pair = (baseline.antenna1.beam_id, baseline.antenna2.beam_id)

        if (not self.current_time == task.time.jd) or (self.sources is not task.sources):
            self.update_positions = True
            self.update_local_coherency = True
            self.update_beams = True
        else:
            self.update_beams = False
            self.update_positions = False
            self.update_local_coherency = False

        if not self.current_freq == task.freq.to('Hz').value:
            self.update_beams = True

        if not self.current_beam_pair == beam_pair:
            self.current_beam_pair = beam_pair
            self.update_beams = True

        self.current_time = task.time.jd
        self.current_freq = task.freq.to("Hz").value
        self.sources = task.sources

    def apply_beam(self):
        """Set apparent coherency from jones matrices and source coherency."""
        beam1_id, beam2_id = self.current_beam_pair

        sources = self.task.sources
        baseline = self.task.baseline

        if sources.alt_az is None:
            sources.update_positions(self.task.time, self.task.telescope.location)

        if self.update_local_coherency:
            self.local_coherency = sources.coherency_calc()

        self.beam1_jones = baseline.antenna1.get_beam_jones(
            self.task.telescope, sources.alt_az[..., sources.above_horizon],
            self.task.freq, reuse_spline=self.reuse_spline
        )

        if beam1_id == beam2_id:
            self.beam2_jones = np.copy(self.beam1_jones)
        else:
            self.beam2_jones = baseline.antenna2.get_beam_jones(
                self.task.telescope, sources.alt_az[..., sources.above_horizon],
                self.task.freq, reuse_spline=self.reuse_spline
            )

        # coherency is a 2x2 matrix
        # [ |Ex|^2, Ex* Ey, Ey* Ex |Ey|^2 ]
        # where x and y vectors along the local alt/az axes.

        # Apparent coherency gives the direction and polarization dependent baseline response to
        # a source.

        coherency = self.local_coherency[:, :, self.task.freq_i, :]

        self.beam2_jones = np.swapaxes(self.beam2_jones, 0, 1).conj()  # Transpose at each component

        self.apparent_coherency = np.einsum(
            "abz,bcz,cdz->adz", self.beam1_jones, coherency, self.beam2_jones
        )

    def make_visibility(self):
        """Calculate visibility contribution from a set of source components."""
        assert (isinstance(self.task.freq, Quantity))
        srcs = self.task.sources
        time = self.task.time
        location = self.task.telescope.location

        if self.update_positions:
            srcs.update_positions(time, location)

        if self.update_beams:
            self.apply_beam()

        pos_lmn = srcs.pos_lmn[..., srcs.above_horizon]

        # need to convert uvws from meters to wavelengths
        uvw_wavelength = self.task.baseline.uvw / speed_of_light * self.task.freq.to('1/s')
        fringe = np.exp(2j * np.pi * np.dot(uvw_wavelength, pos_lmn))
        vij = self.apparent_coherency * fringe

        # Sum over source component axis:
        vij = np.sum(vij, axis=2)

        # Reshape to be [xx, yy, xy, yx]
        vis_vector = np.asarray([vij[0, 0], vij[1, 1], vij[0, 1], vij[1, 0]])
        return vis_vector

def _make_task_inds(Nblts, Nfreqs, Nsrcs, rank, Npus):
    """
    Make iterators defining task and sources computed on rank.

    The splitting depends on the relative sizes of the axes and the number of processing
    units. There are three possible ways the splitting is done:

        - if Npus < Nbltf -- Split by Nbltf, split sources in the task loop for memory's sake.
        - if Nbltf < Npus and Nsrcs > Npus -- Split by Nsrcs only
        - if (Nsrcs, Nbltf) < Npus -- Split by Nbltf

    Parameters
    ----------
    Nblts : int
        Number of blts included in the sim.
    Nfreqs :  int
        Number of frequencies included in the sim.
    Nsrcs : int
        Number of sources included in the sim.
    rank : int
        mpi rank.
    Npus : int
        Number of processing units.

    """
    Nbltf = Nblts * Nfreqs

    split_srcs = False

    if (Nbltf < Npus) and (Npus < Nsrcs):
        split_srcs = True

    if split_srcs:
        src_inds, Nsrcs_local = simutils.iter_array_split(rank, Nsrcs, Npus)
        task_inds = range(Nbltf)
        Ntasks_local = Nbltf
    else:
        task_inds, Ntasks_local = simutils.iter_array_split(rank, Nbltf, Npus)
        src_inds = range(Nsrcs)
        Nsrcs_local = Nsrcs

    return task_inds, src_inds, Ntasks_local, Nsrcs_local

def uvdata_to_task_iter(task_ids, input_uv, catalog, beam_list, beam_dict, Nsky_parts=1):
    """
    Generate UVTask objects.

    Parameters
    ----------
    task_ids : range
        Task indices in the full flattened meshgrid of parameters.
    input_uv : :class:`pyuvdata.UVData`
        UVData object to be filled with data.
    catalog : :class:`pyuvsim.simsetup.SkyModelData`
        Source components.
    beam_list : :class:`pyuvsim.BeamList`
        BeamList carrying beam model (in object mode).
    beam_dict : dict
        Map of antenna numbers to index in beam_list.

    Yields
    ------
        Iterable of UVTask objects.

    """
    # The task_ids refer to tasks on the flattened meshgrid.
    if not isinstance(input_uv, UVData):
        raise TypeError("input_uv must be UVData object.")

    #   Skymodel will now be passed in as a catalog array.
    if not isinstance(catalog, SkyModelData):
        raise TypeError("catalog must be a SkyModelData object.")

    # use future array shapes
    if not input_uv.future_array_shapes:
        input_uv.use_future_array_shapes()

    # Splitting the catalog for memory's sake.
    Nsrcs_total = catalog.Ncomponents
    if Nsky_parts > 1:
        Nsky_parts = int(Nsky_parts)
        src_iter = [simutils.iter_array_split(s, Nsrcs_total, Nsky_parts)[0]
                    for s in range(Nsky_parts)]
    else:
        src_iter = [range(Nsrcs_total)]
    # Build the antenna list.
    antenna_names = input_uv.antenna_names
    antennas = []
    antpos_enu, antnums = input_uv.get_ENU_antpos()
    for num, antname in enumerate(antenna_names):
        if beam_dict is None:
            beam_id = 0
        else:
            beam_id = beam_dict[antname]
        antennas.append(Antenna(antname, num, antpos_enu[num], beam_id))

    baselines = {}
    Nfreqs = input_uv.Nfreqs
    Nblts = input_uv.Nblts

    tasks_shape = (Nblts, Nfreqs)
    blt_ax, freq_ax = range(2)

    # we need to get slice or indices to index into the flat iterators
    # it does not accept range arguments as indices.
    if isinstance(task_ids, range):
        task_slice = slice(task_ids.start, task_ids.stop, task_ids.step)
    else:
        # There is a test where the input ids is an np.arange object
        # direct indexing should be fine.
        task_slice = task_ids

    # broadcast_to returns a view into the array
    # flat returns an iterator object.
    # no overhead from these calls.
    _times = np.broadcast_to(
        input_uv.time_array.reshape(-1, 1),
        (input_uv.Nblts, input_uv.Nfreqs),
    ).flat

    _bls = np.broadcast_to(
        input_uv.baseline_array.reshape(-1, 1),
        (input_uv.Nblts, input_uv.Nfreqs),
    ).flat

    _freqs = np.broadcast_to(
        input_uv.freq_array.flatten().reshape(1, -1),
        (input_uv.Nblts, input_uv.Nfreqs),
    ).flat

    # indexing with the slice should return a view,
    # but the iterator has to be collected into an array.
    # This should incure 64 * Ntasks_local bits per array
    # some additional overhead for numpy arrays as well
    # usually seeing [0.00001, 0.00003] MiB / task memory required
    # based on the reference simulations.
    order = np.lexsort((_bls[task_slice], _freqs[task_slice], _times[task_slice])).flat

    tloc = [np.float64(x) for x in input_uv.telescope_location]

    world = input_uv.extra_keywords.get('world', 'earth')

    if world.lower() == 'earth':
        location = EarthLocation.from_geocentric(*tloc, unit='m')
    elif world.lower() == 'moon':
        if not hasmoon:
            raise ValueError("Need lunarsky module to simulate an array on the Moon.")
        location = MoonLocation.from_selenocentric(*tloc, unit='m')
    else:
        raise ValueError("If world keyword is set, it must be either 'moon' or 'earth'.")
    telescope = Telescope(input_uv.telescope_name, location, beam_list)
    freq_array = input_uv.freq_array * units.Hz
    time_array = Time(input_uv.time_array, scale='utc', format='jd', location=telescope.location)
    for src_i in src_iter:
        sky = catalog.get_skymodel(src_i)
        if (
            sky.spectral_type == 'flat'
            and sky.freq_array is None
            and sky.reference_frequency is None
        ):
            sky.freq_array = freq_array
        if sky.component_type == 'healpix' and hasattr(sky, 'healpix_to_point'):
            sky.healpix_to_point()
        if sky.spectral_type != 'flat':
            sky.at_frequencies(freq_array)

        for index in order:
            task_index = task_ids[index]
            # Shape indicates slowest to fastest index.
            if not isinstance(task_index, tuple):
                task_index = np.unravel_index(task_index, tasks_shape)
            blt_i = task_index[blt_ax]
            freq_i = task_index[freq_ax]

            bl_num = input_uv.baseline_array[blt_i]

            # We reuse a lot of baseline info, so make the baseline list on the first go and reuse.
            if bl_num not in baselines.keys():
                antnum1 = input_uv.ant_1_array[blt_i]
                antnum2 = input_uv.ant_2_array[blt_i]
                index1 = np.where(input_uv.antenna_numbers == antnum1)[0][0]
                index2 = np.where(input_uv.antenna_numbers == antnum2)[0][0]
                baselines[bl_num] = Baseline(antennas[index1], antennas[index2])

            time = time_array[blt_i]
            bl = baselines[bl_num]
            freq = freq_array[freq_i]  # 0 = spw axis

            task = UVTask(sky, time, freq, bl, telescope, freq_i)
            task.uvdata_index = (blt_i, freq_i)    # 0 = spectral window index

            yield task
        del sky

def _check_ntasks_valid(Ntasks_tot):
    """
    Check that the size of the task array won't overflow the gather.

    Parameters
    ----------
    Ntasks_tot : int
        Number of total tasks.

    """
    # Found experimentally, this is the limit on the number of tasks
    # that can be gathered with MPI.
    MAX_NTASKS_GATHER = 10226018

    if Ntasks_tot >= MAX_NTASKS_GATHER:
        raise ValueError(
            f"Too many tasks for MPI to gather successfully: {Ntasks_tot}\n"
            "\t Consider splitting your simulation into multiple smaller jobs "
            "and joining them together with UVData.\n"
            "\t This error will go away when issue #289 is closed."
        )

def run_uvdata_uvsim(
    input_uv,
    beam_list,
    beam_dict=None,
    catalog=None,
    quiet=False,
    block_nonroot_stdout=True
):
    """
    Run uvsim from UVData object.

    Parameters
    ----------
    input_uv : :class:`pyuvdata.UVData` instance
        Provides baseline/time/frequency information.
    beam_list : :class:`pyuvsim.BeamList`
        BeamList carrying beam models.
    beam_dict : dictionary, optional
        {`antenna_name` : `beam_id`}, where `beam_id` is an index in the beam_list.
        This is used to assign beams to antennas. Default: All antennas get the 0th
        beam in the `beam_list`.
    catalog : :class:`pyuvsim.simsetup.SkyModelData`
        Immutable source parameters.
    quiet : bool
        Do not print anything.
    block_nonroot_stdout : bool
        Redirect stdout on nonzero ranks to /dev/null, for cleaner output.

    Returns
    -------
    :class:`pyuvdata.UVData`
        instance containing simulated visibilities.

    """
    if mpi is None:
        raise ImportError("You need mpi4py to use the uvsim module. "
                          "Install it by running pip install pyuvsim[sim] "
                          "or pip install pyuvsim[all] if you also want the "
                          "line_profiler installed.")

    mpi.start_mpi(block_nonroot_stdout=block_nonroot_stdout)
    rank = mpi.get_rank()
    comm = mpi.get_comm()
    Npus = mpi.get_Npus()

    if not isinstance(input_uv, UVData):
        raise TypeError("input_uv must be UVData object")

    if not ((input_uv.Npols == 4) and (input_uv.polarization_array.tolist() == [-5, -6, -7, -8])):
        raise ValueError("input_uv must have XX,YY,XY,YX polarization")

    if not input_uv.future_array_shapes:
        input_uv.use_future_array_shapes()

    input_order = input_uv.blt_order
    if input_order != ("time", "baseline"):
        input_uv.reorder_blts(order="time", minor_order="baseline")

    # The root node will initialize our simulation
    # Read input file and make uvtask list
    if rank == 0 and not quiet:
        print('Nbls:', input_uv.Nbls, flush=True)
        print('Ntimes:', input_uv.Ntimes, flush=True)
        print('Nfreqs:', input_uv.Nfreqs, flush=True)
        print('Nsrcs:', catalog.Ncomponents, flush=True)
    if rank == 0:
        uv_container = simsetup._complete_uvdata(input_uv, inplace=False)
        if 'world' in input_uv.extra_keywords:
            uv_container.extra_keywords['world'] = input_uv.extra_keywords['world']
        vis_data = mpi.MPI.Win.Create(uv_container._data_array.value, comm=mpi.world_comm)
    else:
        vis_data = mpi.MPI.Win.Create(None, comm=mpi.world_comm)

    Nbls = input_uv.Nbls
    Nblts = input_uv.Nblts
    Nfreqs = input_uv.Nfreqs
    Nsrcs = catalog.Ncomponents

    task_inds, src_inds, Ntasks_local, Nsrcs_local = _make_task_inds(
        Nblts, Nfreqs, Nsrcs, rank, Npus
    )

    # Construct beam objects from strings
    beam_list.set_obj_mode(use_shared_mem=True)

    # In case the user created the beam list without checking consistency:
    beam_list.check_consistency()

    if beam_list.beam_type != 'efield':
        raise ValueError("Beam type must be efield!")

    # Estimating required memory to decide how to split source array.
    mem_avail = (simutils.get_avail_memory()
                 - mpi.get_max_node_rss(return_per_node=True) * 2**30)

    Npus_node = mpi.node_comm.Get_size()
    skymodel_mem_footprint = (
        simutils.estimate_skymodel_memory_usage(Nsrcs, catalog.Nfreqs) * Npus_node
    )

    # Allow up to 50% of available memory for SkyModel data.
    skymodel_mem_max = 0.5 * mem_avail

    Nsky_parts = np.ceil(skymodel_mem_footprint / float(skymodel_mem_max))
    Nsky_parts = max(Nsky_parts, 1)
    if Nsky_parts > Nsrcs:
        raise ValueError("Insufficient memory for simulation.")

    local_task_iter = uvdata_to_task_iter(
        task_inds, input_uv, catalog.subselect(src_inds),
        beam_list, beam_dict, Nsky_parts=Nsky_parts
    )

    Ntasks_tot = Ntasks_local * Nsky_parts
    # Sum all the tasks across each node
    Nsky_parts = comm.reduce(Nsky_parts, op=mpi.MPI.MAX, root=0)
    Ntasks_tot = comm.reduce(Ntasks_tot, op=mpi.MPI.SUM, root=0)
    if rank == 0 and not quiet:
        if Nsky_parts > 1:
            print(
                "The source list has been split into Nsky_parts"
                f" <= {Nsky_parts} chunks on some or all nodes"
                " due to memory limitations.",
                flush=True,
            )
        print("Tasks: ", Ntasks_tot, flush=True)
        pbar = simutils.progsteps(maxval=Ntasks_tot)

    engine = UVEngine()
    count = mpi.Counter()
    size_complex = np.ones(1, dtype=complex).nbytes
    data_array_shape = (Nblts, Nfreqs, 4)
    uvdata_indices = []

    for task in local_task_iter:
        engine.set_task(task)
        vis = engine.make_visibility()

        blti, freq_ind = task.uvdata_index

        uvdata_indices.append(task.uvdata_index)

        flat_ind = np.ravel_multi_index(
            (blti, freq_ind, 0), data_array_shape
        )
        offset = flat_ind * size_complex

        vis_data.Lock(0)
        vis_data.Accumulate(vis, 0, target=offset, op=mpi.MPI.SUM)
        vis_data.Unlock(0)

        cval = count.next()
        if rank == 0 and not quiet:
            pbar.update(cval)

    request = comm.Ibarrier()

    while not request.Test():
        if rank == 0 and not quiet:
            cval = count.current_value()
            pbar.update(cval)

    count.free()
    if rank == 0 and not quiet:
        pbar.finish()

    if rank == 0 and not quiet:
        print("Calculations Complete.", flush=True)

    # If profiling is active, save meta data:
    from .profiling import prof     # noqa
    if hasattr(prof, 'meta_file'):  # pragma: nocover
        # Saving axis sizes on current rank (local) and for the whole job (global).
        # These lines are affected by issue 179 of line_profiler, so the nocover
        # above will need to stay until this issue is resolved (see profiling.py).
        task_inds = np.array(uvdata_indices)
        bl_inds = task_inds[:, 0] % Nbls
        time_inds = (task_inds[:, 0] - bl_inds) // Nbls
        Ntimes_loc = np.unique(time_inds).size
        Nbls_loc = np.unique(bl_inds).size
        Nfreqs_loc = np.unique(task_inds[:, 1]).size
        axes_dict = {
            'Ntimes_loc': Ntimes_loc,
            'Nbls_loc': Nbls_loc,
            'Nfreqs_loc': Nfreqs_loc,
            'Nsrcs_loc': Nsky_parts,
            'prof_rank': prof.rank
        }

        with open(prof.meta_file, 'w') as afile:
            for k, v in axes_dict.items():
                afile.write("{} \t {:d}\n".format(k, int(v)))

    vis_data.Free()

    if rank == 0:
        if input_order is not None:
            if len(input_order) < 2:
                input_order = (input_order[0], None)
            # Don't call the function if we're already expecting (time, baseline)
            if input_order != ("time", "baseline"):
                uv_container.reorder_blts(
                    order=input_order[0],
                    minor_order=input_order[1],
                )
        else:
            warnings.warn(
                "The parameter `blt_order` could not be identified for input_uv "
                " so the ordering cannot be restored."
                "The output UVData object will have (time, baseline) ordering."
            )

        # Updating file history.
        history = simutils.get_version_string()
        if catalog.filename is not None:
            history += ' Sources from source list(s): [' + ', '.join(catalog.filename) + '].'
        if 'obsparam' in input_uv.extra_keywords:
            obs_param_file = input_uv.extra_keywords['obsparam']
            telescope_config_file = input_uv.extra_keywords['telecfg']
            antenna_location_file = input_uv.extra_keywords['layout']
            history += (' Based on config files: ' + obs_param_file + ', '
                        + telescope_config_file + ', ' + antenna_location_file)
        elif hasattr(uv_container, "filename") and uv_container.filename is not None:
            history += ' Based on uvdata file(s): [' + ', '.join(uv_container.filename) + '].'
        history += ' Npus = ' + str(mpi.Npus) + '.'

        # add pyradiosky version
        history += catalog.pyradiosky_version_str

        # add pyuvdata version info
        history += uv_container.pyuvdata_version_str

        uv_container.history = history

        return uv_container

def run_uvsim(params, return_uv=False, quiet=False, block_nonroot_stdout=True):
    """
    Run a simulation off of an obsparam yaml file.

    Parameters
    ----------
    params : str
        Path to a parameter yaml file.
    return_uv : bool
        If true, do not write results to file and return uv_out. (Default False)
    quiet : bool
        If True, do not print anything to stdout. (Default False)
    block_nonroot_stdout : bool
        Redirect stdout on nonzero ranks to /dev/null, for cleaner output.

    Returns
    -------
    uv_out : `pyuvdata.UVData`
        Finished simulation results.
        Returned only if return_uv is True.

    """
    if mpi is None:
        raise ImportError("You need mpi4py to use the uvsim module. "
                          "Install it by running pip install pyuvsim[sim] "
                          "or pip install pyuvsim[all] if you also want the "
                          "line_profiler installed.")

    mpi.start_mpi(block_nonroot_stdout=block_nonroot_stdout)
    rank = mpi.get_rank()
    comm = mpi.get_comm()

    input_uv = UVData()
    beam_list = None
    beam_dict = None
    skydata = SkyModelData()

    if rank == 0:
        start = Time.now()
        input_uv, beam_list, beam_dict = simsetup.initialize_uvdata_from_params(
            params, return_beams=True
        )
        skydata = simsetup.initialize_catalog_from_params(
            params, input_uv, return_recarray=False, return_catname=False
        )
        print(f"UVData initialization took {(Time.now() - start).to('minute'):.3f}")
        start = Time.now()
        skydata = simsetup.SkyModelData(skydata)
        print(f"Skymodel setup took {(Time.now() - start).to('minute'):.3f}")

    input_uv = comm.bcast(input_uv, root=0)
    beam_list = comm.bcast(beam_list, root=0)
    beam_dict = comm.bcast(beam_dict, root=0)
    skydata.share(root=0)

    start = Time.now()
    uv_out = run_uvdata_uvsim(
        input_uv, beam_list, beam_dict=beam_dict, catalog=skydata, quiet=quiet
    )
    if rank == 0:
        print(f"Run uvdata uvsim took {(Time.now() - start).to('minute'):.3f}")

    if rank == 0:
        if isinstance(params, str):
            with open(params, 'r') as pfile:
                param_dict = yaml.safe_load(pfile)
        else:
            param_dict = params

        simutils.write_uvdata(uv_out, param_dict, dryrun=return_uv)

    if return_uv:
        return uv_out

import numpy,import astropy,from astropy


#  Author(s):  Anton Deguet
#  Created on: 2016-05

#   (C) Copyright 2016-2022 Johns Hopkins University (JHU), All Rights Reserved.

# --- begin cisst license - do not edit ---

# This software is provided "as is" under an open source license, with
# no warranty.  The complete license can be found in license.txt and
# http://www.cisst.org/cisst/license.txt.

# --- end cisst license ---

from dvrk.arm import *

import numpy

class psm(arm):
    """Simple robot API wrapping around ROS messages
    """

    # class to contain jaw methods
    class __Jaw:
        def __init__(self, ros_node, ros_sub_namespace, expected_interval, operating_state_instance):
            self.__crtk_utils = crtk.utils(self, ros_node, expected_interval, operating_state_instance)
            self.__crtk_utils.add_measured_js(ros_sub_namespace)
            self.__crtk_utils.add_setpoint_js(ros_sub_namespace)
            self.__crtk_utils.add_servo_jp(ros_sub_namespace)
            self.__crtk_utils.add_move_jp(ros_sub_namespace)
            self.__crtk_utils.add_servo_jf(ros_sub_namespace)

        def close(self):
            "Close the tool jaw"
            return self.move_jp(numpy.array([math.radians(-20.0)]))

        def open(self, angle = math.radians(60.0)):
            "Close the tool jaw"
            return self.move_jp(numpy.array([angle]))

    # initialize the robot
    def __init__(self, arm_name, ros_node = None,
                 expected_interval = 0.01):
        # first call base class constructor
        self._arm__init_arm(arm_name, ros_node, expected_interval)
        self.jaw = self.__Jaw(self._arm__ros_node, 'jaw/', expected_interval,
                              operating_state_instance = self)

        # publishers
        self.__set_tool_present_publisher = self._arm__ros_node.create_publisher(std_msgs.msg.Bool,
                                                                                'jaw/set_tool_present',
                                                                                10) # latch = True
        self._arm__pub_list.extend([self.__set_tool_present_publisher])

    def insert_jp(self, depth):
        "insert the tool, by moving it to an absolute depth"
        goal = numpy.copy(self.setpoint_jp())
        goal[2] = depth
        return self.move_jp(goal)

    def set_tool_present(self, tool_present):
        "Set tool inserted.  To be used only for custom tools that can't be detected automatically"
        tp = std_msgs.msg.Bool()
        tp.data = tool_present
        self.__set_tool_present_publisher.publish(tp)


""" Models Module.

    - Builds the models that are then fed into a training loop.
    - This will save the output graphs after every epoch, so they can be used in a different environment.
"""
import os
import tensorflow as tf
import numpy as np
import argparse
import math, time
import scipy
from utility import BatchGenerator, tileImage
from operations import *
from architecture import discriminator, artist_generator, animator_generator

class GAN(object):
    """ Base class of GAN.

        Sets attributes that are shared across both GAN models.

        Args:
            - sess: Tensorflow session to attach to.
            - isTraining: toggles the updating of models, when feeding in examples.
            - imageSize: dimensions of the images used in training.
            - args: extra arguments fed in through the training script.
    """
    def __init__(self, sess, isTraining, imageSize, labelSize, args):
        self.bs = args.batch_size
        self.learning_rate = args.learning_rate
        self.zdim = args.zdim
        self.isTraining = isTraining
        self.imageSize = imageSize
        self.save_folder = args.save_folder
        self.reload = args.reload
        self.epoch = args.epoch
        self.cdim = args.cdim
        self.labelSize = labelSize
        self.sess = sess
        self.gf_dim = args.gf_dim
        self.df_dim = args.df_dim

    def loadModel(self, model_path=None):
        """ Restores an existing checkpoint to use in training. """
        if model_path: self.saver.restore(self.sess, model_path)

class Animator(GAN):
    """ Animator model setup which learns to generate animations, creates models and variables and then runs the training cycle. """
    def __init__(self, sess, isTraining, imageSize, labelSize, args):
        GAN.__init__(self, sess, isTraining, imageSize, labelSize, args)
        self.buildModel()
        return

    def buildModel(self):
        """ Build models networks, and set up loss and variables. """
        self.batch_size = tf.placeholder(tf.int32, [None, 1], name="batch_size") # Enable dynamic batch size.
        self.l = tf.placeholder(tf.float32, [self.batch_size.get_shape()[0], self.labelSize], name="label")

        img_dimensions = [self.imageSize[0], self.imageSize[1], self.cdim]
        self.z = tf.placeholder(tf.float32, [self.batch_size.get_shape()[0]] + img_dimensions, name="base")
        self.g_real = tf.placeholder(tf.float32, [self.batch_size.get_shape()[0]] + img_dimensions, name="images")

        ### GENERATORS ###
        self.g_fake = animator_generator(self.z, self.l, img_dimensions, self.gf_dim, self.cdim, self.batch_size, self.labelSize)
        self.g_sample = animator_generator(self.z, self.l, img_dimensions, self.gf_dim, self.cdim, self.batch_size, self.labelSize, reuse=True, isTraining=False)

        ### DISCRIMINATORS ###
        self.d_real = discriminator(self.z, self.l, self.df_dim, self.cdim, self.batch_size, self.labelSize, isTraining=self.isTraining)
        self.d_fake = discriminator(self.z, self.l, self.df_dim, self.cdim, self.batch_size, self.labelSize, reuse=True, isTraining=self.isTraining)

        print("BUILT MODELS")

        # Define loss
        self.d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.d_real, labels=tf.ones_like (self.d_real)))
        self.d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.d_fake, labels=tf.zeros_like(self.d_fake)))
        self.g_loss      = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.d_fake, labels=tf.ones_like (self.d_fake))) \
                            + 100 * tf.reduce_mean(tf.abs(self.g_real - self.g_fake))
        self.d_loss      = self.d_loss_real + self.d_loss_fake

        print("DEFINED LOSS FUNCTIONS")

        self.g_optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5).minimize(self.g_loss, var_list=[x for x in tf.trainable_variables() if "Generator"     in x.name])
        self.d_optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5).minimize(self.d_loss, var_list=[x for x in tf.trainable_variables() if "Discriminator" in x.name])

        print("DEFINED OPTIMIZERS")

        #############################
        ### saver
        self.saver = tf.train.Saver()
        self.summary = tf.summary.merge_all()
        if self.save_folder: self.writer = tf.summary.FileWriter(self.save_folder, self.sess.graph)

    def train(self, batch_generator):
        """ Runs training loop of model.

            Args:
                - batch_generator: object to use to retrieve batches of images from.
        """"
        if self.save_folder and not os.path.exists(os.path.join(self.save_folder,"images")):
            os.makedirs(os.path.join(self.save_folder,"images"))

        init = tf.global_variables_initializer()
        self.sess.run(init)
        self.loadModel(self.reload)
        start = time.time()
        self.batch_s = np.zeros((self.bs, 1))

        for epoch in range(self.epoch):

            batch_steps = batch_generator.get_file_count() // self.bs

            for step in range(batch_steps):

                batch_z = np.random.uniform(-1., +1., [self.bs, self.zdim])
                batch_images, batch_labels, batch_bases = batch_generator.get_batch(self.bs)

                if step % 5 == 1:
                    feed_dict = {self.z : batch_bases, self.l : batch_labels, self.g_real : batch_images, self.batch_size: self.batch_s}
                    _, d_loss, g_real, summary = self.sess.run([self.d_optimizer, self.d_loss, self.g_real, self.summary], feed_dict = feed_dict)
                else:
                    # Update generators twice.
                    _, g_loss                = self.sess.run([self.g_optimizer, self.g_loss], feed_dict={self.z: batch_bases, self.l: batch_labels, self.g_real: batch_images, self.batch_size: self.batch_s})
                    _, g_loss                = self.sess.run([self.g_optimizer, self.g_loss], feed_dict={self.z: batch_bases, self.l: batch_labels, self.g_real: batch_images, self.batch_size: self.batch_s})
                    feed_dict = {self.z : batch_bases, self.l : batch_labels, self.g_real : batch_images, self.batch_size: self.batch_s}
                    _, d_loss, g_fake, g_real, summary = self.sess.run([self.d_optimizer, self.d_loss, self.g_fake, self.g_real, self.summary], feed_dict = feed_dict)

                if step % 10 == 0:
                    print ("Epoch {}: [{}/{}] loss(D)={:.4f}, loss(G)={:.4f}; time/step = {:.2f} sec".format(epoch, step, batch_steps, d_loss, g_loss, time.time() - start))
                    start = time.time()

                if step % 100 == 0:
                    # Run models outputting images as training is run.
                    self.writer.add_summary(summary, step)
                    scipy.misc.imsave(os.path.join(self.save_folder,"images","img_{}_{}_bases.png".format(epoch, step)), tileImage(batch_bases))
                    self.generate_sample(g_real, batch_z, batch_labels, epoch, step, batch_bases)

            batch_generator.reset_buffer()

            freeze_graph('Generator_1/sprite', 'Animator', self.save_folder)

    def generate_sample(self, real_image, batch_z, batch_labels, epoch, step, bases):
        """ Generate sample images during training of the networks.

            One image is matched to the real_image that is fed into this function.
            In order to show how close the output is to the target output.

            Args:
                - real_image: example of a real image from the batch.
                - batch_z: noise vector used to generate a match to the real image.
                - batch_labels: labels that match the real images labels.
                - epoch: current epoch number.all
                - step: current step of epoch.
                - bases: base frames of the animations to reproduce.
        """
        l0 = np.random.uniform(-1, +1, [self.bs, self.labelSize])
        l1 = np.array([np.random.binomial(1, 0.5, self.labelSize) for x in range(self.bs)])

        binomial_image = self.sess.run(self.g_sample, feed_dict={self.z:bases, self.l:l1, self.batch_size: self.batch_s})
        noise_image = self.sess.run(self.g_sample, feed_dict={self.z:bases, self.l:l0, self.batch_size: self.batch_s})
        matched_image = self.sess.run(self.g_sample, feed_dict={self.z:bases, self.l:batch_labels, self.batch_size: self.batch_s})

        scipy.misc.imsave(os.path.join(self.save_folder,"images","anim_img_{}_{}_real.png".format(epoch, step)), tileImage(real_image))
        scipy.misc.imsave(os.path.join(self.save_folder,"images","anim_img_{}_{}_matched.png".format(epoch, step)), tileImage(matched_image))
        scipy.misc.imsave(os.path.join(self.save_folder,"images","anim_img_{}_{}_noise.png".format(epoch, step)), tileImage(noise_image))
        scipy.misc.imsave(os.path.join(self.save_folder,"images","anim_img_{}_{}_binomial.png".format(epoch, step)), tileImage(binomial_image))

        self.saver.save(self.sess, os.path.join(self.save_folder, "model.ckpt"), step)

class Artist(GAN):
    """ Model for artist network, which learns how to draw sprites.
        Creates models and variables and then runs the training cycle.
    """
    def __init__(self, sess, isTraining, imageSize, labelSize, args):
        GAN.__init__(self, sess, isTraining, imageSize, labelSize, args)
        self.buildModel()
        return

    def buildModel(self):
        """ Build models networks, and set up loss and variables. """
        # define variables
        self.batch_size = tf.placeholder(tf.int32, [None, 1], name="batch_size")
        self.z = tf.placeholder(tf.float32, [self.batch_size.get_shape()[0], self.zdim], name="z")
        self.l = tf.placeholder(tf.float32, [self.batch_size.get_shape()[0], self.labelSize], name="label")

        img_dimensions = [self.imageSize[0], self.imageSize[1], self.cdim]
        self.g_real = tf.placeholder(tf.float32, [self.batch_size.get_shape()[0]] + img_dimensions, name="images")

        ### GENERATORS ###
        self.g_fake = artist_generator(self.z, self.l, img_dimensions, self.gf_dim, self.cdim, self.batch_size, self.labelSize)
        self.g_sample = artist_generator(self.z, self.l, img_dimensions, self.gf_dim, self.cdim, self.batch_size, self.labelSize, reuse=True, isTraining=False)

        ### DISCRIMINATORS ###
        self.d_real = discriminator(self.g_real, self.l, self.df_dim, self.cdim, self.batch_size, self.labelSize, isTraining=self.isTraining)
        self.d_fake = discriminator(self.g_fake, self.l, self.df_dim, self.cdim, self.batch_size, self.labelSize, reuse=True, isTraining=self.isTraining)

        print("BUILT MODELS")

        # define loss
        self.d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.d_real, labels=tf.ones_like (self.d_real)))
        self.d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.d_fake, labels=tf.zeros_like(self.d_fake)))
        self.g_loss      = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.d_fake, labels=tf.ones_like (self.d_fake)))
        self.d_loss      = self.d_loss_real + self.d_loss_fake

        print("DEFINED LOSS FUNCTIONS")

        # define optimizer
        self.g_optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5).minimize(self.g_loss, var_list=[x for x in tf.trainable_variables() if "Generator"     in x.name])
        self.d_optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5).minimize(self.d_loss, var_list=[x for x in tf.trainable_variables() if "Discriminator" in x.name])

        print("DEFINED OPTIMIZERS")

        tf.summary.scalar("d_loss_real"   ,self.d_loss_real)
        tf.summary.scalar("d_loss_fake"   ,self.d_loss_fake)
        tf.summary.scalar("d_loss"      ,self.d_loss)
        tf.summary.scalar("g_loss"      ,self.g_loss)

        #############################
        ### saver
        self.saver = tf.train.Saver()
        self.summary = tf.summary.merge_all()
        if self.save_folder: self.writer = tf.summary.FileWriter(self.save_folder, self.sess.graph)

        return

    def train(self, batch_generator):
        """ Runs training loop of model.

            Args:
                - batch_generator: object to use to retrieve batches of images from.
        """"
        if self.save_folder and not os.path.exists(os.path.join(self.save_folder,"images")):
            os.makedirs(os.path.join(self.save_folder,"images"))

        init = tf.global_variables_initializer()
        self.sess.run(init)
        self.loadModel(self.reload)
        start = time.time()
        self.batch_s = np.zeros((self.bs, 1))

        for epoch in range(self.epoch):

            batch_steps = batch_generator.get_file_count() // self.bs

            for step in range(batch_steps):

                batch_z = np.random.uniform(-1., +1., [self.bs, self.zdim])
                batch_images, batch_labels = batch_generator.get_batch(self.bs)

                # Add some random noise to the labels every 5 steps, to train GAN to generalize.
                if step % 5 == 0:
                    batch_labels = batch_labels * np.random.uniform(0, 1, [self.bs, self.labelSize])

                feed_dict = {self.z : batch_z, self.l : batch_labels, self.g_real : batch_images, self.batch_size : self.batch_s}

                # Every now and again train discriminator model more.
                if step % 5 == 1:
                    _, d_loss, g_real, summary = self.sess.run([self.d_optimizer, self.d_loss, self.g_real, self.summary], feed_dict = feed_dict)
                else:
                    # Update generator
                    _, g_loss              = self.sess.run([self.g_optimizer, self.g_loss],feed_dict={self.z: batch_z, self.l: batch_labels, self.batch_size : self.batch_s})
                    _, g_loss                = self.sess.run([self.g_optimizer, self.g_loss],feed_dict={self.z: batch_z, self.l: batch_labels, self.batch_size : self.batch_s})
                    _, d_loss, g_fake, g_real, summary = self.sess.run([self.d_optimizer, self.d_loss, self.g_fake, self.g_real, self.summary], feed_dict = feed_dict)

                if step % 10 == 0:
                    print ("Epoch {}: [{}/{}] loss(D)={:.4f}, loss(G)={:.4f}; time/step = {:.2f} sec".format(epoch, step, batch_steps, d_loss, g_loss, time.time() - start))
                    start = time.time()

                if step % 100 == 0:
                    # Run models outputting images as training is run.
                    self.writer.add_summary(summary, step)
                    self.generate_sample(g_real, batch_z, batch_labels, epoch, step)

            freeze_graph('Generator_1/sprite', 'Artist', self.save_folder)

            batch_generator.reset_buffer()

    def generate_sample(self, real_image, batch_z, batch_labels, epoch, step):
        """ Generate sample images during training of the networks.

            One image is matched to the real_image that is fed into this function.
            In order to show how close the output is to the target output.

            Args:
                - real_image: example of a real image from the batch.
                - batch_z: noise vector used to generate a match to the real image.
                - batch_labels: labels that match the real images labels.
                - epoch: current epoch number.all
                - step: current step of epoch.
        """
        l0 = np.random.uniform(-1, +1, [self.bs, self.labelSize])
        l1 = np.array([np.random.binomial(1, 0.5, self.labelSize) for x in range(self.bs)])
        z1 = np.random.uniform(-1, +1, [self.bs, self.zdim])

        binomial_image = self.sess.run(self.g_sample, feed_dict={self.z:z1, self.l:l1, self.batch_size : self.batch_s})
        noise_image = self.sess.run(self.g_sample, feed_dict={self.z:z1, self.l:l0, self.batch_size : self.batch_s})
        matched_image = self.sess.run(self.g_sample, feed_dict={self.z:batch_z, self.l:batch_labels, self.batch_size : self.batch_s})

        scipy.misc.imsave(os.path.join(self.save_folder,"images","img_{}_{}_real.png".format(epoch, step)), tileImage(real_image))
        scipy.misc.imsave(os.path.join(self.save_folder,"images","img_{}_{}_matched.png".format(epoch, step)), tileImage(matched_image))
        scipy.misc.imsave(os.path.join(self.save_folder,"images","img_{}_{}_noise.png".format(epoch, step)), tileImage(noise_image))
        scipy.misc.imsave(os.path.join(self.save_folder,"images","img_{}_{}_binomial.png".format(epoch, step)), tileImage(binomial_image))

        self.saver.save(self.sess, os.path.join(self.save_folder, "model.ckpt"), step)


import numpy as np
import cv2
import pygame
from pygame.locals import *
import socket
import time
import struct
import io

class CollectTrainingData(object):

    def __init__(self):
        # Computer Client to send out RC Cart control signals
        # self.broadcastIP = '192.168.0.104'           # IP address to send to, 255 in one or more positions is a broadcast / wild-card
        self.broadcastIP = '192.168.0.100'           # IP address to send to, 255 in one or more positions is a broadcast / wild-card
        self.broadcastIP = '10.3.34.215'             #PI IP. Galvanize's education ip, 255 in one or more positions is a broadcast / wild-card
        self.broadcastIP = '192.168.1.12'            # PI IP. Gordon's house IP address to send to
        self.broadcastIP = '192.168.43.103'
        self.broadcastPort = 9038                    # What message number to send with (LEDB on an LCD)
        self.interval = 0.1                          # Time between keyboard updates in seconds, smaller responds faster but uses more processor time
        self.regularUpdate = True                    # If True we send a command at a regular interval, if False we only send commands when keys are pressed or released

        # Setup the connection for sending on
        self.sender = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)       # Create the socket
        self.sender.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)                        # Enable broadcasting (sending to many IPs based on wild-cards)
        self.sender.bind(('0.0.0.0', 0))                                                         # Set the IP and port number to use locally, IP 0.0.0.0 means all connections and port 0 means assign a number for us (do not care)
        print 'Connected to RC Car Control...'

        # Computer Server to receive PI Image Stream
        self.server_socket = socket.socket()
        self.server_socket.bind(('0.0.0.0', 8000))
        self.server_socket.listen(0)
        self.connection = self.server_socket.accept()[0].makefile('rb')
        self.send_inst = True
        print 'Received PI Image Stream...'

        # Create labels -> k[0] = Left, k[1] = Right, k[2] = Up, k[3] = Down, k[4] = k[]
        self.k = np.zeros((4, 4), 'float')
        for i in range(4):
            self.k[i, i] = 1
        self.temp_label = np.zeros((1, 4), 'float')
        # self.k = np.zeros((6, 6), 'float')
        # for i in range(6):
        #     self.k[i, i] = 1
        # self.temp_label = np.zeros((1, 6), 'float')

        # Initialized pygame module, control car from the pop up window and start collecting images
        pygame.init()
        screen = pygame.display.set_mode([300,300])
        pygame.display.set_caption("RemoteKeyController - Press [ESC] to quit")
        self.collect_image()

    def collect_image(self):

        saved_frame = 0
        total_frame = 0

        # collect images for training
        print 'Start collecting images...'
        e1 = cv2.getTickCount()
        # image_array = 360*240/2  -> Lowerver half of images..

        image_array = np.zeros((1, 38400))
        label_array = np.zeros((1, 4), 'float')

        # stream video frames one by one
        try:
            stream_bytes = ' '
            frame = 1
            while self.send_inst:
                stream_bytes += self.connection.read(1024)
                first = stream_bytes.find('\xff\xd8')
                last = stream_bytes.find('\xff\xd9')
                if first != -1 and last != -1:
                    jpg = stream_bytes[first:last + 2]
                    stream_bytes = stream_bytes[last + 2:]
                    image = cv2.imdecode(np.fromstring(jpg, dtype=np.uint8), cv2.CV_LOAD_IMAGE_GRAYSCALE)
                    # image = cv2.imdecode(np.fromstring(jpg, dtype=np.uint8),cv2.CV_LOAD_IMAGE_UNCHANGED)

                    # select lower half of the image
                    roi = image[120:240, :]

                    # save streamed images
                    cv2.imwrite('training_images/frame{:>05}.jpg'.format(frame), image)

                    #cv2.imshow('roi_image', roi)
                    # cv2.imshow('image', image)

                    # reshape the roi image into one row array
                    temp_array = roi.reshape(1, 38400).astype(np.float32)
                    # temp_array = np.zeros((1,38400))
                    frame += 1
                    total_frame += 1

                    # get input from human driver
                    for event in pygame.event.get():
                        if event.type == pygame.QUIT:
                            send_inst = False
                            command = 'x'
                            self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))
                            exit()
                            break

                        elif event.type == KEYDOWN:
                            key_input = pygame.key.get_pressed()
                            # complex orders
                            if (key_input[pygame.K_UP] and key_input[pygame.K_RIGHT]) or (key_input[pygame.K_w] and key_input[pygame.K_d]):
                                print("Forward Right")
                                image_array = np.vstack((image_array, temp_array))
                                label_array = np.vstack((label_array, self.k[1]))
                                saved_frame += 1
                                command = 'wd'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif (key_input[pygame.K_UP] and key_input[pygame.K_LEFT]) or (key_input[pygame.K_w] and key_input[pygame.K_a]):
                                print("Forward Left")
                                image_array = np.vstack((image_array, temp_array))
                                label_array = np.vstack((label_array, self.k[0]))
                                saved_frame += 1
                                command = 'wa'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif (key_input[pygame.K_DOWN] and key_input[pygame.K_RIGHT]) or (key_input[pygame.K_s] and key_input[pygame.K_d]):
                                print("Reverse Right")
                                command = 'sd'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif (key_input[pygame.K_DOWN] and key_input[pygame.K_LEFT]) or (key_input[pygame.K_s] and key_input[pygame.K_a]):
                                print("Reverse Left")
                                command = 'sa'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            # simple orders
                            elif key_input[pygame.K_UP] or key_input[pygame.K_w]:
                                print("Forward")
                                saved_frame += 1
                                image_array = np.vstack((image_array, temp_array))
                                label_array = np.vstack((label_array, self.k[2]))
                                command = 'w'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif key_input[pygame.K_DOWN] or key_input[pygame.K_s]:
                                print("Reverse")
                                saved_frame += 1
                                image_array = np.vstack((image_array, temp_array))
                                label_array = np.vstack((label_array, self.k[3]))
                                command = 's'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif key_input[pygame.K_RIGHT] or key_input[pygame.K_d]:
                                print("Right")
                                image_array = np.vstack((image_array, temp_array))
                                label_array = np.vstack((label_array, self.k[1]))
                                saved_frame += 1
                                command = 'r'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif key_input[pygame.K_LEFT] or key_input[pygame.K_a]:
                                print("Left")
                                image_array = np.vstack((image_array, temp_array))
                                label_array = np.vstack((label_array, self.k[0]))
                                saved_frame += 1
                                command = 'l'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

                            elif key_input[pygame.K_x] or key_input[pygame.K_q]:
                                print 'exit'
                                self.send_inst = False
                                command = 'x'
                                self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))
                                break

                        elif event.type == pygame.KEYUP:
                            # key_input = pygame.key.get_pressed()
                            # if (key_input[pygame.K_UP]) or  (key_input[pygame.K_w]):
                            #     print("FR or FL to Forward")
                            #     saved_frame += 1
                            #     image_array = np.vstack((image_array, temp_array))
                            #     label_array = np.vstack((label_array, self.k[2]))
                            #     command = 'w0'
                            #     self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))
                            # elif (key_input[pygame.K_DOWN]) or  (key_input[pygame.K_s]):
                            #     print("RR or RL to Reverse")
                            #     saved_frame += 1
                            #     image_array = np.vstack((image_array, temp_array))
                            #     label_array = np.vstack((label_array, self.k[3]))
                            #     command = 's0'
                            #     self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))
                            # if (not key_input[pygame.K_DOWN] and not key_input[pygame.K_UP]) and \
                            #     (not key_input[pygame.K_w] and not key_input[pygame.K_s]):
                            #     print 'no foward/backward key pressed'
                            command = '0'
                            self.sender.sendto(command, (self.broadcastIP, self.broadcastPort))

            # save training images and labels
            train = image_array[1:, :]
            train_labels = label_array[1:, :]

            # save training data as a numpy file
            np.savez('training_data_temp/lvtest06.npz', train=train, train_labels=train_labels)

            e2 = cv2.getTickCount()
            # calculate streaming duration
            time0 = (e2 - e1) / cv2.getTickFrequency()
            print 'Streaming duration:', time0

            print(train.shape)
            print(train_labels.shape)
            print 'Total frame:', total_frame
            print 'Saved frame:', saved_frame
            print 'Dropped frame', total_frame - saved_frame

        finally:
            self.connection.close()
            self.server_socket.close()

if __name__ == '__main__':
    CollectTrainingData()