import numpy as np
from typing import Union, Tuple

from memory.common import ReplayBufferAbstract, ExperienceSamples
from memory.prioritized_memory import PrioritizedExperienceBuffer
from memory.uniform_memory import UniformExperienceBuffer
from utils.wrappers import LazyFrames

class DemoReplayBuffer(ReplayBufferAbstract):
    """
    Replay buffer that keeps 2 separate buffers. One for expert data and another for data obtained by the agent.
    It appends to both and can sample from both.
    Adapted from the demo memory Preferred Networks baselines implementation:
    https://github.com/minerllabs/baselines/blob/master/general/chainerrl/baselines/dqfd.py
    """
    def __init__(
            self, prioritized: bool, capacity, demo_capacity, n_step, eps=0.01, alpha=0.6, beta0=0.4,
            betasteps=1000, gamma=0.99, ir_prob=0.2, n_policies=1, bonus_priority_demo=1.0,
            bonus_priority_agent=0.001, other_policy_factor=0.1):
        self.capacity = capacity
        self.demo_capacity = demo_capacity
        self.prioritized = prioritized
        memory_type = PrioritizedExperienceBuffer if prioritized else UniformExperienceBuffer

        self.agent_memory = memory_type(
            capacity=capacity, n_step=n_step, eps=eps, alpha=alpha, beta0=beta0, betasteps=betasteps,
            gamma=gamma, ir_prob=ir_prob, n_policies=n_policies, normalize_by_max=False)  # Handled in this cls)
        self.demo_memory = memory_type(
            capacity=demo_capacity, n_step=n_step, eps=eps, alpha=alpha, beta0=beta0, betasteps=betasteps,
            gamma=gamma, ir_prob=ir_prob, n_policies=n_policies, normalize_by_max=False)  # Handled in this cls)

        self.agent_samples = 0
        self.demo_samples = 0
        self.policy_idxs = None
        self.pretrain_phase = False
        self.current_policy = 0

        self.bonus_priority_demo = bonus_priority_demo
        self.bonus_priority_agent = bonus_priority_agent

    def set_pretrain_phase(self, pretraining: bool):
        # In this phase, expert data can also be added to the non-expert memory.
        self.pretrain_phase = pretraining

    def append(self, state: Union[LazyFrames, Tuple[LazyFrames, np.ndarray]] = None,
               action: Union[np.ndarray, int] = None, reward: int = 0, done: bool = False,
               next_state: Union[LazyFrames, Tuple[LazyFrames, np.ndarray]] = None,
               skip_step: int = 0, p_idx: int = 0, prob=None, expert: bool = False,
               **kwargs) -> None:
        memory = self.agent_memory
        if expert and (self.pretrain_phase and not self.demo_memory.is_full):
            memory = self.demo_memory
        memory.append(state=state, action=action, reward=reward, done=done, next_state=next_state,
                      skip_step=skip_step, p_idx=p_idx, prob=prob, expert=expert, **kwargs)

    def _sample_from_memory(self, nsample_agent, nsample_demo, p_idx):
        """
        Samples experiences from memory
        :param nsample_agent: Number of RL transitions to sample
        :param nsample_demo: Number of demonstration transitions to sample
        :param p_idx: Index of priorities.
        """
        if nsample_demo > 0:
            sampled_demo, is_weight_demo = self.demo_memory.sample(
                nsample_demo, p_idx=p_idx)
            sampled_demo = sampled_demo.samples
        else:
            sampled_demo, is_weight_demo = [], []

        if nsample_agent > 0:
            sampled_agent, is_weight_agent = self.agent_memory.sample(
                nsample_agent, p_idx=p_idx)
            sampled_agent = sampled_agent.samples
        else:
            sampled_agent, is_weight_agent = [], []
        sampled_demo.extend(sampled_agent)
        sampled_demo = ExperienceSamples(sampled_demo)
        is_weight_demo = np.append(is_weight_demo, is_weight_agent)
        return sampled_demo, is_weight_demo

    def sample(self, n, demo_fraction=None, p_idx=0) -> Tuple[ExperienceSamples, np.ndarray]:
        """
        Sample n experiences from memory.
        :param n: Number of experiences to sample
        :param demo_fraction: Fraction of experiences to come from demo memory. If set will override 50/50
        :param p_idx: Index of priorities.
        """
        if demo_fraction == 1.0 or len(self.agent_memory) == 0:  # Demo only
            samples, is_weight = self._sample_from_memory(
                nsample_agent=0, nsample_demo=n, p_idx=p_idx)
            self.agent_samples = 0
            self.demo_samples = n
            return samples, is_weight
        if demo_fraction is not None:
            nsample_demo = int(demo_fraction*n)
            nsample_agent = n - nsample_demo
        else:
            priority_sums = self.agent_memory.sumtree[p_idx]
            priority_sums_demo = self.demo_memory.sumtree[p_idx]
            psum_agent = priority_sums.total_sum  # / (1 if self.pre_train_phase else len(self.memory))
            psum_demo = priority_sums_demo.total_sum  # / (1 if self.pre_train_phase else len(self.memory_demo))
            psample_agent = psum_agent / (psum_agent + psum_demo)

            # nsample_agent = int(n * (1.0 - demo_fraction))
            nsample_agent = np.random.binomial(n, psample_agent)
            # Increase agent samples with forget percentage -> leads to less demo samples, i.e. forgetting
            # nsample_agent += int((n - nsample_agent) * self.forget_percent)
            # If we don't have enough RL transitions yet, force more demos
            nsample_agent = min(nsample_agent, len(self.agent_memory))
            nsample_demo = n - nsample_agent

        samples, is_weight = self._sample_from_memory(
            nsample_agent, nsample_demo, p_idx=p_idx)
        self.agent_samples = nsample_agent
        self.demo_samples = nsample_demo
        self.policy_idxs = np.array(samples.policies)
        assert len(is_weight) == n
        return samples, is_weight

    def update_priorities(self, errors, p_idx=0):
        """
        Update sumtree priorities
        """
        errors[:self.demo_samples] += self.bonus_priority_demo
        errors[-self.agent_samples:] += self.bonus_priority_agent
        errors = errors / errors.max()
        if self.agent_samples > 0:
            agent_errors = errors[-self.agent_samples:]
            self.agent_memory.update_priorities(agent_errors, p_idx=p_idx)
        if self.demo_samples > 0:
            demo_errors = errors[:self.demo_samples]
            self.demo_memory.update_priorities(demo_errors, p_idx=p_idx)

    def __len__(self):
        return len(self.demo_memory) + len(self.agent_memory)

    def stop_current_episode(self, env_id=0, p_idx=0):
        """
        Manual end to an episode
        """
        if self.pretrain_phase:
            memory = self.demo_memory if not self.demo_memory.is_full else self.agent_memory
        else:
            memory = self.agent_memory
        last_n_transitions = memory.last_n_transitions[env_id]
        memory.add_transition(last_n_transitions, done=True, p_idx=p_idx)

    def clear_transition_buffers(self):
        """
        Manually ensure all
        """
        for memory in [self.demo_memory, self.agent_memory]:
            for val in memory.last_n_transitions.values():
                val.clear()


"""
Kernel Call Details
===================

When calling sas computational kernels with polydispersity there are a
number of details that need to be sent to the caller.  This includes the
list of polydisperse parameters, the number of points in the polydispersity
weight distribution, and which parameter is the "theta" parameter for
polar coordinate integration.  The :class:`CallDetails` object maintains
this data.  Use :func:`make_details` to build a *details* object which
can be passed to one of the computational kernels.
"""

from __future__ import print_function

import numpy as np  # type: ignore
from numpy import cos, sin, radians

try:
    np.meshgrid([])
    meshgrid = np.meshgrid
except Exception:
    # CRUFT: np.meshgrid requires multiple vectors
    def meshgrid(*args):
        """See docs from a recent version of numpy"""
        if len(args) > 1:
            return np.meshgrid(*args)
        else:
            return [np.asarray(v) for v in args]

# pylint: disable=unused-import
try:
    from typing import List, Tuple, Sequence
    from .modelinfo import ModelInfo, ParameterTable
    from .kernel import Kernel
except ImportError:
    pass
# pylint: enable=unused-import

