""" Main implementation of parcel model.
"""
from __future__ import print_function
from __future__ import absolute_import
from builtins import zip
from builtins import range
from builtins import object

import os

from scipy.optimize import bisect

# Parcel model imports
from . import constants as c
from . import output
from . aerosol import AerosolSpecies
# from . integrator import Integrator
from . thermo import *

__all__ = ['ParcelModel', ]

class ParcelModelError(Exception):
    """ Custom exception to throw during parcel model execution.
    """
    def __init__(self, error_str):
        self.error_str = error_str

    def __str__(self):
        return repr(self.error_str)

class ParcelModel(object):
    """ Wrapper class for instantiating and running the parcel model.

    The parcel model has been implemented in an object-oriented format to facilitate
    easy extensibility to different aerosol and meteorological conditions. A
    typical use case would involve specifying the initial conditions such as:

    >>> import pyrcel as pm
    >>> P0 = 80000.
    >>> T0 = 283.15
    >>> S0 = 0.0
    >>> V = 1.0
    >>> aerosol1 = pm.AerosolSpecies('sulfate',
    ...                              Lognorm(mu=0.025, sigma=1.3, N=2000.),
    ...                              bins=200, kappa=0.54)
    >>> initial_aerosols = [aerosol1, ]
    >>> z_top = 50.
    >>> dt = 0.01

    which initializes the model with typical conditions at the top of the boundary
    layer (800 hPa, 283.15 K, 100% Relative Humidity, 1 m/s updraft), and a simple
    sulfate aerosol distribution which will be discretized into 200 size bins to
    track. Furthermore the model was specified to simulate the updraft for 50
    meters (`z_top`) and use a time-discretization of 0.01 seconds. This
    timestep is used in the model output -- the actual ODE solver will generally
    calculate the trace of the model at many more times.

    Running the model and saving the output can be accomplished by invoking::

    >>> model = pm.ParcelModel(initial_aerosols, V, T0, S0, P0)
    >>> par_out, aer_out = pm.run(z_top, dt)

    This will yield `par_out`, a  :class:`pandas.DataFrame` containing the meteorological
    conditions in the parcel, and `aerosols`, a dictionary of :class:`DataFrame` objects
    for each species in `initial_aerosols` with the appropriately tracked size
    bins and their evolution over time.

    Attributes
    ----------
    V, T0, S0, P0, aerosols : floats
        Initial parcel settings (see **Parameters**).
    _r0s : array_like of floats
        Initial equilibrium droplet sizes.
    _r_drys : array_like of floats
        Dry radii of aerosol population.
    _kappas : array_like of floats
        Hygroscopicity of each aerosol size.
    _Nis : array_like of floats
        Number concentration of each aerosol size.
    _nr : int
        Number of aerosol sizes tracked in model.
    _model_set : boolean
        Flag indicating whether or not at any given time the model
        initialization/equilibration routine has been run with the current
        model settings.
    _y0 : array_like
        Initial state vector.

    Methods
    -------
    run(t_end, dt, max_steps=1000, solver="odeint", output_fmt="dataframes",\
        terminate=False, solver_args={})
        Execute model simulation.
    set_initial_conditions(V=None, T0=None, S0=None, P0=None, aerosols=None)
        Re-initialize a model simulation in order to run it.

    See Also
    --------
    _setup_run : companion routine which computes equilibrium droplet sizes
                 and sets the model's state vectors.

    """

    def __init__(self, aerosols, V, T0, S0, P0, console=False, accom=c.ac,
                 truncate_aerosols=False):
        """ Initialize the parcel model.

        Parameters
        ----------
        aerosols : array_like sequence of :class:`AerosolSpecies`
            The aerosols contained in the parcel.
        V, T0, S0, P0 : floats
            The updraft speed and initial temperature (K), pressure (Pa),
            supersaturation (percent, with 0.0 = 100% RH).
        console : boolean, optional
            Enable some basic debugging output to print to the terminal.
        accom : float, optional (default=:const:`constants.ac`)
            Condensation coefficient
        truncate_aerosols : boolean, optional (default=**False**)
            Eliminate extremely small aerosol which will cause numerical problems

        """
        self._model_set = False
        self.console = console
        self.trunc = truncate_aerosols

        self.V = V
        self.T0 = T0
        self.S0 = S0
        self.P0 = P0
        if self.trunc:
            self.aerosols = self._replace_aerosol(aerosols)
        else:
            self.aerosols = aerosols
        self.accom = accom

        # To be set by call to "self._setup_run()"
        self._r0s = None
        self._r_drys = None
        self._kappas = None
        self._Nis = None
        self._nr = None

        self._setup_run()

    def set_initial_conditions(self, V=None, T0=None, S0=None, P0=None, aerosols=None):
        """ Set the initial conditions and parameters for a new parcel
        model run without having to create a new :class:`ParcelModel` instance.

        Based on the aerosol population which has been stored in the model, this
        method will finish initializing the model. This has three major parts:

        1. concatenate the aerosol population information (their dry radii,
           hygroscopicities, etc) into single arrays which can be placed into the
           state vector for forward integration.
        2. Given the initial ambient water vapor concentration (computed from the
           temperature, pressure, and supersaturation), determine how much water
           must already be coated on the aerosol particles in order for their
           size to be in equilibrium.
        3. Set-up the state vector with these initial conditions.

        Once the state vector has been set up, the setup routine will record
        attributes in the parent instance of the :class:`ParcelModel`.

        Parameters
        ----------
        V, T0, S0, P0 : floats
            The updraft speed and initial temperature (K), pressure (Pa),
            supersaturation (percent, with 0.0 = 100% RH).
        aerosols : array_like sequence of :class:`AerosolSpecies`
            The aerosols contained in the parcel.

        Raises
        ------
        ParcelModelError
            If an equilibrium droplet size distribution could not be calculated.

        Notes
        -----
        The actual setup occurs in the private method `_setup_run()`; this
        method is simply an interface that can be used to modify an existing
        :class:`ParcelModel`.

        """

        if V:
            self.V = V
        if T0:
            self.T0 = T0
        if P0:
            self.P0 = P0
        if S0:
            self.S0 = S0
        if aerosols:
            if self.trunc:
                self.aerosols = self._replace_aerosol(aerosols)
            else:
                self.aerosols = aerosols

        if T0 or P0 or S0 or aerosols:
            self._setup_run()

    def _replace_aerosol(self, aerosols, smallest_rdry=1e-10):
        """ Truncates the bins with the smallest particles so that none
        have a dry radius of less than 0.1 nm.

        Parameters
        ----------
        aerosols : array_like sequence of :class:`AerosolSpecies`
            The aerosols contained in the parcel.
        smallest_rdry: float
            Smallest allowable dry radius bin.

        Returns
        -------
        fixed_aerosols : array_like sequence of :class:`AerosolSpecies`
            The modified sequence of aerosols.

        """
        fixed_aerosols = []

        for aer in aerosols:
            N_old = np.sum(aer.Nis)
            bin_count = np.count_nonzero(aer.r_drys[aer.r_drys < smallest_rdry])
            if bin_count > 0:
                aer_new = AerosolSpecies(aer.species,
                                         aer.distribution,
                                         kappa=aer.kappa,
                                         bins=aer.bins,
                                         r_min=smallest_rdry*1e6)
                N_new = np.sum(aer_new.Nis)
                if self.console:
                    print("%s: Removed %03d bins, N change: %5.1f -> %5.1f (%2.1f%%)" %
                          (aer.species, bin_count, N_old, N_new, (N_new-N_old)/N_new))
                fixed_aerosols.append(aer_new)
            else:
                fixed_aerosols.append(aer)

        return fixed_aerosols

    def _setup_run(self):
        """ Perform equilibration and initialization calculations for the
        parcel model simulation.

        .. note:: See `set_initial_conditions` for full details.

        """
        T0, S0, P0 = self.T0, self.S0, self.P0
        z0 = 0.0
        out = dict()

        if self.console:
            print("Setting up parcel model initial conditions...")

        # 1) Setup aerosols
        # a) grab all the initial aerosol size/concentrations
        species = []
        r_drys, Nis, kappas = [], [], []
        for aerosol in self.aerosols:
            r_drys.extend(aerosol.r_drys)
            kappas.extend([aerosol.kappa]*aerosol.nr)
            Nis.extend(aerosol.Nis)
            species.extend([aerosol.species]*aerosol.nr)

        r_drys = np.array(r_drys)
        kappas = np.array(kappas)
        Nis = np.array(Nis)

        out['r_drys'] = r_drys
        out['kappas'] = kappas
        out['Nis'] = Nis

        if self.console:
            # TODO: Fix print table formatting
            print("AEROSOL DISTRIBUTION")
            print("%8s %6s" % ("r", "N"))
            for sp, r, N in zip(species, r_drys, Nis):
                print("%10s %2.2e %4.1f" % (sp, r, N))
            print("\n"+"-"*44)

        # 2) Setup parcel initial conditions
        # a) water vapor
        #        RH    * (Rd/Rv = epsilon) * ( es / P - es )
        wv0 = (S0 + 1.)*(c.epsilon*es(T0-273.15)/(P0-es(T0-273.15)))  # Water Vapor mixing ratio, kg/kg

        # b) find equilibrium wet particle radius
        # wrapper function for quickly computing deviation from chosen
        # equilibrium supersaturation given current size and dry size
        f = lambda r, r_dry, kappa: Seq(r, r_dry, T0, kappa) - S0
        # Compute the equilibrium wet particle radii
        r0s = []
        # r_b, _ = kohler_crit(T0, r_drys[-1], kappas[-1])
        # for r_dry , kappa in zip(r_drys, kappas)[::-1]:
        for r_dry, kappa in zip(reversed(r_drys), reversed(kappas)):
        # Locate the critical value (f(r_crit) > 0), and use bisection from it and
        # r_dry (f(r_dry) < 0) to find equilibrium wet particle radius for a given S0
            r_b, _ = kohler_crit(T0, r_dry, kappa)
            r_a = r_dry

            r0 = bisect(f, r_a, r_b, args=(r_dry, kappa), xtol=1e-30, maxiter=500)
            r0s.append(r0)

        r0s = np.array(r0s[::-1])

        # Console logging output, if requested, of the equilibrium calcuations. Useful for
        # checking if the computations worked
        raised = False
        for (r, r_dry, sp, kappa) in zip(r0s, r_drys, species, kappas):
            ss = Seq(r, r_dry, T0, kappa)
            # rc, _ = kohler_crit(T0, r_dry, kappa)
            if r < 0:
                if self.console:
                    print("Found bad r", r, r_dry, sp)
                raised = True
            if np.abs(ss-S0) > 1e-4:
                if self.console:
                    print("Found S discrepancy", ss, S0, r_dry)
                raised = True
        if raised:
            raise ParcelModelError("Couldn't calculate initial aerosol population wet sizes.")
        out['r0s'] = r0s

        # c) compute equilibrium droplet water content
        water_vol = lambda r0, r_dry, Ni: (4.*np.pi/3.)*rho_w*Ni*(r0**3 - r_dry**3)
        wc0 = np.sum([water_vol(r0, r_dry, Ni) for r0, r_dry, Ni in zip(r0s, r_drys, Nis)])
        wc0 /= rho_air(T0, P0, 0.)

        # d) compute initial ice water content
        wi0 = 0.0

        # e) concatenate into initial conditions arrays
        y0 = [z0, P0, T0, wv0, wc0, wi0, S0]
        if self.console:
            print("PARCEL INITIAL CONDITIONS")
            print(("    " + "{:>9} "*6).format("P (hPa)", "T (K)", "wv (g/kg)",
                                               "wc (g/kg)", "wi (g/kg)", "S"))
            print("    " + "{:9.1f} {:9.2f} {:9.1e} {:9.1e} {:9.1e} {:9.3f}".format(
                           P0/100., T0, wv0*1e3, wc0*1e3, wi0*1e3, S0))
        y0.extend(r0s)
        y0 = np.array(y0)
        out['y0'] = y0
        self.y0 = y0

        # Store the model configuration
        self._r0s = r0s
        self._r_drys = r_drys
        self._kappas = kappas
        self._Nis = Nis
        self._nr = len(r_drys)

        self._model_set = True

        if self.console:
            print("Initial conditions set successfully.")

    def run(self, t_end,
            output_dt=1., solver_dt=None,
            max_steps=1000, solver="odeint", output_fmt="dataframes",
            terminate=False, terminate_depth=100., **solver_args):
        """ Run the parcel model simulation.

        Once the model has been instantiated, a simulation can immediately be
        performed by invoking this method. The numerical details underlying the
        simulation and the times over which to integrate can be flexibly set
        here.

        **Time** -- The user must specify two timesteps: `output_dt`, which is the
        timestep between output snapshots of the state of the parcel model, and
        `solver_dt`, which is the the interval of time before the ODE integrator
        is paused and re-started. It's usually okay to use a very large `solver_dt`,
        as `output_dt` can be interpolated from the simulation. In some cases though
        a small `solver_dt` could be useful to force the solver to use smaller
        internal timesteps.

        **Numerical Solver** -- By default, the model will use the `odeint` wrapper
        of LSODA shipped by default with scipy. Some fine-tuning of the solver tolerances
        is afforded here through the `max_steps`. For other solvers, a set of optional
        arguments `solver_args` can be passed.

        **Solution Output** -- Several different output formats are available by default.
        Additionally, the output arrays are saved with the `ParcelModel` instance so they
        can be used later.

        Parameters
        ----------
        t_end : float
            Total time over interval over which the model should be integrated
        output_dt : float
            Timestep intervals to report model output.
        solver_dt : float
            Timestep interval for calling solver integration routine.
        max_steps : int
            Maximum number of steps allowed by solver to satisfy error tolerances
            per timestep.
        solver : {'odeint', 'lsoda', 'lsode', 'vode', cvode'}
            Choose which numerical solver to use:
            * `'odeint'`: LSODA implementation from ODEPACK via
                SciPy's integrate module
            * `'lsoda'`: LSODA implementation from ODEPACK via odespy
            * `'lsode'`: LSODE implementation from ODEPACK via odespy
            * `'vode'` : VODE implementation from ODEPACK via odespy
            * `'cvode'` : CVODE implementation from Sundials via Assimulo
            * `'lsodar'` : LSODAR implementation from Sundials via Assimulo
        output_fmt : str, one of {'dataframes', 'arrays', 'smax'}
            Choose format of solution output.
        terminate : boolean
            End simulation at or shortly after a maximum supersaturation has been achieved
        terminate_depth : float, optional (default=100.)
            Additional depth (in meters) to integrate after termination criterion
            eached.

        Returns
        -------
        DataFrames, array, or float
            Depending on what was passed to the *output* argument, different
            types of data might be returned:

            - `dataframes': (default) will process the output into
               two pandas DataFrames - the first one containing profiles
               of the meteorological quantities tracked in the model,
               and the second a dictionary of DataFrames with one for
               each AerosolSpecies, tracking the growth in each bin
               for those species.
            - 'arrays': will return the raw output from the solver
               used internally by the parcel model - the state vector
               `y` and the evaluated timesteps converted into height
               coordinates.
            - 'smax': will only return the maximum supersaturation
               value achieved in the simulation.

        Raises
        ------
        ParcelModelError
            The parcel model failed to complete successfully or failed to initialize.

        See Also
        --------
        der : right-hand side derivative evaluated during model integration.

        """
        from . integrator import Integrator

        if output_fmt not in ["dataframes", "arrays", "smax"]:
            raise ParcelModelError("Invalid value ('%s') specified for output format." % output)

        if solver_dt is None:
            solver_dt = 10.*output_dt

        if not self._model_set:
            self._setup_run()

        y0 = self.y0
        r_drys = self._r_drys
        kappas = self._kappas
        Nis = self._Nis
        nr = self._nr

        # Setup/run integrator
        try:
            from .parcel_aux import der as der_fcn
        except ImportError:
            print("Could not load Cython derivative; using Python version.")
            from .parcel import der as der_fcn
        # Hack in Python derivative function
        # der_fcn = der

        # Is the updraft speed a function of time?
        v_is_func = hasattr(self.V, '__call__')
        if v_is_func:  # Re-wrap the function to correctly figure out V
            orig_der_fcn = der_fcn

            def der_fcn(y, t, *args):
                V_t = self.V(t)
                args[3] = V_t
                return orig_der_fcn(y, t, *args)

        # Will the simulation terminate early?
        if not terminate:
            terminate_depth = 0.
        else:
            if terminate_depth <= 0.:
                raise ParcelModelError("`terminate_depth` must be greater than 0!")

        if self.console:
            print()
            print("Integration control")
            print("----------------------------")
            print("        output dt: ", output_dt)
            print("    max solver dt: ", solver_dt)
            print(" solver int steps: ", int(solver_dt/output_dt))
            print("      termination: %r (%5dm)" % (terminate, terminate_depth))

        args = [nr, r_drys, Nis, self.V, kappas, self.accom]
        integrator_type = Integrator.solver(solver)
        integrator = integrator_type(der_fcn, output_dt, solver_dt, y0, args,
                                     terminate=terminate, terminate_depth=terminate_depth,
                                     console=self.console,
                                     **solver_args)
        success = False
        try:
            # Pack args as tuple for solvers
            args = tuple(args)

            if self.console:
                print("\nBEGIN INTEGRATION ->\n")
            x, t, success = integrator.integrate(t_end)
        except ValueError as e:
            raise ParcelModelError("Something failed during model integration: %r" % e)
        finally:
            if not success:
                raise ParcelModelError("Something failed during model integration.")

            # Success if reached this point!
            if self.console:
                print("\nEND INTEGRATION <-\n")

            self.x = x
            self.heights = self.x[:, c.STATE_VAR_MAP['z']]
            self.time = t

            if output_fmt == "dataframes":
                return output.parcel_to_dataframes(self)
            elif output_fmt == "arrays":
                return self.x, self.heights
            elif output_fmt == "smax":
                S = self.x[:, c.STATE_VAR_MAP['S']]
                return S.max()

    def write_summary(self, parcel_data, aerosol_data, out_filename):
        """ Write a quick and dirty summary of given parcel model output to the
        terminal.
        """
        from .activation import lognormal_activation
        # Check if parent dir of out_filename exists, and if not,
        # create it
        out_dir = os.path.dirname(out_filename)
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        # Open a file to write to
        with open(out_filename, 'w') as out_file:
            # Write header
            out_file.write("PARCEL MODEL\n")
            out_file.write("--------------------------------------\n")

            # Write initial conditions
            out_file.write("V = %f\n" % self.V)
            out_file.write("P0 = %f\n" % self.P0)
            out_file.write("T0 = %f\n" % self.T0)
            out_file.write("S0 = %f\n" % self.S0)
            out_file.write("--------------------------------------\n")

            # Write aerosol details
            for aerosol in self.aerosols:
                out_file.write(aerosol.species+" = "+aerosol.summary_str()+"\n")
            out_file.write("--------------------------------------\n")

            # Write simulation summary results
            # 1) Maximum supersaturation in parcel
            S_max = parcel_data['S'].max()
            S_max_idx = np.argmax(parcel_data.S)
            out_file.write("S_max = %f\n" % S_max)

            # 2) Activated fraction of each species
            T_at_S_max = parcel_data['T'].ix[S_max_idx]
            total_number = 0.0
            total_activated = 0.0
            for aerosol in self.aerosols:
                act_frac = lognormal_activation(S_max, aerosol.mu*1e-6, aerosol.sigma,
                                                aerosol.N, aerosol.kappa, T=T_at_S_max)
                act_num = act_frac*aerosol.N
                out_file.write("%s - eq_act_frac = %f (%3.2f/%3.2f)\n" %
                               (aerosol.species, act_frac, act_num, aerosol.N))

                total_number += aerosol.N
                total_activated += act_num
            total_act_frac = total_activated/total_number
            out_file.write("Total activated fraction = %f (%3.2f/%3.2f)\n" %
                           (total_act_frac, total_activated, total_number))

    def save(self, filename=None, format="nc", other_dfs=None):
        output.write_parcel_output(filename=filename, format=format, parcel=self,
                                   other_dfs=other_dfs)

    @staticmethod
    def write_csv(parcel_data, aerosol_data, output_dir=None):
        """Write output to CSV files.

        Utilize pandas fast output procedures to write the model run output to a
        set of CSV files. Written as a static method so that any prior saved
        parcel model output can be saved to disk in batch, even after its
        associated model has been destroyed.

        **Args:
            * *parcel_data* -- Pandas DataFrame of the parcel thermodynamic profile
            * *aerosol_data* -- dictionary of pandas DataFrames with the aerosol radii \
                at each model step
            * *output_dir* -- String to location where output should be saved; if not \
                provided then the model will save to the current path.

        """
        if not output_dir:
            output_dir = os.getcwd()

        # Write parcel data
        parcel_data.to_csv(os.path.join(output_dir, "parcel.csv"))

        # Write aerosol data
        for species, data in list(aerosol_data.items()):
            data.to_csv(os.path.join(output_dir, "%s.csv" % species))