class CallDetails(object):
    """
    Manage the polydispersity information for the kernel call.

    Conceptually, a polydispersity calculation is an integral over a mesh
    in n-D space where n is the number of polydisperse parameters.  In order
    to keep the program responsive, and not crash the GPU, only a portion
    of the mesh is computed at a time.  Meshes with a large number of points
    will therefore require many calls to the polydispersity loop.  Restarting
    a nested loop in the middle requires that the indices of the individual
    mesh dimensions can be computed for the current loop location.  This
    is handled by the *pd_stride* vector, with n//stride giving the loop
    index and n%stride giving the position in the sub loops.

    One of the parameters may be the latitude.  When integrating in polar
    coordinates, the total circumference decreases as latitude varies from
    pi r^2 at the equator to 0 at the pole, and the weight associated
    with a range of latitude values needs to be scaled by this circumference.
    This scale factor needs to be updated each time the theta value
    changes.  *theta_par* indicates which of the values in the parameter
    vector is the latitude parameter, or -1 if there is no latitude
    parameter in the model.  In practice, the normalization term cancels
    if the latitude is not a polydisperse parameter.
    """
    parts = None  # type: List["CallDetails"]
    def __init__(self, model_info):
        # type: (ModelInfo) -> None
        parameters = model_info.parameters
        max_pd = parameters.max_pd

        # Structure of the call details buffer:
        #   pd_par[max_pd]     pd params in order of length
        #   pd_length[max_pd]  length of each pd param
        #   pd_offset[max_pd]  offset of pd values in parameter array
        #   pd_stride[max_pd]  index of pd value in loop = n//stride[k]
        #   num_eval           total length of pd loop
        #   num_weights        total length of the weight vector
        #   num_active         number of pd params
        #   theta_par          parameter number for theta parameter
        self.buffer = np.empty(4*max_pd + 4, 'i4')

        # generate views on different parts of the array
        self._pd_par = self.buffer[0 * max_pd:1 * max_pd]
        self._pd_length = self.buffer[1 * max_pd:2 * max_pd]
        self._pd_offset = self.buffer[2 * max_pd:3 * max_pd]
        self._pd_stride = self.buffer[3 * max_pd:4 * max_pd]

        # theta_par is fixed
        self.theta_par = parameters.theta_offset

        # offset and length are for all parameters, not just pd parameters
        # They are not sent to the kernel function, though they could be.
        # They are used by the composite models (sum and product) to
        # figure out offsets into the combined value list.
        self.offset = None  # type: np.ndarray
        self.length = None  # type: np.ndarray

        # keep hold of ifno show() so we can break a values vector
        # into the individual components
        self.info = model_info

    @property
    def pd_par(self):
        """List of polydisperse parameters"""
        return self._pd_par

    @property
    def pd_length(self):
        """Number of weights for each polydisperse parameter"""
        return self._pd_length

    @property
    def pd_offset(self):
        """Offsets for the individual weight vectors in the set of weights"""
        return self._pd_offset

    @property
    def pd_stride(self):
        """Stride in the pd mesh for each pd dimension"""
        return self._pd_stride

    @property
    def num_eval(self):
        """Total size of the pd mesh"""
        return self.buffer[-4]

    @num_eval.setter
    def num_eval(self, v):
        """Total size of the pd mesh"""
        self.buffer[-4] = v

    @property
    def num_weights(self):
        """Total length of all the weight vectors"""
        return self.buffer[-3]

    @num_weights.setter
    def num_weights(self, v):
        """Total length of all the weight vectors"""
        self.buffer[-3] = v

    @property
    def num_active(self):
        """Number of active polydispersity loops"""
        return self.buffer[-2]

    @num_active.setter
    def num_active(self, v):
        """Number of active polydispersity loops"""
        self.buffer[-2] = v

    @property
    def theta_par(self):
        """Location of the theta parameter in the parameter vector"""
        return self.buffer[-1]

    @theta_par.setter
    def theta_par(self, v):
        """Location of the theta parameter in the parameter vector"""
        self.buffer[-1] = v

    def show(self, values=None):
        """Print the polydispersity call details to the console"""
        print("===== %s details ===="%self.info.name)
        print("num_active:%d  num_eval:%d  num_weights:%d  theta=%d"
              % (self.num_active, self.num_eval, self.num_weights, self.theta_par))
        if self.pd_par.size:
            print("pd_par", self.pd_par)
            print("pd_length", self.pd_length)
            print("pd_offset", self.pd_offset)
            print("pd_stride", self.pd_stride)
        if values is not None:
            nvalues = self.info.parameters.nvalues
            print("scale, background", values[:2])
            print("val", values[2:nvalues])
            print("pd", values[nvalues:nvalues+self.num_weights])
            print("wt", values[nvalues+self.num_weights:nvalues+2*self.num_weights])
            print("offsets", self.offset)

def make_details(model_info, length, offset, num_weights):
    # type: (ModelInfo, np.ndarray, np.ndarray, int) -> CallDetails
    """
    Return a :class:`CallDetails` object for a polydisperse calculation
    of the model defined by *model_info*.  Polydispersity is defined by
    the *length* of the polydispersity distribution for each parameter
    and the *offset* of the distribution in the polydispersity array.
    Monodisperse parameters should use a polydispersity length of one
    with weight 1.0. *num_weights* is the total length of the polydispersity
    array.
    """
    #pars = model_info.parameters.call_parameters[2:model_info.parameters.npars+2]
    #print(", ".join(str(i)+"-"+p.id for i,p in enumerate(pars)))
    #print("len:",length)
    #print("off:",offset)

    # Check that we aren't using too many polydispersity loops
    num_active = np.sum(length > 1)
    max_pd = model_info.parameters.max_pd
    if num_active > max_pd:
        raise ValueError("Too many polydisperse parameters")

    # Decreasing list of polydpersity lengths
    # Note: the reversing view, x[::-1], does not require a copy
    idx = np.argsort(length)[::-1][:max_pd]
    pd_stride = np.cumprod(np.hstack((1, length[idx])))

    call_details = CallDetails(model_info)
    call_details.pd_par[:max_pd] = idx
    call_details.pd_length[:max_pd] = length[idx]
    call_details.pd_offset[:max_pd] = offset[idx]
    call_details.pd_stride[:max_pd] = pd_stride[:-1]
    call_details.num_eval = pd_stride[-1]
    call_details.num_weights = num_weights
    call_details.num_active = num_active
    call_details.length = length
    call_details.offset = offset
    #call_details.show()
    return call_details

ZEROS = tuple([0.]*31)
def make_kernel_args(kernel, mesh):
    # type: (Kernel, Tuple[List[np.ndarray], List[np.ndarray]]) -> Tuple[CallDetails, np.ndarray, bool]
    """
    Converts (value, dispersity, weight) for each parameter into kernel pars.

    Returns a CallDetails object indicating the polydispersity, a data object
    containing the different values, and the magnetic flag indicating whether
    any magnetic magnitudes are non-zero. Magnetic vectors (M0, phi, theta) are
    converted to rectangular coordinates (mx, my, mz).
    """
    npars = kernel.info.parameters.npars
    nvalues = kernel.info.parameters.nvalues
    scalars = [value for value, dispersity, weight in mesh]
    # skipping scale and background when building values and weights
    _, dispersity, weight = zip(*mesh[2:npars+2]) if npars else ((), (), ())
    #weight = correct_theta_weights(kernel.info.parameters, dispersity, weight)
    length = np.array([len(w) for w in weight])
    offset = np.cumsum(np.hstack((0, length)))
    call_details = make_details(kernel.info, length, offset[:-1], offset[-1])
    # Pad value array to a 32 value boundary
    data_len = nvalues + 2*sum(len(v) for v in dispersity)
    extra = (32 - data_len%32)%32
    data = np.hstack((scalars,) + dispersity + weight + ZEROS[:extra])
    data = data.astype(kernel.dtype)
    is_magnetic = convert_magnetism(kernel.info.parameters, data)
    #call_details.show()
    #print("data", data)
    return call_details, data, is_magnetic

def correct_theta_weights(parameters, dispersity, weights):
    # type: (ParameterTable, Sequence[np.ndarray], Sequence[np.ndarray]) -> Sequence[np.ndarray]
    """
    **Deprecated** Theta weights will be computed in the kernel wrapper if
    they are needed.

    If there is a theta parameter, update the weights of that parameter so that
    the cosine weighting required for polar integration is preserved.

    Avoid evaluation strictly at the pole, which would otherwise send the
    weight to zero.  This is probably not a problem in practice (if dispersity
    is +/- 90, then you probably should be using a 1-D model of the circular
    average).

    Note: scale and background parameters are not include in the tuples for
    dispersity and weights, so index is parameters.theta_offset, not
    parameters.theta_offset+2

    Returns updated weights vectors
    """
    # Apparently the parameters.theta_offset similarly skips scale and
    # and background, so the indexing works out, but they are still shipped
    # to the kernel, so we need to add two there.
    if parameters.theta_offset >= 0:
        index = parameters.theta_offset
        theta = dispersity[index]
        theta_weight = abs(cos(radians(theta)))
        weights = tuple(theta_weight*w if k == index else w
                        for k, w in enumerate(weights))
    return weights

def convert_magnetism(parameters, values):
    # type: (ParameterTable, Sequence[np.ndarray]) -> bool
    """
    Convert magnetism values from polar to rectangular coordinates.

    Returns True if any magnetism is present.
    """
    mag = values[parameters.nvalues-3*parameters.nmagnetic:parameters.nvalues]
    mag = mag.reshape(-1, 3)
    if np.any(mag[:, 0] != 0.0):
        M0 = mag[:, 0].copy()
        theta, phi = radians(mag[:, 1]), radians(mag[:, 2])
        mag[:, 0] = M0 * sin(theta) * cos(phi)  # mx
        mag[:, 1] = M0 * sin(theta) * sin(phi)  # my
        mag[:, 2] = M0 * cos(theta) # mz
        return True
    else:
        return False

def dispersion_mesh(model_info, mesh):
    # type: (ModelInfo, List[Tuple[float, np.ndarray, np.ndarray]]) -> Tuple[List[np.ndarray], List[np.ndarray]]
    """
    Create a mesh grid of dispersion parameters and weights.

    *mesh* is a list of (value, dispersity, weights), where the values
    are the individual parameter values, and (dispersity, weights) is
    the distribution of parameter values.

    Only the volume parameters should be included in this list.  Orientation
    parameters do not affect the calculation of effective radius or volume
    ratio.  This is convenient since it avoids the distinction between
    value and dispersity that is present in orientation parameters but not
    shape parameters.

    Returns [p1,p2,...],w where pj is a vector of values for parameter j
    and w is a vector containing the products for weights for each
    parameter set in the vector.
    """
    _, dispersity, weight = zip(*mesh)
    #weight = [w if len(w)>0 else [1.] for w in weight]
    weight = np.vstack([v.flatten() for v in meshgrid(*weight)])
    weight = np.prod(weight, axis=0)
    dispersity = [v.flatten() for v in meshgrid(*dispersity)]
    lengths = [par.length for par in model_info.parameters.kernel_parameters
               if par.type == 'volume']
    if any(n > 1 for n in lengths):
        pars = []
        offset = 0
        for n in lengths:
            pars.append(np.vstack(dispersity[offset:offset+n])
                        if n > 1 else dispersity[offset])
            offset += n
        dispersity = pars
    return dispersity, weight


import numpy as np
from Constants import *

# Simple tests for broad regions
def Europe(lat,lon):
    # stops just before moscow
    return (-17 < lon <  35 and lat >  35 ) and not Anatolia(lat,lon)

def Russia(lat,lon):
    # Central and Western Russia
    return (35 <= lon <= 89 and lat >= 35) and not Anatolia(lat,lon) and not Mesopotamia(lat,lon)

def Britain(lat,lon):
    return (lon <= 1 and lat >= 50 )

def Anatolia(lat,lon):
    # anatolia and the caucus
    return (25 <= lon <= 48 and 35  <= lat <= 41 )

def Egypt(lat,lon):
    return (28 <= lon <= 35 and 16  <= lat <= 31 )