def der(y, t, nr, r_drys, Nis, V, kappas, accom=c.ac):
    """ Calculates the instantaneous time-derivate of the parcel model system.

    Given a current state vector `y` of the parcel model, computes the tendency
    of each term including thermodynamic (pressure, temperature, etc) and aerosol
    terms. The basic aerosol properties used in the model must be passed along
    with the state vector (i.e. if being used as the callback function in an ODE
    solver).

    This function is implemented in NumPy and Python, and is likely *very* slow
    compared to the available Cython version.

    Parameters
    ----------
    y : array_like
        Current state of the parcel model system,
            * y[0] = altitude, m
            * y[1] = Pressure, Pa
            * y[2] = temperature, K
            * y[3] = water vapor mass mixing ratio, kg/kg
            * y[4] = cloud liquid water mass mixing ratio, kg/kg
            * y[5] = cloud ice water mass mixing ratio, kg/kg
            * y[6] = parcel supersaturation
            * y[7:] = aerosol bin sizes (radii), m
    t : float
        Current simulation time, in seconds.
    nr : Integer
        Number of aerosol radii being tracked.
    r_drys : array_like
        Array recording original aerosol dry radii, m.
    Nis : array_like
        Array recording aerosol number concentrations, 1/(m**3).
    V : float
        Updraft velocity, m/s.
    kappas : array_like
        Array recording aerosol hygroscopicities.
    accom : float, optional (default=:const:`constants.ac`)
        Condensation coefficient.

    Returns
    -------
    x : array_like
        Array of shape (``nr``+7, ) containing the evaluated parcel model
        instaneous derivative.

    Notes
    -----
    This Python sketch of the derivative function shouldn't really be used for
    any computational purposes. Instead, see the cythonized version in the auxiliary
    file, **parcel_aux.pyx**. In the default configuration, once the code has been
    built, you can set the environmental variable **OMP_NUM_THREADS** to control
    the parallel for loop which calculates the condensational growth rate for each
    bin.

    """
    z = y[0]
    P = y[1]
    T = y[2]
    wv = y[3]
    wc = y[4]
    wi = y[5]
    S = y[6]
    rs = np.asarray(y[c.N_STATE_VARS:])

    T_c = T - 273.15   # convert temperature to Celsius
    pv_sat = es(T-T_c) # saturation vapor pressure
    wv_sat = wv/(S+1.)  # saturation mixing ratio
    Tv = (1.+0.61*wv)*T  # virtual temperature given parcel humidity
    e = (1. + S)*pv_sat  # water vapor pressure
    rho_air = P/(Rd*Tv)  # current air density accounting for humidity
    rho_air_dry = (P-e)/Rd/T  # dry air density

    # 1) Pressure
    dP_dt = -1.*rho_air*g*V

    # 2/3) Wet particle growth rates and droplet liquid water
    drs_dt = np.zeros(shape=(nr, ))
    dwc_dt = 0.
    for i in range(nr):
        r = rs[i]
        r_dry = r_drys[i]
        kappa = kappas[i]

        dv_r = dv(T, r, P, accom)
        ka_r = ka(T, rho_air, r)

        G_a = (rho_w*R*T)/(pv_sat*dv_r*Mw)
        G_b = (L*rho_w*((L*Mw/(R*T))-1.))/(ka_r*T)
        G = 1./(G_a + G_b)

        delta_S = S - Seq(r, r_dry, T, kappa)
        dr_dt = (G/r)*delta_S
        Ni = Nis[i]
        dwc_dt += Ni*(r*r)*dr_dt
        drs_dt[i] = dr_dt

        # if i == nr-1:
        #    print 1, r, r_dry, Ni
        #    print 2, dv(T, r, P, accom), ka(T, rho_air, r)
        #    print 3, G_a, G_b
        #    print 4, Seq(r, r_dry, T, kappa, 1.0)
        #    print 5, dr_dt

    dwc_dt *= 4.*np.pi*rho_w/rho_air_dry

    # 4) ice water content
    dwi_dt = 0.0

    # 5) Water vapor content
    dwv_dt = -1.*(dwc_dt + dwi_dt)

    # 6) Temperature
    dT_dt = -g*V/Cp - L*dwv_dt/Cp

    # 7) Supersaturation
    alpha = (g*Mw*L)/(Cp*R*(T**2)) - (g*Ma)/(R*T)
    gamma = (P*Ma)/(Mw*es(T-273.15)) + (Mw*L*L)/(Cp*R*T*T)
    dS_dt = alpha*V - gamma*dwc_dt

    dz_dt = V

    # Repackage tendencies for feedback to numerical solver
    x = np.zeros(shape=(nr+c.N_STATE_VARS, ))
    x[0] = dz_dt
    x[1] = dP_dt
    x[2] = dT_dt
    x[3] = dwv_dt
    x[4] = dwc_dt
    x[5] = dwi_dt
    x[6] = dS_dt
    x[c.N_STATE_VARS:] = drs_dt[:]

    return x