def Mesopotamia(lat,lon):
    return (30 < lon <  61 and 16  < lat <= 41) or Anatolia(lat,lon) or Egypt(lat,lon)

def North_Africa(lat,lon):
    return (-14 <= lon <= 28 and 26  <= lat <= 36 )

def sub_Saharan_Africa(lat,lon):
    # -4 lat to match/overlap the start of Agr_Expand1 in Africa
    return (-17 <= lon <= 53 and -3  <  lat <= 16 )

def Africa(lat,lon):
    return North_Africa(lat,lon) or sub_Saharan_Africa(lat,lon) or Egypt(lat,lon)

def India(lat,lon):
    return (61 <= lon <= 89 and 0  <= lat <  35 )

def Asia(lat,lon):
    # China but drop tundra to the north; include Khmer and Indonesia, etc. to the south
    return (lon >  89 and -13  < lat <= 47)

def China(lat,lon):
    # drops Khmer etc to the south; drop tundra to the north
    return (lon >  89 and 17  < lat <= 47)

def LiaoRV(lat,lon):
    return (95 <= lon <= 180 and  40  <  lat <= 70)
# def MYRV(lat,lon):
#     return (95 <= lon <= 180 and  30  <  lat <= 40)
# def PearlRV(lat,lon):
#     return (95 <= lon <= 180 and  19  <  lat <= 30)
def MekongRV(lat,lon):
    return (95 <= lon <= 180 and -15  <  lat <= 19)
def BurmaRV(lat,lon):
    return (85 <= lon <=  95 and  20  <  lat <= 40)

from shapely.geometry import Point
from shapely.geometry.polygon import Polygon
myrv_shape = None
def MYRV(lat,lon):
    global myrv_shape
    if myrv_shape is None:
        # note diagnonal downward base (could be any shape w/ many vertices)
        myrv_shape = Polygon([(95,30),(150,25),(150,40),(95,40)])
    return myrv_shape.contains(Point(lon,lat))

pearlrv_shape = None
def PearlRV(lat,lon):
    global pearlrv_shape
    if pearlrv_shape is None:
        # note diagnonal upward base to abutt MYRV above
        pearlrv_shape = Polygon([(95,19),(150,19),(150,25),(95,30)])
    return pearlrv_shape.contains(Point(lon,lat))

geography_region_list = ['Europe','Anatolia','Egypt','Mesopotamia','North_Africa','sub_Saharan_Africa','Africa','Asia','China','Russia','Britain','India',
                         'LiaoRV','MYRV','PearlRV','MekongRV','BurmaRV']

class Geography:
    def __init__(self,name,n_regions):
        self.name = name
        self.n_regions = n_regions
        self.lat = np.zeros(n_regions)
        self.lon = np.zeros(n_regions)
        self.Elevation = np.zeros(n_regions)
        self.Biome = np.zeros(n_regions)
        self.Biome[:] = GEO_DESERT

        # region sizes and distances (in km)

        # Littoral regions boolean (needed?)
        # immediate von Neuman connections to *all regions* (regardless of biome)
        # when forming territories a subset might be chosen, such as only AG possible, etc.
        # LandConnection {region => [regions]}, including negative numbers to indicate ocean at some VN neighbor
        # SeaConnection {littoral_region => {[[lit_region dist], ...]}) sorted by increasing distance (encode as km)

        # The actual empire history of these regions, if known empire_index[year,region]
        #DEAD self.history = None # up to the creator to provide this, if appropriate
        self.actual_history = False
        self.min_polity_size = 0
        self.years = None
        self.empire_names = None # names per empire index with [0] ensured as '' (e.g., Hinterland, Unoccupied)

    def __repr__(self):
        return "<Geography %s (%d regions)>" % (self.name,self.n_regions)


import numpy as np
import math
import time

import pandas as pd
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

MAX_WAIT_TIME = 10  # s

def craw_mlss(df_mls):
    mlss = list(df_mls['_id'])
    n_mls = len(mlss)
    print('num urls: %d' % n_mls)

    chrome_options = webdriver.ChromeOptions()

    chrome_options.add_argument('--headless')
    driver_url = webdriver.Chrome('/var/qindom/chromedriver', chrome_options=chrome_options)  # for url
    driver = webdriver.Chrome('/var/qindom/chromedriver', chrome_options=chrome_options)  # for content

    '''signin first'''
    print('signin...')
    root_url = 'https://housesigma.com'

    driver_url.get(url=root_url + '/web/en')

    driver.get(url=root_url + '/web/en/login')
    WebDriverWait(driver, MAX_WAIT_TIME).until(
                EC.visibility_of_element_located((By.XPATH, '//*[@id="pane-email"]/form/div[1]/div/div/input'))
            )
    inputElement = driver.find_element(By.XPATH, '//*[@id="pane-email"]/form/div[1]/div/div/input')
    inputElement.send_keys('nuhubam@skymailapp.com')

    time.sleep(1)

    inputElement = driver.find_element(By.XPATH, '//*[@id="pane-email"]/form/div[2]/div/div/input')
    inputElement.send_keys('123456')

    time.sleep(1)

    login_button = driver.find_element(By.XPATH, '//*[@id="login"]/div[1]/div/div[3]/button')
    login_button.click()

    WebDriverWait(driver, MAX_WAIT_TIME).until(
                EC.visibility_of_element_located((By.XPATH, '//*[@id="index"]/div[1]/div[2]/div[1]/div[2]'))
            )

    time.sleep(1)

    # example: W4436597
    start = time.time()
    contents = []
    for ind in range(n_mls):
        mls = mlss[ind]
        '''step 1, get mls url'''
        inputElement_url_search = driver_url.find_element_by_id('search_input')

        inputElement_url_search.clear()
        inputElement_url_search.send_keys(mls)

        mls_href = ''
        try:
            element = WebDriverWait(driver_url, MAX_WAIT_TIME).until(
                EC.visibility_of_element_located(
                    (By.XPATH, '//*[@id="index"]/div[1]/div[2]/div[1]/div[2]/div[2]/div/p[2]'))
            )

            page_source = driver_url.page_source
            mls_ind = page_source.find(mls)
            from_ind = page_source.rfind('href="', 0, mls_ind) + len('href="')
            to_ind = mls_ind + len(mls)

            mls_href = page_source[from_ind + len('/web/en/house/'):to_ind]
        except:
            print('%s href not found. is the max waiting time too short?' % mls)

        '''step 2, get listing contents'''
        url = mls_href
        time_average = (time.time() - start) / (ind + 0.01)
        print('processing %d/%d: %s. (avg time: %0.1f s, est time: %d s)' %
              (ind, n_mls, mls, time_average, int(time_average * n_mls)))

        content_market, content_estimate_selling_price, content_estimate_renting_price = -1, -1, -1
        page_source = ''

        if (not isinstance(url, str)) or (len(url) < 6):
            print('%s not found in housesigma' % mls)
        else:
            full_url = root_url + '/web/en/house/' + url
            print(full_url)
            driver.get(url=full_url)

            try:
                WebDriverWait(driver, MAX_WAIT_TIME).until(
                    EC.visibility_of_element_located((
                        By.XPATH, '//*[@id="room_count_to_estimate"]/div[6]/div[1]/div[1]/div[2]/span'))
                )
                page_source = driver.page_source
            except:
                print('%s page not loaded.' % mls)

            try:
                # buyer/seller market
                anckor = 'margin-left: calc('
                from_ind = page_source.find(anckor)
                to_ind = page_source.find('%', from_ind)

                content_market = float(page_source[from_ind + len(anckor):to_ind])
            except:
                pass
                # print('%s market not found.' % mls)

            try:
                # estimate selling price
                from_ind = page_source.find('>SigmaEstimate<')
                lower_bound = page_source.find('>Rental')
                from_ind = page_source.find('>$', from_ind, lower_bound) + len('>$')
                to_ind = page_source.find('<', from_ind)

                content_estimate_selling_price = int(page_source[from_ind:to_ind].replace(',', ''))
            except:
                pass
                # print('%s predict selling price not found.' % mls)

            try:
                # estimate renting price
                from_ind = page_source.find('>Rental Estimate<')
                lower_bound = page_source.find('>School')
                from_ind = page_source.find('>$', from_ind, lower_bound) + len('>$')
                to_ind = page_source.find('<', from_ind)

                content_estimate_renting_price = int(page_source[from_ind:to_ind].replace(',', ''))
            except:
                pass
                # print('%s predict renting price not found.' % mls)

            contents.append([mls, url, page_source,
                             content_market, content_estimate_selling_price, content_estimate_renting_price])
            print('%s page source len: %d; market: %0.1f; pred selling: %d; pred renting: %d' %
                  (mls, len(page_source),
                   content_market, content_estimate_selling_price, content_estimate_renting_price))

    driver_url.quit()
    driver.quit()

    print('crawling finished.')
    df_contents = pd.DataFrame(contents, columns=['_id', 'url', 'page_source', 'market_bias', 'pred_sell', 'pred_rent'])
    return df_contents

if __name__ == '__main__':
    # house address
    # /media/qindom-cpu/wd1/kai/real_master_crawler/house_sigma/
    df_mls = pd.read_csv('mls_sample_to_20190430.csv',
                         index_col=[0])

    df_res = craw_mlss(df_mls.iloc[:5, :])
    print(df_res[['_id', 'url', 'market_bias', 'pred_sell', 'pred_rent']].to_string())


import cv2
import numpy as np
from pynput.mouse import Button, Controller
import wx

mouse = Controller()
app = wx.App(False)
sx,sy = wx.GetDisplaySize()
(camx, camy) = (400,300)

"""
#RED
lowerBound = np.array([0,80,40])
upperBound = np.array([25,255,255])
"""
"""
#YELLOW
lowerBound = np.array([7,80,80])
upperBound = np.array([43,255,255])
"""

#GREEN
lowerBound = np.array([33,80,40])
upperBound = np.array([102,255,255])

cam = cv2.VideoCapture(0)
cam.set(3,camx)
cam.set(4,camy)
kernalOpen = np.ones((5,5))
kernelClose=np.ones((20,20))

mLoc0ld = np.array([0,0])
mouseLoc = np.array([0,0])
DamplingFactor = 2   #should be between > 1

pinchFlag = 0
openx,openy,openw,openh = (0,0,0,0)

#mouseLoc = mLoc0ld + (targetLoc - mLoc0ld)/DamplingFactor

while True :
       ret, img = cam.read()
       #mg = cv2.resize(img,(340,240))

       #convert BGR to HSV
       imgHSV = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
       # creat the mask
       mask = cv2. inRange(imgHSV,lowerBound,upperBound)
       # morphology
       maskOpen = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernalOpen)
       maskClose = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernalOpen)

       maskFinal  = maskClose
       _,conts,h = cv2.findContours(maskFinal.copy(),cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

       if (len(conts) == 2) :
              if (pinchFlag == 1) :
                     pinchFlag = 0
                     mouse.release(Button.left)
              mouse.release(Button.left)
              x1,y1,w1,h1 = cv2.boundingRect(conts[0])
              x2,y2,w2,h2 = cv2.boundingRect(conts[1])
              cv2.rectangle(img,(x1,y1),(x1+w1,y1+h1),(255,0,0),2)
              cv2.rectangle(img,(x2,y2),(x2+w2,y2+h2),(255,0,0),2)
              cx1=x1+w1/2
              cy1=y1+h1/2
              cx2=x2+w2/2
              cy2=y2+h2/2
              cx = (cx1 + cx2)/2
              cy = (cy2 + cy2)/2
              cv2.line(img, (int(cx1),int(cy1)),(int(cx2),int(cy2)),(255,0,0),2)
              cv2.circle(img, (int(cx),int(cy)),2,(0,0,255),2)
              mouseLoc = mLoc0ld + ( (int(cx),int(cy)) - mLoc0ld)/DamplingFactor
              mouse.position =  (sx-int(mouseLoc[0]*sx/camx),int(mouseLoc[1]*sy/camy))
              while mouse.position != (sx-int(mouseLoc[0]*sx/camx),int(mouseLoc[1]*sy/camy)) :
                     pass
              mLoc0ld = mouseLoc
              openx,openy,openw,openh=cv2.boundingRect(np.array([[[x1,y1],[x1+w1,y1+h1],[x2,y2],[x2+w2, y2+h2]]]))
              #cv2.rectangle(img,(openx,openy),(openx+openw,openy+openh),(255,0,0),2)

       elif (len(conts)==1):
              x,y,w,h = cv2.boundingRect(conts[0])
              if (pinchFlag == 0) :
                     if (abs((w*h - openw*openh)*100/(w*h))<20) :
                            pinchFlag = 1
                            mouse.press(Button.left)
                            openx,openy,openw,openh = (0,0,0,0)
              else :
                     cv2.rectangle(img,(x,y),(x+w,y+h),(255,0,0),2)
                     cx = x + w/2
                     cy = y + h/2
                     cv2.circle(img, (int(cx),int(cy)),int((w+h)/4),(0,0,255),2)
                     mouseLoc = mLoc0ld + ( (int(cx),int(cy)) - mLoc0ld)/DamplingFactor
                     mouse.position =  (sx-int(mouseLoc[0]*sx/camx),int(mouseLoc[1]*sy/camy))
                     while mouse.position != (sx-int(mouseLoc[0]*sx/camx),int(mouseLoc[1]*sy/camy)) :
                                   pass
                     mLoc0ld = mouseLoc

       #cv2.imshow("maskClose", maskClose)
       #cv2.imshow("maskOpen", maskOpen)
       #cv2.imshow("mask", mask)
       cv2.imshow("cam", img)
       if cv2.waitKey(10) &0xFF ==ord('q'):
              cap.release()
              cv2.destroyAllWindows()
              break


#!/usr/bin/env python
# -*- coding: utf-8 -*-

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2002-2017 Prof. William H. Green (whgreen@mit.edu),
#   Prof. Richard H. West (r.west@neu.edu) and the RMG Team (rmg_dev@mit.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

import numpy
import os.path
import logging
import rmgpy.constants as constants
################################################################################

def checkConformerEnergy(Vlist,path):
    """
    Check to see that the starting energy of the species in the potential energy scan calculation
    is not 0.5 kcal/mol (or more) higher than any other energies in the scan. If so, print and
    log a warning message.
    """
    Vlist = numpy.array(Vlist, numpy.float64)
    Vdiff = (Vlist[0] - numpy.min(Vlist))*constants.E_h*constants.Na/1000
    if Vdiff >= 2: #we choose 2 kJ/mol to be the critical energy
        logging.warning('the species corresponding to ' + str(os.path.basename(path)) + ' is different in energy from the lowest energy conformer by ' + "%0.2f" % Vdiff + ' kJ/mol. This can cause significant errors in your computed rate constants. ')


## classify.py -- actually classify a sequence with DeepSpeech
##
## Copyright (C) 2017, Nicholas Carlini <nicholas@carlini.com>.
##
## This program is licenced under the BSD 2-Clause licence,
## contained in the LICENCE file in this directory.

import numpy as np
import tensorflow as tf
import argparse

import scipy.io.wavfile as wav

import time
import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''
import sys
from collections import namedtuple
sys.path.append("DeepSpeech")
import DeepSpeech

try:
    import pydub
    import struct
except:
    print("pydub was not loaded, MP3 compression will not work")

from tf_logits import get_logits

# These are the tokens that we're allowed to use.
# The - token is special and corresponds to the epsilon
# value in CTC decoding, and can not occur in the phrase.
toks = " abcdefghijklmnopqrstuvwxyz'-"

def main_classify(input, restore_path):
    # parser = argparse.ArgumentParser(description=None)
    # parser.add_argument('--in', type=str, dest="input",
    #                     required=True,
    #                     help="Input audio .wav file(s), at 16KHz (separated by spaces)")
    # parser.add_argument('--restore_path', type=str,
    #                     required=True,
    #                     help="Path to the DeepSpeech checkpoint (ending in model0.4.1)")
    # args = parser.parse_args()
    # while len(sys.argv) > 1:
    #     sys.argv.pop()
    with tf.Session() as sess:
        if input.split(".")[-1] == 'mp3':
            raw = pydub.AudioSegment.from_mp3(input)
            audio = np.array([struct.unpack("<h", raw.raw_data[i:i+2])[0] for i in range(0,len(raw.raw_data),2)])
        elif input.split(".")[-1] == 'wav' or input.split(".")[-1] == 'WAV':
            _, audio = wav.read(input)
        else:
            raise Exception("Unknown file format")
        N = len(audio)
        new_input = tf.placeholder(tf.float32, [1, N])
        lengths = tf.placeholder(tf.int32, [1])

        with tf.variable_scope("", reuse=tf.AUTO_REUSE):
            logits = get_logits(new_input, lengths)

        saver = tf.train.Saver()
        saver.restore(sess, restore_path)

        decoded, _ = tf.nn.ctc_beam_search_decoder(logits, lengths, merge_repeated=False, beam_width=500)

        # print('logits shape', logits.shape)
        length = (len(audio)-1)//320
        l = len(audio)
        r = sess.run(decoded, {new_input: [audio],
                               lengths: [length]})

        print("-"*80)
        print("-"*80)

        wav_name = input.split('/')[-1]
        print("Classification of ", wav_name)
        print("".join([toks[x] for x in r[0].values]))
        # print("-"*80)
        # print("-"*80)

        output_text = str("".join([toks[x] for x in r[0].values]))
        return wav_name, output_text

def main():
    parser = argparse.ArgumentParser(description=None)
    parser.add_argument('--in', type=str, dest="input",
                        required=True,
                        help="Input audio .wav file(s), at 16KHz (separated by spaces)")
    parser.add_argument('--restore_path', type=str,
                        required=True,
                        help="Path to the DeepSpeech checkpoint (ending in model0.4.1)")
    args = parser.parse_args()
    while len(sys.argv) > 1:
        sys.argv.pop()
    with tf.Session() as sess:
        if args.input.split(".")[-1] == 'mp3':
            raw = pydub.AudioSegment.from_mp3(args.input)
            audio = np.array([struct.unpack("<h", raw.raw_data[i:i+2])[0] for i in range(0,len(raw.raw_data),2)])
        elif args.input.split(".")[-1] == 'wav' or args.input.split(".")[-1] == 'WAV':
            _, audio = wav.read(args.input)
        else:
            raise Exception("Unknown file format")
        N = len(audio)
        new_input = tf.placeholder(tf.float32, [1, N])
        lengths = tf.placeholder(tf.int32, [1])

        with tf.variable_scope("", reuse=tf.AUTO_REUSE):
            logits = get_logits(new_input, lengths)

        saver = tf.train.Saver()
        saver.restore(sess, args.restore_path)

        decoded, _ = tf.nn.ctc_beam_search_decoder(logits, lengths, merge_repeated=False, beam_width=500)

        print('logits shape', logits.shape)
        length = (len(audio)-1)//320
        l = len(audio)
        r = sess.run(decoded, {new_input: [audio],
                               lengths: [length]})

        print("-"*80)
        print("-"*80)

        print("Classification:")
        print("".join([toks[x] for x in r[0].values]))
        print("-"*80)
        print("-"*80)

        output_text = "".join([toks[x] for x in r[0].values])
        return output_text

if __name__ == "__main__":

#     parser = argparse.ArgumentParser(description=None)
#     parser.add_argument('--in', type=str, dest="input",
#                         required=True,
#                         help="Input audio .wav file(s), at 16KHz (separated by spaces)")
#     parser.add_argument('--restore_path', type=str,
#                         required=True,
#                         help="Path to the DeepSpeech checkpoint (ending in model0.4.1)")
#     args = parser.parse_args()

#     output_text = main_classify(args.input, args.restore_path)
#     print(output_text)

    main()


# from https://github.com/ajbrock/BigGAN-PyTorch (MIT license)
# some modifications in class Generator and G_D
# new class "Unet_Discriminator" based on original class "Discriminator"
import numpy as np
import math
import functools

import torch
import torch.nn as nn
from torch.nn import init
import torch.optim as optim
import torch.nn.functional as F
from torch.nn import Parameter as P

from unetgan import layers
from unetgan import utils

import copy
from matplotlib import pyplot as plt

# Architectures for G
# Attention is passed in in the format '32_64' to mean applying an attention
# block at both resolution 32x32 and 64x64.
def G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
    arch = {}

    arch[256] = {'in_channels': [ch * item for item in [16, 16, 8, 8, 4, 2]],
                 'out_channels': [ch * item for item in [16, 8, 8, 4, 2, 1]],
                 'upsample': [True] * 6,
                 'resolution': [8, 16, 32, 64, 128, 256],
                 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])
                               for i in range(3, 9)}}
    arch[128] = {'in_channels': [ch * item for item in [16, 16, 8, 4, 2]],
                 'out_channels': [ch * item for item in [16, 8, 4, 2, 1]],
                 'upsample': [True] * 5,
                 'resolution': [8, 16, 32, 64, 128],
                 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])
                               for i in range(3, 8)}}
    return arch