"""
Convolutional neural nets driving code for TF2.0

Created on 04/26/2019

@author: RH
"""
from datetime import datetime
import os
import time
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import tensorflow as tf
import Accessory2 as ac
import Panoptes1

# Define an Inception
class INCEPTION:
    def __init__(self, input_dim, d_hyperparams=None,
                 save_graph_def=True, meta_graph=None, transfer=False,
                 log_dir="./log", meta_dir="./meta", weights=tf.constant([1., 1., 1., 1.])):
        if d_hyperparams is None:
            d_hyperparams = {}
        self.input_dim = input_dim
        self.batch_size = d_hyperparams['batch_size']
        self.dropout = d_hyperparams['dropout']
        self.learning_rate = d_hyperparams['learning_rate']
        self.sesh = tf.Session()
        self.weights = weights
        self.transfer = transfer

        if meta_graph:  # load saved graph
            model_name = os.path.basename(meta_graph)
            tf.train.import_meta_graph(meta_dir + '/' + model_name +'.meta').restore(
                self.sesh, meta_dir + '/' + model_name)
            handles = self.sesh.graph.get_collection("cnn_to_restore")

        else:  # build graph from scratch
            self.datetime = datetime.now().strftime(r"%y%m%d_%H%M")
            handles = self._buildGraph()
            for handle in handles:
                tf.add_to_collection("cnn_to_restore", handle)
            self.sesh.run(tf.global_variables_initializer())

        # unpack handles for tensor ops to feed or fetch for lower layers
        (self.xa_in, self.xb_in, self.xc_in, self.is_train, self.y_in, self.logits,
         self.net, self.w, self.pred, self.pred_loss,
         self.global_step, self.train_op, self.merged_summary, self.tumor) = handles

        # vars = []
        # for mm in ['Panoptes1/loss3/classifier/kernel:0', 'Panoptes1/loss3/classifier/bias:0',
        #           'Panoptes1/loss2/classifier_1/kernel:0', 'Panoptes1/loss2/classifier_1/bias:0',
        #           'Panoptes1/loss2/classifier_2/kernel:0', 'Panoptes1/loss2/classifier_2/bias:0',
        #           'Panoptes1/loss2/classifier/kernel:0', 'Panoptes1/loss2/classifier/bias:0']:
        #     vars.extend(tf.trainable_variables(scope=mm))

        if transfer:
            sample_weights = tf.gather_nd(self.weights,
                                          tf.stack([tf.argmax(self.tumor, axis=1),
                                                    tf.argmax(self.y_in, axis=1)], axis=1))

            self.pred_loss = tf.losses.softmax_cross_entropy(
                onehot_labels=self.y_in, logits=self.logits, weights=sample_weights)
            self.train_op = tf.train.AdamOptimizer(
                learning_rate=self.learning_rate, name='trans_Adam').minimize(
                loss=self.pred_loss, global_step=self.global_step)

            uninitialized_vars = []
            for var in tf.global_variables():
                try:
                    self.sesh.run(var)
                except tf.errors.FailedPreconditionError:
                    uninitialized_vars.append(var)
            self.sesh.run(tf.variables_initializer(uninitialized_vars))

        if save_graph_def:  # tensorboard
            try:
                os.mkdir(log_dir + '/training')
                os.mkdir(log_dir + '/validation')

            except FileExistsError:
                pass

            self.train_logger = tf.summary.FileWriter(log_dir + '/training', self.sesh.graph)
            self.valid_logger = tf.summary.FileWriter(log_dir + '/validation', self.sesh.graph)

    @property
    def step(self):
        return self.global_step.eval(session=self.sesh)

    # build graph; choose a structure defined in model
    def _buildGraph(self):
        # image input
        xa_in = tf.placeholder(tf.float32, name="xa")
        xa_in_reshape = tf.reshape(xa_in, [-1, self.input_dim[1], self.input_dim[2], 3])
        xb_in = tf.placeholder(tf.float32, name="xb")
        xb_in_reshape = tf.reshape(xb_in, [-1, self.input_dim[1], self.input_dim[2], 3])
        xc_in = tf.placeholder(tf.float32, name="xc")
        xc_in_reshape = tf.reshape(xc_in, [-1, self.input_dim[1], self.input_dim[2], 3])
        # dropout
        dropout = self.dropout
        # label input
        y_in = tf.placeholder(dtype=tf.float32, name="y")
        # tumor type input
        tumor = tf.placeholder(dtype=tf.float32, name="t")
        # train or test
        is_train = tf.placeholder_with_default(True, shape=[], name="is_train")
        classes = 20

        # other features input
        logits, nett, ww = Panoptes1.Panoptes1(xa_in_reshape, xb_in_reshape, xc_in_reshape,
                                                   num_cls=classes,
                                                   is_train=is_train,
                                                   dropout=dropout,
                                                   scope='Panoptes1')

        pred = tf.nn.softmax(logits, name="prediction")

        global_step = tf.Variable(0, trainable=False)

        sample_weights = tf.gather_nd(self.weights,
                                      tf.stack([tf.argmax(tumor, axis=1), tf.argmax(y_in, axis=1)], axis=1))

        pred_loss = tf.losses.softmax_cross_entropy(
            onehot_labels=y_in, logits=logits, weights=sample_weights)

        tf.summary.scalar("loss", pred_loss)

        tf.summary.tensor_summary("pred", pred)

        # optimizer
        train_op = tf.train.AdamOptimizer(
        learning_rate=self.learning_rate).minimize(
        loss=pred_loss, global_step=global_step)

        merged_summary = tf.summary.merge_all()

        return (xa_in, xb_in, xc_in, is_train,
                y_in, logits, nett, ww, pred, pred_loss,
                global_step, train_op, merged_summary, tumor)

    # inference using trained models
    def inference(self, X, dirr, bs, testset=None, pmd=None, train_status=False, realtest=False):
        now = datetime.now().isoformat()[11:]
        print("------- Testing begin: {} -------\n".format(now))
        rd = 0
        pdx = []
        yl = []
        if realtest:
            itr, file, ph = X.data(realtest=True, train=False)
            next_element = itr.get_next()
            with tf.Session() as sessa:
                sessa.run(itr.initializer, feed_dict={ph: file})
                while True:
                    try:
                        xa, xb, xc = sessa.run(next_element)
                        feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc,
                                     self.is_train: train_status}
                        fetches = [self.pred, self.net, self.w]
                        pred, net, w = self.sesh.run(fetches, feed_dict)
                        neta = net[:, :, :, :int(np.shape(net)[3] / 3)]
                        netb = net[:, :, :, int(np.shape(net)[3] / 3):2 * int(np.shape(net)[3] / 3)]
                        netc = net[:, :, :, 2 * int(np.shape(net)[3] / 3):]
                        wa = w[:int(np.shape(net)[3] / 3), :]
                        wb = w[int(np.shape(net)[3] / 3):2 * int(np.shape(net)[3] / 3), :]
                        wc = w[2 * int(np.shape(net)[3] / 3):, :]
                        cama = ac.CAM_R(neta, wa, pred, xa, dirr, 'Test_level1', bs, rd)
                        camb = ac.CAM_R(netb, wb, pred, xb, dirr, 'Test_level2', bs, rd)
                        camc = ac.CAM_R(netc, wc, pred, xc, dirr, 'Test_level3', bs, rd)
                        pred = np.concatenate((pred, cama, camb, camc), axis=1)
                        if rd == 0:
                            pdx = pred
                        else:
                            pdx = np.concatenate((pdx, pred), axis=0)
                        rd += 1
                    except tf.errors.OutOfRangeError:
                        ac.realout(pdx, dirr, 'Test', pmd)
                        break
        else:
            itr, file, ph = X.data(train=False)
            next_element = itr.get_next()
            with tf.Session() as sessa:
                sessa.run(itr.initializer, feed_dict={ph: file})
                while True:
                    try:
                        xa, xb, xc, y, _ = sessa.run(next_element)
                        feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc,
                                     self.is_train: train_status}
                        fetches = [self.pred, self.net, self.w]
                        pred, net, w = self.sesh.run(fetches, feed_dict)
                        # neta = net[:, :, :, :int(np.shape(net)[3] / 3)]
                        # netb = net[:, :, :, int(np.shape(net)[3] / 3):2 * int(np.shape(net)[3] / 3)]
                        # netc = net[:, :, :, 2 * int(np.shape(net)[3] / 3):]
                        # wa = w[:int(np.shape(net)[3] / 3), :]
                        # wb = w[int(np.shape(net)[3] / 3):2 * int(np.shape(net)[3] / 3), :]
                        # wc = w[2 * int(np.shape(net)[3] / 3):, :]
                        # ac.CAM(neta, wa, pred, xa, y, dirr, 'Test_level1', bs, pmd, rd)
                        # ac.CAM(netb, wb, pred, xb, y, dirr, 'Test_level2', bs, pmd, rd)
                        # ac.CAM(netc, wc, pred, xc, y, dirr, 'Test_level3', bs, pmd, rd)
                        net = np.mean(net, axis=(1, 2))
                        if rd == 0:
                            pdx = pred
                            yl = y
                            netl = net
                        else:
                            pdx = np.concatenate((pdx, pred), axis=0)
                            yl = np.concatenate((yl, y), axis=0)
                            netl = np.concatenate((netl, net), axis=0)
                        rd += 1
                    except tf.errors.OutOfRangeError:
                        ac.metrics(pdx, yl, dirr, 'Test', pmd, testset)
                        ac.tSNE_prep(flatnet=netl, ori_test=testset, y=yl, pred=pdx, path=dirr, pmd=pmd)
                        break
        now = datetime.now().isoformat()[11:]
        print("------- Testing end: {} -------\n".format(now), flush=True)

    # training
    def train(self, X, VAX, ct, bs, dirr, pmd, max_iter=np.inf, save=True, outdir="./out"):
        start_time = time.time()
        svs = 0
        if save:
            saver = tf.train.Saver(tf.global_variables(), max_to_keep=None)

        err_train = 0
        now = datetime.now().isoformat()[11:]
        print("------- Training begin: {} -------\n".format(now))
        itr, file, ph = X.data()
        next_element = itr.get_next()

        vaitr, vafile, vaph = VAX.data(train=False)
        vanext_element = vaitr.get_next()

        init_i = self.step

        with tf.Session() as sessa:
            sessa.run(itr.initializer, feed_dict={ph: file})
            sessa.run(vaitr.initializer, feed_dict={vaph: vafile})
            train_loss = []
            validation_loss = []
            valid_loss = 0
            trloss_plt = []
            valoss_plt = []

            try:
                while True:
                    xa, xb, xc, y, tum = sessa.run(next_element)
                    feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc, self.y_in: y, self.tumor: tum}

                    fetches = [self.merged_summary, self.logits, self.pred,
                               self.pred_loss, self.global_step, self.train_op]

                    summary, logits, pred, loss, i, _ = self.sesh.run(fetches, feed_dict)
                    if self.transfer:
                        i = i - init_i
                    self.train_logger.add_summary(summary, i)
                    err_train += loss

                    train_loss.append(loss)

                    try:
                        mintrain = min(train_loss)
                    except ValueError:
                        mintrain = 0

                    if loss <= mintrain and i > 19999:
                        print("round {} --> loss: ".format(i), loss)
                        temp_valid = []
                        for iii in range(20):
                            xa, xb, xc, y, tum = sessa.run(vanext_element)
                            feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc, self.y_in: y, self.tumor: tum,
                                         self.is_train: False}
                            fetches = [self.pred_loss, self.merged_summary]
                            valid_loss, valid_summary = self.sesh.run(fetches, feed_dict)
                            self.valid_logger.add_summary(valid_summary, i)
                            temp_valid.append(valid_loss)

                        tempminvalid = np.mean(temp_valid)
                        try:
                            minvalid = min(validation_loss)
                        except ValueError:
                            minvalid = 0

                        if save and tempminvalid <= minvalid:
                            validation_loss.append(tempminvalid)
                            print("round {} --> validation loss: ".format(i), tempminvalid)
                            print("New Min loss model found!", flush=True)
                            outfile = os.path.join(os.path.abspath(outdir),
                                                   "{}".format("_".join(['dropout', str(self.dropout)])))
                            saver.save(self.sesh, outfile, global_step=None)
                            svs = i

                    elif i % 1000 == 0:
                        print("round {} --> loss: ".format(i), loss)
                        temp_valid = []
                        for iii in range(100):
                            xa, xb, xc, y, tum = sessa.run(vanext_element)
                            feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc, self.y_in: y, self.tumor: tum,
                                         self.is_train: False}
                            fetches = [self.pred_loss, self.merged_summary]
                            valid_loss, valid_summary = self.sesh.run(fetches, feed_dict)
                            self.valid_logger.add_summary(valid_summary, i)
                            temp_valid.append(valid_loss)
                        tempminvalid = np.mean(temp_valid)
                        try:
                            minvalid = min(validation_loss)
                        except ValueError:
                            minvalid = 0
                        validation_loss.append(tempminvalid)
                        print("round {} --> Step Average validation loss: ".format(i), tempminvalid, flush=True)

                        # Loss figures
                        if len(train_loss) >= 1000:
                            trloss_plt.append(np.mean(train_loss[-1000:]))
                        else:
                            trloss_plt.append(np.mean(train_loss[-len(train_loss):]))
                        valoss_plt.append(tempminvalid)
                        plt.plot(trloss_plt, color='blue')
                        plt.plot(valoss_plt, color='orange')
                        plt.xlabel('thousand iterations')
                        plt.ylabel('loss')
                        plt.title('Average Train & Validation Loss (every 1000 itrs)')
                        plt.legend(['Train', 'Validation'], loc='upper right')
                        plt.savefig(outdir + '/loss.png')

                        if save and tempminvalid <= minvalid:
                            print("New Min loss model found!")
                            print("round {} --> loss: ".format(i), loss)
                            outfile = os.path.join(os.path.abspath(outdir),
                                                   "{}".format("_".join(['dropout', str(self.dropout)])))
                            saver.save(self.sesh, outfile, global_step=None)
                            svs = i

                        if i > 69999:
                            valid_mean_loss = np.mean(validation_loss[-10:-1])
                            print('Mean validation loss: {}'.format(valid_mean_loss))
                            if valid_loss > valid_mean_loss:
                                print("Early stopped! No improvement for at least 10000 iterations", flush=True)
                                break
                            else:
                                print("Passed early stopping evaluation. Continue training!")

                    elif i >= max_iter-2:
                        print("final avg loss (@ step {} = epoch {}): {}".format(
                            i + 1, np.around(i / ct * bs), err_train / i))

                        now = datetime.now().isoformat()[11:]
                        print("------- Training end: {} -------\n".format(now))

                        now = datetime.now().isoformat()[11:]
                        print("------- Final Validation begin: {} -------\n".format(now))
                        xa, xb, xc, y, tum = sessa.run(vanext_element)
                        feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc, self.y_in: y, self.tumor: tum,
                                     self.is_train: False}
                        fetches = [self.pred_loss, self.merged_summary]
                        valid_loss, valid_summary = self.sesh.run(fetches, feed_dict)

                        self.valid_logger.add_summary(valid_summary, i)
                        print("round {} --> Final Last validation loss: ".format(i), valid_loss)
                        now = datetime.now().isoformat()[11:]
                        print("------- Final Validation end: {} -------\n".format(now), flush=True)
                        try:
                            self.train_logger.flush()
                            self.train_logger.close()
                            self.valid_logger.flush()
                            self.valid_logger.close()

                        except AttributeError:  # not logging
                            print('Not logging')
                        break
            except tf.errors.OutOfRangeError:
                pass

            try:
                print("final avg loss (@ step {} = epoch {}): {}".format(
                    i + 1, np.around(i / ct * bs), err_train / i))

                now = datetime.now().isoformat()[11:]
                print("------- Training end: {} -------\n".format(now))

                if svs < 1000 and save:
                        print("Save the last model as the best model.")
                        outfile = os.path.join(os.path.abspath(outdir),
                                               "{}".format("_".join(['dropout', str(self.dropout)])))
                        saver.save(self.sesh, outfile, global_step=None)

                now = datetime.now().isoformat()[11:]
                print("------- Validation begin: {} -------\n".format(now))

                xa, xb, xc, y, tum = sessa.run(vanext_element)
                feed_dict = {self.xa_in: xa, self.xb_in: xb, self.xc_in: xc, self.y_in: y, self.tumor: tum,
                             self.is_train: False}
                fetches = [self.pred_loss, self.merged_summary, self.pred, self.net, self.w]
                valid_loss, valid_summary, pred, net, w = self.sesh.run(fetches, feed_dict)

                self.valid_logger.add_summary(valid_summary, i)
                print("round {} --> Last validation loss: ".format(i), valid_loss)
                # for i in range(3):
                #     neta = net[:, :, :, :int(np.shape(net)[3] / 3)]
                #     netb = net[:, :, :, int(np.shape(net)[3] / 3):2 * int(np.shape(net)[3] / 3)]
                #     netc = net[:, :, :, 2 * int(np.shape(net)[3] / 3):]
                #     wa = w[:int(np.shape(net)[3] / 3), :]
                #     wb = w[int(np.shape(net)[3] / 3):2 * int(np.shape(net)[3] / 3), :]
                #
                #     wc = w[2 * int(np.shape(net)[3] / 3):, :]
                #
                #     ac.CAM(neta, wa, pred, xa, y, dirr, 'Validation_level1', bs, pmd)
                #     ac.CAM(netb, wb, pred, xb, y, dirr, 'Validation_level2', bs, pmd)
                #     ac.CAM(netc, wc, pred, xc, y, dirr, 'Validation_level3', bs, pmd)
                ac.metrics(pred, y, dirr, 'Validation', pmd)
                now = datetime.now().isoformat()[11:]
                print("------- Validation end: {} -------\n".format(now), flush=True)

                try:
                    self.train_logger.flush()
                    self.train_logger.close()
                    self.valid_logger.flush()
                    self.valid_logger.close()

                except AttributeError:  # not logging
                    print('Not logging')

            except tf.errors.OutOfRangeError:
                print("final avg loss (@ step {} = epoch {}): {}".format(
                    i + 1, np.around(i / ct * bs), err_train / i))

                now = datetime.now().isoformat()[11:]
                print("------- Training end: {} -------\n".format(now))
                print('No more validation needed!')

        print("--- %s seconds ---" % (time.time() - start_time))