class Generator(nn.Module):
    def __init__(self, G_ch=64, dim_z=128, bottom_width=4, resolution=128,
                 G_kernel_size=3, G_attn='64', n_classes=1000,
                 num_G_SVs=1, num_G_SV_itrs=1,
                 G_shared=True, shared_dim=0, hier=False,
                 cross_replica=False, mybn=False,
                 G_activation=nn.ReLU(inplace=False),
                 G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8,
                 BN_eps=1e-5, SN_eps=1e-12, G_mixed_precision=False, G_fp16=False,
                 G_init='ortho', skip_init=False, no_optim=False,
                 G_param='SN', norm_style='bn',
                 **kwargs):
        super(Generator, self).__init__()
        # Channel width mulitplier
        self.ch = G_ch
        # Dimensionality of the latent space
        self.dim_z = dim_z
        # The initial spatial dimensions
        self.bottom_width = bottom_width
        # Resolution of the output
        self.resolution = resolution
        # Kernel size?
        self.kernel_size = G_kernel_size
        # Attention?
        self.attention = G_attn
        # number of classes, for use in categorical conditional generation
        self.n_classes = n_classes
        # Use shared embeddings?
        self.G_shared = G_shared
        # Dimensionality of the shared embedding? Unused if not using G_shared
        self.shared_dim = shared_dim if shared_dim > 0 else dim_z
        # Hierarchical latent space?
        self.hier = hier
        # Cross replica batchnorm?
        self.cross_replica = cross_replica
        # Use my batchnorm?
        self.mybn = mybn
        # nonlinearity for residual blocks
        self.activation = G_activation
        # Initialization style
        self.init = G_init
        # Parameterization style
        self.G_param = G_param
        # Normalization style
        self.norm_style = norm_style
        # Epsilon for BatchNorm?
        self.BN_eps = BN_eps
        # Epsilon for Spectral Norm?
        self.SN_eps = SN_eps
        # fp16?
        self.fp16 = G_fp16
        # Architecture dict
        self.arch = G_arch(self.ch, self.attention)[resolution]

        self.unconditional = kwargs["unconditional"]

        # If using hierarchical latents, adjust z
        if self.hier:
            # Number of places z slots into
            self.num_slots = len(self.arch['in_channels']) + 1
            self.z_chunk_size = (self.dim_z // self.num_slots)

            if not self.unconditional:
                self.dim_z = self.z_chunk_size * self.num_slots
        else:
            self.num_slots = 1
            self.z_chunk_size = 0

        # Which convs, batchnorms, and linear layers to use
        if self.G_param == 'SN':
            self.which_conv = functools.partial(layers.SNConv2d,
                                                kernel_size=3, padding=1,
                                                num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
                                                eps=self.SN_eps)
            self.which_linear = functools.partial(layers.SNLinear,
                                                  num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
                                                  eps=self.SN_eps)
        else:
            self.which_conv = functools.partial(nn.Conv2d, kernel_size=3, padding=1)
            self.which_linear = nn.Linear

        # We use a non-spectral-normed embedding here regardless;
        # For some reason applying SN to G's embedding seems to randomly cripple G
        self.which_embedding = nn.Embedding

        if self.unconditional:
            bn_linear = nn.Linear
            input_size = self.dim_z + (self.shared_dim if self.G_shared else 0)
        else:
            bn_linear = (functools.partial(self.which_linear, bias=False) if self.G_shared
                         else self.which_embedding)

            input_size = (self.shared_dim + self.z_chunk_size if self.G_shared
                          else self.n_classes)
        self.which_bn = functools.partial(layers.ccbn,
                                          which_linear=bn_linear,
                                          cross_replica=self.cross_replica,
                                          mybn=self.mybn,
                                          input_size=input_size,
                                          norm_style=self.norm_style,
                                          eps=self.BN_eps,
                                          self_modulation=self.unconditional)

        # Prepare model
        # If not using shared embeddings, self.shared is just a passthrough
        self.shared = (self.which_embedding(n_classes, self.shared_dim) if G_shared
                       else layers.identity())
        # First linear layer
        if self.unconditional:
            self.linear = self.which_linear(self.dim_z, self.arch['in_channels'][0] * (self.bottom_width ** 2))
        else:
            self.linear = self.which_linear(self.dim_z // self.num_slots,
                                            self.arch['in_channels'][0] * (self.bottom_width ** 2))

        # self.blocks is a doubly-nested list of modules, the outer loop intended
        # to be over blocks at a given resolution (resblocks and/or self-attention)
        # while the inner loop is over a given block
        self.blocks = []
        for index in range(len(self.arch['out_channels'])):

            self.blocks += [[layers.GBlock(in_channels=self.arch['in_channels'][index],
                                           out_channels=self.arch['out_channels'][index],
                                           which_conv=self.which_conv,
                                           which_bn=self.which_bn,
                                           activation=self.activation,
                                           upsample=(functools.partial(F.interpolate, scale_factor=2)
                                                     if self.arch['upsample'][index] else None))]]

            # If attention on this block, attach it to the end
            if self.arch['attention'][self.arch['resolution'][index]]:
                print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])
                self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index], self.which_conv)]

        # Turn self.blocks into a ModuleList so that it's all properly registered.
        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])

        # output layer: batchnorm-relu-conv.
        # Consider using a non-spectral conv here
        self.output_layer = nn.Sequential(layers.bn(self.arch['out_channels'][-1],
                                                    cross_replica=self.cross_replica,
                                                    mybn=self.mybn),
                                          self.activation,
                                          self.which_conv(self.arch['out_channels'][-1], 3))

        # Initialize weights. Optionally skip init for testing.
        if not skip_init:
            self.init_weights()

        # Set up optimizer
        # If this is an EMA copy, no need for an optim, so just return now
        if no_optim:
            return
        self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps
        if G_mixed_precision:
            print('Using fp16 adam in G...')
            import utils
            self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
                                      betas=(self.B1, self.B2), weight_decay=0,
                                      eps=self.adam_eps)
        else:
            self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
                                    betas=(self.B1, self.B2), weight_decay=0,
                                    eps=self.adam_eps)

        # LR scheduling, left here for forward compatibility
        # self.lr_sched = {'itr' : 0}# if self.progressive else {}
        # self.j = 0

    # Initialize
    def init_weights(self):
        self.param_count = 0
        for module in self.modules():
            if (isinstance(module, nn.Conv2d)
                    or isinstance(module, nn.Linear)
                    or isinstance(module, nn.Embedding)):
                if self.init == 'ortho':
                    init.orthogonal_(module.weight)
                elif self.init == 'N02':
                    init.normal_(module.weight, 0, 0.02)
                elif self.init in ['glorot', 'xavier']:
                    init.xavier_uniform_(module.weight)
                else:
                    print('Init style not recognized...')
                self.param_count += sum([p.data.nelement() for p in module.parameters()])
        print('Param count for G''s initialized parameters: %d' % self.param_count)

    # Note on this forward function: we pass in a y vector which has
    # already been passed through G.shared to enable easy class-wise
    # interpolation later. If we passed in the one-hot and then ran it through
    # G.shared in this forward function, it would be harder to handle.
    def forward(self, z, y):
        # If hierarchical, concatenate zs and ys
        if self.hier:
            # faces
            if self.unconditional:
                ys = [z for _ in range(self.num_slots)]
            else:
                zs = torch.split(z, self.z_chunk_size, 1)
                z = zs[0]

                ys = [torch.cat([y, item], 1) for item in zs[1:]]
        else:
            if self.unconditional:
                ys = [None] * len(self.blocks)
            else:
                ys = [y] * len(self.blocks)

        # First linear layer
        h = self.linear(z)
        # Reshape
        h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)

        # Loop over blocks
        for index, blocklist in enumerate(self.blocks):
            # Second inner loop in case block has multiple layers
            for block in blocklist:
                h = block(h, ys[index])

        # Apply batchnorm-relu-conv-tanh at output
        return torch.tanh(self.output_layer(h))

# Discriminator architecture, same paradigm as G's above
def D_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
    arch = {}
    arch[256] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 8, 16]],
                 'out_channels': [item * ch for item in [1, 2, 4, 8, 8, 16, 16]],
                 'downsample': [True] * 6 + [False],
                 'resolution': [128, 64, 32, 16, 8, 4, 4],
                 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]
                               for i in range(2, 8)}}
    arch[128] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 16]],
                 'out_channels': [item * ch for item in [1, 2, 4, 8, 16, 16]],
                 'downsample': [True] * 5 + [False],
                 'resolution': [64, 32, 16, 8, 4, 4],
                 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]
                               for i in range(2, 8)}}
    return arch

def D_unet_arch(ch=64, attention='64', ksize='333333', dilation='111111', out_channel_multiplier=1):
    arch = {}

    n = 2

    ocm = out_channel_multiplier

    # covers bigger perceptual fields
    arch[128] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 16, 8 * n, 4 * 2, 2 * 2, 1 * 2, 1]],
                 'out_channels': [item * ch for item in [1, 2, 4, 8, 16, 8, 4, 2, 1, 1]],
                 'downsample': [True] * 5 + [False] * 5,
                 'upsample': [False] * 5 + [True] * 5,
                 'resolution': [64, 32, 16, 8, 4, 8, 16, 32, 64, 128],
                 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]
                               for i in range(2, 11)}}

    arch[256] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 8, 16, 8 * 2, 8 * 2, 4 * 2, 2 * 2, 1 * 2, 1]],
                 'out_channels': [item * ch for item in [1, 2, 4, 8, 8, 16, 8, 8, 4, 2, 1, 1]],
                 'downsample': [True] * 6 + [False] * 6,
                 'upsample': [False] * 6 + [True] * 6,
                 'resolution': [128, 64, 32, 16, 8, 4, 8, 16, 32, 64, 128, 256],
                 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]
                               for i in range(2, 13)}}

    return arch