import pyautogui, win32api, win32con, ctypes, autoit
from PIL import ImageOps, Image, ImageGrab
from numpy import *
import os
import time
import cv2

leftCornerx = 7
leftCornery = 38
x1 = leftCornerx
y1 = leftCornery
x2 = 1390
y2 = 1000

# 967 55
# 1120 62

def getScreen(x1, y1, x2, y2):
    box = (x1, y1, x2, y2)
    screen = ImageGrab.grab(box)
    img =  array(screen.getdata(),dtype=uint8).reshape((screen.size[1],screen.size[0],3))
    # img = cv2.imread('snap__1426174983.png')
    # cv2.imshow('image',img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    return img

def findTarget(img):
    template_tg = cv2.imread('template_target2.png', 0)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret,th1 = cv2.threshold(gray,253,255,cv2.THRESH_TOZERO_INV)
    ret,th3 = cv2.threshold(th1,251,255,cv2.THRESH_BINARY)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 15))
    closed = cv2.morphologyEx(th3, cv2.MORPH_CLOSE, kernel)
    closed = cv2.erode(closed, None, iterations = 3)
    closed = cv2.dilate(closed, None, iterations = 2)
    (cnts, hierarchy) = cv2.findContours(closed,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

    approxes = []
    hulls = []
    for cnt in cnts:
        approxes.append(cv2.approxPolyDP(cnt,0.01*cv2.arcLength(cnt,True),True))
        hulls.append(cv2.convexHull(cnt))
        left = list(cnt[cnt[:,:,0].argmin()][0])
        right = list(cnt[cnt[:,:,0].argmax()][0])
        print 'left x' + str(left[0])+ 'y '+ str(left[1])
        print 'right x' + str(right[0])+ 'y '+ str(right[1])
        center = round((right[0]+left[0])/2)
        center = int(center)
        moveMouse(center-10,left[1]+70)
        if (findFromTargeted(template_tg, left, right)):
            autoit.mouse_click('left', center-10, left[1]+70)
            return True
        pyautogui.moveTo(center,left[1]+70)
        moveMouse(center,left[1]+70)
        if (findFromTargeted(template_tg, left, right)):
            autoit.mouse_click('left', center+10, left[1]+70)
            return True
    # autoit.mouse_click('middle')

def findFromTargeted(template, left, right):
    # print template.shape
    roi = getScreen(left[0]-25+leftCornerx, left[1]-8+leftCornery, right[0]+leftCornerx, right[1]+8+leftCornery)
    cv2.imwrite('roi' + str(int(time.time())) + '.png',roi)
    roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    # roi = img[left[1]-8:right[1]+8, left[0]-25:right[0]];
    ret,th1 = cv2.threshold(roi,200,255,cv2.THRESH_TOZERO_INV)
    ret,th2 = cv2.threshold(th1,100,255,cv2.THRESH_BINARY)
    ret,tp1 = cv2.threshold(template,200,255,cv2.THRESH_TOZERO_INV)
    ret,tp2 = cv2.threshold(tp1,100,255,cv2.THRESH_BINARY)
    wth, hth = th2.shape
    wtp, htp = tp2.shape
    if wth > wtp and hth > htp:
        res = cv2.matchTemplate(th2, tp2, cv2.TM_CCORR_NORMED)
        if (res.any()):
            # cv2.imwrite('th2' + str(int(time.time())) + '.png',th2)
            # cv2.imwrite('tp2' + str(int(time.time())) + '.png',tp2)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
            print 'max_val %.2f'%max_val
            if (max_val > 0.9):
                return True
            else:
                return False
    return False
def grabHP():
    hp = getScreen(leftCornerx + 958,leftCornery + 16,leftCornerx + 1111,leftCornery + 25)

    # img = cv2.imread('snap__1426174990.png')
    # gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # roi = gray[16:25, 958:1111];
    return hp

def findHP(img):
    statuses = {'none': -1, 'dead' : 0,  'lhalf' : 1, 'mhalf' : 2, 'full' : 3}
    # hpcolor = [231, 73, 132]
    hpcolor = [107, 101, 107]
    hp = grabHP()
    gray = cv2.cvtColor(hp, cv2.COLOR_BGR2GRAY)
    ret,th1 = cv2.threshold(gray,120,255,cv2.THRESH_TOZERO_INV)
    ret,th1 = cv2.threshold(th1,100,255,cv2.THRESH_TOZERO)
    # cv2.imwrite('th1' + str(int(time.time())) + '.png',th1)

    (cnts, hierarchy) = cv2.findContours(th1,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    if (len(cnts) == 0):
        return statuses['dead']
    left = list(cnts[0][cnts[0][:,:,0].argmin()][0])
    px = hp[left[1], left[0]];
    # print px
    if (hpcolor != px).any():
        return statuses['none']

    leftx = list(cnts[0][cnts[0][:,:,0].argmin()][0])[0]
    rightx = list(cnts[0][cnts[0][:,:,0].argmax()][0])[0]
    diff = rightx - leftx
    if diff > 140:
        return statuses['full']
    if diff >= 75:
        return statuses['mhalf']
    if diff < 75:
        return statuses['lhalf']
    return statuses['dead']

def moveMouse(x,y):
    autoit.mouse_move(x,y)

def main():
    title = "[CLASS:l2UnrealWWindowsViewportWindow]"
    autoit.win_wait(title, 5)
    counter = 0
    while True:
        img = getScreen(leftCornerx,leftCornery,x2,y2)
        res = findHP(img);
        print res
        if res > 0:
            autoit.control_send(title, '', '{F1}', 0)
            time.sleep(2)
        else:
            if counter == 1:
                autoit.control_send(title, '', '{F3}', 0)

            if counter > 2:
                findTarget(img)
                counter = 0
            counter += 1
        print 'counter ' + str(counter)
    pass

if __name__ == '__main__':
    main()


#!/usr/bin/env python

''' Script to precompute image features using a Caffe ResNet CNN, using 36 discretized views
    at each viewpoint in 30 degree increments, and the provided camera WIDTH, HEIGHT
    and VFOV parameters. '''

import argparse
import numpy as np
import cv2
import json
import math
import base64
import csv
import sys

csv.field_size_limit(sys.maxsize)

# CLIP Support
import torch
import clip
from PIL import Image

# Caffe and MatterSim need to be on the Python path
sys.path.insert(0, 'build')
import MatterSim

sys.path.insert(0, "scripts")
from timer import Timer

TSV_FIELDNAMES = ['scanId', 'viewpointId', 'image_w','image_h', 'vfov', 'features']
VIEWPOINT_SIZE = 36 # Number of discretized views from one viewpoint
BATCH_SIZE = 4  # Some fraction of viewpoint size - batch size 4 equals 11GB memory
GPU_ID = 0

LABEL = False

parser = argparse.ArgumentParser()
parser.add_argument('--arch', default='resnet', help='architecture')
args = parser.parse_args()

# Labels Prob
if args.arch == "resnet":
    FEATURE_SIZE = 1024
    MODEL = "RN50"
    OUTFILE = 'img_features/CLIP-ResNet-50-views.tsv'
elif args.arch == "vit":
    FEATURE_SIZE = 512
    MODEL = "ViT-B/32"
    OUTFILE = 'img_features/CLIP-ViT-B-32-views.tsv'
else:
    assert False

GRAPHS = 'connectivity/'

# Simulator image parameters
WIDTH=640
HEIGHT=480
VFOV=60

def load_viewpointids():
    viewpointIds = []
    with open(GRAPHS+'scans.txt') as f:
        scans = [scan.strip() for scan in f.readlines()]
        for scan in scans:
            with open(GRAPHS+scan+'_connectivity.json')  as j:
                data = json.load(j)
                for item in data:
                    if item['included']:
                        viewpointIds.append((scan, item['image_id']))
    print('Loaded %d viewpoints' % len(viewpointIds))
    return viewpointIds

def build_tsv():
    # Set up the simulator
    sim = MatterSim.Simulator()
    sim.setCameraResolution(WIDTH, HEIGHT)
    sim.setCameraVFOV(math.radians(VFOV))
    sim.setDiscretizedViewingAngles(True)
    #sim.setRenderingEnabled(False)
    sim.init()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(device)
    model, preprocess = clip.load(MODEL, device=device)
    print(model)

    count = 0
    t_render = Timer()
    t_net = Timer()
    max_prob = np.zeros(FEATURE_SIZE, dtype=np.float32)
    features_list = []
    with open(OUTFILE, 'w') as tsvfile:
        writer = csv.DictWriter(tsvfile, delimiter = '\t', fieldnames = TSV_FIELDNAMES)

        # Loop all the viewpoints in the simulator
        viewpointIds = load_viewpointids()
        for scanId,viewpointId in viewpointIds:
            t_render.tic()

            # Loop all discretized views from this location
            blobs = []
            for ix in range(VIEWPOINT_SIZE):
                if ix == 0:
                    sim.newEpisode(scanId, viewpointId, 0, math.radians(-30))
                elif ix % 12 == 0:
                    sim.makeAction(0, 1.0, 1.0)
                else:
                    sim.makeAction(0, 1.0, 0)

                state = sim.getState()
                assert state.viewIndex == ix

                # Transform and save generated image
                blobs.append(Image.fromarray(state.rgb))

            blobs = [
                preprocess(blob).unsqueeze(0)
                for blob in blobs
            ]
            blobs = torch.cat(blobs, 0)
            #print(blobs.shape)
            blobs = blobs.to(device)

            t_render.toc()
            t_net.tic()
            # Run as many forward passes as necessary

            features = model.encode_image(blobs).float()
            #print(features[0])
            #print(features.shape)

            if LABEL:
                for k in range(VIEWPOINT_SIZE):
                    feature = features[k]
                    max_prob = np.maximum(max_prob, feature)
                features_list.append(features.detach().cpu().numpy())
            else:
                features = features.detach().cpu().numpy()
                #print(features.shape, features.dtype)
                writer.writerow({
                    'scanId': scanId,
                    'viewpointId': viewpointId,
                    'image_w': WIDTH,
                    'image_h': HEIGHT,
                    'vfov' : VFOV,
                    'features': base64.b64encode(features).decode(),
                })

            count += 1
            t_net.toc()
            if count % 100 == 0:
                print('Processed %d / %d viewpoints, %.1fs avg render time, %.1fs avg net time, projected %.1f hours' %\
                  (count, len(viewpointIds), t_render.average_time, t_net.average_time,
                  (t_render.average_time+t_net.average_time)*len(viewpointIds)/3600))

        if LABEL:
            for i, (scanId, viewpointId) in enumerate(viewpointIds):
                if LABEL:
                    features = features_list[i] / max_prob
                else:
                    features = features_list[i]
                writer.writerow({
                    'scanId': scanId,
                    'viewpointId': viewpointId,
                    'image_w': WIDTH,
                    'image_h': HEIGHT,
                    'vfov' : VFOV,
                    'features': base64.b64encode(features).decode(),
                })

def read_tsv(infile):
    # Verify we can read a tsv
    in_data = []
    with open(infile, "r") as tsv_in_file:
        reader = csv.DictReader(tsv_in_file, delimiter='\t', fieldnames = TSV_FIELDNAMES)
        for item in reader:
            item['image_h'] = int(item['image_h'])
            item['image_w'] = int(item['image_w'])
            item['vfov'] = int(item['vfov'])
            item['features'] = np.frombuffer(base64.b64decode(item['features']),
                    dtype=np.float32).reshape((VIEWPOINT_SIZE, FEATURE_SIZE))
            in_data.append(item)
    return in_data

if __name__ == "__main__":

    build_tsv()
    data = read_tsv(OUTFILE)
    print('Completed %d viewpoints' % len(data))


from __future__ import print_function
import argparse

import os
import torch
import cv2
from model import *
import torchvision.transforms as transforms
from collections import OrderedDict
import numpy as np
from os.path import join
import time
from network import encoder4, decoder4
import numpy
from dataset import is_image_file, rescale_img
from image_utils import *
from PIL import Image, ImageOps
from os import listdir
import torch.utils.data as utils
import os

# Training settings
parser = argparse.ArgumentParser(description='PyTorch Super Res Example')
parser.add_argument('--testBatchSize', type=int, default=8, help='testing batch size')
parser.add_argument('--up_factor', type=int, default=4, help="super resolution upscale factor")
parser.add_argument('--gpu_mode', type=bool, default=True)
parser.add_argument('--chop_forward', type=bool, default=True)
parser.add_argument('--use_img_self', action='store_true', help='using LR image itself or not')
parser.add_argument('--use_ref', action='store_true', help='using external reference images or not')
parser.add_argument('--num_sample', type=int, default=10, help='number of SR images')
parser.add_argument('--threads', type=int, default=6, help='number of threads for data loader to use')
parser.add_argument('--input_dataset', type=str, default='input')
parser.add_argument('--output_dataset', type=str, default='result')
parser.add_argument('--model_type', type=str, default='ConVAE')
parser.add_argument('--model', default='models/ConVAE_4x.pth', help='sr pretrained base model')
parser.add_argument("--encoder_dir", default='models/vgg_r41.pth', help='pre-trained encoder path')
parser.add_argument("--decoder_dir", default='models/dec_r41.pth', help='pre-trained encoder path')

opt = parser.parse_args()

print(opt)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print('===> Building model ', opt.model_type)

model = VAE_v3_4x(up_factor=opt.up_factor)

enc = encoder4()
dec = decoder4()

if os.path.exists(opt.encoder_dir):
    enc.load_state_dict(torch.load(opt.encoder_dir))
    print('encoder model is loaded!')

if os.path.exists(opt.decoder_dir):
    dec.load_state_dict(torch.load(opt.decoder_dir))
    print('decoder model is loaded!')

for param in enc.parameters():
    param.requires_grad = False

for param in dec.parameters():
    param.requires_grad = False

model = model.to(device)
enc = enc.to(device)
dec = dec.to(device)

print('===> Loading datasets')

def eval():
    model.eval()
    enc.eval()
    dec.eval()

    if os.path.exists(opt.model):
        model.load_state_dict(torch.load(opt.model, map_location=lambda storage, loc: storage))
        print(opt.model)

    Ref_filename = os.path.join(opt.input_dataset, 'Ref')
    LR_filename = os.path.join(opt.input_dataset, 'LR')
    SR_filename = opt.output_dataset

    lr_image = [join(LR_filename, x) for x in listdir(LR_filename) if is_image_file(x)]
    lr_image = sorted(lr_image)
    ref_image = [join(Ref_filename, x) for x in listdir(Ref_filename) if is_image_file(x)]
    ref_image = sorted(ref_image)

    for i in range(len(lr_image)):

        LR = Image.open(lr_image[i]).convert('RGB')
        LR = modcrop(LR, opt.up_factor)

        if len(ref_image) != 0 and opt.use_ref:
            print("using ref images for SR")
            for j in range(len(ref_image)):
                Ref = Image.open(ref_image[j]).convert('RGB')

                with torch.no_grad():
                    prediction = chop_forward(Ref, LR)

                prediction = prediction.data[0].cpu().permute(1, 2, 0)

                prediction = prediction * 255.0
                prediction = prediction.clamp(0, 255)
                lr_name = lr_image[i][-8:-4]
                output_name = SR_filename + '/' + lr_name.zfill(6) + '_sample' + str(j).zfill(5) + '.png'
                Image.fromarray(np.uint8(prediction)).save(output_name)

        else:
            if opt.use_img_self:
                print("using LR images itself for SR")
                Ref = LR.resize((256, 256))

                with torch.no_grad():
                    prediction = chop_forward(Ref, LR)

                prediction = prediction.data[0].cpu().permute(1, 2, 0)

                prediction = prediction * 255.0
                prediction = prediction.clamp(0, 255)
                lr_name = lr_image[i][-8:-4]
                output_name = SR_filename + '/' + lr_name.zfill(6) + '_sample0.png'
                Image.fromarray(np.uint8(prediction)).save(output_name)

            else:
                print("using random noise for SR")
                for j in range(opt.num_sample):
                    a = np.random.rand(256, 256, 3)
                    Ref = Image.fromarray(np.uint8(a * 128))

                    with torch.no_grad():
                        prediction = chop_forward(Ref, LR)

                    prediction = prediction.data[0].cpu().permute(1, 2, 0)

                    prediction = prediction * 255.0
                    prediction = prediction.clamp(0, 255)
                    lr_name = lr_image[i][-8:-4]
                    output_name = SR_filename + '/' + lr_name.zfill(6) + '_sample' + str(j).zfill(5) + '.png'
                    print("random SR: {}".format(j))
                    Image.fromarray(np.uint8(prediction)).save(output_name)

            # pre_LR = F.interpolate(prediction, scale_factor=1 / opt.up_factor, mode='bicubic')
            #
            # prediction = prediction.data[0].cpu().permute(1, 2, 0)
            # pre_LR = pre_LR.data[0].cpu().permute(1, 2, 0)
            #
            # prediction = prediction * 255.0
            # prediction = prediction.clamp(0, 255)
            # lr_name = lr_image[i][-8:-4]
            # output_name = SR_filename + '/' + lr_name.zfill(6) + '_sample0.png'
            # Image.fromarray(np.uint8(prediction)).save(output_name)
            #
            # pre_LR = pre_LR * 255.0
            # pre_LR = pre_LR.clamp(0, 255)
            #
            # LR = np.array(LR).astype(np.float32)
            # LR_Y = rgb2ycbcr(LR)
            # pre_LR = np.array(pre_LR).astype(np.float32)
            # pre_LR_Y = rgb2ycbcr(pre_LR)
            #
            # psnr_predicted = PSNR(pre_LR_Y, LR_Y, shave_border=1)
            # ssim_predicted = SSIM(pre_LR_Y, LR_Y, shave_border=1)

transform = transforms.Compose([
    transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
    ]
)

style_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
    ]
)

def chop_forward(ref, img):

    img = transform(img).unsqueeze(0)
    ref = style_transform(ref).unsqueeze(0)

    testset = utils.TensorDataset(ref, img)
    test_dataloader = utils.DataLoader(testset, num_workers=opt.threads,
                                       drop_last=False, batch_size=opt.testBatchSize, shuffle=False)

    for iteration, batch in enumerate(test_dataloader, 1):
        ref, input = batch[0].to(device), batch[1].to(device)

        LR_feat = enc(F.interpolate(input, scale_factor=opt.up_factor, mode='bicubic'))
        ref_feat = enc(ref)
        SR, _ = model(input, LR_feat['r41'], ref_feat['r41'])

    return SR

##Eval Start!!!!
eval()


import torch
import torch.nn as nn

from fakeGen.Discriminator import Discriminator, RNNclaissfier
from fakeGen.seq2seq import Seq2seq
from fakeGen.Generator import Generator
from fakeGen.evaluate import f1_score, accuracy_score, tensor2list

from SeqModel.trainer import SupervisedTrainer
from SeqModel.models import EncoderRNN, DecoderRNN, ItemEncoder, ContextDecoderRNN
from SeqModel.loss import VAELoss, Perplexity
from SeqModel.optim import Optimizer
from SeqModel.dataset import SourceField, TargetField
from SeqModel.evaluator import Predictor
from SeqModel.util.checkpoint import Checkpoint
import torchtext
import sys
import torch.optim as optim
import argparse
import itertools
import numpy as np