class Unet_Discriminator(nn.Module):

    def __init__(self, D_ch=64, D_wide=True, resolution=128,
                 D_kernel_size=3, D_attn='64', n_classes=1000,
                 num_D_SVs=1, num_D_SV_itrs=1, D_activation=nn.ReLU(inplace=False),
                 D_lr=2e-4, D_B1=0.0, D_B2=0.999, adam_eps=1e-8,
                 SN_eps=1e-12, output_dim=1, D_mixed_precision=False, D_fp16=False,
                 D_init='ortho', skip_init=False, D_param='SN', decoder_skip_connection=True, **kwargs):
        super(Unet_Discriminator, self).__init__()

        # Width multiplier
        self.ch = D_ch
        # Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?
        self.D_wide = D_wide
        # Resolution
        self.resolution = resolution
        # Kernel size
        self.kernel_size = D_kernel_size
        # Attention?
        self.attention = D_attn
        # Number of classes
        self.n_classes = n_classes
        # Activation
        self.activation = D_activation
        # Initialization style
        self.init = D_init
        # Parameterization style
        self.D_param = D_param
        # Epsilon for Spectral Norm?
        self.SN_eps = SN_eps
        # Fp16?
        self.fp16 = D_fp16

        if self.resolution == 128:
            self.save_features = [0, 1, 2, 3, 4]
        elif self.resolution == 256:
            self.save_features = [0, 1, 2, 3, 4, 5]

        self.out_channel_multiplier = 1  # 4
        # Architecture
        self.arch = D_unet_arch(self.ch, self.attention, out_channel_multiplier=self.out_channel_multiplier)[resolution]

        self.unconditional = kwargs["unconditional"]

        # Which convs, batchnorms, and linear layers to use
        # No option to turn off SN in D right now
        if self.D_param == 'SN':
            self.which_conv = functools.partial(layers.SNConv2d,
                                                kernel_size=3, padding=1,
                                                num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
                                                eps=self.SN_eps)
            self.which_linear = functools.partial(layers.SNLinear,
                                                  num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
                                                  eps=self.SN_eps)

            self.which_embedding = functools.partial(layers.SNEmbedding,
                                                     num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
                                                     eps=self.SN_eps)
        # Prepare model
        # self.blocks is a doubly-nested list of modules, the outer loop intended
        # to be over blocks at a given resolution (resblocks and/or self-attention)
        self.blocks = []

        for index in range(len(self.arch['out_channels'])):

            if self.arch["downsample"][index]:
                self.blocks += [[layers.DBlock(in_channels=self.arch['in_channels'][index],
                                               out_channels=self.arch['out_channels'][index],
                                               which_conv=self.which_conv,
                                               wide=self.D_wide,
                                               activation=self.activation,
                                               preactivation=(index > 0),
                                               downsample=(
                                                   nn.AvgPool2d(2) if self.arch['downsample'][index] else None))]]

            elif self.arch["upsample"][index]:
                upsample_function = (
                    functools.partial(F.interpolate, scale_factor=2, mode="nearest")  # mode=nearest is default
                    if self.arch['upsample'][index] else None)

                self.blocks += [[layers.GBlock2(in_channels=self.arch['in_channels'][index],
                                                out_channels=self.arch['out_channels'][index],
                                                which_conv=self.which_conv,
                                                # which_bn=self.which_bn,
                                                activation=self.activation,
                                                upsample=upsample_function, skip_connection=True)]]

            # If attention on this block, attach it to the end
            attention_condition = index < 5
            if self.arch['attention'][self.arch['resolution'][index]] and attention_condition:  # index < 5
                print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])
                print("index = ", index)
                self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index],
                                                     self.which_conv)]

        # Turn self.blocks into a ModuleList so that it's all properly registered.
        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])

        last_layer = nn.Conv2d(self.ch * self.out_channel_multiplier, 1, kernel_size=1)
        self.blocks.append(last_layer)
        #
        # Linear output layer. The output dimension is typically 1, but may be
        # larger if we're e.g. turning this into a VAE with an inference output
        self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)

        self.linear_middle = self.which_linear(16 * self.ch, output_dim)
        # Embedding for projection discrimination
        # if not kwargs["agnostic_unet"] and not kwargs["unconditional"]:
        #    self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1]+extra)
        if not kwargs["unconditional"]:
            self.embed_middle = self.which_embedding(self.n_classes, 16 * self.ch)
            self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])

        # Initialize weights
        if not skip_init:
            self.init_weights()

        ###
        print("_____params______")
        for name, param in self.named_parameters():
            print(name, param.size())

        # Set up optimizer
        self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps
        if D_mixed_precision:
            print('Using fp16 adam in D...')
            import utils
            self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
                                      betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
        else:
            self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
                                    betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
        # LR scheduling, left here for forward compatibility
        # self.lr_sched = {'itr' : 0}# if self.progressive else {}
        # self.j = 0

    # Initialize
    def init_weights(self):
        self.param_count = 0
        for module in self.modules():
            if (isinstance(module, nn.Conv2d)
                    or isinstance(module, nn.Linear)
                    or isinstance(module, nn.Embedding)):
                if self.init == 'ortho':
                    init.orthogonal_(module.weight)
                elif self.init == 'N02':
                    init.normal_(module.weight, 0, 0.02)
                elif self.init in ['glorot', 'xavier']:
                    init.xavier_uniform_(module.weight)
                else:
                    print('Init style not recognized...')
                self.param_count += sum([p.data.nelement() for p in module.parameters()])
        print('Param count for D''s initialized parameters: %d' % self.param_count)

    def forward(self, x, y=None):
        # Stick x into h for cleaner for loops without flow control
        h = x

        residual_features = []
        residual_features.append(x)
        # Loop over blocks

        for index, blocklist in enumerate(self.blocks[:-1]):
            if self.resolution == 128:
                if index == 6:
                    h = torch.cat((h, residual_features[4]), dim=1)
                elif index == 7:
                    h = torch.cat((h, residual_features[3]), dim=1)
                elif index == 8:  #
                    h = torch.cat((h, residual_features[2]), dim=1)
                elif index == 9:  #
                    h = torch.cat((h, residual_features[1]), dim=1)

            if self.resolution == 256:
                if index == 7:
                    h = torch.cat((h, residual_features[5]), dim=1)
                elif index == 8:
                    h = torch.cat((h, residual_features[4]), dim=1)
                elif index == 9:  #
                    h = torch.cat((h, residual_features[3]), dim=1)
                elif index == 10:  #
                    h = torch.cat((h, residual_features[2]), dim=1)
                elif index == 11:
                    h = torch.cat((h, residual_features[1]), dim=1)

            for block in blocklist:
                h = block(h)

            if index in self.save_features[:-1]:
                residual_features.append(h)

            if index == self.save_features[-1]:
                # Apply global sum pooling as in SN-GAN
                h_ = torch.sum(self.activation(h), [2, 3])
                # Get initial class-unconditional output
                bottleneck_out = self.linear_middle(h_)
                # Get projection of final featureset onto class vectors and add to evidence
                if self.unconditional:
                    projection = 0
                else:
                    # this is the bottleneck classifier c
                    emb_mid = self.embed_middle(y)
                    projection = torch.sum(emb_mid * h_, 1, keepdim=True)
                bottleneck_out = bottleneck_out + projection

        out = self.blocks[-1](h)

        if self.unconditional:
            proj = 0
        else:
            emb = self.embed(y)
            emb = emb.view(emb.size(0), emb.size(1), 1, 1).expand_as(h)
            proj = torch.sum(emb * h, 1, keepdim=True)
            ################
        out = out + proj

        out = out.view(out.size(0), 1, self.resolution, self.resolution)

        return out, bottleneck_out

# Parallelized G_D to minimize cross-gpu communication
# Without this, Generator outputs would get all-gathered and then rebroadcast.
class G_D(nn.Module):
    def __init__(self, G, D, config):
        super(G_D, self).__init__()
        self.G = G
        self.D = D

        self.config = config

    def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=False,
                split_D=False, dw1=[], dw2=[], reference_x=None, mixup=False, mixup_only=False, target_map=None):

        if mixup:
            gy = dy
            # why? so the mixup samples consist of same class

        # If training G, enable grad tape
        with torch.set_grad_enabled(train_G):

            G_z = self.G(z, self.G.shared(gy))
            # Cast as necessary
            if self.G.fp16 and not self.D.fp16:
                G_z = G_z.float()
            if self.D.fp16 and not self.G.fp16:
                G_z = G_z.half()

        if mixup:
            initial_x_size = x.size(0)

            mixed = target_map * x + (1 - target_map) * G_z
            mixed_y = dy

        if not mixup_only:
            # we get here in the cutmix cons extra case
            D_input = torch.cat([G_z, x], 0) if x is not None else G_z
            D_class = torch.cat([gy, dy], 0) if dy is not None else gy
            dmap = torch.tensor([])
            if mixup:
                # we get here in the cutmix  "consistency loss and augmentation" case, if "mixup" is true for the current round (depends on p mixup)
                D_input = torch.cat([D_input, mixed], 0)
                if self.config["dataset"] != "coco_animals":
                    D_class = torch.cat([D_class.float(), mixed_y.float()], 0)
                else:
                    D_class = torch.cat([D_class.long(), mixed_y.long()], 0)
        else:
            # not reached in cutmix "consistency loss and augmentation"
            D_input = mixed
            D_class = mixed_y
            dmap = torch.tensor([])

            del G_z
            del x
            G_z = None
            x = None

        D_out, D_middle = self.D(D_input, D_class)

        del D_input
        del D_class

        if x is not None:

            if not mixup:
                out = torch.split(D_out, [G_z.shape[0], x.shape[0]])  # D_fake, D_real
            else:
                out = torch.split(D_out, [G_z.shape[0], x.shape[0], mixed.shape[0]])  # D_fake, D_real, D_mixed
            out = out + (G_z,)
            if mixup:
                out = out + (mixed,)

            if not mixup:
                D_middle = torch.split(D_middle, [G_z.shape[0], x.shape[0]])  # D_middle_fake, D_middle_real
            else:
                D_middle = torch.split(D_middle, [G_z.shape[0], x.shape[0], mixed.shape[0]])
            out = out + D_middle
            ###return target map as well
            if mixup:
                out = out + (target_map,)

            return out

        else:
            # in mixup# you arrive here
            out = (D_out,)

            if return_G_z:
                out = out + (G_z,)
            if mixup_only:
                out = out + (mixed,)

            out = out + (D_middle,)
            ##return target map as well
            if mixup:
                out = out + (target_map,)

            return out