def helper_fusion(pos_data, neg_data):
    pos_batch = pos_data.shape[0]
    neg_batch = neg_data.shape[0]
    label_pos = torch.ones((pos_batch, 1), device=pos_data.device, dtype=torch.long)
    label_neg = torch.ones((neg_batch, 1), device=pos_data.device, dtype=torch.long)
    all_data = torch.cat((pos_data, neg_data), dim=0)
    label = torch.cat((label_pos, label_neg), dim=0)

    perm = torch.randperm(label.size()[0], device=pos_data.device)
    all_data = all_data[perm].contiguous()
    label = label[perm].contiguous()

    return all_data, label

def train_LM(opt, loss, seq2seq, fake_data, optimizer=None):
    # train

    train_lm_data, test_lm_data = fake_data.split(split_ratio=0.2)
    t = SupervisedTrainer(loss=loss, batch_size=opt.batch_size,
                             checkpoint_every=opt.check_point, device=opt.device,
                             print_every=opt.print_every, expt_dir=opt.expt_dir, predic_rate=True)

    seq2seq = t.train(seq2seq, train_lm_data, dev_data=test_lm_data,
                      num_epochs=opt.lm_epochs,
                      optimizer=optimizer,
                      teacher_forcing_ratio=opt.teach_force_ratio,
                      resume=opt.resume)
    return seq2seq

def pre_train_deceptive(rnn_claissfier, classifier_opt, data, opt):
    rnn_claissfier.train()
    for epochs in range(opt.pre_clf_epochs):
        total_loss = 0
        data_iter = data.__iter__()
        for batch in data_iter:
            feature, input_length = getattr(batch, 'src')
            label = getattr(batch, 'label')
            classifier_opt.zero_grad()
            loss, _ = rnn_claissfier(feature, label, input_length)
            loss.backward()
            classifier_opt.step()

            total_loss += loss.item()
        print('[INFO] ---PRE-TRAIN--- clf loss is {}'.format(total_loss))
    return rnn_claissfier

def train_discriminator(discriminator, dis_opt, seq2seq, gen, fake_data, opt):
    for epoch in range(opt.dis_epoch):
        print('epoch %d : ' % (epoch + 1), end='')
        sys.stdout.flush()
        total_loss = 0
        # clf the simulate text and fake text
        fake_iter = fake_data.__iter__()
        true_list = []
        pre_list = []
        for batch in fake_iter:
            feature, input_lengths = getattr(batch, 'src')
            _, hidden = seq2seq.encoder_seq(feature, input_lengths=input_lengths)
            if len(hidden) == 2:
                hidden = hidden[0]
            hidden = hidden.permute(1, 0, 2).contiguous()
            hidden = hidden.view(hidden.shape[0], -1)
            shape = torch.Size((opt.batch_size, opt.z_size))
            if next(discriminator.parameters()).is_cuda:
                z = torch.cuda.FloatTensor(shape)
            else:
                z = torch.FloatTensor(shape)
            torch.randn(shape, out=z)
            # classify the hidden state
            sim_data = gen(z)
            all_data, label = helper_fusion(hidden, sim_data)
            dis_opt.zero_grad()
            loss, out = discriminator.batchBCELoss(all_data, label)
            loss.backward()
            dis_opt.step()
            true_list.append(tensor2list(out['y_true']))
            pre_list.append(tensor2list(out['y_pre']))
            total_loss += loss.item()

        y_true = list(itertools.chain.from_iterable(true_list))
        y_pre = list(itertools.chain.from_iterable(pre_list))

        f1 = f1_score(y_true, y_pre)
        acc = accuracy_score(y_true, y_pre)

        print("[INFO] ---TRAIN---- discriminator loss {}, acc {}, f1 {}".format(total_loss, f1, acc))
        # train the classifier

    # return discriminator
    #TODO: print total loss

# pad one
def pad(tensor, length):
    return torch.cat([tensor, tensor.new_ones(length - tensor.size(0), tensor.size()[1])])

def train_classifier(opt, real_data, gen, seq2seq, rnn_classifier, rnn_opt):
    rnn_classifier.train()
    seq2seq.train()
    gen.train()
    real_iter = real_data.__iter__()
    total_loss = 0
    true_list = []
    pre_list = []
    for batch in real_iter:
        feature, lengths = getattr(batch, 'src')
        label = getattr(batch, 'label')
        # shape = torch.Size((opt.batch_size, opt.z_size))
        shape = torch.Size((10, opt.z_size))
        if next(rnn_classifier.parameters()).is_cuda:
            z = torch.cuda.FloatTensor(shape)
        else:
            z = torch.FloatTensor(shape)

            # sim_seq: distribution of words
        torch.randn(shape, out=z)
        sim_hidden = gen(z)
        sim_hidden = sim_hidden.view(z.shape[0], opt.n_layers * (2 if opt.bidirectional else 1), opt.hidden_size)
        sim_hidden = sim_hidden.permute(1, 0, 2).contiguous()
        sim_seq_list, _, _ = seq2seq.decoder_hidden(sim_hidden, teacher_forcing_ratio=0)

        sim_seq_length = [seq.shape[0] for seq in sim_seq_list]
        length_sort = np.sort(sim_seq_length)
        index_sort = np.argsort(sim_seq_length)
        max_length = np.max(sim_seq_length)
        sim_seq_reorder = [pad(sim_seq_list[index], max_length).unsqueeze(0) for index in index_sort]
        sim_seq = torch.cat(sim_seq_reorder)
        real_label = torch.ones_like(label)
        sim_label = label.new_zeros(sim_seq.shape[0])
        rnn_opt.zero_grad()
        loss, out = rnn_classifier(feature, real_label, lengths, sim_seq, sim_label, length_sort)
        loss.backward()
        rnn_opt.step()
        total_loss += loss.item()
        true_list.append(tensor2list(out['y_true']))
        pre_list.append(tensor2list(out['y_pre']))

    y_pre = [int(i[0]) for i in list(itertools.chain.from_iterable(pre_list))]
    y_true = [i for i in list(itertools.chain.from_iterable(true_list))]

    f1 = f1_score(y_true, y_pre)
    acc = accuracy_score(y_true, y_pre)

    print("[INFO] ---TRAINING--- clf loss {}, f1 {}, acc {}".format(total_loss, f1, acc))

def clf_test(test_clf_data, rnn_classifier):
    rnn_classifier.eval()
    pre_list = []
    true_list = []
    data_iter = test_clf_data.__iter__()
    for batch in data_iter:
        feature, lengths = getattr(batch, 'src')
        label = getattr(batch, 'label')
        # feature and label will be shuffle in clf
        loss, out = rnn_classifier(feature, label, lengths)
        y_pre = tensor2list(out['y_pre'])
        y_true = tensor2list(out['y_true'])
        pre_list.append(y_pre)
        true_list.append(y_true)
    y_pre = [int(i[0]) for i in list(itertools.chain.from_iterable(pre_list))]
    y_true = [ i - 1 for i in list(itertools.chain.from_iterable(true_list))]
    f1 = f1_score(y_true, y_pre)
    acc = accuracy_score(y_true, y_pre)
    print("[INFO] ---TEST--- acc is {}, f1 is {}".format(acc, f1))
    #
        # print the loss and F1 score

def train_gen(gen, gen_opt, dis_gen, opt):
    for epoch in range(opt.gen_epoch):
        for _ in range(opt.batch_count):
            shape = torch.Size((opt.batch_size, opt.z_size))
            if next(gen.parameters()).is_cuda:
                z = torch.cuda.FloatTensor(shape)
            else:
                z = torch.FloatTensor(shape)

            torch.randn(shape, out=z)
            # generate fake review
            sim_data = gen(z)
            gen_opt.zero_grad()
            # fool the discriminator
            # seqGAN will use the prediction logit as reward
            loss, _ = dis_gen.batchBCELoss(sim_data, torch.ones((sim_data.shape[0], 1), device=opt.device, dtype=torch.long))
            loss.backward()
            gen_opt.step()

def prepare_data(opt):
    # seq-label real dataset
    # seq-label all dataset
    # seq-label fake dataset
    tgt = TargetField()
    src = SourceField()
    label = torchtext.data.Field(sequential=False)
    # fake_data_lm = torchtext.data.TabularDataset(
    #     path=opt.fake_data_path, format='csv',
    #     fields=[('src', src), ('label', label), ('tgt', tgt)]
    # )

    # real_data_clf = torchtext.data.TabularDataset(
    #     path=opt.real_data_path, format='csv',
    #     fields=[('src', src), ('label', label)]
    # )
    #
    # train_clf = torchtext.data.TabularDataset(
    #     path=opt.train_data_path, format='csv',
    #     fields=[('src', src), ('label', label)]
    # )
    #
    #
    # test_clf = torchtext.data.TabularDataset(
    #     path=opt.test_data_path, format='csv',
    #     fields=[('src', src), ('label', label)]
    # )
    all_data = torchtext.data.TabularDataset(
        path=opt.text_path, format='csv',
        fields=[('src', src), ('tgt', tgt)]
    )

    src.build_vocab(all_data.src, max_size=opt.max_word)
    tgt.build_vocab(all_data.src, max_size=opt.max_word)
    # label.build_vocab(train_clf)
    input_vocab = src.vocab

    output_vocab = tgt.vocab
    return all_data, tgt

    # # the data is so large?
    # test_clf = torchtext.data.BucketIterator(
    #     dataset=test_clf, batch_size=opt.batch_size,
    #     sort=False, sort_within_batch=True,
    #     sort_key=lambda x: len(x.src),
    #     device=opt.device, repeat=False)
    #
    # train_clf = torchtext.data.BucketIterator(
    #     dataset=train_clf, batch_size=opt.batch_size,
    #     sort=False, sort_within_batch=True,
    #     sort_key=lambda x: len(x.src),
    #     device=opt.device, repeat=False)
    #
    # real_data_clf = torchtext.data.BucketIterator(
    #     dataset=real_data_clf, batch_size=opt.batch_size,
    #     sort=False, sort_within_batch=True,
    #     sort_key=lambda x: len(x.src),
    #     device=opt.device, repeat=False)
    #
    # fake_data_dis = torchtext.data.BucketIterator(
    #     dataset=fake_data_lm, batch_size=opt.batch_size,
    #     sort=False, sort_within_batch=True,
    #     sort_key=lambda x: len(x.src),
    #     device=opt.device, repeat=False)
    fake_data_dis = 1
    # return fake_data_dis, fake_data_lm, real_data_clf, train_clf, test_clf, input_vocab, tgt

def prepare_loss(tgt, opt):
    weight = torch.ones(len(tgt.vocab))
    pad = tgt.vocab.stoi[tgt.pad_token]
    loss = Perplexity(weight, pad)
    if opt.cuda:
        loss.cuda()
    return loss

def prepare_model(opt, vocab_size, tgt):

    dis_hidden_size = opt.hidden_size * opt.n_layers * (2 if opt.bidirectional else 1)
    # Prepare loss
    encoder = EncoderRNN(vocab_size, opt.max_len, opt.hidden_size,
                         bidirectional=opt.bidirectional, n_layers=opt.n_layers, variable_lengths=True)
    decoder = DecoderRNN(vocab_size, opt.max_len, opt.hidden_size * 2 if opt.bidirectional else opt.hidden_size,
                         dropout_p=opt.dropout, n_layers=opt.n_layers, use_attention=False, bidirectional=opt.bidirectional,
                         eos_id=tgt.eos_id, sos_id=tgt.sos_id)
    seq2seq = Seq2seq(encoder, decoder).to(opt.device)

    # gen = Generator(dis_hidden_size, opt.z_size).to(opt.device)
    # encoder_new = EncoderRNN(vocab_size, opt.max_len, opt.hidden_size,
    #                      bidirectional=opt.bidirectional, n_layers=opt.n_layers,variable_lengths=True).to(opt.device)
    #
    #
    # dis_clf = Discriminator(dis_hidden_size, opt.clf_layers).to(opt.device)
    # rnn_clf = RNNclaissfier(encoder_new, dis_clf).to(opt.device)
    #
    # dis_gen = Discriminator(dis_hidden_size, opt.clf_layers).to(opt.device)
    # opt_gen = optim.Adam(gen.parameters(), lr=opt.gen_lr)
    # opt_dis_clf = optim.Adam(rnn_clf.parameters(), lr=opt.dis_dec_lr)
    # opt_dis_gen = optim.Adam(dis_gen.parameters(), lr=opt.dis_gen_lr)

    gen = 1
    opt_gen = 1
    rnn_clf = 1
    opt_dis_clf =1
    dis_gen = 1
    opt_dis_gen = 1
    return seq2seq, gen, opt_gen, rnn_clf, opt_dis_clf, dis_gen, opt_dis_gen

def build_parser():
    parser = argparse.ArgumentParser()
    # data
    parser.add_argument('-fake_data_path', type=str,
                        default='/home/yichuan/course/seq2/data/YelpNYC/text.csv')
    parser.add_argument('-real_data_path', type=str,
                        default='/home/yichuan/course/seq2/data/YelpNYC/real_data.csv')
    parser.add_argument('-train_data_path', type=str,
                        default='/home/yichuan/course/seq2/data/YelpNYC/train_data.csv')
    parser.add_argument('-test_data_path', type=str,
                        default='/home/yichuan/course/seq2/data/YelpNYC/test_data.csv')

    # language model
    parser.add_argument('-max_len', type=int, default=100)
    parser.add_argument('-bidirectional', action='store_false', default=True)
    parser.add_argument('-dropout', type=float, default=0.3)
    parser.add_argument('-hidden_size', type=int, default=128)
    parser.add_argument('-max_word', type=int, default=30000)
    parser.add_argument('-n_layers', type=int, default=2)
    # seq2seq model
    parser.add_argument('-batch_size', type=int, default=20)
    parser.add_argument('-check_point', type=int, default=100)
    parser.add_argument('-print_every', type=int, default=100)
    parser.add_argument('-expt_dir', type=str, default='./experiment')
    parser.add_argument('-teach_force_ratio', type=float, default=0.4)
    parser.add_argument('-resume', action='store_true', default=False)

    # GAN discriminator
    parser.add_argument('-clf_layers', type=int, default=3)
    parser.add_argument('-z_size', type=int, default=128)

    # learning rate
    parser.add_argument('-dis_dec_lr', type=float, default=0.05)
    parser.add_argument('-dis_gen_lr', type=float, default=0.005)
    parser.add_argument('-gen_lr', type=float, default=0.005)

    # epochs
    parser.add_argument('-gan_epoch', type=int, default=10)
    parser.add_argument('-gen_epoch', type=int, default=10)
    parser.add_argument('-clf_epoch', type=int, default=10)
    parser.add_argument('-dis_epoch', type=int, default=1)
    parser.add_argument('-pre_clf_epochs', type=int, default=10)
    parser.add_argument('-lm_epochs', type=int, default=20)
    parser.add_argument('-batch_count', type=int, default=200)

    # cuda
    parser.add_argument('-cuda',  action='store_false')
    return parser

def main(parser):
    opt = parser.parse_args()
    opt.device = torch.device('cuda') if opt.cuda else torch.device('cpu')
    # fake_data_dis, fake_data_lm, real_data_clf, train_clf, test_clf, vocab, tgt = prepare_data(opt)
    fake_data_dis, tgt = prepare_data(opt)
    seq2seq, gen, opt_gen, rnn_claissfier, classifier_opt, dis_gen, opt_dis_gen = \
        prepare_model(opt, len(tgt.vocab), tgt=tgt)

    # pre-train the LM model
    loss_seq = prepare_loss(tgt, opt)
    seq2seq = train_LM(opt, loss_seq, seq2seq, fake_data_dis)

    exit()
    # # pre-train the classify model
    # pre_train_deceptive(rnn_claissfier, classifier_opt, train_clf, opt)
    #
    # # train the generator
    # for epoch in range(opt.gan_epoch):
    #     train_gen(gen, opt_gen, dis_gen, opt)
    #
    # # train the discriminator
    #     train_discriminator(dis_gen, opt_dis_gen, seq2seq, gen, fake_data_dis, opt)
    #
    # # train the classification on simulate data and real review
    # for epoch in range(opt.clf_epoch):
    #     # test the classifier
    #     clf_test(test_clf, rnn_claissfier)
    #     train_classifier(opt, real_data_clf, gen, seq2seq, rnn_claissfier, classifier_opt)

if __name__ == '__main__':
    parser = build_parser()
    main(parser)


import sys
import torch
import numpy
import argparse
import torch.nn as nn

from attack import FWAdampAttackPlus
from attack import EnergyAttack

from model import get_model_for_attack
from eval_model import eval_model_with_attack
from utils import get_test_cifar, get_test_imagenet, get_test_mnist
from utils import print_attack_main_args

from torch.utils.data import DataLoader

torch.set_num_threads(8)

def parse_args():
    parser = argparse.ArgumentParser(
        description='Run Energy Attack benchmark.')
    parser.add_argument(
        '--batch_size', type=int, default=128, metavar='N',
        help='Batch size for training. Default: 128')
    parser.add_argument(
        '--subbatch_size', type=int, default=16, metavar='N',
        help='Sub-batch size for supported attacks (ta, energy). Default: 16')
    parser.add_argument(
        '--epsilon', type=float, default=0.05,
        help='max distance for attack. Default: 0.05')
    parser.add_argument(
        '--perturb_steps', type=int, default=10000,
        help='Maximum queries for attack. Default: 10000')
    parser.add_argument(
        '--model_name', type=str, default="model_vgg16bn",
        help='Model to attack. Default: model_vgg16bn')
    parser.add_argument(
        '--custom_flags', type=str, default=[], action='append',
        help="Customized flags for specific attacks"
    )
    parser.add_argument('--device', type=str, default="cuda:0")
    parser.add_argument(
        '--step_size', type=float, default=None,
        help='Stepsize used in PGD-based attacks.'
        'This argument is not required if using FW or Energy.')
    parser.add_argument(
        '--attacker',
        choices=[
            'fw', 'energy',
        ],
        default='energy')
    parser.add_argument(
        '--dataset', choices=['cifar10', 'imagenet', 'mnist', 'rand_mnist'], default='imagenet',
        help='Dataset on which the attack is performed. Default: ImageNet')
    return parser.parse_args()

def merge_state_dicts(dicts):
    refd = dicts[0]
    result = dict()
    if isinstance(refd, dict):
        for k in refd.keys():
            result[k] = merge_state_dicts([d[k] for d in dicts])
        return result
    else:
        return sum(dicts)

def divide_state_dict(d, n):
    result = dict()
    if isinstance(d, dict):
        for k in d.keys():
            result[k] = divide_state_dict(d[k], n)
        return result
    else:
        return d / n

def get_attacker(attacker, step_size, epsilon, perturb_steps):
    if attacker == 'fw':
        print('Using FW-AdAmp', file=sys.stderr)
        return FWAdampAttackPlus(
            step_size, epsilon, perturb_steps)
    elif attacker == 'energy':
        print('Using Energy Attack', file=sys.stderr)
        return EnergyAttack(
            step_size, epsilon, perturb_steps)

class WrappedModel(nn.Module):
    def __init__(self, wrap, subbatch):
        super().__init__()
        self.wrap = wrap
        self.subbatch = subbatch

    def forward(self, x):
        subbatches = torch.split(x, self.subbatch)
        return torch.cat([self.wrap(sb) for sb in subbatches])

if __name__ == '__main__':
    args = parse_args()
    print_attack_main_args(args)
    device = torch.device(args.device)
    if (
        args.model_name.startswith('model')
    ):
        # load model to attack
        model = get_model_for_attack(args.model_name).to(device)
    else:
        raise ValueError(f'Model {args.model_name} does not exist.')
    model = WrappedModel(model, args.subbatch_size)
    # load attack method
    attack = get_attacker(
        args.attacker, args.step_size, args.epsilon,
        args.perturb_steps
    )
    model.eval()
    if args.dataset == 'cifar10':
        test_loader = get_test_cifar(args.batch_size)
    elif args.dataset == 'mnist':
        test_loader = get_test_mnist(args.batch_size)
    elif args.dataset == 'rand_mnist':
        test_loader = get_test_mnist(1)
        xs = [x for (x,), _ in test_loader]
        ys = numpy.load("rand_mnist.npy")
        ds = list(zip(xs, ys))
        test_loader = DataLoader(ds, args.batch_size)
    else:
        test_loader = get_test_imagenet(args.batch_size)
    # non-targeted attack
    natural_acc, robust_acc, distance = eval_model_with_attack(
        model, test_loader, attack, args.epsilon, device, args.dataset)
    print(
        "Natural Acc: %.5f, Robust acc: %.5f, distance: %.5f" %
        (natural_acc, robust_acc, distance)
    )
    if hasattr(attack, 'consumed_steps'):
        consumed_steps = numpy.array(attack.consumed_steps)
        consumed_steps = consumed_steps[consumed_steps <=
                                        attack.perturb_steps]
        print("Mean", numpy.mean(consumed_steps))
        print("Median", numpy.median(consumed_steps))