import math, os, json, re
import torch
import torch.utils.data
import argparse

from dataset import collate_fn, TranslationDataset
import transformer.Constants as Constants
from transformer.Translator import Translator
from preprocess import read_instances_from_file, convert_instance_to_idx_seq

import sys
sys.path.append(".")
sys.path.append("..")

import numpy as np
from neuronlp2.io import get_logger, conllx_stacked_data
from neuronlp2.models import StackPtrNet
from neuronlp2 import utils
from neuronlp2.io import CoNLLXWriter
from neuronlp2.tasks import parser

from UPosTagMap import *
from SimplifyDepParse3 import SimplifyDepParse

class SentenceAnalyzer(object):
    def __init__(self, morphology_analysis="./models/MorphologicalAnalyzer.chkpt", morphology_vocab="./models/MorphologicalAnalyzerVocab.pt", dependency_parser_path="./models", dependency_parser_name="DependencyParser.pt"):
        self._init_morphology_analysis(morphology_analysis, morphology_vocab)
        self._init_dependency_parser(dependency_parser_path, dependency_parser_name)
        self.upostag_map = UPosTagMap()
        self.oo = SimplifyDepParse()

    def _init_morphology_analysis(self, model, vocab):
        parser = argparse.ArgumentParser(description='translate.py')

        parser.add_argument('-model', default=model,
                            help='Path to model .pt file')
        parser.add_argument('-vocab', default=vocab,
                            help='Source sequence to decode (one line per sequence)')
        parser.add_argument('-beam_size', type=int, default=5,
                            help='Beam size')
        parser.add_argument('-batch_size', type=int, default=1,
                            help='Batch size')
        parser.add_argument('-n_best', type=int, default=1,
                            help="""If verbose is set, will output the n_best
                                decoded sentences""")
        parser.add_argument('-no_cuda', action='store_true')

        opt = parser.parse_args()
        opt.cuda = not opt.no_cuda

        # Prepare DataLoader
        self.preprocess_data = torch.load(opt.vocab)
        self.preprocess_settings = self.preprocess_data['settings']

        self.translator = Translator(opt)

    def _init_dependency_parser(self, model_path="", model_name=""):
        args_parser = argparse.ArgumentParser(description='Testing with stack pointer parser')

        args_parser.add_argument('--model_path', help='path for parser model directory', default=model_path)
        args_parser.add_argument('--model_name', help='parser model file', default=model_name)
        args_parser.add_argument('--test', default="tmp.txt")
        args_parser.add_argument('--beam', type=int, default=1, help='Beam size for decoding')
        args_parser.add_argument('--batch_size', type=int, default=1)

        args = args_parser.parse_args()

        model_path = args.model_path
        model_name = os.path.join(model_path, args.model_name)
        self.test_path = args.test
        beam = args.beam
        batch_size = args.batch_size

        def load_args():
            with open("{}.arg.json".format(model_name)) as f:
                key_parameters = json.loads(f.read())

            return key_parameters['args'], key_parameters['kwargs']

        arguments, kwarguments = load_args()
        mode = arguments[10]
        input_size_decoder = arguments[11]
        hidden_size = arguments[12]
        arc_space = arguments[16]
        type_space = arguments[17]
        encoder_layers = arguments[13]
        decoder_layers = arguments[14]
        char_num_filters = arguments[6]
        eojul_num_filters = arguments[8]
        p_rnn = kwarguments['p_rnn']
        p_in = kwarguments['p_in']
        p_out = kwarguments['p_out']
        self.prior_order = kwarguments['prior_order']
        skipConnect = kwarguments['skipConnect']
        grandPar = kwarguments['grandPar']
        sibling = kwarguments['sibling']
        use_char = kwarguments['char']
        use_pos = kwarguments['pos']
        use_eojul = kwarguments['eojul']

        alphabet_path = os.path.join(model_path, 'alphabets/')
        self.word_alphabet, self.char_alphabet, self.pos_alphabet, self.type_alphabet = conllx_stacked_data.load_alphabets(alphabet_path)
        num_words = self.word_alphabet.size()
        num_chars = self.char_alphabet.size()
        num_pos = self.pos_alphabet.size()
        num_types = self.type_alphabet.size()

        word_table = None
        word_dim = arguments[0]
        char_table = None
        char_dim = arguments[2]
        pos_table = None
        pos_dim = arguments[4]

        char_window = arguments[7]
        eojul_window = arguments[9]

        self.dependency = StackPtrNet(word_dim, num_words, char_dim, num_chars, pos_dim, num_pos, char_num_filters, char_window, eojul_num_filters, eojul_window,
                              mode, input_size_decoder, hidden_size, encoder_layers, decoder_layers,
                              num_types, arc_space, type_space,
                              embedd_word=word_table, embedd_char=char_table, embedd_pos=pos_table, p_in=p_in, p_out=p_out, p_rnn=p_rnn,
                              biaffine=True, pos=use_pos, char=use_char, eojul=use_eojul, prior_order=self.prior_order,
                              skipConnect=skipConnect, grandPar=grandPar, sibling=sibling)

        self.dependency.load_state_dict(torch.load(model_name, map_location='cpu'))
        self.dependency.eval()

    def transform_sentence(self, sentence):
        tokens = sentence.split(" ")

        result = []
        result_org = []
        org_space = []
        for token in tokens:
            result_tmp_tokens = []
            if "[" in token and "]" not in token:
                tmp_token = token.replace("[", " [ ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif "[" not in token and "]" in token:
                tmp_token = token.replace("]", " ] ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif "(" in token and ")" not in token:
                tmp_token = token.replace("(", " ( ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif "(" not in token and ")" in token:
                tmp_token = token.replace(")", " ) ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif "{" in token and "}" not in token:
                tmp_token = token.replace("{", " { ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif "}" not in token and "}" in token:
                tmp_token = token.replace("}", " } ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif token.count("\'") == 1:
                tmp_token = token.replace("\'", " \' ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            elif token.count("\"") == 1:
                tmp_token = token.replace("\"", " \" ")
                tmp_tokens = [x for x in tmp_token.split(" ") if x != ""]
                result_tmp_tokens.extend(tmp_tokens)

            else:
                result_tmp_tokens.append(token)

            reuslt_tmp_tokens_length = len(result_tmp_tokens)

            for idx, t in enumerate(result_tmp_tokens, start=1):
                result.extend(t)
                result_org.append(t)
                if idx == reuslt_tmp_tokens_length:
                    org_space.append(False)
                else:
                    org_space.append(True)

                result.append("<sa>")

        return " ".join(result[:-1]), " ".join(result_org), org_space

    def get_instances(self, sentence, max_sent_len, keep_case):
        word_insts = []
        trimmed_sent_count = 0
        if not keep_case:
            sentence = sentence.lower()
        words = sentence.split()
        if len(words) > max_sent_len:
            trimmed_sent_count += 1
        word_inst = words[:max_sent_len]

        if word_inst:
            word_insts += [[Constants.BOS_WORD] + word_inst + [Constants.EOS_WORD]]
        else:
            word_insts += [None]
        return word_insts

    def convert_instance_to_idx_seq(self, word_insts, word2idx):
        ''' Mapping words to idx sequence. '''
        return [[word2idx.get(w, Constants.UNK) for w in s] for s in word_insts]

    def transform_result(self, sentence):
        result = []

        tokens = sentence.split(" ")
        m_tmp = []
        for token in tokens:
            if token == "<sa>":
                result.append(token)
            elif re.search("<[a-z]+>", token):
                morph = "".join(m_tmp)
                pos = token.upper().replace("<", "").replace(">", "")
                tok = "{}/{}".format(morph, pos)
                result.append(tok)
                m_tmp = []
            else:
                m_tmp.append(token)

        return " ".join(result)

    def convertToCoNLL(self, org_sentence, org_space_info, sentence):
        org_tokens = [x.strip() for x in org_sentence.split(" ")]
        tokens = [x.strip() for x in sentence.split("<sa>")]
        conll = []

        for idx, (token, org_token, space_info) in enumerate(zip(tokens, org_tokens, org_space_info) , start=1):
            only_mo = []
            only_pos = []
            for mo in token.split(" "):
                m, p = mo.split("/")
                if m == "<sf>":
                    m = "."
                only_mo.append(m)
                only_pos.append(p)
            input_xpostag = " ".join(only_pos)
            lemma = " ".join(only_mo)
            if space_info:
                tmp = [str(idx), org_token, " ".join(only_mo), "+".join(self.upostag_map.GetUPOS(input_xpostag, lemma)),
                       "+".join(only_pos), "_", "0", "root", "_", "SpaceAfter=No"]
            else:
                tmp = [str(idx), org_token, " ".join(only_mo), "+".join(self.upostag_map.GetUPOS(input_xpostag, lemma)), "+".join(only_pos), "_", "0", "root", "_", "_"]
            conll.append("\t".join(tmp))
        return "\n".join(conll)

    def morphologyAnalysis(self, sentence):
        sentence, org_sentence, org_space_info = self.transform_sentence(sentence)
        token_length = len(org_sentence.split(" "))
        test_src_word_insts = self.get_instances(sentence, self.preprocess_settings.max_word_seq_len, self.preprocess_settings.keep_case)
        test_src_insts = self.convert_instance_to_idx_seq(test_src_word_insts, self.preprocess_data['dict']['src'])

        test_loader = torch.utils.data.DataLoader(
            TranslationDataset(
                src_word2idx=self.preprocess_data['dict']['src'],
                tgt_word2idx=self.preprocess_data['dict']['tgt'],
                src_insts=test_src_insts),
            num_workers=2,
            batch_size=1,
            collate_fn=collate_fn)

        for batch in test_loader:
            all_hyp, all_scores = self.translator.translate_batch(*batch)
            for idx_seqs in all_hyp:
                for idx_seq in idx_seqs:
                    pred_line = " ".join([test_loader.dataset.tgt_idx2word[idx] for idx in idx_seq])

        result_line = self.transform_result(pred_line)
        result_length = len(result_line.split(" <sa> "))
        print(org_sentence.split(" "))
        print(result_line.split(" <sa> "))
        if token_length == result_length:
            return self.convertToCoNLL(org_sentence, org_space_info, result_line)

        return result_line

    def convert(self, sentences, word, pos, head, type, lengths, symbolic_root=False, symbolic_end=False):
        batch_size, _, lemma_length = word.shape
        start = 1 if symbolic_root else 0
        end = 1 if symbolic_end else 0
        result = []
        for i in range(batch_size):
            for j in range(start, lengths[i] - end):
                if len(sentences[i]) == j-start:
                    break

                t = self.type_alphabet.get_instance(type[i, j])
                h = head[i, j]

                result.append('%d\t%s\t%s\t%s\t%s\t%s\t%d\t%s\t%s\t%s\n' % (j, sentences[i][j-start][1], sentences[i][j-start][2], sentences[i][j-start][3], \
                                                                                sentences[i][j-start][4], sentences[i][j-start][5], h, t, sentences[i][j-start][8], sentences[i][j-start][9]))
        return result

    def dependencyparsing(self):
        pred_writer = CoNLLXWriter(self.word_alphabet, self.char_alphabet, self.pos_alphabet, self.type_alphabet)
        data_test = conllx_stacked_data.read_stacked_data_to_variable(self.test_path, self.word_alphabet, self.char_alphabet, self.pos_alphabet, self.type_alphabet, use_gpu=False,
                                                                      prior_order=self.prior_order)
        for batch in conllx_stacked_data.iterate_batch_stacked_variable(data_test, 1, use_gpu=False):
            input_encoder, _, sentences = batch
            word, char, pos, heads, types, masks, lengths = input_encoder

            heads_pred, types_pred, _, _ = self.dependency.decode(word, char, pos, mask=masks, length=lengths, beam=1, leading_symbolic=conllx_stacked_data.NUM_SYMBOLIC_TAGS)

            word = word.data.cpu().numpy()
            pos = pos.data.cpu().numpy()
            lengths = lengths.cpu().numpy()

            result = pred_writer.write_system(sentences, word, pos, heads_pred, types_pred, lengths, symbolic_root=True)

        return result

    def simplify(self, sentence):
        RESULT = self.oo.simplify_dep_parse(sentence)
        RESULT = self.oo.merge_linear_dep(RESULT, sentence)
        RESULT = self.oo.assign_impt_factor(RESULT, 1)
        result = self.oo.print_dep_parse(sentence, RESULT)
        return result

    def processing(self, sentence):
        conll = a.morphologyAnalysis(sentence)
        with open("tmp.txt", 'w') as f:
            f.write(conll)
        parsing = a.dependencyparsing()
        result = a.simplify(parsing)
        result.setdefault("conll", parsing)
        return result


import numpy as np

from nms_wrapper import *
from generate_anchors import generate_anchors
from config import *

def proposal_layer(rpn_cls_prob, rpn_bbox_pred, im_info, cfg_key, _feat_stride=[16, ],
                   anchor_scales=[8, 16, 32]):
    """
    :param rpn_cls_prob: 每个框的可能性，(1,18,W,H)
    :param rpn_bbox_pred: 预测出的每个框,(1,36,W,H)
    :param im_info: 图片信息，[高，宽，缩放倍数]
    :param cfg_key: TRAIN or TEST
    :param _feat_stride: 降采样倍数
    :param anchor_scales: anchor缩放倍数
    :return: (H*W*A,5),eg:[0,x1,y1,x2,y2]，A为feature map中每个点生成的anchor数，默认为9
    """
    # 根据anchor_scales生成的基础9个anchors
    _anchors = generate_anchors(scales=np.array(anchor_scales))  # (9,4)
    _num_anchors = _anchors.shape[0]

    scores = rpn_cls_prob[:, _num_anchors:, :, :]  # 背景评分
    bbox_deltas = rpn_bbox_pred  # 预测的前景框

    # 1.生成anchors
    # 生成W*H个偏移量，然后将基础anchors加上这些偏移量生成W*H*9个anchors
    height, width = scores.shape[-2:]  # (W,H)
    shift_x = np.arange(0, width) * _feat_stride  # 乘以降采样倍数，默认为16
    shift_y = np.arange(0, height) * _feat_stride
    shift_x, shift_y = np.meshgrid(shift_x, shift_y)  # 构成 width * height 个点， shift_x为x轴取值，shift_y为y轴取值
    shifts = np.vstack((shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel())).transpose()  # (W*H,4)

    # 共有9个anchors，每个anchor有W*H个变换方式
    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    A = _num_anchors  # 9
    K = shifts.shape[0]  # W*H
    anchors = _anchors.reshape((1, A, 4)) + shifts.reshape((1, K, 4)).transpose((1, 0, 2))
    anchors = anchors.reshape((K * A, 4))  # 最终生成K*A个anchors，这些anchors对于每个图片来说都是一样的，需要bbox_deltas来加以修正

    # reshape
    bbox_deltas = bbox_deltas.transpose((0, 2, 3, 1)).reshape((-1, 4))  # 原shape为(1,4,W,H)
    scores = scores.transpose((0, 2, 3, 1)).reshape((-1, 4))

    # 通过传入的bbox_deltas修正得到anchors
    proposals = bbox_transform_inv(anchors, bbox_deltas)

    # 2. 调整包含图像外侧的anchor
    proposals = clip_boxes(proposals, im_info[:2])

    # 3. 过滤掉比min_size小的anchor
    keep = _filter_boxes(proposals, RPN_MIN_SIZE * im_info[2])  # RPN_NMS_THRESH*图片缩放倍数
    proposals = proposals[keep, :]
    scores = scores[keep]

    # 4. 按照score排序获得的proposal
    # 5. 获取前RPN_PRE_NMS_TOP_N个anchor
    order = scores.ravel().argsort()[::-1]  # -1代表获取倒序
    if RPN_PRE_NMS_TOP_N>0:
        order=order[:RPN_PRE_NMS_TOP_N]
    proposals=proposals[order,:]
    scores=scores[order]

    # 6. 应用非极大值抑制算法，去除重合率大的冗余候选框，得到最具代表性的结果，以加快目标检测的效率
    keep=nms(np.hstack((proposals,scores)),RPN_NMS_THRESH)
    # 7. 获取前RPN_NMS_THRESH个应用了nms后的anchor
    if RPN_POST_NMS_TOP_N>0:
        keep=keep[:RPN_POST_NMS_TOP_N]
    proposals=proposals[keep,:]
    scores=scores[keep]

    batch_inds=np.zeros((proposals.shape[0],1),dtype=np.float32)
    blob=np.hstack((batch_inds,proposals.astype(np.float32,copy=False)))
    return blob

def bbox_transform_inv(boxes, deltas):
    """
    :param boxes: 生成的anchors
    :param deltas: 网络计算得出的deltas
    :return: 返回经过运算后偏移完成的anchors
    """
    if boxes.shape[0] == 0:
        return np.zeros((0, 0), dtype=deltas.dtype)
    boxes = boxes.astype(deltas.dtype, copy=False)

    # 得到宽度高度坐标
    widths = boxes[:, 2] - boxes[:, 0] + 1.0
    heights = boxes[:, 3] - boxes[:, 1] + 1.0
    ctr_xs = boxes[:, 0] + widths * 0.5
    ctr_ys = boxes[:, 1] + heights * 0.5

    # 得到deltas
    dx = deltas[:, 0::4]
    dy = deltas[:, 1::4]
    dw = deltas[:, 2::4]
    dh = deltas[:, 3::4]

    # 生成预测的anchors
    pred_ctr_x = dx * widths[:, np.newaxis] + ctr_xs[:, np.newaxis]
    pred_ctr_y = dy * heights[:, np.newaxis] + ctr_ys[:np.newaxis]
    pre_w = np.exp(dw) * widths[:, np.newaxis]
    pre_h = np.exp(dh) * heights[:, np.newaxis]

    pred_boxes = np.zeros(deltas.shape, dtype=deltas.dtype)
    # x1
    pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pre_w
    pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pre_h
    pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pre_w
    pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pre_h

    return pred_boxes

def clip_boxes(boxes, im_shape):
    """
    调整包含图像外侧的anchor
    :param boxes: 所有anchor
    :param im_shape: 图像的shape
    :return: 调整后的anchor
    """
    if boxes.shape[0] == 0:
        return boxes

    boxes[:, 0::4] = np.maximum(np.minimum(boxes[:, 0::4], im_shape[1] - 1), 0)
    boxes[:, 1::4] = np.maximum(np.minimum(boxes[:, 1::4], im_shape[0] - 1), 0)
    boxes[:, 2::4] = np.maximum(np.minimum(boxes[:, 2::4], im_shape[1] - 1), 0)
    boxes[:, 3::4] = np.maximum(np.minimum(boxes[:, 3::4], im_shape[0] - 1), 0)

def _filter_boxes(boxes, min_size):
    """
    剔除所有比min_size小的anchor
    :param boxes: 所有的anchors
    :param min_size: 最小的size
    :return: 剔除后的anchor
    """
    ws = boxes[:, 2] - boxes[:, 0] + 1.0
    hs = boxes[:, 3] - boxes[:, 1] + 1.0
    keep = np.where((ws >= min_size) & (hs >= min_size))[0]
    return keep