"""
SCRIPT TO  WRITE LAMMPS DATA FILES LMP & INP Files
FROM BOSS ZMATRIX
Created on Mon Sep 30 03:31:05 2017
@author: Leela S. Dodda leela.dodda@yale.edu
@author: William L. Jorgensen Lab

REQUIREMENTS:
BOSS (need to set BOSSdir in bashrc and cshrc)
Preferably Anaconda python with following modules
pandas
argparse
numpy
"""

from LigParGenPSP.BOSSReader import bossPdbAtom2Element, ucomb
from LigParGenPSP.BOSSReader import bossElement2Mass, tor_cent
import pickle
import pandas as pd
import numpy as np

def Boss2LammpsLMP(resid, num2typ2symb, Qs, bnd_df, ang_df, tor_df, molecule_data):
    xyz_df = molecule_data.MolData["XYZ"]
    max_mol_size = 50
    prm = open(resid + ".lmp", "w+")
    prm.write("LAMMPS data file Created by - (Written by Leela S. Dodda)\n\n")
    prm.write("%8d atoms\n" % len(Qs))
    prm.write("%8d bonds\n" % len(bnd_df.KIJ))
    prm.write("%8d angles\n" % len(ang_df.K))
    prm.write("%8d dihedrals\n" % len(tor_df[tor_df.TY == "Proper"].index))
    prm.write("%8d impropers\n \n" % len(tor_df[tor_df.TY == "Improper"].index))
    prm.write("%8d atom types\n" % len(Qs))
    prm.write("%8d bond types\n" % len(bnd_df.KIJ))
    prm.write("%8d angle types\n" % len(ang_df.K))
    prm.write("%8d dihedral types\n" % len(tor_df[tor_df.TY == "Proper"].index))
    prm.write("%8d improper types\n \n" % len(tor_df[tor_df.TY == "Improper"].index))
    prm.write(
        "%12.6f %12.6f xlo xhi\n" % (xyz_df.X.min(), xyz_df.X.min() + max_mol_size)
    )
    prm.write(
        "%12.6f %12.6f ylo yhi\n" % (xyz_df.Y.min(), xyz_df.Y.min() + max_mol_size)
    )
    prm.write(
        "%12.6f %12.6f zlo zhi\n" % (xyz_df.Z.min(), xyz_df.Z.min() + max_mol_size)
    )
    # Printing Parameters for ALL BONDS/ANGLES/DIHEDRALS/IMPROPERS/Q/LJ #######
    prm.write("\nMasses\n\n")
    for i in range(len(Qs)):
        prm.write("%8d %10.3f  \n" % (i + 1, float(num2typ2symb[i][4])))
    prm.write("\nPair Coeffs \n\n")
    for i in range(len(Qs)):
        prm.write("%8d%11.3f%11.7f \n" % (i + 1, float(Qs[i][3]), float(Qs[i][2])))
    prm.write("\nBond Coeffs \n\n")
    for i in bnd_df.index:
        prm.write("%8d%11.4f%11.4f \n" % (i + 1, bnd_df.KIJ[i], bnd_df.RIJ[i]))
    prm.write("\nAngle Coeffs \n\n")
    for i in ang_df.index:
        prm.write("%8d%11.3f%11.3f\n" % (i + 1, ang_df.K[i], ang_df.R[i]))
    dihedral_df = tor_df[tor_df.TY == "Proper"]
    dihedral_df.index = range(len(dihedral_df.V1))
    prm.write("\nDihedral Coeffs  \n\n")
    for i, row in dihedral_df.iterrows():
        prm.write(
            "%8d%11.3f%11.3f%11.3f%11.3f \n" % (i + 1, row.V1, row.V2, row.V3, row.V4)
        )
    bndlist = list(bnd_df.UR) + (list(bnd_df.UR))
    improper_df = tor_df[tor_df.TY == "Improper"]
    improper_df.index = range(len(improper_df.V2))
    if len(improper_df.index) > 0:
        prm.write("\nImproper Coeffs \n\n")
        for i, row in improper_df.iterrows():
            prm.write("%8d%11.3f%8d%8d \n" % (i + 1, row.V2 * 0.5, -1, 2))
    # Printing EXPLICITLY ALL BONDS/ANGLES/DIHEDRALS/IMPROPERS/Q/LJ #######
    prm.write("\nAtoms \n\n")
    for i in range(len(xyz_df.index)):
        prm.write(
            "%6d %6d %6d %10.8f %8.3f %8.5f %8.5f\n"
            % (i + 1, 1, i + 1, float(Qs[i][1]), xyz_df.X[i], xyz_df.Y[i], xyz_df.Z[i])
        )
    prm.write("\nBonds \n\n")
    for i in bnd_df.index:
        prm.write(
            "%6d %6d %6d %6d\n" % (i + 1, i + 1, bnd_df.cl1[i] + 1, bnd_df.cl2[i] + 1)
        )
    prm.write("\nAngles \n\n")
    for i in ang_df.index:
        prm.write(
            "%6d %6d %6d %6d %6d\n"
            % (i + 1, i + 1, ang_df.cl1[i] + 1, ang_df.cl2[i] + 1, ang_df.cl3[i] + 1)
        )
    prm.write("\nDihedrals\n\n")
    for i, row in dihedral_df.iterrows():
        prm.write(
            "%6d %6d %6d %6d %6d %6d \n"
            % (i + 1, i + 1, row.I + 1, row.J + 1, row.K + 1, row.L + 1)
        )
    if len(improper_df.index) > 0:
        prm.write("\nImpropers\n\n")
        for row in improper_df.iterrows():
            index, dat = row
            ndata = tor_cent([dat.I, dat.J, dat.K, dat.L], bndlist)
            prm.write(
                "%6d %6d %6d %6d %6d %6d \n"
                % (
                    index + 1,
                    index + 1,
                    ndata[0] + 1,
                    ndata[1] + 1,
                    ndata[2] + 1,
                    ndata[3] + 1,
                )
            )
    return None

def Boss2CharmmTorsion(bnd_df, num2opls, st_no, molecule_data, num2typ2symb):
    dhd = []
    for line in molecule_data.MolData["TORSIONS"]:
        dt = [float(i) for i in line]
        dhd.append(dt)
    dhd = np.array(dhd)
    dhd = dhd  # kcal to kj conversion
    dhd = dhd  # Klammps = Vopls
    dhd_df = pd.DataFrame(dhd, columns=["V1", "V2", "V3", "V4"])
    ats = []
    for line in molecule_data.MolData["ATOMS"][3:]:
        dt = [line.split()[0], line.split()[4], line.split()[6], line.split()[8]]
        dt = [int(d) for d in dt]
        ats.append(dt)
    for line in molecule_data.MolData["ADD_DIHED"]:
        dt = [int(i) for i in line]
        ats.append(dt)
    assert len(ats) == len(
        dhd
    ), "Number of Dihedral angles in Zmatrix and Out file dont match"
    ats = np.array(ats) - st_no
    for i in range(len(ats)):
        for j in range(len(ats[0])):
            if ats[i][j] < 0:
                ats[i][j] = 0
    at_df = pd.DataFrame(ats, columns=["I", "J", "K", "L"])
    # final_df = pd.concat([dhd_df, at_df], axis=1, join_axes=[at_df.index]) backup
    final_df = pd.concat([dhd_df, at_df], axis=1)
    final_df = final_df.reindex(dhd_df.index)

    bndlist = list(bnd_df.UR) + (list(bnd_df.UR))
    final_df["TY"] = [
        "Proper"
        if ucomb(
            list([final_df.I[n], final_df.J[n], final_df.K[n], final_df.L[n]]), bndlist
        )
        == 3
        else "Improper"
        for n in range(len(final_df.I))
    ]
    final_df["TI"] = [num2typ2symb[j][2] for j in final_df.I]
    final_df["TJ"] = [num2typ2symb[j][2] for j in final_df.J]
    final_df["TK"] = [num2typ2symb[j][2] for j in final_df.K]
    final_df["TL"] = [num2typ2symb[j][2] for j in final_df.L]
    final_df["SYMB"] = [
        "-".join(
            [
                num2typ2symb[final_df.I[i]][0],
                num2typ2symb[final_df.J[i]][0],
                num2typ2symb[final_df.K[i]][0],
                num2typ2symb[final_df.L[i]][0],
            ]
        )
        for i in final_df.index
    ]
    if len(final_df.index) > 0:
        final_df["NAME"] = (
            final_df.TI + "-" + final_df.TJ + "-" + final_df.TK + "-" + final_df.TL
        )
    return final_df

def boss2CharmmBond(molecule_data, st_no):
    bdat = molecule_data.MolData["BONDS"]
    bdat["cl1"] = [x - st_no if not x - st_no < 0 else 0 for x in bdat["cl1"]]
    bdat["cl2"] = [x - st_no if not x - st_no < 0 else 0 for x in bdat["cl2"]]
    bnd_df = pd.DataFrame(bdat)
    bnd_df["UF"] = (
        (bnd_df.cl1 + bnd_df.cl2) * (bnd_df.cl1 + bnd_df.cl2 + 1) * 0.5
    ) + bnd_df.cl2
    bnd_df["UR"] = (
        (bnd_df.cl1 + bnd_df.cl2) * (bnd_df.cl1 + bnd_df.cl2 + 1) * 0.5
    ) + bnd_df.cl1
    hb_df = bnd_df.drop(["cl1", "cl2", "UF", "UR"], 1)
    hb_df = hb_df.drop_duplicates()
    return bnd_df

def boss2CharmmAngle(anglefile, num2opls, st_no):
    adat = anglefile
    adat["cl1"] = [x - st_no if not x - st_no < 0 else 0 for x in adat["cl1"]]
    adat["cl2"] = [x - st_no if not x - st_no < 0 else 0 for x in adat["cl2"]]
    adat["cl3"] = [x - st_no if not x - st_no < 0 else 0 for x in adat["cl3"]]
    ang_df = pd.DataFrame(adat)
    ang_df = ang_df[ang_df.K > 0]
    ang_df["TY"] = np.array(
        [
            num2opls[i] + "-" + num2opls[j] + "-" + num2opls[k]
            for i, j, k in zip(ang_df.cl1, ang_df.cl2, ang_df.cl3)
        ]
    )
    return ang_df

def bossData(molecule_data):
    ats_file = molecule_data.MolData["ATOMS"]
    types = []
    for i in enumerate(ats_file):
        types.append([i[1].split()[1], "opls_" + i[1].split()[2]])
    st_no = 3
    Qs = molecule_data.MolData["Q_LJ"]
    assert len(Qs) == len(types), "Please check the at_info and Q_LJ_dat files"
    num2opls = {}
    for i in range(0, len(types)):
        num2opls[i] = Qs[i][0]
    num2typ2symb = {i: types[i] for i in range(len(Qs))}
    for i in range(len(Qs)):
        num2typ2symb[i].append(
            bossPdbAtom2Element(num2typ2symb[i][0]) + num2typ2symb[i][1][-3:]
        )
        num2typ2symb[i].append(bossPdbAtom2Element(num2typ2symb[i][0]))
        num2typ2symb[i].append(bossElement2Mass(num2typ2symb[i][3]))
        num2typ2symb[i].append(Qs[i][0])
    return (types, Qs, num2opls, st_no, num2typ2symb)

def Boss2Lammps(resid, molecule_data):
    types, Qs, num2opls, st_no, num2typ2symb = bossData(molecule_data)
    bnd_df = boss2CharmmBond(molecule_data, st_no)
    ang_df = boss2CharmmAngle(molecule_data.MolData["ANGLES"], num2opls, st_no)
    tor_df = Boss2CharmmTorsion(bnd_df, num2opls, st_no, molecule_data, num2typ2symb)
    Boss2LammpsLMP(resid, num2typ2symb, Qs, bnd_df, ang_df, tor_df, molecule_data)
    return None

def mainBOSS2LAMMPS(resid, clu=False):
    mol = pickle.load(open(resid + ".p", "rb"))
    Boss2Lammps(resid, mol)
    return None